[ { "path": ".git-blame-ignore-revs", "content": "# apply pyink\n40a6e074e5224d733f964be00e21e0a1cb98bd2e" }, { "path": ".github/ISSUE_TEMPLATE/bug_report.md", "content": "---\nname: Bug report\nabout: Create a report to help us improve\ntitle: ''\nlabels: bug\nassignees: ''\n\n---\n\nProvide as much information as possible. At least, this should include a description of your issue and steps to reproduce the problem. If possible also provide a summary of what steps or workarounds you have already tried.\n\n### System information\n- OS Platform and Distribution (e.g., Linux Ubuntu 16.04):\n- Flax, jax, jaxlib versions (obtain with `pip show flax jax jaxlib`:\n- Python version:\n- GPU/TPU model and memory:\n- CUDA version (if applicable):\n\n\n### Problem you have encountered:\n\n\n### What you expected to happen:\n\n\n### Logs, error messages, etc:\n\n\n\n### Steps to reproduce:\nWhenever possible, please provide a *minimal example*. Please consider submitting it as a Colab link.\n" }, { "path": ".github/analytics/README.md", "content": "# Repo Analytics\n\nTo run the repo analytics follow the steps below:\n\n1. You must have a Github token, if you don't have one you can create one by following [this guide](https://docs.github.com/en/enterprise-server@3.4/authentication/keeping-your-account-and-data-secure/creating-a-personal-access-token).\n2. Install the requirements:\n\n```bash\npip install -r .github/analytics/requirements.txt\n```\n3. Run the analytics:\n\n```bash\nGITHUB_TOKEN= \\\npython .github/analytics/get_repo_metrics.py \\\n --repo-owner google \\\n --repo-name flax\n```" }, { "path": ".github/analytics/get_repo_metrics.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport json\nimport os\nfrom datetime import datetime\nfrom collections.abc import Callable\n\nimport matplotlib.dates as mdates\nimport matplotlib.pyplot as plt\nimport pandas as pd\nimport requests\nfrom absl import app, flags\n\ntoken = os.environ['GITHUB_TOKEN']\nendpoint = r'https://api.github.com/graphql'\nheaders = {'Authorization': f'bearer {token}'}\n\n# ------------------------------------------------------------------------------\n# GraphQL\n# ------------------------------------------------------------------------------\n# NOTE: This GraphQL logic was ported and adapted from this script:\n# https://github.com/scientific-python/devstats-data/blob/4c022961abc4ca6061f8719d9c3387e98734b90c/query.py\n# It contains style differences from Google's style guide.\n\n\ndef load_query_from_file(fname, repo_owner, repo_name) -> str:\n with open(fname) as fh:\n query = fh.read()\n # Set target repo from template\n query = query.replace('_REPO_OWNER_', repo_owner)\n query = query.replace('_REPO_NAME_', repo_name)\n return query\n\n\ndef send_query(query, query_type, cursor=None):\n \"\"\"\n Sends a GraphQL to the GitHub API.\n\n No validation is done on the query before sending. GitHub GraphQL is\n supported with the `cursor` argument.\n\n Parameters\n ----------\n query : str\n The GraphQL query to be sent\n query_type : {\"issues\", \"pullRequests\"}\n The object being queried according to the GitHub GraphQL schema.\n Currently only issues and pullRequests are supported\n cursor : str, optional\n If given, then the cursor is injected into the query to support\n GitHub's GraphQL pagination.\n\n Returns\n -------\n dict\n The result of the query (json) parsed by `json.loads`\n\n Notes\n -----\n This is intended mostly for internal use within `get_all_responses`.\n \"\"\"\n # TODO: Expand this, either by parsing the query type from the query\n # directly or manually adding more query_types to the set\n if query_type not in {'issues', 'pullRequests'}:\n raise ValueError(\n \"Only 'issues' and 'pullRequests' queries are currently supported\"\n )\n # TODO: Generalize this\n # WARNING: The cursor injection depends on the specific structure of the\n # query, this is the main reason why query types are limited to issues/PRs\n if cursor is not None:\n cursor_insertion_key = query_type + '('\n cursor_ind = query.find(cursor_insertion_key) + len(cursor_insertion_key)\n query = query[:cursor_ind] + f'after:\"{cursor}\", ' + query[cursor_ind:]\n # Build request payload\n payload = {'query': query}\n response = requests.post(endpoint, json=payload, headers=headers)\n return json.loads(response.content)\n\n\ndef get_all_responses(query, query_type):\n 'Helper function to bypass GitHub GraphQL API node limit.'\n # Get data from a single response\n initial_data = send_query(query, query_type)\n data, last_cursor, total_count = parse_single_query(initial_data, query_type)\n print(f'Retrieving {len(data)} out of {total_count} values...')\n # Continue requesting data (with pagination) until all are acquired\n while len(data) < total_count:\n rdata = send_query(query, query_type, cursor=last_cursor)\n pdata, last_cursor, _ = parse_single_query(rdata, query_type)\n data.extend(pdata)\n print(f'Retrieving {len(data)} out of {total_count} values...')\n print('Done.')\n return data\n\n\ndef parse_single_query(data, query_type):\n \"\"\"\n Parses the data returned by `send_query`\n\n .. warning::\n\n Like `send_query`, the logic here depends on the specific structure\n of the query (e.g. it must be an issue or PR query, and must have a\n total count).\n \"\"\"\n try:\n total_count = data['data']['repository'][query_type]['totalCount']\n data = data['data']['repository'][query_type]['edges']\n last_cursor = data[-1]['cursor']\n except KeyError as e:\n print(data)\n raise e\n return data, last_cursor, total_count\n\n\nclass GithubGrabber:\n \"\"\"\n Pulls down data via the GitHub APIv.4 given a valid GraphQL query.\n \"\"\"\n\n def __init__(self, query_fname, query_type, repo_owner, repo_name):\n \"\"\"\n Create an object to send/recv queries related to the issue tracker\n for the given repository via the GitHub API v.4.\n\n The repository to query against is given by:\n https://github.com//\n\n Parameters\n ----------\n query_fname : str\n Path to a valid GraphQL query conforming to the GitHub GraphQL\n schema\n query_type : {\"issues\", \"pullRequests\"}\n Type of object that is being queried according to the GitHub GraphQL\n schema. Currently only \"issues\" and \"pullRequests\" are supported.\n repo_owner : str\n Repository owner.\n repo_name : str\n Repository name.\n \"\"\"\n self.query_fname = query_fname\n self.query_type = query_type # TODO: Parse this directly from query\n self.repo_owner = repo_owner\n self.repo_name = repo_name\n self.raw_data = None\n self.load_query()\n\n def load_query(self):\n self.query = load_query_from_file(\n self.query_fname, self.repo_owner, self.repo_name\n )\n\n def get(self):\n self.raw_data = get_all_responses(self.query, self.query_type)\n\n\n# ------------------------------------------------------------------------------\n# metrics helpers\n# ------------------------------------------------------------------------------\n\n\ndef _to_datetime(date_str: str) -> datetime:\n return datetime.fromisoformat(date_str.replace('Z', ''))\n\n\ndef _get_issues_features(issues):\n for issue in issues:\n issue = issue['node']\n\n created_at = _to_datetime(issue['createdAt'])\n time_labeled_or_converted = None\n time_issue_closed = None\n\n for event in issue['timelineItems']['edges']:\n event = event['node']\n\n if event['__typename'] in {'LabeledEvent', 'ConvertedToDiscussionEvent'}:\n time_labeled_or_converted = _to_datetime(event['createdAt'])\n\n if event['__typename'] == 'ClosedEvent':\n time_issue_closed = _to_datetime(event['createdAt'])\n\n yield {\n 'created_at': created_at,\n 'time_labeled_or_converted': time_labeled_or_converted,\n 'time_issue_closed': time_issue_closed,\n 'issue_closed': issue['state'] == 'CLOSED',\n }\n\n\ndef _get_pr_features(prs):\n for pr in prs:\n pr = pr['node']\n\n created_at = _to_datetime(pr['createdAt'])\n ready_for_review_at = _to_datetime(pr['createdAt'])\n time_labeled_or_assigned = None\n time_merged_or_closed = None\n time_review = None\n\n if pr['reviews']['nodes']:\n review = pr['reviews']['nodes'][0]\n time_review = _to_datetime(review['createdAt'])\n\n for event in pr['timelineItems']['edges']:\n event = event['node']\n\n if (\n time_labeled_or_assigned is None\n and event['__typename'] == 'LabeledEvent'\n and 'cla:' not in event['label']['name']\n ):\n time_labeled_or_assigned = _to_datetime(event['createdAt'])\n\n if (\n time_labeled_or_assigned is None\n and event['__typename'] == 'AssignedEvent'\n ):\n time_labeled_or_assigned = _to_datetime(event['createdAt'])\n\n if event['__typename'] in {'ClosedEvent', 'MergedEvent'}:\n time_merged_or_closed = _to_datetime(event['createdAt'])\n\n if event['__typename'] == 'ReadyForReviewEvent':\n ready_for_review_at = _to_datetime(event['createdAt'])\n\n yield {\n 'created_at': created_at,\n 'ready_for_review_at': ready_for_review_at,\n 'time_labeled_or_assigned': time_labeled_or_assigned,\n 'time_merged_or_closed': time_merged_or_closed,\n 'time_review': time_review,\n 'pr_closed': pr['state'] != 'OPEN',\n }\n\n\ndef _start_of_month(date: datetime) -> datetime:\n return date.replace(day=1, hour=0, minute=0, second=0, microsecond=0)\n\n\ndef _shift_n_months(date: datetime, n: int) -> datetime:\n month = ((date.month + n - 1) % 12) + 1\n\n # shift to next year if necessary\n if date.month > month:\n date = date.replace(year=date.year + 1)\n\n date = date.replace(month=month)\n\n return date\n\n\ndef _rolling_window(\n df: pd.DataFrame,\n f: Callable[[pd.DataFrame], pd.Series],\n window_size: int = 6,\n step: int = 1,\n) -> pd.DataFrame:\n # start of month of the first issue\n start: datetime = df.iloc[0]['created_at'].replace(\n day=1, hour=0, minute=0, second=0, microsecond=0\n )\n end = _shift_n_months(start, window_size)\n\n last_month = _start_of_month(df.iloc[-1]['created_at'])\n last_month = _shift_n_months(last_month, 1)\n\n rows: list[pd.Series] = []\n while end < last_month:\n row = f(df[(df['created_at'] >= start) & (df['created_at'] < end)])\n row['period_start'] = start\n row['period_end'] = end\n rows.append(row)\n start = _shift_n_months(start, step)\n end = _shift_n_months(end, step)\n\n df = pd.DataFrame(rows)\n df = df[['period_start', 'period_end'] + list(df.columns[:-2])]\n\n return df\n\n\ndef _process_prs(df: pd.DataFrame) -> pd.Series:\n return pd.Series(\n {\n 'pr_response_time': df['pr_response_time'].dt.days.mean(),\n 'pr_resolution_time': df['pr_resolution_time'].dt.days.mean(),\n }\n )\n\n\ndef _process_issues(df: pd.DataFrame) -> pd.Series:\n return pd.Series(\n {\n 'issue_response_time': df['issue_response_time'].dt.days.mean(),\n 'issue_resolution_time': df['issue_resolution_time'].dt.days.mean(),\n }\n )\n\n\n# -----------------------------------------------------------------------------\n# main\n# -----------------------------------------------------------------------------\nFLAGS = flags.FLAGS\nflags.DEFINE_string('repo_owner', 'google', 'User name or organization')\nflags.DEFINE_string('repo_name', 'flax', 'Name of the repository')\n\n\ndef main(_):\n repo_owner: str = FLAGS.repo_owner\n repo_name: str = FLAGS.repo_name\n\n # Download issue data\n issues = GithubGrabber(\n '.github/analytics/issue_activity_since_date.gql',\n 'issues',\n repo_owner=repo_owner,\n repo_name=repo_name,\n )\n issues.get()\n\n df_issues = df_issues0 = pd.DataFrame(\n list(_get_issues_features(issues.raw_data))\n )\n df_issues['issue_response_time'] = (\n df_issues['time_labeled_or_converted'] - df_issues['created_at']\n )\n df_issues['issue_resolution_time'] = (\n df_issues['time_issue_closed'] - df_issues['created_at']\n )\n\n df_issues = _rolling_window(df_issues, _process_issues)\n\n prs = GithubGrabber(\n '.github/analytics/pr_data_query.gql',\n 'pullRequests',\n repo_owner=repo_owner,\n repo_name=repo_name,\n )\n prs.get()\n\n df_prs = df_prs0 = pd.DataFrame(list(_get_pr_features(prs.raw_data)))\n time_response = df_prs[['time_labeled_or_assigned', 'time_review']].min(\n axis=1\n )\n df_prs['pr_response_time'] = time_response - df_prs['ready_for_review_at']\n df_prs['pr_resolution_time'] = (\n df_prs['time_merged_or_closed'] - df_prs['ready_for_review_at']\n )\n\n df_prs = _rolling_window(df_prs, _process_prs)\n\n # get cummulative issues\n df_issues0 = df_issues0.copy()\n df_issues0['number_of_issues'] = 1\n df_issues0['number_of_issues'] = df_issues0['number_of_issues'].cumsum()\n\n # get cummulative prs\n df_prs0 = df_prs0.copy()\n df_prs0['number_of_prs'] = 1\n df_prs0['number_of_prs'] = df_prs0['number_of_prs'].cumsum()\n\n # plot cumulative issues\n plt.figure()\n plt.plot(df_issues0['created_at'], df_issues0['number_of_issues'])\n plt.xlabel('Date')\n plt.ylabel('Number of issues')\n plt.title('Number of issues')\n plt.gca().xaxis.set_major_locator(plt.MaxNLocator(5))\n plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%b %Y'))\n\n # plot cumulative prs\n plt.figure()\n plt.plot(df_prs0['created_at'], df_prs0['number_of_prs'])\n plt.xlabel('Date')\n plt.ylabel('Number of PRs')\n plt.title('Number of PRs')\n plt.gca().xaxis.set_major_locator(plt.MaxNLocator(5))\n plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%b %Y'))\n\n # plot for isssue_response_time\n plt.figure()\n plt.plot(df_issues['period_end'], df_issues['issue_response_time'])\n plt.xlabel('Date')\n plt.ylabel('Issue Response Time (days)')\n plt.title('Issue Response Time')\n plt.gca().xaxis.set_major_locator(plt.MaxNLocator(5))\n plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%b %Y'))\n plt.ylim(0)\n\n # plot for issue_resolution_time\n plt.figure()\n plt.plot(df_issues['period_end'], df_issues['issue_resolution_time'])\n plt.xlabel('Date')\n plt.ylabel('Issue Resolution Time (days)')\n plt.title('Issue Resolution Time')\n plt.gca().xaxis.set_major_locator(plt.MaxNLocator(5))\n plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%b %Y'))\n plt.ylim(0)\n\n # plot for pr_response_time\n plt.figure()\n plt.plot(df_prs['period_end'], df_prs['pr_response_time'])\n plt.xlabel('Date')\n plt.ylabel('Pull Request Response Time (days)')\n plt.title('Pull Request Response Time')\n plt.gca().xaxis.set_major_locator(plt.MaxNLocator(5))\n plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%b %Y'))\n plt.ylim(0)\n\n # plot for pr_resolution_time\n plt.figure()\n plt.plot(df_prs['period_end'], df_prs['pr_resolution_time'])\n plt.xlabel('Date')\n plt.ylabel('Pull Request Resolution Time (days)')\n plt.title('Pull Request Resolution Time')\n plt.gca().xaxis.set_major_locator(plt.MaxNLocator(5))\n plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%b %Y'))\n plt.ylim(0)\n\n # show plots\n plt.show()\n\n\nif __name__ == '__main__':\n app.run(main)\n" }, { "path": ".github/analytics/issue_activity_since_date.gql", "content": "{\n # Queries all the issues in a repo. For each issue, we get some basic data such as\n # the number, state, labels, and title. The most important part is the 'timelineItems'\n # which are the events that happened to the issue, we can use the information about\n # the datetime about certain key events to define some metrics. Note that we are\n # getting more information than is probably needed but its fine for now.\n repository(owner: \"_REPO_OWNER_\", name: \"_REPO_NAME_\") {\n issues(first: 100) {\n totalCount\n edges {\n cursor\n node {\n number\n title\n createdAt\n state\n closedAt\n updatedAt\n url\n labels(first: 100) {\n edges {\n node {\n name\n }\n }\n }\n timelineItems(first: 100, itemTypes: [LABELED_EVENT, CONVERTED_TO_DISCUSSION_EVENT, ISSUE_COMMENT, CLOSED_EVENT]) {\n totalCount\n edges {\n node {\n __typename\n ... on ConvertedToDiscussionEvent {\n createdAt\n }\n ... on IssueComment {\n author {\n login\n }\n createdAt\n }\n ... on ClosedEvent {\n actor {\n login\n }\n createdAt\n }\n ... on LabeledEvent {\n label {\n name\n }\n createdAt\n }\n }\n }\n }\n }\n }\n }\n }\n}\n\n" }, { "path": ".github/analytics/pr_data_query.gql", "content": "query {\n # Queries all the Pull Requests in a repo. For each issue, we get some basic data such as\n # the number, state, reviews, and title. The most important part is the 'timelineItems'\n # which are the events that happened to the issue, we can use the information about\n # the datetime about certain key events to define some metrics. We also use the 'reviews'\n # as indicators for certain metrics. Note that we are getting more information than is\n # probably needed but its fine for now.\n repository(owner:\"_REPO_OWNER_\", name:\"_REPO_NAME_\") {\n pullRequests(first:100) {\n totalCount\n edges {\n cursor\n node {\n number\n state\n title\n createdAt\n author{\n login\n }\n mergedAt\n reviews(first: 100){\n nodes {\n createdAt\n }\n }\n timelineItems(first: 100, itemTypes: [LABELED_EVENT, ASSIGNED_EVENT, MERGED_EVENT, READY_FOR_REVIEW_EVENT, CLOSED_EVENT]) {\n edges {\n node {\n __typename\n ... on ClosedEvent {\n actor {\n login\n }\n createdAt\n }\n ... on LabeledEvent {\n label {\n name\n }\n actor {\n login\n }\n createdAt\n }\n ... on MergedEvent {\n actor {\n login\n }\n createdAt\n }\n ... on ReadyForReviewEvent {\n actor {\n login\n }\n createdAt\n }\n ... on AssignedEvent {\n actor {\n login\n }\n createdAt\n }\n }\n }\n }\n }\n }\n }\n }\n}\n" }, { "path": ".github/analytics/requirements.txt", "content": "pandas\nabsl-py\nrequests\nmatplotlib" }, { "path": ".github/pull_request_template.md", "content": "# What does this PR do?\n\n\n\nFixes # (issue)\n\n## Checklist\n- [ ] This PR fixes a minor issue (e.g.: typo or small bug) or improves the docs (you can dismiss the other checks if that's the case).\n- [ ] This change is discussed in a Github issue/[discussion](https://github.com/google/flax/discussions) (please add a link).\n- [ ] The documentation and docstrings adhere to the [documentation guidelines](https://github.com/google/flax/blob/main/docs/README.md#how-to-write-code-documentation).\n- [ ] This change includes necessary high-coverage tests. (No quality testing = no merge!)\n" }, { "path": ".github/workflows/flax_publish.yml", "content": "# This workflows will upload a Python Package using Twine when a release is created\n# For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries\n\nname: Flax - Build and upload to PyPI\n\non:\n release:\n types: [published]\n\njobs:\n build_package:\n name: Build package\n runs-on: ubuntu-latest\n steps:\n - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n - name: Build package\n run: pipx run build\n - name: List files\n run: ls -l dist/\n - uses: actions/upload-artifact@65c4c4a1ddee5b72f698fdd19549f0f0fb45cf08 # v4.6.0\n with:\n name: distribution\n path: |\n dist/*.tar.gz\n dist/*.whl\n\n upload_pypi:\n name: Release & Upload to PyPI\n # Only publish release to PyPI when a github release is created.\n if: github.event_name == 'release' && github.event.action == 'published'\n needs: build_package\n runs-on: ubuntu-latest\n environment: release\n permissions:\n id-token: write\n steps:\n - uses: actions/download-artifact@fa0a91b85d4f404e444e00e005971372dc801d16 # v4.1.8\n with:\n name: distribution\n path: dist\n - name: List files\n run: ls -l dist/\n - name: Publish package distributions to PyPI\n uses: pypa/gh-action-pypi-publish@76f52bc884231f62b9a034ebfe128415bbaabdfc # v1.12.4\n\n discord_release:\n if: github.repository_owner == 'google'\n runs-on: ubuntu-latest\n steps:\n - name: Get release URL\n id: get-release-url\n run: |\n URL=\"https://github.com/google/flax/releases\"\n echo \"::set-output name=URL::$URL\"\n - name: Get content\n uses: 2428392/gh-truncate-string-action@b3ff790d21cf42af3ca7579146eedb93c8fb0757 # v1.4.1\n id: get-content\n with:\n stringToTruncate: |\n Flax [${{ github.event.release.tag_name }}](<${{ steps.get-release-url.outputs.URL }}>) was just released!\n\n ${{ github.event.release.body }}\n maxLength: 2000\n truncationSymbol: \"...\"\n - name: Discord Webhook Action\n uses: tsickert/discord-webhook@c840d45a03a323fbc3f7507ac7769dbd91bfb164 # v5.3.0\n with:\n webhook-url: ${{ secrets.DISCORD_WEBHOOK_URL }}\n content: ${{ steps.get-content.outputs.string }}\n" }, { "path": ".github/workflows/flax_test.yml", "content": "# This workflow will install Python dependencies, run tests and lint with a variety of Python versions\n# For more information see: https://help.github.com/actions/language-and-framework-guides/using-python-with-github-actions\n\nname: Flax - Test\n\nconcurrency:\n group: ${{ github.workflow }}-${{ github.ref }}\n cancel-in-progress: true\n\non:\n push:\n branches:\n - main\n - 'test_*'\n pull_request:\n branches:\n - main\n\njobs:\n pre-commit:\n name: Test pre-commit hooks\n runs-on: ubuntu-latest\n steps:\n - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n - name: Set up Python\n uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0\n with:\n python-version: '3.11'\n - run: python -m pip install pre-commit\n - uses: actions/cache@1bd1e32a3bdc45362d1e726936510720a7c30a57 # v4.2.0\n with:\n path: ~/.cache/pre-commit\n key: pre-commit-${{ env.pythonLocation }}-${{ hashFiles('.pre-commit-config.yaml', 'pyproject.toml') }}\n - run: pre-commit run --show-diff-on-failure --color=always --all-files\n commit-count:\n name: Check commit count\n runs-on: ubuntu-latest\n steps:\n - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n # We allow at most 5 commits in a branch to ensure our CI doesn't break.\n - name: Check commit count in PR\n if: always()\n shell: bash\n run: |\n set -x\n # $GITHUB_REF is in format `refs/heads/`. We fetch it under\n # the name `commit-count` so we can refer to it below.\n # Do an unshallow fetch so we retrieve all commits (this is necessary\n # because ations/checkout@v2 fetches a shallow copy).\n git fetch origin --unshallow $GITHUB_REF:commit-count\n git fetch origin main\n diff=$(git rev-list --count origin/main...commit-count)\n # $GITHUB_REF adds an additional commit to the commit tree, so $diff is\n # one too high when executing this as a Github Action.\n if (( $diff > 6)); then\n echo \"ERROR! More than 5 commits in PR -- please squash your commits.\"\n url=https://flax.readthedocs.io/en/latest/contributing.html#too-many-commits-in-a-pull-request\n echo \"See $url for help on how to resolve this.\"\n exit 1\n fi\n test-import:\n name: Test import standalone\n runs-on: ubuntu-latest\n strategy:\n matrix:\n python-version: ['3.11', '3.12', '3.13', '3.14']\n steps:\n - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n - name: Set up Python ${{ matrix.python-version }}\n uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0\n with:\n python-version: ${{ matrix.python-version }}\n - uses: astral-sh/setup-uv@887a942a15af3a7626099df99e897a18d9e5ab3a # v5.1.0\n with:\n version: \"0.8.13\"\n - name: Install standalone dependencies only\n run: |\n uv sync\n - name: Test importing Flax\n run: |\n uv run --no-sync python -c \"import flax\"\n\n tests:\n name: Run Tests\n needs: [pre-commit, commit-count]\n runs-on: ubuntu-24.04-16core\n strategy:\n # Make sure to change `github_check_runs` in `copy.bara.sky` if you change the tests here.\n matrix:\n python-version: ['3.11', '3.12', '3.13']\n test-type: [doctest, pytest]\n jax-version: [newest]\n include:\n # keep in sync with internal type checking\n - python-version: '3.12'\n test-type: pytype\n jax-version: 'newest'\n - python-version: '3.12'\n test-type: mypy\n jax-version: 'newest'\n steps:\n - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n - name: Setup uv\n uses: astral-sh/setup-uv@887a942a15af3a7626099df99e897a18d9e5ab3a # v5.1.0\n with:\n version: \"0.8.13\"\n python-version: ${{ matrix.python-version }}\n enable-cache: true\n\n - name: Install dependencies\n run: |\n rm -fr .venv\n uv sync --extra testing --extra docs\n - name: Install JAX\n run: |\n if [[ \"${{ matrix.jax-version }}\" == \"newest\" ]]; then\n uv pip install -U jax jaxlib\n else\n uv pip install \"jax==${{ matrix.jax-version }}\" \"jaxlib==${{ matrix.jax-version }}\"\n fi\n if [[ \"${{ matrix.test-type }}\" == \"pytest\" ]]; then\n uv pip install -U tensorflow-datasets\n fi\n - name: Test with ${{ matrix.test-type }}\n run: |\n if [[ \"${{ matrix.test-type }}\" == \"doctest\" ]]; then\n uv run --no-sync tests/run_all_tests.sh --only-doctest\n elif [[ \"${{ matrix.test-type }}\" == \"pytest\" ]]; then\n uv run --no-sync tests/run_all_tests.sh --only-pytest\n elif [[ \"${{ matrix.test-type }}\" == \"pytype\" ]]; then\n uv run --no-sync tests/run_all_tests.sh --only-pytype\n elif [[ \"${{ matrix.test-type }}\" == \"mypy\" ]]; then\n uv run --no-sync tests/run_all_tests.sh --only-mypy\n else\n echo \"Unknown test type: ${{ matrix.test-type }}\"\n exit 1\n fi\n - name: Upload coverage to Codecov\n if: matrix.test-type == 'pytest'\n uses: codecov/codecov-action@1e68e06f1dbfde0e4cefc87efeba9e4643565303 # v5.1.2\n env:\n CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}\n with:\n file: ./coverage.xml\n # The below step just reports the success or failure of tests as a \"commit status\".\n # This is needed for copybara integration.\n - name: Report success or failure as github status\n if: always()\n shell: bash\n run: |\n status=\"${{ job.status }}\"\n lowercase_status=$(echo $status | tr '[:upper:]' '[:lower:]')\n curl -sS --request POST \\\n --url https://api.github.com/repos/${{ github.repository }}/statuses/${{ github.sha }} \\\n --header 'authorization: Bearer ${{ secrets.GITHUB_TOKEN }}' \\\n --header 'content-type: application/json' \\\n --data '{\n \"state\": \"'$lowercase_status'\",\n \"target_url\": \"https://github.com/${{ github.repository }}/actions/runs/${{ github.run_id }}\",\n \"description\": \"'$status'\",\n \"context\": \"github-actions/Build\"\n }'\n\n # This is a temporary workflow to test flax on Python 3.14 and\n # skipping deps like tensorstore, tensorflow etc\n tests-python314:\n name: Run Tests on Python 3.14\n needs: [pre-commit, commit-count]\n runs-on: ubuntu-24.04-16core\n steps:\n - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n - name: Setup uv\n uses: astral-sh/setup-uv@887a942a15af3a7626099df99e897a18d9e5ab3a # v5.1.0\n with:\n version: \"0.9.2\"\n python-version: \"3.14\"\n activate-environment: true\n enable-cache: true\n\n - name: Install dependencies\n run: |\n rm -fr .venv\n uv sync --extra testing --extra docs\n - name: Test with pytest\n run: |\n export XLA_FLAGS='--xla_force_host_platform_device_count=4'\n find tests/ -name \"*.py\" | grep -vE 'io_test|tensorboard' | xargs pytest -n auto\n\n" }, { "path": ".github/workflows/flaxlib_publish.yml", "content": "name: Flaxlib - Build and upload to PyPI\n\n# for testing only:\non:\n push:\n branches: [main]\n paths: ['flaxlib/**']\n release:\n types: [published]\n\njobs:\n build_wheels:\n if: github.event_name == 'push' && contains(github.event.head_commit.modified, 'flaxlib/')\n name: Build wheels on ${{ matrix.os }}\n runs-on: ${{ matrix.os }}\n strategy:\n matrix:\n # macos-13 is an intel runner, macos-14 is apple silicon\n os: [ubuntu-latest, windows-latest, macos-13, macos-14]\n\n steps:\n - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n\n - uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0\n\n - name: Setup Rust\n uses: actions-rust-lang/setup-rust-toolchain@11df97af8e8102fd60b60a77dfbf58d40cd843b8 # v1.10.1\n\n - name: Install cibuildwheel\n run: python -m pip install cibuildwheel==2.21.0\n\n - name: Build wheels\n run: python -m cibuildwheel --output-dir ./flaxlib/wheelhouse ./flaxlib\n env:\n # rust doesn't seem to be available for musl linux on i686\n CIBW_SKIP: \"*-musllinux_i686\"\n CIBW_ENVIRONMENT: 'PATH=\"$HOME/.cargo/bin:$PATH\" CARGO_TERM_COLOR=\"always\"'\n CIBW_ENVIRONMENT_WINDOWS: 'PATH=\"$UserProfile\\.cargo\\bin;$PATH\"'\n CIBW_BEFORE_BUILD: rustup show\n CIBW_BEFORE_BUILD_LINUX: |\n curl https://sh.rustup.rs -sSf | sh -s -- --default-toolchain=stable --profile=minimal -y &&\n rustup show\n\n - uses: actions/upload-artifact@65c4c4a1ddee5b72f698fdd19549f0f0fb45cf08 # v4.6.0\n with:\n name: cibw-wheels-${{ matrix.os }}-${{ strategy.job-index }}\n path: ./flaxlib/wheelhouse/*.whl\n\n build_sdist:\n if: github.event_name == 'push' && contains(github.event.head_commit.modified, 'flaxlib/')\n name: Build source distribution\n runs-on: ubuntu-latest\n steps:\n - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n\n - name: Setup Rust\n uses: actions-rust-lang/setup-rust-toolchain@11df97af8e8102fd60b60a77dfbf58d40cd843b8 # v1.10.1\n\n - name: Build sdist\n run: pipx run build --sdist flaxlib\n\n - uses: actions/upload-artifact@65c4c4a1ddee5b72f698fdd19549f0f0fb45cf08 # v4.6.0\n with:\n name: cibw-sdist\n path: ./flaxlib/dist/*.tar.gz\n\n upload_pypi:\n if: github.event_name == 'push' && contains(github.event.head_commit.modified, 'flaxlib/')\n name: Upload to PyPI\n needs: [build_wheels, build_sdist]\n runs-on: ubuntu-latest\n permissions:\n id-token: write\n steps:\n - uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0\n with:\n python-version: '3.x'\n - name: Install dependencies\n run: |\n python -m pip install --upgrade pip\n pip install setuptools build wheel twine\n - uses: actions/download-artifact@fa0a91b85d4f404e444e00e005971372dc801d16 # v4.1.8\n with:\n # unpacks all CIBW artifacts into dist/\n pattern: cibw-*\n path: ./flaxlib/dist\n merge-multiple: true\n\n - name: Build and publish\n env:\n TWINE_USERNAME: __token__\n TWINE_PASSWORD: ${{ secrets.FLAXLIB_PYPI_TOKEN }}\n run: |\n twine upload flaxlib/dist/*" }, { "path": ".github/workflows/jax_nightly.yml", "content": "name: CI - with JAX nightly\n\nconcurrency:\n group: ${{ github.workflow }}-${{ github.ref }}\n cancel-in-progress: true\n\non:\n schedule:\n - cron: \"0 12 * * *\" # Daily at 12:00 UTC\n workflow_dispatch: # allows triggering the workflow run manually\n pull_request: # Automatically trigger on pull requests affecting this file\n branches:\n - main\n paths:\n - '**workflows/jax_nightly.yml'\n\njobs:\n jax-nightly:\n runs-on: ubuntu-latest\n permissions:\n contents: read\n issues: write # for failed-build-issue\n strategy:\n fail-fast: false\n matrix:\n python-version: [\"3.11\"]\n steps:\n - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n - name: Set up Python ${{ matrix.python-version }}\n id: setup_python\n uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0\n with:\n python-version: ${{ matrix.python-version }}\n - name: Setup uv\n uses: astral-sh/setup-uv@887a942a15af3a7626099df99e897a18d9e5ab3a # v5.1.0\n with:\n version: \"0.8.13\"\n - name: Install dependencies\n run: |\n uv sync --extra testing --extra docs\n - name: Install JAX\n run: |\n uv pip install -U --pre jax jaxlib -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/\n - name: Run test suite\n if: success()\n run: |\n uv run tests/run_all_tests.sh --only-pytest\n - name: Notify failed build\n uses: jayqi/failed-build-issue-action@1a893bbf43ef1c2a8705e2b115cd4f0fe3c5649b # v1.2.0\n if: failure() && github.event.pull_request == null\n with:\n github-token: ${{ secrets.GITHUB_TOKEN }}" }, { "path": ".gitignore", "content": "*~\n\\#*\\#\n*.pyc\n.tfds\n.DS_Store\ndist/\nbuild/\n*.egg-info\n*.rej\n.pytype\n.vscode/*\n/.devcontainer\ndocs*/**/_autosummary\ndocs*/_build\ndocs*/**/tmp\nflaxlib_src/build\nflaxlib_src/builddir\nflaxlib_src/dist\nflaxlib_src/subprojects\n.venv\nvenv/\nvenv.bak/\n\n# used by direnv\n.envrc\n\n# uv\nuv.lock\n\n# custom\n/tmp-files\n.agent/rules/flax.md\n.github/copilot-instructions.md\n.env\n" }, { "path": ".pre-commit-config.yaml", "content": "# Install the pre-commit hooks below with\n# 'pre-commit install'\n\n# Auto-update the version of the hooks with\n# 'pre-commit autoupdate'\n\n# Run the hooks on all files with\n# 'pre-commit run --all'\n\nrepos:\n- repo: https://github.com/mwouts/jupytext\n rev: v1.13.8\n hooks:\n - id: jupytext\n args: [--sync]\n# diable pyink for now\n# - repo: https://github.com/google/pyink\n# rev: 23.5.0\n# hooks:\n# - id: pyink\n- repo: https://github.com/pre-commit/pre-commit-hooks\n rev: v5.0.0\n hooks:\n - id: check-toml\n - id: trailing-whitespace\n exclude: ^docs.*\\.md$\n- repo: https://github.com/kynan/nbstripout\n rev: 0.6.1\n hooks:\n - id: nbstripout\n exclude: ^examples/.*\n args: [\n --keep-output,\n --keep-count,\n --extra-keys,\n \"cell.metadata.executionInfo cell.metadata.id metadata.kernelspec metadata.vscode metadata.colab cell.metadata.executionInfo.user cell.metadata.executionInfo.user_tz cell.metadata.colab\",\n ]\n- repo: https://github.com/astral-sh/ruff-pre-commit\n # Ruff version.\n rev: v0.1.3\n hooks:\n # Run the Ruff linter.\n - id: ruff\n args: [--fix, --exit-non-zero-on-fix]\n # Disable Ruff formatter for now\n # # Run the Ruff formatter.\n # - id: ruff-format\n- repo: https://github.com/asottile/pyupgrade\n rev: v3.16.0\n hooks:\n - id: pyupgrade\n args: [--py310-plus]\n" }, { "path": ".readthedocs.yml", "content": "# deprecated" }, { "path": "AUTHORS", "content": "# This is the list the Flax authors for copyright purposes.\n#\n# This does not necessarily list everyone who has contributed code, since in\n# some cases, their employer may be the copyright holder. To see the full list\n# of contributors, see the revision history in source control.\nGoogle LLC\n" }, { "path": "CHANGELOG.md", "content": "Changelog\n----------\n\nvNext\n------\n(Add your change to a random empty line to avoid merge conflicts)\n-\n-\n-\n-\n- removed GeGLU simplistic activation, it should be implemented manually.\n-\n-\n-\n- removed FLAX_LAZY_RNG flag support for old non-lazy PRNG derivation mode\n-\n-\n-\n-\n-\n-\n-\n-\n-\n-\n\n0.8.2\n-----\n- fixed rng guide outputs by @chiamp in https://github.com/google/flax/pull/3685\n- enforce mask kwarg in norm layers by @chiamp in https://github.com/google/flax/pull/3663\n- added kwargs to self.param and self.variable by @chiamp in https://github.com/google/flax/pull/3675\n- added nnx normalization tests by @chiamp in https://github.com/google/flax/pull/3689\n- added NNX init_cache docstring example by @chiamp in https://github.com/google/flax/pull/3688\n- added nnx attention equivalence test by @chiamp in https://github.com/google/flax/pull/3687\n- Fix bug that assumed frozen-dict keys were strings. by @copybara-service in https://github.com/google/flax/pull/3692\n- added nnx rmsnorm by @chiamp in https://github.com/google/flax/pull/3691\n- updated nnx compute_stats by @chiamp in https://github.com/google/flax/pull/3693\n- fixed intercept_methods docstring by @chiamp in https://github.com/google/flax/pull/3694\n- [nnx] Add Sphinx Docs by @cgarciae in https://github.com/google/flax/pull/3678\n- Fix pointless docstring example of nn.checkpoint / nn.remat. by @levskaya in https://github.com/google/flax/pull/3703\n- added default params rng to .apply by @chiamp in https://github.com/google/flax/pull/3698\n- [nnx] add partial_init by @cgarciae in https://github.com/google/flax/pull/3674\n- make make_rng default to 'params' by @chiamp in https://github.com/google/flax/pull/3699\n- Add SimpleCell. by @carlosgmartin in https://github.com/google/flax/pull/3697\n- fix Module.module_paths docstring by @cgarciae in https://github.com/google/flax/pull/3709\n- Guarantee the latest JAX version on CI by @cgarciae in https://github.com/google/flax/pull/3705\n- Replace deprecated API `jax.tree.map` by @copybara-service in https://github.com/google/flax/pull/3715\n- Use `jax.tree_util.tree_map` instead of deprecated `jax.tree.map`. by @copybara-service in https://github.com/google/flax/pull/3714\n- [nnx] simplify readme by @cgarciae in https://github.com/google/flax/pull/3707\n- [nnx] add demo.ipynb by @cgarciae in https://github.com/google/flax/pull/3680\n- Fix Tabulate's compute_flops by @cgarciae in https://github.com/google/flax/pull/3721\n- [nnx] simplify TraceState by @cgarciae in https://github.com/google/flax/pull/3724\n- Add broadcast of `strides` and `kernel_dilation` to `nn.ConvTranspose` by @IvyZX in https://github.com/google/flax/pull/3731\n- [nnx] Fix State.__sub__ by @cgarciae in https://github.com/google/flax/pull/3704\n- [nnx] always fold_in on fork + new ForkedKeys return type by @cgarciae in https://github.com/google/flax/pull/3722\n- [nnx] explicit Variables by @cgarciae in https://github.com/google/flax/pull/3720\n- Improves fingerprint definition for Modules in nn.jit. by @copybara-service in https://github.com/google/flax/pull/3736\n- Flax: avoid key reuse in tests by @copybara-service in https://github.com/google/flax/pull/3740\n- added Einsum layer by @chiamp in https://github.com/google/flax/pull/3710\n- nn.jit: automatic fingerprint definition for dataclass attributes by @cgarciae in https://github.com/google/flax/pull/3737\n- [NVIDIA] Use custom grad accumulation for FP8 params by @kaixih in https://github.com/google/flax/pull/3623\n- removed nnx dataclass by @chiamp in https://github.com/google/flax/pull/3742\n- [nnx] cleanup graph_utils by @cgarciae in https://github.com/google/flax/pull/3728\n- Fix doctest and unbreak head by @IvyZX in https://github.com/google/flax/pull/3753\n- [nnx] add pytree support by @cgarciae in https://github.com/google/flax/pull/3732\n- fixed intercept_methods docstring by @chiamp in https://github.com/google/flax/pull/3752\n- Add ConvLSTMCell to docs. by @carlosgmartin in https://github.com/google/flax/pull/3712\n- [nnx] remove flagslib by @cgarciae in https://github.com/google/flax/pull/3733\n- Fix tests after applying JAX key-reuse checker. See: by @copybara-service in https://github.com/google/flax/pull/3748\n\n0.8.1\n-----\n- Added default collection in `make_rng`.\n- Added `InstanceNorm` and renamed `channel_axes` to `feature_axes`.\n- Added norm equivalence tests.\n- Added `Module.module_paths` and doc.\n- make `Sequential.__call__` compact.\n- Added `nn.compact_name_scope` v3.\n- Add explicit control over frozen/slots setting in `flax.struct.dataclass`.\n- Replacing `jax.tree_util.tree_map` with mapping over leafs.\n- Fixed docs and docstrings.\n\n0.8.0\n-----\n- Added [NNX](https://github.com/google/flax/tree/main/flax/nnx#nnx), a neural network library for JAX that provides a simple yet powerful module system that adheres to standard Python semantics. Its aim is to combine the robustness of Linen with a simplified, Pythonic API akin to that of PyTorch.\n- Added `nn.compact_name_scope` decorator that enables methods to act as compact name scopes as with regular Haiku methods. This makes porting Haiku code easier.\n- Add copy() method to Module. This is a user-friendly version of the internal clone() method with better\n defaults for common use cases.\n- Added [`BatchApply`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/layers.html#batchapply) class.\n- Added `sow_weights` option in attention layer.\n- Added [`MultiHeadAttention`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/_autosummary/flax.linen.MultiHeadAttention.html) alias.\n- Added kwargs support for `nn.jit`.\n- Deprecated `normalize` activation function, in favor of `standardize`.\n- Added `GeGLU` activation function.\n- Added `Enum` support for `tabulate` function.\n- Added simple argument-only lifted `nn.grad` function.\n\n0.7.5\n-----\n- Report forward and backward pass FLOPs of modules and submodules in `linen.Module.tabulate` and `summary.tabulate` (in new `flops` and `vjp_flops` table columns). Pass `compute_flops=True` and/or `compute_vjp_flops=True` to include these columns.\n- Re-factored `MultiHeadDotProductAttention`'s call method signature, by adding\n`inputs_k` and `inputs_v` args and switching `inputs_kv`, `mask` and `determistic`\nto keyword arguments. See more details in [#3389](https://github.com/google/flax/discussions/3389).\n- Use new typed PRNG keys throughout flax: this essentially involved changing\n uses of `jax.random.PRNGKey` to `jax.random.key`.\n (See [JEP 9263](https://github.com/jax-ml/jax/pull/17297) for details).\n If you notice dispatch performance regressions after this change, be sure\n you update `jax` to version 0.4.16 or newer.\n- Added `has_improved` field to EarlyStopping and changed the return signature of\n`EarlyStopping.update` from returning a tuple to returning just the updated class.\nSee more details in [#3385](https://github.com/google/flax/pull/3385)\n\n0.7.4\n-----\nNew features:\n- Add QK-normalization to MultiHeadDotProductAttention\n- Allow apply's method argument to accept submodules\n- Add module path to nn.module.\n- [JAX] Generate new type of PRNG keys\n\nBug fixes:\n- Directly call original method if method interceptor stack is empty.\n- fix stackoverflow when loading pickled module\n- Improve kw_only_dataclass.\n- Allow pass-through implementation of state dict\n- Promote dot_general injections from a function to a module.\n\n0.7.2\n-----\nNew features:\n- make `flax.core.copy` `add_or_replace` optional\n- Add `use_fast_variance` option to `GroupNorm` and `BatchNorm` to allow disabling it.\n\nBug fixes:\n- Use `field_specifiers` instead of `field_descriptors` in `@dataclass_transform`.\n- Fix `nn.Module` typing.\n- [JAX] Replace uses of `jax.experimental.pjit.with_sharding_constraint` with `jax.lax.with_sharding_constraint`.\n\n0.7.1\n-----\nBreaking changes:\n- Migrating Flax from returning FrozenDicts to returning regular dicts. More details can be found in this [announcement](https://github.com/google/flax/discussions/3191)\n\nNew features:\n- Use pyink\n- added dict migration guide to index\n- add scan over layers section\n- Expose options to customize rich.Table\n- add support for initializing carry variables in scan\n- Let Flax-Orbax to not port the shape of `target` arrays when they port the `target` shardings.\n\nBug fixes:\n- Use import `orbax.checkpoint` which is a better import pattern.\n- Use import `orbax.checkpoint as ocp` to avoid the verbosity of using 'orbax.checkpoint` every time.\n- [linen] Add alternative, more numerically stable, variance calculation to `LayerNorm`.\n- [linen] Minor cleanup to normalization code.\n- Fix norm calculation bug for 0-rank arrays.\n- [JAX] Remove references to jax.config.jax_array.\n- [linen] Use `stack` instead of `concatenate` in `compute_stats`, to handle scalar stats case.\n- [linen] More minor cleanup in normalization `compute_stats`.\n- Fix warnings from atari gym.\n- Refactor TypeHandler to operate over batches of values, rather than individual ones. This allows more flexibility for implementations that may operate more efficiently on batches.\n- Fix carry slice logic\n- make flax_basics guide use utility fns\n- Fix checkpointing guide error at head\n- Improve scan docs\n\n0.7.0\n-----\n- RNNCellBase refactor.\n\n0.6.11\n-----\n- Set Orbax-as-backend to be the default checkpointing method.\n- Fix setup trigger issue under sharing and transforms.\n- Add collection to self.scope.reserve(name, col) so that sow works with the same name in different collections.\n- Minor improvements for Sequential.\n- Improve the error message in MultiHeadDotProductAttention.\n- Allow manually specifying the rng key for Dropout.\n- RNN refactor.\n- fixed missing separator for rng fold in.\n\n0.6.10\n-----\n- Rudimentary quantization support: some layers can be parametrized with custom dot_general and conv_general_dilated.\n\n0.6.9\n-----\n- Depend on `orbax-checkpoint` package instead of `orbax`.\n- Refactored setup scripts to `project.toml`.\n- Added pretty_repr utility fn.\n- Fix get_partition_spec on replicated array.\n- Updates imagenet.ipynb to use GPU Colab runtime, and fixed config.\n- Upgrade checkpointing code to `jax.sharding`, and with more warnings.\n\n0.6.8\n-----\n- The automatic checkpoint migration was temporarily rolled back due to legacy compatibility issues.\n - We still recommend you to use the [upgrade guide](https://flax.readthedocs.io/en/latest/guides/orbax_upgrade_guide.html) and migrate completely to the Orbax API to ensure stability.\n - Or alternatively, add `flax.config.update('flax_use_orbax_checkpointing', True)` to your project to avoid being impacted by the automatic migration process.\n- Added utility functions to frozen_dict api.\n- Migrated Flax away from `register_keypaths`.\n- Fixes kwargs in convert_to_graphs_tuple_fn.\n- Fixed examples in a few ways:\n - Bumped the TF version\n - Used latest checkpoint formats\n - Other misc fixes.\n\n0.6.7\n-----\n- New checkpoints will be saved using Orbax! Please check out [upgrade guide](https://flax.readthedocs.io/en/latest/guides/orbax_upgrade_guide.html) and consider migrating completely to the Orbax API.\n - You could `flax.config.update('flax_use_orbax_checkpointing', False)` to temporarily disable this migration, but note that Flax legacy checkpointing will be removed 3 months from Mar 10, 2023.\n- Migrating `FrozenDict` to regular dict: utility functions now work on both.\n- Migrated Flax dataclass and `FrozenDict` to JAX pytree keypath API.\n- Fixed pytype and improved typing for `Module`\n- Fixed up uses of PyTree and PyTreeDef types.\n\n0.6.6\n-----\n- 0.6.5 was yanked so this release contains all that was in 0.6.5 as well.\n- Migrated regular dict to FrozenDict, currently controlled by a flag.\n- Refactored and separate out name relaxation policy changes.\n- Added RMS normalization layer.\n\n0.6.5\n-----\n- Added logical partitioning helpers for using pjit with Flax.\n- Add ``Module.lazy_init`` to avoid compute during Module initialization.\n\n0.6.4\n-----\nNew features:\n- Our [ReadTheDoc](https://flax.readthedocs.io/en/latest/index.html) site is a lot more organized now! More improvements on the way.\n- Flax auto-SPMD parallelism API to work seamlessly with `jax.pjit`: https://flax.readthedocs.io/en/latest/guides/flax_on_pjit.html\n- Added new `zeros_init` and `ones_init` initializers.\n- Adds standardize initializer.\n- Allowed specifying method as a string.\n- Allowed runtime overwrite of `flax.config` flags.\n\nBug fixes:\n- Added missing `dataclass.fields` from `__repr__`.\n- Renamed ConvLSTM to ConvLSTMCell.\n- Fix some tiny inconsistencies between scope.py and module.py.\n- Improved many many docstrings, comments and error messages.\n\n0.6.3\n-----\nNew features:\n- Flax checkpointing now uses [Orbax](https://github.com/google/orbax) for more flexiblity and features.\n- Added support for python 3.10 and removed support for 3.7.\n\nBug fixes:\n- Fixed rng generation in DenseGeneral init.\n- Improved support for Mac M1 chip.\n- Bumped package versions for a bunch of examples.\n- Improved many docstrings and error messages.\n\n0.6.2\n-----\nNew features:\n- Add rng_collection argument to Dropout.\n- Fix flax.linen.stochastic.Dropout.\n- Add flag allow_partial_mpa_restoration in checkpointing.\n- Use `gfile.remove` for files because it doesn't work on GCS files.\n- Added guides for: Flax the Sharp Bits, Checkpointing, Extracting Gradients\n- Improved existed documentation pages.\n- Improved errors, error messages and tests.\n- Removed codebase's trailing whitespaces.\n\nBug fixes:\n- Fixes launch_gce.sh with imagenet example.\n- Properly report AttributeErrors from descriptors.\n- Fixes usages of `pmap`.\n- Return None if no _parent_ref is set.\n- Cap dynamic scale to float32 max.\n- no-op when double wrapping with struct.dataclass.\n- Allow variable_with_axes to have empty axes when axes is set to an empty tuple.\n- Don't create reference cycles among Modules.\n\n0.6.1\n-----\n- Adds axis_name and axis_index_groups to LayerNorm and GroupNorm. by @copybara-service in [#2402](https://github.com/google/flax/pull/2402)\n- Plumb spmd_axis_name through transforms.vmap through to JAX vmap by @copybara-service in [#2398](https://github.com/google/flax/pull/2398)\n- Support multiple inputs in flax lifted vjp/custom_vjp by @copybara-service in [#2399](https://github.com/google/flax/pull/2399)\n- Improve tabulate by @cgarciae in [#2316](https://github.com/google/flax/pull/2316)\n- Add path_aware_map function by @cgarciae in [#2371](https://github.com/google/flax/pull/2371)\n- Add static_argnums to nn.checkpoint by @cgarciae in [#2457](https://github.com/google/flax/pull/2457)\n- Adding \"count_include_pad\" argument to flax.linen.pooling.avg_pool by @dslisleedh in [#2451](https://github.com/google/flax/pull/2451)\n- Add perturb() to allow capturing intermediate gradients by @IvyZX in [#2476](https://github.com/google/flax/pull/2476)\n\n0.6.0\n-----\n\n- Removed deprecated optimizers in `flax.optim` package.\n- Moved `flax.optim.dynamic_scale` to `flax.training.dynamic_scale`.\n- Switched to using `jax.named_scope` for all profile naming, cut some pointless\n stack traces out.\n\n0.5.3\n-----\nNew features:\n- Added `nn.switch` as a lifted version of `jax.lax.switch`.\n- Added a method for detecting the use of \"init\" functions.\n- Added checkpointing support for `jax.experimental.GlobalDeviceArray`, a useful array type for multiprocess/multihost computing.\n- Added async option to `save_checkpoints()` on single-process scenario.\n- Improved documentation pages.\n\nBug fixes:\n- Fixed variable aliasing in put_variable\n- Fixed missing passthrough of nn.scan unroll arg\n- Fixed the MNIST example\n\n0.5.2\n-----\n- Fixes missing PyYAML dependency.\n\n0.5.1\n-----\nNew features:\n- Added `nn.tabulate` and `Module.tabulate` to generate rich representations of the network structure.\n\n0.5.0\n-----\n- Added `flax.jax_utils.ad_shard_unpad()` by @lucasb-eyer\n- Implemented [default dtype FLIP](https://github.com/google/flax/blob/main/docs/flip/1777-default-dtype.md).\n This means the default dtype is now inferred from inputs and params rather than being hard-coded to float32.\n This is especially useful for dealing with complex numbers because the standard Modules will no longer truncate\n complex numbers to their real component by default. Instead the complex dtype is preserved by default.\n\n\nBug fixes:\n- Fix support for JAX's experimental_name_stack.\n\nBreaking changes:\n- In rare cases the dtype of a layer can change due to [default dtype FLIP](https://github.com/google/flax/blob/main/docs/flip/1777-default-dtype.md). See the \"Backward compatibility\" section of the proposal for more information.\n\n0.4.2\n-----\n\nNew features:\n- Add lifted conditional `nn.cond`.\n- Improved error messages: parameters not found, loading checkpoints.\n- Replace `jax.tree_multimap` (deprecated) with `jax.tree.map`.\n- Add the \"Module Lifecycle\" design note.\n- Add support for JAX dynamic stack-based named_call\n\nBug fixes:\n- Handle rate==1.0 edgecase in Dropout.\n- Fix bug where Linen Module state is reused.\n- Bug fixes and generalizations of nn.partitioning API.\n\n0.4.1\n-----\n\nNew features:\n- Added locally-connected (unshared CNN) layer `flax.linen.ConvLocal`.\n- Improved seq2seq example: Factored our model and input pipeline code.\n- Added Optax update guide and deprecated `flax.optim`.\n- Added `sep` argument to `flax.traverse_util.flatten_dict()`.\n- Implemented Sequential module, in `flax.linen.combinators`.\n\n0.4.0\n------\nBreaking changes:\n- flax.deprecated.nn is removed. Please pin to flax==0.3.6 if you are still using it.\n- PixelCNN++ example is removed. It was not working well on TPU.\n- linen Normalization layers no longer downcast double and complex floats tofloat32\n when computing the mean and variance.\n\nNew features:\n- Added `flax.linen.custom_vjp` for custom derivatives inside a `Module`.\n- Add `param_dtype` attribute to standard Linen Modules for specifying parameter dtypes.\n\n\n0.3.6\n------\nBreaking changes:\n- Move `flax.nn` to `flax.deprecated.nn`.\n\nNew features:\n- Add experimental checkpoint policy argument. See `flax.linen.checkpoint`\n- Add lifted versions of jvp and vjp.\n- Add lifted transformation for mapping variables. See `flax.linen.map_variables`.\n\n\n0.3.5\n------\n\nBreaking changes:\n - You can no longer pass an int as the `kernel_size` for a `flax.linen.Conv.\n Instead a type error is raised stating that\n a tuple/list should be provided. Stride and dilation arguments do support broadcasting a single int value now because this is not\n ambigious when the kernel rank is known.\n - `flax.linen.enable_named_call` and `flax.linen.disable_named_call` now work anywhere instead of only affecting Modules constructed after the enable/disable call. Additionally, there is now `flax.linen.override_named_call` that provided a context manager to locally disable/enable named_call.\n - NamedTuples are no longer converted to tuples on assignment to a `linen.Module`.\n\nNew features:\n - Flax internal stack frames are now removed from exception state traces.\n - Added `flax.linen.nowrap` to decorate method that should not be transformed\n because they are stateful.\n - Flax no longer uses implicit rank broadcasting. Thus, you can now use Flax with `--jax_numpy_rank_promotion=raise`.\n\nBugfixes:\n - linen Modules and dataclasses made with `flax.struct.dataclass` or `flax.struct.PyTreeNode` are now correctly recognized as dataclasses by static analysis tools like PyLance. Autocomplete of constructors has been verified to work with VSCode.\n - Fixed a bug in FrozenDict which didn't allow copying dicts with reserved names.\n - Fix the serialization of named tuples. Tuple fields are no longer stored in the state dict and the named tuple class is no longer recreated ([bug](https://github.com/google/flax/issues/1429)).\n - Mixed precision training with float16 now works correctly with the attention layers.\n - auto-generated linen Module `__hash__`, `__eq__`, `__repr__` no longer fail by default on non-init attributes.\n\n\n\n0.3.4\n------\n\nPossibly breaking changes:\n - When calling `init` the 'intermediates' collection is no longer mutable.\n Therefore, intermediates will no longer be returned from initialization by default.\n - Don't update batch statistics during initialization.\n - When not using any non-determinism (e.g., dropout), it is not longer necessary to specify the `deterministic` argument in `MultiHeadDotProductAttention`.\n\n\nOther changes:\n - Rewrote various examples to use Optax instead of Flax optimizers (e.g., Imagenet, SST2).\n - Added an NLP text classification example (on the SST-2 dataset) to\n [`examples/sst2`](https://github.com/google/flax/tree/main/examples/sst2).\n that uses a bidirectional LSTM (BiLSTM) to encode the input text.\n - Added `flax.training.train_state` to simplify using Optax optimizers.\n - `mutable` argument is now available on `Module.init` and `Module.init_with_outputs`\n - Bug fix: Correctly handle non-default parameters of Linen Modules with nested inheritance.\n - Expose `dot_product_attention_weights`, allowing access to attention weights.\n - `BatchNorm` instances will behave correctly during init when called multiple times.\n - Added a more extensive \"how to contribute\" guide in `contributing.md`.\n - Add proper cache behavior for [`lift.jit`](https://flax.readthedocs.io/en/latest/_autosummary/flax.linen.jit.html#flax.linen.jit),\nfixing cache misses.\n - Fix bug in Embed layer: make sure it behaves correctly when embedding is np.array.\n - Fix `linen.Module` for deep inheritance chains.\n - Fix bug in DenseGeneral: correctly expand bias to account for batch & noncontracting dimensions.\n - Allow Flax lifted transforms to work on partially applied Modules.\n - Make `MultiOptimizer` use `apply_gradient` instead of `apply_param_gradient`.\n\n0.3.3\n------\n\nPossible breaking changes:\n - Bug Fix: Disallow modifying attributes in Modules after they are initialized.\n - Raise an error when saving a checkpoint which has a smaller step than the\n latest checkpoint already saved.\n - MultiOptimizer now rejects the case where multiple sub optimizers update the\n same parameter.\n\nOther changes:\n - Added custom error classes to many Linen errors. See:\n https://flax.readthedocs.io/en/latest/flax.errors.html\n - Adds `Module.bind` for binding variables and RNGs to an interactive Module.\n - Adds `nn.apply` and `nn.init` for transforming arbitrary functions that take a `linen.Module` as their first argument.\n - Add option to overwrite existing checkpoints in `save_checkpoint`.\n - Remove JAX omnistaging check for forward compatibility.\n - Pathlib compatibility for checkpoint paths.\n - `is_leaf` argument in `traverse_util.flatten_dict`\n\n0.3.2\n------\n\n`flax.nn` deprecation message no longer appears if you import flax directly.\n\nNOTE: You must now explicitly import `flax.nn` if you want to use the old\n pre-Linen `flax.nn.Module`.\n\n0.3.1\n------\n\nMany improvements to Linen, and the old `flax.nn` is officially deprecated!\n\nNotably, there's a clean API for extracting intermediates from modules\ndefined using `@nn.compact`, a more ergonomic API for using Batch Norm and Dropout in modules\ndefined using `setup`, support for `MultiOptimizer` with Linen, and multiple safety, performance\nand error message improvements.\n\nPossible breaking changes:\n - Call setup lazily. See #938 for motivation and more details.\n - Linen `Module` instances are now frozen after `setup` has been called.\n Previously mutations after setup could be dropped silently. Now the stateless requirement\n is enforced by raising a TypeError in `__setattr__` after `setup`.\n - Pytrees of dicts and lists are transformed into FrozenDict and tuples during\n attribute assignment.\n This avoids undetected submodules and inner state.\n - Bug Fix `flax.core.apply` and `Module.apply`. Now it returns a tuple\n containing the output and a frozen empty\n collection when `mutable` is specified as an empty list.\n - `broadcast_dims` is now an attribute to `Dropout` instead of a `__call__`\n argument.\n - `use_running_average` and `deterministic` no longer have a default. They\n should be passed explicitly\n - Bug Fix `Scope.variable` mutability check, before a variable could only be\n initialized if the 'params' collection was mutable.\n\nOther Improvements:\n - Re-introduced the `lm1b` language modeling example\n - Recognizes batch free inputs in pooling layers. (for use with vmap)\n - Add Adadelta optimizer\n - Fully deprecate all \"pre-Linen\" `flax.nn` classes and methods.\n - Some Module arguments can now be passed either as dataclass attribute or\n as argument to `__call__`. See [design note](https://flax.readthedocs.io/en/latest/guides/arguments.html)\n - Add `sow` method to `Module` and `capture_intermediates` argument to `Module.apply`.\n See [howto](https://flax.readthedocs.io/en/latest/howtos/extracting_intermediates.html) for usage patterns.\n - Support passing in modules directly as attributes to other modules, and\n deal with them correctly both in top-level modules and in submodules.\n - Don't require the `variable` argument to `Module.apply` to be a FrozenDict\n - Add support for dict/FrozenDict when using `ModelParamTraversal`\n As a result `MultiOptimizer` can be used properly with linen modules.\n - Added OptimizedLSTM: ~33% faster than the original LSTM when using <=1024 units\n - Fix dtype handling for Adam and LAMB optimizers in 64bit mode.\n - Added `is_mutable()` method to `Variable` and `is_mutable_collection()` to `flax.linen.Module`.\n - Add `axis_name` arg to `flax.linen.vmap`\n - Enable broadcast in `flax.linen.scan`\n - Fix behavior when inner module classes were defined in another module\n - Add automatic giant array chunking in msgpack checkpoints.\n - Log info message when a checkpoint is not found in the directory.\n\nv0.3\n-----\nLinen is now out of Alpha (flax.nn is being deprecated)!\n\n - `flax.core.apply` and linen `Module.apply` will now only return the variables\n collections that were specified as mutable.\n - Fixed handling of multiple separate subclasses of a Module.\n - We now allow assignment of mixed Module pytrees in setup.\n - Refactored collection creation to fail early when modifying an undefined collection as\n before an non-existing non-mutable collection would just be silently ignored.\n - Added the silu activation function.\n - Add offset argument to Adafactor optimizer for fine-tuning schedules.\n - Relaxed limit on calling methods on unbound modules.\n - Relaxed parameter attribute check\n - Added centered version of RMSProp.\n - Added GCE getting started kit.\n - Renamed -gpu_type to -accelerator_type.\n - Fixed bug in MultiOptimizer causing it to throw away empty dictionary\n\n### Improvements\n - Made FrozenDict constructor freeze correctly.\n - Made freeze a synonym of the FrozenDict constructor\n - Optimize freezing FrozenDicts by sharing immutable internal state.\n - We simplified __setattr__ handling of trees with Modules.\n - Minor improvements in dtype handling, broadcast option for dropout.\n - Added a dtype specification to Embed layer, made Adafactor use float32\n state consistently, and added a broadcasting option to the Dropout layer.\n - Improved frozen dict performance.\n - (Massive) docs improvements\n - End to end benchmarks added.\n - Examples were updated to Linen.\n\nv0.2.2\n----\n - Added Reinforcement Learning example (examples/ppo).\n - Fix Adafactor bug that prevented factorization.\n - Fix scan broadcast issue in functional core.\n - Fix initialization RNGs to work with omnistaging for jitted inits.\n - Replaces usage of 'param' kind to 'params' collection.\n - Fix LARS optimizer for zero param initialization.\n - Added various examples in Linen API. See [README.md](https://github.com/google/flax/blob/main/flax/linen/README.md) for more information.\n - Full JAX omnistaging compatibility.\n\nv0.2\n----\n - Added JAX trace-level checks for transforms.\n - BatchNorm added axis_index_groups for control in parallel training.\n - Optimizers broken out into separate directory with base class and implementations.\n - traverse_util added flatten_dict and unflatten_dict utility methods for nested dicts.\n\nv0.1\n----\n\n### API Changes\n - Add ConvTranspose Module to nn.linear\n - Rename the following optional arguments to nn.linear.Conv:\n `lhs_dilation` -> `input_dilation`,\n `rhs_dilation` -> `kernel_dilation`\n - Change default layer names from numbers '0', '1', etc. to\n include the Module class name, e.g. 'Dense_0', 'LayerNorm_1'.\n" }, { "path": "LICENSE", "content": " Version 2.0, January 2004\n http://www.apache.org/licenses/\n\n TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION\n\n 1. Definitions.\n\n \"License\" shall mean the terms and conditions for use, reproduction,\n and distribution as defined by Sections 1 through 9 of this document.\n\n \"Licensor\" shall mean the copyright owner or entity authorized by\n the copyright owner that is granting the License.\n\n \"Legal Entity\" shall mean the union of the acting entity and all\n other entities that control, are controlled by, or are under common\n control with that entity. For the purposes of this definition,\n \"control\" means (i) the power, direct or indirect, to cause the\n direction or management of such entity, whether by contract or\n otherwise, or (ii) ownership of fifty percent (50%) or more of the\n outstanding shares, or (iii) beneficial ownership of such entity.\n\n \"You\" (or \"Your\") shall mean an individual or Legal Entity\n exercising permissions granted by this License.\n\n \"Source\" form shall mean the preferred form for making modifications,\n including but not limited to software source code, documentation\n source, and configuration files.\n\n \"Object\" form shall mean any form resulting from mechanical\n transformation or translation of a Source form, including but\n not limited to compiled object code, generated documentation,\n and conversions to other media types.\n\n \"Work\" shall mean the work of authorship, whether in Source or\n Object form, made available under the License, as indicated by a\n copyright notice that is included in or attached to the work\n (an example is provided in the Appendix below).\n\n \"Derivative Works\" shall mean any work, whether in Source or Object\n form, that is based on (or derived from) the Work and for which the\n editorial revisions, annotations, elaborations, or other modifications\n represent, as a whole, an original work of authorship. For the purposes\n of this License, Derivative Works shall not include works that remain\n separable from, or merely link (or bind by name) to the interfaces of,\n the Work and Derivative Works thereof.\n\n \"Contribution\" shall mean any work of authorship, including\n the original version of the Work and any modifications or additions\n to that Work or Derivative Works thereof, that is intentionally\n submitted to Licensor for inclusion in the Work by the copyright owner\n or by an individual or Legal Entity authorized to submit on behalf of\n the copyright owner. For the purposes of this definition, \"submitted\"\n means any form of electronic, verbal, or written communication sent\n to the Licensor or its representatives, including but not limited to\n communication on electronic mailing lists, source code control systems,\n and issue tracking systems that are managed by, or on behalf of, the\n Licensor for the purpose of discussing and improving the Work, but\n excluding communication that is conspicuously marked or otherwise\n designated in writing by the copyright owner as \"Not a Contribution.\"\n\n \"Contributor\" shall mean Licensor and any individual or Legal Entity\n on behalf of whom a Contribution has been received by Licensor and\n subsequently incorporated within the Work.\n\n 2. Grant of Copyright License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n copyright license to reproduce, prepare Derivative Works of,\n publicly display, publicly perform, sublicense, and distribute the\n Work and such Derivative Works in Source or Object form.\n\n 3. Grant of Patent License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n (except as stated in this section) patent license to make, have made,\n use, offer to sell, sell, import, and otherwise transfer the Work,\n where such license applies only to those patent claims licensable\n by such Contributor that are necessarily infringed by their\n Contribution(s) alone or by combination of their Contribution(s)\n with the Work to which such Contribution(s) was submitted. If You\n institute patent litigation against any entity (including a\n cross-claim or counterclaim in a lawsuit) alleging that the Work\n or a Contribution incorporated within the Work constitutes direct\n or contributory patent infringement, then any patent licenses\n granted to You under this License for that Work shall terminate\n as of the date such litigation is filed.\n\n 4. Redistribution. You may reproduce and distribute copies of the\n Work or Derivative Works thereof in any medium, with or without\n modifications, and in Source or Object form, provided that You\n meet the following conditions:\n\n (a) You must give any other recipients of the Work or\n Derivative Works a copy of this License; and\n\n (b) You must cause any modified files to carry prominent notices\n stating that You changed the files; and\n\n (c) You must retain, in the Source form of any Derivative Works\n that You distribute, all copyright, patent, trademark, and\n attribution notices from the Source form of the Work,\n excluding those notices that do not pertain to any part of\n the Derivative Works; and\n\n (d) If the Work includes a \"NOTICE\" text file as part of its\n distribution, then any Derivative Works that You distribute must\n include a readable copy of the attribution notices contained\n within such NOTICE file, excluding those notices that do not\n pertain to any part of the Derivative Works, in at least one\n of the following places: within a NOTICE text file distributed\n as part of the Derivative Works; within the Source form or\n documentation, if provided along with the Derivative Works; or,\n within a display generated by the Derivative Works, if and\n wherever such third-party notices normally appear. The contents\n of the NOTICE file are for informational purposes only and\n do not modify the License. You may add Your own attribution\n notices within Derivative Works that You distribute, alongside\n or as an addendum to the NOTICE text from the Work, provided\n that such additional attribution notices cannot be construed\n as modifying the License.\n\n You may add Your own copyright statement to Your modifications and\n may provide additional or different license terms and conditions\n for use, reproduction, or distribution of Your modifications, or\n for any such Derivative Works as a whole, provided Your use,\n reproduction, and distribution of the Work otherwise complies with\n the conditions stated in this License.\n\n 5. Submission of Contributions. Unless You explicitly state otherwise,\n any Contribution intentionally submitted for inclusion in the Work\n by You to the Licensor shall be under the terms and conditions of\n this License, without any additional terms or conditions.\n Notwithstanding the above, nothing herein shall supersede or modify\n the terms of any separate license agreement you may have executed\n with Licensor regarding such Contributions.\n\n 6. Trademarks. This License does not grant permission to use the trade\n names, trademarks, service marks, or product names of the Licensor,\n except as required for reasonable and customary use in describing the\n origin of the Work and reproducing the content of the NOTICE file.\n\n 7. Disclaimer of Warranty. Unless required by applicable law or\n agreed to in writing, Licensor provides the Work (and each\n Contributor provides its Contributions) on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or\n implied, including, without limitation, any warranties or conditions\n of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A\n PARTICULAR PURPOSE. You are solely responsible for determining the\n appropriateness of using or redistributing the Work and assume any\n risks associated with Your exercise of permissions under this License.\n\n 8. Limitation of Liability. In no event and under no legal theory,\n whether in tort (including negligence), contract, or otherwise,\n unless required by applicable law (such as deliberate and grossly\n negligent acts) or agreed to in writing, shall any Contributor be\n liable to You for damages, including any direct, indirect, special,\n incidental, or consequential damages of any character arising as a\n result of this License or out of the use or inability to use the\n Work (including but not limited to damages for loss of goodwill,\n work stoppage, computer failure or malfunction, or any and all\n other commercial damages or losses), even if such Contributor\n has been advised of the possibility of such damages.\n\n 9. Accepting Warranty or Additional Liability. While redistributing\n the Work or Derivative Works thereof, You may choose to offer,\n and charge a fee for, acceptance of support, warranty, indemnity,\n or other liability obligations and/or rights consistent with this\n License. However, in accepting such obligations, You may act only\n on Your own behalf and on Your sole responsibility, not on behalf\n of any other Contributor, and only if You agree to indemnify,\n defend, and hold each Contributor harmless for any liability\n incurred by, or claims asserted against, such Contributor by reason\n of your accepting any such warranty or additional liability.\n\n END OF TERMS AND CONDITIONS\n\n APPENDIX: How to apply the Apache License to your work.\n\n To apply the Apache License to your work, attach the following\n boilerplate notice, with the fields enclosed by brackets \"[]\"\n replaced with your own identifying information. (Don't include\n the brackets!) The text should be enclosed in the appropriate\n comment syntax for the file format. We also recommend that a\n file or class name and description of purpose be included on the\n same \"printed page\" as the copyright notice for easier\n identification within third-party archives.\n\n Copyright [yyyy] [name of copyright owner]\n\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n\n http://www.apache.org/licenses/LICENSE-2.0\n\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n" }, { "path": "README.md", "content": "
\n\"logo\"\n
\n\n# Flax: A neural network library and ecosystem for JAX designed for flexibility\n\n[![Flax - Test](https://github.com/google/flax/actions/workflows/flax_test.yml/badge.svg)](https://github.com/google/flax/actions/workflows/flax_test.yml)\n[![PyPI version](https://img.shields.io/pypi/v/flax)](https://pypi.org/project/flax/)\n\n[**Overview**](#overview)\n| [**Quick install**](#quick-install)\n| [**What does Flax look like?**](#what-does-flax-look-like)\n| [**Documentation**](https://flax.readthedocs.io/)\n\nReleased in 2024, Flax NNX is a new simplified Flax API that is designed to make\nit easier to create, inspect, debug, and analyze neural networks in\n[JAX](https://jax.readthedocs.io/). It achieves this by adding first class support\nfor Python reference semantics. This allows users to express their models using\nregular Python objects, enabling reference sharing and mutability.\n\nFlax NNX evolved from the [Flax Linen API](https://flax-linen.readthedocs.io/), which\nwas released in 2020 by engineers and researchers at Google Brain in close collaboration\nwith the JAX team.\n\nYou can learn more about Flax NNX on the [dedicated Flax documentation site](https://flax.readthedocs.io/). Make sure you check out:\n\n* [Flax NNX basics](https://flax.readthedocs.io/en/latest/nnx_basics.html)\n* [MNIST tutorial](https://flax.readthedocs.io/en/latest/mnist_tutorial.html)\n* [Why Flax NNX](https://flax.readthedocs.io/en/latest/why.html)\n* [Evolution from Flax Linen to Flax NNX](https://flax.readthedocs.io/en/latest/guides/linen_to_nnx.html)\n\n**Note:** Flax Linen's [documentation has its own site](https://flax-linen.readthedocs.io/).\n\nThe Flax team's mission is to serve the growing JAX neural network\nresearch ecosystem - both within Alphabet and with the broader community,\nand to explore the use-cases where JAX shines. We use GitHub for almost\nall of our coordination and planning, as well as where we discuss\nupcoming design changes. We welcome feedback on any of our discussion,\nissue and pull request threads.\n\nYou can make feature requests, let us know what you are working on,\nreport issues, ask questions in our [Flax GitHub discussion\nforum](https://github.com/google/flax/discussions).\n\nWe expect to improve Flax, but we don't anticipate significant\nbreaking changes to the core API. We use [Changelog](https://github.com/google/flax/tree/main/CHANGELOG.md)\nentries and deprecation warnings when possible.\n\nIn case you want to reach us directly, we're at flax-dev@google.com.\n\n## Overview\n\nFlax is a high-performance neural network library and ecosystem for\nJAX that is **designed for flexibility**:\nTry new forms of training by forking an example and by modifying the training\nloop, not adding features to a framework.\n\nFlax is being developed in close collaboration with the JAX team and\ncomes with everything you need to start your research, including:\n\n* **Neural network API** (`flax.nnx`): Including [`Linear`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/linear.html#flax.nnx.Linear), [`Conv`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/linear.html#flax.nnx.Conv), [`BatchNorm`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/normalization.html#flax.nnx.BatchNorm), [`LayerNorm`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/normalization.html#flax.nnx.LayerNorm), [`GroupNorm`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/normalization.html#flax.nnx.GroupNorm), [Attention](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/attention.html) ([`MultiHeadAttention`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/attention.html#flax.nnx.MultiHeadAttention)), [`LSTMCell`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/recurrent.html#flax.nnx.nn.recurrent.LSTMCell), [`GRUCell`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/recurrent.html#flax.nnx.nn.recurrent.GRUCell), [`Dropout`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/stochastic.html#flax.nnx.Dropout).\n\n* **Utilities and patterns**: replicated training, serialization and checkpointing, metrics, prefetching on device.\n\n* **Educational examples**: [MNIST](https://flax.readthedocs.io/en/latest/mnist_tutorial.html), [Inference/sampling with the Gemma language model (transformer)](https://github.com/google/flax/tree/main/examples/gemma).\n\n## Quick install\n\nFlax uses JAX, so do check out [JAX installation instructions on CPUs, GPUs and TPUs](https://jax.readthedocs.io/en/latest/installation.html).\n\nYou will need Python 3.8 or later. Install Flax from PyPi:\n\n```\npip install flax\n```\n\nTo upgrade to the latest version of Flax, you can use:\n\n```\npip install --upgrade git+https://github.com/google/flax.git\n```\n\nTo install some additional dependencies (like `matplotlib`) that are required but not included\nby some dependencies, you can use:\n\n```bash\npip install \"flax[all]\"\n```\n\n## What does Flax look like?\n\nWe provide three examples using the Flax API: a simple multi-layer perceptron, a CNN and an auto-encoder.\n\nTo learn more about the `Module` abstraction, check out our [docs](https://flax.readthedocs.io/), our [broad intro to the Module abstraction](https://github.com/google/flax/blob/main/docs/linen_intro.ipynb). For additional concrete demonstrations of best practices, refer to our\n[guides](https://flax.readthedocs.io/en/latest/guides/index.html) and\n[developer notes](https://flax.readthedocs.io/en/latest/developer_notes/index.html).\n\nExample of an MLP:\n\n```py\nclass MLP(nnx.Module):\n def __init__(self, din: int, dmid: int, dout: int, *, rngs: nnx.Rngs):\n self.linear1 = nnx.Linear(din, dmid, rngs=rngs)\n self.dropout = nnx.Dropout(rate=0.1, rngs=rngs)\n self.bn = nnx.BatchNorm(dmid, rngs=rngs)\n self.linear2 = nnx.Linear(dmid, dout, rngs=rngs)\n\n def __call__(self, x: jax.Array):\n x = nnx.gelu(self.dropout(self.bn(self.linear1(x))))\n return self.linear2(x)\n```\n\nExample of a CNN:\n\n```py\nclass CNN(nnx.Module):\n def __init__(self, *, rngs: nnx.Rngs):\n self.conv1 = nnx.Conv(1, 32, kernel_size=(3, 3), rngs=rngs)\n self.conv2 = nnx.Conv(32, 64, kernel_size=(3, 3), rngs=rngs)\n self.avg_pool = partial(nnx.avg_pool, window_shape=(2, 2), strides=(2, 2))\n self.linear1 = nnx.Linear(3136, 256, rngs=rngs)\n self.linear2 = nnx.Linear(256, 10, rngs=rngs)\n\n def __call__(self, x):\n x = self.avg_pool(nnx.relu(self.conv1(x)))\n x = self.avg_pool(nnx.relu(self.conv2(x)))\n x = x.reshape(x.shape[0], -1) # flatten\n x = nnx.relu(self.linear1(x))\n x = self.linear2(x)\n return x\n```\n\nExample of an autoencoder:\n\n\n```py\nEncoder = lambda rngs: nnx.Linear(2, 10, rngs=rngs)\nDecoder = lambda rngs: nnx.Linear(10, 2, rngs=rngs)\n\nclass AutoEncoder(nnx.Module):\n def __init__(self, rngs):\n self.encoder = Encoder(rngs)\n self.decoder = Decoder(rngs)\n\n def __call__(self, x) -> jax.Array:\n return self.decoder(self.encoder(x))\n\n def encode(self, x) -> jax.Array:\n return self.encoder(x)\n```\n\n## Citing Flax\n\nTo cite this repository:\n\n```\n@software{flax2020github,\n author = {Jonathan Heek and Anselm Levskaya and Avital Oliver and Marvin Ritter and Bertrand Rondepierre and Andreas Steiner and Marc van {Z}ee},\n title = {{F}lax: A neural network library and ecosystem for {JAX}},\n url = {http://github.com/google/flax},\n version = {0.12.6},\n year = {2024},\n}\n```\n\nIn the above bibtex entry, names are in alphabetical order, the version number\nis intended to be that from [flax/version.py](https://github.com/google/flax/blob/main/flax/version.py), and the year corresponds to the project's open-source release.\n\n## Note\n\nFlax is an open source project maintained by a dedicated team at Google DeepMind, but is not an official Google product.\n\n" }, { "path": "benchmarks/README.md", "content": "# Benchmarks\n\nThese are mini benchmarks to measure the performance of NNX operations.\n\nSample profile command:\n\n```shell\npython -m cProfile -o ~/tmp/overhead.prof benchmarks/nnx_graph_overhead.py --mode=nnx --depth=100 --total_steps=1000\n```\n\nSample profile inspection:\n\n```shell\nsnakeviz ~/tmp/overhead.prof\n```" }, { "path": "benchmarks/nnx_graph_overhead.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport jax\nimport jax.numpy as jnp\nimport numpy as np\nimport optax\nfrom time import time\n\nfrom flax import nnx\n\nfrom absl import flags\nfrom absl import app\n\nFLAGS = flags.FLAGS\nflags.DEFINE_enum(\n 'mode', 'nnx', ['all', 'nnx', 'jax'], 'Mode to run the script in'\n)\nflags.DEFINE_integer('total_steps', 100, 'Total number of training steps')\nflags.DEFINE_integer('width', 32, 'Hidden layer size')\nflags.DEFINE_integer('depth', 5, 'Depth of the model')\n\n\nclass Linear(nnx.Module):\n def __init__(self, din: int, dout: int, *, rngs: nnx.Rngs):\n self.w = nnx.Param(jax.random.uniform(rngs.params(), (din, dout)))\n self.b = nnx.Param(jnp.zeros((dout,)))\n\n def __call__(self, x):\n return x @ self.w + self.b\n\n\nclass Block(nnx.Module):\n def __init__(self, din: int, dout: int, *, rngs: nnx.Rngs):\n self.linear = Linear(din, dout, rngs=rngs)\n self.bn = nnx.BatchNorm(dout, rngs=rngs)\n\n def __call__(self, x):\n return nnx.relu(self.bn(self.linear(x)))\n\n\nclass Count(nnx.Variable):\n pass\n\n\nclass MLP(nnx.Module):\n def __init__(self, din, dhidden, dout, depth, *, rngs: nnx.Rngs):\n self.count = Count(jnp.array(0))\n self.linear_in = Block(din, dhidden, rngs=rngs)\n self.intermediates = nnx.List([\n Block(dhidden, dhidden, rngs=rngs) for _ in range(depth - 2)\n ])\n self.linear_out = Block(dhidden, dout, rngs=rngs)\n\n def __call__(self, x):\n self.count.value += 1\n x = nnx.relu(self.linear_in(x))\n for layer in self.intermediates:\n x = nnx.relu(layer(x))\n x = self.linear_out(x)\n return x\n\n\ndef main(argv):\n print(argv)\n mode: str = FLAGS.mode\n total_steps: int = FLAGS.total_steps\n width: int = FLAGS.width\n depth: int = FLAGS.depth\n\n print(f'{mode=}, {total_steps=}, {width=}')\n\n X = np.linspace(0, 1, 100)[:, None]\n Y = 0.8 * X**2 + 0.1 + np.random.normal(0, 0.1, size=X.shape)\n\n # ------------------------------------------------------------\n # NNX\n # ------------------------------------------------------------\n if mode in ['all', 'nnx']:\n model = MLP(din=1, dhidden=width, dout=1, depth=depth, rngs=nnx.Rngs(0))\n tx = optax.sgd(1e-3)\n optimizer = nnx.Optimizer(model, tx, wrt=nnx.Param)\n t0 = time()\n\n @nnx.jit\n def step_nnx(model: MLP, optimizer: nnx.Optimizer):\n pass\n\n t0 = time()\n for _ in range(total_steps):\n step_nnx(model, optimizer)\n\n total_time = time() - t0\n time_per_step = total_time / total_steps\n time_per_layer = time_per_step / depth\n print('### NNX ###')\n print('total time:', total_time)\n print(f'time per step: {time_per_step * 1e6:.2f} µs')\n print(f'time per layer: {time_per_layer * 1e6:.2f} µs')\n\n # ------------------------------------------------------------\n # JAX\n # ------------------------------------------------------------\n\n if mode in ['all', 'jax']:\n model = MLP(din=1, dhidden=width, dout=1, depth=depth, rngs=nnx.Rngs(0))\n tx = optax.sgd(1e-3)\n optimizer = nnx.Optimizer(model, tx, wrt=nnx.Param)\n t0 = time()\n\n @jax.jit\n def step_jax(graphdef, state):\n return graphdef, state\n\n graphdef, state = nnx.split((model, optimizer))\n t0 = time()\n for _ in range(total_steps):\n graphdef, state = step_jax(graphdef, state)\n\n total_time = time() - t0\n time_per_step = total_time / total_steps\n time_per_layer = time_per_step / depth\n print('### JAX ###')\n print('total time:', total_time)\n print(f'time per step: {time_per_step * 1e6:.2f} µs')\n print(f'time per layer: {time_per_layer * 1e6:.2f} µs')\n print()\n\n\nif __name__ == '__main__':\n app.run(main)\n" }, { "path": "benchmarks/nnx_mlpmixer_training.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# %%\nfrom functools import partial\nimport jax\nimport jax.numpy as jnp\nfrom flax import nnx\nimport optax\nimport numpy as np\nfrom einop import einop\nfrom time import time\nfrom tqdm import tqdm\n\nfrom flax import nnx\n\nfrom absl import flags\nfrom absl import app\n\nFLAGS = flags.FLAGS\nflags.DEFINE_enum(\n 'mode', 'all', ['all', 'nnx', 'jax'], 'Mode to run the script in'\n)\nflags.DEFINE_integer('total_steps', 10_000, 'Total number of training steps')\nflags.DEFINE_integer('batch_size', 32, 'Batch size')\nflags.DEFINE_integer('width', 32, 'Hidden layer size')\nflags.DEFINE_integer('depth', 4, 'Depth of the model')\n\n\nclass MlpBlock(nnx.Module):\n def __init__(self, din: int, mlp_dim: int, rngs: nnx.Rngs):\n self.din, self.mlp_dim = din, mlp_dim\n self.linear_in = nnx.Linear(din, mlp_dim, rngs=rngs)\n self.linear_out = nnx.Linear(mlp_dim, din, rngs=rngs)\n\n def __call__(self, x):\n return self.linear_out(nnx.gelu(self.linear_in(x)))\n\n\nclass MixerBlock(nnx.Module):\n def __init__(\n self,\n tokens_mlp_dim: int,\n channels_mlp_dim: int,\n hidden_dim: int,\n rngs: nnx.Rngs,\n ):\n self.tokens_mlp_dim = tokens_mlp_dim\n self.channels_mlp_dim = channels_mlp_dim\n self.hidden_dim = hidden_dim\n self.token_mixing = MlpBlock(tokens_mlp_dim, hidden_dim, rngs=rngs)\n self.channel_mixing = MlpBlock(channels_mlp_dim, hidden_dim, rngs=rngs)\n self.ln1 = nnx.LayerNorm(channels_mlp_dim, rngs=rngs)\n self.ln2 = nnx.LayerNorm(channels_mlp_dim, rngs=rngs)\n\n def __call__(self, x):\n y = self.ln1(x)\n y = y.swapaxes(1, 2)\n y = self.token_mixing(y)\n y = y.swapaxes(1, 2)\n x = x + y\n y = self.ln2(x)\n return x + self.channel_mixing(y)\n\n\nclass MlpMixer(nnx.Module):\n def __init__(\n self,\n din: int,\n kernel_size: tuple[int, int],\n strides: tuple[int, int],\n num_blocks: int,\n hidden_dim: int,\n tokens_mlp_dim: int,\n channels_mlp_dim: int,\n rngs: nnx.Rngs,\n ):\n self.din = din\n self.kernel_size = kernel_size\n self.num_blocks = num_blocks\n self.hidden_dim = hidden_dim\n self.tokens_mlp_dim = tokens_mlp_dim\n self.channels_mlp_dim = channels_mlp_dim\n self.stem = nnx.Conv(\n din + 1,\n channels_mlp_dim,\n kernel_size=kernel_size,\n strides=strides,\n rngs=rngs,\n )\n self.blocks = [\n MixerBlock(tokens_mlp_dim, channels_mlp_dim, hidden_dim, rngs=rngs)\n for _ in range(num_blocks)\n ]\n self.pre_head_layer_norm = nnx.LayerNorm(channels_mlp_dim, rngs=rngs)\n self.conv_t = nnx.ConvTranspose(\n channels_mlp_dim, din, kernel_size=kernel_size, strides=strides, rngs=rngs\n )\n\n def __call__(self, *, x, t):\n # add time feature to input\n t = einop(t, 'n -> n h w c', h=x.shape[1], w=x.shape[2], c=1)\n x = jnp.concatenate([x, t], axis=-1)\n # create patches\n x = self.stem(x)\n h, w = x.shape[1], x.shape[2]\n x = einop(x, 'n h w c -> n (h w) c')\n # apply blocks\n for block in self.blocks:\n x = block(x)\n x = self.pre_head_layer_norm(x)\n # recreate image\n x = einop(x, 'n (h w) c -> n h w c', h=h, w=w)\n x = self.conv_t(x)\n return x\n\n\ndef main(argv):\n print(argv)\n mode: str = FLAGS.mode\n total_steps: int = FLAGS.total_steps\n batch_size: int = FLAGS.batch_size\n width: int = FLAGS.width\n depth: int = FLAGS.depth\n\n print(f'{mode=}, {total_steps=}, {batch_size=}, {width=}')\n\n X = np.random.uniform(size=(batch_size, 28, 28, 1))\n\n if mode == 'nnx' or mode == 'all':\n rngs = nnx.Rngs(0)\n flow = MlpMixer(\n din=1,\n kernel_size=(2, 2),\n strides=(2, 2),\n num_blocks=4,\n hidden_dim=512,\n tokens_mlp_dim=196,\n channels_mlp_dim=512,\n rngs=rngs,\n )\n optimizer = nnx.Optimizer(\n flow, tx=optax.adamw(1e-4), wrt=nnx.Param\n )\n t0 = time()\n\n mse = lambda a, b: jnp.mean((a - b) ** 2)\n\n @nnx.jit(donate_argnums=(0, 1, 2))\n def train_step_nnx(flow, optimizer, rngs, x_1):\n print('JITTING NNX')\n x_0 = jax.random.normal(rngs(), x_1.shape)\n t = jax.random.uniform(rngs(), (len(x_1),))\n\n x_t = jax.vmap(lambda x_0, x_1, t: (1 - t) * x_0 + t * x_1)(x_0, x_1, t)\n dx_t = x_1 - x_0\n\n loss, grads = nnx.value_and_grad(\n lambda flow: mse(flow(x=x_t, t=t), dx_t)\n )(flow)\n optimizer.update(flow, grads)\n return loss\n\n losses = []\n t0 = time()\n for step in tqdm(range(total_steps), desc='NNX'):\n loss = train_step_nnx(flow, optimizer, rngs, X)\n losses.append(loss)\n\n total_time = time() - t0\n print('### NNX ###')\n print(f'final loss: {losses[-1]}')\n print('total time:', total_time)\n print(f'time per step: {total_time / total_steps * 1e6:.2f} µs')\n\n if mode == 'jax' or mode == 'all':\n rngs = nnx.Rngs(0)\n flow = MlpMixer(\n din=1,\n kernel_size=(2, 2),\n strides=(2, 2),\n num_blocks=depth,\n hidden_dim=width,\n tokens_mlp_dim=196,\n channels_mlp_dim=width,\n rngs=rngs,\n )\n optimizer = nnx.Optimizer(\n flow, tx=optax.adamw(1e-4), wrt=nnx.Param\n )\n graphdef, state = nnx.split((flow, optimizer, rngs))\n t0 = time()\n\n mse = lambda a, b: jnp.mean((a - b) ** 2)\n\n @partial(nnx.jit, donate_argnums=0)\n def train_step_jax(state, x_1):\n print('JITTING JAX')\n flow, optimizer, rngs = nnx.merge(graphdef, state)\n x_0 = jax.random.normal(rngs(), x_1.shape)\n t = jax.random.uniform(rngs(), (len(x_1),))\n\n x_t = jax.vmap(lambda x_0, x_1, t: (1 - t) * x_0 + t * x_1)(x_0, x_1, t)\n dx_t = x_1 - x_0\n\n loss, grads = nnx.value_and_grad(\n lambda flow: mse(flow(x=x_t, t=t), dx_t)\n )(flow)\n optimizer.update(flow, grads)\n state = nnx.state((flow, optimizer, rngs))\n return loss, state\n\n losses = []\n t0 = time()\n for step in tqdm(range(total_steps), desc='JAX'):\n loss, state = train_step_jax(state, X)\n losses.append(loss)\n\n nnx.update((flow, optimizer, rngs), state)\n total_time = time() - t0\n print('### JAX ###')\n print(f'final loss: {losses[-1]}')\n print('total time:', total_time)\n print(f'time per step: {total_time / total_steps * 1e6:.2f} µs')\n\n\nif __name__ == '__main__':\n app.run(main)\n" }, { "path": "benchmarks/nnx_simple_training.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# %%\nfrom functools import partial\nimport jax\nimport jax.numpy as jnp\nimport numpy as np\nimport optax\nfrom time import time\n\nfrom flax import nnx\n\nfrom absl import flags\nfrom absl import app\n\nFLAGS = flags.FLAGS\nflags.DEFINE_enum(\n 'mode', 'all', ['all', 'nnx', 'jax'], 'Mode to run the script in'\n)\nflags.DEFINE_integer('total_steps', 10_000, 'Total number of training steps')\nflags.DEFINE_integer('batch_size', 32, 'Batch size')\nflags.DEFINE_integer('width', 32, 'Hidden layer size')\nflags.DEFINE_integer('depth', 5, 'Depth of the model')\n\n\ndef dataset(X, Y, batch_size):\n while True:\n idx = np.random.choice(len(X), size=batch_size)\n yield X[idx], Y[idx]\n\n\nclass Linear(nnx.Module):\n def __init__(self, din: int, dout: int, *, rngs: nnx.Rngs):\n self.w = nnx.Param(jax.random.uniform(rngs.params(), (din, dout)))\n self.b = nnx.Param(jnp.zeros((dout,)))\n\n def __call__(self, x):\n return x @ self.w + self.b\n\n\nclass Block(nnx.Module):\n def __init__(self, din: int, dout: int, *, rngs: nnx.Rngs):\n self.linear = Linear(din, dout, rngs=rngs)\n self.bn = nnx.BatchNorm(dout, rngs=rngs)\n\n def __call__(self, x):\n return nnx.relu(self.bn(self.linear(x)))\n\n\nclass Count(nnx.Variable):\n pass\n\n\nclass MLP(nnx.Module):\n def __init__(self, din, dhidden, dout, depth, *, rngs: nnx.Rngs):\n self.count = Count(jnp.array(0))\n self.linear_in = Block(din, dhidden, rngs=rngs)\n self.intermediates = [\n Block(dhidden, dhidden, rngs=rngs) for _ in range(depth - 2)\n ]\n self.linear_out = Block(dhidden, dout, rngs=rngs)\n\n def __call__(self, x):\n self.count.value += 1\n x = nnx.relu(self.linear_in(x))\n for layer in self.intermediates:\n x = nnx.relu(layer(x))\n x = self.linear_out(x)\n return x\n\n\ndef main(argv):\n print(argv)\n mode: str = FLAGS.mode\n total_steps: int = FLAGS.total_steps\n batch_size: int = FLAGS.batch_size\n width: int = FLAGS.width\n depth: int = FLAGS.depth\n\n print(f'{mode=}, {total_steps=}, {batch_size=}, {width=}')\n\n X = np.linspace(0, 1, 100)[:, None]\n Y = 0.8 * X**2 + 0.1 + np.random.normal(0, 0.1, size=X.shape)\n\n if mode == 'nnx' or mode == 'all':\n model = MLP(din=1, dhidden=width, dout=1, depth=depth, rngs=nnx.Rngs(0))\n tx = optax.sgd(1e-3)\n optimizer = nnx.Optimizer(model, tx, wrt=nnx.Param)\n t0 = time()\n\n @nnx.jit(donate_argnums=(0, 1))\n def train_step_nnx(model: MLP, optimizer: nnx.Optimizer, batch):\n x, y = batch\n\n def loss_fn(model: MLP):\n y_pred = model(x)\n return jnp.mean((y - y_pred) ** 2)\n\n grads: nnx.State = nnx.grad(loss_fn)(model)\n optimizer.update(model, grads)\n\n @nnx.jit(donate_argnums=0)\n def test_step_nnx(model: MLP, batch):\n x, y = batch\n y_pred = model(x)\n loss = jnp.mean((y - y_pred) ** 2)\n return {'loss': loss}\n\n for step, batch in enumerate(dataset(X, Y, batch_size)):\n train_step_nnx(model, optimizer, batch)\n\n if step % 1000 == 0:\n logs = test_step_nnx(model, (X, Y))\n\n if step >= total_steps - 1:\n break\n\n print('### NNX ###')\n print(f'final loss: {logs[\"loss\"]}')\n total_time = time() - t0\n print('total time:', total_time)\n print(f'time per step: {total_time / total_steps * 1e6:.2f} µs')\n print('times called:', model.count.value)\n\n if mode == 'jax' or mode == 'all':\n model = MLP(din=1, dhidden=width, dout=1, depth=depth, rngs=nnx.Rngs(0))\n tx = optax.sgd(1e-3)\n optimizer = nnx.Optimizer(model, tx, wrt=nnx.Param)\n t0 = time()\n\n @partial(jax.jit, donate_argnums=0)\n def train_step_jax(state, batch):\n model, optimizer = nnx.merge(graphdef, state)\n x, y = batch\n\n def loss_fn(model: MLP):\n y_pred = model(x)\n return jnp.mean((y - y_pred) ** 2)\n\n grads = nnx.grad(loss_fn)(model)\n optimizer.update(model,grads)\n\n return nnx.state((model, optimizer))\n\n @partial(jax.jit, donate_argnums=0)\n def test_step_jax(state, batch):\n model, optimizer = nnx.merge(graphdef, state)\n x, y = batch\n y_pred = model(x)\n loss = jnp.mean((y - y_pred) ** 2)\n state = nnx.state((model, optimizer))\n return state, {'loss': loss}\n\n graphdef, state = nnx.split((model, optimizer))\n\n for step, batch in enumerate(dataset(X, Y, batch_size)):\n state = train_step_jax(state, batch)\n\n if step % 1000 == 0:\n state, logs = test_step_jax(state, (X, Y))\n\n if step >= total_steps - 1:\n break\n\n model, optimizer = nnx.merge(graphdef, state)\n\n print('### JAX ###')\n print(f'final loss: {logs[\"loss\"]}')\n total_time = time() - t0\n print('total time:', total_time)\n print(f'time per step: {total_time / total_steps * 1e6:.2f} µs')\n print('times called:', model.count.value)\n\n\nif __name__ == '__main__':\n app.run(main)\n" }, { "path": "benchmarks/nnx_state_traversal.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# Example profile command:\n# python -m cProfile -o ~/tmp/overhead.prof benchmarks/nnx_graph_overhead.py --mode=nnx --depth=100 --total_steps=1000\n# View profile (need to install snakeviz):\n# snakeviz ~/tmp/overhead.prof\n\nimport jax\nfrom time import time\n\nfrom flax import nnx\n\nfrom absl import flags\nfrom absl import app\n\nFLAGS = flags.FLAGS\nflags.DEFINE_integer('total_steps', 1000, 'Total number of training steps')\nflags.DEFINE_integer('width', 4, 'Width of each level')\nflags.DEFINE_integer('depth', 4, 'Depth of the model')\n\n\nclass NestedClass(nnx.Module):\n def __init__(self, width, depth):\n self.x = nnx.Variable(jax.numpy.ones((depth+1, )))\n if depth > 0:\n for i in range(width):\n setattr(self, f'child{i}', NestedClass(width, depth-1))\n\n\ndef main(argv):\n print(argv)\n total_steps: int = FLAGS.total_steps\n width: int = FLAGS.width\n depth: int = FLAGS.depth\n\n\n model = NestedClass(width, depth)\n to_test = nnx.state(model)\n\n print(f'{total_steps=}, {width=}')\n\n #------------------------------------------------------------\n # tree_flatten_with_path\n #------------------------------------------------------------\n t0 = time()\n for _ in range(total_steps):\n jax.tree_util.tree_flatten_with_path(to_test)\n\n total_time = time() - t0\n time_per_step = total_time / total_steps\n time_per_layer = time_per_step / depth\n print(\"### tree_flatten_with_path ###\")\n print('total time:', total_time)\n print(f'time per step: {time_per_step * 1e6:.2f} µs')\n print(f'time per layer: {time_per_layer * 1e6:.2f} µs')\n\n\n #------------------------------------------------------------\n # tree_map_with_path\n #------------------------------------------------------------\n\n t0 = time()\n for _ in range(total_steps):\n jax.tree_util.tree_map_with_path(lambda _, x: x, to_test)\n\n total_time = time() - t0\n time_per_step = total_time / total_steps\n time_per_layer = time_per_step / depth\n print(\"### tree_map_with_path ###\")\n print('total time:', total_time)\n print(f'time per step: {time_per_step * 1e6:.2f} µs')\n print(f'time per layer: {time_per_layer * 1e6:.2f} µs')\n\n\n #------------------------------------------------------------\n # tree_flatten\n #------------------------------------------------------------\n\n t0 = time()\n for _ in range(total_steps):\n jax.tree_util.tree_flatten(to_test)\n\n total_time = time() - t0\n time_per_step = total_time / total_steps\n time_per_layer = time_per_step / depth\n print(\"### tree_flatten ###\")\n print('total time:', total_time)\n print(f'time per step: {time_per_step * 1e6:.2f} µs')\n print(f'time per layer: {time_per_layer * 1e6:.2f} µs')\n\n\n\nif __name__ == '__main__':\n app.run(main)\n" }, { "path": "benchmarks/tracing/README.md", "content": "# Tracing and lowering benchmarks for Flax examples\n\nSee Flax\n[documentation](https://flax.readthedocs.io/en/latest/examples/index.html) on\ntheir examples.\n\n## Getting started\nbash\n```\npip install -r benchmarks/tracing/requirements.txt\n\n# Benchmark trace and lower timing for all workloads.\npython tracing_benchmark.py\n\n# Profile a single example.\npython tracing_benchmark.py --example=wmt\n\n# Profile just tracing for a single example.\npython tracing_benchmark.py --example=wmt --mode=trace\n```" }, { "path": "benchmarks/tracing/__init__.py", "content": "# Copyright 2025 The JAX Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n" }, { "path": "benchmarks/tracing/gemma.py", "content": "# Copyright 2025 The JAX Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"Gemma helper functions.\"\"\"\n\nfrom typing import Any\n\nfrom flax import nnx\nfrom flax.benchmarks.tracing import tracing_benchmark\nfrom flax.examples.gemma import transformer as transformer_lib\nfrom flax.examples.gemma import utils\nfrom flax.examples.gemma.configs import default as gemma_config\nfrom flax.training import common_utils\nimport google_benchmark\nimport jax\nimport jax.numpy as jnp\nimport ml_collections\nimport numpy as np\nimport optax\n\n\ndef rsqrt_schedule(init_value: float, shift: int = 0):\n def schedule(count):\n return init_value * (count + shift) ** -0.5 * shift**0.5\n return schedule\n\n\ndef create_learning_rate_schedule(learning_rate: float, warmup_steps: int):\n return optax.join_schedules(\n [\n optax.linear_schedule(\n init_value=0,\n end_value=learning_rate,\n transition_steps=warmup_steps,\n ),\n rsqrt_schedule(init_value=learning_rate, shift=warmup_steps),\n ],\n boundaries=[warmup_steps],\n )\n\n\ndef compute_weighted_cross_entropy(\n logits, targets, weights=None, label_smoothing=0.0\n):\n if logits.ndim != targets.ndim + 1:\n raise ValueError(\n 'Incorrect shapes. Got shape %s logits and %s targets'\n % (str(logits.shape), str(targets.shape))\n )\n vocab_size = logits.shape[-1]\n confidence = 1.0 - label_smoothing\n low_confidence = (1.0 - confidence) / (vocab_size - 1)\n normalizing_constant = -(\n confidence * jnp.log(confidence)\n + (vocab_size - 1) * low_confidence * jnp.log(low_confidence + 1e-20)\n )\n soft_targets = common_utils.onehot(\n targets, vocab_size, on_value=confidence, off_value=low_confidence\n )\n\n loss = -jnp.sum(soft_targets * nnx.log_softmax(logits), axis=-1)\n loss = loss - normalizing_constant\n\n normalizing_factor = np.prod(targets.shape)\n if weights is not None:\n loss = loss * weights\n normalizing_factor = weights.sum()\n\n return loss.sum(), normalizing_factor\n\n\ndef compute_weighted_accuracy(logits, targets, weights=None):\n if logits.ndim != targets.ndim + 1:\n raise ValueError(\n 'Incorrect shapes. Got shape %s logits and %s targets'\n % (str(logits.shape), str(targets.shape))\n )\n loss = jnp.equal(jnp.argmax(logits, axis=-1), targets)\n normalizing_factor = np.prod(logits.shape[:-1])\n if weights is not None:\n loss = loss * weights\n normalizing_factor = weights.sum()\n\n return loss.sum(), normalizing_factor\n\n\ndef compute_metrics(logits, labels, weights, label_smoothing=0.0):\n loss, weight_sum = compute_weighted_cross_entropy(\n logits, labels, weights, label_smoothing\n )\n acc, _ = compute_weighted_accuracy(logits, labels, weights)\n metrics = {\n 'loss': loss,\n 'accuracy': acc,\n 'denominator': weight_sum,\n }\n return metrics\n\n\ndef train_step(\n state: utils.TrainState,\n batch,\n learning_rate_fn,\n label_smoothing=0.0,\n):\n train_keys = ['inputs', 'inputs_position', 'inputs_segmentation', 'targets']\n (inputs, inputs_positions, inputs_segmentation, targets) = (\n batch.get(k, None) for k in train_keys\n )\n\n pad_id = 0\n weights = jnp.where(inputs > pad_id, 1, 0).astype(jnp.float32)\n input_mask = inputs > pad_id\n attention_mask = transformer_lib.make_causal_attn_mask(input_mask)\n mask = (\n inputs_segmentation[:, :, None] == inputs_segmentation[:, None, :]\n )\n attention_mask = jnp.logical_and(mask, attention_mask)\n\n def loss_fn(params):\n module = nnx.merge(state.graphdef, params)\n\n logits, _ = module(\n inputs,\n positions=inputs_positions,\n attention_mask=attention_mask,\n cache=None,\n )\n\n loss, weight_sum = compute_weighted_cross_entropy(\n logits, targets, weights, label_smoothing\n )\n mean_loss = loss / weight_sum\n return mean_loss, logits\n\n step = state.step\n lr = learning_rate_fn(step)\n grad_fn = jax.value_and_grad(loss_fn, has_aux=True)\n (_, logits), grads = grad_fn(state.params)\n new_state = state.apply_gradients(grads=grads)\n metrics = compute_metrics(logits, targets, weights)\n metrics['learning_rate'] = lr\n\n return new_state, metrics\n\n\ndef get_fake_batch(batch_size: int) -> Any:\n rng = jax.random.PRNGKey(0)\n batch = {}\n for k in (\n 'inputs',\n 'inputs_position',\n 'inputs_segmentation',\n 'targets',\n 'targets_position',\n 'targets_segmentation',\n ):\n batch[k] = jax.random.randint(rng, (batch_size, 128), 0, 9999999, jnp.int32)\n return batch\n\n\ndef get_apply_fn_and_args(\n config: ml_collections.ConfigDict,\n vocab_size: int | None = None,\n) -> tuple[Any, tuple[Any, ...], dict[str, Any]]:\n if vocab_size is None:\n vocab_size = config.vocab_size\n\n if config.transformer_name is not None:\n model_config = transformer_lib.TransformerConfig.from_version_name(\n config.transformer_name,\n num_embed=vocab_size,\n dtype=jnp.bfloat16 if config.use_bfloat16 else jnp.float32,\n axis_rules=config.axis_rules,\n )\n else:\n assert config.transformer_params is not None\n model_config = transformer_lib.TransformerConfig.from_dict(\n **config.transformer_params,\n num_embed=vocab_size,\n dtype=jnp.bfloat16 if config.use_bfloat16 else jnp.float32,\n axis_rules=config.axis_rules,\n )\n\n devices_array = utils.create_device_mesh(config)\n mesh = jax.sharding.Mesh(devices_array, config.mesh_axes)\n\n rng = jax.random.PRNGKey(config.seed)\n rng, init_rng = jax.random.split(rng)\n\n def constructor(config: transformer_lib.TransformerConfig, key: jax.Array):\n return transformer_lib.Transformer(config, rngs=nnx.Rngs(params=key))\n\n learning_rate_fn = create_learning_rate_schedule(\n learning_rate=config.learning_rate, warmup_steps=config.warmup_steps\n )\n\n optimizer = optax.adamw(\n learning_rate_fn,\n b1=0.9,\n b2=0.98,\n eps=1e-9,\n weight_decay=config.weight_decay,\n )\n\n state, state_sharding = utils.setup_initial_state(\n constructor, optimizer, model_config, init_rng, mesh\n )\n data_sharding = jax.NamedSharding(mesh, jax.P(config.data_sharding))\n jit_train_step = jax.jit(\n train_step,\n in_shardings=(\n state_sharding,\n data_sharding,\n ), # type: ignore\n out_shardings=(state_sharding, None), # type: ignore\n static_argnames=('learning_rate_fn', 'label_smoothing'),\n donate_argnums=0,\n )\n\n batch = get_fake_batch(config.per_device_batch_size)\n batch = jax.tree.map(lambda x: jnp.asarray(x, device=data_sharding), batch)\n\n return (\n jit_train_step,\n (state, batch, learning_rate_fn, 0.0),\n dict(),\n )\n\n\n@google_benchmark.register\n@google_benchmark.option.unit(google_benchmark.kMillisecond)\ndef test_flax_gemma_trace(state):\n tracing_benchmark.benchmark_tracing(\n get_apply_fn_and_args, gemma_config.get_config, state\n )\n\n\n@google_benchmark.register\n@google_benchmark.option.unit(google_benchmark.kMillisecond)\ndef test_flax_gemma_lower(state):\n tracing_benchmark.benchmark_lowering(\n get_apply_fn_and_args, gemma_config.get_config, state\n )\n\n\nif __name__ == '__main__':\n tracing_benchmark.run_benchmarks()\n" }, { "path": "benchmarks/tracing/imagenet.py", "content": "# Copyright 2025 The JAX Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"ImageNet helper functions for benchmarking.\"\"\"\n\nimport functools\nfrom typing import Any\n\nfrom flax import linen as nn\nfrom flax.benchmarks.tracing import tracing_benchmark\nfrom flax.examples.imagenet import models\nfrom flax.examples.imagenet.configs import default as imagenet_config\nfrom flax.training import common_utils\nfrom flax.training import dynamic_scale as dynamic_scale_lib\nfrom flax.training import train_state\nimport google_benchmark\nimport jax\nimport jax.numpy as jnp\nimport ml_collections\nimport optax\n\nNUM_CLASSES = 1000\n\n\nclass TrainState(train_state.TrainState):\n batch_stats: Any\n dynamic_scale: dynamic_scale_lib.DynamicScale\n\n\ndef create_model(*, model_cls, half_precision, **kwargs):\n platform = jax.local_devices()[0].platform\n if half_precision:\n if platform == 'tpu':\n model_dtype = jnp.bfloat16\n else:\n model_dtype = jnp.float16\n else:\n model_dtype = jnp.float32\n return model_cls(num_classes=NUM_CLASSES, dtype=model_dtype, **kwargs)\n\n\ndef initialized(key, image_size, model):\n input_shape = (1, image_size, image_size, 3)\n\n @jax.jit\n def init(*args):\n return model.init(*args)\n\n variables = init({'params': key}, jnp.ones(input_shape, model.dtype))\n return variables['params'], variables['batch_stats']\n\n\ndef cross_entropy_loss(logits, labels):\n one_hot_labels = common_utils.onehot(labels, num_classes=NUM_CLASSES)\n xentropy = optax.softmax_cross_entropy(logits=logits, labels=one_hot_labels)\n return jnp.mean(xentropy)\n\n\ndef create_train_state(\n rng, config: ml_collections.ConfigDict, model, image_size, learning_rate_fn\n):\n dynamic_scale = None\n platform = jax.local_devices()[0].platform\n if config.half_precision and platform == 'gpu':\n dynamic_scale = dynamic_scale_lib.DynamicScale()\n\n params, batch_stats = initialized(rng, image_size, model)\n tx = optax.sgd(\n learning_rate=learning_rate_fn,\n momentum=config.momentum,\n nesterov=True,\n )\n state = TrainState.create(\n apply_fn=model.apply,\n params=params,\n tx=tx,\n batch_stats=batch_stats,\n dynamic_scale=dynamic_scale,\n )\n return state\n\n\ndef get_fake_batch(batch_size: int = 128) -> dict[str, jnp.ndarray]:\n images = jax.random.uniform(\n jax.random.key(0), (batch_size, 224, 224, 3), dtype=jnp.float32\n )\n labels = jax.random.randint(\n jax.random.key(1), (batch_size,), minval=0, maxval=1000, dtype=jnp.int32\n )\n return {'image': images, 'label': labels}\n\n\nclass BenchmarkResNet(models.ResNet):\n\n @nn.compact\n def __call__(self, x, train: bool = True):\n conv = functools.partial(self.conv, use_bias=False, dtype=self.dtype)\n norm = functools.partial(\n nn.BatchNorm,\n use_running_average=not train,\n momentum=0.9,\n epsilon=1e-5,\n dtype=self.dtype,\n axis_name=None,\n )\n\n x = conv(\n self.num_filters,\n (7, 7),\n (2, 2),\n padding=[(3, 3), (3, 3)],\n name='conv_init',\n )(x)\n x = norm(name='bn_init')(x)\n x = nn.relu(x)\n x = nn.max_pool(x, (3, 3), strides=(2, 2), padding='SAME')\n for i, block_size in enumerate(self.stage_sizes):\n for j in range(block_size):\n strides = (2, 2) if i > 0 and j == 0 else (1, 1)\n x = self.block_cls(\n self.num_filters * 2**i,\n strides=strides,\n conv=conv,\n norm=norm,\n act=self.act,\n )(x)\n x = jnp.mean(x, axis=(1, 2))\n x = nn.Dense(self.num_classes, dtype=self.dtype)(x)\n x = jnp.asarray(x, self.dtype)\n return x\n\n\ndef get_apply_fn_and_args(\n config: ml_collections.ConfigDict,\n) -> tuple[Any, tuple[Any, ...], dict[str, Any]]:\n if config.model == 'ResNet50':\n model_cls = functools.partial(\n BenchmarkResNet,\n stage_sizes=[3, 4, 6, 3],\n block_cls=models.BottleneckResNetBlock,\n )\n else:\n model_cls = getattr(models, config.model)\n\n model = create_model(\n model_cls=model_cls, half_precision=config.half_precision\n )\n\n learning_rate_fn = lambda step: 0.1\n\n rng = jax.random.key(0)\n image_size = 224\n state = create_train_state(\n rng, config, model, image_size, learning_rate_fn\n )\n\n batch = get_fake_batch(config.batch_size)\n\n return (\n bench_train_step,\n (state, batch, learning_rate_fn),\n {},\n )\n\n\n@functools.partial(jax.jit, static_argnums=(2,))\ndef bench_train_step(state, batch, learning_rate_fn):\n\n def compute_metrics(logits, labels):\n loss = cross_entropy_loss(logits, labels)\n accuracy = jnp.mean(jnp.argmax(logits, -1) == labels)\n metrics = {\n 'loss': loss,\n 'accuracy': accuracy,\n }\n return metrics\n\n def loss_fn(params):\n logits, new_model_state = state.apply_fn(\n {'params': params, 'batch_stats': state.batch_stats},\n batch['image'],\n mutable=['batch_stats'],\n )\n loss = cross_entropy_loss(logits, batch['label'])\n weight_penalty_params = jax.tree_util.tree_leaves(params)\n weight_decay = 0.0001\n weight_l2 = sum(jnp.sum(x**2) for x in weight_penalty_params if x.ndim > 1)\n weight_penalty = weight_decay * 0.5 * weight_l2\n loss = loss + weight_penalty\n return loss, (new_model_state, logits)\n\n step = state.step\n dynamic_scale = state.dynamic_scale\n lr = learning_rate_fn(step)\n\n if dynamic_scale:\n grad_fn = dynamic_scale.value_and_grad(loss_fn, has_aux=True)\n dynamic_scale, is_fin, aux, grads = grad_fn(state.params)\n else:\n grad_fn = jax.value_and_grad(loss_fn, has_aux=True)\n aux, grads = grad_fn(state.params)\n\n new_model_state, logits = aux[1]\n metrics = compute_metrics(logits, batch['label'])\n metrics['learning_rate'] = lr\n\n new_state = state.apply_gradients(\n grads=grads,\n batch_stats=new_model_state['batch_stats'],\n )\n if dynamic_scale:\n new_state = new_state.replace(\n opt_state=jax.tree_util.tree_map(\n functools.partial(jnp.where, is_fin),\n new_state.opt_state,\n state.opt_state,\n ),\n params=jax.tree_util.tree_map(\n functools.partial(jnp.where, is_fin), new_state.params, state.params\n ),\n dynamic_scale=dynamic_scale,\n )\n metrics['scale'] = dynamic_scale.scale\n\n return new_state, metrics\n\n\n@google_benchmark.register\n@google_benchmark.option.unit(google_benchmark.kMillisecond)\ndef test_flax_imagenet_trace(state):\n tracing_benchmark.benchmark_tracing(\n get_apply_fn_and_args, imagenet_config.get_config, state\n )\n\n\n@google_benchmark.register\n@google_benchmark.option.unit(google_benchmark.kMillisecond)\ndef test_flax_imagenet_lower(state):\n tracing_benchmark.benchmark_lowering(\n get_apply_fn_and_args, imagenet_config.get_config, state\n )\n\n\nif __name__ == '__main__':\n tracing_benchmark.run_benchmarks()\n" }, { "path": "benchmarks/tracing/lm1b.py", "content": "# Copyright 2025 The JAX Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"LM1B helper functions for benchmarking.\"\"\"\n\nimport functools\nfrom typing import Any\n\nfrom flax import linen as nn\nfrom flax.benchmarks.tracing import tracing_benchmark\nfrom flax.examples.lm1b import models\nfrom flax.examples.lm1b.configs import default as lm1b_config\nfrom flax.training import common_utils\nfrom flax.training import train_state\nimport google_benchmark\nimport jax\nimport jax.numpy as jnp\nimport ml_collections\nimport numpy as np\nimport optax\n\n\ndef rsqrt_schedule(init_value: float, shift: int = 0):\n def schedule(count):\n return init_value * (count + shift) ** -0.5 * shift**0.5\n return schedule\n\n\ndef create_learning_rate_schedule(learning_rate: float, warmup_steps: int):\n return optax.join_schedules(\n [\n optax.linear_schedule(\n init_value=0,\n end_value=learning_rate,\n transition_steps=warmup_steps,\n ),\n rsqrt_schedule(init_value=learning_rate, shift=warmup_steps),\n ],\n boundaries=[warmup_steps],\n )\n\n\ndef compute_weighted_cross_entropy(\n logits, targets, weights=None, label_smoothing=0.0\n):\n if logits.ndim != targets.ndim + 1:\n raise ValueError(\n 'Incorrect shapes. Got shape %s logits and %s targets'\n % (str(logits.shape), str(targets.shape))\n )\n vocab_size = logits.shape[-1]\n confidence = 1.0 - label_smoothing\n low_confidence = (1.0 - confidence) / (vocab_size - 1)\n normalizing_constant = -(\n confidence * jnp.log(confidence)\n + (vocab_size - 1) * low_confidence * jnp.log(low_confidence + 1e-20)\n )\n soft_targets = common_utils.onehot(\n targets, vocab_size, on_value=confidence, off_value=low_confidence\n )\n\n loss = -jnp.sum(soft_targets * nn.log_softmax(logits), axis=-1)\n loss = loss - normalizing_constant\n\n normalizing_factor = np.prod(targets.shape)\n if weights is not None:\n loss = loss * weights\n normalizing_factor = weights.sum()\n\n return loss.sum(), normalizing_factor\n\n\ndef compute_weighted_accuracy(logits, targets, weights=None):\n if logits.ndim != targets.ndim + 1:\n raise ValueError(\n 'Incorrect shapes. Got shape %s logits and %s targets'\n % (str(logits.shape), str(targets.shape))\n )\n loss = jnp.equal(jnp.argmax(logits, axis=-1), targets)\n normalizing_factor = np.prod(logits.shape[:-1])\n if weights is not None:\n loss = loss * weights\n normalizing_factor = weights.sum()\n\n return loss.sum(), normalizing_factor\n\n\ndef get_fake_batch(config: ml_collections.ConfigDict) -> dict[str, jnp.ndarray]:\n batch_size = config.per_device_batch_size\n max_len = config.max_target_length\n\n inputs = jax.random.randint(\n jax.random.key(0),\n (batch_size, max_len),\n minval=0,\n maxval=config.vocab_size,\n dtype=jnp.int32,\n )\n inputs_position = jnp.tile(\n jnp.arange(max_len, dtype=jnp.int32), (batch_size, 1)\n )\n inputs_segmentation = jnp.ones((batch_size, max_len), dtype=jnp.int32)\n\n return {\n 'inputs': inputs,\n 'inputs_position': inputs_position,\n 'inputs_segmentation': inputs_segmentation,\n }\n\n\n@functools.partial(jax.jit, static_argnums=(2, 3))\ndef bench_train_step(state, batch, config, learning_rate_fn):\n\n def compute_metrics(logits, labels, weights):\n loss, weight_sum = compute_weighted_cross_entropy(\n logits, labels, weights, 0.0\n )\n acc, _ = compute_weighted_accuracy(logits, labels, weights)\n metrics = {\n 'loss': loss,\n 'accuracy': acc,\n 'denominator': weight_sum,\n }\n return metrics\n\n inputs = batch['inputs']\n inputs_positions = batch['inputs_position']\n inputs_segmentation = batch['inputs_segmentation']\n\n weights = jnp.where(inputs > 0, 1, 0).astype(jnp.float32)\n dropout_rng = jax.random.fold_in(jax.random.key(0), state.step)\n\n def loss_fn(params):\n logits = models.TransformerLM(config).apply(\n {'params': params},\n inputs,\n inputs_positions=inputs_positions,\n inputs_segmentation=inputs_segmentation,\n rngs={'dropout': dropout_rng},\n )\n\n loss, weight_sum = compute_weighted_cross_entropy(\n logits, inputs, weights, 0.0\n )\n mean_loss = loss / weight_sum\n return mean_loss, logits\n\n step = state.step\n lr = learning_rate_fn(step)\n grad_fn = jax.value_and_grad(loss_fn, has_aux=True)\n (_, logits), grads = grad_fn(state.params)\n new_state = state.apply_gradients(grads=grads)\n\n metrics = compute_metrics(logits, inputs, weights)\n metrics['learning_rate'] = lr\n\n return new_state, metrics\n\n\ndef get_apply_fn_and_args(\n config: ml_collections.ConfigDict,\n) -> tuple[Any, tuple[Any, ...], dict[str, Any]]:\n train_config = models.TransformerConfig(\n vocab_size=config.vocab_size,\n output_vocab_size=config.vocab_size,\n logits_via_embedding=config.logits_via_embedding,\n dtype=jnp.bfloat16 if config.use_bfloat16 else jnp.float32,\n emb_dim=config.emb_dim,\n num_heads=config.num_heads,\n num_layers=config.num_layers,\n qkv_dim=config.qkv_dim,\n mlp_dim=config.mlp_dim,\n max_len=max(config.max_target_length, config.max_eval_target_length),\n dropout_rate=config.dropout_rate,\n attention_dropout_rate=config.attention_dropout_rate,\n deterministic=False,\n decode=False,\n kernel_init=jax.nn.initializers.xavier_uniform(),\n bias_init=jax.nn.initializers.normal(stddev=1e-6),\n )\n\n model = models.TransformerLM(train_config)\n\n learning_rate_fn = create_learning_rate_schedule(\n learning_rate=config.learning_rate, warmup_steps=config.warmup_steps\n )\n\n optimizer = optax.adamw(\n learning_rate_fn,\n b1=0.9,\n b2=0.98,\n eps=1e-9,\n weight_decay=config.weight_decay,\n )\n\n rng = jax.random.key(0)\n init_rng, _ = jax.random.split(rng)\n\n initial_variables = model.init(\n init_rng,\n jnp.ones(\n (config.per_device_batch_size, config.max_target_length), jnp.int32\n ),\n jnp.ones(\n (config.per_device_batch_size, config.max_target_length), jnp.int32\n ),\n jnp.ones(\n (config.per_device_batch_size, config.max_target_length), jnp.int32\n ),\n )\n\n state = train_state.TrainState.create(\n apply_fn=model.apply,\n params=initial_variables['params'],\n tx=optimizer,\n )\n\n batch = get_fake_batch(config)\n\n return (\n bench_train_step,\n (state, batch, train_config, learning_rate_fn),\n {},\n )\n\n\n@google_benchmark.register\n@google_benchmark.option.unit(google_benchmark.kMillisecond)\ndef test_flax_lm1b_trace(state):\n tracing_benchmark.benchmark_tracing(\n get_apply_fn_and_args, lm1b_config.get_config, state\n )\n\n\n@google_benchmark.register\n@google_benchmark.option.unit(google_benchmark.kMillisecond)\ndef test_flax_lm1b_lower(state):\n tracing_benchmark.benchmark_lowering(\n get_apply_fn_and_args, lm1b_config.get_config, state\n )\n\n\nif __name__ == '__main__':\n tracing_benchmark.run_benchmarks()\n" }, { "path": "benchmarks/tracing/mnist.py", "content": "# Copyright 2025 The JAX Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"MNIST helper functions.\"\"\"\n\nfrom functools import partial\nfrom typing import Any\n\nfrom flax import nnx\nfrom flax.benchmarks.tracing import tracing_benchmark\nfrom flax.examples.mnist.configs import default as mnist_config\nimport google_benchmark\nimport jax\nimport jax.numpy as jnp\nimport ml_collections\nimport optax\n\n\nclass CNN(nnx.Module):\n\n def __init__(self, rngs: nnx.Rngs):\n self.conv1 = nnx.Conv(1, 32, kernel_size=(3, 3), rngs=rngs)\n self.batch_norm1 = nnx.BatchNorm(32, rngs=rngs)\n self.dropout1 = nnx.Dropout(rate=0.025)\n self.conv2 = nnx.Conv(32, 64, kernel_size=(3, 3), rngs=rngs)\n self.batch_norm2 = nnx.BatchNorm(64, rngs=rngs)\n self.avg_pool = partial(nnx.avg_pool, window_shape=(2, 2), strides=(2, 2))\n self.linear1 = nnx.Linear(3136, 256, rngs=rngs)\n self.dropout2 = nnx.Dropout(rate=0.025)\n self.linear2 = nnx.Linear(256, 10, rngs=rngs)\n\n def __call__(self, x, rngs: nnx.Rngs):\n x = self.avg_pool(nnx.relu(self.batch_norm1(self.dropout1(self.conv1(x), rngs=rngs))))\n x = self.avg_pool(nnx.relu(self.batch_norm2(self.conv2(x))))\n x = x.reshape(x.shape[0], -1) # flatten\n x = nnx.relu(self.dropout2(self.linear1(x), rngs=rngs))\n x = self.linear2(x)\n return x\n\n\ndef loss_fn(model: CNN, batch, rngs):\n logits = model(batch['image'], rngs)\n loss = optax.softmax_cross_entropy_with_integer_labels(\n logits=logits, labels=batch['label']\n ).mean()\n return loss, logits\n\n\ndef get_fake_batch(batch_size: int) -> dict[str, Any]:\n rng = jax.random.PRNGKey(0)\n images = jax.random.normal(rng, (batch_size, 28, 28, 1), jnp.float32)\n labels = jax.random.randint(rng, (batch_size,), 0, 10, jnp.int32)\n return {'image': images, 'label': labels}\n\n\ndef get_apply_fn_and_args(\n config: ml_collections.ConfigDict,\n) -> tuple[Any, tuple[Any, ...], dict[str, Any]]:\n model = CNN(rngs=nnx.Rngs(0))\n batch = get_fake_batch(config.batch_size)\n rngs = nnx.Rngs(0)\n loss_fn_jit = jax.jit(loss_fn)\n return (\n loss_fn_jit,\n (model, batch, rngs),\n dict(),\n )\n\n\n@google_benchmark.register\n@google_benchmark.option.unit(google_benchmark.kMillisecond)\ndef test_flax_mnist_trace(state):\n tracing_benchmark.benchmark_tracing(\n get_apply_fn_and_args, mnist_config.get_config, state\n )\n\n\n@google_benchmark.register\n@google_benchmark.option.unit(google_benchmark.kMillisecond)\ndef test_flax_mnist_lower(state):\n tracing_benchmark.benchmark_lowering(\n get_apply_fn_and_args, mnist_config.get_config, state\n )\n\n\nif __name__ == '__main__':\n tracing_benchmark.run_benchmarks()\n" }, { "path": "benchmarks/tracing/nlp_seq.py", "content": "# Copyright 2025 The JAX Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"NLP Sequence Tagging helper functions for benchmarking.\"\"\"\n\nimport functools\nfrom typing import Any\n\nfrom flax import linen as nn\nfrom flax.benchmarks.tracing import tracing_benchmark\nfrom flax.examples.nlp_seq import models\nfrom flax.examples.nlp_seq.configs import default as nlp_seq_config\nfrom flax.training import common_utils\nfrom flax.training import train_state\nimport google_benchmark\nimport jax\nimport jax.numpy as jnp\nimport ml_collections\nimport numpy as np\nimport optax\n\n\ndef create_learning_rate_scheduler(\n factors='constant * linear_warmup * rsqrt_decay',\n base_learning_rate=0.5,\n warmup_steps=8000,\n decay_factor=0.5,\n steps_per_decay=20000,\n steps_per_cycle=100000,\n):\n factors = [n.strip() for n in factors.split('*')]\n\n def step_fn(step):\n ret = 1.0\n for name in factors:\n if name == 'constant':\n ret *= base_learning_rate\n elif name == 'linear_warmup':\n ret *= jnp.minimum(1.0, step / warmup_steps)\n elif name == 'rsqrt_decay':\n ret /= jnp.sqrt(jnp.maximum(step, warmup_steps))\n elif name == 'rsqrt_normalized_decay':\n ret *= jnp.sqrt(warmup_steps)\n ret /= jnp.sqrt(jnp.maximum(step, warmup_steps))\n elif name == 'decay_every':\n ret *= decay_factor ** (step // steps_per_decay)\n elif name == 'cosine_decay':\n progress = jnp.maximum(\n 0.0, (step - warmup_steps) / float(steps_per_cycle)\n )\n ret *= jnp.maximum(\n 0.0, 0.5 * (1.0 + jnp.cos(jnp.pi * (progress % 1.0)))\n )\n else:\n raise ValueError('Unknown factor %s.' % name)\n return jnp.asarray(ret, dtype=jnp.float32)\n\n return step_fn\n\n\ndef compute_weighted_cross_entropy(logits, targets, weights=None):\n if logits.ndim != targets.ndim + 1:\n raise ValueError(\n 'Incorrect shapes. Got shape %s logits and %s targets'\n % (str(logits.shape), str(targets.shape))\n )\n onehot_targets = common_utils.onehot(targets, logits.shape[-1])\n loss = -jnp.sum(onehot_targets * nn.log_softmax(logits), axis=-1)\n normalizing_factor = onehot_targets.sum()\n if weights is not None:\n loss = loss * weights\n normalizing_factor = weights.sum()\n\n return loss.sum(), normalizing_factor\n\n\ndef compute_weighted_accuracy(logits, targets, weights=None):\n if logits.ndim != targets.ndim + 1:\n raise ValueError(\n 'Incorrect shapes. Got shape %s logits and %s targets'\n % (str(logits.shape), str(targets.shape))\n )\n loss = jnp.equal(jnp.argmax(logits, axis=-1), targets)\n normalizing_factor = np.prod(logits.shape[:-1])\n if weights is not None:\n loss = loss * weights\n normalizing_factor = weights.sum()\n\n return loss.sum(), normalizing_factor\n\n\ndef get_fake_batch(config: ml_collections.ConfigDict) -> dict[str, jnp.ndarray]:\n batch_size = config.batch_size\n max_len = config.max_length\n vocab_size = config.vocab_size\n\n inputs = jax.random.randint(\n jax.random.key(0),\n (batch_size, max_len),\n minval=0,\n maxval=vocab_size,\n dtype=jnp.int32,\n )\n targets = jax.random.randint(\n jax.random.key(1),\n (batch_size, max_len),\n minval=0,\n maxval=config.output_vocab_size,\n dtype=jnp.int32,\n )\n\n return {\n 'inputs': inputs,\n 'targets': targets,\n }\n\n\n@functools.partial(jax.jit, static_argnums=(2, 3))\ndef bench_train_step(state, batch, config, learning_rate_fn):\n\n def local_compute_metrics(logits, labels, weights):\n loss, weight_sum = compute_weighted_cross_entropy(\n logits, labels, weights\n )\n acc, _ = compute_weighted_accuracy(logits, labels, weights)\n metrics = {\n 'loss': loss,\n 'accuracy': acc,\n 'denominator': weight_sum,\n }\n return metrics\n\n inputs = batch['inputs']\n targets = batch['targets']\n\n weights = jnp.where(targets > 0, 1, 0).astype(jnp.float32)\n dropout_rng = jax.random.fold_in(jax.random.key(0), state.step)\n\n def loss_fn(params):\n model = models.Transformer(config)\n logits = model.apply(\n {'params': params},\n inputs=inputs,\n train=True,\n rngs={'dropout': dropout_rng},\n )\n\n loss, weight_sum = compute_weighted_cross_entropy(\n logits, targets, weights\n )\n mean_loss = loss / weight_sum\n return mean_loss, logits\n\n step = state.step\n lr = learning_rate_fn(step)\n grad_fn = jax.value_and_grad(loss_fn, has_aux=True)\n (_, logits), grads = grad_fn(state.params)\n new_state = state.apply_gradients(grads=grads)\n\n metrics = local_compute_metrics(logits, targets, weights)\n metrics['learning_rate'] = lr\n\n return new_state, metrics\n\n\ndef get_apply_fn_and_args(\n config: ml_collections.ConfigDict,\n) -> tuple[Any, tuple[Any, ...], dict[str, Any]]:\n if not hasattr(config, 'vocab_size'):\n config.vocab_size = 30000\n if not hasattr(config, 'output_vocab_size'):\n config.output_vocab_size = 50\n\n model_config = models.TransformerConfig(\n vocab_size=config.vocab_size,\n output_vocab_size=config.output_vocab_size,\n max_len=config.max_length,\n )\n\n model = models.Transformer(model_config)\n\n learning_rate_fn = create_learning_rate_scheduler(\n base_learning_rate=config.learning_rate\n )\n\n optimizer = optax.adamw(\n learning_rate_fn,\n b1=0.9,\n b2=0.98,\n eps=1e-9,\n weight_decay=config.weight_decay,\n )\n\n rng = jax.random.key(0)\n init_rng, _ = jax.random.split(rng)\n\n init_batch = jnp.ones((config.batch_size, config.max_length), jnp.int32)\n initial_variables = model.init(init_rng, inputs=init_batch, train=False)\n\n state = train_state.TrainState.create(\n apply_fn=model.apply,\n params=initial_variables['params'],\n tx=optimizer,\n )\n\n batch = get_fake_batch(config)\n\n return (\n bench_train_step,\n (state, batch, model_config, learning_rate_fn),\n {},\n )\n\n\n@google_benchmark.register\n@google_benchmark.option.unit(google_benchmark.kMillisecond)\ndef test_flax_nlp_seq_trace(state):\n tracing_benchmark.benchmark_tracing(\n get_apply_fn_and_args, nlp_seq_config.get_config, state\n )\n\n\n@google_benchmark.register\n@google_benchmark.option.unit(google_benchmark.kMillisecond)\ndef test_flax_nlp_seq_lower(state):\n tracing_benchmark.benchmark_lowering(\n get_apply_fn_and_args, nlp_seq_config.get_config, state\n )\n\n\nif __name__ == '__main__':\n tracing_benchmark.run_benchmarks()\n" }, { "path": "benchmarks/tracing/ogbg_molpcba.py", "content": "# Copyright 2025 The JAX Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"OGBG-MolPCBA helper functions for benchmarking.\"\"\"\n\nfrom typing import Any\n\nfrom clu import metrics\nimport flax\nimport flax.linen as nn\nfrom flax.benchmarks.tracing import tracing_benchmark\nfrom flax.examples.ogbg_molpcba import models\nfrom flax.examples.ogbg_molpcba.configs import default as ogbg_config\nfrom flax.training import train_state\nimport google_benchmark\nimport jax\nimport jax.numpy as jnp\nimport jraph\nimport ml_collections\nimport optax\n\n\ndef create_model(\n config: ml_collections.ConfigDict, deterministic: bool\n) -> nn.Module:\n if config.model == 'GraphNet':\n return models.GraphNet(\n latent_size=config.latent_size,\n num_mlp_layers=config.num_mlp_layers,\n message_passing_steps=config.message_passing_steps,\n output_globals_size=config.num_classes,\n dropout_rate=config.dropout_rate,\n skip_connections=config.skip_connections,\n layer_norm=config.layer_norm,\n use_edge_model=config.use_edge_model,\n deterministic=deterministic,\n )\n if config.model == 'GraphConvNet':\n return models.GraphConvNet(\n latent_size=config.latent_size,\n num_mlp_layers=config.num_mlp_layers,\n message_passing_steps=config.message_passing_steps,\n output_globals_size=config.num_classes,\n dropout_rate=config.dropout_rate,\n skip_connections=config.skip_connections,\n layer_norm=config.layer_norm,\n deterministic=deterministic,\n )\n raise ValueError(f'Unsupported model: {config.model}.')\n\n\ndef create_optimizer(\n config: ml_collections.ConfigDict,\n) -> optax.GradientTransformation:\n if config.optimizer == 'adam':\n return optax.adam(learning_rate=config.learning_rate)\n if config.optimizer == 'sgd':\n return optax.sgd(\n learning_rate=config.learning_rate, momentum=config.momentum\n )\n raise ValueError(f'Unsupported optimizer: {config.optimizer}.')\n\n\ndef binary_cross_entropy_with_mask(\n *, logits: jnp.ndarray, labels: jnp.ndarray, mask: jnp.ndarray\n):\n assert logits.shape == labels.shape == mask.shape\n assert len(logits.shape) == 2\n\n labels = jnp.where(mask, labels, -1)\n\n positive_logits = logits >= 0\n relu_logits = jnp.where(positive_logits, logits, 0)\n abs_logits = jnp.where(positive_logits, logits, -logits)\n return relu_logits - (logits * labels) + (jnp.log(1 + jnp.exp(-abs_logits)))\n\n\ndef predictions_match_labels(\n *, logits: jnp.ndarray, labels: jnp.ndarray, **kwargs\n) -> jnp.ndarray:\n del kwargs\n preds = logits > 0\n return (preds == labels).astype(jnp.float32)\n\n\ndef replace_globals(graphs: jraph.GraphsTuple) -> jraph.GraphsTuple:\n return graphs._replace(globals=jnp.ones([graphs.n_node.shape[0], 1]))\n\n\ndef get_predicted_logits(\n state: train_state.TrainState,\n graphs: jraph.GraphsTuple,\n rngs: dict[str, jnp.ndarray] | None,\n) -> jnp.ndarray:\n pred_graphs = state.apply_fn(state.params, graphs, rngs=rngs)\n logits = pred_graphs.globals\n return logits\n\n\ndef get_valid_mask(\n labels: jnp.ndarray, graphs: jraph.GraphsTuple\n) -> jnp.ndarray:\n labels_mask = ~jnp.isnan(labels)\n graph_mask = jraph.get_graph_padding_mask(graphs)\n return labels_mask & graph_mask[:, None]\n\n\n@flax.struct.dataclass\nclass TrainMetrics(metrics.Collection):\n accuracy: metrics.Average.from_fun(predictions_match_labels)\n loss: metrics.Average.from_output('loss')\n\n\n@jax.jit\ndef ogbg_train_step(\n state: train_state.TrainState,\n graphs: jraph.GraphsTuple,\n rngs: dict[str, jnp.ndarray],\n) -> tuple[train_state.TrainState, metrics.Collection]:\n\n def loss_fn(params, graphs):\n curr_state = state.replace(params=params)\n\n labels = graphs.globals\n\n graphs = replace_globals(graphs)\n\n logits = get_predicted_logits(curr_state, graphs, rngs)\n mask = get_valid_mask(labels, graphs)\n loss = binary_cross_entropy_with_mask(\n logits=logits, labels=labels, mask=mask\n )\n mean_loss = jnp.sum(jnp.where(mask, loss, 0)) / jnp.sum(mask)\n\n return mean_loss, (loss, logits, labels, mask)\n\n grad_fn = jax.value_and_grad(loss_fn, has_aux=True)\n (_, (loss, logits, labels, mask)), grads = grad_fn(state.params, graphs)\n state = state.apply_gradients(grads=grads)\n\n metrics_update = TrainMetrics.single_from_model_output(\n loss=loss, logits=logits, labels=labels, mask=mask\n )\n return state, metrics_update\n\n\ndef get_fake_graphs(config: ml_collections.ConfigDict) -> jraph.GraphsTuple:\n rng = jax.random.key(0)\n n_graphs = config.batch_size\n n_total_nodes = n_graphs * 20\n n_total_edges = n_graphs * 40\n\n nodes = jax.random.normal(rng, (n_total_nodes, 9))\n edges = jax.random.normal(rng, (n_total_edges, 3))\n senders = jax.random.randint(rng, (n_total_edges,), 0, n_total_nodes)\n receivers = jax.random.randint(rng, (n_total_edges,), 0, n_total_nodes)\n n_node = jnp.full((n_graphs,), 20, dtype=jnp.int32)\n n_edge = jnp.full((n_graphs,), 40, dtype=jnp.int32)\n globals_ = jax.random.bernoulli(\n rng, shape=(n_graphs, config.num_classes)\n ).astype(jnp.float32)\n\n return jraph.GraphsTuple(\n nodes=nodes,\n edges=edges,\n senders=senders,\n receivers=receivers,\n n_node=n_node,\n n_edge=n_edge,\n globals=globals_,\n )\n\n\ndef get_apply_fn_and_args(\n config: ml_collections.ConfigDict,\n) -> tuple[Any, tuple[Any, ...], dict[str, Any]]:\n rng = jax.random.key(0)\n rng, init_rng, dropout_rng = jax.random.split(rng, 3)\n\n graphs = get_fake_graphs(config)\n\n init_net = create_model(config, deterministic=True)\n init_graphs = replace_globals(graphs)\n params = jax.jit(init_net.init)(init_rng, init_graphs)\n\n tx = create_optimizer(config)\n\n net = create_model(config, deterministic=False)\n state = train_state.TrainState.create(\n apply_fn=net.apply, params=params, tx=tx\n )\n\n return (\n ogbg_train_step,\n (state, graphs, {'dropout': dropout_rng}),\n {},\n )\n\n\n@google_benchmark.register\n@google_benchmark.option.unit(google_benchmark.kMillisecond)\ndef test_flax_ogbg_molpcba_trace(state):\n tracing_benchmark.benchmark_tracing(\n get_apply_fn_and_args, ogbg_config.get_config, state\n )\n\n\n@google_benchmark.register\n@google_benchmark.option.unit(google_benchmark.kMillisecond)\ndef test_flax_ogbg_molpcba_lower(state):\n tracing_benchmark.benchmark_lowering(\n get_apply_fn_and_args, ogbg_config.get_config, state\n )\n\n\nif __name__ == '__main__':\n tracing_benchmark.run_benchmarks()\n" }, { "path": "benchmarks/tracing/ppo.py", "content": "# Copyright 2025 The JAX Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"PPO helper functions for RL benchmarking.\"\"\"\n\nfrom typing import Any\n\nfrom flax.benchmarks.tracing import tracing_benchmark\nfrom flax.examples.ppo import models\nfrom flax.examples.ppo import ppo_lib\nfrom flax.examples.ppo.configs import default as ppo_config\nimport google_benchmark\nimport jax\nimport jax.numpy as jnp\nimport ml_collections\n\n\ndef get_fake_batch(batch_size: int = 256) -> tuple[jnp.ndarray, ...]:\n \"\"\"Generate a batch of fake Atari observations and trajectories.\n\n Args:\n batch_size: Size of the minibatch.\n\n Returns:\n A tuple of (states, actions, old_log_probs, returns, advantages).\n \"\"\"\n # Atari observations: (batch_size, height, width, stacked_frames)\n states = jax.random.randint(\n jax.random.key(0),\n (batch_size, 84, 84, 4),\n minval=0,\n maxval=256,\n dtype=jnp.int32,\n )\n\n # Actions: discrete action space (e.g., 6 actions for Pong)\n actions = jax.random.randint(\n jax.random.key(1), (batch_size,), minval=0, maxval=6, dtype=jnp.int32\n )\n\n # Old log probabilities from behavior policy\n old_log_probs = jax.random.normal(\n jax.random.key(2), (batch_size,), dtype=jnp.float32\n )\n\n # Returns (discounted cumulative rewards)\n returns = jax.random.normal(\n jax.random.key(3), (batch_size,), dtype=jnp.float32\n )\n\n # Advantages (GAE advantages)\n advantages = jax.random.normal(\n jax.random.key(4), (batch_size,), dtype=jnp.float32\n )\n\n return (states, actions, old_log_probs, returns, advantages)\n\n\ndef get_apply_fn_and_args(\n config: ml_collections.ConfigDict,\n) -> tuple[Any, tuple[Any, ...], dict[str, Any]]:\n \"\"\"Returns the apply function and args for the given config.\n\n Args:\n config: The training configuration.\n\n Returns:\n A tuple of the apply function, args, and kwargs.\n \"\"\"\n # Create model (6 actions for Pong)\n num_outputs = 6\n model = models.ActorCritic(num_outputs=num_outputs)\n\n # Initialize model parameters\n rng = jax.random.key(0)\n init_shape = jnp.ones((1, 84, 84, 4), jnp.float32)\n initial_params = model.init(rng, init_shape)['params']\n\n # Create train state\n # For benchmarking, we don't need actual training loops, just one step\n train_steps = 1000 # Dummy value for state creation\n state = ppo_lib.create_train_state(initial_params, model, config, train_steps)\n\n # Generate fake trajectories\n trajectories = get_fake_batch(config.batch_size)\n\n # PPO hyperparameters\n clip_param = config.clip_param\n vf_coeff = config.vf_coeff\n entropy_coeff = config.entropy_coeff\n\n # ppo_lib.train_step is already JIT-compiled with static batch_size\n return (\n ppo_lib.train_step,\n (state, trajectories, config.batch_size),\n {\n 'clip_param': clip_param,\n 'vf_coeff': vf_coeff,\n 'entropy_coeff': entropy_coeff,\n },\n )\n\n\n@google_benchmark.register\n@google_benchmark.option.unit(google_benchmark.kMillisecond)\ndef test_flax_ppo_trace(state):\n tracing_benchmark.benchmark_tracing(\n get_apply_fn_and_args, ppo_config.get_config, state\n )\n\n\n@google_benchmark.register\n@google_benchmark.option.unit(google_benchmark.kMillisecond)\ndef test_flax_ppo_lower(state):\n tracing_benchmark.benchmark_lowering(\n get_apply_fn_and_args, ppo_config.get_config, state\n )\n\n\nif __name__ == '__main__':\n tracing_benchmark.run_benchmarks()\n" }, { "path": "benchmarks/tracing/requirements.txt", "content": "absl-py\nflax\ngoogle-benchmark\njax\nml_collections\nnumpy\noptax" }, { "path": "benchmarks/tracing/run_all_benchmarks.sh", "content": "#!/bin/bash\nset -e\n\nexport XLA_FLAGS=--xla_force_host_platform_device_count=8\n\nTARGETS=(\n mnist\n vae\n sst2\n gemma\n imagenet\n seq2seq\n lm1b\n nlp_seq\n ogbg_molpcba\n wmt\n ppo\n)\n\nfor target in \"${TARGETS[@]}\"; do\n echo \"============================================\"\n echo \"Running benchmark: ${target}\"\n echo \"============================================\"\n benchy \"third_party/py/flax/benchmarks/tracing:${target}\"\n echo \"\"\ndone\n" }, { "path": "benchmarks/tracing/seq2seq.py", "content": "# Copyright 2025 The JAX Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"Seq2Seq helper functions.\"\"\"\n\nfrom typing import Any\n\nfrom flax import linen as nn\nfrom flax.benchmarks.tracing import tracing_benchmark\nfrom flax.examples.seq2seq import models\nfrom flax.examples.seq2seq.configs import default as seq2seq_config\nfrom flax.examples.seq2seq.input_pipeline import CharacterTable\nfrom flax.examples.seq2seq.input_pipeline import get_sequence_lengths\nfrom flax.examples.seq2seq.input_pipeline import mask_sequences\nfrom flax.training import train_state\nimport google_benchmark\nimport jax\nimport jax.numpy as jnp\nimport ml_collections\nimport optax\n\n\ndef cross_entropy_loss(logits, labels, lengths):\n xe = jnp.sum(nn.log_softmax(logits) * labels, axis=-1)\n masked_xe = jnp.mean(mask_sequences(xe, lengths))\n return -masked_xe\n\n\ndef compute_metrics(logits, labels, eos_id):\n lengths = get_sequence_lengths(labels, eos_id)\n loss = cross_entropy_loss(logits, labels, lengths)\n token_accuracy = jnp.argmax(logits, -1) == jnp.argmax(labels, -1)\n sequence_accuracy = (\n jnp.sum(mask_sequences(token_accuracy, lengths), axis=-1) == lengths\n )\n accuracy = jnp.mean(sequence_accuracy)\n metrics = {\n 'loss': loss,\n 'accuracy': accuracy,\n }\n return metrics\n\n\n@jax.jit\ndef seq2seq_train_step(state, batch, lstm_rng, eos_id):\n labels = batch['answer'][:, 1:]\n lstm_key = jax.random.fold_in(lstm_rng, state.step)\n\n def loss_fn(params):\n logits, _ = state.apply_fn(\n {'params': params},\n batch['query'],\n batch['answer'],\n rngs={'lstm': lstm_key},\n )\n loss = cross_entropy_loss(\n logits, labels, get_sequence_lengths(labels, eos_id)\n )\n return loss, logits\n\n grad_fn = jax.value_and_grad(loss_fn, has_aux=True)\n (_, logits), grads = grad_fn(state.params)\n state = state.apply_gradients(grads=grads)\n metrics = compute_metrics(logits, labels, eos_id)\n\n return state, metrics\n\n\ndef get_fake_batch(batch_size: int, ctable: CharacterTable) -> dict[str, Any]:\n return ctable.get_batch(batch_size)\n\n\ndef get_apply_fn_and_args(\n config: ml_collections.ConfigDict,\n) -> tuple[Any, tuple[Any, ...], dict[str, Any]]:\n rng = jax.random.key(0)\n ctable = CharacterTable(\"0123456789+= \", config.max_len_query_digit)\n\n model = models.Seq2seq(\n teacher_force=False,\n hidden_size=config.hidden_size,\n eos_id=ctable.eos_id,\n vocab_size=ctable.vocab_size,\n )\n\n rng1, rng2 = jax.random.split(rng)\n params = model.init(\n {'params': rng1, 'lstm': rng2},\n jnp.ones(ctable.encoder_input_shape, jnp.float32),\n jnp.ones(ctable.decoder_input_shape, jnp.float32),\n )['params']\n\n tx = optax.adam(config.learning_rate)\n state = train_state.TrainState.create(\n apply_fn=model.apply, params=params, tx=tx\n )\n\n batch = get_fake_batch(config.batch_size, ctable)\n return seq2seq_train_step, (state, batch, rng, ctable.eos_id), {}\n\n\n@google_benchmark.register\n@google_benchmark.option.unit(google_benchmark.kMillisecond)\ndef test_flax_seq2seq_trace(state):\n tracing_benchmark.benchmark_tracing(\n get_apply_fn_and_args, seq2seq_config.get_config, state\n )\n\n\n@google_benchmark.register\n@google_benchmark.option.unit(google_benchmark.kMillisecond)\ndef test_flax_seq2seq_lower(state):\n tracing_benchmark.benchmark_lowering(\n get_apply_fn_and_args, seq2seq_config.get_config, state\n )\n\n\nif __name__ == \"__main__\":\n tracing_benchmark.run_benchmarks()\n" }, { "path": "benchmarks/tracing/sst2.py", "content": "# Copyright 2025 The JAX Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"SST2 helper functions.\"\"\"\n\nfrom typing import Any\n\nfrom flax.benchmarks.tracing import tracing_benchmark\nfrom flax.examples.sst2 import models\nfrom flax.examples.sst2.configs import default as sst2_config\nfrom flax.training import train_state as train_state_lib\nimport google_benchmark\nimport jax\nimport jax.numpy as jnp\nimport ml_collections\nimport optax\n\nArray = jnp.ndarray\nTrainState = train_state_lib.TrainState\n\n\n@jax.vmap\ndef sigmoid_cross_entropy_with_logits(*, labels: Array, logits: Array) -> Array:\n zeros = jnp.zeros_like(logits, dtype=logits.dtype)\n condition = logits >= zeros\n relu_logits = jnp.where(condition, logits, zeros)\n neg_abs_logits = jnp.where(condition, -logits, logits)\n return relu_logits - logits * labels + jnp.log1p(jnp.exp(neg_abs_logits))\n\n\ndef model_from_config(config: ml_collections.ConfigDict):\n model = models.TextClassifier(\n embedding_size=config.embedding_size,\n hidden_size=config.hidden_size,\n vocab_size=config.vocab_size,\n output_size=config.output_size,\n dropout_rate=config.dropout_rate,\n word_dropout_rate=config.word_dropout_rate,\n unk_idx=config.unk_idx,\n )\n return model\n\n\ndef get_initial_params(rng, model):\n token_ids = jnp.ones((2, 3), jnp.int32)\n lengths = jnp.ones((2,), dtype=jnp.int32)\n variables = model.init(rng, token_ids, lengths, deterministic=True)\n return variables['params']\n\n\ndef create_train_state(rng, config: ml_collections.ConfigDict, model):\n params = get_initial_params(rng, model)\n tx = optax.chain(\n optax.sgd(learning_rate=config.learning_rate, momentum=config.momentum),\n optax.add_decayed_weights(weight_decay=config.weight_decay),\n )\n state = TrainState.create(apply_fn=model.apply, params=params, tx=tx)\n return state\n\n\ndef compute_metrics(*, labels: Array, logits: Array):\n if labels.ndim == 1:\n labels = jnp.expand_dims(labels, axis=1)\n loss = sigmoid_cross_entropy_with_logits(labels=labels, logits=logits)\n binary_predictions = logits >= 0.0\n binary_accuracy = jnp.equal(binary_predictions, labels)\n return {\n 'loss': jnp.sum(loss),\n 'accuracy': jnp.sum(binary_accuracy),\n 'count': logits.shape[0],\n }\n\n\ndef train_step(\n state: TrainState,\n batch: dict[str, Array],\n rngs: dict[str, Any],\n) -> tuple[TrainState, Any]:\n step = state.step\n rngs = {name: jax.random.fold_in(rng, step) for name, rng in rngs.items()}\n\n def loss_fn(params):\n variables = {'params': params}\n logits = state.apply_fn(\n variables,\n batch['token_ids'],\n batch['length'],\n deterministic=False,\n rngs=rngs,\n )\n\n labels = batch['label']\n if labels.ndim == 1:\n labels = jnp.expand_dims(labels, 1)\n loss = jnp.mean(\n sigmoid_cross_entropy_with_logits(labels=labels, logits=logits)\n )\n return loss, logits\n\n grad_fn = jax.value_and_grad(loss_fn, has_aux=True)\n value, grads = grad_fn(state.params)\n (_, logits) = value\n\n new_state = state.apply_gradients(grads=grads)\n metrics = compute_metrics(labels=batch['label'], logits=logits)\n return new_state, metrics\n\n\ndef get_fake_batch(batch_size: int) -> dict[str, Any]:\n rng = jax.random.key(0)\n max_length = 60\n token_ids = jax.random.randint(\n rng, (batch_size, max_length), 0, 1000, jnp.int32\n )\n lengths = jnp.full((batch_size,), max_length, jnp.int32)\n labels = jax.random.uniform(rng, (batch_size,), jnp.float32)\n return {\n 'token_ids': token_ids,\n 'length': lengths,\n 'label': labels,\n }\n\n\ndef get_apply_fn_and_args(\n config: ml_collections.ConfigDict,\n) -> tuple[Any, tuple[Any, ...], dict[str, Any]]:\n rng = jax.random.key(0)\n config = config.copy_and_resolve_references()\n if config.vocab_size is None:\n config.vocab_size = 1000\n model = model_from_config(config)\n state = create_train_state(rng, config, model)\n batch = get_fake_batch(config.batch_size)\n _, dropout_rng = jax.random.split(rng)\n rngs = {'dropout': dropout_rng}\n train_step_jit = jax.jit(train_step)\n return train_step_jit, (state, batch, rngs), {}\n\n\n@google_benchmark.register\n@google_benchmark.option.unit(google_benchmark.kMillisecond)\ndef test_flax_sst2_trace(state):\n tracing_benchmark.benchmark_tracing(\n get_apply_fn_and_args, sst2_config.get_config, state\n )\n\n\n@google_benchmark.register\n@google_benchmark.option.unit(google_benchmark.kMillisecond)\ndef test_flax_sst2_lower(state):\n tracing_benchmark.benchmark_lowering(\n get_apply_fn_and_args, sst2_config.get_config, state\n )\n\n\nif __name__ == '__main__':\n tracing_benchmark.run_benchmarks()\n" }, { "path": "benchmarks/tracing/tracing_benchmark.py", "content": "# Copyright 2025 The JAX Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"Shared benchmark utilities for Jax tracing flax examples.\"\"\"\n\nfrom collections.abc import Callable\nimport sys\nfrom typing import Any\n\nfrom absl import app\nfrom absl import flags\nfrom absl import logging\nimport google_benchmark\nimport jax\n\nflags.DEFINE_enum(\n \"mode\",\n \"trace_and_lower\",\n [\"trace\", \"lower\", \"trace_and_lower\"],\n \"Measure trace, lower, or trace_and_lower.\",\n)\n\n\ndef clear_caches(state):\n state.pause_timing()\n jax.clear_caches()\n state.resume_timing()\n\n\ndef benchmark_tracing(\n get_apply_fn_and_args: Callable[..., Any],\n get_config: Callable[[], Any],\n state: Any,\n) -> None:\n \"\"\"Benchmark for tracing a flax example.\"\"\"\n config = get_config()\n apply_fn, args, kwargs = get_apply_fn_and_args(config)\n while state:\n if flags.FLAGS.mode == 'trace' or flags.FLAGS.mode == 'trace_and_lower':\n _ = apply_fn.trace(*args, **kwargs)\n clear_caches(state)\n\n\ndef benchmark_lowering(\n get_apply_fn_and_args: Callable[..., Any],\n get_config: Callable[[], Any],\n state: Any,\n platform: str = 'tpu',\n) -> None:\n \"\"\"Benchmark for lowering a flax example.\"\"\"\n config = get_config()\n apply_fn, args, kwargs = get_apply_fn_and_args(config)\n traced = apply_fn.trace(*args, **kwargs)\n while state:\n if flags.FLAGS.mode == 'lower' or flags.FLAGS.mode == 'trace_and_lower':\n _ = traced.lower(lowering_platforms=(platform,))\n clear_caches(state)\n\n\ndef run_single_example(\n get_apply_fn_and_args: Callable[..., Any],\n get_config: Callable[[], Any],\n) -> None:\n \"\"\"Run a single example for profiling.\"\"\"\n\n def main(argv):\n del argv\n if flags.FLAGS.mode == 'lower':\n raise ValueError(\n '`--mode=lower` is not supported when profiling a single example.'\n )\n config = get_config()\n apply_fn, args, kwargs, *_ = get_apply_fn_and_args(config)\n traced = apply_fn.trace(*args, **kwargs)\n lowered = traced.lower(lowering_platforms=('tpu',))\n logging.info('lowered: %s', lowered.as_text('hlo'))\n\n app.run(main)\n\n\ndef run_benchmarks() -> None:\n \"\"\"Run registered google_benchmark benchmarks.\"\"\"\n flags.FLAGS(sys.argv, known_only=True)\n flags.FLAGS.mark_as_parsed()\n google_benchmark.main()\n" }, { "path": "benchmarks/tracing/vae.py", "content": "# Copyright 2025 The JAX Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"VAE helper functions.\"\"\"\n\nfrom typing import Any\n\nfrom flax import linen as nn\nfrom flax.benchmarks.tracing import tracing_benchmark\nfrom flax.examples.vae import models\nfrom flax.examples.vae.configs import default as vae_config\nfrom flax.training import train_state\nimport google_benchmark\nimport jax\nimport jax.numpy as jnp\nimport ml_collections\nimport optax\n\n\n@jax.vmap\ndef binary_cross_entropy_with_logits(logits, labels):\n logits = nn.log_sigmoid(logits)\n return -jnp.sum(\n labels * logits + (1.0 - labels) * jnp.log(-jnp.expm1(logits))\n )\n\n\n@jax.vmap\ndef kl_divergence(mean, logvar):\n return -0.5 * jnp.sum(1 + logvar - jnp.square(mean) - jnp.exp(logvar))\n\n\ndef train_step(state, batch, z_rng, latents):\n def loss_fn(params):\n recon_x, mean, logvar = models.model(latents).apply(\n {'params': params}, batch, z_rng\n )\n bce_loss = binary_cross_entropy_with_logits(recon_x, batch).mean()\n kld_loss = kl_divergence(mean, logvar).mean()\n loss = bce_loss + kld_loss\n return loss\n\n grads = jax.grad(loss_fn)(state.params)\n return state.apply_gradients(grads=grads)\n\n\ndef get_fake_batch(batch_size: int) -> Any:\n return jnp.ones((batch_size, 784), jnp.float32)\n\n\ndef get_apply_fn_and_args(\n config: ml_collections.ConfigDict,\n) -> tuple[Any, tuple[Any, ...], dict[str, Any]]:\n rng = jax.random.key(0)\n rng, key = jax.random.split(rng)\n batch = get_fake_batch(config.batch_size)\n params = models.model(config.latents).init(key, batch, rng)['params']\n state = train_state.TrainState.create(\n apply_fn=models.model(config.latents).apply,\n params=params,\n tx=optax.adam(config.learning_rate),\n )\n train_step_jit = jax.jit(train_step, static_argnames=('latents',))\n return (\n train_step_jit,\n (state, batch, rng, config.latents),\n dict(),\n )\n\n\n@google_benchmark.register\n@google_benchmark.option.unit(google_benchmark.kMillisecond)\ndef test_flax_vae_trace(state):\n tracing_benchmark.benchmark_tracing(\n get_apply_fn_and_args, vae_config.get_config, state\n )\n\n\n@google_benchmark.register\n@google_benchmark.option.unit(google_benchmark.kMillisecond)\ndef test_flax_vae_lower(state):\n tracing_benchmark.benchmark_lowering(\n get_apply_fn_and_args, vae_config.get_config, state\n )\n\n\nif __name__ == '__main__':\n tracing_benchmark.run_benchmarks()\n" }, { "path": "benchmarks/tracing/wmt.py", "content": "# Copyright 2025 The JAX Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"WMT helper functions for benchmarking.\"\"\"\n\nimport functools\nfrom typing import Any\n\nfrom flax import linen as nn\nfrom flax.benchmarks.tracing import tracing_benchmark\nfrom flax.examples.wmt import models\nfrom flax.examples.wmt.configs import default as wmt_config\nfrom flax.training import common_utils\nfrom flax.training import dynamic_scale as dynamic_scale_lib\nfrom flax.training import train_state\nimport google_benchmark\nimport jax\nimport jax.numpy as jnp\nimport ml_collections\nimport numpy as np\nimport optax\n\n\nclass TrainState(train_state.TrainState):\n dynamic_scale: dynamic_scale_lib.DynamicScale\n\n\ndef rsqrt_schedule(init_value: float, shift: int = 0):\n def schedule(count):\n return init_value * (count + shift) ** -0.5 * shift**0.5\n return schedule\n\n\ndef create_learning_rate_schedule(learning_rate: float, warmup_steps: int):\n return optax.join_schedules(\n [\n optax.linear_schedule(\n init_value=0,\n end_value=learning_rate,\n transition_steps=warmup_steps,\n ),\n rsqrt_schedule(init_value=learning_rate, shift=warmup_steps),\n ],\n boundaries=[warmup_steps],\n )\n\n\ndef preferred_dtype(config):\n platform = jax.local_devices()[0].platform\n if config.use_mixed_precision:\n if platform == 'tpu':\n return jnp.bfloat16\n elif platform == 'gpu':\n return jnp.float16\n return jnp.float32\n\n\ndef compute_weighted_cross_entropy(\n logits, targets, weights=None, label_smoothing=0.0\n):\n if logits.ndim != targets.ndim + 1:\n raise ValueError(\n 'Incorrect shapes. Got shape %s logits and %s targets'\n % (str(logits.shape), str(targets.shape))\n )\n vocab_size = logits.shape[-1]\n confidence = 1.0 - label_smoothing\n low_confidence = (1.0 - confidence) / (vocab_size - 1)\n normalizing_constant = -(\n confidence * jnp.log(confidence)\n + (vocab_size - 1) * low_confidence * jnp.log(low_confidence + 1e-20)\n )\n soft_targets = common_utils.onehot(\n targets, vocab_size, on_value=confidence, off_value=low_confidence\n )\n\n loss = -jnp.sum(soft_targets * nn.log_softmax(logits), axis=-1)\n loss = loss - normalizing_constant\n\n normalizing_factor = np.prod(targets.shape)\n if weights is not None:\n loss = loss * weights\n normalizing_factor = weights.sum()\n\n return loss.sum(), normalizing_factor\n\n\ndef compute_weighted_accuracy(logits, targets, weights=None):\n if logits.ndim != targets.ndim + 1:\n raise ValueError(\n 'Incorrect shapes. Got shape %s logits and %s targets'\n % (str(logits.shape), str(targets.shape))\n )\n loss = jnp.equal(jnp.argmax(logits, axis=-1), targets)\n normalizing_factor = np.prod(logits.shape[:-1])\n if weights is not None:\n loss = loss * weights\n normalizing_factor = weights.sum()\n\n return loss.sum(), normalizing_factor\n\n\ndef compute_metrics(logits, labels, weights, label_smoothing=0.0):\n loss, weight_sum = compute_weighted_cross_entropy(\n logits, labels, weights, label_smoothing\n )\n acc, _ = compute_weighted_accuracy(logits, labels, weights)\n metrics = {\n 'loss': loss,\n 'accuracy': acc,\n 'denominator': weight_sum,\n }\n return metrics\n\n\ndef wmt_train_step(\n state,\n batch,\n config,\n learning_rate_fn,\n label_smoothing=0.0,\n dropout_rng=None,\n):\n train_keys = [\n 'inputs',\n 'targets',\n 'inputs_position',\n 'targets_position',\n 'inputs_segmentation',\n 'targets_segmentation',\n ]\n (\n inputs,\n targets,\n inputs_positions,\n targets_positions,\n inputs_segmentation,\n targets_segmentation,\n ) = (batch.get(k, None) for k in train_keys)\n\n weights = jnp.where(targets > 0, 1, 0).astype(jnp.float32)\n\n dropout_rng = jax.random.fold_in(dropout_rng, state.step)\n\n def loss_fn(params):\n logits = models.Transformer(config).apply(\n {'params': params},\n inputs,\n targets,\n inputs_positions=inputs_positions,\n targets_positions=targets_positions,\n inputs_segmentation=inputs_segmentation,\n targets_segmentation=targets_segmentation,\n rngs={'dropout': dropout_rng},\n )\n\n loss, weight_sum = compute_weighted_cross_entropy(\n logits, targets, weights, label_smoothing\n )\n mean_loss = loss / weight_sum\n return mean_loss, logits\n\n step = state.step\n\n if state.dynamic_scale:\n grad_fn = state.dynamic_scale.value_and_grad(loss_fn, has_aux=True)\n dynamic_scale, is_fin, (_, logits), grads = grad_fn(state.params)\n state = state.replace(dynamic_scale=dynamic_scale)\n else:\n grad_fn = jax.value_and_grad(loss_fn, has_aux=True)\n (_, logits), grads = grad_fn(state.params)\n\n new_state = state.apply_gradients(grads=grads)\n metrics = compute_metrics(logits, targets, weights)\n metrics['learning_rate'] = learning_rate_fn(step)\n\n if state.dynamic_scale:\n select_fn = functools.partial(jnp.where, is_fin)\n new_state = new_state.replace(\n opt_state=jax.tree_util.tree_map(\n select_fn, new_state.opt_state, state.opt_state\n ),\n params=jax.tree_util.tree_map(\n select_fn, new_state.params, state.params\n ),\n )\n metrics['loss_scale'] = dynamic_scale.scale * metrics['denominator']\n\n return new_state, metrics\n\n\ndef get_fake_batch(config: ml_collections.ConfigDict) -> dict[str, Any]:\n rng = jax.random.key(0)\n batch_size = config.per_device_batch_size\n max_len = config.max_target_length\n\n inputs = jax.random.randint(\n rng, (batch_size, max_len), 0, config.vocab_size, jnp.int32\n )\n targets = jax.random.randint(\n rng, (batch_size, max_len), 0, config.vocab_size, jnp.int32\n )\n\n return {\n 'inputs': inputs,\n 'targets': targets,\n 'inputs_position': None,\n 'targets_position': None,\n 'inputs_segmentation': None,\n 'targets_segmentation': None,\n }\n\n\ndef get_apply_fn_and_args(\n config: ml_collections.ConfigDict,\n) -> tuple[Any, tuple[Any, ...], dict[str, Any]]:\n dtype = preferred_dtype(config)\n\n train_config = models.TransformerConfig(\n vocab_size=config.vocab_size,\n output_vocab_size=config.vocab_size,\n share_embeddings=config.share_embeddings,\n logits_via_embedding=config.logits_via_embedding,\n dtype=dtype,\n emb_dim=config.emb_dim,\n num_heads=config.num_heads,\n num_layers=config.num_layers,\n qkv_dim=config.qkv_dim,\n mlp_dim=config.mlp_dim,\n max_len=max(config.max_target_length, config.max_eval_target_length),\n dropout_rate=config.dropout_rate,\n attention_dropout_rate=config.attention_dropout_rate,\n deterministic=False,\n decode=False,\n kernel_init=jax.nn.initializers.xavier_uniform(),\n bias_init=jax.nn.initializers.normal(stddev=1e-6),\n )\n\n model = models.Transformer(train_config)\n\n learning_rate_fn = create_learning_rate_schedule(\n learning_rate=config.learning_rate, warmup_steps=config.warmup_steps\n )\n\n optimizer = optax.adamw(\n learning_rate_fn,\n b1=0.9,\n b2=0.98,\n eps=1e-9,\n weight_decay=config.weight_decay,\n )\n\n rng = jax.random.key(0)\n init_rng, dropout_rng = jax.random.split(rng)\n\n batch = get_fake_batch(config)\n inputs = batch['inputs']\n targets = batch['targets']\n\n initial_variables = model.init(init_rng, inputs, targets)\n\n dynamic_scale = None\n platform = jax.local_devices()[0].platform\n if config.use_mixed_precision and platform == 'gpu':\n dynamic_scale = dynamic_scale_lib.DynamicScale()\n\n state = TrainState.create(\n apply_fn=model.apply,\n params=initial_variables['params'],\n tx=optimizer,\n dynamic_scale=dynamic_scale,\n )\n\n jit_train_step = jax.jit(\n wmt_train_step,\n static_argnums=(2, 3, 4),\n )\n\n return (\n jit_train_step,\n (state, batch, train_config, learning_rate_fn, 0.0, dropout_rng),\n {},\n )\n\n\n@google_benchmark.register\n@google_benchmark.option.unit(google_benchmark.kMillisecond)\ndef test_flax_wmt_trace(state):\n tracing_benchmark.benchmark_tracing(\n get_apply_fn_and_args, wmt_config.get_config, state\n )\n\n\n@google_benchmark.register\n@google_benchmark.option.unit(google_benchmark.kMillisecond)\ndef test_flax_wmt_lower(state):\n tracing_benchmark.benchmark_lowering(\n get_apply_fn_and_args, wmt_config.get_config, state\n )\n\n\nif __name__ == '__main__':\n tracing_benchmark.run_benchmarks()\n" }, { "path": "contributing.md", "content": "# How to Contribute\n\nPlease see https://flax.readthedocs.io/en/latest/contributing.html for more information.\n" }, { "path": "docs/.gitignore", "content": "_formatted_howtos\n" }, { "path": "docs/.readthedocs.yaml", "content": "# .readthedocs.yml\n# Read the Docs configuration file\n# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details\n\n# Required\nversion: 2\n\nbuild:\n os: ubuntu-22.04\n tools:\n python: \"3.12\"\n\n# Build documentation in the docs/ directory with Sphinx\nsphinx:\n configuration: docs/conf.py\n\n# Optionally build your docs in additional formats such as PDF and ePub\nformats:\n - htmlzip\n - epub\n # - pdf\n\n# Optionally set the version of Python and requirements required to build your docs\npython:\n install:\n - method: pip\n path: .\n extra_requirements:\n - all\n - testing\n - docs\n" }, { "path": "docs/Makefile", "content": "# Minimal makefile for Sphinx documentation\n#\n\n# You can set these variables from the command line, and also\n# from the environment for the first two.\nSPHINXOPTS ?=\nSPHINXBUILD ?= sphinx-build\nSOURCEDIR = .\nBUILDDIR = _build\n\n# Put it first so that \"make\" without argument is like \"make help\".\nhelp:\n\t@$(SPHINXBUILD) -M help \"$(SOURCEDIR)\" \"$(BUILDDIR)\" $(SPHINXOPTS) $(O)\n\n.PHONY: help Makefile\n\n# Catch-all target: route all unknown targets to Sphinx using the new\n# \"make mode\" option. $(O) is meant as a shortcut for $(SPHINXOPTS).\n%: Makefile\n\t@$(SPHINXBUILD) -M $@ \"$(SOURCEDIR)\" \"$(BUILDDIR)\" $(SPHINXOPTS) $(O)\n" }, { "path": "docs/README.md", "content": "# Deprecation\n\nThis folder contains the deprecated Flax Linen documentation. For the latest Flax NNX docs, check out the `docs_nnx` folder.\n\n# Where to find the docs\n\nThe FLAX Linen documentation can be found here: https://flax-linen.readthedocs.io/en/latest/\n\n# How to build the docs\n\n1. Clone the `flax` repository with `git clone https://github.com/google/flax.git`.\n1. In the main `flax` folder, install the required dependencies using `uv pip install -e .[docs]`.\n1. [Optional] If you need to make any local changes to the docs, create and switch to a branch. Make your changes to the docs in that branch.\n1. To build the docs, in the `flax/docs` folder run the make script: `make html`. Alternatively, install [`entr`](https://github.com/eradman/entr/), which helps run arbitrary commands when files change. Then run `find ../ ! -regex '.*/[\\.|\\_].*' | entr -s 'make html'`.\n1. If the build is successful, you should get the `The HTML pages are in _build/html.` message. You can preview the docs in `flax/docs/_build/html`.\n\n# How to run embedded code tests\n\nWe use `doctest` blocks for embedded code in documents, that are also\ntested. Learn more at https://www.sphinx-doc.org/en/master/usage/extensions/doctest.html\n\nTo run tests locally, run `make doctest`\n\n# How to write code documentation\n\nOur documentation is written in reStructuredText for Sphinx. It is a\nmeta-language that is compiled into online documentation. For more details, \ncheck out\n[Sphinx's documentation](https://www.sphinx-doc.org/en/master/usage/restructuredtext/index.html).\nAs a result, our docstrings adhere to a specific syntax that has to be kept in\nmind. Below we provide some guidelines.\n\nTo learn how to contribute to Jupyter Notebooks or other formats in Flax docs,\nrefer to the dedicated\n[Contributing](https://flax.readthedocs.io/en/latest/contributing.html) page.\n\n## How much information to put in a docstring\n\nDocstring should be informative. We prefer to err on the side of too much\ndocumentation than too little. For instance, providing a one-line explanation\nto a new `Module` which implements new functionality is not sufficient.\n\nFurthermore, we highly encourage adding examples to your docstrings, so users\ncan directly see how code can be used.\n\n## How to write inline tested code\n\nWe use [doctest](https://docs.python.org/3/library/doctest.html) syntax for\nwriting examples in documentation. These examples are ran as tests as part of\nour CI process. In order to write `doctest` code in your documentation, please\nuse the following notation:\n\n```bash\n# Example code::\n#\n# def sum(a, b):\n# return a + b\n#\n# sum(0, 1)\n```\n\nThe `Example code` string at the beginning can be replaced by anything as long\nas there are two semicolons and a newline following it, and the code is\nindented.\n\n## How to use \"code font\"\n\nWhen writing code font in a docstring, please use double backticks. Example:\n\n```bash\n# This returns a ``str`` object.\n```\n\nNote that argument names and objects like True, None or any strings should\nusually be put in `code`.\n\n## How to create cross-references/links\n\nIt is possible to create cross-references to other classes, functions, and\nmethods. In the following, `obj_typ` is either `class`, `func`, or `meth`.\n\n```bash\n# First method:\n# :`path_to_obj`\n\n# Second method:\n# ::`description `\n```\n\nYou can use the second method if the `path_to_obj` is very long. Some examples:\n\n```bash\n# Create: a reference to class flax.linen.Module.\n# :class:`flax.linen.Module`\n\n# Create a reference to local function my_func.\n# :func:`my_func`\n\n# Create a reference \"Module.apply()\" to method flax.linen.Module.apply.\n# :meth:`Module.apply() ` #\n```\n\nTo creata a hyperlink, use the following syntax:\n```bash\n# Note the double underscore at the end:\n# `Link to Google `__\n```\n\n### How to specify arguments for classes and methods\n\n* Class attributes should be specified using the `Attributes:` tag.\n* Method argument should be specified using the `Args:` tags.\n* All attributes and arguments should have types.\n\nHere is an example from our library:\n\n```python\nclass DenseGeneral(Module):\n \"\"\"A linear transformation with flexible axes.\n Attributes:\n features: int or tuple with number of output features.\n axis: int or tuple with axes to apply the transformation on. For instance,\n (-2, -1) will apply the transformation to the last two axes.\n batch_dims: tuple with batch axes.\n use_bias: whether to add a bias to the output (default: True).\n dtype: the dtype of the computation (default: float32).\n kernel_init: initializer function for the weight matrix.\n bias_init: initializer function for the bias.\n precision: numerical precision of the computation see `jax.lax.Precision`\n for details.\n \"\"\"\n features: Union[int, Iterable[int]]\n axis: Union[int, Iterable[int]] = -1\n batch_dims: Iterable[int] = ()\n use_bias: bool = True\n dtype: Dtype = jnp.float32\n kernel_init: Callable[[PRNGKey, Shape, Dtype], Array] = default_kernel_init\n bias_init: Callable[[PRNGKey, Shape, Dtype], Array] = zeros\n precision: Any = None\n\n @compact\n def __call__(self, inputs: Array) -> Array:\n \"\"\"Applies a linear transformation to the inputs along multiple dimensions.\n Args:\n inputs: The nd-array to be transformed.\n Returns:\n The transformed input.\n \"\"\"\n ...\n```" }, { "path": "docs/_ext/codediff.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Sphinx directive for creating code diff tables.\n\nUse directive as follows:\n\n.. codediff::\n :title: , \n\n \n ---\n \n\nIn order to highlight a line of code, append \"#!\" to it.\n\"\"\"\n\n\nimport sphinx\nfrom docutils import nodes\nfrom docutils.parsers.rst import directives\nfrom docutils.statemachine import ViewList\nfrom sphinx.util.docutils import SphinxDirective\n\nMISSING = object()\n\n\nclass CodeDiffParser:\n def parse(\n self,\n lines: list[str],\n title: str,\n groups: list[str] | None = None,\n skip_test: str | None = None,\n code_sep: str = '---',\n sync: object = MISSING,\n ):\n \"\"\"Parse the code diff block and format it so that it\n renders in different tabs and is tested by doctest.\n\n For example:\n\n .. testcode:: tab0, tab2, tab3\n\n \n\n .. codediff::\n :title: Tab 0, Tab 1, Tab 2, Tab 3\n :groups: tab0, tab1, tab2, tab3\n :skip_test: tab1, tab3\n\n \n\n ---\n\n \n\n ---\n\n \n\n ---\n\n \n\n For group tab0: and are executed.\n For group tab1: Nothing is executed.\n For group tab2: and are executed.\n For group tab3: is executed.\n\n Arguments:\n lines: a string list, where each element is a single string code line\n title: a single string that contains the titles of each tab (they should\n be separated by commas)\n groups: a single string that contains the group of each tab (they should\n be separated by commas). Code snippets that are part of the same group\n will be executed together. If groups=None, then the group names will\n default to the tab title names.\n skip_test: a single string denoting which group(s) to skip testing (they\n should be separated by commas). This is useful for legacy code snippets\n that no longer run correctly anymore. If skip_test=None, then no tests\n are skipped.\n code_sep: the separator character(s) used to denote a separate code block\n for a new tab. The default code separator is '---'.\n sync: an option for Sphinx directives, that will sync all tabs together.\n This means that if the user clicks to switch to another tab, all tabs\n will switch to the new tab.\n \"\"\"\n titles = [t.strip() for t in title.split(',')]\n num_tabs = len(titles)\n\n sync = sync is not MISSING\n # skip legacy code snippets in upgrade guides\n if skip_test is not None:\n skip_tests = {index.strip() for index in skip_test.split(',')}\n else:\n skip_tests = set()\n\n code_blocks = '\\n'.join(lines)\n if code_blocks.count(code_sep) != num_tabs - 1:\n raise ValueError(\n f'Expected {num_tabs-1} code separator(s) for {num_tabs} tab(s), but got {code_blocks.count(code_sep)} code separator(s) instead.'\n )\n code_blocks = [\n code_block.split('\\n')\n for code_block in code_blocks.split(code_sep + '\\n')\n ] # list[code_tab_list1[string_line1, ...], ...]\n\n # by default, put each code snippet in a different group denoted by an index number, to be executed separately\n if groups is not None:\n groups = [group_name.strip() for group_name in groups.split(',')]\n else:\n groups = titles\n if len(groups) != num_tabs:\n raise ValueError(\n f'Expected {num_tabs} group assignment(s) for {num_tabs} tab(s), but got {len(groups)} group assignment(s) instead.'\n )\n\n tabs = []\n test_codes = []\n for i, code_block in enumerate(code_blocks):\n if groups[i] not in skip_tests:\n test_codes.append((code_block, groups[i]))\n tabs.append((titles[i], self._code_block(code_block)))\n output = self._tabs(*tabs, sync=sync)\n\n return output, test_codes\n\n def _code_block(self, lines):\n \"\"\"Creates a codeblock.\"\"\"\n # Remove right trailing whitespace so we can detect the comments.\n lines = [x.rstrip() for x in lines]\n highlight = lambda x: x.endswith('#!')\n code = map(lambda x: x[:-2].rstrip() if highlight(x) else x, lines)\n highlights = [i + 1 for i in range(len(lines)) if highlight(lines[i])]\n highlights = ','.join(str(i) for i in highlights)\n\n directive = ['.. code-block:: python']\n if highlights:\n directive += [f' :emphasize-lines: {highlights}']\n\n # Indent code and add empty line so the code is picked up by the directive.\n return directive + [''] + list(map(lambda x: ' ' + x, code))\n\n def _tabs(self, *contents: tuple[str, list[str]], sync):\n output = ['.. tab-set::'] + [' ']\n\n for title, content in contents:\n output += [f' .. tab-item:: {title}']\n\n if sync:\n key = title.strip()\n output += [f' :sync: {key}']\n\n output += [' ']\n output += [' ' + line for line in content]\n\n return output\n\n\nclass CodeDiffDirective(SphinxDirective):\n has_content = True\n option_spec = {\n 'title': directives.unchanged,\n 'groups': directives.unchanged,\n 'skip_test': directives.unchanged,\n 'code_sep': directives.unchanged,\n 'sync': directives.flag,\n }\n\n def run(self):\n table_code, test_codes = CodeDiffParser().parse(\n list(self.content), **self.options\n )\n\n # Create a test node as a comment node so it won't show up in the docs.\n # We add attribute \"testnodetype\" so it is be picked up by the doctest\n # builder. This functionality is not officially documented but can be found\n # in the source code:\n # https://github.com/sphinx-doc/sphinx/blob/master/sphinx/ext/doctest.py\n # (search for 'testnodetype').\n test_nodes = []\n for test_code, group in test_codes:\n test_node = nodes.comment(\n '\\n'.join(test_code),\n '\\n'.join(test_code),\n testnodetype='testcode',\n groups=[group],\n )\n self.set_source_info(test_node)\n test_node['options'] = {}\n test_node['language'] = 'python3'\n test_nodes.append(test_node)\n\n # The table node is the side-by-side diff view that will be shown on RTD.\n table_node = nodes.paragraph()\n self.content = ViewList(table_code, self.content.parent)\n self.state.nested_parse(self.content, self.content_offset, table_node)\n\n return [table_node] + test_nodes\n\n\ndef setup(app):\n app.add_directive('codediff', CodeDiffDirective)\n\n return {\n 'version': sphinx.__display_version__,\n 'parallel_read_safe': True,\n 'parallel_write_safe': True,\n }\n" }, { "path": "docs/_ext/codediff_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for codediff Sphinx extension.\"\"\"\n\nfrom absl.testing import parameterized\nfrom codediff import CodeDiffParser\n\n\nclass CodeDiffTest(parameterized.TestCase):\n def test_parse(self):\n input_text = r\"\"\"@jax.jit #!\ndef get_initial_params(key): #!\n init_val = jnp.ones((1, 28, 28, 1), jnp.float32)\n initial_params = CNN().init(key, init_val)['params']\n extra_line\n return initial_params\n---\n@jax.pmap #!\ndef get_initial_params(key):\n init_val = jnp.ones((1, 28, 28, 1), jnp.float32)\n initial_params = CNN().init(key, init_val)['params']\n return initial_params\"\"\"\n\n expected_table = \"\"\".. tab-set::\\n \\n .. tab-item:: Single device\\n \\n .. code-block:: python\\n :emphasize-lines: 1,2\\n \\n @jax.jit\\n def get_initial_params(key):\\n init_val = jnp.ones((1, 28, 28, 1), jnp.float32)\\n initial_params = CNN().init(key, init_val)['params']\\n extra_line\\n return initial_params\\n \\n .. tab-item:: Ensembling on multiple devices\\n \\n .. code-block:: python\\n :emphasize-lines: 1\\n \\n @jax.pmap\\n def get_initial_params(key):\\n init_val = jnp.ones((1, 28, 28, 1), jnp.float32)\\n initial_params = CNN().init(key, init_val)['params']\\n return initial_params\"\"\"\n\n expected_testcodes = [\n r\"\"\"@jax.jit #!\ndef get_initial_params(key): #!\n init_val = jnp.ones((1, 28, 28, 1), jnp.float32)\n initial_params = CNN().init(key, init_val)['params']\n extra_line\n return initial_params\n\"\"\",\n r\"\"\"@jax.pmap #!\ndef get_initial_params(key):\n init_val = jnp.ones((1, 28, 28, 1), jnp.float32)\n initial_params = CNN().init(key, init_val)['params']\n return initial_params\"\"\",\n ]\n\n title_left = 'Single device'\n title_right = 'Ensembling on multiple devices'\n\n actual_table, actual_testcodes = CodeDiffParser().parse(\n lines=input_text.split('\\n'),\n title=f'{title_left}, {title_right}',\n )\n\n actual_table = '\\n'.join(actual_table)\n actual_testcodes = ['\\n'.join(testcode) for testcode, _ in actual_testcodes]\n\n self.assertEqual(expected_table, actual_table)\n self.assertEqual(expected_testcodes[0], actual_testcodes[0])\n self.assertEqual(expected_testcodes[1], actual_testcodes[1])\n\n @parameterized.parameters(\n {\n 'input_text': r\"\"\"x = 1\n ---\n x = 2\n\"\"\",\n 'title': 'Tab 0, Tab1, Tab2',\n 'groups': None,\n 'error_msg': 'Expected 2 code separator\\\\(s\\\\) for 3 tab\\\\(s\\\\), but got 1 code separator\\\\(s\\\\) instead.',\n },\n {\n 'input_text': r\"\"\"x = 1\n ---\n x = 2\n ---\n x = 3\n ---\n x = 4\n\"\"\",\n 'title': 'Tab 0, Tab1, Tab2',\n 'groups': None,\n 'error_msg': 'Expected 2 code separator\\\\(s\\\\) for 3 tab\\\\(s\\\\), but got 3 code separator\\\\(s\\\\) instead.',\n },\n {\n 'input_text': r\"\"\"x = 1\n ---\n x = 2\n ---\n x = 3\n\"\"\",\n 'title': 'Tab 0, Tab1, Tab2',\n 'groups': 'tab0, tab2',\n 'error_msg': 'Expected 3 group assignment\\\\(s\\\\) for 3 tab\\\\(s\\\\), but got 2 group assignment\\\\(s\\\\) instead.',\n },\n {\n 'input_text': r\"\"\"x = 1\n ---\n x = 2\n ---\n x = 3\n\"\"\",\n 'title': 'Tab 0, Tab1, Tab2',\n 'groups': 'tab0, tab1, tab2, tab3',\n 'error_msg': 'Expected 3 group assignment\\\\(s\\\\) for 3 tab\\\\(s\\\\), but got 4 group assignment\\\\(s\\\\) instead.',\n },\n )\n def test_parse_errors(self, input_text, title, groups, error_msg):\n with self.assertRaisesRegex(ValueError, error_msg):\n _, _ = CodeDiffParser().parse(\n lines=input_text.split('\\n'),\n title=title,\n groups=groups,\n )\n" }, { "path": "docs/_ext/flax_module.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Sphinx directive for visualizing Flax modules.\n\nUse directive as follows:\n\n.. flax_module::\n :module: flax.linen\n :class: Dense\n\"\"\"\n\nimport importlib\n\nimport sphinx\nimport sphinx.ext.autosummary.generate as ag\nfrom docutils import nodes\nfrom docutils.parsers.rst import directives\nfrom docutils.statemachine import ViewList\nfrom sphinx.util.docutils import SphinxDirective\n\nfrom docs.conf_sphinx_patch import generate_autosummary_content\n\n\ndef render_module(modname: str, qualname: str, app):\n parent = importlib.import_module(modname)\n obj = getattr(parent, qualname)\n template = ag.AutosummaryRenderer(app)\n template_name = 'flax_module'\n imported_members = False\n recursive = False\n context = {}\n return generate_autosummary_content(\n qualname,\n obj,\n parent,\n template,\n template_name,\n imported_members,\n app,\n recursive,\n context,\n modname,\n qualname,\n )\n\n\nclass FlaxModuleDirective(SphinxDirective):\n has_content = True\n option_spec = {\n 'module': directives.unchanged,\n 'class': directives.unchanged,\n }\n\n def run(self):\n module_template = render_module(\n self.options['module'], self.options['class'], self.env.app\n )\n module_template = module_template.splitlines()\n\n # Create a container for the rendered nodes\n container_node = nodes.container()\n self.content = ViewList(module_template, self.content.parent)\n self.state.nested_parse(self.content, self.content_offset, container_node)\n\n return [container_node]\n\n\ndef setup(app):\n app.add_directive('flax_module', FlaxModuleDirective)\n\n return {\n 'version': sphinx.__display_version__,\n 'parallel_read_safe': True,\n 'parallel_write_safe': True,\n }\n" }, { "path": "docs/_static/css/flax_theme.css", "content": "@import url(\"theme.css\");\n\n.wy-nav-content {\n max-width: 1290px;\n}\n\n.rst-content table.docutils {\n width: 100%;\n}\n\n.rst-content table.docutils td {\n vertical-align: top;\n padding: 0;\n}\n\n.rst-content table.docutils td p {\n padding: 8px;\n}\n\n.rst-content div[class^=highlight] {\n border: 0;\n margin: 0;\n}\n" }, { "path": "docs/_templates/autosummary/flax_module.rst", "content": "{{ fullname | escape | underline }}\n\n.. currentmodule:: {{ module }}\n\n.. autoclass:: {{ objname }}\n :exclude-members:\n\n .. automethod:: __call__\n\n {% block methods %}\n\n {% for item in methods %}\n {%- if item not in inherited_members and item not in annotations and not item in ['__init__', 'setup'] %}\n .. automethod:: {{ item }}\n {%- endif %}\n {%- endfor %}\n\n {% if methods %}\n .. rubric:: Methods\n\n .. autosummary::\n\n {% for item in methods %}\n {%- if item not in inherited_members and item not in annotations and not item in ['__init__', 'setup'] %}\n ~{{ name }}.{{ item }}\n {%- endif %}\n {%- endfor %}\n {% endif %}\n {% endblock %}" }, { "path": "docs/api_reference/flax.core.frozen_dict.rst", "content": "\nflax.core.frozen_dict package\n=============================\n\n.. currentmodule:: flax.core.frozen_dict\n\n.. autoclass:: FrozenDict\n :members: pretty_repr, copy, pop, unfreeze, tree_flatten\n\n.. autofunction:: freeze\n\n.. autofunction:: unfreeze\n\n.. autofunction:: copy\n\n.. autofunction:: pop\n\n.. autofunction:: pretty_repr\n" }, { "path": "docs/api_reference/flax.cursor.rst", "content": "\nflax.cursor package\n=============================\n\nThe Cursor API allows for mutability of pytrees. This API provides a more\nergonomic solution to making partial-updates of deeply nested immutable\ndata structures, compared to making many nested ``dataclasses.replace`` calls.\n\nTo illustrate, consider the example below::\n\n >>> from flax.cursor import cursor\n >>> import dataclasses\n >>> from typing import Any\n\n >>> @dataclasses.dataclass(frozen=True)\n >>> class A:\n ... x: Any\n\n >>> a = A(A(A(A(A(A(A(0)))))))\n\nTo replace the int ``0`` using ``dataclasses.replace``, we would have to write many nested calls::\n\n >>> a2 = dataclasses.replace(\n ... a,\n ... x=dataclasses.replace(\n ... a.x,\n ... x=dataclasses.replace(\n ... a.x.x,\n ... x=dataclasses.replace(\n ... a.x.x.x,\n ... x=dataclasses.replace(\n ... a.x.x.x.x,\n ... x=dataclasses.replace(\n ... a.x.x.x.x.x,\n ... x=dataclasses.replace(a.x.x.x.x.x.x, x=1),\n ... ),\n ... ),\n ... ),\n ... ),\n ... ),\n ... )\n\nThe equivalent can be achieved much more simply using the Cursor API::\n\n >>> a3 = cursor(a).x.x.x.x.x.x.x.set(1)\n >>> assert a2 == a3\n\nThe Cursor object keeps tracks of changes made to it and when ``.build`` is called,\ngenerates a new object with the accumulated changes. Basic usage involves\nwrapping the object in a Cursor, making changes to the Cursor object and\ngenerating a new copy of the original object with the accumulated changes.\n\n.. currentmodule:: flax.cursor\n\n.. autofunction:: cursor\n\n.. autoclass:: Cursor\n :members: apply_update, build, find, find_all, set\n\n\n" }, { "path": "docs/api_reference/flax.errors.rst", "content": "\nflax.errors package\n===================\n\nFlax has the following classes of errors.\n\n.. automodule:: flax.errors\n :members:\n :exclude-members: FlaxError" }, { "path": "docs/api_reference/flax.jax_utils.rst", "content": "\nflax.jax_utils package\n========================\n\n.. currentmodule:: flax.jax_utils\n\n.. automodule:: flax.jax_utils\n\n\n.. autofunction:: partial_eval_by_shape\n\n\nMulti device utilities\n------------------------\n\n.. autofunction:: replicate\n.. autofunction:: unreplicate\n\n.. autofunction:: prefetch_to_device\n\n.. autofunction:: pmean\n\n.. autofunction:: pad_shard_unpad\n" }, { "path": "docs/api_reference/flax.linen/activation_functions.rst", "content": "\nActivation functions\n------------------------\n\n.. automodule:: flax.linen.activation\n.. currentmodule:: flax.linen.activation\n\n.. autoclass:: PReLU\n :members:\n :special-members: __call__\n\n.. autofunction:: celu\n.. autofunction:: elu\n.. autofunction:: gelu\n.. autofunction:: glu\n.. autofunction:: hard_sigmoid\n.. autofunction:: hard_silu\n.. autofunction:: hard_swish\n.. autofunction:: hard_tanh\n.. autofunction:: leaky_relu\n.. autofunction:: log_sigmoid\n.. autofunction:: log_softmax\n.. autofunction:: logsumexp\n.. autofunction:: one_hot\n.. autofunction:: relu\n.. autofunction:: relu6 as relu6,\n.. autofunction:: selu\n.. autofunction:: sigmoid\n.. autofunction:: silu\n.. autofunction:: soft_sign\n.. autofunction:: softmax\n.. autofunction:: softplus\n.. autofunction:: standardize\n.. autofunction:: swish\n.. autofunction:: tanh\n" }, { "path": "docs/api_reference/flax.linen/decorators.rst", "content": "Decorators\n----------------------\n\n.. currentmodule:: flax.linen\n\n.. autofunction:: compact\n.. autofunction:: nowrap\n" }, { "path": "docs/api_reference/flax.linen/index.rst", "content": "\nflax.linen\n==========\n\nLinen is the Flax Module system. Read more about our design goals in the `Linen README `_.\n\n.. toctree::\n :maxdepth: 2\n\n module\n init_apply\n layers\n activation_functions\n initializers\n transformations\n inspection\n variable\n spmd\n decorators\n profiling" }, { "path": "docs/api_reference/flax.linen/init_apply.rst", "content": "\nInit/Apply\n==============\n\n.. currentmodule:: flax.linen\n\n.. autofunction:: apply\n.. autofunction:: init\n.. autofunction:: init_with_output\n" }, { "path": "docs/api_reference/flax.linen/initializers.rst", "content": "Initializers\n------------------------\n\n.. automodule:: flax.linen.initializers\n.. currentmodule:: flax.linen.initializers\n\n.. autofunction:: constant\n.. autofunction:: delta_orthogonal\n.. autofunction:: glorot_normal\n.. autofunction:: glorot_uniform\n.. autofunction:: he_normal\n.. autofunction:: he_uniform\n.. autofunction:: kaiming_normal\n.. autofunction:: kaiming_uniform\n.. autofunction:: lecun_normal\n.. autofunction:: lecun_uniform\n.. autofunction:: normal\n.. autofunction:: truncated_normal\n.. autofunction:: ones\n.. autofunction:: ones_init\n.. autofunction:: orthogonal\n.. autofunction:: uniform\n.. autofunction:: variance_scaling\n.. autofunction:: xavier_normal\n.. autofunction:: xavier_uniform\n.. autofunction:: zeros\n.. autofunction:: zeros_init\n" }, { "path": "docs/api_reference/flax.linen/inspection.rst", "content": "\nInspection\n----------------------\n\n.. currentmodule:: flax.linen\n\n.. autofunction:: tabulate\n" }, { "path": "docs/api_reference/flax.linen/layers.rst", "content": "Layers\n======\n\n.. currentmodule:: flax.linen\n\nLinear Modules\n------------------------\n\n.. flax_module::\n :module: flax.linen\n :class: Dense\n\n.. flax_module::\n :module: flax.linen\n :class: DenseGeneral\n\n.. flax_module::\n :module: flax.linen\n :class: Conv\n\n.. flax_module::\n :module: flax.linen\n :class: ConvTranspose\n\n.. flax_module::\n :module: flax.linen\n :class: ConvLocal\n\n.. flax_module::\n :module: flax.linen\n :class: Einsum\n\n.. flax_module::\n :module: flax.linen\n :class: Embed\n\nPooling\n------------------------\n\n.. autofunction:: max_pool\n.. autofunction:: avg_pool\n.. autofunction:: pool\n\nNormalization\n------------------------\n\n.. flax_module::\n :module: flax.linen\n :class: BatchNorm\n\n.. flax_module::\n :module: flax.linen\n :class: LayerNorm\n\n.. flax_module::\n :module: flax.linen\n :class: GroupNorm\n\n.. flax_module::\n :module: flax.linen\n :class: RMSNorm\n\n.. flax_module::\n :module: flax.linen\n :class: InstanceNorm\n\n.. flax_module::\n :module: flax.linen\n :class: SpectralNorm\n\n.. flax_module::\n :module: flax.linen\n :class: WeightNorm\n\n\nCombinators\n------------------------\n\n.. flax_module::\n :module: flax.linen\n :class: Sequential\n\nStochastic\n------------------------\n\n.. flax_module::\n :module: flax.linen\n :class: Dropout\n\nAttention\n------------------------\n\n.. flax_module::\n :module: flax.linen\n :class: MultiHeadDotProductAttention\n\n.. flax_module::\n :module: flax.linen\n :class: MultiHeadAttention\n\n.. flax_module::\n :module: flax.linen\n :class: SelfAttention\n\n.. autofunction:: dot_product_attention_weights\n.. autofunction:: dot_product_attention\n.. autofunction:: make_attention_mask\n.. autofunction:: make_causal_mask\n\nRecurrent\n------------------------\n\n.. flax_module::\n :module: flax.linen\n :class: RNNCellBase\n\n.. flax_module::\n :module: flax.linen\n :class: LSTMCell\n\n.. flax_module::\n :module: flax.linen\n :class: OptimizedLSTMCell\n\n.. flax_module::\n :module: flax.linen\n :class: ConvLSTMCell\n\n.. flax_module::\n :module: flax.linen\n :class: SimpleCell\n\n.. flax_module::\n :module: flax.linen\n :class: GRUCell\n\n.. flax_module::\n :module: flax.linen\n :class: MGUCell\n\n.. flax_module::\n :module: flax.linen\n :class: RNN\n\n.. flax_module::\n :module: flax.linen\n :class: Bidirectional\n\nBatchApply\n------------------------\n\n.. flax_module::\n :module: flax.linen\n :class: BatchApply\n" }, { "path": "docs/api_reference/flax.linen/module.rst", "content": "Module\n------------------------\n\n.. automodule:: flax.linen\n.. currentmodule:: flax.linen\n\n.. autoclass:: Module\n :members: setup, variable, param, bind, unbind, apply, init, init_with_output, copy, make_rng, sow, variables, Variable, __setattr__, tabulate, module_paths, is_initializing, perturb, put_variable, has_variable, has_rng, lazy_init, get_variable, path, is_mutable_collection\n\n.. autofunction:: apply\n.. autofunction:: init\n.. autofunction:: init_with_output\n.. autofunction:: intercept_methods\n.. autofunction:: share_scope\n" }, { "path": "docs/api_reference/flax.linen/profiling.rst", "content": "Profiling\n----------------------\n\n.. currentmodule:: flax.linen\n\n.. autofunction:: enable_named_call\n.. autofunction:: disable_named_call\n.. autofunction:: override_named_call\n" }, { "path": "docs/api_reference/flax.linen/spmd.rst", "content": "\nSPMD\n----------------------\n\n.. automodule:: flax.linen.spmd\n.. currentmodule:: flax.linen\n\n.. autofunction:: Partitioned\n.. autofunction:: with_partitioning\n.. autofunction:: get_partition_spec\n.. autofunction:: get_sharding\n.. autofunction:: LogicallyPartitioned\n.. autofunction:: logical_axis_rules\n.. autofunction:: set_logical_axis_rules\n.. autofunction:: get_logical_axis_rules\n.. autofunction:: logical_to_mesh_axes\n.. autofunction:: logical_to_mesh\n.. autofunction:: logical_to_mesh_sharding\n.. autofunction:: with_logical_constraint\n.. autofunction:: with_logical_partitioning\n" }, { "path": "docs/api_reference/flax.linen/transformations.rst", "content": "Transformations\n----------------------\n\n.. automodule:: flax.linen.transforms\n.. currentmodule:: flax.linen\n\n.. autofunction:: vmap\n.. autofunction:: scan\n.. autofunction:: jit\n.. autofunction:: remat\n.. autofunction:: remat_scan\n.. autofunction:: map_variables\n.. autofunction:: jvp\n.. autofunction:: vjp\n.. autofunction:: custom_vjp\n.. autofunction:: while_loop\n.. autofunction:: cond\n.. autofunction:: switch\n" }, { "path": "docs/api_reference/flax.linen/variable.rst", "content": "\nVariable dictionary\n----------------------\n\n.. automodule:: flax.core.variables\n.. autoclass:: flax.linen.Variable\n" }, { "path": "docs/api_reference/flax.serialization.rst", "content": "\nflax.serialization package\n============================\n\n.. currentmodule:: flax.serialization\n\n.. automodule:: flax.serialization\n\n\nState dicts\n------------------------\n\n.. autofunction:: from_state_dict\n.. autofunction:: to_state_dict\n\n.. autofunction:: register_serialization_state\n\n\nSerialization with MessagePack\n--------------------------------\n\n.. autofunction:: msgpack_serialize\n.. autofunction:: msgpack_restore\n\n.. autofunction:: to_bytes\n.. autofunction:: from_bytes\n" }, { "path": "docs/api_reference/flax.struct.rst", "content": "\nflax.struct package\n=====================\n\n.. currentmodule:: flax.struct\n\n.. automodule:: flax.struct\n\n\n.. autofunction:: dataclass\n\n\n.. autoclass:: PyTreeNode" }, { "path": "docs/api_reference/flax.traceback_util.rst", "content": "flax.traceback_util package\n============================\n\n.. currentmodule:: flax.traceback_util\n\n.. automodule:: flax.traceback_util\n\n\nTraceback filtering utils\n--------------------------\n\n.. autofunction:: hide_flax_in_tracebacks\n\n.. autofunction:: show_flax_in_tracebacks\n" }, { "path": "docs/api_reference/flax.training.rst", "content": "\nflax.training package\n=====================\n\nCheckpoints\n------------------------\n\n.. currentmodule:: flax.training.checkpoints\n\n.. automodule:: flax.training.checkpoints\n\n.. autofunction:: save_checkpoint\n\n.. autofunction:: save_checkpoint_multiprocess\n\n.. autofunction:: latest_checkpoint\n\n.. autofunction:: restore_checkpoint\n\n.. autofunction:: convert_pre_linen\n\nLearning rate schedules\n------------------------\n\n.. currentmodule:: flax.training.lr_schedule\n\n.. automodule:: flax.training.lr_schedule\n\n.. autofunction:: create_constant_learning_rate_schedule\n\n.. autofunction:: create_stepped_learning_rate_schedule\n\n.. autofunction:: create_cosine_learning_rate_schedule\n\nTrain state\n------------------------\n\n.. currentmodule:: flax.training.train_state\n\n.. autoclass:: TrainState\n :members: apply_gradients, create\n\nEarly Stopping\n------------------------\n\n.. currentmodule:: flax.training.early_stopping\n\n.. autoclass:: EarlyStopping\n :members: reset, update\n\nCommon Utilities\n------------------------\n\n.. currentmodule:: flax.training.common_utils\n\n.. autofunction:: shard\n\n.. autofunction:: shard_prng_key\n\n.. autofunction:: stack_forest\n\n.. autofunction:: get_metrics\n\n.. autofunction:: onehot\n" }, { "path": "docs/api_reference/index.rst", "content": "API Reference\n=============\n\n.. toctree::\n :maxdepth: 4\n\n flax.config\n flax.core.frozen_dict\n flax.cursor\n flax.errors\n flax.jax_utils\n flax.linen/index\n flax.serialization\n flax.struct\n flax.traceback_util\n flax.training\n flax.traverse_util" }, { "path": "docs/conf.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Configuration file for the Sphinx documentation builder.\"\"\"\n\n\n# This file only contains a selection of the most common options. For a full\n# list see the documentation:\n# https://www.sphinx-doc.org/en/master/usage/configuration.html\n\n# -- Path setup --------------------------------------------------------------\n\n# If extensions (or modules to document with autodoc) are in another directory,\n# add these directories to sys.path here. If the directory is relative to the\n# documentation root, use os.path.abspath to make it absolute, like shown here.\n#\n# import os\n# import sys\n# sys.path.insert(0, os.path.abspath('.'))\n\nimport os\nimport sys\nimport doctest\n\nsys.path.insert(0, os.path.abspath('..'))\n# Include local extension.\nsys.path.append(os.path.abspath('./_ext'))\n\n# patch sphinx\n# -- Project information -----------------------------------------------------\n\nproject = 'Flax'\ncopyright = '2023, The Flax authors' # pylint: disable=redefined-builtin\nauthor = 'The Flax authors'\n\n\n# -- General configuration ---------------------------------------------------\n\n# Add any Sphinx extension module names here, as strings. They can be\n# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom\n# ones.\nextensions = [\n 'sphinx.ext.autodoc',\n 'sphinx.ext.autosummary',\n 'sphinx.ext.autosectionlabel',\n 'sphinx.ext.doctest',\n 'sphinx.ext.intersphinx',\n 'sphinx.ext.mathjax',\n 'sphinx.ext.napoleon',\n 'sphinx.ext.viewcode',\n 'myst_nb',\n 'codediff',\n 'flax_module',\n 'sphinx_design',\n]\n\n# Add any paths that contain templates here, relative to this directory.\ntemplates_path = ['_templates']\n\n# List of patterns, relative to source directory, that match files and\n# directories to ignore when looking for source files.\n# This pattern also affects html_static_path and html_extra_path.\nexclude_patterns = ['_build', 'Thumbs.db', '.DS_Store']\n\n# The suffix(es) of source filenames.\n# Note: important to list ipynb before md here: we have both md and ipynb\n# copies of each notebook, and myst will choose which to convert based on\n# the order in the source_suffix list. Notebooks which are not executed have\n# outputs stored in ipynb but not in md, so we must convert the ipynb.\nsource_suffix = ['.rst', '.ipynb', '.md']\n\nautosummary_generate = True\n\nmaster_doc = 'index'\n\nautodoc_typehints = 'none'\n\n\n# -- Options for HTML output -------------------------------------------------\n\n# The theme to use for HTML and HTML Help pages. See the documentation for\n# a list of builtin themes.\n#\n# html_theme = 'pydata_sphinx_theme'\nhtml_theme = 'sphinx_book_theme'\nhtml_css_files = ['css/flax_theme.css']\n\n# The name of an image file (relative to this directory) to place at the top\n# of the sidebar.\nhtml_logo = './flax.png'\nhtml_favicon = './flax.png'\n\n# title of the website\nhtml_title = ''\n\n# Add any paths that contain custom static files (such as style sheets) here,\n# relative to this directory. They are copied after the builtin static files,\n# so a file named 'default.css' will overwrite the builtin 'default.css'.\nhtml_static_path = ['_static']\n\nhtml_extra_path = ['robots.txt']\n\n# href with no underline and white bold text color\nannouncement = \"\"\"\n\n This site covers the old Flax Linen API. [Explore the new Flax NNX API ✨]\n\n\"\"\"\n\nhtml_theme_options = {\n 'repository_url': 'https://github.com/google/flax',\n 'use_repository_button': True, # add a 'link to repository' button\n 'use_issues_button': False, # add an 'Open an Issue' button\n 'path_to_docs': (\n 'docs'\n ), # used to compute the path to launch notebooks in colab\n 'launch_buttons': {\n 'colab_url': 'https://colab.research.google.com/',\n },\n 'prev_next_buttons_location': None,\n 'show_navbar_depth': 1,\n 'announcement': announcement,\n}\n\n# -- Options for myst ----------------------------------------------\n# uncomment line below to avoid running notebooks during development\n# nb_execution_mode = 'off'\n# Notebook cell execution timeout; defaults to 30.\nnb_execution_timeout = 100\n# List of patterns, relative to source directory, that match notebook\n# files that will not be executed.\nmyst_enable_extensions = ['dollarmath']\nnb_execution_excludepatterns = [\n 'quick_start.ipynb', # <-- times out\n 'transfer_learning.ipynb', # <-- transformers requires flax<=0.7.0\n 'flax/nnx', # exclude nnx\n 'guides/quantization/fp8_basics.ipynb',\n 'guides/training_techniques/use_checkpointing.ipynb', # TODO(IvyZX): needs to be updated\n]\n# raise exceptions on execution so CI can catch errors\nnb_execution_allow_errors = False\nnb_execution_raise_on_error = True\n\n# -- Extension configuration -------------------------------------------------\n\n# Tell sphinx-autodoc-typehints to generate stub parameter annotations including\n# types, even if the parameters aren't explicitly documented.\nalways_document_param_types = True\n\n# -- doctest configuration -------------------------------------------------\ndoctest_default_flags = doctest.NORMALIZE_WHITESPACE\ndoctest_global_setup = \"\"\"\nimport jax\nimport jax.numpy as jnp\nfrom flax import nnx\n\nimport logging as slog\nfrom absl import logging as alog\n\n# Avoid certain absl logging messages to break doctest\nfiltered_message = [\n 'SaveArgs.aggregate is deprecated',\n '',\n]\n\nclass _CustomLogFilter(slog.Formatter):\n def format(self, record):\n message = super(_CustomLogFilter, self).format(record)\n for m in filtered_message:\n if m in message:\n return ''\n return message\n\nalog.use_absl_handler()\nalog.get_absl_handler().setFormatter(_CustomLogFilter())\n\"\"\"\n" }, { "path": "docs/conf_sphinx_patch.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Patch Sphinx to improve documentation aesthetics.\"\"\"\n\n# TODO(cgarciae): Send a PR to sphinx to upstream this fix. Issue: https://github.com/google/flax/issues/2196\n# This patch is needed to make autosummary provide the \"annotations\"\n# variable so we can exclude function attributes from the methods list\n# in flax_module.rst. The patch as such only adds this single line:\n#\n# ns['annotations'] = list(getattr(obj, '__annotations__', {}).keys())'\n#\n# We should consider sending a PR to sphinx so we can get rid of this.\n# Original source: https://github.com/sphinx-doc/sphinx/blob/0aedcc9a916daa92d477226da67d33ce1831822e/sphinx/ext/autosummary/generate.py#L211-L351\nfrom typing import Any\n\nimport sphinx.ext.autodoc\nimport sphinx.ext.autosummary.generate as ag\n\n\ndef generate_autosummary_content(\n name: str,\n obj: Any,\n parent: Any,\n template: ag.AutosummaryRenderer,\n template_name: str,\n imported_members: bool,\n app: Any,\n recursive: bool,\n context: dict,\n modname: str = None,\n qualname: str = None,\n) -> str:\n doc = ag.get_documenter(app, obj, parent)\n\n def skip_member(obj: Any, name: str, objtype: str) -> bool:\n try:\n return app.emit_firstresult(\n 'autodoc-skip-member', objtype, name, obj, False, {}\n )\n except Exception as exc:\n ag.logger.warning(\n __(\n 'autosummary: failed to determine %r to be documented, '\n 'the following exception was raised:\\n%s'\n ),\n name,\n exc,\n type='autosummary',\n )\n return False\n\n def get_class_members(obj: Any) -> dict[str, Any]:\n members = sphinx.ext.autodoc.get_class_members(\n obj, [qualname], ag.safe_getattr\n )\n return {name: member.object for name, member in members.items()}\n\n def get_module_members(obj: Any) -> dict[str, Any]:\n members = {}\n for name in ag.members_of(obj, app.config):\n try:\n members[name] = ag.safe_getattr(obj, name)\n except AttributeError:\n continue\n return members\n\n def get_all_members(obj: Any) -> dict[str, Any]:\n if doc.objtype == 'module':\n return get_module_members(obj)\n elif doc.objtype == 'class':\n return get_class_members(obj)\n return {}\n\n def get_members(\n obj: Any,\n types: set[str],\n include_public: list[str] = [],\n imported: bool = True,\n ) -> tuple[list[str], list[str]]:\n items: list[str] = []\n public: list[str] = []\n\n all_members = get_all_members(obj)\n for name, value in all_members.items():\n documenter = ag.get_documenter(app, value, obj)\n if documenter.objtype in types:\n # skip imported members if expected\n if imported or getattr(value, '__module__', None) == obj.__name__:\n skipped = skip_member(value, name, documenter.objtype)\n if skipped is True:\n pass\n elif skipped is False:\n # show the member forcedly\n items.append(name)\n public.append(name)\n else:\n items.append(name)\n if name in include_public or not name.startswith('_'):\n # considers member as public\n public.append(name)\n return public, items\n\n def get_module_attrs(members: Any) -> tuple[list[str], list[str]]:\n \"\"\"Find module attributes with docstrings.\"\"\"\n attrs, public = [], []\n try:\n analyzer = ag.ModuleAnalyzer.for_module(name)\n attr_docs = analyzer.find_attr_docs()\n for namespace, attr_name in attr_docs:\n if namespace == '' and attr_name in members:\n attrs.append(attr_name)\n if not attr_name.startswith('_'):\n public.append(attr_name)\n except ag.PycodeError:\n pass # give up if ModuleAnalyzer fails to parse code\n return public, attrs\n\n def get_modules(obj: Any) -> tuple[list[str], list[str]]:\n items: list[str] = []\n for _, modname, _ispkg in ag.pkgutil.iter_modules(obj.__path__):\n fullname = name + '.' + modname\n try:\n module = ag.import_module(fullname)\n if module and hasattr(module, '__sphinx_mock__'):\n continue\n except ImportError:\n pass\n\n items.append(fullname)\n public = [x for x in items if not x.split('.')[-1].startswith('_')]\n return public, items\n\n ns: dict[str, Any] = {}\n ns.update(context)\n\n if doc.objtype == 'module':\n scanner = ag.ModuleScanner(app, obj)\n ns['members'] = scanner.scan(imported_members)\n ns['functions'], ns['all_functions'] = get_members(\n obj, {'function'}, imported=imported_members\n )\n ns['classes'], ns['all_classes'] = get_members(\n obj, {'class'}, imported=imported_members\n )\n ns['exceptions'], ns['all_exceptions'] = get_members(\n obj, {'exception'}, imported=imported_members\n )\n ns['attributes'], ns['all_attributes'] = get_module_attrs(ns['members'])\n ispackage = hasattr(obj, '__path__')\n if ispackage and recursive:\n ns['modules'], ns['all_modules'] = get_modules(obj)\n elif doc.objtype == 'class':\n ns['members'] = dir(obj)\n ns['inherited_members'] = set(dir(obj)) - set(obj.__dict__.keys())\n ns['methods'], ns['all_methods'] = get_members(\n obj, {'method'}, ['__init__']\n )\n ns['attributes'], ns['all_attributes'] = get_members(\n obj, {'attribute', 'property'}\n )\n ns['annotations'] = list(getattr(obj, '__annotations__', {}).keys())\n\n if modname is None or qualname is None:\n modname, qualname = ag.split_full_qualified_name(name)\n\n if doc.objtype in ('method', 'attribute', 'property'):\n ns['class'] = qualname.rsplit('.', 1)[0]\n\n if doc.objtype in ('class',):\n shortname = qualname\n else:\n shortname = qualname.rsplit('.', 1)[-1]\n\n ns['fullname'] = name\n ns['module'] = modname\n ns['objname'] = qualname\n ns['name'] = shortname\n\n ns['objtype'] = doc.objtype\n ns['underline'] = len(name) * '='\n\n if template_name:\n return template.render(template_name, ns)\n else:\n return template.render(doc.objtype, ns)\n\n\nag.generate_autosummary_content = generate_autosummary_content\n" }, { "path": "docs/developer_notes/index.rst", "content": "Developer notes\n===============\n\n.. toctree::\n :maxdepth: 1\n\n module_lifecycle\n lift\n FLIPs \n" }, { "path": "docs/developer_notes/lift.md", "content": "# Lifted transformations\n\n⚠️ Advanced topic ⚠️\n\nThis design note explains the underlying implementation of `flax.linen.transform`, which enables JAX transformations inside Flax `Module`s.\n\n\n## Introduction\n\nJAX uses a functional API meaning that it only guarantees correct behavior when using functions without side effects ([JAX docs](https://jax.readthedocs.io/en/latest/jax-101/01-jax-basics.html#differences-from-numpy)).\nTypically, these side effects are the result of mutating an object that lives outside the function.\n\nThe functional paradigm has some advantages like the ability to explicitly reason about state and stochasticity.\nThe function output only changes when an input argument changes.\nTherefore, a function is guaranteed to behave deterministically.\n\nBut pure functions offer another big advantage to JAX: specifically, they enable functional transformations.\nFor example `jax.vmap(f)` will vectorize a function `f`.\nBecause `f` cannot have side effects the vectorized/parallel version of `f` is well-defined. To see why we need this restriction, consider what happens if `f` would increment a counter or draw a random number.\nWould `f` draw the same or a different random number for each item in the vector?\nWould each item in the batch have its own counter or is the counter shared among the items?\nAnd in what order is the counter incremented if `f` is computed in parallel?\nThe answer to all these questions is \"it depends\".\nThe behavior is ambiguous and the functional constraint elegantly avoids this problem.\n\nFlax introduces a safe way to have limited randomness and stateful variables in a JAX-compatible form.\nThe reason why the state in Flax is not problematic is because it is local: inside a Flax `Module` there are variables and PRNG sequences,\nbut on the outside there are only JAX Arrays and PRNG keys.\n\nFor most use cases, Flax is used to define models in a stateful way.\nBecause a `Module` behaves like a pure function externally, we can fully utilize JAX with all of its transformations.\nThere are, however, cases when we want to have the best of both worlds by using transformations and `Module` together.\nThis design note explains how we extend JAX's functional transformation to work on `Module`s that have internal state and randomness.\n\n\n## Functionalization\n\nBefore we jump into the details let's consider a simple example where we would like to use `vmap` inside a `Module`.\n\nFirst, we define a simple MLP without any transformations:\n\n```python\nimport jax\nfrom jax import random, numpy as jnp\nfrom flax import linen as nn\n\nclass MLP(nn.Module):\n @nn.compact\n def __call__(self, xs):\n h = nn.Dense(4, name='hidden')(xs)\n h = nn.relu(h)\n return nn.Dense(1, name='out')(h)\n```\n\nNow what if we want to have separate MLP parameters for each item in `xs`?\nIf this were \"vanilla JAX\" we could imagine writing something like `jax.vmap(apply_mlp)(mlp_params, xs)`.\nBut doing something like this in Linen will actually fail:\n\n```python\nclass NaiveVmapMLP(nn.Module):\n @nn.compact\n def __call__(self, xs):\n mlp = MLP()\n return jax.vmap(lambda mlp, x: mlp(x))(mlp, xs) # fails\n```\n\nJAX will raise an error when `vmap` is used on `mlp` because it's not a JAX array or a simple container of arrays.\nWe can not really blame JAX for refusing to perform this under-specified job.\nAfter all, it's not even clear what should happen here.\nThe parameters inside the MLP are not even initialized yet and we will need a separate PRNG key for each group of parameters.\n`jax.vmap` can only broadcast or map over an axis but it cannot automatically split an PRNG key.\nTherefore, we have to call `jax.random.split` manually.\n\nWe can fix this problem by first turning `MLP` into a pure init and apply function.\nAfterwards, we use the `param` method to store the parameters:\n\n```python\nclass ManualVmapMLP(nn.Module):\n @nn.compact\n def __call__(self, xs):\n mlp = MLP(parent=None)\n init_fn = lambda rng, xs: jax.vmap(mlp.init, in_axes=0)(random.split(rng, xs.shape[0]), xs)['params']\n apply_fn = jax.vmap(mlp.apply, in_axes=0)\n mlp_params = self.param('mlp', init_fn, xs)\n return apply_fn({'params': mlp_params}, xs)\n\nxs = jnp.ones((3, 4))\nvariables = ManualVmapMLP().init(random.key(0), xs)\nprint(jax.tree_util.tree_map(jnp.shape, variables['params']))\n\"\"\"==>\n{\n mlp: {\n hidden: {\n bias: (3, 4),\n kernel: (3, 4, 4),\n },\n out: {\n bias: (3, 1),\n kernel: (3, 4, 1),\n },\n },\n}\n\"\"\"\n```\n\nHere, `MLP(parent=None)` creates a detached instance of `MLP`.\nThis avoids reserving a name for the submodule inside the current module.\nAlthough not strictly necessary, this also ensures we cannot accidentally use the MLP instance in a stateful way and we are forced to use it through either `.init` or `.apply`.\n\nThis example is still relatively concise but it already takes a few extra \"bookkeeping\" statements to make it work.\nHowever, this implementation has a number of limitations:\n1. During initialization, we call the submodule twice through `init_fn` and `apply_fn`. If the submodule used the same trick to do\n functional transformation we will end up executing a lot of code as the number of module calls grows like 2^d where d is the number of\n nested function transformations.\n2. The implementation assumes the submodule only requires the parameter RNG sequence.\n3. The implementation assumes we only create variables in the \"params\" collection during `init`. However, it does not support other variable collections and creating/updating variables in `apply`.\n\nPoint 3 in particular makes manual functionalization cumbersome.\nFeel free to try and extend the above example with a `nn.BatchNorm` layer in the `MLP` module.\nThis will require dealing with some additional complexity like storing the updated batch stats and making sure the batch stats are not mutable inside `vmap` when it should be immutable (e.g.: eval mode).\n\n\nWe call the process of transforming a stateful Module into a pure function \"functionalization\".\nBy temporarily turning a stateful `Module` into a function we make it compatible with JAX's functional transformations.\n\n## Lifting\n\nFlax provides an alternative for manual functionalization which we call lifted transformation.\nLifted transformations are defined in `flax.core.lift`.\nAll the lifted JAX transformations are defined with a single generic lifting API called `pack`.\n\nA number of decisions had to be made in order to define `pack`. The implementation\nof `pack` controls how variables and rngs are lifted and how fine-grained the user control is.\nIt must also decide whether lifting decisions are made at variable or transformation definition.\n\n\n### Lifting granularity\n\n\nWith the Linen API, users can define arbitrary variable collections and PRNG sequences.\nEach variable in a collection is lifted in the same way.\n\nCollections are typically given a semantically meaningful name like \"params\" or \"batch_stats\" rather than a general purpose name like \"state\".\nBecause collections carry semantic meaning we can decide at the transformation level how each collection should be lifted.\nFor example, we want to share all parameter variables when we add a batch dimension to a model.\n\nAt the same time we can write generic code that uses transformations without knowing exactly what kind of variables the submodules will create.\nCollections thus strike a balance between fine-grained control and generality.\nWe also avoid brittle string matching code that loops over all variables and tries to split up collections in an ad-hoc way based on\nnaming conventions like: target all variables with the name prefix \"kernel\".\nIf more fine-grained control is necessary a user can simply split up a set of variables over multiple collections that should be handled differently.\n\n\n### Transformation vs variable control\n\n\nLifting behavior could be defined either at the transformation level or during variable definition.\nWe use transformation level definitions of lifting behavior.\nThe reason for this choice is that there are many different transformations with various behaviors.\nFor example: `vmap` has broadcasted and vectorized arguments, while `scan` has scan, carry, and broadcast arguments.\nA variable would have to define its behavior for all these transformations otherwise a `Module` would not be compatible with\nthese transformations. Alternatively, we would have to make default decisions for how transformations are handled.\nHowever, this could lead to silent bugs because the behavior might not actually be valid given the users intent.\n\nThe lift package also provides a general purpose `transform`, which allows an arbitrary function to transform a variable collection.\nFor example, this can be used to tie the weights in a tied auto-encoder by transposing the weights.\nIt is unclear whether a similar general purpose transform could be defined if lifting decisions were made at variable definition.\n\n\n### Linen\n\nThe lifting module does not know about the Linen `Module` API.\nInstead it operates directly on instances of `flax.core.Scope`.\nA `Scope` instance contains the variables and PRNG sequences of a `Module`.\nEach `Module` instance has a `Scope` instance in the `.scope` field if it has a parent or it was created using `init` or `apply`.\nTypically, the top-level `Module` instance — on which you call `init` or `apply` — is the only `Module` instance that does not have a `Scope` bound to it.\n\nWhen a `Module` is transformed, we use the `flax.core.lift` APIs to lift the scope and use `Module.clone()` to create a new `Module` instance with the lifted scope bound to it.\n\n`flax.linen.transforms` exposes wrappers for the transformations in `flax.core.lift`. The core lifting APIs operate on functions while\nthe Linen wrappers can transform either a `Module` class or a `Module` method.\n\nThus, lifting is implemented independently from the Linen API. This separation of concern simplifies the implementation, while potentially allowing alternative `Module` abstractions to build upon a common core for lifting and state management.\n\n\n### Implementation\n\nThe `pack(fn, in_vars, out_vars, rngs)` API goes through the following stages:\n\n\n1. *Scope de-duplication*\n\n This stage is only relevant if multiple Scopes are lifted together.\n In this case we must first find the set of root scopes.\n A scope is a root if none of its ancestors are in the set of scopes that need to be lifted.\n\n By only lifting roots we avoid lifting the same variables twice.\n\n For non-root scopes we store a reference to its ancestor scope and a path such that we can later reconstruct it (stage 4).\n\n2. *Filter stage*\n\n Variables and PRNG sequences are split up into groups. This way `fn` can lift each group into the transformation separately.\n A group is defined by a filter specified as:\n - a list of collections/prng names\n - `True` (match everything)\n - `False` (match nothing)\n - `DenyList(filter)` (match everything but the specified collections (e.g.: `DenyList(['params'])` matches everything except the 'params' collection.)).\n\n A collection or PRNG sequence can only be put into a single group. If a collection matches multiple filters, it will be put into the first group with a matching filter.\n If a collection or PRNG sequence does not match any filter it will not be lifted.\n This means that it cannot be used inside the transformation and attempting to do this will cause an error to be raised.\n For example, `in_vars = ([\"params\"], True)` will cause the \"params\" collection to be put in the first group and all other collection to be put in the second group.\n\n For each PRNG sequence that is matched we seed a new PRNG sequence by calling `make_rng`.\n This avoids the need to update the PRNG state after the lifted transformation is complete.\n\n3. *Transform-specific lifting*\n\n `fn` is called with the variable and PRNG groups.\n JAX transforms have varying signatures and lifting options. Arguably the cleanest example is `vmap`.\n In the case of vmap the function arguments, PRNGs and variable collections are passed into a `jax.vmap` wrapped function.\n\n4. *Scope reconstruction*\n\n Now that the variables and PRNGs are lifted inside the transformation, we want to recreate the lifted scopes. Pack calls\n `fn` with a `scope_fn` that takes the lifted variables and PRNGs and returns the reconstructed scopes with the lifted variables and rng sequences.\n\n5. *Repack stage*\n\n After we have used the lifted scopes we have to retrieve the updated variables (PRNG sequences can simply be discarded).\n pack passes the `repack_fn` to support this.\n This stage is similar to stage 2 except that we only lift variables and immutable variables are ignored.\n Immutable variables cannot be updated. Therefore, they should not be returned from the transformed function.\n\n6. *Commit stage*\n\n `pack` expects `fn` to return a pair where the first item will simply be returned from pack and the second item should be the repacked variables.\n The updated variables are stored in the original/un-lifted scopes such that the mutations that happen inside the transformation survive after the transformation completes.\n\n\n### Using pack example\n\n\nA minimal example of using `pack` to transpose each matrix in a variable collection:\n\n```python\nfrom flax.core import lift\nfrom flax.core import Scope, init, apply, nn as core_nn\n\ndef lift_transpose(fn, target='params', variables=True, rngs=True):\n # by default we transpose 'params' and simply pass through all other variables.\n def wrapper(scope_fn, repack_fn, variable_groups, rng_groups, *args):\n # normally we would first call into a JAX transformed function here...\n target, rest = variable_groups\n def trans(x):\n if x.ndim == 2:\n return x.T\n return x\n target = jax.tree_util.tree_map(trans, target)\n variable_groups = (target, rest)\n scope = scope_fn(variable_groups, rng_groups)\n y = fn(scope, *args)\n out_variables = repack_fn(scope)\n return y, out_variables\n return lift.pack(\n wrapper,\n in_variable_filters=(target, variables),\n out_variable_filters=(variables,),\n rng_filters=(rngs,))\n\nx = jnp.ones((3, 2))\ny, params = init(lift_transpose(core_nn.dense))(random.key(0), x, 4)\n```\n\nNOTE that most users should not need to interact with `pack` directly.\nPlease open a GitHub issue when you find a use case that is not supported yet by the existing lifted transformations.\n\n### Supported transformations\n\n| Jax Transform | Supported in Linen? | Comments |\n|-|-|-|\n| vmap | ✅ | |\n| scan | ✅ | Carry variables cannot be initialized inside the scan body. |\n| remat | ✅ | |\n| jit | ✅ | Current implementation might cause unnecessary recompilation. |\n| jvp | ✅ | |\n| vjp | ✅ | |\n| custom_vjp | ✅ | |\n| custom_jvp | ❌ | |\n| while_loop | ✅ | Carry variables cannot be initialized inside the while_loop body. |\n| cond | ✅ | Variable initialization / mutation must structurally match across branches. |\n| switch | ✅ | Variable initialization / mutation must structurally match across branches. |\n| pmap | ❌ | |\n| xmap | ❌ | |\n\nReferences:\n- [Linen transforms documentation](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/transformations.html)\n- [Linen transforms source code](https://github.com/google/flax/blob/main/flax/linen/transforms.py)\n- [Core lifting source code](https://github.com/google/flax/blob/main/flax/core/lift.py)\n\n### Linen examples\n\nGoing back to our original example, we can now use `nn.vmap` to simplify our implementation:\n\n```python\nclass LinenVmapMLP(nn.Module):\n @nn.compact\n def __call__(self, xs):\n VmapMLP = nn.vmap(MLP, variable_axes={'params': 0}, split_rngs={'params': True}, in_axes=0)\n return VmapMLP(name='mlp')(xs)\n\nvariables = LinenVmapMLP().init(random.key(0), xs)\nprint(jax.tree_util.tree_map(jnp.shape, variables['params']))\n\"\"\"==>\n{\n mlp: {\n Dense_0: {\n bias: (3, 4),\n kernel: (3, 2, 4),\n },\n Dense_1: {\n bias: (3, 1),\n kernel: (3, 4, 1),\n },\n },\n}\n\"\"\"\n```\n\nHere we use `variable_axes={'params': 0}` to indicate that parameters are vectorized rather than shared and `split_rngs={'params': True}` means each set of parameters is initialized independently.\n\nWe can also extend the example with some inner state by adding a `BatchNorm` layer:\n\n```python\nclass StatefulMLP(nn.Module):\n @nn.compact\n def __call__(self, x, *, train):\n h = nn.Dense(4, name='hidden')(x)\n h = nn.BatchNorm(axis_name='batch')(h, use_running_average=not train)\n h = nn.relu(h)\n return nn.Dense(1, name='out')(h)\n\nclass LinenStatefulVmapMLP(nn.Module):\n @nn.compact\n def __call__(self, xs, *, train):\n VmapMLP = nn.vmap(StatefulMLP, variable_axes={'params': 0, 'batch_stats': 0}, split_rngs={'params': True}, in_axes=0)\n return VmapMLP(name='mlp')(xs, train=train)\nvariables = LinenStatefulVmapMLP().init(random.key(0), xs)\n```\n\nAll we had to add to `nn.vmap` is `'batch_stats': 0`, indicating that the batch stats are vectorized rather than shared along the first axis.\n\n\n## Alternatives\n\nOther numerical computation frameworks consider variables a first-class citizen.\nAn alternative to functionalization would be to use a variable system either integrated or on top of JAX.\nAn advantage of this is that per-variable lifting becomes easier.\nIf variables are part of the JAX IR (JAXPR), we could inspect which variables have to be lifted in a certain computation.\nOptionally, they could be annotated with a collection tag to decide on various lifting options.\n\nThe downside of this approach is that a variable system is more complicated.\nVariables are related references and break a core assumption of Functional Programming (see [referential transparency](https://en.wikipedia.org/wiki/Referential_transparency))\nOther APIs that currently have a functional interface would probably require integration as well (e.g.: checkpointing and optimization APIs).\n" }, { "path": "docs/developer_notes/module_lifecycle.rst", "content": "The Flax Module lifecycle\n#########################\n\n.. testsetup::\n\n from typing import Any, Callable, Iterable\n import flax\n from flax import linen as nn\n from jax import random\n import jax\n\n\nThis design note is intended for users who are already familiar with Flax Linen Modules but want to understand more about the design principles behind the abstraction. This note should give you a good understanding of the assumptions and guarantees the Module API is built upon. If you have no practical experience with Modules yet, check out the `Quickstart guide `_.\n\nFlax Linen Modules offer a Pythonic abstraction on top of Flax core. The `Module `_ abstraction allows you to create classes that have state, parameters and randomness on top of JAX. This is a practical guide to the design and behavior of the ``Module`` class. By the end, you should feel comfortable to go off the beaten track and use Modules in new ways.\n\n\nOverview\n***********\n\nDefinition\n=============\n\nLet's start with a high-level overview of the Module lifecycle. First, define a simple Module:\n\n\n.. testcode::\n\n class MLP(nn.Module):\n # 1. Attribute annotations\n hidden_size: int\n out_size: int\n\n # 2. The ``setup`` method\n def setup(self):\n self.hidden = nn.Dense(self.hidden_size)\n self.out = nn.Dense(self.out_size)\n\n # 3. User methods\n def __call__(self, x):\n a = self.hidden(x)\n h = nn.relu(a)\n return self.out(h)\n\n\nThis Module consists of:\n\n#. **Attribute annotations**, defined as `dataclass `_ fields. These annotations automatically define a constructor.\n#. **The ``setup`` method**, which creates submodules and assigns them to attributes.\n#. **User methods**. By convention, most Modules have just one ``__call__`` method, but you can define multiple methods or use different method names.\n\nConstruction/initialization\n=============================\n\nNow we want to construct and use the ``MLP`` Module:\n\n\n.. testcode::\n\n mlp = MLP(hidden_size=5, out_size=3)\n x = jax.numpy.ones((1, 2))\n variables = mlp.init(random.key(0), x)\n y = mlp.apply(variables, x)\n\n\nFirst, we construct an instance of ``MLP`` and pass the construction attributes. Note that construction here is different from what you might expect if you are not used to Functional Programming patterns. The ``MLP`` constructor does not actually create variables or any internal state whatsoever. It's best to think of it as a specification or template of the Module that contains functionality but no data.\n\nLet's take a closer look at initialization. Surprisingly, there actually is no separate initialization path in Flax. Calling ``init`` is just a special case of ``apply``, which you can also write as:\n\n\n.. testcode::\n\n # equivalent to: variables = mlp.init(random.key(0), x)\n _, variables = mlp.apply({}, x, rngs={\"params\": random.key(0)}, mutable=True)\n\n\nThus, ``init`` is nothing more than a wrapper around ``apply`` where:\n\n#. We call a Module without any initial variables (an empty dict).\n#. A PRNG generator named ``\"params\"`` is always passed for randomly initializing parameters (using the parameter initialization function).\n#. All variable collections are set to mutable (``mutable=True``). When a collection is mutable, existing variables can be updated and new variables can be created. Thus, inside ``init`` variables can be initialized in any variable collection and they are all added to the returned variable dictionary.\n\nLifecycle\n=============\n\n\nNow that you have learned about ``init`` being a special case of ``apply``, let's look at ``.apply(...)`` in more detail. In fact, most of the complexity of Modules resides in the ``apply`` method. The \"Module lifecycle\" consists of constructing and ``apply``-ing a Module. We can summarize the Module lifecycle as follows:\n\n\n#. We construct ``mlp = MLP(hidden_size=5, out_size=3)``, such that ``mlp.hidden_size=5`` and ``mlp.out_size=3``.\n\n#. Then, call ``mlp.apply``, which:\n\n #. Makes a clone of ``mlp``, let's call it ``mlp_copy``.\n\n #. Calls ``mlp_copy.setup()``.\n\n #. Returns the output of ``mlp_copy.__call__()`` and optionally the variable collections that were specified as mutable using the keyword argument ``mutable=``.\n\nNotice that the lifecycle includes cloning the Module instance. This is done to ensure that ``apply`` can be treated as a pure function (i.e., if you pass the same arguments in, it will return the same outputs). You will learn about this in more detail later in the :ref:`Top-level Modules` section.\n\nVariables\n==========\n\nThe word “variable” is ubiquitous in programming and math. However, it's important to have a good understanding of what variables are in the context of JAX and Flax. Inside Flax Modules, `variables `_ act like you expect from Python. They are initialized once, read, and perhaps even updated every so often. However, JAX has no concept of variables. Instead, values are stored in arrays similar to NumPy arrays - with one important difference: they are immutable.\n\nThe ``init`` and ``apply`` methods return the variables as a nested dictionary with string keys and JAX arrays at the leaves. At the top level each key corresponds to a variable collection. Inside each collection the nested dict structure corresponds with the ``Module`` hierarchy. The variable dict is immutable and therefore really just a snapshot of state the variables are in. When ``apply`` is called again, the variable dict is passed as an argument. Such that the variables are in the same state as when the previous ``init`` / ``apply`` call finished.\n\n\n.. note::\n Module fields are declared using the `field_name: TypeHint` syntax (same as dataclasses). Without a type hint, an attribute is considered a static property of the class. In case you cannot specify the type you can use ``typing.Any`` as a wildcard type.\n\n\nCompact Modules\n******************\n\nLinen provides an alternative API for defining modules more compactly. This is especially useful for the common case where the Module consists of only one method that uses parameters and/or sub-modules. Using the compact API the MLP can be rewritten as follows:\n\n\n.. testcode::\n\n class CompactMLP(nn.Module):\n hidden_size: int\n out_size: int\n\n @nn.compact\n def __call__(self, x):\n a = nn.Dense(self.hidden_size)(x)\n h = nn.relu(a)\n return nn.Dense(self.out_size)(h)\n\n\nA compact ``Module`` is similar in spirit to a function. It offers a concise notation and restricts external interaction to the inputs and return values of the function. In this case the concise notation might make it easier for others to understand what the Module does. There is no need to jump back and forth between the ``setup`` and ``__call__`` method to understand what the submodules are doing. Instead, simply reading the ``__call__`` method from top to bottom once should provide a concise overview. This can make a significant difference if you are implementing complex Modules with many hyperparameters. See `setup or compact `_ for a practical guide on deciding between setup and compact.\n\nAnother benefit of defining submodules and/or variables inline is that you can add arguments to your method when constructing variables. The most common example of this is using shape information to determine the shape of a parameter like this:\n\n\n.. testcode::\n\n class CompactScaledMLP(nn.Module):\n hidden_size: int\n out_size: int\n\n @nn.compact\n def __call__(self, x):\n scale = self.param(\"scale\", nn.initializers.ones_init(), x.shape[-1:])\n x *= scale[None]\n a = nn.Dense(self.hidden_size)(x)\n h = nn.relu(a)\n return nn.Dense(self.out_size)(h)\n\n\n.. testcode::\n :hide:\n\n mdl = CompactScaledMLP(hidden_size=4, out_size=5)\n x = jax.numpy.ones((3, 2))\n vars = mdl.init(random.key(0), x)\n assert vars[\"params\"][\"scale\"].shape == (2,)\n\nMany of the standard Linen Modules like ``nn.Dense`` use shape inference already to avoid the need to specify input shapes (like the number of input features to a Dense layer).\n\nCompact control flow\n=====================\n\nThe order in which you define submodules determines the name of a submodule if none is provided explicitly (using the ``name=`` keyword argument passed to the Module's constructor). Because the ``name`` determines how parameters are mapped to submodules, you must be careful about mixing control flow with auto-generated names. Using control flow can change the order or remove certain submodules altogether. This is useful in case a submodule should only exist depending on some construction argument. However, when control flow depends on the input arguments to the Module, you should be careful. For example, the following Module will break:\n\n\n.. testcode::\n\n class WrongModule(nn.Module):\n @nn.compact\n def __call__(self, x, mode):\n if mode == \"encode\":\n return nn.Dense(features=8)(x)\n elif mode == \"decode\":\n return nn.Dense(features=4)(x)\n\n\nThe above Module will break because either the encoder or decoder path will construct a Module named \"Dense_0\". This means the two Modules will share parameters which is not intended here. Actually, the two Modules cannot share parameters because they each have a different number of features.\n\nThis problem can be solved in various ways:\n - Provide explicit names\n - create the modules in ``setup``\n - or move the constructor out of the control flow.\n\nThe latter is done as follows:\n\n.. testcode::\n\n class CorrectModule(nn.Module):\n @nn.compact\n def __call__(self, x, mode):\n encoder = nn.Dense(8)\n decoder = nn.Dense(4)\n if mode == \"encode\":\n return encoder(x)\n elif mode == \"decode\":\n return decoder(x)\n\n.. testcode::\n :hide:\n\n def init_fn(mdl):\n x = jax.numpy.ones((3, 2))\n z = mdl(x, \"encode\")\n return mdl(z, \"decode\")\n\n mdl = CorrectModule()\n vars = nn.init(init_fn, mdl)(random.key(0))\n assert vars[\"params\"][\"Dense_0\"][\"kernel\"].shape == (2, 8)\n assert vars[\"params\"][\"Dense_1\"][\"kernel\"].shape == (8, 4)\n\n\nIn the above example the construction order is fixed. After construction the submodules can be used in an arbitrary order.\n\n.. note::\n compact modules show a strong resemblance to `React hooks `_.\n\n\nTop-level Modules\n*****************\n\nWhen a Module instance is created at the \"top-level\", it will be in an \"unbound\" state - that is, it has no variables attached. \"Top-level\" means it is not constructed as a sub-Module inside another Module class. Apart from calling ``init`` and ``apply``, there is not much you can do with an unbound Module. Note also that ``setup`` is not called on unbound Modules, so you can only access the construction arguments. Refer to the :ref:`Future work` section to learn how this might change in the future.\n\nWhy are top-level Modules always unbound?\n===============================================\n\nWhen we call ``apply``, a copy of the top-level Module is created which will actually hold the variables and PRNG sequences. This stateful, \"bound\", clone only exists while we are executing the apply method. The reason for this is that if you create a stateful object and destroy it before the apply function returns, the ``apply`` function itself behaves like a pure function. A pure function has two constraints:\n\n#. If you put the same arguments in, it will return the same outputs\n#. It does not change anything outside the function. This means you cannot manipulate stateful objects that are accessible outside the pure function.\n\n\nPure functions have many advantages but when using JAX they are often essential. For example, most code requires compilation using ``jax.jit`` to be fast and once you created a Module you probably want to optimize its parameters using ``jax.grad``. However, these APIs expect a pure function and don't work on stateful bound ``Module`` instances directly. Moreover, pure functions allow for flexible interoperability with other libraries. For example, We recommend `Optax `_ for optimizing parameters. The optimizers in Optax expect and return a PyTree of JAX arrays to optimize, just like the ``apply`` function of a Linen Module.\n\nCloning\n===============================================\n\nTo make this approach work reliably we need well-defined cloning behavior. Rather than relying on a complex nested cloning procedure like Python's ``deepcopy``, Flax enforces that a ``Module`` is exactly defined by its construction arguments. Therefore cloning a Module reduces to calling the constructor with its original construction arguments. Because ``Module`` acts as an immutable dataclass, the construction arguments are mapped directly to instance attributes. Non-construction attributes that are computed in ``setup`` or ``__post_init__`` should also depend only on the construciton arguments to ensure a well-defined clone.\n\nBind\n===============================================\n\nSometimes it's useful to have a bound, top-level Module without having to wrap the code in a function. For example: to interact with a Module inside a Jupyter notebook. The `bind `_ method returns a bound clone with an unlimited lifetime. The downside of this is that you cannot combine it with JAX transformations or integrate it into a vanilla JAX codebase that expects stateless code. For example, `Optax `_ can optimize a Pytree of parameters but it cannot directly optimize a bound ``Module`` instance created with ``.bind`` (because that's not a Pytree). Thus, you cannot combine the ``bind`` API with a functional optimizer API like Optax.\n\n\nSetup\n**********\n\nThe ``setup`` method is often used like the constructor hook (``__init__``) in normal Python classes. However, for more advanced use cases it's good to realize that it is not quite the same as a constructor.\n\n``setup`` is only called after a Module becomes bound. Normally, this is not an issue because most Modules are bound (almost) immediately (as part of ``init`` and ``apply``). Inside ``setup``, sub-modules become bound when they are assigned to an attribute. Inside an ``nn.compact`` decorated method, sub-modules are bound immediately when constructed. As explained in the previous section, top-level Modules are never bound and thus setup is not called when they are constructed. This means you cannot access attributes assigned in setup from an unbound, top-level module.\n\n.. testcode::\n\n class TopLevelAccess(nn.Module):\n\n def setup(self):\n self.foo = nn.Dense(2)\n\n mdl = TopLevelAccess()\n assert not hasattr(mdl, \"foo\") # foo is not defined because setup is not called\n\nThe ``setup`` method is not called immediately after the ``Module`` becomes bound but only when you interact with the ``Module`` instance (e.g.: call a method or access an attribute). This should not impact the behavior of a ``Module`` but the lazy execution does sometimes affect log statements and stack traces during debugging. The section on :ref:`Functionalization` will explain why we need ``setup`` to be lazy in the first place.\n\n\nFunctionalization\n******************\n\nSo far we had a pure ``apply`` function that is typically transformed with some JAX transformations and inside ``apply`` we have a stateful Module instance to work with. In other words: Outside of a Module we are in a functional world where we have the power of JAX's functional transformations and inside the Module we get the power of Flax's stateful variables and PRNG sequence, and the ``apply`` method is our bridge between these two worlds.\n\nBut what if we want to use JAX transformations **inside** Modules? The answer to this is functionalization.\n\nThis procedure itself is tedious and error-prone but handled internally by Flax. At a high-level we can summarize it as follows. For a method ``fn`` defined within a Module:\n\n#. Collect the state (variables & PRNG sequences) of the Module(s) that should be available inside the JAX transformation and take a snapshot of it.\n\n#. Call the JAX transformation with the original arguments and the collected state. Then inside the transformation:\n\n #. Unpack the state and recreate the Modules\n\n #. Call the user code ``fn``\n\n #. Collect the updated variables and rng and return it together with the original return values from ``fn``\n\n#. Update the original state with the updated state returned from the transformation.\n\nA more in depth explanation of functionalization and lifting can be found in the `Lifted Transformation `_ design note.\n\nPractical consequences\n==========================\n\nFor the most part functionalization is something that is handled automatically for you. Still there are some constraints that you must take into account. Most importantly, Flax only handles the stateful primitives (Linen variables and RNGs) and not arbitrary stateful Python code. Most importantly: You cannot close over stateful objects and ``Module`` objects because they are invisible to Flax's internals (and to JAX in general).\n\n\n.. testcode::\n\n class Foo(nn.Module):\n @nn.compact\n def __call__(self, x):\n dense = nn.Dense(x.shape[-1])\n fn = lambda x: dense(x) + 1\n # simply calling inner works fine\n # return self.inner(x, fn)\n # but applying a transformation doesn't:\n vmap_inner = nn.vmap(Foo.inner, in_axes=0, variable_axes={\"params\": 0}, split_rngs={\"params\": True})\n return vmap_inner(self, x, fn)\n\n def inner(self, x, fn):\n for i in range(3):\n x = fn(x)\n return x\n\nHere ``inner`` takes a function that closes over a Module instance. In this example, that works fine because we are not transforming the inner method with a lifted transformation. Most methods are not transformed but it is good to know how to make Module methods transformable.\n\nThe main obstacle for transformability are types that JAX does not recognize. JAX only understands `Pytree `_ arguments; i.e. arbitrarily nested Python containers (dict, list, tuple) of (Jax) numpy ndarrays and Python numbers/bools. Flax allows to define dataclasses which are Pytree compatible using the `flax.struct `_ API.\n\nFunction closure is the most common way to accidentally hide a JAX array or Linen Module from a transformation. There is however an easy workaround if you want to pass closures that are also compatible with JAX and Linen transformations:\n\n\n.. testcode::\n\n class Partial(flax.struct.PyTreeNode):\n fn: Callable = flax.struct.field(pytree_node=False)\n args: Iterable[Any]\n\n def __call__(self, *args, **kwargs):\n return self.fn(*(tuple(self.args) + args), **kwargs)\n\n class Foo(nn.Module):\n\n @nn.compact\n def __call__(self, x):\n dense = nn.Dense(x.shape[-1])\n fn = lambda mdl, x: mdl(x) + 1\n vmap_inner = nn.vmap(Foo.inner, in_axes=0, variable_axes={\"params\": 0}, split_rngs={\"params\": True})\n return vmap_inner(self, x, Partial(fn, [dense]))\n\n def inner(self, x, fn):\n for i in range(3):\n x = fn(x)\n return x\n\n\n.. testcode::\n :hide:\n\n x = jax.numpy.ones((3, 2))\n mdl = Foo()\n vars = mdl.init(random.key(0), x)\n assert vars['params']['Dense_0']['kernel'].shape == (3, 2, 2)\n\n\n\nHere the closure is implemented using a Flax dataclass. The function itself is annotated with ``flax.struct.field(pytree_node=False)`` to indicate that it does not contain JAX Arrays or Linen Modules. The partially applied ``args`` on the other hand is treated as a pytree container. We rewrite the closure to use Partial. Now the inner method can be transformed using lifted transformations.\n\n\nFuture work\n*************\n\n\nSetup for unbound Modules\n===========================\n\nThe current Module abstraction is particularly restrictive when it comes to initializing fields after construction. In the current Module API, the ``setup`` method is the place to initialize the fields of the Module instance. Because ``setup`` is only called on a bound Module, the full Module API is available inside ``setup``, including variable declaration. However, oftentimes we don't actually require any stateful API's to initialize a field. In fact, most commonly we simply want to declare a submodule. More importantly, it's often useful to inspect submodules for debugging or to partially run the model. Consider for example:\n\n\n.. testcode::\n\n class AutoEncoder(nn.Module):\n def setup(self):\n self.encoder = Encoder(...)\n self.decoder = Decoder(...)\n\n\nImagine we want to call just the decoder using `auto_encoder.decoder.apply(decoder_variables, x)`. With the current setup API this does not work because we must first bind the variables before setup is called and the decoder attribute is defined. Of course we can manually construct the Decoder Module with the same attributes as in setup but this is not ideal in many cases.\n\nThere are two possible solutions to make this use case more ergonomic. First, setup could be made to run immediately after construction before it becomes bound. This means you can still create sub modules but you can no longer define or manipulate variables. Therefore, this would be a breaking change and it would require a new API for defining variables lazily\n\nAlternatively, an additional special method could be introduced that runs right away after Module construction and before it becomes bound. In this case, the ``setup`` method would preserve its original semantics.\n" }, { "path": "docs/examples/community_examples.rst", "content": "Community examples\n==================\n\nIn addition to the `curated list of official Flax examples on GitHub `__,\nthere is a growing community of people using Flax to build new types of machine\nlearning models. We are happy to showcase any example built by the community here!\n\nIf you want to submit your own Flax example, you can start by forking\none of the `official Flax examples on GitHub `__.\n\nModels\n******\n.. list-table::\n :header-rows: 1\n\n * - Link\n - Author\n - Task type\n - Reference\n * - `matthias-wright/flaxmodels `__\n - `@matthias-wright `__\n - Various\n - GPT-2, ResNet, StyleGAN-2, VGG, ...\n * - `DarshanDeshpande/jax-models `__\n - `@DarshanDeshpande `__\n - Various\n - Segformer, Swin Transformer, ... also some stand-alone layers\n * - `google/vision_transformer `__\n - `@andsteing `__\n - Image classification, image/text\n - https://arxiv.org/abs/2010.11929, https://arxiv.org/abs/2105.01601, https://arxiv.org/abs/2111.07991, ...\n * - `jax-resnet `__\n - `@n2cholas `__\n - Various resnet implementations\n - `torch.hub `__\n * - `Wav2Vec2 finetuning `__\n - `@vasudevgupta7 `__\n - Automatic Speech Recognition\n - https://arxiv.org/abs/2006.11477\n\nExamples\n********\n\n.. list-table::\n :header-rows: 1\n\n * - Link\n - Author\n - Task type\n - Reference\n * - `JAX-RL `__\n - `@henry-prior `__\n - Reinforcement learning\n - N/A\n * - `BigBird Fine-tuning `__\n - `@vasudevgupta7 `__\n - Question-Answering\n - https://arxiv.org/abs/2007.14062\n * - `DCGAN `__\n - `@bkkaggle `__\n - Image Synthesis\n - https://arxiv.org/abs/1511.06434\n * - `denoising-diffusion-flax `__\n - `@yiyixuxu `__\n - Image generation\n - https://arxiv.org/abs/2006.11239\n\nTutorials\n*********\n\n.. currently left empty as a placeholder for tutorials\n.. list-table::\n :header-rows: 1\n\n * - Link\n - Author\n - Task type\n - Reference\n * -\n -\n -\n -\n\nContributing policy\n*******************\n\nIf you are interested in adding a project to the Community Examples section, take the following\ninto consideration:\n\n* **Code examples**: Examples must contain a README that is helpful, clear, and explains\n how to run the code. The code itself should be easy to follow.\n* **Tutorials**: These docs should preferrably be a Jupyter Notebook format\n (refer to `Contributing `__\n to learn how to convert a Jupyter Notebook into a Markdown file with `jupytext`).\n Your tutorial should be well-written, and discuss/describe an interesting topic/task.\n To avoid duplication, the content of these docs must be different from\n `existing docs on the Flax documentation site `__\n or other community examples mentioned in this document.\n* **Models**: repositories with models ported to Flax must provide at least one of the following:\n\n * Metrics that are comparable to the original work when the model is trained to completion. Having\n available plots of the metric's history during training is highly encouraged.\n * Tests to verify numerical equivalence against a well known implementation (same inputs\n + weights = same outputs) preferably using pretrained weights.\n\nIn all cases mentioned above, the code must work with the latest stable versions of the\nfollowing packages: ``jax``, ``flax``, and ``optax``, and make substantial use of Flax.\nNote that both ``jax`` and ``optax`` are `required packages `__\nof ``flax`` (refer to the `installation instructions `__\nfor more details).\n" }, { "path": "docs/examples/core_examples.rst", "content": "Core examples\n=============\n\nCore examples are hosted on the GitHub Flax repository in the `examples `__\ndirectory.\n\nEach example is designed to be **self-contained and easily forkable**, while\nreproducing relevant results in different areas of machine learning.\n\nAs discussed in `#231 `__, we decided\nto go for a standard pattern for all examples including the simplest ones (like MNIST).\nThis makes every example a bit more verbose, but once you know one example, you\nknow the structure of all of them. Having unit tests and integration tests is also\nvery useful when you fork these examples.\n\nSome of the examples below have a link \"Interactive🕹\" that lets you run them\ndirectly in Colab.\n\nImage classification\n********************\n\n- :octicon:`mark-github;0.9em` `MNIST `__ -\n `Interactive🕹 `__:\n Convolutional neural network for MNIST classification (featuring simple\n code).\n\n- :octicon:`mark-github;0.9em` `ImageNet `__ -\n `Interactive🕹 `__:\n Resnet-50 on ImageNet with weight decay (featuring multi-host SPMD, custom\n preprocessing, checkpointing, dynamic scaling, mixed precision).\n\nReinforcement learning\n**********************\n\n- :octicon:`mark-github;0.9em` `Proximal Policy Optimization `__:\n Learning to play Atari games (featuring single host SPMD, RL setup).\n\nNatural language processing\n***************************\n\n- :octicon:`mark-github;0.9em` `Sequence to sequence for number\n addition `__:\n (featuring simple code, LSTM state handling, on the fly data generation).\n- :octicon:`mark-github;0.9em` `Parts-of-speech\n tagging `__: Simple\n transformer encoder model using the universal dependency dataset.\n- :octicon:`mark-github;0.9em` `Sentiment\n classification `__:\n with a LSTM model.\n- :octicon:`mark-github;0.9em` `Transformer encoder/decoder model trained on\n WMT `__:\n Translating English/German (featuring multihost SPMD, dynamic bucketing,\n attention cache, packed sequences, recipe for TPU training on GCP).\n- :octicon:`mark-github;0.9em` `Transformer encoder trained on one billion word\n benchmark `__:\n for autoregressive language modeling, based on the WMT example above.\n\nGenerative models\n*****************\n\n- :octicon:`mark-github;0.9em` `Variational\n auto-encoder `__:\n Trained on binarized MNIST (featuring simple code, vmap).\n\nGraph modeling\n**************\n\n- :octicon:`mark-github;0.9em` `Graph Neural Networks `__:\n Molecular predictions on ogbg-molpcba from the Open Graph Benchmark.\n\nContributing to core Flax examples\n**********************************\n\nMost of the `core Flax examples on GitHub `__\nfollow a structure that the Flax dev team found works well with Flax projects.\nThe team strives to make these examples easy to explore and fork. In particular\n(as per GitHub Issue `#231 `__):\n\n- README: contains links to paper, command line, `TensorBoard `__ metrics.\n- Focus: an example is about a single model/dataset.\n- Configs: we use ``ml_collections.ConfigDict`` stored under ``configs/``.\n- Tests: executable ``main.py`` loads ``train.py`` which has ``train_test.py``.\n- Data: is read from `TensorFlow Datasets `__.\n- Standalone: every directory is self-contained.\n- Requirements: versions are pinned in ``requirements.txt``.\n- Boilerplate: is reduced by using `clu `__.\n- Interactive: the example can be explored with a `Colab `__." }, { "path": "docs/examples/google_research_examples.rst", "content": "########################\nGoogle Research examples\n########################\n\nA collection of research by Google Research made with Flax.\n\nAttention\n*********\n\nFast Attention (FAVOR+) and Rethinking Attention with Performers\n================================================================\n\n- Code on GitHub:\n\n - `Performer's Fast Attention (FAVOR+) module `__\n\n- Research paper:\n\n - `Rethinking Attention with Performers `__ (Choromanski et al., 2020)\n\n - Introduces *\"Performers, Transformer architectures which can estimate regular (softmax) full-rank-attention Transformers with provable accuracy, but using only linear (as opposed to quadratic) space and time complexity, without relying on any priors such as sparsity or low-rankness. To approximate softmax attention-kernels, Performers use a novel Fast Attention Via positive Orthogonal Random features approach (FAVOR+), which may be of independent interest for scalable kernel methods. FAVOR+ can be also used to efficiently model kernelizable attention mechanisms beyond softmax.\"*\n\nSelf-attention Does Not Need O(n^2) Memory\n==========================================\n\n- `Code on GitHub `__\n- `Colab notebook `__\n\n- Research paper:\n\n - `Self-attention Does Not Need O(n^2) Memory `__ (Rabe and Staats, 2021)\n\n - *\"We present a very simple algorithm for attention that requires O(1) memory with respect to sequence length and an extension to self-attention that requires O(log n) memory. This is in contrast with the frequently stated belief that self-attention requires O(n^2) memory. While the time complexity is still O(n^2), device memory rather than compute capability is often the limiting factor on modern accelerators. Thus, reducing the memory requirements of attention allows processing of longer sequences than might otherwise be feasible...\"*\n\nComputer vision\n***************\n\nColorization Transformer (ColTran)\n==================================\n\n- `Code on GitHub `__\n\n- Research paper:\n\n - `Colorization Transformer `__ (Kumar et al., 2020)\n\n - *\"We presented the Colorization Transformer (ColTran), an architecture that entirely relies on self-attention for image colorization. We introduce conditional transformer layers, a novel building block for conditional, generative models based on self-attention. Our ablations show the superiority of employing this mechanism over a number of different baselines. Finally, we demonstrate that ColTran can generate diverse, high-fidelity colorizations on ImageNet, which are largely indistinguishable from the ground-truth even for human raters.\"*\n\nVision Transformer (ViT), MLP-Mixer Architectures *and* Big Vision\n==================================================================\n\n- Code on GitHub:\n\n - `Vision Transformer and MLP-Mixer Architectures `__\n\n - `Big Vision `__\n\n - *\"This codebase is designed for training large-scale vision models using Cloud TPU VMs or GPU machines. It is based on Jax/Flax libraries, and uses tf.data and TensorFlow Datasets for scalable and reproducible input pipelines.\"*\n\n- `Colab notebooks `__:\n\n - The JAX code of Vision Transformers and MLP Mixers\n - More than 50k Vision Transformer and hybrid checkpoints that were used to generate the data of \"How to train your ViT?\"\n\n- Research papers:\n\n - `An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale `__ (Dosovitskiy et al., 2020)\n\n - *\"In vision, attention is either applied in conjunction with convolutional networks, or used to replace certain components of convolutional networks while keeping their overall structure in place. We show that this reliance on CNNs is not necessary and a pure transformer applied directly to sequences of image patches can perform very well on image classification tasks. When pre-trained on large amounts of data and transferred to multiple mid-sized or small image recognition benchmarks (ImageNet, CIFAR-100, VTAB, etc.), Vision Transformer (ViT) attains excellent results compared to state-of-the-art convolutional networks while requiring substantially fewer computational resources to train.\"*\n\n - `MLP-Mixer: An All-MLP Architecture for Vision `__ (Tolstikhin et al., 2021)\n\n - *\"In this paper we show that while convolutions and attention are both sufficient for good performance, neither of them are necessary. We present MLP-Mixer, an architecture based exclusively on multi-layer perceptrons (MLPs). MLP-Mixer contains two types of layers: one with MLPs applied independently to image patches (i.e. \"mixing\" the per-location features), and one with MLPs applied across patches (i.e. \"mixing\" spatial information). When trained on large datasets, or with modern regularization schemes, MLP-Mixer attains competitive scores on image classification benchmarks, with pre-training and inference cost comparable to state-of-the-art models.\"*\n\n - `How to Train Your ViT? Data, Augmentation, and Regularization in Vision Transformers `__ (Steiner et al., 2021)\n\n - *\"Vision Transformers (ViT) have been shown to attain highly competitive performance for a wide range of vision applications, such as image classification, object detection and semantic image segmentation. In comparison to convolutional neural networks, the Vision Transformer's weaker inductive bias is generally found to cause an increased reliance on model regularization or data augmentation (\"AugReg\" for short) when training on smaller training datasets. We conduct a systematic empirical study in order to better understand the interplay between the amount of training data, AugReg, model size and compute budget.\"*\n\n - `When Vision Transformers Outperform ResNets without Pretraining or Strong Data Augmentations `__ (X. Chen et al., 2021)\n\n - *\"Vision Transformers (ViTs) and MLPs signal further efforts on replacing hand-wired features or inductive biases with general-purpose neural architectures. Existing works empower the models by massive data, such as large-scale pre-training and/or repeated strong data augmentations, and still report optimization-related problems (e.g., sensitivity to initialization and learning rates). Hence, this paper investigates ViTs and MLP-Mixers from the lens of loss geometry, intending to improve the models' data efficiency at training and generalization at inference.\"*\n\n - `LiT: Zero-Shot Transfer with Locked-image Text Tuning `__ (X. Zhai et al., 2021)\n\n - *\"This paper presents contrastive-tuning, a simple method employing contrastive training to align image and text models while still taking advantage of their pre-training. In our empirical study we find that locked pre-trained image models with unlocked text models work best. We call this instance of contrastive-tuning \"Locked-image Tuning\" (LiT), which just teaches a text model to read out good representations from a pre-trained image model for new tasks. A LiT model gains the capability of zero-shot transfer to new vision tasks, such as image classification or retrieval. The proposed LiT is widely applicable; it works reliably with multiple pre-training methods (supervised and unsupervised) and across diverse architectures (ResNet, Vision Transformers and MLP-Mixer) using three different image-text datasets.\"*\n\nScaling Vision with Sparse Mixture of Experts (MoE)\n===================================================\n\n- `Code on GitHub `__\n- Research paper:\n\n - `Scaling Vision with Sparse Mixture of Experts `__ (Riquelme et al., 2021)\n\n - *\"Sparsely-gated Mixture of Experts networks (MoEs) have demonstrated excellent scalability in Natural Language Processing. In Computer Vision, however, almost all performant networks are \"dense\", that is, every input is processed by every parameter. We present a Vision MoE (V-MoE), a sparse version of the Vision Transformer, that is scalable and competitive with the largest dense networks... we demonstrate the potential of V-MoE to scale vision models, and train a 15B parameter model that attains 90.35% on ImageNet...\"*\n\nDiffusion\n*********\n\nVariational Diffusion Models\n============================\n\n- `Code on GitHub `__\n- `Colab notebooks `__\n- Research paper:\n\n - `Variational Diffusion Models `__ (Kingma et al., 2021)\n\n - *\"Diffusion-based generative models have demonstrated a capacity for perceptually impressive synthesis, but can they also be great likelihood-based models? We answer this in the affirmative, and introduce a family of diffusion-based generative models that obtain state-of-the-art likelihoods on standard image density estimation benchmarks. Unlike other diffusion-based models, our method allows for efficient optimization of the noise schedule jointly with the rest of the model. We show that the variational lower bound (VLB) simplifies to a remarkably short expression in terms of the signal-to-noise ratio of the diffused data, thereby improving our theoretical understanding of this model class. Using this insight, we prove an equivalence between several models proposed in the literature. In addition, we show that the continuous-time VLB is invariant to the noise schedule, except for the signal-to-noise ratio at its endpoints. This enables us to learn a noise schedule that minimizes the variance of the resulting VLB estimator, leading to faster optimization...\"*\n\nDomain adaptation\n*****************\n\nGIFT (Gradual Interpolation of Features toward Target)\n======================================================\n\n- `Code on GitHub `__\n- Research paper:\n\n - `Gradual Domain Adaptation in the Wild: When Intermediate Distributions are Absent `__ (Abnar et al., 2021)\n\n - *\"We focus on the problem of domain adaptation when the goal is shifting the model towards the target distribution, rather than learning domain invariant representations. It has been shown that under the following two assumptions: (a) access to samples from intermediate distributions, and (b) samples being annotated with the amount of change from the source distribution, self-training can be successfully applied on gradually shifted samples to adapt the model toward the target distribution. We hypothesize having (a) is enough to enable iterative self-training to slowly adapt the model to the target distribution, by making use of an implicit curriculum. In the case where (a) does not hold, we observe that iterative self-training falls short. We propose GIFT, a method that creates virtual samples from intermediate distributions by interpolating representations of examples from source and target domains...\"*\n\nGeneralization\n**************\n\nSurrogate Gap Minimization Improves Sharpness-Aware Training\n============================================================\n\n- `Code on GitHub `__\n- Research paper:\n\n - `Surrogate Gap Minimization Improves Sharpness-Aware Training `__ (J. Zhuang et al., 2022)\n\n - *\"The recently proposed Sharpness-Aware Minimization (SAM) improves generalization by minimizing a perturbed loss defined as the maximum loss within a neighborhood in the parameter space. However, we show that both sharp and flat minima can have a low perturbed loss, implying that SAM does not always prefer flat minima. Instead, we define a surrogate gap, a measure equivalent to the dominant eigenvalue of Hessian at a local minimum when the radius of neighborhood (to derive the perturbed loss) is small. The surrogate gap is easy to compute and feasible for direct minimization during training. Based on the above observations, we propose Surrogate Gap Guided Sharpness-Aware Minimization (GSAM), a novel improvement over SAM with negligible computation overhead...\"*\n\nMeta learning\n*************\n\n``learned_optimization``\n=======================\n\n- Code on GitHub: `learned_optimization `__\n- `Colab notebooks `__\n\n- Research papers:\n\n - `Unbiased Gradient Estimation in Unrolled Computation Graphs with Persistent Evolution Strategies `__ (Vicol et al., 2021)\n\n - *\"We introduce a method called Persistent Evolution Strategies (PES), which divides the computation graph into a series of truncated unrolls, and performs an evolution strategies-based update step after each unroll. PES eliminates bias from these truncations by accumulating correction terms over the entire sequence of unrolls. PES allows for rapid parameter updates, has low memory usage, is unbiased, and has reasonable variance characteristics.\"*\n\n - `Gradients Are Not All You Need `__\t(Metz et al., 2021)\n\n - *\"...In this short report, we discuss a common chaos based failure mode which appears in a variety of differentiable circumstances, ranging from recurrent neural networks and numerical physics simulation to training learned optimizers. We trace this failure to the spectrum of the Jacobian of the system under study, and provide criteria for when a practitioner might expect this failure to spoil their differentiation based optimization algorithms.\"*\n\nModel efficiency\n****************\n\nEfficiently Scaling Transformer Inference\n=========================================\n\n- Code on GitHub:\n\n - `T5X `__\n - `AQT: Accurate Quantized Training `__\n\n- Research paper:\n\n - `Efficiently Scaling Transformer Inference `__ (Pope et al., 2022)\n\n - *\"We develop a simple analytical model for inference efficiency to select the best multi-dimensional partitioning techniques optimized for TPU v4 slices based on the application requirements. We combine these with a suite of low-level optimizations to achieve a new Pareto frontier on the latency and model FLOPS utilization (MFU) tradeoffs on 500B+ parameter models that outperforms the FasterTransformer suite of benchmarks. We further show that with appropriate partitioning, the lower memory requirements of multiquery attention (i.e. multiple query heads share single key/value head) enables scaling up to 32× larger context lengths.\"*\n\nNeural rendering / NeRF\n***********************\n\nGeneralizable Patch-Based Neural Rendering\n==========================================\n\n- `Code on GitHub `__\n- Research paper:\n\n - `Generalizable Patch-Based Neural Rendering `__ (Suhail et al., 2022)\n\n - *\"...We propose a different paradigm, where no deep features and no NeRF-like volume rendering are needed. Our method is capable of predicting the color of a target ray in a novel scene directly, just from a collection of patches sampled from the scene.\"*\n\nVoxel-based Radiance Fields in JAX and Flax\n===========================================\n\n- `Colab notebook `__ (Velez and Dellaert, 2022)\n\n - *\"In this notebook we show how with JAX/Flax, it is relatively easy to quickly get a voxel-based NeRF variant up and running. Specifically, we will develop a simplified version of DVGO that directly regresses color instead of having a small MLP. It works remarkably well.\"*\n\nOptimization\n************\n\nAmos Optimizer *and* JEstimator\n===============================\n\n- Code on GitHub:\n\n - `Amos and JEstimator `__\n\n - *\"... implements Amos, an optimizer compatible with the optax library, and JEstimator, a light-weight library with a tf.Estimator-like interface to manage T5X-compatible checkpoints for machine learning programs in JAX, which we use to run experiments in the paper.\"*\n\n- Research paper:\n\n - `Amos: An Adam-style Optimizer with Adaptive Weight Decay towards Model-Oriented Scale `__ (Tian and Parikh, 2022)\n\n - Presents *\"Amos, an optimizer compatible with the optax library, and JEstimator, a light-weight library with a tf.Estimator-like interface to manage T5X-compatible checkpoints for machine learning programs in JAX.\"* *\"When used for pre-training BERT variants and T5, Amos consistently converges faster than the state-of-the-art settings of AdamW, achieving better validation loss within <=70% training steps and time, while requiring <=51% memory for slot variables.\"*\n\nQuantization\n************\n\nPareto-Optimal Quantized ResNet Is Mostly 4-bit *and* AQT: Accurate Quantized Training\n======================================================================================\n\n- Code on GitHub:\n\n - `AQT: Accurate Quantized Training `__\n\n- Research paper:\n\n - `Pareto-Optimal Quantized ResNet Is Mostly 4-bit `__ (Abdolrashidi et al., 2021)\n\n - *\"In this work, we use ResNet as a case study to systematically investigate the effects of quantization on inference compute cost-quality tradeoff curves. Our results suggest that for each bfloat16 ResNet model, there are quantized models with lower cost and higher accuracy; in other words, the bfloat16 compute cost-quality tradeoff curve is Pareto-dominated by the 4-bit and 8-bit curves, with models primarily quantized to 4-bit yielding the best Pareto curve... The quantization method we used is optimized for practicality: It requires little tuning and is designed with hardware capabilities in mind... As part of this work, we contribute a quantization library written in JAX...\"*\n\nReinforcement learning\n**********************\n\nContinuous Control with Action Quantization from Demonstrations (AQuaDem)\n=========================================================================\n\n- `Code on GitHub `__\n\n- Research paper:\n\n - `Continuous Control with Action Quantization from Demonstrations `__ (Dadashi et al., 2021)\n\n - Proposes *\"a novel Reinforcement Learning (RL) framework for problems with continuous action spaces: Action Quantization from Demonstrations (AQuaDem). The proposed approach consists in learning a discretization of continuous action spaces from human demonstrations. This discretization returns a set of plausible actions (in light of the demonstrations) for each input state, thus capturing the priors of the demonstrator and their multimodal behavior. By discretizing the action space, any discrete action deep RL technique can be readily applied to the continuous control problem. Experiments show that the proposed approach outperforms state-of-the-art methods such as SAC in the RL setup, and GAIL in the Imitation Learning setup.\"*\n\nSequence models / Model parallelism\n***********************************\n\nT5X: Scaling Up Models and Data with ``t5x`` and ``seqio``\n==========================================================\n\n- `Code on GitHub `__\n\n - *\"T5X is a modular, composable, research-friendly framework for high-performance, configurable, self-service training, evaluation, and inference of sequence models (starting with language) at many scales.\"*\n\n- Research paper:\n\n - `T5X: Scaling Up Models and Data with t5x and seqio `__ (Roberts et al., 2022)\n\n - *\"Recent neural network-based language models have benefited greatly from scaling up the size of training datasets and the number of parameters in the models themselves. Scaling can be complicated due to various factors including the need to distribute computation on supercomputer clusters (e.g., TPUs), prevent bottlenecks when infeeding data, and ensure reproducible results. In this work, we present two software libraries that ease these issues: t5x simplifies the process of building and training large language models at scale while maintaining ease of use, and seqio provides a task-based API for simple creation of fast and reproducible training data and evaluation pipelines. These open-source libraries have been used to train models with hundreds of billions of parameters on datasets with multiple terabytes of training data. Along with the libraries, we release configurations and instructions for T5-like encoder-decoder models as well as GPT-like decoder-only architectures.\"*\n\nSimulation\n**********\n\nBrax - A Differentiable Physics Engine for Large Scale Rigid Body Simulation\n============================================================================\n\n- `Code on GitHub `__\n- `Colab notebooks `__\n- Research paper:\n\n - `Brax - A Differentiable Physics Engine for Large Scale Rigid Body Simulation `__ (Freeman et al., 2021)\n\n - *\"We present Brax, an open source library for rigid body simulation with a focus on performance and parallelism on accelerators, written in JAX. We present results on a suite of tasks inspired by the existing reinforcement learning literature, but remade in our engine. Additionally, we provide reimplementations of PPO, SAC, ES, and direct policy optimization in JAX that compile alongside our environments, allowing the learning algorithm and the environment processing to occur on the same device, and to scale seamlessly on accelerators.\"*\n" }, { "path": "docs/examples/index.rst", "content": "Examples\n========\n\n.. toctree::\n :maxdepth: 2\n\n core_examples\n google_research_examples\n repositories_that_use_flax\n community_examples\n\n\n" }, { "path": "docs/examples/repositories_that_use_flax.rst", "content": "Repositories that use Flax\n==========================\n\nThe following code bases use Flax and provide training frameworks and a wealth\nof examples. In many cases, you can also find pre-trained weights:\n\n\n🤗 Hugging Face\n***************\n\n`🤗 Hugging Face `__ is a\nvery popular library for building, training, and deploying state of the art\nmachine learning models.\nThese models can be applied on text, images, and audio. After organizing the\n`JAX/Flax community week `__,\nthey have now over 5,000\n`Flax/JAX models `__ in\ntheir repository.\n\n🥑 DALLE Mini\n*************\n\n`🥑 DALLE Mini `__ is a Transformer-based\ntext-to-image model implemented in JAX/Flax that follows the ideas from the\noriginal `DALLE `__ paper by OpenAI.\n\nScenic\n******\n\n`Scenic `__ is a codebase/library\nfor computer vision research and beyond. Scenic's main focus is around\nattention-based models. Scenic has been successfully used to develop\nclassification, segmentation, and detection models for multiple modalities\nincluding images, video, audio, and multimodal combinations of them.\n\nBig Vision\n**********\n\n`Big Vision `__ is a codebase\ndesigned for training large-scale vision models using Cloud TPU VMs or GPU\nmachines. It is based on Jax/Flax libraries, and uses tf.data and TensorFlow\nDatasets for scalable and reproducible input pipelines. This is the original\ncodebase of ViT, MLP-Mixer, LiT, UViM, and many more models.\n\nT5X\n***\n\n`T5X `__ is a modular, composable,\nresearch-friendly framework for high-performance, configurable, self-service\ntraining, evaluation, and inference of sequence models (starting with\nlanguage) at many scales." }, { "path": "docs/faq.rst", "content": "Frequently Asked Questions (FAQ)\n================================\n\nThis is a collection of answers to frequently asked questions (FAQ). You can contribute to the Flax FAQ by starting a new topic in `GitHub Discussions `__.\n\nWhere to search for an answer to a Flax-related question?\n*********************************************************\n\nThere are a number of official Flax resources to search for information:\n\n- `Flax Documentation on ReadTheDocs `__ (this site): Use the `search bar `__ or the table of contents on the left-hand side.\n- `google/flax GitHub Discussions `__: Search for an existing topic or start a new one. If you can't find what you're looking for, feel free to ask the Flax team or community a question.\n- `google/flax GitHub Issues `__: Use the search bar to look for an existing issue or a feature request, or start a new one.\n\nHow to take the derivative with respect to an intermediate value (using :code:`Module.perturb`)?\n************************************************************************************************\n\nTo take the derivative(s) or gradient(s) of the output with respect to a hidden/intermediate activation inside a model layer, you can use :meth:`flax.linen.Module.perturb`. You define a zero-value :class:`flax.linen.Module` \"perturbation\" parameter – :code:`perturb(...)` – in the forward pass with the same shape as the intermediate activation, define the loss function with :code:`'perturbations'` as an added standalone argument, perform a JAX derivative operation with :code:`jax.grad` on the perturbation argument.\n\nFor full examples and detailed documentation, go to:\n\n- The :meth:`flax.linen.Module.perturb` API docs\n- The `Extracting gradients of intermediate values `_ guide\n- `Flax GitHub Discussions #1152 `__\n\nIs Flax Linen :code:`remat_scan()` the same as :code:`scan(remat(...))`?\n************************************************************************\n\nFlax :code:`remat_scan()` (:meth:`flax.linen.remat_scan()`) and :code:`scan(remat(...))` (:meth:`flax.linen.scan` over :meth:`flax.linen.remat`) are not the same, and :code:`remat_scan()` is limited in cases it supports. Namely, :code:`remat_scan()` treats the inputs and outputs as carries (hidden states that are carried through the training loop). You are recommended to use :code:`scan(remat(...))`, as typically you would need the extra parameters, such as ``in_axes`` (for input array axes) or ``out_axes`` (output array axes), which :meth:`flax.linen.remat_scan` does not expose.\n\nWhat are the recommended training loop libraries?\n*************************************************\n\nConsider using CLU (Common Loop Utils) `google/CommonLoopUtils `__. To get started, go to this `CLU Synopsis Colab `__. You can find answers to common questions about CLU with Flax on `google/flax GitHub Discussions `__.\n\nCheck out the official `google/flax Examples `__ for examples of using the training loop with (CLU) metrics. For example, this is `Flax ImageNet's train.py `__.\n\nFor computer vision research, consider `google-research/scenic `__. Scenic is a set of shared light-weight libraries solving commonly encountered tasks when training large-scale vision models (with examples of several projects). Scenic is developed in JAX with Flax. To get started, go to the `README page on GitHub `__." }, { "path": "docs/flip/0000-template.md", "content": "- Start Date: (fill me in with today's date, YYYY-MM-DD)\n- FLIP PR: [#0000](https://github.com/google/flax/pull/0000)\n- FLIP Issue: [#0000](https://github.com/google/flax/issues/0000)\n\n(Below sections are just a possible structure - please adapt to your FLIP.)\n\n# Summary\n[summary]: #summary\n\nOne paragraph explanation of the FLIP.\n\n# Motivation\n[motivation]: #motivation\n\nWhy are we doing this? What use cases does it support? What is the expected outcome?\n\n# Implementation\n[implementation]: #implementation\n\nThe technical part.\n\n# Discussion\n[discussion]: #discussion\n\nSummarize the discussion from the original issue and from the pull request.\n" }, { "path": "docs/flip/1009-optimizer-api.md", "content": "- Start Date: 2021-02-08\n- FLIP PR: [#1011](https://github.com/google/flax/pull/1011)\n- FLIP Issue: [#1009](https://github.com/google/flax/issues/1009)\n\nTable of contents:\n\n- [Summary]\n- [Motivation]\n- [Using Optax]\n - [Gradient Transformations]\n - [Optax Training Step]\n - [Multi Optimizer]\n - [Train State]\n- [Previous API]\n - [Optimizer and OptimizerDef]\n - [Previous Training Step]\n- [Update Plan]\n- [Appendix]\n - [Setup Code]\n\n# Summary\n[Summary]: #summary\n\nThis FLIP proposes to replace our current `flax.optim` API (referred to as\n[previous API] in this document) with [Optax], DeepMind's optimizer library.\n\n# Motivation\n[motivation]: #motivation\n\nOur current API (referred to as [previous API] in this document) uses a pattern\nwhere an `Optimizer` dataclass is created from a pytree of `target` variables\nand from an `OptimizerDef` that defines how to update optimizer state,\nhyperparameters, and target variables. This pattern is relatively complex for\nimplementing a simple optimizer, while being quite verbose in the typical Linen\ntrain step (especially when using mutable state collections).\n\nThis package `flax.optim` contains some optimizers, but that list is far from\nexhaustive and ideally we would instead use JAX optimizers from a dedicated PyPi\npackage.\n\nDeepMind already has a dedicated library — [Optax] — that implements a wide\nrange of interesting optimizers and provides a framework to compose new\noptimizers from reusable gradient transformations.\n\n[Optax]: https://github.com/deepmind/optax\n\n# Using Optax\n[Using Optax]: #using-optax\n\n## Gradient Transformations\n[Gradient Transformations]: #gradient-transformations\n\nWhile [Optax] does provide predefined optimizers (like `optax.adam`, or\n`optax.sgd` with momentum), it is really a library of *gradient transformations*\nand the idiomatic way of instantiating an optimizer is by providing a\ncombination of these gradient transformations. To emulate the momentum\noptimizer from the example when using the [previous API] we would write:\n\n```python\nimport optax\n\ntx = optax.chain(\n optax.trace(decay=0.9, nesterov=False),\n optax.scale_by_schedule(lambda step: -get_learning_rate(step)),\n)\n```\n\nRemarks:\n\n- Above gradient transformation would be equivalent with the example under\n [Optimizer and OptimizerDef] where we define a Momentum optimizer without\n Nesterov momentum (note that the `beta` parameter corresponds to the `decay`\n parameter of the `optax.trace()` transformation, and the learning rate is\n applied in a second chained transformation).\n- Note that hyper parameters like `decay` or `nesterov` only exist in the inner\n scope of the higher order functions returning the `GradientTransformation`.\n Such a gradient transformation is currently defined as a `NamedTuple` of the\n `init()` and the `update()` function. In principle this pattern could be\n extended to also store hyperparameters, maybe a point to discuss on the\n [Optax] repo.\n- We can use a `get_learning_rate()` that returns the learning rate depending on\n the step number when defining the Optax gradient update transformation. Above\n code illustrates how this can be a drop-in replacement for a function we also\n use in our [previous training step], where this update function already exists\n (notice how we need to invert the sign because we add the gradient update to\n the parameters). In addition, you can use\n [`inject_hyperparams()`](https://github.com/deepmind/optax/pull/48) to\n schedule arbitrary hyper parameters with Optax.\n\n## Optax Training Step\n[Optax Training Step]: #optax-training-step\n\n```python\n@functools.partial(jax.jit, static_argnums=(4, 5))\ndef train_step(opt_state, variables, inputs, labels, apply_fn, tx_update_fn):\n\n def loss_fn(params):\n logits, new_model_state = apply_fn(\n {**variables, 'params': params}, inputs, mutable=['batch_stats'])\n loss = xent_loss(logits, labels)\n return loss, new_model_state\n\n variables, params = variables.pop('params')\n (loss, new_model_state), grads = jax.value_and_grad(loss_fn, has_aux=True)(\n params)\n updates, new_opt_state = tx_update_fn(grads, opt_state, params)\n new_params = optax.apply_updates(params, updates)\n new_variables = {**variables, **new_model_state, 'params': new_params}\n return new_opt_state, new_variables, loss\n\n\nopt_state = tx.init(variables['params'])\nfor batch in ds.as_numpy_iterator():\n opt_state, variables, loss = train_step(\n opt_state, variables, batch['image'], batch['label'], model.apply,\n tx.update)\n print(loss)\n```\n\nRemarks:\n\n- Since `tx.update()` only transforms the gradient, we still need to call\n `optax.apply_updates()` to apply these transformed gradients to the\n parameters.\n- Compared with the [previous API], we can now keep the entire `variables`\n including the `params` as an input and output to the `train_step()`.\n- Splitting `params` from `variables` is still necessary inside the train step\n because we only want to compute gradients with respect to `params` and not the\n entire `variables`.\n- We can still log internal optimizer state, such as the learning rate, as long\n as Optax transformations expose that information in their respective state.\n For example, `optax.scale_by_schedule()` currently only exposes\n `opt_state.count` but could easily be extend to also expose the `step_size`.\n The same is true for internal optimizer states that change over time.\n\n## Multi Optimizer\n[Multi Optimizer]: #multi-optimizer\n\nThe [previous API] defined `flax.optim.MultiOptimizer` for processing different\nparts of the parameter tree with different optimizers:\n\n```python\nbiases_traversal = flax.optim.ModelParamTraversal(\n lambda path, _: path.endswith('/bias'))\nnot_biases_traversal = flax.optim.ModelParamTraversal(\n lambda path, _: not path.endswith('/bias'))\n\noptimizer_def = flax.optim.MultiOptimizer(\n (biases_traversal, flax.optim.GradientDescent(learning_rate=0.1)),\n (not_biases_traversal, flax.optim.GradientDescent(learning_rate=0.05)),\n)\n```\n\nNote how we first define a traversal that selects parameters based on their\npath (which is the concatenation of module scopes and variable name), and then\ncreate a `MultiOptimizer` that binds a different optimizer for each of these\nseparate traversals.\n\nOptax has recently implemented `optax.masked()` that can be used for specifying\ngradient transformations that only applied to a subset of the gradients:\n\n```python\ndef flattened_traversal(fn):\n def mask(data):\n flat = traverse_util.flatten_dict(data)\n return traverse_util.unflatten_dict({k: fn(k, v) for k, v in flat.items()})\n return mask\n\ntx = optax.chain(\n optax.masked(optax.sgd(learning_rate=0.1),\n mask=flattened_traversal(lambda path, _: path[-1] == 'bias')),\n optax.masked(optax.sgd(learning_rate=0.05),\n mask=flattened_traversal(lambda path, _: path[-1] != 'bias')),\n)\n```\n\n## Train State\n[Train State]: #train-state\n\nIn Flax it is common to hand around a `TrainState` object that can then be\nused for checkpointing. This simplifies the above [Optax training step] a bit by\nreducing the number of arguments and getting rid of the `static_argnums`.\n\nWe can define a `TrainState` dataclass that wraps the common pattern of updating\nthe optimizer state and parameters by applying the gradients.\n\n```python\n# Small helper class in flax.training\nclass TrainState(flax.struct.PyTreeNode):\n step: int\n apply_fn: Callable = flax.struct.field(pytree_node=False)\n params: flax.core.FrozenDict[str, Any]\n tx: optax.GradientTransformation = flax.struct.field(pytree_node=False)\n opt_state: optax.OptState\n\n def apply_gradients(self, *, grads, **kwargs):\n updates, new_opt_state = self.tx.update(\n grads, self.opt_state, self.params)\n new_params = optax.apply_updates(self.params, updates)\n return self.replace(\n step=self.step + 1,\n params=new_params,\n opt_state=new_opt_state,\n **kwargs,\n )\n\n @classmethod\n def create(cls, *, apply_fn, params, tx, **kwargs):\n opt_state = tx.init(params)\n return cls(\n step=0,\n apply_fn=apply_fn,\n params=params,\n tx=tx,\n opt_state=opt_state,\n **kwargs,\n )\n```\n\nUsers can then derive from this dataclass and add more fields, for example\nmutable model state:\n\n```python\nfrom flax.training import train_state\n\nclass TrainState(train_state.TrainState):\n batch_stats: flax.core.FrozenDict[str, Any]\n```\n\nWith this the [Optax Training Step] becomes:\n\n```python\n@jax.jit\ndef train_step(state, inputs, labels):\n\n def loss_fn(params):\n outputs, new_model_state = state.apply_fn(\n {'params': params, 'batch_stats': state.batch_stats},\n inputs,\n mutable=['batch_stats'])\n loss = xent_loss(outputs, labels)\n return loss, new_model_state\n\n (loss, new_model_state), grads = jax.value_and_grad(\n loss_fn, has_aux=True)(state.params)\n new_state = state.apply_gradients(\n grads=grads,\n batch_stats=new_model_state['batch_stats'],\n )\n\n return new_state, loss\n\n\nstate = TrainState.create(\n apply_fn=model.apply,\n params=variables['params'],\n tx=tx,\n batch_stats=variables['batch_stats'],\n)\nfor batch in ds.as_numpy_iterator():\n state, loss = train_step(state, batch['image'], batch['label'])\n```\n\nThe train step without mutable state reduces to:\n\n```python\n@jax.jit\ndef train_step(state, inputs, labels):\n\n def loss_fn(params):\n outputs = state.apply_fn({'params': params}, inputs)\n loss = xent_loss(outputs, labels)\n return loss\n\n loss, grads = jax.value_and_grad(loss_fn)(state.params)\n new_state = state.update(grads=grads)\n\n return new_state, loss\n\n\nstate = flax.training.TrainState.create(\n apply_fn=model.apply,\n params=variables['params'],\n tx=tx,\n)\nfor batch in ds.as_numpy_iterator():\n state, loss = train_step(state, batch['image'], batch['label'])\n```\n\nRemarks:\n\n- It is a common pattern in Flax training loops to have a `TrainState` dataclass\n that is updated with new state after every step.\n- The simple solution proposed in `flax.training.train_state` an be extended\n with additional data, but advanced usecases (e.g. multiple different models\n and/or optimizers) are not supported. Users should instead fork the dataclass\n and re-implement it to their needs.\n- As opposed to the `Optimizer` abstraction in the [previous API], the\n `TrainState` now directly contains the `.params`, without having to to through\n `.optimizer`\n\n# Previous API\n[previous API]: #previous-api\n\n## Optimizer and OptimizerDef\n[Optimizer and OptimizerDef]: #optimizer-and-optimizerdef\n\nThe optimizer itself would be implemented by creating a new class derived\nfrom `OpimizerDef`:\n\n```python\n# flax/optim/momentum.py\n\n@flax.struct.dataclass\nclass _MomentumHyperParams:\n learning_rate: jnp.ndarray\n beta: jnp.ndarray\n\n\n@flax.struct.dataclass\nclass _MomentumParamState:\n momentum: np.ndarray\n\n\nclass Momentum(flax.optim.OptimizerDef):\n\n def __init__(self, learning_rate=None, beta=0.9):\n super().__init__(\n _MomentumHyperParams(learning_rate, beta)\n )\n\n def init_param_state(self, param):\n return _MomentumParamState(jnp.zeros_like(param))\n\n def apply_param_gradient(self, step, hyper_params, param, state, grad):\n del step\n assert hyper_params.learning_rate is not None\n new_momentum = state.momentum * hyper_params.beta + grad\n new_params = param - hyper_params.learning_rate * new_momentum\n return new_params, _MomentumParamState(new_momentum)\n```\n\nRemarks:\n\n- Note the relationship between `OptimizerDef` and `Optimizer` : When the\n function `Optimizer.apply_gradient()` is called from the user code, it calls\n into `OptimizerDef.apply_gradient()` (among other things) which in turn will\n call `OptimizerDef.apply_param_gradient()` (implemented by subclasses of\n `OptimizerDef`).\n- The functions `init_param_state()` and `apply_param_gradient()` are called\n for every leaf in the params/grads pytree. This makes it possible to write the\n calculations directly without `jax.tree_util.tree_map()`.\n- The interface was defined in pre-Linen without the distinction of `params` vs.\n other collections in `variables` in mind. The original API was elegant because\n one only needed to pass around the optimizer, which included the parameters,\n optimizer state, optimizer hyperparameters, and a reference to the\n `OptimizerDef` to perform the param/state update.\n\n## Previous Training Step\n[Previous Training Step]: #previous-training-step\n\nAn optimizer would first be constructed from its definition and the pytree of\ntarget params:\n\n```python\noptimizer_def = flax.optim.Momentum(learning_rate=0.1, beta=0.9)\noptimizer = optimizer_def.create(variables['params'])\n```\n\nThen, the target variables would optimized in the train step (assuming a single\nnon-params collection \"batch_stats\"):\n\n```python\ndef make_train_step(apply_fn):\n @jax.jit\n def train_step(optimizer, batch_stats, inputs, labels):\n\n def loss_fn(params):\n variables = {'params': params, 'batch_stats': batch_stats}\n logits, new_model_state = apply_fn(\n variables, inputs, mutable=['batch_stats'])\n loss = xent_loss(logits, labels)\n return loss, new_model_state['batch_stats']\n\n (loss, new_batch_stats), grad = jax.value_and_grad(loss_fn, has_aux=True)(\n optimizer.target)\n lr = get_learning_rate(step)\n new_optimizer = optimizer.apply_gradient(grad, learning_rate=lr)\n return new_optimizer, new_batch_stats, loss\n\n return train_step\n\n\nbatch_stats = variables['batch_stats']\ntrain_step = make_train_step(model.apply)\nfor step, batch in enumerate(ds)\n optimizer, batch_stats, loss = train_step(\n optimizer, batch_stats, batch['image'], batch['label'])\n```\n\nRemarks:\n\n- Notice how `optimizer.apply_gradient()` can take additional arguments to\n update hyperparameters, such as learning rate from an independent function\n `get_learning_rate()` in this case.\n\n\n# Update Plan\n[Update Plan]: #update-plan\n\n1. Finalize discussions on this FLIP\n2. Add [equivalence tests] to Optax that guarantee that existing `flax.optim`\n optimizers return identical values with corresponding `optax` optimizers.\n3. Update examples to use Optax and verify that they reach the same final\n performance with the same computational cost.\n4. Port missing optimizers to Optax (e.g. Adafactor) - and verify above points.\n5. Update all documentation (including README, Flax guided tour, HOWTOs, ...) to\n talk exclusively about Optax optimizers.\n6. Create a transition guide for updating users from `flax.optim` to using\n Optax. This transition guide should also point to Optax's [equivalence tests]\n and the pull requests updating the examples.\n7. Mark optimizers in `flax.optim` as deprecated.\n\n[equivalence tests]: https://github.com/deepmind/optax/blob/master/optax/_src/equivalence_test.py\n\nNote that all current Flax examples use an optimizer that is already available\nin Optax:\n\n| Example | Flax | Optax | Comments |\n| -------- | -------------- | ----------- | ----------------------------------- |\n| imagenet | optim.Momentum | optax.sgd | DynamicScale can be used unchanged. |\n| mnist | optim.Momentum | optax.sgd | |\n| nlp_seq | optim.Adam | optax.adamw | |\n| pixelcnn | optim.Adam | optax.adam | |\n| ppo | optim.Adam | optax.adam | |\n| seq2seq | optim.Adam | optax.adam | |\n| vae | optim.Adam | optax.adam | |\n| wmt | optim.Adam | optax.adamw | |\n\n(Flax's Adam implementation has an optional parameter for weight decay, but in\nOptax Adam with and without weight decay are two different aliases.)\n\n# Appendix\n[Appendix]: #appendix\n\n## Setup Code\n[Setup Code]: #setup-code\n\nThe following setup code can be used for running the code snippets in this\nFLIP:\n\n```python\nimport functools\nfrom typing import Callable, Sequence\n\nimport jax\nimport jax.numpy as jnp\nimport flax\nimport flax.linen as nn\nimport tensorflow as tf\nimport tensorflow_datasets as tfds\n\n\ndef pp(features):\n return {\n 'image': tf.cast(features['image'], tf.float32) / 255 - 0.5,\n 'label': features['label'],\n }\n\n\nclass Model(nn.Module):\n\n @nn.compact\n def __call__(self, inputs):\n x = inputs.reshape([inputs.shape[0], -1])\n x = nn.normalization.BatchNorm(True)(x)\n x = nn.Dense(10)(x)\n x = nn.log_softmax(x)\n return x\n\n\ndef onehot(labels, num_classes, on_value=1.0, off_value=0.0):\n x = (labels[..., None] == jnp.arange(num_classes)[None])\n x = jax.lax.select(\n x, jnp.full(x.shape, on_value), jnp.full(x.shape, off_value))\n return x.astype(jnp.float32)\n\n\ndef xent_loss(logits, labels):\n return -jnp.sum(\n onehot(labels, num_classes=10) * logits) / labels.size\n\n\ndef get_learning_rate(step):\n return 0.1\n\n\nmodel = Model()\nrng = jax.random.key(0)\nds = tfds.load('mnist')['train'].take(160).map(pp).batch(16)\nbatch = next(iter(ds))\nvariables = model.init(rng, jnp.array(batch['image'][:1]))\njax.tree_util.tree_map(jnp.shape, variables)\n```\n" }, { "path": "docs/flip/1777-default-dtype.md", "content": "# FLIP: Default dtypes\n\n\n- Start Date: 2022-01-11\n- FLIP PR: [#1776](https://github.com/google/flax/pull/1776)\n- FLIP Issue: [#1777](https://github.com/google/flax/issues/1777)\n- Status: Implemented\n\n\n## Summary\n\nThis FLIP proposes to replace the default dtype which is currently fixed to float32, and instead use the JAX type promotion results to derive a default dtype from the input and parameters of a layer.\n\n\n## Motivation\n\nCurrently, Linen Modules always produce `module.dtype` (defaults to float32) outputs regardless of input and parameter dtypes. Half-precision types like float16 and bfloat16 are supported by explicitly passing the half-precision type to each Module. The way this is currently implemented is that each Module has a dtype argument with float32 as the default value. The layer guarantees that this dtype will be the return type of the result returned by `__call__`.\n\nThe current behavior is problematic and results in silent bugs, especially for dtypes that do not fit inside float32 (complex, float64). Also, the Linen dtype behavior is significantly different from how NumPy and by extension JAX handle dtypes.\n\n\n### Dtypes in JAX\n\nJAX uses a NumPy-inspired [dtype promotion](https://github.com/jax-ml/jax/blob/main/jax/_src/dtypes.py) mechanism as explained [here](https://jax.readthedocs.io/en/latest/type_promotion.html?highlight=lattice#type-promotion-semantics). The type promotion rules are summarized by the following type lattice:\n\n![JAX type promotion lattice](https://jax.readthedocs.io/en/latest/_images/type_lattice.svg)\n\n\n## Dtypes in Linen\n\nBesides input arguments, state and in particular parameters could affect dtype promotion. For example: we might feed a float64 input to a Dense layer with float32 parameters. Currently, the result would be truncated to float32. If the input is a complex number the result is even worse because the imaginary part will be silently dropped when casting to float32.\n\nBy using the dtype promotion rules already available in JAX we can avoid this issue. A public API is available called `jax.numpy.result_dtype(*args)`, which returns the dtype that JAX would promote the given arguments to, in accordance with the type promotion lattice. For Linen layers the arguments would be the layer inputs together with the parameters. For example, for a linear layer this would be inputs, kernel, and bias.\n\nNote that there is also a `param_dtype` attribute in standard Linen Modules that also defaults to flaot32. This behavior is left untouched and encodes the common case of having float32 parameters.\nThere are a few reasons why float32 is almost always the correct dtype for parameters:\n1. Storing weights in half-precision often leads to underflow during optimization.\n2. Double precision is rarely used because it severely slows down modern accelerators (GPU, TPU). Therefore, such a cost should be explicitly opted-in for.\n3. Complex Modules are relatively uncommon. Even within complex networks, the complex inputs can be projected with a real matrix.\n\n\n# Implementation\n\nA simplified example implementation:\n\n\n```python\ndef promote_arrays(*xs, dtype):\n if dtype is None:\n dtype = jnp.result_type(*jax.tree_util.tree_leaves(xs))\n return jax.tree_util.tree_map(lambda x: jnp.asarray(x, dtype), xs)\n\nDtype = Any\nclass Dense(nn.Module):\n features: int\n kernel_init: Callable\n bias_init: Callable\n dtype: Optional[Dtype] = None\n param_dtype: Dtype = jnp.float32\n\n @nn.compact\n def __call__(self, x):\n kernel = self.param(\"kernel\",\n self.kernel_init,\n (x.shape[-1], self.features), self.param_dtype)\n bias = self.param(\"bias\", self.bias_init, (self.features,), self.param_dtype)\n x, kernel, bias = promote_arrays(x, kernel, bias, dtype=self.dtype)\n return x @ kernel + bias\n```\n\n\n## Half-precision dtypes\n\nSome layers don’t work with half-precision dtypes internally. For example: The normalization layers currently compute mean and variance in float32 even when a half-precision dtype is specified to avoid numerical issues. We can replicate this behavior by calling result_dtype with a dummy argument that has the minimum precision for the sub computation to work correctly.\n\n\n## Backward compatibility\n\nThis proposal causes some layers to behave differently in cases where the dtype is not specified to a Linen Module. By default, parameters are in float32. Therefore, passing in half or float32 precision inputs will cause a float32 dtype and no functional differences with current behavior.\n\nWhen passing complex or float64 precision, the result will no longer truncate the imaginary component or the precision. The silent truncation is problematic and has caused [user complaints](https://github.com/google/flax/issues/805#issuecomment-981468837). Therefore, this change can be considered a bugfix.\n\nThus, although this proposal strictly speaking changes behavior it is unlikely to cause problems for users. There are 2 exceptions to this which should be rare and easy to fix:\n1. A user relies on the enforced float32 to downcast a double precision value.\n2. A user relies on the float32 to explicitly upcast a half precision value even though the weights are in half precision.\n\n\n## Corner cases\n\nIn this section we describe corner cases where the implementation of the proposal is not obvious. The two main concerns are how complex numbers are handled in existing layers and how to determine the dtype of state variables.\n\n**Autoregressive decoding cache**\n\nCurrently, only attention implements autoregressive caching and the stored key and value mirror the dtype of the key and value passed to the layer. Forcing the cache dtype to be the same as the output dtype could result in reduced precision during cached decoding vs uncached. This seems undesirable. Decision: keep the current behavior.\n\n**Batch statistics**\n\nBatchNorm layers are often used with a half precision output dtype. However, calculating statistics is by default always done in float32 to avoid numerical precision issues and over/underflow for float16. With float64 this would actually cause a downcast so we should now use `np.promote_types(float32, dtype)` such that the precision is at least float32. The running batch statistics will be stored with the same dtype for consistency.\n\n**Complex number support**\n\nCurrently, our complex number support is brittle because the default behavior is to truncate the output to the real part. This issue will be fixed by the automatic type promotion proposed in this FLIP. However, some layers require some additional thought to extend to complex numbers correctly:\n\n1. Normalization layers use the complex conjugate to calculate norms instead of normal squaring.\n2. Attention: It’s not exactly clear how the dot product and softmax are defined in this case. Raise an error on complex inputs.\n3. Recurrent layers: might require special gating / activation functions to function correctly, but these can be specified by the user.\n\n\n# Discussion\n\nSummarizing the main points from the discussion:\n\n\n## Consider implicit complex truncation an error\n\nQ:\nI'm wondering if we should always raise an error if one of the xs tree leaves is complex but dtype is not. Users should maybe remove imaginary part by themselves if that's really what they want to do.\n(Maybe it's a contrived example, but I can imagine cases where layers have their dtype set by parent modules based on assumptions without complex numbers in mind)\n\nA:\nThis is worth considering in a follow-up CL but this might as well be solved in JAX directly where the safeguard would apply more generally. In NumPy this was also considered but abandoned because it is not backwards compatible.\n\n\n## Dtype attribute names\n\nQ:\nAre the dtype and param_dtype arguments confusing? In particular, should dtype perhaps be called output_dtype to make the difference between the two dtypes more explicit?\n\nA:\nThis would be a large and orthogonal change wrt to this proposal so leaving it out for now.\nAlso, this breaks with the standard dtype argument in NumPY/JAX.\nAlthough dtype indeed constrains the output dtype it is also a hint for the dtype we would like the computation to happen in.\n\n" }, { "path": "docs/flip/2396-rnn.md", "content": "# RNN Flip\n\n- Start Date: 2022-08-18\n- FLIP PR: [#2604](https://github.com/google/flax/pull/2604)\n- FLIP Issue: [#2396](https://github.com/google/flax/issues/2396)\n- Authors: Jasmijn Bastings (@bastings) and Cristian Garcia (@cgarciae)\n\n## Summary\nThis FLIP adds support for higher-level recurrent layers (RNN, GRU, LSTM) that can help users process input sequences using the recurrent cells already available in Flax.\n\n## Motivation\nImplementing well known recurrent architectures is tricky and prone to user errors, even a simple LSTM layers involves the manual creation and handling of the carry/memory and correctly setting up `nn.scan`:\n\n```python\n@nn.compact\ndef __call__(self, x):\n LSTM = nn.scan(\n nn.LSTMCell, variable_broadcast=\"params\", split_rngs={\"params\": False}\n )\n carry = LSTM.initialize_carry(\n jax.random.key(0), batch_dims=x.shape[:1], size=self.hidden_size\n )\n carry, x = LSTM()(carry, x)\n return x\n```\nSlightly more complicated cases involving padding like in the [seq2seq](https://github.com/google/flax/blob/main/examples/seq2seq/models.py) example require even more work but couple potentially be simplified to a couple of lines with the right abstractions. We propose providing users with clean, correct, and efficient abstractions to use recurrent cells. \n\n## Requirements\n\n* **Masking**: We need to support a batch of sequences that contain padding at the end of each sequence. \n * We do not intend to support non-contiguous padding, i.e. padding that is not at the end of a sequence, for performance reasons, except in the case of packing (see below).\n* **Bidirectionality**: The ability to process a sequence in both the forward and reverse directions, respecting padding (i.e., the reverse direction should start with the actual inputs, not with padding values).\n* **Performance**: The proposed classes should be benchmarked to provide the best performance in terms of step time and/or memory use.\n* **Recurrent Dropout**: Support for recurrent dropout in cells (e.g. dropout on the state of the cell).\n\n## Implementation\n### High-level structure\n\nWe propose to have these 3 levels of abstraction:\n\n* **Cells (unchanged)**: all RNNCellBase subclasses such as LSTMCell and GRUCell, these implement the stepwise logic. These already exist in Flax today.\n* **Layers (new)**: a class (RNN) that takes a cell and scans over a sequence respecting possible padding values and optionally also allows packed sequences.\n* **Bidirectional (new)**: a single class that takes a forward and a backward RNN instance and correctly processes the input sequence in both directions and merges the results.\n\n### Example of proposed API\nWe start with a code example of what you could do with the proposed API, and then we discuss the API in detail below.\n\n```python\ncell = nn.LSTMCell()\n# Encodes a batch of input sequences.\ncarry, outputs = nn.RNN(cell, cell_size)(inputs, seq_lengths)\n```\n\nA Bidirectional layer with a LSTM RNNs for the forward and backward directions respectively would look like this:\n\n```python\nforward_rnn = nn.RNN(nn.LSTMCell(), cell_size=32)\nbackward_rnn = nn.RNN(nn.LSTMCell(), cell_size=32)\n# Bidirectional combinator.\nbi_rnn = nn.Bidirectional(forward_rnn, backward_rnn)\n# Encodes a batch of input sequences in both directions.\ncarry, outputs = bi_rnn(inputs, seq_lengths) \n```\n\nNext we will discuss `RNN`, `Bidirectional`, and proposed changes to `RNNCellBase`.\n\n### RNNBase\nThe `RNNBase` class serves as a base class for the `RNN` class, it specifies\nthe API that all RNN layers should implement to be compatible with the `Bidirectional`.\n`RNNBase` contains the `__call__` and `flip_sequences` methods:\n\n```python\nclass RNNBase(Protocol):\n def __call__(\n self,\n inputs: jax.Array,\n *,\n initial_carry: Optional[Carry] = None,\n init_key: Optional[random.KeyArray] = None,\n seq_lengths: Optional[Array] = None,\n return_carry: Optional[bool] = None,\n time_major: Optional[bool] = None,\n reverse: Optional[bool] = None,\n keep_order: Optional[bool] = None,\n ) -> Union[Output, Tuple[Carry, Output]]:\n ...\n```\nWhere:\n\n* `inputs`: the input sequence.\n* `initial_carry`: the initial carry, if not provided it will be initialized\n using the cell's :meth:`RNNCellBase.initialize_carry` method.\n* `init_key`: a PRNG key used to initialize the carry, if not provided\n ``jax.random.key(0)`` will be used. Most cells will ignore this\n argument.\n* `seq_lengths`: an optional integer array of shape ``(*batch)`` indicating\n the length of each sequence, elements whose index in the time dimension\n is greater than the corresponding length will be considered padding and\n will be ignored.\n* `return_carry`: if ``return_carry=False`` (default) only the output sequence is returned,\n else it will return a tuple of the final carry and the output sequence.\n* `time_major`: if ``time_major=False`` (default) it will expect inputs with shape\n ``(*batch, time, *features)``, else it will expect inputs with shape ``(time, *batch, *features)``.\n* `reverse`: if ``reverse=False`` (default) the sequence is\n processed from left to right and returned in the original order, else it will be processed\n from right to left, and returned in reverse order. If ``seq_lengths`` is passed,\n padding will always remain at the end of the sequence.\n* `keep_order`: if ``keep_order=True``, when ``reverse=True``\n the output will be reversed back to the original order after processing, this is\n useful to align sequences in bidirectional RNNs. If ``keep_order=False`` (default),\n the output will remain in the order specified by ``reverse``.\n* `Returns`: if ``return_carry=False`` (default) only the output sequence is returned,\nelse it will return a tuple of the final carry and the output sequence.\n\n### RNN\nThe `RNN` module inherits from `RNNBase`, it main function is to apply an `RNNCellBase` instance over a batch of input sequences, it can be used with any type of cell (e.g., `GRUCell`, `LSTMCell`, etc). It accepts the following parameters:\n\n```python\nclass RNN(RNNBase):\n cell: RNNCellBase,\n cell_size: int | Tuple[int, ...]\n time_axis: int = -2,\n variable_axes = FrozenDict(),\n variable_broadcast: CollectionFilter = 'params'\n variable_carry: CollectionFilter = False\n split_rngs = FrozenDict({'params': False})\n # implement RNNBase\n ...\n```\n\nAttributes like `variable_axes`, `variable_broadcast`, `variable_carry`, and `split_rngs` are directly passed to `nn.scan`, their default values are set such that common cells like `LSTMCell` and `GRUCell` work out of the box.\n\n### Masking\n`seq_lengths` is defined as an integer array of shape `(*batch,)` indicating the length of each sequence.\n\n
Discussion\n\nThere are various masking formats found in other frameworks, here are some of the most popular ones:\n\n* **Binary masking**: specifies per-sample and timestep whether that data point should be included or not in the computation, it can be non-contigous (e.g., [1, 1, 0, 1]). This is used by Keras.\n* **Sequence length masking**: specifies per-sample the number of non-padding examples contained in the sequence, any padding contained in the sequence should be stacked at the end. This is used by FlaxFormer.\n* **Segmentation Mask**: specifies row and timestep to which sample the data point belongs to, this format allows more than one sample per row which potentially reduces the total amount of padding needed (e.g. [1, 1, 1, 2, 2, 0, 0]). Pytorch uses this representation (see [pack_padded_sequence](https://pytorch.org/docs/stable/generated/torch.nn.utils.rnn.pack_padded_sequence.html)).\n\nWhile Sequence packing (see [LM1B example](https://github.com/google/flax/blob/main/examples/lm1b/input_pipeline.py#L90-L92)) is is more powerful, its implementation is more complex and it is not clear whether it is worth the effort. The simplest format is sequence length masking, which is the one we propose to use.\n\n
\n\n### Bidirectional\nBidirectional processing can be achieved via a Module that accepts a `forward_rnn` Module and a `backward_rnn` Module, both of which should be `RNN` instances, in order to process the input sequence in both directions. Here we present some pseudo code of the implementation:\n\n```python\ndef __call__(self, inputs, seq_lengths):\n # Encode in the forward direction.\n carry_forward, outputs_forward = self.forward_rnn(\n inputs, seq_lengths=seq_lengths, \n return_carry=True, reverse=False,\n )\n # Encode in the reverse order.\n carry_backward, outputs_backward = self.backward_rnn(\n inputs, seq_lengths=seq_lengths,\n return_carry=True, reverse=True, # process in reverse order\n keep_order=True, # but return the sequence in the original order\n )\n # Merge both sequences.\n outputs = jax.tree.map(self.merge_fn, outputs_forward, outputs_backward)\n\n return (carry_forward, carry_backward), outputs\n```\n\nHere `merge_fn` a function that takes both outputs and fuses them (`concat` by default). As showcased in the beginning of this document, usage would look like this:\n\n```python\nforward_rnn = nn.RNN(nn.LSTMCell(), cell_size=32)\nbackward_rnn = nn.RNN(nn.GRUCell(), cell_size=32)\n# Bidirectional combinator.\nbi_rnn = nn.Bidirectional(forward_rnn, backward_rnn)\n# Encodes a batch of input sequences in both directions.\ncarry, outputs = bi_rnn(inputs, seq_lengths) \n```\n\n### Recurrent Dropout\nThere are two main uses of dropout in RNNs:\n1. Input dropout: regular dropout applied to the inputs, different for every step.\n4. Recurrent dropout: applies dropout to a recurrent input/output, same for every step.\n\nFlax's `nn.scan` can easily express both types of dropout via `split_rns`, input dropout would split rngs while recurrent dropout would not. [#2540](https://github.com/google/flax/pull/2540) was introduces such that the `rng_name` in `nn.Dropout` can now be defined by the user, this way Cells could define both types of dropout e.g:\n\n```python\nself.dropout = nn.Dropout(...) # input dropout\nself.recurrent_dropout = nn.Dropout(..., rng_collection='recurrent_dropout')\n```\nBased on this, `nn.scan` / `nn.RNN` can now specify `split_rngs` accordingly e.g:\n```\nnn.scan(scan_fn, ..., split_rngs={'dropout': True, 'recurrent_dropout': False})\n```\n\n# Future ideas\n\n
show\n\n### Sequence Packing\nAllow packing multiple sequences to make efficient use of space/memory. This might result in a trade-off where step time is higher (because at each step we need to check whether we are starting a new sequence and reset the carry/initial state), but where less padding is used increasing efficiency overall.\n\n### RNNCell redesign\n\n#### Make initialize_state an instance method\nFirst altenative is to make `initialize_carry` a instance method. With this change hyperparameters can be passed directly to the cell, it signature would look like this:\n\n```python\ndef initialize_carry(self, sample_input) -> Carry:\n ...\n```\n\nUsage would look like this:\n\n```python\nLSTM = nn.scan(\n nn.LSTMCell, variable_broadcast='params', \n split_rngs={'dropout': True})\nlstm = LSTM(features=32)\ncarry = lstm.initialize_carry(x[:, 0])\ncarry, y = lstm(carry, x)\n```\n\n#### Remove initialize_carry\n\nAn alternative is to remove `initialize_carry` entirely and have the carry state be handled as a carry collection. This would simplify usage quite a bit:\n\n```python\nLSTM = nn.scan(\n nn.LSTMCell, variable_broadcast='params', \n split_rngs={'dropout': True})\ny = LSTM(features=32)(carry, x)\n```\n\nHowever, this would require `nn.scan` to support initialization of carry collections which is currently not possible. Also, users would have to specify that a collection is mutable e.g. `mutable=['carry']`, even if they are not interested in the output carry state.\n\n
\n" }, { "path": "docs/flip/2434-general-metadata.md", "content": "# FLIP: Axis Metadata\n\n\n- Start Date: 2022-08-08\n- FLIP Issue: [#2434](https://github.com/google/flax/issues/2434)\n- FLIP PR: [#2435](https://github.com/google/flax/pull/2435)\n- Status: Proposal\n\n\n## Summary\n\nThis FLIP proposes to extend Flax's variable collections with a generic axis metadata API.\nThe core of the API is an abstract base class that is recognized by lifting transformations that can add an axis (vmap, scan).\nUsers can extend the base class to keep track of per-axis metadata in a way that works with lifted transformations.\n\n\n## Motivation\n\nGenerally, there is no way in Flax to track metadata for variables across lifted transformations.\nAxis metadata is used to keep track of semantic information about axes into other (Flax independent) APIs.\nFor example, optimizers like AdaFactor can be configured on a per-axis level and partitioning APIs\nin JAX like xmap or pjit require per variable annotations to map effectiently to parallel hardware.\n\nCurrently, there is an experimental [API](https://github.com/google/flax/blob/main/flax/linen/partitioning.py)\nsupporting partitioning annotations with wrappers around lifted transforms that change axes (``nn.scan_with_axes``, ``nn.vmap_with_axes``)\nand a special APIs to create variables (``param_with_axes`` and ``variable_with_axes``).\nThe experimental partitioning API stores the metadata in a separate collection named \"[collection]_axes\".\n\n\nThe experimental API has a number of shortcomings that we like to solve:\n1. The current API works for tracking PartitionSpecs but not for other types of metadata like optimizer annotations.\n2. The implementation using an \"xxx_axes\" collection requires error-prone and non-composable string manipulation.\n3. Special, partioning-aware variable creators and lifted transforms are required\n4. The partioning API is hard to use with pre-existing Modules that aren't partioning aware.\n\n\n## Proposal\n\nTo generalize metadata tracking and keep the specific metadata out of core Flax we propose the following abstract base class:\n\n```python\nTAxisMetadata = TypeVar(\"TAxisMetadata\", bound=\"AxisMetadata\")\n\nclass AxisMetadata(metaclass=abc.ABCMeta):\n \"\"\"Abstract base class for boxed Metadata.\n\n ``AxisMetadata`` enables arbitrary, per axis metadata for variables.\n By using ``unbox`` the metadata is stripped away to obtain the original\n variables. By using unboxing, most code handling variables does not need\n to handle ``AxisMetadata`` specifically, but can directly operate on the JAX\n arrays that they wrap.\n\n Additionally, ``AxisMetadata`` supports updating metadata whenever an axis\n is added or removed by a functional transformation\n (e.g.: ``nn.scan`` or ``nn.vmap``) using the ``add_axis`` and ``remove_axis``\n methods.\n\n By extending ``AxisMetadata``, custom metadata can be stored. See\n ``Partitioned`` for a specific implementation.\n \"\"\"\n\n @abc.abstractmethod\n def unbox(self) -> Any:\n \"\"\"Returns the content of the AxisMetadata box.\n\n Note that unlike ``meta.unbox`` the unbox call should recursively unbox\n metadata. It should simply return value that it wraps directly even\n if that value itself is an instance of AxisMetadata.\n\n In practise, AxisMetadata subclasses should be registred as PyTree nodes to\n support passing instances to JAX and Flax APIs. The leaves returned for this\n note should correspond to the value returned by unbox.\n\n Returns:\n The unboxed value.\n \"\"\"\n pass\n\n @abc.abstractmethod\n def add_axis(self: TAxisMetadata, index: int,\n params: Dict[Any, Any]) -> TAxisMetadata:\n \"\"\"Adds a new axis to the axis metadata.\n\n Note that add_axis and remove_axis should act as each other's inverse\n (meaning: ``x.add_axis(i, p).remove_axis(i, p) == x``)\n\n Args:\n index: The position at which the new axis will be inserted\n params: An arbitrary dictionary of parameters passed by the transformation\n that introduces the new axis (e.g.: ``nn.scan`` or ``nn.vmap``). The\n user passes this dictionary as the `metadata_param` argument to the\n transformation.\n Returns:\n A new instance of the same type as self and with the same ``unbox``\n content with updated axis metadata.\n \"\"\"\n pass\n\n @abc.abstractmethod\n def remove_axis(self: TAxisMetadata, index: int,\n params: Dict[Any, Any]) -> TAxisMetadata:\n \"\"\"Removes an axis from the axis metadata.\n\n Note that add_axis and remove_axis should act as each other's inverse\n (meaning: ``x.remove_axis(i, p).add_axis(i, p) == x``)\n\n Args:\n index: The position of the axis that is to be removed\n params: An arbitrary dictionary of parameters passed by the transformation\n that introduced the axis (e.g.: ``nn.scan`` or ``nn.vmap``). The\n user passes this dictionary as the `metadata_param` argument to the\n transformation.\n Returns:\n A new instance of the same type as self and with the same ``unbox``\n content with updated axis metadata.\n \"\"\"\n pass\n```\n\nWe call this type of class wrapping a value and keeping track of some additional data a **box**.\nBy defining an abstract base class for this box, the API does not need to be aware of the specifics of the metadata that is tracked.\nThis should make the API future proof and modular.\n\nThe ``add_axis`` and ``remove_axis`` method return an instance of their own type instead of mutating in-place.\nTypically, an implementation would be a ``flax.struct.PyTreeNode`` because the box should still be a valid JAX value and must therefore be handled by the PyTree API.\nCalling ``jax.tree.map`` on a boxed value will simply map over the value in the box.\nThe lifted transforms that need to handle metadata will call ``jax.tree.map(..., is_leaf=lambda x: isinstance(x, AxisMetadata))`` to find the AxisMetadata instances within a PyTree.\n\nAdvantages of the boxing approach:\n1. Boxing can be used outside of Flax and metadata is automatically \"inherited\". For example, the optimizer state will\n have the same partitioning spec as the parameters, because the state is initialized using a ``jax.tree.map`` over the boxed parameters.\n2. Boxes are composable.\n3. Boxing avoids string manipulation and generally avoids having to handle additional auxiliary collections like \"param_axes\" in the current\n partitioning API.\n4. No need to lift metadata collections separately.\n\n\nDisadvantages:\n1. Adding the boxes changes the PyTree hierarchy and introduces dataclasses within the otherwise plain, nested dict of variables.\n3. Custom Pytree nodes have a small runtime overhead. It's hard to observe this in practise because JAX calls are async.\n\n\n### Init syntax\n\n\nBoxes can be created directly by the init function of a variable. Therefore, we propose to create metadata using higher-order initializers.\nThe main advantage of this is that we can decouple metadata handling completely from the Module definition. Also, most Modules already overwrite\nattributes to override the default initialzers so users can add metadata to existing Modules without requiring any code changes.\n\nTo illustrate this, let's consider a metadata class that keeps track of PartitionSpecs used by ``pjit``:\n\n```python\nclass Partitioned(flax.struct.PyTreeNode, AxisMetadata):\n value: Any\n names: Tuple[Optional[str], ...] = flax.struct.field(pytree_node=False)\n\n def add_axis(self, index: int, params: Dict[Any, Any]) -> TAxisMetadata:\n axis_name = self._get_partition_name(params)\n names = list(self.names)\n names.insert(index, axis_name)\n return self.replace(names=tuple(names))\n\n def remove_axis(self, index: int, params: Dict[Any, Any]) -> TAxisMetadata:\n axis_name = self._get_partition_name(params)\n names = list(self.names)\n assert names.pop(index) == axis_name\n return self.replace(names=tuple(names))\n\ndef with_partitioning(init_fn, names):\n def wrapper(*args, **kwargs):\n return Partitioned(init_fn(*args, **kwargs), names)\n return wrapper\n```\n\nHere we also defined a small utility called ``with_partitioning`` that we can use to wrap existing initialzers to add metadata:\n\n\n```python\n# init kernel with lecun normal and split the output features over the data axis\npartitioned_dense = nn.Dense(features, kernel_init=with_partitioning(nn.initializers.lecun_normal, (None, \"data\")))\n```\n\nInitializing a model that creates partitioned weights would result in the following variable structure:\n\n```python\nvariables = partitioned_dense.init(rng, jnp.ones((4,)))\njax.tree.map(np.shape, variables) # => {\"params\": {\"kernel\": Partitioned(value=(4, 8), names=(None, \"data\")), bias: (8,)}}\n```\n\nThe variable tree with metadata can be used to integrate with other libraries and APIs.\nFor example, we can turn the ``Partitioned`` metadata into ``jax.pjit`` sharding annotations:\n\n```python\ndef to_sharding_spec(x):\n if isinstance(x, Partitioned):\n return PartitionSpec(*x.names)\n else:\n # fully replicated\n return PartitionSpec()\n\n# Result: {\"params\": {\"kernel\": PartitionSpec(None, \"data\"), bias: PartitionSpec()}}\nvariables_pspec = jax.tree.map(to_sharding_spec, variables, is_leaf=lambda x: isinstance(x, Partitioned))\n```\n\n### Unbox syntax\n\n\nMetadata typically doesn't need to be handled by Modules directly. Therefore, we propose to make Modules agnostic to Metadata boxes by default.\nThe ``unbox`` method can be used to unpack a variable such that only the original JAX arrays remain. Users can manually call unbox but to make\nsure Module classes don't have to call it everywhere we add an unbox keyword arg to variable returning APIs (e.g.: ``.param``, ``.variable``, ``.get_variable``).\nThe keyword arg ``unbox`` will default to ``True`` such that a Modules are metadata agnostic by default. This also means existing Modules will be backward compatible\nwith the new API.\n\n```python\nkernel = self.param(\"kernel\", self.kernel_init, shape) # No AxisMetadata instances\nkernel_box = self.get_variable(\"param\", \"kernel\", unbox=False) # AxisMetadata boxes are preserved\n```\n\n\n### Lift syntax\n\nWhen calling a lifted transformation that adds an axis you will now be able to pass a dictionary with arguments.\nThese params will be passed to ``AxisMetadata`` add_axis/remove_axis callbacks:\n\n```python\nnn.scan(..., variable_axes={\"params\": 0}, metadata_params={nn.Partitioned.AXIS_NAME: \"layers\"})\n```\n\nA dict is used such that users can add their own arguments to custom AxisMetadata classes.\n\n" }, { "path": "docs/flip/2974-kw-only-dataclasses.md", "content": "# FLIP: kw_only dataclasses\nAuthors: Brennan Saeta, Ivy Zheng\n\n - Start Date: Mar 23, 2023\n - FLIP Issue: [TBD]\n - FLIP PR: #2974\n - Status: Implementing\n\n\n## Summary\n\nPython 3.10 adds support for `kw_only` dataclasses. Subclasses of `flax.linen.Module` are automatically converted to `dataclasses` on users' behalf, but today, Flax doesn't allow setting the `kw_only` parameter to this dataclass transform, even if users are running Python 3.10. This proposal allows users to use this new feature with `nn.Module`'s.\n\n\n## Motivation\n\nIn larger Flax-based codebases (e.g. [`PaxML`](https://github.com/google/paxml) / [`Praxis`](https://github.com/google/praxis)), it’s not uncommon to define an (abstract) subclass of nn.Module that contains shared functionality that is itself further subclassed for specific implementations (e.g. [`BaseLayer`](https://github.com/google/praxis/blob/main/praxis/base_layer.py), or [`StackedTransformerRepeat`](https://github.com/google/praxis/blob/81479b260fcc13de8549cdbfb0fdf5c3f188ac90/praxis/layers/transformers.py#L1836) which is further subclassed by [`PipelineCompatibleStackedTransformerRepeat`](https://github.com/google/praxis/blob/81479b260fcc13de8549cdbfb0fdf5c3f188ac90/praxis/layers/transformers.py#L2198)).\n\nOften, these parent types define hyperparameters (constructor arguments), often with default values. Without `kw_only` on the `dataclass` transform, default values must be specified for all child layers hyperparameters. This is suboptimal, because users could forget to set them when instantiating the modules. For example, `Child` must set a default value for `num_heads` (because a non-defaulted argument can’t come after a defaulted argument if they are positional), but no reasonable default is available:\n\n```python\nclass BaseLayer(nn.Module):\n mesh: Optional[jax.experimental.mesh.Mesh] = None\n\n def with_sharding(self, some_variable, some_sharding):\n if self.mesh:\n # Do something useful here.\n\nclass Child(BaseLayer):\n num_heads: int # Don't want to have to set a default argument!\n\n def __call__(self, x):\n ...\n```\n\nNote: Flax already has this problem, which is why `nn.Module` has its own fancy `kw_only_dataclasses.dataclass` transform: it moves the `name` and `parent` dataclass fields to the end, so they can have defaults.\n\n\n## Implementation\n\nTo allow modules to optionally opt into this `kw_only` dataclass behavior, we leverage arguments to `__init_subclass__`. This would look as follows:\n\n```python\nclass BaseLayer(nn.Module, kw_only=True):\n ...\n\nclass Child(BaseLayer):\n ...\n```\n\nThe implementation of `nn.Module`’s `__init_subclass__` will be tweaked as follows:\n\n```python\nclass Module(ModuleBase):\n def __init_subclass__(self, kw_only: Optional[bool] = None):\n # ...\n if kw_only:\n if is_python_310_or_above():\n dataclass_transform_args = {'kw_only': True}\n else:\n raise TypeError(\"Can't use `kw_only` before Py3.10.\")\n else:\n dataclass_transform_args = {}\n\n kw_only_dataclasses.dataclass(\n cls, unsafe_hash='__hash__' not in cls.__dict__,\n repr=False,\n **dataclass_transform_args)\n```\n\n### Forward compatibility\n\nFor future simplification, if `kw_only` is requested and the Python version is 3.10 or above, bypass the `kw_only_dataclasses` implementation and just use the regular `dataclasses` transform.\n\nThat means we may one day remove `flax/linen/kw_only_dataclasses.py` when Flax rolls over 3.10.\n\n\n## Discussion\n\n### Aligned with Python `dataclass`\n\nWe prefer to keep the behavior of `nn.Module`’s `kw_only` aligned with the Python dataclasses. Note that this means `kw_only` will not be inheritable, and this could happen:\n\n```python\nclass BaseLayer(nn.Module, kw_only=True):\n base_muliplier: Optional[int] = -1\n\nclass ChildLayer(BaseLayer):\n child_multiplier: int\n\nBaseLayer(2) # This will throw error\nChildLayer(2) # But this will not\n```\n\n### `flax.struct.dataclass`\n\nThere’s a potentially related feature to allow `kw_only` to be specified for `flax.struct.dataclass`. This should be considered an orthogonal decision.\n\n\n" }, { "path": "docs/flip/3099-rnnbase-refactor.md", "content": "# Refactor RNNCellBase in FLIP\n\nAuthors: Cristian Garcia, Marcus Chiam, Jasmijn Bastings\n\n - Start Date: May 1, 2023\n - FLIP Issue: [TBD]\n - FLIP PR: #3053\n - Status: Implemented\n\n## Summary\nThis proposal aims to improve the usability of the `RNNCellBase` class by refactoring the `initialize_carry` method and other relevant components.\n\n## Motivation\n\nCurrently, `initialize_carry` is used to both initialize the carry and pass crucial metadata like the number of features. The API can be unintuitive as it requires users to manually calculate things that could typically be inferred by the modules themselves, such as the shape of batch dimensions and the shape of feature dimensions.\n\n### Example: ConvLSTM\nThe current API can be unintuitive in cases like `ConvLSTM` where a the `size` parameter contains both the input image shape and output feature dimensions:\n\n```python\nx = jnp.ones((2, 4, 4, 3)) # (batch, *image_shape, channels)\n\n# image shape: vvvvvvv\ncarry = nn.ConvLSTMCell.initialize_carry(key1, (16,), (64, 64, 16))\n# batch size: ^^ ^^ :output features\n\nlstm = nn.ConvLSTMCell(features=6, kernel_size=(3, 3))\n(carry, y), initial_params = lstm.init_with_output(key2, carry, x)\n```\n\nThis FLIP will propose some changes to `initialize_carry` such that the previous example can be simplified to:\n\n```python\nx = jnp.ones((2, 4, 4, 3)) # (batch, *image_shape, channels)\n\nlstm = nn.ConvLSTMCell(features=6, kernel_size=(3, 3))\ncarry = lstm.initialize_carry(key1, input_shape=x.shape)\n\n(carry, y), initial_params = lstm.init_with_output(key2, carry, x)\n```\n\n## Implementation\nThe proposal suggests the following changes:\n\n### initialize_carry\n`initialize_carry` should be refactored as an instance method with the following signature:\n\n```python\ndef initialize_carry(self, key, sample_input):\n```\n\n`sample_input` should be an array of the same shape that will be processed by the cell, excluding the time axis.\n\n### Refactor RNNCellBase subclasses\n`RNNCellBase` should be refactored to include the metadata required to initialize the cell and execute its forward pass. For `LSTMCell` and `GRUCell`, this means adding a `features` attribute that should be provided by the user upon construction. This change aligns with the structure of most other `Module`s, making them more familiar to users.\n\n```python\nx = jnp.ones((2, 100, 10)) # (batch, time, features)\n\ncell = nn.LSTMCell(features=32)\ncarry = cell.initialize_carry(PRNGKey(0), x[:, 0]) # sample input\n\n(carry, y), variables = cell.init_with_output(PRNGKey(1), carry, x)\n```\n\n### num_feature_dims\nTo simplify the handling of `RNNCellBase` instances in abstractions like `RNN`, each cell should implement the `num_feature_dims` property. For most cells, such as `LSTMCell` and `GRUCell`, this is always 1. For cells like `ConvLSTM`, this depends on their `kernel_size`.\n\n## Discussion\n### Alternative Approaches\n* To eliminate the need for `num_feature_dims`, `RNN` could support only a single batch dimension, i.e., inputs of the form `(batch, time, *features)`. Currently, it supports both multiple batch dimensions and multiple feature dimensions.\n* Another approach could be a complete redesign of how Flax deals with recurrent states. For example, a `memory` collection could be handled as part of the variables. However, this introduces challenges such as handling stateless cells during training, passing state from one layer to another, and performing initialization inside `scan`.\n\n### Refactor Cost\nInitial TGP results showed 761 broken and 110 failed tests. However, after fixing one test, TGP results in 231 broken and 13 failed tests so there seems to be a lot\nof overlap between the broken tests.\n\nTo minimize refactor costs, the current implementation will be kept for Google internal users under a deprecated name. This will allow users to migrate to the new API at their own pace. For Open Source users we should bump Flax version to\n`0.7.0` so existing users can continue to depend on `0.6.x` versions.\n" }, { "path": "docs/flip/4105-jax-style-nnx-transforms.md", "content": "# JAX-style NNX Transforms\n\n- Authors: Cristian Garcia, Anselm Levskaya\n- Date: Jun/2024\n- FLIP PR: #4107\n- Status: Implementing\n\n## Motivation\n\nNNX allows users to utilize Modules at the top level due to their eager initialization and self-contained state. This naturally leads users to want to use them with transforms and soon start playing with NNX transforms. Since NNX Modules resemble PyTrees in that they contain Arrays, new users often attempt to apply JAX conventions, for example:\n\n```py\n@nnx.vmap(in_axes=(1, 0))\ndef f(m1: Module, m2: Module):\n ...\n```\n\nHowever, this can be misleading. Currently, NNX transforms follow Linen's convention of treating input Modules as a single unit (all Modules are split together to preserve shared references) and provide APIs for transforming that State separately. The previous example effectively translates to:\n\n```py\n# this is what is really happening\n@nnx.vmap(in_axes=(IGNORE, IGNORE), state_axes={BatchStat: None, ...: 0})\ndef f(m1: Module, m2: Module):\n ...\n```\n\nNote that `IGNORE` is not a real symbol, but represents the fact that any value placed here won't affect the outcome, as Modules are replaced by empty PyTree placeholders (similar to `None`). The `state_axes` parameter controls how the State is vectorized through a mapping of high-level `Filter`s to their desired axes. In this example, `...` (ellipsis) is a filter that accepts everything, so by default all States are vectorized on the 0th axis.\n\nTo express their original intention, users must resort to more complex custom filters that guess the index of each Module in the monolith. While this is straightforward in simple cases, users generally need to calculate the index (Modules appear in the order specified by `jax.tree.leaves` over the `args`):\n\n```py\nselect_m1 = lambda path, value: path[0] == 0\nselect_m2 = lambda path, value: path[0] == 1\n\n# To select modules individually, you must create a filter (which can be tricky)\n@nnx.vmap(state_axes={select_m1: 1, select_m2: 0})\ndef f(m1: Module, m2: Module):\n ...\n```\n\n## What if JAX conventions Just Worked™?\n\nThis proposal aims to align NNX transforms with user's expectations based on their JAX experience, making the syntax work as intuitively as possible. The original example would function **as if** `m1` and `m2` were PyTrees vectorized in axes `1` and `0` respectively:\n\n```py\n@nnx.vmap(in_axes=(1, 0))\ndef f(m1: Module, m2: Module):\n ...\n```\n\nThe primary advantage of this approach is that for `vmap` and `scan`, we could eliminate the `state_axes` and `split_rngs` arguments, relying solely on the `in_axes` API. This syntax alone would likely suffice for 80-90% of use cases, as users tend to manage state in predictable ways.\n\n### The Lift symbols\n\nTo enable more fine-grained state control within each Module, we introduce the `Lift` API. By using special types containing State Filters in place of a tree prefix, state lifting can now be done **structurally**. This allows different Filters to be applied to different Modules in the arguments without the need for complex path-based filters. Ideally, each transform would support its own Lift type, adding the desired behavior through existing JAX APIs.\n\nFor example, in `vmap`, we could allow `StateAxes` instances (vmap's Lift type) to be accepted by `in/out_axes` to control how substates are handled by mapping state `Filter`s to an axis specifier:\n\n```py\nstate_axes = StateAxes({Param: 1, BatchStat: None})\n\n@nnx.vmap(in_axes=(state_axes, 0))\ndef f(m1: Module, m2: Module):\n ...\n```\n\nIn this case, `m1`'s `Param`s are vectorized in axis `1` while its `BatchStat`s are broadcasted, and `m2`'s entire state is vectorized in axis `0`.\n\nFor `nnx.grad`, we could allow `DiffState` to be used in the `argnums` parameter to specify both the position of the argument to be differentiated and a Filter specifying the differentiable State of the Module:\n\n```py\ngrads = nnx.grad(loss_fn, argnums=(DiffState(0, LoRAParam),))(model, x, y)\n```\n\n## Rng Handling\n\nTo simplify RNG state handling, we propose removing the separate `split_rngs` parameter in `vmap` and `scan`. Instead, we suggest introducing a new `nnx.split_rngs` API that would manage RNG handling before and after the transformation. This approach provides more explicit control to the user and aligns better with JAX transform behavior.\n\n## Consistent Aliasing\n\nTo ensure the correctness of transformations with objects that obey reference semantics, we must enforce consistent lifting/lowering specifications for all aliases of a reference. Transforms must adhere to two rules:\n\n1. All aliases of a reference must receive the **exact same** lifting/lowering specification.\n2. Captured references are not allowed on the output of transformed functions.\n\nFor example:\n\n```py\n@nnx.vmap(in_axes=(m1_axes, m2_axes, m1_axes), out_axes=m2_axes)\ndef f(m1, m2, m1_alias):\n return m2\n\nm2 = f(m1, m2, m1)\n```\n\nHere, `m1` has two input aliases as it is passed as the first and third input to `f`, but this is acceptable because `m1_axes` is assigned to both in `in_axes`. `m2` is passed as the second input and has an output alias, which is also acceptable because `m2_axes` is assigned in both `in_axes` and `out_axes`.\n\nLet's examine some examples of programs that should be **rejected** based on these criteria:\n\n### Inconsistent input aliases\n\nConsider a function with two arguments `m1` and `m2` being vectorized in axis `0` and `1` respectively. Passing the same Module as both arguments would be inconsistent:\n\n```py\n@nnx.vmap(in_axes=(0, 1))\ndef f(m1: Module, m2: Module):\n ...\n\nf(m, m) # This should be rejected\n```\n\n### Inconsistent input / output aliases\n\nNow consider an identity function `g` under `vmap` with `in_axes=0` and `out_axes=1`. In JAX, this would result in transposing the arrays in the inputs:\n\n```py\n@nnx.vmap(in_axes=0, out_axes=1)\ndef g(m: Module):\n return m\n```\n\nWhile this appears correct, in NNX this behavior is not well-defined because shared mutable references behave as auxiliary outputs. Under the hood, `g` is converted into a function that has the inputs as an extra first output, and `out_axes` is set to the same values as `in_axes` for that output:\n\n```py\n@nnx.vmap(in_axes=0, out_axes=(0, 1))\ndef g_real(m: Module):\n return m, m\n```\n\nThis return structure reveals an inconsistency: we're attempting to lower `m` with both `out_axes=0` and `out_axes=1`.\n\n### Inconsistent aliases in nested structures\n\nSimilar issues can arise in less obvious cases, such as when `m` is contained within another structure:\n\n```py\n@nnx.vmap(in_axes=0, out_axes=1)\ndef f(m: Module):\n return SomeModule(m)\n```\n\nThis means we must traverse the entire graph of both inputs and outputs to check for consistent assignments. The same problem occurs when passing shared reference inputs/outputs with different specifications:\n\n```py\nshared = Shared()\nm1, m2 = Foo(shared), Foo(shared)\n\n@nnx.vmap(in_axes=(0, 1))\ndef f(m1, m2): # shared is passed through both\n ...\n```\n\n### Captured Modules cannot be outputs\n\nFinally, let's consider the second consistent aliasing rule, which states that captured Modules cannot be outputs. The main issue here is that NNX needs to split all input references together to track changes, but captured Modules bypass this process. Treating them as new references would result in **implicit cloning**:\n\n```py\nm = SomeModule()\n\n@nnx.vmap(out_axes=0, axis_size=5)\ndef f():\n return m\n\nassert m is not f() # implicit cloning\n```\n\nTo preserve reference identity, we must disallow captured Modules as outputs. In practice, we can detect captured Modules using the trace level context machinery used to restrict stateful updates on Modules from a different level.\n\n## Recap\n\nIn this document, we have:\n\n* Discussed issues with the current implementation that make it unintuitive for JAX users.\n* Proposed refactoring NNX transforms to allow users to use regular JAX semantics when interacting with objects, removing extra arguments introduced by NNX transforms.\n* Introduced the use of Lift types in JAX APIs to compensate for the lack of a \"prefix\" notion in NNX objects, enabling independent lifting of Module substates.\n* Proposed a new `nnx.split_rngs` API to replace the `split_rngs` arguments in `vmap` and `scan`, making RNG handling an explicit operation and giving users more control.\n* Analyzed edge cases resulting from aliasing shared mutable references and proposed enforcing **consistent aliasing** on all transforms with semantics over the inputs." }, { "path": "docs/flip/README.md", "content": "# FLIP: Flax Improvement Process\n\nMost changes can be discussed with simple issues/discussions and pull requests.\n\nSome changes though are a bit larger in scope or require more discussion, and\nthese should be implemented as FLIPs. This allows for writing longer documents\nthat can be discussed in a pull request themselves.\n\nThe structure of FLIPs is kept as lightweight as possible to start and might\nbe extended later on.\n\n## When you should use a FLIP\n\n- When your change requires a design doc. We prefer collecting the designs as\n FLIPs for better discoverability and further reference.\n\n- When your change requires extensive discussion. It's fine to have relatively\n short discussions on issues or pull requests, but when the discussion gets\n longer this becomes unpractical for later digestion. FLIPs allow to update the\n main document with a summary of the discussion and these updates can be\n discussed themselves in the pull request adding the FLIP.\n\n## How to start a FLIP\n\nFirst, create an issue with the [FLIP label]. All pull requests that relate to\nthe FLIP (i.e. adding the FLIP itself as well as any implementing pull requests)\nshould be linked to this issue.\n\nThen create a pull request that consists of a copy of the `0000-template.md`\nrenamed to `%04d-{short-title}.md` - with the number being the issue number.\n\n[FLIP label]: https://github.com/google/flax/issues?q=label%3AFLIP\n" }, { "path": "docs/glossary.rst", "content": "*********\nGlossary\n*********\n\nFor additional terms, refer to the `Jax glossary `__.\n\n.. glossary::\n\n Bound Module\n When a :class:`Module `\n is created through regular Python object construction (e.g. `module = SomeModule(args...)`, it is in an *unbound* state. This means that only\n dataclass attributes are set, and no variables are bound to the module. When the pure\n functions :meth:`Module.init() `\n or :meth:`Module.apply() `\n are called, Flax clones the Module and binds the variables to it, and the module's method code is\n executed in a locally bound state, allowing things like calling submodules directly without\n providing variables. For more details, refer to the\n `module lifecycle `__.\n\n Compact / Non-compact Module\n Modules with a single method are able to declare submodules and variables inline by\n using the :func:`@nn.compact ` decorator.\n These are referred to as “compact-style modules”,\n whereas modules defining a :meth:`setup() ` method\n (usually but not always with multiple callable methods)\n are referred to as “setup-style modules”. To learn more, refer to the\n `setup vs compact guide `__.\n\n `Folding in `__\n Generating a new PRNG key given an input PRNG key and integer. Typically used when you want to\n generate a new key but still be able to use the original rng key afterwards. You can also do this with\n `jax.random.split `__\n but this will effectively create two RNG keys, which is slower. See how Flax generates new PRNG keys\n automatically within ``Modules`` in our\n `RNG guide `__.\n\n `FrozenDict `__\n An immutable dictionary which can be “`unfrozen `__”\n to a regular, mutable dictionary. Internally, Flax uses FrozenDicts to ensure variable dicts\n aren't accidentally mutated. Note: We are considering returning to regular dicts from our APIs,\n and only using FrozenDicts internally.\n (see `#1223 `__).\n\n Functional core\n The flax core library implements the simple container Scope API for threading\n variables and PRNGs through a model, as well as the lifting machinery needed to\n transform functions passing Scope objects. The python class-based module API\n is built on top of this core library.\n\n Lazy initialization\n Variables in Flax are initialized late, only when needed. That is, during normal\n execution of a module, if a requested variable name isn’t found in the provided\n variable collection data, we call the initializer function to create it. This\n allows us to treat initialization and application under the same code-paths,\n simplifying the use of JAX transforms with layers.\n\n Lifted transformation\n Refer to the `Flax docs `__.\n\n Module\n A dataclass allowing the definition and initialization of parameters in a\n referentially-transparent form. This is responsible for storing and updating variables\n and parameters within itself. Modules can be readily transformed into functions,\n allowing them to be trivially used with JAX transformations like `vmap` and `scan`.\n\n Params / parameters\n \"params\" is the canonical variable collection in the variable dictionary (dict).\n The “params” collection generally contains the trainable weights.\n\n RNG sequences\n Inside Flax :class:`Modules `, you can obtain a new\n `PRNG `__\n key through :meth:`Module.make_rng() `.\n These keys can be used to generate random numbers through\n `JAX's functional random number generators `__.\n Having different RNG sequences (e.g. for \"params\" and \"dropout\") allows fine-grained\n control in a multi-host setup (e.g. initializing parameters identically on different\n hosts, but have different dropout masks) and treating these sequences differently when\n `lifting transformations `__.\n See the `RNG guide `__\n for more details.\n\n Scope\n A container class for holding the variables and PRNG keys for each layer.\n\n Shape inference\n Modules do not need to specify the shape of the input array in their definitions.\n Flax upon initialization inspects the input array, and infers the correct shapes\n for parameters in the model.\n\n TrainState\n Refer to :class:`flax.training.train_state.TrainState`.\n\n Variable\n The `weights / parameters / data / arrays `__\n residing in the leaves of :term:`variable collections`.\n Variables are defined inside modules using :meth:`Module.variable() `.\n A variable of collection \"params\" is simply called a param and can be set using\n :meth:`Module.param() `.\n\n Variable collections\n Entries in the variable dict, containing weights / parameters / data / arrays that\n are used by the model. “params” is the canonical collection in the variable dict.\n They are typically differentiable, updated by an outer SGD-like loop / optimizer,\n rather than modified directly by forward-pass code.\n\n `Variable dictionary `__\n A dictionary containing :term:`variable collections`.\n Each variable collection is a mapping from a string name\n (e.g., \":term:`params`\" or \"batch_stats\") to a (possibly nested)\n dictionary with :term:`Variables` as leaves, matching the submodule tree structure.\n Read more about pytrees and leaves in the `Jax docs `__." }, { "path": "docs/guides/converting_and_upgrading/convert_pytorch_to_flax.rst", "content": "Convert PyTorch models to Flax\n==============================\n\n.. testsetup::\n\n import numpy as np\n import jax\n from jax import random, numpy as jnp\n import flax\n\n from flax import linen as nn\n\n import torch\n\nWe will show how to convert PyTorch models to Flax. We will cover convolutions, fc layers, batch norm, and average pooling.\n\n\nFC Layers\n--------------------------------\n\nLet's start with fc layers. The only thing to be aware of here is that the PyTorch kernel has shape [outC, inC]\nand the Flax kernel has shape [inC, outC]. Transposing the kernel will do the trick.\n\n.. testcode::\n\n t_fc = torch.nn.Linear(in_features=3, out_features=4)\n\n kernel = t_fc.weight.detach().cpu().numpy()\n bias = t_fc.bias.detach().cpu().numpy()\n\n # [outC, inC] -> [inC, outC]\n kernel = jnp.transpose(kernel, (1, 0))\n\n key = random.key(0)\n x = random.normal(key, (1, 3))\n\n variables = {'params': {'kernel': kernel, 'bias': bias}}\n j_fc = nn.Dense(features=4)\n j_out = j_fc.apply(variables, x)\n\n t_x = torch.from_numpy(np.array(x))\n t_out = t_fc(t_x)\n t_out = t_out.detach().cpu().numpy()\n\n np.testing.assert_almost_equal(j_out, t_out, decimal=6)\n\n\nConvolutions\n--------------------------------\n\nLet's now look at 2D convolutions. PyTorch uses the NCHW format and Flax uses NHWC.\nConsequently, the kernels will have different shapes. The kernel in PyTorch has shape [outC, inC, kH, kW]\nand the Flax kernel has shape [kH, kW, inC, outC]. Transposing the kernel will do the trick.\n\n.. testcode::\n\n t_conv = torch.nn.Conv2d(in_channels=3, out_channels=4, kernel_size=2, padding='valid')\n\n kernel = t_conv.weight.detach().cpu().numpy()\n bias = t_conv.bias.detach().cpu().numpy()\n\n # [outC, inC, kH, kW] -> [kH, kW, inC, outC]\n kernel = jnp.transpose(kernel, (2, 3, 1, 0))\n\n key = random.key(0)\n x = random.normal(key, (1, 6, 6, 3))\n\n variables = {'params': {'kernel': kernel, 'bias': bias}}\n j_conv = nn.Conv(features=4, kernel_size=(2, 2), padding='valid')\n j_out = j_conv.apply(variables, x)\n\n # [N, H, W, C] -> [N, C, H, W]\n t_x = torch.from_numpy(np.transpose(np.array(x), (0, 3, 1, 2)))\n t_out = t_conv(t_x)\n # [N, C, H, W] -> [N, H, W, C]\n t_out = np.transpose(t_out.detach().cpu().numpy(), (0, 2, 3, 1))\n\n np.testing.assert_almost_equal(j_out, t_out, decimal=6)\n\n\n\nConvolutions and FC Layers\n--------------------------------\n\nWe have to be careful, when we have a model that uses convolutions followed by fc layers (ResNet, VGG, etc).\nIn PyTorch, the activations will have shape [N, C, H, W] after the convolutions and are then\nreshaped to [N, C * H * W] before being fed to the fc layers.\nWhen we port our weights from PyTorch to Flax, the activations after the convolutions will be of shape [N, H, W, C] in Flax.\nBefore we reshape the activations for the fc layers, we have to transpose them to [N, C, H, W].\n\nConsider this PyTorch model:\n\n.. testcode::\n\n class TModel(torch.nn.Module):\n\n def __init__(self):\n super(TModel, self).__init__()\n self.conv = torch.nn.Conv2d(in_channels=3, out_channels=4, kernel_size=2, padding='valid')\n self.fc = torch.nn.Linear(in_features=100, out_features=2)\n\n def forward(self, x):\n x = self.conv(x)\n x = x.reshape(x.shape[0], -1)\n x = self.fc(x)\n return x\n\n\n t_model = TModel()\n\n\n\nNow, if you want to use the weights from this model in Flax, the corresponding Flax model has to look like this:\n\n\n.. testcode::\n\n class JModel(nn.Module):\n\n @nn.compact\n def __call__(self, x):\n x = nn.Conv(features=4, kernel_size=(2, 2), padding='valid', name='conv')(x)\n # [N, H, W, C] -> [N, C, H, W]\n x = jnp.transpose(x, (0, 3, 1, 2))\n x = jnp.reshape(x, (x.shape[0], -1))\n x = nn.Dense(features=2, name='fc')(x)\n return x\n\n\n j_model = JModel()\n\n\n\nThe model looks very similar to the PyTorch model, except that we included a transpose operation before\nreshaping our activations for the fc layer.\nWe can omit the transpose operation if we apply pooling before reshaping such that the spatial dimensions are 1x1.\n\nOther than the transpose operation before reshaping, we can convert the weights the same way as we did before:\n\n\n.. testcode::\n\n conv_kernel = t_model.state_dict()['conv.weight'].detach().cpu().numpy()\n conv_bias = t_model.state_dict()['conv.bias'].detach().cpu().numpy()\n fc_kernel = t_model.state_dict()['fc.weight'].detach().cpu().numpy()\n fc_bias = t_model.state_dict()['fc.bias'].detach().cpu().numpy()\n\n # [outC, inC, kH, kW] -> [kH, kW, inC, outC]\n conv_kernel = jnp.transpose(conv_kernel, (2, 3, 1, 0))\n\n # [outC, inC] -> [inC, outC]\n fc_kernel = jnp.transpose(fc_kernel, (1, 0))\n\n variables = {'params': {'conv': {'kernel': conv_kernel, 'bias': conv_bias},\n 'fc': {'kernel': fc_kernel, 'bias': fc_bias}}}\n\n key = random.key(0)\n x = random.normal(key, (1, 6, 6, 3))\n\n j_out = j_model.apply(variables, x)\n\n # [N, H, W, C] -> [N, C, H, W]\n t_x = torch.from_numpy(np.transpose(np.array(x), (0, 3, 1, 2)))\n t_out = t_model(t_x)\n t_out = t_out.detach().cpu().numpy()\n\n np.testing.assert_almost_equal(j_out, t_out, decimal=6)\n\n\n\nBatch Norm\n--------------------------------\n\n``torch.nn.BatchNorm2d`` uses ``0.1`` as the default value for the ``momentum`` parameter while\n|nn.BatchNorm|_ uses ``0.9``. However, this corresponds to the same computation, because PyTorch multiplies\nthe estimated statistic with ``(1 − momentum)`` and the new observed value with ``momentum``,\nwhile Flax multiplies the estimated statistic with ``momentum`` and the new observed value with ``(1 − momentum)``.\n\n.. |nn.BatchNorm| replace:: ``nn.BatchNorm``\n.. _nn.BatchNorm: https://flax.readthedocs.io/en/latest/api_reference/flax.linen/layers.html#flax.linen.BatchNorm\n\n.. testcode::\n\n t_bn = torch.nn.BatchNorm2d(num_features=3, momentum=0.1)\n t_bn.eval()\n\n scale = t_bn.weight.detach().cpu().numpy()\n bias = t_bn.bias.detach().cpu().numpy()\n mean = t_bn.running_mean.detach().cpu().numpy()\n var = t_bn.running_var.detach().cpu().numpy()\n\n variables = {'params': {'scale': scale, 'bias': bias},\n 'batch_stats': {'mean': mean, 'var': var}}\n\n key = random.key(0)\n x = random.normal(key, (1, 6, 6, 3))\n\n j_bn = nn.BatchNorm(momentum=0.9, use_running_average=True)\n\n j_out = j_bn.apply(variables, x)\n\n # [N, H, W, C] -> [N, C, H, W]\n t_x = torch.from_numpy(np.transpose(np.array(x), (0, 3, 1, 2)))\n t_out = t_bn(t_x)\n # [N, C, H, W] -> [N, H, W, C]\n t_out = np.transpose(t_out.detach().cpu().numpy(), (0, 2, 3, 1))\n\n np.testing.assert_almost_equal(j_out, t_out, decimal=6)\n\n\n\nAverage Pooling\n--------------------------------\n\n``torch.nn.AvgPool2d`` and |nn.avg_pool()|_ are compatible when using default parameters.\nHowever, ``torch.nn.AvgPool2d`` has a parameter ``count_include_pad``. When ``count_include_pad=False``,\nthe zero-padding will not be considered for the average calculation. There does not exist a similar\nparameter for |nn.avg_pool()|_. However, we can easily implement a wrapper around the pooling\noperation. ``nn.pool()`` is the core function behind |nn.avg_pool()|_ and |nn.max_pool()|_.\n\n.. |nn.avg_pool()| replace:: ``nn.avg_pool()``\n.. _nn.avg_pool(): https://flax.readthedocs.io/en/latest/api_reference/flax.linen/layers.html#flax.linen.avg_pool\n\n.. |nn.max_pool()| replace:: ``nn.max_pool()``\n.. _nn.max_pool(): https://flax.readthedocs.io/en/latest/api_reference/flax.linen/layers.html#flax.linen.max_pool\n\n\n.. testcode::\n\n def avg_pool(inputs, window_shape, strides=None, padding='VALID'):\n \"\"\"\n Pools the input by taking the average over a window.\n In comparison to nn.avg_pool(), this pooling operation does not\n consider the padded zero's for the average computation.\n \"\"\"\n assert len(window_shape) == 2\n\n y = nn.pool(inputs, 0., jax.lax.add, window_shape, strides, padding)\n counts = nn.pool(jnp.ones_like(inputs), 0., jax.lax.add, window_shape, strides, padding)\n y = y / counts\n return y\n\n\n key = random.key(0)\n x = random.normal(key, (1, 6, 6, 3))\n\n j_out = avg_pool(x, window_shape=(2, 2), strides=(1, 1), padding=((1, 1), (1, 1)))\n t_pool = torch.nn.AvgPool2d(kernel_size=2, stride=1, padding=1, count_include_pad=False)\n\n # [N, H, W, C] -> [N, C, H, W]\n t_x = torch.from_numpy(np.transpose(np.array(x), (0, 3, 1, 2)))\n t_out = t_pool(t_x)\n # [N, C, H, W] -> [N, H, W, C]\n t_out = np.transpose(t_out.detach().cpu().numpy(), (0, 2, 3, 1))\n\n np.testing.assert_almost_equal(j_out, t_out, decimal=6)\n\n\n\nTransposed Convolutions\n--------------------------------\n\n``torch.nn.ConvTranspose2d`` and |nn.ConvTranspose|_ are not compatible.\n|nn.ConvTranspose|_ is a wrapper around |jax.lax.conv_transpose|_ which computes a fractionally strided convolution,\nwhile ``torch.nn.ConvTranspose2d`` computes a gradient based transposed convolution. Currently, there is no\nimplementation of a gradient based transposed convolution is ``Jax``. However, there is a pending `pull request`_\nthat contains an implementation.\n\nTo load ``torch.nn.ConvTranspose2d`` parameters into Flax, we need to use the ``transpose_kernel`` arg in Flax's\n``nn.ConvTranspose`` layer.\n\n.. testcode::\n\n # padding is inverted\n torch_padding = 0\n flax_padding = 1 - torch_padding\n\n t_conv = torch.nn.ConvTranspose2d(in_channels=3, out_channels=4, kernel_size=2, padding=torch_padding)\n\n kernel = t_conv.weight.detach().cpu().numpy()\n bias = t_conv.bias.detach().cpu().numpy()\n\n # [inC, outC, kH, kW] -> [kH, kW, outC, inC]\n kernel = jnp.transpose(kernel, (2, 3, 1, 0))\n\n key = random.key(0)\n x = random.normal(key, (1, 6, 6, 3))\n\n variables = {'params': {'kernel': kernel, 'bias': bias}}\n # ConvTranspose expects the kernel to be [kH, kW, inC, outC],\n # but with `transpose_kernel=True`, it expects [kH, kW, outC, inC] instead\n j_conv = nn.ConvTranspose(features=4, kernel_size=(2, 2), padding=flax_padding, transpose_kernel=True)\n j_out = j_conv.apply(variables, x)\n\n # [N, H, W, C] -> [N, C, H, W]\n t_x = torch.from_numpy(np.transpose(np.array(x), (0, 3, 1, 2)))\n t_out = t_conv(t_x)\n # [N, C, H, W] -> [N, H, W, C]\n t_out = np.transpose(t_out.detach().cpu().numpy(), (0, 2, 3, 1))\n np.testing.assert_almost_equal(j_out, t_out, decimal=6)\n\n\n.. _`pull request`: https://github.com/jax-ml/jax/pull/5772\n\n.. |nn.ConvTranspose| replace:: ``nn.ConvTranspose``\n.. _nn.ConvTranspose: https://flax.readthedocs.io/en/latest/api_reference/flax.linen/layers.html#flax.linen.ConvTranspose\n\n.. |jax.lax.conv_transpose| replace:: ``jax.lax.conv_transpose``\n.. _jax.lax.conv_transpose: https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.conv_transpose.html\n\n" }, { "path": "docs/guides/converting_and_upgrading/haiku_migration_guide.rst", "content": "\nMigrating from Haiku to Flax\n============================\n\nThis guide will walk through the process of migrating Haiku models to Flax,\nand highlight the differences between the two libraries.\n\n.. testsetup:: Haiku, Flax\n\n import jax\n import jax.numpy as jnp\n from jax import random\n import optax\n import flax.linen as nn\n import haiku as hk\n\nBasic Example\n-----------------\n\nTo create custom Modules you subclass from a ``Module`` base class in\nboth Haiku and Flax. However, Haiku classes use a regular ``__init__`` method\nwhereas Flax classes are ``dataclasses``, meaning you define some class\nattributes that are used to automatically generate a constructor. Also,\nall Flax Modules accept a ``name`` argument without needing to define it,\nwhereas in Haiku ``name`` must be explicitly defined in the constructor\nsignature and passed to the superclass constructor.\n\n.. codediff::\n :title: Haiku, Flax\n :sync:\n\n import haiku as hk\n\n class Block(hk.Module):\n def __init__(self, features: int, name=None):\n super().__init__(name=name)\n self.features = features\n\n def __call__(self, x, training: bool):\n x = hk.Linear(self.features)(x)\n x = hk.dropout(hk.next_rng_key(), 0.5 if training else 0, x)\n x = jax.nn.relu(x)\n return x\n\n class Model(hk.Module):\n def __init__(self, dmid: int, dout: int, name=None):\n super().__init__(name=name)\n self.dmid = dmid\n self.dout = dout\n\n def __call__(self, x, training: bool):\n x = Block(self.dmid)(x, training)\n x = hk.Linear(self.dout)(x)\n return x\n\n ---\n\n import flax.linen as nn\n\n class Block(nn.Module):\n features: int\n\n\n @nn.compact\n def __call__(self, x, training: bool):\n x = nn.Dense(self.features)(x)\n x = nn.Dropout(0.5, deterministic=not training)(x)\n x = jax.nn.relu(x)\n return x\n\n class Model(nn.Module):\n dmid: int\n dout: int\n\n\n @nn.compact\n def __call__(self, x, training: bool):\n x = Block(self.dmid)(x, training)\n x = nn.Dense(self.dout)(x)\n return x\n\nThe ``__call__`` method looks very similar in both libraries, however, in Flax\nyou have to use the ``@nn.compact`` decorator in order to be able to define\nsubmodules inline. In Haiku, this is the default behavior.\n\nNow, a place where Haiku and Flax differ substantially is in how you construct\nthe model. In Haiku, you use ``hk.transform`` over a function\nthat calls your Module, ``transform`` will return an object with ``init``\nand ``apply`` methods. In Flax, you simply instantiate your Module.\n\n.. codediff::\n :title: Haiku, Flax\n :sync:\n\n def forward(x, training: bool):\n return Model(256, 10)(x, training)\n\n model = hk.transform(forward)\n\n ---\n\n ...\n\n\n model = Model(256, 10)\n\nTo get the model parameters in both libraries you use the ``init`` method\nwith a ``random.key`` plus some inputs to run the model. The main difference here is\nthat Flax returns a mapping from collection names to nested array dictionaries,\n``params`` is just one of these possible collections. In Haiku, you get the ``params``\nstructure directly.\n\n.. codediff::\n :title: Haiku, Flax\n :sync:\n\n sample_x = jax.numpy.ones((1, 784))\n params = model.init(\n random.key(0),\n sample_x, training=False # <== inputs\n )\n ...\n\n ---\n\n sample_x = jax.numpy.ones((1, 784))\n variables = model.init(\n random.key(0),\n sample_x, training=False # <== inputs\n )\n params = variables[\"params\"]\n\nOne very important thing to note is that in Flax the parameters structure is\nhierarchical, with one level per nested module and a final level for the\nparameter name.\nIn Haiku the parameters structure is a python dictionary with a two level hierarchy:\nthe fully qualified module name mapping to the parameter name. The module name\nconsists of a ``/`` separated string path of all the nested Modules.\n\n.. tab-set::\n\n .. tab-item:: Haiku\n :sync: Haiku\n\n .. code-block:: python\n\n ...\n {\n 'model/block/linear': {\n 'b': (256,),\n 'w': (784, 256),\n },\n 'model/linear': {\n 'b': (10,),\n 'w': (256, 10),\n }\n }\n ...\n\n\n .. tab-item:: Flax\n :sync: Flax\n\n .. code-block:: python\n\n FrozenDict({\n Block_0: {\n Dense_0: {\n bias: (256,),\n kernel: (784, 256),\n },\n },\n Dense_0: {\n bias: (10,),\n kernel: (256, 10),\n },\n })\n\nDuring training in both frameworks you pass the parameters structure to the\n``apply`` method to run the forward pass. Since we are using dropout, in\nboth cases we must provide a ``key`` to ``apply`` in order to generate\nthe random dropout masks.\n\n.. codediff::\n :title: Haiku, Flax\n :sync:\n\n def train_step(key, params, inputs, labels):\n def loss_fn(params):\n logits = model.apply(\n params,\n key,\n inputs, training=True # <== inputs\n )\n return optax.softmax_cross_entropy_with_integer_labels(logits, labels).mean()\n\n grads = jax.grad(loss_fn)(params)\n params = jax.tree_util.tree_map(lambda p, g: p - 0.1 * g, params, grads)\n\n return params\n\n ---\n\n def train_step(key, params, inputs, labels):\n def loss_fn(params):\n logits = model.apply(\n {'params': params},\n inputs, training=True, # <== inputs\n rngs={'dropout': key}\n )\n return optax.softmax_cross_entropy_with_integer_labels(logits, labels).mean()\n\n grads = jax.grad(loss_fn)(params)\n params = jax.tree_util.tree_map(lambda p, g: p - 0.1 * g, params, grads)\n\n return params\n\n.. testcode:: Haiku, Flax\n :hide:\n\n train_step(random.key(0), params, sample_x, jnp.ones((1,), dtype=jnp.int32))\n\nThe most notable differences is that in Flax you have to\npass the parameters inside a dictionary with a ``params`` key, and the\nkey inside a dictionary with a ``dropout`` key. This is because in Flax\nyou can have many types of model state and random state. In Haiku, you\njust pass the parameters and the key directly.\n\nHandling State\n-----------------\n\nNow let's see how mutable state is handled in both libraries. We will take\nthe same model as before, but now we will replace Dropout with BatchNorm.\n\n.. codediff::\n :title: Haiku, Flax\n :sync:\n\n class Block(hk.Module):\n def __init__(self, features: int, name=None):\n super().__init__(name=name)\n self.features = features\n\n def __call__(self, x, training: bool):\n x = hk.Linear(self.features)(x)\n x = hk.BatchNorm(\n create_scale=True, create_offset=True, decay_rate=0.99\n )(x, is_training=training)\n x = jax.nn.relu(x)\n return x\n\n ---\n\n class Block(nn.Module):\n features: int\n\n\n @nn.compact\n def __call__(self, x, training: bool):\n x = nn.Dense(self.features)(x)\n x = nn.BatchNorm(\n momentum=0.99\n )(x, use_running_average=not training)\n x = jax.nn.relu(x)\n return x\n\nThe code is very similar in this case as both libraries provide a BatchNorm\nlayer. The most notable difference is that Haiku uses ``is_training`` to\ncontrol whether or not to update the running statistics, whereas Flax uses\n``use_running_average`` for the same purpose.\n\nTo instantiate a stateful model in Haiku you use ``hk.transform_with_state``,\nwhich changes the signature for ``init`` and ``apply`` to accept and return\nstate. As before, in Flax you construct the Module directly.\n\n.. codediff::\n :title: Haiku, Flax\n :sync:\n\n def forward(x, training: bool):\n return Model(256, 10)(x, training)\n\n model = hk.transform_with_state(forward)\n\n ---\n\n ...\n\n\n model = Model(256, 10)\n\n\nTo initialize both the parameters and state you just call the ``init`` method\nas before. However, in Haiku you now get ``state`` as a second return value, and\nin Flax you get a new ``batch_stats`` collection in the ``variables`` dictionary.\nNote that since ``hk.BatchNorm`` only initializes batch statistics when\n``is_training=True``, we must set ``training=True`` when initializing parameters\nof a Haiku model with an ``hk.BatchNorm`` layer. In Flax, we can set\n``training=False`` as usual.\n\n.. codediff::\n :title: Haiku, Flax\n :sync:\n\n sample_x = jax.numpy.ones((1, 784))\n params, state = model.init(\n random.key(0),\n sample_x, training=True # <== inputs #!\n )\n ...\n\n ---\n\n sample_x = jax.numpy.ones((1, 784))\n variables = model.init(\n random.key(0), #!\n sample_x, training=False # <== inputs\n )\n params, batch_stats = variables[\"params\"], variables[\"batch_stats\"]\n\n\nIn general, in Flax you might find other state collections in the ``variables``\ndictionary such as ``cache`` for auto-regressive transformers models,\n``intermediates`` for intermediate values added using ``Module.sow``, or other\ncollection names defined by custom layers. Haiku only makes a distinction\nbetween ``params`` (variables which do not change while running ``apply``) and\n``state`` (variables which can change while running ``apply``).\n\nNow, training looks very similar in both frameworks as you use the same\n``apply`` method to run the forward pass. In Haiku, now pass the ``state``\nas the second argument to ``apply``, and get the new state as the second\nreturn value. In Flax, you instead add ``batch_stats`` as a new key to the\ninput dictionary, and get the ``updates`` variables dictionary as the second\nreturn value.\n\n.. codediff::\n :title: Haiku, Flax\n :sync:\n\n def train_step(params, state, inputs, labels):\n def loss_fn(params):\n logits, new_state = model.apply(\n params, state,\n None, # <== rng\n inputs, training=True # <== inputs\n )\n loss = optax.softmax_cross_entropy_with_integer_labels(logits, labels).mean()\n return loss, new_state\n\n grads, new_state = jax.grad(loss_fn, has_aux=True)(params)\n params = jax.tree_util.tree_map(lambda p, g: p - 0.1 * g, params, grads)\n\n return params, new_state\n ---\n\n def train_step(params, batch_stats, inputs, labels):\n def loss_fn(params):\n logits, updates = model.apply(\n {'params': params, 'batch_stats': batch_stats},\n inputs, training=True, # <== inputs\n mutable='batch_stats',\n )\n loss = optax.softmax_cross_entropy_with_integer_labels(logits, labels).mean()\n return loss, updates[\"batch_stats\"]\n\n grads, batch_stats = jax.grad(loss_fn, has_aux=True)(params)\n params = jax.tree_util.tree_map(lambda p, g: p - 0.1 * g, params, grads)\n\n return params, batch_stats\n\n.. testcode:: Flax\n :hide:\n\n train_step(params, batch_stats, sample_x, jnp.ones((1,), dtype=jnp.int32))\n\nOne major difference is that in Flax a state collection can be mutable or immutable.\nDuring ``init`` all collections are mutable by default, however, during ``apply``\nyou have to explicitly specify which collections are mutable. In this example,\nwe specify that ``batch_stats`` is mutable. Here a single string is passed but a list\ncan also be given if there are more mutable collections. If this is not done an\nerror will be raised at runtime when trying to mutate ``batch_stats``.\nAlso, when ``mutable`` is anything other than ``False``, the ``updates``\ndictionary is returned as the second return value of ``apply``, else only the\nmodel output is returned.\nHaiku makes the mutable/immutable distinction through having ``params``\n(immutable) and ``state`` (mutable) and using either ``hk.transform`` or\n``hk.transform_with_state``\n\nUsing Multiple Methods\n-----------------------\n\nIn this section we will take a look at how to use multiple methods in Haiku and Flax.\nAs an example, we will implement an auto-encoder model with three methods:\n``encode``, ``decode``, and ``__call__``.\n\nIn Haiku, we can just define the submodules that ``encode`` and ``decode`` need\ndirectly in ``__init__``, in this case each will just use a ``Linear`` layer.\nIn Flax, we will define an ``encoder`` and a ``decoder`` Module ahead of time\nin ``setup``, and use them in the ``encode`` and ``decode`` respectively.\n\n.. codediff::\n :title: Haiku, Flax\n :sync:\n\n class AutoEncoder(hk.Module):\n\n\n def __init__(self, embed_dim: int, output_dim: int, name=None):\n super().__init__(name=name)\n self.encoder = hk.Linear(embed_dim, name=\"encoder\")\n self.decoder = hk.Linear(output_dim, name=\"decoder\")\n\n def encode(self, x):\n return self.encoder(x)\n\n def decode(self, x):\n return self.decoder(x)\n\n def __call__(self, x):\n x = self.encode(x)\n x = self.decode(x)\n return x\n\n ---\n\n class AutoEncoder(nn.Module):\n embed_dim: int\n output_dim: int\n\n def setup(self):\n self.encoder = nn.Dense(self.embed_dim)\n self.decoder = nn.Dense(self.output_dim)\n\n def encode(self, x):\n return self.encoder(x)\n\n def decode(self, x):\n return self.decoder(x)\n\n def __call__(self, x):\n x = self.encode(x)\n x = self.decode(x)\n return x\n\nNote that in Flax ``setup`` doesn't run after ``__init__``, instead it runs\nwhen ``init`` or ``apply`` are called.\n\nNow, we want to be able to call any method from our ``AutoEncoder`` model. In Haiku we\ncan define multiple ``apply`` methods for a module through ``hk.multi_transform``. The\nfunction passed to ``multi_transform`` defines how to initialize the module and which\ndifferent apply methods to generate.\n\n.. codediff::\n :title: Haiku, Flax\n :sync:\n\n def forward():\n module = AutoEncoder(256, 784)\n init = lambda x: module(x)\n return init, (module.encode, module.decode)\n\n model = hk.multi_transform(forward)\n\n ---\n\n ...\n\n\n\n\n model = AutoEncoder(256, 784)\n\n\nTo initialize the parameters of our model, ``init`` can be used to trigger the\n``__call__`` method, which uses both the ``encode`` and ``decode``\nmethod. This will create all the necessary parameters for the model.\n\n.. codediff::\n :title: Haiku, Flax\n :sync:\n\n params = model.init(\n random.key(0),\n x=jax.numpy.ones((1, 784)),\n )\n ...\n\n ---\n\n variables = model.init(\n random.key(0),\n x=jax.numpy.ones((1, 784)),\n )\n params = variables[\"params\"]\n\nThis generates the following parameter structure.\n\n.. tab-set::\n\n .. tab-item:: Haiku\n :sync: Haiku\n\n .. code-block:: python\n\n {\n 'auto_encoder/~/decoder': {\n 'b': (784,),\n 'w': (256, 784)\n },\n 'auto_encoder/~/encoder': {\n 'b': (256,),\n 'w': (784, 256)\n }\n }\n\n .. tab-item:: Flax\n :sync: Flax\n\n .. code-block:: python\n\n FrozenDict({\n decoder: {\n bias: (784,),\n kernel: (256, 784),\n },\n encoder: {\n bias: (256,),\n kernel: (784, 256),\n },\n })\n\n\nFinally, let's explore how we can employ the ``apply`` function to invoke the ``encode`` method:\n\n.. codediff::\n :title: Haiku, Flax\n :sync:\n\n encode, decode = model.apply\n z = encode(\n params,\n None, # <== rng\n x=jax.numpy.ones((1, 784)),\n\n )\n\n ---\n\n ...\n z = model.apply(\n {\"params\": params},\n\n x=jax.numpy.ones((1, 784)),\n method=\"encode\",\n )\n\nBecause the Haiku ``apply`` function is generated through\n``hk.multi_transform``, it's a tuple of two functions which we can unpack into\nan ``encode`` and ``decode`` function which correspond to the methods on the\n``AutoEncoder`` module. In Flax we call the ``encode`` method through passing\nthe method name as a string.\nAnother noteworthy distinction here is that in Haiku, ``rng`` needs to be\nexplicitly passed, even though the module does not use any stochastic\noperations during ``apply``. In Flax this is not necessary (check out\n`Randomness and PRNGs in Flax `_).\nThe Haiku ``rng`` is set to ``None`` here, but you could also use\n``hk.without_apply_rng`` on the ``apply`` function to remove the ``rng`` argument.\n\n\nLifted Transforms\n-----------------\n\nBoth Flax and Haiku provide a set of transforms, which we will refer to as lifted transforms,\nthat wrap JAX transformations in such a way that they can be used with Modules and sometimes\nprovide additional functionality. In this section we will take a look at how to use the\nlifted version of ``scan`` in both Flax and Haiku to implement a simple RNN layer.\n\nTo begin, we will first define a ``RNNCell`` module that will contain the logic for a single\nstep of the RNN. We will also define a ``initial_state`` method that will be used to initialize\nthe state (a.k.a. ``carry``) of the RNN. Like with ``jax.lax.scan``, the ``RNNCell.__call__``\nmethod will be a function that takes the carry and input, and returns the new\ncarry and output. In this case, the carry and the output are the same.\n\n.. codediff::\n :title: Haiku, Flax\n :sync:\n\n class RNNCell(hk.Module):\n def __init__(self, hidden_size: int, name=None):\n super().__init__(name=name)\n self.hidden_size = hidden_size\n\n def __call__(self, carry, x):\n x = jnp.concatenate([carry, x], axis=-1)\n x = hk.Linear(self.hidden_size)(x)\n x = jax.nn.relu(x)\n return x, x\n\n def initial_state(self, batch_size: int):\n return jnp.zeros((batch_size, self.hidden_size))\n\n ---\n\n class RNNCell(nn.Module):\n hidden_size: int\n\n\n @nn.compact\n def __call__(self, carry, x):\n x = jnp.concatenate([carry, x], axis=-1)\n x = nn.Dense(self.hidden_size)(x)\n x = jax.nn.relu(x)\n return x, x\n\n def initial_state(self, batch_size: int):\n return jnp.zeros((batch_size, self.hidden_size))\n\nNext, we will define a ``RNN`` Module that will contain the logic for the entire RNN.\nIn Haiku, we will first initialze the ``RNNCell``, then use it to construct the ``carry``,\nand finally use ``hk.scan`` to run the ``RNNCell`` over the input sequence. In Flax its\ndone a bit differently, we will use ``nn.scan`` to define a new temporary type that wraps\n``RNNCell``. During this process we will also specify instruct ``nn.scan`` to broadcast\nthe ``params`` collection (all steps share the same parameters) and to not split the\n``params`` rng stream (so all steps intialize with the same parameters), and finally\nwe will specify that we want scan to run over the second axis of the input and stack\nthe outputs along the second axis as well. We will then use this temporary type immediately\nto create an instance of the lifted ``RNNCell`` and use it to create the ``carry`` and\nthe run the ``__call__`` method which will ``scan`` over the sequence.\n\n.. codediff::\n :title: Haiku, Flax\n :sync:\n\n class RNN(hk.Module):\n def __init__(self, hidden_size: int, name=None):\n super().__init__(name=name)\n self.hidden_size = hidden_size\n\n def __call__(self, x):\n cell = RNNCell(self.hidden_size)\n carry = cell.initial_state(x.shape[0])\n carry, y = hk.scan(cell, carry, jnp.swapaxes(x, 1, 0))\n y = jnp.swapaxes(y, 0, 1)\n return y\n\n ---\n\n class RNN(nn.Module):\n hidden_size: int\n\n\n @nn.compact\n def __call__(self, x):\n rnn = nn.scan(RNNCell, variable_broadcast='params', split_rngs={'params': False},\n in_axes=1, out_axes=1)(self.hidden_size)\n carry = rnn.initial_state(x.shape[0])\n carry, y = rnn(carry, x)\n return y\n\nIn general, the main difference between lifted transforms between Flax and Haiku is that\nin Haiku the lifted transforms don't operate over the state, that is, Haiku will handle the\n``params`` and ``state`` in such a way that it keeps the same shape inside and outside of the\ntransform. In Flax, the lifted transforms can operate over both variable collections and rng\nstreams, the user must define how different collections are treated by each transform\naccording to the transform's semantics.\n\nFinally, let's quickly view how the ``RNN`` Module would be used in both Haiku and Flax.\n\n.. codediff::\n :title: Haiku, Flax\n :sync:\n\n def forward(x):\n return RNN(64)(x)\n\n model = hk.without_apply_rng(hk.transform(forward))\n\n params = model.init(\n random.key(0),\n x=jax.numpy.ones((3, 12, 32)),\n )\n\n y = model.apply(\n params,\n x=jax.numpy.ones((3, 12, 32)),\n )\n\n ---\n\n ...\n\n\n model = RNN(64)\n\n variables = model.init(\n random.key(0),\n x=jax.numpy.ones((3, 12, 32)),\n )\n params = variables['params']\n y = model.apply(\n {'params': params},\n x=jax.numpy.ones((3, 12, 32)),\n )\n\nThe only notable change with respect to the examples in the previous sections is that\nthis time around we used ``hk.without_apply_rng`` in Haiku so we didn't have to\npass the ``rng`` argument as ``None`` to the ``apply`` method.\n\nScan over layers\n----------------\nOne very important application of ``scan`` is apply a sequence of layers iteratively\nover an input, passing the output of each layer as the input to the next layer. This\nis very useful to reduce compilation time for big models. As an example we will create\na simple ``Block`` Module, and then use it inside an ``MLP`` Module that will apply\nthe ``Block`` Module ``num_layers`` times.\n\nIn Haiku, we define the ``Block`` Module as usual, and then inside ``MLP`` we will\nuse ``hk.experimental.layer_stack`` over a ``stack_block`` function to create a stack\nof ``Block`` Modules. In Flax, the definition of ``Block`` is a little different,\n``__call__`` will accept and return a second dummy input/output that in both cases will\nbe ``None``. In ``MLP``, we will use ``nn.scan`` as in the previous example, but\nby setting ``split_rngs={'params': True}`` and ``variable_axes={'params': 0}``\nwe are telling ``nn.scan`` create different parameters for each step and slice the\n``params`` collection along the first axis, effectively implementing a stack of\n``Block`` Modules as in Haiku.\n\n\n.. codediff::\n :title: Haiku, Flax\n :sync:\n\n class Block(hk.Module):\n def __init__(self, features: int, name=None):\n super().__init__(name=name)\n self.features = features\n\n def __call__(self, x, training: bool):\n x = hk.Linear(self.features)(x)\n x = hk.dropout(hk.next_rng_key(), 0.5 if training else 0, x)\n x = jax.nn.relu(x)\n return x\n\n class MLP(hk.Module):\n def __init__(self, features: int, num_layers: int, name=None):\n super().__init__(name=name)\n self.features = features\n self.num_layers = num_layers\n\n def __call__(self, x, training: bool):\n @hk.experimental.layer_stack(self.num_layers)\n def stack_block(x):\n return Block(self.features)(x, training)\n\n stack = hk.experimental.layer_stack(self.num_layers)\n return stack_block(x)\n\n ---\n\n class Block(nn.Module):\n features: int\n training: bool\n\n @nn.compact\n def __call__(self, x, _):\n x = nn.Dense(self.features)(x)\n x = nn.Dropout(0.5)(x, deterministic=not self.training)\n x = jax.nn.relu(x)\n return x, None\n\n class MLP(nn.Module):\n features: int\n num_layers: int\n\n @nn.compact\n def __call__(self, x, training: bool):\n ScanBlock = nn.scan(\n Block, variable_axes={'params': 0}, split_rngs={'params': True},\n length=self.num_layers)\n\n y, _ = ScanBlock(self.features, training)(x, None)\n return y\n\nNotice how in Flax we pass ``None`` as the second argument to ``ScanBlock`` and ignore\nits second output. These represent the inputs/outputs per-step but they are ``None``\nbecause in this case we don't have any.\n\nInitializing each model is the same as in previous examples. In this case,\nwe will be specifying that we want to use ``5`` layers each with ``64`` features.\n\n.. codediff::\n :title: Haiku, Flax\n :sync:\n\n def forward(x, training: bool):\n return MLP(64, num_layers=5)(x, training)\n\n model = hk.transform(forward)\n\n sample_x = jax.numpy.ones((1, 64))\n params = model.init(\n random.key(0),\n sample_x, training=False # <== inputs\n )\n ...\n\n ---\n\n ...\n\n\n model = MLP(64, num_layers=5)\n\n sample_x = jax.numpy.ones((1, 64))\n variables = model.init(\n random.key(0),\n sample_x, training=False # <== inputs\n )\n params = variables['params']\n\nWhen using scan over layers the one thing you should notice is that all layers\nare fused into a single layer whose parameters have an extra \"layer\" dimension on\nthe first axis. In this case, the shape of all parameters will start with ``(5, ...)``\nas we are using ``5`` layers.\n\n.. tab-set::\n\n .. tab-item:: Haiku\n :sync: Haiku\n\n .. code-block:: python\n\n ...\n {\n 'mlp/__layer_stack_no_per_layer/block/linear': {\n 'b': (5, 64),\n 'w': (5, 64, 64)\n }\n }\n ...\n\n .. tab-item:: Flax\n :sync: Flax\n\n .. code-block:: python\n\n FrozenDict({\n ScanBlock_0: {\n Dense_0: {\n bias: (5, 64),\n kernel: (5, 64, 64),\n },\n },\n })\n\nTop-level Haiku functions vs top-level Flax modules\n-----------------------------------\n\nIn Haiku, it is possible to write the entire model as a single function by using the raw ``hk.{get,set}_{parameter,state}`` to define/access model parameters and states. It very common to write the top-level \"Module\" as a function instead:\n\nThe Flax team recommends a more Module-centric approach that uses `__call__` to define the forward function. The corresponding accessor will be `nn.module.param` and `nn.module.variable` (go to `Handling State <#handling-state>`__ for an explanaion on collections).\n\n.. codediff::\n :title: Haiku, Flax\n :sync:\n\n def forward(x):\n\n\n counter = hk.get_state('counter', shape=[], dtype=jnp.int32, init=jnp.ones)\n multiplier = hk.get_parameter('multiplier', shape=[1,], dtype=x.dtype, init=jnp.ones)\n output = x + multiplier * counter\n hk.set_state(\"counter\", counter + 1)\n\n return output\n\n model = hk.transform_with_state(forward)\n\n params, state = model.init(random.key(0), jax.numpy.ones((1, 64)))\n\n ---\n\n class FooModule(nn.Module):\n @nn.compact\n def __call__(self, x):\n counter = self.variable('counter', 'count', lambda: jnp.ones((), jnp.int32))\n multiplier = self.param('multiplier', nn.initializers.ones_init(), [1,], x.dtype)\n output = x + multiplier * counter.value\n if not self.is_initializing(): # otherwise model.init() also increases it\n counter.value += 1\n return output\n\n model = FooModule()\n variables = model.init(random.key(0), jax.numpy.ones((1, 64)))\n params, counter = variables['params'], variables['counter']" }, { "path": "docs/guides/converting_and_upgrading/index.rst", "content": "Converting and upgrading\n========================\n\n.. toctree::\n :maxdepth: 1\n\n haiku_migration_guide\n convert_pytorch_to_flax\n orbax_upgrade_guide\n optax_update_guide\n linen_upgrade_guide\n rnncell_upgrade_guide\n regular_dict_upgrade_guide" }, { "path": "docs/guides/converting_and_upgrading/linen_upgrade_guide.rst", "content": "Upgrading my codebase to Linen\n==============================\n\nAs of Flax v0.4.0, ``flax.nn`` no longer exists, and is replaced with the new\nLinen API at ``flax.linen``. If your codebase is still using the old API, you\ncan use this upgrade guide to upgrade it to Linen.\n\n.. testsetup:: Linen\n\n from flax.training import train_state\n from jax import random\n import optax\n import jax\n import flax.linen as nn\n from flax.linen import initializers\n\n from jax import lax\n import jax.numpy as jnp\n import numpy as np\n from typing import Any, Callable, Sequence, Tuple\n\n PRNGKey = Any\n Shape = Tuple[int, ...]\n Dtype = Any\n Array = Any\n\n default_kernel_init = initializers.lecun_normal()\n\nDefining simple Flax Modules\n----------------------------\n\n.. codediff::\n :title: Old Flax, Linen\n :skip_test: Old Flax\n :sync:\n\n from flax import nn\n\n class Dense(base.Module):\n def apply(self,\n inputs,\n features,\n use_bias=True,\n kernel_init=default_kernel_init,\n bias_init=initializers.zeros_init()):\n\n kernel = self.param('kernel',\n (inputs.shape[-1], features), kernel_init)\n y = jnp.dot(inputs, kernel)\n if use_bias:\n bias = self.param(\n 'bias', (features,), bias_init)\n y = y + bias\n return y\n ---\n from flax import linen as nn # [1] #!\n\n class Dense(nn.Module):\n features: int # [2] #!\n use_bias: bool = True\n kernel_init: Callable[[PRNGKey, Shape, Dtype], Array] = default_kernel_init\n bias_init: Callable[[PRNGKey, Shape, Dtype], Array] = initializers.zeros_init()\n\n @nn.compact\n def __call__(self, inputs): # [3] #!\n kernel = self.param('kernel',\n self.kernel_init, (inputs.shape[-1], self.features)) # [4] #!\n y = jnp.dot(inputs, kernel)\n if self.use_bias:\n bias = self.param(\n 'bias', self.bias_init, (self.features,)) # [5] #!\n y = y + bias\n return y\n\n1. Replace ``from flax import nn`` with ``from flax import linen as nn``.\n\n2. Move arguments to ``apply`` into dataclass attributes. Add type annotations\n (or use type ``Any`` to bypass).\n\n3. Rename method ``apply`` to ``__call__`` and (optionally) wrap with\n |@compact|_. Methods wrapped in |@compact|_ can define submodules directly\n within the method (like in old Flax). You can only wrap a single method with\n |@compact|_. Alternatively, you can define a ``setup`` method. For more\n details, please see our other HOWTO `Should I use setup or nn.compact?`_.\n\n4. Access dataclass attributes values by ``self.`` inside methods, e.g.\n ``self.features``.\n\n5. Move shape to the end of the arguments to |self.param|_ (initializer functions\n can take arbitrary argument lists).\n\n\nUsing Flax Modules inside other Modules\n---------------------------------------\n\n.. codediff::\n :title: Old Flax, Linen\n :skip_test: Old Flax\n :sync:\n\n class Encoder(nn.Module):\n\n def apply(self, x):\n x = nn.Dense(x, 500)\n x = nn.relu(x)\n z = nn.Dense(x, 500, name=\"latents\")\n return z\n ---\n class Encoder(nn.Module):\n @nn.compact\n def __call__(self, x):\n x = nn.Dense(500)(x) # [1] #!\n x = nn.relu(x)\n z = nn.Dense(500, name='latents')(x) # [2] #!\n return z\n\n1. Module constructors no longer return the outputs. Instead, they work like\n normal constructors and return module instances. These instances can be\n shared like in normal Python (instead of using ``.shared()`` in old Flax).\n Since most modules implement ``__call__``, you can retain the conciseness of\n old Flax.\n\n2. Names can be optionally passed to all module constructors.\n\nSharing submodules and defining multiple methods\n--------------------------------\n\n.. codediff::\n :title: Old Flax, Linen\n :skip_test: Old Flax\n :sync:\n\n class AutoEncoder(nn.Module):\n def _create_submodules(self):\n return Decoder.shared(name=\"encoder\")\n\n def apply(self, x, z_rng, latents=20):\n decoder = self._create_decoder()\n z = Encoder(x, latents, name=\"encoder\")\n return decoder(z)\n\n @nn.module_method\n def generate(self, z, **unused_kwargs):\n decoder = self._create_decoder()\n return nn.sigmoid(decoder(z))\n ---\n class AutoEncoder(nn.Module):\n latents: int = 20\n\n def setup(self): # [1] #!\n self.encoder = Encoder(self.latents) # [2] #!\n self.decoder = Decoder()\n\n def __call__(self, x): # [3] #!\n z = self.encoder(x)\n return self.decoder(z)\n\n def generate(self, z): # [4] #!\n return nn.sigmoid(self.decoder(z))\n\n\n1. Use |setup|_ instead of ``__init__``, which is already defined in\n the dataclasses library. Flax calls setup right after modules are ready to be\n used. (You can do this for all modules if you like instead of using\n |@compact|, but we like how |@compact| co-locates where modules are defined\n and used, especially if you have loops or conditionals).\n\n2. Like regular Python, share submodules by assigning to self during\n initialization. Similar to PyTorch, ``self.encoder`` automatically has the\n name ``\"encoder\"``.\n\n3. We don't use |@compact|_ here because we're not defining any inline\n submodules (all submodules are defined in setup).\n\n4. Define additional methods just like in regular Python.\n\n``Module.partial`` inside other modules\n---------------------------------------\n\n.. codediff::\n :title: Old Flax, Linen\n :skip_test: Old Flax\n :sync:\n\n # no import #!\n\n class ResNet(nn.Module):\n \"\"\"ResNetV1.\"\"\"\n\n\n def apply(self, x,\n stage_sizes,\n num_filters=64,\n train=True):\n conv = nn.Conv.partial(bias=False)\n norm = nn.BatchNorm.partial(\n use_running_average=not train,\n momentum=0.9, epsilon=1e-5)\n\n x = conv(x, num_filters, (7, 7), (2, 2),\n padding=[(3, 3), (3, 3)],\n name='conv_init')\n x = norm(x, name='bn_init')\n\n # [...]\n return x\n ---\n from functools import partial #!\n\n class ResNet(nn.Module):\n \"\"\"ResNetV1.\"\"\"\n stage_sizes: Sequence[int]\n num_filters: int = 64\n train: bool = True\n\n @nn.compact\n def __call__(self, x):\n conv = partial(nn.Conv, use_bias=False) #!\n norm = partial(nn.BatchNorm, #!\n use_running_average=not self.train, #!\n momentum=0.9, epsilon=1e-5) #!\n\n x = conv(self.num_filters, (7, 7), (2, 2),\n padding=[(3, 3), (3, 3)],\n name='conv_init')(x)\n x = norm(name='bn_init')(x)\n\n # [...]\n return x\n\nUse normal ``functools.partial`` instead of ``Module.partial``. The rest stays\nthe same.\n\nTop-level training code patterns\n--------------------------------\n\n.. codediff::\n :title: Old Flax, Linen\n :skip_test: Old Flax\n :sync:\n\n def create_model(key):\n _, initial_params = CNN.init_by_shape(\n key, [((1, 28, 28, 1), jnp.float32)])\n model = nn.Model(CNN, initial_params)\n return model\n\n def create_optimizer(model, learning_rate):\n optimizer_def = optim.Momentum(learning_rate=learning_rate)\n optimizer = optimizer_def.create(model)\n return optimizer\n\n def cross_entropy_loss(*, logits, labels):\n one_hot_labels = jax.nn.one_hot(labels, num_classes=10)\n return -jnp.mean(jnp.sum(one_hot_labels * logits, axis=-1))\n\n def loss_fn(model):\n logits = model(batch['image'])\n one_hot = jax.nn.one_hot(batch['label'], num_classes=10)\n loss = -jnp.mean(jnp.sum(one_hot_labels * batch['label'],\n axis=-1))\n return loss, logits\n ---\n def create_train_state(rng, config): # [1] #!\n variables = CNN().init(rng, jnp.ones([1, 28, 28, 1])) # [2] #!\n params = variables['params'] # [3] #!\n tx = optax.sgd(config.learning_rate, config.momentum) # [4] #!\n return train_state.TrainState.create(\n apply_fn=CNN.apply, params=params, tx=tx)\n\n\n def loss_fn(params):\n logits = CNN().apply({'params': params}, batch['image']) # [5] #!\n one_hot = jax.nn.one_hot(batch['label'], 10)\n loss = jnp.mean(optax.softmax_cross_entropy(logits=logits,\n labels=one_hot))\n return loss, logits\n\n\n1. We no longer use the ``Model`` abstraction -- instead we pass parameters\n around directly, usually encapsulated in a `TrainState`_ object, which can\n directly be passed to JAX transformations.\n\n2. To compute initial parameters, construct a module instance and call |init|_\n or |init_with_output|_. We haven't ported over ``init_by_shape`` because this\n function did some magic we did not like (it evaluated the function by shape.\n but returned real values anyway). Therefore, you should now pass concrete\n values to the initializer functions, and you can optimize the initialization\n by wrapping it with |jax.jit|_, which is highly recommended to avoid running\n a full forward pass.\n\n3. Linen generalizes parameters into variables. Parameters are one\n \"collection\" of variables. Variables are nested dicts, where the top-level\n keys reflect the different variable collections, of which \"param\" is one of.\n See the `Variables documentation`_ for more details.\n\n4. We recommend using Optax optimizers. See our separate HOWTO called\n `Upgrading my codebase to Optax`_ for more details.\n\n5. To make predictions with your model, make an instance at the top level (this\n is free -- just a wrapper around constructor attributes) and call the\n ``apply`` method (which will call ``__call__`` internally).\n\nNon-trainable variables (\"state\"): Use within Modules\n-----------------------------------------------------\n\n.. codediff::\n :title: Old Flax, Linen\n :skip_test: Old Flax\n :sync:\n\n class BatchNorm(nn.Module):\n def apply(self, x):\n # [...]\n ra_mean = self.state(\n 'mean', (x.shape[-1], ), initializers.zeros_init())\n ra_var = self.state(\n 'var', (x.shape[-1], ), initializers.ones_init())\n # [...]\n ---\n class BatchNorm(nn.Module):\n def __call__(self, x):\n # [...]\n ra_mean = self.variable( #!\n 'batch_stats', 'mean', initializers.zeros_init(), (x.shape[-1], ))\n ra_var = self.variable(\n 'batch_stats', 'var', initializers.ones_init(), (x.shape[-1], ))\n # [...]\n\nThe first argument is the name of the variable collection (\"param\" is the only\nvariable collection that's always available). Some colllections may be treated\nas mutable, and others as immutable at top-level training code (see next section\nfor details). Flax also lets you treat each variable collection differently when\nusing JAX transformations inside modules.\n\nNon-trainable variables (\"state\"): Top-level training code patterns\n-------------------------------------------------------------------\n\n.. codediff::\n :title: Old Flax, Linen\n :skip_test: Old Flax\n :sync:\n\n # initial params and state\n def initial_model(key, init_batch):\n with nn.stateful() as initial_state:\n _, initial_params = ResNet.init(key, init_batch)\n model = nn.Model(ResNet, initial_params)\n return model, init_state\n\n\n # updates batch statistics during training\n def loss_fn(model, model_state):\n with nn.stateful(model_state) as new_model_state:\n logits = model(batch['image'])\n # [...]\n\n\n\n # reads immutable batch statistics during evaluation\n def eval_step(model, model_state, batch):\n with nn.stateful(model_state, mutable=False):\n logits = model(batch['image'], train=False)\n return compute_metrics(logits, batch['label'])\n ---\n # initial variables ({\"param\": ..., \"batch_stats\": ...})\n def initial_variables(key, init_batch):\n return ResNet().init(key, init_batch) # [1] #!\n\n\n\n # updates batch statistics during training\n def loss_fn(params, batch_stats):\n variables = {'params': params, 'batch_stats': batch_stats} # [2] #!\n logits, new_variables = ResNet(train=true).apply(\n variables, batch['image'], mutable=['batch_stats']) # [3] #!\n new_batch_stats = new_variables['batch_stats']\n # [...]\n\n\n # reads immutable batch statistics during evaluation\n def eval_step(params, batch_stats, batch):\n variables = {'params': params, 'batch_stats': batch_stats}\n logits = ResNet(train=False).apply(\n variables, batch['image'], mutable=False) # [4] #!\n return compute_metrics(logits, batch['label'])\n\n1. |init|_ returns a variable dict, e.g. ``{\"param\": ..., \"batch_stats\": ...}``\n (see `Variables documentation`_).\n\n2. Combine the different variable collections into a variable dict.\n\n3. During training, the ``batch_stats`` variable collection changes. Since we\n specify that in the mutable argument, the return value from ``module.apply``\n becomes an ordered pair of ``output, new_variables``.\n\n4. During evaluation, we want to raise an error if we're accidentally applying\n Batch Norm in training mode. By passing ``mutable=False`` into\n ``module.apply`` we enforce that. Since no variables are mutated, the return\n value is once again just the output.\n\nLoading pre-Linen checkpoints\n-----------------------------\n\nWhile most Linen modules should be able to use pre-Linen weights without any\nmodification, there is one catch: In pre-Linen API submodules were numbered\nincrementally, independent of the submodule class. With Linen this behavior has\nchanged to keep separate submodule counts per module class.\n\nIn pre-Linen, params have the following structure:\n\n``{'Conv_0': { ... }, 'Dense_1': { ... } }``\n\nIn Linen this is instead:\n\n``{'Conv_0': { ... }, 'Dense_0': { ... } }``\n\nTODO: Add an example here how to load a new ``TrainState`` object.\n\nRandomness\n----------\n\n.. codediff::\n :title: Old Flax, Linen\n :skip_test: Old Flax\n :sync:\n\n def dropout(inputs, rate, deterministic=False):\n keep_prob = 1. - rate\n if deterministic:\n return inputs\n else:\n mask = random.bernoulli(\n make_rng(), p=keep_prob, shape=inputs.shape)\n return lax.select(\n mask, inputs / keep_prob, jnp.zeros_like(inputs))\n\n\n def loss_fn(model, dropout_rng):\n with nn.stochastic(dropout_rng):\n logits = model(inputs)\n ---\n class Dropout(nn.Module):\n rate: float\n\n @nn.compact\n def __call__(self, inputs, deterministic=False):\n keep_prob = 1. - self.rate\n if deterministic:\n return inputs\n else:\n mask = random.bernoulli(\n self.make_rng('dropout'), p=keep_prob, shape=inputs.shape) # [1] #!\n return lax.select(\n mask, inputs / keep_prob, jnp.zeros_like(inputs))\n\n\n def loss_fn(params, dropout_rng):\n logits = Transformer().apply(\n {'params': params}, inputs, rngs={'dropout': dropout_rng}) # [2] #!\n\n1. RNGs in Linen have \"kinds\" -- in this case ``'dropout'``. Different kinds can\n be treated different in JAX transformations (for example, do you want the\n same dropout mask for each timestep in a sequence model or a different one?)\n\n2. Instead of using the ``nn.stochastic`` context manager, you pass in RNGs\n explicitly to ``module.apply``. During evaluation you wouldn't pass any RNGs\n -- then if you accidentally use dropout in non-deterministic mode,\n ``self.make_rng('dropout')`` would raise an error.\n\n\nLifted transformations\n----------------------\n\nIn Linen, rather than using JAX transformation directly, we are using\n\"lifted transforms\", which are JAX transformations applied to Flax Modules.\n\nFor more information, please see the design note on `Lifted transformations`_.\n\nTODO: Given an example of ``jax.scan_in_dim`` (pre-Linen) vs. ``nn.scan``\n(Linen).\n\n.. _`Should I use setup or nn.compact?`: https://flax.readthedocs.io/en/latest/guides/flax_fundamentals/setup_or_nncompact.html\n.. _`Variables documentation`: https://flax.readthedocs.io/en/latest/api_reference/flax.linen/variable.html\n.. _`TrainState`: https://flax.readthedocs.io/en/latest/flax.training.html#train-state\n.. _`Upgrading my codebase to Optax`: https://flax.readthedocs.io/en/latest/guides/converting_and_upgrading/optax_update_guide.html\n.. _`Lifted transformations`: https://flax.readthedocs.io/en/latest/developer_notes/lift.html\n\n\n.. |@compact| replace:: ``@compact``\n.. _@compact: https://flax.readthedocs.io/en/latest/api_reference/flax.linen/decorators.html#flax.linen.compact\n\n.. |init| replace:: ``init``\n.. _init: https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.init\n\n.. |init_with_output| replace:: ``init_with_output``\n.. _init_with_output: https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.init_with_output\n\n.. |jax.jit| replace:: ``jax.jit``\n.. _jax.jit: https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html#jax.jit\n\n.. |self.param| replace:: ``self.param``\n.. _self.param: https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.param\n\n.. |setup| replace:: ``setup``\n.. _setup: https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.setup\n\n.. |@flax.struct.dataclass| replace:: ``@flax.struct.dataclass``\n.. _@flax.struct.dataclass: https://flax.readthedocs.io/en/latest/flax.struct.html#flax.struct.dataclass\n\n.. |checkpoints.convert_pre_linen()| replace:: ``checkpoints.convert_pre_linen()``\n.. _checkpoints.convert_pre_linen(): https://flax.readthedocs.io/en/latest/flax.training.html#flax.training.checkpoints.convert_pre_linen\n" }, { "path": "docs/guides/converting_and_upgrading/optax_update_guide.rst", "content": "Upgrading my codebase to Optax\n==============================\n\nWe have proposed to replace :py:mod:`flax.optim` with `Optax\n`_ in 2021 with `FLIP #1009\n`_ and\nthe Flax optimizers have been removed in v0.6.0 - this guide is targeted\ntowards :py:mod:`flax.optim` users to help them update their code to Optax.\n\nSee also Optax's quick start documentation:\nhttps://optax.readthedocs.io/en/latest/getting_started.html\n\n.. testsetup:: default, flax.optim, optax\n\n import flax\n import jax\n import jax.numpy as jnp\n import flax.linen as nn\n import optax\n\n # Note: this is the minimal code required to make below code run. See in the\n # Colab linked above for a more meaningful definition of datasets etc.\n batch = {'image': jnp.ones([1, 28, 28, 1]), 'label': jnp.array([0])}\n ds_train = [batch]\n get_ds_train = lambda: [batch]\n model = nn.Dense(1)\n variables = model.init(jax.random.key(0), batch['image'])\n learning_rate, momentum, weight_decay, grad_clip_norm = .1, .9, 1e-3, 1.\n loss = lambda params, batch: jnp.array(0.)\n\nReplacing ``flax.optim`` with ``optax``\n---------------------------------------\n\nOptax has drop-in replacements for all of Flax's optimizers. Refer to Optax's\ndocumentation `Common Optimizers `_\nfor API details.\n\nThe usage is very similar, with the difference that ``optax`` does not keep a\ncopy of the ``params``, so they need to be passed around separately. Flax\nprovides the utility :py:class:`~flax.training.train_state.TrainState` to store\noptimizer state, parameters, and other associated data in a single dataclass\n(not used in code below).\n\n.. codediff::\n :title: flax.optim, optax\n :skip_test: flax.optim\n :sync:\n\n @jax.jit\n def train_step(optimizer, batch):\n grads = jax.grad(loss)(optimizer.target, batch)\n\n\n return optimizer.apply_gradient(grads)\n\n optimizer_def = flax.optim.Momentum(\n learning_rate, momentum)\n optimizer = optimizer_def.create(variables['params'])\n\n for batch in get_ds_train():\n optimizer = train_step(optimizer, batch)\n\n ---\n\n @jax.jit\n def train_step(params, opt_state, batch):\n grads = jax.grad(loss)(params, batch)\n updates, opt_state = tx.update(grads, opt_state)\n params = optax.apply_updates(params, updates)\n return params, opt_state\n\n tx = optax.sgd(learning_rate, momentum)\n params = variables['params']\n opt_state = tx.init(params)\n\n for batch in ds_train:\n params, opt_state = train_step(params, opt_state, batch)\n\n\nComposable Gradient Transformations\n-----------------------------------\n\nThe function |optax.sgd()|_ used in the code snippet above is simply a wrapper\nfor the sequential application of two gradient transformations. Instead of using\nthis alias, it is common to use |optax.chain()|_ to combine multiple of these\ngeneric building blocks.\n\n.. |optax.sgd()| replace:: ``optax.sgd()``\n.. _optax.sgd(): https://optax.readthedocs.io/en/latest/api/optimizers.html#optax.sgd\n.. |optax.chain()| replace:: ``optax.chain()``\n.. _optax.chain(): https://optax.readthedocs.io/en/latest/api/combining_optimizers.html#optax.chain\n\n.. codediff::\n :title: Pre-defined alias, Combining transformations\n :groups: default, default\n\n # Note that the aliases follow the convention to use positive\n # values for the learning rate by default.\n tx = optax.sgd(learning_rate, momentum)\n\n ---\n\n #\n\n tx = optax.chain(\n # 1. Step: keep a trace of past updates and add to gradients.\n optax.trace(decay=momentum),\n # 2. Step: multiply result from step 1 with negative learning rate.\n # Note that `optax.apply_updates()` simply adds the final updates to the\n # parameters, so we must make sure to flip the sign here for gradient\n # descent.\n optax.scale(-learning_rate),\n )\n\nWeight Decay\n------------\n\nSome of Flax's optimizers also include a weight decay. In Optax, some optimizers\nalso have a weight decay parameter (such as |optax.adamw()|_), and to others the\nweight decay can be added as another \"gradient transformation\"\n|optax.add_decayed_weights()|_ that adds an update derived from the parameters.\n\n.. |optax.adamw()| replace:: ``optax.adamw()``\n.. _optax.adamw(): https://optax.readthedocs.io/en/latest/api/optimizers.html#optax.adamw\n.. |optax.add_decayed_weights()| replace:: ``optax.add_decayed_weights()``\n.. _optax.add_decayed_weights(): https://optax.readthedocs.io/en/latest/api/transformations.html#optax.add_decayed_weights\n\n.. codediff::\n :title: flax.optim, optax\n :skip_test: flax.optim\n :sync:\n\n optimizer_def = flax.optim.Adam(\n learning_rate, weight_decay=weight_decay)\n optimizer = optimizer_def.create(variables['params'])\n\n ---\n\n # (Note that you could also use `optax.adamw()` in this case)\n tx = optax.chain(\n optax.scale_by_adam(),\n optax.add_decayed_weights(weight_decay),\n # params -= learning_rate * (adam(grads) + params * weight_decay)\n optax.scale(-learning_rate),\n )\n # Note that you'll need to specify `params` when computing the udpates:\n # tx.update(grads, opt_state, params)\n\nGradient Clipping\n-----------------\n\nTraining can be stabilized by clipping gradients to a global norm (`Pascanu et\nal, 2012 `_). In Flax this is often done by\nprocessing the gradients before passing them to the optimizer. With Optax this\nbecomes just another gradient transformation |optax.clip_by_global_norm()|_.\n\n.. |optax.clip_by_global_norm()| replace:: ``optax.clip_by_global_norm()``\n.. _optax.clip_by_global_norm(): https://optax.readthedocs.io/en/latest/api/transformations.html#optax.clip_by_global_norm\n\n.. codediff::\n :title: flax.optim, optax\n :skip_test: flax.optim\n :sync:\n\n def train_step(optimizer, batch):\n grads = jax.grad(loss)(optimizer.target, batch)\n grads_flat, _ = jax.tree_util.tree_flatten(grads)\n global_l2 = jnp.sqrt(sum([jnp.vdot(p, p) for p in grads_flat]))\n g_factor = jnp.minimum(1.0, grad_clip_norm / global_l2)\n grads = jax.tree_util.tree_map(lambda g: g * g_factor, grads)\n return optimizer.apply_gradient(grads)\n\n ---\n\n tx = optax.chain(\n optax.clip_by_global_norm(grad_clip_norm),\n optax.trace(decay=momentum),\n optax.scale(-learning_rate),\n )\n\nLearning Rate Schedules\n-----------------------\n\nFor learning rate schedules, Flax allows overwriting hyper parameters when\napplying the gradients. Optax maintains a step counter and provides this as an\nargument to a function for scaling the updates added with\n|optax.scale_by_schedule()|_. Optax also allows specifying a functions to\ninject arbitrary scalar values for other gradient updates via\n|optax.inject_hyperparams()|_.\n\nRead more about learning rate schedules in the :doc:`lr_schedule` guide.\n\nRead more about schedules defined in Optax under `Optimizer Schedules\n`_. the\nstandard optimizers (like ``optax.adam()``, ``optax.sgd()`` etc.) also accept a\nlearning rate schedule as a parameter for ``learning_rate``.\n\n\n.. |optax.scale_by_schedule()| replace:: ``optax.scale_by_schedule()``\n.. _optax.scale_by_schedule(): https://optax.readthedocs.io/en/latest/api/transformations.html#optax.scale_by_schedule\n.. |optax.inject_hyperparams()| replace:: ``optax.inject_hyperparams()``\n.. _optax.inject_hyperparams(): https://optax.readthedocs.io/en/latest/api/optimizer_schedules.html#optax.inject_hyperparams\n\n.. codediff::\n :title: flax.optim, optax\n :skip_test: flax.optim\n :sync:\n\n def train_step(step, optimizer, batch):\n grads = jax.grad(loss)(optimizer.target, batch)\n return step + 1, optimizer.apply_gradient(grads, learning_rate=schedule(step))\n\n ---\n\n tx = optax.chain(\n optax.trace(decay=momentum),\n # Note that we still want a negative value for scaling the updates!\n optax.scale_by_schedule(lambda step: -schedule(step)),\n )\n\nMultiple Optimizers / Updating a Subset of Parameters\n-----------------------------------------------------\n\nIn Flax, traversals are used to specify which parameters should be updated by an\noptimizer. And you can combine traversals using\n:py:class:`flax.optim.MultiOptimizer` to apply different optimizers on different\nparameters. The equivalent in Optax is |optax.masked()|_ and |optax.chain()|_.\n\nNote that the example below is using :py:mod:`flax.traverse_util` to create the\nboolean masks required by |optax.masked()|_ - alternatively you could also\ncreate them manually, or use |optax.multi_transform()|_ that takes a\nmultivalent pytree to specify gradient transformations.\n\nBeware that |optax.masked()|_ flattens the pytree internally and the inner\ngradient transformations will only be called with that partial flattened view of\nthe params/gradients. This is not a problem usually, but it makes it hard to\nnest multiple levels of masked gradient transformations (because the inner\nmasks will expect the mask to be defined in terms of the partial flattened view\nthat is not readily available outside the outer mask).\n\n.. |optax.masked()| replace:: ``optax.masked()``\n.. _optax.masked(): https://optax.readthedocs.io/en/latest/api/optimizer_wrappers.html#optax.masked\n.. |optax.multi_transform()| replace:: ``optax.multi_transform()``\n.. _optax.multi_transform(): https://optax.readthedocs.io/en/latest/api/combining_optimizers.html#optax.multi_transform\n\n.. codediff::\n :title: flax.optim, optax\n :skip_test: flax.optim\n :sync:\n\n kernels = flax.traverse_util.ModelParamTraversal(lambda p, _: 'kernel' in p)\n biases = flax.traverse_util.ModelParamTraversal(lambda p, _: 'bias' in p)\n\n kernel_opt = flax.optim.Momentum(learning_rate, momentum)\n bias_opt = flax.optim.Momentum(learning_rate * 0.1, momentum)\n\n\n optimizer = flax.optim.MultiOptimizer(\n (kernels, kernel_opt),\n (biases, bias_opt)\n ).create(variables['params'])\n\n ---\n\n kernels = flax.traverse_util.ModelParamTraversal(lambda p, _: 'kernel' in p)\n biases = flax.traverse_util.ModelParamTraversal(lambda p, _: 'bias' in p)\n\n all_false = jax.tree_util.tree_map(lambda _: False, params)\n kernels_mask = kernels.update(lambda _: True, all_false)\n biases_mask = biases.update(lambda _: True, all_false)\n\n tx = optax.chain(\n optax.trace(decay=momentum),\n optax.masked(optax.scale(-learning_rate), kernels_mask),\n optax.masked(optax.scale(-learning_rate * 0.1), biases_mask),\n )\n\nFinal Words\n-----------\n\nAll above patterns can of course also be mixed and Optax makes it possible to\nencapsulate all these transformations into a single place outside the main\ntraining loop, which makes testing much easier.\n" }, { "path": "docs/guides/converting_and_upgrading/orbax_upgrade_guide.rst", "content": "Migrate checkpointing to Orbax\n==============================\n\nThis guide shows how to convert Flax's checkpoint saving and restoring calls — `flax.training.checkpoints.save_checkpoint `__ and `restore_checkpoint `__ — to the equivalent `Orbax `__ methods. Orbax provides a flexible and customizable API for managing checkpoints for various objects. Note that as Flax's checkpointing is being migrated to Orbax from ``flax.training.checkpoints``, all existing features in the Flax API will continue to be supported, but the API will change.\n\nYou will learn how to migrate to Orbax through the following scenarios:\n\n* The most common use case: Saving/loading and managing checkpoints\n* A \"lightweight\" use case: \"Pure\" saving/loading without the top-level checkpoint manager\n* Restoring checkpoints without a target pytree\n* Async checkpointing\n* Saving/loading a single JAX or NumPy Array\n\nTo learn more about Orbax, check out the `quick start introductory Colab notebook `__ and `the official Orbax documentation `_.\n\nYou can click on \"Open in Colab\" above to run the code from this guide.\n\nThroughout the guide, you will be able to compare code examples with and without the Orbax code.\n\n.. testsetup:: orbax.checkpoint\n\n import flax\n from flax.training import checkpoints, orbax_utils\n import orbax\n import jax\n import jax.numpy as jnp\n import numpy as np\n\n # Orbax needs to have asyncio enabled in the Colab environment.\n import nest_asyncio\n nest_asyncio.apply()\n\n # Set up the directory.\n import os\n import shutil\n if os.path.exists('/tmp/orbax_upgrade'):\n shutil.rmtree('/tmp/orbax_upgrade')\n os.makedirs('/tmp/orbax_upgrade')\n\n\nSetup\n*****\n\n.. testcode:: orbax.checkpoint\n\n # Create some dummy variables for this example.\n MAX_STEPS = 5\n CKPT_PYTREE = [12, {'bar': np.array((2, 3))}, [1, 4, 10]]\n TARGET_PYTREE = [0, {'bar': np.array((0))}, [0, 0, 0]]\n\nMost common use case: Saving/loading and managing checkpoints\n*************************************************************\n\nThis section covers the following scenario:\n\n* Your original Flax ``save_checkpoint()`` or ``save_checkpoint_multiprocess()`` call contains the following arguments: ``prefix``, ``keep``, ``keep_every_n_steps``; or\n* You want to use some automatic management logic for your checkpoints (for example, for deleting old data, deleting data based on metrics/loss, and so on).\n\nIn this case, you need to use ``orbax.CheckpointManager``. This allows you to not only save and load your model, but also manage your checkpoints and delete outdated checkpoints *automatically*.\n\nTo upgrade your code:\n\n1. Create and keep an ``orbax.CheckpointManager`` instance at the top level, customized with ``orbax.CheckpointManagerOptions``.\n\n2. At runtime, call ``orbax.CheckpointManager.save()`` to save your data.\n\n3. Then, call ``orbax.CheckpointManager.restore()`` to restore your data.\n\n4. And, if your checkpoint includes some multi-host/multi-process array, pass the correct ``mesh`` into ``flax.training.orbax_utils.restore_args_from_target()`` to generate the correct ``restore_args`` before restoring.\n\nFor example:\n\n.. codediff::\n :title: flax.checkpoints, orbax.checkpoint\n :skip_test: flax.checkpoints\n :sync:\n\n CKPT_DIR = '/tmp/orbax_upgrade/'\n flax.config.update('flax_use_orbax_checkpointing', False)\n\n # Inside your training loop\n for step in range(MAX_STEPS):\n # do training\n checkpoints.save_checkpoint(CKPT_DIR, CKPT_PYTREE, step=step,\n prefix='test_', keep=3, keep_every_n_steps=2)\n\n\n checkpoints.restore_checkpoint(CKPT_DIR, target=TARGET_PYTREE, step=4, prefix='test_')\n\n ---\n\n CKPT_DIR = '/tmp/orbax_upgrade/orbax'\n\n # At the top level\n mgr_options = orbax.checkpoint.CheckpointManagerOptions(\n create=True, max_to_keep=3, keep_period=2, step_prefix='test')\n ckpt_mgr = orbax.checkpoint.CheckpointManager(\n CKPT_DIR,\n orbax.checkpoint.Checkpointer(orbax.checkpoint.PyTreeCheckpointHandler()), mgr_options)\n\n # Inside your training loop\n for step in range(MAX_STEPS):\n # do training\n save_args = flax.training.orbax_utils.save_args_from_target(CKPT_PYTREE)\n ckpt_mgr.save(step, CKPT_PYTREE, save_kwargs={'save_args': save_args})\n\n\n restore_args = flax.training.orbax_utils.restore_args_from_target(TARGET_PYTREE, mesh=None)\n ckpt_mgr.restore(4, items=TARGET_PYTREE, restore_kwargs={'restore_args': restore_args})\n\n\nA \"lightweight\" use case: \"Pure\" saving/loading without the top-level checkpoint manager\n****************************************************************************************\n\nIf you prefer to not maintain a top-level checkpoint manager, you can still save and restore any individual checkpoint with an ``orbax.checkpoint.Checkpointer``. Note that this means you cannot use all the Orbax management features.\n\nTo migrate to Orbax code, instead of using the ``overwrite`` argument in ``flax.save_checkpoint()`` use the ``force`` argument in ``orbax.checkpoint.Checkpointer.save()``.\n\nFor example:\n\n.. codediff::\n :title: flax.checkpoints, orbax.checkpoint\n :skip_test: flax.checkpoints\n :sync:\n\n PURE_CKPT_DIR = '/tmp/orbax_upgrade/pure'\n flax.config.update('flax_use_orbax_checkpointing', False)\n\n checkpoints.save_checkpoint(PURE_CKPT_DIR, CKPT_PYTREE, step=0, overwrite=True)\n checkpoints.restore_checkpoint(PURE_CKPT_DIR, target=TARGET_PYTREE)\n\n ---\n\n PURE_CKPT_DIR = '/tmp/orbax_upgrade/pure'\n\n ckptr = orbax.checkpoint.Checkpointer(orbax.checkpoint.PyTreeCheckpointHandler()) # A stateless object, can be created on the fly.\n ckptr.save(PURE_CKPT_DIR, CKPT_PYTREE,\n save_args=flax.training.orbax_utils.save_args_from_target(CKPT_PYTREE), force=True)\n ckptr.restore(PURE_CKPT_DIR, item=TARGET_PYTREE,\n restore_args=flax.training.orbax_utils.restore_args_from_target(TARGET_PYTREE, mesh=None))\n\n\n\nRestoring checkpoints without a target pytree\n*********************************************\n\nIf you need to restore your checkpoints without a target pytree, pass ``item=None`` to ``orbax.checkpoint.Checkpointer`` or ``items=None`` to ``orbax.CheckpointManager``'s ``.restore()`` method, which should trigger the restoration.\n\nFor example:\n\n.. codediff::\n :title: flax.checkpoints, orbax.checkpoint\n :skip_test: flax.checkpoints\n :sync:\n\n NOTARGET_CKPT_DIR = '/tmp/orbax_upgrade/no_target'\n flax.config.update('flax_use_orbax_checkpointing', False)\n\n checkpoints.save_checkpoint(NOTARGET_CKPT_DIR, CKPT_PYTREE, step=0)\n checkpoints.restore_checkpoint(NOTARGET_CKPT_DIR, target=None)\n\n ---\n\n NOTARGET_CKPT_DIR = '/tmp/orbax_upgrade/no_target'\n\n # A stateless object, can be created on the fly.\n ckptr = orbax.checkpoint.Checkpointer(orbax.checkpoint.PyTreeCheckpointHandler())\n ckptr.save(NOTARGET_CKPT_DIR, CKPT_PYTREE,\n save_args=flax.training.orbax_utils.save_args_from_target(CKPT_PYTREE))\n ckptr.restore(NOTARGET_CKPT_DIR, item=None)\n\n\nAsync checkpointing\n*******************\n\nTo make your checkpoint-saving asynchronous, substitute ``orbax.checkpoint.Checkpointer`` with ``orbax.checkpoint.AsyncCheckpointer``.\n\nThen, you can call ``orbax.checkpoint.AsyncCheckpointer.wait_until_finished()`` or Orbax's ``CheckpointerManager.wait_until_finished()`` to wait for the save the complete.\n\nFor more details, read the `checkpoint guide `_.\n\nYou can also use Orbax AsyncCheckpointer with Flax APIs through async manager. Async manager internally calls wait_until_finished(). This solution is not actively maintained and the recommedation is to use Orbax async checkpointing.\n\nFor example:\n\n.. codediff::\n :title: flax.checkpoints, orbax.checkpoint\n :skip_test: flax.checkpoints\n :sync:\n\n ASYNC_CKPT_DIR = '/tmp/orbax_upgrade/async'\n flax.config.update('flax_use_orbax_checkpointing', True)\n async_manager = checkpoints.AsyncManager()\n\n checkpoints.save_checkpoint(ASYNC_CKPT_DIR, CKPT_PYTREE, step=0, overwrite=True, async_manager=async_manager)\n checkpoints.restore_checkpoint(ASYNC_CKPT_DIR, target=TARGET_PYTREE)\n ---\n\n ASYNC_CKPT_DIR = '/tmp/orbax_upgrade/async'\n\n import orbax.checkpoint as ocp\n ckptr = ocp.AsyncCheckpointer(ocp.StandardCheckpointHandler())\n ckptr.save(ASYNC_CKPT_DIR, args=ocp.args.StandardSave(CKPT_PYTREE))\n # ... Continue with your work...\n # ... Until a time when you want to wait until the save completes:\n ckptr.wait_until_finished() # Blocks until the checkpoint saving is completed.\n ckptr.restore(ASYNC_CKPT_DIR, args=ocp.args.StandardRestore(TARGET_PYTREE))\n\n\nSaving/loading a single JAX or NumPy Array\n******************************************\n\nThe ``orbax.checkpoint.PyTreeCheckpointHandler`` class, as the name suggests, can only be used for pytrees. Therefore, if you need to save/restore a single pytree leaf (for example, an array), use ``orbax.checkpoint.ArrayCheckpointHandler`` instead.\n\nFor example:\n\n.. codediff::\n :title: flax.checkpoints, orbax.checkpoint\n :skip_test: flax.checkpoints\n :sync:\n\n ARR_CKPT_DIR = '/tmp/orbax_upgrade/singleton'\n flax.config.update('flax_use_orbax_checkpointing', False)\n\n checkpoints.save_checkpoint(ARR_CKPT_DIR, jnp.arange(10), step=0)\n checkpoints.restore_checkpoint(ARR_CKPT_DIR, target=None)\n\n ---\n\n ARR_CKPT_DIR = '/tmp/orbax_upgrade/singleton'\n\n ckptr = orbax.checkpoint.Checkpointer(orbax.checkpoint.ArrayCheckpointHandler())\n ckptr.save(ARR_CKPT_DIR, jnp.arange(10))\n ckptr.restore(ARR_CKPT_DIR, item=None)\n\n\nFinal words\n***********\n\nThis guide provides an overview of how to migrate from the \"legacy\" Flax checkpointing API to the Orbax API. Orbax provides more functionalities and the Orbax team is actively developing new features. Stay tuned and follow the `official Orbax GitHub repository `__ for more!\n" }, { "path": "docs/guides/converting_and_upgrading/regular_dict_upgrade_guide.rst", "content": "Migrate to regular dicts\n========================\n\nFlax will migrate from returning ``FrozenDicts`` to regular dicts when calling\n:meth:`.init `, :meth:`.init_with_output ` and\n:meth:`.apply ` ``Module`` methods.\n\nThe original issue is outlined `here `__.\n\nThis guide shows some common upgrade patterns.\n\n\nUtility functions\n-----------------\n\n``FrozenDicts`` are immutable dictionaries that implement an additional 4 methods:\n\n* :meth:`copy `\n* :meth:`pop `\n* :meth:`pretty_repr `\n* :meth:`unfreeze `\n\nTo accommodate the regular dict change, replace usage of ``FrozenDict`` methods with their utility function equivalent from ``flax.core.frozen_dict``.\nThese utility functions mimic the behavior of their corresponding ``FrozenDict`` method, and can be called on either ``FrozenDicts`` or regular dicts.\nThe following are the utility functions and example upgrade patterns:\n\n.. testsetup:: default, Only ``FrozenDict``, Both ``FrozenDict`` and regular dict\n\n import flax\n import flax.linen as nn\n import jax\n import jax.numpy as jnp\n\n x = jnp.empty((1,3))\n variables = flax.core.freeze(nn.Dense(5).init(jax.random.key(0), x))\n\n other_variables = jnp.array([1, 1, 1, 1, 1], dtype=jnp.float32)\n\n:meth:`copy `\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n\n.. codediff::\n :title: Only ``FrozenDict``, Both ``FrozenDict`` and regular dict\n :sync:\n\n variables = variables.copy(add_or_replace={'other_variables': other_variables})\n\n ---\n\n variables = flax.core.copy(variables, add_or_replace={'other_variables': other_variables})\n\n:meth:`pop `\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n\n.. codediff::\n :title: Only ``FrozenDict``, Both ``FrozenDict`` and regular dict\n :sync:\n\n state, params = variables.pop('params')\n\n ---\n\n state, params = flax.core.pop(variables, 'params')\n\n:meth:`pretty_repr `\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n\n.. codediff::\n :title: Only ``FrozenDict``, Both ``FrozenDict`` and regular dict\n :sync:\n\n str_repr = variables.pretty_repr()\n\n ---\n\n str_repr = flax.core.pretty_repr(variables)\n\n:meth:`unfreeze `\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n\n.. codediff::\n :title: Only ``FrozenDict``, Both ``FrozenDict`` and regular dict\n :sync:\n\n variables = variables.unfreeze()\n\n ---\n\n variables = flax.core.unfreeze(variables)\n\n\nModifying config values\n-----------------------\n\nA temporary feature flag ``flax_return_frozendict`` is set up to help with the migration.\nTo toggle behavior between returning FrozenDict and regular dict variables at runtime,\nrun ``flax.config.update('flax_return_frozendict', )`` in your code.\n\nFor example:\n\n.. testcode::\n\n x = jnp.empty((1,3))\n\n flax.config.update('flax_return_frozendict', True) # set Flax to return FrozenDicts\n variables = nn.Dense(5).init(jax.random.key(0), x)\n\n assert isinstance(variables, flax.core.FrozenDict)\n\n flax.config.update('flax_return_frozendict', False) # set Flax to return regular dicts\n variables = nn.Dense(5).init(jax.random.key(0), x)\n\n assert isinstance(variables, dict)\n\nAlternatively, the environment variable ``flax_return_frozendict``\n(found `here `__) can be directly modified in the Flax source code.\n\n\nMigration status\n--------------\n\nAs of July 19th, 2023, ``flax_return_frozendict`` is set to ``False`` (see\n`#3193 `__), meaning Flax will default to\nreturning regular dicts from version `0.7.1 `__\nonward. This flag can be flipped to ``True`` temporarily to have Flax return\n``Frozendicts``. However this feature flag will eventually be removed in the future." }, { "path": "docs/guides/converting_and_upgrading/rnncell_upgrade_guide.rst", "content": "RNNCellBase Upgrade Guide\n=========================\n\nThe ``RNNCellBase`` API has undergone some key updates aimed at enhancing usability:\n\n- The ``initialize_carry`` method has transitioned from a class method to an instance method, simplifying its application.\n- All necessary metadata is now stored directly within the cell instance, providing a streamlined method signature.\n\nThis guide will walk you through these changes, demonstrating how to update your existing code to align with these enhancements.\n\nBasic Usage\n-----------\n\n.. testsetup:: New\n\n import flax.linen as nn\n import jax.numpy as jnp\n import jax\n import functools\n\nLet's begin by defining some variables and a sample input that represents\na batch of sequences:\n\n.. testcode:: New\n\n batch_size = 32\n seq_len = 10\n in_features = 64\n out_features = 128\n\n x = jnp.ones((batch_size, seq_len, in_features))\n\nFirst and foremost, it's important to note that all metadata, including the number of features,\ncarry initializer, and so on, is now stored within the cell instance:\n\n.. codediff::\n :title: Legacy, New\n :skip_test: Legacy\n :sync:\n\n cell = nn.LSTMCell()\n\n ---\n\n cell = nn.LSTMCell(features=out_features)\n\nA significant change is that ``initialize_carry`` has been transitioned into an instance method. Given that\nthe cell instance now contains all metadata, the ``initialize_carry`` method's\nsignature only requires a PRNG key and a sample input:\n\n.. codediff::\n :title: Legacy, New\n :skip_test: Legacy\n :sync:\n\n carry = nn.LSTMCell.initialize_carry(jax.random.key(0), (batch_size,), out_features)\n\n ---\n\n carry = cell.initialize_carry(jax.random.key(0), x[:, 0].shape)\n\nHere, ``x[:, 0].shape`` represents the input for the cell (without the time dimension).\nYou can also just create the input shape directly when its more convenient:\n\n.. testcode:: New\n\n carry = cell.initialize_carry(jax.random.key(0), (batch_size, in_features))\n\n\nUpgrade Patterns\n-----------------\n\nThe following sections will demonstrate some useful\npatterns for updating your code to align with the new API.\n\nFirst, we will show how to upgrade a ``Module`` that wraps\na cell, applies the scan logic during ``__call__``, and\nhas a static ``initialize_carry`` method. Here, we will try\nto make the minimal amount of changes to the code to get\nit working, albeit not in the most idiomatic way:\n\n.. codediff::\n :title: Legacy, New\n :skip_test: Legacy\n :sync:\n\n class SimpleLSTM(nn.Module):\n\n @functools.partial(\n nn.transforms.scan,\n variable_broadcast='params',\n in_axes=1, out_axes=1,\n split_rngs={'params': False})\n @nn.compact\n def __call__(self, carry, x):\n\n return nn.OptimizedLSTMCell()(carry, x)\n\n @staticmethod\n def initialize_carry(batch_dims, hidden_size):\n return nn.OptimizedLSTMCell.initialize_carry(\n jax.random.key(0), batch_dims, hidden_size)\n\n ---\n\n class SimpleLSTM(nn.Module):\n\n @functools.partial(\n nn.transforms.scan,\n variable_broadcast='params',\n in_axes=1, out_axes=1,\n split_rngs={'params': False})\n @nn.compact\n def __call__(self, carry, x):\n features = carry[0].shape[-1]\n return nn.OptimizedLSTMCell(features)(carry, x)\n\n @staticmethod\n def initialize_carry(batch_dims, hidden_size):\n return nn.OptimizedLSTMCell(hidden_size, parent=None).initialize_carry(\n jax.random.key(0), (*batch_dims, hidden_size))\n\nNotice how in the new version, we have to extract the number of features from the carry\nduring ``__call__``, and use ``parent=None`` during ``initialize_carry`` to avoid some potential\nside effects.\n\nNext, we will show a more idiomatic way of writing a similar LSTM module. The main change\nhere will be that we will add a ``features`` attribute to the module and use it to initialize\na ``nn.scan``-ed version of the cell in the ``setup`` method:\n\n.. codediff::\n :title: Legacy, New\n :skip_test: Legacy\n :sync:\n\n class SimpleLSTM(nn.Module):\n\n @functools.partial(\n nn.transforms.scan,\n variable_broadcast='params',\n in_axes=1, out_axes=1,\n split_rngs={'params': False})\n @nn.compact\n def __call__(self, carry, x):\n return nn.OptimizedLSTMCell()(carry, x)\n\n @staticmethod\n def initialize_carry(batch_dims, hidden_size):\n return nn.OptimizedLSTMCell.initialize_carry(\n jax.random.key(0), batch_dims, hidden_size)\n\n model = SimpleLSTM()\n carry = SimpleLSTM.initialize_carry((batch_size,), out_features)\n variables = model.init(jax.random.key(0), carry, x)\n\n ---\n\n class SimpleLSTM(nn.Module):\n features: int\n\n def setup(self):\n self.scan_cell = nn.transforms.scan(\n nn.OptimizedLSTMCell,\n variable_broadcast='params',\n in_axes=1, out_axes=1,\n split_rngs={'params': False})(self.features)\n\n\n @nn.compact\n def __call__(self, x):\n carry = self.scan_cell.initialize_carry(jax.random.key(0), x[:, 0].shape)\n return self.scan_cell(carry, x)[1] # only return the output\n\n\n model = SimpleLSTM(features=out_features)\n variables = model.init(jax.random.key(0), x)\n\nBecause the ``carry`` can be easily initialized from the sample input, we can move the\ncall to ``initialize_carry`` into the ``__call__`` method, somewhat simplifying the code.\n\nDevelopment Notes\n-----------------\n\nWhen developing a new cell, consider the following:\n\n* Include necessary metadata as instance attributes.\n* The ``initialize_carry`` now only requires a PRNG key and a sample input.\n* A new ``num_feature_axes`` property is required to specify the number of\n feature dimensions.\n\n.. code-block::\n\n class LSTMCell(nn.RNNCellBase):\n features: int # ← All metadata is now stored within the cell instance\n ... # ↓\n carry_init: Initializer\n\n def initialize_carry(self, rng, input_shape) -> Carry:\n ...\n\n @property\n def num_feature_axes(self):\n return 1\n\n``num_feature_axes`` is a new API feature that allows code handling arbitrary ``RNNCellBase``\ninstances, such as the ``RNN`` Module, to infer the number of batch dimensions and\ndetermine the position of the time axis." }, { "path": "docs/guides/data_preprocessing/full_eval.rst", "content": "Processing the entire Dataset\n=============================\n\nFor efficiency reasons, we form batches that contain multiple examples and\nprocess them in parallel. Especially when evaluating a model, it is important\nthat we process all examples and **avoid losing the remainder** of examples that\ndoes not form a complete batch at the end.\n\n\nThe problem\n-----------\n\nWhen evaluating on a single device, one can either drop the last incomplete\nbatch, or one can form a last batch with a shape different from the preceding\nbatches. Doing the latter has the disadvantage that this will trigger a\n**recompilation** of the ``eval_step()`` because XLA is not shape polymorphic.\n\n.. code-block:: python\n\n collections.Counter(\n tuple(batch['image'].shape)\n for batch in tfds.load('mnist', split='test').batch(per_device_batch_size)\n )\n # output:\n # Counter({(272, 28, 28, 1): 1, (512, 28, 28, 1): 19})\n\nThe problem is accentuated when using multiple devices for data parallelism. If\nthe batch size is not **divisible by the number devices**, then that last step\nmust be executed on a single device (or a subset of devices). Usually one would\ndrop the last batch, but this will lead to incorrect results.\n\n\n.. code-block:: python\n\n sum(\n np.prod(batch['label'].shape)\n for batch in tfds.load('mnist', split='test')\n .batch(per_device_batch_size, drop_remainder=True)\n .batch(jax.local_device_count())\n )\n # output:\n # 9728\n\nUsing multiple hosts further complicates the situation because JAX uses the SPMD\nparadigm and every host must execute the same program. We would usually form\nnon-overlapping splits for different hosts with |tfds.split_for_jax_process()|_,\nbut this can lead to **different numbers for different hosts**, resulting in\ndifferent JAX programs when all examples are to be processed.\n\n.. code-block:: python\n\n process_count = 6\n [\n len(tfds.load(dataset_name, split=tfds.split_for_jax_process(\n 'test', process_index=process_index, process_count=process_count)))\n for process_index in range(process_count)\n ]\n # output:\n # [1667, 1667, 1667, 1667, 1666, 1666]\n\n\n\n.. |tfds.split_for_jax_process()| replace:: ``tfds.split_for_jax_process()``\n.. _tfds.split_for_jax_process(): https://www.tensorflow.org/datasets/api_docs/python/tfds/split_for_jax_process\n\n\nThe solution: padding\n---------------------\n\nEven though it's possible to solve this problem by cleverly adjusting the number\nof batches executed by different devices on different hosts, such a solution\nquickly becomes complicated and makes the main eval loop hard to read with a lot\nof cumbersome logic.\n\nThe more straightforward solution to this problem is to use padding at the end\nof the dataset to make sure that the last batch has the same size as the\npreceding batches.\n\n\nManual implementation\n~~~~~~~~~~~~~~~~~~~~~\n\nThe last batch is manually padded to contain the same number of examples as in\nthe preceding batches. The predictions for the padded examples are discarded\nfrom the computation.\n\n.. code-block:: python\n\n shard = lambda x: einops.rearrange(\n x, '(d b) ... -> d b ...', d=jax.local_device_count())\n unshard = lambda x: einops.rearrange(x, 'd b ... -> (d b) ...')\n\n correct = total = 0\n for batch in ds.as_numpy_iterator():\n images = batch['image']\n n = len(images)\n padding = np.zeros([per_host_batch_size - n, *images.shape[1:]], images.dtype)\n padded_images = np.concatenate([images, padding])\n preds = unshard(get_preds(variables, shard(padded_images)))[:n]\n total += n\n correct += (batch['label'] == preds.argmax(axis=-1)).sum()\n\n\nUsing ``pad_shard_unpad()``\n~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\nThe above pattern, namely the pad→shard→predict→unshard→unpad sequence, can be\nextracted into a utility wrapper ``pad_shard_unpad()``, which greatly simplifies\nabove evaluation loop.\n\n.. code-block:: python\n\n correct = total = 0\n for batch in ds.as_numpy_iterator():\n preds = flax.jax_utils.pad_shard_unpad(get_preds)(\n vs, batch['image'], min_device_batch=per_device_batch_size)\n total += len(batch['image'])\n correct += (batch['label'] == preds.argmax(axis=-1)).sum()\n\n\nComputing metrics in ``eval_step()``\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\nInstead of returning the predictions and computing the metrics in the main\nevaluation loop, we would often want to make the metric computation part of the\nevaluation step, especially when using libraries like |jax_metrics|_, or\n|clu.metrics|_.\n\nIn that case we would want to pass the metrics as a ``static_argnums`` (i.e. do\nnot shard/pad it), and treat the return value as ``static_return`` too (i.e. no\nun-sharding or un-padding):\n\n.. code-block:: python\n\n def eval_step(metrics, variables, batch):\n print('retrigger compilation', {k: v.shape for k, v in batch.items()})\n preds = model.apply(variables, batch['image'])\n correct = (batch['mask'] & (batch['label'] == preds.argmax(axis=-1))).sum()\n total = batch['mask'].sum()\n return dict(\n correct=metrics['correct'] + jax.lax.psum(correct, axis_name='batch'),\n total=metrics['total'] + jax.lax.psum(total, axis_name='batch'),\n )\n\n eval_step = jax.pmap(eval_step, axis_name='batch')\n eval_step = flax.jax_utils.pad_shard_unpad(\n eval_step, static_argnums=(0, 1), static_return=True)\n\n.. |jax_metrics| replace:: ``clu.metrics``\n.. _jax_metrics: https://github.com/cgarciae/jax_metrics\n\n\n.. |clu.metrics| replace:: ``clu.metrics``\n.. _clu.metrics: https://github.com/google/CommonLoopUtils/blob/main/clu/metrics.py\n\n\nAdding \"infinite padding\"\n~~~~~~~~~~~~~~~~~~~~~~~~~\n\nThe above solution works in most cases, but it has some limitations:\n\n1. In the rare case where even splitting of the dataset on multiple hosts leads\n to a different number of batches. Imagine having a dataset of ``n=4097``\n examples, and evaluating this on ``h=8``, each having ``d=8`` local devices,\n and forming on-device batch sizes of ``b=128``. With even dataset splitting,\n the first host would get ``4096/8+1==513`` examples, and all other hosts\n would get ``4096/8==512`` examples. Forming per-host batches of ``d*b==512``\n this would lead to two batches on the first host, and a single batch on all\n other hosts, violating SPMD principles and hanging the multi-host setup in\n the last ``psum()`` directive (which would only be executed by the first\n host, but not the others).\n\n2. When dropping examples dynamically by using ``ds.filter()``.\n\nIn these more complicated cases we could add \"infinite padding\" to the dataset,\non each of the hosts independently, and continuing processing examples until\n*all* hosts run out of unpadded examples.\n\n.. code-block:: python\n\n correct = total = 0\n for batch in ds.as_numpy_iterator():\n n = count_p(batch['mask'])[0].item() # adds sync barrier\n if not n: break\n\n preds = get_preds(vs, batch['image']).argmax(axis=-1)\n total += n\n correct += count_correct_p(batch['label'], preds, batch['mask'])[0]\n" }, { "path": "docs/guides/data_preprocessing/index.rst", "content": "Data preprocessing\n=================\n\n.. toctree::\n :maxdepth: 1\n\n full_eval\n loading_datasets\n" }, { "path": "docs/guides/data_preprocessing/loading_datasets.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"# Loading datasets\\n\",\n \"\\n\",\n \"[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google/flax/blob/main/docs/guides/data_preprocessing/loading_datasets.ipynb)\\n\",\n \"\\n\",\n \"A neural net written in Jax+Flax expects its input data as `jax.numpy` array instances. Therefore, loading a dataset from any source is as simple as converting it to `jax.numpy` types and reshaping it to the appropriate dimensions for your network.\\n\",\n \"\\n\",\n \"As an example, this guide demonstrates how to import [MNIST](http://yann.lecun.com/exdb/mnist/) using the APIs from Torchvision, Tensorflow, and Hugging Face. We'll load the whole dataset into memory. For datasets that don't fit into memory the process is analogous but should be done in a batchwise fashion.\\n\",\n \"\\n\",\n \"The MNIST dataset consists of greyscale images of 28x28 pixels of handwritten digits, and has a designated 60k/10k train/test split. The task is to predict the correct class (digit 0, ..., 9) of each image.\\n\",\n \"\\n\",\n \"Assuming a CNN-based classifier, the input data should have shape `(B, 28, 28, 1)`, where the trailing singleton dimension denotes the greyscale image channel.\\n\",\n \"\\n\",\n \"The labels are simply the integer denoting the digit corresponding to the image. Labels should therefore have shape `(B,)`, to enable loss computation with [`optax.softmax_cross_entropy_with_integer_labels`](https://optax.readthedocs.io/en/latest/api.html#optax.softmax_cross_entropy_with_integer_labels).\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 2,\n \"metadata\": {\n \"tags\": [\n \"skip-execution\"\n ]\n },\n \"outputs\": [],\n \"source\": [\n \"import numpy as np\\n\",\n \"import jax.numpy as jnp\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Loading from `torchvision.datasets`\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 3,\n \"metadata\": {\n \"tags\": [\n \"skip-execution\"\n ]\n },\n \"outputs\": [],\n \"source\": [\n \"import torchvision\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 4,\n \"metadata\": {\n \"tags\": [\n \"skip-execution\"\n ]\n },\n \"outputs\": [],\n \"source\": [\n \"def get_dataset_torch():\\n\",\n \" mnist = {\\n\",\n \" 'train': torchvision.datasets.MNIST('./data', train=True, download=True),\\n\",\n \" 'test': torchvision.datasets.MNIST('./data', train=False, download=True)\\n\",\n \" }\\n\",\n \"\\n\",\n \" ds = {}\\n\",\n \"\\n\",\n \" for split in ['train', 'test']:\\n\",\n \" ds[split] = {\\n\",\n \" 'image': mnist[split].data.numpy(),\\n\",\n \" 'label': mnist[split].targets.numpy()\\n\",\n \" }\\n\",\n \"\\n\",\n \" # cast from np to jnp and rescale the pixel values from [0,255] to [0,1]\\n\",\n \" ds[split]['image'] = jnp.float32(ds[split]['image']) / 255\\n\",\n \" ds[split]['label'] = jnp.int16(ds[split]['label'])\\n\",\n \"\\n\",\n \" # torchvision returns shape (B, 28, 28).\\n\",\n \" # hence, append the trailing channel dimension.\\n\",\n \" ds[split]['image'] = jnp.expand_dims(ds[split]['image'], 3)\\n\",\n \"\\n\",\n \" return ds['train'], ds['test']\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 5,\n \"metadata\": {\n \"outputId\": \"be39b756-d13e-4380-b99e-a5cbf61458cc\",\n \"tags\": [\n \"skip-execution\"\n ]\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"(60000, 28, 28, 1) float32\\n\",\n \"(60000,) int16\\n\",\n \"(10000, 28, 28, 1) float32\\n\",\n \"(10000,) int16\\n\"\n ]\n }\n ],\n \"source\": [\n \"train, test = get_dataset_torch()\\n\",\n \"print(train['image'].shape, train['image'].dtype)\\n\",\n \"print(train['label'].shape, train['label'].dtype)\\n\",\n \"print(test['image'].shape, test['image'].dtype)\\n\",\n \"print(test['label'].shape, test['label'].dtype)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Loading from `tensorflow_datasets`\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 5,\n \"metadata\": {\n \"tags\": [\n \"skip-execution\"\n ]\n },\n \"outputs\": [],\n \"source\": [\n \"import tensorflow_datasets as tfds\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 6,\n \"metadata\": {\n \"tags\": [\n \"skip-execution\"\n ]\n },\n \"outputs\": [],\n \"source\": [\n \"def get_dataset_tf():\\n\",\n \" mnist = tfds.builder('mnist')\\n\",\n \" mnist.download_and_prepare()\\n\",\n \"\\n\",\n \" ds = {}\\n\",\n \"\\n\",\n \" for split in ['train', 'test']:\\n\",\n \" ds[split] = tfds.as_numpy(mnist.as_dataset(split=split, batch_size=-1))\\n\",\n \"\\n\",\n \" # cast to jnp and rescale pixel values\\n\",\n \" ds[split]['image'] = jnp.float32(ds[split]['image']) / 255\\n\",\n \" ds[split]['label'] = jnp.int16(ds[split]['label'])\\n\",\n \"\\n\",\n \" return ds['train'], ds['test']\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 7,\n \"metadata\": {\n \"outputId\": \"25d2c468-cbc8-4971-a738-1295ce8c6f16\",\n \"tags\": [\n \"skip-execution\"\n ]\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"(60000, 28, 28, 1) float32\\n\",\n \"(60000,) int16\\n\",\n \"(10000, 28, 28, 1) float32\\n\",\n \"(10000,) int16\\n\"\n ]\n }\n ],\n \"source\": [\n \"train, test = get_dataset_tf()\\n\",\n \"print(train['image'].shape, train['image'].dtype)\\n\",\n \"print(train['label'].shape, train['label'].dtype)\\n\",\n \"print(test['image'].shape, test['image'].dtype)\\n\",\n \"print(test['label'].shape, test['label'].dtype)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Loading from 🤗 Hugging Face `datasets`\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {\n \"tags\": [\n \"skip-execution\"\n ]\n },\n \"outputs\": [],\n \"source\": [\n \"#!pip install datasets # datasets isn't preinstalled on Colab; uncomment to install\\n\",\n \"from datasets import load_dataset\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 9,\n \"metadata\": {\n \"tags\": [\n \"skip-execution\"\n ]\n },\n \"outputs\": [],\n \"source\": [\n \"def get_dataset_hf():\\n\",\n \" mnist = load_dataset(\\\"mnist\\\")\\n\",\n \"\\n\",\n \" ds = {}\\n\",\n \"\\n\",\n \" for split in ['train', 'test']:\\n\",\n \" ds[split] = {\\n\",\n \" 'image': np.array([np.array(im) for im in mnist[split]['image']]),\\n\",\n \" 'label': np.array(mnist[split]['label'])\\n\",\n \" }\\n\",\n \"\\n\",\n \" # cast to jnp and rescale pixel values\\n\",\n \" ds[split]['image'] = jnp.float32(ds[split]['image']) / 255\\n\",\n \" ds[split]['label'] = jnp.int16(ds[split]['label'])\\n\",\n \"\\n\",\n \" # append trailing channel dimension\\n\",\n \" ds[split]['image'] = jnp.expand_dims(ds[split]['image'], 3)\\n\",\n \"\\n\",\n \" return ds['train'], ds['test']\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 10,\n \"metadata\": {\n \"outputId\": \"b026b33f-3bdd-4d26-867c-49400fff1c96\",\n \"tags\": [\n \"skip-execution\"\n ]\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"(60000, 28, 28, 1) float32\\n\",\n \"(60000,) int16\\n\",\n \"(10000, 28, 28, 1) float32\\n\",\n \"(10000,) int16\\n\"\n ]\n }\n ],\n \"source\": [\n \"train, test = get_dataset_hf()\\n\",\n \"print(train['image'].shape, train['image'].dtype)\\n\",\n \"print(train['label'].shape, train['label'].dtype)\\n\",\n \"print(test['image'].shape, test['image'].dtype)\\n\",\n \"print(test['label'].shape, test['label'].dtype)\"\n ]\n }\n ],\n \"metadata\": {\n \"jupytext\": {\n \"formats\": \"ipynb,md:myst\"\n },\n \"language_info\": {\n \"codemirror_mode\": {\n \"name\": \"ipython\",\n \"version\": 3\n },\n \"file_extension\": \".py\",\n \"mimetype\": \"text/x-python\",\n \"name\": \"python\",\n \"nbconvert_exporter\": \"python\",\n \"pygments_lexer\": \"ipython3\",\n \"version\": \"3.11.5\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 0\n}\n" }, { "path": "docs/guides/data_preprocessing/loading_datasets.md", "content": "---\njupytext:\n formats: ipynb,md:myst\n text_representation:\n extension: .md\n format_name: myst\n format_version: 0.13\n jupytext_version: 1.13.8\n---\n\n# Loading datasets\n\n[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google/flax/blob/main/docs/guides/data_preprocessing/loading_datasets.ipynb)\n\nA neural net written in Jax+Flax expects its input data as `jax.numpy` array instances. Therefore, loading a dataset from any source is as simple as converting it to `jax.numpy` types and reshaping it to the appropriate dimensions for your network.\n\nAs an example, this guide demonstrates how to import [MNIST](http://yann.lecun.com/exdb/mnist/) using the APIs from Torchvision, Tensorflow, and Hugging Face. We'll load the whole dataset into memory. For datasets that don't fit into memory the process is analogous but should be done in a batchwise fashion.\n\nThe MNIST dataset consists of greyscale images of 28x28 pixels of handwritten digits, and has a designated 60k/10k train/test split. The task is to predict the correct class (digit 0, ..., 9) of each image.\n\nAssuming a CNN-based classifier, the input data should have shape `(B, 28, 28, 1)`, where the trailing singleton dimension denotes the greyscale image channel.\n\nThe labels are simply the integer denoting the digit corresponding to the image. Labels should therefore have shape `(B,)`, to enable loss computation with [`optax.softmax_cross_entropy_with_integer_labels`](https://optax.readthedocs.io/en/latest/api.html#optax.softmax_cross_entropy_with_integer_labels).\n\n```{code-cell} ipython3\n:tags: [skip-execution]\n\nimport numpy as np\nimport jax.numpy as jnp\n```\n\n## Loading from `torchvision.datasets`\n\n```{code-cell} ipython3\n:tags: [skip-execution]\n\nimport torchvision\n```\n\n```{code-cell} ipython3\n:tags: [skip-execution]\n\ndef get_dataset_torch():\n mnist = {\n 'train': torchvision.datasets.MNIST('./data', train=True, download=True),\n 'test': torchvision.datasets.MNIST('./data', train=False, download=True)\n }\n\n ds = {}\n\n for split in ['train', 'test']:\n ds[split] = {\n 'image': mnist[split].data.numpy(),\n 'label': mnist[split].targets.numpy()\n }\n\n # cast from np to jnp and rescale the pixel values from [0,255] to [0,1]\n ds[split]['image'] = jnp.float32(ds[split]['image']) / 255\n ds[split]['label'] = jnp.int16(ds[split]['label'])\n\n # torchvision returns shape (B, 28, 28).\n # hence, append the trailing channel dimension.\n ds[split]['image'] = jnp.expand_dims(ds[split]['image'], 3)\n\n return ds['train'], ds['test']\n```\n\n```{code-cell} ipython3\n:outputId: be39b756-d13e-4380-b99e-a5cbf61458cc\n:tags: [skip-execution]\n\ntrain, test = get_dataset_torch()\nprint(train['image'].shape, train['image'].dtype)\nprint(train['label'].shape, train['label'].dtype)\nprint(test['image'].shape, test['image'].dtype)\nprint(test['label'].shape, test['label'].dtype)\n```\n\n## Loading from `tensorflow_datasets`\n\n```{code-cell} ipython3\n:tags: [skip-execution]\n\nimport tensorflow_datasets as tfds\n```\n\n```{code-cell} ipython3\n:tags: [skip-execution]\n\ndef get_dataset_tf():\n mnist = tfds.builder('mnist')\n mnist.download_and_prepare()\n\n ds = {}\n\n for split in ['train', 'test']:\n ds[split] = tfds.as_numpy(mnist.as_dataset(split=split, batch_size=-1))\n\n # cast to jnp and rescale pixel values\n ds[split]['image'] = jnp.float32(ds[split]['image']) / 255\n ds[split]['label'] = jnp.int16(ds[split]['label'])\n\n return ds['train'], ds['test']\n```\n\n```{code-cell} ipython3\n:outputId: 25d2c468-cbc8-4971-a738-1295ce8c6f16\n:tags: [skip-execution]\n\ntrain, test = get_dataset_tf()\nprint(train['image'].shape, train['image'].dtype)\nprint(train['label'].shape, train['label'].dtype)\nprint(test['image'].shape, test['image'].dtype)\nprint(test['label'].shape, test['label'].dtype)\n```\n\n## Loading from 🤗 Hugging Face `datasets`\n\n```{code-cell} ipython3\n:tags: [skip-execution]\n\n#!pip install datasets # datasets isn't preinstalled on Colab; uncomment to install\nfrom datasets import load_dataset\n```\n\n```{code-cell} ipython3\n:tags: [skip-execution]\n\ndef get_dataset_hf():\n mnist = load_dataset(\"mnist\")\n\n ds = {}\n\n for split in ['train', 'test']:\n ds[split] = {\n 'image': np.array([np.array(im) for im in mnist[split]['image']]),\n 'label': np.array(mnist[split]['label'])\n }\n\n # cast to jnp and rescale pixel values\n ds[split]['image'] = jnp.float32(ds[split]['image']) / 255\n ds[split]['label'] = jnp.int16(ds[split]['label'])\n\n # append trailing channel dimension\n ds[split]['image'] = jnp.expand_dims(ds[split]['image'], 3)\n\n return ds['train'], ds['test']\n```\n\n```{code-cell} ipython3\n:outputId: b026b33f-3bdd-4d26-867c-49400fff1c96\n:tags: [skip-execution]\n\ntrain, test = get_dataset_hf()\nprint(train['image'].shape, train['image'].dtype)\nprint(train['label'].shape, train['label'].dtype)\nprint(test['image'].shape, test['image'].dtype)\nprint(test['label'].shape, test['label'].dtype)\n```\n" }, { "path": "docs/guides/flax_fundamentals/arguments.md", "content": "# Dealing with Flax Module arguments\n\n## Introduction\n\nIn Flax Linen we can define `Module` arguments either as dataclass attributes or as arguments to methods (usually `__call__`).\nTypically the distinction is clear:\n* Completely fixed properties, such as the choice of kernel initializer or number of output features, are hyperparameters and should be defined as dataclass attributes. Typically two Module instances with different hyperparameters cannot share in a meaningful way.\n* Dynamic properties, such as input data and top-level \"mode switches\" like `train=True/False`, should be passed as arguments to `__call__` or another method.\n\nSome cases are however less clear cut. Take for example the `Dropout` module.\nWe have a number of clear hyperparameters:\n\n1. The dropout rate\n2. The axes for which a dropout mask is generated\n\nAnd some clear call time arguments:\n\n1. The input that should be masked using dropout\n2. The (optional) rng used to sample the random mask\n\nThere is however one property that is ambiguous -- the `deterministic` property in a Dropout module.\n\nIf `deterministic` is `True` no dropout mask is sampled. This is typically used during model evaluation.\nHowever, if we pass `eval=True` or `train=False` to a top-level Module. The `deterministic` argument needs\nto be applied everywhere and the boolean argument needs to be passed down to all the layers that might use `Dropout`.\nIf instead `deterministic` is a dataclass attribute, we might do the following:\n\n```python\nfrom functools import partial\nfrom flax import linen as nn\n\nclass ResidualModel(nn.Module):\n drop_rate: float\n\n @nn.compact\n def __call__(self, x, *, train):\n dropout = partial(nn.Dropout, rate=self.drop_rate, deterministic=not train)\n for i in range(10):\n x += ResidualBlock(dropout=dropout, ...)(x)\n```\n\nIt makes sense to pass `determinstic` to the constructor here because this way we can pass the dropout template to the sub-modules.\nNow the sub-module no longer needs to take care of train vs eval mode and can simply use the `dropout` argument.\nNote that because the dropout layer can only be constructed in the sub-module we can only partially apply `deterministic` to the constructor but not to `__call__`.\n\nHowever, if `deterministic` is a dataclass attribute we run into trouble when using the setup pattern. We would **want** to write our module code like this:\n\n```python\nclass SomeModule(nn.Module):\n drop_rate: float\n\n def setup(self):\n self.dropout = nn.Dropout(rate=self.drop_rate)\n\n @nn.compact\n def __call__(self, x, *, train):\n # ...\n x = self.dropout(x, deterministic=not train)\n # ...\n```\n\nBut, as defined above, `deterministic` would be an attribute, so this doesn't work.\nHere it makes sense to pass `deterministic` during `__call__` because it depends on the `train` argument.\n\n## Solution\n\nWe can support both use cases described before by allowing certain properties to be passed\nas dataclass attributes or as method argument (but not both!).\nThis can be implemented as follows:\n```python\nclass MyDropout(nn.Module):\n drop_rate: float\n deterministic: Optional[bool] = None\n\n @nn.compact\n def __call__(self, x, deterministic=None):\n deterministic = nn.merge_param('deterministic', self.deterministic, deterministic)\n # ...\n```\n\nIn this example `nn.merge_param` will ensure that either `self.deterministic` or `deterministic` is set but not both.\nAn error is raised if both values are `None` or both values are not `None`.\nThis avoids confusing behavior where 2 different parts of the code set the same parameter and one is overruled by the other.\nIt also avoids a default value which would probably cause either the train step or eval step of a training procedure to be broken by default.\n\n\n\n## Functional Core\n\nFunctional core defines functions rather than classes.\nTherefore, there is no clear distinction between hyperparameters and call-time arguments.\nThe only way to pre-determine the hyperparameters is by using `partial`.\nOn the upside, there are no ambiguous cases where method arguments could also be attributes.\n" }, { "path": "docs/guides/flax_fundamentals/flax_basics.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google/flax/blob/main/docs/guides/flax_fundamentals/flax_basics.ipynb)\\n\",\n \"[![Open On GitHub](https://img.shields.io/badge/Open-on%20GitHub-blue?logo=GitHub)](https://github.com/google/flax/blob/main/docs/guides/flax_fundamentals/flax_basics.ipynb)\\n\",\n \"\\n\",\n \"# Flax Basics\\n\",\n \"\\n\",\n \"This notebook will walk you through the following workflow:\\n\",\n \"\\n\",\n \"* Instantiating a model from Flax built-in layers or third-party models.\\n\",\n \"* Initializing parameters of the model and manually written training.\\n\",\n \"* Using optimizers provided by Flax to ease training.\\n\",\n \"* Serialization of parameters and other objects.\\n\",\n \"* Creating your own models and managing state.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Setting up our environment\\n\",\n \"\\n\",\n \"Here we provide the code needed to set up the environment for our notebook.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 1,\n \"metadata\": {\n \"outputId\": \"e30aa464-fa52-4f35-df96-716c68a4b3ee\",\n \"tags\": [\n \"skip-execution\"\n ]\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"\\u001b[33mWARNING: Running pip as root will break packages and permissions. You should install packages reliably by using venv: https://pip.pypa.io/warnings/venv\\u001b[0m\\n\",\n \"\\u001b[33mWARNING: Running pip as root will break packages and permissions. You should install packages reliably by using venv: https://pip.pypa.io/warnings/venv\\u001b[0m\\n\"\n ]\n }\n ],\n \"source\": [\n \"# Install the latest JAXlib version.\\n\",\n \"!pip install --upgrade -q pip jax jaxlib\\n\",\n \"# Install Flax at head:\\n\",\n \"!pip install --upgrade -q git+https://github.com/google/flax.git\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 2,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"import jax\\n\",\n \"from typing import Any, Callable, Sequence\\n\",\n \"from jax import random, numpy as jnp\\n\",\n \"import flax\\n\",\n \"from flax import linen as nn\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Linear regression with Flax\\n\",\n \"\\n\",\n \"In the previous *JAX for the impatient* notebook, we finished up with a linear regression example. As we know, linear regression can also be written as a single dense neural network layer, which we will show in the following so that we can compare how it's done.\\n\",\n \"\\n\",\n \"A dense layer is a layer that has a kernel parameter $W\\\\in\\\\mathcal{M}_{m,n}(\\\\mathbb{R})$ where $m$ is the number of features as an output of the model, and $n$ the dimensionality of the input, and a bias parameter $b\\\\in\\\\mathbb{R}^m$. The dense layers returns $Wx+b$ from an input $x\\\\in\\\\mathbb{R}^n$.\\n\",\n \"\\n\",\n \"This dense layer is already provided by Flax in the `flax.linen` module (here imported as `nn`).\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 3,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# We create one dense layer instance (taking 'features' parameter as input)\\n\",\n \"model = nn.Dense(features=5)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Layers (and models in general, we'll use that word from now on) are subclasses of the `linen.Module` class.\\n\",\n \"\\n\",\n \"### Model parameters & initialization\\n\",\n \"\\n\",\n \"Parameters are not stored with the models themselves. You need to initialize parameters by calling the `init` function, using a PRNGKey and dummy input data.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 4,\n \"metadata\": {\n \"outputId\": \"06feb9d2-db50-4f41-c169-6df4336f43a5\"\n },\n \"outputs\": [\n {\n \"name\": \"stderr\",\n \"output_type\": \"stream\",\n \"text\": [\n \"WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\\n\"\n ]\n },\n {\n \"data\": {\n \"text/plain\": [\n \"FrozenDict({\\n\",\n \" params: {\\n\",\n \" bias: (5,),\\n\",\n \" kernel: (10, 5),\\n\",\n \" },\\n\",\n \"})\"\n ]\n },\n \"execution_count\": 4,\n \"metadata\": {\n \"tags\": []\n },\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"key1, key2 = random.split(random.key(0))\\n\",\n \"x = random.normal(key1, (10,)) # Dummy input data\\n\",\n \"params = model.init(key2, x) # Initialization call\\n\",\n \"jax.tree_util.tree_map(lambda x: x.shape, params) # Checking output shapes\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"*Note: JAX and Flax, like NumPy, are row-based systems, meaning that vectors are represented as row vectors and not column vectors. This can be seen in the shape of the kernel here.*\\n\",\n \"\\n\",\n \"The result is what we expect: bias and kernel parameters of the correct size. Under the hood:\\n\",\n \"\\n\",\n \"* The dummy input data `x` is used to trigger shape inference: we only declared the number of features we wanted in the output of the model, not the size of the input. Flax finds out by itself the correct size of the kernel.\\n\",\n \"* The random PRNG key is used to trigger the initialization functions (those have default values provided by the module here).\\n\",\n \"* Initialization functions are called to generate the initial set of parameters that the model will use. Those are functions that take as arguments `(PRNG Key, shape, dtype)` and return an Array of shape `shape`.\\n\",\n \"* The init function returns the initialized set of parameters (you can also get the output of the forward pass on the dummy input with the same syntax by using the `init_with_output` method instead of `init`.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"To conduct a forward pass with the model with a given set of parameters (which are never stored with the model), we just use the `apply` method by providing it the parameters to use as well as the input:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 6,\n \"metadata\": {\n \"outputId\": \"7bbe6bb4-94d5-4574-fbb5-aa0fcd1c84ae\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"DeviceArray([-0.7358944, 1.3583755, -0.7976872, 0.8168598, 0.6297793], dtype=float32)\"\n ]\n },\n \"execution_count\": 6,\n \"metadata\": {\n \"tags\": []\n },\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"model.apply(params, x)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Gradient descent\\n\",\n \"\\n\",\n \"If you jumped here directly without going through the JAX part, here is the linear regression formulation we're going to use: from a set of data points $\\\\{(x_i,y_i), i\\\\in \\\\{1,\\\\ldots, k\\\\}, x_i\\\\in\\\\mathbb{R}^n,y_i\\\\in\\\\mathbb{R}^m\\\\}$, we try to find a set of parameters $W\\\\in \\\\mathcal{M}_{m,n}(\\\\mathbb{R}), b\\\\in\\\\mathbb{R}^m$ such that the function $f_{W,b}(x)=Wx+b$ minimizes the mean squared error:\\n\",\n \"\\n\",\n \"$$\\\\mathcal{L}(W,b)\\\\rightarrow\\\\frac{1}{k}\\\\sum_{i=1}^{k} \\\\frac{1}{2}\\\\|y_i-f_{W,b}(x_i)\\\\|^2_2$$\\n\",\n \"\\n\",\n \"Here, we see that the tuple $(W,b)$ matches the parameters of the Dense layer. We'll perform gradient descent using those. Let's first generate the fake data we'll use. The data is exactly the same as in the JAX part's linear regression pytree example.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 7,\n \"metadata\": {\n \"outputId\": \"6eae59dc-0632-4f53-eac8-c22a7c646a52\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"x shape: (20, 10) ; y shape: (20, 5)\\n\"\n ]\n }\n ],\n \"source\": [\n \"# Set problem dimensions.\\n\",\n \"n_samples = 20\\n\",\n \"x_dim = 10\\n\",\n \"y_dim = 5\\n\",\n \"\\n\",\n \"# Generate random ground truth W and b.\\n\",\n \"key = random.key(0)\\n\",\n \"k1, k2 = random.split(key)\\n\",\n \"W = random.normal(k1, (x_dim, y_dim))\\n\",\n \"b = random.normal(k2, (y_dim,))\\n\",\n \"# Store the parameters in a FrozenDict pytree.\\n\",\n \"true_params = flax.core.freeze({'params': {'bias': b, 'kernel': W}})\\n\",\n \"\\n\",\n \"# Generate samples with additional noise.\\n\",\n \"key_sample, key_noise = random.split(k1)\\n\",\n \"x_samples = random.normal(key_sample, (n_samples, x_dim))\\n\",\n \"y_samples = jnp.dot(x_samples, W) + b + 0.1 * random.normal(key_noise,(n_samples, y_dim))\\n\",\n \"print('x shape:', x_samples.shape, '; y shape:', y_samples.shape)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"We copy the same training loop that we used in the JAX pytree linear regression example with `jax.value_and_grad()`, but here we can use `model.apply()` instead of having to define our own feed-forward function (`predict_pytree()` in the [JAX example](https://flax.readthedocs.io/en/latest/guides/jax_for_the_impatient.html#linear-regression-with-pytrees)).\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 8,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# Same as JAX version but using model.apply().\\n\",\n \"@jax.jit\\n\",\n \"def mse(params, x_batched, y_batched):\\n\",\n \" # Define the squared loss for a single pair (x,y)\\n\",\n \" def squared_error(x, y):\\n\",\n \" pred = model.apply(params, x)\\n\",\n \" return jnp.inner(y-pred, y-pred) / 2.0\\n\",\n \" # Vectorize the previous to compute the average of the loss on all samples.\\n\",\n \" return jnp.mean(jax.vmap(squared_error)(x_batched,y_batched), axis=0)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"And finally perform the gradient descent.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 9,\n \"metadata\": {\n \"outputId\": \"50d975b3-4706-4d8a-c4b8-2629ab8e3ac4\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"Loss for \\\"true\\\" W,b: 0.023639778\\n\",\n \"Loss step 0: 38.094772\\n\",\n \"Loss step 10: 0.44692168\\n\",\n \"Loss step 20: 0.10053458\\n\",\n \"Loss step 30: 0.035822745\\n\",\n \"Loss step 40: 0.018846875\\n\",\n \"Loss step 50: 0.013864839\\n\",\n \"Loss step 60: 0.012312559\\n\",\n \"Loss step 70: 0.011812928\\n\",\n \"Loss step 80: 0.011649306\\n\",\n \"Loss step 90: 0.011595251\\n\",\n \"Loss step 100: 0.0115773035\\n\"\n ]\n }\n ],\n \"source\": [\n \"learning_rate = 0.3 # Gradient step size.\\n\",\n \"print('Loss for \\\"true\\\" W,b: ', mse(true_params, x_samples, y_samples))\\n\",\n \"loss_grad_fn = jax.value_and_grad(mse)\\n\",\n \"\\n\",\n \"@jax.jit\\n\",\n \"def update_params(params, learning_rate, grads):\\n\",\n \" params = jax.tree_util.tree_map(\\n\",\n \" lambda p, g: p - learning_rate * g, params, grads)\\n\",\n \" return params\\n\",\n \"\\n\",\n \"for i in range(101):\\n\",\n \" # Perform one gradient update.\\n\",\n \" loss_val, grads = loss_grad_fn(params, x_samples, y_samples)\\n\",\n \" params = update_params(params, learning_rate, grads)\\n\",\n \" if i % 10 == 0:\\n\",\n \" print(f'Loss step {i}: ', loss_val)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Optimizing with Optax\\n\",\n \"\\n\",\n \"Flax used to use its own `flax.optim` package for optimization, but with\\n\",\n \"[FLIP #1009](https://github.com/google/flax/blob/main/docs/flip/1009-optimizer-api.md)\\n\",\n \"this was deprecated in favor of\\n\",\n \"[Optax](https://github.com/deepmind/optax).\\n\",\n \"\\n\",\n \"Basic usage of Optax is straightforward:\\n\",\n \"\\n\",\n \"1. Choose an optimization method (e.g. `optax.adam`).\\n\",\n \"2. Create optimizer state from parameters (for the Adam optimizer, this state will contain the [momentum values](https://optax.readthedocs.io/en/latest/api.html#optax.adam)).\\n\",\n \"3. Compute the gradients of your loss with `jax.value_and_grad()`.\\n\",\n \"4. At every iteration, call the Optax `update` function to update the internal\\n\",\n \" optimizer state and create an update to the parameters. Then add the update\\n\",\n \" to the parameters with Optax's `apply_updates` method.\\n\",\n \"\\n\",\n \"Note that Optax can do a lot more: it's designed for composing simple gradient\\n\",\n \"transformations into more complex transformations that allows to implement a\\n\",\n \"wide range of optimizers. There is also support for changing optimizer\\n\",\n \"hyperparameters over time (\\\"schedules\\\"), applying different updates to different\\n\",\n \"parts of the parameter tree (\\\"masking\\\") and much more. For details please refer\\n\",\n \"to the\\n\",\n \"[official documentation](https://optax.readthedocs.io/en/latest/).\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 10,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"import optax\\n\",\n \"tx = optax.adam(learning_rate=learning_rate)\\n\",\n \"opt_state = tx.init(params)\\n\",\n \"loss_grad_fn = jax.value_and_grad(mse)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 11,\n \"metadata\": {\n \"outputId\": \"eec0c096-1d9e-4b3c-f8e5-942ee63828ec\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"Loss step 0: 0.011576377\\n\",\n \"Loss step 10: 0.0115710115\\n\",\n \"Loss step 20: 0.011569244\\n\",\n \"Loss step 30: 0.011568661\\n\",\n \"Loss step 40: 0.011568454\\n\",\n \"Loss step 50: 0.011568379\\n\",\n \"Loss step 60: 0.011568358\\n\",\n \"Loss step 70: 0.01156836\\n\",\n \"Loss step 80: 0.01156835\\n\",\n \"Loss step 90: 0.011568353\\n\",\n \"Loss step 100: 0.011568348\\n\"\n ]\n }\n ],\n \"source\": [\n \"for i in range(101):\\n\",\n \" loss_val, grads = loss_grad_fn(params, x_samples, y_samples)\\n\",\n \" updates, opt_state = tx.update(grads, opt_state)\\n\",\n \" params = optax.apply_updates(params, updates)\\n\",\n \" if i % 10 == 0:\\n\",\n \" print('Loss step {}: '.format(i), loss_val)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Serializing the result\\n\",\n \"\\n\",\n \"Now that we're happy with the result of our training, we might want to save the model parameters to load them back later. Flax provides a serialization package to enable you to do that.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 13,\n \"metadata\": {\n \"outputId\": \"b97e7d83-3e40-4a80-b1fe-1f6ceff30a0c\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"Dict output\\n\",\n \"{'params': {'bias': DeviceArray([-1.4540135, -2.0262308, 2.0806582, 1.2201802, -0.9964547], dtype=float32), 'kernel': DeviceArray([[ 1.0106664 , 0.19014716, 0.04533899, -0.92722285,\\n\",\n \" 0.34720102],\\n\",\n \" [ 1.7320251 , 0.9901233 , 1.1662225 , 1.1027892 ,\\n\",\n \" -0.10574618],\\n\",\n \" [-1.2009128 , 0.28837162, 1.4176372 , 0.12073109,\\n\",\n \" -1.3132601 ],\\n\",\n \" [-1.1944956 , -0.18993308, 0.03379077, 1.3165942 ,\\n\",\n \" 0.07996067],\\n\",\n \" [ 0.14103189, 1.3737966 , -1.3162128 , 0.53401774,\\n\",\n \" -2.239638 ],\\n\",\n \" [ 0.5643044 , 0.813604 , 0.31888172, 0.5359193 ,\\n\",\n \" 0.90352124],\\n\",\n \" [-0.37948322, 1.7408353 , 1.0788013 , -0.5041964 ,\\n\",\n \" 0.9286919 ],\\n\",\n \" [ 0.9701384 , -1.3158673 , 0.33630812, 0.80941117,\\n\",\n \" -1.202457 ],\\n\",\n \" [ 1.0198247 , -0.6198277 , 1.0822718 , -1.8385581 ,\\n\",\n \" -0.45790705],\\n\",\n \" [-0.64384323, 0.4564892 , -1.1331053 , -0.68556863,\\n\",\n \" 0.17010891]], dtype=float32)}}\\n\",\n \"Bytes output\\n\",\n \"b'\\\\x81\\\\xa6params\\\\x82\\\\xa4bias\\\\xc7!\\\\x01\\\\x93\\\\x91\\\\x05\\\\xa7float32\\\\xc4\\\\x14\\\\x1d\\\\x1d\\\\xba\\\\xbf\\\\xc4\\\\xad\\\\x01\\\\xc0\\\\x81)\\\\x05@\\\\xdd.\\\\x9c?\\\\xa8\\\\x17\\\\x7f\\\\xbf\\\\xa6kernel\\\\xc7\\\\xd6\\\\x01\\\\x93\\\\x92\\\\n\\\\x05\\\\xa7float32\\\\xc4\\\\xc8\\\\x84]\\\\x81?\\\\xf0\\\\xb5B>`\\\\xb59=z^m\\\\xbfU\\\\xc4\\\\xb1>\\\\x00\\\\xb3\\\\xdd?\\\\xb8x}?\\\\xc7F\\\\x95?2(\\\\x8d?t\\\\x91\\\\xd8\\\\xbd\\\\x83\\\\xb7\\\\x99\\\\xbfr\\\\xa5\\\\x93>#u\\\\xb5?\\\\xdcA\\\\xf7=\\\\xe8\\\\x18\\\\xa8\\\\xbf;\\\\xe5\\\\x98\\\\xbf\\\\xd1}B\\\\xbe0h\\\\n=)\\\\x86\\\\xa8?k\\\\xc2\\\\xa3=\\\\xaaj\\\\x10>\\\\x91\\\\xd8\\\\xaf?\\\\xa9y\\\\xa8\\\\xbfc\\\\xb5\\\\x08?;V\\\\x0f\\\\xc0Av\\\\x10?ZHP?wD\\\\xa3>\\\\x022\\\\t?+Mg?\\\\xa0K\\\\xc2\\\\xbe\\\\xb1\\\\xd3\\\\xde?)\\\\x16\\\\x8a?\\\\x04\\\\x13\\\\x01\\\\xbf\\\\xc1\\\\xbem?\\\\xfdZx?Wn\\\\xa8\\\\xbf\\\\x940\\\\xac>\\\\x925O?\\\\x1c\\\\xea\\\\x99\\\\xbf\\\\x9e\\\\x89\\\\x82?\\\\x07\\\\xad\\\\x1e\\\\xbf\\\\xe2\\\\x87\\\\x8a?\\\\xdfU\\\\xeb\\\\xbf\\\\xcbr\\\\xea\\\\xbe\\\\xe9\\\\xd2$\\\\xbf\\\\xf4\\\\xb8\\\\xe9>\\\\x98\\\\t\\\\x91\\\\xbfm\\\\x81/\\\\xbf\\\\x081.>'\\n\"\n ]\n }\n ],\n \"source\": [\n \"from flax import serialization\\n\",\n \"bytes_output = serialization.to_bytes(params)\\n\",\n \"dict_output = serialization.to_state_dict(params)\\n\",\n \"print('Dict output')\\n\",\n \"print(dict_output)\\n\",\n \"print('Bytes output')\\n\",\n \"print(bytes_output)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"To load the model back, you'll need to use a template of the model parameter structure, like the one you would get from the model initialization. Here, we use the previously generated `params` as a template. Note that this will produce a new variable structure, and not mutate in-place.\\n\",\n \"\\n\",\n \"*The point of enforcing structure through template is to avoid users issues downstream, so you need to first have the right model that generates the parameters structure.*\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 14,\n \"metadata\": {\n \"outputId\": \"13acc4e1-8757-4554-e2c8-d594ba6e67dc\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"FrozenDict({\\n\",\n \" params: {\\n\",\n \" bias: array([-1.4540135, -2.0262308, 2.0806582, 1.2201802, -0.9964547],\\n\",\n \" dtype=float32),\\n\",\n \" kernel: array([[ 1.0106664 , 0.19014716, 0.04533899, -0.92722285, 0.34720102],\\n\",\n \" [ 1.7320251 , 0.9901233 , 1.1662225 , 1.1027892 , -0.10574618],\\n\",\n \" [-1.2009128 , 0.28837162, 1.4176372 , 0.12073109, -1.3132601 ],\\n\",\n \" [-1.1944956 , -0.18993308, 0.03379077, 1.3165942 , 0.07996067],\\n\",\n \" [ 0.14103189, 1.3737966 , -1.3162128 , 0.53401774, -2.239638 ],\\n\",\n \" [ 0.5643044 , 0.813604 , 0.31888172, 0.5359193 , 0.90352124],\\n\",\n \" [-0.37948322, 1.7408353 , 1.0788013 , -0.5041964 , 0.9286919 ],\\n\",\n \" [ 0.9701384 , -1.3158673 , 0.33630812, 0.80941117, -1.202457 ],\\n\",\n \" [ 1.0198247 , -0.6198277 , 1.0822718 , -1.8385581 , -0.45790705],\\n\",\n \" [-0.64384323, 0.4564892 , -1.1331053 , -0.68556863, 0.17010891]],\\n\",\n \" dtype=float32),\\n\",\n \" },\\n\",\n \"})\"\n ]\n },\n \"execution_count\": 14,\n \"metadata\": {\n \"tags\": []\n },\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"serialization.from_bytes(params, bytes_output)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Defining your own models\\n\",\n \"\\n\",\n \"Flax allows you to define your own models, which should be a bit more complicated than a linear regression. In this section, we'll show you how to build simple models. To do so, you'll need to create subclasses of the base `nn.Module` class.\\n\",\n \"\\n\",\n \"*Keep in mind that we imported* `linen as nn` *and this only works with the new linen API*\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Module basics\\n\",\n \"\\n\",\n \"The base abstraction for models is the `nn.Module` class, and every type of predefined layers in Flax (like the previous `Dense`) is a subclass of `nn.Module`. Let's take a look and start by defining a simple but custom multi-layer perceptron i.e. a sequence of Dense layers interleaved with calls to a non-linear activation function.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 17,\n \"metadata\": {\n \"outputId\": \"b59c679c-d164-4fd6-92db-b50f0d310ec3\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"initialized parameter shapes:\\n\",\n \" {'params': {'layers_0': {'bias': (3,), 'kernel': (4, 3)}, 'layers_1': {'bias': (4,), 'kernel': (3, 4)}, 'layers_2': {'bias': (5,), 'kernel': (4, 5)}}}\\n\",\n \"output:\\n\",\n \" [[ 4.2292815e-02 -4.3807115e-02 2.9323792e-02 6.5492536e-03\\n\",\n \" -1.7147182e-02]\\n\",\n \" [ 1.2967806e-01 -1.4551792e-01 9.4432183e-02 1.2521387e-02\\n\",\n \" -4.5417298e-02]\\n\",\n \" [ 0.0000000e+00 0.0000000e+00 0.0000000e+00 0.0000000e+00\\n\",\n \" 0.0000000e+00]\\n\",\n \" [ 9.3024032e-04 2.7864395e-05 2.4478821e-04 8.1344310e-04\\n\",\n \" -1.0110770e-03]]\\n\"\n ]\n }\n ],\n \"source\": [\n \"class ExplicitMLP(nn.Module):\\n\",\n \" features: Sequence[int]\\n\",\n \"\\n\",\n \" def setup(self):\\n\",\n \" # we automatically know what to do with lists, dicts of submodules\\n\",\n \" self.layers = [nn.Dense(feat) for feat in self.features]\\n\",\n \" # for single submodules, we would just write:\\n\",\n \" # self.layer1 = nn.Dense(feat1)\\n\",\n \"\\n\",\n \" def __call__(self, inputs):\\n\",\n \" x = inputs\\n\",\n \" for i, lyr in enumerate(self.layers):\\n\",\n \" x = lyr(x)\\n\",\n \" if i != len(self.layers) - 1:\\n\",\n \" x = nn.relu(x)\\n\",\n \" return x\\n\",\n \"\\n\",\n \"key1, key2 = random.split(random.key(0), 2)\\n\",\n \"x = random.uniform(key1, (4,4))\\n\",\n \"\\n\",\n \"model = ExplicitMLP(features=[3,4,5])\\n\",\n \"params = model.init(key2, x)\\n\",\n \"y = model.apply(params, x)\\n\",\n \"\\n\",\n \"print('initialized parameter shapes:\\\\n', jax.tree_util.tree_map(jnp.shape, flax.core.unfreeze(params)))\\n\",\n \"print('output:\\\\n', y)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"As we can see, a `nn.Module` subclass is made of:\\n\",\n \"\\n\",\n \"* A collection of data fields (`nn.Module` are Python dataclasses) - here we only have the `features` field of type `Sequence[int]`.\\n\",\n \"* A `setup()` method that is being called at the end of the `__postinit__` where you can register submodules, variables, parameters you will need in your model.\\n\",\n \"* A `__call__` function that returns the output of the model from a given input.\\n\",\n \"* The model structure defines a pytree of parameters following the same tree structure as the model: the params tree contains one `layers_n` sub dict per layer, and each of those contain the parameters of the associated Dense layer. The layout is very explicit.\\n\",\n \"\\n\",\n \"*Note: lists are mostly managed as you would expect (WIP), there are corner cases you should be aware of as pointed out* [here](https://github.com/google/flax/issues/524)\\n\",\n \"\\n\",\n \"Since the module structure and its parameters are not tied to each other, you can't directly call `model(x)` on a given input as it will return an error. The `__call__` function is being wrapped up in the `apply` one, which is the one to call on an input:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 19,\n \"metadata\": {\n \"outputId\": \"4af16ec5-b52a-43b0-fc47-1f8ab25e7058\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"\\\"ExplicitMLP\\\" object has no attribute \\\"layers\\\"\\n\"\n ]\n }\n ],\n \"source\": [\n \"try:\\n\",\n \" y = model(x) # Returns an error\\n\",\n \"except AttributeError as e:\\n\",\n \" print(e)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Since here we have a very simple model, we could have used an alternative (but equivalent) way of declaring the submodules inline in the `__call__` using the `@nn.compact` annotation like so:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 20,\n \"metadata\": {\n \"outputId\": \"183a74ef-f54e-4848-99bf-fee4c174ba6d\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"initialized parameter shapes:\\n\",\n \" {'params': {'layers_0': {'bias': (3,), 'kernel': (4, 3)}, 'layers_1': {'bias': (4,), 'kernel': (3, 4)}, 'layers_2': {'bias': (5,), 'kernel': (4, 5)}}}\\n\",\n \"output:\\n\",\n \" [[ 4.2292815e-02 -4.3807115e-02 2.9323792e-02 6.5492536e-03\\n\",\n \" -1.7147182e-02]\\n\",\n \" [ 1.2967806e-01 -1.4551792e-01 9.4432183e-02 1.2521387e-02\\n\",\n \" -4.5417298e-02]\\n\",\n \" [ 0.0000000e+00 0.0000000e+00 0.0000000e+00 0.0000000e+00\\n\",\n \" 0.0000000e+00]\\n\",\n \" [ 9.3024032e-04 2.7864395e-05 2.4478821e-04 8.1344310e-04\\n\",\n \" -1.0110770e-03]]\\n\"\n ]\n }\n ],\n \"source\": [\n \"class SimpleMLP(nn.Module):\\n\",\n \" features: Sequence[int]\\n\",\n \"\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self, inputs):\\n\",\n \" x = inputs\\n\",\n \" for i, feat in enumerate(self.features):\\n\",\n \" x = nn.Dense(feat, name=f'layers_{i}')(x)\\n\",\n \" if i != len(self.features) - 1:\\n\",\n \" x = nn.relu(x)\\n\",\n \" # providing a name is optional though!\\n\",\n \" # the default autonames would be \\\"Dense_0\\\", \\\"Dense_1\\\", ...\\n\",\n \" return x\\n\",\n \"\\n\",\n \"key1, key2 = random.split(random.key(0), 2)\\n\",\n \"x = random.uniform(key1, (4,4))\\n\",\n \"\\n\",\n \"model = SimpleMLP(features=[3,4,5])\\n\",\n \"params = model.init(key2, x)\\n\",\n \"y = model.apply(params, x)\\n\",\n \"\\n\",\n \"print('initialized parameter shapes:\\\\n', jax.tree_util.tree_map(jnp.shape, flax.core.unfreeze(params)))\\n\",\n \"print('output:\\\\n', y)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"There are, however, a few differences you should be aware of between the two declaration modes:\\n\",\n \"\\n\",\n \"* In `setup`, you are able to name some sublayers and keep them around for further use (e.g. encoder/decoder methods in autoencoders).\\n\",\n \"* If you want to have multiple methods, then you **need** to declare the module using `setup`, as the `@nn.compact` annotation only allows one method to be annotated.\\n\",\n \"* The last initialization will be handled differently. See these notes for more details (TODO: add notes link).\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Module parameters\\n\",\n \"\\n\",\n \"In the previous MLP example, we relied only on predefined layers and operators (`Dense`, `relu`). Let's imagine that you didn't have a Dense layer provided by Flax and you wanted to write it on your own. Here is what it would look like using the `@nn.compact` way to declare a new modules:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 21,\n \"metadata\": {\n \"outputId\": \"83b5fea4-071e-4ea0-8fa8-610e69fb5fd5\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"initialized parameters:\\n\",\n \" FrozenDict({\\n\",\n \" params: {\\n\",\n \" kernel: DeviceArray([[ 0.6503669 , 0.86789787, 0.4604268 ],\\n\",\n \" [ 0.05673932, 0.9909285 , -0.63536596],\\n\",\n \" [ 0.76134115, -0.3250529 , -0.65221626],\\n\",\n \" [-0.82430327, 0.4150194 , 0.19405058]], dtype=float32),\\n\",\n \" bias: DeviceArray([0., 0., 0.], dtype=float32),\\n\",\n \" },\\n\",\n \"})\\n\",\n \"output:\\n\",\n \" [[ 0.5035518 1.8548558 -0.4270195 ]\\n\",\n \" [ 0.0279097 0.5589246 -0.43061772]\\n\",\n \" [ 0.3547128 1.5740999 -0.32865518]\\n\",\n \" [ 0.5264864 1.2928858 0.10089308]]\\n\"\n ]\n }\n ],\n \"source\": [\n \"class SimpleDense(nn.Module):\\n\",\n \" features: int\\n\",\n \" kernel_init: Callable = nn.initializers.lecun_normal()\\n\",\n \" bias_init: Callable = nn.initializers.zeros_init()\\n\",\n \"\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self, inputs):\\n\",\n \" kernel = self.param('kernel',\\n\",\n \" self.kernel_init, # Initialization function\\n\",\n \" (inputs.shape[-1], self.features)) # shape info.\\n\",\n \" y = jnp.dot(inputs, kernel)\\n\",\n \" bias = self.param('bias', self.bias_init, (self.features,))\\n\",\n \" y = y + bias\\n\",\n \" return y\\n\",\n \"\\n\",\n \"key1, key2 = random.split(random.key(0), 2)\\n\",\n \"x = random.uniform(key1, (4,4))\\n\",\n \"\\n\",\n \"model = SimpleDense(features=3)\\n\",\n \"params = model.init(key2, x)\\n\",\n \"y = model.apply(params, x)\\n\",\n \"\\n\",\n \"print('initialized parameters:\\\\n', params)\\n\",\n \"print('output:\\\\n', y)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Here, we see how to both declare and assign a parameter to the model using the `self.param` method. It takes as input `(name, init_fn, *init_args, **init_kwargs)` :\\n\",\n \"\\n\",\n \"* `name` is simply the name of the parameter that will end up in the parameter structure.\\n\",\n \"* `init_fn` is a function with input `(PRNGKey, *init_args, **init_kwargs)` returning an Array, with `init_args` and `init_kwargs` being the arguments needed to call the initialisation function.\\n\",\n \"* `init_args` and `init_kwargs` are the arguments to provide to the initialization function.\\n\",\n \"\\n\",\n \"Such params can also be declared in the `setup` method; it won't be able to use shape inference because Flax is using lazy initialization at the first call site.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Variables and collections of variables\\n\",\n \"\\n\",\n \"As we've seen so far, working with models means working with:\\n\",\n \"\\n\",\n \"* A subclass of `nn.Module`;\\n\",\n \"* A pytree of parameters for the model (typically from `model.init()`);\\n\",\n \"\\n\",\n \"However this is not enough to cover everything that we would need for machine learning, especially neural networks. In some cases, you might want your neural network to keep track of some internal state while it runs (e.g. batch normalization layers). There is a way to declare variables beyond the parameters of the model with the `variable` method.\\n\",\n \"\\n\",\n \"For demonstration purposes, we'll implement a simplified but similar mechanism to batch normalization: we'll store running averages and subtract those to the input at training time. For proper batchnorm, you should use (and look at) the implementation [here](https://github.com/google/flax/blob/main/flax/linen/normalization.py).\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 22,\n \"metadata\": {\n \"outputId\": \"75465fd6-cdc8-497c-a3ec-7f709b5dde7a\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"initialized variables:\\n\",\n \" FrozenDict({\\n\",\n \" batch_stats: {\\n\",\n \" mean: DeviceArray([0., 0., 0., 0., 0.], dtype=float32),\\n\",\n \" },\\n\",\n \" params: {\\n\",\n \" bias: DeviceArray([0., 0., 0., 0., 0.], dtype=float32),\\n\",\n \" },\\n\",\n \"})\\n\",\n \"updated state:\\n\",\n \" FrozenDict({\\n\",\n \" batch_stats: {\\n\",\n \" mean: DeviceArray([[0.01, 0.01, 0.01, 0.01, 0.01]], dtype=float32),\\n\",\n \" },\\n\",\n \"})\\n\"\n ]\n }\n ],\n \"source\": [\n \"class BiasAdderWithRunningMean(nn.Module):\\n\",\n \" decay: float = 0.99\\n\",\n \"\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self, x):\\n\",\n \" # easy pattern to detect if we're initializing via empty variable tree\\n\",\n \" is_initialized = self.has_variable('batch_stats', 'mean')\\n\",\n \" ra_mean = self.variable('batch_stats', 'mean',\\n\",\n \" lambda s: jnp.zeros(s),\\n\",\n \" x.shape[1:])\\n\",\n \" bias = self.param('bias', lambda rng, shape: jnp.zeros(shape), x.shape[1:])\\n\",\n \" if is_initialized:\\n\",\n \" ra_mean.value = self.decay * ra_mean.value + (1.0 - self.decay) * jnp.mean(x, axis=0, keepdims=True)\\n\",\n \"\\n\",\n \" return x - ra_mean.value + bias\\n\",\n \"\\n\",\n \"\\n\",\n \"key1, key2 = random.split(random.key(0), 2)\\n\",\n \"x = jnp.ones((10,5))\\n\",\n \"model = BiasAdderWithRunningMean()\\n\",\n \"variables = model.init(key1, x)\\n\",\n \"print('initialized variables:\\\\n', variables)\\n\",\n \"y, updated_state = model.apply(variables, x, mutable=['batch_stats'])\\n\",\n \"print('updated state:\\\\n', updated_state)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Here, `updated_state` returns only the state variables that are being mutated by the model while applying it on data. To update the variables and get the new parameters of the model, we can use the following pattern:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 23,\n \"metadata\": {\n \"outputId\": \"09a8bdd1-eaf8-401a-cf7c-386a7a5aa87b\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"updated state:\\n\",\n \" FrozenDict({\\n\",\n \" batch_stats: {\\n\",\n \" mean: DeviceArray([[0.01, 0.01, 0.01, 0.01, 0.01]], dtype=float32),\\n\",\n \" },\\n\",\n \"})\\n\",\n \"updated state:\\n\",\n \" FrozenDict({\\n\",\n \" batch_stats: {\\n\",\n \" mean: DeviceArray([[0.0299, 0.0299, 0.0299, 0.0299, 0.0299]], dtype=float32),\\n\",\n \" },\\n\",\n \"})\\n\",\n \"updated state:\\n\",\n \" FrozenDict({\\n\",\n \" batch_stats: {\\n\",\n \" mean: DeviceArray([[0.059601, 0.059601, 0.059601, 0.059601, 0.059601]], dtype=float32),\\n\",\n \" },\\n\",\n \"})\\n\"\n ]\n }\n ],\n \"source\": [\n \"for val in [1.0, 2.0, 3.0]:\\n\",\n \" x = val * jnp.ones((10,5))\\n\",\n \" y, updated_state = model.apply(variables, x, mutable=['batch_stats'])\\n\",\n \" old_state, params = flax.core.pop(variables, 'params')\\n\",\n \" variables = flax.core.freeze({'params': params, **updated_state})\\n\",\n \" print('updated state:\\\\n', updated_state) # Shows only the mutable part\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"From this simplified example, you should be able to derive a full BatchNorm implementation, or any layer involving a state. To finish, let's add an optimizer to see how to play with both parameters updated by an optimizer and state variables.\\n\",\n \"\\n\",\n \"*This example isn't doing anything and is only for demonstration purposes.*\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 29,\n \"metadata\": {\n \"outputId\": \"0906fbab-b866-4956-d231-b1374415d448\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"Updated state: FrozenDict({\\n\",\n \" batch_stats: {\\n\",\n \" mean: DeviceArray([[0.01, 0.01, 0.01, 0.01, 0.01]], dtype=float32),\\n\",\n \" },\\n\",\n \"})\\n\",\n \"Updated state: FrozenDict({\\n\",\n \" batch_stats: {\\n\",\n \" mean: DeviceArray([[0.0199, 0.0199, 0.0199, 0.0199, 0.0199]], dtype=float32),\\n\",\n \" },\\n\",\n \"})\\n\",\n \"Updated state: FrozenDict({\\n\",\n \" batch_stats: {\\n\",\n \" mean: DeviceArray([[0.029701, 0.029701, 0.029701, 0.029701, 0.029701]], dtype=float32),\\n\",\n \" },\\n\",\n \"})\\n\"\n ]\n }\n ],\n \"source\": [\n \"from functools import partial\\n\",\n \"\\n\",\n \"@partial(jax.jit, static_argnums=(0, 1))\\n\",\n \"def update_step(tx, apply_fn, x, opt_state, params, state):\\n\",\n \"\\n\",\n \" def loss(params):\\n\",\n \" y, updated_state = apply_fn({'params': params, **state},\\n\",\n \" x, mutable=list(state.keys()))\\n\",\n \" l = ((x - y) ** 2).sum()\\n\",\n \" return l, updated_state\\n\",\n \"\\n\",\n \" (l, state), grads = jax.value_and_grad(loss, has_aux=True)(params)\\n\",\n \" updates, opt_state = tx.update(grads, opt_state)\\n\",\n \" params = optax.apply_updates(params, updates)\\n\",\n \" return opt_state, params, state\\n\",\n \"\\n\",\n \"x = jnp.ones((10,5))\\n\",\n \"variables = model.init(random.key(0), x)\\n\",\n \"state, params = flax.core.pop(variables, 'params')\\n\",\n \"del variables\\n\",\n \"tx = optax.sgd(learning_rate=0.02)\\n\",\n \"opt_state = tx.init(params)\\n\",\n \"\\n\",\n \"for _ in range(3):\\n\",\n \" opt_state, params, state = update_step(tx, model.apply, x, opt_state, params, state)\\n\",\n \" print('Updated state: ', state)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Note that the above function has a quite verbose signature and it would not actually\\n\",\n \"work with `jax.jit()` because the function arguments are not \\\"valid JAX types\\\".\\n\",\n \"\\n\",\n \"Flax provides a handy wrapper - `TrainState` - that simplifies the above code. Check out [`flax.training.train_state.TrainState`](https://flax.readthedocs.io/en/latest/api_reference/flax.training.html#flax.training.train_state.TrainState) to learn more.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Exporting to Tensorflow's SavedModel with jax2tf\\n\",\n \"\\n\",\n \"JAX released an experimental converter called [jax2tf](https://github.com/jax-ml/jax/tree/main/jax/jax2tf), which allows converting trained Flax models into Tensorflow's SavedModel format (so it can be used for [TF Hub](https://www.tensorflow.org/hub), [TF.lite](https://www.tensorflow.org/lite), [TF.js](https://www.tensorflow.org/js), or other downstream applications). The repository contains more documentation and has various examples for Flax.\"\n ]\n }\n ],\n \"metadata\": {\n \"jupytext\": {\n \"formats\": \"ipynb,md:myst\"\n },\n \"language_info\": {\n \"name\": \"python\",\n \"version\": \"3.8.15\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 0\n}\n" }, { "path": "docs/guides/flax_fundamentals/flax_basics.md", "content": "---\njupytext:\n formats: ipynb,md:myst\n text_representation:\n extension: .md\n format_name: myst\n format_version: 0.13\n jupytext_version: 1.13.8\n---\n\n[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google/flax/blob/main/docs/guides/flax_fundamentals/flax_basics.ipynb)\n[![Open On GitHub](https://img.shields.io/badge/Open-on%20GitHub-blue?logo=GitHub)](https://github.com/google/flax/blob/main/docs/guides/flax_fundamentals/flax_basics.ipynb)\n\n# Flax Basics\n\nThis notebook will walk you through the following workflow:\n\n* Instantiating a model from Flax built-in layers or third-party models.\n* Initializing parameters of the model and manually written training.\n* Using optimizers provided by Flax to ease training.\n* Serialization of parameters and other objects.\n* Creating your own models and managing state.\n\n+++\n\n## Setting up our environment\n\nHere we provide the code needed to set up the environment for our notebook.\n\n```{code-cell}\n:outputId: e30aa464-fa52-4f35-df96-716c68a4b3ee\n:tags: [skip-execution]\n\n# Install the latest JAXlib version.\n!pip install --upgrade -q pip jax jaxlib\n# Install Flax at head:\n!pip install --upgrade -q git+https://github.com/google/flax.git\n```\n\n```{code-cell}\nimport jax\nfrom typing import Any, Callable, Sequence\nfrom jax import random, numpy as jnp\nimport flax\nfrom flax import linen as nn\n```\n\n## Linear regression with Flax\n\nIn the previous *JAX for the impatient* notebook, we finished up with a linear regression example. As we know, linear regression can also be written as a single dense neural network layer, which we will show in the following so that we can compare how it's done.\n\nA dense layer is a layer that has a kernel parameter $W\\in\\mathcal{M}_{m,n}(\\mathbb{R})$ where $m$ is the number of features as an output of the model, and $n$ the dimensionality of the input, and a bias parameter $b\\in\\mathbb{R}^m$. The dense layers returns $Wx+b$ from an input $x\\in\\mathbb{R}^n$.\n\nThis dense layer is already provided by Flax in the `flax.linen` module (here imported as `nn`).\n\n```{code-cell}\n# We create one dense layer instance (taking 'features' parameter as input)\nmodel = nn.Dense(features=5)\n```\n\nLayers (and models in general, we'll use that word from now on) are subclasses of the `linen.Module` class.\n\n### Model parameters & initialization\n\nParameters are not stored with the models themselves. You need to initialize parameters by calling the `init` function, using a PRNGKey and dummy input data.\n\n```{code-cell}\n:outputId: 06feb9d2-db50-4f41-c169-6df4336f43a5\n\nkey1, key2 = random.split(random.key(0))\nx = random.normal(key1, (10,)) # Dummy input data\nparams = model.init(key2, x) # Initialization call\njax.tree_util.tree_map(lambda x: x.shape, params) # Checking output shapes\n```\n\n*Note: JAX and Flax, like NumPy, are row-based systems, meaning that vectors are represented as row vectors and not column vectors. This can be seen in the shape of the kernel here.*\n\nThe result is what we expect: bias and kernel parameters of the correct size. Under the hood:\n\n* The dummy input data `x` is used to trigger shape inference: we only declared the number of features we wanted in the output of the model, not the size of the input. Flax finds out by itself the correct size of the kernel.\n* The random PRNG key is used to trigger the initialization functions (those have default values provided by the module here).\n* Initialization functions are called to generate the initial set of parameters that the model will use. Those are functions that take as arguments `(PRNG Key, shape, dtype)` and return an Array of shape `shape`.\n* The init function returns the initialized set of parameters (you can also get the output of the forward pass on the dummy input with the same syntax by using the `init_with_output` method instead of `init`.\n\n+++\n\nTo conduct a forward pass with the model with a given set of parameters (which are never stored with the model), we just use the `apply` method by providing it the parameters to use as well as the input:\n\n```{code-cell}\n:outputId: 7bbe6bb4-94d5-4574-fbb5-aa0fcd1c84ae\n\nmodel.apply(params, x)\n```\n\n### Gradient descent\n\nIf you jumped here directly without going through the JAX part, here is the linear regression formulation we're going to use: from a set of data points $\\{(x_i,y_i), i\\in \\{1,\\ldots, k\\}, x_i\\in\\mathbb{R}^n,y_i\\in\\mathbb{R}^m\\}$, we try to find a set of parameters $W\\in \\mathcal{M}_{m,n}(\\mathbb{R}), b\\in\\mathbb{R}^m$ such that the function $f_{W,b}(x)=Wx+b$ minimizes the mean squared error:\n\n$$\\mathcal{L}(W,b)\\rightarrow\\frac{1}{k}\\sum_{i=1}^{k} \\frac{1}{2}\\|y_i-f_{W,b}(x_i)\\|^2_2$$\n\nHere, we see that the tuple $(W,b)$ matches the parameters of the Dense layer. We'll perform gradient descent using those. Let's first generate the fake data we'll use. The data is exactly the same as in the JAX part's linear regression pytree example.\n\n```{code-cell}\n:outputId: 6eae59dc-0632-4f53-eac8-c22a7c646a52\n\n# Set problem dimensions.\nn_samples = 20\nx_dim = 10\ny_dim = 5\n\n# Generate random ground truth W and b.\nkey = random.key(0)\nk1, k2 = random.split(key)\nW = random.normal(k1, (x_dim, y_dim))\nb = random.normal(k2, (y_dim,))\n# Store the parameters in a FrozenDict pytree.\ntrue_params = flax.core.freeze({'params': {'bias': b, 'kernel': W}})\n\n# Generate samples with additional noise.\nkey_sample, key_noise = random.split(k1)\nx_samples = random.normal(key_sample, (n_samples, x_dim))\ny_samples = jnp.dot(x_samples, W) + b + 0.1 * random.normal(key_noise,(n_samples, y_dim))\nprint('x shape:', x_samples.shape, '; y shape:', y_samples.shape)\n```\n\nWe copy the same training loop that we used in the JAX pytree linear regression example with `jax.value_and_grad()`, but here we can use `model.apply()` instead of having to define our own feed-forward function (`predict_pytree()` in the [JAX example](https://flax.readthedocs.io/en/latest/guides/jax_for_the_impatient.html#linear-regression-with-pytrees)).\n\n```{code-cell}\n# Same as JAX version but using model.apply().\n@jax.jit\ndef mse(params, x_batched, y_batched):\n # Define the squared loss for a single pair (x,y)\n def squared_error(x, y):\n pred = model.apply(params, x)\n return jnp.inner(y-pred, y-pred) / 2.0\n # Vectorize the previous to compute the average of the loss on all samples.\n return jnp.mean(jax.vmap(squared_error)(x_batched,y_batched), axis=0)\n```\n\nAnd finally perform the gradient descent.\n\n```{code-cell}\n:outputId: 50d975b3-4706-4d8a-c4b8-2629ab8e3ac4\n\nlearning_rate = 0.3 # Gradient step size.\nprint('Loss for \"true\" W,b: ', mse(true_params, x_samples, y_samples))\nloss_grad_fn = jax.value_and_grad(mse)\n\n@jax.jit\ndef update_params(params, learning_rate, grads):\n params = jax.tree_util.tree_map(\n lambda p, g: p - learning_rate * g, params, grads)\n return params\n\nfor i in range(101):\n # Perform one gradient update.\n loss_val, grads = loss_grad_fn(params, x_samples, y_samples)\n params = update_params(params, learning_rate, grads)\n if i % 10 == 0:\n print(f'Loss step {i}: ', loss_val)\n```\n\n### Optimizing with Optax\n\nFlax used to use its own `flax.optim` package for optimization, but with\n[FLIP #1009](https://github.com/google/flax/blob/main/docs/flip/1009-optimizer-api.md)\nthis was deprecated in favor of\n[Optax](https://github.com/deepmind/optax).\n\nBasic usage of Optax is straightforward:\n\n1. Choose an optimization method (e.g. `optax.adam`).\n2. Create optimizer state from parameters (for the Adam optimizer, this state will contain the [momentum values](https://optax.readthedocs.io/en/latest/api.html#optax.adam)).\n3. Compute the gradients of your loss with `jax.value_and_grad()`.\n4. At every iteration, call the Optax `update` function to update the internal\n optimizer state and create an update to the parameters. Then add the update\n to the parameters with Optax's `apply_updates` method.\n\nNote that Optax can do a lot more: it's designed for composing simple gradient\ntransformations into more complex transformations that allows to implement a\nwide range of optimizers. There is also support for changing optimizer\nhyperparameters over time (\"schedules\"), applying different updates to different\nparts of the parameter tree (\"masking\") and much more. For details please refer\nto the\n[official documentation](https://optax.readthedocs.io/en/latest/).\n\n```{code-cell}\nimport optax\ntx = optax.adam(learning_rate=learning_rate)\nopt_state = tx.init(params)\nloss_grad_fn = jax.value_and_grad(mse)\n```\n\n```{code-cell}\n:outputId: eec0c096-1d9e-4b3c-f8e5-942ee63828ec\n\nfor i in range(101):\n loss_val, grads = loss_grad_fn(params, x_samples, y_samples)\n updates, opt_state = tx.update(grads, opt_state)\n params = optax.apply_updates(params, updates)\n if i % 10 == 0:\n print('Loss step {}: '.format(i), loss_val)\n```\n\n### Serializing the result\n\nNow that we're happy with the result of our training, we might want to save the model parameters to load them back later. Flax provides a serialization package to enable you to do that.\n\n```{code-cell}\n:outputId: b97e7d83-3e40-4a80-b1fe-1f6ceff30a0c\n\nfrom flax import serialization\nbytes_output = serialization.to_bytes(params)\ndict_output = serialization.to_state_dict(params)\nprint('Dict output')\nprint(dict_output)\nprint('Bytes output')\nprint(bytes_output)\n```\n\nTo load the model back, you'll need to use a template of the model parameter structure, like the one you would get from the model initialization. Here, we use the previously generated `params` as a template. Note that this will produce a new variable structure, and not mutate in-place.\n\n*The point of enforcing structure through template is to avoid users issues downstream, so you need to first have the right model that generates the parameters structure.*\n\n```{code-cell}\n:outputId: 13acc4e1-8757-4554-e2c8-d594ba6e67dc\n\nserialization.from_bytes(params, bytes_output)\n```\n\n## Defining your own models\n\nFlax allows you to define your own models, which should be a bit more complicated than a linear regression. In this section, we'll show you how to build simple models. To do so, you'll need to create subclasses of the base `nn.Module` class.\n\n*Keep in mind that we imported* `linen as nn` *and this only works with the new linen API*\n\n+++\n\n### Module basics\n\nThe base abstraction for models is the `nn.Module` class, and every type of predefined layers in Flax (like the previous `Dense`) is a subclass of `nn.Module`. Let's take a look and start by defining a simple but custom multi-layer perceptron i.e. a sequence of Dense layers interleaved with calls to a non-linear activation function.\n\n```{code-cell}\n:outputId: b59c679c-d164-4fd6-92db-b50f0d310ec3\n\nclass ExplicitMLP(nn.Module):\n features: Sequence[int]\n\n def setup(self):\n # we automatically know what to do with lists, dicts of submodules\n self.layers = [nn.Dense(feat) for feat in self.features]\n # for single submodules, we would just write:\n # self.layer1 = nn.Dense(feat1)\n\n def __call__(self, inputs):\n x = inputs\n for i, lyr in enumerate(self.layers):\n x = lyr(x)\n if i != len(self.layers) - 1:\n x = nn.relu(x)\n return x\n\nkey1, key2 = random.split(random.key(0), 2)\nx = random.uniform(key1, (4,4))\n\nmodel = ExplicitMLP(features=[3,4,5])\nparams = model.init(key2, x)\ny = model.apply(params, x)\n\nprint('initialized parameter shapes:\\n', jax.tree_util.tree_map(jnp.shape, flax.core.unfreeze(params)))\nprint('output:\\n', y)\n```\n\nAs we can see, a `nn.Module` subclass is made of:\n\n* A collection of data fields (`nn.Module` are Python dataclasses) - here we only have the `features` field of type `Sequence[int]`.\n* A `setup()` method that is being called at the end of the `__postinit__` where you can register submodules, variables, parameters you will need in your model.\n* A `__call__` function that returns the output of the model from a given input.\n* The model structure defines a pytree of parameters following the same tree structure as the model: the params tree contains one `layers_n` sub dict per layer, and each of those contain the parameters of the associated Dense layer. The layout is very explicit.\n\n*Note: lists are mostly managed as you would expect (WIP), there are corner cases you should be aware of as pointed out* [here](https://github.com/google/flax/issues/524)\n\nSince the module structure and its parameters are not tied to each other, you can't directly call `model(x)` on a given input as it will return an error. The `__call__` function is being wrapped up in the `apply` one, which is the one to call on an input:\n\n```{code-cell}\n:outputId: 4af16ec5-b52a-43b0-fc47-1f8ab25e7058\n\ntry:\n y = model(x) # Returns an error\nexcept AttributeError as e:\n print(e)\n```\n\nSince here we have a very simple model, we could have used an alternative (but equivalent) way of declaring the submodules inline in the `__call__` using the `@nn.compact` annotation like so:\n\n```{code-cell}\n:outputId: 183a74ef-f54e-4848-99bf-fee4c174ba6d\n\nclass SimpleMLP(nn.Module):\n features: Sequence[int]\n\n @nn.compact\n def __call__(self, inputs):\n x = inputs\n for i, feat in enumerate(self.features):\n x = nn.Dense(feat, name=f'layers_{i}')(x)\n if i != len(self.features) - 1:\n x = nn.relu(x)\n # providing a name is optional though!\n # the default autonames would be \"Dense_0\", \"Dense_1\", ...\n return x\n\nkey1, key2 = random.split(random.key(0), 2)\nx = random.uniform(key1, (4,4))\n\nmodel = SimpleMLP(features=[3,4,5])\nparams = model.init(key2, x)\ny = model.apply(params, x)\n\nprint('initialized parameter shapes:\\n', jax.tree_util.tree_map(jnp.shape, flax.core.unfreeze(params)))\nprint('output:\\n', y)\n```\n\nThere are, however, a few differences you should be aware of between the two declaration modes:\n\n* In `setup`, you are able to name some sublayers and keep them around for further use (e.g. encoder/decoder methods in autoencoders).\n* If you want to have multiple methods, then you **need** to declare the module using `setup`, as the `@nn.compact` annotation only allows one method to be annotated.\n* The last initialization will be handled differently. See these notes for more details (TODO: add notes link).\n\n+++\n\n### Module parameters\n\nIn the previous MLP example, we relied only on predefined layers and operators (`Dense`, `relu`). Let's imagine that you didn't have a Dense layer provided by Flax and you wanted to write it on your own. Here is what it would look like using the `@nn.compact` way to declare a new modules:\n\n```{code-cell}\n:outputId: 83b5fea4-071e-4ea0-8fa8-610e69fb5fd5\n\nclass SimpleDense(nn.Module):\n features: int\n kernel_init: Callable = nn.initializers.lecun_normal()\n bias_init: Callable = nn.initializers.zeros_init()\n\n @nn.compact\n def __call__(self, inputs):\n kernel = self.param('kernel',\n self.kernel_init, # Initialization function\n (inputs.shape[-1], self.features)) # shape info.\n y = jnp.dot(inputs, kernel)\n bias = self.param('bias', self.bias_init, (self.features,))\n y = y + bias\n return y\n\nkey1, key2 = random.split(random.key(0), 2)\nx = random.uniform(key1, (4,4))\n\nmodel = SimpleDense(features=3)\nparams = model.init(key2, x)\ny = model.apply(params, x)\n\nprint('initialized parameters:\\n', params)\nprint('output:\\n', y)\n```\n\nHere, we see how to both declare and assign a parameter to the model using the `self.param` method. It takes as input `(name, init_fn, *init_args, **init_kwargs)` :\n\n* `name` is simply the name of the parameter that will end up in the parameter structure.\n* `init_fn` is a function with input `(PRNGKey, *init_args, **init_kwargs)` returning an Array, with `init_args` and `init_kwargs` being the arguments needed to call the initialisation function.\n* `init_args` and `init_kwargs` are the arguments to provide to the initialization function.\n\nSuch params can also be declared in the `setup` method; it won't be able to use shape inference because Flax is using lazy initialization at the first call site.\n\n+++\n\n### Variables and collections of variables\n\nAs we've seen so far, working with models means working with:\n\n* A subclass of `nn.Module`;\n* A pytree of parameters for the model (typically from `model.init()`);\n\nHowever this is not enough to cover everything that we would need for machine learning, especially neural networks. In some cases, you might want your neural network to keep track of some internal state while it runs (e.g. batch normalization layers). There is a way to declare variables beyond the parameters of the model with the `variable` method.\n\nFor demonstration purposes, we'll implement a simplified but similar mechanism to batch normalization: we'll store running averages and subtract those to the input at training time. For proper batchnorm, you should use (and look at) the implementation [here](https://github.com/google/flax/blob/main/flax/linen/normalization.py).\n\n```{code-cell}\n:outputId: 75465fd6-cdc8-497c-a3ec-7f709b5dde7a\n\nclass BiasAdderWithRunningMean(nn.Module):\n decay: float = 0.99\n\n @nn.compact\n def __call__(self, x):\n # easy pattern to detect if we're initializing via empty variable tree\n is_initialized = self.has_variable('batch_stats', 'mean')\n ra_mean = self.variable('batch_stats', 'mean',\n lambda s: jnp.zeros(s),\n x.shape[1:])\n bias = self.param('bias', lambda rng, shape: jnp.zeros(shape), x.shape[1:])\n if is_initialized:\n ra_mean.value = self.decay * ra_mean.value + (1.0 - self.decay) * jnp.mean(x, axis=0, keepdims=True)\n\n return x - ra_mean.value + bias\n\n\nkey1, key2 = random.split(random.key(0), 2)\nx = jnp.ones((10,5))\nmodel = BiasAdderWithRunningMean()\nvariables = model.init(key1, x)\nprint('initialized variables:\\n', variables)\ny, updated_state = model.apply(variables, x, mutable=['batch_stats'])\nprint('updated state:\\n', updated_state)\n```\n\nHere, `updated_state` returns only the state variables that are being mutated by the model while applying it on data. To update the variables and get the new parameters of the model, we can use the following pattern:\n\n```{code-cell}\n:outputId: 09a8bdd1-eaf8-401a-cf7c-386a7a5aa87b\n\nfor val in [1.0, 2.0, 3.0]:\n x = val * jnp.ones((10,5))\n y, updated_state = model.apply(variables, x, mutable=['batch_stats'])\n old_state, params = flax.core.pop(variables, 'params')\n variables = flax.core.freeze({'params': params, **updated_state})\n print('updated state:\\n', updated_state) # Shows only the mutable part\n```\n\nFrom this simplified example, you should be able to derive a full BatchNorm implementation, or any layer involving a state. To finish, let's add an optimizer to see how to play with both parameters updated by an optimizer and state variables.\n\n*This example isn't doing anything and is only for demonstration purposes.*\n\n```{code-cell}\n:outputId: 0906fbab-b866-4956-d231-b1374415d448\n\nfrom functools import partial\n\n@partial(jax.jit, static_argnums=(0, 1))\ndef update_step(tx, apply_fn, x, opt_state, params, state):\n\n def loss(params):\n y, updated_state = apply_fn({'params': params, **state},\n x, mutable=list(state.keys()))\n l = ((x - y) ** 2).sum()\n return l, updated_state\n\n (l, state), grads = jax.value_and_grad(loss, has_aux=True)(params)\n updates, opt_state = tx.update(grads, opt_state)\n params = optax.apply_updates(params, updates)\n return opt_state, params, state\n\nx = jnp.ones((10,5))\nvariables = model.init(random.key(0), x)\nstate, params = flax.core.pop(variables, 'params')\ndel variables\ntx = optax.sgd(learning_rate=0.02)\nopt_state = tx.init(params)\n\nfor _ in range(3):\n opt_state, params, state = update_step(tx, model.apply, x, opt_state, params, state)\n print('Updated state: ', state)\n```\n\nNote that the above function has a quite verbose signature and it would not actually\nwork with `jax.jit()` because the function arguments are not \"valid JAX types\".\n\nFlax provides a handy wrapper - `TrainState` - that simplifies the above code. Check out [`flax.training.train_state.TrainState`](https://flax.readthedocs.io/en/latest/api_reference/flax.training.html#flax.training.train_state.TrainState) to learn more.\n\n+++\n\n### Exporting to Tensorflow's SavedModel with jax2tf\n\nJAX released an experimental converter called [jax2tf](https://github.com/jax-ml/jax/tree/main/jax/jax2tf), which allows converting trained Flax models into Tensorflow's SavedModel format (so it can be used for [TF Hub](https://www.tensorflow.org/hub), [TF.lite](https://www.tensorflow.org/lite), [TF.js](https://www.tensorflow.org/js), or other downstream applications). The repository contains more documentation and has various examples for Flax.\n" }, { "path": "docs/guides/flax_fundamentals/index.rst", "content": "Flax fundamentals\n=================\n\n.. toctree::\n :maxdepth: 1\n\n JAX 101 \n flax_basics\n state_params\n setup_or_nncompact\n arguments\n rng_guide\n" }, { "path": "docs/guides/flax_fundamentals/rng_guide.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"# Randomness and PRNGs in Flax\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"In this guide, you will learn how Flax uses [JAX's explicit pseudorandom number generator (PRNG) keys](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#jax-prng) to emulate randomness, and adds some additional features to make it easier for users to thread PRNG keys through different Flax `Module`s.\\n\",\n \"\\n\",\n \"If you are new to JAX PRNG keys or need a refresher, check out:\\n\",\n \"- [JAX 101: PRNGs in JAX](https://jax.readthedocs.io/en/latest/jax-101/05-random-numbers.html)\\n\",\n \"- [🔪 JAX - The Sharp Bits 🔪: Random Numbers](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#random-numbers).\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Setup\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Install or upgrade Flax, and then import some necessary dependencies.\\n\",\n \"\\n\",\n \"**Note:** This guide uses the `--xla_force_host_platform_device_count=8` flag to emulate multiple devices in a CPU environment in a Google Colab/Jupyter Notebook. You don’t need this if you are already using a multi-device Google Cloud TPU environment, for example, on Google Cloud or in a Kaggle VM with a TPU.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 1,\n \"metadata\": {\n \"tags\": [\n \"skip-execution\"\n ]\n },\n \"outputs\": [],\n \"source\": [\n \"!pip install -q flax\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 2,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"import os\\n\",\n \"os.environ[\\\"XLA_FLAGS\\\"] = '--xla_force_host_platform_device_count=8'\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 3,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"import flax, flax.linen as nn\\n\",\n \"import jax, jax.numpy as jnp\\n\",\n \"from jax.sharding import Mesh, PartitionSpec, NamedSharding\\n\",\n \"from jax.experimental import mesh_utils\\n\",\n \"from jax.experimental.shard_map import shard_map\\n\",\n \"\\n\",\n \"import hashlib\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 4,\n \"metadata\": {\n \"outputId\": \"ec904f6b-0e87-4efe-87c4-fea0f8e8ec23\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"[CpuDevice(id=0),\\n\",\n \" CpuDevice(id=1),\\n\",\n \" CpuDevice(id=2),\\n\",\n \" CpuDevice(id=3),\\n\",\n \" CpuDevice(id=4),\\n\",\n \" CpuDevice(id=5),\\n\",\n \" CpuDevice(id=6),\\n\",\n \" CpuDevice(id=7)]\"\n ]\n },\n \"execution_count\": 4,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"jax.devices()\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Set the JAX config variable `jax_threefry_partitionable` to `True`. This will be the default value in the future and makes the PRNG more efficiently auto-parallelizable under `jax.jit`. Refer to [JAX discussion](https://github.com/jax-ml/jax/discussions/18480) for more details.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 5,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"jax.config.update('jax_threefry_partitionable', True)\\n\",\n \"assert jax.config.jax_threefry_partitionable == True\\n\",\n \"assert jax.config.jax_default_prng_impl == 'threefry2x32'\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Receiving, manipulating and creating PRNG keys with `Module.make_rng`\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"The primary method Flax uses to receive, manipulate and create PRNG keys is via the `Module` method [`self.make_rng`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.make_rng). It is a method that accepts a string name that represents an \\\"RNG stream\\\". Each RNG stream has an initial starting seed PRNG key, which the user passes in as a dictionary argument (i.e. into an [`.init`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.init) or [`.apply`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.apply) function), and the starting seed is used by `self.make_rng` to generate more PRNG keys for that stream. If `self.make_rng` is called on a string name that does not have an initial starting seed PRNG key (i.e. the user did not pass a key with the corresponding name into `.init` or `.apply`), then `self.make_rng` will use the `'params'` key as the initial starting seed by default.\\n\",\n \"\\n\",\n \"Note that this method can only be called with bounded modules (see [The Flax Module lifecycle](https://flax.readthedocs.io/en/latest/developer_notes/module_lifecycle.html#top-level-modules)).\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 6,\n \"metadata\": {\n \"outputId\": \"2a16435b-e92a-480a-f9fb-e6effc42c4c2\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"Array((), dtype=key) overlaying:\\n\",\n \"[1428664606 3351135085]\\n\",\n \"Array((), dtype=key) overlaying:\\n\",\n \"[3456700291 3873160899]\\n\",\n \"Array((), dtype=key) overlaying:\\n\",\n \"[2411773124 4124888837]\\n\"\n ]\n }\n ],\n \"source\": [\n \"class RNGModule(nn.Module):\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self):\\n\",\n \" print(self.make_rng('rng_stream'))\\n\",\n \" print(self.make_rng('rng_stream'))\\n\",\n \" print(self.make_rng('rng_stream'))\\n\",\n \"\\n\",\n \"rng_module = RNGModule()\\n\",\n \"variables = rng_module.init({'rng_stream': jax.random.key(0)})\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Now if we use a different starting seed PRNG key, we will generate different values (as intended).\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 7,\n \"metadata\": {\n \"outputId\": \"985b0f62-dfde-4f0f-fad4-a31927fc9f59\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"Array((), dtype=key) overlaying:\\n\",\n \"[3077990774 2166202870]\\n\",\n \"Array((), dtype=key) overlaying:\\n\",\n \"[3825832496 2886313970]\\n\",\n \"Array((), dtype=key) overlaying:\\n\",\n \"[ 791337683 1373966058]\\n\"\n ]\n }\n ],\n \"source\": [\n \"variables = rng_module.init({'rng_stream': jax.random.key(1)})\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Calling `self.make_rng` for one stream will not affect the random values generated from another stream; i.e. the call order doesn't matter.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 8,\n \"metadata\": {\n \"outputId\": \"7e8ce538-e380-4db9-db23-bc4a8da577da\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"rng_stream1: Array((), dtype=key) overlaying:\\n\",\n \"[1428664606 3351135085]\\n\",\n \"rng_stream2: Array((), dtype=key) overlaying:\\n\",\n \"[3077990774 2166202870]\\n\",\n \"rng_stream1: Array((), dtype=key) overlaying:\\n\",\n \"[3456700291 3873160899]\\n\",\n \"rng_stream2: Array((), dtype=key) overlaying:\\n\",\n \"[3825832496 2886313970]\\n\",\n \"rng_stream1: Array((), dtype=key) overlaying:\\n\",\n \"[2411773124 4124888837]\\n\",\n \"rng_stream2: Array((), dtype=key) overlaying:\\n\",\n \"[ 791337683 1373966058]\\n\"\n ]\n }\n ],\n \"source\": [\n \"class RNGModuleTwoStreams(nn.Module):\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self):\\n\",\n \" # same value as first code snippet above\\n\",\n \" print(f\\\"rng_stream1: {self.make_rng('rng_stream1')}\\\")\\n\",\n \" # same value as second code snippet above\\n\",\n \" print(f\\\"rng_stream2: {self.make_rng('rng_stream2')}\\\")\\n\",\n \" # same value as first code snippet above\\n\",\n \" print(f\\\"rng_stream1: {self.make_rng('rng_stream1')}\\\")\\n\",\n \" # same value as second code snippet above\\n\",\n \" print(f\\\"rng_stream2: {self.make_rng('rng_stream2')}\\\")\\n\",\n \" # same value as first code snippet above\\n\",\n \" print(f\\\"rng_stream1: {self.make_rng('rng_stream1')}\\\")\\n\",\n \" # same value as second code snippet above\\n\",\n \" print(f\\\"rng_stream2: {self.make_rng('rng_stream2')}\\\")\\n\",\n \"\\n\",\n \"rng_module_two_streams = RNGModuleTwoStreams()\\n\",\n \"variables = rng_module_two_streams.init(\\n\",\n \" {'rng_stream1': jax.random.key(0), 'rng_stream2': jax.random.key(1)}\\n\",\n \")\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Providing the same seed PRNG key will result in the same values being generated (provided that the same operations are used for those keys).\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 9,\n \"metadata\": {\n \"outputId\": \"b70be039-589a-48f7-dc54-65e78c449c65\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"rng_stream1: Array((), dtype=key) overlaying:\\n\",\n \"[1428664606 3351135085]\\n\",\n \"rng_stream2: Array((), dtype=key) overlaying:\\n\",\n \"[1428664606 3351135085]\\n\",\n \"rng_stream1: Array((), dtype=key) overlaying:\\n\",\n \"[3456700291 3873160899]\\n\",\n \"rng_stream2: Array((), dtype=key) overlaying:\\n\",\n \"[3456700291 3873160899]\\n\",\n \"rng_stream1: Array((), dtype=key) overlaying:\\n\",\n \"[2411773124 4124888837]\\n\",\n \"rng_stream2: Array((), dtype=key) overlaying:\\n\",\n \"[2411773124 4124888837]\\n\"\n ]\n }\n ],\n \"source\": [\n \"variables = rng_module_two_streams.init(\\n\",\n \" {'rng_stream1': jax.random.key(0), 'rng_stream2': jax.random.key(0)}\\n\",\n \")\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### How `self.make_rng` works under the hood\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"This is what happens when `self.make_rng` (`flax.linen.Module.make_rng`) is called:\\n\",\n \"* The following data is collected:\\n\",\n \" * The path of the `Module` as provided by `self.scope.path` (the top-level root module has an empty path `()`).\\n\",\n \" * The `self.make_rng` call count. That is, the number of times `self.make_rng` has been called for this specific stream (including this call).\\n\",\n \" * **Note:** Each sub-`Module` will have its own individual call count that's separate from other `Module`s. For example, a `Module` that has called `self.make_rng('params')` twice and contains a sub-`Module` that has called `self.make_rng('params')` once, will have a call count of 2 and 1 for each of the RNG stream `'params'`, respectively.\\n\",\n \"* The data is bundled into a tuple and fed into a hash function and produces an integer.\\n\",\n \"* The generated integer is folded into the RNG stream's starting seed PRNG key to generate a new, unique PRNG key.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Below is a slightly simplified version of the hash function that Flax uses for `self.make_rng`:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 10,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"def produce_hash(data):\\n\",\n \" m = hashlib.sha1()\\n\",\n \" for x in data:\\n\",\n \" if isinstance(x, str):\\n\",\n \" m.update(x.encode('utf-8'))\\n\",\n \" elif isinstance(x, int):\\n\",\n \" m.update(x.to_bytes((x.bit_length() + 7) // 8, byteorder='big'))\\n\",\n \" else:\\n\",\n \" raise ValueError(f'Expected int or string, got: {x}')\\n\",\n \" d = m.digest()\\n\",\n \" hash_int = int.from_bytes(d[:4], byteorder='big')\\n\",\n \" return hash_int\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"And now you can manually reproduce the PRNG keys generated from `self.make_rng`:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 11,\n \"metadata\": {\n \"outputId\": \"d26b7355-9e8b-4954-b2f4-cf7520d5c5a3\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"Array((), dtype=key) overlaying:\\n\",\n \"[1428664606 3351135085]\\n\",\n \"Array((), dtype=key) overlaying:\\n\",\n \"[3456700291 3873160899]\\n\",\n \"Array((), dtype=key) overlaying:\\n\",\n \"[2411773124 4124888837]\\n\"\n ]\n }\n ],\n \"source\": [\n \"stream_seed = jax.random.key(0)\\n\",\n \"for call_count in range(1, 4):\\n\",\n \" hash_int = produce_hash(data=(call_count,))\\n\",\n \" print(jax.random.fold_in(stream_seed, jnp.uint32(hash_int)))\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 12,\n \"metadata\": {\n \"outputId\": \"dec627a6-4c5a-4e3e-ce11-ce4f72775261\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"Array((), dtype=key) overlaying:\\n\",\n \"[1428664606 3351135085]\\n\",\n \"Array((), dtype=key) overlaying:\\n\",\n \"[3456700291 3873160899]\\n\",\n \"Array((), dtype=key) overlaying:\\n\",\n \"[2411773124 4124888837]\\n\"\n ]\n }\n ],\n \"source\": [\n \"variables = rng_module.init({'rng_stream': jax.random.key(0)})\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Sub-`Module`s and `self.make_rng`\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"This section explores how `self.make_rng` (`flax.linen.Module.make_rng`) behaves with sub-`Module`s.\\n\",\n \"\\n\",\n \"Consider the following example:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 13,\n \"metadata\": {\n \"outputId\": \"5b7a9ae9-ca49-4ac0-d007-5caeee739ff0\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"RNGModule, count 1: Array((), dtype=key) overlaying:\\n\",\n \"[1428664606 3351135085]\\n\",\n \"RNGModule, count 2: Array((), dtype=key) overlaying:\\n\",\n \"[3456700291 3873160899]\\n\",\n \"RNGSubModule_0, count 1: Array((), dtype=key) overlaying:\\n\",\n \"[3858825717 2323087578]\\n\",\n \"RNGSubModule_0, count 2: Array((), dtype=key) overlaying:\\n\",\n \"[ 601859108 3782857444]\\n\",\n \"RNGSubSubModule_0, count 1: Array((), dtype=key) overlaying:\\n\",\n \"[ 234240654 1028548813]\\n\",\n \"RNGSubSubModule_0, count 2: Array((), dtype=key) overlaying:\\n\",\n \"[3650462303 2124609379]\\n\"\n ]\n }\n ],\n \"source\": [\n \"class RNGSubSubModule(nn.Module):\\n\",\n \" def __call__(self):\\n\",\n \" print(f\\\"{self.name}, count 1: {self.make_rng('rng_stream')}\\\")\\n\",\n \" print(f\\\"{self.name}, count 2: {self.make_rng('rng_stream')}\\\")\\n\",\n \"\\n\",\n \"class RNGSubModule(nn.Module):\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self):\\n\",\n \" print(f\\\"{self.name}, count 1: {self.make_rng('rng_stream')}\\\")\\n\",\n \" print(f\\\"{self.name}, count 2: {self.make_rng('rng_stream')}\\\")\\n\",\n \" RNGSubSubModule()()\\n\",\n \"\\n\",\n \"class RNGModule(nn.Module):\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self):\\n\",\n \" print(f\\\"RNGModule, count 1: {self.make_rng('rng_stream')}\\\")\\n\",\n \" print(f\\\"RNGModule, count 2: {self.make_rng('rng_stream')}\\\")\\n\",\n \" RNGSubModule()()\\n\",\n \"\\n\",\n \"rng_module = RNGModule()\\n\",\n \"variables = rng_module.init({'rng_stream': jax.random.key(0)})\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"As previously discussed, the data that is fed into the Flax hash function consists of:\\n\",\n \"\\n\",\n \" * The path of the `Module`, provided by `self.scope.path` (the top-level root module has an empty path `()`); and\\n\",\n \" * The call count for the specific RNG stream.\\n\",\n \"\\n\",\n \"In addition, note that each Flax `Module` and sub-`Module` have their own individual call counts, even for the same RNG stream. The convention for sub-`Module` names is: `f'{module_name}_{module_number}'`. For example, the first `Dense` sub-`Module` will be called `Dense_0`, the second one will be called `Dense_1`, and so on.\\n\",\n \"\\n\",\n \"Therefore, the following data will be fed into the hash function:\\n\",\n \"\\n\",\n \" * For `RNGModule`: The data is just the call count, such as `(1,)` and `(2,)`, since the root `Module` has an empty path.\\n\",\n \" * For `RNGSubModule`: The data is `('RNGSubModule_0', 1)` and `('RNGSubModule_0', 2)`.\\n\",\n \" * For `RNGSubSubModule`: The data is `('RNGSubModule_0', 'RNGSubSubModule_0', 1)` and `('RNGSubModule_0', 'RNGSubSubModule_0', 2)`.\\n\",\n \"\\n\",\n \"With this data, you can manually reproduce the PRNG keys generated from the `Module` and sub-`Module`s using `self.make_rng`.\\n\",\n \"\\n\",\n \"For example:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 14,\n \"metadata\": {\n \"outputId\": \"c0de4d37-0f00-4e58-bdfd-e8a6454ed681\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"RNGModule, count 1: Array((), dtype=key) overlaying:\\n\",\n \"[1428664606 3351135085]\\n\",\n \"RNGModule, count 2: Array((), dtype=key) overlaying:\\n\",\n \"[3456700291 3873160899]\\n\",\n \"RNGSubModule_0, count 1: Array((), dtype=key) overlaying:\\n\",\n \"[3858825717 2323087578]\\n\",\n \"RNGSubModule_0, count 2: Array((), dtype=key) overlaying:\\n\",\n \"[ 601859108 3782857444]\\n\",\n \"RNGSubSubModule_0, count 1: Array((), dtype=key) overlaying:\\n\",\n \"[ 234240654 1028548813]\\n\",\n \"RNGSubSubModule_0, count 2: Array((), dtype=key) overlaying:\\n\",\n \"[3650462303 2124609379]\\n\"\n ]\n }\n ],\n \"source\": [\n \"stream_seed = jax.random.key(0)\\n\",\n \"for initial_data in ((), ('RNGSubModule_0',), ('RNGSubModule_0', 'RNGSubSubModule_0')):\\n\",\n \" if initial_data:\\n\",\n \" module_name = initial_data[-1]\\n\",\n \" else:\\n\",\n \" module_name = 'RNGModule'\\n\",\n \" for call_count in (1, 2):\\n\",\n \" hash_int = produce_hash(data=initial_data+(call_count,))\\n\",\n \" rng_key = jax.random.fold_in(stream_seed, jnp.uint32(hash_int))\\n\",\n \" print(f\\\"{module_name}, count {call_count}: {rng_key}\\\")\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"If the same sub-`Module` class is used multiple times, you can increment the suffix of the sub-`Module` name accordingly. For example: `RNGSubModule_0`, `RNGSubModule_1`, and so on.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 15,\n \"metadata\": {\n \"outputId\": \"0b77a038-7000-407b-c5b8-a28dea7951d1\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"RNGModule, count 1: Array((), dtype=key) overlaying:\\n\",\n \"[1428664606 3351135085]\\n\",\n \"RNGModule, count 2: Array((), dtype=key) overlaying:\\n\",\n \"[3456700291 3873160899]\\n\",\n \"RNGSubModule_0, count 1: Array((), dtype=key) overlaying:\\n\",\n \"[3858825717 2323087578]\\n\",\n \"RNGSubModule_0, count 2: Array((), dtype=key) overlaying:\\n\",\n \"[ 601859108 3782857444]\\n\",\n \"RNGSubModule_1, count 1: Array((), dtype=key) overlaying:\\n\",\n \"[ 426957352 2006350344]\\n\",\n \"RNGSubModule_1, count 2: Array((), dtype=key) overlaying:\\n\",\n \"[4006253729 4205356731]\\n\"\n ]\n }\n ],\n \"source\": [\n \"class RNGSubModule(nn.Module):\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self):\\n\",\n \" print(f\\\"{self.name}, count 1: {self.make_rng('rng_stream')}\\\")\\n\",\n \" print(f\\\"{self.name}, count 2: {self.make_rng('rng_stream')}\\\")\\n\",\n \"\\n\",\n \"class RNGModule(nn.Module):\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self):\\n\",\n \" print(f\\\"RNGModule, count 1: {self.make_rng('rng_stream')}\\\")\\n\",\n \" print(f\\\"RNGModule, count 2: {self.make_rng('rng_stream')}\\\")\\n\",\n \" RNGSubModule()()\\n\",\n \" RNGSubModule()()\\n\",\n \"\\n\",\n \"rng_module = RNGModule()\\n\",\n \"variables = rng_module.init({'rng_stream': jax.random.key(0)})\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 16,\n \"metadata\": {\n \"outputId\": \"d189d25e-425d-4fd7-fe18-2dfd63f28b87\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"RNGModule, count 1: Array((), dtype=key) overlaying:\\n\",\n \"[1428664606 3351135085]\\n\",\n \"RNGModule, count 2: Array((), dtype=key) overlaying:\\n\",\n \"[3456700291 3873160899]\\n\",\n \"RNGSubModule_0, count 1: Array((), dtype=key) overlaying:\\n\",\n \"[3858825717 2323087578]\\n\",\n \"RNGSubModule_0, count 2: Array((), dtype=key) overlaying:\\n\",\n \"[ 601859108 3782857444]\\n\",\n \"RNGSubModule_1, count 1: Array((), dtype=key) overlaying:\\n\",\n \"[ 426957352 2006350344]\\n\",\n \"RNGSubModule_1, count 2: Array((), dtype=key) overlaying:\\n\",\n \"[4006253729 4205356731]\\n\"\n ]\n }\n ],\n \"source\": [\n \"stream_seed = jax.random.key(0)\\n\",\n \"for initial_data in ((), ('RNGSubModule_0',), ('RNGSubModule_1',)):\\n\",\n \" if initial_data:\\n\",\n \" module_name = initial_data[-1]\\n\",\n \" else:\\n\",\n \" module_name = 'RNGModule'\\n\",\n \" for call_count in (1, 2):\\n\",\n \" hash_int = produce_hash(data=initial_data+(call_count,))\\n\",\n \" rng_key = jax.random.fold_in(stream_seed, jnp.uint32(hash_int))\\n\",\n \" print(f\\\"{module_name}, count {call_count}: {rng_key}\\\")\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Using `self.param` and `self.variable`\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Flax users have the option of creating additional parameters and variables in their modules by using the [`self.param`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.param) and [`self.variable`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.variable) `Module` methods. An `init_fn` argument must be passed to these methods so that it can generate the initial value of the parameter/variable. `self.make_rng` is commonly used implicitly or explicitly in this `init_fn`, since many initializer functions are stochastic in nature and require a PRNG key. See the full list of Flax initializers [here](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/initializers.html).\\n\",\n \"\\n\",\n \"There are a couple of differences between the two methods that the user should take note of:\\n\",\n \"* `self.param` always creates a parameter in the `'params'` [collection](https://flax.readthedocs.io/en/latest/glossary.html#term-Params-parameters), whereas `self.variable` creates a variable in any [collection](https://flax.readthedocs.io/en/latest/glossary.html#term-Variable-collections) the user specifies\\n\",\n \"* `self.param` will automatically call `self.make_rng('params')` and pass in the generated PRNG key implicitly to the `init_fn` of the parameter you instantiated (it will be passed in as the first argument), whereas users will have to manually specify what RNG stream to call `self.make_rng` on in the `init_fn` of `self.variable` (it could be `'params'` or something different).\\n\",\n \"\\n\",\n \"Below is an example using both `self.param` and `self.variable`:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 17,\n \"metadata\": {\n \"outputId\": \"a7816385-0e08-48e2-dc51-055d7bcd0bab\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"[[-1.6185919 0.700908 ]\\n\",\n \" [-1.3146383 -0.79342234]]\\n\",\n \"[[ 0.0761425 -1.6157459]\\n\",\n \" [-1.6857724 0.7126891]]\\n\",\n \"[[ 0.60175574 0.2553228 ]\\n\",\n \" [ 0.27367848 -2.1975214 ]]\\n\",\n \"[[1.6249592 0.30813068]\\n\",\n \" [1.6613585 1.0404155 ]]\\n\",\n \"[[ 0.0030665 0.29551846]\\n\",\n \" [ 0.16670242 -0.78252524]]\\n\",\n \"[1.582462 0.15216611]\\n\"\n ]\n }\n ],\n \"source\": [\n \"class Model(nn.Module):\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self, x):\\n\",\n \" # kernel will use 'params' seed, initial data will include 'Dense_0', call count 1\\n\",\n \" x = nn.Dense(2, kernel_init=jax.random.normal, use_bias=False)(x)\\n\",\n \" # model_param will use 'params' seed, call count 1\\n\",\n \" model_param = self.param('model_param', jax.random.normal, x.shape)\\n\",\n \" # model_variable1 will use 'params' seed, call count 2\\n\",\n \" model_variable1 = self.variable(\\n\",\n \" 'other_collection',\\n\",\n \" 'model_variable1',\\n\",\n \" lambda: jax.random.normal(self.make_rng('params'), x.shape),\\n\",\n \" )\\n\",\n \" # model_variable2 will use 'other' seed, call count 1\\n\",\n \" model_variable2 = self.variable(\\n\",\n \" 'other_collection',\\n\",\n \" 'model_variable2',\\n\",\n \" lambda: jax.random.normal(self.make_rng('other'), x.shape),\\n\",\n \" )\\n\",\n \" # kernel will use 'params' seed, initial data will include 'Dense_1', call count 1\\n\",\n \" # bias will use 'params' seed, initial data will include 'Dense_1', call count 2\\n\",\n \" x = nn.Dense(2, kernel_init=jax.random.normal, bias_init=jax.random.normal)(\\n\",\n \" x\\n\",\n \" )\\n\",\n \" return x\\n\",\n \"\\n\",\n \"model = Model()\\n\",\n \"variables = model.init(\\n\",\n \" {'params': jax.random.key(0), 'other': jax.random.key(1)}, jnp.ones((2, 2))\\n\",\n \")\\n\",\n \"print(variables['params']['Dense_0']['kernel'])\\n\",\n \"print(variables['params']['model_param'])\\n\",\n \"print(variables['other_collection']['model_variable1'])\\n\",\n \"print(variables['other_collection']['model_variable2'])\\n\",\n \"print(variables['params']['Dense_1']['kernel'])\\n\",\n \"print(variables['params']['Dense_1']['bias'])\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Remember:\\n\",\n \"* there is a separate count for each RNG stream; this is why the count for `self.make_rng('other')` starts at 1 even though there were earlier calls of `self.make_rng('params')`\\n\",\n \"* each submodule has their own separate count for each rng stream; this is why each `Dense` layer has their own separate count for `self.make_rng('params')` and why `model_param` and `model_variable1` share the same count (since they are defined within the same top-level parent module)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 18,\n \"metadata\": {\n \"outputId\": \"ccec9d64-9a27-47f7-adaf-b36a5ea655db\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"[[-1.6185919 0.700908 ]\\n\",\n \" [-1.3146383 -0.79342234]]\\n\",\n \"[[ 0.0761425 -1.6157459]\\n\",\n \" [-1.6857724 0.7126891]]\\n\",\n \"[[ 0.60175574 0.2553228 ]\\n\",\n \" [ 0.27367848 -2.1975214 ]]\\n\",\n \"[[1.6249592 0.30813068]\\n\",\n \" [1.6613585 1.0404155 ]]\\n\",\n \"[[ 0.0030665 0.29551846]\\n\",\n \" [ 0.16670242 -0.78252524]]\\n\",\n \"[[1.582462 0.15216611]]\\n\"\n ]\n }\n ],\n \"source\": [\n \"params_seed = jax.random.key(0)\\n\",\n \"other_seed = jax.random.key(1)\\n\",\n \"for initial_data, count, seed, shape in (\\n\",\n \" (('Dense_0',), 1, params_seed, (2, 2)),\\n\",\n \" ((), 1, params_seed, (2, 2)),\\n\",\n \" ((), 2, params_seed, (2, 2)),\\n\",\n \" ((), 1, other_seed, (2, 2)),\\n\",\n \" (('Dense_1',), 1, params_seed, (2, 2)),\\n\",\n \" (('Dense_1',), 2, params_seed, (1, 2)),\\n\",\n \"):\\n\",\n \" hash_int = produce_hash(data=(*initial_data, count))\\n\",\n \" rng_key = jax.random.fold_in(seed, jnp.uint32(hash_int))\\n\",\n \" print(jax.random.normal(rng_key, shape))\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Managing RNG streams inside a training loop\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Below is an example of managing RNG streams from `self.make_rng`, `self.param`, `self.variable` and `nn.Dropout` in a training loop (note: `nn.Dropout` requires a seed PRNG key to be passed in the `'dropout'` RNG stream, since it implicitly calls `self.make_rng('dropout')`):\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 19,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"class SubModule(nn.Module):\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self, x, train):\\n\",\n \" # variables created using `self.param` will use `self.make_rng('params')`\\n\",\n \" kernel = self.param('submodule_kernel', jax.random.normal, x.shape)\\n\",\n \" x = x + kernel\\n\",\n \" # `nn.Dropout` will use self.make_rng('dropout')\\n\",\n \" x = nn.Dropout(0.2)(x, deterministic=not train)\\n\",\n \" # `nn.Dense` will use self.make_rng('params')\\n\",\n \" x = nn.Dense(3)(x)\\n\",\n \" return x\\n\",\n \"\\n\",\n \"class Model(nn.Module):\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self, x, train):\\n\",\n \" # make kernel use `self.make_rng('other')`\\n\",\n \" kernel = self.variable(\\n\",\n \" 'other_collection',\\n\",\n \" 'module_kernel',\\n\",\n \" lambda: jax.random.normal(self.make_rng('other'), x.shape),\\n\",\n \" )\\n\",\n \" x = (\\n\",\n \" x + kernel.value\\n\",\n \" ) # `.value` will extract the underlying value of the variable\\n\",\n \" x = SubModule()(x, train)\\n\",\n \" # `nn.Dropout` will use self.make_rng('dropout')\\n\",\n \" x = nn.Dropout(0.2)(x, deterministic=not train)\\n\",\n \" # `nn.Dense` will use self.make_rng('params')\\n\",\n \" x = nn.Dense(2)(x)\\n\",\n \" return x\\n\",\n \"\\n\",\n \"params_rng, other_rng, train_rng = jax.random.split(jax.random.key(0), 3)\\n\",\n \"init_rngs = {'params': params_rng, 'other': other_rng}\\n\",\n \"\\n\",\n \"x = jnp.ones((1, 3))\\n\",\n \"y = jnp.ones((1, 2))\\n\",\n \"\\n\",\n \"module = Model()\\n\",\n \"variables = module.init(init_rngs, x, train=False)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 20,\n \"metadata\": {\n \"outputId\": \"e9da8228-acba-403d-bcb5-33a39d4d530d\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"2.518454\\n\",\n \"2.4859657\\n\",\n \"2.4171872\\n\",\n \"2.412684\\n\",\n \"2.3435805\\n\",\n \"2.2773488\\n\",\n \"2.2592616\\n\",\n \"2.2009292\\n\",\n \"2.1839895\\n\",\n \"2.1707344\\n\"\n ]\n }\n ],\n \"source\": [\n \"def update(variables, rng):\\n\",\n \" # we don't need to provide a 'params' or 'other' rng, as only 'dropout' rng will be used during training\\n\",\n \" # split the rng to get a dropout_rng to be used for this training iteration,\\n\",\n \" # and to get another rng key to be used for the next training iteration\\n\",\n \" dropout_rng, next_rng = jax.random.split(rng)\\n\",\n \" def loss(params):\\n\",\n \" out = module.apply(\\n\",\n \" {'params': params, 'other_collection': variables['other_collection']},\\n\",\n \" x,\\n\",\n \" train=True,\\n\",\n \" rngs={'dropout': dropout_rng},\\n\",\n \" )\\n\",\n \" return jnp.mean((y - out) ** 2)\\n\",\n \" grads = jax.grad(loss)(variables['params'])\\n\",\n \" params = jax.tree_util.tree_map(lambda p, g: p - 1e-3 * g, variables['params'], grads)\\n\",\n \" return {\\n\",\n \" 'params': params,\\n\",\n \" 'other_collection': variables['other_collection'],\\n\",\n \" }, next_rng\\n\",\n \"\\n\",\n \"for _ in range(10):\\n\",\n \" variables, train_rng = update(variables, train_rng)\\n\",\n \" out = module.apply(variables, x, train=False)\\n\",\n \" print(jnp.mean((y - out)**2))\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### 🔪 Sharp edge 🔪 - unintentionally generating the same values\\n\",\n \"\\n\",\n \"There is an edge case where the same value can be unintentionally generated.\\n\",\n \"See the [Flax issue](https://github.com/google/flax/issues/2157) for more details.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 21,\n \"metadata\": {\n \"outputId\": \"887142ff-c9ca-4aae-d9fa-cc9993d809c5\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"Array(True, dtype=bool)\"\n ]\n },\n \"execution_count\": 21,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"class Leaf(nn.Module):\\n\",\n \" def __call__(self, x):\\n\",\n \" return x + jax.random.randint(self.make_rng(\\\"rng\\\"), (), 0, 100)\\n\",\n \"\\n\",\n \"class Node(nn.Module):\\n\",\n \" leaf_name: str\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self, x):\\n\",\n \" return Leaf(name=self.leaf_name)(x)\\n\",\n \"\\n\",\n \"class Model(nn.Module):\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self, x):\\n\",\n \" return (Node(name=\\\"ab\\\", leaf_name=\\\"cdef\\\")(x),\\n\",\n \" Node(name=\\\"abc\\\", leaf_name=\\\"def\\\")(x),\\n\",\n \" )\\n\",\n \"\\n\",\n \"out1, out2 = Model().apply({}, 0, rngs={\\\"rng\\\": jax.random.key(33)})\\n\",\n \"out1 == out2 # same output, despite having different submodule names\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"This occurs because the hash function [concatenates strings together](https://docs.python.org/3/library/hashlib.html#hashlib.hash.update), so the data `('AB', 'C')` is equivalent to data `('A', 'BC')` when fed into the hash function, therefore producing the same hash int.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 22,\n \"metadata\": {\n \"outputId\": \"001cbd49-129b-4474-c6a1-3255a4ee3dfe\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"947574064\\n\",\n \"947574064\\n\",\n \"947574064\\n\",\n \"947574064\\n\"\n ]\n }\n ],\n \"source\": [\n \"print(produce_hash(data=('A', 'B', 'C', 1)))\\n\",\n \"print(produce_hash(data=('AB', 'C', 1)))\\n\",\n \"print(produce_hash(data=('A', 'BC', 1)))\\n\",\n \"print(produce_hash(data=('ABC', 1)))\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"To avoid this edge case, users can flip the `flax_fix_rng_separator` [configuration flag](https://flax.readthedocs.io/en/latest/api_reference/flax.config.html#flax.configurations.Config.flax_fix_rng_separator) to `True`.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 23,\n \"metadata\": {\n \"outputId\": \"35a2b204-bdfd-4f83-8e98-ba723963cb0c\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"Array(False, dtype=bool)\"\n ]\n },\n \"execution_count\": 23,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"flax.config.update('flax_fix_rng_separator', True)\\n\",\n \"out1, out2 = Model().apply({}, 0, rngs={\\\"rng\\\": jax.random.key(33)})\\n\",\n \"out1 == out2 # different output\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Managing RNG's on multiple devices\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"This section will show examples on how to use `jit` and `shard_map` to use RNG's in multi-device settings.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Using `jax.jit`\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"When using [`jax.jit`](https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html), we can use RNG's as we did before, but we now include `in_shardings` and `out_shardings` arguments to specify how to shard input and output data. The RNG key itself gets replicated (not sharded); `jax.jit` makes each device use it as appropriate for its shard of the data.\\n\",\n \"\\n\",\n \"For more details on training on multiple devices in Flax using `jax.jit`, see our [Scale up Flax Modules on multiple devices guide](https://flax.readthedocs.io/en/latest/guides/parallel_training/flax_on_pjit.html#) and [lm1b example](https://github.com/google/flax/tree/main/examples/lm1b).\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 24,\n \"metadata\": {\n \"outputId\": \"6c280522-4b43-4b82-f40a-b73986659b2c\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"[CpuDevice(id=0) CpuDevice(id=1) CpuDevice(id=2) CpuDevice(id=3)\\n\",\n \" CpuDevice(id=4) CpuDevice(id=5) CpuDevice(id=6) CpuDevice(id=7)]\\n\",\n \"Mesh(device_ids=array([0, 1, 2, 3, 4, 5, 6, 7]), axis_names=('data',))\\n\",\n \"NamedSharding(mesh=Mesh('data': 8), spec=PartitionSpec('data',))\\n\"\n ]\n }\n ],\n \"source\": [\n \"# Create a mesh and annotate the axis with a name.\\n\",\n \"device_mesh = mesh_utils.create_device_mesh((8,))\\n\",\n \"print(device_mesh)\\n\",\n \"\\n\",\n \"mesh = Mesh(devices=device_mesh, axis_names=('data',))\\n\",\n \"print(mesh)\\n\",\n \"\\n\",\n \"data_sharding = NamedSharding(mesh, PartitionSpec('data',))\\n\",\n \"print(data_sharding)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 25,\n \"metadata\": {\n \"outputId\": \"d1bfbcad-e28a-4fae-8136-98bd1efb9332\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"Array([[-2.2187614 ],\\n\",\n \" [-2.8055234 ],\\n\",\n \" [-2.5464187 ],\\n\",\n \" [ 1.0270392 ],\\n\",\n \" [-3.5243359 ],\\n\",\n \" [-2.2795477 ],\\n\",\n \" [-0.6504516 ],\\n\",\n \" [ 0.17373264]], dtype=float32)\"\n ]\n },\n \"execution_count\": 25,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"class Model(nn.Module):\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self, x, add_noise):\\n\",\n \" x = nn.Dense(1)(x)\\n\",\n \" # use jnp.where for control flow; for more details see: https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerBoolConversionError\\n\",\n \" return jnp.where(\\n\",\n \" add_noise, x + jax.random.normal(self.make_rng('params'), x.shape), x\\n\",\n \" )\\n\",\n \"\\n\",\n \"module = Model()\\n\",\n \"init_rng, apply_rng = jax.random.split(jax.random.key(0))\\n\",\n \"x = jnp.ones((8, 1))\\n\",\n \"variables = module.init(init_rng, x, False)\\n\",\n \"\\n\",\n \"# create custom forward function, since jit does not support kwargs when in_shardings is specified\\n\",\n \"def forward(variables, x, add_noise, rng):\\n\",\n \" return module.apply(variables, x, add_noise, rngs={'params': rng})\\n\",\n \"\\n\",\n \"# shard the inputs x across devices\\n\",\n \"# replicate the variables, add_noise boolean and rng key across devices\\n\",\n \"# shard the output across devices\\n\",\n \"jit_forward = jax.jit(\\n\",\n \" forward,\\n\",\n \" in_shardings=(None, data_sharding, None, None),\\n\",\n \" out_shardings=data_sharding,\\n\",\n \")\\n\",\n \"out = jit_forward(variables, x, True, apply_rng)\\n\",\n \"out\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"The output is different given the same input, meaning the RNG key was used to add noise to the output.\\n\",\n \"\\n\",\n \"We can also confirm that the output is sharded across devices:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 26,\n \"metadata\": {\n \"outputId\": \"b672b85f-7a2d-44b5-afc1-bbf9426655ed\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"[Shard(device=CpuDevice(id=0), index=(slice(0, 1, None), slice(None, None, None)), replica_id=0, data=[[-2.2187614]]),\\n\",\n \" Shard(device=CpuDevice(id=1), index=(slice(1, 2, None), slice(None, None, None)), replica_id=0, data=[[-2.8055234]]),\\n\",\n \" Shard(device=CpuDevice(id=2), index=(slice(2, 3, None), slice(None, None, None)), replica_id=0, data=[[-2.5464187]]),\\n\",\n \" Shard(device=CpuDevice(id=3), index=(slice(3, 4, None), slice(None, None, None)), replica_id=0, data=[[1.0270392]]),\\n\",\n \" Shard(device=CpuDevice(id=4), index=(slice(4, 5, None), slice(None, None, None)), replica_id=0, data=[[-3.5243359]]),\\n\",\n \" Shard(device=CpuDevice(id=5), index=(slice(5, 6, None), slice(None, None, None)), replica_id=0, data=[[-2.2795477]]),\\n\",\n \" Shard(device=CpuDevice(id=6), index=(slice(6, 7, None), slice(None, None, None)), replica_id=0, data=[[-0.6504516]]),\\n\",\n \" Shard(device=CpuDevice(id=7), index=(slice(7, 8, None), slice(None, None, None)), replica_id=0, data=[[0.17373264]])]\"\n ]\n },\n \"execution_count\": 26,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"out.addressable_shards\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Another way to visualize the output sharding:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 27,\n \"metadata\": {\n \"outputId\": \"1c0a16ce-fa3f-4b95-d794-58464bbaa9ae\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/html\": [\n \"
  CPU 0  \\n\",\n       \"         \\n\",\n       \"  CPU 1  \\n\",\n       \"         \\n\",\n       \"  CPU 2  \\n\",\n       \"         \\n\",\n       \"  CPU 3  \\n\",\n       \"         \\n\",\n       \"  CPU 4  \\n\",\n       \"         \\n\",\n       \"  CPU 5  \\n\",\n       \"         \\n\",\n       \"  CPU 6  \\n\",\n       \"         \\n\",\n       \"  CPU 7  \\n\",\n       \"         \\n\",\n       \"
\\n\"\n ],\n \"text/plain\": [\n \"\\u001b[38;2;255;255;255;48;2;57;59;121m \\u001b[0m\\u001b[38;2;255;255;255;48;2;57;59;121mCPU 0\\u001b[0m\\u001b[38;2;255;255;255;48;2;57;59;121m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;57;59;121m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;214;97;107m \\u001b[0m\\u001b[38;2;255;255;255;48;2;214;97;107mCPU 1\\u001b[0m\\u001b[38;2;255;255;255;48;2;214;97;107m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;214;97;107m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;140;162;82m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;162;82mCPU 2\\u001b[0m\\u001b[38;2;255;255;255;48;2;140;162;82m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;140;162;82m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;222;158;214m \\u001b[0m\\u001b[38;2;255;255;255;48;2;222;158;214mCPU 3\\u001b[0m\\u001b[38;2;255;255;255;48;2;222;158;214m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;222;158;214m \\u001b[0m\\n\",\n \"\\u001b[38;2;0;0;0;48;2;231;203;148m \\u001b[0m\\u001b[38;2;0;0;0;48;2;231;203;148mCPU 4\\u001b[0m\\u001b[38;2;0;0;0;48;2;231;203;148m \\u001b[0m\\n\",\n \"\\u001b[38;2;0;0;0;48;2;231;203;148m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;107;110;207m \\u001b[0m\\u001b[38;2;255;255;255;48;2;107;110;207mCPU 5\\u001b[0m\\u001b[38;2;255;255;255;48;2;107;110;207m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;107;110;207m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;165;81;148m \\u001b[0m\\u001b[38;2;255;255;255;48;2;165;81;148mCPU 6\\u001b[0m\\u001b[38;2;255;255;255;48;2;165;81;148m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;165;81;148m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;140;109;49m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;109;49mCPU 7\\u001b[0m\\u001b[38;2;255;255;255;48;2;140;109;49m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;140;109;49m \\u001b[0m\\n\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"jax.debug.visualize_array_sharding(out)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"If we choose not to add noise, then the output is the same across all batches (as expected, since the input is the same for all batches):\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 28,\n \"metadata\": {\n \"outputId\": \"fe9ec875-3e7f-4861-babc-f07064737276\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"Array([[-1.2839764],\\n\",\n \" [-1.2839764],\\n\",\n \" [-1.2839764],\\n\",\n \" [-1.2839764],\\n\",\n \" [-1.2839764],\\n\",\n \" [-1.2839764],\\n\",\n \" [-1.2839764],\\n\",\n \" [-1.2839764]], dtype=float32)\"\n ]\n },\n \"execution_count\": 28,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"out = jit_forward(variables, x, False, apply_rng)\\n\",\n \"out\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"We can confirm the un-jitted function produces the same values, albeit unsharded (note there may be small numerical differences due to compiler optimizations from jitting):\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 29,\n \"metadata\": {\n \"outputId\": \"0a9e5f2c-d4bf-4051-bf71-f32a9c32dc06\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"Array([[-2.2187614 ],\\n\",\n \" [-2.8055234 ],\\n\",\n \" [-2.5464187 ],\\n\",\n \" [ 1.0270392 ],\\n\",\n \" [-3.5243359 ],\\n\",\n \" [-2.2795477 ],\\n\",\n \" [-0.6504516 ],\\n\",\n \" [ 0.17373264]], dtype=float32)\"\n ]\n },\n \"execution_count\": 29,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"out = forward(variables, x, True, apply_rng)\\n\",\n \"out\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 30,\n \"metadata\": {\n \"outputId\": \"772a5063-1bd5-46b4-f6f6-cae9b4b81a26\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"Array([[-1.2839764],\\n\",\n \" [-1.2839764],\\n\",\n \" [-1.2839764],\\n\",\n \" [-1.2839764],\\n\",\n \" [-1.2839764],\\n\",\n \" [-1.2839764],\\n\",\n \" [-1.2839764],\\n\",\n \" [-1.2839764]], dtype=float32)\"\n ]\n },\n \"execution_count\": 30,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"out = forward(variables, x, False, apply_rng)\\n\",\n \"out\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Using `shard_map`\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"When using [`jax.experimental.shard_map.shard_map`](https://jax.readthedocs.io/en/latest/jep/14273-shard-map.html), the important parts to remember are to:\\n\",\n \"* split your PRNG key to produce a different key for each device\\n\",\n \"* the PRNG keys will be sharded automatically to each device (provided you use the correct partition specification), but the [**rank of the original batched PRNG key array will not be reduced**](https://jax.readthedocs.io/en/latest/jep/14273-shard-map.html#rank-reducing-vs-rank-preserving-maps-over-array-axes); e.g.\\n\",\n \"with a batch of 8 PRNG keys and 8 devices, each device will see a PRNG key batch of size 1 within the `shard_map`-ed function\\n\",\n \" * therefore to access the PRNG key itself, we need to index slice into it (see the example below)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 31,\n \"metadata\": {\n \"outputId\": \"aa00b9a3-24ba-4048-ed8c-afbb9070f039\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"Array([[-1.2605132 ],\\n\",\n \" [-1.2405176 ],\\n\",\n \" [-0.99350417],\\n\",\n \" [-1.0277128 ],\\n\",\n \" [-1.4154483 ],\\n\",\n \" [-0.3905797 ],\\n\",\n \" [-2.417677 ],\\n\",\n \" [ 0.9023453 ]], dtype=float32)\"\n ]\n },\n \"execution_count\": 31,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"def forward(variables, x, add_noise, rng_key_batch):\\n\",\n \" # rng_key_batch is a batch of size 1 containing 1 PRNG key\\n\",\n \" # index slice into the rng_key_batch to access the PRNG key\\n\",\n \" return module.apply(\\n\",\n \" variables, x, add_noise, rngs={'params': rng_key_batch[0]}\\n\",\n \" )\\n\",\n \"\\n\",\n \"# define partition specifications\\n\",\n \"data_pspec = PartitionSpec('data')\\n\",\n \"no_pspec = PartitionSpec()\\n\",\n \"\\n\",\n \"# shard the inputs x and rng keys across devices\\n\",\n \"# replicate the variables and add_noise boolean across devices\\n\",\n \"# shard the output across devices\\n\",\n \"shmap_forward = shard_map(\\n\",\n \" forward,\\n\",\n \" mesh=mesh,\\n\",\n \" in_specs=(no_pspec, data_pspec, no_pspec, data_pspec),\\n\",\n \" out_specs=data_pspec,\\n\",\n \")\\n\",\n \"# get 8 different rng's that will be used by the 8 devices when doing forward inference\\n\",\n \"apply_rngs = jax.random.split(apply_rng, 8)\\n\",\n \"out = shmap_forward(variables, x, True, apply_rngs)\\n\",\n \"out\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Confirm that the output is sharded across devices:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 32,\n \"metadata\": {\n \"outputId\": \"e304289b-ef1c-4e4a-d4c1-4c41613bfa62\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"[Shard(device=CpuDevice(id=0), index=(slice(0, 1, None), slice(None, None, None)), replica_id=0, data=[[-1.2605132]]),\\n\",\n \" Shard(device=CpuDevice(id=1), index=(slice(1, 2, None), slice(None, None, None)), replica_id=0, data=[[-1.2405176]]),\\n\",\n \" Shard(device=CpuDevice(id=2), index=(slice(2, 3, None), slice(None, None, None)), replica_id=0, data=[[-0.99350417]]),\\n\",\n \" Shard(device=CpuDevice(id=3), index=(slice(3, 4, None), slice(None, None, None)), replica_id=0, data=[[-1.0277128]]),\\n\",\n \" Shard(device=CpuDevice(id=4), index=(slice(4, 5, None), slice(None, None, None)), replica_id=0, data=[[-1.4154483]]),\\n\",\n \" Shard(device=CpuDevice(id=5), index=(slice(5, 6, None), slice(None, None, None)), replica_id=0, data=[[-0.3905797]]),\\n\",\n \" Shard(device=CpuDevice(id=6), index=(slice(6, 7, None), slice(None, None, None)), replica_id=0, data=[[-2.417677]]),\\n\",\n \" Shard(device=CpuDevice(id=7), index=(slice(7, 8, None), slice(None, None, None)), replica_id=0, data=[[0.9023453]])]\"\n ]\n },\n \"execution_count\": 32,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"out.addressable_shards\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 33,\n \"metadata\": {\n \"outputId\": \"52fdb6d2-4c4f-44b3-feee-4bc5363c8f2f\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/html\": [\n \"
  CPU 0  \\n\",\n       \"         \\n\",\n       \"  CPU 1  \\n\",\n       \"         \\n\",\n       \"  CPU 2  \\n\",\n       \"         \\n\",\n       \"  CPU 3  \\n\",\n       \"         \\n\",\n       \"  CPU 4  \\n\",\n       \"         \\n\",\n       \"  CPU 5  \\n\",\n       \"         \\n\",\n       \"  CPU 6  \\n\",\n       \"         \\n\",\n       \"  CPU 7  \\n\",\n       \"         \\n\",\n       \"
\\n\"\n ],\n \"text/plain\": [\n \"\\u001b[38;2;255;255;255;48;2;57;59;121m \\u001b[0m\\u001b[38;2;255;255;255;48;2;57;59;121mCPU 0\\u001b[0m\\u001b[38;2;255;255;255;48;2;57;59;121m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;57;59;121m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;214;97;107m \\u001b[0m\\u001b[38;2;255;255;255;48;2;214;97;107mCPU 1\\u001b[0m\\u001b[38;2;255;255;255;48;2;214;97;107m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;214;97;107m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;140;162;82m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;162;82mCPU 2\\u001b[0m\\u001b[38;2;255;255;255;48;2;140;162;82m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;140;162;82m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;222;158;214m \\u001b[0m\\u001b[38;2;255;255;255;48;2;222;158;214mCPU 3\\u001b[0m\\u001b[38;2;255;255;255;48;2;222;158;214m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;222;158;214m \\u001b[0m\\n\",\n \"\\u001b[38;2;0;0;0;48;2;231;203;148m \\u001b[0m\\u001b[38;2;0;0;0;48;2;231;203;148mCPU 4\\u001b[0m\\u001b[38;2;0;0;0;48;2;231;203;148m \\u001b[0m\\n\",\n \"\\u001b[38;2;0;0;0;48;2;231;203;148m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;107;110;207m \\u001b[0m\\u001b[38;2;255;255;255;48;2;107;110;207mCPU 5\\u001b[0m\\u001b[38;2;255;255;255;48;2;107;110;207m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;107;110;207m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;165;81;148m \\u001b[0m\\u001b[38;2;255;255;255;48;2;165;81;148mCPU 6\\u001b[0m\\u001b[38;2;255;255;255;48;2;165;81;148m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;165;81;148m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;140;109;49m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;109;49mCPU 7\\u001b[0m\\u001b[38;2;255;255;255;48;2;140;109;49m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;140;109;49m \\u001b[0m\\n\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"jax.debug.visualize_array_sharding(out)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Lifted transforms\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"[Flax lifted transforms](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/transformations.html) allow you to use [JAX transforms](https://github.com/jax-ml/jax#transformations) with `Module` arguments. This section will show you how to control how PRNG keys are split in Flax lifted transforms.\\n\",\n \"\\n\",\n \"Refer to [Lifted transformations](https://flax.readthedocs.io/en/latest/developer_notes/lift.html) for more detail.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### `nn.vmap`\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"We can use [`nn.vmap`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/_autosummary/flax.linen.vmap.html) to create a batched `Dense` layer:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 34,\n \"metadata\": {\n \"outputId\": \"f0830f6b-659c-446f-c933-7b2a430f8004\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"{'params': {'bias': Array([0., 0.], dtype=float32),\\n\",\n \" 'kernel': Array([[-1.2488099 , -0.6127134 ],\\n\",\n \" [-0.07084481, 0.60130936]], dtype=float32)}}\"\n ]\n },\n \"execution_count\": 34,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"x = jnp.ones((3, 2))\\n\",\n \"\\n\",\n \"BatchDense = nn.vmap(\\n\",\n \" nn.Dense,\\n\",\n \" in_axes=0, out_axes=0,\\n\",\n \" variable_axes={'params': None},\\n\",\n \" split_rngs={'params': False})\\n\",\n \"\\n\",\n \"BatchDense(2).init(jax.random.key(0), x)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"By denoting `variable_axes={'params': 0}'`, we vectorize the `params` Arrays on the first axis. However the parameter values generated are all identical to each other:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 35,\n \"metadata\": {\n \"outputId\": \"eef5c0ca-f8d5-4f25-8ce6-9f2f60622daf\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"{'params': {'bias': Array([[0., 0.],\\n\",\n \" [0., 0.],\\n\",\n \" [0., 0.]], dtype=float32),\\n\",\n \" 'kernel': Array([[[-1.2488099 , -0.6127134 ],\\n\",\n \" [-0.07084481, 0.60130936]],\\n\",\n \" \\n\",\n \" [[-1.2488099 , -0.6127134 ],\\n\",\n \" [-0.07084481, 0.60130936]],\\n\",\n \" \\n\",\n \" [[-1.2488099 , -0.6127134 ],\\n\",\n \" [-0.07084481, 0.60130936]]], dtype=float32)}}\"\n ]\n },\n \"execution_count\": 35,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"BatchDense = nn.vmap(\\n\",\n \" nn.Dense,\\n\",\n \" in_axes=0, out_axes=0,\\n\",\n \" variable_axes={'params': 0},\\n\",\n \" split_rngs={'params': False})\\n\",\n \"\\n\",\n \"BatchDense(2).init(jax.random.key(0), x)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"If we also make `split_rngs={'params': True}`, then the PRNG key we provide is split across the variable axis (in this case, the batch axis 0), and we can generate different parameters for each batch input:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 36,\n \"metadata\": {\n \"outputId\": \"275699c3-ba48-403e-877d-07b65981cff5\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"{'params': {'bias': Array([[0., 0.],\\n\",\n \" [0., 0.],\\n\",\n \" [0., 0.]], dtype=float32),\\n\",\n \" 'kernel': Array([[[-0.2526208 , -0.15088455],\\n\",\n \" [-1.1987205 , -0.40843305]],\\n\",\n \" \\n\",\n \" [[-0.7064888 , -1.108805 ],\\n\",\n \" [-0.938775 , 1.4812315 ]],\\n\",\n \" \\n\",\n \" [[-0.59468937, -0.2502723 ],\\n\",\n \" [-1.33515 , 0.5067442 ]]], dtype=float32)}}\"\n ]\n },\n \"execution_count\": 36,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"BatchDense = nn.vmap(\\n\",\n \" nn.Dense,\\n\",\n \" in_axes=0, out_axes=0,\\n\",\n \" variable_axes={'params': 0},\\n\",\n \" split_rngs={'params': True})\\n\",\n \"\\n\",\n \"BatchDense(2).init(jax.random.key(0), x)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Adding a variable via `self.variable` is straightforward:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 37,\n \"metadata\": {\n \"outputId\": \"c11a80bc-d865-4e2e-e059-4d6bcea79e09\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"{'params': {'Dense_0': {'bias': Array([[0., 0.],\\n\",\n \" [0., 0.],\\n\",\n \" [0., 0.]], dtype=float32),\\n\",\n \" 'kernel': Array([[[-0.9079084 , 0.76390624],\\n\",\n \" [-0.01285526, 0.4320353 ]],\\n\",\n \" \\n\",\n \" [[ 0.12398645, 0.7884565 ],\\n\",\n \" [ 1.5344163 , 1.3186085 ]],\\n\",\n \" \\n\",\n \" [[-0.44171348, 0.43430036],\\n\",\n \" [-0.40732604, 0.29774475]]], dtype=float32)}},\\n\",\n \" 'other_collection': {'kernel': Array([[-0.8193048 , 0.711106 ],\\n\",\n \" [-0.37802765, -0.66705877],\\n\",\n \" [-0.44808003, 0.93031347]], dtype=float32)}}\"\n ]\n },\n \"execution_count\": 37,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"class Model(nn.Module):\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self, x):\\n\",\n \" x = nn.Dense(2)(x)\\n\",\n \" kernel = self.variable(\\n\",\n \" 'other_collection',\\n\",\n \" 'kernel',\\n\",\n \" lambda: jax.random.normal(self.make_rng('other'), x.shape),\\n\",\n \" )\\n\",\n \" return x + kernel.value\\n\",\n \"\\n\",\n \"BatchModel = nn.vmap(\\n\",\n \" Model,\\n\",\n \" in_axes=0,\\n\",\n \" out_axes=0,\\n\",\n \" variable_axes={'params': 0, 'other_collection': 0},\\n\",\n \" split_rngs={'params': True, 'other': True},\\n\",\n \")\\n\",\n \"\\n\",\n \"BatchModel().init({'params': jax.random.key(0), 'other': jax.random.key(1)}, x)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"We can control which RNG stream to split, for example, if we only wanted to split the `'params'` RNG stream, then the variables generated from `self.variable` will be the same for each batch input:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 38,\n \"metadata\": {\n \"outputId\": \"fb16619c-c975-497d-c867-6fd5143b4507\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"{'params': {'Dense_0': {'bias': Array([[0., 0.],\\n\",\n \" [0., 0.],\\n\",\n \" [0., 0.]], dtype=float32),\\n\",\n \" 'kernel': Array([[[-0.9079084 , 0.76390624],\\n\",\n \" [-0.01285526, 0.4320353 ]],\\n\",\n \" \\n\",\n \" [[ 0.12398645, 0.7884565 ],\\n\",\n \" [ 1.5344163 , 1.3186085 ]],\\n\",\n \" \\n\",\n \" [[-0.44171348, 0.43430036],\\n\",\n \" [-0.40732604, 0.29774475]]], dtype=float32)}},\\n\",\n \" 'other_collection': {'kernel': Array([[ 0.44956833, -1.1854612 ],\\n\",\n \" [ 0.44956833, -1.1854612 ],\\n\",\n \" [ 0.44956833, -1.1854612 ]], dtype=float32)}}\"\n ]\n },\n \"execution_count\": 38,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"BatchModel = nn.vmap(\\n\",\n \" Model,\\n\",\n \" in_axes=0, out_axes=0,\\n\",\n \" variable_axes={'params': 0, 'other_collection': 0},\\n\",\n \" split_rngs={'params': True, 'other': False})\\n\",\n \"\\n\",\n \"BatchModel().init({'params': jax.random.key(0), 'other': jax.random.key(1)}, x)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"We can also control which parameters / variables should be generated for each batch input, for example, if we only wanted `'params'` to generate separate parameters for each batch input:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 39,\n \"metadata\": {\n \"outputId\": \"f3a17d59-6f75-4408-caba-5769d4589263\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"{'params': {'Dense_0': {'bias': Array([[0., 0.],\\n\",\n \" [0., 0.],\\n\",\n \" [0., 0.]], dtype=float32),\\n\",\n \" 'kernel': Array([[[-0.9079084 , 0.76390624],\\n\",\n \" [-0.01285526, 0.4320353 ]],\\n\",\n \" \\n\",\n \" [[ 0.12398645, 0.7884565 ],\\n\",\n \" [ 1.5344163 , 1.3186085 ]],\\n\",\n \" \\n\",\n \" [[-0.44171348, 0.43430036],\\n\",\n \" [-0.40732604, 0.29774475]]], dtype=float32)}},\\n\",\n \" 'other_collection': {'kernel': Array([ 0.44956833, -1.1854612 ], dtype=float32)}}\"\n ]\n },\n \"execution_count\": 39,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"BatchModel = nn.vmap(\\n\",\n \" Model,\\n\",\n \" in_axes=0, out_axes=0,\\n\",\n \" variable_axes={'params': 0, 'other_collection': None},\\n\",\n \" split_rngs={'params': True, 'other': False})\\n\",\n \"\\n\",\n \"BatchModel().init({'params': jax.random.key(0), 'other': jax.random.key(1)}, x)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### `nn.scan`\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"We can use [`nn.scan`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/_autosummary/flax.linen.scan.html) to create a scanned `Module` layer (this is useful for simplifying repetitively stacked submodules):\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 40,\n \"metadata\": {\n \"outputId\": \"29d1863b-809f-42ce-894c-1b0810faa41e\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"{'params': {'Dense_0': {'bias': Array([[0., 0.],\\n\",\n \" [0., 0.],\\n\",\n \" [0., 0.]], dtype=float32),\\n\",\n \" 'kernel': Array([[[-0.07838312, -0.7422982 ],\\n\",\n \" [ 0.87488323, 0.13773395]],\\n\",\n \" \\n\",\n \" [[ 0.97309333, 0.9087693 ],\\n\",\n \" [-0.12564984, -1.0920651 ]],\\n\",\n \" \\n\",\n \" [[-0.99055105, 1.1499453 ],\\n\",\n \" [-0.15721127, -0.62520015]]], dtype=float32)}}}\"\n ]\n },\n \"execution_count\": 40,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"x = jnp.ones((3, 2))\\n\",\n \"\\n\",\n \"class ResidualMLPBlock(nn.Module):\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self, x, _):\\n\",\n \" h = nn.Dense(features=2)(x)\\n\",\n \" h = nn.relu(h)\\n\",\n \" return x + h, None # return an empty carry\\n\",\n \"\\n\",\n \"ScanMLP = nn.scan(\\n\",\n \" ResidualMLPBlock, variable_axes={'params': 0},\\n\",\n \" variable_broadcast=False, split_rngs={'params': True},\\n\",\n \" length=3)\\n\",\n \"\\n\",\n \"ScanMLP().init(jax.random.key(0), x, None) # pass in an empty carry\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Similar to before, we can control whether to split the RNG stream or not, for example, if we wanted all the stacked modules to be initialized to the same parameter values, we can pass in `split_rngs={'params': False}`:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 41,\n \"metadata\": {\n \"outputId\": \"6a825bcd-9c3b-43c2-afd2-42500d89fb26\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"{'params': {'Dense_0': {'bias': Array([[0., 0.],\\n\",\n \" [0., 0.],\\n\",\n \" [0., 0.]], dtype=float32),\\n\",\n \" 'kernel': Array([[[-0.66715515, -0.0484313 ],\\n\",\n \" [ 0.9867164 , 0.75408363]],\\n\",\n \" \\n\",\n \" [[-0.66715515, -0.0484313 ],\\n\",\n \" [ 0.9867164 , 0.75408363]],\\n\",\n \" \\n\",\n \" [[-0.66715515, -0.0484313 ],\\n\",\n \" [ 0.9867164 , 0.75408363]]], dtype=float32)}}}\"\n ]\n },\n \"execution_count\": 41,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"ScanMLP = nn.scan(\\n\",\n \" ResidualMLPBlock, variable_axes={'params': 0},\\n\",\n \" variable_broadcast=False, split_rngs={'params': False},\\n\",\n \" length=3)\\n\",\n \"\\n\",\n \"ScanMLP().init(jax.random.key(0), x, None)\"\n ]\n }\n ],\n \"metadata\": {\n \"jupytext\": {\n \"formats\": \"ipynb,md:myst\"\n },\n \"language_info\": {\n \"name\": \"python\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 0\n}\n" }, { "path": "docs/guides/flax_fundamentals/rng_guide.md", "content": "---\njupytext:\n formats: ipynb,md:myst\n text_representation:\n extension: .md\n format_name: myst\n format_version: 0.13\n jupytext_version: 1.13.8\n---\n\n# Randomness and PRNGs in Flax\n\n+++\n\nIn this guide, you will learn how Flax uses [JAX's explicit pseudorandom number generator (PRNG) keys](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#jax-prng) to emulate randomness, and adds some additional features to make it easier for users to thread PRNG keys through different Flax `Module`s.\n\nIf you are new to JAX PRNG keys or need a refresher, check out:\n- [JAX 101: PRNGs in JAX](https://jax.readthedocs.io/en/latest/jax-101/05-random-numbers.html)\n- [🔪 JAX - The Sharp Bits 🔪: Random Numbers](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#random-numbers).\n\n+++\n\n## Setup\n\n+++\n\nInstall or upgrade Flax, and then import some necessary dependencies.\n\n**Note:** This guide uses the `--xla_force_host_platform_device_count=8` flag to emulate multiple devices in a CPU environment in a Google Colab/Jupyter Notebook. You don’t need this if you are already using a multi-device Google Cloud TPU environment, for example, on Google Cloud or in a Kaggle VM with a TPU.\n\n```{code-cell}\n:tags: [skip-execution]\n\n!pip install -q flax\n```\n\n```{code-cell}\nimport os\nos.environ[\"XLA_FLAGS\"] = '--xla_force_host_platform_device_count=8'\n```\n\n```{code-cell}\nimport flax, flax.linen as nn\nimport jax, jax.numpy as jnp\nfrom jax.sharding import Mesh, PartitionSpec, NamedSharding\nfrom jax.experimental import mesh_utils\nfrom jax.experimental.shard_map import shard_map\n\nimport hashlib\n```\n\n```{code-cell}\n:outputId: ec904f6b-0e87-4efe-87c4-fea0f8e8ec23\n\njax.devices()\n```\n\nSet the JAX config variable `jax_threefry_partitionable` to `True`. This will be the default value in the future and makes the PRNG more efficiently auto-parallelizable under `jax.jit`. Refer to [JAX discussion](https://github.com/jax-ml/jax/discussions/18480) for more details.\n\n```{code-cell}\njax.config.update('jax_threefry_partitionable', True)\nassert jax.config.jax_threefry_partitionable == True\nassert jax.config.jax_default_prng_impl == 'threefry2x32'\n```\n\n## Receiving, manipulating and creating PRNG keys with `Module.make_rng`\n\n+++\n\nThe primary method Flax uses to receive, manipulate and create PRNG keys is via the `Module` method [`self.make_rng`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.make_rng). It is a method that accepts a string name that represents an \"RNG stream\". Each RNG stream has an initial starting seed PRNG key, which the user passes in as a dictionary argument (i.e. into an [`.init`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.init) or [`.apply`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.apply) function), and the starting seed is used by `self.make_rng` to generate more PRNG keys for that stream. If `self.make_rng` is called on a string name that does not have an initial starting seed PRNG key (i.e. the user did not pass a key with the corresponding name into `.init` or `.apply`), then `self.make_rng` will use the `'params'` key as the initial starting seed by default.\n\nNote that this method can only be called with bounded modules (see [The Flax Module lifecycle](https://flax.readthedocs.io/en/latest/developer_notes/module_lifecycle.html#top-level-modules)).\n\n```{code-cell}\n:outputId: 2a16435b-e92a-480a-f9fb-e6effc42c4c2\n\nclass RNGModule(nn.Module):\n @nn.compact\n def __call__(self):\n print(self.make_rng('rng_stream'))\n print(self.make_rng('rng_stream'))\n print(self.make_rng('rng_stream'))\n\nrng_module = RNGModule()\nvariables = rng_module.init({'rng_stream': jax.random.key(0)})\n```\n\nNow if we use a different starting seed PRNG key, we will generate different values (as intended).\n\n```{code-cell}\n:outputId: 985b0f62-dfde-4f0f-fad4-a31927fc9f59\n\nvariables = rng_module.init({'rng_stream': jax.random.key(1)})\n```\n\nCalling `self.make_rng` for one stream will not affect the random values generated from another stream; i.e. the call order doesn't matter.\n\n```{code-cell}\n:outputId: 7e8ce538-e380-4db9-db23-bc4a8da577da\n\nclass RNGModuleTwoStreams(nn.Module):\n @nn.compact\n def __call__(self):\n # same value as first code snippet above\n print(f\"rng_stream1: {self.make_rng('rng_stream1')}\")\n # same value as second code snippet above\n print(f\"rng_stream2: {self.make_rng('rng_stream2')}\")\n # same value as first code snippet above\n print(f\"rng_stream1: {self.make_rng('rng_stream1')}\")\n # same value as second code snippet above\n print(f\"rng_stream2: {self.make_rng('rng_stream2')}\")\n # same value as first code snippet above\n print(f\"rng_stream1: {self.make_rng('rng_stream1')}\")\n # same value as second code snippet above\n print(f\"rng_stream2: {self.make_rng('rng_stream2')}\")\n\nrng_module_two_streams = RNGModuleTwoStreams()\nvariables = rng_module_two_streams.init(\n {'rng_stream1': jax.random.key(0), 'rng_stream2': jax.random.key(1)}\n)\n```\n\nProviding the same seed PRNG key will result in the same values being generated (provided that the same operations are used for those keys).\n\n```{code-cell}\n:outputId: b70be039-589a-48f7-dc54-65e78c449c65\n\nvariables = rng_module_two_streams.init(\n {'rng_stream1': jax.random.key(0), 'rng_stream2': jax.random.key(0)}\n)\n```\n\n### How `self.make_rng` works under the hood\n\n+++\n\nThis is what happens when `self.make_rng` (`flax.linen.Module.make_rng`) is called:\n* The following data is collected:\n * The path of the `Module` as provided by `self.scope.path` (the top-level root module has an empty path `()`).\n * The `self.make_rng` call count. That is, the number of times `self.make_rng` has been called for this specific stream (including this call).\n * **Note:** Each sub-`Module` will have its own individual call count that's separate from other `Module`s. For example, a `Module` that has called `self.make_rng('params')` twice and contains a sub-`Module` that has called `self.make_rng('params')` once, will have a call count of 2 and 1 for each of the RNG stream `'params'`, respectively.\n* The data is bundled into a tuple and fed into a hash function and produces an integer.\n* The generated integer is folded into the RNG stream's starting seed PRNG key to generate a new, unique PRNG key.\n\n+++\n\nBelow is a slightly simplified version of the hash function that Flax uses for `self.make_rng`:\n\n```{code-cell}\ndef produce_hash(data):\n m = hashlib.sha1()\n for x in data:\n if isinstance(x, str):\n m.update(x.encode('utf-8'))\n elif isinstance(x, int):\n m.update(x.to_bytes((x.bit_length() + 7) // 8, byteorder='big'))\n else:\n raise ValueError(f'Expected int or string, got: {x}')\n d = m.digest()\n hash_int = int.from_bytes(d[:4], byteorder='big')\n return hash_int\n```\n\nAnd now you can manually reproduce the PRNG keys generated from `self.make_rng`:\n\n```{code-cell}\n:outputId: d26b7355-9e8b-4954-b2f4-cf7520d5c5a3\n\nstream_seed = jax.random.key(0)\nfor call_count in range(1, 4):\n hash_int = produce_hash(data=(call_count,))\n print(jax.random.fold_in(stream_seed, jnp.uint32(hash_int)))\n```\n\n```{code-cell}\n:outputId: dec627a6-4c5a-4e3e-ce11-ce4f72775261\n\nvariables = rng_module.init({'rng_stream': jax.random.key(0)})\n```\n\n### Sub-`Module`s and `self.make_rng`\n\n+++\n\nThis section explores how `self.make_rng` (`flax.linen.Module.make_rng`) behaves with sub-`Module`s.\n\nConsider the following example:\n\n```{code-cell}\n:outputId: 5b7a9ae9-ca49-4ac0-d007-5caeee739ff0\n\nclass RNGSubSubModule(nn.Module):\n def __call__(self):\n print(f\"{self.name}, count 1: {self.make_rng('rng_stream')}\")\n print(f\"{self.name}, count 2: {self.make_rng('rng_stream')}\")\n\nclass RNGSubModule(nn.Module):\n @nn.compact\n def __call__(self):\n print(f\"{self.name}, count 1: {self.make_rng('rng_stream')}\")\n print(f\"{self.name}, count 2: {self.make_rng('rng_stream')}\")\n RNGSubSubModule()()\n\nclass RNGModule(nn.Module):\n @nn.compact\n def __call__(self):\n print(f\"RNGModule, count 1: {self.make_rng('rng_stream')}\")\n print(f\"RNGModule, count 2: {self.make_rng('rng_stream')}\")\n RNGSubModule()()\n\nrng_module = RNGModule()\nvariables = rng_module.init({'rng_stream': jax.random.key(0)})\n```\n\nAs previously discussed, the data that is fed into the Flax hash function consists of:\n\n * The path of the `Module`, provided by `self.scope.path` (the top-level root module has an empty path `()`); and\n * The call count for the specific RNG stream.\n\nIn addition, note that each Flax `Module` and sub-`Module` have their own individual call counts, even for the same RNG stream. The convention for sub-`Module` names is: `f'{module_name}_{module_number}'`. For example, the first `Dense` sub-`Module` will be called `Dense_0`, the second one will be called `Dense_1`, and so on.\n\nTherefore, the following data will be fed into the hash function:\n\n * For `RNGModule`: The data is just the call count, such as `(1,)` and `(2,)`, since the root `Module` has an empty path.\n * For `RNGSubModule`: The data is `('RNGSubModule_0', 1)` and `('RNGSubModule_0', 2)`.\n * For `RNGSubSubModule`: The data is `('RNGSubModule_0', 'RNGSubSubModule_0', 1)` and `('RNGSubModule_0', 'RNGSubSubModule_0', 2)`.\n\nWith this data, you can manually reproduce the PRNG keys generated from the `Module` and sub-`Module`s using `self.make_rng`.\n\nFor example:\n\n```{code-cell}\n:outputId: c0de4d37-0f00-4e58-bdfd-e8a6454ed681\n\nstream_seed = jax.random.key(0)\nfor initial_data in ((), ('RNGSubModule_0',), ('RNGSubModule_0', 'RNGSubSubModule_0')):\n if initial_data:\n module_name = initial_data[-1]\n else:\n module_name = 'RNGModule'\n for call_count in (1, 2):\n hash_int = produce_hash(data=initial_data+(call_count,))\n rng_key = jax.random.fold_in(stream_seed, jnp.uint32(hash_int))\n print(f\"{module_name}, count {call_count}: {rng_key}\")\n```\n\nIf the same sub-`Module` class is used multiple times, you can increment the suffix of the sub-`Module` name accordingly. For example: `RNGSubModule_0`, `RNGSubModule_1`, and so on.\n\n```{code-cell}\n:outputId: 0b77a038-7000-407b-c5b8-a28dea7951d1\n\nclass RNGSubModule(nn.Module):\n @nn.compact\n def __call__(self):\n print(f\"{self.name}, count 1: {self.make_rng('rng_stream')}\")\n print(f\"{self.name}, count 2: {self.make_rng('rng_stream')}\")\n\nclass RNGModule(nn.Module):\n @nn.compact\n def __call__(self):\n print(f\"RNGModule, count 1: {self.make_rng('rng_stream')}\")\n print(f\"RNGModule, count 2: {self.make_rng('rng_stream')}\")\n RNGSubModule()()\n RNGSubModule()()\n\nrng_module = RNGModule()\nvariables = rng_module.init({'rng_stream': jax.random.key(0)})\n```\n\n```{code-cell}\n:outputId: d189d25e-425d-4fd7-fe18-2dfd63f28b87\n\nstream_seed = jax.random.key(0)\nfor initial_data in ((), ('RNGSubModule_0',), ('RNGSubModule_1',)):\n if initial_data:\n module_name = initial_data[-1]\n else:\n module_name = 'RNGModule'\n for call_count in (1, 2):\n hash_int = produce_hash(data=initial_data+(call_count,))\n rng_key = jax.random.fold_in(stream_seed, jnp.uint32(hash_int))\n print(f\"{module_name}, count {call_count}: {rng_key}\")\n```\n\n### Using `self.param` and `self.variable`\n\n+++\n\nFlax users have the option of creating additional parameters and variables in their modules by using the [`self.param`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.param) and [`self.variable`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.variable) `Module` methods. An `init_fn` argument must be passed to these methods so that it can generate the initial value of the parameter/variable. `self.make_rng` is commonly used implicitly or explicitly in this `init_fn`, since many initializer functions are stochastic in nature and require a PRNG key. See the full list of Flax initializers [here](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/initializers.html).\n\nThere are a couple of differences between the two methods that the user should take note of:\n* `self.param` always creates a parameter in the `'params'` [collection](https://flax.readthedocs.io/en/latest/glossary.html#term-Params-parameters), whereas `self.variable` creates a variable in any [collection](https://flax.readthedocs.io/en/latest/glossary.html#term-Variable-collections) the user specifies\n* `self.param` will automatically call `self.make_rng('params')` and pass in the generated PRNG key implicitly to the `init_fn` of the parameter you instantiated (it will be passed in as the first argument), whereas users will have to manually specify what RNG stream to call `self.make_rng` on in the `init_fn` of `self.variable` (it could be `'params'` or something different).\n\nBelow is an example using both `self.param` and `self.variable`:\n\n```{code-cell}\n:outputId: a7816385-0e08-48e2-dc51-055d7bcd0bab\n\nclass Model(nn.Module):\n @nn.compact\n def __call__(self, x):\n # kernel will use 'params' seed, initial data will include 'Dense_0', call count 1\n x = nn.Dense(2, kernel_init=jax.random.normal, use_bias=False)(x)\n # model_param will use 'params' seed, call count 1\n model_param = self.param('model_param', jax.random.normal, x.shape)\n # model_variable1 will use 'params' seed, call count 2\n model_variable1 = self.variable(\n 'other_collection',\n 'model_variable1',\n lambda: jax.random.normal(self.make_rng('params'), x.shape),\n )\n # model_variable2 will use 'other' seed, call count 1\n model_variable2 = self.variable(\n 'other_collection',\n 'model_variable2',\n lambda: jax.random.normal(self.make_rng('other'), x.shape),\n )\n # kernel will use 'params' seed, initial data will include 'Dense_1', call count 1\n # bias will use 'params' seed, initial data will include 'Dense_1', call count 2\n x = nn.Dense(2, kernel_init=jax.random.normal, bias_init=jax.random.normal)(\n x\n )\n return x\n\nmodel = Model()\nvariables = model.init(\n {'params': jax.random.key(0), 'other': jax.random.key(1)}, jnp.ones((2, 2))\n)\nprint(variables['params']['Dense_0']['kernel'])\nprint(variables['params']['model_param'])\nprint(variables['other_collection']['model_variable1'])\nprint(variables['other_collection']['model_variable2'])\nprint(variables['params']['Dense_1']['kernel'])\nprint(variables['params']['Dense_1']['bias'])\n```\n\nRemember:\n* there is a separate count for each RNG stream; this is why the count for `self.make_rng('other')` starts at 1 even though there were earlier calls of `self.make_rng('params')`\n* each submodule has their own separate count for each rng stream; this is why each `Dense` layer has their own separate count for `self.make_rng('params')` and why `model_param` and `model_variable1` share the same count (since they are defined within the same top-level parent module)\n\n```{code-cell}\n:outputId: ccec9d64-9a27-47f7-adaf-b36a5ea655db\n\nparams_seed = jax.random.key(0)\nother_seed = jax.random.key(1)\nfor initial_data, count, seed, shape in (\n (('Dense_0',), 1, params_seed, (2, 2)),\n ((), 1, params_seed, (2, 2)),\n ((), 2, params_seed, (2, 2)),\n ((), 1, other_seed, (2, 2)),\n (('Dense_1',), 1, params_seed, (2, 2)),\n (('Dense_1',), 2, params_seed, (1, 2)),\n):\n hash_int = produce_hash(data=(*initial_data, count))\n rng_key = jax.random.fold_in(seed, jnp.uint32(hash_int))\n print(jax.random.normal(rng_key, shape))\n```\n\n### Managing RNG streams inside a training loop\n\n+++\n\nBelow is an example of managing RNG streams from `self.make_rng`, `self.param`, `self.variable` and `nn.Dropout` in a training loop (note: `nn.Dropout` requires a seed PRNG key to be passed in the `'dropout'` RNG stream, since it implicitly calls `self.make_rng('dropout')`):\n\n```{code-cell}\nclass SubModule(nn.Module):\n @nn.compact\n def __call__(self, x, train):\n # variables created using `self.param` will use `self.make_rng('params')`\n kernel = self.param('submodule_kernel', jax.random.normal, x.shape)\n x = x + kernel\n # `nn.Dropout` will use self.make_rng('dropout')\n x = nn.Dropout(0.2)(x, deterministic=not train)\n # `nn.Dense` will use self.make_rng('params')\n x = nn.Dense(3)(x)\n return x\n\nclass Model(nn.Module):\n @nn.compact\n def __call__(self, x, train):\n # make kernel use `self.make_rng('other')`\n kernel = self.variable(\n 'other_collection',\n 'module_kernel',\n lambda: jax.random.normal(self.make_rng('other'), x.shape),\n )\n x = (\n x + kernel.value\n ) # `.value` will extract the underlying value of the variable\n x = SubModule()(x, train)\n # `nn.Dropout` will use self.make_rng('dropout')\n x = nn.Dropout(0.2)(x, deterministic=not train)\n # `nn.Dense` will use self.make_rng('params')\n x = nn.Dense(2)(x)\n return x\n\nparams_rng, other_rng, train_rng = jax.random.split(jax.random.key(0), 3)\ninit_rngs = {'params': params_rng, 'other': other_rng}\n\nx = jnp.ones((1, 3))\ny = jnp.ones((1, 2))\n\nmodule = Model()\nvariables = module.init(init_rngs, x, train=False)\n```\n\n```{code-cell}\n:outputId: e9da8228-acba-403d-bcb5-33a39d4d530d\n\ndef update(variables, rng):\n # we don't need to provide a 'params' or 'other' rng, as only 'dropout' rng will be used during training\n # split the rng to get a dropout_rng to be used for this training iteration,\n # and to get another rng key to be used for the next training iteration\n dropout_rng, next_rng = jax.random.split(rng)\n def loss(params):\n out = module.apply(\n {'params': params, 'other_collection': variables['other_collection']},\n x,\n train=True,\n rngs={'dropout': dropout_rng},\n )\n return jnp.mean((y - out) ** 2)\n grads = jax.grad(loss)(variables['params'])\n params = jax.tree_util.tree_map(lambda p, g: p - 1e-3 * g, variables['params'], grads)\n return {\n 'params': params,\n 'other_collection': variables['other_collection'],\n }, next_rng\n\nfor _ in range(10):\n variables, train_rng = update(variables, train_rng)\n out = module.apply(variables, x, train=False)\n print(jnp.mean((y - out)**2))\n```\n\n### 🔪 Sharp edge 🔪 - unintentionally generating the same values\n\nThere is an edge case where the same value can be unintentionally generated.\nSee the [Flax issue](https://github.com/google/flax/issues/2157) for more details.\n\n```{code-cell}\n:outputId: 887142ff-c9ca-4aae-d9fa-cc9993d809c5\n\nclass Leaf(nn.Module):\n def __call__(self, x):\n return x + jax.random.randint(self.make_rng(\"rng\"), (), 0, 100)\n\nclass Node(nn.Module):\n leaf_name: str\n @nn.compact\n def __call__(self, x):\n return Leaf(name=self.leaf_name)(x)\n\nclass Model(nn.Module):\n @nn.compact\n def __call__(self, x):\n return (Node(name=\"ab\", leaf_name=\"cdef\")(x),\n Node(name=\"abc\", leaf_name=\"def\")(x),\n )\n\nout1, out2 = Model().apply({}, 0, rngs={\"rng\": jax.random.key(33)})\nout1 == out2 # same output, despite having different submodule names\n```\n\nThis occurs because the hash function [concatenates strings together](https://docs.python.org/3/library/hashlib.html#hashlib.hash.update), so the data `('AB', 'C')` is equivalent to data `('A', 'BC')` when fed into the hash function, therefore producing the same hash int.\n\n```{code-cell}\n:outputId: 001cbd49-129b-4474-c6a1-3255a4ee3dfe\n\nprint(produce_hash(data=('A', 'B', 'C', 1)))\nprint(produce_hash(data=('AB', 'C', 1)))\nprint(produce_hash(data=('A', 'BC', 1)))\nprint(produce_hash(data=('ABC', 1)))\n```\n\nTo avoid this edge case, users can flip the `flax_fix_rng_separator` [configuration flag](https://flax.readthedocs.io/en/latest/api_reference/flax.config.html#flax.configurations.Config.flax_fix_rng_separator) to `True`.\n\n```{code-cell}\n:outputId: 35a2b204-bdfd-4f83-8e98-ba723963cb0c\n\nflax.config.update('flax_fix_rng_separator', True)\nout1, out2 = Model().apply({}, 0, rngs={\"rng\": jax.random.key(33)})\nout1 == out2 # different output\n```\n\n## Managing RNG's on multiple devices\n\n+++\n\nThis section will show examples on how to use `jit` and `shard_map` to use RNG's in multi-device settings.\n\n+++\n\n### Using `jax.jit`\n\n+++\n\nWhen using [`jax.jit`](https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html), we can use RNG's as we did before, but we now include `in_shardings` and `out_shardings` arguments to specify how to shard input and output data. The RNG key itself gets replicated (not sharded); `jax.jit` makes each device use it as appropriate for its shard of the data.\n\nFor more details on training on multiple devices in Flax using `jax.jit`, see our [Scale up Flax Modules on multiple devices guide](https://flax.readthedocs.io/en/latest/guides/parallel_training/flax_on_pjit.html#) and [lm1b example](https://github.com/google/flax/tree/main/examples/lm1b).\n\n```{code-cell}\n:outputId: 6c280522-4b43-4b82-f40a-b73986659b2c\n\n# Create a mesh and annotate the axis with a name.\ndevice_mesh = mesh_utils.create_device_mesh((8,))\nprint(device_mesh)\n\nmesh = Mesh(devices=device_mesh, axis_names=('data',))\nprint(mesh)\n\ndata_sharding = NamedSharding(mesh, PartitionSpec('data',))\nprint(data_sharding)\n```\n\n```{code-cell}\n:outputId: d1bfbcad-e28a-4fae-8136-98bd1efb9332\n\nclass Model(nn.Module):\n @nn.compact\n def __call__(self, x, add_noise):\n x = nn.Dense(1)(x)\n # use jnp.where for control flow; for more details see: https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerBoolConversionError\n return jnp.where(\n add_noise, x + jax.random.normal(self.make_rng('params'), x.shape), x\n )\n\nmodule = Model()\ninit_rng, apply_rng = jax.random.split(jax.random.key(0))\nx = jnp.ones((8, 1))\nvariables = module.init(init_rng, x, False)\n\n# create custom forward function, since jit does not support kwargs when in_shardings is specified\ndef forward(variables, x, add_noise, rng):\n return module.apply(variables, x, add_noise, rngs={'params': rng})\n\n# shard the inputs x across devices\n# replicate the variables, add_noise boolean and rng key across devices\n# shard the output across devices\njit_forward = jax.jit(\n forward,\n in_shardings=(None, data_sharding, None, None),\n out_shardings=data_sharding,\n)\nout = jit_forward(variables, x, True, apply_rng)\nout\n```\n\nThe output is different given the same input, meaning the RNG key was used to add noise to the output.\n\nWe can also confirm that the output is sharded across devices:\n\n```{code-cell}\n:outputId: b672b85f-7a2d-44b5-afc1-bbf9426655ed\n\nout.addressable_shards\n```\n\nAnother way to visualize the output sharding:\n\n```{code-cell}\n:outputId: 1c0a16ce-fa3f-4b95-d794-58464bbaa9ae\n\njax.debug.visualize_array_sharding(out)\n```\n\nIf we choose not to add noise, then the output is the same across all batches (as expected, since the input is the same for all batches):\n\n```{code-cell}\n:outputId: fe9ec875-3e7f-4861-babc-f07064737276\n\nout = jit_forward(variables, x, False, apply_rng)\nout\n```\n\nWe can confirm the un-jitted function produces the same values, albeit unsharded (note there may be small numerical differences due to compiler optimizations from jitting):\n\n```{code-cell}\n:outputId: 0a9e5f2c-d4bf-4051-bf71-f32a9c32dc06\n\nout = forward(variables, x, True, apply_rng)\nout\n```\n\n```{code-cell}\n:outputId: 772a5063-1bd5-46b4-f6f6-cae9b4b81a26\n\nout = forward(variables, x, False, apply_rng)\nout\n```\n\n### Using `shard_map`\n\n+++\n\nWhen using [`jax.experimental.shard_map.shard_map`](https://jax.readthedocs.io/en/latest/jep/14273-shard-map.html), the important parts to remember are to:\n* split your PRNG key to produce a different key for each device\n* the PRNG keys will be sharded automatically to each device (provided you use the correct partition specification), but the [**rank of the original batched PRNG key array will not be reduced**](https://jax.readthedocs.io/en/latest/jep/14273-shard-map.html#rank-reducing-vs-rank-preserving-maps-over-array-axes); e.g.\nwith a batch of 8 PRNG keys and 8 devices, each device will see a PRNG key batch of size 1 within the `shard_map`-ed function\n * therefore to access the PRNG key itself, we need to index slice into it (see the example below)\n\n```{code-cell}\n:outputId: aa00b9a3-24ba-4048-ed8c-afbb9070f039\n\ndef forward(variables, x, add_noise, rng_key_batch):\n # rng_key_batch is a batch of size 1 containing 1 PRNG key\n # index slice into the rng_key_batch to access the PRNG key\n return module.apply(\n variables, x, add_noise, rngs={'params': rng_key_batch[0]}\n )\n\n# define partition specifications\ndata_pspec = PartitionSpec('data')\nno_pspec = PartitionSpec()\n\n# shard the inputs x and rng keys across devices\n# replicate the variables and add_noise boolean across devices\n# shard the output across devices\nshmap_forward = shard_map(\n forward,\n mesh=mesh,\n in_specs=(no_pspec, data_pspec, no_pspec, data_pspec),\n out_specs=data_pspec,\n)\n# get 8 different rng's that will be used by the 8 devices when doing forward inference\napply_rngs = jax.random.split(apply_rng, 8)\nout = shmap_forward(variables, x, True, apply_rngs)\nout\n```\n\nConfirm that the output is sharded across devices:\n\n```{code-cell}\n:outputId: e304289b-ef1c-4e4a-d4c1-4c41613bfa62\n\nout.addressable_shards\n```\n\n```{code-cell}\n:outputId: 52fdb6d2-4c4f-44b3-feee-4bc5363c8f2f\n\njax.debug.visualize_array_sharding(out)\n```\n\n## Lifted transforms\n\n+++\n\n[Flax lifted transforms](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/transformations.html) allow you to use [JAX transforms](https://github.com/jax-ml/jax#transformations) with `Module` arguments. This section will show you how to control how PRNG keys are split in Flax lifted transforms.\n\nRefer to [Lifted transformations](https://flax.readthedocs.io/en/latest/developer_notes/lift.html) for more detail.\n\n+++\n\n### `nn.vmap`\n\n+++\n\nWe can use [`nn.vmap`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/_autosummary/flax.linen.vmap.html) to create a batched `Dense` layer:\n\n```{code-cell}\n:outputId: f0830f6b-659c-446f-c933-7b2a430f8004\n\nx = jnp.ones((3, 2))\n\nBatchDense = nn.vmap(\n nn.Dense,\n in_axes=0, out_axes=0,\n variable_axes={'params': None},\n split_rngs={'params': False})\n\nBatchDense(2).init(jax.random.key(0), x)\n```\n\nBy denoting `variable_axes={'params': 0}'`, we vectorize the `params` Arrays on the first axis. However the parameter values generated are all identical to each other:\n\n```{code-cell}\n:outputId: eef5c0ca-f8d5-4f25-8ce6-9f2f60622daf\n\nBatchDense = nn.vmap(\n nn.Dense,\n in_axes=0, out_axes=0,\n variable_axes={'params': 0},\n split_rngs={'params': False})\n\nBatchDense(2).init(jax.random.key(0), x)\n```\n\nIf we also make `split_rngs={'params': True}`, then the PRNG key we provide is split across the variable axis (in this case, the batch axis 0), and we can generate different parameters for each batch input:\n\n```{code-cell}\n:outputId: 275699c3-ba48-403e-877d-07b65981cff5\n\nBatchDense = nn.vmap(\n nn.Dense,\n in_axes=0, out_axes=0,\n variable_axes={'params': 0},\n split_rngs={'params': True})\n\nBatchDense(2).init(jax.random.key(0), x)\n```\n\nAdding a variable via `self.variable` is straightforward:\n\n```{code-cell}\n:outputId: c11a80bc-d865-4e2e-e059-4d6bcea79e09\n\nclass Model(nn.Module):\n @nn.compact\n def __call__(self, x):\n x = nn.Dense(2)(x)\n kernel = self.variable(\n 'other_collection',\n 'kernel',\n lambda: jax.random.normal(self.make_rng('other'), x.shape),\n )\n return x + kernel.value\n\nBatchModel = nn.vmap(\n Model,\n in_axes=0,\n out_axes=0,\n variable_axes={'params': 0, 'other_collection': 0},\n split_rngs={'params': True, 'other': True},\n)\n\nBatchModel().init({'params': jax.random.key(0), 'other': jax.random.key(1)}, x)\n```\n\nWe can control which RNG stream to split, for example, if we only wanted to split the `'params'` RNG stream, then the variables generated from `self.variable` will be the same for each batch input:\n\n```{code-cell}\n:outputId: fb16619c-c975-497d-c867-6fd5143b4507\n\nBatchModel = nn.vmap(\n Model,\n in_axes=0, out_axes=0,\n variable_axes={'params': 0, 'other_collection': 0},\n split_rngs={'params': True, 'other': False})\n\nBatchModel().init({'params': jax.random.key(0), 'other': jax.random.key(1)}, x)\n```\n\nWe can also control which parameters / variables should be generated for each batch input, for example, if we only wanted `'params'` to generate separate parameters for each batch input:\n\n```{code-cell}\n:outputId: f3a17d59-6f75-4408-caba-5769d4589263\n\nBatchModel = nn.vmap(\n Model,\n in_axes=0, out_axes=0,\n variable_axes={'params': 0, 'other_collection': None},\n split_rngs={'params': True, 'other': False})\n\nBatchModel().init({'params': jax.random.key(0), 'other': jax.random.key(1)}, x)\n```\n\n### `nn.scan`\n\n+++\n\nWe can use [`nn.scan`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/_autosummary/flax.linen.scan.html) to create a scanned `Module` layer (this is useful for simplifying repetitively stacked submodules):\n\n```{code-cell}\n:outputId: 29d1863b-809f-42ce-894c-1b0810faa41e\n\nx = jnp.ones((3, 2))\n\nclass ResidualMLPBlock(nn.Module):\n @nn.compact\n def __call__(self, x, _):\n h = nn.Dense(features=2)(x)\n h = nn.relu(h)\n return x + h, None # return an empty carry\n\nScanMLP = nn.scan(\n ResidualMLPBlock, variable_axes={'params': 0},\n variable_broadcast=False, split_rngs={'params': True},\n length=3)\n\nScanMLP().init(jax.random.key(0), x, None) # pass in an empty carry\n```\n\nSimilar to before, we can control whether to split the RNG stream or not, for example, if we wanted all the stacked modules to be initialized to the same parameter values, we can pass in `split_rngs={'params': False}`:\n\n```{code-cell}\n:outputId: 6a825bcd-9c3b-43c2-afd2-42500d89fb26\n\nScanMLP = nn.scan(\n ResidualMLPBlock, variable_axes={'params': 0},\n variable_broadcast=False, split_rngs={'params': False},\n length=3)\n\nScanMLP().init(jax.random.key(0), x, None)\n```\n" }, { "path": "docs/guides/flax_fundamentals/setup_or_nncompact.rst", "content": "``setup`` vs ``compact``\n=========================================\n\nIn Flax's module system (named `Linen`_), submodules and variables (parameters or others)\ncan be defined in two ways:\n\n1. **Explicitly** (using ``setup``):\n\n Assign submodules or variables to ``self.`` inside a\n :meth:`setup ` method. Then use the submodules\n and variables assigned to ``self.`` in ``setup`` from\n any \"forward pass\" method defined on the class.\n This resembles how modules are defined in PyTorch.\n\n2. **In-line** (using ``nn.compact``):\n\n Write your network's logic directly within a single \"forward pass\" method annotated\n with :meth:`nn.compact `. This allows you to define your whole module\n in a single method, and \"co-locate\" submodules and variables next to\n where they are used.\n\n**Both of these approaches are perfectly valid, behave the same way, and interoperate with all of Flax**.\n\nHere is a short example of a module defined in both ways, with exactly\nthe same functionality.\n\n.. testsetup:: Using ``setup``, Using ``nn.compact``\n\n import flax.linen as nn\n\n.. codediff::\n :title: Using ``setup``, Using ``nn.compact``\n\n class MLP(nn.Module):\n def setup(self):\n # Submodule names are derived by the attributes you assign to. In this\n # case, \"dense1\" and \"dense2\". This follows the logic in PyTorch.\n self.dense1 = nn.Dense(32)\n self.dense2 = nn.Dense(32)\n\n def __call__(self, x):\n x = self.dense1(x)\n x = nn.relu(x)\n x = self.dense2(x)\n return x\n ---\n class MLP(nn.Module):\n\n @nn.compact #!\n def __call__(self, x):\n x = nn.Dense(32, name=\"dense1\")(x) #!\n x = nn.relu(x)\n x = nn.Dense(32, name=\"dense2\")(x) #!\n return x\n\nSo, how would you decide which style to use? It can be a matter of taste, but here are some pros and cons:\n\nReasons to prefer using ``nn.compact``:\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n\n1. Allows defining submodules, parameters and other variables next to where they are used: less\n scrolling up/down to see how everything is defined.\n2. Reduces code duplication when there are conditionals or for loops that conditionally define\n submodules, parameters or variables.\n3. Code typically looks more like mathematical notation: ``y = self.param('W', ...) @ x + self.param('b', ...)``\n looks similar to :math:`y=Wx+b``)\n4. If you are using shape inference, i.e. using parameters whose shape/value depend on shapes of\n the inputs (which are unknown at initialization), this is not possible using ``setup``.\n\nReasons to prefer using ``setup``:\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n\n1. Closer to the PyTorch convention, thus easier when porting models\n from PyTorch\n2. Some people find it more natural to explicitly separate the definition\n of submodules and variables from where they are used\n3. Allows defining more than one \"forward pass\" method\n (see :class:`MultipleMethodsCompactError `)\n\n\n\n\n.. _`Linen`: https://jax.readthedocs.io/en/latest/notebooks/thinking_in_jax.html#JIT-mechanics:-tracing-and-static-variables\n" }, { "path": "docs/guides/flax_fundamentals/state_params.rst", "content": "Managing Parameters and State\n=============================\n\nWe will show you how to...\n\n* manage the variables from initialization to updates.\n* split and re-assemble parameters and state.\n* use :code:`vmap` with batch-dependant state.\n\n.. testsetup::\n\n import flax\n from flax import linen as nn\n from jax import random\n import jax.numpy as jnp\n import jax\n import optax\n\n # Create some fake data and run only for one epoch for testing.\n dummy_input = jnp.ones((3, 4))\n num_epochs = 1\n\n.. testcode::\n\n class BiasAdderWithRunningMean(nn.Module):\n momentum: float = 0.9\n\n @nn.compact\n def __call__(self, x):\n is_initialized = self.has_variable('batch_stats', 'mean')\n mean = self.variable('batch_stats', 'mean', jnp.zeros, x.shape[1:])\n bias = self.param('bias', lambda rng, shape: jnp.zeros(shape), x.shape[1:])\n if is_initialized:\n mean.value = (self.momentum * mean.value +\n (1.0 - self.momentum) * jnp.mean(x, axis=0, keepdims=True))\n return mean.value + bias\n\nThis example model is a minimal example that contains both parameters (declared\nwith :code:`self.param`) and state variables (declared with\n:code:`self.variable`).\n\nThe tricky part with initialization here is that we need to split the state\nvariables and the parameters we're going to optimize for.\n\nFirst we define ``update_step`` as follows (with a dummy loss that should be\nreplaced with yours):\n\n.. testcode::\n\n def update_step(apply_fn, x, opt_state, params, state):\n def loss(params):\n y, updated_state = apply_fn({'params': params, **state},\n x, mutable=list(state.keys()))\n l = ((x - y) ** 2).sum() # Replace with your loss here.\n return l, updated_state\n\n (l, updated_state), grads = jax.value_and_grad(\n loss, has_aux=True)(params)\n updates, opt_state = tx.update(grads, opt_state) # Defined below.\n params = optax.apply_updates(params, updates)\n return opt_state, params, updated_state\n\nThen we can write the actual training code.\n\n.. testcode::\n\n model = BiasAdderWithRunningMean()\n variables = model.init(random.key(0), dummy_input)\n # Split state and params (which are updated by optimizer).\n state, params = flax.core.pop(variables, 'params')\n del variables # Delete variables to avoid wasting resources\n tx = optax.sgd(learning_rate=0.02)\n opt_state = tx.init(params)\n\n for _ in range(num_epochs):\n opt_state, params, state = update_step(\n model.apply, dummy_input, opt_state, params, state)\n\n\n:code:`vmap` accross the batch dimension\n----------------------------------------\nWhen using :code:`vmap` and managing state that depends on the batch dimension,\nfor example when using :code:`BatchNorm`, the setup above must be modified\nslightly. This is because any layer whose state depends on the batch dimension\nis not strictly vectorizable. In the case of :code:`BatchNorm`,\n:code:`lax.pmean()` must be used to average the statistics over the batch\ndimension so that the state is in sync for each item in the batch.\n\nThis requires two small changes. Firstly, we need to name the batch axis in our\nmodel definition. Here, this is done by specifying the :code:`axis_name`\nargument of :code:`BatchNorm`. In your own code this might require specifying\nthe :code:`axis_name` argument of :code:`lax.pmean()` directly.\n\n.. testsetup::\n\n from functools import partial\n from flax import linen as nn\n from jax import random\n import jax.numpy as jnp\n import jax\n import optax\n\n # Create some fake data and run only for one epoch for testing.\n dummy_input = jnp.ones((100,))\n key1, key2 = random.split(random.key(0), num=2)\n batch_size = 64\n X = random.normal(key1, (batch_size, 100))\n Y = random.normal(key2, (batch_size, 1))\n num_epochs = 1\n\n.. testcode::\n\n class MLP(nn.Module):\n hidden_size: int\n out_size: int\n\n @nn.compact\n def __call__(self, x, train=False):\n norm = partial(\n nn.BatchNorm,\n use_running_average=not train,\n momentum=0.9,\n epsilon=1e-5,\n axis_name=\"batch\", # Name batch dim\n )\n\n x = nn.Dense(self.hidden_size)(x)\n x = norm()(x)\n x = nn.relu(x)\n x = nn.Dense(self.hidden_size)(x)\n x = norm()(x)\n x = nn.relu(x)\n y = nn.Dense(self.out_size)(x)\n\n return y\n\nSecondly, we need to specify the same name when calling :code:`vmap` in our training code:\n\n.. testcode::\n\n def update_step(apply_fn, x_batch, y_batch, opt_state, params, state):\n\n def batch_loss(params):\n def loss_fn(x, y):\n pred, updated_state = apply_fn(\n {'params': params, **state},\n x, mutable=list(state.keys())\n )\n return (pred - y) ** 2, updated_state\n\n loss, updated_state = jax.vmap(\n loss_fn, out_axes=(0, None), # Do not vmap `updated_state`.\n axis_name='batch' # Name batch dim\n )(x_batch, y_batch) # vmap only `x`, `y`, but not `state`.\n return jnp.mean(loss), updated_state\n\n (loss, updated_state), grads = jax.value_and_grad(\n batch_loss, has_aux=True\n )(params)\n\n updates, opt_state = tx.update(grads, opt_state) # Defined below.\n params = optax.apply_updates(params, updates)\n return opt_state, params, updated_state, loss\n\nNote that we also need to specify that the model state does not have a batch\ndimension. Now we are able to train the model:\n\n.. testcode::\n\n model = MLP(hidden_size=10, out_size=1)\n variables = model.init(random.key(0), dummy_input)\n # Split state and params (which are updated by optimizer).\n state, params = flax.core.pop(variables, 'params')\n del variables # Delete variables to avoid wasting resources\n tx = optax.sgd(learning_rate=0.02)\n opt_state = tx.init(params)\n\n for _ in range(num_epochs):\n opt_state, params, state, loss = update_step(\n model.apply, X, Y, opt_state, params, state)\n" }, { "path": "docs/guides/flax_sharp_bits.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"# 🔪 Flax - The Sharp Bits 🔪\\n\",\n \"\\n\",\n \"Flax exposes the full power of JAX. And just like when using JAX, there are certain _[\\\"sharp bits\\\"](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html)_ you may experience when working with Flax. This evolving document is designed to assist you with them.\\n\",\n \"\\n\",\n \"First, install and/or update Flax:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 1,\n \"metadata\": {\n \"tags\": [\n \"skip-execution\"\n ]\n },\n \"outputs\": [],\n \"source\": [\n \"! pip install -qq flax\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## 🔪 `flax.linen.Dropout` layer and randomness\\n\",\n \"\\n\",\n \"### TL;DR\\n\",\n \"\\n\",\n \"When working on a model with dropout (subclassed from [Flax `Module`](https://flax.readthedocs.io/en/latest/guides/flax_fundamentals/flax_basics.html#module-basics)), add the `'dropout'` PRNGkey only during the forward pass.\\n\",\n \"\\n\",\n \"1. Start with [`jax.random.split()`](https://jax.readthedocs.io/en/latest/_autosummary/jax.random.split.html#jax-random-split) to explicitly create PRNG keys for `'params'` and `'dropout'`.\\n\",\n \"2. Add the [`flax.linen.Dropout`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/layers.html#flax.linen.Dropout) layer(s) to your model (subclassed from Flax [`Module`](https://flax.readthedocs.io/en/latest/guides/flax_fundamentals/flax_basics.html#module-basics)).\\n\",\n \"3. When initializing the model ([`flax.linen.init()`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/init_apply.html)), there's no need to pass in an extra `'dropout'` PRNG key—just the `'params'` key like in a \\\"simpler\\\" model.\\n\",\n \"4. During the forward pass with [`flax.linen.apply()`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/init_apply.html), pass in `rngs={'dropout': dropout_key}`.\\n\",\n \"\\n\",\n \"Check out a full example below.\\n\",\n \"\\n\",\n \"### Why this works\\n\",\n \"\\n\",\n \"- Internally, `flax.linen.Dropout` makes use of [`flax.linen.Module.make_rng`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.make_rng) to create a key for dropout (check out the [source code](https://github.com/google/flax/blob/5714e57a0dc8146eb58a7a06ed768ed3a17672f9/flax/linen/stochastic.py#L72)).\\n\",\n \"- Every time `make_rng` is called (in this case, it's done implicitly in `Dropout`), you get a new PRNG key split from the main/root PRNG key.\\n\",\n \"- `make_rng` still _guarantees full reproducibility_.\\n\",\n \"\\n\",\n \"### Background \\n\",\n \"\\n\",\n \"The [dropout](https://jmlr.org/papers/volume15/srivastava14a/srivastava14a.pdf) stochastic regularization technique randomly removes hidden and visible units in a network. Dropout is a random operation, requiring a PRNG state, and Flax (like JAX) uses [Threefry](https://github.com/jax-ml/jax/blob/main/docs/jep/263-prng.md) PRNG that is splittable. \\n\",\n \"\\n\",\n \"> Note: Recall that JAX has an explicit way of giving you PRNG keys: you can fork the main PRNG state (such as `key = jax.random.key(seed=0)`) into multiple new PRNG keys with `key, subkey = jax.random.split(key)`. Refresh your memory in [🔪 JAX - The Sharp Bits 🔪 Randomness and PRNG keys](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#random-numbers).\\n\",\n \"\\n\",\n \"Flax provides an _implicit_ way of handling PRNG key streams via [Flax `Module`](https://flax.readthedocs.io/en/latest/guides/flax_fundamentals/flax_basics.html#module-basics)'s [`flax.linen.Module.make_rng`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.make_rng) helper function. It allows the code in Flax `Module`s (or its sub-`Module`s) to \\\"pull PRNG keys\\\". `make_rng` guarantees to provide a unique key each time you call it. See the [RNG guide](https://flax.readthedocs.io/en/latest/guides/flax_fundamentals/rng_guide.html) for more details.\\n\",\n \"\\n\",\n \"> Note: Recall that [`flax.linen.Module`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html) is the base class for all neural network modules. All layers and models are subclassed from it.\\n\",\n \"\\n\",\n \"### Example\\n\",\n \"\\n\",\n \"Remember that each of the Flax PRNG streams has a name. The example below uses the `'params'` stream for initializing parameters, as well as the `'dropout'` stream. The PRNG key provided to [`flax.linen.init()`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/init_apply.html) is the one that seeds the `'params'` PRNG key stream. To draw PRNG keys during the forward pass (with dropout), provide a PRNG key to seed that stream (`'dropout'`) when you call `Module.apply()`.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 2,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# Setup.\\n\",\n \"import jax\\n\",\n \"import jax.numpy as jnp\\n\",\n \"import flax.linen as nn\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 3,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# Randomness.\\n\",\n \"seed = 0\\n\",\n \"root_key = jax.random.key(seed=seed)\\n\",\n \"main_key, params_key, dropout_key = jax.random.split(key=root_key, num=3)\\n\",\n \"\\n\",\n \"# A simple network.\\n\",\n \"class MyModel(nn.Module):\\n\",\n \" num_neurons: int\\n\",\n \" training: bool\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self, x):\\n\",\n \" x = nn.Dense(self.num_neurons)(x)\\n\",\n \" # Set the dropout layer with a rate of 50% .\\n\",\n \" # When the `deterministic` flag is `True`, dropout is turned off.\\n\",\n \" x = nn.Dropout(rate=0.5, deterministic=not self.training)(x)\\n\",\n \" return x\\n\",\n \"\\n\",\n \"# Instantiate `MyModel` (you don't need to set `training=True` to\\n\",\n \"# avoid performing the forward pass computation).\\n\",\n \"my_model = MyModel(num_neurons=3, training=False)\\n\",\n \"\\n\",\n \"x = jax.random.uniform(key=main_key, shape=(3, 4, 4))\\n\",\n \"\\n\",\n \"# Initialize with `flax.linen.init()`.\\n\",\n \"# The `params_key` is equivalent to a dictionary of PRNGs.\\n\",\n \"# (Here, you are providing only one PRNG key.) \\n\",\n \"variables = my_model.init(params_key, x)\\n\",\n \"\\n\",\n \"# Perform the forward pass with `flax.linen.apply()`.\\n\",\n \"y = my_model.apply(variables, x, rngs={'dropout': dropout_key})\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Real-life examples:\\n\",\n \"\\n\",\n \"* Applying word dropout to a batch of input IDs (in a [text classification](https://github.com/google/flax/blob/main/examples/sst2/models.py) context).\\n\",\n \"* Defining a prediction token in a decoder of a [sequence-to-sequence model](https://github.com/google/flax/blob/main/examples/seq2seq/models.py).\"\n ]\n }\n ],\n \"metadata\": {\n \"accelerator\": \"GPU\",\n \"jupytext\": {\n \"formats\": \"ipynb,md:myst\"\n },\n \"language_info\": {\n \"codemirror_mode\": {\n \"name\": \"ipython\",\n \"version\": 3\n },\n \"file_extension\": \".py\",\n \"mimetype\": \"text/x-python\",\n \"name\": \"python\",\n \"nbconvert_exporter\": \"python\",\n \"pygments_lexer\": \"ipython3\",\n \"version\": \"3.8.10\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 0\n}\n" }, { "path": "docs/guides/flax_sharp_bits.md", "content": "---\njupytext:\n formats: ipynb,md:myst\n text_representation:\n extension: .md\n format_name: myst\n format_version: 0.13\n jupytext_version: 1.13.8\n---\n\n# 🔪 Flax - The Sharp Bits 🔪\n\nFlax exposes the full power of JAX. And just like when using JAX, there are certain _[\"sharp bits\"](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html)_ you may experience when working with Flax. This evolving document is designed to assist you with them.\n\nFirst, install and/or update Flax:\n\n```{code-cell} ipython3\n:tags: [skip-execution]\n\n! pip install -qq flax\n```\n\n## 🔪 `flax.linen.Dropout` layer and randomness\n\n### TL;DR\n\nWhen working on a model with dropout (subclassed from [Flax `Module`](https://flax.readthedocs.io/en/latest/guides/flax_fundamentals/flax_basics.html#module-basics)), add the `'dropout'` PRNGkey only during the forward pass.\n\n1. Start with [`jax.random.split()`](https://jax.readthedocs.io/en/latest/_autosummary/jax.random.split.html#jax-random-split) to explicitly create PRNG keys for `'params'` and `'dropout'`.\n2. Add the [`flax.linen.Dropout`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/layers.html#flax.linen.Dropout) layer(s) to your model (subclassed from Flax [`Module`](https://flax.readthedocs.io/en/latest/guides/flax_fundamentals/flax_basics.html#module-basics)).\n3. When initializing the model ([`flax.linen.init()`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/init_apply.html)), there's no need to pass in an extra `'dropout'` PRNG key—just the `'params'` key like in a \"simpler\" model.\n4. During the forward pass with [`flax.linen.apply()`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/init_apply.html), pass in `rngs={'dropout': dropout_key}`.\n\nCheck out a full example below.\n\n### Why this works\n\n- Internally, `flax.linen.Dropout` makes use of [`flax.linen.Module.make_rng`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.make_rng) to create a key for dropout (check out the [source code](https://github.com/google/flax/blob/5714e57a0dc8146eb58a7a06ed768ed3a17672f9/flax/linen/stochastic.py#L72)).\n- Every time `make_rng` is called (in this case, it's done implicitly in `Dropout`), you get a new PRNG key split from the main/root PRNG key.\n- `make_rng` still _guarantees full reproducibility_.\n\n### Background \n\nThe [dropout](https://jmlr.org/papers/volume15/srivastava14a/srivastava14a.pdf) stochastic regularization technique randomly removes hidden and visible units in a network. Dropout is a random operation, requiring a PRNG state, and Flax (like JAX) uses [Threefry](https://github.com/jax-ml/jax/blob/main/docs/jep/263-prng.md) PRNG that is splittable. \n\n> Note: Recall that JAX has an explicit way of giving you PRNG keys: you can fork the main PRNG state (such as `key = jax.random.key(seed=0)`) into multiple new PRNG keys with `key, subkey = jax.random.split(key)`. Refresh your memory in [🔪 JAX - The Sharp Bits 🔪 Randomness and PRNG keys](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#random-numbers).\n\nFlax provides an _implicit_ way of handling PRNG key streams via [Flax `Module`](https://flax.readthedocs.io/en/latest/guides/flax_fundamentals/flax_basics.html#module-basics)'s [`flax.linen.Module.make_rng`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.make_rng) helper function. It allows the code in Flax `Module`s (or its sub-`Module`s) to \"pull PRNG keys\". `make_rng` guarantees to provide a unique key each time you call it. See the [RNG guide](https://flax.readthedocs.io/en/latest/guides/flax_fundamentals/rng_guide.html) for more details.\n\n> Note: Recall that [`flax.linen.Module`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html) is the base class for all neural network modules. All layers and models are subclassed from it.\n\n### Example\n\nRemember that each of the Flax PRNG streams has a name. The example below uses the `'params'` stream for initializing parameters, as well as the `'dropout'` stream. The PRNG key provided to [`flax.linen.init()`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/init_apply.html) is the one that seeds the `'params'` PRNG key stream. To draw PRNG keys during the forward pass (with dropout), provide a PRNG key to seed that stream (`'dropout'`) when you call `Module.apply()`.\n\n```{code-cell} ipython3\n# Setup.\nimport jax\nimport jax.numpy as jnp\nimport flax.linen as nn\n```\n\n```{code-cell} ipython3\n# Randomness.\nseed = 0\nroot_key = jax.random.key(seed=seed)\nmain_key, params_key, dropout_key = jax.random.split(key=root_key, num=3)\n\n# A simple network.\nclass MyModel(nn.Module):\n num_neurons: int\n training: bool\n @nn.compact\n def __call__(self, x):\n x = nn.Dense(self.num_neurons)(x)\n # Set the dropout layer with a rate of 50% .\n # When the `deterministic` flag is `True`, dropout is turned off.\n x = nn.Dropout(rate=0.5, deterministic=not self.training)(x)\n return x\n\n# Instantiate `MyModel` (you don't need to set `training=True` to\n# avoid performing the forward pass computation).\nmy_model = MyModel(num_neurons=3, training=False)\n\nx = jax.random.uniform(key=main_key, shape=(3, 4, 4))\n\n# Initialize with `flax.linen.init()`.\n# The `params_key` is equivalent to a dictionary of PRNGs.\n# (Here, you are providing only one PRNG key.) \nvariables = my_model.init(params_key, x)\n\n# Perform the forward pass with `flax.linen.apply()`.\ny = my_model.apply(variables, x, rngs={'dropout': dropout_key})\n```\n\nReal-life examples:\n\n* Applying word dropout to a batch of input IDs (in a [text classification](https://github.com/google/flax/blob/main/examples/sst2/models.py) context).\n* Defining a prediction token in a decoder of a [sequence-to-sequence model](https://github.com/google/flax/blob/main/examples/seq2seq/models.py).\n" }, { "path": "docs/guides/index.rst", "content": "Guides\n======\n\n.. toctree::\n :maxdepth: 2\n\n flax_fundamentals/index\n data_preprocessing/index\n training_techniques/index\n parallel_training/index\n model_inspection/index\n converting_and_upgrading/index\n quantization/index\n The Sharp Bits \n" }, { "path": "docs/guides/model_inspection/extracting_intermediates.rst", "content": "Extracting intermediate values\n==============================\n\nThis guide will show you how to extract intermediate values from a module.\nLet's start with this simple CNN that uses :code:`nn.compact`.\n\n.. testsetup:: default, sow\n\n import flax\n import flax.linen as nn\n import jax\n import jax.numpy as jnp\n from flax.core import FrozenDict\n from typing import Sequence\n\n batch = jnp.ones((4, 32, 32, 3))\n\n.. testcode::\n\n from flax import linen as nn\n import jax\n import jax.numpy as jnp\n from typing import Sequence\n\n class CNN(nn.Module):\n @nn.compact\n def __call__(self, x):\n x = nn.Conv(features=32, kernel_size=(3, 3))(x)\n x = nn.relu(x)\n x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))\n x = nn.Conv(features=64, kernel_size=(3, 3))(x)\n x = nn.relu(x)\n x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))\n x = x.reshape((x.shape[0], -1)) # flatten\n x = nn.Dense(features=256)(x)\n x = nn.relu(x)\n x = nn.Dense(features=10)(x)\n x = nn.log_softmax(x)\n return x\n\n\nBecause this module uses ``nn.compact``, we don't have direct access to\nintermediate values. There are a few ways to expose them:\n\n\nStore intermediate values in a new variable collection\n------------------------------------------------------\n\nThe CNN can be augmented with calls to ``sow`` to store intermediates as following:\n\n\n.. codediff::\n :title: Default CNN, CNN using sow API\n :groups: default, sow\n\n class CNN(nn.Module):\n @nn.compact\n def __call__(self, x):\n x = nn.Conv(features=32, kernel_size=(3, 3))(x)\n x = nn.relu(x)\n x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))\n x = nn.Conv(features=64, kernel_size=(3, 3))(x)\n x = nn.relu(x)\n x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))\n x = x.reshape((x.shape[0], -1)) # flatten\n\n x = nn.Dense(features=256)(x)\n x = nn.relu(x)\n x = nn.Dense(features=10)(x)\n x = nn.log_softmax(x)\n return x\n ---\n class SowCNN(nn.Module):\n @nn.compact\n def __call__(self, x):\n x = nn.Conv(features=32, kernel_size=(3, 3))(x)\n x = nn.relu(x)\n x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))\n x = nn.Conv(features=64, kernel_size=(3, 3))(x)\n x = nn.relu(x)\n x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))\n x = x.reshape((x.shape[0], -1)) # flatten\n self.sow('intermediates', 'features', x) #!\n x = nn.Dense(features=256)(x)\n x = nn.relu(x)\n x = nn.Dense(features=10)(x)\n x = nn.log_softmax(x)\n return x\n\n``sow`` acts as a no-op when the variable collection is not mutable.\nTherefore, it works perfectly for debugging and optional tracking of intermediates.\nThe 'intermediates' collection is also used by the ``capture_intermediates`` API (see the :ref:`Use ``capture_intermediates``` section).\n\nNote that, by default ``sow`` appends values every time it is called:\n\n* This is necessary because once instantiated, a module could be called multiple\n times in its parent module, and we want to catch all the sowed values.\n* Therefore you want to make sure that you **do not** feed intermediate values back\n into ``variables``. Otherwise every call will increase the length of that tuple\n and trigger a recompile.\n* To override the default append behavior, specify ``init_fn`` and ``reduce_fn``\n - see :meth:`Module.sow() `.\n\n.. testcode:: sow\n\n class SowCNN2(nn.Module):\n @nn.compact\n def __call__(self, x):\n mod = SowCNN(name='SowCNN')\n return mod(x) + mod(x) # Calling same module instance twice.\n\n @jax.jit\n def init(key, x):\n variables = SowCNN2().init(key, x)\n # By default the 'intermediates' collection is not mutable during init.\n # So variables will only contain 'params' here.\n return variables\n\n @jax.jit\n def predict(variables, x):\n # If mutable='intermediates' is not specified, then .sow() acts as a noop.\n output, mod_vars = SowCNN2().apply(variables, x, mutable='intermediates')\n features = mod_vars['intermediates']['SowCNN']['features']\n return output, features\n\n batch = jnp.ones((1,28,28,1))\n variables = init(jax.random.key(0), batch)\n preds, feats = predict(variables, batch)\n\n assert len(feats) == 2 # Tuple with two values since module was called twice.\n\nRefactor module into submodules\n-------------------------------\n\nThis is a useful pattern for cases where it's clear in what particular\nway you want to split your submodules. Any submodule you expose in ``setup`` can\nbe used directly. In the limit, you can define all submodules in ``setup`` and\navoid using ``nn.compact`` altogether.\n\n.. testcode::\n\n class RefactoredCNN(nn.Module):\n def setup(self):\n self.features = Features()\n self.classifier = Classifier()\n\n def __call__(self, x):\n x = self.features(x)\n x = self.classifier(x)\n return x\n\n class Features(nn.Module):\n @nn.compact\n def __call__(self, x):\n x = nn.Conv(features=32, kernel_size=(3, 3))(x)\n x = nn.relu(x)\n x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))\n x = nn.Conv(features=64, kernel_size=(3, 3))(x)\n x = nn.relu(x)\n x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))\n x = x.reshape((x.shape[0], -1)) # flatten\n return x\n\n class Classifier(nn.Module):\n @nn.compact\n def __call__(self, x):\n x = nn.Dense(features=256)(x)\n x = nn.relu(x)\n x = nn.Dense(features=10)(x)\n x = nn.log_softmax(x)\n return x\n\n @jax.jit\n def init(key, x):\n variables = RefactoredCNN().init(key, x)\n return variables['params']\n\n @jax.jit\n def features(params, x):\n return RefactoredCNN().apply({\"params\": params}, x,\n method=lambda module, x: module.features(x))\n\n params = init(jax.random.key(0), batch)\n\n features(params, batch)\n\n\nUse ``capture_intermediates``\n-----------------------------\n\nLinen supports the capture of intermediate return values from submodules automatically without any code changes.\nThis pattern should be considered the \"sledge hammer\" approach to capturing intermediates.\nAs a debugging and inspection tool it is very useful, but using the other patterns described in this guide\nwill give you more fine-grained control over what intermediates you want to extract.\n\nIn the following code example we check if any intermediate activations are non-finite (NaN or infinite):\n\n.. testcode::\n\n @jax.jit\n def init(key, x):\n variables = CNN().init(key, x)\n return variables\n\n @jax.jit\n def predict(variables, x):\n y, state = CNN().apply(variables, x, capture_intermediates=True, mutable=[\"intermediates\"])\n intermediates = state['intermediates']\n fin = jax.tree_util.tree_map(lambda xs: jnp.all(jnp.isfinite(xs)), intermediates)\n return y, fin\n\n variables = init(jax.random.key(0), batch)\n y, is_finite = predict(variables, batch)\n all_finite = all(jax.tree_util.tree_leaves(is_finite))\n assert all_finite, \"non-finite intermediate detected!\"\n\nBy default only the intermediates of ``__call__`` methods are collected.\nAlternatively, you can pass a custom filter function based on the ``Module`` instance and the method name.\n\n.. testcode::\n\n filter_Dense = lambda mdl, method_name: isinstance(mdl, nn.Dense)\n filter_encodings = lambda mdl, method_name: method_name == \"encode\"\n\n y, state = CNN().apply(variables, batch, capture_intermediates=filter_Dense, mutable=[\"intermediates\"])\n dense_intermediates = state['intermediates']\n\nNote that ``capture_intermediates`` will only apply to layers. You can use ``self.sow`` to manually store\nnon-layer intermediates, but the filter function won't be applied to it.\n\n.. codediff::\n :title: Capturing all layer intermediates, Using filter function and ``self.sow()``\n :groups: default, sow\n\n class Model(nn.Module):\n @nn.compact\n def __call__(self, x):\n a = nn.Dense(4)(x) # Dense_0\n b = nn.Dense(4)(x) # Dense_1\n c = a + b # not a Flax layer, so won't be stored as an intermediate\n d = nn.Dense(4)(c) # Dense_2\n return d\n\n @jax.jit\n def init(key, x):\n variables = Model().init(key, x)\n return variables['params']\n\n @jax.jit\n def predict(params, x):\n return Model().apply({\"params\": params}, x, capture_intermediates=True)\n\n batch = jax.random.uniform(jax.random.key(1), (1,3))\n params = init(jax.random.key(0), batch)\n preds, feats = predict(params, batch)\n feats # intermediate c in Model was not stored because it's not a Flax layer\n ---\n class Model(nn.Module):\n @nn.compact\n def __call__(self, x):\n a = nn.Dense(4)(x) # Dense_0\n b = nn.Dense(4)(x) # Dense_1\n c = a + b\n self.sow('intermediates', 'c', c) # store intermediate c #!\n d = nn.Dense(4)(c) # Dense_2\n return d\n\n @jax.jit\n def init(key, x):\n variables = Model().init(key, x)\n return variables['params']\n\n @jax.jit\n def predict(params, x):\n # filter specifically for only the Dense_0 and Dense_2 layer #!\n filter_fn = lambda mdl, method_name: isinstance(mdl.name, str) and (mdl.name in {'Dense_0', 'Dense_2'}) #!\n return Model().apply({\"params\": params}, x, capture_intermediates=filter_fn) #!\n\n batch = jax.random.uniform(jax.random.key(1), (1,3))\n params = init(jax.random.key(0), batch)\n preds, feats = predict(params, batch)\n feats # intermediate c in Model is stored and isn't filtered out by the filter function #!\n\nTo separate the intermediates extracted from ``self.sow`` from the intermediates extracted from ``capture_intermediates``,\nwe can either define a separate collection like ``self.sow('sow_intermediates', 'c', c)``, or manually filter out\nthe intermediates after calling ``.apply()``. For example:\n\n.. testcode:: sow\n\n flattened_dict = flax.traverse_util.flatten_dict(feats['intermediates'], sep='/')\n flattened_dict['c']\n\nIn terms of efficiency, as long as everything is jitted, then any intermediates you don't end up using\nshould be optimized away by XLA.\n\nUse ``Sequential``\n---------------------\n\nYou could also define ``CNN`` using a simple implementation of a ``Sequential`` combinator (this is quite common in more stateful approaches). This may be useful\nfor very simple models and gives you arbitrary model\nsurgery. But it can be very limiting -- if you even want to add one conditional, you are\nforced to refactor away from ``Sequential`` and structure\nyour model more explicitly.\n\n.. testcode::\n\n class Sequential(nn.Module):\n layers: Sequence[nn.Module]\n\n def __call__(self, x):\n for layer in self.layers:\n x = layer(x)\n return x\n\n def SeqCNN():\n return Sequential([\n nn.Conv(features=32, kernel_size=(3, 3)),\n nn.relu,\n lambda x: nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2)),\n nn.Conv(features=64, kernel_size=(3, 3)),\n nn.relu,\n lambda x: nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2)),\n lambda x: x.reshape((x.shape[0], -1)), # flatten\n nn.Dense(features=256),\n nn.relu,\n nn.Dense(features=10),\n nn.log_softmax,\n ])\n\n @jax.jit\n def init(key, x):\n variables = SeqCNN().init(key, x)\n return variables['params']\n\n @jax.jit\n def features(params, x):\n return Sequential(SeqCNN().layers[0:7]).apply({\"params\": params}, x)\n\n batch = jnp.ones((1,28,28,1))\n params = init(jax.random.key(0), batch)\n features(params, batch)\n\nExtracting gradients of intermediate values\n===========================================\nFor debugging purposes, it can be useful to extract the gradients of intermediate values.\nThis can be done by using the :meth:`Module.perturb() ` method over the desired values.\n\n.. testcode::\n\n class Model(nn.Module):\n @nn.compact\n def __call__(self, x):\n x = nn.relu(nn.Dense(8)(x))\n x = self.perturb('hidden', x)\n x = nn.Dense(2)(x)\n x = self.perturb('logits', x)\n return x\n\n``perturb`` adds a variable to a ``perturbations`` collection by default,\nit behaves like an identity function and the gradient of the perturbation\nmatches the gradient of the input. To get the perturbations just initialize\nthe model:\n\n.. testcode::\n\n x = jnp.empty((1, 4)) # random data\n y = jnp.empty((1, 2)) # random data\n\n model = Model()\n variables = model.init(jax.random.key(1), x)\n params, perturbations = variables['params'], variables['perturbations']\n\nFinally compute the gradients of the loss with respect to the perturbations,\nthese will match the gradients of the intermediates:\n\n.. testcode::\n\n def loss_fn(params, perturbations, x, y):\n y_pred = model.apply({'params': params, 'perturbations': perturbations}, x)\n return jnp.mean((y_pred - y) ** 2)\n\n intermediate_grads = jax.grad(loss_fn, argnums=1)(params, perturbations, x, y)" }, { "path": "docs/guides/model_inspection/index.rst", "content": "Model inspection\n================\n\n.. toctree::\n :maxdepth: 1\n\n model_surgery\n extracting_intermediates\n" }, { "path": "docs/guides/model_inspection/model_surgery.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"id\": \"120e57f5\",\n \"metadata\": {},\n \"source\": [\n \"Model surgery\\n\",\n \"==============================\\n\",\n \"\\n\",\n \"Usually, Flax modules and optimizers track and update the params for you. But there may be some time when you want to do some model surgery and tweak the param tensors yourself. This guide shows you how to do the trick.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"9c3bfb0e\",\n \"metadata\": {},\n \"source\": [\n \"## Setup\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"413f8b2d\",\n \"metadata\": {\n \"tags\": [\n \"skip-execution\"\n ]\n },\n \"outputs\": [],\n \"source\": [\n \"!pip install --upgrade -q pip jax jaxlib flax\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"5b002c8d\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"import functools\\n\",\n \"\\n\",\n \"import jax\\n\",\n \"import jax.numpy as jnp\\n\",\n \"from flax import traverse_util\\n\",\n \"from flax import linen as nn\\n\",\n \"from flax.core import freeze\\n\",\n \"import jax\\n\",\n \"import optax\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"1060b519\",\n \"metadata\": {},\n \"source\": [\n \"Surgery with Flax Modules\\n\",\n \"--------------------------------\\n\",\n \"\\n\",\n \"Let's create a small convolutional neural network model for our demo.\\n\",\n \"\\n\",\n \"As usual, you can run `CNN.init(...)['params']` to get the `params` to pass and modify it in every step of your training.\\n\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"755ae323\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"class CNN(nn.Module):\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self, x):\\n\",\n \" x = nn.Conv(features=32, kernel_size=(3, 3))(x)\\n\",\n \" x = nn.relu(x)\\n\",\n \" x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))\\n\",\n \" x = nn.Conv(features=64, kernel_size=(3, 3))(x)\\n\",\n \" x = nn.relu(x)\\n\",\n \" x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))\\n\",\n \" x = x.reshape((x.shape[0], -1))\\n\",\n \" x = nn.Dense(features=256)(x)\\n\",\n \" x = nn.relu(x)\\n\",\n \" x = nn.Dense(features=10)(x)\\n\",\n \" x = nn.log_softmax(x)\\n\",\n \" return x\\n\",\n \"\\n\",\n \"def get_initial_params(key):\\n\",\n \" init_shape = jnp.ones((1, 28, 28, 1), jnp.float32)\\n\",\n \" initial_params = CNN().init(key, init_shape)['params']\\n\",\n \" return initial_params\\n\",\n \"\\n\",\n \"key = jax.random.key(0)\\n\",\n \"params = get_initial_params(key)\\n\",\n \"\\n\",\n \"jax.tree_util.tree_map(jnp.shape, params)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"170273f8\",\n \"metadata\": {},\n \"source\": [\n \"Note that what returned as `params` is a `FrozenDict`, which contains a few JAX arrays as kernel and bias. \\n\",\n \"\\n\",\n \"A `FrozenDict` is nothing more than a read-only dict, and Flax made it read-only because of the functional nature of JAX: JAX arrays are immutable, and the new `params` need to replace the old `params`. Making the dict read-only ensures that no in-place mutation of the dict can happen accidentally during the training and updating.\\n\",\n \"\\n\",\n \"One way to actually modify the params outside of a Flax module is to explicitly flatten it and creates a mutable dict. Note that you can use a separator `sep` to join all nested keys. If no `sep` is given, the key will be a tuple of all nested keys.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"c7ec7741\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# Get a flattened key-value list.\\n\",\n \"flat_params = traverse_util.flatten_dict(params, sep='/')\\n\",\n \"\\n\",\n \"jax.tree_util.tree_map(jnp.shape, flat_params)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"2adda656\",\n \"metadata\": {},\n \"source\": [\n \"Now you can do whatever you want with the params. When you are done, unflatten it back and use it in future training.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"bb975feb\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# Somehow modify a layer\\n\",\n \"dense_kernel = flat_params['Dense_1/kernel']\\n\",\n \"flat_params['Dense_1/kernel'] = dense_kernel / jnp.linalg.norm(dense_kernel)\\n\",\n \"\\n\",\n \"# Unflatten.\\n\",\n \"unflat_params = traverse_util.unflatten_dict(flat_params, sep='/')\\n\",\n \"# Refreeze.\\n\",\n \"unflat_params = freeze(unflat_params)\\n\",\n \"jax.tree_util.tree_map(jnp.shape, unflat_params)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"f3462cd8\",\n \"metadata\": {},\n \"source\": [\n \"Surgery with Optimizers\\n\",\n \"--------------------------------\\n\",\n \"\\n\",\n \"When using `Optax` as an optimizer, the ``opt_state`` is actually a nested tuple\\n\",\n \"of the states of individual gradient transformations that compose the optimizer.\\n\",\n \"These states contain pytrees that mirror the parameter tree, and can be modified\\n\",\n \"the same way: flattening, modifying, unflattening, and then recreating a new\\n\",\n \"optimizer state that mirrors the original state.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"3cbecb63\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"tx = optax.adam(1.0)\\n\",\n \"opt_state = tx.init(params)\\n\",\n \"\\n\",\n \"# The optimizer state is a tuple of gradient transformation states.\\n\",\n \"jax.tree_util.tree_map(jnp.shape, opt_state)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"18f1cebb\",\n \"metadata\": {},\n \"source\": [\n \"The pytrees inside the optimizer state follow the same structure as the\\n\",\n \"parameters and can be flattened / modified exactly the same way.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"13b5e25f\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"flat_mu = traverse_util.flatten_dict(opt_state[0].mu, sep='/')\\n\",\n \"flat_nu = traverse_util.flatten_dict(opt_state[0].nu, sep='/')\\n\",\n \"\\n\",\n \"jax.tree_util.tree_map(jnp.shape, flat_mu)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"e5c4479e\",\n \"metadata\": {},\n \"source\": [\n \"After modification, re-create optimizer state. Use this for future training.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"9dcac8cd\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"opt_state = (\\n\",\n \" opt_state[0]._replace(\\n\",\n \" mu=traverse_util.unflatten_dict(flat_mu, sep='/'),\\n\",\n \" nu=traverse_util.unflatten_dict(flat_nu, sep='/'),\\n\",\n \" ),\\n\",\n \") + opt_state[1:]\\n\",\n \"jax.tree_util.tree_map(jnp.shape, opt_state)\"\n ]\n }\n ],\n \"metadata\": {\n \"jupytext\": {\n \"formats\": \"md,ipynb\",\n \"main_language\": \"python\"\n },\n \"language_info\": {\n \"codemirror_mode\": {\n \"name\": \"ipython\",\n \"version\": 3\n },\n \"file_extension\": \".py\",\n \"mimetype\": \"text/x-python\",\n \"name\": \"python\",\n \"nbconvert_exporter\": \"python\",\n \"pygments_lexer\": \"ipython3\",\n \"version\": \"3.9.15\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 5\n}\n" }, { "path": "docs/guides/model_inspection/model_surgery.md", "content": "---\njupyter:\n jupytext:\n formats: md,ipynb\n main_language: python\n text_representation:\n extension: .md\n format_name: markdown\n format_version: '1.3'\n jupytext_version: 1.13.8\n---\n\nModel surgery\n==============================\n\nUsually, Flax modules and optimizers track and update the params for you. But there may be some time when you want to do some model surgery and tweak the param tensors yourself. This guide shows you how to do the trick.\n\n\n## Setup\n\n```python tags=[\"skip-execution\"]\n!pip install --upgrade -q pip jax jaxlib flax\n```\n\n```python\nimport functools\n\nimport jax\nimport jax.numpy as jnp\nfrom flax import traverse_util\nfrom flax import linen as nn\nfrom flax.core import freeze\nimport jax\nimport optax\n```\n\nSurgery with Flax Modules\n--------------------------------\n\nLet's create a small convolutional neural network model for our demo.\n\nAs usual, you can run `CNN.init(...)['params']` to get the `params` to pass and modify it in every step of your training.\n\n\n```python\nclass CNN(nn.Module):\n @nn.compact\n def __call__(self, x):\n x = nn.Conv(features=32, kernel_size=(3, 3))(x)\n x = nn.relu(x)\n x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))\n x = nn.Conv(features=64, kernel_size=(3, 3))(x)\n x = nn.relu(x)\n x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))\n x = x.reshape((x.shape[0], -1))\n x = nn.Dense(features=256)(x)\n x = nn.relu(x)\n x = nn.Dense(features=10)(x)\n x = nn.log_softmax(x)\n return x\n\ndef get_initial_params(key):\n init_shape = jnp.ones((1, 28, 28, 1), jnp.float32)\n initial_params = CNN().init(key, init_shape)['params']\n return initial_params\n\nkey = jax.random.key(0)\nparams = get_initial_params(key)\n\njax.tree_util.tree_map(jnp.shape, params)\n```\n\nNote that what returned as `params` is a `FrozenDict`, which contains a few JAX arrays as kernel and bias. \n\nA `FrozenDict` is nothing more than a read-only dict, and Flax made it read-only because of the functional nature of JAX: JAX arrays are immutable, and the new `params` need to replace the old `params`. Making the dict read-only ensures that no in-place mutation of the dict can happen accidentally during the training and updating.\n\nOne way to actually modify the params outside of a Flax module is to explicitly flatten it and creates a mutable dict. Note that you can use a separator `sep` to join all nested keys. If no `sep` is given, the key will be a tuple of all nested keys.\n\n```python\n# Get a flattened key-value list.\nflat_params = traverse_util.flatten_dict(params, sep='/')\n\njax.tree_util.tree_map(jnp.shape, flat_params)\n```\n\nNow you can do whatever you want with the params. When you are done, unflatten it back and use it in future training.\n\n```python\n# Somehow modify a layer\ndense_kernel = flat_params['Dense_1/kernel']\nflat_params['Dense_1/kernel'] = dense_kernel / jnp.linalg.norm(dense_kernel)\n\n# Unflatten.\nunflat_params = traverse_util.unflatten_dict(flat_params, sep='/')\n# Refreeze.\nunflat_params = freeze(unflat_params)\njax.tree_util.tree_map(jnp.shape, unflat_params)\n```\n\nSurgery with Optimizers\n--------------------------------\n\nWhen using `Optax` as an optimizer, the ``opt_state`` is actually a nested tuple\nof the states of individual gradient transformations that compose the optimizer.\nThese states contain pytrees that mirror the parameter tree, and can be modified\nthe same way: flattening, modifying, unflattening, and then recreating a new\noptimizer state that mirrors the original state.\n\n```python\ntx = optax.adam(1.0)\nopt_state = tx.init(params)\n\n# The optimizer state is a tuple of gradient transformation states.\njax.tree_util.tree_map(jnp.shape, opt_state)\n```\n\nThe pytrees inside the optimizer state follow the same structure as the\nparameters and can be flattened / modified exactly the same way.\n\n```python\nflat_mu = traverse_util.flatten_dict(opt_state[0].mu, sep='/')\nflat_nu = traverse_util.flatten_dict(opt_state[0].nu, sep='/')\n\njax.tree_util.tree_map(jnp.shape, flat_mu)\n```\n\nAfter modification, re-create optimizer state. Use this for future training.\n\n```python\nopt_state = (\n opt_state[0]._replace(\n mu=traverse_util.unflatten_dict(flat_mu, sep='/'),\n nu=traverse_util.unflatten_dict(flat_nu, sep='/'),\n ),\n) + opt_state[1:]\njax.tree_util.tree_map(jnp.shape, opt_state)\n```\n" }, { "path": "docs/guides/parallel_training/ensembling.rst", "content": "Ensembling on multiple devices\n==============================\n\nWe show how to train an ensemble of CNNs on the MNIST dataset, where the size of\nthe ensemble is equal to the number of available devices. In short, this change\nbe described as:\n\n* make a number of functions parallel using |jax.pmap()|_,\n* split the random seed to obtain different parameter initialization,\n* replicate the inputs and unreplicate the outputs where necessary,\n* average probabilities across devices to compute the predictions.\n\nIn this HOWTO we omit some of the code such as imports, the CNN module, and\nmetrics computation, but they can be found in the `MNIST example`_.\n\n.. testsetup:: Single-model, Ensemble\n\n import functools\n from flax import jax_utils\n\n # Copied from examples/mnist/train.py\n from absl import logging\n from flax import linen as nn\n from flax.training import train_state\n import jax\n import jax.numpy as jnp\n import numpy as np\n import optax\n\n class CNN(nn.Module):\n \"\"\"A simple CNN model.\"\"\"\n\n @nn.compact\n def __call__(self, x):\n x = nn.Conv(features=32, kernel_size=(3, 3))(x)\n x = nn.relu(x)\n x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))\n x = nn.Conv(features=64, kernel_size=(3, 3))(x)\n x = nn.relu(x)\n x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))\n x = x.reshape((x.shape[0], -1)) # flatten\n x = nn.Dense(features=256)(x)\n x = nn.relu(x)\n x = nn.Dense(features=10)(x)\n return x\n\n # Fake data for faster execution.\n def get_datasets():\n train_ds = test_ds = {\n 'image': jnp.zeros([64, 28, 28, 1]),\n 'label': jnp.zeros([64], jnp.int32),\n }\n return train_ds, test_ds\n\n # Modified from examples/mnist/configs.default.py\n learning_rate = 0.1\n momentum = 0.9\n batch_size = 32\n num_epochs = 1\n\n\nParallel functions\n------------------\n\nWe start by creating a parallel version of ``create_train_state()``, which\nretrieves the initial parameters of the models. We do this using |jax.pmap()|_.\nThe effect of \"pmapping\" a function is that it will compile the function with\nXLA (similar to |jax.jit()|_), but execute it in parallel on XLA devices (e.g.,\nGPUs/TPUs).\n\n.. codediff::\n :title: Single-model, Ensemble\n :sync:\n\n #!\n def create_train_state(rng, learning_rate, momentum):\n cnn = CNN()\n params = cnn.init(rng, jnp.ones([1, 28, 28, 1]))['params']\n tx = optax.sgd(learning_rate, momentum)\n return train_state.TrainState.create(\n apply_fn=cnn.apply, params=params, tx=tx)\n ---\n @functools.partial(jax.pmap, static_broadcasted_argnums=(1, 2)) #!\n def create_train_state(rng, learning_rate, momentum):\n cnn = CNN()\n params = cnn.init(rng, jnp.ones([1, 28, 28, 1]))['params']\n tx = optax.sgd(learning_rate, momentum)\n return train_state.TrainState.create(\n apply_fn=cnn.apply, params=params, tx=tx)\n\nNote that for the single-model code above, we use |jax.jit()|_ to lazily\ninitialize the model (see `Module.init`_'s documentation for more details).\nFor the ensembling case, |jax.pmap()|_ will map over the first axis of the\nprovided argument ``rng`` by default, so we should make sure that we provide\na different value for each device when we call this function later on.\n\nNote also how we specify that ``learning_rate`` and ``momentum`` are static\narguments, which means the concrete values of these arguments will be used,\nrather than abstract shapes. This is necessary because the provided arguments\nwill be scalar values. For more details see `JIT mechanics: tracing and static\nvariables`_.\n\nNext we simply do the same for the functions ``apply_model()`` and\n``update_model()``. To compute the predictions from the ensemble, we take the\naverage of the individual probabilities. We use |jax.lax.pmean()|_ to compute\nthe average *across devices*. This also requires us to specify the\n``axis_name`` to both |jax.pmap()|_ and |jax.lax.pmean()|_.\n\n.. codediff::\n :title: Single-model, Ensemble\n :sync:\n\n @jax.jit #!\n def apply_model(state, images, labels):\n def loss_fn(params):\n logits = CNN().apply({'params': params}, images)\n one_hot = jax.nn.one_hot(labels, 10)\n loss = optax.softmax_cross_entropy(logits=logits, labels=one_hot).mean()\n return loss, logits\n\n grad_fn = jax.value_and_grad(loss_fn, has_aux=True)\n (loss, logits), grads = grad_fn(state.params)\n #!\n accuracy = jnp.mean(jnp.argmax(logits, -1) == labels) #!\n return grads, loss, accuracy\n\n @jax.jit #!\n def update_model(state, grads):\n return state.apply_gradients(grads=grads)\n ---\n @functools.partial(jax.pmap, axis_name='ensemble') #!\n def apply_model(state, images, labels):\n def loss_fn(params):\n logits = CNN().apply({'params': params}, images)\n one_hot = jax.nn.one_hot(labels, 10)\n loss = optax.softmax_cross_entropy(logits=logits, labels=one_hot).mean()\n return loss, logits\n\n grad_fn = jax.value_and_grad(loss_fn, has_aux=True)\n (loss, logits), grads = grad_fn(state.params)\n probs = jax.lax.pmean(jax.nn.softmax(logits), axis_name='ensemble') #!\n accuracy = jnp.mean(jnp.argmax(probs, -1) == labels) #!\n return grads, loss, accuracy\n\n @jax.pmap #!\n def update_model(state, grads):\n return state.apply_gradients(grads=grads)\n\nTraining the Ensemble\n---------------------\n\nNext we transform the ``train_epoch()`` function. When calling the pmapped\nfunctions from above, we mainly need to take care of duplicating the arguments\nfor all devices where necessary, and de-duplicating the return values.\n\n.. codediff::\n :title: Single-model, Ensemble\n :sync:\n\n def train_epoch(state, train_ds, batch_size, rng):\n train_ds_size = len(train_ds['image'])\n steps_per_epoch = train_ds_size // batch_size\n\n perms = jax.random.permutation(rng, len(train_ds['image']))\n perms = perms[:steps_per_epoch * batch_size]\n perms = perms.reshape((steps_per_epoch, batch_size))\n\n epoch_loss = []\n epoch_accuracy = []\n\n for perm in perms:\n batch_images = train_ds['image'][perm, ...] #!\n batch_labels = train_ds['label'][perm, ...] #!\n grads, loss, accuracy = apply_model(state, batch_images, batch_labels)\n state = update_model(state, grads)\n epoch_loss.append(loss) #!\n epoch_accuracy.append(accuracy) #!\n train_loss = np.mean(epoch_loss)\n train_accuracy = np.mean(epoch_accuracy)\n return state, train_loss, train_accuracy\n ---\n def train_epoch(state, train_ds, batch_size, rng):\n train_ds_size = len(train_ds['image'])\n steps_per_epoch = train_ds_size // batch_size\n\n perms = jax.random.permutation(rng, len(train_ds['image']))\n perms = perms[:steps_per_epoch * batch_size]\n perms = perms.reshape((steps_per_epoch, batch_size))\n\n epoch_loss = []\n epoch_accuracy = []\n\n for perm in perms:\n batch_images = jax_utils.replicate(train_ds['image'][perm, ...]) #!\n batch_labels = jax_utils.replicate(train_ds['label'][perm, ...]) #!\n grads, loss, accuracy = apply_model(state, batch_images, batch_labels)\n state = update_model(state, grads)\n epoch_loss.append(jax_utils.unreplicate(loss)) #!\n epoch_accuracy.append(jax_utils.unreplicate(accuracy)) #!\n train_loss = np.mean(epoch_loss)\n train_accuracy = np.mean(epoch_accuracy)\n return state, train_loss, train_accuracy\n\nAs can be seen, we do not have to make any changes to the logic around the\n``state``. This is because, as we will see below in our training code,\nthe train state is replicated already, so when we pass it to ``train_step()``,\nthings will just work fine since ``train_step()`` is pmapped. However,\nthe train dataset is not yet replicated, so we do that here. Since replicating\nthe entire train dataset is too memory intensive we do it at the batch level.\n\nWe can now rewrite the actual training logic. This consists of two simple\nchanges: making sure the RNGs are replicated when we pass them to\n``create_train_state()``, and replicating the test dataset, which is much\nsmaller than the train dataset so we can do this for the entire dataset\ndirectly.\n\n.. codediff::\n :title: Single-model, Ensemble\n :sync:\n\n train_ds, test_ds = get_datasets()\n #!\n rng = jax.random.key(0)\n\n rng, init_rng = jax.random.split(rng)\n state = create_train_state(init_rng, learning_rate, momentum) #!\n #!\n\n for epoch in range(1, num_epochs + 1):\n rng, input_rng = jax.random.split(rng)\n state, train_loss, train_accuracy = train_epoch(\n state, train_ds, batch_size, input_rng)\n\n _, test_loss, test_accuracy = apply_model( #!\n state, test_ds['image'], test_ds['label']) #!\n\n logging.info(\n 'epoch:% 3d, train_loss: %.4f, train_accuracy: %.2f, '\n 'test_loss: %.4f, test_accuracy: %.2f'\n % (epoch, train_loss, train_accuracy * 100, test_loss,\n test_accuracy * 100))\n ---\n train_ds, test_ds = get_datasets()\n test_ds = jax_utils.replicate(test_ds) #!\n rng = jax.random.key(0)\n\n rng, init_rng = jax.random.split(rng)\n state = create_train_state(jax.random.split(init_rng, jax.device_count()), #!\n learning_rate, momentum) #!\n\n for epoch in range(1, num_epochs + 1):\n rng, input_rng = jax.random.split(rng)\n state, train_loss, train_accuracy = train_epoch(\n state, train_ds, batch_size, input_rng)\n\n _, test_loss, test_accuracy = jax_utils.unreplicate( #!\n apply_model(state, test_ds['image'], test_ds['label'])) #!\n\n logging.info(\n 'epoch:% 3d, train_loss: %.4f, train_accuracy: %.2f, '\n 'test_loss: %.4f, test_accuracy: %.2f'\n % (epoch, train_loss, train_accuracy * 100, test_loss,\n test_accuracy * 100))\n\n\n.. |jax.jit()| replace:: ``jax.jit()``\n.. _jax.jit(): https://jax.readthedocs.io/en/latest/notebooks/thinking_in_jax.html#To-JIT-or-not-to-JIT\n.. |jax.pmap()| replace:: ``jax.pmap()``\n.. _jax.pmap(): https://jax.readthedocs.io/en/latest/jax.html#jax.pmap\n.. |jax.lax.pmean()| replace:: ``jax.lax.pmean()``\n.. _jax.lax.pmean(): https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.pmean.html\n.. _Module.init: https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.init\n.. _`JIT mechanics: tracing and static variables`: https://jax.readthedocs.io/en/latest/notebooks/thinking_in_jax.html#JIT-mechanics:-tracing-and-static-variables\n.. _`MNIST example`: https://github.com/google/flax/blob/main/examples/mnist/train.py\n" }, { "path": "docs/guides/parallel_training/flax_on_pjit.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"# Scale up Flax Modules on multiple devices\\n\",\n \"\\n\",\n \"This guide shows how to scale up [Flax Modules](https://flax.readthedocs.io/en/latest/developer_notes/module_lifecycle.html) on multiple devices and hosts using [`jax.jit`](https://jax.readthedocs.io/en/latest/jax-101/02-jitting.html) (formerly [`experimental.pjit`](https://jax.readthedocs.io/en/latest/jax.experimental.pjit.html#module-jax.experimental.pjit)) and [`flax.linen`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/index.html).\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Flax and `jax.jit` scaled up\\n\",\n \"\\n\",\n \"[`jax.jit`](https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html) follows the [Single Program Multi Data (SPMD)](https://jax.readthedocs.io/en/latest/glossary.html#term-SPMD) paradigm and automatically compiles your code to run it on multiple devices. You need to only specify how you want the input and output of your code to be partitioned, and the compiler will figure out how to: 1) partition everything inside; and 2) compile inter-device communications.\\n\",\n \"\\n\",\n \"Flax provides several functionalities that can help you use auto-SPMD on [Flax Modules](https://flax.readthedocs.io/en/latest/developer_notes/module_lifecycle.html), including:\\n\",\n \"\\n\",\n \"1. An interface to specify partitions of your data when defining [`flax.linen.Module`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html).\\n\",\n \"2. Utility functions to generate the sharding information that `jax.jit` requires to run.\\n\",\n \"3. An interface to customize your axis names called \\\"logical axis annotations\\\" to decouple both your Module code and partition plan to experiment with different partition layouts more easily.\\n\",\n \"\\n\",\n \"You can learn more about `jax.jit` APIs for scaling up in [JAX in multi-process environments](https://jax.readthedocs.io/en/latest/multi_process.html) and [Distributed arrays and automatic parallelization](https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html) on JAX's documentation site.\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Setup\\n\",\n \"\\n\",\n \"Import some necessary dependencies.\\n\",\n \"\\n\",\n \"**Note:** This guide uses the `--xla_force_host_platform_device_count=8` flag to emulate multiple devices in a CPU environment in a Google Colab/Jupyter Notebook. You don't need this if you are already using a multi-device TPU environment.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 2,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"import os\\n\",\n \"os.environ[\\\"XLA_FLAGS\\\"] = '--xla_force_host_platform_device_count=8'\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 3,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stderr\",\n \"output_type\": \"stream\",\n \"text\": [\n \"WARNING:absl:Tensorflow library not found, tensorflow.io.gfile operations will use native shim calls. GCS paths (i.e. 'gs://...') cannot be accessed.\\n\"\n ]\n }\n ],\n \"source\": [\n \"import functools\\n\",\n \"from typing import Optional, Callable\\n\",\n \"\\n\",\n \"import numpy as np\\n\",\n \"import jax\\n\",\n \"from jax import lax, random, numpy as jnp\\n\",\n \"\\n\",\n \"import flax\\n\",\n \"from flax import struct, traverse_util, linen as nn\\n\",\n \"from flax.core import freeze, unfreeze\\n\",\n \"from flax.training import train_state, checkpoints\\n\",\n \"\\n\",\n \"import optax # Optax for common losses and optimizers.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 4,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"We have 8 fake JAX devices now: [CpuDevice(id=0), CpuDevice(id=1), CpuDevice(id=2), CpuDevice(id=3), CpuDevice(id=4), CpuDevice(id=5), CpuDevice(id=6), CpuDevice(id=7)]\\n\"\n ]\n }\n ],\n \"source\": [\n \"print(f'We have 8 fake JAX devices now: {jax.devices()}')\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"The code below shows how to import and set up the JAX-level device API, following JAX's [Distributed arrays and automatic parallelization](https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html) guide:\\n\",\n \"\\n\",\n \"1. Start a 2x4 device `mesh` (8 devices) using JAX's `mesh_utils.create_device_mesh`. This layout is the same as the one of a [TPU v3-8](https://cloud.google.com/tpu/docs/system-architecture-tpu-vm#single_tpu_board).\\n\",\n \"\\n\",\n \"2. Annotate each axis with a name using the `axis_names` parameter in `jax.sharding.Mesh`. A typical way to annotate axis names is `axis_name=('data', 'model')`, where:\\n\",\n \" * `'data'`: the mesh dimension used for data-parallel sharding of the batch dimension of inputs and activations.\\n\",\n \" * `'model'`: the mesh dimension used for sharding parameters of the model across devices.\\n\",\n \"\\n\",\n \"3. Make a simple utility function `mesh_sharding` for generating a sharding object from the mesh and any layout.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 5,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"from jax.sharding import Mesh, PartitionSpec, NamedSharding\\n\",\n \"from jax.lax import with_sharding_constraint\\n\",\n \"from jax.experimental import mesh_utils\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 6,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"[[CpuDevice(id=0) CpuDevice(id=1) CpuDevice(id=2) CpuDevice(id=3)]\\n\",\n \" [CpuDevice(id=4) CpuDevice(id=5) CpuDevice(id=6) CpuDevice(id=7)]]\\n\",\n \"Mesh('data': 2, 'model': 4)\\n\"\n ]\n }\n ],\n \"source\": [\n \"# Create a mesh and annotate each axis with a name.\\n\",\n \"device_mesh = mesh_utils.create_device_mesh((2, 4))\\n\",\n \"print(device_mesh)\\n\",\n \"\\n\",\n \"mesh = Mesh(devices=device_mesh, axis_names=('data', 'model'))\\n\",\n \"print(mesh)\\n\",\n \"\\n\",\n \"def mesh_sharding(pspec: PartitionSpec) -> NamedSharding:\\n\",\n \" return NamedSharding(mesh, pspec)\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Define a layer\\n\",\n \"\\n\",\n \"Before defining a simple model, create an example layer called `DotReluDot` (by subclassing `flax.linen.Module`). The layer creates two parameters `W1` and `W2` for dot product multiplication, and uses the `jax.nn.relu` (ReLU) activation function in-between.\\n\",\n \"\\n\",\n \"To shard the parameters efficiently, apply the following APIs to annotate the parameters and intermediate variables:\\n\",\n \"\\n\",\n \"1. Use [`flax.linen.with_partitioning`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/spmd.html#flax.linen.with_partitioning) to decorate the initializer function when creating sub-layers or raw parameters.\\n\",\n \"\\n\",\n \"2. Apply [`jax.lax.with_sharding_constraint`](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.with_sharding_constraint.html) (formerly, `pjit.with_sharding_constraint`) to annotate intermediate variables like `y` and `z` to force a particular sharding pattern when the ideal constraint is known.\\n\",\n \"\\n\",\n \" * This step is optional, but can sometimes help auto-SPMD to partition efficiently. In the example below, the call is not required, because XLA will figure out the same sharding layout for `y` and `z` regardless.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 7,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"class DotReluDot(nn.Module):\\n\",\n \" depth: int\\n\",\n \" dense_init: Callable = nn.initializers.xavier_normal()\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self, x):\\n\",\n \"\\n\",\n \" y = nn.Dense(self.depth,\\n\",\n \" kernel_init=nn.with_partitioning(self.dense_init, (None, 'model')),\\n\",\n \" use_bias=False, # or overwrite with `bias_init`\\n\",\n \" )(x)\\n\",\n \"\\n\",\n \" y = jax.nn.relu(y)\\n\",\n \" # Force a local sharding annotation.\\n\",\n \" y = with_sharding_constraint(y, mesh_sharding(PartitionSpec('data', 'model')))\\n\",\n \"\\n\",\n \" W2 = self.param(\\n\",\n \" 'W2',\\n\",\n \" nn.with_partitioning(self.dense_init, ('model', None)),\\n\",\n \" (self.depth, x.shape[-1]))\\n\",\n \"\\n\",\n \" z = jnp.dot(y, W2)\\n\",\n \" # Force a local sharding annotation.\\n\",\n \" z = with_sharding_constraint(z, mesh_sharding(PartitionSpec('data', None)))\\n\",\n \"\\n\",\n \" # Return a tuple to conform with the API `flax.linen.scan` as shown in the cell below.\\n\",\n \" return z, None\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Note that device axis names like `'data'`, `'model'` or `None` are passed into both [`flax.linen.with_partitioning`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/spmd.html#flax.linen.with_partitioning) and [`jax.lax.with_sharding_constraint`](https://github.com/jax-ml/jax/blob/main/jax/_src/pjit.py#L1516) API calls. This refers to how each dimension of this data should be sharded — either across one of the device mesh dimensions, or not sharded at all.\\n\",\n \"\\n\",\n \"For example:\\n\",\n \"\\n\",\n \"* When you define `W1` with shape `(x.shape[-1], self.depth)` and annotate as `(None, 'model')`:\\n\",\n \"\\n\",\n \" * The first dimension (of length `x.shape[-1]`) will be replicated across all devices.\\n\",\n \" * The second dimension (of length `self.depth`) will be sharded over the `'model'` axis of the device mesh. This means `W1` will be sharded 4-way on devices `(0, 4)`, `(1, 5)`, `(2, 6)` and `(3, 7)`, on this dimension.\\n\",\n \"\\n\",\n \"* When you annotate the output `z` as `('data', None)`:\\n\",\n \"\\n\",\n \" * The first dimension — the batch dimension — will be sharded over the `'data'` axis. This means half of the batch will be processed on devices `0-3` (first four devices), and another half on devices `4-7` (the remaining four devices).\\n\",\n \" * The second dimension — the data depth dimension — will be replicated across all devices.\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Define a model with `flax.linen.scan` lifted transformation\\n\",\n \"\\n\",\n \"Having created `DotReluDot`, you can now define the `MLP` model (by subclassing [`flax.linen.Module`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module)) as multiple layers of `DotReluDot`.\\n\",\n \"\\n\",\n \"To replicate identical layers, you can either use [`flax.linen.scan`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/transformations.html#flax.linen.scan), or a for-loop:\\n\",\n \"\\n\",\n \"* `flax.linen.scan` can provide faster compilation times.\\n\",\n \"* The for-loop can be faster on runtime.\\n\",\n \"\\n\",\n \"The code below shows how to apply both methods, and default with the for-loop, so that all the parameters are two-dimensional and you can visualize their sharding.\\n\",\n \"\\n\",\n \"The `flax.linen.scan` code is just to show that this API works with [Flax lifted transforms](https://flax.readthedocs.io/en/latest/developer_notes/lift.html#supported-transformations).\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 8,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"class MLP(nn.Module):\\n\",\n \" num_layers: int\\n\",\n \" depth: int\\n\",\n \" use_scan: bool\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self, x):\\n\",\n \" if self.use_scan:\\n\",\n \" x, _ = nn.scan(DotReluDot, length=self.num_layers,\\n\",\n \" variable_axes={\\\"params\\\": 0},\\n\",\n \" split_rngs={\\\"params\\\": True},\\n\",\n \" metadata_params={nn.PARTITION_NAME: None}\\n\",\n \" )(self.depth)(x)\\n\",\n \" else:\\n\",\n \" for i in range(self.num_layers):\\n\",\n \" x, _ = DotReluDot(self.depth)(x)\\n\",\n \" return x\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Now, create a `model` instance, and a sample input `x`.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 9,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# MLP hyperparameters.\\n\",\n \"BATCH, LAYERS, DEPTH, USE_SCAN = 8, 4, 1024, False\\n\",\n \"# Create fake inputs.\\n\",\n \"x = jnp.ones((BATCH, DEPTH))\\n\",\n \"# Initialize a PRNG key.\\n\",\n \"k = random.key(0)\\n\",\n \"\\n\",\n \"# Create an Optax optimizer.\\n\",\n \"optimizer = optax.adam(learning_rate=0.001)\\n\",\n \"# Instantiate the model.\\n\",\n \"model = MLP(LAYERS, DEPTH, USE_SCAN)\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Specify sharding\\n\",\n \"\\n\",\n \"Next, you need to tell `jax.jit` how to shard our data across devices.\\n\",\n \"\\n\",\n \"### The input's sharding\\n\",\n \"\\n\",\n \"For data parallelism, you can shard the batched _input_ `x` across the `data` axis by denoting the batch axis as `'data'`. Then, use [`jax.device_put`](https://jax.readthedocs.io/en/latest/_autosummary/jax.device_put.html) to place it onto the correct `device`s.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 10,\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/html\": [\n \"
┌──────────────────────────────────────────────────────────────────────────────┐\\n\",\n       \"│                                                                              │\\n\",\n       \"│                                 CPU 0,1,2,3                                  │\\n\",\n       \"│                                                                              │\\n\",\n       \"│                                                                              │\\n\",\n       \"├──────────────────────────────────────────────────────────────────────────────┤\\n\",\n       \"│                                                                              │\\n\",\n       \"│                                 CPU 4,5,6,7                                  │\\n\",\n       \"│                                                                              │\\n\",\n       \"│                                                                              │\\n\",\n       \"└──────────────────────────────────────────────────────────────────────────────┘\\n\",\n       \"
\\n\"\n ],\n \"text/plain\": [\n \"┌──────────────────────────────────────────────────────────────────────────────┐\\n\",\n \"│ │\\n\",\n \"│ CPU \\u001b[1;36m0\\u001b[0m,\\u001b[1;36m1\\u001b[0m,\\u001b[1;36m2\\u001b[0m,\\u001b[1;36m3\\u001b[0m │\\n\",\n \"│ │\\n\",\n \"│ │\\n\",\n \"├──────────────────────────────────────────────────────────────────────────────┤\\n\",\n \"│ │\\n\",\n \"│ CPU \\u001b[1;36m4\\u001b[0m,\\u001b[1;36m5\\u001b[0m,\\u001b[1;36m6\\u001b[0m,\\u001b[1;36m7\\u001b[0m │\\n\",\n \"│ │\\n\",\n \"│ │\\n\",\n \"└──────────────────────────────────────────────────────────────────────────────┘\\n\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"x_sharding = mesh_sharding(PartitionSpec('data', None)) # dimensions: (batch, length)\\n\",\n \"x = jax.device_put(x, x_sharding)\\n\",\n \"jax.debug.visualize_array_sharding(x)\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### The output's sharding\\n\",\n \"\\n\",\n \"You need to compile `model.init()` (that is, [`flax.linen.Module.init()`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.init)), and its output as a pytree of parameters. Additionally, you may sometimes need wrap it with a [`flax.training.train_state`](https://flax.readthedocs.io/en/latest/api_reference/flax.training.html#flax.training.train_state.TrainState) to track other variables, such as optimizer states, and that would make the output an even more complex pytree.\\n\",\n \"\\n\",\n \"To achieve this, luckily, you don't have to hardcode the output's sharding by hand. Instead, you can:\\n\",\n \"\\n\",\n \"1. Evaluate `model.init` (in this case, a wrapper of it) abstractly using [`jax.eval_shape`](https://jax.readthedocs.io/en/latest/_autosummary/jax.eval_shape.html).\\n\",\n \"\\n\",\n \"1. Use [`flax.linen.get_sharding`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/spmd.html#flax.linen.get_sharding) to automatically generate the `jax.sharding.NamedSharding`.\\n\",\n \" * This step utilizes the [`flax.linen.with_partitioning`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/spmd.html#flax.linen.with_partitioning) annotations in the earlier definition to generate the correct sharding for the parameters.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 11,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"def init_fn(k, x, model, optimizer):\\n\",\n \" variables = model.init(k, x) # Initialize the model.\\n\",\n \" state = train_state.TrainState.create( # Create a `TrainState`.\\n\",\n \" apply_fn=model.apply,\\n\",\n \" params=variables['params'],\\n\",\n \" tx=optimizer)\\n\",\n \" return state\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 12,\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"TrainState(step=NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec()), apply_fn=, params={'DotReluDot_0': {'Dense_0': {'kernel': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec(None, 'model'))}, 'W2': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec('model', None))}, 'DotReluDot_1': {'Dense_0': {'kernel': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec(None, 'model'))}, 'W2': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec('model', None))}, 'DotReluDot_2': {'Dense_0': {'kernel': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec(None, 'model'))}, 'W2': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec('model', None))}, 'DotReluDot_3': {'Dense_0': {'kernel': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec(None, 'model'))}, 'W2': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec('model', None))}}, tx=GradientTransformationExtraArgs(init=.init_fn at 0x33e134280>, update=.update_fn at 0x33e134430>), opt_state=(ScaleByAdamState(count=NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec()), mu={'DotReluDot_0': {'Dense_0': {'kernel': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec(None, 'model'))}, 'W2': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec('model', None))}, 'DotReluDot_1': {'Dense_0': {'kernel': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec(None, 'model'))}, 'W2': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec('model', None))}, 'DotReluDot_2': {'Dense_0': {'kernel': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec(None, 'model'))}, 'W2': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec('model', None))}, 'DotReluDot_3': {'Dense_0': {'kernel': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec(None, 'model'))}, 'W2': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec('model', None))}}, nu={'DotReluDot_0': {'Dense_0': {'kernel': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec(None, 'model'))}, 'W2': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec('model', None))}, 'DotReluDot_1': {'Dense_0': {'kernel': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec(None, 'model'))}, 'W2': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec('model', None))}, 'DotReluDot_2': {'Dense_0': {'kernel': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec(None, 'model'))}, 'W2': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec('model', None))}, 'DotReluDot_3': {'Dense_0': {'kernel': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec(None, 'model'))}, 'W2': NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec('model', None))}}), EmptyState()))\"\n ]\n },\n \"execution_count\": 12,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"# Create an abstract closure to wrap the function before feeding it in\\n\",\n \"# because `jax.eval_shape` only takes pytrees as arguments.\\n\",\n \"abstract_variables = jax.eval_shape(\\n\",\n \" functools.partial(init_fn, model=model, optimizer=optimizer), k, x)\\n\",\n \"\\n\",\n \"# This `state_sharding` has the same pytree structure as `state`, the output\\n\",\n \"# of the `init_fn`.\\n\",\n \"state_sharding = nn.get_sharding(abstract_variables, mesh)\\n\",\n \"state_sharding\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Compile the code\\n\",\n \"\\n\",\n \"Now you can apply [`jax.jit`](https://jax.readthedocs.io/en/latest/jax-101/02-jitting.html) to your `init_fn`, but with two extra arguments: `in_shardings` and `out_shardings`.\\n\",\n \"\\n\",\n \"Run it to get the `initialized_state`, in which parameters are sharded exactly as instructed:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 16,\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/html\": [\n \"
┌───────┬───────┬───────┬───────┐\\n\",\n       \"│       │       │       │       │\\n\",\n       \"│       │       │       │       │\\n\",\n       \"│       │       │       │       │\\n\",\n       \"│       │       │       │       │\\n\",\n       \"│CPU 0,4│CPU 1,5│CPU 2,6│CPU 3,7│\\n\",\n       \"│       │       │       │       │\\n\",\n       \"│       │       │       │       │\\n\",\n       \"│       │       │       │       │\\n\",\n       \"│       │       │       │       │\\n\",\n       \"└───────┴───────┴───────┴───────┘\\n\",\n       \"
\\n\"\n ],\n \"text/plain\": [\n \"┌───────┬───────┬───────┬───────┐\\n\",\n \"│ │ │ │ │\\n\",\n \"│ │ │ │ │\\n\",\n \"│ │ │ │ │\\n\",\n \"│ │ │ │ │\\n\",\n \"│CPU \\u001b[1;36m0\\u001b[0m,\\u001b[1;36m4\\u001b[0m│CPU \\u001b[1;36m1\\u001b[0m,\\u001b[1;36m5\\u001b[0m│CPU \\u001b[1;36m2\\u001b[0m,\\u001b[1;36m6\\u001b[0m│CPU \\u001b[1;36m3\\u001b[0m,\\u001b[1;36m7\\u001b[0m│\\n\",\n \"│ │ │ │ │\\n\",\n \"│ │ │ │ │\\n\",\n \"│ │ │ │ │\\n\",\n \"│ │ │ │ │\\n\",\n \"└───────┴───────┴───────┴───────┘\\n\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"text/html\": [\n \"
┌───────────────────────┐\\n\",\n       \"│        CPU 0,4        │\\n\",\n       \"├───────────────────────┤\\n\",\n       \"│        CPU 1,5        │\\n\",\n       \"├───────────────────────┤\\n\",\n       \"│        CPU 2,6        │\\n\",\n       \"├───────────────────────┤\\n\",\n       \"│        CPU 3,7        │\\n\",\n       \"└───────────────────────┘\\n\",\n       \"
\\n\"\n ],\n \"text/plain\": [\n \"┌───────────────────────┐\\n\",\n \"│ CPU \\u001b[1;36m0\\u001b[0m,\\u001b[1;36m4\\u001b[0m │\\n\",\n \"├───────────────────────┤\\n\",\n \"│ CPU \\u001b[1;36m1\\u001b[0m,\\u001b[1;36m5\\u001b[0m │\\n\",\n \"├───────────────────────┤\\n\",\n \"│ CPU \\u001b[1;36m2\\u001b[0m,\\u001b[1;36m6\\u001b[0m │\\n\",\n \"├───────────────────────┤\\n\",\n \"│ CPU \\u001b[1;36m3\\u001b[0m,\\u001b[1;36m7\\u001b[0m │\\n\",\n \"└───────────────────────┘\\n\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"jit_init_fn = jax.jit(init_fn, static_argnums=(2, 3),\\n\",\n \" in_shardings=(mesh_sharding(PartitionSpec()), x_sharding), # PRNG key and x\\n\",\n \" out_shardings=state_sharding)\\n\",\n \"\\n\",\n \"initialized_state = jit_init_fn(k, x, model, optimizer)\\n\",\n \"\\n\",\n \"# for weight, partitioned in initialized_state.params['DotReluDot_0'].items():\\n\",\n \"# print(f'Sharding of {weight}: {partitioned.names}')\\n\",\n \"jax.debug.visualize_array_sharding(initialized_state.params['DotReluDot_0']['Dense_0']['kernel'].value)\\n\",\n \"jax.debug.visualize_array_sharding(initialized_state.params['DotReluDot_0']['W2'].value)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Inspect the Module output\\n\",\n \"\\n\",\n \"Note that in the output of `initialized_state`, the `params` `W1` and `W2` are of type [`flax.linen.Partitioned`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/spmd.html#flax.linen.Partitioned). This is a wrapper around the actual `jax.Array` that allows Flax to record the axis names associated with it.\\n\",\n \"\\n\",\n \"You can access the raw `jax.Array`s by calling `flax.linen.meta.unbox()` upon the dictionary, or call `.value` upon individual variable. You can also use `flax.linen.meta.replace_boxed()` to change the underlying `jax.Array` without modifying the sharding annotations.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 18,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"\\n\",\n \"\\n\",\n \"(None, 'model')\\n\",\n \"(1024, 1024)\\n\"\n ]\n }\n ],\n \"source\": [\n \"print(type(initialized_state.params['DotReluDot_0']['Dense_0']['kernel']))\\n\",\n \"print(type(initialized_state.params['DotReluDot_0']['Dense_0']['kernel'].value))\\n\",\n \"print(initialized_state.params['DotReluDot_0']['Dense_0']['kernel'].names)\\n\",\n \"print(initialized_state.params['DotReluDot_0']['Dense_0']['kernel'].value.shape)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 19,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# Say for some unknown reason you want to make the whole param tree all-zero\\n\",\n \"unboxed_params = nn.meta.unbox(initialized_state.params)\\n\",\n \"all_zero = jax.tree.map(jnp.zeros_like, unboxed_params)\\n\",\n \"all_zero_params = nn.meta.replace_boxed(initialized_state.params, all_zero)\\n\",\n \"assert jnp.sum(nn.meta.unbox(all_zero_params['DotReluDot_0']['Dense_0']['kernel'])) == 0\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"You can also check the underlying [`jax.sharding`](https://jax.readthedocs.io/en/latest/jax.sharding.html) of each parameter, which is now more internal than `NamedSharding`. Note that numbers like `initialized_state.step` are replicated across all devices.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 20,\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec(None, 'model'))\"\n ]\n },\n \"execution_count\": 20,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"initialized_state.params['DotReluDot_0']['Dense_0']['kernel'].value.sharding\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 21,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"0\\n\"\n ]\n },\n {\n \"data\": {\n \"text/plain\": [\n \"NamedSharding(mesh=Mesh('data': 2, 'model': 4), spec=PartitionSpec())\"\n ]\n },\n \"execution_count\": 21,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"print(initialized_state.step)\\n\",\n \"initialized_state.step.sharding\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"You can use [`jax.tree_util.tree_map`](https://jax.readthedocs.io/en/latest/_autosummary/jax.tree_util.tree_map.html) to perform mass computation on a dict of boxed params, in the same way as on a dict of JAX arrays.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 22,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"{'Dense_0': {'kernel': Partitioned(value=(1024, 1024), names=(None, 'model'), mesh=None)}, 'W2': Partitioned(value=(1024, 1024), names=('model', None), mesh=None)}\\n\",\n \"\\n\",\n \"(1024, 1024)\\n\"\n ]\n }\n ],\n \"source\": [\n \"diff = jax.tree_util.tree_map(\\n\",\n \" lambda a, b: a - b,\\n\",\n \" initialized_state.params['DotReluDot_0'], initialized_state.params['DotReluDot_0'])\\n\",\n \"print(jax.tree_util.tree_map(jnp.shape, diff))\\n\",\n \"diff_array = diff['Dense_0']['kernel'].value\\n\",\n \"print(type(diff_array))\\n\",\n \"print(diff_array.shape)\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Compile the train step and inference\\n\",\n \"\\n\",\n \"Create a `jit`ted training step as follows:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 23,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"@functools.partial(jax.jit, in_shardings=(state_sharding, x_sharding),\\n\",\n \" out_shardings=state_sharding)\\n\",\n \"def train_step(state, x):\\n\",\n \" # A fake loss function.\\n\",\n \" def loss_unrolled(params):\\n\",\n \" y = model.apply({'params': params}, x)\\n\",\n \" return y.sum()\\n\",\n \" grad_fn = jax.grad(loss_unrolled)\\n\",\n \" grads = grad_fn(state.params)\\n\",\n \" state = state.apply_gradients(grads=grads)\\n\",\n \" return state\\n\",\n \"\\n\",\n \"with jax.set_mesh(mesh):\\n\",\n \" new_state = train_step(initialized_state, x)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 24,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"Sharding of Weight 1:\\n\"\n ]\n },\n {\n \"data\": {\n \"text/html\": [\n \"
┌───────┬───────┬───────┬───────┐\\n\",\n       \"│       │       │       │       │\\n\",\n       \"│       │       │       │       │\\n\",\n       \"│       │       │       │       │\\n\",\n       \"│       │       │       │       │\\n\",\n       \"│CPU 0,4│CPU 1,5│CPU 2,6│CPU 3,7│\\n\",\n       \"│       │       │       │       │\\n\",\n       \"│       │       │       │       │\\n\",\n       \"│       │       │       │       │\\n\",\n       \"│       │       │       │       │\\n\",\n       \"└───────┴───────┴───────┴───────┘\\n\",\n       \"
\\n\"\n ],\n \"text/plain\": [\n \"┌───────┬───────┬───────┬───────┐\\n\",\n \"│ │ │ │ │\\n\",\n \"│ │ │ │ │\\n\",\n \"│ │ │ │ │\\n\",\n \"│ │ │ │ │\\n\",\n \"│CPU \\u001b[1;36m0\\u001b[0m,\\u001b[1;36m4\\u001b[0m│CPU \\u001b[1;36m1\\u001b[0m,\\u001b[1;36m5\\u001b[0m│CPU \\u001b[1;36m2\\u001b[0m,\\u001b[1;36m6\\u001b[0m│CPU \\u001b[1;36m3\\u001b[0m,\\u001b[1;36m7\\u001b[0m│\\n\",\n \"│ │ │ │ │\\n\",\n \"│ │ │ │ │\\n\",\n \"│ │ │ │ │\\n\",\n \"│ │ │ │ │\\n\",\n \"└───────┴───────┴───────┴───────┘\\n\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"Sharding of Weight 2:\\n\"\n ]\n },\n {\n \"data\": {\n \"text/html\": [\n \"
┌───────────────────────┐\\n\",\n       \"│        CPU 0,4        │\\n\",\n       \"├───────────────────────┤\\n\",\n       \"│        CPU 1,5        │\\n\",\n       \"├───────────────────────┤\\n\",\n       \"│        CPU 2,6        │\\n\",\n       \"├───────────────────────┤\\n\",\n       \"│        CPU 3,7        │\\n\",\n       \"└───────────────────────┘\\n\",\n       \"
\\n\"\n ],\n \"text/plain\": [\n \"┌───────────────────────┐\\n\",\n \"│ CPU \\u001b[1;36m0\\u001b[0m,\\u001b[1;36m4\\u001b[0m │\\n\",\n \"├───────────────────────┤\\n\",\n \"│ CPU \\u001b[1;36m1\\u001b[0m,\\u001b[1;36m5\\u001b[0m │\\n\",\n \"├───────────────────────┤\\n\",\n \"│ CPU \\u001b[1;36m2\\u001b[0m,\\u001b[1;36m6\\u001b[0m │\\n\",\n \"├───────────────────────┤\\n\",\n \"│ CPU \\u001b[1;36m3\\u001b[0m,\\u001b[1;36m7\\u001b[0m │\\n\",\n \"└───────────────────────┘\\n\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"print(f'Sharding of Weight 1:')\\n\",\n \"jax.debug.visualize_array_sharding(initialized_state.params['DotReluDot_0']['Dense_0']['kernel'].value)\\n\",\n \"print(f'Sharding of Weight 2:')\\n\",\n \"jax.debug.visualize_array_sharding(initialized_state.params['DotReluDot_0']['W2'].value)\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Then, create a compiled inference step. Note that the output is also sharded along `(data, None)`.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 25,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"\\n\",\n \"float32\\n\",\n \"(8, 1024)\\n\"\n ]\n },\n {\n \"data\": {\n \"text/html\": [\n \"
┌──────────────────────────────────────────────────────────────────────────────┐\\n\",\n       \"│                                                                              │\\n\",\n       \"│                                 CPU 0,1,2,3                                  │\\n\",\n       \"│                                                                              │\\n\",\n       \"│                                                                              │\\n\",\n       \"├──────────────────────────────────────────────────────────────────────────────┤\\n\",\n       \"│                                                                              │\\n\",\n       \"│                                 CPU 4,5,6,7                                  │\\n\",\n       \"│                                                                              │\\n\",\n       \"│                                                                              │\\n\",\n       \"└──────────────────────────────────────────────────────────────────────────────┘\\n\",\n       \"
\\n\"\n ],\n \"text/plain\": [\n \"┌──────────────────────────────────────────────────────────────────────────────┐\\n\",\n \"│ │\\n\",\n \"│ CPU \\u001b[1;36m0\\u001b[0m,\\u001b[1;36m1\\u001b[0m,\\u001b[1;36m2\\u001b[0m,\\u001b[1;36m3\\u001b[0m │\\n\",\n \"│ │\\n\",\n \"│ │\\n\",\n \"├──────────────────────────────────────────────────────────────────────────────┤\\n\",\n \"│ │\\n\",\n \"│ CPU \\u001b[1;36m4\\u001b[0m,\\u001b[1;36m5\\u001b[0m,\\u001b[1;36m6\\u001b[0m,\\u001b[1;36m7\\u001b[0m │\\n\",\n \"│ │\\n\",\n \"│ │\\n\",\n \"└──────────────────────────────────────────────────────────────────────────────┘\\n\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"@functools.partial(jax.jit, in_shardings=(state_sharding, x_sharding),\\n\",\n \" out_shardings=x_sharding)\\n\",\n \"def apply_fn(state, x):\\n\",\n \" return state.apply_fn({'params': state.params}, x)\\n\",\n \"\\n\",\n \"with jax.set_mesh(mesh):\\n\",\n \" y = apply_fn(new_state, x)\\n\",\n \"print(type(y))\\n\",\n \"print(y.dtype)\\n\",\n \"print(y.shape)\\n\",\n \"jax.debug.visualize_array_sharding(y)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Profiling\\n\",\n \"\\n\",\n \"If you are running on a TPU pod or a pod slice, you can use a custom `block_all` utility function, as defined below, to measure the performance:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 26,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"20.9 ms ± 319 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)\\n\"\n ]\n }\n ],\n \"source\": [\n \"%%timeit\\n\",\n \"\\n\",\n \"def block_all(xs):\\n\",\n \" jax.tree_util.tree_map(lambda x: x.block_until_ready(), xs)\\n\",\n \" return xs\\n\",\n \"\\n\",\n \"with jax.set_mesh(mesh):\\n\",\n \" new_state = block_all(train_step(initialized_state, x))\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Logical axis annotation\\n\",\n \"\\n\",\n \"JAX's automatic SPMD encourages users to explore different sharding layouts to find the optimal one. To this end, in Flax you actually can annotate the dimensions of any data with more descriptive axis names (not just device mesh axis names like `'data'` and `'model'`).\\n\",\n \"\\n\",\n \"The `LogicalDotReluDot` and `LogicalMLP` Module definition below are similar to the Modules you created earlier, except for the following:\\n\",\n \"\\n\",\n \"1. All axes are annotated with more concrete, meaningful names, such as `'embed'`, `'hidden'`, `'batch'` and `'layer'`. These names are referred to as _logical axis names_ in Flax. They make the dimensional changes inside model definitions more readable.\\n\",\n \"\\n\",\n \"2. [`flax.linen.with_logical_partitioning`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/spmd.html#flax.linen.with_logical_partitioning) replaces `flax.linen.with_partitioning`; and [`flax.linen.with_logical_constraint`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/spmd.html#flax.linen.with_logical_constraint) replaces `jax.lax.with_sharding_constraint`, to recognize the logical axis names.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 27,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"class LogicalDotReluDot(nn.Module):\\n\",\n \" depth: int\\n\",\n \" dense_init: Callable = nn.initializers.xavier_normal()\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self, x):\\n\",\n \" y = nn.Dense(self.depth,\\n\",\n \" kernel_init=nn.with_logical_partitioning(self.dense_init, ('embed', 'hidden')),\\n\",\n \" use_bias=False, # or overwrite with `bias_init`\\n\",\n \" )(x)\\n\",\n \"\\n\",\n \" y = jax.nn.relu(y)\\n\",\n \" # Force a local sharding annotation.\\n\",\n \" y = with_sharding_constraint(y, mesh_sharding(PartitionSpec('data', 'model')))\\n\",\n \"\\n\",\n \" W2 = self.param(\\n\",\n \" 'W2',\\n\",\n \" nn.with_logical_partitioning(self.dense_init, ('hidden', 'embed')),\\n\",\n \" (self.depth, x.shape[-1]))\\n\",\n \"\\n\",\n \" z = jnp.dot(y, W2)\\n\",\n \" # Force a local sharding annotation.\\n\",\n \" z = nn.with_logical_constraint(z, ('batch', 'embed'))\\n\",\n \" return z, None\\n\",\n \"\\n\",\n \"class LogicalMLP(nn.Module):\\n\",\n \" num_layers: int\\n\",\n \" depth: int\\n\",\n \" use_scan: bool\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self, x):\\n\",\n \" if self.use_scan:\\n\",\n \" x, _ = nn.scan(LogicalDotReluDot, length=self.num_layers,\\n\",\n \" variable_axes={\\\"params\\\": 0},\\n\",\n \" split_rngs={\\\"params\\\": True},\\n\",\n \" metadata_params={nn.PARTITION_NAME: 'layer'}\\n\",\n \" )(self.depth)(x)\\n\",\n \" else:\\n\",\n \" for i in range(self.num_layers):\\n\",\n \" x, _ = LogicalDotReluDot(self.depth)(x)\\n\",\n \" return x\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Now, initiate a model and try to figure out what sharding its `state` should have.\\n\",\n \"\\n\",\n \"To allow the device mesh to take your model correctly, you need to decide which of these logical axis names are mapped to the device axis `'data'` or `'model'`. This rule is a list of (`logical_axis_name`, `device_axis_name`) tuples, and [`flax.linen.logical_to_mesh_sharding`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/spmd.html#flax.linen.logical_to_mesh_sharding) will convert them to the kind of sharding that the device mesh can understand.\\n\",\n \"\\n\",\n \"This allows you to change the rules and try out new partition layouts without modifying the model definition.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 28,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"annotations are logical, not mesh-specific: PartitionSpec('embed', 'hidden')\\n\",\n \"sharding annotations are mesh-specific: PartitionSpec(None, 'model')\\n\"\n ]\n }\n ],\n \"source\": [\n \"# Unspecified rule means unsharded by default, so no need to specify `('embed', None)` and `('layer', None)`.\\n\",\n \"rules = (('batch', 'data'),\\n\",\n \" ('hidden', 'model'))\\n\",\n \"\\n\",\n \"logical_model = LogicalMLP(LAYERS, DEPTH, USE_SCAN)\\n\",\n \"\\n\",\n \"logical_abstract_variables = jax.eval_shape(\\n\",\n \" functools.partial(init_fn, model=logical_model, optimizer=optimizer), k, x)\\n\",\n \"logical_state_spec = nn.get_partition_spec(logical_abstract_variables)\\n\",\n \"print('annotations are logical, not mesh-specific: ',\\n\",\n \" logical_state_spec.params['LogicalDotReluDot_0']['Dense_0']['kernel'])\\n\",\n \"\\n\",\n \"logical_state_sharding = nn.logical_to_mesh_sharding(logical_state_spec, mesh, rules)\\n\",\n \"print('sharding annotations are mesh-specific: ',\\n\",\n \" logical_state_sharding.params['LogicalDotReluDot_0']['Dense_0']['kernel'].spec)\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"You can verify that the `logical_state_spec` here has the same content as `state_spec` in the previous (\\\"non-logical\\\") example. This allows you to `jax.jit` your Module's [`flax.linen.Module.init`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.init) and [`flax.linen.Module.apply`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.apply) the same way in the above above.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 29,\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"True\"\n ]\n },\n \"execution_count\": 29,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"state_sharding.params['DotReluDot_0'] == logical_state_sharding.params['LogicalDotReluDot_0']\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 31,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"logical_jit_init_fn = jax.jit(init_fn, static_argnums=(2, 3),\\n\",\n \" in_shardings=(mesh_sharding(PartitionSpec()), x_sharding), # PRNG key and x\\n\",\n \" out_shardings=logical_state_sharding)\\n\",\n \"\\n\",\n \"logical_initialized_state = logical_jit_init_fn(k, x, logical_model, optimizer)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 32,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"Sharding of Weight 1:\\n\"\n ]\n },\n {\n \"data\": {\n \"text/html\": [\n \"
┌───────┬───────┬───────┬───────┐\\n\",\n       \"│       │       │       │       │\\n\",\n       \"│       │       │       │       │\\n\",\n       \"│       │       │       │       │\\n\",\n       \"│       │       │       │       │\\n\",\n       \"│CPU 0,4│CPU 1,5│CPU 2,6│CPU 3,7│\\n\",\n       \"│       │       │       │       │\\n\",\n       \"│       │       │       │       │\\n\",\n       \"│       │       │       │       │\\n\",\n       \"│       │       │       │       │\\n\",\n       \"└───────┴───────┴───────┴───────┘\\n\",\n       \"
\\n\"\n ],\n \"text/plain\": [\n \"┌───────┬───────┬───────┬───────┐\\n\",\n \"│ │ │ │ │\\n\",\n \"│ │ │ │ │\\n\",\n \"│ │ │ │ │\\n\",\n \"│ │ │ │ │\\n\",\n \"│CPU \\u001b[1;36m0\\u001b[0m,\\u001b[1;36m4\\u001b[0m│CPU \\u001b[1;36m1\\u001b[0m,\\u001b[1;36m5\\u001b[0m│CPU \\u001b[1;36m2\\u001b[0m,\\u001b[1;36m6\\u001b[0m│CPU \\u001b[1;36m3\\u001b[0m,\\u001b[1;36m7\\u001b[0m│\\n\",\n \"│ │ │ │ │\\n\",\n \"│ │ │ │ │\\n\",\n \"│ │ │ │ │\\n\",\n \"│ │ │ │ │\\n\",\n \"└───────┴───────┴───────┴───────┘\\n\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"Sharding of Weight 2:\\n\"\n ]\n },\n {\n \"data\": {\n \"text/html\": [\n \"
┌───────────────────────┐\\n\",\n       \"│        CPU 0,4        │\\n\",\n       \"├───────────────────────┤\\n\",\n       \"│        CPU 1,5        │\\n\",\n       \"├───────────────────────┤\\n\",\n       \"│        CPU 2,6        │\\n\",\n       \"├───────────────────────┤\\n\",\n       \"│        CPU 3,7        │\\n\",\n       \"└───────────────────────┘\\n\",\n       \"
\\n\"\n ],\n \"text/plain\": [\n \"┌───────────────────────┐\\n\",\n \"│ CPU \\u001b[1;36m0\\u001b[0m,\\u001b[1;36m4\\u001b[0m │\\n\",\n \"├───────────────────────┤\\n\",\n \"│ CPU \\u001b[1;36m1\\u001b[0m,\\u001b[1;36m5\\u001b[0m │\\n\",\n \"├───────────────────────┤\\n\",\n \"│ CPU \\u001b[1;36m2\\u001b[0m,\\u001b[1;36m6\\u001b[0m │\\n\",\n \"├───────────────────────┤\\n\",\n \"│ CPU \\u001b[1;36m3\\u001b[0m,\\u001b[1;36m7\\u001b[0m │\\n\",\n \"└───────────────────────┘\\n\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"print(f'Sharding of Weight 1:')\\n\",\n \"jax.debug.visualize_array_sharding(logical_initialized_state.params['LogicalDotReluDot_0']['Dense_0']['kernel'].value)\\n\",\n \"print(f'Sharding of Weight 2:')\\n\",\n \"jax.debug.visualize_array_sharding(logical_initialized_state.params['LogicalDotReluDot_0']['W2'].value)\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## When to use device axis / logical axis\\n\",\n \"\\n\",\n \"Choosing when to use a device or logical axis depends on how much you want to control the partitioning of your model:\\n\",\n \"\\n\",\n \"* **Device mesh axis**: If you want a very simple model, or you are very confident of your way of partitioning, defining it with __device mesh axis__ can potentially save you a few extra lines of code of converting the logical naming back to the device naming.\\n\",\n \"\\n\",\n \"* **Logical naming**: On the other hand, the __logical naming__ helpers can be useful for exploring different sharding layouts. Use this if you want to experiment around and find the most optimal partition layout for your model.\\n\",\n \"\\n\",\n \"* **Device axis names**: In really advanced use cases, you may have more complicated sharding patterns that require annotating *activation* dimension names differently from *parameter* dimension names. If you wish to have more fine-grained control on manual mesh assignments, directly using __device axis names__ could be more helpful.\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Save the data\\n\",\n \"\\n\",\n \"To save the cross-device array, you can use Orbax as shown in the [Save and load checkpoints guide - Multi-host/multi-process checkpointing](https://flax.readthedocs.io/en/latest/guides/training_techniques/use_checkpointing.html#multi-host-multi-process-checkpointing). This is especially required if you are running on a multi-host environment (for example, a TPU pod).\\n\",\n \"\\n\",\n \"In practice, you might want to save the raw `jax.Array` pytree as checkpoint, instead of the wrapped `Partitioned` values, to reduce complexity. You can restore it as-is and put it back into an annotated pytree with `flax.linen.meta.replace_boxed()`.\\n\",\n \"\\n\",\n \"Keep in mind that to restore the arrays to the desired partition, you need to provide a sample `target` pytree that has the same structure and has the desired [`jax.sharding.Sharding`](https://jax.readthedocs.io/en/latest/jax.sharding.html#jax.sharding.Sharding) in place for each JAX array. The sharding you use to restore the array doesn't necessarily need to be the same as the ones you used to store the array.\"\n ]\n }\n ],\n \"metadata\": {\n \"jupytext\": {\n \"formats\": \"ipynb,md:myst\"\n },\n \"language_info\": {\n \"codemirror_mode\": {\n \"name\": \"ipython\",\n \"version\": 3\n },\n \"file_extension\": \".py\",\n \"mimetype\": \"text/x-python\",\n \"name\": \"python\",\n \"nbconvert_exporter\": \"python\",\n \"pygments_lexer\": \"ipython3\",\n \"version\": \"3.10.13\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 0\n}\n" }, { "path": "docs/guides/parallel_training/flax_on_pjit.md", "content": "---\njupytext:\n formats: ipynb,md:myst\n text_representation:\n extension: .md\n format_name: myst\n format_version: 0.13\n jupytext_version: 1.13.8\n---\n\n# Scale up Flax Modules on multiple devices\n\nThis guide shows how to scale up [Flax Modules](https://flax.readthedocs.io/en/latest/developer_notes/module_lifecycle.html) on multiple devices and hosts using [`jax.jit`](https://jax.readthedocs.io/en/latest/jax-101/02-jitting.html) (formerly [`experimental.pjit`](https://jax.readthedocs.io/en/latest/jax.experimental.pjit.html#module-jax.experimental.pjit)) and [`flax.linen`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/index.html).\n\n+++\n\n## Flax and `jax.jit` scaled up\n\n[`jax.jit`](https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html) follows the [Single Program Multi Data (SPMD)](https://jax.readthedocs.io/en/latest/glossary.html#term-SPMD) paradigm and automatically compiles your code to run it on multiple devices. You need to only specify how you want the input and output of your code to be partitioned, and the compiler will figure out how to: 1) partition everything inside; and 2) compile inter-device communications.\n\nFlax provides several functionalities that can help you use auto-SPMD on [Flax Modules](https://flax.readthedocs.io/en/latest/developer_notes/module_lifecycle.html), including:\n\n1. An interface to specify partitions of your data when defining [`flax.linen.Module`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html).\n2. Utility functions to generate the sharding information that `jax.jit` requires to run.\n3. An interface to customize your axis names called \"logical axis annotations\" to decouple both your Module code and partition plan to experiment with different partition layouts more easily.\n\nYou can learn more about `jax.jit` APIs for scaling up in [JAX in multi-process environments](https://jax.readthedocs.io/en/latest/multi_process.html) and [Distributed arrays and automatic parallelization](https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html) on JAX's documentation site.\n\n+++\n\n## Setup\n\nImport some necessary dependencies.\n\n**Note:** This guide uses the `--xla_force_host_platform_device_count=8` flag to emulate multiple devices in a CPU environment in a Google Colab/Jupyter Notebook. You don't need this if you are already using a multi-device TPU environment.\n\n```{code-cell} ipython3\nimport os\nos.environ[\"XLA_FLAGS\"] = '--xla_force_host_platform_device_count=8'\n```\n\n```{code-cell} ipython3\nimport functools\nfrom typing import Optional, Callable\n\nimport numpy as np\nimport jax\nfrom jax import lax, random, numpy as jnp\n\nimport flax\nfrom flax import struct, traverse_util, linen as nn\nfrom flax.core import freeze, unfreeze\nfrom flax.training import train_state, checkpoints\n\nimport optax # Optax for common losses and optimizers.\n```\n\n```{code-cell} ipython3\nprint(f'We have 8 fake JAX devices now: {jax.devices()}')\n```\n\nThe code below shows how to import and set up the JAX-level device API, following JAX's [Distributed arrays and automatic parallelization](https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html) guide:\n\n1. Start a 2x4 device `mesh` (8 devices) using JAX's `mesh_utils.create_device_mesh`. This layout is the same as the one of a [TPU v3-8](https://cloud.google.com/tpu/docs/system-architecture-tpu-vm#single_tpu_board).\n\n2. Annotate each axis with a name using the `axis_names` parameter in `jax.sharding.Mesh`. A typical way to annotate axis names is `axis_name=('data', 'model')`, where:\n * `'data'`: the mesh dimension used for data-parallel sharding of the batch dimension of inputs and activations.\n * `'model'`: the mesh dimension used for sharding parameters of the model across devices.\n\n3. Make a simple utility function `mesh_sharding` for generating a sharding object from the mesh and any layout.\n\n```{code-cell} ipython3\nfrom jax.sharding import Mesh, PartitionSpec, NamedSharding\nfrom jax.lax import with_sharding_constraint\nfrom jax.experimental import mesh_utils\n```\n\n```{code-cell} ipython3\n# Create a mesh and annotate each axis with a name.\ndevice_mesh = mesh_utils.create_device_mesh((2, 4))\nprint(device_mesh)\n\nmesh = Mesh(devices=device_mesh, axis_names=('data', 'model'))\nprint(mesh)\n\ndef mesh_sharding(pspec: PartitionSpec) -> NamedSharding:\n return NamedSharding(mesh, pspec)\n```\n\n## Define a layer\n\nBefore defining a simple model, create an example layer called `DotReluDot` (by subclassing `flax.linen.Module`). The layer creates two parameters `W1` and `W2` for dot product multiplication, and uses the `jax.nn.relu` (ReLU) activation function in-between.\n\nTo shard the parameters efficiently, apply the following APIs to annotate the parameters and intermediate variables:\n\n1. Use [`flax.linen.with_partitioning`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/spmd.html#flax.linen.with_partitioning) to decorate the initializer function when creating sub-layers or raw parameters.\n\n2. Apply [`jax.lax.with_sharding_constraint`](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.with_sharding_constraint.html) (formerly, `pjit.with_sharding_constraint`) to annotate intermediate variables like `y` and `z` to force a particular sharding pattern when the ideal constraint is known.\n\n * This step is optional, but can sometimes help auto-SPMD to partition efficiently. In the example below, the call is not required, because XLA will figure out the same sharding layout for `y` and `z` regardless.\n\n```{code-cell} ipython3\nclass DotReluDot(nn.Module):\n depth: int\n dense_init: Callable = nn.initializers.xavier_normal()\n @nn.compact\n def __call__(self, x):\n\n y = nn.Dense(self.depth,\n kernel_init=nn.with_partitioning(self.dense_init, (None, 'model')),\n use_bias=False, # or overwrite with `bias_init`\n )(x)\n\n y = jax.nn.relu(y)\n # Force a local sharding annotation.\n y = with_sharding_constraint(y, mesh_sharding(PartitionSpec('data', 'model')))\n\n W2 = self.param(\n 'W2',\n nn.with_partitioning(self.dense_init, ('model', None)),\n (self.depth, x.shape[-1]))\n\n z = jnp.dot(y, W2)\n # Force a local sharding annotation.\n z = with_sharding_constraint(z, mesh_sharding(PartitionSpec('data', None)))\n\n # Return a tuple to conform with the API `flax.linen.scan` as shown in the cell below.\n return z, None\n```\n\nNote that device axis names like `'data'`, `'model'` or `None` are passed into both [`flax.linen.with_partitioning`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/spmd.html#flax.linen.with_partitioning) and [`jax.lax.with_sharding_constraint`](https://github.com/jax-ml/jax/blob/main/jax/_src/pjit.py#L1516) API calls. This refers to how each dimension of this data should be sharded — either across one of the device mesh dimensions, or not sharded at all.\n\nFor example:\n\n* When you define `W1` with shape `(x.shape[-1], self.depth)` and annotate as `(None, 'model')`:\n\n * The first dimension (of length `x.shape[-1]`) will be replicated across all devices.\n * The second dimension (of length `self.depth`) will be sharded over the `'model'` axis of the device mesh. This means `W1` will be sharded 4-way on devices `(0, 4)`, `(1, 5)`, `(2, 6)` and `(3, 7)`, on this dimension.\n\n* When you annotate the output `z` as `('data', None)`:\n\n * The first dimension — the batch dimension — will be sharded over the `'data'` axis. This means half of the batch will be processed on devices `0-3` (first four devices), and another half on devices `4-7` (the remaining four devices).\n * The second dimension — the data depth dimension — will be replicated across all devices.\n\n+++\n\n## Define a model with `flax.linen.scan` lifted transformation\n\nHaving created `DotReluDot`, you can now define the `MLP` model (by subclassing [`flax.linen.Module`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module)) as multiple layers of `DotReluDot`.\n\nTo replicate identical layers, you can either use [`flax.linen.scan`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/transformations.html#flax.linen.scan), or a for-loop:\n\n* `flax.linen.scan` can provide faster compilation times.\n* The for-loop can be faster on runtime.\n\nThe code below shows how to apply both methods, and default with the for-loop, so that all the parameters are two-dimensional and you can visualize their sharding.\n\nThe `flax.linen.scan` code is just to show that this API works with [Flax lifted transforms](https://flax.readthedocs.io/en/latest/developer_notes/lift.html#supported-transformations).\n\n```{code-cell} ipython3\nclass MLP(nn.Module):\n num_layers: int\n depth: int\n use_scan: bool\n @nn.compact\n def __call__(self, x):\n if self.use_scan:\n x, _ = nn.scan(DotReluDot, length=self.num_layers,\n variable_axes={\"params\": 0},\n split_rngs={\"params\": True},\n metadata_params={nn.PARTITION_NAME: None}\n )(self.depth)(x)\n else:\n for i in range(self.num_layers):\n x, _ = DotReluDot(self.depth)(x)\n return x\n```\n\nNow, create a `model` instance, and a sample input `x`.\n\n```{code-cell} ipython3\n# MLP hyperparameters.\nBATCH, LAYERS, DEPTH, USE_SCAN = 8, 4, 1024, False\n# Create fake inputs.\nx = jnp.ones((BATCH, DEPTH))\n# Initialize a PRNG key.\nk = random.key(0)\n\n# Create an Optax optimizer.\noptimizer = optax.adam(learning_rate=0.001)\n# Instantiate the model.\nmodel = MLP(LAYERS, DEPTH, USE_SCAN)\n```\n\n## Specify sharding\n\nNext, you need to tell `jax.jit` how to shard our data across devices.\n\n### The input's sharding\n\nFor data parallelism, you can shard the batched _input_ `x` across the `data` axis by denoting the batch axis as `'data'`. Then, use [`jax.device_put`](https://jax.readthedocs.io/en/latest/_autosummary/jax.device_put.html) to place it onto the correct `device`s.\n\n```{code-cell} ipython3\nx_sharding = mesh_sharding(PartitionSpec('data', None)) # dimensions: (batch, length)\nx = jax.device_put(x, x_sharding)\njax.debug.visualize_array_sharding(x)\n```\n\n### The output's sharding\n\nYou need to compile `model.init()` (that is, [`flax.linen.Module.init()`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.init)), and its output as a pytree of parameters. Additionally, you may sometimes need wrap it with a [`flax.training.train_state`](https://flax.readthedocs.io/en/latest/api_reference/flax.training.html#flax.training.train_state.TrainState) to track other variables, such as optimizer states, and that would make the output an even more complex pytree.\n\nTo achieve this, luckily, you don't have to hardcode the output's sharding by hand. Instead, you can:\n\n1. Evaluate `model.init` (in this case, a wrapper of it) abstractly using [`jax.eval_shape`](https://jax.readthedocs.io/en/latest/_autosummary/jax.eval_shape.html).\n\n1. Use [`flax.linen.get_sharding`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/spmd.html#flax.linen.get_sharding) to automatically generate the `jax.sharding.NamedSharding`.\n * This step utilizes the [`flax.linen.with_partitioning`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/spmd.html#flax.linen.with_partitioning) annotations in the earlier definition to generate the correct sharding for the parameters.\n\n```{code-cell} ipython3\ndef init_fn(k, x, model, optimizer):\n variables = model.init(k, x) # Initialize the model.\n state = train_state.TrainState.create( # Create a `TrainState`.\n apply_fn=model.apply,\n params=variables['params'],\n tx=optimizer)\n return state\n```\n\n```{code-cell} ipython3\n# Create an abstract closure to wrap the function before feeding it in\n# because `jax.eval_shape` only takes pytrees as arguments.\nabstract_variables = jax.eval_shape(\n functools.partial(init_fn, model=model, optimizer=optimizer), k, x)\n\n# This `state_sharding` has the same pytree structure as `state`, the output\n# of the `init_fn`.\nstate_sharding = nn.get_sharding(abstract_variables, mesh)\nstate_sharding\n```\n\n## Compile the code\n\nNow you can apply [`jax.jit`](https://jax.readthedocs.io/en/latest/jax-101/02-jitting.html) to your `init_fn`, but with two extra arguments: `in_shardings` and `out_shardings`.\n\nRun it to get the `initialized_state`, in which parameters are sharded exactly as instructed:\n\n```{code-cell} ipython3\njit_init_fn = jax.jit(init_fn, static_argnums=(2, 3),\n in_shardings=(mesh_sharding(PartitionSpec()), x_sharding), # PRNG key and x\n out_shardings=state_sharding)\n\ninitialized_state = jit_init_fn(k, x, model, optimizer)\n\n# for weight, partitioned in initialized_state.params['DotReluDot_0'].items():\n# print(f'Sharding of {weight}: {partitioned.names}')\njax.debug.visualize_array_sharding(initialized_state.params['DotReluDot_0']['Dense_0']['kernel'].value)\njax.debug.visualize_array_sharding(initialized_state.params['DotReluDot_0']['W2'].value)\n```\n\n## Inspect the Module output\n\nNote that in the output of `initialized_state`, the `params` `W1` and `W2` are of type [`flax.linen.Partitioned`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/spmd.html#flax.linen.Partitioned). This is a wrapper around the actual `jax.Array` that allows Flax to record the axis names associated with it.\n\nYou can access the raw `jax.Array`s by calling `flax.linen.meta.unbox()` upon the dictionary, or call `.value` upon individual variable. You can also use `flax.linen.meta.replace_boxed()` to change the underlying `jax.Array` without modifying the sharding annotations.\n\n```{code-cell} ipython3\nprint(type(initialized_state.params['DotReluDot_0']['Dense_0']['kernel']))\nprint(type(initialized_state.params['DotReluDot_0']['Dense_0']['kernel'].value))\nprint(initialized_state.params['DotReluDot_0']['Dense_0']['kernel'].names)\nprint(initialized_state.params['DotReluDot_0']['Dense_0']['kernel'].value.shape)\n```\n\n```{code-cell} ipython3\n# Say for some unknown reason you want to make the whole param tree all-zero\nunboxed_params = nn.meta.unbox(initialized_state.params)\nall_zero = jax.tree.map(jnp.zeros_like, unboxed_params)\nall_zero_params = nn.meta.replace_boxed(initialized_state.params, all_zero)\nassert jnp.sum(nn.meta.unbox(all_zero_params['DotReluDot_0']['Dense_0']['kernel'])) == 0\n```\n\nYou can also check the underlying [`jax.sharding`](https://jax.readthedocs.io/en/latest/jax.sharding.html) of each parameter, which is now more internal than `NamedSharding`. Note that numbers like `initialized_state.step` are replicated across all devices.\n\n```{code-cell} ipython3\ninitialized_state.params['DotReluDot_0']['Dense_0']['kernel'].value.sharding\n```\n\n```{code-cell} ipython3\nprint(initialized_state.step)\ninitialized_state.step.sharding\n```\n\nYou can use [`jax.tree_util.tree_map`](https://jax.readthedocs.io/en/latest/_autosummary/jax.tree_util.tree_map.html) to perform mass computation on a dict of boxed params, in the same way as on a dict of JAX arrays.\n\n```{code-cell} ipython3\ndiff = jax.tree_util.tree_map(\n lambda a, b: a - b,\n initialized_state.params['DotReluDot_0'], initialized_state.params['DotReluDot_0'])\nprint(jax.tree_util.tree_map(jnp.shape, diff))\ndiff_array = diff['Dense_0']['kernel'].value\nprint(type(diff_array))\nprint(diff_array.shape)\n```\n\n## Compile the train step and inference\n\nCreate a `jit`ted training step as follows:\n\n```{code-cell} ipython3\n@functools.partial(jax.jit, in_shardings=(state_sharding, x_sharding),\n out_shardings=state_sharding)\ndef train_step(state, x):\n # A fake loss function.\n def loss_unrolled(params):\n y = model.apply({'params': params}, x)\n return y.sum()\n grad_fn = jax.grad(loss_unrolled)\n grads = grad_fn(state.params)\n state = state.apply_gradients(grads=grads)\n return state\n\nwith jax.set_mesh(mesh):\n new_state = train_step(initialized_state, x)\n```\n\n```{code-cell} ipython3\nprint(f'Sharding of Weight 1:')\njax.debug.visualize_array_sharding(initialized_state.params['DotReluDot_0']['Dense_0']['kernel'].value)\nprint(f'Sharding of Weight 2:')\njax.debug.visualize_array_sharding(initialized_state.params['DotReluDot_0']['W2'].value)\n```\n\nThen, create a compiled inference step. Note that the output is also sharded along `(data, None)`.\n\n```{code-cell} ipython3\n@functools.partial(jax.jit, in_shardings=(state_sharding, x_sharding),\n out_shardings=x_sharding)\ndef apply_fn(state, x):\n return state.apply_fn({'params': state.params}, x)\n\nwith jax.set_mesh(mesh):\n y = apply_fn(new_state, x)\nprint(type(y))\nprint(y.dtype)\nprint(y.shape)\njax.debug.visualize_array_sharding(y)\n```\n\n## Profiling\n\nIf you are running on a TPU pod or a pod slice, you can use a custom `block_all` utility function, as defined below, to measure the performance:\n\n```{code-cell} ipython3\n%%timeit\n\ndef block_all(xs):\n jax.tree_util.tree_map(lambda x: x.block_until_ready(), xs)\n return xs\n\nwith jax.set_mesh(mesh):\n new_state = block_all(train_step(initialized_state, x))\n```\n\n## Logical axis annotation\n\nJAX's automatic SPMD encourages users to explore different sharding layouts to find the optimal one. To this end, in Flax you actually can annotate the dimensions of any data with more descriptive axis names (not just device mesh axis names like `'data'` and `'model'`).\n\nThe `LogicalDotReluDot` and `LogicalMLP` Module definition below are similar to the Modules you created earlier, except for the following:\n\n1. All axes are annotated with more concrete, meaningful names, such as `'embed'`, `'hidden'`, `'batch'` and `'layer'`. These names are referred to as _logical axis names_ in Flax. They make the dimensional changes inside model definitions more readable.\n\n2. [`flax.linen.with_logical_partitioning`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/spmd.html#flax.linen.with_logical_partitioning) replaces `flax.linen.with_partitioning`; and [`flax.linen.with_logical_constraint`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/spmd.html#flax.linen.with_logical_constraint) replaces `jax.lax.with_sharding_constraint`, to recognize the logical axis names.\n\n```{code-cell} ipython3\nclass LogicalDotReluDot(nn.Module):\n depth: int\n dense_init: Callable = nn.initializers.xavier_normal()\n @nn.compact\n def __call__(self, x):\n y = nn.Dense(self.depth,\n kernel_init=nn.with_logical_partitioning(self.dense_init, ('embed', 'hidden')),\n use_bias=False, # or overwrite with `bias_init`\n )(x)\n\n y = jax.nn.relu(y)\n # Force a local sharding annotation.\n y = with_sharding_constraint(y, mesh_sharding(PartitionSpec('data', 'model')))\n\n W2 = self.param(\n 'W2',\n nn.with_logical_partitioning(self.dense_init, ('hidden', 'embed')),\n (self.depth, x.shape[-1]))\n\n z = jnp.dot(y, W2)\n # Force a local sharding annotation.\n z = nn.with_logical_constraint(z, ('batch', 'embed'))\n return z, None\n\nclass LogicalMLP(nn.Module):\n num_layers: int\n depth: int\n use_scan: bool\n @nn.compact\n def __call__(self, x):\n if self.use_scan:\n x, _ = nn.scan(LogicalDotReluDot, length=self.num_layers,\n variable_axes={\"params\": 0},\n split_rngs={\"params\": True},\n metadata_params={nn.PARTITION_NAME: 'layer'}\n )(self.depth)(x)\n else:\n for i in range(self.num_layers):\n x, _ = LogicalDotReluDot(self.depth)(x)\n return x\n```\n\nNow, initiate a model and try to figure out what sharding its `state` should have.\n\nTo allow the device mesh to take your model correctly, you need to decide which of these logical axis names are mapped to the device axis `'data'` or `'model'`. This rule is a list of (`logical_axis_name`, `device_axis_name`) tuples, and [`flax.linen.logical_to_mesh_sharding`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/spmd.html#flax.linen.logical_to_mesh_sharding) will convert them to the kind of sharding that the device mesh can understand.\n\nThis allows you to change the rules and try out new partition layouts without modifying the model definition.\n\n```{code-cell} ipython3\n# Unspecified rule means unsharded by default, so no need to specify `('embed', None)` and `('layer', None)`.\nrules = (('batch', 'data'),\n ('hidden', 'model'))\n\nlogical_model = LogicalMLP(LAYERS, DEPTH, USE_SCAN)\n\nlogical_abstract_variables = jax.eval_shape(\n functools.partial(init_fn, model=logical_model, optimizer=optimizer), k, x)\nlogical_state_spec = nn.get_partition_spec(logical_abstract_variables)\nprint('annotations are logical, not mesh-specific: ',\n logical_state_spec.params['LogicalDotReluDot_0']['Dense_0']['kernel'])\n\nlogical_state_sharding = nn.logical_to_mesh_sharding(logical_state_spec, mesh, rules)\nprint('sharding annotations are mesh-specific: ',\n logical_state_sharding.params['LogicalDotReluDot_0']['Dense_0']['kernel'].spec)\n```\n\nYou can verify that the `logical_state_spec` here has the same content as `state_spec` in the previous (\"non-logical\") example. This allows you to `jax.jit` your Module's [`flax.linen.Module.init`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.init) and [`flax.linen.Module.apply`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.apply) the same way in the above above.\n\n```{code-cell} ipython3\nstate_sharding.params['DotReluDot_0'] == logical_state_sharding.params['LogicalDotReluDot_0']\n```\n\n```{code-cell} ipython3\nlogical_jit_init_fn = jax.jit(init_fn, static_argnums=(2, 3),\n in_shardings=(mesh_sharding(PartitionSpec()), x_sharding), # PRNG key and x\n out_shardings=logical_state_sharding)\n\nlogical_initialized_state = logical_jit_init_fn(k, x, logical_model, optimizer)\n```\n\n```{code-cell} ipython3\nprint(f'Sharding of Weight 1:')\njax.debug.visualize_array_sharding(logical_initialized_state.params['LogicalDotReluDot_0']['Dense_0']['kernel'].value)\nprint(f'Sharding of Weight 2:')\njax.debug.visualize_array_sharding(logical_initialized_state.params['LogicalDotReluDot_0']['W2'].value)\n```\n\n## When to use device axis / logical axis\n\nChoosing when to use a device or logical axis depends on how much you want to control the partitioning of your model:\n\n* **Device mesh axis**: If you want a very simple model, or you are very confident of your way of partitioning, defining it with __device mesh axis__ can potentially save you a few extra lines of code of converting the logical naming back to the device naming.\n\n* **Logical naming**: On the other hand, the __logical naming__ helpers can be useful for exploring different sharding layouts. Use this if you want to experiment around and find the most optimal partition layout for your model.\n\n* **Device axis names**: In really advanced use cases, you may have more complicated sharding patterns that require annotating *activation* dimension names differently from *parameter* dimension names. If you wish to have more fine-grained control on manual mesh assignments, directly using __device axis names__ could be more helpful.\n\n+++\n\n## Save the data\n\nTo save the cross-device array, you can use Orbax as shown in the [Save and load checkpoints guide - Multi-host/multi-process checkpointing](https://flax.readthedocs.io/en/latest/guides/training_techniques/use_checkpointing.html#multi-host-multi-process-checkpointing). This is especially required if you are running on a multi-host environment (for example, a TPU pod).\n\nIn practice, you might want to save the raw `jax.Array` pytree as checkpoint, instead of the wrapped `Partitioned` values, to reduce complexity. You can restore it as-is and put it back into an annotated pytree with `flax.linen.meta.replace_boxed()`.\n\nKeep in mind that to restore the arrays to the desired partition, you need to provide a sample `target` pytree that has the same structure and has the desired [`jax.sharding.Sharding`](https://jax.readthedocs.io/en/latest/jax.sharding.html#jax.sharding.Sharding) in place for each JAX array. The sharding you use to restore the array doesn't necessarily need to be the same as the ones you used to store the array.\n" }, { "path": "docs/guides/parallel_training/index.rst", "content": "Parallel training\n=================\n\n.. toctree::\n :maxdepth: 1\n\n ensembling\n flax_on_pjit\n" }, { "path": "docs/guides/quantization/fp8_basics.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"id\": \"ca360491\",\n \"metadata\": {},\n \"source\": [\n \"# User Guide on Using FP8\\n\",\n \"\\n\",\n \"JAX supports various FP8 formats, including E4M3 (jnp.float8_e4m3fn) and E5M2\\n\",\n \"(jnp.float8_e5m2). Due to the limited range of FP8 data types, higher-precision\\n\",\n \"data must be scaled to fit within the FP8 representable range, a process known\\n\",\n \"as quantization (Q). Conversely, de-quantization (DQ) rescales the FP8 data back\\n\",\n \"to its original type.\\n\",\n \"\\n\",\n \"While jnp.dot supports FP8 inputs directly, proper quantization and\\n\",\n \"dequantization is needed for optimal performance. Flax provides\\n\",\n \"nn.fp8_ops.Fp8DotGeneral and nn.fp8_ops.Fp8Einsum modules that handle\\n\",\n \"this automatically and can be used with existing layers like nn.Dense.\\n\",\n \"\\n\",\n \"This tutorial will walk you through the basics of how to use it.\\n\",\n \"\\n\",\n \"## Setting up our environment\\n\",\n \"\\n\",\n \"Here, we provide the code necessary to set up the environment for our notebook.\\n\",\n \"Additionally, we define a function to check if the XLA-optimized HLO will indeed\\n\",\n \"call an FP8 dot operation under the hood.\\n\",\n \"\\n\",\n \"*Note: This tutorial relies on the XLA-FP8 feature, which is only supported on\\n\",\n \"NVIDIA Hopper GPUs or later.*\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"177b91c4\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"import flax\\n\",\n \"import jax\\n\",\n \"import re\\n\",\n \"import pprint\\n\",\n \"from jax import random\\n\",\n \"from jax import numpy as jnp\\n\",\n \"from jax._src import test_util as jtu\\n\",\n \"from flax import linen as nn\\n\",\n \"from flax.linen import fp8_ops\\n\",\n \"\\n\",\n \"e4m3 = jnp.float8_e4m3fn\\n\",\n \"f32 = jnp.float32\\n\",\n \"E4M3_MAX = jnp.finfo(e4m3).max.astype(f32)\\n\",\n \"\\n\",\n \"assert jtu.is_cuda_compute_capability_at_least(\\\"9.0\\\")\\n\",\n \"\\n\",\n \"def check_fp8_call(lowered):\\n\",\n \" hlo = lowered.compile()\\n\",\n \" if re.search(r\\\"custom-call\\\\(f8e4m3fn.*, f8e4m3fn.*\\\", hlo.as_text()):\\n\",\n \" print(\\\"Fp8 call detected!\\\")\\n\",\n \" else:\\n\",\n \" print(\\\"No Fp8 call!\\\")\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"4adc021f\",\n \"metadata\": {},\n \"source\": [\n \"## FLAX Low Level API\\n\",\n \"\\n\",\n \"The JAX dot operations (e.g. `jnp.dot`) support the FP8 dtype inputs. So it is\\n\",\n \"legal to do the following call:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"c54c374e\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"k0, k1 = random.split(random.key(0), 2)\\n\",\n \"a = random.uniform(k0, (16, 32))\\n\",\n \"b = random.uniform(k1, (32, 64))\\n\",\n \"@jax.jit\\n\",\n \"def dot_fp8(a, b):\\n\",\n \" return jnp.dot(a.astype(e4m3), b.astype(e4m3), preferred_element_type=f32)\\n\",\n \"check_fp8_call(dot_fp8.lower(a, b))\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"adb22878\",\n \"metadata\": {},\n \"source\": [\n \"However, this approach has two key limitations:\\n\",\n \"\\n\",\n \"1. `jnp.dot` does not support custom scaling factors for operands, defaulting to\\n\",\n \" a scale of 1.0\\n\",\n \"2. The autodiff does not automatically use E5M2 for gradients and E4M3 for\\n\",\n \" activations/weights during training, which is the recommended practice\\n\",\n \"\\n\",\n \"To overcome these limitations and implement proper FP8 matrix multiplication, we\\n\",\n \"recommend using the Flax FP8 APIs. Let's start with a basic scaling approach.\\n\",\n \"\\n\",\n \"\\n\",\n \"### Current Scaling\\n\",\n \"\\n\",\n \"Scaling factors are usually defined as `scale = amax(x) / MAX`, where `amax` is\\n\",\n \"an operation to find the absolute maximum value of the tensor, and `MAX` is the\\n\",\n \"maximum value of the representable range of the target dtype. This scaling\\n\",\n \"approach allows us to derive the scaling factors directly from the current\\n\",\n \"operand tensors of the dot product.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"f0e746e3\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"@jax.jit\\n\",\n \"def dot_fp8(a, b):\\n\",\n \" a_scale = jnp.max(jnp.abs(A)) / E4M3_MAX\\n\",\n \" b_scale = jnp.max(jnp.abs(B)) / E4M3_MAX\\n\",\n \" a = fp8_ops.quantize(a, e4m3, a_scale, f32)\\n\",\n \" b = fp8_ops.quantize(b, e4m3, b_scale, f32)\\n\",\n \"\\n\",\n \" c = jnp.dot(a, b, preferred_element_type=f32)\\n\",\n \" c = fp8_ops.dequantize(c, f32, a_scale * b_scale)\\n\",\n \" return c\\n\",\n \"\\n\",\n \"c = dot_fp8(a, b)\\n\",\n \"check_fp8_call(dot_fp8.lower(a, b))\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"59aca6fe\",\n \"metadata\": {},\n \"source\": [\n \"As shown in the code, we perform quantization (`fp8_ops.quantize`) on the\\n\",\n \"tensors to get the lower precision operands. The `jnp.dot` processes them and\\n\",\n \"accumulates the output in high precision (i.e., the `preferred_element_type`).\\n\",\n \"After that, we multiply the result by the scaling factors to dequantize back to\\n\",\n \"the original range (`fp8_ops.dequantize`). Note that while this example uses\\n\",\n \"E4M3 for both inputs, it is possible to use different FP8 dtypes like E4M3 and\\n\",\n \"E5M2 for the inputs. The quantization method and the scaling factors can also be\\n\",\n \"customized based on application needs.\\n\",\n \"\\n\",\n \"One major issue with the current scaling method is the performance overhead\\n\",\n \"introduced by computing `a_scale` and `b_scale`, which requires additional\\n\",\n \"loading of the operand tensors. To overcome this issue, we recommend the delayed\\n\",\n \"scaling.\\n\",\n \"\\n\",\n \"### Delayed Scaling\\n\",\n \"\\n\",\n \"In delayed scaling, we use a scaling factor associated with an amax history. The\\n\",\n \"scaling factor remains a scalar, but the amax history is a list that stores amax\\n\",\n \"values from recent steps (e.g., 1024 steps). Both tensors are computed from\\n\",\n \"previous steps and maintained in the model parameters.\\n\",\n \"\\n\",\n \"The quantization and dequantization operations for delayed scaling are provided\\n\",\n \"by `fp8_ops.in_q` and `fp8_ops.out_dq` respectively. `fp8_ops.in_q` handles\\n\",\n \"input quantization and update the amax history and scaling factor, while\\n\",\n \"`fp8_ops.out_dq` performs output dequantization.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"cf466308\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"a_scale = jnp.array(1.0)\\n\",\n \"b_scale = jnp.array(1.0)\\n\",\n \"a_amax_hist = jnp.zeros((1024,))\\n\",\n \"b_amax_hist = jnp.zeros((1024,))\\n\",\n \"\\n\",\n \"@jax.jit\\n\",\n \"def dot_fp8(a, a_scale, a_amax_hist, b, b_scale, b_amax_hist):\\n\",\n \" a, a_scale = fp8_ops.in_q(f32, e4m3, a, a_scale, a_amax_hist)\\n\",\n \" b, b_scale = fp8_ops.in_q(f32, e4m3, b, b_scale, b_amax_hist)\\n\",\n \" \\n\",\n \" c = jnp.dot(a, b, preferred_element_type=f32)\\n\",\n \" c = fp8_ops.out_dq(f32, a_scale, b_scale, c)\\n\",\n \" return c\\n\",\n \"\\n\",\n \"c = dot_fp8(a, a_scale, a_amax_hist, b, b_scale, b_amax_hist)\\n\",\n \"check_fp8_call(dot_fp8.lower(a, a_scale, a_amax_hist, b, b_scale, b_amax_hist))\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"b3bdc038\",\n \"metadata\": {},\n \"source\": [\n \"In this example, we first prepare three pairs of scaling factors and amax\\n\",\n \"histories, treating them as results computed from previous steps. Then, we apply\\n\",\n \"`fp8_ops.in_q` to the input operands of `jnp.dot`, followed by `fp8_ops.out_dq`\\n\",\n \"to the output of `jnp.dot`.\\n\",\n \"\\n\",\n \"\\n\",\n \"## FLAX High Level API\\n\",\n \"Flax provides high-level operations to seamlessly integrate FP8 quantization\\n\",\n \"into existing layers. Instead of manually handling quantization of the delayed\\n\",\n \"scaling (e.g., the maintanence of the amax history and scaling factors), users\\n\",\n \"can simply use these drop-in replacements:\\n\",\n \"\\n\",\n \"* `fp8_ops.Fp8DotGeneral` for `lax.dot_general` operations\\n\",\n \"* `fp8_ops.Fp8Einsum` for `jnp.einsum` operations \\n\",\n \"\\n\",\n \"These operations automatically handle all FP8-related functionality, including\\n\",\n \"quantization/dequantization, scale factor updates, and FP8 dtype selection for\\n\",\n \"both forward and backward passes.\\n\",\n \"\\n\",\n \"Consider the following example:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"bd8d9dba\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"model = nn.Dense(features=64, dot_general_cls=fp8_ops.Fp8DotGeneral)\\n\",\n \"params = model.init(k0, A)\\n\",\n \"\\n\",\n \"@jax.jit\\n\",\n \"def train_step(var, a): \\n\",\n \" c = model.apply(var, a)\\n\",\n \" return jnp.sum(c)\\n\",\n \"\\n\",\n \"check_fp8_call(train_step.lower(params, A))\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"ba280e79\",\n \"metadata\": {},\n \"source\": [\n \"By setting `dot_general_cls=fp8_ops.Fp8DotGeneral`, we replace the\\n\",\n \"default `lax.dot_general` operation in `nn.Dense` with an FP8-enabled version.\\n\",\n \"The model usage remains similar, but now includes additional parameters for FP8\\n\",\n \"quantization: scaling factors and amax history values. The next section explains\\n\",\n \"how to update these FP8-specific parameters.\\n\",\n \"\\n\",\n \"For models that use `jnp.einsum` operations, such as Mixture of Experts (MoE)\\n\",\n \"layers, users can replace them with `fp8_ops.Fp8Einsum` to enable FP8\\n\",\n \"quantization. Here's an example:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"961b4549\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"from typing import Any\\n\",\n \"class FooModule(nn.Module):\\n\",\n \" einsum: Any = None\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self, a, b):\\n\",\n \" if self.einsum is not None:\\n\",\n \" einsum_fn = self.einsum()\\n\",\n \" elif self.einsum is None:\\n\",\n \" einsum_fn = jnp.einsum\\n\",\n \" c = einsum_fn(\\\"mk,kn->mn\\\", a, b)\\n\",\n \" return c\\n\",\n \"\\n\",\n \"model = FooModule(einsum=fp8_ops.Fp8Einsum)\\n\",\n \"params = model.init(k0, a, b)\\n\",\n \"\\n\",\n \"@jax.jit\\n\",\n \"def train_step(var, a, b):\\n\",\n \" c = model.apply(var, a, b)\\n\",\n \" return jnp.sum(c)\\n\",\n \"\\n\",\n \"check_fp8_call(train_step.lower(params, a, b))\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"a83b0851\",\n \"metadata\": {},\n \"source\": [\n \"## Manipulate FP8 params\\n\",\n \"\\n\",\n \"The following sections explain the internal FP8 parameters managed by\\n\",\n \"`fp8_ops.Fp8DotGeneral` and `fp8_ops.Fp8Einsum`. These parameters\\n\",\n \"include scaling factors and amax history values that control the FP8\\n\",\n \"quantization process. While most users don't need to interact with these\\n\",\n \"directly, understanding them can be valuable for advanced optimization and\\n\",\n \"debugging.\\n\",\n \"\\n\",\n \"Let's first examine the data structure of `params`. In the code below, we redact\\n\",\n \"the parameter values and then display the PyTree structure.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"873799fe\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"params_structure = flax.core.unfreeze(params).copy()\\n\",\n \"params_structure = flax.traverse_util.flatten_dict(params_structure, sep='/')\\n\",\n \"for key, value in params_structure.items():\\n\",\n \" params_structure[key] = '*'\\n\",\n \"params_structure = flax.traverse_util.unflatten_dict(params_structure, sep='/')\\n\",\n \"pprint.pprint(params_structure)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"031894dc\",\n \"metadata\": {},\n \"source\": [\n \"The output is as follows:\\n\",\n \"\\n\",\n \"```plaintext\\n\",\n \"{'_overwrite_with_gradient': {'Fp8Einsum_0': {'input_amax_history': '*',\\n\",\n \" 'input_scale': '*',\\n\",\n \" 'kernel_amax_history': '*',\\n\",\n \" 'kernel_scale': '*',\\n\",\n \" 'output_grad_amax_history': '*',\\n\",\n \" 'output_grad_scale': '*'}}}\\n\",\n \"```\\n\",\n \"\\n\",\n \"In addition to the expected `params`, there is an additional category called\\n\",\n \"`_overwrite_with_gradient`. This category includes three pairs of `amax_history`\\n\",\n \"and `scale` for the activation, kernel, and dot gradient, respectively.\\n\",\n \"\\n\",\n \"### Update gradient of FP8 params\\n\",\n \"Now, we perform one training step to obtain the gradients and see how to use\\n\",\n \"them to update the parameters.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"593fc35f\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"step_fn = jax.jit(jax.grad(train_step, (0, 1)))\\n\",\n \"\\n\",\n \"grads = step_fn(params, A)\\n\",\n \"\\n\",\n \"params = flax.core.unfreeze(params)\\n\",\n \"params = flax.traverse_util.flatten_dict(params, sep='/')\\n\",\n \"grads = flax.traverse_util.flatten_dict(grads[0], sep='/')\\n\",\n \"\\n\",\n \"for key, value in params.items():\\n\",\n \" if key.startswith('params'):\\n\",\n \" params[key] = value + 0.01 * grads[key]\\n\",\n \" if key.startswith('_overwrite_with_gradient'):\\n\",\n \" params[key] = grads[key]\\n\",\n \"\\n\",\n \"params = flax.traverse_util.unflatten_dict(params, sep='/')\\n\",\n \"params = flax.core.freeze(params)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"1a8e2153\",\n \"metadata\": {},\n \"source\": [\n \"The above code demonstrates how to update both `params` and\\n\",\n \"`_overwrite_with_gradient`. For `params`, we use the formula `new_param =\\n\",\n \"old_param + 0.01 * grads`, where `0.01` is the learning rate (or users can use\\n\",\n \"whatever optimizers from `optax`). For `_overwrite_with_gradient`, we simply use\\n\",\n \"the gradient to overwrite the old values.\\n\",\n \"\\n\",\n \"Note that `flax.training.train_state.TrainState` conveniently supports the\\n\",\n \"category of `_overwrite_with_gradient`, so users do not need to modify their\\n\",\n \"scripts if they don't use custom `TrainState`.\\n\",\n \"\\n\",\n \"## Accumulate gradient of FP8 params\\n\",\n \"When the same parameter is used in a branched manner, the autograd mechanism\\n\",\n \"will add their gradients from these branches. This is common in scenarios like\\n\",\n \"pipeline parallelism, where each microbatch shares the same set of parameters\\n\",\n \"for the minibatch. However, for the `_overwrite_with_gradient` parameters, this\\n\",\n \"accumulation by addition is not meaningful. Instead, we prefer custom\\n\",\n \"accumulation by taking the maximum value.\\n\",\n \"\\n\",\n \"To address this, we introduce a custom dtype `fp8_ops.fp32_max_grad`. The basic\\n\",\n \"usage is demonstrated below:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"2d3a86e9\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"fmax32 = fp8_ops.fp32_max_grad\\n\",\n \"\\n\",\n \"def reuse_fp8_param(x, y, scale, amax_history):\\n\",\n \" scale = scale.astype(fmax32)\\n\",\n \" amax_history = amax_history.astype(fmax32)\\n\",\n \"\\n\",\n \" x = fp8_ops.in_qdq(f32, e4m3, x, scale, amax_history)\\n\",\n \" y = fp8_ops.in_qdq(f32, e4m3, y, scale, amax_history)\\n\",\n \" return x + y\\n\",\n \"\\n\",\n \"reuse_fp8_param_fn = jax.grad(reuse_fp8_param, (0, 1, 2, 3))\\n\",\n \"reuse_fp8_param_fn = jax.jit(reuse_fp8_param_fn)\\n\",\n \"\\n\",\n \"_, _, new_ah, new_sf = reuse_fp8_param_fn(2.0, 3.0, a_scale, a_amax_hist)\\n\",\n \"print(new_ah, new_sf)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"2321a9bb\",\n \"metadata\": {},\n \"source\": [\n \"In this example, we first cast the `scale` and `amax_history` to\\n\",\n \"`fp8_ops.fp32_max_grad` and then call `fp8_ops.in_qdq` twice using the same pair\\n\",\n \"of `scale` and `amax_history`. During autograd, their gradients from each branch\\n\",\n \"will be taken as the maximum, giving us the correct results of:\\n\",\n \"\\n\",\n \"```plaintext\\n\",\n \"1.0 [3. 0. 0. ... 0. 0. 0.]\\n\",\n \"```\\n\",\n \"\\n\",\n \"If we do not perform the type casting, we get the following result, meaning the\\n\",\n \"gradients of the two branches are added:\\n\",\n \"\\n\",\n \"```plaintext\\n\",\n \"2.0 [5. 0. 0. ... 0. 0. 0.]\\n\",\n \"```\\n\",\n \"\\n\",\n \"This casting is already included if users choose to use the high-level APIs.\\n\",\n \"\\n\",\n \"## Deprecated APIs\\n\",\n \"Previously, we provided APIs like `fp8_ops.quantize_dequantize` for current\\n\",\n \"scaling and `fp8_ops.[in|out]_qdq` for delayed scaling. These were used with\\n\",\n \"high precision dot operations, leveraging an XLA-FP8 feature that\\n\",\n \"pattern-matched QDQ->dot sequences to Q->fp8_cublas_gemm. The corresponding\\n\",\n \"high-level API was called `fp8_ops.Fp8DotGeneralOp`. However, this pattern\\n\",\n \"matching-based solution proved brittle, as the patterns could be easily broken\\n\",\n \"by other XLA optimizations. We recommend users migrate from these deprecated\\n\",\n \"APIs to the newer ones described above.\\n\",\n \"\\n\",\n \"For migration, users should replace:\\n\",\n \"* `fp8_ops.quantize_dequantize -> jnp.dot` with `fp8_ops.quantize -> jnp.dot ->\\n\",\n \" fp8_ops.dequantize`\\n\",\n \"* `fp8_ops.in_qdq -> jnp.dot -> fp8_ops.out_qdq` with `fp8_ops.in_q -> jnp.dot\\n\",\n \" -> fp8_ops.out_dq`\\n\",\n \"* `fp8_ops.Fp8DotGeneralOp` with `fp8_ops.Fp8DotGeneral`\\n\",\n \"\\n\",\n \"Additionally, we provide an einsum variant through `fp8_ops.Fp8Einsum`.\"\n ]\n }\n ],\n \"metadata\": {\n \"jupytext\": {\n \"formats\": \"ipynb,md:myst\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 5\n}\n" }, { "path": "docs/guides/quantization/fp8_basics.md", "content": "---\njupytext:\n formats: ipynb,md:myst\n text_representation:\n extension: .md\n format_name: myst\n format_version: 0.13\n jupytext_version: 1.13.8\n---\n\n# User Guide on Using FP8\n\nJAX supports various FP8 formats, including E4M3 (jnp.float8_e4m3fn) and E5M2\n(jnp.float8_e5m2). Due to the limited range of FP8 data types, higher-precision\ndata must be scaled to fit within the FP8 representable range, a process known\nas quantization (Q). Conversely, de-quantization (DQ) rescales the FP8 data back\nto its original type.\n\nWhile jnp.dot supports FP8 inputs directly, proper quantization and\ndequantization is needed for optimal performance. Flax provides\nnn.fp8_ops.Fp8DotGeneral and nn.fp8_ops.Fp8Einsum modules that handle\nthis automatically and can be used with existing layers like nn.Dense.\n\nThis tutorial will walk you through the basics of how to use it.\n\n## Setting up our environment\n\nHere, we provide the code necessary to set up the environment for our notebook.\nAdditionally, we define a function to check if the XLA-optimized HLO will indeed\ncall an FP8 dot operation under the hood.\n\n*Note: This tutorial relies on the XLA-FP8 feature, which is only supported on\nNVIDIA Hopper GPUs or later.*\n\n```{code-cell}\nimport flax\nimport jax\nimport re\nimport pprint\nfrom jax import random\nfrom jax import numpy as jnp\nfrom jax._src import test_util as jtu\nfrom flax import linen as nn\nfrom flax.linen import fp8_ops\n\ne4m3 = jnp.float8_e4m3fn\nf32 = jnp.float32\nE4M3_MAX = jnp.finfo(e4m3).max.astype(f32)\n\nassert jtu.is_cuda_compute_capability_at_least(\"9.0\")\n\ndef check_fp8_call(lowered):\n hlo = lowered.compile()\n if re.search(r\"custom-call\\(f8e4m3fn.*, f8e4m3fn.*\", hlo.as_text()):\n print(\"Fp8 call detected!\")\n else:\n print(\"No Fp8 call!\")\n```\n\n## FLAX Low Level API\n\nThe JAX dot operations (e.g. `jnp.dot`) support the FP8 dtype inputs. So it is\nlegal to do the following call:\n\n```{code-cell}\nk0, k1 = random.split(random.key(0), 2)\na = random.uniform(k0, (16, 32))\nb = random.uniform(k1, (32, 64))\n@jax.jit\ndef dot_fp8(a, b):\n return jnp.dot(a.astype(e4m3), b.astype(e4m3), preferred_element_type=f32)\ncheck_fp8_call(dot_fp8.lower(a, b))\n```\n\nHowever, this approach has two key limitations:\n\n1. `jnp.dot` does not support custom scaling factors for operands, defaulting to\n a scale of 1.0\n2. The autodiff does not automatically use E5M2 for gradients and E4M3 for\n activations/weights during training, which is the recommended practice\n\nTo overcome these limitations and implement proper FP8 matrix multiplication, we\nrecommend using the Flax FP8 APIs. Let's start with a basic scaling approach.\n\n\n### Current Scaling\n\nScaling factors are usually defined as `scale = amax(x) / MAX`, where `amax` is\nan operation to find the absolute maximum value of the tensor, and `MAX` is the\nmaximum value of the representable range of the target dtype. This scaling\napproach allows us to derive the scaling factors directly from the current\noperand tensors of the dot product.\n\n```{code-cell}\n@jax.jit\ndef dot_fp8(a, b):\n a_scale = jnp.max(jnp.abs(A)) / E4M3_MAX\n b_scale = jnp.max(jnp.abs(B)) / E4M3_MAX\n a = fp8_ops.quantize(a, e4m3, a_scale, f32)\n b = fp8_ops.quantize(b, e4m3, b_scale, f32)\n\n c = jnp.dot(a, b, preferred_element_type=f32)\n c = fp8_ops.dequantize(c, f32, a_scale * b_scale)\n return c\n\nc = dot_fp8(a, b)\ncheck_fp8_call(dot_fp8.lower(a, b))\n```\n\nAs shown in the code, we perform quantization (`fp8_ops.quantize`) on the\ntensors to get the lower precision operands. The `jnp.dot` processes them and\naccumulates the output in high precision (i.e., the `preferred_element_type`).\nAfter that, we multiply the result by the scaling factors to dequantize back to\nthe original range (`fp8_ops.dequantize`). Note that while this example uses\nE4M3 for both inputs, it is possible to use different FP8 dtypes like E4M3 and\nE5M2 for the inputs. The quantization method and the scaling factors can also be\ncustomized based on application needs.\n\nOne major issue with the current scaling method is the performance overhead\nintroduced by computing `a_scale` and `b_scale`, which requires additional\nloading of the operand tensors. To overcome this issue, we recommend the delayed\nscaling.\n\n### Delayed Scaling\n\nIn delayed scaling, we use a scaling factor associated with an amax history. The\nscaling factor remains a scalar, but the amax history is a list that stores amax\nvalues from recent steps (e.g., 1024 steps). Both tensors are computed from\nprevious steps and maintained in the model parameters.\n\nThe quantization and dequantization operations for delayed scaling are provided\nby `fp8_ops.in_q` and `fp8_ops.out_dq` respectively. `fp8_ops.in_q` handles\ninput quantization and update the amax history and scaling factor, while\n`fp8_ops.out_dq` performs output dequantization.\n\n```{code-cell}\na_scale = jnp.array(1.0)\nb_scale = jnp.array(1.0)\na_amax_hist = jnp.zeros((1024,))\nb_amax_hist = jnp.zeros((1024,))\n\n@jax.jit\ndef dot_fp8(a, a_scale, a_amax_hist, b, b_scale, b_amax_hist):\n a, a_scale = fp8_ops.in_q(f32, e4m3, a, a_scale, a_amax_hist)\n b, b_scale = fp8_ops.in_q(f32, e4m3, b, b_scale, b_amax_hist)\n \n c = jnp.dot(a, b, preferred_element_type=f32)\n c = fp8_ops.out_dq(f32, a_scale, b_scale, c)\n return c\n\nc = dot_fp8(a, a_scale, a_amax_hist, b, b_scale, b_amax_hist)\ncheck_fp8_call(dot_fp8.lower(a, a_scale, a_amax_hist, b, b_scale, b_amax_hist))\n```\n\nIn this example, we first prepare three pairs of scaling factors and amax\nhistories, treating them as results computed from previous steps. Then, we apply\n`fp8_ops.in_q` to the input operands of `jnp.dot`, followed by `fp8_ops.out_dq`\nto the output of `jnp.dot`.\n\n\n## FLAX High Level API\nFlax provides high-level operations to seamlessly integrate FP8 quantization\ninto existing layers. Instead of manually handling quantization of the delayed\nscaling (e.g., the maintanence of the amax history and scaling factors), users\ncan simply use these drop-in replacements:\n\n* `fp8_ops.Fp8DotGeneral` for `lax.dot_general` operations\n* `fp8_ops.Fp8Einsum` for `jnp.einsum` operations \n\nThese operations automatically handle all FP8-related functionality, including\nquantization/dequantization, scale factor updates, and FP8 dtype selection for\nboth forward and backward passes.\n\nConsider the following example:\n\n```{code-cell}\nmodel = nn.Dense(features=64, dot_general_cls=fp8_ops.Fp8DotGeneral)\nparams = model.init(k0, A)\n\n@jax.jit\ndef train_step(var, a): \n c = model.apply(var, a)\n return jnp.sum(c)\n\ncheck_fp8_call(train_step.lower(params, A))\n```\n\nBy setting `dot_general_cls=fp8_ops.Fp8DotGeneral`, we replace the\ndefault `lax.dot_general` operation in `nn.Dense` with an FP8-enabled version.\nThe model usage remains similar, but now includes additional parameters for FP8\nquantization: scaling factors and amax history values. The next section explains\nhow to update these FP8-specific parameters.\n\nFor models that use `jnp.einsum` operations, such as Mixture of Experts (MoE)\nlayers, users can replace them with `fp8_ops.Fp8Einsum` to enable FP8\nquantization. Here's an example:\n\n```{code-cell}\nfrom typing import Any\nclass FooModule(nn.Module):\n einsum: Any = None\n @nn.compact\n def __call__(self, a, b):\n if self.einsum is not None:\n einsum_fn = self.einsum()\n elif self.einsum is None:\n einsum_fn = jnp.einsum\n c = einsum_fn(\"mk,kn->mn\", a, b)\n return c\n\nmodel = FooModule(einsum=fp8_ops.Fp8Einsum)\nparams = model.init(k0, a, b)\n\n@jax.jit\ndef train_step(var, a, b):\n c = model.apply(var, a, b)\n return jnp.sum(c)\n\ncheck_fp8_call(train_step.lower(params, a, b))\n```\n\n## Manipulate FP8 params\n\nThe following sections explain the internal FP8 parameters managed by\n`fp8_ops.Fp8DotGeneral` and `fp8_ops.Fp8Einsum`. These parameters\ninclude scaling factors and amax history values that control the FP8\nquantization process. While most users don't need to interact with these\ndirectly, understanding them can be valuable for advanced optimization and\ndebugging.\n\nLet's first examine the data structure of `params`. In the code below, we redact\nthe parameter values and then display the PyTree structure.\n\n```{code-cell}\nparams_structure = flax.core.unfreeze(params).copy()\nparams_structure = flax.traverse_util.flatten_dict(params_structure, sep='/')\nfor key, value in params_structure.items():\n params_structure[key] = '*'\nparams_structure = flax.traverse_util.unflatten_dict(params_structure, sep='/')\npprint.pprint(params_structure)\n```\n\nThe output is as follows:\n\n```plaintext\n{'_overwrite_with_gradient': {'Fp8Einsum_0': {'input_amax_history': '*',\n 'input_scale': '*',\n 'kernel_amax_history': '*',\n 'kernel_scale': '*',\n 'output_grad_amax_history': '*',\n 'output_grad_scale': '*'}}}\n```\n\nIn addition to the expected `params`, there is an additional category called\n`_overwrite_with_gradient`. This category includes three pairs of `amax_history`\nand `scale` for the activation, kernel, and dot gradient, respectively.\n\n### Update gradient of FP8 params\nNow, we perform one training step to obtain the gradients and see how to use\nthem to update the parameters.\n\n```{code-cell}\nstep_fn = jax.jit(jax.grad(train_step, (0, 1)))\n\ngrads = step_fn(params, A)\n\nparams = flax.core.unfreeze(params)\nparams = flax.traverse_util.flatten_dict(params, sep='/')\ngrads = flax.traverse_util.flatten_dict(grads[0], sep='/')\n\nfor key, value in params.items():\n if key.startswith('params'):\n params[key] = value + 0.01 * grads[key]\n if key.startswith('_overwrite_with_gradient'):\n params[key] = grads[key]\n\nparams = flax.traverse_util.unflatten_dict(params, sep='/')\nparams = flax.core.freeze(params)\n```\n\nThe above code demonstrates how to update both `params` and\n`_overwrite_with_gradient`. For `params`, we use the formula `new_param =\nold_param + 0.01 * grads`, where `0.01` is the learning rate (or users can use\nwhatever optimizers from `optax`). For `_overwrite_with_gradient`, we simply use\nthe gradient to overwrite the old values.\n\nNote that `flax.training.train_state.TrainState` conveniently supports the\ncategory of `_overwrite_with_gradient`, so users do not need to modify their\nscripts if they don't use custom `TrainState`.\n\n## Accumulate gradient of FP8 params\nWhen the same parameter is used in a branched manner, the autograd mechanism\nwill add their gradients from these branches. This is common in scenarios like\npipeline parallelism, where each microbatch shares the same set of parameters\nfor the minibatch. However, for the `_overwrite_with_gradient` parameters, this\naccumulation by addition is not meaningful. Instead, we prefer custom\naccumulation by taking the maximum value.\n\nTo address this, we introduce a custom dtype `fp8_ops.fp32_max_grad`. The basic\nusage is demonstrated below:\n\n```{code-cell}\nfmax32 = fp8_ops.fp32_max_grad\n\ndef reuse_fp8_param(x, y, scale, amax_history):\n scale = scale.astype(fmax32)\n amax_history = amax_history.astype(fmax32)\n\n x = fp8_ops.in_qdq(f32, e4m3, x, scale, amax_history)\n y = fp8_ops.in_qdq(f32, e4m3, y, scale, amax_history)\n return x + y\n\nreuse_fp8_param_fn = jax.grad(reuse_fp8_param, (0, 1, 2, 3))\nreuse_fp8_param_fn = jax.jit(reuse_fp8_param_fn)\n\n_, _, new_ah, new_sf = reuse_fp8_param_fn(2.0, 3.0, a_scale, a_amax_hist)\nprint(new_ah, new_sf)\n```\n\nIn this example, we first cast the `scale` and `amax_history` to\n`fp8_ops.fp32_max_grad` and then call `fp8_ops.in_qdq` twice using the same pair\nof `scale` and `amax_history`. During autograd, their gradients from each branch\nwill be taken as the maximum, giving us the correct results of:\n\n```plaintext\n1.0 [3. 0. 0. ... 0. 0. 0.]\n```\n\nIf we do not perform the type casting, we get the following result, meaning the\ngradients of the two branches are added:\n\n```plaintext\n2.0 [5. 0. 0. ... 0. 0. 0.]\n```\n\nThis casting is already included if users choose to use the high-level APIs.\n\n## Deprecated APIs\nPreviously, we provided APIs like `fp8_ops.quantize_dequantize` for current\nscaling and `fp8_ops.[in|out]_qdq` for delayed scaling. These were used with\nhigh precision dot operations, leveraging an XLA-FP8 feature that\npattern-matched QDQ->dot sequences to Q->fp8_cublas_gemm. The corresponding\nhigh-level API was called `fp8_ops.Fp8DotGeneralOp`. However, this pattern\nmatching-based solution proved brittle, as the patterns could be easily broken\nby other XLA optimizations. We recommend users migrate from these deprecated\nAPIs to the newer ones described above.\n\nFor migration, users should replace:\n* `fp8_ops.quantize_dequantize -> jnp.dot` with `fp8_ops.quantize -> jnp.dot ->\n fp8_ops.dequantize`\n* `fp8_ops.in_qdq -> jnp.dot -> fp8_ops.out_qdq` with `fp8_ops.in_q -> jnp.dot\n -> fp8_ops.out_dq`\n* `fp8_ops.Fp8DotGeneralOp` with `fp8_ops.Fp8DotGeneral`\n\nAdditionally, we provide an einsum variant through `fp8_ops.Fp8Einsum`.\n" }, { "path": "docs/guides/quantization/index.rst", "content": "Quantization\n============\n\n.. toctree::\n :maxdepth: 1\n\n fp8_basics\n" }, { "path": "docs/guides/training_techniques/batch_norm.rst", "content": "Batch normalization\n===================\n\nIn this guide, you will learn how to apply `batch normalization `__\nusing :meth:`flax.linen.BatchNorm `.\n\nBatch normalization is a regularization technique used to speed up training and improve convergence.\nDuring training, it computes running averages over feature dimensions. This adds a new form\nof non-differentiable state that must be handled appropriately.\n\nThroughout the guide, you will be able to compare code examples with and without Flax ``BatchNorm``.\n\n.. testsetup:: No BatchNorm, With BatchNorm\n\n import flax.linen as nn\n import jax.numpy as jnp\n import jax\n import optax\n from typing import Any\n from flax.core import FrozenDict\n\nDefining the model with ``BatchNorm``\n*************************************\n\nIn Flax, ``BatchNorm`` is a :meth:`flax.linen.Module ` that exhibits different runtime\nbehavior between training and inference. You explicitly specify it via the ``use_running_average`` argument,\nas demonstrated below.\n\nA common pattern is to accept a ``train`` (``training``) argument in the parent Flax ``Module``, and use\nit to define ``BatchNorm``'s ``use_running_average`` argument.\n\nNote: In other machine learning frameworks, like PyTorch or\nTensorFlow (Keras), this is specified via a mutable state or a call flag (for example, in\n`torch.nn.Module.eval `__\nor ``tf.keras.Model`` by setting the\n`training `__ flag).\n\n.. codediff::\n :title: No BatchNorm, With BatchNorm\n :sync:\n\n class MLP(nn.Module):\n @nn.compact\n def __call__(self, x):\n x = nn.Dense(features=4)(x)\n\n x = nn.relu(x)\n x = nn.Dense(features=1)(x)\n return x\n\n ---\n class MLP(nn.Module):\n @nn.compact\n def __call__(self, x, train: bool): #!\n x = nn.Dense(features=4)(x)\n x = nn.BatchNorm(use_running_average=not train)(x) #!\n x = nn.relu(x)\n x = nn.Dense(features=1)(x)\n return x\n\nOnce you create your model, initialize it by calling :meth:`flax.linen.init() ` to\nget the ``variables`` structure. Here, the main difference between the code without ``BatchNorm``\nand with ``BatchNorm`` is that the ``train`` argument must be provided.\n\nThe ``batch_stats`` collection\n******************************\n\nIn addition to the ``params`` collection, ``BatchNorm`` also adds a ``batch_stats`` collection\nthat contains the running average of the batch statistics.\n\nNote: You can learn more in the ``flax.linen`` `variables `__\nAPI documentation.\n\nThe ``batch_stats`` collection must be extracted from the ``variables`` for later use.\n\n.. codediff::\n :title: No BatchNorm, With BatchNorm\n :sync:\n\n mlp = MLP()\n x = jnp.ones((1, 3))\n variables = mlp.init(jax.random.key(0), x)\n params = variables['params']\n\n\n jax.tree_util.tree_map(jnp.shape, variables)\n ---\n mlp = MLP()\n x = jnp.ones((1, 3))\n variables = mlp.init(jax.random.key(0), x, train=False) #!\n params = variables['params']\n batch_stats = variables['batch_stats'] #!\n\n jax.tree_util.tree_map(jnp.shape, variables)\n\n\nFlax ``BatchNorm`` adds a total of 4 variables: ``mean`` and ``var`` that live in the\n``batch_stats`` collection, and ``scale`` and ``bias`` that live in the ``params``\ncollection.\n\n.. codediff::\n :title: No BatchNorm, With BatchNorm\n :sync:\n\n FrozenDict({\n 'params': {\n 'Dense_0': {\n 'bias': (4,),\n 'kernel': (3, 4),\n },\n 'Dense_1': {\n 'bias': (1,),\n 'kernel': (4, 1),\n },\n },\n })\n ---\n FrozenDict({\n 'batch_stats': { #!\n 'BatchNorm_0': { #!\n 'mean': (4,), #!\n 'var': (4,), #!\n }, #!\n }, #!\n 'params': {\n 'BatchNorm_0': { #!\n 'bias': (4,), #!\n 'scale': (4,), #!\n }, #!\n 'Dense_0': {\n 'bias': (4,),\n 'kernel': (3, 4),\n },\n 'Dense_1': {\n 'bias': (1,),\n 'kernel': (4, 1),\n },\n },\n })\n\nModifying ``flax.linen.apply``\n******************************\n\nWhen using :meth:`flax.linen.apply ` to run your model with the ``train=True``\nargument (that is, you have ``use_running_average=False`` in the call to ``BatchNorm``), you\nneed to consider the following:\n\n* ``batch_stats`` must be passed as an input variable.\n* The ``batch_stats`` collection needs to be marked as mutable by setting ``mutable=['batch_stats']``.\n* The mutated variables are returned as a second output.\n The updated ``batch_stats`` must be extracted from here.\n\n.. codediff::\n :title: No BatchNorm, With BatchNorm\n :sync:\n\n y = mlp.apply(\n {'params': params},\n x,\n )\n ...\n\n ---\n y, updates = mlp.apply( #!\n {'params': params, 'batch_stats': batch_stats}, #!\n x,\n train=True, mutable=['batch_stats'] #!\n )\n batch_stats = updates['batch_stats'] #!\n\nTraining and evaluation\n***********************\n\nWhen integrating models that use ``BatchNorm`` into a training loop, the main challenge\nis handling the additional ``batch_stats`` state. To do this, you need to:\n\n* Add a ``batch_stats`` field to a custom :meth:`flax.training.train_state.TrainState ` class.\n* Pass the ``batch_stats`` values to the :meth:`train_state.TrainState.create ` method.\n\n.. codediff::\n :title: No BatchNorm, With BatchNorm\n :sync:\n\n from flax.training import train_state\n\n\n state = train_state.TrainState.create(\n apply_fn=mlp.apply,\n params=params,\n\n tx=optax.adam(1e-3),\n )\n ---\n from flax.training import train_state\n\n class TrainState(train_state.TrainState): #!\n batch_stats: Any #!\n\n state = TrainState.create( #!\n apply_fn=mlp.apply,\n params=params,\n batch_stats=batch_stats, #!\n tx=optax.adam(1e-3),\n )\n\nIn addition, update your ``train_step`` function to reflect these changes:\n\n* Pass all new parameters to ``flax.linen.apply`` (as previously discussed).\n* The ``updates`` to the ``batch_stats`` must be propagated out of the ``loss_fn``.\n* The ``batch_stats`` from the ``TrainState`` must be updated.\n\n.. codediff::\n :title: No BatchNorm, With BatchNorm\n :sync:\n\n @jax.jit\n def train_step(state: train_state.TrainState, batch):\n \"\"\"Train for a single step.\"\"\"\n def loss_fn(params):\n logits = state.apply_fn(\n {'params': params},\n x=batch['image'])\n loss = optax.softmax_cross_entropy_with_integer_labels(\n logits=logits, labels=batch['label']).mean()\n return loss, logits\n grad_fn = jax.value_and_grad(loss_fn, has_aux=True)\n (loss, logits), grads = grad_fn(state.params)\n state = state.apply_gradients(grads=grads)\n\n metrics = {\n 'loss': loss,\n 'accuracy': jnp.mean(jnp.argmax(logits, -1) == batch['label']),\n }\n return state, metrics\n ---\n @jax.jit\n def train_step(state: TrainState, batch):\n \"\"\"Train for a single step.\"\"\"\n def loss_fn(params):\n logits, updates = state.apply_fn( #!\n {'params': params, 'batch_stats': state.batch_stats}, #!\n x=batch['image'], train=True, mutable=['batch_stats']) #!\n loss = optax.softmax_cross_entropy_with_integer_labels(\n logits=logits, labels=batch['label']).mean()\n return loss, (logits, updates) #!\n grad_fn = jax.value_and_grad(loss_fn, has_aux=True)\n (loss, (logits, updates)), grads = grad_fn(state.params) #!\n state = state.apply_gradients(grads=grads)\n state = state.replace(batch_stats=updates['batch_stats']) #!\n metrics = {\n 'loss': loss,\n 'accuracy': jnp.mean(jnp.argmax(logits, -1) == batch['label']),\n }\n return state, metrics\n\nThe ``eval_step`` is much simpler. Because ``batch_stats`` is not mutable, no\nupdates\nneed to be propagated. Make sure you pass the ``batch_stats`` to ``flax.linen.apply``,\nand the ``train`` argument is set to ``False``:\n\n.. codediff::\n :title: No BatchNorm, With BatchNorm\n :sync:\n\n @jax.jit\n def eval_step(state: train_state.TrainState, batch):\n \"\"\"Train for a single step.\"\"\"\n logits = state.apply_fn(\n {'params': params},\n x=batch['image'])\n loss = optax.softmax_cross_entropy_with_integer_labels(\n logits=logits, labels=batch['label']).mean()\n metrics = {\n 'loss': loss,\n 'accuracy': jnp.mean(jnp.argmax(logits, -1) == batch['label']),\n }\n return state, metrics\n ---\n @jax.jit\n def eval_step(state: TrainState, batch):\n \"\"\"Evaluate for a single step.\"\"\"\n logits = state.apply_fn(\n {'params': state.params, 'batch_stats': state.batch_stats}, #!\n x=batch['image'], train=False) #!\n loss = optax.softmax_cross_entropy_with_integer_labels(\n logits=logits, labels=batch['label']).mean()\n metrics = {\n 'loss': loss,\n 'accuracy': jnp.mean(jnp.argmax(logits, -1) == batch['label']),\n }\n return state, metrics\n" }, { "path": "docs/guides/training_techniques/dropout.rst", "content": "Dropout\n=======\n\nThis guide provides an overview of how to apply\n`dropout `__\nusing :meth:`flax.linen.Dropout`.\n\nDropout is a stochastic regularization technique that randomly removes hidden\nand visible units in a network.\n\nThroughout the guide, you will be able to compare code examples with and without\nFlax ``Dropout``.\n\n.. testsetup:: No Dropout, With Dropout\n\n import flax.linen as nn\n import jax.numpy as jnp\n import jax\n import optax\n\nSplit the PRNG key\n******************\n\nSince dropout is a random operation, it requires a pseudorandom number generator\n(PRNG) state. Flax uses JAX's (splittable) PRNG keys, which have a number of\ndesirable properties for neural networks. To learn more, refer to the\n`Pseudorandom numbers in JAX tutorial `__.\n\n**Note:** Recall that JAX has an explicit way of giving you PRNG keys:\nyou can fork the main PRNG state (such as ``key = jax.random.key(seed=0)``)\ninto multiple new PRNG keys with ``key, subkey = jax.random.split(key)``. You\ncan refresh your memory in\n`🔪 JAX - The Sharp Bits 🔪 Randomness and PRNG keys `__.\n\nBegin by splitting the PRNG key using\n`jax.random.split() `__\ninto three keys, including one for Flax Linen ``Dropout``.\n\n.. codediff::\n :title: No Dropout, With Dropout\n :sync:\n\n root_key = jax.random.key(seed=0)\n main_key, params_key = jax.random.split(key=root_key)\n ---\n root_key = jax.random.key(seed=0)\n main_key, params_key, dropout_key = jax.random.split(key=root_key, num=3) #!\n\n**Note:** In Flax, you provide *PRNG streams* with *names*, so that you can use them later\nin your :meth:`flax.linen.Module`. For example, you pass the stream ``'params'``\nfor initializing parameters, and ``'dropout'`` for applying\n:meth:`flax.linen.Dropout`.\n\nDefine your model with ``Dropout``\n**********************************\n\nTo create a model with dropout:\n\n* Subclass :meth:`flax.linen.Module`, and then use\n :meth:`flax.linen.Dropout` to add a dropout layer. Recall that\n :meth:`flax.linen.Module` is the\n `base class for all neural network Modules `__,\n and all layers and models are subclassed from it.\n\n* In :meth:`flax.linen.Dropout`, the ``deterministic`` argument is required to\n be passed as a keyword argument, either:\n\n * When constructing the :meth:`flax.linen.Module`; or\n * When calling :meth:`flax.linen.init()` or :meth:`flax.linen.apply()` on a constructed ``Module``. (Refer to :meth:`flax.linen.module.merge_param` for more details.)\n\n* Because ``deterministic`` is a boolean:\n\n * If it's set to ``False``, the inputs are masked (that is, set to zero) with\n a probability set by ``rate``. And the remaining inputs are scaled by\n ``1 / (1 - rate)``, which ensures that the means of the inputs are\n preserved.\n * If it's set to ``True``, no mask is applied (the dropout is turned off),\n and the inputs are returned as-is.\n\nA common pattern is to accept a ``training`` (or ``train``) argument (a boolean)\nin the parent Flax ``Module``, and use it to enable or disable dropout (as\ndemonstrated in later sections of this guide). In other machine learning\nframeworks, like PyTorch or TensorFlow (Keras), this is specified via a\nmutable state or a call flag (for example, in\n`torch.nn.Module.eval `__\nor ``tf.keras.Model`` by setting the\n`training `__ flag).\n\n**Note:** Flax provides an implicit way of handling PRNG key streams via Flax\n:meth:`flax.linen.Module`'s :meth:`flax.linen.Module.make_rng` method.\nThis allows you to split off a fresh PRNG key inside Flax Modules (or their\nsub-Modules) from the PRNG stream. The ``make_rng`` method guarantees to provide a\nunique key each time you call it. Internally, :meth:`flax.linen.Dropout` makes\nuse of :meth:`flax.linen.Module.make_rng` to create a key for dropout. You can\ncheck out the\n`source code `__.\nIn short, :meth:`flax.linen.Module.make_rng` *guarantees full reproducibility*.\n\n.. codediff::\n :title: No Dropout, With Dropout\n :sync:\n\n class MyModel(nn.Module):\n num_neurons: int\n\n @nn.compact\n def __call__(self, x):\n x = nn.Dense(self.num_neurons)(x)\n\n return x\n ---\n class MyModel(nn.Module):\n num_neurons: int\n\n @nn.compact\n def __call__(self, x, training: bool): #!\n x = nn.Dense(self.num_neurons)(x)\n # Set the dropout layer with a `rate` of 50%. #!\n # When the `deterministic` flag is `True`, dropout is turned off. #!\n x = nn.Dropout(rate=0.5, deterministic=not training)(x) #!\n return x\n\nInitialize the model\n********************\n\nAfter creating your model:\n\n* Instantiate the model.\n* Then, in the :meth:`flax.linen.init()` call, set ``training=False``.\n* Finally, extract the ``params`` from the\n `variable dictionary `__.\n\nHere, the main difference between the code without Flax ``Dropout``\nand with ``Dropout`` is that the ``training`` (or ``train``) argument must be\nprovided if you need dropout enabled.\n\n.. codediff::\n :title: No Dropout, With Dropout\n :sync:\n\n my_model = MyModel(num_neurons=3)\n x = jnp.empty((3, 4, 4))\n\n variables = my_model.init(params_key, x)\n params = variables['params']\n ---\n my_model = MyModel(num_neurons=3)\n x = jnp.empty((3, 4, 4))\n # Dropout is disabled with `training=False` (that is, `deterministic=True`). #!\n variables = my_model.init(params_key, x, training=False) #!\n params = variables['params']\n\nPerform the forward pass during training\n****************************************\n\nWhen using :meth:`flax.linen.apply()` to run your model:\n\n* Pass ``training=True`` to :meth:`flax.linen.apply()`.\n* Then, to draw PRNG keys during the forward pass (with dropout), provide a PRNG key\n to seed the ``'dropout'`` stream when you call :meth:`flax.linen.apply()`.\n\n.. codediff::\n :title: No Dropout, With Dropout\n :sync:\n\n # No need to pass the `training` and `rngs` flags.\n y = my_model.apply({'params': params}, x)\n ---\n # Dropout is enabled with `training=True` (that is, `deterministic=False`). #!\n y = my_model.apply({'params': params}, x, training=True, rngs={'dropout': dropout_key}) #!\n\nHere, the main difference between the code without Flax ``Dropout``\nand with ``Dropout`` is that the ``training`` (or ``train``) and ``rngs``\narguments must be provided if you need dropout enabled.\n\nDuring evaluation, use the above code with no dropout enabled (this means you do\nnot have to pass a RNG either).\n\n``TrainState`` and the training step\n************************************\n\nThis section explains how to amend your code inside the training step function if\nyou have dropout enabled.\n\n**Note:** Recall that Flax has a common pattern where you create a dataclass\nthat represents the whole training state, including parameters and the optimizer\nstate. Then, you can pass a single parameter, ``state: TrainState``, to\nthe training step function. Refer to the\n:meth:`flax.training.train_state.TrainState` API docs to learn more.\n\n* First, add a ``key`` field to a custom :meth:`flax.training.train_state.TrainState` class.\n* Then, pass the ``key`` value—in this case, the ``dropout_key``—to the :meth:`train_state.TrainState.create` method.\n\n.. codediff::\n :title: No Dropout, With Dropout\n :sync:\n\n from flax.training import train_state\n\n state = train_state.TrainState.create(\n apply_fn=my_model.apply,\n params=params,\n\n tx=optax.adam(1e-3)\n )\n ---\n from flax.training import train_state\n\n class TrainState(train_state.TrainState): #!\n key: jax.Array #!\n\n state = TrainState.create( #!\n apply_fn=my_model.apply,\n params=params,\n key=dropout_key, #!\n tx=optax.adam(1e-3)\n )\n\n* Next, in the Flax training step function, ``train_step``, generate a new PRNG\n key from the ``dropout_key`` to apply dropout at each step. This can be done with one of the following:\n\n * `jax.random.split() `__; or\n * `jax.random.fold_in() `__\n\n Using ``jax.random.fold_in()`` is generally faster. When you use\n ``jax.random.split()`` you split off a PRNG key that can be reused\n afterwards. However, using ``jax.random.fold_in()`` makes sure to 1) fold in\n unique data; and 2) can result in longer sequences of PRNG streams.\n\n* Finally, when performing the forward pass, pass the new PRNG key to ``state.apply_fn()``\n as an extra parameter.\n\n.. codediff::\n :title: No Dropout, With Dropout\n :sync:\n\n @jax.jit\n def train_step(state: train_state.TrainState, batch):\n\n def loss_fn(params):\n logits = state.apply_fn(\n {'params': params},\n x=batch['image'],\n\n\n )\n loss = optax.softmax_cross_entropy_with_integer_labels(\n logits=logits, labels=batch['label'])\n return loss, logits\n grad_fn = jax.value_and_grad(loss_fn, has_aux=True)\n (loss, logits), grads = grad_fn(state.params)\n state = state.apply_gradients(grads=grads)\n return state\n\n ---\n @jax.jit\n def train_step(state: TrainState, batch, dropout_key): #!\n dropout_train_key = jax.random.fold_in(key=dropout_key, data=state.step) #!\n def loss_fn(params):\n logits = state.apply_fn(\n {'params': params},\n x=batch['image'],\n training=True, #!\n rngs={'dropout': dropout_train_key} #!\n )\n loss = optax.softmax_cross_entropy_with_integer_labels(\n logits=logits, labels=batch['label'])\n return loss, logits\n grad_fn = jax.value_and_grad(loss_fn, has_aux=True)\n (loss, logits), grads = grad_fn(state.params)\n state = state.apply_gradients(grads=grads)\n return state\n\nFlax examples with dropout\n**************************\n\n* A `Transformer-based model `__\n trained on the WMT Machine Translation dataset. This example uses dropout and attention dropout.\n\n* Applying word dropout to a batch of input IDs in a\n `text classification `__\n context. This example uses a custom :meth:`flax.linen.Dropout` layer.\n\nMore Flax examples that use Module ``make_rng()``\n*************************************************\n\n* Defining a prediction token in a decoder of a\n `sequence-to-sequence model `__.\n" }, { "path": "docs/guides/training_techniques/index.rst", "content": "Training techniques\n===================\n\n.. toctree::\n :maxdepth: 1\n\n batch_norm\n dropout\n lr_schedule\n transfer_learning\n use_checkpointing" }, { "path": "docs/guides/training_techniques/lr_schedule.rst", "content": "Learning rate scheduling\n=============================\n\nThe learning rate is considered one of the most important hyperparameters for\ntraining deep neural networks, but choosing it can be quite hard.\nRather than simply using a fixed learning rate, it is common to use a learning rate scheduler.\nIn this example, we will use the *cosine scheduler*.\nBefore the cosine scheduler comes into play, we start with a so-called *warmup* period in which the\nlearning rate increases linearly for ``warmup_epochs`` epochs.\nFor more information about the cosine scheduler, check out the paper\n`\"SGDR: Stochastic Gradient Descent with Warm Restarts\" `_.\n\nWe will show you how to...\n\n* define a learning rate schedule\n* train a simple model using that schedule\n\n\n.. testsetup:: Default learning rate, Learning rate schedule\n\n import jax\n import jax.numpy as jnp\n import flax.linen as nn\n from flax.training import train_state\n import optax\n import numpy as np\n import tensorflow_datasets as tfds\n import functools\n import ml_collections\n from absl import logging\n\n\n class CNN(nn.Module):\n \"\"\"A simple CNN model.\"\"\"\n\n @nn.compact\n def __call__(self, x):\n x = nn.Conv(features=32, kernel_size=(3, 3))(x)\n x = nn.relu(x)\n x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))\n x = nn.Conv(features=64, kernel_size=(3, 3))(x)\n x = nn.relu(x)\n x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))\n x = x.reshape((x.shape[0], -1)) # flatten\n x = nn.Dense(features=256)(x)\n x = nn.relu(x)\n x = nn.Dense(features=10)(x)\n return x\n\n def get_dummy_data(ds_size):\n image = np.random.rand(ds_size, 28, 28, 1)\n label = np.random.randint(low=0, high=10, size=(ds_size,))\n return {'image': image, 'label': label}\n\n def get_config():\n \"\"\"Get the default hyperparameter configuration.\"\"\"\n config = ml_collections.ConfigDict()\n\n config.learning_rate = 0.001\n config.momentum = 0.9\n config.batch_size = 128\n config.num_epochs = 10\n config.warmup_epochs = 2\n config.train_ds_size = 128\n return config\n\n def compute_metrics(logits, labels):\n one_hot = jax.nn.one_hot(labels, 10)\n loss = jnp.mean(optax.softmax_cross_entropy(logits, one_hot))\n accuracy = jnp.mean(jnp.argmax(logits, -1) == labels)\n metrics = {\n 'loss': loss,\n 'accuracy': accuracy,\n }\n return metrics\n\n\n.. testcode:: Default learning rate, Learning rate schedule\n\n def create_learning_rate_fn(config, base_learning_rate, steps_per_epoch):\n \"\"\"Creates learning rate schedule.\"\"\"\n warmup_fn = optax.linear_schedule(\n init_value=0., end_value=base_learning_rate,\n transition_steps=config.warmup_epochs * steps_per_epoch)\n cosine_epochs = max(config.num_epochs - config.warmup_epochs, 1)\n cosine_fn = optax.cosine_decay_schedule(\n init_value=base_learning_rate,\n decay_steps=cosine_epochs * steps_per_epoch)\n schedule_fn = optax.join_schedules(\n schedules=[warmup_fn, cosine_fn],\n boundaries=[config.warmup_epochs * steps_per_epoch])\n return schedule_fn\n\nTo use the schedule, we must create a learning rate function by passing the hyperparameters to the\n``create_learning_rate_fn`` function and then pass the function to your |Optax|_ optimizer.\nFor example using this schedule on MNIST would require changing the ``train_step`` function:\n\n.. |Optax| replace:: ``Optax``\n.. _Optax: https://optax.readthedocs.io/en/latest/api.html#optimizer-schedules\n\n.. codediff::\n :title: Default learning rate, Learning rate schedule\n :sync:\n\n @jax.jit\n def train_step(state, batch):\n def loss_fn(params):\n logits = CNN().apply({'params': params}, batch['image'])\n one_hot = jax.nn.one_hot(batch['label'], 10)\n loss = jnp.mean(optax.softmax_cross_entropy(logits, one_hot))\n return loss, logits\n grad_fn = jax.value_and_grad(loss_fn, has_aux=True)\n (_, logits), grads = grad_fn(state.params)\n new_state = state.apply_gradients(grads=grads)\n metrics = compute_metrics(logits, batch['label'])\n\n\n return new_state, metrics\n ---\n @functools.partial(jax.jit, static_argnums=2) #!\n def train_step(state, batch, learning_rate_fn): #!\n def loss_fn(params):\n logits = CNN().apply({'params': params}, batch['image'])\n one_hot = jax.nn.one_hot(batch['label'], 10)\n loss = jnp.mean(optax.softmax_cross_entropy(logits, one_hot))\n return loss, logits\n grad_fn = jax.value_and_grad(loss_fn, has_aux=True)\n (_, logits), grads = grad_fn(state.params)\n new_state = state.apply_gradients(grads=grads)\n metrics = compute_metrics(logits, batch['label'])\n lr = learning_rate_fn(state.step) #!\n metrics['learning_rate'] = lr #!\n return new_state, metrics\n\nAnd the ``train_epoch`` function:\n\n.. codediff::\n :title: Default learning rate, Learning rate schedule\n :sync:\n\n def train_epoch(state, train_ds, batch_size, epoch, rng):\n \"\"\"Trains for a single epoch.\"\"\"\n train_ds_size = len(train_ds['image'])\n steps_per_epoch = train_ds_size // batch_size\n perms = jax.random.permutation(rng, len(train_ds['image']))\n perms = perms[:steps_per_epoch * batch_size]\n perms = perms.reshape((steps_per_epoch, batch_size))\n batch_metrics = []\n for perm in perms:\n batch = {k: v[perm, ...] for k, v in train_ds.items()}\n state, metrics = train_step(state, batch)\n batch_metrics.append(metrics)\n\n # compute mean of metrics across each batch in epoch.\n batch_metrics = jax.device_get(batch_metrics)\n epoch_metrics = {\n k: np.mean([metrics[k] for metrics in batch_metrics])\n for k in batch_metrics[0]}\n\n logging.info('train epoch: %d, loss: %.4f, accuracy: %.2f', epoch,\n epoch_metrics['loss'], epoch_metrics['accuracy'] * 100)\n\n return state, epoch_metrics\n ---\n def train_epoch(state, train_ds, batch_size, epoch, learning_rate_fn, rng): #!\n \"\"\"Trains for a single epoch.\"\"\"\n train_ds_size = len(train_ds['image'])\n steps_per_epoch = train_ds_size // batch_size\n perms = jax.random.permutation(rng, len(train_ds['image']))\n perms = perms[:steps_per_epoch * batch_size]\n perms = perms.reshape((steps_per_epoch, batch_size))\n batch_metrics = []\n for perm in perms:\n batch = {k: v[perm, ...] for k, v in train_ds.items()}\n state, metrics = train_step(state, batch, learning_rate_fn) #!\n batch_metrics.append(metrics)\n\n # compute mean of metrics across each batch in epoch.\n batch_metrics = jax.device_get(batch_metrics)\n epoch_metrics = {\n k: np.mean([metrics[k] for metrics in batch_metrics])\n for k in batch_metrics[0]}\n\n logging.info('train epoch: %d, loss: %.4f, accuracy: %.2f', epoch,\n epoch_metrics['loss'], epoch_metrics['accuracy'] * 100)\n\n return state, epoch_metrics\n\n\nAnd the ``create_train_state`` function:\n\n\n.. codediff::\n :title: Default learning rate, Learning rate schedule\n :sync:\n\n def create_train_state(rng, config):\n \"\"\"Creates initial `TrainState`.\"\"\"\n cnn = CNN()\n params = cnn.init(rng, jnp.ones([1, 28, 28, 1]))['params']\n tx = optax.sgd(config.learning_rate, config.momentum)\n return train_state.TrainState.create(\n apply_fn=cnn.apply, params=params, tx=tx)\n ---\n def create_train_state(rng, config, learning_rate_fn): #!\n \"\"\"Creates initial `TrainState`.\"\"\"\n cnn = CNN()\n params = cnn.init(rng, jnp.ones([1, 28, 28, 1]))['params']\n tx = optax.sgd(learning_rate_fn, config.momentum) #!\n return train_state.TrainState.create(\n apply_fn=cnn.apply, params=params, tx=tx)\n\n\n.. testcleanup:: Learning rate schedule\n\n config = get_config()\n\n train_ds_size = config.train_ds_size\n steps_per_epoch = train_ds_size // config.batch_size\n learning_rate_fn = create_learning_rate_fn(config, config.learning_rate, steps_per_epoch)\n\n rng = jax.random.key(0)\n state = create_train_state(rng, config, learning_rate_fn)\n\n train_ds = get_dummy_data(config.train_ds_size)\n rng, _ = jax.random.split(rng)\n state, epoch_metrics = train_epoch(state, train_ds, config.batch_size, 0, learning_rate_fn, rng)\n\n assert 'accuracy' in epoch_metrics and 'learning_rate' in epoch_metrics\n\n\n\n" }, { "path": "docs/guides/training_techniques/transfer_learning.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"# Transfer learning\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"This guide demonstrates various parts of the transfer learning workflow with Flax. Depending on the task, a pretrained model can be used just as a feature extractor or it can be fine-tuned as part of a larger model.\\n\",\n \"\\n\",\n \"This guide demonstrates how to:\\n\",\n \"\\n\",\n \"* Load a pretrained model from HuggingFace [Transformers](https://huggingface.co/docs/transformers/index) and extract a specific sub-module from that pretrained model.\\n\",\n \"* Create a classifier model.\\n\",\n \"* Transfer the pretrained parameters to the new model structure.\\n\",\n \"* Create an optimizer for training different parts of the model separately with [Optax](https://optax.readthedocs.io/).\\n\",\n \"* Set up the model for training.\\n\",\n \"\\n\",\n \"
Performance Note\\n\",\n \"\\n\",\n \"Depending on your task, some of the content in this guide may be suboptimal. For example, if you are only going to train a linear classifier on top of a pretrained model, it may be better to just extract the feature embeddings once, which can result in much faster training, and you can use specialized algorithms for linear regression or logistic classification. This guide shows how to do transfer learning with all the model parameters.\\n\",\n \"\\n\",\n \"

\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Setup\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {\n \"tags\": [\n \"skip-execution\"\n ]\n },\n \"outputs\": [],\n \"source\": [\n \"# Note that the Transformers library doesn't use the latest Flax version.\\n\",\n \"! pip install -q \\\"transformers[flax]\\\"\\n\",\n \"# Install/upgrade Flax and JAX. For JAX installation with GPU/TPU support,\\n\",\n \"# visit https://github.com/jax-ml/jax#installation.\\n\",\n \"! pip install -U -q flax jax jaxlib\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Create a function for model loading\\n\",\n \"\\n\",\n \"To load a pre-trained classifier, for convenience first create a function that returns a [Flax `Module`](https://flax.readthedocs.io/en/latest/guides/flax_fundamentals/flax_basics.html#module-basics) and its pretrained variables.\\n\",\n \"\\n\",\n \"In the code below, the `load_model` function uses HuggingFace's `FlaxCLIPVisionModel` model from the [Transformers](https://huggingface.co/docs/transformers/index) library and extracts a `FlaxCLIPModule` module.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 3,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"%%capture\\n\",\n \"from IPython.display import clear_output\\n\",\n \"from transformers import FlaxCLIPModel\\n\",\n \"\\n\",\n \"# Note: FlaxCLIPModel is not a Flax Module\\n\",\n \"def load_model():\\n\",\n \" clip = FlaxCLIPModel.from_pretrained('openai/clip-vit-base-patch32')\\n\",\n \" clear_output(wait=False) # Clear the loading messages\\n\",\n \" module = clip.module # Extract the Flax Module\\n\",\n \" variables = {'params': clip.params} # Extract the parameters\\n\",\n \" return module, variables\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Note that `FlaxCLIPVisionModel` itself is not a Flax `Module` which is why we need to do this extra step.\\n\",\n \"\\n\",\n \"### Extracting a submodule\\n\",\n \"\\n\",\n \"Calling `load_model` from the snippet above returns the `FlaxCLIPModule`, which is composed of `text_model` and `vision_model` submodules.\\n\",\n \"\\n\",\n \"An easy way to extract the `vision_model` sub-Module defined inside `.setup()` and its variables is to use [`flax.linen.Module.bind`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.bind) on the `clip` Module immediately followed by [`flax.linen.Module.unbind`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.unbind) on the `vision_model` sub-Module.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 8,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"import flax.linen as nn\\n\",\n \"\\n\",\n \"clip, clip_variables = load_model()\\n\",\n \"vision_model, vision_model_vars = clip.bind(clip_variables).vision_model.unbind()\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Creating a classifier\\n\",\n \"\\n\",\n \"To create a classifier define a new Flax [`Module`](https://flax.readthedocs.io/en/latest/guides/flax_fundamentals/flax_basics.html#module-basics) consisting of a `backbone` (the pretrained vision model) and a `head` (the classifier) submodules.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 7,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"from typing import Callable\\n\",\n \"import jax.numpy as jnp\\n\",\n \"import jax\\n\",\n \"\\n\",\n \"class Classifier(nn.Module):\\n\",\n \" num_classes: int\\n\",\n \" backbone: nn.Module\\n\",\n \" \\n\",\n \"\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self, x):\\n\",\n \" x = self.backbone(x).pooler_output\\n\",\n \" x = nn.Dense(\\n\",\n \" self.num_classes, name='head', kernel_init=nn.zeros)(x)\\n\",\n \" return x\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"To construct a classifier `model`, the `vision_model` Module is passed as the `backbone` to `Classifier`. Then the model's `params` can be randomly initialized by passing fake data that is used to infer the parameter shapes.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 8,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"num_classes = 3\\n\",\n \"model = Classifier(num_classes=num_classes, backbone=vision_model)\\n\",\n \"\\n\",\n \"x = jnp.empty((1, 224, 224, 3))\\n\",\n \"variables = model.init(jax.random.key(1), x)\\n\",\n \"params = variables['params']\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Transfering the parameters\\n\",\n \"Since `params` are currently random, the pretrained parameters from `vision_model_vars` have to be transfered to the `params` structure at the appropriate location (i.e. the `backbone`):\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 9,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"params['backbone'] = vision_model_vars['params']\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"**Note:** if the model contains other variable collections such as `batch_stats`, these have to be transfered as well.\\n\",\n \"\\n\",\n \"## Optimization\\n\",\n \"\\n\",\n \"If you need to to train different parts of the model separately, you have three options:\\n\",\n \"\\n\",\n \"1. Use `stop_gradient`.\\n\",\n \"2. Filter the parameters for `jax.grad`.\\n\",\n \"3. Use multiple optimizers for different parameters.\\n\",\n \"\\n\",\n \"For most situations we recommend using multiple optimizers via [Optax](https://optax.readthedocs.io/)'s [`multi_transform`](https://optax.readthedocs.io/en/latest/api.html#optax.multi_transform) as its both efficient and can be easily extended to implement many fine-tunning strategies. \\n\",\n \"\\n\",\n \"### **optax.multi_transform**\\n\",\n \"\\n\",\n \"To use `optax.multi_transform` following must be defined:\\n\",\n \"\\n\",\n \"1. The parameter partitions.\\n\",\n \"2. A mapping between partitions and their optimizer.\\n\",\n \"3. A pytree with the same shape as the parameters but its leaves containing the corresponding partition label.\\n\",\n \"\\n\",\n \"To freeze layers with `optax.multi_transform` for the model above, the following setup can be used:\\n\",\n \"\\n\",\n \"* Define the `trainable` and `frozen` parameter partitions.\\n\",\n \"* For the `trainable` parameters select the Adam (`optax.adam`) optimizer.\\n\",\n \"- For the `frozen` parameters select the `optax.set_to_zero` optimizer. This dummy optimizer zeros-out the gradients so no training is done.\\n\",\n \"- Map parameters to partitions using [`flax.traverse_util.path_aware_map`](https://flax.readthedocs.io/en/latest/api_reference/flax.traverse_util.html#flax.traverse_util.path_aware_map), mark the leaves from the `backbone` as `frozen`, and the rest as `trainable`.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 10,\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"FrozenDict({\\n\",\n \" backbone: {\\n\",\n \" embeddings: {\\n\",\n \" class_embedding: 'frozen',\\n\",\n \" patch_embedding: {\\n\",\n \" kernel: 'frozen',\\n\",\n \" },\\n\",\n \" },\\n\",\n \" },\\n\",\n \" head: {\\n\",\n \" bias: 'trainable',\\n\",\n \" kernel: 'trainable',\\n\",\n \" },\\n\",\n \"})\"\n ]\n },\n \"execution_count\": 10,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"from flax import traverse_util\\n\",\n \"import optax\\n\",\n \"\\n\",\n \"partition_optimizers = {'trainable': optax.adam(5e-3), 'frozen': optax.set_to_zero()}\\n\",\n \"param_partitions = traverse_util.path_aware_map(\\n\",\n \" lambda path, v: 'frozen' if 'backbone' in path else 'trainable', params)\\n\",\n \"tx = optax.multi_transform(partition_optimizers, param_partitions)\\n\",\n \"\\n\",\n \"# visualize a subset of the param_partitions structure\\n\",\n \"flat = list(traverse_util.flatten_dict(param_partitions).items())\\n\",\n \"traverse_util.unflatten_dict(dict(flat[:2] + flat[-2:]))\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"To implement [differential learning rates](https://blog.slavv.com/differential-learning-rates-59eff5209a4f), the `optax.set_to_zero` can be replaced with any other optimizer, different optimizers and partitioning schemes can be selected depending on the task. For more information on advanced optimizers, refer to Optax's [Combining Optimizers](https://optax.readthedocs.io/en/latest/api.html#combining-optimizers) documentation.\\n\",\n \"\\n\",\n \"## Creating the `TrainState`\\n\",\n \"\\n\",\n \"Once the module, params, and optimizer are defined, the `TrainState` can be constructed as usual:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 12,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"from flax.training.train_state import TrainState\\n\",\n \"\\n\",\n \"state = TrainState.create(\\n\",\n \" apply_fn=model.apply,\\n\",\n \" params=params,\\n\",\n \" tx=tx)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"083d8854\",\n \"metadata\": {},\n \"source\": [\n \"Since the optimizer takes care of the freezing or fine-tunning strategy, the `train_step` requires no additional changes, training can proceed normally.\"\n ]\n }\n ],\n \"metadata\": {\n \"jupytext\": {\n \"formats\": \"ipynb,md:myst\"\n },\n \"language_info\": {\n \"codemirror_mode\": {\n \"name\": \"ipython\",\n \"version\": 3\n },\n \"file_extension\": \".py\",\n \"mimetype\": \"text/x-python\",\n \"name\": \"python\",\n \"nbconvert_exporter\": \"python\",\n \"pygments_lexer\": \"ipython3\",\n \"version\": \"3.9.14\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 2\n}\n" }, { "path": "docs/guides/training_techniques/transfer_learning.md", "content": "---\njupytext:\n formats: ipynb,md:myst\n text_representation:\n extension: .md\n format_name: myst\n format_version: 0.13\n jupytext_version: 1.13.8\n---\n\n# Transfer learning\n\n+++\n\nThis guide demonstrates various parts of the transfer learning workflow with Flax. Depending on the task, a pretrained model can be used just as a feature extractor or it can be fine-tuned as part of a larger model.\n\nThis guide demonstrates how to:\n\n* Load a pretrained model from HuggingFace [Transformers](https://huggingface.co/docs/transformers/index) and extract a specific sub-module from that pretrained model.\n* Create a classifier model.\n* Transfer the pretrained parameters to the new model structure.\n* Create an optimizer for training different parts of the model separately with [Optax](https://optax.readthedocs.io/).\n* Set up the model for training.\n\n
Performance Note\n\nDepending on your task, some of the content in this guide may be suboptimal. For example, if you are only going to train a linear classifier on top of a pretrained model, it may be better to just extract the feature embeddings once, which can result in much faster training, and you can use specialized algorithms for linear regression or logistic classification. This guide shows how to do transfer learning with all the model parameters.\n\n

\n\n+++\n\n## Setup\n\n```{code-cell} ipython3\n:tags: [skip-execution]\n\n# Note that the Transformers library doesn't use the latest Flax version.\n! pip install -q \"transformers[flax]\"\n# Install/upgrade Flax and JAX. For JAX installation with GPU/TPU support,\n# visit https://github.com/jax-ml/jax#installation.\n! pip install -U -q flax jax jaxlib\n```\n\n## Create a function for model loading\n\nTo load a pre-trained classifier, for convenience first create a function that returns a [Flax `Module`](https://flax.readthedocs.io/en/latest/guides/flax_fundamentals/flax_basics.html#module-basics) and its pretrained variables.\n\nIn the code below, the `load_model` function uses HuggingFace's `FlaxCLIPVisionModel` model from the [Transformers](https://huggingface.co/docs/transformers/index) library and extracts a `FlaxCLIPModule` module.\n\n```{code-cell} ipython3\n%%capture\nfrom IPython.display import clear_output\nfrom transformers import FlaxCLIPModel\n\n# Note: FlaxCLIPModel is not a Flax Module\ndef load_model():\n clip = FlaxCLIPModel.from_pretrained('openai/clip-vit-base-patch32')\n clear_output(wait=False) # Clear the loading messages\n module = clip.module # Extract the Flax Module\n variables = {'params': clip.params} # Extract the parameters\n return module, variables\n```\n\nNote that `FlaxCLIPVisionModel` itself is not a Flax `Module` which is why we need to do this extra step.\n\n### Extracting a submodule\n\nCalling `load_model` from the snippet above returns the `FlaxCLIPModule`, which is composed of `text_model` and `vision_model` submodules.\n\nAn easy way to extract the `vision_model` sub-Module defined inside `.setup()` and its variables is to use [`flax.linen.Module.bind`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.bind) on the `clip` Module immediately followed by [`flax.linen.Module.unbind`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.unbind) on the `vision_model` sub-Module.\n\n```{code-cell} ipython3\nimport flax.linen as nn\n\nclip, clip_variables = load_model()\nvision_model, vision_model_vars = clip.bind(clip_variables).vision_model.unbind()\n```\n\n### Creating a classifier\n\nTo create a classifier define a new Flax [`Module`](https://flax.readthedocs.io/en/latest/guides/flax_fundamentals/flax_basics.html#module-basics) consisting of a `backbone` (the pretrained vision model) and a `head` (the classifier) submodules.\n\n```{code-cell} ipython3\nfrom typing import Callable\nimport jax.numpy as jnp\nimport jax\n\nclass Classifier(nn.Module):\n num_classes: int\n backbone: nn.Module\n \n\n @nn.compact\n def __call__(self, x):\n x = self.backbone(x).pooler_output\n x = nn.Dense(\n self.num_classes, name='head', kernel_init=nn.zeros)(x)\n return x\n```\n\nTo construct a classifier `model`, the `vision_model` Module is passed as the `backbone` to `Classifier`. Then the model's `params` can be randomly initialized by passing fake data that is used to infer the parameter shapes.\n\n```{code-cell} ipython3\nnum_classes = 3\nmodel = Classifier(num_classes=num_classes, backbone=vision_model)\n\nx = jnp.empty((1, 224, 224, 3))\nvariables = model.init(jax.random.key(1), x)\nparams = variables['params']\n```\n\n## Transfering the parameters\nSince `params` are currently random, the pretrained parameters from `vision_model_vars` have to be transfered to the `params` structure at the appropriate location (i.e. the `backbone`):\n\n```{code-cell} ipython3\nparams['backbone'] = vision_model_vars['params']\n```\n\n**Note:** if the model contains other variable collections such as `batch_stats`, these have to be transfered as well.\n\n## Optimization\n\nIf you need to to train different parts of the model separately, you have three options:\n\n1. Use `stop_gradient`.\n2. Filter the parameters for `jax.grad`.\n3. Use multiple optimizers for different parameters.\n\nFor most situations we recommend using multiple optimizers via [Optax](https://optax.readthedocs.io/)'s [`multi_transform`](https://optax.readthedocs.io/en/latest/api.html#optax.multi_transform) as its both efficient and can be easily extended to implement many fine-tunning strategies. \n\n### **optax.multi_transform**\n\nTo use `optax.multi_transform` following must be defined:\n\n1. The parameter partitions.\n2. A mapping between partitions and their optimizer.\n3. A pytree with the same shape as the parameters but its leaves containing the corresponding partition label.\n\nTo freeze layers with `optax.multi_transform` for the model above, the following setup can be used:\n\n* Define the `trainable` and `frozen` parameter partitions.\n* For the `trainable` parameters select the Adam (`optax.adam`) optimizer.\n- For the `frozen` parameters select the `optax.set_to_zero` optimizer. This dummy optimizer zeros-out the gradients so no training is done.\n- Map parameters to partitions using [`flax.traverse_util.path_aware_map`](https://flax.readthedocs.io/en/latest/api_reference/flax.traverse_util.html#flax.traverse_util.path_aware_map), mark the leaves from the `backbone` as `frozen`, and the rest as `trainable`.\n\n```{code-cell} ipython3\nfrom flax import traverse_util\nimport optax\n\npartition_optimizers = {'trainable': optax.adam(5e-3), 'frozen': optax.set_to_zero()}\nparam_partitions = traverse_util.path_aware_map(\n lambda path, v: 'frozen' if 'backbone' in path else 'trainable', params)\ntx = optax.multi_transform(partition_optimizers, param_partitions)\n\n# visualize a subset of the param_partitions structure\nflat = list(traverse_util.flatten_dict(param_partitions).items())\ntraverse_util.unflatten_dict(dict(flat[:2] + flat[-2:]))\n```\n\nTo implement [differential learning rates](https://blog.slavv.com/differential-learning-rates-59eff5209a4f), the `optax.set_to_zero` can be replaced with any other optimizer, different optimizers and partitioning schemes can be selected depending on the task. For more information on advanced optimizers, refer to Optax's [Combining Optimizers](https://optax.readthedocs.io/en/latest/api.html#combining-optimizers) documentation.\n\n## Creating the `TrainState`\n\nOnce the module, params, and optimizer are defined, the `TrainState` can be constructed as usual:\n\n```{code-cell} ipython3\nfrom flax.training.train_state import TrainState\n\nstate = TrainState.create(\n apply_fn=model.apply,\n params=params,\n tx=tx)\n```\n\nSince the optimizer takes care of the freezing or fine-tunning strategy, the `train_step` requires no additional changes, training can proceed normally.\n" }, { "path": "docs/guides/training_techniques/use_checkpointing.ipynb", "content": "{\n \"cells\": [\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"id\": \"6e9134fa\",\n \"metadata\": {},\n \"source\": [\n \"# Save and load checkpoints\\n\",\n \"\\n\",\n \"This guide demonstrates how to save and load Flax checkpoints with [Orbax](https://github.com/google/orbax).\\n\",\n \"\\n\",\n \"Orbax provides a variety of features for saving and loading model data, which you will learn about in this doc:\\n\",\n \"\\n\",\n \"* Support for various array types and storage formats\\n\",\n \"* Asynchronous saving to reduce training wait time\\n\",\n \"* Versioning and automatic bookkeeping of past checkpoints\\n\",\n \"* Flexible [`transformations`](https://orbax.readthedocs.io/en/latest/transformations.html) to tweak and load old checkpoints\\n\",\n \"* [`jax.sharding`](https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html)-based API to save and load in multi-host scenarios\\n\",\n \"\\n\",\n \"---\\n\",\n \"**_Ongoing migration to Orbax:_**\\n\",\n \"\\n\",\n \"After July 30 2023, Flax's legacy `flax.training.checkpoints` API will be deprecated in favor of [Orbax](https://github.com/google/orbax).\\n\",\n \"\\n\",\n \"* **If you are a new Flax user**: Use the new `orbax.checkpoint` API, as demonstrated in this guide.\\n\",\n \"\\n\",\n \"* **If you have legacy `flax.training.checkpoints` code in your project**: Consider the following options:\\n\",\n \"\\n\",\n \" * **Migrating your code to Orbax (Recommended)**: Migrate your API calls to `orbax.checkpoint` API by following this [migration guide](https://flax.readthedocs.io/en/latest/guides/converting_and_upgrading/orbax_upgrade_guide.html).\\n\",\n \"\\n\",\n \" * **Automatically use the Orbax backend**: Add `flax.config.update('flax_use_orbax_checkpointing', True)` to your project, which will let your `flax.training.checkpoints` calls automatically use the Orbax backend to save your checkpoints.\\n\",\n \"\\n\",\n \" * **Scheduled flip**: This will become the default mode after **May 2023** (tentative date).\\n\",\n \"\\n\",\n \" * Visit [Orbax-as-backend troubleshooting section](https://flax.readthedocs.io/en/latest/guides/training_techniques/use_checkpointing.html#orbax-as-backend-troubleshooting) if you meet any issue in the automatic migration.\\n\",\n \"---\\n\",\n \"\\n\",\n \"For backward-compatibility, this guide shows the Orbax-equivalent calls in the Flax legacy `flax.training.checkpoints` API.\\n\",\n \"\\n\",\n \"If you need to learn more about `orbax.checkpoint`, refer to the [Orbax docs](https://orbax.readthedocs.io/en/latest/).\\n\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"5a2f6aae\",\n \"metadata\": {},\n \"source\": [\n \"## Setup\\n\",\n \"\\n\",\n \"Install/upgrade Flax and [Orbax](https://github.com/google/orbax). For JAX installation with GPU/TPU support, visit [this section on GitHub](https://github.com/jax-ml/jax#installation).\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"id\": \"-icO30rwmKYj\",\n \"metadata\": {},\n \"source\": [\n \"Note: Before running `import jax`, create eight fake devices to mimic a [multi-host environment](https://jax.readthedocs.io/en/latest/jax-101/06-parallelism.html?#aside-hosts-and-devices-in-jax) in this notebook. Note that the order of imports is important here. The `os.environ[\\\"XLA_FLAGS\\\"] = '--xla_force_host_platform_device_count=8'` command works only with the CPU backend, which means it won't work with GPU/TPU acceleration on if you're running this notebook in Google Colab. If you are already running the code on multiple devices (for example, in a 4x2 TPU environment), you can skip running the next cell.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 1,\n \"id\": \"ArKLnsyGRxGv\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"import os\\n\",\n \"os.environ[\\\"XLA_FLAGS\\\"] = '--xla_force_host_platform_device_count=8'\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 2,\n \"id\": \"SJT9DTxTytjn\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stderr\",\n \"output_type\": \"stream\",\n \"text\": [\n \"WARNING:absl:Tensorflow library not found, tensorflow.io.gfile operations will use native shim calls. GCS paths (i.e. 'gs://...') cannot be accessed.\\n\"\n ]\n }\n ],\n \"source\": [\n \"from typing import Optional, Any\\n\",\n \"import shutil\\n\",\n \"\\n\",\n \"import numpy as np\\n\",\n \"import jax\\n\",\n \"from jax import random, numpy as jnp\\n\",\n \"\\n\",\n \"import flax\\n\",\n \"from flax import linen as nn\\n\",\n \"from flax.training import checkpoints, train_state\\n\",\n \"from flax import struct, serialization\\n\",\n \"import orbax.checkpoint\\n\",\n \"\\n\",\n \"import optax\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 3,\n \"id\": \"afd6db30\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"ckpt_dir = '/tmp/flax_ckpt'\\n\",\n \"\\n\",\n \"if os.path.exists(ckpt_dir):\\n\",\n \" shutil.rmtree(ckpt_dir) # Remove any existing checkpoints from the last notebook run.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"40d434cd\",\n \"metadata\": {},\n \"source\": [\n \"## Save checkpoints\\n\",\n \"\\n\",\n \"In Orbax and Flax, you can save and load any given JAX [pytree](https://jax.readthedocs.io/en/latest/pytrees.html). This includes not only typical Python and NumPy containers, but also customized classes extended from [`flax.struct.dataclass`](https://flax.readthedocs.io/en/latest/api_reference/flax.struct.html#flax.struct.dataclass). That means you can store almost any data generated — not only your model parameters, but any arrays/dictionaries, metadata/configs, and so on.\\n\",\n \"\\n\",\n \"First, create a pytree with many data structures and containers, and play with it:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 4,\n \"id\": \"56dec3f6\",\n \"metadata\": {\n \"outputId\": \"f1856d96-1961-48ed-bb7c-cb63fbaa7567\"\n },\n \"outputs\": [\n {\n \"name\": \"stderr\",\n \"output_type\": \"stream\",\n \"text\": [\n \"WARNING: All log messages before absl::InitializeLog() is called are written to STDERR\\n\",\n \"I0000 00:00:1695322343.254588 1 tfrt_cpu_pjrt_client.cc:349] TfrtCpuClient created.\\n\"\n ]\n },\n {\n \"data\": {\n \"text/plain\": [\n \"{'model': TrainState(step=1, apply_fn=, params={'bias': Array([-0.001, -0.001, -0.001], dtype=float32), 'kernel': Array([[ 0.26048955, -0.61399287, -0.23458514],\\n\",\n \" [ 0.11050402, -0.8765793 , 0.9800635 ],\\n\",\n \" [ 0.36260957, 0.18276349, -0.6856061 ],\\n\",\n \" [-0.8519373 , -0.6416717 , -0.4818122 ],\\n\",\n \" [-0.6886102 , -0.33987316, -0.05898903]], dtype=float32)}, tx=GradientTransformationExtraArgs(init=.init_fn at 0x13d5d83a0>, update=.update_fn at 0x13d5d8dc0>), opt_state=(EmptyState(), EmptyState())),\\n\",\n \" 'config': {'dimensions': array([5, 3])},\\n\",\n \" 'data': [Array([0.59902626, 0.2172144 , 2.4202902 , 0.03266738, 1.2164948 ], dtype=float32)]}\"\n ]\n },\n \"execution_count\": 4,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"# A simple model with one linear layer.\\n\",\n \"key1, key2 = random.split(random.key(0))\\n\",\n \"x1 = random.normal(key1, (5,)) # A simple JAX array.\\n\",\n \"model = nn.Dense(features=3)\\n\",\n \"variables = model.init(key2, x1)\\n\",\n \"\\n\",\n \"# Flax's TrainState is a pytree dataclass and is supported in checkpointing.\\n\",\n \"# Define your class with `@flax.struct.dataclass` decorator to make it compatible.\\n\",\n \"tx = optax.sgd(learning_rate=0.001) # An Optax SGD optimizer.\\n\",\n \"state = train_state.TrainState.create(\\n\",\n \" apply_fn=model.apply,\\n\",\n \" params=variables['params'],\\n\",\n \" tx=tx)\\n\",\n \"# Perform a simple gradient update similar to the one during a normal training workflow.\\n\",\n \"state = state.apply_gradients(grads=jax.tree_util.tree_map(jnp.ones_like, state.params))\\n\",\n \"\\n\",\n \"# Some arbitrary nested pytree with a dictionary and a NumPy array.\\n\",\n \"config = {'dimensions': np.array([5, 3])}\\n\",\n \"\\n\",\n \"# Bundle everything together.\\n\",\n \"ckpt = {'model': state, 'config': config, 'data': [x1]}\\n\",\n \"ckpt\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"8c715b95\",\n \"metadata\": {},\n \"source\": [\n \"### With Orbax\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"6fc59dfa\",\n \"metadata\": {},\n \"source\": [\n \"Save the checkpoint with `orbax.checkpoint.PyTreeCheckpointer`, directly to the `tmp/orbax/single_save` directory.\\n\",\n \"\\n\",\n \"Note: An optional `save_args` is provided. This is recommended for performance speedups, as it bundles smaller arrays in your pytree to a single large file instead of multiple smaller files.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 5,\n \"id\": \"61b12da2\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"from flax.training import orbax_utils\\n\",\n \"\\n\",\n \"orbax_checkpointer = orbax.checkpoint.PyTreeCheckpointer()\\n\",\n \"save_args = orbax_utils.save_args_from_target(ckpt)\\n\",\n \"orbax_checkpointer.save('/tmp/flax_ckpt/orbax/single_save', ckpt, save_args=save_args)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"07d4de1a\",\n \"metadata\": {},\n \"source\": [\n \"Next, to use versioning and automatic bookkeeping features, you need to wrap `orbax.checkpoint.CheckpointManager` over `orbax.checkpoint.PyTreeCheckpointer`.\\n\",\n \"\\n\",\n \"In addition, provide `orbax.checkpoint.CheckpointManagerOptions` that customizes your needs, such as how often and on what criteria you prefer old checkpoints be deleted. See [documentation](https://orbax.readthedocs.io/en/latest/guides/checkpoint/api_refactor.html) for a full list of options offered.\\n\",\n \"\\n\",\n \"`orbax.checkpoint.CheckpointManager` should be placed at the top-level outside your training steps to manage your saves.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 6,\n \"id\": \"d3686ea5\",\n \"metadata\": {\n \"outputId\": \"b7132933-566d-440d-c34e-c5468d87cbdc\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"['4', '3']\"\n ]\n },\n \"execution_count\": 6,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"options = orbax.checkpoint.CheckpointManagerOptions(max_to_keep=2, create=True)\\n\",\n \"checkpoint_manager = orbax.checkpoint.CheckpointManager(\\n\",\n \" '/tmp/flax_ckpt/orbax/managed', orbax_checkpointer, options)\\n\",\n \"\\n\",\n \"# Inside a training loop\\n\",\n \"for step in range(5):\\n\",\n \" # ... do your training\\n\",\n \" checkpoint_manager.save(step, ckpt, save_kwargs={'save_args': save_args})\\n\",\n \"\\n\",\n \"os.listdir('/tmp/flax_ckpt/orbax/managed') # Because max_to_keep=2, only step 3 and 4 are retained\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"8ecbc4cc\",\n \"metadata\": {},\n \"source\": [\n \"### With the legacy API\\n\",\n \"\\n\",\n \"And here's how to save with the legacy Flax checkpointing utilities (note that this provides less management features compared with `orbax.checkpoint.CheckpointManagerOptions`):\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 7,\n \"id\": \"4cdb35ef\",\n \"metadata\": {\n \"outputId\": \"6d849273-15ce-4480-8864-726d1838ac1f\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"'/tmp/flax_ckpt/flax-checkpointing/checkpoint_0'\"\n ]\n },\n \"execution_count\": 7,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"# Import Flax Checkpoints.\\n\",\n \"from flax.training import checkpoints\\n\",\n \"\\n\",\n \"checkpoints.save_checkpoint(ckpt_dir='/tmp/flax_ckpt/flax-checkpointing',\\n\",\n \" target=ckpt,\\n\",\n \" step=0,\\n\",\n \" overwrite=True,\\n\",\n \" keep=2)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"6b658bd1\",\n \"metadata\": {},\n \"source\": [\n \"## Restore checkpoints\\n\",\n \"\\n\",\n \"### With Orbax\\n\",\n \"\\n\",\n \"In Orbax, call `.restore()` for either `orbax.checkpoint.PyTreeCheckpointer` or `orbax.checkpoint.CheckpointManager` to restore your checkpoint in the raw pytree format.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 8,\n \"id\": \"a807a9c1\",\n \"metadata\": {\n \"outputId\": \"b4af1ef4-f22f-459b-bdca-2e6bfa16c08b\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"{'config': {'dimensions': array([5, 3])},\\n\",\n \" 'data': [array([0.59902626, 0.2172144 , 2.4202902 , 0.03266738, 1.2164948 ],\\n\",\n \" dtype=float32)],\\n\",\n \" 'model': {'opt_state': [None, None],\\n\",\n \" 'params': {'bias': array([-0.001, -0.001, -0.001], dtype=float32),\\n\",\n \" 'kernel': array([[ 0.26048955, -0.61399287, -0.23458514],\\n\",\n \" [ 0.11050402, -0.8765793 , 0.9800635 ],\\n\",\n \" [ 0.36260957, 0.18276349, -0.6856061 ],\\n\",\n \" [-0.8519373 , -0.6416717 , -0.4818122 ],\\n\",\n \" [-0.6886102 , -0.33987316, -0.05898903]], dtype=float32)},\\n\",\n \" 'step': 1}}\"\n ]\n },\n \"execution_count\": 8,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"raw_restored = orbax_checkpointer.restore('/tmp/flax_ckpt/orbax/single_save')\\n\",\n \"raw_restored\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"8c015a22\",\n \"metadata\": {},\n \"source\": [\n \"Note that the `step` number is required for `CheckpointManger`. You can also use `.latest_step()` to find the latest step available.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 9,\n \"id\": \"251d7085\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"{'config': {'dimensions': array([5, 3])},\\n\",\n \" 'data': [array([0.59902626, 0.2172144 , 2.4202902 , 0.03266738, 1.2164948 ],\\n\",\n \" dtype=float32)],\\n\",\n \" 'model': {'opt_state': [None, None],\\n\",\n \" 'params': {'bias': array([-0.001, -0.001, -0.001], dtype=float32),\\n\",\n \" 'kernel': array([[ 0.26048955, -0.61399287, -0.23458514],\\n\",\n \" [ 0.11050402, -0.8765793 , 0.9800635 ],\\n\",\n \" [ 0.36260957, 0.18276349, -0.6856061 ],\\n\",\n \" [-0.8519373 , -0.6416717 , -0.4818122 ],\\n\",\n \" [-0.6886102 , -0.33987316, -0.05898903]], dtype=float32)},\\n\",\n \" 'step': 1}}\"\n ]\n },\n \"execution_count\": 9,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"step = checkpoint_manager.latest_step() # step = 4\\n\",\n \"checkpoint_manager.restore(step)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"c7fe3bc8\",\n \"metadata\": {},\n \"source\": [\n \"### With the legacy API\\n\",\n \"\\n\",\n \"Note that with the migration to Orbax in progress, `flax.training.checkpointing.restore_checkpoint` can automatically identify whether a checkpoint is saved in the legacy Flax format or with an Orbax backend, and restore the pytree correctly. Therefore, adding `flax.config.update('flax_use_orbax_checkpointing', True)` won't hurt your ability to restore old checkpoints.\\n\",\n \"\\n\",\n \"Here's how to restore checkpoints using the legacy API:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 10,\n \"id\": \"150b20a0\",\n \"metadata\": {\n \"outputId\": \"85ffceca-f38d-46b8-e567-d9d38b7885f9\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"{'config': {'dimensions': array([5, 3])},\\n\",\n \" 'data': {'0': array([0.59902626, 0.2172144 , 2.4202902 , 0.03266738, 1.2164948 ],\\n\",\n \" dtype=float32)},\\n\",\n \" 'model': {'opt_state': {'0': None, '1': None},\\n\",\n \" 'params': {'bias': array([-0.001, -0.001, -0.001], dtype=float32),\\n\",\n \" 'kernel': array([[ 0.26048955, -0.61399287, -0.23458514],\\n\",\n \" [ 0.11050402, -0.8765793 , 0.9800635 ],\\n\",\n \" [ 0.36260957, 0.18276349, -0.6856061 ],\\n\",\n \" [-0.8519373 , -0.6416717 , -0.4818122 ],\\n\",\n \" [-0.6886102 , -0.33987316, -0.05898903]], dtype=float32)},\\n\",\n \" 'step': 1}}\"\n ]\n },\n \"execution_count\": 10,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"raw_restored = checkpoints.restore_checkpoint(ckpt_dir='/tmp/flax_ckpt/flax-checkpointing', target=None)\\n\",\n \"raw_restored\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"987b981f\",\n \"metadata\": {},\n \"source\": [\n \"## Restore with custom dataclasses\\n\",\n \"\\n\",\n \"### With Orbax\\n\",\n \"\\n\",\n \"* The pytrees restored in the previous examples are in the form of raw dictionaries. Original pytrees contain custom dataclasses like [`TrainState`](https://flax.readthedocs.io/en/latest/flip/1009-optimizer-api.html?#train-state) and `optax` states.\\n\",\n \"* This is because when restoring a pytree, the program does not yet know which structure it once belonged to.\\n\",\n \"* To resolve this, you should first provide an example pytree to let Orbax or Flax know exactly which structure to restore to.\\n\",\n \"\\n\",\n \"This section demonstrates how to set up any custom Flax dataclass explicitly, and have the same structure as a saved checkpoint.\\n\",\n \"\\n\",\n \"Note: Data that was a JAX NumPy array (`jnp.array`) format will be restored as a NumPy array (`numpy.array`). This would not affect your work because JAX will [automatically convert](https://jax.readthedocs.io/en/latest/jax-101/01-jax-basics.html) NumPy arrays to JAX arrays once the computation starts.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 11,\n \"id\": \"58f42513\",\n \"metadata\": {\n \"outputId\": \"110c6b6e-fe42-4179-e5d8-6b92d355e11b\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"{'config': {'dimensions': array([5, 3])},\\n\",\n \" 'data': [array([0.59902626, 0.2172144 , 2.4202902 , 0.03266738, 1.2164948 ],\\n\",\n \" dtype=float32)],\\n\",\n \" 'model': TrainState(step=1, apply_fn=, params={'bias': array([-0.001, -0.001, -0.001], dtype=float32), 'kernel': array([[ 0.26048955, -0.61399287, -0.23458514],\\n\",\n \" [ 0.11050402, -0.8765793 , 0.9800635 ],\\n\",\n \" [ 0.36260957, 0.18276349, -0.6856061 ],\\n\",\n \" [-0.8519373 , -0.6416717 , -0.4818122 ],\\n\",\n \" [-0.6886102 , -0.33987316, -0.05898903]], dtype=float32)}, tx=GradientTransformationExtraArgs(init=.init_fn at 0x13d5d83a0>, update=.update_fn at 0x13d5d8dc0>), opt_state=(EmptyState(), EmptyState()))}\"\n ]\n },\n \"execution_count\": 11,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"empty_state = train_state.TrainState.create(\\n\",\n \" apply_fn=model.apply,\\n\",\n \" params=jax.tree_util.tree_map(np.zeros_like, variables['params']), # values of the tree leaf doesn't matter\\n\",\n \" tx=tx,\\n\",\n \")\\n\",\n \"empty_config = {'dimensions': np.array([0, 0])}\\n\",\n \"target = {'model': empty_state, 'config': empty_config, 'data': [jnp.zeros_like(x1)]}\\n\",\n \"state_restored = orbax_checkpointer.restore('/tmp/flax_ckpt/orbax/single_save', item=target)\\n\",\n \"state_restored\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"f1c18bc6\",\n \"metadata\": {},\n \"source\": [\n \"### With the legacy API\\n\",\n \"\\n\",\n \"Alternatively, you can restore from Orbax `CheckpointManager` and from the legacy Flax code as follows:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 12,\n \"id\": \"a61e9a66\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"{'config': {'dimensions': array([5, 3])},\\n\",\n \" 'data': [array([0.59902626, 0.2172144 , 2.4202902 , 0.03266738, 1.2164948 ],\\n\",\n \" dtype=float32)],\\n\",\n \" 'model': TrainState(step=1, apply_fn=, params={'bias': array([-0.001, -0.001, -0.001], dtype=float32), 'kernel': array([[ 0.26048955, -0.61399287, -0.23458514],\\n\",\n \" [ 0.11050402, -0.8765793 , 0.9800635 ],\\n\",\n \" [ 0.36260957, 0.18276349, -0.6856061 ],\\n\",\n \" [-0.8519373 , -0.6416717 , -0.4818122 ],\\n\",\n \" [-0.6886102 , -0.33987316, -0.05898903]], dtype=float32)}, tx=GradientTransformationExtraArgs(init=.init_fn at 0x13d5d83a0>, update=.update_fn at 0x13d5d8dc0>), opt_state=(EmptyState(), EmptyState()))}\"\n ]\n },\n \"execution_count\": 12,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"checkpoint_manager.restore(4, items=target)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 13,\n \"id\": \"412af50e\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stderr\",\n \"output_type\": \"stream\",\n \"text\": [\n \"WARNING:absl:The transformations API will eventually be replaced by an upgraded design. The current API will not be removed until this point, but it will no longer be actively worked on.\\n\"\n ]\n },\n {\n \"data\": {\n \"text/plain\": [\n \"{'model': TrainState(step=1, apply_fn=, params={'bias': array([-0.001, -0.001, -0.001], dtype=float32), 'kernel': array([[ 0.26048955, -0.61399287, -0.23458514],\\n\",\n \" [ 0.11050402, -0.8765793 , 0.9800635 ],\\n\",\n \" [ 0.36260957, 0.18276349, -0.6856061 ],\\n\",\n \" [-0.8519373 , -0.6416717 , -0.4818122 ],\\n\",\n \" [-0.6886102 , -0.33987316, -0.05898903]], dtype=float32)}, tx=GradientTransformationExtraArgs(init=.init_fn at 0x13d5d83a0>, update=.update_fn at 0x13d5d8dc0>), opt_state=(EmptyState(), EmptyState())),\\n\",\n \" 'config': {'dimensions': array([5, 3])},\\n\",\n \" 'data': [Array([0.59902626, 0.2172144 , 2.4202902 , 0.03266738, 1.2164948 ], dtype=float32)]}\"\n ]\n },\n \"execution_count\": 13,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"checkpoints.restore_checkpoint(ckpt_dir='/tmp/flax_ckpt/flax-checkpointing', target=target)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"27461ac8\",\n \"metadata\": {},\n \"source\": [\n \"It's often recommended to refactor out the process of initializing a checkpoint's structure (for example, a [`TrainState`](https://flax.readthedocs.io/en/latest/flip/1009-optimizer-api.html?#train-state)), so that saving/loading is easier and less error-prone. This is because functions and complex objects like `apply_fn` and `tx` (optimizer) cannot be serialized into the checkpoint file and must be initialized by code.\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"id\": \"136a300a\",\n \"metadata\": {},\n \"source\": [\n \"## Restore when checkpoint structures differ\\n\",\n \"\\n\",\n \"During your development, your checkpoint structure will change when changing the model, adding/removing fields during tweaking, and so on.\\n\",\n \"\\n\",\n \"This section explains how to load old data to your new code.\\n\",\n \"\\n\",\n \"Below is a simple example — a `CustomTrainState` extended from `flax.training.train_state.TrainState` that contains an extra field called `batch_stats`. When working on a real-world model, you may need this when applying [batch normalization](https://flax.readthedocs.io/en/latest/guides/training_techniques/batch_norm.html).\\n\",\n \"\\n\",\n \"Here, you store the new `CustomTrainState` as step 5, while step 4 contains the old/previous `TrainState`.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 14,\n \"id\": \"be65d4af\",\n \"metadata\": {\n \"outputId\": \"4fe776f0-65f8-4fc4-d64a-990520b36dce\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"True\"\n ]\n },\n \"execution_count\": 14,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"class CustomTrainState(train_state.TrainState):\\n\",\n \" batch_stats: Any = None\\n\",\n \"\\n\",\n \"custom_state = CustomTrainState.create(\\n\",\n \" apply_fn=state.apply_fn,\\n\",\n \" params=state.params,\\n\",\n \" tx=state.tx,\\n\",\n \" batch_stats=np.arange(10),\\n\",\n \")\\n\",\n \"\\n\",\n \"custom_ckpt = {'model': custom_state, 'config': config, 'data': [x1]}\\n\",\n \"# Use a custom state to read the old `TrainState` checkpoint.\\n\",\n \"custom_target = {'model': custom_state, 'config': None, 'data': [jnp.zeros_like(x1)]}\\n\",\n \"\\n\",\n \"# Save it in Orbax.\\n\",\n \"custom_save_args = orbax_utils.save_args_from_target(custom_ckpt)\\n\",\n \"checkpoint_manager.save(5, custom_ckpt, save_kwargs={'save_args': custom_save_args})\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"379c2255\",\n \"metadata\": {},\n \"source\": [\n \"It is recommended to keep your checkpoints up-to-date with your pytree dataclass definitions. However, you might be forced to restore the checkpoints with incompatible reference objects at runtime. When this happens, the checkpoint restoration will try to respect the structure of the reference when given.\\n\",\n \"\\n\",\n \"Below are examples of a few common scenarios.\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"id\": \"d5fa9652\",\n \"metadata\": {},\n \"source\": [\n \"### Scenario 1: When a reference object is partial\\n\",\n \"\\n\",\n \"If your reference object is a subtree of your checkpoint, the restoration will ignore the additional field(s) and restore a checkpoint with the same structure as the reference.\\n\",\n \"\\n\",\n \"Like in the example below, the `batch_stats` field in `CustomTrainState` was ignored, and the checkpoint was restored as a `TrainState`.\\n\",\n \"\\n\",\n \"This can also be useful for reading only part of your checkpoint.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 15,\n \"id\": \"68828029\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"{'config': {'dimensions': array([5, 3])},\\n\",\n \" 'data': [array([0.59902626, 0.2172144 , 2.4202902 , 0.03266738, 1.2164948 ],\\n\",\n \" dtype=float32)],\\n\",\n \" 'model': TrainState(step=0, apply_fn=, params={'bias': array([-0.001, -0.001, -0.001], dtype=float32), 'kernel': array([[ 0.26048955, -0.61399287, -0.23458514],\\n\",\n \" [ 0.11050402, -0.8765793 , 0.9800635 ],\\n\",\n \" [ 0.36260957, 0.18276349, -0.6856061 ],\\n\",\n \" [-0.8519373 , -0.6416717 , -0.4818122 ],\\n\",\n \" [-0.6886102 , -0.33987316, -0.05898903]], dtype=float32)}, tx=GradientTransformationExtraArgs(init=.init_fn at 0x13d5d83a0>, update=.update_fn at 0x13d5d8dc0>), opt_state=(EmptyState(), EmptyState()))}\"\n ]\n },\n \"execution_count\": 15,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"restored = checkpoint_manager.restore(5, items=target)\\n\",\n \"assert not hasattr(restored, 'batch_stats')\\n\",\n \"assert type(restored['model']) == train_state.TrainState\\n\",\n \"restored\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"id\": \"5c6822c6\",\n \"metadata\": {},\n \"source\": [\n \"### Scenario 2: When a checkpoint is partial\\n\",\n \"\\n\",\n \"On the other hand, if the reference object contains a value that is not available in the checkpoint, the checkpointing code will by default warn that some data is not compatible.\\n\",\n \"\\n\",\n \"To bypass the error, you need to pass an Orbax [`transform`](https://orbax.readthedocs.io/en/latest/guides/checkpoint/transformations.html) that teaches Orbax how to conform this checkpoint into the structure of the `custom_target`.\\n\",\n \"\\n\",\n \"In this case, pass a default `{}` that lets Orbax use values in the `custom_target` to fill in the blank. This allows you to restore an old checkpoint into a new data structure, the `CustomTrainState`.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 16,\n \"id\": \"a5d14c9f\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stderr\",\n \"output_type\": \"stream\",\n \"text\": [\n \"WARNING:absl:The transformations API will eventually be replaced by an upgraded design. The current API will not be removed until this point, but it will no longer be actively worked on.\\n\"\n ]\n },\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"KeyError when target state has an unmentioned field: 'batch_stats'\\n\",\n \"\\n\"\n ]\n },\n {\n \"data\": {\n \"text/plain\": [\n \"{'config': None,\\n\",\n \" 'data': [Array([0.59902626, 0.2172144 , 2.4202902 , 0.03266738, 1.2164948 ], dtype=float32)],\\n\",\n \" 'model': CustomTrainState(step=1, apply_fn=, params={'bias': Array([-0.001, -0.001, -0.001], dtype=float32), 'kernel': Array([[ 0.26048955, -0.61399287, -0.23458514],\\n\",\n \" [ 0.11050402, -0.8765793 , 0.9800635 ],\\n\",\n \" [ 0.36260957, 0.18276349, -0.6856061 ],\\n\",\n \" [-0.8519373 , -0.6416717 , -0.4818122 ],\\n\",\n \" [-0.6886102 , -0.33987316, -0.05898903]], dtype=float32)}, tx=GradientTransformationExtraArgs(init=.init_fn at 0x13d5d83a0>, update=.update_fn at 0x13d5d8dc0>), opt_state=(EmptyState(), EmptyState()), batch_stats=array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]))}\"\n ]\n },\n \"execution_count\": 16,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"try:\\n\",\n \" checkpoint_manager.restore(4, items=custom_target)\\n\",\n \"except KeyError as e:\\n\",\n \" print(f'KeyError when target state has an unmentioned field: {e}')\\n\",\n \" print('')\\n\",\n \"\\n\",\n \"# Step 4 is an original `TrainState`, without the `batch_stats`\\n\",\n \"custom_restore_args = orbax_utils.restore_args_from_target(custom_target)\\n\",\n \"restored = checkpoint_manager.restore(4, items=custom_target,\\n\",\n \" restore_kwargs={'transforms': {}, 'restore_args': custom_restore_args})\\n\",\n \"assert type(restored['model']) == CustomTrainState\\n\",\n \"np.testing.assert_equal(restored['model'].batch_stats,\\n\",\n \" custom_target['model'].batch_stats)\\n\",\n \"restored\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"74a4b0fd\",\n \"metadata\": {},\n \"source\": [\n \"##### With Orbax\\n\",\n \"\\n\",\n \"If you have already saved your checkpoints with the Orbax backend, you can use `orbax_transforms` to access this `transforms` argument in the Flax API.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 17,\n \"id\": \"29fd1e33\",\n \"metadata\": {\n \"outputId\": \"cdbb9247-d1eb-4458-aa83-8db0332af7cb\"\n },\n \"outputs\": [\n {\n \"name\": \"stderr\",\n \"output_type\": \"stream\",\n \"text\": [\n \"WARNING:absl:The transformations API will eventually be replaced by an upgraded design. The current API will not be removed until this point, but it will no longer be actively worked on.\\n\"\n ]\n },\n {\n \"data\": {\n \"text/plain\": [\n \"{'model': CustomTrainState(step=1, apply_fn=, params={'bias': Array([-0.001, -0.001, -0.001], dtype=float32), 'kernel': Array([[ 0.26048955, -0.61399287, -0.23458514],\\n\",\n \" [ 0.11050402, -0.8765793 , 0.9800635 ],\\n\",\n \" [ 0.36260957, 0.18276349, -0.6856061 ],\\n\",\n \" [-0.8519373 , -0.6416717 , -0.4818122 ],\\n\",\n \" [-0.6886102 , -0.33987316, -0.05898903]], dtype=float32)}, tx=GradientTransformationExtraArgs(init=.init_fn at 0x13d5d83a0>, update=.update_fn at 0x13d5d8dc0>), opt_state=(EmptyState(), EmptyState()), batch_stats=array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])),\\n\",\n \" 'config': None,\\n\",\n \" 'data': [Array([0.59902626, 0.2172144 , 2.4202902 , 0.03266738, 1.2164948 ], dtype=float32)]}\"\n ]\n },\n \"execution_count\": 17,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"# Save in the \\\"Flax-with-Orbax\\\" backend.\\n\",\n \"flax.config.update('flax_use_orbax_checkpointing', True)\\n\",\n \"checkpoints.save_checkpoint(ckpt_dir='/tmp/flax_ckpt/flax-checkpointing',\\n\",\n \" target=ckpt,\\n\",\n \" step=4,\\n\",\n \" overwrite=True,\\n\",\n \" keep=2)\\n\",\n \"\\n\",\n \"checkpoints.restore_checkpoint('/tmp/flax_ckpt/flax-checkpointing', target=custom_target, step=4,\\n\",\n \" orbax_transforms={})\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"830ef07c\",\n \"metadata\": {},\n \"source\": [\n \"##### With the legacy API\\n\",\n \"\\n\",\n \"Using the legacy `flax.training.checkpoints` API, similar things are doable too, but they are not as flexible as the [Orbax Transformations](https://orbax.readthedocs.io/en/latest/guides/checkpoint/transformations.html).\\n\",\n \"\\n\",\n \"You need to restore the checkpoint to a raw dict with `target=None`, modify the structure accordingly, and then deserialize it back to the original target.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 18,\n \"id\": \"051e7a16\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"{'model': CustomTrainState(step=1, apply_fn=, params={'bias': array([-0.001, -0.001, -0.001], dtype=float32), 'kernel': array([[ 0.26048955, -0.61399287, -0.23458514],\\n\",\n \" [ 0.11050402, -0.8765793 , 0.9800635 ],\\n\",\n \" [ 0.36260957, 0.18276349, -0.6856061 ],\\n\",\n \" [-0.8519373 , -0.6416717 , -0.4818122 ],\\n\",\n \" [-0.6886102 , -0.33987316, -0.05898903]], dtype=float32)}, tx=GradientTransformationExtraArgs(init=.init_fn at 0x13d5d83a0>, update=.update_fn at 0x13d5d8dc0>), opt_state=(EmptyState(), EmptyState()), batch_stats=array([9, 8, 7, 6, 5, 4, 3, 2, 1, 0])),\\n\",\n \" 'config': {'dimensions': array([5, 3])},\\n\",\n \" 'data': [array([0.59902626, 0.2172144 , 2.4202902 , 0.03266738, 1.2164948 ],\\n\",\n \" dtype=float32)]}\"\n ]\n },\n \"execution_count\": 18,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"# Save using the legacy Flax `checkpoints` API.\\n\",\n \"flax.config.update('flax_use_orbax_checkpointing', False)\\n\",\n \"checkpoints.save_checkpoint(ckpt_dir='/tmp/flax_ckpt/flax-checkpointing',\\n\",\n \" target=ckpt,\\n\",\n \" step=5,\\n\",\n \" overwrite=True,\\n\",\n \" keep=2)\\n\",\n \"\\n\",\n \"# Pass no target to get a raw state dictionary first.\\n\",\n \"raw_state_dict = checkpoints.restore_checkpoint('/tmp/flax_ckpt/flax-checkpointing', target=None, step=5)\\n\",\n \"# Add/remove fields as needed.\\n\",\n \"raw_state_dict['model']['batch_stats'] = np.flip(np.arange(10))\\n\",\n \"# Restore the classes with correct target now\\n\",\n \"flax.serialization.from_state_dict(custom_target, raw_state_dict)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"a6b39501\",\n \"metadata\": {},\n \"source\": [\n \"## Asynchronized checkpointing\\n\",\n \"\\n\",\n \"Checkpointing is I/O heavy, and if you have a large amount of data to save, it may be worthwhile to put it into a background thread, while continuing with your training.\\n\",\n \"\\n\",\n \"You can do this by creating an [`orbax.checkpoint.AsyncCheckpointer`](https://github.com/google/orbax/blob/main/checkpoint/orbax/checkpoint/async_checkpointer.py) in place of the `orbax.checkpoint.PyTreeCheckpointer`.\\n\",\n \"\\n\",\n \"Note: You should use the same `async_checkpointer` to handle all your async saves across your training steps, so that it can make sure that a previous async save is done before the next one begins. This enables bookkeeping, such as `keep` (the number of checkpoints) and `overwrite` to be consistent across steps.\\n\",\n \"\\n\",\n \"Whenever you want to explicitly wait until an async save is done, you can call `async_checkpointer.wait_until_finished()`.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 19,\n \"id\": \"85be68a6\",\n \"metadata\": {\n \"outputId\": \"aefce94c-8bae-4355-c142-05f2b61c39e2\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"{'config': {'dimensions': array([5, 3])},\\n\",\n \" 'data': [array([0.59902626, 0.2172144 , 2.4202902 , 0.03266738, 1.2164948 ],\\n\",\n \" dtype=float32)],\\n\",\n \" 'model': TrainState(step=1, apply_fn=, params={'bias': array([-0.001, -0.001, -0.001], dtype=float32), 'kernel': array([[ 0.26048955, -0.61399287, -0.23458514],\\n\",\n \" [ 0.11050402, -0.8765793 , 0.9800635 ],\\n\",\n \" [ 0.36260957, 0.18276349, -0.6856061 ],\\n\",\n \" [-0.8519373 , -0.6416717 , -0.4818122 ],\\n\",\n \" [-0.6886102 , -0.33987316, -0.05898903]], dtype=float32)}, tx=GradientTransformationExtraArgs(init=.init_fn at 0x13d5d83a0>, update=.update_fn at 0x13d5d8dc0>), opt_state=(EmptyState(), EmptyState()))}\"\n ]\n },\n \"execution_count\": 19,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"# `orbax.checkpoint.AsyncCheckpointer` needs some multi-process initialization, because it was\\n\",\n \"# originally designed for multi-process large model checkpointing.\\n\",\n \"# For Python notebooks or other single-process settings, just set up with `num_processes=1`.\\n\",\n \"# Refer to https://jax.readthedocs.io/en/latest/multi_process.html#initializing-the-cluster\\n\",\n \"# for how to set it up in multi-process scenarios.\\n\",\n \"jax.distributed.initialize(\\\"localhost:8889\\\", num_processes=1, process_id=0)\\n\",\n \"\\n\",\n \"async_checkpointer = orbax.checkpoint.AsyncCheckpointer(\\n\",\n \" orbax.checkpoint.PyTreeCheckpointHandler(), timeout_secs=50)\\n\",\n \"\\n\",\n \"# Save your job:\\n\",\n \"async_checkpointer.save('/tmp/flax_ckpt/orbax/single_save_async', ckpt, save_args=save_args)\\n\",\n \"# ... Continue with your work...\\n\",\n \"\\n\",\n \"# ... Until a time when you want to wait until the save completes:\\n\",\n \"async_checkpointer.wait_until_finished() # Blocks until the checkpoint saving is completed.\\n\",\n \"async_checkpointer.restore('/tmp/flax_ckpt/orbax/single_save_async', item=target)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"13e93db6\",\n \"metadata\": {},\n \"source\": [\n \"If you are using Orbax `CheckpointManager`, just pass in the async_checkpointer when initializing it. Then, in practice, call `async_checkpoint_manager.wait_until_finished()` instead.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 20,\n \"id\": \"af33b138\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"async_checkpoint_manager = orbax.checkpoint.CheckpointManager(\\n\",\n \" '/tmp/flax_ckpt/orbax/managed_async', async_checkpointer, options)\\n\",\n \"async_checkpoint_manager.wait_until_finished()\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"bb0e03cd\",\n \"metadata\": {},\n \"source\": [\n \"## Multi-host/multi-process checkpointing\\n\",\n \"\\n\",\n \"JAX provides a few ways to scale up your code on multiple hosts at the same time. This usually happens when the number of devices (CPU/GPU/TPU) is so large that different devices are managed by different hosts (CPU). To get started on JAX in multi-process settings, check out [Using JAX in multi-host and multi-process environments](https://jax.readthedocs.io/en/latest/multi_process.html) and the [distributed array guide](https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html).\\n\",\n \"\\n\",\n \"In the [Single Program Multi Data (SPMD)](https://jax.readthedocs.io/en/latest/glossary.html#term-SPMD) paradigm with JAX `jit`, a large multi-process array can have its data sharded across different devices. (Note that JAX `pjit` and `jit` have been merged into a single unified interface. To learn about compiling and executing JAX functions in multi-host or multi-core environments, refer to [this guide](https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html) and the [jax.Array migration guide](https://jax.readthedocs.io/en/latest/jax_array_migration.html).) When a multi-process array is serialized, each host dumps its data shards to a single shared storage, such as a Google Cloud bucket.\\n\",\n \"\\n\",\n \"Orbax supports saving and loading pytrees with multi-process arrays in the same fashion as single-process pytrees. However, it's recommended to use the asynchronized [`orbax.AsyncCheckpointer`](https://github.com/google/orbax/blob/main/checkpoint/orbax/checkpoint/async_checkpointer.py) to save large multi-process arrays on another thread, so that you can perform computation alongside the saves. With pure Orbax, saving checkpoints in a multi-process context uses the same API as in a single-process context.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 21,\n \"id\": \"ubdUvyMrhD-1\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"from jax.sharding import PartitionSpec, NamedSharding\\n\",\n \"\\n\",\n \"# Create an array sharded across multiple devices.\\n\",\n \"mesh_shape = (4, 2)\\n\",\n \"devices = np.asarray(jax.devices()).reshape(*mesh_shape)\\n\",\n \"mesh = jax.sharding.Mesh(devices, ('x', 'y'))\\n\",\n \"\\n\",\n \"mp_array = jax.device_put(np.arange(8 * 2).reshape(8, 2),\\n\",\n \" NamedSharding(mesh, PartitionSpec('x', 'y')))\\n\",\n \"\\n\",\n \"# Make it a pytree.\\n\",\n \"mp_ckpt = {'model': mp_array}\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 22,\n \"id\": \"a669bc05\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"async_checkpoint_manager.save(0, mp_ckpt)\\n\",\n \"async_checkpoint_manager.wait_until_finished()\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"4deee32e\",\n \"metadata\": {},\n \"source\": [\n \"When restoring a checkpoint with multi-process arrays, you need to specify what `sharding` each array should be restored back to. Otherwise, they will be restored as large `np.array`s on process 0, costing time and memory.\\n\",\n \"\\n\",\n \"(In this notebook, since we are on single-process, it will be restored as `np.array` even if we provide shardings.)\\n\",\n \"\\n\",\n \"### With Orbax\\n\",\n \"\\n\",\n \"Orbax allows you to specify this by passing a pytree of `sharding`s in `restore_args`. If you already have a reference pytree that has all the arrays with the right sharding, you can use `orbax_utils.restore_args_from_target` to transform it into the `restore_args` that Orbax needs.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 23,\n \"id\": \"b8e7daaa\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"{'model': Array([[ 0, 1],\\n\",\n \" [ 2, 3],\\n\",\n \" [ 4, 5],\\n\",\n \" [ 6, 7],\\n\",\n \" [ 8, 9],\\n\",\n \" [10, 11],\\n\",\n \" [12, 13],\\n\",\n \" [14, 15]], dtype=int32)}\"\n ]\n },\n \"execution_count\": 23,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"# The reference doesn't need to be as large as your checkpoint!\\n\",\n \"# Just make sure it has the `.sharding` you want.\\n\",\n \"mp_smaller = jax.device_put(np.arange(8).reshape(4, 2),\\n\",\n \" NamedSharding(mesh, PartitionSpec('x', 'y')))\\n\",\n \"ref_ckpt = {'model': mp_smaller}\\n\",\n \"\\n\",\n \"restore_args = orbax_utils.restore_args_from_target(ref_ckpt)\\n\",\n \"async_checkpoint_manager.restore(\\n\",\n \" 0, items=ref_ckpt, restore_kwargs={'restore_args': restore_args})\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"edc355ce\",\n \"metadata\": {},\n \"source\": [\n \"### With the legacy Flax: use `save_checkpoint_multiprocess`\\n\",\n \"\\n\",\n \"In legacy Flax, to save multi-process arrays, use [`flax.training.checkpoints.save_checkpoint_multiprocess()`](https://flax.readthedocs.io/en/latest/api_reference/flax.training.html#flax.training.checkpoints.save_checkpoint_multiprocess) in place of `save_checkpoint()` and with the same arguments.\\n\",\n \"\\n\",\n \"If your checkpoint is too large, you can specify `timeout_secs` in the manager and give it more time to finish writing.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 24,\n \"id\": \"5d10039b\",\n \"metadata\": {\n \"outputId\": \"901bb097-0899-479d-b9ae-61dae79e7057\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"'/tmp/flax_ckpt/checkpoint_3'\"\n ]\n },\n \"execution_count\": 24,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"async_checkpointer = orbax.checkpoint.AsyncCheckpointer(orbax.checkpoint.PyTreeCheckpointHandler(), timeout_secs=50)\\n\",\n \"checkpoints.save_checkpoint_multiprocess(ckpt_dir,\\n\",\n \" mp_ckpt,\\n\",\n \" step=3,\\n\",\n \" overwrite=True,\\n\",\n \" keep=4,\\n\",\n \" orbax_checkpointer=async_checkpointer)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 25,\n \"id\": \"a9f9724c\",\n \"metadata\": {\n \"outputId\": \"393c4a0e-8a8c-4ca6-c609-93c8bab38e75\"\n },\n \"outputs\": [\n {\n \"name\": \"stderr\",\n \"output_type\": \"stream\",\n \"text\": [\n \"WARNING:absl:The transformations API will eventually be replaced by an upgraded design. The current API will not be removed until this point, but it will no longer be actively worked on.\\n\"\n ]\n },\n {\n \"data\": {\n \"text/plain\": [\n \"{'model': Array([[ 0, 1],\\n\",\n \" [ 2, 3],\\n\",\n \" [ 4, 5],\\n\",\n \" [ 6, 7],\\n\",\n \" [ 8, 9],\\n\",\n \" [10, 11],\\n\",\n \" [12, 13],\\n\",\n \" [14, 15]], dtype=int32)}\"\n ]\n },\n \"execution_count\": 25,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"mp_restored = checkpoints.restore_checkpoint(ckpt_dir,\\n\",\n \" target=ref_ckpt,\\n\",\n \" step=3,\\n\",\n \" orbax_checkpointer=async_checkpointer)\\n\",\n \"mp_restored\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"65cfdd59\",\n \"metadata\": {},\n \"source\": [\n \"## Orbax-as-backend troubleshooting\\n\",\n \"\\n\",\n \"As an intermediate stage of the migration (to Orbax from the legacy Flax `checkpoints` API), `flax.training.checkpoints` APIs will start to use Orbax as their backend when saving checkpoints starting from May 15, 2023.\\n\",\n \"\\n\",\n \"Checkpoints saved with the Orbax backend can be readable by either `flax.training.checkpoints.restore_checkpoint` or `orbax.checkpoint.PyTreeCheckpointer`.\\n\",\n \"\\n\",\n \"Code-wise, this is equivalent to setting the config flag [`flax.config.flax_use_orbax_checkpointing`](https://github.com/google/flax/blob/main/flax/configurations.py#L103) default to `True`. You can overwrite this value in your project with `flax.config.update('flax_use_orbax_checkpointing', )` at any time.\\n\",\n \"\\n\",\n \"In general, this automatic migration will not affect most users. However, you may encounter issues if your API usage follows some specific pattern. Check out the sections below for troubleshooting.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"415bceb1\",\n \"metadata\": {},\n \"source\": [\n \"### If your devices hang when writing checkpoints\\n\",\n \"\\n\",\n \"If you are running in a multi-host environment (usually anything larger than 8 TPU devices) and your devices hang when writing checkpoints, check if your code is in the following pattern (that is, the `save_checkpoint` only ran on host `0`):\\n\",\n \"\\n\",\n \"```\\n\",\n \"if jax.process_index() == 0:\\n\",\n \" flax.training.checkpoints.save_checkpoint(...)\\n\",\n \"```\\n\",\n \"\\n\",\n \"Unfortunately this is a legacy pattern that will be deprecated and won't be supported, because in a multi-process environment, the checkpointing code should coordinate among hosts instead of being triggered only on the host `0`. Replacing the code above with the following should resolve the hang issue:\\n\",\n \"\\n\",\n \"```\\n\",\n \"flax.training.checkpoints.save_checkpoint_multiprocess(...)\\n\",\n \"```\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"70e0ebb3\",\n \"metadata\": {},\n \"source\": [\n \"### If you don't save pytrees\\n\",\n \"\\n\",\n \"Orbax uses `orbax.checkpoint.PyTreeCheckpointHandler` to save checkpoints, which means they only save pytrees.\\n\",\n \"\\n\",\n \"If you want to save singular arrays or numbers, you have two options:\\n\",\n \"\\n\",\n \"1. Use `orbax.ArrayCheckpointHandler` to save them following [this migration section](https://flax.readthedocs.io/en/latest/guides/converting_and_upgrading/orbax_upgrade_guide.html#saving-loading-a-single-jax-or-numpy-array).\\n\",\n \"\\n\",\n \"1. Wrap it inside a pytree and save as usual.\"\n ]\n }\n ],\n \"metadata\": {\n \"gpuClass\": \"standard\",\n \"jupytext\": {\n \"formats\": \"ipynb,md\",\n \"main_language\": \"python\"\n },\n \"language_info\": {\n \"codemirror_mode\": {\n \"name\": \"ipython\",\n \"version\": 3\n },\n \"file_extension\": \".py\",\n \"mimetype\": \"text/x-python\",\n \"name\": \"python\",\n \"nbconvert_exporter\": \"python\",\n \"pygments_lexer\": \"ipython3\",\n \"version\": \"3.10.11\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 5\n}\n" }, { "path": "docs/guides/training_techniques/use_checkpointing.md", "content": "---\njupyter:\n jupytext:\n formats: ipynb,md\n main_language: python\n text_representation:\n extension: .md\n format_name: markdown\n format_version: '1.3'\n jupytext_version: 1.13.8\n---\n\n# Save and load checkpoints\n\nThis guide demonstrates how to save and load Flax checkpoints with [Orbax](https://github.com/google/orbax).\n\nOrbax provides a variety of features for saving and loading model data, which you will learn about in this doc:\n\n* Support for various array types and storage formats\n* Asynchronous saving to reduce training wait time\n* Versioning and automatic bookkeeping of past checkpoints\n* Flexible [`transformations`](https://orbax.readthedocs.io/en/latest/transformations.html) to tweak and load old checkpoints\n* [`jax.sharding`](https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html)-based API to save and load in multi-host scenarios\n\n---\n**_Ongoing migration to Orbax:_**\n\nAfter July 30 2023, Flax's legacy `flax.training.checkpoints` API will be deprecated in favor of [Orbax](https://github.com/google/orbax).\n\n* **If you are a new Flax user**: Use the new `orbax.checkpoint` API, as demonstrated in this guide.\n\n* **If you have legacy `flax.training.checkpoints` code in your project**: Consider the following options:\n\n * **Migrating your code to Orbax (Recommended)**: Migrate your API calls to `orbax.checkpoint` API by following this [migration guide](https://flax.readthedocs.io/en/latest/guides/converting_and_upgrading/orbax_upgrade_guide.html).\n\n * **Automatically use the Orbax backend**: Add `flax.config.update('flax_use_orbax_checkpointing', True)` to your project, which will let your `flax.training.checkpoints` calls automatically use the Orbax backend to save your checkpoints.\n\n * **Scheduled flip**: This will become the default mode after **May 2023** (tentative date).\n\n * Visit [Orbax-as-backend troubleshooting section](https://flax.readthedocs.io/en/latest/guides/training_techniques/use_checkpointing.html#orbax-as-backend-troubleshooting) if you meet any issue in the automatic migration.\n---\n\nFor backward-compatibility, this guide shows the Orbax-equivalent calls in the Flax legacy `flax.training.checkpoints` API.\n\nIf you need to learn more about `orbax.checkpoint`, refer to the [Orbax docs](https://orbax.readthedocs.io/en/latest/).\n\n\n\n## Setup\n\nInstall/upgrade Flax and [Orbax](https://github.com/google/orbax). For JAX installation with GPU/TPU support, visit [this section on GitHub](https://github.com/jax-ml/jax#installation).\n\n\nNote: Before running `import jax`, create eight fake devices to mimic a [multi-host environment](https://jax.readthedocs.io/en/latest/jax-101/06-parallelism.html?#aside-hosts-and-devices-in-jax) in this notebook. Note that the order of imports is important here. The `os.environ[\"XLA_FLAGS\"] = '--xla_force_host_platform_device_count=8'` command works only with the CPU backend, which means it won't work with GPU/TPU acceleration on if you're running this notebook in Google Colab. If you are already running the code on multiple devices (for example, in a 4x2 TPU environment), you can skip running the next cell.\n\n```python\nimport os\nos.environ[\"XLA_FLAGS\"] = '--xla_force_host_platform_device_count=8'\n```\n\n```python\nfrom typing import Optional, Any\nimport shutil\n\nimport numpy as np\nimport jax\nfrom jax import random, numpy as jnp\n\nimport flax\nfrom flax import linen as nn\nfrom flax.training import checkpoints, train_state\nfrom flax import struct, serialization\nimport orbax.checkpoint\n\nimport optax\n```\n\n```python\nckpt_dir = '/tmp/flax_ckpt'\n\nif os.path.exists(ckpt_dir):\n shutil.rmtree(ckpt_dir) # Remove any existing checkpoints from the last notebook run.\n```\n\n## Save checkpoints\n\nIn Orbax and Flax, you can save and load any given JAX [pytree](https://jax.readthedocs.io/en/latest/pytrees.html). This includes not only typical Python and NumPy containers, but also customized classes extended from [`flax.struct.dataclass`](https://flax.readthedocs.io/en/latest/api_reference/flax.struct.html#flax.struct.dataclass). That means you can store almost any data generated — not only your model parameters, but any arrays/dictionaries, metadata/configs, and so on.\n\nFirst, create a pytree with many data structures and containers, and play with it:\n\n```python outputId=\"f1856d96-1961-48ed-bb7c-cb63fbaa7567\"\n# A simple model with one linear layer.\nkey1, key2 = random.split(random.key(0))\nx1 = random.normal(key1, (5,)) # A simple JAX array.\nmodel = nn.Dense(features=3)\nvariables = model.init(key2, x1)\n\n# Flax's TrainState is a pytree dataclass and is supported in checkpointing.\n# Define your class with `@flax.struct.dataclass` decorator to make it compatible.\ntx = optax.sgd(learning_rate=0.001) # An Optax SGD optimizer.\nstate = train_state.TrainState.create(\n apply_fn=model.apply,\n params=variables['params'],\n tx=tx)\n# Perform a simple gradient update similar to the one during a normal training workflow.\nstate = state.apply_gradients(grads=jax.tree_util.tree_map(jnp.ones_like, state.params))\n\n# Some arbitrary nested pytree with a dictionary and a NumPy array.\nconfig = {'dimensions': np.array([5, 3])}\n\n# Bundle everything together.\nckpt = {'model': state, 'config': config, 'data': [x1]}\nckpt\n```\n\n### With Orbax\n\n\nSave the checkpoint with `orbax.checkpoint.PyTreeCheckpointer`, directly to the `tmp/orbax/single_save` directory.\n\nNote: An optional `save_args` is provided. This is recommended for performance speedups, as it bundles smaller arrays in your pytree to a single large file instead of multiple smaller files.\n\n```python\nfrom flax.training import orbax_utils\n\norbax_checkpointer = orbax.checkpoint.PyTreeCheckpointer()\nsave_args = orbax_utils.save_args_from_target(ckpt)\norbax_checkpointer.save('/tmp/flax_ckpt/orbax/single_save', ckpt, save_args=save_args)\n```\n\nNext, to use versioning and automatic bookkeeping features, you need to wrap `orbax.checkpoint.CheckpointManager` over `orbax.checkpoint.PyTreeCheckpointer`.\n\nIn addition, provide `orbax.checkpoint.CheckpointManagerOptions` that customizes your needs, such as how often and on what criteria you prefer old checkpoints be deleted. See [documentation](https://orbax.readthedocs.io/en/latest/guides/checkpoint/api_refactor.html) for a full list of options offered.\n\n`orbax.checkpoint.CheckpointManager` should be placed at the top-level outside your training steps to manage your saves.\n\n```python outputId=\"b7132933-566d-440d-c34e-c5468d87cbdc\"\noptions = orbax.checkpoint.CheckpointManagerOptions(max_to_keep=2, create=True)\ncheckpoint_manager = orbax.checkpoint.CheckpointManager(\n '/tmp/flax_ckpt/orbax/managed', orbax_checkpointer, options)\n\n# Inside a training loop\nfor step in range(5):\n # ... do your training\n checkpoint_manager.save(step, ckpt, save_kwargs={'save_args': save_args})\n\nos.listdir('/tmp/flax_ckpt/orbax/managed') # Because max_to_keep=2, only step 3 and 4 are retained\n```\n\n### With the legacy API\n\nAnd here's how to save with the legacy Flax checkpointing utilities (note that this provides less management features compared with `orbax.checkpoint.CheckpointManagerOptions`):\n\n```python outputId=\"6d849273-15ce-4480-8864-726d1838ac1f\"\n# Import Flax Checkpoints.\nfrom flax.training import checkpoints\n\ncheckpoints.save_checkpoint(ckpt_dir='/tmp/flax_ckpt/flax-checkpointing',\n target=ckpt,\n step=0,\n overwrite=True,\n keep=2)\n```\n\n## Restore checkpoints\n\n### With Orbax\n\nIn Orbax, call `.restore()` for either `orbax.checkpoint.PyTreeCheckpointer` or `orbax.checkpoint.CheckpointManager` to restore your checkpoint in the raw pytree format.\n\n```python outputId=\"b4af1ef4-f22f-459b-bdca-2e6bfa16c08b\"\nraw_restored = orbax_checkpointer.restore('/tmp/flax_ckpt/orbax/single_save')\nraw_restored\n```\n\nNote that the `step` number is required for `CheckpointManger`. You can also use `.latest_step()` to find the latest step available.\n\n```python\nstep = checkpoint_manager.latest_step() # step = 4\ncheckpoint_manager.restore(step)\n```\n\n### With the legacy API\n\nNote that with the migration to Orbax in progress, `flax.training.checkpointing.restore_checkpoint` can automatically identify whether a checkpoint is saved in the legacy Flax format or with an Orbax backend, and restore the pytree correctly. Therefore, adding `flax.config.update('flax_use_orbax_checkpointing', True)` won't hurt your ability to restore old checkpoints.\n\nHere's how to restore checkpoints using the legacy API:\n\n```python outputId=\"85ffceca-f38d-46b8-e567-d9d38b7885f9\"\nraw_restored = checkpoints.restore_checkpoint(ckpt_dir='/tmp/flax_ckpt/flax-checkpointing', target=None)\nraw_restored\n```\n\n## Restore with custom dataclasses\n\n### With Orbax\n\n* The pytrees restored in the previous examples are in the form of raw dictionaries. Original pytrees contain custom dataclasses like [`TrainState`](https://flax.readthedocs.io/en/latest/flip/1009-optimizer-api.html?#train-state) and `optax` states.\n* This is because when restoring a pytree, the program does not yet know which structure it once belonged to.\n* To resolve this, you should first provide an example pytree to let Orbax or Flax know exactly which structure to restore to.\n\nThis section demonstrates how to set up any custom Flax dataclass explicitly, and have the same structure as a saved checkpoint.\n\nNote: Data that was a JAX NumPy array (`jnp.array`) format will be restored as a NumPy array (`numpy.array`). This would not affect your work because JAX will [automatically convert](https://jax.readthedocs.io/en/latest/jax-101/01-jax-basics.html) NumPy arrays to JAX arrays once the computation starts.\n\n```python outputId=\"110c6b6e-fe42-4179-e5d8-6b92d355e11b\"\nempty_state = train_state.TrainState.create(\n apply_fn=model.apply,\n params=jax.tree_util.tree_map(np.zeros_like, variables['params']), # values of the tree leaf doesn't matter\n tx=tx,\n)\nempty_config = {'dimensions': np.array([0, 0])}\ntarget = {'model': empty_state, 'config': empty_config, 'data': [jnp.zeros_like(x1)]}\nstate_restored = orbax_checkpointer.restore('/tmp/flax_ckpt/orbax/single_save', item=target)\nstate_restored\n```\n\n### With the legacy API\n\nAlternatively, you can restore from Orbax `CheckpointManager` and from the legacy Flax code as follows:\n\n```python\ncheckpoint_manager.restore(4, items=target)\n```\n\n```python\ncheckpoints.restore_checkpoint(ckpt_dir='/tmp/flax_ckpt/flax-checkpointing', target=target)\n```\n\nIt's often recommended to refactor out the process of initializing a checkpoint's structure (for example, a [`TrainState`](https://flax.readthedocs.io/en/latest/flip/1009-optimizer-api.html?#train-state)), so that saving/loading is easier and less error-prone. This is because functions and complex objects like `apply_fn` and `tx` (optimizer) cannot be serialized into the checkpoint file and must be initialized by code.\n\n\n## Restore when checkpoint structures differ\n\nDuring your development, your checkpoint structure will change when changing the model, adding/removing fields during tweaking, and so on.\n\nThis section explains how to load old data to your new code.\n\nBelow is a simple example — a `CustomTrainState` extended from `flax.training.train_state.TrainState` that contains an extra field called `batch_stats`. When working on a real-world model, you may need this when applying [batch normalization](https://flax.readthedocs.io/en/latest/guides/training_techniques/batch_norm.html).\n\nHere, you store the new `CustomTrainState` as step 5, while step 4 contains the old/previous `TrainState`.\n\n```python outputId=\"4fe776f0-65f8-4fc4-d64a-990520b36dce\"\nclass CustomTrainState(train_state.TrainState):\n batch_stats: Any = None\n\ncustom_state = CustomTrainState.create(\n apply_fn=state.apply_fn,\n params=state.params,\n tx=state.tx,\n batch_stats=np.arange(10),\n)\n\ncustom_ckpt = {'model': custom_state, 'config': config, 'data': [x1]}\n# Use a custom state to read the old `TrainState` checkpoint.\ncustom_target = {'model': custom_state, 'config': None, 'data': [jnp.zeros_like(x1)]}\n\n# Save it in Orbax.\ncustom_save_args = orbax_utils.save_args_from_target(custom_ckpt)\ncheckpoint_manager.save(5, custom_ckpt, save_kwargs={'save_args': custom_save_args})\n```\n\nIt is recommended to keep your checkpoints up-to-date with your pytree dataclass definitions. However, you might be forced to restore the checkpoints with incompatible reference objects at runtime. When this happens, the checkpoint restoration will try to respect the structure of the reference when given.\n\nBelow are examples of a few common scenarios.\n\n\n### Scenario 1: When a reference object is partial\n\nIf your reference object is a subtree of your checkpoint, the restoration will ignore the additional field(s) and restore a checkpoint with the same structure as the reference.\n\nLike in the example below, the `batch_stats` field in `CustomTrainState` was ignored, and the checkpoint was restored as a `TrainState`.\n\nThis can also be useful for reading only part of your checkpoint.\n\n```python\nrestored = checkpoint_manager.restore(5, items=target)\nassert not hasattr(restored, 'batch_stats')\nassert type(restored['model']) == train_state.TrainState\nrestored\n```\n\n### Scenario 2: When a checkpoint is partial\n\nOn the other hand, if the reference object contains a value that is not available in the checkpoint, the checkpointing code will by default warn that some data is not compatible.\n\nTo bypass the error, you need to pass an Orbax [`transform`](https://orbax.readthedocs.io/en/latest/guides/checkpoint/transformations.html) that teaches Orbax how to conform this checkpoint into the structure of the `custom_target`.\n\nIn this case, pass a default `{}` that lets Orbax use values in the `custom_target` to fill in the blank. This allows you to restore an old checkpoint into a new data structure, the `CustomTrainState`.\n\n```python\ntry:\n checkpoint_manager.restore(4, items=custom_target)\nexcept KeyError as e:\n print(f'KeyError when target state has an unmentioned field: {e}')\n print('')\n\n# Step 4 is an original `TrainState`, without the `batch_stats`\ncustom_restore_args = orbax_utils.restore_args_from_target(custom_target)\nrestored = checkpoint_manager.restore(4, items=custom_target,\n restore_kwargs={'transforms': {}, 'restore_args': custom_restore_args})\nassert type(restored['model']) == CustomTrainState\nnp.testing.assert_equal(restored['model'].batch_stats,\n custom_target['model'].batch_stats)\nrestored\n```\n\n##### With Orbax\n\nIf you have already saved your checkpoints with the Orbax backend, you can use `orbax_transforms` to access this `transforms` argument in the Flax API.\n\n```python outputId=\"cdbb9247-d1eb-4458-aa83-8db0332af7cb\"\n# Save in the \"Flax-with-Orbax\" backend.\nflax.config.update('flax_use_orbax_checkpointing', True)\ncheckpoints.save_checkpoint(ckpt_dir='/tmp/flax_ckpt/flax-checkpointing',\n target=ckpt,\n step=4,\n overwrite=True,\n keep=2)\n\ncheckpoints.restore_checkpoint('/tmp/flax_ckpt/flax-checkpointing', target=custom_target, step=4,\n orbax_transforms={})\n```\n\n##### With the legacy API\n\nUsing the legacy `flax.training.checkpoints` API, similar things are doable too, but they are not as flexible as the [Orbax Transformations](https://orbax.readthedocs.io/en/latest/guides/checkpoint/transformations.html).\n\nYou need to restore the checkpoint to a raw dict with `target=None`, modify the structure accordingly, and then deserialize it back to the original target.\n\n```python\n# Save using the legacy Flax `checkpoints` API.\nflax.config.update('flax_use_orbax_checkpointing', False)\ncheckpoints.save_checkpoint(ckpt_dir='/tmp/flax_ckpt/flax-checkpointing',\n target=ckpt,\n step=5,\n overwrite=True,\n keep=2)\n\n# Pass no target to get a raw state dictionary first.\nraw_state_dict = checkpoints.restore_checkpoint('/tmp/flax_ckpt/flax-checkpointing', target=None, step=5)\n# Add/remove fields as needed.\nraw_state_dict['model']['batch_stats'] = np.flip(np.arange(10))\n# Restore the classes with correct target now\nflax.serialization.from_state_dict(custom_target, raw_state_dict)\n```\n\n## Asynchronized checkpointing\n\nCheckpointing is I/O heavy, and if you have a large amount of data to save, it may be worthwhile to put it into a background thread, while continuing with your training.\n\nYou can do this by creating an [`orbax.checkpoint.AsyncCheckpointer`](https://github.com/google/orbax/blob/main/checkpoint/orbax/checkpoint/async_checkpointer.py) in place of the `orbax.checkpoint.PyTreeCheckpointer`.\n\nNote: You should use the same `async_checkpointer` to handle all your async saves across your training steps, so that it can make sure that a previous async save is done before the next one begins. This enables bookkeeping, such as `keep` (the number of checkpoints) and `overwrite` to be consistent across steps.\n\nWhenever you want to explicitly wait until an async save is done, you can call `async_checkpointer.wait_until_finished()`.\n\n```python outputId=\"aefce94c-8bae-4355-c142-05f2b61c39e2\"\n# `orbax.checkpoint.AsyncCheckpointer` needs some multi-process initialization, because it was\n# originally designed for multi-process large model checkpointing.\n# For Python notebooks or other single-process settings, just set up with `num_processes=1`.\n# Refer to https://jax.readthedocs.io/en/latest/multi_process.html#initializing-the-cluster\n# for how to set it up in multi-process scenarios.\njax.distributed.initialize(\"localhost:8889\", num_processes=1, process_id=0)\n\nasync_checkpointer = orbax.checkpoint.AsyncCheckpointer(\n orbax.checkpoint.PyTreeCheckpointHandler(), timeout_secs=50)\n\n# Save your job:\nasync_checkpointer.save('/tmp/flax_ckpt/orbax/single_save_async', ckpt, save_args=save_args)\n# ... Continue with your work...\n\n# ... Until a time when you want to wait until the save completes:\nasync_checkpointer.wait_until_finished() # Blocks until the checkpoint saving is completed.\nasync_checkpointer.restore('/tmp/flax_ckpt/orbax/single_save_async', item=target)\n```\n\nIf you are using Orbax `CheckpointManager`, just pass in the async_checkpointer when initializing it. Then, in practice, call `async_checkpoint_manager.wait_until_finished()` instead.\n\n```python\nasync_checkpoint_manager = orbax.checkpoint.CheckpointManager(\n '/tmp/flax_ckpt/orbax/managed_async', async_checkpointer, options)\nasync_checkpoint_manager.wait_until_finished()\n```\n\n## Multi-host/multi-process checkpointing\n\nJAX provides a few ways to scale up your code on multiple hosts at the same time. This usually happens when the number of devices (CPU/GPU/TPU) is so large that different devices are managed by different hosts (CPU). To get started on JAX in multi-process settings, check out [Using JAX in multi-host and multi-process environments](https://jax.readthedocs.io/en/latest/multi_process.html) and the [distributed array guide](https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html).\n\nIn the [Single Program Multi Data (SPMD)](https://jax.readthedocs.io/en/latest/glossary.html#term-SPMD) paradigm with JAX `jit`, a large multi-process array can have its data sharded across different devices. (Note that JAX `pjit` and `jit` have been merged into a single unified interface. To learn about compiling and executing JAX functions in multi-host or multi-core environments, refer to [this guide](https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html) and the [jax.Array migration guide](https://jax.readthedocs.io/en/latest/jax_array_migration.html).) When a multi-process array is serialized, each host dumps its data shards to a single shared storage, such as a Google Cloud bucket.\n\nOrbax supports saving and loading pytrees with multi-process arrays in the same fashion as single-process pytrees. However, it's recommended to use the asynchronized [`orbax.AsyncCheckpointer`](https://github.com/google/orbax/blob/main/checkpoint/orbax/checkpoint/async_checkpointer.py) to save large multi-process arrays on another thread, so that you can perform computation alongside the saves. With pure Orbax, saving checkpoints in a multi-process context uses the same API as in a single-process context.\n\n```python\nfrom jax.sharding import PartitionSpec, NamedSharding\n\n# Create an array sharded across multiple devices.\nmesh_shape = (4, 2)\ndevices = np.asarray(jax.devices()).reshape(*mesh_shape)\nmesh = jax.sharding.Mesh(devices, ('x', 'y'))\n\nmp_array = jax.device_put(np.arange(8 * 2).reshape(8, 2),\n NamedSharding(mesh, PartitionSpec('x', 'y')))\n\n# Make it a pytree.\nmp_ckpt = {'model': mp_array}\n```\n\n```python\nasync_checkpoint_manager.save(0, mp_ckpt)\nasync_checkpoint_manager.wait_until_finished()\n```\n\nWhen restoring a checkpoint with multi-process arrays, you need to specify what `sharding` each array should be restored back to. Otherwise, they will be restored as large `np.array`s on process 0, costing time and memory.\n\n(In this notebook, since we are on single-process, it will be restored as `np.array` even if we provide shardings.)\n\n### With Orbax\n\nOrbax allows you to specify this by passing a pytree of `sharding`s in `restore_args`. If you already have a reference pytree that has all the arrays with the right sharding, you can use `orbax_utils.restore_args_from_target` to transform it into the `restore_args` that Orbax needs.\n\n```python\n# The reference doesn't need to be as large as your checkpoint!\n# Just make sure it has the `.sharding` you want.\nmp_smaller = jax.device_put(np.arange(8).reshape(4, 2),\n NamedSharding(mesh, PartitionSpec('x', 'y')))\nref_ckpt = {'model': mp_smaller}\n\nrestore_args = orbax_utils.restore_args_from_target(ref_ckpt)\nasync_checkpoint_manager.restore(\n 0, items=ref_ckpt, restore_kwargs={'restore_args': restore_args})\n```\n\n### With the legacy Flax: use `save_checkpoint_multiprocess`\n\nIn legacy Flax, to save multi-process arrays, use [`flax.training.checkpoints.save_checkpoint_multiprocess()`](https://flax.readthedocs.io/en/latest/api_reference/flax.training.html#flax.training.checkpoints.save_checkpoint_multiprocess) in place of `save_checkpoint()` and with the same arguments.\n\nIf your checkpoint is too large, you can specify `timeout_secs` in the manager and give it more time to finish writing.\n\n```python outputId=\"901bb097-0899-479d-b9ae-61dae79e7057\"\nasync_checkpointer = orbax.checkpoint.AsyncCheckpointer(orbax.checkpoint.PyTreeCheckpointHandler(), timeout_secs=50)\ncheckpoints.save_checkpoint_multiprocess(ckpt_dir,\n mp_ckpt,\n step=3,\n overwrite=True,\n keep=4,\n orbax_checkpointer=async_checkpointer)\n```\n\n```python outputId=\"393c4a0e-8a8c-4ca6-c609-93c8bab38e75\"\nmp_restored = checkpoints.restore_checkpoint(ckpt_dir,\n target=ref_ckpt,\n step=3,\n orbax_checkpointer=async_checkpointer)\nmp_restored\n```\n\n## Orbax-as-backend troubleshooting\n\nAs an intermediate stage of the migration (to Orbax from the legacy Flax `checkpoints` API), `flax.training.checkpoints` APIs will start to use Orbax as their backend when saving checkpoints starting from May 15, 2023.\n\nCheckpoints saved with the Orbax backend can be readable by either `flax.training.checkpoints.restore_checkpoint` or `orbax.checkpoint.PyTreeCheckpointer`.\n\nCode-wise, this is equivalent to setting the config flag [`flax.config.flax_use_orbax_checkpointing`](https://github.com/google/flax/blob/main/flax/configurations.py#L103) default to `True`. You can overwrite this value in your project with `flax.config.update('flax_use_orbax_checkpointing', )` at any time.\n\nIn general, this automatic migration will not affect most users. However, you may encounter issues if your API usage follows some specific pattern. Check out the sections below for troubleshooting.\n\n\n### If your devices hang when writing checkpoints\n\nIf you are running in a multi-host environment (usually anything larger than 8 TPU devices) and your devices hang when writing checkpoints, check if your code is in the following pattern (that is, the `save_checkpoint` only ran on host `0`):\n\n```\nif jax.process_index() == 0:\n flax.training.checkpoints.save_checkpoint(...)\n```\n\nUnfortunately this is a legacy pattern that will be deprecated and won't be supported, because in a multi-process environment, the checkpointing code should coordinate among hosts instead of being triggered only on the host `0`. Replacing the code above with the following should resolve the hang issue:\n\n```\nflax.training.checkpoints.save_checkpoint_multiprocess(...)\n```\n\n\n### If you don't save pytrees\n\nOrbax uses `orbax.checkpoint.PyTreeCheckpointHandler` to save checkpoints, which means they only save pytrees.\n\nIf you want to save singular arrays or numbers, you have two options:\n\n1. Use `orbax.ArrayCheckpointHandler` to save them following [this migration section](https://flax.readthedocs.io/en/latest/guides/converting_and_upgrading/orbax_upgrade_guide.html#saving-loading-a-single-jax-or-numpy-array).\n\n1. Wrap it inside a pytree and save as usual.\n" }, { "path": "docs/index.rst", "content": ".. Flax documentation main file, created by\n sphinx-quickstart on Mon Feb 17 11:41:38 2020.\n You can adapt this file completely to your liking, but it should at least\n contain the root `toctree` directive.\n\n******************************\nFlax Linen\n******************************\n\n\n.. div:: sd-text-left sd-font-italic\n\n Neural networks with JAX\n\n\n----\n\nFlax Linen delivers an **end-to-end and flexible user experience for researchers\nwho use JAX with neural networks**. Flax\nexposes the full power of `JAX `__. It is made up of\nloosely coupled libraries, which are showcased with end-to-end integrated\n`guides `__\nand `examples `__.\n\nFlax Linen is used by\n`hundreds of projects (and growing) `__,\nboth in the open source community\n(like `Hugging Face `__)\nand at Google\n(like\n`Gemini `__,\n`Imagen `__,\n`Scenic `__,\nand `Big Vision `__).\n\n\nFeatures\n^^^^^^^^^\n\n.. grid::\n\n .. grid-item::\n :columns: 12 12 12 6\n\n .. card:: Safety\n :class-card: sd-border-0\n :shadow: none\n :class-title: sd-fs-5\n\n .. div:: sd-font-normal\n\n Flax is designed for correctness and safety. Thanks to its immutable Modules\n and Functional API, Flax helps mitigate bugs that arise when handling state\n in JAX.\n\n .. grid-item::\n :columns: 12 12 12 6\n\n .. card:: Control\n :class-card: sd-border-0\n :shadow: none\n :class-title: sd-fs-5\n\n .. div:: sd-font-normal\n\n Flax grants more fine-grained control and expressivity than most Neural Network\n frameworks via its Variable Collections, RNG Collections and Mutability conditions.\n\n .. grid-item::\n :columns: 12 12 12 6\n\n .. card:: Functional API\n :class-card: sd-border-0\n :shadow: none\n :class-title: sd-fs-5\n\n .. div:: sd-font-normal\n\n Flax's functional API radically redefines what Modules can do via lifted transformations like vmap, scan, etc, while also enabling seamless integration with other JAX libraries like Optax and Chex.\n\n .. grid-item::\n :columns: 12 12 12 6\n\n .. card:: Terse code\n :class-card: sd-border-0\n :shadow: none\n :class-title: sd-fs-5\n\n .. div:: sd-font-normal\n\n Flax's :meth:`compact ` Modules enables submodules to be defined directly at their callsite, leading to code that is easier to read and avoids repetition.\n\n\n----\n\nInstallation\n^^^^^^^^^^^^\n\n.. code-block:: bash\n\n pip install flax\n # or to install the latest version of Flax:\n pip install --upgrade git+https://github.com/google/flax.git\n\nFlax installs the vanilla CPU version of JAX, if you need a custom version please check out `JAX's installation page `__.\n\nBasic usage\n^^^^^^^^^^^^\n\n.. testsetup::\n\n import jax\n from jax import random\n import flax.linen as nn\n import jax.numpy as jnp\n\n.. testcode::\n\n class MLP(nn.Module): # create a Flax Module dataclass\n out_dims: int\n\n @nn.compact\n def __call__(self, x):\n x = x.reshape((x.shape[0], -1))\n x = nn.Dense(128)(x) # create inline Flax Module submodules\n x = nn.relu(x)\n x = nn.Dense(self.out_dims)(x) # shape inference\n return x\n\n model = MLP(out_dims=10) # instantiate the MLP model\n\n x = jnp.empty((4, 28, 28, 1)) # generate random data\n variables = model.init(random.key(42), x)# initialize the weights\n y = model.apply(variables, x) # make forward pass\n\n----\n\nLearn more\n^^^^^^^^^^\n\n.. grid::\n\n .. grid-item::\n :columns: 6 6 6 4\n\n .. card:: :material-regular:`rocket_launch;2em` Quickstart\n :class-card: sd-text-black sd-bg-light\n :link: quick_start.html\n\n .. grid-item::\n :columns: 6 6 6 4\n\n .. card:: :material-regular:`library_books;2em` Guides\n :class-card: sd-text-black sd-bg-light\n :link: guides/index.html\n\n .. grid-item::\n :columns: 6 6 6 4\n\n .. card:: :material-regular:`science;2em` Examples\n :class-card: sd-text-black sd-bg-light\n :link: examples.html\n\n .. grid-item::\n :columns: 6 6 6 4\n\n .. card:: :material-regular:`import_contacts;2em` Glossary\n :class-card: sd-text-black sd-bg-light\n :link: glossary.html\n\n .. grid-item::\n :columns: 6 6 6 4\n\n .. card:: :material-regular:`settings;2em` Developer notes\n :class-card: sd-text-black sd-bg-light\n :link: developer_notes/index.html\n\n .. grid-item::\n :columns: 6 6 6 4\n\n .. card:: :material-regular:`history_edu;2em` The Flax philosophy\n :class-card: sd-text-black sd-bg-light\n :link: https://flax.readthedocs.io/en/latest/philosophy.html\n\n .. grid-item::\n :columns: 6 6 6 4\n\n .. card:: :material-regular:`menu_book;2em` API reference\n :class-card: sd-text-black sd-bg-light\n :link: api_reference/index.html\n\n----\n\nEcosystem\n^^^^^^^^^\n\nNotable examples in Flax include:\n\n\n.. grid::\n\n .. grid-item::\n :columns: 6 6 6 4\n\n .. card:: `🤗 Hugging Face `__\n :class-card: sd-border-0\n :shadow: none\n :class-title: sd-text-center sd-fs-5\n\n .. div:: sd-text-center sd-font-italic\n\n NLP and computer vision models\n\n .. grid-item::\n :columns: 6 6 6 4\n\n .. card:: `🥑 DALLE Mini `__\n :class-card: sd-border-0\n :shadow: none\n :class-title: sd-text-center sd-fs-5\n\n .. div:: sd-text-center sd-font-italic\n\n Model for text-to-image generation\n\n .. grid-item::\n :columns: 6 6 6 4\n\n .. card:: `PaLM `__\n :class-card: sd-border-0\n :shadow: none\n :class-title: sd-text-center sd-fs-5\n\n .. div:: sd-text-center sd-font-italic\n\n 540-billion parameter model for text generation\n\n .. grid-item::\n :columns: 6 6 6 4\n\n .. card:: `Imagen `__\n :class-card: sd-border-0\n :shadow: none\n :class-title: sd-text-center sd-fs-5\n\n .. div:: sd-text-center sd-font-italic\n\n Text-to-image diffusion models\n\n .. grid-item::\n :columns: 6 6 6 4\n\n .. card:: `Scenic `__\n :class-card: sd-border-0\n :shadow: none\n :class-title: sd-text-center sd-fs-5\n\n .. div:: sd-text-center sd-font-italic\n\n Libraries for large-scale computer vision\n\n .. grid-item::\n :columns: 6 6 6 4\n\n .. card:: `Big Vision `__\n :class-card: sd-border-0\n :shadow: none\n :class-title: sd-text-center sd-fs-5\n\n .. div:: sd-text-center sd-font-italic\n\n Large-scale computer vision models\n\n .. grid-item::\n :columns: 6 6 6 4\n\n .. card:: `MaxText `__\n :class-card: sd-border-0\n :shadow: none\n :class-title: sd-text-center sd-fs-5\n\n .. div:: sd-text-center sd-font-italic\n\n Open source high performance LLM\n\n .. grid-item::\n :columns: 6 6 6 4\n\n .. card:: `T5x `__\n :class-card: sd-border-0\n :shadow: none\n :class-title: sd-text-center sd-fs-5\n\n .. div:: sd-text-center sd-font-italic\n\n Large language models\n\n .. grid-item::\n :columns: 6 6 6 4\n\n .. card:: `Brax `__\n :class-card: sd-border-0\n :shadow: none\n :class-title: sd-text-center sd-fs-5\n\n .. div:: sd-text-center sd-font-italic\n\n On-device differentiable reinforcement learning environments\n\n\n\n.. role:: bold\n :class: bold\n\n.. toctree::\n :hidden:\n :maxdepth: 2\n\n Quick start \n guides/flax_fundamentals/flax_basics\n guides/index\n examples/index\n glossary\n faq\n developer_notes/index\n The Flax philosophy \n How to contribute \n api_reference/index\n Flax NNX \n" }, { "path": "docs/linen_intro.ipynb", "content": "{\n \"cells\": [\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google/flax/blob/main/docs/notebooks/linen_intro.ipynb)\\n\",\n \"[![Open On GitHub](https://img.shields.io/badge/Open-on%20GitHub-blue?logo=GitHub)](https://github.com/google/flax/blob/main/docs/notebooks/linen_intro.ipynb)\\n\",\n \"\\n\",\n \"# Preface\\n\",\n \"\\n\",\n \"
\\n\",\n \"
CAVEAT PROGRAMMER
\\n\",\n \"\\n\",\n \"The below is an alpha API preview and things might break. The surface syntax of the features of the API are not fixed in stone, and we welcome feedback on any points.\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Useful links\\n\",\n \"\\n\",\n \"⟶ [Slides](https://docs.google.com/presentation/d/1ngKWUwsSqAwPRvATG8sAxMzu9ujv4N__cKsUofdNno0/edit?usp=sharing) for the core ideas of the new Functional Core and Linen\\n\",\n \"\\n\",\n \"⟶ \\\"Design tests\\\" guided our design process. Many are available for [functional core](https://github.com/google/flax/tree/main/examples/core_design_test) and some for the [proposed Module abstraction](https://github.com/google/flax/tree/main/examples/linen_design_test/)\\n\",\n \"\\n\",\n \"⟶ Ported examples: [ImageNet](https://github.com/google/flax/tree/main/examples/imagenet) and [WMT](https://github.com/google/flax/tree/main/examples/wmt) (to the proposed Module abstraction). TODO: Port to functional core.\\n\",\n \"\\n\",\n \"⟶ Our new [discussion forums](https://github.com/google/flax/discussions/)\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"# Install and Import\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {\n \"tags\": [\n \"skip-execution\"\n ]\n },\n \"outputs\": [],\n \"source\": [\n \"# Install the newest JAXlib version.\\n\",\n \"!pip install --upgrade -q pip jax jaxlib\\n\",\n \"# Install Flax at head:\\n\",\n \"!pip install --upgrade -q git+https://github.com/google/flax.git\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 2,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"import functools\\n\",\n \"from typing import Any, Callable, Sequence, Optional\\n\",\n \"import jax\\n\",\n \"from jax import lax, random, numpy as jnp\\n\",\n \"import flax\\n\",\n \"from flax import linen as nn\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"# Invoking Modules\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Let's instantiate a `Dense` layer.\\n\",\n \" - Modules are actually objects in this API, so we provide _constructor arguments_ when initializing the Module. In this case, we only have to provide the output `features` dimension.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 3,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"model = nn.Dense(features=3)\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"We need to initialize the Module variables, these include the parameters of the Module as well as any other state variables.\\n\",\n \"\\n\",\n \"We call the `init` method on the instantiated Module. If the Module `__call__` method has args `(self, *args, **kwargs)` then we call `init` with `(rngs, *args, **kwargs)` so in this case, just `(rng, input)`:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 4,\n \"metadata\": {\n \"outputId\": \"3adfaeaf-977e-4e82-8adf-d254fae6eb91\"\n },\n \"outputs\": [\n {\n \"name\": \"stderr\",\n \"output_type\": \"stream\",\n \"text\": [\n \"WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\\n\"\n ]\n },\n {\n \"data\": {\n \"text/plain\": [\n \"FrozenDict({\\n\",\n \" params: {\\n\",\n \" kernel: DeviceArray([[ 0.6503669 , 0.8678979 , 0.46042678],\\n\",\n \" [ 0.05673932, 0.9909285 , -0.63536596],\\n\",\n \" [ 0.76134115, -0.3250529 , -0.6522163 ],\\n\",\n \" [-0.8243032 , 0.4150194 , 0.19405058]], dtype=float32),\\n\",\n \" bias: DeviceArray([0., 0., 0.], dtype=float32),\\n\",\n \" },\\n\",\n \"})\"\n ]\n },\n \"execution_count\": 4,\n \"metadata\": {\n \"tags\": []\n },\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"# Make RNG Keys and a fake input.\\n\",\n \"key1, key2 = random.split(random.key(0), 2)\\n\",\n \"x = random.uniform(key1, (4,4))\\n\",\n \"\\n\",\n \"# provide key and fake input to get initialized variables\\n\",\n \"init_variables = model.init(key2, x)\\n\",\n \"\\n\",\n \"init_variables\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"We call the `apply` method on the instantiated Module. If the Module `__call__` method has args `(self, *args, **kwargs)` then we call `apply` with `(variables, *args, rngs=, mutable=, **kwargs)` where\\n\",\n \" - `` are the optional _call time_ RNGs for things like dropout. For simple Modules this is just a single key, but if your module has multiple __kinds__ of data, it's a dictionary of rng-keys per-kind, e.g. `{'params': key0, 'dropout': key1}` for a Module with dropout layers.\\n\",\n \" - `` is an optional list of names of __kinds__ that are expected to be mutated during the call. e.g. `['batch_stats']` for a layer updating batchnorm statistics.\\n\",\n \"\\n\",\n \"So in this case, just `(variables, input)`:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 5,\n \"metadata\": {\n \"outputId\": \"e8c389a6-29f3-4f93-97ea-703e85a8b811\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"DeviceArray([[ 0.5035518 , 1.8548559 , -0.4270196 ],\\n\",\n \" [ 0.0279097 , 0.5589246 , -0.43061775],\\n\",\n \" [ 0.35471284, 1.5741 , -0.3286552 ],\\n\",\n \" [ 0.5264864 , 1.2928858 , 0.10089308]], dtype=float32)\"\n ]\n },\n \"execution_count\": 5,\n \"metadata\": {\n \"tags\": []\n },\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"y = model.apply(init_variables, x)\\n\",\n \"y\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Additional points:\\n\",\n \" - If you want to `init` or `apply` a Module using a method other than call, you need to provide the `method=` kwarg to `init` and `apply` to use it instead of the default `__call__`, e.g. `method='encode'`, `method='decode'` to apply the encode/decode methods of an autoencoder.\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"# Defining Basic Modules\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Composing submodules\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"We support declaring modules in `setup()` that can still benefit from shape inference by using __Lazy Initialization__ that sets up variables the first time the Module is called.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 6,\n \"metadata\": {\n \"outputId\": \"1a6c6a17-0b95-42c2-b5bf-b9ad80fd7758\",\n \"tags\": []\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"initialized parameter shapes:\\n\",\n \" {'params': {'layers_0': {'bias': (3,), 'kernel': (4, 3)}, 'layers_1': {'bias': (4,), 'kernel': (3, 4)}, 'layers_2': {'bias': (5,), 'kernel': (4, 5)}}}\\n\",\n \"output:\\n\",\n \" [[ 4.2292815e-02 -4.3807115e-02 2.9323792e-02 6.5492536e-03\\n\",\n \" -1.7147182e-02]\\n\",\n \" [ 1.2967804e-01 -1.4551792e-01 9.4432175e-02 1.2521386e-02\\n\",\n \" -4.5417294e-02]\\n\",\n \" [ 0.0000000e+00 0.0000000e+00 0.0000000e+00 0.0000000e+00\\n\",\n \" 0.0000000e+00]\\n\",\n \" [ 9.3024090e-04 2.7864411e-05 2.4478839e-04 8.1344356e-04\\n\",\n \" -1.0110775e-03]]\\n\"\n ]\n }\n ],\n \"source\": [\n \"class ExplicitMLP(nn.Module):\\n\",\n \" features: Sequence[int]\\n\",\n \"\\n\",\n \" def setup(self):\\n\",\n \" # we automatically know what to do with lists, dicts of submodules\\n\",\n \" self.layers = [nn.Dense(feat) for feat in self.features]\\n\",\n \" # for single submodules, we would just write:\\n\",\n \" # self.layer1 = nn.Dense(feat1)\\n\",\n \"\\n\",\n \" def __call__(self, inputs):\\n\",\n \" x = inputs\\n\",\n \" for i, lyr in enumerate(self.layers):\\n\",\n \" x = lyr(x)\\n\",\n \" if i != len(self.layers) - 1:\\n\",\n \" x = nn.relu(x)\\n\",\n \" return x\\n\",\n \"\\n\",\n \"key1, key2 = random.split(random.key(0), 2)\\n\",\n \"x = random.uniform(key1, (4,4))\\n\",\n \"\\n\",\n \"model = ExplicitMLP(features=[3,4,5])\\n\",\n \"init_variables = model.init(key2, x)\\n\",\n \"y = model.apply(init_variables, x)\\n\",\n \"\\n\",\n \"print('initialized parameter shapes:\\\\n', jax.tree_util.tree_map(jnp.shape, flax.core.unfreeze(init_variables)))\\n\",\n \"print('output:\\\\n', y)\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Here we show the equivalent compact form of the MLP that declares the submodules inline using the `@compact` decorator.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 7,\n \"metadata\": {\n \"outputId\": \"b3709789-e66e-4e20-f6b2-04022f8a62bb\",\n \"tags\": []\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"initialized parameter shapes:\\n\",\n \" {'params': {'layers_0': {'bias': (3,), 'kernel': (4, 3)}, 'layers_1': {'bias': (4,), 'kernel': (3, 4)}, 'layers_2': {'bias': (5,), 'kernel': (4, 5)}}}\\n\",\n \"output:\\n\",\n \" [[ 4.2292815e-02 -4.3807115e-02 2.9323792e-02 6.5492536e-03\\n\",\n \" -1.7147182e-02]\\n\",\n \" [ 1.2967804e-01 -1.4551792e-01 9.4432175e-02 1.2521386e-02\\n\",\n \" -4.5417294e-02]\\n\",\n \" [ 0.0000000e+00 0.0000000e+00 0.0000000e+00 0.0000000e+00\\n\",\n \" 0.0000000e+00]\\n\",\n \" [ 9.3024090e-04 2.7864411e-05 2.4478839e-04 8.1344356e-04\\n\",\n \" -1.0110775e-03]]\\n\"\n ]\n }\n ],\n \"source\": [\n \"class SimpleMLP(nn.Module):\\n\",\n \" features: Sequence[int]\\n\",\n \"\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self, inputs):\\n\",\n \" x = inputs\\n\",\n \" for i, feat in enumerate(self.features):\\n\",\n \" x = nn.Dense(feat, name=f'layers_{i}')(x)\\n\",\n \" if i != len(self.features) - 1:\\n\",\n \" x = nn.relu(x)\\n\",\n \" # providing a name is optional though!\\n\",\n \" # the default autonames would be \\\"Dense_0\\\", \\\"Dense_1\\\", ...\\n\",\n \" # x = nn.Dense(feat)(x)\\n\",\n \" return x\\n\",\n \"\\n\",\n \"key1, key2 = random.split(random.key(0), 2)\\n\",\n \"x = random.uniform(key1, (4,4))\\n\",\n \"\\n\",\n \"model = SimpleMLP(features=[3,4,5])\\n\",\n \"init_variables = model.init(key2, x)\\n\",\n \"y = model.apply(init_variables, x)\\n\",\n \"\\n\",\n \"print('initialized parameter shapes:\\\\n', jax.tree_util.tree_map(jnp.shape, flax.core.unfreeze(init_variables)))\\n\",\n \"print('output:\\\\n', y)\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Declaring and using variables\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Flax uses lazy initialization, which allows declared variables to be initialized only at the first site of their use, using whatever shape information is available a the local call site for shape inference. Once a variable has been initialized, a reference to the data is kept for use in subsequent calls.\\n\",\n \"\\n\",\n \"For declaring parameters that aren't mutated inside the model, but rather by gradient descent, we use the syntax:\\n\",\n \"\\n\",\n \" `self.param(parameter_name, parameter_init_fn, *init_args, **init_kwargs)`\\n\",\n \"\\n\",\n \"with arguments:\\n\",\n \" - `parameter_name` just the name, a string\\n\",\n \" - `parameter_init_fn` a function taking an RNG key and a variable number of other arguments, i.e. `fn(rng, *args)`. typically those in `nn.initializers` take an `rng` and a `shape` argument.\\n\",\n \" - the remaining arguments to feed to the init function when initializing.\\n\",\n \"\\n\",\n \"Again, we'll demonstrate declaring things inline as we typically do using the `@compact` decorator.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 8,\n \"metadata\": {\n \"outputId\": \"bc5cb1f2-c5e9-4159-d131-73247009e32f\",\n \"tags\": []\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"initialized parameters:\\n\",\n \" FrozenDict({\\n\",\n \" params: {\\n\",\n \" kernel: DeviceArray([[ 0.6503669 , 0.8678979 , 0.46042678],\\n\",\n \" [ 0.05673932, 0.9909285 , -0.63536596],\\n\",\n \" [ 0.76134115, -0.3250529 , -0.6522163 ],\\n\",\n \" [-0.8243032 , 0.4150194 , 0.19405058]], dtype=float32),\\n\",\n \" bias: DeviceArray([0., 0., 0.], dtype=float32),\\n\",\n \" },\\n\",\n \"})\\n\",\n \"output:\\n\",\n \" [[ 0.5035518 1.8548559 -0.4270196 ]\\n\",\n \" [ 0.0279097 0.5589246 -0.43061775]\\n\",\n \" [ 0.35471284 1.5741 -0.3286552 ]\\n\",\n \" [ 0.5264864 1.2928858 0.10089308]]\\n\"\n ]\n }\n ],\n \"source\": [\n \"class SimpleDense(nn.Module):\\n\",\n \" features: int\\n\",\n \" kernel_init: Callable = nn.initializers.lecun_normal()\\n\",\n \" bias_init: Callable = nn.initializers.zeros_init()\\n\",\n \"\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self, inputs):\\n\",\n \" kernel = self.param('kernel',\\n\",\n \" self.kernel_init, # RNG passed implicitly.\\n\",\n \" (inputs.shape[-1], self.features)) # shape info.\\n\",\n \" y = lax.dot_general(inputs, kernel,\\n\",\n \" (((inputs.ndim - 1,), (0,)), ((), ())),)\\n\",\n \" bias = self.param('bias', self.bias_init, (self.features,))\\n\",\n \" y = y + bias\\n\",\n \" return y\\n\",\n \"\\n\",\n \"key1, key2 = random.split(random.key(0), 2)\\n\",\n \"x = random.uniform(key1, (4,4))\\n\",\n \"\\n\",\n \"model = SimpleDense(features=3)\\n\",\n \"init_variables = model.init(key2, x)\\n\",\n \"y = model.apply(init_variables, x)\\n\",\n \"\\n\",\n \"print('initialized parameters:\\\\n', init_variables)\\n\",\n \"print('output:\\\\n', y)\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"We can also declare variables in setup, though in doing so you can't take advantage of shape inference and have to provide explicit shape information at initialization. The syntax is a little repetitive in this case right now, but we do force agreement of the assigned names.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 9,\n \"metadata\": {\n \"outputId\": \"1e822bd8-7a08-4e80-e0e6-a86637c46772\",\n \"tags\": []\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"initialized parameters:\\n\",\n \" FrozenDict({\\n\",\n \" params: {\\n\",\n \" kernel: DeviceArray([[ 0.6503669 , 0.8678979 , 0.46042678],\\n\",\n \" [ 0.05673932, 0.9909285 , -0.63536596],\\n\",\n \" [ 0.76134115, -0.3250529 , -0.6522163 ],\\n\",\n \" [-0.8243032 , 0.4150194 , 0.19405058]], dtype=float32),\\n\",\n \" bias: DeviceArray([0., 0., 0.], dtype=float32),\\n\",\n \" },\\n\",\n \"})\\n\",\n \"output:\\n\",\n \" [[ 0.5035518 1.8548559 -0.4270196 ]\\n\",\n \" [ 0.0279097 0.5589246 -0.43061775]\\n\",\n \" [ 0.35471284 1.5741 -0.3286552 ]\\n\",\n \" [ 0.5264864 1.2928858 0.10089308]]\\n\"\n ]\n }\n ],\n \"source\": [\n \"class ExplicitDense(nn.Module):\\n\",\n \" features_in: int # <-- explicit input shape\\n\",\n \" features: int\\n\",\n \" kernel_init: Callable = nn.initializers.lecun_normal()\\n\",\n \" bias_init: Callable = nn.initializers.zeros_init()\\n\",\n \"\\n\",\n \" def setup(self):\\n\",\n \" self.kernel = self.param('kernel',\\n\",\n \" self.kernel_init,\\n\",\n \" (self.features_in, self.features))\\n\",\n \" self.bias = self.param('bias', self.bias_init, (self.features,))\\n\",\n \"\\n\",\n \" def __call__(self, inputs):\\n\",\n \" y = lax.dot_general(inputs, self.kernel,\\n\",\n \" (((inputs.ndim - 1,), (0,)), ((), ())),)\\n\",\n \" y = y + self.bias\\n\",\n \" return y\\n\",\n \"\\n\",\n \"key1, key2 = random.split(random.key(0), 2)\\n\",\n \"x = random.uniform(key1, (4,4))\\n\",\n \"\\n\",\n \"model = ExplicitDense(features_in=4, features=3)\\n\",\n \"init_variables = model.init(key2, x)\\n\",\n \"y = model.apply(init_variables, x)\\n\",\n \"\\n\",\n \"print('initialized parameters:\\\\n', init_variables)\\n\",\n \"print('output:\\\\n', y)\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## General Variables\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"For declaring generally mutable _variables_ that may be mutated inside the model we use the call:\\n\",\n \"\\n\",\n \" `self.variable(variable_kind, variable_name, variable_init_fn, *init_args, **init_kwargs)`\\n\",\n \"\\n\",\n \"with arguments:\\n\",\n \" - `variable_kind` the \\\"kind\\\" of state this variable is, i.e. the name of the nested-dict collection that this will be stored in inside the top Modules variables. e.g. `batch_stats` for the moving statistics for a batch norm layer or `cache` for autoregressive cache data. Note that parameters also have a kind, but they're set to the default `param` kind.\\n\",\n \" - `variable_name` just the name, a string\\n\",\n \" - `variable_init_fn` a function taking a variable number of other arguments, i.e. `fn(*args)`. Note that we __don't__ assume the need for an RNG, if you _do_ want an RNG, provide it via a `self.make_rng(variable_kind)` call in the provided arguments.\\n\",\n \" - the remaining arguments to feed to the init function when initializing.\\n\",\n \"\\n\",\n \"⚠️ Unlike parameters, we expect these to be mutated, so `self.variable` returns not a constant, but a _reference_ to the variable. To __get__ the raw value, you'd write `myvariable.value` and to __set__ it `myvariable.value = new_value`.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 10,\n \"metadata\": {\n \"outputId\": \"2a8f5453-81b1-44dc-a431-d14b372c5710\",\n \"tags\": []\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"initialized variables:\\n\",\n \" FrozenDict({\\n\",\n \" counter: {\\n\",\n \" count: DeviceArray(0, dtype=int32),\\n\",\n \" },\\n\",\n \"})\\n\",\n \"mutated variables:\\n\",\n \" FrozenDict({\\n\",\n \" counter: {\\n\",\n \" count: DeviceArray(1, dtype=int32),\\n\",\n \" },\\n\",\n \"})\\n\",\n \"output:\\n\",\n \" 1\\n\"\n ]\n }\n ],\n \"source\": [\n \"class Counter(nn.Module):\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self):\\n\",\n \" # easy pattern to detect if we're initializing\\n\",\n \" is_initialized = self.has_variable('counter', 'count')\\n\",\n \" counter = self.variable('counter', 'count', lambda: jnp.zeros((), jnp.int32))\\n\",\n \" if is_initialized:\\n\",\n \" counter.value += 1\\n\",\n \" return counter.value\\n\",\n \"\\n\",\n \"\\n\",\n \"key1 = random.key(0)\\n\",\n \"\\n\",\n \"model = Counter()\\n\",\n \"init_variables = model.init(key1)\\n\",\n \"print('initialized variables:\\\\n', init_variables)\\n\",\n \"\\n\",\n \"y, mutated_variables = model.apply(init_variables, mutable=['counter'])\\n\",\n \"\\n\",\n \"print('mutated variables:\\\\n', mutated_variables)\\n\",\n \"print('output:\\\\n', y)\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Another Mutability and RNGs Example\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Let's make an artificial, goofy example that mixes differentiable parameters, stochastic layers, and mutable variables:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 11,\n \"metadata\": {\n \"outputId\": \"8f299a5c-74c8-476c-93fa-e5543901ec45\",\n \"tags\": []\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"updated variables:\\n\",\n \" FrozenDict({\\n\",\n \" params: {\\n\",\n \" Dense_0: {\\n\",\n \" kernel: DeviceArray([[ 0.6498898 , -0.5000124 , 0.78573596],\\n\",\n \" [-0.25609785, -0.7132329 , 0.2500864 ],\\n\",\n \" [-0.64630085, 0.39321756, -1.0203307 ],\\n\",\n \" [ 0.38721725, 0.86828285, 0.10860055]], dtype=float32),\\n\",\n \" bias: DeviceArray([0., 0., 0.], dtype=float32),\\n\",\n \" },\\n\",\n \" BatchNorm_0: {\\n\",\n \" scale: DeviceArray([1., 1., 1.], dtype=float32),\\n\",\n \" bias: DeviceArray([0., 0., 0.], dtype=float32),\\n\",\n \" },\\n\",\n \" },\\n\",\n \" batch_stats: {\\n\",\n \" BatchNorm_0: {\\n\",\n \" mean: DeviceArray([ 0.00059601, -0.00103457, 0.00166948], dtype=float32),\\n\",\n \" var: DeviceArray([0.9907686, 0.9923046, 0.992195 ], dtype=float32),\\n\",\n \" },\\n\",\n \" },\\n\",\n \"})\\n\",\n \"initialized variable shapes:\\n\",\n \" FrozenDict({\\n\",\n \" batch_stats: {\\n\",\n \" BatchNorm_0: {\\n\",\n \" mean: (3,),\\n\",\n \" var: (3,),\\n\",\n \" },\\n\",\n \" },\\n\",\n \" params: {\\n\",\n \" BatchNorm_0: {\\n\",\n \" bias: (3,),\\n\",\n \" scale: (3,),\\n\",\n \" },\\n\",\n \" Dense_0: {\\n\",\n \" bias: (3,),\\n\",\n \" kernel: (4, 3),\\n\",\n \" },\\n\",\n \" },\\n\",\n \"})\\n\",\n \"output:\\n\",\n \" [[[-0.21496922 0.21550177 -0.35633382]\\n\",\n \" [-0.21496922 -2.0458 1.3015485 ]\\n\",\n \" [-0.21496922 -0.925116 -0.35633382]\\n\",\n \" [-0.6595459 0.21550177 0.3749205 ]]\\n\",\n \"\\n\",\n \" [[-0.21496922 1.642865 -0.35633382]\\n\",\n \" [-0.21496922 1.3094063 -0.88034123]\\n\",\n \" [ 2.5726683 0.21550177 0.34353197]\\n\",\n \" [-0.21496922 0.21550177 1.6778195 ]]\\n\",\n \"\\n\",\n \" [[-1.6060593 0.21550177 -1.9460517 ]\\n\",\n \" [ 1.4126908 -1.4898677 1.2790381 ]\\n\",\n \" [-0.21496922 0.21550177 -0.35633382]\\n\",\n \" [-0.21496922 0.21550177 -0.7251308 ]]]\\n\",\n \"eval output:\\n\",\n \" [[[ 3.2246590e-01 2.6108384e-02 4.4821960e-01]\\n\",\n \" [ 8.5726947e-02 -5.4385906e-01 3.8821870e-01]\\n\",\n \" [-2.3933809e-01 -2.7381191e-01 -1.7526165e-01]\\n\",\n \" [-6.2515378e-02 -5.2414006e-01 1.7029770e-01]]\\n\",\n \"\\n\",\n \" [[ 1.5014435e-01 3.4498507e-01 -1.3554120e-01]\\n\",\n \" [-3.6971044e-04 2.6463276e-01 -1.2491019e-01]\\n\",\n \" [ 3.8763803e-01 2.9023719e-01 1.6291586e-01]\\n\",\n \" [ 4.1320035e-01 4.1468274e-02 4.7670874e-01]]\\n\",\n \"\\n\",\n \" [[-1.9433719e-01 5.2831882e-01 -3.7554008e-01]\\n\",\n \" [ 2.2608691e-01 -4.0989807e-01 3.8292480e-01]\\n\",\n \" [-2.4945706e-01 1.6170470e-01 -2.5247774e-01]\\n\",\n \" [-7.2220474e-02 1.2077977e-01 -8.8408351e-02]]]\\n\"\n ]\n }\n ],\n \"source\": [\n \"class Block(nn.Module):\\n\",\n \" features: int\\n\",\n \" training: bool\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self, inputs):\\n\",\n \" x = nn.Dense(self.features)(inputs)\\n\",\n \" x = nn.Dropout(rate=0.5)(x, deterministic=not self.training)\\n\",\n \" x = nn.BatchNorm(use_running_average=not self.training)(x)\\n\",\n \" return x\\n\",\n \"\\n\",\n \"key1, key2, key3, key4 = random.split(random.key(0), 4)\\n\",\n \"x = random.uniform(key1, (3,4,4))\\n\",\n \"\\n\",\n \"model = Block(features=3, training=True)\\n\",\n \"\\n\",\n \"init_variables = model.init({'params': key2, 'dropout': key3}, x)\\n\",\n \"_, init_params = flax.core.pop(init_variables, 'params')\\n\",\n \"\\n\",\n \"# When calling `apply` with mutable kinds, returns a pair of output,\\n\",\n \"# mutated_variables.\\n\",\n \"y, mutated_variables = model.apply(\\n\",\n \" init_variables, x, rngs={'dropout': key4}, mutable=['batch_stats'])\\n\",\n \"\\n\",\n \"# Now we reassemble the full variables from the updates (in a real training\\n\",\n \"# loop, with the updated params from an optimizer).\\n\",\n \"updated_variables = flax.core.freeze(dict(params=init_params,\\n\",\n \" **mutated_variables))\\n\",\n \"\\n\",\n \"print('updated variables:\\\\n', updated_variables)\\n\",\n \"print('initialized variable shapes:\\\\n',\\n\",\n \" jax.tree_util.tree_map(jnp.shape, init_variables))\\n\",\n \"print('output:\\\\n', y)\\n\",\n \"\\n\",\n \"# Let's run these model variables during \\\"evaluation\\\":\\n\",\n \"eval_model = Block(features=3, training=False)\\n\",\n \"y = eval_model.apply(updated_variables, x) # Nothing mutable; single return value.\\n\",\n \"print('eval output:\\\\n', y)\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"# JAX transformations inside modules\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## JIT\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"It's not immediately clear what use this has, but you can compile specific submodules if there's a reason to.\\n\",\n \"\\n\",\n \"_Known Gotcha_: at the moment, the decorator changes the RNG stream slightly, so comparing jitted an unjitted initializations will look different.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 12,\n \"metadata\": {\n \"outputId\": \"3f324d0f-259f-40f0-8273-103f7fc281c5\",\n \"tags\": []\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"initialized parameter shapes:\\n\",\n \" {'params': {'layers_0': {'bias': (3,), 'kernel': (4, 3)}, 'layers_1': {'bias': (4,), 'kernel': (3, 4)}, 'layers_2': {'bias': (5,), 'kernel': (4, 5)}}}\\n\",\n \"output:\\n\",\n \" [[ 0.2524199 0.11621253 0.5246693 0.19144788 0.2096542 ]\\n\",\n \" [ 0.08557513 -0.04126885 0.2502836 0.03910369 0.16575359]\\n\",\n \" [ 0.2804383 0.27751124 0.44969672 0.26016283 0.05875347]\\n\",\n \" [ 0.2440843 0.17069656 0.45499086 0.20377949 0.13428023]]\\n\"\n ]\n }\n ],\n \"source\": [\n \"class MLP(nn.Module):\\n\",\n \" features: Sequence[int]\\n\",\n \"\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self, inputs):\\n\",\n \" x = inputs\\n\",\n \" for i, feat in enumerate(self.features):\\n\",\n \" # JIT the Module (it's __call__ fn by default.)\\n\",\n \" x = nn.jit(nn.Dense)(feat, name=f'layers_{i}')(x)\\n\",\n \" if i != len(self.features) - 1:\\n\",\n \" x = nn.relu(x)\\n\",\n \" return x\\n\",\n \"\\n\",\n \"key1, key2 = random.split(random.key(3), 2)\\n\",\n \"x = random.uniform(key1, (4,4))\\n\",\n \"\\n\",\n \"model = MLP(features=[3,4,5])\\n\",\n \"init_variables = model.init(key2, x)\\n\",\n \"y = model.apply(init_variables, x)\\n\",\n \"\\n\",\n \"print('initialized parameter shapes:\\\\n', jax.tree_util.tree_map(jnp.shape, flax.core.unfreeze(init_variables)))\\n\",\n \"print('output:\\\\n', y)\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Remat\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"For memory-expensive computations, we can `remat` our method to recompute a Module's output during a backwards pass.\\n\",\n \"\\n\",\n \"_Known Gotcha_: at the moment, the decorator changes the RNG stream slightly, so comparing remat'd and undecorated initializations will look different.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 13,\n \"metadata\": {\n \"outputId\": \"7fe8e13b-7dd6-4e55-ee50-ce334e8ed178\",\n \"tags\": []\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"initialized parameter shapes:\\n\",\n \" {'params': {'layers_0': {'bias': (3,), 'kernel': (4, 3)}, 'layers_1': {'bias': (4,), 'kernel': (3, 4)}, 'layers_2': {'bias': (5,), 'kernel': (4, 5)}}}\\n\",\n \"output:\\n\",\n \" [[-0.14814317 0.06889858 -0.19695625 0.12019286 0.02068037]\\n\",\n \" [-0.04439102 -0.06698258 -0.11579747 -0.19906905 -0.04342325]\\n\",\n \" [-0.08875751 -0.13392815 -0.23153095 -0.39802808 -0.0868225 ]\\n\",\n \" [-0.01606487 -0.02424064 -0.04190649 -0.07204203 -0.01571464]]\\n\"\n ]\n }\n ],\n \"source\": [\n \"class RematMLP(nn.Module):\\n\",\n \" features: Sequence[int]\\n\",\n \" # For all transforms, we can annotate a method, or wrap an existing\\n\",\n \" # Module class. Here we annotate the method.\\n\",\n \" @nn.remat\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self, inputs):\\n\",\n \" x = inputs\\n\",\n \" for i, feat in enumerate(self.features):\\n\",\n \" x = nn.Dense(feat, name=f'layers_{i}')(x)\\n\",\n \" if i != len(self.features) - 1:\\n\",\n \" x = nn.relu(x)\\n\",\n \" return x\\n\",\n \"\\n\",\n \"key1, key2 = random.split(random.key(3), 2)\\n\",\n \"x = random.uniform(key1, (4,4))\\n\",\n \"\\n\",\n \"model = RematMLP(features=[3,4,5])\\n\",\n \"init_variables = model.init(key2, x)\\n\",\n \"y = model.apply(init_variables, x)\\n\",\n \"\\n\",\n \"print('initialized parameter shapes:\\\\n', jax.tree_util.tree_map(jnp.shape, flax.core.unfreeze(init_variables)))\\n\",\n \"print('output:\\\\n', y)\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Vmap\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"You can now `vmap` Modules inside. The transform has a lot of arguments, they have the usual jax vmap args:\\n\",\n \" - `in_axes` - an integer or `None` for each input argument\\n\",\n \" - `out_axes` - an integer or `None` for each output argument\\n\",\n \" - `axis_size` - the axis size if you need to give it explicitly\\n\",\n \"\\n\",\n \"In addition, we provide for each __kind__ of variable it's axis rules:\\n\",\n \"\\n\",\n \" - `variable_in_axes` - a dict from kinds to a single integer or `None` specifying the input axes to map\\n\",\n \" - `variable_out_axes` - a dict from kinds to a single integer or `None` specifying the output axes to map\\n\",\n \" - `split_rngs` - a dict from RNG-kinds to a bool, specifying whether to split the rng along the axis.\\n\",\n \"\\n\",\n \"\\n\",\n \"Below we show an example defining a batched, multiheaded attention module from a single-headed unbatched attention implementation.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 14,\n \"metadata\": {\n \"outputId\": \"223d880e-c7b2-4210-ebb5-dbfcdd9aed09\",\n \"tags\": []\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"initialized parameter shapes:\\n\",\n \" {'params': {'attention': {'key': {'kernel': (2, 64, 32)}, 'out': {'bias': (2, 64), 'kernel': (2, 32, 64)}, 'query': {'kernel': (2, 64, 32)}, 'value': {'kernel': (2, 64, 32)}}}}\\n\",\n \"output:\\n\",\n \" (3, 13, 2)\\n\"\n ]\n }\n ],\n \"source\": [\n \"class RawDotProductAttention(nn.Module):\\n\",\n \" attn_dropout_rate: float = 0.1\\n\",\n \" train: bool = False\\n\",\n \"\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self, query, key, value, bias=None, dtype=jnp.float32):\\n\",\n \" assert key.ndim == query.ndim\\n\",\n \" assert key.ndim == value.ndim\\n\",\n \"\\n\",\n \" n = query.ndim\\n\",\n \" attn_weights = lax.dot_general(\\n\",\n \" query, key,\\n\",\n \" (((n-1,), (n - 1,)), ((), ())))\\n\",\n \" if bias is not None:\\n\",\n \" attn_weights += bias\\n\",\n \" norm_dims = tuple(range(attn_weights.ndim // 2, attn_weights.ndim))\\n\",\n \" attn_weights = jax.nn.softmax(attn_weights, axis=norm_dims)\\n\",\n \" attn_weights = nn.Dropout(self.attn_dropout_rate)(attn_weights,\\n\",\n \" deterministic=not self.train)\\n\",\n \" attn_weights = attn_weights.astype(dtype)\\n\",\n \"\\n\",\n \" contract_dims = (\\n\",\n \" tuple(range(n - 1, attn_weights.ndim)),\\n\",\n \" tuple(range(0, n - 1)))\\n\",\n \" y = lax.dot_general(\\n\",\n \" attn_weights, value,\\n\",\n \" (contract_dims, ((), ())))\\n\",\n \" return y\\n\",\n \"\\n\",\n \"class DotProductAttention(nn.Module):\\n\",\n \" qkv_features: Optional[int] = None\\n\",\n \" out_features: Optional[int] = None\\n\",\n \" train: bool = False\\n\",\n \"\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self, inputs_q, inputs_kv, bias=None, dtype=jnp.float32):\\n\",\n \" qkv_features = self.qkv_features or inputs_q.shape[-1]\\n\",\n \" out_features = self.out_features or inputs_q.shape[-1]\\n\",\n \"\\n\",\n \" QKVDense = functools.partial(\\n\",\n \" nn.Dense, features=qkv_features, use_bias=False, dtype=dtype)\\n\",\n \" query = QKVDense(name='query')(inputs_q)\\n\",\n \" key = QKVDense(name='key')(inputs_kv)\\n\",\n \" value = QKVDense(name='value')(inputs_kv)\\n\",\n \"\\n\",\n \" y = RawDotProductAttention(train=self.train)(\\n\",\n \" query, key, value, bias=bias, dtype=dtype)\\n\",\n \"\\n\",\n \" y = nn.Dense(features=out_features, dtype=dtype, name='out')(y)\\n\",\n \" return y\\n\",\n \"\\n\",\n \"class MultiHeadDotProductAttention(nn.Module):\\n\",\n \" qkv_features: Optional[int] = None\\n\",\n \" out_features: Optional[int] = None\\n\",\n \" batch_axes: Sequence[int] = (0,)\\n\",\n \" num_heads: int = 1\\n\",\n \" broadcast_dropout: bool = False\\n\",\n \" train: bool = False\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self, inputs_q, inputs_kv, bias=None, dtype=jnp.float32):\\n\",\n \" qkv_features = self.qkv_features or inputs_q.shape[-1]\\n\",\n \" out_features = self.out_features or inputs_q.shape[-1]\\n\",\n \"\\n\",\n \" # Make multiheaded attention from single-headed dimension.\\n\",\n \" Attn = nn.vmap(DotProductAttention,\\n\",\n \" in_axes=(None, None, None),\\n\",\n \" out_axes=2,\\n\",\n \" axis_size=self.num_heads,\\n\",\n \" variable_axes={'params': 0},\\n\",\n \" split_rngs={'params': True,\\n\",\n \" 'dropout': not self.broadcast_dropout})\\n\",\n \"\\n\",\n \" # Vmap across batch dimensions.\\n\",\n \" for axis in reversed(sorted(self.batch_axes)):\\n\",\n \" Attn = nn.vmap(Attn,\\n\",\n \" in_axes=(axis, axis, axis),\\n\",\n \" out_axes=axis,\\n\",\n \" variable_axes={'params': None},\\n\",\n \" split_rngs={'params': False, 'dropout': False})\\n\",\n \"\\n\",\n \" # Run the vmap'd class on inputs.\\n\",\n \" y = Attn(qkv_features=qkv_features // self.num_heads,\\n\",\n \" out_features=out_features,\\n\",\n \" train=self.train,\\n\",\n \" name='attention')(inputs_q, inputs_kv, bias)\\n\",\n \"\\n\",\n \" return y.mean(axis=-2)\\n\",\n \"\\n\",\n \"\\n\",\n \"key1, key2, key3, key4 = random.split(random.key(0), 4)\\n\",\n \"x = random.uniform(key1, (3, 13, 64))\\n\",\n \"\\n\",\n \"model = functools.partial(\\n\",\n \" MultiHeadDotProductAttention,\\n\",\n \" broadcast_dropout=False,\\n\",\n \" num_heads=2,\\n\",\n \" batch_axes=(0,))\\n\",\n \"\\n\",\n \"init_variables = model(train=False).init({'params': key2}, x, x)\\n\",\n \"print('initialized parameter shapes:\\\\n', jax.tree_util.tree_map(jnp.shape, flax.core.unfreeze(init_variables)))\\n\",\n \"\\n\",\n \"y = model(train=True).apply(init_variables, x, x, rngs={'dropout': key4})\\n\",\n \"print('output:\\\\n', y.shape)\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Scan\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Scan allows us to apply `lax.scan` to Modules, including their parameters and mutable variables. To use it we have to specify how we want each \\\"kind\\\" of variable to be transformed. For scanned variables we specify similar to vmap via in `variable_in_axes`, `variable_out_axes`:\\n\",\n \" - `nn.broadcast` broadcast the variable kind across the scan steps as a constant\\n\",\n \" - `` scan along `axis` for e.g. unique parameters at each step\\n\",\n \"\\n\",\n \"OR we specify that the variable kind is to be treated like a \\\"carry\\\" by passing to the `variable_carry` argument.\\n\",\n \"\\n\",\n \"Further, for `scan`'d variable kinds, we further specify whether or not to split the rng at each step.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 15,\n \"metadata\": {\n \"outputId\": \"7d9ebed3-64de-4ca8-9dce-4b09ba9e31a1\",\n \"tags\": []\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"initialized parameter shapes:\\n\",\n \" {'params': {'lstm_cell': {'hf': {'bias': (2,), 'kernel': (2, 2)}, 'hg': {'bias': (2,), 'kernel': (2, 2)}, 'hi': {'bias': (2,), 'kernel': (2, 2)}, 'ho': {'bias': (2,), 'kernel': (2, 2)}, 'if': {'kernel': (2, 2)}, 'ig': {'kernel': (2, 2)}, 'ii': {'kernel': (2, 2)}, 'io': {'kernel': (2, 2)}}}}\\n\",\n \"output:\\n\",\n \" ((DeviceArray([[-0.562219 , 0.92847174]], dtype=float32), DeviceArray([[-0.31570646, 0.2885693 ]], dtype=float32)), DeviceArray([[[-0.08265854, 0.01302483],\\n\",\n \" [-0.10249066, 0.21991298],\\n\",\n \" [-0.26609066, 0.22519003],\\n\",\n \" [-0.27982554, 0.28393182],\\n\",\n \" [-0.31570646, 0.2885693 ]]], dtype=float32))\\n\"\n ]\n }\n ],\n \"source\": [\n \"class SimpleScan(nn.Module):\\n\",\n \" features: int\\n\",\n \"\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self, xs):\\n\",\n \" LSTM = nn.scan(nn.LSTMCell,\\n\",\n \" in_axes=1, out_axes=1,\\n\",\n \" variable_broadcast='params',\\n\",\n \" split_rngs={'params': False})\\n\",\n \" lstm = LSTM(self.features, name=\\\"lstm_cell\\\")\\n\",\n \"\\n\",\n \" dummy_rng = random.key(0)\\n\",\n \" input_shape = xs[:, 0].shape\\n\",\n \" init_carry = lstm.initialize_carry(dummy_rng, input_shape)\\n\",\n \"\\n\",\n \" return lstm(init_carry, xs)\\n\",\n \"\\n\",\n \"key1, key2 = random.split(random.key(0), 2)\\n\",\n \"xs = random.uniform(key1, (1, 5, 2))\\n\",\n \"\\n\",\n \"model = SimpleScan(2)\\n\",\n \"init_variables = model.init(key2, xs)\\n\",\n \"\\n\",\n \"print('initialized parameter shapes:\\\\n', jax.tree_util.tree_map(jnp.shape, flax.core.unfreeze(init_variables)))\\n\",\n \"\\n\",\n \"y = model.apply(init_variables, xs)\\n\",\n \"print('output:\\\\n', y)\"\n ]\n }\n ],\n \"metadata\": {\n \"jupytext\": {\n \"formats\": \"ipynb,md:myst\"\n },\n \"language_info\": {\n \"name\": \"python\",\n \"version\": \"3.8.10\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 0\n}\n" }, { "path": "docs/linen_intro.md", "content": "---\njupytext:\n formats: ipynb,md:myst\n text_representation:\n extension: .md\n format_name: myst\n format_version: 0.13\n jupytext_version: 1.13.8\n---\n\n[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google/flax/blob/main/docs/notebooks/linen_intro.ipynb)\n[![Open On GitHub](https://img.shields.io/badge/Open-on%20GitHub-blue?logo=GitHub)](https://github.com/google/flax/blob/main/docs/notebooks/linen_intro.ipynb)\n\n# Preface\n\n
\n
CAVEAT PROGRAMMER
\n\nThe below is an alpha API preview and things might break. The surface syntax of the features of the API are not fixed in stone, and we welcome feedback on any points.\n\n+++\n\n## Useful links\n\n⟶ [Slides](https://docs.google.com/presentation/d/1ngKWUwsSqAwPRvATG8sAxMzu9ujv4N__cKsUofdNno0/edit?usp=sharing) for the core ideas of the new Functional Core and Linen\n\n⟶ \"Design tests\" guided our design process. Many are available for [functional core](https://github.com/google/flax/tree/main/examples/core_design_test) and some for the [proposed Module abstraction](https://github.com/google/flax/tree/main/examples/linen_design_test/)\n\n⟶ Ported examples: [ImageNet](https://github.com/google/flax/tree/main/examples/imagenet) and [WMT](https://github.com/google/flax/tree/main/examples/wmt) (to the proposed Module abstraction). TODO: Port to functional core.\n\n⟶ Our new [discussion forums](https://github.com/google/flax/discussions/)\n\n+++\n\n# Install and Import\n\n```{code-cell}\n:tags: [skip-execution]\n\n# Install the newest JAXlib version.\n!pip install --upgrade -q pip jax jaxlib\n# Install Flax at head:\n!pip install --upgrade -q git+https://github.com/google/flax.git\n```\n\n```{code-cell}\nimport functools\nfrom typing import Any, Callable, Sequence, Optional\nimport jax\nfrom jax import lax, random, numpy as jnp\nimport flax\nfrom flax import linen as nn\n```\n\n# Invoking Modules\n\n+++\n\nLet's instantiate a `Dense` layer.\n - Modules are actually objects in this API, so we provide _constructor arguments_ when initializing the Module. In this case, we only have to provide the output `features` dimension.\n\n```{code-cell}\nmodel = nn.Dense(features=3)\n```\n\nWe need to initialize the Module variables, these include the parameters of the Module as well as any other state variables.\n\nWe call the `init` method on the instantiated Module. If the Module `__call__` method has args `(self, *args, **kwargs)` then we call `init` with `(rngs, *args, **kwargs)` so in this case, just `(rng, input)`:\n\n```{code-cell}\n:outputId: 3adfaeaf-977e-4e82-8adf-d254fae6eb91\n\n# Make RNG Keys and a fake input.\nkey1, key2 = random.split(random.key(0), 2)\nx = random.uniform(key1, (4,4))\n\n# provide key and fake input to get initialized variables\ninit_variables = model.init(key2, x)\n\ninit_variables\n```\n\nWe call the `apply` method on the instantiated Module. If the Module `__call__` method has args `(self, *args, **kwargs)` then we call `apply` with `(variables, *args, rngs=, mutable=, **kwargs)` where\n - `` are the optional _call time_ RNGs for things like dropout. For simple Modules this is just a single key, but if your module has multiple __kinds__ of data, it's a dictionary of rng-keys per-kind, e.g. `{'params': key0, 'dropout': key1}` for a Module with dropout layers.\n - `` is an optional list of names of __kinds__ that are expected to be mutated during the call. e.g. `['batch_stats']` for a layer updating batchnorm statistics.\n\nSo in this case, just `(variables, input)`:\n\n```{code-cell}\n:outputId: e8c389a6-29f3-4f93-97ea-703e85a8b811\n\ny = model.apply(init_variables, x)\ny\n```\n\nAdditional points:\n - If you want to `init` or `apply` a Module using a method other than call, you need to provide the `method=` kwarg to `init` and `apply` to use it instead of the default `__call__`, e.g. `method='encode'`, `method='decode'` to apply the encode/decode methods of an autoencoder.\n\n+++\n\n# Defining Basic Modules\n\n+++\n\n## Composing submodules\n\n+++\n\nWe support declaring modules in `setup()` that can still benefit from shape inference by using __Lazy Initialization__ that sets up variables the first time the Module is called.\n\n```{code-cell}\n:outputId: 1a6c6a17-0b95-42c2-b5bf-b9ad80fd7758\n:tags: []\n\nclass ExplicitMLP(nn.Module):\n features: Sequence[int]\n\n def setup(self):\n # we automatically know what to do with lists, dicts of submodules\n self.layers = [nn.Dense(feat) for feat in self.features]\n # for single submodules, we would just write:\n # self.layer1 = nn.Dense(feat1)\n\n def __call__(self, inputs):\n x = inputs\n for i, lyr in enumerate(self.layers):\n x = lyr(x)\n if i != len(self.layers) - 1:\n x = nn.relu(x)\n return x\n\nkey1, key2 = random.split(random.key(0), 2)\nx = random.uniform(key1, (4,4))\n\nmodel = ExplicitMLP(features=[3,4,5])\ninit_variables = model.init(key2, x)\ny = model.apply(init_variables, x)\n\nprint('initialized parameter shapes:\\n', jax.tree_util.tree_map(jnp.shape, flax.core.unfreeze(init_variables)))\nprint('output:\\n', y)\n```\n\nHere we show the equivalent compact form of the MLP that declares the submodules inline using the `@compact` decorator.\n\n```{code-cell}\n:outputId: b3709789-e66e-4e20-f6b2-04022f8a62bb\n:tags: []\n\nclass SimpleMLP(nn.Module):\n features: Sequence[int]\n\n @nn.compact\n def __call__(self, inputs):\n x = inputs\n for i, feat in enumerate(self.features):\n x = nn.Dense(feat, name=f'layers_{i}')(x)\n if i != len(self.features) - 1:\n x = nn.relu(x)\n # providing a name is optional though!\n # the default autonames would be \"Dense_0\", \"Dense_1\", ...\n # x = nn.Dense(feat)(x)\n return x\n\nkey1, key2 = random.split(random.key(0), 2)\nx = random.uniform(key1, (4,4))\n\nmodel = SimpleMLP(features=[3,4,5])\ninit_variables = model.init(key2, x)\ny = model.apply(init_variables, x)\n\nprint('initialized parameter shapes:\\n', jax.tree_util.tree_map(jnp.shape, flax.core.unfreeze(init_variables)))\nprint('output:\\n', y)\n```\n\n## Declaring and using variables\n\n+++\n\nFlax uses lazy initialization, which allows declared variables to be initialized only at the first site of their use, using whatever shape information is available a the local call site for shape inference. Once a variable has been initialized, a reference to the data is kept for use in subsequent calls.\n\nFor declaring parameters that aren't mutated inside the model, but rather by gradient descent, we use the syntax:\n\n `self.param(parameter_name, parameter_init_fn, *init_args, **init_kwargs)`\n\nwith arguments:\n - `parameter_name` just the name, a string\n - `parameter_init_fn` a function taking an RNG key and a variable number of other arguments, i.e. `fn(rng, *args)`. typically those in `nn.initializers` take an `rng` and a `shape` argument.\n - the remaining arguments to feed to the init function when initializing.\n\nAgain, we'll demonstrate declaring things inline as we typically do using the `@compact` decorator.\n\n```{code-cell}\n:outputId: bc5cb1f2-c5e9-4159-d131-73247009e32f\n:tags: []\n\nclass SimpleDense(nn.Module):\n features: int\n kernel_init: Callable = nn.initializers.lecun_normal()\n bias_init: Callable = nn.initializers.zeros_init()\n\n @nn.compact\n def __call__(self, inputs):\n kernel = self.param('kernel',\n self.kernel_init, # RNG passed implicitly.\n (inputs.shape[-1], self.features)) # shape info.\n y = lax.dot_general(inputs, kernel,\n (((inputs.ndim - 1,), (0,)), ((), ())),)\n bias = self.param('bias', self.bias_init, (self.features,))\n y = y + bias\n return y\n\nkey1, key2 = random.split(random.key(0), 2)\nx = random.uniform(key1, (4,4))\n\nmodel = SimpleDense(features=3)\ninit_variables = model.init(key2, x)\ny = model.apply(init_variables, x)\n\nprint('initialized parameters:\\n', init_variables)\nprint('output:\\n', y)\n```\n\nWe can also declare variables in setup, though in doing so you can't take advantage of shape inference and have to provide explicit shape information at initialization. The syntax is a little repetitive in this case right now, but we do force agreement of the assigned names.\n\n```{code-cell}\n:outputId: 1e822bd8-7a08-4e80-e0e6-a86637c46772\n:tags: []\n\nclass ExplicitDense(nn.Module):\n features_in: int # <-- explicit input shape\n features: int\n kernel_init: Callable = nn.initializers.lecun_normal()\n bias_init: Callable = nn.initializers.zeros_init()\n\n def setup(self):\n self.kernel = self.param('kernel',\n self.kernel_init,\n (self.features_in, self.features))\n self.bias = self.param('bias', self.bias_init, (self.features,))\n\n def __call__(self, inputs):\n y = lax.dot_general(inputs, self.kernel,\n (((inputs.ndim - 1,), (0,)), ((), ())),)\n y = y + self.bias\n return y\n\nkey1, key2 = random.split(random.key(0), 2)\nx = random.uniform(key1, (4,4))\n\nmodel = ExplicitDense(features_in=4, features=3)\ninit_variables = model.init(key2, x)\ny = model.apply(init_variables, x)\n\nprint('initialized parameters:\\n', init_variables)\nprint('output:\\n', y)\n```\n\n## General Variables\n\n+++\n\nFor declaring generally mutable _variables_ that may be mutated inside the model we use the call:\n\n `self.variable(variable_kind, variable_name, variable_init_fn, *init_args, **init_kwargs)`\n\nwith arguments:\n - `variable_kind` the \"kind\" of state this variable is, i.e. the name of the nested-dict collection that this will be stored in inside the top Modules variables. e.g. `batch_stats` for the moving statistics for a batch norm layer or `cache` for autoregressive cache data. Note that parameters also have a kind, but they're set to the default `param` kind.\n - `variable_name` just the name, a string\n - `variable_init_fn` a function taking a variable number of other arguments, i.e. `fn(*args)`. Note that we __don't__ assume the need for an RNG, if you _do_ want an RNG, provide it via a `self.make_rng(variable_kind)` call in the provided arguments.\n - the remaining arguments to feed to the init function when initializing.\n\n⚠️ Unlike parameters, we expect these to be mutated, so `self.variable` returns not a constant, but a _reference_ to the variable. To __get__ the raw value, you'd write `myvariable.value` and to __set__ it `myvariable.value = new_value`.\n\n```{code-cell}\n:outputId: 2a8f5453-81b1-44dc-a431-d14b372c5710\n:tags: []\n\nclass Counter(nn.Module):\n @nn.compact\n def __call__(self):\n # easy pattern to detect if we're initializing\n is_initialized = self.has_variable('counter', 'count')\n counter = self.variable('counter', 'count', lambda: jnp.zeros((), jnp.int32))\n if is_initialized:\n counter.value += 1\n return counter.value\n\n\nkey1 = random.key(0)\n\nmodel = Counter()\ninit_variables = model.init(key1)\nprint('initialized variables:\\n', init_variables)\n\ny, mutated_variables = model.apply(init_variables, mutable=['counter'])\n\nprint('mutated variables:\\n', mutated_variables)\nprint('output:\\n', y)\n```\n\n## Another Mutability and RNGs Example\n\n+++\n\nLet's make an artificial, goofy example that mixes differentiable parameters, stochastic layers, and mutable variables:\n\n```{code-cell}\n:outputId: 8f299a5c-74c8-476c-93fa-e5543901ec45\n:tags: []\n\nclass Block(nn.Module):\n features: int\n training: bool\n @nn.compact\n def __call__(self, inputs):\n x = nn.Dense(self.features)(inputs)\n x = nn.Dropout(rate=0.5)(x, deterministic=not self.training)\n x = nn.BatchNorm(use_running_average=not self.training)(x)\n return x\n\nkey1, key2, key3, key4 = random.split(random.key(0), 4)\nx = random.uniform(key1, (3,4,4))\n\nmodel = Block(features=3, training=True)\n\ninit_variables = model.init({'params': key2, 'dropout': key3}, x)\n_, init_params = flax.core.pop(init_variables, 'params')\n\n# When calling `apply` with mutable kinds, returns a pair of output,\n# mutated_variables.\ny, mutated_variables = model.apply(\n init_variables, x, rngs={'dropout': key4}, mutable=['batch_stats'])\n\n# Now we reassemble the full variables from the updates (in a real training\n# loop, with the updated params from an optimizer).\nupdated_variables = flax.core.freeze(dict(params=init_params,\n **mutated_variables))\n\nprint('updated variables:\\n', updated_variables)\nprint('initialized variable shapes:\\n',\n jax.tree_util.tree_map(jnp.shape, init_variables))\nprint('output:\\n', y)\n\n# Let's run these model variables during \"evaluation\":\neval_model = Block(features=3, training=False)\ny = eval_model.apply(updated_variables, x) # Nothing mutable; single return value.\nprint('eval output:\\n', y)\n```\n\n# JAX transformations inside modules\n\n+++\n\n## JIT\n\n+++\n\nIt's not immediately clear what use this has, but you can compile specific submodules if there's a reason to.\n\n_Known Gotcha_: at the moment, the decorator changes the RNG stream slightly, so comparing jitted an unjitted initializations will look different.\n\n```{code-cell}\n:outputId: 3f324d0f-259f-40f0-8273-103f7fc281c5\n:tags: []\n\nclass MLP(nn.Module):\n features: Sequence[int]\n\n @nn.compact\n def __call__(self, inputs):\n x = inputs\n for i, feat in enumerate(self.features):\n # JIT the Module (it's __call__ fn by default.)\n x = nn.jit(nn.Dense)(feat, name=f'layers_{i}')(x)\n if i != len(self.features) - 1:\n x = nn.relu(x)\n return x\n\nkey1, key2 = random.split(random.key(3), 2)\nx = random.uniform(key1, (4,4))\n\nmodel = MLP(features=[3,4,5])\ninit_variables = model.init(key2, x)\ny = model.apply(init_variables, x)\n\nprint('initialized parameter shapes:\\n', jax.tree_util.tree_map(jnp.shape, flax.core.unfreeze(init_variables)))\nprint('output:\\n', y)\n```\n\n## Remat\n\n+++\n\nFor memory-expensive computations, we can `remat` our method to recompute a Module's output during a backwards pass.\n\n_Known Gotcha_: at the moment, the decorator changes the RNG stream slightly, so comparing remat'd and undecorated initializations will look different.\n\n```{code-cell}\n:outputId: 7fe8e13b-7dd6-4e55-ee50-ce334e8ed178\n:tags: []\n\nclass RematMLP(nn.Module):\n features: Sequence[int]\n # For all transforms, we can annotate a method, or wrap an existing\n # Module class. Here we annotate the method.\n @nn.remat\n @nn.compact\n def __call__(self, inputs):\n x = inputs\n for i, feat in enumerate(self.features):\n x = nn.Dense(feat, name=f'layers_{i}')(x)\n if i != len(self.features) - 1:\n x = nn.relu(x)\n return x\n\nkey1, key2 = random.split(random.key(3), 2)\nx = random.uniform(key1, (4,4))\n\nmodel = RematMLP(features=[3,4,5])\ninit_variables = model.init(key2, x)\ny = model.apply(init_variables, x)\n\nprint('initialized parameter shapes:\\n', jax.tree_util.tree_map(jnp.shape, flax.core.unfreeze(init_variables)))\nprint('output:\\n', y)\n```\n\n## Vmap\n\n+++\n\nYou can now `vmap` Modules inside. The transform has a lot of arguments, they have the usual jax vmap args:\n - `in_axes` - an integer or `None` for each input argument\n - `out_axes` - an integer or `None` for each output argument\n - `axis_size` - the axis size if you need to give it explicitly\n\nIn addition, we provide for each __kind__ of variable it's axis rules:\n\n - `variable_in_axes` - a dict from kinds to a single integer or `None` specifying the input axes to map\n - `variable_out_axes` - a dict from kinds to a single integer or `None` specifying the output axes to map\n - `split_rngs` - a dict from RNG-kinds to a bool, specifying whether to split the rng along the axis.\n\n\nBelow we show an example defining a batched, multiheaded attention module from a single-headed unbatched attention implementation.\n\n```{code-cell}\n:outputId: 223d880e-c7b2-4210-ebb5-dbfcdd9aed09\n:tags: []\n\nclass RawDotProductAttention(nn.Module):\n attn_dropout_rate: float = 0.1\n train: bool = False\n\n @nn.compact\n def __call__(self, query, key, value, bias=None, dtype=jnp.float32):\n assert key.ndim == query.ndim\n assert key.ndim == value.ndim\n\n n = query.ndim\n attn_weights = lax.dot_general(\n query, key,\n (((n-1,), (n - 1,)), ((), ())))\n if bias is not None:\n attn_weights += bias\n norm_dims = tuple(range(attn_weights.ndim // 2, attn_weights.ndim))\n attn_weights = jax.nn.softmax(attn_weights, axis=norm_dims)\n attn_weights = nn.Dropout(self.attn_dropout_rate)(attn_weights,\n deterministic=not self.train)\n attn_weights = attn_weights.astype(dtype)\n\n contract_dims = (\n tuple(range(n - 1, attn_weights.ndim)),\n tuple(range(0, n - 1)))\n y = lax.dot_general(\n attn_weights, value,\n (contract_dims, ((), ())))\n return y\n\nclass DotProductAttention(nn.Module):\n qkv_features: Optional[int] = None\n out_features: Optional[int] = None\n train: bool = False\n\n @nn.compact\n def __call__(self, inputs_q, inputs_kv, bias=None, dtype=jnp.float32):\n qkv_features = self.qkv_features or inputs_q.shape[-1]\n out_features = self.out_features or inputs_q.shape[-1]\n\n QKVDense = functools.partial(\n nn.Dense, features=qkv_features, use_bias=False, dtype=dtype)\n query = QKVDense(name='query')(inputs_q)\n key = QKVDense(name='key')(inputs_kv)\n value = QKVDense(name='value')(inputs_kv)\n\n y = RawDotProductAttention(train=self.train)(\n query, key, value, bias=bias, dtype=dtype)\n\n y = nn.Dense(features=out_features, dtype=dtype, name='out')(y)\n return y\n\nclass MultiHeadDotProductAttention(nn.Module):\n qkv_features: Optional[int] = None\n out_features: Optional[int] = None\n batch_axes: Sequence[int] = (0,)\n num_heads: int = 1\n broadcast_dropout: bool = False\n train: bool = False\n @nn.compact\n def __call__(self, inputs_q, inputs_kv, bias=None, dtype=jnp.float32):\n qkv_features = self.qkv_features or inputs_q.shape[-1]\n out_features = self.out_features or inputs_q.shape[-1]\n\n # Make multiheaded attention from single-headed dimension.\n Attn = nn.vmap(DotProductAttention,\n in_axes=(None, None, None),\n out_axes=2,\n axis_size=self.num_heads,\n variable_axes={'params': 0},\n split_rngs={'params': True,\n 'dropout': not self.broadcast_dropout})\n\n # Vmap across batch dimensions.\n for axis in reversed(sorted(self.batch_axes)):\n Attn = nn.vmap(Attn,\n in_axes=(axis, axis, axis),\n out_axes=axis,\n variable_axes={'params': None},\n split_rngs={'params': False, 'dropout': False})\n\n # Run the vmap'd class on inputs.\n y = Attn(qkv_features=qkv_features // self.num_heads,\n out_features=out_features,\n train=self.train,\n name='attention')(inputs_q, inputs_kv, bias)\n\n return y.mean(axis=-2)\n\n\nkey1, key2, key3, key4 = random.split(random.key(0), 4)\nx = random.uniform(key1, (3, 13, 64))\n\nmodel = functools.partial(\n MultiHeadDotProductAttention,\n broadcast_dropout=False,\n num_heads=2,\n batch_axes=(0,))\n\ninit_variables = model(train=False).init({'params': key2}, x, x)\nprint('initialized parameter shapes:\\n', jax.tree_util.tree_map(jnp.shape, flax.core.unfreeze(init_variables)))\n\ny = model(train=True).apply(init_variables, x, x, rngs={'dropout': key4})\nprint('output:\\n', y.shape)\n```\n\n## Scan\n\n+++\n\nScan allows us to apply `lax.scan` to Modules, including their parameters and mutable variables. To use it we have to specify how we want each \"kind\" of variable to be transformed. For scanned variables we specify similar to vmap via in `variable_in_axes`, `variable_out_axes`:\n - `nn.broadcast` broadcast the variable kind across the scan steps as a constant\n - `` scan along `axis` for e.g. unique parameters at each step\n\nOR we specify that the variable kind is to be treated like a \"carry\" by passing to the `variable_carry` argument.\n\nFurther, for `scan`'d variable kinds, we further specify whether or not to split the rng at each step.\n\n```{code-cell}\n:outputId: 7d9ebed3-64de-4ca8-9dce-4b09ba9e31a1\n:tags: []\n\nclass SimpleScan(nn.Module):\n features: int\n\n @nn.compact\n def __call__(self, xs):\n LSTM = nn.scan(nn.LSTMCell,\n in_axes=1, out_axes=1,\n variable_broadcast='params',\n split_rngs={'params': False})\n lstm = LSTM(self.features, name=\"lstm_cell\")\n\n dummy_rng = random.key(0)\n input_shape = xs[:, 0].shape\n init_carry = lstm.initialize_carry(dummy_rng, input_shape)\n\n return lstm(init_carry, xs)\n\nkey1, key2 = random.split(random.key(0), 2)\nxs = random.uniform(key1, (1, 5, 2))\n\nmodel = SimpleScan(2)\ninit_variables = model.init(key2, xs)\n\nprint('initialized parameter shapes:\\n', jax.tree_util.tree_map(jnp.shape, flax.core.unfreeze(init_variables)))\n\ny = model.apply(init_variables, xs)\nprint('output:\\n', y)\n```\n" }, { "path": "docs/quick_start.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"id\": \"6eea21b3\",\n \"metadata\": {},\n \"source\": [\n \"[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google/flax/blob/main/docs/quick_start.ipynb)\\n\",\n \"[![Open On GitHub](https://img.shields.io/badge/Open-on%20GitHub-blue?logo=GitHub)](https://github.com/google/flax/blob/main/docs/quick_start.ipynb)\\n\",\n \"\\n\",\n \"# Quick start\\n\",\n \"\\n\",\n \"Welcome to Flax!\\n\",\n \"\\n\",\n \"Flax is an open source Python neural network library built on top of [JAX](https://github.com/jax-ml/jax). This tutorial demonstrates how to construct a simple convolutional neural\\n\",\n \"network (CNN) using the [Flax](https://flax.readthedocs.io) Linen API and train\\n\",\n \"the network for image classification on the MNIST dataset.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"nwJWKIhdwxDo\",\n \"metadata\": {},\n \"source\": [\n \"## 1. Install Flax\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"bb81587e\",\n \"metadata\": {\n \"tags\": [\n \"skip-execution\"\n ]\n },\n \"outputs\": [],\n \"source\": [\n \"!pip install -q flax>=0.7.5\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"b529fbef\",\n \"metadata\": {},\n \"source\": [\n \"## 2. Loading data\\n\",\n \"\\n\",\n \"Flax can use any\\n\",\n \"data-loading pipeline and this example demonstrates how to utilize TFDS. Define a function that loads and prepares the MNIST dataset and converts the\\n\",\n \"samples to floating-point numbers.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 48,\n \"id\": \"bRlrHqZVXZvk\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"import tensorflow_datasets as tfds # TFDS for MNIST\\n\",\n \"import tensorflow as tf # TensorFlow operations\\n\",\n \"\\n\",\n \"def get_datasets(num_epochs, batch_size):\\n\",\n \" \\\"\\\"\\\"Load MNIST train and test datasets into memory.\\\"\\\"\\\"\\n\",\n \" train_ds = tfds.load('mnist', split='train')\\n\",\n \" test_ds = tfds.load('mnist', split='test')\\n\",\n \"\\n\",\n \" train_ds = train_ds.map(lambda sample: {'image': tf.cast(sample['image'],\\n\",\n \" tf.float32) / 255.,\\n\",\n \" 'label': sample['label']}) # normalize train set\\n\",\n \" test_ds = test_ds.map(lambda sample: {'image': tf.cast(sample['image'],\\n\",\n \" tf.float32) / 255.,\\n\",\n \" 'label': sample['label']}) # normalize test set\\n\",\n \"\\n\",\n \" train_ds = train_ds.repeat(num_epochs).shuffle(1024) # create shuffled dataset by allocating a buffer size of 1024 to randomly draw elements from\\n\",\n \" train_ds = train_ds.batch(batch_size, drop_remainder=True).prefetch(1) # group into batches of batch_size and skip incomplete batch, prefetch the next sample to improve latency\\n\",\n \" test_ds = test_ds.shuffle(1024) # create shuffled dataset by allocating a buffer size of 1024 to randomly draw elements from\\n\",\n \" test_ds = test_ds.batch(batch_size, drop_remainder=True).prefetch(1) # group into batches of batch_size and skip incomplete batch, prefetch the next sample to improve latency\\n\",\n \"\\n\",\n \" return train_ds, test_ds\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"7057395a\",\n \"metadata\": {},\n \"source\": [\n \"## 3. Define network\\n\",\n \"\\n\",\n \"Create a convolutional neural network with the Linen API by subclassing\\n\",\n \"[Flax Module](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html).\\n\",\n \"Because the architecture in this example is relatively simple—you're just\\n\",\n \"stacking layers—you can define the inlined submodules directly within the\\n\",\n \"`__call__` method and wrap it with the\\n\",\n \"[`@compact`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/decorators.html#flax.linen.compact)\\n\",\n \"decorator. To learn more about the Flax Linen `@compact` decorator, refer to the [`setup` vs `compact`](https://flax.readthedocs.io/en/latest/guides/setup_or_nncompact.html) guide.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 49,\n \"id\": \"cbc079cd\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"from flax import linen as nn # Linen API\\n\",\n \"\\n\",\n \"class CNN(nn.Module):\\n\",\n \" \\\"\\\"\\\"A simple CNN model.\\\"\\\"\\\"\\n\",\n \"\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self, x):\\n\",\n \" x = nn.Conv(features=32, kernel_size=(3, 3))(x)\\n\",\n \" x = nn.relu(x)\\n\",\n \" x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))\\n\",\n \" x = nn.Conv(features=64, kernel_size=(3, 3))(x)\\n\",\n \" x = nn.relu(x)\\n\",\n \" x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))\\n\",\n \" x = x.reshape((x.shape[0], -1)) # flatten\\n\",\n \" x = nn.Dense(features=256)(x)\\n\",\n \" x = nn.relu(x)\\n\",\n \" x = nn.Dense(features=10)(x)\\n\",\n \" return x\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"hy7iRu7_zlx-\",\n \"metadata\": {},\n \"source\": [\n \"### View model layers\\n\",\n \"\\n\",\n \"Create an instance of the Flax Module and use the [`Module.tabulate`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.tabulate) method to visualize a table of the model layers by passing an RNG key and template image input.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 50,\n \"id\": \"lDHfog81zLQa\",\n \"metadata\": {\n \"outputId\": \"2c580f41-bf5d-40ec-f1cf-ab7f319a84da\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"\\n\",\n \"\\u001b[3m CNN Summary \\u001b[0m\\n\",\n \"┏━━━━━━━━━┳━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━━━━━━┳━━━━━━━━━┳━━━━━━━━━━━┳━━━━━━━━━━━━┓\\n\",\n \"┃\\u001b[1m \\u001b[0m\\u001b[1mpath \\u001b[0m\\u001b[1m \\u001b[0m┃\\u001b[1m \\u001b[0m\\u001b[1mmodule\\u001b[0m\\u001b[1m \\u001b[0m┃\\u001b[1m \\u001b[0m\\u001b[1minputs \\u001b[0m\\u001b[1m \\u001b[0m┃\\u001b[1m \\u001b[0m\\u001b[1moutputs \\u001b[0m\\u001b[1m \\u001b[0m┃\\u001b[1m \\u001b[0m\\u001b[1mflops \\u001b[0m\\u001b[1m \\u001b[0m┃\\u001b[1m \\u001b[0m\\u001b[1mvjp_flops\\u001b[0m\\u001b[1m \\u001b[0m┃\\u001b[1m \\u001b[0m\\u001b[1mparams \\u001b[0m\\u001b[1m \\u001b[0m┃\\n\",\n \"┡━━━━━━━━━╇━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━━━━━━╇━━━━━━━━━╇━━━━━━━━━━━╇━━━━━━━━━━━━┩\\n\",\n \"│ │ CNN │ \\u001b[2mfloat32\\u001b[0m[1… │ \\u001b[2mfloat32\\u001b[0m[… │ 8708106 │ 26957556 │ │\\n\",\n \"├─────────┼────────┼────────────┼───────────┼─────────┼───────────┼────────────┤\\n\",\n \"│ Conv_0 │ Conv │ \\u001b[2mfloat32\\u001b[0m[1… │ \\u001b[2mfloat32\\u001b[0m[… │ 455424 │ 1341472 │ bias: │\\n\",\n \"│ │ │ │ │ │ │ \\u001b[2mfloat32\\u001b[0m[3… │\\n\",\n \"│ │ │ │ │ │ │ kernel: │\\n\",\n \"│ │ │ │ │ │ │ \\u001b[2mfloat32\\u001b[0m[3… │\\n\",\n \"│ │ │ │ │ │ │ │\\n\",\n \"│ │ │ │ │ │ │ \\u001b[1m320 \\u001b[0m\\u001b[1;2m(1.3 \\u001b[0m │\\n\",\n \"│ │ │ │ │ │ │ \\u001b[1;2mKB)\\u001b[0m │\\n\",\n \"├─────────┼────────┼────────────┼───────────┼─────────┼───────────┼────────────┤\\n\",\n \"│ Conv_1 │ Conv │ \\u001b[2mfloat32\\u001b[0m[1… │ \\u001b[2mfloat32\\u001b[0m[… │ 6566144 │ 19704320 │ bias: │\\n\",\n \"│ │ │ │ │ │ │ \\u001b[2mfloat32\\u001b[0m[6… │\\n\",\n \"│ │ │ │ │ │ │ kernel: │\\n\",\n \"│ │ │ │ │ │ │ \\u001b[2mfloat32\\u001b[0m[3… │\\n\",\n \"│ │ │ │ │ │ │ │\\n\",\n \"│ │ │ │ │ │ │ \\u001b[1m18,496 \\u001b[0m │\\n\",\n \"│ │ │ │ │ │ │ \\u001b[1;2m(74.0 KB)\\u001b[0m │\\n\",\n \"├─────────┼────────┼────────────┼───────────┼─────────┼───────────┼────────────┤\\n\",\n \"│ Dense_0 │ Dense │ \\u001b[2mfloat32\\u001b[0m[1… │ \\u001b[2mfloat32\\u001b[0m[… │ 1605888 │ 5620224 │ bias: │\\n\",\n \"│ │ │ │ │ │ │ \\u001b[2mfloat32\\u001b[0m[2… │\\n\",\n \"│ │ │ │ │ │ │ kernel: │\\n\",\n \"│ │ │ │ │ │ │ \\u001b[2mfloat32\\u001b[0m[3… │\\n\",\n \"│ │ │ │ │ │ │ │\\n\",\n \"│ │ │ │ │ │ │ \\u001b[1m803,072 \\u001b[0m │\\n\",\n \"│ │ │ │ │ │ │ \\u001b[1;2m(3.2 MB)\\u001b[0m │\\n\",\n \"├─────────┼────────┼────────────┼───────────┼─────────┼───────────┼────────────┤\\n\",\n \"│ Dense_1 │ Dense │ \\u001b[2mfloat32\\u001b[0m[1… │ \\u001b[2mfloat32\\u001b[0m[… │ 5130 │ 17940 │ bias: │\\n\",\n \"│ │ │ │ │ │ │ \\u001b[2mfloat32\\u001b[0m[1… │\\n\",\n \"│ │ │ │ │ │ │ kernel: │\\n\",\n \"│ │ │ │ │ │ │ \\u001b[2mfloat32\\u001b[0m[2… │\\n\",\n \"│ │ │ │ │ │ │ │\\n\",\n \"│ │ │ │ │ │ │ \\u001b[1m2,570 \\u001b[0m │\\n\",\n \"│ │ │ │ │ │ │ \\u001b[1;2m(10.3 KB)\\u001b[0m │\\n\",\n \"├─────────┼────────┼────────────┼───────────┼─────────┼───────────┼────────────┤\\n\",\n \"│\\u001b[1m \\u001b[0m\\u001b[1m \\u001b[0m\\u001b[1m \\u001b[0m│\\u001b[1m \\u001b[0m\\u001b[1m \\u001b[0m\\u001b[1m \\u001b[0m│\\u001b[1m \\u001b[0m\\u001b[1m \\u001b[0m\\u001b[1m \\u001b[0m│\\u001b[1m \\u001b[0m\\u001b[1m \\u001b[0m\\u001b[1m \\u001b[0m│\\u001b[1m \\u001b[0m\\u001b[1m \\u001b[0m\\u001b[1m \\u001b[0m│\\u001b[1m \\u001b[0m\\u001b[1m Total\\u001b[0m\\u001b[1m \\u001b[0m│\\u001b[1m \\u001b[0m\\u001b[1m824,458 \\u001b[0m\\u001b[1m \\u001b[0m│\\n\",\n \"│\\u001b[1m \\u001b[0m│\\u001b[1m \\u001b[0m│\\u001b[1m \\u001b[0m│\\u001b[1m \\u001b[0m│\\u001b[1m \\u001b[0m│\\u001b[1m \\u001b[0m│\\u001b[1m \\u001b[0m\\u001b[1;2m(3.3 MB)\\u001b[0m\\u001b[1m \\u001b[0m\\u001b[1m \\u001b[0m│\\n\",\n \"└─────────┴────────┴────────────┴───────────┴─────────┴───────────┴────────────┘\\n\",\n \"\\u001b[1m \\u001b[0m\\n\",\n \"\\u001b[1m Total Parameters: 824,458 \\u001b[0m\\u001b[1;2m(3.3 MB)\\u001b[0m\\u001b[1m \\u001b[0m\\n\",\n \"\\n\",\n \"\\n\"\n ]\n }\n ],\n \"source\": [\n \"import jax\\n\",\n \"import jax.numpy as jnp # JAX NumPy\\n\",\n \"\\n\",\n \"cnn = CNN()\\n\",\n \"print(cnn.tabulate(jax.random.key(0), jnp.ones((1, 28, 28, 1)),\\n\",\n \" compute_flops=True, compute_vjp_flops=True))\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"4b5ac16e\",\n \"metadata\": {},\n \"source\": [\n \"## 4. Create a `TrainState`\\n\",\n \"\\n\",\n \"A common pattern in Flax is to create a single dataclass that represents the\\n\",\n \"entire training state, including step number, parameters, and optimizer state.\\n\",\n \"\\n\",\n \"Because this is such a common pattern, Flax provides the class\\n\",\n \"[`flax.training.train_state.TrainState`](https://flax.readthedocs.io/en/latest/flax.training.html#train-state)\\n\",\n \"that serves most basic usecases.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"qXr7JDpIxGNZ\",\n \"metadata\": {\n \"outputId\": \"1249b7fb-6787-41eb-b34c-61d736300844\"\n },\n \"outputs\": [],\n \"source\": [\n \"!pip install -q clu\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 52,\n \"id\": \"CJDaJNijyOji\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"from clu import metrics\\n\",\n \"from flax.training import train_state # Useful dataclass to keep train state\\n\",\n \"from flax import struct # Flax dataclasses\\n\",\n \"import optax # Common loss functions and optimizers\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"8b86b5f1\",\n \"metadata\": {},\n \"source\": [\n \"We will be using the `clu` library for computing metrics. For more information on `clu`, refer to the [repo](https://github.com/google/CommonLoopUtils) and [notebook](https://colab.research.google.com/github/google/CommonLoopUtils/blob/master/clu_synopsis.ipynb#scrollTo=ueom-uBWLbeQ).\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 53,\n \"id\": \"7W0qf7FC9uG5\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"@struct.dataclass\\n\",\n \"class Metrics(metrics.Collection):\\n\",\n \" accuracy: metrics.Accuracy\\n\",\n \" loss: metrics.Average.from_output('loss')\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"f3ce5e4c\",\n \"metadata\": {},\n \"source\": [\n \"You can then subclass `train_state.TrainState` so that it also contains metrics. This has the advantage that we only need\\n\",\n \"to pass around a single argument to functions like `train_step()` (see below) to calculate the loss, update the parameters and compute the metrics all at once.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 54,\n \"id\": \"e0102447\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"class TrainState(train_state.TrainState):\\n\",\n \" metrics: Metrics\\n\",\n \"\\n\",\n \"def create_train_state(module, rng, learning_rate, momentum):\\n\",\n \" \\\"\\\"\\\"Creates an initial `TrainState`.\\\"\\\"\\\"\\n\",\n \" params = module.init(rng, jnp.ones([1, 28, 28, 1]))['params'] # initialize parameters by passing a template image\\n\",\n \" tx = optax.sgd(learning_rate, momentum)\\n\",\n \" return TrainState.create(\\n\",\n \" apply_fn=module.apply, params=params, tx=tx,\\n\",\n \" metrics=Metrics.empty())\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"a15de484\",\n \"metadata\": {},\n \"source\": [\n \"## 5. Training step\\n\",\n \"\\n\",\n \"A function that:\\n\",\n \"\\n\",\n \"- Evaluates the neural network given the parameters and a batch of input images\\n\",\n \" with [`TrainState.apply_fn`](https://flax.readthedocs.io/en/latest/api_reference/flax.training.html#flax.training.train_state.TrainState) (which contains the [`Module.apply`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.apply)\\n\",\n \" method (forward pass)).\\n\",\n \"- Computes the cross entropy loss, using the predefined [`optax.softmax_cross_entropy_with_integer_labels()`](https://optax.readthedocs.io/en/latest/api.html#optax.softmax_cross_entropy_with_integer_labels). Note that this function expects integer labels, so there is no need to convert labels to onehot encoding.\\n\",\n \"- Evaluates the gradient of the loss function using\\n\",\n \" [`jax.grad`](https://jax.readthedocs.io/en/latest/jax.html#jax.grad).\\n\",\n \"- Applies a\\n\",\n \" [pytree](https://jax.readthedocs.io/en/latest/pytrees.html#pytrees-and-jax-functions)\\n\",\n \" of gradients to the optimizer to update the model's parameters.\\n\",\n \"\\n\",\n \"Use JAX's [@jit](https://jax.readthedocs.io/en/latest/jax.html#jax.jit)\\n\",\n \"decorator to trace the entire `train_step` function and just-in-time compile\\n\",\n \"it with [XLA](https://www.tensorflow.org/xla) into fused device operations\\n\",\n \"that run faster and more efficiently on hardware accelerators.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 55,\n \"id\": \"9b0af486\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"@jax.jit\\n\",\n \"def train_step(state, batch):\\n\",\n \" \\\"\\\"\\\"Train for a single step.\\\"\\\"\\\"\\n\",\n \" def loss_fn(params):\\n\",\n \" logits = state.apply_fn({'params': params}, batch['image'])\\n\",\n \" loss = optax.softmax_cross_entropy_with_integer_labels(\\n\",\n \" logits=logits, labels=batch['label']).mean()\\n\",\n \" return loss\\n\",\n \" grad_fn = jax.grad(loss_fn)\\n\",\n \" grads = grad_fn(state.params)\\n\",\n \" state = state.apply_gradients(grads=grads)\\n\",\n \" return state\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"0ff5145f\",\n \"metadata\": {},\n \"source\": [\n \"## 6. Metric computation\\n\",\n \"\\n\",\n \"Create a separate function for loss and accuracy metrics. Loss is calculated using the `optax.softmax_cross_entropy_with_integer_labels` function, while accuracy is calculated using `clu.metrics`.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 56,\n \"id\": \"961bf70b\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"@jax.jit\\n\",\n \"def compute_metrics(*, state, batch):\\n\",\n \" logits = state.apply_fn({'params': state.params}, batch['image'])\\n\",\n \" loss = optax.softmax_cross_entropy_with_integer_labels(\\n\",\n \" logits=logits, labels=batch['label']).mean()\\n\",\n \" metric_updates = state.metrics.single_from_model_output(\\n\",\n \" logits=logits, labels=batch['label'], loss=loss)\\n\",\n \" metrics = state.metrics.merge(metric_updates)\\n\",\n \" state = state.replace(metrics=metrics)\\n\",\n \" return state\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"497241c3\",\n \"metadata\": {},\n \"source\": [\n \"## 7. Download data\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 57,\n \"id\": \"bff5393e\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"num_epochs = 10\\n\",\n \"batch_size = 32\\n\",\n \"\\n\",\n \"train_ds, test_ds = get_datasets(num_epochs, batch_size)\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"id\": \"809ae1a0\",\n \"metadata\": {},\n \"source\": [\n \"## 8. Seed randomness\\n\",\n \"\\n\",\n \"- Set the TF random seed to ensure dataset shuffling (with `tf.data.Dataset.shuffle`) is reproducible.\\n\",\n \"- Get one\\n\",\n \" [PRNGKey](https://jax.readthedocs.io/en/latest/_autosummary/jax.random.PRNGKey.html#jax.random.PRNGKey)\\n\",\n \" and use it for parameter initialization. (Learn\\n\",\n \" more about\\n\",\n \" [JAX PRNG design](https://jax.readthedocs.io/en/latest/jax-101/05-random-numbers.html)\\n\",\n \" and [PRNG chains](https://flax.readthedocs.io/en/latest/philosophy.html#how-are-parameters-represented-and-how-do-we-handle-general-differentiable-algorithms-that-update-stateful-variables).)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 58,\n \"id\": \"xC4MFyBsfT-U\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"tf.random.set_seed(0)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 59,\n \"id\": \"e4f6f4d3\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"init_rng = jax.random.key(0)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"80fbb60b\",\n \"metadata\": {},\n \"source\": [\n \"## 9. Initialize the `TrainState`\\n\",\n \"\\n\",\n \"Remember that the function `create_train_state` initializes the model parameters, optimizer and metrics\\n\",\n \"and puts them into the training state dataclass that is returned.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 60,\n \"id\": \"445fcab0\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"learning_rate = 0.01\\n\",\n \"momentum = 0.9\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 61,\n \"id\": \"5221eafd\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"state = create_train_state(cnn, init_rng, learning_rate, momentum)\\n\",\n \"del init_rng # Must not be used anymore.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"b1c00230\",\n \"metadata\": {},\n \"source\": [\n \"## 10. Train and evaluate\\n\",\n \"\\n\",\n \"Create a \\\"shuffled\\\" dataset by:\\n\",\n \"- Repeating the dataset equal to the number of training epochs\\n\",\n \"- Allocating a buffer of size 1024 (containing the first 1024 samples in the dataset) of which to randomly sample batches from\\n\",\n \" - Everytime a sample is randomly drawn from the buffer, the next sample in the dataset is loaded into the buffer\\n\",\n \"\\n\",\n \"Define a training loop that:\\n\",\n \"- Randomly samples batches from the dataset.\\n\",\n \"- Runs an optimization step for each training batch.\\n\",\n \"- Computes the mean training metrics across each batch in an epoch.\\n\",\n \"- Computes the metrics for the test set using the updated parameters.\\n\",\n \"- Records the train and test metrics for visualization.\\n\",\n \"\\n\",\n \"Once the training and testing is done after 10 epochs, the output should show that your model was able to achieve approximately 99% accuracy.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 62,\n \"id\": \"74295360\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# since train_ds is replicated num_epochs times in get_datasets(), we divide by num_epochs\\n\",\n \"num_steps_per_epoch = train_ds.cardinality().numpy() // num_epochs\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 63,\n \"id\": \"cRtnMZuQFlKl\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"metrics_history = {'train_loss': [],\\n\",\n \" 'train_accuracy': [],\\n\",\n \" 'test_loss': [],\\n\",\n \" 'test_accuracy': []}\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 64,\n \"id\": \"2c40ce90\",\n \"metadata\": {\n \"outputId\": \"258a2c76-2c8f-4a9e-d48b-dde57c342a87\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"train epoch: 1, loss: 0.20290373265743256, accuracy: 93.87000274658203\\n\",\n \"test epoch: 1, loss: 0.07591685652732849, accuracy: 97.60617065429688\\n\",\n \"train epoch: 2, loss: 0.05760224163532257, accuracy: 98.28500366210938\\n\",\n \"test epoch: 2, loss: 0.050395529717206955, accuracy: 98.3974380493164\\n\",\n \"train epoch: 3, loss: 0.03897436335682869, accuracy: 98.83000183105469\\n\",\n \"test epoch: 3, loss: 0.04574578255414963, accuracy: 98.54767608642578\\n\",\n \"train epoch: 4, loss: 0.028721099719405174, accuracy: 99.15166473388672\\n\",\n \"test epoch: 4, loss: 0.035722777247428894, accuracy: 98.91827392578125\\n\",\n \"train epoch: 5, loss: 0.021948494017124176, accuracy: 99.37999725341797\\n\",\n \"test epoch: 5, loss: 0.035723842680454254, accuracy: 98.87820434570312\\n\",\n \"train epoch: 6, loss: 0.01705147698521614, accuracy: 99.54833221435547\\n\",\n \"test epoch: 6, loss: 0.03456473350524902, accuracy: 98.96835327148438\\n\",\n \"train epoch: 7, loss: 0.014007646590471268, accuracy: 99.6116714477539\\n\",\n \"test epoch: 7, loss: 0.04089202359318733, accuracy: 98.7880630493164\\n\",\n \"train epoch: 8, loss: 0.011265480890870094, accuracy: 99.73333740234375\\n\",\n \"test epoch: 8, loss: 0.03337760642170906, accuracy: 98.93830108642578\\n\",\n \"train epoch: 9, loss: 0.00918484665453434, accuracy: 99.78334045410156\\n\",\n \"test epoch: 9, loss: 0.034478139132261276, accuracy: 98.96835327148438\\n\",\n \"train epoch: 10, loss: 0.007260234095156193, accuracy: 99.84166717529297\\n\",\n \"test epoch: 10, loss: 0.032822880893945694, accuracy: 99.07852172851562\\n\"\n ]\n }\n ],\n \"source\": [\n \"for step,batch in enumerate(train_ds.as_numpy_iterator()):\\n\",\n \"\\n\",\n \" # Run optimization steps over training batches and compute batch metrics\\n\",\n \" state = train_step(state, batch) # get updated train state (which contains the updated parameters)\\n\",\n \" state = compute_metrics(state=state, batch=batch) # aggregate batch metrics\\n\",\n \"\\n\",\n \" if (step+1) % num_steps_per_epoch == 0: # one training epoch has passed\\n\",\n \" for metric,value in state.metrics.compute().items(): # compute metrics\\n\",\n \" metrics_history[f'train_{metric}'].append(value) # record metrics\\n\",\n \" state = state.replace(metrics=state.metrics.empty()) # reset train_metrics for next training epoch\\n\",\n \"\\n\",\n \" # Compute metrics on the test set after each training epoch\\n\",\n \" test_state = state\\n\",\n \" for test_batch in test_ds.as_numpy_iterator():\\n\",\n \" test_state = compute_metrics(state=test_state, batch=test_batch)\\n\",\n \"\\n\",\n \" for metric,value in test_state.metrics.compute().items():\\n\",\n \" metrics_history[f'test_{metric}'].append(value)\\n\",\n \"\\n\",\n \" print(f\\\"train epoch: {(step+1) // num_steps_per_epoch}, \\\"\\n\",\n \" f\\\"loss: {metrics_history['train_loss'][-1]}, \\\"\\n\",\n \" f\\\"accuracy: {metrics_history['train_accuracy'][-1] * 100}\\\")\\n\",\n \" print(f\\\"test epoch: {(step+1) // num_steps_per_epoch}, \\\"\\n\",\n \" f\\\"loss: {metrics_history['test_loss'][-1]}, \\\"\\n\",\n \" f\\\"accuracy: {metrics_history['test_accuracy'][-1] * 100}\\\")\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"gfsecJzvzgCT\",\n \"metadata\": {},\n \"source\": [\n \"## 11. Visualize metrics\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 65,\n \"id\": \"Zs5atiqIG9Kz\",\n \"metadata\": {\n \"outputId\": \"431a2fcd-44fa-4202-f55a-906555f060ac\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"image/png\": \"iVBORw0KGgoAAAANSUhEUgAAA3cAAAE/CAYAAADlpzo+AAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjYuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/av/WaAAAACXBIWXMAAAsTAAALEwEAmpwYAABsiElEQVR4nO3dd3yddd3/8dcneyfNaJs26aJ7JAFKyxLUKlBWAQEBmQK9uRUEb/EWt7hufooDFcEyBVFUFK1QLFBEZLeFpLt00qRJ23Rk7+T7++O6kp6maXPSjJPxfj4e53HONc/3nJ7mOu/zXeacQ0RERERERAa2sFAXQERERERERLpP4U5ERERERGQQULgTEREREREZBBTuREREREREBgGFOxERERERkUFA4U5ERERERGQQULgTEREREREZBBTuRPqYmW03s0+EuhwiIiK9ycxeNbMDZhYd6rKIDBUKdyIiIiLSo8xsHPARwAEX9uHzRvTVc4n0Rwp3Iv2AmUWb2c/NrNi//bz1l04zSzez58yszMz2m9l/zCzM3/YVM9tpZpVmttHM5oX2lYiIiABwLfA28DhwXetKM8s2s7+aWamZ7TOzXwVsu9nM1vvXtHVmdoK/3pnZxID9Hjez7/uPP2pmRf71cBfwmJkN86+bpX7N4XNmlhVwfKqZPeZfbw+Y2d/89WvM7IKA/SLNbK+Z5fXSeyTS4xTuRPqHrwMnA3lALjAH+Ia/7UtAEZABjAC+BjgzmwLcCpzknEsEzga292mpRUREOnYt8JR/O9vMRphZOPAc8CEwDhgNPA1gZpcB3/GPS8Kr7dsX5HONBFKBscBCvO+3j/nLY4Ba4FcB+z8JxAEzgOHAz/z1TwBXB+x3LlDinMsPshwiIaeqa5H+4TPAbc65PQBmdjfwG+CbQCOQCYx1zm0G/uPv0wxEA9PNrNQ5tz0UBRcREQlkZqfjBas/Oef2mtkW4Cq8mrxRwJedc03+7q/79zcBP3LOLfeXN3fhKVuAbzvn6v3lWuAvAeX5AfAv/3EmMB9Ic84d8Hf5t3//O+CbZpbknKsArsELgiIDhmruRPqHUXi/ZLb60F8H8GO8i9yLZrbVzO4C8IPeHXi/dO4xs6fNbBQiIiKhdR3wonNur7/8e39dNvBhQLALlA1sOcbnK3XO1bUumFmcmf3GzD40swrgNSDFrznMBvYHBLs2zrli4A3gU2aWghcCnzrGMomEhMKdSP9QjPcrZ6sx/jqcc5XOuS855yYAFwD/09q3zjn3e+dc6y+kDvh/fVtsERGRg8wsFrgcONPMdvn94L6I1+VgNzDmCIOeFALHHeG0NXjNKFuNbLfdtVv+EjAFmOucSwLOaC2e/zypfnjryG/xmmZeBrzlnNt5hP1E+iWFO5HQiDSzmNYb8AfgG2aWYWbpwLfwmodgZueb2UQzM6ACaAaazWyKmX3cH3ilDq8ZSnNoXo6IiAgAF+Fdi6bj9SPPA6bhdSm4CCgB7jGzeP8aeJp/3MPAnWZ2onkmmlnrj575wFVmFm5m5wBndlKGRLxrYpmZpQLfbt3gnCsBXgB+7Q+8EmlmZwQc+zfgBOB2vD54IgOKwp1IaCzBu/C03mKAFcAqYDXwHvB9f99JwMtAFfAW8Gvn3Kt4/e3uAfYCu/A6hX+tz16BiIjI4a4DHnPO7XDO7Wq94Q1ociVeC5SJwA68wcI+DeCc+zPwA7wmnJV4ISvVP+ft/nFleH3U/9ZJGX4OxOJdH98G/tlu+zV4/dk3AHvwujjgl6O1v9544K/Bv2yR/sGca1+TLSIiIiIyNJnZt4DJzrmrO91ZpJ/RaJkiIiIiInhz4AE34tXuiQw4apYpIiIiIkOemd2MN+DKC86510JdHpFjoWaZIiIiIiIig4Bq7kRERERERAYBhTsREREREZFBYEANqJKenu7GjRsX6mKIiEgvW7ly5V7nXEaoyzFQ6PooIjJ0HO0aOaDC3bhx41ixYkWoiyEiIr3MzD4MdRkGEl0fRUSGjqNdI9UsU0REpIeZ2aNmtsfM1hxhu5nZL8xss5mtMrMTAradY2Yb/W139V2pRURkoFO4ExER6XmPA+ccZft8YJJ/Wwg8AGBm4cD9/vbpwJVmNr1XSyoiIoOGwp2IiEgP8+fI2n+UXRYATzjP20CKmWUCc4DNzrmtzrkG4Gl/XxERkU4NqD53IiL9RWNjI0VFRdTV1YW6KANaTEwMWVlZREZGhroofW003mTJrYr8dR2tn3ssT6DP6MAzhP8/iEgPUbgTETkGRUVFJCYmMm7cOMws1MUZkJxz7Nu3j6KiIsaPHx/q4vS1jj407ijrDz+B2UK8Jp2MGTPmsO36jA4sQ/z/g4j0EDXLFBE5BnV1daSlpelLczeYGWlpaUO1ZqkIyA5YzgKKj7L+MM65Rc652c652RkZh4+Irc/owDLE/z+ISA9RuBMROUb60tx9Q/g9XAxc64+aeTJQ7pwrAZYDk8xsvJlFAVf4+x6TIfz+Dkj69xKR7lKzTBERkR5mZn8APgqkm1kR8G0gEsA59yCwBDgX2AzUADf425rM7FZgKRAOPOqcW9vnL0BERAYk1dyJiAxAZWVl/PrXv+7yceeeey5lZWVdPu7666/nmWee6fJxQ5Vz7krnXKZzLtI5l+Wce8Q596Af7PBHyfy8c+4459ws59yKgGOXOOcm+9t+ELpX0X19/TkVERnqhlS4W1tczlPvHHFCdxGRAeNIX5qbm5uPetySJUtISUnppVKJHGqwfk47K7+IDG1NzS0cqG7gw33VrC4q5/VNe1myuoSn393Bpt2VvfrcQTXLNLNzgPvwmog87Jy7p932zwBf8RergP92zhUc7VgzSwX+CIwDtgOXO+cOdPP1HNVL63Zz37JNLMgbTUK0WqSKyMB11113sWXLFvLy8oiMjCQhIYHMzEzy8/NZt24dF110EYWFhdTV1XH77bezcOFCAMaNG8eKFSuoqqpi/vz5nH766bz55puMHj2av//978TGxnb63MuWLePOO++kqamJk046iQceeIDo6GjuuusuFi9eTEREBGeddRb33nsvf/7zn7n77rsJDw8nOTmZ1157rbffGulH+vpz+tBDD7Fo0SIaGhqYOHEiTz75JHFxcezevZtbbrmFrVu3AvDAAw9w6qmn8sQTT3DvvfdiZuTk5PDkk09y/fXXc/7553PppZcCkJCQQFVVFa+++ip33313UOX/5z//yde+9jWam5tJT0/npZdeYsqUKbz55ptkZGTQ0tLC5MmTefvtt0lPT++DfwkR6YrmFkdlXSMVtU1U1DVSUdvo3zdR3va4kYq6pkO2ta6vbjjyD0DfXTCDSSMSe63snSYcMwsH7gc+iTeK13IzW+ycWxew2zbgTOfcATObDywC5nZy7F3AMufcPWZ2l7/8FXpRbnYKzsHqonJOOS6tN59KRIaQu/+xlnXFFT16zumjkvj2BTOOuP2ee+5hzZo15Ofn8+qrr3LeeeexZs2atiHUH330UVJTU6mtreWkk07iU5/6FGlph/7d27RpE3/4wx946KGHuPzyy/nLX/7C1VdffdRy1dXVcf3117Ns2TImT57MtddeywMPPMC1117Ls88+y4YNGzCztiZ13/3ud1m6dCmjR49WM7sQCsVnFPr+c3rJJZdw8803A/CNb3yDRx55hNtuu40vfOELnHnmmTz77LM0NzdTVVXF2rVr+cEPfsAbb7xBeno6+/cfbc55z7vvvttp+VtaWrj55pt57bXXGD9+PPv37ycsLIyrr76ap556ijvuuIOXX36Z3NxcBTuRXtLc4qiq88JWeQfh69BQdvi2qvqmo54/zCAxJpKk2AiSYiJJiolkXHocybHe46TYSJJiIvx7f9nfNzU+qldfezDVV3OAzc65rQBm9jSwAGgLd865NwP2fxtv6ObOjl2A19kc4LfAq/R2uMtKAWBVUZnCnYgMKnPmzDlkbqxf/OIXPPvsswAUFhayadOmw740jx8/nry8PABOPPFEtm/f3unzbNy4kfHjxzN58mQArrvuOu6//35uvfVWYmJiuOmmmzjvvPM4//zzATjttNO4/vrrufzyy7nkkkt64JXKQNbbn9M1a9bwjW98g7KyMqqqqjj77LMBeOWVV3jiiScA2mqRn3jiCS699NK2gJWamtoj5S8tLeWMM85o26/1vJ/97GdZsGABd9xxB48++ig33HBDp88nMpQ456hvaqGqvonq+iYq/ZBV1Xpff+hyZZ23X1V9E5X1TVTVNVJd39y239GYQWJ0YPiKYExq3CHLHQY1/3F8VARhYf1zdNtgwt1ooDBguQiYe5T9bwReCOLYEf6wzzjnSsxseFAl7obU+CiyU2MpKCrr7acSkSGks9qLvhAfH9/2+NVXX+Xll1/mrbfeIi4ujo9+9KMdzp0VHR3d9jg8PJza2tpOn8e5DufTJiIignfffZdly5bx9NNP86tf/YpXXnmFBx98kHfeeYfnn3+evLw88vPzD/vyLr2vP3xGofc/p9dffz1/+9vfyM3N5fHHH+fVV1894r7OuQ6nHoiIiKClpaVtn4aGhi6V/0jnzc7OZsSIEbzyyiu88847PPXUU0csm8hA0tzi2gLV4aGskar65oDHTf5yY1tAaz2uqr6JxuaOrzGBwgwSoiNIjIkkITqChBgviGWlxLYtx0e3hrODAS65tfYsNpKEfhzOuiuYcNfRK+/wnTezj+GFu9O7euwRn9xsIbAQYMyYMV05tEO5WSm8v6Os2+cREQmlxMREKis77pRdXl7OsGHDiIuLY8OGDbz99ts99rxTp05l+/btbN68ua1P05lnnklVVRU1NTWce+65nHzyyUycOBGALVu2MHfuXObOncs//vEPCgsLFe6GkL7+nFZWVpKZmUljYyNPPfUUo0ePBmDevHk88MAD3HHHHTQ3N1NdXc28efO4+OKL+eIXv0haWhr79+8nNTWVcePGsXLlSi6//HL+/ve/09jY2KXyn3LKKXz+859n27Ztbc0yW2vvbrrpJq6++mquueYawsPDu/16RXqbc47SynoKD9SwY38NO/bVtj3eeaCWAzUN1Bylf1mg2MhwEmIiSIz2wldCdATZqXEkBgQyL7R59223mIP3idGRxESGaU7Iowgm3BUB2QHLWUBx+53MLAd4GJjvnNsXxLG7zSzTr7XLBPZ09OTOuUV4ffiYPXt2l4JhR3KzUnhuVQmllfVkJEZ3foCISD+UlpbGaaedxsyZM4mNjWXEiBFt28455xwefPBBcnJymDJlCieffHKPPW9MTAyPPfYYl112WduAKrfccgv79+9nwYIFbTUXP/vZzwD48pe/zKZNm3DOMW/ePHJzc3usLNL/9fXn9Hvf+x5z585l7NixzJo1qy1Y3nfffSxcuJBHHnmE8PBwHnjgAU455RS+/vWvc+aZZxIeHs7xxx/P448/zs0338yCBQuYM2cO8+bNO6S2LtCRyp+RkcGiRYu45JJLaGlpYfjw4bz00ksAXHjhhdxwww1qkin9SlV9E4X7vcBW2Ho7UMuO/TUUHaihrrHlkP1HJEWTPSyOOeNTSYuPagtfiYcFtEhvW1QE8dHhRIQPqUH6Q8aO1MSmbQezCOADYB6wE1gOXBU4qaqZjQFeAa4N7H93tGPN7MfAvoABVVKdc/97tLLMnj3brVix4mi7dOrdbfu5/Ddv8ch1s5k3bUTnB4iIdGD9+vVMmzYt1MUYFDp6L81spXNudoiKNOB0dH3UZ7T/WbFiBV/84hf5z3/+c8R99O8mPa2xuYWSsrq2GrdDgtyBWvZXNxyyf2uN2pjUWLKHxTEmLY7sYXFkp8aRNSyWmEjVOofa0a6RndbcOeeazOxWYCnedAaP+uHsFn/7g8C3gDTg1341aZNzbvaRjvVPfQ/wJzO7EdgBXNatVxmkmaOTCDMoKCpXuBMREZE+cc899/DAAw+or530OOcc+6ob2kJb0YFaduyraQtzJeV1NLccrMyJCDNGD4tlTGocZ49KZkxqHNmp3nL2sDhS4iLV7HEAC2qyN+fcEmBJu3UPBjy+Cbgp2GP99fvwavT6VFxUBJNHJFJQWNbXTy0i0u99/vOf54033jhk3e23365mZNKvDMTP6V133cVdd90V6mLIAFXb0EzhgcBat4PNJnfsrzms31t6QjRjUmM5ceywttCW7Ye4kUkxaiI5iA3Jmbxzs1JYum7XEUe0EhEZqu6///5QF0GkU/qcykDV3OKoqG2krNabf628tpGymgZvXY2/3LrNX95X3cDeqvpDzhMXFd4W2E45Lo0xqXF+DZzXdDIuakh+xReGaLjLyU7mjysKKdxfy5i0uFAXR0REREQGCOcc1Q3NlNU0HBLCytrCmj9xdm0jZbUNh6yrrDv6/GtxUeEkx0a23calx3H8mJS20NYa4NLio1RBIR0akuGudTLz/KIyhTsRERGRIayusZnNe6rYVV4XUGvWcFhgqwiobWtqOfKAhJHhdkhAG54Yw6ThiW3LKXGH3ifHRrVti4pQc0npniEZ7qaMTCQ6IoxVhWVcmDsq1MURERERkV7mnGN3RT3rSypYv6uCDSWVrC+pYOve6kMGHAEwg8ToCFLiotqC2OhhsaTEHjmYta6LiwpXrZqEzJAMd5HhYcwYlURBUVmoiyIiIjJolZWV8fvf/57Pfe5zXT725z//OQsXLiQuTi1spOvqGpvZtLuqLcitL6lgw65KymoOTko/OiWWaZmJnD1jJNMyk8gaFtsW0BJjIgkPU0CTgWdIhjuA3OwU/vDuDpqaWzRikIgMOL39pXncuHGsWLGC9PT07hRThriysjJ+/etfH/Pn9Oqrr+4X4a6pqYmIiCH7lalfc85RUl7Hhl0VrPdr4jbsqmRraRWtlXGxkeFMGZnI/JkjmToyiWmZSUwZ6TWTFBlshuxfqtysFB57Yzub9lQxLTMp1MUREemSwfKlWQa3u+66iy1btpCXl8cnP/lJhg8fzp/+9Cfq6+u5+OKLufvuu6murubyyy+nqKiI5uZmvvnNb7J7926Ki4v52Mc+Rnp6Ov/61786PP9///d/s3z5cmpra7n00ku5++67AVi+fDm333471dXVREdHs2zZMuLi4vjKV77C0qVLMTNuvvlmbrvttkN+yFixYgV33nknr776Kt/5zncoLi5m+/btpKen88Mf/pBrrrmG6upqAH71q19x6qmnAvCjH/2IJ598krCwMObPn8/NN9/MZZddxnvvvQfApk2buOKKK1i5cmUfvOuDV21DMx/srjwsyJXXHqyNyxoWy7TMJM6dOZKpmV6QG5Map1o4GTKGbrjLTgGgoLBM4U5EuueFu2DX6p4958hZMP+eI27u7S/NgX7605/y6KOPAnDTTTdxxx13dHjuT3/609x1110sXryYiIgIzjrrLO69994ee0ukG0LwGQVv4u41a9aQn5/Piy++yDPPPMO7776Lc44LL7yQ1157jdLSUkaNGsXzzz8PQHl5OcnJyfz0pz/lX//611Frj3/wgx+QmppKc3Mz8+bNY9WqVUydOpVPf/rT/PGPf+Skk06ioqKC2NhYFi1axLZt23j//feJiIhg//79nb7ElStX8vrrrxMbG0tNTQ0vvfQSMTExbNq0iSuvvJIVK1bwwgsv8Le//Y133nmHuLg49u/fT2pqKsnJyeTn55OXl8djjz3G9ddf36W3dyhzzrGzrJYNJQFBblcF2/dWt9XGxUV5tXHn5WQybWQiU/3auKQY1cbJ0DZkw924tDiSYiIoKCrjijljQl0cEZEu6e0vza1WrlzJY489xjvvvINzjrlz53LmmWeydevWw869f/9+nn32WTZs2ICZUVZW1ptvgQwwL774Ii+++CLHH388AFVVVWzatImPfOQj3HnnnXzlK1/h/PPP5yMf+UjQ5/zTn/7EokWLaGpqoqSkhHXr1mFmZGZmctJJJwGQlOT9gPvyyy9zyy23tDWvTE1N7fT8F154IbGxsQA0NjZy6623kp+fT3h4OB988EHbeW+44Ya2mvDW895000089thj/PSnP+WPf/wj7777btCvayipaWhi465KNuyqZEPJwSAXOGXAmNQ4po5M5IKcUUzLTGTqSK82Lky1cSKHGbLhzszIzU6hoLA81EURkYGuk9qL3tYbX5pbvf7661x88cXEx8cDcMkll/Cf//yHc84557BzNzU1ERMTw0033cR5553H+eef36OvU7ohxJ9R8GpjvvrVr/Jf//Vfh21buXIlS5Ys4atf/SpnnXUW3/rWtzo937Zt27j33ntZvnw5w4YN4/rrr6eurg7nXIcjFR5pfUREBC0tLQDU1dUdsq31cw/ws5/9jBEjRlBQUEBLSwsxMTFHPe+nPvUp7r77bj7+8Y9z4oknkpaW1ulrGuz2Vzew8sMDfnNKL8ht31eN82vj4qPCmZqZxIW5o5iWmcS0zESmjEwiIXrIfl0V6bIh/b8lNyuFB/69hdqGZmKjwkNdHBGRY9LTX5rbn7sjkydP7vDc7777LsuWLePpp5/mV7/6Fa+88soxvSYZHBITE6msrATg7LPP5pvf/Caf+cxnSEhIYOfOnURGRtLU1ERqaipXX301CQkJPP7444cce6Qa5oqKCuLj40lOTmb37t288MILfPSjH2Xq1KkUFxezfPlyTjrpJCorK4mNjeWss87iwQcf5KMf/Whbs8zU1FTGjRvHypUrmT9/Pn/5y1+O+FrKy8vJysoiLCyM3/72tzQ3NwNw1lln8d3vfperrrrqkGaZMTExnH322fz3f/83jzzySM++sQNEaWU972zbxztb9/Putv1s3F3Ztm1cWhxTRyaxIM8PciO90SpVGyfSPUM63OVkJdPc4lhXUs6JYztvniEi0l/05pfmQGeccQbXX389d911F845nn32WZ588kmKi4sPO3dVVRU1NTWce+65nHzyyUycOLE33wIZANLS0jjttNOYOXMm8+fP56qrruKUU04BICEhgd/97nds3ryZL3/5y4SFhREZGckDDzwAwMKFC5k/fz6ZmZkd9g3Nzc3l+OOPZ8aMGUyYMIHTTjsNgKioKP74xz9y2223UVtbS2xsLC+//DI33XQTH3zwATk5OURGRnLzzTdz66238u1vf5sbb7yRH/7wh8ydO/eIr+Vzn/scn/rUp/jzn//Mxz72sbZavXPOOYf8/Hxmz55NVFQU5557Lj/84Q8B+MxnPsNf//pXzjrrrB59X/urXeV1vLNtH29v3c872/axtdQbfCYuKpwTxw7jwrxRzBmfyvTMJOJVGyfSK+xIv8r2R7Nnz3YrVqzosfPtqahjzg+X8c3zp3Pj6eN77LwiMvitX7+eadOmhbQMV111FatWrWL+/PlkZWXx8MMPA0f/0jx79mx++ctfcv/99x/xSzMcOhVCRwOqLF269LBzjx49mgULFrQ1jbvzzju57rrrOn0dHb2XZrbSOTe7m2/RkNHR9bE/fEaHunvvvZfy8nK+973vBX3MQPp3KzpQwzt+kHtn234+3FcDeJN/zx43jLkT0pg7PpWZo5OJ1LRTIj3maNfIIR3uAE75v2WcNC6VX1x5fI+eV0QGt4H0Bay/U7jrPoW7/ufiiy9my5YtvPLKK12aL7K//rs55/hwXw3vbtvP235Ty51ltQAkx0Zy0rhUTp6QytzxaUwflaSpB0R60dGukUO+TjwnK5lVRWWhLoaIiIgcwdy5c6mvrz9k3ZNPPsmsWbNCVKLOPfvss6EuQrc459hSWn1In7ldFd6AM2nxUcwZn8rNHxnP3AlpTBmRqL5yIv3EkA93udkpLF27m7KaBlLiokJdHBGRPjUQvzTL0PPOO++EugiDXkuLY9OeqrYw9862/eyt8v42ZCRGM3d8KnMnpHHy+FQmDk/ocIRQEQk9hbusFABWFZVzxuSM0BZGRKSP6UuzyNDU3OLYsKuirc/cu9v2c6CmEYDM5BhOn5jW1mdufHq8wpzIADHkw92srGQACgrLFO5EpEuONL+VBG8g9fseiPQZHVh68/9DU3MLa4sr2mrmlm/fT4U/UXh2aizzpo1gzvhUTh6fRnZqrD43IgPUkA93STGRHJcRT4H63YlIF8TExLBv3z7S0tL0JegYOefYt29f22TQ0rP0GR1Yevr/Q2NzC6uKytvC3MoPD1BV74W58enxnDsrk7n+ACijUmJ75DlFJPSGfLgDr2nma5v26hdOEQlaVlYWRUVFlJaWhrooA1pMTAxZWVmhLsagpM/owNPd/w81DU28tG43/ygo4Y3Ne6lt9CZanzg8gQV5o9qaWY5I0g8qIoOVwh3eoCp/fX8nJeV1+vVKRIISGRnJ+PGaH1P6L31Gh4b6pmb+vbGUxQXFLFu/h9rGZkYmxXDZ7CxOnpDGnPGppCdEh7qYItJHFO7wpkMAWFVUpnAnIiIi/VpTcwtvbd3H4vxi/rl2F5V1TQyLi+SSE0ZzYe4oThqXqqkJRIYohTtgWmYSkeFGfmE558zMDHVxRERkEDCzc4D7gHDgYefcPe22DwMeBY4D6oDPOufW+NtuB24GDHjIOffzPiy69EMtLY73dhxgcUExS1aXsLeqgYToCM6aMYILc0dx2sR0IsPDQl1MEWlpgeo9UF4EZTu8+/IiKC/0bmd8GaYv6LWnDyrcBXGBmgo8BpwAfN05d6+/fgrwx4BdJwDfcs793My+g3fhau0M8DXn3JJuvJZjFhMZzrTMJAoKy0Lx9CIiMsiYWThwP/BJoAhYbmaLnXPrAnb7GpDvnLvYv47eD8wzs5l418c5QAPwTzN73jm3qW9fhYSac461xRX8o6CY51aVsLOsluiIMOZNG86FuaP46JThxESGh7qYIkNLY92hYa01vLUGuYqd0Nxw6DHRyZCc5d0i43u1eJ2GuyAvUPuBLwAXBR7rnNsI5AWcZyfwbMAuP2sNgqGWk5XM394vpqXFqSmDiIh01xxgs3NuK4CZPQ0sAAKvndOB/wNwzm0ws3FmNgKYBrztnKvxj/03cDHwoz4sv4TQltIqFucX849VxWwtrSYizPjIpHTuPHsyn5g2gsSYyFAXUWRwcg5q9kO5H9TKCg8PctXtBqmyMEjM9ILb6BO8WrnkLEjOhpRs73FMcp+9hGBq7jq9QDnn9gB7zOy8o5xnHrDFOfdhN8rba3KzUvjd2zvYureKicMTQ10cEREZ2EYDhQHLRcDcdvsUAJcAr5vZHGAskAWsAX5gZmlALXAusKLXSywhtbOsln8UFPOPgmLWFldgBnPHp3LT6RM4Z+ZIUuOjQl1EkYGvqQEqiwNCW9HhQa6p9tBjIuMOhrWROX5gyz64LmkUhPefH1yCCXfBXKCCcQXwh3brbjWza/EuWl9yzh04hvP2iLzsFAAKCssV7kREpLs6agLSfobqe4D7zCwfWA28DzQ559ab2f8DXgKq8EJg02FPYLYQWAgwZsyYniu59JnSynqWrC5hcUExKz/0vgLlZqfwzfOnc96sTEYma8qCXtdYB1W7oGoPRCdC2iQI15AU/V5LMzTVef9+Te1v9VBXfrDGLbD2rXIXh/0pjh/uBbUR02Hy2QHBLQtSxkDsMBhAU6UF8+kN5gJ19BOYRQEXAl8NWP0A8D3/XN8DfgJ8toNj++TiNSEjgfiocAqKyvjUiZpzSUREuqUIyA5YzgKKA3dwzlUANwCYN8nqNv+Gc+4R4BF/2w/989Hu+EXAIoDZs2d36bosoVNe28jSNbtYXFDMm1v20uJgyohEvnz2FC7IGcWYtLhQF3Hgc877cl+1xwtulbv9+11Qtdu/97fVlR96bEQsjJzp1dBk5nq34dMgQtNJHKalxavlaqr3g1bA46b6Q7c11bfbHhDEDglpQR7XctjvXR0LjzoY1I6b5we2wFq30RA5uH5ECSbcdXqBCsJ84D3n3O7WFYGPzewh4LmODuyri1d4mDErK1mDqoiISE9YDkwys/F4/c2vAK4K3MHMUoAa51wDcBPwmh/4MLPhzrk9ZjYGr+nmKX1ZeOlZNQ1NvLx+D/8oKObfG0tpaG5hTGocn/voRC7IHcWUkT3cYqiuwqupiIyFqHivWVlkHIQN8NE0W1qgZm+7kLb70Metwa190zqAiBhIGAGJIyFjCkw4ExKGQ8JIb33tASgp8G6r/wwrHvGOC4uE4VO9oDfSD3wjZ3rv7VDQVA/7tkDpBtj7gXdf+gHs23T4wCFdER7l/ZtERHuhOiL64HJkrFdjFhETcPO3R7ZbDtwe6Z8nKtELcPEZA/9z30XBhLtOL1BBuJJ2TTLNLNM5V+IvXozXxyCkcrNSeOyN7dQ3NRMdodGnRETk2DjnmszsVmAp3kjTjzrn1prZLf72B/EGTnnCzJrx+rHfGHCKv/h97hqBz4ey24Icm/qmZl77YC//KCjmpXW7qW1sZkRSNNecMpYLckeRm5WM9WRTL+egaAWsfBzW/KXjcNMa8qLivBH7olqX4ztY39n2OIhKOPg4rBvfm5rq/Zq01nAWWNu2+2CAq9oDrvnw42OS/YA2HLLnHAxwCSMhccTBbTHJnTevy/20d9/SAmXbD4a9klWw8QV4/3f+jgbpkw7W7mXmwshZXiAZqOqrvPAWGOBKN8CB7QHvu8GwsZA+BSZ+3AtPnQWtIwW47nxm5Ig6DXfBXKDMbCRev7kkoMXM7gCmO+cqzCwOb6TN/2p36h+ZWR5es8ztHWzvc7nZKTQ0t7ChpJJcvw+eiIjIsfCn91nSbt2DAY/fAiYd4diP9G7ppDc0tzje2rKPxQU7+eeaXVT4k4tffMJoLsgZxZzxqYT39IjctWVeLdPKx2H3Gi9w5X4axp/hhaaGamisgYYaaKz2ltse13jbKorb7VMDLY1dK0dETCdB0F8fEQ3Vew8NcLUd/XZhXnBoDWcjZ/phbeTB2rbEEV6Qi4ztgTeynbAwSJ3g3WZc7K1zznuvSgpg1yrv/sM3vfe/VcpYP+zlQGae9zhheM+Xrztq9kPpRti78WCA2/uB1yetVVgEpB4HI2bAzEsgYyqkT/YCbW+839JjguoxGsQFahdec82Ojq0B0jpYf02XStoHWgPdqqIyhTsRERHplHP+5OL5xTy/ehd7q+qJjwrn7BkjuSB3FKdP6oXJxZ2DouV+Ld1fvVq6UcfDBffBzE95A4N0V3NjB8GwXShsqAp43NG+NV5tW+D6pgaIS/OCWdpxMPZUP7CNOPQ+Lr3/DWxiBsmjvdvUcw+ur94bUMPnB7/1iw9uTxh5aA1fZo7X36s3B+lwzqsF3bvRC3KlGw/WyAUO5R8R6wW2MSdDxnVejVzGVEgd369GgJTg9bP/NaE1KjmG9IQo8gvLuUa9G0REROQIymsaWfSfLfzt/WJ2ltUSFRHGvKnDuSB3FB+f2kuTi9cegFV/8kLdnnVev6K8K+GE62BUXs8+V3gkxKZ4Nzm6+HSYOM+7taorh12rveacraFv80vgWrztscP8ppytA7fkebWEXe0f1tLiDeXfFuBaw9wHUB8wWEx0MmRM9kaDbA1wGZMhecyQ65M22CncBTAzcrNSKCgqC3VRREREpB9qbnH8cXkhP166gfLaRs6YnMH/fHIyZ83opcnFnYPCd7xAt/ZZb6TAUSfABb/wa+kSev45pftikmHc6d6tVUONF8pL8g/243vnwYODkkQleP32AkNfxhQvaDc3wv6th4e4vZsO7V8ZP9w7ZtalBwNcxlSvRnQADecvx07hrp2crBRe2biHyrrG3vkjLSIiIgPSiu37+fbitawtrmDO+FS+c8EMpo9K6p0nqz0ABX/0Ql3per+W7jNw4nXel34ZeKLiIGu2d2vV1OA1lQzsx/feE15TVoDwaG+S7PLCQ4f/T872+sCN+8jBAJc+GeJS+/Y1Sb+jcNdObnYyzsHqneWcelx6qIsjIiIiIba7oo57XtjAs+/vZGRSDL+48nguyMns2dEuwaul2/G2F+jW/c2rpRt9Ilz4S5hxiWrpBqOIKH/wlZyD61qavakHSgq8Wr7yIphxUcCgJpP1WZAjUrhrJzcrBYCCQoU7ERGRoay+qZlHX9/OL1/ZRFOz49aPTeRzHzuOuKge/vpUsx8KnvZC3d6NEJ0Ex1/t9aUL/NIvQ0NYuF8bNxlyLgt1aWSAUbhrZ1h8FGNS41ilfnciIiJD1isbdvPdf6xj+74aPjl9BN84bxpj03pw0mrnYMdbsOIxWPd3aK6H0bNhwf3e0PtDZYJsEelRCncdyM1OYeX2/aEuhoiIiPSxbXur+d5z63hlwx4mZMTz+A0n8dEpPThPWc1+KPiDX0v3gVdLd8K1Xl+6kbN67nlEZEhSuOtAblYy/ygoZk9lHcMTY0JdHBEREell1fVN/PKVzTzy+laiI8L5+rnTuO7UcURF9MAw8c7Bh2/4fen+7o2OmHUSLPi115dKtXQi0kMU7jrQNpl5YTmfmK5wJyIiMlg55/h7fjH/98J6dlfUc+mJWfzvOVN65sfd6n1Q8HtY+VvYt8mba+zE672+dCNndv/8IiLtKNx1YMaoJMLDjIKiMj4xfUSoiyMiIiK9YM3Ocr69eC0rPzxATlYyD1x9IieMGda9kzoH21/3aunWL/Zq6bLnwkcegOkXecPhi4j0EoW7DsRFRTBpeAIFReWhLoqIiIj0sP3VDfx46UaeXr6D1LgofvSpHC49MYuwsG5MbVC9F/J/D+/9FvZt9iaxnv1Zr5ZuxPSeK7yIyFEo3B1BXnYK/1y7C+dcz89jIyIiIn2uqbmFp97ZwU9e3Eh1QzOfPW08X5g3ieTYyGM7oXOw7TW/lu4f0NII2SfDR+6E6QtUSycifU7h7ghyslJ4enkhO/bX9OzQxyIiItLn3tyyl7sXr2Pj7kpOm5jGdy6YwaQRicGfoKEa9m6C0o3eXHSlG6FkFZTv8GrpTrrJG/Fy+LTeexEiIp1QuDuC3OxkAPILyxTuREREBqidZbX88Pn1PL+6hKxhsTx49YmcPWPEkVvl1B6A0g8OBrjWW/mOg/tYOKQdB6Ny4eNf92rpImP75gWJiByFwt0RTB6RSExkGAWF5SzIGx3q4oiIiEgX1DU2s+i1rfz61c0A/M8nJ7PwjAnERIZ7zSkrd0PpBm+uudKNBx9X7T54kogYSJsE2XPghGsgfTJkTIXUCRARFaJXJiJyZAp3RxAZHsaMUcmsKioLdVFEREQkSM45lq7dzfefX8fOA9VcPTWc23ObSa9dCktaw9wGqAsYNC0qETKmwMRPePfpUyBjMqSMhbDw0L0YEZEuUrg7itysFH7/7oc0NbcQEd4Dk5iKiIhIz2tuggPbKNn8Pv95800iDmziscgSJsQXE769Frb7+8Wle+Ft5qcOBriMqZCYCRo8TUQGAYW7o8jNTubRN1r4YHcV00clhbo4IiIiQ1tjnTcZeGs/uL0bofQD3L7NWEsjmcDlQHXcCGJHzyAs4ywvzLXWxsWnhfoViIj0KoW7o8jNSgGgoKhM4U5ERCQUtvwL3vmN15Sy7ENwLd56C8MNG0dxxBheclNY3TiSCdNO4Mr580hNSw9tmUVEQkTh7ijGpsWRHBtJQWEZV84ZE+riiIiIDC3vPgQv/C8kjoKs2ZBzeVstXH5NGt9esoWCHeWcOHYYd184g5mjk0NdYhGRkFK4OwozIycrmYKi8s53FhERkZ7R0gxLvw7vPACTzoZLH4Fob066PZV1/OifG3lm5XsMT4zm55/OY0HeqCNPbSAiMoQo3HUiLzuFX7+6hdqGZmKjNGKWiIhIr6qvgr/cCB/8E07+HJz1fQgLp6Gphd++uZ37lm2ivqmZW848jls/PpGEaH2VERFpFdRfRDM7B7gPCAceds7d0277VOAx4ATg6865ewO2bQcqgWagyTk321+fCvwRGIc3jtXlzrkD3Xs5PS83K4XmFsfa4nJmj0sNdXFEREQGr/Kd8IdPw+61cO69MOdmAF77oJS7/7GWLaXVfHzqcL55/nTGp8eHuLAiIv1Pp+P7m1k4cD8wH5gOXGlm09vtth/4AnAvHfuYcy6vNdj57gKWOecmAcv85X4nJ9trv59fWBbagoiIiAxmxe/DQx+H/dvhqj+3Bbs/vLuDax99l+YWx6PXz+bR609SsBMROYJgJm+bA2x2zm11zjUATwMLAndwzu1xzi0HGrvw3AuA3/qPfwtc1IVj+8zwxBhGJceo352IiEhvWf8cPHYuhEfCjUth0ifaNr24dhcT0uNZ+sUz+PjUESEspIhI/xdMuBsNFAYsF/nrguWAF81spZktDFg/wjlXAuDfD+/COftUTlYKq4rKQl0MERGRwcU5ePOX8MerYfg0uGkZjJgRsNmxqsgbDTM6Qv3eRUQ6E0y462j4KdeF5zjNOXcCXrPOz5vZGV04FjNbaGYrzGxFaWlpVw7tMbnZKXy4r4YD1Q0heX4REZFBp7kRnrsDXvwGTL8QrnsOEg+tmSs6UMu+6gZys1NCUkQRkYEmmHBXBGQHLGcBxcE+gXOu2L/fAzyL18wTYLeZZQL493uOcPwi59xs59zsjIyMYJ+2R+X6/e5W7VTTTBERkW6rLYOnLoWVj8Pp/wOXPg5RcYftVuC3mslTuBMRCUow4W45MMnMxptZFHAFsDiYk5tZvJkltj4GzgLW+JsXA9f5j68D/t6VgvelWaOTMYMCDaoiIiLSPQe2wyNnwfbXYcH98IlvQ1jHX0cKCsuIighjysjEvi2jiMgA1elUCM65JjO7FViKNxXCo865tWZ2i7/9QTMbCawAkoAWM7sDb2TNdOBZf2LRCOD3zrl/+qe+B/iTmd0I7AAu69FX1oMSYyI5LiNB4U5ERKQ7Ct+FP1wJLU1wzbMw/ug9NQoKy5k5KonI8GB+ixYRkaDmuXPOLQGWtFv3YMDjXXjNNdurAHKPcM59wLygSxpiOVnJvPbBXpxz+GFVREREgrX6Gfjb5yBpFHzmz5A+6ai7NzW3sHpnOZ8+Kfuo+4mIyEH6KSxIedkp7K2qp7i8LtRFERERGTicg3//GP5yI4w+0RsRs5NgB7C5tIraxmb1txMR6QKFuyDlZqUA6ncnIiLBMbNzzGyjmW02s7s62D7MzJ41s1Vm9q6ZzQzY9kUzW2tma8zsD2YW07el7yFN9fDsLfCv70POp+Hav0F8WlCHtl5vNVKmiEjwFO6CNDUzkchwaxu5S0RE5EjMLBy4H28aoOnAlWY2vd1uXwPynXM5wLXAff6xo4EvALOdczPx+rtf0Vdl7zE1++GJi2DV0/Cxr8PFv4GI6KAPzy8sJykmgnFph4+iKSIiHVO4C1J0RDjTM5NUcyciIsGYA2x2zm11zjUATwML2u0zHVgG4JzbAIwzs9aJ3iKAWDOLAOLowhRE/cLeTfDwPNi5Ej71CJz5v9DF/uoFhWXkZqeon7uISBco3HVBTlYKa3ZW0NzSlTncRURkCBoNFAYsF/nrAhUAlwCY2RxgLJDlnNsJ3Is3knQJUO6ce7HXS9xTtv0HHv4E1FXAdf+AWZd2+RS1Dc1s3F2p/nYiIl2kcNcFudkpVNU3sbW0KtRFERGR/q2j6qb2vwzeAwwzs3zgNuB9oMnMhuHV8o0HRgHxZnb1YU9gttDMVpjZitLS0h4t/DF7/3fw5MWQMAJuehnGzD2m06wtLqe5xbX1dxcRkeAo3HVBXnYyAPlqmikiIkdXBASO4Z9Fu6aVzrkK59wNzrk8vD53GcA24BPANudcqXOuEfgrcGr7J3DOLXLOzXbOzc7IyOillxGklhZ4+W74++dh3Glw44uQOv6YT9d6nc3xr7siIhIchbsumJCeQEJ0BKuKykNdFBER6d+WA5PMbLyZReENiLI4cAczS/G3AdwEvOacq8BrjnmymcWZ1+FsHrC+D8veNY218MwN8PpP4YTr4DPPQGxKt065qqicUckxDE8cmIOEioiESlCTmIsnLMyYNTpZI2aKiMhROeeazOxWYCneaJePOufWmtkt/vYHgWnAE2bWDKwDbvS3vWNmzwDvAU14zTUXheBldK5qD/zhSm/glLO+D6fc2uWBUzpSUFSmKRBERI6Bwl0X5WQn8+jr26hvaiY6IjzUxRERkX7KObcEWNJu3YMBj98COpzN2zn3beDbvVrA7tq9Dn7/aajZC5/+HUw7v0dOe6C6gQ/31XDlnDE9cj4RkaFEzTK7KC8rhcZmx/qSylAXRUREJDQ2vwyPnAXNDXDDkh4LdkBb6xgNpiIi0nUKd13U2kxE892JiMiQtPxheOpyGDYObn4FRh3fo6cvKCzHDGZlaTAVEZGuUrPMLspMjiE9IVr97kREZGhpaYYXvwFv/xomnQ2XPgLRiT3+NAVFZUwa7g1gJiIiXaO/nF1kZuRlJ6vmTkREho76KvjLTfDBCzD3v+HsH0BYz/c7d85RUFjGx6YO7/Fzi4gMBWqWeQxys1LYureairrGUBdFRESkd5XvhMfOgU1L4dx7Yf49vRLsAHaW1bKvukEjZYqIHCOFu2OQk52Cc7BG892JiMhgVpwPD8+D/dvhqj/DnJt79ekKCr3rap4GUxEROSYKd8cg1+/kna9+dyIiMlhteB4emw9hEXDjUpj0iV5/yoKiMqIiwpgysuf78omIDAUKd8cgJS6KsWlxrCpUzZ2IiAwyzsGbv4KnPwMZU+GmZTBiRp88dX5hGTNGJREVoa8nIiLHQn89j1FuVopGzBQRkcGluRGe+yK8+HWYfiFc/zwkjuiTp25qbmF1UbnmtxMR6QaFu2OUm51CSXkdeyrqQl0UERGR7qsrh6cug5WPwelfhEsfh6i4Pnv6zaVV1DY2k5ut+e1ERI6Vwt0xau13V6BBVUREZKCr2Q+PnAXb/wML7odPfAfC+vYrQusUQ6q5ExE5dgp3x2jGqGTCw0zz3YmIyMAXOwzGngbXPAvHXx2SIhQUlZMUE8G4tPiQPL+IyGCgScyPUWxUOJNHJKrfnYiIDHxmcP5PQ1qEgsIycrNTCAuzkJZDRGQgC6rmzszOMbONZrbZzO7qYPtUM3vLzOrN7M6A9dlm9i8zW29ma83s9oBt3zGznWaW79/O7ZmX1HfyspMpKCzDORfqooiIiAxYdY3NbNhVqSaZIiLd1Gm4M7Nw4H5gPjAduNLMprfbbT/wBeDeduubgC8556YBJwOfb3fsz5xzef5tybG+iFDJzUqhoq6J7ftqQl0UERGRAWttcTnNLY7c7JRQF0VEZEALpuZuDrDZObfVOdcAPA0sCNzBObfHObccaGy3vsQ5957/uBJYD4zukZL3Azn+L4yr1DRTRETkmOX788a2DlYmIiLHJphwNxooDFgu4hgCmpmNA44H3glYfauZrTKzR81sWFfPGWqTRyQQExlGvgZVEREROWYFhWVkJscwPCkm1EURERnQggl3HfVs7lInMzNLAP4C3OGcq/BXPwAcB+QBJcBPjnDsQjNbYWYrSktLu/K0vS4iPIyZo5JZpekQREREjtmqojL1txMR6QHBhLsiIDtgOQsoDvYJzCwSL9g95Zz7a+t659xu51yzc64FeAiv+edhnHOLnHOznXOzMzIygn3aPpObncKaneU0NreEuigiIiIDTllNA9v31ai/nYhIDwgm3C0HJpnZeDOLAq4AFgdzcjMz4BFgvXPup+22ZQYsXgysCa7I/Utudgr1TS1s3FUZ6qKIiIgMOAV+65fcbPW3ExHprk7nuXPONZnZrcBSIBx41Dm31sxu8bc/aGYjgRVAEtBiZnfgjayZA1wDrDazfP+UX/NHxvyRmeXhNfHcDvxXD76uPtPa+XtVUTkzR+vCJCIi0hUFhWWYwSxdQ0VEui2oScz9MLak3boHAx7vwmuu2d7rdNxnD+fcNcEXs/8akxpHSlwkBYVlXDV3TKiLIyIiMqAUFJYxMSOBxJjIUBdFRGTAC2oSczkyMyMnK4UCTYcgIiLSJc45CorK1N9ORKSHKNz1gLysZD7YXUlNQ1OoiyIiIjJg7CyrZW9Vg+a3ExHpIQp3PSA3O4UWB2t2VnS+s4iIiAC0TSWkmjsRkZ6hcNcDcvy5eVapaaaIiEjQCgrLiAoPY+rIpFAXRURkUFC46wEZidGMToklv7As1EUREREZMPILy5g+KomoCH0dERHpCfpr2kNyspLbmpeIiIjI0TW3OFbvLCdPTTJFRHqMwl0Pyc1OYcf+GvZXN4S6KCIiIv3e5j1V1DQ0a/JyEZEepHDXQ3L9fneaEkFERKRzBX5Xhtbrp4iIdJ/CXQ+ZlZWMGawqVNNMERGRzuQXlZEYE8G4tPhQF0VEZNBQuOshCdERTMxIUM2diIhIEFYVlZGblUJYmIW6KCIig4bCXQ/KzU5hVVEZzrlQF0VERELMzM4xs41mttnM7upg+zAze9bMVpnZu2Y2018/xczyA24VZnZHn7+AXlTX2MyGkkr1txMR6WEKdz0oNyuZvVUN7CyrDXVRREQkhMwsHLgfmA9MB640s+ntdvsakO+cywGuBe4DcM5tdM7lOefygBOBGuDZvip7X1hbXEFTi1N/OxGRHqZw14Ny/eGcC9TvTkRkqJsDbHbObXXONQBPAwva7TMdWAbgnNsAjDOzEe32mQdscc592NsF7kutg6loGgQRkZ6lcNeDpo5MIio8jFXqdyciMtSNBgoDlov8dYEKgEsAzGwOMBbIarfPFcAfeqmMIVNQVEZmcgzDk2JCXRQRkUFF4a4HRUWEMW1UEvn+L5IiIjJkdTRKSPsO2fcAw8wsH7gNeB9oajuBWRRwIfDnDp/AbKGZrTCzFaWlpT1S6L5SUFimJpkiIr1A4a6H5WUls2ZnOc0tGlRFRGQIKwKyA5azgOLAHZxzFc65G/y+ddcCGcC2gF3mA+8553Z39ATOuUXOudnOudkZGRk9WvjeVFbTwPZ9NeRoMBURkR6ncNfDcrJSqG5oZktpVaiLIiIiobMcmGRm4/0auCuAxYE7mFmKvw3gJuA151xFwC5XMgibZK4q8vql56nmTkSkxync9bDWQVXUNFNEZOhyzjUBtwJLgfXAn5xza83sFjO7xd9tGrDWzDbg1dLd3nq8mcUBnwT+2rcl730FhWWYwcws1dyJiPS0iFAXYLCZkB5PYnQEq4rKuHx2ducHiIjIoOScWwIsabfuwYDHbwGTjnBsDZDWqwUMkYKiMo7LSCApJjLURRERGXRUc9fDwsKMWVnJmg5BRESkHecc+YXlGkxFRKSXKNz1gtzsFNaXVFDX2BzqooiIiPQbxeV17K2qJ0+DqYiI9AqFu16Qm5VMU4tjfUlF5zuLiIgMEa2Tl+dq8nIRkV6hcNcLWi9aBRpURUREpE1BURlR4WFMHZkU6qKIiAxKQYU7MzvHzDaa2WYzu6uD7VPN7C0zqzezO4M51sxSzewlM9vk3w/r/svpH0YmxZCRGN023LOIiIh4P3pOG5VEVIR+WxYR6Q2d/nU1s3DgfrxhmqcDV5rZ9Ha77Qe+ANzbhWPvApY55yYBy/zlQcHMyM1KIb+oLNRFERER6ReaWxyri8rJ0xQIIiK9JpifzuYAm51zW51zDcDTwILAHZxze5xzy4HGLhy7APit//i3wEXH9hL6p7zsZLaWVlNe2/4tERERGXq2lFZR3dCs/nYiIr0omHA3GigMWC7y1wXjaMeOcM6VAPj3wzs6gZktNLMVZraitLQ0yKcNvRx/mOc1O9U0U0REJF+DqYiI9Lpgwp11sM4Fef7uHOvt7Nwi59xs59zsjIyMrhwaUjl+s5N8DaoiIiJCQWEZiTERjE+LD3VRREQGrWDCXRGQHbCcBRQHef6jHbvbzDIB/Ps9QZ5zQEiJi2JcWhyr1O9ORESEgqIycrKSCQvr6HdfERHpCcGEu+XAJDMbb2ZRwBXA4iDPf7RjFwPX+Y+vA/4efLEHhtzsFAoK1SxTRESGtrrGZjaUVJLrd1kQEZHe0Wm4c841AbcCS4H1wJ+cc2vN7BYzuwXAzEaaWRHwP8A3zKzIzJKOdKx/6nuAT5rZJuCT/vKgkpuVwq6KOnZX1IW6KCIiIiGzrqSCphan/nYiIr0sIpidnHNLgCXt1j0Y8HgXXpPLoI711+8D5nWlsANNbrbX766gsIyzZowMcWlERERCo8Dvf56ncCci0qs0i2gvmjEqmfAwo0D97kREZAgrKCxjZFIMI5JiQl0UEZFBTeGuF8VEhjN1ZCKritTvTkREhq6CovK21iwiItJ7FO56WU5WCgWFZbS0dGkGCBERkUGhvKaRbXur1d9ORKQPKNz1srzsZCrqmti+rzrURREREelzq3aWAWikTBGRPqBw18ty/IuZmmaKiMhQ1DqYyqwsNcsUEeltCne9bNLwBGIjw8n3L24iIiJDSX5hOcdlxJMUExnqooiIDHoKd70sIjyMWaOTWaURM0VEZIhxzpFfWKb+diIifUThrg/kZCWzpriCxuaWUBdFRESkz5SU17G3ql7z24mI9BGFuz6Qm51CQ1MLG3dVhrooIiIifaa1v50GUxER6RsKd32g9aKmycxFRGQoyS8qIyo8jKmZiaEuiojIkKBw1weyU2MZFhfZ9gumiIjIULCqsJxpmYlER4SHuigiIkOCwl0fMDNys1M0HYKIiAwZzS2O1TvLNZiKiEgfUrjrIzlZKXywu5Lq+qZQF0VERKTXbS2toqq+Sf3tRET6kMJdH8nLTqbFwZqdqr0TEZHBr3V+V9XciYj0HYW7PpLj/3KpppkiIjIUFBSVkRgdwYT0+FAXRURkyFC46yPpCdGMToklXyNmiojIEFBQWE5OdjJhYRbqooiIDBkKd30oLztFI2aKiMigV9fYzPqSCvW3ExHpYwp3fSgnK5miA7Xsq6oPdVFERER6zfqSCppaXFuXBBER6RsKd32otVO5+t2JiMhg1tpKJU+DqYiI9CmFuz40c3QyZl4ncxERkcGqoKicEUnRjEyOCXVRRESGFIW7PpQQHcGk4QnqdyciMgSY2TlmttHMNpvZXR1sH2Zmz5rZKjN718xmBmxLMbNnzGyDma03s1P6tvTdU1BYpv52IiIhoHDXx3KzUigoKsc5F+qiiIhILzGzcOB+YD4wHbjSzKa32+1rQL5zLge4FrgvYNt9wD+dc1OBXGB975e6Z5TXNLJ1b7XmtxMRCYGgwl0Qvz6amf3C377KzE7w108xs/yAW4WZ3eFv+46Z7QzYdm6PvrJ+Kic7hf3VDRQdqA11UUREpPfMATY757Y65xqAp4EF7faZDiwDcM5tAMaZ2QgzSwLOAB7xtzU458r6rOTdtGpnGaD+diIiodBpuAvy18f5wCT/thB4AMA5t9E5l+ecywNOBGqAZwOO+1nrdufcku6+mE61NEPlrl5/mqPJ85upqN+diMigNhooDFgu8tcFKgAuATCzOcBYIAuYAJQCj5nZ+2b2sJkdNhO4mS00sxVmtqK0tLQ3XsMxae16MCsrObQFEREZgoKpuQvm18cFwBPO8zaQYmaZ7faZB2xxzn3Y7VIfq1fvgQdPh8J3Q1aEKSMTiQoP04iZIiKDW0czd7dvj38PMMzM8oHbgPeBJiACOAF4wDl3PFANHNZqxjm3yDk32zk3OyMjoyfL3i0FReVMyIgnKSYy1EURERlyggl3wfz6GMw+VwB/aLfuVr8Z56NmNiyIsnTPrMsgKgEePx9W/anXn64jURFhTB+VRL4GVRERGcyKgOyA5SygOHAH51yFc+4Gv3XLtUAGsM0/tsg5946/6zN4Ya/fc86RX1jW1kpFRET6VjDhLphfH4+6j5lFARcCfw7Y/gBwHJAHlAA/6fDJe7LZScZkuPkVyDoJ/nozLPsetLR075zHIC87hTU7y2lu0aAqIiKD1HJgkpmN96+BVwCLA3fwR8SM8hdvAl7zA98uoNDMpvjb5gHr+qrg3bGroo7SynoNpiIiEiLBhLtOf30MYp/5wHvOud2tK5xzu51zzc65FuAhvOafh+nxZidxqXDNs3D8NfCfe+GZ66Ghpvvn7YKcrGRqGprZvKeqT59XRET6hnOuCbgVWIo30uWfnHNrzewWM7vF320asNbMNuBdJ28POMVtwFNmtgrvR9Af9lnhu6G1v53CnYhIaEQEsU/br4/ATrxfH69qt89ivCaWTwNzgXLnXEnA9itp1yTTzDID9rkYWHMM5T82EVFw4S8hYyq8+A048CFc+QdIGtUnT9960SsoLGPKyMQ+eU4REelb/kBhS9qtezDg8Vt4A5F1dGw+MLs3y9cb8gvLiQw3pmXq2iYiEgqd1twF+evjEmArsBmvFu5zrcebWRzwSeCv7U79IzNb7f8q+THgi919MV1iBqfeClc+Dfs2w0Mfh+L3++Spx6fFkxgToREzRURkUCkoLGN6ZhLREeGhLoqIyJAUTM1dML8+OuDzRzi2BkjrYP01XSppb5lyDnx2KfzhCnh0Plz8IMy4qFefMizMyMlKVrgTEZFBo6XFsXpnORcf3348NRER6StBTWI+6I2c6Q20MnIW/Pk6eO3H4Hp3sJPcrBQ2lFRS19jcq88jIiLSF7buraKqvkn97UREQkjhrlXCcLjuHzDrcnjl+/DXhdBY12tPl5OVQlOLY11JRa89h4iISF/JL/Tmb83L1uTlIiKhonAXKDIGLlkEH/8mrP4T/PYCqNrTK0+VFzCoioiIyEBXUFhGQnQEE9ITQl0UEZEhS+GuPTM44064/AnYtdobaGVXzw/kOTI5hhFJ0awqKu/xc4uIiPS1gqIycrKSCQvraOpbERHpCwp3RzJ9AXz2BWhpgkfPho0v9PhT5GSlqOZOREQGvLrGZtaXVKi/nYhIiCncHc2o472BVtImwh+uhDd/2aMDreRlp7B1bzXltY09dk4REZG+tr6kgsZmR25WSqiLIiIypCncdSZpFNzwAky/0JvwfPFt0NTQI6fOyfI6na9W00wRERnAWrsY5GowFRGRkFK4C0ZUHFz6OJzxZXj/SXjyYqjZ3+3T5oxOAdB8dyIiMqAVFJYxPDGakUkxoS6KiMiQpnAXrLAw+Pg34JKHoGi5N9BK6QfdOmVyXCQT0uPV705ERAa0/KIycrNTMNNgKiIioaRw11U5l8P1z0FDFTz8Cdi8rHuny0pWzZ2IiAxY5bWNbC2tbpviR0REQkfh7lhkz/EGWknOgqcug3cfOuZT5WansLuinl3lvTdhuoiISG9p7TeuwVREREJP4e5YpYyBG5fCpE/Ckjvh+TuhuanLp8nxL4aqvRMRkYGo9fo1K0uDqYiIhJrCXXdEJ8IVv4dTb4PlD8HvL4Pasi6dYsaoJCLCTP3uRERkQMovLGNCRjzJsZGhLoqIyJCncNddYeFw1vfhwl/CttfgkU/Cvi1BHx4TGc7UzETV3ImIyIC0qqhMTTJFRPoJhbuecsK1cO3foboUHp4H218P+tCcrBRWFZXT0tJzE6SLiIj0tl3ldeyuqCdXTTJFRPoFhbueNO50uGkZxGfAExfBe08GdVheVgqVdU1s21fdu+UTERHpQfl+l4JcjZQpItIvKNz1tLTj4MaXYPxHYPGt8OI3oKX5qIfkZHu/eD76+jbqGo++r4iISH9RUFRGZLgxLTMp1EUREREU7npHbApc9Wc46WZ485fw9GegvvKIu08ensgVJ2Xz1Ds7OOtnr/GvjXv6rqwiIiLHqKCwjGmZScREhoe6KCIigsJd7wmPgPPuhXPvhU0vwqPnQNmODncNCzPu+VQOv795LpHhxg2PLee/f7eSkvLaPi60iIhIcFpaHKuLyjWYiohIP6Jw19vm3Ayf+TOUFcJDH4fCd4+466nHpfPC7Wfw5bOn8MqGPcz7yb956LWtNDa39GGBRUREOrd1bzWV9U3qbyci0o8o3PWFifPgppchKgEePx9W/emIu0ZFhPH5j03k5f85k5MnpPGDJeu54Jevs2L7/j4ssIiIyNG1zs+qkTJFRPoPhbu+kjEZbn4Fsk6Cv94Mr3wfWo5cI5edGscj183mN9ecSEVtI5c++Bb/+0wB+6sb+rDQIiIiHSsoKiMhOoIJGQmhLoqIiPgU7vpSXCpc8ywcfw289mN45npoqDni7mbG2TNG8vKXzuS/zpzAX9/bycd/8ip/XL5Dc+KJiEhIFRSWMWt0MuFhFuqiiIiIL6hwZ2bnmNlGM9tsZnd1sN3M7Bf+9lVmdkLAtu1mttrM8s1sRcD6VDN7ycw2+ffDeuYl9XMRUXDhL+Gs78O6xfDYfKgoPuohcVERfHX+NJ7/wkeYPDyRr/xlNZf95i3Wl1T0UaFFREQOqm9qZl1JhfrbiYj0M52GOzMLB+4H5gPTgSvNbHq73eYDk/zbQuCBdts/5pzLc87NDlh3F7DMOTcJWOYvDw1mcOptcOXTsG+zN9BK8fudHjZlZCJ//K+TufeyXLbtreb8X77O959bR1V9Ux8UWkRExLO+pJLGZkdetvrbiYj0J8HU3M0BNjvntjrnGoCngQXt9lkAPOE8bwMpZpbZyXkXAL/1H/8WuCj4Yg8SU86Bzy6FsAh4dD786VpY9j0oeBqKVkJd+WGHmBmXnpjFK186k8tnZ/Pw69v4xE/+zZLVJTinppoiItL72gZTUc2diEi/EhHEPqOBwoDlImBuEPuMBkoAB7xoZg74jXNukb/PCOdcCYBzrsTMhnf05Ga2EK82kDFjxgRR3AFm5ExvoJWlX4Od78H658A1H9yeMALSJkH6RP9+EqRNJCVlLP93ySwum53F159dw+eeeo8zJ2fw3QUzGJsWH7rXIyIig15BURnDE6MZmRQT6qKIiEiAYMJdRz2l21cRHW2f05xzxX54e8nMNjjnXgu2gH4YXAQwe/bswVk1lTAcPvWw97ipAQ5sh32bYK9/27fJ659XGzAdQlgkpE7ghPRJPD9tIq9nDuM3a7dy6c8+5OqP5nHLRycQHREekpcjIiKDW0FhGTlZKZhpMBURkf4kmHBXBGQHLGcB7UcAOeI+zrnW+z1m9ixeM8/XgN1mlunX2mUCe47tJQwyEVHetAkZkw/fVrP/YNjbu8nrr7d3E2EfLOWMlkbOMCAC9v0nkQ/ezCJj/ExGjp/p1falT4Zh4yA8sq9fkYiIDCIVdY1sKa3m4uNHh7ooIiLSTjDhbjkwyczGAzuBK4Cr2u2zGLjVzJ7Ga7JZ7oe2eCDMOVfpPz4L+G7AMdcB9/j3f+/2qxns4lJhzFzvFqi5Cco+bAt+DVtW07xtDeGbX4Qtfz64X1iEF/AOa+Y5CeLTvYFeREREjmJ1kdcfXP3tRET6n07DnXOuycxuBZYC4cCjzrm1ZnaLv/1BYAlwLrAZqAFu8A8fATzrN9uIAH7vnPunv+0e4E9mdiOwA7isx17VUBMeAWnHeTfOIfNUGNbYzIP/3sKTr65iUvhuPjezmdNTDhC2f7NX47flFWiuP3iOmJSDQS8w+KVOgIjoUL0yEZEBy8zOAe7Du3Y+7Jy7p932YcCjwHFAHfBZ59waf9t2oBJoBprajTYdUvn+YCo5o1NCWg4RETlcMDV3OOeW4AW4wHUPBjx2wOc7OG4rkHuEc+4D5nWlsBK8mMhw7vjEZC7KG823Fq/l2pWlzBiVxPcv+gLHjxkGLc1QXgh7N8PeDw429dz6Lyj4/cETWRgMnwHTLoAZF0HGlJC9JhGRgSJgGqFP4nVdWG5mi51z6wJ2+xqQ75y72Mym+vsHXhc/5pzb22eFDlJBYRkT0uNJjlMzfxGR/iaocCcD17j0eH57w0ksWb2L7z63lkseeJMr54zhK2dPJXnYOK+Z5qRPHHpQfaXfn88Pfttfh1f/D179IWRMhekLYPpFMHyamnKKiHSsbRohAL/bwgIgMNxNB/4PwDm3wczGmdkI59zuPi9tFxQUlXHqcemhLoaIiHRA4W4IMDPOy8nkjMnp/OylTTz+5jaWrtnFV8+dxqdOGH34aGfRiTDqeO/WqqIENjwH6/4Or/0Y/v3/vKabMy7ywt6ImQp6IiIHBTONUAFwCfC6mc0BxuINSLabI08jFFK7yuvYXVFPbpYmLxcR6Y+CmcRcBonEmEi+dcF0/nHb6YxNi+POPxfw6UVv88Huys4PTsqEOTfD9c/BlzbCeT+FpFHwn5/Ag6fDL0+Al++G4nzQZOoiIsFMI3QPMMzM8oHbgPeBJn/bac65E4D5wOfN7IzDnsBsoZmtMLMVpaWlPVfyoygoKgMgR4OpiIj0Swp3Q9CMUck8c8up3HPJLD7YXcm59/2He17YQE1DU+cHgzcv30k3wnWL4c5NcMF9XvPON+6DRWfCL/LgpW/BzpUKeiIyVHU6jZBzrsI5d4NzLg+4FsgAtvnb2qYRAlqnEaLd8Yucc7Odc7MzMjJ65UW0V1BYRkSYMT0zqU+eT0REukbhbogKCzOumDOGZf9zJhcfP5oH/72FT/70NV5cu6trJ4pPhxOvh2uehS9vhgt/5TXXfOt+eOjj8PMcWPp1KFwOLS298lpERPqhtmmEzCwKbxqhxYE7mFmKvw3gJuA151yFmcWbWaK/T+s0Qmv6sOxHVFBUxrTMJGIiw0NdFBER6YD63A1xaQnR/PiyXC4/KZtvPLuGhU+u5BPThvPtC2aQnRrXtZPFpcIJ13i32gOw8QWvj967i+CtX0HSaH8wlgWQNQfC9NvCkOUcVBR7tbvF70HJKu+HgrGnwbjTvSk41IdTBrAgpxGaBjxhZs14A63c6B9+tGmEQqalxbGqsJwFx48KdVFEROQIzA2gZnOzZ892K1asCHUxBq3G5hYee2MbP395Ey3OcdvHJ3HzRyYQFdHNEFZXDhv/6QW9zS978+sljITpF3qjbo45GcL0K/CgVrPfC3E7/Vvxe1DlDwgYFgEZ06BqF1T7/YYSM/2gdxqMPd2bc1Fhb0gxs5X9aW63/q4vro+b91TxiZ/+mx9fmsNls7M7P0BERHrF0a6RqrmTNpHhYSw84zjOzxnF3f9Yy4+XbuSv7xXx9fOmccakDCLCjzHkxSRD7qe9W10FbHoR1v0N3nvCq9WLH35wHr0xp3qTssvA1VANJQV+kPNr5g5sP7g9fTJM+BiMPgFGnQAjZ0FkjFebt3cTfPg6bH/Dm4JjzTPeMfHD/aDn1+xlTFXYE+ljBf7k5XkaTEVEpN9SzZ0c0SsbdvOtv6+l6EAtqfFRnDNzJOfPymTuhDTCw3rgi3V9lR/0/u7dN9ZAXDpMO9+r0Rv3EQW9/q6pAfasPVgbt/M9KN0Azu9fmZTlhbjWIDcqzwv7wXAO9m/1Qt6Hftir2Olti0s7GPTGngbDp6uZ7yCjmruu6Yvr47f/voZnVhax6jtn98w1QEREjsnRrpEKd3JUdY3NvLqxlOdXl7Bs/W5qGppJT4hi/sxMzsvJ5KRxqT1zkW+o9ppsrvu714SzsRpiU2HqeV6N3vgzITyy+88jx66lBfZtCghyK2HXGq+ZLXj/XqNPPBjkRp/gjazaU5zzagA/fMOr2fvwdSjb4T/3MK/Wd9zpXg3fiJlq6jvAKdx1TV9cHxfc/waxkWE8vfCUXn0eERE5OoU76RG1Dc28unEPz60qYdmG3dQ1tjA8MZpzZ3lB78QxwwjriaDXWAubl/lB7wVoqISYFC/oTb8IJnwUIqI6OYl0i3NQXnhojVxxvvdvARAZ79XCBQa5lLF931SybMfBoLf9DTiwzVsfnQxjTzlYuzcyR7XAA4zCXdf09vWxvqmZWd9+kRtOH8dX50/rtecREZHOqc+d9IjYqHDmz8pk/qxMahqaeGXDHp5fVcIf3t3B429uZ0SSF/TOz8nk+OxuBL3IWK9p5rTzobEOtv7LC3rrn4P8p7wv7lPmw9RzvRqb/igswnsdkXHefUSsvxzbP2uUqve2C3LvHRzcJCwSRs6EnMsP1sylT+4fryNlDOSNgbwrveXynQebcH74BnzgDzAYlegN3DPuNK+5b2auaoJFumBDSSUNzS3kZaWEuigiInIUCndyTOKiIjg/ZxTn54yiqr6JZet38/yqEp56ZwePvbGdUckxbTV6edkp2LHW6ETGeEFuynxoqoet//aC3obnYNXTPfui+kp49MGgd0gAjDn4ODLOe+1ty4EBMa7d8f66iHb7Hym81Fd6tXCBQa61eSMGGVNg4icP9pUbMRMiovvq3eme5NFeCM253Fuu3BXQZ+8NePklb31kPIyZe7Bmb9QJqg0WOYqCojIAcjWYiohIv6ZmmdKjKusaedkPev/+oJTGZsfolFjOz/GC3qzRycce9AI1N3oBpbW/V3/T3Og1L22sgaa6g48bawNu/nJTXcC2Gq+28pD9a4Bj+H8aFnF4OGxpgn2bD54vZczBZpWtA55EJ/bgG9HPVO0J6LP3BuxZ562PiIXsOQcHaMma3TuB1jlobmj3714HTbUH/91bHzcFfFbaPkOB2+u8QB+dGHBLarfcbl1k7IAZZVTNMrumt6+P//OnfP6zaS/vfm1ez/wNFxGRY6ZmmdJnEmMiufj4LC4+Povy2kZeWreb51cV88jr2/jNa1sZkxrHeTmZnDcrkxmjko79S0J4JGSf1LOF76/aAkH7cBgYENstN7YLjK1BwjmYdenBQBefHupX17cShsOMi70bQPU+2PGmV7u3/Q341w8B59WuZs/xgt7Imf77Hxi4AgNZQIAPJrC1jiTaVeFRATW4MV6wa6rzamLrKqClsfNzWDhEJxwhBB4pHLZbH5Xg3feHZrnSZ1YVlZOb1Y1WGCIi0icU7qTXJMdGcumJWVx6YhZlNQ28uG43z60qYdFrW3ng1S2MS/OC3vk5o5g6MlFfGo7EzKtFiojuv30MB6r4NG+OxWkXeMu1B+DDt/zavf/Aaz86chiz8HbNaf3AFRkLUXHedA2tTWtb17feH/Y49gj7Bpy3szDVVO8FvfoK/77Sm27ksHWVh66r2Q8HPjy4vrE6uPcuMv7IQTB5NHz8G8H/O0i/VlHXyJbSKhbkjgp1UUREpBMKd9InUuKiuHx2NpfPzmZ/dQMvrt3Fc6tKeODVLdz/ry1MyIjn/FmZnJcziikjB3GzQOnfYod5A/VMPddbriv35tprC1wBIay/DcjS+gNAd2tjW5qhoarjINhhQAzYt3qvd5+QoXA3iKwpKsc59bcTERkIFO6kz6XGR3HFnDFcMWcM+6rq+efaXTy/qoRf/Wszv3hlM5OGJ/g1eplMHK6gJyEUkwyjjg91KfpWWLj3uoOdbF4GvXx/MJWcLH0mRET6O4U7Cam0hGg+M3csn5k7ltLKev65poTnVpVw37JN/PzlTUwdmch5/qibEzISQl1cEZEhp6CwjPHp8aTEaURZEZH+TuFO+o2MxGiuOWUc15wyjj0VdSxZXcLzq0v4yUsf8JOXPmBaZpI36uasTMalx4e6uCIiQ0JBYTknT0gNdTFERCQICnfSLw1PiuH608Zz/WnjKSmv5YXVu3huVTE/XrqRHy/dyMzRSZw7K5NTJqQxY1QyURFhoS6yiMigs7uijl0VdepvJyIyQCjcSb+XmRzLZ08fz2dPH8/OslpeWO013fzRPzcCEBURxqzRyZwwJoUTxgzjhLHDGJEUE+JSi4gMfAWFZYAGUxERGSiCCndmdg5wHxAOPOycu6fddvO3nwvUANc7594zs2zgCWAk0AIscs7d5x/zHeBmoNQ/zdecc0u6/YpkUBudEstNH5nATR+ZwO6KOt778ADv7TjAezvK+O1bH/LQf7a17Xf8mBSOHzOME8akqHZPROQYFBSVERFmTM9MCnVRREQkCJ2GOzMLB+4HPgkUAcvNbLFzbl3AbvOBSf5tLvCAf98EfMkPeonASjN7KeDYnznn7u25lyNDyYikGObPymT+rEwA6puaWVdcwXs7yrzA9+EBnltVAqh2T0TkWBQUljM1M5GYSE1aLyIyEARTczcH2Oyc2wpgZk8DC4DAcLcAeMI554C3zSzFzDKdcyVACYBzrtLM1gOj2x0r0iOiI8I5fswwjh8zjBsZD8Cu8rq2oPfejgP89s3Da/daw970zCTV7omI+FpaHAVFZVyoyctFRAaMYMLdaKAwYLkIr1aus31G4wc7ADMbBxwPvBOw361mdi2wAq+G70DQJRcJwsjkGM6dlcm5AbV7a4sreO/DA7y/o4yVAbV70a21e2OHtdXwDVftnogMUdv2VVNZ16T+diIiA0gw4c46WOe6so+ZJQB/Ae5wzlX4qx8Avufv9z3gJ8BnD3tys4XAQoAxY8YEUVyRI4uOCPdq6sYMa1tXUl7Lex/6TTl3HODxN7az6LUWwKvdCwx700clERmu2j0RGfxaB1PJU7gTERkwggl3RUB2wHIWUBzsPmYWiRfsnnLO/bV1B+fc7tbHZvYQ8FxHT+6cWwQsApg9e3b7UCnSbZnJsZyXE8t5OQdr99bsrOB9P+wt37affxR4H/noiDByspI5wW/+ecLYFIYnqnZPRAafgsIy4qPCOS4jIdRFERGRIAUT7pYDk8xsPLATuAK4qt0+i/GaWD6N12Sz3DlX4o+i+Qiw3jn308ADAvrkAVwMrOnG6xDpMdER4Zw4dhgnjj1Yu1dcVuv33fNq+B59YxuNr20FIGtYrF8bmMIJY4cxLVO1eyIy8BUUlTMrK5nwsI4a54iISH/UabhzzjWZ2a3AUrypEB51zq01s1v87Q8CS/CmQdiMNxXCDf7hpwHXAKvNLN9f1zrlwY/MLA+vWeZ24L966DWJ9LhRKbGMSonl/BxvYIG6xmbWFpe3hb13tu1jsV+7FxMZxpSRSUwansDE4QlMzPDus1Pj9CVJRAaEhqYW1hVXcMNp40JdFBER6YKg5rnzw9iSduseDHjsgM93cNzrdNwfD+fcNV0qqUg/EhMZzoljUzlxbCoAzjmKyw/Ou7ehpJJ/f1DKMyuL2o6JighjQnq8F/j826ThiYxLjyM6QsOMi0j/sWFXBQ3NLRpMRURkgAkq3InI0ZkZo1NiGZ0SywUBw4aX1zSyubSSzXuq2m4FRWU8v7oE5/cgDQ8zxqTGcVxGApNGHKzpO254AgnR+i8qIn2vdTAVhTsRkYFF3xxFelFyXOQhNXytahua2VJaxZbSg6Fv054qXt24h6aWg+MGZSbHHFLT1xr80hKi+/qliMgQkl9YTnpCNKOSNWCUiMhAonAnEgKxUeHMHJ3MzNHJh6xvbG7hw301bN7jBb9NuyvZXFrF0+8WUtvY3LZfanwUEzO82j2vead3n5kcgzeOkYjIsSsoKiMvO1l/T0REBhiFO5F+JDI8rK2WLlBLi6O4vPaQ5p2b91TxwpoSymoa2/aLjwr3Al9GAhMDmniOSY0jQiN4ikgQKusa2VJaxYKAJuYiIjIwKNyJDABhYUbWsDiyhsXx0SnD29Y759hX3dDWrHOLH/re3LKPv76/s22/qPAwxqV7/fpGJMUwIimG4YnR/uNohifGkBQboV/pRYTVO8txTv3tREQGIoU7kQHMzEhPiCY9IZqTJ6Qdsq2irrEt7G0u9YLfxt2V/GfTXqrqmw47V3REGMOTohmR6IW/DD/8KQSKDC0FheUA5GQld7KniIj0Nwp3IoNUUkwkx48ZxvFjhh22rbq+iT2V9eypqGO3f7+nsp7dFXXsqahn/a4K/v1BvUKgyBBUUFjGuLQ4UuKiQl0UERHpIoU7kSEoPjqC8dERjE+PP+p+rSFwtx/+AkPg7oo6hUCRQaigqIw541M731FERPodhTsROaLuhMDdAUEwmBCY4TcvTU/07jMSorz7xIPr46PCFQRlwDCzc4D7gHDgYefcPe22DwMeBY4D6oDPOufWBGwPB1YAO51z5/dFmXdX1FFSXkduVkpfPJ2IiPQwhTsR6bbuhsDdFfXsrapn+75qVnx4gP3VDR0eHxMZ1tbH0At+Ue2Wo0lPiCI9MZrEaNUISuj4wex+4JNAEbDczBY759YF7PY1IN85d7GZTfX3nxew/XZgPZDUR8XW5OUiIgOcwp2I9JlgQ2BTcwv7qxvYU+mFvr1VDd59wHLRgRryC70gGDDve5uoiDC/NjAgAAaEwYy2GsJoNQ2V3jAH2Oyc2wpgZk8DC4DAcDcd+D8A59wGMxtnZiOcc7vNLAs4D/gB8D99VehVReVEhBkzRvVZnhQRkR6kcCci/U5EeBjDk2IYnhTT6b7NLY791X74a71VesulfhAsKa9j9c5y9lU30NxBEowKDyOtLQRGHdI8ND0hiuTYSFLi/PvYSJJiIwkPUxiUoxoNFAYsFwFz2+1TAFwCvG5mc4CxQBawG/g58L9AYq+XNLBARWVMGZlITGR4Xz6tiIj0EIU7ERnQwsOMjESvJq4zLS2OAzUNB2sCq+opraw/dLmqnvUlleyrrqexuYMqQV9iTATJsZF+8Iv0H0cdspzib09uXY6LUr/BoaOjf+T2H6h7gPvMLB9YDbwPNJnZ+cAe59xKM/voEZ/AbCGwEGDMmDHdLnBLi6OgsIzzNXm5iMiApXAnIkNGWJiRlhBNWkI0UzqpEHHOUV7byN6qBsprG6mobaSstoHymkbKahspr22kvMa7L6ttZHdFFWU1jZTXNhw1FEaEWVsoTD4sBEYdstwWGv376AjVpgwgRUB2wHIWUBy4g3OuArgBwLzEv82/XQFcaGbnAjFAkpn9zjl3dbvjFwGLAGbPnn3kD12Qtu+rpqKuiTwNpiIiMmAp3ImIdMDMSImL6vJcX845ahubvdDXGv5qAsJhwPry2kb2VzewtbTaC5B1jbijfEWPjQxvC31JfgBMiokkKTaCpBh/OTaSpJiIg9v95QQNMNPXlgOTzGw8sBMvsF0VuIOZpQA1zrkG4CbgNT/wfdW/4dfc3dk+2PWGgqIyQIOpiIgMZAp3IiI9yMyIi4ogLiqCzOTYLh3b3OKoqms6LASWtdYc1hxcX1bbSOH+GirrmqiobaSyg2kmAoUZJLYFQC8MtgbDgyGxo6DorYuNVHPSrnDONZnZrcBSvKkQHnXOrTWzW/ztDwLTgCfMrBlvoJUbQ1ZgoKCwnLiocCYOTwhlMUREpBsU7kRE+onwMPOaYMZFdvnY5hZHZV0jFbVNVNR5YbCirtFvUhq4rqmtmenWvVVt22oamo96/shwOxgA/ZrBpHY1h621henxUZw6Mf1Y34ZBwzm3BFjSbt2DAY/fAiZ1co5XgVd7oXiHyS8sY9boZA0WJCIygCnciYgMAuFhx9aMtFVDU4sXDv2awI6DYSPltQe3F5fVti03NLe0nWtcWhyvfvljPfXSpA80NLWwrqSCG04dF+qiiIhINyjciYgIURFhbYPNHIu6xmY/BDbRGBD0ZGCIDDf+eftHiIoIC3VRRESkGxTuRESk22Iiw4mJDGd4n87KJj3FzJiQob52IiIDnX6iExERERERGQQU7kRERERERAYBhTsREREREZFBIKhwZ2bnmNlGM9tsZnd1sN3M7Bf+9lVmdkJnx5pZqpm9ZGab/PthPfOSREREREREhp5Ow52ZhQP3A/OB6cCVZja93W7z8ebqmQQsBB4I4ti7gGXOuUnAMn9ZREREREREjkEwNXdzgM3Oua3OuQbgaWBBu30WAE84z9tAiplldnLsAuC3/uPfAhd176WIiIiIiIgMXcGEu9FAYcBykb8umH2OduwI51wJgH8/PPhii4iIiIiISKBgwp11sM4FuU8wxx79yc0WmtkKM1tRWlralUNFRERERESGjGDCXRGQHbCcBRQHuc/Rjt3tN93Ev9/T0ZM75xY552Y752ZnZGQEUVwREREREZGhJ5hwtxyYZGbjzSwKuAJY3G6fxcC1/qiZJwPlflPLox27GLjOf3wd8PduvhYREREREZEhy5zrvJWkmZ0L/BwIBx51zv3AzG4BcM49aGYG/Ao4B6gBbnDOrTjSsf76NOBPwBhgB3CZc25/J+UoBT7s+ss8RDqwt5vnGGr0nnWd3rOu03vWdYP5PRvrnFNzjSD10PURBvdnqrfoPes6vWddo/er6wb7e3bEa2RQ4W4wMbMVzrnZoS7HQKL3rOv0nnWd3rOu03smPU2fqa7Te9Z1es+6Ru9X1w3l9yyoScxFRERERESkf1O4ExERERERGQSGYrhbFOoCDEB6z7pO71nX6T3rOr1n0tP0meo6vWddp/esa/R+dd2Qfc+GXJ87ERERERGRwWgo1tyJiIiIiIgMOkMq3JnZOWa20cw2m9ldoS5Pf2dm2Wb2LzNbb2Zrzez2UJdpIDCzcDN738yeC3VZBgIzSzGzZ8xsg/9ZOyXUZervzOyL/v/JNWb2BzOLCXWZZGDT9bFrdH08drpGdo2ukV031K+RQybcmVk4cD8wH5gOXGlm00Nbqn6vCfiSc24acDLweb1nQbkdWB/qQgwg9wH/dM5NBXLRe3dUZjYa+AIw2zk3E28O0StCWyoZyHR9PCa6Ph47XSO7RtfILtA1cgiFO2AOsNk5t9U51wA8DSwIcZn6NedciXPuPf9xJd4flNGhLVX/ZmZZwHnAw6Euy0BgZknAGcAjAM65BudcWUgLNTBEALFmFgHEAcUhLo8MbLo+dpGuj8dG18iu0TXymA3pa+RQCnejgcKA5SL0hzhoZjYOOB54J8RF6e9+Dvwv0BLicgwUE4BS4DG/mc7DZhYf6kL1Z865ncC9wA6gBCh3zr0Y2lLJAKfrYzfo+tglP0fXyK7QNbKLdI0cWuHOOlinoUKDYGYJwF+AO5xzFaEuT39lZucDe5xzK0NdlgEkAjgBeMA5dzxQDai/z1GY2TC8WpXxwCgg3syuDm2pZIDT9fEY6foYPF0jj4mukV2ka+TQCndFQHbAchZDrJr2WJhZJN6F6ynn3F9DXZ5+7jTgQjPbjtes6eNm9rvQFqnfKwKKnHOtv3g/g3chkyP7BLDNOVfqnGsE/gqcGuIyycCm6+Mx0PWxy3SN7DpdI7tuyF8jh1K4Ww5MMrPxZhaF17lycYjL1K+ZmeG1817vnPtpqMvT3znnvuqcy3LOjcP7fL3inBtSvxZ1lXNuF1BoZlP8VfOAdSEs0kCwAzjZzOL8/6PzUAd76R5dH7tI18eu0zWy63SNPCZD/hoZEeoC9BXnXJOZ3QosxRs551Hn3NoQF6u/Ow24BlhtZvn+uq8555aErkgyCN0GPOV/qdwK3BDi8vRrzrl3zOwZ4D28EfveBxaFtlQykOn6eEx0fZS+omtkF+gaCeacmtWLiIiIiIgMdEOpWaaIiIiIiMigpXAnIiIiIiIyCCjciYiIiIiIDAIKdyIiIiIiIoOAwp2IiIiIiMggoHAnIiIiIiIyCCjciYiIiIiIDAIKdyIiIiIiIoPA/wdADUHMxal/GAAAAABJRU5ErkJggg==\",\n \"text/plain\": [\n \"
\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"text/plain\": [\n \"
\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"import matplotlib.pyplot as plt # Visualization\\n\",\n \"\\n\",\n \"# Plot loss and accuracy in subplots\\n\",\n \"fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))\\n\",\n \"ax1.set_title('Loss')\\n\",\n \"ax2.set_title('Accuracy')\\n\",\n \"for dataset in ('train','test'):\\n\",\n \" ax1.plot(metrics_history[f'{dataset}_loss'], label=f'{dataset}_loss')\\n\",\n \" ax2.plot(metrics_history[f'{dataset}_accuracy'], label=f'{dataset}_accuracy')\\n\",\n \"ax1.legend()\\n\",\n \"ax2.legend()\\n\",\n \"plt.show()\\n\",\n \"plt.clf()\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"qQbKS0tV3sZ1\",\n \"metadata\": {},\n \"source\": [\n \"## 12. Perform inference on test set\\n\",\n \"\\n\",\n \"Define a jitted inference function `pred_step`. Use the learned parameters to do model inference on the test set and visualize the images and their corresponding predicted labels.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 66,\n \"id\": \"DFwxgBQf44ks\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"@jax.jit\\n\",\n \"def pred_step(state, batch):\\n\",\n \" logits = state.apply_fn({'params': state.params}, test_batch['image'])\\n\",\n \" return logits.argmax(axis=1)\\n\",\n \"\\n\",\n \"test_batch = test_ds.as_numpy_iterator().next()\\n\",\n \"pred = pred_step(state, test_batch)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 67,\n \"id\": \"5d5nF3u44JFI\",\n \"metadata\": {\n \"outputId\": \"1db5a01c-9d70-4f7d-8c0d-0a3ad8252d3e\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"image/png\": \"iVBORw0KGgoAAAANSUhEUgAAAqkAAAKqCAYAAAAZssdpAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjYuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/av/WaAAAACXBIWXMAAAsTAAALEwEAmpwYAABhcUlEQVR4nO3debxV8/7H8c+neZ7k0qDipktRcRMqDcqQuBVFbshM5Ip0yVS5dCV0dQ0ZKq6hIopKoSRjUt1QJNU9NKGRSnPr98c5Hr/z+e5jD2dP33XO6/l47MfjvPdee63vOefb2p+z+uzv1iAIBAAAAPBJiWwPAAAAAHBRpAIAAMA7FKkAAADwDkUqAAAAvEORCgAAAO9QpAIAAMA7oShSVTVHVTvFsV2gqg0LeYxCPxf+YK4gHswTxIu5gngwT9IjFEWqr1S1rKqOVtUfVXWzqk5V1TrZHhf8o6odVHWOqv6sqjnZHg/8pKr9VXWVqv6iqutUdaSqlsr2uOAfzimIh6oOUdW9qro93+2IbI8rXhSpyblRRE4WkaYiUltEtorIv7M5IHhrh4iMFZGB2R4IvDZVRI4PgqCKiBwjIs1E5G/ZHRI8xTkF8ZoYBEGlfLdV2R5QvEJVpKpqS1X9RFW3qup6VX1UVcs4m52VdyVio6qOUNUS+Z5/uap+rapbVPUtVa2f5JAOF5G3giD4MQiCXSIyQUSaJLlPpIBvcyUIgvlBEDwvIqE5ORQHHs6TlUEQbP1t9yJyQESK1X/v+crDucI5xUO+zZOwC1WRKiL7ReQmEakpuVcwO4rIdc423UWkhYgcLyJdReRyERFV7SYit4vIuSJysIh8ICLjCzqIqt6WN8EKvOXbdIyItFbV2qpaQUR6i8iMlHynSJZvcwV+8m6eqOpfVfUXEdkouVdSn0zFN4qkeTdX4CUf58k5mtuSuFRV+6bim8yYIAi8v4lIjoh0KuD+/iIyOV8OROTMfPk6EZmd9/UMEbki32MlRORXEamf77kNExxXFcmdQIGI7BOR/4pIjWz/vIrzzde5km9fnUQkJ9s/p+J+832e5D3/SBH5h4gcmu2fV3G++T5XOKf4cfN1nohIY8ltRywpIq1EZL2IXJjtn1e8t1BdSVXVRqo6TVV/yLvSMExy/1rJb3W+r7+T3F+OiEh9EXkk318ZmyX3v9OSeaPTEyJSTkQOEpGKIvKacCXVCx7OFXjI53kSBMG3IrJURB5Pxf6QHJ/nCvzh2zwJguCrIAjWBUGwPwiCj0XkERHpUdj9ZVqoilTJLQqXiciRQe4bC26X3F9gfofl+7qeiKzL+3q1iFwTBEG1fLfyeb80Q1VvV/tOOHPLt2kzEXk2CILNQRDsltw3TbVUVXdCIvN8myvwk+/zpJSI/LHQ3x1Syfe5Aj/4Pk+CAsbjrbAVqZVF5BcR2a6qR4lIQb0VA1W1uqoeJrnvvp+Yd/9oERmkqk1ERFS1qqr2LOggQRAMC+w74cwt36aficglefsqLbmX7dcFQbAxNd8ukuDVXFHVEqpaTkRK50Ytp5HN9Mg83+bJlar6h7yvG4vIIBGZnapvFknxba5wTvGTb/Oka96xVFVbSu5qIa+n7ttNr7AVqbeIyF9FZJuIPC3//4vN73URWSgii0VkuuS+uUmCIJgsIsNFZELeJfglItI5BePZJSLfisgGETlLchuikX2+zZW2IrJTRN6U3L+cd4rI20nuE8nzbZ60FpEvVXWH5M6VNyX3Sgyyz7e5wjnFT77Nk14isiJvPP8RkeFBEDyX5D4zRvMaawEAAABvhO1KKgAAAIoBilQAAAB4hyIVAAAA3qFIBQAAgHdKRXtQVXlXVcgFQZCR9dCYK+GXibnCPAk/zimIF+cUxCPaPOFKKgAAALxDkQoAAADvUKQCAADAOxSpAAAA8A5FKgAAALxDkQoAAADvUKQCAADAOxSpAAAA8A5FKgAAALxDkQoAAADvUKQCAADAOxSpAAAA8A5FKgAAALxDkQoAAADvUKQCAADAOxSpAAAA8A5FKgAAALxDkQoAAADvUKQCAADAO6WyPQCgqLj00ktNHjdunMmzZs0y+bTTTkv3kIq92rVrm1yrVi2TDzrooIT2d+qpp0bdfxAEEc+ZPn26ybNnzzZ506ZNCY0B4fTBBx+YXNDvvXfv3ibv2LEjrWMCfMeVVAAAAHiHIhUAAADeoUgFAACAd+hJdbRp08bkbt26mVyjRg2Tt27davJ9991n8vjx4012+xDfeOMNk7t27RrvUOGZ008/3eQDBw6Y3LZtW5M7dOhg8pw5c9IzsCLs+eefN7ljx44mly9f3uRy5cqZXLZsWZML6imNRlVjPv+iiy4yecuWLSa/9957Ji9cuNDkf/3rXybv3LkzoTHCD7t37zb5zDPPjNjmyCOPNHnx4sXpHBI81LlzZ5Pr1q0bsc2DDz5ocpUqVUyeNm2ayY8//rjJM2bMSGaIGcWVVAAAAHiHIhUAAADeoUgFAACAdzRaD5aqJtag5blOnTqZfOedd0Zs4/akliiRWB2/fv16k911GV179+412e2RS1YQBBp7q+QVtbkSj+rVq5u8YMECkxs0aGDyrl27TG7atKnJK1euTN3gCiETcyXV88Tt+3XPZ+7jGzZsSOXhI5QqFdnmH2st1lh9rVOmTDF5wIABJufk5MQ/wBTgnFI4vXr1Mrl///4R29x8880mf/zxx+kcUtqF8ZySbu7ayi+99JLJzZo1M9ntNy2MX375xWS319mdi8uWLTPZ7adOtWjzhCupAAAA8A5FKgAAALxDkQoAAADvFOl1Um+44QaT3TVMK1WqFHMfbm+G23d43HHHmdykSZNEhpj1PkQU3iWXXGKy24Pq+vbbb03md5+8M844I+rje/bsMXnu3LnpHE5EP5mIyNNPP23yHXfcYbLbt/7EE0+Y7K7VPHr0aJMz3ZOKwrn++utNPvHEEyO26d27t8lh70mFyCGHHGLy1KlTTW7evHnax+D2tbprdi9atMjke+65x+ShQ4emZ2Bx4EoqAAAAvEORCgAAAO9QpAIAAMA7oe5JLVOmjMlPPvmkyW7PoLse4aZNmyL26X6esrue2P79+01210CcPn26yS1btow4Rn7Dhw+P+jj81aNHj4S2z2ZfT1H1zjvvZHsIxueffx5xX6xzwOmnnx71cfe8hXCaNWuWya1bt87SSJBJ1apVMzkTPajJcteQr1Gjhsk33nhjxsbClVQAAAB4hyIVAAAA3qFIBQAAgHdC3ZPavn17k/v06RN1e7cH9eyzz47YZuHChQmNoUQJW+dXrFgx6vZr1qwxec6cOQkdD9nhrmX5e/dFM23atFQNByFSr149k/v162eyu35muXLlTH7xxRdNTvdar0iPr776KttDQAZUrVrV5Lvvvjvlx9i+fbvJa9euNdl9r0zNmjVNHjx4sMnu+3t27txp8gUXXFCocaYCV1IBAADgHYpUAAAAeIciFQAAAN7RIAh+/0HV33/QA7Nnzza5Q4cOUbd310B9++23Ez5mnTp1THbXE7vmmmuiPv+EE04wOdEe2EQFQZCRRRZ9nyvJateuXcR97777btTnfPrppyafcsopJrtr7mZbJuZK2OdJ+fLlTXb7z6666qqI51x99dUm165d2+Q9e/aY7K6d7Ga3XyzTOKekhvt7F4l8TevcuXOmhpMWxfGc8vrrr5tc0HtfEjFz5syI+55++mmTp0yZYvJJJ51ksrsm79SpU01evnx5EiNMXrR5wpVUAAAAeIciFQAAAN6hSAUAAIB3QrVOauXKlU3+4x//GHV7t5fD/ezkeBx22GEmP/744yZ36dLF5AMHDph88803m7xo0aKEx4DMO/TQQ00eM2ZMwvu47777TPatBxWRevbsafJ5551n8tFHH23ysccea3K0Hv/fc/HFF5s8adKkhPeB8HF71kUi3/OA8Dn99NOTev7kyZNN7t27d8Q2u3fvjrqPefPmRc1hwpVUAAAAeIciFQAAAN6hSAUAAIB3QtWT6vZ/uZ+J7XLXAnP7RQvi7nP69OkmN2nSxOR9+/aZfMcdd5g8atSomMeEf2rUqGHy4YcfHvM5H374ocmx1lFF6rl96/fee6/Jxx9/vMnu+oHJUk18WchHHnnE5BtvvNHkVatWRX3+iy++aPLcuXNNjtW/huz48ssvI+5z19lt2LChyStWrEjrmJC8jRs3muyui+x6//33TXbXWi/u/365kgoAAADvUKQCAADAOxSpAAAA8E6oelITFatntVWrVhH3jR071uRGjRpF3ceTTz5p8ogRI+IcHXxWt27dhJ/j9hZl+zPWi6MjjjjC5H79+iX0/LVr15q8YcMGk//3v/+Z/PHHH8fcZ7ly5UyuVq2ayZ06dTK5YsWKJp977rkmV6hQwWR3nVV3fWh3vd6PPvoo+oCRNSVLljTZ7ZmmJ9V/t99+u8nPPvts1O1r1qxpcvXq1U3etGlTSsYVVlxJBQAAgHcoUgEAAOAdilQAAAB4hyIVAAAA3gnVG6fmz58fNbds2dLks88+2+SlS5eaPHTo0IhjuIu2u2+k+Nvf/mbylClTfn/ACA33zS1///vfYz7nxx9/NPmpp55K6ZiQOPeNTg8++GDU7d2F8NevXx91f9ngvnmzS5cuJt95550mn3HGGSZ37NjR5AceeMDku+66K9khAiikxo0bm9ytWzeTY53DijqupAIAAMA7FKkAAADwDkUqAAAAvKNBEPz+g6q//6AHBgwYYHIqFtJ/5513oh5jyZIlSR8jk4Ig0Ewcx/e5Eovb11dQv7Jr6tSpJru9RGGTibkS9nnio4MOOsjkxx9/3OQePXqYvG7dOpMPO+ywhI7HOSU13N+TiEjfvn1NvvTSS01+7rnn0jmklCuO5xT3g2BmzJhhstuD6lqzZo3JTZo0idhm+/bthRydn6LNE66kAgAAwDsUqQAAAPAORSoAAAC8E6p1Ul0vvfSSyYn2pL7yyisR91100UUm7927N/GBIXRq1KiR8HMee+yxNIwESMymTZtMHjZsmMk9e/Y0uU6dOmkfEwon2ntEEA5uT+m//vUvky+//HKTTzrpJJPdntbXXnst4hjPPPOMyS+//HKiwwwNrqQCAADAOxSpAAAA8A5FKgAAALwT6p7U0047LaHtN2/ebPLFF18csQ09qMVDpUqVTL7hhhuibn/gwIGI+7Zt25bSMQGpcOWVV5rs9jkuXLgwk8NBAvbv32/y3LlzszQSpMqYMWNMdntMx44da3Lr1q1N7tixY8Q+K1asaPKcOXNM3rBhQ8Lj9BVXUgEAAOAdilQAAAB4hyIVAAAA3glVT+rxxx9v8ujRoxN6fpUqVUxu2bJlxDYffvhh4gND6Nx9990mlygR/e+1mTNnRtw3b968lI4JiatatarJ+/btM3nHjh2ZHE5G/PnPfzb5jjvuMLlLly4mu/3U48ePT8/AkDT3d5WTk5OdgSBttmzZYnL37t2j5kmTJkXsw11b1e1rdddG3rVrV8Lj9AVXUgEAAOAdilQAAAB4hyIVAAAA3glVT+o//vEPk1XV5MWLF5vcvHlzk0uVst9u9erVUzY2hEvfvn2jPr57926TR4wYkc7hoJC++uorkx9++GGTH3rooUwOJyXcfrLjjjvOZHcd1Jo1a5rsrovqnjdHjhyZ7BABpMl7771nckHvfXB7Us866yyTBw0aZPLgwYNTM7gs4EoqAAAAvEORCgAAAO9QpAIAAMA7Xvekuj2lZ5xxhslvvvmmyS+++KLJrAeI39StW9fkWOuirly50uT3338/5WNC8mrXrm3y7bffbnLp0qVNXrRokclvv/22ye3btze5TJkySY4wct1Sd71n97O6E/XOO++YfNNNN5ns9u3CDxdffHG2hwAP7dy50+QVK1ZEbOP2pBZlXEkFAACAdyhSAQAA4B2KVAAAAHjH657UY445xmS3j7BevXqZHA5CrGvXriaXK1cu6vYvvPBCOoeDFOnXr5/J7jqp9913X9Tn//TTTya7a47G6l1212p21yiNx+eff26y+9ner776qsnuZ3n/8ssvJrs9bfBT2bJlI+6bNm1aFkaCVCpZsqTJbl98LHfccYfJF110UdJjCjOupAIAAMA7FKkAAADwDkUqAAAAvON1T2os7tqX3bp1y85A4L0TTzwx6uO//vqryXPnzk3ncJAijz32mMmLFy82+amnnjK5Vq1aJlerVs3kDRs2RD2eu37uJ598YnJBPam7du0y2e0x/eKLL6IeE8XH+vXrsz0EJOnPf/6zye6/d3dt51RwX7/mz5+f8mNkC1dSAQAA4B2KVAAAAHiHIhUAAADe8bon1e0vW7BggcktWrQw+YILLoi6v++//97kWbNmFX5wCJXHH3/c5PPPP99kdz3NefPmpX1MSL2PPvrI5CZNmpjs9qS6edGiRekZGIBiwe0HXbVqlcmp6El97733TL7++utNXrZsWdLH8AVXUgEAAOAdilQAAAB4hyIVAAAA3tFonzWtqol/EHUatWnTxuSpU6eaXLVqVZPdNQ87d+5scnHoPwuCQGNvlTzf5goSl4m5wjwJP84piBfnFJGzzz7bZLd/9PTTTzd506ZNJg8aNChin+77a955551khph10eYJV1IBAADgHYpUAAAAeIciFQAAAN4JVU8qEkf/GOJF/xjiwTkF8eKcgnjQkwoAAIBQoUgFAACAdyhSAQAA4B2KVAAAAHiHIhUAAADeoUgFAACAdyhSAQAA4J2o66QCAAAA2cCVVAAAAHiHIhUAAADeoUgFAACAdyhSAQAA4B2KVAAAAHiHIhUAAADeCUWRqqo5qtopju0CVW1YyGMU+rnwB3MF8WCeIF7MFcSDeZIeoShSfaeqZVR1maquyfZY4CdV7aCqc1T1Z1XNyfZ44CdVHaKqe1V1e77bEdkeF/zDOQXxUtXjVfX9vPPJj6p6Y7bHFC+K1NQYKCI/ZXsQ8NoOERkruXMFiGZiEASV8t1WZXtA8BLnFMSkqjVFZKaIPCkiB4lIQxF5O6uDSkCoilRVbamqn6jqVlVdr6qPqmoZZ7OzVHWVqm5U1RGqWiLf8y9X1a9VdYuqvqWq9VMwpsNF5CIR+Wey+0Lq+DZXgiCYHwTB8yJCweER3+YJ/OXbXOGc4iff5omI3CwibwVB8GIQBLuDINgWBMHXSe4zY0JVpIrIfhG5SURqisjJItJRRK5ztukuIi1E5HgR6Soil4uIqGo3EbldRM4VkYNF5AMRGV/QQVT1trwJVuDN2fzfefvdmfy3hxTyca7APz7Ok3NUdbOqLlXVvqn4JpESPs4V+Me3eXKSiGxW1Y9V9SdVnaqq9VL0vaZfEATe30QkR0Q6FXB/fxGZnC8HInJmvnydiMzO+3qGiFyR77ESIvKriNTP99yGCY6ru4jMzPu6vYisyfbPqrjffJ0r+fbVSURysv1zKu43X+eJiDQWkdoiUlJEWonIehG5MNs/r+J883Wu5NsX5xQPbr7OExFZLiJbReQEESknIqNE5KNs/7zivYXqSqqqNlLVaar6g6r+IiLDJPevlfxW5/v6O8k94YuI1BeRR/L9lbFZRFRE6hRyLBVF5AERuaEwz0d6+TRX4C/f5kkQBF8FQbAuCIL9QRB8LCKPiEiPwu4PqePbXIGfPJwnOyW3SP4sCIJdIjJURFqpatUk9pkxoSpSReQJEVkmIkcGQVBFci+Lq7PNYfm+rici6/K+Xi0i1wRBUC3frXzeC4GhqrerfXetueVtdqSINBCRD1T1BxF5TURq5U3MBqn6hlFoPs0V+Mv3eRIUMB5kh+9zBX7wbZ58Ibnnkd/89nUozithK1Iri8gvIrJdVY8SkYL6tQaqanVVPUxEbhSRiXn3jxaRQaraREREVauqas+CDhIEwbDAvrvW3PI2WyK5E6153u1KEfkx7+vVBewWmeXTXBFVLaGq5USkdG7UchrZTI/M822edM07lqpqSxH5m4i8nrpvF0nwba5wTvGTV/NERMaJSHdVba6qpUXkLhH5MAiCrSn5btMsbEXqLSLyVxHZJiJPy///YvN7XUQWishiEZkuImNERIIgmCwiw0VkQt4l+CUi0rmwAwmCYF8QBD/8dpPcy/IH8vL+wu4XKePNXMnTVnL/2+VNyf3LeaeEaBmQIsy3edJLRFbkjec/IjI8CILnktwnUsO3ucI5xU9ezZMgCN6V3Ku50yV3qcyGeeMLBQ2CIPZWAAAAQAaF7UoqAAAAigGKVAAAAHiHIhUAAADeoUgFAACAd0pFe1BVeVdVyAVBkJG10Jgr4ZeJucI8CT/OKYgX5xTEI9o84UoqAAAAvEORCgAAAO9QpAIAAMA7FKkAAADwDkUqAAAAvEORCgAAAO9QpAIAAMA7UddJBQAAQHa0b98+4r45c+aYPHToUJOHDBmSxhFlFldSAQAA4B2KVAAAAHiHIhUAAADeoScVAADAA24Pqtt/WtxwJRUAAADeoUgFAACAdyhSAQAA4B16UgEASNJDDz1kcv/+/U1+9dVXTT7//PPTPSSEUEHrosby3nvvpXwcvuBKKgAAALxDkQoAAADvUKQCAADAOxoEwe8/qPr7DxZTkyZNMrl69eomd+zYMZPDiSkIAs3EccI2VypXrmzywoULTd65c6fJN9xwQ8Q+3n///dQPLIsyMVfCNk8QiXNKwfbv32/ygQMHTF63bp3JF1xwQcQ+5s2bl/qBZRHnlNiGDBli8uDBg2M+x+1B7dChQwpHlHnR5glXUgEAAOAdilQAAAB4hyIVAAAA3qEnNYbWrVub7PaCzJ071+ROnTqle0gJoX+sYGXKlDF5xowZJrdr187k2bNnR+zjjDPOSP3Asoj+McSDc0rB3HVPx48fb3KJEvaakNuzKiJSsmTJ1A8sizinRHLXQZ0zZ07C+1DNyD/BjKEnFQAAAKFCkQoAAADvUKQCAADAO6WyPYBUqlGjhsmbN29Oep8NGzY0uVSpIvUjK7b27Nlj8saNG6NuX69evYj73L5Wd58Aig/3/R0F9Zwm8jiKhmR7UMO+BmqyuJIKAAAA71CkAgAAwDsUqQAAAPBOqBssmzRpYvKzzz5rctu2bU12P489Hsccc0zUxydMmJDwPhE+jRo1irjv5JNPNtldMxfFT7Vq1SLuGz58uMlffvmlyY8++mg6h4QMcdeudNdFdXNBXn75ZZPdtVcRPm5PaizuWuxuLm64kgoAAADvUKQCAADAOxSpAAAA8E6oe1Jvvvlmk1u0aGFy+fLlTY6nJ7Vq1aomX3311Sbv3r3b5BdeeCHmPgEUTe46yvPnz4/Yxu1TddfTHDFihMnbt283+ZVXXjF56tSpJn/66acmp2J9aCQuFeuksnZq0dOuXbuEtue9DRZXUgEAAOAdilQAAAB4hyIVAAAA3glVT2qpUna4xx9/fNTt3c9bj6dX68gjjzS5SpUqJk+ZMsXkXbt2xdwn/Pfhhx+a3KNHD5PdNRBFRPr27WsyvURFX5s2bUx2zwcFrZPq+uKLL0xu1qyZyeXKlTP52muvjZrdXvv9+/ebvGrVKpMXLFhgstsH6a7b6o4XBUvFOqnua1bdunVNXrNmTSFHh2xJdJ3UIUOGpGUcYcWVVAAAAHiHIhUAAADeoUgFAACAd0LVk3rZZZeZ3Lx586jbf//99wkfg89KLp7cz1N31zxE8VCjRg2TH3/8cZO7detmcpkyZUwu6Jxz6623mjx58mST//nPf5p80003xTXW37jrQbuaNm0aNbu9lH369DG5bNmyCY2nuPr444+j5latWplc0JqoJ554YtRMT6r/Eu0p7dChQ3oGUkRwJRUAAADeoUgFAACAdyhSAQAA4B2ve1LdXqjrr78+6vbPPvusyVu2bIm6fUG9XGeffXZ8g0ORsnfvXpPdtSbdNXpFRBo3bmxyxYoVTd6xY0eKRod0cXtQZ86caXKLFi2iPn/btm0mF9SPNnHixKj7cHtW77//fpPdeZXsOapkyZIm/+EPfzB5/vz5Se2/uHL7RV977TWTW7dubXJB66a6/cEvv/yyye7vDv4ZPHhwQtu/99576RlIPnPmzIn6uLvGt09rtXIlFQAAAN6hSAUAAIB3KFIBAADgHY22HqSqZnWxyEsvvdTkcePGRd1+9uzZJn/zzTdRtz/mmGMi7mvbtm3U57i9iytXrjR57NixJo8YMSLq/tItCILID51Pg2zPlVRz1011+09FItdSrVWrlskbNmxI/cDSKBNzJdvzJNkeVHf7e++912R3bcyiiHNK4bh97gWtk+r2qbrblC5dOvUDS6PicE5xJbrGttuHXBjt27ePmhPtk03FmBIRbZ5wJRUAAADeoUgFAACAdyhSAQAA4B2v1kl110UdMGBAQs/v2LFj1JwKbk/QUUcdZXLPnj1NznZPKlCcVa9e3eREe1BfeOEFky+77DKT3T5D4Pd8+umnJp944okR27i9gAWtpQq/JLqm6NChQ1M+Brfn1O1JDTP+BQAAAMA7FKkAAADwDkUqAAAAvEORCgAAAO949cYpdzH0Ro0aZWkk/++HH34wOdYHCvzxj39M53CQJQW9gaGgxbjhlxtvvNHkWG+Uuu+++0x23xTBG6VQWCNHjjT5pZdeitgm1mL+N910U9R9wn+JvtGqIHPmzDG5KL1RysWVVAAAAHiHIhUAAADeoUgFAACAd7zqSc3JyTH5zjvvNLlevXoJ7e/VV1812V3c392/iMjGjRtNPvroo03eunVrQmNA0VBQ/2kQBFkYCRJx8803R338/fffN9ldFJu+Y6RLQX3usRbzP+mkk9I6JqSf2z/63nvvJb2PoowrqQAAAPAORSoAAAC8Q5EKAAAA73jVk+oaMWJESvd3xhlnxNxmwoQJJtODCoTXzJkzTe7Ro4fJ7trMBfWp57dhwwaT33zzTZO/++67RIeIYqqgfudY66TSB+8ft6fU7Wt3uY/H6klNxbqqsRSmLzZTuJIKAAAA71CkAgAAwDsUqQAAAPCO1z2pqdamTZtsDwEhsWTJEpMbN26cpZEgGVdffbXJNWrUMLlDhw4mJ9r/tWPHDpMfe+yxiG1uu+22hPaJomn16tUmr1u3LmKbww47zGS3R9VdRxXZl2g/p7vGqdtnPHToUJNj9bimgntMn3AlFQAAAN6hSAUAAIB3KFIBAADgnWLVk1qvXr2Y20ycODEDI4HvjjnmmGwPASngrnPcqVMnk48//niT3XVTjzjiCJPdz06/8MILTT711FMLM0wUA/PmzTP5k08+idimbt26JrvrpLrzz83uMZB5bt/wnDlzTHZ7Ul2Z6EF1e/FZJxUAAABIAEUqAAAAvEORCgAAAO8U6Z7UBg0amFy9enWTt2zZEvGcVatWpXNICCl3vUKRgj97G+GyaNGihLY/6KCD0jQSFDcFrXnq3ueed9x1VN0eVvgn1hqksXpUUyHM6+tyJRUAAADeoUgFAACAdyhSAQAA4J0i3ZNap04dkytXrmzyd999F/Gcgj5PGcXPlClTTG7cuHHENu5nLqPocdfLveKKK7I0EhQ1I0eOjLivR48eJrt9726P6o033mjypEmTUjQ6pIq7BmmsNUmHDBkSc5/uWqruPt11UMOMK6kAAADwDkUqAAAAvEORCgAAAO8U6Z7Uo48+OurjM2bMyNBIEDb0JkNEpGLFiib/4Q9/iLr9mDFj0jkcFCHz5s2LuC/WOqlhXu8S8YmnJzWebYoKrqQCAADAOxSpAAAA8A5FKgAAALxTpHtSjzvuuKiPb9myJUMjARBGd999t8mlS5c22V3Hcvr06WkfE4quhx56yOT+/fub7PaxXnjhhekeEpBVXEkFAACAdyhSAQAA4B2KVAAAAHinSPekup9jfOmll5r83XffZXA0CJOJEyeafO2110Zss3btWpPpcS565s6da3KnTp1Mds8xa9asSfuYUHQNHDgwagaKG66kAgAAwDsUqQAAAPAORSoAAAC8o0EQ/P6Dqr//IEIhCIKMfNgzcyX8MjFXwj5P6tSpE/Vxt0+5KOKcgnhxTkE8os0TrqQCAADAOxSpAAAA8A5FKgAAALxDT2oRR/8Y4kX/GOLBOQXx4pyCeNCTCgAAgFChSAUAAIB3KFIBAADgHYpUAAAAeIciFQAAAN6hSAUAAIB3KFIBAADgnajrpAIAAADZwJVUAAAAeIciFQAAAN6hSAUAAIB3KFIBAADgHYpUAAAAeIciFQAAAN4JRZGqqjmq2imO7QJVbVjIYxT6ufAHcwXxYJ4gXswVxIN5kh6hKFJ9paoDVXWJqm5T1f+p6sBsjwl+UtUZqro9322Pqn6Z7XHBL5pruKpuyrs9oKqa7XHBP7z+IB5hnyelsj2AkFMRuUREvhCRP4rI26q6OgiCCdkdFnwTBEHn/FlV3xORd7MzGnjsahHpJiLNRCQQkXdEZJWIjM7imOAnXn8Qj1DPk1BdSVXVlqr6iapuVdX1qvqoqpZxNjtLVVep6kZVHaGqJfI9/3JV/VpVt6jqW6paP5nxBEHwQBAEi4Ig2BcEwTci8rqItE5mn0gN3+aKM7YGInKKiDyfqn2icDycJ31E5KEgCNYEQbBWRB4SkUuT3CdSwLe5wuuPn5gnqRWqIlVE9ovITSJSU0ROFpGOInKds013EWkhIseLSFcRuVxERFW7icjtInKuiBwsIh+IyPiCDqKqt+VNsAJvv/McldzCY2lS3yFSxdu5Irl/1X4QBMH/kvj+kBq+zZMmIvJ5vvx53n3IPt/mSv7n8PrjD+ZJKgVB4P1NRHJEpFMB9/cXkcn5ciAiZ+bL14nI7LyvZ4jIFfkeKyEiv4pI/XzPbZjEGIdK7gtK2Wz/vIrzLSRzZYWIXJrtn1Vxvvk6TyT3Be6ofPnIvP1otn9mxfXm61xxxsLrD/OkSM6TUF1JVdVGqjpNVX9Q1V9EZJjk/rWS3+p8X38nIrXzvq4vIo/k+ytjs+T2atRJwbj6Se7VsS5BEOxOdn9InsdzpY2IHCoik5LdF5Ln4TzZLiJV8uUqIrI9yHuFQfZ4OFd+GxevPx5hnqRWqIpUEXlCRJaJyJFBEFSR3Mvi7jtfD8v3dT0RWZf39WoRuSYIgmr5buWDIPjYPYiq3q72ndjm5mx7uYjcJiIdgyBYk6LvE8nzbq7k6SMirwVBUNBjyDzf5slSyX3T1G+aSZj+a65o822u8PrjJ+ZJCoWtSK0sIr+IyHZVPUpE+hawzUBVra6qh4nIjSIyMe/+0SIySFWbiIioalVV7VnQQYIgGBYEQaXfu/22nar2lty/kk4LgmBV6r5NpIBXcyVvP+VFpKeIPJuS7xCp4Ns8+Y+I3KyqdVS1togMEOaLL7yaK7z+eIt5kkJhK1JvEZG/isg2EXla/v8Xm9/rIrJQRBaLyHQRGSMiEgTBZBEZLiIT8i7BLxGRzgU8PxH3ishBIvJZvr9gWCrGD77NFZHcpYV+FpE5KdgXUsO3efKkiEwVkS/z9jc97z5kn29zhdcfPzFPUkhpdQIAAIBvwnYlFQAAAMUARSoAAAC8Q5EKAAAA71CkAgAAwDuloj2oqryrKuSCIHDXZ0sL5kr4ZWKuME/Cj3MK4sU5BfGINk+4kgoAAADvUKQCAADAOxSpAAAA8A5FKgAAALxDkQoAAADvUKQCAADAOxSpAAAA8A5FKgAAALxDkQoAAADvUKQCAADAOxSpAAAA8E6pbA8AAAAAIg0bNjT5pptuitimb9++Ufcxbdo0k6+66iqTf/zxx0KOLvO4kgoAAADvUKQCAADAOxSpAAAA8I4GQfD7D6r+/oMZcNhhh5ncqlUrk9u0aWNyt27dTK5Ro4bJa9euNfmzzz6LOOYNN9xg8ubNm+Maq6+CINBMHCfbcwXJy8RcKerzxO0nExE5+uijk9rn22+/bfLu3buT2l+yiuo5pUKFCibfeeedJh977LEmd+nSJanj/fTTTxH3ub2ErhdffNHkhQsXmvzLL78kNaZU45wSqVQp+1agwYMHm9yvXz+Tq1SpkvQxv/jiC5Pdubtu3bqkj5GMaPOEK6kAAADwDkUqAAAAvEORCgAAAO9ktSf1yiuvNPmvf/2ryY0bNza5Zs2aJqvaNoZo30u8HnnkEZMHDBiQ9D6zqaj2jyH16B+LrWrVqiY//PDDJrvnMBGRsmXLJnXM7777zuR77rnH5HHjxiW1/0QV1XNKz549TZ4wYYI7HpOTfb1x91eYfc6aNcvkgQMHmuz2ImYa55RI119/vcmjRo0yOZ559umnn5p83HHHmVymTJmo+xw+fLjJgwYNijLi9KMnFQAAAKFCkQoAAADvUKQCAADAO6Vib5I6F198scn//ve/TS5durTJ7ppvGzZsMNnts/j2229Nfuedd0yuXbu2yZdddlnEGHv37h11jDk5ORHPAVA0HXXUUSa755Q6deqkfQz169c3efTo0SZXrlzZZLfHDfG59dZbU7q/bdu2mbxy5UqT3T7CwujUqZPJbq+h22e7ffv2pI+J5Ljrv8fyn//8J+K+a665xmR3jfgxY8aYXLFixYSO6ROupAIAAMA7FKkAAADwDkUqAAAAvJPRntTu3bubvHXrVpOfe+45kx999FGT16xZk9LxtG7dOuI+d23Wq666yuQ77rgjpWNAOBTU0+N+lrerR48eJp977rkmH3744QmNwe1pc+fqnj17EtofIntK3X/vbt+6u/3y5ctNdnvDCuOGG24w+cILLzS5WrVqJo8YMcLkL7/80uQ5c+YkPabioF69elEf37Fjh8kjR440eenSpSa/9dZbJu/atcvkeM4pJ5xwgsnu57j379/f5DPOOMPkSZMmmeyek+hRzbwuXboktL37718kcu3k+++/3+TFixebXFCtExZcSQUAAIB3KFIBAADgHYpUAAAAeEejfVZwqj8T99prrzV50aJFJs+fPz+Vh4vJ7d0Siezze+WVV0zu1atXWseUakX1c7YbNGhgsruWpPvZxiVK2L/HOnfubPJ5551n8jHHHGNy+fLlI8bwxz/+Ma6xpovb07Zz586k9lccPmd7yJAhJrs9fW7PXyxun2JB54fp06cntE+X2yt58803m9y3b1+Tf/rpJ5PbtWtn8qpVq5IaT1E9p3zwwQcmt2rVyuRmzZqZvGTJkrSPKRb3POX2oB555JEm//3vfzf5oYceSs/A8hSHc0qiBgwYYLLbU+6u/15QjbZp0yaTGzZsaPK0adNMbtOmjcnuerq33XZblBGnX7R5wpVUAAAAeIciFQAAAN6hSAUAAIB3MrpOqvuZ0z5yexe3bNmSpZEgv7Jly5rsrv3o9qSuXbvW5IMOOsjkcuXKJT2mBQsWmOx+Vrfb8+z2ybrrXT722GNRj/fjjz+afODAgXiGWawccsghJrufx3755Zeb7Pag7t+/3+ScnByT3Z6/WbNmRd0+Fb7//nuT3T5at3eyRYsWJrtrbybbk1pUnXPOOSa7c8eHHlSXO6bXXnvNZLfX0O1NTHdPKiL9+9//Ntk9j7vrae/evTtiH3feeafJpUrZUs5dz9nta432XiTfcCUVAAAA3qFIBQAAgHcoUgEAAOCdjPakZlvz5s1NdvsYRUS2bt1qsts/guxwe2h++OEHk2vWrGmy22vofka1m//zn/+Y7PZ2uccTiex7Lah3KJpYa9O5c9Fd2zXR4xUHr7/+usktW7aMur3bg+quI+l+PruP3L7C8ePHm+x+Xrv7M0Iu99/boEGDsjOQJPzjH/8w+fzzzzfZ7V1034NBn3v67dmzx2T3HFOYc87pp59uckG1TVhxJRUAAADeoUgFAACAdyhSAQAA4J0i3ZPq9tu4695VqFAh4jnuNl999VXqB4aEuX08J598ssmNGjUy2V2zdP369ekZWALcfse77ror6vYTJ040efHixakeUugtWrTIZPfz1V3uOqbu52hPnjw5JePKpMMOOyzq4xdeeKHJF198cTqHgyzauXOnye5586yzzjK5adOmJnOOCacxY8ZEfXzHjh0mf/LJJ+kcTkpxJRUAAADeoUgFAACAdyhSAQAA4J0i3ZPq9tv07Nkz5nNWrlyZruEgjZYvX57tIURwPwt86NChJpctW9Zk93PhBw4cmJ6BhdhFF11kstuDqqomL1261GS3J2/16tUpHF12tGjRIurjbm8+8JvTTjvNZHpSw8l9LXG9/PLLJr/xxhvpHE5KcfYCAACAdyhSAQAA4B2KVAAAAHinSPekup9h7SponcopU6akaTQo6g4++GCTR48ebbLbN/Tiiy+afPPNN5u8ffv2FI6uaOjcubPJsXpQ+/bta3JR6EEtV66cyW7vveu5555L53DgEXcuxFpDF+F09dVXm1yzZs2o24dx/effcCUVAAAA3qFIBQAAgHcoUgEAAOCdUPekur1Zjz32mMnu57m7Xn/99ZSPCcVHyZIlTX777bdNrlWrlsnu5yc///zzJm/YsCGFoyua3HVOXW7/5YcffpjO4WTF3//+d5OPOuqoqNu/+uqr6RwOPOL++6hYsWKWRoJUcV9HRERGjBhhchAEJn/22WcmT5s2LfUDyxCupAIAAMA7FKkAAADwDkUqAAAAvEORCgAAAO+E+o1TAwYMMLlPnz4mu83E7hsIvvvuu/QMDMXCwIEDTW7WrJnJe/fuNfnCCy802X2jFWKrVq2aye6/8aKoYcOGJl9//fVRt9+zZ4/JOTk5qR5SKFWvXt3kq666ymT3zSZffvmlye6Ha+zatSuFo4tPqVL2Jbtly5Ym33rrrSbH+veRje8B0R100EEmDxs2LGKbSpUqRd2H+8EwYcaVVAAAAHiHIhUAAADeoUgFAACAd0LVk9q2bVuTb7/99qjbr1mzxuRLLrnE5N27d6dmYCjyhg4dGnHfXXfdFfU5vXr1MjnMCyr7Ys6cOSa3b9/e5Pfeey9zg0mTOnXqmLxgwQKTq1SpEvX5TzzxhMlLlixJzcBCzu1B/ec//5nQ8z///HOTly9fbvL48eNNXrp0qckrVqxI6HgFue+++0y+5ZZbTFZVk92e1HXr1pn87LPPJj0mROf20bu2bt1qcr9+/Ux265aCfPvttyYXpffbcCUVAAAA3qFIBQAAgHcoUgEAAOAdr3tS3V6OcePGmVyuXDmT3f4btyeQHlTEq2vXribH6j8VEXnxxRdNnjFjRkrHBJH169dHfdztUXX7OX1Uu3Ztk92+21g9qK7JkycnPaaiqEWLFkk9v3nz5ia76yL37NnT5F9//dXkxYsXmzxlypSIY8yaNctkd91T9xixbNq0yeSrr77a5G3btiW0P4g0atTI5DFjxkTd/g9/+EPUx3/66SeTW7dubXI8a0F369bN5LVr18Z8TlhwJRUAAADeoUgFAACAdyhSAQAA4B2ve1Ivuugik+vXrx91+5UrV5p87LHHmjxv3rzUDAxFzoUXXmjyv/71r5jPcdfhvfjii1M5JBQRpUuXNvmyyy4zedSoUSaXKVMmof3fdtttJn/44YcJPb+oOuSQQ0x2e/1c7nqVr732WtT9denSJer+KlSoYHKrVq1ijiee/sNEDBo0yGT65BN3xx13mPyPf/zD5Fhr08Zy5JFHRt1fPNwxunP3448/NvnHH39M+BjZwpVUAAAAeIciFQAAAN6hSAUAAIB3NFr/hKqmtkEmQaNHjzbZ/ezlEiVsjX3gwIGo+3M/w3rChAkR27zxxhtR9/HNN9+YvG/fvqjbJ+qwww4z2V1DLdG1XoMgSLzBpRCyPVcSdcwxx5j8zjvvmOz2n7n9pyIinTt3Ntn9rO6wycRcSXaeuGuGur+XLVu2mOz2ar3wwgvJHD6C+++1U6dOEdu4a1ueeeaZCR1j+/btJj/11FMmu32He/fuTWj/iQrLOcXt+fzggw+ibu/2Cj/33HNRt3fX6XbXpz3jjDOiPr+g3sNke1Ldfe7cudPk4cOHm/zMM8+YnOr1NcNwTonFfQ0+6KCD3OOb/PTTT5vsrt3csGHDqMdLtse1IG4P6sMPP2zygw8+mPQxkhFtnnAlFQAAAN6hSAUAAIB3KFIBAADgHa97UqtVq2bysGHDTHb7v4444oikjxmrH+STTz4xOScnx2S3pzXW/sqXL2/yAw88EHX/J510UuSgowhL/1i6uT9ndy3J4447zuS5c+eafOWVV0bs012XN+zC2D/m9lLdcMMNUbdfvny5ye56grH06NHD5EMPPdTk6tWrJ7S/grh977fccovJ06dPT/oYyQjLOcX9N//ll1+afPjhh5vsrnvqrqtdqpRdVvy8884z2V3X210X1ZWJntRY+3v//fdN7tChQ1LHd4XxnOJasGCBye5rhfszf+KJJ0zu1q2bye45w/X222+bPHbs2Iht3HW93fdHuGstu2NcuHChySeccELUMaUbPakAAAAIFYpUAAAAeIciFQAAAN7xuic1lkqVKpncq1cvk90e1WuuucbkqlWrRuwz1WuUpXp/bl9ULGHpH0s3d03c888/3+T9+/eb7PY2umv2FkVFoX/soYceMrlfv34mly5dOp2HL5C7frO7FqU75kcffTTq87MtrOeUmTNnmnzaaaeZ7K5Hu2fPHpPd9TGTPZdv3rw54r7HH3/cZLd30P2c92+//dZkd21Yd/vZs2eb/NZbb5m8YsWKKCNOXFE4p7h96BMnTnSPb3Ki82LWrFkmd+3a1eRdu3bF3EeDBg1MHjdunMktWrQwedKkSSa7awRnGj2pAAAACBWKVAAAAHiHIhUAAADeCXVPaqIOPvhgk9u1axexzSmnnGKy+xnvbu9H/fr1ox4z2X4Vt7fkqquuSuj5Ye0fS5a7nuybb75psrsG76hRo0zu379/OobltaLQP+Zy130cM2aMyXXr1jU50Z5vt49x2rRpEdu4/dDuWsphE9ZzStu2bU2eOnWqye57HAoYj8mxzuVuT6vb/9mnT5+I5/z8889R9xk2ReGc4q63+9xzz5ns9qy68+K7774zefjw4SaPHz/e5F9++aVQ44zGrWOWLFmS8mMkg55UAAAAhApFKgAAALxDkQoAAADvFKue1FRw+1rd7Orbt29C+1+3bp3JDz/8sMm7d+9OaH9h7R9LlNsr7H7+ubs+5kcffWTy2WefbXJR6w2LR1HoH0uU27tcs2bNhJ7vrju5c+fOpMfku6JyTjnjjDNMvuKKK0w+77zz3PGY7K41+dprr5m8bNkykxcvXlyYYYZacTynIHH0pAIAACBUKFIBAADgHYpUAAAAeIee1CKuqPSPuQ499FCTZ8yYYXKzZs1M3rFjh8knn3yyyb6tG5cN9I8hHkX1nILU45yCeNCTCgAAgFChSAUAAIB3KFIBAADgncQ+qBrwhLumoduD6jr22GNNzsnJSfWQAABACnElFQAAAN6hSAUAAIB3KFIBAADgHXpSEQrNmzc3+aabboq6/dixY01et25dqocEAADSiCupAAAA8A5FKgAAALxDkQoAAADvaBD8/sfe8pm44cfnbCNefM424sE5BfHinIJ4RJsnXEkFAACAdyhSAQAA4B2KVAAAAHgnak8qAAAAkA1cSQUAAIB3KFIBAADgHYpUAAAAeIciFQAAAN6hSAUAAIB3KFIBAADgnVAUqaqao6qd4tguUNWGhTxGoZ8LfzBXEA/mCeLFXEE8mCfpEYoi1VeqOkRV96rq9ny3I7I9LvhLVcuo6jJVXZPtscA/qlpNVZ9T1Z/ybkOyPSb4SVXLqupoVf1RVTer6lRVrZPtccEvqtpBVeeo6s+qmpPt8SSKIjV5E4MgqJTvtirbA4LXBorIT9keBLw1UkQqiEgDEWkpIher6mVZHRF8daOInCwiTUWktohsFZF/Z3NA8NIOERkrua89oROqIlVVW6rqJ6q6VVXXq+qjqlrG2ewsVV2lqhtVdYSqlsj3/MtV9WtV3aKqb6lq/Qx/C8gQH+eKqh4uIheJyD+T3RdSw8N5co6IPBAEwa9BEOSIyBgRuTzJfSIFPJwrh4vIW0EQ/BgEwS4RmSAiTZLcJ5Lk2zwJgmB+EATPi0goL6CFqkgVkf0icpOI1JTcvyA7ish1zjbdRaSFiBwvIl0l7wSvqt1E5HYROVdEDhaRD0RkfEEHUdXb8iZYgTdn83Py/qtlqar2TcU3iZTwca78O2+/O5P/9pAiPs4Tdb4+pvDfHlLIt7kyRkRaq2ptVa0gIr1FZEZKvlMkw7d5Em5BEHh/E5EcEelUwP39RWRyvhyIyJn58nUiMjvv6xkickW+x0qIyK8iUj/fcxsmOK7GkvvfLCVFpJWIrBeRC7P98yrON4/nSncRmZn3dXsRWZPtn1Vxvnk8T14QkddEpLKINBSRlSKyO9s/r+J883iuVJHcAiYQkX0i8l8RqZHtn1dxvfk6T/Ltq5OI5GT755ToLVRXUlW1kapOU9UfVPUXERkmuX+t5Lc639ffSW4RKSJSX0QeyfdXxmbJvUpR6EbzIAi+CoJgXRAE+4Mg+FhEHhGRHoXdH1LHp7miqhVF5AERuaEwz0f6+DRP8vxNcq+0fysir0tuEcKb7Dzg4Vx5QkTKichBIlJRcv+44Upqlnk4T0ItVEWq5P6jXCYiRwZBUEVyL4urs81h+b6uJyLr8r5eLSLXBEFQLd+tfF5xaajq7WrfsW9uUcYXFDAeZIdPc+VIyX0jzAeq+oPkvpjUyjuJNUjVN4xC8WmeSBAEm4Mg6B0EwaFBEDSR3HP0/BR+vyg8r+aKiDQTkWfz5sxuyW0naqmqbkGEzPJtnoRa2IrUyiLyi4hsV9WjRKSgHtCBqlpdVQ+T3Hc/Tsy7f7SIDFLVJiIiqlpVVXsWdJAgCIYF9h375vbbdqraNe9YqqotJfcqyOup+3aRBJ/myhLJPSk1z7tdKSI/5n29uoDdInN8mieiqn9U1YNUtaSqdhaRq0Xk3tR9u0iCV3NFRD4TkUvy9lVacv/beF0QBBtT8+2ikLyaJ6paQlXLiUjp3KjlNPKNXN4KW5F6i4j8VUS2icjT8v+/2PxeF5GFIrJYRKZLbnO5BEEwWUSGi8iEvEvwS0Skc5Lj6SUiK/LG8x8RGR4EwXNJ7hOp4c1cCYJgXxAEP/x2k9z/wjmQl/cXdr9ICW/mSZ4/i8iXeeP5p4j0DoJgaZL7RGr4NlduEZFdktsaskFEzpLc3ndkl2/zpK3kthC9KblXbXeKyNtJ7jNjNMhtqAUAAAC8EbYrqQAAACgGKFIBAADgHYpUAAAAeIciFQAAAN4pFe1BVeVdVSEXBEFG1m1lroRfJuYK8yT8OKcgXpxTEI9o84QrqQAAAPAORSoAAAC8Q5EKAAAA71CkAgAAwDsUqQAAAPAORSoAAAC8Q5EKAAAA71CkAgAAwDsUqQAAAPAORSoAAAC8Q5EKAAAA71CkAgAAwDsUqQAAAPAORSoAAAC8Q5EKAAAA71CkAgAAwDsUqQAAAPAORSoAAAC8Q5EKAAAA75TK9gCiGTdunMknn3yyyR9++KHJS5YsMfmzzz4zOScnx+R9+/ZFHLNVq1Ymn3rqqSYPHjzY5M2bN0fsAwBQtFWqVMnkOnXqmDxp0iSTmzRpYvLixYsj9vn+++9H3WbevHkmL1u2LJ6hophr0aKFyW5tdODAAZP/+te/mjxx4sT0DCwOXEkFAACAdyhSAQAA4B2KVAAAAHhHgyD4/QdVf//BNGjTpo3Js2bNMrl06dImq6rJ0b4XEZENGzaYvH///ohtatWqFXWf3bt3N/mNN96IesxsC4JAY2+VvEzPlTBasWKFyQ888IDJTz31VCaHEyETc6WozZNy5cqZ3Lp164ht3PPa4YcfbvJZZ51l8tKlS012+xLvvvtuk7dt2xbXWFOFc0qul156yeQLLrgg6X3Gek37/vvvTR41apTJI0eOTHoMqcQ5JTvatWtn8tixY01u0KCByW5PqvtadfTRR6ducAWINk+4kgoAAADvUKQCAADAOxSpAAAA8I5XPamu+++/3+SBAwea7K5R+sEHH5js9pcecsghJrv9PSIiW7ZsMfmLL74w+YknnjD5hx9+iNiHT+gfy3XppZeavHbtWpPfeeedlB/TXZvu008/NfmFF14wuU+fPikfQyLoH4vkrn15xx13mNy1a1eT3XNOOsycOdPk888/3+Tt27en9fjF5ZxSvnx5k59//nmTzzjjDJMrVKiQ9DETfZ+F+76K9957z+SLLrrI5J9++qnwgysEzinpcfDBB5vcuHFjkydMmGByzZo1Td65c6fJ7uuh+3rpvnalGj2pAAAACBWKVAAAAHiHIhUAAADeKZXtAUQzefJkk92eVPdzi88991yTy5QpEzXv3r074ph79+5NeJzwT8+ePU121yA99dRT0z6GE044wWS336ygnmhkVokS9u/0k08+2WR3HeTq1atH3V9B5xT389hXrlyZyBAj1t90e50rVqxocrp7UouLKlWqmOyukZ2sTZs2Rdznzh93/rmvcX/4wx9M7tixo8m9e/c22bd1VFE4bdu2NfnJJ580uWrVqlGf/91335ns9i5//vnnSYwutbiSCgAAAO9QpAIAAMA7FKkAAADwjtc9qccee2xSz9+zZ0/UjKLDXaPQ/XzzHTt2mJyTk5PuIUmPHj1MXr9+vclunywyz12L+ZZbbkno+dOmTTP5yiuvjNgm2bUpBw0aZLLbqz9//nyTO3ToYPKqVauSOn5xdeaZZ0Z9/OOPPzb5scceM7l58+YmL1682GR37ohErsv7zTffmHz99deb/O2335p8xBFHmOyu5froo4+azHswwsntRS5dunRCz3fXVW3Tpo3J9KQCAAAAUVCkAgAAwDsUqQAAAPCO1z2py5cvN9ldZ9JVt25dk+vXr2+yu25l2bJlI/YxY8YMk7/44ouY40T23XjjjSY3adLEZHcduTVr1qR8DC1btjT5lFNOMdld13f16tUpHwOsQw45xORRo0aZfN5550V9/tatW03+y1/+YvInn3xisvtZ6oXhrnF43XXXmdy+ffuoz3/mmWdMzsSawEVRo0aNoj7urnPsfl66m+Ph9qDG8vLLL5t82223JXxM+O2uu+6KuG/IkCFJ7fO+++4z2e2n9glXUgEAAOAdilQAAAB4hyIVAAAA3vG6J/Wss84yOQgCk9015WbPnm1yw4YNEz6mu77mww8/bLK7RqG79t2BAwcSPiaS5/YebtmyxeR///vfaR9DuXLlTC5Vyut/XsVCr169TO7Zs2fU7X/88UeTTzrpJJPdz7xOhWbNmpnsfr56rB5U9/Pe//Wvf6ViWIjB/fddsWJFk921meNRooS9buT2xd5zzz0mx+qt//nnn03et29fwmNCZvXt29fkgvpPY9UZc+fONdmtW3zuQXVxJRUAAADeoUgFAACAdyhSAQAA4B2vm+Y6deoU9fEGDRqY7Pasrlu3zuRZs2aZvGTJkoh9uusg3n777VGzuy7diBEjfn/ASJkyZcqY3LVrV5Nfeuklk7/66qu0j+m0006L+rjbw+b20br9kEjeBRdcEPVxdx3UHj16mJyOHtRLLrnE5AEDBph87LHHJrS//v37m/zGG28Ualyw3NcL91zvzpVu3bqZ3KdPH5PdPsGCXt/ctb7vvffeuMb6m19//dVkd21xt29+586dCe0fqVetWjWTY/XNx8OtjebNm5f0PrOFK6kAAADwDkUqAAAAvEORCgAAAO943ZNar169hLZ3+3cefPBBk7dt2xZzH+5ne3fs2NHkiRMnmux+Bq67buo777wT85hIXNOmTU2uX7++yW6vYTqULl3aZHdNTddRRx1l8owZM0w+/vjjUzOwYqx8+fImV6pUKer2OTk5Jn/00UcJHa958+Ymu73RIpG9i+7al+48imXTpk0mP/PMMwk9H/H573//a7L7HofatWub7Pacv/jiiya7a5q6/acikT2k7vssYpkyZYrJF198cULPR/q5Pelnn322yaecckrC+3Rf7wYOHGjywoULE96nL7iSCgAAAO9QpAIAAMA7FKkAAADwjtc9qdOnTzfZ7eW4/PLLTXb7RQuzBtzevXtNnjlzpsnnn3++ya+99prJzz33nMnu524vX7484TEhktur5X6W8TnnnGPyo48+anKsNUnLli1rcps2bSK2ueaaa0x2+5fdMbm/+7Zt20YdAxK3f/9+k2N9VvnRRx9t8jfffJPQ8dxeaHf93nR44YUXTHa/Z6SG2+fXpUsXk93XBnfdY1dBPajJmj17tsnu577DP7t37za5e/fuSe9z48aNJk+ePDnpffqCK6kAAADwDkUqAAAAvEORCgAAAO9QpAIAAMA7Gm2xYFVNbCXhFKtWrZrJboPxuHHjMjiago0ePdrkq6++2mT3jVXu4sqFeXNXIoIg0NhbJS/bc8V9A0GHDh1MXrFihcmvv/66ye5ixzfffLPJLVq0SHhMkyZNMtl9051vMjFXMj1PhgwZYvLdd9+dycOLSPILtLtz1/3Qh+3btxduYIVUXM4psfzlL38x2T3XlyiR+DUg9wNnKleuHHV79/XDXRh+zpw5CY8hlYriOSVRjRs3Nvmtt94y2f1QCFdB82jJkiUmn3766SavX78+kSFmXbR5wpVUAAAAeIciFQAAAN6hSAUAAIB3vO5JDQN3AXe3T9ZdwLl58+Ymf/HFF2kZ12+KS/9Y06ZNTX7yySdNPvHEE6M+3+0b/PTTT00uaHHkM8880+R27dqZPGDAAJNHjhwZdQzZVhT7x0qVsp9Xcuutt5rs9i67/Z6LFi0yecOGDSa7C7oXZM2aNSZPmTLF5IoVK5q8adMmk0877TSTFy9eHPOY6VRczimxuB/G4fYaxvpgh8cffzziPve8dfjhh5s8atQok+vVq2eye95q1apV1DGkW1E8p8Ti9qC6H75x7LHHJrS/ZcuWRdzXu3dvk9NdR6QbPakAAAAIFYpUAAAAeIciFQAAAN4pFXsTRPPtt9+a/MEHH5h84YUXRs1h7yXxhftzvPbaa02+4YYbTF69erXJq1atMvmll14yef/+/RHHdHtQXfPnz4/6ONJv3759Jt93331R88EHH2yy24NaGO76mW4Pqsvtc812DyoKVqdOHZNj9aA++uijJg8cODBimz179pjsrof59ddfm/z555+b7PY7uq8348ePjzpGJK9+/fomJ9qDum7dOpPdNVBFwrcOajK4kgoAAADvUKQCAADAOxSpAAAA8E6oe1Jr1aplsru+oNvfkw7ff/991IzscHu1rrzyyqT216xZs4j7OnXqZPJHH31k8ieffJLUMZF5yfag3nTTTRH3devWLepz3H7q/v37JzUGZMZ1110X9fFnnnnG5JtvvtnkgvrcY1mxYoXJbp/rLbfcYvLZZ59tMj2p6XfOOecktL37WnXRRReZXJz6TwvClVQAAAB4hyIVAAAA3qFIBQAAgHdC1ZPaqFEjk+fOnWvy2rVrTXZ7uz788MO0jCs/d4woGtzP0BaJ/Fz4BQsWmHzgwIG0jgnZV758eZP79OmT8D7++c9/muz21sNPGzdujPq4u05yYXpQY3HXAXa5r0cVKlQw+ddff035mIq7a665xuRYrwPvv/++ycuWLUv5mMKMK6kAAADwDkUqAAAAvEORCgAAAO+Eqif1sssuM/nQQw81ee/evZkcToGWL19usqqa/Kc//SmTw0EGub97FH133nmnyU2bNo35nAkTJpj86quvpnRMyIyXX37Z5L/85S8mZ+L9Ce4x3H5md91eelCT97e//c3kkSNHmlyiRPRrfy+99JLJ7vq5sLiSCgAAAO9QpAIAAMA7FKkAAADwTqh6UpcsWWJyEAQmu2sWnnzyySYvXLjQ5J07dyY9ptq1a5t85plnmuyOkf6zcGrdunXMbb766qsMjATZ1LJlS5Pj6Sdz17K87777oj6OcHB/b+65/uKLLzZ53LhxJq9ZsyZin1WqVDHZXXfX/Vz3o446yuRPP/3U5EysDV7cuL/nWOuguo8PGTIk1UMq0riSCgAAAO9QpAIAAMA7FKkAAADwTqh6Ul988UWT27dvb/Lll19u8v3332/yJZdcYvKoUaMijvH0009HHUOtWrWi7sNdJ/Hrr782+fXXX4+6f/jpo48+irhvwIABJrvr9qLoef75500uW7Zsws9ZunRpSseE7HjllVdM7tatm8m9evUyOZ7fu7vGZqx+R6RfnTp1TL766qujbr9gwQKT+/bta/L333+fmoEVE1xJBQAAgHcoUgEAAOAdilQAAAB4J1Q9qa7rrrvOZHcN0qeeesrkxo0bmzx69OiIfQ4bNsxkd020MmXKmFy5cmWTt27darK7zt327dsjjomiwV1Dc+LEiVkaCVLl2muvNblhw4YJ7+Oaa65J1XDgscGDB5vcqlUrk+vVqxdzH+7rTaJWrlyZ1PMRae3atSa7dcXDDz9scs2aNU2uUKGCyXv37k3h6Io+rqQCAADAOxSpAAAA8A5FKgAAALyj0XpgVDW5Bpksq1GjhslDhw41+bzzzot4jrvWZaweoc8++8xk97O8P/7445jjTKcgCDQTxwn7XInl+OOPj7jP/d3OmzfPZHcdX99lYq74Pk+qV69usrumYcWKFaM+/4UXXoi4z12fOew4p8TnqKOOMvmtt94yuW7duhHPUbU/Wnf+TZ482WT3HDRjxgyTs/0eiKJwTnnmmWdMPvHEE01+6aWXTH777bdNXrhwYXoGVoREmydcSQUAAIB3KFIBAADgHYpUAAAAeKdI96SC/rF0GjJkiMnr1q0z2V1Pz3dFoX8sWaVLlzZ57NixJvfu3dvke+65x+RHHnkkYp9btmxJ0ej8wDkF8eKcgnjQkwoAAIBQoUgFAACAdyhSAQAA4B16Uos4+scQL/rHEA/OKYgX5xTEg55UAAAAhApFKgAAALxDkQoAAADvUKQCAADAOxSpAAAA8A5FKgAAALxDkQoAAADvRF0nFQAAAMgGrqQCAADAOxSpAAAA8A5FKgAAALxDkQoAAADvUKQCAADAOxSpAAAA8E4oilRVzVHVTnFsF6hqw0Ieo9DPhT+YK4gH8wTxYq4gHsyT9AhFkeorzTVcVTfl3R5QVc32uOAfVR2oqktUdZuq/k9VB2Z7TPCPqnZQ1Tmq+rOq5mR7PPCXqg5R1b2quj3f7Yhsjwt+Cfs8oUhNztUi0k1EmolIUxE5W0SuyeaA4C0VkUtEpLqInCki/VS1V3aHBA/tEJGxIsIfMYjHxCAIKuW7rcr2gOCl0M6TUBWpqtpSVT9R1a2qul5VH1XVMs5mZ6nqKlXdqKojVLVEvudfrqpfq+oWVX1LVesnOaQ+IvJQEARrgiBYKyIPicilSe4TKeDbXAmC4IEgCBYFQbAvCIJvROR1EWmdzD6RPA/nyfwgCJ4XkdC8iBQXvs0V+Il5klqhKlJFZL+I3CQiNUXkZBHpKCLXOdt0F5EWInK8iHQVkctFRFS1m4jcLiLnisjBIvKBiIwv6CCqelveBCvwlm/TJiLyeb78ed59yD7f5kr+56iInCIiS5P6DpEK3s4TeMfHuXKOqm5W1aWq2jcV3ySSxjxJpSAIvL+JSI6IdCrg/v4iMjlfDkTkzHz5OhGZnff1DBG5It9jJUTkVxGpn++5DRMc134ROSpfPjJvP5rtn1lxvfk6V5yxDJXcP2jKZvvnVVxvvs8TEekkIjnZ/jlx83euiEhjEaktIiVFpJWIrBeRC7P98yquN+ZJem6hupKqqo1UdZqq/qCqv4jIMMn9ayW/1fm+/k5yfzkiIvVF5JF8f2Vsltw+wTpJDGm7iFTJl6uIyPYgb2YgezycK7+Nq5/k9qZ2CYJgd7L7Q3J8nSfwj29zJQiCr4IgWBcEwf4gCD4WkUdEpEdh94fUYJ6kVqiKVBF5QkSWiciRQRBUkdzL4u676Q/L93U9EVmX9/VqEbkmCIJq+W7l835phqrervadcOaWb9Olkvumqd80E/4L1xe+zRVR1ctF5DYR6RgEwZoUfZ9IjnfzBN7yfa4EBYwHmcc8SaGwFamVReQXEdmuqkeJSEG9FQNVtbqqHiYiN4rIxLz7R4vIIFVtIiKiqlVVtWdBBwmCYFhg3wlnbvk2/Y+I3KyqdVS1togMEJFnU/KdIllezRVV7S25f1GfFoTonZXFgG/zpISqlhOR0rlRy2nkmy6QHb7Nla55x1JVbSkif5PcN2Qiu5gnKRS2IvUWEfmriGwTkafl/3+x+b0uIgtFZLGITBeRMSIiQRBMFpHhIjIh7xL8EhHpnOR4nhSRqSLyZd7+pufdh+zzba7cKyIHichn+f7aHZ3kPpE83+ZJWxHZKSJvSu4Vlp0i8naS+0Rq+DZXeonIirzx/EdEhgdB8FyS+0TymCcppLRPAgAAwDdhu5IKAACAYoAiFQAAAN6hSAUAAIB3KFIBAADgnVLRHlRV3lUVckEQZGQ9NOZK+GVirjBPwo9zCuLFOQXxiDZPuJIKAAAA71CkAgAAwDsUqQAAAPAORSoAAAC8Q5EKAAAA71CkAgAAwDsUqQAAAPAORSoAAAC8Q5EKAAAA71CkAgAAwDsUqQAAAPBOqWwPAAAAoCgqUcJeCyxdurTJBw4cMLlcuXJRny8iUrJkyaTGtG3bNpP37t2b1P7SiSupAAAA8A5FKgAAALxDkQoAAADv0JOKYqlChQomn3baaSa3bds24X0uWrTI5Dlz5pi8bt26hPcJuHP1zjvvNPlPf/qTyeedd17axwQgPpdccknU7PaHNm3a1OQqVapE7LNatWpJjWnmzJkmu69V//vf/6Juv2PHjqSOnwiupAIAAMA7FKkAAADwDkUqAAAAvKNBEPz+g6q//2Ax1b59e5MHDx4c9fEOHTqY/N5776VhVL8vCALNxHF8nyuVK1c2edy4cSZ369bNZFX7Y4v27+T3bNiwweRjjjnG5E2bNiW8z3TKxFzxfZ74wJ0nEydONHnChAkmjxo1yuSff/45PQPLwzmlcNy1LUuVinxLSKtWrUyuX79+1H3WqlXL5PXr10fd3l2js1+/fibXrFnT5I0bN5p8wQUXmLxs2bKoxyuO5xR3XdPnn3/e5F69emVyOCkxbdo0k7t27ZrS/UebJ1xJBQAAgHcoUgEAAOAdilQAAAB4h3VSHYn2nMbirj+W7R7V4qpJkyYmuz2orldeecXkgtY4ffnll00+7rjjTH700Uej7rNnz54m+9ajisxwe1Dfeustk59++mmTH3vsMZPT3YOK+Li9iO65/rbbbjO5Xr16Efto0KCByW4PaWF64/NLtNf+0EMPNfnYY481OVZPanHUqFEjk8PYg+ravn171o7NlVQAAAB4hyIVAAAA3qFIBQAAgHeKVU/qkCFDTHb7TTPBPSY9qZnxww8/mPz++++b7Pb9jR8/PuFjzJs3z+QDBw6Y7PYSnn/++SY/8cQTCR8T4eP2HU6fPt3kkSNHmvzwww+b7M4rZIfbU/rnP//Z5NGjR5vsrkGaCV9++aXJ7txxe1Ld91B89NFHJr/xxhspHF3R5L7XIBPctZT3799vsnuOqVKlStT9ffrppyb/9NNPSYwuOVxJBQAAgHcoUgEAAOAdilQAAAB4p0j3pLr9NYmucZoJ7pjoUU2PnJwck0899dS0H/PNN9802V2jsHPnzia7fbH79u1Lz8CQVSNGjDD5+++/N/nBBx/M5HAQp4YNG5o8aNAgky+77DKTY61B6v57FxH5+uuvTXZ7A9evXx9znPmtXr3aZPqZ/eP2i7777rsmv/rqqzH3sWbNGpOTXU/XJ1xJBQAAgHcoUgEAAOAdilQAAAB4p0j1pLr9nZnoQR06dGjUx9u1a2dyrDG6+3PXdkV4uX1C1atXN9n9nG56UouGK664wuRmzZqZfMIJJ2RyOCik8847z+RLL7006vYrV640edKkSSbPnDkz4jnu+s0In5tuuinq42PHjjW5b9++JnPet7iSCgAAAO9QpAIAAMA7FKkAAADwTqh7Ut1+zcGDByf0/IL6Sd19utldxzTRdU2L0vpliK5Xr15RH1+7dq3JO3fuTOdwkAFdu3aNuO+ee+4x+d577zX5559/TuuYUDglS5Y0uXv37ia76x5PmTLF5HPPPTct44Lf3PcWuOrWrWvySSedZPLSpUtN3rJlS2oGFlJcSQUAAIB3KFIBAADgHYpUAAAAeEej9UiqqlcNlMn2oHbo0MHkRPtJUyHR78Hte0pUEATJ7SBOvs0VH7i9hhUrVjTZXTcxVg9rumVirhT1ebJgwYKI+2rUqGHyEUcckanhpEVxOadUqlTJ5Llz55rcvHlzk/v162fyE088kZZxhUlxPKdMmDDB5J49eyb0/DVr1pg8Z86cmM/58MMPTXZrmxUrViQ0hkyLNk+4kgoAAADvUKQCAADAOxSpAAAA8E6o1klt165dQtsnu6ZpOrg9qe731L59e5PdfhS3rxbZU6FCBZMnTpxocuXKlU12+7/dHjeEz913321ys2bNIrZxexURDtu3bzd58eLFJrs9qaNGjTL51FNPNfnVV181uaD+Zd97BxHbzJkzTW7durXJtWvXjvp8dx3Viy++OOYx3W1+/fVXk999912Tn3nmGZOnTp0a8xjZwpVUAAAAeIciFQAAAN6hSAUAAIB3vF4nNdHPufdhHdRYYvWcuoYOHWqy29MaS3FZ09Dt86lXr57JXbp0ifr8H374weSFCxfGPOYtt9xicrdu3Ux217h15+P5559v8qZNm2IeM52K45qGiSpXrpzJH330kcnVq1ePeI7bu/jLL7+kfFyZVFzOKa7OnTubPG3aNJPdf++xXr927doVcZ/b9zp//nyTH3vsMZN972HlnCLSoEEDk6+44gqTjz32WJPbtGljckHnlGS5a3i/9tprJt94440m79ixI+VjyI91UgEAABAqFKkAAADwDkUqAAAAvOPVOqmJ9lu6fOxBTVZR/J7i4fbxtW3b1uQLLrjA5Fq1aplcv359kxPtb060vywes2bNMjnbPahI3L333mvycccdZ/Ktt94a8Zyw96Ai1zvvvGNykyZNTJ49e7bJhx56aNT9uf3NIiInnXSSySeffLLJXbt2Ndmdf26vIbIvJyfH5Lvuuivq9m4P6wknnBCxjdsf7c6Dpk2bRj1G1apVTb7sssuiPn7RRReZvHv37qj7TyWupAIAAMA7FKkAAADwDkUqAAAAvEORCgAAAO9kdTH/RBe2d4Vh8X6X+z26PwNXst9jWBfedt8AULFixUTHY/K6detMnjBhgslXXXWVyZUrVza5MG+ccsfgNpu7C8Gfe+65Jm/bti3hYyaDhbcjlS5d2uQlS5aY7H5oxNFHHx2xD/eNE2EX1nNKprmLsru5UaNGEc85/vjjTW7WrJnJ7nno1VdfNblPnz4m//rrr/ENNk04p2RGpUqVTHbftHf33Xeb3Lt374T2X7NmTZO3bNmS0PNjYTF/AAAAhApFKgAAALxDkQoAAADvZHQx/+LQg5rs9+j2MRZXVapUMfnAgQMmu71W7kL57qLrbj/osGHDTHZ7UN3fw549eyLG+OGHH5r85ptvRmyTX5cuXUx2P7Bg69atJrs9qu5i4tnuNysOnnjiCZMbNmxo8t/+9jeTi1r/KQrPPT+4OR6nn366yc8//7zJ7utNhQoVTOYcUTy4r0+rV682+frrrze5du3aJru1lU+4kgoAAADvUKQCAADAOxSpAAAA8E5We1JjGTp0qMlFsQfV516QbHJ7UN31AadPn27yiBEjTL711ltNPvHEE02uU6dO1P2762H2798/YoyJ/q5Hjhxpsvu7f+aZZ0x+7bXXoubLL7/c5Eyvq1oUVatWzWT3Z7xy5UqTn3766XQPCcWY+5q3adMmkw8++OAMjqZocN/vsHPnTpNr1aqV0uO5fcJuf2gquOvruuumHnPMMSk/ZqZwJRUAAADeoUgFAACAdyhSAQAA4J2M9qS2a9cuoe196EF1e04HDx4c9XGX21c7ZMiQFIyq6Pv2229NdtenPPXUU03+y1/+YnLZsmVNdntOXW6P67XXXmvy+vXroz6/MNyeVvezvdesWWNy9+7do+7Pt8/tDqNYP+P777/f5ILWzwVSxf03/ac//cnkzZs3Z3I4RcIDDzxgsrtmqLueNbKLK6kAAADwDkUqAAAAvEORCgAAAO94vU5qrO1j9ay6z4/VX1oY7hhY9zQ1OnXqZPLw4cNNvuCCCxLan9tz+o9//MPk//73vybv378/of2ngtv32qRJE5MnT55ssts/6fbhun26iOSuiThs2DCTVdXkwnz+OvzUrVs3k93f7caNGzM4mlxXXnmlyU8++aTJbm/96NGjTc7GmMPmmmuuMdldkxsi69atM3nfvn1ZGglXUgEAAOAhilQAAAB4hyIVAAAA3tFo60eqavTFJZMUa+1KH7k9pz6s5RpNEAQae6vkpXuuuBo1amSyu7bdyJEjMzmcjOjXr5/JjzzySNTtS5YsmdD+MzFXMj1PYjnqqKNMXrp0qckff/yxyR07djS5OK6TWlTOKU2bNjXZXT/ztttuM3nx4sVJHa9y5com9+rVK2IbtwfV7Yl210U97rjjTP7++++TGWLK+XhOGTRokMn33ntvSscTBm7/9fvvv2/ymDFjTM7JyUnreKLNE66kAgAAwDsUqQAAAPAORSoAAAC8k9F1Ul1uP2ei66gme7y5c+fG3Mb3ntPiavny5VFzUeT2cIexpztsxo4da3Jx7EEtLk477TSTjzzySJOnTZtmstuv/Omnn5p8+umnmzxixAiTK1WqFDEG99/09u3bTb7kkktM9q0HNQzc3uOXX37ZZHf9XPf9D7FceumlJpcqlf4ya8GCBSa7/dNTp041+d133zX5119/Tcu4UoErqQAAAPAORSoAAAC8Q5EKAAAA72R1ndRYhgwZktbti4OisqYhItdZdPslu3fvbnKivVA+rmmYbrHWSX3rrbdM7tOnj8kbNmxIz8A8VlTOKTVq1DB55syZJv/5z39OaH/umqaF6Rl3e0zdtVTdvlffFcdzSrVq1Ux250U67Nq1y+SdO3em/ZipxDqpAAAACBWKVAAAAHiHIhUAAADe8bonFckrKv1jiFShQgWT3R67NWvWJLS/4tg/duihh5o8ceJEk8eNGxf18bD1fqVCUT2nuOukTpo0yeSC1jXNL9GeVHfdVZHINTa3bNkSdR++K47nFCSOnlQAAACECkUqAAAAvEORCgAAAO/Qk1rEFdX+MaQe/WOIR3E5p9StW9fkCy64wOR+/fqZ7PakLly40OQ33njD5Oeffz7imAcOHEh4nD7jnIJ40JMKAACAUKFIBQAAgHcoUgEAAOAdelKLuOLSP4bk0T+GeHBOQbw4pyAe9KQCAAAgVChSAQAA4B2KVAAAAHiHIhUAAADeoUgFAACAdyhSAQAA4B2KVAAAAHiHIhUAAADeoUgFAACAdyhSAQAA4B2KVAAAAHhHg4CPvQUAAIBfuJIKAAAA71CkAgAAwDsUqQAAAPAORSoAAAC8Q5EKAAAA71CkAgAAwDv/Bx6z9iwB7yj7AAAAAElFTkSuQmCC\",\n \"text/plain\": [\n \"
\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"fig, axs = plt.subplots(5, 5, figsize=(12, 12))\\n\",\n \"for i, ax in enumerate(axs.flatten()):\\n\",\n \" ax.imshow(test_batch['image'][i, ..., 0], cmap='gray')\\n\",\n \" ax.set_title(f\\\"label={pred[i]}\\\")\\n\",\n \" ax.axis('off')\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"edb528b6\",\n \"metadata\": {},\n \"source\": [\n \"Congratulations! You made it to the end of the annotated MNIST example. You can revisit\\n\",\n \"the same example, but structured differently as a couple of Python modules, test\\n\",\n \"modules, config files, another Colab, and documentation in Flax's Git repo:\\n\",\n \"\\n\",\n \"[https://github.com/google/flax/tree/main/examples/mnist](https://github.com/google/flax/tree/main/examples/mnist)\"\n ]\n }\n ],\n \"metadata\": {\n \"jupytext\": {\n \"formats\": \"ipynb,md:myst\",\n \"main_language\": \"python\"\n },\n \"language_info\": {\n \"name\": \"python\",\n \"version\": \"3.9.6\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 5\n}\n" }, { "path": "docs/quick_start.md", "content": "---\njupytext:\n formats: ipynb,md:myst\n main_language: python\n text_representation:\n extension: .md\n format_name: myst\n format_version: 0.13\n jupytext_version: 1.13.8\n---\n\n[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google/flax/blob/main/docs/quick_start.ipynb)\n[![Open On GitHub](https://img.shields.io/badge/Open-on%20GitHub-blue?logo=GitHub)](https://github.com/google/flax/blob/main/docs/quick_start.ipynb)\n\n# Quick start\n\nWelcome to Flax!\n\nFlax is an open source Python neural network library built on top of [JAX](https://github.com/jax-ml/jax). This tutorial demonstrates how to construct a simple convolutional neural\nnetwork (CNN) using the [Flax](https://flax.readthedocs.io) Linen API and train\nthe network for image classification on the MNIST dataset.\n\n+++\n\n## 1. Install Flax\n\n```{code-cell}\n:tags: [skip-execution]\n\n!pip install -q flax>=0.7.5\n```\n\n## 2. Loading data\n\nFlax can use any\ndata-loading pipeline and this example demonstrates how to utilize TFDS. Define a function that loads and prepares the MNIST dataset and converts the\nsamples to floating-point numbers.\n\n```{code-cell}\nimport tensorflow_datasets as tfds # TFDS for MNIST\nimport tensorflow as tf # TensorFlow operations\n\ndef get_datasets(num_epochs, batch_size):\n \"\"\"Load MNIST train and test datasets into memory.\"\"\"\n train_ds = tfds.load('mnist', split='train')\n test_ds = tfds.load('mnist', split='test')\n\n train_ds = train_ds.map(lambda sample: {'image': tf.cast(sample['image'],\n tf.float32) / 255.,\n 'label': sample['label']}) # normalize train set\n test_ds = test_ds.map(lambda sample: {'image': tf.cast(sample['image'],\n tf.float32) / 255.,\n 'label': sample['label']}) # normalize test set\n\n train_ds = train_ds.repeat(num_epochs).shuffle(1024) # create shuffled dataset by allocating a buffer size of 1024 to randomly draw elements from\n train_ds = train_ds.batch(batch_size, drop_remainder=True).prefetch(1) # group into batches of batch_size and skip incomplete batch, prefetch the next sample to improve latency\n test_ds = test_ds.shuffle(1024) # create shuffled dataset by allocating a buffer size of 1024 to randomly draw elements from\n test_ds = test_ds.batch(batch_size, drop_remainder=True).prefetch(1) # group into batches of batch_size and skip incomplete batch, prefetch the next sample to improve latency\n\n return train_ds, test_ds\n```\n\n## 3. Define network\n\nCreate a convolutional neural network with the Linen API by subclassing\n[Flax Module](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html).\nBecause the architecture in this example is relatively simple—you're just\nstacking layers—you can define the inlined submodules directly within the\n`__call__` method and wrap it with the\n[`@compact`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/decorators.html#flax.linen.compact)\ndecorator. To learn more about the Flax Linen `@compact` decorator, refer to the [`setup` vs `compact`](https://flax.readthedocs.io/en/latest/guides/setup_or_nncompact.html) guide.\n\n```{code-cell}\nfrom flax import linen as nn # Linen API\n\nclass CNN(nn.Module):\n \"\"\"A simple CNN model.\"\"\"\n\n @nn.compact\n def __call__(self, x):\n x = nn.Conv(features=32, kernel_size=(3, 3))(x)\n x = nn.relu(x)\n x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))\n x = nn.Conv(features=64, kernel_size=(3, 3))(x)\n x = nn.relu(x)\n x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))\n x = x.reshape((x.shape[0], -1)) # flatten\n x = nn.Dense(features=256)(x)\n x = nn.relu(x)\n x = nn.Dense(features=10)(x)\n return x\n```\n\n### View model layers\n\nCreate an instance of the Flax Module and use the [`Module.tabulate`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.tabulate) method to visualize a table of the model layers by passing an RNG key and template image input.\n\n```{code-cell}\n:outputId: 2c580f41-bf5d-40ec-f1cf-ab7f319a84da\n\nimport jax\nimport jax.numpy as jnp # JAX NumPy\n\ncnn = CNN()\nprint(cnn.tabulate(jax.random.key(0), jnp.ones((1, 28, 28, 1)),\n compute_flops=True, compute_vjp_flops=True))\n```\n\n## 4. Create a `TrainState`\n\nA common pattern in Flax is to create a single dataclass that represents the\nentire training state, including step number, parameters, and optimizer state.\n\nBecause this is such a common pattern, Flax provides the class\n[`flax.training.train_state.TrainState`](https://flax.readthedocs.io/en/latest/flax.training.html#train-state)\nthat serves most basic usecases.\n\n```{code-cell}\n:outputId: 1249b7fb-6787-41eb-b34c-61d736300844\n\n!pip install -q clu\n```\n\n```{code-cell}\nfrom clu import metrics\nfrom flax.training import train_state # Useful dataclass to keep train state\nfrom flax import struct # Flax dataclasses\nimport optax # Common loss functions and optimizers\n```\n\nWe will be using the `clu` library for computing metrics. For more information on `clu`, refer to the [repo](https://github.com/google/CommonLoopUtils) and [notebook](https://colab.research.google.com/github/google/CommonLoopUtils/blob/master/clu_synopsis.ipynb#scrollTo=ueom-uBWLbeQ).\n\n```{code-cell}\n@struct.dataclass\nclass Metrics(metrics.Collection):\n accuracy: metrics.Accuracy\n loss: metrics.Average.from_output('loss')\n```\n\nYou can then subclass `train_state.TrainState` so that it also contains metrics. This has the advantage that we only need\nto pass around a single argument to functions like `train_step()` (see below) to calculate the loss, update the parameters and compute the metrics all at once.\n\n```{code-cell}\nclass TrainState(train_state.TrainState):\n metrics: Metrics\n\ndef create_train_state(module, rng, learning_rate, momentum):\n \"\"\"Creates an initial `TrainState`.\"\"\"\n params = module.init(rng, jnp.ones([1, 28, 28, 1]))['params'] # initialize parameters by passing a template image\n tx = optax.sgd(learning_rate, momentum)\n return TrainState.create(\n apply_fn=module.apply, params=params, tx=tx,\n metrics=Metrics.empty())\n```\n\n## 5. Training step\n\nA function that:\n\n- Evaluates the neural network given the parameters and a batch of input images\n with [`TrainState.apply_fn`](https://flax.readthedocs.io/en/latest/api_reference/flax.training.html#flax.training.train_state.TrainState) (which contains the [`Module.apply`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.apply)\n method (forward pass)).\n- Computes the cross entropy loss, using the predefined [`optax.softmax_cross_entropy_with_integer_labels()`](https://optax.readthedocs.io/en/latest/api.html#optax.softmax_cross_entropy_with_integer_labels). Note that this function expects integer labels, so there is no need to convert labels to onehot encoding.\n- Evaluates the gradient of the loss function using\n [`jax.grad`](https://jax.readthedocs.io/en/latest/jax.html#jax.grad).\n- Applies a\n [pytree](https://jax.readthedocs.io/en/latest/pytrees.html#pytrees-and-jax-functions)\n of gradients to the optimizer to update the model's parameters.\n\nUse JAX's [@jit](https://jax.readthedocs.io/en/latest/jax.html#jax.jit)\ndecorator to trace the entire `train_step` function and just-in-time compile\nit with [XLA](https://www.tensorflow.org/xla) into fused device operations\nthat run faster and more efficiently on hardware accelerators.\n\n```{code-cell}\n@jax.jit\ndef train_step(state, batch):\n \"\"\"Train for a single step.\"\"\"\n def loss_fn(params):\n logits = state.apply_fn({'params': params}, batch['image'])\n loss = optax.softmax_cross_entropy_with_integer_labels(\n logits=logits, labels=batch['label']).mean()\n return loss\n grad_fn = jax.grad(loss_fn)\n grads = grad_fn(state.params)\n state = state.apply_gradients(grads=grads)\n return state\n```\n\n## 6. Metric computation\n\nCreate a separate function for loss and accuracy metrics. Loss is calculated using the `optax.softmax_cross_entropy_with_integer_labels` function, while accuracy is calculated using `clu.metrics`.\n\n```{code-cell}\n@jax.jit\ndef compute_metrics(*, state, batch):\n logits = state.apply_fn({'params': state.params}, batch['image'])\n loss = optax.softmax_cross_entropy_with_integer_labels(\n logits=logits, labels=batch['label']).mean()\n metric_updates = state.metrics.single_from_model_output(\n logits=logits, labels=batch['label'], loss=loss)\n metrics = state.metrics.merge(metric_updates)\n state = state.replace(metrics=metrics)\n return state\n```\n\n## 7. Download data\n\n```{code-cell}\nnum_epochs = 10\nbatch_size = 32\n\ntrain_ds, test_ds = get_datasets(num_epochs, batch_size)\n```\n\n## 8. Seed randomness\n\n- Set the TF random seed to ensure dataset shuffling (with `tf.data.Dataset.shuffle`) is reproducible.\n- Get one\n [PRNGKey](https://jax.readthedocs.io/en/latest/_autosummary/jax.random.PRNGKey.html#jax.random.PRNGKey)\n and use it for parameter initialization. (Learn\n more about\n [JAX PRNG design](https://jax.readthedocs.io/en/latest/jax-101/05-random-numbers.html)\n and [PRNG chains](https://flax.readthedocs.io/en/latest/philosophy.html#how-are-parameters-represented-and-how-do-we-handle-general-differentiable-algorithms-that-update-stateful-variables).)\n\n```{code-cell}\ntf.random.set_seed(0)\n```\n\n```{code-cell}\ninit_rng = jax.random.key(0)\n```\n\n## 9. Initialize the `TrainState`\n\nRemember that the function `create_train_state` initializes the model parameters, optimizer and metrics\nand puts them into the training state dataclass that is returned.\n\n```{code-cell}\nlearning_rate = 0.01\nmomentum = 0.9\n```\n\n```{code-cell}\nstate = create_train_state(cnn, init_rng, learning_rate, momentum)\ndel init_rng # Must not be used anymore.\n```\n\n## 10. Train and evaluate\n\nCreate a \"shuffled\" dataset by:\n- Repeating the dataset equal to the number of training epochs\n- Allocating a buffer of size 1024 (containing the first 1024 samples in the dataset) of which to randomly sample batches from\n - Everytime a sample is randomly drawn from the buffer, the next sample in the dataset is loaded into the buffer\n\nDefine a training loop that:\n- Randomly samples batches from the dataset.\n- Runs an optimization step for each training batch.\n- Computes the mean training metrics across each batch in an epoch.\n- Computes the metrics for the test set using the updated parameters.\n- Records the train and test metrics for visualization.\n\nOnce the training and testing is done after 10 epochs, the output should show that your model was able to achieve approximately 99% accuracy.\n\n```{code-cell}\n# since train_ds is replicated num_epochs times in get_datasets(), we divide by num_epochs\nnum_steps_per_epoch = train_ds.cardinality().numpy() // num_epochs\n```\n\n```{code-cell}\nmetrics_history = {'train_loss': [],\n 'train_accuracy': [],\n 'test_loss': [],\n 'test_accuracy': []}\n```\n\n```{code-cell}\n:outputId: 258a2c76-2c8f-4a9e-d48b-dde57c342a87\n\nfor step,batch in enumerate(train_ds.as_numpy_iterator()):\n\n # Run optimization steps over training batches and compute batch metrics\n state = train_step(state, batch) # get updated train state (which contains the updated parameters)\n state = compute_metrics(state=state, batch=batch) # aggregate batch metrics\n\n if (step+1) % num_steps_per_epoch == 0: # one training epoch has passed\n for metric,value in state.metrics.compute().items(): # compute metrics\n metrics_history[f'train_{metric}'].append(value) # record metrics\n state = state.replace(metrics=state.metrics.empty()) # reset train_metrics for next training epoch\n\n # Compute metrics on the test set after each training epoch\n test_state = state\n for test_batch in test_ds.as_numpy_iterator():\n test_state = compute_metrics(state=test_state, batch=test_batch)\n\n for metric,value in test_state.metrics.compute().items():\n metrics_history[f'test_{metric}'].append(value)\n\n print(f\"train epoch: {(step+1) // num_steps_per_epoch}, \"\n f\"loss: {metrics_history['train_loss'][-1]}, \"\n f\"accuracy: {metrics_history['train_accuracy'][-1] * 100}\")\n print(f\"test epoch: {(step+1) // num_steps_per_epoch}, \"\n f\"loss: {metrics_history['test_loss'][-1]}, \"\n f\"accuracy: {metrics_history['test_accuracy'][-1] * 100}\")\n```\n\n## 11. Visualize metrics\n\n```{code-cell}\n:outputId: 431a2fcd-44fa-4202-f55a-906555f060ac\n\nimport matplotlib.pyplot as plt # Visualization\n\n# Plot loss and accuracy in subplots\nfig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))\nax1.set_title('Loss')\nax2.set_title('Accuracy')\nfor dataset in ('train','test'):\n ax1.plot(metrics_history[f'{dataset}_loss'], label=f'{dataset}_loss')\n ax2.plot(metrics_history[f'{dataset}_accuracy'], label=f'{dataset}_accuracy')\nax1.legend()\nax2.legend()\nplt.show()\nplt.clf()\n```\n\n## 12. Perform inference on test set\n\nDefine a jitted inference function `pred_step`. Use the learned parameters to do model inference on the test set and visualize the images and their corresponding predicted labels.\n\n```{code-cell}\n@jax.jit\ndef pred_step(state, batch):\n logits = state.apply_fn({'params': state.params}, test_batch['image'])\n return logits.argmax(axis=1)\n\ntest_batch = test_ds.as_numpy_iterator().next()\npred = pred_step(state, test_batch)\n```\n\n```{code-cell}\n:outputId: 1db5a01c-9d70-4f7d-8c0d-0a3ad8252d3e\n\nfig, axs = plt.subplots(5, 5, figsize=(12, 12))\nfor i, ax in enumerate(axs.flatten()):\n ax.imshow(test_batch['image'][i, ..., 0], cmap='gray')\n ax.set_title(f\"label={pred[i]}\")\n ax.axis('off')\n```\n\nCongratulations! You made it to the end of the annotated MNIST example. You can revisit\nthe same example, but structured differently as a couple of Python modules, test\nmodules, config files, another Colab, and documentation in Flax's Git repo:\n\n[https://github.com/google/flax/tree/main/examples/mnist](https://github.com/google/flax/tree/main/examples/mnist)\n" }, { "path": "docs/robots.txt", "content": "User-agent: *\n\nDisallow: /api_reference/flax.linen/_autosummary/ # for SEO, since Google still indexes this deprecated link\n\nSitemap: https://flax.readthedocs.io/sitemap.xml\n" }, { "path": "docs_nnx/.gitignore", "content": "_formatted_howtos\n" }, { "path": "docs_nnx/.readthedocs.yaml", "content": "# .readthedocs.yml\n# Read the Docs configuration file\n# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details\n\n# Required\nversion: 2\n\nbuild:\n os: ubuntu-22.04\n tools:\n python: \"3.12\"\n jobs:\n pre_build:\n - pip install \".[all, testing, docs]\"\n - pip install -U --pre jax jaxlib -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/\n\n# Build documentation in the docs/ directory with Sphinx\nsphinx:\n configuration: docs_nnx/conf.py\n\n# Optionally build your docs in additional formats such as PDF and ePub\nformats:\n - htmlzip\n - epub\n # - pdf\n" }, { "path": "docs_nnx/Makefile", "content": "# Minimal makefile for Sphinx documentation\n#\n\n# You can set these variables from the command line, and also\n# from the environment for the first two.\nSPHINXOPTS ?=\nSPHINXBUILD ?= sphinx-build\nSOURCEDIR = .\nBUILDDIR = _build\n\n# Put it first so that \"make\" without argument is like \"make help\".\nhelp:\n\t@$(SPHINXBUILD) -M help \"$(SOURCEDIR)\" \"$(BUILDDIR)\" $(SPHINXOPTS) $(O)\n\n.PHONY: help Makefile\n\n# Catch-all target: route all unknown targets to Sphinx using the new\n# \"make mode\" option. $(O) is meant as a shortcut for $(SPHINXOPTS).\n%: Makefile\n\t@$(SPHINXBUILD) -M $@ \"$(SOURCEDIR)\" \"$(BUILDDIR)\" $(SPHINXOPTS) $(O)\n" }, { "path": "docs_nnx/README.md", "content": "# Where to find the docs\n\nThe FLAX documentation can be found here:\nhttps://flax.readthedocs.io/en/latest/\n\n# How to build the docs\n\n1. Clone the `flax` repository with `git clone https://github.com/google/flax.git`.\n1. In the main `flax` folder, install the required dependencies using `uv pip install -e .[docs]`.\n1. [Optional] If you need to make any local changes to the docs, create and switch to a branch. Make your changes to the docs in that branch.\n1. To build the docs, in the `flax/docs_nnx` folder run the make script: `make html`. Alternatively, install [`entr`](https://github.com/eradman/entr/), which helps run arbitrary commands when files change. Then run `find ../ ! -regex '.*/[\\.|\\_].*' | entr -s 'make html'`.\n1. If the build is successful, you should get the `The HTML pages are in _build/html.` message. You can preview the docs in `flax/docs/_build/html`.\n\n# How to run embedded code tests\n\nWe use `doctest` blocks for embedded code in documents, that are also\ntested. Learn more at https://www.sphinx-doc.org/en/master/usage/extensions/doctest.html\n\nTo run tests locally, run `make doctest`\n\n# How to write code documentation\n\nOur documentation is written in reStructuredText for Sphinx. It is a\nmeta-language that is compiled into online documentation. For more details,\ncheck out\n[Sphinx's documentation](https://www.sphinx-doc.org/en/master/usage/restructuredtext/index.html).\nAs a result, our docstrings adhere to a specific syntax that has to be kept in\nmind. Below we provide some guidelines.\n\nTo learn how to contribute to Jupyter Notebooks or other formats in Flax docs,\nrefer to the dedicated\n[Contributing](https://flax.readthedocs.io/en/latest/contributing.html) page.\n\n## How much information to put in a docstring\n\nDocstring should be informative. We prefer to err on the side of too much\ndocumentation than too little. For instance, providing a one-line explanation\nto a new `Module` which implements new functionality is not sufficient.\n\nFurthermore, we highly encourage adding examples to your docstrings, so users\ncan directly see how code can be used.\n\n## How to write inline tested code\n\nWe use [doctest](https://docs.python.org/3/library/doctest.html) syntax for\nwriting examples in documentation. These examples are ran as tests as part of\nour CI process. In order to write `doctest` code in your documentation, please\nuse the following notation:\n\n```bash\n# Example code::\n#\n# def sum(a, b):\n# return a + b\n#\n# sum(0, 1)\n```\n\nThe `Example code` string at the beginning can be replaced by anything as long\nas there are two semicolons and a newline following it, and the code is\nindented.\n\n## How to use \"code font\"\n\nWhen writing code font in a docstring, please use double backticks. Example:\n\n```bash\n# This returns a ``str`` object.\n```\n\nNote that argument names and objects like True, None or any strings should\nusually be put in `code`.\n\n## How to create cross-references/links\n\nIt is possible to create cross-references to other classes, functions, and\nmethods. In the following, `obj_typ` is either `class`, `func`, or `meth`.\n\n```bash\n# First method:\n# :`path_to_obj`\n\n# Second method:\n# ::`description `\n```\n\nYou can use the second method if the `path_to_obj` is very long. Some examples:\n\n```bash\n# Create: a reference to class flax.linen.Module.\n# :class:`flax.linen.Module`\n\n# Create a reference to local function my_func.\n# :func:`my_func`\n\n# Create a reference \"Module.apply()\" to method flax.linen.Module.apply.\n# :meth:`Module.apply() ` #\n```\n\nTo creata a hyperlink, use the following syntax:\n```bash\n# Note the double underscore at the end:\n# `Link to Google `__\n```\n\n### How to specify arguments for classes and methods\n\n* Class attributes should be specified using the `Attributes:` tag.\n* Method argument should be specified using the `Args:` tags.\n* All attributes and arguments should have types.\n\nHere is an example from our library:\n\n```python\nclass DenseGeneral(Module):\n \"\"\"A linear transformation with flexible axes.\n Attributes:\n features: int or tuple with number of output features.\n axis: int or tuple with axes to apply the transformation on. For instance,\n (-2, -1) will apply the transformation to the last two axes.\n batch_dims: tuple with batch axes.\n use_bias: whether to add a bias to the output (default: True).\n dtype: the dtype of the computation (default: float32).\n kernel_init: initializer function for the weight matrix.\n bias_init: initializer function for the bias.\n precision: numerical precision of the computation see `jax.lax.Precision`\n for details.\n \"\"\"\n features: Union[int, Iterable[int]]\n axis: Union[int, Iterable[int]] = -1\n batch_dims: Iterable[int] = ()\n use_bias: bool = True\n dtype: Dtype = jnp.float32\n kernel_init: Callable[[PRNGKey, Shape, Dtype], Array] = default_kernel_init\n bias_init: Callable[[PRNGKey, Shape, Dtype], Array] = zeros\n precision: Any = None\n\n @compact\n def __call__(self, inputs: Array) -> Array:\n \"\"\"Applies a linear transformation to the inputs along multiple dimensions.\n Args:\n inputs: The nd-array to be transformed.\n Returns:\n The transformed input.\n \"\"\"\n ...\n```" }, { "path": "docs_nnx/_ext/codediff.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Sphinx directive for creating code diff tables.\n\nUse directive as follows:\n\n.. codediff::\n :title: , \n\n \n ---\n \n\nIn order to highlight a line of code, append \"#!\" to it.\n\"\"\"\n\n\nimport sphinx\nfrom docutils import nodes\nfrom docutils.parsers.rst import directives\nfrom docutils.statemachine import ViewList\nfrom sphinx.util.docutils import SphinxDirective\n\nMISSING = object()\n\n\nclass CodeDiffParser:\n def parse(\n self,\n lines: list[str],\n title: str,\n groups: list[str] | None = None,\n skip_test: str | None = None,\n code_sep: str = '---',\n sync: object = MISSING,\n ):\n \"\"\"Parse the code diff block and format it so that it\n renders in different tabs and is tested by doctest.\n\n For example:\n\n .. testcode:: tab0, tab2, tab3\n\n \n\n .. codediff::\n :title: Tab 0, Tab 1, Tab 2, Tab 3\n :groups: tab0, tab1, tab2, tab3\n :skip_test: tab1, tab3\n\n \n\n ---\n\n \n\n ---\n\n \n\n ---\n\n \n\n For group tab0: and are executed.\n For group tab1: Nothing is executed.\n For group tab2: and are executed.\n For group tab3: is executed.\n\n Arguments:\n lines: a string list, where each element is a single string code line\n title: a single string that contains the titles of each tab (they should\n be separated by commas)\n groups: a single string that contains the group of each tab (they should\n be separated by commas). Code snippets that are part of the same group\n will be executed together. If groups=None, then the group names will\n default to the tab title names.\n skip_test: a single string denoting which group(s) to skip testing (they\n should be separated by commas). This is useful for legacy code snippets\n that no longer run correctly anymore. If skip_test=None, then no tests\n are skipped.\n code_sep: the separator character(s) used to denote a separate code block\n for a new tab. The default code separator is '---'.\n sync: an option for Sphinx directives, that will sync all tabs together.\n This means that if the user clicks to switch to another tab, all tabs\n will switch to the new tab.\n \"\"\"\n titles = [t.strip() for t in title.split(',')]\n num_tabs = len(titles)\n\n sync = sync is not MISSING\n # skip legacy code snippets in upgrade guides\n if skip_test is not None:\n skip_tests = {index.strip() for index in skip_test.split(',')}\n else:\n skip_tests = set()\n\n code_blocks = '\\n'.join(lines)\n if code_blocks.count(code_sep) != num_tabs - 1:\n raise ValueError(\n f'Expected {num_tabs-1} code separator(s) for {num_tabs} tab(s), but got {code_blocks.count(code_sep)} code separator(s) instead.'\n )\n code_blocks = [\n code_block.split('\\n')\n for code_block in code_blocks.split(code_sep + '\\n')\n ] # list[code_tab_list1[string_line1, ...], ...]\n\n # by default, put each code snippet in a different group denoted by an index number, to be executed separately\n if groups is not None:\n groups = [group_name.strip() for group_name in groups.split(',')]\n else:\n groups = titles\n if len(groups) != num_tabs:\n raise ValueError(\n f'Expected {num_tabs} group assignment(s) for {num_tabs} tab(s), but got {len(groups)} group assignment(s) instead.'\n )\n\n tabs = []\n test_codes = []\n for i, code_block in enumerate(code_blocks):\n if groups[i] not in skip_tests:\n test_codes.append((code_block, groups[i]))\n tabs.append((titles[i], self._code_block(code_block)))\n output = self._tabs(*tabs, sync=sync)\n\n return output, test_codes\n\n def _code_block(self, lines):\n \"\"\"Creates a codeblock.\"\"\"\n # Remove right trailing whitespace so we can detect the comments.\n lines = [x.rstrip() for x in lines]\n highlight = lambda x: x.endswith('#!')\n code = map(lambda x: x[:-2].rstrip() if highlight(x) else x, lines)\n highlights = [i + 1 for i in range(len(lines)) if highlight(lines[i])]\n highlights = ','.join(str(i) for i in highlights)\n\n directive = ['.. code-block:: python']\n if highlights:\n directive += [f' :emphasize-lines: {highlights}']\n\n # Indent code and add empty line so the code is picked up by the directive.\n return directive + [''] + list(map(lambda x: ' ' + x, code))\n\n def _tabs(self, *contents: tuple[str, list[str]], sync):\n output = ['.. tab-set::'] + [' ']\n\n for title, content in contents:\n output += [f' .. tab-item:: {title}']\n\n if sync:\n key = title.strip()\n output += [f' :sync: {key}']\n\n output += [' ']\n output += [' ' + line for line in content]\n\n return output\n\n\nclass CodeDiffDirective(SphinxDirective):\n has_content = True\n option_spec = {\n 'title': directives.unchanged,\n 'groups': directives.unchanged,\n 'skip_test': directives.unchanged,\n 'code_sep': directives.unchanged,\n 'sync': directives.flag,\n }\n\n def run(self):\n table_code, test_codes = CodeDiffParser().parse(\n list(self.content), **self.options\n )\n\n # Create a test node as a comment node so it won't show up in the docs.\n # We add attribute \"testnodetype\" so it is be picked up by the doctest\n # builder. This functionality is not officially documented but can be found\n # in the source code:\n # https://github.com/sphinx-doc/sphinx/blob/master/sphinx/ext/doctest.py\n # (search for 'testnodetype').\n test_nodes = []\n for test_code, group in test_codes:\n test_node = nodes.comment(\n '\\n'.join(test_code),\n '\\n'.join(test_code),\n testnodetype='testcode',\n groups=[group],\n )\n self.set_source_info(test_node)\n test_node['options'] = {}\n test_node['language'] = 'python3'\n test_nodes.append(test_node)\n\n # The table node is the side-by-side diff view that will be shown on RTD.\n table_node = nodes.paragraph()\n self.content = ViewList(table_code, self.content.parent)\n self.state.nested_parse(self.content, self.content_offset, table_node)\n\n return [table_node] + test_nodes\n\n\ndef setup(app):\n app.add_directive('codediff', CodeDiffDirective)\n\n return {\n 'version': sphinx.__display_version__,\n 'parallel_read_safe': True,\n 'parallel_write_safe': True,\n }\n" }, { "path": "docs_nnx/_ext/codediff_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for codediff Sphinx extension.\"\"\"\n\nfrom absl.testing import parameterized\nfrom codediff import CodeDiffParser\n\n\nclass CodeDiffTest(parameterized.TestCase):\n def test_parse(self):\n input_text = r\"\"\"@jax.jit #!\ndef get_initial_params(key): #!\n init_val = jnp.ones((1, 28, 28, 1), jnp.float32)\n initial_params = CNN().init(key, init_val)['params']\n extra_line\n return initial_params\n---\n@jax.pmap #!\ndef get_initial_params(key):\n init_val = jnp.ones((1, 28, 28, 1), jnp.float32)\n initial_params = CNN().init(key, init_val)['params']\n return initial_params\"\"\"\n\n expected_table = \"\"\".. tab-set::\\n \\n .. tab-item:: Single device\\n \\n .. code-block:: python\\n :emphasize-lines: 1,2\\n \\n @jax.jit\\n def get_initial_params(key):\\n init_val = jnp.ones((1, 28, 28, 1), jnp.float32)\\n initial_params = CNN().init(key, init_val)['params']\\n extra_line\\n return initial_params\\n \\n .. tab-item:: Ensembling on multiple devices\\n \\n .. code-block:: python\\n :emphasize-lines: 1\\n \\n @jax.pmap\\n def get_initial_params(key):\\n init_val = jnp.ones((1, 28, 28, 1), jnp.float32)\\n initial_params = CNN().init(key, init_val)['params']\\n return initial_params\"\"\"\n\n expected_testcodes = [\n r\"\"\"@jax.jit #!\ndef get_initial_params(key): #!\n init_val = jnp.ones((1, 28, 28, 1), jnp.float32)\n initial_params = CNN().init(key, init_val)['params']\n extra_line\n return initial_params\n\"\"\",\n r\"\"\"@jax.pmap #!\ndef get_initial_params(key):\n init_val = jnp.ones((1, 28, 28, 1), jnp.float32)\n initial_params = CNN().init(key, init_val)['params']\n return initial_params\"\"\",\n ]\n\n title_left = 'Single device'\n title_right = 'Ensembling on multiple devices'\n\n actual_table, actual_testcodes = CodeDiffParser().parse(\n lines=input_text.split('\\n'),\n title=f'{title_left}, {title_right}',\n )\n\n actual_table = '\\n'.join(actual_table)\n actual_testcodes = ['\\n'.join(testcode) for testcode, _ in actual_testcodes]\n\n self.assertEqual(expected_table, actual_table)\n self.assertEqual(expected_testcodes[0], actual_testcodes[0])\n self.assertEqual(expected_testcodes[1], actual_testcodes[1])\n\n @parameterized.parameters(\n {\n 'input_text': r\"\"\"x = 1\n ---\n x = 2\n\"\"\",\n 'title': 'Tab 0, Tab1, Tab2',\n 'groups': None,\n 'error_msg': 'Expected 2 code separator\\\\(s\\\\) for 3 tab\\\\(s\\\\), but got 1 code separator\\\\(s\\\\) instead.',\n },\n {\n 'input_text': r\"\"\"x = 1\n ---\n x = 2\n ---\n x = 3\n ---\n x = 4\n\"\"\",\n 'title': 'Tab 0, Tab1, Tab2',\n 'groups': None,\n 'error_msg': 'Expected 2 code separator\\\\(s\\\\) for 3 tab\\\\(s\\\\), but got 3 code separator\\\\(s\\\\) instead.',\n },\n {\n 'input_text': r\"\"\"x = 1\n ---\n x = 2\n ---\n x = 3\n\"\"\",\n 'title': 'Tab 0, Tab1, Tab2',\n 'groups': 'tab0, tab2',\n 'error_msg': 'Expected 3 group assignment\\\\(s\\\\) for 3 tab\\\\(s\\\\), but got 2 group assignment\\\\(s\\\\) instead.',\n },\n {\n 'input_text': r\"\"\"x = 1\n ---\n x = 2\n ---\n x = 3\n\"\"\",\n 'title': 'Tab 0, Tab1, Tab2',\n 'groups': 'tab0, tab1, tab2, tab3',\n 'error_msg': 'Expected 3 group assignment\\\\(s\\\\) for 3 tab\\\\(s\\\\), but got 4 group assignment\\\\(s\\\\) instead.',\n },\n )\n def test_parse_errors(self, input_text, title, groups, error_msg):\n with self.assertRaisesRegex(ValueError, error_msg):\n _, _ = CodeDiffParser().parse(\n lines=input_text.split('\\n'),\n title=title,\n groups=groups,\n )\n" }, { "path": "docs_nnx/_ext/flax_module.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Sphinx directive for visualizing Flax modules.\n\nUse directive as follows:\n\n.. flax_module::\n :module: flax.linen\n :class: Dense\n\"\"\"\n\nimport importlib\n\nimport sphinx\nimport sphinx.ext.autosummary.generate as ag\nfrom docutils import nodes\nfrom docutils.parsers.rst import directives\nfrom docutils.statemachine import ViewList\nfrom sphinx.util.docutils import SphinxDirective\n\nfrom docs.conf_sphinx_patch import generate_autosummary_content\n\n\ndef render_module(modname: str, qualname: str, app):\n parent = importlib.import_module(modname)\n obj = getattr(parent, qualname)\n template = ag.AutosummaryRenderer(app)\n template_name = 'flax_module'\n imported_members = False\n recursive = False\n context = {}\n return generate_autosummary_content(\n qualname,\n obj,\n parent,\n template,\n template_name,\n imported_members,\n app,\n recursive,\n context,\n modname,\n qualname,\n )\n\n\nclass FlaxModuleDirective(SphinxDirective):\n has_content = True\n option_spec = {\n 'module': directives.unchanged,\n 'class': directives.unchanged,\n }\n\n def run(self):\n module_template = render_module(\n self.options['module'], self.options['class'], self.env.app\n )\n module_template = module_template.splitlines()\n\n # Create a container for the rendered nodes\n container_node = nodes.container()\n self.content = ViewList(module_template, self.content.parent)\n self.state.nested_parse(self.content, self.content_offset, container_node)\n\n return [container_node]\n\n\ndef setup(app):\n app.add_directive('flax_module', FlaxModuleDirective)\n\n return {\n 'version': sphinx.__display_version__,\n 'parallel_read_safe': True,\n 'parallel_write_safe': True,\n }\n" }, { "path": "docs_nnx/_static/css/flax_theme.css", "content": "@import url(\"theme.css\");\n\n.wy-nav-content {\n max-width: 1290px;\n}\n\n.rst-content table.docutils {\n width: 100%;\n}\n\n.rst-content table.docutils td {\n vertical-align: top;\n padding: 0;\n}\n\n.rst-content table.docutils td p {\n padding: 8px;\n}\n\n.rst-content div[class^=highlight] {\n border: 0;\n margin: 0;\n}\n" }, { "path": "docs_nnx/_templates/autosummary/flax_module.rst", "content": "{{ fullname | escape | underline }}\n\n.. currentmodule:: {{ module }}\n\n.. autoclass:: {{ objname }}\n :exclude-members:\n\n .. automethod:: __call__\n\n {% block methods %}\n\n {% for item in methods %}\n {%- if item not in inherited_members and item not in annotations and not item in ['__init__', 'setup'] %}\n .. automethod:: {{ item }}\n {%- endif %}\n {%- endfor %}\n\n {% if methods %}\n .. rubric:: Methods\n\n .. autosummary::\n\n {% for item in methods %}\n {%- if item not in inherited_members and item not in annotations and not item in ['__init__', 'setup'] %}\n ~{{ name }}.{{ item }}\n {%- endif %}\n {%- endfor %}\n {% endif %}\n {% endblock %}" }, { "path": "docs_nnx/api_reference/flax.config.rst", "content": "\nflax.config package\n====================\n\n.. automodule:: flax.configurations\n :members:\n :undoc-members:\n :exclude-members: FlagHolder, bool_flag, static_bool_env\n" }, { "path": "docs_nnx/api_reference/flax.core.frozen_dict.rst", "content": "\nflax.core.frozen_dict package\n=============================\n\n.. currentmodule:: flax.core.frozen_dict\n\n.. autoclass:: FrozenDict\n :members: pretty_repr, copy, pop, unfreeze, tree_flatten\n\n.. autofunction:: freeze\n\n.. autofunction:: unfreeze\n\n.. autofunction:: copy\n\n.. autofunction:: pop\n\n.. autofunction:: pretty_repr\n" }, { "path": "docs_nnx/api_reference/flax.nnx/bridge.rst", "content": "bridge\n------------------------\n\n.. automodule:: flax.nnx.bridge\n.. currentmodule:: flax.nnx.bridge\n\n.. flax_module::\n :module: flax.nnx.bridge\n :class: ToNNX\n\n.. flax_module::\n :module: flax.nnx.bridge\n :class: ToLinen\n\n.. autofunction:: to_linen\n\n.. flax_module::\n :module: flax.nnx.bridge\n :class: NNXMeta\n" }, { "path": "docs_nnx/api_reference/flax.nnx/filterlib.rst", "content": "filterlib\n------------------------\n\n.. automodule:: flax.nnx\n.. currentmodule:: flax.nnx\n\n\n.. autofunction:: flax.nnx.filterlib.to_predicate\n.. autoclass:: WithTag\n.. autoclass:: PathContains\n.. autoclass:: OfType\n.. autoclass:: Any\n.. autoclass:: All\n.. autoclass:: Not\n.. autoclass:: Everything\n.. autoclass:: Nothing" }, { "path": "docs_nnx/api_reference/flax.nnx/graph.rst", "content": "graph\n------------------------\n\n.. automodule:: flax.nnx\n.. currentmodule:: flax.nnx\n\n\n.. autofunction:: split\n.. autofunction:: merge\n.. autofunction:: update\n.. autofunction:: pop\n.. autofunction:: state\n.. autofunction:: variables\n.. autofunction:: graph\n.. autofunction:: graphdef\n.. autofunction:: iter_graph\n.. autofunction:: recursive_map\n.. autofunction:: clone\n.. autofunction:: call\n.. autofunction:: set_metadata\n.. autofunction:: cached_partial\n\n.. autoclass:: GraphDef\n :members:\n\n.. autoclass:: UpdateContext\n :members:\n\n.. autofunction:: update_context\n.. autofunction:: current_update_context\n\n.. autofunction:: find_duplicates\n.. autofunction:: pure\n.. autofunction:: flatten\n.. autofunction:: unflatten\n" }, { "path": "docs_nnx/api_reference/flax.nnx/helpers.rst", "content": "helpers\n------------------------\n\n.. automodule:: flax.nnx\n.. currentmodule:: flax.nnx\n\n\n.. autoclass:: Sequential\n :members:\n.. autoclass:: List\n :members:\n.. autoclass:: Dict\n :members:\n.. autoclass:: TrainState\n :members:" }, { "path": "docs_nnx/api_reference/flax.nnx/index.rst", "content": "flax.nnx\n------------------------\n\nExperimental API. See the `NNX page `__ for more details.\n\n.. toctree::\n :maxdepth: 3\n\n graph\n object\n module\n nn/index\n rnglib\n spmd\n state\n training/index\n transforms\n variables\n helpers\n visualization\n filterlib\n bridge\n" }, { "path": "docs_nnx/api_reference/flax.nnx/module.rst", "content": "module\n------------------------\n\n.. automodule:: flax.nnx\n :members: iter_children, iter_modules\n.. currentmodule:: flax.nnx\n.. autoclass:: Module\n :members:\n.. autofunction:: view\n.. autofunction:: view_info\n.. autofunction:: with_attributes\n" }, { "path": "docs_nnx/api_reference/flax.nnx/nn/activations.rst", "content": "Activation functions\n------------------------\n\n.. automodule:: flax.nnx\n.. currentmodule:: flax.nnx\n\n.. autofunction:: celu\n.. autofunction:: elu\n.. autofunction:: gelu\n.. autofunction:: glu\n.. autofunction:: hard_sigmoid\n.. autofunction:: hard_silu\n.. autofunction:: hard_swish\n.. autofunction:: hard_tanh\n.. autofunction:: leaky_relu\n.. autofunction:: log_sigmoid\n.. autofunction:: log_softmax\n.. autofunction:: logsumexp\n.. autofunction:: one_hot\n.. autofunction:: relu\n.. autofunction:: relu6 as relu6,\n.. autofunction:: selu\n.. autofunction:: sigmoid\n.. autofunction:: identity\n.. autofunction:: silu\n.. autofunction:: soft_sign\n.. autofunction:: softmax\n.. autofunction:: softplus\n.. autofunction:: standardize\n.. autofunction:: swish\n.. autofunction:: tanh" }, { "path": "docs_nnx/api_reference/flax.nnx/nn/attention.rst", "content": "Attention\n------------------------\n\n.. automodule:: flax.nnx\n.. currentmodule:: flax.nnx\n\n.. flax_module::\n :module: flax.nnx\n :class: MultiHeadAttention\n\n.. autofunction:: combine_masks\n.. autofunction:: dot_product_attention\n.. autofunction:: make_attention_mask\n.. autofunction:: make_causal_mask" }, { "path": "docs_nnx/api_reference/flax.nnx/nn/dtypes.rst", "content": "Dtypes\n------------------------\n\n.. automodule:: flax.nnx.nn.dtypes\n.. currentmodule:: flax.nnx.nn.dtypes\n\n.. autofunction:: canonicalize_dtype\n.. autofunction:: promote_dtype" }, { "path": "docs_nnx/api_reference/flax.nnx/nn/index.rst", "content": "nn\n----------------------------\n\nNeural network layers and activation functions used in NNX :class:`Module`'s.\nSee the `NNX page `__ for more details.\n\n.. toctree::\n :maxdepth: 3\n\n activations\n attention\n dtypes\n initializers\n linear\n lora\n normalization\n recurrent\n stochastic\n\n" }, { "path": "docs_nnx/api_reference/flax.nnx/nn/initializers.rst", "content": "Initializers\n------------------------\n\n.. automodule:: flax.nnx.initializers\n.. currentmodule:: flax.nnx.initializers\n\n.. autofunction:: constant\n.. autofunction:: delta_orthogonal\n.. autofunction:: glorot_normal\n.. autofunction:: glorot_uniform\n.. autofunction:: he_normal\n.. autofunction:: he_uniform\n.. autofunction:: kaiming_normal\n.. autofunction:: kaiming_uniform\n.. autofunction:: lecun_normal\n.. autofunction:: lecun_uniform\n.. autofunction:: normal\n.. autofunction:: truncated_normal\n.. autofunction:: ones\n.. autofunction:: ones_init\n.. autofunction:: orthogonal\n.. autofunction:: uniform\n.. autofunction:: variance_scaling\n.. autofunction:: xavier_normal\n.. autofunction:: xavier_uniform\n.. autofunction:: zeros\n.. autofunction:: zeros_init\n" }, { "path": "docs_nnx/api_reference/flax.nnx/nn/linear.rst", "content": "Linear\n------------------------\n\nNNX linear layer classes.\n\n.. automodule:: flax.nnx\n.. currentmodule:: flax.nnx\n\n.. flax_module::\n :module: flax.nnx\n :class: Conv\n\n.. flax_module::\n :module: flax.nnx\n :class: ConvTranspose\n\n.. flax_module::\n :module: flax.nnx\n :class: Embed\n\n.. flax_module::\n :module: flax.nnx\n :class: Linear\n\n.. flax_module::\n :module: flax.nnx\n :class: LinearGeneral\n\n.. flax_module::\n :module: flax.nnx\n :class: Einsum" }, { "path": "docs_nnx/api_reference/flax.nnx/nn/lora.rst", "content": "LoRA\n------------------------\n\nNNX LoRA classes.\n\n.. automodule:: flax.nnx\n.. currentmodule:: flax.nnx\n\n.. flax_module::\n :module: flax.nnx\n :class: LoRA\n\n.. flax_module::\n :module: flax.nnx\n :class: LoRALinear\n" }, { "path": "docs_nnx/api_reference/flax.nnx/nn/normalization.rst", "content": "Normalization\n------------------------\n\n.. automodule:: flax.nnx\n.. currentmodule:: flax.nnx\n\n.. flax_module::\n :module: flax.nnx\n :class: BatchNorm\n\n.. flax_module::\n :module: flax.nnx\n :class: LayerNorm\n\n.. flax_module::\n :module: flax.nnx\n :class: RMSNorm\n\n.. flax_module::\n :module: flax.nnx\n :class: GroupNorm\n\n.. flax_module::\n :module: flax.nnx\n :class: InstanceNorm\n\n.. flax_module::\n :module: flax.nnx\n :class: SpectralNorm\n\n.. flax_module::\n :module: flax.nnx\n :class: WeightNorm\n" }, { "path": "docs_nnx/api_reference/flax.nnx/nn/recurrent.rst", "content": "Recurrent\n------------------------\n\n.. automodule:: flax.nnx.nn.recurrent\n.. currentmodule:: flax.nnx.nn.recurrent\n\n.. flax_module::\n :module: flax.nnx.nn.recurrent\n :class: LSTMCell\n\n.. flax_module::\n :module: flax.nnx.nn.recurrent\n :class: OptimizedLSTMCell\n\n.. flax_module::\n :module: flax.nnx.nn.recurrent\n :class: SimpleCell\n\n.. flax_module::\n :module: flax.nnx.nn.recurrent\n :class: GRUCell\n\n.. flax_module::\n :module: flax.nnx.nn.recurrent\n :class: RNN\n\n.. flax_module::\n :module: flax.nnx.nn.recurrent\n :class: Bidirectional\n\n\n.. autofunction:: flip_sequences" }, { "path": "docs_nnx/api_reference/flax.nnx/nn/stochastic.rst", "content": "Stochastic\n------------------------\n\n.. automodule:: flax.nnx\n.. currentmodule:: flax.nnx\n\n.. autoclass:: Dropout\n :members:" }, { "path": "docs_nnx/api_reference/flax.nnx/object.rst", "content": "object\n------------------------\n\n.. automodule:: flax.nnx\n.. currentmodule:: flax.nnx\n\n.. autoclass:: Pytree\n :members:\n.. autoclass:: Object\n :members:\n.. autofunction:: data\n.. autodata:: Data\n :annotation:\n.. autofunction:: static\n.. autodata:: Static\n :annotation:\n.. autofunction:: is_data\n.. autofunction:: register_data_type\n.. autofunction:: check_pytree" }, { "path": "docs_nnx/api_reference/flax.nnx/rnglib.rst", "content": "rnglib\n------------------------\n\n.. automodule:: flax.nnx\n.. currentmodule:: flax.nnx\n\n.. autoclass:: Rngs\n :members: __init__\n.. autoclass:: RngStream\n :members:\n.. autofunction:: split_rngs\n.. autofunction:: fork_rngs\n.. autofunction:: reseed\n" }, { "path": "docs_nnx/api_reference/flax.nnx/spmd.rst", "content": "spmd\n------------------------\n\n.. automodule:: flax.nnx\n.. currentmodule:: flax.nnx\n\n.. autofunction:: get_partition_spec\n.. autofunction:: get_named_sharding\n.. autofunction:: with_partitioning\n" }, { "path": "docs_nnx/api_reference/flax.nnx/state.rst", "content": "state\n------------------------\n\n.. automodule:: flax.nnx\n.. currentmodule:: flax.nnx\n\n\n.. autoclass:: State\n :members:\n\n.. autoclass:: FlatState\n :members:\n\n.. autofunction:: filter_state\n.. autofunction:: from_flat_state\n.. autofunction:: map_state\n.. autofunction:: merge_state\n.. autofunction:: replace_by_pure_dict\n.. autofunction:: restore_int_paths\n.. autofunction:: to_flat_state\n.. autofunction:: to_pure_dict\n.. autofunction:: split_state\n" }, { "path": "docs_nnx/api_reference/flax.nnx/summary.rst", "content": "summary\n------------------------\n\n.. automodule:: flax.nnx\n.. currentmodule:: flax.nnx\n\n.. autofunction:: tabulate" }, { "path": "docs_nnx/api_reference/flax.nnx/training/index.rst", "content": "training\n----------------------------\n\nExperimental API. See the `NNX page `__ for more details.\n\n.. toctree::\n :maxdepth: 3\n\n metrics\n optimizer\n\n" }, { "path": "docs_nnx/api_reference/flax.nnx/training/metrics.rst", "content": "Metrics\n------------------------\n\n.. automodule:: flax.nnx.metrics\n.. currentmodule:: flax.nnx.metrics\n\n\n.. autoclass:: Metric\n :members: __init__, reset, update, compute\n\n.. autoclass:: Average\n :members: __init__, reset, update, compute\n\n.. autoclass:: Accuracy\n :members: update\n\n.. autoclass:: Welford\n :members: __init__, reset, update, compute\n\n.. autoclass:: MultiMetric\n :members: __init__, reset, update, compute\n\n" }, { "path": "docs_nnx/api_reference/flax.nnx/training/optimizer.rst", "content": "Optimizer\n------------------------\n\n.. automodule:: flax.nnx.optimizer\n.. currentmodule:: flax.nnx.optimizer\n\n.. autoclass:: Optimizer\n :members: __init__, update\n" }, { "path": "docs_nnx/api_reference/flax.nnx/transforms.rst", "content": "transforms\n------------------------\n\n.. automodule:: flax.nnx\n.. currentmodule:: flax.nnx\n.. autofunction:: grad\n.. autofunction:: jit\n.. autofunction:: shard_map\n.. autofunction:: remat\n.. autofunction:: scan\n.. autoclass:: Carry\n.. autofunction:: value_and_grad\n.. autofunction:: vmap\n.. autofunction:: eval_shape\n.. autofunction:: custom_vjp\n.. autofunction:: vjp\n.. autofunction:: jvp\n.. autofunction:: cond\n.. autofunction:: switch\n.. autofunction:: while_loop\n.. autofunction:: fori_loop\n" }, { "path": "docs_nnx/api_reference/flax.nnx/variables.rst", "content": "variables\n------------------------\n\n.. automodule:: flax.nnx\n.. currentmodule:: flax.nnx\n\n.. autoclass:: BatchStat\n :members:\n.. autoclass:: Cache\n :members:\n.. autoclass:: Intermediate\n :members:\n.. autoclass:: Param\n :members:\n.. autoclass:: Variable\n :members:\n.. autoclass:: VariableMetadata\n :members:\n\n.. autofunction:: with_metadata\n\n.. autofunction:: variable_name_from_type\n.. autofunction:: variable_type_from_name\n.. autofunction:: register_variable_name\n" }, { "path": "docs_nnx/api_reference/flax.nnx/visualization.rst", "content": "visualization\n------------------------\n\n.. automodule:: flax.nnx\n.. currentmodule:: flax.nnx\n\n.. autofunction:: display" }, { "path": "docs_nnx/api_reference/flax.struct.rst", "content": "\nflax.struct package\n=====================\n\n.. currentmodule:: flax.struct\n\n.. automodule:: flax.struct\n\n\n.. autofunction:: dataclass\n\n\n.. autoclass:: PyTreeNode" }, { "path": "docs_nnx/api_reference/flax.training.rst", "content": "\nflax.training package\n=====================\n\nTrain state\n------------------------\n\n.. currentmodule:: flax.training.train_state\n\n.. autoclass:: TrainState\n :members: apply_gradients, create\n\n" }, { "path": "docs_nnx/api_reference/flax.traverse_util.rst", "content": "flax.traverse_util package\n============================\n\n.. currentmodule:: flax.traverse_util\n\n.. automodule:: flax.traverse_util\n\nDict utils\n------------\n\n.. autofunction:: flatten_dict\n\n.. autofunction:: unflatten_dict\n\n.. autofunction:: path_aware_map\n" }, { "path": "docs_nnx/api_reference/index.rst", "content": "API Reference\n=============\n\n.. toctree::\n :maxdepth: 4\n\n flax.nnx/index\n flax.core.frozen_dict\n flax.struct\n flax.training" }, { "path": "docs_nnx/conf.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Configuration file for the Sphinx documentation builder.\"\"\"\n\n\n# This file only contains a selection of the most common options. For a full\n# list see the documentation:\n# https://www.sphinx-doc.org/en/master/usage/configuration.html\n\n# -- Path setup --------------------------------------------------------------\n\n# If extensions (or modules to document with autodoc) are in another directory,\n# add these directories to sys.path here. If the directory is relative to the\n# documentation root, use os.path.abspath to make it absolute, like shown here.\n#\n# import os\n# import sys\n# sys.path.insert(0, os.path.abspath('.'))\n\nimport os\nimport sys\nimport doctest\n\nsys.path.insert(0, os.path.abspath('..'))\n# Include local extension.\nsys.path.append(os.path.abspath('./_ext'))\n# Set environment variable to indicate that we are building the docs.\nos.environ['FLAX_DOC_BUILD'] = 'true'\n\n# patch sphinx\n# -- Project information -----------------------------------------------------\n\nproject = 'Flax'\ncopyright = '2023, The Flax authors' # pylint: disable=redefined-builtin\nauthor = 'The Flax authors'\n\n\n# -- General configuration ---------------------------------------------------\n\n# Add any Sphinx extension module names here, as strings. They can be\n# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom\n# ones.\nextensions = [\n 'sphinx.ext.autodoc',\n 'sphinx.ext.autosummary',\n 'sphinx.ext.autosectionlabel',\n 'sphinx.ext.doctest',\n 'sphinx.ext.intersphinx',\n 'sphinx.ext.mathjax',\n 'sphinx.ext.napoleon',\n 'sphinx.ext.viewcode',\n 'myst_nb',\n 'codediff',\n 'flax_module',\n 'sphinx_design',\n]\n\n# Add any paths that contain templates here, relative to this directory.\ntemplates_path = ['_templates']\n\n# List of patterns, relative to source directory, that match files and\n# directories to ignore when looking for source files.\n# This pattern also affects html_static_path and html_extra_path.\nexclude_patterns = ['_build', 'Thumbs.db', '.DS_Store']\n\n# The suffix(es) of source filenames.\n# Note: important to list ipynb before md here: we have both md and ipynb\n# copies of each notebook, and myst will choose which to convert based on\n# the order in the source_suffix list. Notebooks which are not executed have\n# outputs stored in ipynb but not in md, so we must convert the ipynb.\nsource_suffix = ['.rst', '.ipynb', '.md']\n\nautosummary_generate = True\n\nmaster_doc = 'index'\n\nautodoc_typehints = 'none'\n\n\n# -- Options for HTML output -------------------------------------------------\n\n# The theme to use for HTML and HTML Help pages. See the documentation for\n# a list of builtin themes.\n#\n# html_theme = 'pydata_sphinx_theme'\nhtml_theme = 'sphinx_book_theme'\nhtml_css_files = ['css/flax_theme.css']\n\n# The name of an image file (relative to this directory) to place at the top\n# of the sidebar.\nhtml_logo = './flax.png'\nhtml_favicon = './flax.png'\n\n# title of the website\nhtml_title = ''\n\n# Add any paths that contain custom static files (such as style sheets) here,\n# relative to this directory. They are copied after the builtin static files,\n# so a file named 'default.css' will overwrite the builtin 'default.css'.\nhtml_static_path = ['_static']\n\nhtml_extra_path = ['robots.txt']\n\n# href with no underline and white bold text color\nannouncement = \"\"\"\n\n This site covers the new Flax NNX API. [Click here for the old Flax Linen API]\n\n\"\"\"\n\nhtml_theme_options = {\n 'repository_url': 'https://github.com/google/flax',\n 'use_repository_button': True, # add a 'link to repository' button\n 'use_issues_button': False, # add an 'Open an Issue' button\n 'path_to_docs': (\n 'docs_nnx'\n ), # used to compute the path to launch notebooks in colab\n 'launch_buttons': {\n 'colab_url': 'https://colab.research.google.com/',\n },\n 'prev_next_buttons_location': None,\n 'show_navbar_depth': 1,\n 'announcement': announcement,\n}\n\n# -- Options for myst ----------------------------------------------\n# uncomment line below to avoid running notebooks during development\nnb_execution_mode = os.environ.get(\"NB_EXECUTION_MODE\", 'off')\n# Notebook cell execution timeout; defaults to 30.\nnb_execution_timeout = 100\n# List of patterns, relative to source directory, that match notebook\n# files that will not be executed.\nmyst_enable_extensions = ['dollarmath']\nnb_execution_excludepatterns = [\n 'mnist_tutorial.ipynb', # <-- times out\n 'transfer_learning.ipynb', # <-- transformers requires flax<=0.7.0\n 'flax/nnx', # exclude nnx\n 'guides/demo.ipynb', # TODO(cgarciae): broken, remove or update\n 'examples/gemma.ipynb',\n 'guides/bridge_guide.ipynb', # TODO(cgarciae): broken, bridge doesn't support Linen sow yet\n]\n# raise exceptions on execution so CI can catch errors\nnb_execution_allow_errors = False\nnb_execution_raise_on_error = True\n\n# -- Extension configuration -------------------------------------------------\n\n# Tell sphinx-autodoc-typehints to generate stub parameter annotations including\n# types, even if the parameters aren't explicitly documented.\nalways_document_param_types = True\n\n# -- doctest configuration -------------------------------------------------\n\ndoctest_default_flags = doctest.NORMALIZE_WHITESPACE\ndoctest_global_setup = \"\"\"\nimport jax\nimport jax.numpy as jnp\nfrom flax import nnx\n\nimport logging as slog\nfrom absl import logging as alog\n\n# Avoid certain absl logging messages to break doctest\nfiltered_message = [\n 'SaveArgs.aggregate is deprecated',\n '',\n]\n\nclass _CustomLogFilter(slog.Formatter):\n def format(self, record):\n message = super(_CustomLogFilter, self).format(record)\n for m in filtered_message:\n if m in message:\n return ''\n return message\n\nalog.use_absl_handler()\nalog.get_absl_handler().setFormatter(_CustomLogFilter())\n\"\"\"\n" }, { "path": "docs_nnx/conf_sphinx_patch.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Patch Sphinx to improve documentation aesthetics.\"\"\"\n\n# TODO(cgarciae): Send a PR to sphinx to upstream this fix. Issue: https://github.com/google/flax/issues/2196\n# This patch is needed to make autosummary provide the \"annotations\"\n# variable so we can exclude function attributes from the methods list\n# in flax_module.rst. The patch as such only adds this single line:\n#\n# ns['annotations'] = list(getattr(obj, '__annotations__', {}).keys())'\n#\n# We should consider sending a PR to sphinx so we can get rid of this.\n# Original source: https://github.com/sphinx-doc/sphinx/blob/0aedcc9a916daa92d477226da67d33ce1831822e/sphinx/ext/autosummary/generate.py#L211-L351\nfrom typing import Any\n\nimport sphinx.ext.autodoc\nimport sphinx.ext.autosummary.generate as ag\n\n\ndef generate_autosummary_content(\n name: str,\n obj: Any,\n parent: Any,\n template: ag.AutosummaryRenderer,\n template_name: str,\n imported_members: bool,\n app: Any,\n recursive: bool,\n context: dict,\n modname: str = None,\n qualname: str = None,\n) -> str:\n doc = ag.get_documenter(app, obj, parent)\n\n def skip_member(obj: Any, name: str, objtype: str) -> bool:\n try:\n return app.emit_firstresult(\n 'autodoc-skip-member', objtype, name, obj, False, {}\n )\n except Exception as exc:\n ag.logger.warning(\n __(\n 'autosummary: failed to determine %r to be documented, '\n 'the following exception was raised:\\n%s'\n ),\n name,\n exc,\n type='autosummary',\n )\n return False\n\n def get_class_members(obj: Any) -> dict[str, Any]:\n members = sphinx.ext.autodoc.get_class_members(\n obj, [qualname], ag.safe_getattr\n )\n return {name: member.object for name, member in members.items()}\n\n def get_module_members(obj: Any) -> dict[str, Any]:\n members = {}\n for name in ag.members_of(obj, app.config):\n try:\n members[name] = ag.safe_getattr(obj, name)\n except AttributeError:\n continue\n return members\n\n def get_all_members(obj: Any) -> dict[str, Any]:\n if doc.objtype == 'module':\n return get_module_members(obj)\n elif doc.objtype == 'class':\n return get_class_members(obj)\n return {}\n\n def get_members(\n obj: Any,\n types: set[str],\n include_public: list[str] = [],\n imported: bool = True,\n ) -> tuple[list[str], list[str]]:\n items: list[str] = []\n public: list[str] = []\n\n all_members = get_all_members(obj)\n for name, value in all_members.items():\n documenter = ag.get_documenter(app, value, obj)\n if documenter.objtype in types:\n # skip imported members if expected\n if imported or getattr(value, '__module__', None) == obj.__name__:\n skipped = skip_member(value, name, documenter.objtype)\n if skipped is True:\n pass\n elif skipped is False:\n # show the member forcedly\n items.append(name)\n public.append(name)\n else:\n items.append(name)\n if name in include_public or not name.startswith('_'):\n # considers member as public\n public.append(name)\n return public, items\n\n def get_module_attrs(members: Any) -> tuple[list[str], list[str]]:\n \"\"\"Find module attributes with docstrings.\"\"\"\n attrs, public = [], []\n try:\n analyzer = ag.ModuleAnalyzer.for_module(name)\n attr_docs = analyzer.find_attr_docs()\n for namespace, attr_name in attr_docs:\n if namespace == '' and attr_name in members:\n attrs.append(attr_name)\n if not attr_name.startswith('_'):\n public.append(attr_name)\n except ag.PycodeError:\n pass # give up if ModuleAnalyzer fails to parse code\n return public, attrs\n\n def get_modules(obj: Any) -> tuple[list[str], list[str]]:\n items: list[str] = []\n for _, modname, _ispkg in ag.pkgutil.iter_modules(obj.__path__):\n fullname = name + '.' + modname\n try:\n module = ag.import_module(fullname)\n if module and hasattr(module, '__sphinx_mock__'):\n continue\n except ImportError:\n pass\n\n items.append(fullname)\n public = [x for x in items if not x.split('.')[-1].startswith('_')]\n return public, items\n\n ns: dict[str, Any] = {}\n ns.update(context)\n\n if doc.objtype == 'module':\n scanner = ag.ModuleScanner(app, obj)\n ns['members'] = scanner.scan(imported_members)\n ns['functions'], ns['all_functions'] = get_members(\n obj, {'function'}, imported=imported_members\n )\n ns['classes'], ns['all_classes'] = get_members(\n obj, {'class'}, imported=imported_members\n )\n ns['exceptions'], ns['all_exceptions'] = get_members(\n obj, {'exception'}, imported=imported_members\n )\n ns['attributes'], ns['all_attributes'] = get_module_attrs(ns['members'])\n ispackage = hasattr(obj, '__path__')\n if ispackage and recursive:\n ns['modules'], ns['all_modules'] = get_modules(obj)\n elif doc.objtype == 'class':\n ns['members'] = dir(obj)\n ns['inherited_members'] = set(dir(obj)) - set(obj.__dict__.keys())\n ns['methods'], ns['all_methods'] = get_members(\n obj, {'method'}, ['__init__']\n )\n ns['attributes'], ns['all_attributes'] = get_members(\n obj, {'attribute', 'property'}\n )\n ns['annotations'] = list(getattr(obj, '__annotations__', {}).keys())\n\n if modname is None or qualname is None:\n modname, qualname = ag.split_full_qualified_name(name)\n\n if doc.objtype in ('method', 'attribute', 'property'):\n ns['class'] = qualname.rsplit('.', 1)[0]\n\n if doc.objtype in ('class',):\n shortname = qualname\n else:\n shortname = qualname.rsplit('.', 1)[-1]\n\n ns['fullname'] = name\n ns['module'] = modname\n ns['objname'] = qualname\n ns['name'] = shortname\n\n ns['objtype'] = doc.objtype\n ns['underline'] = len(name) * '='\n\n if template_name:\n return template.render(template_name, ns)\n else:\n return template.render(doc.objtype, ns)\n\n\nag.generate_autosummary_content = generate_autosummary_content\n" }, { "path": "docs_nnx/contributing.md", "content": "# How to contribute\n\nEveryone can contribute to Flax, and the Flax development team values everyone's contributions!\nYou can contribute in many more ways than just writing code. Answering questions\non the [Flax GitHub Discussions page](https://github.com/google/flax/discussions), helping\neach other, and improving Flax documentation are extremely valuable to the Flax\necosystem.\n\nWe also appreciate if you spread the word, for instance by starring the [Flax GitHub repository](https://github.com/google/flax),\nor referencing Flax in blog posts of projects that used it.\n\nThis project follows\n[Google's Open Source Community Guidelines](https://opensource.google/conduct/).\n\n## Ways to contribute\n\nWe welcome pull requests (PRs), in particular for those issues\n[marked as PR-ready](https://github.com/google/flax/issues?q=is%3Aopen+is%3Aissue+label%3A%22Status%3A+pull+requests+welcome%22).\nFor other proposals, you should first open a GitHub Issue or a GitHub Discussion to\nstart a conversation about your planned contribution.\n\n## Contributing code using pull requests\n\nThe Flax development team performs all development using [Git](https://git-scm.com/). To contribute,\nyou should have basic knowledge of [Git](https://git-scm.com/) and [GitHub](https://docs.github.com).\n(You can learn how to set up Git by following Git's official\n[Getting Started - First-Time Git Setup](https://git-scm.com/book/en/v2/Getting-Started-First-Time-Git-Setup)\nand GitHub's [Set Up Git](https://docs.github.com/en/get-started/quickstart/set-up-git) guides.)\n\nTo contribute code to Flax on GitHub, follow these steps:\n\n### To create a pull request from a fork\n\n1. Using GitHub's web UI, fork the Flax repository by clicking the 'Fork' button on the\n [`github.com/google/flax` repository page](http://www.github.com/google/flax). This creates a\n fork (a copy) of the Flax repository in your own GitHub.\n\n Reference: [Creating a pull request from a fork](https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/proposing-changes-to-your-work-with-pull-requests/creating-a-pull-request-from-a-fork).\n\n2. Install [Python >=3.7](https://www.python.org/downloads/).\n\n3. (Optional) Create a virtual environment or a Docker container. See\n [`dev/README.md`](https://github.com/google/flax/blob/main/dev/README.md)\n for details on how to set up a Docker Container. To set up a virtual environment,\n run the following:\n\n ```bash\n python3 -m virtualenv env\n . env/bin/activate\n ```\n\n This ensures all your dependencies are installed in this environment.\n\n4. Clone your local forked Flax repo with `git clone`. Then, install the required packages\n with [PyPi](https://pip.pypa.io/en/stable/cli/pip_install/). This enables you to immediately\n test the code after modifying it:\n\n ```bash\n git clone https://github.com/YOUR_USERNAME/flax\n cd flax\n pip install -e \".[all,testing,docs]\"\n ```\n\n You can also use [uv](https://docs.astral.sh/uv/) to setup\n the development environment:\n\n ```bash\n uv sync --all-extras\n ```\n\n5. Set up pre-commit hooks, this will run some automated checks during each `git` commit and\n possibly update some files that require changes.\n\n ```bash\n pip install pre-commit\n pre-commit install\n ```\n\n6. Add the Google Flax repo (not your fork) as an upstream remote, so you can use it to sync your\n changes.\n\n ```bash\n git remote add upstream http://www.github.com/google/flax\n ```\n\n\n7. Create a branch, such as `my_development_branch`, you will develop from:\n\n ```bash\n git checkout -b my_development_branch\n ```\n\n8. Implement your changes using your favorite editor (we recommend\n [Visual Studio Code](https://code.visualstudio.com/)).\n\n Make sure the tests pass by running the following command from the top of\n the repository:\n\n ```bash\n ./tests/run_all_tests.sh\n ```\n\n9. Once you finish making changes, don't forget to create commits\n ([learn how to write a commit message](https://chris.beams.io/posts/git-commit/)):\n\n ```bash\n git add file1.py file2.py ...\n # or use `git add .` to add all changed files\n git commit -m \"Your commit message\"\n ```\n\n Then sync your code with the main repository:\n\n ```bash\n git fetch upstream\n git rebase upstream/main\n ```\n\n10. Finally, push your commit on your `my_development_branch`, and create a remote\n branch in your fork that you can use to create a pull request from:\n\n ```bash\n git push --set-upstream origin my_development_branch\n ```\n\n After running the command, you should get a GitHub link in your (VS Code) terminal output for creating a pull request.\n If you don't receive a link after `git push`, use the [GitHub web UI](https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/proposing-changes-to-your-work-with-pull-requests/creating-a-pull-request?tool=webui) to create a pull request.\n\n11. Make sure your pull request passes the\n [Flax PR checklist](https://github.com/google/flax/blob/main/.github/pull_request_template.md#checklist).\n If so, create a pull request from the Flax repository and send it for review.\n Consult [GitHub Help](https://help.github.com/articles/about-pull-requests/)\n for more information on using pull requests.\n\nYou can learn more in GitHub's [Creating a pull request from a fork\n](https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/proposing-changes-to-your-work-with-pull-requests/creating-a-pull-request-from-a-fork). documentation.\n\n### Adding or updating dependencies\n\nTo add or update dependencies, you must use `uv` after\nupdating the `pyproject.toml` file to ensure that the `uv.lock` file is up-to-date.\n\n```bash\nuv sync --all-extras\n```\nAlternatively use can use `uv add` to add or update the dependencies automatically, for example:\n\n```bash\nuv add 'some-package>=1.2.3'\n```\n\n### Updating Jupyter Notebooks\n\nWe use [jupytext](https://jupytext.readthedocs.io/) to maintain two synced copies of docs\nin `docs/notebooks`: one in the Jupyter Notebook (`.ipynb`) format, and one in Markdown (`.md`).\n\nThe former can be opened and executed directly in [Google Colab](https://colab.research.google.com/).\nMarkdown makes it easier to track changes/diffs within version control and, for example, GitHub\nweb UI, since `.ipynb` files are based on JSON.\n\n#### Editing Jupyter Notebooks (`.ipynb`)\n\nFor making large changes that substantially modify code and outputs, it's recommended to edit\nthe notebooks in [Jupyter](https://jupyter.org/install) or in [Colab](https://colab.research.google.com/).\n\nIf you choose to work in Colab, go to **File** and click **Upload notebook**, then pick your file.\nAfter loading it into Colab and editing it, make sure you run the cells, and that there aren't any errors.\nClick on **Runtime**, then select **Run all**. After you finish, click **File** > **Download** > **Download ipynb**.\nYou may also want to test that the file executes properly by using `sphinx-build`, as explained above.\n\nAfter you make changes in your Jupyter Notebook, follow the steps _Syncing notebooks_ below.\n\n#### Editing Markdown files (`.md`)\n\nFor making smaller changes to the text content of the notebooks, it is easiest to edit the\n`.md` versions using a text editor.\n\nAfter you make changes in your Markdown file, follow the steps _Syncing notebooks_ below.\n\n#### Syncing notebooks\n\nAfter editing either the `.ipynb` or `.md` versions of the docs, sync the two versions\nusing [jupytext](https://jupytext.readthedocs.io/) by running `jupytext --sync` on the updated\nnotebooks.\n\nFirst, make sure you have jupytext installed. The jupytext version should match\nthe one specified in [.pre-commit-config.yaml](https://github.com/google/flax/blob/main/.pre-commit-config.yaml)\n(currently, it is v1.13.8).\n\n```bash\npip install jupytext==1.13.8\n```\n\nThen, after you have made your changes in the Jupyter Notebook, sync the contents with its Markdown-equivalent\nfile by running the following command:\n\n```bash\njupytext --sync path/to/the/file.ipynb\n```\n\nSimilarly, to sync your Markdown file with its Jupyter Notebook version, run:\n\n```bash\njupytext --sync path/to/the/file.md\n```\n\nNote that if you receive an error, and it is the first time you worked in a Jupyter Notebook, you may need\nto (re)create a synced copy of the document (which is explained in detail in _Creating new notebooks_ section below):\n\n```bash\njupytext --set-formats ipynb,md:myst path/to/the/notebook.ipynb\n```\n\nOnce you're finished with syncing the `.md` and `.ipynb` files, you can check that they are properly synced using the\n[pre-commit](https://pre-commit.com/) framework to perform the same checks used\nin the Flax GitHub CI:\n\n```bash\ngit add docs -u # pre-commit runs on files in git staging.\npre-commit run jupytext\n```\n\n#### Creating new notebooks\n\nIf you are adding a new Jupyter Notebook to the documentation, you can use `jupytext --set-formats`.\nIt can set up both the Jupyter Notebook (`.ipynb`) and Markdown (`.md`) versions of the file:\n\n```bash\njupytext --set-formats ipynb,md:myst path/to/the/notebook.ipynb\n```\n\nThis works by adding a `\"jupytext\"` metadata field to the notebook file which specifies the\ndesired formats. The `jupytext --sync` command can then recognize them when invoked.\n\nAfter you make changes in your file(s), follow the steps from the _Syncing notebooks_\nsection above to keep the contents of both Markdown and Jupyter Notebook files in sync.\n\n#### Notebooks within the Sphinx build\n\nSome of the notebooks are built automatically as part of the pre-submit checks and\nas part of the [Read the Docs](https://flax.readthedocs.io/en/latest) build.\nThe build will fail if cells raise errors. If the errors are intentional, you can either catch them,\nor tag the cell with `raises-exceptions` metadata ([example PR](https://github.com/jax-ml/jax/pull/2402/files)).\nYou have to add this metadata by hand in the `.ipynb` file. It will be preserved when somebody else\nre-saves the notebook.\n\nWe exclude some notebooks from the build because, for example, they contain long computations.\nSee `exclude_patterns` in [`conf.py`](https://github.com/google/flax/blob/main/docs/conf.py).\n\n### Updating the pull request contents\n\nEvery pull request should ideally be limited to just one commit, so if you have multiple commits please squash them.\n\nAssuming you now have only one commit in your pull request, and want to add changes requested during review:\n\n1. Make the changes locally in your editor.\n2. Run `git commit -a --amend`. This updates the commit contents and allows you to edit the commit message.\n3. At this point, `git push` alone will result in an error. Instead, use `git push --force`.\n4. Check that it's done: The changes to your commit should be immediately reflected in the Github web UI.\n\n## Troubleshooting\n\n### Too many commits in a pull request\n\nIf your PR has too many commits associated with it (for example, more than five),\nyou need to squash them. Otherwise, the Flax docs build process may fail with an\nerror message. This is because of the following reasons:\n\n* There are more than five commits in your pull request; and\n* The Flax source sync process fails when the commit tree is too large.\n\nTo squash your commits, you can rebase your branch to `main` and create a new\ncommit containing all your changes, run the following command:\n\n```bash\ngit rebase main && git reset --soft main && git commit\n```\n\nThis will apply all your changes to the main branch. Note that if you had to\nresolve any conflicts while working on your change (for instance, you did a\n`pull upstream main` which led to conflict), then you will have to resolve these\nconflicts again.\n\nAfter you have successfully rebased your branch, you should push your changes.\nAnd because you changed the commit history, you may have to use `git push --force`.\n\n## Contributor License Agreement\n\nContributions to this project must be accompanied by a Contributor License\nAgreement. You (or your employer) retain the copyright to your contribution;\nthis simply gives us permission to use and redistribute your contributions as\npart of the project. Head over to to see\nyour current agreements on file or to sign a new one.\n\nYou generally only need to submit a CLA once, so if you've already submitted one\n(even if it was for a different project), you probably don't need to do it\nagain.\n\nAll submissions to Google Open Source projects need to follow Google’s Contributor License Agreement (CLA), in which contributors agree that their contribution is an original work of authorship. This doesn’t prohibit the use of coding assistance tools, but what’s submitted does need to be a contributor’s original creation.\n" }, { "path": "docs_nnx/examples/core_examples.rst", "content": "Core examples\n=============\n\nCore examples are hosted on the GitHub Flax repository in the `examples `__\ndirectory.\n\nEach example is designed to be **self-contained and easily forkable**, while\nreproducing relevant results in different areas of machine learning.\n\nSome of the examples below have a link \"Interactive🕹\" that lets you run them\ndirectly in Colab.\n\nTransformers\n********************\n\n- :octicon:`mark-github;0.9em` `Gemma `__ :\n A family of open-weights Large Language Model (LLM) by Google DeepMind, based on Gemini research and technology.\n\n- :octicon:`mark-github;0.9em` `LM1B `__ :\n Transformer encoder trained on the One Billion Word Benchmark.\n\nToy examples\n********************\n\n`NNX toy examples `__\ndirectory contains a few smaller, standalone toy examples for simple training scenarios.\n" }, { "path": "docs_nnx/examples/gemma.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"# Example: Using pretrained Gemma for inference with Flax NNX\\n\",\n \"\\n\",\n \"This example shows how to use Flax NNX to load the [Gemma](https://ai.google.dev/gemma) open model files and use them to perform sampling/inference for generating text. You will use [Flax NNX `gemma` modules](https://github.com/google/flax/tree/main/examples/gemma) written with Flax and JAX for model parameter configuration and inference.\\n\",\n \"\\n\",\n \"> Gemma is a family of lightweight, state-of-the-art open models based on Google DeepMind’s [Gemini](https://deepmind.google/technologies/gemini/#introduction). Read more about [Gemma](https://blog.google/technology/developers/gemma-open-models/) and [Gemma 2](https://blog.google/technology/developers/google-gemma-2/).\\n\",\n \"\\n\",\n \"You are recommended to use [Google Colab](https://colab.research.google.com/) with access to A100 GPU acceleration to run the code.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Installation\\n\",\n \"\\n\",\n \"Install the necessary dependencies, including `kagglehub`.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 1,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"! pip install --no-deps -U flax\\n\",\n \"! pip install jaxtyping kagglehub treescope\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Download the model\\n\",\n \"\\n\",\n \"To use Gemma model, you'll need a [Kaggle](https://www.kaggle.com/models/google/gemma/) account and API key:\\n\",\n \"\\n\",\n \"1. To create an account, visit [Kaggle](https://www.kaggle.com/) and click on 'Register'.\\n\",\n \"2. If/once you have an account, you need to sign in, go to your ['Settings'](https://www.kaggle.com/settings), and under 'API' click on 'Create New Token' to generate and download your Kaggle API key.\\n\",\n \"3. In [Google Colab](https://colab.research.google.com/), under 'Secrets' add your Kaggle username and API key, storing the username as `KAGGLE_USERNAME` and the key as `KAGGLE_KEY`. If you are using a [Kaggle Notebook](https://www.kaggle.com/code) for free TPU or other hardware acceleration, it has a key storage feature under 'Add-ons' > 'Secrets', along with instructions for accessing stored keys.\\n\",\n \"\\n\",\n \"Then run the cell below.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 2,\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"application/vnd.jupyter.widget-view+json\": {\n \"model_id\": \"2e7cf9f0345845f1a3edc72fa4411eb4\",\n \"version_major\": 2,\n \"version_minor\": 0\n },\n \"text/plain\": [\n \"VBox(children=(HTML(value='
(()=>{ if (customElements.get('treescope-container') === undefined) { class TreescopeContainer extends HTMLElement { constructor() { super(); this.attachShadow({mode: \\\"open\\\"}); this.defns = {}; this.state = {}; } } customElements.define(\\\"treescope-container\\\", TreescopeContainer); } if (customElements.get('treescope-run-here') === undefined) { class RunHere extends HTMLElement { constructor() { super() } connectedCallback() { const run = child => { const fn = new Function(child.textContent); child.textContent = \\\"\\\"; fn.call(this); this.remove(); }; const child = this.querySelector(\\\"script\\\"); if (child) { run(child); } else { new MutationObserver(()=>{ run(this.querySelector(\\\"script\\\")); }).observe(this, {childList: true}); } } } customElements.define(\\\"treescope-run-here\\\", RunHere); } })();
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"transformer = transformer_lib.Transformer.from_params(params)\\n\",\n \"nnx.display(transformer)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Perform sampling/inference\\n\",\n \"\\n\",\n \"Build a Flax NNX [`gemma.Sampler`](https://github.com/google/flax/blob/main/examples/gemma/sampler.py) on top of your model and tokenizer with the right parameter shapes.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 9,\n \"metadata\": {\n \"cellView\": \"form\"\n },\n \"outputs\": [],\n \"source\": [\n \"sampler = sampler_lib.Sampler(\\n\",\n \" transformer=transformer,\\n\",\n \" vocab=vocab,\\n\",\n \")\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"You're ready to start sampling!\\n\",\n \"\\n\",\n \"**Note:** This Flax NNX [`gemma.Sampler`](https://github.com/google/flax/blob/main/examples/gemma/sampler.py) uses JAX’s [just-in-time (JIT) compilation](https://jax.readthedocs.io/en/latest/jit-compilation.html), so changing the input shape triggers recompilation, which can slow things down. For the fastest and most efficient results, keep your batch size consistent.\\n\",\n \"\\n\",\n \"Write a prompt in `input_batch` and perform inference. Feel free to tweak `total_generation_steps` (the number of steps performed when generating a response).\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 10,\n \"metadata\": {\n \"cellView\": \"form\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"Prompt:\\n\",\n \"\\n\",\n \"# Python program for implementation of Bubble Sort\\n\",\n \"\\n\",\n \"def bubbleSort(arr):\\n\",\n \"Output:\\n\",\n \"\\n\",\n \" for i in range(len(arr)):\\n\",\n \" for j in range(len(arr) - i - 1):\\n\",\n \" if arr[j] > arr[j + 1]:\\n\",\n \" swap(arr, j, j + 1)\\n\",\n \"\\n\",\n \"\\n\",\n \"def swap(arr, i, j):\\n\",\n \" temp = arr[i]\\n\",\n \" arr[i] = arr[j]\\n\",\n \" arr[j] = temp\\n\",\n \"\\n\",\n \"\\n\",\n \"# Driver code\\n\",\n \"arr = [5, 2, 8, 3, 1, 9]\\n\",\n \"print(\\\"Unsorted array:\\\")\\n\",\n \"print(arr)\\n\",\n \"bubbleSort(arr)\\n\",\n \"print(\\\"Sorted array:\\\")\\n\",\n \"print(arr)\\n\",\n \"\\n\",\n \"\\n\",\n \"# Time complexity of Bubble sort O(n^2)\\n\",\n \"# where n is the length of the array\\n\",\n \"\\n\",\n \"\\n\",\n \"# Space complexity of Bubble sort O(1)\\n\",\n \"# as it only requires constant extra space for the swap operation\\n\",\n \"\\n\",\n \"\\n\",\n \"# This program uses the bubble sort algorithm to sort the given array in ascending order.\\n\",\n \"\\n\",\n \"```python\\n\",\n \"# This program uses the bubble sort algorithm to sort the given array in ascending order.\\n\",\n \"\\n\",\n \"def bubbleSort(arr):\\n\",\n \" for i in range(len(arr)):\\n\",\n \" for j in range(len(arr) - i - 1):\\n\",\n \" if arr[j] > arr[j + 1]:\\n\",\n \" swap(arr, j, j + 1)\\n\",\n \"\\n\",\n \"\\n\",\n \"def swap(\\n\",\n \"\\n\",\n \"##########\\n\"\n ]\n }\n ],\n \"source\": [\n \"input_batch = [\\n\",\n \" \\\"\\\\n# Python program for implementation of Bubble Sort\\\\n\\\\ndef bubbleSort(arr):\\\",\\n\",\n \" ]\\n\",\n \"\\n\",\n \"out_data = sampler(\\n\",\n \" input_strings=input_batch,\\n\",\n \" total_generation_steps=300, # The number of steps performed when generating a response.\\n\",\n \" )\\n\",\n \"\\n\",\n \"for input_string, out_string in zip(input_batch, out_data.text):\\n\",\n \" print(f\\\"Prompt:\\\\n{input_string}\\\\nOutput:\\\\n{out_string}\\\")\\n\",\n \" print()\\n\",\n \" print(10*'#')\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"You should get a Python implementation of the bubble sort algorithm.\"\n ]\n }\n ],\n \"metadata\": {\n \"jupytext\": {\n \"formats\": \"ipynb,md:myst\"\n },\n \"language_info\": {\n \"codemirror_mode\": {\n \"name\": \"ipython\",\n \"version\": 3\n },\n \"file_extension\": \".py\",\n \"mimetype\": \"text/x-python\",\n \"name\": \"python\",\n \"nbconvert_exporter\": \"python\",\n \"pygments_lexer\": \"ipython3\",\n \"version\": \"3.11.9\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 0\n}\n" }, { "path": "docs_nnx/examples/gemma.md", "content": "---\njupytext:\n formats: ipynb,md:myst\n text_representation:\n extension: .md\n format_name: myst\n format_version: 0.13\n jupytext_version: 1.13.8\n---\n\n# Example: Using pretrained Gemma for inference with Flax NNX\n\nThis example shows how to use Flax NNX to load the [Gemma](https://ai.google.dev/gemma) open model files and use them to perform sampling/inference for generating text. You will use [Flax NNX `gemma` modules](https://github.com/google/flax/tree/main/examples/gemma) written with Flax and JAX for model parameter configuration and inference.\n\n> Gemma is a family of lightweight, state-of-the-art open models based on Google DeepMind’s [Gemini](https://deepmind.google/technologies/gemini/#introduction). Read more about [Gemma](https://blog.google/technology/developers/gemma-open-models/) and [Gemma 2](https://blog.google/technology/developers/google-gemma-2/).\n\nYou are recommended to use [Google Colab](https://colab.research.google.com/) with access to A100 GPU acceleration to run the code.\n\n+++\n\n## Installation\n\nInstall the necessary dependencies, including `kagglehub`.\n\n```{code-cell} ipython3\n! pip install --no-deps -U flax\n! pip install jaxtyping kagglehub treescope\n```\n\n## Download the model\n\nTo use Gemma model, you'll need a [Kaggle](https://www.kaggle.com/models/google/gemma/) account and API key:\n\n1. To create an account, visit [Kaggle](https://www.kaggle.com/) and click on 'Register'.\n2. If/once you have an account, you need to sign in, go to your ['Settings'](https://www.kaggle.com/settings), and under 'API' click on 'Create New Token' to generate and download your Kaggle API key.\n3. In [Google Colab](https://colab.research.google.com/), under 'Secrets' add your Kaggle username and API key, storing the username as `KAGGLE_USERNAME` and the key as `KAGGLE_KEY`. If you are using a [Kaggle Notebook](https://www.kaggle.com/code) for free TPU or other hardware acceleration, it has a key storage feature under 'Add-ons' > 'Secrets', along with instructions for accessing stored keys.\n\nThen run the cell below.\n\n```{code-cell} ipython3\nimport kagglehub\nkagglehub.login()\n```\n\nIf everything went well, it should say `Kaggle credentials set. Kaggle credentials successfully validated.`.\n\n**Note:** In Google Colab, you can instead authenticate into Kaggle using the code below after following the optional step 3 from above.\n\n```\nimport os\nfrom google.colab import userdata # `userdata` is a Colab API.\n\nos.environ[\"KAGGLE_USERNAME\"] = userdata.get('KAGGLE_USERNAME')\nos.environ[\"KAGGLE_KEY\"] = userdata.get('KAGGLE_KEY')\n``` \n\nNow, load the Gemma model you want to try. The code in the next cell utilizes [`kagglehub.model_download`](https://github.com/Kaggle/kagglehub/blob/8efe3e99477aa4f41885840de6903e61a49df4aa/src/kagglehub/models.py#L16) to download model files.\n\n**Note:** For larger models, such as `gemma 7b` and `gemma 7b-it` (instruct), you may require a hardware accelerator with plenty of memory, such as the NVIDIA A100.\n\n**Note:** To avoid 403 error when downloading the model, you need to consent to the license for Gemma models on Kaggle. To do that, open https://www.kaggle.com/models/google/gemma/flax/ in the browser and click on \"Download\" button choosing any version of Gemma model. In the next window you will be proposed to agree with Gemma models usage license. Once, this step is done, you will be able to download the model using the code below.\n\n```{code-cell} ipython3\nfrom IPython.display import clear_output\n\nVARIANT = '2b-it' # @param ['2b', '2b-it', '7b', '7b-it'] {type:\"string\"}\nweights_dir = kagglehub.model_download(f'google/gemma/Flax/{VARIANT}')\nckpt_path = f'{weights_dir}/{VARIANT}'\nvocab_path = f'{weights_dir}/tokenizer.model'\n```\n\n## Python imports\n\n```{code-cell} ipython3\nfrom flax import nnx\nimport sentencepiece as spm\n```\n\nTo interact with the Gemma model, you will use the Flax NNX `gemma` code from [`google/flax` examples on GitHub](https://github.com/google/flax/tree/main/examples/gemma). Since it is not exposed as a package, you need to use the following workaround to import from the Flax NNX `examples/gemma` on GitHub.\n\n```{code-cell} ipython3\nimport sys\nimport tempfile\nwith tempfile.TemporaryDirectory() as tmp:\n # Create a temporary directory and clone the `flax` repo.\n # Then, append the `examples/gemma` folder to the path for loading the `gemma` modules.\n ! git clone https://github.com/google/flax.git {tmp}/flax\n sys.path.append(f\"{tmp}/flax/examples/gemma\")\n import params as params_lib\n import sampler as sampler_lib\n import transformer as transformer_lib\n sys.path.pop();\n```\n\n## Load and prepare the Gemma model\n\nFirst, load the Gemma model parameters for use with Flax.\n\n```{code-cell} ipython3\n:cellView: form\n\nparams = params_lib.load_and_format_params(ckpt_path)\n```\n\nNext, load the tokenizer file constructed using the [SentencePiece](https://github.com/google/sentencepiece) library.\n\n```{code-cell} ipython3\n:cellView: form\n\nvocab = spm.SentencePieceProcessor()\nvocab.Load(vocab_path)\n```\n\nThen, use the Flax NNX [`gemma.transformer.TransformerConfig.from_params`](https://github.com/google/flax/blob/3f3c03b23d4fd3d85d1c5d4d97381a8a2c48b475/examples/gemma/transformer.py#L193) function to automatically load the correct configuration from a checkpoint.\n\n**Note:** The vocabulary size is smaller than the number of input embeddings due to unused tokens in this release.\n\n```{code-cell} ipython3\ntransformer = transformer_lib.Transformer.from_params(params)\nnnx.display(transformer)\n```\n\n## Perform sampling/inference\n\nBuild a Flax NNX [`gemma.Sampler`](https://github.com/google/flax/blob/main/examples/gemma/sampler.py) on top of your model and tokenizer with the right parameter shapes.\n\n```{code-cell} ipython3\n:cellView: form\n\nsampler = sampler_lib.Sampler(\n transformer=transformer,\n vocab=vocab,\n)\n```\n\nYou're ready to start sampling!\n\n**Note:** This Flax NNX [`gemma.Sampler`](https://github.com/google/flax/blob/main/examples/gemma/sampler.py) uses JAX’s [just-in-time (JIT) compilation](https://jax.readthedocs.io/en/latest/jit-compilation.html), so changing the input shape triggers recompilation, which can slow things down. For the fastest and most efficient results, keep your batch size consistent.\n\nWrite a prompt in `input_batch` and perform inference. Feel free to tweak `total_generation_steps` (the number of steps performed when generating a response).\n\n```{code-cell} ipython3\n:cellView: form\n\ninput_batch = [\n \"\\n# Python program for implementation of Bubble Sort\\n\\ndef bubbleSort(arr):\",\n ]\n\nout_data = sampler(\n input_strings=input_batch,\n total_generation_steps=300, # The number of steps performed when generating a response.\n )\n\nfor input_string, out_string in zip(input_batch, out_data.text):\n print(f\"Prompt:\\n{input_string}\\nOutput:\\n{out_string}\")\n print()\n print(10*'#')\n```\n\nYou should get a Python implementation of the bubble sort algorithm.\n" }, { "path": "docs_nnx/examples/index.rst", "content": "Examples\n========\n\n.. toctree::\n :maxdepth: 2\n\n gemma\n core_examples\n\n\n" }, { "path": "docs_nnx/faq.rst", "content": "Frequently Asked Questions (FAQ)\n================================\n\nThis is a collection of answers to frequently asked questions (FAQ). You can contribute to the Flax FAQ by starting a new topic in `GitHub Discussions `__.\n\nWhere to search for an answer to a Flax-related question?\n*********************************************************\n\nThere are a number of official Flax resources to search for information:\n\n- `Flax Documentation on ReadTheDocs `__ (this site): Use the `search bar `__ or the table of contents on the left-hand side.\n- `google/flax GitHub Discussions `__: Search for an existing topic or start a new one. If you can't find what you're looking for, feel free to ask the Flax team or community a question.\n- `google/flax GitHub Issues `__: Use the search bar to look for an existing issue or a feature request, or start a new one.\n\nHow to take the derivative with respect to an intermediate value (using :code:`Module.perturb`)?\n************************************************************************************************\n\nTo take the derivative(s) or gradient(s) of the output with respect to a hidden/intermediate activation inside a model layer, you can use :meth:`flax.linen.Module.perturb`. You define a zero-value :class:`flax.linen.Module` \"perturbation\" parameter – :code:`perturb(...)` – in the forward pass with the same shape as the intermediate activation, define the loss function with :code:`'perturbations'` as an added standalone argument, perform a JAX derivative operation with :code:`jax.grad` on the perturbation argument.\n\nFor full examples and detailed documentation, go to:\n\n- The :meth:`flax.linen.Module.perturb` API docs\n- The `Extracting gradients of intermediate values `_ guide\n- `Flax GitHub Discussions #1152 `__\n\nIs Flax Linen :code:`remat_scan()` the same as :code:`scan(remat(...))`?\n************************************************************************\n\nFlax :code:`remat_scan()` (:meth:`flax.linen.remat_scan()`) and :code:`scan(remat(...))` (:meth:`flax.linen.scan` over :meth:`flax.linen.remat`) are not the same, and :code:`remat_scan()` is limited in cases it supports. Namely, :code:`remat_scan()` treats the inputs and outputs as carries (hidden states that are carried through the training loop). You are recommended to use :code:`scan(remat(...))`, as typically you would need the extra parameters, such as ``in_axes`` (for input array axes) or ``out_axes`` (output array axes), which :meth:`flax.linen.remat_scan` does not expose.\n\nWhat are the recommended training loop libraries?\n*************************************************\n\nConsider using CLU (Common Loop Utils) `google/CommonLoopUtils `__. To get started, go to this `CLU Synopsis Colab `__. You can find answers to common questions about CLU with Flax on `google/flax GitHub Discussions `__.\n\nCheck out the official `google/flax Examples `__ for examples of using the training loop with (CLU) metrics. For example, this is `Flax ImageNet's train.py `__.\n\nFor computer vision research, consider `google-research/scenic `__. Scenic is a set of shared light-weight libraries solving commonly encountered tasks when training large-scale vision models (with examples of several projects). Scenic is developed in JAX with Flax. To get started, go to the `README page on GitHub `__." }, { "path": "docs_nnx/flip/0000-template.md", "content": "- Start Date: (fill me in with today's date, YYYY-MM-DD)\n- FLIP PR: [#0000](https://github.com/google/flax/pull/0000)\n- FLIP Issue: [#0000](https://github.com/google/flax/issues/0000)\n\n(Below sections are just a possible structure - please adapt to your FLIP.)\n\n# Summary\n[summary]: #summary\n\nOne paragraph explanation of the FLIP.\n\n# Motivation\n[motivation]: #motivation\n\nWhy are we doing this? What use cases does it support? What is the expected outcome?\n\n# Implementation\n[implementation]: #implementation\n\nThe technical part.\n\n# Discussion\n[discussion]: #discussion\n\nSummarize the discussion from the original issue and from the pull request.\n" }, { "path": "docs_nnx/flip/1009-optimizer-api.md", "content": "- Start Date: 2021-02-08\n- FLIP PR: [#1011](https://github.com/google/flax/pull/1011)\n- FLIP Issue: [#1009](https://github.com/google/flax/issues/1009)\n\nTable of contents:\n\n- [Summary]\n- [Motivation]\n- [Using Optax]\n - [Gradient Transformations]\n - [Optax Training Step]\n - [Multi Optimizer]\n - [Train State]\n- [Previous API]\n - [Optimizer and OptimizerDef]\n - [Previous Training Step]\n- [Update Plan]\n- [Appendix]\n - [Setup Code]\n\n# Summary\n[Summary]: #summary\n\nThis FLIP proposes to replace our current `flax.optim` API (referred to as\n[previous API] in this document) with [Optax], DeepMind's optimizer library.\n\n# Motivation\n[motivation]: #motivation\n\nOur current API (referred to as [previous API] in this document) uses a pattern\nwhere an `Optimizer` dataclass is created from a pytree of `target` variables\nand from an `OptimizerDef` that defines how to update optimizer state,\nhyperparameters, and target variables. This pattern is relatively complex for\nimplementing a simple optimizer, while being quite verbose in the typical Linen\ntrain step (especially when using mutable state collections).\n\nThis package `flax.optim` contains some optimizers, but that list is far from\nexhaustive and ideally we would instead use JAX optimizers from a dedicated PyPi\npackage.\n\nDeepMind already has a dedicated library — [Optax] — that implements a wide\nrange of interesting optimizers and provides a framework to compose new\noptimizers from reusable gradient transformations.\n\n[Optax]: https://github.com/deepmind/optax\n\n# Using Optax\n[Using Optax]: #using-optax\n\n## Gradient Transformations\n[Gradient Transformations]: #gradient-transformations\n\nWhile [Optax] does provide predefined optimizers (like `optax.adam`, or\n`optax.sgd` with momentum), it is really a library of *gradient transformations*\nand the idiomatic way of instantiating an optimizer is by providing a\ncombination of these gradient transformations. To emulate the momentum\noptimizer from the example when using the [previous API] we would write:\n\n```python\nimport optax\n\ntx = optax.chain(\n optax.trace(decay=0.9, nesterov=False),\n optax.scale_by_schedule(lambda step: -get_learning_rate(step)),\n)\n```\n\nRemarks:\n\n- Above gradient transformation would be equivalent with the example under\n [Optimizer and OptimizerDef] where we define a Momentum optimizer without\n Nesterov momentum (note that the `beta` parameter corresponds to the `decay`\n parameter of the `optax.trace()` transformation, and the learning rate is\n applied in a second chained transformation).\n- Note that hyper parameters like `decay` or `nesterov` only exist in the inner\n scope of the higher order functions returning the `GradientTransformation`.\n Such a gradient transformation is currently defined as a `NamedTuple` of the\n `init()` and the `update()` function. In principle this pattern could be\n extended to also store hyperparameters, maybe a point to discuss on the\n [Optax] repo.\n- We can use a `get_learning_rate()` that returns the learning rate depending on\n the step number when defining the Optax gradient update transformation. Above\n code illustrates how this can be a drop-in replacement for a function we also\n use in our [previous training step], where this update function already exists\n (notice how we need to invert the sign because we add the gradient update to\n the parameters). In addition, you can use\n [`inject_hyperparams()`](https://github.com/deepmind/optax/pull/48) to\n schedule arbitrary hyper parameters with Optax.\n\n## Optax Training Step\n[Optax Training Step]: #optax-training-step\n\n```python\n@functools.partial(jax.jit, static_argnums=(4, 5))\ndef train_step(opt_state, variables, inputs, labels, apply_fn, tx_update_fn):\n\n def loss_fn(params):\n logits, new_model_state = apply_fn(\n {**variables, 'params': params}, inputs, mutable=['batch_stats'])\n loss = xent_loss(logits, labels)\n return loss, new_model_state\n\n variables, params = variables.pop('params')\n (loss, new_model_state), grads = jax.value_and_grad(loss_fn, has_aux=True)(\n params)\n updates, new_opt_state = tx_update_fn(grads, opt_state, params)\n new_params = optax.apply_updates(params, updates)\n new_variables = {**variables, **new_model_state, 'params': new_params}\n return new_opt_state, new_variables, loss\n\n\nopt_state = tx.init(variables['params'])\nfor batch in ds.as_numpy_iterator():\n opt_state, variables, loss = train_step(\n opt_state, variables, batch['image'], batch['label'], model.apply,\n tx.update)\n print(loss)\n```\n\nRemarks:\n\n- Since `tx.update()` only transforms the gradient, we still need to call\n `optax.apply_updates()` to apply these transformed gradients to the\n parameters.\n- Compared with the [previous API], we can now keep the entire `variables`\n including the `params` as an input and output to the `train_step()`.\n- Splitting `params` from `variables` is still necessary inside the train step\n because we only want to compute gradients with respect to `params` and not the\n entire `variables`.\n- We can still log internal optimizer state, such as the learning rate, as long\n as Optax transformations expose that information in their respective state.\n For example, `optax.scale_by_schedule()` currently only exposes\n `opt_state.count` but could easily be extend to also expose the `step_size`.\n The same is true for internal optimizer states that change over time.\n\n## Multi Optimizer\n[Multi Optimizer]: #multi-optimizer\n\nThe [previous API] defined `flax.optim.MultiOptimizer` for processing different\nparts of the parameter tree with different optimizers:\n\n```python\nbiases_traversal = flax.optim.ModelParamTraversal(\n lambda path, _: path.endswith('/bias'))\nnot_biases_traversal = flax.optim.ModelParamTraversal(\n lambda path, _: not path.endswith('/bias'))\n\noptimizer_def = flax.optim.MultiOptimizer(\n (biases_traversal, flax.optim.GradientDescent(learning_rate=0.1)),\n (not_biases_traversal, flax.optim.GradientDescent(learning_rate=0.05)),\n)\n```\n\nNote how we first define a traversal that selects parameters based on their\npath (which is the concatenation of module scopes and variable name), and then\ncreate a `MultiOptimizer` that binds a different optimizer for each of these\nseparate traversals.\n\nOptax has recently implemented `optax.masked()` that can be used for specifying\ngradient transformations that only applied to a subset of the gradients:\n\n```python\ndef flattened_traversal(fn):\n def mask(data):\n flat = traverse_util.flatten_dict(data)\n return traverse_util.unflatten_dict({k: fn(k, v) for k, v in flat.items()})\n return mask\n\ntx = optax.chain(\n optax.masked(optax.sgd(learning_rate=0.1),\n mask=flattened_traversal(lambda path, _: path[-1] == 'bias')),\n optax.masked(optax.sgd(learning_rate=0.05),\n mask=flattened_traversal(lambda path, _: path[-1] != 'bias')),\n)\n```\n\n## Train State\n[Train State]: #train-state\n\nIn Flax it is common to hand around a `TrainState` object that can then be\nused for checkpointing. This simplifies the above [Optax training step] a bit by\nreducing the number of arguments and getting rid of the `static_argnums`.\n\nWe can define a `TrainState` dataclass that wraps the common pattern of updating\nthe optimizer state and parameters by applying the gradients.\n\n```python\n# Small helper class in flax.training\nclass TrainState(flax.struct.PyTreeNode):\n step: int\n apply_fn: Callable = flax.struct.field(pytree_node=False)\n params: flax.core.FrozenDict[str, Any]\n tx: optax.GradientTransformation = flax.struct.field(pytree_node=False)\n opt_state: optax.OptState\n\n def apply_gradients(self, *, grads, **kwargs):\n updates, new_opt_state = self.tx.update(\n grads, self.opt_state, self.params)\n new_params = optax.apply_updates(self.params, updates)\n return self.replace(\n step=self.step + 1,\n params=new_params,\n opt_state=new_opt_state,\n **kwargs,\n )\n\n @classmethod\n def create(cls, *, apply_fn, params, tx, **kwargs):\n opt_state = tx.init(params)\n return cls(\n step=0,\n apply_fn=apply_fn,\n params=params,\n tx=tx,\n opt_state=opt_state,\n **kwargs,\n )\n```\n\nUsers can then derive from this dataclass and add more fields, for example\nmutable model state:\n\n```python\nfrom flax.training import train_state\n\nclass TrainState(train_state.TrainState):\n batch_stats: flax.core.FrozenDict[str, Any]\n```\n\nWith this the [Optax Training Step] becomes:\n\n```python\n@jax.jit\ndef train_step(state, inputs, labels):\n\n def loss_fn(params):\n outputs, new_model_state = state.apply_fn(\n {'params': params, 'batch_stats': state.batch_stats},\n inputs,\n mutable=['batch_stats'])\n loss = xent_loss(outputs, labels)\n return loss, new_model_state\n\n (loss, new_model_state), grads = jax.value_and_grad(\n loss_fn, has_aux=True)(state.params)\n new_state = state.apply_gradients(\n grads=grads,\n batch_stats=new_model_state['batch_stats'],\n )\n\n return new_state, loss\n\n\nstate = TrainState.create(\n apply_fn=model.apply,\n params=variables['params'],\n tx=tx,\n batch_stats=variables['batch_stats'],\n)\nfor batch in ds.as_numpy_iterator():\n state, loss = train_step(state, batch['image'], batch['label'])\n```\n\nThe train step without mutable state reduces to:\n\n```python\n@jax.jit\ndef train_step(state, inputs, labels):\n\n def loss_fn(params):\n outputs = state.apply_fn({'params': params}, inputs)\n loss = xent_loss(outputs, labels)\n return loss\n\n loss, grads = jax.value_and_grad(loss_fn)(state.params)\n new_state = state.update(grads=grads)\n\n return new_state, loss\n\n\nstate = flax.training.TrainState.create(\n apply_fn=model.apply,\n params=variables['params'],\n tx=tx,\n)\nfor batch in ds.as_numpy_iterator():\n state, loss = train_step(state, batch['image'], batch['label'])\n```\n\nRemarks:\n\n- It is a common pattern in Flax training loops to have a `TrainState` dataclass\n that is updated with new state after every step.\n- The simple solution proposed in `flax.training.train_state` an be extended\n with additional data, but advanced usecases (e.g. multiple different models\n and/or optimizers) are not supported. Users should instead fork the dataclass\n and re-implement it to their needs.\n- As opposed to the `Optimizer` abstraction in the [previous API], the\n `TrainState` now directly contains the `.params`, without having to to through\n `.optimizer`\n\n# Previous API\n[previous API]: #previous-api\n\n## Optimizer and OptimizerDef\n[Optimizer and OptimizerDef]: #optimizer-and-optimizerdef\n\nThe optimizer itself would be implemented by creating a new class derived\nfrom `OpimizerDef`:\n\n```python\n# flax/optim/momentum.py\n\n@flax.struct.dataclass\nclass _MomentumHyperParams:\n learning_rate: jnp.ndarray\n beta: jnp.ndarray\n\n\n@flax.struct.dataclass\nclass _MomentumParamState:\n momentum: np.ndarray\n\n\nclass Momentum(flax.optim.OptimizerDef):\n\n def __init__(self, learning_rate=None, beta=0.9):\n super().__init__(\n _MomentumHyperParams(learning_rate, beta)\n )\n\n def init_param_state(self, param):\n return _MomentumParamState(jnp.zeros_like(param))\n\n def apply_param_gradient(self, step, hyper_params, param, state, grad):\n del step\n assert hyper_params.learning_rate is not None\n new_momentum = state.momentum * hyper_params.beta + grad\n new_params = param - hyper_params.learning_rate * new_momentum\n return new_params, _MomentumParamState(new_momentum)\n```\n\nRemarks:\n\n- Note the relationship between `OptimizerDef` and `Optimizer` : When the\n function `Optimizer.apply_gradient()` is called from the user code, it calls\n into `OptimizerDef.apply_gradient()` (among other things) which in turn will\n call `OptimizerDef.apply_param_gradient()` (implemented by subclasses of\n `OptimizerDef`).\n- The functions `init_param_state()` and `apply_param_gradient()` are called\n for every leaf in the params/grads pytree. This makes it possible to write the\n calculations directly without `jax.tree_util.tree_map()`.\n- The interface was defined in pre-Linen without the distinction of `params` vs.\n other collections in `variables` in mind. The original API was elegant because\n one only needed to pass around the optimizer, which included the parameters,\n optimizer state, optimizer hyperparameters, and a reference to the\n `OptimizerDef` to perform the param/state update.\n\n## Previous Training Step\n[Previous Training Step]: #previous-training-step\n\nAn optimizer would first be constructed from its definition and the pytree of\ntarget params:\n\n```python\noptimizer_def = flax.optim.Momentum(learning_rate=0.1, beta=0.9)\noptimizer = optimizer_def.create(variables['params'])\n```\n\nThen, the target variables would optimized in the train step (assuming a single\nnon-params collection \"batch_stats\"):\n\n```python\ndef make_train_step(apply_fn):\n @jax.jit\n def train_step(optimizer, batch_stats, inputs, labels):\n\n def loss_fn(params):\n variables = {'params': params, 'batch_stats': batch_stats}\n logits, new_model_state = apply_fn(\n variables, inputs, mutable=['batch_stats'])\n loss = xent_loss(logits, labels)\n return loss, new_model_state['batch_stats']\n\n (loss, new_batch_stats), grad = jax.value_and_grad(loss_fn, has_aux=True)(\n optimizer.target)\n lr = get_learning_rate(step)\n new_optimizer = optimizer.apply_gradient(grad, learning_rate=lr)\n return new_optimizer, new_batch_stats, loss\n\n return train_step\n\n\nbatch_stats = variables['batch_stats']\ntrain_step = make_train_step(model.apply)\nfor step, batch in enumerate(ds)\n optimizer, batch_stats, loss = train_step(\n optimizer, batch_stats, batch['image'], batch['label'])\n```\n\nRemarks:\n\n- Notice how `optimizer.apply_gradient()` can take additional arguments to\n update hyperparameters, such as learning rate from an independent function\n `get_learning_rate()` in this case.\n\n\n# Update Plan\n[Update Plan]: #update-plan\n\n1. Finalize discussions on this FLIP\n2. Add [equivalence tests] to Optax that guarantee that existing `flax.optim`\n optimizers return identical values with corresponding `optax` optimizers.\n3. Update examples to use Optax and verify that they reach the same final\n performance with the same computational cost.\n4. Port missing optimizers to Optax (e.g. Adafactor) - and verify above points.\n5. Update all documentation (including README, Flax guided tour, HOWTOs, ...) to\n talk exclusively about Optax optimizers.\n6. Create a transition guide for updating users from `flax.optim` to using\n Optax. This transition guide should also point to Optax's [equivalence tests]\n and the pull requests updating the examples.\n7. Mark optimizers in `flax.optim` as deprecated.\n\n[equivalence tests]: https://github.com/deepmind/optax/blob/master/optax/_src/equivalence_test.py\n\nNote that all current Flax examples use an optimizer that is already available\nin Optax:\n\n| Example | Flax | Optax | Comments |\n| -------- | -------------- | ----------- | ----------------------------------- |\n| imagenet | optim.Momentum | optax.sgd | DynamicScale can be used unchanged. |\n| mnist | optim.Momentum | optax.sgd | |\n| nlp_seq | optim.Adam | optax.adamw | |\n| pixelcnn | optim.Adam | optax.adam | |\n| ppo | optim.Adam | optax.adam | |\n| seq2seq | optim.Adam | optax.adam | |\n| vae | optim.Adam | optax.adam | |\n| wmt | optim.Adam | optax.adamw | |\n\n(Flax's Adam implementation has an optional parameter for weight decay, but in\nOptax Adam with and without weight decay are two different aliases.)\n\n# Appendix\n[Appendix]: #appendix\n\n## Setup Code\n[Setup Code]: #setup-code\n\nThe following setup code can be used for running the code snippets in this\nFLIP:\n\n```python\nimport functools\nfrom typing import Callable, Sequence\n\nimport jax\nimport jax.numpy as jnp\nimport flax\nimport flax.linen as nn\nimport tensorflow as tf\nimport tensorflow_datasets as tfds\n\n\ndef pp(features):\n return {\n 'image': tf.cast(features['image'], tf.float32) / 255 - 0.5,\n 'label': features['label'],\n }\n\n\nclass Model(nn.Module):\n\n @nn.compact\n def __call__(self, inputs):\n x = inputs.reshape([inputs.shape[0], -1])\n x = nn.normalization.BatchNorm(True)(x)\n x = nn.Dense(10)(x)\n x = nn.log_softmax(x)\n return x\n\n\ndef onehot(labels, num_classes, on_value=1.0, off_value=0.0):\n x = (labels[..., None] == jnp.arange(num_classes)[None])\n x = jax.lax.select(\n x, jnp.full(x.shape, on_value), jnp.full(x.shape, off_value))\n return x.astype(jnp.float32)\n\n\ndef xent_loss(logits, labels):\n return -jnp.sum(\n onehot(labels, num_classes=10) * logits) / labels.size\n\n\ndef get_learning_rate(step):\n return 0.1\n\n\nmodel = Model()\nrng = jax.random.key(0)\nds = tfds.load('mnist')['train'].take(160).map(pp).batch(16)\nbatch = next(iter(ds))\nvariables = model.init(rng, jnp.array(batch['image'][:1]))\njax.tree_util.tree_map(jnp.shape, variables)\n```\n" }, { "path": "docs_nnx/flip/1777-default-dtype.md", "content": "# FLIP: Default dtypes\n\n\n- Start Date: 2022-01-11\n- FLIP PR: [#1776](https://github.com/google/flax/pull/1776)\n- FLIP Issue: [#1777](https://github.com/google/flax/issues/1777)\n- Status: Implemented\n\n\n## Summary\n\nThis FLIP proposes to replace the default dtype which is currently fixed to float32, and instead use the JAX type promotion results to derive a default dtype from the input and parameters of a layer.\n\n\n## Motivation\n\nCurrently, Linen Modules always produce `module.dtype` (defaults to float32) outputs regardless of input and parameter dtypes. Half-precision types like float16 and bfloat16 are supported by explicitly passing the half-precision type to each Module. The way this is currently implemented is that each Module has a dtype argument with float32 as the default value. The layer guarantees that this dtype will be the return type of the result returned by `__call__`.\n\nThe current behavior is problematic and results in silent bugs, especially for dtypes that do not fit inside float32 (complex, float64). Also, the Linen dtype behavior is significantly different from how NumPy and by extension JAX handle dtypes.\n\n\n### Dtypes in JAX\n\nJAX uses a NumPy-inspired [dtype promotion](https://github.com/jax-ml/jax/blob/main/jax/_src/dtypes.py) mechanism as explained [here](https://jax.readthedocs.io/en/latest/type_promotion.html?highlight=lattice#type-promotion-semantics). The type promotion rules are summarized by the following type lattice:\n\n![JAX type promotion lattice](https://jax.readthedocs.io/en/latest/_images/type_lattice.svg)\n\n\n## Dtypes in Linen\n\nBesides input arguments, state and in particular parameters could affect dtype promotion. For example: we might feed a float64 input to a Dense layer with float32 parameters. Currently, the result would be truncated to float32. If the input is a complex number the result is even worse because the imaginary part will be silently dropped when casting to float32.\n\nBy using the dtype promotion rules already available in JAX we can avoid this issue. A public API is available called `jax.numpy.result_dtype(*args)`, which returns the dtype that JAX would promote the given arguments to, in accordance with the type promotion lattice. For Linen layers the arguments would be the layer inputs together with the parameters. For example, for a linear layer this would be inputs, kernel, and bias.\n\nNote that there is also a `param_dtype` attribute in standard Linen Modules that also defaults to flaot32. This behavior is left untouched and encodes the common case of having float32 parameters.\nThere are a few reasons why float32 is almost always the correct dtype for parameters:\n1. Storing weights in half-precision often leads to underflow during optimization.\n2. Double precision is rarely used because it severely slows down modern accelerators (GPU, TPU). Therefore, such a cost should be explicitly opted-in for.\n3. Complex Modules are relatively uncommon. Even within complex networks, the complex inputs can be projected with a real matrix.\n\n\n# Implementation\n\nA simplified example implementation:\n\n\n```python\ndef promote_arrays(*xs, dtype):\n if dtype is None:\n dtype = jnp.result_type(*jax.tree_util.tree_leaves(xs))\n return jax.tree_util.tree_map(lambda x: jnp.asarray(x, dtype), xs)\n\nDtype = Any\nclass Dense(nn.Module):\n features: int\n kernel_init: Callable\n bias_init: Callable\n dtype: Optional[Dtype] = None\n param_dtype: Dtype = jnp.float32\n\n @nn.compact\n def __call__(self, x):\n kernel = self.param(\"kernel\",\n self.kernel_init,\n (x.shape[-1], self.features), self.param_dtype)\n bias = self.param(\"bias\", self.bias_init, (self.features,), self.param_dtype)\n x, kernel, bias = promote_arrays(x, kernel, bias, dtype=self.dtype)\n return x @ kernel + bias\n```\n\n\n## Half-precision dtypes\n\nSome layers don’t work with half-precision dtypes internally. For example: The normalization layers currently compute mean and variance in float32 even when a half-precision dtype is specified to avoid numerical issues. We can replicate this behavior by calling result_dtype with a dummy argument that has the minimum precision for the sub computation to work correctly.\n\n\n## Backward compatibility\n\nThis proposal causes some layers to behave differently in cases where the dtype is not specified to a Linen Module. By default, parameters are in float32. Therefore, passing in half or float32 precision inputs will cause a float32 dtype and no functional differences with current behavior.\n\nWhen passing complex or float64 precision, the result will no longer truncate the imaginary component or the precision. The silent truncation is problematic and has caused [user complaints](https://github.com/google/flax/issues/805#issuecomment-981468837). Therefore, this change can be considered a bugfix.\n\nThus, although this proposal strictly speaking changes behavior it is unlikely to cause problems for users. There are 2 exceptions to this which should be rare and easy to fix:\n1. A user relies on the enforced float32 to downcast a double precision value.\n2. A user relies on the float32 to explicitly upcast a half precision value even though the weights are in half precision.\n\n\n## Corner cases\n\nIn this section we describe corner cases where the implementation of the proposal is not obvious. The two main concerns are how complex numbers are handled in existing layers and how to determine the dtype of state variables.\n\n**Autoregressive decoding cache**\n\nCurrently, only attention implements autoregressive caching and the stored key and value mirror the dtype of the key and value passed to the layer. Forcing the cache dtype to be the same as the output dtype could result in reduced precision during cached decoding vs uncached. This seems undesirable. Decision: keep the current behavior.\n\n**Batch statistics**\n\nBatchNorm layers are often used with a half precision output dtype. However, calculating statistics is by default always done in float32 to avoid numerical precision issues and over/underflow for float16. With float64 this would actually cause a downcast so we should now use `np.promote_types(float32, dtype)` such that the precision is at least float32. The running batch statistics will be stored with the same dtype for consistency.\n\n**Complex number support**\n\nCurrently, our complex number support is brittle because the default behavior is to truncate the output to the real part. This issue will be fixed by the automatic type promotion proposed in this FLIP. However, some layers require some additional thought to extend to complex numbers correctly:\n\n1. Normalization layers use the complex conjugate to calculate norms instead of normal squaring.\n2. Attention: It’s not exactly clear how the dot product and softmax are defined in this case. Raise an error on complex inputs.\n3. Recurrent layers: might require special gating / activation functions to function correctly, but these can be specified by the user.\n\n\n# Discussion\n\nSummarizing the main points from the discussion:\n\n\n## Consider implicit complex truncation an error\n\nQ:\nI'm wondering if we should always raise an error if one of the xs tree leaves is complex but dtype is not. Users should maybe remove imaginary part by themselves if that's really what they want to do.\n(Maybe it's a contrived example, but I can imagine cases where layers have their dtype set by parent modules based on assumptions without complex numbers in mind)\n\nA:\nThis is worth considering in a follow-up CL but this might as well be solved in JAX directly where the safeguard would apply more generally. In NumPy this was also considered but abandoned because it is not backwards compatible.\n\n\n## Dtype attribute names\n\nQ:\nAre the dtype and param_dtype arguments confusing? In particular, should dtype perhaps be called output_dtype to make the difference between the two dtypes more explicit?\n\nA:\nThis would be a large and orthogonal change wrt to this proposal so leaving it out for now.\nAlso, this breaks with the standard dtype argument in NumPY/JAX.\nAlthough dtype indeed constrains the output dtype it is also a hint for the dtype we would like the computation to happen in.\n\n" }, { "path": "docs_nnx/flip/2396-rnn.md", "content": "# RNN Flip\n\n- Start Date: 2022-08-18\n- FLIP PR: [#2604](https://github.com/google/flax/pull/2604)\n- FLIP Issue: [#2396](https://github.com/google/flax/issues/2396)\n- Authors: Jasmijn Bastings (@bastings) and Cristian Garcia (@cgarciae)\n\n## Summary\nThis FLIP adds support for higher-level recurrent layers (RNN, GRU, LSTM) that can help users process input sequences using the recurrent cells already available in Flax.\n\n## Motivation\nImplementing well known recurrent architectures is tricky and prone to user errors, even a simple LSTM layers involves the manual creation and handling of the carry/memory and correctly setting up `nn.scan`:\n\n```python\n@nn.compact\ndef __call__(self, x):\n LSTM = nn.scan(\n nn.LSTMCell, variable_broadcast=\"params\", split_rngs={\"params\": False}\n )\n carry = LSTM.initialize_carry(\n jax.random.key(0), batch_dims=x.shape[:1], size=self.hidden_size\n )\n carry, x = LSTM()(carry, x)\n return x\n```\nSlightly more complicated cases involving padding like in the [seq2seq](https://github.com/google/flax/blob/main/examples/seq2seq/models.py) example require even more work but couple potentially be simplified to a couple of lines with the right abstractions. We propose providing users with clean, correct, and efficient abstractions to use recurrent cells.\n\n## Requirements\n\n* **Masking**: We need to support a batch of sequences that contain padding at the end of each sequence.\n * We do not intend to support non-contiguous padding, i.e. padding that is not at the end of a sequence, for performance reasons, except in the case of packing (see below).\n* **Bidirectionality**: The ability to process a sequence in both the forward and reverse directions, respecting padding (i.e., the reverse direction should start with the actual inputs, not with padding values).\n* **Performance**: The proposed classes should be benchmarked to provide the best performance in terms of step time and/or memory use.\n* **Recurrent Dropout**: Support for recurrent dropout in cells (e.g. dropout on the state of the cell).\n\n## Implementation\n### High-level structure\n\nWe propose to have these 3 levels of abstraction:\n\n* **Cells (unchanged)**: all RNNCellBase subclasses such as LSTMCell and GRUCell, these implement the stepwise logic. These already exist in Flax today.\n* **Layers (new)**: a class (RNN) that takes a cell and scans over a sequence respecting possible padding values and optionally also allows packed sequences.\n* **Bidirectional (new)**: a single class that takes a forward and a backward RNN instance and correctly processes the input sequence in both directions and merges the results.\n\n### Example of proposed API\nWe start with a code example of what you could do with the proposed API, and then we discuss the API in detail below.\n\n```python\ncell = nn.LSTMCell()\n# Encodes a batch of input sequences.\ncarry, outputs = nn.RNN(cell, cell_size)(inputs, seq_lengths)\n```\n\nA Bidirectional layer with a LSTM RNNs for the forward and backward directions respectively would look like this:\n\n```python\nforward_rnn = nn.RNN(nn.LSTMCell(), cell_size=32)\nbackward_rnn = nn.RNN(nn.LSTMCell(), cell_size=32)\n# Bidirectional combinator.\nbi_rnn = nn.Bidirectional(forward_rnn, backward_rnn)\n# Encodes a batch of input sequences in both directions.\ncarry, outputs = bi_rnn(inputs, seq_lengths)\n```\n\nNext we will discuss `RNN`, `Bidirectional`, and proposed changes to `RNNCellBase`.\n\n### RNNBase\nThe `RNNBase` class serves as a base class for the `RNN` class, it specifies\nthe API that all RNN layers should implement to be compatible with the `Bidirectional`.\n`RNNBase` contains the `__call__` and `flip_sequences` methods:\n\n```python\nclass RNNBase(Protocol):\n def __call__(\n self,\n inputs: jax.Array,\n *,\n initial_carry: Optional[Carry] = None,\n init_key: Optional[random.KeyArray] = None,\n seq_lengths: Optional[Array] = None,\n return_carry: Optional[bool] = None,\n time_major: Optional[bool] = None,\n reverse: Optional[bool] = None,\n keep_order: Optional[bool] = None,\n ) -> Union[Output, Tuple[Carry, Output]]:\n ...\n```\nWhere:\n\n* `inputs`: the input sequence.\n* `initial_carry`: the initial carry, if not provided it will be initialized\n using the cell's :meth:`RNNCellBase.initialize_carry` method.\n* `init_key`: a PRNG key used to initialize the carry, if not provided\n ``jax.random.key(0)`` will be used. Most cells will ignore this\n argument.\n* `seq_lengths`: an optional integer array of shape ``(*batch)`` indicating\n the length of each sequence, elements whose index in the time dimension\n is greater than the corresponding length will be considered padding and\n will be ignored.\n* `return_carry`: if ``return_carry=False`` (default) only the output sequence is returned,\n else it will return a tuple of the final carry and the output sequence.\n* `time_major`: if ``time_major=False`` (default) it will expect inputs with shape\n ``(*batch, time, *features)``, else it will expect inputs with shape ``(time, *batch, *features)``.\n* `reverse`: if ``reverse=False`` (default) the sequence is\n processed from left to right and returned in the original order, else it will be processed\n from right to left, and returned in reverse order. If ``seq_lengths`` is passed,\n padding will always remain at the end of the sequence.\n* `keep_order`: if ``keep_order=True``, when ``reverse=True``\n the output will be reversed back to the original order after processing, this is\n useful to align sequences in bidirectional RNNs. If ``keep_order=False`` (default),\n the output will remain in the order specified by ``reverse``.\n* `Returns`: if ``return_carry=False`` (default) only the output sequence is returned,\nelse it will return a tuple of the final carry and the output sequence.\n\n### RNN\nThe `RNN` module inherits from `RNNBase`, it main function is to apply an `RNNCellBase` instance over a batch of input sequences, it can be used with any type of cell (e.g., `GRUCell`, `LSTMCell`, etc). It accepts the following parameters:\n\n```python\nclass RNN(RNNBase):\n cell: RNNCellBase,\n cell_size: int | Tuple[int, ...]\n time_axis: int = -2,\n variable_axes = FrozenDict(),\n variable_broadcast: CollectionFilter = 'params'\n variable_carry: CollectionFilter = False\n split_rngs = FrozenDict({'params': False})\n # implement RNNBase\n ...\n```\n\nAttributes like `variable_axes`, `variable_broadcast`, `variable_carry`, and `split_rngs` are directly passed to `nn.scan`, their default values are set such that common cells like `LSTMCell` and `GRUCell` work out of the box.\n\n### Masking\n`seq_lengths` is defined as an integer array of shape `(*batch,)` indicating the length of each sequence.\n\n
Discussion\n\nThere are various masking formats found in other frameworks, here are some of the most popular ones:\n\n* **Binary masking**: specifies per-sample and timestep whether that data point should be included or not in the computation, it can be non-contigous (e.g., [1, 1, 0, 1]). This is used by Keras.\n* **Sequence length masking**: specifies per-sample the number of non-padding examples contained in the sequence, any padding contained in the sequence should be stacked at the end. This is used by FlaxFormer.\n* **Segmentation Mask**: specifies row and timestep to which sample the data point belongs to, this format allows more than one sample per row which potentially reduces the total amount of padding needed (e.g. [1, 1, 1, 2, 2, 0, 0]). Pytorch uses this representation (see [pack_padded_sequence](https://pytorch.org/docs/stable/generated/torch.nn.utils.rnn.pack_padded_sequence.html)).\n\nWhile Sequence packing (see [LM1B example](https://github.com/google/flax/blob/main/examples/lm1b/input_pipeline.py#L90-L92)) is is more powerful, its implementation is more complex and it is not clear whether it is worth the effort. The simplest format is sequence length masking, which is the one we propose to use.\n\n
\n\n### Bidirectional\nBidirectional processing can be achieved via a Module that accepts a `forward_rnn` Module and a `backward_rnn` Module, both of which should be `RNN` instances, in order to process the input sequence in both directions. Here we present some pseudo code of the implementation:\n\n```python\ndef __call__(self, inputs, seq_lengths):\n # Encode in the forward direction.\n carry_forward, outputs_forward = self.forward_rnn(\n inputs, seq_lengths=seq_lengths,\n return_carry=True, reverse=False,\n )\n # Encode in the reverse order.\n carry_backward, outputs_backward = self.backward_rnn(\n inputs, seq_lengths=seq_lengths,\n return_carry=True, reverse=True, # process in reverse order\n keep_order=True, # but return the sequence in the original order\n )\n # Merge both sequences.\n outputs = jax.tree.map(self.merge_fn, outputs_forward, outputs_backward)\n\n return (carry_forward, carry_backward), outputs\n```\n\nHere `merge_fn` a function that takes both outputs and fuses them (`concat` by default). As showcased in the beginning of this document, usage would look like this:\n\n```python\nforward_rnn = nn.RNN(nn.LSTMCell(), cell_size=32)\nbackward_rnn = nn.RNN(nn.GRUCell(), cell_size=32)\n# Bidirectional combinator.\nbi_rnn = nn.Bidirectional(forward_rnn, backward_rnn)\n# Encodes a batch of input sequences in both directions.\ncarry, outputs = bi_rnn(inputs, seq_lengths)\n```\n\n### Recurrent Dropout\nThere are two main uses of dropout in RNNs:\n1. Input dropout: regular dropout applied to the inputs, different for every step.\n4. Recurrent dropout: applies dropout to a recurrent input/output, same for every step.\n\nFlax's `nn.scan` can easily express both types of dropout via `split_rns`, input dropout would split rngs while recurrent dropout would not. [#2540](https://github.com/google/flax/pull/2540) was introduces such that the `rng_name` in `nn.Dropout` can now be defined by the user, this way Cells could define both types of dropout e.g:\n\n```python\nself.dropout = nn.Dropout(...) # input dropout\nself.recurrent_dropout = nn.Dropout(..., rng_collection='recurrent_dropout')\n```\nBased on this, `nn.scan` / `nn.RNN` can now specify `split_rngs` accordingly e.g:\n```\nnn.scan(scan_fn, ..., split_rngs={'dropout': True, 'recurrent_dropout': False})\n```\n\n# Future ideas\n\n
show\n\n### Sequence Packing\nAllow packing multiple sequences to make efficient use of space/memory. This might result in a trade-off where step time is higher (because at each step we need to check whether we are starting a new sequence and reset the carry/initial state), but where less padding is used increasing efficiency overall.\n\n### RNNCell redesign\n\n#### Make initialize_state an instance method\nFirst altenative is to make `initialize_carry` a instance method. With this change hyperparameters can be passed directly to the cell, it signature would look like this:\n\n```python\ndef initialize_carry(self, sample_input) -> Carry:\n ...\n```\n\nUsage would look like this:\n\n```python\nLSTM = nn.scan(\n nn.LSTMCell, variable_broadcast='params',\n split_rngs={'dropout': True})\nlstm = LSTM(features=32)\ncarry = lstm.initialize_carry(x[:, 0])\ncarry, y = lstm(carry, x)\n```\n\n#### Remove initialize_carry\n\nAn alternative is to remove `initialize_carry` entirely and have the carry state be handled as a carry collection. This would simplify usage quite a bit:\n\n```python\nLSTM = nn.scan(\n nn.LSTMCell, variable_broadcast='params',\n split_rngs={'dropout': True})\ny = LSTM(features=32)(carry, x)\n```\n\nHowever, this would require `nn.scan` to support initialization of carry collections which is currently not possible. Also, users would have to specify that a collection is mutable e.g. `mutable=['carry']`, even if they are not interested in the output carry state.\n\n
\n" }, { "path": "docs_nnx/flip/2434-general-metadata.md", "content": "# FLIP: Axis Metadata\n\n\n- Start Date: 2022-08-08\n- FLIP Issue: [#2434](https://github.com/google/flax/issues/2434)\n- FLIP PR: [#2435](https://github.com/google/flax/pull/2435)\n- Status: Proposal\n\n\n## Summary\n\nThis FLIP proposes to extend Flax's variable collections with a generic axis metadata API.\nThe core of the API is an abstract base class that is recognized by lifting transformations that can add an axis (vmap, scan).\nUsers can extend the base class to keep track of per-axis metadata in a way that works with lifted transformations.\n\n\n## Motivation\n\nGenerally, there is no way in Flax to track metadata for variables across lifted transformations.\nAxis metadata is used to keep track of semantic information about axes into other (Flax independent) APIs.\nFor example, optimizers like AdaFactor can be configured on a per-axis level and partitioning APIs\nin JAX like xmap or pjit require per variable annotations to map effectiently to parallel hardware.\n\nCurrently, there is an experimental [API](https://github.com/google/flax/blob/main/flax/linen/partitioning.py)\nsupporting partitioning annotations with wrappers around lifted transforms that change axes (``nn.scan_with_axes``, ``nn.vmap_with_axes``)\nand a special APIs to create variables (``param_with_axes`` and ``variable_with_axes``).\nThe experimental partitioning API stores the metadata in a separate collection named \"[collection]_axes\".\n\n\nThe experimental API has a number of shortcomings that we like to solve:\n1. The current API works for tracking PartitionSpecs but not for other types of metadata like optimizer annotations.\n2. The implementation using an \"xxx_axes\" collection requires error-prone and non-composable string manipulation.\n3. Special, partioning-aware variable creators and lifted transforms are required\n4. The partioning API is hard to use with pre-existing Modules that aren't partioning aware.\n\n\n## Proposal\n\nTo generalize metadata tracking and keep the specific metadata out of core Flax we propose the following abstract base class:\n\n```python\nTAxisMetadata = TypeVar(\"TAxisMetadata\", bound=\"AxisMetadata\")\n\nclass AxisMetadata(metaclass=abc.ABCMeta):\n \"\"\"Abstract base class for boxed Metadata.\n\n ``AxisMetadata`` enables arbitrary, per axis metadata for variables.\n By using ``unbox`` the metadata is stripped away to obtain the original\n variables. By using unboxing, most code handling variables does not need\n to handle ``AxisMetadata`` specifically, but can directly operate on the JAX\n arrays that they wrap.\n\n Additionally, ``AxisMetadata`` supports updating metadata whenever an axis\n is added or removed by a functional transformation\n (e.g.: ``nn.scan`` or ``nn.vmap``) using the ``add_axis`` and ``remove_axis``\n methods.\n\n By extending ``AxisMetadata``, custom metadata can be stored. See\n ``Partitioned`` for a specific implementation.\n \"\"\"\n\n @abc.abstractmethod\n def unbox(self) -> Any:\n \"\"\"Returns the content of the AxisMetadata box.\n\n Note that unlike ``meta.unbox`` the unbox call should recursively unbox\n metadata. It should simply return value that it wraps directly even\n if that value itself is an instance of AxisMetadata.\n\n In practise, AxisMetadata subclasses should be registred as PyTree nodes to\n support passing instances to JAX and Flax APIs. The leaves returned for this\n note should correspond to the value returned by unbox.\n\n Returns:\n The unboxed value.\n \"\"\"\n pass\n\n @abc.abstractmethod\n def add_axis(self: TAxisMetadata, index: int,\n params: Dict[Any, Any]) -> TAxisMetadata:\n \"\"\"Adds a new axis to the axis metadata.\n\n Note that add_axis and remove_axis should act as each other's inverse\n (meaning: ``x.add_axis(i, p).remove_axis(i, p) == x``)\n\n Args:\n index: The position at which the new axis will be inserted\n params: An arbitrary dictionary of parameters passed by the transformation\n that introduces the new axis (e.g.: ``nn.scan`` or ``nn.vmap``). The\n user passes this dictionary as the `metadata_param` argument to the\n transformation.\n Returns:\n A new instance of the same type as self and with the same ``unbox``\n content with updated axis metadata.\n \"\"\"\n pass\n\n @abc.abstractmethod\n def remove_axis(self: TAxisMetadata, index: int,\n params: Dict[Any, Any]) -> TAxisMetadata:\n \"\"\"Removes an axis from the axis metadata.\n\n Note that add_axis and remove_axis should act as each other's inverse\n (meaning: ``x.remove_axis(i, p).add_axis(i, p) == x``)\n\n Args:\n index: The position of the axis that is to be removed\n params: An arbitrary dictionary of parameters passed by the transformation\n that introduced the axis (e.g.: ``nn.scan`` or ``nn.vmap``). The\n user passes this dictionary as the `metadata_param` argument to the\n transformation.\n Returns:\n A new instance of the same type as self and with the same ``unbox``\n content with updated axis metadata.\n \"\"\"\n pass\n```\n\nWe call this type of class wrapping a value and keeping track of some additional data a **box**.\nBy defining an abstract base class for this box, the API does not need to be aware of the specifics of the metadata that is tracked.\nThis should make the API future proof and modular.\n\nThe ``add_axis`` and ``remove_axis`` method return an instance of their own type instead of mutating in-place.\nTypically, an implementation would be a ``flax.struct.PyTreeNode`` because the box should still be a valid JAX value and must therefore be handled by the PyTree API.\nCalling ``jax.tree.map`` on a boxed value will simply map over the value in the box.\nThe lifted transforms that need to handle metadata will call ``jax.tree.map(..., is_leaf=lambda x: isinstance(x, AxisMetadata))`` to find the AxisMetadata instances within a PyTree.\n\nAdvantages of the boxing approach:\n1. Boxing can be used outside of Flax and metadata is automatically \"inherited\". For example, the optimizer state will\n have the same partitioning spec as the parameters, because the state is initialized using a ``jax.tree.map`` over the boxed parameters.\n2. Boxes are composable.\n3. Boxing avoids string manipulation and generally avoids having to handle additional auxiliary collections like \"param_axes\" in the current\n partitioning API.\n4. No need to lift metadata collections separately.\n\n\nDisadvantages:\n1. Adding the boxes changes the PyTree hierarchy and introduces dataclasses within the otherwise plain, nested dict of variables.\n3. Custom Pytree nodes have a small runtime overhead. It's hard to observe this in practise because JAX calls are async.\n\n\n### Init syntax\n\n\nBoxes can be created directly by the init function of a variable. Therefore, we propose to create metadata using higher-order initializers.\nThe main advantage of this is that we can decouple metadata handling completely from the Module definition. Also, most Modules already overwrite\nattributes to override the default initialzers so users can add metadata to existing Modules without requiring any code changes.\n\nTo illustrate this, let's consider a metadata class that keeps track of PartitionSpecs used by ``pjit``:\n\n```python\nclass Partitioned(flax.struct.PyTreeNode, AxisMetadata):\n value: Any\n names: Tuple[Optional[str], ...] = flax.struct.field(pytree_node=False)\n\n def add_axis(self, index: int, params: Dict[Any, Any]) -> TAxisMetadata:\n axis_name = self._get_partition_name(params)\n names = list(self.names)\n names.insert(index, axis_name)\n return self.replace(names=tuple(names))\n\n def remove_axis(self, index: int, params: Dict[Any, Any]) -> TAxisMetadata:\n axis_name = self._get_partition_name(params)\n names = list(self.names)\n assert names.pop(index) == axis_name\n return self.replace(names=tuple(names))\n\ndef with_partitioning(init_fn, names):\n def wrapper(*args, **kwargs):\n return Partitioned(init_fn(*args, **kwargs), names)\n return wrapper\n```\n\nHere we also defined a small utility called ``with_partitioning`` that we can use to wrap existing initialzers to add metadata:\n\n\n```python\n# init kernel with lecun normal and split the output features over the data axis\npartitioned_dense = nn.Dense(features, kernel_init=with_partitioning(nn.initializers.lecun_normal, (None, \"data\")))\n```\n\nInitializing a model that creates partitioned weights would result in the following variable structure:\n\n```python\nvariables = partitioned_dense.init(rng, jnp.ones((4,)))\njax.tree.map(np.shape, variables) # => {\"params\": {\"kernel\": Partitioned(value=(4, 8), names=(None, \"data\")), bias: (8,)}}\n```\n\nThe variable tree with metadata can be used to integrate with other libraries and APIs.\nFor example, we can turn the ``Partitioned`` metadata into ``jax.pjit`` sharding annotations:\n\n```python\ndef to_sharding_spec(x):\n if isinstance(x, Partitioned):\n return PartitionSpec(*x.names)\n else:\n # fully replicated\n return PartitionSpec()\n\n# Result: {\"params\": {\"kernel\": PartitionSpec(None, \"data\"), bias: PartitionSpec()}}\nvariables_pspec = jax.tree.map(to_sharding_spec, variables, is_leaf=lambda x: isinstance(x, Partitioned))\n```\n\n### Unbox syntax\n\n\nMetadata typically doesn't need to be handled by Modules directly. Therefore, we propose to make Modules agnostic to Metadata boxes by default.\nThe ``unbox`` method can be used to unpack a variable such that only the original JAX arrays remain. Users can manually call unbox but to make\nsure Module classes don't have to call it everywhere we add an unbox keyword arg to variable returning APIs (e.g.: ``.param``, ``.variable``, ``.get_variable``).\nThe keyword arg ``unbox`` will default to ``True`` such that a Modules are metadata agnostic by default. This also means existing Modules will be backward compatible\nwith the new API.\n\n```python\nkernel = self.param(\"kernel\", self.kernel_init, shape) # No AxisMetadata instances\nkernel_box = self.get_variable(\"param\", \"kernel\", unbox=False) # AxisMetadata boxes are preserved\n```\n\n\n### Lift syntax\n\nWhen calling a lifted transformation that adds an axis you will now be able to pass a dictionary with arguments.\nThese params will be passed to ``AxisMetadata`` add_axis/remove_axis callbacks:\n\n```python\nnn.scan(..., variable_axes={\"params\": 0}, metadata_params={nn.Partitioned.AXIS_NAME: \"layers\"})\n```\n\nA dict is used such that users can add their own arguments to custom AxisMetadata classes.\n\n" }, { "path": "docs_nnx/flip/2974-kw-only-dataclasses.md", "content": "# FLIP: kw_only dataclasses\nAuthors: Brennan Saeta, Ivy Zheng\n\n - Start Date: Mar 23, 2023\n - FLIP Issue: [TBD]\n - FLIP PR: #2974\n - Status: Implementing\n\n\n## Summary\n\nPython 3.10 adds support for `kw_only` dataclasses. Subclasses of `flax.linen.Module` are automatically converted to `dataclasses` on users' behalf, but today, Flax doesn't allow setting the `kw_only` parameter to this dataclass transform, even if users are running Python 3.10. This proposal allows users to use this new feature with `nn.Module`'s.\n\n\n## Motivation\n\nIn larger Flax-based codebases (e.g. [`PaxML`](https://github.com/google/paxml) / [`Praxis`](https://github.com/google/praxis)), it’s not uncommon to define an (abstract) subclass of nn.Module that contains shared functionality that is itself further subclassed for specific implementations (e.g. [`BaseLayer`](https://github.com/google/praxis/blob/main/praxis/base_layer.py), or [`StackedTransformerRepeat`](https://github.com/google/praxis/blob/81479b260fcc13de8549cdbfb0fdf5c3f188ac90/praxis/layers/transformers.py#L1836) which is further subclassed by [`PipelineCompatibleStackedTransformerRepeat`](https://github.com/google/praxis/blob/81479b260fcc13de8549cdbfb0fdf5c3f188ac90/praxis/layers/transformers.py#L2198)).\n\nOften, these parent types define hyperparameters (constructor arguments), often with default values. Without `kw_only` on the `dataclass` transform, default values must be specified for all child layers hyperparameters. This is suboptimal, because users could forget to set them when instantiating the modules. For example, `Child` must set a default value for `num_heads` (because a non-defaulted argument can’t come after a defaulted argument if they are positional), but no reasonable default is available:\n\n```python\nclass BaseLayer(nn.Module):\n mesh: Optional[jax.experimental.mesh.Mesh] = None\n\n def with_sharding(self, some_variable, some_sharding):\n if self.mesh:\n # Do something useful here.\n\nclass Child(BaseLayer):\n num_heads: int # Don't want to have to set a default argument!\n\n def __call__(self, x):\n ...\n```\n\nNote: Flax already has this problem, which is why `nn.Module` has its own fancy `kw_only_dataclasses.dataclass` transform: it moves the `name` and `parent` dataclass fields to the end, so they can have defaults.\n\n\n## Implementation\n\nTo allow modules to optionally opt into this `kw_only` dataclass behavior, we leverage arguments to `__init_subclass__`. This would look as follows:\n\n```python\nclass BaseLayer(nn.Module, kw_only=True):\n ...\n\nclass Child(BaseLayer):\n ...\n```\n\nThe implementation of `nn.Module`’s `__init_subclass__` will be tweaked as follows:\n\n```python\nclass Module(ModuleBase):\n def __init_subclass__(self, kw_only: Optional[bool] = None):\n # ...\n if kw_only:\n if is_python_310_or_above():\n dataclass_transform_args = {'kw_only': True}\n else:\n raise TypeError(\"Can't use `kw_only` before Py3.10.\")\n else:\n dataclass_transform_args = {}\n\n kw_only_dataclasses.dataclass(\n cls, unsafe_hash='__hash__' not in cls.__dict__,\n repr=False,\n **dataclass_transform_args)\n```\n\n### Forward compatibility\n\nFor future simplification, if `kw_only` is requested and the Python version is 3.10 or above, bypass the `kw_only_dataclasses` implementation and just use the regular `dataclasses` transform.\n\nThat means we may one day remove `flax/linen/kw_only_dataclasses.py` when Flax rolls over 3.10.\n\n\n## Discussion\n\n### Aligned with Python `dataclass`\n\nWe prefer to keep the behavior of `nn.Module`’s `kw_only` aligned with the Python dataclasses. Note that this means `kw_only` will not be inheritable, and this could happen:\n\n```python\nclass BaseLayer(nn.Module, kw_only=True):\n base_muliplier: Optional[int] = -1\n\nclass ChildLayer(BaseLayer):\n child_multiplier: int\n\nBaseLayer(2) # This will throw error\nChildLayer(2) # But this will not\n```\n\n### `flax.struct.dataclass`\n\nThere’s a potentially related feature to allow `kw_only` to be specified for `flax.struct.dataclass`. This should be considered an orthogonal decision.\n\n\n" }, { "path": "docs_nnx/flip/3099-rnnbase-refactor.md", "content": "# Refactor RNNCellBase in FLIP\n\nAuthors: Cristian Garcia, Marcus Chiam, Jasmijn Bastings\n\n - Start Date: May 1, 2023\n - FLIP Issue: [TBD]\n - FLIP PR: #3053\n - Status: Implemented\n\n## Summary\nThis proposal aims to improve the usability of the `RNNCellBase` class by refactoring the `initialize_carry` method and other relevant components.\n\n## Motivation\n\nCurrently, `initialize_carry` is used to both initialize the carry and pass crucial metadata like the number of features. The API can be unintuitive as it requires users to manually calculate things that could typically be inferred by the modules themselves, such as the shape of batch dimensions and the shape of feature dimensions.\n\n### Example: ConvLSTM\nThe current API can be unintuitive in cases like `ConvLSTM` where a the `size` parameter contains both the input image shape and output feature dimensions:\n\n```python\nx = jnp.ones((2, 4, 4, 3)) # (batch, *image_shape, channels)\n\n# image shape: vvvvvvv\ncarry = nn.ConvLSTMCell.initialize_carry(key1, (16,), (64, 64, 16))\n# batch size: ^^ ^^ :output features\n\nlstm = nn.ConvLSTMCell(features=6, kernel_size=(3, 3))\n(carry, y), initial_params = lstm.init_with_output(key2, carry, x)\n```\n\nThis FLIP will propose some changes to `initialize_carry` such that the previous example can be simplified to:\n\n```python\nx = jnp.ones((2, 4, 4, 3)) # (batch, *image_shape, channels)\n\nlstm = nn.ConvLSTMCell(features=6, kernel_size=(3, 3))\ncarry = lstm.initialize_carry(key1, input_shape=x.shape)\n\n(carry, y), initial_params = lstm.init_with_output(key2, carry, x)\n```\n\n## Implementation\nThe proposal suggests the following changes:\n\n### initialize_carry\n`initialize_carry` should be refactored as an instance method with the following signature:\n\n```python\ndef initialize_carry(self, key, sample_input):\n```\n\n`sample_input` should be an array of the same shape that will be processed by the cell, excluding the time axis.\n\n### Refactor RNNCellBase subclasses\n`RNNCellBase` should be refactored to include the metadata required to initialize the cell and execute its forward pass. For `LSTMCell` and `GRUCell`, this means adding a `features` attribute that should be provided by the user upon construction. This change aligns with the structure of most other `Module`s, making them more familiar to users.\n\n```python\nx = jnp.ones((2, 100, 10)) # (batch, time, features)\n\ncell = nn.LSTMCell(features=32)\ncarry = cell.initialize_carry(PRNGKey(0), x[:, 0]) # sample input\n\n(carry, y), variables = cell.init_with_output(PRNGKey(1), carry, x)\n```\n\n### num_feature_dims\nTo simplify the handling of `RNNCellBase` instances in abstractions like `RNN`, each cell should implement the `num_feature_dims` property. For most cells, such as `LSTMCell` and `GRUCell`, this is always 1. For cells like `ConvLSTM`, this depends on their `kernel_size`.\n\n## Discussion\n### Alternative Approaches\n* To eliminate the need for `num_feature_dims`, `RNN` could support only a single batch dimension, i.e., inputs of the form `(batch, time, *features)`. Currently, it supports both multiple batch dimensions and multiple feature dimensions.\n* Another approach could be a complete redesign of how Flax deals with recurrent states. For example, a `memory` collection could be handled as part of the variables. However, this introduces challenges such as handling stateless cells during training, passing state from one layer to another, and performing initialization inside `scan`.\n\n### Refactor Cost\nInitial TGP results showed 761 broken and 110 failed tests. However, after fixing one test, TGP results in 231 broken and 13 failed tests so there seems to be a lot\nof overlap between the broken tests.\n\nTo minimize refactor costs, the current implementation will be kept for Google internal users under a deprecated name. This will allow users to migrate to the new API at their own pace. For Open Source users we should bump Flax version to\n`0.7.0` so existing users can continue to depend on `0.6.x` versions.\n" }, { "path": "docs_nnx/flip/4105-jax-style-nnx-transforms.md", "content": "# JAX-style NNX Transforms\n\n- Authors: Cristian Garcia, Anselm Levskaya\n- Date: Jun/2024\n- FLIP PR: #4107\n- Status: Implementing\n\n## Motivation\n\nNNX allows users to utilize Modules at the top level due to their eager initialization and self-contained state. This naturally leads users to want to use them with transforms and soon start playing with NNX transforms. Since NNX Modules resemble PyTrees in that they contain Arrays, new users often attempt to apply JAX conventions, for example:\n\n```py\n@nnx.vmap(in_axes=(1, 0))\ndef f(m1: Module, m2: Module):\n ...\n```\n\nHowever, this can be misleading. Currently, NNX transforms follow Linen's convention of treating input Modules as a single unit (all Modules are split together to preserve shared references) and provide APIs for transforming that State separately. The previous example effectively translates to:\n\n```py\n# this is what is really happening\n@nnx.vmap(in_axes=(IGNORE, IGNORE), state_axes={BatchStat: None, ...: 0})\ndef f(m1: Module, m2: Module):\n ...\n```\n\nNote that `IGNORE` is not a real symbol, but represents the fact that any value placed here won't affect the outcome, as Modules are replaced by empty PyTree placeholders (similar to `None`). The `state_axes` parameter controls how the State is vectorized through a mapping of high-level `Filter`s to their desired axes. In this example, `...` (ellipsis) is a filter that accepts everything, so by default all States are vectorized on the 0th axis.\n\nTo express their original intention, users must resort to more complex custom filters that guess the index of each Module in the monolith. While this is straightforward in simple cases, users generally need to calculate the index (Modules appear in the order specified by `jax.tree.leaves` over the `args`):\n\n```py\nselect_m1 = lambda path, value: path[0] == 0\nselect_m2 = lambda path, value: path[0] == 1\n\n# To select modules individually, you must create a filter (which can be tricky)\n@nnx.vmap(state_axes={select_m1: 1, select_m2: 0})\ndef f(m1: Module, m2: Module):\n ...\n```\n\n## What if JAX conventions Just Worked™?\n\nThis proposal aims to align NNX transforms with user's expectations based on their JAX experience, making the syntax work as intuitively as possible. The original example would function **as if** `m1` and `m2` were PyTrees vectorized in axes `1` and `0` respectively:\n\n```py\n@nnx.vmap(in_axes=(1, 0))\ndef f(m1: Module, m2: Module):\n ...\n```\n\nThe primary advantage of this approach is that for `vmap` and `scan`, we could eliminate the `state_axes` and `split_rngs` arguments, relying solely on the `in_axes` API. This syntax alone would likely suffice for 80-90% of use cases, as users tend to manage state in predictable ways.\n\n### The Lift symbols\n\nTo enable more fine-grained state control within each Module, we introduce the `Lift` API. By using special types containing State Filters in place of a tree prefix, state lifting can now be done **structurally**. This allows different Filters to be applied to different Modules in the arguments without the need for complex path-based filters. Ideally, each transform would support its own Lift type, adding the desired behavior through existing JAX APIs.\n\nFor example, in `vmap`, we could allow `StateAxes` instances (vmap's Lift type) to be accepted by `in/out_axes` to control how substates are handled by mapping state `Filter`s to an axis specifier:\n\n```py\nstate_axes = StateAxes({Param: 1, BatchStat: None})\n\n@nnx.vmap(in_axes=(state_axes, 0))\ndef f(m1: Module, m2: Module):\n ...\n```\n\nIn this case, `m1`'s `Param`s are vectorized in axis `1` while its `BatchStat`s are broadcasted, and `m2`'s entire state is vectorized in axis `0`.\n\nFor `nnx.grad`, we could allow `DiffState` to be used in the `argnums` parameter to specify both the position of the argument to be differentiated and a Filter specifying the differentiable State of the Module:\n\n```py\ngrads = nnx.grad(loss_fn, argnums=(DiffState(0, LoRAParam),))(model, x, y)\n```\n\n## Rng Handling\n\nTo simplify RNG state handling, we propose removing the separate `split_rngs` parameter in `vmap` and `scan`. Instead, we suggest introducing a new `nnx.split_rngs` API that would manage RNG handling before and after the transformation. This approach provides more explicit control to the user and aligns better with JAX transform behavior.\n\n## Consistent Aliasing\n\nTo ensure the correctness of transformations with objects that obey reference semantics, we must enforce consistent lifting/lowering specifications for all aliases of a reference. Transforms must adhere to two rules:\n\n1. All aliases of a reference must receive the **exact same** lifting/lowering specification.\n2. Captured references are not allowed on the output of transformed functions.\n\nFor example:\n\n```py\n@nnx.vmap(in_axes=(m1_axes, m2_axes, m1_axes), out_axes=m2_axes)\ndef f(m1, m2, m1_alias):\n return m2\n\nm2 = f(m1, m2, m1)\n```\n\nHere, `m1` has two input aliases as it is passed as the first and third input to `f`, but this is acceptable because `m1_axes` is assigned to both in `in_axes`. `m2` is passed as the second input and has an output alias, which is also acceptable because `m2_axes` is assigned in both `in_axes` and `out_axes`.\n\nLet's examine some examples of programs that should be **rejected** based on these criteria:\n\n### Inconsistent input aliases\n\nConsider a function with two arguments `m1` and `m2` being vectorized in axis `0` and `1` respectively. Passing the same Module as both arguments would be inconsistent:\n\n```py\n@nnx.vmap(in_axes=(0, 1))\ndef f(m1: Module, m2: Module):\n ...\n\nf(m, m) # This should be rejected\n```\n\n### Inconsistent input / output aliases\n\nNow consider an identity function `g` under `vmap` with `in_axes=0` and `out_axes=1`. In JAX, this would result in transposing the arrays in the inputs:\n\n```py\n@nnx.vmap(in_axes=0, out_axes=1)\ndef g(m: Module):\n return m\n```\n\nWhile this appears correct, in NNX this behavior is not well-defined because shared mutable references behave as auxiliary outputs. Under the hood, `g` is converted into a function that has the inputs as an extra first output, and `out_axes` is set to the same values as `in_axes` for that output:\n\n```py\n@nnx.vmap(in_axes=0, out_axes=(0, 1))\ndef g_real(m: Module):\n return m, m\n```\n\nThis return structure reveals an inconsistency: we're attempting to lower `m` with both `out_axes=0` and `out_axes=1`.\n\n### Inconsistent aliases in nested structures\n\nSimilar issues can arise in less obvious cases, such as when `m` is contained within another structure:\n\n```py\n@nnx.vmap(in_axes=0, out_axes=1)\ndef f(m: Module):\n return SomeModule(m)\n```\n\nThis means we must traverse the entire graph of both inputs and outputs to check for consistent assignments. The same problem occurs when passing shared reference inputs/outputs with different specifications:\n\n```py\nshared = Shared()\nm1, m2 = Foo(shared), Foo(shared)\n\n@nnx.vmap(in_axes=(0, 1))\ndef f(m1, m2): # shared is passed through both\n ...\n```\n\n### Captured Modules cannot be outputs\n\nFinally, let's consider the second consistent aliasing rule, which states that captured Modules cannot be outputs. The main issue here is that NNX needs to split all input references together to track changes, but captured Modules bypass this process. Treating them as new references would result in **implicit cloning**:\n\n```py\nm = SomeModule()\n\n@nnx.vmap(out_axes=0, axis_size=5)\ndef f():\n return m\n\nassert m is not f() # implicit cloning\n```\n\nTo preserve reference identity, we must disallow captured Modules as outputs. In practice, we can detect captured Modules using the trace level context machinery used to restrict stateful updates on Modules from a different level.\n\n## Recap\n\nIn this document, we have:\n\n* Discussed issues with the current implementation that make it unintuitive for JAX users.\n* Proposed refactoring NNX transforms to allow users to use regular JAX semantics when interacting with objects, removing extra arguments introduced by NNX transforms.\n* Introduced the use of Lift types in JAX APIs to compensate for the lack of a \"prefix\" notion in NNX objects, enabling independent lifting of Module substates.\n* Proposed a new `nnx.split_rngs` API to replace the `split_rngs` arguments in `vmap` and `scan`, making RNG handling an explicit operation and giving users more control.\n* Analyzed edge cases resulting from aliasing shared mutable references and proposed enforcing **consistent aliasing** on all transforms with semantics over the inputs." }, { "path": "docs_nnx/flip/4844-var-eager-sharding.md", "content": "- Start Date: 2025-09-12\n- FLIP PR: [#4844](https://github.com/google/flax/pull/4844)\n\n# FLIP 4844: Variable eager sharding\n\n## Summary\n[summary]: #summary\n\nSimplify the creation of sharded NNX models. When a sharding annotation is provided, all `nnx.Variable` creation will **require a mesh context** and automatically be sharded as annotated.\n\nSee [GSPMD Guide](https://flax.readthedocs.io/en/latest/guides/flax_gspmd.html) for a comprehensive guide on how to make sharded NNX models.\n\n# Motivation\n\nTo create a sharded model, user should only need to do this:\n\n```python\nmesh = jax.make_mesh(((2, 4)), (\"data\", \"model\"))\nwith jax.set_mesh(mesh):\n model = YourModelWithShardingAnnotations()\n```\n\nInstead of the current boilerplate combo of `nnx.jit`, `nnx.get_partition_spec`, `with_sharding_constraint` and `nnx.update`:\n\n```python\n@nnx.jit\ndef create_sharded_model():\n model = YourModelWithShardingAnnotations() # Unsharded at this moment.\n state = nnx.state(model) # The model's state, a pure pytree.\n pspecs = nnx.get_partition_spec(state) # Strip out the annotations from state.\n sharded_state = jax.lax.with_sharding_constraint(state, pspecs)\n nnx.update(model, sharded_state) # The model is sharded now!\n return model\n\nmesh = jax.make_mesh(((2, 4)), (\"data\", \"model\"))\nwith jax.set_mesh(mesh):\n sharded_model = create_sharded_model()\n```\n\n# Backward compatibility\n\nUser can turn off this feature in two ways:\n\n* **Global config flag**: Run `flax.config.update('flax_always_shard_variable', False)` before running any NNX model initialization.\n\n* **Variable-specific flag**: Create a specific variable with metadata `eager_sharding=False`, such as: `nnx.Param(..., eager_sharding=False)`.\n\n\n# Flexibility options\n\nFor debugging in a CPU environment, make a dummy mesh to run the model:\n\n```python\nmesh = jax.make_mesh(((1, 1, 1)), ('your', 'axes', 'names'))\nwith jax.set_mesh(mesh):\n ...\n```\n\nFor JAX explicit mode, remove the `out_sharding=` annotation on the `nnx.Variable`.\n\n\n# Implementation\n[implementation]: #implementation\n\nWhen an `nnx.Variable` is created, check for the metadata `out_sharding`, and if present, check if under a valid global mesh context of was supplied with a valid mesh. If no, throw error; if yes, call `jax.lax.with_sharding_constraint` to apply sharding constraint on the value.\n\nNote that this only works in auto sharding mode. User should use JAX-level APIs to annotate shardings for explicit mode." }, { "path": "docs_nnx/flip/5310-tree-mode-nnx.md", "content": "# Tree Mode NNX\n\nMar 4, 2026\nCristian Garcia, Samuel Anklesaria, Flax Team\n\n## Motivation\n\nCurrent NNX APIs allow general graph structures and graph transformations, this includes:\n\n1. Tracking Variable state updates \n2. Handling shared references (graphs) \n3. Supporting prefix filters (StateAxes, DiffState, StateSharding) \n4. Propagating graph updates (static state and structure changes)\n\nWhile powerful, some of these capabilities (**3** and **4**) are beyond what JAX transform APIs offer and supporting them results in both internal complexity, harder to reason about code, and a larger set of APIs a user must learn. We wish to tackle all these issues by simplifying NNX.\n\n## Proposal\n\nTo do this we propose two things. First, the introduction of **Tree Mode NNX**: a reimplementation of the NNX APIs that only handles trees, assumes referential transparency, and has a more limited support for state updates. Concretely, this means:\n\n* Automatic state updates only for Variables in NNX transforms. \n* Tree structure assumed and enforced on all APIs (no sharing) \n* Modules treated as stateless pytrees (no graph updates). \n* Full JAX transform compatibility (remove [prefix filters](#prefix-filters): StateAxes, DiffState, StateSharding).\n\nSecond, simplifying graph support. Graphs stand out as an important feature for some NNX users. However, we will be limiting support to **1** and **2**, meaning that prefix filters and graph updates will be dropped. This will make it such that tree and graph transforms can share the same underlying implementation and semantics while still allowing for a great deal of expressivity.\n\n## Implementation\n\nTree mode will be implemented on top of the current APIs by introducing a `graph` argument, when `True` graph support is enabled, when `False` only trees are supported and internals rely on `jax.tree.*` APIs. Additionally, a `graph_updates` argument will be added to NNX transforms, when `False` transforms will no longer propagate graph structure update (**4**) or support prefix filters (**3**).\n\n```py\ndef split(..., graph: bool | None = None)\n...\ndef jit(..., graph: bool | None = None, graph_updates: bool | None = None)\n...\n```\n\nIf `graph` or `graph_updates` are not provided, their default values will be taken from the `nnx_graph_mode` and `nnx_graph_updates` config flags respectively. These can be easily fetched and updated via `set_graph_mode` and `set_graph_updates`.\n\n```py\n# status\nprint(nnx.set_graph_mode.current_value())\nprint(nnx.set_graph_updates.current_value())\n\n# set value\nnnx.set_graph_mode(True/False)\nnnx.set_graph_updates(True/False)\n\n# via env vars \n# NNX_GRAPH_MODE=true/false\n# NNX_GRAPH_UPDATES=true/false\n\n# context managers\nwith nnx.set_graph_mode(True/False):\n ...\nwith nnx.set_graph_updates(True/False):\n ...\n```\n\nThe goal will be to have the default value for `nnx_graph_mode` and `nnx_graph_updates` to be set to `False`, thus enabling tree mode for new projects. Users that don’t want to migrate can use these flags to make sure their code continues to work with current features.\n\n### Simple transforms\n\nThese new transforms are highly simplified compared to current transforms, they are easier to implement and optimize, while supporting both trees and graphs. Given a user function f, most simplified transforms follow this pattern:\n\n```py\ndef transform_wrapper(*args):\n if graph: args = to_tree(args)\n check_no_aliases(args=args)\n \n @jax_transform\n def transformed_f(*args):\n updates, snapshot = updates_and_snapshot(args)\n if graph: args = from_tree(args)\n out = f(*args)\n if graph: out = to_tree(out)\n check_no_aliases(args=updates, out=out)\n updates = mask_variable_updates(updates, snapshot)\n return out, updates\n \n out, updates = transformed_f(*args)\n apply_variable_updates(args, updates)\n if graph: out = from_tree(out)\n return out\n```\n\nThe transformed function tracks input Variable `updates`, applies f, and masks Variable updates (no updates for Variables that didn’t change). It also checks that there are no Variable aliases between the inputs and outputs (no shared references), and returns the user output plus Variable updates. The wrapper function calls the transformed function, applies the Variable updates to the input Variables, and returns the user output. To support graphs, we simply convert objects to a tree representation before passing them to jax, and back to graphs before passing them to the user code.\n\n## Backward Compatibility\n\nWhen tree mode is on by default, code that relies on graphs, graph updates, and prefix filters will stop working. There are two ways to port existing code, the first is reverting the defaults config via `set_graph_mode` and `set_graph_updates` somewhere in the after the imports:\n\n```py\nfrom flax import nnx\n...\nnnx.set_graph_mode(True)\nnnx.set_graph_updates(True)\n```\n\nThe previous implementation of the transform APIs will also be accessible via the `nnx.compat` module. They are implemented as partials that set `graph=True` and `graph_updates=True`:\n\n```py\nnnx.compat.split = partial(nnx.split, graph=True)\n...\nnnx.compat.jit = partial(nnx.jit, graph=True, graph_updates=True)\n...\n```\n\n \nThe above shortcuts will make it such that porting existing code (if needed) is as simple as performing some rewrites:\n\n`nnx.split` → `nnx.compat.split` \n`nnx.jit` → `nnx.compat.jit` \n…\n\n## Breaking changes\n\n### Prefix filters {#prefix-filters}\n\nCode that relies on prefix filters such as StateAxes, StateSharding, and DiffState will require some restructuring as JAX has no equivalent mechanisms (these were added to make Linen migration easier). The solution is to use `split` and `merge` to create state groups, and pass each group through their corresponding tree prefix on the jax transform. For example:\n\n```py\n# previous code\nstate_axes = nnx.StateAxes({some_filter: 0, ...: None})\n\n@nnx.vmap(in_axis=state_axes, graph=True, graph_updates=True)\ndef f(model):\n ...\n```\n\nThis can be rewritten to `split` the model into two state groups using the previous filter, passing the groups as separate arguments, one vectorized and the other broadcasted, and using `merge` to reconstruct the model inside the transform.\n\n```py\n# new code\ngraphdef, vectorized, broadcasted = nnx.split(model, some_filter, ...)\n\n@nnx.vmap(in_axis=(0, None))\ndef f(vectorized, broadcasted):\n model = nnx.merge(graphdef, vectorized, broadcasted)\n ...\n```\n\nThis is roughly how prefix filters were implemented under the hood.\n\n### nnx.grad\n\nCode that uses `nnx.grad` will change in two ways:\n\n1. The first argument will no longer be differentiated w.r.t. to `Param`s only, this is because `grad` used this prefix filter by default: `DiffState(0, Param)`. \n2. The gradients of NNX Pytree/Module types will no longer be `State` types. Now they just follow JAX and return the same input type.\n\nConcretely it means that code like this:\n\n```py\n# previous code\ndef loss_fn(model: Foo):\n ...\n\n# uses argnums=nnx.DiffState(0, nnx.Param)\ngrads = nnx.grad(loss_fn)(model)\n```\n\nNow has to explicitly use `split` and `merge` if to avoid calculating gradients for the non-differentiable state:\n\n```py\n# new code\ndef graphdef, params, nondiff = nnx.split(model, nnx.Param, ...)\n\ndef loss_fn(params, nondiff):\n model = nnx.merge(graphdef, params, nondiff)\n ...\n\n# uses argnums=0\ngrads = nnx.grad(loss_fn)(params, nondiff)\n```\n\nIf there is no non-differentiable the `model` can be passed in directly but the gradients will now be of the same type:\n\n```py\n# new code\ndef loss_fn(model: Foo):\n ...\n\n# uses argnums=0\ngrads: Foo = nnx.grad(loss_fn)(model)\n```\n\n### nnx.custom_vjp\n\nPreviously `nnx.custom_vjp` did two particular things:\n\n1. The backward function returned the gradients of the Variable updates (`m_updates_g`) along with the output gradient. \n2. The tangent for nnx.Pytree/Module objects were of type `nnx.State`.\n\nFor a `Foo` Module with `x: Param` and `y: Param` attributes, a simple example could look like this:\n\n```py\n# previous code\n@nnx.custom_vjp\ndef f(m: Foo):\n return jnp.sin(m.x) * m.y\n\ndef f_fwd(m: Foo):\n return f(m), (jnp.cos(m.x), jnp.sin(m.x), m)\n\ndef f_bwd(res, g):\n (m_updates_g,), out_g = g\n cos_x, sin_x, m = res\n m_g: nnx.State = nnx.clone(m_updates_g) # create copy\n m_g['x'][...] = cos_x * out_g * m.y\n m_g['y'][...] = sin_x * out_g\n return (m_g,) # State gradient\n```\n\nIn the new implementation gradients for Variable updates are not returned, and the tangent type is the same as the input type (`Foo`), this matches the behavior of `jax.custom_vjp`: \n\n```py\n# new code\n@nnx.custom_vjp\ndef f(m: Foo):\n return jnp.sin(m.x) * m.y\n\ndef f_fwd(m: Foo):\n return f(m), (jnp.cos(m.x), jnp.sin(m.x), m)\n\ndef f_bwd(res, g): # no gradients for updates\n cos_x, sin_x, m = res\n m_g: Foo = nnx.clone(m) # create copy\n m_g.x[...] = cos_x * g * m.y\n m_g.y[...] = sin_x * g\n return (m_g,) # Foo gradient\n```\n\nNote that to avoid losing information, now differentiable Variables are not allowed to be updated inside `nnx.custom_vjp`.\n\n### transform\\_metadata\n\nPreviously NNX transforms like `vmap` and `scan` had a `transform_metadata` metadata argument that allowed them to update the sharding metadata.\n\n```py\n# old code\n@nnx.split_rngs(8)\n@nnx.vmap(in_axes=0, out_axes=0, transform_metadata={nnx.PARTITION_NAME: 'din'})\nclass create_stack(rngs): # 'din' added to out_sharding metadata\n return nnx.Variable(rngs.uniform((16,)), out_sharding=('dout',))\n\nv_stack = create_stack(nnx.Rngs(0))\nassert v_stack.shape == (8, 16)\nassert v_stack.out_shardings == ('din', 'dout')\n```\n\nThe new simplified NNX transform implementations don’t support this argument. However, to keep supporting the behavior, a new `nnx.transform_metadata` transform is introduced that can be inserted to get back the same results. TODO: mention it works on `jax.vmap`.\n\n```py\n# new code\n@nnx.split_rngs(8)\n@nnx.vmap(in_axes=0, out_axes=0)\n@nnx.transform_metadata(in_axes=0, out_axes=0, partition='din')\nclass create_stack(rngs): # 'din' added to out_sharding metadata\n return nnx.Variable(rngs.uniform((16,)), out_sharding=('dout',))\n\nv_stack = create_stack(nnx.Rngs(0))\nassert v_stack.shape == (8, 16)\nassert v_stack.out_shardings == ('din', 'dout')\n```\n\n`transform_metada` accepts `in_axes` and `out_axes`, these should match the values passed to the corresponding transform.\n\n### Module.sow\n\nPreviously, `Module.sow` used graph updates to capture intermediate values during computations and propagate them outside, it was used in conjunction with `nnx.pop` to log and extract intermediates:\n\n```py\n# old code\nclass Foo(nnx.Module):\n def __call__(self, x):\n self.sow(nnx.Intermediate, \"y_mean\", jnp.mean(x))\n return x\n\nmodel = Foo()\nresult = model(x)\nintermediates = nnx.pop(model, nnx.Intermediate) # extract intermediate values\n```\n\nTo achieve the same without graph updates we’ve added a new `nnx.capture` API which allows for a similar workflow.\n\n```py\n# New Code\nclass Foo(nnx.Module):\n def __call__(self, x):\n self.sow(nnx.Intermediate, \"y_mean\", jnp.mean(x))\n return x\n\nmodel = Foo()\nresult, intermediates = nnx.capture(model, nnx.Intermediate)(x)\n```\n\nIn general, `nnx.capture` takes a function or Module to be transformed, a `nnx.Variable` subclass to collect, and an optional `init` argument to initialize the collected state, which will be stored within `nnx.Variable` objects. `nnx.capture` creates a `__captures__: tuple[Variable, ...]` attribute on each `Module` instance, each Variable in `__captures__` contains a dictionary which `sow` and `perturb` populate.\n\n### Module.perturb\n\nSimilarly, `Module.perturb` was previously used to extract the gradients of intermediate values. This was done in two steps: initializing a perturbation state by running a module once, and then passing the perturbation state as a differentiable target to `grad`.\n\n```py\nclass Model(nnx.Module):\n def __call__(self, x):\n x = self.perturb('grad_of_x', x)\n ...\n return y\n\n# old code\n@nnx.jit\ndef train_step(model, optimizer, x, y):\n model(x) # Initialize perturbation state\n def loss_fn(model):\n y_pred = model(x)\n return jnp.mean((y_pred - y) ** 2)\n diff_state = nnx.DiffState(0, (nnx.Param, nnx.Perturbation))\n grads = nnx.grad(loss_fn, argnums=diff_state)(model)\n grads, interm_grads = nnx.state(grads, nnx.Param, nnx.Perturbation)\n optimizer.update(model, grads)\n nnx.pop(model, nnx.Perturbation) # clean up perturbations\n return interm_grads\n```\n\nSimilar pattern can be used with `nnx.capture` during both perturbation initialization and when running the forward pass to insert the differentiable perturbations state. In this version explicitly pass the `perturbs` state as a separate argument and use `argnums` to specify that both arguments are differentiable:\n\n```py\n# new code\n@nnx.jit\ndef train_step(model, optimizer, x, y):\n _, perturbs = nnx.capture(model, nnx.Perturbation)(x) # init perturbations\n def loss_fn(model, perturbs):\n y_pred = nnx.capture(model, init=perturbs)(x)\n return jnp.mean((y_pred - y) ** 2)\n grads, interm_grads = nnx.grad(loss_fn, argnums=(0, 1))(model, perturbs)\n optimizer.update(model, grads)\n return interm_grads\n```\n" }, { "path": "docs_nnx/flip/README.md", "content": "# FLIP: Flax Improvement Process\n\nMost changes can be discussed with simple issues/discussions and pull requests.\n\nSome changes though are a bit larger in scope or require more discussion, and\nthese should be implemented as FLIPs. This allows for writing longer documents\nthat can be discussed in a pull request themselves.\n\nThe structure of FLIPs is kept as lightweight as possible to start and might\nbe extended later on.\n\n## When you should use a FLIP\n\n- When your change requires a design doc. We prefer collecting the designs as\n FLIPs for better discoverability and further reference.\n\n- When your change requires extensive discussion. It's fine to have relatively\n short discussions on issues or pull requests, but when the discussion gets\n longer this becomes unpractical for later digestion. FLIPs allow to update the\n main document with a summary of the discussion and these updates can be\n discussed themselves in the pull request adding the FLIP.\n\n## How to start a FLIP\n\nFirst, create an issue with the [FLIP label]. All pull requests that relate to\nthe FLIP (i.e. adding the FLIP itself as well as any implementing pull requests)\nshould be linked to this issue.\n\nThen create a pull request that consists of a copy of the `0000-template.md`\nrenamed to `%04d-{short-title}.md` - with the number being the issue number.\n\n[FLIP label]: https://github.com/google/flax/issues?q=label%3AFLIP\n" }, { "path": "docs_nnx/guides/blog.md", "content": "### Do we need another JAX NN library?\n\nHello, today I want to talk to you about a new JAX library that I have been working on, but before I do that, I wanted to discuss the topic: Do we need another JAX NN library?\n\n### JAX Libraries\n\nJAX NN libraries come in a wide variety ranging from functional like Flax and Haiku, to Pytree-based like Equinox." }, { "path": "docs_nnx/guides/bridge_guide.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"# Use Flax NNX and Linen together\\n\",\n \"\\n\",\n \"This guide is for existing Flax users who want to make their codebase a mixture of Flax Linen and Flax NNX `Module`s, which is made possible thanks to the `flax.nnx.bridge` API.\\n\",\n \"\\n\",\n \"This will be helpful if you:\\n\",\n \"\\n\",\n \"* Want to migrate your codebase to NNX gradually, one module at a time;\\n\",\n \"* Have external dependency that already moved to NNX but you haven't, or is still in Linen while you've moved to NNX.\\n\",\n \"\\n\",\n \"We hope this allows you to move and try out NNX at your own pace, and leverage the best of both worlds. We will also talk about how to resolve the caveats of interoperating the two APIs, on a few aspects that they are fundamentally different.\\n\",\n \"\\n\",\n \"**Note**:\\n\",\n \"\\n\",\n \"This guide is about glueing Linen and NNX modules. To migrate an existing Linen module to NNX, check out the [Migrate from Flax Linen to Flax NNX](https://flax.readthedocs.io/en/latest/guides/linen_to_nnx.html) guide.\\n\",\n \"\\n\",\n \"And all built-in Linen layers should have equivalent NNX versions! Check out the list of [Built-in NNX layers](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/index.html).\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 1,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"import os\\n\",\n \"os.environ[\\\"XLA_FLAGS\\\"] = '--xla_force_host_platform_device_count=8'\\n\",\n \"\\n\",\n \"from flax import nnx\\n\",\n \"from flax import linen as nn\\n\",\n \"from flax.nnx import bridge\\n\",\n \"import jax\\n\",\n \"from jax import numpy as jnp\\n\",\n \"from jax.experimental import mesh_utils\\n\",\n \"from typing import *\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Submodule is all you need\\n\",\n \"\\n\",\n \"A Flax model is always a tree of modules - either old Linen modules (`flax.linen.Module`, usually written as `nn.Module`) or NNX modules (`nnx.Module`).\\n\",\n \"\\n\",\n \"An `nnx.bridge` wrapper glues the two types together, in both ways:\\n\",\n \"\\n\",\n \"* `nnx.bridge.ToNNX`: Convert a Linen module to NNX, so that it can be a submodule of another NNX module, or stand alone to be trained in NNX-style training loops.\\n\",\n \"* `nnx.bridge.ToLinen`: Vice versa, convert a NNX module to Linen.\\n\",\n \"\\n\",\n \"This means you can move in either top-down or bottom-up behavior: convert the whole Linen module to NNX, then gradually move down, or convert all the lower level modules to NNX then move up.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## The Basics\\n\",\n \"\\n\",\n \"There are two fundamental difference between Linen and NNX modules:\\n\",\n \"\\n\",\n \"* **Stateless vs. stateful**: Linen module instances are stateless: variables are returned from a purely functional `.init()` call and managed separately. NNX modules, however, owns its variables as instance attributes.\\n\",\n \"\\n\",\n \"* **Lazy vs. eager**: Linen modules only allocate space to create variables when they actually see their input. Whereas NNX module instances create their variables the moment they are instantiated, without seeing a sample input.\\n\",\n \"\\n\",\n \"With that in mind, let's look at how the `nnx.bridge` wrappers tackle the differences.\\n\",\n \"\\n\",\n \"### Linen -> NNX\\n\",\n \"\\n\",\n \"Since Linen modules may require an input to create variables, we semi-formally supported lazy initialization in the NNX modules converted from Linen. The Linen variables are created when you give it a sample input.\\n\",\n \"\\n\",\n \"For you, it's calling `nnx.bridge.lazy_init()` where you call `module.init()` in Linen code.\\n\",\n \"\\n\",\n \"(Note: you can call `nnx.display` upon any NNX module to inspect all its variables and state.)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 2,\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"class LinenDot(nn.Module):\\n\",\n \" out_dim: int\\n\",\n \" w_init: Callable[..., Any] = nn.initializers.lecun_normal()\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self, x):\\n\",\n \" # Linen might need the input shape to create the weight!\\n\",\n \" w = self.param('w', self.w_init, (x.shape[-1], self.out_dim))\\n\",\n \" return x @ w\\n\",\n \"\\n\",\n \"x = jax.random.normal(jax.random.key(42), (4, 32))\\n\",\n \"model = bridge.ToNNX(LinenDot(64),\\n\",\n \" rngs=nnx.Rngs(0)) # => `model = LinenDot(64)` in Linen\\n\",\n \"bridge.lazy_init(model, x) # => `var = model.init(key, x)` in Linen\\n\",\n \"y = model(x) # => `y = model.apply(var, x)` in Linen\\n\",\n \"\\n\",\n \"nnx.display(model)\\n\",\n \"\\n\",\n \"# In-place swap your weight array and the model still works!\\n\",\n \"model.w.value = jax.random.normal(jax.random.key(1), (32, 64))\\n\",\n \"assert not jnp.allclose(y, model(x))\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"`nnx.bridge.lazy_init` also works even if the top-level module is a pure-NNX one, so you can do sub-moduling as you wish:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 3,\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"class NNXOuter(nnx.Module):\\n\",\n \" def __init__(self, out_dim: int, rngs: nnx.Rngs):\\n\",\n \" self.dot = nnx.bridge.ToNNX(LinenDot(out_dim), rngs=rngs)\\n\",\n \" self.b = nnx.Param(jax.random.uniform(rngs.params(), (1, out_dim,)))\\n\",\n \"\\n\",\n \" def __call__(self, x):\\n\",\n \" return self.dot(x) + self.b\\n\",\n \"\\n\",\n \"x = jax.random.normal(jax.random.key(42), (4, 32))\\n\",\n \"model = bridge.lazy_init(NNXOuter(64, rngs=nnx.Rngs(0)), x) # Can fit into one line\\n\",\n \"nnx.display(model)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"The Linen weight is already converted to a typical NNX variable, which is a thin wrapper of the actual JAX array value within. Here, `w` is an `nnx.Param` because it belongs to the `params` collection of `LinenDot` module.\\n\",\n \"\\n\",\n \"We will talk more about different collections and types in the [NNX Variable <-> Linen Collections](#variable-types-vs-collections) section. Right now, just know that they are converted to NNX variables like native ones.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 4,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"assert isinstance(model.dot.w, nnx.Param)\\n\",\n \"assert isinstance(model.dot.w.value, jax.Array)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"If you create this model witout using `nnx.bridge.lazy_init`, the NNX variables defined outside will be initialized as usual, but the Linen part (wrapped inside `ToNNX`) will not.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 5,\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"partial_model = NNXOuter(64, rngs=nnx.Rngs(0))\\n\",\n \"nnx.display(partial_model)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 6,\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"full_model = bridge.lazy_init(partial_model, x)\\n\",\n \"nnx.display(full_model)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### NNX -> Linen\\n\",\n \"\\n\",\n \"To convert an NNX module to Linen, you should forward your creation arguments to `bridge.ToLinen` and let it handle the actual creation process.\\n\",\n \"\\n\",\n \"This is because NNX module instance initializes all its variables eagerly when it is created, which consumes memory and compute. On the other hand, Linen modules are stateless, and the typical `init` and `apply` process involves multiple creation of them. So `bridge.to_linen` will handle the actual module creation and make sure no memory is allocated twice.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 7,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"['params']\\n\",\n \"(32, 64)\\n\",\n \"(4, 64)\\n\"\n ]\n }\n ],\n \"source\": [\n \"class NNXDot(nnx.Module):\\n\",\n \" def __init__(self, in_dim: int, out_dim: int, rngs: nnx.Rngs):\\n\",\n \" self.w = nnx.Param(nnx.initializers.lecun_normal()(\\n\",\n \" rngs.params(), (in_dim, out_dim)))\\n\",\n \" def __call__(self, x: jax.Array):\\n\",\n \" return x @ self.w\\n\",\n \"\\n\",\n \"x = jax.random.normal(jax.random.key(42), (4, 32))\\n\",\n \"# Pass in the arguments, not an actual module\\n\",\n \"model = bridge.to_linen(NNXDot, 32, out_dim=64)\\n\",\n \"variables = model.init(jax.random.key(0), x)\\n\",\n \"y = model.apply(variables, x)\\n\",\n \"\\n\",\n \"print(list(variables.keys()))\\n\",\n \"print(variables['params']['w'].shape) # => (32, 64)\\n\",\n \"print(y.shape) # => (4, 64)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"`bridge.to_linen` is actually a convenience wrapper around the Linen module `bridge.ToLinen`. Most likely you won't need to use `ToLinen` directly at all, unless you are using one of the built-in arguments of `ToLinen`. For example, if your NNX module doesn't want to be initialized with RNG handling:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 8,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"class NNXAddConstant(nnx.Module):\\n\",\n \" def __init__(self):\\n\",\n \" self.constant = nnx.Variable(jnp.array(1))\\n\",\n \" def __call__(self, x):\\n\",\n \" return x + self.constant\\n\",\n \"\\n\",\n \"# You have to use `skip_rng=True` because this module's `__init__` don't\\n\",\n \"# take `rng` as argument\\n\",\n \"model = bridge.ToLinen(NNXAddConstant, skip_rng=True)\\n\",\n \"y, var = model.init_with_output(jax.random.key(0), x)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Similar to `ToNNX`, you can use `ToLinen` to create a submodule of another Linen module.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 9,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"(32, 64) (1, 64) (4, 64)\\n\"\n ]\n }\n ],\n \"source\": [\n \"class LinenOuter(nn.Module):\\n\",\n \" out_dim: int\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self, x):\\n\",\n \" dot = bridge.to_linen(NNXDot, x.shape[-1], self.out_dim)\\n\",\n \" b = self.param('b', nn.initializers.lecun_normal(), (1, self.out_dim))\\n\",\n \" return dot(x) + b\\n\",\n \"\\n\",\n \"x = jax.random.normal(jax.random.key(42), (4, 32))\\n\",\n \"model = LinenOuter(out_dim=64)\\n\",\n \"y, variables = model.init_with_output(jax.random.key(0), x)\\n\",\n \"w, b = variables['params']['ToLinen_0']['w'], variables['params']['b']\\n\",\n \"print(w.shape, b.shape, y.shape)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Handling RNG keys\\n\",\n \"\\n\",\n \"All Flax modules, Linen or NNX, automatically handle the RNG keys for variable creation and random layers like dropouts. However, the specific logics of RNG key splitting are different, so you cannot generate the same params between Linen and NNX modules, even if you pass in same keys.\\n\",\n \"\\n\",\n \"Another difference is that NNX modules are stateful, so they can track and update the RNG keys within themselves.\\n\",\n \"\\n\",\n \"### Linen to NNX\\n\",\n \"\\n\",\n \"If you convert a Linen module to NNX, you enjoy the stateful benefit and don't need to pass in extra RNG keys on every module call. You can use always `nnx.reseed` to reset the RNG state within.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 10,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"x = jax.random.normal(jax.random.key(42), (4, 32))\\n\",\n \"model = bridge.ToNNX(nn.Dropout(rate=0.5, deterministic=False), rngs=nnx.Rngs(dropout=0))\\n\",\n \"# We don't really need to call lazy_init because no extra params were created here,\\n\",\n \"# but it's a good practice to always add this line.\\n\",\n \"bridge.lazy_init(model, x)\\n\",\n \"y1, y2 = model(x), model(x)\\n\",\n \"assert not jnp.allclose(y1, y2) # Two runs yield different outputs!\\n\",\n \"\\n\",\n \"# Reset the dropout RNG seed, so that next model run will be the same as the first.\\n\",\n \"nnx.reseed(model, dropout=0)\\n\",\n \"y1 = model(x)\\n\",\n \"nnx.reseed(model, dropout=0)\\n\",\n \"y2 = model(x)\\n\",\n \"assert jnp.allclose(y1, y2) # Two runs yield the same output!\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### NNX to Linen\\n\",\n \"\\n\",\n \"`to_linen` will automatically take the `rngs` dict argument and create a `Rngs` object that is passed to the underlying NNX module via the `rngs` keyword argument. If the module holds internal `RngState`, `to_linen` will always call reseed using the `rngs` dict to reset the RNG state.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 12,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"x = jax.random.normal(jax.random.key(42), (4, 32))\\n\",\n \"model = bridge.to_linen(nnx.Dropout, rate=0.5)\\n\",\n \"variables = model.init({'dropout': jax.random.key(0)}, x)\\n\",\n \"\\n\",\n \"# Just pass different RNG keys for every `apply()` call.\\n\",\n \"y1 = model.apply(variables, x, rngs={'dropout': jax.random.key(1)})\\n\",\n \"y2 = model.apply(variables, x, rngs={'dropout': jax.random.key(2)})\\n\",\n \"assert not jnp.allclose(y1, y2) # Every call yields different output!\\n\",\n \"y3 = model.apply(variables, x, rngs={'dropout': jax.random.key(1)})\\n\",\n \"assert jnp.allclose(y1, y3) # When you use same top-level RNG, outputs are same\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## NNX variable types vs. Linen collections\\n\",\n \"\\n\",\n \"When you want to group some variables as one category, in Linen you use different collections; in NNX, since all variables shall be top-level Python attributes, you use different variable types.\\n\",\n \"\\n\",\n \"Therefore, when mixing Linen and NNX modules, Flax must know the 1-to-1 mapping between Linen collections and NNX variable types, so that `ToNNX` and `ToLinen` can do the conversion automatically.\\n\",\n \"\\n\",\n \"Flax keeps a registry for this, and it already covers all Flax's built-in Linen collections. You can register extra mapping of NNX variable type and Linen collection names using `nnx.register_variable_name_type_pair`.\\n\",\n \"\\n\",\n \"### Linen to NNX\\n\",\n \"\\n\",\n \"For any collection of your Linen module, `ToNNX` will convert all its endpoint arrays (aka. leaves) to a subtype of `nnx.Variable`, either from registry or automatically created on-the-fly.\\n\",\n \"\\n\",\n \"(However, we still keep the whole collection as one class attribute, because Linen modules may have duplicated names over different collections.)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 13,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"\\u001b[38;2;79;201;177mParam\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;105;105;105m # 12 (48 B)\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254mvalue\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0mArray([[ 0.53824717, 0.7668343 , -0.38585317],\\n\",\n \" [-0.35335615, -0.5244857 , -0.43152452],\\n\",\n \" [-1.0662307 , 0.14089198, -0.16519307],\\n\",\n \" [ 0.3971692 , 0.43213558, -0.461545 ]], dtype=float32)\\n\",\n \"\\u001b[38;2;255;213;3m)\\u001b[0m\\n\",\n \"\\u001b[38;2;79;201;177mParam\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;105;105;105m # 3 (12 B)\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254mvalue\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0mArray([0., 0., 0.], dtype=float32)\\n\",\n \"\\u001b[38;2;255;213;3m)\\u001b[0m\\n\",\n \"\\u001b[38;2;79;201;177mcounter\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;105;105;105m # 1 (4 B)\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254mvalue\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0mArray(0, dtype=int32)\\n\",\n \"\\u001b[38;2;255;213;3m)\\u001b[0m\\n\",\n \"\\n\",\n \"(Intermediate( # 1 (4 B)\\n\",\n \" value=Array(0.5475821, dtype=float32)\\n\",\n \"),)\\n\"\n ]\n }\n ],\n \"source\": [\n \"class LinenMultiCollections(nn.Module):\\n\",\n \" out_dim: int\\n\",\n \" def setup(self):\\n\",\n \" self.w = self.param('w', nn.initializers.lecun_normal(), (x.shape[-1], self.out_dim))\\n\",\n \" self.b = self.param('b', nn.zeros_init(), (self.out_dim,))\\n\",\n \" self.count = self.variable('counter', 'count', lambda: jnp.zeros((), jnp.int32))\\n\",\n \"\\n\",\n \" def __call__(self, x):\\n\",\n \" if not self.is_initializing():\\n\",\n \" self.count.value += 1\\n\",\n \" y = x @ self.w + self.b\\n\",\n \" self.sow('intermediates', 'dot_sum', jnp.sum(y))\\n\",\n \" return y\\n\",\n \"\\n\",\n \"x = jax.random.normal(jax.random.key(42), (2, 4))\\n\",\n \"model = bridge.lazy_init(bridge.ToNNX(LinenMultiCollections(3), rngs=nnx.Rngs(0)), x)\\n\",\n \"print(model.w) # Of type `nnx.Param` - note this is still under attribute `params`\\n\",\n \"print(model.b) # Of type `nnx.Param`\\n\",\n \"print(model.count) # Of type `counter` - auto-created type from the collection name\\n\",\n \"print(type(model.count))\\n\",\n \"\\n\",\n \"y = model(x, mutable=True) # Linen's `sow()` needs `mutable=True` to trigger\\n\",\n \"print(model.dot_sum) # Of type `nnx.Intermediates`\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"You can quickly separate different types of NNX variables apart using `nnx.split`.\\n\",\n \"\\n\",\n \"This can be handy when you only want to set some variables as trainable.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 14,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"All Params: ['b', 'w']\\n\",\n \"All Counters: ['count']\\n\",\n \"All the rest (intermediates and RNG keys): ['dot_sum', 'rngs']\\n\"\n ]\n }\n ],\n \"source\": [\n \"# Separate variables of different types with nnx.split\\n\",\n \"CountType = type(model.count)\\n\",\n \"static, params, counter, the_rest = nnx.split(model, nnx.Param, CountType, ...)\\n\",\n \"print('All Params:', list(params.keys()))\\n\",\n \"print('All Counters:', list(counter.keys()))\\n\",\n \"print('All the rest (intermediates and RNG keys):', list(the_rest.keys()))\\n\",\n \"\\n\",\n \"model = nnx.merge(static, params, counter, the_rest) # You can merge them back at any time\\n\",\n \"y = model(x, mutable=True) # still works!\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"cc9d78ed\",\n \"metadata\": {},\n \"source\": [\n \" All Params: ['b', 'w']\\n\",\n \" All Counters: ['count']\\n\",\n \" All the rest (intermediates and RNG keys): ['dot_sum', 'rngs']\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### NNX to Linen\\n\",\n \"\\n\",\n \"If you define custom NNX variable types, you should register their names with `nnx.register_variable_name` so that they go to the desired collections.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 15,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"All Linen collections: ['nnx', 'LoRAParam', 'params', 'counts']\\n\",\n \"{'w': Array([[ 0.2916921 , 0.22780475, 0.06553137],\\n\",\n \" [ 0.17487915, -0.34043145, 0.24764155],\\n\",\n \" [ 0.6420431 , 0.6220095 , -0.44769976],\\n\",\n \" [ 0.11161668, 0.83873135, -0.7446058 ]], dtype=float32)}\\n\"\n ]\n }\n ],\n \"source\": [\n \"@nnx.register_variable_name('counts', overwrite=True)\\n\",\n \"class Count(nnx.Variable): pass\\n\",\n \"\\n\",\n \"\\n\",\n \"class NNXMultiCollections(nnx.Module):\\n\",\n \" def __init__(self, din, dout, rngs):\\n\",\n \" self.w = nnx.Param(nnx.initializers.lecun_normal()(rngs.params(), (din, dout)))\\n\",\n \" self.lora = nnx.LoRA(din, 3, dout, rngs=rngs)\\n\",\n \" self.count = Count(jnp.array(0))\\n\",\n \"\\n\",\n \" def __call__(self, x):\\n\",\n \" self.count.value += 1\\n\",\n \" return (x @ self.w.value) + self.lora(x)\\n\",\n \"\\n\",\n \"xkey, pkey, dkey = jax.random.split(jax.random.key(0), 3)\\n\",\n \"x = jax.random.normal(xkey, (2, 4))\\n\",\n \"model = bridge.to_linen(NNXMultiCollections, 4, 3)\\n\",\n \"var = model.init({'params': pkey, 'dropout': dkey}, x)\\n\",\n \"print('All Linen collections:', list(var.keys()))\\n\",\n \"print(var['params'])\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \" All Linen collections: ['LoRAParam', 'params', 'counts']\\n\",\n \" {'w': Array([[ 0.2916921 , 0.22780475, 0.06553137],\\n\",\n \" [ 0.17487915, -0.34043145, 0.24764155],\\n\",\n \" [ 0.6420431 , 0.6220095 , -0.44769976],\\n\",\n \" [ 0.11161668, 0.83873135, -0.7446058 ]], dtype=float32)}\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Partition metadata\\n\",\n \"\\n\",\n \"Flax uses a metadata wrapper box over the raw JAX array to annotate how a variable should be sharded.\\n\",\n \"\\n\",\n \"In Linen, this is an optional feature that triggered by using `nn.with_partitioning` on initializers (see more on [Linen partition metadata guide](https://flax.readthedocs.io/en/latest/guides/parallel_training/flax_on_pjit.html)). In NNX, since all NNX variables are wrapped by `nnx.Variable` class anyway, that class will hold the sharding annotations too.\\n\",\n \"\\n\",\n \"The `bridge.ToNNX` and `bridge.ToLinen` API will automatically convert the sharding annotations, if you use the built-in annotation methods (aka. `nn.with_partitioning` for Linen and `nnx.with_partitioning` for NNX).\\n\",\n \"\\n\",\n \"### Linen to NNX\\n\",\n \"\\n\",\n \"Even if you are not using any partition metadata in your Linen module, the variable JAX arrays will be converted to `nnx.Variable`s that wraps the true JAX array within.\\n\",\n \"\\n\",\n \"If you use `nn.with_partitioning` to annotate your Linen module's variables, the annotation will be converted to a `.sharding` field in the corresponding `nnx.Variable`.\\n\",\n \"\\n\",\n \"You can then use `nnx.with_sharding_constraint` to explicitly put the arrays into the annotated partitions within a `jax.jit`-compiled function, to initialize the whole model with every array at the right sharding.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 16,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"We have 8 fake JAX devices now to partition this model...\\n\",\n \"\\n\",\n \"('in', 'out')\\n\",\n \"GSPMDSharding({devices=[2,4]<=[8]})\\n\"\n ]\n }\n ],\n \"source\": [\n \"class LinenDotWithPartitioning(nn.Module):\\n\",\n \" out_dim: int\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self, x):\\n\",\n \" w = self.param('w', nn.with_partitioning(nn.initializers.lecun_normal(),\\n\",\n \" ('in', 'out')),\\n\",\n \" (x.shape[-1], self.out_dim))\\n\",\n \" return x @ w\\n\",\n \"\\n\",\n \"@nnx.jit\\n\",\n \"def create_sharded_nnx_module(x):\\n\",\n \" model = bridge.lazy_init(\\n\",\n \" bridge.ToNNX(LinenDotWithPartitioning(64), rngs=nnx.Rngs(0)), x)\\n\",\n \" state = nnx.state(model)\\n\",\n \" sharded_state = nnx.with_sharding_constraint(state, nnx.get_partition_spec(state))\\n\",\n \" nnx.update(model, sharded_state)\\n\",\n \" return model\\n\",\n \"\\n\",\n \"\\n\",\n \"print(f'We have {len(jax.devices())} fake JAX devices now to partition this model...')\\n\",\n \"mesh = jax.sharding.Mesh(devices=mesh_utils.create_device_mesh((2, 4)),\\n\",\n \" axis_names=('in', 'out'))\\n\",\n \"x = jax.random.normal(jax.random.key(42), (4, 32))\\n\",\n \"with jax.set_mesh(mesh):\\n\",\n \" model = create_sharded_nnx_module(x)\\n\",\n \"\\n\",\n \"print(type(model.w)) # `nnx.Param`\\n\",\n \"print(model.w.sharding) # The partition annotation attached with `w`\\n\",\n \"print(model.w.value.sharding) # The underlying JAX array is sharded across the 2x4 mesh\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"08555a06\",\n \"metadata\": {},\n \"source\": [\n \" We have 8 fake JAX devices now to partition this model...\\n\",\n \" \\n\",\n \" ('in', 'out')\\n\",\n \" GSPMDSharding({devices=[2,4]<=[8]})\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### NNX to Linen\\n\",\n \"\\n\",\n \"If you are not using any metadata feature of the `nnx.Variable` (i.e., no sharding annotation, no registered hooks), the converted Linen module will not add a metadata wrapper to your NNX variable, and you don't need to worry about it.\\n\",\n \"\\n\",\n \"But if you did add sharding annotations to your NNX variables, `ToLinen` will convert them to a default Linen partition metadata class called `bridge.NNXMeta`, retaining all the metadata you put into the NNX variable.\\n\",\n \"\\n\",\n \"Like with any Linen metadata wrappers, you can use `linen.unbox()` to get the raw JAX array tree.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 17,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"GSPMDSharding({devices=[2,4]<=[8]})\\n\"\n ]\n }\n ],\n \"source\": [\n \"class NNXDotWithParititioning(nnx.Module):\\n\",\n \" def __init__(self, in_dim: int, out_dim: int, rngs: nnx.Rngs):\\n\",\n \" init_fn = nnx.with_partitioning(nnx.initializers.lecun_normal(), ('in', 'out'))\\n\",\n \" self.w = nnx.Param(init_fn(rngs.params(), (in_dim, out_dim)))\\n\",\n \" def __call__(self, x: jax.Array):\\n\",\n \" return x @ self.w\\n\",\n \"\\n\",\n \"x = jax.random.normal(jax.random.key(42), (4, 32))\\n\",\n \"\\n\",\n \"@jax.jit\\n\",\n \"def create_sharded_variables(key, x):\\n\",\n \" model = bridge.to_linen(NNXDotWithParititioning, 32, 64)\\n\",\n \" variables = model.init(key, x)\\n\",\n \" # A `NNXMeta` wrapper of the underlying `nnx.Param`\\n\",\n \" assert type(variables['params']['w']) == bridge.NNXMeta\\n\",\n \" # The annotation coming from the `nnx.Param` => (in, out)\\n\",\n \" assert variables['params']['w'].metadata['sharding'] == ('in', 'out')\\n\",\n \"\\n\",\n \" unboxed_variables = nn.unbox(variables)\\n\",\n \" variable_pspecs = nn.get_partition_spec(variables)\\n\",\n \" assert isinstance(unboxed_variables['params']['w'], jax.Array)\\n\",\n \" assert variable_pspecs['params']['w'] == jax.sharding.PartitionSpec('in', 'out')\\n\",\n \"\\n\",\n \" sharded_vars = jax.tree.map(jax.lax.with_sharding_constraint,\\n\",\n \" nn.unbox(variables),\\n\",\n \" nn.get_partition_spec(variables))\\n\",\n \" return sharded_vars\\n\",\n \"\\n\",\n \"with jax.set_mesh(mesh):\\n\",\n \" variables = create_sharded_variables(jax.random.key(0), x)\\n\",\n \"\\n\",\n \"# The underlying JAX array is sharded across the 2x4 mesh\\n\",\n \"print(variables['params']['w'].sharding)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \" GSPMDSharding({devices=[2,4]<=[8]})\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Lifted transforms\\n\",\n \"\\n\",\n \"In general, if you want to apply Linen/NNX-style lifted transforms upon an `nnx.bridge`-converted module, just go ahead and do it in the usual Linen/NNX syntax.\\n\",\n \"\\n\",\n \"For Linen-style transforms, note that `bridge.ToLinen` is the top level module class, so you may want to just use it as the first argument of your transforms (which needs to be a `linen.Module` class in most cases)\\n\",\n \"\\n\",\n \"### Linen to NNX\\n\",\n \"\\n\",\n \"NNX style lifted transforms are similar to JAX transforms, and they work on functions.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 18,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"(4, 32, 64)\\n\",\n \"(4, 64)\\n\"\n ]\n }\n ],\n \"source\": [\n \"class NNXVmapped(nnx.Module):\\n\",\n \" def __init__(self, out_dim: int, vmap_axis_size: int, rngs: nnx.Rngs):\\n\",\n \" self.linen_dot = nnx.bridge.ToNNX(nn.Dense(out_dim, use_bias=False), rngs=rngs)\\n\",\n \" self.vmap_axis_size = vmap_axis_size\\n\",\n \"\\n\",\n \" def __call__(self, x):\\n\",\n \"\\n\",\n \" @nnx.split_rngs(splits=self.vmap_axis_size)\\n\",\n \" @nnx.vmap(in_axes=(0, 0), axis_size=self.vmap_axis_size)\\n\",\n \" def vmap_fn(submodule, x):\\n\",\n \" return submodule(x)\\n\",\n \"\\n\",\n \" return vmap_fn(self.linen_dot, x)\\n\",\n \"\\n\",\n \"x = jax.random.normal(jax.random.key(0), (4, 32))\\n\",\n \"model = bridge.lazy_init(NNXVmapped(64, 4, rngs=nnx.Rngs(0)), x)\\n\",\n \"\\n\",\n \"print(model.linen_dot.kernel.shape) # (4, 32, 64) - first axis with dim 4 got vmapped\\n\",\n \"y = model(x)\\n\",\n \"print(y.shape)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \" (4, 32, 64)\\n\",\n \" (4, 64)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"61a1ac21\",\n \"metadata\": {},\n \"source\": [\n \"### NNX to Linen\\n\",\n \"\\n\",\n \"Note that `bridge.ToLinen` is the top level module class, so you may want to just use it as the first argument of your transforms (which needs to be a `linen.Module` class in most cases).\\n\",\n \"\\n\",\n \"`ToLien` can naturally be used with Linen transforms like `nn.vmap` or `nn.scan`.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 19,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"(4, 32, 64)\\n\",\n \"(4, 64)\\n\"\n ]\n }\n ],\n \"source\": [\n \"class LinenVmapped(nn.Module):\\n\",\n \" dout: int\\n\",\n \" @nn.compact\\n\",\n \" def __call__(self, x):\\n\",\n \" inner = nn.vmap(bridge.ToLinen, variable_axes={'params': 0}, split_rngs={'params': True}\\n\",\n \" )(nnx.Linear, args=(x.shape[-1], self.dout))\\n\",\n \" return inner(x)\\n\",\n \"\\n\",\n \"x = jax.random.normal(jax.random.key(42), (4, 32))\\n\",\n \"model = LinenVmapped(64)\\n\",\n \"var = model.init(jax.random.key(0), x)\\n\",\n \"print(var['params']['VmapToLinen_0']['kernel'].shape) # (4, 32, 64) - leading dim 4 got vmapped\\n\",\n \"y = model.apply(var, x)\\n\",\n \"print(y.shape)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"178d2b2f\",\n \"metadata\": {},\n \"source\": [\n \" (4, 32, 64)\\n\",\n \" (4, 64)\"\n ]\n }\n ],\n \"metadata\": {\n \"jupytext\": {\n \"cell_metadata_filter\": \"-all\",\n \"formats\": \"ipynb,md:myst\",\n \"main_language\": \"python\"\n },\n \"language_info\": {\n \"codemirror_mode\": {\n \"name\": \"ipython\",\n \"version\": 3\n },\n \"file_extension\": \".py\",\n \"mimetype\": \"text/x-python\",\n \"name\": \"python\",\n \"nbconvert_exporter\": \"python\",\n \"pygments_lexer\": \"ipython3\",\n \"version\": \"3.11.9\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 2\n}\n" }, { "path": "docs_nnx/guides/bridge_guide.md", "content": "---\njupytext:\n cell_metadata_filter: -all\n formats: ipynb,md:myst\n main_language: python\n text_representation:\n extension: .md\n format_name: myst\n format_version: 0.13\n jupytext_version: 1.13.8\n---\n\n# Use Flax NNX and Linen together\n\nThis guide is for existing Flax users who want to make their codebase a mixture of Flax Linen and Flax NNX `Module`s, which is made possible thanks to the `flax.nnx.bridge` API.\n\nThis will be helpful if you:\n\n* Want to migrate your codebase to NNX gradually, one module at a time;\n* Have external dependency that already moved to NNX but you haven't, or is still in Linen while you've moved to NNX.\n\nWe hope this allows you to move and try out NNX at your own pace, and leverage the best of both worlds. We will also talk about how to resolve the caveats of interoperating the two APIs, on a few aspects that they are fundamentally different.\n\n**Note**:\n\nThis guide is about glueing Linen and NNX modules. To migrate an existing Linen module to NNX, check out the [Migrate from Flax Linen to Flax NNX](https://flax.readthedocs.io/en/latest/guides/linen_to_nnx.html) guide.\n\nAnd all built-in Linen layers should have equivalent NNX versions! Check out the list of [Built-in NNX layers](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/index.html).\n\n```{code-cell} ipython3\nimport os\nos.environ[\"XLA_FLAGS\"] = '--xla_force_host_platform_device_count=8'\n\nfrom flax import nnx\nfrom flax import linen as nn\nfrom flax.nnx import bridge\nimport jax\nfrom jax import numpy as jnp\nfrom jax.experimental import mesh_utils\nfrom typing import *\n```\n\n## Submodule is all you need\n\nA Flax model is always a tree of modules - either old Linen modules (`flax.linen.Module`, usually written as `nn.Module`) or NNX modules (`nnx.Module`).\n\nAn `nnx.bridge` wrapper glues the two types together, in both ways:\n\n* `nnx.bridge.ToNNX`: Convert a Linen module to NNX, so that it can be a submodule of another NNX module, or stand alone to be trained in NNX-style training loops.\n* `nnx.bridge.ToLinen`: Vice versa, convert a NNX module to Linen.\n\nThis means you can move in either top-down or bottom-up behavior: convert the whole Linen module to NNX, then gradually move down, or convert all the lower level modules to NNX then move up.\n\n+++\n\n## The Basics\n\nThere are two fundamental difference between Linen and NNX modules:\n\n* **Stateless vs. stateful**: Linen module instances are stateless: variables are returned from a purely functional `.init()` call and managed separately. NNX modules, however, owns its variables as instance attributes.\n\n* **Lazy vs. eager**: Linen modules only allocate space to create variables when they actually see their input. Whereas NNX module instances create their variables the moment they are instantiated, without seeing a sample input.\n\nWith that in mind, let's look at how the `nnx.bridge` wrappers tackle the differences.\n\n### Linen -> NNX\n\nSince Linen modules may require an input to create variables, we semi-formally supported lazy initialization in the NNX modules converted from Linen. The Linen variables are created when you give it a sample input.\n\nFor you, it's calling `nnx.bridge.lazy_init()` where you call `module.init()` in Linen code.\n\n(Note: you can call `nnx.display` upon any NNX module to inspect all its variables and state.)\n\n```{code-cell} ipython3\nclass LinenDot(nn.Module):\n out_dim: int\n w_init: Callable[..., Any] = nn.initializers.lecun_normal()\n @nn.compact\n def __call__(self, x):\n # Linen might need the input shape to create the weight!\n w = self.param('w', self.w_init, (x.shape[-1], self.out_dim))\n return x @ w\n\nx = jax.random.normal(jax.random.key(42), (4, 32))\nmodel = bridge.ToNNX(LinenDot(64),\n rngs=nnx.Rngs(0)) # => `model = LinenDot(64)` in Linen\nbridge.lazy_init(model, x) # => `var = model.init(key, x)` in Linen\ny = model(x) # => `y = model.apply(var, x)` in Linen\n\nnnx.display(model)\n\n# In-place swap your weight array and the model still works!\nmodel.w.value = jax.random.normal(jax.random.key(1), (32, 64))\nassert not jnp.allclose(y, model(x))\n```\n\n`nnx.bridge.lazy_init` also works even if the top-level module is a pure-NNX one, so you can do sub-moduling as you wish:\n\n```{code-cell} ipython3\nclass NNXOuter(nnx.Module):\n def __init__(self, out_dim: int, rngs: nnx.Rngs):\n self.dot = nnx.bridge.ToNNX(LinenDot(out_dim), rngs=rngs)\n self.b = nnx.Param(jax.random.uniform(rngs.params(), (1, out_dim,)))\n\n def __call__(self, x):\n return self.dot(x) + self.b\n\nx = jax.random.normal(jax.random.key(42), (4, 32))\nmodel = bridge.lazy_init(NNXOuter(64, rngs=nnx.Rngs(0)), x) # Can fit into one line\nnnx.display(model)\n```\n\nThe Linen weight is already converted to a typical NNX variable, which is a thin wrapper of the actual JAX array value within. Here, `w` is an `nnx.Param` because it belongs to the `params` collection of `LinenDot` module.\n\nWe will talk more about different collections and types in the [NNX Variable <-> Linen Collections](#variable-types-vs-collections) section. Right now, just know that they are converted to NNX variables like native ones.\n\n```{code-cell} ipython3\nassert isinstance(model.dot.w, nnx.Param)\nassert isinstance(model.dot.w.value, jax.Array)\n```\n\nIf you create this model witout using `nnx.bridge.lazy_init`, the NNX variables defined outside will be initialized as usual, but the Linen part (wrapped inside `ToNNX`) will not.\n\n```{code-cell} ipython3\npartial_model = NNXOuter(64, rngs=nnx.Rngs(0))\nnnx.display(partial_model)\n```\n\n```{code-cell} ipython3\nfull_model = bridge.lazy_init(partial_model, x)\nnnx.display(full_model)\n```\n\n### NNX -> Linen\n\nTo convert an NNX module to Linen, you should forward your creation arguments to `bridge.ToLinen` and let it handle the actual creation process.\n\nThis is because NNX module instance initializes all its variables eagerly when it is created, which consumes memory and compute. On the other hand, Linen modules are stateless, and the typical `init` and `apply` process involves multiple creation of them. So `bridge.to_linen` will handle the actual module creation and make sure no memory is allocated twice.\n\n```{code-cell} ipython3\nclass NNXDot(nnx.Module):\n def __init__(self, in_dim: int, out_dim: int, rngs: nnx.Rngs):\n self.w = nnx.Param(nnx.initializers.lecun_normal()(\n rngs.params(), (in_dim, out_dim)))\n def __call__(self, x: jax.Array):\n return x @ self.w\n\nx = jax.random.normal(jax.random.key(42), (4, 32))\n# Pass in the arguments, not an actual module\nmodel = bridge.to_linen(NNXDot, 32, out_dim=64)\nvariables = model.init(jax.random.key(0), x)\ny = model.apply(variables, x)\n\nprint(list(variables.keys()))\nprint(variables['params']['w'].shape) # => (32, 64)\nprint(y.shape) # => (4, 64)\n```\n\n`bridge.to_linen` is actually a convenience wrapper around the Linen module `bridge.ToLinen`. Most likely you won't need to use `ToLinen` directly at all, unless you are using one of the built-in arguments of `ToLinen`. For example, if your NNX module doesn't want to be initialized with RNG handling:\n\n```{code-cell} ipython3\nclass NNXAddConstant(nnx.Module):\n def __init__(self):\n self.constant = nnx.Variable(jnp.array(1))\n def __call__(self, x):\n return x + self.constant\n\n# You have to use `skip_rng=True` because this module's `__init__` don't\n# take `rng` as argument\nmodel = bridge.ToLinen(NNXAddConstant, skip_rng=True)\ny, var = model.init_with_output(jax.random.key(0), x)\n```\n\nSimilar to `ToNNX`, you can use `ToLinen` to create a submodule of another Linen module.\n\n```{code-cell} ipython3\nclass LinenOuter(nn.Module):\n out_dim: int\n @nn.compact\n def __call__(self, x):\n dot = bridge.to_linen(NNXDot, x.shape[-1], self.out_dim)\n b = self.param('b', nn.initializers.lecun_normal(), (1, self.out_dim))\n return dot(x) + b\n\nx = jax.random.normal(jax.random.key(42), (4, 32))\nmodel = LinenOuter(out_dim=64)\ny, variables = model.init_with_output(jax.random.key(0), x)\nw, b = variables['params']['ToLinen_0']['w'], variables['params']['b']\nprint(w.shape, b.shape, y.shape)\n```\n\n## Handling RNG keys\n\nAll Flax modules, Linen or NNX, automatically handle the RNG keys for variable creation and random layers like dropouts. However, the specific logics of RNG key splitting are different, so you cannot generate the same params between Linen and NNX modules, even if you pass in same keys.\n\nAnother difference is that NNX modules are stateful, so they can track and update the RNG keys within themselves.\n\n### Linen to NNX\n\nIf you convert a Linen module to NNX, you enjoy the stateful benefit and don't need to pass in extra RNG keys on every module call. You can use always `nnx.reseed` to reset the RNG state within.\n\n```{code-cell} ipython3\nx = jax.random.normal(jax.random.key(42), (4, 32))\nmodel = bridge.ToNNX(nn.Dropout(rate=0.5, deterministic=False), rngs=nnx.Rngs(dropout=0))\n# We don't really need to call lazy_init because no extra params were created here,\n# but it's a good practice to always add this line.\nbridge.lazy_init(model, x)\ny1, y2 = model(x), model(x)\nassert not jnp.allclose(y1, y2) # Two runs yield different outputs!\n\n# Reset the dropout RNG seed, so that next model run will be the same as the first.\nnnx.reseed(model, dropout=0)\ny1 = model(x)\nnnx.reseed(model, dropout=0)\ny2 = model(x)\nassert jnp.allclose(y1, y2) # Two runs yield the same output!\n```\n\n### NNX to Linen\n\n`to_linen` will automatically take the `rngs` dict argument and create a `Rngs` object that is passed to the underlying NNX module via the `rngs` keyword argument. If the module holds internal `RngState`, `to_linen` will always call reseed using the `rngs` dict to reset the RNG state.\n\n```{code-cell} ipython3\nx = jax.random.normal(jax.random.key(42), (4, 32))\nmodel = bridge.to_linen(nnx.Dropout, rate=0.5)\nvariables = model.init({'dropout': jax.random.key(0)}, x)\n\n# Just pass different RNG keys for every `apply()` call.\ny1 = model.apply(variables, x, rngs={'dropout': jax.random.key(1)})\ny2 = model.apply(variables, x, rngs={'dropout': jax.random.key(2)})\nassert not jnp.allclose(y1, y2) # Every call yields different output!\ny3 = model.apply(variables, x, rngs={'dropout': jax.random.key(1)})\nassert jnp.allclose(y1, y3) # When you use same top-level RNG, outputs are same\n```\n\n## NNX variable types vs. Linen collections\n\nWhen you want to group some variables as one category, in Linen you use different collections; in NNX, since all variables shall be top-level Python attributes, you use different variable types.\n\nTherefore, when mixing Linen and NNX modules, Flax must know the 1-to-1 mapping between Linen collections and NNX variable types, so that `ToNNX` and `ToLinen` can do the conversion automatically.\n\nFlax keeps a registry for this, and it already covers all Flax's built-in Linen collections. You can register extra mapping of NNX variable type and Linen collection names using `nnx.register_variable_name_type_pair`.\n\n### Linen to NNX\n\nFor any collection of your Linen module, `ToNNX` will convert all its endpoint arrays (aka. leaves) to a subtype of `nnx.Variable`, either from registry or automatically created on-the-fly.\n\n(However, we still keep the whole collection as one class attribute, because Linen modules may have duplicated names over different collections.)\n\n```{code-cell} ipython3\nclass LinenMultiCollections(nn.Module):\n out_dim: int\n def setup(self):\n self.w = self.param('w', nn.initializers.lecun_normal(), (x.shape[-1], self.out_dim))\n self.b = self.param('b', nn.zeros_init(), (self.out_dim,))\n self.count = self.variable('counter', 'count', lambda: jnp.zeros((), jnp.int32))\n\n def __call__(self, x):\n if not self.is_initializing():\n self.count.value += 1\n y = x @ self.w + self.b\n self.sow('intermediates', 'dot_sum', jnp.sum(y))\n return y\n\nx = jax.random.normal(jax.random.key(42), (2, 4))\nmodel = bridge.lazy_init(bridge.ToNNX(LinenMultiCollections(3), rngs=nnx.Rngs(0)), x)\nprint(model.w) # Of type `nnx.Param` - note this is still under attribute `params`\nprint(model.b) # Of type `nnx.Param`\nprint(model.count) # Of type `counter` - auto-created type from the collection name\nprint(type(model.count))\n\ny = model(x, mutable=True) # Linen's `sow()` needs `mutable=True` to trigger\nprint(model.dot_sum) # Of type `nnx.Intermediates`\n```\n\nYou can quickly separate different types of NNX variables apart using `nnx.split`.\n\nThis can be handy when you only want to set some variables as trainable.\n\n```{code-cell} ipython3\n# Separate variables of different types with nnx.split\nCountType = type(model.count)\nstatic, params, counter, the_rest = nnx.split(model, nnx.Param, CountType, ...)\nprint('All Params:', list(params.keys()))\nprint('All Counters:', list(counter.keys()))\nprint('All the rest (intermediates and RNG keys):', list(the_rest.keys()))\n\nmodel = nnx.merge(static, params, counter, the_rest) # You can merge them back at any time\ny = model(x, mutable=True) # still works!\n```\n\n All Params: ['b', 'w']\n All Counters: ['count']\n All the rest (intermediates and RNG keys): ['dot_sum', 'rngs']\n\n+++\n\n### NNX to Linen\n\nIf you define custom NNX variable types, you should register their names with `nnx.register_variable_name` so that they go to the desired collections.\n\n```{code-cell} ipython3\n@nnx.register_variable_name('counts', overwrite=True)\nclass Count(nnx.Variable): pass\n\n\nclass NNXMultiCollections(nnx.Module):\n def __init__(self, din, dout, rngs):\n self.w = nnx.Param(nnx.initializers.lecun_normal()(rngs.params(), (din, dout)))\n self.lora = nnx.LoRA(din, 3, dout, rngs=rngs)\n self.count = Count(jnp.array(0))\n\n def __call__(self, x):\n self.count.value += 1\n return (x @ self.w.value) + self.lora(x)\n\nxkey, pkey, dkey = jax.random.split(jax.random.key(0), 3)\nx = jax.random.normal(xkey, (2, 4))\nmodel = bridge.to_linen(NNXMultiCollections, 4, 3)\nvar = model.init({'params': pkey, 'dropout': dkey}, x)\nprint('All Linen collections:', list(var.keys()))\nprint(var['params'])\n```\n\n All Linen collections: ['LoRAParam', 'params', 'counts']\n {'w': Array([[ 0.2916921 , 0.22780475, 0.06553137],\n [ 0.17487915, -0.34043145, 0.24764155],\n [ 0.6420431 , 0.6220095 , -0.44769976],\n [ 0.11161668, 0.83873135, -0.7446058 ]], dtype=float32)}\n\n+++\n\n## Partition metadata\n\nFlax uses a metadata wrapper box over the raw JAX array to annotate how a variable should be sharded.\n\nIn Linen, this is an optional feature that triggered by using `nn.with_partitioning` on initializers (see more on [Linen partition metadata guide](https://flax.readthedocs.io/en/latest/guides/parallel_training/flax_on_pjit.html)). In NNX, since all NNX variables are wrapped by `nnx.Variable` class anyway, that class will hold the sharding annotations too.\n\nThe `bridge.ToNNX` and `bridge.ToLinen` API will automatically convert the sharding annotations, if you use the built-in annotation methods (aka. `nn.with_partitioning` for Linen and `nnx.with_partitioning` for NNX).\n\n### Linen to NNX\n\nEven if you are not using any partition metadata in your Linen module, the variable JAX arrays will be converted to `nnx.Variable`s that wraps the true JAX array within.\n\nIf you use `nn.with_partitioning` to annotate your Linen module's variables, the annotation will be converted to a `.sharding` field in the corresponding `nnx.Variable`.\n\nYou can then use `nnx.with_sharding_constraint` to explicitly put the arrays into the annotated partitions within a `jax.jit`-compiled function, to initialize the whole model with every array at the right sharding.\n\n```{code-cell} ipython3\nclass LinenDotWithPartitioning(nn.Module):\n out_dim: int\n @nn.compact\n def __call__(self, x):\n w = self.param('w', nn.with_partitioning(nn.initializers.lecun_normal(),\n ('in', 'out')),\n (x.shape[-1], self.out_dim))\n return x @ w\n\n@nnx.jit\ndef create_sharded_nnx_module(x):\n model = bridge.lazy_init(\n bridge.ToNNX(LinenDotWithPartitioning(64), rngs=nnx.Rngs(0)), x)\n state = nnx.state(model)\n sharded_state = nnx.with_sharding_constraint(state, nnx.get_partition_spec(state))\n nnx.update(model, sharded_state)\n return model\n\n\nprint(f'We have {len(jax.devices())} fake JAX devices now to partition this model...')\nmesh = jax.sharding.Mesh(devices=mesh_utils.create_device_mesh((2, 4)),\n axis_names=('in', 'out'))\nx = jax.random.normal(jax.random.key(42), (4, 32))\nwith jax.set_mesh(mesh):\n model = create_sharded_nnx_module(x)\n\nprint(type(model.w)) # `nnx.Param`\nprint(model.w.sharding) # The partition annotation attached with `w`\nprint(model.w.value.sharding) # The underlying JAX array is sharded across the 2x4 mesh\n```\n\n We have 8 fake JAX devices now to partition this model...\n \n ('in', 'out')\n GSPMDSharding({devices=[2,4]<=[8]})\n\n+++\n\n### NNX to Linen\n\nIf you are not using any metadata feature of the `nnx.Variable` (i.e., no sharding annotation, no registered hooks), the converted Linen module will not add a metadata wrapper to your NNX variable, and you don't need to worry about it.\n\nBut if you did add sharding annotations to your NNX variables, `ToLinen` will convert them to a default Linen partition metadata class called `bridge.NNXMeta`, retaining all the metadata you put into the NNX variable.\n\nLike with any Linen metadata wrappers, you can use `linen.unbox()` to get the raw JAX array tree.\n\n```{code-cell} ipython3\nclass NNXDotWithParititioning(nnx.Module):\n def __init__(self, in_dim: int, out_dim: int, rngs: nnx.Rngs):\n init_fn = nnx.with_partitioning(nnx.initializers.lecun_normal(), ('in', 'out'))\n self.w = nnx.Param(init_fn(rngs.params(), (in_dim, out_dim)))\n def __call__(self, x: jax.Array):\n return x @ self.w\n\nx = jax.random.normal(jax.random.key(42), (4, 32))\n\n@jax.jit\ndef create_sharded_variables(key, x):\n model = bridge.to_linen(NNXDotWithParititioning, 32, 64)\n variables = model.init(key, x)\n # A `NNXMeta` wrapper of the underlying `nnx.Param`\n assert type(variables['params']['w']) == bridge.NNXMeta\n # The annotation coming from the `nnx.Param` => (in, out)\n assert variables['params']['w'].metadata['sharding'] == ('in', 'out')\n\n unboxed_variables = nn.unbox(variables)\n variable_pspecs = nn.get_partition_spec(variables)\n assert isinstance(unboxed_variables['params']['w'], jax.Array)\n assert variable_pspecs['params']['w'] == jax.sharding.PartitionSpec('in', 'out')\n\n sharded_vars = jax.tree.map(jax.lax.with_sharding_constraint,\n nn.unbox(variables),\n nn.get_partition_spec(variables))\n return sharded_vars\n\nwith jax.set_mesh(mesh):\n variables = create_sharded_variables(jax.random.key(0), x)\n\n# The underlying JAX array is sharded across the 2x4 mesh\nprint(variables['params']['w'].sharding)\n```\n\n GSPMDSharding({devices=[2,4]<=[8]})\n\n+++\n\n## Lifted transforms\n\nIn general, if you want to apply Linen/NNX-style lifted transforms upon an `nnx.bridge`-converted module, just go ahead and do it in the usual Linen/NNX syntax.\n\nFor Linen-style transforms, note that `bridge.ToLinen` is the top level module class, so you may want to just use it as the first argument of your transforms (which needs to be a `linen.Module` class in most cases)\n\n### Linen to NNX\n\nNNX style lifted transforms are similar to JAX transforms, and they work on functions.\n\n```{code-cell} ipython3\nclass NNXVmapped(nnx.Module):\n def __init__(self, out_dim: int, vmap_axis_size: int, rngs: nnx.Rngs):\n self.linen_dot = nnx.bridge.ToNNX(nn.Dense(out_dim, use_bias=False), rngs=rngs)\n self.vmap_axis_size = vmap_axis_size\n\n def __call__(self, x):\n\n @nnx.split_rngs(splits=self.vmap_axis_size)\n @nnx.vmap(in_axes=(0, 0), axis_size=self.vmap_axis_size)\n def vmap_fn(submodule, x):\n return submodule(x)\n\n return vmap_fn(self.linen_dot, x)\n\nx = jax.random.normal(jax.random.key(0), (4, 32))\nmodel = bridge.lazy_init(NNXVmapped(64, 4, rngs=nnx.Rngs(0)), x)\n\nprint(model.linen_dot.kernel.shape) # (4, 32, 64) - first axis with dim 4 got vmapped\ny = model(x)\nprint(y.shape)\n```\n\n (4, 32, 64)\n (4, 64)\n\n+++\n\n### NNX to Linen\n\nNote that `bridge.ToLinen` is the top level module class, so you may want to just use it as the first argument of your transforms (which needs to be a `linen.Module` class in most cases).\n\n`ToLien` can naturally be used with Linen transforms like `nn.vmap` or `nn.scan`.\n\n```{code-cell} ipython3\nclass LinenVmapped(nn.Module):\n dout: int\n @nn.compact\n def __call__(self, x):\n inner = nn.vmap(bridge.ToLinen, variable_axes={'params': 0}, split_rngs={'params': True}\n )(nnx.Linear, args=(x.shape[-1], self.dout))\n return inner(x)\n\nx = jax.random.normal(jax.random.key(42), (4, 32))\nmodel = LinenVmapped(64)\nvar = model.init(jax.random.key(0), x)\nprint(var['params']['VmapToLinen_0']['kernel'].shape) # (4, 32, 64) - leading dim 4 got vmapped\ny = model.apply(var, x)\nprint(y.shape)\n```\n\n (4, 32, 64)\n (4, 64)\n" }, { "path": "docs_nnx/guides/checkpointing.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"# Save and load checkpoints\\n\",\n \"\\n\",\n \"This guide demonstrates how to save and load Flax NNX model checkpoints with [Orbax](https://orbax.readthedocs.io/).\\n\",\n \"\\n\",\n \"> **Note:** The Flax team does not actively maintain a library for saving and loading model checkpoints to disk. Therefore, it is recommended you use external libraries like [Orbax](https://orbax.readthedocs.io/en/latest/index.html) to do it.\\n\",\n \"\\n\",\n \"In this guide you will learn how to:\\n\",\n \"\\n\",\n \"* Save checkpoints.\\n\",\n \"* Restore checkpoints.\\n\",\n \"* Restore checkpoints if checkpoint structures differ. \\n\",\n \"* Perform multi-process checkpointing. \\n\",\n \"\\n\",\n \"The Orbax API examples used throughout the guide are for demonstration purposes, and for the most up-to-date recommended APIs refer to the [Orbax website](https://orbax.readthedocs.io/).\\n\",\n \"\\n\",\n \"> **Note:** The Flax team recommends using [Orbax](https://orbax.readthedocs.io/en/latest/index.html) for saving and loading checkpoints to disk, as we do not actively maintain a library for these functionalities.\\n\",\n \"\\n\",\n \"> **Note:** If you are looking for Flax Linen's legacy `flax.training.checkpoints` package, it was deprecated in 2023 in favor of Orbax. The documentation resides on the [Flax Linen site](https://flax-linen.readthedocs.io/en/latest/guides/training_techniques/use_checkpointing.html).\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Setup\\n\",\n \"\\n\",\n \"Import the necessary dependencies, set up a checkpoint directory and an example Flax NNX model - `TwoLayerMLP` - by subclassing [`nnx.Module`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/module.html).\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 1,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"from flax import nnx\\n\",\n \"import orbax.checkpoint as ocp\\n\",\n \"import jax\\n\",\n \"from jax import numpy as jnp\\n\",\n \"import numpy as np\\n\",\n \"\\n\",\n \"ckpt_dir = ocp.test_utils.erase_and_create_empty('/tmp/my-checkpoints/')\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 2,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"class TwoLayerMLP(nnx.Module):\\n\",\n \" def __init__(self, dim, rngs: nnx.Rngs):\\n\",\n \" self.linear1 = nnx.Linear(dim, dim, rngs=rngs, use_bias=False)\\n\",\n \" self.linear2 = nnx.Linear(dim, dim, rngs=rngs, use_bias=False)\\n\",\n \"\\n\",\n \" def __call__(self, x):\\n\",\n \" x = self.linear1(x)\\n\",\n \" return self.linear2(x)\\n\",\n \"\\n\",\n \"# Instantiate the model and show we can run it.\\n\",\n \"model = TwoLayerMLP(4, rngs=nnx.Rngs(0))\\n\",\n \"x = jax.random.normal(jax.random.key(42), (3, 4))\\n\",\n \"assert model(x).shape == (3, 4)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Save checkpoints\\n\",\n \"\\n\",\n \"JAX checkpointing libraries, such as Orbax, can save and load any given JAX [pytree](https://jax.readthedocs.io/en/latest/pytrees.html), which is a pure, possibly nested container of [`jax.Array`s)](https://jax.readthedocs.io/en/latest/key-concepts.html#jax-arrays-jax-array) (or, \\\"tensors\\\" as some other frameworks would put it). In the context of machine learning, the checkpoint is usually a pytree of model parameters and other data, such as optimizer states.\\n\",\n \"\\n\",\n \"In Flax NNX, you can obtain such a pytree from an [`nnx.Module`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/module.html) by calling [`nnx.split`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/graph.html#flax.nnx.split), and picking up the returned [`nnx.State`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/state.html#flax.nnx.State).\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 3,\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"_, state = nnx.split(model)\\n\",\n \"nnx.display(state)\\n\",\n \"\\n\",\n \"checkpointer = ocp.StandardCheckpointer()\\n\",\n \"checkpointer.save(ckpt_dir / 'state', state)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"
\\n\",\n \"\\n\",\n \"\\n\",\n \"\\n\",\n \"
\\n\",\n \"\\n\",\n \"\\n\",\n \"## Restore checkpoints\\n\",\n \"\\n\",\n \"Note that you saved the checkpoint as a Flax class of [`nnx.State`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/state.html#flax.nnx.State), which is also nested with [`nnx.Variable`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Variable) and [`nnx.Param`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Param) classes.\\n\",\n \"\\n\",\n \"At checkpoint restoration time, you need to have these classes ready in your runtime, and instruct the checkpointing library (Orbax) to restore your pytree back to that structure. This can be achieved as follows:\\n\",\n \"- First, create an abstract Flax NNX model (without allocating any memory for arrays), and show its abstract variable state to the checkpointing library.\\n\",\n \"- Once you have the state, use [`nnx.merge`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/graph.html#flax.nnx.merge) to obtain your Flax NNX model, and use it as usual.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 4,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"The abstract NNX state (all leaves are abstract arrays):\\n\"\n ]\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"NNX State restored: \\n\"\n ]\n },\n {\n \"name\": \"stderr\",\n \"output_type\": \"stream\",\n \"text\": [\n \"/Users/cgarciae/repos/flax/.venv/lib/python3.11/site-packages/orbax/checkpoint/_src/serialization/type_handlers.py:1251: UserWarning: Sharding info not provided when restoring. Populating sharding info from sharding file. Please note restoration time will be slightly increased due to reading from file. Note also that this option is unsafe when restoring on a different topology than the checkpoint was saved with.\\n\",\n \" warnings.warn(\\n\"\n ]\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"# Restore the checkpoint back to its `nnx.State` structure - need an abstract reference.\\n\",\n \"abstract_model = nnx.eval_shape(lambda: TwoLayerMLP(4, rngs=nnx.Rngs(0)))\\n\",\n \"graphdef, abstract_state = nnx.split(abstract_model)\\n\",\n \"print('The abstract NNX state (all leaves are abstract arrays):')\\n\",\n \"nnx.display(abstract_state)\\n\",\n \"\\n\",\n \"state_restored = checkpointer.restore(ckpt_dir / 'state', abstract_state)\\n\",\n \"jax.tree.map(np.testing.assert_array_equal, state, state_restored)\\n\",\n \"print('NNX State restored: ')\\n\",\n \"nnx.display(state_restored)\\n\",\n \"\\n\",\n \"# The model is now good to use!\\n\",\n \"model = nnx.merge(graphdef, state_restored)\\n\",\n \"assert model(x).shape == (3, 4)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \" The abstract NNX state (all leaves are abstract arrays):\\n\",\n \"\\n\",\n \"\\n\",\n \"\\n\",\n \"
\\n\",\n \"\\n\",\n \"\\n\",\n \" NNX State restored: \\n\",\n \"\\n\",\n \"\\n\",\n \" /Users/ivyzheng/envs/flax-head/lib/python3.11/site-packages/orbax/checkpoint/type_handlers.py:1439: UserWarning: Couldn't find sharding info under RestoreArgs. Populating sharding info from sharding file. Please note restoration time will be slightly increased due to reading from file instead of directly from RestoreArgs. Note also that this option is unsafe when restoring on a different topology than the checkpoint was saved with.\\n\",\n \" warnings.warn(\\n\",\n \"\\n\",\n \"\\n\",\n \"\\n\",\n \"
\\n\",\n \"\\n\",\n \"\\n\",\n \"\\n\",\n \"
\\n\",\n \"\\n\",\n \"\\n\",\n \"## Save and restore as pure dictionaries\\n\",\n \"\\n\",\n \"When interacting with checkpoint libraries (like Orbax), you may prefer to work with Python built-in container types. In this case, you can use the [`nnx.State.to_pure_dict`](https://github.com/google/flax/blob/764e1732dcd3b8bf178b9ba73ddecf125709b5d7/flax/nnx/statelib.py#L170) and [`nnx.State.replace_by_pure_dict`](https://github.com/google/flax/blob/764e1732dcd3b8bf178b9ba73ddecf125709b5d7/flax/nnx/statelib.py#L179) API to convert an [`nnx.State`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/state.html#flax.nnx.State) to and from pure nested dictionaries.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 5,\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"name\": \"stderr\",\n \"output_type\": \"stream\",\n \"text\": [\n \"WARNING:absl:`StandardCheckpointHandler` expects a target tree to be provided for restore. Not doing so is generally UNSAFE unless you know the present topology to be the same one as the checkpoint was saved under.\\n\"\n ]\n }\n ],\n \"source\": [\n \"# Save as pure dict\\n\",\n \"pure_dict_state = nnx.to_pure_dict(state)\\n\",\n \"nnx.display(pure_dict_state)\\n\",\n \"checkpointer.save(ckpt_dir / 'pure_dict', pure_dict_state)\\n\",\n \"\\n\",\n \"# Restore as a pure dictionary.\\n\",\n \"restored_pure_dict = checkpointer.restore(ckpt_dir / 'pure_dict')\\n\",\n \"abstract_model = nnx.eval_shape(lambda: TwoLayerMLP(4, rngs=nnx.Rngs(0)))\\n\",\n \"graphdef, abstract_state = nnx.split(abstract_model)\\n\",\n \"nnx.replace_by_pure_dict(abstract_state, restored_pure_dict)\\n\",\n \"model = nnx.merge(graphdef, abstract_state)\\n\",\n \"assert model(x).shape == (3, 4) # The model still works!\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"
\\n\",\n \"\\n\",\n \"\\n\",\n \"\\n\",\n \"
\\n\",\n \"\\n\",\n \"\\n\",\n \" WARNING:absl:`StandardCheckpointHandler` expects a target tree to be provided for restore. Not doing so is generally UNSAFE unless you know the present topology to be the same one as the checkpoint was saved under.\\n\",\n \"\\n\",\n \"\\n\",\n \"## Restore when checkpoint structures differ\\n\",\n \"\\n\",\n \"The ability to load a checkpoint as a pure nested dictionary can come in handy when you want to load some outdated checkpoints that no longer match with your current model code. Check out this simple example below.\\n\",\n \"\\n\",\n \"This pattern also works if you save the checkpoint as an [`nnx.State`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/state.html#flax.nnx.State) instead of a pure dictionary. Check out the [Checkpoint surgery section](https://flax.readthedocs.io/en/latest/guides/surgery.html#checkpoint-surgery) of the [Model Surgery](https://flax.readthedocs.io/en/latest/guides/surgery.html) guide for an example with code. The only difference is you need to reprocess your raw dictionary a bit before calling [`nnx.State.replace_by_pure_dict`](https://github.com/google/flax/blob/764e1732dcd3b8bf178b9ba73ddecf125709b5d7/flax/nnx/statelib.py#L179).\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 6,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stderr\",\n \"output_type\": \"stream\",\n \"text\": [\n \"WARNING:absl:`StandardCheckpointHandler` expects a target tree to be provided for restore. Not doing so is generally UNSAFE unless you know the present topology to be the same one as the checkpoint was saved under.\\n\"\n ]\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"class ModifiedTwoLayerMLP(nnx.Module):\\n\",\n \" \\\"\\\"\\\"A modified version of TwoLayerMLP, which requires bias arrays.\\\"\\\"\\\"\\n\",\n \" def __init__(self, dim, rngs: nnx.Rngs):\\n\",\n \" self.linear1 = nnx.Linear(dim, dim, rngs=rngs, use_bias=True) # We need bias now!\\n\",\n \" self.linear2 = nnx.Linear(dim, dim, rngs=rngs, use_bias=True) # We need bias now!\\n\",\n \"\\n\",\n \" def __call__(self, x):\\n\",\n \" x = self.linear1(x)\\n\",\n \" return self.linear2(x)\\n\",\n \"\\n\",\n \"# Accommodate your old checkpoint to the new code.\\n\",\n \"restored_pure_dict = checkpointer.restore(ckpt_dir / 'pure_dict')\\n\",\n \"restored_pure_dict['linear1']['bias'] = jnp.zeros((4,))\\n\",\n \"restored_pure_dict['linear2']['bias'] = jnp.zeros((4,))\\n\",\n \"\\n\",\n \"# Same restore code as above.\\n\",\n \"abstract_model = nnx.eval_shape(lambda: ModifiedTwoLayerMLP(4, rngs=nnx.Rngs(0)))\\n\",\n \"graphdef, abstract_state = nnx.split(abstract_model)\\n\",\n \"nnx.replace_by_pure_dict(abstract_state, restored_pure_dict)\\n\",\n \"model = nnx.merge(graphdef, abstract_state)\\n\",\n \"assert model(x).shape == (3, 4) # The new model works!\\n\",\n \"\\n\",\n \"nnx.display(model.linear1)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \" WARNING:absl:`StandardCheckpointHandler` expects a target tree to be provided for restore. Not doing so is generally UNSAFE unless you know the present topology to be the same one as the checkpoint was saved under.\\n\",\n \"\\n\",\n \"\\n\",\n \"\\n\",\n \"
\\n\",\n \"\\n\",\n \"\\n\",\n \"\\n\",\n \"
\\n\",\n \"\\n\",\n \"\\n\",\n \"## Multi-process checkpointing\\n\",\n \"\\n\",\n \"In a multi-host/multi-process environment, you would want to restore your checkpoint as sharded across multiple devices. Check out the [Load sharded model from a checkpoint](https://flax.readthedocs.io/en/latest/guides/flax_gspmd.html#load-sharded-model-from-a-checkpoint) section in the Flax [Scale up on multiple devices](https://flax.readthedocs.io/en/latest/guides/flax_gspmd.html) guide to learn how to derive a sharding pytree and use it to load your checkpoint.\\n\",\n \"\\n\",\n \"> **Note:** JAX provides several ways to scale up your code on multiple hosts at the same time. This usually happens when the number of devices (CPU/GPU/TPU) is so large that different devices are managed by different hosts (CPU). Check out JAX’s [Introduction to parallel programming](https://jax.readthedocs.io/en/latest/sharded-computation.html), [Using JAX in multi-host and multi-process environments](https://jax.readthedocs.io/en/latest/multi_process.html), [Distributed arrays and automatic parallelization](https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html), and [Manual parallelism with `shard_map`](https://jax.readthedocs.io/en/latest/notebooks/shard_map.html).\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Other checkpointing features\\n\",\n \"\\n\",\n \"This guide only uses the simplest [`orbax.checkpoint.StandardCheckpointer`](https://orbax.readthedocs.io/en/latest/api_reference/checkpoint.checkpointers.html#standardcheckpointer) API to show how to save and load on the Flax modeling side. Feel free to use other tools or libraries as you see fit.\\n\",\n \"\\n\",\n \"In addition, check out [the Orbax website](https://orbax.readthedocs.io/en/latest/index.html) for other commonly used features, such as:\\n\",\n \"\\n\",\n \"* [`CheckpointManager`](https://orbax.readthedocs.io/en/latest/guides/checkpoint/api_refactor.html) to track checkpoints from different steps.\\n\",\n \"\\n\",\n \"* [Asynchronous checkpointing](https://orbax.readthedocs.io/en/latest/guides/checkpoint/async_checkpointing.html).\\n\",\n \"\\n\",\n \"* [Orbax transformations](https://orbax.readthedocs.io/en/latest/guides/checkpoint/transformations.html): A way to modify pytree structure during loading time, instead of after loading time, which is demonstrated in this guide.\"\n ]\n }\n ],\n \"metadata\": {\n \"jupytext\": {\n \"formats\": \"ipynb,md:myst\"\n },\n \"language_info\": {\n \"codemirror_mode\": {\n \"name\": \"ipython\",\n \"version\": 3\n },\n \"file_extension\": \".py\",\n \"mimetype\": \"text/x-python\",\n \"name\": \"python\",\n \"nbconvert_exporter\": \"python\",\n \"pygments_lexer\": \"ipython3\",\n \"version\": \"3.11.9\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 2\n}\n" }, { "path": "docs_nnx/guides/checkpointing.md", "content": "---\njupytext:\n formats: ipynb,md:myst\n text_representation:\n extension: .md\n format_name: myst\n format_version: 0.13\n jupytext_version: 1.13.8\n---\n\n# Save and load checkpoints\n\nThis guide demonstrates how to save and load Flax NNX model checkpoints with [Orbax](https://orbax.readthedocs.io/).\n\n> **Note:** The Flax team does not actively maintain a library for saving and loading model checkpoints to disk. Therefore, it is recommended you use external libraries like [Orbax](https://orbax.readthedocs.io/en/latest/index.html) to do it.\n\nIn this guide you will learn how to:\n\n* Save checkpoints.\n* Restore checkpoints.\n* Restore checkpoints if checkpoint structures differ. \n* Perform multi-process checkpointing. \n\nThe Orbax API examples used throughout the guide are for demonstration purposes, and for the most up-to-date recommended APIs refer to the [Orbax website](https://orbax.readthedocs.io/).\n\n> **Note:** The Flax team recommends using [Orbax](https://orbax.readthedocs.io/en/latest/index.html) for saving and loading checkpoints to disk, as we do not actively maintain a library for these functionalities.\n\n> **Note:** If you are looking for Flax Linen's legacy `flax.training.checkpoints` package, it was deprecated in 2023 in favor of Orbax. The documentation resides on the [Flax Linen site](https://flax-linen.readthedocs.io/en/latest/guides/training_techniques/use_checkpointing.html).\n\n+++\n\n### Setup\n\nImport the necessary dependencies, set up a checkpoint directory and an example Flax NNX model - `TwoLayerMLP` - by subclassing [`nnx.Module`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/module.html).\n\n```{code-cell} ipython3\nfrom flax import nnx\nimport orbax.checkpoint as ocp\nimport jax\nfrom jax import numpy as jnp\nimport numpy as np\n\nckpt_dir = ocp.test_utils.erase_and_create_empty('/tmp/my-checkpoints/')\n```\n\n```{code-cell} ipython3\nclass TwoLayerMLP(nnx.Module):\n def __init__(self, dim, rngs: nnx.Rngs):\n self.linear1 = nnx.Linear(dim, dim, rngs=rngs, use_bias=False)\n self.linear2 = nnx.Linear(dim, dim, rngs=rngs, use_bias=False)\n\n def __call__(self, x):\n x = self.linear1(x)\n return self.linear2(x)\n\n# Instantiate the model and show we can run it.\nmodel = TwoLayerMLP(4, rngs=nnx.Rngs(0))\nx = jax.random.normal(jax.random.key(42), (3, 4))\nassert model(x).shape == (3, 4)\n```\n\n## Save checkpoints\n\nJAX checkpointing libraries, such as Orbax, can save and load any given JAX [pytree](https://jax.readthedocs.io/en/latest/pytrees.html), which is a pure, possibly nested container of [`jax.Array`s)](https://jax.readthedocs.io/en/latest/key-concepts.html#jax-arrays-jax-array) (or, \"tensors\" as some other frameworks would put it). In the context of machine learning, the checkpoint is usually a pytree of model parameters and other data, such as optimizer states.\n\nIn Flax NNX, you can obtain such a pytree from an [`nnx.Module`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/module.html) by calling [`nnx.split`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/graph.html#flax.nnx.split), and picking up the returned [`nnx.State`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/state.html#flax.nnx.State).\n\n```{code-cell} ipython3\n_, state = nnx.split(model)\nnnx.display(state)\n\ncheckpointer = ocp.StandardCheckpointer()\ncheckpointer.save(ckpt_dir / 'state', state)\n```\n\n
\n\n\n\n
\n\n\n## Restore checkpoints\n\nNote that you saved the checkpoint as a Flax class of [`nnx.State`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/state.html#flax.nnx.State), which is also nested with [`nnx.Variable`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Variable) and [`nnx.Param`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Param) classes.\n\nAt checkpoint restoration time, you need to have these classes ready in your runtime, and instruct the checkpointing library (Orbax) to restore your pytree back to that structure. This can be achieved as follows:\n- First, create an abstract Flax NNX model (without allocating any memory for arrays), and show its abstract variable state to the checkpointing library.\n- Once you have the state, use [`nnx.merge`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/graph.html#flax.nnx.merge) to obtain your Flax NNX model, and use it as usual.\n\n```{code-cell} ipython3\n# Restore the checkpoint back to its `nnx.State` structure - need an abstract reference.\nabstract_model = nnx.eval_shape(lambda: TwoLayerMLP(4, rngs=nnx.Rngs(0)))\ngraphdef, abstract_state = nnx.split(abstract_model)\nprint('The abstract NNX state (all leaves are abstract arrays):')\nnnx.display(abstract_state)\n\nstate_restored = checkpointer.restore(ckpt_dir / 'state', abstract_state)\njax.tree.map(np.testing.assert_array_equal, state, state_restored)\nprint('NNX State restored: ')\nnnx.display(state_restored)\n\n# The model is now good to use!\nmodel = nnx.merge(graphdef, state_restored)\nassert model(x).shape == (3, 4)\n```\n\n The abstract NNX state (all leaves are abstract arrays):\n\n\n\n
\n\n\n NNX State restored: \n\n\n /Users/ivyzheng/envs/flax-head/lib/python3.11/site-packages/orbax/checkpoint/type_handlers.py:1439: UserWarning: Couldn't find sharding info under RestoreArgs. Populating sharding info from sharding file. Please note restoration time will be slightly increased due to reading from file instead of directly from RestoreArgs. Note also that this option is unsafe when restoring on a different topology than the checkpoint was saved with.\n warnings.warn(\n\n\n\n
\n\n\n\n
\n\n\n## Save and restore as pure dictionaries\n\nWhen interacting with checkpoint libraries (like Orbax), you may prefer to work with Python built-in container types. In this case, you can use the [`nnx.State.to_pure_dict`](https://github.com/google/flax/blob/764e1732dcd3b8bf178b9ba73ddecf125709b5d7/flax/nnx/statelib.py#L170) and [`nnx.State.replace_by_pure_dict`](https://github.com/google/flax/blob/764e1732dcd3b8bf178b9ba73ddecf125709b5d7/flax/nnx/statelib.py#L179) API to convert an [`nnx.State`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/state.html#flax.nnx.State) to and from pure nested dictionaries.\n\n```{code-cell} ipython3\n# Save as pure dict\npure_dict_state = nnx.to_pure_dict(state)\nnnx.display(pure_dict_state)\ncheckpointer.save(ckpt_dir / 'pure_dict', pure_dict_state)\n\n# Restore as a pure dictionary.\nrestored_pure_dict = checkpointer.restore(ckpt_dir / 'pure_dict')\nabstract_model = nnx.eval_shape(lambda: TwoLayerMLP(4, rngs=nnx.Rngs(0)))\ngraphdef, abstract_state = nnx.split(abstract_model)\nnnx.replace_by_pure_dict(abstract_state, restored_pure_dict)\nmodel = nnx.merge(graphdef, abstract_state)\nassert model(x).shape == (3, 4) # The model still works!\n```\n\n
\n\n\n\n
\n\n\n WARNING:absl:`StandardCheckpointHandler` expects a target tree to be provided for restore. Not doing so is generally UNSAFE unless you know the present topology to be the same one as the checkpoint was saved under.\n\n\n## Restore when checkpoint structures differ\n\nThe ability to load a checkpoint as a pure nested dictionary can come in handy when you want to load some outdated checkpoints that no longer match with your current model code. Check out this simple example below.\n\nThis pattern also works if you save the checkpoint as an [`nnx.State`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/state.html#flax.nnx.State) instead of a pure dictionary. Check out the [Checkpoint surgery section](https://flax.readthedocs.io/en/latest/guides/surgery.html#checkpoint-surgery) of the [Model Surgery](https://flax.readthedocs.io/en/latest/guides/surgery.html) guide for an example with code. The only difference is you need to reprocess your raw dictionary a bit before calling [`nnx.State.replace_by_pure_dict`](https://github.com/google/flax/blob/764e1732dcd3b8bf178b9ba73ddecf125709b5d7/flax/nnx/statelib.py#L179).\n\n```{code-cell} ipython3\nclass ModifiedTwoLayerMLP(nnx.Module):\n \"\"\"A modified version of TwoLayerMLP, which requires bias arrays.\"\"\"\n def __init__(self, dim, rngs: nnx.Rngs):\n self.linear1 = nnx.Linear(dim, dim, rngs=rngs, use_bias=True) # We need bias now!\n self.linear2 = nnx.Linear(dim, dim, rngs=rngs, use_bias=True) # We need bias now!\n\n def __call__(self, x):\n x = self.linear1(x)\n return self.linear2(x)\n\n# Accommodate your old checkpoint to the new code.\nrestored_pure_dict = checkpointer.restore(ckpt_dir / 'pure_dict')\nrestored_pure_dict['linear1']['bias'] = jnp.zeros((4,))\nrestored_pure_dict['linear2']['bias'] = jnp.zeros((4,))\n\n# Same restore code as above.\nabstract_model = nnx.eval_shape(lambda: ModifiedTwoLayerMLP(4, rngs=nnx.Rngs(0)))\ngraphdef, abstract_state = nnx.split(abstract_model)\nnnx.replace_by_pure_dict(abstract_state, restored_pure_dict)\nmodel = nnx.merge(graphdef, abstract_state)\nassert model(x).shape == (3, 4) # The new model works!\n\nnnx.display(model.linear1)\n```\n\n WARNING:absl:`StandardCheckpointHandler` expects a target tree to be provided for restore. Not doing so is generally UNSAFE unless you know the present topology to be the same one as the checkpoint was saved under.\n\n\n\n
\n\n\n\n
\n\n\n## Multi-process checkpointing\n\nIn a multi-host/multi-process environment, you would want to restore your checkpoint as sharded across multiple devices. Check out the [Load sharded model from a checkpoint](https://flax.readthedocs.io/en/latest/guides/flax_gspmd.html#load-sharded-model-from-a-checkpoint) section in the Flax [Scale up on multiple devices](https://flax.readthedocs.io/en/latest/guides/flax_gspmd.html) guide to learn how to derive a sharding pytree and use it to load your checkpoint.\n\n> **Note:** JAX provides several ways to scale up your code on multiple hosts at the same time. This usually happens when the number of devices (CPU/GPU/TPU) is so large that different devices are managed by different hosts (CPU). Check out JAX’s [Introduction to parallel programming](https://jax.readthedocs.io/en/latest/sharded-computation.html), [Using JAX in multi-host and multi-process environments](https://jax.readthedocs.io/en/latest/multi_process.html), [Distributed arrays and automatic parallelization](https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html), and [Manual parallelism with `shard_map`](https://jax.readthedocs.io/en/latest/notebooks/shard_map.html).\n\n+++\n\n## Other checkpointing features\n\nThis guide only uses the simplest [`orbax.checkpoint.StandardCheckpointer`](https://orbax.readthedocs.io/en/latest/api_reference/checkpoint.checkpointers.html#standardcheckpointer) API to show how to save and load on the Flax modeling side. Feel free to use other tools or libraries as you see fit.\n\nIn addition, check out [the Orbax website](https://orbax.readthedocs.io/en/latest/index.html) for other commonly used features, such as:\n\n* [`CheckpointManager`](https://orbax.readthedocs.io/en/latest/guides/checkpoint/api_refactor.html) to track checkpoints from different steps.\n\n* [Asynchronous checkpointing](https://orbax.readthedocs.io/en/latest/guides/checkpoint/async_checkpointing.html).\n\n* [Orbax transformations](https://orbax.readthedocs.io/en/latest/guides/checkpoint/transformations.html): A way to modify pytree structure during loading time, instead of after loading time, which is demonstrated in this guide.\n" }, { "path": "docs_nnx/guides/demo.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"id\": \"a1b37dff\",\n \"metadata\": {},\n \"source\": [\n \"# NNX Demo\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 1,\n \"id\": \"e8099a6f\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"import jax\\n\",\n \"from jax import numpy as jnp\\n\",\n \"from flax import nnx\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"bcc5cffe\",\n \"metadata\": {},\n \"source\": [\n \"### [1] NNX is Pythonic\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 7,\n \"id\": \"d99b73af\",\n \"metadata\": {\n \"outputId\": \"d8ef66d5-6866-4d5c-94c2-d22512bfe718\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"model = MLP(\\n\",\n \" blocks=[Block(\\n\",\n \" linear=Linear(\\n\",\n \" in_features=4,\\n\",\n \" out_features=4,\\n\",\n \" use_bias=True,\\n\",\n \" dtype=None,\\n\",\n \" param_dtype=,\\n\",\n \" precision=None,\\n\",\n \" kernel_init=.init at 0x28ae86dc0>,\\n\",\n \" bias_init=,\\n\",\n \" dot_general=\\n\",\n \" ),\\n\",\n \" bn=BatchNorm(\\n\",\n \" num_features=4,\\n\",\n \" \\n\",\n \"...\\n\"\n ]\n }\n ],\n \"source\": [\n \"\\n\",\n \"class Block(nnx.Module):\\n\",\n \" def __init__(self, din, dout, *, rngs):\\n\",\n \" self.linear = nnx.Linear(din, dout, rngs=rngs)\\n\",\n \" self.bn = nnx.BatchNorm(dout, rngs=rngs)\\n\",\n \"\\n\",\n \" def __call__(self, x):\\n\",\n \" return nnx.relu(self.bn(self.linear(x)))\\n\",\n \"\\n\",\n \"\\n\",\n \"class MLP(nnx.Module):\\n\",\n \" def __init__(self, nlayers, dim, *, rngs): # explicit RNG threading\\n\",\n \" self.blocks = [\\n\",\n \" Block(dim, dim, rngs=rngs) for _ in range(nlayers)\\n\",\n \" ]\\n\",\n \" self.count = Count(0) # stateful variables are defined as attributes\\n\",\n \"\\n\",\n \" def __call__(self, x):\\n\",\n \" self.count.value += 1 # in-place stateful updates\\n\",\n \" for block in self.blocks:\\n\",\n \" x = block(x)\\n\",\n \" return x\\n\",\n \"\\n\",\n \"class Count(nnx.Variable): # custom Variable types define the \\\"collections\\\"\\n\",\n \" pass\\n\",\n \"\\n\",\n \"model = MLP(5, 4, rngs=nnx.Rngs(0)) # no special `init` method\\n\",\n \"model.set_attributes(use_running_average=False) # set flags\\n\",\n \"y = model(jnp.ones((2, 4))) # call methods directly\\n\",\n \"\\n\",\n \"print(f'{model = }'[:500] + '\\\\n...')\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"523aa27c\",\n \"metadata\": {},\n \"source\": [\n \"Because NNX Modules contain their own state, they are very easily to inspect:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 9,\n \"id\": \"6f278ec4\",\n \"metadata\": {\n \"outputId\": \"10a46b0f-2993-4677-c26d-36a4ddf33449\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"model.count = Count(\\n\",\n \" raw_value=1\\n\",\n \")\\n\",\n \"model.blocks[0].linear.kernel = Param(\\n\",\n \" raw_value=Array([[-0.80345297, -0.34071913, -0.9408296 , 0.01005968],\\n\",\n \" [ 0.26146442, 1.1247735 , 0.54563737, -0.374164 ],\\n\",\n \" [ 1.0281805 , -0.6798804 , -0.1488401 , 0.05694951],\\n\",\n \" [-0.44308168, -0.60587114, 0.434087 , -0.40541083]], dtype=float32)\\n\",\n \")\\n\"\n ]\n }\n ],\n \"source\": [\n \"print(f'{model.count = }')\\n\",\n \"print(f'{model.blocks[0].linear.kernel = }')\\n\",\n \"# print(f'{model.blocks.sdf.kernel = }') # typesafe inspection\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"95f389f2\",\n \"metadata\": {},\n \"source\": [\n \"### [2] Model Surgery is Intuitive\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 10,\n \"id\": \"96f61108\",\n \"metadata\": {\n \"outputId\": \"e6f86be8-3537-4c48-f471-316ee0fb6c45\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"y.shape = (2, 4)\\n\"\n ]\n }\n ],\n \"source\": [\n \"# Module sharing\\n\",\n \"model.blocks[1] = model.blocks[3]\\n\",\n \"# Weight tying\\n\",\n \"model.blocks[0].linear.kernel = model.blocks[-1].linear.kernel\\n\",\n \"# Monkey patching\\n\",\n \"def my_optimized_layer(x): return x\\n\",\n \"model.blocks[2] = my_optimized_layer\\n\",\n \"\\n\",\n \"y = model(jnp.ones((2, 4))) # still works\\n\",\n \"print(f'{y.shape = }')\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"aca5a6cd\",\n \"metadata\": {},\n \"source\": [\n \"### [3] Interacting with JAX is easy\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 11,\n \"id\": \"c166dcc7\",\n \"metadata\": {\n \"outputId\": \"9a3f378b-739e-4f45-9968-574651200ede\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"state = State({\\n\",\n \" 'blocks': {\\n\",\n \" '0': {\\n\",\n \" 'linear': {\\n\",\n \" 'kernel': Param(\\n\",\n \" raw_value=Array([[-0.33095378, 0.67149884, 0.33700302, 0.30972847],\\n\",\n \" [ 0.8662822 , -0.11225506, -1.0820619 , -0.9906892 ],\\n\",\n \" [ 0.88298297, -0.2143851 , 0.48143268, 0.6474548 ],\\n\",\n \" [-0.7710582 , 0.3372276 , 0.15487202, 0.6219269 ]], dtype=float32)\\n\",\n \" ),\\n\",\n \" 'bias': Param(\\n\",\n \" raw_value=Array([0., 0., 0., 0.], dtype=float32)\\n\",\n \" \\n\",\n \"...\\n\",\n \"\\n\",\n \"graphdef = GraphDef(\\n\",\n \" type=MLP,\\n\",\n \" index=0,\\n\",\n \" attributes=('blocks', 'count'),\\n\",\n \" subgraphs={\\n\",\n \" 'blocks': GraphDef(\\n\",\n \" type=list,\\n\",\n \" index=1,\\n\",\n \" attributes=('0', '1', '2', '3', '4'),\\n\",\n \" subgraphs={\\n\",\n \" '0': GraphDef(\\n\",\n \" type=Block,\\n\",\n \" index=2,\\n\",\n \" attributes=('line\\n\",\n \"...\\n\"\n ]\n }\n ],\n \"source\": [\n \"graphdef, state = model.split()\\n\",\n \"\\n\",\n \"# state is a dictionary-like JAX pytree\\n\",\n \"print(f'{state = }'[:500] + '\\\\n...')\\n\",\n \"\\n\",\n \"# graphdef is also a JAX pytree, but just metadata\\n\",\n \"print(f'\\\\n{graphdefefefefefef = }'[:300] + '\\\\n...')\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 12,\n \"id\": \"9f03e3af\",\n \"metadata\": {\n \"outputId\": \"0007d357-152a-449e-bcb9-b1b5a91d2d8d\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"y.shape = (2, 4)\\n\",\n \"model.count.value = Array(3, dtype=int32, weak_type=True)\\n\"\n ]\n }\n ],\n \"source\": [\n \"graphdef, state = model.split()\\n\",\n \"\\n\",\n \"@jax.jit\\n\",\n \"def forward(graphdef: nnx.GraphDef, state: nnx.State, x: jax.Array):\\n\",\n \" model = graphdef.merge(state)\\n\",\n \" y = model(x)\\n\",\n \" state, _ = model.split()\\n\",\n \" return y, state\\n\",\n \"\\n\",\n \"x = jnp.ones((2, 4))\\n\",\n \"y, state = forward(graphdef,state, x)\\n\",\n \"\\n\",\n \"model.update(state)\\n\",\n \"\\n\",\n \"print(f'{y.shape = }')\\n\",\n \"print(f'{model.count.value = }')\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 7,\n \"id\": \"9e23dbb4\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"y.shape = (2, 4)\\n\",\n \"model.count = Array(4, dtype=int32, weak_type=True)\\n\"\n ]\n }\n ],\n \"source\": [\n \"params, batch_stats, counts, graphdef = model.split(nnx.Param, nnx.BatchStat, Count)\\n\",\n \"\\n\",\n \"@jax.jit\\n\",\n \"def forward(graphdef: nnx.GraphDef, params, batch_stats, counts, x: jax.Array):\\n\",\n \" model = graphdef.merge(params, batch_stats, counts)\\n\",\n \" y = model(x, train=True)\\n\",\n \" params, batch_stats, counts, _ = model.split(nnx.Param, nnx.BatchStat, Count)\\n\",\n \" return y, params, batch_stats, counts\\n\",\n \"\\n\",\n \"x = jnp.ones((2, 4))\\n\",\n \"y, params, batch_stats, counts = forward(graphdef, params, batch_stats, counts, x)\\n\",\n \"\\n\",\n \"model.update(params, batch_stats, counts)\\n\",\n \"\\n\",\n \"print(f'{y.shape = }')\\n\",\n \"print(f'{model.count = }')\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 14,\n \"id\": \"2461bfe8\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"y.shape = (2, 4)\\n\",\n \"parent.model.count.value = Array(4, dtype=int32, weak_type=True)\\n\"\n ]\n }\n ],\n \"source\": [\n \"class Parent(nnx.Module):\\n\",\n \" def __init__(self, model: MLP):\\n\",\n \" self.model = model\\n\",\n \"\\n\",\n \" def __call__(self, x):\\n\",\n \" params, batch_stats, counts, graphdef = self.model.split(nnx.Param, nnx.BatchStat, Count)\\n\",\n \"\\n\",\n \" @jax.jit\\n\",\n \" def forward(graphdef: nnx.GraphDef, params, batch_stats, counts, x: jax.Array):\\n\",\n \" model = graphdef.merge(params, batch_stats, counts)\\n\",\n \" y = model(x)\\n\",\n \" params, batch_stats, counts, _ = model.split(nnx.Param, nnx.BatchStat, Count)\\n\",\n \" return y, params, batch_stats, counts\\n\",\n \"\\n\",\n \" y, params, batch_stats, counts = forward(graphdef, params, batch_stats, counts, x)\\n\",\n \"\\n\",\n \" self.model.update(params, batch_stats, counts)\\n\",\n \" return y\\n\",\n \"\\n\",\n \"parent = Parent(model)\\n\",\n \"\\n\",\n \"y = parent(jnp.ones((2, 4)))\\n\",\n \"\\n\",\n \"print(f'{y.shape = }')\\n\",\n \"print(f'{parent.model.count.value = }')\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"2e340bcb\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": []\n }\n ],\n \"metadata\": {\n \"jupytext\": {\n \"formats\": \"ipynb,md:myst\"\n },\n \"language_info\": {\n \"codemirror_mode\": {\n \"name\": \"ipython\",\n \"version\": 3\n },\n \"file_extension\": \".py\",\n \"mimetype\": \"text/x-python\",\n \"name\": \"python\",\n \"nbconvert_exporter\": \"python\",\n \"pygments_lexer\": \"ipython3\",\n \"version\": \"3.9.18\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 5\n}\n" }, { "path": "docs_nnx/guides/demo.md", "content": "---\njupytext:\n formats: ipynb,md:myst\n text_representation:\n extension: .md\n format_name: myst\n format_version: 0.13\n jupytext_version: 1.13.8\n---\n\n# NNX Demo\n\n```{code-cell} ipython3\nimport jax\nfrom jax import numpy as jnp\nfrom flax import nnx\n```\n\n### [1] NNX is Pythonic\n\n```{code-cell} ipython3\n:outputId: d8ef66d5-6866-4d5c-94c2-d22512bfe718\n\n\nclass Block(nnx.Module):\n def __init__(self, din, dout, *, rngs):\n self.linear = nnx.Linear(din, dout, rngs=rngs)\n self.bn = nnx.BatchNorm(dout, rngs=rngs)\n\n def __call__(self, x):\n return nnx.relu(self.bn(self.linear(x)))\n\n\nclass MLP(nnx.Module):\n def __init__(self, nlayers, dim, *, rngs): # explicit RNG threading\n self.blocks = [\n Block(dim, dim, rngs=rngs) for _ in range(nlayers)\n ]\n self.count = Count(0) # stateful variables are defined as attributes\n\n def __call__(self, x):\n self.count.value += 1 # in-place stateful updates\n for block in self.blocks:\n x = block(x)\n return x\n\nclass Count(nnx.Variable): # custom Variable types define the \"collections\"\n pass\n\nmodel = MLP(5, 4, rngs=nnx.Rngs(0)) # no special `init` method\nmodel.set_attributes(use_running_average=False) # set flags\ny = model(jnp.ones((2, 4))) # call methods directly\n\nprint(f'{model = }'[:500] + '\\n...')\n```\n\nBecause NNX Modules contain their own state, they are very easily to inspect:\n\n```{code-cell} ipython3\n:outputId: 10a46b0f-2993-4677-c26d-36a4ddf33449\n\nprint(f'{model.count = }')\nprint(f'{model.blocks[0].linear.kernel = }')\n# print(f'{model.blocks.sdf.kernel = }') # typesafe inspection\n```\n\n### [2] Model Surgery is Intuitive\n\n```{code-cell} ipython3\n:outputId: e6f86be8-3537-4c48-f471-316ee0fb6c45\n\n# Module sharing\nmodel.blocks[1] = model.blocks[3]\n# Weight tying\nmodel.blocks[0].linear.kernel = model.blocks[-1].linear.kernel\n# Monkey patching\ndef my_optimized_layer(x): return x\nmodel.blocks[2] = my_optimized_layer\n\ny = model(jnp.ones((2, 4))) # still works\nprint(f'{y.shape = }')\n```\n\n### [3] Interacting with JAX is easy\n\n```{code-cell} ipython3\n:outputId: 9a3f378b-739e-4f45-9968-574651200ede\n\ngraphdef, state = model.split()\n\n# state is a dictionary-like JAX pytree\nprint(f'{state = }'[:500] + '\\n...')\n\n# graphdef is also a JAX pytree, but just metadata\nprint(f'\\n{graphdefefefefefef = }'[:300] + '\\n...')\n```\n\n```{code-cell} ipython3\n:outputId: 0007d357-152a-449e-bcb9-b1b5a91d2d8d\n\ngraphdef, state = model.split()\n\n@jax.jit\ndef forward(graphdef: nnx.GraphDef, state: nnx.State, x: jax.Array):\n model = graphdef.merge(state)\n y = model(x)\n state, _ = model.split()\n return y, state\n\nx = jnp.ones((2, 4))\ny, state = forward(graphdef,state, x)\n\nmodel.update(state)\n\nprint(f'{y.shape = }')\nprint(f'{model.count.value = }')\n```\n\n```{code-cell} ipython3\nparams, batch_stats, counts, graphdef = model.split(nnx.Param, nnx.BatchStat, Count)\n\n@jax.jit\ndef forward(graphdef: nnx.GraphDef, params, batch_stats, counts, x: jax.Array):\n model = graphdef.merge(params, batch_stats, counts)\n y = model(x, train=True)\n params, batch_stats, counts, _ = model.split(nnx.Param, nnx.BatchStat, Count)\n return y, params, batch_stats, counts\n\nx = jnp.ones((2, 4))\ny, params, batch_stats, counts = forward(graphdef, params, batch_stats, counts, x)\n\nmodel.update(params, batch_stats, counts)\n\nprint(f'{y.shape = }')\nprint(f'{model.count = }')\n```\n\n```{code-cell} ipython3\nclass Parent(nnx.Module):\n def __init__(self, model: MLP):\n self.model = model\n\n def __call__(self, x):\n params, batch_stats, counts, graphdef = self.model.split(nnx.Param, nnx.BatchStat, Count)\n\n @jax.jit\n def forward(graphdef: nnx.GraphDef, params, batch_stats, counts, x: jax.Array):\n model = graphdef.merge(params, batch_stats, counts)\n y = model(x)\n params, batch_stats, counts, _ = model.split(nnx.Param, nnx.BatchStat, Count)\n return y, params, batch_stats, counts\n\n y, params, batch_stats, counts = forward(graphdef, params, batch_stats, counts, x)\n\n self.model.update(params, batch_stats, counts)\n return y\n\nparent = Parent(model)\n\ny = parent(jnp.ones((2, 4)))\n\nprint(f'{y.shape = }')\nprint(f'{parent.model.count.value = }')\n```\n\n```{code-cell} ipython3\n\n```\n" }, { "path": "docs_nnx/guides/extracting_intermediates.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"id\": \"0fa991e8\",\n \"metadata\": {},\n \"source\": [\n \"# Extracting intermediate values\\n\",\n \"\\n\",\n \"This guide will show you how to extract intermediate values from a module.\\n\",\n \"Consider a toy neural network with two pieces: a \\\"feature\\\" component that embeds\\n\",\n \"inputs in some feature space, and a \\\"loss\\\" component that operates on those features.\\n\",\n \"We'll want to log these feature components during training to identify any issues with\\n\",\n \"the feature extraction. To do this, we can use the `Module.sow` method.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 2,\n \"id\": \"e4c7c65b\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stderr\",\n \"output_type\": \"stream\",\n \"text\": [\n \"W0317 18:04:12.704562 2028538 cpp_gen_intrinsics.cc:74] Empty bitcode string provided for eigen. Optimizations relying on this IR will be disabled.\\n\"\n ]\n }\n ],\n \"source\": [\n \"from flax import nnx\\n\",\n \"import jax\\n\",\n \"import jax.numpy as jnp\\n\",\n \"from functools import partial\\n\",\n \"\\n\",\n \"class Foo(nnx.Module):\\n\",\n \" def __init__(self, *, rngs: nnx.Rngs):\\n\",\n \" self.dense1 = nnx.Linear(8, 32, rngs=rngs)\\n\",\n \" self.dense2 = nnx.Linear(32, 1, rngs=rngs)\\n\",\n \"\\n\",\n \" def features(self, x, rngs= None):\\n\",\n \" feature = nnx.relu(self.dense1(x))\\n\",\n \" self.sow(nnx.Intermediate, 'features', feature)\\n\",\n \" return feature\\n\",\n \"\\n\",\n \" def loss(self, x_features, y_features):\\n\",\n \" return jnp.sum((x_features - y_features)**2)\\n\",\n \"\\n\",\n \" def __call__(self, x, y):\\n\",\n \" return self.loss(self.features(x), self.features(y))\\n\",\n \"\\n\",\n \"# Instantiate the model.\\n\",\n \"rngs = nnx.Rngs(0)\\n\",\n \"model = Foo(rngs=rngs)\\n\",\n \"\\n\",\n \"# Dummy input for testing\\n\",\n \"x, y = rngs.normal((2,8))\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"c56dd826\",\n \"metadata\": {},\n \"source\": [\n \"Here, `self.sow` will store intermediate values under the key `'features'` in a collection associated with the\\n\",\n \"`nnx.Intermediate` type. If you want to log values to multiple different collections, you can use different subclasses of `nnx.Intermediate`\\n\",\n \"for each collection.\\n\",\n \"\\n\",\n \"Now, we can wrap the module with the `nnx.capture` decorator, which wraps any `Callable` accepting a module as its argument (which includes `nnx.Module`s, their methods, or ordinary functions) to return both the resulting loss as well as any intermediate values stored to the `nnx.Intermediate` collection:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 2,\n \"id\": \"c508f8f3\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"State({\\n\",\n \" 'features': Intermediate(\\n\",\n \" value=((32,), (32,))\\n\",\n \" )\\n\",\n \"})\"\n ]\n },\n \"execution_count\": 2,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"capturing_model = nnx.capture(model, nnx.Intermediate)\\n\",\n \"result, intms = capturing_model(x, y)\\n\",\n \"jax.tree.map(lambda a: a.shape, intms)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"d2f2f609\",\n \"metadata\": {},\n \"source\": [\n \"Note that, by default, sow appends values every time it is called. We can see\\n\",\n \"this in the *features* intermediate values logged above. It contains a tuple with one element for the call on `x` and one for the call on `y`. To override the default append behavior, specify `init_fn` and `reduce_fn` - see `Module.sow()`.\\n\",\n \"\\n\",\n \"## How `nnx.capture` Works\\n\",\n \"\\n\",\n \"`nnx.capture` works by temporarily installing a set of mutable capture buffers on every module in the graph before calling the wrapped function, then harvesting those buffers afterward. Before calling the wrapped function, `capture` walks the entire module graph with `iter_modules`. For each module it sets a `__captures__` attribute: a tuple of Variable instances, one per requested `var_type`. Each Variable holds a plain `dict` that maps sow-key → accumulated value.\\n\",\n \"\\n\",\n \"We can see this `__captures__` tuple by printing out the module contents during a `nnx.capture` call:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 15,\n \"id\": \"1e5c2ae7\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"Captures: (Intermediate(\\n\",\n \" value={}\\n\",\n \"),)\\n\"\n ]\n }\n ],\n \"source\": [\n \"@nnx.capture(nnx.Intermediate)\\n\",\n \"def print_captures(model):\\n\",\n \" print(\\\"Captures:\\\", model.__captures__)\\n\",\n \"_, intms = print_captures(nnx.Module())\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"2daab2c9\",\n \"metadata\": {},\n \"source\": [\n \"`Module.sow` looks for the Variable in the `__captures__` tuple whose type matches `variable_type`, then writes its value into that dict using `reduce_fn`.\\n\",\n \"\\n\",\n \"If no matching type is found, `sow` silently returns `False` without logging the value. This can be used to capture only a subset of the sown values. For example:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 16,\n \"id\": \"159a909b-0c3a-411e-9ccb-98ddecd5720e\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"State({\\n\",\n \" 'gets_sown': Metric1(\\n\",\n \" value=((2,),)\\n\",\n \" )\\n\",\n \"})\"\n ]\n },\n \"execution_count\": 16,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"class Metric1(nnx.Intermediate):\\n\",\n \" pass\\n\",\n \"\\n\",\n \"class Metric2(nnx.Intermediate):\\n\",\n \" pass\\n\",\n \"\\n\",\n \"@nnx.capture(Metric1)\\n\",\n \"def get_captures(model):\\n\",\n \" model.sow(Metric1, 'gets_sown', jnp.ones(2))\\n\",\n \" model.sow(Metric2, 'gets_ignored', jnp.ones(2))\\n\",\n \"_, intms = get_captures(nnx.Module())\\n\",\n \"jax.tree.map(lambda a: a.shape, intms)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"1328ce66\",\n \"metadata\": {},\n \"source\": [\n \"## Capturing all intermediate values\\n\",\n \"\\n\",\n \"To observe the output of each method without manually adding calls to `sow`, we can call `nnx.capture` with the `method_outputs` argument. This will automatically `sow` the output of each method using the given variable type, including methods of sub-modules.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 7,\n \"id\": \"47781215\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"State({\\n\",\n \" '__call__': Intermediate(\\n\",\n \" value=((),)\\n\",\n \" ),\\n\",\n \" 'dense1': {\\n\",\n \" '__call__': Intermediate(\\n\",\n \" value=((32,), (32,))\\n\",\n \" )\\n\",\n \" },\\n\",\n \" 'features': Intermediate(\\n\",\n \" value=((32,), (32,))\\n\",\n \" ),\\n\",\n \" 'loss': Intermediate(\\n\",\n \" value=((),)\\n\",\n \" )\\n\",\n \"})\"\n ]\n },\n \"execution_count\": 7,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"class Foo(nnx.Module):\\n\",\n \" def __init__(self, *, rngs: nnx.Rngs):\\n\",\n \" self.dense1 = nnx.Linear(8, 32, rngs=rngs)\\n\",\n \" self.dense2 = nnx.Linear(32, 1, rngs=rngs)\\n\",\n \"\\n\",\n \" def features(self, x, rngs= None):\\n\",\n \" feature = nnx.relu(self.dense1(x))\\n\",\n \" return feature\\n\",\n \"\\n\",\n \" def loss(self, x_features, y_features):\\n\",\n \" return jnp.sum((x_features - y_features)**2)\\n\",\n \"\\n\",\n \" def __call__(self, x, y):\\n\",\n \" return self.loss(self.features(x), self.features(y))\\n\",\n \"\\n\",\n \"model = Foo(rngs=nnx.Rngs(0))\\n\",\n \"capturing_model = nnx.capture(model, nnx.Intermediate, method_outputs=nnx.Intermediate)\\n\",\n \"result, intms = capturing_model(x, y)\\n\",\n \"jax.tree.map(lambda a: a.shape, intms)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"eee2809a\",\n \"metadata\": {},\n \"source\": [\n \"This pattern should be considered the \\\"sledge hammer\\\" approach to capturing intermediates. As a debugging and inspection tool it is very useful, but using the other patterns described in this guide will give you more fine-grained control over what intermediates you want to extract. We can also combine the `method_output_type` argument with manual calls to sow to capture both layer outputs and computations mid-layer.\\n\",\n \"\\n\",\n \"## Extracting gradients of intermediate values\\n\",\n \"\\n\",\n \"For debugging purposes, it can be useful to extract the gradients of intermediate values. This is a little tricky: jax doesn't have a stable mechanism for sowing information from the backward pass into to objects from the forward pass. Instead, we record the gradients of intermediate values using the `Module.perturb` method as follows:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 30,\n \"id\": \"84911b66\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"State({\\n\",\n \" 'grad_of_x': Perturbation( # 1 (4 B)\\n\",\n \" value=Array(3., dtype=float32, weak_type=True)\\n\",\n \" ),\\n\",\n \" 'activations': Intermediate( # 1 (4 B)\\n\",\n \" value=(Array(1., dtype=float32, weak_type=True),)\\n\",\n \" )\\n\",\n \"})\"\n ]\n },\n \"execution_count\": 30,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"class Model(nnx.Module):\\n\",\n \" def __call__(self, x):\\n\",\n \" x2 = self.perturb('grad_of_x', x)\\n\",\n \" self.sow(nnx.Intermediate, 'activations', x2)\\n\",\n \" return 3 * x2\\n\",\n \"\\n\",\n \"model = Model()\\n\",\n \"\\n\",\n \"def train_step(model, x):\\n\",\n \" _, perturbations = nnx.capture(model, nnx.Perturbation)(x)\\n\",\n \" def loss(model, perturbations, x):\\n\",\n \" return nnx.capture(model, nnx.Intermediate, init=perturbations)(x)\\n\",\n \"\\n\",\n \" (grads, perturb_grads), sowed = nnx.grad(loss, argnums=(0, 1), has_aux=True)(model, perturbations, x)\\n\",\n \" return nnx.merge_state(perturb_grads, sowed)\\n\",\n \"\\n\",\n \"train_step(model, 1.0)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"7dc978af\",\n \"metadata\": {},\n \"source\": [\n \"There are four steps:\\n\",\n \"\\n\",\n \"**Step One: Initialize *perturbations* of the model**.\\n\",\n \"\\n\",\n \"We do this with a call to `nnx.capture(model, nnx.Perturbation)`. Before the call, `capture` installs `__captures__` on the module — a tuple containing one empty `Perturbation` buffer (as described in \\\"How `nnx.capture` Works\\\" above). When `self.perturb` runs, it checks `__captures__` for a matching `Perturbation` Variable, initialises the slot to `zeros_like(value)`, and returns `zeros + x`. After the call, `__captures__` is removed and the filled buffer is returned as `perturbations`.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 12,\n \"id\": \"5f4fda80\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"before perturb: (Perturbation(\\n\",\n \" value={}\\n\",\n \"),)\\n\",\n \"after perturb: (Perturbation( # 1 (4 B)\\n\",\n \" value={'grad_of_x': Array(0., dtype=float32, weak_type=True)}\\n\",\n \"),)\\n\",\n \"\\u001b[38;2;79;201;177mState\\u001b[0m\\u001b[38;2;255;213;3m({\\u001b[0m\\u001b[38;2;105;105;105m\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254m'grad_of_x'\\u001b[0m\\u001b[38;2;212;212;212m: \\u001b[0m\\u001b[38;2;79;201;177mPerturbation\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;105;105;105m # 1 (4 B)\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254mvalue\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0mArray(0., dtype=float32, weak_type=True)\\n\",\n \" \\u001b[38;2;255;213;3m)\\u001b[0m\\n\",\n \"\\u001b[38;2;255;213;3m})\\u001b[0m\\n\"\n ]\n }\n ],\n \"source\": [\n \"class Model(nnx.Module):\\n\",\n \" def __call__(self, x):\\n\",\n \" print(\\\"before perturb:\\\", self.__captures__)\\n\",\n \" x2 = self.perturb('grad_of_x', x)\\n\",\n \" print(\\\"after perturb:\\\", self.__captures__)\\n\",\n \" self.sow(nnx.Intermediate, 'activations', x2)\\n\",\n \" # sow is a no-op: Intermediate is not in __captures__, so it returns False silently\\n\",\n \" return 3 * x2\\n\",\n \"\\n\",\n \"model = Model()\\n\",\n \"_, perturbations = nnx.capture(model, nnx.Perturbation)(1.0)\\n\",\n \"print(perturbations)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"2b4e1d98\",\n \"metadata\": {},\n \"source\": [\n \"There are only two differences between `sow` and `perturb`:\\n\",\n \"\\n\",\n \"- The `nnx.Variable` tag used for values written with `self.perturb` is `nnx.Perturbation` rather than `nnx.Intermediate`.\\n\",\n \" \\n\",\n \"- `perturb` returns the logged value. You must use this returned value rather than the original value for the gradient capturing machinery to work.\\n\",\n \"\\n\",\n \"The `var_types` argument to `capture` restricts which of the logged values we want to return. Because we only want the intermediates logged with `self.perturb` statements, we only capture `nnx.Perturbation` types.\\n\",\n \"\\n\",\n \"**Step Two: Run the model again, but add in these perturbations**.\\n\",\n \"\\n\",\n \"Call `capture` again with `init=perturbations`. `capture` first builds a mapping from module path to the Variables in `init`, then uses it to pre-populate `__captures__`. Now `__captures__` has *two* buffers: an empty `Intermediate` buffer (from `var_types`) and a `Perturbation` buffer pre-populated from `init`. `self.perturb` finds the pre-populated buffer and returns `x + perturbation`; `self.sow` writes into the `Intermediate` buffer as normal.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 13,\n \"id\": \"a4087d73\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"before perturb: (Intermediate(\\n\",\n \" value={}\\n\",\n \"), Perturbation( # 1 (4 B)\\n\",\n \" value={'grad_of_x': Array(0., dtype=float32, weak_type=True)}\\n\",\n \"))\\n\",\n \"after sow: (Intermediate( # 1 (4 B)\\n\",\n \" value={'activations': (Array(1., dtype=float32, weak_type=True),)}\\n\",\n \"), Perturbation( # 1 (4 B)\\n\",\n \" value={'grad_of_x': Array(0., dtype=float32, weak_type=True)}\\n\",\n \"))\\n\"\n ]\n }\n ],\n \"source\": [\n \"class Model(nnx.Module):\\n\",\n \" def __call__(self, x):\\n\",\n \" print(\\\"before perturb:\\\", self.__captures__)\\n\",\n \" x2 = self.perturb('grad_of_x', x)\\n\",\n \" self.sow(nnx.Intermediate, 'activations', x2)\\n\",\n \" print(\\\"after sow: \\\", self.__captures__)\\n\",\n \" return 3 * x2\\n\",\n \"\\n\",\n \"model = Model()\\n\",\n \"_, interms = nnx.capture(model, nnx.Intermediate, init=perturbations)(1.0)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"63e2ba59\",\n \"metadata\": {},\n \"source\": [\n \"This changes the behavior of `x2 = self.perturb('name', x)` to essentially be `x2 = x + perturbations['name']`. The gradient of our output with respect to `x` will be the same as the gradient with respect to the perturbation, because JAX can differentiate through the addition with respect to the perturbation value stored in the capture dict.\\n\",\n \"\\n\",\n \"**Step Three: Take gradients**.\\n\",\n \"\\n\",\n \"Specifically, take the gradient of this second `capture` call with respect to the perturbation arguments. JAX traces through exactly the same `__captures__` setup as Step Two, but with abstract (traced) array values instead of concrete ones. This will give us the same values as the gradients with respect to the intermediate variables. If we want to track intermediate variables in the forward pass at the same time, we'll need to return the intermediate values output of the `capture` call as well, so we'll need to pass `has_aux=True` to `nnx.grad`.\\n\",\n \"\\n\",\n \"**Step Four: Combine intermediate states**\\n\",\n \"\\n\",\n \"Merge the `State` object we get from the perturbation gradients with the `State` object for forward intermediates with `nnx.merge_state(perturb_grads, sowed)`. At this point `__captures__` no longer exists on any module — it was cleaned up at the end of the `capture` call in Step Three.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"23ccf952\",\n \"metadata\": {},\n \"source\": [\n \"## NNX Transforms and Capturing\\n\",\n \"\\n\",\n \"`nnx.capture` composes with NNX transforms such as `nnx.vmap`. The main thing to keep in mind is that perturbations must be initialized with a run that has the same batch structure as the training step that will consume them.\\n\",\n \"\\n\",\n \"Consider a model that calls both `sow` and `perturb`:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 34,\n \"id\": \"7c8f8d83\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"class Foo(nnx.Module):\\n\",\n \" def __init__(self, dim):\\n\",\n \" self.w = nnx.Param(jax.random.normal(jax.random.key(0), dim))\\n\",\n \"\\n\",\n \" def __call__(self, x):\\n\",\n \" x = self.perturb('grad_of_x', x)\\n\",\n \" y = jnp.dot(x, self.w)\\n\",\n \" self.sow(nnx.Intermediate, 'y', y)\\n\",\n \" return y\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"786278eb\",\n \"metadata\": {},\n \"source\": [\n \"The training step vmaps `loss_grad` over a batch of inputs and perturbations, while the model weights are shared across the batch (`in_axes=None`):\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 35,\n \"id\": \"8ac86ed1\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"@nnx.jit\\n\",\n \"def train_step(model, x):\\n\",\n \" _, perturbations = init_perturbations(model, x)\\n\",\n \" def loss_grad(model, perturbations, x):\\n\",\n \" def loss(model, perturbations, x):\\n\",\n \" loss, interms = nnx.capture(model, nnx.Intermediate, init=perturbations)(x)\\n\",\n \" return loss, interms\\n\",\n \" (grads, perturb_grads), sowed = nnx.grad(loss, argnums=(0, 1), has_aux=True)(model, perturbations, x)\\n\",\n \" return grads, nnx.merge_state(perturb_grads, sowed)\\n\",\n \" return nnx.vmap(loss_grad, in_axes=(None, 0, 0))(model, perturbations, x)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"3de42b97-9923-436f-9db0-d9aeedc259ad\",\n \"metadata\": {},\n \"source\": [\n \"After every training step, we can sum the gradients and pass them to an `Optimizer` to adjust the model, as usual. But we can also look at the full batch of sown values and perturbations.\\n\",\n \"\\n\",\n \"Because `train_step` expects `perturbations` to have a leading batch axis (axis 0), the perturbation initialization run must also produce a batched `perturbations` state. We do this inside an `init_perturbations` method that splits the model and vmaps the run with `in_axes=(0, None, 0)` for `(intermediates, params, x)`.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 36,\n \"id\": \"76c291c8\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"@nnx.capture(nnx.Perturbation)\\n\",\n \"def init_perturbations(model, x):\\n\",\n \" graphdef, intms, params = nnx.split(model, nnx.Intermediate, nnx.Param)\\n\",\n \" def forward(intms, params, x):\\n\",\n \" return nnx.merge(graphdef, intms, params)(x)\\n\",\n \" return nnx.vmap(forward, in_axes=(0, None, 0))(intms, params, x)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"981642f6\",\n \"metadata\": {},\n \"source\": [\n \"Putting it together:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 37,\n \"id\": \"7a741ca4\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"State({\\n\",\n \" 'grad_of_x': Perturbation(\\n\",\n \" value=(3, 4)\\n\",\n \" ),\\n\",\n \" 'y': Intermediate(\\n\",\n \" value=((3,),)\\n\",\n \" )\\n\",\n \"})\"\n ]\n },\n \"execution_count\": 37,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"model, x = Foo(4), jnp.ones((3, 4))\\n\",\n \"_, intermediates = train_step(model, x)\\n\",\n \"jax.tree.map(lambda a: a.shape, intermediates)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"027216c7\",\n \"metadata\": {},\n \"source\": [\n \"The pattern generalises: whenever a transform introduces a new batch axis over which `capture` runs, initialize perturbations with a matching vmapped pre-run so that the `init=perturbations` argument inside the transform has the correct shape.\"\n ]\n }\n ],\n \"metadata\": {\n \"jupytext\": {\n \"formats\": \"ipynb,md\",\n \"main_language\": \"python\"\n },\n \"language_info\": {\n \"codemirror_mode\": {\n \"name\": \"ipython\",\n \"version\": 3\n },\n \"file_extension\": \".py\",\n \"mimetype\": \"text/x-python\",\n \"name\": \"python\",\n \"nbconvert_exporter\": \"python\",\n \"pygments_lexer\": \"ipython3\",\n \"version\": \"3.13.5\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 5\n}\n" }, { "path": "docs_nnx/guides/extracting_intermediates.md", "content": "---\njupyter:\n jupytext:\n formats: ipynb,md\n main_language: python\n text_representation:\n extension: .md\n format_name: markdown\n format_version: '1.3'\n jupytext_version: 1.13.8\n---\n\n# Extracting intermediate values\n\nThis guide will show you how to extract intermediate values from a module.\nConsider a toy neural network with two pieces: a \"feature\" component that embeds\ninputs in some feature space, and a \"loss\" component that operates on those features.\nWe'll want to log these feature components during training to identify any issues with\nthe feature extraction. To do this, we can use the `Module.sow` method.\n\n```python\nfrom flax import nnx\nimport jax\nimport jax.numpy as jnp\nfrom functools import partial\n\nclass Foo(nnx.Module):\n def __init__(self, *, rngs: nnx.Rngs):\n self.dense1 = nnx.Linear(8, 32, rngs=rngs)\n self.dense2 = nnx.Linear(32, 1, rngs=rngs)\n\n def features(self, x, rngs= None):\n feature = nnx.relu(self.dense1(x))\n self.sow(nnx.Intermediate, 'features', feature)\n return feature\n\n def loss(self, x_features, y_features):\n return jnp.sum((x_features - y_features)**2)\n\n def __call__(self, x, y):\n return self.loss(self.features(x), self.features(y))\n\n# Instantiate the model.\nrngs = nnx.Rngs(0)\nmodel = Foo(rngs=rngs)\n\n# Dummy input for testing\nx, y = rngs.normal((2,8))\n```\n\nHere, `self.sow` will store intermediate values under the key `'features'` in a collection associated with the\n`nnx.Intermediate` type. If you want to log values to multiple different collections, you can use different subclasses of `nnx.Intermediate`\nfor each collection.\n\nNow, we can wrap the module with the `nnx.capture` decorator, which wraps any `Callable` accepting a module as its argument (which includes `nnx.Module`s, their methods, or ordinary functions) to return both the resulting loss as well as any intermediate values stored to the `nnx.Intermediate` collection:\n\n```python\ncapturing_model = nnx.capture(model, nnx.Intermediate)\nresult, intms = capturing_model(x, y)\njax.tree.map(lambda a: a.shape, intms)\n```\n\nNote that, by default, sow appends values every time it is called. We can see\nthis in the *features* intermediate values logged above. It contains a tuple with one element for the call on `x` and one for the call on `y`. To override the default append behavior, specify `init_fn` and `reduce_fn` - see `Module.sow()`.\n\n## How `nnx.capture` Works\n\n`nnx.capture` works by temporarily installing a set of mutable capture buffers on every module in the graph before calling the wrapped function, then harvesting those buffers afterward. Before calling the wrapped function, `capture` walks the entire module graph with `iter_modules`. For each module it sets a `__captures__` attribute: a tuple of Variable instances, one per requested `var_type`. Each Variable holds a plain `dict` that maps sow-key → accumulated value.\n\nWe can see this `__captures__` tuple by printing out the module contents during a `nnx.capture` call:\n\n```python\n@nnx.capture(nnx.Intermediate)\ndef print_captures(model):\n print(\"Captures:\", model.__captures__)\n_, intms = print_captures(nnx.Module())\n```\n\n`Module.sow` looks for the Variable in the `__captures__` tuple whose type matches `variable_type`, then writes its value into that dict using `reduce_fn`.\n\nIf no matching type is found, `sow` silently returns `False` without logging the value. This can be used to capture only a subset of the sown values. For example:\n\n```python\nclass Metric1(nnx.Intermediate):\n pass\n\nclass Metric2(nnx.Intermediate):\n pass\n\n@nnx.capture(Metric1)\ndef get_captures(model):\n model.sow(Metric1, 'gets_sown', jnp.ones(2))\n model.sow(Metric2, 'gets_ignored', jnp.ones(2))\n_, intms = get_captures(nnx.Module())\njax.tree.map(lambda a: a.shape, intms)\n```\n\n## Capturing all intermediate values\n\nTo observe the output of each method without manually adding calls to `sow`, we can call `nnx.capture` with the `method_outputs` argument. This will automatically `sow` the output of each method using the given variable type, including methods of sub-modules.\n\n```python\nclass Foo(nnx.Module):\n def __init__(self, *, rngs: nnx.Rngs):\n self.dense1 = nnx.Linear(8, 32, rngs=rngs)\n self.dense2 = nnx.Linear(32, 1, rngs=rngs)\n\n def features(self, x, rngs= None):\n feature = nnx.relu(self.dense1(x))\n return feature\n\n def loss(self, x_features, y_features):\n return jnp.sum((x_features - y_features)**2)\n\n def __call__(self, x, y):\n return self.loss(self.features(x), self.features(y))\n\nmodel = Foo(rngs=nnx.Rngs(0))\ncapturing_model = nnx.capture(model, nnx.Intermediate, method_outputs=nnx.Intermediate)\nresult, intms = capturing_model(x, y)\njax.tree.map(lambda a: a.shape, intms)\n```\n\nThis pattern should be considered the \"sledge hammer\" approach to capturing intermediates. As a debugging and inspection tool it is very useful, but using the other patterns described in this guide will give you more fine-grained control over what intermediates you want to extract. We can also combine the `method_output_type` argument with manual calls to sow to capture both layer outputs and computations mid-layer.\n\n## Extracting gradients of intermediate values\n\nFor debugging purposes, it can be useful to extract the gradients of intermediate values. This is a little tricky: jax doesn't have a stable mechanism for sowing information from the backward pass into to objects from the forward pass. Instead, we record the gradients of intermediate values using the `Module.perturb` method as follows:\n\n```python\nclass Model(nnx.Module):\n def __call__(self, x):\n x2 = self.perturb('grad_of_x', x)\n self.sow(nnx.Intermediate, 'activations', x2)\n return 3 * x2\n\nmodel = Model()\n\ndef train_step(model, x):\n _, perturbations = nnx.capture(model, nnx.Perturbation)(x)\n def loss(model, perturbations, x):\n return nnx.capture(model, nnx.Intermediate, init=perturbations)(x)\n\n (grads, perturb_grads), sowed = nnx.grad(loss, argnums=(0, 1), has_aux=True)(model, perturbations, x)\n return nnx.merge_state(perturb_grads, sowed)\n\ntrain_step(model, 1.0)\n```\n\nThere are four steps:\n\n**Step One: Initialize *perturbations* of the model**.\n\nWe do this with a call to `nnx.capture(model, nnx.Perturbation)`. Before the call, `capture` installs `__captures__` on the module — a tuple containing one empty `Perturbation` buffer (as described in \"How `nnx.capture` Works\" above). When `self.perturb` runs, it checks `__captures__` for a matching `Perturbation` Variable, initialises the slot to `zeros_like(value)`, and returns `zeros + x`. After the call, `__captures__` is removed and the filled buffer is returned as `perturbations`.\n\n```python\nclass Model(nnx.Module):\n def __call__(self, x):\n print(\"before perturb:\", self.__captures__)\n x2 = self.perturb('grad_of_x', x)\n print(\"after perturb:\", self.__captures__)\n self.sow(nnx.Intermediate, 'activations', x2)\n # sow is a no-op: Intermediate is not in __captures__, so it returns False silently\n return 3 * x2\n\nmodel = Model()\n_, perturbations = nnx.capture(model, nnx.Perturbation)(1.0)\nprint(perturbations)\n```\n\nThere are only two differences between `sow` and `perturb`:\n\n- The `nnx.Variable` tag used for values written with `self.perturb` is `nnx.Perturbation` rather than `nnx.Intermediate`.\n \n- `perturb` returns the logged value. You must use this returned value rather than the original value for the gradient capturing machinery to work.\n\nThe `var_types` argument to `capture` restricts which of the logged values we want to return. Because we only want the intermediates logged with `self.perturb` statements, we only capture `nnx.Perturbation` types.\n\n**Step Two: Run the model again, but add in these perturbations**.\n\nCall `capture` again with `init=perturbations`. `capture` first builds a mapping from module path to the Variables in `init`, then uses it to pre-populate `__captures__`. Now `__captures__` has *two* buffers: an empty `Intermediate` buffer (from `var_types`) and a `Perturbation` buffer pre-populated from `init`. `self.perturb` finds the pre-populated buffer and returns `x + perturbation`; `self.sow` writes into the `Intermediate` buffer as normal.\n\n```python\nclass Model(nnx.Module):\n def __call__(self, x):\n print(\"before perturb:\", self.__captures__)\n x2 = self.perturb('grad_of_x', x)\n self.sow(nnx.Intermediate, 'activations', x2)\n print(\"after sow: \", self.__captures__)\n return 3 * x2\n\nmodel = Model()\n_, interms = nnx.capture(model, nnx.Intermediate, init=perturbations)(1.0)\n```\n\nThis changes the behavior of `x2 = self.perturb('name', x)` to essentially be `x2 = x + perturbations['name']`. The gradient of our output with respect to `x` will be the same as the gradient with respect to the perturbation, because JAX can differentiate through the addition with respect to the perturbation value stored in the capture dict.\n\n**Step Three: Take gradients**.\n\nSpecifically, take the gradient of this second `capture` call with respect to the perturbation arguments. JAX traces through exactly the same `__captures__` setup as Step Two, but with abstract (traced) array values instead of concrete ones. This will give us the same values as the gradients with respect to the intermediate variables. If we want to track intermediate variables in the forward pass at the same time, we'll need to return the intermediate values output of the `capture` call as well, so we'll need to pass `has_aux=True` to `nnx.grad`.\n\n**Step Four: Combine intermediate states**\n\nMerge the `State` object we get from the perturbation gradients with the `State` object for forward intermediates with `nnx.merge_state(perturb_grads, sowed)`. At this point `__captures__` no longer exists on any module — it was cleaned up at the end of the `capture` call in Step Three.\n\n\n## NNX Transforms and Capturing\n\n`nnx.capture` composes with NNX transforms such as `nnx.vmap`. The main thing to keep in mind is that perturbations must be initialized with a run that has the same batch structure as the training step that will consume them.\n\nConsider a model that calls both `sow` and `perturb`:\n\n```python\nclass Foo(nnx.Module):\n def __init__(self, dim):\n self.w = nnx.Param(jax.random.normal(jax.random.key(0), dim))\n\n def __call__(self, x):\n x = self.perturb('grad_of_x', x)\n y = jnp.dot(x, self.w)\n self.sow(nnx.Intermediate, 'y', y)\n return y\n```\n\nThe training step vmaps `loss_grad` over a batch of inputs and perturbations, while the model weights are shared across the batch (`in_axes=None`):\n\n```python\n@nnx.jit\ndef train_step(model, x):\n _, perturbations = init_perturbations(model, x)\n def loss_grad(model, perturbations, x):\n def loss(model, perturbations, x):\n loss, interms = nnx.capture(model, nnx.Intermediate, init=perturbations)(x)\n return loss, interms\n (grads, perturb_grads), sowed = nnx.grad(loss, argnums=(0, 1), has_aux=True)(model, perturbations, x)\n return grads, nnx.merge_state(perturb_grads, sowed)\n return nnx.vmap(loss_grad, in_axes=(None, 0, 0))(model, perturbations, x)\n```\n\nAfter every training step, we can sum the gradients and pass them to an `Optimizer` to adjust the model, as usual. But we can also look at the full batch of sown values and perturbations.\n\nBecause `train_step` expects `perturbations` to have a leading batch axis (axis 0), the perturbation initialization run must also produce a batched `perturbations` state. We do this inside an `init_perturbations` method that splits the model and vmaps the run with `in_axes=(0, None, 0)` for `(intermediates, params, x)`.\n\n```python\n@nnx.capture(nnx.Perturbation)\ndef init_perturbations(model, x):\n graphdef, intms, params = nnx.split(model, nnx.Intermediate, nnx.Param)\n def forward(intms, params, x):\n return nnx.merge(graphdef, intms, params)(x)\n return nnx.vmap(forward, in_axes=(0, None, 0))(intms, params, x)\n```\n\nPutting it together:\n\n```python\nmodel, x = Foo(4), jnp.ones((3, 4))\n_, intermediates = train_step(model, x)\njax.tree.map(lambda a: a.shape, intermediates)\n```\n\nThe pattern generalises: whenever a transform introduces a new batch axis over which `capture` runs, initialize perturbations with a matching vmapped pre-run so that the `init=perturbations` argument inside the transform has the correct shape.\n" }, { "path": "docs_nnx/guides/filters_guide.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"id\": \"95b08e64\",\n \"metadata\": {},\n \"source\": [\n \"# Filters\\n\",\n \"\\n\",\n \"Flax NNX uses [`Filter`s](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/filterlib.html) extensively as a way to create [`nnx.State`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/state.html#flax.nnx.State) groups in APIs, such as [`nnx.split`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/graph.html#flax.nnx.split), [`nnx.state()`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/graph.html#flax.nnx.state), and many of the [Flax NNX transformations (transforms)](https://flax.readthedocs.io/en/latest/guides/jax_and_nnx_transforms.html).\\n\",\n \"\\n\",\n \"In this guide you will learn how to:\\n\",\n \"\\n\",\n \"* Use [`Filter`s](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/filterlib.html) to group Flax NNX variables and states into subgroups;\\n\",\n \"* Understand relationships between types, such as [`nnx.Param`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Param) or [`nnx.BatchStat`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.BatchStat), and [`Filter`s](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/filterlib.html);\\n\",\n \"* Express your `Filter`s flexibly with [`nnx.filterlib.Filter`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/filterlib.html) language.\\n\",\n \"\\n\",\n \"In the following example [`nnx.Param`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Param) and [`nnx.BatchStat`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.BatchStat) are used as `Filter`s to split the model into two groups: one with the parameters and the other with the batch statistics:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 1,\n \"id\": \"45485345\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"params = State({\\n\",\n \" 'a': Param(\\n\",\n \" value=0\\n\",\n \" )\\n\",\n \"})\\n\",\n \"batch_stats = State({\\n\",\n \" 'b': BatchStat(\\n\",\n \" value=True\\n\",\n \" )\\n\",\n \"})\\n\"\n ]\n }\n ],\n \"source\": [\n \"from flax import nnx\\n\",\n \"\\n\",\n \"class Foo(nnx.Module):\\n\",\n \" def __init__(self):\\n\",\n \" self.a = nnx.Param(0)\\n\",\n \" self.b = nnx.BatchStat(True)\\n\",\n \"\\n\",\n \"foo = Foo()\\n\",\n \"\\n\",\n \"graphdef, params, batch_stats = nnx.split(foo, nnx.Param, nnx.BatchStat)\\n\",\n \"\\n\",\n \"print(f'{params = }')\\n\",\n \"print(f'{batch_stats = }')\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"8f77e99a\",\n \"metadata\": {},\n \"source\": [\n \"Let's dive deeper into `Filter`s.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"a0413d64\",\n \"metadata\": {},\n \"source\": [\n \"## The `Filter` Protocol\\n\",\n \"\\n\",\n \"In general, Flax `Filter`s are predicate functions of the form:\\n\",\n \"\\n\",\n \"```python\\n\",\n \"\\n\",\n \"(path: tuple[Key, ...], value: Any) -> bool\\n\",\n \"\\n\",\n \"```\\n\",\n \"\\n\",\n \"where:\\n\",\n \"\\n\",\n \"- `Key` is a hashable and comparable type;\\n\",\n \"- `path` is a tuple of `Key`s representing the path to the value in a nested structure; and\\n\",\n \"- `value` is the value at the path.\\n\",\n \"\\n\",\n \"The function returns `True` if the value should be included in the group, and `False` otherwise.\\n\",\n \"\\n\",\n \"Types are not functions of this form. They are treated as `Filter`s because, as you will learn in the next section, types and some other literals are converted to _predicates_. For example, [`nnx.Param`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Param) is roughly converted to a predicate like this:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 2,\n \"id\": \"30f4c868\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"is_param((), nnx.Param(0)) = True\\n\"\n ]\n }\n ],\n \"source\": [\n \"def is_param(path, value) -> bool:\\n\",\n \" return isinstance(value, nnx.Param)\\n\",\n \"\\n\",\n \"print(f'{is_param((), nnx.Param(0)) = }')\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"a8a2641e\",\n \"metadata\": {},\n \"source\": [\n \"Such function matches any value that is an instance of [`nnx.Param`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Param). Internally Flax NNX uses `OfType` which defines a callable of this form for a given type:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 3,\n \"id\": \"b3095221\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"is_param((), nnx.Param(0)) = True\\n\"\n ]\n }\n ],\n \"source\": [\n \"is_param = nnx.OfType(nnx.Param)\\n\",\n \"\\n\",\n \"print(f'{is_param((), nnx.Param(0)) = }')\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"87c06e39\",\n \"metadata\": {},\n \"source\": [\n \"## The `Filter` DSL\\n\",\n \"\\n\",\n \"Flax NNX exposes a small domain specific language ([DSL](https://en.wikipedia.org/wiki/Domain-specific_language)), formalized as the [`nnx.filterlib.Filter`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/filterlib.html) type. This means users don't have to create functions like in the previous section.\\n\",\n \"\\n\",\n \"Here is a list of all the callable `Filter`s included in Flax NNX, and their corresponding DSL literals (when available):\\n\",\n \"\\n\",\n \"\\n\",\n \"| Literal | Callable | Description |\\n\",\n \"|--------|----------------------|-------------|\\n\",\n \"| `...` or `True` | `Everything()` | Matches all values |\\n\",\n \"| `None` or `False` | `Nothing()` | Matches no values |\\n\",\n \"| `type` | `OfType(type)` | Matches values that are instances of `type` or have a `type` attribute that is an instance of `type` |\\n\",\n \"| | `PathContains(key)` | Matches values that have an associated `path` that contains the given `key` |\\n\",\n \"| `'{filter}'` str | `WithTag('{filter}')` | Matches values that have string `tag` attribute equal to `'{filter}'`. Used by `RngKey` and `RngCount`. |\\n\",\n \"| `(*filters)` tuple or `[*filters]` list | `Any(*filters)` | Matches values that match any of the inner `filters` |\\n\",\n \"| | `All(*filters)` | Matches values that match all of the inner `filters` |\\n\",\n \"| | `Not(filter)` | Matches values that do not match the inner `filter` |\\n\",\n \"\\n\",\n \"\\n\",\n \"Let's check out the DSL in action by using [`nnx.vmap`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.vmap) as an example. Consider the following:\\n\",\n \"\\n\",\n \"1) You want to vectorize all parameters;\\n\",\n \"2) Apply `'dropout'` `Rng(Keys|Counts)` on the `0`th axis; and\\n\",\n \"3) Broadcast the rest.\\n\",\n \"\\n\",\n \"To do this, you can use the following `Filter`s to define a `nnx.StateAxes` object that you can pass to `nnx.vmap`'s `in_axes` to specify how the `model`'s various sub-states should be vectorized:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 4,\n \"id\": \"d38b7694\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"state_axes = nnx.StateAxes({(nnx.Param, 'dropout'): 0, ...: None})\\n\",\n \"\\n\",\n \"@nnx.vmap(in_axes=(state_axes, 0))\\n\",\n \"def forward(model, x):\\n\",\n \" ...\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"bd60f0e1\",\n \"metadata\": {},\n \"source\": [\n \"Here `(nnx.Param, 'dropout')` expands to `Any(OfType(nnx.Param), WithTag('dropout'))` and `...` expands to `Everything()`.\\n\",\n \"\\n\",\n \"If you wish to manually convert literal into a predicate, you can use [`nnx.filterlib.to_predicate`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/filterlib.html#flax.nnx.filterlib.to_predicate):\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 5,\n \"id\": \"7e065fa9\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"is_param = OfType()\\n\",\n \"everything = Everything()\\n\",\n \"nothing = Nothing()\\n\",\n \"params_or_dropout = Any(OfType(), WithTag('dropout'))\\n\"\n ]\n }\n ],\n \"source\": [\n \"is_param = nnx.filterlib.to_predicate(nnx.Param)\\n\",\n \"everything = nnx.filterlib.to_predicate(...)\\n\",\n \"nothing = nnx.filterlib.to_predicate(False)\\n\",\n \"params_or_dropout = nnx.filterlib.to_predicate((nnx.Param, 'dropout'))\\n\",\n \"\\n\",\n \"print(f'{is_param = }')\\n\",\n \"print(f'{everything = }')\\n\",\n \"print(f'{nothing = }')\\n\",\n \"print(f'{params_or_dropout = }')\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"db9b4cf3\",\n \"metadata\": {},\n \"source\": [\n \"## Grouping `State`s\\n\",\n \"\\n\",\n \"With the knowledge of `Filter`s from previous sections at hand, let's learn how to roughly implement [`nnx.split`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/graph.html#flax.nnx.split). Here are the key ideas:\\n\",\n \"\\n\",\n \"* Use `nnx.graph.flatten` to get the [`GraphDef`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/graph.html#flax.nnx.GraphDef) and [`nnx.State`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/state.html#flax.nnx.State) representation of the node.\\n\",\n \"* Convert all the `Filter`s to predicates.\\n\",\n \"* Use `State.flat_state` to get the flat representation of the state.\\n\",\n \"* Traverse all the `(path, value)` pairs in the flat state and group them according to the predicates.\\n\",\n \"* Use `State.from_flat_state` to convert the flat states to nested [`nnx.State`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/state.html#flax.nnx.State)s.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 6,\n \"id\": \"068208fc\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"params = State({\\n\",\n \" 'a': Param(\\n\",\n \" value=0\\n\",\n \" )\\n\",\n \"})\\n\",\n \"batch_stats = State({\\n\",\n \" 'b': BatchStat(\\n\",\n \" value=True\\n\",\n \" )\\n\",\n \"})\\n\"\n ]\n }\n ],\n \"source\": [\n \"from typing import Any\\n\",\n \"KeyPath = tuple[nnx.graph.Key, ...]\\n\",\n \"\\n\",\n \"def split(node, *filters):\\n\",\n \" graphdef, state = nnx.graph.flatten(node)\\n\",\n \" predicates = [nnx.filterlib.to_predicate(f) for f in filters]\\n\",\n \" flat_states: list[dict[KeyPath, Any]] = [{} for p in predicates]\\n\",\n \"\\n\",\n \" for path, value in state:\\n\",\n \" for i, predicate in enumerate(predicates):\\n\",\n \" if predicate(path, value):\\n\",\n \" flat_states[i][path] = value\\n\",\n \" break\\n\",\n \" else:\\n\",\n \" raise ValueError(f'No filter matched {path = } {value = }')\\n\",\n \"\\n\",\n \" states: tuple[nnx.GraphState, ...] = tuple(\\n\",\n \" nnx.State.from_flat_path(flat_state) for flat_state in flat_states\\n\",\n \" )\\n\",\n \" return graphdef, *states\\n\",\n \"\\n\",\n \"# Let's test it.\\n\",\n \"foo = Foo()\\n\",\n \"\\n\",\n \"graphdef, params, batch_stats = split(foo, nnx.Param, nnx.BatchStat)\\n\",\n \"\\n\",\n \"print(f'{params = }')\\n\",\n \"print(f'{batch_stats = }')\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"7b3aeac8\",\n \"metadata\": {},\n \"source\": [\n \"**Note:*** It's very important to know that **filtering is order-dependent**. The first `Filter` that matches a value will keep it, and therefore you should place more specific `Filter`s before more general `Filter`s.\\n\",\n \"\\n\",\n \"For example, as demonstrated below, if you:\\n\",\n \"\\n\",\n \"1) Create a `SpecialParam` type that is a subclass of [`nnx.Param`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Param), and a `Bar` object (subclassing [`nnx.Module`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/module.html)) that contains both types of parameters; and\\n\",\n \"2) Try to split the [`nnx.Param`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Param)s before the `SpecialParam`s\\n\",\n \"\\n\",\n \"then all the values will be placed in the [`nnx.Param`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Param) group, and the `SpecialParam` group will be empty because all `SpecialParam`s are also [`nnx.Param`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Param)s:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 7,\n \"id\": \"014da4d4\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"params = State({\\n\",\n \" 'a': Param(\\n\",\n \" value=0\\n\",\n \" ),\\n\",\n \" 'b': SpecialParam(\\n\",\n \" value=0\\n\",\n \" )\\n\",\n \"})\\n\",\n \"special_params = State({})\\n\"\n ]\n }\n ],\n \"source\": [\n \"class SpecialParam(nnx.Param):\\n\",\n \" pass\\n\",\n \"\\n\",\n \"class Bar(nnx.Module):\\n\",\n \" def __init__(self):\\n\",\n \" self.a = nnx.Param(0)\\n\",\n \" self.b = SpecialParam(0)\\n\",\n \"\\n\",\n \"bar = Bar()\\n\",\n \"\\n\",\n \"graphdef, params, special_params = split(bar, nnx.Param, SpecialParam) # wrong!\\n\",\n \"print(f'{params = }')\\n\",\n \"print(f'{special_params = }')\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"a9f0b7b8\",\n \"metadata\": {},\n \"source\": [\n \"And reversing the order will ensure that the `SpecialParam` are captured first:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 8,\n \"id\": \"a2ebf5b2\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"params = State({\\n\",\n \" 'a': Param(\\n\",\n \" value=0\\n\",\n \" )\\n\",\n \"})\\n\",\n \"special_params = State({\\n\",\n \" 'b': SpecialParam(\\n\",\n \" value=0\\n\",\n \" )\\n\",\n \"})\\n\"\n ]\n }\n ],\n \"source\": [\n \"graphdef, special_params, params = split(bar, SpecialParam, nnx.Param) # correct!\\n\",\n \"print(f'{params = }')\\n\",\n \"print(f'{special_params = }')\"\n ]\n }\n ],\n \"metadata\": {\n \"jupytext\": {\n \"formats\": \"ipynb,md:myst\"\n },\n \"language_info\": {\n \"codemirror_mode\": {\n \"name\": \"ipython\",\n \"version\": 3\n },\n \"file_extension\": \".py\",\n \"mimetype\": \"text/x-python\",\n \"name\": \"python\",\n \"nbconvert_exporter\": \"python\",\n \"pygments_lexer\": \"ipython3\",\n \"version\": \"3.11.9\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 5\n}\n" }, { "path": "docs_nnx/guides/filters_guide.md", "content": "---\njupytext:\n formats: ipynb,md:myst\n text_representation:\n extension: .md\n format_name: myst\n format_version: 0.13\n jupytext_version: 1.13.8\n---\n\n# Filters\n\nFlax NNX uses [`Filter`s](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/filterlib.html) extensively as a way to create [`nnx.State`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/state.html#flax.nnx.State) groups in APIs, such as [`nnx.split`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/graph.html#flax.nnx.split), [`nnx.state()`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/graph.html#flax.nnx.state), and many of the [Flax NNX transformations (transforms)](https://flax.readthedocs.io/en/latest/guides/jax_and_nnx_transforms.html).\n\nIn this guide you will learn how to:\n\n* Use [`Filter`s](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/filterlib.html) to group Flax NNX variables and states into subgroups;\n* Understand relationships between types, such as [`nnx.Param`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Param) or [`nnx.BatchStat`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.BatchStat), and [`Filter`s](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/filterlib.html);\n* Express your `Filter`s flexibly with [`nnx.filterlib.Filter`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/filterlib.html) language.\n\nIn the following example [`nnx.Param`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Param) and [`nnx.BatchStat`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.BatchStat) are used as `Filter`s to split the model into two groups: one with the parameters and the other with the batch statistics:\n\n```{code-cell} ipython3\nfrom flax import nnx\n\nclass Foo(nnx.Module):\n def __init__(self):\n self.a = nnx.Param(0)\n self.b = nnx.BatchStat(True)\n\nfoo = Foo()\n\ngraphdef, params, batch_stats = nnx.split(foo, nnx.Param, nnx.BatchStat)\n\nprint(f'{params = }')\nprint(f'{batch_stats = }')\n```\n\nLet's dive deeper into `Filter`s.\n\n+++\n\n## The `Filter` Protocol\n\nIn general, Flax `Filter`s are predicate functions of the form:\n\n```python\n\n(path: tuple[Key, ...], value: Any) -> bool\n\n```\n\nwhere:\n\n- `Key` is a hashable and comparable type;\n- `path` is a tuple of `Key`s representing the path to the value in a nested structure; and\n- `value` is the value at the path.\n\nThe function returns `True` if the value should be included in the group, and `False` otherwise.\n\nTypes are not functions of this form. They are treated as `Filter`s because, as you will learn in the next section, types and some other literals are converted to _predicates_. For example, [`nnx.Param`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Param) is roughly converted to a predicate like this:\n\n```{code-cell} ipython3\ndef is_param(path, value) -> bool:\n return isinstance(value, nnx.Param)\n\nprint(f'{is_param((), nnx.Param(0)) = }')\n```\n\nSuch function matches any value that is an instance of [`nnx.Param`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Param). Internally Flax NNX uses `OfType` which defines a callable of this form for a given type:\n\n```{code-cell} ipython3\nis_param = nnx.OfType(nnx.Param)\n\nprint(f'{is_param((), nnx.Param(0)) = }')\n```\n\n## The `Filter` DSL\n\nFlax NNX exposes a small domain specific language ([DSL](https://en.wikipedia.org/wiki/Domain-specific_language)), formalized as the [`nnx.filterlib.Filter`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/filterlib.html) type. This means users don't have to create functions like in the previous section.\n\nHere is a list of all the callable `Filter`s included in Flax NNX, and their corresponding DSL literals (when available):\n\n\n| Literal | Callable | Description |\n|--------|----------------------|-------------|\n| `...` or `True` | `Everything()` | Matches all values |\n| `None` or `False` | `Nothing()` | Matches no values |\n| `type` | `OfType(type)` | Matches values that are instances of `type` or have a `type` attribute that is an instance of `type` |\n| | `PathContains(key)` | Matches values that have an associated `path` that contains the given `key` |\n| `'{filter}'` str | `WithTag('{filter}')` | Matches values that have string `tag` attribute equal to `'{filter}'`. Used by `RngKey` and `RngCount`. |\n| `(*filters)` tuple or `[*filters]` list | `Any(*filters)` | Matches values that match any of the inner `filters` |\n| | `All(*filters)` | Matches values that match all of the inner `filters` |\n| | `Not(filter)` | Matches values that do not match the inner `filter` |\n\n\nLet's check out the DSL in action by using [`nnx.vmap`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.vmap) as an example. Consider the following:\n\n1) You want to vectorize all parameters;\n2) Apply `'dropout'` `Rng(Keys|Counts)` on the `0`th axis; and\n3) Broadcast the rest.\n\nTo do this, you can use the following `Filter`s to define a `nnx.StateAxes` object that you can pass to `nnx.vmap`'s `in_axes` to specify how the `model`'s various sub-states should be vectorized:\n\n```{code-cell} ipython3\nstate_axes = nnx.StateAxes({(nnx.Param, 'dropout'): 0, ...: None})\n\n@nnx.vmap(in_axes=(state_axes, 0))\ndef forward(model, x):\n ...\n```\n\nHere `(nnx.Param, 'dropout')` expands to `Any(OfType(nnx.Param), WithTag('dropout'))` and `...` expands to `Everything()`.\n\nIf you wish to manually convert literal into a predicate, you can use [`nnx.filterlib.to_predicate`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/filterlib.html#flax.nnx.filterlib.to_predicate):\n\n```{code-cell} ipython3\nis_param = nnx.filterlib.to_predicate(nnx.Param)\neverything = nnx.filterlib.to_predicate(...)\nnothing = nnx.filterlib.to_predicate(False)\nparams_or_dropout = nnx.filterlib.to_predicate((nnx.Param, 'dropout'))\n\nprint(f'{is_param = }')\nprint(f'{everything = }')\nprint(f'{nothing = }')\nprint(f'{params_or_dropout = }')\n```\n\n## Grouping `State`s\n\nWith the knowledge of `Filter`s from previous sections at hand, let's learn how to roughly implement [`nnx.split`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/graph.html#flax.nnx.split). Here are the key ideas:\n\n* Use `nnx.graph.flatten` to get the [`GraphDef`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/graph.html#flax.nnx.GraphDef) and [`nnx.State`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/state.html#flax.nnx.State) representation of the node.\n* Convert all the `Filter`s to predicates.\n* Use `State.flat_state` to get the flat representation of the state.\n* Traverse all the `(path, value)` pairs in the flat state and group them according to the predicates.\n* Use `State.from_flat_state` to convert the flat states to nested [`nnx.State`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/state.html#flax.nnx.State)s.\n\n```{code-cell} ipython3\nfrom typing import Any\nKeyPath = tuple[nnx.graph.Key, ...]\n\ndef split(node, *filters):\n graphdef, state = nnx.graph.flatten(node)\n predicates = [nnx.filterlib.to_predicate(f) for f in filters]\n flat_states: list[dict[KeyPath, Any]] = [{} for p in predicates]\n\n for path, value in state:\n for i, predicate in enumerate(predicates):\n if predicate(path, value):\n flat_states[i][path] = value\n break\n else:\n raise ValueError(f'No filter matched {path = } {value = }')\n\n states: tuple[nnx.GraphState, ...] = tuple(\n nnx.State.from_flat_path(flat_state) for flat_state in flat_states\n )\n return graphdef, *states\n\n# Let's test it.\nfoo = Foo()\n\ngraphdef, params, batch_stats = split(foo, nnx.Param, nnx.BatchStat)\n\nprint(f'{params = }')\nprint(f'{batch_stats = }')\n```\n\n**Note:*** It's very important to know that **filtering is order-dependent**. The first `Filter` that matches a value will keep it, and therefore you should place more specific `Filter`s before more general `Filter`s.\n\nFor example, as demonstrated below, if you:\n\n1) Create a `SpecialParam` type that is a subclass of [`nnx.Param`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Param), and a `Bar` object (subclassing [`nnx.Module`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/module.html)) that contains both types of parameters; and\n2) Try to split the [`nnx.Param`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Param)s before the `SpecialParam`s\n\nthen all the values will be placed in the [`nnx.Param`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Param) group, and the `SpecialParam` group will be empty because all `SpecialParam`s are also [`nnx.Param`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Param)s:\n\n```{code-cell} ipython3\nclass SpecialParam(nnx.Param):\n pass\n\nclass Bar(nnx.Module):\n def __init__(self):\n self.a = nnx.Param(0)\n self.b = SpecialParam(0)\n\nbar = Bar()\n\ngraphdef, params, special_params = split(bar, nnx.Param, SpecialParam) # wrong!\nprint(f'{params = }')\nprint(f'{special_params = }')\n```\n\nAnd reversing the order will ensure that the `SpecialParam` are captured first:\n\n```{code-cell} ipython3\ngraphdef, special_params, params = split(bar, SpecialParam, nnx.Param) # correct!\nprint(f'{params = }')\nprint(f'{special_params = }')\n```\n" }, { "path": "docs_nnx/guides/flax_gspmd.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"# Scale up on multiple devices\\n\",\n \"\\n\",\n \"This guide demonstrates how to scale up a Flax NNX model on multiple accelerators (GPUs or Google TPUs) using JAX's parallel programming APIs.\\n\",\n \"\\n\",\n \"[Introduction to Parallel Programming](https://docs.jax.dev/en/latest/sharded-computation.html) is a fantastic guide to learn about the distributed programming essentials of JAX. It describes three parallelism APIs - automatic, explicit and manual - for different levels of control.\\n\",\n \"\\n\",\n \"This guide will primarily cover the automatic scenario, which use the [`jax.jit`](https://jax.readthedocs.io/en/latest/jit-compilation.html) to compile your single-device code as multi-device. You will use [`flax.nnx.spmd`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/spmd.html) APIs to annotate your model variables with how it should be sharded.\\n\",\n \"\\n\",\n \"If you want to follow explicit sharding style, follow [JAX Explicit Sharding](https://docs.jax.dev/en/latest/notebooks/explicit-sharding.html) guide and use JAX's relevant APIs. No API on Flax side is needed.\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Setup\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 1,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"You have 8 “fake” JAX devices now: [CpuDevice(id=0), CpuDevice(id=1), CpuDevice(id=2), CpuDevice(id=3), CpuDevice(id=4), CpuDevice(id=5), CpuDevice(id=6), CpuDevice(id=7)]\\n\"\n ]\n }\n ],\n \"source\": [\n \"from functools import partial\\n\",\n \"\\n\",\n \"import jax\\n\",\n \"from jax import numpy as jnp\\n\",\n \"from jax.sharding import PartitionSpec as P, NamedSharding, AxisType\\n\",\n \"import optax\\n\",\n \"import flax\\n\",\n \"from flax import nnx\\n\",\n \"\\n\",\n \"# Ignore this if you are already running on a TPU or GPU\\n\",\n \"if not jax._src.xla_bridge.backends_are_initialized():\\n\",\n \" jax.config.update('jax_num_cpu_devices', 8)\\n\",\n \"print(f'You have 8 “fake” JAX devices now: {jax.devices()}')\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Set up a `2x4` device mesh as the [JAX data sharding tutorial](https://docs.jax.dev/en/latest/sharded-computation.html#key-concept-data-sharding) instructs.\\n\",\n \"\\n\",\n \"In this guide we use a standard FSDP layout and shard our devices on two axes - `data` and `model`, for doing batch data parallelism and tensor parallelism.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 2,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# Create an auto-mode mesh of two dimensions and annotate each axis with a name.\\n\",\n \"auto_mesh = jax.make_mesh((2, 4), ('data', 'model'))\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"> Compatibility Note: This guide covers the [eager sharding feature](https://flax.readthedocs.io/en/latest/flip/4844-var-eager-sharding.html) that greatly simplifies creating sharded model. If your project already used Flax GSPMD API on version `flax<0.12`, you might have turned the feature off to keep your code working. Users can toggle this feature using the `nnx.use_eager_sharding` function.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 3,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"nnx.use_eager_sharding(True)\\n\",\n \"assert nnx.using_eager_sharding()\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"c24144d8\",\n \"metadata\": {},\n \"source\": [\n \"The `nnx.use_eager_sharding` function can also be used as a context manager to toggle the eager sharding feature within a specific scope.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"2d849e2e\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"with nnx.use_eager_sharding(False):\\n\",\n \" assert not nnx.using_eager_sharding()\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"c9f808ec\",\n \"metadata\": {},\n \"source\": [\n \"You can also enable eager sharding on a per-variable basis by passing `eager_sharding=False` during variable initialization. The mesh can also be passed this way.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"67bbd440\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"nnx.Param(jnp.ones(4,4), out_sharding=(None, 'model'), eager_sharding=True, mesh=auto_mesh)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Shard a single-array model\\n\",\n \"\\n\",\n \"Let's begin by sharding the simplest component possible - a Flax variable.\\n\",\n \"\\n\",\n \"When you define a Flax variable, you can pass in a metadata field called `out_sharding`, to specify how the underlying JAX array should be sharded. This field should be a tuple of names, each of which refer to how an axis of the array should be sharded.\\n\",\n \"\\n\",\n \"**You must have an existing device mesh** and create a sharding-annotated `nnx.Variable` within its scope. This allows the result variable to be sharded accordingly on those devices. The device mesh can be your actual accelerator mesh, or a dummy fake CPU mesh like in this notebook.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 4,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"PartitionSpec(None, 'model')\\n\"\n ]\n },\n {\n \"data\": {\n \"text/html\": [\n \"
                                                \\n\",\n       \"                                                \\n\",\n       \"                                                \\n\",\n       \"                                                \\n\",\n       \"                                                \\n\",\n       \"  CPU 0,4     CPU 1,5     CPU 2,6     CPU 3,7   \\n\",\n       \"                                                \\n\",\n       \"                                                \\n\",\n       \"                                                \\n\",\n       \"                                                \\n\",\n       \"                                                \\n\",\n       \"
\\n\"\n ],\n \"text/plain\": [\n \"\\u001b[38;2;255;255;255;48;2;57;59;121m \\u001b[0m\\u001b[38;2;255;255;255;48;2;222;158;214m \\u001b[0m\\u001b[38;2;255;255;255;48;2;173;73;74m \\u001b[0m\\u001b[38;2;0;0;0;48;2;181;207;107m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;57;59;121m \\u001b[0m\\u001b[38;2;255;255;255;48;2;222;158;214m \\u001b[0m\\u001b[38;2;255;255;255;48;2;173;73;74m \\u001b[0m\\u001b[38;2;0;0;0;48;2;181;207;107m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;57;59;121m \\u001b[0m\\u001b[38;2;255;255;255;48;2;222;158;214m \\u001b[0m\\u001b[38;2;255;255;255;48;2;173;73;74m \\u001b[0m\\u001b[38;2;0;0;0;48;2;181;207;107m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;57;59;121m \\u001b[0m\\u001b[38;2;255;255;255;48;2;222;158;214m \\u001b[0m\\u001b[38;2;255;255;255;48;2;173;73;74m \\u001b[0m\\u001b[38;2;0;0;0;48;2;181;207;107m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;57;59;121m \\u001b[0m\\u001b[38;2;255;255;255;48;2;222;158;214m \\u001b[0m\\u001b[38;2;255;255;255;48;2;173;73;74m \\u001b[0m\\u001b[38;2;0;0;0;48;2;181;207;107m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;57;59;121m \\u001b[0m\\u001b[38;2;255;255;255;48;2;57;59;121mCPU 0,4\\u001b[0m\\u001b[38;2;255;255;255;48;2;57;59;121m \\u001b[0m\\u001b[38;2;255;255;255;48;2;222;158;214m \\u001b[0m\\u001b[38;2;255;255;255;48;2;222;158;214mCPU 1,5\\u001b[0m\\u001b[38;2;255;255;255;48;2;222;158;214m \\u001b[0m\\u001b[38;2;255;255;255;48;2;173;73;74m \\u001b[0m\\u001b[38;2;255;255;255;48;2;173;73;74mCPU 2,6\\u001b[0m\\u001b[38;2;255;255;255;48;2;173;73;74m \\u001b[0m\\u001b[38;2;0;0;0;48;2;181;207;107m \\u001b[0m\\u001b[38;2;0;0;0;48;2;181;207;107mCPU 3,7\\u001b[0m\\u001b[38;2;0;0;0;48;2;181;207;107m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;57;59;121m \\u001b[0m\\u001b[38;2;255;255;255;48;2;222;158;214m \\u001b[0m\\u001b[38;2;255;255;255;48;2;173;73;74m \\u001b[0m\\u001b[38;2;0;0;0;48;2;181;207;107m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;57;59;121m \\u001b[0m\\u001b[38;2;255;255;255;48;2;222;158;214m \\u001b[0m\\u001b[38;2;255;255;255;48;2;173;73;74m \\u001b[0m\\u001b[38;2;0;0;0;48;2;181;207;107m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;57;59;121m \\u001b[0m\\u001b[38;2;255;255;255;48;2;222;158;214m \\u001b[0m\\u001b[38;2;255;255;255;48;2;173;73;74m \\u001b[0m\\u001b[38;2;0;0;0;48;2;181;207;107m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;57;59;121m \\u001b[0m\\u001b[38;2;255;255;255;48;2;222;158;214m \\u001b[0m\\u001b[38;2;255;255;255;48;2;173;73;74m \\u001b[0m\\u001b[38;2;0;0;0;48;2;181;207;107m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;57;59;121m \\u001b[0m\\u001b[38;2;255;255;255;48;2;222;158;214m \\u001b[0m\\u001b[38;2;255;255;255;48;2;173;73;74m \\u001b[0m\\u001b[38;2;0;0;0;48;2;181;207;107m \\u001b[0m\\n\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"rngs = nnx.Rngs(0)\\n\",\n \"\\n\",\n \"with jax.set_mesh(auto_mesh):\\n\",\n \" w = nnx.Param(\\n\",\n \" rngs.lecun_normal()((4, 8)),\\n\",\n \" out_sharding=(None, 'model')\\n\",\n \" )\\n\",\n \" print(w.sharding.spec)\\n\",\n \" jax.debug.visualize_array_sharding(w) # already sharded!\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Initialize with style\\n\",\n \"\\n\",\n \"When using existing modules, you can apply [`flax.nnx.with_partitioning`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/spmd.html#flax.nnx.with_partitioning) on initializers to achieve the same effect. Here we create a sharded `nnx.Linear` module with only the kernel weight.\\n\",\n \"\\n\",\n \"Also, you should use `jax.jit` for the whole initialization for maximum performance. This is because without `jax.jit`, a single-device variable must be created first before we apply sharding constraints and then make it sharded, which is wasteful. `jax.jit` will automatically optimize this out.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"@jax.jit\\n\",\n \"def init_sharded_linear(key):\\n\",\n \" init_fn = nnx.nn.linear.default_kernel_init\\n\",\n \" # Shard your parameter along `model` dimension, as in model/tensor parallelism\\n\",\n \" return nnx.Linear(4, 8, use_bias=False, rngs=nnx.Rngs(key),\\n\",\n \" kernel_init=nnx.with_partitioning(init_fn, (None, 'model')))\\n\",\n \"\\n\",\n \"with jax.set_mesh(auto_mesh):\\n\",\n \" key= rngs()\\n\",\n \" linear = init_sharded_linear(key)\\n\",\n \" assert linear.kernel.sharding.spec == P(None, 'model') # already sharded!\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Run the model\\n\",\n \"\\n\",\n \"If you also shard your input correctly, JAX would be able to carry out the most natural and optimized computation and produce your output as sharded.\\n\",\n \"\\n\",\n \"You should still make sure to `jax.jit` for maximum performance, and also to explicitly control how each array is sharded when you want to. We will give an example of that control in the next section.\\n\",\n \"\\n\",\n \"> Note: You need to `jax.jit` a pure function that takes the model as an argument, instead of jitting the callable model directly.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 6,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"PartitionSpec('data', 'model')\\n\"\n ]\n },\n {\n \"data\": {\n \"text/html\": [\n \"
                                    \\n\",\n       \"                                    \\n\",\n       \"  CPU 0    CPU 1    CPU 2    CPU 3  \\n\",\n       \"                                    \\n\",\n       \"                                    \\n\",\n       \"                                    \\n\",\n       \"                                    \\n\",\n       \"                                    \\n\",\n       \"  CPU 4    CPU 5    CPU 6    CPU 7  \\n\",\n       \"                                    \\n\",\n       \"                                    \\n\",\n       \"                                    \\n\",\n       \"
\\n\"\n ],\n \"text/plain\": [\n \"\\u001b[38;2;255;255;255;48;2;57;59;121m \\u001b[0m\\u001b[38;2;255;255;255;48;2;214;97;107m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;162;82m \\u001b[0m\\u001b[38;2;255;255;255;48;2;222;158;214m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;57;59;121m \\u001b[0m\\u001b[38;2;255;255;255;48;2;214;97;107m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;162;82m \\u001b[0m\\u001b[38;2;255;255;255;48;2;222;158;214m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;57;59;121m \\u001b[0m\\u001b[38;2;255;255;255;48;2;57;59;121mCPU 0\\u001b[0m\\u001b[38;2;255;255;255;48;2;57;59;121m \\u001b[0m\\u001b[38;2;255;255;255;48;2;214;97;107m \\u001b[0m\\u001b[38;2;255;255;255;48;2;214;97;107mCPU 1\\u001b[0m\\u001b[38;2;255;255;255;48;2;214;97;107m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;162;82m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;162;82mCPU 2\\u001b[0m\\u001b[38;2;255;255;255;48;2;140;162;82m \\u001b[0m\\u001b[38;2;255;255;255;48;2;222;158;214m \\u001b[0m\\u001b[38;2;255;255;255;48;2;222;158;214mCPU 3\\u001b[0m\\u001b[38;2;255;255;255;48;2;222;158;214m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;57;59;121m \\u001b[0m\\u001b[38;2;255;255;255;48;2;214;97;107m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;162;82m \\u001b[0m\\u001b[38;2;255;255;255;48;2;222;158;214m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;57;59;121m \\u001b[0m\\u001b[38;2;255;255;255;48;2;214;97;107m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;162;82m \\u001b[0m\\u001b[38;2;255;255;255;48;2;222;158;214m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;57;59;121m \\u001b[0m\\u001b[38;2;255;255;255;48;2;214;97;107m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;162;82m \\u001b[0m\\u001b[38;2;255;255;255;48;2;222;158;214m \\u001b[0m\\n\",\n \"\\u001b[38;2;0;0;0;48;2;231;203;148m \\u001b[0m\\u001b[38;2;255;255;255;48;2;107;110;207m \\u001b[0m\\u001b[38;2;255;255;255;48;2;165;81;148m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;109;49m \\u001b[0m\\n\",\n \"\\u001b[38;2;0;0;0;48;2;231;203;148m \\u001b[0m\\u001b[38;2;255;255;255;48;2;107;110;207m \\u001b[0m\\u001b[38;2;255;255;255;48;2;165;81;148m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;109;49m \\u001b[0m\\n\",\n \"\\u001b[38;2;0;0;0;48;2;231;203;148m \\u001b[0m\\u001b[38;2;0;0;0;48;2;231;203;148mCPU 4\\u001b[0m\\u001b[38;2;0;0;0;48;2;231;203;148m \\u001b[0m\\u001b[38;2;255;255;255;48;2;107;110;207m \\u001b[0m\\u001b[38;2;255;255;255;48;2;107;110;207mCPU 5\\u001b[0m\\u001b[38;2;255;255;255;48;2;107;110;207m \\u001b[0m\\u001b[38;2;255;255;255;48;2;165;81;148m \\u001b[0m\\u001b[38;2;255;255;255;48;2;165;81;148mCPU 6\\u001b[0m\\u001b[38;2;255;255;255;48;2;165;81;148m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;109;49m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;109;49mCPU 7\\u001b[0m\\u001b[38;2;255;255;255;48;2;140;109;49m \\u001b[0m\\n\",\n \"\\u001b[38;2;0;0;0;48;2;231;203;148m \\u001b[0m\\u001b[38;2;255;255;255;48;2;107;110;207m \\u001b[0m\\u001b[38;2;255;255;255;48;2;165;81;148m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;109;49m \\u001b[0m\\n\",\n \"\\u001b[38;2;0;0;0;48;2;231;203;148m \\u001b[0m\\u001b[38;2;255;255;255;48;2;107;110;207m \\u001b[0m\\u001b[38;2;255;255;255;48;2;165;81;148m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;109;49m \\u001b[0m\\n\",\n \"\\u001b[38;2;0;0;0;48;2;231;203;148m \\u001b[0m\\u001b[38;2;255;255;255;48;2;107;110;207m \\u001b[0m\\u001b[38;2;255;255;255;48;2;165;81;148m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;109;49m \\u001b[0m\\n\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"# For simple computations, you can get correctly-sharded output without jitting\\n\",\n \"# In this case, ('data', None) @ (None, 'model') = ('data', 'model')\\n\",\n \"with jax.set_mesh(auto_mesh):\\n\",\n \" # Create your input data, sharded along `data` dimension, as in data parallelism\\n\",\n \" x = jax.device_put(jnp.ones((16, 4)), P('data', None))\\n\",\n \"\\n\",\n \" # Run the model forward function, jitted\\n\",\n \" y = jax.jit(lambda m, x: m(x))(linear, x)\\n\",\n \" print(y.sharding.spec) # sharded: ('data', 'model')\\n\",\n \" jax.debug.visualize_array_sharding(y)\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Shard a wholesome model\\n\",\n \"\\n\",\n \"Now we construct a more wholesome model to show a few advanced tricks. Check out this simple `DotReluDot` module that does two matmuls, and the `MultiDotReluDot` module that creates an arbitrary stack of `DotReluDot` sublayers.\\n\",\n \"\\n\",\n \"Make note of the following:\\n\",\n \"\\n\",\n \"* **Additional axis annotation**: Transforms like `vmap` and `scan` will add additional dimensions to the JAX arrays. Unfortunately, in auto sharding mode you will need to use `nnx.vmap` and `nnx.scan` instead of raw JAX transforms, so that both JAX and Flax knows how to shard this dimension. You won't need this in [explicit sharding mode](#explicit-sharding).\\n\",\n \"\\n\",\n \"* [`jax.lax.with_sharding_constraint`](https://docs.jax.dev/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html#constraining-shardings-of-intermediates-in-jitted-code): They can help you to enforce specific shardings on intermediate activations. Only works under an auto mode mesh context.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 7,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"class DotReluDot(nnx.Module):\\n\",\n \" def __init__(self, depth: int, rngs: nnx.Rngs):\\n\",\n \" init_fn = nnx.initializers.lecun_normal()\\n\",\n \" self.dot1 = nnx.Linear(\\n\",\n \" depth, depth,\\n\",\n \" kernel_init=nnx.with_partitioning(init_fn, (None, 'model')),\\n\",\n \" use_bias=False, # or use `bias_init` to give it annotation too\\n\",\n \" rngs=rngs)\\n\",\n \" self.w2 = nnx.Param(\\n\",\n \" init_fn(rngs.params(), (depth, depth)), # RNG key and shape for W2 creation\\n\",\n \" sharding=('model', None),\\n\",\n \" )\\n\",\n \"\\n\",\n \" def __call__(self, x: jax.Array):\\n\",\n \" y = self.dot1(x)\\n\",\n \" y = jax.nn.relu(y)\\n\",\n \" y = jax.lax.with_sharding_constraint(y, P('data', 'model'))\\n\",\n \" z = jnp.dot(y, self.w2[...])\\n\",\n \" return z\\n\",\n \"\\n\",\n \"class MultiDotReluDot(nnx.Module):\\n\",\n \" def __init__(self, depth: int, num_layers: int, rngs: nnx.Rngs):\\n\",\n \" # Annotate the additional axis with sharding=None, meaning it will be\\n\",\n \" # replicated across all devices.\\n\",\n \" @nnx.vmap(transform_metadata={nnx.PARTITION_NAME: None})\\n\",\n \" def create_sublayers(r):\\n\",\n \" return DotReluDot(depth, r)\\n\",\n \" self.layers = create_sublayers(rngs.fork(split=num_layers))\\n\",\n \"\\n\",\n \" def __call__(self, x):\\n\",\n \" def scan_over_layers(x, layer):\\n\",\n \" return layer(x), None\\n\",\n \" x, _ = jax.lax.scan(scan_over_layers, x, self.layers)\\n\",\n \" return x\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Now a sample training loop, using `jax.jit`.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 8,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"1.251457\\n\",\n \"0.8495563\\n\",\n \"0.6590716\\n\",\n \"0.5399748\\n\",\n \"0.39150265\\n\"\n ]\n }\n ],\n \"source\": [\n \"@jax.jit\\n\",\n \"def train_step(model, optimizer, x, y):\\n\",\n \" def loss_fn(model: DotReluDot):\\n\",\n \" y_pred = model(x)\\n\",\n \" return jnp.mean((y_pred - y) ** 2)\\n\",\n \"\\n\",\n \" loss, grads = jax.value_and_grad(loss_fn)(model)\\n\",\n \" optimizer.update(model, grads)\\n\",\n \" return model, loss\\n\",\n \"\\n\",\n \"\\n\",\n \"with jax.set_mesh(auto_mesh):\\n\",\n \" # Training data\\n\",\n \" input = jax.device_put(rngs.normal((8, 1024)), P('data', None))\\n\",\n \" label = jax.device_put(rngs.normal((8, 1024)), P('data', None))\\n\",\n \" # Model and optimizer\\n\",\n \" model = MultiDotReluDot(1024, 2, rngs=nnx.Rngs(0))\\n\",\n \" optimizer = nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param)\\n\",\n \"\\n\",\n \" # The loop\\n\",\n \" for i in range(5):\\n\",\n \" model, loss = train_step(model, optimizer, input, label)\\n\",\n \" print(loss) # Model (over-)fitting to the labels quickly.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Profiling\\n\",\n \"\\n\",\n \"If you are using a Google TPU pod or a pod slice, you can create a custom `block_all()` utility function, as defined below, to measure the performance:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 9,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"13 ms ± 588 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\\n\"\n ]\n }\n ],\n \"source\": [\n \"%%timeit\\n\",\n \"\\n\",\n \"def block_all(xs):\\n\",\n \" jax.tree_util.tree_map(lambda x: x.block_until_ready(), xs)\\n\",\n \" return xs\\n\",\n \"\\n\",\n \"with jax.set_mesh(auto_mesh):\\n\",\n \" new_state = block_all(train_step(model, optimizer, input, label))\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Load a sharded model from a checkpoint\\n\",\n \"\\n\",\n \"Now you learned how to initialize a sharded model without OOM, but what about saving and loading it from a checkpoint on disk? JAX checkpointing libraries, such as [Orbax](https://orbax.readthedocs.io/en/latest/), support loading a model distributedly if a sharding pytree is provided. Below is an example that uses Orbax's `StandardCheckpointer` API.\\n\",\n \"\\n\",\n \"Make sure you save a model's state, especially if your model shares some variables across modules. Given a You can generate an identical abstract pytree with shardings using Flax’s `nnx.get_abstract_model`.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 10,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"PartitionSpec(None, None, 'model')\\n\",\n \"(2, 1024, 1024)\\n\"\n ]\n }\n ],\n \"source\": [\n \"import orbax.checkpoint as ocp\\n\",\n \"\\n\",\n \"# Save the sharded state.\\n\",\n \"sharded_state = nnx.state(model)\\n\",\n \"path = ocp.test_utils.erase_and_create_empty('/tmp/my-checkpoints/')\\n\",\n \"checkpointer = ocp.StandardCheckpointer()\\n\",\n \"checkpointer.save(path / 'checkpoint_name', sharded_state)\\n\",\n \"\\n\",\n \"# Load a sharded state from the checkpoint.\\n\",\n \"graphdef, abs_state = nnx.get_abstract_model(\\n\",\n \" lambda: MultiDotReluDot(1024, 2, rngs=nnx.Rngs(0)), auto_mesh)\\n\",\n \"restored_state = checkpointer.restore(path / 'checkpoint_name',\\n\",\n \" target=abs_state)\\n\",\n \"restored_model = nnx.merge(graphdef, abs_state)\\n\",\n \"print(restored_model.layers.dot1.kernel.sharding.spec)\\n\",\n \"print(restored_model.layers.dot1.kernel.shape)\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Logical axis annotation\\n\",\n \"\\n\",\n \"JAX's [automatic](https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html) [SPMD](https://jax.readthedocs.io/en/latest/glossary.html#term-SPMD) encourages users to explore different sharding layouts to find the optimal one. To this end, in Flax you have the option to annotate with more descriptive axis names (not just device mesh axis names like `'data'` and `'model'`), as long as you provide a mapping from your alias to the device mesh axes.\\n\",\n \"\\n\",\n \"You can provide the mapping along with the annotation as another metadata of the corresponding [`nnx.Variable`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Variable), or overwrite it at top-level. Check out the `LogicalDotReluDot` example below.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 11,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# The mapping from alias annotation to the device mesh.\\n\",\n \"sharding_rules = (('batch', 'data'), ('hidden', 'model'), ('embed', None))\\n\",\n \"\\n\",\n \"class LogicalDotReluDot(nnx.Module):\\n\",\n \" def __init__(self, depth: int, rngs: nnx.Rngs):\\n\",\n \" init_fn = nnx.initializers.lecun_normal()\\n\",\n \" self.dot1 = nnx.Linear(\\n\",\n \" depth, depth,\\n\",\n \" kernel_init=nnx.with_partitioning(init_fn, ('embed', 'hidden')),\\n\",\n \" use_bias=False, # or use `bias_init` to give it annotation too\\n\",\n \" rngs=rngs)\\n\",\n \" self.w2 = nnx.Param(\\n\",\n \" init_fn(rngs.params(), (depth, depth)), # RNG key and shape for W2 creation\\n\",\n \" sharding=('hidden', 'embed'),\\n\",\n \" )\\n\",\n \"\\n\",\n \" def __call__(self, x: jax.Array):\\n\",\n \" y = self.dot1(x)\\n\",\n \" y = jax.nn.relu(y)\\n\",\n \" # Unfortunately the logical aliasing doesn't work on lower-level JAX calls.\\n\",\n \" y = jax.lax.with_sharding_constraint(y, P('data', None))\\n\",\n \" z = jnp.dot(y, self.w2[...])\\n\",\n \" return z\\n\",\n \"\\n\",\n \"class LogicalMultiDotReluDot(nnx.Module):\\n\",\n \" def __init__(self, depth: int, num_layers: int, rngs: nnx.Rngs):\\n\",\n \" @nnx.vmap(transform_metadata={nnx.PARTITION_NAME: None})\\n\",\n \" def create_sublayers(r):\\n\",\n \" return LogicalDotReluDot(depth, r)\\n\",\n \" self.layers = create_sublayers(rngs.fork(split=num_layers))\\n\",\n \"\\n\",\n \" def __call__(self, x):\\n\",\n \" def scan_over_layers(x, layer):\\n\",\n \" return layer(x), None\\n\",\n \" x, _ = jax.lax.scan(scan_over_layers, x, self.layers)\\n\",\n \" return x\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"If you didn't provide all `sharding_rule` annotations in the model definition, you can apply them at top level by put them into the context via `nnx.logical_axis_rules`.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 12,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"with jax.set_mesh(auto_mesh), nnx.logical_axis_rules(sharding_rules):\\n\",\n \" # Model and optimizer\\n\",\n \" logical_model = LogicalMultiDotReluDot(1024, 2, rngs=nnx.Rngs(0))\\n\",\n \" logical_output = logical_model(input)\\n\",\n \"\\n\",\n \"# Check out their equivalency with some easier-to-read sharding descriptions.\\n\",\n \"assert logical_model.layers.dot1.kernel.sharding.is_equivalent_to(\\n\",\n \" NamedSharding(auto_mesh, P(None, None, 'model')), ndim=3\\n\",\n \")\\n\",\n \"assert logical_model.layers.w2.sharding.is_equivalent_to(\\n\",\n \" NamedSharding(auto_mesh, P(None, 'model', None)), ndim=3\\n\",\n \")\\n\",\n \"assert logical_output.sharding.is_equivalent_to(\\n\",\n \" NamedSharding(auto_mesh, P('data', None)), ndim=2\\n\",\n \")\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### When to use device axis / logical axis\\n\",\n \"\\n\",\n \"Choosing when to use a device or logical axis depends on how much you want to control the partitioning of your model:\\n\",\n \"\\n\",\n \"* **Device mesh axis**:\\n\",\n \"\\n\",\n \" * For a simpler model, this can save you a few extra lines of code of converting the logical naming back to the device naming.\\n\",\n \"\\n\",\n \" * Shardings of intermediate *activation* values can only be done via [`jax.lax.with_sharding_constraint`](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.with_sharding_constraint.html) and device mesh axis. Therefore, if you want super fine-grained control over your model's sharding, directly using device mesh axis names everywhere might be less confusing.\\n\",\n \"\\n\",\n \"* **Logical naming**: This is helpful if you want to experiment around and find the most optimal partition layout for your *model weights*.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Explicit sharding\\n\",\n \"\\n\",\n \"[Explicit sharding](https://docs.jax.dev/en/latest/notebooks/explicit-sharding.html), also called \\\"sharding-in-types\\\", is a new JAX sharding feature that allows every sharding of every array to be deterministic and explicit. Instead of letting XLA compiler figure out the shardings, you as user would explicitly state the shardings via JAX APIs.\\n\",\n \"\\n\",\n \"For education purposes, we provide a simple Flax model example using explicit sharding. Note how you specify shardings for this model:\\n\",\n \"\\n\",\n \"* Parameters: `out_sharding` argument passed into JAX initializers.\\n\",\n \"\\n\",\n \"* Ambigious computations like `jnp.dot`: provide `out_sharding` argument to specify the output sharding.\\n\",\n \"\\n\",\n \"* Additional dimension from transforms: use `jax.vmap`'s argument `spmd_axis_name`, instead of Flax lifted transforms.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"PartitionSpec(None, None, 'model')\\n\",\n \"PartitionSpec(None, 'model', None)\\n\"\n ]\n }\n ],\n \"source\": [\n \"# Explicit axis mesh\\n\",\n \"explicit_mesh = jax.make_mesh((2, 4), ('data', 'model'),\\n\",\n \" axis_types=(AxisType.Explicit, AxisType.Explicit))\\n\",\n \"\\n\",\n \"class ExplicitDotReluDot(nnx.Module):\\n\",\n \" def __init__(self, depth: int, rngs: nnx.Rngs):\\n\",\n \" init_fn = nnx.initializers.lecun_normal()\\n\",\n \" self.dot1 = nnx.Linear(\\n\",\n \" depth, depth,\\n\",\n \" kernel_init=partial(init_fn, out_sharding=P(None, 'model')),\\n\",\n \" use_bias=False,\\n\",\n \" rngs=rngs)\\n\",\n \" self.w2 = nnx.Param(\\n\",\n \" init_fn(rngs.params(), (depth, depth), out_sharding=P('model', None)),\\n\",\n \" )\\n\",\n \" self.b2 = nnx.Param(jnp.zeros((depth, ), out_sharding=P(None,)))\\n\",\n \"\\n\",\n \" def __call__(self, x: jax.Array):\\n\",\n \" y = self.dot1(x)\\n\",\n \" y = jax.nn.relu(y)\\n\",\n \" z = jnp.dot(y, self.w2[...], out_sharding=P('data', None))\\n\",\n \" z = z + self.b2\\n\",\n \" return z\\n\",\n \"\\n\",\n \"\\n\",\n \"class ExplicitMultiDotReluDot(nnx.Module):\\n\",\n \" def __init__(self, depth: int, num_layers: int, rngs: nnx.Rngs):\\n\",\n \" # Annotate the additional axis with sharding=None, meaning it will be\\n\",\n \" # replicated across all devices.\\n\",\n \" @partial(jax.vmap, spmd_axis_name=None)\\n\",\n \" def create_sublayers(r):\\n\",\n \" return ExplicitDotReluDot(depth, r)\\n\",\n \" self.layers = create_sublayers(rngs.fork(split=num_layers))\\n\",\n \"\\n\",\n \" def __call__(self, x):\\n\",\n \" def scan_over_layers(x, layer):\\n\",\n \" return layer(x), None\\n\",\n \" x, _ = jax.lax.scan(scan_over_layers, x, self.layers)\\n\",\n \" return x\\n\",\n \"\\n\",\n \"\\n\",\n \"with jax.set_mesh(explicit_mesh):\\n\",\n \" model = ExplicitMultiDotReluDot(1024, 2, rngs=nnx.Rngs(0))\\n\",\n \" x = jax.device_put(rngs.normal((8, 1024)),\\n\",\n \" NamedSharding(explicit_mesh, P('data', None)))\\n\",\n \" y = model(x)\\n\",\n \"\\n\",\n \"print(model.layers.dot1.kernel.sharding.spec)\\n\",\n \"print(model.layers.w2.sharding.spec)\\n\",\n \"assert x.sharding.is_equivalent_to(y.sharding, ndim=2)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"One thing easier in explicit mode is that you can obtain the abstract array tree with shardings via `jax.eval_shape`, instead of calling `nnx.get_abstract_sharding`. This is not possible in auto mode.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 14,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"PartitionSpec(None, None, 'model')\\n\",\n \"PartitionSpec(None, 'model', None)\\n\"\n ]\n }\n ],\n \"source\": [\n \"# Get the sharding tree to load checkpoint with\\n\",\n \"with jax.set_mesh(explicit_mesh):\\n\",\n \" abs_model = jax.eval_shape(\\n\",\n \" lambda: ExplicitMultiDotReluDot(1024, 2, rngs=nnx.Rngs(0)))\\n\",\n \" print(abs_model.layers.dot1.kernel.sharding.spec)\\n\",\n \" print(abs_model.layers.w2.sharding.spec)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Further readings\\n\",\n \"\\n\",\n \"JAX has abundant documentation on scaled computing.\\n\",\n \"\\n\",\n \"- [Introduction to parallel programming](https://jax.readthedocs.io/en/latest/sharded-computation.html): A 101 level tutorial covering the basics of automatic parallelization with [`jax.jit`](https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html#jax.jit), semi-automatic parallelization with [`jax.jit`](https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html) and [`jax.lax.with_sharding_constraint`](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.with_sharding_constraint.html), and manual sharding with [`shard_map`](https://jax.readthedocs.io/en/latest/_autosummary/jax.experimental.shard_map.shard_map.html#jax.experimental.shard_map.shard_map).\\n\",\n \"- [JAX in multi-process environments](https://jax.readthedocs.io/en/latest/multi_process.html).\\n\",\n \"- [Distributed arrays and automatic parallelization](https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html): A more detailed tutorial about parallelization with [`jax.jit`](https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html#jax.jit) and [`jax.lax.with_sharding_constraint`](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.with_sharding_constraint.html).\"\n ]\n }\n ],\n \"metadata\": {\n \"jupytext\": {\n \"formats\": \"ipynb,md:myst\"\n },\n \"language_info\": {\n \"codemirror_mode\": {\n \"name\": \"ipython\",\n \"version\": 3\n },\n \"file_extension\": \".py\",\n \"mimetype\": \"text/x-python\",\n \"name\": \"python\",\n \"nbconvert_exporter\": \"python\",\n \"pygments_lexer\": \"ipython3\",\n \"version\": \"3.11.9\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 0\n}\n" }, { "path": "docs_nnx/guides/flax_gspmd.md", "content": "---\njupytext:\n formats: ipynb,md:myst\n text_representation:\n extension: .md\n format_name: myst\n format_version: 0.13\n jupytext_version: 1.13.8\n---\n\n# Scale up on multiple devices\n\nThis guide demonstrates how to scale up a Flax NNX model on multiple accelerators (GPUs or Google TPUs) using JAX's parallel programming APIs.\n\n[Introduction to Parallel Programming](https://docs.jax.dev/en/latest/sharded-computation.html) is a fantastic guide to learn about the distributed programming essentials of JAX. It describes three parallelism APIs - automatic, explicit and manual - for different levels of control.\n\nThis guide will primarily cover the automatic scenario, which use the [`jax.jit`](https://jax.readthedocs.io/en/latest/jit-compilation.html) to compile your single-device code as multi-device. You will use [`flax.nnx.spmd`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/spmd.html) APIs to annotate your model variables with how it should be sharded.\n\nIf you want to follow explicit sharding style, follow [JAX Explicit Sharding](https://docs.jax.dev/en/latest/notebooks/explicit-sharding.html) guide and use JAX's relevant APIs. No API on Flax side is needed.\n\n+++\n\n### Setup\n\n```{code-cell} ipython3\nfrom functools import partial\n\nimport jax\nfrom jax import numpy as jnp\nfrom jax.sharding import PartitionSpec as P, NamedSharding, AxisType\nimport optax\nimport flax\nfrom flax import nnx\n\n# Ignore this if you are already running on a TPU or GPU\nif not jax._src.xla_bridge.backends_are_initialized():\n jax.config.update('jax_num_cpu_devices', 8)\nprint(f'You have 8 “fake” JAX devices now: {jax.devices()}')\n```\n\nSet up a `2x4` device mesh as the [JAX data sharding tutorial](https://docs.jax.dev/en/latest/sharded-computation.html#key-concept-data-sharding) instructs.\n\nIn this guide we use a standard FSDP layout and shard our devices on two axes - `data` and `model`, for doing batch data parallelism and tensor parallelism.\n\n```{code-cell} ipython3\n# Create an auto-mode mesh of two dimensions and annotate each axis with a name.\nauto_mesh = jax.make_mesh((2, 4), ('data', 'model'))\n```\n\n> Compatibility Note: This guide covers the [eager sharding feature](https://flax.readthedocs.io/en/latest/flip/4844-var-eager-sharding.html) that greatly simplifies creating sharded model. If your project already used Flax GSPMD API on version `flax<0.12`, you might have turned the feature off to keep your code working. Users can toggle this feature using the `nnx.use_eager_sharding` function.\n\n```{code-cell} ipython3\nnnx.use_eager_sharding(True)\nassert nnx.using_eager_sharding()\n```\n\nThe `nnx.use_eager_sharding` function can also be used as a context manager to toggle the eager sharding feature within a specific scope.\n\n```{code-cell} ipython3\nwith nnx.use_eager_sharding(False):\n assert not nnx.using_eager_sharding()\n```\n\nYou can also enable eager sharding on a per-variable basis by passing `eager_sharding=False` during variable initialization. The mesh can also be passed this way.\n\n```{code-cell} ipython3\nnnx.Param(jnp.ones(4,4), out_sharding=(None, 'model'), eager_sharding=True, mesh=auto_mesh)\n```\n\n## Shard a single-array model\n\nLet's begin by sharding the simplest component possible - a Flax variable.\n\nWhen you define a Flax variable, you can pass in a metadata field called `out_sharding`, to specify how the underlying JAX array should be sharded. This field should be a tuple of names, each of which refer to how an axis of the array should be sharded.\n\n**You must have an existing device mesh** and create a sharding-annotated `nnx.Variable` within its scope. This allows the result variable to be sharded accordingly on those devices. The device mesh can be your actual accelerator mesh, or a dummy fake CPU mesh like in this notebook.\n\n```{code-cell} ipython3\nrngs = nnx.Rngs(0)\n\nwith jax.set_mesh(auto_mesh):\n w = nnx.Param(\n rngs.lecun_normal()((4, 8)),\n out_sharding=(None, 'model')\n )\n print(w.sharding.spec)\n jax.debug.visualize_array_sharding(w) # already sharded!\n```\n\n### Initialize with style\n\nWhen using existing modules, you can apply [`flax.nnx.with_partitioning`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/spmd.html#flax.nnx.with_partitioning) on initializers to achieve the same effect. Here we create a sharded `nnx.Linear` module with only the kernel weight.\n\nAlso, you should use `jax.jit` for the whole initialization for maximum performance. This is because without `jax.jit`, a single-device variable must be created first before we apply sharding constraints and then make it sharded, which is wasteful. `jax.jit` will automatically optimize this out.\n\n```{code-cell} ipython3\n@jax.jit\ndef init_sharded_linear(key):\n init_fn = nnx.nn.linear.default_kernel_init\n # Shard your parameter along `model` dimension, as in model/tensor parallelism\n return nnx.Linear(4, 8, use_bias=False, rngs=nnx.Rngs(key),\n kernel_init=nnx.with_partitioning(init_fn, (None, 'model')))\n\nwith jax.set_mesh(auto_mesh):\n key= rngs()\n linear = init_sharded_linear(key)\n assert linear.kernel.sharding.spec == P(None, 'model') # already sharded!\n```\n\n### Run the model\n\nIf you also shard your input correctly, JAX would be able to carry out the most natural and optimized computation and produce your output as sharded.\n\nYou should still make sure to `jax.jit` for maximum performance, and also to explicitly control how each array is sharded when you want to. We will give an example of that control in the next section.\n\n> Note: You need to `jax.jit` a pure function that takes the model as an argument, instead of jitting the callable model directly.\n\n```{code-cell} ipython3\n# For simple computations, you can get correctly-sharded output without jitting\n# In this case, ('data', None) @ (None, 'model') = ('data', 'model')\nwith jax.set_mesh(auto_mesh):\n # Create your input data, sharded along `data` dimension, as in data parallelism\n x = jax.device_put(jnp.ones((16, 4)), P('data', None))\n\n # Run the model forward function, jitted\n y = jax.jit(lambda m, x: m(x))(linear, x)\n print(y.sharding.spec) # sharded: ('data', 'model')\n jax.debug.visualize_array_sharding(y)\n```\n\n## Shard a wholesome model\n\nNow we construct a more wholesome model to show a few advanced tricks. Check out this simple `DotReluDot` module that does two matmuls, and the `MultiDotReluDot` module that creates an arbitrary stack of `DotReluDot` sublayers.\n\nMake note of the following:\n\n* **Additional axis annotation**: Transforms like `vmap` and `scan` will add additional dimensions to the JAX arrays. Unfortunately, in auto sharding mode you will need to use `nnx.vmap` and `nnx.scan` instead of raw JAX transforms, so that both JAX and Flax knows how to shard this dimension. You won't need this in [explicit sharding mode](#explicit-sharding).\n\n* [`jax.lax.with_sharding_constraint`](https://docs.jax.dev/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html#constraining-shardings-of-intermediates-in-jitted-code): They can help you to enforce specific shardings on intermediate activations. Only works under an auto mode mesh context.\n\n```{code-cell} ipython3\nclass DotReluDot(nnx.Module):\n def __init__(self, depth: int, rngs: nnx.Rngs):\n init_fn = nnx.initializers.lecun_normal()\n self.dot1 = nnx.Linear(\n depth, depth,\n kernel_init=nnx.with_partitioning(init_fn, (None, 'model')),\n use_bias=False, # or use `bias_init` to give it annotation too\n rngs=rngs)\n self.w2 = nnx.Param(\n init_fn(rngs.params(), (depth, depth)), # RNG key and shape for W2 creation\n sharding=('model', None),\n )\n\n def __call__(self, x: jax.Array):\n y = self.dot1(x)\n y = jax.nn.relu(y)\n y = jax.lax.with_sharding_constraint(y, P('data', 'model'))\n z = jnp.dot(y, self.w2[...])\n return z\n\nclass MultiDotReluDot(nnx.Module):\n def __init__(self, depth: int, num_layers: int, rngs: nnx.Rngs):\n # Annotate the additional axis with sharding=None, meaning it will be\n # replicated across all devices.\n @nnx.vmap(transform_metadata={nnx.PARTITION_NAME: None})\n def create_sublayers(r):\n return DotReluDot(depth, r)\n self.layers = create_sublayers(rngs.fork(split=num_layers))\n\n def __call__(self, x):\n def scan_over_layers(x, layer):\n return layer(x), None\n x, _ = jax.lax.scan(scan_over_layers, x, self.layers)\n return x\n```\n\nNow a sample training loop, using `jax.jit`.\n\n```{code-cell} ipython3\n@jax.jit\ndef train_step(model, optimizer, x, y):\n def loss_fn(model: DotReluDot):\n y_pred = model(x)\n return jnp.mean((y_pred - y) ** 2)\n\n loss, grads = jax.value_and_grad(loss_fn)(model)\n optimizer.update(model, grads)\n return model, loss\n\n\nwith jax.set_mesh(auto_mesh):\n # Training data\n input = jax.device_put(rngs.normal((8, 1024)), P('data', None))\n label = jax.device_put(rngs.normal((8, 1024)), P('data', None))\n # Model and optimizer\n model = MultiDotReluDot(1024, 2, rngs=nnx.Rngs(0))\n optimizer = nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param)\n\n # The loop\n for i in range(5):\n model, loss = train_step(model, optimizer, input, label)\n print(loss) # Model (over-)fitting to the labels quickly.\n```\n\n## Profiling\n\nIf you are using a Google TPU pod or a pod slice, you can create a custom `block_all()` utility function, as defined below, to measure the performance:\n\n```{code-cell} ipython3\n%%timeit\n\ndef block_all(xs):\n jax.tree_util.tree_map(lambda x: x.block_until_ready(), xs)\n return xs\n\nwith jax.set_mesh(auto_mesh):\n new_state = block_all(train_step(model, optimizer, input, label))\n```\n\n## Load a sharded model from a checkpoint\n\nNow you learned how to initialize a sharded model without OOM, but what about saving and loading it from a checkpoint on disk? JAX checkpointing libraries, such as [Orbax](https://orbax.readthedocs.io/en/latest/), support loading a model distributedly if a sharding pytree is provided. Below is an example that uses Orbax's `StandardCheckpointer` API.\n\nMake sure you save a model's state, especially if your model shares some variables across modules. Given a You can generate an identical abstract pytree with shardings using Flax’s `nnx.get_abstract_model`.\n\n```{code-cell} ipython3\nimport orbax.checkpoint as ocp\n\n# Save the sharded state.\nsharded_state = nnx.state(model)\npath = ocp.test_utils.erase_and_create_empty('/tmp/my-checkpoints/')\ncheckpointer = ocp.StandardCheckpointer()\ncheckpointer.save(path / 'checkpoint_name', sharded_state)\n\n# Load a sharded state from the checkpoint.\ngraphdef, abs_state = nnx.get_abstract_model(\n lambda: MultiDotReluDot(1024, 2, rngs=nnx.Rngs(0)), auto_mesh)\nrestored_state = checkpointer.restore(path / 'checkpoint_name',\n target=abs_state)\nrestored_model = nnx.merge(graphdef, abs_state)\nprint(restored_model.layers.dot1.kernel.sharding.spec)\nprint(restored_model.layers.dot1.kernel.shape)\n```\n\n## Logical axis annotation\n\nJAX's [automatic](https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html) [SPMD](https://jax.readthedocs.io/en/latest/glossary.html#term-SPMD) encourages users to explore different sharding layouts to find the optimal one. To this end, in Flax you have the option to annotate with more descriptive axis names (not just device mesh axis names like `'data'` and `'model'`), as long as you provide a mapping from your alias to the device mesh axes.\n\nYou can provide the mapping along with the annotation as another metadata of the corresponding [`nnx.Variable`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/variables.html#flax.nnx.Variable), or overwrite it at top-level. Check out the `LogicalDotReluDot` example below.\n\n```{code-cell} ipython3\n# The mapping from alias annotation to the device mesh.\nsharding_rules = (('batch', 'data'), ('hidden', 'model'), ('embed', None))\n\nclass LogicalDotReluDot(nnx.Module):\n def __init__(self, depth: int, rngs: nnx.Rngs):\n init_fn = nnx.initializers.lecun_normal()\n self.dot1 = nnx.Linear(\n depth, depth,\n kernel_init=nnx.with_partitioning(init_fn, ('embed', 'hidden')),\n use_bias=False, # or use `bias_init` to give it annotation too\n rngs=rngs)\n self.w2 = nnx.Param(\n init_fn(rngs.params(), (depth, depth)), # RNG key and shape for W2 creation\n sharding=('hidden', 'embed'),\n )\n\n def __call__(self, x: jax.Array):\n y = self.dot1(x)\n y = jax.nn.relu(y)\n # Unfortunately the logical aliasing doesn't work on lower-level JAX calls.\n y = jax.lax.with_sharding_constraint(y, P('data', None))\n z = jnp.dot(y, self.w2[...])\n return z\n\nclass LogicalMultiDotReluDot(nnx.Module):\n def __init__(self, depth: int, num_layers: int, rngs: nnx.Rngs):\n @nnx.vmap(transform_metadata={nnx.PARTITION_NAME: None})\n def create_sublayers(r):\n return LogicalDotReluDot(depth, r)\n self.layers = create_sublayers(rngs.fork(split=num_layers))\n\n def __call__(self, x):\n def scan_over_layers(x, layer):\n return layer(x), None\n x, _ = jax.lax.scan(scan_over_layers, x, self.layers)\n return x\n```\n\nIf you didn't provide all `sharding_rule` annotations in the model definition, you can apply them at top level by put them into the context via `nnx.logical_axis_rules`.\n\n```{code-cell} ipython3\nwith jax.set_mesh(auto_mesh), nnx.logical_axis_rules(sharding_rules):\n # Model and optimizer\n logical_model = LogicalMultiDotReluDot(1024, 2, rngs=nnx.Rngs(0))\n logical_output = logical_model(input)\n\n# Check out their equivalency with some easier-to-read sharding descriptions.\nassert logical_model.layers.dot1.kernel.sharding.is_equivalent_to(\n NamedSharding(auto_mesh, P(None, None, 'model')), ndim=3\n)\nassert logical_model.layers.w2.sharding.is_equivalent_to(\n NamedSharding(auto_mesh, P(None, 'model', None)), ndim=3\n)\nassert logical_output.sharding.is_equivalent_to(\n NamedSharding(auto_mesh, P('data', None)), ndim=2\n)\n```\n\n### When to use device axis / logical axis\n\nChoosing when to use a device or logical axis depends on how much you want to control the partitioning of your model:\n\n* **Device mesh axis**:\n\n * For a simpler model, this can save you a few extra lines of code of converting the logical naming back to the device naming.\n\n * Shardings of intermediate *activation* values can only be done via [`jax.lax.with_sharding_constraint`](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.with_sharding_constraint.html) and device mesh axis. Therefore, if you want super fine-grained control over your model's sharding, directly using device mesh axis names everywhere might be less confusing.\n\n* **Logical naming**: This is helpful if you want to experiment around and find the most optimal partition layout for your *model weights*.\n\n+++\n\n## Explicit sharding\n\n[Explicit sharding](https://docs.jax.dev/en/latest/notebooks/explicit-sharding.html), also called \"sharding-in-types\", is a new JAX sharding feature that allows every sharding of every array to be deterministic and explicit. Instead of letting XLA compiler figure out the shardings, you as user would explicitly state the shardings via JAX APIs.\n\nFor education purposes, we provide a simple Flax model example using explicit sharding. Note how you specify shardings for this model:\n\n* Parameters: `out_sharding` argument passed into JAX initializers.\n\n* Ambigious computations like `jnp.dot`: provide `out_sharding` argument to specify the output sharding.\n\n* Additional dimension from transforms: use `jax.vmap`'s argument `spmd_axis_name`, instead of Flax lifted transforms.\n\n```{code-cell} ipython3\n# Explicit axis mesh\nexplicit_mesh = jax.make_mesh((2, 4), ('data', 'model'),\n axis_types=(AxisType.Explicit, AxisType.Explicit))\n\nclass ExplicitDotReluDot(nnx.Module):\n def __init__(self, depth: int, rngs: nnx.Rngs):\n init_fn = nnx.initializers.lecun_normal()\n self.dot1 = nnx.Linear(\n depth, depth,\n kernel_init=partial(init_fn, out_sharding=P(None, 'model')),\n use_bias=False,\n rngs=rngs)\n self.w2 = nnx.Param(\n init_fn(rngs.params(), (depth, depth), out_sharding=P('model', None)),\n )\n self.b2 = nnx.Param(jnp.zeros((depth, ), out_sharding=P(None,)))\n\n def __call__(self, x: jax.Array):\n y = self.dot1(x)\n y = jax.nn.relu(y)\n z = jnp.dot(y, self.w2[...], out_sharding=P('data', None))\n z = z + self.b2\n return z\n\n\nclass ExplicitMultiDotReluDot(nnx.Module):\n def __init__(self, depth: int, num_layers: int, rngs: nnx.Rngs):\n # Annotate the additional axis with sharding=None, meaning it will be\n # replicated across all devices.\n @partial(jax.vmap, spmd_axis_name=None)\n def create_sublayers(r):\n return ExplicitDotReluDot(depth, r)\n self.layers = create_sublayers(rngs.fork(split=num_layers))\n\n def __call__(self, x):\n def scan_over_layers(x, layer):\n return layer(x), None\n x, _ = jax.lax.scan(scan_over_layers, x, self.layers)\n return x\n\n\nwith jax.set_mesh(explicit_mesh):\n model = ExplicitMultiDotReluDot(1024, 2, rngs=nnx.Rngs(0))\n x = jax.device_put(rngs.normal((8, 1024)),\n NamedSharding(explicit_mesh, P('data', None)))\n y = model(x)\n\nprint(model.layers.dot1.kernel.sharding.spec)\nprint(model.layers.w2.sharding.spec)\nassert x.sharding.is_equivalent_to(y.sharding, ndim=2)\n```\n\nOne thing easier in explicit mode is that you can obtain the abstract array tree with shardings via `jax.eval_shape`, instead of calling `nnx.get_abstract_sharding`. This is not possible in auto mode.\n\n```{code-cell} ipython3\n# Get the sharding tree to load checkpoint with\nwith jax.set_mesh(explicit_mesh):\n abs_model = jax.eval_shape(\n lambda: ExplicitMultiDotReluDot(1024, 2, rngs=nnx.Rngs(0)))\n print(abs_model.layers.dot1.kernel.sharding.spec)\n print(abs_model.layers.w2.sharding.spec)\n```\n\n## Further readings\n\nJAX has abundant documentation on scaled computing.\n\n- [Introduction to parallel programming](https://jax.readthedocs.io/en/latest/sharded-computation.html): A 101 level tutorial covering the basics of automatic parallelization with [`jax.jit`](https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html#jax.jit), semi-automatic parallelization with [`jax.jit`](https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html) and [`jax.lax.with_sharding_constraint`](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.with_sharding_constraint.html), and manual sharding with [`shard_map`](https://jax.readthedocs.io/en/latest/_autosummary/jax.experimental.shard_map.shard_map.html#jax.experimental.shard_map.shard_map).\n- [JAX in multi-process environments](https://jax.readthedocs.io/en/latest/multi_process.html).\n- [Distributed arrays and automatic parallelization](https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html): A more detailed tutorial about parallelization with [`jax.jit`](https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html#jax.jit) and [`jax.lax.with_sharding_constraint`](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.with_sharding_constraint.html).\n" }, { "path": "docs_nnx/guides/index.rst", "content": "Guides\n======\n\n.. toctree::\n :maxdepth: 2\n\n ../guides_basic\n ../guides_advanced\n" }, { "path": "docs_nnx/guides/jax_and_nnx_transforms.rst", "content": "Flax NNX vs JAX transformations\n===============================\n\nThis guide describes the differences between\n`Flax NNX transformations `__\nand `JAX transformations `__,\nand how to seamlessly switch between them or use them side-by-side. The examples here will focus on\n``nnx.jit``, ``jax.jit``, ``nnx.grad`` and ``jax.grad`` function transformations (transforms).\n\nFirst, let's set up imports and generate some dummy data:\n\n.. testcode:: Flax NNX, JAX\n\n from flax import nnx\n import jax\n\n x = jax.random.normal(jax.random.key(0), (1, 2))\n y = jax.random.normal(jax.random.key(1), (1, 3))\n\nDifferences\n***********\n\nFlax NNX transformations can transform functions that are not pure and make mutations and\nside-effects:\n- Flax NNX transforms enable you to transform functions that take in Flax NNX graph objects as\narguments - such as ``nnx.Module``, ``nnx.Rngs``, ``nnx.Optimizer``, and so on - even those whose state\nwill be mutated.\n- In comparison, these kinds of objects aren't recognized in JAX transformations.\n\nThe Flax NNX `Functional API `_\nprovides a way to convert graph structures to `pytrees `__\nand back. By doing this at every function boundary you can effectively use graph structures with any\nJAX transforms and propagate state updates in a way consistent with functional purity.\n\nFlax NNX custom transforms, such as ``nnx.jit`` and ``nnx.grad``, simply remove the boilerplate, and\nas a result the code looks stateful.\n\nBelow is an example of using the ``nnx.jit`` and ``nnx.grad`` transforms compared to the\nthe code that uses ``jax.jit`` and ``jax.grad`` transforms.\n\nNotice that:\n\n- The function signature of Flax NNX-transformed functions can accept the ``nnx.Linear``\n ``nnx.Module`` instances directly and make stateful updates to the ``Module``.\n- The function signature of JAX-transformed functions can only accept the pytree-registered\n ``nnx.State`` and ``nnx.GraphDef`` objects, and must return an updated copy of them to maintain the\n purity of the transformed function.\n\n.. codediff::\n :title: Flax NNX transforms, JAX transforms\n :groups: Flax NNX, JAX\n :sync:\n\n @nnx.jit\n def train_step(model, x, y):\n def loss_fn(model):\n return ((model(x) - y) ** 2).mean()\n grads = nnx.grad(loss_fn)(model)\n params = nnx.state(model, nnx.Param)\n params = jax.tree_util.tree_map(\n lambda p, g: p - 0.1 * g, params, grads\n )\n nnx.update(model, params)\n\n model = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n train_step(model, x, y)\n\n ---\n @jax.jit #!\n def train_step(graphdef, state, x, y): #!\n def loss_fn(graphdef, state): #!\n model = nnx.merge(graphdef, state) #!\n return ((model(x) - y) ** 2).mean()\n grads = jax.grad(loss_fn, argnums=1)(graphdef, state) #!\n\n model = nnx.merge(graphdef, state) #!\n params = nnx.state(model, nnx.Param)\n params = jax.tree_util.tree_map(\n lambda p, g: p - 0.1 * g, params, grads\n )\n nnx.update(model, params)\n return nnx.split(model) #!\n\n graphdef, state = nnx.split(nnx.Linear(2, 3, rngs=nnx.Rngs(0))) #!\n graphdef, state = train_step(graphdef, state, x, y) #!\n\n\nMixing Flax NNX and JAX transforms\n**********************************\n\nBoth Flax NNX transforms and JAX transforms can be mixed together, so long as the JAX-transformed function\nin your code is pure and has valid argument types that are recognized by JAX.\n\n.. codediff::\n :title: Using ``nnx.jit`` with ``jax.grad``, Using ``jax.jit`` with ``nnx.grad``\n :groups: Flax NNX, JAX\n :sync:\n\n @nnx.jit\n def train_step(model, x, y):\n def loss_fn(graphdef, state): #!\n model = nnx.merge(graphdef, state)\n return ((model(x) - y) ** 2).mean()\n grads = jax.grad(loss_fn, 1)(*nnx.split(model)) #!\n params = nnx.state(model, nnx.Param)\n params = jax.tree_util.tree_map(\n lambda p, g: p - 0.1 * g, params, grads\n )\n nnx.update(model, params)\n\n model = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n train_step(model, x, y)\n\n ---\n @jax.jit #!\n def train_step(graphdef, state, x, y): #!\n model = nnx.merge(graphdef, state)\n def loss_fn(model):\n return ((model(x) - y) ** 2).mean()\n grads = nnx.grad(loss_fn)(model)\n params = nnx.state(model, nnx.Param)\n params = jax.tree_util.tree_map(\n lambda p, g: p - 0.1 * g, params, grads\n )\n nnx.update(model, params)\n return nnx.split(model)\n\n graphdef, state = nnx.split(nnx.Linear(2, 3, rngs=nnx.Rngs(0)))\n graphdef, state = train_step(graphdef, state, x, y)\n" }, { "path": "docs_nnx/guides/performance.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"# Performance considerations\\n\",\n \"\\n\",\n \"Currently, Flax [`nnx.jit`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.jit) traverses the object graph in pure Python which can add overhead. This overhead mostly affects small to medium models and can be mitigated in the following ways:\\n\",\n \"* By leveraging JAX's [Asynchronous dispatch](#asynchronous-dispatch).\\n\",\n \"* By using [nnx.cached_partial](#caching-graph-node-traversals) to cache the graph node traversals.\\n\",\n \"* By using a [Functional training loop](#functional-training-loop) which stages out the graph traversals.\\n\",\n \"\\n\",\n \"A full resolution _might_ involve developing a C extension (e.g. `flaxlib`) to speed up some of the traversal logic in [`graph.py`](https://github.com/google/flax/blob/main/flax/nnx/graph.py). Before we continue lets an example of a model and a simple training loop:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"from flax import nnx\\n\",\n \"import jax\\n\",\n \"import jax.numpy as jnp\\n\",\n \"import optax\\n\",\n \"\\n\",\n \"class Model(nnx.Module):\\n\",\n \" def __init__(self, din, dmid, dout, rngs: nnx.Rngs):\\n\",\n \" self.linear = nnx.Linear(din, dmid, rngs=rngs)\\n\",\n \" self.bn = nnx.BatchNorm(dmid, rngs=rngs)\\n\",\n \" self.dropout = nnx.Dropout(0.2, rngs=rngs)\\n\",\n \" self.linear_out = nnx.Linear(dmid, dout, rngs=rngs)\\n\",\n \"\\n\",\n \" def __call__(self, x):\\n\",\n \" x = nnx.relu(self.dropout(self.bn(self.linear(x))))\\n\",\n \" return self.linear_out(x)\\n\",\n \" \\n\",\n \"model = Model(2, 64, 3, rngs=nnx.Rngs(0)) # eager initialization\\n\",\n \"optimizer = nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param)\\n\",\n \"metrics = nnx.MultiMetric(\\n\",\n \" loss=nnx.metrics.Average('loss'),\\n\",\n \")\\n\",\n \"\\n\",\n \"@nnx.jit # <== currently slow\\n\",\n \"def train_step(model, optimizer, metrics, x, y):\\n\",\n \" def loss_fn(model):\\n\",\n \" y_pred = model(x) # call methods directly\\n\",\n \" return ((y_pred - y) ** 2).mean()\\n\",\n \"\\n\",\n \" loss, grads = nnx.value_and_grad(loss_fn)(model)\\n\",\n \" optimizer.update(model, grads) # in-place updates\\n\",\n \" metrics.update(loss=loss)\\n\",\n \"\\n\",\n \" return loss\\n\",\n \" \\n\",\n \"for _ in range(10):\\n\",\n \" x, y = jnp.ones((32, 2)), jnp.zeros((32, 3))\\n\",\n \" loss = train_step(model, optimizer, metrics, x, y)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Important thing here is that we created a `train_step()` function that uses `nnx.jit` and takes in a `model`, `optimizer`, and `metrics` arguments, all of which are Flax NNX objects. We'll later see how to improve this.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Asynchronous dispatch\\n\",\n \"\\n\",\n \"Asynchronous dispatch is a feature of JAX where it runs operations in the background whenever possible so Python can continue executing other code. This can be use to absorb the cost of data loading and in this case the overhead of `nnx.jit` and similar transforms. In general, as the amount of computation JAX has to perform per iteration increases the more it is able to absorb the python overhead since eventually the JAX computation will be the main blocker and programs with different overhead will have the same performance. This could be achieved in a couple of ways:\\n\",\n \"\\n\",\n \"* Increasing the batch size.\\n\",\n \"* Increasing the model size.\\n\",\n \"* Performing more JAX steps per python step if data loading is fast enough.\\n\",\n \"\\n\",\n \"To demonstrate this, the graph below which shows total time of running [benchmarks/nnx_simple_training.py](https://github.com/google/flax/blob/main/benchmarks/nnx_simple_training.py) for both `jax.jit` and `nnx.jit` with different model sizes:\\n\",\n \"\\n\",\n \"![performance-graph](images/performance-graph.png)\\n\",\n \"\\n\",\n \"As we can observe, after a certain model size both `jax.jit` and `nnx.jit` converge to the same runtime cost. This means we don't have to modify our training loop above.\\n\",\n \"\\n\",\n \"## Caching graph node traversals\\n\",\n \"\\n\",\n \"The simplest way to get rid of the traversal overhead entirely is by using `nnx.cached_partial` to convert a transformed function and the input graph objects into a partial function which caches the graph object and just expects the remaining arguments. In this example we use `nnx.cached_partial` over `train_step` and partially apply `model`, `optimizer`, and `metrics`, to create `cached_train_step`. Then we simply update our training loop to use `cached_train_step` which only expects the `x` and `y` inputs:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"cached_train_step = nnx.cached_partial(train_step, model, optimizer, metrics)\\n\",\n \"\\n\",\n \"for _ in range(10):\\n\",\n \" x, y = jnp.ones((32, 2)), jnp.zeros((32, 3))\\n\",\n \" loss = cached_train_step(x, y)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Note that `cached_partial` will enforce that the structure of the graph nodes doesn't change during `train_step` (no mutations except for `Variable` state update) so the cache is guaranteed to be up-to-date and we can avoid costly checks which require traversals. This is actually what is expected for most step functions as making any change here would imply costly recompilation, so enforcing this might be a secondary feature that could be useful for this purpose.\\n\",\n \"\\n\",\n \"Similarly, to prevent the user from mutating the cached objects outside, `cached_partial` creates a copy of all the graph nodes but, to allow state to be propagated to the original objects, they share references to the same `Variable`s.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Functional training loop\\n\",\n \"\\n\",\n \"To remove the Python overhead we can create a functional training loop that uses regular `jax.jit` in combination with `nnx.split` and `nnx.merge` to stage out the traversal logic. Concretely we can use [`nnx.split`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/graph.html#flax.nnx.split) before the training loop to create a single `graphdef` and `state` pytrees for all the graph nodes. Then we change `train_step()` to accept `graphdef` and `state`, and use [`nnx.merge`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/graph.html#flax.nnx.merge) to recreate the objects inside, and either `nnx.split` or `nnx.state` at the end to get the output `state`. At the end of the training loop or whenever needed we can use [`nnx.update`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/graph.html#flax.nnx.update) to update the objects to the current `state`.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# split before training loop\\n\",\n \"graphdef, state = nnx.split((model, optimizer, metrics))\\n\",\n \"\\n\",\n \"@jax.jit # regular JAX\\n\",\n \"def jax_train_step(graphdef, state, x, y):\\n\",\n \" # merge at the beginning of the function\\n\",\n \" model, optimizer, metrics = nnx.merge(graphdef, state)\\n\",\n \"\\n\",\n \" def loss_fn(model):\\n\",\n \" y_pred = model(x) # call methods directly\\n\",\n \" return ((y_pred - y) ** 2).mean()\\n\",\n \"\\n\",\n \" loss, grads = nnx.value_and_grad(loss_fn)(model)\\n\",\n \" optimizer.update(model, grads)\\n\",\n \" metrics.update(loss=loss)\\n\",\n \"\\n\",\n \" state = nnx.state((model, optimizer, metrics))\\n\",\n \" return loss, state\\n\",\n \"\\n\",\n \"for _ in range(10):\\n\",\n \" x, y = jnp.ones((32, 2)), jnp.zeros((32, 3))\\n\",\n \" loss, state = jax_train_step(graphdef, state, x, y)\\n\",\n \"\\n\",\n \"# update objects after training\\n\",\n \"nnx.update((model, optimizer, metrics), state)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Notice that we only need to do this for `jit`, the use of other Flax transforms like [`nnx.value_and_grad`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.value_and_grad) inside `train_step` doesn't have any performance cost since `jit` will make sure this only traced once.\"\n ]\n }\n ],\n \"metadata\": {\n \"jupytext\": {\n \"formats\": \"ipynb,md:myst\"\n },\n \"language_info\": {\n \"name\": \"python\",\n \"version\": \"3.11.9\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 2\n}\n" }, { "path": "docs_nnx/guides/performance.md", "content": "---\njupytext:\n formats: ipynb,md:myst\n text_representation:\n extension: .md\n format_name: myst\n format_version: 0.13\n jupytext_version: 1.13.8\n---\n\n# Performance considerations\n\nCurrently, Flax [`nnx.jit`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.jit) traverses the object graph in pure Python which can add overhead. This overhead mostly affects small to medium models and can be mitigated in the following ways:\n* By leveraging JAX's [Asynchronous dispatch](#asynchronous-dispatch).\n* By using [nnx.cached_partial](#caching-graph-node-traversals) to cache the graph node traversals.\n* By using a [Functional training loop](#functional-training-loop) which stages out the graph traversals.\n\nA full resolution _might_ involve developing a C extension (e.g. `flaxlib`) to speed up some of the traversal logic in [`graph.py`](https://github.com/google/flax/blob/main/flax/nnx/graph.py). Before we continue lets an example of a model and a simple training loop:\n\n```{code-cell}\nfrom flax import nnx\nimport jax\nimport jax.numpy as jnp\nimport optax\n\nclass Model(nnx.Module):\n def __init__(self, din, dmid, dout, rngs: nnx.Rngs):\n self.linear = nnx.Linear(din, dmid, rngs=rngs)\n self.bn = nnx.BatchNorm(dmid, rngs=rngs)\n self.dropout = nnx.Dropout(0.2, rngs=rngs)\n self.linear_out = nnx.Linear(dmid, dout, rngs=rngs)\n\n def __call__(self, x):\n x = nnx.relu(self.dropout(self.bn(self.linear(x))))\n return self.linear_out(x)\n \nmodel = Model(2, 64, 3, rngs=nnx.Rngs(0)) # eager initialization\noptimizer = nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param)\nmetrics = nnx.MultiMetric(\n loss=nnx.metrics.Average('loss'),\n)\n\n@nnx.jit # <== currently slow\ndef train_step(model, optimizer, metrics, x, y):\n def loss_fn(model):\n y_pred = model(x) # call methods directly\n return ((y_pred - y) ** 2).mean()\n\n loss, grads = nnx.value_and_grad(loss_fn)(model)\n optimizer.update(model, grads) # in-place updates\n metrics.update(loss=loss)\n\n return loss\n \nfor _ in range(10):\n x, y = jnp.ones((32, 2)), jnp.zeros((32, 3))\n loss = train_step(model, optimizer, metrics, x, y)\n```\n\nImportant thing here is that we created a `train_step()` function that uses `nnx.jit` and takes in a `model`, `optimizer`, and `metrics` arguments, all of which are Flax NNX objects. We'll later see how to improve this.\n\n+++\n\n## Asynchronous dispatch\n\nAsynchronous dispatch is a feature of JAX where it runs operations in the background whenever possible so Python can continue executing other code. This can be use to absorb the cost of data loading and in this case the overhead of `nnx.jit` and similar transforms. In general, as the amount of computation JAX has to perform per iteration increases the more it is able to absorb the python overhead since eventually the JAX computation will be the main blocker and programs with different overhead will have the same performance. This could be achieved in a couple of ways:\n\n* Increasing the batch size.\n* Increasing the model size.\n* Performing more JAX steps per python step if data loading is fast enough.\n\nTo demonstrate this, the graph below which shows total time of running [benchmarks/nnx_simple_training.py](https://github.com/google/flax/blob/main/benchmarks/nnx_simple_training.py) for both `jax.jit` and `nnx.jit` with different model sizes:\n\n![performance-graph](images/performance-graph.png)\n\nAs we can observe, after a certain model size both `jax.jit` and `nnx.jit` converge to the same runtime cost. This means we don't have to modify our training loop above.\n\n## Caching graph node traversals\n\nThe simplest way to get rid of the traversal overhead entirely is by using `nnx.cached_partial` to convert a transformed function and the input graph objects into a partial function which caches the graph object and just expects the remaining arguments. In this example we use `nnx.cached_partial` over `train_step` and partially apply `model`, `optimizer`, and `metrics`, to create `cached_train_step`. Then we simply update our training loop to use `cached_train_step` which only expects the `x` and `y` inputs:\n\n```{code-cell}\ncached_train_step = nnx.cached_partial(train_step, model, optimizer, metrics)\n\nfor _ in range(10):\n x, y = jnp.ones((32, 2)), jnp.zeros((32, 3))\n loss = cached_train_step(x, y)\n```\n\nNote that `cached_partial` will enforce that the structure of the graph nodes doesn't change during `train_step` (no mutations except for `Variable` state update) so the cache is guaranteed to be up-to-date and we can avoid costly checks which require traversals. This is actually what is expected for most step functions as making any change here would imply costly recompilation, so enforcing this might be a secondary feature that could be useful for this purpose.\n\nSimilarly, to prevent the user from mutating the cached objects outside, `cached_partial` creates a copy of all the graph nodes but, to allow state to be propagated to the original objects, they share references to the same `Variable`s.\n\n+++\n\n## Functional training loop\n\nTo remove the Python overhead we can create a functional training loop that uses regular `jax.jit` in combination with `nnx.split` and `nnx.merge` to stage out the traversal logic. Concretely we can use [`nnx.split`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/graph.html#flax.nnx.split) before the training loop to create a single `graphdef` and `state` pytrees for all the graph nodes. Then we change `train_step()` to accept `graphdef` and `state`, and use [`nnx.merge`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/graph.html#flax.nnx.merge) to recreate the objects inside, and either `nnx.split` or `nnx.state` at the end to get the output `state`. At the end of the training loop or whenever needed we can use [`nnx.update`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/graph.html#flax.nnx.update) to update the objects to the current `state`.\n\n```{code-cell}\n# split before training loop\ngraphdef, state = nnx.split((model, optimizer, metrics))\n\n@jax.jit # regular JAX\ndef jax_train_step(graphdef, state, x, y):\n # merge at the beginning of the function\n model, optimizer, metrics = nnx.merge(graphdef, state)\n\n def loss_fn(model):\n y_pred = model(x) # call methods directly\n return ((y_pred - y) ** 2).mean()\n\n loss, grads = nnx.value_and_grad(loss_fn)(model)\n optimizer.update(model, grads)\n metrics.update(loss=loss)\n\n state = nnx.state((model, optimizer, metrics))\n return loss, state\n\nfor _ in range(10):\n x, y = jnp.ones((32, 2)), jnp.zeros((32, 3))\n loss, state = jax_train_step(graphdef, state, x, y)\n\n# update objects after training\nnnx.update((model, optimizer, metrics), state)\n```\n\nNotice that we only need to do this for `jit`, the use of other Flax transforms like [`nnx.value_and_grad`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.value_and_grad) inside `train_step` doesn't have any performance cost since `jit` will make sure this only traced once.\n" }, { "path": "docs_nnx/guides/pytree.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"id\": \"90ad74ee\",\n \"metadata\": {},\n \"source\": [\n \"# Module & Pytree\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"5c0691d0\",\n \"metadata\": {},\n \"source\": [\n \"Flax NNX's Modules are by default registered as JAX Pytrees, this allows using them throughout most of JAX APIs but in particular JAX transforms and the `jax.tree.*` functions. Thanks to the pytree protocol a simple NNX program might look like this:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 1,\n \"id\": \"9b2b929d\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"y.shape = (5, 3)\\n\"\n ]\n }\n ],\n \"source\": [\n \"from flax import nnx\\n\",\n \"import jax\\n\",\n \"import jax.numpy as jnp\\n\",\n \"\\n\",\n \"class Linear(nnx.Module):\\n\",\n \" def __init__(self, din, dout, rngs: nnx.Rngs):\\n\",\n \" self.din, self.dout = din, dout\\n\",\n \" self.kernel = nnx.Param(rngs.normal((din, dout)))\\n\",\n \"\\n\",\n \"rngs = nnx.Rngs(0)\\n\",\n \"weights = Linear(2, 3, rngs=rngs)\\n\",\n \"\\n\",\n \"@jax.jit\\n\",\n \"def forward(weights, x):\\n\",\n \" return x @ weights.kernel\\n\",\n \"\\n\",\n \"y = forward(weights, x=rngs.uniform((5, 2)))\\n\",\n \"print(f\\\"{y.shape = }\\\")\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"bcfbbdb2\",\n \"metadata\": {},\n \"source\": [\n \"Here `weights`, of type `Linear`, was able to be passed directly to the `jit`-ed function `forward`. Throughout the rest of this guide we will try to answer the questions:\\n\",\n \"1. What are pytrees? \\n\",\n \"2. How does NNX implement pytrees?\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"4f610bb3\",\n \"metadata\": {},\n \"source\": [\n \"## Pytrees 101\\n\",\n \"Most modern ML models have too many Arrays for users to pass around individually, to deal with this JAX developed a way to track Array data in nested structures that still allowed caching for compilation: Pytrees. JAX pytrees are tree structures made of python objects that can be recursively traversed in order to collect an ordered list of leaves and a definition of the tree structure, this is done via the `jax.tree.flatten` function. Most common pytrees are native python containers like `list`, `dict`, and `tuple`, but interestingly it also include `None`. The example bellow shows how to collect all the integer leaves from a nested structure using `flatten`:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 2,\n \"id\": \"c3529274\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"leaves = [1, 2, 3, 4]\\n\",\n \"treedef = PyTreeDef([{'a': *}, {'b': *, 'c': (*, *), 'd': None}])\\n\"\n ]\n }\n ],\n \"source\": [\n \"pytree = [\\n\",\n \" {'a': 1},\\n\",\n \" {\\n\",\n \" 'b': 2,\\n\",\n \" 'c': (3, 4),\\n\",\n \" 'd': None,\\n\",\n \" }\\n\",\n \"]\\n\",\n \"\\n\",\n \"leaves, treedef = jax.tree.flatten(pytree)\\n\",\n \"print(f\\\"leaves = {leaves}\\\")\\n\",\n \"print(f\\\"treedef = {treedef}\\\")\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"9c037b5e\",\n \"metadata\": {},\n \"source\": [\n \"Note that `None` is not a leaf because its defined as a pytree with no children. The main purpose of being able to flatten, apart from collecting the leaves, is being able reconstruct the pytree structure from the tree definition from any sequence of leaves of the same length via the `jax.tree.unflatten` function:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 3,\n \"id\": \"d8237524\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"old pytree = [{'a': 1}, {'b': 2, 'c': (3, 4), 'd': None}]\\n\",\n \"new pytree = [{'a': 10}, {'b': 20, 'c': (30, 40), 'd': None}]\\n\"\n ]\n }\n ],\n \"source\": [\n \"new_leaves = [x * 10 for x in leaves]\\n\",\n \"new_pytree = jax.tree.unflatten(treedef, new_leaves)\\n\",\n \"\\n\",\n \"print(f\\\"old pytree = {pytree}\\\")\\n\",\n \"print(f\\\"new pytree = {new_pytree}\\\")\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"117c8b2d\",\n \"metadata\": {},\n \"source\": [\n \"### Custom Pytrees\\n\",\n \"JAX allows us to register custom pytree node type by using the `jax.tree_util.register_pytree_node` utility. For any type we are able to define a flatten that decomposes the object into a a sequence of nodes / children and a static (hashable) structure, and a unflatten function which takes the sequence of nodes and the static structure and creates a new instance. In the following example we create a simple type `Foo` with the attributes `a`, `b`, and `c`, and define `a` and `b` as nodes and `c` as static.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 4,\n \"id\": \"0c46905c\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"leaves = [1, 2]\\n\",\n \"treedef = PyTreeDef(CustomNode(Foo[('hi',)], [*, *]))\\n\"\n ]\n }\n ],\n \"source\": [\n \"class Foo:\\n\",\n \" def __init__(self):\\n\",\n \" self.a = 1\\n\",\n \" self.b = 2\\n\",\n \" self.c = \\\"hi\\\"\\n\",\n \"\\n\",\n \"def flatten_foo(foo: Foo):\\n\",\n \" nodes = [foo.a, foo.b] # sequence of nodes\\n\",\n \" static = (foo.c,) # hashable & equatable structure\\n\",\n \" return nodes, static\\n\",\n \"\\n\",\n \"def unflatten_foo(static, nodes):\\n\",\n \" foo = object.__new__(Foo) # create uninitialized instance\\n\",\n \" foo.a = nodes[0]\\n\",\n \" foo.b = nodes[1]\\n\",\n \" foo.c = static[0]\\n\",\n \" return foo\\n\",\n \"\\n\",\n \"jax.tree_util.register_pytree_node(Foo, flatten_foo, unflatten_foo)\\n\",\n \"\\n\",\n \"foo = Foo()\\n\",\n \"leaves, treedef = jax.tree.flatten(foo)\\n\",\n \"print(f\\\"leaves = {leaves}\\\")\\n\",\n \"print(f\\\"treedef = {treedef}\\\")\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"15647b13\",\n \"metadata\": {},\n \"source\": [\n \"Notice that `'hi'` does not appear in the leaves because `c` is defined as static, but you can see it as part of the `PyTreeDef` structure.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"d603fa09\",\n \"metadata\": {},\n \"source\": [\n \"## nnx.Pytree\\n\",\n \"In general it would be cumbersome for users to manually register the pytree definition for every type they create. To automate this process NNX provides the `nnx.Pytree` base type that offers a simple API: users annotate attributes using either `nnx.static` or `nnx.data`, and Pytree will register some flatten and unflatten functions that will take the annotations into account. The `nnx.data` and `nnx.static` annotations must only be assigned to `Pytree` attributes directly.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 5,\n \"id\": \"95016a94\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"pytree structure:\\n\",\n \" - pytree.layers[0].b = Array([0.], dtype=float32)\\n\",\n \" - pytree.layers[0].w = Array([[1.]], dtype=float32)\\n\",\n \" - pytree.layers[1].b = Array([0.], dtype=float32)\\n\",\n \" - pytree.layers[1].w = Array([[1.]], dtype=float32)\\n\"\n ]\n }\n ],\n \"source\": [\n \"class Linear(nnx.Pytree):\\n\",\n \" def __init__(self, din: int, dout: int):\\n\",\n \" self.din = nnx.static(din)\\n\",\n \" self.dout = nnx.static(dout)\\n\",\n \" self.w = nnx.data(jnp.ones((din, dout)))\\n\",\n \" self.b = nnx.data(jnp.zeros((dout,)))\\n\",\n \"\\n\",\n \"class MLP(nnx.Pytree):\\n\",\n \" def __init__(self, num_layers, dim):\\n\",\n \" self.num_layers = nnx.static(num_layers)\\n\",\n \" self.layers = nnx.data([\\n\",\n \" Linear(dim, dim) for _ in range(num_layers)\\n\",\n \" ])\\n\",\n \"\\n\",\n \"pytree = MLP(num_layers=2, dim=1)\\n\",\n \"\\n\",\n \"def pytree_structure(pytree, title='pytree structure'):\\n\",\n \" print(f\\\"{title}:\\\")\\n\",\n \" path_leaves, treedef = jax.tree.flatten_with_path(pytree)\\n\",\n \" for path, value in path_leaves:\\n\",\n \" print(f\\\" - pytree{jax.tree_util.keystr(path)} = {value!r}\\\")\\n\",\n \"\\n\",\n \"pytree_structure(pytree)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"3a2db214\",\n \"metadata\": {},\n \"source\": [\n \"As you can see above, only the `data` paths appear in the leaves. However, its very verbose to have to define `static` and `data` for each attribute, so `Pytree` has sensible defaults. You can remove most of them and it will just work:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 6,\n \"id\": \"a8665146\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"pytree structure:\\n\",\n \" - pytree.layers[0].b = Array([0.], dtype=float32)\\n\",\n \" - pytree.layers[0].w = Array([[1.]], dtype=float32)\\n\",\n \" - pytree.layers[1].b = Array([0.], dtype=float32)\\n\",\n \" - pytree.layers[1].w = Array([[1.]], dtype=float32)\\n\"\n ]\n }\n ],\n \"source\": [\n \"class Linear(nnx.Pytree):\\n\",\n \" def __init__(self, din: int, dout: int):\\n\",\n \" self.din = din # static\\n\",\n \" self.dout = dout # static\\n\",\n \" self.w = jnp.ones((din, dout)) # data\\n\",\n \" self.b = jnp.zeros((dout,)) # data\\n\",\n \"\\n\",\n \"class MLP(nnx.Pytree):\\n\",\n \" def __init__(self, num_layers, dim):\\n\",\n \" self.num_layers = num_layers # static\\n\",\n \" self.layers = nnx.List([ # data\\n\",\n \" Linear(dim, dim) for _ in range(num_layers)\\n\",\n \" ])\\n\",\n \"\\n\",\n \"pytree = MLP(num_layers=2, dim=1)\\n\",\n \"pytree_structure(pytree)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"a249c5f0\",\n \"metadata\": {},\n \"source\": [\n \"The only change we had to do here is use `nnx.List` to signal that `layers` contains `data`, the status of the rest of the attributes can be correctly inferred. The rules that determine if a value is data or not are the following:\\n\",\n \"\\n\",\n \"* `Array`s, `Variable`s, `ArrayRef`s, and `nnx.Pytree`s are data.\\n\",\n \"* Types registered using `nnx.register_data_type` are data.\\n\",\n \"* All other types are static.\\n\",\n \"\\n\",\n \"To check if a value is data use the `nnx.is_data` function which will return its status:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 7,\n \"id\": \"27682936\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"\\n\",\n \"# ------ DATA ------------\\n\",\n \"nnx.is_data( jnp.array(0) ) = True # Arrays are data\\n\",\n \"nnx.is_data( nnx.Param(1) ) = True # Variables are data\\n\",\n \"nnx.is_data( nnx.Rngs(2) ) = True # nnx.Pytrees are data\\n\",\n \"\\n\",\n \"# ------ STATIC ------------\\n\",\n \"nnx.is_data( 'hello' ) = False # strings, arbitrary objects\\n\",\n \"nnx.is_data( 42 ) = False # int, float, bool, complex, etc.\\n\",\n \"nnx.is_data( [1, 2.0, 3j, jnp.array(1)] ) = False # list, dict, tuple, regular pytrees\\n\",\n \"\\n\"\n ]\n }\n ],\n \"source\": [\n \"print(f\\\"\\\"\\\"\\n\",\n \"# ------ DATA ------------\\n\",\n \"{nnx.is_data( jnp.array(0) ) = } # Arrays are data\\n\",\n \"{nnx.is_data( nnx.Param(1) ) = } # Variables are data\\n\",\n \"{nnx.is_data( nnx.Rngs(2) ) = } # nnx.Pytrees are data\\n\",\n \"\\n\",\n \"# ------ STATIC ------------\\n\",\n \"{nnx.is_data( 'hello' ) = } # strings, arbitrary objects\\n\",\n \"{nnx.is_data( 42 ) = } # int, float, bool, complex, etc.\\n\",\n \"{nnx.is_data( [1, 2.0, 3j, jnp.array(1)] ) = } # list, dict, tuple, regular pytrees\\n\",\n \"\\\"\\\"\\\")\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"5ab06e79\",\n \"metadata\": {},\n \"source\": [\n \"### When to use explicit annotations?\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"9dead73d\",\n \"metadata\": {},\n \"source\": [\n \"There are cases were you do want to explicitely annotate the attributes to avoid ambiguity or protect yourself against possible edge cases. These include constraining input arguments which might have unexpected types, forcing attributes as data when their type is not treated as data by default, or using `nnx.static` as a way to assert the attribute should not contain data.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 8,\n \"id\": \"9e064461\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"pytree structure:\\n\",\n \" - pytree.bias.value = Array(0., dtype=float32, weak_type=True)\\n\",\n \" - pytree.ls[0] = Array(0, dtype=int32, weak_type=True)\\n\",\n \" - pytree.ls[1] = Array(1, dtype=int32, weak_type=True)\\n\",\n \" - pytree.ls[2] = Array(2, dtype=int32, weak_type=True)\\n\",\n \" - pytree.x = 1.0\\n\",\n \" - pytree.y = 42\\n\"\n ]\n }\n ],\n \"source\": [\n \"class Bar(nnx.Pytree):\\n\",\n \" def __init__(self, x, use_bias: bool):\\n\",\n \" self.x = nnx.data(x) # constrain inputs (e.g. user could pass Array or float)\\n\",\n \" self.y = nnx.data(42) # force types that are not data by default\\n\",\n \" self.ls = nnx.List([jnp.array(i) for i in range(3)]) # use nnx.List for lists of data\\n\",\n \" self.bias = nnx.data(None) # optional values that can be data\\n\",\n \" if use_bias:\\n\",\n \" self.bias = nnx.Param(jnp.array(0.0))\\n\",\n \"\\n\",\n \"pytree = Bar(1.0, True)\\n\",\n \"pytree_structure(pytree)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"8055b72c\",\n \"metadata\": {},\n \"source\": [\n \"### Dataclasses\\n\",\n \"`nnx.Pytree` dataclasses can be created by using the `nnx.dataclass` decorator. To control the status of each field, `nnx.static` and `nnx.data` can be used as `field` specifiers.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 9,\n \"id\": \"9ca77d65\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"pytree structure:\\n\",\n \" - pytree.ls[0].i = 0\\n\",\n \" - pytree.ls[0].x = Array(0, dtype=int32, weak_type=True)\\n\",\n \" - pytree.ls[1].i = 1\\n\",\n \" - pytree.ls[1].x = Array(42, dtype=int32, weak_type=True)\\n\"\n ]\n }\n ],\n \"source\": [\n \"import dataclasses\\n\",\n \"\\n\",\n \"@nnx.dataclass\\n\",\n \"class Foo(nnx.Pytree):\\n\",\n \" i: int = nnx.data()\\n\",\n \" x: jax.Array\\n\",\n \" a: int\\n\",\n \" s: str = nnx.static(default='hi', kw_only=True)\\n\",\n \"\\n\",\n \"@nnx.dataclass\\n\",\n \"class Bar(nnx.Pytree):\\n\",\n \" ls: list[Foo] = nnx.data()\\n\",\n \" shapes: list[int]\\n\",\n \"\\n\",\n \"pytree = Bar(\\n\",\n \" ls=[Foo(i, jnp.array(42 * i), hash(i)) for i in range(2)],\\n\",\n \" shapes=[8, 16, 32]\\n\",\n \")\\n\",\n \"pytree_structure(pytree)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"fca51f65\",\n \"metadata\": {},\n \"source\": [\n \"`dataclasses.dataclass` can also be used directly, however type checkers will not handle `nnx.static` and `nnx.data` correctly. To solve this `dataclasses.field` can be used by setting `metadata` with the appropriate entry for `static`.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 10,\n \"id\": \"ff54e732\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"dataclass pytree structure:\\n\",\n \" - pytree.a = 10\\n\"\n ]\n }\n ],\n \"source\": [\n \"@dataclasses.dataclass\\n\",\n \"class Bar(nnx.Pytree):\\n\",\n \" a: int = dataclasses.field(metadata={'static': False}) # data\\n\",\n \" b: str = dataclasses.field(metadata={'static': True}) # static\\n\",\n \"\\n\",\n \"pytree = Bar(a=10, b=\\\"hello\\\")\\n\",\n \"pytree_structure(pytree, title='dataclass pytree structure')\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"d6036a0e\",\n \"metadata\": {},\n \"source\": [\n \"### Attribute Updates\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"4e0a1f63\",\n \"metadata\": {},\n \"source\": [\n \"The status of an attribute is defined during its first assignment and will not change upon reassignment. However, it is possible to override the status by explicitly using `nnx.data` or `nnx.static` on reassignment.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 11,\n \"id\": \"509a517e\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"original:\\n\",\n \" - pytree.a = Array(1., dtype=float32, weak_type=True)\\n\",\n \" - pytree.c = 3.14\\n\",\n \"updated:\\n\",\n \" - pytree.a = '🤔'\\n\",\n \" - pytree.b = 42\\n\"\n ]\n }\n ],\n \"source\": [\n \"class Foo(nnx.Pytree):\\n\",\n \" def __init__(self):\\n\",\n \" self.a = jnp.array(1.0) # data\\n\",\n \" self.b = \\\"Hello, world!\\\" # static\\n\",\n \" self.c = nnx.data(3.14) # data\\n\",\n \"\\n\",\n \"pytree = Foo()\\n\",\n \"pytree_structure(pytree, \\\"original\\\")\\n\",\n \"\\n\",\n \"pytree.a = \\\"🤔\\\" # data status doesn't change\\n\",\n \"pytree.b = nnx.data(42) # explicit annotation overrides status to data\\n\",\n \"pytree.c = nnx.static(0.5) # explicit annotation overrides status to static\\n\",\n \"pytree_structure(pytree, \\\"updated\\\")\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"17fd41f5\",\n \"metadata\": {},\n \"source\": [\n \"### Attribute checks\\n\",\n \"`Pytree` has a variety of checks to prevent a common class of errors in JAX. This includes checking for Arrays being assigned to new `static` attributes:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 12,\n \"id\": \"98e04ff9\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"ValueError: Found Arrays in value of type '' annotated with nnx.static(...) when setting attribute 'name' of Pytree type ''.\\n\"\n ]\n }\n ],\n \"source\": [\n \"class Foo(nnx.Pytree):\\n\",\n \" def __init__(self, name):\\n\",\n \" self.name = nnx.static(name)\\n\",\n \"\\n\",\n \"try:\\n\",\n \" foo = Foo(name=jnp.array(123))\\n\",\n \"except ValueError as e:\\n\",\n \" print(\\\"ValueError:\\\", e)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"070a07f0\",\n \"metadata\": {},\n \"source\": [\n \"Checking for Arrays being assigned to known `static` attributes:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 13,\n \"id\": \"c864d5b1\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"ValueError: Cannot assign data value of type '' to static attribute 'name' of Pytree type ''. To override the status explicitly wrap the value with nnx.data on assignment:\\n\",\n \"\\n\",\n \" _.name = nnx.data(...)\\n\",\n \"\\n\",\n \"\\n\"\n ]\n }\n ],\n \"source\": [\n \"try:\\n\",\n \" foo = Foo(name=\\\"mattjj\\\")\\n\",\n \" foo.name = jnp.array(123)\\n\",\n \"except ValueError as e:\\n\",\n \" print(\\\"ValueError:\\\", e)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"c3b17f07\",\n \"metadata\": {},\n \"source\": [\n \"Checking for Arrays after `__init__` on `static` attributes that could've been inserted via mutation. This check can be manually trigger via `nnx.check_pytree` at any time.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 14,\n \"id\": \"628698dd\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"ValueError: Found unexpected Arrays on value of type in static attribute 'ls' of Pytree type ''. This is an error starting from Flax version 0.12.0.\\n\",\n \"Consider one of the following options:\\n\",\n \"\\n\",\n \"1. If the attribute is meant to be static, either remove the Array value or wrap it in a static container.\\n\",\n \"2. Wrap the value with nnx.data on assignment:\\n\",\n \"\\n\",\n \" _.ls = nnx.data(...)\\n\",\n \"\\n\",\n \"3. Annotate the class attribute with nnx.Data:\\n\",\n \"\\n\",\n \" class Foo(Pytree):\\n\",\n \" ls: nnx.Data[list]\\n\",\n \"\\n\",\n \"4. If the container is a list or dict, try using nnx.List(...) or nnx.Dict(...) instead.\\n\",\n \"5. Disable pytree for this class:\\n\",\n \"\\n\",\n \" class Foo(Pytree, pytree=False):\\n\",\n \"\\n\",\n \"\\n\"\n ]\n }\n ],\n \"source\": [\n \"class Foo(nnx.Pytree):\\n\",\n \" def __init__(self):\\n\",\n \" self.ls = [] # treated as static\\n\",\n \" for i in range(5):\\n\",\n \" self.ls.append(jnp.array(i)) # append arrays into static attribute\\n\",\n \"\\n\",\n \"try:\\n\",\n \" foo = Foo() # error: Array found in static attribute after `__init__`\\n\",\n \"except ValueError as e:\\n\",\n \" print(\\\"ValueError:\\\", e)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"37ee2429\",\n \"metadata\": {},\n \"source\": [\n \"Checking for `nnx.data` or `nnx.static` annotations stored inside nested structures that are not `nnx.Pytree` instances:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 15,\n \"id\": \"f9d69634\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"ValueError: Found unexpected tags {'static', 'data'} on attribute 'Foo.a'. Values from nnx.data(...) and\\n\",\n \"nnx.static(...) should be assigned to nnx.Pytree attributes directly, they should not be inside other structures. Got value of type '' on Pytree of type ''.\\n\"\n ]\n }\n ],\n \"source\": [\n \"class Foo(nnx.Pytree):\\n\",\n \" def __init__(self):\\n\",\n \" self.a = [nnx.data(1), nnx.static(2)] # annotations in sub-pytree\\n\",\n \"\\n\",\n \"try:\\n\",\n \" foo = Foo()\\n\",\n \"except ValueError as e:\\n\",\n \" print(\\\"ValueError:\\\", e)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"452915ac\",\n \"metadata\": {},\n \"source\": [\n \"### Trace-level awareness\\n\",\n \"To prevent tracer leakage NNX will raise an error when trying to update the attribute of a `Pytree` or the value of a `Variable` on instances that are passed as captures to functions called by JAX transforms:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 16,\n \"id\": \"668db479\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"Error: Cannot mutate 'Foo' from different trace level (https://flax.readthedocs.io/en/latest/api_reference/flax.errors.html#flax.errors.TraceContextError)\\n\"\n ]\n }\n ],\n \"source\": [\n \"class Foo(nnx.Pytree):\\n\",\n \" def __init__(self):\\n\",\n \" self.count = nnx.data(0)\\n\",\n \"\\n\",\n \"foo = Foo()\\n\",\n \"\\n\",\n \"@jax.vmap # or jit, grad, shard_map, pmap, scan, etc.\\n\",\n \"def increment(n):\\n\",\n \" # foo passed as capture\\n\",\n \" foo.count += 1 # error!\\n\",\n \"\\n\",\n \"try:\\n\",\n \" increment(jnp.arange(5))\\n\",\n \"except Exception as e:\\n\",\n \" print(f\\\"Error: {e}\\\")\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"536d77b0\",\n \"metadata\": {},\n \"source\": [\n \"### Reference Sharing\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"2449a3c5\",\n \"metadata\": {},\n \"source\": [\n \"As the name implies Pytrees should be trees. To check if a structure is a well-defined tree you can use the `nnx.find_duplicates` functions which will return a list of duplicates, where each duplicate is a list of path tuples. In the example below we see that `left` and `right` are shared references therefore `find_duplicates` returns a non-empty list with the paths:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 17,\n \"id\": \"32c46ce8\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"nnx.find_duplicates(m) = [[('left',), ('right',)]] # not a tree\\n\"\n ]\n }\n ],\n \"source\": [\n \"class Shared(nnx.Pytree):\\n\",\n \" def __init__(self):\\n\",\n \" self.x = jnp.array(1.0)\\n\",\n \"\\n\",\n \"class Parent(nnx.Pytree):\\n\",\n \" def __init__(self):\\n\",\n \" self.left = Shared()\\n\",\n \" self.right = self.left # reference sharing\\n\",\n \"\\n\",\n \"m = Parent()\\n\",\n \"\\n\",\n \"print(f\\\"{nnx.find_duplicates(m) = } # not a tree\\\")\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"49b2e5f0\",\n \"metadata\": {},\n \"source\": [\n \"The main issue is that sharing is not preserved across pytree operations including JAX transforms, and this results in unintended state duplication:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 18,\n \"id\": \"c33e4862\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"Before: m.left is m.right = True\\n\",\n \"Inside: m.left is m.right = False\\n\",\n \"After: m.left is m.right = False\\n\"\n ]\n }\n ],\n \"source\": [\n \"m = Parent()\\n\",\n \"print(f\\\"Before: {m.left is m.right = }\\\")\\n\",\n \"\\n\",\n \"@jax.jit\\n\",\n \"def f(m):\\n\",\n \" print(f\\\"Inside: {m.left is m.right = }\\\")\\n\",\n \" return m\\n\",\n \"\\n\",\n \"m = f(m)\\n\",\n \"print(f\\\"After: {m.left is m.right = }\\\")\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"d60953f4\",\n \"metadata\": {},\n \"source\": [\n \"Reference sharing is rare in most Machine Learning applications, however if it is required you can either use the `nnx.{split, merge, state, update}` APIs to move the deduplicated state and graph definiton across the JAX transforms:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 19,\n \"id\": \"dda51b67\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"Before: m.left is m.right = True\\n\",\n \"Inside: m.left is m.right = True\\n\",\n \"After: m.left is m.right = True\\n\",\n \"state = State({\\n\",\n \" 'left': {\\n\",\n \" 'x': Array(1., dtype=float32, weak_type=True)\\n\",\n \" }\\n\",\n \"})\\n\"\n ]\n }\n ],\n \"source\": [\n \"m = Parent()\\n\",\n \"print(f\\\"Before: {m.left is m.right = }\\\")\\n\",\n \"graphdef, state = nnx.split(m)\\n\",\n \"\\n\",\n \"@jax.jit\\n\",\n \"def f(graphdef, state):\\n\",\n \" m = nnx.merge(graphdef, state)\\n\",\n \" print(f\\\"Inside: {m.left is m.right = }\\\")\\n\",\n \" return nnx.state(m)\\n\",\n \"\\n\",\n \"state = f(graphdef, state)\\n\",\n \"nnx.update(m, state)\\n\",\n \"\\n\",\n \"print(f\\\"After: {m.left is m.right = }\\\")\\n\",\n \"print(f\\\"{state = }\\\") # deduplicated state\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"0afee781\",\n \"metadata\": {},\n \"source\": [\n \"Or alternatively you can use the NNX transforms which preserve shared references:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 20,\n \"id\": \"caa01e3b\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"Before: m.left is m.right = True\\n\",\n \"Inside: m.left is m.right = True\\n\",\n \"After: m.left is m.right = True\\n\"\n ]\n }\n ],\n \"source\": [\n \"m = Parent()\\n\",\n \"print(f\\\"Before: {m.left is m.right = }\\\")\\n\",\n \"\\n\",\n \"@nnx.jit\\n\",\n \"def f(m):\\n\",\n \" print(f\\\"Inside: {m.left is m.right = }\\\")\\n\",\n \" return m\\n\",\n \"\\n\",\n \"m = f(m)\\n\",\n \"\\n\",\n \"print(f\\\"After: {m.left is m.right = }\\\")\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"734aa5b3\",\n \"metadata\": {},\n \"source\": [\n \"### Turning off pytree registration\\n\",\n \"`nnx.Pytree` allows you to turn off the pytree registration along with the attribute checks for subtypes by setting `pytree` type attribute option to `False`. This can be useful when upgrading to previous NNX code to newer Flax verions as you will still be able to use the NNX APIs or when creating types that should not be treated as pytree because e.g. they shared references.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 21,\n \"id\": \"d2e03753\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \" nnx.state(foo) = State({\\n\",\n \" 'a': {\\n\",\n \" 0: Array(2, dtype=int32, weak_type=True),\\n\",\n \" 1: Array(4, dtype=int32, weak_type=True)\\n\",\n \" },\\n\",\n \" 'b': Array(6, dtype=int32, weak_type=True)\\n\",\n \"})\\n\",\n \" jax.tree_util.all_leaves([foo]) = True\\n\"\n ]\n }\n ],\n \"source\": [\n \"class Foo(nnx.Pytree, pytree=False):\\n\",\n \" def __init__(self):\\n\",\n \" self.a = [jnp.array(1), jnp.array(2)] # no checks\\n\",\n \" self.b = \\\"hello\\\" \\n\",\n \" self.b = jnp.array(3) # no checks\\n\",\n \"\\n\",\n \"foo = Foo()\\n\",\n \"\\n\",\n \"@nnx.jit # can use in NNX transformations\\n\",\n \"def double(foo: Foo):\\n\",\n \" foo.a = [x * 2 for x in foo.a]\\n\",\n \" foo.b *= 2\\n\",\n \"\\n\",\n \"double(foo)\\n\",\n \"print(f\\\"{ nnx.state(foo) = }\\\") # can be used with NNX APIs\\n\",\n \"print(f\\\"{ jax.tree_util.all_leaves([foo]) = }\\\") # not a pytree\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"bfaf17de\",\n \"metadata\": {},\n \"source\": [\n \"## Module\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"67000b88\",\n \"metadata\": {},\n \"source\": [\n \"NNX Modules are `Pytree`s that have two additional methods for traking intermediate values: `sow` and `perturb`.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"cc5afc70\",\n \"metadata\": {},\n \"source\": [\n \"### sow\\n\",\n \"`sow` receives a `Variable` type, a `name`, and a `value`, and stores it in the `Module` so it can be retrieved at a later time. As the following example shows, NNX APIs such as `nnx.state` or `nnx.pop` are a good way of retrieving the sowed state, however `pop` is recommended because it explicitly removes the temporary state from the Module.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 22,\n \"id\": \"ca9f58a2\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"\\u001b[38;2;79;201;177mState\\u001b[0m\\u001b[38;2;255;213;3m({\\u001b[0m\\u001b[38;2;105;105;105m\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254m'blocks'\\u001b[0m\\u001b[38;2;212;212;212m: \\u001b[0m\\u001b[38;2;255;213;3m{\\u001b[0m\\u001b[38;2;105;105;105m\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254m0\\u001b[0m\\u001b[38;2;212;212;212m: \\u001b[0m\\u001b[38;2;255;213;3m{\\u001b[0m\\u001b[38;2;105;105;105m\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254m'y_mean'\\u001b[0m\\u001b[38;2;212;212;212m: \\u001b[0m\\u001b[38;2;79;201;177mIntermediate\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;105;105;105m # 1 (4 B)\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254mvalue\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0mArray(4.659754e-06, dtype=float32),\\u001b[38;2;255;213;3m)\\u001b[0m\\n\",\n \" \\u001b[38;2;255;213;3m)\\u001b[0m\\n\",\n \" \\u001b[38;2;255;213;3m}\\u001b[0m,\\n\",\n \" \\u001b[38;2;156;220;254m1\\u001b[0m\\u001b[38;2;212;212;212m: \\u001b[0m\\u001b[38;2;255;213;3m{\\u001b[0m\\u001b[38;2;105;105;105m\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254m'y_mean'\\u001b[0m\\u001b[38;2;212;212;212m: \\u001b[0m\\u001b[38;2;79;201;177mIntermediate\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;105;105;105m # 1 (4 B)\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254mvalue\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0mArray(0.00025933, dtype=float32),\\u001b[38;2;255;213;3m)\\u001b[0m\\n\",\n \" \\u001b[38;2;255;213;3m)\\u001b[0m\\n\",\n \" \\u001b[38;2;255;213;3m}\\u001b[0m,\\n\",\n \" \\u001b[38;2;156;220;254m2\\u001b[0m\\u001b[38;2;212;212;212m: \\u001b[0m\\u001b[38;2;255;213;3m{\\u001b[0m\\u001b[38;2;105;105;105m\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254m'y_mean'\\u001b[0m\\u001b[38;2;212;212;212m: \\u001b[0m\\u001b[38;2;79;201;177mIntermediate\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;105;105;105m # 1 (4 B)\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254mvalue\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0mArray(0.05561922, dtype=float32),\\u001b[38;2;255;213;3m)\\u001b[0m\\n\",\n \" \\u001b[38;2;255;213;3m)\\u001b[0m\\n\",\n \" \\u001b[38;2;255;213;3m}\\u001b[0m\\n\",\n \" \\u001b[38;2;255;213;3m}\\u001b[0m\\n\",\n \"\\u001b[38;2;255;213;3m})\\u001b[0m\\n\"\n ]\n }\n ],\n \"source\": [\n \"class Block(nnx.Module):\\n\",\n \" def __init__(self, din: int, dout: int, rngs: nnx.Rngs):\\n\",\n \" self.linear = nnx.Linear(din, dout, rngs=rngs)\\n\",\n \" self.bn = nnx.BatchNorm(dout, rngs=rngs)\\n\",\n \" self.dropout = nnx.Dropout(0.1, rngs=rngs)\\n\",\n \"\\n\",\n \" def __call__(self, x):\\n\",\n \" y = nnx.relu(self.dropout(self.bn(self.linear(x))))\\n\",\n \" self.sow(nnx.Intermediate, \\\"y_mean\\\", jnp.mean(y))\\n\",\n \" return y\\n\",\n \"\\n\",\n \"class MLP(nnx.Module):\\n\",\n \" def __init__(self, num_layers, dim, rngs: nnx.Rngs):\\n\",\n \" self.blocks = nnx.List([Block(dim, dim, rngs) for _ in range(num_layers)])\\n\",\n \"\\n\",\n \" def __call__(self, x):\\n\",\n \" for block in self.blocks:\\n\",\n \" x = block(x)\\n\",\n \" return x\\n\",\n \"\\n\",\n \"\\n\",\n \"model = MLP(num_layers=3, dim=20, rngs=nnx.Rngs(0))\\n\",\n \"x = jnp.ones((10, 20))\\n\",\n \"y = model(x)\\n\",\n \"intermediates = nnx.pop(model, nnx.Intermediate) # extract intermediate values\\n\",\n \"print(intermediates)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"6e266e5f\",\n \"metadata\": {},\n \"source\": [\n \"### perturb\\n\",\n \"`perturb` is similar to `sow` but it aims to capture the gradient of a value, currently this is a two step process although it might be simplified in the future:\\n\",\n \"1. Initialize the pertubation state by running the model once.\\n\",\n \"2. Pass the perturbation state as a differentiable target to `grad`.\\n\",\n \"\\n\",\n \"As an example lets create a simple model and use `perturb` to get the intermediate gradient `xgrad` for the variable `x`, and initialize the perturbations:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 23,\n \"id\": \"41398e14\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"nnx.state(model, nnx.Perturbation) = \\u001b[38;2;79;201;177mState\\u001b[0m\\u001b[38;2;255;213;3m({\\u001b[0m\\u001b[38;2;105;105;105m\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254m'xgrad'\\u001b[0m\\u001b[38;2;212;212;212m: \\u001b[0m\\u001b[38;2;79;201;177mPerturbation\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;105;105;105m # 3 (12 B)\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254mvalue\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0mArray([[0., 0., 0.]], dtype=float32)\\n\",\n \" \\u001b[38;2;255;213;3m)\\u001b[0m\\n\",\n \"\\u001b[38;2;255;213;3m})\\u001b[0m\\n\"\n ]\n }\n ],\n \"source\": [\n \"import optax\\n\",\n \"\\n\",\n \"class Model(nnx.Module):\\n\",\n \" def __init__(self, rngs):\\n\",\n \" self.linear1 = nnx.Linear(2, 3, rngs=rngs)\\n\",\n \" self.linear2 = nnx.Linear(3, 4, rngs=rngs)\\n\",\n \" def __call__(self, x):\\n\",\n \" x = nnx.gelu(self.linear1(x))\\n\",\n \" x = self.perturb('xgrad', x)\\n\",\n \" x = self.linear2(x)\\n\",\n \" return x\\n\",\n \"\\n\",\n \"rngs = nnx.Rngs(0)\\n\",\n \"model = Model(rngs)\\n\",\n \"optimizer = nnx.Optimizer(model, tx=optax.sgd(1e-1), wrt=nnx.Param)\\n\",\n \"x, y = rngs.uniform((1, 2)), rngs.uniform((1, 4))\\n\",\n \"_ = model(x) # initialize perturbations\\n\",\n \"print(f\\\"{nnx.state(model, nnx.Perturbation) = !s}\\\")\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"c9221005\",\n \"metadata\": {},\n \"source\": [\n \"Next we'll create a training step function that differentiates w.r.t. both the parameters of the model and the perturbations, the later will be the gradients for the intermediate values. `nnx.jit` and `nnx.value_and_grad` will be use to automatically propagate state updates. We'll return the `loss` function and the itermediate gradients.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 24,\n \"id\": \"d10effba\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"step = 0, loss = Array(0.7326511, dtype=float32), iterm_grads = \\u001b[38;2;79;201;177mState\\u001b[0m\\u001b[38;2;255;213;3m({\\u001b[0m\\u001b[38;2;105;105;105m\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254m'xgrad'\\u001b[0m\\u001b[38;2;212;212;212m: \\u001b[0m\\u001b[38;2;79;201;177mPerturbation\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;105;105;105m # 3 (12 B)\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254mvalue\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0mArray([[-0.430146 , -0.14356601, 0.2935633 ]], dtype=float32)\\n\",\n \" \\u001b[38;2;255;213;3m)\\u001b[0m\\n\",\n \"\\u001b[38;2;255;213;3m})\\u001b[0m\\n\",\n \"step = 1, loss = Array(0.65039134, dtype=float32), iterm_grads = \\u001b[38;2;79;201;177mState\\u001b[0m\\u001b[38;2;255;213;3m({\\u001b[0m\\u001b[38;2;105;105;105m\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254m'xgrad'\\u001b[0m\\u001b[38;2;212;212;212m: \\u001b[0m\\u001b[38;2;79;201;177mPerturbation\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;105;105;105m # 3 (12 B)\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254mvalue\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0mArray([[-0.38535568, -0.11745065, 0.24441527]], dtype=float32)\\n\",\n \" \\u001b[38;2;255;213;3m)\\u001b[0m\\n\",\n \"\\u001b[38;2;255;213;3m})\\u001b[0m\\n\"\n ]\n }\n ],\n \"source\": [\n \"@nnx.jit\\n\",\n \"def train_step(model, optimizer, x, y):\\n\",\n \" graphdef, params, perturbations = nnx.split(model, nnx.Param, nnx.Perturbation)\\n\",\n \"\\n\",\n \" def loss_fn(params, perturbations):\\n\",\n \" model = nnx.merge(graphdef, params, perturbations)\\n\",\n \" return jnp.mean((model(x) - y) ** 2)\\n\",\n \"\\n\",\n \" loss, (grads, iterm_grads) = nnx.value_and_grad(loss_fn, argnums=(0, 1))(params, perturbations)\\n\",\n \" optimizer.update(model, grads)\\n\",\n \"\\n\",\n \" return loss, iterm_grads\\n\",\n \"\\n\",\n \"for step in range(2):\\n\",\n \" loss, iterm_grads = train_step(model, optimizer, x, y)\\n\",\n \" print(f\\\"{step = }, {loss = }, {iterm_grads = !s}\\\")\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"d8511c3d\",\n \"metadata\": {},\n \"source\": [\n \"## Object\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"1c2fd61c\",\n \"metadata\": {},\n \"source\": [\n \"`Object` are NNX types that are **not** registered as JAX pytrees. Formally, any `Object` subclass is a `nnx.Pytree` with `pytree=False`.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 25,\n \"id\": \"a9cab639\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \" nnx.state(foo) = State({\\n\",\n \" 'a': {\\n\",\n \" 0: Array(2, dtype=int32, weak_type=True),\\n\",\n \" 1: Array(4, dtype=int32, weak_type=True)\\n\",\n \" },\\n\",\n \" 'b': Array(6, dtype=int32, weak_type=True)\\n\",\n \"})\\n\",\n \" jax.tree_util.all_leaves([foo]) = True\\n\"\n ]\n }\n ],\n \"source\": [\n \"class Foo(nnx.Object): # instead of Foo(nnx.Pytree, pytree=False)\\n\",\n \" def __init__(self):\\n\",\n \" self.a = [jnp.array(1), jnp.array(2)] # no checks\\n\",\n \" self.b = \\\"hello\\\" \\n\",\n \" self.b = jnp.array(3) # no checks\\n\",\n \"\\n\",\n \"foo = Foo()\\n\",\n \"\\n\",\n \"@nnx.jit # can use in NNX transformations\\n\",\n \"def double(foo: Foo):\\n\",\n \" foo.a = [x * 2 for x in foo.a]\\n\",\n \" foo.b *= 2\\n\",\n \"\\n\",\n \"double(foo)\\n\",\n \"print(f\\\"{ nnx.state(foo) = }\\\") # can be used with NNX APIs\\n\",\n \"print(f\\\"{ jax.tree_util.all_leaves([foo]) = }\\\") # not a pytree\"\n ]\n }\n ],\n \"metadata\": {\n \"jupytext\": {\n \"formats\": \"ipynb,md:myst\"\n },\n \"language_info\": {\n \"codemirror_mode\": {\n \"name\": \"ipython\",\n \"version\": 3\n },\n \"file_extension\": \".py\",\n \"mimetype\": \"text/x-python\",\n \"name\": \"python\",\n \"nbconvert_exporter\": \"python\",\n \"pygments_lexer\": \"ipython3\",\n \"version\": \"3.11.9\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 5\n}\n" }, { "path": "docs_nnx/guides/pytree.md", "content": "---\njupytext:\n formats: ipynb,md:myst\n text_representation:\n extension: .md\n format_name: myst\n format_version: 0.13\n jupytext_version: 1.13.8\n---\n\n# Module & Pytree\n\n+++\n\nFlax NNX's Modules are by default registered as JAX Pytrees, this allows using them throughout most of JAX APIs but in particular JAX transforms and the `jax.tree.*` functions. Thanks to the pytree protocol a simple NNX program might look like this:\n\n```{code-cell} ipython3\nfrom flax import nnx\nimport jax\nimport jax.numpy as jnp\n\nclass Linear(nnx.Module):\n def __init__(self, din, dout, rngs: nnx.Rngs):\n self.din, self.dout = din, dout\n self.kernel = nnx.Param(rngs.normal((din, dout)))\n\nrngs = nnx.Rngs(0)\nweights = Linear(2, 3, rngs=rngs)\n\n@jax.jit\ndef forward(weights, x):\n return x @ weights.kernel\n\ny = forward(weights, x=rngs.uniform((5, 2)))\nprint(f\"{y.shape = }\")\n```\n\nHere `weights`, of type `Linear`, was able to be passed directly to the `jit`-ed function `forward`. Throughout the rest of this guide we will try to answer the questions:\n1. What are pytrees? \n2. How does NNX implement pytrees?\n\n+++\n\n## Pytrees 101\nMost modern ML models have too many Arrays for users to pass around individually, to deal with this JAX developed a way to track Array data in nested structures that still allowed caching for compilation: Pytrees. JAX pytrees are tree structures made of python objects that can be recursively traversed in order to collect an ordered list of leaves and a definition of the tree structure, this is done via the `jax.tree.flatten` function. Most common pytrees are native python containers like `list`, `dict`, and `tuple`, but interestingly it also include `None`. The example bellow shows how to collect all the integer leaves from a nested structure using `flatten`:\n\n```{code-cell} ipython3\npytree = [\n {'a': 1},\n {\n 'b': 2,\n 'c': (3, 4),\n 'd': None,\n }\n]\n\nleaves, treedef = jax.tree.flatten(pytree)\nprint(f\"leaves = {leaves}\")\nprint(f\"treedef = {treedef}\")\n```\n\nNote that `None` is not a leaf because its defined as a pytree with no children. The main purpose of being able to flatten, apart from collecting the leaves, is being able reconstruct the pytree structure from the tree definition from any sequence of leaves of the same length via the `jax.tree.unflatten` function:\n\n```{code-cell} ipython3\nnew_leaves = [x * 10 for x in leaves]\nnew_pytree = jax.tree.unflatten(treedef, new_leaves)\n\nprint(f\"old pytree = {pytree}\")\nprint(f\"new pytree = {new_pytree}\")\n```\n\n### Custom Pytrees\nJAX allows us to register custom pytree node type by using the `jax.tree_util.register_pytree_node` utility. For any type we are able to define a flatten that decomposes the object into a a sequence of nodes / children and a static (hashable) structure, and a unflatten function which takes the sequence of nodes and the static structure and creates a new instance. In the following example we create a simple type `Foo` with the attributes `a`, `b`, and `c`, and define `a` and `b` as nodes and `c` as static.\n\n```{code-cell} ipython3\nclass Foo:\n def __init__(self):\n self.a = 1\n self.b = 2\n self.c = \"hi\"\n\ndef flatten_foo(foo: Foo):\n nodes = [foo.a, foo.b] # sequence of nodes\n static = (foo.c,) # hashable & equatable structure\n return nodes, static\n\ndef unflatten_foo(static, nodes):\n foo = object.__new__(Foo) # create uninitialized instance\n foo.a = nodes[0]\n foo.b = nodes[1]\n foo.c = static[0]\n return foo\n\njax.tree_util.register_pytree_node(Foo, flatten_foo, unflatten_foo)\n\nfoo = Foo()\nleaves, treedef = jax.tree.flatten(foo)\nprint(f\"leaves = {leaves}\")\nprint(f\"treedef = {treedef}\")\n```\n\nNotice that `'hi'` does not appear in the leaves because `c` is defined as static, but you can see it as part of the `PyTreeDef` structure.\n\n+++\n\n## nnx.Pytree\nIn general it would be cumbersome for users to manually register the pytree definition for every type they create. To automate this process NNX provides the `nnx.Pytree` base type that offers a simple API: users annotate attributes using either `nnx.static` or `nnx.data`, and Pytree will register some flatten and unflatten functions that will take the annotations into account. The `nnx.data` and `nnx.static` annotations must only be assigned to `Pytree` attributes directly.\n\n```{code-cell} ipython3\nclass Linear(nnx.Pytree):\n def __init__(self, din: int, dout: int):\n self.din = nnx.static(din)\n self.dout = nnx.static(dout)\n self.w = nnx.data(jnp.ones((din, dout)))\n self.b = nnx.data(jnp.zeros((dout,)))\n\nclass MLP(nnx.Pytree):\n def __init__(self, num_layers, dim):\n self.num_layers = nnx.static(num_layers)\n self.layers = nnx.data([\n Linear(dim, dim) for _ in range(num_layers)\n ])\n\npytree = MLP(num_layers=2, dim=1)\n\ndef pytree_structure(pytree, title='pytree structure'):\n print(f\"{title}:\")\n path_leaves, treedef = jax.tree.flatten_with_path(pytree)\n for path, value in path_leaves:\n print(f\" - pytree{jax.tree_util.keystr(path)} = {value!r}\")\n\npytree_structure(pytree)\n```\n\nAs you can see above, only the `data` paths appear in the leaves. However, its very verbose to have to define `static` and `data` for each attribute, so `Pytree` has sensible defaults. You can remove most of them and it will just work:\n\n```{code-cell} ipython3\nclass Linear(nnx.Pytree):\n def __init__(self, din: int, dout: int):\n self.din = din # static\n self.dout = dout # static\n self.w = jnp.ones((din, dout)) # data\n self.b = jnp.zeros((dout,)) # data\n\nclass MLP(nnx.Pytree):\n def __init__(self, num_layers, dim):\n self.num_layers = num_layers # static\n self.layers = nnx.List([ # data\n Linear(dim, dim) for _ in range(num_layers)\n ])\n\npytree = MLP(num_layers=2, dim=1)\npytree_structure(pytree)\n```\n\nThe only change we had to do here is use `nnx.List` to signal that `layers` contains `data`, the status of the rest of the attributes can be correctly inferred. The rules that determine if a value is data or not are the following:\n\n* `Array`s, `Variable`s, `ArrayRef`s, and `nnx.Pytree`s are data.\n* Types registered using `nnx.register_data_type` are data.\n* All other types are static.\n\nTo check if a value is data use the `nnx.is_data` function which will return its status:\n\n```{code-cell} ipython3\nprint(f\"\"\"\n# ------ DATA ------------\n{nnx.is_data( jnp.array(0) ) = } # Arrays are data\n{nnx.is_data( nnx.Param(1) ) = } # Variables are data\n{nnx.is_data( nnx.Rngs(2) ) = } # nnx.Pytrees are data\n\n# ------ STATIC ------------\n{nnx.is_data( 'hello' ) = } # strings, arbitrary objects\n{nnx.is_data( 42 ) = } # int, float, bool, complex, etc.\n{nnx.is_data( [1, 2.0, 3j, jnp.array(1)] ) = } # list, dict, tuple, regular pytrees\n\"\"\")\n```\n\n### When to use explicit annotations?\n\n+++\n\nThere are cases were you do want to explicitely annotate the attributes to avoid ambiguity or protect yourself against possible edge cases. These include constraining input arguments which might have unexpected types, forcing attributes as data when their type is not treated as data by default, or using `nnx.static` as a way to assert the attribute should not contain data.\n\n```{code-cell} ipython3\nclass Bar(nnx.Pytree):\n def __init__(self, x, use_bias: bool):\n self.x = nnx.data(x) # constrain inputs (e.g. user could pass Array or float)\n self.y = nnx.data(42) # force types that are not data by default\n self.ls = nnx.List([jnp.array(i) for i in range(3)]) # use nnx.List for lists of data\n self.bias = nnx.data(None) # optional values that can be data\n if use_bias:\n self.bias = nnx.Param(jnp.array(0.0))\n\npytree = Bar(1.0, True)\npytree_structure(pytree)\n```\n\n### Dataclasses\n`nnx.Pytree` dataclasses can be created by using the `nnx.dataclass` decorator. To control the status of each field, `nnx.static` and `nnx.data` can be used as `field` specifiers.\n\n```{code-cell} ipython3\nimport dataclasses\n\n@nnx.dataclass\nclass Foo(nnx.Pytree):\n i: int = nnx.data()\n x: jax.Array\n a: int\n s: str = nnx.static(default='hi', kw_only=True)\n\n@nnx.dataclass\nclass Bar(nnx.Pytree):\n ls: list[Foo] = nnx.data()\n shapes: list[int]\n\npytree = Bar(\n ls=[Foo(i, jnp.array(42 * i), hash(i)) for i in range(2)],\n shapes=[8, 16, 32]\n)\npytree_structure(pytree)\n```\n\n`dataclasses.dataclass` can also be used directly, however type checkers will not handle `nnx.static` and `nnx.data` correctly. To solve this `dataclasses.field` can be used by setting `metadata` with the appropriate entry for `static`.\n\n```{code-cell} ipython3\n@dataclasses.dataclass\nclass Bar(nnx.Pytree):\n a: int = dataclasses.field(metadata={'static': False}) # data\n b: str = dataclasses.field(metadata={'static': True}) # static\n\npytree = Bar(a=10, b=\"hello\")\npytree_structure(pytree, title='dataclass pytree structure')\n```\n\n### Attribute Updates\n\n+++\n\nThe status of an attribute is defined during its first assignment and will not change upon reassignment. However, it is possible to override the status by explicitly using `nnx.data` or `nnx.static` on reassignment.\n\n```{code-cell} ipython3\nclass Foo(nnx.Pytree):\n def __init__(self):\n self.a = jnp.array(1.0) # data\n self.b = \"Hello, world!\" # static\n self.c = nnx.data(3.14) # data\n\npytree = Foo()\npytree_structure(pytree, \"original\")\n\npytree.a = \"🤔\" # data status doesn't change\npytree.b = nnx.data(42) # explicit annotation overrides status to data\npytree.c = nnx.static(0.5) # explicit annotation overrides status to static\npytree_structure(pytree, \"updated\")\n```\n\n### Attribute checks\n`Pytree` has a variety of checks to prevent a common class of errors in JAX. This includes checking for Arrays being assigned to new `static` attributes:\n\n```{code-cell} ipython3\nclass Foo(nnx.Pytree):\n def __init__(self, name):\n self.name = nnx.static(name)\n\ntry:\n foo = Foo(name=jnp.array(123))\nexcept ValueError as e:\n print(\"ValueError:\", e)\n```\n\nChecking for Arrays being assigned to known `static` attributes:\n\n```{code-cell} ipython3\ntry:\n foo = Foo(name=\"mattjj\")\n foo.name = jnp.array(123)\nexcept ValueError as e:\n print(\"ValueError:\", e)\n```\n\nChecking for Arrays after `__init__` on `static` attributes that could've been inserted via mutation. This check can be manually trigger via `nnx.check_pytree` at any time.\n\n```{code-cell} ipython3\nclass Foo(nnx.Pytree):\n def __init__(self):\n self.ls = [] # treated as static\n for i in range(5):\n self.ls.append(jnp.array(i)) # append arrays into static attribute\n\ntry:\n foo = Foo() # error: Array found in static attribute after `__init__`\nexcept ValueError as e:\n print(\"ValueError:\", e)\n```\n\nChecking for `nnx.data` or `nnx.static` annotations stored inside nested structures that are not `nnx.Pytree` instances:\n\n```{code-cell} ipython3\nclass Foo(nnx.Pytree):\n def __init__(self):\n self.a = [nnx.data(1), nnx.static(2)] # annotations in sub-pytree\n\ntry:\n foo = Foo()\nexcept ValueError as e:\n print(\"ValueError:\", e)\n```\n\n### Trace-level awareness\nTo prevent tracer leakage NNX will raise an error when trying to update the attribute of a `Pytree` or the value of a `Variable` on instances that are passed as captures to functions called by JAX transforms:\n\n```{code-cell} ipython3\nclass Foo(nnx.Pytree):\n def __init__(self):\n self.count = nnx.data(0)\n\nfoo = Foo()\n\n@jax.vmap # or jit, grad, shard_map, pmap, scan, etc.\ndef increment(n):\n # foo passed as capture\n foo.count += 1 # error!\n\ntry:\n increment(jnp.arange(5))\nexcept Exception as e:\n print(f\"Error: {e}\")\n```\n\n### Reference Sharing\n\n+++\n\nAs the name implies Pytrees should be trees. To check if a structure is a well-defined tree you can use the `nnx.find_duplicates` functions which will return a list of duplicates, where each duplicate is a list of path tuples. In the example below we see that `left` and `right` are shared references therefore `find_duplicates` returns a non-empty list with the paths:\n\n```{code-cell} ipython3\nclass Shared(nnx.Pytree):\n def __init__(self):\n self.x = jnp.array(1.0)\n\nclass Parent(nnx.Pytree):\n def __init__(self):\n self.left = Shared()\n self.right = self.left # reference sharing\n\nm = Parent()\n\nprint(f\"{nnx.find_duplicates(m) = } # not a tree\")\n```\n\nThe main issue is that sharing is not preserved across pytree operations including JAX transforms, and this results in unintended state duplication:\n\n```{code-cell} ipython3\nm = Parent()\nprint(f\"Before: {m.left is m.right = }\")\n\n@jax.jit\ndef f(m):\n print(f\"Inside: {m.left is m.right = }\")\n return m\n\nm = f(m)\nprint(f\"After: {m.left is m.right = }\")\n```\n\nReference sharing is rare in most Machine Learning applications, however if it is required you can either use the `nnx.{split, merge, state, update}` APIs to move the deduplicated state and graph definiton across the JAX transforms:\n\n```{code-cell} ipython3\nm = Parent()\nprint(f\"Before: {m.left is m.right = }\")\ngraphdef, state = nnx.split(m)\n\n@jax.jit\ndef f(graphdef, state):\n m = nnx.merge(graphdef, state)\n print(f\"Inside: {m.left is m.right = }\")\n return nnx.state(m)\n\nstate = f(graphdef, state)\nnnx.update(m, state)\n\nprint(f\"After: {m.left is m.right = }\")\nprint(f\"{state = }\") # deduplicated state\n```\n\nOr alternatively you can use the NNX transforms which preserve shared references:\n\n```{code-cell} ipython3\nm = Parent()\nprint(f\"Before: {m.left is m.right = }\")\n\n@nnx.jit\ndef f(m):\n print(f\"Inside: {m.left is m.right = }\")\n return m\n\nm = f(m)\n\nprint(f\"After: {m.left is m.right = }\")\n```\n\n### Turning off pytree registration\n`nnx.Pytree` allows you to turn off the pytree registration along with the attribute checks for subtypes by setting `pytree` type attribute option to `False`. This can be useful when upgrading to previous NNX code to newer Flax verions as you will still be able to use the NNX APIs or when creating types that should not be treated as pytree because e.g. they shared references.\n\n```{code-cell} ipython3\nclass Foo(nnx.Pytree, pytree=False):\n def __init__(self):\n self.a = [jnp.array(1), jnp.array(2)] # no checks\n self.b = \"hello\" \n self.b = jnp.array(3) # no checks\n\nfoo = Foo()\n\n@nnx.jit # can use in NNX transformations\ndef double(foo: Foo):\n foo.a = [x * 2 for x in foo.a]\n foo.b *= 2\n\ndouble(foo)\nprint(f\"{ nnx.state(foo) = }\") # can be used with NNX APIs\nprint(f\"{ jax.tree_util.all_leaves([foo]) = }\") # not a pytree\n```\n\n## Module\n\n+++\n\nNNX Modules are `Pytree`s that have two additional methods for traking intermediate values: `sow` and `perturb`.\n\n+++\n\n### sow\n`sow` receives a `Variable` type, a `name`, and a `value`, and stores it in the `Module` so it can be retrieved at a later time. As the following example shows, NNX APIs such as `nnx.state` or `nnx.pop` are a good way of retrieving the sowed state, however `pop` is recommended because it explicitly removes the temporary state from the Module.\n\n```{code-cell} ipython3\nclass Block(nnx.Module):\n def __init__(self, din: int, dout: int, rngs: nnx.Rngs):\n self.linear = nnx.Linear(din, dout, rngs=rngs)\n self.bn = nnx.BatchNorm(dout, rngs=rngs)\n self.dropout = nnx.Dropout(0.1, rngs=rngs)\n\n def __call__(self, x):\n y = nnx.relu(self.dropout(self.bn(self.linear(x))))\n self.sow(nnx.Intermediate, \"y_mean\", jnp.mean(y))\n return y\n\nclass MLP(nnx.Module):\n def __init__(self, num_layers, dim, rngs: nnx.Rngs):\n self.blocks = nnx.List([Block(dim, dim, rngs) for _ in range(num_layers)])\n\n def __call__(self, x):\n for block in self.blocks:\n x = block(x)\n return x\n\n\nmodel = MLP(num_layers=3, dim=20, rngs=nnx.Rngs(0))\nx = jnp.ones((10, 20))\ny = model(x)\nintermediates = nnx.pop(model, nnx.Intermediate) # extract intermediate values\nprint(intermediates)\n```\n\n### perturb\n`perturb` is similar to `sow` but it aims to capture the gradient of a value, currently this is a two step process although it might be simplified in the future:\n1. Initialize the pertubation state by running the model once.\n2. Pass the perturbation state as a differentiable target to `grad`.\n\nAs an example lets create a simple model and use `perturb` to get the intermediate gradient `xgrad` for the variable `x`, and initialize the perturbations:\n\n```{code-cell} ipython3\nimport optax\n\nclass Model(nnx.Module):\n def __init__(self, rngs):\n self.linear1 = nnx.Linear(2, 3, rngs=rngs)\n self.linear2 = nnx.Linear(3, 4, rngs=rngs)\n def __call__(self, x):\n x = nnx.gelu(self.linear1(x))\n x = self.perturb('xgrad', x)\n x = self.linear2(x)\n return x\n\nrngs = nnx.Rngs(0)\nmodel = Model(rngs)\noptimizer = nnx.Optimizer(model, tx=optax.sgd(1e-1), wrt=nnx.Param)\nx, y = rngs.uniform((1, 2)), rngs.uniform((1, 4))\n_ = model(x) # initialize perturbations\nprint(f\"{nnx.state(model, nnx.Perturbation) = !s}\")\n```\n\nNext we'll create a training step function that differentiates w.r.t. both the parameters of the model and the perturbations, the later will be the gradients for the intermediate values. `nnx.jit` and `nnx.value_and_grad` will be use to automatically propagate state updates. We'll return the `loss` function and the itermediate gradients.\n\n```{code-cell} ipython3\n@nnx.jit\ndef train_step(model, optimizer, x, y):\n graphdef, params, perturbations = nnx.split(model, nnx.Param, nnx.Perturbation)\n\n def loss_fn(params, perturbations):\n model = nnx.merge(graphdef, params, perturbations)\n return jnp.mean((model(x) - y) ** 2)\n\n loss, (grads, iterm_grads) = nnx.value_and_grad(loss_fn, argnums=(0, 1))(params, perturbations)\n optimizer.update(model, grads)\n\n return loss, iterm_grads\n\nfor step in range(2):\n loss, iterm_grads = train_step(model, optimizer, x, y)\n print(f\"{step = }, {loss = }, {iterm_grads = !s}\")\n```\n\n## Object\n\n+++\n\n`Object` are NNX types that are **not** registered as JAX pytrees. Formally, any `Object` subclass is a `nnx.Pytree` with `pytree=False`.\n\n```{code-cell} ipython3\nclass Foo(nnx.Object): # instead of Foo(nnx.Pytree, pytree=False)\n def __init__(self):\n self.a = [jnp.array(1), jnp.array(2)] # no checks\n self.b = \"hello\" \n self.b = jnp.array(3) # no checks\n\nfoo = Foo()\n\n@nnx.jit # can use in NNX transformations\ndef double(foo: Foo):\n foo.a = [x * 2 for x in foo.a]\n foo.b *= 2\n\ndouble(foo)\nprint(f\"{ nnx.state(foo) = }\") # can be used with NNX APIs\nprint(f\"{ jax.tree_util.all_leaves([foo]) = }\") # not a pytree\n```\n" }, { "path": "docs_nnx/guides/randomness.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"# Randomness\\n\",\n \"\\n\",\n \"Flax NNX uses the stateful `nnx.Rngs` class to simplify Jax's handling of random states. For example, the code below uses a `nnx.Rngs` object to define a simple linear model with dropout:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 2,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"from flax import nnx\\n\",\n \"\\n\",\n \"class Model(nnx.Module):\\n\",\n \" def __init__(self, *, rngs: nnx.Rngs):\\n\",\n \" self.linear = nnx.Linear(20, 10, rngs=rngs)\\n\",\n \" self.drop = nnx.Dropout(0.1)\\n\",\n \"\\n\",\n \" def __call__(self, x, *, rngs):\\n\",\n \" return nnx.relu(self.drop(self.linear(x), rngs=rngs))\\n\",\n \"\\n\",\n \"rngs = nnx.Rngs(0)\\n\",\n \"model = Model(rngs=rngs) # pass rngs to initialize parameters\\n\",\n \"x = rngs.normal((32, 20)) # convenient jax.random methods\\n\",\n \"y = model(x, rngs=rngs) # pass rngs for dropout masks\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"We always pass `nnx.Rngs` objects to models at initialization (to initialize parameters). For models with nondeterministic outputs like the one above, we also pass `nnx.Rngs` objects to the model's `__call__` method.\\n\",\n \"\\n\",\n \"The Flax NNX [pseudorandom number generator (PRNG)](https://flax.readthedocs.io/en/latest/glossary.html#term-RNG-sequences) system has the following main characteristics:\\n\",\n \"\\n\",\n \"- It is **explicit**.\\n\",\n \"- It is **order-based**.\\n\",\n \"- It uses **dynamic counters**.\\n\",\n \"\\n\",\n \"> **Note:** To learn more about random number generation in JAX, the `jax.random` API, and PRNG-generated sequences, check out this [JAX PRNG tutorial](https://jax.readthedocs.io/en/latest/random-numbers.html).\\n\",\n \"\\n\",\n \"## `Rngs`, `RngStream`, and `RngState`\\n\",\n \"\\n\",\n \"In Flax NNX, the `nnx.Rngs` type is the primary convenience API for managing the random state(s). Following Flax Linen's footsteps, `nnx.Rngs` have the ability to create multiple named PRNG key [streams](https://jax.readthedocs.io/en/latest/jep/263-prng.html), each with its own state, for the purpose of having tight control over randomness in the context of [JAX transformations (transforms)](https://jax.readthedocs.io/en/latest/key-concepts.html#transformations).\\n\",\n \"\\n\",\n \"Here are the main PRNG-related types in Flax NNX:\\n\",\n \"\\n\",\n \"* **`nnx.Rngs`**: The main user interface. It defines a set of named `nnx.RngStream` objects.\\n\",\n \"* **`nnx.RngStream`**: An object that can generate a stream of PRNG keys. It holds a root `key` and a `count` inside an `nnx.RngKey` and `nnx.RngCount` `nnx.Variable`s, respectively. When a new key is generated, the count is incremented.\\n\",\n \"* **`nnx.RngState`**: The base type for all RNG-related states.\\n\",\n \" * **`nnx.RngKey`**: NNX Variable type for holding PRNG keys. It includes a `tag` attribute containing the name of the PRNG key stream.\\n\",\n \" * **`nnx.RngCount`**: NNX Variable type for holding PRNG counts. It includes a `tag` attribute containing the PRNG key stream name.\\n\",\n \"\\n\",\n \"To create an `nnx.Rngs` object you can simply pass an integer seed or `jax.random.key` instance to any keyword argument of your choice in the constructor.\\n\",\n \"\\n\",\n \"Here's an example:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 3,\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"rngs = nnx.Rngs(params=0, dropout=random.key(1))\\n\",\n \"nnx.display(rngs)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Notice that the `key` and `count` `nnx.Variable`s contain the PRNG key stream name in a `tag` attribute. This is primarily used for filtering as we'll see later.\\n\",\n \"\\n\",\n \"To generate new keys, you can access one of the streams and use its `__call__` method with no arguments. This will return a new key by using `random.fold_in` with the current `key` and `count`. The `count` is then incremented so that subsequent calls will return new keys.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 4,\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"params_key = rngs.params()\\n\",\n \"dropout_key = rngs.dropout()\\n\",\n \"\\n\",\n \"nnx.display(rngs)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Note that the `key` attribute does not change when new PRNG keys are generated.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Using random state with flax Modules.\\n\",\n \"\\n\",\n \"Almost all flax Modules require a random state for initialization. In a `Linear` layer, for example, we need to sample the weights and biases from the appropriate Normal distribution. Random state is provided using the `rngs` keyword argument at initialization.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 5,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"linear = nnx.Linear(20, 10, rngs=rngs)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Specifically, this will use the RngSteam `rngs.params` for weight initialization. The `params` stream is also used for initialization of `nnx.Conv`, `nnx.ConvTranspose`, and `nnx.Embed`.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"The `nnx.Dropout` module also requires a random state, but it requires this state at *call* time rather than initialization. Once again, we can pass it random state using the `rngs` keyword argument.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 6,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"dropout = nnx.Dropout(0.5)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 7,\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"Array([2., 0., 2., 2.], dtype=float32)\"\n ]\n },\n \"execution_count\": 7,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"import jax.numpy as jnp\\n\",\n \"dropout(jnp.ones(4), rngs=rngs)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"The `nnx.Dropout` layer will use the rng's `dropout` stream. This also applies to Modules that use `Dropout` as a sub-Module, like `nnx.MultiHeadAttention`.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"To summarize, there are only two standard PRNG key stream names used by Flax NNX's built-in layers, shown in the table below:\\n\",\n \"\\n\",\n \"| PRNG key stream name | Description |\\n\",\n \"|----------------------|-----------------------------------------------|\\n\",\n \"| `params` | Used for parameter initialization |\\n\",\n \"| `dropout` | Used by `nnx.Dropout` to create dropout masks |\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Default PRNG key stream\\n\",\n \"\\n\",\n \"One of the downsides of having named streams is that the user needs to know all the possible names that a model will use when creating the `nnx.Rngs` object. While this could be solved with some documentation, Flax NNX provides a `default` stream that can be\\n\",\n \"be used as a fallback when a stream is not found. To use the default PRNG key stream, you can simply pass an integer seed or `jax.random.key` as the first positional argument.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 8,\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"rngs = nnx.Rngs(0, params=1)\\n\",\n \"\\n\",\n \"key1 = rngs.params() # Call params.\\n\",\n \"key2 = rngs.dropout() # Fallback to the default stream.\\n\",\n \"key3 = rngs() # Call the default stream directly.\\n\",\n \"\\n\",\n \"nnx.display(rngs)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"As shown above, a PRNG key from the `default` stream can also be generated by calling the `nnx.Rngs` object itself.\\n\",\n \"\\n\",\n \"> **Note**\\n\",\n \">
For large projects it is recommended to use named streams to avoid potential conflicts. For small projects or quick prototyping just using the `default` stream is a good choice.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### jax.random shorthand methods\\n\",\n \"Since a very common pattern is to sample a key and immediately pass it to a function from `jax.random`, both `Rngs` and `RngStream` expose the same functions as methods with the same signature except they don't require a key:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 9,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"import jax\\n\",\n \"rngs = nnx.Rngs(0, params=1)\\n\",\n \"\\n\",\n \"# using jax.random\\n\",\n \"z1 = jax.random.normal(rngs(), (2, 3))\\n\",\n \"z2 = jax.random.bernoulli(rngs.params(), 0.5, (10,))\\n\",\n \"\\n\",\n \"# shorthand methods\\n\",\n \"z1 = rngs.normal((2, 3)) # generates key from rngs.default\\n\",\n \"z2 = rngs.params.bernoulli(0.5, (10,)) # generates key from rngs.params\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Forking random state\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Say you want to train a model that uses dropout on a batch of data. You don't want to use the same random state for every dropout mask in your batch. Instead, you want to fork the random state into separate pieces for each layer. This can be accomplished with the `fork` method, as shown below.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 10,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"class Model(nnx.Module):\\n\",\n \" def __init__(self, rngs: nnx.Rngs):\\n\",\n \" self.linear = nnx.Linear(20, 10, rngs=rngs)\\n\",\n \" self.drop = nnx.Dropout(0.1)\\n\",\n \"\\n\",\n \" def __call__(self, x, rngs):\\n\",\n \" return nnx.relu(self.drop(self.linear(x), rngs=rngs))\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 11,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"model = Model(rngs=nnx.Rngs(0))\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 12,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"@nnx.vmap(in_axes=(None, 0, 0), out_axes=0)\\n\",\n \"def model_forward(model, x, rngs):\\n\",\n \" return model(x, rngs=rngs)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 13,\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"(Rngs( # RngState: 2 (12 B)\\n\",\n \" default=RngStream( # RngState: 2 (12 B)\\n\",\n \" tag='default',\\n\",\n \" key=RngKey( # 1 (8 B)\\n\",\n \" value=Array((), dtype=key) overlaying:\\n\",\n \" [0 1],\\n\",\n \" tag='default'\\n\",\n \" ),\\n\",\n \" count=RngCount( # 1 (4 B)\\n\",\n \" value=Array(1, dtype=uint32),\\n\",\n \" tag='default'\\n\",\n \" )\\n\",\n \" )\\n\",\n \" ),\\n\",\n \" Rngs( # RngState: 10 (60 B)\\n\",\n \" default=RngStream( # RngState: 10 (60 B)\\n\",\n \" tag='default',\\n\",\n \" key=RngKey( # 5 (40 B)\\n\",\n \" value=Array(shape=(5,), dtype=key),\\n\",\n \" tag='default'\\n\",\n \" ),\\n\",\n \" count=RngCount( # 5 (20 B)\\n\",\n \" value=Array(shape=(5,), dtype=dtype('uint32')),\\n\",\n \" tag='default'\\n\",\n \" )\\n\",\n \" )\\n\",\n \" ))\"\n ]\n },\n \"execution_count\": 13,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"dropout_rngs = nnx.Rngs(1)\\n\",\n \"forked_rngs = dropout_rngs.fork(split=5)\\n\",\n \"(dropout_rngs, forked_rngs)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 14,\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"(5, 10)\"\n ]\n },\n \"execution_count\": 14,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"model_forward(model, jnp.ones((5, 20)), forked_rngs).shape\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"The output of `rng.fork` is another `Rng` with keys and counts that have an expanded shape. In the example above, the `RngKey` and `RngCount` of `dropout_rngs` have shape `()`, but in `forked_rngs` they have shape `(5,)`.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"# Implicit Random State\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"So far, we have looked at passing random state directly to each Module when it gets called. But there's another way to handle call-time randomness in flax: we can bundle the random state into the Module itself. This makes the random state is just another type of Module state. Using implicit random state requires passing the `rngs` keyward argument when initializing the module rather than when calling it. For example, here is how we might construct the simple `Module` we defined earlier using an implicit style.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 18,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"y.shape = (1, 10)\\n\"\n ]\n }\n ],\n \"source\": [\n \"class Model(nnx.Module):\\n\",\n \" def __init__(self, rngs: nnx.Rngs):\\n\",\n \" self.linear = nnx.Linear(20, 10, rngs=rngs)\\n\",\n \" self.drop = nnx.Dropout(0.1, rngs=rngs)\\n\",\n \"\\n\",\n \" def __call__(self, x):\\n\",\n \" return nnx.relu(self.drop(self.linear(x)))\\n\",\n \"\\n\",\n \"model = Model(nnx.Rngs(params=0, dropout=1))\\n\",\n \"\\n\",\n \"y = model(x=jnp.ones((1, 20)))\\n\",\n \"print(f'{y.shape = }')\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"This implicit state handling style is less verbose than passing RNGs explicitly, and more closely resembles code in other deep learning frameworks like PyTorch. However, as we'll see in the following sections, using implicit state makes it less obvious how to apply jax transformations to your Modules. With explicit state, you can usually use tranforms like `jax.vmap` directly. With implicit state, you'll need to some extra tricks with `nnx.vmap` to make everything work. Because of this additional complexity, we recommend that new flax projects stick to the explicit style.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Filtering random state\\n\",\n \"\\n\",\n \"Implicit random state can be manipulated using [Filters](https://flax.readthedocs.io/en/latest/guides/filters_guide.html) just like any other type of state. It can be filtered using types (`nnx.RngState`, `nnx.RngKey`, `nnx.RngCount`) or using strings corresponding to the stream names (refer to [the Flax NNX `Filter` DSL](https://flax.readthedocs.io/en/latest/guides/filters_guide.html#the-filter-dsl)). Here's an example using `nnx.state` with various filters to select different substates of the `Rngs` inside a `Model`:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 21,\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"model = Model(nnx.Rngs(params=0, dropout=1))\\n\",\n \"\\n\",\n \"rng_state = nnx.state(model, nnx.RngState) # All random states.\\n\",\n \"key_state = nnx.state(model, nnx.RngKey) # Only PRNG keys.\\n\",\n \"count_state = nnx.state(model, nnx.RngCount) # Only counts.\\n\",\n \"rng_dropout_state = nnx.state(model, 'dropout') # Only `dropout`.\\n\",\n \"\\n\",\n \"nnx.display(rng_dropout_state)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Reseeding\\n\",\n \"\\n\",\n \"In Haiku and Flax Linen, random states are explicitly passed to `Module.apply` each time before you call the model. This makes it easy to control the randomness of the model when needed (for example, for reproducibility).\\n\",\n \"\\n\",\n \"In Flax NNX, there are two ways to approach this:\\n\",\n \"\\n\",\n \"1. By passing an `nnx.Rngs` object through the `__call__` stack manually, as shown previously.\\n\",\n \"2. By using `nnx.reseed` to set the random state of the model to a specific configuration. This option is less intrusive and can be used even if the model is not designed to enable manual control over the random state.\\n\",\n \"\\n\",\n \"`nnx.reseed` is a function that accepts an arbitrary graph node (this includes [pytrees](https://jax.readthedocs.io/en/latest/working-with-pytrees.html#working-with-pytrees) of `nnx.Module`s) and some keyword arguments containing the new seed or key value for the `nnx.RngStream`s specified by the argument names. `nnx.reseed` will then traverse the graph and update the random state of the matching `nnx.RngStream`s, this includes both setting the `key` to a possibly new value and resetting the `count` to zero.\\n\",\n \"\\n\",\n \"Here's an example of how to use `nnx.reseed` to reset the random state of the `nnx.Dropout` layer and verify that the computation is identical to the first time the model was called:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 22,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"model = Model(nnx.Rngs(params=0, dropout=1))\\n\",\n \"x = jnp.ones((1, 20))\\n\",\n \"\\n\",\n \"y1 = model(x)\\n\",\n \"y2 = model(x)\\n\",\n \"\\n\",\n \"nnx.reseed(model, dropout=1) # reset dropout RngState\\n\",\n \"y3 = model(x)\\n\",\n \"\\n\",\n \"assert not jnp.allclose(y1, y2) # different\\n\",\n \"assert jnp.allclose(y1, y3) # same\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Forking implicit random state\\n\",\n \"\\n\",\n \"We saw above how to use `rng.fork` when passing explicit random state through [Flax NNX transforms](https://flax.readthedocs.io/en/latest/guides/transforms.html) like `nnx.vmap` or `nnx.pmap`. The decorator `nnx.fork_rngs` allows this for implicit random state. Consider the example below, which generates a batch of samples from the nondeterministic model we defined above.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 23,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"(5, 1, 10)\\n\"\n ]\n }\n ],\n \"source\": [\n \"rng_axes = nnx.StateAxes({'dropout': 0, ...: None})\\n\",\n \"\\n\",\n \"@nnx.fork_rngs(split={'dropout': 5})\\n\",\n \"@nnx.vmap(in_axes=(rng_axes, None), out_axes=0)\\n\",\n \"def sample_from_model(model, x):\\n\",\n \" return model(x)\\n\",\n \"\\n\",\n \"print(sample_from_model(model, x).shape)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Here `sample_from_model` is modified by two decorators:\\n\",\n \"- The function we get from the `nnx.vmap` decorator expects that the random state of the `model` argument has already been split into 5 pieces. It runs the model once for each random key.\\n\",\n \"- The function we get from the `nnx.fork_rngs` decorator splits the random state of its `model` argument into five pieces before passing it on to the inner function.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Transforming implicit state\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"In the previous section, we showed how to use `nnx.vmap` with a module that contained implicit random state. But we can use other `nnx` transformations too! Remember: implicit random state isn't different from any other type of Model state, and this applies to Flax NNX transforms too. This means you can use the Flax NNX state handling APIs of each transform to get the results you want. For a more involved example, let’s explore how to implement recurrent dropout on an `RNNCell` using `nnx.scan`.\\n\",\n \"\\n\",\n \"We'll start by constructing the `RNNCell` class:\\n\",\n \"\\n\",\n \"- First, create an `nnx.Dropout` layer that will sample PRNG keys from a custom `recurrent_dropout` stream.\\n\",\n \"- Apply dropout (`drop`) to the hidden state `h` of the `RNNCell`.\\n\",\n \"- Then, define an `initial_state` function to create the initial state of the `RNNCell`.\\n\",\n \"- Finally, instantiate `RNNCell`.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 24,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"class Count(nnx.Variable): pass\\n\",\n \"\\n\",\n \"class RNNCell(nnx.Module):\\n\",\n \" def __init__(self, din, dout, rngs):\\n\",\n \" self.linear = nnx.Linear(dout + din, dout, rngs=rngs)\\n\",\n \" self.drop = nnx.Dropout(0.1, rngs=rngs, rng_collection='recurrent_dropout')\\n\",\n \" self.dout = dout\\n\",\n \" self.count = Count(jnp.array(0, jnp.uint32))\\n\",\n \"\\n\",\n \" def __call__(self, h, x) -> tuple[jax.Array, jax.Array]:\\n\",\n \" h = self.drop(h) # Recurrent dropout.\\n\",\n \" y = nnx.relu(self.linear(jnp.concatenate([h, x], axis=-1)))\\n\",\n \" self.count.value += 1\\n\",\n \" return y, y\\n\",\n \"\\n\",\n \" def initial_state(self, batch_size: int):\\n\",\n \" return jnp.zeros((batch_size, self.dout))\\n\",\n \"\\n\",\n \"cell = RNNCell(8, 16, nnx.Rngs(params=0, recurrent_dropout=1))\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Next, use `nnx.scan` over an `unroll` function to implement the `rnn_forward` operation:\\n\",\n \"- The key ingredient of recurrent dropout is to apply the same dropout mask across all time steps. Therefore, to achieve this you will pass `nnx.StateAxes` to `nnx.scan`'s `in_axes`, specifying that the `cell`'s `recurrent_dropout` PRNG stream will be broadcast, and the rest of the `RNNCell`'s state will be carried over.\\n\",\n \"- Also, the hidden state `h` will be the `nnx.scan`'s `Carry` variable, and the sequence `x` will be `scan`ned over its axis `1`.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 25,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"y.shape = (4, 20, 16)\\n\",\n \"cell.count.value = Array(20, dtype=uint32)\\n\"\n ]\n }\n ],\n \"source\": [\n \"@nnx.jit\\n\",\n \"def rnn_forward(cell: RNNCell, x: jax.Array):\\n\",\n \" h = cell.initial_state(batch_size=x.shape[0])\\n\",\n \"\\n\",\n \" # Broadcast the 'recurrent_dropout' PRNG state to have the same mask on every step.\\n\",\n \" state_axes = nnx.StateAxes({'recurrent_dropout': None, ...: nnx.Carry})\\n\",\n \" @nnx.scan(in_axes=(state_axes, nnx.Carry, 1), out_axes=(nnx.Carry, 1))\\n\",\n \" def unroll(cell: RNNCell, h, x) -> tuple[jax.Array, jax.Array]:\\n\",\n \" h, y = cell(h, x)\\n\",\n \" return h, y\\n\",\n \"\\n\",\n \" h, y = unroll(cell, h, x)\\n\",\n \" return y\\n\",\n \"\\n\",\n \"x = jnp.ones((4, 20, 8))\\n\",\n \"y = rnn_forward(cell, x)\\n\",\n \"\\n\",\n \"print(f'{y.shape = }')\\n\",\n \"print(f'{cell.count.value = }')\"\n ]\n }\n ],\n \"metadata\": {\n \"jupytext\": {\n \"formats\": \"ipynb,md:myst\"\n },\n \"language_info\": {\n \"codemirror_mode\": {\n \"name\": \"ipython\",\n \"version\": 3\n },\n \"file_extension\": \".py\",\n \"mimetype\": \"text/x-python\",\n \"name\": \"python\",\n \"nbconvert_exporter\": \"python\",\n \"pygments_lexer\": \"ipython3\",\n \"version\": \"3.11.6\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 4\n}\n" }, { "path": "docs_nnx/guides/randomness.md", "content": "---\njupytext:\n formats: ipynb,md:myst\n text_representation:\n extension: .md\n format_name: myst\n format_version: 0.13\n jupytext_version: 1.13.8\n---\n\n# Randomness\n\nFlax NNX uses the stateful `nnx.Rngs` class to simplify Jax's handling of random states. For example, the code below uses a `nnx.Rngs` object to define a simple linear model with dropout:\n\n```{code-cell} ipython3\nfrom flax import nnx\n\nclass Model(nnx.Module):\n def __init__(self, *, rngs: nnx.Rngs):\n self.linear = nnx.Linear(20, 10, rngs=rngs)\n self.drop = nnx.Dropout(0.1)\n\n def __call__(self, x, *, rngs):\n return nnx.relu(self.drop(self.linear(x), rngs=rngs))\n\nrngs = nnx.Rngs(0)\nmodel = Model(rngs=rngs) # pass rngs to initialize parameters\nx = rngs.normal((32, 20)) # convenient jax.random methods\ny = model(x, rngs=rngs) # pass rngs for dropout masks\n```\n\nWe always pass `nnx.Rngs` objects to models at initialization (to initialize parameters). For models with nondeterministic outputs like the one above, we also pass `nnx.Rngs` objects to the model's `__call__` method.\n\nThe Flax NNX [pseudorandom number generator (PRNG)](https://flax.readthedocs.io/en/latest/glossary.html#term-RNG-sequences) system has the following main characteristics:\n\n- It is **explicit**.\n- It is **order-based**.\n- It uses **dynamic counters**.\n\n> **Note:** To learn more about random number generation in JAX, the `jax.random` API, and PRNG-generated sequences, check out this [JAX PRNG tutorial](https://jax.readthedocs.io/en/latest/random-numbers.html).\n\n## `Rngs`, `RngStream`, and `RngState`\n\nIn Flax NNX, the `nnx.Rngs` type is the primary convenience API for managing the random state(s). Following Flax Linen's footsteps, `nnx.Rngs` have the ability to create multiple named PRNG key [streams](https://jax.readthedocs.io/en/latest/jep/263-prng.html), each with its own state, for the purpose of having tight control over randomness in the context of [JAX transformations (transforms)](https://jax.readthedocs.io/en/latest/key-concepts.html#transformations).\n\nHere are the main PRNG-related types in Flax NNX:\n\n* **`nnx.Rngs`**: The main user interface. It defines a set of named `nnx.RngStream` objects.\n* **`nnx.RngStream`**: An object that can generate a stream of PRNG keys. It holds a root `key` and a `count` inside an `nnx.RngKey` and `nnx.RngCount` `nnx.Variable`s, respectively. When a new key is generated, the count is incremented.\n* **`nnx.RngState`**: The base type for all RNG-related states.\n * **`nnx.RngKey`**: NNX Variable type for holding PRNG keys. It includes a `tag` attribute containing the name of the PRNG key stream.\n * **`nnx.RngCount`**: NNX Variable type for holding PRNG counts. It includes a `tag` attribute containing the PRNG key stream name.\n\nTo create an `nnx.Rngs` object you can simply pass an integer seed or `jax.random.key` instance to any keyword argument of your choice in the constructor.\n\nHere's an example:\n\n```{code-cell} ipython3\nrngs = nnx.Rngs(params=0, dropout=random.key(1))\nnnx.display(rngs)\n```\n\nNotice that the `key` and `count` `nnx.Variable`s contain the PRNG key stream name in a `tag` attribute. This is primarily used for filtering as we'll see later.\n\nTo generate new keys, you can access one of the streams and use its `__call__` method with no arguments. This will return a new key by using `random.fold_in` with the current `key` and `count`. The `count` is then incremented so that subsequent calls will return new keys.\n\n```{code-cell} ipython3\nparams_key = rngs.params()\ndropout_key = rngs.dropout()\n\nnnx.display(rngs)\n```\n\nNote that the `key` attribute does not change when new PRNG keys are generated.\n\n+++\n\n### Using random state with flax Modules.\n\nAlmost all flax Modules require a random state for initialization. In a `Linear` layer, for example, we need to sample the weights and biases from the appropriate Normal distribution. Random state is provided using the `rngs` keyword argument at initialization.\n\n```{code-cell} ipython3\nlinear = nnx.Linear(20, 10, rngs=rngs)\n```\n\nSpecifically, this will use the RngSteam `rngs.params` for weight initialization. The `params` stream is also used for initialization of `nnx.Conv`, `nnx.ConvTranspose`, and `nnx.Embed`.\n\n+++\n\nThe `nnx.Dropout` module also requires a random state, but it requires this state at *call* time rather than initialization. Once again, we can pass it random state using the `rngs` keyword argument.\n\n```{code-cell} ipython3\ndropout = nnx.Dropout(0.5)\n```\n\n```{code-cell} ipython3\nimport jax.numpy as jnp\ndropout(jnp.ones(4), rngs=rngs)\n```\n\nThe `nnx.Dropout` layer will use the rng's `dropout` stream. This also applies to Modules that use `Dropout` as a sub-Module, like `nnx.MultiHeadAttention`.\n\n+++\n\nTo summarize, there are only two standard PRNG key stream names used by Flax NNX's built-in layers, shown in the table below:\n\n| PRNG key stream name | Description |\n|----------------------|-----------------------------------------------|\n| `params` | Used for parameter initialization |\n| `dropout` | Used by `nnx.Dropout` to create dropout masks |\n\n+++\n\n### Default PRNG key stream\n\nOne of the downsides of having named streams is that the user needs to know all the possible names that a model will use when creating the `nnx.Rngs` object. While this could be solved with some documentation, Flax NNX provides a `default` stream that can be\nbe used as a fallback when a stream is not found. To use the default PRNG key stream, you can simply pass an integer seed or `jax.random.key` as the first positional argument.\n\n```{code-cell} ipython3\nrngs = nnx.Rngs(0, params=1)\n\nkey1 = rngs.params() # Call params.\nkey2 = rngs.dropout() # Fallback to the default stream.\nkey3 = rngs() # Call the default stream directly.\n\nnnx.display(rngs)\n```\n\nAs shown above, a PRNG key from the `default` stream can also be generated by calling the `nnx.Rngs` object itself.\n\n> **Note**\n>
For large projects it is recommended to use named streams to avoid potential conflicts. For small projects or quick prototyping just using the `default` stream is a good choice.\n\n+++\n\n### jax.random shorthand methods\nSince a very common pattern is to sample a key and immediately pass it to a function from `jax.random`, both `Rngs` and `RngStream` expose the same functions as methods with the same signature except they don't require a key:\n\n```{code-cell} ipython3\nimport jax\nrngs = nnx.Rngs(0, params=1)\n\n# using jax.random\nz1 = jax.random.normal(rngs(), (2, 3))\nz2 = jax.random.bernoulli(rngs.params(), 0.5, (10,))\n\n# shorthand methods\nz1 = rngs.normal((2, 3)) # generates key from rngs.default\nz2 = rngs.params.bernoulli(0.5, (10,)) # generates key from rngs.params\n```\n\n## Forking random state\n\n+++\n\nSay you want to train a model that uses dropout on a batch of data. You don't want to use the same random state for every dropout mask in your batch. Instead, you want to fork the random state into separate pieces for each layer. This can be accomplished with the `fork` method, as shown below.\n\n```{code-cell} ipython3\nclass Model(nnx.Module):\n def __init__(self, rngs: nnx.Rngs):\n self.linear = nnx.Linear(20, 10, rngs=rngs)\n self.drop = nnx.Dropout(0.1)\n\n def __call__(self, x, rngs):\n return nnx.relu(self.drop(self.linear(x), rngs=rngs))\n```\n\n```{code-cell} ipython3\nmodel = Model(rngs=nnx.Rngs(0))\n```\n\n```{code-cell} ipython3\n@nnx.vmap(in_axes=(None, 0, 0), out_axes=0)\ndef model_forward(model, x, rngs):\n return model(x, rngs=rngs)\n```\n\n```{code-cell} ipython3\ndropout_rngs = nnx.Rngs(1)\nforked_rngs = dropout_rngs.fork(split=5)\n(dropout_rngs, forked_rngs)\n```\n\n```{code-cell} ipython3\nmodel_forward(model, jnp.ones((5, 20)), forked_rngs).shape\n```\n\nThe output of `rng.fork` is another `Rng` with keys and counts that have an expanded shape. In the example above, the `RngKey` and `RngCount` of `dropout_rngs` have shape `()`, but in `forked_rngs` they have shape `(5,)`.\n\n+++\n\n# Implicit Random State\n\n+++\n\nSo far, we have looked at passing random state directly to each Module when it gets called. But there's another way to handle call-time randomness in flax: we can bundle the random state into the Module itself. This makes the random state is just another type of Module state. Using implicit random state requires passing the `rngs` keyward argument when initializing the module rather than when calling it. For example, here is how we might construct the simple `Module` we defined earlier using an implicit style.\n\n```{code-cell} ipython3\nclass Model(nnx.Module):\n def __init__(self, rngs: nnx.Rngs):\n self.linear = nnx.Linear(20, 10, rngs=rngs)\n self.drop = nnx.Dropout(0.1, rngs=rngs)\n\n def __call__(self, x):\n return nnx.relu(self.drop(self.linear(x)))\n\nmodel = Model(nnx.Rngs(params=0, dropout=1))\n\ny = model(x=jnp.ones((1, 20)))\nprint(f'{y.shape = }')\n```\n\nThis implicit state handling style is less verbose than passing RNGs explicitly, and more closely resembles code in other deep learning frameworks like PyTorch. However, as we'll see in the following sections, using implicit state makes it less obvious how to apply jax transformations to your Modules. With explicit state, you can usually use tranforms like `jax.vmap` directly. With implicit state, you'll need to some extra tricks with `nnx.vmap` to make everything work. Because of this additional complexity, we recommend that new flax projects stick to the explicit style.\n\n+++\n\n## Filtering random state\n\nImplicit random state can be manipulated using [Filters](https://flax.readthedocs.io/en/latest/guides/filters_guide.html) just like any other type of state. It can be filtered using types (`nnx.RngState`, `nnx.RngKey`, `nnx.RngCount`) or using strings corresponding to the stream names (refer to [the Flax NNX `Filter` DSL](https://flax.readthedocs.io/en/latest/guides/filters_guide.html#the-filter-dsl)). Here's an example using `nnx.state` with various filters to select different substates of the `Rngs` inside a `Model`:\n\n```{code-cell} ipython3\nmodel = Model(nnx.Rngs(params=0, dropout=1))\n\nrng_state = nnx.state(model, nnx.RngState) # All random states.\nkey_state = nnx.state(model, nnx.RngKey) # Only PRNG keys.\ncount_state = nnx.state(model, nnx.RngCount) # Only counts.\nrng_dropout_state = nnx.state(model, 'dropout') # Only `dropout`.\n\nnnx.display(rng_dropout_state)\n```\n\n## Reseeding\n\nIn Haiku and Flax Linen, random states are explicitly passed to `Module.apply` each time before you call the model. This makes it easy to control the randomness of the model when needed (for example, for reproducibility).\n\nIn Flax NNX, there are two ways to approach this:\n\n1. By passing an `nnx.Rngs` object through the `__call__` stack manually, as shown previously.\n2. By using `nnx.reseed` to set the random state of the model to a specific configuration. This option is less intrusive and can be used even if the model is not designed to enable manual control over the random state.\n\n`nnx.reseed` is a function that accepts an arbitrary graph node (this includes [pytrees](https://jax.readthedocs.io/en/latest/working-with-pytrees.html#working-with-pytrees) of `nnx.Module`s) and some keyword arguments containing the new seed or key value for the `nnx.RngStream`s specified by the argument names. `nnx.reseed` will then traverse the graph and update the random state of the matching `nnx.RngStream`s, this includes both setting the `key` to a possibly new value and resetting the `count` to zero.\n\nHere's an example of how to use `nnx.reseed` to reset the random state of the `nnx.Dropout` layer and verify that the computation is identical to the first time the model was called:\n\n```{code-cell} ipython3\nmodel = Model(nnx.Rngs(params=0, dropout=1))\nx = jnp.ones((1, 20))\n\ny1 = model(x)\ny2 = model(x)\n\nnnx.reseed(model, dropout=1) # reset dropout RngState\ny3 = model(x)\n\nassert not jnp.allclose(y1, y2) # different\nassert jnp.allclose(y1, y3) # same\n```\n\n## Forking implicit random state\n\nWe saw above how to use `rng.fork` when passing explicit random state through [Flax NNX transforms](https://flax.readthedocs.io/en/latest/guides/transforms.html) like `nnx.vmap` or `nnx.pmap`. The decorator `nnx.fork_rngs` allows this for implicit random state. Consider the example below, which generates a batch of samples from the nondeterministic model we defined above.\n\n```{code-cell} ipython3\nrng_axes = nnx.StateAxes({'dropout': 0, ...: None})\n\n@nnx.fork_rngs(split={'dropout': 5})\n@nnx.vmap(in_axes=(rng_axes, None), out_axes=0)\ndef sample_from_model(model, x):\n return model(x)\n\nprint(sample_from_model(model, x).shape)\n```\n\nHere `sample_from_model` is modified by two decorators:\n- The function we get from the `nnx.vmap` decorator expects that the random state of the `model` argument has already been split into 5 pieces. It runs the model once for each random key.\n- The function we get from the `nnx.fork_rngs` decorator splits the random state of its `model` argument into five pieces before passing it on to the inner function.\n\n+++\n\n## Transforming implicit state\n\n+++\n\nIn the previous section, we showed how to use `nnx.vmap` with a module that contained implicit random state. But we can use other `nnx` transformations too! Remember: implicit random state isn't different from any other type of Model state, and this applies to Flax NNX transforms too. This means you can use the Flax NNX state handling APIs of each transform to get the results you want. For a more involved example, let’s explore how to implement recurrent dropout on an `RNNCell` using `nnx.scan`.\n\nWe'll start by constructing the `RNNCell` class:\n\n- First, create an `nnx.Dropout` layer that will sample PRNG keys from a custom `recurrent_dropout` stream.\n- Apply dropout (`drop`) to the hidden state `h` of the `RNNCell`.\n- Then, define an `initial_state` function to create the initial state of the `RNNCell`.\n- Finally, instantiate `RNNCell`.\n\n```{code-cell} ipython3\nclass Count(nnx.Variable): pass\n\nclass RNNCell(nnx.Module):\n def __init__(self, din, dout, rngs):\n self.linear = nnx.Linear(dout + din, dout, rngs=rngs)\n self.drop = nnx.Dropout(0.1, rngs=rngs, rng_collection='recurrent_dropout')\n self.dout = dout\n self.count = Count(jnp.array(0, jnp.uint32))\n\n def __call__(self, h, x) -> tuple[jax.Array, jax.Array]:\n h = self.drop(h) # Recurrent dropout.\n y = nnx.relu(self.linear(jnp.concatenate([h, x], axis=-1)))\n self.count.value += 1\n return y, y\n\n def initial_state(self, batch_size: int):\n return jnp.zeros((batch_size, self.dout))\n\ncell = RNNCell(8, 16, nnx.Rngs(params=0, recurrent_dropout=1))\n```\n\nNext, use `nnx.scan` over an `unroll` function to implement the `rnn_forward` operation:\n- The key ingredient of recurrent dropout is to apply the same dropout mask across all time steps. Therefore, to achieve this you will pass `nnx.StateAxes` to `nnx.scan`'s `in_axes`, specifying that the `cell`'s `recurrent_dropout` PRNG stream will be broadcast, and the rest of the `RNNCell`'s state will be carried over.\n- Also, the hidden state `h` will be the `nnx.scan`'s `Carry` variable, and the sequence `x` will be `scan`ned over its axis `1`.\n\n```{code-cell} ipython3\n@nnx.jit\ndef rnn_forward(cell: RNNCell, x: jax.Array):\n h = cell.initial_state(batch_size=x.shape[0])\n\n # Broadcast the 'recurrent_dropout' PRNG state to have the same mask on every step.\n state_axes = nnx.StateAxes({'recurrent_dropout': None, ...: nnx.Carry})\n @nnx.scan(in_axes=(state_axes, nnx.Carry, 1), out_axes=(nnx.Carry, 1))\n def unroll(cell: RNNCell, h, x) -> tuple[jax.Array, jax.Array]:\n h, y = cell(h, x)\n return h, y\n\n h, y = unroll(cell, h, x)\n return y\n\nx = jnp.ones((4, 20, 8))\ny = rnn_forward(cell, x)\n\nprint(f'{y.shape = }')\nprint(f'{cell.count.value = }')\n```\n" }, { "path": "docs_nnx/guides/surgery.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"# Model surgery\\n\",\n \"\\n\",\n \"Model surgery is an act of making modifications on an existing neural network's building blocks and parameters, such as layer replacement, parameter or state manipulation, or even \\\"monkey patching\\\". In this guide, you will learn how to perform model surgery in Flax NNX using several real-world scenarios:\\n\",\n \"\\n\",\n \"* __Pythonic `nnx.Module` manipulation__: Using Pythonic ways to manipulate sub-`Module`s given a model.\\n\",\n \"\\n\",\n \"* __Manipulation of an abstract model or state__: A key trick for playing with `flax.nnx.Module`s and states without memory allocation.\\n\",\n \"\\n\",\n \"* __Checkpoint surgery from a raw state to model__: How to manipulate parameter states when they are incompatible with existing model code.\\n\",\n \"\\n\",\n \"* __Partial initialization__: How to initialize only a part of the model from scratch using a naive method or a memory-efficient method.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 1,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"from typing import *\\n\",\n \"from pprint import pprint\\n\",\n \"import functools\\n\",\n \"\\n\",\n \"import jax\\n\",\n \"from jax import lax, numpy as jnp, tree_util as jtu\\n\",\n \"\\n\",\n \"from jax.sharding import PartitionSpec, Mesh, NamedSharding\\n\",\n \"from jax.experimental import mesh_utils\\n\",\n \"import flax\\n\",\n \"from flax import nnx\\n\",\n \"import flax.traverse_util\\n\",\n \"import numpy as np\\n\",\n \"import orbax.checkpoint as orbax\\n\",\n \"\\n\",\n \"key = jax.random.key(0)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 2,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"class TwoLayerMLP(nnx.Module):\\n\",\n \" def __init__(self, dim, rngs: nnx.Rngs):\\n\",\n \" self.linear1 = nnx.Linear(dim, dim, rngs=rngs)\\n\",\n \" self.linear2 = nnx.Linear(dim, dim, rngs=rngs)\\n\",\n \"\\n\",\n \" def __call__(self, x):\\n\",\n \" x = self.linear1(x)\\n\",\n \" return self.linear2(x)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Pythonic `nnx.Module` manipulation\\n\",\n \"\\n\",\n \"It is easier to perform model surgery when:\\n\",\n \"\\n\",\n \"1) You already have a fully fleshed-out model loaded with correct parameters; and\\n\",\n \"2) You don't intend to change your model definition code.\\n\",\n \"\\n\",\n \"You can perform a variety of Pythonic operations on its sub-`Module`s, such as sub-`Module` swapping, `Module` sharing, variable sharing, and monkey-patching:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 3,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"model = TwoLayerMLP(4, rngs=nnx.Rngs(0))\\n\",\n \"x = jax.random.normal(jax.random.key(42), (3, 4))\\n\",\n \"np.testing.assert_allclose(model(x), model.linear2(model.linear1(x)))\\n\",\n \"\\n\",\n \"# Sub-`Module` swapping.\\n\",\n \"original1, original2 = model.linear1, model.linear2\\n\",\n \"model.linear1, model.linear2 = model.linear2, model.linear1\\n\",\n \"np.testing.assert_allclose(model(x), original1(original2(x)))\\n\",\n \"\\n\",\n \"# `Module` sharing (tying all weights together).\\n\",\n \"model = TwoLayerMLP(4, rngs=nnx.Rngs(0))\\n\",\n \"model.linear2 = model.linear1\\n\",\n \"assert not hasattr(nnx.state(model), 'linear2')\\n\",\n \"np.testing.assert_allclose(model(x), model.linear1(model.linear1(x)))\\n\",\n \"\\n\",\n \"# Variable sharing (weight-tying).\\n\",\n \"model = TwoLayerMLP(4, rngs=nnx.Rngs(0))\\n\",\n \"model.linear1.kernel = model.linear2.kernel # the bias parameter is kept separate\\n\",\n \"assert 'linear2' in nnx.state(model)\\n\",\n \"assert 'bias' in nnx.state(model)['linear2']\\n\",\n \"assert not hasattr(nnx.state(model)['linear2'], 'kernel')\\n\",\n \"\\n\",\n \"# Monkey-patching.\\n\",\n \"model = TwoLayerMLP(4, rngs=nnx.Rngs(0))\\n\",\n \"def awesome_layer(x): return x\\n\",\n \"model.linear2 = awesome_layer\\n\",\n \"np.testing.assert_allclose(model(x), model.linear1(x))\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Creating an abstract model or state without memory allocation\\n\",\n \"\\n\",\n \"To do more complex model surgery, the key technique you can use is creating and manipulating an abstract model or state without allocating any real parameter data. This makes trial iteration faster and removes any concern on memory constraints.\\n\",\n \"\\n\",\n \"To create an abstract model:\\n\",\n \"\\n\",\n \"* Create a function that returns a valid Flax NNX model; and\\n\",\n \"* Run `nnx.eval_shape` (not `jax.eval_shape`) upon it.\\n\",\n \"\\n\",\n \"Now you can use `nnx.split` as usual to get its abstract state. Note that all fields that should be `jax.Array`s in a real model are now of an abstract `jax.ShapeDtypeStruct` type with only shape/dtype/sharding information.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 4,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"State({\\n\",\n \" 'linear1': {\\n\",\n \" 'bias': Param( # 4 (16 B)\\n\",\n \" value=ShapeDtypeStruct(shape=(4,), dtype=float32)\\n\",\n \" ),\\n\",\n \" 'kernel': Param( # 16 (64 B)\\n\",\n \" value=ShapeDtypeStruct(shape=(4, 4), dtype=float32)\\n\",\n \" )\\n\",\n \" },\\n\",\n \" 'linear2': {\\n\",\n \" 'bias': Param( # 4 (16 B)\\n\",\n \" value=ShapeDtypeStruct(shape=(4,), dtype=float32)\\n\",\n \" ),\\n\",\n \" 'kernel': Param( # 16 (64 B)\\n\",\n \" value=ShapeDtypeStruct(shape=(4, 4), dtype=float32)\\n\",\n \" )\\n\",\n \" }\\n\",\n \"})\\n\"\n ]\n }\n ],\n \"source\": [\n \"abs_model = nnx.eval_shape(lambda: TwoLayerMLP(4, rngs=nnx.Rngs(0)))\\n\",\n \"gdef, abs_state = nnx.split(abs_model)\\n\",\n \"pprint(abs_state)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"When you fill every `nnx.Variable` pytree leaf's `value` attributes with real `jax.Array`s, the abstract model becomes equivalent to a real model.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 5,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"model = TwoLayerMLP(4, rngs=nnx.Rngs(0))\\n\",\n \"abs_state['linear1']['kernel'].value = model.linear1.kernel.value\\n\",\n \"abs_state['linear1']['bias'].value = model.linear1.bias.value\\n\",\n \"abs_state['linear2']['kernel'].value = model.linear2.kernel.value\\n\",\n \"abs_state['linear2']['bias'].value = model.linear2.bias.value\\n\",\n \"nnx.update(abs_model, abs_state)\\n\",\n \"np.testing.assert_allclose(abs_model(x), model(x)) # They are equivalent now!\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Checkpoint surgery\\n\",\n \"\\n\",\n \"With the abstract state technique in hand, you can perform arbitrary manipulation on any checkpoint - or runtime parameter pytree - to make them fit with your given model code, and then call `nnx.update` to merge them.\\n\",\n \"\\n\",\n \"This can be helpful if you are trying to significantly change the model code - for example, when migrating from Flax Linen to Flax NNX - and old weights are no longer naturally compatible.\\n\",\n \"\\n\",\n \"Let's run a simple example here:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 6,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# Save a version of model into a checkpoint\\n\",\n \"checkpointer = orbax.PyTreeCheckpointer()\\n\",\n \"old_model = TwoLayerMLP(4, rngs=nnx.Rngs(0))\\n\",\n \"checkpointer.save(f'/tmp/nnx-surgery-state', nnx.state(model), force=True)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"In this new model, the sub-`Module`s are renamed from `linear(1|2)` to `layer(1|2)`. Since the pytree structure has changed, it is impossible to directly load the old checkpoint with the new model state structure:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 7,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"This will throw error: : User-provided restore item and on-disk value metadata tree structures do not match: {'layer1': Diff(lhs={'bias': {'value': ShapeDtypeStruct(shape=(4,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(4, 4), dtype=float32)}}, rhs=None), 'layer2': Diff(lhs={'bias': {'value': ShapeDtypeStruct(shape=(4,), dtype=float32)}, 'kernel': {'value': ShapeDtypeStruct(shape=(4, 4), dtype=float32)}}, rhs=None), 'linear1': Diff(lhs=None, rhs={'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(4,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(4, 4))}}), 'linear2': Diff(lhs=None, rhs={'bias': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(4,))}, 'kernel': {'value': ValueMetadataEntry(value_type='jax.Array', skip_deserialize=False, write_shape=(4, 4))}})}\\n\"\n ]\n }\n ],\n \"source\": [\n \"class ModifiedTwoLayerMLP(nnx.Module):\\n\",\n \" def __init__(self, dim, rngs: nnx.Rngs):\\n\",\n \" self.layer1 = nnx.Linear(dim, dim, rngs=rngs) # no longer linear1!\\n\",\n \" self.layer2 = nnx.Linear(dim, dim, rngs=rngs)\\n\",\n \"\\n\",\n \" def __call__(self, x):\\n\",\n \" x = self.layer1(x)\\n\",\n \" return self.layer2(x)\\n\",\n \"\\n\",\n \"abs_model = nnx.eval_shape(lambda: ModifiedTwoLayerMLP(4, rngs=nnx.Rngs(0)))\\n\",\n \"try:\\n\",\n \" with_item = checkpointer.restore('/tmp/nnx-surgery-state', item=nnx.state(abs_model))\\n\",\n \" print(with_item)\\n\",\n \"except Exception as e:\\n\",\n \" print(f'This will throw error: {type(e)}: {e}')\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"However, you can load the parameter pytree as a raw dictionary, perform the renames, and generate a new state that is guaranteed to be compatible with your new model definition.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 8,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"{'linear1': {'bias': {'value': Array([0., 0., 0., 0.], dtype=float32)},\\n\",\n \" 'kernel': {'value': Array([[ 0.5350889 , -0.48486355, -0.4022262 , -0.61925626],\\n\",\n \" [-0.46665004, 0.31773907, 0.38944173, -0.54608804],\\n\",\n \" [ 0.84378934, -0.93099 , -0.67658 , 0.0724705 ],\\n\",\n \" [-0.6101737 , 0.12972134, 0.877074 , 0.27292168]], dtype=float32)}},\\n\",\n \" 'linear2': {'bias': {'value': Array([0., 0., 0., 0.], dtype=float32)},\\n\",\n \" 'kernel': {'value': Array([[ 0.67979455, 0.7079946 , -0.22166717, -0.4147039 ],\\n\",\n \" [ 0.20622818, 0.01024843, 0.31011865, -0.40491563],\\n\",\n \" [ 0.12478007, -0.7697264 , -0.48899388, 0.8853114 ],\\n\",\n \" [-0.5123713 , -0.23335123, 0.4374407 , 0.63321066]], dtype=float32)}}}\\n\"\n ]\n },\n {\n \"name\": \"stderr\",\n \"output_type\": \"stream\",\n \"text\": [\n \"/Users/cgarciae/repos/flax/.venv/lib/python3.11/site-packages/orbax/checkpoint/_src/serialization/type_handlers.py:1251: UserWarning: Sharding info not provided when restoring. Populating sharding info from sharding file. Please note restoration time will be slightly increased due to reading from file. Note also that this option is unsafe when restoring on a different topology than the checkpoint was saved with.\\n\",\n \" warnings.warn(\\n\"\n ]\n }\n ],\n \"source\": [\n \"def process_raw_dict(raw_state_dict):\\n\",\n \" flattened = nnx.traversals.flatten_mapping(raw_state_dict)\\n\",\n \" # Cut the '.value' postfix on every leaf path.\\n\",\n \" flattened = {(path[:-1] if path[-1] == 'value' else path): value\\n\",\n \" for path, value in flattened.items()}\\n\",\n \" return nnx.traversals.unflatten_mapping(flattened)\\n\",\n \"\\n\",\n \"# Make your local change on the checkpoint dictionary.\\n\",\n \"raw_dict = checkpointer.restore('/tmp/nnx-surgery-state')\\n\",\n \"pprint(raw_dict)\\n\",\n \"raw_dict['layer1'] = raw_dict.pop('linear1')\\n\",\n \"raw_dict['layer2'] = raw_dict.pop('linear2')\\n\",\n \"\\n\",\n \"# Fit it into the model state.\\n\",\n \"abs_model = nnx.eval_shape(lambda: ModifiedTwoLayerMLP(4, rngs=nnx.Rngs(0)))\\n\",\n \"graph_def, state = nnx.split(abs_model)\\n\",\n \"nnx.replace_by_pure_dict(state, process_raw_dict(raw_dict))\\n\",\n \"restored_model = nnx.merge(graph_def, state)\\n\",\n \"\\n\",\n \"np.testing.assert_allclose(restored_model(jnp.ones((3, 4))), old_model(jnp.ones((3, 4))))\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Partial initialization\\n\",\n \"\\n\",\n \"In some cases - such as with LoRA (Low-Rank Adaption) - you may want to randomly-initialize only *part of* your model parameters. This can be achieved through:\\n\",\n \"\\n\",\n \"- Naive partial initialization; or\\n\",\n \"- Memory-efficient partial initialization.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Naive partial initialization\\n\",\n \"\\n\",\n \"To do naive partial initialization, you can just initialize the whole model, then swap the pre-trained parameters in. However, this approach may allocate additional memory midway if your modification requires re-creating module parameters that you will later discard. Below is an example of this.\\n\",\n \"\\n\",\n \"> **Note:** You can use `jax.live_arrays()` to check all the arrays live in memory at any given time. This call can be “messed up” when you run a single Jupyter notebook cell multiple times (due to garbage-collection of old Python variables). However, restarting the Python kernel in the notebook and running the code from scratch will always yield the same output.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 9,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"Number of jax arrays in memory at start: 38\\n\",\n \"Number of jax arrays in memory midway: 42 (4 new created in LoRALinear - kernel, bias, lora_a & lora_b)\\n\",\n \"Number of jax arrays in memory at end: 40 (2 discarded - only lora_a & lora_b are used in model)\\n\"\n ]\n }\n ],\n \"source\": [\n \"# Some pretrained model state\\n\",\n \"old_state = nnx.state(TwoLayerMLP(4, rngs=nnx.Rngs(0)))\\n\",\n \"\\n\",\n \"simple_model = nnx.eval_shape(lambda: TwoLayerMLP(4, rngs=nnx.Rngs(42)))\\n\",\n \"print(f'Number of jax arrays in memory at start: {len(jax.live_arrays())}')\\n\",\n \"# In this line, extra kernel and bias is created inside the new LoRALinear!\\n\",\n \"# They are wasted, because you are going to use the kernel and bias in `old_state` anyway.\\n\",\n \"simple_model.linear1 = nnx.LoRALinear(4, 4, lora_rank=3, rngs=nnx.Rngs(42))\\n\",\n \"print(f'Number of jax arrays in memory midway: {len(jax.live_arrays())}'\\n\",\n \" ' (4 new created in LoRALinear - kernel, bias, lora_a & lora_b)')\\n\",\n \"nnx.update(simple_model, old_state)\\n\",\n \"print(f'Number of jax arrays in memory at end: {len(jax.live_arrays())}'\\n\",\n \" ' (2 discarded - only lora_a & lora_b are used in model)')\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Memory-efficient partial initialization\\n\",\n \"\\n\",\n \"To do memory-efficient partial initialization, use `nnx.jit`'s efficiently compiled code to make sure only the state parameters you need are initialized:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 10,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"Number of JAX Arrays in memory at start: 44\\n\",\n \"Number of JAX Arrays in memory at end: 50 (2 new created - lora_a and lora_b)\\n\"\n ]\n }\n ],\n \"source\": [\n \"# Some pretrained model state\\n\",\n \"old_state = nnx.state(TwoLayerMLP(4, rngs=nnx.Rngs(0)))\\n\",\n \"\\n\",\n \"# Use `nnx.jit` (which wraps `jax.jit`) to automatically skip unused arrays - memory efficient!\\n\",\n \"@nnx.jit(donate_argnums=0)\\n\",\n \"def partial_init(old_state, rngs):\\n\",\n \" model = TwoLayerMLP(4, rngs=rngs)\\n\",\n \" # Create a new state.\\n\",\n \" model.linear1 = nnx.LoRALinear(4, 4, lora_rank=3, rngs=rngs)\\n\",\n \" # Add the existing state.\\n\",\n \" nnx.update(model, old_state)\\n\",\n \" return model\\n\",\n \"\\n\",\n \"print(f'Number of JAX Arrays in memory at start: {len(jax.live_arrays())}')\\n\",\n \"# Note that `old_state` will be deleted after this `partial_init` call.\\n\",\n \"good_model = partial_init(old_state, nnx.Rngs(42))\\n\",\n \"print(f'Number of JAX Arrays in memory at end: {len(jax.live_arrays())}'\\n\",\n \" ' (2 new created - lora_a and lora_b)')\"\n ]\n }\n ],\n \"metadata\": {\n \"jupytext\": {\n \"formats\": \"ipynb,md:myst\"\n },\n \"language_info\": {\n \"codemirror_mode\": {\n \"name\": \"ipython\",\n \"version\": 3\n },\n \"file_extension\": \".py\",\n \"mimetype\": \"text/x-python\",\n \"name\": \"python\",\n \"nbconvert_exporter\": \"python\",\n \"pygments_lexer\": \"ipython3\",\n \"version\": \"3.11.9\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 2\n}\n" }, { "path": "docs_nnx/guides/surgery.md", "content": "---\njupytext:\n formats: ipynb,md:myst\n text_representation:\n extension: .md\n format_name: myst\n format_version: 0.13\n jupytext_version: 1.13.8\n---\n\n# Model surgery\n\nModel surgery is an act of making modifications on an existing neural network's building blocks and parameters, such as layer replacement, parameter or state manipulation, or even \"monkey patching\". In this guide, you will learn how to perform model surgery in Flax NNX using several real-world scenarios:\n\n* __Pythonic `nnx.Module` manipulation__: Using Pythonic ways to manipulate sub-`Module`s given a model.\n\n* __Manipulation of an abstract model or state__: A key trick for playing with `flax.nnx.Module`s and states without memory allocation.\n\n* __Checkpoint surgery from a raw state to model__: How to manipulate parameter states when they are incompatible with existing model code.\n\n* __Partial initialization__: How to initialize only a part of the model from scratch using a naive method or a memory-efficient method.\n\n```{code-cell} ipython3\nfrom typing import *\nfrom pprint import pprint\nimport functools\n\nimport jax\nfrom jax import lax, numpy as jnp, tree_util as jtu\n\nfrom jax.sharding import PartitionSpec, Mesh, NamedSharding\nfrom jax.experimental import mesh_utils\nimport flax\nfrom flax import nnx\nimport flax.traverse_util\nimport numpy as np\nimport orbax.checkpoint as orbax\n\nkey = jax.random.key(0)\n```\n\n```{code-cell} ipython3\nclass TwoLayerMLP(nnx.Module):\n def __init__(self, dim, rngs: nnx.Rngs):\n self.linear1 = nnx.Linear(dim, dim, rngs=rngs)\n self.linear2 = nnx.Linear(dim, dim, rngs=rngs)\n\n def __call__(self, x):\n x = self.linear1(x)\n return self.linear2(x)\n```\n\n## Pythonic `nnx.Module` manipulation\n\nIt is easier to perform model surgery when:\n\n1) You already have a fully fleshed-out model loaded with correct parameters; and\n2) You don't intend to change your model definition code.\n\nYou can perform a variety of Pythonic operations on its sub-`Module`s, such as sub-`Module` swapping, `Module` sharing, variable sharing, and monkey-patching:\n\n```{code-cell} ipython3\nmodel = TwoLayerMLP(4, rngs=nnx.Rngs(0))\nx = jax.random.normal(jax.random.key(42), (3, 4))\nnp.testing.assert_allclose(model(x), model.linear2(model.linear1(x)))\n\n# Sub-`Module` swapping.\noriginal1, original2 = model.linear1, model.linear2\nmodel.linear1, model.linear2 = model.linear2, model.linear1\nnp.testing.assert_allclose(model(x), original1(original2(x)))\n\n# `Module` sharing (tying all weights together).\nmodel = TwoLayerMLP(4, rngs=nnx.Rngs(0))\nmodel.linear2 = model.linear1\nassert not hasattr(nnx.state(model), 'linear2')\nnp.testing.assert_allclose(model(x), model.linear1(model.linear1(x)))\n\n# Variable sharing (weight-tying).\nmodel = TwoLayerMLP(4, rngs=nnx.Rngs(0))\nmodel.linear1.kernel = model.linear2.kernel # the bias parameter is kept separate\nassert 'linear2' in nnx.state(model)\nassert 'bias' in nnx.state(model)['linear2']\nassert not hasattr(nnx.state(model)['linear2'], 'kernel')\n\n# Monkey-patching.\nmodel = TwoLayerMLP(4, rngs=nnx.Rngs(0))\ndef awesome_layer(x): return x\nmodel.linear2 = awesome_layer\nnp.testing.assert_allclose(model(x), model.linear1(x))\n```\n\n## Creating an abstract model or state without memory allocation\n\nTo do more complex model surgery, the key technique you can use is creating and manipulating an abstract model or state without allocating any real parameter data. This makes trial iteration faster and removes any concern on memory constraints.\n\nTo create an abstract model:\n\n* Create a function that returns a valid Flax NNX model; and\n* Run `nnx.eval_shape` (not `jax.eval_shape`) upon it.\n\nNow you can use `nnx.split` as usual to get its abstract state. Note that all fields that should be `jax.Array`s in a real model are now of an abstract `jax.ShapeDtypeStruct` type with only shape/dtype/sharding information.\n\n```{code-cell} ipython3\nabs_model = nnx.eval_shape(lambda: TwoLayerMLP(4, rngs=nnx.Rngs(0)))\ngdef, abs_state = nnx.split(abs_model)\npprint(abs_state)\n```\n\nWhen you fill every `nnx.Variable` pytree leaf's `value` attributes with real `jax.Array`s, the abstract model becomes equivalent to a real model.\n\n```{code-cell} ipython3\nmodel = TwoLayerMLP(4, rngs=nnx.Rngs(0))\nabs_state['linear1']['kernel'].value = model.linear1.kernel.value\nabs_state['linear1']['bias'].value = model.linear1.bias.value\nabs_state['linear2']['kernel'].value = model.linear2.kernel.value\nabs_state['linear2']['bias'].value = model.linear2.bias.value\nnnx.update(abs_model, abs_state)\nnp.testing.assert_allclose(abs_model(x), model(x)) # They are equivalent now!\n```\n\n## Checkpoint surgery\n\nWith the abstract state technique in hand, you can perform arbitrary manipulation on any checkpoint - or runtime parameter pytree - to make them fit with your given model code, and then call `nnx.update` to merge them.\n\nThis can be helpful if you are trying to significantly change the model code - for example, when migrating from Flax Linen to Flax NNX - and old weights are no longer naturally compatible.\n\nLet's run a simple example here:\n\n```{code-cell} ipython3\n# Save a version of model into a checkpoint\ncheckpointer = orbax.PyTreeCheckpointer()\nold_model = TwoLayerMLP(4, rngs=nnx.Rngs(0))\ncheckpointer.save(f'/tmp/nnx-surgery-state', nnx.state(model), force=True)\n```\n\nIn this new model, the sub-`Module`s are renamed from `linear(1|2)` to `layer(1|2)`. Since the pytree structure has changed, it is impossible to directly load the old checkpoint with the new model state structure:\n\n```{code-cell} ipython3\nclass ModifiedTwoLayerMLP(nnx.Module):\n def __init__(self, dim, rngs: nnx.Rngs):\n self.layer1 = nnx.Linear(dim, dim, rngs=rngs) # no longer linear1!\n self.layer2 = nnx.Linear(dim, dim, rngs=rngs)\n\n def __call__(self, x):\n x = self.layer1(x)\n return self.layer2(x)\n\nabs_model = nnx.eval_shape(lambda: ModifiedTwoLayerMLP(4, rngs=nnx.Rngs(0)))\ntry:\n with_item = checkpointer.restore('/tmp/nnx-surgery-state', item=nnx.state(abs_model))\n print(with_item)\nexcept Exception as e:\n print(f'This will throw error: {type(e)}: {e}')\n```\n\nHowever, you can load the parameter pytree as a raw dictionary, perform the renames, and generate a new state that is guaranteed to be compatible with your new model definition.\n\n```{code-cell} ipython3\ndef process_raw_dict(raw_state_dict):\n flattened = nnx.traversals.flatten_mapping(raw_state_dict)\n # Cut the '.value' postfix on every leaf path.\n flattened = {(path[:-1] if path[-1] == 'value' else path): value\n for path, value in flattened.items()}\n return nnx.traversals.unflatten_mapping(flattened)\n\n# Make your local change on the checkpoint dictionary.\nraw_dict = checkpointer.restore('/tmp/nnx-surgery-state')\npprint(raw_dict)\nraw_dict['layer1'] = raw_dict.pop('linear1')\nraw_dict['layer2'] = raw_dict.pop('linear2')\n\n# Fit it into the model state.\nabs_model = nnx.eval_shape(lambda: ModifiedTwoLayerMLP(4, rngs=nnx.Rngs(0)))\ngraph_def, state = nnx.split(abs_model)\nnnx.replace_by_pure_dict(state, process_raw_dict(raw_dict))\nrestored_model = nnx.merge(graph_def, state)\n\nnp.testing.assert_allclose(restored_model(jnp.ones((3, 4))), old_model(jnp.ones((3, 4))))\n```\n\n## Partial initialization\n\nIn some cases - such as with LoRA (Low-Rank Adaption) - you may want to randomly-initialize only *part of* your model parameters. This can be achieved through:\n\n- Naive partial initialization; or\n- Memory-efficient partial initialization.\n\n+++\n\n### Naive partial initialization\n\nTo do naive partial initialization, you can just initialize the whole model, then swap the pre-trained parameters in. However, this approach may allocate additional memory midway if your modification requires re-creating module parameters that you will later discard. Below is an example of this.\n\n> **Note:** You can use `jax.live_arrays()` to check all the arrays live in memory at any given time. This call can be “messed up” when you run a single Jupyter notebook cell multiple times (due to garbage-collection of old Python variables). However, restarting the Python kernel in the notebook and running the code from scratch will always yield the same output.\n\n```{code-cell} ipython3\n# Some pretrained model state\nold_state = nnx.state(TwoLayerMLP(4, rngs=nnx.Rngs(0)))\n\nsimple_model = nnx.eval_shape(lambda: TwoLayerMLP(4, rngs=nnx.Rngs(42)))\nprint(f'Number of jax arrays in memory at start: {len(jax.live_arrays())}')\n# In this line, extra kernel and bias is created inside the new LoRALinear!\n# They are wasted, because you are going to use the kernel and bias in `old_state` anyway.\nsimple_model.linear1 = nnx.LoRALinear(4, 4, lora_rank=3, rngs=nnx.Rngs(42))\nprint(f'Number of jax arrays in memory midway: {len(jax.live_arrays())}'\n ' (4 new created in LoRALinear - kernel, bias, lora_a & lora_b)')\nnnx.update(simple_model, old_state)\nprint(f'Number of jax arrays in memory at end: {len(jax.live_arrays())}'\n ' (2 discarded - only lora_a & lora_b are used in model)')\n```\n\n### Memory-efficient partial initialization\n\nTo do memory-efficient partial initialization, use `nnx.jit`'s efficiently compiled code to make sure only the state parameters you need are initialized:\n\n```{code-cell} ipython3\n# Some pretrained model state\nold_state = nnx.state(TwoLayerMLP(4, rngs=nnx.Rngs(0)))\n\n# Use `nnx.jit` (which wraps `jax.jit`) to automatically skip unused arrays - memory efficient!\n@nnx.jit(donate_argnums=0)\ndef partial_init(old_state, rngs):\n model = TwoLayerMLP(4, rngs=rngs)\n # Create a new state.\n model.linear1 = nnx.LoRALinear(4, 4, lora_rank=3, rngs=rngs)\n # Add the existing state.\n nnx.update(model, old_state)\n return model\n\nprint(f'Number of JAX Arrays in memory at start: {len(jax.live_arrays())}')\n# Note that `old_state` will be deleted after this `partial_init` call.\ngood_model = partial_init(old_state, nnx.Rngs(42))\nprint(f'Number of JAX Arrays in memory at end: {len(jax.live_arrays())}'\n ' (2 new created - lora_a and lora_b)')\n```\n" }, { "path": "docs_nnx/guides/tiny_nnx.ipynb", "content": "{\n \"cells\": [\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"# Tiny NNX\\n\",\n \"[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/cgarciae/nnx/blob/main/docs/tiny_nnx.ipynb)\\n\",\n \"\\n\",\n \"A pedagogical implementation of NNX's core APIs.\\n\",\n \"\\n\",\n \"## Core APIs\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 1,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"import dataclasses\\n\",\n \"import hashlib\\n\",\n \"import typing as tp\\n\",\n \"\\n\",\n \"import jax\\n\",\n \"import jax.numpy as jnp\\n\",\n \"from jax import random\\n\",\n \"\\n\",\n \"A = tp.TypeVar(\\\"A\\\")\\n\",\n \"M = tp.TypeVar(\\\"M\\\", bound=\\\"Module\\\")\\n\",\n \"Sharding = tp.Tuple[tp.Optional[str], ...]\\n\",\n \"Array = jax.Array\\n\",\n \"\\n\",\n \"\\n\",\n \"class Variable(tp.Generic[A]):\\n\",\n \"\\n\",\n \" def __init__(\\n\",\n \" self,\\n\",\n \" value: A,\\n\",\n \" *,\\n\",\n \" sharding: tp.Optional[Sharding] = None,\\n\",\n \" ):\\n\",\n \" self.value = value\\n\",\n \" self.sharding = sharding\\n\",\n \"\\n\",\n \" def __repr__(self) -> str:\\n\",\n \" return (\\n\",\n \" f\\\"{type(self).__name__}(value={self.value}, sharding={self.sharding})\\\"\\n\",\n \" )\\n\",\n \"\\n\",\n \" def __init_subclass__(cls):\\n\",\n \" super().__init_subclass__()\\n\",\n \" jax.tree_util.register_pytree_node(\\n\",\n \" cls,\\n\",\n \" lambda x: ((x.value,), (x.sharding,)),\\n\",\n \" lambda metadata, value: cls(value[0], sharding=metadata[0]),\\n\",\n \" )\\n\",\n \"\\n\",\n \"\\n\",\n \"class State(dict[str, Variable[tp.Any]]):\\n\",\n \"\\n\",\n \" def extract(self, variable_type: tp.Type[Variable]) -> \\\"State\\\":\\n\",\n \" return State(\\n\",\n \" {\\n\",\n \" path: variable\\n\",\n \" for path, variable in self.items()\\n\",\n \" if isinstance(variable, variable_type)\\n\",\n \" }\\n\",\n \" )\\n\",\n \"\\n\",\n \" def __repr__(self) -> str:\\n\",\n \" elems = \\\",\\\\n \\\".join(\\n\",\n \" f\\\"'{path}': {variable}\\\".replace(\\\"\\\\n\\\", \\\"\\\\n \\\")\\n\",\n \" for path, variable in self.items()\\n\",\n \" )\\n\",\n \" return f\\\"State({{\\\\n {elems}\\\\n}})\\\"\\n\",\n \"\\n\",\n \"\\n\",\n \"jax.tree_util.register_pytree_node(\\n\",\n \" State,\\n\",\n \" # in reality, values and paths should be sorted by path\\n\",\n \" lambda x: (tuple(x.values()), tuple(x.keys())),\\n\",\n \" lambda paths, values: State(dict(zip(paths, values))),\\n\",\n \")\\n\",\n \"\\n\",\n \"\\n\",\n \"@dataclasses.dataclass\\n\",\n \"class GraphDef(tp.Generic[M]):\\n\",\n \" type: tp.Type[M]\\n\",\n \" index: int\\n\",\n \" submodules: dict[str, tp.Union[\\\"GraphDef[Module]\\\", int]]\\n\",\n \" static_fields: dict[str, tp.Any]\\n\",\n \"\\n\",\n \" def merge(self, state: State) -> M:\\n\",\n \" module = GraphDef._build_module_recursive(self, {})\\n\",\n \" module.update(state)\\n\",\n \" return module\\n\",\n \"\\n\",\n \" @staticmethod\\n\",\n \" def _build_module_recursive(\\n\",\n \" graphdef: tp.Union[\\\"GraphDef[M]\\\", int],\\n\",\n \" index_to_module: dict[int, \\\"Module\\\"],\\n\",\n \" ) -> M:\\n\",\n \" if isinstance(graphdef, int):\\n\",\n \" return index_to_module[graphdef] # type: ignore\\n\",\n \"\\n\",\n \" assert graphdef.index not in index_to_module\\n\",\n \"\\n\",\n \" # add a dummy module to the index to avoid infinite recursion\\n\",\n \" module = object.__new__(graphdef.type)\\n\",\n \" index_to_module[graphdef.index] = module\\n\",\n \"\\n\",\n \" submodules = {\\n\",\n \" name: GraphDef._build_module_recursive(submodule, index_to_module)\\n\",\n \" for name, submodule in graphdef.submodules.items()\\n\",\n \" }\\n\",\n \" vars(module).update(graphdef.static_fields)\\n\",\n \" vars(module).update(submodules)\\n\",\n \" return module\\n\",\n \"\\n\",\n \" def apply(\\n\",\n \" self, state: State\\n\",\n \" ) -> tp.Callable[..., tuple[tp.Any, tuple[State, \\\"GraphDef[M]\\\"]]]:\\n\",\n \" def _apply(*args, **kwargs):\\n\",\n \" module = self.merge(state)\\n\",\n \" out = module(*args, **kwargs) # type: ignore\\n\",\n \" return out, module.split()\\n\",\n \"\\n\",\n \" return _apply\\n\",\n \"\\n\",\n \"\\n\",\n \"class Module:\\n\",\n \"\\n\",\n \" def split(self: M) -> tp.Tuple[State, GraphDef[M]]:\\n\",\n \" state = State()\\n\",\n \" graphdef = Module._partition_recursive(\\n\",\n \" module=self, module_id_to_index={}, path_parts=(), state=state\\n\",\n \" )\\n\",\n \" assert isinstance(graphdef, GraphDef)\\n\",\n \" return state, graphdef\\n\",\n \"\\n\",\n \" @staticmethod\\n\",\n \" def _partition_recursive(\\n\",\n \" module: M,\\n\",\n \" module_id_to_index: dict[int, int],\\n\",\n \" path_parts: tp.Tuple[str, ...],\\n\",\n \" state: State,\\n\",\n \" ) -> tp.Union[GraphDef[M], int]:\\n\",\n \" if id(module) in module_id_to_index:\\n\",\n \" return module_id_to_index[id(module)]\\n\",\n \"\\n\",\n \" index = len(module_id_to_index)\\n\",\n \" module_id_to_index[id(module)] = index\\n\",\n \"\\n\",\n \" submodules = {}\\n\",\n \" static_fields = {}\\n\",\n \"\\n\",\n \" # iterate fields sorted by name to ensure deterministic order\\n\",\n \" for name, value in sorted(vars(module).items(), key=lambda x: x[0]):\\n\",\n \" value_path = (*path_parts, name)\\n\",\n \" # if value is a Module, recurse\\n\",\n \" if isinstance(value, Module):\\n\",\n \" submoduledef = Module._partition_recursive(\\n\",\n \" value, module_id_to_index, value_path, state\\n\",\n \" )\\n\",\n \" submodules[name] = submoduledef\\n\",\n \" # if value is a Variable, add to state\\n\",\n \" elif isinstance(value, Variable):\\n\",\n \" state[\\\"/\\\".join(value_path)] = value\\n\",\n \" else: # otherwise, add to graphdef fields\\n\",\n \" static_fields[name] = value\\n\",\n \"\\n\",\n \" return GraphDef(\\n\",\n \" type=type(module),\\n\",\n \" index=index,\\n\",\n \" submodules=submodules,\\n\",\n \" static_fields=static_fields,\\n\",\n \" )\\n\",\n \"\\n\",\n \" def update(self, state: State) -> None:\\n\",\n \" for path, value in state.items():\\n\",\n \" path_parts = path.split(\\\"/\\\")\\n\",\n \" Module._set_value_at_path(self, path_parts, value)\\n\",\n \"\\n\",\n \" @staticmethod\\n\",\n \" def _set_value_at_path(\\n\",\n \" module: \\\"Module\\\", path_parts: tp.Sequence[str], value: Variable[tp.Any]\\n\",\n \" ) -> None:\\n\",\n \" if len(path_parts) == 1:\\n\",\n \" setattr(module, path_parts[0], value)\\n\",\n \" else:\\n\",\n \" Module._set_value_at_path(\\n\",\n \" getattr(module, path_parts[0]), path_parts[1:], value\\n\",\n \" )\\n\",\n \"\\n\",\n \"\\n\",\n \"@dataclasses.dataclass\\n\",\n \"class Rngs:\\n\",\n \" key: jax.Array\\n\",\n \" count: int = 0\\n\",\n \" count_path: tuple[int, ...] = ()\\n\",\n \"\\n\",\n \" def fork(self) -> \\\"Rngs\\\":\\n\",\n \" \\\"\\\"\\\"Forks the context, guaranteeing that all the random numbers generated\\n\",\n \" will be different from the ones generated in the original context. Fork is\\n\",\n \" used to create a new Rngs that can be passed to a JAX transform\\\"\\\"\\\"\\n\",\n \" count_path = self.count_path + (self.count,)\\n\",\n \" self.count += 1\\n\",\n \" return Rngs(self.key, count_path=count_path)\\n\",\n \"\\n\",\n \" def make_rng(self) -> jax.Array:\\n\",\n \" fold_data = self._stable_hash(self.count_path + (self.count,))\\n\",\n \" self.count += 1\\n\",\n \" return random.fold_in(self.key, fold_data) # type: ignore\\n\",\n \"\\n\",\n \" @staticmethod\\n\",\n \" def _stable_hash(data: tuple[int, ...]) -> int:\\n\",\n \" hash_str = \\\" \\\".join(str(x) for x in data)\\n\",\n \" _hash = hashlib.blake2s(hash_str.encode())\\n\",\n \" hash_bytes = _hash.digest()\\n\",\n \" # uint32 is represented as 4 bytes in big endian\\n\",\n \" return int.from_bytes(hash_bytes[:4], byteorder=\\\"big\\\")\\n\",\n \"\\n\",\n \"\\n\",\n \"# in the real NNX Rngs is not a pytree, instead\\n\",\n \"# it has a split/merge API similar to Module\\n\",\n \"# but for simplicity we use a pytree here\\n\",\n \"jax.tree_util.register_pytree_node(\\n\",\n \" Rngs,\\n\",\n \" lambda x: ((x.key,), (x.count, x.count_path)),\\n\",\n \" lambda metadata, value: Rngs(value[0], *metadata),\\n\",\n \")\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Basic Layers\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 2,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"class Param(Variable[A]):\\n\",\n \" pass\\n\",\n \"\\n\",\n \"\\n\",\n \"class BatchStat(Variable[A]):\\n\",\n \" pass\\n\",\n \"\\n\",\n \"\\n\",\n \"class Linear(Module):\\n\",\n \"\\n\",\n \" def __init__(self, din: int, dout: int, *, rngs: Rngs):\\n\",\n \" self.din = din\\n\",\n \" self.dout = dout\\n\",\n \" key = rngs.make_rng()\\n\",\n \" self.w = Param(random.uniform(key, (din, dout)))\\n\",\n \" self.b = Param(jnp.zeros((dout,)))\\n\",\n \"\\n\",\n \" def __call__(self, x: jax.Array) -> jax.Array:\\n\",\n \" return x @ self.w.value + self.b.value\\n\",\n \"\\n\",\n \"\\n\",\n \"class BatchNorm(Module):\\n\",\n \"\\n\",\n \" def __init__(self, din: int, mu: float = 0.95):\\n\",\n \" self.mu = mu\\n\",\n \" self.scale = Param(jax.numpy.ones((din,)))\\n\",\n \" self.bias = Param(jax.numpy.zeros((din,)))\\n\",\n \" self.mean = BatchStat(jax.numpy.zeros((din,)))\\n\",\n \" self.var = BatchStat(jax.numpy.ones((din,)))\\n\",\n \"\\n\",\n \" def __call__(self, x, train: bool) -> jax.Array:\\n\",\n \" if train:\\n\",\n \" axis = tuple(range(x.ndim - 1))\\n\",\n \" mean = jax.numpy.mean(x, axis=axis)\\n\",\n \" var = jax.numpy.var(x, axis=axis)\\n\",\n \" # ema update\\n\",\n \" self.mean.value = self.mu * self.mean.value + (1 - self.mu) * mean\\n\",\n \" self.var.value = self.mu * self.var.value + (1 - self.mu) * var\\n\",\n \" else:\\n\",\n \" mean, var = self.mean.value, self.var.value\\n\",\n \"\\n\",\n \" scale, bias = self.scale.value, self.bias.value\\n\",\n \" x = (x - mean) / jax.numpy.sqrt(var + 1e-5) * scale + bias\\n\",\n \" return x\\n\",\n \"\\n\",\n \"\\n\",\n \"class Dropout(Module):\\n\",\n \"\\n\",\n \" def __init__(self, rate: float):\\n\",\n \" self.rate = rate\\n\",\n \"\\n\",\n \" def __call__(self, x: jax.Array, *, train: bool, rngs: Rngs) -> jax.Array:\\n\",\n \" if train:\\n\",\n \" mask = random.bernoulli(rngs.make_rng(), (1 - self.rate), x.shape)\\n\",\n \" x = x * mask / (1 - self.rate)\\n\",\n \" return x\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Scan Over Layers Example\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 3,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"class Block(Module):\\n\",\n \"\\n\",\n \" def __init__(self, din: int, dout: int, *, rngs: Rngs):\\n\",\n \" self.linear = Linear(din, dout, rngs=rngs)\\n\",\n \" self.bn = BatchNorm(dout)\\n\",\n \" self.dropout = Dropout(0.1)\\n\",\n \"\\n\",\n \" def __call__(self, x: jax.Array, *, train: bool, rngs: Rngs) -> jax.Array:\\n\",\n \" x = self.linear(x)\\n\",\n \" x = self.bn(x, train=train)\\n\",\n \" x = jax.nn.gelu(x)\\n\",\n \" x = self.dropout(x, train=train, rngs=rngs)\\n\",\n \" return x\\n\",\n \"\\n\",\n \"\\n\",\n \"class ScanMLP(Module):\\n\",\n \"\\n\",\n \" def __init__(self, hidden_size: int, n_layers: int, *, rngs: Rngs):\\n\",\n \" self.n_layers = n_layers\\n\",\n \"\\n\",\n \" # lift init\\n\",\n \" key = random.split(rngs.make_rng(), n_layers - 1)\\n\",\n \" graphdef: GraphDef[Block] = None # type: ignore\\n\",\n \"\\n\",\n \" def init_fn(key):\\n\",\n \" nonlocal graphdef\\n\",\n \" state, graphdef = Block(\\n\",\n \" hidden_size, hidden_size, rngs=Rngs(key)\\n\",\n \" ).split()\\n\",\n \" return state\\n\",\n \"\\n\",\n \" state = jax.vmap(init_fn)(key)\\n\",\n \" self.layers = graphdef.merge(state)\\n\",\n \" self.linear = Linear(hidden_size, hidden_size, rngs=rngs)\\n\",\n \"\\n\",\n \" def __call__(self, x: jax.Array, *, train: bool, rngs: Rngs) -> jax.Array:\\n\",\n \" # lift call\\n\",\n \" key: jax.Array = random.split(rngs.make_rng(), self.n_layers - 1) # type: ignore\\n\",\n \" state, graphdef = self.layers.split()\\n\",\n \"\\n\",\n \" def scan_fn(x, inputs: tuple[jax.Array, State]):\\n\",\n \" key, state = inputs\\n\",\n \" x, (state, _) = graphdef.apply(state)(x, train=train, rngs=Rngs(key))\\n\",\n \" return x, state\\n\",\n \"\\n\",\n \" x, state = jax.lax.scan(scan_fn, x, (key, state))\\n\",\n \" self.layers.update(state)\\n\",\n \" x = self.linear(x)\\n\",\n \" return x\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 4,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"state = State({\\n\",\n \" 'layers/bn/bias': Param(value=(4, 10), sharding=None),\\n\",\n \" 'layers/bn/mean': BatchStat(value=(4, 10), sharding=None),\\n\",\n \" 'layers/bn/scale': Param(value=(4, 10), sharding=None),\\n\",\n \" 'layers/bn/var': BatchStat(value=(4, 10), sharding=None),\\n\",\n \" 'layers/linear/b': Param(value=(4, 10), sharding=None),\\n\",\n \" 'layers/linear/w': Param(value=(4, 10, 10), sharding=None),\\n\",\n \" 'linear/b': Param(value=(10,), sharding=None),\\n\",\n \" 'linear/w': Param(value=(10, 10), sharding=None)\\n\",\n \"})\\n\",\n \"graphdef = GraphDef(type=, index=0, submodules={'layers': GraphDef(type=, index=1, submodules={'bn': GraphDef(type=, index=2, submodules={}, static_fields={'mu': 0.95}), 'dropout': GraphDef(type=, index=3, submodules={}, static_fields={'rate': 0.1}), 'linear': GraphDef(type=, index=4, submodules={}, static_fields={'din': 10, 'dout': 10})}, static_fields={}), 'linear': GraphDef(type=, index=5, submodules={}, static_fields={'din': 10, 'dout': 10})}, static_fields={'n_layers': 5})\\n\"\n ]\n }\n ],\n \"source\": [\n \"module = ScanMLP(hidden_size=10, n_layers=5, rngs=Rngs(random.key(0)))\\n\",\n \"x = jax.random.normal(random.key(0), (2, 10))\\n\",\n \"y = module(x, train=True, rngs=Rngs(random.key(1)))\\n\",\n \"\\n\",\n \"state, graphdef = module.split()\\n\",\n \"print(\\\"state =\\\", jax.tree.map(jnp.shape, state))\\n\",\n \"print(\\\"graphdef =\\\", graphdef)\"\n ]\n },\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Filtering State\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 5,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"params = State({\\n\",\n \" 'layers/bn/bias': Param(value=(4, 10), sharding=None),\\n\",\n \" 'layers/bn/scale': Param(value=(4, 10), sharding=None),\\n\",\n \" 'layers/linear/b': Param(value=(4, 10), sharding=None),\\n\",\n \" 'layers/linear/w': Param(value=(4, 10, 10), sharding=None),\\n\",\n \" 'linear/b': Param(value=(10,), sharding=None),\\n\",\n \" 'linear/w': Param(value=(10, 10), sharding=None)\\n\",\n \"})\\n\",\n \"batch_stats = State({\\n\",\n \" 'layers/bn/mean': BatchStat(value=(4, 10), sharding=None),\\n\",\n \" 'layers/bn/var': BatchStat(value=(4, 10), sharding=None)\\n\",\n \"})\\n\"\n ]\n }\n ],\n \"source\": [\n \"# split\\n\",\n \"params = state.extract(Param)\\n\",\n \"batch_stats = state.extract(BatchStat)\\n\",\n \"# merge\\n\",\n \"state = State({**params, **batch_stats})\\n\",\n \"\\n\",\n \"print(\\\"params =\\\", jax.tree.map(jnp.shape, params))\\n\",\n \"print(\\\"batch_stats =\\\", jax.tree.map(jnp.shape, batch_stats))\"\n ]\n }\n ],\n \"metadata\": {\n \"language_info\": {\n \"codemirror_mode\": {\n \"name\": \"ipython\",\n \"version\": 3\n },\n \"file_extension\": \".py\",\n \"mimetype\": \"text/x-python\",\n \"name\": \"python\",\n \"nbconvert_exporter\": \"python\",\n \"pygments_lexer\": \"ipython3\",\n \"version\": \"3.10.13\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 2\n}\n" }, { "path": "docs_nnx/guides/transforms.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"id\": \"962be290\",\n \"metadata\": {},\n \"source\": [\n \"# Transformations\\n\",\n \"\\n\",\n \"In general, JAX transformations (transforms) operate on [pytrees](https://jax.readthedocs.io/en/latest/pytrees.html) of `jax.Array`s\\n\",\n \"and abide by value semantics. This presents a challenge for Flax NNX, which represents `nnx.Module`s as regular Python objects\\n\",\n \"that follow reference semantics. To address this, Flax NNX introduced its own set of transforms that extend JAX\\n\",\n \"transforms to allow `nnx.Module`s and other Flax NNX objects to be passed in and out of transforms while preserving\\n\",\n \"reference semantics.\\n\",\n \"\\n\",\n \"Flax NNX transforms should feel quite familiar if you have used JAX transforms before. They use the\\n\",\n \"same APIs and behave like the JAX transforms when only working with pytrees of `jax.Array`s. However, when working with\\n\",\n \"Flax NNX objects, they allow Python's reference semantics to be preserved for these objects, this includes:\\n\",\n \"\\n\",\n \"* Preserving shared references across multiple objects in the inputs and outputs of the transformation.\\n\",\n \"* Propagating any state changes made to the objects inside the transformation to the objects outside the transformation.\\n\",\n \"* Enforcing consistency of how objects are transformed when aliases are present across multiple inputs and outputs.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"8d645146\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"import jax\\n\",\n \"from jax import numpy as jnp, random\\n\",\n \"from flax import nnx\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"b44fb248\",\n \"metadata\": {},\n \"source\": [\n \"Throughout this guide, `nnx.vmap` is used as a case study to demonstrate how Flax NNX transforms work. However, the principles\\n\",\n \"outlined in this document extends to all transforms.\\n\",\n \"\\n\",\n \"## Basic example\\n\",\n \"\\n\",\n \"To begin, let's look at a simple example of using `nnx.vmap` to extend an element wise `vector_dot` function to work on\\n\",\n \"batched inputs. We will define a `Weights` Module with no methods to hold some parameters, these weights will be passed\\n\",\n \"as an input to the `vector_dot` function along with some data. Both the weights and data will be batched on axis `0` and we will use\\n\",\n \"`nnx.vmap` to apply `vector_dot` to each batch element, and the result will be a batched on axis `1`:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 2,\n \"id\": \"4eab27a4\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"y.shape = (3, 10)\\n\"\n ]\n },\n {\n \"name\": \"stderr\",\n \"output_type\": \"stream\",\n \"text\": [\n \"/Users/ivyzheng/envs/f1/lib/python3.12/site-packages/treescope/renderers.py:251: UserWarning: Ignoring error inside wrapper hook :\\n\",\n \"Traceback (most recent call last):\\n\",\n \" File \\\"/Users/ivyzheng/envs/f1/lib/python3.12/site-packages/treescope/renderers.py\\\", line 225, in _render_subtree\\n\",\n \" postprocessed_result = hook(\\n\",\n \" ^^^^^\\n\",\n \" File \\\"/Users/ivyzheng/envs/f1/lib/python3.12/site-packages/treescope/_internal/handlers/autovisualizer_hook.py\\\", line 47, in use_autovisualizer_if_present\\n\",\n \" result = autoviz(node, path)\\n\",\n \" ^^^^^^^^^^^^^^^^^^^\\n\",\n \" File \\\"/Users/ivyzheng/envs/f1/lib/python3.12/site-packages/treescope/_internal/api/array_autovisualizer.py\\\", line 306, in __call__\\n\",\n \" jax.sharding.PositionalSharding\\n\",\n \" File \\\"/Users/ivyzheng/envs/f1/lib/python3.12/site-packages/jax/_src/deprecations.py\\\", line 54, in getattr\\n\",\n \" raise AttributeError(message)\\n\",\n \"AttributeError: jax.sharding.PositionalSharding was deprecated in JAX v0.6.0 and removed in JAX v0.7.0\\n\",\n \"\\n\",\n \" warnings.warn(\\n\"\n ]\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"class Weights(nnx.Module):\\n\",\n \" def __init__(self, kernel: jax.Array, bias: jax.Array):\\n\",\n \" self.kernel, self.bias = nnx.Param(kernel), nnx.Param(bias)\\n\",\n \"\\n\",\n \"weights = Weights(\\n\",\n \" kernel=random.uniform(random.key(0), (10, 2, 3)),\\n\",\n \" bias=jnp.zeros((10, 3)),\\n\",\n \")\\n\",\n \"x = jax.random.normal(random.key(1), (10, 2))\\n\",\n \"\\n\",\n \"def vector_dot(weights: Weights, x: jax.Array):\\n\",\n \" assert weights.kernel.ndim == 2, 'Batch dimensions not allowed'\\n\",\n \" assert x.ndim == 1, 'Batch dimensions not allowed'\\n\",\n \" return x @ weights.kernel + weights.bias\\n\",\n \"\\n\",\n \"y = nnx.vmap(vector_dot, in_axes=0, out_axes=1)(weights, x)\\n\",\n \"\\n\",\n \"print(f'{y.shape = }')\\n\",\n \"nnx.display(weights)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"d2b222eb\",\n \"metadata\": {},\n \"source\": [\n \"Notice that `in_axes` interacts naturally with the `Weights` Module, treating it as if it were a pytree of `jax.Array`s. Prefix patterns are also allowed, so `in_axes=(0, 0)` would have also worked in this case.\\n\",\n \"\\n\",\n \"Objects are also allowed as outputs of Flax NNX transforms, which can be useful to transform initializers. For example,\\n\",\n \"you can define a `create_weights` function to create an single `Weights` `nnx.Module`, and use `nnx.vmap` to create a stack of\\n\",\n \"`Weights` with the same shapes as before:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 3,\n \"id\": \"0b076a0f\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"def create_weights(seed: jax.Array):\\n\",\n \" return Weights(\\n\",\n \" kernel=random.uniform(random.key(seed), (2, 3)),\\n\",\n \" bias=jnp.zeros((3,)),\\n\",\n \" )\\n\",\n \"\\n\",\n \"seeds = jnp.arange(10)\\n\",\n \"weights = nnx.vmap(create_weights)(seeds)\\n\",\n \"nnx.display(weights)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"fac3dca9\",\n \"metadata\": {},\n \"source\": [\n \"## Transforming methods\\n\",\n \"\\n\",\n \"Methods in Python are just functions that take the instance as the first argument, this means that you can decorate methods from `Module` and other Flax NNX subtypes. For example, we can refactor `Weights` from the previous example and decorate `__init__` with `vmap` to do the work of `create_weights`, and add a `__call__` method and decorate it with `@nnx.vmap` to do the work of `vector_dot`:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 4,\n \"id\": \"5d9a55fd\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"y.shape = (3, 10)\\n\"\n ]\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"class WeightStack(nnx.Module):\\n\",\n \" @nnx.vmap\\n\",\n \" def __init__(self, seed: jax.Array):\\n\",\n \" self.kernel = nnx.Param(random.uniform(random.key(seed), (2, 3)))\\n\",\n \" self.bias = nnx.Param(jnp.zeros((3,)))\\n\",\n \"\\n\",\n \" @nnx.vmap(in_axes=0, out_axes=1)\\n\",\n \" def __call__(self, x: jax.Array):\\n\",\n \" assert self.kernel.ndim == 2, 'Batch dimensions not allowed'\\n\",\n \" assert x.ndim == 1, 'Batch dimensions not allowed'\\n\",\n \" return x @ self.kernel + self.bias\\n\",\n \"\\n\",\n \"weights = WeightStack(jnp.arange(10))\\n\",\n \"\\n\",\n \"x = jax.random.normal(random.key(1), (10, 2))\\n\",\n \"y = weights(x)\\n\",\n \"\\n\",\n \"print(f'{y.shape = }')\\n\",\n \"nnx.display(weights)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"13b52d61\",\n \"metadata\": {},\n \"source\": [\n \"The rest of the guide will focus on transforming individual functions. But do note that all examples can be written in this method style.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"0251e7db\",\n \"metadata\": {},\n \"source\": [\n \"## State propagation\\n\",\n \"\\n\",\n \"So far our functions have been stateless. However, the real power of Flax NNX transforms comes when you have stateful functions, because one of their main features is to propagate state changes to preserve reference semantics. Let's update the previous example by adding\\n\",\n \"a `count` attribute to `Weights` and incrementing it in the new `stateful_vector_dot` function:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 5,\n \"id\": \"a4fbadb3\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"Count( # 10 (40 B)\\n\",\n \" value=Array([ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10], dtype=int32)\\n\",\n \")\"\n ]\n },\n \"execution_count\": 5,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"class Count(nnx.Variable): pass\\n\",\n \"\\n\",\n \"class Weights(nnx.Module):\\n\",\n \" def __init__(self, kernel: jax.Array, bias: jax.Array, count: jax.Array):\\n\",\n \" self.kernel, self.bias = nnx.Param(kernel), nnx.Param(bias)\\n\",\n \" self.count = Count(count)\\n\",\n \"\\n\",\n \"weights = Weights(\\n\",\n \" kernel=random.uniform(random.key(0), (10, 2, 3)),\\n\",\n \" bias=jnp.zeros((10, 3)),\\n\",\n \" count=jnp.arange(10),\\n\",\n \")\\n\",\n \"x = jax.random.normal(random.key(1), (10, 2))\\n\",\n \"\\n\",\n \"def stateful_vector_dot(weights: Weights, x: jax.Array):\\n\",\n \" assert weights.kernel.ndim == 2, 'Batch dimensions not allowed'\\n\",\n \" assert x.ndim == 1, 'Batch dimensions not allowed'\\n\",\n \" weights.count += 1\\n\",\n \" return x @ weights.kernel + weights.bias\\n\",\n \"\\n\",\n \"\\n\",\n \"y = nnx.vmap(stateful_vector_dot, in_axes=0, out_axes=1)(weights, x)\\n\",\n \"\\n\",\n \"weights.count\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"322312ee\",\n \"metadata\": {},\n \"source\": [\n \"After running `stateful_vector_dot` once, you verified that the `count` attribute was correctly updated. Because `Weights` was vectorized, `count` was initialized as an `arange(10)`, and all of its elements were incremented by `1` inside the transformation. The most important part is that updates were propagated to the original `Weights` object outside the transformation. Nice!\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"7294661f\",\n \"metadata\": {},\n \"source\": [\n \"### Graph updates propagation\\n\",\n \"\\n\",\n \"JAX transforms see inputs as pytrees of `jax.Array`s, and Flax NNX sees inputs as pytrees of `jax.Array`s and Python references, where references form a graph. Flax NNX's state propagation machinery can track arbitrary updates to the objects as long as they're local to the inputs (updates to globals inside transforms are not supported).\\n\",\n \"\\n\",\n \"This means that you can modify graph structure as needed, including updating existing attributes, adding/deleting attributes, swapping attributes, sharing (new) references between objects, sharing `nnx.Variable`s between objects, etc. Sky is the limit!\\n\",\n \"\\n\",\n \"The following example demonstrates performing some arbitrary updates to the `Weights` object inside `nnx.vmap`, and verifying that the updates are correctly propagated to the original `Weights` object outside the transformation:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 6,\n \"id\": \"76c58a29\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"class Count(nnx.Variable): pass\\n\",\n \"\\n\",\n \"class Weights(nnx.Module):\\n\",\n \" def __init__(self, kernel: jax.Array, bias: jax.Array, count: jax.Array):\\n\",\n \" self.kernel, self.bias = nnx.Param(kernel), nnx.Param(bias)\\n\",\n \" self.count = Count(count)\\n\",\n \"\\n\",\n \"weights = Weights(\\n\",\n \" kernel=random.uniform(random.key(0), (10, 2, 3)),\\n\",\n \" bias=jnp.zeros((10, 3)),\\n\",\n \" count=jnp.arange(10),\\n\",\n \")\\n\",\n \"x = jax.random.normal(random.key(1), (10, 2))\\n\",\n \"\\n\",\n \"def crazy_vector_dot(weights: Weights, x: jax.Array):\\n\",\n \" assert weights.kernel.ndim == 2, 'Batch dimensions not allowed'\\n\",\n \" assert x.ndim == 1, 'Batch dimensions not allowed'\\n\",\n \" weights.count += 1\\n\",\n \" y = x @ weights.kernel + weights.bias\\n\",\n \" weights.some_property = ['a', 2, False] # add attribute\\n\",\n \" del weights.bias # delete attribute\\n\",\n \" weights.new_param = weights.kernel # share reference\\n\",\n \" return y\\n\",\n \"\\n\",\n \"y = nnx.vmap(crazy_vector_dot, in_axes=0, out_axes=1)(weights, x)\\n\",\n \"\\n\",\n \"nnx.display(weights)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"743bcc34\",\n \"metadata\": {},\n \"source\": [\n \"> With great power comes great responsibility.\\n\",\n \">
\\\\- Uncle Ben\\n\",\n \"\\n\",\n \"While this feature is very powerful, it must be used with care because it can clash with JAX's underlying assumptions for certain transforms. For example, `jit` expects the structure of the inputs to be stable in order to cache the compiled function, so changing the graph structure inside an `nnx.jit`-ed function causes continuous recompilations and performance degradation. On the other hand, `scan` only allows a fixed `carry` structure, so adding/removing sub-states declared as carry will cause an error.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"0d11d191\",\n \"metadata\": {},\n \"source\": [\n \"## Transforming sub-states (lift types)\\n\",\n \"\\n\",\n \"Certain JAX transforms allow the use of pytree prefixes to specify how different parts of the inputs/outputs should be transformed. Flax NNX supports pytree prefixes for pytree structures but currently it doesn't have the notion of a prefix for graph objects. Instead, Flax NNX introduces the concept of “lift types” which allow specifying how different sub-states of an object should be transformed. Different transforms support different lift types, here is the list of currently supported FLax NNX lift types for each JAX transformation:\\n\",\n \"\\n\",\n \"| Lift type | JAX transforms |\\n\",\n \"|------------------|-----------------------------------------|\\n\",\n \"| `StateAxes` | `vmap`, `pmap`, `scan` |\\n\",\n \"| `StateSharding` | `jit`, `shard_map` |\\n\",\n \"| `DiffState` | `grad`, `value_and_grad`, `custom_vjp` |\\n\",\n \"\\n\",\n \"To specify how to vectorize different sub-states of an object in `nnx.vmap`, the Flax team created a `nnx.StateAxes`. `StateAxes` maps a set of sub-states via Flax NNX [Filters](https://flax.readthedocs.io/en/latest/guides/filters_guide.html) to their corresponding axes, and you can pass the `nnx.StateAxes` to `in_axes` and `out_axes` as if it/they were a pytree prefix.\\n\",\n \"\\n\",\n \"Let's use the previous `stateful_vector_dot` example and vectorize only the `nnx.Param` variables and broadcast the `count` variable so we only keep a single count for all the batch elements.\\n\",\n \"To do this we will define a `nnx.StateAxes` with a filter that matches the `nnx.Param` variables and maps them to axis `0`, and all the `Count` variables to `None`, and pass this `nnx.StateAxes` to `in_axes` for the `Weights` object.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 7,\n \"id\": \"d10aee8a\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"Count( # 1 (4 B)\\n\",\n \" value=Array(1, dtype=int32, weak_type=True)\\n\",\n \")\"\n ]\n },\n \"execution_count\": 7,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"class Weights(nnx.Module):\\n\",\n \" def __init__(self, kernel: jax.Array, bias: jax.Array, count: jax.Array):\\n\",\n \" self.kernel, self.bias = nnx.Param(kernel), nnx.Param(bias)\\n\",\n \" self.count = Count(count)\\n\",\n \"\\n\",\n \"weights = Weights(\\n\",\n \" kernel=random.uniform(random.key(0), (10, 2, 3)),\\n\",\n \" bias=jnp.zeros((10, 3)),\\n\",\n \" count=jnp.array(0),\\n\",\n \")\\n\",\n \"x = jax.random.normal(random.key(1), (10, 2))\\n\",\n \"\\n\",\n \"\\n\",\n \"def stateful_vector_dot(weights: Weights, x: jax.Array):\\n\",\n \" assert weights.kernel.ndim == 2, 'Batch dimensions not allowed'\\n\",\n \" assert x.ndim == 1, 'Batch dimensions not allowed'\\n\",\n \" weights.count += 1\\n\",\n \" return x @ weights.kernel + weights.bias\\n\",\n \"\\n\",\n \"state_axes = nnx.StateAxes({nnx.Param: 0, Count: None}) # broadcast Count\\n\",\n \"y = nnx.vmap(stateful_vector_dot, in_axes=(state_axes, 0), out_axes=1)(weights, x)\\n\",\n \"\\n\",\n \"weights.count\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"1cfd87e1\",\n \"metadata\": {},\n \"source\": [\n \"Here, `count` is now a scalar since it's not being vectorized. Also, note that `nnx.StateAxes` can only be used directly on Flax NNX objects, and it cannot be used as a prefix for a pytree of objects.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"1c8bb104\",\n \"metadata\": {},\n \"source\": [\n \"### Random state\\n\",\n \"\\n\",\n \"In Flax NNX, a random state is just a regular state. This means that it is stored inside `nnx.Module`s that need it, and it is treated as any other type of state. This is a simplification over Flax Linen, where a random state was handled by a separate mechanism. In practice `nnx.Module`s simply need to keep a reference to a `Rngs` object that is passed to them during initialization, and use it to generate a unique key for each random operation. For the purposes of this guide, this means that random state can be transformed like any other type of state but we also need to be aware of how the state is laid out so we can transform it correctly.\\n\",\n \"\\n\",\n \"Suppose you want to change things up a bit and apply the same weights to all elements in the batch. But you also want to add different random noise to each element.\\n\",\n \"\\n\",\n \"To do this, you will add an `Rngs` attribute to `Weights`, created from a `seed` key argument passed during construction. This seed key must be `split` beforehand, so that you can vectorize it successfully. For pedagogical reasons, you will assign the seed key to a `noise` “stream” and sample from it. To vectorize the PRNG state, you must configure `nnx.StateAxes` to map all `RngState`s (a base class for all variables in `Rngs`) to axis `0`, and `nnx.Param` and `Count` to `None`.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 8,\n \"id\": \"33c284b6\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"False\\n\"\n ]\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"class Weights(nnx.Module):\\n\",\n \" def __init__(self, kernel, bias, count, seed):\\n\",\n \" self.kernel, self.bias = nnx.Param(kernel), nnx.Param(bias)\\n\",\n \" self.count = Count(count)\\n\",\n \" self.rngs = nnx.Rngs(noise=seed)\\n\",\n \"\\n\",\n \"weights = Weights(\\n\",\n \" kernel=random.uniform(random.key(0), (2, 3)),\\n\",\n \" bias=jnp.zeros((3,)),\\n\",\n \" count=jnp.array(0),\\n\",\n \" seed=random.split(random.key(0), num=10),\\n\",\n \")\\n\",\n \"x = random.normal(random.key(1), (10, 2))\\n\",\n \"\\n\",\n \"def noisy_vector_dot(weights: Weights, x: jax.Array):\\n\",\n \" assert weights.kernel.ndim == 2, 'Batch dimensions not allowed'\\n\",\n \" assert x.ndim == 1, 'Batch dimensions not allowed'\\n\",\n \" weights.count += 1\\n\",\n \" y = x @ weights.kernel + weights.bias\\n\",\n \" return y + random.normal(weights.rngs.noise(), y.shape)\\n\",\n \"\\n\",\n \"state_axes = nnx.StateAxes({nnx.RngState: 0, (nnx.Param, Count): None})\\n\",\n \"y1 = nnx.vmap(noisy_vector_dot, in_axes=(state_axes, 0))(weights, x)\\n\",\n \"y2 = nnx.vmap(noisy_vector_dot, in_axes=(state_axes, 0))(weights, x)\\n\",\n \"\\n\",\n \"print(jnp.allclose(y1, y2))\\n\",\n \"nnx.display(weights)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"6f26b99f\",\n \"metadata\": {},\n \"source\": [\n \"Because `Rngs`'s state is updated in place and automatically propagated by `nnx.vmap`, we will get a different result every time that `noisy_vector_dot` is called.\\n\",\n \"\\n\",\n \"In the example above, you manually split the random state during construction. This is fine, as it makes the intention clear, but it also doesn't let you use `Rngs` outside of `nnx.vmap` because its state is always split. To solve this, you can pass an unsplit seed and use the `nnx.split_rngs` decorator before `nnx.vmap` to split the `RngState` right before each call to the function, and then \\\"lower\\\" it back so that it becomes usable.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 9,\n \"id\": \"8c9e5858\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"False\\n\"\n ]\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"weights = Weights(\\n\",\n \" kernel=random.uniform(random.key(0), (2, 3)),\\n\",\n \" bias=jnp.zeros((3,)),\\n\",\n \" count=jnp.array(0),\\n\",\n \" seed=0,\\n\",\n \")\\n\",\n \"x = random.normal(random.key(1), (10, 2))\\n\",\n \"\\n\",\n \"state_axes = nnx.StateAxes({nnx.RngState: 0, (nnx.Param, Count): None})\\n\",\n \"\\n\",\n \"@nnx.split_rngs(splits=10)\\n\",\n \"@nnx.vmap(in_axes=(state_axes, 0))\\n\",\n \"def noisy_vector_dot(weights: Weights, x: jax.Array):\\n\",\n \" assert weights.kernel.ndim == 2, 'Batch dimensions not allowed'\\n\",\n \" assert x.ndim == 1, 'Batch dimensions not allowed'\\n\",\n \" weights.count += 1\\n\",\n \" y = x @ weights.kernel + weights.bias\\n\",\n \" return y + random.normal(weights.rngs.noise(), y.shape)\\n\",\n \"\\n\",\n \"y1 = noisy_vector_dot(weights, x)\\n\",\n \"y2 = noisy_vector_dot(weights, x)\\n\",\n \"\\n\",\n \"print(jnp.allclose(y1, y2))\\n\",\n \"nnx.display(weights)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"60eee7f9\",\n \"metadata\": {},\n \"source\": [\n \"## Rules and limitations\\n\",\n \"In this section we will cover some rules and limitations apply when using Modules inside transformations.\\n\",\n \"\\n\",\n \"### Mutable Module cannot be passed by closure\\n\",\n \"\\n\",\n \"While Python allows for passing objects as closures to functions, this is generally not supported by Flax NNX transforms. The reason is that because Modules are mutable it is very easy to capture tracer into a Module created outside of the transform, this is silent error in JAX. To avoid this, Flax NNX checks that the Modules and Variables being mutated are passed as arguments to the transformed function.\\n\",\n \"\\n\",\n \"For example, if we have a stateful Module such as `Counter` that increments a counter every time it is called, and we try to pass it as a closure to a function decorated with `nnx.jit`, we would be leaking the tracer. However Flax NNX will raise an error instead to prevent this:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 10,\n \"id\": \"f8b95c03\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"Cannot mutate Param from a different trace level (https://flax.readthedocs.io/en/latest/api_reference/flax.errors.html#flax.errors.TraceContextError)\\n\"\n ]\n }\n ],\n \"source\": [\n \"class Counter(nnx.Module):\\n\",\n \" def __init__(self):\\n\",\n \" self.count = nnx.Param(jnp.array(0))\\n\",\n \"\\n\",\n \" def increment(self):\\n\",\n \" self.count += jnp.array(1)\\n\",\n \"\\n\",\n \"counter = Counter()\\n\",\n \"\\n\",\n \"@nnx.jit\\n\",\n \"def f(x):\\n\",\n \" counter.increment()\\n\",\n \" return 2 * x\\n\",\n \"\\n\",\n \"try:\\n\",\n \" y = f(3)\\n\",\n \"except Exception as e:\\n\",\n \" print(e)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"6f37e23b\",\n \"metadata\": {},\n \"source\": [\n \"To solve this issue pass all Module as arguments to the functions being transformed. In this case `f` should accept `counter` as an argument.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"75edf7a8\",\n \"metadata\": {},\n \"source\": [\n \"### Consistent aliasing\\n\",\n \"\\n\",\n \"The main issue with allowing for reference semantics in transforms is that references can be shared across inputs and outputs. This can be problematic if it is not taken care of because it would lead to ill-defined or inconsistent behavior. In the example below you have a single `Weights` Module `m` whose reference appears in multiple places in `arg1` and `arg2`. The problem here is that you also specify that you want to vectorize `arg1` in axis `0` and `arg2` in axis `1`. This would be fine in JAX because of referential transparency of pytrees. But this would be problematic in Flax NNX because you are trying to vectorize `m` in two different ways. Flax NNX will enforce consistency by raising an error.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 11,\n \"id\": \"46b1cc25\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"Inconsistent aliasing detected. The following nodes have different prefixes:\\n\",\n \"Node: Param\\n\",\n \" param: 0\\n\",\n \" param: 0\\n\",\n \" param: 1\\n\"\n ]\n }\n ],\n \"source\": [\n \"class Weights(nnx.Module):\\n\",\n \" def __init__(self, array: jax.Array):\\n\",\n \" self.param = nnx.Param(array)\\n\",\n \"\\n\",\n \"m = Weights(jnp.arange(10))\\n\",\n \"arg1 = {'a': {'b': m}, 'c': m}\\n\",\n \"arg2 = [(m, m), m]\\n\",\n \"\\n\",\n \"@nnx.vmap(in_axes=(0, 1))\\n\",\n \"def f(arg1, arg2):\\n\",\n \" ...\\n\",\n \"\\n\",\n \"try:\\n\",\n \" f(arg1, arg2)\\n\",\n \"except ValueError as e:\\n\",\n \" print(e)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"46aa978c\",\n \"metadata\": {},\n \"source\": [\n \"Inconsistent aliasing can also happen between inputs and outputs. In the next example you have a trivial function that accepts and immediately returns `arg1`. However, `arg1` is vectorized on axis `0` on the input, and axis `1` on the output. As expected, this is problematic and Flax NNX will raise an error.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 12,\n \"id\": \"cca9cf31\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"Inconsistent aliasing detected. The following nodes have different prefixes:\\n\",\n \"Node: Param\\n\",\n \" param: 0\\n\",\n \" param: 0\\n\",\n \" param: 1\\n\"\n ]\n }\n ],\n \"source\": [\n \"@nnx.vmap(in_axes=0, out_axes=1)\\n\",\n \"def f(arg1):\\n\",\n \" return arg1\\n\",\n \"\\n\",\n \"try:\\n\",\n \" f(arg1)\\n\",\n \"except ValueError as e:\\n\",\n \" print(e)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"13f9aeea\",\n \"metadata\": {},\n \"source\": [\n \"## Axis metadata\\n\",\n \"\\n\",\n \"Flax NNX `Variable`s can hold arbitrary metadata, which can be added by simply passing it as keyword arguments to its constructor. This is often used to store `sharding` information, as used by the `nnx.spmd` APIs (like `nnx.get_partition_spec` and `nnx.get_named_sharding`).\\n\",\n \"\\n\",\n \"However, it is often important to keep this axes-related information in sync to what the actual state of the axes is when transforms are involved. For example, if you vectorize a variable on axis `1`, you should remove the `sharding` information at position `1` when inside a `vmap` or `scan` to reflect the fact that the axes are temporarily removed.\\n\",\n \"\\n\",\n \"To achieve this, Flax NNX transforms provide a non-standard `transform_metadata` dictionary argument. And when the `nnx.PARTITION_NAME` key is present, the `sharding` metadata will be updated as specified by `in_axes` and `out_axes`.\\n\",\n \"\\n\",\n \"Let's see an example of this in action:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"d85c772c\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"Inner m.param.shape = (3, 5)\\n\",\n \"Inner m.param.out_sharding = ('a', None)\\n\",\n \"Outter m.param.shape = (3, 4, 5)\\n\",\n \"Outter m.param.out_sharding = ('a', 'b', None)\\n\"\n ]\n }\n ],\n \"source\": [\n \"mesh = jax.make_mesh((1, 1), ('a', 'b'))\\n\",\n \"\\n\",\n \"class Weights(nnx.Module):\\n\",\n \" def __init__(self, array: jax.Array, out_sharding: tuple[str | None, ...]):\\n\",\n \" self.param = nnx.Param(array, out_sharding=out_sharding)\\n\",\n \"\\n\",\n \"@nnx.vmap(in_axes=1, transform_metadata={nnx.PARTITION_NAME: 'b'})\\n\",\n \"def f(m: Weights):\\n\",\n \" print(f'Inner {m.param.shape = }')\\n\",\n \" print(f'Inner {m.param.out_sharding = }')\\n\",\n \"\\n\",\n \"with jax.set_mesh(mesh):\\n\",\n \" m = Weights(jnp.ones((3, 4, 5)), out_sharding=('a', 'b', None))\\n\",\n \" f(m)\\n\",\n \"\\n\",\n \"print(f'Outter {m.param.shape = }')\\n\",\n \"print(f'Outter {m.param.out_sharding = }')\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"a23bda09\",\n \"metadata\": {},\n \"source\": [\n \"Here, you added a `sharding` metadata to the `nnx.Param` variables, and used `transform_metadata` to update the `sharding` metadata to reflect the axis changes. Specifically, you can see that the first axis `b` was removed from the `sharding` metadata when inside of `nnx.vmap`, and then added back when outside of `nnx.vmap`.\\n\",\n \"\\n\",\n \"You can verify that this also works when `nnx.Module`s are created inside the transformation - the new `sharding` axes will be added to the `nnx.Module` `nnx.Variable`s outside the transformation, matching the axes of the transformed `nnx.Variable`s.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"358e51f7\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"Outter m.param.shape = (3, 4, 5)\\n\",\n \"Outter m.param.out_sharding = ('a', 'b', None)\\n\"\n ]\n }\n ],\n \"source\": [\n \"@nnx.vmap(out_axes=1, axis_size=4, transform_metadata={nnx.PARTITION_NAME: 'b'})\\n\",\n \"def init_vmap():\\n\",\n \" return Weights(jnp.ones((3, 5)), out_sharding=('a', None))\\n\",\n \"\\n\",\n \"with jax.set_mesh(mesh):\\n\",\n \" m = init_vmap()\\n\",\n \"print(f'Outter {m.param.shape = }')\\n\",\n \"print(f'Outter {m.param.out_sharding = }')\"\n ]\n }\n ],\n \"metadata\": {\n \"jupytext\": {\n \"cell_metadata_filter\": \"-all\",\n \"formats\": \"ipynb,md:myst\",\n \"main_language\": \"python\"\n },\n \"language_info\": {\n \"codemirror_mode\": {\n \"name\": \"ipython\",\n \"version\": 3\n },\n \"file_extension\": \".py\",\n \"mimetype\": \"text/x-python\",\n \"name\": \"python\",\n \"nbconvert_exporter\": \"python\",\n \"pygments_lexer\": \"ipython3\",\n \"version\": \"3.12.7\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 5\n}\n" }, { "path": "docs_nnx/guides/transforms.md", "content": "---\njupytext:\n cell_metadata_filter: -all\n formats: ipynb,md:myst\n main_language: python\n text_representation:\n extension: .md\n format_name: myst\n format_version: 0.13\n jupytext_version: 1.13.8\n---\n\n# Transformations\n\nIn general, JAX transformations (transforms) operate on [pytrees](https://jax.readthedocs.io/en/latest/pytrees.html) of `jax.Array`s\nand abide by value semantics. This presents a challenge for Flax NNX, which represents `nnx.Module`s as regular Python objects\nthat follow reference semantics. To address this, Flax NNX introduced its own set of transforms that extend JAX\ntransforms to allow `nnx.Module`s and other Flax NNX objects to be passed in and out of transforms while preserving\nreference semantics.\n\nFlax NNX transforms should feel quite familiar if you have used JAX transforms before. They use the\nsame APIs and behave like the JAX transforms when only working with pytrees of `jax.Array`s. However, when working with\nFlax NNX objects, they allow Python's reference semantics to be preserved for these objects, this includes:\n\n* Preserving shared references across multiple objects in the inputs and outputs of the transformation.\n* Propagating any state changes made to the objects inside the transformation to the objects outside the transformation.\n* Enforcing consistency of how objects are transformed when aliases are present across multiple inputs and outputs.\n\n```{code-cell} ipython3\nimport jax\nfrom jax import numpy as jnp, random\nfrom flax import nnx\n```\n\nThroughout this guide, `nnx.vmap` is used as a case study to demonstrate how Flax NNX transforms work. However, the principles\noutlined in this document extends to all transforms.\n\n## Basic example\n\nTo begin, let's look at a simple example of using `nnx.vmap` to extend an element wise `vector_dot` function to work on\nbatched inputs. We will define a `Weights` Module with no methods to hold some parameters, these weights will be passed\nas an input to the `vector_dot` function along with some data. Both the weights and data will be batched on axis `0` and we will use\n`nnx.vmap` to apply `vector_dot` to each batch element, and the result will be a batched on axis `1`:\n\n```{code-cell} ipython3\nclass Weights(nnx.Module):\n def __init__(self, kernel: jax.Array, bias: jax.Array):\n self.kernel, self.bias = nnx.Param(kernel), nnx.Param(bias)\n\nweights = Weights(\n kernel=random.uniform(random.key(0), (10, 2, 3)),\n bias=jnp.zeros((10, 3)),\n)\nx = jax.random.normal(random.key(1), (10, 2))\n\ndef vector_dot(weights: Weights, x: jax.Array):\n assert weights.kernel.ndim == 2, 'Batch dimensions not allowed'\n assert x.ndim == 1, 'Batch dimensions not allowed'\n return x @ weights.kernel + weights.bias\n\ny = nnx.vmap(vector_dot, in_axes=0, out_axes=1)(weights, x)\n\nprint(f'{y.shape = }')\nnnx.display(weights)\n```\n\nNotice that `in_axes` interacts naturally with the `Weights` Module, treating it as if it were a pytree of `jax.Array`s. Prefix patterns are also allowed, so `in_axes=(0, 0)` would have also worked in this case.\n\nObjects are also allowed as outputs of Flax NNX transforms, which can be useful to transform initializers. For example,\nyou can define a `create_weights` function to create an single `Weights` `nnx.Module`, and use `nnx.vmap` to create a stack of\n`Weights` with the same shapes as before:\n\n```{code-cell} ipython3\ndef create_weights(seed: jax.Array):\n return Weights(\n kernel=random.uniform(random.key(seed), (2, 3)),\n bias=jnp.zeros((3,)),\n )\n\nseeds = jnp.arange(10)\nweights = nnx.vmap(create_weights)(seeds)\nnnx.display(weights)\n```\n\n## Transforming methods\n\nMethods in Python are just functions that take the instance as the first argument, this means that you can decorate methods from `Module` and other Flax NNX subtypes. For example, we can refactor `Weights` from the previous example and decorate `__init__` with `vmap` to do the work of `create_weights`, and add a `__call__` method and decorate it with `@nnx.vmap` to do the work of `vector_dot`:\n\n```{code-cell} ipython3\nclass WeightStack(nnx.Module):\n @nnx.vmap\n def __init__(self, seed: jax.Array):\n self.kernel = nnx.Param(random.uniform(random.key(seed), (2, 3)))\n self.bias = nnx.Param(jnp.zeros((3,)))\n\n @nnx.vmap(in_axes=0, out_axes=1)\n def __call__(self, x: jax.Array):\n assert self.kernel.ndim == 2, 'Batch dimensions not allowed'\n assert x.ndim == 1, 'Batch dimensions not allowed'\n return x @ self.kernel + self.bias\n\nweights = WeightStack(jnp.arange(10))\n\nx = jax.random.normal(random.key(1), (10, 2))\ny = weights(x)\n\nprint(f'{y.shape = }')\nnnx.display(weights)\n```\n\nThe rest of the guide will focus on transforming individual functions. But do note that all examples can be written in this method style.\n\n+++\n\n## State propagation\n\nSo far our functions have been stateless. However, the real power of Flax NNX transforms comes when you have stateful functions, because one of their main features is to propagate state changes to preserve reference semantics. Let's update the previous example by adding\na `count` attribute to `Weights` and incrementing it in the new `stateful_vector_dot` function:\n\n```{code-cell} ipython3\nclass Count(nnx.Variable): pass\n\nclass Weights(nnx.Module):\n def __init__(self, kernel: jax.Array, bias: jax.Array, count: jax.Array):\n self.kernel, self.bias = nnx.Param(kernel), nnx.Param(bias)\n self.count = Count(count)\n\nweights = Weights(\n kernel=random.uniform(random.key(0), (10, 2, 3)),\n bias=jnp.zeros((10, 3)),\n count=jnp.arange(10),\n)\nx = jax.random.normal(random.key(1), (10, 2))\n\ndef stateful_vector_dot(weights: Weights, x: jax.Array):\n assert weights.kernel.ndim == 2, 'Batch dimensions not allowed'\n assert x.ndim == 1, 'Batch dimensions not allowed'\n weights.count += 1\n return x @ weights.kernel + weights.bias\n\n\ny = nnx.vmap(stateful_vector_dot, in_axes=0, out_axes=1)(weights, x)\n\nweights.count\n```\n\nAfter running `stateful_vector_dot` once, you verified that the `count` attribute was correctly updated. Because `Weights` was vectorized, `count` was initialized as an `arange(10)`, and all of its elements were incremented by `1` inside the transformation. The most important part is that updates were propagated to the original `Weights` object outside the transformation. Nice!\n\n+++\n\n### Graph updates propagation\n\nJAX transforms see inputs as pytrees of `jax.Array`s, and Flax NNX sees inputs as pytrees of `jax.Array`s and Python references, where references form a graph. Flax NNX's state propagation machinery can track arbitrary updates to the objects as long as they're local to the inputs (updates to globals inside transforms are not supported).\n\nThis means that you can modify graph structure as needed, including updating existing attributes, adding/deleting attributes, swapping attributes, sharing (new) references between objects, sharing `nnx.Variable`s between objects, etc. Sky is the limit!\n\nThe following example demonstrates performing some arbitrary updates to the `Weights` object inside `nnx.vmap`, and verifying that the updates are correctly propagated to the original `Weights` object outside the transformation:\n\n```{code-cell} ipython3\nclass Count(nnx.Variable): pass\n\nclass Weights(nnx.Module):\n def __init__(self, kernel: jax.Array, bias: jax.Array, count: jax.Array):\n self.kernel, self.bias = nnx.Param(kernel), nnx.Param(bias)\n self.count = Count(count)\n\nweights = Weights(\n kernel=random.uniform(random.key(0), (10, 2, 3)),\n bias=jnp.zeros((10, 3)),\n count=jnp.arange(10),\n)\nx = jax.random.normal(random.key(1), (10, 2))\n\ndef crazy_vector_dot(weights: Weights, x: jax.Array):\n assert weights.kernel.ndim == 2, 'Batch dimensions not allowed'\n assert x.ndim == 1, 'Batch dimensions not allowed'\n weights.count += 1\n y = x @ weights.kernel + weights.bias\n weights.some_property = ['a', 2, False] # add attribute\n del weights.bias # delete attribute\n weights.new_param = weights.kernel # share reference\n return y\n\ny = nnx.vmap(crazy_vector_dot, in_axes=0, out_axes=1)(weights, x)\n\nnnx.display(weights)\n```\n\n> With great power comes great responsibility.\n>
\\- Uncle Ben\n\nWhile this feature is very powerful, it must be used with care because it can clash with JAX's underlying assumptions for certain transforms. For example, `jit` expects the structure of the inputs to be stable in order to cache the compiled function, so changing the graph structure inside an `nnx.jit`-ed function causes continuous recompilations and performance degradation. On the other hand, `scan` only allows a fixed `carry` structure, so adding/removing sub-states declared as carry will cause an error.\n\n+++\n\n## Transforming sub-states (lift types)\n\nCertain JAX transforms allow the use of pytree prefixes to specify how different parts of the inputs/outputs should be transformed. Flax NNX supports pytree prefixes for pytree structures but currently it doesn't have the notion of a prefix for graph objects. Instead, Flax NNX introduces the concept of “lift types” which allow specifying how different sub-states of an object should be transformed. Different transforms support different lift types, here is the list of currently supported FLax NNX lift types for each JAX transformation:\n\n| Lift type | JAX transforms |\n|------------------|-----------------------------------------|\n| `StateAxes` | `vmap`, `pmap`, `scan` |\n| `StateSharding` | `jit`, `shard_map` |\n| `DiffState` | `grad`, `value_and_grad`, `custom_vjp` |\n\nTo specify how to vectorize different sub-states of an object in `nnx.vmap`, the Flax team created a `nnx.StateAxes`. `StateAxes` maps a set of sub-states via Flax NNX [Filters](https://flax.readthedocs.io/en/latest/guides/filters_guide.html) to their corresponding axes, and you can pass the `nnx.StateAxes` to `in_axes` and `out_axes` as if it/they were a pytree prefix.\n\nLet's use the previous `stateful_vector_dot` example and vectorize only the `nnx.Param` variables and broadcast the `count` variable so we only keep a single count for all the batch elements.\nTo do this we will define a `nnx.StateAxes` with a filter that matches the `nnx.Param` variables and maps them to axis `0`, and all the `Count` variables to `None`, and pass this `nnx.StateAxes` to `in_axes` for the `Weights` object.\n\n```{code-cell} ipython3\nclass Weights(nnx.Module):\n def __init__(self, kernel: jax.Array, bias: jax.Array, count: jax.Array):\n self.kernel, self.bias = nnx.Param(kernel), nnx.Param(bias)\n self.count = Count(count)\n\nweights = Weights(\n kernel=random.uniform(random.key(0), (10, 2, 3)),\n bias=jnp.zeros((10, 3)),\n count=jnp.array(0),\n)\nx = jax.random.normal(random.key(1), (10, 2))\n\n\ndef stateful_vector_dot(weights: Weights, x: jax.Array):\n assert weights.kernel.ndim == 2, 'Batch dimensions not allowed'\n assert x.ndim == 1, 'Batch dimensions not allowed'\n weights.count += 1\n return x @ weights.kernel + weights.bias\n\nstate_axes = nnx.StateAxes({nnx.Param: 0, Count: None}) # broadcast Count\ny = nnx.vmap(stateful_vector_dot, in_axes=(state_axes, 0), out_axes=1)(weights, x)\n\nweights.count\n```\n\nHere, `count` is now a scalar since it's not being vectorized. Also, note that `nnx.StateAxes` can only be used directly on Flax NNX objects, and it cannot be used as a prefix for a pytree of objects.\n\n+++\n\n### Random state\n\nIn Flax NNX, a random state is just a regular state. This means that it is stored inside `nnx.Module`s that need it, and it is treated as any other type of state. This is a simplification over Flax Linen, where a random state was handled by a separate mechanism. In practice `nnx.Module`s simply need to keep a reference to a `Rngs` object that is passed to them during initialization, and use it to generate a unique key for each random operation. For the purposes of this guide, this means that random state can be transformed like any other type of state but we also need to be aware of how the state is laid out so we can transform it correctly.\n\nSuppose you want to change things up a bit and apply the same weights to all elements in the batch. But you also want to add different random noise to each element.\n\nTo do this, you will add an `Rngs` attribute to `Weights`, created from a `seed` key argument passed during construction. This seed key must be `split` beforehand, so that you can vectorize it successfully. For pedagogical reasons, you will assign the seed key to a `noise` “stream” and sample from it. To vectorize the PRNG state, you must configure `nnx.StateAxes` to map all `RngState`s (a base class for all variables in `Rngs`) to axis `0`, and `nnx.Param` and `Count` to `None`.\n\n```{code-cell} ipython3\nclass Weights(nnx.Module):\n def __init__(self, kernel, bias, count, seed):\n self.kernel, self.bias = nnx.Param(kernel), nnx.Param(bias)\n self.count = Count(count)\n self.rngs = nnx.Rngs(noise=seed)\n\nweights = Weights(\n kernel=random.uniform(random.key(0), (2, 3)),\n bias=jnp.zeros((3,)),\n count=jnp.array(0),\n seed=random.split(random.key(0), num=10),\n)\nx = random.normal(random.key(1), (10, 2))\n\ndef noisy_vector_dot(weights: Weights, x: jax.Array):\n assert weights.kernel.ndim == 2, 'Batch dimensions not allowed'\n assert x.ndim == 1, 'Batch dimensions not allowed'\n weights.count += 1\n y = x @ weights.kernel + weights.bias\n return y + random.normal(weights.rngs.noise(), y.shape)\n\nstate_axes = nnx.StateAxes({nnx.RngState: 0, (nnx.Param, Count): None})\ny1 = nnx.vmap(noisy_vector_dot, in_axes=(state_axes, 0))(weights, x)\ny2 = nnx.vmap(noisy_vector_dot, in_axes=(state_axes, 0))(weights, x)\n\nprint(jnp.allclose(y1, y2))\nnnx.display(weights)\n```\n\nBecause `Rngs`'s state is updated in place and automatically propagated by `nnx.vmap`, we will get a different result every time that `noisy_vector_dot` is called.\n\nIn the example above, you manually split the random state during construction. This is fine, as it makes the intention clear, but it also doesn't let you use `Rngs` outside of `nnx.vmap` because its state is always split. To solve this, you can pass an unsplit seed and use the `nnx.split_rngs` decorator before `nnx.vmap` to split the `RngState` right before each call to the function, and then \"lower\" it back so that it becomes usable.\n\n```{code-cell} ipython3\nweights = Weights(\n kernel=random.uniform(random.key(0), (2, 3)),\n bias=jnp.zeros((3,)),\n count=jnp.array(0),\n seed=0,\n)\nx = random.normal(random.key(1), (10, 2))\n\nstate_axes = nnx.StateAxes({nnx.RngState: 0, (nnx.Param, Count): None})\n\n@nnx.split_rngs(splits=10)\n@nnx.vmap(in_axes=(state_axes, 0))\ndef noisy_vector_dot(weights: Weights, x: jax.Array):\n assert weights.kernel.ndim == 2, 'Batch dimensions not allowed'\n assert x.ndim == 1, 'Batch dimensions not allowed'\n weights.count += 1\n y = x @ weights.kernel + weights.bias\n return y + random.normal(weights.rngs.noise(), y.shape)\n\ny1 = noisy_vector_dot(weights, x)\ny2 = noisy_vector_dot(weights, x)\n\nprint(jnp.allclose(y1, y2))\nnnx.display(weights)\n```\n\n## Rules and limitations\nIn this section we will cover some rules and limitations apply when using Modules inside transformations.\n\n### Mutable Module cannot be passed by closure\n\nWhile Python allows for passing objects as closures to functions, this is generally not supported by Flax NNX transforms. The reason is that because Modules are mutable it is very easy to capture tracer into a Module created outside of the transform, this is silent error in JAX. To avoid this, Flax NNX checks that the Modules and Variables being mutated are passed as arguments to the transformed function.\n\nFor example, if we have a stateful Module such as `Counter` that increments a counter every time it is called, and we try to pass it as a closure to a function decorated with `nnx.jit`, we would be leaking the tracer. However Flax NNX will raise an error instead to prevent this:\n\n```{code-cell} ipython3\nclass Counter(nnx.Module):\n def __init__(self):\n self.count = nnx.Param(jnp.array(0))\n\n def increment(self):\n self.count += jnp.array(1)\n\ncounter = Counter()\n\n@nnx.jit\ndef f(x):\n counter.increment()\n return 2 * x\n\ntry:\n y = f(3)\nexcept Exception as e:\n print(e)\n```\n\nTo solve this issue pass all Module as arguments to the functions being transformed. In this case `f` should accept `counter` as an argument.\n\n+++\n\n### Consistent aliasing\n\nThe main issue with allowing for reference semantics in transforms is that references can be shared across inputs and outputs. This can be problematic if it is not taken care of because it would lead to ill-defined or inconsistent behavior. In the example below you have a single `Weights` Module `m` whose reference appears in multiple places in `arg1` and `arg2`. The problem here is that you also specify that you want to vectorize `arg1` in axis `0` and `arg2` in axis `1`. This would be fine in JAX because of referential transparency of pytrees. But this would be problematic in Flax NNX because you are trying to vectorize `m` in two different ways. Flax NNX will enforce consistency by raising an error.\n\n```{code-cell} ipython3\nclass Weights(nnx.Module):\n def __init__(self, array: jax.Array):\n self.param = nnx.Param(array)\n\nm = Weights(jnp.arange(10))\narg1 = {'a': {'b': m}, 'c': m}\narg2 = [(m, m), m]\n\n@nnx.vmap(in_axes=(0, 1))\ndef f(arg1, arg2):\n ...\n\ntry:\n f(arg1, arg2)\nexcept ValueError as e:\n print(e)\n```\n\nInconsistent aliasing can also happen between inputs and outputs. In the next example you have a trivial function that accepts and immediately returns `arg1`. However, `arg1` is vectorized on axis `0` on the input, and axis `1` on the output. As expected, this is problematic and Flax NNX will raise an error.\n\n```{code-cell} ipython3\n@nnx.vmap(in_axes=0, out_axes=1)\ndef f(arg1):\n return arg1\n\ntry:\n f(arg1)\nexcept ValueError as e:\n print(e)\n```\n\n## Axis metadata\n\nFlax NNX `Variable`s can hold arbitrary metadata, which can be added by simply passing it as keyword arguments to its constructor. This is often used to store `sharding` information, as used by the `nnx.spmd` APIs (like `nnx.get_partition_spec` and `nnx.get_named_sharding`).\n\nHowever, it is often important to keep this axes-related information in sync to what the actual state of the axes is when transforms are involved. For example, if you vectorize a variable on axis `1`, you should remove the `sharding` information at position `1` when inside a `vmap` or `scan` to reflect the fact that the axes are temporarily removed.\n\nTo achieve this, Flax NNX transforms provide a non-standard `transform_metadata` dictionary argument. And when the `nnx.PARTITION_NAME` key is present, the `sharding` metadata will be updated as specified by `in_axes` and `out_axes`.\n\nLet's see an example of this in action:\n\n```{code-cell} ipython3\nmesh = jax.make_mesh((1, 1), ('a', 'b'))\n\nclass Weights(nnx.Module):\n def __init__(self, array: jax.Array, out_sharding: tuple[str | None, ...]):\n self.param = nnx.Param(array, out_sharding=out_sharding)\n\n@nnx.vmap(in_axes=1, transform_metadata={nnx.PARTITION_NAME: 'b'})\ndef f(m: Weights):\n print(f'Inner {m.param.shape = }')\n print(f'Inner {m.param.out_sharding = }')\n\nwith jax.set_mesh(mesh):\n m = Weights(jnp.ones((3, 4, 5)), out_sharding=('a', 'b', None))\n f(m)\n\nprint(f'Outter {m.param.shape = }')\nprint(f'Outter {m.param.out_sharding = }')\n```\n\nHere, you added a `sharding` metadata to the `nnx.Param` variables, and used `transform_metadata` to update the `sharding` metadata to reflect the axis changes. Specifically, you can see that the first axis `b` was removed from the `sharding` metadata when inside of `nnx.vmap`, and then added back when outside of `nnx.vmap`.\n\nYou can verify that this also works when `nnx.Module`s are created inside the transformation - the new `sharding` axes will be added to the `nnx.Module` `nnx.Variable`s outside the transformation, matching the axes of the transformed `nnx.Variable`s.\n\n```{code-cell} ipython3\n@nnx.vmap(out_axes=1, axis_size=4, transform_metadata={nnx.PARTITION_NAME: 'b'})\ndef init_vmap():\n return Weights(jnp.ones((3, 5)), out_sharding=('a', None))\n\nwith jax.set_mesh(mesh):\n m = init_vmap()\nprint(f'Outter {m.param.shape = }')\nprint(f'Outter {m.param.out_sharding = }')\n```\n" }, { "path": "docs_nnx/guides/view.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"id\": \"75afc9f3\",\n \"metadata\": {},\n \"source\": [\n \"# Model Views\\n\",\n \"This guide covers how to use the `nnx.view` function. This function is useful for handling state in layers like `Dropout` and `BatchNorm`, which behave differently in training and evaluation. Similar to `.view` for numpy arrays, `nnx.view` allows you to set modes of the model while still sharing the same data. For a quick intro to how this function works, refer to the following example:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"8e333aab\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"from flax import nnx\\n\",\n \"\\n\",\n \"# example model with different train/eval behavior\\n\",\n \"rngs = nnx.Rngs(0)\\n\",\n \"model = nnx.Sequential(\\n\",\n \" nnx.Linear(2, 4, rngs=rngs), nnx.BatchNorm(4, rngs=rngs), nnx.Dropout(0.1)\\n\",\n \")\\n\",\n \"\\n\",\n \"# set train and eval modes\\n\",\n \"train_model = nnx.view(model, deterministic=False, use_running_average=False)\\n\",\n \"eval_model = nnx.view(model, deterministic=True, use_running_average=True)\\n\",\n \"\\n\",\n \"# Can see deterministic is different between train_model and eval_model\\n\",\n \"assert train_model.layers[2].deterministic == False\\n\",\n \"assert eval_model.layers[2].deterministic == True\\n\",\n \"\\n\",\n \"# Weights are shared between the models\\n\",\n \"assert train_model.layers[0].kernel is eval_model.layers[0].kernel\\n\",\n \"\\n\",\n \"# Print information about kwargs for nnx.view with nnx.view_info\\n\",\n \"print(nnx.view_info(model))\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"f70f9353\",\n \"metadata\": {},\n \"source\": [\n \"## Motivation\\n\",\n \"\\n\",\n \"Some layers in ML inherently involve state. Consider for example the `nnx.Dropout` layer, which behaves differently during training and evaluation. In these different scenarios, we need a simple way to ensure that the model behaves as intended to avoid silent bugs. A common pattern in other frameworks is to mutate a single `model` object to switch between training and evaluation modes. This requires the programmer to remember to toggle modes in many places throughout the code, which can hurt readability and lead to subtle bugs when a mode switch is forgotten.\\n\",\n \"\\n\",\n \"`nnx.view` offers a cleaner alternative: you declare the different model configurations once at the beginning of your code and then simply use the appropriate view wherever needed. Each view shares the same underlying weights, so parameter updates are automatically reflected across all views. We demonstrate this with a simple example below.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"886c7479\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"import jax\\n\",\n \"import optax\\n\",\n \"import matplotlib.pyplot as plt\\n\",\n \"\\n\",\n \"in_dim, hidden_dim, out_dim = 16, 32, 2\\n\",\n \"\\n\",\n \"\\n\",\n \"class MyModel(nnx.Module):\\n\",\n \" def __init__(\\n\",\n \" self,\\n\",\n \" in_dim: int,\\n\",\n \" hidden_dim: int,\\n\",\n \" out_dim: int,\\n\",\n \" dropout_rate: float,\\n\",\n \" *,\\n\",\n \" rngs: nnx.Rngs,\\n\",\n \" ):\\n\",\n \" self.lin1 = nnx.Linear(in_dim, hidden_dim, rngs=rngs)\\n\",\n \" self.do = nnx.Dropout(dropout_rate)\\n\",\n \" self.bn = nnx.BatchNorm(hidden_dim, rngs=rngs)\\n\",\n \" self.lin2 = nnx.Linear(hidden_dim, out_dim, rngs=rngs)\\n\",\n \"\\n\",\n \" def __call__(self, x, *, rngs=None):\\n\",\n \" x = nnx.relu(self.do(self.bn(self.lin1(x)), rngs=rngs))\\n\",\n \" return self.lin2(x)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"0568e6ce\",\n \"metadata\": {},\n \"source\": [\n \"Lets take a look at the model to see what is going on.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"da2c8f80\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# can display to inspect state\\n\",\n \"model = MyModel(in_dim, hidden_dim, out_dim, 0.1, rngs=nnx.Rngs(0))\\n\",\n \"nnx.display(model)\\n\",\n \"\\n\",\n \"# can assert to inspect state\\n\",\n \"assert model.do.deterministic == False\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"1a075c12\",\n \"metadata\": {},\n \"source\": [\n \"From the model display, we can see that `Dropout` has `deterministic == False`, suggesting that the model is in training mode. In order to know this, we had to display the model and/or know that `Dropout` is set to training mode by default. It is not clear what state the model is in just by looking at the code without additional inspection. We instead want to be very explicit about what state the model is in. \\n\",\n \"\\n\",\n \"This is where `nnx.view` comes in. This function updates the modes for each submodule of a neural network based on the kwargs passed into the function. The underlying model weights are then shared between different views. We set up a training and evaluation version of the model below.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"11e59178\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"train_model = nnx.view(model, deterministic=False)\\n\",\n \"eval_model = nnx.view(model, deterministic=True)\\n\",\n \"\\n\",\n \"# weights are references to the same data\\n\",\n \"assert train_model.lin1.kernel is eval_model.lin1.kernel\\n\",\n \"\\n\",\n \"# Dropout.deterministic is different in each model\\n\",\n \"assert train_model.do.deterministic is False\\n\",\n \"assert eval_model.do.deterministic is True\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"5c1ee1db\",\n \"metadata\": {},\n \"source\": [\n \"## Example with `nnx.view`\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"e35d1bfd\",\n \"metadata\": {},\n \"source\": [\n \"We first set up data generators and define train/eval step functions. The `train_step` receives an `nnx.Rngs` object for dropout randomness, while `eval_step` doesn't since dropout is disabled in `eval_model`.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"a0f72d8d\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"ndata, batch_size, total_epochs, lr = 2048, 32, 100, 1e-3\\n\",\n \"rngs = nnx.Rngs(0)\\n\",\n \"x = rngs.normal((ndata, in_dim))\\n\",\n \"y = rngs.normal((ndata, out_dim))\\n\",\n \"\\n\",\n \"\\n\",\n \"@nnx.jit\\n\",\n \"def train_step(model, optimizer, x, y, rngs):\\n\",\n \" def loss_fn(model, rngs):\\n\",\n \" return ((model(x, rngs=rngs) - y) ** 2).mean()\\n\",\n \"\\n\",\n \" grads = nnx.grad(loss_fn)(model, rngs)\\n\",\n \" optimizer.update(model, grads)\\n\",\n \"\\n\",\n \"\\n\",\n \"@nnx.jit\\n\",\n \"def eval_step(model, x, y):\\n\",\n \" return ((model(x) - y) ** 2).mean()\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"70c05c4d\",\n \"metadata\": {},\n \"source\": [\n \"Now we create `train_model` and `eval_model` views up front. During the training loop we simply use the appropriate view — there is no need to call `.train()` or `.eval()`, and it is always clear from the code which mode the model is in.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"175db8e7\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"model = MyModel(in_dim, hidden_dim, out_dim, 0.1, rngs=rngs)\\n\",\n \"optimizer = nnx.Optimizer(model, optax.adam(lr), wrt=nnx.Param)\\n\",\n \"train_model = nnx.view(model, deterministic=False) # training view\\n\",\n \"eval_model = nnx.view(model, deterministic=True) # eval view\\n\",\n \"\\n\",\n \"eval_results = []\\n\",\n \"for epoch in range(total_epochs):\\n\",\n \" for i in range(ndata // batch_size):\\n\",\n \" idx = slice(i * batch_size, (i + 1) * batch_size)\\n\",\n \" train_step(train_model, optimizer, x[idx], y[idx], rngs) # use train_model\\n\",\n \"\\n\",\n \" eval_results.append(eval_step(eval_model, x, y)) # use eval_model\\n\",\n \"plt.plot(eval_results)\\n\",\n \"plt.show()\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"3d666cdb\",\n \"metadata\": {},\n \"source\": [\n \"## Getting information with `nnx.view_info`\\n\",\n \"To see more information about the options for `nnx.view`, we can use the `nnx.view_info` function to display information about the arguments. This will display each submodule which contains a `set_view` method. It also provides information about the keyword arguments accepted by each submodule, including type information, default values, and docstring descriptions.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"54a65a31\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"print(nnx.view_info(model))\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"47479be6\",\n \"metadata\": {},\n \"source\": [\n \"## Writing modules compatible with `nnx.view`\\n\",\n \"\\n\",\n \"You can make any custom module work with `nnx.view` by defining a `set_view` method. When `nnx.view` is called, it traverses the module tree and calls `set_view` on every submodule that defines one. `nnx.view` inspects the signature of each `set_view` method and only passes the keyword arguments that match the method's declared parameters. This means each module only receives the kwargs it cares about.\\n\",\n \"\\n\",\n \"Your `set_view` method should follow these conventions:\\n\",\n \"\\n\",\n \"1. **Accept keyword arguments with `None` defaults.** Each kwarg represents a configurable mode for this module. A `None` default means \\\"leave unchanged\\\", so views only override the modes you explicitly set.\\n\",\n \"2. **Only update the attribute when the kwarg is not `None`.** This ensures that unrelated views don't accidentally reset each other's settings.\\n\",\n \"3. **Include a Google-style docstring.** The `nnx.view_info` function parses these docstrings to display human-readable information about available view options.\\n\",\n \"\\n\",\n \"The general pattern looks like this:\\n\",\n \"\\n\",\n \"```python\\n\",\n \"class MyLayer(nnx.Module):\\n\",\n \" ...\\n\",\n \"\\n\",\n \" def set_view(self, kwarg1: type1 = None, ..., kwargN: typeN = None):\\n\",\n \" \\\"\\\"\\\"Description of the module's configurable modes.\\n\",\n \"\\n\",\n \" Args:\\n\",\n \" kwarg1: description of kwarg1.\\n\",\n \" ...\\n\",\n \" kwargN: description of kwargN.\\n\",\n \" \\\"\\\"\\\"\\n\",\n \" if kwarg1 is not None:\\n\",\n \" self.kwarg1 = kwarg1\\n\",\n \" ...\\n\",\n \"```\\n\",\n \"\\n\",\n \"Here is a concrete example — a `PrintLayer` that can be toggled to print a message during its forward pass:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"2dfdfd64\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"class PrintLayer(nnx.Module):\\n\",\n \" def __init__(self, msg: str | None = None):\\n\",\n \" self.msg = msg\\n\",\n \"\\n\",\n \" def __call__(self, *args, **kwargs):\\n\",\n \" if self.msg:\\n\",\n \" print(self.msg)\\n\",\n \"\\n\",\n \" def set_view(self, msg: bool | None = None):\\n\",\n \" \\\"\\\"\\\"Example set_view docstring. This follows Google style docstrings.\\n\",\n \"\\n\",\n \" Args:\\n\",\n \" msg: bool indicating if a message should be printed.\\n\",\n \" If True, the `__call__` method prints the message.\\n\",\n \" \\\"\\\"\\\"\\n\",\n \" if msg is not None:\\n\",\n \" self.msg = msg\\n\",\n \"\\n\",\n \"\\n\",\n \"model = PrintLayer()\\n\",\n \"model_print = nnx.view(model, msg='Hello, World!')\\n\",\n \"\\n\",\n \"model() # nothing printed\\n\",\n \"model_print() # prints \\\"Hello, World!\\\"\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"c7b261b8\",\n \"metadata\": {},\n \"source\": [\n \"We can use `nnx.view_info` to inspect what view options `PrintLayer` exposes. This is especially handy when working with unfamiliar models — it lists every submodule that defines `set_view`, along with the accepted kwargs, their types, defaults, and docstring descriptions.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"a5e3bc03\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# Display the information for nnx.view\\n\",\n \"print(nnx.view_info(model))\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"1acbcc09\",\n \"metadata\": {},\n \"source\": [\n \"The output shows that `PrintLayer` accepts a `msg` kwarg of type `bool` in its `set_view` method. When building larger models composed of many custom submodules, `nnx.view_info` gives you a quick summary of all the configurable modes across the entire module tree.\\n\",\n \"\\n\",\n \"## Using `with_attributes`\\n\",\n \"\\n\",\n \"If you are working with modules that don't implement the `set_view` API, you can use {func}`nnx.with_attributes ` to create views by directly replacing their attributes. Like `nnx.view`, it returns a new instance that shares jax arrays with the original, leaving the original unchanged.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"62b5c185\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"class NoisyLinear(nnx.Module):\\n\",\n \" def __init__(self, din, dout, *, training=None, rngs: nnx.Rngs):\\n\",\n \" self.linear = nnx.Linear(din, dout, rngs=rngs)\\n\",\n \" self.training = training\\n\",\n \"\\n\",\n \" def __call__(self, x, rngs=None):\\n\",\n \" assert self.training is not None\\n\",\n \" x = self.linear(x)\\n\",\n \" if self.training:\\n\",\n \" x = x + rngs.normal(x.shape) * 0.1\\n\",\n \" return x\\n\",\n \"\\n\",\n \"rngs = nnx.Rngs(0)\\n\",\n \"model = nnx.Sequential(\\n\",\n \" NoisyLinear(4, 8, rngs=rngs),\\n\",\n \" NoisyLinear(8, 2, rngs=rngs),\\n\",\n \")\\n\",\n \"\\n\",\n \"train_model = nnx.with_attributes(model, training=True)\\n\",\n \"eval_model = nnx.with_attributes(model, training=False)\\n\",\n \"\\n\",\n \"print(f'{train_model.layer1.training=}')\\n\",\n \"y1 = train_model(jnp.ones((1, 4)), rngs=rngs)\\n\",\n \"\\n\",\n \"print(f'{eval_model.layer1.training=}')\\n\",\n \"y2 = eval_model(jnp.ones((1, 4)))\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"0cc37a57\",\n \"metadata\": {},\n \"source\": [\n \"## Using `recursive_map`\\n\",\n \"\\n\",\n \"For more advanced transformations — such as replacing submodules — you can use {func}`nnx.recursive_map `. This function traverses the entire module tree bottom-up, calling a user-defined function `f(path, node)` on every node and leaf. Whatever `f` returns is used as the replacement for that node in the new tree. The resulting model view shares the Variables with the original (unless instructed otherwise).\\n\",\n \"\\n\",\n \"In the example below, we use `recursive_map` to replace every `nnx.Linear` layer with a `NoisyLinear` version (reusing the class defined earlier) that adds random noise during training:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"2e77b49d\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"import jax.numpy as jnp\\n\",\n \"\\n\",\n \"def add_noise(path, node):\\n\",\n \" if isinstance(node, nnx.Linear):\\n\",\n \" noisy = nnx.eval_shape(\\n\",\n \" lambda: NoisyLinear(node.in_features, node.out_features, rngs=nnx.Rngs(0))\\n\",\n \" )\\n\",\n \" noisy.linear = node\\n\",\n \" return noisy\\n\",\n \" return node\\n\",\n \"\\n\",\n \"rngs = nnx.Rngs(0)\\n\",\n \"model = nnx.Sequential(\\n\",\n \" nnx.Linear(4, 8, rngs=rngs),\\n\",\n \" nnx.Linear(8, 2, rngs=rngs),\\n\",\n \")\\n\",\n \"\\n\",\n \"noisy_model = nnx.recursive_map(add_noise, model)\\n\",\n \"\\n\",\n \"y = noisy_model(jnp.ones((1, 4)), rngs=rngs)\\n\",\n \"print(noisy_model)s\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"bf521e45\",\n \"metadata\": {},\n \"source\": [\n \"Here `recursive_map` visited each node, and when it found an `nnx.Linear` instance it created a `NoisyLinear`, swapped in the original `Linear` as its inner layer, and returned it. The original `model` is unchanged and its weights are shared with `noisy_model`.\"\n ]\n }\n ],\n \"metadata\": {\n \"jupytext\": {\n \"formats\": \"ipynb,md:myst\"\n },\n \"language_info\": {\n \"name\": \"python\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 5\n}\n" }, { "path": "docs_nnx/guides/view.md", "content": "---\njupytext:\n formats: ipynb,md:myst\n text_representation:\n extension: .md\n format_name: myst\n format_version: 0.13\n jupytext_version: 1.13.8\n---\n\n# Model Views\nThis guide covers how to use the `nnx.view` function. This function is useful for handling state in layers like `Dropout` and `BatchNorm`, which behave differently in training and evaluation. Similar to `.view` for numpy arrays, `nnx.view` allows you to set modes of the model while still sharing the same data. For a quick intro to how this function works, refer to the following example:\n\n```{code-cell}\nfrom flax import nnx\n\n# example model with different train/eval behavior\nrngs = nnx.Rngs(0)\nmodel = nnx.Sequential(\n nnx.Linear(2, 4, rngs=rngs), nnx.BatchNorm(4, rngs=rngs), nnx.Dropout(0.1)\n)\n\n# set train and eval modes\ntrain_model = nnx.view(model, deterministic=False, use_running_average=False)\neval_model = nnx.view(model, deterministic=True, use_running_average=True)\n\n# Can see deterministic is different between train_model and eval_model\nassert train_model.layers[2].deterministic == False\nassert eval_model.layers[2].deterministic == True\n\n# Weights are shared between the models\nassert train_model.layers[0].kernel is eval_model.layers[0].kernel\n\n# Print information about kwargs for nnx.view with nnx.view_info\nprint(nnx.view_info(model))\n```\n\n## Motivation\n\nSome layers in ML inherently involve state. Consider for example the `nnx.Dropout` layer, which behaves differently during training and evaluation. In these different scenarios, we need a simple way to ensure that the model behaves as intended to avoid silent bugs. A common pattern in other frameworks is to mutate a single `model` object to switch between training and evaluation modes. This requires the programmer to remember to toggle modes in many places throughout the code, which can hurt readability and lead to subtle bugs when a mode switch is forgotten.\n\n`nnx.view` offers a cleaner alternative: you declare the different model configurations once at the beginning of your code and then simply use the appropriate view wherever needed. Each view shares the same underlying weights, so parameter updates are automatically reflected across all views. We demonstrate this with a simple example below.\n\n```{code-cell}\nimport jax\nimport optax\nimport matplotlib.pyplot as plt\n\nin_dim, hidden_dim, out_dim = 16, 32, 2\n\n\nclass MyModel(nnx.Module):\n def __init__(\n self,\n in_dim: int,\n hidden_dim: int,\n out_dim: int,\n dropout_rate: float,\n *,\n rngs: nnx.Rngs,\n ):\n self.lin1 = nnx.Linear(in_dim, hidden_dim, rngs=rngs)\n self.do = nnx.Dropout(dropout_rate)\n self.bn = nnx.BatchNorm(hidden_dim, rngs=rngs)\n self.lin2 = nnx.Linear(hidden_dim, out_dim, rngs=rngs)\n\n def __call__(self, x, *, rngs=None):\n x = nnx.relu(self.do(self.bn(self.lin1(x)), rngs=rngs))\n return self.lin2(x)\n```\n\nLets take a look at the model to see what is going on.\n\n```{code-cell}\n# can display to inspect state\nmodel = MyModel(in_dim, hidden_dim, out_dim, 0.1, rngs=nnx.Rngs(0))\nnnx.display(model)\n\n# can assert to inspect state\nassert model.do.deterministic == False\n```\n\nFrom the model display, we can see that `Dropout` has `deterministic == False`, suggesting that the model is in training mode. In order to know this, we had to display the model and/or know that `Dropout` is set to training mode by default. It is not clear what state the model is in just by looking at the code without additional inspection. We instead want to be very explicit about what state the model is in. \n\nThis is where `nnx.view` comes in. This function updates the modes for each submodule of a neural network based on the kwargs passed into the function. The underlying model weights are then shared between different views. We set up a training and evaluation version of the model below.\n\n```{code-cell}\ntrain_model = nnx.view(model, deterministic=False)\neval_model = nnx.view(model, deterministic=True)\n\n# weights are references to the same data\nassert train_model.lin1.kernel is eval_model.lin1.kernel\n\n# Dropout.deterministic is different in each model\nassert train_model.do.deterministic is False\nassert eval_model.do.deterministic is True\n```\n\n## Example with `nnx.view`\n\n+++\n\nWe first set up data generators and define train/eval step functions. The `train_step` receives an `nnx.Rngs` object for dropout randomness, while `eval_step` doesn't since dropout is disabled in `eval_model`.\n\n```{code-cell}\nndata, batch_size, total_epochs, lr = 2048, 32, 100, 1e-3\nrngs = nnx.Rngs(0)\nx = rngs.normal((ndata, in_dim))\ny = rngs.normal((ndata, out_dim))\n\n\n@nnx.jit\ndef train_step(model, optimizer, x, y, rngs):\n def loss_fn(model, rngs):\n return ((model(x, rngs=rngs) - y) ** 2).mean()\n\n grads = nnx.grad(loss_fn)(model, rngs)\n optimizer.update(model, grads)\n\n\n@nnx.jit\ndef eval_step(model, x, y):\n return ((model(x) - y) ** 2).mean()\n```\n\nNow we create `train_model` and `eval_model` views up front. During the training loop we simply use the appropriate view — there is no need to call `.train()` or `.eval()`, and it is always clear from the code which mode the model is in.\n\n```{code-cell}\nmodel = MyModel(in_dim, hidden_dim, out_dim, 0.1, rngs=rngs)\noptimizer = nnx.Optimizer(model, optax.adam(lr), wrt=nnx.Param)\ntrain_model = nnx.view(model, deterministic=False) # training view\neval_model = nnx.view(model, deterministic=True) # eval view\n\neval_results = []\nfor epoch in range(total_epochs):\n for i in range(ndata // batch_size):\n idx = slice(i * batch_size, (i + 1) * batch_size)\n train_step(train_model, optimizer, x[idx], y[idx], rngs) # use train_model\n\n eval_results.append(eval_step(eval_model, x, y)) # use eval_model\nplt.plot(eval_results)\nplt.show()\n```\n\n## Getting information with `nnx.view_info`\nTo see more information about the options for `nnx.view`, we can use the `nnx.view_info` function to display information about the arguments. This will display each submodule which contains a `set_view` method. It also provides information about the keyword arguments accepted by each submodule, including type information, default values, and docstring descriptions.\n\n```{code-cell}\nprint(nnx.view_info(model))\n```\n\n## Writing modules compatible with `nnx.view`\n\nYou can make any custom module work with `nnx.view` by defining a `set_view` method. When `nnx.view` is called, it traverses the module tree and calls `set_view` on every submodule that defines one. `nnx.view` inspects the signature of each `set_view` method and only passes the keyword arguments that match the method's declared parameters. This means each module only receives the kwargs it cares about.\n\nYour `set_view` method should follow these conventions:\n\n1. **Accept keyword arguments with `None` defaults.** Each kwarg represents a configurable mode for this module. A `None` default means \"leave unchanged\", so views only override the modes you explicitly set.\n2. **Only update the attribute when the kwarg is not `None`.** This ensures that unrelated views don't accidentally reset each other's settings.\n3. **Include a Google-style docstring.** The `nnx.view_info` function parses these docstrings to display human-readable information about available view options.\n\nThe general pattern looks like this:\n\n```python\nclass MyLayer(nnx.Module):\n ...\n\n def set_view(self, kwarg1: type1 = None, ..., kwargN: typeN = None):\n \"\"\"Description of the module's configurable modes.\n\n Args:\n kwarg1: description of kwarg1.\n ...\n kwargN: description of kwargN.\n \"\"\"\n if kwarg1 is not None:\n self.kwarg1 = kwarg1\n ...\n```\n\nHere is a concrete example — a `PrintLayer` that can be toggled to print a message during its forward pass:\n\n```{code-cell}\nclass PrintLayer(nnx.Module):\n def __init__(self, msg: str | None = None):\n self.msg = msg\n\n def __call__(self, *args, **kwargs):\n if self.msg:\n print(self.msg)\n\n def set_view(self, msg: bool | None = None):\n \"\"\"Example set_view docstring. This follows Google style docstrings.\n\n Args:\n msg: bool indicating if a message should be printed.\n If True, the `__call__` method prints the message.\n \"\"\"\n if msg is not None:\n self.msg = msg\n\n\nmodel = PrintLayer()\nmodel_print = nnx.view(model, msg='Hello, World!')\n\nmodel() # nothing printed\nmodel_print() # prints \"Hello, World!\"\n```\n\nWe can use `nnx.view_info` to inspect what view options `PrintLayer` exposes. This is especially handy when working with unfamiliar models — it lists every submodule that defines `set_view`, along with the accepted kwargs, their types, defaults, and docstring descriptions.\n\n```{code-cell}\n# Display the information for nnx.view\nprint(nnx.view_info(model))\n```\n\nThe output shows that `PrintLayer` accepts a `msg` kwarg of type `bool` in its `set_view` method. When building larger models composed of many custom submodules, `nnx.view_info` gives you a quick summary of all the configurable modes across the entire module tree.\n\n## Using `with_attributes`\n\nIf you are working with modules that don't implement the `set_view` API, you can use {func}`nnx.with_attributes ` to create views by directly replacing their attributes. Like `nnx.view`, it returns a new instance that shares jax arrays with the original, leaving the original unchanged.\n\n```{code-cell}\nclass NoisyLinear(nnx.Module):\n def __init__(self, din, dout, *, training=None, rngs: nnx.Rngs):\n self.linear = nnx.Linear(din, dout, rngs=rngs)\n self.training = training\n\n def __call__(self, x, rngs=None):\n assert self.training is not None\n x = self.linear(x)\n if self.training:\n x = x + rngs.normal(x.shape) * 0.1\n return x\n\nrngs = nnx.Rngs(0)\nmodel = nnx.Sequential(\n NoisyLinear(4, 8, rngs=rngs),\n NoisyLinear(8, 2, rngs=rngs),\n)\n\ntrain_model = nnx.with_attributes(model, training=True)\neval_model = nnx.with_attributes(model, training=False)\n\nprint(f'{train_model.layer1.training=}')\ny1 = train_model(jnp.ones((1, 4)), rngs=rngs)\n\nprint(f'{eval_model.layer1.training=}')\ny2 = eval_model(jnp.ones((1, 4)))\n```\n\n## Using `recursive_map`\n\nFor more advanced transformations — such as replacing submodules — you can use {func}`nnx.recursive_map `. This function traverses the entire module tree bottom-up, calling a user-defined function `f(path, node)` on every node and leaf. Whatever `f` returns is used as the replacement for that node in the new tree. The resulting model view shares the Variables with the original (unless instructed otherwise).\n\nIn the example below, we use `recursive_map` to replace every `nnx.Linear` layer with a `NoisyLinear` version (reusing the class defined earlier) that adds random noise during training:\n\n```{code-cell}\nimport jax.numpy as jnp\n\ndef add_noise(path, node):\n if isinstance(node, nnx.Linear):\n noisy = nnx.eval_shape(\n lambda: NoisyLinear(node.in_features, node.out_features, rngs=nnx.Rngs(0))\n )\n noisy.linear = node\n return noisy\n return node\n\nrngs = nnx.Rngs(0)\nmodel = nnx.Sequential(\n nnx.Linear(4, 8, rngs=rngs),\n nnx.Linear(8, 2, rngs=rngs),\n)\n\nnoisy_model = nnx.recursive_map(add_noise, model)\n\ny = noisy_model(jnp.ones((1, 4)), rngs=rngs)\nprint(noisy_model)s\n```\n\nHere `recursive_map` visited each node, and when it found an `nnx.Linear` instance it created a `NoisyLinear`, swapped in the original `Linear` as its inner layer, and returned it. The original `model` is unchanged and its weights are shared with `noisy_model`.\n" }, { "path": "docs_nnx/guides_advanced.rst", "content": "Advanced Guides\n======\n\n.. toctree::\n :maxdepth: 1\n :caption: Advanced\n\n guides/flax_gspmd\n guides/performance\n guides/bridge_guide\n guides/surgery\n guides/extracting_intermediates\n" }, { "path": "docs_nnx/guides_basic.rst", "content": "Basic Guides\n======\n\n.. toctree::\n :maxdepth: 1\n :caption: Basic\n\n guides/pytree\n guides/transforms\n guides/view\n guides/filters_guide\n guides/randomness\n guides/checkpointing\n guides/jax_and_nnx_transforms" }, { "path": "docs_nnx/hijax/hijax.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"id\": \"15c2d208\",\n \"metadata\": {},\n \"source\": [\n \"# Hijax\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 1,\n \"id\": \"99809892\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"from flax import nnx\\n\",\n \"import jax\\n\",\n \"import jax.numpy as jnp\\n\",\n \"import optax\\n\",\n \"\\n\",\n \"current_mode = nnx.var_defaults().hijax # ignore: only needed for testing\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 2,\n \"id\": \"d1aaa0ec\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"0.85250294\\n\",\n \"0.8165137\\n\",\n \"0.7814907\\n\"\n ]\n }\n ],\n \"source\": [\n \"nnx.var_defaults(hijax=True)\\n\",\n \"\\n\",\n \"rngs = nnx.Rngs(0)\\n\",\n \"model = nnx.Linear(2, 3, rngs=rngs)\\n\",\n \"optimizer = nnx.Optimizer(model, optax.adamw(1e-2), wrt=nnx.Param)\\n\",\n \"\\n\",\n \"@jax.jit\\n\",\n \"def train_step(x, y):\\n\",\n \" loss_fn = lambda m: jnp.mean((m(x) - y) ** 2)\\n\",\n \" loss, grads = jax.value_and_grad(loss_fn)(nnx.vars_as(model, mutable=False)) # tmp fix for jax.grad\\n\",\n \" optimizer.update(model, grads)\\n\",\n \" return loss\\n\",\n \"\\n\",\n \"x, y = rngs.uniform((4, 2)), rngs.uniform((4, 3))\\n\",\n \"for _ in range(3):\\n\",\n \" print(train_step(x, y))\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"04458d66\",\n \"metadata\": {},\n \"source\": [\n \"## Hijax Variable\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"f4220c6f\",\n \"metadata\": {},\n \"source\": [\n \"State propagation:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 3,\n \"id\": \"396a07a3\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"0\\n\",\n \"1\\n\"\n ]\n }\n ],\n \"source\": [\n \"v = nnx.Variable(jnp.array(0), hijax=True)\\n\",\n \"\\n\",\n \"@jax.jit\\n\",\n \"def inc(v):\\n\",\n \" v[...] += 1\\n\",\n \"\\n\",\n \"print(v[...]); inc(v); print(v[...])\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 4,\n \"id\": \"2ab7d801\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"{ \\u001b[34;1mlambda \\u001b[39;22m; a\\u001b[35m:Variable()\\u001b[39m. \\u001b[34;1mlet\\n\",\n \" \\u001b[39;22mjit[\\n\",\n \" name=inc\\n\",\n \" jaxpr={ \\u001b[34;1mlambda \\u001b[39;22m; a\\u001b[35m:Variable()\\u001b[39m. \\u001b[34;1mlet\\n\",\n \" \\u001b[39;22mb\\u001b[35m:i32[]\\u001b[39m = get_variable[\\n\",\n \" avals=(ShapedArray(int32[], weak_type=True),)\\n\",\n \" has_qdd=True\\n\",\n \" treedef=PyTreeDef(CustomNode(Variable[(('eager_sharding', True), ('hijax', True), ('mutable', True), ('ref', False))], [*]))\\n\",\n \" var_type=\\n\",\n \" ] a\\n\",\n \" c\\u001b[35m:i32[]\\u001b[39m = add b 1:i32[]\\n\",\n \" _\\u001b[35m:i32[]\\u001b[39m = get_variable[\\n\",\n \" avals=(ShapedArray(int32[], weak_type=True),)\\n\",\n \" has_qdd=True\\n\",\n \" treedef=PyTreeDef(CustomNode(Variable[(('eager_sharding', True), ('hijax', True), ('mutable', True), ('ref', False))], [*]))\\n\",\n \" var_type=\\n\",\n \" ] a\\n\",\n \" set_variable[\\n\",\n \" treedef=PyTreeDef(CustomNode(Variable[(('eager_sharding', True), ('hijax', True), ('mutable', True), ('ref', False))], [*]))\\n\",\n \" var_type=\\n\",\n \" ] a c\\n\",\n \" \\u001b[34;1min \\u001b[39;22m() }\\n\",\n \" ] a\\n\",\n \" \\u001b[34;1min \\u001b[39;22m() }\\n\"\n ]\n }\n ],\n \"source\": [\n \"v = nnx.Variable(jnp.array(0), hijax=True)\\n\",\n \"print(jax.make_jaxpr(inc)(v))\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"39070460\",\n \"metadata\": {},\n \"source\": [\n \"Pytree values:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 5,\n \"id\": \"fcd0de3f\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"\\u001b[38;2;79;201;177mVariable\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;105;105;105m # 2 (8 B)\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254mvalue\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;255;213;3m{\\u001b[0m\\u001b[38;2;207;144;120m'a'\\u001b[0m: Array(0, dtype=int32, weak_type=True), \\u001b[38;2;207;144;120m'b'\\u001b[0m: Array(2, dtype=int32, weak_type=True)\\u001b[38;2;255;213;3m}\\u001b[0m,\\n\",\n \" \\u001b[38;2;156;220;254mhijax\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;86;156;214mTrue\\u001b[0m\\n\",\n \"\\u001b[38;2;255;213;3m)\\u001b[0m\\n\",\n \"\\u001b[38;2;79;201;177mVariable\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;105;105;105m # 2 (8 B)\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254mvalue\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;255;213;3m{\\u001b[0m\\u001b[38;2;207;144;120m'a'\\u001b[0m: Array(1, dtype=int32, weak_type=True), \\u001b[38;2;207;144;120m'b'\\u001b[0m: Array(4, dtype=int32, weak_type=True)\\u001b[38;2;255;213;3m}\\u001b[0m,\\n\",\n \" \\u001b[38;2;156;220;254mhijax\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;86;156;214mTrue\\u001b[0m\\n\",\n \"\\u001b[38;2;255;213;3m)\\u001b[0m\\n\"\n ]\n }\n ],\n \"source\": [\n \"v = nnx.Variable({'a': jnp.array(0), 'b': jnp.array(2)}, hijax=True)\\n\",\n \"\\n\",\n \"@jax.jit\\n\",\n \"def inc_and_double(v):\\n\",\n \" v['a'] += 1\\n\",\n \" v['b'] *= 2\\n\",\n \"\\n\",\n \"print(v); inc_and_double(v); print(v)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"f0cfe954\",\n \"metadata\": {},\n \"source\": [\n \"Dynamic state structure:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 6,\n \"id\": \"0d83a130\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"Before: \\u001b[38;2;79;201;177mVariable\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;105;105;105m\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254mvalue\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;255;213;3m{\\u001b[0m\\u001b[38;2;255;213;3m}\\u001b[0m,\\n\",\n \" \\u001b[38;2;156;220;254mhijax\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;86;156;214mTrue\\u001b[0m\\n\",\n \"\\u001b[38;2;255;213;3m)\\u001b[0m\\n\",\n \"After: \\u001b[38;2;79;201;177mVariable\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;105;105;105m # 1 (4 B)\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254mvalue\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;255;213;3m{\\u001b[0m\\u001b[38;2;207;144;120m'y_mean'\\u001b[0m: Array(-1.1782329, dtype=float32)\\u001b[38;2;255;213;3m}\\u001b[0m,\\n\",\n \" \\u001b[38;2;156;220;254mhijax\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;86;156;214mTrue\\u001b[0m\\n\",\n \"\\u001b[38;2;255;213;3m)\\u001b[0m\\n\"\n ]\n }\n ],\n \"source\": [\n \"rngs = nnx.Rngs(0)\\n\",\n \"x = rngs.uniform((4, 5))\\n\",\n \"w = rngs.normal((5, 3))\\n\",\n \"metrics = nnx.Variable({}, hijax=True)\\n\",\n \"\\n\",\n \"@jax.jit\\n\",\n \"def linear(x, w, metrics: nnx.Variable):\\n\",\n \" y = x @ w\\n\",\n \" metrics['y_mean'] = jnp.mean(y)\\n\",\n \" return y\\n\",\n \"\\n\",\n \"print(\\\"Before:\\\", metrics)\\n\",\n \"y = linear(x, w, metrics)\\n\",\n \"print(\\\"After:\\\", metrics)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 7,\n \"id\": \"0a55df94\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"\\u001b[38;2;79;201;177mVariable\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;105;105;105m # 3 (12 B)\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254mvalue\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0mArray([1, 2, 3], dtype=int32),\\n\",\n \" \\u001b[38;2;156;220;254mhijax\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;86;156;214mTrue\\u001b[0m\\n\",\n \"\\u001b[38;2;255;213;3m)\\u001b[0m\\n\"\n ]\n }\n ],\n \"source\": [\n \"# set default Variable mode for the rest of the guide\\n\",\n \"nnx.var_defaults(hijax=True)\\n\",\n \"\\n\",\n \"variable = nnx.Variable(jnp.array([1, 2, 3]))\\n\",\n \"\\n\",\n \"print(variable)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"1b2632f1\",\n \"metadata\": {},\n \"source\": [\n \"### Mutability\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 8,\n \"id\": \"b7b1f421\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"nnx.vars_as(model, mutable=False) = \\u001b[38;2;79;201;177mLinear\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;105;105;105m # Param: 3 (12 B)\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254mkernel\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;79;201;177mParam\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;105;105;105m # 3 (12 B)\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254mvalue\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;79;201;177mArray\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;156;220;254mshape\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;182;207;169m1\\u001b[0m, \\u001b[38;2;182;207;169m3\\u001b[0m\\u001b[38;2;255;213;3m)\\u001b[0m, \\u001b[38;2;156;220;254mdtype\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0mdtype('float32')\\u001b[38;2;255;213;3m)\\u001b[0m,\\n\",\n \" \\u001b[38;2;156;220;254mhijax\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;86;156;214mTrue\\u001b[0m,\\n\",\n \" \\u001b[38;2;156;220;254mmutable\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;86;156;214mFalse\\u001b[0m\\n\",\n \" \\u001b[38;2;255;213;3m)\\u001b[0m\\n\",\n \"\\u001b[38;2;255;213;3m)\\u001b[0m\\n\",\n \"nnx.vars_as(model, mutable=True) = \\u001b[38;2;79;201;177mLinear\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;105;105;105m # Param: 3 (12 B)\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254mkernel\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;79;201;177mParam\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;105;105;105m # 3 (12 B)\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254mvalue\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;79;201;177mArray\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;156;220;254mshape\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;182;207;169m1\\u001b[0m, \\u001b[38;2;182;207;169m3\\u001b[0m\\u001b[38;2;255;213;3m)\\u001b[0m, \\u001b[38;2;156;220;254mdtype\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0mdtype('float32')\\u001b[38;2;255;213;3m)\\u001b[0m,\\n\",\n \" \\u001b[38;2;156;220;254mhijax\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;86;156;214mTrue\\u001b[0m\\n\",\n \" \\u001b[38;2;255;213;3m)\\u001b[0m\\n\",\n \"\\u001b[38;2;255;213;3m)\\u001b[0m\\n\"\n ]\n }\n ],\n \"source\": [\n \"class Linear(nnx.Module):\\n\",\n \" def __init__(self, in_features, out_features, rngs: nnx.Rngs):\\n\",\n \" self.kernel = nnx.Param(rngs.normal((in_features, out_features)))\\n\",\n \"\\n\",\n \" def __call__(self, x):\\n\",\n \" return x @ self.kernel\\n\",\n \"\\n\",\n \"model = Linear(1, 3, rngs=nnx.Rngs(0))\\n\",\n \"\\n\",\n \"print(f\\\"{nnx.vars_as(model, mutable=False) = !s}\\\")\\n\",\n \"print(f\\\"{nnx.vars_as(model, mutable=True) = !s}\\\")\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 9,\n \"id\": \"594cb65e\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"ImmutableVariableError: Cannot mutate Variable as it is marked as immutable. (https://flax.readthedocs.io/en/latest/api_reference/flax.errors.html#flax.errors.ImmutableVariableError)\\n\"\n ]\n }\n ],\n \"source\": [\n \"v = nnx.Variable(jnp.array(0))\\n\",\n \"v_immut = nnx.vars_as(v, mutable=False)\\n\",\n \"assert not v_immut.mutable\\n\",\n \"\\n\",\n \"try:\\n\",\n \" v_immut[...] += 1 # raises an error\\n\",\n \"except Exception as e:\\n\",\n \" print(f\\\"{type(e).__name__}: {e}\\\")\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"58692a37\",\n \"metadata\": {},\n \"source\": [\n \"### Ref support\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 10,\n \"id\": \"fcd4fb4f\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"\\u001b[38;2;79;201;177mVariable\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;105;105;105m # 1 (4 B)\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254mvalue\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0mArray(0, dtype=int32, weak_type=True),\\n\",\n \" \\u001b[38;2;156;220;254mhijax\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;86;156;214mTrue\\u001b[0m,\\n\",\n \" \\u001b[38;2;156;220;254mref\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;86;156;214mTrue\\u001b[0m\\n\",\n \"\\u001b[38;2;255;213;3m)\\u001b[0m\\n\",\n \"Ref(0, dtype=int32, weak_type=True)\\n\"\n ]\n }\n ],\n \"source\": [\n \"v = nnx.Variable(jnp.array(0))\\n\",\n \"v_ref = nnx.vars_as(v, ref=True)\\n\",\n \"assert v_ref.ref\\n\",\n \"print(v_ref)\\n\",\n \"print(v_ref.get_raw_value())\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 11,\n \"id\": \"18256668\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"immutable = \\u001b[38;2;79;201;177mVariable\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;105;105;105m # 1 (4 B)\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254mvalue\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0mArray(0, dtype=int32, weak_type=True),\\n\",\n \" \\u001b[38;2;156;220;254mhad_ref\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;86;156;214mTrue\\u001b[0m,\\n\",\n \" \\u001b[38;2;156;220;254mhijax\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;86;156;214mTrue\\u001b[0m,\\n\",\n \" \\u001b[38;2;156;220;254mmutable\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;86;156;214mFalse\\u001b[0m\\n\",\n \"\\u001b[38;2;255;213;3m)\\u001b[0m\\n\",\n \"mutable = \\u001b[38;2;79;201;177mVariable\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;105;105;105m # 1 (4 B)\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254mvalue\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0mArray(0, dtype=int32, weak_type=True),\\n\",\n \" \\u001b[38;2;156;220;254mhijax\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;86;156;214mTrue\\u001b[0m,\\n\",\n \" \\u001b[38;2;156;220;254mref\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;86;156;214mTrue\\u001b[0m\\n\",\n \"\\u001b[38;2;255;213;3m)\\u001b[0m\\n\"\n ]\n }\n ],\n \"source\": [\n \"v_immut = nnx.vars_as(v_ref, mutable=False)\\n\",\n \"assert not v_immut.ref\\n\",\n \"print(\\\"immutable =\\\", v_immut)\\n\",\n \"\\n\",\n \"v_ref = nnx.vars_as(v_immut, mutable=True)\\n\",\n \"assert v_ref.ref\\n\",\n \"print(\\\"mutable =\\\", v_ref)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"f4e35e75\",\n \"metadata\": {},\n \"source\": [\n \"### Examples\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 12,\n \"id\": \"5400fe58\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"class Block(nnx.Module):\\n\",\n \" def __init__(self, din, dmid, dout, rngs: nnx.Rngs):\\n\",\n \" self.linear = Linear(din, dmid, rngs=rngs)\\n\",\n \" self.bn = nnx.BatchNorm(dmid, rngs=rngs)\\n\",\n \" self.dropout = nnx.Dropout(0.1, rngs=rngs)\\n\",\n \" self.linear_out = Linear(dmid, dout, rngs=rngs)\\n\",\n \"\\n\",\n \" def __call__(self, x):\\n\",\n \" x = nnx.gelu(self.dropout(self.bn(self.linear(x))))\\n\",\n \" return self.linear_out(x)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"ba980b6b\",\n \"metadata\": {},\n \"source\": [\n \"#### Training Loop\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 13,\n \"id\": \"566c4249\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"loss = 1.000178\\n\",\n \"loss = 0.9700456\\n\",\n \"loss = 0.93967044\\n\"\n ]\n }\n ],\n \"source\": [\n \"# hijax Variables by default\\n\",\n \"model = Block(2, 64, 3, rngs=nnx.Rngs(0))\\n\",\n \"optimizer = nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param)\\n\",\n \"\\n\",\n \"@jax.jit\\n\",\n \"def train_step(model, optimizer, x, y):\\n\",\n \" graphdef, params, nondiff = nnx.split(model, nnx.Param, ...)\\n\",\n \" def loss_fn(params):\\n\",\n \" model = nnx.merge(graphdef, params, nondiff)\\n\",\n \" return ((model(x) - y) ** 2).mean()\\n\",\n \"\\n\",\n \" loss, grads = jax.value_and_grad(loss_fn)(nnx.vars_as(params, mutable=False)) # immutable for jax.grad\\n\",\n \" optimizer.update(model, grads)\\n\",\n \"\\n\",\n \" return loss\\n\",\n \"\\n\",\n \"for _ in range(3):\\n\",\n \" loss = train_step(model, optimizer, x=jnp.ones((10, 2)), y=jnp.ones((10, 3)))\\n\",\n \" print(f\\\"{loss = !s}\\\")\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"1dea99c1\",\n \"metadata\": {},\n \"source\": [\n \"#### Scan Over Layers\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 14,\n \"id\": \"d8136be4\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# TODO: does not work with hijax yet\\n\",\n \"# @jax.vmap\\n\",\n \"# def create_stack(rngs):\\n\",\n \"# return nnx.as_immutable_vars(Block(2, 64, 2, rngs=rngs))\\n\",\n \"\\n\",\n \"# block_stack = nnx.as_mutable_vars(create_stack(nnx.Rngs(0).fork(split=8)))\\n\",\n \"\\n\",\n \"# def scan_fn(x, block):\\n\",\n \"# x = block(x)\\n\",\n \"# return x, None\\n\",\n \"\\n\",\n \"# x = jax.random.uniform(jax.random.key(0), (3, 2))\\n\",\n \"# y, _ = jax.lax.scan(scan_fn, x, block_stack)\\n\",\n \"\\n\",\n \"# print(\\\"y = \\\", y)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"7ca18a0d\",\n \"metadata\": {},\n \"source\": [\n \"### Limitations\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"1dd39c79\",\n \"metadata\": {},\n \"source\": [\n \"#### Mutable Outputs\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 15,\n \"id\": \"c6062d19\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"Error: mutable hitypes should use lo_ty_qdd instead\\n\"\n ]\n }\n ],\n \"source\": [\n \"@jax.jit\\n\",\n \"def create_model(rngs):\\n\",\n \" return Block(2, 64, 3, rngs=rngs)\\n\",\n \"\\n\",\n \"try:\\n\",\n \" model = create_model(nnx.Rngs(0))\\n\",\n \"except Exception as e:\\n\",\n \" print(f\\\"Error:\\\", e)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 16,\n \"id\": \"8bb1e9e7\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"model.linear = \\u001b[38;2;79;201;177mLinear\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;105;105;105m # Param: 128 (512 B)\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254mkernel\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;79;201;177mParam\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;105;105;105m # 128 (512 B)\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254mvalue\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;79;201;177mArray\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;156;220;254mshape\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;182;207;169m2\\u001b[0m, \\u001b[38;2;182;207;169m64\\u001b[0m\\u001b[38;2;255;213;3m)\\u001b[0m, \\u001b[38;2;156;220;254mdtype\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0mdtype('float32')\\u001b[38;2;255;213;3m)\\u001b[0m,\\n\",\n \" \\u001b[38;2;156;220;254mhijax\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;86;156;214mTrue\\u001b[0m\\n\",\n \" \\u001b[38;2;255;213;3m)\\u001b[0m\\n\",\n \"\\u001b[38;2;255;213;3m)\\u001b[0m\\n\"\n ]\n }\n ],\n \"source\": [\n \"@jax.jit\\n\",\n \"def create_model(rngs):\\n\",\n \" return nnx.vars_as((Block(2, 64, 3, rngs=rngs)), hijax=False)\\n\",\n \"\\n\",\n \"model = nnx.vars_as(create_model(nnx.Rngs(0)), hijax=True)\\n\",\n \"\\n\",\n \"print(\\\"model.linear =\\\", model.linear)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"609bed7c\",\n \"metadata\": {},\n \"source\": [\n \"#### Reference Sharing (aliasing)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 17,\n \"id\": \"045d03c1\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"None\\n\"\n ]\n }\n ],\n \"source\": [\n \"# NOTE: doesn't currently fail on the jax side\\n\",\n \"def get_error(f, *args):\\n\",\n \" try:\\n\",\n \" return f(*args)\\n\",\n \" except Exception as e:\\n\",\n \" return f\\\"{type(e).__name__}: {e}\\\"\\n\",\n \"\\n\",\n \"x = nnx.Variable(jnp.array(0))\\n\",\n \"\\n\",\n \"@jax.jit\\n\",\n \"def f(a, b):\\n\",\n \" ...\\n\",\n \"\\n\",\n \"print(get_error(f, x, x))\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 18,\n \"id\": \"bc2e87e5\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"None\\n\",\n \"\\u001b[38;2;79;201;177mHasShared\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;105;105;105m # Variable: 1 (4 B)\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254ma\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;79;201;177mVariable\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;105;105;105m # 1 (4 B)\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254mvalue\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0mArray(0, dtype=int32, weak_type=True),\\n\",\n \" \\u001b[38;2;156;220;254mhijax\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;86;156;214mTrue\\u001b[0m\\n\",\n \" \\u001b[38;2;255;213;3m)\\u001b[0m,\\n\",\n \" \\u001b[38;2;156;220;254mb\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;79;201;177mVariable\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;105;105;105m # 1 (4 B)\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254mvalue\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0mArray(0, dtype=int32, weak_type=True),\\n\",\n \" \\u001b[38;2;156;220;254mhijax\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;86;156;214mTrue\\u001b[0m\\n\",\n \" \\u001b[38;2;255;213;3m)\\u001b[0m\\n\",\n \"\\u001b[38;2;255;213;3m)\\u001b[0m\\n\"\n ]\n }\n ],\n \"source\": [\n \"# NOTE: doesn't currently fail on the jax side\\n\",\n \"class HasShared(nnx.Pytree):\\n\",\n \" def __init__(self):\\n\",\n \" self.a = nnx.Variable(jnp.array(0))\\n\",\n \" self.b = self.a\\n\",\n \"\\n\",\n \"@jax.jit\\n\",\n \"def g(has_shared):\\n\",\n \" has_shared.a[...] = 5\\n\",\n \"\\n\",\n \"has_shared = HasShared()\\n\",\n \"\\n\",\n \"print(get_error(g, has_shared))\\n\",\n \"print(has_shared) # updates don't propagate\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 19,\n \"id\": \"6298f3d9\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"Duplicates found:\\n\",\n \"- [('a',), ('b',)]\\n\"\n ]\n }\n ],\n \"source\": [\n \"print(\\\"Duplicates found:\\\")\\n\",\n \"if (all_duplicates := nnx.find_duplicates(has_shared)):\\n\",\n \" for duplicates in all_duplicates:\\n\",\n \" print(\\\"-\\\", duplicates)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 20,\n \"id\": \"00854d38\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"\\u001b[38;2;79;201;177mHasShared\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;105;105;105m # Variable: 1 (4 B)\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254ma\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;79;201;177mVariable\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;105;105;105m # 1 (4 B)\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254mvalue\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0mArray(5, dtype=int32, weak_type=True),\\n\",\n \" \\u001b[38;2;156;220;254mhijax\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;86;156;214mTrue\\u001b[0m\\n\",\n \" \\u001b[38;2;255;213;3m)\\u001b[0m,\\n\",\n \" \\u001b[38;2;156;220;254mb\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;79;201;177mVariable\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;105;105;105m # 1 (4 B)\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254mvalue\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0mArray(5, dtype=int32, weak_type=True),\\n\",\n \" \\u001b[38;2;156;220;254mhijax\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;86;156;214mTrue\\u001b[0m\\n\",\n \" \\u001b[38;2;255;213;3m)\\u001b[0m\\n\",\n \"\\u001b[38;2;255;213;3m)\\u001b[0m\\n\"\n ]\n }\n ],\n \"source\": [\n \"@jax.jit\\n\",\n \"def h(graphdef, state):\\n\",\n \" has_shared = nnx.merge(graphdef, state)\\n\",\n \" has_shared.a[...] = 5\\n\",\n \"\\n\",\n \"graphdef, state = nnx.split(has_shared)\\n\",\n \"h(graphdef, state)\\n\",\n \"print(has_shared)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 21,\n \"id\": \"195296c8\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# clean up for CI tests\\n\",\n \"_ = nnx.var_defaults(hijax=current_mode)\"\n ]\n }\n ],\n \"metadata\": {\n \"jupytext\": {\n \"formats\": \"ipynb,md:myst\"\n },\n \"language_info\": {\n \"codemirror_mode\": {\n \"name\": \"ipython\",\n \"version\": 3\n },\n \"file_extension\": \".py\",\n \"mimetype\": \"text/x-python\",\n \"name\": \"python\",\n \"nbconvert_exporter\": \"python\",\n \"pygments_lexer\": \"ipython3\",\n \"version\": \"3.11.9\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 5\n}\n" }, { "path": "docs_nnx/hijax/hijax.md", "content": "---\njupytext:\n formats: ipynb,md:myst\n text_representation:\n extension: .md\n format_name: myst\n format_version: 0.13\n jupytext_version: 1.13.8\n---\n\n# Hijax\n\n```{code-cell} ipython3\nfrom flax import nnx\nimport jax\nimport jax.numpy as jnp\nimport optax\n\ncurrent_mode = nnx.var_defaults().hijax # ignore: only needed for testing\n```\n\n```{code-cell} ipython3\nnnx.var_defaults(hijax=True)\n\nrngs = nnx.Rngs(0)\nmodel = nnx.Linear(2, 3, rngs=rngs)\noptimizer = nnx.Optimizer(model, optax.adamw(1e-2), wrt=nnx.Param)\n\n@jax.jit\ndef train_step(x, y):\n loss_fn = lambda m: jnp.mean((m(x) - y) ** 2)\n loss, grads = jax.value_and_grad(loss_fn)(nnx.vars_as(model, mutable=False)) # tmp fix for jax.grad\n optimizer.update(model, grads)\n return loss\n\nx, y = rngs.uniform((4, 2)), rngs.uniform((4, 3))\nfor _ in range(3):\n print(train_step(x, y))\n```\n\n## Hijax Variable\n\n+++\n\nState propagation:\n\n```{code-cell} ipython3\nv = nnx.Variable(jnp.array(0), hijax=True)\n\n@jax.jit\ndef inc(v):\n v[...] += 1\n\nprint(v[...]); inc(v); print(v[...])\n```\n\n```{code-cell} ipython3\nv = nnx.Variable(jnp.array(0), hijax=True)\nprint(jax.make_jaxpr(inc)(v))\n```\n\nPytree values:\n\n```{code-cell} ipython3\nv = nnx.Variable({'a': jnp.array(0), 'b': jnp.array(2)}, hijax=True)\n\n@jax.jit\ndef inc_and_double(v):\n v['a'] += 1\n v['b'] *= 2\n\nprint(v); inc_and_double(v); print(v)\n```\n\nDynamic state structure:\n\n```{code-cell} ipython3\nrngs = nnx.Rngs(0)\nx = rngs.uniform((4, 5))\nw = rngs.normal((5, 3))\nmetrics = nnx.Variable({}, hijax=True)\n\n@jax.jit\ndef linear(x, w, metrics: nnx.Variable):\n y = x @ w\n metrics['y_mean'] = jnp.mean(y)\n return y\n\nprint(\"Before:\", metrics)\ny = linear(x, w, metrics)\nprint(\"After:\", metrics)\n```\n\n```{code-cell} ipython3\n# set default Variable mode for the rest of the guide\nnnx.var_defaults(hijax=True)\n\nvariable = nnx.Variable(jnp.array([1, 2, 3]))\n\nprint(variable)\n```\n\n### Mutability\n\n```{code-cell} ipython3\nclass Linear(nnx.Module):\n def __init__(self, in_features, out_features, rngs: nnx.Rngs):\n self.kernel = nnx.Param(rngs.normal((in_features, out_features)))\n\n def __call__(self, x):\n return x @ self.kernel\n\nmodel = Linear(1, 3, rngs=nnx.Rngs(0))\n\nprint(f\"{nnx.vars_as(model, mutable=False) = !s}\")\nprint(f\"{nnx.vars_as(model, mutable=True) = !s}\")\n```\n\n```{code-cell} ipython3\nv = nnx.Variable(jnp.array(0))\nv_immut = nnx.vars_as(v, mutable=False)\nassert not v_immut.mutable\n\ntry:\n v_immut[...] += 1 # raises an error\nexcept Exception as e:\n print(f\"{type(e).__name__}: {e}\")\n```\n\n### Ref support\n\n```{code-cell} ipython3\nv = nnx.Variable(jnp.array(0))\nv_ref = nnx.vars_as(v, ref=True)\nassert v_ref.ref\nprint(v_ref)\nprint(v_ref.get_raw_value())\n```\n\n```{code-cell} ipython3\nv_immut = nnx.vars_as(v_ref, mutable=False)\nassert not v_immut.ref\nprint(\"immutable =\", v_immut)\n\nv_ref = nnx.vars_as(v_immut, mutable=True)\nassert v_ref.ref\nprint(\"mutable =\", v_ref)\n```\n\n### Examples\n\n```{code-cell} ipython3\nclass Block(nnx.Module):\n def __init__(self, din, dmid, dout, rngs: nnx.Rngs):\n self.linear = Linear(din, dmid, rngs=rngs)\n self.bn = nnx.BatchNorm(dmid, rngs=rngs)\n self.dropout = nnx.Dropout(0.1, rngs=rngs)\n self.linear_out = Linear(dmid, dout, rngs=rngs)\n\n def __call__(self, x):\n x = nnx.gelu(self.dropout(self.bn(self.linear(x))))\n return self.linear_out(x)\n```\n\n#### Training Loop\n\n```{code-cell} ipython3\n# hijax Variables by default\nmodel = Block(2, 64, 3, rngs=nnx.Rngs(0))\noptimizer = nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param)\n\n@jax.jit\ndef train_step(model, optimizer, x, y):\n graphdef, params, nondiff = nnx.split(model, nnx.Param, ...)\n def loss_fn(params):\n model = nnx.merge(graphdef, params, nondiff)\n return ((model(x) - y) ** 2).mean()\n\n loss, grads = jax.value_and_grad(loss_fn)(nnx.vars_as(params, mutable=False)) # immutable for jax.grad\n optimizer.update(model, grads)\n\n return loss\n\nfor _ in range(3):\n loss = train_step(model, optimizer, x=jnp.ones((10, 2)), y=jnp.ones((10, 3)))\n print(f\"{loss = !s}\")\n```\n\n#### Scan Over Layers\n\n```{code-cell} ipython3\n# TODO: does not work with hijax yet\n# @jax.vmap\n# def create_stack(rngs):\n# return nnx.as_immutable_vars(Block(2, 64, 2, rngs=rngs))\n\n# block_stack = nnx.as_mutable_vars(create_stack(nnx.Rngs(0).fork(split=8)))\n\n# def scan_fn(x, block):\n# x = block(x)\n# return x, None\n\n# x = jax.random.uniform(jax.random.key(0), (3, 2))\n# y, _ = jax.lax.scan(scan_fn, x, block_stack)\n\n# print(\"y = \", y)\n```\n\n### Limitations\n\n+++\n\n#### Mutable Outputs\n\n```{code-cell} ipython3\n@jax.jit\ndef create_model(rngs):\n return Block(2, 64, 3, rngs=rngs)\n\ntry:\n model = create_model(nnx.Rngs(0))\nexcept Exception as e:\n print(f\"Error:\", e)\n```\n\n```{code-cell} ipython3\n@jax.jit\ndef create_model(rngs):\n return nnx.vars_as((Block(2, 64, 3, rngs=rngs)), hijax=False)\n\nmodel = nnx.vars_as(create_model(nnx.Rngs(0)), hijax=True)\n\nprint(\"model.linear =\", model.linear)\n```\n\n#### Reference Sharing (aliasing)\n\n```{code-cell} ipython3\n# NOTE: doesn't currently fail on the jax side\ndef get_error(f, *args):\n try:\n return f(*args)\n except Exception as e:\n return f\"{type(e).__name__}: {e}\"\n\nx = nnx.Variable(jnp.array(0))\n\n@jax.jit\ndef f(a, b):\n ...\n\nprint(get_error(f, x, x))\n```\n\n```{code-cell} ipython3\n# NOTE: doesn't currently fail on the jax side\nclass HasShared(nnx.Pytree):\n def __init__(self):\n self.a = nnx.Variable(jnp.array(0))\n self.b = self.a\n\n@jax.jit\ndef g(has_shared):\n has_shared.a[...] = 5\n\nhas_shared = HasShared()\n\nprint(get_error(g, has_shared))\nprint(has_shared) # updates don't propagate\n```\n\n```{code-cell} ipython3\nprint(\"Duplicates found:\")\nif (all_duplicates := nnx.find_duplicates(has_shared)):\n for duplicates in all_duplicates:\n print(\"-\", duplicates)\n```\n\n```{code-cell} ipython3\n@jax.jit\ndef h(graphdef, state):\n has_shared = nnx.merge(graphdef, state)\n has_shared.a[...] = 5\n\ngraphdef, state = nnx.split(has_shared)\nh(graphdef, state)\nprint(has_shared)\n```\n\n```{code-cell} ipython3\n# clean up for CI tests\n_ = nnx.var_defaults(hijax=current_mode)\n```\n" }, { "path": "docs_nnx/hijax/index.rst", "content": "Hijax (experimental)\n====================\n\nBasic usage\n^^^^^^^^^^^^\n\n.. testsetup::\n\n import jax\n import jax.numpy as jnp\n\n current_mode = nnx.var_defaults().hijax\n\n.. testcode::\n\n from flax import nnx\n import optax\n\n nnx.var_defaults(hijax=True)\n\n class Model(nnx.Module):\n def __init__(self, din, dmid, dout, rngs: nnx.Rngs):\n self.linear = nnx.Linear(din, dmid, rngs=rngs)\n self.bn = nnx.BatchNorm(dmid, rngs=rngs)\n self.dropout = nnx.Dropout(0.2)\n self.linear_out = nnx.Linear(dmid, dout, rngs=rngs)\n\n def __call__(self, x, rngs):\n x = nnx.relu(self.dropout(self.bn(self.linear(x)), rngs=rngs))\n return self.linear_out(x)\n\n model = Model(2, 64, 3, rngs=nnx.Rngs(0)) # eager initialization\n optimizer = nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param)\n\n @jax.jit\n def train_step(model, optimizer, rngs, x, y):\n graphdef, params, nondiff = nnx.split(model, nnx.Param, ...)\n def loss_fn(params):\n model = nnx.merge(graphdef, params, nondiff)\n return ((model(x, rngs) - y) ** 2).mean()\n loss, grads = jax.value_and_grad(loss_fn)(nnx.vars_as(params, mutable=False))\n optimizer.update(model, grads) # in-place updates\n return loss\n\n nnx.var_defaults(hijax=current_mode) # clean up for CI tests\n\n\n----\n\n.. toctree::\n :hidden:\n :maxdepth: 2\n\n hijax\n" }, { "path": "docs_nnx/index.rst", "content": "\nFlax\n====\n.. div:: sd-text-left sd-font-italic\n\n **N**\\ eural **N**\\ etworks for JA\\ **X**\n\n\n----\n\nFlax provides a **flexible end-to-end user experience for researchers and developers who use JAX for neural networks**. Flax enables you to use the full power of `JAX `__.\n\nAt the core of Flax is **NNX - a simplified API that makes it easier to create, inspect,\ndebug, and analyze neural networks in JAX.** Flax NNX has first class support\nfor Python reference semantics, enabling users to express their models using regular\nPython objects. Flax NNX is an evolution of the previous `Flax Linen `__\nAPI, and it took years of experience to bring a simpler and more user-friendly API.\n\n.. note::\n Flax Linen API is not going to be deprecated in the near future as most of Flax users still rely on this API. However, new users are encouraged to use Flax NNX. Check out `Why Flax NNX `_ for a comparison between Flax NNX and Linen, and our reasoning to make the new API.\n\n To move your Flax Linen codebase to Flax NNX, get familiarized with the API in `NNX Basics `_ and then start your move following the `evolution guide `_.\n\nFeatures\n^^^^^^^^^\n\n.. grid::\n\n .. grid-item::\n :columns: 12 12 12 6\n\n .. card:: Pythonic\n :class-card: sd-border-0\n :shadow: none\n :class-title: sd-fs-5\n\n .. div:: sd-font-normal\n\n Flax NNX supports the use of regular Python objects, providing an intuitive\n and predictable development experience.\n\n .. grid-item::\n :columns: 12 12 12 6\n\n .. card:: Simple\n :class-card: sd-border-0\n :shadow: none\n :class-title: sd-fs-5\n\n .. div:: sd-font-normal\n\n Flax NNX relies on Python's object model, which results in simplicity for\n the user and increases development speed.\n\n .. grid-item::\n :columns: 12 12 12 6\n\n .. card:: Expressive\n :class-card: sd-border-0\n :shadow: none\n :class-title: sd-fs-5\n\n .. div:: sd-font-normal\n\n Flax NNX allows fine-grained control of the model's state via\n its `Filter `__\n system.\n\n .. grid-item::\n :columns: 12 12 12 6\n\n .. card:: Familiar\n :class-card: sd-border-0\n :shadow: none\n :class-title: sd-fs-5\n\n .. div:: sd-font-normal\n\n Flax NNX makes it very easy to integrate objects with regular JAX code\n via the `Functional API `__.\n\nBasic usage\n^^^^^^^^^^^^\n\n.. testsetup::\n\n import jax\n import jax.numpy as jnp\n\n.. testcode::\n\n from flax import nnx\n import optax\n\n\n class Model(nnx.Module):\n def __init__(self, din, dmid, dout, rngs: nnx.Rngs):\n self.linear = nnx.Linear(din, dmid, rngs=rngs)\n self.bn = nnx.BatchNorm(dmid, rngs=rngs)\n self.dropout = nnx.Dropout(0.2)\n self.linear_out = nnx.Linear(dmid, dout, rngs=rngs)\n\n def __call__(self, x, rngs):\n x = nnx.relu(self.dropout(self.bn(self.linear(x)), rngs=rngs))\n return self.linear_out(x)\n\n model = Model(2, 64, 3, rngs=nnx.Rngs(0)) # eager initialization\n optimizer = nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param)\n\n @nnx.jit # automatic state propagation\n def train_step(model, optimizer, x, y):\n loss_fn = lambda model: ((model(x) - y) ** 2).mean()\n loss, grads = nnx.value_and_grad(loss_fn)(model)\n optimizer.update(model, grads) # in-place updates\n return loss\n\n\nInstallation\n^^^^^^^^^^^^\n\nInstall via pip:\n\n.. code-block:: bash\n\n pip install flax\n\nOr install the latest version from the repository:\n\n.. code-block:: bash\n\n pip install git+https://github.com/google/flax.git\n\n\n----\n\nLearn more\n^^^^^^^^^^\n\n.. grid::\n\n .. grid-item::\n :columns: 6 6 6 4\n\n .. card:: :material-regular:`rocket_launch;2em` Flax NNX Basics\n :class-card: sd-text-black sd-bg-light\n :link: nnx_basics.html\n\n .. grid-item::\n :columns: 6 6 6 4\n\n .. card:: :material-regular:`library_books;2em` MNIST Tutorial\n :class-card: sd-text-black sd-bg-light\n :link: mnist_tutorial.html\n\n .. grid-item::\n :columns: 6 6 6 4\n\n .. card:: :material-regular:`library_books;2em` Guides\n :class-card: sd-text-black sd-bg-light\n :link: guides/index.html\n\n .. grid-item::\n :columns: 6 6 6 4\n\n .. card:: :material-regular:`transform;2em` Flax Linen to Flax NNX\n :class-card: sd-text-black sd-bg-light\n :link: guides/linen_to_nnx.html\n\n .. grid-item::\n :columns: 6 6 6 4\n\n .. card:: :material-regular:`menu_book;2em` API reference\n :class-card: sd-text-black sd-bg-light\n :link: api_reference/index.html\n\n .. grid-item::\n :columns: 6 6 6 4\n\n .. card:: :material-regular:`import_contacts;2em` Glossary\n :class-card: sd-text-black sd-bg-light\n :link: nnx_glossary.html\n\n\n----\n\n.. toctree::\n :hidden:\n :maxdepth: 3\n\n nnx_basics\n mnist_tutorial\n why\n key_concepts\n guides_basic\n guides_advanced\n Models (Bonsai) \n Post-training (Tunix) \n hijax/index\n migrating/index\n examples/index\n nnx_glossary\n philosophy\n contributing\n api_reference/index\n" }, { "path": "docs_nnx/key_concepts.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"id\": \"929920d0\",\n \"metadata\": {},\n \"source\": [\n \"# JAX/Flax Key Concepts\\n\",\n \"\\n\",\n \"Flax is a **neural network library** built on top of JAX, a language for **accelerated numerical computations**. In effect, Flax is a pretty thin layer, and you likely will use some JAX APIs directly to do anything more than using the built-in Flax modules.\\n\",\n \"\\n\",\n \"This means a **basic understanding on JAX helps you to use Flax well**. You would have better a mental model to understand what's happening underneath and how to debug a confusing error. This doc aims to clarify a few key concepts and help you build that uniquely-JAX mental model as a practical model developer (pun intended).\\n\",\n \"\\n\",\n \"[JAX documentations](https://docs.jax.dev/en/latest/index.html) are great sources to learn more. We recommend all Flax users to at least read the [JAX Key Concepts](https://docs.jax.dev/en/latest/key-concepts.html) doc.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 1,\n \"id\": \"3515d62b\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"import jax\\n\",\n \"import jax.numpy as jnp\\n\",\n \"import flax\\n\",\n \"from flax import nnx\\n\",\n \"from functools import partial\\n\",\n \"\\n\",\n \"# For simulating multi-device environment\\n\",\n \"jax.config.update('jax_num_cpu_devices', 8)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"be2cad4a\",\n \"metadata\": {},\n \"source\": [\n \"## What is JAX?\\n\",\n \"\\n\",\n \"JAX is the lower level library that does **all the large-scale data computations**. It provides the singular data container, aka the `jax.Array`, and all the ways we possibly deal with them:\\n\",\n \"\\n\",\n \"* **Make arithmetic operations upon the arrays**, including: the `jax.numpy` ops, automatic differentiation (`jax.grad`), batching (`jax.vmap`), and more.\\n\",\n \"\\n\",\n \"* **Run computation on accelerators**, including: interface with various accelerator platforms and layouts; allocating buffers for arrays; compile and execute computation programs across accelerators.\\n\",\n \"\\n\",\n \"* **Bundle multiple arrays together** using a simple concept called [pytrees](#pytrees).\\n\",\n \"\\n\",\n \"This implies that any error related with accelerators and numericals are probably a JAX issue, or an issue with Flax built-in layers.\\n\",\n \"\\n\",\n \"It also means you *can* build a neural network model with JAX alone, especially if you are comfortable with functional programming. JAX docsite have some [simple examples](https://docs.jax.dev/en/latest/notebooks/neural_network_with_tfds_data.html). The article [GPT in 60 Lines of NumPy](https://jaykmody.com/blog/gpt-from-scratch/) also shows how to implement all the key elements of a GPT using JAX.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 2,\n \"id\": \"a3769631\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"def jax_linear(x, kernel, bias):\\n\",\n \" return jnp.dot(x, kernel) + bias\\n\",\n \"\\n\",\n \"params = {'kernel': jax.random.normal(jax.random.key(42), (4, 2)), \\n\",\n \" 'bias': jnp.zeros((2,))}\\n\",\n \"x = jax.random.normal(jax.random.key(0), (2, 4))\\n\",\n \"y = jax_linear(x, params['kernel'], params['bias'])\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"ee6f86e7\",\n \"metadata\": {},\n \"source\": [\n \"## What is Flax?\\n\",\n \"\\n\",\n \"Flax is a **neural network toolkit**, offering higher level abstractions that are handy for model developers. Such as:\\n\",\n \"\\n\",\n \"* **Object-oriented `Module` class** to represent layers/models and bookkeep parameters.\\n\",\n \"\\n\",\n \"* **Modeling utilities** like random number handling, model traversal and surgery, optimizers, advanced parameter bookkeeping, sharding annotations, and more.\\n\",\n \"\\n\",\n \"* **Some built-in commonly-used** layers, initializers, and model examples.\\n\",\n \"\\n\",\n \"Take the example below: A Flax layer `Linear`, during initialization, takes one RNG key and automatically initialize all internal parameters as `jax.Array`s. In forward pass, it carries out the exact same computation via JAX APIs.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 3,\n \"id\": \"14caace1\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"\\n\",\n \"\\n\"\n ]\n }\n ],\n \"source\": [\n \"# Eligible parameters were created inside `linear`, using one RNG key 42\\n\",\n \"linear = nnx.Linear(in_features=4, out_features=2, rngs=nnx.Rngs(42))\\n\",\n \"\\n\",\n \"# Flax created a `Param` wrapper over the actual `jax.Array` parameter to track metadata\\n\",\n \"print(type(linear.kernel)) # flax.nnx.Param\\n\",\n \"print(type(linear.kernel.value)) # jax.Array\\n\",\n \"\\n\",\n \"# The computation of the two are the same\\n\",\n \"x = jax.random.normal(jax.random.key(0), (2, 4))\\n\",\n \"flax_y = linear(x)\\n\",\n \"jax_y = jax_linear(x, linear.kernel.value, linear.bias.value)\\n\",\n \"assert jnp.array_equal(flax_y, jax_y)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"09989bf7\",\n \"metadata\": {},\n \"source\": [\n \"## Pytrees\\n\",\n \"\\n\",\n \"Your code likely needs more than one `jax.Array`. A **pytree** is a container structure of multiple pytrees, possibly nested. It is a key and handly concept in the JAX world.\\n\",\n \"\\n\",\n \"Many things are pytrees: Python dicts, lists, tuples, dataclasses, and more. The key is that a pytree can be \\\"flattened\\\" into multiple children, which are either pytrees or individual leaves - a `jax.Array` counts as a leaf. Other metadata of a pytree are stored in the `PyTreeDef` object, allowing \\\"unflattening\\\" to restore the old pytree.\\n\",\n \"\\n\",\n \"Pytree is the primary data holder in JAX. When JAX transforms see a pytree argument, they automatically trace its internal `jax.Array`s when compiling. Therefore, it's crucial to organize your data as pytrees. You can use [`flax.struct.dataclass`](https://flax.readthedocs.io/en/latest/api_reference/flax.struct.html#flax.struct.dataclass) to quickly construct a pytree node dataclass, or register your own classes via JAX API. [JAX pytree documentation](https://docs.jax.dev/en/latest/working-with-pytrees.html) has a thorough overview on pytrees and JAX APIs to manipulate them. \\n\",\n \"\\n\",\n \"In Flax, a `Module` is a pytree, and variables are its flattenable data. This means you can directly run JAX transforms upon a Flax model.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 4,\n \"id\": \"a2059c47\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"linear.bias.value: [0. 0.]\\n\",\n \"linear.kernel.value: [[ 0.04119061 -0.2629074 ]\\n\",\n \" [ 0.6772455 0.2807398 ]\\n\",\n \" [ 0.16276604 0.16813846]\\n\",\n \" [ 0.310975 -0.43336964]]\\n\",\n \"treedef = PyTreeDef(CustomNode(Linear[(('_pytree__state', 'bias', 'kernel'), (('_object__nodes', frozenset({'kernel', '_pytree__state', 'bias'})), ('bias_init', ), ('dot_general', ), ('dtype', None), ('in_features', 4), ('kernel_init', .init at 0x120f45260>), ('out_features', 2), ('param_dtype', ), ('precision', None), ('promote_dtype', ), ('use_bias', True)))], [CustomNode(ObjectState[(False, False)], []), CustomNode(Param[()], [*]), CustomNode(Param[()], [*])]))\\n\"\n ]\n }\n ],\n \"source\": [\n \"# Flatten allows you to see all the content inside a pytree\\n\",\n \"arrays, treedef = jax.tree.flatten_with_path(linear)\\n\",\n \"assert len(arrays) > 1\\n\",\n \"for kp, value in arrays:\\n\",\n \" print(f'linear{jax.tree_util.keystr(kp)}: {value}')\\n\",\n \"print(f'{treedef = }')\\n\",\n \"\\n\",\n \"# Unflatten brings the pytree back intact\\n\",\n \"linear = jax.tree.unflatten(treedef, [value for _, value in arrays])\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 5,\n \"id\": \"4ea2f351\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"y = jax.jit(linear)(x) # JAX transforms works on Flax modules\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"723b3f42\",\n \"metadata\": {},\n \"source\": [\n \"## Traced vs. static data\\n\",\n \"\\n\",\n \"A pytree *contains* JAX arrays, but a pytree is *more than* its JAX arrays. For example, a dictionary keeps information like the key of every array, and it might contain entries that are not JAX arrays. From JAX's standpoint, all data are one of the two types:\\n\",\n \"\\n\",\n \"* **Traced** (\\\"dynamic\\\") data: JAX will trace them during compilation and optimize the operations upon them. If they stay inside a pytree argument, `jax.tree.flatten` must return them as leaves. They must be data values (`jax.Array`, Numpy array, scalar, etc), and implement basic functionalities like `__eq__` and `__hash__`.\\n\",\n \"\\n\",\n \"* **\\\"Static\\\"** data: They stay as simple Python objects that don't get traced by JAX.\\n\",\n \"\\n\",\n \"In practice, you would want to control what data goes into dynamic, and what to static. Dynamic data and their computation will be optimized by JAX, but you cannot base your code control flow upon its values. Non-data values like strings must stay static.\\n\",\n \"\\n\",\n \"Take a Flax model: you would want JAX to only track and optimize its parameters, and the RNG keys. For trivial things like the model hyperparameters (e.g., the param shape, the initializer function), they can stay static to save compilation bandwidth and allow code path customization.\\n\",\n \"\\n\",\n \"Current Flax module automatically classifies this for you. Only the `jax.Array` attributes are treated as dynamic data, unless you explicitly wrap a data value using `nnx.Variable` classes.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 6,\n \"id\": \"7b6cbc3a\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"['rng']['default']['count'].value: 1\\n\",\n \"['rng']['default']['key'].value: Array((), dtype=key) overlaying:\\n\",\n \"[0 0]\\n\",\n \"['traced_dim'].value: 4\\n\",\n \"['w'].value: [[ 1.0040143 -0.9063372 -0.7481722 -1.1713669 ]\\n\",\n \" [-0.8712328 0.5888381 0.72392994 -1.0255982 ]\\n\",\n \" [ 1.661628 -1.8910251 -1.2889339 0.13360691]\\n\",\n \" [-1.1530392 0.23929629 1.7448074 0.5050189 ]]\\n\"\n ]\n }\n ],\n \"source\": [\n \"class Foo(nnx.Module):\\n\",\n \" def __init__(self, dim, rngs):\\n\",\n \" self.w = nnx.Param(jax.random.normal(rngs.param(), (dim, dim)))\\n\",\n \" self.dim = dim\\n\",\n \" self.traced_dim = nnx.Param(dim) # This became traced!\\n\",\n \" self.rng = rngs\\n\",\n \"\\n\",\n \"foo = Foo(4, nnx.Rngs(0))\\n\",\n \"for kp, x in jax.tree.flatten_with_path(nnx.state(foo))[0]:\\n\",\n \" print(f'{jax.tree_util.keystr(kp)}: {x}')\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"0b10383c\",\n \"metadata\": {},\n \"source\": [\n \"When compiling a function using this pytree, you'll notice the difference between traced and static values. You can only use static ones in control flows.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 7,\n \"id\": \"395c9d79\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"model.dim = 4\\n\",\n \"model.traced_dim.value = JitTracer<~int32[]>\\n\",\n \"Code path based on static data value works fine.\\n\",\n \"Code path based on JAX data value throws error: Attempted boolean conversion of traced array with shape bool[].\\n\",\n \"The error occurred while tracing the function jitted at /var/folders/4c/ylxxyg_n67957jf6616c7z5000gbn1/T/ipykernel_69242/584946237.py:1 for jit. This concrete value was not available in Python because it depends on the value of the argument model.traced_dim.value.\\n\",\n \"See https://docs.jax.dev/en/latest/errors.html#jax.errors.TracerBoolConversionError\\n\"\n ]\n }\n ],\n \"source\": [\n \"@jax.jit\\n\",\n \"def jitted(model):\\n\",\n \" print(f'{model.dim = }')\\n\",\n \" print(f'{model.traced_dim.value = }') # This is being traced\\n\",\n \" if model.dim == 4:\\n\",\n \" print('Code path based on static data value works fine.')\\n\",\n \" try:\\n\",\n \" if model.traced_dim.value == 4:\\n\",\n \" print('This will never run :(')\\n\",\n \" except jax.errors.TracerBoolConversionError as e:\\n\",\n \" print(f'Code path based on JAX data value throws error: {e}')\\n\",\n \"\\n\",\n \"jitted(foo)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"202bf52b\",\n \"metadata\": {},\n \"source\": [\n \"## Abstract arrays\\n\",\n \"\\n\",\n \"Abstract array is a JAX class to represent an array not by its value, but simply by its metadata information like shape, dtype and sharding. It is fast and handy because it doesn't allocate any memory for the array data.\\n\",\n \"\\n\",\n \"You can construct an abstract array by calling [`jax.ShapeDtypeStruct`](https://docs.jax.dev/en/latest/_autosummary/jax.ShapeDtypeStruct.html) on your own, or use [`jax.eval_shape`](https://docs.jax.dev/en/latest/_autosummary/jax.eval_shape.html), which takes a function and arguments and returns the abstract version of its output.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 8,\n \"id\": \"21ebeebf\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"[[ 1.0040143 -0.9063372 -0.7481722 -1.1713669 ]\\n\",\n \" [-0.8712328 0.5888381 0.72392994 -1.0255982 ]\\n\",\n \" [ 1.661628 -1.8910251 -1.2889339 0.13360691]\\n\",\n \" [-1.1530392 0.23929629 1.7448074 0.5050189 ]]\\n\",\n \"ShapeDtypeStruct(shape=(4, 4), dtype=float32)\\n\"\n ]\n }\n ],\n \"source\": [\n \"print(x)\\n\",\n \"abs_x = jax.eval_shape(lambda x: x, x)\\n\",\n \"print(abs_x)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"e8345d12\",\n \"metadata\": {},\n \"source\": [\n \"It is a good way to dry-run your code and debug a model without any actual compute and memory cost. For example, you can have an overview of the parameters inside this very large model.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"f9b1b308\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"\\u001b[38;2;79;201;177mLinear\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;105;105;105m # Param: 67,084,290 (268.3 MB)\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254mbias\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;79;201;177mParam\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;105;105;105m # 8,190 (32.8 KB)\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254mvalue\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0mShapeDtypeStruct(shape=(8190,), dtype=float32)\\n\",\n \" \\u001b[38;2;255;213;3m)\\u001b[0m,\\n\",\n \" \\u001b[38;2;156;220;254mkernel\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;79;201;177mParam\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;105;105;105m # 67,076,100 (268.3 MB)\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254mvalue\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0mShapeDtypeStruct(shape=(8190, 8190), dtype=float32)\\n\",\n \" \\u001b[38;2;255;213;3m)\\u001b[0m,\\n\",\n \" \\u001b[38;2;156;220;254mbias_init\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m,\\n\",\n \" \\u001b[38;2;156;220;254mdot_general\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m,\\n\",\n \" \\u001b[38;2;156;220;254mdtype\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;86;156;214mNone\\u001b[0m,\\n\",\n \" \\u001b[38;2;156;220;254min_features\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;182;207;169m8190\\u001b[0m,\\n\",\n \" \\u001b[38;2;156;220;254mkernel_init\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m.init at 0x120f45260>,\\n\",\n \" \\u001b[38;2;156;220;254mout_features\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;182;207;169m8190\\u001b[0m,\\n\",\n \" \\u001b[38;2;156;220;254mparam_dtype\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;79;201;177mfloat32\\u001b[0m,\\n\",\n \" \\u001b[38;2;156;220;254mprecision\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;86;156;214mNone\\u001b[0m,\\n\",\n \" \\u001b[38;2;156;220;254mpromote_dtype\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m,\\n\",\n \" \\u001b[38;2;156;220;254muse_bias\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0m\\u001b[38;2;86;156;214mTrue\\u001b[0m\\n\",\n \"\\u001b[38;2;255;213;3m)\\u001b[0m\\n\"\n ]\n }\n ],\n \"source\": [\n \"class MLP(nnx.Module):\\n\",\n \" def __init__(self, dim, nlayers, rngs):\\n\",\n \" self.blocks = nnx.List([nnx.Linear(dim, dim, rngs=rngs) for _ in range(nlayers)])\\n\",\n \" self.activation = jax.nn.relu\\n\",\n \" self.nlayers = nlayers\\n\",\n \" def __call__(self, x):\\n\",\n \" for block in self.blocks:\\n\",\n \" x = self.activation(block(x))\\n\",\n \" return x\\n\",\n \"\\n\",\n \"dim, nlayers = 8190, 64 # Some very big numbers\\n\",\n \"@partial(jax.jit, static_argnums=(0, 1))\\n\",\n \"def init_state(dim, nlayers):\\n\",\n \" return MLP(dim, nlayers, nnx.Rngs(0))\\n\",\n \"abstract_model = jax.eval_shape(partial(init_state, dim, nlayers))\\n\",\n \"print(abstract_model.blocks[0])\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"8894cbc6\",\n \"metadata\": {},\n \"source\": [\n \"Once you have an abstract pytree for your function input or output, it's easier to describe how you want your data to be sharded. You should use such a pytree with sharding information to instruct your checkpoint loading library to load your arrays distributedly. Our checkpointing guide contains [an example of how to do this](https://flax.readthedocs.io/en/latest/guides/flax_gspmd.html#load-a-sharded-model-from-a-checkpoint).\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"b98b5184\",\n \"metadata\": {},\n \"source\": [\n \"## Distributed computing\\n\",\n \"\\n\",\n \"Another big use case for abstract pytrees is to tell JAX machinery how you want each array to be sharded during any point of your computation.\\n\",\n \"\\n\",\n \"Remember what we mentioned earlier? JAX handles the actual computation and data allocation on accelerators. This means you **must** use some `jax.jit`-compiled function to run any distributed computation task.\\n\",\n \"\\n\",\n \"There are a few ways to tell `jax.jit` of your model sharding. The simplest way is to call `jax.lax.with_sharding_constraint` to constraint the so-to-be model with your predetermined model sharding.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 10,\n \"id\": \"9b289c02\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"\\u001b[38;2;79;201;177mParam\\u001b[0m\\u001b[38;2;255;213;3m(\\u001b[0m\\u001b[38;2;105;105;105m\\u001b[0m\\n\",\n \" \\u001b[38;2;156;220;254mvalue\\u001b[0m\\u001b[38;2;212;212;212m=\\u001b[0mNamedSharding(mesh=Mesh('model': 8, axis_types=(Auto,)), spec=PartitionSpec(None, 'model'), memory_kind=unpinned_host)\\n\",\n \"\\u001b[38;2;255;213;3m)\\u001b[0m\\n\"\n ]\n },\n {\n \"data\": {\n \"text/html\": [\n \"
                                                                        \\n\",\n       \"                                                                        \\n\",\n       \"                                                                        \\n\",\n       \"                                                                        \\n\",\n       \"                                                                        \\n\",\n       \"  CPU 0    CPU 1    CPU 2    CPU 3    CPU 4    CPU 5    CPU 6    CPU 7  \\n\",\n       \"                                                                        \\n\",\n       \"                                                                        \\n\",\n       \"                                                                        \\n\",\n       \"                                                                        \\n\",\n       \"                                                                        \\n\",\n       \"
\\n\"\n ],\n \"text/plain\": [\n \"\\u001b[38;2;255;255;255;48;2;57;59;121m \\u001b[0m\\u001b[38;2;255;255;255;48;2;214;97;107m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;162;82m \\u001b[0m\\u001b[38;2;255;255;255;48;2;222;158;214m \\u001b[0m\\u001b[38;2;0;0;0;48;2;231;203;148m \\u001b[0m\\u001b[38;2;255;255;255;48;2;107;110;207m \\u001b[0m\\u001b[38;2;255;255;255;48;2;165;81;148m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;109;49m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;57;59;121m \\u001b[0m\\u001b[38;2;255;255;255;48;2;214;97;107m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;162;82m \\u001b[0m\\u001b[38;2;255;255;255;48;2;222;158;214m \\u001b[0m\\u001b[38;2;0;0;0;48;2;231;203;148m \\u001b[0m\\u001b[38;2;255;255;255;48;2;107;110;207m \\u001b[0m\\u001b[38;2;255;255;255;48;2;165;81;148m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;109;49m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;57;59;121m \\u001b[0m\\u001b[38;2;255;255;255;48;2;214;97;107m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;162;82m \\u001b[0m\\u001b[38;2;255;255;255;48;2;222;158;214m \\u001b[0m\\u001b[38;2;0;0;0;48;2;231;203;148m \\u001b[0m\\u001b[38;2;255;255;255;48;2;107;110;207m \\u001b[0m\\u001b[38;2;255;255;255;48;2;165;81;148m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;109;49m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;57;59;121m \\u001b[0m\\u001b[38;2;255;255;255;48;2;214;97;107m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;162;82m \\u001b[0m\\u001b[38;2;255;255;255;48;2;222;158;214m \\u001b[0m\\u001b[38;2;0;0;0;48;2;231;203;148m \\u001b[0m\\u001b[38;2;255;255;255;48;2;107;110;207m \\u001b[0m\\u001b[38;2;255;255;255;48;2;165;81;148m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;109;49m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;57;59;121m \\u001b[0m\\u001b[38;2;255;255;255;48;2;214;97;107m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;162;82m \\u001b[0m\\u001b[38;2;255;255;255;48;2;222;158;214m \\u001b[0m\\u001b[38;2;0;0;0;48;2;231;203;148m \\u001b[0m\\u001b[38;2;255;255;255;48;2;107;110;207m \\u001b[0m\\u001b[38;2;255;255;255;48;2;165;81;148m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;109;49m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;57;59;121m \\u001b[0m\\u001b[38;2;255;255;255;48;2;57;59;121mCPU 0\\u001b[0m\\u001b[38;2;255;255;255;48;2;57;59;121m \\u001b[0m\\u001b[38;2;255;255;255;48;2;214;97;107m \\u001b[0m\\u001b[38;2;255;255;255;48;2;214;97;107mCPU 1\\u001b[0m\\u001b[38;2;255;255;255;48;2;214;97;107m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;162;82m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;162;82mCPU 2\\u001b[0m\\u001b[38;2;255;255;255;48;2;140;162;82m \\u001b[0m\\u001b[38;2;255;255;255;48;2;222;158;214m \\u001b[0m\\u001b[38;2;255;255;255;48;2;222;158;214mCPU 3\\u001b[0m\\u001b[38;2;255;255;255;48;2;222;158;214m \\u001b[0m\\u001b[38;2;0;0;0;48;2;231;203;148m \\u001b[0m\\u001b[38;2;0;0;0;48;2;231;203;148mCPU 4\\u001b[0m\\u001b[38;2;0;0;0;48;2;231;203;148m \\u001b[0m\\u001b[38;2;255;255;255;48;2;107;110;207m \\u001b[0m\\u001b[38;2;255;255;255;48;2;107;110;207mCPU 5\\u001b[0m\\u001b[38;2;255;255;255;48;2;107;110;207m \\u001b[0m\\u001b[38;2;255;255;255;48;2;165;81;148m \\u001b[0m\\u001b[38;2;255;255;255;48;2;165;81;148mCPU 6\\u001b[0m\\u001b[38;2;255;255;255;48;2;165;81;148m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;109;49m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;109;49mCPU 7\\u001b[0m\\u001b[38;2;255;255;255;48;2;140;109;49m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;57;59;121m \\u001b[0m\\u001b[38;2;255;255;255;48;2;214;97;107m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;162;82m \\u001b[0m\\u001b[38;2;255;255;255;48;2;222;158;214m \\u001b[0m\\u001b[38;2;0;0;0;48;2;231;203;148m \\u001b[0m\\u001b[38;2;255;255;255;48;2;107;110;207m \\u001b[0m\\u001b[38;2;255;255;255;48;2;165;81;148m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;109;49m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;57;59;121m \\u001b[0m\\u001b[38;2;255;255;255;48;2;214;97;107m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;162;82m \\u001b[0m\\u001b[38;2;255;255;255;48;2;222;158;214m \\u001b[0m\\u001b[38;2;0;0;0;48;2;231;203;148m \\u001b[0m\\u001b[38;2;255;255;255;48;2;107;110;207m \\u001b[0m\\u001b[38;2;255;255;255;48;2;165;81;148m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;109;49m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;57;59;121m \\u001b[0m\\u001b[38;2;255;255;255;48;2;214;97;107m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;162;82m \\u001b[0m\\u001b[38;2;255;255;255;48;2;222;158;214m \\u001b[0m\\u001b[38;2;0;0;0;48;2;231;203;148m \\u001b[0m\\u001b[38;2;255;255;255;48;2;107;110;207m \\u001b[0m\\u001b[38;2;255;255;255;48;2;165;81;148m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;109;49m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;57;59;121m \\u001b[0m\\u001b[38;2;255;255;255;48;2;214;97;107m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;162;82m \\u001b[0m\\u001b[38;2;255;255;255;48;2;222;158;214m \\u001b[0m\\u001b[38;2;0;0;0;48;2;231;203;148m \\u001b[0m\\u001b[38;2;255;255;255;48;2;107;110;207m \\u001b[0m\\u001b[38;2;255;255;255;48;2;165;81;148m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;109;49m \\u001b[0m\\n\",\n \"\\u001b[38;2;255;255;255;48;2;57;59;121m \\u001b[0m\\u001b[38;2;255;255;255;48;2;214;97;107m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;162;82m \\u001b[0m\\u001b[38;2;255;255;255;48;2;222;158;214m \\u001b[0m\\u001b[38;2;0;0;0;48;2;231;203;148m \\u001b[0m\\u001b[38;2;255;255;255;48;2;107;110;207m \\u001b[0m\\u001b[38;2;255;255;255;48;2;165;81;148m \\u001b[0m\\u001b[38;2;255;255;255;48;2;140;109;49m \\u001b[0m\\n\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"# Some smaller numbers so that we actually can run it\\n\",\n \"dim, nlayers = 1024, 2\\n\",\n \"abstract_model = jax.eval_shape(partial(init_state, dim, nlayers))\\n\",\n \"mesh = jax.make_mesh((jax.device_count(), ), 'model')\\n\",\n \"\\n\",\n \"# Generate sharding for each of your params manually, sharded along the last axis.\\n\",\n \"def make_sharding(abs_x):\\n\",\n \" if len(abs_x.shape) > 1:\\n\",\n \" pspec = jax.sharding.PartitionSpec(None, 'model') # kernel\\n\",\n \" else:\\n\",\n \" pspec = jax.sharding.PartitionSpec('model',) # bias\\n\",\n \" return jax.sharding.NamedSharding(mesh, pspec)\\n\",\n \"model_shardings = jax.tree.map(make_sharding, abstract_model)\\n\",\n \"print(model_shardings.blocks[0].kernel)\\n\",\n \"\\n\",\n \"@partial(jax.jit, static_argnums=(0, 1))\\n\",\n \"def sharded_init(dim, nlayers):\\n\",\n \" model = MLP(dim, nlayers, nnx.Rngs(0))\\n\",\n \" return jax.lax.with_sharding_constraint(model, model_shardings)\\n\",\n \"model = sharded_init(dim, nlayers)\\n\",\n \"jax.debug.visualize_array_sharding(model.blocks[0].kernel.value)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"3b7ac7ea\",\n \"metadata\": {},\n \"source\": [\n \"The example below are just to showcase how to do sharding in pure JAX API. Flax offers a small API to annotate the sharding when you define a parameter, so that you don't have to write an arbitrary `make_sharding()` function at top level. Check out our [GSPMD guide](https://flax.readthedocs.io/en/latest/guides/flax_gspmd.html) to learn more.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"7dbf421a\",\n \"metadata\": {},\n \"source\": [\n \"## Transformations\\n\",\n \"\\n\",\n \"For Flax transforms and their relation with JAX transforms, refer to [Flax Transforms](https://flax.readthedocs.io/en/latest/guides/transforms.html) doc. This should be a rarer use case now that Flax NNX modules are JAX pytrees.\"\n ]\n }\n ],\n \"metadata\": {\n \"language_info\": {\n \"codemirror_mode\": {\n \"name\": \"ipython\",\n \"version\": 3\n },\n \"file_extension\": \".py\",\n \"mimetype\": \"text/x-python\",\n \"name\": \"python\",\n \"nbconvert_exporter\": \"python\",\n \"pygments_lexer\": \"ipython3\",\n \"version\": \"3.12.7\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 5\n}\n" }, { "path": "docs_nnx/key_concepts.md", "content": "# JAX/Flax Key Concepts\n\nFlax is a **neural network library** built on top of JAX, a language for **accelerated numerical computations**. In effect, Flax is a pretty thin layer, and you likely will use some JAX APIs directly to do anything more than using the built-in Flax modules.\n\nThis means a **basic understanding on JAX helps you to use Flax well**. You would have better a mental model to understand what's happening underneath and how to debug a confusing error. This doc aims to clarify a few key concepts and help you build that uniquely-JAX mental model as a practical model developer (pun intended).\n\n[JAX documentations](https://docs.jax.dev/en/latest/index.html) are great sources to learn more. We recommend all Flax users to at least read the [JAX Key Concepts](https://docs.jax.dev/en/latest/key-concepts.html) doc.\n\n\n```python\nimport jax\nimport jax.numpy as jnp\nimport flax\nfrom flax import nnx\nfrom functools import partial\n\n# For simulating multi-device environment\njax.config.update('jax_num_cpu_devices', 8)\n```\n\n## What is JAX?\n\nJAX is the lower level library that does **all the large-scale data computations**. It provides the singular data container, aka the `jax.Array`, and all the ways we possibly deal with them:\n\n* **Make arithmetic operations upon the arrays**, including: the `jax.numpy` ops, automatic differentiation (`jax.grad`), batching (`jax.vmap`), and more.\n\n* **Run computation on accelerators**, including: interface with various accelerator platforms and layouts; allocating buffers for arrays; compile and execute computation programs across accelerators.\n\n* **Bundle multiple arrays together** using a simple concept called [pytrees](#pytrees).\n\nThis implies that any error related with accelerators and numericals are probably a JAX issue, or an issue with Flax built-in layers.\n\nIt also means you *can* build a neural network model with JAX alone, especially if you are comfortable with functional programming. JAX docsite have some [simple examples](https://docs.jax.dev/en/latest/notebooks/neural_network_with_tfds_data.html). The article [GPT in 60 Lines of NumPy](https://jaykmody.com/blog/gpt-from-scratch/) also shows how to implement all the key elements of a GPT using JAX.\n\n\n```python\ndef jax_linear(x, kernel, bias):\n return jnp.dot(x, kernel) + bias\n\nparams = {'kernel': jax.random.normal(jax.random.key(42), (4, 2)), \n 'bias': jnp.zeros((2,))}\nx = jax.random.normal(jax.random.key(0), (2, 4))\ny = jax_linear(x, params['kernel'], params['bias'])\n```\n\n## What is Flax?\n\nFlax is a **neural network toolkit**, offering higher level abstractions that are handy for model developers. Such as:\n\n* **Object-oriented `Module` class** to represent layers/models and bookkeep parameters.\n\n* **Modeling utilities** like random number handling, model traversal and surgery, optimizers, advanced parameter bookkeeping, sharding annotations, and more.\n\n* **Some built-in commonly-used** layers, initializers, and model examples.\n\nTake the example below: A Flax layer `Linear`, during initialization, takes one RNG key and automatically initialize all internal parameters as `jax.Array`s. In forward pass, it carries out the exact same computation via JAX APIs.\n\n\n```python\n# Eligible parameters were created inside `linear`, using one RNG key 42\nlinear = nnx.Linear(in_features=4, out_features=2, rngs=nnx.Rngs(42))\n\n# Flax created a `Param` wrapper over the actual `jax.Array` parameter to track metadata\nprint(type(linear.kernel)) # flax.nnx.Param\nprint(type(linear.kernel.value)) # jax.Array\n\n# The computation of the two are the same\nx = jax.random.normal(jax.random.key(0), (2, 4))\nflax_y = linear(x)\njax_y = jax_linear(x, linear.kernel.value, linear.bias.value)\nassert jnp.array_equal(flax_y, jax_y)\n```\n\n \n \n\n\n## Pytrees\n\nYour code likely needs more than one `jax.Array`. A **pytree** is a container structure of multiple pytrees, possibly nested. It is a key and handly concept in the JAX world.\n\nMany things are pytrees: Python dicts, lists, tuples, dataclasses, and more. The key is that a pytree can be \"flattened\" into multiple children, which are either pytrees or individual leaves - a `jax.Array` counts as a leaf. Other metadata of a pytree are stored in the `PyTreeDef` object, allowing \"unflattening\" to restore the old pytree.\n\nPytree is the primary data holder in JAX. When JAX transforms see a pytree argument, they automatically trace its internal `jax.Array`s when compiling. Therefore, it's crucial to organize your data as pytrees. You can use [`flax.struct.dataclass`](https://flax.readthedocs.io/en/latest/api_reference/flax.struct.html#flax.struct.dataclass) to quickly construct a pytree node dataclass, or register your own classes via JAX API. [JAX pytree documentation](https://docs.jax.dev/en/latest/working-with-pytrees.html) has a thorough overview on pytrees and JAX APIs to manipulate them. \n\nIn Flax, a `Module` is a pytree, and variables are its flattenable data. This means you can directly run JAX transforms upon a Flax model.\n\n\n```python\n# Flatten allows you to see all the content inside a pytree\narrays, treedef = jax.tree.flatten_with_path(linear)\nassert len(arrays) > 1\nfor kp, value in arrays:\n print(f'linear{jax.tree_util.keystr(kp)}: {value}')\nprint(f'{treedef = }')\n\n# Unflatten brings the pytree back intact\nlinear = jax.tree.unflatten(treedef, [value for _, value in arrays])\n```\n\n linear.bias.value: [0. 0.]\n linear.kernel.value: [[ 0.04119061 -0.2629074 ]\n [ 0.6772455 0.2807398 ]\n [ 0.16276604 0.16813846]\n [ 0.310975 -0.43336964]]\n treedef = PyTreeDef(CustomNode(Linear[(('_pytree__state', 'bias', 'kernel'), (('_object__nodes', frozenset({'kernel', '_pytree__state', 'bias'})), ('bias_init', ), ('dot_general', ), ('dtype', None), ('in_features', 4), ('kernel_init', .init at 0x120f45260>), ('out_features', 2), ('param_dtype', ), ('precision', None), ('promote_dtype', ), ('use_bias', True)))], [CustomNode(ObjectState[(False, False)], []), CustomNode(Param[()], [*]), CustomNode(Param[()], [*])]))\n\n\n\n```python\ny = jax.jit(linear)(x) # JAX transforms works on Flax modules\n```\n\n## Traced vs. static data\n\nA pytree *contains* JAX arrays, but a pytree is *more than* its JAX arrays. For example, a dictionary keeps information like the key of every array, and it might contain entries that are not JAX arrays. From JAX's standpoint, all data are one of the two types:\n\n* **Traced** (\"dynamic\") data: JAX will trace them during compilation and optimize the operations upon them. If they stay inside a pytree argument, `jax.tree.flatten` must return them as leaves. They must be data values (`jax.Array`, Numpy array, scalar, etc), and implement basic functionalities like `__eq__` and `__hash__`.\n\n* **\"Static\"** data: They stay as simple Python objects that don't get traced by JAX.\n\nIn practice, you would want to control what data goes into dynamic, and what to static. Dynamic data and their computation will be optimized by JAX, but you cannot base your code control flow upon its values. Non-data values like strings must stay static.\n\nTake a Flax model: you would want JAX to only track and optimize its parameters, and the RNG keys. For trivial things like the model hyperparameters (e.g., the param shape, the initializer function), they can stay static to save compilation bandwidth and allow code path customization.\n\nCurrent Flax module automatically classifies this for you. Only the `jax.Array` attributes are treated as dynamic data, unless you explicitly wrap a data value using `nnx.Variable` classes.\n\n\n```python\nclass Foo(nnx.Module):\n def __init__(self, dim, rngs):\n self.w = nnx.Param(jax.random.normal(rngs.param(), (dim, dim)))\n self.dim = dim\n self.traced_dim = nnx.Param(dim) # This became traced!\n self.rng = rngs\n\nfoo = Foo(4, nnx.Rngs(0))\nfor kp, x in jax.tree.flatten_with_path(nnx.state(foo))[0]:\n print(f'{jax.tree_util.keystr(kp)}: {x}')\n```\n\n ['rng']['default']['count'].value: 1\n ['rng']['default']['key'].value: Array((), dtype=key) overlaying:\n [0 0]\n ['traced_dim'].value: 4\n ['w'].value: [[ 1.0040143 -0.9063372 -0.7481722 -1.1713669 ]\n [-0.8712328 0.5888381 0.72392994 -1.0255982 ]\n [ 1.661628 -1.8910251 -1.2889339 0.13360691]\n [-1.1530392 0.23929629 1.7448074 0.5050189 ]]\n\n\nWhen compiling a function using this pytree, you'll notice the difference between traced and static values. You can only use static ones in control flows.\n\n\n```python\n@jax.jit\ndef jitted(model):\n print(f'{model.dim = }')\n print(f'{model.traced_dim.value = }') # This is being traced\n if model.dim == 4:\n print('Code path based on static data value works fine.')\n try:\n if model.traced_dim.value == 4:\n print('This will never run :(')\n except jax.errors.TracerBoolConversionError as e:\n print(f'Code path based on JAX data value throws error: {e}')\n\njitted(foo)\n```\n\n model.dim = 4\n model.traced_dim.value = JitTracer<~int32[]>\n Code path based on static data value works fine.\n Code path based on JAX data value throws error: Attempted boolean conversion of traced array with shape bool[].\n The error occurred while tracing the function jitted at /var/folders/4c/ylxxyg_n67957jf6616c7z5000gbn1/T/ipykernel_69242/584946237.py:1 for jit. This concrete value was not available in Python because it depends on the value of the argument model.traced_dim.value.\n See https://docs.jax.dev/en/latest/errors.html#jax.errors.TracerBoolConversionError\n\n\n## Abstract arrays\n\nAbstract array is a JAX class to represent an array not by its value, but simply by its metadata information like shape, dtype and sharding. It is fast and handy because it doesn't allocate any memory for the array data.\n\nYou can construct an abstract array by calling [`jax.ShapeDtypeStruct`](https://docs.jax.dev/en/latest/_autosummary/jax.ShapeDtypeStruct.html) on your own, or use [`jax.eval_shape`](https://docs.jax.dev/en/latest/_autosummary/jax.eval_shape.html), which takes a function and arguments and returns the abstract version of its output.\n\n\n```python\nprint(x)\nabs_x = jax.eval_shape(lambda x: x, x)\nprint(abs_x)\n```\n\n [[ 1.0040143 -0.9063372 -0.7481722 -1.1713669 ]\n [-0.8712328 0.5888381 0.72392994 -1.0255982 ]\n [ 1.661628 -1.8910251 -1.2889339 0.13360691]\n [-1.1530392 0.23929629 1.7448074 0.5050189 ]]\n ShapeDtypeStruct(shape=(4, 4), dtype=float32)\n\n\nIt is a good way to dry-run your code and debug a model without any actual compute and memory cost. For example, you can have an overview of the parameters inside this very large model.\n\n\n```python\nclass MLP(nnx.Module):\n def __init__(self, dim, nlayers, rngs):\n self.blocks = [nnx.Linear(dim, dim, rngs=rngs) for _ in range(nlayers)]\n self.activation = jax.nn.relu\n self.nlayers = nlayers\n def __call__(self, x):\n for block in self.blocks:\n x = self.activation(block(x))\n return x\n\ndim, nlayers = 8190, 64 # Some very big numbers\n@partial(jax.jit, static_argnums=(0, 1))\ndef init_state(dim, nlayers):\n return MLP(dim, nlayers, nnx.Rngs(0))\nabstract_model = jax.eval_shape(partial(init_state, dim, nlayers))\nprint(abstract_model.blocks[0])\n```\n\n \u001b[38;2;79;201;177mLinear\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;105;105;105m # Param: 67,084,290 (268.3 MB)\u001b[0m\n \u001b[38;2;156;220;254mbias\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;79;201;177mParam\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;105;105;105m # 8,190 (32.8 KB)\u001b[0m\n \u001b[38;2;156;220;254mvalue\u001b[0m\u001b[38;2;212;212;212m=\u001b[0mShapeDtypeStruct(shape=(8190,), dtype=float32)\n \u001b[38;2;255;213;3m)\u001b[0m,\n \u001b[38;2;156;220;254mkernel\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;79;201;177mParam\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;105;105;105m # 67,076,100 (268.3 MB)\u001b[0m\n \u001b[38;2;156;220;254mvalue\u001b[0m\u001b[38;2;212;212;212m=\u001b[0mShapeDtypeStruct(shape=(8190, 8190), dtype=float32)\n \u001b[38;2;255;213;3m)\u001b[0m,\n \u001b[38;2;156;220;254mbias_init\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m,\n \u001b[38;2;156;220;254mdot_general\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m,\n \u001b[38;2;156;220;254mdtype\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;86;156;214mNone\u001b[0m,\n \u001b[38;2;156;220;254min_features\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;182;207;169m8190\u001b[0m,\n \u001b[38;2;156;220;254mkernel_init\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m.init at 0x120f45260>,\n \u001b[38;2;156;220;254mout_features\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;182;207;169m8190\u001b[0m,\n \u001b[38;2;156;220;254mparam_dtype\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;79;201;177mfloat32\u001b[0m,\n \u001b[38;2;156;220;254mprecision\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;86;156;214mNone\u001b[0m,\n \u001b[38;2;156;220;254mpromote_dtype\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m,\n \u001b[38;2;156;220;254muse_bias\u001b[0m\u001b[38;2;212;212;212m=\u001b[0m\u001b[38;2;86;156;214mTrue\u001b[0m\n \u001b[38;2;255;213;3m)\u001b[0m\n\n\nOnce you have an abstract pytree for your function input or output, it's easier to describe how you want your data to be sharded. You should use such a pytree with sharding information to instruct your checkpoint loading library to load your arrays distributedly. Our checkpointing guide contains [an example of how to do this](https://flax.readthedocs.io/en/latest/guides/flax_gspmd.html#load-a-sharded-model-from-a-checkpoint).\n\n## Distributed computing\n\nAnother big use case for abstract pytrees is to tell JAX machinery how you want each array to be sharded during any point of your computation.\n\nRemember what we mentioned earlier? JAX handles the actual computation and data allocation on accelerators. This means you **must** use some `jax.jit`-compiled function to run any distributed computation task.\n\nThere are a few ways to tell `jax.jit` of your model sharding. The simplest way is to call `jax.lax.with_sharding_constraint` to constraint the so-to-be model with your predetermined model sharding.\n\n\n```python\n# Some smaller numbers so that we actually can run it\ndim, nlayers = 1024, 2\nabstract_model = jax.eval_shape(partial(init_state, dim, nlayers))\nmesh = jax.make_mesh((jax.device_count(), ), 'model')\n\n# Generate sharding for each of your params manually, sharded along the last axis.\ndef make_sharding(abs_x):\n if len(abs_x.shape) > 1:\n pspec = jax.sharding.PartitionSpec(None, 'model') # kernel\n else:\n pspec = jax.sharding.PartitionSpec('model',) # bias\n return jax.sharding.NamedSharding(mesh, pspec)\nmodel_shardings = jax.tree.map(make_sharding, abstract_model)\nprint(model_shardings.blocks[0].kernel)\n\n@partial(jax.jit, static_argnums=(0, 1))\ndef sharded_init(dim, nlayers):\n model = MLP(dim, nlayers, nnx.Rngs(0))\n return jax.lax.with_sharding_constraint(model, model_shardings)\nmodel = sharded_init(dim, nlayers)\njax.debug.visualize_array_sharding(model.blocks[0].kernel.value)\n```\n\n \u001b[38;2;79;201;177mParam\u001b[0m\u001b[38;2;255;213;3m(\u001b[0m\u001b[38;2;105;105;105m\u001b[0m\n \u001b[38;2;156;220;254mvalue\u001b[0m\u001b[38;2;212;212;212m=\u001b[0mNamedSharding(mesh=Mesh('model': 8, axis_types=(Auto,)), spec=PartitionSpec(None, 'model'), memory_kind=unpinned_host)\n \u001b[38;2;255;213;3m)\u001b[0m\n\n\n\n
                                                                        \n                                                                        \n                                                                        \n                                                                        \n                                                                        \n  CPU 0    CPU 1    CPU 2    CPU 3    CPU 4    CPU 5    CPU 6    CPU 7  \n                                                                        \n                                                                        \n                                                                        \n                                                                        \n                                                                        \n
\n\n\n\nThe example below are just to showcase how to do sharding in pure JAX API. Flax offers a small API to annotate the sharding when you define a parameter, so that you don't have to write an arbitrary `make_sharding()` function at top level. Check out our [GSPMD guide](https://flax.readthedocs.io/en/latest/guides/flax_gspmd.html) to learn more.\n\n## Transformations\n\nFor Flax transforms and their relation with JAX transforms, refer to [Flax Transforms](https://flax.readthedocs.io/en/latest/guides/transforms.html) doc. This should be a rarer use case now that Flax NNX modules are JAX pytrees.\n" }, { "path": "docs_nnx/migrating/convert_pytorch_to_flax.rst", "content": "Convert PyTorch models to Flax\n==============================\n\n.. testsetup::\n\n import numpy as np\n import jax\n from jax import random, numpy as jnp\n from flax import nnx\n\n import torch\n\nWe will show how to convert PyTorch models to Flax. We will cover convolutions, fc layers, batch norm, and average pooling.\n\n\nFC Layers\n--------------------------------\n\nLet's start with fc layers. The only thing to be aware of here is that the PyTorch kernel has shape [outC, inC]\nand the Flax kernel has shape [inC, outC]. Transposing the kernel will do the trick.\n\n.. testcode::\n\n t_fc = torch.nn.Linear(in_features=3, out_features=4)\n\n kernel = t_fc.weight.detach().cpu().numpy()\n bias = t_fc.bias.detach().cpu().numpy()\n\n # [outC, inC] -> [inC, outC]\n kernel = jnp.transpose(kernel, (1, 0))\n\n key = random.key(0)\n x = random.normal(key, (1, 3))\n\n j_fc = nnx.Linear(in_features=3, out_features=4, rngs=nnx.Rngs(0))\n j_fc.kernel.value = kernel\n j_fc.bias.value = jnp.array(bias)\n j_out = j_fc(x)\n\n t_x = torch.from_numpy(np.array(x))\n t_out = t_fc(t_x)\n t_out = t_out.detach().cpu().numpy()\n\n np.testing.assert_almost_equal(j_out, t_out, decimal=6)\n\n\nConvolutions\n--------------------------------\n\nLet's now look at 2D convolutions. PyTorch uses the NCHW format and Flax uses NHWC.\nConsequently, the kernels will have different shapes. The kernel in PyTorch has shape [outC, inC, kH, kW]\nand the Flax kernel has shape [kH, kW, inC, outC]. Transposing the kernel will do the trick.\n\n.. testcode::\n\n t_conv = torch.nn.Conv2d(in_channels=3, out_channels=4, kernel_size=2, padding='valid')\n\n kernel = t_conv.weight.detach().cpu().numpy()\n bias = t_conv.bias.detach().cpu().numpy()\n\n # [outC, inC, kH, kW] -> [kH, kW, inC, outC]\n kernel = jnp.transpose(kernel, (2, 3, 1, 0))\n\n key = random.key(0)\n x = random.normal(key, (1, 6, 6, 3))\n\n j_conv = nnx.Conv(3, 4, kernel_size=(2, 2), padding='valid', rngs=nnx.Rngs(0))\n j_conv.kernel.value = kernel\n j_conv.bias.value = bias\n j_out = j_conv(x)\n\n # [N, H, W, C] -> [N, C, H, W]\n t_x = torch.from_numpy(np.transpose(np.array(x), (0, 3, 1, 2)))\n t_out = t_conv(t_x)\n # [N, C, H, W] -> [N, H, W, C]\n t_out = np.transpose(t_out.detach().cpu().numpy(), (0, 2, 3, 1))\n\n np.testing.assert_almost_equal(j_out, t_out, decimal=6)\n\n\n\nConvolutions and FC Layers\n--------------------------------\n\nWe have to be careful, when we have a model that uses convolutions followed by fc layers (ResNet, VGG, etc).\nIn PyTorch, the activations will have shape [N, C, H, W] after the convolutions and are then\nreshaped to [N, C * H * W] before being fed to the fc layers.\nWhen we port our weights from PyTorch to Flax, the activations after the convolutions will be of shape [N, H, W, C] in Flax.\nBefore we reshape the activations for the fc layers, we have to transpose them to [N, C, H, W].\n\nConsider this PyTorch model:\n\n.. testcode::\n\n class TModel(torch.nn.Module):\n\n def __init__(self):\n super(TModel, self).__init__()\n self.conv = torch.nn.Conv2d(in_channels=3, out_channels=4, kernel_size=2, padding='valid')\n self.fc = torch.nn.Linear(in_features=100, out_features=2)\n\n def forward(self, x):\n x = self.conv(x)\n x = x.reshape(x.shape[0], -1)\n x = self.fc(x)\n return x\n\n\n t_model = TModel()\n\n\n\nNow, if you want to use the weights from this model in Flax, the corresponding Flax model has to look like this:\n\n\n.. testcode::\n\n class JModel(nnx.Module):\n def __init__(self, rngs):\n self.conv = nnx.Conv(3, 4, kernel_size=(2, 2), padding='valid', rngs=rngs)\n self.linear = nnx.Linear(100, 2, rngs=rngs)\n\n def __call__(self, x):\n x = self.conv(x)\n # [N, H, W, C] -> [N, C, H, W]\n x = jnp.transpose(x, (0, 3, 1, 2))\n x = jnp.reshape(x, (x.shape[0], -1))\n x = self.linear(x)\n return x\n\n j_model = JModel(nnx.Rngs(0))\n\n\n\nThe model looks very similar to the PyTorch model, except that we included a transpose operation before\nreshaping our activations for the fc layer.\nWe can omit the transpose operation if we apply pooling before reshaping such that the spatial dimensions are 1x1.\n\nOther than the transpose operation before reshaping, we can convert the weights the same way as we did before:\n\n\n.. testcode::\n\n conv_kernel = t_model.state_dict()['conv.weight'].detach().cpu().numpy()\n conv_bias = t_model.state_dict()['conv.bias'].detach().cpu().numpy()\n fc_kernel = t_model.state_dict()['fc.weight'].detach().cpu().numpy()\n fc_bias = t_model.state_dict()['fc.bias'].detach().cpu().numpy()\n\n # [outC, inC, kH, kW] -> [kH, kW, inC, outC]\n conv_kernel = jnp.transpose(conv_kernel, (2, 3, 1, 0))\n\n # [outC, inC] -> [inC, outC]\n fc_kernel = jnp.transpose(fc_kernel, (1, 0))\n\n j_model.conv.kernel.value = conv_kernel\n j_model.conv.bias.value = conv_bias\n j_model.linear.kernel.value = fc_kernel\n j_model.linear.bias.value = fc_bias\n\n key = random.key(0)\n x = random.normal(key, (1, 6, 6, 3))\n j_out = j_model(x)\n\n t_x = torch.from_numpy(np.transpose(np.array(x), (0, 3, 1, 2)))\n t_out = t_model(t_x)\n t_out = t_out.detach().cpu().numpy()\n\n np.testing.assert_almost_equal(j_out, t_out, decimal=6)\n\n\n\nBatch Norm\n--------------------------------\n\n``torch.nn.BatchNorm2d`` uses ``0.1`` as the default value for the ``momentum`` parameter while\n|nnx.BatchNorm|_ uses ``0.9``. However, this corresponds to the same computation, because PyTorch multiplies\nthe estimated statistic with ``(1 − momentum)`` and the new observed value with ``momentum``,\nwhile Flax multiplies the estimated statistic with ``momentum`` and the new observed value with ``(1 − momentum)``.\n\n.. |nnx.BatchNorm| replace:: ``nnx.BatchNorm``\n.. _nnx.BatchNorm: https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/normalization.html#flax.nnx.BatchNorm\n\n.. testcode::\n\n t_bn = torch.nn.BatchNorm2d(num_features=3, momentum=0.1)\n t_bn.eval()\n\n key = random.key(0)\n x = random.normal(key, (1, 6, 6, 3))\n\n j_bn = nnx.BatchNorm(num_features=3, momentum=0.9, use_running_average=True, rngs=nnx.Rngs(0))\n\n j_out = j_bn(x)\n\n # [N, H, W, C] -> [N, C, H, W]\n t_x = torch.from_numpy(np.transpose(np.array(x), (0, 3, 1, 2)))\n t_out = t_bn(t_x)\n # [N, C, H, W] -> [N, H, W, C]\n t_out = np.transpose(t_out.detach().cpu().numpy(), (0, 2, 3, 1))\n\n np.testing.assert_almost_equal(j_out, t_out, decimal=6)\n\n\nAverage Pooling\n--------------------------------\n\n``torch.nn.AvgPool2d`` and |nnx.avg_pool()|_ are compatible when using default parameters.\nHowever, ``torch.nn.AvgPool2d`` has a parameter ``count_include_pad``. When ``count_include_pad=False``,\nthe zero-padding will not be considered for the average calculation. There does not exist a similar\nparameter for |nnx.avg_pool()|_. However, we can easily implement a wrapper around the pooling\noperation. ``nnx.pool()`` is the core function behind |nnx.avg_pool()|_ and |nnx.max_pool()|_.\n\n.. |nnx.avg_pool()| replace:: ``nnx.avg_pool()``\n.. _nnx.avg_pool(): https://flax.readthedocs.io/en/latest/api_reference/flax.linen/layers.html#flax.linen.avg_pool\n\n.. |nnx.max_pool()| replace:: ``nnx.max_pool()``\n.. _nnx.max_pool(): https://flax.readthedocs.io/en/latest/api_reference/flax.linen/layers.html#flax.linen.max_pool\n\n\n.. testcode::\n\n def avg_pool(inputs, window_shape, strides=None, padding='VALID'):\n \"\"\"\n Pools the input by taking the average over a window.\n In comparison to nnx.avg_pool(), this pooling operation does not\n consider the padded zero's for the average computation.\n \"\"\"\n assert len(window_shape) == 2\n\n y = nnx.pool(inputs, 0., jax.lax.add, window_shape, strides, padding)\n counts = nnx.pool(jnp.ones_like(inputs), 0., jax.lax.add, window_shape, strides, padding)\n y = y / counts\n return y\n\n\n key = random.key(0)\n x = random.normal(key, (1, 6, 6, 3))\n\n j_out = avg_pool(x, window_shape=(2, 2), strides=(1, 1), padding=((1, 1), (1, 1)))\n t_pool = torch.nn.AvgPool2d(kernel_size=2, stride=1, padding=1, count_include_pad=False)\n\n # [N, H, W, C] -> [N, C, H, W]\n t_x = torch.from_numpy(np.transpose(np.array(x), (0, 3, 1, 2)))\n t_out = t_pool(t_x)\n # [N, C, H, W] -> [N, H, W, C]\n t_out = np.transpose(t_out.detach().cpu().numpy(), (0, 2, 3, 1))\n\n np.testing.assert_almost_equal(j_out, t_out, decimal=6)\n\n\n\nTransposed Convolutions\n--------------------------------\n\n``torch.nn.ConvTranspose2d`` and |nnx.ConvTranspose|_ are not compatible.\n|nnx.ConvTranspose|_ is a wrapper around |jax.lax.conv_transpose|_ which computes a fractionally strided convolution,\nwhile ``torch.nn.ConvTranspose2d`` computes a gradient based transposed convolution. Currently, there is no\nimplementation of a gradient based transposed convolution is ``Jax``. However, there is a pending `pull request`_\nthat contains an implementation.\n\nTo load ``torch.nn.ConvTranspose2d`` parameters into Flax, we need to use the ``transpose_kernel`` arg in Flax's\n``nnx.ConvTranspose`` layer.\n\n.. testcode::\n\n t_conv = torch.nn.ConvTranspose2d(in_channels=3, out_channels=4, kernel_size=2, padding=0)\n\n kernel = t_conv.weight.detach().cpu().numpy()\n bias = t_conv.bias.detach().cpu().numpy()\n\n # [inC, outC, kH, kW] -> [kH, kW, outC, inC]\n kernel = jnp.transpose(kernel, (2, 3, 1, 0))\n\n key = random.key(0)\n x = random.normal(key, (1, 6, 6, 3))\n\n # ConvTranspose expects the kernel to be [kH, kW, inC, outC],\n # but with `transpose_kernel=True`, it expects [kH, kW, outC, inC] instead\n j_conv = nnx.ConvTranspose(3, 4, kernel_size=(2, 2), padding='VALID', transpose_kernel=True, rngs=nnx.Rngs(0))\n j_conv.kernel.value = kernel\n j_conv.bias.value = bias\n j_out = j_conv(x)\n\n # [N, H, W, C] -> [N, C, H, W]\n t_x = torch.from_numpy(np.transpose(np.array(x), (0, 3, 1, 2)))\n t_out = t_conv(t_x)\n # [N, C, H, W] -> [N, H, W, C]\n t_out = np.transpose(t_out.detach().cpu().numpy(), (0, 2, 3, 1))\n np.testing.assert_almost_equal(j_out, t_out, decimal=6)\n\n.. _`pull request`: https://github.com/jax-ml/jax/pull/5772\n\n.. |nnx.ConvTranspose| replace:: ``nnx.ConvTranspose``\n.. _nnx.ConvTranspose: https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/linear.html#flax.nnx.ConvTranspose\n\n.. |jax.lax.conv_transpose| replace:: ``jax.lax.conv_transpose``\n.. _jax.lax.conv_transpose: https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.conv_transpose.html\n" }, { "path": "docs_nnx/migrating/haiku_to_flax.rst", "content": "Haiku to Flax NNX\n############################\n\nThis guide demonstrates the differences between Haiku and Flax NNX models, providing side-by-side example code to help you migrate to the Flax NNX API from Haiku.\n\nIf you are new to Flax NNX, make sure you become familiarized with `Flax NNX basics `__, which covers the :class:`nnx.Module` system, `Flax transformations `__, and the `Functional API `__ with examples.\n\nLet’s start with some imports.\n\n.. testsetup:: Haiku, Flax NNX\n\n import jax\n import jax.numpy as jnp\n import optax\n from typing import Any\n\n\nBasic Module definition\n=======================\n\nBoth Haiku and Flax use the ``Module`` class as the default unit to express a neural network library layer. For example, to create a one-layer network with dropout and a ReLU activation function, you:\n\n* First, create a ``Block`` (by subclassing ``Module``) composed of one linear layer with dropout and a ReLU activation function.\n* Then, use ``Block`` as a sub-``Module`` when creating a ``Model`` (also by subclassing ``Module``), which is made up of ``Block`` and a linear layer.\n\nThere are two fundamental differences between Haiku and Flax ``Module`` objects:\n\n* **Stateless vs. stateful**:\n\n * A ``haiku.Module`` instance is stateless. This means, the variables are returned from a purely functional ``Module.init()`` call and managed separately.\n * A :class:`flax.nnx.Module`, however, owns its variables as attributes of this Python object.\n\n* **Lazy vs. eager**:\n\n * A ``haiku.Module`` only allocates space to create variables when they actually see the input when the user calls the model (lazy).\n * A ``flax.nnx.Module`` instance creates variables the moment they are instantiated, before seeing a sample input (eager).\n\n\n.. codediff::\n :title: Haiku, Flax NNX\n :sync:\n\n import haiku as hk\n\n class Block(hk.Module):\n def __init__(self, features: int, name=None):\n super().__init__(name=name)\n self.features = features\n\n def __call__(self, x, training: bool):\n x = hk.Linear(self.features)(x)\n x = hk.dropout(hk.next_rng_key(), 0.5 if training else 0, x)\n x = jax.nn.relu(x)\n return x\n\n class Model(hk.Module):\n def __init__(self, dmid: int, dout: int, name=None):\n super().__init__(name=name)\n self.dmid = dmid\n self.dout = dout\n\n def __call__(self, x, training: bool):\n x = Block(self.dmid)(x, training)\n x = hk.Linear(self.dout)(x)\n return x\n\n ---\n\n from flax import nnx\n\n class Block(nnx.Module):\n def __init__(self, in_features: int , out_features: int, rngs: nnx.Rngs):\n self.linear = nnx.Linear(in_features, out_features, rngs=rngs)\n self.dropout = nnx.Dropout(0.5, rngs=rngs)\n\n def __call__(self, x):\n x = self.linear(x)\n x = self.dropout(x)\n x = jax.nn.relu(x)\n return x\n\n class Model(nnx.Module):\n def __init__(self, din: int, dmid: int, dout: int, rngs: nnx.Rngs):\n self.block = Block(din, dmid, rngs=rngs)\n self.linear = nnx.Linear(dmid, dout, rngs=rngs)\n\n\n def __call__(self, x):\n x = self.block(x)\n x = self.linear(x)\n return x\n\n\nVariable creation\n=================\n\nThis section is about instantiating a model and initializing its parameters.\n\n* To generate model parameters for a Haiku model, you need to put it inside a forward function and use ``haiku.transform`` to make it purely functional. This results in a nested dictionary of `JAX Arrays `__ (``jax.Array`` data types) to be carried around and maintained separately.\n\n* In Flax NNX, the model parameters are automatically initialized when you instantiate the model, and the variables (:class:`nnx.Variable` objects) are stored inside the :class:`nnx.Module` (or its sub-Module) as attributes. You still need to provide it with a `pseudorandom number generator (PRNG) `__ key, but that key will be wrapped inside an :class:`nnx.Rngs` class and stored inside, generating more PRNG keys when needed.\n\nIf you want to access Flax model parameters in the stateless, dictionary-like fashion for checkpoint saving or model surgery, check out the `Flax NNX split/merge API `__ (:func:`nnx.split` / :func:`nnx.merge`).\n\n\n.. codediff::\n :title: Haiku, Flax NNX\n :sync:\n\n def forward(x, training: bool):\n return Model(256, 10)(x, training)\n\n model = hk.transform(forward)\n sample_x = jnp.ones((1, 784))\n params = model.init(jax.random.key(0), sample_x, training=False)\n\n\n assert params['model/linear']['b'].shape == (10,)\n assert params['model/block/linear']['w'].shape == (784, 256)\n\n ---\n\n ...\n\n\n model = Model(784, 256, 10, rngs=nnx.Rngs(0))\n\n\n # Parameters were already initialized during model instantiation.\n\n assert model.linear.bias.value.shape == (10,)\n assert model.block.linear.kernel.value.shape == (784, 256)\n\nTraining step and compilation\n=============================\n\nThis section covers writing a training step and compiling it using the `JAX just-in-time compilation `__.\n\nWhen compiling the training step:\n\n* Haiku uses ``@jax.jit`` - a `JAX transformation `__ - to compile a purely functional training step.\n* Flax NNX uses :meth:`@nnx.jit` - a `Flax NNX transformation `__ (one of several transform APIs that behave similarly to JAX transforms, but also `work well with Flax objects `__). While ``jax.jit`` only accepts functions with pure stateless arguments, ``flax.nnx.jit`` allows the arguments to be stateful Modules. This greatly reduces the number of lines needed for a train step.\n\nWhen taking gradients:\n\n* Similarly, Haiku uses ``jax.grad`` (a JAX transformation for `automatic differentiation `__) to return a raw dictionary of gradients.\n* Meanwhile, Flax NNX uses :meth:`flax.nnx.grad` (a Flax NNX transformation) to return the gradients of Flax NNX Modules as :class:`flax.nnx.State` dictionaries. If you want to use regular ``jax.grad`` with Flax NNX, you need to use the `split/merge API `__.\n\nFor optimizers:\n\n* If you are already using `Optax `__ optimizers like ``optax.adamw`` (instead of the raw ``jax.tree.map`` computation shown here) with Haiku, check out the :class:`flax.nnx.Optimizer` example in the `Flax basics `__ guide for a much more concise way of training and updating your model.\n\nModel updates during each training step:\n\n* The Haiku training step needs to return a `JAX pytree `__ of parameters as the input of the next step.\n* The Flax NNX training step does not need to return anything, because the ``model`` was already updated in-place within :meth:`nnx.jit`.\n* In addition, :class:`nnx.Module` objects are stateful, and ``Module`` automatically tracks several things within it, such as PRNG keys and ``flax.nnx.BatchNorm`` stats. That is why you don't need to explicitly pass a PRNG key in at every step. Also note that you can use :meth:`flax.nnx.reseed` to reset its underlying PRNG state.\n\nThe dropout behavior:\n\n* In Haiku, you need to explicitly define and pass in the ``training`` argument to toggle ``haiku.dropout`` and make sure that random dropout only happens if ``training=True``.\n* In Flax NNX, you can call ``model.train()`` (:meth:`flax.nnx.Module.train`) to automatically switch :class:`flax.nnx.Dropout` to the training mode. Conversely, you can call ``model.eval()`` (:meth:`flax.nnx.Module.eval`) to turn off the training mode. You can learn more about what ``flax.nnx.Module.train`` does in its `API reference `__.\n\n.. codediff::\n :title: Haiku, Flax NNX\n :sync:\n\n ...\n\n @jax.jit\n def train_step(key, params, inputs, labels):\n def loss_fn(params):\n logits = model.apply(\n params, key,\n inputs, training=True # <== inputs\n\n )\n return optax.softmax_cross_entropy_with_integer_labels(logits, labels).mean()\n\n grads = jax.grad(loss_fn)(params)\n\n\n params = jax.tree_util.tree_map(lambda p, g: p - 0.1 * g, params, grads)\n\n return params\n\n ---\n\n model.train() # set deterministic=False\n\n @nnx.jit\n def train_step(model, inputs, labels):\n def loss_fn(model):\n logits = model(\n\n inputs, # <== inputs\n\n )\n return optax.softmax_cross_entropy_with_integer_labels(logits, labels).mean()\n\n grads = nnx.grad(loss_fn)(model)\n _, params, rest = nnx.split(model, nnx.Param, ...)\n params = jax.tree.map(lambda p, g: p - 0.1 * g, params, grads)\n nnx.update(model, nnx.merge_state(params, rest))\n\n.. testcode:: Haiku\n :hide:\n\n train_step(jax.random.key(0), params, sample_x, jnp.ones((1,), dtype=jnp.int32))\n\n.. testcode:: Flax NNX\n :hide:\n\n sample_x = jnp.ones((1, 784))\n train_step(model, sample_x, jnp.ones((1,), dtype=jnp.int32))\n\n\n\nHandling non-parameter states\n=============================\n\nHaiku makes a distinction between trainable parameters and all other data (\"states\") that the model tracks. For example, the batch stats used in batch norm is considered a state. Models with states needs to be transformed with ``hk.transform_with_state`` so that their ``.init()`` returns both params and states.\n\nIn Flax, there isn't such a strong distinction - they are all subclasses of ``nnx.Variable`` and seen by a module as its attributes. Parameters are instances of a subclass called ``nnx.Param``, and batch stats can be of another subclass called ``nnx.BatchStat``. You can use :func:`nnx.split` to quickly extract all data of a certain variable type.\n\nLet's see an example of this by taking the ``Block`` definition above but replace dropout with ``BatchNorm``.\n\n.. codediff::\n :title: Haiku, Flax NNX\n :sync:\n\n class Block(hk.Module):\n def __init__(self, features: int, name=None):\n super().__init__(name=name)\n self.features = features\n\n\n\n def __call__(self, x, training: bool):\n x = hk.Linear(self.features)(x)\n x = hk.BatchNorm(\n create_scale=True, create_offset=True, decay_rate=0.99\n )(x, is_training=training)\n x = jax.nn.relu(x)\n return x\n\n def forward(x, training: bool):\n return Model(256, 10)(x, training)\n model = hk.transform_with_state(forward)\n\n sample_x = jnp.ones((1, 784))\n params, batch_stats = model.init(jax.random.key(0), sample_x, training=True)\n\n ---\n\n class Block(nnx.Module):\n def __init__(self, in_features: int , out_features: int, rngs: nnx.Rngs):\n self.linear = nnx.Linear(in_features, out_features, rngs=rngs)\n self.batchnorm = nnx.BatchNorm(\n num_features=out_features, momentum=0.99, rngs=rngs\n )\n\n def __call__(self, x):\n x = self.linear(x)\n x = self.batchnorm(x)\n\n\n x = jax.nn.relu(x)\n return x\n\n\n\n model = Block(4, 4, rngs=nnx.Rngs(0))\n\n model.linear.kernel # Param(value=...)\n model.batchnorm.mean # BatchStat(value=...)\n\n\nFlax takes the difference of trainable params and other data into account. ``nnx.grad`` will only take gradients on the ``nnx.Param`` variables, thus skipping the ``batchnorm`` arrays automatically. Therefore, the training step will look the same for Flax NNX with this model.\n\n\nUsing multiple methods\n======================\n\nIn this section you will learn how to use multiple methods in Haiku and Flax. As an example, you will implement an auto-encoder model with three methods: ``encode``, ``decode``, and ``__call__``.\n\nIn Haiku, you need to use ``hk.multi_transform`` to explicitly define how the model shall be initialized and what methods (``encode`` and ``decode`` here) it can call. Note that you still need to define a ``__call__`` that activates both layers for the lazy initialization of all model parameters.\n\nIn Flax, it's simpler as you initialized parameters in ``__init__`` and the :class:`nnx.Module` methods ``encode`` and ``decode`` can be used directly.\n\n.. codediff::\n :title: Haiku, Flax NNX\n :sync:\n\n class AutoEncoder(hk.Module):\n\n def __init__(self, embed_dim: int, output_dim: int, name=None):\n super().__init__(name=name)\n self.encoder = hk.Linear(embed_dim, name=\"encoder\")\n self.decoder = hk.Linear(output_dim, name=\"decoder\")\n\n def encode(self, x):\n return self.encoder(x)\n\n def decode(self, x):\n return self.decoder(x)\n\n def __call__(self, x):\n x = self.encode(x)\n x = self.decode(x)\n return x\n\n def forward():\n module = AutoEncoder(256, 784)\n init = lambda x: module(x)\n return init, (module.encode, module.decode)\n\n model = hk.multi_transform(forward)\n params = model.init(jax.random.key(0), x=jnp.ones((1, 784)))\n\n ---\n\n class AutoEncoder(nnx.Module):\n\n def __init__(self, in_dim: int, embed_dim: int, output_dim: int, rngs):\n\n self.encoder = nnx.Linear(in_dim, embed_dim, rngs=rngs)\n self.decoder = nnx.Linear(embed_dim, output_dim, rngs=rngs)\n\n def encode(self, x):\n return self.encoder(x)\n\n def decode(self, x):\n return self.decoder(x)\n\n\n\n\n\n\n\n\n\n\n\n model = AutoEncoder(784, 256, 784, rngs=nnx.Rngs(0))\n ...\n\n\nThe parameter structure is as follows:\n\n.. tab-set::\n\n .. tab-item:: Haiku\n :sync: Haiku\n\n .. code-block:: python\n\n ...\n\n\n {\n 'auto_encoder/~/decoder': {\n 'b': (784,),\n 'w': (256, 784)\n },\n 'auto_encoder/~/encoder': {\n 'b': (256,),\n 'w': (784, 256)\n }\n }\n\n .. tab-item:: Flax NNX\n :sync: Flax NNX\n\n .. code-block:: python\n\n _, params, _ = nnx.split(model, nnx.Param, ...)\n\n params\n {\n 'decoder': {\n 'bias': Param(value=(784,)),\n 'kernel': Param(value=(256, 784))\n },\n 'encoder': {\n 'bias': Param(value=(256,)),\n 'kernel': Param(value=(784, 256))\n }\n }\n\n\nTo call those custom methods:\n\n* In Haiku, you need to decouple the `.apply` function to extract your method before calling it.\n* In Flax, you can simply call the method directly.\n\n.. codediff::\n :title: Haiku, Flax NNX\n :sync:\n\n encode, decode = model.apply\n z = encode(params, None, x=jnp.ones((1, 784)))\n\n ---\n\n ...\n z = model.encode(jnp.ones((1, 784)))\n\n\n\nTransformations\n=======================\n\nBoth Haiku and `Flax transformations `__ provide their own set of transforms that wrap `JAX transforms `__ in a way that they can be used with ``Module`` objects.\n\nFor more information on Flax transforms, check out the `Transforms guide `__.\n\nLet's start with an example:\n\n* First, define an ``RNNCell`` ``Module`` that will contain the logic for a single step of the RNN.\n* Define a ``initial_state`` method that will be used to initialize the state (a.k.a. ``carry``) of the RNN. Like with ``jax.lax.scan`` (`API doc `__), the ``RNNCell.__call__`` method will be a function that takes the carry and input, and returns the new carry and output. In this case, the carry and the output are the same.\n\n\n.. codediff::\n :title: Haiku, Flax NNX\n :sync:\n\n class RNNCell(hk.Module):\n def __init__(self, hidden_size: int, name=None):\n super().__init__(name=name)\n self.hidden_size = hidden_size\n\n def __call__(self, carry, x):\n x = jnp.concatenate([carry, x], axis=-1)\n x = hk.Linear(self.hidden_size)(x)\n x = jax.nn.relu(x)\n return x, x\n\n def initial_state(self, batch_size: int):\n return jnp.zeros((batch_size, self.hidden_size))\n\n ---\n\n class RNNCell(nnx.Module):\n def __init__(self, input_size, hidden_size, rngs):\n self.linear = nnx.Linear(hidden_size + input_size, hidden_size, rngs=rngs)\n self.hidden_size = hidden_size\n\n def __call__(self, carry, x):\n x = jnp.concatenate([carry, x], axis=-1)\n x = self.linear(x)\n x = jax.nn.relu(x)\n return x, x\n\n def initial_state(self, batch_size: int):\n return jnp.zeros((batch_size, self.hidden_size))\n\nNext, we will define a ``RNN`` Module that will contain the logic for the entire RNN. In both cases, we use the library's ``scan`` call to run the ``RNNCell`` over the input sequence.\n\nThe only difference is that Flax ``nnx.scan`` allows you to specify which axis to repeat over in arguments ``in_axes`` and ``out_axes``, which will be forwarded to the underlying `jax.lax.scan `__, whereas in Haiku you need to transpose the input and output explicitly.\n\n.. codediff::\n :title: Haiku, Flax NNX\n :sync:\n\n class RNN(hk.Module):\n def __init__(self, hidden_size: int, name=None):\n super().__init__(name=name)\n self.hidden_size = hidden_size\n\n def __call__(self, x):\n cell = RNNCell(self.hidden_size)\n carry = cell.initial_state(x.shape[0])\n carry, y = hk.scan(\n cell, carry,\n jnp.swapaxes(x, 1, 0)\n )\n y = jnp.swapaxes(y, 0, 1)\n return y\n\n ---\n\n class RNN(nnx.Module):\n def __init__(self, input_size: int, hidden_size: int, rngs: nnx.Rngs):\n self.hidden_size = hidden_size\n self.cell = RNNCell(input_size, self.hidden_size, rngs=rngs)\n\n def __call__(self, x):\n scan_fn = lambda carry, cell, x: cell(carry, x)\n carry = self.cell.initial_state(x.shape[0])\n carry, y = nnx.scan(\n scan_fn, in_axes=(nnx.Carry, None, 1), out_axes=(nnx.Carry, 1)\n )(carry, self.cell, x)\n\n return y\n\n\nScan over layers\n=======================\n\nMost Haiku transforms should look similar with Flax, since they all wraps their JAX counterparts, but the scan-over-layers use case is an exception.\n\nScan-over-layers is a technique where you run an input through a sequence of N repeated layers, passing the output of each layer as the input to the next layer. This pattern can significantly reduce compilation time for large models. In the example below, you will repeat the ``Block`` ``Module`` 5 times in the top-level ``MLP`` ``Module``.\n\nIn Haiku, we define the ``Block`` Module as usual, and then inside ``MLP`` we will\nuse ``hk.experimental.layer_stack`` over a ``stack_block`` function to create a stack\nof ``Block`` Modules. The same code will create 5 layers of parameters in initialization time, and run the input through them in call time.\n\nIn Flax, model initialization and calling code are completely decoupled, so we use the :func:`nnx.vmap` transform to initialize the underlying ``Block`` parameters, and the :func:`nnx.scan` transform to run the model input through them.\n\n.. codediff::\n :title: Haiku, Flax NNX\n :sync:\n\n class Block(hk.Module):\n def __init__(self, features: int, name=None):\n super().__init__(name=name)\n self.features = features\n\n def __call__(self, x, training: bool):\n x = hk.Linear(self.features)(x)\n x = hk.dropout(hk.next_rng_key(), 0.5 if training else 0, x)\n x = jax.nn.relu(x)\n return x\n\n class MLP(hk.Module):\n def __init__(self, features: int, num_layers: int, name=None):\n super().__init__(name=name)\n self.features = features\n self.num_layers = num_layers\n\n\n\n\n\n def __call__(self, x, training: bool):\n\n @hk.experimental.layer_stack(self.num_layers)\n def stack_block(x):\n return Block(self.features)(x, training)\n\n stack = hk.experimental.layer_stack(self.num_layers)\n return stack_block(x)\n\n def forward(x, training: bool):\n return MLP(64, num_layers=5)(x, training)\n model = hk.transform(forward)\n\n sample_x = jnp.ones((1, 64))\n params = model.init(jax.random.key(0), sample_x, training=False)\n\n ---\n\n class Block(nnx.Module):\n def __init__(self, input_dim, features, rngs):\n self.linear = nnx.Linear(input_dim, features, rngs=rngs)\n self.dropout = nnx.Dropout(0.5, rngs=rngs)\n\n def __call__(self, x: jax.Array): # No need to require a second input!\n x = self.linear(x)\n x = self.dropout(x)\n x = jax.nn.relu(x)\n return x # No need to return a second output!\n\n class MLP(nnx.Module):\n def __init__(self, features, num_layers, rngs):\n @nnx.split_rngs(splits=num_layers)\n @nnx.vmap(in_axes=(0,), out_axes=0)\n def create_block(rngs: nnx.Rngs):\n return Block(features, features, rngs=rngs)\n\n self.blocks = create_block(rngs)\n self.num_layers = num_layers\n\n def __call__(self, x):\n @nnx.split_rngs(splits=self.num_layers)\n @nnx.scan(in_axes=(nnx.Carry, 0), out_axes=nnx.Carry)\n def forward(x, model):\n x = model(x)\n return x\n\n return forward(x, self.blocks)\n\n\n\n model = MLP(64, num_layers=5, rngs=nnx.Rngs(0))\n\n\nThere are a few other details to explain in the Flax example above:\n\n* **The `@nnx.split_rngs` decorator:** Flax transforms, like their JAX counterparts, are completely agnostic of the PRNG state and rely on input for PRNG keys. The ``nnx.split_rngs`` decorator allows you to split the ``nnx.Rngs`` before passing them to the decorated function and 'lower' them afterwards, so they can be used outside.\n\n * Here, you split the PRNG keys because ``jax.vmap`` and ``jax.lax.scan`` require a list of PRNG keys if each of its internal operations needs its own key. So for the 5 layers inside the ``MLP``, you split and provide 5 different PRNG keys from its arguments before going down to the JAX transform.\n\n * Note that actually ``create_block()`` knows it needs to create 5 layers *precisely because* it sees 5 PRNG keys, because ``in_axes=(0,)`` indicates that ``vmap`` will look into the first argument's first dimension to know the size it will map over.\n\n * Same goes for ``forward()``, which looks at the variables inside the first argument (aka. ``model``) to find out how many times it needs to scan. ``nnx.split_rngs`` here actually splits the PRNG state inside the ``model``. (If the ``Block`` ``Module`` doesn't have dropout, you don't need the :meth:`nnx.split_rngs` line as it would not consume any PRNG key anyway.)\n\n* **Why the Block Module in Flax doesn't need to take and return that extra dummy value:** ``jax.lax.scan`` `(API doc `__ requires its function to return two inputs - the carry and the stacked output. In this case, we didn't use the latter. Flax simplifies this, so that you can now choose to ignore the second output if you set ``out_axes=nnx.Carry`` instead of the default ``(nnx.Carry, 0)``.\n\n * This is one of the rare cases where Flax NNX transforms diverge from the `JAX transforms `__ APIs.\n\nThere are more lines of code in the Flax example above, but they express what happens at each time more precisely. Since Flax transforms become way closer to the JAX transform APIs, it is recommended to have a good understanding of the underlying `JAX transforms `__ before using their `Flax NNX equivalents `__\n\nNow inspect the variable pytree on both sides:\n\n.. tab-set::\n\n .. tab-item:: Haiku\n :sync: Haiku\n\n .. code-block:: python\n\n ...\n\n\n {\n 'mlp/__layer_stack_no_per_layer/block/linear': {\n 'b': (5, 64),\n 'w': (5, 64, 64)\n }\n }\n\n\n\n ...\n\n .. tab-item:: Flax NNX\n :sync: Flax NNX\n\n .. code-block:: python\n\n _, params, _ = nnx.split(model, nnx.Param, ...)\n\n params\n {\n 'blocks': {\n 'linear': {\n 'bias': Param(value=(5, 64)),\n 'kernel': Param(value=(5, 64, 64))\n }\n }\n }\n\n\nTop-level Haiku functions vs top-level Flax modules\n=======================\n\nIn Haiku, it is possible to write the entire model as a single function by using\nthe raw ``hk.{get,set}_{parameter,state}`` to define/access model parameters and\nstates. It is very common to write the top-level \"Module\" as a function instead.\n\nThe Flax team recommends a more Module-centric approach that uses ``__call__`` to\ndefine the forward function. In Flax modules, the parameters and variables can\nbe set and accessed as normal using regular Python class semantics.\n\n.. codediff::\n :title: Haiku, Flax NNX\n :sync:\n\n ...\n\n\n def forward(x):\n\n\n counter = hk.get_state('counter', shape=[], dtype=jnp.int32, init=jnp.ones)\n multiplier = hk.get_parameter(\n 'multiplier', shape=[1,], dtype=x.dtype, init=jnp.ones\n )\n\n output = x + multiplier * counter\n\n hk.set_state(\"counter\", counter + 1)\n return output\n\n model = hk.transform_with_state(forward)\n\n params, state = model.init(jax.random.key(0), jnp.ones((1, 64)))\n\n ---\n\n class Counter(nnx.Variable):\n pass\n\n class FooModule(nnx.Module):\n\n def __init__(self, rngs):\n self.counter = Counter(jnp.ones((), jnp.int32))\n self.multiplier = nnx.Param(\n nnx.initializers.ones(rngs.params(), [1,], jnp.float32)\n )\n def __call__(self, x):\n output = x + self.multiplier * self.counter.value\n\n self.counter.value += 1\n return output\n\n model = FooModule(rngs=nnx.Rngs(0))\n\n _, params, counter = nnx.split(model, nnx.Param, Counter)\n\n\n\n\n" }, { "path": "docs_nnx/migrating/index.rst", "content": "Migrating\n------------------------\n\n.. toctree::\n :maxdepth: 2\n\n convert_pytorch_to_flax\n nnx_010_to_nnx_011\n linen_to_nnx\n haiku_to_flax\n" }, { "path": "docs_nnx/migrating/linen_to_nnx.rst", "content": "Flax Linen to Flax NNX\n################################\n\nThis guide demonstrates the differences between Flax Linen and Flax NNX models, providing side-by-side example code to help you migrate to the Flax NNX API from Flax Linen.\n\nThis document mainly teaches how to convert arbitrary Flax Linen code to Flax NNX. If you want to play it “safe” and convert your codebase iteratively, check out the `Use Flax NNX and Linen together via nnx.bridge `__ guide.\n\nTo get the most out of this guide, it is highly recommended to get go through `Flax NNX basics `__ document, which covers the :class:`nnx.Module` system, `Flax transformations `__, and the `Functional API `__ with examples.\n\n.. testsetup:: Linen, NNX\n\n import jax\n import jax.numpy as jnp\n import optax\n import flax.linen as nn\n from typing import Any\n\nBasic ``Module`` definition\n===========================\n\nBoth Flax Linen and Flax NNX use the ``Module`` class as the default unit to express a neural network library layer. In the example below, you first create a ``Block`` (by subclassing ``Module``) composed of one linear layer with dropout and a ReLU activation function; then you use it as a sub-``Module`` when creating a ``Model`` (also by subclassing ``Module``), which is made up of ``Block`` and a linear layer.\n\nThere are two fundamental differences between Flax Linen and Flax NNX ``Module`` objects:\n\n* **Stateless vs. stateful**: A ``flax.linen.Module`` (``nn.Module``) instance is stateless - the variables are returned from a purely functional ``Module.init()`` call and managed separately. A :class:`flax.nnx.Module`, however, owns its variables as attributes of this Python object.\n\n* **Lazy vs. eager**: A ``flax.linen.Module`` only allocates space to create variables when they actually see their input (lazy). A :class:`flax.nnx.Module` instance creates variables the moment they are instantiated before seeing a sample input (eager).\n\n* Flax Linen can use the ``@nn.compact`` decorator to define the model in a single method, and use shape inference from the input sample. A Flax NNX ``Module`` generally requests additional shape information to create all parameters during ``__init__`` , and separately defines the computation in the ``__call__`` method.\n\n.. codediff::\n :title: Linen, NNX\n :sync:\n\n import flax.linen as nn\n\n class Block(nn.Module):\n features: int\n\n\n @nn.compact\n def __call__(self, x, training: bool):\n x = nn.Dense(self.features)(x)\n x = nn.Dropout(0.5, deterministic=not training)(x)\n x = jax.nn.relu(x)\n return x\n\n class Model(nn.Module):\n dmid: int\n dout: int\n\n @nn.compact\n def __call__(self, x, training: bool):\n x = Block(self.dmid)(x, training)\n x = nn.Dense(self.dout)(x)\n return x\n\n ---\n\n from flax import nnx\n\n class Block(nnx.Module):\n def __init__(self, in_features: int , out_features: int, rngs: nnx.Rngs):\n self.linear = nnx.Linear(in_features, out_features, rngs=rngs)\n self.dropout = nnx.Dropout(0.5, rngs=rngs)\n\n def __call__(self, x):\n x = self.linear(x)\n x = self.dropout(x)\n x = jax.nn.relu(x)\n return x\n\n class Model(nnx.Module):\n def __init__(self, din: int, dmid: int, dout: int, rngs: nnx.Rngs):\n self.block = Block(din, dmid, rngs=rngs)\n self.linear = nnx.Linear(dmid, dout, rngs=rngs)\n\n def __call__(self, x):\n x = self.block(x)\n x = self.linear(x)\n return x\n\n\nVariable creation\n=================\n\nNext, let’s discuss instantiating the model and initializing its parameters:\n\n* To generate model parameters for a Flax Linen model, you call the ``flax.linen.Module.init`` (``nn.Module.init``) method with a ``jax.random.key`` (`doc `__) plus some sample inputs that the model shall take. This results in a nested dictionary of `JAX Arrays `__ (``jax.Array`` data types) to be carried around and maintained separately.\n\n* In Flax NNX, the model parameters are automatically initialized when you instantiate the model, and the variables (:class:`nnx.Variable` objects) are stored inside the :class:`nnx.Module` (or its sub-``Module``) as attributes. You still need to provide it with a `pseudorandom number generator (PRNG) `__ key, but that key will be wrapped inside an :class:`nnx.Rngs` class and stored inside, generating more PRNG keys when needed.\n\nIf you want to access Flax NNX model parameters in the stateless, dictionary-like fashion for checkpoint saving or model surgery, check out the `Flax NNX split/merge API `__ (:func:`nnx.split` / :func:`nnx.merge`).\n\n.. codediff::\n :title: Linen, NNX\n :sync:\n\n model = Model(256, 10)\n sample_x = jnp.ones((1, 784))\n variables = model.init(jax.random.key(0), sample_x, training=False)\n params = variables[\"params\"]\n\n assert params['Dense_0']['bias'].shape == (10,)\n assert params['Block_0']['Dense_0']['kernel'].shape == (784, 256)\n\n ---\n\n model = Model(784, 256, 10, rngs=nnx.Rngs(0))\n\n\n # Parameters were already initialized during model instantiation.\n\n assert model.linear.bias.value.shape == (10,)\n assert model.block.linear.kernel.value.shape == (784, 256)\n\n\nTraining step and compilation\n=============================\n\nNow, let’s proceed to writing a training step and compiling it using `JAX just-in-time compilation `__. Below are certain differences between Flax Linen and Flax NNX approaches.\n\nCompiling the training step:\n\n* Flax Linen uses ``@jax.jit`` - a `JAX transform `__ - to compile the training step.\n* Flax NNX uses :meth:`@nnx.jit` - a `Flax NNX transform `__ (one of several transform APIs that behave similarly to JAX transforms, but also `work well with Flax NNX objects `__). So, while ``jax.jit`` only accepts functions pure stateless arguments, ``nnx.jit`` allows the arguments to be stateful NNX Modules. This greatly reduced the number of lines needed for a train step.\n\nTaking gradients:\n\n* Similarly, Flax Linen uses ``jax.grad`` (a JAX transform for `automatic differentiation `__) to return a raw dictionary of gradients.\n* Flax NNX uses :meth:`nnx.grad` (a Flax NNX transform) to return the gradients of NNX Modules as :class:`nnx.State` dictionaries. If you want to use regular ``jax.grad`` with Flax NNX you need to use the `Flax NNX split/merge API `__.\n\nOptimizers:\n\n* If you are already using `Optax `__ optimizers like ``optax.adamw`` (instead of the raw ``jax.tree.map`` computation shown here) with Flax Linen, check out the :class:`nnx.Optimizer` example in the `Flax NNX basics `__ guide for a much more concise way of training and updating your model.\n\nModel updates during each training step:\n\n* The Flax Linen training step needs to return a `pytree `__ of parameters as the input of the next step.\n* The Flax NNX training step doesn't need to return anything, because the ``model`` was already updated in-place within :meth:`nnx.jit`.\n* In addition, :class:`nnx.Module` objects are stateful, and ``Module`` automatically tracks several things within it, such as PRNG keys and ``BatchNorm`` stats. That is why you don't need to explicitly pass an PRNG key in on every step. Also note that you can use :meth:`nnx.reseed` to reset its underlying PRNG state.\n\nDropout behavior:\n\n* In Flax Linen, you need to explicitly define and pass in the ``training`` argument to control the behavior of ``flax.linen.Dropout`` (``nn.Dropout``), namely, its ``deterministic`` flag, which means random dropout only happens if ``training=True``.\n* In Flax NNX, you can call ``model.train()`` (:meth:`flax.nnx.Module.train`) to automatically switch :class:`nnx.Dropout` to the training mode. Conversely, you can call ``model.eval()`` (:meth:`flax.nnx.Module.eval`) to turn off the training mode. You can learn more about what ``nnx.Module.train`` does in its `API reference `__.\n\n\n.. codediff::\n :title: Linen, NNX\n :sync:\n\n ...\n\n @jax.jit\n def train_step(key, params, inputs, labels):\n def loss_fn(params):\n logits = model.apply(\n {'params': params},\n inputs, training=True, # <== inputs\n rngs={'dropout': key}\n )\n return optax.softmax_cross_entropy_with_integer_labels(logits, labels).mean()\n\n grads = jax.grad(loss_fn)(params)\n\n params = jax.tree.map(lambda p, g: p - 0.1 * g, params, grads)\n return params\n\n ---\n\n model.train() # Sets ``deterministic=False` under the hood for nnx.Dropout\n\n @nnx.jit\n def train_step(model, inputs, labels):\n def loss_fn(model):\n logits = model(inputs)\n\n\n\n\n return optax.softmax_cross_entropy_with_integer_labels(logits, labels).mean()\n\n grads = nnx.grad(loss_fn)(model)\n _, params, rest = nnx.split(model, nnx.Param, ...)\n params = jax.tree.map(lambda p, g: p - 0.1 * g, params, grads)\n nnx.update(model, nnx.merge_state(params, rest))\n\n.. testcode:: Linen\n :hide:\n\n train_step(jax.random.key(0), params, sample_x, jnp.ones((1,), dtype=jnp.int32))\n\n.. testcode:: NNX\n :hide:\n\n sample_x = jnp.ones((1, 784))\n train_step(model, sample_x, jnp.ones((1,), dtype=jnp.int32))\n\n\nCollections and variable types\n==============================\n\nOne key difference between Flax Linen and NNX APIs is how they group variables into categories. Flax Linen uses different collections, while Flax NNX, since all variables shall be top-level Python attributes, you use different variable types.\n\nIn Flax NNX, you can freely create your own variable types as subclasses of ``nnx.Variable``.\n\nFor all the built-in Flax Linen layers and collections, Flax NNX already creates the corresponding layers and variable types. For example:\n\n* ``flax.linen.Dense`` (``nn.Dense``) creates ``params`` -> :class:`nnx.Linear` creates :class:nnx.Param`.\n* ``flax.linen.BatchNorm`` (``nn.BatchNorm``) creates ``batch_stats`` -> :class:`nnx.BatchNorm` creates :class:`nnx.BatchStats`.\n* ``flax.linen.Module.sow()`` creates ``intermediates`` -> :class:`nnx.Module.sow()` creates :class:`nnx.Intermediaries`.\n* In Flax NNX, you can also simply obtain the intermediates by assigning it to an ``nnx.Module`` attribute - for example, ``self.sowed = nnx.Intermediates(x)``. This will be similar to Flax Linen's ``self.variable('intermediates' 'sowed', lambda: x)``.\n\n.. codediff::\n :title: Linen, NNX\n :sync:\n\n class Block(nn.Module):\n features: int\n def setup(self):\n self.dense = nn.Dense(self.features)\n self.batchnorm = nn.BatchNorm(momentum=0.99)\n self.count = self.variable('counter', 'count',\n lambda: jnp.zeros((), jnp.int32))\n\n\n @nn.compact\n def __call__(self, x, training: bool):\n x = self.dense(x)\n x = self.batchnorm(x, use_running_average=not training)\n self.count.value += 1\n x = jax.nn.relu(x)\n return x\n\n x = jax.random.normal(jax.random.key(0), (2, 4))\n model = Block(4)\n variables = model.init(jax.random.key(0), x, training=True)\n variables['params']['dense']['kernel'].shape # (4, 4)\n variables['batch_stats']['batchnorm']['mean'].shape # (4, )\n variables['counter']['count'] # 1\n\n ---\n\n class Counter(nnx.Variable): pass\n\n class Block(nnx.Module):\n def __init__(self, in_features: int , out_features: int, rngs: nnx.Rngs):\n self.linear = nnx.Linear(in_features, out_features, rngs=rngs)\n self.batchnorm = nnx.BatchNorm(\n num_features=out_features, momentum=0.99, rngs=rngs\n )\n self.count = Counter(jnp.array(0))\n\n def __call__(self, x):\n x = self.linear(x)\n x = self.batchnorm(x)\n self.count.value += 1\n x = jax.nn.relu(x)\n return x\n\n\n\n model = Block(4, 4, rngs=nnx.Rngs(0))\n\n model.linear.kernel # Param(value=...)\n model.batchnorm.mean # BatchStat(value=...)\n model.count # Counter(value=...)\n\nIf you want to extract certain arrays from the pytree of variables:\n\n* In Flax Linen, you can access the specific dictionary path.\n* In Flax NNX, you can use :func:`nnx.split` to distinguish the types apart in Flax NNX. The code below is a simple example that splits up the variables by their types - check out the `Flax NNX Filters `__ guide for more sophisticated filtering expressions.\n\n.. codediff::\n :title: Linen, NNX\n :sync:\n\n params, batch_stats, counter = (\n variables['params'], variables['batch_stats'], variables['counter'])\n params.keys() # ['dense', 'batchnorm']\n batch_stats.keys() # ['batchnorm']\n counter.keys() # ['count']\n\n # ... make arbitrary modifications ...\n # Merge back with raw dict to carry on:\n variables = {'params': params, 'batch_stats': batch_stats, 'counter': counter}\n\n ---\n\n graphdef, params, batch_stats, count = nnx.split(\n model, nnx.Param, nnx.BatchStat, Counter)\n params.keys() # ['batchnorm', 'linear']\n batch_stats.keys() # ['batchnorm']\n count.keys() # ['count']\n\n # ... make arbitrary modifications ...\n # Merge back with ``nnx.merge`` to carry on:\n model = nnx.merge(graphdef, params, batch_stats, count)\n\n\n\nUsing multiple methods\n======================\n\nIn this section you will learn how to use multiple methods in both Flax Linen and Flax NNX. As an example, you will implement an auto-encoder model with three methods: ``encode``, ``decode``, and ``__call__``.\n\nDefining the encoder and decoder layers:\n\n* In Flax Linen, as before, define the layers without having to pass in the input shape, since the ``flax.linen.Module`` parameters will be initialized lazily using shape inference.\n* In Flax NNX, you must pass in the input shape since the :class:`nnx.Module` parameters will be initialized eagerly without shape inference.\n\n.. codediff::\n :title: Linen, NNX\n :sync:\n\n class AutoEncoder(nn.Module):\n embed_dim: int\n output_dim: int\n\n def setup(self):\n self.encoder = nn.Dense(self.embed_dim)\n self.decoder = nn.Dense(self.output_dim)\n\n def encode(self, x):\n return self.encoder(x)\n\n def decode(self, x):\n return self.decoder(x)\n\n def __call__(self, x):\n x = self.encode(x)\n x = self.decode(x)\n return x\n\n model = AutoEncoder(256, 784)\n variables = model.init(jax.random.key(0), x=jnp.ones((1, 784)))\n\n ---\n\n class AutoEncoder(nnx.Module):\n\n\n\n def __init__(self, in_dim: int, embed_dim: int, output_dim: int, rngs):\n self.encoder = nnx.Linear(in_dim, embed_dim, rngs=rngs)\n self.decoder = nnx.Linear(embed_dim, output_dim, rngs=rngs)\n\n def encode(self, x):\n return self.encoder(x)\n\n def decode(self, x):\n return self.decoder(x)\n\n def __call__(self, x):\n x = self.encode(x)\n x = self.decode(x)\n return x\n\n model = AutoEncoder(784, 256, 784, rngs=nnx.Rngs(0))\n\n\nThe variable structure is as follows:\n\n.. tab-set::\n\n .. tab-item:: Linen\n :sync: Linen\n\n .. code-block:: python\n\n\n # variables['params']\n {\n decoder: {\n bias: (784,),\n kernel: (256, 784),\n },\n encoder: {\n bias: (256,),\n kernel: (784, 256),\n },\n }\n\n .. tab-item:: NNX\n :sync: NNX\n\n .. code-block:: python\n\n # _, params, _ = nnx.split(model, nnx.Param, ...)\n # params\n {\n 'decoder': {\n 'bias': Param(value=(784,)),\n 'kernel': Param(value=(256, 784))\n },\n 'encoder': {\n 'bias': Param(value=(256,)),\n 'kernel': Param(value=(784, 256))\n }\n }\n\nTo call methods other than ``__call__``:\n\n* In Flax Linen, you still need to use the ``apply`` API.\n* In Flax NNX, you can simply call the method directly.\n\n.. codediff::\n :title: Linen, NNX\n :sync:\n\n z = model.apply(variables, x=jnp.ones((1, 784)), method=\"encode\")\n\n ---\n\n z = model.encode(jnp.ones((1, 784)))\n\n\n\nTransformations\n===============\n\nBoth Flax Linen and `Flax NNX transformations `__ provide their own set of transforms that wrap `JAX transforms `__ in a way that they can be used with ``Module`` objects.\n\nMost of the transforms in Flax Linen, such as ``grad`` or ``jit``, don't change much in Flax NNX. But, for example, if you try to do ``scan`` over layers, as described in the next section, the code differs by a lot.\n\nLet’s start with an example:\n\n* First, define an ``RNNCell`` ``Module`` that will contain the logic for a single step of the RNN.\n* Define a ``initial_state`` method that will be used to initialize the state (a.k.a. ``carry``) of the RNN. Like with ``jax.lax.scan`` (`API doc `__), the ``RNNCell.__call__`` method will be a function that takes the carry and input, and returns the new carry and output. In this case, the carry and the output are the same.\n\n.. codediff::\n :title: Linen, NNX\n :sync:\n\n class RNNCell(nn.Module):\n hidden_size: int\n\n\n @nn.compact\n def __call__(self, carry, x):\n x = jnp.concatenate([carry, x], axis=-1)\n x = nn.Dense(self.hidden_size)(x)\n x = jax.nn.relu(x)\n return x, x\n\n def initial_state(self, batch_size: int):\n return jnp.zeros((batch_size, self.hidden_size))\n\n ---\n\n class RNNCell(nnx.Module):\n def __init__(self, input_size, hidden_size, rngs):\n self.linear = nnx.Linear(hidden_size + input_size, hidden_size, rngs=rngs)\n self.hidden_size = hidden_size\n\n def __call__(self, carry, x):\n x = jnp.concatenate([carry, x], axis=-1)\n x = self.linear(x)\n x = jax.nn.relu(x)\n return x, x\n\n def initial_state(self, batch_size: int):\n return jnp.zeros((batch_size, self.hidden_size))\n\nNext, define an ``RNN`` ``Module`` that will contain the logic for the entire RNN.\n\nIn Flax Linen:\n\n* You will use ``flax.linen.scan`` (``nn.scan``) to define a new temporary type that wraps ``RNNCell``. During this process you will also: 1) instruct ``nn.scan`` to broadcast the ``params`` collection (all steps share the same parameters) and to not split the ``params`` PRNG stream (so that all steps initialize with the same parameters); and, finally, 2) specify that you want scan to run over the second axis of the input and stack outputs along the second axis as well.\n* You will then use this temporary type immediately to create an instance of the “lifted” ``RNNCell`` and use it to create the ``carry``, and the run the ``__call__`` method, which will ``scan`` over the sequence.\n\nIn Flax NNX:\n\n* You will create a ``scan`` function (``scan_fn``) that will use the ``RNNCell`` defined in ``__init__`` to scan over the sequence, and explicitly set ``in_axes=(nnx.Carry, None, 1)``. ``nnx.Carry`` means that the ``carry`` argument will be the carry, ``None`` means that ``cell`` will be broadcasted to all steps, and ``1`` means ``x`` will be scanned across axis `1`.\n\n.. codediff::\n :title: Linen, NNX\n :sync:\n\n class RNN(nn.Module):\n hidden_size: int\n\n @nn.compact\n def __call__(self, x):\n rnn = nn.scan(\n RNNCell, variable_broadcast='params',\n split_rngs={'params': False}, in_axes=1, out_axes=1\n )(self.hidden_size)\n carry = rnn.initial_state(x.shape[0])\n carry, y = rnn(carry, x)\n\n return y\n\n x = jnp.ones((3, 12, 32))\n model = RNN(64)\n variables = model.init(jax.random.key(0), x=jnp.ones((3, 12, 32)))\n y = model.apply(variables, x=jnp.ones((3, 12, 32)))\n\n ---\n\n class RNN(nnx.Module):\n def __init__(self, input_size: int, hidden_size: int, rngs: nnx.Rngs):\n self.hidden_size = hidden_size\n self.cell = RNNCell(input_size, self.hidden_size, rngs=rngs)\n\n def __call__(self, x):\n scan_fn = lambda carry, cell, x: cell(carry, x)\n carry = self.cell.initial_state(x.shape[0])\n carry, y = nnx.scan(\n scan_fn, in_axes=(nnx.Carry, None, 1), out_axes=(nnx.Carry, 1)\n )(carry, self.cell, x)\n\n return y\n\n x = jnp.ones((3, 12, 32))\n model = RNN(x.shape[2], 64, rngs=nnx.Rngs(0))\n\n y = model(x)\n\n\n\nScan over layers\n================\n\nIn general, transforms of Flax Linen and Flax NNX should look the same. However, `Flax NNX transforms `__ are designed to be closer to their lower-level `JAX counterparts `__, and thus we throw away some assumptions in certain Linen lifted transforms. This scan-over-layers use case will be a good example to showcase it.\n\nScan-over-layers is a technique where you run an input through a sequence of N repeated layers, passing the output of each layer as the input to the next layer. This pattern can significantly reduce compilation time for large models. In the example below, you will repeat the ``Block`` ``Module`` 5 times in the top-level ``MLP`` ``Module``.\n\n* In Flax Linen, you apply the ``flax.linen.scan`` (``nn.scan``) transforms upon the ``Block`` ``nn.Module`` to create a larger ``ScanBlock`` ``nn.Module`` that contains 5 ``Block`` ``nn.Module`` objects. It will automatically create a large parameter of shape ``(5, 64, 64)`` at initialization time, and iterate over at call time every ``(64, 64)`` slice for a total of 5 times, like a ``jax.lax.scan`` (`API doc `__) would.\n* Up close, in the logic of this model there actually is no need for the ``jax.lax.scan`` operation at initialization time. What happens there is more like a ``jax.vmap`` operation - you are given a ``Block`` sub-``Module`` that accepts ``(in_dim, out_dim)``, and you \"vmap\" it over ``num_layers`` of times to create a larger array.\n* In Flax NNX, you take advantage of the fact that model initialization and running code are completely decoupled, and instead use the :func:`nnx.vmap` transform to initialize the underlying ``Block`` parameters, and the :func:`nnx.scan` transform to run the model input through them.\n\nFor more information on Flax NNX transforms, check out the `Transforms guide `__.\n\n.. codediff::\n :title: Linen, NNX\n :sync:\n\n class Block(nn.Module):\n features: int\n training: bool\n\n @nn.compact\n def __call__(self, x, _):\n x = nn.Dense(self.features)(x)\n x = nn.Dropout(0.5)(x, deterministic=not self.training)\n x = jax.nn.relu(x)\n return x, None\n\n class MLP(nn.Module):\n features: int\n num_layers: int\n\n\n\n\n @nn.compact\n def __call__(self, x, training: bool):\n ScanBlock = nn.scan(\n Block, variable_axes={'params': 0}, split_rngs={'params': True},\n length=self.num_layers)\n\n y, _ = ScanBlock(self.features, training)(x, None)\n return y\n\n model = MLP(64, num_layers=5)\n\n ---\n\n class Block(nnx.Module):\n def __init__(self, input_dim, features, rngs):\n self.linear = nnx.Linear(input_dim, features, rngs=rngs)\n self.dropout = nnx.Dropout(0.5, rngs=rngs)\n\n def __call__(self, x: jax.Array): # No need to require a second input!\n x = self.linear(x)\n x = self.dropout(x)\n x = jax.nn.relu(x)\n return x # No need to return a second output!\n\n class MLP(nnx.Module):\n def __init__(self, features, num_layers, rngs):\n @nnx.split_rngs(splits=num_layers)\n @nnx.vmap(in_axes=(0,), out_axes=0)\n def create_block(rngs: nnx.Rngs):\n return Block(features, features, rngs=rngs)\n\n self.blocks = create_block(rngs)\n self.num_layers = num_layers\n\n def __call__(self, x):\n @nnx.scan(in_axes=(nnx.Carry, 0), out_axes=nnx.Carry)\n def forward(x, model):\n x = model(x)\n return x\n\n return forward(x, self.blocks)\n\n model = MLP(64, num_layers=5, rngs=nnx.Rngs(0))\n\n\nThere are a few other details to explain in the Flax NNX example above:\n\n* **The `@nnx.split_rngs` decorator:** Flax NNX transforms are completely agnostic of PRNG state, which makes them behave more like JAX transforms but diverge from the Flax Linen transforms that handle PRNG state. To regain this functionality, the ``nnx.split_rngs`` decorator allows you to split the ``nnx.Rngs`` before passing them to the decorated function and 'lower' them afterwards, so they can be used outside.\n\n * Here, you split the PRNG keys because ``jax.vmap`` and ``jax.lax.scan`` require a list of PRNG keys if each of its internal operations needs its own key. So for the 5 layers inside the ``MLP``, you split and provide 5 different PRNG keys from its arguments before going down to the JAX transform.\n\n * Note that actually ``create_block()`` knows it needs to create 5 layers *precisely because* it sees 5 PRNG keys, because ``in_axes=(0,)`` indicates that ``vmap`` will look into the first argument's first dimension to know the size it will map over.\n\n * Same goes for ``forward()``, which looks at the variables inside the first argument (aka. ``model``) to find out how many times it needs to scan. ``nnx.split_rngs`` here actually splits the PRNG state inside the ``model``. (If the ``Block`` ``Module`` doesn't have dropout, you don't need the :meth:`nnx.split_rngs` line as it would not consume any PRNG key anyway.)\n\n* **Why the Block Module in Flax NNX doesn't need to take and return that extra dummy value:** This is a requirement from ``jax.lax.scan`` `(API doc `__. Flax NNX simplifies this, so that you can now choose to ignore the second output if you set ``out_axes=nnx.Carry`` instead of the default ``(nnx.Carry, 0)``.\n\n * This is one of the rare cases where Flax NNX transforms diverge from the `JAX transforms `__ APIs.\n\nThere are more lines of code in the Flax NNX example above, but they express what happens at each time more precisely. Since Flax NNX transforms become way closer to the JAX transform APIs, it is recommended to have a good understanding of the underlying `JAX transforms `__ before using their `Flax NNX equivalents `__\n\nNow inspect the variable pytree on both sides:\n\n.. tab-set::\n\n .. tab-item:: Linen\n :sync: Linen\n\n .. code-block:: python\n\n # variables = model.init(key, x=jnp.ones((1, 64)), training=True)\n # variables['params']\n {\n ScanBlock_0: {\n Dense_0: {\n bias: (5, 64),\n kernel: (5, 64, 64),\n },\n },\n }\n\n .. tab-item:: NNX\n :sync: NNX\n\n .. code-block:: python\n\n # _, params, _ = nnx.split(model, nnx.Param, ...)\n # params\n {\n 'blocks': {\n 'linear': {\n 'bias': Param(value=(5, 64)),\n 'kernel': Param(value=(5, 64, 64))\n }\n }\n }\n\n\nUsing ``TrainState`` in Flax NNX\n================================\n\nFlax Linen has a convenient ``TrainState`` data class to bundle the model,\nparameters and optimizer. In Flax NNX, this is not really necessary. In this section,\nyou will learn how to construct your Flax NNX code around ``TrainState`` for any backward\ncompatibility needs.\n\nIn Flax NNX:\n\n* You must first call :meth:`nnx.split` on the model to get the\n separate :class:`nnx.GraphDef` and :class:`nnx.State`\n objects.\n* You can pass in :class:`nnx.Param` to filter all trainable parameters\n into a single :class:`nnx.State`, and pass in ``...`` for the remaining\n variables.\n* You also need to subclass ``TrainState`` to add a field for the other variables.\n* Then, you can pass in :meth:`nnx.GraphDef.apply` as the ``apply`` function,\n :class:`nnx.State` as the parameters and other variables, and an optimizer as arguments to the\n ``TrainState`` constructor.\n\nNote that :class:`nnx.GraphDef.apply` will take in :class:`nnx.State` objects as arguments and\nreturn a callable function. This function can be called on the inputs to output the\nmodel's logits, as well as the updated :class:`nnx.GraphDef` and :class:`nnx.State` objects.\nNotice below the use of ``@jax.jit`` since you aren't passing in Flax NNX Modules into\nthe ``train_step``.\n\n.. codediff::\n :title: Linen, NNX\n :sync:\n\n from flax.training import train_state\n\n sample_x = jnp.ones((1, 784))\n model = nn.Dense(features=10)\n params = model.init(jax.random.key(0), sample_x)['params']\n\n\n\n\n state = train_state.TrainState.create(\n apply_fn=model.apply,\n params=params,\n\n tx=optax.adam(1e-3)\n )\n\n @jax.jit\n def train_step(key, state, inputs, labels):\n def loss_fn(params):\n logits = state.apply_fn(\n {'params': params},\n inputs, # <== inputs\n rngs={'dropout': key}\n )\n return optax.softmax_cross_entropy_with_integer_labels(logits, labels).mean()\n\n grads = jax.grad(loss_fn)(state.params)\n\n\n state = state.apply_gradients(grads=grads)\n\n return state\n\n ---\n\n from flax.training import train_state\n\n model = nnx.Linear(784, 10, rngs=nnx.Rngs(0))\n model.train() # set deterministic=False\n graphdef, params, other_variables = nnx.split(model, nnx.Param, ...)\n\n class TrainState(train_state.TrainState):\n other_variables: nnx.State\n\n state = TrainState.create(\n apply_fn=graphdef.apply,\n params=params,\n other_variables=other_variables,\n tx=optax.adam(1e-3)\n )\n\n @jax.jit\n def train_step(state, inputs, labels):\n def loss_fn(params, other_variables):\n logits, (graphdef, new_state) = state.apply_fn(\n params,\n other_variables\n\n )(inputs) # <== inputs\n return optax.softmax_cross_entropy_with_integer_labels(logits, labels).mean()\n\n grads = jax.grad(loss_fn)(state.params, state.other_variables)\n\n\n state = state.apply_gradients(grads=grads)\n\n return state\n\n.. testcode:: Linen\n :hide:\n\n train_step(jax.random.key(0), state, sample_x, jnp.ones((1,), dtype=jnp.int32))\n\n.. testcode:: NNX\n :hide:\n\n sample_x = jnp.ones((1, 784))\n train_step(state, sample_x, jnp.ones((1,), dtype=jnp.int32))\n" }, { "path": "docs_nnx/migrating/nnx_010_to_nnx_011.rst", "content": "NNX 0.10 to NNX 0.11\n#########################################\n\nIn this guide we present the code changes required when we update Flax NNX code from Flax version\n``0.10.x`` to ``0.11.x``.\n\n\nUsing Rngs in NNX Transforms\n====================================\n\nNNX layers that use RNGs like Dropout or MultiHeadAttention now hold a ``fork``-ed copy of the ``Rngs``\nobject given at construction time instead of a shared reference to the original ``Rngs`` object. This has\ntwo consequences:\n* It changes the checkpoint structure, as each layer will have unique RNG state.\n* It changes how ``nnx.split_rngs`` interacts with transforms like ``nnx.vmap`` and ``nnx.scan``,\n as the resulting RNG state will now not be stored in scalar form.\n\nHere is how a \"scan over layers\" looks like in the new version:\n\n.. tab-set::\n\n .. tab-item:: v0.11\n :sync: v0.11\n\n .. code-block:: python\n\n import flax.nnx as nnx\n\n class MLP(nnx.Module):\n @nnx.split_rngs(splits=5)\n @nnx.vmap(in_axes=(0, 0))\n def __init__(self, rngs: nnx.Rngs):\n self.linear = nnx.Linear(3, 3, rngs=rngs)\n self.bn = nnx.BatchNorm(3, rngs=rngs)\n self.dropout = nnx.Dropout(0.5, rngs=rngs)\n self.node = nnx.Param(jnp.ones((2,)))\n\n\n @nnx.scan(in_axes=(0, nnx.Carry), out_axes=nnx.Carry)\n def __call__(self, x: jax.Array):\n return nnx.gelu(self.dropout(self.bn(self.linear(x))))\n\n .. tab-item:: v0.10\n :sync: v0.10\n\n .. code-block:: python\n :emphasize-lines: 12\n\n import flax.nnx as nnx\n\n class MLP(nnx.Module):\n @nnx.split_rngs(splits=5)\n @nnx.vmap(in_axes=(0, 0))\n def __init__(self, rngs: nnx.Rngs):\n self.linear = nnx.Linear(3, 3, rngs=rngs)\n self.bn = nnx.BatchNorm(3, rngs=rngs)\n self.dropout = nnx.Dropout(0.5, rngs=rngs)\n self.node = nnx.Param(jnp.ones((2,)))\n\n @nnx.split_rngs(splits=5)\n @nnx.scan(in_axes=(0, nnx.Carry), out_axes=nnx.Carry)\n def __call__(self, x: jax.Array):\n return nnx.gelu(self.dropout(self.bn(self.linear(x))))\n\n\nThe main thing to note is that the ``nnx.split_rngs`` over ``scan`` is not needed anymore, as the RNGs produced\nby ``__init__`` are no longer in scalar form (they keep the additional dimension) and thus can be used directly\nin ``scan`` without the need to split them again. Alternatively, can even remove the ``nnx.split_rngs`` decorator\nfrom the ``__init__`` method and use ``Rngs.fork`` directly before passing the RNGs to the module.\n\n.. code-block:: python\n\n class MLP(nnx.Module):\n @nnx.vmap(in_axes=(0, 0))\n def __init__(self, rngs: nnx.Rngs):\n self.linear = nnx.Linear(3, 3, rngs=rngs)\n self.bn = nnx.BatchNorm(3, rngs=rngs)\n self.dropout = nnx.Dropout(0.5, rngs=rngs)\n self.node = nnx.Param(jnp.ones((2,)))\n\n @nnx.scan(in_axes=(0, nnx.Carry), out_axes=nnx.Carry)\n def __call__(self, x: jax.Array):\n return nnx.gelu(self.dropout(self.bn(self.linear(x))))\n\n rngs = nnx.Rngs(0)\n mlp = MLP(rngs=rngs.fork(splits=5))\n\nLoading Checkpoints with RNGs\n==================================================\n\nWhen loading checkpoints in the new version, you need to drop the old RNGs structure and\npartially reinitialize the model with new RNGs. To do this, you can use ``nnx.jit`` to\n\n1. Remove the RNGs from the checkpoint.\n2. Perform partial initialization of the model with new RNGs.\n\n.. code-block:: python\n\n # load checkpoint\n checkpointer = ocp.StandardCheckpointer()\n checkpoint = checkpointer.restore(path / \"state\")\n\n @jax.jit\n def fix_checkpoint(checkpoint, rngs: nnx.Rngs):\n # drop rngs keys\n flat_paths = nnx.traversals.flatten_mapping(checkpoint)\n flat_paths = {\n path[:-1] if path[-1] == \"value\" else path: value # remove \"value\" suffix\n for path, value in flat_paths.items()\n if \"rngs\" not in path # remove rngs paths\n }\n checkpoint = nnx.traversals.unflatten_mapping(flat_paths)\n\n # initialize new model with given rngs\n model = MyModel(rngs=rngs)\n # overwrite model parameters with checkpoint\n nnx.update(model, checkpoint)\n # get full checkpoint with new rngs\n new_checkpoint = nnx.state(model)\n\n return new_checkpoint\n\n checkpoint = fix_checkpoint(checkpoint, rngs=nnx.Rngs(params=0, dropout=1))\n checkpointer.save(path.with_name(path.name + \"_new\"), checkpoint)\n\nThe previous code is efficient because ``jit`` performs dead code elimination (DCE) so it will not\nactually initialize the existing model parameters in memory.\n\nOptimizer Updates\n====================================\n\nOptimizer has been updated to not hold a reference to the model anymore. Instead, it now\ntakes the model and gradients as arguments in the ``update`` method. Concretely, these are the\nthe new changes:\n\n1. The ``wrt`` constructor argument is now required.\n2. The ``model`` attribute has been removed.\n3. The ``update`` method now takes ``(model, grads)`` instead of only ``(grads)``.\n\n.. tab-set::\n\n .. tab-item:: v0.11\n :sync: v0.11\n\n .. code-block:: python\n :emphasize-lines: 17, 26\n\n from flax import nnx\n import optax\n\n\n class Model(nnx.Module):\n def __init__(self, din, dmid, dout, rngs: nnx.Rngs):\n self.linear = nnx.Linear(din, dmid, rngs=rngs)\n self.bn = nnx.BatchNorm(dmid, rngs=rngs)\n self.dropout = nnx.Dropout(0.2, rngs=rngs)\n self.linear_out = nnx.Linear(dmid, dout, rngs=rngs)\n\n def __call__(self, x):\n x = nnx.relu(self.dropout(self.bn(self.linear(x))))\n return self.linear_out(x)\n\n model = Model(2, 64, 3, rngs=nnx.Rngs(0))\n optimizer = nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param)\n\n @nnx.jit\n def train_step(model, optimizer, x, y):\n def loss_fn(model):\n y_pred = model(x)\n return ((y_pred - y) ** 2).mean()\n\n loss, grads = nnx.value_and_grad(loss_fn)(model)\n optimizer.update(model, grads)\n\n return loss\n\n .. tab-item:: v0.10\n :sync: v0.10\n\n .. code-block:: python\n :emphasize-lines: 17, 26\n\n from flax import nnx\n import optax\n\n\n class Model(nnx.Module):\n def __init__(self, din, dmid, dout, rngs: nnx.Rngs):\n self.linear = nnx.Linear(din, dmid, rngs=rngs)\n self.bn = nnx.BatchNorm(dmid, rngs=rngs)\n self.dropout = nnx.Dropout(0.2, rngs=rngs)\n self.linear_out = nnx.Linear(dmid, dout, rngs=rngs)\n\n def __call__(self, x):\n x = nnx.relu(self.dropout(self.bn(self.linear(x))))\n return self.linear_out(x)\n\n model = Model(2, 64, 3, rngs=nnx.Rngs(0))\n optimizer = nnx.Optimizer(model, optax.adam(1e-3))\n\n @nnx.jit\n def train_step(model, optimizer, x, y):\n def loss_fn(model):\n y_pred = model(x)\n return ((y_pred - y) ** 2).mean()\n\n loss, grads = nnx.value_and_grad(loss_fn)(model)\n optimizer.update(grads)\n\n return loss\n\nPytrees containing NNX Objects\n====================================\n\nIn the new version, NNX modules are now Pytrees. This means that you can use them with JAX transforms\nlike ``jax.vmap`` and ``jax.jit`` directly (more documentation on this will be available soon). However,\nthis also means that code using ``jax.tree.*`` functions on structures that contain NNX modules will\nneed to take this into account to maintain the current behavior. In these cases, the solution is to\nuse the ``is_leaf`` argument to specify that NNX modules and other NNX objects should be treated as leaves.\n\n\n.. code-block:: python\n\n modules = [nnx.Linear(3, 3, rngs=nnx.Rngs(0)), nnx.BatchNorm(3, rngs=nnx.Rngs(1))]\n\n type_names = jax.tree.map(\n lambda x: type(x).__name__,\n modules,\n is_leaf=lambda x: isinstance(x, nnx.Pytree) # <-- specify that NNX objects are leaves\n )\n" }, { "path": "docs_nnx/mnist_tutorial.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"id\": \"0\",\n \"metadata\": {},\n \"source\": [\n \"[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google/flax/blob/main/docs_nnx/mnist_tutorial.ipynb)\\n\",\n \"[![Open On GitHub](https://img.shields.io/badge/Open-on%20GitHub-blue?logo=GitHub)](https://github.com/google/flax/blob/main/docs_nnx/mnist_tutorial.ipynb)\\n\",\n \"\\n\",\n \"# MNIST tutorial\\n\",\n \"\\n\",\n \"Welcome to Flax NNX! In this tutorial you will learn how to build and train a simple convolutional neural network (CNN) to classify handwritten digits on the MNIST dataset using the Flax NNX API.\\n\",\n \"\\n\",\n \"Flax NNX is a Python neural network library built upon [JAX](https://github.com/jax-ml/jax). If you have used the Flax Linen API before, check out [Why Flax NNX](https://flax.readthedocs.io/en/latest/why.html). You should have some knowledge of the main concepts of deep learning.\\n\",\n \"\\n\",\n \"Let’s get started!\\n\",\n \"\\n\",\n \"## 1. Install Flax\\n\",\n \"\\n\",\n \"If `flax` is not installed in your Python environment, use `pip` to install the package from PyPI (below, just uncomment the code in the cell if you are working from Google Colab/Jupyter Notebook):\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"2\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# !pip install -U \\\"jax[cuda12]\\\"\\n\",\n \"# !pip install -U flax\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"3\",\n \"metadata\": {},\n \"source\": [\n \"## 2. Load the MNIST dataset\\n\",\n \"\\n\",\n \"First, you need to load the MNIST dataset and then prepare the training and testing sets via Tensorflow Datasets (TFDS). You normalize image values, shuffle the data and divide it into batches, and prefetch samples to enhance performance.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 1,\n \"id\": \"4\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"import tensorflow_datasets as tfds # TFDS to download MNIST.\\n\",\n \"import tensorflow as tf # TensorFlow / `tf.data` operations.\\n\",\n \"\\n\",\n \"tf.random.set_seed(0) # Set the random seed for reproducibility.\\n\",\n \"\\n\",\n \"train_steps = 1200\\n\",\n \"eval_every = 200\\n\",\n \"batch_size = 32\\n\",\n \"\\n\",\n \"train_ds: tf.data.Dataset = tfds.load('mnist', split='train')\\n\",\n \"test_ds: tf.data.Dataset = tfds.load('mnist', split='test')\\n\",\n \"\\n\",\n \"train_ds = train_ds.map(\\n\",\n \" lambda sample: {\\n\",\n \" 'image': tf.cast(sample['image'], tf.float32) / 255,\\n\",\n \" 'label': sample['label'],\\n\",\n \" }\\n\",\n \") # normalize train set\\n\",\n \"test_ds = test_ds.map(\\n\",\n \" lambda sample: {\\n\",\n \" 'image': tf.cast(sample['image'], tf.float32) / 255,\\n\",\n \" 'label': sample['label'],\\n\",\n \" }\\n\",\n \") # Normalize the test set.\\n\",\n \"\\n\",\n \"# Create a shuffled dataset by allocating a buffer size of 1024 to randomly draw elements from.\\n\",\n \"train_ds = train_ds.repeat().shuffle(1024)\\n\",\n \"# Group into batches of `batch_size` and skip incomplete batches, prefetch the next sample to improve latency.\\n\",\n \"train_ds = train_ds.batch(batch_size, drop_remainder=True).take(train_steps).prefetch(1)\\n\",\n \"# Group into batches of `batch_size` and skip incomplete batches, prefetch the next sample to improve latency.\\n\",\n \"test_ds = test_ds.batch(batch_size, drop_remainder=True).prefetch(1)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"5\",\n \"metadata\": {},\n \"source\": [\n \"## 3. Define the model with Flax NNX\\n\",\n \"\\n\",\n \"Create a CNN for classification with Flax NNX by subclassing `nnx.Module`:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"6\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"from flax import nnx # The Flax NNX API.\\n\",\n \"from functools import partial\\n\",\n \"from typing import Optional\\n\",\n \"\\n\",\n \"class CNN(nnx.Module):\\n\",\n \" \\\"\\\"\\\"A simple CNN model.\\\"\\\"\\\"\\n\",\n \"\\n\",\n \" def __init__(self, *, rngs: nnx.Rngs):\\n\",\n \" self.conv1 = nnx.Conv(1, 32, kernel_size=(3, 3), rngs=rngs)\\n\",\n \" self.batch_norm1 = nnx.BatchNorm(32, rngs=rngs)\\n\",\n \" self.dropout1 = nnx.Dropout(rate=0.025)\\n\",\n \" self.conv2 = nnx.Conv(32, 64, kernel_size=(3, 3), rngs=rngs)\\n\",\n \" self.batch_norm2 = nnx.BatchNorm(64, rngs=rngs)\\n\",\n \" self.avg_pool = partial(nnx.avg_pool, window_shape=(2, 2), strides=(2, 2))\\n\",\n \" self.linear1 = nnx.Linear(3136, 256, rngs=rngs)\\n\",\n \" self.dropout2 = nnx.Dropout(rate=0.025)\\n\",\n \" self.linear2 = nnx.Linear(256, 10, rngs=rngs)\\n\",\n \"\\n\",\n \" def __call__(self, x, rngs: nnx.Rngs | None = None):\\n\",\n \" x = self.avg_pool(nnx.relu(self.batch_norm1(self.dropout1(self.conv1(x), rngs=rngs))))\\n\",\n \" x = self.avg_pool(nnx.relu(self.batch_norm2(self.conv2(x))))\\n\",\n \" x = x.reshape(x.shape[0], -1) # flatten\\n\",\n \" x = nnx.relu(self.dropout2(self.linear1(x), rngs=rngs))\\n\",\n \" x = self.linear2(x)\\n\",\n \" return x\\n\",\n \"\\n\",\n \"# Instantiate the model.\\n\",\n \"model = CNN(rngs=nnx.Rngs(0))\\n\",\n \"# Visualize it.\\n\",\n \"nnx.display(model)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"7\",\n \"metadata\": {},\n \"source\": [\n \"### Run the model\\n\",\n \"\\n\",\n \"Let's put the CNN model to the test! Here, you’ll perform a forward pass with arbitrary data and print the results.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 3,\n \"id\": \"8\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"Array([[ 0.11409501, 0.4546129 , -0.6421267 , -0.12122799, -0.22859162,\\n\",\n \" 0.13616608, 1.0126765 , -0.03625144, 0.6132787 , -0.06018351]], dtype=float32)\"\n ]\n },\n \"execution_count\": 3,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"import jax.numpy as jnp # JAX NumPy\\n\",\n \"\\n\",\n \"y = model(jnp.ones((1, 28, 28, 1)), nnx.Rngs(0))\\n\",\n \"y\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"9\",\n \"metadata\": {},\n \"source\": [\n \"## 4. Create the optimizer and define some metrics\\n\",\n \"\\n\",\n \"In Flax NNX, you need to create an `nnx.Optimizer` object to manage the model's parameters and apply gradients during training. `nnx.Optimizer` receives the model's reference, so that it can update its parameters, and an [Optax](https://optax.readthedocs.io/) optimizer to define the update rules. Additionally, you will define an `nnx.MultiMetric` object to keep track of the `Accuracy` and the `Average` loss.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 4,\n \"id\": \"12\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"import optax\\n\",\n \"\\n\",\n \"learning_rate = 0.005\\n\",\n \"momentum = 0.9\\n\",\n \"\\n\",\n \"optimizer = nnx.Optimizer(\\n\",\n \" model, optax.adamw(learning_rate, momentum), wrt=nnx.Param\\n\",\n \")\\n\",\n \"metrics = nnx.MultiMetric(\\n\",\n \" accuracy=nnx.metrics.Accuracy(),\\n\",\n \" loss=nnx.metrics.Average('loss'),\\n\",\n \")\\n\",\n \"\\n\",\n \"nnx.display(optimizer)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"13\",\n \"metadata\": {},\n \"source\": [\n \"## 5. Define training step functions\\n\",\n \"\\n\",\n \"In this section, you will define a loss function using the cross entropy loss ([`optax.softmax_cross_entropy_with_integer_labels()`](https://optax.readthedocs.io/en/latest/api/losses.html#optax.softmax_cross_entropy_with_integer_labels)) that the CNN model will optimize over.\\n\",\n \"\\n\",\n \"In addition to the `loss`, during training and testing you will also get the `logits`, which will be used to calculate the accuracy metric.\\n\",\n \"\\n\",\n \"During training — the `train_step` — you will use `nnx.value_and_grad` to compute the gradients and update the model's parameters using the `optimizer` you have already defined. The `train_step` also receives an `nnx.Rngs` object to provide randomness for dropout. The `eval_step` omits `rngs` because the eval view already has `deterministic=True`, so dropout is disabled and no random key is needed. During both steps, the `loss` and `logits` are used to update the metrics.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"14\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"def loss_fn(model: CNN, batch, rngs: nnx.Rngs | None = None):\\n\",\n \" logits = model(batch['image'], rngs)\\n\",\n \" loss = optax.softmax_cross_entropy_with_integer_labels(\\n\",\n \" logits=logits, labels=batch['label']\\n\",\n \" ).mean()\\n\",\n \" return loss, logits\\n\",\n \"\\n\",\n \"@nnx.jit\\n\",\n \"def train_step(model: CNN, optimizer: nnx.Optimizer, metrics: nnx.MultiMetric, rngs: nnx.Rngs, batch):\\n\",\n \" \\\"\\\"\\\"Train for a single step.\\\"\\\"\\\"\\n\",\n \" grad_fn = nnx.value_and_grad(loss_fn, has_aux=True)\\n\",\n \" (loss, logits), grads = grad_fn(model, batch, rngs)\\n\",\n \" metrics.update(loss=loss, logits=logits, labels=batch['label']) # In-place updates.\\n\",\n \" optimizer.update(model, grads) # In-place updates.\\n\",\n \"\\n\",\n \"@nnx.jit\\n\",\n \"def eval_step(model: CNN, metrics: nnx.MultiMetric, batch):\\n\",\n \" loss, logits = loss_fn(model, batch)\\n\",\n \" metrics.update(loss=loss, logits=logits, labels=batch['label']) # In-place updates.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"17\",\n \"metadata\": {},\n \"source\": [\n \"In the code above, the [`nnx.jit`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.jit) transformation decorator traces the `train_step` function for just-in-time compilation with [XLA](https://www.tensorflow.org/xla), optimizing performance on hardware accelerators, such as Google TPUs and GPUs. `nnx.jit` is a stateful version of the `jax.jit` transform that allows its function input and outputs to be Flax NNX objects. Similarly, `nnx.value_and_grad` is a stateful version of `jax.value_and_grad`. Check out [the transforms guide](https://flax.readthedocs.io/en/latest/guides/transforms.html) to learn more.\\n\",\n \"\\n\",\n \"> **Note:** The code shows how to perform several in-place updates to the model, the optimizer, the RNG streams, and the metrics, but _state updates_ were not explicitly returned. This is because Flax NNX transformations respect _reference semantics_ for Flax NNX objects, and will propagate the state updates of the objects passed as input arguments. This is a key feature of Flax NNX that allows for a more concise and readable code. You can learn more in [Why Flax NNX](https://flax.readthedocs.io/en/latest/why.html).\\n\",\n \"\\n\",\n \"\\n\",\n \"## 6. Train and evaluate the model\\n\",\n \"\\n\",\n \"Now, you can train the CNN model. Before the training loop, we use [`nnx.view`](https://flax.readthedocs.io/en/latest/guides/view.html) to create a `train_model` (with dropout enabled and batch norm in training mode) and an `eval_model` (with dropout disabled and batch norm using running statistics). These views share the same underlying weights, so updates during training are automatically reflected during evaluation.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"22\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"image/png\": \"iVBORw0KGgoAAAANSUhEUgAABL4AAAHDCAYAAAAqZtO0AAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjUsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvWftoOwAAAAlwSFlzAAAPYQAAD2EBqD+naQAApvxJREFUeJzs3Qd4VFX6BvB3Mukd0hMCoXcSepMmCIrS7QXL2sWy/NUVReyyuitrwYq9NzoqiDRBmhBC7530AOmkzcz/+c7NpEACCSS5U97f81wyc+dm5oRkkjvvfOc7BovFYgEREREREREREZGDcdF7AERERERERERERPWBwRcRERERERERETkkBl9EREREREREROSQGHwREREREREREZFDYvBFREREREREREQOicEXERERERERERE5JAZfRERERERERETkkBh8ERERERERERGRQ2LwRUREREREREREDonBFxEREREREREROSQGX0Skq88//xwGgwGbNm3SeyhEREREVOq9995T52i9e/fWeyhERJeEwRcRERERERFV8s033yAmJgYbN27EgQMH9B4OEdFFY/BFREREREREZQ4fPoy1a9dixowZCAkJUSGYLcrLy9N7CERkBxh8EZHN27JlC6666ir4+/vD19cXQ4cOxfr16ysdU1xcjBdeeAGtW7eGp6cngoKCcNlll2Hp0qVlx6SkpODOO+9EkyZN4OHhgYiICIwZMwZHjhzR4asiIiIisk0SdDVq1AhXX301rr322iqDr8zMTPzzn/9UVWFyXiXnVxMnTkRGRkbZMQUFBXj++efRpk0bdX4m517jx4/HwYMH1e0rV65U0ynlY0Vybib7pSWG1R133KHOA+VzR44cCT8/P9xyyy3qttWrV+O6665D06ZN1Viio6PV2M6cOXPOuPfs2YPrr79eBXpeXl5o27YtnnnmGXXbihUr1OPOnTv3nM/79ttv1W3r1q27pP9bImp4rjo8JhFRje3cuRMDBgxQodeTTz4JNzc3fPjhhxg8eDBWrVpV1ndCTqqmT5+Ou+++G7169UJ2drbqGxYfH48rrrhCHTNhwgR1fw8//LA6SUtLS1PB2LFjx9R1IiIiItKCLwmo3N3dcdNNN+H999/H33//jZ49e6rbc3Nz1fnZ7t27cdddd6Fbt24q8FqwYAFOnDiB4OBgmEwmXHPNNVi2bBluvPFGPProo8jJyVHnXjt27EDLli1rPa6SkhKMGDFCvbn53//+F97e3mr/Tz/9hPz8fDzwwAPqzU+ZnvnOO++oschtVtu2bVPjlvPJe++9V53/SZC2cOFCvPLKK+r8UkIz+frHjRt3zv+JjLlv376X/P9LRA3MQkSko88++8wiv4r+/vvvKm8fO3asxd3d3XLw4MGyfUlJSRY/Pz/LwIEDy/bFxsZarr766mof5/Tp0+px/vOf/9TxV0BERETkODZt2qTOmZYuXaqum81mS5MmTSyPPvpo2THTpk1Tx8yZM+ecz5fjxaeffqqOmTFjRrXHrFixQh0jHys6fPiw2i/niVa333672vfUU0+dc3/5+fnn7Js+fbrFYDBYjh49WrZPzh3lHLLivorjEVOmTLF4eHhYMjMzy/alpaVZXF1dLc8991wV/2NEZOs41ZGIbJa8U/j7779j7NixaNGiRdl+KZO/+eabsWbNGlXZJQIDA1U11/79+6u8Lylll3ctpZT+9OnTDfY1EBEREdkTqWwKCwvDkCFD1HWZ3nfDDTfg+++/V+dmYvbs2YiNjT2nKsp6vPUYqfySSvvqjrkYUtVV1Xlexb5fUn3Wr18/KfJQLTNEeno6/vzzT1WhJlMiqxuPTNcsLCzEzz//XLbvhx9+UNVmt95660WPm4j0w+CLiGyWnKBI2br0Xjhb+/btYTabcfz4cXX9xRdfVL0mpIdE586d8cQTT6hydivp9/Daa6/ht99+UydzAwcOxOuvv676fhERERGR9qajBFwSekmDe1nNUTZpLZGamqqmLQqZHtipU6fz3pccI+dwrq51111H7kt6iZ1N2lZID7DGjRurPmDSv2vQoEHqtqysLPXx0KFD6uOFxt2uXTs1pbNiXzO53KdPH7Rq1arOvhYiajgMvojIIUiQJSdYn376qTqh+fjjj1W/Cflo9dhjj2Hfvn2qF5g0WH322WdVgGZ9J5CIiIjImS1fvhzJyckq/JIFg6ybNIMXdb26Y3WVX9bKsrPJG5kuLi7nHCv9XH/55Rf861//wrx581QfMWtjfHmjtLak6kt6yUqPMDm/lEWVWO1FZL/Y3J6IbJa8WydNS/fu3Vvlijxy4iMNSK3kXT5ZtVE2aboqYZg0vZeG91bSlPT//u//1CbTIuPi4vDGG2/g66+/brCvi4iIiMgWSbAVGhqKd99995zb5syZo1Y7/OCDD9T5lDSoPx85ZsOGDWrlbWkmXxVZOVJI1X5FR48erfGYt2/frt7Y/OKLL1RgZVVxZW9hbZtxoXELacY/efJkfPfdd2plSBm/TPckIvvEii8isllGoxHDhw/H/Pnz1bLWVlJqL0tKy4o+stqjOHnyZKXPlTJ3KUeXHg1CpkzKktpnn5DJUtjWY4iIiIiclQQ8Em7JSozXXnvtOdukSZPUqoyycqOslL1161YVhJ1N+moJOUZ6bc2cObPaY5o1a6bO96T3VkXvvfdejcctn1/xPq2X33rrrXPeUJU3RWV2gEyNrGo8VtKb7KqrrlJvjEoYeOWVV6p9RGSfWPFFRDZBTkIWL158zn6p2JJ37CTkevDBB1Vvhw8//FCFVdKjy6pDhw5qCeru3buryq9NmzappqRykibkncChQ4eqUn05Vu5HTtYkRJN39YiIiIicmQRaEmyNHj26ytulx5WERxIEyRuQcp513XXXqWbxcv516tQpdR9SESaN76X66ssvv1SVUxs3bsSAAQNU4/k//vhDndONGTMGAQEB6j7eeecdNe1R3pRctGgR0tLSajxu6ckln/f4448jMTFRvSkqjfWrWszo7bffVueU0g7j3nvvRfPmzdWbqzJNMiEhodKxMn4J/MRLL71U6/9PIrIhei8rSUTOTZapll9F1W3Hjx+3xMfHW0aMGGHx9fW1eHt7W4YMGWJZu3Ztpft5+eWXLb169bIEBgZavLy8LO3atbO88sorlqKiInV7RkaG5aGHHlL7fXx8LAEBAZbevXtbfvzxR52+ciIiIiLbMWrUKIunp6clLy+v2mPuuOMOi5ubmzqvOnnypGXSpEmWqKgoi7u7u6VJkyaW22+/Xd1mlZ+fb3nmmWcszZs3V58XHh5uufbaay0HDx4sOyY9Pd0yYcIEdY7XqFEjy3333WfZsWOHOg+U80QruW85h6vKrl27LMOGDVPnisHBwZZ77rnHsnXr1nPuQ8h9jxs3Tp0zytfbtm1by7PPPnvOfRYWFqrxyDnjmTNnav3/SUS2wyD/6B2+EREREREREdmKkpISREZGYtSoUfjkk0/0Hg4RXQL2+CIiIiIiIiKqQFaHTE9Pr9Qwn4jsEyu+iIiIiIiIiAC1EuW2bdtUXy9paB8fH6/3kIjoErHii4iIiIiIiAjA+++/jwceeAChoaGqOT8R2T9WfBERERERERERkUNixRcRERERERERETkkBl9EREREREREROSQXGEHzGYzkpKS4OfnB4PBoPdwiIiIyA5IN4ecnBy1HL2LC9/rs1U8zyMiIqL6PM+zi+BLToaio6P1HgYRERHZoePHj6NJkyZ6D4OqwfM8IiIiqs/zPLsIvuQdQOsX5O/vr/dwiIiIyA5kZ2erQMV6HkG2ied5REREVJ/neXYRfFnL3uVkiCdEREREVBucPmfbeJ5HRERE9Xmex4YXRERERERERETkkBh8ERERERERERGRQ2LwRUREREREREREDskuenwRERHVB5PJhOLiYr2HQRfJzc0NRqNR72FQA+HzleoSf38QETkPBl9EROR0LBYLUlJSkJmZqfdQ6BIFBgYiPDycDewdGJ+vVF/4+4OIyDkw+CIiIqdjfREdGhoKb29vvuixQxKG5OfnIy0tTV2PiIjQe0hUT/h8pbrG3x9ERM6FwRcRETnddCnri+igoCC9h0OXwMvLS32UF6/y/eS0JcfD5yvVF/7+ICJyHmxuT0RETsXaI0gqR8j+Wb+P7P3kmPh8pfrE3x9ERM6BwRcRETklTpdyDPw+Ogd+n6k+8OeKiMg5MPgiIiIiIiIiIiKHdFHB17vvvouYmBh4enqid+/e2LhxY7XHfv755+rdlIqbfB4RERHpR/6Ov/nmm3VyXytXrlR/37nqHpHtP1+JiIicTa2b2//www+YPHkyPvjgAxV6yR/hESNGYO/evaoxZFX8/f3V7VYsKyYiIqq9wYMHIy4urk5eAP/999/w8fGpk3ER0bn4fCUiIrLTiq8ZM2bgnnvuwZ133okOHTqoAEwaQ3766afVfo4EXeHh4WVbWFjYpY6biIiIzmKxWFBSUlKjY0NCQtgwnEhHfL6WKyoq0nsIRETkwFxq+0dp8+bNGDZsWPkduLio6+vWrav283Jzc9GsWTNER0djzJgx2Llz53kfp7CwENnZ2ZW2+mI2W7D+0ElsPHyq3h6DiIjoUt1xxx1YtWoV3nrrrbLWAdZ2Ar/99hu6d+8ODw8PrFmzBgcPHlR/b+WNJl9fX/Ts2RN//PHHeadOyf18/PHHGDdunHqB3bp1ayxYsOCixzt79mx07NhRjUke64033qh0+3vvvaceQ9ofyDivvfbastt+/vlndO7cGV5eXggKClLnGXl5eRc9FqKGZsvPV5PJhH/84x9o3ry5eo61bdtWjfNs8qa29TkcERGBSZMmld0m05rvu+8+NWZ5Dnfq1AmLFi1Stz3//POq0q0iGbt8DRX/f8aOHYtXXnkFkZGRagziq6++Qo8ePeDn56feLL/55puRlpZW6b7kdcQ111yjZpTIcQMGDFD/h3/++Sfc3NyQkpJS6fjHHntMHUNERA0sLwPY9BnsLvjKyMhQfyzPrtiS62f/kbGSP2Tyh3P+/Pn4+uuvYTab0a9fP5w4caLax5k+fToCAgLKNgnM6ssX647gxo/WY8bS8qmYRETkfJUX+UUlDb7J49aUvDDt27evqrpOTk5Wm/Xv41NPPYV///vf2L17N7p06aLecBo5ciSWLVuGLVu24Morr8SoUaNw7Nix8z7GCy+8gOuvvx7btm1Tn3/LLbfg1KnavzEkb5LJ/dx4443Yvn27eiH87LPPqhf+YtOmTXjkkUfw4osvqlYIixcvxsCBA9Vt8nXddNNNuOuuu9TXI/3Dxo8fX6v/K3Jcej1XHen5KufiTZo0wU8//YRdu3Zh2rRpePrpp/Hjjz+WHfP+++/joYcewr333quewxKqtWrVquzzr7rqKvz111/q3F7uQ74eo9GI2pCvV57/S5cuLQvNiouL8dJLL2Hr1q2YN28ejhw5okIyq8TERPW7QsK45cuXq9818rtCKudkf4sWLVR4ZiX3980336hjiIiogZQUAn+9BbzdFVj0GHBkDeyux1dtyR992awk9Grfvj0+/PBD9YetKlOmTFF9xKyk4qu+wq/hHcPxwsJd2HD4FFKyChAewMb7RETO5kyxCR2mLWnwx9314gh4u9fsT7G8EeTu7q6qO6QSQuzZs0d9lADpiiuuKDu2cePGiI2NLbsuf2/nzp2rXrxWrNo4m7zAlNBJvPrqq3j77bfVAjbyQry2bRGGDh2qwi7Rpk0b9eL4P//5j3oMeUEv/YqkakMqNqQqvGvXrupYCQjkRayEXbJfSPUX1Q1ZoEi+D/KGpfyMvPPOO+jVq1eVx0poIG9GfvHFFypwkDczX3vttUo/D/KGqASbEoDIfUr1jnyPp06dWi89XfV6rjrS81WqoiQ0s5LKL5m5IcGXBGni5Zdfxv/93//h0UcfLTtOKtGEVKPJ40hwJ89tIYFTbcnvAKlak/8nq4oBldynfE3yuBIOSjWc/PzK/+3333+vvg5hHYOQSrbPPvsMTzzxhLq+cOFCFBQUlH1dRERUj+QNot0LgKXTgNNHtH3hXQBXL9hVxVdwcLB6Nyc1NbXSfrlu/aN+IfJHSk5uDxw4UO0x8i6OlC9X3OpLVKAXesY0Ut+jRduS6u1xiIiI6otMDapIXiQ+/vjj6o2mwMBA9YJRXqReqIJEqk8qviiVv79nTzOqCXms/v37V9on1/fv36+CEnnRL6GWvLC97bbbVEVGfn6+Ok4CAAnNJOy67rrrMGvWLJw+fbrWY6DqFyh67rnnEB8fr/6vZYGi6r7HEl7JG5USjklwef/996updVKVZCVBmFQHzZw5U33f5frrr7+uPods9/kqAZJMt5TeYfJ4H330UdnjyX0kJSWp52FVEhISVMVYxcDpYshzvGLoJaSCS6rdmjZtqkLxQYMGqf3Wscljy7RFa+hVVRgorzHWr1+vrkuVqYReXBiAiKieJSUAn18N/DhRC718w4Ex7wH3rgKadIddVXzJHyf5IymlyTIv31ruLNfP945URXLCKyXTUpJtK0bHRuLvI6exYGsS7h5Q+3esiIjIvnm5GVU1hx6PWxfOflEnL6Jl+tB///tfNT1J+vhID60LNZA++8WkVOzI3/m6Ji9oJXiRaYy///67mmolVUOycp288Jexr127Vt0mAcozzzyDDRs2qMoUungVFygSskDRL7/8olpSyPS7s8mUMfm/t56zPfDAA6raR/q1SYWXkO+T9Ke6+uqr1XXp4/Tdd9+piiBHeq5aH9sRnq9SLSWPKd9HmZUhz0epApTnmJDHP58L3S79f8+eFirVgxf6f5A+fhLEyiZhuIRyEnjJdev/xYUeW1aYl+BMqr7k94X0U5PfM0REVE+yk4HlLwEJ30rJF+DqCfR7BOj/KODhC1tR66mO8k7h7bffrt6tktJ4aVYpf6isJ1ETJ05EVFSUKo23lnP36dNH/SGXRpjyh/Xo0aO4++67YStGdo7A8wt3YduJLBzOyEPzYL4rRETkTOQFY02nMOlJ3oCSN5AuRHrvSOWDVOdYK0qkV05DkcoVGcPZY5IKEWsfIFdXV9W0XjapQJLAS3r2yBRH+X5IhZhsEopJdZhM/arYBoFqx7pAkbSTqOkCRbLYkDQur0iCB2nIXrGFhVQL7du3T31/pTeT3C4hmzM/V235+SqPJ9+3Bx98sGyfNIe3kiBMAkx5Y3vIkCFVVppJr17r9/xsEljJtFcJv6zTXaVS60JkKujJkydVvzBrixPpB3j2Y8vUWwnSqqv6ktcYMgVUqtJatmx5TvUpERHVgaJ8YN1MYM2bQHHpAkSdrwOGPgcE1l+P9otV6zOHG264Aenp6epEVP6oyaot0pTW2vBe3pmREykrmZ4g7y7KsY0aNVIVY/LuYIcOHWArgnw9cFmrYKzal44FCUl4dFhrvYdERER0DnkxKlUZ8qJYpidVV90hK7zNmTNHVT7IC0/ptVUflVvVkd5A0pdHehXJeYMEKzIVTlZyFNLI+tChQ6oZtZwb/Prrr2p80kNKvj55wT18+HBVvSHX5bxDwjS6eOdboMjae+psUmkjAZZ8nyRAkO+L/FxVDHOkUkx6sbZr106FmnKbrNQnjdarI4GabFb1uXq3nmz1+SqP9+WXX2LJkiWqKkoq+6TasmJFpVRgytRWeQ5KI/ucnBwVmD388MNq+qH8TEyYMEH9fMib2/IzJGOX/mKDBw9Wz1mZ8iqVa/I6QSqvLtS6RKY3SlgoVZ7y2Dt27DinH7DMMJHbZeEMCXGl35dMa5Q3460rQ8rPrTyW9CmTN+CJiKgOSUXv9p+BP54HsksXLGzSE7jy30CTylP5bUmtenxV/KMjVVty0iJ/0Hv37l12m5QTW1dtEv/73//KjpXwS0rqrQ1sbcmYuEj1cf7WRK4cRURENkmmJ0m4IG8eWacBVUVejEqgJFUd8mJaXgh269atwcYpjyWNsmVKVadOndSbZfIC1Lo6m1R3yQv9yy+/XAVaMuVOpsd17NhRvWD9888/1fQ6qSaRPlMyJUtefFPDkpUJJSSRUEsCCTn/kwr/im9wyvdZpqV9++23avqqVOPIlD35aAurd+vJVp+v9913n6qslFBazuGlyqpi9ZeQ2R0yq0PCanleykIU0qPPavbs2Srclsoq+fqefPLJskBUntPyedJHTPrIybRX+b+4EPk/ktcQstqk3KdUfsnPUkVBQUGqMlSq4iSAkzfUpQ9gxeov+fmU3zUyHpmJQkREdeT4RuDjYcCcu7XQKyAamPAJ8I+lNh16CYPFDlIeeSdQToyysrLqrdF9bmEJur+0FIUlZix6+DJ0igqol8chIiJ9yQpfhw8fVtUNZ0/jIsf6fjbE+YM9TXWUFQZ//vnnsj6t1oBDWlHMnz//vP/HEo7Iio1S4SUVezt37lS3SWgl+x566KGy46XSRnqAVVdJVlXFl9xPVd8nPl/pYsjqjlJ1Jitjng9/voiIaiDzmFbhtWO2dt3NBxgwGej7EOCm34qNtTnPu6iKL0fk6+GKYe218n9pck9ERETkKCouUGRlXaBIGpyfjwQC0r+1pKREVfpIM3srWY2zYgWYkCqn803Va8jVu8m5yIsf6TEnFYgyLZOIiC5BYQ6w7EVgZs/S0MsAdL0VeCQeGPi4rqFXbTH4qmBUrDbdceHWJJjNNl8IR0RE1CCk3470KKpqk9vIPsjiADItTKYh7t69W63SePYCRRWb30s7C5mSKv3YVq9erfo3SaAl09qsZGqe9PSSVhbSy0oWIZCpe9ZG7dTwnPn5KqGs9AeUr/OKK67QezhERPbJbALivwLe6Q6sfgMoKQBiBgD3rQLGvAv4hcPe2MeyOA1kcNsQ+Hm6IjmrAH8fOYXeLYL0HhIREZHupD9XdT16WK1jP2q7QJFMA5MeaxJ8SWgifdekEbr0aLOSRuPSjF16RKWlpanpkNJDSh6D9OHMz1fpNUxERJfg8GpgyRQgZbt2vVFzYPjLQLurZWll2Cv2+DrLkz9vxY+bTuDm3k3x6rjO9fpYRETU8NjTxbGwx5f9O9/3ic9Xqk/8+SIiKnXyILB0GrBnkXbdIwAY9ATQ617A1QO2iD2+LsHo2Cj18dftySgqabil34mIiIiIiIiIGsyZTGDJM8C7vbXQy2AEet6t9fHq97DNhl61xamOZ+nbMgjBvh7IyC3EmgPpuLydVv5PRERERERERGT3TCXA5s+AldOB/JPavlbDgOGvAKHt4GhY8XUWo4sB13SJUJcXJHB1RyIiIiIiIiJyEPv/AD7oD/z6uBZ6BbcFbpkN3DrbIUMvwYqvKoyJi8Tna4/g912pOFNkgpe7Ue8hERERERERERFdnPS92rTGA0u1616NgSFPA93vBIyOHQ059ld3keKiA9G0sTeOncrHH7tTMSo2Uu8hERERERERkT05dUj72LiF3iMhZ5Z3UpvSuOlTwGICXNyA3vcBAx8HvBrBGTD4qoLBYMDo2EjMXHEA8xOSGHwREREBOHLkiFr9bMuWLYiLi9N7OERERLaluAA4ugbYvxTY/3t58NWkJ9BtItBxPODhq/coyVmUFAEbPwJWvQ4UZmn72l0DXPEiENQSzoTBVzVGx2nB16p9acjKL0aAt5veQyIiIic3ePBgFTi9+eabdXJ/d9xxBzIzMzFv3rw6uT8iKsfnK5GTOH1UmzomYdehVUDJmfLbpLLGYgZO/K1ti6cAncYDXScCTXpIxYWeIydHZbEAe34Blj5bHr6GdQaufBVoPrBeH7rEZMaelBzEHzuNLccy1ccf7+uLMH9P6InBVzXahPmhXbif+qb9tiMZN/ZqqveQiIiIiIjsVlFREdzd3fUeBtGlV9EcX69VdEnYlb6n8u3+UdrqeK2HAy0GAUX5wNbvgPgvgVMHtY+yhbTXqsC63AD4BOn11ZCjSd4GLHkaOLJau+4TCgx9Foi7BXCp+97l6TmFlUKu7SeycKbYVOmYLcdO48pO2gKCeuGqjheo+hILtnJ1RyIi0pdUe6xatQpvvfWWmpIvm0w93LFjB6666ir4+voiLCwMt912GzIyMso+7+eff0bnzp3h5eWFoKAgDBs2DHl5eXj++efxxRdfYP78+WX3t3LlylqPS8bUq1cveHh4ICIiAk899RRKSkou+PhCHk8+18fHB4GBgejfvz+OHj1aR/9jRPqxlefrv/71L7Rp0wbe3t5o0aIFnn32WRQXF1c6ZuHChejZsyc8PT0RHByMcePGld1WWFio7iM6Olo9x1u1aoVPPvlE3fb555+r521FUo0mY7OScUvV28cff6ymSctjiMWLF+Oyyy5Tny9f5zXXXIODBw9Wuq8TJ07gpptuQuPGjdXviB49emDDhg3q/9HFxQWbNm2qdLxU1jVr1gxms7mG3yWiWshO1sKqH24FXm8BfDEKWPuOFnoZjECz/sCw54EH1gL/3AmMfhtofw3g4Qf4hQGXPQY8vBm441egy42AqyeQvhtYMgWY0Q746U7g4AqAP790sXJSgfmTgA8HaqGX0QMY8H/AI/FawFoHoVdRiRlbj2fi878O45HvtuCy15aj5yt/4L6vNuODVQex8fApFXr5ebpiYJsQPDq0Nb64qxf6twqG3ljxdR6jukTi9cV7se7QSaRmF+henkdERPVYEl6c3/CP6+Zd42kO8gJ637596NSpE1588UXt093cVHB0991343//+x/OnDmjXqRef/31WL58OZKTk9ULx9dff129mM3JycHq1athsVjw+OOPY/fu3cjOzsZnn32m7k9eYNZGYmIiRo4cqV7kf/nll9izZw/uuece9eJWXvCe7/ElHBs7dqw6/rvvvlOVIBs3bqz0opnIZp6rdvp89fPzUwFVZGQktm/frp5vsu/JJ59Ut//yyy/qsZ555hn1HJbn4a+//lr2+RMnTsS6devw9ttvIzY2FocPH64U1NXEgQMHMHv2bMyZMwdGo/bCS8K8yZMno0uXLsjNzcW0adPUOBISElSoJfsGDRqEqKgoLFiwAOHh4YiPj1ehVkxMjAoE5f9BwjAruS6/i+TziS6ZqQRI3FRa1fU7kLK98u1SRdP6Cm1rMQTwqhwCV0l+f8T017arXgN2/KyFaclbgZ1ztC2wqTYNMu5mICCq3r48ciDFZ4B17wJr/gcU5Wr7Ok3Qglj5eboEqdkFqlorXqq5jp7G9sQsFJaYz/mxbhPqh65NA9GtaSP1sWWIL1xcbOt8jsHXeUQ39kb3Zo2w+ehpLNqWjH9c1lzvIRERUX2QF9Kv6rCQydNJgLtPjQ4NCAhQU4SkckNeBIqXX34ZXbt2xauvvlp23KeffqqqM+RFt7x4lIBp/PjxqhJCSDWJlVSVSEWH9f5q67333lOPNXPmTBVYtWvXDklJSerFvLyQlRfy1T3+qVOnkJWVpSo9WrbUGqy2b9/+osZBTkSv56qdPl+nTp1adlkCIwnQvv/++7Lg65VXXsGNN96IF154oew4CbiEjOnHH3/E0qVLVdAkpGqstiRMk1AtJCSkbN+ECRMqHSP/D3L7rl27VFj47bffIj09HX///XdZwCfVZlYSHt5///2YMWOGqkSTUEyCPamII7pouenAgT+0oOvgcqAgs8KNBq0nl0xflLArPBa4lJBVgrKed2tbUgKw5Stg209A5jFgxcvAyle16ZJSqdPmSsDIftNUxRtBO2YDf7wAZB3T9kV1B0ZMB5r2rvXdFZaYsCspWwu5jp1GwrFMJGZW6FdXKsDLrSzkkq1LdAD8PW3/55PB1wWMiYtUwdeChEQGX0REZFO2bt2KFStWqGlTZ5NpQ8OHD8fQoUPVi+cRI0ao69deey0aNaqbpaulAqVv376VqrRkuqK8gJdpSvICurrHlxezUp0h+6+44gr1wloqX2S6JJEj0uP5+sMPP6hqLbl/a7Dm7+9fdrtUWEkVWFXkNqnQksqrSyEhXsXQS+zfv1+F4zJ1USrIrNMTjx07poIveWwJCaurapNq0Yceeghz585VwZ1UtQ0ZMkSFe0Q1Jj93SVvKq7rkMizlt3s1Ku/V1XJo/fXhiozTtiteAnYv0KrAjv5VPi6fEK0CTCrBgssDYHJiJzZr02SPbyjvKycVXp2urXEgm5x1BvFHtZBLqrp2JGWrqYwVSdFW23D/StVcLYJ97LI6n8HXBYzsHIEXFu7C1hNZOJKRh5jgmr3TR0REdkSmMEk1hx6PewnkheyoUaPw2muvnXObBEjyolWqNdauXYvff/8d77zzjprSJC82pd9OfbvQ48vUpEceeUT1+5EX6FKdIsf36dOn3sdGdkqv56r1se3o+SpTFG+55RZVzSVBmlShSbXXG2+8UamKrDrnu03IlEKZhlnR2f3DhPTnOpv8P0ggNmvWLDUNU4IvCbykOqwmjy3VdDINU36HSIWcVIjJ9FKiC8o/pVVzSVN6qe7KP2vqbkRsaVXXcK16ph6agVfL3RuIvVHbMg5oVWAJ3wJ5acBfb2mb9BKTKrD2o7XjyblkndAqvLb/WP536bJ/An0nnffnoaDYhJ1JWSro2nL8tPqYkl1wznGNfdzRNToQ3Zo1Uh+7RAfC18MxIiPH+CrqUbCvh2rG9ue+dNXk/pGhrfUeEhER1TV556qGU5j0JC/2TKbylXK6deumeudIlYOra9V/0uVdOanCkk0qLOTFplRJSH+ds++vtmRqojy+vPi1vvv3119/qR5CTZo0ueDjC6nqkG3KlCmqekxewDL4Int/rtrC81UCNPl8Cc+szl48QnpsLVu2DHfeeec5ny+VZxJISZN+61THiqSKS/qQSb8ua7gllVoXcvLkSezdu1eFXgMGDFD71qxZc864pCG+TImurupLpjtKWCZTrq1TRInOIeGs9OeyrsB4YiNgqVDV4uEPtByiBV1S3eV3cVP/65xUdl3xAnD5VGDfEq0K7MBSrRJMtl+fADpfp4VgUi1Gjq0wVws+ZUGFktLph7JK4+XPAv6VK+XlnEymKMqURWt/rl1JWSg2VX6jwuhiQLtwv7JKLvnYLMjbLqu5aoLBVw2Mjo1Uwdf8hEQ8fHkrh/1hICIi2yYvmK2rmsl0KZnqIy8epSG29OyRF4jSSFqqOuRFo6x6Ji9qZcpUaGio+lzpm2PtpSX3t2TJEvUiVFZWk4oQacBdUw8++KBaSe3hhx/GpEmT1P0899xz6kW6VIPI41X3+NIk+6OPPsLo0aNVxYd8rkx/kioOIkeg9/O1devWauqg3L+s2iiN7CVEq0ierzK9UvrsyZRBCZCkub306ZPHu/3223HXXXeVNbeX4CwtLU1NS+7du7fqYfb000+ryk0Zr0w5vBCZuinjl+e/VLrJGGU12Irk/0h6ocmUxunTp6vjtmzZon5XSEAu5P9FQnIZq4zxQlVi5EQKsoFDK8vDrtyUyreHdihtTD8ciO5t2/2zZGyyOqRsWYlaBdiWL7VeYJs+0bbwLloAJkFYTZrsk31Nx936HbDsxfKf46b9gCtfBSK7llVzbTuRVTZlUYKu9JzCc+4q2NcdXSuEXF2aBMDb3YniIIsdyMrKknhSfdRD9pkiS+tnfrU0+9ciy47ETF3GQEREdePMmTOWXbt2qY/2Zu/evZY+ffpYvLy81N/Fw4cPW/bt22cZN26cJTAwUO1v166d5bHHHrOYzWb1dY4YMcISEhJi8fDwsLRp08byzjvvlN1fWlqa5YorrrD4+vqq+1uxYsV5H18eT47bsmVL2b6VK1daevbsaXF3d7eEh4db/vWvf1mKi4vVbed7/JSUFMvYsWMtERER6nObNWtmmTZtmsVkMtXZ91Pv8weqmfN9n/h8vfjnq3jiiScsQUFB6nNuuOEGy//+9z9LQEBApWNmz55tiYuLU8/D4OBgy/jx48tuk//3f/7zn2XP01atWlk+/fTTstvnzp2r9snXcs0111g++ugjNTar5557zhIbG3vOuJYuXWpp3769+jq7dOmifo/I58n9WR05csQyYcIEi7+/v8Xb29vSo0cPy4YNGyrdzyeffKI+b+PGjZaLYc8/X1SB2WyxpO62WNa8ZbF8drXF8kJji+U5//Lt5XCL5dsbLZa/P7FYTh+z2D35O3lwhcXy050Wy4vB5V/nS6EWy+x7LZbDq7X/E7JvR/6yWD4YWP79/V9ni3nHXMvR9FzL3PgTlmnztluueXu1peWUX1ROUXGTfaPeWa2OmbflhOXYyTz1d8bR1OY8zyD/wMbJ0s3yrpas/lSxIWdDeuDrzfhtRwruG9QCU67iqlNERPaqoKBAVRtJzxxPT0+9h0P1+P20hfMHurDzfZ/4fKXzeemll/DTTz9h27ZtF/X5/PmyY0V5wOE/tYou2ayr2lkFtSpfgVH6Yrl6wGF7lm37QZsKmbarfH/jlkC324DYmwG/MD1HSLV16jCwdJq20AGAEjdfbIz+B740X4lNJ/KQkav1QqwoxM8D3awrLTZrhE6RAfByb8D+dDqpzXmeE9W2XfrqjhJ8LUxIwr9GtIOLLHFAREREREQNShYKkCmkM2fOxMsvv6z3cKihnDxYvtLhkb8AU4XpXK6eQMyA0rBrGNC4BZyCd2OgzwNA7/uBxM1aALZjNnDqIPDH88Cyl4C2V2lTIWVlSiNf/tsqy5lMZP3+b/ht/RhGczFMcMH3pssxo2ACTu4KAHBaHedmNKBjZECllRajAr3YjukC+JNfQ4PbhsLPwxVJWQXYdPQ0ejWvutEmERGRvZKeOrJVRZpQ//bbbw0+JiKqmjM/X6Wn4Hfffad6gEl/L3JQxQXA0TWlVV2/A6cOVb49sCnQeoQWdsVc5tyrHEro0aSHto14Fdg5VwvBpJn/nkXa5hehNUTveivQuP5Xdqbzyy0swbbjmdhyNB1+O7/DNac+Q2Nkq9v+NHXGKyW3YK+lKSICPDGyLORqhI6R/vB0c/xqrrrGqY618PhPW/Hz5hO4tU9TvDy2s27jICKii8epLdWTFdRkq4o0jo6KioKt4VRH+8epjs7zfLU1/PmyQdK03dqUXqYyFueX3+biBjTrV1rVNRwIbq0FPlS9tN1A/Fdag/QzFX5fNB+kVYG1uwZw489+fZPI5VBGHuKPnsaW45nq477UHPQzbMdU16/RzuW4Ou6QJRLfBd4LS6vh6BbTWFVzRQRw4Y7qcKpjPa7uKMHXL9uS8dyojnAzuug9JCIiojojq8zJRkS2j89XcgglRcDx9eVhV/qeyrf7RZavwNhiEODhB3sLPKTMRLc2OaHttRUAhz0H7P1VqwI7uAI4vErbvBoBXW7U+oGFddRnjA4ou6AYW6Wa61hm6WqLmcg6U1x2e0tDIma5fouhxi3qeoFbANK7/xNRQx7AMx4MIusDg69a6NcySC0DKg3l1hzIwJC2oXoPiYiIiIiIyH5kJwMHSqcvHlwJFOWU32YwAtG9y8MuCWNsqKpLgiyZonYqr0i9JpSPJ3MLcVJ9LMLJvMLSj9p+ud1kscDf0w0BXm4I9NY++stlL+2ydZPb/Msuu6uPPu7GuundJM39O47TttNHgYRvgC1fA9mJwIb3tS2qu1YF1mmC3QWMejKbLTiYnlsWcMnH/Wm5KvCsyMPVBf0iDXjI8DO6pc2Fi6UEcHEFet0Lz4FPIFr6tVG9YfBVC65GF1zTJRKfrz2CBQlJDL6IiOyYHcz0pxrg99E58PtM9YE/Vw3EVAIkbipvTJ+yvfLtPiFAKwm6rgBaDtGqkBrQmSITMkpDqrODq4qB1qncImTkFaGoxFzrx5BqH9mOVT07uVquLoayYKw8FKscmAWcFZZZj6m2D1SjZsCQp4FB/wIOLgfivwD2/qY1x5dt8dNAp3FAt9uBJj1tKni0BVn5xdhyvDzkSjieiZyCknOOi27spfXlig5Et2hfdDjxE1z/fA0oyNQOaHMVMPwlbcou1TsGX7U0KlYLvn7fmaJ+STrDMqFERI7Ezc1NfczPz1d9cMi+yfex4veVHAufr1Sf+PujHuWmAweXaUHXgWXlL/YVg1ZdpHp1XQFExMlcwDp76MISU2klVuXKK6nQKrtcYX9+kanWj+HtbkRjH3cE+Xog2Me9/LJv+eUg9dEdRhcDskuDL9ky8ytftt6WWeEYCVeKTGaUmC3a15BXVOsxuru6lAdhFUOyCsFZoHd7BMS+gaDYZxF1dAEC93wP19MHtGow2ULaaVVgMh3SJwjOxmS2YH9ajhZyHT2tgq6D6XnnHOflZkSXJrLSYiN0axqIuKaBCPXzlHQd2LcYmD8VOHlAOzi0ozb9tMXghv+CnBiDr1qSH+Qmjbxw4vQZLNuTqirAiIjIfhiNRgQGBiItLU1d9/b25hLQdlqpIS9a5fso30/5vpLj4fOV6gN/f9QDsxlI3lK+AmNivPxPl98uVVwth2phV6uhgE9wje+6xGTGqXwtyNICLK0qy1qhVWnKYW4RcgrPrb6pSUikAixfdwT5eKjAKuisAKt8v0etix9UCFLLn9GCYnOFgKyoPBSrJkSruElgI5Vp6TmFaquZDgBeQA/DXtzitgpXuayDp/RcW/I0SpZMw07/AdgZNganwvshwNvjnGmZqirN01XNkrJXp/OKVAWXBFyybT2epaa2ni0myLss5JKP7cL9zv26U3ao/zvVS81a2Xj5VKDrbYALf+c0NAZftSQnW9Lk/r2VB9V0RwZfRET2Jzw8XH20vpgm+yUvWq3fT3JMfL5SfeHvj0t05rQ2VU6FXUuB/IzKt0fElq/AKBVepS/2pSdSZmlQlXFOgFU61dDaLyuvSIU7tSVTBM+uvDo70JLbrRVavh6uNhWqy1gkXJMtPKD2oZmENZXCsIpVZtXtz9dCw02WdthU1A7TcBtGG9fieuNKxLocQmz2CrWd2BeMH0sG4z3TICTj3CowPw9XrYfZ2VMyK1aaeVWelinHy+c15CIAEqjuS9V6c6kpi8cy1cqLZ5M+a7HREnAFqqmLcdGB6uenWrlpwPKXgS1fARYzYHQH+jwIDPg/wJMrTOvFYLGDye22thz53pQcjHjzT7gbXfD31GHqCUtERPbHZDKhuLj2J9RkG2R60vkqNWzt/IEu7fvE5ys15O8PqoK8bJT+XNYVGE9s1F7YlzK7+yEz4jIkhlyGfX59kFgScG6frDwt6DLX8hWo5CGNvMsDLKnMCj4rwCq77CPVSLYVZNkLqRTLLShB5pnKFWaGlO2IPjIbrVN/hZdJW4zADAM2u3bDXMNQ/FoUh8yaFpad53tsrSKruodZxeulwZm3No1Tpp5e6PstP4sVV1nceiKzymmuLUJ80DW6Ebo1C1Qf24b7qemqF1RcoC0S8Ocb5Qs2dBgLXPEC0Cjmov9fCHVynseKr4sgP/xtw/ywNzUHS3ak4Pqe0XoPiYiILoK86OELHyL7wOcrUcOQuoi8IpMKCk6fPgnDoZXwPbYCYWmr4VuUXunYA2iK5aZYLCuJw+aC1ijJdgX2yi2ppVv1JMiwBlUSaJ3TJ6tChZZMp6tR+ECXRP6PVWWW91mFHWqW0wig+AyweyEQ/yVcjqxGz5LN6InNeNUnGKZ+NyK73U047R1TVlWWfcG+ZlrAJtM6JQyV2y62wq/qyjI3ZBeUqLDr6Emtp19FUuknFVzWKYtyuZGPe+0D4V3zgKXTgMxj2r7IrsCI6UCzvrX+Wqh+MPi6SKPjIvGfJXsxf2sigy8iIiIiIrJZsiiXdcXCsj5ZeVX0zMopQGD+YfS3xGOISwJ6uuyFm6G8Kibf4oG/zJ2wwhyHlaZYJKG8V5dMVTt/gFUeZEm44GbHvaCclpsX0OV6bTt5UJvOl/AtkJsK47qZaCRb075aQ/wOYwB3nxrdbUGxqdoFAM7ta2atRpPpnEUoNllqvAhAq1DfspBLpi3K9UsKVKWXnfTxOrZOu+4XCQx7Duh8fZ0u2ECXjlMdL9LxU/kY8PoKVZK5fspQhPrXbu41EREROd/5A52L3yeiuicv8VbsTcP7Kw9iZ1L2eVcu9EIB+rrsUkHXEGMCmhgq9+pKNEZhp09vHGnUH1mhvRDo51sWcAVLD63Syx6urMh0SqZibepr/JfA/iXl0189/IFOE7QQTCqg6mHqqfycnyk2nRuWVbgsAaussijVXHXWoig7CVj2IrD1O+26qxfQ/1Gg/yM1Dvvo0nGqYwOIbuyt0uL4Y5lYtC0Zd13WXO8hERERERGRE5MgYNnuNLy9fD+2nciqdJv0J1YVV77uaO+egX7mzehyZiOa5W6Bq7m8UsZi9IC52WUwth0BtBqGqKCWiNLhayE7YXQD2o3UNgmEpAJMKsFOHwE2f6ZtYZ21AKzLddoKn3VE+np5u7uqLSLAC/WuKA9Y+w7w11tAcenUyS43AkOnAQF8ltgyBl+XQFZ3lOBrwdYkBl9ERERERKRb4LV0V6oKvHYkZqt9Xm5GTOzbDNf1iEaYtwW+KRthUCsw/g4kH6x8B4FNy1ZgNMQMgNHdW58vhOybfyQw8HHgssnAkdVaALZrAZC6HfjtCeD3qdoUyG63Ac0us5/pgGYzsP1H4I8XgJwkbV90H+DKV7UVS8nmMfi6BFd3icSLi3Yh4Xgmjp7MQ7MgljUSEREREVHDMJst+F0Crz/24WhKGoIM2ejjnotr23nhyuau8C3ZCfzxN3B4VXmFinBxA5r1A1pfoQVewW3qZSoaOSkJtFoM0rarTgHbfwI2fwGk7dQCJNkaNdcCsLhbAL9w2Kxj64HFU4Ck+PKQeNgLQMdxfM7YEfb4ukS3fbIBq/dn4PHhbTDp8tZ6D4eIiIjs4PyByvH7RFQFeYlWmAPkpQP5J4G8DCA/Q7uedxKWvHRkpCUhKyMJ3iWZCEIOPAwXWA1PGm9bgy4JJDz8GuqrIdJ+piU8kl5g22cDRTnafoMRaDNCmwrZ6grAaCO1OaePAn88B+ycq1139wMGTAb6PAi4sb+3LWCPrwae7ijB1/yEJDw0pJWaZ0xERERERFQ5yMrWAqyyECuj2mBLXTZVv0KdvOIIKd3UFSs3b8A7GPAp3eRycGst7ArryAoV0o/87Mm0QNlGvArsnKeFYMfXA3t/1TbfcKDrLUDXW4HGLfQZZ0E2sGYGsO49wFQIGFyArrcBl08FfEP1GRNdMgZfl2hEp3A8M28H9qflYk9KDtpH8J1KIiIiIiKHD7IKskpDq/Sqg6uzQy7zBSqyquLmA/gEweIdjFSTH7aecsXhM144afFHvmsgenRogyt6dIBvUIQWcrE3F9kDWflQBVy3AOl7tQBMVkjMTQFWv6FtMQOAbrcD7Uc1TIWV2aT1JFv+svY8Fs0HaSFdeKf6f3yqVwy+LpG/pxsubxuKxTtTVNUXgy8iIiIiInsMsjK10EpVYZ1dmWW9bL395MUFWe6+gHdQaUVWSGl1VlDpx5DSKq3S272DYXL1wqJtSXhn+QEcSMtVd+Hv6aoW1prUvzkCvNzq/v+CqCGFtAVGvAIMfU6r+pIQ7OByrTm+bJ6BQJcbtKmQ9RVAHVoFLHkaSN2hXW/cEhj+MtD2KlZJOggGX3VgdFykCr4Wbk3CkyPawsWFTw4iIiIiIl1XYZMg65yKrLODrYpBVkntH0f6/pQFVxWmF5Z9DKl8u5tXje62xGTGwtLA61B6ntonIdc/LmuOO/rHqDffiRyKqzvQcay2ZR4HEr4BtnwNZB0HNn6obZHdtACs0wTAsw4KTjIOAEuf1QI34RkADHoK6Hm3Nh5yGAy+6sDl7ULh6+GKxMwziD92Gj1iGus9JCIiIiIixwuyKk0prBBcVVWRZTFdZJBVMbgKqlCZVUWwVcdTsCTwklkkM1ccwOEMLfAK9HbD3Zc1x+39YuDHwIucQWA0MPgpYOATwKEVWhXYnl+15viySXWWrKooIVh079pXZZ05Dax6Hdj4kRZ4S4N9CbvkMb35Wt4RMfiqA55uRgzvGIY58YnqDxWDLyIiIiKiCwRZ8uKzUpP30uDqnGAr4+KDLA//s8Kq0iCrqmBLbtNptTYJvOZuScS7Kw7gyMl8ta+RBF4DWqjAS95kJ3I6Lkag1TBty00Htn0PxH8FZOzVKsJkC26jBWBdbgR81XIP1TMVA5s+BVZO137/iNYjtGmNIW0a5EsifRgsFpnQbtvsYZnrVfvScfunGxHk4471Tw+Fm9FF7yERERE5NXs4fyB+nxyarI627QetefXZqxeqIMtc+/v0CDirCquq/lgVLrt6wJYVS+AVn6gqvI6d0gKvxj7uuGdAC0zs2ww+DLyIKpP44vhGrQps5xygWHvewMUNaDcS6DoRaDlEC80qfs7+pcDvzwAZ+7R9Ie213mKthurzdVCDnj/wN2kd6d8ySIVeJ/OK8NeBDAxuy6VOiYiIiMgJZSUCGz4ANn8OFGaf/1jpqVNVY/eqgi0VZDlG352iEjNmx59QFV4nTp9R++S1xL0DW+DWPgy8iKol0xqb9ta2K6cDO2ZrIZhMgdw1X9v8mwBdb9VWjSzM1QIvaZgv5PfIkGe0FSONfJ45C36n64ir0QVXd4nAl+uOYsHWJAZfRERERORckrcB62ZqL0StjeJlGlK7awDf0PLwyjrN0IGCrNoEXj9tPo73VhxU/YFFsK877hvYErf0aQpvd748I6oxaXDf405tS9kBbPkK2Po9kH0CWPVvYNVrWlAm1aVGd6D3/cDAx7XAnZwKf7PWoTFxkSr4WrIjBQXjTKr3FxERERGRw5IpRAeWAWvfBg6vKt8fMwDoOwloPRxwYQuQwhITftx0Au+vOICkrAK1L8TPA/cPaombezWFlztfNxBdkvBOwFWvAcNeAPYsAuK/AA7/qf2Oaj8auOIFoHELvUdJOmHwVYe6NW2EqEAv9e7N8j1pGNk5Qu8hERERERHVvZJCYPtPwNqZQPpubZ+sjCYrrfWbBER21XuENqGgWAKv43h/5UEklwZeoX4eeGBwS9zUqynfKCeqa7JARedrtS3zmNbQPqil3qMinTH4qkMGgwGj4yLVH7b5CYkMvoiIiIjIseSf0lZF2/gRkJuq7XP3BbrfAfS+DwhsqvcIbSbw+n7jMby/6iBSswvVvnB/TxV43dAzmoEXUUPg7yMqxeCrjo2O1YKvFXvTkXWmGAFebnoPiYiIiIjo0pw6BKx/H9jydfkqan6RQJ/7tSbRXoF6j9BmAq9vNxzDB6sOIi1HC7wiAjzx4OCWuK4HAy8iIj0w+Kpj7cL90CbMF/tSc7FkZwqu7xGt95CIiIiIiC7O8b+1/l3SM0caRIvwzkDfh7VpjU7WnL46Z4pM+GbDUXyw6hAycrXAS1qgSIXXdT2awMOVgRcRkV4YfNXHdMfYSPz3931YuDWJwRcRERER2RezCdj7K7D2HeD4hvL9rYYB/R4Gmg/SVkoj5BeV4Ov1R/HRnxJ4FZUFXg8NaYVruzeBuysb+xMR6Y3BVz0YHRulgq+/DmQgLacAoX6eeg+JiIiIiOj8ivKBhG+A9e9pUxuF0R3ofD3Q9yEgrIPeI7QZeYUl+Gr9Ucz68xBO5mmBV3RjL0wa0grjujLwIiKyJQy+6kHTIG/ERQci4Xgmft2WjDv6N9d7SEREREREVctN05rV//0JcOaUts8zEOj5D6DXvYBfuN4jtBm5hSX4ct0RfLz6ME6VBl7NgrxVhde4rlFwMzLwIiKyNQy+6smYuEgVfM3fmsTgi4iIiIhsT9oeYN1MYNuPgEnrS4XAZkDfSUDXWwB3H71HaDNyCorx5bqjmLX6EDLzi9W+mCBvTLq8NcbGRcKVgRcRkc1i8FVPru4SgZcW7cKWY5k4djJfVYEREREREenKYgGOrAbWzgT2Lynf36Sn1r+r3TWACxuxW2UXFOPzv47gkzWH1YrtokWwDyZd3kr19WXgRURk+xh81RPp69WvZTDWHMjAwm1JqvyZiIiIiEgXpmJg5zxg3TtA8tbSnQag3dVAv0eApr11HqBtkZDrs78O49M1h5FdUKL2tQjxwSOXt8ao2EgYXdjcn4jIXjD4qkfyLpAEXwsSGHwRERERkQ4KsoH4L4H17wPZJ7R9rl7aVMY+DwJBLfUeoU3Jyi/GJ38dVqFXTmng1SrUFw9f3grXdGHgRURkjxh81aMRncIxdd4O7E3NwZ6UbLQL99d7SERERETkDLJOABs+ADZ/ARRma/t8QoBe9wE97gJ8gvQeoU3JzC9S0xllWmNOoRZ4tQmTwKs1RnaOYOBFRGTHGHzVowAvNwxuG4Lfd6Wqqq92VzL4IiIiIqJ6JNMYpX/XzjmAWQtwENwW6DcJ6Hw94Oap9whtyum8Iny85hC+WHtUrdgo2ob54ZGhrXFVp3C4MPAiIrJ7DL7q2Zi4KBV8zU9IwhMj2sJg4B9PIiIiIqrjhvUH/gDWvg0c/rN8f8wArWF9qysAFzZhr+hUXpFaofHLtUeQV2RS+9pH+OPRoa0wvAMDLyIiR8Lgq54NbR8KH3cjEjPPIP7YaXRv1ljvIRERERGRIygpBLb9CKybCaTv0fYZjECn8UDfSUBknN4jtDkncwvx0epD+GrdUeSXBl4dJPAa1hpXtA9j4EVE5ID41k8983QzYkTHcHVZpjsSERER6eXdd99FTEwMPD090bt3b2zcuLHaY4uLi/Hiiy+iZcuW6vjY2FgsXrz4nOMSExNx6623IigoCF5eXujcuTM2bdpUz1+Jk8s/Bfz5H+B/nYAFk7TQy91PC7se3QpM+Jih11nScwrx6q+7cdlrK/DhqkMq9OoU5Y9ZE3vgl0cuU+frDL2IiBwTK74awKi4SMzZkohftifj2Ws6wNXIvJGIiIga1g8//IDJkyfjgw8+UKHXm2++iREjRmDv3r0IDQ095/ipU6fi66+/xqxZs9CuXTssWbIE48aNw9q1a9G1a1d1zOnTp9G/f38MGTIEv/32G0JCQrB//340atRIh6/QCZw6BKx7D0j4BijO1/b5RwG97we63w54Bug9QpuTllOggq5vNhxFQbFZ7evSJACPDm2Ny9uFsg0JEZETMFgs0hTAtmVnZyMgIABZWVnw97e/BvHFJjN6v7pM9RL48q5eGNgmRO8hEREROTx7P3+oaxJ29ezZEzNnzlTXzWYzoqOj8fDDD+Opp5465/jIyEg888wzeOihh8r2TZgwQVV1SSAm5PP++usvrF69+qLHxe9TDRzfqPXv2r1IGnpp+8I7A/0eATqOA4xueo/Q5qRlF+D9VQfx7YZjKCzRAq/Y6EA8NrS1WnyKgRcRkX2rzfkDK74agJvRBSM7h+Pr9cdUk3sGX0RERNSQioqKsHnzZkyZMqVsn4uLC4YNG4Z169ZV+TmFhYVqimNFEnqtWbOm7PqCBQtU1dh1112HVatWISoqCg8++CDuueeeasci9ytbxRNXqoLZBOz5BVj7DnCiwpRUaVQvDeubDwQY3pwjJasAH0jgtfEYikoDr65NA1WF16A2DLyIiJyRS333h6jo+++/V39sxo4dC2dc3VEs2ZmCgmKtkSYRERFRQ8jIyIDJZEJYWFil/XI9JSWlys+RQGvGjBlq6qJUhy1duhRz5sxBcnJy2TGHDh3C+++/j9atW6upkA888AAeeeQRfPHFF9WOZfr06eodWusmVWdUQVEesHEW8E534MfbtNDL6A50vRV4cD1w689Ai0EMvc6SnHUG0+bvwMD/rMDna4+o0KtHs0b46h+9MOeBfhjcltMaiYiclWt994ewOnLkCB5//HEMGDAAzqh700aIDPBEUlYBVuxJw1WdI/QeEhEREVG13nrrLVW5Jf29JDCQJvd33nknPv3007JjJBDr0aMHXn31VXVden/t2LFDnSfefvvtVd6vVJ3JuWTFii+GXwByUoGNHwGbPgHOnNb2eQYCPe8Get0L+FUOLUkjK6e/v/IAfvz7BIpMWoVXr5jGapXGfi2DGHYREVHtK77knT85CZITnw4dOqgTG29v70onQWeTdxhvueUWvPDCC2jRogWckawSI03uxYKtXN2RiIiIGk5wcDCMRiNSU1Mr7Zfr4eHa6tNnk0b18+bNQ15eHo4ePYo9e/bA19e30rlcRESEOh+sqH379jh27Fi1Y/Hw8FC9OCpuTi1tNzD/IeDNTsDq/2qhV6MYYOR/gcm7gKHPMvSqwonT+Xh67nYM/s8K1U5EQq/ezRvj23t644f7+qB/q2CGXkREVPuKr4vpDyFkKWypBvvHP/5Ro+anjtr7YUxslFpVZtmeNGQXFMPfk41IiYiIqP65u7uje/fuWLZsWVnLCanWkuuTJk067+dKawvp3VVcXIzZs2fj+uuvL7tNVnSUqv+K9u3bh2bNmtXTV+IgZG2pw39q/bsOLC3f36SX1r+r3dWAi1HPEdqs46fy8d7KA/h58wkUm7RG/31bBKkKrz4tgvQeHhER2Xvwdb7+EPIuYFWkAeonn3yChISEGj+O9H6Q6jBH0z7CD61CfXEgLRe/70zFtd2b6D0kIiIichIyvVCmH8rUxF69eql2FVLNJVX8YuLEiSrgkvMwsWHDBiQmJiIuLk59fP7551VY9uSTT5bd5z//+U/069dPTXWUQEz6vn700UdqoyqYioGdc7UVGlO2l+40AO2vAfo+DDTtrfMAbdexk/mYuWI/5sQnosSsBV79WwXh0aFt0Kt5Y72HR0RENqxeV3XMycnBbbfdhlmzZqkS+5py1N4PUm49JjYSbyzdh/kJiQy+iIiIqMHccMMNSE9Px7Rp01RDewm0Fi9eXPaGpkxPlEp+q4KCAkydOlU1sJcpjiNHjsRXX32FwMDAsmN69uyJuXPnqnM3qfBv3ry5CtSkxQVVUJAFbP4C2PABkJ2o7XP10hrW93kACGqp9wht1pGMPMxccQBztyTCVBp4DWgdrFZp7BHDwIuIiC7MYLFIrXXNpzpKP6+ff/650sqM8u5hZmYm5s+fX+l4qfKSJqfSU8JK3ikUcmIlpfHSKPVCJPiSVX+ysrLsvg+E/PEe/N+VMLoYsH7KUIT4eeg9JCIiIofkSOcPjsyhv0+Zx7WwS0Kvohxtn08o0PteoMc/AG8GN9U5nJGHd5bvx/yEpLLAa2CbEBV4dW/WSO/hERGRHZ0/uNZnfwhZBWj7dmsZt0beOZRKMFkpyBGquGorJtgHsdGB2Ho8E79uT8bt/WL0HhIRERER1aWkBK1/l0xrtJi0fcFtgX6TgM7XA26eeo/QZh1Mz8XM5QfU7IjSvAuD22qBV9emDLyIiKgBpjrWpj+ENEPt1KlTpc+3lsefvd+ZjI6NVMGXrO7I4IuIiIjIAcisBmlUL4HXkQqLOTUfqPXvajVMpjzoOUKbdiAtB+8sP4CFW5PKAq/L24XikaGtERddPr2WiIio3oOv2vaHoHON6hKBl3/Zhc1HT6uVaaIbe+s9JCIiIiK6GMUFwLYfgHXvAhmlK1wajECn8UDfSUBknN4jtGn7U3Pw9vIDWLQtSS12KYa1D1MVXp2bBOg9PCIicrYeX3pxxN4PN89aj7UHT+LJK9viwcGt9B4OERGRw3HE8wdHZLffp/xTwN+fABs/BPLStX3ufkD324He9wOBztfSozb2puTg7WX78euO5LLAa3iHMFXh1SmKgRcREenU44vqzpi4SBV8LUhIYvBFREREZC9OHgTWvwds+QYoOaPt828C9Lkf6DYR8GRocz67k7NV4PXbjpSyfVd2DMfDQ1uhYyT/74iIqO4x+NLJlR0jMHXeDuxJyVHveLUN99N7SERERERUnWMbgLVvA3t+AVBaohTeBej3CNBxLGB003uENm1nUpYKvJbsTC3bN7JzOB6+vDXaR9hRpR8REdkdBl86CfB2w+C2oVi6KxULtibiifB2eg+JiIiIiCoym4A9i7SG9Sf+Lt/fejjQ72EgZgBgMOg5Qruo8JqxdJ865xXy3zWycwQeubw13/glIqIGweBL59UdteArCY8PbwsDT5yIiIiI9FeUp01lXP8ucPqIts/oDnS5QWtYH8o3LGsiOesMJry/FvlFJhV4XdMlEo9c3gqtwxh4ERFRw2HwpSNZscbb3Yjjp85gy/FMdGvaSO8hERERETmvnBRg40da0/qCTG2fVyOg591Az3sAP20Vc6qZOfGJKvRqF+6HmTd3RatQBl5ERNTwGHzpyMvdqFavmZeQpJrcM/giIiIi0kHabmDtTGD7j4CpSNvXqDnQ9yEg7mbA3UfvEdodWTh+TvwJdfmu/s0ZehERkW4YfOlsTFyUCr4WbUvG1Kvbw9XooveQiIiIiByfxQIcXqX17zrwR/n+6N5a/662IwEXo54jtGvbTmThYHoePN1ccFXncL2HQ0RETozBl84uax2MRt5uyMgtxLpDJzGgdYjeQyIiIiJyXKZiYMccYN07QMr20p0GoP0oLfCK7qXzAB3D7NJqrxEdw+HnyRUviYhIPwy+dOZmdFEr23yz4Zia7sjgi4iIiKgeFGQBmz8HNnwIZCdq+9y8ga63An0eABq30HuEDqOoxKwWbxLjuzXRezhEROTkGHzZyOqOEnwt3pGCl8Z2gqcby+qJiIiI6kTmMWD9B0D8l0BRjrbPJxTofS/Q4x+Ad2O9R+hwVuxNQ2Z+MUL9PHBZq2C9h0NERE6OwZcN6BnTGBEBnkjOKsDKvem4shP7IBARERFdsl3zgZ/uBCwm7XpIO6DvJKDL9YCrh96jc1jWpvbjukbB6GLQezhEROTk2EndBri4GDAqNlJdXrC1tPSeiIiIiC5Ns8sAozvQfBBwy8/Ag+uBbrcx9KpHp/OKsHxPmrrMaY5ERGQLWPFlQ9MdP/rzEJbtTkNOQTGbgBIRERFdKp8g4JEtgH+E3iNxGgu3JaHYZEHHSH+0DffTezhERESs+LIVcnLQIsQHhSVm/L4zVe/hEBERETkGhl4Nana8NnuB1V5ERGQrGHzZCIPBgDGxUeqydRUcIiIiIiJ7cSAtF1uPZ6q+XmPitDYeREREemPwZUNGl54grDmQgZO5hXoPh4iIiIio1k3tB7cJQbAv+6gREZFtYPBlQ5oH+6BLkwCYzBb8uj1Z7+EQEREREdWI2WzB3C2c5khERLaHwZcNNrkX8xM43ZGIiIiI7MP6QyeRnFUAP09XDG0fqvdwiIiIyjD4sjGjYiNhMACbjp7GidP5eg+HiIiIiKjGTe2v6RIJTzej3sMhIiIqw+DLxoT5e6JP8yB1eeFWTnckIiIiItuWX1SC33Zo563XdtcWayIiIrIVDL5suMk9V3ckIiIiIlu3eEcK8otMaBbkjW5NG+k9HCIiokoYfNmgqzqFw81owO7kbOxPzdF7OERERERE1ZpTOs1xfNcmMEjPDiIiIhvC4MsGBXq7Y1CbEHWZVV9EREREZKuSs87gr4MZ6vL4bpzmSEREtofBl40aHRdVtrqjxWLRezhEREREROeYt0XOVYFezRsjurG33sMhIiI6B4MvGzWsfSi83Iw4diofW09k6T0cIiIiIqJK5M3Z2fEn1OUJrPYiIiIbxeDLRnm7u2J4xzB1eX6C1jeBiIiIiMhWbE/MwoG0XHi4uuCqzhF6D4eIiKhKDL5s2OhYbXXHRduSYTJzuiMRERER2V5T++Edw+Hv6ab3cIiIiKrE4MuGDWgdgkBvN6TnFGL9oZN6D4eIiIiISCkqMZctwsRpjkREZMsYfNkwdykb76SVjXO6IxERERHZipV703Aqrwghfh64rFWw3sMhIiKqFoMvGzcmTpvu+NuOFBSWmPQeDhERERFR2TTHsXGRcDXyJQUREdku/pWycb1iGiPc3xM5BSVYuTdd7+EQERERkZPLzC/Csj2p6vL4bk30Hg4REdF5MfiycS4uBoyK1aY7WvsoEBERERHpZeG2ZBSbLOgQ4Y/2Ef56D4eIiOi8GHzZgdGxWsPQP3alIrewRO/hEBEREZETmxN/Qn0cz6b2RERkBxh82YFOUf5oEeyDwhIzlu5K0Xs4REREROSkDqbnYsuxTBhdDBhd2ouWiIjIljH4sgMGg0x31E4s5idwuiMRERER6WNuaVP7ga2DEernqfdwiIiILojBl52wvqO2en8GTuYW6j0cIiIiInIyZrMFc7dowReb2hMRkb1g8GUnWob4onNUAExmC37dwemORERERNSwNhw+hcTMM/DzdMUVHcL0Hg4REVGNMPiyI6NLpzsu5HRHIiIiImpgs0ub2l/TJQKebka9h0NERFQjDL7syDWxETAYgI1HtHfbiIiIiIgaQn5RCX7bnqwuc5ojERHZEwZfdiQiwAu9Yhqry4u2suqLiIiIiBrG7ztTkVdkQtPG3ujRrJHewyEiIqoxBl92ZkxclPrI1R2JiIiIqKGnOY7vFqVWHCciIrIXDL7szFWdwuHqYsCu5GwcSMvRezhERERE5OBSsgrw14EMdXl8V05zJCIi+8Lgy8408nHHoDYh6vICVn0RERERUT2bl5AIswXoGdMITYO89R4OERFRrTD4skOj47TVHRdsTYLFYtF7OERERETkoORcc/Zm6zRHVnsREZH9YfBlh4a1D4OXmxFHTuZj24ksvYdDRERERA5qZ1I29qflwt3VBVd3idB7OERERLXG4MsO+Xi4YliHsLKqLyIiIiKi+mxqP7xDGPw93fQeDhERUa0x+LJTY2K16Y4LtybBJE0XiIiIiIjqULHJXNZTdgKnORIRkZ1i8GWnBrYJQYCXG9JyCrHh0Em9h0NEREREDmbV3nSczCtCsK8HBrQO1ns4REREF4XBl52SPgsjO4ery5zuSERERER1bc4WbZrj2LhIuBr5soGIiOwT/4LZsVGl0x1/3Z6MwhKT3sMhIiIiIgeRlV+MP3alqctczZGIiOwZgy871rt5EML8PZBdUII/92XoPRwiIiIichALtyWhyGRGu3A/dIj013s4REREF43Blx0zuhhwTRet6mt+QqLewyEiIiIiBzGndDVHNrUnIiJ7x+DLzo2J04KvP3anIq+wRO/hEBEREZGdO5yRh/hjmXAxlJ9rEhER2SsGX3auc1QAYoK8UVBsxtJdqXoPh4iIiIjs3NzSai9ZRTzU31Pv4RAREV0SBl92zmAwYHRclLrM1R2JiIjofN59913ExMTA09MTvXv3xsaNG6s9tri4GC+++CJatmypjo+NjcXixYurPf7f//63Oi957LHH6mn01BDMZgtmx2stNNjUnoiIHAGDLwcwunR1xz/3peN0XpHewyEiIiIb9MMPP2Dy5Ml47rnnEB8fr4KsESNGIC1NW7nvbFOnTsWHH36Id955B7t27cL999+PcePGYcuWLecc+/fff6tju3Tp0gBfCdWnjUdOITHzDPw8XDG8Q5jewyEiIrpkDL4cQKtQX3SM9EeJ2YJfdyTrPRwiIiKyQTNmzMA999yDO++8Ex06dMAHH3wAb29vfPrpp1Ue/9VXX+Hpp5/GyJEj0aJFCzzwwAPq8htvvFHpuNzcXNxyyy2YNWsWGjVq1EBfDdV3U/uRnSPg6WbUezhERESXjMGXg1V9zU/gdEciIiKqrKioCJs3b8awYcPK9rm4uKjr69atq/JzCgsL1RTHiry8vLBmzZpK+x566CFcffXVle77fOR+s7OzK21kG84UmfDr9hR1eUJ3TnMkIiLHwODLQYwqDb7+PnIKSZln9B4OERER2ZCMjAyYTCaEhVWeuibXU1K0oONsMg1SqsT2798Ps9mMpUuXYs6cOUhOLq8u//7779W0yenTp9d4LHJsQEBA2RYdHX0JXxnVpd93pSC3sATRjb3Qoxmr94iIyDEw+HIQkYFe6NW8MSwWYNE2Vn0RERHRpXnrrbfQunVrtGvXDu7u7pg0aZKaJimVYuL48eN49NFH8c0335xTGXY+U6ZMQVZWVtkm90O2wdrUflzXJnBxMeg9HCIiojrB4MsBpztydUciIiKqKDg4GEajEampqZX2y/Xw8PAqPyckJATz5s1DXl4ejh49ij179sDX11f1+xIydVIa43fr1g2urq5qW7VqFd5++211WSrMquLh4QF/f/9KG+kvNbsAa/anq8vju2orhhMRETkCBl8ORJqQuroYsCMxGwfTc/UeDhEREdkIqdjq3r07li1bVrZPpi/K9b59+573c6WaKyoqCiUlJZg9ezbGjBmj9g8dOhTbt29HQkJC2dajRw/V6F4uS9BG9mN+QiLMFqgpjjHBPnoPh4iIqM641t1dkd4a+7hjQOtgrNibjgUJSfjnFW30HhIRERHZiMmTJ+P2229X4VSvXr3w5ptvqmoumb4oJk6cqAIua7+uDRs2IDExEXFxcerj888/r8KyJ598Ut3u5+eHTp06VXoMHx8fBAUFnbOfbJvFYsHszdo0x/Hd2NSeiIgcC4MvBzMmLkoLvrYm4bFhrWEwsD8DERERATfccAPS09Mxbdo01dBeAq3FixeXNbw/duxYWf8uUVBQgKlTp+LQoUNqiuPIkSPx1VdfITAwUMevgurDzqRs7E3NgburC67uHKH3cIiIiOoUgy8Hc0WHMHi6ueBwRp6a8ti5SYDeQyIiIiIbIQ3qZavKypUrK10fNGgQdu3aVav7P/s+yD7MKW1qf0X7MAR4u+k9HCIiojrFHl8OxsfDFcPah5X1aiAiIiIiqk6xyYwFW7Vzxgnd2dSeiIgcD4MvB17dceG2JJikSykRERERURVW709HRm4Rgn2lV2yI3sMhIiKqcwy+HNCgtiHw93RFanYhNh4+pfdwiIiIiMhGWZvaj46NgpuRLw2IiMjxXNRft3fffRcxMTFqeevevXtj48aN1R47Z84ctXqQNEKVlX6kkao0RqX64+FqxFWdtMak1tJ1IiIiIqKKsvKLsXR3qro8vhunORIRkWOqdfD1ww8/qOWwn3vuOcTHxyM2NhYjRoxAWlpalcc3btwYzzzzDNatW4dt27apJbNlW7JkSV2Mn6oxJk6b7vjr9hQUlZj1Hg4RERER2Zhftier88S2YX7oGOmv93CIiIhsI/iaMWMG7rnnHhVedejQAR988AG8vb3x6aefVnn84MGDMW7cOLRv3x4tW7bEo48+ii5dumDNmjV1MX6qRu8WQQj180DWmWL8uS9d7+EQERERkY2ZE3+irKm9wWDQezhERET6B19FRUXYvHkzhg0bVn4HLi7qulR0XYjFYsGyZcuwd+9eDBw48OJGTDVidDHgmi5a1deCrUl6D4eIiIiIbMiRjDxsOnoaLgaZKcBpjkRE5Lhca3NwRkYGTCYTwsLCKu2X63v27Kn287KyshAVFYXCwkIYjUa89957uOKKK6o9Xo6TzSo7O7s2w6RSo+Mi8elfh7F0Vyryi0rg7V6rbzcREREROag5W7Q+sJe1DkGYv6fewyEiIqo3DbJ0i5+fHxISEvD333/jlVdeUT3CVq5cWe3x06dPR0BAQNkWHR3dEMN0OLFNAtAsyBtnik0q/CIiIiIiMpst5dMc2dSeiIgcXK2Cr+DgYFWxlZpaOUSR6+Hh4dU/iIsLWrVqpVZ0/L//+z9ce+21KtyqzpQpU1SVmHU7fvx4bYZJpaRXw+jY0umOCZzuSERERERQUxxPnD4DXw9XDO9Q/Tk8ERGR0wVf7u7u6N69u+rTZWU2m9X1vn371vh+5HMqTmU8m4eHB/z9/SttdGmrO67al47TeUV6D4eIiIiIdGat9hrZORxe7ka9h0NERGRbUx1lmuKsWbPwxRdfYPfu3XjggQeQl5enVnkUEydOVBVbVlLZtXTpUhw6dEgd/8Ybb+Crr77CrbfeWrdfCVWpVagf2kf4o8RswW87UvQeDhERERHpqKDYhF+2JavL47s10Xs4RERE9a7W3c5vuOEGpKenY9q0aUhJSVHTFxcvXlzW8P7YsWNqaqOVhGIPPvggTpw4AS8vL7Rr1w5ff/21uh9quKqv3cnZWLA1ETf3bqr3cIiIiIhIJ7/vSkVOYQmiAr3QK6ax3sMhIiKqdwaLxWKBjZNVHaXJvfT74rTH2kvMPIP+/14OgwFY99RQhAdw5R4iInJ8PH+wD/w+Naw7PtuIlXvT8cjlrTB5eFu9h0NERFTv5w8Nsqoj6Uve0esZ0wgScS7axib3RERERM4oLacAf+5LV5fHcZojERE5CQZfTsK6uuN8ru5IRERE5JTmb0mC2QJ0axqI5sE+eg+HiIioQTD4chIjO0fA6GLA9sQsHErP1Xs4RERERNTAZpeu5sim9kRE5EwYfDmJIF8PDGgdrC4v2MqqLyIiIiJnsispG3tScuBudMGoLtpMACIiImfA4MsJpztK8GUHaxoQERERUR2ZU1rtNaxDKAK83fQeDhERUYNh8OVEhncMh4erCw6l52FnUrbewyEiIiKiBlBiMmNeaZ/X8V05zZGIiJwLgy8n4uvhimHtw9RlTnckIiIicg6r92cgI7cQjX3cMahtiN7DISIialAMvpzM6LjS6Y4JSTDLsj5ERERE5BRN7aXthZuRp/9ERORc+JfPyQxuGwI/T1ekZBdg45FTeg+HiIiIiOpR1pli/L4rVV2+tjunORIRkfNh8OVkPFyNuKpTuLrM6Y5EREREju237ckoKjGjTZgvOkb66z0cIiKiBsfgywmNjo1SH38tPREiIiIiIsee5ji+WxMYDAa9h0NERNTgGHw5ob4tgxDs64HM/GKsOZCu93CIiIiIqB4cPZmHv4+chosBGBunvfFJRETkbBh8OSGjiwHXdIlQl+eXLm1NRERERI5l7pZE9bF/q2CEB3jqPRwiIiJdMPhyUmNKV3dcuisV+UUleg+HiIiIiOqQxWLBnHgt+JrQjU3tiYjIeTH4clJx0YFo2tgb+UUm/LE7Te/hEBEREVEd2nT0NI6dyoePuxHDO4bpPRwiIiLdMPhyUtLcdHSsVvW1gNMdiYiIiBzKnNKm9ld1joC3u6vewyEiItINgy8nNrp0uuOqfWnIzC/SezhEREREVAcKik1YtC1ZXeY0RyIicnYMvpxYmzA/tAv3Q7HJgsU7UvQeDhERERHVgT92pyKnoARRgV7o3byx3sMhIiLSFYMvJ2et+uLqjkRERESOYfZmbZrjuK5RcHEx6D0cIiIiXTH4cnKjumjB1/rDJ5GaXaD3cIiIiIjoEqTlFODP/Rnq8rhuUXoPh4iISHcMvpxcdGNvdG/WCBYLsHArq76IiIiI7JksWmQyW9C1aSBahvjqPRwiIiLdMfgijCmd7sjgi4iIiMi+zYlPVB/Hs6k9ERGRwuCLMLJzBIwuBmw9kYXDGXl6D4eIiIiILsLu5GzsSs6Gm9GAUV0i9B4OERGRTWDwRQj29UD/VsHqMqu+iIiIiOzTnHitqf3QdmEI9HbXezhEREQ2gcEXKWNiras7JsIiDb+IiIiIyG6UmMyYV7pK93g2tSciIirD4IuU4R3D4OHqgoPpeapEnoiIiIjsx5oDGUjPKURjH3cMbhuq93CIiIhsBoMvUvw83TC0fWjZakBEREREZH9N7UfHRsLdlaf4REREVvyrSGXkRMna58ts5nRHIiIiInuQXVCMJTtT1GVOcyQiIqqMwReVkbJ4Pw9XJGUVYNPR03oPh4iIiIhq4LftySgsMaNVqC86RwXoPRwiIiKbwuCLyni6GTGiU3hZk3siIiIisn2zS6c5TujWBAaDQe/hEBER2RQGX1TJmDhtuuOv25NRbDLrPRwiIiIiOo/jp/Kx8fApSN41tqt2HkdERETlGHxRJX1bBCHY1x2n84uxZn+G3sMhIiIioho0te/fMhgRAV56D4eIiMjmMPiiSlyNLrimi/Zu4YKtXN2RiIiIyFZZLBbM2XJCXWZTeyIioqox+KJzjCpd3VFWBzpTZNJ7OERERERUhfhjp3H0ZD683Y24srRPKxEREVXG4IvO0a1pIJo08kJ+kQnL9qTqPRwiIiIiOk9T+6s6RcDb3VXv4RAREdkkBl90DlkNaHRp1df8BE53JCIiIrI1BcUmLCptSzGB0xyJiIiqxeCLqjQmTjuBWrU3HVn5xXoPh4iIiIgqWLY7DdkFJYgM8ESfFkF6D4eIiMhmMfiiKrUN90PbMD8UmcxYvDNZ7+EQERERUQVz4rWm9uO6RcHFxaD3cIiIiGwWgy+q1ug4ru5IREREZGsycguxcl+6ujyuaxO9h0NERGTTGHxRtax9vtYePIm07AK9h0NERERE8qZkQhJMZgtiowPRKtRX7+EQERHZNAZfVK3oxt5qhUeLBVi0jdMdiYiIiGzB7NJpjmxqT0REdGEMvui8ylZ35HRHIiIiIt3tScnGzqRsuBkNGNVFO08jIiKi6jH4ovO6ukskpF/q1uOZOHoyT+/hEBERETm1ufGJ6uPl7ULRyMdd7+EQERHZPAZfdF4hfh7o3yq4rJ8EEREREelD+nrN3aIFX+O7sak9ERFRTTD4olpNd7RIwy8iIiKyS++++y5iYmLg6emJ3r17Y+PGjdUeW1xcjBdffBEtW7ZUx8fGxmLx4sWVjpk+fTp69uwJPz8/hIaGYuzYsdi7d28DfCXOac2BDKTlFCLQ2w1D2obqPRwiIiK7wOCLLmhEp3C4u7rgQFoudifn6D0cIiIiugg//PADJk+ejOeeew7x8fEqyBoxYgTS0tKqPH7q1Kn48MMP8c4772DXrl24//77MW7cOGzZsqXsmFWrVuGhhx7C+vXrsXTpUhWWDR8+HHl5bI9QH+aUNrWXNyXl3IyIiIguzGCxgxKe7OxsBAQEICsrC/7+/noPxynd/9VmLN6ZgvsHtcRTV7XTezhEREQXxPOHyqTCS6qzZs6cqa6bzWZER0fj4YcfxlNPPXXO8ZGRkXjmmWdUsGU1YcIEeHl54euvv67yMdLT01XllwRiAwcOrNG4+H2qmZyCYvR85Q8UFJsx/6H+iI0O1HtIREREuqnN+QPfKqIaGROnTXdcuDUJZrPNZ6VERERUQVFRETZv3oxhw4aV7XNxcVHX161bV+XnFBYWqimOFUnotWbNmmofR04+RePGjets7KT5bUeKCr1ahvigS5MAvYdDRERkNxh8UY0MaRcKXw9XJGaeQfyx03oPh4iIiGohIyMDJpMJYWFhlfbL9ZSUlCo/R6ZBzpgxA/v371fVYTKVcc6cOUhOTq7yeDnmscceQ//+/dGpU6dqxyKBmrxLW3GjC5u9+URZU3uDwaD3cIiIiOwGgy+qEU83I0Z0DFeX53N1RyIiIof31ltvoXXr1mjXrh3c3d0xadIk3HnnnapSrCoyJXLHjh34/vvvz3u/0hBfpiZYN5luSed3/FQ+Nhw+Bcm7xnWN0ns4REREdoXBF9XY6NLpjr9sT0axyaz3cIiIiKiGgoODYTQakZqaWmm/XA8P197YOltISAjmzZunGtUfPXoUe/bsga+vL1q0aHHOsRKKLVq0CCtWrECTJk3OO5YpU6aoKZHW7fjx45f41Tm+eVsS1cd+LYMQGeil93CIiIjsCoMvqrH+LYMQ5OOOU3lF+OtAht7DISIiohqSiq3u3btj2bJllaYmyvW+ffue93Olz1dUVBRKSkowe/ZsjBkzpuw2WSNJQq+5c+di+fLlaN68+QXH4uHhoZrQVtyoevJ/PKc0+Brf9fyhIhEREZ2LwRfVmKvRBVd3iVCXF3C6IxERkV2ZPHkyZs2ahS+++AK7d+/GAw88oKq5ZPqimDhxoqrGstqwYYPq6XXo0CGsXr0aV155pQrLnnzyyUrTG2WFx2+//RZ+fn6qX5hsZ86c0eVrdERbjmficEYevNyMuLJT1dV5REREVD3X89xGVOXqjl+uO4olO2VlIZPq/UVERES274YbbkB6ejqmTZumwqm4uDgsXry4rOH9sWPHKvXvKigowNSpU1XwJVMcR44cia+++gqBgYFlx7z//vvq4+DBgys91meffYY77rijwb42Z2hqf1WncPh48NSdiIiotvjXk2qlW9NGiAr0Uqs7LtudVlYBRkRERLZPpiXKVpWVK1dWuj5o0CDs2rXrgtPwqP4UlpiwcKtWZT+hO6c5EhERXQxOdaRakeWzrU3uF2zV+k0QERERUd1bvjsN2QUliAjwRJ8WQXoPh4iIyC4x+KJaGx2rBV8r9qQj60yx3sMhIiIickiz47U3Gcd2jYLRxaD3cIiIiOwSgy+qtXbhfmgT5osik1n1+iIiIiKiunUytxAr96apy+O7Ruk9HCIiIrvF4IsubrpjadUXV3ckIiIiqnsLtiahxGxBlyYBaB3mp/dwiIiI7BaDL7ooo2O1dx7XHsxAWk6B3sMhIiIicihzSqc5TujGpvZERESXgsEXXZSmQd6Iiw6E2QL8si1Z7+EQEREROYx9qTnYnpgFVxcDRpVW2RMREdHFYfAl9i4Gis/oPQq7M6ZsdUdOdyQiIiKqK7PjT6iPQ9qForGPu97DISIismsMvvYtAb67EfjsKiCblUu1cXWXCMgCQ1uOZeLYyXy9h0NERERk90xmC+ZtsU5zZFN7IiKiS8Xgy80b8AoEkrYAs4YAiZv1HpHdCPXzRL+Wwerywm2s+iIiIiK6VNI/NTW7EIHebqrii4iIiC4Ng6/mA4B7VgAh7YCcZOCzkcD2n/Ueld2wru44P0F7Z5KIiIiILr2p/agukfBwNeo9HCIiIrvH4Es0bg78YynQegRQUgDM/gew/GXAbNZ7ZDZvRKdwuBtdsC81F3tSsvUeDhEREZHdyi0sweIdKeryeE5zJCIiqhMMvqw8/YGbvgP6PaJd//M/wI+3AYW5eo/MpgV4uWFw2xB1eX4CpzsSERERXazftifjTLEJLYJ91OrZREREdOkYfFXkYgSGvwSMfR8wugN7FgGfXglkHtN7ZDZtTJz2juSChCRYLBa9h0NERERk19McJ3RvAoPBoPdwiIiInDf4evfddxETEwNPT0/07t0bGzdurPbYWbNmYcCAAWjUqJHahg0bdt7jbULczcDtiwCfECB1OzDrcuDYBr1HZbOGtg+Fj7sRiZlnEH/stN7DISIiIrI7J07nY92hk+ry2K6c5khERKRb8PXDDz9g8uTJeO655xAfH4/Y2FiMGDECaWlpVR6/cuVK3HTTTVixYgXWrVuH6OhoDB8+HImJNt4MvWlvrel9eGcgLx344hpgyzd6j8omeboZMaJjeFnVFxERERHVjrVlRN8WQYgK9NJ7OERERM4bfM2YMQP33HMP7rzzTnTo0AEffPABvL298emnn1Z5/DfffIMHH3wQcXFxaNeuHT7++GOYzWYsW7YMNi8wGrhrCdB+FGAqAuY/CPw+FTCb9B6ZzRkdp63uuGhbMkpMXBSAiIiIqKakVcTszSfUZTa1JyIi0jH4KioqwubNm9V0xbI7cHFR16Waqyby8/NRXFyMxo0bwy64+wDXfQkMfFK7vvYd4LsbgQKuYFhR/1bBaOzjjpN5RfjroFamT0REREQXlnA8E4cy8uDlZsRVnSP0Hg4REZHzBl8ZGRkwmUwICwurtF+up6RoSy9fyL/+9S9ERkZWCs/OVlhYiOzs7EqbrlxcgMufAa79FHD1BPb/DnxyBXDqkL7jsiFuRhdcXXqixumORERERLVvan9lp3D4erjqPRwiIiKH0qCrOv773//G999/j7lz56rG+NWZPn06AgICyjbpC2YTOk0A7vwN8IsA0vdoTe8P/6n3qGxuuuOSnSkoKOZ0UCIiIqILKSwxYeE27U1DTnMkIiLSOfgKDg6G0WhEampqpf1yPTxca25enf/+978q+Pr999/RpUuX8x47ZcoUZGVllW3Hjx+HzYjqpjW9j+wGnDkNfDUO2FR1fzNn071pI9WMNbewBCv2VL3YARERERGVk3OmzPxihPl7oF/LYL2HQ0RE5NzBl7u7O7p3716pMb21UX3fvn2r/bzXX38dL730EhYvXowePXpc8HE8PDzg7+9fabMp/hHAnb8Cna8DzCXAon8Cvz4BmErgzFxcDLgmNqLSykREREREVL3ZpdMcx3aNgtHFoPdwiIiIHE6tpzpOnjwZs2bNwhdffIHdu3fjgQceQF5enlrlUUycOFFVbFm99tprePbZZ9WqjzExMaoXmGy5ubmwa25ewPhZwNBp2vWNHwHfTNCqwJzYmFitRH/53jRkFxTrPRwiIiIim3Uqr6isSn5CtyZ6D4eIiMgh1Tr4uuGGG9S0xWnTpiEuLg4JCQmqksva8P7YsWNITk4uO/79999Xq0Fee+21iIiIKNvkPuyewQAM+D/ghm8ANx/g0Epg1lAgfR+cVfsIP7QK9UVRiRlLdtRswQMiIiIiZ7RwaxJKzBZ0jgpAmzA/vYdDRETkkAwWi8UCGyerOkqTe+n3ZXPTHq1SdgDf3QRkHQM8ArQVIFtXv3KlI3tn2X68sXQfBrQOxlf/6K33cIiIyEnZxfkDOfX3aczMNdh6IgvPjeqAO/s313s4REREDnn+0KCrOjq08E7APcuBpn2Bwizg2+uAde8Btp8r1rlRsdrqjn8dyEB6TqHewyEiIiKyOQfSclTo5epiwOjScyciIiKqewy+6pJvCDBxPhB3K2AxA0umAAsfAUqK4Exign0QGx0IswX4dXv5tFciIiIiqtzUfnDbUAT5eug9HCIiIofF4KuuuXpI3Tow4lXA4ALEfwl8OQbIy4Azsb5zOT9BO6kjIiIiIo3JbMG8Ldo50oRu2sJAREREVD8YfNVX0/u+DwE3/wh4+APH1gKzhgCpO+EsRnWJUP8N8ccycfxUvt7DISIiIrIZ6w+dRHJWAfw9XXF5+1C9h0NEROTQGHzVp9ZXAHf/ATRqDmQeAz4ZDuz5Fc4g1N8TfVsEqcsLtibpPRwiIiIimzF784myvqgerka9h0NEROTQGHzVt5C2WtP75gOBolzg+5uB1TOcoun9mLjIsqW6iYiIiAjIKyzBbztS1OUJ3ZvoPRwiIiKHx+CrIXg3Bm6dA/S8G4AFWPYCMOdeoLgAjuzKjhFwMxqwJyUHe1Ny9B4OERERke4W70jBmWITmgf7oGt0oN7DISIicngMvhqK0Q24+g1g5H8BgxHY/iPw+dVAjvaOnyMK8HZTKxWJBVvZ5J6IiIhozhZtmuP4rlEwSENUIiIiqlcMvhpar3uA2+YCnoFA4iZg1uVAUgIcfXVH6fNlcYLpnURERETVSco8g7UHT6rLY7tyNUciIqKGwOBLDy0GaX2/gtsA2YnAp1cCO+fCEQ1rHwZvdyOOnzqDLccz9R4OERERkW7mbklUbV77tGiM6Mbeeg+HiIjIKTD40ktQS23Fx1ZXACVngJ/uAFZMB8xmOBIvdyOGdwhTlxcksMk9EREROSepfJ8TXzrNsRub2hMRETUUBl968gwAbv4B6DtJu77q38DPdwBFeXAkY+K0Uv5F25JRYnKsYI+IiIioJradyMLB9Dx4urngqk7heg+HiIjIaTD40puLERjxCjDmXcDFDdg1X5v6mKW9I+gILmsdjEbebsjILcS6Q1pfCyIiIiJnYq32GtExHH6ebnoPh4iIyGkw+LIVXW8Fbl8IeAcDKduAj4YAx/+GI3AzumBk5wh1eT6nOxIREZGTKSoxq4V+BKc5EhERNSwGX7akWV/g3hVAWCcgLw34/Gpg6/dwpOmOS3akoKDYpPdwiIiIiBrMir1pOJ1fjFA/D1zWKljv4RARETkVBl+2JrApcNcSoN01gKkQmHsfsPQ5wGzfYVGPZo0QEeCJnMISrNybpvdwiIiIiBp8muO4rlEwuhj0Hg4REZFTYfBlizx8geu/AgY8rl3/603g+1uAwhzYKxcXA0bHRqrL1lJ/IiIiIkd3Oq8Iy/dob/pxmiMREVHDY/Blq1xcgKHPAuM/BowewL7fgE+GA6ePwF6NKg2+/tidhpyCYr2HQ0RERFTvFm5LQrHJgo6R/mgb7qf3cIiIiJwOgy9b1+U64M7fAN9wIG2X1vT+yF+wR3LC1zLERzV4/X1nqt7DISIiIqp3s+MT1ccJrPYiIiLSBYMve9Cku9b0PrIrcOYU8OVoYPMXsDcGg0x31Jrcz+d0RyIiInJwB9JysfV4purrNTpOq3wnIiKihsXgy174RwJ3/Ap0HA+YS4CFjwC/PQWYSmBPrCd9fx3IQEZuod7DISIiIqo3c7doTe0HtwlBsK+H3sMhIiJySgy+7Im7N3Dtp8CQqdr1De8D31wLnDkNe9E82AddmgTAZLbg1+3Jeg+HiIiIqF6YzRbMLZ3myKb2RERE+mHwZW8MBmDQE9qqj27ewKEVwMfDgIwDsBdlqzsmcLojEREROab1h04iKasA/p6uGNo+VO/hEBEROS0GX/aqw2jgriWAfxPg5AHg48uBg8thL6s7Sn636ehpnDidr/dwiIiIiOqtqf01sZHwdDPqPRwiIiKnxeDLnkV00ZreN+kFFGQBX18LbPgIsFhgy8L8PdGneZC6vHArpzsSERGRY8kvKsFvO7RznAndtIV9iIiISB8MvuydbyhwxyIg9mbAYgJ+ewJY9BhQUgR7aHI/P0F7N5SIiIjIUSzekYL8IhNigrzRrWkjvYdDRETk1Bh8OQJXD2Dse8AVL0kTMGDz58BX44C8k7BVV3UKh5vRgD0pOdiXmqP3cIiIiIjqzJwKTe0N0t+BiIiIdMPgy1HISVX/R4CbfwDc/YCja4BZQ4C03bBFgd7uGNQmRF1mk3siIiJyFMlZZ/DXwQx1eVxXTnMkIiLSG4MvR9NmBHD3UqBRDJB5FPj4CmDvYtii0XHayeCCrUmw2HhfMiIiIqKamLdFzmuAXs0bI7qxt97DISIicnoMvhxRaHvg7uVAs8uAohzguxuBv96yuab3w9qHwsvNiGOn8pFwPFPv4RARERFdEnkjb078CXWZTe2JiIhsA4MvR+UTBNw2F+h+p5yGAUunAfMeBEoKYSu83V0xvGOYuvz1+mMoNpn1HhIRERHRRduemIX9abnwcHXBVZ0j9B4OERERMfhycK7uwDX/A0b+FzAYga3fAp9fA+SkwlaMLZ3uODv+BPr9ezlmLN2HlKwCvYdFREREdNFN7Ud0DIe/p5vewyEiIiIGX07S9L7XPcCtswHPAODERmDW5UDyVtiCwW1D8PTIdgj29UB6TiHeXrYf/V9bjge+3oy1BzLY+4uIiIjsQlGJWfUtFeM5zZGIiMhmMPhyFi2HaH2/gloD2SeAT68Eds3Xe1Rqie97B7bE2qcuxzs3dVWNYE1mC37bkYKbP96AYTNW4bO/DiO7oFjvoRIRERFVa9W+dJzKK0KInwcuaxWs93CIiIioFIMvZxLcCrj7D6Dl5UBxPvDjRGDV6zbR9N7d1QWjYiPx4319seSxgbi1T1P4uBtxMD0PLyzchd6vLMOUOduxKylb76ESERERnWP2Zq2p/di4SLgaeYpNRERkK/hX2dl4BQI3/wT0eVC7vuIV4Oe7gKJ82Iq24X54eWxnrH96KF4c0xGtQ31xptiE7zYew8i3V+Pa99difkIiCktMeg+ViIjIrrz77ruIiYmBp6cnevfujY0bN1Z7bHFxMV588UW0bNlSHR8bG4vFixdf0n06qsz8Iizbo/VQndC9id7DISIiogoYfDkjoytw5XRg1NuAixuwcw7w2VVAttaXwlb4ebphYt8Y/P7Pgfj+3j64uksEXF0M2HT0NB79PgH9/70c/1myB4mZZ/QeKhERkc374YcfMHnyZDz33HOIj49XQdaIESOQlpZW5fFTp07Fhx9+iHfeeQe7du3C/fffj3HjxmHLli0XfZ+OauG2ZBSbLOgQ4Y924f56D4eIiIgqMFjsoHt4dnY2AgICkJWVBX9/nkzUqSN/AT/eBuSfBHzDgRu/BZp0h61KzS7A9xuP49uNR5GaXaj2uRiAoe3DcFufZqqnhovsICIip8fzh8qkGqtnz56YOXOmum42mxEdHY2HH34YTz311DnHR0ZG4plnnsFDDz1Utm/ChAnw8vLC119/fVH36ajfp3Hv/YUtxzIx9er2uHtAC72HQ0RE5PCya3H+wIovZxfTH7hnORDaAchN0Sq/tv0EWxXm74lHh7XGmn9djvdv6YZ+LYNgtgBLd6Vi4qcbMXTGKny8+hCy8tkMn4iIyKqoqAibN2/GsGHDyva5uLio6+vWravycwoLC9X0xYok9FqzZs1F36f1fuVkteJmzw6m56rQy+hiwJg4ruZIRERkaxh8EdAoBvjH70CbqwBTITDnbmDZi/K2LWyVm9EFV3WOwLf39MEfkwfijn4x8PNwxeGMPLz8y270nv4Hnvx5K7afyNJ7qERERLrLyMiAyWRCWFhYpf1yPSUlpcrPkSmLM2bMwP79+1Ul19KlSzFnzhwkJydf9H2K6dOnq3dorZtUiNmzufGJ6uOgNiFqRUciIiKyLQy+SOPhB9z4DXDZP7Xrq98AfrgVKMyFrWsV6ofnR3dUzfBfHdcZ7cL9UFBsxo+bTmDUzDUY++5faqWlgmI2wyciIqqpt956C61bt0a7du3g7u6OSZMm4c4771RVXZdiypQpalqCdTt+/DjsldlswdwtWvA1vhurvYiIiGwRgy8q52IEhj0PjPsIMHoAe38BPh0BnD4Ke+Dj4YqbezfFb48OwM/398Xo2Ei4GQ1IOJ6J//tpK/pOX4bpv+3G8VO2s4IlERFRQwgODobRaERqqrbyoJVcDw8Pr/JzQkJCMG/ePOTl5eHo0aPYs2cPfH190aJFi4u+T+Hh4aF6cVTc7NWGw6fUIjt+nq4Y1r5y5RsRERHZBgZfdK7YG4A7fgF8QoHUHcCsIcDR6nt12BqDwYAeMY3x9k1dsfapoXhiRFtEBnjidH4xPlx1CAP/swJ3frYRK/akwSQNwoiIiBycVGx1794dy5YtK9sn0xflet++fc/7udLnKyoqCiUlJZg9ezbGjBlzyffpKObEn1Afr+kSAU83o97DISIioiow+KKqRfcE7l0BhHfRVnz8YhQQ/xXsjfTaeGhIK/z55BB8dFt3DGgdDFnHdMXedNz5+d8Y/N8V+HDVQZzKK9J7qERERPVq8uTJmDVrFr744gvs3r0bDzzwgKrmkumLYuLEiWoaotWGDRtUT69Dhw5h9erVuPLKK1Ww9eSTT9b4Ph1ZflEJft2u9Tub0K2J3sMhIiKiarhWdwMRApoAdy0G5j0A7JoPLJgEpO8BrnhRmxZpR1yNLhjeMVxt0gD/6/VH8dOm4zh+6gym/7YHbyzdp96tva1PM8RFB6qqMSIiIkdyww03ID09HdOmTVPN5+Pi4rB48eKy5vTHjh2r1L+roKAAU6dOVcGXTHEcOXIkvvrqKwQGBtb4Ph3Z7ztTkVdkQtPG3ujerJHewyEiIqJqGCwWqX+xbbLMtaz6Iw1Q7bkPhN2S1R3/fB1YOV273moYcO2ngGcA7NmZIhMWbk3Cl+uPYEdi+VLqnaMCVAA2KjYSXu72FfAREVE5nj/YB3v9Pt32yQas3p+Bx4a1xmPD2ug9HCIiIqeSXYvzB051pAuTd38HPwVc9wXg6gUc+AP4eBhw8iDsmYRa1/eMxsJJl2Hug/3Uakzuri7YnpiFJ2dvQ5/py/Dyol2qQoyIiIjIKiWrAH8dyFCXx3flNEciIiJbxuCLaq7jWG3qo38UkLEPmHU5cGgV7J1Ma+zatBFmXB+H9VOG4qmr2qFJIy9knSnGx2sOY8h/V2LipxuxdFcqm+ETERER5iUkQk4JesY0QtMgb72HQ0REROfB4ItqJzIOuGc5ENUDKMgEvhoHbJwFR9HYxx33D2qJVU8Mwad39MCQtiGQdl9/7kvHPV9uwsDXV+DdFQeQkVuo91CJiIhIB9IlZPZmbTVHNrUnIiKyfQy+qPb8woE7fgG63ABYTMCvjwOLJgOmYjgKo4sBl7cLw2d39sKqx4fgvoEtEOjthsTMM/jPkr3oO30ZHv1+CzYfPaVOgImIiMg57EzKxv60XNUeYWSXCL2HQ0RERBfA4IsujpsnMO5DYNjzMlkQ2PQJ8PV4IP8UHI1MYZgysr2aBvnGdbGIjQ5EscmC+QlJmPD+Oox8ew2+3XAMeYUleg+ViIiI6tnseK3aa3iHMPh7uuk9HCIiIroABl908WQO4GX/BG76DnD3BQ7/qfX9St8LR+TpZsSE7k0w/6H+WDCpP67v0QQeri7YnZyNp+duR59Xl+H5BTtxIC1X76ESERFRPSg2mbEgIUld5jRHIiIi+8Dgiy5d26uAfywFApsCpw9rKz7uXwpH1qVJIF6/NhYbnh6KqVe3R0yQN3IKS/D52iMYNmMVbp61Hot3JKPEZNZ7qERERFRHVu1Nx8m8IgT7emBA62C9h0NEREQ1wOCL6kZYB+CeFUCz/kBhNvDt9cDamdIBFo4s0Nsddw9ogeX/Nxhf3NULw9qHwcUArD14Evd/HY/LXluBt5ftR1p2gd5DJSIioks0Z4s2zXFsXCRcjTyNJiIisgf8i011xycYuG0e0G0iYDEDvz8DzJ8ElDj+CoguLgYMahOCj2/vgT+fHIIHB7dEkI87UrILMGPpPvT793I89G081h86yWb4REREdigrvxh/7EpTl8dzmiMREZHdYPBFdcvVHRj1NnDla4DBBUj4GvhiNJCbDmfRpJE3nryyHdZOuRxv3RiH7s0aocRswS/bknHjR+sx4s0/8dW6I8hlM3wiIiK7sWh7EopMZrQL90OHSH+9h0NEREQ1xOCL6qfpfZ/7gVt+BjwCgOPrgVlDgJQdcCYerkaMiYvC7Af64ZdHLsNNvZrCy82Ifam5eHb+TvR+5Q88O28H9qbk6D1UIiIiuoDZm7Vpjtd2Z7UXERGRPWHwRfWn1VDgnmVA45ZA1nHgk+HA7kVwRh0jAzB9fGesf3oonhvVAS1CfJBXZMJX64+qCrDrP1yHRduSUFTCZvhERES25nBGHuKPZao+nqPjIvUeDhEREdUCgy+qX8GttfCrxWCgOA/44Rbgz/86fNP76gR4ueHO/s2xbPIgfHN3b1zZMRxGFwM2Hj6FSd9uQf/XlmPG73uRnHVG76ESERFRqbnxWrXXwDYhCPXz1Hs4REREVAsMvqj+eTUCbpkN9LpPu778JWD23UCx84Y7BoMB/VsF44PbumPNv4bgkctbqaXR03MK8fbyA2o1yPu/2oy1BzLYDJ+IiEhHZrMFc7Ykqstsak9ERGR/GHxRwzC6AiNfB675H+DiCuz4GfhsJJCdDGcXEeCFycPbYu1Tl+Odm7qiV/PGMJktWLwzBTd/vAHDZqzCZ38dRtaZYr2HSkRE5HQ2HjmFE6fPwM/DFcM7hOk9HCIiIqolBl/UsHrcBdw2T6sCS4rXmt4nxus9Kpvg7uqCUbGR+PG+vljy2EDc1qcZfNyNOJiehxcW7kKfV5dhypzt2JWUrfdQiYiInMac0mmOV3eJgKebUe/hEBERUS0ZLHYwjyo7OxsBAQHIysqCvz+Xj3YIpw4D390IpO8BXD2B7ncCfuGATzDgEwJ4y8cg7aO7j7ZSpBPKKSjGvC2J+HLdUexPyy3b36NZI9zWtxmu7BSuVo8kIqJz8fzBPtjy9+lMkQk9X/kDuYUl6o0pqcomIiIi+zp/cG2wURFV1Lg58I+lWq+v/UuADe9Xf6wEYyoMC9KCMRWKBVe+bA3K5Dh3X4cJyvw83XBb3xjc2qcZNhw+pVaBXLIjBZuOnlZbsK87bugZjZt7N0NUoJfewyUiInIov+9KUaFXdGMv9aYTERER2R8GX6QfT3/gpu+Ard8BabuBvAwgLx3Il48ntY8lBdqWdVzbasLoURqGBVVdQVa2r/R2D3+bD8qkGX6fFkFqS80uwPcbj+PbjUeRml2Id1ccxPsrD+LydmGqCmxAq2C4yHrrREREdElmx2tN7cd1bcK/rURERHaKwRfpy8UIdL216ttkFm5RXmkYdlILxlQollHFPrmcDpScAUyFQHaittWE0f3carKKFWRn7/MM1DUoC/P3xKPDWuPBIS2xbHeqmga59uBJ/LE7VW0xQd6qQuy67tEI8HbTbZxERET2TN5oWrM/XV2e0C1K7+EQERHRRWLwRbZLwiUPX22TqZE1oYKys8Iwa1hmDcoqVpUV5wGmIiAnWdtqwsWtQlBWg6oyCcpc6n4dCTejC67sFKG2A2k5+Hr9MczefAJHTubj5V9247+/78Xo2Ejc1icGnZsE1PnjExERObL5CYkwW7S+ms2CfPQeDhEREV0kBl/kWKQRvmyNmtXs+OIz1VSQWcOyih9PAkU5gLkYyE3RtpowGCsEZDWoKpMVL2sZlLUK9cPzozviiRFtMT8hCV+uO4I9KTn4cdMJtcVGB2Jin2ZckYqIiKgGZO2n2Zu1yvHx3ZroPRwiIiK6BAy+yLm5eQGB0dpWE8UF54Zh1VaVnQQKswGLCchL07aaMLhoQVlZGFaxquzsBv8hpUGZFmb5eLji5t5NcVOvaGw+elo1w/91ezK2Hs/E/x3PxMu/7ML1PaNxa+9miG7sfQn/cUR2rjAXOLYeOLpGq+LscgMQ3ErvURGRjdiVnI29qTlwd3VRbxoRERGRkwVf7777Lv7zn/8gJSUFsbGxeOedd9CrV68qj925cyemTZuGzZs34+jRo/jf//6Hxx577FLHTaQPN08goIm21URJ4blh2PmqygqyAItZO1Y2rbXIBRgA78aVVrs0eAejh08wejQPxkstA7D8uBk/7z6DPTme+HhVAT768xAGtwnBxL4xGNgmBEY27CVHJ9WdxzcCR1YDh/8EEjcD5pLy2/98HWjaD+g2EegwBnBnMEzkzKzVXld0CEOAF/tlEhEROVXw9cMPP2Dy5Mn44IMP0Lt3b7z55psYMWIE9u7di9DQ0HOOz8/PR4sWLXDdddfhn//8Z12Nm8g+uHoA/pHaVhMlRVo4Vm1fsrP2FWTKhIzSzzkJZOw95y79AYwt3eAJmGFAlsUHJw/7q221WyAQ3Bqng7qjMLw7/AJD0NjHHcG+7upjoLc7gzGyP/JcknDLGnRJ6CULX1QU2BSIGag9t/b/Dhxbq22/PQl0vhboehsQ2dXmV30lorpVbDJjwVYt+GJTeyIiIvtnsEgTg1qQsKtnz56YOXOmum42mxEdHY2HH34YTz311Hk/NyYmRlV71bbiKzs7GwEBAcjKyoK/v7yMJyLFVAzknzrPapdnVZWdOa0FZdUwWwzYa4nGRnNbbDK3xd/mtkgzBKGRtxaCBfm6I8jHo8Jl2e9RdjnI1wOBXm5c8p0anqkESN4KHPlTC7pkGmNxfuVj/CKAmAFA84FA8wFAo5jy27KTgIRvgC1fA6ePlO8P66xVgXW5TptWTHaF5w/2wda+T8v3pOKuzzepN4DWTRmqFpMhIiIi+z1/qFXFV1FRkZqyOGXKlLJ9Li4uGDZsGNatW4e6UlhYqLaKXxARVcHoBviFaVtNwwEJv0oryIqy0rD74EEYU7cjPHMLgotOoL3hGNq7HMPtWKo+5YQlGH8XtcWmgrbYmN4OGyyRsKD6FwGSeUlQJmGYCshKgzEtLPMoDcusFWUMyugimc1A6o7yiq6ja7WeehVJT7yyoGsgENSq+uotqcoc+ARw2f9p97nlK2DXAiB1O/DbE8DvU4EOo7UqMLnPeliplYhsw+x4rdprdGwUQy8iIiIHUKvgKyMjAyaTCWFhlV9ky/U9e/bU2aCmT5+OF154oc7uj4hKGV0B3xBtk0UwAcTGVbg9JxU4vl6rljm2DpbkbWiCDDQxZmCc8S91SKGrP074dcF+z07YZmiPLaYYpOYBJ/OKkHWmWC39LpdlqwnJvCQIs27WcExVllkryUorzBiUOTEpTk7fWxp0rQKOrCmtYKzAMwBodll5RVdI+9oHVHJ8i0HadtUpYPtPQPyXWsgml2WTSjEJwOJurvk0ZiKyC1n5xVi6K1VdntCd0xyJiIgcgU2u6igVZdJHrGLFl0ynJKJ6JpVj0thbNpkLLSvfnfi7LAiTyx7F2Wh5eg1aYg2ulIOMHkBUdyCuD0qa9Mbpxl2RYfLCqbwiZOQWqo/aZflYiJPqY1GloExuk60mpN9YI2+3smoyCceCS6dcll8un5YpTYkZlNlp0HXqUHlF1+HV566M6u4LNO1bHnSFdylb4bROyKIRve8Det0LJG3RqsC2/6xNhVz+ErDiFaD1cC0EazNCq8AkIrv2y/ZkFJWY0S7cDx0i9J92SURERA0cfAUHB8NoNCI1VXsnzEquh4eHo654eHiojYh05uELtByibdaeYinbS4MwaQS+Xps2WdoUXH6hhMCAkNAOQLO+WijRqQ8Q0LzaBsKnS0MwCcRO5mlBmXZZPlYMzgqRXVACk9lSISjLrXFQdr7eZBUrzRiU6SjzeOWgK/tE5dtdPYHo3uVTF6XxfEOETTI9Mqqbtg1/Bdg1X6sCk5/7fYu1zSdUqwCTECy4Vf2PiYjqxZx47ffO+G5RMHBhCyIiIucLvtzd3dG9e3csW7YMY8eOLWtuL9cnTZpUX2MkIlshIYM1AOj7YHlVjlSDyXZ0HXDqIJC2U9v+/lj7vIBooGmf0q1v2RQ06Z0S6u+ptpqwBmVa9ZgWlJVXkFWuJjtZZVCGGgZlFadXljfuP/uy9Cnz92RQdtFkaq116qIEXacPV77dxQ1o0rO8oksuy0qpenL3BuJu0raM/VoVWMK3WjXaX29qW7P+WgAmlZNyPBHZhaMn87Dp6Gk1BX9sHKc5EhEROe1UR5mCePvtt6NHjx7o1asX3nzzTeTl5eHOO+9Ut0+cOBFRUVGqT5e1If6uXbvKLicmJiIhIQG+vr5o1YrvihPZNXk3PKiltnW9VduXm1ZaEVY6PVJW2ss6DmyX7afyXkzRFYIwqdxxu3D4VdugTKarnM6vqpqscmXZqXOCskK1oXJx63mDMq1Zv7Vxf/XVZU4dlMkKpCroKq3qythb+XaDUftZsAZd8jNiy8FRcGvgiheBy58F9i3RqsAOLAWO/qVtvz0JdL4O6HYbEBFXfWN9IrKppvYDWofU+O8MERER2T6DxSIlG7Uzc+ZM/Oc//0FKSgri4uLw9ttvo3fv3uq2wYMHIyYmBp9//rm6fuTIETRvfu40p0GDBmHlypV2ucw1EdWC9AlL3FQehB3/GyjOq3yM0V3rE2YNwqJ7AV6NGnyo1qCsqt5k5Ze1kEwCs5yCklo/hgRl1sox6xRLadjv4+EKP09X+Hq4qsu+1s2z8mVvN6P9BGcFWdpqi9agS1ZIrMQAhHcun7oo33tPO/8dn50EJHwDxH8FZB4t3y9fZ9eJQJfrdPnZdlY8f7APtvB9MpstGPifFThx+gzeujEOY1jxRURE5DDnDxcVfDnjCRER1RFTiRaASBB21Non7Kym5UL6hKkgrJ/2MdD2FrgoLDHhdF5xFRVkhedUk11sUHY2KRrycbcGZEb4errBz3rZww2+ap/rOZfldr/Sj7JPLnu6udRtD5uiPC3ctAZdyQmAxVz5GJnmKtVcEnTJlEBpIO+IzGatuk2qwHYvBEyF2n5ZDKLDaKDbRG0FytquOkm1wvMH+2AL36eNh0/h+g/Xqd+tfz8zDF7udbhQBhEREel6/mCTqzoSkQMzumrT2WTr80CFPmGlFWHy8eR+IG2Xtm36VPs8/yblfcKa9SvrE6YnD1cjwgNk86x1UGbtR2Zt2p9bUIK8whLkVtwKKl+XaZjy32W9fqmkcOzsyrKzK8+0UO2syjPrdWMJAk8mwCdpLYxH18CQuBkwF1d+kMYty4OumAGAbyicgvxsthikbTLFU6b5SgiWukO7LFujGK0XWNwtgH+E3iMmcmrWpvYjO4cz9CIiInIwrPgiItuTmw4cP6tPmPmsoMcjAGjau0KfsG416hNmr+RXdWGJWVWNWQOynAuFZXJ70VnHye1FJSpAqy03lKCL4SD6uuxCP5ed6O6yHx6GykFXiiEE29y6YK9XVxzy7YZin4gqp21aw7SyYK00cJPL0svNIcl/etIWLQDb/jNQlKPtN7gArYdrVWDysSFWqnQSPH+wD3p/nwqKTej58h/IKSzB9/f2QZ8WQQ0+BiIiIqodTnUkIsciU+ikmkhWjZQg7MTfQFHuuX3CJPyq2CfMUafRXSL5tZ9fZKockEmQdnZYVlAI/9O7EJX5N5rnxKNVwXZ4Wgoq3VeqJRDrzB2w1txRfTxukYquS5s+6eHqcm4VWun1SlVoVfRAqxSsebiqnmo2+zO9a74WgsnPtJVPKBB3sxaCyaIRdEl4/mAf9P4+LdiahEe+24Imjbzw5xND7KePIhERkRPLZvBFRI7fJ2xHhemR64DcKpZglOmQ1iCsWV8gIJor612oL1XazvIeXdKDrTCr8jHeQUDMZaVTFwfC1LgV8opNWkVZgRaeWS9XqkIrKt+nQrbSgK3icQXFZ/UDqwPe7sYLhmUStBkNBvVi10U+GrRFCKT/mdGACvsNkGI0bb+h9BjtWOvt1s9Vl9VHqGPV55Ret96fth/wyjqIRnt/gP++n+F6JqNs7IVR/9/enUBHWd57HP/NBANhSVjCHpDFAhKRVSOLC0rhAEW91lvFyOZ+3cvVHtQWxI32aC29VqjneIq2RwrUKvQi4lUULVe4LAFEFBBFCIRVLWExJEzmnud5mYQhwQ0yz8z7fj/nvGcW3sk8eWcy8+f//p//c4EOn5Ov0k4/USi9rjfG439uNc97Wnu2+QDxQ2pw/TqNnbFcizfu1d2XnqXxgzsn/PkBAMD3R+ILQLCYj7GvtsT3Cdu3qep+ma0rE2Hm0jTQD6cF+7iZ42SSXGb7fIn09ZdVp5S261/Zo8ses5qZilgWKY+buhlLkMUly46YSrWyY/tEdLDkuOtHyuy/HygpU1kk6b/aqqilo7osvFo/S1usS8JrlBbyfofiaIbmRfprVmSg1kerrpJ8IpMAiyXVKpNjXsKtMglnerzFEmmqklSrTPRVn6w7eXLv2GMr9tdJk3UVicNQSLde3FFNG9Q+7ceU+CE1uHyd9hwo0QVPLFJ5VHrnvkvUPrteQp8fAAD8MDS3BxAs5n/PjTt4m5kmZhzaF58IMysMFu+QPvy7t8WSOmZKZCwZ1tr0CcuQ7xOEsYous+rgiZVyZ9TzquNiia6W3ROWHDS9vRrWTbfbqTILCdgkma1CKzsuYebdd/x1s295NGoL3iLm0l43l95tc34ocuy2d39UJq9Web/3WO/+E/Y7tiBBlZ9r7z9x/zO0tLyv3o/2VXZ0ny6PvqurQ++oTWiPRtV6y27ry8+0CbB5kX4qVv1qf3fzsyNKrcTfyLy2NZL4Ar7NP9YU2b/BXm0bkvQCAMCnSHwB8Kd62dLZP/E2o/Sw1ycsNjWycLk3jW/zm95mhM/wVps0iR/bJywv9fuE7d/uJbo+P5bs2l8Y/++16njJv2NTF23yzwfN1c2Km2ZrXO/Uk2ju5HvTTz9/Tyr4i6If/0O52qpHwy/okYxZinT+icq6X6/SnH6KKlSZnIuekHSLJeeOJdmqS9Z5STxVSfp9e7Ku+uRe1bHoGxOKDTNS/z2H1PTyKm81x5/2znE9FAAAUEOY6ggguH3CTD8rUw1melmZy4O7qu7XtMtx0yP7Sg3bJnefsIN7Kqu5zOWXn8X/e7iWlHOeV81lkl3muo9Xw/SVw19K6/4mrXrRe+/GNGov9bxe6pEvZbZ0OcKkQ/yQGly9Th8VFWvYf/1T6WlhrXhokLLqkoAFACBV0OMLAH7QNMDPT+gTtrHqfg1axfcJa57rtk+YSYaY3lyxRNfeDfH/Hgp7VWyxRJcZczrTeVL+vVpUYKvAtO5lqfRA5Wv9oyFSr1HSjwb7onLvVBE/pAZXr9Nj8z/S80u2aFi3FpqW3zthzwsAAE4diS8AOB1Mn7DC/6tMhBWtlsqPxu9TO/OEPmG9a7ZPWEmxV6FmE13vSrs+NJmQ+H1adPOmLZpEl5m2WSer5sYDt0oPSR/Nkwr+7L1PY+o3l7qPlHqNlpp0VFARP6QGF6/T0Ui5LpjytvYdPKLnR/fRoK7NE/K8AADg9CDxBQA1oaJP2LLKPmGxapsY2yesR+XUSNMnrF6TU3jOQ97zxSq6itZI0UjV6Zixiq52A1K/Lxl+mL2bpNV/kdb+VTq0t/L+M/t7CbCzL5fS6ypIiB9Sg4vX6Z0NezTuhRVqUi9dyx68zC6uAQAAUgeJLwBIhPKItPtYn7Bt70tbl1bfJyy7c/z0yEbtTt4nrKxE2r6iMtG1faVUXha/j1m9siLRdaHUgEoFHCdSJm1a6FWBbX5LipZXrmLa7WovCWaSswFA/JAaXLxOd84s0PwPdmpc/3aaNCI3Ic8JAABOHxJfAOCC+Tj919bKijCTCKu2T1jL+ESYSXaZlftMostUkR0tid8/M8dLctntQimL1cfwHe3fIa2ZKa3+s/SvbfHTYXuN8RJhGY3kV8QPqSHRr9P+r8t03uNvqfRouebfNUDntGY6OAAAqYbEFwAki0NfVNMn7IQKrhOZ/kyxii6T6DKr9iXzSpJIfuXlXnLVVIF9/N9SpNS7v1YdbwqkqQIz02R99j4jfkgNiX6dZi3fpgmvrFOn5vX1xr0XKeSz9z0AAEFQ/D3ih1oJGxUABJHp79VlmLcZZV8f6xN2LBFmKrzMqpAVia6LpOxOvktAwLFwWOpwibeZlUA/mOMlwfasl9bN8TaTYDUrQna/Tsps6XrEQI35e8F2e/nTXjkkvQAACAASXwCQSGbFR1NZYzbDFN3yHy8kkln84ILbpLxbpaICLwG27u/SV1ukRY9Ibz8u/WiwVwVmLtMIFeAfW784pBWff6VwSLqyZ2vXwwEAAAlANAsALpH0gsv3Xuve3jbkCWn9XG9VSFONuOl1bzPTbntcJ/UcJTXp6HrEwCl7dfUOe9n/rGw1z6zjejgAACABWLsZAICgS68n9cyXblgo3bFC6neXVDdbOrhbWvI76Zle0ozh0tpZUulh16MFfhDT1vaVgh0V0xwBAEAwkPgCAACVmnaSBj8mjf9Y+tlfvOmOobC0dYn06q3Sb7tIr/2nVLTG9UiB72XV1q+07cvDqpeepsG5zV0PBwAAJAhTHQEAQFW10qWul3vb/u3SmpneVMh/bZNWPO9tLc71eoF1u1rKaOR6xMB3amo/rFtL1U0nBAYAICio+AIAAN8sK0e6+BfS3WulUXOlc34qpaVLuz6QFtznVYG9cou05Z/egg1Akikpi2j+Bzvt9auY5ggAQKBwugsAAHw34bDUcaC3Hf5S+mCOtyrknvXSB7O9rXEHqef1Uo98qUEL1yMGrLc+3q0DJUfVumGG8to3dj0cAACQQFR8AQCA769uY+mC26T/+F/p5rel3mOl9AbSl59Jix6Rnu4q/XWktGGBFDnqerQIuFhT+3/r2VrhMKvpAgAQJFR8AQCAHy4Uklr39rYhT0jr53pVYIXLpI0LvK1+C6nHSKnnKKlJR9cjRsDsPXBE727aa69f1au16+EAAIAEo+ILAACcHun1pJ750o1vSHeskPrdJdXNlg7ukpb8TnqmlzRjuLR2tlT2tevRBtKzzz6rdu3aqU6dOsrLy9Py5cu/cf+pU6eqc+fOysjIUJs2bfTzn/9cJSUlFf8eiUT0q1/9Su3bt7f7dOzYUY8++qiiSdTrbd6aHYqUR9WzbUN1aFrf9XAAAECCUfEFAABOv6adpMGPSZdOlDYt9KrAPl0kbV3ibQvul879d29VyJbdXY82EGbPnq3x48frj3/8o016maTWkCFDtHHjRjVr1qzK/jNnztSECRP0pz/9Sf369dOmTZs0duxYhUIhPf3003af3/zmN5o+fbpefPFF5ebmauXKlRo3bpyysrJ09913K5mmOdLUHgCAYKLiCwAA1Jxa6VLXy6XrX5buXScNfEjKaisd2S+teF567iJp0/+4HmUgmGTVzTffbBNTXbt2tQmwunXr2sRWdd5//331799f1113na0SGzx4sEaOHBlXJWb2ueKKKzR8+HC7z9VXX233+7ZKskT5eGexPtpZrPS0sEac29L1cAAAgAMkvgAAQGJk5UgX/0K6Z600aq6Ue5XUoKXU4WLXI/O90tJSrVq1SoMGDaq4LxwO29tLly6t9jGmyss8JpbE+uyzz7RgwQINGzYsbp9FixbZajBj7dq1WrJkiYYOHXrSsRw5ckTFxcVxW035uiyi89o10mVnN1PDuuk19jwAACB5MdURAAAkVjgsdRzobUdLvaow1Kh9+/bZflzNmzePu9/c3rBhQ7WPMZVe5nEDBgywPbuOHj2q2267TQ8++GDFPmYqpElcdenSRWlpafY5Hn/8ceXn5590LFOmTNHkyZOVCL3aNtLfbuun0qPlCXk+AACQfKj4AgAA7pD0SlqLFy/WE088oWnTpqmgoECvvPKKXnvtNdu8PmbOnDl66aWXbD8ws4/p9fXUU0/Zy5N54IEHtH///oqtsLCwxn+X9FqEvAAABBUVXwAAAD6XnZ1tK7J2794dd7+53aJFi2ofY1ZrHDVqlG666SZ7u1u3bjp06JBuueUWPfTQQ3aq5P3332+rvq699tqKfbZu3WqrusaMGVPtz61du7bdAAAAEoHTXwAAAD6Xnp6u3r17235cMeXl5fZ23759q33M4cOHbXLreCZ5Zpipj9+0j/nZAAAAyYCKLwAAgAAYP368rcLq06ePzj//fE2dOtVWcJlVHo3Ro0erdevWtlrLGDFihF0JsmfPnsrLy9PmzZttFZi5P5YAM9dNT6+2bdsqNzdXq1evto+54YYbnP6uAAAAMSS+AAAAAuCaa67R3r17NXHiRO3atUs9evTQwoULKxreb9u2La5665e//KVCoZC93LFjh5o2bVqR6Ip55plnbDLs9ttv1549e9SqVSvdeuut9jkAAACSQSgaq1VPYma1oKysLNsANTMz0/VwAABACiB+SA28TgAAoCbjB3p8AQAAAAAAwJdIfAEAAAAAAMCXSHwBAAAAAADAl0h8AQAAAAAAwJdIfAEAAAAAAMCXSHwBAAAAAADAl0h8AQAAAAAAwJdIfAEAAAAAAMCXaikFRKNRe1lcXOx6KAAAIEXE4oZYHIHkRJwHAABqMs5LicTXgQMH7GWbNm1cDwUAAKQYE0dkZWW5HgZOgjgPAADUZJwXiqbAadDy8nIVFRWpQYMGCoVCNZIpNMFWYWGhMjMzT/vPxzfj+LvF8XeL4+8Wx9/fx9+EOCYYatWqlcJhujskK+I8f+P4u8Xxd4vj7xbH363iJIrzUqLiy/wSOTk5Nf485sXgD8Idjr9bHH+3OP5ucfz9e/yp9Ep+xHnBwPF3i+PvFsffLY6/W8kQ53H6EwAAAAAAAL5E4gsAAAAAAAC+ROJLUu3atTVp0iR7icTj+LvF8XeL4+8Wx98tjj8SgfeZWxx/tzj+bnH83eL4u1U7iY5/SjS3BwAAAAAAAL4vKr4AAAAAAADgSyS+AAAAAAAA4EskvgAAAAAAAOBLJL4AAAAAAADgS4FPfD377LNq166d6tSpo7y8PC1fvtz1kALjvffe04gRI9SqVSuFQiHNnTvX9ZACZcqUKTrvvPPUoEEDNWvWTFdeeaU2btzoeliBMX36dJ177rnKzMy0W9++ffX666+7HlYg/frXv7afQffee6/roQTGww8/bI/58VuXLl1cDws+RJznDnGeW8R5bhHnJQ/ivMRLxjgv0Imv2bNna/z48XaJzYKCAnXv3l1DhgzRnj17XA8tEA4dOmSPuQlKkXjvvvuu7rjjDi1btkxvvvmmysrKNHjwYPu6oObl5OTYL+JVq1Zp5cqVuvTSS3XFFVdo/fr1rocWKCtWrNBzzz1ng1MkVm5urnbu3FmxLVmyxPWQ4DPEeW4R57lFnOcWcV5yIM5zJzfJ4rxQNBqNKqDMmT9zJuQPf/iDvV1eXq42bdrorrvu0oQJE1wPL1BMFvjVV1+1Z6Pgxt69e+0ZQRMoXXTRRa6HE0iNGzfWk08+qRtvvNH1UALh4MGD6tWrl6ZNm6bHHntMPXr00NSpU10PKzBnAk31x5o1a1wPBT5GnJc8iPPcI85zjzgvsYjz3Hk4CeO8wFZ8lZaW2gz8oEGDKu4Lh8P29tKlS52ODXBh//79FV/KSKxIJKJZs2bZs7CmFB6JYc6EDx8+PO57AInzySef2ClQHTp0UH5+vrZt2+Z6SPAR4jwgHnGeO8R5bhDnufVJksV5tRRQ+/btsx9CzZs3j7vf3N6wYYOzcQEumLPgZt57//79dc4557geTmCsW7fOBkAlJSWqX7++PRvetWtX18MKBBOAmqlPpgQebipxXnjhBXXu3NmWv0+ePFkXXnihPvzwQ9uPBjhVxHlAJeI8N4jz3CHOcysvCeO8wCa+AMSfETEfRK7nXgeN+TIwJcDmLOzLL7+sMWPG2CkIBEU1q7CwUPfcc4/teWIaXiPxhg4dWnHd9N0wAdKZZ56pOXPmMAUEAE4z4jw3iPPcIM5zb2gSxnmBTXxlZ2crLS1Nu3fvjrvf3G7RooWzcQGJduedd2r+/Pl29SXTiBOJk56errPOOste7927tz0r9fvf/9424UTNMdOfTHNr0/chxlSGmL8B0wvoyJEj9vsBidOwYUN16tRJmzdvdj0U+ARxHuAhznOHOM8N4rzk0zAJ4rxwkD+IzAfQokWL4sqAzW3mXiMIzLoWJhgyZddvv/222rdv73pIgWc+g8yXMWrWZZddZqcfmLOwsa1Pnz62/4C5TjDkpgHtp59+qpYtW7oeCnyCOA9BR5yXfIjzEoM4L/kcTII4L7AVX4ZZ4tqUnJo/hPPPP9+u8mCaDo4bN8710ALzB3B81nfLli32w8g03Wzbtq3TsQWl7H3mzJmaN2+enWu9a9cue39WVpYyMjJcD8/3HnjgAVsGbN7rBw4csK/F4sWL9cYbb7gemu+Z9/uJPU7q1aunJk2a0PskQe677z6NGDHClr0XFRVp0qRJNhAdOXKk66HBR4jz3CLOc4s4zy3iPHeI89y7LwnjvEAnvq655hq7tO/EiRPtl4FZ4nThwoVVGqGiZqxcuVIDBw6MC1ANE6SaZnioWdOnT7eXl1xySdz9M2bM0NixYx2NKjhMCfbo0aNtw0cThJr57yYY+vGPf+x6aECN2759uw1+vvjiCzVt2lQDBgzQsmXL7HXgdCHOc4s4zy3iPLeI8xBk25MwzgtFTR0sAAAAAAAA4DOB7fEFAAAAAAAAfyPxBQAAAAAAAF8i8QUAAAAAAABfIvEFAAAAAAAAXyLxBQAAAAAAAF8i8QUAAAAAAABfIvEFAAAAAAAAXyLxBQAAAAAAAF8i8QUAAAAAAABfIvEFAAAAAAAAXyLxBQAAAAAAAF8i8QUAAAAAAAD50f8DPdpQHPUYsXYAAAAASUVORK5CYII=\",\n \"text/plain\": [\n \"
\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"from IPython.display import clear_output\\n\",\n \"import matplotlib.pyplot as plt\\n\",\n \"\\n\",\n \"metrics_history = {\\n\",\n \" 'train_loss': [],\\n\",\n \" 'train_accuracy': [],\\n\",\n \" 'test_loss': [],\\n\",\n \" 'test_accuracy': [],\\n\",\n \"}\\n\",\n \"\\n\",\n \"rngs = nnx.Rngs(0)\\n\",\n \"train_model = nnx.view(model, deterministic=False, use_running_average=False)\\n\",\n \"eval_model = nnx.view(model, deterministic=True, use_running_average=True)\\n\",\n \"\\n\",\n \"for step, batch in enumerate(train_ds.as_numpy_iterator()):\\n\",\n \" # Run the optimization for one step and make a stateful update to the following:\\n\",\n \" # - The train state's model parameters\\n\",\n \" # - The optimizer state\\n\",\n \" # - The training loss and accuracy batch metrics\\n\",\n \" train_step(train_model, optimizer, metrics, rngs, batch)\\n\",\n \"\\n\",\n \" if step > 0 and (step % eval_every == 0 or step == train_steps - 1): # One training epoch has passed.\\n\",\n \" # Log the training metrics.\\n\",\n \" for metric, value in metrics.compute().items(): # Compute the metrics.\\n\",\n \" metrics_history[f'train_{metric}'].append(value) # Record the metrics.\\n\",\n \" metrics.reset() # Reset the metrics for the test set.\\n\",\n \"\\n\",\n \" # Compute the metrics on the test set after each training epoch.\\n\",\n \" for test_batch in test_ds.as_numpy_iterator():\\n\",\n \" eval_step(eval_model, metrics, test_batch)\\n\",\n \"\\n\",\n \" # Log the test metrics.\\n\",\n \" for metric, value in metrics.compute().items():\\n\",\n \" metrics_history[f'test_{metric}'].append(value)\\n\",\n \" metrics.reset() # Reset the metrics for the next training epoch.\\n\",\n \"\\n\",\n \" clear_output(wait=True)\\n\",\n \" # Plot loss and accuracy in subplots\\n\",\n \" fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))\\n\",\n \" ax1.set_title('Loss')\\n\",\n \" ax2.set_title('Accuracy')\\n\",\n \" for dataset in ('train', 'test'):\\n\",\n \" ax1.plot(metrics_history[f'{dataset}_loss'], label=f'{dataset}_loss')\\n\",\n \" ax2.plot(metrics_history[f'{dataset}_accuracy'], label=f'{dataset}_accuracy')\\n\",\n \" ax1.legend()\\n\",\n \" ax2.legend()\\n\",\n \" plt.show()\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"25\",\n \"metadata\": {},\n \"source\": [\n \"## 7. Perform inference on the test set\\n\",\n \"\\n\",\n \"Create a `jit`-compiled model inference function (with `nnx.jit`) - `pred_step` - to generate predictions on the test set using the learned model parameters. Since we already have `eval_model` (an `nnx.view` with `deterministic=True` and `use_running_average=True`), we can use it directly for inference. This will enable you to visualize test images alongside their predicted labels for a qualitative assessment of model performance.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"26\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"@nnx.jit\\n\",\n \"def pred_step(model: CNN, batch):\\n\",\n \" logits = model(batch['image'], None)\\n\",\n \" return logits.argmax(axis=1)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"1d6cb81f\",\n \"metadata\": {},\n \"source\": [\n \"We reuse the `eval_model` view created earlier so that `Dropout` is disabled and `BatchNorm` uses stored running stats during inference.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"27\",\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"image/png\": \"iVBORw0KGgoAAAANSUhEUgAAA7QAAAPGCAYAAADTLdZkAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjUsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvWftoOwAAAAlwSFlzAAAPYQAAD2EBqD+naQAAjY5JREFUeJzt/QeYFFX6P27XINnAoqJgxIgBFUXBhCCirjlnFDNiXgOurgqYMfvVNbuIGDDnrIA5KwaUNbvGFSUYQIL0e1X/XvijddrtdmboOd33fV2sy4dD9ZnhnJ5+qqqfrsnlcrkEAAAAItOo3BMAAACAP0NBCwAAQJQUtAAAAERJQQsAAECUFLQAAABESUELAABAlBS0AAAARElBCwAAQJQUtAAAAESpqgra66+/PqmpqUk+/fTTkv5ejx49ko4dO9bpXNq3b5/su+++dXpM+CPWP9XOHqCaWf9UO3ugclVVQVuJXnnlleTwww9PVl111WTeeedNllpqqWTXXXdN3n///XJPDeaKqVOnJieccEKy2GKLJS1atEi6du2aPP744+WeFpTFmWeemX/BVtcvvqCh+uCDD5Ldd989WWKJJZKWLVsmK620UnLaaaclkydPLvfUoN6lRXFNTU3BX19++WVSDRqXewLUzuDBg5Pnnnsu2WWXXZLVV189+eabb5LLLrssWWuttZIXX3zRixqq4sn8jjvuSI4++uhkhRVWyJ+B3XLLLZORI0cmG264YbmnB3PNF198kZx11ln5k5tQDT7//POkS5cuSatWrfIn9xdccMHkhRdeSAYMGJC89tpryb333lvuKUK96tu3b9KrV6/fZLlcLjnkkEPyV4EXX3zxpBooaCN3zDHHJDfffHPStGnT2dluu+2WrLbaask555yT3HjjjWWdH9Snl19+ORk+fHhy3nnnJccdd1w+22efffIncvr37588//zz5Z4izDXpHlh33XWTX3/9Nfnuu+/KPR2od8OGDUsmTpyYPPvss/k71VIHH3xwMnPmzOSGG25IJkyYkLRu3brc04R6s9566+V/zSndD+kdCnvttVdSLar6luP0zN1WW22Vv1WxWbNmyXLLLZecfvrp+RcDIenZvvXXXz9/W+MyyyyTXHnllcHbH9Mzg8svv3z+mEsuuWT+hXWa14d0PnMWs6n0KlX6xP7ee+/Vy2NSGSph/adXZueZZ578C5hZmjdvnhxwwAH5s/Tp2Xuo5D0wy9NPP53fDxdffHG9Pg6VoxLW/w8//JD/76KLLvqbvF27dkmjRo0yr4+g0vZAyM0335y/3XjPPfdMqkVVX6FNb02cb7758lc50/+OGDEiOfXUU/NPkOkVnzmlZ/nS2xjT96fuscceyW233Zb069cv/2S5//7758ekZwS33Xbb/JmR9AX2yiuvnLz99tvJRRddlH9P6z333FNwLunfHT9+fFHzTm+tadKkScE/T281+O9//zv7bCVU6vp/4403khVXXDFZYIEFfjMmvQUtNXr06PwPE6jUPZBKX3wdccQRyYEHHpi/OweqZf2nzXrSt16lJzEHDRqULLTQQvk7c6644orkyCOPdPs9Fb8Hfm/69On5uaWFd3rLcdXIVZEhQ4bk0i/5k08+yf9+8uTJmTF9+/bNtWzZMvfLL7/Mzrp3757/exdccMHsbOrUqblOnTrlFllkkdy0adPy2bBhw3KNGjXKPfPMM7855pVXXpn/+88999zsbOmll8716dNn9u/TOaVjivk1cuTIP/w603mk46677ro/9X2iMlXi+l911VVzPXv2zHwdY8aMyY9NHxsqeQ+kLrvsslyrVq1y33777ez5pnsDqmH9n3766bkWLVr8Zsw//vGPOvmeUVkqdQ/M6f7778+Pufzyy3PVpKqv0Ka3DMzy448/5m8H6NatW3LVVVclY8eOTdZYY43Zf964ceP8G69nSc/IpL9Pz86ktyCk71u6/fbb82dj0g57c75/qWfPnvn/pk1q0jMmIW3bti26M+uc8/q9dN6HHXZY/n76Pn36FHU8qlMlrP8pU6bkb+n5vfS241l/DpW8B77//vv8FYVTTjkladOmTYnfAapZJaz/VHoVaqONNkp22mmn/BXaBx98MN8cLT1m2igKKn0P/P524/TqbXoluZpUdUE7ZsyY5OSTT87fYjDrfRizTJo06Te/T++v//2tK+mtjqn086zShZy2jk/ft1roRcW3335bcC7pC/DfdykrVdrhOH0vQHorwqz3FkIlr//0h1HofSm//PLL7D+HSt4D6fzTzq7pLcdQbes/bQqY3tqZ3s6ZfmxPascdd8zfvpl+nFt6a2ha5EKl7oE5/fTTT/n3BW+++eZVt+6rtqBNu+J17949/9679PPK0jeCp4vp9ddfzz8Jpk+GpUr/Tvr+pQsvvDD453/0Xr70PVDjxo0r6nHSFy+/b3SQbrwtttgi/3U988wz+Y0Hlb7+08Yfoc9Y+/rrr/P/tQ+o5D2Qvni6+uqr842gvvrqq9+c0EnfR5W+yEq/vnQ8VNr6T11++eXJmmuuObuYnSV9H2P6/si0z0JtiwQqU6XsgTndc889VdfdOKn2gnbUqFH5W7Xuuuuu/K0qs3zyySfB8emLhZ9//vk3Z2fSM4KpWW+6TjfDm2++mWyyySb57mKlSLuxph3TipHespA2Qpjzxcs222yTn88TTzyRrLLKKiU9NtWnUtZ/p06d8r9Pz6zO2RjqpZdemv3nUKl7ID2Zk76ASpvfpL9+Lz3eUUcdpfMxFbn+U2kDzNDH8qQndFIzZswoaR5Uj0rZA3O66aab8s2t0hM61aZqC9pZt+OmHYFnmTZtWv5sX0j6pJjeU592Qps1Nv19eltB586d81l6v/pDDz2UXHPNNb/5GJFZ7+VLX3gU6rj3Z++dT8/opJ87m35ESXqbwe8/iwoqef3vvPPOyfnnn5+/SjXrc2jTW5CHDBmSdO3aVYdjKnoPpJ+3fPfdd2f+PL2FLn0/2CWXXJJ/gQWVuP5n3fL52GOP5QuLWbd/pm655Zb8x/asvvrqRR2T6lMpe2CWcePG5S9qpbfZt2zZMqk2VVvQpm/KTs/qpY2T0jPb6ZmU9AO651zYc0pvXUxbw6e3cKVPmrfeemv+I0HSF9KzWmfvvffe+VbZhxxySP7syQYbbJAvONM3lqf5o48+mqy99tp1eu/8sccem9x33335K7Rpu+8bb7zxN3/eu3fvko9J5auU9Z8Wrbvsskty4okn5t+bkn7u29ChQ/PzvO6660o+HtWjEvbAwgsvnGy//faZfNYV2dCfQaWs/9Txxx+fPPzww/lGPmkDqPR9gw888EA+Sz/GyttOqPQ9MMutt96aL7qr8XbjvFwVt+tO22evu+66+Xbviy22WK5///65Rx99NNMSe9ZHILz66qu59dZbL9e8efN8u+30oxJ+L23dPXjw4Pz4Zs2a5Vq3bp3r3LlzbtCgQblJkyYVbNf9Z81qJV7oF1Ty+k9NmTIld9xxx+Xatm2bf8x11lkn98gjj9TJsakslboHfs/H9lBN6/+ll17KbbHFFvmfAU2aNMmtuOKKuTPPPDM3ffr0Ojk+laNS90Aq/TrSjxCaMWNGrhrVpP9T7qIaAAAAStWo5L8BAAAADYCCFgAAgCgpaAEAAIiSghYAAIAoKWgBAACIkoIWAACAKCloAQAAiFLjYgfW1NTU70zgD5T745Ktf6p5/afsAap5D1j/VPP6T9kDNOQ94AotAAAAUVLQAgAAECUFLQAAAFFS0AIAABAlBS0AAABRUtACAAAQJQUtAAAAUVLQAgAAECUFLQAAAFFS0AIAABAlBS0AAABRUtACAAAQJQUtAAAAUVLQAgAAECUFLQAAAFFS0AIAABAlBS0AAABRUtACAAAQJQUtAAAAUVLQAgAAECUFLQAAAFFqXO4JUHvt2rUL5gsuuGAmmzFjRnDsv//97zqfFw3TWmutFcwPOOCAYN6vX79gfu+992ayxx57rJazS5J33303mD/11FO1PjYAAJXFFVoAAACipKAFAAAgSgpaAAAAoqSgBQAAIEoKWgAAAKJUk8vlckUNrKmp/9nwh5ZffvlgPnLkyKK7H0+fPj049oorrgjmxxxzTNIQFLlM602s679Tp06Z7KGHHgqOXXTRRZOGYMKECcH86aefDuYXXnhhMP/iiy8y2aeffprEqNzrP+Y9QGUo9x6w/qnm9Z+yB2jIe8AVWgAAAKKkoAUAACBKCloAAACipKAFAAAgSppC1cJGG22UyW6//fbg2ELf5iFDhhR13FTHjh2D+XzzzVfSY4YUahb13HPPZbJevXol1dYQoaGv/1Dzp9Rdd92VyZZeeumkISv0vS51Dbz77ruZ7Oabbw6OPf/880vaF9W2/hvSHij0PDhixIhgftVVV2WyU045JakkvXv3Dua77LJLJtt///2DY7///vukISv3Hmgo67+aFWpcuOeee5b0czHk0ksvDeavvvpq0hCUe/2n7AHKSVMoAAAAKpKCFgAAgCgpaAEAAIiSghYAAIAoKWgBAACIki7HRfjLX/4SzF977bVM1r59+3rrUPfVV18F82OOOaboYwwYMCCYr7zyysH8sccey2RbbrllUm0d/hr6+n/zzTdL6ghbDV2OS1Gow+XRRx+dNATlXv8NaQ9ceOGFwfxvf/tbMH/rrbcy2XbbbRcc++mnnyYxGjNmTDBfZZVVMtkdd9xRdEfkhqTce6ChrP9KM88882Sy/v37l/Rap9C/zYILLlj0PJ588slgvummmyYNQbnXf0PaAyuttFIwv/jii4P54osvXnT36kLHKPQai7lHl2MAAAAqkoIWAACAKCloAQAAiJKCFgAAgCgpaAEAAIhS43JPoCHp0qVLMD/jjDOC+dJLL13rxxwyZEgm+/jjj4sem/rmm2+KfrzTTz+9hNklyUcffVTSeKrH2LFji+4eO3Xq1GC+xx57ZLJu3bqV1G18/fXXT2rr0EMPLbqr47HHHhscO2PGjFrPg//9b77EEkvU+hjNmjVLYlSoi37Lli2LPsYmm2xShzOC4qyxxhrBfODAgUX/HBk6dGgwHzRoUDD//PPPM9kNN9wQHNuzZ8+kttq2bVvr12j8b4suumgw33zzzYs+RqFPgOjdu3cwf//994P5s88+m9TWQw89lMmmTJkSHLvjjjsG81tuuaXW8yjU5f+zzz5LYuAKLQAAAFFS0AIAABAlBS0AAABRUtACAAAQpZpcLpcramCgOUqlCTUnSJ1yyilFH+O5554ruvlN6ssvv0zmpv/+97/BfOGFFy66IdaAAQOSua3IZVpvGsr632qrrYL5TTfdFMznn3/+Wj/muHHjgvkGG2wwV5uILbjggsF84403DuZXX3110Y2lSrHccsuV1FChEtZ/ufbARhttlMmeeuqpko4Reg4r5Tm9ITnzzDOD+UknnVT0MSZMmFDS/mooyr0HGsrPgIZu3XXXDebXX3990c+nhxxySEmNMWfOnFn0/BZffPFg/vDDDwfz/fbbr+jXQG+++WYwr4vnm3Kv/4a0Bwo19Su0Pgq9/ua3fvzxxyTk5ZdfzmS9evVK5rb/tQdcoQUAACBKCloAAACipKAFAAAgSgpaAAAAoqSgBQAAIEqNyz2BhmTMmDHB/Pbbbw/m77zzTlEdNcvlwAMPzGQLLLBASd3Dbr311jqfF3/eUkstVW/djAu55ZZbgnl9djQOGT9+fDC/8847g/kKK6xQdJfYUtx///3BfJtttpnr3Y8rXahTdTVYY401gvmhhx5a62N/9tlntT4GFHLssccG8w4dOgTz7bbbLpPdd999SX35+eefg/liiy0WzF955ZVMduqppwbHXnjhhbWcHcWYOnVqMN9///2D+WmnnZbJNt988+DYH374IZjvs88+wXzJJZdM5qZ27dqV1Ol7vvnmK/rYhV5HvvHGG0kMXKEFAAAgSgpaAAAAoqSgBQAAIEoKWgAAAKKkoAUAACBKNblC7W1/P7Cmpv5nQ50aMWJEJttoo42CY5988slgvtVWW2WyGTNmJHNbkcu03jSU9T9lypRg3rRp03p7zPfffz+Yr7zyyklD1qxZs0y27bbbBscOHz681o8X6oaZWnfddaNf/+XaAxMnTsxkrVq1KukYoc7zp5xyStKQdenSJZi/9NJLtT52t27dgvmzzz6bNGTl3gMN5WdAQ9K+ffuiu99fc801wbxfv3719m8d+lSASy+9NDh26623LrrT+t/+9rfg2F9++SWp1PWfsgfKb8UVVyzp9dhdd92VyRo1Cl/L/PXXX4P5AQcckMmGDh2azG3/aw+4QgsAAECUFLQAAABESUELAABAlBS0AAAARElBCwAAQJQal3sC1F7Xrl2D+SqrrFL0MQp1ICxHR2NK69xb3x0Ql1566WDeu3fvTHbjjTcmDcXUqVOL7ub9/PPPB/P111+/6Mdr3rx5CbNjToMGDQrm8803X9HHKNRd9corr/zT8wIKa9u2bdGdcJ966qmif3Y1bty46I7IqZ49ewbzv/71r5nsww8/DI7deeedg/ndd98dzKEcPvjgg2B+zjnnBPNQR+NCrxePP/74YF6OjsZ/hiu0AAAARElBCwAAQJQUtAAAAERJQQsAAECUNIWKSMeOHYP5gw8+GMz/8pe/ZLKnn346OPaxxx6r5eyotkZUiy++eBKb8ePHB/OJEyfO9bnwvxuPzTPPPEUfo2XLlsF8iSWWyGRffvllCbMDQjp16lT02O+++y6YH3LIIZnssMMOC45dddVVg/mECROC+eDBgzPZpZdeGhz7/fffB3NoSHr06BHMd9hhh6KPceGFFwbziy66KImZK7QAAABESUELAABAlBS0AAAARElBCwAAQJQUtAAAAESp4rsch7ribb/99sGx2267bTBfe+21i368Ro3C5whmzpwZzF955ZWistQee+wRzBdaaKGiO7cOHDgwOPaHH34I5jQszz77bDDfcMMN5/pcampqkkpx+OGHB/NPPvmk6K979dVXD+b9+vUL5ldccUVJc6xk559/ftHPya1btw6ObdeuXTC/5ZZbMtmHH36YNGStWrWqt2Ofdtppwfyvf/1rMJ82bVq9zYW4FXrtEfLAAw8E88aNsy9D33jjjeDY/fbbL5gPHz48mE+dOrXo+UFDcuCBBwbza665ptaf7HDmmWcmlcgVWgAAAKKkoAUAACBKCloAAACipKAFAAAgSgpaAAAAohRdl+Odd945mB966KHBvHv37pksl8uV9JiljC/UzbjQMUIdlEvpqvxHjxn6njz99NMlHZuGJdSxNbXBBhvU+tiFumt//fXXwfy6665LKsWyyy5b9L6tz+ePavXOO+8E8/XXXz+T3XPPPcGxHTp0CObLLLNMUVm12HjjjYP5lVdeGcz333//ep4RDd1mm20WzE844YSij1GoW/Z2222XyR555JESZgdxW2KJJTLZUUcdVSfH7tu3byabMGFCUolcoQUAACBKCloAAACipKAFAAAgSgpaAAAAotSgm0LtsMMOmeyGG24Ijm3atGkwHzduXNFNWoYMGRLMf/nll2A+fPjwot9sfdpppwXzgw46KKkvX331Vb0dm8qzyy67BPPPP/88qXTHHHNMrY9R6Pv0xBNP1PrY1Wrs2LGZbPfddw+O7dWrVzA/77zz6nxeMfvpp59KagpF9TjggAOC+dVXXx3MP/zww0z27bffBsd27tw5mDdp0qSkOUKlufPOOzNZx44dSzrGlQWevws1UaxErtACAAAQJQUtAAAAUVLQAgAAECUFLQAAAFFS0AIAABClBtHleOeddw7moY7GhboZF+pQXJ9dhENOPfXUojs217e99tork73wwgvBsdOmTZsLM4LyWn755YP5csstV+tjT5w4sehOoPx5o0ePDuZvvfVWML/ssssy2QUXXBAc+/777wfzq666Kph369Ytkx133HFJbfXo0SOYF/r5V8gll1ySyU444YTg2KlTp5Z0bOKw6KKLZrJzzz03OHbLLbcsqfvxzTffnMmWWmqpkl6jhfbnK6+8Ehz7zTffBHOIwYYbbhjM11hjjaKP8fzzzwfzfv36JdXOFVoAAACipKAFAAAgSgpaAAAAoqSgBQAAIEoKWgAAAKJUk8vlckUNrKmpt0mMGDEimG+00UZFd8o7/PDD661z4+KLLx7M//GPf2Syvn37BscW+jaHuvmdddZZwbH77bdfMN9uu+2Kfsy//e1vwbGXXnpp0pAVuUzrTX2u/1LMN998wfzll18O5h06dCj62DfeeGMw79OnT1IpHY0feOCB4NgVVlih1o8X6viZ2nvvvaNf/w1pD1SDr7/+Opi3bds2mH/33XdF/2wo1CWzoSv3Hmjo679x4/CHVnz//fdFfy09e/YM5q+++motZ5cku+66azAfPnx40Z8Kce+99ybVqtzrP4Y90FCsvfbawfy5554runv9LbfcEhx76KGHlvQpC9W0B1yhBQAAIEoKWgAAAKKkoAUAACBKCloAAACipKAFAAAgSuG2ePVkww03DObdu3cP5v/+978z2UEHHVTrebRv3z6Y9+jRI5ifdNJJwXy55ZbLZNOmTQuOPf/884vu2leoo+D9999fdBfD1F/+8pdMtuOOOwbHDh06NJj/8MMPwZzy+Omnn4L59OnTa33szTbbLJjfcMMNwfyII47IZJMmTUrqS/PmzYP50ksvHczvvvvueulm/MUXXwTzSy65pNbHhj+j0L6LtaMxhTVp0iSYP/3000V/0kOh5/rRo0cn9WWhhRYqemyhrt3Q0DRq1Kjo10yhbsapl156KZNVczfjP8sVWgAAAKKkoAUAACBKCloAAACipKAFAAAgSnO1KdQ//vGPYJ7L5YL58OHDiz728ssvH8w32WSTTHbWWWcFx7Zq1SopxaOPPprJTj311ODYQo2e6sKWW24ZzO+5555M1q1bt+DYf/7zn8F87733ruXsmBtCzcVSHTt2LPoYiyyySDDfa6+9gvkSSyyRyV588cXg2Pvuuy+Yb7vttpmspqam6MdL7bnnnsnctNpqqwVzDdSAurTwwgtnstNPPz04tmvXrsF8/fXXn6vNn5o1a1bSa4lQI83333+/zucF9WHIkCGZbOWVVy7pNcJxxx2XyTR/Kp0rtAAAAERJQQsAAECUFLQAAABESUELAABAlBS0AAAARGmudjnebLPNSupy3L1790z23HPPldTNdb755stkv/zyS3Dsf/7zn5K6qIY6F8+YMSOZ21566aVg/sILL2SybbbZpuhOiKktttgikz388MMlz5H6ddpppwXzH3/8MZOdc845dfKYof0ZylJHHXVUMG/evHkma9QofJ5t5syZSX256667gvkBBxxQ1PcU/qxQZ/xQd1uqz3fffZfJWrZsGRw7fvz4op9jGzcu7aVfp06dgvmSSy6ZyS688MKixxb62TVu3LiS5gf17bDDDgvm++yzT9HH+L//+79g/uyzz/7pefH/cYUWAACAKCloAQAAiJKCFgAAgCgpaAEAAIiSghYAAIAozdUux0OGDAnm++67bzAPdUx99913g2Ovv/76YP7MM89ksi+++CI49sUXX0wqyY477pjJhg4dGhy71157Fd3dUJfjhqdQd+2LLrqoqM7fqRNOOCGYN2nSpJazC3faLKRQ1/NShTplPv7448GxRx55ZDD/4Ycf6mQuUEi7du1q3YX2nnvuqcMZ0ZCFntP/6JMeRowYUW9zCXWef+qpp4Jjt95662A+ZsyYOp8X/FktWrQI5oW6d4c89thjwfy888770/Pif3OFFgAAgCgpaAEAAIiSghYAAIAoKWgBAACIUk2uyA4sNTU1tX6wZs2aBfPllluu6GMUauikeUtx2rRpU1L+0UcfZbKpU6cmc1tdNQr6s+pi/Td0vXv3DuZLLrlkMD/jjDPqZR6NGoXPs73//vslNUl54403MtlLL72UxKjc679a9kA5XHHFFZnskEMOKekYhRoCVVLDnXLvgYa+/tu2bRvMN9lkk1of+7PPPgvmY8eOzWTfffddrR+Phrf+Y9gDdeHMM88M5ieddFIw//DDDzPZ6quvHhw7ZcqUWs6uuuX+xx5whRYAAIAoKWgBAACIkoIWAACAKCloAQAAiJKCFgAAgCjN1S7HEGuHP+ufal7/KXugfuhyHMcesP6p5vVfaXtgoYUWCuaffvppMJ9vvvmC+eabb57JHnvssVrOjhBdjgEAAKhICloAAACipKAFAAAgSgpaAAAAoqSgBQAAIEqNyz0BAKhWoS7Ha621VnDsmWeeGcz/85//1Pm8ACrVNttsU1I340KeeeaZOpoRteUKLQAAAFFS0AIAABAlBS0AAABRUtACAAAQJQUtAAAAUdLlGADK5K233spkXbt2LctcAKpBqd2MCznuuOMy2emnn14nx6Y0rtACAAAQJQUtAAAAUVLQAgAAECUFLQAAAFGqyeVyuaIG1tTU/2yggCKXab2x/qnm9Z+yB6jmPWD9U83rP2UP0JD3gCu0AAAARElBCwAAQJQUtAAAAERJQQsAAECUFLQAAABUdpdjAAAAaEhcoQUAACBKCloAAACipKAFAAAgSgpaAAAAoqSgBQAAIEoKWgAAAKKkoAUAACBKCloAAACipKAFAAAgSlVV0F5//fVJTU1N8umnn5b093r06JF07NixTufSvn37ZN99963TY8Ifsf6pdvYA1cz6p9rZA5WrqgraSvb6668n2267bbLgggsmLVu2zG+8//u//yv3tKDeTZ06NTnhhBOSxRZbLGnRokXStWvX5PHHHy/3tGCu8jOAajRmzJhkl112SZZddtn8ul944YWTjTbaKLn//vvLPTWYK3766adkwIAByV//+tf8839asKeFe7VpXO4JUHuPPfZYss022yRrrrlmcsoppyTzzTdf8tFHHyVffPFFuacG9S49w3nHHXckRx99dLLCCivkn8i33HLLZOTIkcmGG25Y7ulBvfMzgGr12WefJT/++GPSp0+f/EnNyZMnJ3feeWf+5M5VV12VHHzwweWeItSr7777LjnttNOSpZZaKlljjTWSUaNGJdVIQRu5H374Idlnn32SrbbaKv+ivlEjF92pHi+//HIyfPjw5LzzzkuOO+64fJbuh/TqVP/+/ZPnn3++3FOEeuVnANUsPXmZ/prT4YcfnnTu3Dm58MILFbRUvHbt2iVff/110rZt2+TVV19N1llnnaQaVfVPvnvvvTf/IiA9q9esWbNkueWWS04//fTk119/DY5/7bXXkvXXXz9/W+MyyyyTXHnllcHbH9NL/8svv3z+mEsuuWT+hXWa14ebb745+e9//5uceeaZ+RcyP//8czJz5sx6eSwqSyWs//QF/DzzzPObFy3NmzdPDjjggOSFF15IPv/883p5XCpDJewBPwOo5vUfkv5MSB934sSJc+0xiVMl7IH0Mdq2bZtUu6q+QpvempjemnXMMcfk/ztixIjk1FNPzZ/xTq/4zGnChAn5s4C77rprssceeyS33XZb0q9fv6Rp06bJ/vvvnx+TvohIb3N59tln8y+wV1555eTtt99OLrroouT9999P7rnnnoJzSf/u+PHji5p3q1atkiZNmuT//xNPPJEssMACyZdffplsv/32+ceZd955k7333jv/uOmLe6jU9f/GG28kK664Yn4PzKlLly75/44ePTr/wwQqdQ/4GUA1r/9Z0hM5U6ZMSSZNmpTcd999ycMPP5zstttuf+r7QvWopD1Q9XJVZMiQIbn0S/7kk0/yv588eXJmTN++fXMtW7bM/fLLL7Oz7t275//eBRdcMDubOnVqrlOnTrlFFlkkN23atHw2bNiwXKNGjXLPPPPMb4555ZVX5v/+c889Nztbeumlc3369Jn9+3RO6Zhifo0cOXL231t99dXz801/HXHEEbk777wz/9903O67715n3zviV4nrf9VVV8317Nkz83WMGTMmPzZ9bKjkPeBnANW8/uec96w/T+ew884758aPH1/r7xmVpZL3QOqVV17J/3n6dVabqr5Cm94yMEvaVCC9HaBbt275RgJjx47Nv7l6lsaNGyd9+/ad/fv0jEz6+/TsTHoLwrrrrpvcfvvt+bMxK620Uv5N2rP07Nkz/9+0SU16q0JIertAsZ1Z55xX2t0sbYJwyCGHzO5oueOOOybTpk3Lfx3pG8XTRjlQies/PSOf3m7ze7OuSqV/DpW8B/wMoJrX/yxpU8Cdd945+eqrr/JXztJbRtM9ANWyB6pd42pv937yySfnbzFIby+YU3rbypzS++vT27jmlN7qmEo/zypdyB988EHy3nvvJW3atAk+3rfffltwLukL8F69ev3pzZje/jCnPffcM78h0/cRejFDJa//0PtSfvnll9l/DpW+B1J+BlCN63+WtIBIf6XSJmmbbbZZvvP3Sy+9lP8YE6j0PVDtqragTZsFdO/ePf/eo/QMdvpG8HQxpZ/ll36m5Z9pqpH+ndVWWy3fWS/kj97Ll55NHDduXFGPk37OVHpmaNYGSzfkoosu+psxiyyyyOx7/qFS13/a3S997+DvpR3/Zu0PqOQ94GcA1bz+C0mv1qZXz9L3LXbo0KGo41JdKn0PVJuqLWjTz2n6/vvvk7vuuiv/IdyzfPLJJ8Hx6W0sadOBOc/OpE+Uqfbt2+f/m26GN998M9lkk01KPiOYdmNNO6YVI71loUePHvn/n7amT29RSF/Uz/mknc43VegsEdWtUtZ/p06d8r9Pz6zO2RgqPSs/68+hkveAnwFU8/ovZNbbTX5/lQ2qZQ9Um6otaNO27qlcLn3/9P+Tvt/i8ssvD46fMWNG/vattBParLHp79MXC+kLilTa+eyhhx5Krrnmmsxnn6VPrumZm9/frlDbe+fTxzznnHOS6667bvY9+qlrr702f7+/BU8lr//0LPz555+fXH311bM/hza9BXnIkCFJ165ddTim4veAnwFU8/pPb+GcdTfCLNOnT09uuOGG/O34q6yySlHHpPpUyh6gygva9E3ZrVu3Tvr06ZMceeSR+TMpw4YN+83CnlN6W9fgwYPz98mn98zfeuut+Y8ESV9Iz2qdnX5MQtqMIG3OkZ492WCDDfK3EKRvLE/zRx99NFl77bXr9N75NddcM98u/F//+ld+s6W3T6RnndI3pp944oluuaSi139atO6yyy75tZ6+sEk/923o0KH5eaYv8KHS94CfAVTz+k9vK07v0EmvsC2++OLJN998k9x00035x7zgggvyH8UClbwHUpdddln+FupZd+bcf//9yRdffJH//0cccUT+Y34qXq6K23Wn7bPXXXfdXIsWLXKLLbZYrn///rlHH3000xI7bdedfjzIq6++mltvvfVyzZs3z7fbvuyyyzKPkbbuHjx4cH58s2bNcq1bt8517tw5N2jQoNykSZMKtuuujfQxBw4cmD9mkyZNcssvv3zuoosuqpNjUzkqdf1PmTIld9xxx+Xatm2bf8x11lkn98gjj9TJsakslboH/AygWtf/LbfckuvVq1du0UUXzTVu3Dj/eOnv77333lofm8pTiXtg1rGSAh/xM+trrXQ16f+Uu6gGAACAUjUq+W8AAABAA6CgBQAAIEoKWgAAAKKkoAUAACBKCloAAACipKAFAAAgSgpaAAAAotS42IE1NTX1OxP4A+X+uGTrn2pe/yl7gGreA9Y/1bz+U/YADXkPuEILAABAlBS0AAAARElBCwAAQJQUtAAAAERJQQsAAECUFLQAAABESUELAABAlBS0AAAARElBCwAAQJQUtAAAAERJQQsAAECUFLQAAABESUELAABAlBS0AAAARElBCwAAQJQal3sCFG/TTTcN5ocddlgw33bbbTPZueeeGxz797//vZazAwAAmLtcoQUAACBKCloAAACipKAFAAAgSgpaAAAAoqSgBQAAIEo1uVwuV9TAmpr6n00VateuXSbbfPPNg2MvvPDCYN6qVauiH2/69OkldUq+7rrrkoagyGVab6x/qnn9p+yBuWe++eYL5tdcc00w33333YP5iy++WPTPlx9++CFpyMq9B6z/4jRt2jSYN2vWrOhj9OrVK5gPGDAgmK+22mpFH7vQMc4444ykISv3+k/ZA+X/9x40aFAwHzhwYFLp/tcecIUWAACAKCloAQAAiJKCFgAAgCgpaAEAAIiSghYAAIAoNS73BKqlO2Xv3r2D+f7775/JOnfunNSXeeaZJ5jPP//89faYNCyNG4e3/YEHHhjMV1hhhaKP/dNPPwXza6+9Nph/++23mWzq1KlFPx7EbqWVVspkDz30UHBs+/btS+r+2LVr10y29957B8f+85///B8zpSEo9DO8Q4cOwbxv377J3LT66qsH827duhXdObfUjr6ljA/tCSinUjoUd+/evV7nEjNXaAEAAIiSghYAAIAoKWgBAACIkoIWAACAKGkKVccKNfPYYIMNat0QoVCznIsuuiiYH3bYYZlswoQJwbEXX3xxMKfynHzyySXlpQit59Q//vGPYD5y5MhM9sQTTwTHFspfe+21kuYI5dCuXbtg/uijj2ayJZdcMjj26quvDuannXZaMP/www+LbgpHHBZZZJFg/tZbb831uTR0U6ZMyWR33XVXWeYCdaFHjx7lnkKD5QotAAAAUVLQAgAAECUFLQAAAFFS0AIAABAlBS0AAABR0u6wCCuttFIwv/fee4vuTlmK8ePHB/ODDjoomN9zzz1Fd9W85ZZbajk7YrLHHntkslNOOSU4tlB37fq08cYbF5WlBg4cGMxff/31YH7rrbdmsqeeeio49s033/wfM4XitGjRoqRu9KGfGY888khw7LHHHhvMf/7552D+wAMPZLJ33nknOBYqTehn3ZAhQ8oyFyike/fu5Z5CRXCFFgAAgCgpaAEAAIiSghYAAIAoKWgBAACIkoIWAACAKNXkimxtWlNTk1S6xo3DTZ8vueSSYH7IIYfU+jE///zzTPa3v/0tOPbuu+9OqlU5OvBWwvofM2ZM0V276+J7XOj71FCO/dNPPwXzQt2/+/XrlzQE5V7/Me+Bua3QmvnnP/8ZzD/55JNMtsYaa5S0fgtp3759Jvvyyy+DY6dPn540ZOXeAw1l/Tdr1iyYX3bZZcF8v/32q/VjvvHGG8E89LOkUJfvUr6vhf6tp0yZEswLde6/6aabMtm4ceOSGJV7/TekPRCrHj16BPORI0cWfYxBgwaV9CkQleR/7QFXaAEAAIiSghYAAIAoKWgBAACIkoIWAACAKIW7IFW4Qk1xjjjiiHpr/lRK0w4o1aWXXlr0Wm/UKHwea+bMmbWeR6FjfPXVV8F8+PDhmeyhhx4Kjn3qqaeC+WKLLRbMd9ttt6IbrvXt2zeYb7311plshx12CI4dPXp0MJ8xY0YwJ35rr712Jrv44ouDY8ePHx/Md91111o3fyrk008/rZPj0HBMnTo1mB955JHBfOjQobV+zELPba+99lomW2655eb61zhkyJBaPyaUqykUdcMVWgAAAKKkoAUAACBKCloAAACipKAFAAAgSgpaAAAAolSVXY4LdeGri27Gjz/+eEldaKEU888/fzDfaKONgnkulyu6E/GPP/5YUpfMtdZaK5M99thjwbGnn356Ul8KdVC+6KKLMtnXX38dHHvTTTcF83bt2mWyF198MTj2sMMOC+ZXXXVVMCd+oa6rTZo0CY594YUXiu4UC6WaMmVKMH/22WdrfexC3YWXXHLJWh871AX+0EMPrbeOzRCzgQMHlnsKDZYrtAAAAERJQQsAAECUFLQAAABESUELAABAlBS0AAAARKkmF2qDGhpYU1P/s5lLRowYEcy7d+9e0nEmTpyYyTbZZJPg2NGjR5d0bH6ryGVabxrK+u/Tp08wv+6662r9tRx99NFV26G7UJfj3XbbrehjPPjgg8F8u+22S2Jf/w1pD5RDly5dgvnzzz+fyT766KPg2LXXXruk7uI0rD1QDev/iCOOCOaDBw8O5k2bNq31Y+67776Z7MYbb6z1cStNudd/teyBhvJvOGrUqGC+8cYbJ9Uq9z++f67QAgAAECUFLQAAAFFS0AIAABAlBS0AAABRUtACAAAQpcZJFVp22WXr5Dj77LNPJtPNmLqy2GKLZbLLLrus1sf96quvgvm1116bVKtvvvmm1sdo165dncyF8inUtfX6668P5o0aZc8JDxs2rKRuxs2bNy96Hj/88EMwh1Icdthhwfzcc88N5k2aNKm3uehoTKUZOHBgrY9Rzd2M/yxXaAEAAIiSghYAAIAoKWgBAACIkoIWAACAKFV8U6gTTzwxky211FJ1cuxnnnmm1sfo2LFjJuvWrVtJx9h8882D+bbbblv0Me69995gvttuu2WyadOmlTA7/qyePXtmspYtW9b6uIWa00yZMiWpVvPPP38wr6mpKfoYTz/9dB3OiHLYcccdg/lKK61U9DFWXHHFYP7JJ58E88aNsz+G55lnnuDYX375JZgPHz48mA8YMCCTTZ8+PTiWyrTDDjtkssMPP3yuN38q5TVaqe6+++5gPnbs2FofG0oVet6l/rlCCwAAQJQUtAAAAERJQQsAAECUFLQAAABESUELAABAlCqmy3GoU2Shjsa5XK6kY1988cXB/Oeff85ka6yxRkldVG+99dZM1rZt26QulPJ1FuqI3Lx580ymy/Hcseaaa9Z67YZcc801SbXaeuutg/kBBxwQzEv5ftfFvw3ltfbaa9f6GL179w7mhZ43Q/uxUDfjPn36BPO///3vwfyRRx7JZLpxV6bll18+mN9xxx1JQ3bWWWdlspkzZ5Z0jDPOOCOY33bbbZnslFNOCY798MMPS3pMoGFxhRYAAIAoKWgBAACIkoIWAACAKCloAQAAiJKCFgAAgChVTJfjeeedN5gffPDBtT72Dz/8EMx79uyZyW688cbg2IUXXjiY19TU1Lpb6tSpU4N5kyZNMlmjRs5hVLvhw4cn1WrLLbest2PrkhmPli1bBvOtttqq1sf+7LPPgvlJJ50UzG+55Zaij33nnXcG8+effz6YX3XVVZmsc+fOwbGTJ08ueh7Eo6F3Xw91NK6rOe+yyy6ZrEuXLsGxO+64YzAfM2ZMJpsxY0YdzI5KMHDgwFofY9CgQXUyl2qnugEAACBKCloAAACipKAFAAAgSgpaAAAAolSTK/Ld96HmRQ1Jq1atgvn48eOThqyUplD33XdfML/yyiuLbgiy5JJLljS/1q1bF90kq5IbW5Rj/Y8cOTKTdevWraRjvP7660U3xag0p556aib7xz/+ERzbuHHjotfd+++/Hxy73nrrBfNJkyYlsa//GH4GlGK33XardYOm1JdffpnJNt5443prGhZq9PdHjQFD2rZtG8y//fbbpCEr9x5o6Ot/gQUWCOb9+vXLZPvtt19JzdIKHbtZs2aZ7Oeffw6O/e6774r+vhZqolnodV59WnfddTPZq6++WnXrP4Y9UJ969OhR9Ou0UlXz97Uu94ArtAAAAERJQQsAAECUFLQAAABESUELAABAlBS0AAAARCnc2pO55sEHH8xk//znP4Nj559//mC+1VZbBfPFFlus6HmMHTs2mM+YMaPoY1C3unfvXutOh08//XRS6Tp27BjM+/btW3Q340JdBqdNm5bJevfuXW/djJk72rVrVyfHefjhh+ulmzGUqtCnDwwePLio7I86YLdv3z6Y/+Uvf8lk33zzTXDs6NGjk2J16tQpmK+zzjrB/Oijjw7mHTp0SGrrpJNOKrpL+vTp02v9eMTV5bgUgwYNqpO5EOYKLQAAAFFS0AIAABAlBS0AAABRUtACAAAQJQUtAAAAUdLluI4V6nJ34YUXBvOzzz47k2222WbBscOHD6/l7JLk3//+dzDfdtttg/nkyZNr/Zj8OaGOxqV2OS51fIzdjEOdwlOLLrpo0d+PUDfjQt0zX3/99f8xU6rFHXfcMVcfr1AX2kLGjBmTyX788cc6nBGVpFCH4kJ5fSnUEblQXuhnwKhRozLZsssuW9JcQq+NFlxwweDY//73vyUdm7g/daJUofVI3XGFFgAAgCgpaAEAAIiSghYAAIAoKWgBAACIkoIWAACAKOlyXMcmTpxY0vjbb789k2266aZJfTn22GOD+UcffVRvj8mf88EHH2Sy5ZdfPql0p556ajDv27dv0d2MS3XEEUcE82uvvbbWx6bh+f777+vkOCNGjEjqQ+PG4R/NQ4cOLek4w4YNy2RTpkz50/Oi/Jo1axbMd9xxx2B+yCGHZLL//Oc/wbGXXHJJMH/11VeThmD11VcP5scff3wwL7WjccgXX3xRdFd84tejR4+S8lLocly/XKEFAAAgSgpaAAAAoqSgBQAAIEoKWgAAAKJUMU2hampqkoagTZs2wfyEE04I5o0aZc8pzJw5s6THHDNmTDC/+eabM9njjz9e0rEpnwcffDCTHXXUUUmMtt5662B+8sknZ7I111yzpEY5uVyu6HkceuihwVzzp+ry2GOP1clxFlhggUw2fvz4ko7RpEmTohv8FGpM8uWXX5bU5Id4HXfcccF80KBBRR9jgw02KOl5+uOPPw7mb731ViZ76KGHklKceOKJRT+nL7nkksF8wQUXTOrLnnvumckmTJhQb49HedVF8yfKwxVaAAAAoqSgBQAAIEoKWgAAAKKkoAUAACBKCloAAACiVDFdjn/66adgvtFGGxXd+bFQd9X6VEqH1vfffz+Yb7PNNsH8s88++9PzomGu6VK7eYe6sJaqZcuWwXyhhRbKZKecckpw7AEHHFDreRT62qdNmxbMjzjiiEymmzF/1In4qaeeCubdu3cvuuPsSSedVHQ340IdjW+55ZaSfs5ttdVWwXzq1KnBnHgtssgi9Xbs+eefP5ivscYaRed77713rZ/XS3ldVKovvvgimF922WXB/JVXXqm3udDwFHqurwsDBw4sKac0rtACAAAQJQUtAAAAUVLQAgAAECUFLQAAAFFS0AIAABClmlyR7eRK7a7akC299NLB/P777w/mq666ar3N5Zlnnslkw4cPD4594okngvmHH36YVLr67HrYUNd/qJvlO++8Exy74IILFn3cO++8s6R5LLHEEsG8a9euRX+f6uLfr9D6Hzx4cDAfOXJkUinKvf4r7WdAIaGu+KmHH344mE+ZMqXoPTrvvPMG886dOxfdzXjbbbcN5qNGjUoqXbn3QENZ/4U+peGwww5LYlSfXY7vu+++THbqqacGxxbatw1Fudd/Q9oDsX6fN95446p9/p4b/zau0AIAABAlBS0AAABRUtACAAAQJQUtAAAAUarKplDEp9wNERrK+l9xxRWDeb9+/YL5gQcemMlatmxZb9/jUptCjRgxoujmT+eee25Srcq9/hvSHiiHxRZbLJjfcMMNmaxnz57BsRMnTgzmt99+eya79NJLo2xcU8l7oKGs/2bNmgXzxo0bF32MXXfdNZgvu+yyJc3lkEMOyWStW7cu6RhPP/10JnvuuedK2kNXXnllMJ86dWommzFjRhKjcq//hrQH6lN9vg6idjSFAgAAoCIpaAEAAIiSghYAAIAoKWgBAACIkoIWAACAKOlyTBTK3eEv1vXfrl27oruwdurUqdaP9/PPPwfza6+9Nph/++23mWzatGm1nkelKff6j3kPUBnKvQesf6p5/afsAcpJl2MAAAAqkoIWAACAKCloAQAAiJKCFgAAgCgpaAEAAIiSLsdEodwd/qx/qnn9p+wBqnkPWP9U8/pP2QOUky7HAAAAVCQFLQAAAFFS0AIAABAlBS0AAABRUtACAAAQJQUtAAAAUVLQAgAAECUFLQAAAFFS0AIAABAlBS0AAABRUtACAAAQJQUtAAAAUVLQAgAAECUFLQAAAFFS0AIAABAlBS0AAABRqsnlcrlyTwIAAABK5QotAAAAUVLQAgAAECUFLQAAAFFS0AIAABAlBS0AAABRUtACAAAQJQUtAAAAUVLQAgAAECUFLQAAAFGqqoL2+uuvT2pqapJPP/20pL/Xo0ePpGPHjnU6l/bt2yf77rtvnR4T/oj1T7WzB6hm1j/Vzh6oXFVV0FaqDz74INl9992TJZZYImnZsmWy0korJaeddloyefLkck8N6t3UqVOTE044IVlsscWSFi1aJF27dk0ef/zxck8L5or0BVH6Aq3Qry+//LLcU4R69dprryV//etfkwUWWCCZf/75k8022ywZPXp0uacFc80H6oCkcbknQO18/vnnSZcuXZJWrVolhx9+eLLgggsmL7zwQjJgwID8k/y9995b7ilCvb+gv+OOO5Kjjz46WWGFFfJnYLfccstk5MiRyYYbblju6UG96tu3b9KrV6/fZLlcLjnkkEPyVwAWX3zxss0N6tvrr7+ef55fcskl8697Zs6cmVx++eVJ9+7dk5dffjnp0KFDuacI9Uod8P8oaCM3bNiwZOLEicmzzz6brLrqqvns4IMPzj+p33DDDcmECROS1q1bl3uaUC/SFyzDhw9PzjvvvOS4447LZ/vss0/+1qD+/fsnzz//fLmnCPVqvfXWy/+aU/rzID0zv9dee5VtXjA3nHLKKfk7c9IX8AsttFA+6927d7LiiismJ510UnLnnXeWe4pQr9QB/09V33KcnrXYaqut8rcqNmvWLFluueWS008/Pfn111+D49MzHeuvv37+yXOZZZZJrrzyyuDtj+lZkeWXXz5/zPSsYfrCOs3rww8//JD/76KLLvqbvF27dkmjRo2Spk2b1svjEr9KWP/pldl55pkn/+Q9S/PmzZMDDjgg/wInPXMJlbwHQm6++eb87cZ77rnnXHtM4lMJ6/+ZZ57J36Ewq5id9fonvUL7wAMPJD/99FO9PC6VoRL2gDrg/6nqK7TprYnzzTdfcswxx+T/O2LEiOTUU0/NL470is+c0jMc6W2Mu+66a7LHHnskt912W9KvX7/8Qtl///3zY9KzIdtuu23+LEn6AnvllVdO3n777eSiiy5K3n///eSee+4pOJf0744fP76oeae3FTRp0mT2G9UHDx6cfwE/aNCg/JN6elXqiiuuSI488shk3nnnrdX3iMpVCev/jTfeyJ+JT987Naf09ptU+j6q9IcJVOoe+L3p06fn55a+6EpvOYZKXv9pkZAWF7+Xvo9w2rRpyTvvvJOsu+66JX5nqBaVsAfUAf9/uSoyZMiQXPolf/LJJ/nfT548OTOmb9++uZYtW+Z++eWX2Vn37t3zf++CCy6YnU2dOjXXqVOn3CKLLJKbNm1aPhs2bFiuUaNGuWeeeeY3x7zyyivzf/+5556bnS299NK5Pn36zP59Oqd0TDG/Ro4c+Zvjn3766bkWLVr8Zsw//vGPOvmeUTkqcf2vuuqquZ49e2a+jjFjxuTHpo8NlbwHfu/+++/Pj7n88sv/9PeJylSJ63+11VbLrbjiirkZM2b8Zm5LLbVUfuwdd9xRB985KkUl7oHU6eqAXFVfoZ3zrN6PP/6YP9PXrVu35KqrrkrGjh2brLHGGrP/vHHjxvnmG7OkZ2TS36dnZ9JbENIzgLfffnv+bEzaXey7776bPbZnz575/6ZNatKz5iFt27YtujPrnPNKpWfhN9poo2SnnXbKn5l58MEHk7POOit/zPQN4lCp63/KlCn5W3p+L73teNafQyXvgdDtxumZ+/QqAlT6+j/00EPzc0ivTqW3daZXuc4444zk66+/zv+5nwFU+h5ItVcHVPctx2PGjElOPvnk/C0Gs+5Bn2XSpEm/+X16f/3vL9untzqm0s+zShdy2jb7vffeS9q0aRN8vG+//bbgXNIX4L/vVFmMtCFOeltDeitD2q47teOOO+af1NOPMklvi5jzvSVQSes//WEUel/KL7/8MvvPoZL3wJzS9wum7wnbfPPNPe9TFes/7ead9kpIbw8dOnRoPlt77bXzxe2ZZ56Zv40UKnkPqAOqvKBNO4KlTQPS996ln9WUvhE8XUxpC/h0AaQLoVTp31lttdWSCy+8MPjnf/RevvQN6OPGjSvqcdKW3LPe5J22p19zzTVnL+JZ0nv40/cGpO8xrO2LJCpPpaz/tOlB6HM2Z52dT38AQSXvgTml78/S3ZhqW/9p4Zp2uU+Lk/S9hekc0g7HcxYcUKl7QB1Q5QXtqFGjku+//z6566678pfpZ/nkk0+C47/66qvk559//s3ZmfRsSGpW4410M7z55pvJJptsku8wWYr0DGPaMa0Y6S0L6ZvAU//973+D7bjTxiCpGTNmlDQPqkOlrP9OnTrlf5+eWZ2zMdRLL700+8+hkvfAnG666ab8Fan0hQxU0/pPXwfN+bnjTzzxRP4FfnrrJ1TyHlAHVHlBm37Ux6wPoJ8l7YiXnukISRdEek992glt1tj09+ltBZ07d85n6XuWHnrooeSaa675zceIzHofR3rmplC3sT9773x69vGxxx7Lb6o5z0Tecsst+Xbdq6++elHHpLpUyvrfeeedk/PPPz+5+uqrZ38ObXoL8pAhQ5KuXbvqcEzF74FZ0jP76Yv49PaytMMrVNP6n9Ott96avPLKK/mfDenrIKjkPaAOqPKCNn1TdnpGo0+fPvm21umZlPTDiedc2HNKb11M22Kn98mnCyZ9wkw/EiR9IT2rdfbee++db+OdvqcjPXuywQYb5G8hSN9YnuaPPvpo/r0ddXnv/PHHH588/PDD+Texp2/8Tu+TTz97Lc0OPPBAt1xS0es/LVp32WWX5MQTT8y/NyX93Lf0fVTpPK+77rqSj0f1qJQ9MEs6n/QFl9uNqab1//TTT+dvF91ss83yr39efPHF/AnNv/71r8lRRx1V8vGoHpWyB9QB/3+5Km7XnbbPXnfddfOtrhdbbLFc//79c48++mimJXbarjv9eJBXX301t9566+WaN2+eb7d92WWXZR4jbd09ePDg/PhmzZrlWrdunevcuXNu0KBBuUmTJhVs110bL730Um6LLbbItW3bNtekSZN8C/szzzwzN3369Do5PpWhUtf/lClTcscdd1x+/aePuc466+QeeeSROjk2laVS90Aq/TrSj4+Y8+NLoNLX/4cffpjbbLPNcgsvvHD+8VZaaaXc2Wefnf9IFaiGPZB6SR2Qq0n/Z1ZxCwAAALHw5gIAAACipKAFAAAgSgpaAAAAoqSgBQAAIEoKWgAAAKKkoAUAACBKCloAAACi1LjYgTU1NfU7E/gD5f64ZOufal7/KXuAat4D1j/VvP5T9gANeQ+4QgsAAECUFLQAAABESUELAABAlBS0AAAARElBCwAAQJQUtAAAAERJQQsAAECUFLQAAABESUELAABAlBS0AAAARElBCwAAQJQUtAAAAERJQQsAAECUFLQAAABESUELAABAlBS0AAAARKlxuSdQ7f72t79lsgsuuCA4dr/99gvmQ4cOrfN5AQAANHSu0AIAABAlBS0AAABRUtACAAAQJQUtAAAAUVLQAgAAECVdjueShx9+OJhvsskmmWzUqFHBsXfccUedzwvmhkIduk8++eRMtswyywTH1tTUBPNcLlf0fjnrrLOCY0ePHh3MAQBo2FyhBQAAIEoKWgAAAKKkoAUAACBKCloAAACipKAFAAAgSjW5Qi1Ci+wwWs0WWmihTHb//fcHx3bp0iWYT5gwIZNtuOGGwbH//ve/k2pV5DKtN9Z/1uDBgzPZUUcdFRzbuHHjBvF9De23VK9evRp09+Nyr/+UPUA17wHrn2pe/yl7YO6Zf/75g3m/fv1KOs5pp52WyZo1axYce8IJJwTzc889N4lhD7hCCwAAQJQUtAAAAERJQQsAAECUFLQAAABESVOoIhRq0nTppZdmsjXWWCM4dujQocH8yCOPzGQ//vhjyXOsdOVuiFDN6/+1114L5quvvnoma9QoznNkw4YNC+b77rtv0hCUe/1X+x6g/Mq9B6z/4nTu3DmYb7/99sG8TZs2mWyHHXYoemzqvffeC+Z33XVXJjv77LODYydPnpw0ZOVe/yl7oHa22267YN6/f/9M1qFDh+DY1q1bJ/Vl+vTpwfziiy/OZNdee21w7IcffpjUF02hAAAAqEgKWgAAAKKkoAUAACBKCloAAACipKAFAAAgSlXZ5bhQJ9ZzzjknmB9++OHBvHHjxpns2GOPDY697LLLGmznuhiU+/tUSeu/VIU6SK644opJpfj555+L7nD+1ltvJdW2/su1B0KPucIKKwTH7rjjjsF8scUWK/rxdtppp2Derl27oudX6r/ho48+msk++OCD4NgzzjgjmH/77bdJpSv3HqjmnwEbbbRRMD/xxBMz2WabbVbSv1/o+1rK2FLH77PPPsGxN910U9KQlXv9V/seKOSss84qupvxMsssE8ybNWuW1JdHHnmk6G7hhTqUh7z77rvBfLXVVkvqiy7HAAAAVCQFLQAAAFFS0AIAABAlBS0AAABRUtACAAAQpWyb3grTtm3bTDZo0KDg2IMOOiiYf/7558F8wIABmez6668veY7Vqnnz5pnsl19+KctcKOyJJ56o+C7H8847bzDv0aNHg+hyXK2aNGlSdNft+lSou2JddB4NdYUt1Cl2vvnmC+b/+Mc/gvnXX39dy9lRiQo9391www3BfIcddih6/ZfaCbeU8XVx7EJf42OPPRbMx40bV9JjEreVVlopmF955ZVFfxJCqet08uTJmeztt98Ojr3vvvuC+bPPPhvMX3jhhUx25JFH1rrL8ZQpU5KGxhVaAAAAoqSgBQAAIEoKWgAAAKKkoAUAACBKFdMUatFFFw3mjzzySCZbffXVg2O//PLLYL755psH87Fjx5Y0x2q18847F93IZM0115wLMyJk2WWXDeY77rhjUulCTRlSL7744lyfC//7ubq+FGry9PPPPwfzTz/9NJN16NCh1o2vCunTp08w/89//hPMBw4cWNJcqA5///vfg/l2221Xb03R7rrrrmB+zz33FN2EqpTmVIUUGlvo2FdffXXRxyYe/fv3D+YHH3xwMF9mmWWKPvZPP/0UzE866aRg/u6772aykSNHJnWhVatWmezoo48u6RjTp0/PZIMHD04aGldoAQAAiJKCFgAAgCgpaAEAAIiSghYAAIAoKWgBAACIUsV0OT7rrLOK7pL59ttvB8eus846wXzatGlJtQp14OzcuXNw7GWXXRbMV1lllWDet2/fWs6Ouuw6fckllwTztm3bJg3ZO++8E8w7duxY9DFatmwZzNddd91M9vLLL5cwO2rjjTfeyGS33nprSf/eoQ6NV111VXDsxx9/HMyfeOKJpLZatGgRzO++++5Mtummm5Z07EKd+M8555xM9ssvv5R0bOIW6t578sknB8fOnDkzmNfU1BS9dgv9fClFoS6sheZRSKnjqUzt27evl27GqYceeiiTXXjhhcGxddW5uBS33XZbJltiiSVKOkaoo/Gdd96ZNDSu0AIAABAlBS0AAABRUtACAAAQJQUtAAAAUVLQAgAAEKWaXC6Xq4RucTvttFMwv/TSSzPZIossEhz71FNPBfMzzjijwXQsC2nVqlUwb926dSbba6+9gmN32223ojtzFnq8a6+9NpgX6kr65ptvJsUqcpnWm4a+/kvp7vfggw8Gx6600krJ3HbzzTcX1VHvjxTqzPnwww/Xurvfhx9+mMk6dOiQzG3lXv8x74GG7i9/+UvRnbSXW265ko49YMCAon+eNXTl3gOxrv9XXnklk6211lolfY/vueeeYL7PPvtkssmTJyf1Mec/M+/Qv1mhsYW6+X/33XdJQ1Du9R/zHgh9WsFzzz1X0jGefPLJYL7FFltksl9//TWZ2zbeeOOiuzA3bdo0OPbTTz8tuot+6LVRufeAK7QAAABESUELAABAlBS0AAAARElBCwAAQJQUtAAAAESpcVIh7rzzzmD+0UcfZbLLL7+8pC5ha6+9djAfOnRoJjv33HODY7/44otgPu+882aynXfeueiOgqllllmm6A63n3/+eXDsiBEjgvl7772Xyf71r3816G6A/H/uvffeBtHNePz48cH8vPPOy2TvvPNOnTzm008/ncn23HPPko6x7LLLZrK99947OHbYsGElHRtSEydOLLrjfqldjrfaaqtMdvbZZwfHlqMzJ3F0qy00fuWVVy76GNtvv33Rn1BRqJN8Xcx7nXXWCY71+oU/Uug5eW4/b15wwQXB/Igjjgjm88wzT9EdikMdm1Mff/xxEgNXaAEAAIiSghYAAIAoKWgBAACIkoIWAACAKNXkcrlcUQNLfDN+Q9a4cbgX1gknnBDMDz744GC+5JJLFv2YN910UzDfdtttM9n8889fdPOQ1HXXXVd0o6wXX3wxiVGRy7TeNPT1v+uuuxbdqKjQ+q8Lr776ajAfOHBgMH/44Ydr/ZhLL710MH/55Zcz2cILL1zSsSdPnlx0U60vv/wyqdT1H8MeqCT77bdfML/22mtrfezmzZsH8+nTpycNWbn3QKzr/5VXXslka621Vknf40Jfe2h8KWMLja+LeaTuueeeoptrhp7rG5Jyr/+Y90DTpk2LbiS75ZZbBvMff/wxmG+22WZFvfb4I7179y66eV+rVq2KbjBbyKmnnhrMzzzzzCTmPeAKLQAAAFFS0AIAABAlBS0AAABRUtACAAAQJQUtAAAAUaq/dqcN2HzzzRfMb7/99mC+3HLLBfN999236Mfca6+9ih5baB4XXXRRMI+1czGl23333YP5oEGDgnl9djS++eabM9mhhx5aUofAutChQ4dgXmpH45Bff/11rnYzhrp07733FrWmqVxjx47NZJ07d6637raldsKti2N/9913wXznnXcuaS5UpmnTpmWyt99+u6Qux4U+fST06ShXXnllcOzRRx8dzDfccMOi65RCPv7446JrjzfffDOpRK7QAgAAECUFLQAAAFFS0AIAABAlBS0AAABRUtACAAAQpYrvcty1a9dMdskllwTHdunSJZjPnDkzmE+YMCGT3XrrrcGxbdq0CeY77LBDJuvRo0dw7GmnnRbMqR5LLbVUMF9++eXr7TEPPPDAYH7HHXfM1W7GoU6Aqeuvv77eHvOGG26ot2NDIc2bN6+T43zzzTdF/zyjMu29995Ff2LCyiuvHMzfe++9oh+v0DGGDh1a9DFyuVxSirPOOquk8XDGGWcE844dOwbzrbbaKphvv/32RWWlmj59ejA/++yzg/mNN94YzD/66KOkWrhCCwAAQJQUtAAAAERJQQsAAECUFLQAAABEqeKbQp1//vlFN3/673//G8zPOeecYF6ouVQpTj311Ew2cODA4NgRI0YE88022yyYv/nmm7WcHdVk4sSJwXzkyJHBvL4aQBVq/nTbbbcF80UXXbTWj1noa7n44otrfWwo1UEHHVQnx/n222/r5DhUltdff72kvBShRpepmpqakvKQxx57rN5ei1G5mjZtmslWWGGF4NiVVlopmdtCTc1effXV4Nh77713LswoTq7QAgAAECUFLQAAAFFS0AIAABAlBS0AAABRUtACAAAQpYrpcnzssccG8/XWWy+TzZw5s6TOkg888EBSX84888xMtskmmwTHduvWLZj37t07mOtyTCkefPDBYP7pp5/W+thLL710MO/QoUMmu/766+utm3Eh06ZNC+Yff/xxvT0mpBZZZJFM1rp165KOUahD/zXXXPOn5wV/xkknnRTMc7lc0ccoNHbvvff+0/Oi8nXs2DGYn3zyyZlsl112qZPHnDFjRiZr3Li00uqpp57KZE888USt5lWNXKEFAAAgSgpaAAAAoqSgBQAAIEoKWgAAAKKkoAUAACBKFdPlePvttw/mjRpla/Zbb711rnczLuTXX3/NZNOnTy/pGHvttVcwP/fcczPZuHHjSjo21WPeeecN5k2aNAnmTZs2zWQbbLBBcOywYcOC+cILL5w0BBMmTCj3FKhS++yzTyZbaqmlSjrGa6+9Fsy//PLLPz0v+F9Cz+s1NTUlHSM0/uqrrw6O/e6770o6NtXlgAMOCOaldDSeOnVqMD/vvPOC+cSJEzPZ+eefn5SiZ8+emUyX49K5QgsAAECUFLQAAABESUELAABAlBS0AAAARElBCwAAQJQqpsvxRx99FMxDXVdj7WhaqHvgW2+9Fcx1NKYuOoXffPPNwbxVq1aZbJNNNkkasiFDhgTzUrsSQqlC+yV1+OGHF32MadOmldSBE+rCSiutVPTPjFwuFxxbKA91Lr7mmmtKniPV49JLLw3mhxxySNHHKPS6pl+/fsH8p59+CuZHHXVUUlubb755JjvppJNqfdxq4wotAAAAUVLQAgAAECUFLQAAAFFS0AIAABClimkK9dJLLwXzffbZJ5O1adMmmdu6du0azHv37p3JunfvHhw7adKkYH7GGWfUcnZQ2I477pg0ZJ988knRjZ5GjRoVHDt27Ng6nxfMqU+fPsF8ySWXLPoYTz/9dEk51IUtttgimLds2bLo5pWF3HTTTZns9ddfL+kYVJfddtstmDdqFL5GN3r06KIbSP3888/J3Pb222/P9cesRK7QAgAAECUFLQAAAFFS0AIAABAlBS0AAABRUtACAAAQpYrpcjx06NBgvummm2ayHXbYITj2ySefDOYjRowI5k2bNs1ku+++e3DscsstV3RXtm+//bakzm7PPvtsMKfyFFqj33//fTBfcMEFa92FsqF3M/7rX/8azD/88MN6nhFkrbXWWvXWjf6+++6r9TGgVNtvv30wz+VyRR+j0NizzjrrT88LijF16tR662bcvn37Wh/jxhtvrJO5VDtXaAEAAIiSghYAAIAoKWgBAACIkoIWAACAKCloAQAAiFLFdDmePHlyMO/du3cm22effYJj+/fvH8xPP/30euti9sYbb2Sy66+/Pjh2woQJtZ4HcXvttdeC+SKLLBLMDz/88Ew2cODA4NjWrVvXcnZJMnPmzGBeqLPyjBkzil7/559/fjDXzZiGZPPNNw/m8847b9HH+Oqrr4L5dddd96fnBf9L3759g/lGG21U9PN96JMbCr0WS3333XclzREuuuiiYH7qqacG8w4dOmSyXXfdNTj2nXfeKel5/YgjjkiK9eijjwbz0aNHF30MCnOFFgAAgCgpaAEAAIiSghYAAIAoKWgBAACIUk0ul8sVNbBAUxeYG4pcpvWmktZ/p06dgvlWW20VzI866qhgPm7cuEx2xhlnBMc2a9YsmI8aNSqTffrpp8Gx1azc67/S9kBd2WOPPTLZtddeGxzbvHnzoo/bq1evYD5y5MikWpV7D1TS+m/Tpk0wf+ihh4L5WmutVfS/SaHv0zrrrBPMX3/99T+YKQ1l/cewB0488cRgPmDAgEzWpEmTepvH1KlTg3mXLl1KakRFaXvAFVoAAACipKAFAAAgSgpaAAAAoqSgBQAAIEoKWgAAAKKkyzFRKHeHP+ufal7/1b4H5plnnmB+++23Z7LtttuupGM///zzmWyjjTZqsOugXMr9tVfS+l977bWD+UsvvRTMGzUKX/uYOXNm0V2Lt9hii2D+3Xff/cFMaSjrP+Y9EOpG//e//z04tmPHjiUd+9lnn81k5557bnDsgw8+WNKx+S1djgEAAKhICloAAACipKAFAAAgSgpaAAAAoqSgBQAAIEq6HBOFcnf4s/6p5vVf7XugdevW9dah9bnnniu6y3E1K/ceqKT137Jly5K6HK+yyirB/K677spk/fr1C47VzTju9V9pe4D46HIMAABARVLQAgAAECUFLQAAAFFS0AIAABAlBS0AAABR0uWYKJS7w5/1TzWv/2rfA/PMM08wP/XUUzPZySefHBw7cuTIYL7ffvtlss8//7zkOVa6cu+Bal7/lF+513/KHqCcdDkGAACgIiloAQAAiJKCFgAAgCgpaAEAAIiSplBEodwNEax/qnn9p+wBqnkPWP9U8/pP2QOUk6ZQAAAAVCQFLQAAAFFS0AIAABAlBS0AAABRUtACAABQ2V2OAQAAoCFxhRYAAIAoKWgBAACIkoIWAACAKCloAQAAiJKCFgAAgCgpaAEAAIiSghYAAIAoKWgBAACIkoIWAACAKFVVQXv99dcnNTU1yaefflrS3+vRo0fSsWPHOp1L+/btk3333bdOjwl/xPqn2tkDVDPrn2pnD1Suqipoq8GZZ56Z36x1vfGgIRo1alR+vYd+vfjii+WeHtS7MWPGJLvsskuy7LLLJi1btkwWXnjhZKONNkruv//+ck8N5jqvgahGU6dOTU444YRkscUWS1q0aJF07do1efzxx5Nq0rjcE6DufPHFF8lZZ52VzDvvvOWeCsxVRx55ZLLOOuv8Jlt++eXLNh+YWz777LPkxx9/TPr06ZN/MTN58uTkzjvvTLbddtvkqquuSg4++OByTxHmCq+BqFb77rtvcscddyRHH310ssIKK+SvRG+55ZbJyJEjkw033DCpBgraCnLccccl6667bvLrr78m3333XbmnA3NNt27dkp133rnc04C5Ln3Rkv6a0+GHH5507tw5ufDCCxW0VA2vgahGL7/8cjJ8+PDkvPPOy++B1D777JO/S6F///7J888/n1SDqr7l+N5770222mqr/FntZs2aJcstt1xy+umn558MQ1577bVk/fXXz1/OX2aZZZIrr7wyeNl/wIAB+atD6TGXXHLJ/IJK8/r09NNP58/OXHzxxfX6OFSOSlr/qfQq1YwZM+r9cagclbYHZplnnnnyjztx4sS59pjEp5LWv9dAVOseSNf9PPPM85uTl82bN08OOOCA5IUXXkg+//zzpBpU9RXa9JL8fPPNlxxzzDH5/44YMSI59dRTkx9++CF/pmNOEyZMyJ8F33XXXZM99tgjue2225J+/folTZs2Tfbff//8mJkzZ+Zv83r22WfzC2vllVdO3n777eSiiy5K3n///eSee+4pOJf0744fP76oebdq1Spp0qTJ7N+nG++II45IDjzwwGS11Vb7098PqkulrP/Ufvvtl/z000/5J/X0am06/7XXXvtPfV+oHpW0B37++edkypQpyaRJk5L77rsvefjhh5PddtvtT31fqA6Vsv69BqKa98Abb7yRrLjiiskCCyzwmzFdunTJ/3f06NH5orri5arIkCFDcumX/Mknn+R/P3ny5MyYvn375lq2bJn75ZdfZmfdu3fP/70LLrhgdjZ16tRcp06dcossskhu2rRp+WzYsGG5Ro0a5Z555pnfHPPKK6/M//3nnntudrb00kvn+vTpM/v36ZzSMcX8Gjly5G+Of9lll+VatWqV+/bbb2fPd9VVV62D7xiVpBLXf3rMnXbaKXfdddfl7r333tzZZ5+dW2ihhXLNmzfPvf7663X2vaMyVOIemHPes/48ncPOO++cGz9+fK2/Z1SOSl3/XgNRzXsgXes9e/bMfB1jxozJj00fuxpU9RXa9JaBOW9XTG8HSK/upI00xo4dm6yxxhqz/7xx48ZJ3759Z/8+PSOT/j49O5PegpC+b+P222/Pn41ZaaWVfvP+jZ49e+b/m745O71VIaRt27ZFdySbc17ff/99/mzSKaeckrRp06bE7wDVrBLWf3q8OY+ZnhlN30u7+uqrJyeeeGLyyCOPFP39oPpUwh6YJW0Gkq79r776Kn/lIL1qNW3atCK/E1SjSlj/XgNR7XsgvTOnWbNmmTHpbcez/rwaNK72jzs4+eST87cYpLcXzCm9bWtO6f31v++cl17iT6WfZ5Uu5A8++CB57733Cj6pfvvttwXnki68Xr16lfw1pPNfcMEF87fbQLWt/5D0fSvbbbddctddd+Vf1Ke3IUOl74H0BVT6a1ZDkM022yzZZpttkpdeein/MSZQievfayCqfQ+kRfnUwPtzf/nll9l/Xg2qtqBNm2V07949f8/5aaedln8jeLqYXn/99fxnOaX3spcq/Tvp+zfSzpIhf3QPe/rCe9y4cUU9TvrknZ4ZSjfO1VdfnW+CkJ6Vn3MRT58+Pb/B0q8vHQ+Vtv7/SPpY6dWp9H2Fv39fCVTDHkiv1qZXD9L3bXXo0KGo41I9KmH9ew1Ete+BVLt27ZIvv/wyM+brr7+eXYhXg6otaEeNGpW/VSW9ipN+CP0sn3zySXB8+mSZvjie8+xM+kIh1b59+/x/083w5ptvJptssknJZ8TTLmRpx7RipLcs9OjRI7+A082TfgZn+uv30uMdddRRuv5Rkev/j3z88cf5H0xpkweoxj0w6zaz319lgEpZ/14DUe17INWpU6f879MrzHOewE/vzpn159WgagvaWbch5nLpe6b/n/SKzuWXXx4cn34cSHpPfdoJbdbY9PfpbQXp5/2l0s5nDz30UHLNNddkPvsvfXGRPvEW+sDvP3PvfPoZU3fffXfmz9PbJ9L3AlxyySX5zQWVuP5T6dnM39/ak/4wSbu8brHFFkmjRlX9yWRUwR5Ib2FbZJFFfvPn6dWpG264IX+r2SqrrFLUMakulbD+vQai2vfArLtxzj///PzdCrM+hza9BXnIkCFJ165dq6PDcTUXtOmbslu3bp306dMnf2YvPZMybNiw3yzsOaWX7AcPHpy/hSW9Z/7WW2/Nt8JOF9Cs1tl77713vhnHIYcckj9bssEGG+RvIUjfWJ7mjz76aMGPEvkz984vvPDCyfbbb5/JZ52NDP0ZVMr6T6UfS5K+aE+/nvRF/bvvvpufU8uWLZNzzjmn5ONRPSplD6S3Fadn5tMrDIsvvnjyzTffJDfddFP+MS+44AJ3KVCx699rIKp9D6TSonWXXXbJN8JMT3CmfUSGDh2an+d1112XVI1cFbfrTttnr7vuurkWLVrkFltssVz//v1zjz76aKYl9qwW8K+++mpuvfXWy38kSNpuO20V/3tp6+7Bgwfnxzdr1izXunXrXOfOnXODBg3KTZo0qWC77rqkZT3Vsv4vueSSXJcuXXILLrhgrnHjxrl27drlevfunfvggw9qfWwqTyXugVtuuSXXq1ev3KKLLprfA+njpb9PP8YKKn39h3gNRLXtgSlTpuSOO+64XNu2bfOPuc466+QeeeSRXDWpSf+n3EU1AAAAlMobzAAAAIiSghYAAIAoKWgBAACIkoIWAACAKCloAQAAiJKCFgAAgCgpaAEAAIhS42IH1tTU1O9M4A+U++OSrX+qef2n7AGqeQ9Y/1Tz+k/ZAzTkPeAKLQAAAFFS0AIAABAlBS0AAABRUtACAAAQJQUtAAAAUVLQAgAAECUFLQAAAFFS0AIAABAlBS0AAABRUtACAAAQJQUtAAAAUVLQAgAAECUFLQAAAFFS0AIAABAlBS0AAABRUtACAAAQJQUtAAAAUVLQAgAAECUFLQAAAFFS0AIAABAlBS0AAABRalzuCQBAQ7bZZpsF87333juT7bXXXsGxo0ePDuaffvppJttxxx1LniMAVCtXaAEAAIiSghYAAIAoKWgBAACIkoIWAACAKCloAQAAiFJNLpfLFTWwpqb+Z1PBBg4cGMxHjRpVVFbtilym9cb6p5rXf7XvgRdeeCGYd+nSpdbHnjx5cibbd999g2PvvPPOpFqVew9U8/qf2y699NJgftNNNwXzF198Mal05V7/KXug4TryyCOD+UorrZTJ+vbtW9KxGzXKXvtcfvnlg2M/+uijpFx7wBVaAAAAoqSgBQAAIEoKWgAAAKKkoAUAACBKmkLVQo8ePTLZgAEDih4b6/e60NdSKC/U5KqU5lflbojQ0P9NqGzlXv/Vvge++eabYN6mTZt6+b4++eSTwbGbbrppUq3KvQcqaf2vuuqqwfzvf/97Sc3I7rnnnlrP5dhjj81k5557bknNnzbYYIOk0pV7/VfaHmhIQj9HrrnmmuDYlVdeOZgvX6BJU12sm9C/+0knnRQcO3jw4KS+aAoFAABARVLQAgAAECUFLQAAAFFS0AIAABAlBS0AAABRalzuCcQs1NW3lG7GMQh9PSNHjizpGIU6P+uYR12Yf/75g/m8886byaZMmRIc26JFi1rPY8KECcF86tSptT425XXBBRcE83POOadeHq9p06bBvHHj8I/sGTNm1Ms8qEwdO3YM5r179w7m22+/fdEdV7/44ouS5hLqUNyoUfhay7rrrlsnn7AA5XDAAQcE84MPPjiTde7cOWnI7rjjjqShcYUWAACAKCloAQAAiJKCFgAAgCgpaAEAAIiSghYAAIAo6XJcRKe8Qnmh7r2lGDRoUDI3VcPXSPmccMIJwXzjjTeu9bELdcVeYYUVgvnSSy+dyb766qvg2MUXXzyY53K5ouf3+uuvB/N11lmn6GPQMF100UXB/K9//Wu9dLrfcMMNg/lKK60UzN95551aPyYUMt988wXziy++OJPtvPPO9TaPQt2P+/fvH8x1Oaa+tWnTJpP94x//CI494ogjav06g8JcoQUAACBKCloAAACipKAFAAAgSgpaAAAAolSVTaFKbYBUF00+CjVGGjhwYDI31efXWKgBw9z+Gimfs88+O2nICjV/KtRwqhRt27at9TFomDp27BjMV1111Xp5vFdffTWYf/zxx/XyeFSX1VdfvU6OU8rzZqGGTs2bN6/1PH755ZdaHwP+yD777BPMjz/++Ey28sor1/rx7rvvvqIbsaWefvrpoo99wAEHBPOrrroqiZkrtAAAAERJQQsAAECUFLQAAABESUELAABAlBS0AAAARKniuxyHOux27969oroZl9K1uT67GW+88ca1PjYNz/zzzx/Mn3rqqaKPkcvlgvn777+fyaZMmRIc+9xzzwXzTz/9NJi3b9++6LGlePfdd4P5888/X+tjU16FuhY/+uijwXzhhReul3msscYawfy0004L5qeeemownzx5cp3Oi8rQuXPnBtNZeYsttqj1sYcPH17rY1BdCnXoPvPMM4P53/72t2DepEmToh/zs88+C+Z77LFHJnv77bfr5Dm9e6Deueiii0o6xg033JDJ/vOf/yQNjSu0AAAARElBCwAAQJQUtAAAAERJQQsAAECUFLQAAABEqeK7HIc6/dan+uxmXOrXWBcdjUN0M65MiyyySDA/77zzSurEGnLggQcW3Z2yUJdjKFcn1vrqZlxIoc6ZhTptrr322sF8++23z2QTJ06s5ezg/3nxxRczWZcuXYJjhw4dOhdmBMU55JBDgnn//v3r7TGHDBkSzF966aV6e8y/BX5mtGzZMjh23LhxwXzw4MGZbPr06UlD4wotAAAAUVLQAgAAECUFLQAAAFFS0AIAABAlBS0AAABRiq7LcaHOvXO7m3F9dvst9DWOHDkyqS+jRo0K5oMGDaq3x6RhOeWUU4J57969a33sQl31mjdvnsl0OYbSdOvWLZhfddVVmWy33XabCzOiIXvnnXeC+eabb17Scc4999ykIZg2bVq5p0BkOnToUG/HvuOOO4L56aefXm+P2b1796J/NhTqZlxo/48dOzaJgSu0AAAARElBCwAAQJQUtAAAAERJQQsAAECUanK5XK6ogTU1SUMwcODAud4UqlDDpLpoChVqAFWfzZ8KKdT8KfT9LvRvUCivC0Uu03rTUNZ/XQg1ikkddNBBSUPwxBNPBPPTTjstmD/77LNJpSv3+q+0PVDIHnvsEcxvvPHGufp9rat/76lTp2ay9ddfPzh29OjRSUNW7j1QSeu/UGOw4cOHJw1ZoYaBhZoOVpJyr/9K2wOFvpYZM2aUdJwPP/xwrjacKnV9zJw5M5Ndd911wbEHH3xwEvMecIUWAACAKCloAQAAiJKCFgAAgCgpaAEAAIiSghYAAIAoNU4i07179wbR6bcuuhnXd3fmUhSaRygv1PW50NdYaDz1r1evXkV3M66LLoqFOge+8847wXyZZZbJZJtssklw7BJLLBHMu3btGsx//PHHP5gpZL3++uvB/Jhjjin6GHfddVcwX2ihhYp+3tx6662DY7t16xbMmzRpEsybN29e9H5p6F2OqTtvvPFGMH/vvfeC+corr5w0BB988EG5p0CFOPnkk0t6HTRu3LhgfuSRRyb1Yd555w3mF198cdHdjAu9/j766KOTSuQKLQAAAFFS0AIAABAlBS0AAABRUtACAAAQJQUtAAAAUarJFdnatFD30rmtUMfh+uwWXKj7cajjcqFOv5Vk4403nuvdjOuiA29tNJT1Xxf69esXzA888MBg/sADDwTzO++8s+jHfPfdd4vu2vrEE08Ex06bNi2YL7XUUiV1JYxRudd/pe2BWD344IPB/K9//WvRx1h//fWD+UsvvZQ0ZOXeA9Ww/gs9l2677bbBfPnll89kvXv3Do79+uuvg3nHjh2Lnt+UKVOCecuWLZNKV+71H/MeaNq0aSa7+uqrg2MLrd8bbrghmO+///5JfTj//PODeaEOxTUF/m222267ol/Txb4HXKEFAAAgSgpaAAAAoqSgBQAAIEoKWgAAAKIUXVOokSNHBvNqaMZUnwo1vgo1eqrP5k8NtSFCQ1n/lea4447LZIMHDw6OHTNmTDDv2rVrSQ1EYlTu9Z+yBxruz7+NNtqo6GNoCvXnWP+1U6iZzUUXXVT0MTSFKq9Y90CoednYsWNLOkaHDh2C+UcffZTU1iqrrJLJ7r///uDYpZdeOpg/88wzwXz77bfPZJMmTUpipCkUAAAAFUlBCwAAQJQUtAAAAERJQQsAAECUFLQAAABEqXESmY033jiYV3P340IdikMGDhxYr3Oh7nTs2DGT3XLLLcGxDz/8cDDv379/0pD17t27YjopEo9WrVoF8+bNmwfzcePGBfOZM2fWei4LLrhgJrv44ouDY7t161bSsWfMmJHJpk+fXtIxACpNOV5nrLHGGsH8sccey2QLL7xwcOzTTz9dUm1UTVyhBQAAIEoKWgAAAKKkoAUAACBKCloAAACipKAFAAAgStF1OS6kUIevUFff7t27N5iOyKNGjcpkTz31VHCsDsXV5d57781k7du3D4497rjjav14LVq0COZt2rQJ5q1bt85kO+20U3DsgQceGMwXWmihTJbL5YJjt99++2A+ZcqUYA6FXHHFFcF8t912C+YnnnhiMD/33HOL2hepDh06BPOTTjopk2211VZJXXjttdcy2euvv14nx4ZS3H///cH8oosuKvoYhbqQb7rppsH88ccfL/rYVK5TTjml6NcZN9xwQzD/z3/+U+t59O3bt+jXQQ899FDRnwzB/+MKLQAAAFFS0AIAABAlBS0AAABRUtACAAAQJQUtAAAAUaqYLseFhDoDF+oWXI4ux4MGDSqq8zHVJ9ShtFCX46FDhwbzr7/+OpONGDEiOLZQZ9WFF144mIe6udbU1ATHFuooOHPmzEz2/PPPB8dOmDAhmEN9O/PMM4P5DjvsUHSX4xVWWCGYh/ZMof1SyFtvvRXMzzvvvJKOA/WlWbNmtT5GoZ8vhfYc1WXttdcO5ptvvnnRx5g0aVIwnz59ejBv2rRpJltqqaVK6nL8008/ZbLBgweXND9coQUAACBSCloAAACipKAFAAAgSgpaAAAAolTxTaFCjZ4GDBhQb49XqKFTqPnTH42H8ePHFz22TZs2RedrrLFGcGypjWhKUaih0wEHHJDJ7r333nqbB/wZjRqFz/126dJlrs5jxowZwfz4448P5k888UQ9zwjq/ucZ/BmFXgcVamxZF0INoMaOHVvSMY4++uhM9uyzz9ZqXtXIFVoAAACipKAFAAAgSgpaAAAAoqSgBQAAIEoKWgAAAKJU8V2OR44cOVcf76mnngrmuhlTqkMOOSSTnXjiicGxW2+9dTBfdtllM9n6669fb12Ob7rppmB+xx13BPMpU6bU+jEhZjU1NZnsp59+Co7de++9g7luxjR03333XdGvmbp3714nXcih0HNsqWPPPvvsYN6/f/+ij73bbruV9PqI0ngWAAAAIEoKWgAAAKKkoAUAACBKCloAAACipKAFAAAgShXT5bhHjx5z/TFDnYsHDhw41+dB9ZgwYUIwHzZs2FyfC8TqrrvuKqkLZV345JNPgvnzzz+fyS688MLg2NGjR9f5vKCc5plnnlofo0+fPsF8+PDhtT428SvlExwKdZJv2bJlMA91pH/66aeDY3Uzrl+u0AIAABAlBS0AAABRUtACAAAQJQUtAAAAUaqYplChBk31bdCgQXP9MQGonQceeCCYDx06tKSmM6U0nDr++OOD+aefflr0sSFWLVq0COZrrbXWXJ8LlWny5MnB/Oeff85k8847b3Bsq1atSnrMV199NZNts802JR2DuuEKLQAAAFFS0AIAABAlBS0AAABRUtACAAAQJQUtAAAAUarJ5XK5ogbW1CQx6tGjR1HZHxk4cGAdzog/o8hlWm9iXf9UhnKv/5Q9QDXvAeu/foS6gu+www4lHeORRx4J5ltssUVSKcq9/mPeA/vtt18mu+aaa0o6xhlnnBHMhwwZksk+++yzko5N3ewBV2gBAACIkoIWAACAKCloAQAAiJKCFgAAgCgpaAEAAIhSxXc5pjKUu8Of9U81r/+UPUA17wHrn2pe/yl7gHLS5RgAAICKpKAFAAAgSgpaAAAAoqSgBQAAIEoKWgAAAKKkoAUAACBKCloAAACipKAFAAAgSgpaAAAAoqSgBQAAIEoKWgAAAKKkoAUAACBKCloAAACipKAFAAAgSgpaAAAAoqSgBQAAIEo1uVwuV+5JAAAAQKlcoQUAACBKCloAAACipKAFAAAgSgpaAAAAoqSgBQAAIEoKWgAAAKKkoAUAACBKCloAAACipKAFAAAgSlVV0F5//fVJTU1N8umnn5b093r06JF07NixTufSvn37ZN99963TY8Ifsf6pdvYA1cz6p9rZA5WrqgraSvb6668n2267bbLgggsmLVu2zG+8//u//yv3tGCusP6pVmPGjEl22WWXZNlll82v/YUXXjjZaKONkvvvv7/cU4O5YurUqckJJ5yQLLbYYkmLFi2Srl27Jo8//ni5pwVzxSuvvJIcfvjhyaqrrprMO++8yVJLLZXsuuuuyfvvv59Uk8blngC199hjjyXbbLNNsuaaayannHJKMt988yUfffRR8sUXX5R7alDvrH+q2WeffZb8+OOPSZ8+ffIv6CdPnpzceeed+RM8V111VXLwwQeXe4pQr9KrXHfccUdy9NFHJyussEL+KtyWW26ZjBw5Mtlwww3LPT2oV4MHD06ee+65/InN1VdfPfnmm2+Syy67LFlrrbWSF198sc6vLDdUCtrI/fDDD8k+++yTbLXVVvkn9EaNXHSnelj/VLv0hXv6a07p2frOnTsnF154oYKWivbyyy8nw4cPT84777zkuOOOy2fpz4T0RXz//v2T559/vtxThHp1zDHHJDfffHPStGnT2dluu+2WrLbaask555yT3HjjjUk1qOpXf/fee2/+hXB6VrtZs2bJcsstl5x++unJr7/+Ghz/2muvJeuvv37+lpZlllkmufLKK4O3vgwYMCBZfvnl88dccskl80+qaV4f0kX83//+NznzzDPzL+Z//vnnZObMmfXyWFQW659qVwl7IGSeeebJP+7EiRPn2mMSn0pY/+mJzHS9z3nipnnz5skBBxyQvPDCC8nnn39eL49LZaiEPZDOp+kcxWwqvVMhvQX5vffeS6pFVV+hTW9LSW9PTM9upP8dMWJEcuqpp+av+qRn++Y0YcKE/Fnw9L70PfbYI7ntttuSfv365RfR/vvvnx+TvpBOb/N69tln80+uK6+8cvL2228nF110Uf5e9nvuuafgXNK/O378+KLm3apVq6RJkyb5///EE08kCyywQPLll18m22+/ff5x0nvo99577/zjpk/sEGL9U+0qYQ/Mkp7MmTJlSjJp0qTkvvvuSx5++OH8WXqo5PX/xhtvJCuuuGL+58CcunTpkv/v6NGj8wUFVOoeCMnlcvmT/WlRWzVyVWTIkCG59Ev+5JNP8r+fPHlyZkzfvn1zLVu2zP3yyy+zs+7du+f/3gUXXDA7mzp1aq5Tp065RRZZJDdt2rR8NmzYsFyjRo1yzzzzzG+OeeWVV+b//nPPPTc7W3rppXN9+vSZ/ft0TumYYn6NHDly9t9bffXV8/NNfx1xxBG5O++8M//fdNzuu+9eZ9874mf9U+0qcQ/MOe9Zf57OYeedd86NHz++1t8zKkclrv9VV10117Nnz8zXMWbMmPzY9LGhkvdAyLBhw/Ljrrvuuly1qOortOktA7OkTTXS2wG6deuWb6QxduzYZI011pj9540bN0769u07+/fpGZn09+nZmfQWhHXXXTe5/fbb82djVlpppeS7776bPbZnz575/6YNCtJbA0Latm1bdFe+Oef1008/5ZuAHHLIIbO7uu64447JtGnT8l/Haaedlr/1AH7P+qfaVcIemCVtiLPzzjsnX331Vf7KQXrLXLoPoJLXf3pXQnpb5+/Nujsn/XOo5D3we2PHjk0OO+ywZL311ss3C6wWjav94w5OPvnk/C0G6e0Fc0pv25pTen99eivjnNLbXFLp51mlC/mDDz7I36/epk2b4ON9++23BeeSPvn26tXrT2/G9PaHOe255575DZm+h8QLekKsf6pdJeyBWdIXUOmvWU1xNttss3z375deein/uYtQqT8DQu9N/OWXX2b/OVTyHpjTN998k39PcHpL8qz3l1eLqi1o02YZ3bt3z7/vIr2Kk74RPF1M6edZpp9n9mcay6R/J+0qlnaWDPmj93GkZ9PHjRtX1OOkn7U56w3g6QZLN+Siiy76mzGLLLLI7Hv+4fesf6pdpeyBQtKrtenVg/R9Wx06dCjquFSPSln/7dq1y/dQ+L2vv/569s8IqOQ9MGcBvsUWW+S/rmeeeabq1n7VFrSjRo1Kvv/+++Suu+7Kfwj9LJ988klwfHobV9p0Y86zM7M+tLh9+/b5/6ab4c0330w22WSTks+Ip5340o5pxUhvWejRo0f+/6cfzZDeopA+oc/5oiWdb6rQWSKqm/VPtauUPVDIrFstf3+VASpp/Xfq1Cn/+/Tq2pyNodI7E2b9OVTyHph1R8I222yTn0/aLHOVVVZJqk3VFrSzLsOnncBmSd9vdPnllwfHz5gxI38LY9oJbdbY9PfpC+b0RXUq7Xz20EMPJddcc03ms//SFxfpmZvf365Q23vn08dMP2fquuuum32Pfuraa6/N3+//v170UJ2sf6pdpeyB9Ba2WXckzDJ9+vTkhhtuyN9uWY0vbKie9Z/eiXD++ecnV1999ezPoU1vQR4yZEjStWtXHY6p+D2QXtndbbfd8m+xSj+GKH3vbDWq2oI2fVN269at82+YPvLII/NnUoYNG/abhT2n9NL94MGD8/fJp/fM33rrrfl28OmT6KzW2elHhaTNONIGNenZkw022CC/0NI3aKf5o48+mqy99tp1eu/8mmuumW8X/q9//Su/2dLbJ9KzTukb00888cSqu+WA4lj/VLtK2QPpbcXp1an0CsPiiy+efw/VTTfdlH/MCy64IP9RFFCp6z8tWnfZZZf88316cif97M+hQ4fm55me6IRK3wPHHnts/qPa0iu06cf+3Hjjjb/58969eydVIVfF7brT9tnrrrturkWLFrnFFlss179//9yjjz6aaYmdtutOW8O/+uqrufXWWy/XvHnzfLvtyy67LPMYaevuwYMH58c3a9Ys17p161znzp1zgwYNyk2aNKlgu+7aSB9z4MCB+WM2adIkt/zyy+cuuuiiOjk2lcP6p9pV4h645ZZbcr169cotuuiiucaNG+cfL/39vffeW+tjU1kqcf2npkyZkjvuuONybdu2zT/mOuusk3vkkUfq5NhUlkrcA7M+Uigp8Kta1KT/U+6iGgAAAErVqOS/AQAAAA2AghYAAIAoKWgBAACIkoIWAACAKCloAQAAiJKCFgAAgCgpaAEAAIhS42IH1tTU1O9M4A+U++OSrX+qef2n7AGqeQ9Y/1Tz+k/ZAzTkPeAKLQAAAFFS0AIAABAlBS0AAABRUtACAAAQJQUtAAAAUVLQAgAAECUFLQAAAFFS0AIAABAlBS0AAABRalzuCZC12WabBfPDDz88mG+66abBfIMNNshkr7/+ei1nB6Vr3759ML/lllsy2VlnnRUce//999f5vAAAiJsrtAAAAERJQQsAAECUFLQAAABESUELAABAlBS0AAAAREmX4zLr1atXJrv77ruDYz/66KNgvtpqqwXzDz/8sJazg9I0b948mA8bNiyYv/fee5nswQcfrPN5AQCVa/755w/mxx9/fCbbeuutg2PXXHPNYP7tt98G86uuuiqTffXVV8Gx1113XTCfPn16MKc0rtACAAAQJQUtAAAAUVLQAgAAECUFLQAAAFFS0AIAABClmlwulytqYE1N/c+mgi2xxBLB/J133slkzzzzTHDs/vvvH8zHjRuXVLoil2m9sf6Lc9BBBwXzY445Jpivs846meynn36q83nFrtzrP4Y9sPzyywfzZs2aZbLlllsuOHbbbbcN5vvtt1/R85g4cWIwP+OMM4L5TTfdVHRHzWpW7j3Q0Nd/Ka89br311uDY9dZbr6Rjv/LKK5nsgQceKHqdpz7//PNMputrw1v/DWkPhJ7T/+i1c+fOnZOG4Isvvgjmt9xySya79tprg2Or+dNLcv9jD7hCCwAAQJQUtAAAAERJQQsAAECUFLQAAABESUELAABAlHQ5rmOFuqlddtllwXz06NGZrF+/fnU+r9iVu8Of9V9cV9lXX301OPbss88O5oMHD67zeVWicq//cu2B008/PZNtsMEGwbFrr712MJ933nmL/n5OnTo1mN9///3BfPPNN89kCyywQHBsocd88803G2xXzoak3Hsg1p8Bb731ViZbccUVg2M//fTTYL7ooosG80JrvRQjRozIZAcccEBw7H/+85+kWpV7/TekPdC0adNgfueddwbzVVddNZNdcsklJT3mwgsvHMwPOeSQTDb//PMHxzZp0qToxyu0F0M/c6ql+3FOl2MAAAAqkYIWAACAKCloAQAAiJKCFgAAgChpClXHhg0bFsxXWWWVYK75RxwNEaz/rJtvvjmTLb300sGxG220UTD/9ddf63xelajc67++90CowVjq2WefLbo5RyGff/55JhsyZEhw7OTJk4P5Aw88EMxHjRqVydq0aVPSv+H7779f9M+LalbuPRDrz4Dtt9++6KY1Dz30UDBv1apVML/qqqsyWc+ePZPa+u9//xvM99xzz6L3YaUp9/qPeQ/MbVtttVUw//vf/150Q8NCja8K7dEdd9wxmE+fPj2pFJpCAQAAUJEUtAAAAERJQQsAAECUFLQAAABESUELAABAlHQ5LkLr1q2D+f/93/9lsu222y44dsCAAcH8oosuquXsqkO5O/xV8/rfcMMNg/mTTz5ZdHfWjz76qM7nVU3Kvf7LtQcOOuigTHbFFVcExx533HHB/P777y96PRbq5vrWW28F88UXX7zo79Pzzz8fzHv16pXJpk6dGhxbzcq9B6rhZ8BSSy0VzC+44IKiO6v++OOPwbFjxowJ5qGu5SuuuGJw7JQpU4L5kUceGcyvvfbapFKUe/1Xyx4oh1CX7m7dupV0jNVXX72kfRcjXY4BAACoSApaAAAAoqSgBQAAIEoKWgAAAKKkoAUAACBKjcs9gYZknnnmCebXXXddMN9iiy0y2b777hsce+utt9ZydlAehx56aDC/5ZZbMpluxtSlIUOGZLLnnnsuOLbQ2iulY3CzZs2K7mZcqkcffTSY62jM3NakSZNgfvzxxxfdzTj11VdfZbLDDjssOPa+++4ren59+/YN5meddVYwP+GEE4L522+/ncleeumloucBc8P48ePLPYWK4AotAAAAUVLQAgAAECUFLQAAAFFS0AIAABAlTaHmcPXVVwfz7bbbLpifeOKJmUzzJ2K1ww47BPNdd901mPfq1aueZ0S1mzFjRiZ799136+3xvv3222B+0EEHBfPLL7+86MZSG2ywQTBfZJFFip4H1IV+/fqV1ACwkFBjzHfeeSeprauuuiqYL7fccsH82GOPDeaPPPJIJuvUqVNw7GeffVbSHKGufPPNN+WeQkVwhRYAAIAoKWgBAACIkoIWAACAKCloAQAAiJKCFgAAgChVZZfjYcOGBfO99tormJ977rkl5RCj5s2bB/NCXWWfeeaZep4RNAxDhgwJ5quuumom+9vf/hYcu+mmmwbzJ598MpMdccQRwbGjRo36HzOF3zr44IMz2aBBg4JjZ86cGcxPOumkYF6fHcdDTj755GC+4IILBvP99tuvqM7Hf9S1/8svvyxpjlBI6OdFaueddy76GBMnTgzmU6ZMSaqdK7QAAABESUELAABAlBS0AAAARElBCwAAQJQUtAAAAESpJpfL5YoaWFOTxGjFFVfMZK+99lpw7D333BPMDzjggGA+bdq0Ws6OYhW5TOtNrOu/FFdccUVJXR7POOOMep4RDWX9V8seqAsHHnhgSV3xF1hggUw2ffr04NhC3Y+vvfbapNKVew/Euv7feuutorutvvTSS8F8/fXXT2IUek23zTbbBMeeeuqpwfzMM89MGoJyr/+Y98DcNv/88wfzK6+8MpjvvvvuRR/7sssuC+ZHHXVUUun+1x5whRYAAIAoKWgBAACIkoIWAACAKCloAQAAiJKCFgAAgChVfJfjjz76KJMtvvjiwbFrrbVWMH/33XeThtqxObXtttsWfYyRI0cG80KdnxuKcnf4i3X9l2L06NHB/I477qi3LsddunTJZKecckpwbI8ePYL5r7/+Gsw333zzort4NnTlXv/Vsgfq0xJLLBHM//Wvf2Wynj17lnTsgw46KJgPGTIkqRTl3gMNff23adMmmL/88suZbKmllgqO/fvf/x7MzzvvvCRGnTt3Lur7kfr222+D+VZbbZXJXn/99aTa1n8Me6ChuPTSS4P5oYceWvQxbr/99pKe63/88cek0ulyDAAAQEVS0AIAABAlBS0AAABRUtACAAAQpcZJhVtmmWUy2WGHHdagmz/tsssuwfyGG24I5k2bNi362GPGjAnmXbt2DeZTpkwp+tjE7bPPPqu3Y6+99tpFNz6Yd955g2MLNfNYf/31g3mfPn0qpikU8fviiy+C+ZZbbpnJzjrrrODYY445JphffvnlRc+jkhpF8f9ZffXVg3moAdTkyZODY5988smk0hsdnn766cGxhZoR7rDDDg2iKRTlNd988wXzCy64IJPttNNOJR17/PjxmWzgwIFV2/zpz3KFFgAAgCgpaAEAAIiSghYAAIAoKWgBAACIkoIWAACAKEXX5bhZs2bBfNiwYUV3Dxs+fHjSUOa933771apjZeqFF14ouhva5ptvHhy7wAILBHNdjqvHIossEsxXXHHFWq/z8847L5j/+9//zmR77rlncOx3330XzM8888xg3rp16z+YKTQMM2bMyGT9+/cPjl133XWDebdu3YL5tddem8nmmWeeoscSv5qamkz2888/B8dWWvfeX3/9NZNdddVVwbH77rtvMO/Ro0edz4uGq9CnLFx99dXBfLfddiv62KF6JLXXXntlsrFjxxZ9XP4fV2gBAACIkoIWAACAKCloAQAAiJKCFgAAgCgpaAEAAIhSdF2OC3Uu3WmnnYL54MGDM9mECROSuW3nnXcO5v/85z8z2X/+85+iv5bUNddcE8znn3/+TPb999//j5lSre65555g/ve//z2YN2nSJJN16tQpOLZQvs466xTdzbiQp59+Ophvt912JR0HGro99tgjmBf6mZHL5TLZ6aefHhyry3FlCq2BUFYtvv7665JeRw0YMCCTbbrppsGxjz/+eC1nx9zSsmXLkp4Hd91113rpZpx67LHHij42hblCCwAAQJQUtAAAAERJQQsAAECUFLQAAABEKbqmUKU677zz6u3YzZo1y2SXX355cOwuu+wSzEMNOi655JLg2ELNrDbbbLNgfvXVV2eyESNGBMdqFsWHH34YzFu1ahXMt95660zWvHnz4NiXX365pMcsxQ477BDMZ86cWetjQ0Py5Zdf1voYLVq0CObLLbdcMP/oo49q/Zg0LAsttFAw32KLLYL5ww8/nFS6jz/+OJjPM888mezEE08MjtUUqmEKPef961//Kum1eiGh1+V77rlncKz1Ub9coQUAACBKCloAAACipKAFAAAgSgpaAAAAoqSgBQAAIErRdTku1On39ddfD+adO3fOZE888USdzKV79+6ZbN999w2OnTRpUjC/6qqriv4aW7duXXQ349T06dMz2YABA4JjZ8yYEcypHm+//XYwHz9+fDA//vjjM9l1112X1JemTZsG83322afo+UHMFl988Vofo0mTJsG8TZs2wVyX4ziMGjUqmL/77ruZbJVVVgmObdu2bZ3Pq5q6RNMwbbTRRrXuZjxx4sRgvscee2Qy3YzLwxVaAAAAoqSgBQAAIEoKWgAAAKKkoAUAACBKCloAAACiFF2X46lTp5bUibU+uxyHuqh+/vnnwbFrrrlmMA91NO7Tp09w7HHHHRfMF1100WC+9957Z7LnnnsuOBbef//9YP7www8H87322iuTNW/ePDh23LhxtZxduJvgHx372muvrfVjQkPqZvzII4/U+tiTJ08uuhsu8fj111+DeS6XK/oYBx10UDAfMmRIUul0eI5foU8Cue2222p97L///e/BXEfjhsMVWgAAAKKkoAUAACBKCloAAACipKAFAAAgSgpaAAAAohRdl+NCrrnmmmB+4403ZrJvv/02OHb48OHBfJVVVgnmm2yySSbbaqutgmPnn3/+YP7kk09msk6dOgXHPvvss8F8tdVWC+YffvhhMIdS/POf/wzm22yzTdFrd+bMmcF80KBBmaxZs2bBsf369Qvmp556akkd0aku7dq1y2SNG4d/9BXqUl+fQnumUFfO5ZZbLpg3atSo6H132mmnBcf+8MMP/2OmxOjss8/OZEOHDg2O7dixYzDfYYcdgvndd9+dVIpdd9216LF10TWXulfoeXC++ear9bG7dOkSzCdNmtSg10eHDh0yWc+ePYNj//a3vwXzm2++OZMNHDgwaWhcoQUAACBKCloAAACipKAFAAAgSgpaAAAAolSTy+VyRQ2sqUliFGo6U+iNz19++WUwb9GiRTBfaqmlMtmYMWOCY//yl78E8+eee67o5lQPPPBAMJ8xY0ZS6YpcpvUm1vVfn/bbb79MduGFFwbHtmrVqujva6F/60J7q1BTtEpS7vUfwx5Yfvnlg/nIkSMz2aeffhoc26tXr1o3GCvUXKZQQ6dQU7MmTZokdfFvE1o3oSZZqXHjxiUNWbn3QENf/6U455xzgvkxxxxT0muMCy64oOjXKa+++mow//XXX5P6EnrdVagpWt++fYP5zz//nMkWX3zx4NgpU6Yklbr+Y9gDhZpCHXjggZnsiiuuqJPHDK3fhtRgr2nTppls3nnnLekYF198cSY79thjk7ntf+0BV2gBAACIkoIWAACAKCloAQAAiJKCFgAAgCgpaAEAAIhSxXc5DunatWsw79evX0nHCXWt/Oijj4JjH3zwwWA+YsSITPb999+XNI9qUO4Of5W0/svRafaAAw4I5gcddFAmmzhxYkn7thr2S7nXfwx74Omnnw7mG2ywQdHfz/PPPz+Yt27dOphvt912maxNmzZz/d8w1Ik1dffdd2ey/fffPzh25syZSUNW7j3Q0Nd/Xdhjjz2C+dVXXx3MW7ZsWfSx//nPfxZ97OnTpwfH/vvf/w7mO+ywQzA/5ZRTMtkaa6wRHDtp0qRgPmDAgEx26aWXJtW2/mPeA6Hux6HOx3XZ/biSXKzLMQAAANQfBS0AAABRUtACAAAQJQUtAAAAUVLQAgAAEKWq7HJMfMrd4c/6p5rXfwx7INR1PjVq1KhM1q5du7n+faqLf8MhQ4YE83POOSeYF+q6H6Ny74GGvv7r09prrx3M+/fvn8l22mmnWj/e1KlTg/nzzz8fzLt06RLM55133qIfc9dddw3md955Z9IQlHv9V9oeKPS1LLroorX+FJTFF188mO+33361fq7/8ssvk9oaPXp0MH/ggQeK7oD/66+/JnObLscAAABUJAUtAAAAUVLQAgAAECUFLQAAAFHSFIoolLshgvVPNa//mPfAKqusksnOP//84NjNNtus1o83ceLEYH7GGWcE88cee6zoYxdq8lSoiU4lKfceiHX916dGjbLXRJZaaqmSmrZtvfXWRT/eNttsE8wnT54czJ988slM9sgjjxTdPK4h7a1yr/+UPUA5aQoFAABARVLQAgAAECUFLQAAAFFS0AIAABAlBS0AAABR0uWYKJS7w5/1TzWv/5Q9QDXvAeufal7/KXuActLlGAAAgIqkoAUAACBKCloAAACipKAFAAAgSgpaAAAAoqSgBQAAIEoKWgAAAKKkoAUAACBKCloAAACipKAFAAAgSgpaAAAAoqSgBQAAIEoKWgAAAKKkoAUAACBKCloAAACipKAFAAAgSgpaAAAAoqSgBQAAIEoKWgAAAKKkoAUAACBKCloAAACipKAFAAAgSjW5XC5X7kkAAABAqVyhBQAAIEoKWgAAAKKkoAUAACBKCloAAACipKAFAAAgSgpaAAAAoqSgBQAAIEoKWgAAAKKkoAUAACCJ0f8PVuZNhiYH2eMAAAAASUVORK5CYII=\",\n \"text/plain\": [\n \"
\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"test_batch = test_ds.as_numpy_iterator().next()\\n\",\n \"pred = pred_step(eval_model, test_batch)\\n\",\n \"\\n\",\n \"fig, axs = plt.subplots(5, 5, figsize=(12, 12))\\n\",\n \"for i, ax in enumerate(axs.flatten()):\\n\",\n \" ax.imshow(test_batch['image'][i, ..., 0], cmap='gray')\\n\",\n \" ax.set_title(f'label={pred[i]}')\\n\",\n \" ax.axis('off')\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"65342ab4\",\n \"metadata\": {},\n \"source\": [\n \"# 8. Export the model\\n\",\n \"\\n\",\n \"Flax models are great for research, but aren't meant to be deployed directly. Instead, high performance inference runtimes like LiteRT or TensorFlow Serving operate on a special [SavedModel](https://www.tensorflow.org/guide/saved_model) format. The [Orbax](https://orbax.readthedocs.io/en/latest/guides/export/orbax_export_101.html) library makes it easy to export Flax models to this format. First, we must create a `JaxModule` object wrapping a model and its prediction method.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"49cace09\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"from orbax.export import JaxModule, ExportManager, ServingConfig\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"421309d4\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"def exported_predict(model, y):\\n\",\n \" return model(y, None)\\n\",\n \"\\n\",\n \"jax_module = JaxModule(eval_model, exported_predict)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"787136af\",\n \"metadata\": {},\n \"source\": [\n \"We also need to tell Tensorflow Serving what input type `exported_predict` expects in its second argument. The export machinery expects type signature arguments to be PyTrees of `tf.TensorSpec`.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"9f2ad72e\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"sig = [tf.TensorSpec(shape=(1, 28, 28, 1), dtype=tf.float32)]\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"31e9668a\",\n \"metadata\": {},\n \"source\": [\n \"Finally, we can bundle up the input signature and the `JaxModule` together using the `ExportManager` class.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"id\": \"18cdf9ad\",\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"export_mgr = ExportManager(jax_module, [\\n\",\n \" ServingConfig('mnist_server', input_signature=sig)\\n\",\n \"])\\n\",\n \"\\n\",\n \"output_dir='/tmp/mnist_export'\\n\",\n \"export_mgr.save(output_dir)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"id\": \"28\",\n \"metadata\": {},\n \"source\": [\n \"Congratulations! You have learned how to use Flax NNX to build and train a simple classification model end-to-end on the MNIST dataset.\\n\",\n \"\\n\",\n \"Next, check out [Why Flax NNX?](https://flax.readthedocs.io/en/latest/why.html) and get started with a series of [Flax NNX Guides](https://flax.readthedocs.io/en/latest/guides/index.html).\"\n ]\n }\n ],\n \"metadata\": {\n \"accelerator\": \"GPU\",\n \"jupytext\": {\n \"formats\": \"ipynb,md:myst\",\n \"main_language\": \"python\"\n },\n \"language_info\": {\n \"codemirror_mode\": {\n \"name\": \"ipython\",\n \"version\": 3\n },\n \"file_extension\": \".py\",\n \"mimetype\": \"text/x-python\",\n \"name\": \"python\",\n \"nbconvert_exporter\": \"python\",\n \"pygments_lexer\": \"ipython3\",\n \"version\": \"3.11.9\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 5\n}\n" }, { "path": "docs_nnx/mnist_tutorial.md", "content": "---\njupytext:\n formats: ipynb,md:myst\n main_language: python\n text_representation:\n extension: .md\n format_name: myst\n format_version: 0.13\n jupytext_version: 1.13.8\n---\n\n[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google/flax/blob/main/docs_nnx/mnist_tutorial.ipynb)\n[![Open On GitHub](https://img.shields.io/badge/Open-on%20GitHub-blue?logo=GitHub)](https://github.com/google/flax/blob/main/docs_nnx/mnist_tutorial.ipynb)\n\n# MNIST tutorial\n\nWelcome to Flax NNX! In this tutorial you will learn how to build and train a simple convolutional neural network (CNN) to classify handwritten digits on the MNIST dataset using the Flax NNX API.\n\nFlax NNX is a Python neural network library built upon [JAX](https://github.com/jax-ml/jax). If you have used the Flax Linen API before, check out [Why Flax NNX](https://flax.readthedocs.io/en/latest/why.html). You should have some knowledge of the main concepts of deep learning.\n\nLet’s get started!\n\n## 1. Install Flax\n\nIf `flax` is not installed in your Python environment, use `pip` to install the package from PyPI (below, just uncomment the code in the cell if you are working from Google Colab/Jupyter Notebook):\n\n```{code-cell} ipython3\n# !pip install -U \"jax[cuda12]\"\n# !pip install -U flax\n```\n\n## 2. Load the MNIST dataset\n\nFirst, you need to load the MNIST dataset and then prepare the training and testing sets via Tensorflow Datasets (TFDS). You normalize image values, shuffle the data and divide it into batches, and prefetch samples to enhance performance.\n\n```{code-cell} ipython3\nimport tensorflow_datasets as tfds # TFDS to download MNIST.\nimport tensorflow as tf # TensorFlow / `tf.data` operations.\n\ntf.random.set_seed(0) # Set the random seed for reproducibility.\n\ntrain_steps = 1200\neval_every = 200\nbatch_size = 32\n\ntrain_ds: tf.data.Dataset = tfds.load('mnist', split='train')\ntest_ds: tf.data.Dataset = tfds.load('mnist', split='test')\n\ntrain_ds = train_ds.map(\n lambda sample: {\n 'image': tf.cast(sample['image'], tf.float32) / 255,\n 'label': sample['label'],\n }\n) # normalize train set\ntest_ds = test_ds.map(\n lambda sample: {\n 'image': tf.cast(sample['image'], tf.float32) / 255,\n 'label': sample['label'],\n }\n) # Normalize the test set.\n\n# Create a shuffled dataset by allocating a buffer size of 1024 to randomly draw elements from.\ntrain_ds = train_ds.repeat().shuffle(1024)\n# Group into batches of `batch_size` and skip incomplete batches, prefetch the next sample to improve latency.\ntrain_ds = train_ds.batch(batch_size, drop_remainder=True).take(train_steps).prefetch(1)\n# Group into batches of `batch_size` and skip incomplete batches, prefetch the next sample to improve latency.\ntest_ds = test_ds.batch(batch_size, drop_remainder=True).prefetch(1)\n```\n\n## 3. Define the model with Flax NNX\n\nCreate a CNN for classification with Flax NNX by subclassing `nnx.Module`:\n\n```{code-cell} ipython3\nfrom flax import nnx # The Flax NNX API.\nfrom functools import partial\nfrom typing import Optional\n\nclass CNN(nnx.Module):\n \"\"\"A simple CNN model.\"\"\"\n\n def __init__(self, *, rngs: nnx.Rngs):\n self.conv1 = nnx.Conv(1, 32, kernel_size=(3, 3), rngs=rngs)\n self.batch_norm1 = nnx.BatchNorm(32, rngs=rngs)\n self.dropout1 = nnx.Dropout(rate=0.025)\n self.conv2 = nnx.Conv(32, 64, kernel_size=(3, 3), rngs=rngs)\n self.batch_norm2 = nnx.BatchNorm(64, rngs=rngs)\n self.avg_pool = partial(nnx.avg_pool, window_shape=(2, 2), strides=(2, 2))\n self.linear1 = nnx.Linear(3136, 256, rngs=rngs)\n self.dropout2 = nnx.Dropout(rate=0.025)\n self.linear2 = nnx.Linear(256, 10, rngs=rngs)\n\n def __call__(self, x, rngs: nnx.Rngs | None = None):\n x = self.avg_pool(nnx.relu(self.batch_norm1(self.dropout1(self.conv1(x), rngs=rngs))))\n x = self.avg_pool(nnx.relu(self.batch_norm2(self.conv2(x))))\n x = x.reshape(x.shape[0], -1) # flatten\n x = nnx.relu(self.dropout2(self.linear1(x), rngs=rngs))\n x = self.linear2(x)\n return x\n\n# Instantiate the model.\nmodel = CNN(rngs=nnx.Rngs(0))\n# Visualize it.\nnnx.display(model)\n```\n\n### Run the model\n\nLet's put the CNN model to the test! Here, you’ll perform a forward pass with arbitrary data and print the results.\n\n```{code-cell} ipython3\nimport jax.numpy as jnp # JAX NumPy\n\ny = model(jnp.ones((1, 28, 28, 1)), nnx.Rngs(0))\ny\n```\n\n## 4. Create the optimizer and define some metrics\n\nIn Flax NNX, you need to create an `nnx.Optimizer` object to manage the model's parameters and apply gradients during training. `nnx.Optimizer` receives the model's reference, so that it can update its parameters, and an [Optax](https://optax.readthedocs.io/) optimizer to define the update rules. Additionally, you will define an `nnx.MultiMetric` object to keep track of the `Accuracy` and the `Average` loss.\n\n```{code-cell} ipython3\nimport optax\n\nlearning_rate = 0.005\nmomentum = 0.9\n\noptimizer = nnx.Optimizer(\n model, optax.adamw(learning_rate, momentum), wrt=nnx.Param\n)\nmetrics = nnx.MultiMetric(\n accuracy=nnx.metrics.Accuracy(),\n loss=nnx.metrics.Average('loss'),\n)\n\nnnx.display(optimizer)\n```\n\n## 5. Define training step functions\n\nIn this section, you will define a loss function using the cross entropy loss ([`optax.softmax_cross_entropy_with_integer_labels()`](https://optax.readthedocs.io/en/latest/api/losses.html#optax.softmax_cross_entropy_with_integer_labels)) that the CNN model will optimize over.\n\nIn addition to the `loss`, during training and testing you will also get the `logits`, which will be used to calculate the accuracy metric.\n\nDuring training — the `train_step` — you will use `nnx.value_and_grad` to compute the gradients and update the model's parameters using the `optimizer` you have already defined. The `train_step` also receives an `nnx.Rngs` object to provide randomness for dropout. The `eval_step` omits `rngs` because the eval view already has `deterministic=True`, so dropout is disabled and no random key is needed. During both steps, the `loss` and `logits` are used to update the metrics.\n\n```{code-cell} ipython3\ndef loss_fn(model: CNN, batch, rngs: nnx.Rngs | None = None):\n logits = model(batch['image'], rngs)\n loss = optax.softmax_cross_entropy_with_integer_labels(\n logits=logits, labels=batch['label']\n ).mean()\n return loss, logits\n\n@nnx.jit\ndef train_step(model: CNN, optimizer: nnx.Optimizer, metrics: nnx.MultiMetric, rngs: nnx.Rngs, batch):\n \"\"\"Train for a single step.\"\"\"\n grad_fn = nnx.value_and_grad(loss_fn, has_aux=True)\n (loss, logits), grads = grad_fn(model, batch, rngs)\n metrics.update(loss=loss, logits=logits, labels=batch['label']) # In-place updates.\n optimizer.update(model, grads) # In-place updates.\n\n@nnx.jit\ndef eval_step(model: CNN, metrics: nnx.MultiMetric, batch):\n loss, logits = loss_fn(model, batch)\n metrics.update(loss=loss, logits=logits, labels=batch['label']) # In-place updates.\n```\n\nIn the code above, the [`nnx.jit`](https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/transforms.html#flax.nnx.jit) transformation decorator traces the `train_step` function for just-in-time compilation with [XLA](https://www.tensorflow.org/xla), optimizing performance on hardware accelerators, such as Google TPUs and GPUs. `nnx.jit` is a stateful version of the `jax.jit` transform that allows its function input and outputs to be Flax NNX objects. Similarly, `nnx.value_and_grad` is a stateful version of `jax.value_and_grad`. Check out [the transforms guide](https://flax.readthedocs.io/en/latest/guides/transforms.html) to learn more.\n\n> **Note:** The code shows how to perform several in-place updates to the model, the optimizer, the RNG streams, and the metrics, but _state updates_ were not explicitly returned. This is because Flax NNX transformations respect _reference semantics_ for Flax NNX objects, and will propagate the state updates of the objects passed as input arguments. This is a key feature of Flax NNX that allows for a more concise and readable code. You can learn more in [Why Flax NNX](https://flax.readthedocs.io/en/latest/why.html).\n\n\n## 6. Train and evaluate the model\n\nNow, you can train the CNN model. Before the training loop, we use [`nnx.view`](https://flax.readthedocs.io/en/latest/guides/view.html) to create a `train_model` (with dropout enabled and batch norm in training mode) and an `eval_model` (with dropout disabled and batch norm using running statistics). These views share the same underlying weights, so updates during training are automatically reflected during evaluation.\n\n```{code-cell} ipython3\nfrom IPython.display import clear_output\nimport matplotlib.pyplot as plt\n\nmetrics_history = {\n 'train_loss': [],\n 'train_accuracy': [],\n 'test_loss': [],\n 'test_accuracy': [],\n}\n\nrngs = nnx.Rngs(0)\ntrain_model = nnx.view(model, deterministic=False, use_running_average=False)\neval_model = nnx.view(model, deterministic=True, use_running_average=True)\n\nfor step, batch in enumerate(train_ds.as_numpy_iterator()):\n # Run the optimization for one step and make a stateful update to the following:\n # - The train state's model parameters\n # - The optimizer state\n # - The training loss and accuracy batch metrics\n train_step(train_model, optimizer, metrics, rngs, batch)\n\n if step > 0 and (step % eval_every == 0 or step == train_steps - 1): # One training epoch has passed.\n # Log the training metrics.\n for metric, value in metrics.compute().items(): # Compute the metrics.\n metrics_history[f'train_{metric}'].append(value) # Record the metrics.\n metrics.reset() # Reset the metrics for the test set.\n\n # Compute the metrics on the test set after each training epoch.\n for test_batch in test_ds.as_numpy_iterator():\n eval_step(eval_model, metrics, test_batch)\n\n # Log the test metrics.\n for metric, value in metrics.compute().items():\n metrics_history[f'test_{metric}'].append(value)\n metrics.reset() # Reset the metrics for the next training epoch.\n\n clear_output(wait=True)\n # Plot loss and accuracy in subplots\n fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))\n ax1.set_title('Loss')\n ax2.set_title('Accuracy')\n for dataset in ('train', 'test'):\n ax1.plot(metrics_history[f'{dataset}_loss'], label=f'{dataset}_loss')\n ax2.plot(metrics_history[f'{dataset}_accuracy'], label=f'{dataset}_accuracy')\n ax1.legend()\n ax2.legend()\n plt.show()\n```\n\n## 7. Perform inference on the test set\n\nCreate a `jit`-compiled model inference function (with `nnx.jit`) - `pred_step` - to generate predictions on the test set using the learned model parameters. Since we already have `eval_model` (an `nnx.view` with `deterministic=True` and `use_running_average=True`), we can use it directly for inference. This will enable you to visualize test images alongside their predicted labels for a qualitative assessment of model performance.\n\n```{code-cell} ipython3\n@nnx.jit\ndef pred_step(model: CNN, batch):\n logits = model(batch['image'], None)\n return logits.argmax(axis=1)\n```\n\nWe reuse the `eval_model` view created earlier so that `Dropout` is disabled and `BatchNorm` uses stored running stats during inference.\n\n```{code-cell} ipython3\ntest_batch = test_ds.as_numpy_iterator().next()\npred = pred_step(eval_model, test_batch)\n\nfig, axs = plt.subplots(5, 5, figsize=(12, 12))\nfor i, ax in enumerate(axs.flatten()):\n ax.imshow(test_batch['image'][i, ..., 0], cmap='gray')\n ax.set_title(f'label={pred[i]}')\n ax.axis('off')\n```\n\n# 8. Export the model\n\nFlax models are great for research, but aren't meant to be deployed directly. Instead, high performance inference runtimes like LiteRT or TensorFlow Serving operate on a special [SavedModel](https://www.tensorflow.org/guide/saved_model) format. The [Orbax](https://orbax.readthedocs.io/en/latest/guides/export/orbax_export_101.html) library makes it easy to export Flax models to this format. First, we must create a `JaxModule` object wrapping a model and its prediction method.\n\n```{code-cell} ipython3\nfrom orbax.export import JaxModule, ExportManager, ServingConfig\n```\n\n```{code-cell} ipython3\ndef exported_predict(model, y):\n return model(y, None)\n\njax_module = JaxModule(eval_model, exported_predict)\n```\n\nWe also need to tell Tensorflow Serving what input type `exported_predict` expects in its second argument. The export machinery expects type signature arguments to be PyTrees of `tf.TensorSpec`.\n\n```{code-cell} ipython3\nsig = [tf.TensorSpec(shape=(1, 28, 28, 1), dtype=tf.float32)]\n```\n\nFinally, we can bundle up the input signature and the `JaxModule` together using the `ExportManager` class.\n\n```{code-cell} ipython3\nexport_mgr = ExportManager(jax_module, [\n ServingConfig('mnist_server', input_signature=sig)\n])\n\noutput_dir='/tmp/mnist_export'\nexport_mgr.save(output_dir)\n```\n\nCongratulations! You have learned how to use Flax NNX to build and train a simple classification model end-to-end on the MNIST dataset.\n\nNext, check out [Why Flax NNX?](https://flax.readthedocs.io/en/latest/why.html) and get started with a series of [Flax NNX Guides](https://flax.readthedocs.io/en/latest/guides/index.html).\n" }, { "path": "docs_nnx/nnx_basics.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"# Flax basics\\n\",\n \"\\n\",\n \"Flax NNX is a new simplified API that is designed to make it easier to create, inspect, debug, and analyze neural networks in [JAX](https://jax.readthedocs.io/). It achieves this by adding first class support for Python reference semantics. This allows users to express their models using regular Python objects, which are modeled as PyGraphs (instead of pytrees), enabling reference sharing and mutability. Such API design should make PyTorch or Keras users feel at home.\\n\",\n \"\\n\",\n \"To begin, install Flax with `pip` and import necessary dependencies:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 1,\n \"metadata\": {\n \"tags\": [\n \"skip-execution\"\n ]\n },\n \"outputs\": [],\n \"source\": [\n \"# ! pip install -U flax\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 1,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"from flax import nnx\\n\",\n \"import jax\\n\",\n \"import jax.numpy as jnp\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## The Flax NNX Module system\\n\",\n \"\\n\",\n \"The main difference between the Flax `Module` and other Module systems in [Flax Linen](https://flax-linen.readthedocs.io/en/latest/api_reference/flax.linen/module.html) or [Haiku](https://dm-haiku.readthedocs.io/en/latest/notebooks/basics.html#Built-in-Haiku-nets-and-nested-modules) is that in NNX everything is **explicit**. This means, among other things, that the NNX Module itself holds the state (such as parameters) directly, the [PRNG](https://jax.readthedocs.io/en/latest/random-numbers.html) state is threaded by the user, and all shape information must be provided on initialization (no shape inference).\\n\",\n \"\\n\",\n \"Let's begin by creating a Linear `Module`. As shown next, dynamic state is usually stored in `Param`s, and static state (all types not handled by NNX) such as integers or strings are stored directly. Attributes of type `jax.Array` and `numpy.ndarray` are also treated as dynamic states, although storing them inside Variables, such as Param, is preferred. Also the `Rngs` object can be used to get new unique keys based on a root PRNG key passed to the constructor.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 3,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"class Linear(nnx.Module):\\n\",\n \" def __init__(self, din: int, dout: int, *, rngs: nnx.Rngs):\\n\",\n \" self.w = nnx.Param(rngs.params.uniform((din, dout)))\\n\",\n \" self.b = nnx.Param(jnp.zeros((dout,)))\\n\",\n \" self.din, self.dout = din, dout\\n\",\n \"\\n\",\n \" def __call__(self, x: jax.Array):\\n\",\n \" return x @ self.w + self.b[None]\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Also note that the inner values of `Variable`s can be accessed using the `value` property, but for convenience they implement all numeric operators and can be used directly in arithmetic expressions (as shown in the code above).\\n\",\n \"\\n\",\n \"To initialize a Flax `Module`, you just call the constructor, and all the parameters of a Module are usually created eagerly. Since Modules hold their own state methods, you can call them directly without the need for a separate apply method.\\n\",\n \"This can be very convenient for debugging, allowing you to directly inspect the entire structure of the model.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 3,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"[[1.5643291 0.94782424 0.37971854 1.0724319 0.22112393]]\\n\"\n ]\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"model = Linear(2, 5, rngs=nnx.Rngs(params=0))\\n\",\n \"y = model(x=jnp.ones((1, 2)))\\n\",\n \"\\n\",\n \"print(y)\\n\",\n \"nnx.display(model)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"The above visualization by `nnx.display` is generated using the awesome\\n\",\n \"[Treescope](https://treescope.readthedocs.io/en/stable/index.html#) library.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Stateful computation\\n\",\n \"\\n\",\n \"Implementing layers, such as `BatchNorm`, requires performing state updates during a forward pass. In Flax NNX, you just need to create a `Variable` and update its `.value` during the forward pass.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 4,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"counter.count[...] = Array(0, dtype=int32, weak_type=True)\\n\",\n \"counter.count[...] = Array(1, dtype=int32, weak_type=True)\\n\"\n ]\n }\n ],\n \"source\": [\n \"class Count(nnx.Variable): pass\\n\",\n \"\\n\",\n \"class Counter(nnx.Module):\\n\",\n \" def __init__(self):\\n\",\n \" self.count = Count(jnp.array(0))\\n\",\n \"\\n\",\n \" def __call__(self):\\n\",\n \" self.count[...] += 1\\n\",\n \"\\n\",\n \"counter = Counter()\\n\",\n \"print(f'{counter.count[...] = }')\\n\",\n \"counter()\\n\",\n \"print(f'{counter.count[...] = }')\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Mutable references are usually avoided in JAX. But Flax NNX provides sound mechanisms\\n\",\n \"to handle them, as demonstrated in later sections of this guide.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Nested Modules\\n\",\n \"\\n\",\n \"Flax `Module`s can be used to compose other Modules in a nested structure. These can be assigned directly as attributes, or inside an attribute of any (nested) pytree type, such as a `list`, `dict`, `tuple`, and so on.\\n\",\n \"\\n\",\n \"The example below shows how to define a simple `MLP` by subclassing `Module`. The model consists of two `Linear` layers, a `Dropout` layer, and a `BatchNorm` layer. Note that we need to pass the `__call__` method the RNG state that we want the `Dropout` layer to use.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 5,\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"class MLP(nnx.Module):\\n\",\n \" def __init__(self, din: int, dmid: int, dout: int, *, rngs: nnx.Rngs):\\n\",\n \" self.linear1 = Linear(din, dmid, rngs=rngs)\\n\",\n \" self.dropout = nnx.Dropout(rate=0.1)\\n\",\n \" self.bn = nnx.BatchNorm(dmid, rngs=rngs)\\n\",\n \" self.linear2 = Linear(dmid, dout, rngs=rngs)\\n\",\n \"\\n\",\n \" def __call__(self, x: jax.Array, rngs: nnx.Rngs):\\n\",\n \" x = nnx.gelu(self.dropout(self.bn(self.linear1(x)), rngs=rngs))\\n\",\n \" return self.linear2(x)\\n\",\n \"\\n\",\n \"model = MLP(2, 16, 5, rngs=nnx.Rngs(0))\\n\",\n \"\\n\",\n \"y = model(x=jnp.ones((3, 2)), rngs=nnx.Rngs(1))\\n\",\n \"\\n\",\n \"nnx.display(model)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Model surgery\\n\",\n \"\\n\",\n \"Flax `Module`s are mutable by default. This means that their structure can be changed at any time, which makes [model surgery](https://flax.readthedocs.io/en/latest/guides/surgery.html) quite easy, as any sub-Module attribute can be replaced with anything else, such as new Modules, existing shared Modules, Modules of different types, and so on. Moreover, `Variable`s can also be modified or replaced/shared.\\n\",\n \"\\n\",\n \"The following example shows how to replace the `Linear` layers in the `MLP` model from the previous example with `LoraLinear` layers:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 7,\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"class LoraParam(nnx.Param): pass\\n\",\n \"\\n\",\n \"class LoraLinear(nnx.Module):\\n\",\n \" def __init__(self, linear: Linear, rank: int, rngs: nnx.Rngs):\\n\",\n \" self.linear = linear\\n\",\n \" self.A = LoraParam(rngs.normal((linear.din, rank)))\\n\",\n \" self.B = LoraParam(rngs.normal((rank, linear.dout)))\\n\",\n \"\\n\",\n \" def __call__(self, x: jax.Array):\\n\",\n \" return self.linear(x) + x @ self.A @ self.B\\n\",\n \"\\n\",\n \"rngs = nnx.Rngs(0)\\n\",\n \"model = MLP(2, 32, 5, rngs=rngs)\\n\",\n \"\\n\",\n \"# Model surgery.\\n\",\n \"model.linear1 = LoraLinear(model.linear1, 4, rngs=rngs)\\n\",\n \"model.linear2 = LoraLinear(model.linear2, 4, rngs=rngs)\\n\",\n \"\\n\",\n \"y = model(x=jnp.ones((3, 2)), rngs=rngs)\\n\",\n \"\\n\",\n \"nnx.display(model)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Flax transformations\\n\",\n \"\\n\",\n \"[Flax NNX transformations (transforms)](https://flax.readthedocs.io/en/latest/guides/jax_and_nnx_transforms.html) extend [JAX transforms](https://jax.readthedocs.io/en/latest/key-concepts.html#transformations) to support `Module`s and other objects. They serve as supersets of their equivalent JAX counterparts with the addition of being aware of the object's state and providing additional APIs to transform it.\\n\",\n \"\\n\",\n \"One of the main features of Flax Transforms is the preservation of reference semantics, meaning that any mutation of the object graph that occurs inside the transform is propagated outside as long as it is legal within the transform rules. In practice this means that Flax programs can be expressed using imperative code, highly simplifying the user experience.\\n\",\n \"\\n\",\n \"In the following example, you define a `train_step` function that takes a `MLP` model, an `Optimizer`, and a batch of data, and returns the loss for that step. The loss and the gradients are computed using the `nnx.value_and_grad` transform over the `loss_fn`. The gradients are passed to the optimizer's `update` method to update the model's parameters.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 8,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"loss = Array(1.0000616, dtype=float32)\\n\",\n \"optimizer.step.value = Array(1, dtype=uint32)\\n\"\n ]\n }\n ],\n \"source\": [\n \"import optax\\n\",\n \"\\n\",\n \"# An MLP containing 2 custom `Linear` layers, 1 `nnx.Dropout` layer, 1 `nnx.BatchNorm` layer.\\n\",\n \"model = MLP(2, 16, 10, rngs=nnx.Rngs(0))\\n\",\n \"optimizer = nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param)\\n\",\n \"\\n\",\n \"@nnx.jit # Automatic state management\\n\",\n \"def train_step(model, optimizer, x, y, rngs):\\n\",\n \" def loss_fn(model: MLP, rngs: nnx.Rngs):\\n\",\n \" y_pred = model(x, rngs)\\n\",\n \" return jnp.mean((y_pred - y) ** 2)\\n\",\n \"\\n\",\n \" loss, grads = nnx.value_and_grad(loss_fn)(model, rngs)\\n\",\n \" optimizer.update(model, grads) # In place updates.\\n\",\n \"\\n\",\n \" return loss\\n\",\n \"\\n\",\n \"x, y = jnp.ones((5, 2)), jnp.ones((5, 10))\\n\",\n \"loss = train_step(model, optimizer, x, y, rngs)\\n\",\n \"\\n\",\n \"print(f'{loss = }')\\n\",\n \"print(f'{optimizer.step.value = }')\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"There are two things happening in this example that are worth mentioning:\\n\",\n \"\\n\",\n \"1. The updates to each of the `BatchNorm` and `Dropout` layer's state is automatically propagated from within `loss_fn` to `train_step` all the way to the `model` reference outside.\\n\",\n \"2. The `optimizer` holds a mutable reference to the model - this relationship is preserved inside the train_step function making it possible to update the model's parameters using the optimizer alone.\\n\",\n \"\\n\",\n \"> **Note**
`nnx.jit` has performance overhead for small models, check the [Performance Considerations](https://flax.readthedocs.io/en/latest/guides/performance.html) guide for more information.\\n\",\n \"\\n\",\n \"### Scan over layers\\n\",\n \"\\n\",\n \"The next example uses Flax `nnx.vmap` to create a stack of multiple MLP layers and `nnx.scan` to iteratively apply each layer of the stack to the input.\\n\",\n \"\\n\",\n \"In the code below notice the following:\\n\",\n \"\\n\",\n \"1. The custom `create_model` function takes in a key and returns an `MLP` object, since you create five keys and use `nnx.vmap` over `create_model` a stack of 5 `MLP` objects is created.\\n\",\n \"2. The `nnx.scan` is used to iteratively apply each `MLP` in the stack to the input `x`.\\n\",\n \"3. The nnx.scan (consciously) deviates from `jax.lax.scan` and instead mimics nnx.vmap, which is more expressive. nnx.scan allows specifying multiple inputs, the scan axes of each input/output, and the position of the carry.\\n\",\n \"4. `State` updates for `BatchNorm` layers are automatically propagated by nnx.scan.\\n\",\n \"5. The `rngs` object is split into separate streams for each layer using the `fork` method.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 8,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"@nnx.vmap(in_axes=0, out_axes=0)\\n\",\n \"def create_model(rngs):\\n\",\n \" return MLP(10, 32, 10, rngs=rngs)\\n\",\n \"\\n\",\n \"@nnx.scan(in_axes=(0, 0, nnx.Carry), out_axes=nnx.Carry)\\n\",\n \"def forward(model: MLP, rngs: nnx.Rngs, x):\\n\",\n \" x = model(x, rngs)\\n\",\n \" return x\\n\",\n \" \\n\",\n \"param_rngs = nnx.Rngs(0).fork(split=5)\\n\",\n \"model = create_model(param_rngs)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 10,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"y.shape = (3, 10)\\n\"\n ]\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"x = jnp.ones((3, 10))\\n\",\n \"dropout_rngs = nnx.Rngs(1).fork(split=5)\\n\",\n \"y = forward(model, dropout_rngs, x)\\n\",\n \"\\n\",\n \"print(f'{y.shape = }')\\n\",\n \"nnx.display(model)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"How do Flax NNX transforms achieve this? To understand how Flax NNX objects interact with JAX transforms, the next section explains the Flax NNX Functional API.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## The Flax Functional API\\n\",\n \"\\n\",\n \"The Flax NNX Functional API establishes a clear boundary between reference/object semantics and value/pytree semantics. It also allows the same amount of fine-grained control over the state that Flax Linen and Haiku users are used to. The Flax NNX Functional API consists of three basic methods: `nnx.split`, `nnx.merge`, and `nnx.update`.\\n\",\n \"\\n\",\n \"Below is an example of `StatefulLinear` `Module` that uses the Functional API. It contains:\\n\",\n \"\\n\",\n \"- Some `Param` Variables; and\\n\",\n \"- A custom `Count` Variable type, which is used to track the integer scalar state that increases on every forward pass.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 10,\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"class Count(nnx.Variable): pass\\n\",\n \"\\n\",\n \"class StatefulLinear(nnx.Module):\\n\",\n \" def __init__(self, din: int, dout: int, *, rngs: nnx.Rngs):\\n\",\n \" self.w = nnx.Param(rngs.uniform((din, dout)))\\n\",\n \" self.b = nnx.Param(jnp.zeros((dout,)))\\n\",\n \" self.count = Count(jnp.array(0, dtype=jnp.uint32))\\n\",\n \"\\n\",\n \" def __call__(self, x: jax.Array):\\n\",\n \" self.count.value += 1\\n\",\n \" return x @ self.w + self.b\\n\",\n \"\\n\",\n \"model = StatefulLinear(din=3, dout=5, rngs=nnx.Rngs(0))\\n\",\n \"y = model(jnp.ones((1, 3)))\\n\",\n \"\\n\",\n \"nnx.display(model)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### State and GraphDef\\n\",\n \"\\n\",\n \"A Flax `Module` can be decomposed into `State` and `GraphDef` using the `nnx.split` function:\\n\",\n \"\\n\",\n \"- `State` is a `Mapping` from strings to `Variable`s or nested `State`s.\\n\",\n \"- `GraphDef` contains all the static information needed to reconstruct a `Module` graph, it is analogous to [JAX's `PyTreeDef`](https://jax.readthedocs.io/en/latest/pytrees.html#internal-pytree-handling).\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 12,\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"graphdef, state = nnx.split(model)\\n\",\n \"\\n\",\n \"nnx.display(graphdef, state)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Split, merge, and update\\n\",\n \"\\n\",\n \"Flax's `nnx.merge` is the reverse of `nnx.split`. It takes the `GraphDef` + `State` and reconstructs the `Module`. The example below demonstrates this as follows:\\n\",\n \"\\n\",\n \"- By using `nnx.split` and `nnx.merge` in sequence any `Module` can be lifted to be used in any JAX transform.\\n\",\n \"- `nnx.update` can update an object in place with the content of a given `State`.\\n\",\n \"- This pattern is used to propagate the state from a transform back to the source object outside.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 13,\n \"metadata\": {},\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"model.count.value = Array(1, dtype=uint32)\\n\",\n \"model.count.value = Array(2, dtype=uint32)\\n\"\n ]\n }\n ],\n \"source\": [\n \"print(f'{model.count.value = }')\\n\",\n \"\\n\",\n \"# 1. Use `nnx.split` to create a pytree representation of the `nnx.Module`.\\n\",\n \"graphdef, state = nnx.split(model)\\n\",\n \"\\n\",\n \"@jax.jit\\n\",\n \"def forward(graphdef: nnx.GraphDef, state: nnx.State, x: jax.Array) -> tuple[jax.Array, nnx.State]:\\n\",\n \" # 2. Use `nnx.merge` to create a new model inside the JAX transformation.\\n\",\n \" model = nnx.merge(graphdef, state)\\n\",\n \" # 3. Call the `nnx.Module`\\n\",\n \" y = model(x)\\n\",\n \" # 4. Use `nnx.split` to propagate `nnx.State` updates.\\n\",\n \" _, state = nnx.split(model)\\n\",\n \" return y, state\\n\",\n \"\\n\",\n \"y, state = forward(graphdef, state, x=jnp.ones((1, 3)))\\n\",\n \"# 5. Update the state of the original `nnx.Module`.\\n\",\n \"nnx.update(model, state)\\n\",\n \"\\n\",\n \"print(f'{model.count.value = }')\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"The key insight of this pattern is that using mutable references is fine within a transform context (including the base eager interpreter) but it is necessary to use the Functional API when crossing boundaries.\\n\",\n \"\\n\",\n \"**Why aren't Modules just pytrees?** The main reason is that it is very easy to lose track of shared references by accident this way, for example if you pass two `Module`s that have a shared Module through a JAX boundary, you will silently lose that sharing. Flax's Functional API makes this behavior explicit, and thus it is much easier to reason about.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"### Fine-grained State control\\n\",\n \"\\n\",\n \"Experienced [Flax Linen](https://flax-linen.readthedocs.io/) or [Haiku](https://dm-haiku.readthedocs.io/) API users may recognize that having all the states in a single structure is not always the best choice as there are cases in which you may want to handle different subsets of the state differently. This is a common occurrence when interacting with [JAX transforms](https://jax.readthedocs.io/en/latest/key-concepts.html#transformations).\\n\",\n \"\\n\",\n \"For example:\\n\",\n \"\\n\",\n \"- Not every model state can or should be differentiated when interacting with `jax.grad`.\\n\",\n \"- Or, sometimes, there is a need to specify what part of the model's state is a carry and what part is not when using `jax.lax.scan`.\\n\",\n \"\\n\",\n \"To address this, the Flax NNX API has `nnx.split`, which allows you to pass one or more `Filter`s to partition the `Variable`s into mutually exclusive `State`s. Flax NNx uses `Filter` create `State` groups in APIs (such as `nnx.split`, `nnx.state`, and many of NNX transforms).\\n\",\n \"\\n\",\n \"The example below shows the most common `Filter`s:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 14,\n \"metadata\": {},\n \"outputs\": [\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"text/html\": [\n \"
\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"# Use `nnx.Variable` type `Filter`s to split into multiple `nnx.State`s.\\n\",\n \"graphdef, params, counts = nnx.split(model, nnx.Param, Count)\\n\",\n \"\\n\",\n \"nnx.display(params, counts)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"**Note:** `Filter`s must be exhaustive, if a value is not matched an error will be raised.\\n\",\n \"\\n\",\n \"As expected, the `nnx.merge` and `nnx.update` methods naturally consume multiple `State`s:\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 15,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# Merge multiple `State`s\\n\",\n \"model = nnx.merge(graphdef, params, counts)\\n\",\n \"# Update with multiple `State`s\\n\",\n \"nnx.update(model, params, counts)\"\n ]\n }\n ],\n \"metadata\": {\n \"jupytext\": {\n \"formats\": \"ipynb,md:myst\"\n },\n \"language_info\": {\n \"codemirror_mode\": {\n \"name\": \"ipython\",\n \"version\": 3\n },\n \"file_extension\": \".py\",\n \"mimetype\": \"text/x-python\",\n \"name\": \"python\",\n \"nbconvert_exporter\": \"python\",\n \"pygments_lexer\": \"ipython3\",\n \"version\": \"3.11.6\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 4\n}\n" }, { "path": "docs_nnx/nnx_basics.md", "content": "---\njupytext:\n formats: ipynb,md:myst\n text_representation:\n extension: .md\n format_name: myst\n format_version: 0.13\n jupytext_version: 1.13.8\n---\n\n# Flax basics\n\nFlax NNX is a new simplified API that is designed to make it easier to create, inspect, debug, and analyze neural networks in [JAX](https://jax.readthedocs.io/). It achieves this by adding first class support for Python reference semantics. This allows users to express their models using regular Python objects, which are modeled as PyGraphs (instead of pytrees), enabling reference sharing and mutability. Such API design should make PyTorch or Keras users feel at home.\n\nTo begin, install Flax with `pip` and import necessary dependencies:\n\n```{code-cell} ipython3\n:tags: [skip-execution]\n\n# ! pip install -U flax\n```\n\n```{code-cell} ipython3\nfrom flax import nnx\nimport jax\nimport jax.numpy as jnp\n```\n\n## The Flax NNX Module system\n\nThe main difference between the Flax `Module` and other Module systems in [Flax Linen](https://flax-linen.readthedocs.io/en/latest/api_reference/flax.linen/module.html) or [Haiku](https://dm-haiku.readthedocs.io/en/latest/notebooks/basics.html#Built-in-Haiku-nets-and-nested-modules) is that in NNX everything is **explicit**. This means, among other things, that the NNX Module itself holds the state (such as parameters) directly, the [PRNG](https://jax.readthedocs.io/en/latest/random-numbers.html) state is threaded by the user, and all shape information must be provided on initialization (no shape inference).\n\nLet's begin by creating a Linear `Module`. As shown next, dynamic state is usually stored in `Param`s, and static state (all types not handled by NNX) such as integers or strings are stored directly. Attributes of type `jax.Array` and `numpy.ndarray` are also treated as dynamic states, although storing them inside Variables, such as Param, is preferred. Also the `Rngs` object can be used to get new unique keys based on a root PRNG key passed to the constructor.\n\n```{code-cell} ipython3\nclass Linear(nnx.Module):\n def __init__(self, din: int, dout: int, *, rngs: nnx.Rngs):\n self.w = nnx.Param(rngs.params.uniform((din, dout)))\n self.b = nnx.Param(jnp.zeros((dout,)))\n self.din, self.dout = din, dout\n\n def __call__(self, x: jax.Array):\n return x @ self.w + self.b[None]\n```\n\nAlso note that the inner values of `Variable`s can be accessed using the `value` property, but for convenience they implement all numeric operators and can be used directly in arithmetic expressions (as shown in the code above).\n\nTo initialize a Flax `Module`, you just call the constructor, and all the parameters of a Module are usually created eagerly. Since Modules hold their own state methods, you can call them directly without the need for a separate apply method.\nThis can be very convenient for debugging, allowing you to directly inspect the entire structure of the model.\n\n```{code-cell} ipython3\nmodel = Linear(2, 5, rngs=nnx.Rngs(params=0))\ny = model(x=jnp.ones((1, 2)))\n\nprint(y)\nnnx.display(model)\n```\n\nThe above visualization by `nnx.display` is generated using the awesome\n[Treescope](https://treescope.readthedocs.io/en/stable/index.html#) library.\n\n+++\n\n### Stateful computation\n\nImplementing layers, such as `BatchNorm`, requires performing state updates during a forward pass. In Flax NNX, you just need to create a `Variable` and update its `.value` during the forward pass.\n\n```{code-cell} ipython3\nclass Count(nnx.Variable): pass\n\nclass Counter(nnx.Module):\n def __init__(self):\n self.count = Count(jnp.array(0))\n\n def __call__(self):\n self.count[...] += 1\n\ncounter = Counter()\nprint(f'{counter.count[...] = }')\ncounter()\nprint(f'{counter.count[...] = }')\n```\n\nMutable references are usually avoided in JAX. But Flax NNX provides sound mechanisms\nto handle them, as demonstrated in later sections of this guide.\n\n+++\n\n### Nested Modules\n\nFlax `Module`s can be used to compose other Modules in a nested structure. These can be assigned directly as attributes, or inside an attribute of any (nested) pytree type, such as a `list`, `dict`, `tuple`, and so on.\n\nThe example below shows how to define a simple `MLP` by subclassing `Module`. The model consists of two `Linear` layers, a `Dropout` layer, and a `BatchNorm` layer. Note that we need to pass the `__call__` method the RNG state that we want the `Dropout` layer to use.\n\n```{code-cell} ipython3\nclass MLP(nnx.Module):\n def __init__(self, din: int, dmid: int, dout: int, *, rngs: nnx.Rngs):\n self.linear1 = Linear(din, dmid, rngs=rngs)\n self.dropout = nnx.Dropout(rate=0.1)\n self.bn = nnx.BatchNorm(dmid, rngs=rngs)\n self.linear2 = Linear(dmid, dout, rngs=rngs)\n\n def __call__(self, x: jax.Array, rngs: nnx.Rngs):\n x = nnx.gelu(self.dropout(self.bn(self.linear1(x)), rngs=rngs))\n return self.linear2(x)\n\nmodel = MLP(2, 16, 5, rngs=nnx.Rngs(0))\n\ny = model(x=jnp.ones((3, 2)), rngs=nnx.Rngs(1))\n\nnnx.display(model)\n```\n\n### Model surgery\n\nFlax `Module`s are mutable by default. This means that their structure can be changed at any time, which makes [model surgery](https://flax.readthedocs.io/en/latest/guides/surgery.html) quite easy, as any sub-Module attribute can be replaced with anything else, such as new Modules, existing shared Modules, Modules of different types, and so on. Moreover, `Variable`s can also be modified or replaced/shared.\n\nThe following example shows how to replace the `Linear` layers in the `MLP` model from the previous example with `LoraLinear` layers:\n\n```{code-cell} ipython3\nclass LoraParam(nnx.Param): pass\n\nclass LoraLinear(nnx.Module):\n def __init__(self, linear: Linear, rank: int, rngs: nnx.Rngs):\n self.linear = linear\n self.A = LoraParam(rngs.normal((linear.din, rank)))\n self.B = LoraParam(rngs.normal((rank, linear.dout)))\n\n def __call__(self, x: jax.Array):\n return self.linear(x) + x @ self.A @ self.B\n\nrngs = nnx.Rngs(0)\nmodel = MLP(2, 32, 5, rngs=rngs)\n\n# Model surgery.\nmodel.linear1 = LoraLinear(model.linear1, 4, rngs=rngs)\nmodel.linear2 = LoraLinear(model.linear2, 4, rngs=rngs)\n\ny = model(x=jnp.ones((3, 2)), rngs=rngs)\n\nnnx.display(model)\n```\n\n## Flax transformations\n\n[Flax NNX transformations (transforms)](https://flax.readthedocs.io/en/latest/guides/jax_and_nnx_transforms.html) extend [JAX transforms](https://jax.readthedocs.io/en/latest/key-concepts.html#transformations) to support `Module`s and other objects. They serve as supersets of their equivalent JAX counterparts with the addition of being aware of the object's state and providing additional APIs to transform it.\n\nOne of the main features of Flax Transforms is the preservation of reference semantics, meaning that any mutation of the object graph that occurs inside the transform is propagated outside as long as it is legal within the transform rules. In practice this means that Flax programs can be expressed using imperative code, highly simplifying the user experience.\n\nIn the following example, you define a `train_step` function that takes a `MLP` model, an `Optimizer`, and a batch of data, and returns the loss for that step. The loss and the gradients are computed using the `nnx.value_and_grad` transform over the `loss_fn`. The gradients are passed to the optimizer's `update` method to update the model's parameters.\n\n```{code-cell} ipython3\nimport optax\n\n# An MLP containing 2 custom `Linear` layers, 1 `nnx.Dropout` layer, 1 `nnx.BatchNorm` layer.\nmodel = MLP(2, 16, 10, rngs=nnx.Rngs(0))\noptimizer = nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param)\n\n@nnx.jit # Automatic state management\ndef train_step(model, optimizer, x, y, rngs):\n def loss_fn(model: MLP, rngs: nnx.Rngs):\n y_pred = model(x, rngs)\n return jnp.mean((y_pred - y) ** 2)\n\n loss, grads = nnx.value_and_grad(loss_fn)(model, rngs)\n optimizer.update(model, grads) # In place updates.\n\n return loss\n\nx, y = jnp.ones((5, 2)), jnp.ones((5, 10))\nloss = train_step(model, optimizer, x, y, rngs)\n\nprint(f'{loss = }')\nprint(f'{optimizer.step.value = }')\n```\n\nThere are two things happening in this example that are worth mentioning:\n\n1. The updates to each of the `BatchNorm` and `Dropout` layer's state is automatically propagated from within `loss_fn` to `train_step` all the way to the `model` reference outside.\n2. The `optimizer` holds a mutable reference to the model - this relationship is preserved inside the train_step function making it possible to update the model's parameters using the optimizer alone.\n\n> **Note**
`nnx.jit` has performance overhead for small models, check the [Performance Considerations](https://flax.readthedocs.io/en/latest/guides/performance.html) guide for more information.\n\n### Scan over layers\n\nThe next example uses Flax `nnx.vmap` to create a stack of multiple MLP layers and `nnx.scan` to iteratively apply each layer of the stack to the input.\n\nIn the code below notice the following:\n\n1. The custom `create_model` function takes in a key and returns an `MLP` object, since you create five keys and use `nnx.vmap` over `create_model` a stack of 5 `MLP` objects is created.\n2. The `nnx.scan` is used to iteratively apply each `MLP` in the stack to the input `x`.\n3. The nnx.scan (consciously) deviates from `jax.lax.scan` and instead mimics nnx.vmap, which is more expressive. nnx.scan allows specifying multiple inputs, the scan axes of each input/output, and the position of the carry.\n4. `State` updates for `BatchNorm` layers are automatically propagated by nnx.scan.\n5. The `rngs` object is split into separate streams for each layer using the `fork` method.\n\n```{code-cell} ipython3\n@nnx.vmap(in_axes=0, out_axes=0)\ndef create_model(rngs):\n return MLP(10, 32, 10, rngs=rngs)\n\n@nnx.scan(in_axes=(0, 0, nnx.Carry), out_axes=nnx.Carry)\ndef forward(model: MLP, rngs: nnx.Rngs, x):\n x = model(x, rngs)\n return x\n \nparam_rngs = nnx.Rngs(0).fork(split=5)\nmodel = create_model(param_rngs)\n```\n\n```{code-cell} ipython3\nx = jnp.ones((3, 10))\ndropout_rngs = nnx.Rngs(1).fork(split=5)\ny = forward(model, dropout_rngs, x)\n\nprint(f'{y.shape = }')\nnnx.display(model)\n```\n\nHow do Flax NNX transforms achieve this? To understand how Flax NNX objects interact with JAX transforms, the next section explains the Flax NNX Functional API.\n\n+++\n\n## The Flax Functional API\n\nThe Flax NNX Functional API establishes a clear boundary between reference/object semantics and value/pytree semantics. It also allows the same amount of fine-grained control over the state that Flax Linen and Haiku users are used to. The Flax NNX Functional API consists of three basic methods: `nnx.split`, `nnx.merge`, and `nnx.update`.\n\nBelow is an example of `StatefulLinear` `Module` that uses the Functional API. It contains:\n\n- Some `Param` Variables; and\n- A custom `Count` Variable type, which is used to track the integer scalar state that increases on every forward pass.\n\n```{code-cell} ipython3\nclass Count(nnx.Variable): pass\n\nclass StatefulLinear(nnx.Module):\n def __init__(self, din: int, dout: int, *, rngs: nnx.Rngs):\n self.w = nnx.Param(rngs.uniform((din, dout)))\n self.b = nnx.Param(jnp.zeros((dout,)))\n self.count = Count(jnp.array(0, dtype=jnp.uint32))\n\n def __call__(self, x: jax.Array):\n self.count.value += 1\n return x @ self.w + self.b\n\nmodel = StatefulLinear(din=3, dout=5, rngs=nnx.Rngs(0))\ny = model(jnp.ones((1, 3)))\n\nnnx.display(model)\n```\n\n### State and GraphDef\n\nA Flax `Module` can be decomposed into `State` and `GraphDef` using the `nnx.split` function:\n\n- `State` is a `Mapping` from strings to `Variable`s or nested `State`s.\n- `GraphDef` contains all the static information needed to reconstruct a `Module` graph, it is analogous to [JAX's `PyTreeDef`](https://jax.readthedocs.io/en/latest/pytrees.html#internal-pytree-handling).\n\n```{code-cell} ipython3\ngraphdef, state = nnx.split(model)\n\nnnx.display(graphdef, state)\n```\n\n### Split, merge, and update\n\nFlax's `nnx.merge` is the reverse of `nnx.split`. It takes the `GraphDef` + `State` and reconstructs the `Module`. The example below demonstrates this as follows:\n\n- By using `nnx.split` and `nnx.merge` in sequence any `Module` can be lifted to be used in any JAX transform.\n- `nnx.update` can update an object in place with the content of a given `State`.\n- This pattern is used to propagate the state from a transform back to the source object outside.\n\n```{code-cell} ipython3\nprint(f'{model.count.value = }')\n\n# 1. Use `nnx.split` to create a pytree representation of the `nnx.Module`.\ngraphdef, state = nnx.split(model)\n\n@jax.jit\ndef forward(graphdef: nnx.GraphDef, state: nnx.State, x: jax.Array) -> tuple[jax.Array, nnx.State]:\n # 2. Use `nnx.merge` to create a new model inside the JAX transformation.\n model = nnx.merge(graphdef, state)\n # 3. Call the `nnx.Module`\n y = model(x)\n # 4. Use `nnx.split` to propagate `nnx.State` updates.\n _, state = nnx.split(model)\n return y, state\n\ny, state = forward(graphdef, state, x=jnp.ones((1, 3)))\n# 5. Update the state of the original `nnx.Module`.\nnnx.update(model, state)\n\nprint(f'{model.count.value = }')\n```\n\nThe key insight of this pattern is that using mutable references is fine within a transform context (including the base eager interpreter) but it is necessary to use the Functional API when crossing boundaries.\n\n**Why aren't Modules just pytrees?** The main reason is that it is very easy to lose track of shared references by accident this way, for example if you pass two `Module`s that have a shared Module through a JAX boundary, you will silently lose that sharing. Flax's Functional API makes this behavior explicit, and thus it is much easier to reason about.\n\n+++\n\n### Fine-grained State control\n\nExperienced [Flax Linen](https://flax-linen.readthedocs.io/) or [Haiku](https://dm-haiku.readthedocs.io/) API users may recognize that having all the states in a single structure is not always the best choice as there are cases in which you may want to handle different subsets of the state differently. This is a common occurrence when interacting with [JAX transforms](https://jax.readthedocs.io/en/latest/key-concepts.html#transformations).\n\nFor example:\n\n- Not every model state can or should be differentiated when interacting with `jax.grad`.\n- Or, sometimes, there is a need to specify what part of the model's state is a carry and what part is not when using `jax.lax.scan`.\n\nTo address this, the Flax NNX API has `nnx.split`, which allows you to pass one or more `Filter`s to partition the `Variable`s into mutually exclusive `State`s. Flax NNx uses `Filter` create `State` groups in APIs (such as `nnx.split`, `nnx.state`, and many of NNX transforms).\n\nThe example below shows the most common `Filter`s:\n\n```{code-cell} ipython3\n# Use `nnx.Variable` type `Filter`s to split into multiple `nnx.State`s.\ngraphdef, params, counts = nnx.split(model, nnx.Param, Count)\n\nnnx.display(params, counts)\n```\n\n**Note:** `Filter`s must be exhaustive, if a value is not matched an error will be raised.\n\nAs expected, the `nnx.merge` and `nnx.update` methods naturally consume multiple `State`s:\n\n```{code-cell} ipython3\n# Merge multiple `State`s\nmodel = nnx.merge(graphdef, params, counts)\n# Update with multiple `State`s\nnnx.update(model, params, counts)\n```\n" }, { "path": "docs_nnx/nnx_glossary.rst", "content": "*****************\nFlax NNX glossary\n*****************\n\nFor additional terms, refer to the `JAX glossary `__.\n\n.. glossary::\n\n Filter\n A way to extract only certain :term:`nnx.Variable` objects out of a Flax NNX :term:`Module` (``nnx.Module``). This is usually done by calling :meth:`nnx.split ` upon the :class:`nnx.Module`. Refer to the `Filter guide `__ to learn more.\n\n Folding in\n In Flax, `folding in `__ means generating a new `JAX pseudorandom number generator (PRNG) `__ key, given an input PRNG key and integer. This is typically used when you want to generate a new key but still be able to use the original PRNG key afterwards. You can also do this in JAX with `jax.random.split `__, but this method will effectively create two PRNG keys, which is slower. Learn how Flax generates new PRNG keys automatically in the `Randomness/PRNG guide `__.\n\n GraphDef\n :class:`nnx.GraphDef` is a class that represents all the static, stateless, and Pythonic parts of a Flax :term:`Module` (:class:`nnx.Module`).\n\n Merge\n Refer to :term:`Split and merge`.\n\n Module\n :class:`nnx.Module ` is a dataclass that enables defining and initializing parameters in a referentially-transparent form. It is responsible for storing and updating :term:`Variable objects and parameters within itself.\n\n Params / parameters\n :class:`nnx.Param ` is a particular subclass of :class:`nnx.Variable ` that generally contains the trainable weights.\n\n PRNG states\n A Flax :class:`nnx.Module ` can keep a reference of a `pseudorandom number generator (PRNG) `__ state object :class:`nnx.Rngs ` that can generate new `JAX PRNG `__ keys. These keys are used to generate random JAX arrays through `JAX's functional PRNGs `__.\n You can use a PRNG state with different seeds to add more fine-grained control to your model (for example, to have independent random numbers for parameters and dropout masks).\n Refer to the Flax `Randomness/PRNG guide `__\n for more details.\n\n Split and merge\n :meth:`nnx.split ` is a way to represent an :class:`nnx.Module ` by two parts: 1) a static Flax NNX :term:`GraphDef ` that captures its Pythonic static information; and 2) one or more :term:`Variable state(s)` that capture its `JAX arrays `__ (``jax.Array``) in the form of `JAX pytrees `__. They can be merged back to the original ``nnx.Module`` using :meth:`nnx.merge `.\n\n Transformation\n A Flax NNX transformation (transform) is a wrapped version of a `JAX transformation `__ that allows the function that is being transformed to take the Flax NNX :term:`Module` (``nnx.Module``) as input or output. For example, a \"lifted\" version of `jax.jit `__ is :meth:`nnx.jit `. Check out the `Flax NNX transforms guide `__ to learn more.\n\n Variable\n The weights / parameters / data / array :class:`nnx.Variable ` residing in a Flax :term:`Module`. Variables are defined inside modules as :class:`nnx.Variable ` or its subclasses.\n" }, { "path": "docs_nnx/philosophy.md", "content": "# The Flax philosophy\n\nIn no particular order:\n\n* Library code should be easy to read and understand.\n* Prefer duplicating code over a bad abstraction.\n* Generally, prefer duplicating code over adding options to functions.\n* Comment-driven design: If it's hard to document your code, consider\n changing the design.\n* Unit test-driven design: If it's hard to test your code, consider\n changing the design.\n* People start projects by copying an existing implementation — make\n base implementations excellent.\n* If we expose an abstraction to our developers, we own the mental\n overhead.\n* Developer-facing functional programming abstractions confuse some users,\n expose them where the benefit is high.\n* \"Read the manual\" is not an appropriate response to developer confusion.\n The framework should guide developers\n towards good solutions, such as through assertions and error messages.\n* An unhelpful error message is a bug.\n* \"Debugging is twice as hard as writing the code in the first\n place. Therefore, if you write the code as cleverly as possible, you\n are, by definition, not smart enough to debug it.\" — Brian Kernighan\n\n## Design principles\n\nFlax is a neural network library built on [JAX](https://jax.readthedocs.io) that has been adopted by a\ngrowing set of users, most notably in the JAX submissions for the MLPerf\n0.7 benchmark. Our experience over the last year (and many conversations\nwith users and JAX core devs) has guided a redesign of the API called\n[Linen](https://github.com/google/flax/blob/main/flax/linen/README.md) ([`flax.linen`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/index.html)) in response to the following basic design questions.\n\n### How does a neural network library benefit from being built on JAX and leverage JAX’s unique strengths?\n\nThe world already has TensorFlow and PyTorch, and there’s little need to\nbuild a clone of either. We believe that the composable\nfunction-transformation approach that JAX takes opens up new frontiers\nfor making neural net code more maintainable, more scalable and more\nperformant than existing libraries. While we strive to offer an API\nfamiliar to those experienced with Keras/Sonnet/PyTorch, Linen is\nfundamentally a functional system for defining neural nets in JAX. Just\na few examples of what we believe a JAX-targeted library can enable:\n\n- Write models as “single-example” code and introduce batching\n automatically with [`jax.vmap`](https://jax.readthedocs.io/en/latest/_autosummary/jax.vmap.html).\n- Automatically handle ragged batches in NLP and other masking issues.\n- Create efficient compile-time and runtime models by utilizing\n rematerialized `scan` for massive convolutional networks.\n- Remove memory headaches by enabling easy rematerialization,\n reversibility, and model-parallel data sharding.\n\n### How does one interoperate with JAX transformations?\n\nArguably, the entire point of a neural net library is to offer an\nimplicit variable management API to save the user from having to\nmanually thread thousands of variables through a complex tree of\nfunctions. However, JAX operates on pure functions. To handle both\ncurrent and future JAX transforms (configured and composed in any way),\nLinen Modules are directly “functionalized”, that is, automatically cast\nin-place as explicit functions of the form:\n\n$$f \\left( v_{in}, x \\right) \\rightarrow v_{out}, y$$\n\nWhere $v_{in}$ is the variable collections and [PRNG](https://jax.readthedocs.io/en/latest/jep/263-prng.html) state used by\nthe model, $v_{out}$ the mutated output variable collections,\n$x$ the input data and $y$ the output data. Applying JAX\ntransformations then simply reduces to specifying any argument-specific\ntransform options to the various variable collections and PRNG state.\nThis unleashes the flexibility and strength of [JAX transformations](https://jax.readthedocs.io/en/latest/notebooks/quickstart.html) – for\nexample, one can achieve either device-parallel training or per-device\nensembling by using [`jax.pmap`](https://jax.readthedocs.io/en/latest/_autosummary/jax.pmap.html) in different ways, without any explicit\nlibrary support. Moreover, **within [Modules](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module)**, we expose lightweight\nwrappers around the complex JAX transforms such as [`jax.vmap`](https://jax.readthedocs.io/en/latest/_autosummary/jax.vmap.html) and [`jax.lax.scan`](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.scan.html)\nthat annotate how each variable collection is to be transformed by JAX.\nImportantly, we handle the nontrivial cases of creating new variables\nand transformed variables under mapping and loop transforms correctly\nfor initialization and application.\n\n### How are parameters represented, and how do we handle general “differentiable algorithms” that update stateful variables?\n\nWe follow the JAX functional conventions of storing data in “pytrees”:\nJAX arrays contained in nested tuples, lists, dictionaries. Because\nresearchers inevitably manually interact with this data, we use nested\ndictionaries with meaningful default keys and offer several utilities\n(traversals, etc.) for handling them directly. Linen uses an accelerated\nversion of a Python frozen dictionary that caches its JAX-flattened form\nto speed up [`jit`](https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html)ted function call overheads.\n\nFlax generalizes the operation of a neural net by allowing models to\naccept collections of several different “kinds”: parameters, batch-norm\nstats, autoregressive caches, debug information, fine-grained\nhyperparameters, etc. Each collection is stored in a nested dictionary\nof the same structure as the model. Importantly, we do *not* conflate\nthese various kinds under the single vague rubric of “state”, but keep\ndifferent logical types of variables separate that can be treated\ndifferently under JAX transformations and under mutations (e.g. training\nvs prediction). Similarly, we allow for multiple separate named PRNG\nchains inside [Modules](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module) for separate treatment of randomness for different\napplications such as initialization, dropout, sampling, etc.\n\nAt every stage the data associated with a neural net is not kept in a\ncustom object hierarchy, but left in an explicit, Python and JAX native\nform that is easy to introspect and modify. Users have utilized this to\nmap TF and PyTorch checkpoints to Flax, to implement submodel-specific\nloss terms, and to perform fast model surgery, etc. For saving this\ndata, most Flax examples store these nested dictionaries via the\nefficient “msgpack” binary format – but as variables are simply Python\ndicts, you can use any (non-JAX-aware) serialization library directly.\n\n### How does one interoperate with purely functional JAX code?\n\nTo be broadly useful to the JAX ecosystem, users shouldn’t need to\nheavily refactor their code in order to add “trainability” for a given\nnumerical task. _“The library should not get in the way.”_ Utilizing\npurely functional code from within Linen is trivial: [Module](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module)\nimplementations are just JAX code with named variables. Using Linen\nModules inside otherwise purely functional code can be as simple as\nusing a single top-level Module transformation to allow initialization\nand pure application of any JAX program that might contain various\ntrainable sections.\n" }, { "path": "docs_nnx/robots.txt", "content": "User-agent: *\n\nDisallow: /api_reference/flax.linen/_autosummary/ # for SEO, since Google still indexes this deprecated link\n\nSitemap: https://flax.readthedocs.io/sitemap.xml\n" }, { "path": "docs_nnx/why.rst", "content": "Why Flax NNX?\n=============\n\nIn 2020, the Flax team released the Flax Linen API to support modeling research on JAX, with a focus on scaling\nand performance. We have learned a lot from users since then. The team introduced certain ideas that have proven to be beneficial to users, such as:\n\n* Organizing variables into `collections `_.\n* Automatic and efficient `pseudorandom number generator (PRNG) management `_.\n* `Variable metadata `_\n for `Single Program Multi Data (SPMD) `_ annotations, optimizer metadata, and other use cases.\n\nOne of the choices the Flax team made was to use functional (``compact``) semantics for neural network programming via lazy initialization of parameters.\nThis made for concise implementation code and aligned the Flax Linen API with Haiku.\n\nHowever, this also meant that the semantics of Modules and variables in Flax were non-Pythonic and often surprising. It also led to implementation\ncomplexity and obscured the core ideas of `transformations (transforms) `_ on neural networks.\n\n.. testsetup:: Linen, NNX\n\n import jax\n from jax import random, numpy as jnp\n from flax import nnx\n import flax.linen as nn\n\nIntroducing Flax NNX\n--------------------\n\nFast forward to 2024, the Flax team developed Flax NNX - an attempt to retain the features that made Flax Linen useful for users, while introducing some new principles.\nThe central idea behind Flax NNX is to introduce reference semantics into JAX. The following are its main features:\n\n- **NNX is Pythonic**: Regular Python semantics for Modules, including support for mutability and shared references.\n- **NNX is simple**: Many of the complex APIs in Flax Linen are either simplified using Python idioms or completely removed.\n- **Better JAX integration**: Custom NNX transforms adopt the same APIs as the JAX transforms. And with NNX\n it is easier to use `JAX transforms (higher-order functions) `_ directly.\n\nHere is an example of a simple Flax NNX program that illustrates many of the points from above:\n\n.. testcode:: NNX\n\n from flax import nnx\n import optax\n\n\n class Model(nnx.Module):\n def __init__(self, din, dmid, dout, rngs: nnx.Rngs):\n self.linear = nnx.Linear(din, dmid, rngs=rngs)\n self.bn = nnx.BatchNorm(dmid, rngs=rngs)\n self.dropout = nnx.Dropout(0.2, rngs=rngs)\n self.linear_out = nnx.Linear(dmid, dout, rngs=rngs)\n\n def __call__(self, x):\n x = nnx.relu(self.dropout(self.bn(self.linear(x))))\n return self.linear_out(x)\n\n model = Model(2, 64, 3, rngs=nnx.Rngs(0)) # Eager initialization\n optimizer = nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param)\n\n @nnx.jit # Automatic state management for JAX transforms.\n def train_step(model, optimizer, x, y):\n def loss_fn(model):\n y_pred = model(x) # call methods directly\n return ((y_pred - y) ** 2).mean()\n\n loss, grads = nnx.value_and_grad(loss_fn)(model)\n optimizer.update(model, grads) # in-place updates\n\n return loss\n\nFlax NNX's improvements on Linen\n--------------------------------\n\nThe rest of this document uses various examples that demonstrate how Flax NNX improves on Flax Linen.\n\nInspection\n^^^^^^^^^^\n\nThe first improvement is that Flax NNX Modules are regular Python objects. This means that you can easily\nconstruct and inspect ``Module`` objects.\n\nOn the other hand, Flax Linen Modules are not easy to inspect and debug because they are lazy, which means some attributes are not available upon construction and are only accessible at runtime.\n\n.. codediff::\n :title: Linen, NNX\n :sync:\n\n class Block(nn.Module):\n def setup(self):\n self.linear = nn.Dense(10)\n\n block = Block()\n\n try:\n block.linear # AttributeError: \"Block\" object has no attribute \"linear\".\n except AttributeError as e:\n pass\n\n\n\n\n\n ...\n\n ---\n\n class Block(nnx.Module):\n def __init__(self, rngs):\n self.linear = nnx.Linear(5, 10, rngs=rngs)\n\n block = Block(nnx.Rngs(0))\n\n\n block.linear\n # Linear(\n # kernel=Param(\n # value=Array(shape=(5, 10), dtype=float32)\n # ),\n # bias=Param(\n # value=Array(shape=(10,), dtype=float32)\n # ),\n # ...\n\nNotice that in the Flax NNX example above, there is no shape inference - both the input and output shapes must be provided\nto the ``Linear`` ``nnx.Module``. This is a tradeoff that allows for more explicit and predictable behavior.\n\nRunning computation\n^^^^^^^^^^^^^^^^^^^\n\nIn Flax Linen, all top-level computation must be done through the ``flax.linen.Module.init`` or ``flax.linen.Module.apply`` methods, and the\nparameters or any other type of state are handled as a separate structure. This creates an asymmetry between: 1) code that runs inside\n``apply`` that can run methods and other ``Module`` objects directly; and 2) code that runs outside of ``apply`` that must use the ``apply`` method.\n\nIn Flax NNX, there's no special context because parameters are held as attributes and methods can be called directly. That means your NNX Module's ``__init__`` and ``__call__`` methods are not treated differently from other class methods, whereas Flax Linen Module's ``setup()`` and ``__call__`` methods are special.\n\n.. codediff::\n :title: Linen, NNX\n :sync:\n\n Encoder = lambda: nn.Dense(10)\n Decoder = lambda: nn.Dense(2)\n\n class AutoEncoder(nn.Module):\n def setup(self):\n self.encoder = Encoder()\n self.decoder = Decoder()\n\n def __call__(self, x) -> jax.Array:\n return self.decoder(self.encoder(x))\n\n def encode(self, x) -> jax.Array:\n return self.encoder(x)\n\n x = jnp.ones((1, 2))\n model = AutoEncoder()\n params = model.init(random.key(0), x)['params']\n\n y = model.apply({'params': params}, x)\n z = model.apply({'params': params}, x, method='encode')\n y = Decoder().apply({'params': params['decoder']}, z)\n\n ---\n\n Encoder = lambda rngs: nnx.Linear(2, 10, rngs=rngs)\n Decoder = lambda rngs: nnx.Linear(10, 2, rngs=rngs)\n\n class AutoEncoder(nnx.Module):\n def __init__(self, rngs):\n self.encoder = Encoder(rngs)\n self.decoder = Decoder(rngs)\n\n def __call__(self, x) -> jax.Array:\n return self.decoder(self.encoder(x))\n\n def encode(self, x) -> jax.Array:\n return self.encoder(x)\n\n x = jnp.ones((1, 2))\n model = AutoEncoder(nnx.Rngs(0))\n\n\n y = model(x)\n z = model.encode(x)\n y = model.decoder(z)\n\nIn Flax Linen, calling sub-Modules directly is not possible because they are not initialized.\nTherefore, what you must do is construct a new instance and then provide a proper parameter structure.\n\nBut in Flax NNX you can call sub-Modules directly without any issues.\n\nState handling\n^^^^^^^^^^^^^^\n\nOne of the areas where Flax Linen is notoriously complex is in state handling. When you use either a\n`Dropout` layer, a `BatchNorm` layer, or both, you suddenly have to handle the new state and use it to\nconfigure the ``flax.linen.Module.apply`` method.\n\nIn Flax NNX, state is kept inside an ``nnx.Module`` and is mutable, which means it can just be called directly.\n\n.. codediff::\n :title: Linen, NNX\n :sync:\n\n class Block(nn.Module):\n train: bool\n\n def setup(self):\n self.linear = nn.Dense(10)\n self.bn = nn.BatchNorm(use_running_average=not self.train)\n self.dropout = nn.Dropout(0.1, deterministic=not self.train)\n\n def __call__(self, x):\n return nn.relu(self.dropout(self.bn(self.linear(x))))\n\n x = jnp.ones((1, 5))\n model = Block(train=True)\n vs = model.init(random.key(0), x)\n params, batch_stats = vs['params'], vs['batch_stats']\n\n y, updates = model.apply(\n {'params': params, 'batch_stats': batch_stats},\n x,\n rngs={'dropout': random.key(1)},\n mutable=['batch_stats'],\n )\n batch_stats = updates['batch_stats']\n\n ---\n\n class Block(nnx.Module):\n\n\n def __init__(self, rngs):\n self.linear = nnx.Linear(5, 10, rngs=rngs)\n self.bn = nnx.BatchNorm(10, rngs=rngs)\n self.dropout = nnx.Dropout(0.1, rngs=rngs)\n\n def __call__(self, x):\n return nnx.relu(self.dropout(self.bn(self.linear(x))))\n\n x = jnp.ones((1, 5))\n model = Block(nnx.Rngs(0))\n\n\n\n y = model(x)\n\n\n\n\n\n ...\n\nThe main benefit of Flax NNX's state handling is that you don't have to change the training code when you add a new stateful layer.\n\nIn addition, in Flax NNX, layers that handle state are also very easy to implement. Below\nis a simplified version of a ``BatchNorm`` layer that updates the mean and variance every time it is called.\n\n.. testcode:: NNX\n\n class BatchNorm(nnx.Module):\n def __init__(self, features: int, mu: float = 0.95):\n # Variables\n self.scale = nnx.Param(jax.numpy.ones((features,)))\n self.bias = nnx.Param(jax.numpy.zeros((features,)))\n self.mean = nnx.BatchStat(jax.numpy.zeros((features,)))\n self.var = nnx.BatchStat(jax.numpy.ones((features,)))\n self.mu = mu # Static\n\n def __call__(self, x):\n mean = jax.numpy.mean(x, axis=-1)\n var = jax.numpy.var(x, axis=-1)\n # ema updates\n self.mean.value = self.mu * self.mean + (1 - self.mu) * mean\n self.var.value = self.mu * self.var + (1 - self.mu) * var\n # normalize and scale\n x = (x - mean) / jax.numpy.sqrt(var + 1e-5)\n return x * self.scale + self.bias\n\n\nModel surgery\n^^^^^^^^^^^^^\n\nIn Flax Linen, `model surgery `_ has historically been challenging because of two reasons:\n\n1. Due to lazy initialization, it is not guaranteed that you can replace a sub-``Module`` with a new one.\n2. The parameter structure is separated from the ``flax.linen.Module`` structure, which means you have to manually keep them in sync.\n\nIn Flax NNX, you can replace sub-Modules directly as per the Python semantics. Since parameters are\npart of the ``nnx.Module`` structure, they are never out of sync. Below is an example of how you can\nimplement a LoRA layer, and then use it to replace a ``Linear`` layer in an existing model.\n\n.. codediff::\n :title: Linen, NNX\n :sync:\n\n class LoraLinear(nn.Module):\n linear: nn.Dense\n rank: int\n\n @nn.compact\n def __call__(self, x: jax.Array):\n A = self.param(random.normal, (x.shape[-1], self.rank))\n B = self.param(random.normal, (self.rank, self.linear.features))\n\n return self.linear(x) + x @ A @ B\n\n try:\n model = Block(train=True)\n model.linear = LoraLinear(model.linear, rank=5) # <-- ERROR\n\n lora_params = model.linear.init(random.key(1), x)\n lora_params['linear'] = params['linear']\n params['linear'] = lora_params\n\n except AttributeError as e:\n pass\n\n ---\n\n class LoraParam(nnx.Param): pass\n\n class LoraLinear(nnx.Module):\n def __init__(self, linear, rank, rngs):\n self.linear = linear\n self.A = LoraParam(random.normal(rngs(), (linear.in_features, rank)))\n self.B = LoraParam(random.normal(rngs(), (rank, linear.out_features)))\n\n def __call__(self, x: jax.Array):\n return self.linear(x) + x @ self.A @ self.B\n\n rngs = nnx.Rngs(0)\n model = Block(rngs)\n model.linear = LoraLinear(model.linear, rank=5, rngs=rngs)\n\n\n\n\n\n\n ...\n\nAs shown above, in Flax Linen this doesn't really work in this case because the ``linear`` sub-``Module``\nis not available. However, the rest of the code provides an idea of how the ``params`` structure must be manually updated.\n\nPerforming arbitrary model surgery is not easy in Flax Linen, and currently the\n`intercept_methods `_\nAPI is the only way to do generic patching of methods. But this API is not very ergonomic.\n\nIn Flax NNX, to do generic model surgery you can just use ``nnx.iter_graph``, which is much simpler and easier than in Linen. Below is an example of replacing all ``nnx.Linear`` layers in a model with custom-made ``LoraLinear`` NNX layers.\n\n.. testcode:: NNX\n\n rngs = nnx.Rngs(0)\n model = Block(rngs)\n\n for path, module in nnx.iter_graph(model):\n if isinstance(module, nnx.Module):\n for name, value in vars(module).items():\n if isinstance(value, nnx.Linear):\n setattr(module, name, LoraLinear(value, rank=5, rngs=rngs))\n\nTransforms\n^^^^^^^^^^\n\nFlax Linen transforms are very powerful in that they enable fine-grained control over the model's state.\nHowever, Flax Linen transforms have drawbacks, such as:\n\n1. They expose additional APIs that are not part of JAX, making their behavior confusing and sometimes divergent from their JAX counterparts. This also constrains your ways to interact with `JAX transforms `_ and keep up with JAX API changes.\n2. They work on functions with very specific signatures, namely:\n - A ``flax.linen.Module`` must be the first argument.\n - They accept other ``Module`` objects as arguments but not as return values.\n3. They can only be used inside ``flax.linen.Module.apply``.\n\nOn the other hand, `Flax NNX transforms `_\nare intented to be equivalent to their corresponding `JAX transforms `_\nwith an exception - they can be used on Flax NNX Modules. This means that Flax transforms:\n\n1) Have the same API as JAX transforms.\n2) Can accept Flax NNX Modules on any argument, and ``nnx.Module`` objects can be returned from it/them.\n3) Can be used anywhere including the training loop.\n\nBelow is an example of using ``vmap`` with Flax NNX to both create a stack of weights by transforming the\n``create_weights`` function, which returns some ``Weights``, and to apply that stack of weights to a batch\nof inputs individually by transforming the ``vector_dot`` function, which takes ``Weights`` as the first\nargument and a batch of inputs as the second argument.\n\n.. testcode:: NNX\n\n class Weights(nnx.Module):\n def __init__(self, kernel: jax.Array, bias: jax.Array):\n self.kernel, self.bias = nnx.Param(kernel), nnx.Param(bias)\n\n def create_weights(seed: jax.Array):\n return Weights(\n kernel=random.uniform(random.key(seed), (2, 3)),\n bias=jnp.zeros((3,)),\n )\n\n def vector_dot(weights: Weights, x: jax.Array):\n assert weights.kernel.ndim == 2, 'Batch dimensions not allowed'\n assert x.ndim == 1, 'Batch dimensions not allowed'\n return x @ weights.kernel + weights.bias\n\n seeds = jnp.arange(10)\n weights = nnx.vmap(create_weights, in_axes=0, out_axes=0)(seeds)\n\n x = jax.random.normal(random.key(1), (10, 2))\n y = nnx.vmap(vector_dot, in_axes=(0, 0), out_axes=1)(weights, x)\n\nContrary to Flax Linen transforms, the ``in_axes`` argument and other APIs do affect how the ``nnx.Module`` state is transformed.\n\nIn addition, Flax NNX transforms can be used as method decorators, because ``nnx.Module`` methods are simply\nfunctions that take a ``Module`` as the first argument. This means that the previous example can be\nrewritten as follows:\n\n.. testcode:: NNX\n\n class WeightStack(nnx.Module):\n @nnx.vmap(in_axes=(0, 0), out_axes=0)\n def __init__(self, seed: jax.Array):\n self.kernel = nnx.Param(random.uniform(random.key(seed), (2, 3)))\n self.bias = nnx.Param(jnp.zeros((3,)))\n\n @nnx.vmap(in_axes=(0, 0), out_axes=1)\n def __call__(self, x: jax.Array):\n assert self.kernel.ndim == 2, 'Batch dimensions not allowed'\n assert x.ndim == 1, 'Batch dimensions not allowed'\n return x @ self.kernel + self.bias\n\n weights = WeightStack(jnp.arange(10))\n\n x = jax.random.normal(random.key(1), (10, 2))\n y = weights(x)\n\n\n" }, { "path": "examples/README.md", "content": "# Flax Examples\n\nEach example is designed to be **self-contained and easily forkable**, while\nreproducing relevant results in different areas of machine learning.\n\nAs discussed in [#231](https://github.com/google/flax/issues/231), we decided\nto go for a standard pattern for all examples including the simplest ones\n(like MNIST). This makes every example a bit more verbose, but once you know\none example, you know the structure of all of them. Having unit tests and\nintegration tests is also very useful when you fork these examples.\n\nFor more examples including contributions from the community and other projects currently using Flax see the **[Examples](https://flax.readthedocs.io/en/latest/examples.html)** section in the documentation.\n" }, { "path": "examples/__init__.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n" }, { "path": "examples/cloud/README.md", "content": "# Launching jobs on Google Cloud\n\nThis directory provides a simple script that can be used to create a new VM\non Google Cloud, train an example on that VM and then shutting it down.\n\nThe training is implemented in a shell that is run on the VM after startup by\nsetting the `startup_script_file` in the metadata. The script opens a TMUX\nsession, installs Flax repository from Github with all dependencies, and then\nruns the training in parallel with `gcloud storage rsync` that copies the training\nartifacts in a storage bucket.\n\nThe advantage of this approach is that every training is run in a single VM\nthat contains all code and configuration, so it is easy to run multiple\nexperiments in parallel without interference. The individual trainings can also\nbe inspected by logging into the VM via SSH and attaching to the tmux session.\n\nThe script `launch_gce.py` launches the VM and prints out the relevant commands\nto track the progress update and to log into the machine.\n\nNote that the VM also shuts down if an error is encountered, after waiting for\nfive minutes.\n\n## Preparation\n\nPrerequisites:\n\n1. Create a Google Cloud account.\n2. Set up billing: https://console.cloud.google.com/billing\n3. Create a storage bucket (GCS).\n4. Optional: Get quota for accelerators. This is usually granted with a short\n delay: https://console.cloud.google.com/iam-admin/quotas\n\n## Setting up your environment\n\nThe commands below use the same set of pre-defined environment variables.\n\nMandatory environment variables:\n\n- `$PROJECT`: The name of your Google Cloud project.\n- `$GCS_BUCKET`: The name of the Google Cloud Storage bucket where the model\n output (artifacts, final checkpoint) are to be stored.\n- `$ZONE`: Compute zone (e.g. `central1-a`).\n\nOptional environment variables:\n\n- `$REPO`: Alternative repo to use instead of the default\n https://github.com/google/flax - this is useful for development.\n- `$BRANCH`: Alternative branch to use instead of the default `main`.\n\n## Training the MNIST example\n\nUse the following command to launch the MNIST example on cloud (make sure to set\n`$PROJECT` and `$GCS_BUCKET` accordingly):\n\n```shell\npython examples/cloud/launch_gce.py \\\n --project=$PROJECT \\\n --zone=us-west1-a \\\n --machine_type=n2-standard-2 \\\n --gcs_workdir_base=gs://$GCS_BUCKET/workdir_base \\\n --repo=${REPO:-https://github.com/google/flax} \\\n --branch=${BRANCH:-main} \\\n --example=mnist \\\n --args='--config=configs/default.py' \\\n --name=default\n```\n\n## Training the imagenet example\n\nNote that you need to first prepare the `imagenet2012` dataset. For this,\ndownload the data from http://image-net.org/ as described in the\n[tensorflow_datasets catalog](https://www.tensorflow.org/datasets/catalog/imagenet2012).\nThen point the environment variable `$IMAGENET_DOWNLOAD_PATH` to the directory\nwhere the downloads are stored and prepare the dataset by running\n\n```shell\npython -c \"\nimport tensorflow_datasets as tfds\ntfds.builder('imagenet2012').download_and_prepare(\n download_config=tfds.download.DownloadConfig(\n manual_dir='$IMAGENET_DOWNLOAD_PATH'))\n\"\n```\n\nThen copy the contents of the directory `~/tensorflow_datasets` into the\ndirectory `gs://$GCS_TFDS_BUCKET/datasets` (note that `$GCS_TFDS_BUCKET` and\n`$GCS_BUCKET` can be identical).\n\nAfter this preparation you can run the imagenet example with the following\ncommand (make sure to set `$PROJECT`, `$GCS_BUCKET` and `$GCS_TFDS_BUCKET`\naccordingly):\n\n```shell\npython examples/cloud/launch_gce.py \\\n --project=$PROJECT \\\n --zone=us-west1-a \\\n --machine_type=n1-standard-96 \\\n --accelerator_type=nvidia-tesla-v100 --accelerator_count=8 \\\n --gcs_workdir_base=gs://$GCS_BUCKET/workdir_base \\\n --tfds_data_dir=gs://$GCS_TFDS_BUCKET/datasets \\\n --repo=${REPO:-https://github.com/google/flax} \\\n --branch=${BRANCH:-main} \\\n --example=imagenet \\\n --args='--config=configs/v100_x8_mixed_precision.py' \\\n --name=v100_x8_mixed_precision\n```\n\n## Tips\n\nYou can add `--connect` to above commands to directly land in the training\nsession once the VM is ready. This is very helpful for debugging when changing\nthings. Note that the VM automatically shuts down after 5 minutes of inactivity,\nboth in case of success as in case of failure. On OS X this could be combined\nwith `VM_READY_CMD=\"osascript -e 'display notification \\\"VM ready\\\"'\"` so get\nundistracted when the VM is up and running.\n\nWhen tweaking the startup script or individual arguments, it is often helpful to\nconnect to the VM, stop the scripts and end the tmux session, and then copy and\npaste the contents of the generated `flax-...-startup_script.sh`, after\nmodifying these contents accordingly." }, { "path": "examples/cloud/launch_gce.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# Script for creating VM on Google cloud and run a Flax example inside.\n# See ./README.md for instructions.\n\nimport datetime\nimport os\nimport re\nimport subprocess\nimport time\nfrom collections.abc import Sequence\n\nfrom absl import app\nfrom absl import flags\n\n# General options.\nflags.DEFINE_bool(\n 'dry_run',\n False,\n help=(\n 'If set, then the command to launch the GCE instance will only be '\n 'printed to stdout.'\n ),\n)\nflags.DEFINE_bool(\n 'connect',\n False,\n help='Same as --wait, but directly connect to VM once it is ready.',\n)\nflags.DEFINE_bool(\n 'wait',\n False,\n help=(\n 'If set, then the script will wait until VM is ready. If VM_READY_CMD '\n 'is set in environment, then that command will be executed once the VM '\n 'is ready. Useful for sending a notification, e.g. \"osascript\" (mac).'\n ),\n)\n\n# Machine configuration.\nflags.DEFINE_string('project', None, help='Name of the Google Cloud project.')\nflags.DEFINE_string('zone', None, help='Zone in which the VM will be created.')\nflags.DEFINE_string(\n 'machine_type',\n None,\n help='Machine type to use for VM. See \"gcloud compute machine-types list\".',\n)\nflags.DEFINE_string(\n 'accelerator_type',\n '',\n help=(\n 'Type of accelerator to use, or empty. '\n 'See \"gcloud compute accelerator-types list\".'\n ),\n)\nflags.DEFINE_integer(\n 'shutdown_secs',\n 300,\n help=(\n 'How long to wait (after successful/failed training) before shutting '\n 'down the VM. Set to 0 to disable.'\n ),\n)\nflags.DEFINE_integer(\n 'accelerator_count', 8, help='Number of accelerators to use.'\n)\n\n# GCS configuration.\nflags.DEFINE_string(\n 'gcs_workdir_base',\n None,\n help=(\n 'GCS base directory for model output. The --workdir argument will be '\n 'constructed from {gcs_workdir_base}/{example}/{name}/{timestamp} .'\n ),\n)\nflags.DEFINE_string(\n 'tfds_data_dir',\n '',\n help=(\n 'Optional tfds data directory. This can be useful to prepare datasets'\n ' on GCS and then point the jobs to this preloaded directory. Dataset'\n ' will be downloaded from the web if not specified.'\n ),\n)\n\n# Repo configuration.\nflags.DEFINE_string(\n 'repo', 'https://github.com/google/flax', help='Git repository'\n)\nflags.DEFINE_string('branch', 'main', help='Git repository')\n\n# Example configuration.\nflags.DEFINE_string(\n 'example', None, help='Name of Flax example (e.g. \"imagenet\").'\n)\nflags.DEFINE_string(\n 'args',\n '',\n help=(\n 'Any additional command line arguments for {example}_main.py, like '\n 'for example --config. Note that --workdir will be provided by the '\n 'script.'\n ),\n)\n\n# Run configuration.\nflags.DEFINE_string(\n 'name',\n None,\n help=(\n 'Name of the experiment. Note that the provided name will be '\n 'extended to {example}/{name}/{timestamp}'\n ),\n)\n\nFLAGS = flags.FLAGS\nflags.mark_flags_as_required(\n ['project', 'zone', 'machine_type', 'gcs_workdir_base', 'example', 'name']\n)\n\ntimestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')\n\n\ndef generate_startup_file(vm_name: str) -> str:\n directory = os.path.dirname(os.path.abspath(__file__))\n startup_script_src = os.path.join(directory, 'startup_script.sh')\n startup_script_dst = os.path.join(directory, f'{vm_name}-startup_script.sh')\n assert not os.path.exists(startup_script_dst)\n with open(startup_script_src, encoding='utf8') as f:\n startup_script_content = f.read()\n for from_str, to_str in (\n ('__REPO__', FLAGS.repo),\n ('__BRANCH__', FLAGS.branch),\n ('__EXAMPLE__', FLAGS.example),\n ('__TIMESTAMP__', timestamp),\n ('__NAME__', FLAGS.name),\n ('__ARGS__', FLAGS.args),\n ('__GCS_WORKDIR_BASE__', FLAGS.gcs_workdir_base),\n ('__TFDS_DATA_DIR__', FLAGS.tfds_data_dir),\n ('__ACCELERATOR_TYPE__', FLAGS.accelerator_type),\n ('__SHUTDOWN_SECS__', str(FLAGS.shutdown_secs)),\n ):\n startup_script_content = startup_script_content.replace(from_str, to_str)\n with open(startup_script_dst, 'w', encoding='utf8') as f:\n f.write(startup_script_content)\n return startup_script_dst\n\n\ndef launch_gce(*, vm_name: str, startup_script: str):\n # Note : Use `gcloud compute images list --project ml-images` to get a list\n # of available VM images.\n args = [\n 'gcloud',\n 'compute',\n 'instances',\n 'create',\n vm_name,\n f'--project={FLAGS.project}',\n f'--zone={FLAGS.zone}',\n '--image=c1-deeplearning-tf-2-10-cu113-v20221107-debian-10',\n '--image-project=ml-images',\n f'--machine-type={FLAGS.machine_type}',\n '--scopes=cloud-platform,storage-full',\n '--boot-disk-size=256GB',\n '--boot-disk-type=pd-ssd',\n '--metadata=install-nvidia-driver=True',\n f'--metadata-from-file=startup-script={startup_script}',\n ]\n if FLAGS.accelerator_type and FLAGS.accelerator_count:\n args.extend([\n '--maintenance-policy=TERMINATE',\n f'--accelerator=type={FLAGS.accelerator_type},count={FLAGS.accelerator_count}',\n ])\n\n if FLAGS.dry_run:\n print()\n print('Would run the following command without --dry-run:')\n print()\n print(' \\\\\\n '.join(args))\n print()\n return\n\n print()\n print('Creating instance on GCE... This will take some minutes...')\n print()\n result = subprocess.run(args)\n if result.returncode:\n raise RuntimeError('Could not create VM!')\n\n\ndef print_howto(login_args: Sequence[str]):\n print(f\"\"\"\n###############################################################################\n###############################################################################\n\nYou can start/stop the instace via the web UI:\nhttps://console.cloud.google.com/compute/instances?project={FLAGS.project}\n\nOnce the VM has started, you can login and connect to the training session:\n\n{' '.join(login_args)}\n\nNote that you can disconnect from the tmux session without stopping the training\nwith the keystrokes 'CTRL-B A'. See \"man tmux\" for help about tmux.\n\nTo observe the training via Tensorboard, simply run in your local computer:\n\n$ tensorboard --logdir={FLAGS.gcs_workdir_base}\n\nYou can also browse the files at\n\nhttps://console.cloud.google.com/storage/browser/{FLAGS.gcs_workdir_base.replace('gs://', '')}\n\n###############################################################################\n###############################################################################\n\"\"\")\n\n\ndef main(_):\n for name in ('repo', 'branch', 'example', 'name', 'gcs_workdir_base'):\n value = getattr(FLAGS, name)\n if re.match(r'[^\\w:/_-]', value):\n raise ValueError(f'Invalid flag value: --{name}=\"{value}\"')\n example_base_directory = os.path.join(\n os.path.dirname(os.path.abspath(__file__)),\n os.path.pardir,\n )\n if not os.path.isdir(os.path.join(example_base_directory, FLAGS.example)):\n raise ValueError(f'Could not find --example={FLAGS.example}')\n if FLAGS.connect and FLAGS.dry_run:\n raise ValueError('Cannot --connect to VM with --dry_run')\n\n vm_name = '-'.join([\n 'flax',\n FLAGS.example,\n timestamp,\n ])\n vm_name = re.sub(r'[^a-z0-9-]', '-', vm_name)\n\n startup_script = generate_startup_file(vm_name)\n launch_gce(vm_name=vm_name, startup_script=startup_script)\n\n login_args = [\n 'gcloud',\n 'compute',\n 'ssh',\n '--project',\n FLAGS.project,\n '--zone',\n FLAGS.zone,\n vm_name,\n '--',\n '/sudo_tmux_a.sh',\n ]\n\n print('Your instance is being started...')\n\n print_howto(login_args)\n\n if FLAGS.connect or FLAGS.wait:\n login_true_args = login_args[:-1] + ['true']\n while True:\n try:\n result = subprocess.run(\n login_true_args, timeout=10, stderr=subprocess.PIPE\n )\n if result.returncode == 0:\n break\n stderr = result.stderr.decode('utf8')\n if 'connection refused' in stderr.lower():\n print('(Connection refused - waiting a little longer...)')\n time.sleep(20)\n else:\n raise ValueError(f'Unknown error: {stderr}')\n except ValueError as e:\n if 'HTTP 502' not in str(e):\n raise e\n print('(Bad Gateway - waiting a little longer...)')\n time.sleep(20)\n except subprocess.TimeoutExpired:\n print('(Timeout - waiting a little longer...)')\n time.sleep(20)\n if 'VM_READY_CMD' in os.environ:\n os.system(os.environ['VM_READY_CMD'])\n if FLAGS.connect:\n result = subprocess.run(login_args)\n # SSH session has cleared previous message, print it again.\n print_howto(login_args)\n\n\nif __name__ == '__main__':\n app.run(main)\n" }, { "path": "examples/cloud/startup_script.sh", "content": "#!/bin/bash\n\n# Note that all __XYZ__ strings are replaced by launch_gce.py\n\nWORKDIR=\"/train/workdir_base/__EXAMPLE__/__NAME__/__TIMESTAMP__\"\n\nmkdir -p /train\ncd /train\n\n# Login directly with:\n# gcloud compute ssh $VM -- /sudo_tmux_a.sh\necho -e '#!/bin/bash\\nsudo /tmux_a.sh' > /sudo_tmux_a.sh\nchmod a+x /sudo_tmux_a.sh\necho -e '#!/bin/bash\\ntmux a' > /tmux_a.sh\nchmod a+x /tmux_a.sh\n\n# Main script running in bottom left tmux pane.\ncat >/install_train_stop.sh <&1 | tee -a $WORKDIR/setup_train_log_${TIMESTAMP}.txt\n\nif [ __SHUTDOWN_SECS__ -gt 0 ]; then\n echo\n echo WILL SHUT DOWN IN $((__SHUTDOWN_SECS__/60)) MIN ...\n sleep __SHUTDOWN_SECS__ && shutdown now\nfi\n\nEOF\n\n\n# Set up TMUX panes:\ntmux new-session -s flax -d\n# - top left: htop\ntmux send 'htop\n'\ntmux split-window\ntmux selectp -U\ntmux split-window -h\n# - top right: htop\ntmux send 'watch nvidia-smi\n'\ntmux selectp -D\n# - bottom left: main script\ntmux send '. /install_train_stop.sh\n'\ntmux split-window -h\n# - bottom right: rsync files to GCS bucket.\ntmux send \"\nwhile true; do\n gcloud storage rsync --recursive workdir_base __GCS_WORKDIR_BASE__\n sleep 60\ndone 2>&1 | tee -a $WORKDIR/gcs_rsync_'__TIMESTAMP__'.txt\n\"\n" }, { "path": "examples/gemma/README.md", "content": "\n## Language modeling\nTrains Gemma model on the One Billion Word Benchmark (lm1b; Chelba *et al.*, 2013).\n\nThis example is based on lm1b and similarly uses linear learning rate warmup and inverse square root learning rate schedule.\n\n\n### Requirements\n\n* TensorFlow datasets `lm1b` need to be downloaded and prepared (see below).\n A sentencepiece tokenizer vocabulary will be automatically generated\n and saved on each training run.\n* This example additionally depends on the `sentencepiece` and `tensorflow-text` packages.\n\n### Downloading the LM1B Datasets\n\nWe recommend downloading and preparing the TFDS datasets beforehand. You can download and prepare LM1B datasets using TFDS directly: `python -m tensorflow_datasets.scripts.download_and_prepare --datasets=lm1b`.\n\n#### Using Cloud Storage FUSE for TPUs\n\nFor Cloud TPUs, we recommend using a cheap standard instance and saving the prepared TFDS\ndata on a storage bucket, from where it can be mounted to the TPU VM using [Cloud Storage FUSE](https://cloud.google.com/storage/docs/cloud-storage-fuse/quickstart-mount-bucket).\n\n##### Copy the preprocessed dataset to the Cloud Storage\n\nWe assume that the dataset was downloaded and prepared. We also assume we have configured `gcloud` CLI. The following commands helps to setup the storage and copy the dataset:\n\n```bash\n# Install gcsfuse CLI\nexport GCSFUSE_REPO=gcsfuse-`lsb_release -c -s`\n# For example, GCSFUSE_REPO=gcsfuse-noble for Ubuntu 24.04\n\necho \"deb https://packages.cloud.google.com/apt $GCSFUSE_REPO main\" | sudo tee /etc/apt/sources.list.d/gcsfuse.list\ncurl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -\nsudo apt-get update\nsudo apt-get install -y fuse gcsfuse --no-install-recommends\n\ngcsfuse -v\n# gcsfuse version 2.12.2 (Go version go1.24.0)\n```\n\nLet's get where LM1B dataset was locally stored:\n```bash\npython -c \"import tensorflow_datasets as tfds; b=tfds.builder('lm1b'); print(b.info.data_dir)\"\n# For example: /home/user/tensorflow_datasets/lm1b/1.1.0\n```\n\nLet's create a GCS bucket for the dataset and link the bucket to a local folder. We choose the bucket name \"flax-lm1b-tfdataset\" but this can be changed.\n```bash\ngcloud storage buckets create gs://flax-lm1b-tfdataset\n\nmkdir -p $HOME/data\ngcsfuse flax-lm1b-tfdataset $HOME/data\n```\n\nNow let's copy the data to the bucket:\n```bash\n# Let's assume that prepared dataset is at $HOME/tensorflow_datasets/lm1b/\ncp -R $HOME/tensorflow_datasets/lm1b $HOME/data\n```\n\n##### Setup the dataset on TPU VM\n\nWe previously have choosen the bucket name \"flax-lm1b-tfdataset\" where stored the dataset, adapt this name to your situation.\n\n```bash\n# On the TPU VM\ngcsfuse flax-lm1b-tfdataset $HOME/tensorflow_datasets\n\nls $HOME/tensorflow_datasets/lm1b/1.1.0/\n```\n\n### How to run on GPU(s)\n\nInstall Jax with CUDA support, Flax and the example dependencies with the following command:\n```bash\npip install jax[cuda12]\n# Check whether GPUs are available:\n# python3 -c \"import jax; print(jax.devices())\"\n\ngit clone --depth=1 --branch=main https://github.com/google/flax\ncd flax\npip install -e .\ncd examples/gemma\npip install -r requirements.txt\n```\n\nStart the training:\n\n- train a small transformer model:\n```bash\npython3 main.py --workdir=$HOME/logs/small_gemma_lm1b --config=configs/small.py\n```\n\n- train Gemma3-4B model:\n```bash\npython3 main.py --workdir=$HOME/logs/gemma3-4b_lm1b --config=configs/gemma3_4b.py\n```\n\nTo monitor the trainings with the TensorBoard:\n```bash\ntensorboard --logdir=$HOME/logs\n```\n\n\n### How to run on Cloud TPUs\n\nSetup the TPU VM and install the Flax dependencies on it as described\n[here](https://cloud.google.com/tpu/docs/jax-pods) for creating pod slices, or\n[here](https://cloud.google.com/tpu/docs/jax-quickstart-tpu-vm) for a single\nv4-8 TPU.\n\n\nFirst create a single TPUv4-8 VM and connect to it (you can find more detailed\ninstructions [here](https://cloud.google.com/tpu/docs/jax-quickstart-tpu-vm)):\n\n```bash\nZONE=us-central1-a\nTPU_TYPE=v4-8\nTPU_NAME=$USER-flax-gemma-lm1b\ngcloud compute tpus tpu-vm create $TPU_NAME \\\n --zone $ZONE \\\n --accelerator-type $TPU_TYPE \\\n --version tpu-ubuntu2204-base\n\ngcloud compute tpus tpu-vm ssh $TPU_NAME --zone $ZONE -- \\\n -L 6006:localhost:6006\n```\n\nWhen connected install JAX:\n\n```bash\npip install \"jax[tpu]>=0.2.16\" \\\n -f https://storage.googleapis.com/jax-releases/libtpu_releases.html\n```\n\nThen install Flax + the example dependencies:\n\n```bash\ngit clone --depth=1 --branch=main https://github.com/google/flax\ncd flax\npip install -e .\ncd examples/gemma\npip install -r requirements.txt\n```\n\nIn case of errors when installing example dependencies, try to upgrade existing `pip` package and downgrade `setuptools` and repeat the installation command\n```bash\n# Optionally\n# pip install -U pip\n# pip install -U \"setuptools<70\"\n# pip install -r requirements.txt\n```\n\nAnd finally start the training:\n\n```bash\npython3 main.py --workdir=$HOME/logs/gemma_lm1b_256 --config.per_device_batch_size=32\n```\n\nNote that you might want to set `TFDS_DATA_DIR` as explained below. You probably\nalso want to start the long-running command above in a `tmux` session and start\nsome monitoring in a separate pane (note that we forwarded port 6006 locally\nabove):\n\n```bash\ntensorboard --logdir=$HOME/logs\n```\n" }, { "path": "examples/gemma/configs/default.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Default Hyperparameter configuration.\"\"\"\n\nimport dataclasses\n\nfrom train import MeshRules, TrainConfig\n\n\n@dataclasses.dataclass(unsafe_hash=True)\nclass Config:\n # Path to load or store sentencepiece vocab file.\n vocab_path: str | None = None\n # Vocabulary size if `vocab_path` is not given.\n vocab_size: int = 35_000 # lm1b dataset vocab size: 35913 (Gemma expected vocab size: 262_144)\n # Maximum number of characters to use for training.\n max_corpus_chars: int = 10**7\n # Name of TFDS translation dataset to use.\n dataset_name: str = 'lm1b'\n # Optional name of TFDS translation dataset to use for evaluation.\n eval_dataset_name: str = 'lm1b'\n # Optional name of TFDS split to use for evaluation.\n eval_split: str = 'test'\n # Per device batch size for training.\n per_device_batch_size: int = 32\n # Per device batch size for training.\n eval_per_device_batch_size: int = 32\n\n # Prompt for language model sampling\n prompts: tuple[str, ...] = (\n 'Paris is a the capital',\n 'Flax is a',\n # From train set:\n 'The shutdown was aimed at creating efficiencies as',\n # -> the plant was already operating at its maximum capacity of 3,000 tonnes of cellulose paste per day\n 'A big theme of this hire is that there are parts of',\n # -> our operations that to use a pretty trite phrase , need to be taken to the next level ...\n\n # From test set:\n 'Because of Bear Stearns , many analysts are',\n # -> raising the odds that a 2008 recession could be worse than expected\n 'Next month , the Brazilian bourse',\n # -> opens a London office',\n )\n # Temperature for top_p sampling.\n sampling_temperature: float = 0.0\n # Top-p sampling threshold.\n sampling_top_p: float = 0.95\n\n # Number of steps to take during training.\n num_train_steps: int = 500_000\n # Number of steps to take during evaluation.\n # Large enough to evaluate all samples: 306_688 / (32 * 8) = 1198\n num_eval_steps: int = 2_000\n # Number of steps to generate predictions.\n # -1 will use the whole eval dataset.\n num_predict_steps: int = 50\n # Base learning rate.\n learning_rate: float = 0.0016\n # Linear learning rate warmup.\n warmup_steps: int = 1000\n # Cross entropy loss label smoothing.\n label_smoothing: float = 0.0\n # Decay factor for AdamW style weight decay.\n weight_decay: float = 0.1\n # Maximum length cutoff for training examples.\n max_target_length: int = 128\n # Maximum length cutoff for eval examples.\n max_eval_target_length: int = 512\n\n # Gemma transformer name.\n # Possible values defined in transformer.TransformerConfig:\n # (gemma_2b, gemma_7b, gemma2_2b, gemma2_9b, gemma2_27b, gemma3_1b, gemma3_4b, ...)\n transformer_name: str | None = \"gemma3_1b\"\n # or alternatively define the model using the dict of parameters\n transformer_params: dict | None = None\n\n # Whether to save model checkpoints.\n save_checkpoints: bool = True\n # Whether to restore from existing model checkpoints.\n restore_checkpoints: bool = True\n # Save a checkpoint every these number of steps.\n checkpoint_every_steps: int = 10_000\n # Frequency of eval during training, e.g. every 1_000 steps.\n eval_every_steps: int = 5_000\n # Use bfloat16 mixed precision training instead of float32.\n use_bfloat16: bool = True\n # Integer for PRNG random seed.\n seed: int = 0\n\n # Parallelism\n mesh_axes: tuple[str, ...] = ('data', 'fsdp', 'tensor')\n axis_rules: MeshRules = MeshRules(\n embed='fsdp',\n mlp='tensor',\n kv='tensor',\n vocab='tensor',\n )\n data_sharding: tuple[str, ...] = ('data', 'fsdp')\n\n # One axis for each parallelism type may hold a placeholder (-1)\n # value to auto-shard based on available slices and devices.\n # By default, product of the DCN axes should equal number of slices\n # and product of the ICI axes should equal number of devices per slice.\n # ICI (Inter-Chip Interconnection): A high-speed connection between\n # sets of TPU chips, which form the TPU network.\n # DCN (Data Center Network): A connection between the TPU networks;\n # not as fast as ICI.\n # ICI has around 100x the bandwidth of DCN, but it is not a general\n # purpose connection, which is why DCN is necessary for scaling to\n # extremely large ML models.\n dcn_data_parallelism: int = -1\n dcn_fsdp_parallelism: int = 1\n dcn_tensor_parallelism: int = 1\n ici_data_parallelism: int = 1\n ici_fsdp_parallelism: int = -1\n ici_tensor_parallelism: int = 1\n\n\ndef get_config() -> TrainConfig:\n \"\"\"Get the default hyperparameter configuration.\"\"\"\n config = Config()\n return TrainConfig(**dataclasses.asdict(config))\n" }, { "path": "examples/gemma/configs/gemma3_4b.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Default Hyperparameter configuration.\"\"\"\n\nimport dataclasses\n\nfrom train import MeshRules, TrainConfig\n\n\n@dataclasses.dataclass(unsafe_hash=True)\nclass Config:\n # Path to load or store sentencepiece vocab file.\n vocab_path: str | None = None\n # Vocabulary size if `vocab_path` is not given.\n vocab_size: int = 35_000 # lm1b dataset vocab size: 35913 (Gemma expected vocab size: 262_144)\n # Maximum number of characters to use for training.\n max_corpus_chars: int = 10**7\n # Name of TFDS translation dataset to use.\n dataset_name: str = 'lm1b'\n # Optional name of TFDS translation dataset to use for evaluation.\n eval_dataset_name: str = 'lm1b'\n # Optional name of TFDS split to use for evaluation.\n eval_split: str = 'test'\n # Per device batch size for training.\n per_device_batch_size: int = 32\n # Per device batch size for training.\n eval_per_device_batch_size: int = 32\n\n # Prompt for language model sampling\n prompts: tuple[str, ...] = (\n 'Paris is a the capital',\n 'Flax is a',\n # From train set:\n 'The shutdown was aimed at creating efficiencies as',\n # -> the plant was already operating at its maximum capacity of 3,000 tonnes of cellulose paste per day\n 'A big theme of this hire is that there are parts of',\n # -> our operations that to use a pretty trite phrase , need to be taken to the next level ...\n\n # From test set:\n 'Because of Bear Stearns , many analysts are',\n # -> raising the odds that a 2008 recession could be worse than expected\n 'Next month , the Brazilian bourse',\n # -> opens a London office',\n )\n # Temperature for top_p sampling.\n sampling_temperature: float = 0.0\n # Top-p sampling threshold.\n sampling_top_p: float = 0.95\n\n # Number of steps to take during training.\n num_train_steps: int = 500_000\n # Number of steps to take during evaluation.\n # Large enough to evaluate all samples: 306_688 / (32 * 8) = 1198\n num_eval_steps: int = 2_000\n # Number of steps to generate predictions.\n # -1 will use the whole eval dataset.\n num_predict_steps: int = 50\n # Base learning rate.\n learning_rate: float = 0.0016\n # Linear learning rate warmup.\n warmup_steps: int = 1000\n # Cross entropy loss label smoothing.\n label_smoothing: float = 0.0\n # Decay factor for AdamW style weight decay.\n weight_decay: float = 0.1\n # Maximum length cutoff for training examples.\n max_target_length: int = 128\n # Maximum length cutoff for eval examples.\n max_eval_target_length: int = 512\n\n # Gemma transformer name.\n # Possible values defined in transformer.TransformerConfig:\n # (gemma_2b, gemma_7b, gemma2_2b, gemma2_9b, gemma2_27b, gemma3_1b, gemma3_4b, ...)\n transformer_name: str | None = \"gemma3_4b\"\n # or alternatively define the model using the dict of parameters\n transformer_params: dict | None = None\n\n # Whether to save model checkpoints.\n save_checkpoints: bool = True\n # Whether to restore from existing model checkpoints.\n restore_checkpoints: bool = True\n # Save a checkpoint every these number of steps.\n checkpoint_every_steps: int = 10_000\n # Frequency of eval during training, e.g. every 1_000 steps.\n eval_every_steps: int = 5_000\n # Use bfloat16 mixed precision training instead of float32.\n use_bfloat16: bool = True\n # Integer for PRNG random seed.\n seed: int = 0\n\n # Parallelism\n mesh_axes: tuple[str, ...] = ('data', 'fsdp', 'tensor')\n axis_rules: MeshRules = MeshRules(\n embed='fsdp',\n mlp='tensor',\n kv='tensor',\n vocab='tensor',\n )\n data_sharding: tuple[str, ...] = ('data', 'fsdp')\n\n # One axis for each parallelism type may hold a placeholder (-1)\n # value to auto-shard based on available slices and devices.\n # By default, product of the DCN axes should equal number of slices\n # and product of the ICI axes should equal number of devices per slice.\n # ICI (Inter-Chip Interconnection): A high-speed connection between\n # sets of TPU chips, which form the TPU network.\n # DCN (Data Center Network): A connection between the TPU networks;\n # not as fast as ICI.\n # ICI has around 100x the bandwidth of DCN, but it is not a general\n # purpose connection, which is why DCN is necessary for scaling to\n # extremely large ML models.\n dcn_data_parallelism: int = -1\n dcn_fsdp_parallelism: int = 1\n dcn_tensor_parallelism: int = 1\n ici_data_parallelism: int = 1\n ici_fsdp_parallelism: int = -1\n ici_tensor_parallelism: int = 1\n\n def replace(self, **kwargs):\n return dataclasses.replace(self, **kwargs)\n\n\ndef get_config() -> TrainConfig:\n \"\"\"Get the default hyperparameter configuration.\"\"\"\n config = Config()\n return TrainConfig(**dataclasses.asdict(config))\n" }, { "path": "examples/gemma/configs/small.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Default Hyperparameter configuration.\"\"\"\n\nimport dataclasses\n\nfrom train import MeshRules, TrainConfig\n\n\n@dataclasses.dataclass(unsafe_hash=True)\nclass Config:\n # Path to load or store sentencepiece vocab file.\n vocab_path: str | None = None\n # Vocabulary size if `vocab_path` is not given.\n vocab_size: int = 35_000 # lm1b dataset vocab size: 35913 (Gemma expected vocab size: 262_144)\n # Maximum number of characters to use for training.\n max_corpus_chars: int = 10**7\n # Name of TFDS translation dataset to use.\n dataset_name: str = 'lm1b'\n # Optional name of TFDS translation dataset to use for evaluation.\n eval_dataset_name: str = 'lm1b'\n # Optional name of TFDS split to use for evaluation.\n eval_split: str = 'test'\n # Per device batch size for training.\n per_device_batch_size: int = 32\n # Per device batch size for training.\n eval_per_device_batch_size: int = 32\n\n # Prompt for language model sampling\n prompts: tuple[str, ...] = (\n 'Paris is a the capital',\n 'Flax is a',\n # From train set:\n 'The shutdown was aimed at creating efficiencies as',\n # -> the plant was already operating at its maximum capacity of 3,000 tonnes of cellulose paste per day\n 'A big theme of this hire is that there are parts of',\n # -> our operations that to use a pretty trite phrase , need to be taken to the next level ...\n\n # From test set:\n 'Because of Bear Stearns , many analysts are',\n # -> raising the odds that a 2008 recession could be worse than expected\n 'Next month , the Brazilian bourse',\n # -> opens a London office',\n )\n # Temperature for top_p sampling.\n sampling_temperature: float = 0.0\n # Top-p sampling threshold.\n sampling_top_p: float = 0.95\n\n # Number of steps to take during training.\n num_train_steps: int = 500_000\n # Number of steps to take during evaluation.\n num_eval_steps: int = 500\n # Number of steps to generate predictions.\n # -1 will use the whole eval dataset.\n num_predict_steps: int = 50\n # Base learning rate.\n learning_rate: float = 0.0016\n # Linear learning rate warmup.\n warmup_steps: int = 1000\n # Cross entropy loss label smoothing.\n label_smoothing: float = 0.0\n # Decay factor for AdamW style weight decay.\n weight_decay: float = 0.1\n # Maximum length cutoff for training examples.\n max_target_length: int = 128\n # Maximum length cutoff for eval examples.\n max_eval_target_length: int = 512\n\n # Gemma transformer name.\n # Possible values defined in transformer.TransformerConfig:\n # (gemma_2b, gemma_7b, gemma2_2b, gemma2_9b, gemma2_27b, gemma3_1b, gemma3_4b, ...)\n transformer_name: str | None = None\n # or alternatively define the model using the dict of parameters\n transformer_params: dict | None = dataclasses.field(\n default_factory=lambda: {\n \"num_layers\": 6,\n \"embed_dim\": 512,\n \"hidden_dim\": 2048,\n \"num_heads\": 4,\n \"head_dim\": 256,\n \"num_kv_heads\": 1,\n \"use_post_attn_norm\": True,\n \"use_post_ffw_norm\": True,\n \"use_qk_norm\": True,\n\n \"attention_types\": (2, 2, 2, 2, 2, 1), # local_sliding, ..., local_sliding, global\n \"query_pre_attn_norm\": 1, # QueryPreAttentionNormalisation.BY_ONE_OVER_SQRT_HEAD_DIM\n\n \"attn_logits_soft_cap\": None,\n \"final_logit_softcap\": None,\n \"sliding_window_size\": 128,\n \"transpose_gating_einsum\": True,\n \"local_base_frequency\": 10_000,\n \"global_base_frequency\": 1_000_000,\n }\n )\n\n # Whether to save model checkpoints.\n save_checkpoints: bool = True\n # Whether to restore from existing model checkpoints.\n restore_checkpoints: bool = True\n # Save a checkpoint every these number of steps.\n checkpoint_every_steps: int = 10_000\n # Frequency of eval during training, e.g. every 1_000 steps.\n eval_every_steps: int = 5_000\n # Use bfloat16 mixed precision training instead of float32.\n use_bfloat16: bool = True\n # Integer for PRNG random seed.\n seed: int = 0\n\n # Parallelism\n mesh_axes: tuple[str, ...] = ('data', 'fsdp', 'tensor')\n axis_rules: MeshRules = MeshRules(\n embed='fsdp',\n mlp='tensor',\n kv='tensor',\n vocab='tensor',\n )\n data_sharding: tuple[str, ...] = ('data', 'fsdp')\n\n # One axis for each parallelism type may hold a placeholder (-1)\n # value to auto-shard based on available slices and devices.\n # By default, product of the DCN axes should equal number of slices\n # and product of the ICI axes should equal number of devices per slice.\n # ICI (Inter-Chip Interconnection): A high-speed connection between\n # sets of TPU chips, which form the TPU network.\n # DCN (Data Center Network): A connection between the TPU networks;\n # not as fast as ICI.\n # ICI has around 100x the bandwidth of DCN, but it is not a general\n # purpose connection, which is why DCN is necessary for scaling to\n # extremely large ML models.\n dcn_data_parallelism: int = -1\n dcn_fsdp_parallelism: int = 1\n dcn_tensor_parallelism: int = 1\n ici_data_parallelism: int = 1\n ici_fsdp_parallelism: int = -1\n ici_tensor_parallelism: int = 1\n\n def replace(self, **kwargs):\n return dataclasses.replace(self, **kwargs)\n\n\ndef get_config() -> TrainConfig:\n \"\"\"Get the default hyperparameter configuration.\"\"\"\n config = Config()\n return TrainConfig(**dataclasses.asdict(config))\n" }, { "path": "examples/gemma/configs/tiny.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Default Hyperparameter configuration.\"\"\"\n\nimport dataclasses\n\nfrom train import MeshRules, TrainConfig\n\n\n@dataclasses.dataclass(unsafe_hash=True)\nclass Config:\n # Path to load or store sentencepiece vocab file.\n vocab_path: str | None = None\n # Vocabulary size if `vocab_path` is not given.\n vocab_size: int = 35_000 # lm1b dataset vocab size: 35913 (Gemma expected vocab size: 262_144)\n # Maximum number of characters to use for training.\n max_corpus_chars: int = 10**7\n # Name of TFDS translation dataset to use.\n dataset_name: str = 'lm1b'\n # Optional name of TFDS translation dataset to use for evaluation.\n eval_dataset_name: str = 'lm1b'\n # Optional name of TFDS split to use for evaluation.\n eval_split: str = 'test'\n # Per device batch size for training.\n per_device_batch_size: int = 32\n # Per device batch size for training.\n eval_per_device_batch_size: int = 32\n\n # Prompt for language model sampling\n prompts: tuple[str, ...] = (\n 'Paris is a the capital',\n 'Flax is a',\n # From train set:\n 'The shutdown was aimed at creating efficiencies as',\n # -> the plant was already operating at its maximum capacity of 3,000 tonnes of cellulose paste per day\n 'A big theme of this hire is that there are parts of',\n # -> our operations that to use a pretty trite phrase , need to be taken to the next level ...\n\n # From test set:\n 'Because of Bear Stearns , many analysts are',\n # -> raising the odds that a 2008 recession could be worse than expected\n 'Next month , the Brazilian bourse',\n # -> opens a London office',\n )\n # Temperature for top_p sampling.\n sampling_temperature: float = 0.0\n # Top-p sampling threshold.\n sampling_top_p: float = 0.95\n\n # Number of steps to take during training.\n num_train_steps: int = 500_000\n # Number of steps to take during evaluation.\n num_eval_steps: int = 500\n # Number of steps to generate predictions.\n # -1 will use the whole eval dataset.\n num_predict_steps: int = 20\n # Base learning rate.\n learning_rate: float = 0.0016\n # Linear learning rate warmup.\n warmup_steps: int = 1000\n # Cross entropy loss label smoothing.\n label_smoothing: float = 0.0\n # Decay factor for AdamW style weight decay.\n weight_decay: float = 0.1\n # Maximum length cutoff for training examples.\n max_target_length: int = 128\n # Maximum length cutoff for eval examples.\n max_eval_target_length: int = 512\n\n # Gemma transformer name.\n # Possible values defined in transformer.TransformerConfig:\n # (gemma_2b, gemma_7b, gemma2_2b, gemma2_9b, gemma2_27b, gemma3_1b, gemma3_4b, ...)\n transformer_name: str | None = None\n # or alternatively define the model using the dict of parameters\n transformer_params: dict | None = dataclasses.field(\n default_factory=lambda: {\n \"num_layers\": 4,\n \"embed_dim\": 256,\n \"hidden_dim\": 256 * 4 // 2, # embed_dim * num_heads // 2\n \"num_heads\": 4,\n \"head_dim\": 128,\n \"num_kv_heads\": 1,\n \"use_post_attn_norm\": False,\n \"use_post_ffw_norm\": False,\n \"attention_types\": (1, 1, 1, 1), # global * num_layers\n \"final_logit_softcap\": None,\n }\n )\n\n # Whether to save model checkpoints.\n save_checkpoints: bool = True\n # Whether to restore from existing model checkpoints.\n restore_checkpoints: bool = True\n # Save a checkpoint every these number of steps.\n checkpoint_every_steps: int = 10_000\n # Frequency of eval during training, e.g. every 1_000 steps.\n eval_every_steps: int = 5_000\n # Use bfloat16 mixed precision training instead of float32.\n use_bfloat16: bool = True\n # Integer for PRNG random seed.\n seed: int = 0\n\n # Parallelism\n mesh_axes: tuple[str, ...] = ('data', 'fsdp', 'tensor')\n axis_rules: MeshRules = MeshRules(\n embed='fsdp',\n mlp='tensor',\n kv='tensor',\n vocab='tensor',\n )\n data_sharding: tuple[str, ...] = ('data', 'fsdp')\n\n # One axis for each parallelism type may hold a placeholder (-1)\n # value to auto-shard based on available slices and devices.\n # By default, product of the DCN axes should equal number of slices\n # and product of the ICI axes should equal number of devices per slice.\n # ICI (Inter-Chip Interconnection): A high-speed connection between\n # sets of TPU chips, which form the TPU network.\n # DCN (Data Center Network): A connection between the TPU networks;\n # not as fast as ICI.\n # ICI has around 100x the bandwidth of DCN, but it is not a general\n # purpose connection, which is why DCN is necessary for scaling to\n # extremely large ML models.\n dcn_data_parallelism: int = -1\n dcn_fsdp_parallelism: int = 1\n dcn_tensor_parallelism: int = 1\n ici_data_parallelism: int = 1\n ici_fsdp_parallelism: int = -1\n ici_tensor_parallelism: int = 1\n\n def replace(self, **kwargs):\n return dataclasses.replace(self, **kwargs)\n\n\ndef get_config() -> TrainConfig:\n \"\"\"Get the default hyperparameter configuration.\"\"\"\n config = Config()\n return TrainConfig(**dataclasses.asdict(config))\n" }, { "path": "examples/gemma/helpers.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ============================================================================\n\"\"\"Helper functions.\"\"\"\n\nfrom __future__ import annotations\n\nfrom collections.abc import Callable\nfrom typing import Any, TypeVar\n\nimport flax\nfrom flax import nnx\nfrom flax.typing import VariableDict # pylint: disable=g-importing-member,g-multiple-import\n\nM = TypeVar('M', bound='nnx.Module')\n\n\ndef _flatten_path(path: tuple[str | int, ...]) -> str:\n def f(item) -> str:\n if isinstance(item, str):\n return f'{item}'\n elif isinstance(item, int):\n return f'[{item}]'\n else:\n raise ValueError(f'Unexpected type {type(item)}')\n\n return '.'.join([f(item) for item in path]).replace('.[', '[')\n\n\ndef module_from_linen_variables(\n module_factory: Callable[[], M],\n variables: VariableDict,\n map_key_fn: None | (\n Callable[[tuple[str, ...]], tuple[str | int, ...]]\n ) = None,\n assign_val_fn: None | (\n Callable[\n [dict[tuple[str, ...], Any], tuple[str | int, ...], VariableDict],\n dict[tuple[str, ...], Any],\n ]\n ) = None,\n) -> M:\n \"\"\"Returns an `nnx.Module` initialized with the `variables` of a linen module.\n\n Args:\n module_factory: A no-args callable that returns an `nnx.Module`.\n variables: A dictionary of variables.\n map_key_fn: An optional function for mapping keys in the `variables`\n dictionary to keys in the `nnx.Module`'s state. If not provided it is\n assumed that after removing the collection name the keys in the\n `variables` dictionary are the same as the keys in the `nnx.Module`'s\n state.\n \"\"\"\n if map_key_fn is None:\n\n def map_key_fn(path: tuple[str, ...]) -> tuple[str | int, ...]:\n return path[1:] if 'params' in variables else path\n\n if assign_val_fn is None:\n\n def assign_val_fn(\n state: dict[tuple[str, ...], Any],\n mapped_path: tuple[str | int, ...],\n val: Any,\n ) -> dict[tuple[str, ...], Any]:\n state[mapped_path].set_value(val)\n return state\n\n mdl: M = nnx.eval_shape(module_factory)\n graph_def, state = nnx.split(mdl)\n state = dict(nnx.to_flat_state(state))\n for path, val in flax.traverse_util.flatten_dict(variables).items():\n mapped_path = map_key_fn(path)\n if mapped_path not in state:\n raise ValueError(\n f\"'{mdl.__class__.__name__}.{_flatten_path(mapped_path)}' doesn't \"\n f' exist (original path={path}).'\n )\n state = assign_val_fn(state, mapped_path, val)\n state = nnx.from_flat_state(state)\n\n return nnx.merge(graph_def, state)\n" }, { "path": "examples/gemma/helpers_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ============================================================================\n\"\"\"Tests for helpers.\"\"\"\n\nfrom __future__ import annotations\n\n\nfrom absl.testing import absltest\nfrom absl.testing import parameterized\nfrom flax import nnx\nimport flax.linen as nn\nimport helpers\nimport jax\nimport jax.numpy as jnp\nimport numpy as np\n\n\nclass ModuleFromLinenVariablesTest(parameterized.TestCase):\n\n @parameterized.parameters(\n dict(\n inputs_shape=(1, 2, 3, 4),\n num_features=10,\n use_bias=True,\n ),\n dict(\n inputs_shape=(10, 5),\n num_features=4,\n use_bias=False,\n ),\n )\n def test_same_structure(self, inputs_shape, num_features, use_bias):\n rng_key = jax.random.PRNGKey(0)\n rng_inputs, rng_params = jax.random.split(rng_key)\n inputs = jax.random.normal(rng_inputs, inputs_shape)\n\n linen_mdl = nn.Dense(\n features=num_features,\n use_bias=use_bias,\n )\n linen_init_vars = linen_mdl.init(rng_params, jnp.zeros(inputs_shape))\n linen_output = linen_mdl.apply(\n linen_init_vars,\n inputs,\n )\n\n mdl = helpers.module_from_linen_variables(\n module_factory=lambda: nnx.Linear(\n in_features=inputs_shape[-1],\n out_features=num_features,\n use_bias=use_bias,\n rngs=nnx.Rngs(params=rng_params),\n ),\n variables=linen_init_vars,\n )\n output = mdl(inputs)\n\n np.testing.assert_array_equal(output, linen_output)\n\n @parameterized.parameters(\n dict(\n inputs_shape=(1, 2, 3, 4),\n num_features=(10, 20, 7),\n use_bias=(False, True, False),\n ),\n )\n def test_different_structure(self, inputs_shape, num_features, use_bias):\n rng_key = jax.random.PRNGKey(0)\n rng_inputs, rng_params = jax.random.split(rng_key)\n inputs = jax.random.normal(rng_inputs, inputs_shape)\n\n linen_mdl = nn.Sequential([\n nn.Sequential([\n nn.BatchNorm(use_running_average=False),\n nn.Dense(features=f, use_bias=b),\n ])\n for f, b in zip(num_features, use_bias)\n ])\n linen_init_vars = linen_mdl.init(rng_key, jnp.zeros(inputs_shape))\n linen_output, linen_vars = linen_mdl.apply(\n linen_init_vars,\n inputs,\n mutable=['batch_stats'],\n )\n\n module_factory = lambda: nnx.Sequential(*[\n nnx.Sequential(\n nnx.BatchNorm(\n num_features=in_f,\n use_running_average=False,\n rngs=nnx.Rngs(params=rng_params),\n ),\n nnx.Linear(\n in_features=in_f,\n out_features=out_f,\n use_bias=b,\n rngs=nnx.Rngs(params=rng_params),\n ),\n )\n for in_f, out_f, b in zip(in_features, out_features, use_bias)\n ])\n\n def _map_key_fn(key: tuple[str, ...]) -> tuple[str | int, ...]:\n new_key = []\n for k in key[1:]:\n if k.startswith('layers_'):\n prefix, suffix = k.split('layers_')\n assert not prefix, prefix\n new_key.append('layers')\n new_key.append(int(suffix))\n else:\n new_key.append(k)\n\n return tuple(new_key)\n\n in_features = (inputs_shape[-1], *num_features[:-1])\n out_features = num_features\n mdl = helpers.module_from_linen_variables(\n module_factory=module_factory,\n variables=linen_init_vars,\n map_key_fn=_map_key_fn,\n )\n output = mdl(inputs)\n\n np.testing.assert_array_equal(output, linen_output)\n for i in range(len(num_features)):\n np.testing.assert_array_equal(\n mdl.layers[i].layers[0].mean[...],\n linen_vars['batch_stats'][f'layers_{i}']['layers_0']['mean'],\n )\n np.testing.assert_array_equal(\n mdl.layers[i].layers[0].var[...],\n linen_vars['batch_stats'][f'layers_{i}']['layers_0']['var'],\n )\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "examples/gemma/input_pipeline.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Input pipeline for a LM1B dataset.\"\"\"\n\nimport os\nfrom typing import Any\n\nimport tokenizer\nimport tensorflow as tf\nimport tensorflow_datasets as tfds\n\nAUTOTUNE = tf.data.experimental.AUTOTUNE\nFeatures = dict[str, tf.Tensor]\n\n\nclass NormalizeFeatureNamesOp:\n \"\"\"Normalizes feature names to 'inputs' and 'targets'.\"\"\"\n\n def __call__(self, features: Features) -> Features:\n features['inputs'] = features.pop('text')\n # Unnecessary step used for uniformizing with examples/wmt.\n features['targets'] = features['inputs']\n return features\n\n\ndef get_raw_dataset(dataset_name: str, split: str) -> tf.data.Dataset:\n \"\"\"Loads a raw text dataset and normalizes feature keys.\n\n Args:\n dataset_name: TFDS dataset name.\n split: Split to use. This must be the full split. We shard the split across\n multiple hosts and currently don't support sharding subsplits.\n\n Returns:\n Dataset with source and target language features mapped to 'inputs' and\n 'targets'.\n \"\"\"\n split = tfds.split_for_jax_process(split, drop_remainder=True)\n ds = tfds.load(dataset_name, split=split)\n ds = ds.map(NormalizeFeatureNamesOp(), num_parallel_calls=AUTOTUNE)\n return ds\n\n\ndef pack_dataset(\n dataset: tf.data.Dataset,\n key2length: int | dict[str, int],\n keys: list[str] | None = None,\n) -> tf.data.Dataset:\n \"\"\"Creates a 'packed' version of a dataset on-the-fly.\n\n Adapted from the mesh-tf implementation.\n\n This is meant to replace the irritation of having to create a separate\n \"packed\" version of a dataset to train efficiently on TPU.\n Each example in the output dataset represents several examples in the\n input dataset.\n For each key in the input dataset, two additional keys are created:\n _segmentation: an int32 tensor identifying the parts\n representing the original example.\n _position: an int32 tensor identifying the position within the original\n example.\n Example:\n Two input examples get combined to form an output example.\n The input examples are:\n {\"inputs\": [8, 7, 1, 0], \"targets\":[4, 1, 0]}\n {\"inputs\": [2, 3, 4, 1], \"targets\":[5, 6, 1]}\n The output example is:\n {\n \"inputs\": [8, 7, 1, 2, 3, 4, 1, 0, 0, 0]\n \"inputs_segmentation\": [1, 1, 1, 2, 2, 2, 2, 0, 0, 0]\n \"inputs_position\": [0, 1, 2, 0, 1, 2, 3, 0, 0, 0]\n \"targets\": [4, 1, 5, 6, 1, 0, 0, 0, 0, 0]\n \"targets_segmentation\": [1, 1, 2, 2, 2, 0, 0, 0, 0, 0]\n \"targets_position\": [0, 1, 0, 1, 2, 0, 0, 0, 0, 0]\n }\n 0 represents padding in both the inputs and the outputs.\n Sequences in the incoming examples are truncated to length \"length\", and the\n sequences in the output examples all have fixed (padded) length \"length\".\n\n Args:\n dataset: a tf.data.Dataset\n key2length: an integer, or a dict from feature-key to integer\n keys: a list of strings (e.g. [\"inputs\", \"targets\"])\n\n Returns:\n a tf.data.Dataset\n \"\"\"\n shapes = tf.nest.map_structure(lambda spec: spec.shape, dataset.element_spec)\n if keys is None:\n keys = list(shapes.keys())\n for k in keys:\n if k not in shapes:\n raise ValueError(\n 'Key %s not found in dataset. Available keys are %s'\n % (k, shapes.keys())\n )\n if not shapes[k].is_compatible_with(tf.TensorShape([None])): # type: ignore[wrong-arg-types]\n raise ValueError('Tensors to be packed must be one-dimensional.')\n # make sure that the length dictionary contains all keys as well as the\n # keys suffixed by \"_segmentation\" and \"_position\"\n if isinstance(key2length, int):\n key2length = {k: key2length for k in keys}\n for k in keys:\n for suffix in ['_segmentation', '_position']:\n key2length[k + suffix] = key2length[k]\n\n # trim to length\n dataset = dataset.map(\n lambda x: {k: x[k][: key2length[k]] for k in keys},\n num_parallel_calls=AUTOTUNE,\n )\n # Setting batch_size=length ensures that the concatenated sequences (if they\n # have length >=1) are sufficient to fill at least one packed example.\n batch_size = max(key2length.values())\n dataset = dataset.padded_batch(\n batch_size, padded_shapes={k: [-1] for k in keys}\n )\n dataset = _pack_with_tf_ops(dataset, keys, key2length)\n\n # Set the Tensor shapes correctly since they get lost in the process.\n def my_fn(x):\n return {k: tf.reshape(v, [key2length[k]]) for k, v in x.items()}\n\n return dataset.map(my_fn, num_parallel_calls=AUTOTUNE)\n\n\ndef _pack_with_tf_ops(\n dataset: tf.data.Dataset, keys: list[str], key2length: dict[str, int]\n) -> tf.data.Dataset:\n \"\"\"Helper-function for packing a dataset which has already been batched.\n\n Helper for pack_dataset() Uses tf.while_loop.\n\n Args:\n dataset: a dataset containing padded batches of examples.\n keys: a list of strings\n key2length: a dict from feature-key to integer\n\n Returns:\n a dataset.\n \"\"\"\n empty_example = {}\n for k in keys:\n empty_example[k] = tf.zeros([0], dtype=tf.int32)\n empty_example[k + '_position'] = tf.zeros([0], dtype=tf.int32)\n keys_etc = empty_example.keys()\n\n def write_packed_example(partial, outputs):\n new_partial = empty_example.copy()\n new_outputs = {}\n for k in keys_etc:\n new_outputs[k] = outputs[k].write(\n outputs[k].size(),\n tf.pad(partial[k], [[0, key2length[k] - tf.size(partial[k])]]),\n )\n return new_partial, new_outputs\n\n def map_fn(x):\n \"\"\"Internal function to flat_map over.\n\n Consumes a batch of input examples and produces a variable number of output\n examples.\n Args:\n x: a single example\n\n Returns:\n a tf.data.Dataset\n \"\"\"\n partial = empty_example.copy()\n i = tf.zeros([], dtype=tf.int32)\n dynamic_batch_size = tf.shape(x[keys[0]])[0]\n outputs = {}\n for k in keys:\n outputs[k] = tf.TensorArray(\n tf.int32, size=0, dynamic_size=True, element_shape=[key2length[k]]\n )\n outputs[k + '_position'] = tf.TensorArray(\n tf.int32, size=0, dynamic_size=True, element_shape=[key2length[k]]\n )\n\n def body_fn(i, partial, outputs):\n \"\"\"Body function for while_loop.\n\n Args:\n i: integer scalar\n partial: dictionary of Tensor (partially-constructed example)\n outputs: dictionary of TensorArray\n\n Returns:\n A triple containing the new values of the inputs.\n \"\"\"\n can_append = True\n one_example = {}\n for k in keys:\n val = tf.cast(x[k][i], tf.int32)\n val = val[: tf.reduce_sum(tf.cast(tf.not_equal(val, 0), tf.int32))]\n one_example[k] = val\n for k in keys:\n can_append = tf.logical_and(\n can_append,\n tf.less_equal(\n tf.size(partial[k]) + tf.size(one_example[k]), key2length[k]\n ),\n )\n\n def false_fn():\n return write_packed_example(partial, outputs)\n\n def true_fn():\n return partial, outputs\n\n partial, outputs = tf.cond(can_append, true_fn, false_fn)\n new_partial = {}\n for k in keys:\n new_seq = one_example[k][: key2length[k]]\n new_seq_len = tf.size(new_seq)\n new_partial[k] = tf.concat([partial[k], new_seq], 0)\n new_partial[k + '_position'] = tf.concat(\n [partial[k + '_position'], tf.range(new_seq_len)], 0\n )\n partial = new_partial\n return i + 1, partial, outputs\n\n # For loop over all examples in the batch.\n _, partial, outputs = tf.while_loop(\n cond=lambda *_: True,\n body=body_fn,\n loop_vars=(i, partial, outputs),\n shape_invariants=(\n tf.TensorShape([]),\n {k: tf.TensorShape([None]) for k in keys_etc}, # type: ignore[wrong-arg-types]\n {k: tf.TensorShape(None) for k in keys_etc}, # type: ignore[wrong-arg-types]\n ),\n maximum_iterations=dynamic_batch_size,\n )\n _, outputs = write_packed_example(partial, outputs)\n packed = {k: outputs[k].stack() for k in keys_etc}\n for k in keys:\n packed[k + '_segmentation'] = tf.cumsum(\n tf.cast(tf.equal(packed[k + '_position'], 0), tf.int32), axis=1\n ) * tf.cast(tf.not_equal(packed[k], 0), tf.int32)\n return packed\n\n dataset = dataset.map(map_fn, num_parallel_calls=AUTOTUNE)\n return dataset.unbatch()\n\n\ndef shift_data_by_truncation(x):\n # https://github.com/AI-Hypercomputer/maxtext/blob/7fe1de75b3919c0fda00d23ad6cb29def9098362/MaxText/input_pipeline/_input_pipeline_utils.py#L53\n x['inputs'] = x['inputs'][:-1]\n x['targets'] = x['targets'][1:]\n return x\n\n\n# -----------------------------------------------------------------------------\n# Main dataset prep routines.\n# -----------------------------------------------------------------------------\ndef preprocess_data(\n dataset,\n shuffle: bool,\n num_epochs: int | None = 1,\n pack_examples: bool = True,\n shuffle_buffer_size: int = 1024,\n max_length: int = 512,\n batch_size: int = 256,\n drop_remainder: bool = True,\n prefetch_size: int = AUTOTUNE,\n shift: bool = True,\n):\n \"\"\"Shuffle and batch/pack the given dataset.\"\"\"\n\n def length_filter(max_len):\n def filter_fn(x):\n source, target = x['inputs'], x['targets']\n l = tf.maximum(tf.shape(source)[0], tf.shape(target)[0])\n return tf.less(l, max_len + 1)\n\n return filter_fn\n\n if max_length > 0:\n dataset = dataset.filter(length_filter(max_length))\n\n if shuffle:\n dataset = dataset.shuffle(shuffle_buffer_size)\n dataset = dataset.repeat(num_epochs)\n\n # Shift inputs for teacher-forced training\n if shift:\n dataset = dataset.map(\n shift_data_by_truncation,\n num_parallel_calls=AUTOTUNE,\n deterministic=True,\n )\n\n if pack_examples:\n dataset = pack_dataset(dataset, max_length)\n dataset = dataset.batch(batch_size, drop_remainder=drop_remainder)\n else: # simple (static-shape) padded batching\n dataset = dataset.padded_batch(\n batch_size,\n padded_shapes={'inputs': max_length, 'targets': max_length},\n padding_values={'inputs': 0, 'targets': 0},\n drop_remainder=drop_remainder,\n )\n\n if prefetch_size:\n dataset = dataset.prefetch(prefetch_size)\n\n return dataset\n\n\ndef get_datasets(\n config: Any,\n *,\n n_devices: int,\n vocab_path: str | None = None,\n):\n \"\"\"Load and return dataset of batched examples for use during training.\"\"\"\n if vocab_path is None:\n vocab_path = os.path.expanduser('~/lm1b_sentencepiece_model')\n\n train_data = get_raw_dataset(config.dataset_name, 'train')\n\n if config.eval_dataset_name:\n eval_dataset_name = config.eval_dataset_name\n else:\n eval_dataset_name = config.dataset_name\n eval_data = get_raw_dataset(eval_dataset_name, config.eval_split)\n\n # Tokenize data.\n sp_processor = tokenizer.load_or_train_tokenizer(\n train_data,\n vocab_path=vocab_path,\n vocab_size=config.vocab_size,\n max_corpus_chars=config.max_corpus_chars,\n )\n train_data = train_data.map(\n tokenizer.TokenizeOp(sp_processor), num_parallel_calls=AUTOTUNE\n )\n eval_data = eval_data.map(\n tokenizer.TokenizeOp(sp_processor), num_parallel_calls=AUTOTUNE\n )\n\n batch_size = config.per_device_batch_size * n_devices\n if config.eval_per_device_batch_size > 0:\n eval_batch_size = config.eval_per_device_batch_size * n_devices\n else:\n eval_batch_size = batch_size\n\n train_ds = preprocess_data(\n train_data,\n shuffle=True,\n num_epochs=None,\n pack_examples=True,\n batch_size=batch_size,\n max_length=config.max_target_length,\n )\n\n eval_ds = preprocess_data(\n eval_data,\n shuffle=False,\n pack_examples=False,\n batch_size=eval_batch_size,\n max_length=config.max_eval_target_length,\n )\n\n return train_ds, eval_ds, sp_processor\n" }, { "path": "examples/gemma/input_pipeline_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nimport pathlib\nimport sys\nimport tempfile\n\nfrom absl.testing import absltest\nimport tensorflow_datasets as tfds\n\nfrom configs import default\nimport input_pipeline\n\n# We just use different values here to verify that the input pipeline uses the\n# the correct value for the 3 different datasets.\n_TARGET_LENGTH = 32\n_EVAL_TARGET_LENGTH = 48\n\n\nclass InputPipelineTest(absltest.TestCase):\n\n def setUp(self):\n super().setUp()\n if sys.version_info >= (3, 13):\n self.skipTest('Test (and tensorflow-text) does not suport Python 3.13+')\n self.train_ds, self.eval_ds = self._get_datasets()\n\n def _get_datasets(self):\n config = default.get_config()\n config.per_device_batch_size = 1\n config.eval_per_device_batch_size = 2\n config.vocab_size = 32\n config.max_corpus_chars = 1000\n config.max_target_length = _TARGET_LENGTH\n config.max_eval_target_length = _EVAL_TARGET_LENGTH\n\n vocab_path = os.path.join(tempfile.mkdtemp(), 'sentencepiece_model')\n\n # Go two directories up to the root of the flax directory.\n # \"/path/to/flax/examples/gemma/input_pipeline_test.py\" -> \"/path/to/flax\"\n flax_root_dir = pathlib.Path(__file__).absolute().parents[2]\n data_dir = str(flax_root_dir) + '/.tfds/metadata' # pylint: disable=unused-variable\n\n with tfds.testing.mock_data(num_examples=128, data_dir=data_dir):\n train_ds, eval_ds, _ = input_pipeline.get_datasets(\n n_devices=2, config=config, vocab_path=vocab_path\n )\n return train_ds, eval_ds\n\n def test_train_ds(self):\n expected_shape = [2, _TARGET_LENGTH] # 2 devices.\n # For training we pack multiple short examples in one example.\n # *_position and *_segmentation indicate the boundaries.\n for batch in self.train_ds.take(3):\n self.assertEqual(\n {k: v.shape.as_list() for k, v in batch.items()},\n {\n 'inputs': expected_shape,\n 'inputs_position': expected_shape,\n 'inputs_segmentation': expected_shape,\n 'targets': expected_shape,\n 'targets_position': expected_shape,\n 'targets_segmentation': expected_shape,\n },\n )\n\n def test_eval_ds(self):\n expected_shape = [4, _EVAL_TARGET_LENGTH] # 2 devices.\n for batch in self.eval_ds.take(3):\n self.assertEqual(\n {k: v.shape.as_list() for k, v in batch.items()},\n {\n 'inputs': expected_shape,\n 'targets': expected_shape,\n },\n )\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "examples/gemma/layers.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ============================================================================\n\"\"\"Base layers.\"\"\"\n\nfrom __future__ import annotations\n\nfrom collections.abc import Sequence\nfrom typing import Any, Union\n\nfrom flax import nnx\nimport jax\nimport jax.numpy as jnp\nfrom jaxtyping import Array, ArrayLike # pylint: disable=g-importing-member,g-multiple-import\n\n\nShape = Sequence[Union[int, Any]]\n\n\nclass Einsum(nnx.Module):\n \"\"\"Einsum is a convenience module for parameterized tensor multiplication.\"\"\"\n\n def __init__(\n self,\n einsum_str: str,\n shape: Shape,\n *,\n kernel_init: nnx.Initializer = nnx.initializers.normal(),\n rngs: nnx.Rngs,\n dtype: Any = jnp.float32,\n ):\n self.einsum_str = einsum_str\n self.w = nnx.Param(kernel_init(rngs.params(), shape, dtype))\n\n def __call__(self, x: ArrayLike) -> Array:\n return jnp.einsum(self.einsum_str, x, self.w[...])\n\n @property\n def shape(self) -> Shape:\n return self.w.shape\n\n\nclass RMSNorm(nnx.Module):\n \"\"\"RMSNorm layer.\"\"\"\n\n def __init__(\n self,\n dim: int,\n *,\n scale_init: nnx.Initializer = nnx.initializers.zeros_init(),\n rngs: nnx.Rngs,\n dtype: Any = jnp.float32,\n ):\n self.scale = nnx.Param(scale_init(rngs.params(), dim, dtype))\n\n def __call__(self, x: Array) -> Array:\n dtype = self.scale.dtype\n var = jnp.mean(jnp.square(x), axis=-1, keepdims=True)\n normed_inputs = jnp.asarray(x * jax.lax.rsqrt(var + 1e-06), dtype=dtype)\n # normed_inputs is a rank-K tensor, K > 1 (K is typically 2 or 3). scale is\n # a rank-1 tensor. To avoid implicit rank-promotion, reshape scale to\n # a (1, ..., 1, D) tensor, so the rank of scale matches normed_inputs.\n scale = jnp.expand_dims(self.scale, axis=range(len(x.shape) - 1))\n normed_inputs = normed_inputs * (1 + scale)\n return normed_inputs\n" }, { "path": "examples/gemma/layers_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ============================================================================\n\"\"\"Tests for transformer layers.\"\"\"\n\nfrom absl.testing import absltest\nfrom absl.testing import parameterized\nfrom flax import nnx\nimport layers\nimport jax.numpy as jnp\nimport numpy as np\n\n\nclass EinsumTest(parameterized.TestCase):\n @parameterized.parameters(\n dict(\n inputs_shape=(1, 4),\n params_shape=(3, 2, 4, 3),\n eqn='TD,SNDH->STNH',\n expected_shape=(3, 1, 2, 3),\n ),\n dict(\n inputs_shape=(1, 2, 4),\n params_shape=(2, 4, 8),\n eqn='ANH,NHD->AD',\n expected_shape=(1, 8),\n ),\n )\n def test_einsum(self, inputs_shape, params_shape, eqn, expected_shape):\n einsum = layers.Einsum(eqn, params_shape, rngs=nnx.Rngs(params=0))\n output = einsum(\n jnp.ones(inputs_shape),\n )\n self.assertEqual(output.shape, expected_shape)\n\n @parameterized.parameters(\n dict(\n shape=(1, 4),\n ),\n dict(\n shape=(2, 5, 4, 7),\n ),\n )\n def test_shape(self, shape):\n einsum = layers.Einsum('ij->ji', shape, rngs=nnx.Rngs(params=0))\n self.assertEqual(einsum.shape, shape)\n\n\nclass RMSNormTest(parameterized.TestCase):\n @parameterized.parameters(dict(x=[0.1, 0.2], expected=[0.6324429, 1.2648858]))\n def test_rmsnorm(self, x, expected):\n x = jnp.array([x])\n rmsnorm = layers.RMSNorm(x.shape[-1], rngs=nnx.Rngs(params=0))\n output = rmsnorm(x)\n np.testing.assert_array_equal(output, jnp.array([expected]))\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "examples/gemma/main.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Main file for training Gemma model on the One Billion Word Benchmark dataset.\n\nThis file is intentionally kept short. The majority for logic is in libraries\nthat can be easily tested and imported in Colab.\n\"\"\"\n\nfrom absl import app\nfrom absl import flags\nfrom absl import logging\nfrom clu import platform\nimport train\nimport jax\nfrom ml_collections import config_flags\nimport tensorflow as tf\n\n\nFLAGS = flags.FLAGS\n\nflags.DEFINE_string('workdir', None, 'Directory to store model data.')\nconfig_flags.DEFINE_config_file(\n 'config',\n 'configs/default.py',\n 'File path to the training hyperparameter configuration.',\n lock_config=True,\n)\nflags.mark_flags_as_required(['workdir'])\n\n\ndef main(argv):\n if len(argv) > 1:\n raise app.UsageError('Too many command-line arguments.')\n\n # Hide any GPUs from TensorFlow. Otherwise TF might reserve memory and make\n # it unavailable to JAX.\n tf.config.experimental.set_visible_devices([], 'GPU')\n\n logging.info('JAX process: %d / %d', jax.process_index(), jax.process_count())\n logging.info('JAX local devices: %r', jax.local_devices())\n\n # Add a note so that we can tell which task is which JAX host.\n # (Depending on the platform task 0 is not guaranteed to be host 0)\n platform.work_unit().set_task_status(\n f'process_index: {jax.process_index()}, '\n f'process_count: {jax.process_count()}'\n )\n platform.work_unit().create_artifact(\n platform.ArtifactType.DIRECTORY, FLAGS.workdir, 'workdir'\n )\n\n train.train_and_evaluate(FLAGS.config, FLAGS.workdir)\n\n\nif __name__ == '__main__':\n jax.config.config_with_absl()\n app.run(main)\n" }, { "path": "examples/gemma/modules.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ============================================================================\n\"\"\"Transformer sub-modules.\"\"\"\n\nfrom __future__ import annotations\n\nfrom collections.abc import Sequence\nimport enum\nfrom typing import Any, Union\n\nfrom flax import nnx\nimport layers\nimport positional_embeddings\nimport sow_lib\nimport jax\nimport jax.numpy as jnp\nfrom jaxtyping import Array, ArrayLike # pylint: disable=g-importing-member,g-multiple-import\n\nLayerCache = dict[str, Array]\nShape = Sequence[Union[int, Any]]\n\nK_MASK = -2.3819763e38 # Set to a large negative number.\nDEFAULT_ROPE_BASE_FREQUENCY = 10_000\nDEFAULT_ROPE_SCALE_FACTOR = 1.0\n\n\nclass AttentionType(enum.Enum):\n GLOBAL = 1\n LOCAL_SLIDING = 2\n\n\nclass Embedder(nnx.Module):\n \"\"\"Embedder module.\"\"\"\n\n def __init__(\n self,\n vocab_size: int,\n embed_dim: int,\n *,\n embedding_init: nnx.Initializer = nnx.initializers.normal(),\n dtype: Any = jnp.float32,\n rngs: nnx.Rngs,\n ):\n self.input_embedding = nnx.Param(\n embedding_init(rngs.params(), (vocab_size, embed_dim), dtype)\n )\n\n def encode(self, x: ArrayLike) -> Array:\n x = self.input_embedding[(x,)]\n x *= jnp.sqrt(x.shape[-1]).astype(x.dtype)\n return x\n\n def decode(self, x: ArrayLike) -> Array:\n return jnp.dot(x, self.input_embedding.T)\n\n @property\n def embed_dim(self):\n return self.input_embedding.shape[1]\n\n @property\n def num_embed(self):\n return self.input_embedding.shape[0]\n\n\nclass Attention(nnx.Module):\n \"\"\"Attention module.\"\"\"\n\n def __init__(\n self,\n num_heads: int,\n num_kv_heads: int,\n features: int,\n head_dim: int,\n query_pre_attn_scalar: float,\n attn_type: AttentionType,\n *,\n rngs: nnx.Rngs,\n rope_base_frequency: int = DEFAULT_ROPE_BASE_FREQUENCY,\n rope_scale_factor: float = DEFAULT_ROPE_SCALE_FACTOR,\n attn_logits_soft_cap: float | None = None,\n sliding_window_size: int | None = None,\n use_qk_norm: bool = False,\n sow_config: sow_lib.SowConfig = sow_lib.SowConfig(),\n dtype: Any = jnp.float16,\n kernel_init: nnx.Initializer = nnx.initializers.normal(),\n scale_init: nnx.Initializer = nnx.initializers.zeros_init(),\n attn_vec_einsum_kernel_init: nnx.Initializer | None = None,\n qkv_einsum_kernel_init: nnx.Initializer | None = None,\n q_einsum_kernel_init: nnx.Initializer | None = None,\n kv_einsum_kernel_init: nnx.Initializer | None = None,\n ):\n if attn_type == AttentionType.LOCAL_SLIDING and sliding_window_size is None:\n raise ValueError(\n '`sliding_window_size` must be set if `attn_type` is Local Sliding.'\n )\n\n self.query_pre_attn_scalar = query_pre_attn_scalar\n self.attn_type = attn_type\n self.sliding_window_size = sliding_window_size\n self.attn_logits_soft_cap = attn_logits_soft_cap\n attn_vec_einsum_kernel_init = attn_vec_einsum_kernel_init if attn_vec_einsum_kernel_init else kernel_init\n self.attn_vec_einsum = layers.Einsum(\n einsum_str='BTNH,NHD->BTD',\n shape=(num_heads, head_dim, features),\n kernel_init=attn_vec_einsum_kernel_init,\n dtype=dtype,\n rngs=rngs,\n )\n self.rope_base_frequency = rope_base_frequency\n self.rope_scale_factor = rope_scale_factor\n self.use_qk_norm = use_qk_norm\n self.sow_config = sow_config\n\n if num_heads == num_kv_heads:\n qkv_einsum_kernel_init = qkv_einsum_kernel_init if qkv_einsum_kernel_init else kernel_init\n self.qkv_einsum = layers.Einsum(\n einsum_str='BTD,SNDH->SBTNH',\n shape=(3, num_heads, features, head_dim),\n kernel_init=qkv_einsum_kernel_init,\n dtype=dtype,\n rngs=rngs,\n )\n else:\n if num_heads % num_kv_heads != 0:\n raise ValueError(\n f\"Number of query heads ({num_heads}) must be divisible by \"\n f\"number of key/value heads ({num_kv_heads}).\"\n )\n\n q_einsum_kernel_init = q_einsum_kernel_init if q_einsum_kernel_init else kernel_init\n self.q_einsum = layers.Einsum(\n einsum_str='BTD,NDH->BTNH',\n shape=(num_heads, features, head_dim),\n kernel_init=q_einsum_kernel_init,\n dtype=dtype,\n rngs=rngs,\n )\n kv_einsum_kernel_init = kv_einsum_kernel_init if kv_einsum_kernel_init else kernel_init\n self.kv_einsum = layers.Einsum(\n einsum_str='BSD,CKDH->CBSKH',\n shape=(2, num_kv_heads, features, head_dim),\n kernel_init=kv_einsum_kernel_init,\n dtype=dtype,\n rngs=rngs,\n )\n\n if self.use_qk_norm:\n self._query_norm = layers.RMSNorm(\n head_dim,\n scale_init=scale_init,\n dtype=dtype,\n rngs=rngs,\n )\n self._key_norm = layers.RMSNorm(\n head_dim,\n scale_init=scale_init,\n dtype=dtype,\n rngs=rngs,\n )\n\n def __call__(\n self,\n x: Array,\n segment_pos: Array,\n cache: LayerCache | None,\n attn_mask: Array,\n ) -> tuple[LayerCache | None, Array]:\n seq_len = x.shape[1]\n\n if self.use_qkv_einsum:\n query_proj, key_proj, value_proj = self.qkv_einsum(x)\n else:\n query_proj = self.q_einsum(x)\n key_proj, value_proj = self.kv_einsum(x)\n\n if self.use_qk_norm:\n query_proj = self._query_norm(query_proj)\n key_proj = self._key_norm(key_proj)\n\n query_proj = positional_embeddings.apply_rope(\n query_proj,\n segment_pos,\n head_dim=self.head_dim,\n max_wavelength=self.rope_base_frequency,\n scale_factor=self.rope_scale_factor,\n )\n query_scaled = query_proj * self.query_pre_attn_scalar\n key_proj = positional_embeddings.apply_rope(\n key_proj,\n segment_pos,\n head_dim=self.head_dim,\n max_wavelength=self.rope_base_frequency,\n scale_factor=self.rope_scale_factor,\n )\n\n # Cache is left aligned.\n if cache is not None:\n end_index = cache['end_index'][0]\n slice_indices = (0, end_index % cache['v'].shape[1], 0, 0)\n value_proj = jax.lax.dynamic_update_slice(\n cache['v'],\n value_proj,\n slice_indices,\n )\n key_proj = jax.lax.dynamic_update_slice(\n cache['k'], key_proj, slice_indices\n )\n\n use_gqa = self.num_heads > self.num_kv_heads and self.num_kv_heads > 1\n if use_gqa:\n # Reshape matrices to enable einsums over groups.\n num_groups = self.num_heads // self.num_kv_heads\n batch_size, seq_size, _, head_dim = query_scaled.shape\n query_scaled = query_scaled.reshape(\n (batch_size, seq_size, self.num_kv_heads, num_groups, head_dim)\n )\n logits = jnp.einsum('BTKGH,BSKH->BTKGS', query_scaled, key_proj)\n logits = logits.reshape(\n (batch_size, seq_size, self.num_heads, -1)\n )\n else:\n logits = jnp.einsum('BTNH,BSNH->BTNS', query_scaled, key_proj)\n\n if self.attn_logits_soft_cap is not None:\n logits = jnp.tanh(logits / self.attn_logits_soft_cap)\n logits = logits * self.attn_logits_soft_cap\n if self.attn_type == AttentionType.LOCAL_SLIDING:\n if self.sliding_window_size is None:\n raise ValueError(\n 'sliding_window_size must be set if attn_type is Local Sliding.'\n )\n\n all_ones = jnp.ones_like(attn_mask)\n sliding_mask = jnp.triu(\n all_ones, -1 * self.sliding_window_size + 1\n ) * jnp.tril(all_ones, self.sliding_window_size - 1)\n attn_mask = sliding_mask * attn_mask\n\n padded_logits = jnp.where((jnp.expand_dims(attn_mask, -2)), logits, K_MASK)\n self.sow_config.maybe_sow_attn_logits_topk(padded_logits, self)\n probs = jax.nn.softmax(padded_logits, axis=-1).astype(key_proj.dtype)\n\n if use_gqa:\n # Reshape matrices to enable einsums over groups.\n num_groups = self.num_heads // self.num_kv_heads\n batch_size, seq_size1, _, _ = probs.shape\n probs = probs.reshape(\n (batch_size, seq_size1, self.num_kv_heads, num_groups, -1)\n )\n encoded = jnp.einsum('BTKGS,BSKH->BTKGH', probs, value_proj)\n encoded = encoded.reshape(\n (batch_size, seq_size, self.num_heads, head_dim)\n )\n else:\n encoded = jnp.einsum('BTNS,BSNH->BTNH', probs, value_proj)\n attn_output = self.attn_vec_einsum(encoded)\n\n if cache is not None:\n new_cache = {\n 'v': value_proj,\n 'k': key_proj,\n 'end_index': cache['end_index'] + seq_len,\n }\n else:\n new_cache = None\n\n return new_cache, attn_output\n\n @property\n def head_dim(self):\n return self.attn_vec_einsum.shape[1]\n\n @property\n def num_heads(self):\n return (\n self.qkv_einsum.shape[1]\n if self.use_qkv_einsum\n else self.q_einsum.shape[0]\n )\n\n @property\n def num_kv_heads(self):\n return (\n self.qkv_einsum.shape[1]\n if self.use_qkv_einsum\n else self.kv_einsum.shape[1]\n )\n\n @property\n def use_qkv_einsum(self):\n return hasattr(self, 'qkv_einsum') and self.qkv_einsum is not None\n\n def init_cache(\n self,\n cache_size: int,\n batch_size: int,\n dtype: jnp.dtype = jnp.bfloat16,\n ) -> LayerCache:\n return {\n 'v': jnp.zeros(\n (batch_size, cache_size, self.num_kv_heads, self.head_dim),\n dtype=dtype,\n ),\n 'k': jnp.zeros(\n (batch_size, cache_size, self.num_kv_heads, self.head_dim),\n dtype=dtype,\n ),\n 'end_index': jnp.zeros((batch_size,), dtype=jnp.int32),\n }\n\n\nclass FeedForward(nnx.Module):\n \"\"\"Feed forward module.\"\"\"\n\n def __init__(\n self,\n features: int,\n hidden_dim: int,\n *,\n kernel_init: nnx.Initializer = nnx.initializers.normal(),\n rngs: nnx.Rngs,\n sow_config: sow_lib.SowConfig = sow_lib.SowConfig(),\n dtype: Any = jnp.float32,\n ):\n self.gate_proj = nnx.Linear(\n in_features=features,\n out_features=hidden_dim,\n use_bias=False,\n rngs=rngs,\n kernel_init=kernel_init,\n dtype=dtype,\n )\n self.up_proj = nnx.Linear(\n in_features=features,\n out_features=hidden_dim,\n use_bias=False,\n rngs=rngs,\n kernel_init=kernel_init,\n dtype=dtype,\n )\n self.down_proj = nnx.Linear(\n in_features=hidden_dim,\n out_features=features,\n use_bias=False,\n rngs=rngs,\n kernel_init=kernel_init,\n dtype=dtype,\n )\n self.sow_config = sow_config\n\n def __call__(self, x: ArrayLike) -> Array:\n ff_gate = self.gate_proj(x)\n gate_value = nnx.gelu(ff_gate)\n\n ff1 = self.up_proj(x)\n activations = gate_value * ff1\n self.sow_config.maybe_sow_mlp_hidden_topk(activations, self)\n\n outputs = self.down_proj(activations)\n return outputs\n\n\nclass Block(nnx.Module):\n \"\"\"Transformer block.\"\"\"\n\n def __init__(\n self,\n config, # TransformerConfig\n attn_type: AttentionType,\n *,\n rngs: nnx.Rngs,\n sow_config: sow_lib.SowConfig = sow_lib.SowConfig(),\n ):\n num_heads = config.num_heads\n num_kv_heads = config.num_kv_heads\n embed_dim = config.embed_dim\n head_dim = config.head_dim\n hidden_dim = config.hidden_dim\n sliding_window_size = config.sliding_window_size\n use_post_attn_norm = config.use_post_attn_norm\n use_post_ffw_norm = config.use_post_ffw_norm\n query_pre_attn_scalar = config.query_pre_attn_scalar()\n if attn_type == AttentionType.LOCAL_SLIDING:\n rope_base_frequency = config.local_base_frequency\n rope_scale_factor = config.local_scale_factor\n else:\n rope_base_frequency = config.global_base_frequency\n rope_scale_factor = config.global_scale_factor\n\n attn_logits_soft_cap = config.attn_logits_soft_cap\n use_qk_norm = config.use_qk_norm\n dtype = config.dtype\n\n self.pre_attention_norm = layers.RMSNorm(\n embed_dim,\n scale_init=maybe_with_partitioning(\n nnx.initializers.zeros_init(),\n config.axis_rules,\n (\"embed\", ),\n ),\n rngs=rngs,\n dtype=dtype,\n )\n self.attn = Attention(\n num_heads=num_heads,\n num_kv_heads=num_kv_heads,\n features=embed_dim,\n head_dim=head_dim,\n query_pre_attn_scalar=query_pre_attn_scalar,\n attn_type=attn_type,\n rope_base_frequency=rope_base_frequency,\n rope_scale_factor=rope_scale_factor,\n attn_logits_soft_cap=attn_logits_soft_cap,\n sliding_window_size=sliding_window_size,\n rngs=rngs,\n use_qk_norm=use_qk_norm,\n sow_config=sow_config,\n attn_vec_einsum_kernel_init=maybe_with_partitioning(\n nnx.initializers.normal(),\n config.axis_rules,\n (None, \"embed\", \"kv\"), # sharded array shape: (num_heads, head_dim, features)\n ),\n qkv_einsum_kernel_init=maybe_with_partitioning(\n nnx.initializers.normal(),\n config.axis_rules,\n (None, None, \"embed\", \"kv\"), # sharded array shape: (3, num_heads, features, head_dim)\n ),\n q_einsum_kernel_init=maybe_with_partitioning(\n nnx.initializers.normal(),\n config.axis_rules,\n (None, \"embed\", \"kv\"), # sharded array shape: (num_heads, features, head_dim)\n ),\n kv_einsum_kernel_init=maybe_with_partitioning(\n nnx.initializers.normal(),\n config.axis_rules,\n (None, None, \"embed\", \"kv\"), # sharded array shape: (2, num_kv_heads, features, head_dim)\n ),\n scale_init=maybe_with_partitioning(\n nnx.initializers.zeros_init(),\n config.axis_rules,\n (\"embed\", ),\n ),\n dtype=dtype,\n )\n if use_post_attn_norm:\n self.post_attention_norm = layers.RMSNorm(\n embed_dim,\n scale_init=maybe_with_partitioning(\n nnx.initializers.zeros_init(),\n config.axis_rules,\n (\"embed\", ),\n ),\n rngs=rngs,\n dtype=dtype,\n )\n else:\n self.post_attention_norm = None\n\n self.pre_ffw_norm = layers.RMSNorm(\n embed_dim,\n scale_init=maybe_with_partitioning(\n nnx.initializers.zeros_init(),\n config.axis_rules,\n (\"embed\", ),\n ),\n rngs=rngs,\n dtype=dtype,\n )\n self.mlp = FeedForward(\n features=embed_dim,\n hidden_dim=hidden_dim,\n kernel_init=maybe_with_partitioning(\n nnx.initializers.normal(),\n config.axis_rules,\n (\"embed\", \"mlp\"),\n ),\n rngs=rngs,\n sow_config=sow_config,\n )\n if use_post_ffw_norm:\n self.post_ffw_norm = layers.RMSNorm(\n embed_dim,\n scale_init=maybe_with_partitioning(\n nnx.initializers.zeros_init(),\n config.axis_rules,\n (\"embed\", ),\n ),\n rngs=rngs,\n dtype=dtype,\n )\n else:\n self.post_ffw_norm = None\n self.sow_config = sow_config\n\n def __call__(\n self,\n x: jax.Array,\n segment_pos: jax.Array,\n cache: LayerCache | None,\n attn_mask: jax.Array,\n ) -> tuple[LayerCache | None, jax.Array]:\n\n # Attention.\n attn_inputs = self.pre_attention_norm(x)\n cache, attn_output = self.attn(\n attn_inputs,\n segment_pos,\n cache,\n attn_mask,\n )\n if self.post_attention_norm is not None:\n attn_output = self.post_attention_norm(attn_output)\n x += attn_output\n self.sow_config.maybe_sow_rs_after_attention(x, self)\n\n # Feed forward.\n ffw_inputs = self.pre_ffw_norm(x)\n ffw_outputs = self.mlp(ffw_inputs)\n if self.post_ffw_norm is not None:\n ffw_outputs = self.post_ffw_norm(ffw_outputs)\n x += ffw_outputs\n self.sow_config.maybe_sow_rs_after_ffw(x, self)\n\n return cache, x\n\n def init_cache(\n self,\n cache_size: int,\n batch_size: int,\n dtype: jnp.dtype = jnp.bfloat16,\n ) -> LayerCache:\n return self.attn.init_cache(\n cache_size=cache_size,\n batch_size=batch_size,\n dtype=dtype,\n )\n\n\ndef maybe_with_partitioning(fn, axis_rules, axis_rules_args=()):\n if axis_rules is None:\n return fn\n return nnx.with_partitioning(fn, axis_rules(*axis_rules_args))\n" }, { "path": "examples/gemma/modules_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ============================================================================\n\"\"\"Tests for transformer modules.\"\"\"\n\nfrom absl.testing import absltest\nfrom absl.testing import parameterized\nfrom flax import nnx\nimport modules\nimport transformer as transformer_lib\nimport jax\nimport jax.numpy as jnp\nimport numpy as np\n\n\nclass EmbedderTest(parameterized.TestCase):\n\n @parameterized.parameters(\n dict(\n vocab_size=10,\n embed_dim=4,\n inputs=[2, 3],\n expected=[[2.0, 2.0, 2.0, 2.0], [2.0, 2.0, 2.0, 2.0]],\n ),\n )\n def test_encode(self, vocab_size, embed_dim, inputs, expected):\n embedder = modules.Embedder(\n vocab_size=vocab_size,\n embed_dim=embed_dim,\n rngs=nnx.Rngs(params=0),\n )\n embedder.input_embedding[...] = jnp.ones((vocab_size, embed_dim))\n output = embedder.encode(inputs)\n np.testing.assert_array_equal(output, jnp.array(expected))\n\n @parameterized.parameters(\n dict(\n vocab_size=5,\n embed_dim=2,\n inputs=[[1, 2]],\n expected=[[3.0, 3.0, 3.0, 3.0, 3.0]],\n ),\n )\n def test_decode(self, vocab_size, embed_dim, inputs, expected):\n embedder = modules.Embedder(\n vocab_size=vocab_size,\n embed_dim=embed_dim,\n rngs=nnx.Rngs(params=0),\n )\n embedder.input_embedding[...] = jnp.ones((vocab_size, embed_dim))\n output = embedder.decode(jnp.array(inputs))\n np.testing.assert_array_equal(output, jnp.array(expected))\n\n\nclass AttentionTest(parameterized.TestCase):\n\n @parameterized.parameters(\n dict(\n head_dim=2,\n ),\n dict(\n head_dim=20,\n ),\n )\n def test_head_dim(self, head_dim):\n attn = modules.Attention(\n num_heads=4,\n num_kv_heads=2,\n features=5,\n head_dim=head_dim,\n query_pre_attn_scalar=1.0,\n attn_type=modules.AttentionType.GLOBAL,\n rngs=nnx.Rngs(params=0),\n )\n\n self.assertEqual(attn.head_dim, head_dim)\n\n @parameterized.parameters(\n dict(\n num_heads=4,\n num_kv_heads=2,\n expected_use_qkv_einsum=False,\n ),\n dict(\n num_heads=3,\n num_kv_heads=3,\n expected_use_qkv_einsum=True,\n ),\n )\n def test_use_qkv_einsum(\n self,\n num_heads,\n num_kv_heads,\n expected_use_qkv_einsum,\n ):\n attn = modules.Attention(\n num_heads=num_heads,\n num_kv_heads=num_kv_heads,\n features=5,\n head_dim=8,\n query_pre_attn_scalar=1.0,\n attn_type=modules.AttentionType.GLOBAL,\n rngs=nnx.Rngs(params=0),\n )\n\n self.assertEqual(attn.use_qkv_einsum, expected_use_qkv_einsum)\n self.assertEqual(hasattr(attn, 'q_einsum'), not expected_use_qkv_einsum)\n self.assertEqual(hasattr(attn, 'kv_einsum'), not expected_use_qkv_einsum)\n\n @parameterized.parameters(\n dict(\n num_heads=2,\n head_dim=4,\n features=8,\n segment_pos=0,\n cache_size=3,\n batch_size=2,\n expected_cache_shape=(2, 3, 2, 4),\n expected_output_shape=(2, 1, 8),\n ),\n )\n def test_attention(\n self,\n num_heads,\n head_dim,\n features,\n segment_pos,\n cache_size,\n batch_size,\n expected_cache_shape,\n expected_output_shape,\n ):\n attn_mask = jnp.ones((batch_size, 1, cache_size))\n attn = modules.Attention(\n num_heads,\n num_heads,\n features,\n head_dim,\n query_pre_attn_scalar=1.0,\n attn_type=modules.AttentionType.GLOBAL,\n rngs=nnx.Rngs(params=0),\n )\n cache = attn.init_cache(\n cache_size=cache_size,\n batch_size=batch_size,\n dtype=jnp.float32,\n )\n x = jnp.ones((batch_size, 1, features))\n cache, output = attn(x, jnp.array([[segment_pos]]), cache, attn_mask)\n\n self.assertEqual(cache['k'].shape, expected_cache_shape)\n self.assertEqual(output.shape, expected_output_shape)\n\n @parameterized.parameters(\n dict(\n sliding_window_size=2,\n ),\n )\n def test_sliding_window(self, sliding_window_size):\n num_heads = 2\n head_dim = 4\n features = 8\n segment_pos = 0\n cache_size = 3\n batch_size = 2\n attn_mask = jnp.ones((batch_size, 1, cache_size))\n x = jnp.ones((batch_size, 1, features))\n attn = modules.Attention(\n num_heads,\n num_heads,\n features,\n head_dim,\n query_pre_attn_scalar=1.0,\n attn_type=modules.AttentionType.GLOBAL,\n rngs=nnx.Rngs(params=0),\n )\n cache = attn.init_cache(\n cache_size=cache_size,\n batch_size=batch_size,\n dtype=jnp.float32,\n )\n _, output = attn(x, jnp.array([[segment_pos]]), cache, attn_mask)\n sliding_attn = modules.Attention(\n num_heads,\n num_heads,\n features,\n head_dim,\n query_pre_attn_scalar=1.0,\n attn_type=modules.AttentionType.LOCAL_SLIDING,\n sliding_window_size=sliding_window_size,\n rngs=nnx.Rngs(params=0),\n )\n _, sliding_output = sliding_attn(\n x, jnp.array([[segment_pos]]), cache, attn_mask\n )\n\n self.assertFalse((output == sliding_output).all())\n\n\nclass FeedForwardTest(parameterized.TestCase):\n\n @parameterized.parameters(\n dict(\n features=2,\n hidden_dim=3,\n batch_size=2,\n expected_val=[11.72758674, 47.99916],\n expected_shape=(2, 1, 2),\n ),\n )\n def test_ffw(\n self, features, hidden_dim, batch_size, expected_val, expected_shape\n ):\n inputs = jnp.arange(1, batch_size+1)[:, None, None]\n inputs = jnp.repeat(inputs, features, axis=-1)\n ffw = modules.FeedForward(\n features=features,\n hidden_dim=hidden_dim,\n rngs=nnx.Rngs(params=0),\n )\n ffw.gate_proj.kernel[...] = jnp.ones((features, hidden_dim))\n ffw.up_proj.kernel[...] = jnp.ones((features, hidden_dim))\n ffw.down_proj.kernel[...] = jnp.ones((hidden_dim, features))\n\n with jax.default_matmul_precision('float32'):\n outputs = ffw(inputs)\n\n np.testing.assert_array_almost_equal(outputs[:, 0, 0], expected_val)\n self.assertEqual(outputs.shape, expected_shape)\n\n\nclass BlockTest(parameterized.TestCase):\n\n @parameterized.parameters(\n dict(\n num_heads=2,\n embed_dim=4,\n head_dim=6,\n cache_size=3,\n batch_size=2,\n use_post_attn_norm=False,\n use_post_ffw_norm=False,\n expected_cache_shape=(2, 3, 2, 6),\n expected_output_shape=(2, 1, 4),\n ),\n )\n def test_block(\n self,\n num_heads,\n embed_dim,\n head_dim,\n cache_size,\n batch_size,\n use_post_attn_norm,\n use_post_ffw_norm,\n expected_cache_shape,\n expected_output_shape,\n ):\n inputs = jnp.ones((batch_size, 1, embed_dim))\n attn_mask = jnp.ones((batch_size, 1, cache_size))\n\n config = transformer_lib.TransformerConfig(\n num_heads=num_heads,\n num_kv_heads=num_heads,\n embed_dim=embed_dim,\n head_dim=head_dim,\n hidden_dim=1,\n use_post_attn_norm=use_post_attn_norm,\n use_post_ffw_norm=use_post_ffw_norm,\n final_logit_softcap=None,\n num_layers=-1,\n num_embed=-1,\n attention_types=[],\n )\n\n block = modules.Block(\n config,\n attn_type=modules.AttentionType.GLOBAL,\n rngs=nnx.Rngs(params=0),\n )\n cache = block.init_cache(\n cache_size=cache_size,\n batch_size=batch_size,\n dtype=jnp.float32,\n )\n\n new_cache, outputs = block(inputs, jnp.array([[0]]), cache, attn_mask)\n\n self.assertEqual(block.post_attention_norm is not None, use_post_attn_norm)\n self.assertEqual(new_cache['k'].shape, expected_cache_shape)\n self.assertEqual(outputs.shape, expected_output_shape)\n\n @parameterized.parameters(\n dict(\n num_heads=1,\n embed_dim=1,\n head_dim=2,\n cache_size=1,\n batch_size=1,\n ),\n )\n def test_post_attention_norm(\n self,\n num_heads,\n embed_dim,\n head_dim,\n cache_size,\n batch_size,\n ):\n inputs = jnp.ones((batch_size, 1, embed_dim))\n attn_mask = jnp.ones((batch_size, 1, cache_size))\n\n normed_block_config = transformer_lib.TransformerConfig(\n num_heads=num_heads,\n num_kv_heads=num_heads,\n embed_dim=embed_dim,\n head_dim=head_dim,\n hidden_dim=1,\n use_post_attn_norm=True,\n use_post_ffw_norm=False,\n final_logit_softcap=None,\n num_layers=-1,\n num_embed=-1,\n attention_types=[],\n )\n normed_block = modules.Block(\n normed_block_config,\n attn_type=modules.AttentionType.GLOBAL,\n rngs=nnx.Rngs(params=0),\n )\n\n unnormed_block_config = transformer_lib.TransformerConfig(\n num_heads=num_heads,\n num_kv_heads=num_heads,\n embed_dim=embed_dim,\n head_dim=head_dim,\n hidden_dim=1,\n use_post_attn_norm=False,\n use_post_ffw_norm=False,\n final_logit_softcap=None,\n num_layers=-1,\n num_embed=-1,\n attention_types=[],\n )\n unnormed_block = modules.Block(\n unnormed_block_config,\n attn_type=modules.AttentionType.GLOBAL,\n rngs=nnx.Rngs(params=0),\n )\n\n # Ok to use the same cache for both blocks.\n cache = normed_block.init_cache(\n cache_size=cache_size,\n batch_size=batch_size,\n dtype=jnp.float32,\n )\n\n all_outputs = []\n for block in (normed_block, unnormed_block):\n _, outputs = block(inputs, jnp.array([[0]]), cache, attn_mask)\n all_outputs.append(outputs)\n\n normed_output, unnormed_output = all_outputs # pylint: disable=unbalanced-tuple-unpacking\n self.assertTrue(jnp.not_equal(normed_output, unnormed_output).all())\n\n @parameterized.parameters(\n dict(\n num_heads=1,\n embed_dim=1,\n head_dim=2,\n cache_size=1,\n batch_size=1,\n ),\n )\n def test_post_ffw_norm(\n self,\n num_heads,\n embed_dim,\n head_dim,\n cache_size,\n batch_size,\n ):\n inputs = jnp.ones((batch_size, 1, embed_dim))\n attn_mask = jnp.ones((batch_size, 1, cache_size))\n\n normed_block_config = transformer_lib.TransformerConfig(\n num_heads=num_heads,\n num_kv_heads=num_heads,\n embed_dim=embed_dim,\n head_dim=head_dim,\n hidden_dim=1,\n use_post_attn_norm=False,\n use_post_ffw_norm=True,\n final_logit_softcap=None,\n num_layers=-1,\n num_embed=-1,\n attention_types=[],\n )\n normed_block = modules.Block(\n normed_block_config,\n attn_type=modules.AttentionType.GLOBAL,\n rngs=nnx.Rngs(params=0),\n )\n\n unnormed_block_config = transformer_lib.TransformerConfig(\n num_heads=num_heads,\n num_kv_heads=num_heads,\n embed_dim=embed_dim,\n head_dim=head_dim,\n hidden_dim=1,\n use_post_attn_norm=False,\n use_post_ffw_norm=False,\n final_logit_softcap=None,\n num_layers=-1,\n num_embed=-1,\n attention_types=[],\n )\n unnormed_block = modules.Block(\n unnormed_block_config,\n attn_type=modules.AttentionType.GLOBAL,\n rngs=nnx.Rngs(params=0),\n )\n\n # Ok to use the same cache for both blocks.\n cache = normed_block.init_cache(\n cache_size=cache_size,\n batch_size=batch_size,\n dtype=jnp.float32,\n )\n\n all_outputs = []\n for block in (normed_block, unnormed_block):\n _, outputs = block(inputs, jnp.array([[0]]), cache, attn_mask)\n all_outputs.append(outputs)\n\n normed_output, unnormed_output = all_outputs # pylint: disable=unbalanced-tuple-unpacking\n print(normed_output.shape, unnormed_output.shape)\n print(f\"{normed_output=}\")\n print(f\"{unnormed_output=}\")\n self.assertTrue(jnp.not_equal(normed_output, unnormed_output).all())\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "examples/gemma/params.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ============================================================================\n\"\"\"Utils for loading Gemma params.\"\"\"\n\nfrom collections.abc import Mapping\nimport functools\nfrom typing import Any\n\nimport jax\nimport jax.numpy as jnp\nimport orbax.checkpoint\n\nParams = Mapping[str, Any]\n\n\ndef load_and_format_params(path: str) -> Params:\n \"\"\"Loads parameters and formats them for compatibility.\"\"\"\n params = load_params(path)\n param_state = jax.tree.map(jnp.array, params)\n remapped_params = param_remapper(param_state)\n nested_params = nest_params(remapped_params)\n return nested_params\n\n\ndef load_metadata(path: str) -> Any | None:\n \"\"\"Loads metadata from a checkpoint path.\"\"\"\n checkpointer = orbax.checkpoint.PyTreeCheckpointer()\n metadata = checkpointer.metadata(path)\n return metadata\n\n\n@functools.cache\ndef load_params(path: str) -> Params:\n \"\"\"Loads parameters from a checkpoint path.\"\"\"\n checkpointer = orbax.checkpoint.PyTreeCheckpointer()\n params = checkpointer.restore(path)\n return params\n\n\ndef param_remapper(orig_params: Params) -> Params:\n \"\"\"Remaps params to new module layout.\n\n This is needed here because the model definition does not have a separate\n `mlp` module.\n\n Args:\n orig_params: original dict of parameters in Gemma format.\n\n Returns:\n dict of params with different names.\n \"\"\"\n new_params = {}\n for k, v in orig_params.items():\n if 'mlp/' in k:\n layer_name, param = k.rsplit('/', maxsplit=1)\n if layer_name not in new_params:\n new_params[layer_name] = {}\n if 'w' in v:\n new_params[layer_name][param] = v['w']\n else:\n new_params[k] = v\n return new_params\n\n\ndef nest_params(params: Params) -> Params:\n \"\"\"Nests params as a dict of dicts rather than a flat dict.\"\"\"\n nested_params = {}\n for path, param in params.items():\n *path, leaf = path.split('/')\n subdict = nested_params\n for key in path:\n subdict = subdict.setdefault(key, {})\n subdict[leaf] = param\n return nested_params\n" }, { "path": "examples/gemma/positional_embeddings.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ============================================================================\n\"\"\"Utils for positional embeddings (including RoPE).\"\"\"\n\nimport jax\nimport jax.numpy as jnp\n\n_MAX_WAVELENGTH = 10_000\n\n\ndef add_positional_embedding(\n input_embedding: jax.Array,\n position: int,\n max_wavelength: int = _MAX_WAVELENGTH,\n) -> jax.Array:\n \"\"\"Adds positional embeddings to input embeddings.\"\"\"\n embed_dim = input_embedding.shape[-1]\n num_timescales = embed_dim // 2\n log_timescale_increment = jnp.log(float(max_wavelength)) / jnp.maximum(\n jnp.asarray(num_timescales, dtype=jnp.float32) - 1, 1\n )\n inv_timescales = jnp.exp(\n jnp.arange(num_timescales, dtype=jnp.float32) * -log_timescale_increment\n )\n scaled_time = position * inv_timescales\n signal = jnp.concatenate([jnp.sin(scaled_time), jnp.cos(scaled_time)])\n signal = jnp.pad(signal, [[0, jnp.mod(embed_dim, 2)]])\n position_embedding = signal.astype(jnp.float32)\n\n return input_embedding + position_embedding\n\n\ndef apply_rope(\n inputs: jax.Array, # [B, L]\n positions: jax.Array, # [B, L]\n head_dim: int,\n max_wavelength: int = _MAX_WAVELENGTH,\n scale_factor: float = 1.0,\n) -> jax.Array:\n \"\"\"Applies RoPE.\"\"\"\n fraction = 2 * jnp.arange(0, head_dim // 2) / head_dim\n timescale = max_wavelength**fraction\n\n sinusoid_inp = (\n positions[..., jnp.newaxis] / timescale[jnp.newaxis, jnp.newaxis, :]\n )\n sinusoid_inp = sinusoid_inp[..., jnp.newaxis, :]\n if scale_factor < 1.0:\n raise ValueError(f'scale_factor must be >= 1.0, got {scale_factor}')\n sinusoid_inp /= scale_factor\n\n sin = jnp.sin(sinusoid_inp)\n cos = jnp.cos(sinusoid_inp)\n\n first_half, second_half = jnp.split(inputs, 2, axis=-1)\n first_part = first_half * cos - second_half * sin\n second_part = second_half * cos + first_half * sin\n out = jnp.concatenate([first_part, second_part], axis=-1)\n return out.astype(inputs.dtype)\n" }, { "path": "examples/gemma/positional_embeddings_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ============================================================================\n\"\"\"Tests for the positional embeddings utilities.\"\"\"\n\nfrom absl.testing import absltest\nfrom absl.testing import parameterized\nimport positional_embeddings\nimport jax\nimport jax.numpy as jnp\nimport numpy as np\n\n# positional_embeddings.py uses implicit rank broadcast and needs this config to\n# be 'allow', while the rest of Flax can use jax_numpy_rank_promotion=raise.\njax.config.update('jax_numpy_rank_promotion', 'allow')\n\n\nclass PositionalEmbeddingsTest(parameterized.TestCase):\n\n @parameterized.parameters(\n dict(\n input_embedding_shape=(2, 1, 1, 5),\n position=3,\n max_wavelength=100,\n expected=[[[[1.1411201, 1.0299965, 0.0100075, 1.99955, 1.0]]],\n [[[1.1411201, 1.0299965, 0.0100075, 1.99955, 1.0]]]]\n )\n )\n def test_adds_positional_embeddings(\n self, input_embedding_shape, position, max_wavelength, expected\n ):\n outputs = positional_embeddings.add_positional_embedding(\n jnp.ones(input_embedding_shape), position, max_wavelength\n )\n np.testing.assert_array_almost_equal(outputs, jnp.array(expected))\n\n @parameterized.parameters(\n dict(\n input_embedding_shape=(2, 1, 2, 4),\n position=3,\n head_dim=4,\n max_wavelength=100,\n expected=[\n [[\n [-1.1311126, 0.6598157, -0.8488725, 1.2508571],\n [-1.1311126, 0.6598157, -0.8488725, 1.2508571],\n ]],\n [[\n [-1.1311126, 0.6598157, -0.8488725, 1.2508571],\n [-1.1311126, 0.6598157, -0.8488725, 1.2508571],\n ]],\n ],\n )\n )\n def test_rope_positional_embeddings(\n self, input_embedding_shape, position, head_dim, max_wavelength, expected\n ):\n outputs = positional_embeddings.apply_rope(\n jnp.ones(input_embedding_shape),\n jnp.array([[position]]),\n head_dim,\n max_wavelength,\n )\n np.testing.assert_array_almost_equal(outputs, jnp.array(expected))\n\n\nif __name__ == \"__main__\":\n absltest.main()\n" }, { "path": "examples/gemma/requirements.txt", "content": "absl-py~=2.2\nclu==0.0.12\nflax~=0.10\njax~=0.6\nmlcroissant~=1.0\nnumpy~=2.1\noptax~=0.2\nsentencepiece~=0.2\njaxtyping~=0.3\ntensorflow~=2.19\ntensorflow-datasets~=4.9\ntensorflow-text~=2.19" }, { "path": "examples/gemma/sampler.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ============================================================================\n\"\"\"Sampler for Gemma transformer.\n\nAn example of a sampling class for a Gemma model.\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom collections.abc import Sequence\nimport dataclasses\n\nimport flax\nfrom flax import nnx\nimport modules\nimport sow_lib\nimport transformer as transformer_lib\nfrom flax.nnx import graph\nfrom flax.nnx import statelib\nimport jax\nimport jax.numpy as jnp\n\nimport sentencepiece as spm\n\n\ndef _sample_top_p(probs: jnp.ndarray, p: float, key: jax.Array) -> jnp.ndarray:\n \"\"\"Sample a token using top-p sampling.\"\"\"\n probs_sorted, indices = jax.lax.top_k(probs, k=probs.shape[-1])\n cumsum_probs = jnp.cumsum(probs_sorted, axis=-1)\n mask = cumsum_probs - probs_sorted > p\n probs_sorted = jnp.where(mask, 0.0, probs_sorted)\n probs_sorted /= jnp.sum(probs_sorted, axis=-1, keepdims=True)\n\n next_token = jax.random.categorical(key, logits=jnp.log(probs_sorted))\n\n next_token = jnp.take_along_axis(indices, next_token[..., None], axis=-1)\n next_token = jnp.squeeze(next_token, axis=-1)\n return next_token\n\n\ndef _compute_attention_masks(\n time_step: jax.Array, seq_len: int, input_mask: jax.Array\n) -> jax.Array:\n \"\"\"Computes causal attention mask.\"\"\"\n batch_size = input_mask.shape[0]\n batch_time_step = jnp.full((batch_size, 1), time_step, dtype=jnp.uint32)\n causal_padding = jnp.greater(\n jnp.expand_dims(jnp.arange(seq_len), 0), batch_time_step\n )\n max_seq_len = min(input_mask.shape[-1], seq_len)\n input_mask = jax.lax.dynamic_slice(\n input_mask,\n (0, jnp.maximum(time_step - seq_len + 1, 0)),\n (batch_size, max_seq_len),\n )\n input_mask = (\n jnp.zeros((batch_size, seq_len), dtype=jnp.bool_)\n .at[:, :max_seq_len]\n .set(input_mask)\n )\n\n causal_padding = jnp.logical_or(causal_padding, input_mask)\n attention_mask = causal_padding[:, jnp.newaxis, :].astype(jnp.bool_)\n\n return ~attention_mask\n\n\n@flax.struct.dataclass\nclass _SamplingState:\n \"\"\"Internal sampling state.\"\"\"\n\n # Decoding step.\n decoding_step: jnp.int32\n\n # Number of tokens in the prompt.\n num_input_tokens: jnp.ndarray # [B]\n\n # Fixed-size buffer for accumulating the output tokens.\n token_buffer: jnp.ndarray # [B, L]\n\n # Position indices, based on ignoring pad tokens.\n positions: jnp.ndarray # [B, L]\n\n # Model state for conditioning the model on autoregressively.\n cache: dict[str, modules.LayerCache]\n\n # Is decoding done on the given sequence?\n done: jnp.ndarray # [B]\n\n # Total sampling steps (including the prompt).\n total_sampling_steps: int\n\n # Fixed-size buffer for accumulating the output logits.\n logits_buffer: jnp.ndarray | None # [B, L, V]\n\n # List of tokens that are forbidden to be generated.\n forbidden_token_ids: Sequence[int] | None\n\n # Intermediate activations from the model if requested.\n intermediates: sow_lib.TransformerIntermediates | None\n\n # Random seed for sampling.\n seed: jax.Array\n\n # Tempurature for top_p sampling.\n temperature: float = flax.struct.field(pytree_node=False)\n\n # Top-p sampling threshold.\n top_p: float = flax.struct.field(pytree_node=False)\n\n\n@dataclasses.dataclass\nclass SamplerOutput:\n \"\"\"Output of the sampler.\"\"\"\n\n # Decoded samples from the model.\n text: list[str]\n\n # Per-step logits used during sampling.\n logits: list[list[float]]\n\n # Tokens corresponding to the generated samples.\n tokens: list[list[int]]\n\n # Intermediate activations from the model if requested.\n intermediates: sow_lib.TransformerIntermediates | None = None\n\n\nclass Sampler:\n \"\"\"Sampler for gemma transformer.\"\"\"\n\n def __init__(\n self,\n transformer: transformer_lib.Transformer,\n vocab: spm.SentencePieceProcessor,\n cache_size: int = 1024,\n ):\n \"\"\"Initializes a sampler for a Gemma model.\n\n Args:\n transformer: an instance of the Gemma transformer.\n vocab: vocabulary of the given model.\n cache_size: size of the cache for the transformer.\n \"\"\"\n self.vocab = vocab\n self.cache_size = cache_size\n graphdef, state = nnx.split(transformer)\n self._transformer_graphdef: graph.NodeDef = graphdef\n self._transformer_state: statelib.State = state\n # we separate out state and graph def so that the state can be passed as an\n # argument to _sample_fn, resulting in it not being treated as a static\n # arg. This greatly reduces the size of the HLO and reduces compile time\n self._compiled_sample_fn = jax.jit(self._sample_fn)\n\n @property\n def transformer(self) -> transformer_lib.Transformer:\n return nnx.merge(self._transformer_graphdef, self._transformer_state)\n\n @property\n def transformer_state(self) -> statelib.State:\n return self._transformer_state\n\n @transformer_state.setter\n def transformer_state(self, state: statelib.State) -> statelib.State:\n def check_tree_structure(tree1, tree2):\n if jax.tree_util.tree_structure(tree1) != jax.tree_util.tree_structure(\n tree2\n ):\n raise ValueError(\n \"New state must have the same structure as the old state.\"\n )\n\n def check_shape_dtype(x, y):\n return jnp.shape(x) == jnp.shape(y) and jnp.dtype(x) == jnp.dtype(y)\n\n if not all(\n jax.tree_util.tree_leaves(\n jax.tree_util.tree_map(check_shape_dtype, tree1, tree2)\n )\n ):\n raise ValueError(\n \"New state must have the same shape and dtype as the old state.\"\n )\n\n check_tree_structure(self._transformer_state, state)\n self._transformer_state = state\n\n @property\n def dtype(self) -> jnp.dtype:\n return jax.tree_util.tree_leaves(\n nnx.to_flat_state(self._transformer_state)\n )[0].dtype\n\n def _sample_step(\n self, params: statelib.State, sampler_state: _SamplingState\n ) -> _SamplingState:\n \"\"\"Performs a single sampling step.\"\"\"\n batch_size = sampler_state.token_buffer.shape[0]\n decoding_step = jnp.asarray(sampler_state.decoding_step, dtype=jnp.int32)\n last_token = sampler_state.token_buffer[:, decoding_step]\n input_mask = sampler_state.token_buffer == self.vocab.pad_id()\n attention_mask = _compute_attention_masks(\n decoding_step, self.cache_size, input_mask\n )\n step_positions = jnp.expand_dims(\n sampler_state.positions[:, decoding_step], -1\n )\n last_token = last_token.reshape((batch_size, 1))\n\n transformer = nnx.merge(self._transformer_graphdef, params)\n logits, cache = transformer(\n last_token,\n step_positions,\n sampler_state.cache,\n attention_mask,\n )\n if sampler_state.forbidden_token_ids:\n logits = logits.at[:, :, sampler_state.forbidden_token_ids].set(-jnp.inf)\n\n def sample_top_p(logits, key):\n probs = jax.nn.softmax(logits[:, -1] / sampler_state.temperature, axis=-1)\n next_token = _sample_top_p(probs, sampler_state.top_p, key)\n return next_token\n\n def sample_best(logits):\n next_token = jnp.argmax(logits, axis=-1)\n next_token = next_token[:, 0]\n return next_token\n\n if sampler_state.temperature > 0:\n key = jax.random.fold_in(sampler_state.seed, decoding_step)\n next_token_candidate = sample_top_p(logits, key)\n else:\n next_token_candidate = sample_best(logits)\n\n next_token_candidate = jnp.where(\n decoding_step < sampler_state.num_input_tokens - 1,\n sampler_state.token_buffer[:, decoding_step + 1],\n next_token_candidate,\n )\n\n token_buffer = sampler_state.token_buffer.at[:, decoding_step + 1].set(\n next_token_candidate\n )\n\n if sampler_state.logits_buffer is not None:\n next_logits = jnp.squeeze(logits, 1)\n logits_buffer = sampler_state.logits_buffer.at[:, decoding_step + 1].set(\n next_logits\n )\n else:\n logits_buffer = sampler_state.logits_buffer\n\n if sampler_state.intermediates is not None:\n sampler_state.intermediates.merge(decoding_step, transformer)\n\n done = sampler_state.done | jnp.equal(\n token_buffer[:, decoding_step + 1], self.vocab.eos_id()\n )\n\n return _SamplingState(\n decoding_step=sampler_state.decoding_step + 1,\n num_input_tokens=sampler_state.num_input_tokens,\n token_buffer=token_buffer,\n positions=sampler_state.positions,\n logits_buffer=logits_buffer,\n cache=cache,\n done=done,\n total_sampling_steps=sampler_state.total_sampling_steps,\n forbidden_token_ids=sampler_state.forbidden_token_ids,\n intermediates=sampler_state.intermediates,\n temperature=sampler_state.temperature,\n top_p=sampler_state.top_p,\n seed=sampler_state.seed,\n )\n\n def init_sample_state(\n self,\n all_input_ids: list[jax.Array],\n total_sampling_steps: int,\n include_logits: bool,\n forbidden_token_ids: Sequence[int] | None,\n temperature: float,\n top_p: float,\n seed: jax.Array,\n ) -> _SamplingState:\n \"\"\"Initializes the sampling state given input prompts.\"\"\"\n batch_size = len(all_input_ids)\n num_input_tokens = [len(input_ids) for input_ids in all_input_ids]\n buffer_size = total_sampling_steps + 1\n\n token_buffer = jnp.full(\n (\n batch_size,\n buffer_size,\n ),\n self.vocab.pad_id(),\n dtype=jnp.int32,\n )\n input_mask = jnp.ones_like(token_buffer, dtype=jnp.bool_)\n for i, (input_ids, num_tokens) in enumerate(\n zip(all_input_ids, num_input_tokens)\n ):\n token_buffer = token_buffer.at[i, :num_tokens].set(input_ids)\n input_mask = input_mask.at[i, :num_tokens].set(\n input_ids != self.vocab.pad_id()\n )\n positions = transformer_lib.build_positions_from_mask(input_mask)\n\n done = jnp.zeros((batch_size,), dtype=jnp.bool_)\n\n if include_logits:\n logits_buffer = jnp.zeros(\n (batch_size, buffer_size, self.transformer.num_embed),\n dtype=jnp.float32,\n )\n else:\n logits_buffer = None\n\n return _SamplingState(\n decoding_step=0,\n num_input_tokens=jnp.array(num_input_tokens, dtype=jnp.int32),\n token_buffer=token_buffer,\n positions=positions,\n logits_buffer=logits_buffer,\n cache=self.transformer.init_cache(\n cache_size=self.cache_size,\n batch_size=batch_size,\n dtype=self.dtype,\n ),\n done=done,\n total_sampling_steps=total_sampling_steps,\n forbidden_token_ids=forbidden_token_ids,\n intermediates=self.transformer.init_intermediates(\n batch_size, buffer_size, self.transformer.sow_config\n ),\n temperature=temperature,\n top_p=top_p,\n seed=seed,\n )\n\n def tokenize(self, input_string: str) -> jax.Array:\n \"\"\"Tokenizes the input string.\"\"\"\n input_ids = self.vocab.EncodeAsIds(input_string)\n input_ids = jnp.array(\n [self.vocab.bos_id()] + jnp.array(input_ids).tolist(), dtype=jnp.int32\n )\n return input_ids\n\n def mask_tokens_after_eos_ids(self, token_buffer):\n \"\"\"Mask token IDs after the EOS token with the padding ID.\"\"\"\n eos_id = self.vocab.eos_id()\n eos_exists = jnp.any(jnp.equal(token_buffer, eos_id), axis=-1)\n eos_indices = jnp.where(\n eos_exists,\n jnp.argmax(jnp.equal(token_buffer, eos_id), axis=-1),\n token_buffer.shape[-1],\n )\n mask = jnp.less_equal(\n jnp.arange(token_buffer.shape[-1]), eos_indices[:, None]\n )\n masked_token_buffer = token_buffer * mask + self.vocab.pad_id() * (1 - mask)\n\n return masked_token_buffer\n\n def _sample_fn(\n self,\n params: statelib.State,\n initial_sampling_state: _SamplingState,\n ) -> _SamplingState:\n \"\"\"Internal sampling function (to be jitted).\"\"\"\n\n def sample_with_params(sampler_state: _SamplingState):\n return self._sample_step(params, sampler_state)\n\n def cond_fn(sampler_state: _SamplingState):\n return (\n sampler_state.decoding_step < sampler_state.total_sampling_steps\n ) & jnp.any(jnp.logical_not(sampler_state.done))\n\n return jax.lax.while_loop(\n cond_fn, sample_with_params, initial_sampling_state\n )\n\n def __call__(\n self,\n input_strings: Sequence[str],\n total_generation_steps: int,\n echo: bool = False,\n return_logits: bool = True,\n forbidden_tokens: Sequence[str] | None = None,\n temperature: float = 0.0,\n top_p: float = 0.95,\n seed: jax.Array | None = None,\n ) -> SamplerOutput:\n \"\"\"Samples a completion of the input string.\n\n Args:\n input_strings: input prompts to feed to the model for sampling.\n total_generation_steps: number of generation steps. will correspond to the\n longest prompt in the batch.\n echo: whether to return the prompt as part of the output sample.\n return_logits: whether to return per-step logits used during generation.\n forbidden_tokens: list of tokens that are forbidden to be generated. Each\n token must map to a single token id in the vocab.\n temperature: temperature for sampling.\n top_p: top-p sampling threshold.\n seed: random seed for sampling.\n\n Returns:\n sampler_output: A SamplerOutput object containing the generated samples.\n \"\"\"\n forbidden_token_ids = None\n if forbidden_tokens is not None:\n forbidden_token_ids = []\n for token in forbidden_tokens:\n token_id = self.vocab.EncodeAsIds(token)\n if len(token_id) != 1:\n raise ValueError(\n \"Forbidden tokens must map to single token ids in the vocab.\"\n )\n forbidden_token_ids.extend(token_id)\n forbidden_token_ids = tuple(forbidden_token_ids)\n all_input_ids = [self.tokenize(x) for x in input_strings]\n max_input_length = max(len(input_ids) for input_ids in all_input_ids)\n total_sampling_steps = max_input_length + total_generation_steps\n\n if seed is None:\n seed = jax.random.PRNGKey(0)\n initial_sampling_state = self.init_sample_state(\n all_input_ids,\n include_logits=return_logits,\n total_sampling_steps=total_sampling_steps,\n forbidden_token_ids=forbidden_token_ids,\n temperature=temperature,\n top_p=top_p,\n seed=seed,\n )\n\n sampling_state = self._compiled_sample_fn(\n self._transformer_state, initial_sampling_state\n )\n\n masked_token_buffer = self.mask_tokens_after_eos_ids(\n sampling_state.token_buffer\n )\n\n out_tokens = []\n out_logits = []\n for i, (token_buffer, num_tokens) in enumerate(\n zip(\n masked_token_buffer,\n sampling_state.num_input_tokens,\n )\n ):\n start_idx = 0 if echo else num_tokens\n out_tokens.append(token_buffer[start_idx:total_sampling_steps].tolist())\n if return_logits:\n logits_buffer = sampling_state.logits_buffer[i]\n out_logits.append(\n logits_buffer[start_idx:total_sampling_steps].tolist()\n )\n\n decoded_outputs = [self.vocab.DecodeIds(tokens) for tokens in out_tokens]\n\n if sampling_state.intermediates is not None:\n sampling_state.intermediates.trim(total_sampling_steps)\n\n result = SamplerOutput(\n text=decoded_outputs,\n logits=out_logits,\n tokens=out_tokens,\n intermediates=sampling_state.intermediates,\n )\n return result\n" }, { "path": "examples/gemma/sampler_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ============================================================================\n\"\"\"Minimal test for sampler.\"\"\"\n\nimport os\nfrom collections.abc import Iterable\n\nfrom absl.testing import absltest\nfrom absl.testing import parameterized\nfrom flax import nnx\nimport modules\nimport params as params_lib\nimport sampler as sampler_lib\nimport sow_lib\nimport transformer as transformer_lib\nimport jax\nimport jax.numpy as jnp\nimport numpy as np\n\nimport sentencepiece as spm\n\n\nclass MockVocab(spm.SentencePieceProcessor):\n\n def __init__(self):\n super().__init__()\n self._start_id = 3\n self._mapping_text_to_id = {\n '': 0,\n '': 1,\n '': 2,\n 'input': 3,\n 'string': 4,\n 'hello': 5,\n 'world': 6,\n 'Hello': 7,\n 'there': 8,\n '!': 9,\n 'My': 10,\n 'name': 11,\n 'is': 12,\n 'Morgane': 13,\n }\n self._vocab_size = len(self._mapping_text_to_id)\n\n def pad_id(self) -> int:\n return 0\n\n def bos_id(self) -> int:\n return 1\n\n def eos_id(self) -> int:\n return 2\n\n def GetPieceSize(self) -> int: # pylint: disable=invalid-name\n return self._vocab_size\n\n def DecodeIds(self, ids: Iterable[int]) -> str: # pylint: disable=invalid-name\n reverse_mapping = {v: k for k, v in self._mapping_text_to_id.items()}\n return ' '.join(reverse_mapping[e] for e in ids)\n\n def EncodeAsIds(self, text: str) -> list[int]: # pylint: disable=invalid-name\n words = text.split(' ')\n return [self._mapping_text_to_id[word] for word in words]\n\n\nclass SamplerTest(parameterized.TestCase):\n\n def assertReasonableTensor(self, array, expected_shape=None):\n self.assertIsNotNone(array)\n if expected_shape is not None:\n self.assertEqual(array.shape, expected_shape)\n\n def test_samples(self):\n vocab = MockVocab()\n num_layers = 6\n transformer_config = transformer_lib.TransformerConfig( # pytype: disable=wrong-arg-types\n num_layers=num_layers,\n num_embed=vocab.GetPieceSize(),\n embed_dim=768,\n hidden_dim=6144,\n num_heads=4,\n num_kv_heads=4,\n head_dim=256,\n final_logit_softcap=None,\n attention_types=[modules.AttentionType.GLOBAL] * num_layers,\n attn_logits_soft_cap=None,\n use_post_attn_norm=None,\n use_post_ffw_norm=None,\n )\n transformer = transformer_lib.Transformer(\n transformer_config, rngs=nnx.Rngs(params=0)\n )\n sampler = sampler_lib.Sampler(\n transformer=transformer,\n vocab=vocab,\n cache_size=1024,\n )\n\n result = sampler(['input string', 'hello world'], total_generation_steps=10)\n self.assertIsNotNone(result)\n\n top_p_result = sampler(\n ['input string', 'hello world'],\n total_generation_steps=10,\n temperature=9,\n top_p=0.95,\n )\n self.assertIsNotNone(top_p_result)\n self.assertNotEqual(result.text, top_p_result.text)\n\n top_p_result_2 = sampler(\n ['input string', 'hello world'],\n total_generation_steps=10,\n temperature=9,\n top_p=0.95,\n seed=jax.random.PRNGKey(42),\n )\n self.assertIsNotNone(top_p_result_2)\n self.assertNotEqual(top_p_result.text, top_p_result_2.text)\n\n def test_state_update(self):\n vocab = MockVocab()\n num_layers = 6\n transformer_config = transformer_lib.TransformerConfig( # pytype: disable=wrong-arg-types\n num_layers=num_layers,\n num_embed=vocab.GetPieceSize(),\n embed_dim=768,\n hidden_dim=6144,\n num_heads=4,\n num_kv_heads=4,\n head_dim=256,\n final_logit_softcap=None,\n attention_types=[modules.AttentionType.GLOBAL] * num_layers,\n attn_logits_soft_cap=None,\n use_post_attn_norm=None,\n use_post_ffw_norm=None,\n )\n transformer = transformer_lib.Transformer(\n transformer_config, rngs=nnx.Rngs(params=0)\n )\n sampler = sampler_lib.Sampler(\n transformer=transformer,\n vocab=vocab,\n cache_size=1024,\n )\n input_strings = ['input string', 'hello world']\n original_logits = sampler(input_strings, total_generation_steps=10).logits\n\n new_transformer = transformer_lib.Transformer(\n transformer_config, rngs=nnx.Rngs(params=42)\n )\n sampler.transformer_state = nnx.state(new_transformer, nnx.Param)\n new_logits = sampler(input_strings, total_generation_steps=10).logits\n with self.assertRaises(AssertionError):\n np.testing.assert_allclose(\n original_logits, new_logits, atol=1e-1, rtol=1e-1\n )\n\n def test_invalid_state_update(self):\n vocab = MockVocab()\n\n def make_config(num_layers, embed_dim):\n return transformer_lib.TransformerConfig( # pytype: disable=wrong-arg-types\n num_layers=num_layers,\n num_embed=vocab.GetPieceSize(),\n embed_dim=embed_dim,\n hidden_dim=6144,\n num_heads=4,\n num_kv_heads=4,\n head_dim=256,\n final_logit_softcap=None,\n attention_types=[modules.AttentionType.GLOBAL] * num_layers,\n attn_logits_soft_cap=None,\n use_post_attn_norm=None,\n use_post_ffw_norm=None,\n )\n\n transformer = transformer_lib.Transformer(\n make_config(num_layers=6, embed_dim=768), rngs=nnx.Rngs(params=0)\n )\n sampler = sampler_lib.Sampler(\n transformer=transformer,\n vocab=vocab,\n cache_size=1024,\n )\n\n new_transformer = transformer_lib.Transformer(\n make_config(num_layers=3, embed_dim=768), rngs=nnx.Rngs(params=42)\n )\n with self.assertRaisesRegex(\n ValueError, '.*must have the same structure.*'\n ):\n sampler.transformer_state = nnx.state(new_transformer, nnx.Param)\n\n new_transformer = transformer_lib.Transformer(\n make_config(num_layers=6, embed_dim=1024), rngs=nnx.Rngs(params=42)\n )\n with self.assertRaisesRegex(\n ValueError, '.*must have the same shape and dtype.*'\n ):\n sampler.transformer_state = nnx.state(new_transformer, nnx.Param)\n\n def test_forbidden_tokens(self):\n vocab = MockVocab()\n transformer_config = transformer_lib.TransformerConfig( # pytype: disable=wrong-arg-types\n num_layers=0,\n num_embed=vocab.GetPieceSize(),\n embed_dim=32,\n hidden_dim=64,\n num_heads=4,\n num_kv_heads=1,\n head_dim=64,\n final_logit_softcap=None,\n attention_types=[],\n use_post_attn_norm=None,\n use_post_ffw_norm=None,\n )\n transformer = transformer_lib.Transformer(\n transformer_config, rngs=nnx.Rngs(params=0)\n )\n # Pre-cook the embedding matrix so that the output is deterministic.\n transformer.embedder.input_embedding.set_value(jnp.eye(\n vocab.GetPieceSize(), 32\n ))\n sampler = sampler_lib.Sampler(\n transformer=transformer,\n vocab=vocab,\n cache_size=8,\n )\n\n # First, we check that the sampler would produce the tokens that we are\n # trying to forbid.\n result1 = sampler(\n ['input string', 'hello world'],\n total_generation_steps=10,\n forbidden_tokens=None,\n )\n self.assertIn('string', result1.text[0])\n self.assertIn('world', result1.text[1])\n\n # Then, we check that the sampler does not produce the forbidden tokens.\n result2 = sampler(\n ['input string', 'hello world'],\n total_generation_steps=10,\n forbidden_tokens=['string', 'world'],\n )\n for output in result2.text:\n self.assertNotIn('string', output)\n self.assertNotIn('world', output)\n\n def test_forward_equivalence(self):\n vocab = MockVocab()\n num_layers = 2\n transformer_config = transformer_lib.TransformerConfig( # pytype: disable=wrong-arg-types\n num_layers=num_layers,\n num_embed=vocab.GetPieceSize(),\n embed_dim=32,\n hidden_dim=64,\n num_heads=4,\n num_kv_heads=1,\n head_dim=64,\n final_logit_softcap=None,\n attention_types=[modules.AttentionType.GLOBAL] * num_layers,\n use_post_attn_norm=None,\n use_post_ffw_norm=None,\n )\n\n transformer = transformer_lib.Transformer(\n transformer_config, rngs=nnx.Rngs(params=0)\n )\n raw_input = 'Hello there ! My name is Morgane '\n token_input = jnp.asarray(\n [vocab.bos_id()] + vocab.EncodeAsIds(raw_input)\n ).reshape((1, -1))\n batch_size = 1\n cache_size = 9\n cache = transformer.init_cache(\n cache_size=cache_size,\n batch_size=batch_size,\n dtype=jnp.float32,\n )\n input_mask = token_input != vocab.pad_id()\n positions = transformer_lib.build_positions_from_mask(input_mask)\n attention_mask = transformer_lib.make_causal_attn_mask(input_mask)\n\n n_input_tokens = token_input.shape[1]\n\n output_forward, _ = transformer(\n last_tokens=token_input,\n positions=positions,\n cache=cache,\n attention_mask=attention_mask,\n )\n output_forward = output_forward[0, :n_input_tokens]\n\n sampler = sampler_lib.Sampler(\n transformer=transformer,\n vocab=vocab,\n cache_size=cache_size,\n )\n\n output_transformer = sampler(\n [raw_input],\n total_generation_steps=10,\n echo=True,\n )\n out_logits = np.array(output_transformer.logits)[0, 1 : n_input_tokens + 1]\n\n np.testing.assert_almost_equal(output_forward, out_logits)\n\n def test_sampler_init_sample_state(self):\n vocab = MockVocab()\n transformer_config = transformer_lib.TransformerConfig( # pytype: disable=wrong-arg-types\n num_layers=0,\n num_embed=vocab.GetPieceSize(),\n embed_dim=32,\n hidden_dim=64,\n num_heads=4,\n num_kv_heads=1,\n head_dim=64,\n final_logit_softcap=None,\n attention_types=[],\n use_post_attn_norm=None,\n use_post_ffw_norm=None,\n )\n transformer = transformer_lib.Transformer(\n transformer_config, rngs=nnx.Rngs(params=0)\n )\n sampler = sampler_lib.Sampler(\n transformer=transformer,\n vocab=vocab,\n cache_size=8,\n )\n\n input_strings = [' hello world', 'input string ']\n all_input_ids = [sampler.tokenize(x) for x in input_strings]\n total_sampling_steps = 5\n sample_state = sampler.init_sample_state(\n all_input_ids,\n total_sampling_steps=total_sampling_steps,\n include_logits=True,\n forbidden_token_ids=None,\n temperature=0.0,\n top_p=0.95,\n seed=jax.random.PRNGKey(0),\n )\n\n # Check that the position indices correctly ignore padding\n self.assertListEqual(list(sample_state.positions[0]), [0, 0, 1, 2, 3, 4])\n self.assertListEqual(list(sample_state.positions[1]), [0, 1, 2, 2, 3, 4])\n\n def test_sampler_mask_tokens_after_eos_ids(self):\n vocab = MockVocab()\n transformer_config = transformer_lib.TransformerConfig( # pytype: disable=wrong-arg-types\n num_layers=0,\n num_embed=vocab.GetPieceSize(),\n embed_dim=32,\n hidden_dim=64,\n num_heads=4,\n num_kv_heads=1,\n head_dim=64,\n final_logit_softcap=None,\n attention_types=[],\n use_post_attn_norm=None,\n use_post_ffw_norm=None,\n )\n transformer = transformer_lib.Transformer(\n transformer_config, rngs=nnx.Rngs(params=0)\n )\n sampler = sampler_lib.Sampler(\n transformer=transformer,\n vocab=vocab,\n cache_size=8,\n )\n\n input_strings = ['hello world hello world', 'input string hello']\n all_input_ids = [sampler.tokenize(x) for x in input_strings]\n total_sampling_steps = 5\n sample_state = sampler.init_sample_state(\n all_input_ids,\n total_sampling_steps=total_sampling_steps,\n include_logits=True,\n forbidden_token_ids=None,\n temperature=0.0,\n top_p=0.95,\n seed=jax.random.PRNGKey(0),\n )\n\n masked_token_buffer = sampler.mask_tokens_after_eos_ids(\n sample_state.token_buffer\n )\n\n self.assertListEqual(list(masked_token_buffer[0]), [1, 5, 6, 2, 0, 0])\n self.assertListEqual(list(masked_token_buffer[1]), [1, 3, 4, 2, 0, 0])\n\n def test_sampler_sows_intermediates(self):\n vocab = MockVocab()\n num_layers = 3\n config = transformer_lib.TransformerConfig( # pytype: disable=wrong-arg-types\n num_layers=num_layers,\n num_embed=vocab.GetPieceSize(),\n embed_dim=64,\n hidden_dim=128,\n num_heads=2,\n num_kv_heads=1,\n head_dim=64,\n final_logit_softcap=None,\n attention_types=[modules.AttentionType.GLOBAL] * num_layers,\n use_post_attn_norm=None,\n attn_logits_soft_cap=None,\n use_post_ffw_norm=None,\n )\n sow_config = sow_lib.SowConfig(\n embeddings=True,\n rs_after_attention=False, # This should results in a None value.\n rs_after_ffw=True,\n attn_logits_topk=5,\n mlp_hidden_topk=11,\n )\n transformer = transformer_lib.Transformer(\n config, rngs=nnx.Rngs(params=0), sow_config=sow_config\n )\n sampler = sampler_lib.Sampler(\n transformer=transformer,\n vocab=vocab,\n )\n raw_input = ['input string', 'hello world']\n\n result = sampler(raw_input, total_generation_steps=10)\n input_length = max([len(vocab.EncodeAsIds(i)) for i in raw_input])\n input_length += 1 # +1 for BOS token\n output_length = max(len(tokens) for tokens in result.tokens)\n length = input_length + output_length\n self.assertIsNotNone(result)\n intermediates = result.intermediates\n self.assertIsNotNone(intermediates)\n self.assertReasonableTensor(\n intermediates.embeddings,\n expected_shape=(2, length, config.embed_dim),\n )\n # Verify that the intermediates are different for two different steps.\n self.assertNotAlmostEqual(\n jnp.sum(intermediates.embeddings[:, 1, ...]),\n jnp.sum(intermediates.embeddings[:, 2, ...]),\n )\n # Verify that the intermediates are filled in for each layer.\n self.assertLen(intermediates.layers, config.num_layers)\n for layer in intermediates.layers:\n # For the requested intermediates we check the shape and that values are\n # not all zeros, which was the initial value.\n self.assertReasonableTensor(\n layer.rs_after_ffw,\n expected_shape=(2, length, config.embed_dim),\n )\n self.assertReasonableTensor(\n layer.attn_logits_topk_values,\n expected_shape=(\n 2,\n length,\n config.num_heads,\n sow_config.attn_logits_topk,\n ),\n )\n self.assertReasonableTensor(\n layer.attn_logits_topk_indices,\n expected_shape=(\n 2,\n length,\n config.num_heads,\n sow_config.attn_logits_topk,\n ),\n )\n self.assertReasonableTensor(\n layer.mlp_hidden_topk_values,\n expected_shape=(2, length, sow_config.mlp_hidden_topk),\n )\n self.assertReasonableTensor(\n layer.mlp_hidden_topk_indices,\n expected_shape=(2, length, sow_config.mlp_hidden_topk),\n )\n # For the none requested intermediates we want to have None values.\n self.assertIsNone(layer.rs_after_attention)\n\n def test_compute_attention_mask(self):\n # Check that the input mask is correctly applied when total sampling steps\n # is lower than the max cache length.\n input_mask = jnp.array([[1, 1, 0, 0, 0], [1, 1, 0, 1, 0]], dtype=jnp.bool_)\n seq_len = 8\n time_step = jnp.asarray(4, dtype=jnp.int32)\n attn_mask = sampler_lib._compute_attention_masks(\n time_step, seq_len, input_mask\n )\n expected_attn_mask = jnp.array(\n [[0, 0, 1, 1, 1, 0, 0, 0], [0, 0, 1, 0, 1, 0, 0, 0]], dtype=jnp.bool_\n )\n\n self.assertTrue((attn_mask.squeeze(1) == expected_attn_mask).all())\n\n # Check that the input mask is correctly applied when total sampling steps\n # is *longer* than the max cache length.\n seq_len = 4\n time_step = jnp.asarray(4, dtype=jnp.int32)\n attn_mask = sampler_lib._compute_attention_masks(\n time_step, seq_len, input_mask\n )\n expected_attn_mask = jnp.array(\n [[0, 1, 1, 1], [0, 1, 0, 1]], dtype=jnp.bool_\n )\n\n self.assertTrue((attn_mask.squeeze(1) == expected_attn_mask).all())\n\n @parameterized.parameters(\n {\"url\": \"google/gemma/flax/2b\"},\n {\"url\": \"google/gemma-2/flax/gemma2-2b\"},\n {\"url\": \"google/gemma-3/flax/gemma3-1b\"},\n )\n def test_models_from_kaggle(self, url):\n # A smoke test based on guide/gemma.md to ensure models are working correctly\n # Check Kaggle creds as env var otherwise skip the test\n has_kaggle_creds = all(k in os.environ for k in [\"KAGGLE_USERNAME\", \"KAGGLE_KEY\"])\n try:\n import kagglehub\n\n has_kagglehub_dep = True\n except ModuleNotFoundError:\n has_kagglehub_dep = False\n\n if not (has_kaggle_creds and has_kagglehub_dep):\n self.skipTest('Skip the test as no Kaggle deps/creds')\n\n variant = url.split(\"/\")[-1]\n weights_dir = kagglehub.model_download(url)\n ckpt_path = f\"{weights_dir}/{variant}\"\n vocab_path = f\"{weights_dir}/tokenizer.model\"\n\n vocab = spm.SentencePieceProcessor()\n vocab.Load(vocab_path)\n\n params = params_lib.load_and_format_params(ckpt_path)\n transformer = transformer_lib.Transformer.from_params(params)\n sampler = sampler_lib.Sampler(\n transformer=transformer,\n vocab=vocab,\n )\n\n input_batch = [\n \"# Python function to compute a square of the input number\",\n ]\n out_data = sampler(\n input_strings=input_batch,\n total_generation_steps=50,\n )\n assert \"def square(\" in out_data.text[0], out_data.text[0]\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "examples/gemma/sow_lib.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ============================================================================\n\"\"\"Utilities for sowing intermediate activations.\"\"\"\n\nimport dataclasses\nfrom flax import nnx\nimport jax\nimport jax.numpy as jnp\n\n\n@jax.tree_util.register_dataclass\n@dataclasses.dataclass\nclass LayerIntermediates:\n \"\"\"Intermediate activations for a single layer.\"\"\"\n\n # Dense residual stream activations.\n rs_after_attention: jax.Array | None = None\n rs_after_ffw: jax.Array | None = None\n\n # Sparse representations for large activations.\n mlp_hidden_topk_values: jax.Array | None = None\n mlp_hidden_topk_indices: jax.Array | None = None\n attn_logits_topk_values: jax.Array | None = None\n attn_logits_topk_indices: jax.Array | None = None\n\n def merge(self, decoding_step, layer: nnx.Module):\n \"\"\"Merges the intermediate activations from one step.\"\"\"\n\n for field in dataclasses.fields(self.__class__):\n value = getattr(self, field.name)\n if value is None:\n continue\n # We put mlp and attn intermediates into this class without any further\n # nesting. So we have to retrieve the intermediates from the correct\n # sub-module.\n try:\n if field.name.startswith('attn_'):\n step_value = getattr(\n layer.attn, field.name.replace('attn_', '')\n )[0]\n elif field.name.startswith('mlp_'):\n step_value = getattr(layer.mlp, field.name.replace('mlp_', ''))[0]\n else:\n step_value = getattr(layer, field.name)[0]\n except AttributeError as exc:\n raise ValueError(\n f'Intermediate {field.name} is not in the step intermediates.'\n ) from exc\n # This logic is the same for all intermediates. The second dimenions is\n # the length dimension, where we want to merge the intermediates from\n # multiple steps.\n setattr(\n self,\n field.name,\n value.at[:, decoding_step + 1].set(step_value[:, 0, ...]),\n )\n\n def trim(self, max_length: int):\n \"\"\"Trims the intermediate activations to the given length.\"\"\"\n for field in dataclasses.fields(self.__class__):\n value = getattr(self, field.name)\n if value is not None:\n setattr(self, field.name, value[:, :max_length, ...])\n\n\n@jax.tree_util.register_dataclass\n@dataclasses.dataclass\nclass TransformerIntermediates:\n \"\"\"Intermediate activations of a transformer network.\"\"\"\n\n # Embeddings of the input tokens.\n embeddings: jax.Array | None = None\n\n # Intermediate activations of each layer.\n layers: list[LayerIntermediates] = dataclasses.field(default_factory=list)\n\n def merge(self, decoding_step, transformer: nnx.Module):\n \"\"\"Merges the intermediate activations from one step.\"\"\"\n if self.embeddings is not None:\n try:\n self.embeddings = self.embeddings.at[:, decoding_step + 1, ...].set(\n transformer.embeddings[0][:, 0, ...]\n )\n except AttributeError as exc:\n raise ValueError(\n 'Embeddings are not in the step intermediates.'\n ) from exc\n if len(self.layers) != len(transformer.layers):\n raise ValueError(\n 'Number of layers in the transformer and intermediates do not match.'\n )\n for layer_intermediates, layer_module in zip(\n self.layers, transformer.layers\n ):\n layer_intermediates.merge(decoding_step, layer_module)\n\n def trim(self, max_length: int):\n \"\"\"Trims the intermediate activations to the given length.\"\"\"\n if self.embeddings is not None:\n self.embeddings = self.embeddings[:, :max_length, ...]\n for layer in self.layers:\n layer.trim(max_length)\n\n\n@dataclasses.dataclass(frozen=True)\nclass SowConfig:\n \"\"\"Module for sowing intermediate activations.\"\"\"\n\n # Whether to sow embeddings.\n embeddings: bool = False\n\n # Whether to sow activations after each attention block (in residual stream).\n rs_after_attention: bool = False\n\n # Whether to sow activations after each FFW block (in residual stream).\n # This is the same as the residual stream activations after a whole layer.\n rs_after_ffw: bool = False\n\n # If non-zero, top-k activations in a ffw hidden layer are sowed.\n # We use a sparse representation here to save memory.\n mlp_hidden_topk: int = 0\n\n # If non-zero, top-k attention logits are sowed.\n # We use a sparse representation here to save memory.\n attn_logits_topk: int = 0\n\n def maybe_sow_embeddings(\n self,\n embeddings: jax.Array,\n module: nnx.Module,\n ):\n \"\"\"Sows embeddings if configured.\"\"\"\n if self.embeddings:\n module.sow(nnx.Intermediate, 'embeddings', embeddings)\n\n def maybe_sow_rs_after_attention(\n self,\n activations: jax.Array,\n module: nnx.Module,\n ):\n \"\"\"Sows activations after attention if configured.\"\"\"\n if self.rs_after_attention:\n module.sow(nnx.Intermediate, 'rs_after_attention', activations)\n\n def maybe_sow_rs_after_ffw(\n self,\n activations: jax.Array,\n module: nnx.Module,\n ):\n \"\"\"Sows activations after FFW if configured.\"\"\"\n if self.rs_after_ffw:\n module.sow(nnx.Intermediate, 'rs_after_ffw', activations)\n\n def maybe_sow_mlp_hidden_topk(\n self,\n activations: jax.Array,\n module: nnx.Module,\n ):\n \"\"\"Sows top-absolute-k activations in a mlp hidden layer if configured.\"\"\"\n if self.mlp_hidden_topk:\n _, indices = jax.lax.top_k(jnp.abs(activations), self.mlp_hidden_topk)\n values = jnp.take_along_axis(activations, indices, axis=-1)\n module.sow(nnx.Intermediate, 'hidden_topk_values', values)\n module.sow(nnx.Intermediate, 'hidden_topk_indices', indices)\n\n def maybe_sow_attn_logits_topk(\n self,\n logits: jax.Array,\n module: nnx.Module,\n ):\n \"\"\"Sows top-k attention logits if configured.\"\"\"\n if self.attn_logits_topk:\n values, indices = jax.lax.top_k(logits, self.attn_logits_topk)\n module.sow(nnx.Intermediate, 'logits_topk_values', values)\n module.sow(nnx.Intermediate, 'logits_topk_indices', indices)\n" }, { "path": "examples/gemma/tokenizer.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Provides op for tokenizing a dataset.\"\"\"\n\nfrom collections.abc import Iterable\nimport dataclasses\nimport os\nimport sys\nimport tempfile\nimport time\nfrom typing import Any\n\nfrom absl import logging\nimport jax\nimport tensorflow as tf\n\nfrom sentencepiece import SentencePieceProcessor # pylint: disable=g-importing-member\nfrom sentencepiece import SentencePieceTrainer # pylint: disable=g-importing-member\n\nif sys.version_info < (3, 13):\n import tensorflow_text as tftxt\n\nFeatures = dict[str, tf.Tensor]\n\n\ndef _dump_chars_to_textfile(\n dataset: tf.data.Dataset,\n maxchars: int = int(1e7),\n data_keys=('inputs', 'targets'),\n) -> tuple[str, int]:\n \"\"\"Write part of a TFDS sentence dataset to lines in a text file.\n\n Args:\n dataset: tf.dataset containing string-data.\n maxchars: int: approximate number of characters to save from dataset.\n data_keys: Tuple[str]: what keys in dataset to dump from.\n\n Returns:\n name of temp file with dataset bytes, exact number of characters dumped.\n \"\"\"\n char_count = 0\n ds_iter = dataset.as_numpy_iterator()\n with tempfile.NamedTemporaryFile(\n delete=False, prefix='/tmp/ds_chars'\n ) as outfp:\n while char_count < maxchars:\n example = next(ds_iter)\n for k in data_keys:\n line = example[k] + b'\\n'\n char_count += len(line)\n outfp.write(line)\n return outfp.name, char_count\n\n\ndef _train_sentencepiece(\n dataset: tf.data.Dataset,\n *,\n vocab_size: int,\n maxchars: int = int(1e7),\n model_path: str,\n model_type: str = 'unigram',\n character_coverage: float = 1.0,\n data_keys=('inputs', 'targets'),\n pad_id: int = 0,\n eos_id: int = 1,\n bos_id: int = 2,\n unk_id: int = 3,\n):\n \"\"\"Train SentencePiece tokenizer from subset of tf dataset.\n\n Args:\n dataset: tf.dataset\n vocab_size: int: size of vocab tokens to train.\n maxchars: int: number of characters to use for sentencepiece training.\n model_path: str: path of model file to save vocab model to.\n model_type: str: type of sentencepiece vocab to train.\n character_coverage: amount of characters covered by the model, good defaults\n are 0.9995 for languages with rich character set like Japanese or Chinese\n and 1.0 for other languages with small character set.\n data_keys: Tuple[str]: keys of dataset to use for training.\n pad_id: int: pad piece id\n eos_id: int: end of sentence piece id\n bos_id: int: begin of sentence piece id\n unk_id: int: unknown piece id\n\n Returns:\n path to the trained sentencepiece vocabulary model.\n \"\"\"\n if model_path.startswith('gs://'):\n abs_model_path = model_path\n else:\n abs_model_path = os.path.abspath(os.path.expanduser(model_path))\n fname, _ = _dump_chars_to_textfile(\n dataset, maxchars=maxchars, data_keys=data_keys\n )\n with tempfile.NamedTemporaryFile(\n delete=False, prefix='/tmp/sp_tmp'\n ) as model_fp:\n pass # we just want a prefix'd tmp-filename\n argstr = ' '.join([\n f'--input={fname}',\n f'--vocab_size={vocab_size}',\n f'--character_coverage={character_coverage}',\n f'--model_prefix={model_fp.name}',\n f'--model_type={model_type}',\n # Setup ids for PAD, EOS, BOS, UNK as 0, 1, 2, 3\n # Default values:\n # --unk_id (Override UNK () id.) type: int32 default: 0\n # --bos_id (Override BOS () id. Set -1 to disable BOS.) type: int32 default: 1\n # --eos_id (Override EOS () id. Set -1 to disable EOS.) type: int32 default: 2\n # --pad_id (Override PAD () id. Set -1 to disable PAD.) type: int32 default: -1\n # https://github.com/google/sentencepiece/blob/master/doc/options.md\n f'--pad_id={pad_id}',\n f'--bos_id={bos_id}',\n f'--eos_id={eos_id}',\n f'--unk_id={unk_id}',\n ])\n SentencePieceTrainer.Train(argstr)\n if jax.process_index() == 0:\n # Use an intermediate filename that is renamed to the target name to address\n # create and fill delays.\n copy_rename_path = abs_model_path + '.rntmp'\n tf.io.gfile.copy(model_fp.name + '.model', copy_rename_path, overwrite=True)\n tf.io.gfile.rename(copy_rename_path, abs_model_path, overwrite=True)\n logging.info('copied %s to %s', model_fp.name + '.model', abs_model_path)\n else:\n while not tf.io.gfile.exists(abs_model_path):\n time.sleep(1)\n time.sleep(1)\n return abs_model_path\n\n\ndef _load_sentencepiece_tokenizer(\n model_path: str,\n add_bos: bool = False,\n add_eos: bool = True,\n reverse: bool = False,\n):\n \"\"\"Load a tf-text SentencePiece tokenizer from given model filepath.\"\"\"\n with tf.io.gfile.GFile(model_path, 'rb') as model_fp:\n sp_model = model_fp.read()\n sp_tokenizer = tftxt.SentencepieceTokenizer(\n model=sp_model, add_bos=add_bos, add_eos=add_eos, reverse=reverse\n )\n return sp_tokenizer\n\n\ndef load_or_train_tokenizer(\n dataset: tf.data.Dataset,\n *,\n vocab_path: str,\n vocab_size: int,\n max_corpus_chars: int,\n data_keys: tuple[str, str] = ('inputs', 'targets'),\n):\n \"\"\"Loads the tokenizer at `vocab_path` or trains a one from `dataset`.\"\"\"\n try:\n return _load_sentencepiece_tokenizer(vocab_path)\n except tf.errors.NotFoundError:\n logging.info('SentencePiece vocab not found, building one from data.')\n vocab_path = _train_sentencepiece(\n dataset,\n vocab_size=vocab_size,\n maxchars=max_corpus_chars,\n model_path=vocab_path,\n data_keys=data_keys,\n )\n return _load_sentencepiece_tokenizer(vocab_path)\n\n\n@dataclasses.dataclass\nclass TokenizeOp:\n sp_tokenizer: Any\n data_keys: Iterable[str] = ('inputs', 'targets')\n\n def __call__(self, features: Features) -> Features:\n for k in self.data_keys:\n features[k] = self.sp_tokenizer.tokenize(features[k])\n return features\n\n\ndef load_sentencepiece_processor(vocab_path: str):\n spp = SentencePieceProcessor()\n spp.Load(vocab_path)\n return spp\n" }, { "path": "examples/gemma/train.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Language Modeling example.\n\nThis script trains a Transformer on a LM1B dataset.\n\"\"\"\n\n# pytype: disable=wrong-arg-count\n# pytype: disable=attribute-error\n\nimport dataclasses\nimport os\nfrom typing import Any\n\nfrom absl import logging\nfrom clu import metric_writers\nfrom clu import periodic_actions\nfrom flax import nnx\nimport input_pipeline\nimport sampler as sampler_lib\nimport tokenizer\nimport transformer as transformer_lib\nimport utils\nfrom flax.training import checkpoints\nfrom flax.training import common_utils\nimport jax\nfrom jax import random\nimport jax.numpy as jnp\nimport numpy as np\nimport optax\nimport tensorflow as tf\n\n\n@dataclasses.dataclass(unsafe_hash=True)\nclass MeshRules:\n embed: str | None = None\n mlp: str | None = None\n kv: str | None = None\n vocab: str | None = None\n\n def __call__(self, *keys: str) -> tuple[str, ...]:\n return tuple(\n getattr(self, key) if key is not None else None for key in keys\n )\n\n\n@dataclasses.dataclass(unsafe_hash=True)\nclass TrainConfig:\n \"\"\"Configuration for training a gemma model.\"\"\"\n\n # Path to load or store sentencepiece vocab file.\n vocab_path: str | None\n # Vocabulary size if `vocab_path` is not given.\n vocab_size: int\n # Maximum number of characters to use for training.\n max_corpus_chars: int\n # Name of TFDS translation dataset to use.\n dataset_name: str\n # Optional name of TFDS translation dataset to use for evaluation.\n eval_dataset_name: str\n # Optional name of TFDS split to use for evaluation.\n eval_split: str\n # Per device batch size for training.\n per_device_batch_size: int\n # Per device batch size for training.\n eval_per_device_batch_size: int\n\n # Prompt for language model sampling\n prompts: tuple[str, ...]\n # Temperature for top_p sampling.\n sampling_temperature: float\n # Top-p sampling threshold.\n sampling_top_p: float\n\n # Number of steps to take during training.\n num_train_steps: int\n # Number of steps to take during evaluation.\n # Large enough to evaluate all samples: 306_688 / (32 * 8) = 1198\n num_eval_steps: int\n # Number of steps to generate predictions.\n # -1 will use the whole eval dataset.\n num_predict_steps: int\n # Base learning rate.\n learning_rate: float\n # Linear learning rate warmup.\n warmup_steps: int\n # Cross entropy loss label smoothing.\n label_smoothing: float\n # Decay factor for AdamW style weight decay.\n weight_decay: float\n # Maximum length cutoff for training examples.\n max_target_length: int\n # Maximum length cutoff for eval examples.\n max_eval_target_length: int\n\n # Gemma transformer name.\n # Possible values defined in transformer.TransformerConfig:\n # (gemma_2b, gemma_7b, gemma2_2b, gemma2_9b, gemma2_27b, gemma3_1b, gemma3_4b,\n # ...)\n transformer_name: str | None\n # or alternatively define the model using the dict of parameters\n transformer_params: dict[Any, Any] | None\n\n # Whether to save model checkpoints.\n save_checkpoints: bool\n # Whether to restore from existing model checkpoints.\n restore_checkpoints: bool\n # Save a checkpoint every these number of steps.\n checkpoint_every_steps: int\n # Frequency of eval during training, e.g. every 1_000 steps.\n eval_every_steps: int\n # Use bfloat16 mixed precision training instead of float32.\n use_bfloat16: bool\n # Integer for PRNG random seed.\n seed: int\n\n # Parallelism\n mesh_axes: tuple[str, ...]\n axis_rules: MeshRules\n data_sharding: tuple[str, ...]\n\n # One axis for each parallelism type may hold a placeholder (-1)\n # value to auto-shard based on available slices and devices.\n # By default, product of the DCN axes should equal number of slices\n # and product of the ICI axes should equal number of devices per slice.\n # ICI (Inter-Chip Interconnection): A high-speed connection between\n # sets of TPU chips, which form the TPU network.\n # DCN (Data Center Network): A connection between the TPU networks;\n # not as fast as ICI.\n # ICI has around 100x the bandwidth of DCN, but it is not a general\n # purpose connection, which is why DCN is necessary for scaling to\n # extremely large ML models.\n dcn_data_parallelism: int = -1\n dcn_fsdp_parallelism: int = 1\n dcn_tensor_parallelism: int = 1\n ici_data_parallelism: int = 1\n ici_fsdp_parallelism: int = -1\n ici_tensor_parallelism: int = 1\n\n def replace(self, **kwargs):\n return dataclasses.replace(self, **kwargs)\n\n def __post_init__(self):\n if isinstance(self.axis_rules, dict):\n self.axis_rules = MeshRules(**self.axis_rules)\n\n\ndef rsqrt_schedule(\n init_value: float,\n shift: int = 0,\n):\n \"\"\"Applies a reverse square-root schedule.\n\n The reverse square root schedule is simply `lr = init_value / sqrt(step)`.\n\n Args:\n init_value: Base learning rate (before applying the rsqrt schedule).\n shift: How many steps the rsqrt should be shifted. Shifting the rsqrt\n schedule makes it less steep in the beginning (close to 0).\n\n Returns:\n A schedule that applies the reverse square root.\n \"\"\"\n\n def schedule(count):\n return init_value * (count + shift) ** -0.5 * shift**0.5\n\n return schedule\n\n\ndef create_learning_rate_schedule(learning_rate: float, warmup_steps: int):\n \"\"\"Creates a rsqrt schedule with linear warmup.\"\"\"\n return optax.join_schedules(\n [\n optax.linear_schedule(\n init_value=0,\n end_value=learning_rate,\n transition_steps=warmup_steps,\n ),\n rsqrt_schedule(init_value=learning_rate, shift=warmup_steps),\n ],\n boundaries=[warmup_steps],\n )\n\n\ndef compute_weighted_cross_entropy(\n logits, targets, weights=None, label_smoothing=0.0\n):\n \"\"\"Compute weighted cross entropy and entropy for log probs and targets.\n\n Args:\n logits: [batch, length, num_classes] float array.\n targets: categorical targets [batch, length] int array.\n weights: None or array of shape [batch, length].\n label_smoothing: label smoothing constant, used to determine the on and off\n values.\n\n Returns:\n Tuple of scalar loss and batch normalizing factor.\n \"\"\"\n if logits.ndim != targets.ndim + 1:\n raise ValueError(\n 'Incorrect shapes. Got shape %s logits and %s targets'\n % (str(logits.shape), str(targets.shape))\n )\n vocab_size = logits.shape[-1]\n confidence = 1.0 - label_smoothing\n low_confidence = (1.0 - confidence) / (vocab_size - 1)\n normalizing_constant = -(\n confidence * jnp.log(confidence)\n + (vocab_size - 1) * low_confidence * jnp.log(low_confidence + 1e-20)\n )\n soft_targets = common_utils.onehot(\n targets, vocab_size, on_value=confidence, off_value=low_confidence\n )\n\n loss = -jnp.sum(soft_targets * nnx.log_softmax(logits), axis=-1)\n loss = loss - normalizing_constant\n\n normalizing_factor = np.prod(targets.shape)\n if weights is not None:\n loss = loss * weights\n normalizing_factor = weights.sum()\n\n return loss.sum(), normalizing_factor\n\n\ndef compute_weighted_accuracy(logits, targets, weights=None):\n \"\"\"Compute weighted accuracy for log probs and targets.\n\n Args:\n logits: [batch, length, num_classes] float array.\n targets: categorical targets [batch, length] int array.\n weights: None or array of shape [batch, length]\n\n Returns:\n Tuple of scalar loss and batch normalizing factor.\n \"\"\"\n if logits.ndim != targets.ndim + 1:\n raise ValueError(\n 'Incorrect shapes. Got shape %s logits and %s targets'\n % (str(logits.shape), str(targets.shape))\n )\n loss = jnp.equal(jnp.argmax(logits, axis=-1), targets)\n normalizing_factor = np.prod(logits.shape[:-1])\n if weights is not None:\n loss = loss * weights\n normalizing_factor = weights.sum()\n\n return loss.sum(), normalizing_factor\n\n\ndef compute_metrics(logits, labels, weights, label_smoothing=0.0):\n \"\"\"Compute summary metrics.\"\"\"\n loss, weight_sum = compute_weighted_cross_entropy(\n logits, labels, weights, label_smoothing\n )\n acc, _ = compute_weighted_accuracy(logits, labels, weights)\n metrics = {\n 'loss': loss,\n 'accuracy': acc,\n 'denominator': weight_sum,\n }\n return metrics\n\n\n# Primary training / eval / decode step functions.\n# -----------------------------------------------------------------------------\n\n\ndef train_step(\n state: utils.TrainState,\n batch,\n learning_rate_fn,\n label_smoothing=0.0,\n):\n \"\"\"Perform a single training step.\"\"\"\n # X_position and X_segmentation are needed only when using \"packed examples\"\n # where multiple sequences are packed into the same example with this\n # metadata.\n # if such features are not present they are ignored and the example is treated\n # like a normal, unpacked sequence example.\n train_keys = ['inputs', 'inputs_position', 'inputs_segmentation', 'targets']\n (inputs, inputs_positions, inputs_segmentation, targets) = (\n batch.get(k, None) for k in train_keys\n )\n\n # TODO: this should be defined globally\n pad_id = 0\n weights = jnp.where(inputs > pad_id, 1, 0).astype(jnp.float32)\n input_mask = inputs > pad_id\n attention_mask = transformer_lib.make_causal_attn_mask(\n input_mask\n ) # (B, L, L)\n # inputs_segmentation: (B, L)\n mask = (\n inputs_segmentation[:, :, None] == inputs_segmentation[:, None, :]\n ) # (B, L, L)\n attention_mask = jnp.logical_and(mask, attention_mask)\n\n def loss_fn(params):\n \"\"\"loss function used for training.\"\"\"\n module = nnx.merge(state.graphdef, params)\n\n logits, _ = module(\n inputs,\n positions=inputs_positions,\n attention_mask=attention_mask,\n cache=None,\n )\n\n loss, weight_sum = compute_weighted_cross_entropy(\n logits, targets, weights, label_smoothing\n )\n mean_loss = loss / weight_sum\n return mean_loss, logits\n\n step = state.step\n lr = learning_rate_fn(step)\n grad_fn = jax.value_and_grad(loss_fn, has_aux=True)\n (_, logits), grads = grad_fn(state.params)\n new_state = state.apply_gradients(grads=grads)\n metrics = compute_metrics(logits, targets, weights)\n metrics['learning_rate'] = lr\n\n return new_state, metrics\n\n\ndef eval_step(\n params: nnx.State,\n batch,\n graphdef: nnx.GraphDef[transformer_lib.Transformer],\n label_smoothing=0.0,\n):\n \"\"\"Calculate evaluation metrics on a batch.\"\"\"\n inputs, targets = batch['inputs'], batch['targets']\n\n # TODO: this should be defined globally\n pad_id = 0\n weights = jnp.where(inputs > pad_id, 1, 0).astype(jnp.float32)\n input_mask = inputs > pad_id\n inputs_positions = transformer_lib.build_positions_from_mask(input_mask)\n attention_mask = transformer_lib.make_causal_attn_mask(input_mask)\n\n module = nnx.merge(graphdef, params)\n logits, _ = module(\n inputs,\n positions=inputs_positions,\n attention_mask=attention_mask,\n cache=None,\n )\n\n return compute_metrics(logits, targets, weights, label_smoothing)\n\n\ndef evaluate(\n *,\n jit_eval_step,\n state: utils.TrainState,\n eval_ds: tf.data.Dataset,\n num_eval_steps: int,\n):\n \"\"\"Evaluate the target an return a dictionary with the metrics.\"\"\"\n logging.info('Gathering evaluation metrics.')\n eval_metrics = []\n eval_iter = iter(eval_ds) # pytype: disable=wrong-arg-types\n for _, eval_batch in zip(range(num_eval_steps), eval_iter):\n eval_batch = jax.tree.map(lambda x: x._numpy(), eval_batch) # pylint: disable=protected-access\n metrics = jit_eval_step(state.params, eval_batch, state.graphdef)\n eval_metrics.append(metrics)\n eval_metrics = common_utils.stack_forest(eval_metrics)\n eval_metrics_sums = jax.tree.map(jnp.sum, eval_metrics)\n eval_denominator = eval_metrics_sums.pop('denominator')\n eval_summary = jax.tree.map(\n lambda x: x / eval_denominator, # pylint: disable=cell-var-from-loop\n eval_metrics_sums,\n )\n return eval_summary\n\n\ndef train_and_evaluate(config: TrainConfig, workdir: str):\n \"\"\"Runs a training and evaluation loop.\n\n Args:\n config: Configuration to use.\n workdir: Working directory for checkpoints and TF summaries. If this\n contains checkpoint training will be resumed from the latest checkpoint.\n \"\"\"\n workdir = os.path.abspath(workdir)\n tf.io.gfile.makedirs(workdir)\n\n vocab_path = config.vocab_path\n if vocab_path is None:\n vocab_path = os.path.join(workdir, 'sentencepiece_model')\n config.vocab_path = vocab_path\n tf.io.gfile.makedirs(os.path.split(vocab_path)[0])\n\n # Load Dataset\n # ---------------------------------------------------------------------------\n logging.info('Initializing dataset.')\n train_ds, eval_ds, encoder = input_pipeline.get_datasets(\n n_devices=jax.local_device_count(), config=config, vocab_path=vocab_path\n )\n\n train_iter = iter(train_ds)\n vocab_size = int(encoder.vocab_size())\n\n logging.info('Initializing model, optimizer, and step functions.')\n # Build Model and Optimizer\n # ---------------------------------------------------------------------------\n if config.transformer_name is not None:\n model_config = transformer_lib.TransformerConfig.from_version_name(\n config.transformer_name,\n num_embed=vocab_size,\n dtype=jnp.bfloat16 if config.use_bfloat16 else jnp.float32,\n axis_rules=config.axis_rules,\n )\n else:\n assert config.transformer_params is not None\n model_config = transformer_lib.TransformerConfig.from_dict(\n **config.transformer_params,\n num_embed=vocab_size,\n dtype=jnp.bfloat16 if config.use_bfloat16 else jnp.float32,\n axis_rules=config.axis_rules,\n )\n\n # Mesh definition\n devices_array = utils.create_device_mesh(config)\n mesh = jax.sharding.Mesh(devices_array, config.mesh_axes)\n\n start_step = 0\n rng = jax.random.PRNGKey(config.seed)\n rng, init_rng = jax.random.split(rng)\n _, inference_rng = random.split(rng)\n\n def constructor(config: transformer_lib.TransformerConfig, key: jax.Array):\n return transformer_lib.Transformer(config, rngs=nnx.Rngs(params=key))\n\n learning_rate_fn = create_learning_rate_schedule(\n learning_rate=config.learning_rate, warmup_steps=config.warmup_steps\n )\n\n optimizer = optax.adamw(\n learning_rate_fn,\n b1=0.9,\n b2=0.98,\n eps=1e-9,\n weight_decay=config.weight_decay,\n )\n\n state, state_sharding = utils.setup_initial_state(\n constructor, optimizer, model_config, init_rng, mesh\n )\n data_sharding = jax.NamedSharding(mesh, jax.P(config.data_sharding))\n\n if config.restore_checkpoints:\n # Restore unreplicated optimizer + model state from last checkpoint.\n state = checkpoints.restore_checkpoint(workdir, state)\n # Grab last step.\n start_step = int(state.step)\n\n writer = metric_writers.create_default_writer(\n workdir, just_logging=jax.process_index() > 0\n )\n if start_step == 0:\n writer.write_hparams(dataclasses.asdict(config))\n\n # compile multidevice versions of train/eval/predict step fn.\n jit_train_step = jax.jit(\n train_step,\n in_shardings=(\n state_sharding,\n data_sharding,\n ), # type: ignore\n out_shardings=(state_sharding, None), # type: ignore\n static_argnames=('learning_rate_fn', 'label_smoothing'),\n donate_argnums=0,\n )\n\n jit_eval_step = jax.jit(\n eval_step,\n in_shardings=(\n state_sharding.params,\n data_sharding,\n ), # type: ignore\n out_shardings=None, # type: ignore\n static_argnames=('graphdef', 'label_smoothing'),\n )\n\n vocab = tokenizer.load_sentencepiece_processor(vocab_path)\n sampler = sampler_lib.Sampler(\n transformer=nnx.merge(state.graphdef, state.params),\n vocab=vocab,\n cache_size=1024,\n )\n\n # Main Train Loop\n # ---------------------------------------------------------------------------\n\n # We init the first set of dropout PRNG keys, but update it afterwards inside\n # the main pmap'd training update for performance.\n logging.info('Starting training loop.')\n hooks = []\n report_progress = periodic_actions.ReportProgress(\n num_train_steps=config.num_train_steps, writer=writer\n )\n if jax.process_index() == 0:\n hooks += [\n report_progress,\n periodic_actions.Profile(logdir=workdir, num_profile_steps=5),\n ]\n train_metrics = []\n with metric_writers.ensure_flushes(writer):\n for step in range(start_step, config.num_train_steps):\n is_last_step = step == config.num_train_steps - 1\n\n # Shard data to devices and do a training step.\n with jax.profiler.StepTraceAnnotation('train', step_num=step):\n with report_progress.timed('data'):\n batch = next(train_iter)\n batch = jax.tree.map(\n lambda x: jnp.asarray(x, device=data_sharding), batch\n )\n\n with report_progress.timed('train_step'):\n state, metrics = jit_train_step(state, batch, learning_rate_fn, 0.0)\n train_metrics.append(metrics)\n\n # Quick indication that training is happening.\n logging.log_first_n(logging.INFO, 'Finished training step %d.', 5, step)\n for h in hooks:\n h(step)\n\n # Write batch loss and lr every step to TB\n # without overwhelming the stdout:\n if jax.process_index() == 0:\n tb_writer = writer._writers[-1] # pylint: disable=protected-access\n lr = train_metrics[-1]['learning_rate']\n train_batch_loss = train_metrics[-1]['loss']\n denominator = train_metrics[-1]['denominator']\n tb_writer.write_scalars(\n step,\n {\n 'train_learning_rate': lr,\n 'train_loss': train_batch_loss / denominator,\n },\n )\n\n # Periodic metric handling.\n if (step > 0 and step % config.eval_every_steps == 0) or is_last_step:\n with report_progress.timed('training_metrics'):\n logging.info('Gathering training metrics.')\n train_metrics = common_utils.stack_forest(train_metrics)\n # Remove learning_rate from the summary\n _ = train_metrics.pop('learning_rate')\n metrics_sums = jax.tree.map(jnp.sum, train_metrics)\n denominator = metrics_sums.pop('denominator')\n summary = jax.tree.map(lambda x: x / denominator, metrics_sums) # pylint: disable=cell-var-from-loop\n summary['perplexity'] = jnp.clip(jnp.exp(summary['loss']), max=1.0e4)\n summary = {'train_' + k: v for k, v in summary.items()}\n writer.write_scalars(step, summary)\n train_metrics = []\n\n with report_progress.timed('generate_text'):\n # update sampler's transformer state:\n sampler.transformer_state = state.params\n exemplars = sampler(\n config.prompts,\n total_generation_steps=config.num_predict_steps,\n temperature=config.sampling_temperature,\n top_p=config.sampling_top_p,\n seed=inference_rng,\n echo=True,\n )\n writer.write_texts(step, {'samples': exemplars.text[0]})\n\n with report_progress.timed('eval'):\n eval_results = evaluate(\n jit_eval_step=jit_eval_step,\n state=state,\n eval_ds=eval_ds,\n num_eval_steps=config.num_eval_steps,\n )\n # (clipped) perplexity after averaging log-perplexity\n eval_results['perplexity'] = jnp.clip(\n jnp.exp(eval_results['loss']), max=1.0e4\n )\n writer.write_scalars(\n step, {'eval_' + k: v for k, v in eval_results.items()}\n )\n\n # Save a checkpoint on one host after every checkpoint_freq steps.\n save_checkpoint = (\n step % config.checkpoint_every_steps == 0 or is_last_step\n )\n if config.save_checkpoints and save_checkpoint:\n logging.info('Saving checkpoint step %d.', step)\n with report_progress.timed('checkpoint'):\n checkpoints.save_checkpoint_multiprocess(workdir, state, step)\n" }, { "path": "examples/gemma/transformer.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ============================================================================\n\"\"\"Gemma transformer.\"\"\"\n\nfrom __future__ import annotations\n\nfrom collections.abc import Iterable\nimport dataclasses\nimport enum\nimport functools\nfrom typing import Any\n\nfrom flax import nnx\nimport helpers\nimport layers\nimport modules\nimport params as params_lib\nimport sow_lib\nimport jax.numpy as jnp\nfrom jaxtyping import Array # pylint: disable=g-importing-member,g-multiple-import\n\nCache = dict[str, modules.LayerCache]\n\n\ndef make_attention_layers_types(\n pattern: tuple[modules.AttentionType, ...],\n num_layers: int,\n) -> tuple[modules.AttentionType, ...]:\n \"\"\"Returns the list of attention types for every layers.\"\"\"\n\n pattern_size = len(pattern)\n out = pattern * (num_layers // pattern_size)\n if num_layers % pattern_size != 0:\n out += pattern[: num_layers % pattern_size]\n return tuple(out)\n\n\nclass QueryPreAttentionNormalisation(enum.Enum):\n \"\"\"Initialization strategy.\"\"\"\n\n # Whether to scale the query by 1/sqrt(head_dim)\n BY_ONE_OVER_SQRT_HEAD_DIM = enum.auto()\n\n # Whether to scale the query by `embed_dim // num_heads`\n BY_EMBED_DIM_DIV_NUM_HEADS = enum.auto()\n\n # Whether to scale the query by `1/sqrt(embed_dim // num_heads)`\n BY_ONE_OVER_SQRT_EMBED_DIM_DIV_NUM_HEADS = enum.auto()\n\n\n_NUM_LAYERS_GEMMA_2B = 18\n_NUM_LAYERS_GEMMA_7B = 28\n_NUM_LAYERS_GEMMA2_2B = 26\n_NUM_LAYERS_GEMMA2_9B = 42\n_NUM_LAYERS_GEMMA2_27B = 46\n_NUM_LAYERS_GEMMA3_1B = 26\n_NUM_LAYERS_GEMMA3_4B = 34\n_NUM_LAYERS_GEMMA3_12B = 48\n_NUM_LAYERS_GEMMA3_27B = 62\nGEMMA3_ATTENTION_PATTERN = (\n modules.AttentionType.LOCAL_SLIDING,\n modules.AttentionType.LOCAL_SLIDING,\n modules.AttentionType.LOCAL_SLIDING,\n modules.AttentionType.LOCAL_SLIDING,\n modules.AttentionType.LOCAL_SLIDING,\n modules.AttentionType.GLOBAL,\n)\n\n\n@dataclasses.dataclass(frozen=True)\nclass TransformerConfig:\n \"\"\"Configuration for the gemma transformer.\"\"\"\n\n num_layers: int\n num_embed: int\n embed_dim: int\n hidden_dim: int\n num_heads: int\n head_dim: int\n num_kv_heads: int\n final_logit_softcap: float | None\n use_post_attn_norm: bool\n use_post_ffw_norm: bool\n attention_types: Iterable[modules.AttentionType]\n query_pre_attn_norm: QueryPreAttentionNormalisation = (\n QueryPreAttentionNormalisation.BY_ONE_OVER_SQRT_HEAD_DIM\n )\n attn_logits_soft_cap: float | None = None\n transpose_gating_einsum: bool = False\n local_base_frequency: int = modules.DEFAULT_ROPE_BASE_FREQUENCY\n global_base_frequency: int = modules.DEFAULT_ROPE_BASE_FREQUENCY\n local_scale_factor: float = modules.DEFAULT_ROPE_SCALE_FACTOR\n global_scale_factor: float = modules.DEFAULT_ROPE_SCALE_FACTOR\n use_qk_norm: bool = False\n sliding_window_size: int | None = None\n dtype: Any = jnp.float32\n axis_rules: Any | None = None\n\n def query_pre_attn_scalar(self) -> float:\n \"\"\"Returns the scalar to multiply the query by before attention.\"\"\"\n match self.query_pre_attn_norm:\n case QueryPreAttentionNormalisation.BY_EMBED_DIM_DIV_NUM_HEADS:\n return self.embed_dim // self.num_heads\n case QueryPreAttentionNormalisation.BY_ONE_OVER_SQRT_EMBED_DIM_DIV_NUM_HEADS: # pylint: disable=line-too-long\n return (self.embed_dim // self.num_heads) ** -0.5\n case QueryPreAttentionNormalisation.BY_ONE_OVER_SQRT_HEAD_DIM | _:\n return self.head_dim**-0.5\n\n @classmethod\n def from_path(cls, path: str) -> TransformerConfig:\n \"\"\"Creates a TransformerConfig from loaded parameters.\"\"\"\n params = params_lib.load_params(path)\n\n return cls.from_params(params)\n\n @classmethod\n def from_params(cls, params: params_lib.Params) -> TransformerConfig:\n \"\"\"Creates a TransformerConfig from loaded parameters.\n\n Args:\n params: Model parameters\n\n Returns:\n TransformerConfig.\n \"\"\"\n\n # Post Attn Norm is only used starting from Gemma 2.\n use_post_attn_norm = (\n 'post_attention_norm' in params['transformer']['layer_0']\n )\n\n # QK Norm is only used starting from Gemma 3.\n use_qk_norm = '_query_norm' in params['transformer']['layer_0']['attn']\n\n # Num layers will give use the model size.\n layer_names = [\n name for name in params['transformer'].keys() if 'layer' in name\n ]\n layer_names = [name.replace('layer_', '') for name in layer_names]\n num_layers = max([int(layer) for layer in layer_names]) + 1\n\n if not use_post_attn_norm: # Gemma 1.\n if num_layers == _NUM_LAYERS_GEMMA_2B:\n return cls.gemma_2b()\n if num_layers == _NUM_LAYERS_GEMMA_7B:\n return cls.gemma_7b()\n raise ValueError(\n 'Guessing Gemma 1 model, but could not determine size from params.'\n )\n elif not use_qk_norm: # Gemma 2.\n if num_layers == _NUM_LAYERS_GEMMA2_2B:\n return cls.gemma2_2b()\n if num_layers == _NUM_LAYERS_GEMMA2_9B:\n return cls.gemma2_9b()\n if num_layers == _NUM_LAYERS_GEMMA2_27B:\n return cls.gemma2_27b()\n raise ValueError(\n 'Guessing Gemma 2 model but could not determine size from params.'\n )\n else: # Gemma 3.\n if num_layers == _NUM_LAYERS_GEMMA3_1B:\n return cls.gemma3_1b()\n if num_layers == _NUM_LAYERS_GEMMA3_4B:\n return cls.gemma3_4b()\n if num_layers == _NUM_LAYERS_GEMMA3_12B:\n return cls.gemma3_12b()\n if num_layers == _NUM_LAYERS_GEMMA3_27B:\n return cls.gemma3_27b()\n\n raise ValueError('Could not determine Gemma variant from params.')\n\n @classmethod\n def from_version_name(cls, name: str, **override) -> TransformerConfig:\n possible_names = (\n \"gemma_2b\", \"gemma_7b\",\n \"gemma2_2b\", \"gemma2_9b\", \"gemma2_27b\",\n \"gemma3_1b\", \"gemma3_4b\", \"gemma3_12b\", \"gemma3_27b\",\n )\n if name not in possible_names:\n raise ValueError(\n f'Unknown version name: {name}. '\n f'Please choose one of the following: {possible_names}'\n )\n if hasattr(cls, name):\n model_config = getattr(cls, name)(**override)\n return model_config\n else:\n raise RuntimeError(\n 'Something wrong in TransformerConfig code. '\n f'No attribute {name} in TransformerConfig'\n )\n\n @classmethod\n def from_dict(cls, **config: Any) -> TransformerConfig:\n # Deserialize query_pre_attn_norm values:\n if \"query_pre_attn_norm\" in config:\n config[\"query_pre_attn_norm\"] = QueryPreAttentionNormalisation(config[\"query_pre_attn_norm\"])\n else:\n config[\"query_pre_attn_norm\"] = QueryPreAttentionNormalisation.BY_ONE_OVER_SQRT_HEAD_DIM\n return cls(**config)\n\n @classmethod\n def gemma_2b(cls, **override) -> TransformerConfig:\n num_layers = _NUM_LAYERS_GEMMA_2B\n config = {\n 'num_layers': num_layers,\n 'num_embed': 256128,\n 'embed_dim': 2048,\n 'hidden_dim': 16384,\n 'num_heads': 8,\n 'head_dim': 256,\n 'num_kv_heads': 1,\n 'final_logit_softcap': None,\n 'attention_types': (modules.AttentionType.GLOBAL,) * num_layers,\n 'use_post_attn_norm': False,\n 'use_post_ffw_norm': False,\n }\n for key, value in override.items():\n config[key] = value\n return cls(**config)\n\n @classmethod\n def gemma_7b(cls, **override):\n num_layers = _NUM_LAYERS_GEMMA_7B\n config = {\n \"num_layers\": num_layers,\n \"num_embed\": 256128,\n \"embed_dim\": 3072,\n \"hidden_dim\": 24576,\n \"num_heads\": 16,\n \"head_dim\": 256,\n \"num_kv_heads\": 16,\n \"final_logit_softcap\": None,\n \"attention_types\": (modules.AttentionType.GLOBAL,) * num_layers,\n \"use_post_attn_norm\": False,\n \"use_post_ffw_norm\": False,\n }\n for key, value in override.items():\n config[key] = value\n return cls(**config)\n\n @classmethod\n def gemma2_2b(cls, **override):\n num_layers = _NUM_LAYERS_GEMMA2_2B\n config = {\n 'num_layers': num_layers,\n 'num_embed': 256128,\n 'embed_dim': 2304,\n 'hidden_dim': 9216,\n 'num_heads': 8,\n 'head_dim': 256,\n 'num_kv_heads': 4,\n 'final_logit_softcap': 30.0,\n 'attention_types': (\n modules.AttentionType.LOCAL_SLIDING,\n modules.AttentionType.GLOBAL,\n )\n * int(num_layers / 2),\n 'use_post_attn_norm': True,\n 'use_post_ffw_norm': True,\n 'query_pre_attn_norm': QueryPreAttentionNormalisation.BY_ONE_OVER_SQRT_HEAD_DIM,\n 'attn_logits_soft_cap': 50.0,\n 'sliding_window_size': 4096,\n }\n for key, value in override.items():\n config[key] = value\n return cls(**config)\n\n @classmethod\n def gemma2_9b(cls, **override):\n num_layers = _NUM_LAYERS_GEMMA2_9B\n config = {\n \"num_layers\": num_layers,\n \"num_embed\": 256128,\n \"embed_dim\": 3584,\n \"hidden_dim\": 28672,\n \"num_heads\": 16,\n \"head_dim\": 256,\n \"num_kv_heads\": 8,\n \"final_logit_softcap\": 30.0,\n \"attention_types\": (\n modules.AttentionType.LOCAL_SLIDING,\n modules.AttentionType.GLOBAL,\n ) * int(num_layers / 2),\n \"use_post_attn_norm\": True,\n \"use_post_ffw_norm\": True,\n \"attn_logits_soft_cap\": 50.0,\n \"sliding_window_size\": 4096,\n }\n for key, value in override.items():\n config[key] = value\n return cls(**config)\n\n @classmethod\n def gemma2_27b(cls, **override):\n num_layers = _NUM_LAYERS_GEMMA2_27B\n config = {\n \"num_layers\": num_layers,\n \"num_embed\": 256128,\n \"embed_dim\": 4608,\n \"hidden_dim\": 72728,\n \"num_heads\": 32,\n \"head_dim\": 128,\n \"num_kv_heads\": 16,\n \"final_logit_softcap\": 30.0,\n \"use_post_attn_norm\": True,\n \"use_post_ffw_norm\": True,\n \"attention_types\": (\n modules.AttentionType.LOCAL_SLIDING,\n modules.AttentionType.GLOBAL,\n ) * int(num_layers / 2),\n \"attn_logits_soft_cap\": 50.0,\n \"sliding_window_size\": 4096,\n }\n for key, value in override.items():\n config[key] = value\n return cls(**config)\n\n @classmethod\n def gemma3_1b(cls, **override):\n num_layers = _NUM_LAYERS_GEMMA3_1B\n config = {\n \"num_layers\": num_layers,\n \"final_logit_softcap\": None,\n \"num_embed\": 262144,\n \"embed_dim\": 1152,\n \"hidden_dim\": 6 * 1152,\n \"num_heads\": 4,\n \"head_dim\": 256,\n \"num_kv_heads\": 1,\n \"use_post_attn_norm\": True,\n \"use_post_ffw_norm\": True,\n \"use_qk_norm\": True,\n \"attention_types\": make_attention_layers_types(\n GEMMA3_ATTENTION_PATTERN, num_layers\n ),\n \"query_pre_attn_norm\": QueryPreAttentionNormalisation.BY_ONE_OVER_SQRT_HEAD_DIM,\n \"attn_logits_soft_cap\": None,\n \"sliding_window_size\": 512,\n \"transpose_gating_einsum\": True,\n \"local_base_frequency\": 10_000,\n \"global_base_frequency\": 1_000_000,\n }\n for key, value in override.items():\n config[key] = value\n return cls(**config)\n\n @classmethod\n def gemma3_4b(cls, **override):\n num_layers = _NUM_LAYERS_GEMMA3_4B\n config = {\n \"num_layers\": num_layers,\n \"final_logit_softcap\": None,\n \"num_embed\": 262_144,\n \"embed_dim\": 2560,\n \"hidden_dim\": 2560 * 8 // 2,\n \"num_heads\": 8,\n \"head_dim\": 256,\n \"num_kv_heads\": 4,\n \"use_post_attn_norm\": True,\n \"use_post_ffw_norm\": True,\n \"use_qk_norm\": True,\n \"attention_types\": make_attention_layers_types(\n GEMMA3_ATTENTION_PATTERN, num_layers\n ),\n \"query_pre_attn_norm\": QueryPreAttentionNormalisation.BY_ONE_OVER_SQRT_HEAD_DIM,\n \"attn_logits_soft_cap\": None,\n \"sliding_window_size\": 1024,\n \"transpose_gating_einsum\": True,\n \"local_base_frequency\": 10_000,\n \"global_base_frequency\": 1_000_000,\n \"global_scale_factor\": 8.0,\n }\n for key, value in override.items():\n config[key] = value\n return cls(**config)\n\n @classmethod\n def gemma3_12b(cls, **override):\n num_layers = _NUM_LAYERS_GEMMA3_12B\n config = {\n \"num_layers\": num_layers,\n \"final_logit_softcap\": None,\n \"num_embed\": 262144,\n \"embed_dim\": 30 * 128,\n \"hidden_dim\": 8 * 30 * 128 // 2,\n \"num_heads\": 16,\n \"head_dim\": 256,\n \"num_kv_heads\": 8,\n \"use_post_attn_norm\": True,\n \"use_post_ffw_norm\": True,\n \"use_qk_norm\": True,\n \"attention_types\": make_attention_layers_types(\n GEMMA3_ATTENTION_PATTERN, num_layers\n ),\n \"query_pre_attn_norm\": QueryPreAttentionNormalisation.BY_ONE_OVER_SQRT_HEAD_DIM,\n \"attn_logits_soft_cap\": None,\n \"sliding_window_size\": 1024,\n \"transpose_gating_einsum\": True,\n \"local_base_frequency\": 10_000,\n \"global_base_frequency\": 1_000_000,\n \"global_scale_factor\": 8.0,\n }\n for key, value in override.items():\n config[key] = value\n return cls(**config)\n\n @classmethod\n def gemma3_27b(cls, **override):\n num_layers = _NUM_LAYERS_GEMMA3_27B\n config = {\n \"num_layers\": num_layers,\n \"final_logit_softcap\": None,\n \"num_embed\": 262144,\n \"embed_dim\": 5376,\n \"hidden_dim\": 5376 * 8 // 2,\n \"num_heads\": 32,\n \"head_dim\": 128,\n \"num_kv_heads\": 16,\n \"use_post_attn_norm\": True,\n \"use_post_ffw_norm\": True,\n \"use_qk_norm\": True,\n \"attention_types\": make_attention_layers_types(\n GEMMA3_ATTENTION_PATTERN, num_layers\n ),\n \"query_pre_attn_norm\": QueryPreAttentionNormalisation.BY_ONE_OVER_SQRT_EMBED_DIM_DIV_NUM_HEADS,\n \"attn_logits_soft_cap\": None,\n \"sliding_window_size\": 1024,\n \"transpose_gating_einsum\": True,\n \"local_base_frequency\": 10_000,\n \"global_base_frequency\": 1_000_000,\n \"global_scale_factor\": 8.0,\n }\n for key, value in override.items():\n config[key] = value\n return cls(**config)\n\n def __post_init__(self):\n if self.num_heads != self.num_kv_heads:\n if self.num_heads % self.num_kv_heads != 0:\n raise ValueError(\n f\"Number of query heads ({self.num_heads}) must be divisible by \"\n f\"number of key/value heads ({self.num_kv_heads}).\"\n )\n\n\ndef _map_linen_var_names(key: tuple[str, ...]) -> tuple[str | int, ...]:\n \"\"\"Maps linen variable names to nnx variable names.\"\"\"\n new_key = []\n for k in key:\n if k.startswith('layer_'):\n prefix, suffix = k.split('layer_')\n assert not prefix, prefix\n new_key.append('layers')\n new_key.append(int(suffix))\n elif k == 'gating_einsum':\n new_key.append('gate_proj')\n new_key.append('kernel')\n elif k == 'linear':\n new_key.append('down_proj')\n new_key.append('kernel')\n else:\n new_key.append(k)\n\n return tuple(new_key)\n\n\ndef _assign_linen_params_to_nnx_state(\n state: dict[tuple[str, ...], Any],\n mapped_path: tuple[str | int, ...],\n val: Any,\n transpose_gating_einsum: bool,\n) -> dict[tuple[str, ...], Any]:\n \"\"\"Splits and maybe transposes gate_proj.\"\"\"\n if 'gate_proj' in mapped_path:\n if transpose_gating_einsum:\n val = jnp.swapaxes(val, 1, 2)\n state[mapped_path].set_value(val[0])\n state[mapped_path[:-2] + ('up_proj', 'kernel')].set_value(val[1])\n else:\n state[mapped_path].set_value(val)\n return state\n\n\nclass Transformer(nnx.Module):\n \"\"\"Gemma transformer.\"\"\"\n\n @classmethod\n def from_params(\n cls,\n params: params_lib.Params,\n config: None | TransformerConfig = None,\n sow_config: sow_lib.SowConfig = sow_lib.SowConfig(),\n ) -> Transformer:\n if config is None:\n config = TransformerConfig.from_params(params)\n assign_val_fn = functools.partial(\n _assign_linen_params_to_nnx_state,\n transpose_gating_einsum=config.transpose_gating_einsum,\n )\n return helpers.module_from_linen_variables(\n module_factory=lambda: cls(\n config, rngs=nnx.Rngs(params=0), sow_config=sow_config\n ),\n variables=params['transformer'],\n map_key_fn=_map_linen_var_names,\n assign_val_fn=assign_val_fn,\n )\n\n def __init__(\n self,\n config: TransformerConfig,\n *,\n rngs: nnx.Rngs,\n sow_config: sow_lib.SowConfig = sow_lib.SowConfig(),\n ):\n self.embedder = modules.Embedder(\n vocab_size=config.num_embed,\n embed_dim=config.embed_dim,\n embedding_init=modules.maybe_with_partitioning(\n nnx.initializers.normal(),\n config.axis_rules,\n (\"vocab\", \"embed\"),\n ),\n dtype=config.dtype,\n rngs=rngs,\n )\n self.layers = nnx.List([\n modules.Block(\n config=config,\n attn_type=attn_type,\n sow_config=sow_config,\n rngs=rngs,\n )\n for _, attn_type in zip(\n range(config.num_layers), config.attention_types\n )\n ])\n self.final_norm = layers.RMSNorm(\n config.embed_dim,\n scale_init=modules.maybe_with_partitioning(\n nnx.initializers.zeros_init(),\n config.axis_rules,\n (\"embed\", ),\n ),\n rngs=rngs,\n )\n self.final_logits_softcap = config.final_logit_softcap\n self.sow_config = sow_config\n\n def __call__(\n self,\n last_tokens: Array, # [B, L]\n positions: Array, # [B, L]\n cache: Cache | None, # (sequence length L')\n attention_mask: Array, # [B, L, L']\n ) -> tuple[Array, Cache | None]:\n \"\"\"Transformer forward pass.\n\n You can run this forward pass two ways: with or without an attention kv\n cache.\n\n Args:\n last_tokens: input sequence of tokens.\n positions: input absolute positions.\n cache: Attention KV cache or None.\n attention_mask: transformer input mask.\n\n Returns:\n predicted_logits, new_cache\n\n predicted_logits: output logits predicted by the model\n new_cache: updated cache if the input cache is not None, None elsewhere.\n \"\"\"\n new_cache = None if cache is None else {}\n x = self.embedder.encode(last_tokens)\n self.sow_config.maybe_sow_embeddings(x, self)\n for i, layer in enumerate(self.layers):\n layer_name = f'layer_{i}'\n layer_cache = cache[layer_name] if cache else None\n layer_cache, x = layer(\n x,\n positions,\n layer_cache,\n attention_mask,\n )\n if cache is not None:\n new_cache[layer_name] = layer_cache # pytype: disable=container-type-mismatch\n\n x = self.final_norm(x)\n logits = self.embedder.decode(x)\n\n if self.final_logits_softcap is not None:\n logits /= self.final_logits_softcap\n logits = jnp.tanh(logits) * self.final_logits_softcap\n\n return logits, new_cache # pytype: disable=bad-return-type\n\n @property\n def embed_dim(self) -> int:\n return self.embedder.embed_dim\n\n @property\n def num_embed(self) -> int:\n return self.embedder.num_embed\n\n @property\n def num_layers(self) -> int:\n return len(self.layers)\n\n def init_cache(\n self,\n cache_size: int,\n batch_size: int,\n dtype: jnp.dtype = jnp.bfloat16,\n ) -> Cache:\n \"\"\"Initializes a new Transformer cache.\"\"\"\n return {\n f'layer_{i}': self.layers[i].init_cache(\n cache_size=cache_size,\n batch_size=batch_size,\n dtype=dtype,\n )\n for i in range(self.num_layers)\n }\n\n def init_intermediates(\n self,\n batch_size: int,\n buffer_size: int,\n sow_config: sow_lib.SowConfig,\n dtype: jnp.dtype = jnp.float32,\n ) -> sow_lib.TransformerIntermediates:\n \"\"\"Initializes the intermediate activations that will be filled.\"\"\"\n intermediates = sow_lib.TransformerIntermediates()\n residual_stream_dummy = jnp.zeros(\n (batch_size, buffer_size, self.embed_dim),\n dtype=dtype,\n )\n if sow_config.embeddings:\n intermediates.embeddings = residual_stream_dummy\n for layer in self.layers:\n layer_intermediates = sow_lib.LayerIntermediates()\n if sow_config.rs_after_attention:\n layer_intermediates.rs_after_attention = residual_stream_dummy\n if sow_config.rs_after_ffw:\n layer_intermediates.rs_after_ffw = residual_stream_dummy\n if sow_config.attn_logits_topk:\n shape = (\n batch_size,\n buffer_size,\n layer.attn.num_heads,\n sow_config.attn_logits_topk,\n )\n layer_intermediates.attn_logits_topk_values = jnp.zeros(\n shape,\n dtype=dtype,\n )\n layer_intermediates.attn_logits_topk_indices = jnp.zeros(\n shape,\n dtype=jnp.int32,\n )\n if sow_config.mlp_hidden_topk:\n shape = (\n batch_size,\n buffer_size,\n sow_config.mlp_hidden_topk,\n )\n layer_intermediates.mlp_hidden_topk_values = jnp.zeros(\n shape,\n dtype=dtype,\n )\n layer_intermediates.mlp_hidden_topk_indices = jnp.zeros(\n shape,\n dtype=jnp.int32,\n )\n intermediates.layers.append(layer_intermediates)\n return intermediates\n\n\ndef make_causal_attn_mask(\n input_mask: Array,\n) -> Array:\n \"\"\"Attention mask in batch mode.\n\n Args:\n input_mask: Input mask for the input. True for non-padded tokens only, else\n False.\n\n Returns:\n Attention mask.\n \"\"\"\n seq_len = input_mask.shape[-1]\n attn_mask = input_mask[..., None, :]\n causal_mask = jnp.tril(jnp.ones((seq_len, seq_len), dtype=jnp.bool_))\n # Prefixes can be attended by all tokens\n attn_mask *= causal_mask[None, ...]\n return attn_mask\n\n\ndef build_positions_from_mask(input_mask: Array) -> Array:\n \"\"\"Computes the `positions` from the `input_mask`.\n\n Args:\n input_mask: The tokens `input_mask`, True for non-padded tokens only.\n\n Returns:\n The indices to use for RoPE and absolute position encodings for the given\n input mask.\n \"\"\"\n positions = jnp.cumsum(input_mask, axis=-1)\n # Subtract one for all positions from the first valid one as they are\n # 0-indexed\n return positions - (positions >= 1)\n" }, { "path": "examples/gemma/transformer_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ============================================================================\n\"\"\"Tests for the Gemma transformer.\"\"\"\n\nfrom collections import defaultdict\nfrom absl.testing import absltest\nfrom absl.testing import parameterized\nfrom flax import nnx\nimport modules\nimport sow_lib\nimport transformer as transformer_lib\nimport jax.numpy as jnp\nimport numpy as np\n\n\ndef create_fake_params(config: transformer_lib.TransformerConfig):\n def nested_defaultdict():\n return defaultdict(nested_defaultdict)\n\n res = nested_defaultdict()\n res['transformer'] = nested_defaultdict()\n params = res['transformer']\n # 1. embedding params\n params['embedder']['input_embedding'] = jnp.ones(\n (config.num_embed, config.embed_dim)\n )\n # 2. final norm params\n params['final_norm'] = {'scale': jnp.ones((config.embed_dim,))}\n\n # 3. attention block params\n for layer_idx in range(config.num_layers):\n params[f'layer_{layer_idx}']['attn']['attn_vec_einsum']['w'] = jnp.ones(\n (config.num_heads, config.head_dim, config.embed_dim)\n )\n if config.num_heads == config.num_kv_heads:\n params[f'layer_{layer_idx}']['attn']['qkv_einsum']['w'] = jnp.ones(\n (3, config.num_heads, config.embed_dim, config.head_dim)\n )\n else:\n params[f'layer_{layer_idx}']['attn']['q_einsum']['w'] = jnp.ones(\n (config.num_heads, config.embed_dim, config.head_dim)\n )\n params[f'layer_{layer_idx}']['attn']['kv_einsum']['w'] = jnp.ones(\n (2, config.num_kv_heads, config.embed_dim, config.head_dim)\n )\n\n # 4. feedforward block params\n params[f'layer_{layer_idx}']['mlp']['gating_einsum'] = jnp.ones(\n (2, config.embed_dim, config.hidden_dim)\n )\n params[f'layer_{layer_idx}']['mlp']['linear'] = jnp.ones(\n (config.hidden_dim, config.embed_dim)\n )\n\n # 5. layer norm params\n params[f'layer_{layer_idx}']['pre_attention_norm']['scale'] = jnp.ones((\n config.embed_dim,\n ))\n params[f'layer_{layer_idx}']['pre_ffw_norm']['scale'] = jnp.ones((\n config.embed_dim,\n ))\n\n if config.use_post_attn_norm:\n params[f'layer_{layer_idx}']['post_attention_norm']['scale'] = jnp.ones((\n config.embed_dim,\n ))\n if config.use_post_ffw_norm:\n params[f'layer_{layer_idx}']['post_ffw_norm']['scale'] = jnp.ones((\n config.embed_dim,\n ))\n return res\n\n\nclass TransformerTest(parameterized.TestCase):\n\n @parameterized.parameters(\n # Prime number to ease shape tracing\n dict(\n num_layers=3,\n num_embed=17,\n embed_dim=2,\n num_heads=2,\n num_kv_heads=2,\n hidden_dim=11,\n head_dim=8,\n cache_size=29,\n batch_size=7,\n sequence_length=18,\n expected_outputs_shape=(7, 18, 17), # batch_size, seq_size, num_embed\n expected_cache_shape=(7, 29, 2, 8), # batch_size, cache_size, num_kv_heads, head_dim\n ),\n dict(\n num_layers=3,\n num_embed=4,\n embed_dim=2,\n num_heads=2,\n num_kv_heads=1,\n hidden_dim=4,\n head_dim=4,\n cache_size=2,\n batch_size=1,\n sequence_length=1,\n expected_outputs_shape=(1, 1, 4), # batch_size, seq_size, num_embed\n expected_cache_shape=(1, 2, 1, 4), # batch_size, cache_size, num_kv_heads, head_dim\n ),\n dict(\n num_layers=3,\n num_embed=7,\n embed_dim=5,\n num_heads=4,\n num_kv_heads=2,\n hidden_dim=6,\n head_dim=8,\n cache_size=9,\n batch_size=1,\n sequence_length=1,\n expected_outputs_shape=(1, 1, 7), # batch_size, seq_size, num_embed\n expected_cache_shape=(1, 9, 2, 8), # batch_size, cache_size, num_kv_heads, head_dim\n ),\n )\n def test_transformer(\n self,\n num_layers,\n num_embed,\n embed_dim,\n num_heads,\n num_kv_heads,\n hidden_dim,\n head_dim,\n cache_size,\n batch_size,\n sequence_length,\n expected_outputs_shape,\n expected_cache_shape,\n ):\n\n config = transformer_lib.TransformerConfig(\n num_layers=num_layers,\n num_embed=num_embed,\n embed_dim=embed_dim,\n hidden_dim=hidden_dim,\n num_heads=num_heads,\n head_dim=head_dim,\n num_kv_heads=num_kv_heads,\n final_logit_softcap=None,\n attention_types=[modules.AttentionType.GLOBAL] * num_layers,\n use_post_attn_norm=False,\n use_post_ffw_norm=False,\n )\n attention_mask = jnp.ones((batch_size, 1, cache_size), dtype=jnp.bool)\n transformer = transformer_lib.Transformer(\n config=config, rngs=nnx.Rngs(params=0)\n )\n cache = transformer.init_cache(\n cache_size=cache_size,\n batch_size=batch_size,\n dtype=jnp.float32,\n )\n\n outputs, cache = transformer(\n jnp.tile(jnp.arange(sequence_length), (batch_size, 1)),\n jnp.tile(jnp.arange(sequence_length), (batch_size, 1)),\n cache,\n attention_mask,\n )\n\n self.assertEqual(outputs.shape, expected_outputs_shape)\n self.assertEqual(cache['layer_0']['v'].shape, expected_cache_shape)\n\n @parameterized.parameters(\n ('final_logit_softcap',),\n ('attn_logits_soft_cap',),\n )\n def test_logit_softcap(\n self,\n soft_cap_arg,\n ):\n cache_size = 2\n batch_size = 1\n soft_cap_val = 0.001\n\n attention_mask = jnp.ones((batch_size, 1, cache_size), dtype=jnp.bool)\n\n params = dict(\n num_layers=3,\n num_embed=4,\n embed_dim=2,\n num_heads=2,\n num_kv_heads=1,\n hidden_dim=4,\n head_dim=4,\n attention_types=[modules.AttentionType.GLOBAL] * 3,\n use_post_attn_norm=False,\n use_post_ffw_norm=False,\n )\n\n no_soft_cap_args = {\n 'final_logit_softcap': None,\n 'attn_logits_soft_cap': None,\n }\n\n soft_cap_args = no_soft_cap_args.copy()\n soft_cap_args[soft_cap_arg] = soft_cap_val\n\n config_soft_cap = transformer_lib.TransformerConfig(\n **(params | soft_cap_args)\n )\n config_no_soft_cap = transformer_lib.TransformerConfig(\n **(params | no_soft_cap_args)\n )\n\n all_outputs = []\n for config in [config_soft_cap, config_no_soft_cap]:\n transformer = transformer_lib.Transformer(\n config=config, rngs=nnx.Rngs(params=1)\n )\n cache = transformer.init_cache(\n cache_size=cache_size,\n batch_size=batch_size,\n dtype=jnp.float32,\n )\n\n outputs, _ = transformer(\n jnp.array([[1]]), jnp.array([[1]]), cache, attention_mask\n )\n all_outputs.append(outputs)\n\n soft_cap_outputs, no_soft_cap_outputs = all_outputs # pylint: disable=unbalanced-tuple-unpacking\n\n # Ensure that values aren't equal coming out of computation\n self.assertFalse((soft_cap_outputs == no_soft_cap_outputs).all())\n\n # Run soft capping manually\n manual_soft_cap_logits = no_soft_cap_outputs / soft_cap_val\n manual_soft_cap_logits = jnp.tanh(manual_soft_cap_logits) * soft_cap_val\n\n np.testing.assert_array_almost_equal(\n manual_soft_cap_logits, soft_cap_outputs, 1e-5\n )\n\n @parameterized.parameters([\n dict(\n config=transformer_lib.TransformerConfig(\n num_layers=2,\n num_embed=0, # unused\n embed_dim=0, # unused\n hidden_dim=0, # unused\n num_heads=3,\n head_dim=4,\n num_kv_heads=3,\n final_logit_softcap=None,\n attention_types=[modules.AttentionType.GLOBAL] * 2,\n use_post_attn_norm=False,\n use_post_ffw_norm=False,\n ),\n cache_size=2,\n keys=['layer_0', 'layer_1'],\n k_shape=(1, 2, 3, 4),\n v_shape=(1, 2, 3, 4),\n )\n ])\n def test_creates_cache(self, config, cache_size, keys, k_shape, v_shape):\n transformer = transformer_lib.Transformer(\n config=config, rngs=nnx.Rngs(params=0)\n )\n cache = transformer.init_cache(\n cache_size=cache_size,\n batch_size=1,\n dtype=jnp.float32,\n )\n self.assertEqual(list(cache.keys()), keys)\n self.assertEqual(cache['layer_0']['k'].shape, k_shape)\n self.assertEqual(cache['layer_0']['v'].shape, v_shape)\n\n @parameterized.parameters([\n dict(\n batch_size=1,\n seq_size=4,\n config=transformer_lib.TransformerConfig(\n num_layers=2,\n num_embed=4, # unused\n embed_dim=2,\n hidden_dim=12, # unused\n num_heads=3,\n head_dim=4,\n num_kv_heads=3,\n final_logit_softcap=None,\n attention_types=[modules.AttentionType.GLOBAL] * 2,\n use_post_attn_norm=False,\n use_post_ffw_norm=False,\n ),\n )\n ])\n def test_forward_no_cache(\n self,\n batch_size: int,\n seq_size: int,\n config: transformer_lib.TransformerConfig,\n ):\n cache_size = 6\n\n token_input = jnp.ones((batch_size, seq_size), dtype=jnp.int32)\n transformer = transformer_lib.Transformer(\n config=config, rngs=nnx.Rngs(params=0)\n )\n empty_cache = transformer.init_cache(\n cache_size=cache_size,\n batch_size=batch_size,\n dtype=jnp.float32,\n )\n attention_mask = jnp.ones(\n (batch_size, seq_size, cache_size), dtype=jnp.bool\n )\n positions = transformer_lib.build_positions_from_mask(token_input != 0)\n\n output_cache, _ = transformer(\n token_input, positions, empty_cache, attention_mask\n )\n\n attention_mask = jnp.ones((batch_size, seq_size, seq_size), dtype=jnp.bool)\n output_none, cache_none = transformer(\n token_input, positions, None, attention_mask\n )\n\n self.assertIsNone(cache_none)\n np.testing.assert_array_almost_equal(output_cache, output_none, 1e-5)\n\n def test_attention_types(\n self,\n ):\n\n config = transformer_lib.TransformerConfig(\n num_layers=2,\n num_embed=4,\n embed_dim=2,\n hidden_dim=12,\n num_heads=3,\n head_dim=4,\n num_kv_heads=3,\n final_logit_softcap=None,\n attention_types=[modules.AttentionType.GLOBAL] * 2,\n use_post_attn_norm=False,\n use_post_ffw_norm=False,\n )\n transformer = transformer_lib.Transformer(\n config=config, rngs=nnx.Rngs(params=0)\n )\n cache = transformer.init_cache(\n cache_size=6,\n batch_size=1,\n dtype=jnp.float32,\n )\n self.assertTrue(cache)\n\n @parameterized.parameters(\n dict(\n config=transformer_lib.TransformerConfig(\n num_layers=2,\n num_embed=4,\n embed_dim=2,\n hidden_dim=12,\n num_heads=3,\n head_dim=4,\n num_kv_heads=3,\n final_logit_softcap=None,\n attention_types=[modules.AttentionType.GLOBAL] * 2,\n use_post_attn_norm=False,\n use_post_ffw_norm=False,\n ),\n ),\n dict(\n config=transformer_lib.TransformerConfig(\n num_layers=2,\n num_embed=4,\n embed_dim=2,\n hidden_dim=12,\n num_heads=3,\n head_dim=4,\n num_kv_heads=3,\n final_logit_softcap=None,\n attention_types=[modules.AttentionType.GLOBAL] * 2,\n use_post_attn_norm=True,\n use_post_ffw_norm=True,\n ),\n ),\n dict(\n config=transformer_lib.TransformerConfig(\n num_layers=2,\n num_embed=4,\n embed_dim=5,\n hidden_dim=12,\n num_heads=6,\n head_dim=4,\n num_kv_heads=3,\n final_logit_softcap=None,\n attention_types=[modules.AttentionType.GLOBAL] * 2,\n use_post_attn_norm=True,\n use_post_ffw_norm=True,\n ),\n ),\n )\n def test_load_from_params(self, config):\n params = create_fake_params(config)\n transformer = transformer_lib.Transformer.from_params(params, config)\n logits, _ = transformer(\n last_tokens=jnp.tile(jnp.arange(3), (2, 1)),\n positions=jnp.tile(jnp.arange(3), (2, 1)),\n cache=None,\n attention_mask=jnp.ones((2, 1, 3), dtype=jnp.bool),\n )\n self.assertEqual(logits.shape, (2, 3, 4))\n\n @parameterized.parameters([\n sow_lib.SowConfig(embeddings=True),\n sow_lib.SowConfig(rs_after_attention=True),\n sow_lib.SowConfig(rs_after_ffw=True),\n sow_lib.SowConfig(attn_logits_topk=5),\n sow_lib.SowConfig(mlp_hidden_topk=11),\n ])\n def test_sow_intermediates(self, sow_config):\n batch_size = 3\n sequence_length = 7\n num_layers = 2\n config = transformer_lib.TransformerConfig(\n num_layers=num_layers,\n num_embed=4,\n embed_dim=48,\n hidden_dim=12,\n num_heads=1,\n head_dim=4,\n num_kv_heads=1,\n final_logit_softcap=None,\n use_post_attn_norm=False,\n use_post_ffw_norm=False,\n attention_types=[modules.AttentionType.GLOBAL] * num_layers,\n )\n attention_mask = jnp.ones(\n (batch_size, sequence_length, sequence_length), dtype=jnp.bool\n )\n transformer = transformer_lib.Transformer(\n config=config, rngs=nnx.Rngs(params=0), sow_config=sow_config\n )\n transformer(\n jnp.tile(jnp.arange(sequence_length), (batch_size, 1)),\n jnp.tile(jnp.arange(sequence_length), (batch_size, 1)),\n None,\n attention_mask,\n )\n\n if sow_config.embeddings:\n self.assertTrue(hasattr(transformer, 'embeddings'))\n embeddings = transformer.embeddings[0]\n self.assertEqual(\n embeddings.shape,\n (batch_size, sequence_length, config.embed_dim),\n )\n else:\n self.assertFalse(hasattr(transformer, 'embeddings'))\n\n for layer in transformer.layers:\n if sow_config.rs_after_attention:\n self.assertTrue(hasattr(layer, 'rs_after_attention'))\n rs_after_attention = layer.rs_after_attention[0]\n self.assertIsNotNone(rs_after_attention)\n self.assertEqual(\n rs_after_attention.shape,\n (batch_size, sequence_length, config.embed_dim),\n )\n else:\n self.assertFalse(hasattr(layer, 'rs_after_attention'))\n if sow_config.rs_after_ffw:\n self.assertTrue(hasattr(layer, 'rs_after_ffw'))\n rs_after_ffw = layer.rs_after_ffw[0]\n self.assertIsNotNone(rs_after_ffw)\n self.assertEqual(\n rs_after_ffw.shape,\n (batch_size, sequence_length, config.embed_dim),\n )\n else:\n self.assertFalse(hasattr(layer, 'rs_after_ffw'))\n if sow_config.attn_logits_topk:\n self.assertTrue(hasattr(layer.attn, 'logits_topk_values'))\n attn_logits_topk_values = layer.attn.logits_topk_values[0]\n self.assertIsNotNone(attn_logits_topk_values)\n self.assertEqual(\n attn_logits_topk_values.shape,\n (\n batch_size,\n sequence_length,\n config.num_heads,\n sow_config.attn_logits_topk,\n ),\n )\n self.assertTrue(hasattr(layer.attn, 'logits_topk_indices'))\n attn_logits_topk_indices = layer.attn.logits_topk_indices[0]\n self.assertIsNotNone(attn_logits_topk_indices)\n self.assertEqual(\n attn_logits_topk_indices.shape,\n (\n batch_size,\n sequence_length,\n config.num_heads,\n sow_config.attn_logits_topk,\n ),\n )\n else:\n self.assertFalse(hasattr(layer.attn, 'logits_topk_values'))\n self.assertFalse(hasattr(layer.attn, 'logits_topk_indices'))\n if sow_config.mlp_hidden_topk:\n self.assertTrue(hasattr(layer.mlp, 'hidden_topk_values'))\n ffw_hidden_topk_values = layer.mlp.hidden_topk_values[0]\n self.assertIsNotNone(ffw_hidden_topk_values)\n self.assertEqual(\n ffw_hidden_topk_values.shape,\n (\n batch_size,\n sequence_length,\n sow_config.mlp_hidden_topk,\n ),\n )\n self.assertTrue(hasattr(layer.mlp, 'hidden_topk_indices'))\n ffw_hidden_topk_indices = layer.mlp.hidden_topk_indices[0]\n self.assertIsNotNone(ffw_hidden_topk_indices)\n self.assertEqual(\n ffw_hidden_topk_indices.shape,\n (\n batch_size,\n sequence_length,\n sow_config.mlp_hidden_topk,\n ),\n )\n else:\n self.assertFalse(hasattr(layer.mlp, 'hidden_topk_values'))\n self.assertFalse(hasattr(layer.mlp, 'hidden_topk_indices'))\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "examples/gemma/utils.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# Copied over from MaxText\n# (https://github.com/google/maxtext/blob/main/MaxText/max_utils.py).\n\"\"\"Provides utilities for training the Flax gemma example.\"\"\"\n\nfrom collections.abc import Callable\nimport logging\nfrom typing import Any\n\nfrom flax import nnx\nimport transformer\nfrom flax.training import train_state\nimport jax\nfrom jax.experimental import mesh_utils\nimport jax.numpy as jnp\nimport numpy as np\n\n\nDtype = Any\nShape = tuple[int, ...]\n\n\nclass TrainState(train_state.TrainState):\n graphdef: nnx.GraphDef[transformer.Transformer]\n\n\n# Mesh utils.\n# -----------------------------------------------------------------------------\n\n\ndef create_device_mesh(config: Any):\n \"\"\"Creates a device mesh with each slice in its own data parallel group.\n\n If there is only one slice, uses two replicas.\n\n Args:\n config: The training configuration.\n Returns:\n The device mesh.\n \"\"\"\n devices = jax.devices()\n num_devices = len(devices)\n try:\n num_slices = 1 + max([d.slice_index for d in devices])\n except AttributeError:\n num_slices = 1\n num_devices_per_slice = num_devices // num_slices\n logging.info(f'Devices: {devices}') # pylint: disable=logging-fstring-interpolation\n logging.info(f'Number of devices: {num_devices}') # pylint: disable=logging-fstring-interpolation\n\n multi_slice_env = hasattr(jax.devices()[0], 'slice_index')\n\n dcn_parallelism = [\n config.dcn_data_parallelism,\n config.dcn_fsdp_parallelism,\n config.dcn_tensor_parallelism,\n ]\n ici_parallelism = [\n config.ici_data_parallelism,\n config.ici_fsdp_parallelism,\n config.ici_tensor_parallelism,\n ]\n\n # Find possible unspecified parallelisms\n dcn_parallelism = fill_unspecified_mesh_axes(\n dcn_parallelism, num_slices, 'DCN'\n )\n ici_parallelism = fill_unspecified_mesh_axes(\n ici_parallelism, num_devices_per_slice, 'ICI'\n )\n\n if multi_slice_env:\n mesh = mesh_utils.create_hybrid_device_mesh(\n ici_parallelism, dcn_parallelism\n )\n else:\n mesh = mesh_utils.create_device_mesh(ici_parallelism)\n\n logging.info(f'Decided on mesh: {mesh}') # pylint: disable=logging-fstring-interpolation\n logging.info(f'Mesh shape: {mesh.shape}') # pylint: disable=logging-fstring-interpolation\n\n return mesh\n\n\ndef fill_unspecified_mesh_axes(\n parallelism_vals, target_product, parallelism_type\n):\n \"\"\"Evaluates unspecified DCN/ICI parallelism values.\"\"\"\n if -1 in parallelism_vals:\n assert parallelism_vals.count(-1) == 1, (\n f'Found unspecified values (-1) for more than one {parallelism_type} '\n ' parallelism axis. At most one axis can be unspecified.'\n )\n\n determined_val = target_product / np.prod(parallelism_vals) * -1\n\n assert determined_val >= 1 and determined_val.is_integer, (\n 'Unspecified value unable to be determined with the given '\n f' {parallelism_type} parallelism values'\n )\n\n parallelism_vals[parallelism_vals.index(-1)] = int(determined_val)\n\n target_type = 'slices' if parallelism_type == 'DCN' else 'devices per slice'\n\n assert np.prod(parallelism_vals) == target_product, (\n f'Number of {target_type} {target_product} does not match the product'\n f' of the {parallelism_type} parallelism {np.prod(parallelism_vals)}'\n )\n\n return parallelism_vals\n\n\n# State initialization utils.\n# -----------------------------------------------------------------------------\n\n\ndef _to_array(x):\n if not isinstance(x, jax.Array):\n x = jnp.asarray(x)\n return x\n\n\ndef setup_initial_state(\n constructor: Callable[\n [transformer.TransformerConfig, jax.Array], transformer.Transformer\n ],\n tx,\n config: transformer.TransformerConfig,\n rng: jax.Array,\n mesh: jax.sharding.Mesh,\n) -> tuple[TrainState, TrainState]:\n \"\"\"We initialize train state, optionally loading from checkpoint.\n\n Args:\n constructor: the model constructor\n tx: the optax.GradientTransformation\n config: config object\n rng: jax.prng key\n mesh: jax.devices() mesh\n\n Returns:\n state: the initialized train state\n state_mesh_annotations: the mesh annotations for the train state\n \"\"\"\n\n @jax.jit\n def sharded_init():\n model = constructor(config, rng)\n graphdef, params = nnx.split(model, nnx.Param)\n state = TrainState.create(\n apply_fn=graphdef.apply,\n params=params,\n tx=tx,\n graphdef=graphdef,\n )\n state = jax.tree.map(_to_array, state)\n state_spec = nnx.get_partition_spec(state)\n state = jax.lax.with_sharding_constraint(state, state_spec)\n return state\n\n # Initialization\n with jax.set_mesh(mesh):\n state = sharded_init()\n\n state_sharding = nnx.get_named_sharding(state, mesh)\n return state, state_sharding\n" }, { "path": "examples/imagenet/README.md", "content": "## ImageNet classification\n\nTrains a ResNet50 model ([He *et al.*, 2016]) for the ImageNet classification task\n([Russakovsky *et al.*, 2015]).\n\nThis example uses linear learning rate warmup and cosine learning rate schedule.\n\n[He *et al.*, 2016]: https://arxiv.org/abs/1512.03385\n[Russakovsky *et al.*, 2015]: https://arxiv.org/abs/1409.0575\n\nYou can run this code and even modify it directly in Google Colab, no\ninstallation required:\n\nhttps://colab.research.google.com/github/google/flax/blob/main/examples/imagenet/imagenet.ipynb\n\nThe Colab also demonstrates how to load pretrained checkpoints from Cloud\nstorage at\n[gs://flax_public/examples/imagenet/](https://console.cloud.google.com/storage/browser/flax_public/examples/imagenet)\n\nTable of contents:\n\n- [Requirements](#requirements)\n- [Example runs](#example-runs)\n- [Running locally](#running-locally)\n - [Overriding parameters on the command line](#overriding-parameters-on-the-command-line)\n- [Running fake data benchmarks](#running-fake-data-benchmarks)\n- [Running on Cloud](#running-on-cloud)\n - [Preparing the dataset](#preparing-the-dataset)\n - [Google Cloud TPU](#google-cloud-tpu)\n - [Google Cloud GPU](#google-cloud-gpu)\n\n### Requirements\n\n* TensorFlow dataset `imagenet2012:5.*.*`\n* `≈180GB` of RAM if you want to cache the dataset in memory for faster IO\n\n### Example runs\n\nWhile the example should run on a variety of hardware,\nwe have tested the following GPU and TPU configurations:\n\n| Name | Steps | Walltime | Top-1 accuracy | Metrics | Workdir |\n| :---------------------- | -----: | :------- | :------------- | :--------------------------------------------------------------------------------------------------------------------------------------------------- | :------------------------------------------------------------------------------------------------------------------------------------------------------------------- |\n| TPU v3-32 | 125100 | 2.1h | 76.54% | [tfhub.dev](https://tensorboard.dev/experiment/GhPHRoLzTqu7c8vynTk6bg/) | [gs://flax_public/examples/imagenet/tpu_v3_32](https://console.cloud.google.com/storage/browser/flax_public/examples/imagenet/tpu_v3_32) |\n| TPU v2-32 | 125100 | 2.5h | 76.67% | [tfhub.dev](https://tensorboard.dev/experiment/qBJ7T9VPSgO5yeb0HAKbIA/) | [gs://flax_public/examples/imagenet/tpu_v2_32](https://console.cloud.google.com/storage/browser/flax_public/examples/imagenet/tpu_v2_32) |\n| TPU v3-8 | 125100 | 4.4h | 76.37% | [tfhub.dev](https://tensorboard.dev/experiment/JwxRMYrsR4O6V6fnkn3dmg/) | [gs://flax_public/examples/imagenet/tpu](https://console.cloud.google.com/storage/browser/flax_public/examples/imagenet/tpu) |\n| v100_x8 | 250200 | 13.2h | 76.72% | [tfhub.dev](https://tensorboard.dev/experiment/venzpsNXR421XLkvvzSkqQ/#scalars&_smoothingWeight=0®exInput=%5Eimagenet/v100_x8%24) | [gs://flax_public/examples/imagenet/v100_x8](https://console.cloud.google.com/storage/browser/flax_public/examples/imagenet/v100_x8) |\n| v100_x8_mixed_precision | 62500 | 4.3h | 76.27% | [tfhub.dev](https://tensorboard.dev/experiment/venzpsNXR421XLkvvzSkqQ/#scalars&_smoothingWeight=0®exInput=%5Eimagenet/v100_x8_mixed_precision%24) | [gs://flax_public/examples/imagenet/v100_x8_mixed_precision](https://console.cloud.google.com/storage/browser/flax_public/examples/imagenet/v100_x8_mixed_precision) |\n\n\n### Running locally\n\n```shell\npython main.py --workdir=./imagenet --config=configs/default.py\n```\n\n#### Overriding parameters on the command line\n\nSpecify a hyperparameter configuration by the means of setting `--config` flag.\nConfiguration flag is defined using\n[config_flags](https://github.com/google/ml_collections/tree/master#config-flags).\n`config_flags` allows overriding configuration fields. This can be done as\nfollows:\n\n```shell\npython main.py --workdir=./imagenet_default --config=configs/default.py \\\n--config.num_epochs=100\n```\n\n### Running fake data benchmarks\n\nExecute the following code with `flax/examples/imagenet` as your current directory:\n```shell\nbash ../../tests/download_dataset_metadata.sh\npython imagenet_fake_data_benchmark.py\n```\n\nIf you get an error like this:\n```shell\nCloning into 'datasets'...\nfatal: cannot change to 'https://github.com/tensorflow/datasets/': No such file or directory\nerror: failed to initialize sparse-checkout\n```\nThis mean your git version is outdated. Just update it and re-run.\n\n### Running on Cloud\n\n#### Preparing the dataset\n\nFor running the ResNet50 model on imagenet dataset,\nyou first need to prepare the `imagenet2012` dataset.\nDownload the data from http://image-net.org/ as described in the\n[tensorflow_datasets catalog](https://www.tensorflow.org/datasets/catalog/imagenet2012).\nThen point the environment variable `$IMAGENET_DOWNLOAD_PATH`\nto the directory where the downloads are stored and prepare the dataset\nby running\n\n```shell\npython -c \"\nimport tensorflow_datasets as tfds\ntfds.builder('imagenet2012').download_and_prepare(\n download_config=tfds.download.DownloadConfig(\n manual_dir='$IMAGENET_DOWNLOAD_PATH'))\n\"\n```\n\nThe contents of the directory `~/tensorflow_datasets` should be copied to your\ngcs bucket. Point the environment variable `GCS_TFDS_BUCKET` to your bucket and\nrun the following command:\n\n```shell\ngcloud storage cp --recursive ~/tensorflow_datasets gs://$GCS_TFDS_BUCKET/datasets\n```\n\n#### Google Cloud TPU\n\nSee below for commands to set up a single VM with 8 TPUs attached\n(`--accelerator-type v3-8`), or for an entire TPU slice spanning multiple\nVMs (e.g. `--accelerator-type v3-32`). For more details about how to set up and\nuse TPUs, refer to Cloud docs for\n[single VM setup](https://cloud.google.com/tpu/docs/jax-quickstart-tpu-vm)\nand [pod slice setup](https://cloud.google.com/tpu/docs/jax-quickstart-tpu-vm).\n\nFirst create a single TPUv3-8 VM and connect to it:\n\n```\nZONE=us-central1-a\nTPU_TYPE=v3-8\nVM_NAME=imagenet\n\ngcloud alpha compute tpus tpu-vm create $VM_NAME \\\n --zone $ZONE \\\n --accelerator-type $TPU_TYPE \\\n --version v2-alpha\n\ngcloud alpha compute tpus tpu-vm ssh $VM_NAME --zone $ZONE -- \\\n -L 6006:localhost:6006\n```\n\nWhen connected install JAX:\n\n```\npip install \"jax[tpu]>=0.2.21\" \\\n -f https://storage.googleapis.com/jax-releases/libtpu_releases.html\n```\n\nThen install Flax + the example dependencies:\n\n```\ngit clone --depth=1 --branch=main https://github.com/google/flax\ncd flax\npip install -e .\ncd examples/imagenet\npip install -r requirements.txt\n```\n\nAnd finally start the training:\n\n```\nexport TFDS_DATA_DIR=gs://$GCS_TFDS_BUCKET/datasets\npython3 main.py --workdir=$HOME/logs/imagenet_tpu --config=configs/tpu.py \\\n --jax_backend_target=\"grpc://192.168.0.2:8470\"\n```\n\nNote that you might want to set `TFDS_DATA_DIR` as explained above. You probably\nalso want to start the long-running command above in a `tmux` session and start\nsome monitoring in a separate pane (note that we forwarded port 6006 locally\nabove):\n\n```\ntensorboard --logdir=$HOME/logs\n```\n\nWhen running on pod slices, after creating the TPU VM, there are different ways\nof running the training in SPMD fashion on the hosts connected to the TPUs that\nmake up the slice. We simply send the same installation/execution shell commands\nto all hosts in parallel with the command below. If anything fails it's\nusually a good idea to connect to a single host and execute the commands\ninteractively.\n\nFor convenience, the TPU creation commands are inlined below. Please note that\nwe define `GCS_TFDS_BUCKET` to where your data stands in your cloud bucket.\nAlso `YOUR_BUCKET` is the work directory you are experimenting in. You should\nchoose ZONE based on where your TPU and work directory is. [Here](https://cloud.google.com/tpu/docs/types-zones)\nhas some useful information on which zones you can have different types of TPUs.\n\n```shell\nVM_NAME=imagenet\nREPO=https://github.com/google/flax\nBRANCH=main\nWORKDIR=gs://$YOUR_BUCKET/flax/examples/imagenet/$(date +%Y%m%d_%H%M)\n\ngcloud alpha compute tpus tpu-vm create $VM_NAME \\\n --zone=$ZONE \\\n --version v2-alpha --accelerator-type v3-32\nFLAGS=\"--config.batch_size=$((32*256))\"\n\ngcloud alpha compute tpus tpu-vm ssh $VM_NAME --zone $ZONE \\\n--worker=all --command \"\npip install 'jax[tpu]>=0.2.21' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html &&\npip install --user git+$REPO.git &&\ngit clone --depth=1 -b $BRANCH $REPO &&\ncd flax/examples/imagenet &&\npip install -r requirements.txt &&\nexport TFDS_DATA_DIR=gs://$GCS_TFDS_BUCKET/datasets &&\npython3 main.py --workdir=$WORKDIR --config=configs/tpu.py $FLAGS\n\"\n```\n\nPlease don't forget to disconnect and delete your vm after you are done:\n\n```\ngcloud alpha compute tpus tpu-vm delete $VM_NAME \\\n --zone $ZONE\n```\n\n#### Google Cloud GPU\n\nCan be launched with utility script described in\n[../cloud/README.md](../cloud/README.md)\n\nThere are two configurations available:\n\n- `configs/v100_x8.py` : Full precision GPU training\n- `configs/v100_x8_mixed_precision.py` : Mixed precision GPU training. Note that\n mixed precision handling is implemented manually with\n [`training.dynamic_scale`](https://github.com/google/flax/blob/main/flax/training/dynamic_scale.py)\n" }, { "path": "examples/imagenet/configs/default.py", "content": "# Copyright 2021 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"Default Hyperparameter configuration.\"\"\"\n\nimport ml_collections\n\n\ndef get_config():\n \"\"\"Get the default hyperparameter configuration.\"\"\"\n config = ml_collections.ConfigDict()\n\n # As defined in the `models` module.\n config.model = 'ResNet50'\n # `name` argument of tensorflow_datasets.builder()\n config.dataset = 'imagenet2012:5.*.*'\n\n config.learning_rate = 0.1\n config.warmup_epochs = 5.0\n config.momentum = 0.9\n config.batch_size = 128\n config.shuffle_buffer_size = 16 * 128\n config.prefetch = 10\n\n config.num_epochs = 100.0\n config.log_every_steps = 100\n\n config.cache = False\n config.half_precision = False\n\n # If num_train_steps==-1 then the number of training steps is calculated from\n # num_epochs using the entire dataset. Similarly for steps_per_eval.\n config.num_train_steps = -1\n config.steps_per_eval = -1\n\n # whether to profile the training loop\n config.profile = True\n\n return config\n\n\ndef metrics():\n return [\n 'train_loss',\n 'eval_loss',\n 'train_accuracy',\n 'eval_accuracy',\n 'steps_per_second',\n 'train_learning_rate',\n ]\n" }, { "path": "examples/imagenet/configs/fake_data_benchmark.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Hyperparameter configuration for Fake data benchmark.\"\"\"\n\nimport jax\n\nfrom configs import default as default_lib\n\n\ndef get_config():\n \"\"\"Get the hyperparameter configuration for Fake data benchmark.\"\"\"\n # Override default configuration to avoid duplication of field definition.\n config = default_lib.get_config()\n config.batch_size = 256 * jax.device_count()\n config.half_precision = True\n config.num_epochs = 5\n\n # Run for a single step:\n config.num_train_steps = 1\n config.steps_per_eval = 1\n\n return config\n" }, { "path": "examples/imagenet/configs/tpu.py", "content": "# Copyright 2021 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"Hyperparameter configuration to run the example on TPUs.\"\"\"\n\nfrom configs import default as default_lib\n\n\ndef get_config():\n \"\"\"Get the hyperparameter configuration to train on TPUs.\"\"\"\n config = default_lib.get_config()\n\n # Consider setting the batch size to max(tpu_chips * 256, 8 * 1024) if you\n # train on a larger pod slice.\n config.batch_size = 1024\n config.shuffle_buffer_size = 16 * 1024\n config.cache = True\n config.half_precision = True\n\n return config\n\n\nmetrics = default_lib.metrics\n" }, { "path": "examples/imagenet/configs/v100_x8.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Hyperparameter configuration to run the example on 8 x Nvidia V100 GPUs.\"\"\"\n\nfrom configs import default as default_lib\n\n\ndef get_config():\n \"\"\"Get the hyperparameter configuration to train on 8 x Nvidia V100 GPUs.\"\"\"\n # Override default configuration to avoid duplication of field definition.\n config = default_lib.get_config()\n\n config.batch_size = 512\n config.shuffle_buffer_size = 16 * 512\n config.cache = True\n\n return config\n\n\nmetrics = default_lib.metrics\n" }, { "path": "examples/imagenet/configs/v100_x8_mixed_precision.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Hyperparameter configuration to run the example on 8 x Nvidia V100 GPUs.\"\"\"\n\nfrom configs import default as default_lib\n\n\ndef get_config():\n \"\"\"Get the hyperparameter configuration to train on 8 x Nvidia V100 GPUs.\"\"\"\n # Override default configuration to avoid duplication of field definition.\n config = default_lib.get_config()\n\n config.batch_size = 2048\n config.shuffle_buffer_size = 16 * 2048\n config.cache = True\n config.half_precision = True\n\n return config\n\n\nmetrics = default_lib.metrics\n" }, { "path": "examples/imagenet/imagenet.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"# Flax Imagenet Example\\n\",\n \"\\n\",\n \"\\\"Open\\n\",\n \"\\n\",\n \"Demonstration notebook for\\n\",\n \"https://github.com/google/flax/tree/main/examples/imagenet\\n\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"The **Flax Notebook Workflow**:\\n\",\n \"\\n\",\n \"1. Run the entire notebook end-to-end and check out the outputs.\\n\",\n \" - This will open Python files in the right-hand editor!\\n\",\n \" - You'll be able to interactively explore metrics in TensorBoard.\\n\",\n \"2. Change `config` and train for different hyperparameters. Check out the\\n\",\n \" updated TensorBoard plots.\\n\",\n \"3. Update the code in `train.py`. Thanks to `%autoreload`, any changes you\\n\",\n \" make in the file will automatically appear in the notebook. Some ideas to\\n\",\n \" get you started:\\n\",\n \" - Change the model.\\n\",\n \" - Log some per-batch metrics during training.\\n\",\n \" - Add new hyperparameters to `configs/default.py` and use them in\\n\",\n \" `train.py`.\\n\",\n \"4. At any time, feel free to paste code from `train.py` into the notebook\\n\",\n \" and modify it directly there!\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Setup\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 28,\n \"metadata\": {\n \"outputId\": \"cb862d1a-2f71-444f-9770-9f0d53b11389\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"GPU 0: Tesla T4 (UUID: GPU-d1652fa8-88b9-3d02-7a65-e7ebabeb0372)\\n\"\n ]\n }\n ],\n \"source\": [\n \"# Tested with a T4 GPU.\\n\",\n \"!nvidia-smi -L\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 29,\n \"metadata\": {\n \"outputId\": \"80340396-77c2-4654-cc6d-67040f227eb9\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \" Preparing metadata (setup.py) ... \\u001b[?25l\\u001b[?25hdone\\n\"\n ]\n }\n ],\n \"source\": [\n \"# Install ml-collections & latest Flax version from Github.\\n\",\n \"!pip install -q clu ml-collections git+https://github.com/google/flax\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 30,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"example_directory = 'examples/imagenet'\\n\",\n \"editor_relpaths = ('configs/default.py', 'input_pipeline.py', 'models.py', 'train.py')\\n\",\n \"\\n\",\n \"repo, branch = 'https://github.com/google/flax', 'main'\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 31,\n \"metadata\": {\n \"cellView\": \"form\",\n \"outputId\": \"9449a7b4-8a5d-4446-abe0-7886435ebd1c\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/html\": [\n \"

WARNING : Editing in VM - changes lost after reboot!!

\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"application/javascript\": [\n \"\\n\",\n \" ((filepath) => {{\\n\",\n \" if (!google.colab.kernel.accessAllowed) {{\\n\",\n \" return;\\n\",\n \" }}\\n\",\n \" google.colab.files.view(filepath);\\n\",\n \" }})(\\\"/content/examples/imagenet/configs/default.py\\\")\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"application/javascript\": [\n \"\\n\",\n \" ((filepath) => {{\\n\",\n \" if (!google.colab.kernel.accessAllowed) {{\\n\",\n \" return;\\n\",\n \" }}\\n\",\n \" google.colab.files.view(filepath);\\n\",\n \" }})(\\\"/content/examples/imagenet/input_pipeline.py\\\")\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"application/javascript\": [\n \"\\n\",\n \" ((filepath) => {{\\n\",\n \" if (!google.colab.kernel.accessAllowed) {{\\n\",\n \" return;\\n\",\n \" }}\\n\",\n \" google.colab.files.view(filepath);\\n\",\n \" }})(\\\"/content/examples/imagenet/models.py\\\")\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"application/javascript\": [\n \"\\n\",\n \" ((filepath) => {{\\n\",\n \" if (!google.colab.kernel.accessAllowed) {{\\n\",\n \" return;\\n\",\n \" }}\\n\",\n \" google.colab.files.view(filepath);\\n\",\n \" }})(\\\"/content/examples/imagenet/train.py\\\")\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"# (If you run this code in Jupyter[lab], then you're already in the\\n\",\n \"# example directory and nothing needs to be done.)\\n\",\n \"\\n\",\n \"#@markdown **Fetch newest Flax, copy example code**\\n\",\n \"#@markdown\\n\",\n \"#@markdown **If you select no** below, then the files will be stored on the\\n\",\n \"#@markdown *ephemeral* Colab VM. **After some time of inactivity, this VM will\\n\",\n \"#@markdown be restarted and any changes are lost**.\\n\",\n \"#@markdown\\n\",\n \"#@markdown **If you select yes** below, then you will be asked for your\\n\",\n \"#@markdown credentials to mount your personal Google Drive. In this case, all\\n\",\n \"#@markdown changes you make will be *persisted*, and even if you re-run the\\n\",\n \"#@markdown Colab later on, the files will still be the same (you can of course\\n\",\n \"#@markdown remove directories inside your Drive's `flax/` root if you want to\\n\",\n \"#@markdown manually revert these files).\\n\",\n \"\\n\",\n \"if 'google.colab' in str(get_ipython()):\\n\",\n \" import os\\n\",\n \" os.chdir('/content')\\n\",\n \" # Download Flax repo from Github.\\n\",\n \" if not os.path.isdir('flaxrepo'):\\n\",\n \" !git clone --depth=1 -b $branch $repo flaxrepo\\n\",\n \" # Copy example files & change directory.\\n\",\n \" mount_gdrive = 'no' #@param ['yes', 'no']\\n\",\n \" if mount_gdrive == 'yes':\\n\",\n \" DISCLAIMER = 'Note : Editing in your Google Drive, changes will persist.'\\n\",\n \" from google.colab import drive\\n\",\n \" drive.mount('/content/gdrive')\\n\",\n \" example_root_path = f'/content/gdrive/My Drive/flax/{example_directory}'\\n\",\n \" else:\\n\",\n \" DISCLAIMER = 'WARNING : Editing in VM - changes lost after reboot!!'\\n\",\n \" example_root_path = f'/content/{example_directory}'\\n\",\n \" from IPython import display\\n\",\n \" display.display(display.HTML(\\n\",\n \" f'

{DISCLAIMER}

'))\\n\",\n \" if not os.path.isdir(example_root_path):\\n\",\n \" os.makedirs(example_root_path)\\n\",\n \" !cp -r flaxrepo/$example_directory/* \\\"$example_root_path\\\"\\n\",\n \" os.chdir(example_root_path)\\n\",\n \" from google.colab import files\\n\",\n \" for relpath in editor_relpaths:\\n\",\n \" s = open(f'{example_root_path}/{relpath}').read()\\n\",\n \" open(f'{example_root_path}/{relpath}', 'w').write(\\n\",\n \" f'## {DISCLAIMER}\\\\n' + '#' * (len(DISCLAIMER) + 3) + '\\\\n\\\\n' + s)\\n\",\n \" files.view(f'{example_root_path}/{relpath}')\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {\n \"outputId\": \"acc1f45d-5062-4ff3-e6d4-10b4ffe0f8ef\"\n },\n \"outputs\": [],\n \"source\": [\n \"# Note : In Colab, above cell changed the working directory.\\n\",\n \"!pwd\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Imports / Helpers\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 33,\n \"metadata\": {\n \"cellView\": \"form\",\n \"outputId\": \"9dc7fb32-331e-44a6-b6e8-830f6a64d845\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"[StreamExecutorGpuDevice(id=0, process_index=0, slice_index=0)]\"\n ]\n },\n \"execution_count\": 33,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"# It's possible to run this Colab with TPUs:\\n\",\n \"# 1. change runtime type to TPU\\n\",\n \"# 2. install compatible flax version: `!pip install flax==0.6.4 jax==0.3.25 jaxlib==0.3.25`\\n\",\n \"# 3. uncomment lines below\\n\",\n \"\\n\",\n \"# import flax, jax, jax.tools.colab_tpu\\n\",\n \"# jax.tools.colab_tpu.setup_tpu()\\n\",\n \"\\n\",\n \"import jax\\n\",\n \"jax.devices()\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 34,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"import json\\n\",\n \"from absl import logging\\n\",\n \"import flax\\n\",\n \"import jax\\n\",\n \"import jax.numpy as jnp\\n\",\n \"from matplotlib import pyplot as plt\\n\",\n \"import numpy as np\\n\",\n \"import tensorflow as tf\\n\",\n \"import tensorflow_datasets as tfds\\n\",\n \"\\n\",\n \"logging.set_verbosity(logging.INFO)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 35,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# Helper functions for images.\\n\",\n \"\\n\",\n \"def show_img(img, ax=None, title=None):\\n\",\n \" \\\"\\\"\\\"Shows a single image.\\\"\\\"\\\"\\n\",\n \" if ax is None:\\n\",\n \" ax = plt.gca()\\n\",\n \" img *= tf.constant(input_pipeline.STDDEV_RGB, shape=[1, 1, 3], dtype=img.dtype)\\n\",\n \" img += tf.constant(input_pipeline.MEAN_RGB, shape=[1, 1, 3], dtype=img.dtype)\\n\",\n \" img = np.clip(img.numpy().astype(int), 0, 255)\\n\",\n \" ax.imshow(img)\\n\",\n \" ax.set_xticks([])\\n\",\n \" ax.set_yticks([])\\n\",\n \" if title:\\n\",\n \" ax.set_title(title)\\n\",\n \"\\n\",\n \"def show_img_grid(imgs, titles):\\n\",\n \" \\\"\\\"\\\"Shows a grid of images.\\\"\\\"\\\"\\n\",\n \" n = int(np.ceil(len(imgs)**.5))\\n\",\n \" _, axs = plt.subplots(n, n, figsize=(3 * n, 3 * n))\\n\",\n \" for i, (img, title) in enumerate(zip(imgs, titles)):\\n\",\n \" show_img(img, axs[i // n][i % n], title)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 36,\n \"metadata\": {\n \"outputId\": \"f943d165-b953-4a70-9f93-96eb857c3d53\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"The autoreload extension is already loaded. To reload it, use:\\n\",\n \" %reload_ext autoreload\\n\"\n ]\n }\n ],\n \"source\": [\n \"# Local imports from current directory - auto reload.\\n\",\n \"# Any changes you make to train.py will appear automatically.\\n\",\n \"%load_ext autoreload\\n\",\n \"%autoreload 2\\n\",\n \"import input_pipeline\\n\",\n \"import models\\n\",\n \"import train\\n\",\n \"from configs import default as config_lib\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Dataset\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 37,\n \"metadata\": {\n \"outputId\": \"a9b6cfe9-cc1c-451a-f8f7-69356cb7bdd2\"\n },\n \"outputs\": [\n {\n \"name\": \"stderr\",\n \"output_type\": \"stream\",\n \"text\": [\n \"INFO:absl:No config specified, defaulting to config: imagenette/full-size-v2\\n\",\n \"INFO:absl:Load dataset info from /root/tensorflow_datasets/imagenette/full-size-v2/1.0.0\\n\",\n \"INFO:absl:Reusing dataset imagenette (/root/tensorflow_datasets/imagenette/full-size-v2/1.0.0)\\n\",\n \"INFO:absl:Constructing tf.data.Dataset imagenette for split train, from /root/tensorflow_datasets/imagenette/full-size-v2/1.0.0\\n\"\n ]\n },\n {\n \"data\": {\n \"text/plain\": [\n \"tfds.core.DatasetInfo(\\n\",\n \" name='imagenette',\\n\",\n \" full_name='imagenette/full-size-v2/1.0.0',\\n\",\n \" description=\\\"\\\"\\\"\\n\",\n \" Imagenette is a subset of 10 easily classified classes from the Imagenet\\n\",\n \" dataset. It was originally prepared by Jeremy Howard of FastAI. The objective\\n\",\n \" behind putting together a small version of the Imagenet dataset was mainly\\n\",\n \" because running new ideas/algorithms/experiments on the whole Imagenet take a\\n\",\n \" lot of time.\\n\",\n \" \\n\",\n \" This version of the dataset allows researchers/practitioners to quickly try out\\n\",\n \" ideas and share with others. The dataset comes in three variants:\\n\",\n \" \\n\",\n \" * Full size\\n\",\n \" * 320 px\\n\",\n \" * 160 px\\n\",\n \" \\n\",\n \" Note: The v2 config correspond to the new 70/30 train/valid split (released in\\n\",\n \" Dec 6 2019).\\n\",\n \" \\\"\\\"\\\",\\n\",\n \" config_description=\\\"\\\"\\\"\\n\",\n \" full-size variant.\\n\",\n \" \\\"\\\"\\\",\\n\",\n \" homepage='https://github.com/fastai/imagenette',\\n\",\n \" data_path='/root/tensorflow_datasets/imagenette/full-size-v2/1.0.0',\\n\",\n \" file_format=tfrecord,\\n\",\n \" download_size=1.45 GiB,\\n\",\n \" dataset_size=1.46 GiB,\\n\",\n \" features=FeaturesDict({\\n\",\n \" 'image': Image(shape=(None, None, 3), dtype=uint8),\\n\",\n \" 'label': ClassLabel(shape=(), dtype=int64, num_classes=10),\\n\",\n \" }),\\n\",\n \" supervised_keys=('image', 'label'),\\n\",\n \" disable_shuffling=False,\\n\",\n \" splits={\\n\",\n \" 'train': ,\\n\",\n \" 'validation': ,\\n\",\n \" },\\n\",\n \" citation=\\\"\\\"\\\"@misc{imagenette,\\n\",\n \" author = \\\"Jeremy Howard\\\",\\n\",\n \" title = \\\"imagenette\\\",\\n\",\n \" url = \\\"https://github.com/fastai/imagenette/\\\"\\n\",\n \" }\\\"\\\"\\\",\\n\",\n \")\"\n ]\n },\n \"execution_count\": 37,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"# We load \\\"imagenette\\\" that has similar pictures to \\\"imagenet2012\\\" but can be\\n\",\n \"# downloaded automatically and is much smaller.\\n\",\n \"dataset_builder = tfds.builder('imagenette')\\n\",\n \"dataset_builder.download_and_prepare()\\n\",\n \"ds = dataset_builder.as_dataset('train')\\n\",\n \"dataset_builder.info\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 38,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# Utilities to help with Imagenette labels.\\n\",\n \"\\n\",\n \"![ ! -f mapping_imagenet.json ] && wget --no-check-certificate https://raw.githubusercontent.com/ozendelait/wordnet-to-json/master/mapping_imagenet.json\\n\",\n \"\\n\",\n \"with open('mapping_imagenet.json') as f:\\n\",\n \" mapping_imagenet = json.load(f)\\n\",\n \"# Mapping imagenette label name to imagenet label index.\\n\",\n \"imagenette_labels = {\\n\",\n \" d['v3p0']: d['label']\\n\",\n \" for d in mapping_imagenet\\n\",\n \"}\\n\",\n \"# Mapping imagenette label name to human-readable label.\\n\",\n \"imagenette_idx = {\\n\",\n \" d['v3p0']: idx\\n\",\n \" for idx, d in enumerate(mapping_imagenet)\\n\",\n \"}\\n\",\n \"\\n\",\n \"def imagenette_label(idx):\\n\",\n \" \\\"\\\"\\\"Returns a short human-readable string for provided imagenette index.\\\"\\\"\\\"\\n\",\n \" net = dataset_builder.info.features['label'].int2str(idx)\\n\",\n \" return imagenette_labels[net].split(',')[0]\\n\",\n \"\\n\",\n \"def imagenette_imagenet2012(idx):\\n\",\n \" \\\"\\\"\\\"Returns the imagenet2012 index for provided imagenette index.\\\"\\\"\\\"\\n\",\n \" net = dataset_builder.info.features['label'].int2str(idx)\\n\",\n \" return imagenette_idx[net]\\n\",\n \"\\n\",\n \"def imagenet2012_label(idx):\\n\",\n \" \\\"\\\"\\\"Returns a short human-readable string for provided imagenet2012 index.\\\"\\\"\\\"\\n\",\n \" return mapping_imagenet[idx]['label'].split(',')[0]\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 39,\n \"metadata\": {\n \"outputId\": \"78142300-cc8b-4a6c-f781-5ab29578d828\"\n },\n \"outputs\": [\n {\n \"name\": \"stderr\",\n \"output_type\": \"stream\",\n \"text\": [\n \"INFO:absl:Constructing tf.data.Dataset imagenette for split train[0:9469], from /root/tensorflow_datasets/imagenette/full-size-v2/1.0.0\\n\",\n \"WARNING:absl:options.experimental_threading is deprecated. Use options.threading instead.\\n\",\n \"INFO:absl:Constructing tf.data.Dataset imagenette for split validation[0:3925], from /root/tensorflow_datasets/imagenette/full-size-v2/1.0.0\\n\",\n \"WARNING:absl:options.experimental_threading is deprecated. Use options.threading instead.\\n\"\n ]\n }\n ],\n \"source\": [\n \"train_ds = input_pipeline.create_split(\\n\",\n \" dataset_builder, 128, train=True,\\n\",\n \")\\n\",\n \"eval_ds = input_pipeline.create_split(\\n\",\n \" dataset_builder, 128, train=False,\\n\",\n \")\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 40,\n \"metadata\": {\n \"outputId\": \"8b3b9cf2-7649-4953-99bb-32a689fe0a29\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"{'image': (TensorShape([128, 224, 224, 3]), tf.float32),\\n\",\n \" 'label': (TensorShape([128]), tf.int64)}\"\n ]\n },\n \"execution_count\": 40,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"train_batch = next(iter(train_ds))\\n\",\n \"{k: (v.shape, v.dtype) for k, v in train_batch.items()}\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Training from scratch\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# Get a live update during training - use the \\\"refresh\\\" button!\\n\",\n \"# (In Jupyter[lab] start \\\"tensorbaord\\\" in the local directory instead.)\\n\",\n \"if 'google.colab' in str(get_ipython()):\\n\",\n \" %load_ext tensorboard\\n\",\n \" %tensorboard --logdir=.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 42,\n \"metadata\": {\n \"outputId\": \"2d0bc789-213d-4a34-a7b1-e7852b40f375\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"batch_size: 32\\n\",\n \"cache: false\\n\",\n \"dataset: imagenette\\n\",\n \"half_precision: true\\n\",\n \"learning_rate: 0.1\\n\",\n \"log_every_steps: 100\\n\",\n \"model: ResNet18\\n\",\n \"momentum: 0.9\\n\",\n \"num_epochs: 5.0\\n\",\n \"num_train_steps: -1\\n\",\n \"prefetch: 1\\n\",\n \"shuffle_buffer_size: 128\\n\",\n \"steps_per_eval: -1\\n\",\n \"warmup_epochs: 0.5\"\n ]\n },\n \"execution_count\": 42,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"config = config_lib.get_config()\\n\",\n \"config.dataset = 'imagenette'\\n\",\n \"config.model = 'ResNet18'\\n\",\n \"config.half_precision = True\\n\",\n \"# Reduce batch size, shuffle buffer and prefetch to avoid Colab runtime OOM.\\n\",\n \"config.batch_size = 32\\n\",\n \"config.shuffle_buffer_size = 128\\n\",\n \"config.prefetch = 1\\n\",\n \"# Reduce epochs so this Colab doesn't take forever.\\n\",\n \"config.num_epochs = 5.0\\n\",\n \"config.warmup_epochs = 0.5\\n\",\n \"config\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 43,\n \"metadata\": {\n \"outputId\": \"de56d320-c336-459b-f258-5d6ae41ce0af\"\n },\n \"outputs\": [\n {\n \"name\": \"stderr\",\n \"output_type\": \"stream\",\n \"text\": [\n \"INFO:absl:Constructing tf.data.Dataset imagenette for split train[0:9469], from /root/tensorflow_datasets/imagenette/full-size-v2/1.0.0\\n\",\n \"WARNING:absl:options.experimental_threading is deprecated. Use options.threading instead.\\n\",\n \"INFO:absl:Constructing tf.data.Dataset imagenette for split validation[0:3925], from /root/tensorflow_datasets/imagenette/full-size-v2/1.0.0\\n\",\n \"WARNING:absl:options.experimental_threading is deprecated. Use options.threading instead.\\n\"\n ]\n }\n ],\n \"source\": [\n \"# Regenerate datasets with updated batch_size.\\n\",\n \"train_ds = input_pipeline.create_split(\\n\",\n \" dataset_builder, config.batch_size, train=True,\\n\",\n \")\\n\",\n \"eval_ds = input_pipeline.create_split(\\n\",\n \" dataset_builder, config.batch_size, train=False,\\n\",\n \")\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 44,\n \"metadata\": {\n \"outputId\": \"018da2c5-c6f0-42ac-843f-7ac855a6bf14\"\n },\n \"outputs\": [\n {\n \"name\": \"stderr\",\n \"output_type\": \"stream\",\n \"text\": [\n \"INFO:absl:No config specified, defaulting to config: imagenette/full-size-v2\\n\",\n \"INFO:absl:Load dataset info from /root/tensorflow_datasets/imagenette/full-size-v2/1.0.0\\n\",\n \"INFO:absl:Constructing tf.data.Dataset imagenette for split train[0:9469], from /root/tensorflow_datasets/imagenette/full-size-v2/1.0.0\\n\",\n \"WARNING:absl:options.experimental_threading is deprecated. Use options.threading instead.\\n\"\n ]\n },\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"\\n\",\n \"\\n\",\n \"ResNet18_lr=0.03\\n\"\n ]\n },\n {\n \"name\": \"stderr\",\n \"output_type\": \"stream\",\n \"text\": [\n \"INFO:absl:Constructing tf.data.Dataset imagenette for split validation[0:3925], from /root/tensorflow_datasets/imagenette/full-size-v2/1.0.0\\n\",\n \"WARNING:absl:options.experimental_threading is deprecated. Use options.threading instead.\\n\",\n \"INFO:absl:Restoring legacy Flax checkpoint from models/ResNet18_lr=0.03/checkpoint_1475\\n\",\n \"INFO:absl:Initial compilation, this might take some minutes...\\n\",\n \"INFO:absl:No config specified, defaulting to config: imagenette/full-size-v2\\n\",\n \"INFO:absl:Load dataset info from /root/tensorflow_datasets/imagenette/full-size-v2/1.0.0\\n\",\n \"INFO:absl:Constructing tf.data.Dataset imagenette for split train[0:9469], from /root/tensorflow_datasets/imagenette/full-size-v2/1.0.0\\n\",\n \"WARNING:absl:options.experimental_threading is deprecated. Use options.threading instead.\\n\"\n ]\n },\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"\\n\",\n \"\\n\",\n \"ResNet18_lr=0.1\\n\"\n ]\n },\n {\n \"name\": \"stderr\",\n \"output_type\": \"stream\",\n \"text\": [\n \"INFO:absl:Constructing tf.data.Dataset imagenette for split validation[0:3925], from /root/tensorflow_datasets/imagenette/full-size-v2/1.0.0\\n\",\n \"WARNING:absl:options.experimental_threading is deprecated. Use options.threading instead.\\n\",\n \"INFO:absl:Restoring legacy Flax checkpoint from models/ResNet18_lr=0.1/checkpoint_1475\\n\",\n \"INFO:absl:Initial compilation, this might take some minutes...\\n\",\n \"INFO:absl:No config specified, defaulting to config: imagenette/full-size-v2\\n\",\n \"INFO:absl:Load dataset info from /root/tensorflow_datasets/imagenette/full-size-v2/1.0.0\\n\",\n \"INFO:absl:Constructing tf.data.Dataset imagenette for split train[0:9469], from /root/tensorflow_datasets/imagenette/full-size-v2/1.0.0\\n\",\n \"WARNING:absl:options.experimental_threading is deprecated. Use options.threading instead.\\n\"\n ]\n },\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"\\n\",\n \"\\n\",\n \"ResNet18_lr=0.3\\n\"\n ]\n },\n {\n \"name\": \"stderr\",\n \"output_type\": \"stream\",\n \"text\": [\n \"INFO:absl:Constructing tf.data.Dataset imagenette for split validation[0:3925], from /root/tensorflow_datasets/imagenette/full-size-v2/1.0.0\\n\",\n \"WARNING:absl:options.experimental_threading is deprecated. Use options.threading instead.\\n\",\n \"INFO:absl:Restoring legacy Flax checkpoint from models/ResNet18_lr=0.3/checkpoint_1475\\n\",\n \"INFO:absl:Initial compilation, this might take some minutes...\\n\"\n ]\n }\n ],\n \"source\": [\n \"# Takes ~1.5 min / epoch.\\n\",\n \"for lr in (0.03, 0.1, 0.3):\\n\",\n \" config.learning_rate = lr\\n\",\n \" name = f'{config.model}_lr={config.learning_rate}'\\n\",\n \" print(f'\\\\n\\\\n{name}')\\n\",\n \" state = train.train_and_evaluate(config, workdir=f'./models/{name}')\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 45,\n \"metadata\": {\n \"cellView\": \"form\"\n },\n \"outputs\": [],\n \"source\": [\n \"if 'google.colab' in str(get_ipython()):\\n\",\n \" #@markdown You can upload the training results directly to https://tensorbaord.dev\\n\",\n \" #@markdown\\n\",\n \" #@markdown Note that everybody with the link will be able to see the data.\\n\",\n \" upload_data = 'no' #@param ['yes', 'no']\\n\",\n \" if upload_data == 'yes':\\n\",\n \" !tensorboard dev upload --one_shot --logdir ./models --name 'Flax examples/mnist'\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Load pre-trained model\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 46,\n \"metadata\": {\n \"outputId\": \"b06fa3d8-a950-46d2-e03e-fc6c971bdbd0\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"-rw-r--r-- 1 root root 196M Apr 3 13:38 checkpoint_250200\\n\"\n ]\n }\n ],\n \"source\": [\n \"# Load model checkpoint from cloud.\\n\",\n \"from flax.training import checkpoints\\n\",\n \"\\n\",\n \"config_name = 'v100_x8'\\n\",\n \"pretrained_path = f'gs://flax_public/examples/imagenet/{config_name}'\\n\",\n \"latest_checkpoint = checkpoints.natural_sort(\\n\",\n \" tf.io.gfile.glob(f'{pretrained_path}/checkpoint_*'))[0]\\n\",\n \"if not os.path.exists(os.path.basename(latest_checkpoint)):\\n\",\n \" tf.io.gfile.copy(latest_checkpoint, os.path.basename(latest_checkpoint))\\n\",\n \"!ls -lh checkpoint_*\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 47,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# Load config that was used to train checkpoint.\\n\",\n \"import importlib\\n\",\n \"config = importlib.import_module(f'configs.{config_name}').get_config()\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 48,\n \"metadata\": {\n \"outputId\": \"57777298-4b4b-4a82-b0f2-4b6ff3b949af\"\n },\n \"outputs\": [\n {\n \"name\": \"stderr\",\n \"output_type\": \"stream\",\n \"text\": [\n \"INFO:absl:Restoring legacy Flax checkpoint from ./checkpoint_250200\\n\"\n ]\n }\n ],\n \"source\": [\n \"# Load models & state (takes ~1 min to load the model).\\n\",\n \"model_cls = getattr(models, config.model)\\n\",\n \"model = train.create_model(\\n\",\n \" model_cls=model_cls, half_precision=config.half_precision)\\n\",\n \"base_learning_rate = config.learning_rate * config.batch_size / 256.\\n\",\n \"steps_per_epoch = (\\n\",\n \" dataset_builder.info.splits['train'].num_examples // config.batch_size\\n\",\n \")\\n\",\n \"learning_rate_fn = train.create_learning_rate_fn(\\n\",\n \" config, base_learning_rate, steps_per_epoch)\\n\",\n \"state = train.create_train_state(\\n\",\n \" jax.random.key(0), config, model, image_size=input_pipeline.IMAGE_SIZE,\\n\",\n \" learning_rate_fn=learning_rate_fn)\\n\",\n \"state = train.restore_checkpoint(state, './')\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Inference\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 49,\n \"metadata\": {\n \"outputId\": \"793a656b-f3ad-4596-ad4f-44c686e5e885\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"{'image': TensorShape([32, 224, 224, 3]), 'label': TensorShape([32])}\"\n ]\n },\n \"execution_count\": 49,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"# Load batch from imagenette eval set.\\n\",\n \"batch = next(iter(eval_ds))\\n\",\n \"{k: v.shape for k, v in batch.items()}\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 50,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# Evaluate using model trained on imagenet.\\n\",\n \"logits = model.apply({'params': state.params, 'batch_stats': state.batch_stats}, batch['image'][:128], train=False)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 51,\n \"metadata\": {\n \"outputId\": \"6ab4bb0b-2c03-4663-d7ac-e51b979d121f\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"[16, 24]\"\n ]\n },\n \"execution_count\": 51,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"# Find classification mistakes.\\n\",\n \"preds_labels = list(zip(logits.argmax(axis=-1), map(imagenette_imagenet2012, batch['label'])))\\n\",\n \"error_idxs = [idx for idx, (pred, label) in enumerate(preds_labels) if pred != label]\\n\",\n \"error_idxs\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 52,\n \"metadata\": {\n \"outputId\": \"142c1acf-037e-4ab0-9ca3-bdf0829c51c4\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"image/png\": \"iVBORw0KGgoAAAANSUhEUgAAAk0AAAJKCAYAAAAx/3HgAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAA9hAAAPYQGoP6dpAAEAAElEQVR4nOz9d7wlVZX3j7/3rnzSzfd2387d5IyNiGRBRQQUEyM+Chhx1NHRnzpRUcfw+HMcnXFGMcyAaQIC84g6ijqAARAEJDYNTedwu28+sXLt7x+7zrn30k1wBLGZ83nZck+dOlW70tqfWuuz1hJKKUUXXXTRRRdddNFFF48L+UwPoIsuuuiiiy666GJ/QJc0ddFFF1100UUXXTwJdElTF1100UUXXXTRxZNAlzR10UUXXXTRRRddPAl0SVMXXXTRRRdddNHFk0CXNHXRRRdddNFFF108CXRJUxdddNFFF1100cWTQJc0ddFFF1100UUXXTwJdElTF1100UUXXXTRxZNAlzQ9i3HTTTchhOCmm256pofyB4OPfOQjCCGe6WF00UUXvwf8IdvA008/ndNPP73zecuWLQghuPLKK5+xMXXxxOiSpi4eE//6r//K5z//+Wd6GP8rcfrpp3PJJZc808PooosuuuhiHsxnegBd/OHiX//1X7n//vv50z/902d6KF100UUXz2qsWLEC3/exLOuZHkoXj4Oup+kPEFmWEQTBMz2MLp4BtFqtZ3oIXXTxjOMP2QY+Xc+oEALXdTEM42nZfhdPDbqk6WlCWzuzfv16LrjgAiqVCgMDA7znPe/ZyxgIIXjXu97Ft7/9bQ4//HAcx+FHP/oRADt37uRNb3oTIyMjOI7D4Ycfzr/8y7/stb8dO3Zw/vnnUywWGR4e5r3vfS9hGO61XqvVYv369UxOTj7u+E8//XR+8IMfsHXrVoQQCCFYuXIlAFEU8eEPf5i1a9fS09NDsVjklFNO4cYbb1ywjXaM/m//9m/53Oc+x4oVK/A8j9NOO437779/r32uX7+eV7/61fT39+O6LscddxzXXXfd447zf7KfR+OKK67gjDPOYHh4GMdxOOyww/jSl760YJ2LL76YwcFB4jje6/cvfvGLOfjggxcs+9a3vsXatWvxPI/+/n5e+9rXsn379gXrnH766RxxxBHceeednHrqqRQKBf7yL//yCcfbRRf7A/Z3GwiP/4x+97vf5ZxzzmF0dBTHcVizZg1/8zd/Q5qme23nK1/5CmvWrMHzPI4//nh+8Ytf7LXOY2mabrjhBk455RSKxSK9vb28/OUv58EHH3zCsXfx9KAbnnuaccEFF7By5Uo+9alP8atf/Yp/+Id/YGZmhm984xsL1rvhhhu46qqreNe73sXg4CArV65kz549nHDCCR2DMjQ0xA9/+EPe/OY3U6vVOmEz3/c588wz2bZtG+9+97sZHR3lm9/8JjfccMNe47n99tt5wQtewGWXXcZHPvKRxxz3X/3VX1GtVtmxYwef+9znACiVSgDUajW+9rWvceGFF/LWt76Ver3OP//zP3PWWWdx++23c8wxxyzY1je+8Q3q9TrvfOc7CYKAv//7v+eMM87gvvvuY2RkBIAHHniAk046iSVLlvDnf/7nFItFrrrqKs4//3yuueYaXvGKVzzhuX4y+9kXvvSlL3H44Yfzspe9DNM0+d73vsc73vEOsizjne98JwBveMMb+MY3vsH111/Pueee2/nt7t27ueGGG7jssss6yz7xiU/woQ99iAsuuIC3vOUtTExM8IUvfIFTTz2V3/zmN/T29nbWnZqa4uyzz+a1r30tr3/96x93nF10sT9if7WBbTzWM3rllVdSKpV43/veR6lU4oYbbuDDH/4wtVqNz3zmM53f//M//zOXXnopJ554In/6p3/Kpk2beNnLXkZ/fz/Lli173H3/9Kc/5eyzz2b16tV85CMfwfd9vvCFL3DSSSdx1113dV5ku/g9QnXxtOCyyy5TgHrZy162YPk73vEOBah77rmnswxQUkr1wAMPLFj3zW9+s1q8eLGanJxcsPy1r32t6unpUa1WSyml1Oc//3kFqKuuuqqzTrPZVAcccIAC1I033thZfuONNypAXXbZZU94DOecc45asWLFXsuTJFFhGC5YNjMzo0ZGRtSb3vSmzrLNmzcrQHmep3bs2NFZfttttylAvfe97+0sO/PMM9WRRx6pgiDoLMuyTJ144onqwAMPfNxx/jb7aV+X+Wifx/k466yz1OrVqzuf0zRVS5cuVX/0R3+0YL2/+7u/U0IItWnTJqWUUlu2bFGGYahPfOITC9a77777lGmaC5afdtppClCXX3754x5fF13sj3g22MDHe0b3ZTcuvfRSVSgUOnYsiiI1PDysjjnmmAU28ytf+YoC1GmnndZZ1rZjV1xxRWfZMccco4aHh9XU1FRn2T333KOklOqiiy56wvF38dSjG557mtH2VLTxJ3/yJwD813/914Llp512Gocddljns1KKa665hvPOOw+lFJOTk51/Z511FtVqlbvuuquzrcWLF/PqV7+68/tCocDb3va2vcZz+umno5R6Um9YjwXDMLBtG9Dag+npaZIk4bjjjuuMaT7OP/98lixZ0vl8/PHH87znPa9zDqanp7nhhhu44IILqNfrneOcmprirLPOYsOGDezcufMJx/VE+3kseJ7X+btarTI5Oclpp53Gpk2bqFarAEgp+T//5/9w3XXXUa/XO+t/+9vf5sQTT2TVqlUAXHvttWRZxgUXXLDgmi1atIgDDzxwrxCm4zi88Y1vfMJj66KL/RX7uw18rGd0vt1o261TTjmlE/4DuOOOOxgfH+ftb397x2YCXHLJJfT09DzufsfGxrj77ru55JJL6O/v7yw/6qijeNGLXvSEdq2Lpwfd8NzTjAMPPHDB5zVr1iClZMuWLQuWtyfdNiYmJpidneUrX/kKX/nKV/a57fHxcQC2bt3KAQccsFf9oUfrbJ5KfP3rX+ezn/0s69evX6DzefRxwN7nAOCggw7iqquuAuCRRx5BKcWHPvQhPvShD+1zf+Pj4wsI0b7wRPt5LNx8881cdtll3HrrrXuJPKvVase4XXTRRXz605/mP//zP7nooot46KGHuPPOO7n88ss762/YsAGl1D7HAuyVGbNkyZIFxrSLLp5t2N9t4GM9ow888AB//dd/zQ033ECtVlvwXftla+vWrcDe58CyLFavXv24+23/dl/HcOihh3L99dfTbDYpFotP/mC6+J3RJU2/ZzxWYcX5by2gPTgAr3/967n44ov3+ZujjjrqqR3ck8S3vvUtLrnkEs4//3w+8IEPMDw8jGEYfOpTn2Ljxo2/9fbax/r+97+fs846a5/rHHDAAb/TmB8LGzdu5Mwzz+SQQw7h7/7u71i2bBm2bfNf//VffO5zn+uMDeCwww5j7dq1fOtb3+Kiiy7iW9/6FrZtc8EFFyw4FiEEP/zhD/eZBdPWhbXx6OveRRfPduxvNnBfz+js7CynnXYalUqFj33sY6xZswbXdbnrrrv4sz/7swV2o4tnF7qk6WnGhg0bFrxBPfLII2RZ9oQCvqGhIcrlMmma8sIXvvBx112xYgX3338/SqkFBumhhx76ncb+WMbt6quvZvXq1Vx77bUL1pkvhp6PDRs27LXs4Ycf7pyD9huXZVlPeKyPhyfaz77wve99jzAMue6661i+fHln+aPDaG1cdNFFvO9972NsbIx//dd/5ZxzzqGvr6/z/Zo1a1BKsWrVKg466KD/8bF00cWzBfuzDXws3HTTTUxNTXHttddy6qmndpZv3rx5r3GBPgdnnHFGZ3kcx2zevJmjjz76MffR/u2+jmH9+vUMDg52vUzPALqapqcZ//RP/7Tg8xe+8AUAzj777Mf9nWEYvOpVr+Kaa67ZZ9r8xMRE5++XvvSl7Nq1i6uvvrqzrNVq7dOl/duk2xaLxY6b+dFjA605aOO2227j1ltv3ed2/t//+38LNEm33347t912W+ccDA8Pc/rpp/PlL3+ZsbGxvX4//1gfD0+0n31hX8dSrVa54oor9rn+hRdeiBCC97znPWzatInXv/71C75/5StfiWEYfPSjH12wzfY+pqamntSxdNHFswX7sw18vLHBQrsRRRFf/OIXF6x33HHHMTQ0xOWXX04URZ3lV155JbOzs4+7j8WLF3PMMcfw9a9/fcG6999/Pz/+8Y956Utf+j8efxf/c3Q9TU8zNm/ezMte9jJe8pKXcOutt/Ktb32L173udY/7htHG//2//5cbb7yR5z3vebz1rW/lsMMOY3p6mrvuuouf/vSnTE9PA/DWt76Vf/zHf+Siiy7izjvvZPHixXzzm9+kUCjstc3fJt127dq1/Md//Afve9/7eO5zn0upVOK8887j3HPP5dprr+UVr3gF55xzDps3b+byyy/nsMMOo9Fo7LWdAw44gJNPPpk//uM/JgxDPv/5zzMwMMAHP/jBzjr/9E//xMknn8yRRx7JW9/6VlavXs2ePXu49dZb2bFjB/fcc88Tnq8ns59H48UvfjG2bXPeeedx6aWX0mg0+OpXv8rw8PA+CdzQ0BAveclL+M53vkNvby/nnHPOgu/XrFnDxz/+cf7iL/6CLVu2cP7551Mul9m8eTP/+Z//ydve9jbe//73P+GxdNHFswX7sw18LJx44on09fVx8cUX8+53vxshBN/85jf3elGyLIuPf/zjXHrppZxxxhn80R/9EZs3b+aKK654Qk0TwGc+8xnOPvtsnv/85/PmN7+5U3Kgp6fnd0rm6eJ3wO89X+9/CdrptuvWrVOvfvWrVblcVn19fepd73qX8n1/wbqAeuc737nP7ezZs0e9853vVMuWLVOWZalFixapM888U33lK19ZsN7WrVvVy172MlUoFNTg4KB6z3veo370ox/9Tum2jUZDve51r1O9vb0K6JQfyLJMffKTn1QrVqxQjuOoY489Vn3/+99XF1988YISBe0U2s985jPqs5/9rFq2bJlyHEedcsopC9KN29i4caO66KKL1KJFi5RlWWrJkiXq3HPPVVdfffXjjvO32c++Sg5cd9116qijjlKu66qVK1eqT3/60+pf/uVfFKA2b9681/6uuuoqBai3ve1tjzmma665Rp188smqWCyqYrGoDjnkEPXOd75TPfTQQ511TjvtNHX44Yc/7rF10cX+imeDDXy8Z/Tmm29WJ5xwgvI8T42OjqoPfvCD6vrrr99rf0op9cUvflGtWrVKOY6jjjvuOPXzn/9cnXbaaU9YckAppX7605+qk046SXmepyqVijrvvPPUunXrnnDsXTw9EEo9ihp38ZTgIx/5CB/96EeZmJhgcHDwmR7OM4ItW7awatUqPvOZzzyt3pXf137a+O53v8v555/Pz3/+c0455ZSnfX9ddLE/omsDu3g2oqtp6qKL3xJf/epXWb16NSeffPIzPZQuuuiiiy5+j+hqmrro4kni3//937n33nv5wQ9+wN///d8/ZnZhF1100UUXz050SVMXXTxJXHjhhZRKJd785jfzjne845keThdddNFFF79ndDVNXXTRRRdddNFFF08CXU1TF1100UUXXXTRxZNAlzR10UUXXXTRRRddPAl0SdOjcOWVVyKE2KuZ5JPB6aefzhFHHPGUjmflypVccsklT+k2u/jdcNNNNyGE4Kabbnqmh9LF/wJ0bVIXT4SuTfr9oUuauviDxLp16/jIRz6yz4nii1/8IldeeeXvfUz7Gz7ykY88YX+vLrro4smha5N+dzwbbFKXNHXxB4l169bx0Y9+tGuguuiiiz8IdG1SF9AlTV108axDs9l8pofQRRdddNHBs8kmdUnTk8B3v/tdzjnnHEZHR3EchzVr1vA3f/M3pGm6z/XvvPNOTjzxRDzPY9WqVVx++eV7rROGIZdddhkHHHAAjuOwbNkyPvjBDxKG4ROOZ+PGjWzcuPFJjX12dpb3vve9rFy5EsdxWLp0KRdddFGnw3cURXz4wx9m7dq19PT0UCwWOeWUU7jxxhv32ta///u/s3btWsrlMpVKhSOPPJK///u/73wfxzEf/ehHOfDAA3Fdl4GBAU4++WR+8pOfLNjO+vXrefWrX01/fz+u63Lcccdx3XXXdb6/8sorec1rXgPAC17wAoQQnXj9ypUreeCBB/jZz37WWX766acvON4//dM/ZdmyZTiOwwEHHMCnP/1psix7wnO1cuVKzj33XH784x9zzDHH4Louhx12GNdee+0T/vYXv/gFr3nNa1i+fHnner73ve/F9/3OOldccQVCCH7zm9/s9ftPfvKTGIbBzp07O8tuu+02XvKSl9DT00OhUOC0007j5ptvXvC7j3zkIwghWLduHa973evo6+vrVir/X4CuTdLo2qTHRtcmPT3oFrd8ErjyyisplUq8733vo1QqccMNN/DhD3+YWq3GZz7zmQXrzszM8NKXvpQLLriACy+8kKuuuoo//uM/xrZt3vSmNwGQZRkve9nL+OUvf8nb3vY2Dj30UO677z4+97nP8fDDD/P//t//e9zxnHnmmQBPKAxtNBqccsopPPjgg7zpTW/iOc95DpOTk1x33XXs2LGDwcFBarUaX/va17jwwgt561vfSr1e55//+Z8566yzuP322znmmGMA+MlPfsKFF17ImWeeyac//WkAHnzwQW6++Wbe8573APph+dSnPsVb3vIWjj/+eGq1GnfccQd33XUXL3rRiwB44IEHOOmkk1iyZAl//ud/TrFY5KqrruL888/nmmuu4RWveAWnnnoq7373u/mHf/gH/vIv/5JDDz0UgEMPPZTPf/7z/Mmf/AmlUom/+qu/AmBkZASAVqvFaaedxs6dO7n00ktZvnw5t9xyC3/xF3/B2NgYn//855/wWm/YsIE/+qM/4u1vfzsXX3wxV1xxBa95zWv40Y9+1DmGfeE73/kOrVaLP/7jP2ZgYIDbb7+dL3zhC+zYsYPvfOc7ALz61a/mne98J9/+9rc59thjF/z+29/+NqeffjpLliwB4IYbbuDss89m7dq1XHbZZUgpueKKKzjjjDP4xS9+wfHHH7/g9695zWs48MAD+eQnP7lXp/Uunn3o2qSuTerapGcIz2S34D9EXHHFFXt1t2+1Wnutd+mll6pCoaCCIOgsO+200xSgPvvZz3aWhWGojjnmGDU8PKyiKFJKKfXNb35TSSnVL37xiwXbvPzyyxWgbr755s6yFStWqIsvvnjBeitWrFArVqx4wmP58Ic/rAB17bXX7vVdlmVKKaWSJFFhGC74bmZmRo2MjKg3velNnWXvec97VKVSUUmSPOb+jj76aHXOOec87pjOPPNMdeSRRy44b1mWqRNPPFEdeOCBnWXf+c539tktXCmlDj/88AXdwdv4m7/5G1UsFtXDDz+8YPmf//mfK8Mw1LZt2x53bCtWrFCAuuaaazrLqtWqWrx4sTr22GM7y9pd0uePbV/3yKc+9SklhFBbt27tLLvwwgvV6OioStO0s+yuu+5a0N08yzJ14IEHqrPOOqtzndr7WLVqlXrRi17UWdbuJH/hhRc+7rF1sf+ia5O6Nqlrk/5w0A3PPQl4ntf5u16vMzk5ySmnnEKr1WL9+vUL1jVNk0svvbTz2bZtLr30UsbHx7nzzjsB/QZw6KGHcsghhzA5Odn5d8YZZwDs0w09H1u2bHlS6cfXXHMNRx99NK94xSv2+q7dN80wDGzbBvTb5vT0NEmScNxxx3HXXXd11u/t7aXZbO7l1p6P3t5eHnjgATZs2LDP76enp7nhhhu44IILOudxcnKSqakpzjrrLDZs2LDAFfzb4jvf+Q6nnHIKfX19C87rC1/4QtI05ec///kTbmN0dHTB+apUKlx00UX85je/Yffu3Y/5u/n3SLPZZHJykhNPPBGl1ALX90UXXcSuXbsWXONvf/vbeJ7Hq171KgDuvvtuNmzYwOte9zqmpqY6x9FsNjnzzDP5+c9/vpdr/+1vf/sTn6AunjXo2qSuTerapGcG3fDck8ADDzzAX//1X3PDDTdQq9UWfFetVhd8Hh0dpVgsLlh20EEHAdqwnHDCCWzYsIEHH3yQoaGhfe5vfHz8KRn3xo0bOzf94+HrX/86n/3sZ1m/fj1xHHeWr1q1qvP3O97xDq666irOPvtslixZwotf/GIuuOACXvKSl3TW+djHPsbLX/5yDjroII444ghe8pKX8IY3vIGjjjoKgEceeQSlFB/60If40Ic+tM+xjI+Pd9zBvy02bNjAvffe+zud1wMOOGCvRrzzr9+iRYv2+btt27bx4Q9/mOuuu46ZmZkF382/R170ohexePFivv3tb3PmmWeSZRn/9m//xstf/nLK5XLnOAAuvvjixxxntVqlr6+v83n+teri2Y+uTeraJOjapGcCXdL0BJidneW0006jUqnwsY99jDVr1uC6LnfddRd/9md/9qTEfI9GlmUceeSR/N3f/d0+v1+2bNnvOuwnjW9961tccsklnH/++XzgAx9geHgYwzD41Kc+tUDYOTw8zN13383111/PD3/4Q374wx9yxRVXcNFFF/H1r38dgFNPPZWNGzfy3e9+lx//+Md87Wtf43Of+xyXX345b3nLWzrn6v3vfz9nnXXWPsdzwAEH/I+PJcsyXvSiF/HBD35wn9+3Dc1TjTRNedGLXsT09DR/9md/xiGHHEKxWGTnzp1ccsklC+4RwzB43etex1e/+lW++MUvcvPNN7Nr1y5e//rXLzgOgM985jMd/cajUSqVFnye/1bZxbMbXZuk0bVJj42uTXr60CVNT4CbbrqJqakprr32Wk499dTO8s2bN+9z/V27dtFsNhe82T388MMAnaJea9as4Z577uHMM8/c6w3iqcSaNWu4//77H3edq6++mtWrV3PttdcuGMtll12217q2bXPeeedx3nnnkWUZ73jHO/jyl7/Mhz70oY5h6e/v541vfCNvfOMbaTQanHrqqXzkIx/hLW95C6tXrwbAsixe+MIXPu64Hu+8PNZ3a9asodFoPOG2Hw/tN8/5+3j09Xs07rvvPh5++GG+/vWvc9FFF3WWP1bY4KKLLuKzn/0s3/ve9/jhD3/I0NDQAoO9Zs0aQLvhf5dj6eLZia5NmkPXJq3c52+6NunpQ1fT9AQwDANggfo/iiK++MUv7nP9JEn48pe/vGDdL3/5ywwNDbF27VoALrjgAnbu3MlXv/rVvX7v+/4T1rR4sum9r3rVq7jnnnv4z//8z72+ax/Pvo7vtttu49Zbb12w/tTU1ILPUsqOi7udkvzodUqlEgcccEDn++HhYU4//XS+/OUvMzY2tteYJiYmOn+3Dfzs7Oxe6xWLxX0uv+CCC7j11lu5/vrr9/pudnaWJEn2Wv5o7Nq1a8H5qtVqfOMb3+CYY455TDf4vs6hUmpB6vN8HHXUURx11FF87Wtf45prruG1r30tpjn3/rJ27VrWrFnD3/7t39JoNPb6/fzz1MX/PnRtkkbXJnVt0jOBrqfpCXDiiSfS19fHxRdfzLvf/W6EEHzzm998zBTK0dFRPv3pT7NlyxYOOugg/uM//oO7776br3zlK1iWBcAb3vAGrrrqKt7+9rdz4403ctJJJ5GmKevXr+eqq67i+uuv57jjjnvMMT3Z9N4PfOADXH311bzmNa/hTW96E2vXrmV6eprrrruOyy+/nKOPPppzzz2Xa6+9lle84hWcc845bN68mcsvv5zDDjtswcPxlre8henpac444wyWLl3K1q1b+cIXvsAxxxzTSb897LDDOP3001m7di39/f3ccccdXH311bzrXe/qbOef/umfOPnkkznyyCN561vfyurVq9mzZw+33norO3bs4J577gHgmGOOwTAMPv3pT1OtVnEchzPOOIPh4WHWrl3Ll770JT7+8Y9zwAEHMDw8zBlnnMEHPvABrrvuOs4991wuueQS1q5dS7PZ5L777uPqq69my5YtDA4OPu45O+igg3jzm9/Mr3/9a0ZGRviXf/kX9uzZwxVXXPGYvznkkENYs2YN73//+9m5cyeVSoVrrrlmLx3BfFx00UW8//3vB1jgBgdt/L/2ta9x9tlnc/jhh/PGN76RJUuWsHPnTm688UYqlQrf+973Hvc4unj2omuTNLo2qWuTnhH8vtP1/tCxr/Tem2++WZ1wwgnK8zw1OjqqPvjBD6rrr79+rxTP0047TR1++OHqjjvuUM9//vOV67pqxYoV6h//8R/32k8URerTn/60Ovzww5XjOKqvr0+tXbtWffSjH1XVarWz3u+S3quUUlNTU+pd73qXWrJkibJtWy1dulRdfPHFanJyUimlU0k/+clPqhUrVijHcdSxxx6rvv/976uLL754wT6uvvpq9eIXv1gNDw8r27bV8uXL1aWXXqrGxsY663z84x9Xxx9/vOrt7VWe56lDDjlEfeITn+ikNbexceNGddFFF6lFixYpy7LUkiVL1LnnnquuvvrqBet99atfVatXr1aGYSw417t371bnnHOOKpfLCliQ6luv19Vf/MVfqAMOOEDZtq0GBwfViSeeqP72b/92r3E8GitWrFDnnHOOuv7669VRRx2lHMdRhxxyiPrOd76zYL19pfeuW7dOvfCFL1SlUkkNDg6qt771reqee+5ZkLY7H2NjY8owDHXQQQc95nh+85vfqFe+8pVqYGBAOY6jVqxYoS644AL13//935112um9ExMTj3tsXey/6Nqkrk3q2qQ/HAilnk1Vp7ro4n+OlStXcsQRR/D973//ad/X5OQkixcv5sMf/vBjZu100UUX/7vRtUl/eOhqmrro4hnAlVdeSZqmvOENb3imh9JFF1100bVJTxJdTVMXXfweccMNN7Bu3To+8YlPcP755z9m9ksXXXTRxe8DXZv026FLmrro4veIj33sY9xyyy2cdNJJfOELX3imh9NFF138L0fXJv126Gqauuiiiy666KKLLp4EupqmLrrooosuuuiiiyeB/SI8l2UZu3btolwuP63VarvooovfDUop6vU6o6OjSPnsfCfr2qMuutg/8HTYo/2CNO3atev32vuoiy66+N2wfft2li5d+kwP42lB1x510cX+hafSHu0XpKndafnoV3yJoUWLyLKQLGjiT4yz8e5fMTu5gySaZGS4Ql9fkULRolgoksQRjZkZiq7N6OLFjE9OYNgWrTAgVSmlSoVyTwXXtenpqZDFKdKALFOYhoUQgpmpGcZ276avt4/VK1dQKRdZuXIlwyPDeEWPRYtGKFcqTIzvZunS5WzduImVK9ewZ2ycNIyJgpB6vUGxXCZJYkzToNlqMT49jumYDA4OYgmTuBWQRCHjE+MsXb6Uhu/z67vu4mc/v4Vt23eQphnFQoEXvPglTDaarDr4KJTdz1QTxusxzdjAcnpxcDASA0vYgESRIWxJkDVp+jUwFQ2/Rd/AIDMzVYLAJ01S+nv7OfTAg+ktutzwo+9x2vPWcvNPf8xvfvVLVBqyeOkyjlx7Agcd/hwMu8LYZJVNW3diug5KZmT4bH9kHSuHHD7/yfcSVcfwZ8e55+47SdOYZcuWMDIyQk9vD4NDg5imSb1eJ45ipJS0Ap9Wo4EhJYaUNJpNess9CCmJg4hmq8mOnbu46ec/45777mFoYJAkCTnhucdz6imnYBgGM9UqrltgeHiEar3Of3znar76z/+CaUqKpRKNZpM0TjEtUECSgOGapGECGRiGwJAS23aQ0qDVaqGUQilFphSmAAFkuQrQACwDwhRS9DZBbwdApQoBmBKkAWkyt5JCb0ugY+QCvY0s/26+0FAApimIEoUBCClIlSKdt5Jh6HGpTO/LNOY2ZhgCKSVpmmKaJr6fYJsSIRTSMLAsizRNkULQbIYMD/Vx5plnMrpkMRsefojbbr+LsT2TGPkxq32MMcvHbwPhvGf22Yj2sf3whz9c0M/tt4VSiizLUEphWRaGYSCEIIoikiRBSolpmhiGQRzHpGmGEBIpDQQGUhr6eymJk5Tp6Sm2bNnCAw+sY8MjDzO2a4xWq4lSCmlIAj8kTROEEFiWhWnqVhtJkubbTwHdR01KiUKRqoRUxQhDYFomQkhQgvYdrFRGHCdkqW4H4noeju3osVsWpWKR/v4B+vv7sWwbpTKSOMErFBkcHMB1XcIgYsvmLZjYDA0tplKuEPg+O3bsYHx8nEaziWVamKZECIkQAtM0MU2r0+qj3VBWqYxMKbI0JU1TkiQhSVKSNCFNEtI01edDCizbplQs4DgeUkriOCLwfaLcJpUqZYaHhli8eBFDQ0OYps3k1AStZoskSQgCn1bLJwh8fD8kjqPO9kEhJAipkBKkVBiGpFgs0N8/wMDAAD09PRiGgWGYmKaB63o4jo2UEtu2sSwLy7Kwbbvz2TCM3GPS9nAulCS37dW+/u0L7Wv96G3s616d/xto37u6WrgQYoHX9dH7fPR3j172dKHZbHL22Wc/pfZovyBN7ZNbKpWIogbbtj5MdXyMokhZsqSfV7z4OFJ/mj3jWwiyJikJcexj2yZLl47gGSajI8N4xQKNIODg0cUMj44yMDzE8MgwvX092JbJiqVLME0jn8j0jRSFCb7vY1kGpWIB13VxXZfZ6gzTM9NgSqaqM2zaupWZWoNmvU4ritmyYTM9pR6Wji6hWCpQnZ1hemaKnTt3glAUy2WEZbJzx27CVghpSpqEtFo1JmbG+dUdd3DfuvVEMQwMLWJ4eIienhKTk3u45bbbuX/9I6w86Dj6Rw/BNUuEiUQqG8vpQRomcZyhUgUiRSpAWEjHQ1kK1zRJpaQVRqxYtpqVo8spOh7j27fyvWuvYtuDd+I2dvDac0/lE/+/8/nvn/6Ub//bv3PTD+7nwfvXMrz8cA466nkcccThPPDQI0hpsGhkGbXJScK4ge1U8MpN8CdZuWIJAwMDLFmyhMnJCe789a/JUBx77LGEYcBDD62nv7+f/v5+fvazn7Hh4Q0MD41QKhQ5bu1x/P0X/pF3vO1tfPPb/0r/0CCHHnYYaw5YQ7Hg8YPvfY8kSSgUCqRJgkpSkigijiMc26Gnpzc37fqBti2LRIDKMtJUgQAVZ5RLRZqNFmmqiRFpPhF0JjT9mCRpglAgxZzJivNm4WrevZopIFMdYqQySPL1TCkx20ZPgSIjy7ShF4AhIFWahLi2hedYhIFPphSVgkWl1MNMtU4QhXoMEkzTIM0yBIpM6G24ph5UGEEcKUwjJckgiBNMIEgyPeY4gSBBovcJUK1WufXmm+nt6SXLEvrLBVaOHsHS0cU4lsXs7CyGNEAKCoUCjmNjug7SMLjlF7/kno07ntVhq/axFYvFvTq7/zaYP6G1J/80TTuTmOM4neVRFJIkCUoJpLSwDAfDsJDSIMsyWtNTbNq0hVtvuYWbb76VqelJTNNgaGgQKQ3CMCRNM5IkQwhNMqJo7hrNESkTKTRhypS+Z4TKkG3ShCBNFUmS5JO9iWWZJIl+XmzbwvVcTFMT8UazSZYpklQ/p4ZhEEUJQRChlKJUKpEmGX4QsGx0mMMOOwTHdtk9tpudO3dQq9XYs2cPcRx3mtY6jp0TDAfXdent7aVYLKDU3LURUiCFRH+cIxjteTxJIgBc1wEEvu9Tr/s0mg0Mw6B/YJCDDz6EI488kiOOOJxly5Zj2zY7duyg0WiQpgmtlk+9Xmd2dpbZ2Vmq1Sq1Wo1ms0kQBCRJTBSHtFoNfL9BHIdMTEzyyCObUSrDMCSO0z4OB8/1cFwX27bo7e2hv3+QUqlET0+Z3t5+enp6KJUKWJad3xeauKRpm8Bo0iulzMmVgZRiL0LTRpt86d+2SefCXnX640Li1b4Oj7Xd+dt5NGkTQnT64v0+bcRTua/9gjS1ce/P/ovSoj6couTwQxezetEgroo4csUovYWDsJwTKPR4WI5FBpimRcFxcQyJWyiSklH3A/oHh+gfGmRiepodO3cgBWzbuZPx8T1IFFEQYJkmtu3QbLQYH58gigJIU1q1Okma0PIbTE9PowSUKiWEECwaXkTBcbAtjzSICRs+UbOFaWjDFvpNhgb66OnrZXjRCMK0aLRaGBjYpsHY7p1s2PQQ1aCJsiRHHHcsPb3DtJohY7t2MLltM5OTk4zt2MRrXn8xy9c8h1ZSZPdMSrSnRaMZIZIAgUESJQhhEEYB0pJYrokQJbJIQZSQRQaHLDuaxQN9BNM1Hnnk16y/+1Z2rv815bLgvl//mLLYyZLBP+K8l70AtyL4t+98l4e3/Iqp5h6E6XPeH72JR7YmSNvBMCRCGvh+xsZNYxyyVFCrz1AsFjoP5caNm7jiyq8zMTnBCSc8j5GRYR566CGmpqbp7e3h4YceJk0zRhePsntsN7OzVW785c2UXJet27ezdcd2ZmdnOeaYozjkoENIXxyT5G93QRASxwmeV8RzPJIswzKM3IsjEZnSNDhToEC0eYvQBEka4LkOKoMwikizjAwwTIk0JGmWkikQc/Y3Jz0ahiFJ0zbJMrBM7dlJ0wzTlFiWQxBEIHKPUr79thdLCT0mlYEwQKVg2QaFok2j3iIGwiQmTWYJo6TjlSIDolR7eSR4BniOpOAYKBSqYNDbP8DQcD+7xycYXrSURn2WRaNLcQsllADLtBEqQylBT6WMISVbNm0h8H16ymVcxwbAc12yLMGPA00gDYNm5JOkCiUlpgl7qns38+xiIR79Bi6E6Hid2g1cTdPU3p7OhCOQ0gSl/5umGdPTU+zYsYONGzfxyCMb2Lp1C+Pj40RRgmN7QEaj3qJQ8CgUCti2TZIkpGlClqVkWZZPYrLj5VJKkaba85SkMRkKwzZAQRzFZGlKkmgvTqlUplQo4RUKxHFMq9UCwJAG5VIZKQ2UUlQqFRYvWkRPby9KKcbGdrNnzzjbt20nyb1B9XqdrZu2sXXLFnrKPQS+z/btO5iZmUER43pz5wMgjn2CoMHsrGJ6ZgIrJ3RCas+JmXtwDMNAGoY+RimRhh5Ts9kkyzI8z8WQBlEc02w0CIIg9/4IduzYCqTs3rMLz3WxHZtSqaTJrGFi2TbFkku5soTly5d0PFz6mZ7zfEVRhO83qdWqjI3tYuPGTezYsYPZ6gxKSVSmSfHk5DRxFINQOI5LwSvguDae51HwingFD89zcR0Xx3GxHQvHsbEsTa4dx2JoaIiBgQEsy8Z1nfw7KydI+tpm2UIio71xyYL7UXs050jRQs6hcoKqf9/28rXRJm3z7/d9kaf9FfsVaTr7hYez7JAVjC5fzAErlzHSU8afmcBD0eN5DI0MEMYRO3eNUa3ViZOYielpQr9Bgr6Zwzgl27QBQ5ps27GNhx56iCgM2bVjJ416FZWlmjQZJr09vViWRbPRRCBwTZOw0cC2TUrlIq3AxzANlixdxkB/HyKOMW0HM83o7+vDsmzIUtJ8aq2UK/T0VdixayePbNmoJ8FUUfAKeI7Lnond3L/+fmbqNVpxxNDiRYRTu9m9e5Lt27cxO6XJm2FZDA70Mb1nJw8/MkFqDlLpW41lOMSpIIpDlFCYjkngB2QJZE1Fve5jGS695QrFUoVl5VG2PvgA6+77Fds23kMwswWvmOIWJa3qLDfd9H0m9mzklNNOAgu8SoZVDEBMsOWhX7Jn7ASk0cSQZv6GbIPhEkQJpUoPwpAIJZiYGKfS04NlO/T19SOkDkEYhsGqVatw3QJ9fT0sXbqU4eERSoUS999/P8cceyxvf8sbGR4c5uxzz2FyYopWELBy1SpWr1zJiqVLuf/ee3BMi8HRQZYtXU6xVKLS00u90aSUu2SlAENKsiQlSzKQ2gi0p64gCJFSP8ipyhAGmIZBnGb5srSzHEUnLCYNMKUgSRRGHmvL2nErzYi0lynNEDIjVYooyRaQrQ4UkIKU4BUtHMumr6fCQG+ZSqWMxOCgAw9g99geykUP13FIMrR7PI/DmSKDLCKJQvxWk0azTiYU/YMlyuUCUxMGI/0eSVDHtixcxwIpsEyLtB1KNDRBsmybJI2JkphGs0V1tkoUBYRRyNT0NJlKkZZBFKcEgQ5BOB5E8bNT/P10oT15zA/VSSmRUnYmoza5MU2JUpI0zqjVWmzfvpN16x7kgQfuZ+PGjczOzuTeHgfLMskyTUh83ydJEyzLyvcnyDKVP7MCKecmMO25SDv7llI/KJnK8tBTRpqkxElMlmYY0sRzC5hGTBwlRFFEmqRkqY4VJ0lCGIQEQUQpTkEIkjilXquzc+cubQekgWVbTEyMMzmxh77eXoQQ1Kq1PFQmsG0D0xRk2dwY4yQijhOCsImQUr8MCdE5JiHniJJhGBjSQBp6Pd/383NlYxhmHmbUIUrDMKg3BDt2ZkzPTGrvWBxhGiajo6OUy2Ucx6FYLFEqaW9joVDA9fRz6TgunuN1PEKmaaBURqPRoFypEEX6JS9JEwzDwjJN0iwh8JtUq1XSNEEppUmfqUO0lmlhWhaWZWJZNrZl5x43B9f1sG0b13MZGRliYGAQ13UpFosUi0Vc18W27QVhtDZJNgwDy7I16cxDnzInnvsiNm3vkyZMCqWyDtl+rHu7/btHL9tfsV+RprPOeT79Aw6GVBjJBLPTE1hKEaQZszO7eXjTetav38CvbruDHTt20QpbtFpVfL8+5y6ct7041g+5UhlkCtu2UGmKZZqoNKM5Pc3w0Ai9Pb309/XSUy6TJjGWKRkaGSJTKUIa9PX26wkzTRECZmdnkEhKpRJBGLJnzzgTU5M4rk1Kyi9vu43t27fpbfb24Td9JicnCLOUYtFh0dJlFCpFfn37bbSmmxR6enE8l3K5iN8Cx/a48p+/gmv1MjOrGFx2JEccv5iiV6DWCkkznyTLIE0w7ZRCsUAcxQiVsGTIoafQS3Nqhg13buTWX/6E6tRWHMenVIwwzBZBs0Ffr01W6mPLlo1s3LSBTGakUlAo92GYGY3mJPfc9UsGhg4gzASGMrAMl1RGlCpD+q1neAlSZVTKFRYtXszg0BB9fT2kWYppWfRUyhSKBQzDplSukEQ+CnAdj4k9exgcXsTZLzsfv15DoCi6BTIBM7Mz1KtV+vv6QCnGx3aTJCmZkOwZn2BgYIjegX76+vpQ7YuuFEJIDKmJRpwBUmAYJikJZArfj5BSe41AILKMJE0hYU541CZEQks72h6iNMkwANsEKTIMdDhQ2gaGLREYDPZXQJooBKZlYDsupmXl4TltrA1L4toWpiGoFIuUikUa9SpSCZ5/0nPZsP5h+vt6KRVLoCSZSonDCJQiDkOajTrVmSkm04xWs0UUhOzZOc7WzbuYnA4I6rM0mjHr1u0gkxJD6PBlgsIElNAULE0SHZZBQO59QGlypARAikhiokQzQNs2sG0X2zJo1MPfj0F4FqEd8mjrmIQQuadSh1y0nsUkSxXVZoNdu8bYunUrExMTZJmiv78/ly+EhGHU0Z0kScz09DSTU1NUKnqyT5Kko78xDDGnacuJWpu4ubYNEsIwIFXocJxhkhgJabNJkqSEkQ77ZZkiyxRxnJCmGVAnjmMajQZCCMbGdjPQP4jnFWg2G9RrDZI4xTJtent7WLxoEbXZWXbt2kkY+gghSNKIONH3tlIZhmnM83gIpAGOYXa8Ipp46nXTLIMsIZ7nHRYIaHtS9AJafkx7VhBCYEhD6yiTkJnZaaamJ4njiDDU88TGTRs618K226TFxfMKFIuFPGRbplQq5lIOj1KpTLFQ7Og4hRAUCkV6Kr15CE0SRRGW1cI0zdwLmHuFVEaWZiRxigjCBWExKUVHE6XDcQIhwDANCoWFYykUCnieh+u6FAqFjufR8zwWLWprtozOfdi+Jx99j7bvDyG0nRRijty312+vM/937X/70lDtb9ivSNMn/uKvMI1aHn/NwywChNJagCzNCPyYViMkTlMykaLSiCzTbwkAfQMDFLwCQRhiSoOyV4QsxbVsBgf6SMOI3p4KWZrSaraI44Ta+DhRvc6M61Bv1IiSiCVLFlPuKWPbNrOzMzi2Q5ImxGHC9u07WLViNcVSmYce3sBDDz+EYVuYtkmqFNVGjaOOOYpXvuIVHHvMMWx+ZDNXX3MNj2zZxMvPP48LXvt6gijkT97/Hu6790H6+vsJmg3CRoO+cplatcaWqXX8yXv+gkXLDmHTrhZTjRlmp6ZJhInpODrUIyKWLO1j1crlOLbEM5ocfdAgD97d5L+//zNu/9lPmJnYSdFT9BQMFC1qjd2UioZev1yhXKjQbPoEaYS0LUy7QBClqBSmd2znhOeczkRVMNNICIMUkUj8IKPZUiSpgSUEg4PD2JaFUIqVK1ZSKBexTJMojklyd/bszBTbt23DdV1WrVhJtVrlll/ewgvOOIPrvvtdAJ573HEoBT/4wfe5+667ePWrXsX2rVtIk4TZao2NW7YQJxknn3QyLz33XEzT0oZSaW+SKbUiO0rSXD9kUHBdWs2WJnKGfiOM4rTzNqX5k8S09VtfluTiR1No71QmsERG1EoxFXgelMsuxaKXu9AdbNfFNG16e/vwiiWiJMN2XEZHl9A/MEArCGi1WhTLRZp+iz27x9m9eze2ZRJEMXfc9QAig/vXbyP2A8olE6szsWY0gwyRKcjDAyKNEWmMShPSJH8zlIqCzKhPJTiWQEUJSaqwLQPDFARRguu6SMMkjGPCJCZJMwzbxHEdsgySKMI2LRzXpdVsEscJZu7tyvyURssn6ZbK/R+hPVG1vUzzSVPbK5CmKfV6k7GxMbZv30Gj0WJocIiRkUXak6oywjCgXm8QhmFOoHxqtRr1eg3fb2kPUZYRhiFxHHRImVKaYCVpQpZq3RPCplQuMTw8jJNr1jzXIwgC9uzZQxRH2KaNaWgdlGu7ZEmKkBLPc3EsmyxJ8YOARq0BqaBQDAD97NmmrbVVYUzgByRpjBAZSRJiGNqzpJSelIWRdcJF7cnXMLVuR2UZWdbW87Qn53kanPY5bv8lBELqkKMmKPp5NwyJlJBlEMdpx6PVDkeiFI0kRs/5UieO5KSl7QFq66x0aMzOvUIOrquJixCCZrPJbHWWVsufR1I0ee3p6SXLUuIoJopzmUCakaZzHkBNUuYIiCZNMtdTJiiVdUTy88XkbS1YmzS19bl9fb1UKj1Ylg7jzV/fdhwc28F2tHerHTa2bZu+Pv2b9pjmk6pHh+LmE7H5BOrJYe4K/nZ4dLrKU4P9ijRt2vAQphVi2RmQoESK7TgUCxWCVkRvqY9KsYxreRTKJQollyCoE0c+o4sXU683CaKIYrGMFAaGkKgkZmp8Aoliao/+bxaFpHFCkr99hWHE5OQkpm3ieC5uwSGIQpKqjgXbpkVvby9ZBgWvxED/AJZj0Qp8LMfm9BeewTkvfzmVnjKpSlFC0NtTZvHICD3lCgcefAhHrX0OjWaDxYsX09NT4aqrv8P4rjH8ap2GXcUyBFHTx+vrp5XBkuFFrD3mUJ53ymls2R3yi9u3cttvNuB6NouWDlOslBhaNEip5JCEdR568E72bN/A9vsrbF63lZJpMDqU0OPaZGmLNKsTZw1KBahUHLI0IAospPIgExApUCl+q8psrYHCZkzdR2t6HMsaQgpwXIdEJYyP76Z03Ap+eMcd1GenUGlEb28PWZYShAGLF4+CgHXrHuDhDRtotVrEUUS5XOYFp7+A3koP9993P1/+0uXc9qtfsXPnTsrlMsVCASklDz30UOeNbe3atUxPTzM9O8vo0mWUKxWOOOIoRhYtIr3/PoQUWJahM2iUfqBNKTCBFEUShhhCh9kEWkxtCUEmACmwXQuv4LJ4yWL6+vtoNZskaYJpGtimhWfbNOs1pscncG2bgmshhDbiUgpsQ6DSmHq9weT4OGGY4YdauFTwXCzHJkq0R8syDcI0JQgCWn6AaUokivpsi8HBCrt27MKQMD2jQ46WJTFNSYbEtkwc2wKZkWUJhqmQpsByDFzLwRCCmZkqnqe9qD1lB78Vo7IUI5PIBEyV4De10bVFhilARSlpGmCY2likYUoctzAyhVACU2hVfAbECgqWIkieUTOx32F+eOPRE0x7YhRCEPgBExMT7Ny5k4nJCQQmw0Mj9PX1USwWQEAURbRaTXy/he/rzK40TYjjmLHdu2g06oD2HrVaTYLAJ060ximJY8Iw6EzWlmlQKhbzkFQF0zKplCs0my0MaTA9PQMoPNdDCokpTWzLRghBqVTGMLS2qdXyieIIKQwkBqZlojwIg5BGs0GjrrNms0yH4jKVIFWGZUukMecFk4b2EKVpHhaSEsOEJE5JM50ViNBeEAFzKap7oe0u1iRrvm5HkZKkWce7pfLsU8vWL2DtzDiVv7ArlZEkEUkS4ft+Tuy0DlIKnVKkMoFh6JcNx3Ex8kzWKIqIkwTDkBQ8TWLK5QoAcRwTRqEmbIm+fnEcd/RH7bEJIZBChyL1cZsdbVwYRgu0a21C5DjOvIw8s+Op0t4z/X3B8ygUihTz8GOxVKRYKOJ5HqZpUiwWiaJFeJ63z/v20Z5TnQUqUbmD48k2Ilm4XvsZefQ6j/f7J7Wb3wr7FWlyHAPTsDBEpF21lkGiMjIylIAgiMi0/JUsM8kUZMQIAyzPYffmTZTKPVTrVWrVOgaCguvS9FskYYBKUkxD4vstsvxBtR0HhQDTQElBnCXErZQwjbBtnQLqFVyKSQmVgtfvkdkZURTTbAUIQ9Lb349Ccfsdv+bBh9azcvVKXvrSs9m4ZQt33/kbDjjgQM466yWoLGP3tm385Qc+wC233oZqhqxatATHdlFpTOwViap1Sq5L1Gryta98mauvvZYjjzuZZYc+l4e2SFpBRLlgMjW2jbtu+i9iv0oUzDIxsYXazBiuTIn8Br2FImGroV3wZFgmWIZDEMckSUytWkdmEUW3D9uw8UoVDENQrc2StmbJhEWrajC+czPuiInj9mFYim1btnHHXb/mgpesJEl8du/eydbNm3MLk1HpqeA4LlmWEUcxhx58MEuWLMUrFojDiMGBQZIkZWBgkAsuuICjjz4aaVnYjsOi4WGklKxcs4aZ8XEOPPBARhctIk4T4kwxOz1LvVZjYHAQ13PzqJzCMAyWLl2CEGBbFgXXwrIMiuUC/b09xFHE9i1bGBwcJo1jlALbcZianiZKExrNJo3JKTJfi0QbDZ8oiLANA8s0aIUhvt/CMAxsy0DKjDRre6t0PDAKE+IkBiWwbA/LdqiFAXEU6RIXtk4Bz5RACYEKI2ozLQwpqBRd4mbAyGCFSrnM7j1TRHGIMARJlpGpjCTNCOMUoRRpmiFRWFIhVUYz0KUKWjH4KiPOMow41il9mSJL9FgjPyFNtVGwJPRXbPoqRTzXxnFM+vp6KJVKCDIqvT3UazWK5RLlSoVm4BOGEaZl8rff/Okzayj2I+wrVXs+TFOHgQAazSZju3YxMTGpkxeEYGp6mkajiWmbOkSVtQm7xDANypVK7gFI6evvQ+QhnTgKafktfL+Zi75Toiik1WzSbDUJgwCF0uG4JKPR9PE8l7QAKhMkqSJJMgzDpFAodzJY/SAgCmOUyvBcj6GhRQghCMOQIAiJIp21Zhp+/nLga69KGCENpcPVKiHNtBc0UwrRyeDSRCFOYtIkIVMpClsnaZBqIThAmzS1/0/pOWEuE4zcc5R7Vcy5LMQ40fZvTiRvYJoGRu7ZbYfO6GyrXZIEreFK8+9ila8itYCfiHqjgWGYeJ6LaZooFGEYYpkWpWKJUv5iGIYhpmV1Mvvanpw0axMoHV4No4goDInCKB+3Xp7EMRk6ZJokGVmaoAApJGmSEUcJUvo5ycyF29LIRfQGpmF07jvb1pl9bi5Atx1Hz3mey9DQIMViAcsy8TyvQ8baHi7HcbSIfZ5ny8i3PT+TUc0jqI/1HORL20/NPj4/1t9PvYZqvyJNtmlQsA2EEmRERHFCEMVEYR2VGSANUhIMCX4rJVZN4jRGINmxaycz1SqG7UAGzVaTJIrxHSfXbQBSaNG2ACVzb4PKMEwL07ZRZERpTJZkhHHYieu6LYdmtUESaxFmpdxDkqSajccJD9x3P5u3buE3997Dtm3beM7xz+HII4+g6BVZ/9BD3HfvvTy87kGmJ6eoTk5y1b9dRbXR1LFo28MtSVzTpIUBcUZPb5kwy7j37t8w2/CxShWOeN5z6e818Mea7N7xCPf8+tdsvO9uDDPFcVJM0WSgKBjoKTAz2cCv7SFNEyLhkKUKw7KxpIXCII6h0YjxLBdlK5I0Ikwj4jQijAIOO/BAzn/lBRhuDxMtlz21WYrDfdiuQaNVZWZ2AsuEU046gSMPPZjJqWmiMCRTip6eCitWriRNYuI4oVwu4To2tVqN++9/AN/38VyPJUuX4JgmrusyNj5OFEXc8vOfIw2DAw48EGlIrrnmGp53/PMYXTpKmkGz0dBvWWlKy/dBKYYG+ukpl+gf6MOUoNKUKApoNJpAhGvFRK2EZsPHlrMkiZbt26ZFq9YgShP8VoNavUnTqVPwCjTqPkEroOA4OI5DM/RRhiDOUrTWVWk3eV7QSSLJUqWzKFNQSYoSsQ6nJSlk+q1MZTqDTSEwkwQ7yyADESVEQYSBIJIBJCkihTjMiHNhpsoyJAmeZ1D0TIrFAkP9fbiWw9TENLXqLCcceyyr16whyTKkIUmiENex6entJYlCWo0aRc/DkpAGASXXoFx0UGmMymIGB/oolzxmZqfo7etnanoKaZoUigWt/VKKaqP1DFqIZwc64adOTR4w8qyv2WqVam0WIRwEgmZrJj/32ZzWxRBaZ1Pw8FxXE4ZMUekpaW1NLhKPo4gwCsky7RpMU53pFYYBYRjq0Fkc0Ww0UQi8QoFisQfbKjAyElDwKkgp6B8Y0lqpKMayNREK/RDDcvAKZRzHJk1SwijQxCkvfyAME4SBUoIoTrDQHppMu3nJlARhdLK0VD4hZgqtWUqSduUO7XlSuacpdwPNTZ8q/99cuC7NEgQSMBFSIHIvlui4qPK0eqnLFwipl88vDwLopBIlkGqOBLQTPYTSG5SGJFOCNElI4gQhwVa2rreWJZAKwiTCz4mq7/u4rotpmUjDQBhSyxukpoUK/WKUxLG+XnmZlSSvl5UmKWmq9xXnnqmOiD/VpLod+k0TnaiUiHaYFuaTjnYmnGkaGIZ+sRNCYNsW5UoJ13VwHAvPK+jMPtfNiZaL53m5rkoL0jWxcjthQCnn6m1pj5fWk0nRDnsaefLCfGKV27sFGaiPfn5AexLlvGN56rBfkSYpBCKDXLeNEBJT2MSpvpls08YVFhBDmhAFIanKQJiMj4/jOg5hEGpBo2WjMkizvNCZY5MmOqXWtmydGp5lZDIvqJanu2bZ3IMXRzFkKX6jRdDShsAxbZYsWkIYxlR6+5itVtm4ZTOGbTE1PUUaRuzcup0ffO8HVCoV1q1bx9bNm/nJ9T+hOlvVcfNmk2VLlnL8c5/LuvvuZ+uunSyu9OGZrg4NSYtIRcRpjFIpvt8kiQNGhstMjk+w8cEH2brhHiQ+xYKLaaQUbY8li/roqzjEjQZpK8aQFimCNEuJwogEheuVcB2XrK9AHKa0WiFB1CJLAvr7ejj00CN5zjHP5eSTXkBL2dxw+8M0ZkMcZWB5HrbrYJoCxwLSiF27dlCvt7BsG8d2aDSa3HP3PdRqtU52TxwGgGLLlq1UKhU8z8MwDDZsfISdu8ZoNVusWLEcPwyoVWsUCh6jo4vZtm0bi0cXo4TC9wOkYdDb16ejiXHUeVtEKGq1KiqJCXy/o/kwHYlXlGQh+H7KxPQsKIHKlHatC128T6FAmqQp+VuadjkrNfcmYwhBpjJ9T6mMNMkfXqFLjKoMyBRxokiyACFCHQrWvAgzDJGAkZtyoRRlAbYjSJOIHk/iOII4bLCoZKKEheV59A4O0ttbyUeR0ddXxnW0wHvRyAi9pR7Gdu5i5/btHHnUkRx04IFMz0zT29dDGAVARm9vD61mg6npcYoFj2LRIYl1yCFLY8JAkCbgFBVKBYR+jZpIydKQZjOkVlVYefG9sbFdv1ebsL/j0W/U80N1hmGQpilhGNLb24vneSRJQrPZwjQVpuWQZClJlqCyPDsBkKkgzRRhHNNq+Z1JxQ8DHLuuM6XyjLksyxASnUJv2ZTKBcqVPMySk60sUwgpMU0Lx3bIsozhxUvxWy3iKEQBYRTiixZJBq6S2JaLYWpvfxTp0JnjFrAdj1azRcsPkNLENG1MS+vrlEgxdQ4nSgjSTCAwkYauG4VSZCikYeZFVrWHlfxFt13NVaBAiVyXpR9CIefVcFICU2mylpERpzEy0yEjhUKaEqlkp9aTylP1FYooiUCRh9/ybDQp9Iu8yjpJJ20vhwAyUjIlOqQnVanejsg1VTkZbrS0ZzlOEkxDl7wxpMQw5zLmLNvGtq28ppaF5bg4ntcRw9uGiSGNOdsax0SRJlda4xYSBEH+Oeqs066zpQnJnEZJh3Z1gpMQoSbxQmFISbPZ1CVZOuRnrm6XlWf3OY7Wd9m29kK5jpvX8dJZe8VigXK5QqVSxnVdbUdMq1Mmo1Qqag987lFXWbtkgr4/hQApjHx8HcV/h/w9Hcl6+xdpMkzSNCZLFMI0saw8YyJRSGVgGTYmkiyNUUmCUok+0SiSKMY0LFr1OtIwyVIdN04BZRg6Wwj9xh5nGULqIoUqzUhJMIXWvCRxkr/VtQXoWiCrUoUlTSbHxiFURGHM1J4p6s0m9WqVUrmkdSAIdu/azS2/uIVVq1dRLlc44ICD8olaP4hSSp5zxFG84pWv5Kaf/jdfuuJrTIyNsXxwKWEY0vJDmnHAopERJmemabVq2FJx9GEr2PjgfWx96A6C+gSuZRE3m/hRndS16HENspZN0BI0W2BZHsKUSCMlSQLSVGCaHipzKRZLTLR20KhX6R+ocNghR3L82udw+OFHIc0Sv1m3iZ3TAdtnEmKjjwCDVJq4hQJJootAPvTgOn70X9fjhwmLFi+hUimTxBFhEFJvNIhjHXP3PJelS5cwPDyM53lEcUxPoUDfwABNP+DgQw7lucetpdzbx/ZtW6jXqqxZvZpGo8GypcsolIo0Gk1d4NHUQmYHEIZgemaWMPTpKRdIwoCgFRAnCUpobUTUUog008Uns0x7hbIUEHkmizbDCZI0SonSSGebKYUfRzTTBJUlOJk2ipZtYJg2wtLGzvM8zHbdEqWwXRclBSrNsAyZpzwb+p7KUgqmIGr4VGdmcWVGT8UhCAKKRRvbc6jVatiOjZKSYqXMAQet4sCDDiJTKb7folAsorKUZr1GuVSiUiwx6CxmpChxCQindjCxZQtBqYhpG6RZwsx2hR/51JqzCENQKHsYjkkQBQSRr8M0wmCm6mAlgmajiVWdRVoGcZyQZAmGaaKA2drMM2ki9ks8Vjp2+3Oapjkp0kUigzDETAU2sqNrwTDmvB9CkCmdyZbkld6FEIRxgiH9jgerTZyEEJiGkU/IdscTQJ727roFDNOkXavDMAU9PS7lci9JHBGEejIuFEIqcUymMk0ogCzPvBSAaZkgoO7WSdKUpu8jpaGrcCcRcRQQxxFtR0GSJB3PjyLrhNny6RKldGRgfjacyP8QeSbZ/FPa9liBwBQmKH3+siwjzQtD6gui/6NrqWWQzl2jBQLmeZlsSuXeLKWJZiccKPJx5GGv9jpJlsyF+DJFnLZQ/ryMs0yv1/b0WKbduT5apK0F2u2XlXZILbUdHMvKhyfz8hMOhYLK63NlHc9Tp3RDrpdauCwhSeZ0VO2K9DqjT3updF0uXQRVtcmiYC7EZ7WF6HZH02SZmvBJQ4dGC4Ui5XKJcrncCf/ZtkWpVNIlenoruK6Tl5AgL4dgYJkmlq1JvCZklk4MEEbnnsmUIrKeeoqzX5Emw7CQQiGljeOaWJ5FPc8IMQ0LyzAxMlBKIJTOjZapQhqCUrmX2WoVmSm9DWGgpO49YZs2aRpj2w5K5VlUCIQpQErt1TIMDAxiEeVMXruGwzgmihOGenvpKZaJA107pFKq4PsBQaNFyfEYGhikTA8jZCAEJ518EieceCKWbZFGMaVCEZUp/HqLgaEhlq1cwcGHHsLI6CJ2jI/xD//4Dyy3LEhSZmdnGF6yiLNf9hJuu/vXGCLDIqLHUbQmtzK5/cG8+q9DsVxicLBMGEbsHttNX6UPx60gzRikhTQNTFuhQkWShvitFo3GDFKYFIoui5cMc/Lzn89LXvJCFi1exPhElUe2TvLwzik27WqQeP24gxVCJcmkxC04NJs1UOA5NitXrmDx6AqWr1qNbVs0qjVsx8kfXp3uazsWfT09OqsnivSDY1ksX76cRaOj9PT047gutbquhN0KfMYnJjBtG8uxqfT2UiyXqdcbNBpNojhCGNotHiUJtBQlzwIyhEgxpS7GmCLIUoMsyYiTDKGSThTcMiVpEiFk7o42JNLVWUJSKWTRwrBMTCmQWUbBhELRxXMsiuVCRxMwPDyMFOB5Llma0dfXh2FI4jDUQttSEcexSJJYe5RKHtNjk9z7m/sJ/DqDAyVC36dQcoiiANe2EIYOGSjhE/sT2GIJSiga/jTV1iRKpQS+j1+V1GxH65Yin9l6iEGTsDXF7h0P09vXQ7FcIoxD4iwhUTEiE8zM1LEKWtQfJiEIsIRBmrSwMy0YjbIQI5EgJYYUpMTEcYrj7Fcm5Q8W8wmUEAI/8KnWavi+r6tjYyKMmDZTEB2iMOchaQvI217TtkgZkZOFNMsneAiJEX64gAQolSGlgesUkHn2cbsQ5nzPgmGYlEpuHlLRngcjr0athC5focXLmvQVi0UdwvE8fN8nCnWYsFarMjs7Q5omnRYoaaqLbKZpQprpkq4qr/8h8oKNbb9Yzkh0GRAdN1sgaVFoT7CWqJoYGHmpBB2+aBd9bOttHl2Usd2+xZBGrp8SeWg8r8ydZTrE3iZNKJAy10QZeoS5l0QTQDpide0xyUvxCl3wMsvaNeW0FyuOo5w0+8hGXkdJdJRcyPxFzzbNnERZnYw+/W8uu68dDtMFQ+no3vS8kXQ8UUEQEoYBvh8QRWHunYpzD1QuNFdaK6VJl95GmmjiFZkxphljGgbSmDtfbS+RIS0sWwvQpdDEyrQMXcKhUKJUKlAo6qKetm1h29a8sJ8mW506Wa7bqYslhI4M+S3/KX829ysLl2YKie7p4xU9Kj0VlJCEQRXHsJBCIAXYlk0mIAl1obVi2Wbp8CLiVkCxXEFIAyWMXAsAxUKBNIlwHZcg9HEdmzjWWW6GbeGHmigZUmiGaxk0m3WCwEeaJkWvwMDgIEMDwyRhRBIlDPQNUa/VidOU3t5ejnv+81h96IH5jZNx/InPp9zXy5133kG91WTxyCJavs/23bs46KgjMIsuVb+B11dh6ZoVxEAsoR74BFHCkiXLOOH5JzDTnGXz9m1s2rCewA/Zs/0R+is606Ovt58Tjj+Ro44+kgfXr+dHP/wx1WaDnsoAllsgjhMMBKVyEcNMmJicIYoDHNcFFfPis17Jc487nqOOOpxKX4Wbb/4VP/rJzQSyh8LASsy+HlKzSGI5xALsgoNXcKjVdiGEYnhokKOPPpqDDzmKvsFBwiBkYs8eTNMkikPaQsooF0P6rSY7tm1jbGy3rvHSrDM4NIyUFnEUsH3HDqIkZvHoCJs2bWR8fIJao84xzjFUKj26zozK8rcd7aOVAqRQRGGEIcGQ2rWPlGQIogwsaePYip6CRRKneK5JpeIRxQE9vZW8ordJb08ZQwpUmlD0XP1ZQBzEmFJQ9DxUllEoF5BSG9K+/j5afksL0zOF4xmY0iCOUpIkwjRreiIjwrQy3KJNZdghcxW1WpOiEAQEWDKjFlUxi4JUJViOQbXWZPP2LfQOFbFti3qtntcKU8RRRBLFuLaN7di65IJh0IgN3D6brJXiCx/XLiAtA5lkWFh4BQ8/r5NjGhaOqQMhaZSg0gTDsnAdG99vgVQgVC4Ylpi2TUT0TJqI/RKPlUk0v11Fo6E7ENTrDeIkxrCcOS+KyFv7yHb9Hpl7l/R3SmgykaTtrK+2hyT/O9fJZPlknimlBeVKv1wmST6hC7GgWGSbNNm2mddwspB5Gr72hlianNiaaAOd4pulYplFixZ3yGEUxXkPt5b2Rkchvu/TbDU64nTfb+ahpVB7PrKcVMUxSRrpauWZ9oSgEhBZ7p2Q80gkICRCGFqgLRRK+6sQ7cy7trdqXjbYfOEyOVmaa1+yUMis68HpF+42sUnTpEPv6Jx3HUZUSjBXLFKH7EVe1VyL2ufCpVmaaHF31t73XI0soXTYzDRMDCE6bW465QNyUXe7+KkuH2As6HvYPnYpRa5XslGq1MnC08Qy7Xio0jTRpXbywqZRHBJFMUnuqdKlgBRhEqPQGqok0V7TttdItAusKtk5/20Ple1oomTZFpZpdGpjeV4hr0VVyGtQFTpC9LYeql0B/6nGfkWaWr5P2dONRlthgOFrN6QlJQYZSRQAQr/5pookyjQDDiKiVgtDgSl0tY4kyzCk0B4nlM6eyKvepllGkqWkSvciS+KIWGmNQZQmpGQEcUICHHToIaxZtYrZqWkMYdDb34dtOpjCoG+gH7fogRCsXrOak08+mT3ju6lWZ5FCMjszy9TEJLMz05SLRcIoZmJmirHJcTbt2kalr8LU5AQ/vuEnuK5LLWjhJwm2WyAII+644w7Gxnbh2RYqDBFJzHHHPofnPed4lFIM9A9w6qkns3h0MYYl+dnPbmLn2DhRlCKESZwkuG6RcrGIbSdMTKYgMl728vMoFsu87LyXY1kuD2/cyAPrN3DLr+5h845ZRg48jlXLF+OVCqSJwE9BZimeZWKakrAREAYxvh8zM13lv2+4EYGRV71t0Ww2aNRqlMq6/Uyz0WDZ0iUMDvSzdfNm7rvvPqRp6oriQuC6Jd2axLap9PWyfPkKVJbmGowyMzMzug5X/mabJDFZkmCZkrKnywAUPBfPsXEME9e2KBa1dmtqdpa+ngpZ7LNk8SB+s45hQF9/iSBsMjTUhzAUgpSenhKolCyJqRQL9Pb2kIQR9VoLWxYwpEO9XkUYkjSLCGKfVjVmamYG07JAkjdtNpECnRoethA5+ZASenr7kIbLjqlJZmozOH0h1WqDULo0WwGDAxXSNEZZJsJRRFFII6xRkB6mK7EMC1SG39Q6ErtgYDoGAoVTcqm16hRLJQaWD2sdg60wTZuoqV8gMgVCSWQmMIWBKSVKZQRJRqo0eZQGKKGFvFmmtR6mZYEQnZ5eXfzPMefpUZ2Ck/V6ncmJCWq1GnEU4bhaIpBHgDrhKJlXu9dkAFCqs147BPPoVHBpyEd5HHRCAkrbRyHbSjvybFDtpYrSiDiKabWytpOHLMt7qrm6VpnrOnlLjfY+tQfEqxRwHZdiqairz8dxnnGlva5BGOD7Le09rteo1mrUalXq9RqNZgO/5efC8lb+uZVnlAVEYUAUh3ll7RSV6eSKNNPeJ4HQ1fRlnpWWzhVs1AJlo5OC366Wro8tJYriBbWI5lqSzBHWtnBaGjocmSQ6u60dzpvriSc7+pv2uWt7fAwMbQMFuVdKh0vn6h+1W3vrsKPMCbIWrKd5f8AIP4C2Mr1d6VuHZtvEWofJLNvKdW1mJ/zn5iUSHNfF9bT3v02w2jfV/FpWugioDtX6fkAQBAR5Vm2cZ2XrsF+Se+20R6gt9G4XJtVRiJQk8fF9Xy9TbV2VnHeN5ryeHZJl21iWvnY9lYq2S08x9ivSFAYRQwP9KAJ279lDtVZj7dq1lA4pouIUmUGr0dIakEKBYtEjCBpM7BljbGwXqZCM7dpJLYyIsrYuSZLmNTuE0Hr7Tkbp40AIQblSZPmqlRyzdi2/+PnP2bljB73lXlYsXU6z0aSnXMH0bLZu28pNv7yJe9ffy623/Yo94+Oc/oIXcM6553LowQdrAbpKEYbJ6PLlNAKf2+/8Ndt2buWWX/4cIxOcdvopbFy/hd7+QUwMxscn+Y//uArshDPPOJPjjn0Oo4uWoTB01eBqjWarwczsLDf97CZ+9vOf47fquJYkDptkKTiui2dbOJaJaRUpFEoo4MwXns7o6FK2bdnFf1x1Hb++/dcUekZYfeRpnLr2SKQ7hOUNsH1sksw0kKSIOCPJQtIoxVQ2k9M+rUARRhl33HEX09PTLF26lMWLFhGGEZbj0t8/gGWZNAoevf19HHXMMRxz7LEcu3Yt/QMDLF+xgmKhQIZkdmaSYrmM43nUarPUqlXSJEEaBrfecgthEDKyaBEoRavRIPBbeI7F6KJ+Bvt7WbViKUXXwiTDcy2GBofo6+1n0+ZNLFuymKmpCVYuXUJtdprZ2Sls1yRTZQoFhyBq4vsBzdoUYRiQRCEN22JmwiFotmg0fSqlYWzL061L0JV8lUqxHItmK8C0225jfW8Ztk3JKuGmNoYlcWxT6zmEyOu2SLLMpadSJvKbWKaRNwrWb6SB38KSUO4rYAmFyhsMx0mitVlKZwAqIM4SpGPqGmEqI4pCiq5HFMaYGNjSJpUpnufRajV1wU7TwjG0xzJJY0xpYzkOUlpkGXhusZMsIaQgS6HlNzHkU2+kns3YlxC8/d/5+pl6vc7k1FRepNLHMB30pDs3yesMJxPTpFNpen4PMcPQsoIsU52JR4ebJArIMp1tleaiYP0CamFaul5YJ1WftrdKLpiIdYHF/BgUJElGHKekeR2yudYsOhvMcR2aTR/LsojiqJO6rvvICgxp09fbT29PH6NZRpJq70WS5BliWUIY+dSqVRqNBlEcEwQtWs0GTb9Jq9mg0WzQbObeqqYucxBHIX4U67C1SjsC6HaFbV0YEkxDdcI8mY4raXIkjZyktMlP+2LmQle9qs7wEyCliW0b+eQ/P+TX7tk2d31Ae/WyXIQ9d588Wu8GmovqbgNtja1SqtPLKVMg2tW5hSJTgtyNhVRG7vnKCMKIlj9HGvU9pHVDZk5G2s2cLdPMPY0mruPpwqYdr6PWKJVK5U5oeL4Xrp3U0Gq1OmUU4jjG9zVB1gVXo04x0XaItn3P6ENTxCpBZVHn/Akp8wiQo8mdrUXzWao1pk819ivSJISFMCT1WgvHdXnpOS/lU5/6v6xcvpK06UOUsHnzVsbGdjO6ZAnLVy5nYmKMn914Pffedy8pgu07d/Hw5q20gkC/GadZ521CS5xy92hbXDjPJTs3Ds3os0zx61/fxa233pZX0k3ZPTnJAw+tZ3JySrdkkQYIxYOPrCdKtWDWcSy+//3v8aMf/YhKqYRt6d5GqUKXOpASs2CTSYXtmLzpDRfz0heew+X/8FXuvO1ePMult79CRJNGOMumjZv5wue/SKsZYVoeSapY99B6tmzdCCJGSknBs1m2aBGHHrCGsZ27mZ2u4ZgC1wTXNgkSA9vymJqd5B3vejuNWR+Jw8jigzn2+PNYc/jzUM4AY1MxrToUyoqgkVHpK+HKBE8YSKWwEpM4UPiRyfI1hzG0ZAXnnP8qBvv78mafMZZr47daRC39FpEonSLbrNUIgoCVq1ZTKJeYna1x152/YcuWrezctY1Vq1YxODzEunXr+M2ddwJQKBSYnJzk2GOO5fknPp/evt5crwBJEFD0PJaOLuaIgw8mDupMT+7Br80yHjZpVvdQmx7n4douZmemEfEs1ZkZZmenKZWLSAOkKYmSCCEUnueRqQy/FRH6EZQNioUepFGg3vCxXYhVhlf0sCyTIAo1GbYt7ILu/5SmKVGSkEUpjmsjpcgLDKZ4rk7NLhYK9JUcMl9hpZKi6VCyPGKjhcygYhdpNn2yTFDqqVAwPaQ0iNOENMkwpYU0beI0BsC2XYQUNOtNPMcjbAZI4ZA1QhphA9vxSNOMHqdMnGnDmCYZQZLoSQMJ0galiEIt+iwUC4RRhFKCJIxptZoYpknBKf3e7MH+ike3qXiswpbzlzWbLWZnZmg0GpogJBlOy9fEpV1Xx7LyTKu8V1l7gpN6ojfNtoYlJy+5d3POm6K9D2qeF0JKA0NKhJK6Bt68lHtp6N8I5jxX7YSHLNO6lkQmKJV7uZJEE4k2QbEMXKeBaZiEsW69Y5pmnjlm43oOnuvklfVtSk6ho8XRDcJ15m/g+4RhqJ+tWFdBDwKfZrNJvdGgWp2lWp1hdrZOq1HXticMieOQONfuRNFcfSatPRLEmUJlcV6VW3tUDcPCyDPmpNEmjLmuqCNaz70m5B7kfP12mE97aOamlk6YtO2pklITw6zdtVQseItvh/x0WDYfSyfcp3dBJueyzDriqTx0aLTJdN7OpFOGICNJ5tez0oKrthZufqFV03QoFkpYlg71tVvGeJ7XKT3gOE6HbLV/lyRpp3WZbrsT0Wr5eUHWIM/ia2f+tdv95F5VRS5GTzpFPtvidaV0PaskThGkkIvTBf/LSVNv3wBBFFCrVrngda/hc3/7eUaGBsniUMejTVh98CpWHrwaqbRQcGTpKK98zWs4//yXoWxdtl86uq5MnKa6A7SpBYVSiDwOm3UyHeIkRmBgW04u0MtQZGTo9gbSNJitzfLwww8zvnsPK1asYGLPBB//6McY27WboYEBkLph7K6xXSxbvpQ1a9YwMznLyPAiXvqSsxkZHmRsbCeGaWA5Ll65RCuJEKbiBz/4Lpd/4R/5xF9/EpFIRGZRtEuEmY+f+aQq1m8dSAQWCBuEQSbAMB0s06avr8SRhx3EUKXCPXfdRckpsnTNKsrlPrAdSn09zDSqjIwuZsWBK6k3Z9mxfYzDDnkeLzzzVRR7lrN5V53xiQamVWGwf4BKpY+kVcchYffYdsyeIiMD/ZRdl6kphekUSLISv/r5z1m6ZBEDfT1s27qVwG+xauVKbrjhv7nv3nsgyxgf34NpmRRLBYYGhhBScO+997FzzzimlCxevJhzzn0pQ8PDhFFI4PuUy2XWPuc5rFq9mmq1qlOxG03SNGV4aIhSqaQ7klsWrVqN+++9G0ukZJFPlsa0DEW9ZmLbkvrsDGXPwRSKLIkoFz0G+nqZqc4SBiHS1A+e74dIaeB5JQwpiRNBPY1QQpKgKBcdRCxIRUYU+ggJvX19+gFPFSqvraJSRZxpLYYQujmwY1iIVNGq+ahQ4tdDslDpgqvCwET/85sBsZkQBDHVWkKUWIyOLsO0LIIwQZoWluWQxDrEoeIU07ZxCy6u7WIKg3owi907gHJSJBLXc2i1WkShLkGQpinNwCdRWZ6dY2p9i9IlP5Ik0W+nzRaVnpJuswIUix6z9W6dpt8Gj1Vy4NEZdWEY0mg08VstavUqabWGkGZHq9PWcbRbarTbaBi5h9IwTHp7e3EcV4f98to/7cnNNE0dunYsQHs7TcMCJFGkyNI5T1OWqTy9Xs/IaaqzlNNEl+NQ6D6chmGQxBYC7TkReWq60JXxco4hEEgsaRMnMX7Dz0OO8/RZub7HMA1MK/dq5MTJzGu52baNECaebVEu9HQIoxS6wniSJoSB1kDFsS7i2WrUaNRrzM7OMj09w/TMDNXqLPV6Qwufw1z8HCekmd5GFKWdTL6OEDv3EhnSwDDzayG1vkuRESUhcRCBSvOaT0YnU7qtIUJp+5CmGalAE0vR7tGm3Uid+0IpVKpD5Fl+DtticsvQnkPyUK12AMwXtutQ2rw7Dsuy8sreqhOqbGuY5u7FOc2VUoI4CmlkIITf0XK1r9VciNLsEOn2cs/zdKZcfs1c16W316Wvr79DvHVIWgvJwzCk2WzSaDQ6lcTb2X1RFOYhQF1CIU0TokhXQpfSoLe3l0LhqX+J269IU39vL7XWHgBc29UGJs1oTkyTRSFuTxFpSZIghlR/Nz09zuz0OKNLRtm6bTtDo6MsXroCt1AAKRCG1WHUkAIOQhi0anUaM7O6VIFl41YMvGJ+AUwDskT7Qw3BooEhBo/vJUkSLMMijmKOvvo7pFGKSlNUrvvYvGUTo6OLaDQbbN+2i7AVEvot1t13L+O7dzEzPU0YJfQNDbJ5+2bCOGDzlo3UJxogJJ4tsLAJVUhsRuimBA4Lg4laU9Lm1wqYmZnill/eip3f3aUsYydbEKZNLAyUaeBnMQEx0rEhVZz9kpdz0Jpj6e9ZRLF3EOkOUKmFKKNAsdyH32gRFmNKns0JhzyXYrFArd5ABsuo9Qh2T8/QXzG4+tprefD+eygVXFrNBsPDQ7zuwgvp7+/j4IMOZvmypaxcuYIkSTEtk5HhETzPZXpmhmK5jBCCnr5e4jiiWp3FMCTPfc5aBDA9NUWz2WTFsmU0my2q1SpxnFfJjUKkITGFwLEtiq6LbSmaswEFy6Gnt0IQNDEMsItlfN+nOl0ljmIcVxtg03KwHBdhGDq0hZ7QoiRBxVpTJU2HLEuxXItm0CROQgqeh1ew85o6Nd32IdO3l22YGLaNzH34hkS3P1EZYdPnkBUH0PIjttlbUZ6FKS1UHuYwbO1JkIZEmQYybiIcg2K/bkwcNrWI1s90o127VIQsJg4ikjBianKKvr5erFxf5ZYKpKJJM2wSxAFZCI5n06jX8cpaIO4HPvWgjlACQ5gITAquR7Hs0YqaxEQIaRMlAXE9pFzpe5qtwLMP8wXfj0WadEaTvq/TJEahhbuZysgy/bI332PR9hBIqYX8UhoEgY/jOLSr5GsNkZ2LZ61O5pv2dph5MUMby3KRwuzodgxDYgpjgedh/r80TQkXhFb0uqY082rQuhZaOzTWJmKGNJCOi25Wm3UK1UJKLBI9EVsGlpXloUgD08iI4wwpg87YdIq+Di11CJZp4ti2PhcFxcjQIiSKKAzwfT/3drRotZr6b79Fq9Wi2WzSarZ02xnfp1arEwYBcf5Sonv86YzpLE2JI0VMDCRIqTP2UhIylSCkwkCHOjHoeG/aep4sr+Y+V2ogzxLM9U9Szrfzmry0a0u1mUuWppop5V4nPbXN9bebP1e0PUppXu8v3+qC30khUPN/l9emU0pXM2+HDdt6pDbm67vmkyZ9z7l56QRdNbxNsOaKYrqdptVKqZzUW3MC+rwmVpLk9adC3aev7fBo1+crlyvziOdTh/2KNMVxgt9s6hslj49neR8wKW3CZo0wjTGkiWW4WNKkt1zBdUyKvX2MximFUpnZqWlSoNRbwbAspiYnkYbQrQOSmJHhRUzuHqdeazC8eJSeSgWpFHGtShhFCFMLBKMkwrQtipWSTvE0LO79zd0MDQ2xZHQJmA76lUB33F6+ajnSMoj8gIPWHEwcpURxTBT6xIGvq2YLA1NAM2ppN2UaEwcBKhaYdi9C2aRxRmJlJNJHGHkn7EwhpYWQJlknBRWESlFxSFirEdYapI0WD92/jj3jM/z3z37BWGuWNFRzj5Kv3eSeU+CWX97Kxo3XoKQLZoEoATCxHQ+VxBgGHHHYYfziBzvIMoFX7GHPxBiVIgwPvpElIxaf//z/n4ldOxEqQxoGjusyNDCA32wyOTFOb28Po6OjtJracGVZiu8Huj2D67FixQq27thOrTpDHIVI2yaJQrJUYRkmnuOSxqnOFMur0BqG1EJwKXFsmyDwURUXA0lPuYgpMuqzk4ztHmP16lUYpkUYztI/MEChWKDZamhipARRFJOqME8dNkFIDFNPMpZp0S79Io2Ull9HCIhTA5nqiaxWq1GpVHAMGyUUQqW55gBQma4VlSlElhK2fKrTMyAt4hharYjAD4lSKID2LEgz19xl2LZBlARs2LQR1ytqV7hlEscxAoVhmtjSRYURcRjhWEUKXg+xEzFdm9XhBtNCug6WIfHjkJHhxfhZRCozwihBWCYHrT6cgw88CMsw2bp5Gzu2bWemPoPtWQgpaTZ1Jl2h7GG5T707/H8DHt1KJcuyDiFpL9NNdXW7ENd1MPNCk+3sqfaE2/YQ6CyldjhOJw1o3cxc/Z+5/xod7dOcFkriuSX6+4ewbQ+BkWddWdi2/m+7SrQWTevftEW/WVs3KubrhfSUkyQJKtHeKZU3m7by4odpphbUDHp0HzNdPNHsVElvNf18AledEgoqDyG2vVvtprpSCBzbZGR4kJ6eHopFl56egY4Aur3NJEl09l6zSb1ex/d9rSubnKTVauH7AbValVqtTq1WpdFo0mw28PNQoW5No2v62dJAGDYYqpPhOJ/E5HIpsnQuKy6TOrMuT1qcFyJbeN9oJ1CbJOnmvpnSWYOCvDAzbUmJnFeeAbRCUrdVaRNckSdH6V5/ufBaiLn5YY6fdRboawOwkKDIPDlhTtOUEYY+09NTC+759gtC2+NVKunmxrbdbtlSzMtU2LrHYTvcJ73OudOlKLI8A0+ftyjQob+nGvsXaYoiRBQjAFslCBR+q4YfNPKGkhIRx6RJRq1eJ00SKj099C9aQhgEDI2MImwLP4pIsizPaBK4ntNJUUyjCMuxMSxTN1gslbArFeJmQ2c+uTZKgWnpGkdSStIoIUoinEJZu4s9j1ajjiGbWI6NtEyyNKE2O4Mgo1jpoTI0qFm7QvtRZS7SUwKM+U+GLqqpfqdLpYAYohjVCti1YxeN2TpnXfBqpmp1QhIiMlKh3b6hH3H0UWuZna6yc8cE45M1mq0I2/ZQSHZs38n6B9axe2wHI70HUTQL7Ny5iyMPWY46uMRXv/Yx/vxPJSefeCyL+npZPrqYnt4eCqUiXrGEKSWDfWUOPGAVCN2407ZsBgcHyLKUIND6BCkE1ZlpbClxDBPTUpgGhM06vh/heAWGh0cYHFlMtVZj19guqtUqtoQwSfCjgEKliClivIIDaUir1SQNtWs+CpukaUQKebZRhh/6VOt1EjFX9Vcb4gwpwXEs0kzpFgZZihQGcRzgeBLTkBQLRYQhtLg/S3BsC8+xMYWh4/NJpt3zRl5LBKUrsicRQRYzWZulVKhoIqwyEgSxglasCMMUr2CSpGnHWNiGyeKhYXr6epmamcKPQqSp6+GkQiAtG8f2qCxeysjQSjKRMVPfzc5dm8nI6O8fZGTRUly3zLbtWymVPTZt20Zv7wAj/cMksWJmKuCmHfcwM9Uii2JWrCwyMDDAbK2KUgaeo0Oh/UMFdnUrgv/WmC/+hjmCNP9zm5AopbRkIJJ0Wkqgp992BtVc+4jcU4GR/y4ijtpERnZ0m21Rt578RF70UqeAm0adeq2ee5rIiYgOl2nPQFvH4uF6ukWGzr5ysS1H9wjNg3FGHjrKdBEitD80F1hnGVEQEAVBTnrIs8baqh6JFAapTIjnecSATjVrQ855NbQYPfeKoCCFRGUIBGkcsmPHGLvH9uiSJFK3b2lnYc1v66FF6y6O5dHXM8DKFWsQQocpkzghyov1BkHQ0VA1Gg2aTd0MudlqMFOdpFqb0WG6OCYMwrxvXEAcJ3mWHx2PkJmXZ1CZ6rRiykTW0S8tLOYpOsQIYaBETn5oX+M5IiPy+0JzFDG3VBqY5tw9OJdE0A4X6yiMtjlzhMu02qSs7RFrl05oR27nsguTJF5Ajhbq99rlGhS+rxO5pJzufK8THbQ43ZDmvNBzOwyYi9QtA8tyOl7I3p5eevt7fruH8UlgvyJN5UqF1nTKGcc9h5efcxaloiSNA6QjUKaBgYMgIUt9LbYDMimJlaBabTC8ZDEYEjdPQ8xydlrprWBYNuTF+ZPQxysV8QpFLDKmd+8iCkOKhYLuo5SmREmsFfuOQ6oUaeQjPMXSpUup9PQQRRFZlhBGAVmg3zCSWBf5isOQNMnyOiGCONINLguFQscIGoZBGseddOAoatBotrCkgVcuIYUkikMUiigICVthnkavw0KpUlh5o0kpBVmqPThplrLkiCMBOPjkE/d5nuM46fQYAvCbCb6fYNsmSim2b9/Dww8/yNZNmznm6KMpl4pMTI2zavVKTEMQNh+h2Zhl5yPruX3LZpYuWYybV/q2XV2YTAlJoVgkTRWzs7OdFgHlcpmeSg9SCPr7+tm+fTtHHX4Etm0TRy0c06V/YAClJEoYjE9Oc8uv72TFqlWsWLGCUm8PtXoVw7HpGRigb3iAPbu24JVdolZMqVJkaGAJnm3SqNXo6+vL+2DptOhWFFEWErfg5RknWrweBgFBGBInuhlqFCZkaYbnFUizPHwmde2vgutBpoiSkJ6eihZ6RzFpkmLZWvwdphLbsTEMk0yA29PDyLJlxHFGT7GXZc2AUn8fwyMj+IGPaRjMzE5R6fHwWy2SKEJSp1wqs3r5KoYWDfHAg+tIZmfAkCRBkzhLsJwCruVSbwYEwTQ7d03g9TQREvoHB6jV4cGfPUJPTy/jE5OcePLBWE6FUnmEwcGVbNmyhxtuvJt779mEyDKOOmSAoaFD6PUqgEkYZAjLpqdSwrUNNm3e+nSbgWcVHkvT1GkBlK/T1isBnawiTRryyUdIJG3CA/MTV9qfdb0l5nmwRO6dAkh1KENJRB72SVOIlCYEqsNk5ode5kIqrufh2LZuteI4FEtlCp6Xewf0C0I7w2pBqxjbxpCSNJ8025N8W8DcJoaCufpAadqe4OeOUwiRS3sEMvd4tD1Onb9ThWFoElCdrRNHAXOekjZ5MDttSnT6up2n2evq/aVSCSlNpGni2B5lOef90Wn3cd6qxCeKIvywRbU2Q71RJYp1K5PADzqhwGZL69SCXD/VLuIZBgF+4Ofep7wMTpK3N+mIusHoiNHzOkuaPs/dF53wnj5X82+3NnkSwuhk7rUJ9fz7Rq+sX/Dne0Tne0IffbeRi+Dn15/SwnU6QnSRi9lRBu2GxFma6fZNeQ0qrY/L61CJOQLVDl+2SZNl5qUHbLtDsNo9/J5q7FekKc0UjaTJK1/9al7w0lfh2BHK9ChW+knijDTMyASYtsvAsIdp6VoZUZIibBtsm7DVRKEFrvotP8NQgiQNMWxLZ1P4gS6W5RZIoojW1CQqU4hSiSTVlW2jIECZFkrolgW6SqzCcV1UlmkjZ0pdsC1o4ToOA4tGkNIg8n3CIECYGUKaBHGEAArSQKI9aqZpkcYJCBtpGiRxTKPZoFgs4JBhmJLI101fWy2fRrWhXeXSIGi1SNMEUSxocXucEYURfqtJKwxJgWq9xuBAP65XRJgShW4bI9A99aRyQUjSJMU0BH29rq4BksJhBy3lsCOWzl2YhLk7SWVc8a1vEMyOs3vXLn71q5tJ87GPT0ziBwFTs7Ps3r2HWq2m3bYKDGkQhBHlcpnhgUGUUixbupR77r6HerVOf18vjfoMhgH9/f0MDA1zxFHHsnX7Fj7xfz/JqpWruPTtb+fAgw/EMLUHJwPG9kywbt1GZCZQccBgb4XevqUUrSJKSiarWnBpewXSxAezSLm3QqFUwrJNkiSk5Tep9OkWAGmmcG0Xx3bzDBiFkClTU7to+VWqtRqV3j5M06Ve9zHtAo6rK9bqzuqCRlPXn4kDsB0Dy3Ip9w2xdPkK6o0G5WKF/mZIKxsHsxev1IeUkmQ2ptqQDPT1kkRNwlAQhHDPPZsobZmlWmvS2zeAV3aRYgbblli2w7btk9x+20Ps3FZjZzXmBc8f4tDDl3LYYUt5ZMMm/vFbN2BkChc47bQTOWDNMYDF5i1T3HzzA9y7fjfC9jjh6KUceXA/s9Mz7Bnfw7JlKxBGQuCHIPqozqY8cP/s788gPEsw36v06JBFG53aPxIW1uiZH05bmH3XTvFve670JGMt0Iu0w18wJ+QVQmiNkaFFw2lee2vue5ln3GVEsU8U+1Rr051srfaE7dg2nlfotGbRXqkCjuN2MqxKpYqe2ITCcVwKXqmTjt7OlGprnuY3a+3IUGGezkZnfyUknRBlOxNa62lsbMvUjn2010RnwbXF20YuQlYEga411B7L/HM0X6RuWia2ZefZY3NZY729feiebAZewcV2Lf0SluuWoijG931qtRozM9NMTU8xNTXF7OwMzWaDmZkZZmamCUMtcg7CgMAPOqE/XSyzXSBSYhg6nGubWouGyDr3R/tEZdlcaK19HvUxzXnt5vdXbUMf71z4dqFHNOsQnLn7jQXbFVJgmfO3n3QIWCfkKOZqjOn+dm1BuU5M0IkIc16uDhFDEcchUaQeNV7Brl07OyUynkrsV6Rp48YNQEKq9EnLMkiVQiSJbu5oSBzLQUhJkuhCWyJOkJZFpVxGRRH1Wl2nhKMzSjIUU5MTjO8Zp39ggM2bN9Pf30dvXz9GrY4hJb39/Z2u05j6yhYrPagkoVWv02zlPb+kxHEcAOLAR5oGxVKJgufRptkqy5uberpOTuAHuvuz52kilWZYnkeaJtiu1izEQYDtFVje16fv/DyUV+ntQylFuacXlkiyJCWJQ9xyUVdHlxIlJfVGA9KMwZERwjjm3rt/wwPr1nHueefqrECVFytDaPGkZ+ZaHUGU6HAnrk4f1b3+MtIw1YdkCsJWC5FlWLaJYQiEbeD19bGqp8KKQ48EFYIwQOWNN3OZ+qPmhQ4EulzA7MwM/3nNtYwuXsz4+G7GJibZvHkjmzdvZmhwiA9d9mFGRkZ4ztFH86tf384XL/8CzzvheSxdOkqlUmL9+nXcfc/9eHaJhx6ZxrMMtm/fw62/2kIShpimIEjBkjG7mzFrFpUplx3cQoFVq1aycuUKEA6tMGFk0SJWrVyFZVpkmaJUKGKYFhMTkwgRY3suvl/DnPj/2PvTZ02y+84P+5yTe+az3aWWW0tXVe/oRjcAEgsxnOFwqJFnhpalkDQeWeGQ5Qi9cMiO8L9g/QseO/xCctiWxxFWyJY0MofD4XDTEARIAAQHQDcINHqp7qrq2u72rLmfc/zid/K5t2qaJMABQDXBDFxU3a57nyXzyczv+f6+y2PibOq1GREqDLn17AvcevY5lgvRRHRtx2ZTcnK68Ks6zft3Trn34R9LoGqS8O577/H9792mKu026ffw8ZJjC3//33idSe44njdsVg2/95Uvc7jp2Z+l/P1/+29z+eJVkmDMeJrxwZ07/MbvvMH33l6Sobi8u8PNGy8zLlKuXH6el58zXA2hbCFNYm7deJVVOeedd9/ny1/5Dr/7je+hnOFvfuFV/ie/9LOMk46v/8GXOX50xM3rz7Izy+k7S55NuXPvmPv3fqyXgL902/mVO5wBqEFbM3x//u+DQ07hcFbOwyAUN5mzMr4TdxYyyhXkIHUqztK1zVmwpQ63Fu6u6+Txh/EMUjlltzdg59lmaJqaJE4IAi3VRa1UsKRpRhRFwr72FipDXWs/glGiN92O8BIfSijJ02mWk6VSPhtHEXGS+iBYGUWGoUgm5DliSRnaisYHNCDXlW2dBxKRYPoe03c09QalRLM0BPJZ5/vl1JkjbhgNnjnWkH0YCeJyg+6qd5jGiH4mLOX9hD5hWyt0qAnjgCQR4Bgnks4dBAFZWhCFCePRhEuXLtO1nV+US5ZaWUp2UV1VbMqNH/ut2Wy8QL3cUFalr6JpfCGvaKoGPdgQ1iwgRUarZwnpQxWNH2Uq/CjuDJxseSQl0RJnoMkhximQ9gOJr3DD86jzn21hCeW4eM0VA9j3z7N9TCU1UVsNm6Lv7TmQeMaSDg/vo7G295Mtu2Utzv2UJ4KTjrlcjNi9eIDWjmqzoTU9m+qELB/TN4bAgnKOo8ePOD09IYoTprMZWVZw6eAie3u7qEBTNzVdueH4+JAv/e6X+Oa//Je8/ulP8+Uvf9lnjYBG7O7PP/88L7/8MleuHBDFIdmoIAgTtFPkoxH5dIrDsZ6fkhWF6J205KegNcqftH3To53XZoUBXW8pNxvMaiVM0mrFbDIlzTJR/49GrFYr5osFxXjE3t7emY7KGJQ2bMqSrusIwghnxO2QJjHGWuqmIgpCcJa2a1kvl2R5zs9+4ef4zOc+T4BDRRGb9RpQ5MUIOgM6wA3zaeVQ+oxmHWhXERC2rE/XvPv9t7h86SKXDy5SLkpGo4I0E2ZKhz2L4xPCOCErRgT+99HBln7VQQDWOz+UnMhxHHPx4ID/5H/3v0Ubg9XgmgYbaKrNhvfffou8yInShH/4f/mHGNPz3u13+OrXvsqb3/kWFy9e5IUXnuf/8J/9Z7z4/HN85vVPE49m2GbNe998gze/8x2SLMH0hjt37vHf/ff/mA/vPeTRSQ0s+ePvP8LxVQCUlouJseCsBzEKHIrWOYpQ80t/+3X+3t/726xLuP3BA569dYtnX3idr/7BH/KP/l//OUeLjs449meaf+ff+iU+97nPMd45wBj45pvv8p//l/+cRIHTCkKNDoSO71tNjO+VigJGVU/ddFjXMZrssruX8+Dwe6jVks2i4oMPHnFw5QpBBN9/6wO+/eZ3ODxckkaaa1cu8Z/+b/4T4mTF9773DT744A5Hxwt6qVEkDgL2nnmOe3/4VX7t13+Lr/7hd1kuaw4O9rhwcAEVWpbrE27dvIZG0VUVTbniwsXLFMWM77/9BvMfve7yL912nkV62jF3PgDy/M8NbIfy4Gc7gnIi9tZK2JPeGs9UywjD9K3/HVmoDMyS1kpG3M5hjcH6G6xSChXK2GcY1Yh13XoWwMcOmIYgiInjAJBydKnk0CTJoLtzOHpMLyClazufuaO2eqG+l+8HoW8QSExClmZMxjOyTEAYQBwnFMWI8XhCngu4stb5G/758ZHcgIfxFdZhnFSQyDjSEqgBGCJAyzgMQ3AnnB/74Rl4rSVxf8i6GkIvre9aa2kYbuYKGRdaNQizNWkizFte5KR+ETowiFlSMJvE3oofCVvupHS5bRoqb62vPUjalCXr9ZrVasVqtZIAz3LN6fxEUuO7bvt7TdNsnWVPxBZsP5CST3i2nQF5NwAlJwdf+/iD7S+eYz3B7+/tmPDs8y7PfVY1s93xnH8s5/f9oIuyvmRYgkzlsGo0Z2zf8FjbobQ6B+yiAGt/yt1zcRDy+PQd7j94gLWafDQmNC2j6Q5BEMs83jroeiajnLa9igPCJEGHIVGR09U1kQpI8xG62hAozeuvvcYv/0//LZI44pf/7t9FhSEqjLC9odxsOHz8mHffe49f/dV/wjvvvo2xPVcvX+Pzn/scX/z5n+fyzVvQteAc6/WaNE19BYis/KyRE7vve5I0I0tiQJEoTTEeI6yZ9RUEEbbvCKIYrTU7Wcr0wj7WrwTzJAag7XsePHjA48ePGU8m3Lhxi2JSyKqx6yDQ5OlIUhGShNOjI07mpzwzm2HaltOTE4o8IyvG5KOxRCO0LUpp6rLm6OSYIssZTackceLHgxsCrcnyHJ2GRIkmK1Kmo4wwCgnjkCzPPL3qVw6dZTSe4YC+7bFAFGmsleBEHWjpFFQBKgyksqYt0TogizIiL3wNVYgqCrkQas0zt54ly1OiNCaMBRiOxq/y4svP0Xe978IK6TtDEsZkUSrxEkHCJ37mizz/+udEx2R7eqf4j/7T/z193WDskKWiZXVkew/+Ik5O5/zxG2/wR3/4DX7zt/8F3/zjt3DAAvjVf/oNfv233qQzPc5asvCrhKGmrDvKppOaNuDhI/j//uN/we9+5Y+4efM6n/70ZyQJ3BpcmvDFv/Yq/+F//B/w3vvv8/j+Ia+89EnqdcXudEoaab72td/Hupq63jCZzAiCiJ07d2ldyXO3nuXf+g/+HV555SX+xW/9Jr/zP3yTt96+R2dgOs757Oc/hXE133v3O5TdChVDnAcUU3h0Cq7vufvuu3z1D77E2++8j3ENP/dzz/PX/8bnuXRhxKY6Juo3GFNx+fIezlq5dSvDh/c/5P6Hj8kUzP8EBvGvtj99e1oA/qf9XG974iDeMkbOOdHc+VZ6Y8VD27at7+QKhXlXijSVcl1jDHUtgYKDEWYQ6eLPMwFXw6iso6okC01rhQ4cWluSNCCMUtFcKUNZrbdZQMEg4o0C4ihF1EnnRMIIWGvaDU1b+QoNAWB5mpNmGWEYeXCovGYqJUlir5OKmIynYsDwz7V1+EWJ1IOEoWd3UrnOAL0P28Q9OfY7X2WyBRceHOlee0OIPgOUnlYZAkT19qatvRpbfF0gP9+2HcasWa82W9ZE+xFhFAUEYUQYCmXTGyNgLQgIA3Eu5vmI2WyPOI6f0J06ZL+t1wKghuTt1WrFYrFksViwXC5YLERCsFouKb1MpK5r2q7dOg8FRYr0YGDYBHhbP06zoAyoXv5EDEwSgeEBvV90i1TtDBwNrr6n9U4OK+ND/5kYRq7gwWp07r06zsaCw9h2q03T55ygevu8P8rtYwWaknxMuzEYJydeEMXEYXBWAAiAQkUhcZgTxJEEXhnD4nRFWldMxmOqzYbjoyOqzZp8VPDSq6+SxAkqCCimk+3qwjkwu7vsX7zIMzee4VOf/iSbciM0rRH67/j4iNV6RZqkXL3+DMYa380DKC+8U4rOGE7nC6ZjR5okcnBDOVlQklvadT3lZu0F5p5+tE7KF+NIftavjpIkYX9/n6IoPO0bnZ2kTuy2xvSUm5JISw9PFsesFwuK8diL2hOaquLw7l0ODw/J0pTnXnyBOEm4cPECQRjKisn1hLEmiDKfltYLVRsE2K4ljkVwjr94r5dLNlXF7t6udEplmVDVviC5N3Zb+tl2HV3XkSTivNGhJkTm2A5w1lCWmy3t75zcDEbjMX3fimarqdFhKEWdOFnJxamsbhKNViGrxZokUsRh4lekAVmaygWtbaFq2b9yEaXA9B1xIu5IFBgfQnnp2g1u3HqWn/v5v8E/+F/+RyyWpVSW6ADVdazXa5LxiFCBsj1hFNJ0Hffv3mUyyfmNf/7P+fVf/w3uH55wOC957/0TvvwHb9Eay9pIxMCXfv+PeeuDf8imKjFtx+7oD7DGMokj0BFNdYS1HZevXuTv/L1PMJrscfKbf8TjdU/5/Q84PlnR9z3vvX+Hw+NjXnzpKn/rb/11XvnE61RVy6hIKPtLnC4M6XhMmK0IMs3FsODgyjXeeecd3n7vXR6fbphMpnzxiz/DL/z1z/HB7be4c+cBF6cx63IJvQT4jadjHPDO7Xt883v3qP4KMP3Q29P6pT9r00pJa7wHNH0nye9t12KNsBqRZ3IA/5nusV4fJOM5iQoIg4AolPG8lLb6G1Iv+lCtFUGocM54pkaeK4oijGlpXIdS1vd9BXRdh7PGM1Vy49TaEeozRs0IZevt/8IGWOcQYj6iLDtM31BWlt60RFF0Tld0Nj50ToTBeZ6TJBlaSxJ6FEckcUIcp1JQm5wJusPAO7C0xA94Tmigl7ZuQrkub/kmGOTXTlg5GXM53OAeUyDwaLhRO89QeTec1iLm7kTPxFM3+iAYOuGESRwSs7WPTJCRZkSXJqSZxRonhch+FBjogCiJSZOM3d29Lfgc3Hp1XZ9FKHiGqqokOX29WlOWG9q281U00hlXlVI70zSNNFb4ehM5ho62q3GulUWl0h60DsdH9Heyz85nj/lyXs4WCWek1MDSiaZp2O9nWjp/viAkhHUKq/Aoans2IU6/87qnH+32sQJNxih0eJHxbA+tNE3dEEaS0l21FVGQEOlQ4vWNYzV0kI3H5OMR1WZN7Vchs709OWi9JYlSlLXYvkeHAdZ3/uhQVkhKK5IkZvfCHs7IRUGh6dqe5WLBcr6gbRp0GIJhCygcnvZUijCKGY3GxEmKtZaq2qCjiDhJMEY6n+I4JkkTgsD3kFlLGMUEvmk+DAKaptmuUEajkU9yFX3WsEqVEDc8uEqxnThtiqkUzGItaZKikwRHw6goZH6cJChkhRnoWADQekVdV6RZQj4aEcQhtutoyoowlkoBsCwXa3QYUYxHRHFMoUXDpINQBJcogjAi0IGAOt+vBLKS6LoW1GAb9qWgfYcONXEs+UMq1GAcCk0QpugWbNehAsBZYct6g9MGFwkw6nsD2pFkBZoAdIAKQpRz4LvgQixojRuyUAIRkZuuwzgpdlbgU741l69e5cr1G37F46lg5zC9EaDVNqzmp0gOkmZvtsPNZ29y69bz/K1/8++w2GwwvRT/Wgc6FMDurMHZFoNBadisljy++yHL4zm2Lvna1/+Iz3zmU3zv+9/nu9+7z+PFrxHGKe9/+JhNCcuq4b/4R/8N//if/DrvvnebR4clZR/x6OEpLzzX8C//8FtMdgqSoieIC/YPnuHduzUfHDvyoME8PuJXf+M3efvdu+B65qcrTo9bZrPL2Ks1y6MPKQrF/PiIvq44ONjj9U99mgeHp3z3rd8nCBL+wb/9d/i//jf/3U/2wvAx257OZPpBt4EHCAIR/CqlaNp2O9IzxpDECdPZlNF4LBEqXU/n2eemaVit1ywXC/qgIy8KxqN8C3YkuV40UMNNFxxRHCDONUugz4p+27aRQN8okmuUzwsKAuUZKykbFi3MkB0lN3Ln8NlIiZhj2kZ0TWkmsSI+ODEMIQjOXG+ymDVY2/vaD8Nms/T7NdgyzBIbIAGKUSw6ozCK0EpcVtPxjCLLt/qewYk1ZEkNY8IBzGjtYxu6VpxcyBfOCmPiBiCgfEq3737TmsgzUc66M+DBsK+E+RJDkT2LRvDsjgW6Vvr7mrphs95so1C2onTk6QMdkBVSY3IWXpow2hsTRbGANyWgdXu8+86HeIrVvyxLlktJST89PWU+n/ui5LUHWQ3GdrRtxWIJfa+lcPycxkgqTc70TYM5QP7xzMF3pkECpQVsncUWnIEmkXF4EKq90kr5oFAvnZDHOtO1yR9nTsMf5faxAk29g9c//TluPfs8dVWzWjzm4sV9gjih6hpcD4RyYvdd550IG57ZmVGMCt+1pAnjmCgbcfToEUePHnPt2efBBwyiIlQQ0rWNWP7DoZ7AtzErJ8+jFC7UjMYjkiSha1qWp6cU0wl9bwiUEq0OA7rWpKl0hPW90I/a6xOsOnOy5EUBQFdKsnVgDLEVANN0HdYY8jzfriK1t4pKIaaRpuow3CrjkiTF6M6PkM8+rW3X0bctWZYx2dtjpBTWdKyWSwoLVVmK8N0N4k/DanGCQwSZtu+Zz48xvWFvf58wCnxqcSv7N4/BGNqqIwoANIv5krppmEyniJ1f0l5FjGq2NK4w3g7TWzbLJVmS4bAS0+AzSUzfEwYCnAOEJo91SJiI9mgxn9N1hiwrCNJILkgomqZGKU1SZKClECdMQpIslhPQGnrbYXvlT0gLKqA33RY0ORtKT6BzxGHsX29AtdlQhCMIQ+Ispa4bNuuKvf0LFKMZL73yGrdefhWDX2V5Onw4NioItloJlMN0S07vP2S1WNGuNnz+c9/hhU9+gvffv83tu/dYNx1Wa37h39AYAzglrKQzfPKTn+Hk6JhH9+5ycjjnzW99l1/5td8hz0LGexlBonh4At9/5yFla6lpWbTHHH3lX9I0NfSW3sKXv/JNJsWIg0sZ9abncVkThiOyiWI8nlHWhu9+/wHff/eYMIxp2r8SNf24tkEnYrYjLPVEAGQYhly9doWf/dnPcuvWzW01RZKIoeSdd97h9u3bPHz4gNP5nKqspDJns9mWqIY+oqTvZXTXda0kWQeKOE7Y3d0lSRO6tts6uYqiIAxDuk7CHOVa1m/HXUEQorSUR4tGctAIKeIkJNABbVdLPIvr6U2Hw/hxj/NaK29T1z5s0TP4cqceQhv94sX2mLan62rq+py7MAhkwaU0h8FDYZsCLdfMKPSi9NinVSc+ekAAmPY6y8lkSpJE2wXqoDHjHBCWG71CewLEDgXDfn/Yc9e5rQh9OyYUkDCAHuUkZ8p6a7MxiO6U80+pvN5MUVaVzy+S6IQhMytJYqLQV7s4iUfoexGdJ0nChQueZTdnDFXbtluGqiyFlaqqmt601PWG45OHbDYL6qZks6l8L+LKC9hb3xvny58dXgultiwbuO09WQc+LkMPuVAS9umckfcWCO0nLrozbdZ2hLj9OAhbN+iqBJz+aLePFWgCzSdefY0rV66LU6yVNFmZ98Yop7xIUcY0cZzgELeaMYZ8NPLjFIUKlAiZN6XURZteqMjNhnyUY52EpgVIiaS4DKS12iE9QdVmjVIBk50d2rLi+PiU8e6uuDmGOa4/0PiE3+HPNM8IIp9sGgbCUnkdwiCO1kGwLS8EmcOHUYQxhjCSeASlpeARawX4+RRX5wPjcI4winHWYrpWxlFa01sjq0qdEyYxAY6m7jw7ZumtQVtFmuc+iqGmXC+8XkgTRgHSq9TSm54sy1G6p2tbiGLiaBBnui3Aa5uGqiwp8pzOyGvP0gRjDabvRJPmxOkjrhlD37UEoxHrpQgeR+MJ4+kUNdigfUGo8t1ubmCxOkNXtxRZAU6SaAMVol0gPVRWM54UaCRsMstigkDTmtYzjYN7Q2h40xmiMCCKYw/aZF+ZXioqcIq2rgm1JooD0X1FAnYnkylBFELnCIysMK1zoP2KHEXTNJKImyagHU25RkUzdl68iKk7rHG88InPkF3Y4eVXHtFi6ZxEXgSBREbIwnfoBDMsjx/xzpvf5vjxQwId8+lPn8h4NwYVwXIdMN29zN/95V/eMuTOGJbzOZvFgjgKUMbyja9+i0+8dI1x7nhw/x7jPGZvOuK99x/wxvfu8vVv3ePR0pCFDb/9+1/7iV0Nflo3ayyd67YaJBl5C1N9/foz/OLf/EVeePEFsixjPp8DwkLs7e7x/HPPc3x8zHe/912++8ff5d69e5RlSRzF7O7tsr+/T5ZmWGto2pbNZs1ieSqGFSNWcdNLno4ABmGTurajaVthqb3zz/pUb4WUAA/6nSgKPAiwktsTKrR2vpLkXG6SCjxDJVUxCr3NrZPzEz8GCtBqeE6HsWfgxPT+pun8DVUpsNA1IowWluk8aDoPnKLtGFMrTZKlXNjf91U0PMFOBf56HZzT0wxWffyiKpQ361+O24KHwYZvjXT6OV8QbI050+cMovVh8wyTjPMkHNlaMRzZZhDbD68v9JUzgccszoMjYfxmsxmz2UwytbQmjhOyLD83rnXepdjT9wbnDG1XMV8cUZZLqVparVkuliyXc5bLlfQkVpI/1bW9d/H1Hmy324J701upqHKii7LG65u26eiyOAgR0kJckudcjeDHqH7UOojQnUToDCzsj3L7WIEmhWJ3d48sL8iznJAdQq3BOpIowlnoOxHiaqWYTiY4RoCjqjaMx1MWyyV5lpPogCRJmE4mWwamrEqWyxVKXyDLUiLf9aWCQS2Frywx2/GSRpitodTV+qRxiw9lU9rrAYxc4IyT80Z78IWMm4aiSxAXWeQj/a0VjYJSyveOBdRVRYLXLlnrdQmhn/ErAUh9j2nFORPGiRfmaaJUSovDJCY+V444uGNG0ymB1ozDAOczLpyV95hnOTZzhFEkoWpxQt9bAiXvOwhDnJELVt+0KAVJmhKEAcb0FKOcrMiJ4gjXGEkL9nqLpq78BdRfgFSKwpLEkeSaNDXr1YLedIDFtQ1aadLRiBAZeZi+o2k7ATs6JE81oZYwUY3DmoYwSqg2CzYnJUVxk7quiOLEj1wtWjmiQISKbVORpono5zQEOkJWR5Yg0linUNavnJwjSSPatkbphCRJSLUm3NnBti3Wj3UbfyFp+x6lFPmoIE0z+q6hrkrC3RmOjuPDR4RByt7sCqbqyUZTjGoo4pw6jMnzmHyyA0FM77TcxDxO1ziU62l3p1y9dNGPcDSf/hv/Jg4lFvJhleub7L2yA5zl8ME97r5/myJJ2CzX3L97l0sX98hSWCwqVstjIuVYLk9orcYSc+XKAVcOLvLWe3d/MheDn+JNVtiKKBTbfqOFAYjDmAu7F3jlE69w6eIlHj56yLf+6FvcuXuHQAd85mc+w2uffE1ACZoPbn9AXdVEQcQrn3iFz3/+87z88stMZ1OUAxUolssFb7/zNm9977u8/f3vc//BA9pWgnQno4kI0L3uxRrLeDIm9aXA0kIvadfOX18CLdqXHrOtWwkDTag1vWc6wmiIQ/A3+aHI1nYYI6Obs3JY77xFbrJbYfCwcPUMh3y+B15XAJtS1rM4lrbr6U1L02oBPoH2r/V8LpPi7p3bWz1mEsc+QkC0U2makibpVkMlWqSELCu8ED3cjtasrzzpei/c70Xnpfw1vO87qs3mLEPqXNK58mKrbXRAIF9Ka4JQ9FqDkF2Ogdm+/uH9nzGUsFot6ftOFl9uuG8Ni8azxO0tkIxCorggTkKcu4D2FSoD+GvbjrpuqarNtnJmCPycz+ccHZ2IC7CpWK9Kcf5VG9pa2K3edvJ5YQh4FZH4wM4NC9otSFJPnx/y58CY/ai3jxVoAgmPjGPfE2RzMAbb9bRdh7HOU5Carm4wXQta0ZUlvTUkTcM777zDszdvEXsB3XgyRvkckzhJmE6VZC35VVTgR13WGhEBWq9b8fZ+HQRYZ2mbmjRN6dqGyOeXoAaxooMgODsBkUAuY+UCAdJzZKxlOp1ijZEuOy+ABDlBAq3pjNn2N8kq023zRrAW28snpu972trrn4baDn8hqNuWOE3RYUjXttiul5myF3LL48lJ1TS1b0QPtsJ0cCLC15YkEReO5G4EBE5WPH0rWTFxKiWZOEuaJigd0HYNcSLuD6xBK0egFX3XYvpWLhB+DFg3NaY3xFHI/qVLrJdz3n3ruygHRV5w7ep1wjzHGIfWMYFWlKWscgZBaG8NURRieotzrYgXkcyWe3fvsb9/kbZtmU13BAAHAcb1VOsNgbMC7rybripLsjQlyTIwUrXiEJdk31th4qKY2Mnq3tDz+PFDLl266Kl+SAhJUqnwCQKFsy1BIMet6yqapqSsNiQBlMGacr4iCRNOj45JRhFJGNC3DZvVKTrKMcSs1w0YS+AckQZcQ10uiSPNhf0LoEIuXLpOEMdESUxVbqSbLk39R8iPBTQ09cvMjx8TBSHlquTR/QfkWU42Srnx/Mt8+N5bBIhTZ7Z/kaN5xfG65vN//Rf4tX/yq/yf/4v/+0/6wvBTtQ0jIYWUkp/PKIqiiFExIslSQq9z7NqO1rUsF0uaumF3d9fHERjyLOeZZ57hF3/xF/mFv/ELXL9+HWsty+XSszrwwvMv8uytZ5lNd/jy732Z+WLB9avXefbZZ8nyjLfffptHjx5Rlhv/fKJziqMYo30XXdf6UZ3Ur9jeUJel6AWdmGqiMCD0N38F4lxVQ2VKSGBFwzS4pgYxjziPYbC0DyJg78V5ou9N+XFe4J3NcKYrk98XION6touhYQQkLjphzobxWeR1qYnPnUpTAU1xFBOEIVEUk2UFWVYQRcmWHRyE7GZYsALaa7KiMEIrPGjw3YAebJ1phPwBHxgrP9mwzqKs3K+sMzg7aITUOeA0FOzKOV/XFX3fbffB8HkaBPHb8WU8vC9h4MNQHH06jkiTVP49DresnzECjNuuk9qprmW9WjOfL2hbAVKbTcmmlN6+zVpKkstqQ1U1tE1D17eSPdWW9L2A7L4TDbDxY08BxwKJh1GpUpLvZM1PuaZJo7YWWRyYrsfalljLzR8dEEQiTnZWhMG96Wj7Vhrbyw3GdCL6BYIwlDRab19M05TRZIIKQymcPKcBcsbQNBVd1zCejDEYyrIkjGKK8Zi+60SzFPhRWXCu88dnNTV1Q6DPLgx4OhXYfsAG94Y1ZmsLxqN+ccAY0QFxJihtq4a6rDw69+nBsZy0Q/FkkiYypux7mqYhyTLZj36KpgO5uK1XK8bFmPV6Ix0+kQTp6UCLcc43bXdty3KxIE0zUHIhlFGZQmMhCL2YVE4YrRRRHDG4ZoJAGK62E0t/mkRYz0gpv4qSsk05nl3bUkwmRBE05QalAvKswLQdve44OT0FHTCaTEjTXKjtSJg6rRTVZi35J1FIMUpJk0hCc/34sq0aurpBx4l0EhKIXsnKjUn5nBssOKswrax+kyQCHaJwREkGVjE/PqUb5exdvABtQxgEqCCgaRt0qCjSDJAohbaraVuDDkR7onRPmGiyIiFSAToy6ARUAsXOiKZZQehwtqfvLMo6ur4hVNKXqJ0lVAbXWSItrBhK+jNc14ANaDZrjh89Ioojpjs7hGFIVTUcHj1g/8IOxajg4MplnHNMpxOm0wlV3ZCNCq7euMHy8DOYrqbvGyY7u6zrlrJpef6lV9nfnf0VaPoxb1rLTb/vpN5HtClykxpSpseTCTs7O3zuc5/jlVdeoes6jo6OePDgAdZajo6OqOuaa9eu8YUvfIEv/vxf49Zzz1JtSt555x3eefcdrLVcu3qVT3360/zsz3yOJEqpq5YP79/nmevX+cW/+Ytce+Ya3/jG13nnnXf58MMPeeedt1mtVhRFwWQyIQi0vzmuqaoS5yDPCy9ClnGe82MUrbVnpQ1dJyJlsZtHJHGKUiFhKP8m40HRPtqBTfIgagjmlf8/62sDzpgan8T4hJbIsh2PCWN1lgbunEMrLSDB28Ks62lbS9e3VKW3uwdn4nE1sLg6JAwkKiHNMrI0I0piuS5oYbbiITE9zYjThDiKyfL0DCCfG885NxToOv/+/KRCi2zAmF4AEw5nRait1dnIcNBcyXs7K0bWgdfZenei1sq72KDrO7q+ZWMV1hmCQJPn2TYPTKpLQt8JF3gx/iCqD0jTmDwfszPb5/p1OD9i63tD27VU5VCQvGS5XG37+1brBcvFiQdTpddOSb9fVdVbZ9+QPzaQE1oFdPyUM01hFLJcLVlvSqzd2bIiOgxIEmFORCOCABgH7aqhNx3T8Yi6qrh+/Rrj8Yi6qVFaMZ5NccpRVxVxksjJhPN2VJnLDymxMjMXhIse2AdB8UJfykrC9nZ7Qlp3ljLbtCJaJhFHR9u0bNYrwigS15x6UtxtrdiJlfb0tQdSA6hqmgacw/rVF+eEhpF3onSdaA3quqZqGl+xYihGI3HsJbFny5xP8RUQF8QRSRKTpKks6JwRrVPX+ZA6w3Kx8oW1PX3SkeWebWk7oiihqipQlrZtt2WTYRgyGo9RKPqmozcdUSyrGB0pnBPn35D1NJnNMF6Qao30/V24eIE0yQh0RLMWAWtT1RjnSNOMvMgoRgUqDDCmxTnD8dEjxqOcOBGgrDKNdiEX9y4yGU9pNzXr5Rqb9eR5QRAGFFmB0mBNj3bil4mC2JeXamxv6ehRkWZxumC6MyPRAavTE+qypKkkA2VnZ4Z1hqPDx4zHBePpGNN3tE2JU44wHhriRfSY5zFtlxCokGKcEUSKPmgZ7RWcnoqA1SmIgoAg0vRtTZbF5Gkq7Yl9g20NRTbFaocJRcPXdy19D7Q9YSy1P31doqKQtlwzP3zAONO0mwWj2RQdinM0HsWsmw3r5YLi4iX2Lh3408TSO0O+JxfGejPn4Orln+xF4adw0559sNZu3bRDWv9iseC9997bCrQvXrzIaDQCDe+98x5lWQo4V5LXdP36dV577TVuPHODpm345je/yZe+9CXefPNNAF566WW0Dnn55Zf51Gc+w3pT8sa336QsN+ztX+D11z/NeDLmhRdf4v3b79P3hvv37/P8c8/z/IsvsLMzw1nL6ekJjx49YL2WKAtjeo6ODrddbcfHR7Rts2XWRb9ifCij8SxavP03AQSiXVXqLF5hcLGdOa8GsOSzmHxq0hC/qBQyuPOTK+1EfyWjJgkAHW7wWmnCOPLXQ7m2C8Cy9MMoUF4Ew2rU4utLUERRTBLL6D7wYEIpKT8Ow0iaIZKUOI6F2c3zrUZqq50KxR04gJQoEiF9EPmFOL78Fy+o1hrlk86Hwl7r7xPCckliuDb99nnCMCQgwKE5K+D1Ces+nVtYyJAwcE8wPQObF4ahT3VPfE5YvGUfozjajhq11qSSAsN4JOO0rm1p2pauben7jrotqas1TVvRNDWbTbkFTUNCelXJGLBtWw/IhQlfr9ccHj38kZ5/HyvQ1NueosiJEhlthGFA21uqspQLvPJuJwum7en9zT0YHCTGMBqNiMKQxXyB0prJbEpnDIvFgktXDkSD5DVLCkHVSil0GJLnBUnsK0CMZTIei122a2malmIceebIr2KU3a4CwzAUy+tACRuD8amtAMWo2Io6tSgcpd/OU9rWg7AojjGdiJCrshJdTF6Q5SMwIpwegJq1hiCKyIKAuq5Zr9ZUdUWe5b5/0Q1mBi861BRFIWDFjwml50guTkEQ0BtL3bSEUURejEmzAkeNQ9G24p7rux7rFE3Xys1VKdquoxnqG+JYRnWBJiDAWmjbjsCvhMq6YnU6lwT0OEY5x3i6gw5geXrEarXEWXDdhlBHBHHEhYPL29XGZr0mTBMSnUmVA4YojHDWsVkuidOUYrKD7WXcS29oqoq+l5/r/fs1vSFMIkBTbUoJm9MhcRoQhjG2s5SbiqxQrBYnbNanXD64zKWDixjT0dUbqs2SPE0wXUVVrtHKCoUdyCo0iCLiNMEYaXlfrE4pRhnGSQp63Veo0NFs5h5Iiu7M+BTiSOZ9lMsTgr4jjiLqzVqypvIEFwQ0bYuxLU3ZUFcNo3REUYwIkcb3znZkScyVS5fJ04zDo0OJ24gjrJK4jCRL6MuWcrEk1IowDqjahuV6QTrKmc52mM9PWC8XP/kLw8dsezpm4HyB7dOVKR/5e8rfnLQiCAPfqSYO0aZpWCzk2hYEAZPxmPF4TFKkzGYzDg4OiKKIi5cuceHCBa5fv87lgwPSJOWDDz7g61//Or/7u7/L97//fdI0ZT5f0LU9Ogz5m7/4N/jCF77AerXhV3/1V/ngg7v8zGd/llu3nuXK1WvcuiFVQW+//TbPPPMMn/3cZ3nhpRcpipyT4yPe8WO8vd09tC+gnp/OuX//Pn/8x9/h+PgQY+UaHfhr1mq9otyUtG0nrmQvMh+AjLUiyNZKFjXyPw90UH4Ed7aAHSpRBgixDbdWZyDKJ7aeAbDhODhfByLf4KU8WxPG9jhx9nsaESTLoTN0fU1vGj9GdNsRIUg9zsDypJ6VCraMzVn1zMBIZZl8pamMxoIo3DJJSp0BmLO+vACQOpteWfBifhFcG/+aRLN7PsBzCIkc9E4MdTNOgfOs2zlGDixda2mDnrDuCMNGwF0QctaPqJ44D7af5yAgCiPiKCIZ5YShIoqkTgVtGUqbxand0TT1NsRzuRzS0TdUVcXR0SHHx8e88+73foiz88/ePlagqV4tePbZG1y+tIcONZ0zLBZzbKCJs4x8MoFAetOGZOeiKDDOsJyfkCY5putpVSslgkEsEQBNQ902LFYrmrYHaySmLAjJkgitFF3X0nU1URQwHo+oNpvtSK5aLLh3/wF7l6/StjVRlMgJiqP2XUDT8YQkiZ44UfNRQZZnQq8GgV+ZOLSVlU/fddukXgl/87PiridKU2G/kEg1aw0K6S0DEcS3TYsKNGEYk2UjkqwAZHUUhuF2Vbd1R3TiRgxCybqqS9ETdb0I84oiF+Do+6HyfCQsWVbIBcU6YXH86HE0HcEgwfSrMulHaghDeV9dI0LS3vTkeY5SmocPHvHB++/zqddeI4xCsmJMGAK9IUtykos5Sgcsjk6Jk0CYxjDailLLVUXvw/gWywXj8YjpdJc4ClivV3RND50I7OMsAx/jkOQFcZyANXRN7ROURZi5KSu6vmf/4kXCJMEauUgb54izlCvXLvG1r32F0SxjMhrjGkuUpMTZHtVqSZamHFy5SlVWNI0kMDd1j6MnLMWN1JmOw0enGAPT2QTrP7dxEJElKVprZvE+2mnCJKVVHW3VEgSRdIFFMQrHelPinCUaT0jGBX21oSs3jCcTuvaEzWbNcj5nPBoReYfmaDZjtrNHGEfMZo4wjuhMT2c6ut7Qdx3O9tRNCUCsUg5Pjjg+OWT/wi5pFtM0Jfu7k5/0ZeGnYjsPoKwVVjgMQ+IkQTmJCAijgN29PV566SXQmg/ef59vf+tbnM7nRHHEJ15+mddff52bz95iZ1ccU3mRy2JNQWc66uZM1C3apgXf+c53+NRnPk0Qhlx/5jqXL1/i0aPHvP/++xw+OuTgygHFqCCJEv7W3/oldnf3+eD92xweHvHSKy+zs7sLzlEUD0jTJRcvXuLCpUtorVivV7x/+zbWWh492iFJI27cuMF0OqWuGz744APee+89Hj58uI1IkKb7gIFNGm6iZwWwA0j4U/bn1oYz7NThL+7s/93ZwlJkQU/GC5w9x3Zw5n9vYLW89AElukf/oOdLkofcIhmvdlgnq9iyLInW0fY5hvc8sDUSI5CQpAlxnAgoiWPyXCpaOA9A4mg7BRmkFpEKCEPRabJ1ng1vZRgFDizT8PczKb21lrKsCPSgO9PnxqDC41kLfWexRgCvVt35vbQVtJ//e+DF7ENIaRgq4jgkSUOp8xqOB9q/p4RRMWFv9wLGGNFPeXH9er1mfnLCr/yT//bPPL9+mO1jBZqCOGL/0pS8CGmqkraqSbOUZDImSlKsUvSmJwwiXKCZHy/ZrBbkoxxjDeWqZL5YcHDtOscnx4Q6ZHf/AnGSMNnZ4d133uXdd98hjmOhseOY69evkUQRp6cnLJan7OxM+fRnPk1dllKJUhTUdSMfkq4jimPqusY6S5pmpInMpUM/NrR9jzMySpNUcL86MIau7+mtIUtTHI44jrcrT6UUgY8ZAGGq+q6j3Gyoy5LRaESe5z5lu6XpGqQLKkKh/UUFwjiSvCe/auuHks4gIEiSLTgbnBhBEJJk4oYJgxClQwIlWiylxHW1vXBokGWc5GOoQNLadRD6DCKZXbet2KUb7/LL8ly0WL5j6sqVK+zs7jEpCqxzHD58zIWLF4kjTW8k9yVQAdM9qZdB+3oCn4gcpz79NwoZTcYUE3EEqq4jy72WrO1xGrq+JwliitmEKIhoa8nnstbQdS1hFBAqzWgyJowTwiBgs1qhdSRjPA8SozTh6tWrJEFAGEf0fUez2RDGkbgZjSEtBuAqYFcHLX3rg1SDEB3G3HzuBbIsJtQO07a0dIBGRyGb9YYkjimrasuctW1HoA1hGKOjmCjPmQaiDSNQlKuSqm7QQUie58TXUqr5gnK5Jhvn5HnO6viEh/fvc/XmswLGTlc+EDRH9Q1d35HlOWEKm/kS4xyT8ZhLaUKSptLiniReoN7+pC8LH7vtrA7C/Yn/7YluLc7cTHAWCRIqCZXsW7lJROmY8XTMpcuX0EHI8ckRp6enfPet79J2PaEHPTe4SVEUJGnCYrXk6OSI6+11ptMpz9y8wfUbz7BYLsBJ/tvRyYlEFxhQoWI0HlNWJY8fP+bBw4dSY0KOsT2vvPoqQRBw5877PH78mHJd0rUdJ8dz3nvvNrdv30apgMl0xq3nnyUKI8ajCR9+eJ80S5lOJ3zxi1/k2rVnaLuGb3z9D1EoyrKibVoPmjRBEHmdpRX2puu9niXYZv882U/mAc+wWbtlTZyTRer52pSzHKCzY6KV9lEowsCg/9VjpgemcBgzIroguY4LQBrUrmqrVVJPPKfx4mVre6+xck99VvRWPxX6sZ0KNGEUyX2gyNGBr5yJBVhlaUqaFqRpTJJEBGHsJQGIO+4cCyTC+3MgyIH12idr8dE+0LYN0EtwaBhJ1dUguOcc8LRWMqow2/15ngGT/SRmlM4N+01LnpeSgOAw8pE/HjRpPVTLRH78lxDHGUUx9cYdORab1fqHPT3/zO1jBZqMM9x/8IhV2RDUFZvVkguXL5BkGTqK6ABjHSrQRE5uahpHNsrRUUDftBhriaKQyXgiKdqAwTGdzXx1iNDdXScurulkQgCMihEX+4vkRYZSAfP5gp3dABVHZKMRFy9e3MYGBF4QrZTaNnlrrWUUZ6VCRAOm7eia1gvoAtG5dB1ZmtLUNbGffQ/ZO85IF16WSgr4eDQSp4XWosdSkiqrPeA5WwFIRsgwcx5yp5RSW+egVr4WZRgNwJbJiqKz+ANr8fkpkoHiAuWrG0SL5az1jzkcNUXbSZGy1hBGCaOJ5BOZrgOlRc+kg60gMc0yRpMpCtmfk854110PKiCIYmGJgpDlfI7qGn/SGZxGHCyeidNRRBD6kDhnOX58SrUp2d+bkY1jTN/SWktdGUZ5QV1LWneWZYwmBWEsuVhJFBDnGU3bsak2BHoobpbQPqdDLl66CtbSNYZAx1hl6FtDEMSsVmvakwWj0WirH0vjxNffaJqu5f69e1y4uEcRKapqw3q5pm4NnXFcvHSBJEkJIkWzruj6GtNZutaQ5zOMhbJqifIC0ohyVVLPlzRVRxSkXLl+gMYQhUCekGcJBqjaEqsdaM3p6ZzZbBdrFdYoMGA7Q1uVaJuSxImMA61D6YhsVNC3xvc9avq2o67/Ktzyh9nO33DPfz8ApOH7QWsiN1+xyhvbb88ZpaXQ+eGjh/zWb/8WL770Eteeuca//w/+fY5PTtiUJbdu3ODmzZs477Iqq5Lbd94nTmOeufkMN2/c5K/9/F/j0eNHvP/BbR49fEzVNFy+PCLLU29AkBykNEtByyKVQHE6n/Po8WOuXLnC5StX+MVf+jeI45idnT3Wqw3vf/ABX/36H/K1r32NN77zXR49PuLvxTHPPHOdLC9kzO8Cmtaws7vPpStXMX3L9773FsZBby0GJZl6Pu5FRluKIIxROtyaaIZR23D9+qhNwIrb3tnVwCptpUnujGEZWCP8bdsJXgqUFKe3bbt1H0dhKNfhIR9oq786G0EpHW6PrQ89ECCgJBU7jIKzY6+2XL28jqdAXe98pIBV9Kalqta4Q+d1TBqlAoIQwiAiihKSVKQRWZpv4wOyLCdNM+JIXOlZJoupLE+JosS7MIfgXZEVWANV1WGNiMAldFh0Zj4dgGGvDUDsibGckmmQ//AzBLcqPyNVauCrFFjlXXDDiBVQPa0y6KZnU7YEwVqAlvLHSDlCHdL+GK5HHyvQ5PqO3/qNf8Ff++zr/MwnbhAlETqKsH0vH77AI3qrCJRmurPDdHcqBgDvahtPp6ggZDydAvjcHwjjkBvP3uLGs8/7Jxue1Hgln5UlgtbM/coryzJwPil3ZwfnReRJkm5dGMYYWX14R8YQNKa95sh2LdYqCGLCOJYDrjVd24FSJIEEqzkrr0MrfLItZMWIrBgx1BcYO7RC262jASeOCu1zPtq2lZGM76CKk4QginBDKKgWhG+3jdTDBx2hkrcfblml9L2USjrrMKb1Fy5HU7ckaYbWEWBRXrAZxYnPvXKEOiYiBqW3tLZcGxSm6yVfCtjZ36eva5q6JwxjnF955lnms6sECDotJZdBFIAegtAUy+MTQh2RhDFaRwRRgvM6q6RI6eqWpt6QxBFaK2ygpL4lS4UxKisRSXoBY5an4CTaoK5qdnZ3WS9XZGkOQFNVBEFImk+pyjXGtERJRtevpSdM+Q+YF4GrMKB1HUo7Nss51GvCQBOgCVHoKCCIcvIiw9FSTHJwDtNBZBRpmNEbceM1bYvTltZ2qDAgDGVlaDvDsjyl7SqM6blw8QLaQm8NNgqo2hq7XBIoEZgqNM5Y4iBEp5ksqo2REFkL5WZDbzfUZUsahnRNh9YBSZz+WM79n5bto7RMZ0WknmmyojtxyGdfawFTfd/x3u13+ZV/2vDqB7f5xCde5srBFcaTMcVYmKXHh4ccnxzzrTe+zXvvv8eDhw9J0oRXXnmFixcvcvXaVX7+5/8adV1z584dlNI899zzvPap1wBYr0o+vP8hKMV4MpYOSNPz7nvv8e1vf5tPf/rTvPjii3zilVe2QO/k5IT5YsnR8Ql3797j9HRBkqRcPrgsWXnTKc89/7y0NgSaJMtYzOe8887bvPHmd7h790PWa3HeaR0KO4+MJIfy4CGG5Yl9+SfvZJxnUNiO1wbxNluh9/nDoAbWydrtr4k0QgCdcgqlndcvOXoji9FAD2yV8Mv6qbGh8noox5mwfOjC27rmlNqahLavwbkn3t/AolVVRds15zKZBOqdL7OVIM94u1hP05QkSb2ZSVx8IsXIt52ggyNQNFYxUZgSRyJUZ8vYPbnXB+fiWRXNILC32591w85h0NcOLJUPRQWMkf253Ufb2SfnRqbyjEMsxEAQ/NSDJhVGHD56zPzkhLR4lQjD8vEjRpMxURbjDPRWaFTnXV1d35JEgdggw5DeGpw1MvpCoxwkceJBjcF0lQipnZwo2sfvW2twtsPYnuXxIxlrBcISmLZjfnJK/fAxKgyZ7eySZRkq0DR1Rde0JD7TIwgjsJKqGoSaOIu3KDuIItIgkI6nIMAZh+t6ej/vDnVAEkdESUy9XqN9ZYozhq5vPYJX3mYvDI5B0WxKPyIMWZ7OmU4nvm27F1G9szR1JaOyQsZHOG9tDgLJUOrFNSdOjVhKOVVA11YikA6EfVJIQGXXdpLQbRAQA2iffG5NL+PHKKJtG+YnS4YE4UAHFKMxcRRRblZYY7ZRDnEmo86uaSjXa9IkoZhMOT56JPRskhJp0TqIKFzeR93URIElCkIObt4E7ejWc6p+TZSldF3PaDIh8JEENC1t3+O6jvVyQVVuKMYFbdfQtT1xlJCkGat1yWq5ZFyMWByfkB4cEMUxD0/uEwSaq89cJ4wSlssV070dZru7chx9FMR6vaZsKiIrx+D5F1+g3mx4dO8Ou7s7XH7mKs5A1XbUxrJYzNGBozGdHwOIYSAIFcUoo+k6FNLLFWgtn8FUs16sOXz0ALTBYpmfHGKMJS8K8vEYghBjLNk4oWkqrLGkaYx1oqmI44Qo0H7UYLB9z3q1ou17kjiR420MxXgmi4y/2n6s29YC7w0ugZYKo7IqOZ3Peffd9/jmN78pIOjqVS5durRNfd5sNrz99tu89dZb3L9/34/S7vDlL38Zay2vvPIKn/r0p3np5ZeZz8V8cOOZG2R5Rl02fOMb3+CNN94A4MqVAw4ODui6ju9///t85StfQRLwJ1y7do0sy7ZGkjAM2ZntsLe3T1mW3L59m9///d9nNtvl537uC3z+858jjEK6ruPevXv8xj//DX7t136Nt956i5OTE9GgRpFc10L5XErNR+VDF+Mffj/+SX8/NwrbLuaG/+7BidwjBHxt3c5uWLiK1EKAS4hSwTmR9Vmv3HkHmbWW3vTbMEqtA+I42v7sk4Lus/oWrc5ricDYGB0gcTHnSpvPRn9Guk67s87C4fHOs5xDfMDwJYXBsYjOo4TJeIcbN16gyMY4dy4ZfQgG9Qnpkuc0hJX69+nNUecF/VIp5bVgHkMZBdKVJ/vFbn/+bCEhqecSczCYa7a9fdZiuuaH/lz8WdvHCjSNxwW/9It/i0+89LJky/SWu/c+5NlnnyHNRGRt+xbnND2ao+Njlos5+/szspEImJeLJeWmZHdnj8lkJiK9XmLc21osi6brwWt7Yh+pH0QaZwPmh8c8fPCA2c4es50dHLBpGt5+9x3mp0uss7z00sscXL7MZHeHujOsTk8x3vGge19DEoAKFEEkgrmub5Ai2pC+bsjzQgTVpmfjA7+SOCLQijjLOD46JEsSJtOZUO3rFV3fMZpMyIsC0wmQ0sGZmNA5h+17Ecp3HdOdKUopTFNTlxuariOIIsLY08HKC++spWsaNqu1zMlnMwIcpmsIo5DVakmSZYzGY+g6HIrp3j7Lo0PmxyeEgawd4ixlsjMjzbItSJ0fH/L++++jjBWBexRzcPkK09mU08WczlmuPXONJMlRvbyHMEmY7e5LnpLSvP/+XcbjMdefuUGapPIefU6KM5bLV67gla4409I1Lcb1THZmdG1Dby2T8RQVJrRNg6laX6rZUFUNQRSTjyY0TcPjx48psoL9C5cY5yPSSwkYy3g8xrQ1bbnm9PiIJE1xRtx3YRyyOD1mZ29PIhy0pi4rTuZznLPEcYR1jjTPyWe7qMePqRpDZy3GWR48vM/pqTzmzu6Me3fvogPls2wsL734EsVoRJoExGlIVXVUqw31qpTSVu3QoWJnZ59sPKGuGx4+OkSpYy7uX+TKlQNeeeUTVKuSsqpYlhvSLKSqKuanJ/R9x85sxu7+HmEaU84rMD17O1OiKKKuG9q2JQg0Vf1XmqZ/3e1pd9357exm4Tsn+x6rhXkIw9AH0jbM53M2mw2Hh4eMx2MvEE6p63qb0TRk6BweHvKlL32Ju3fv8tprr/Hqq69y48YNf4OMKcuK2+99wJtvvsnXvvo1Hj56xKc+9Slee+019i/ugNU899xz3L59mz/+4z+maRp++Zd/meeev0XoNY/Gx6XkeU7f95yenvKHf/gNRqMxSRLzmc+8zigeodD80R/9Eb/9W7/Nd978DsvVkqH3bYhbaZrG3/yNZ0GC7T57Qgf2J+3fP+dxOf/Y25v++WP1FPAYeuv63tJ1AoYGl9kQEzFY8Ye/D+9rOMZKnY3mnngPnrGxSHfdsAmQTs+u9/ZslGe32qghi+lsDPgEUHSSzO6c1KZAs33/A9uUJEc8fHBEHIvDb7hPSrhnQpomZ0npSUIUxVLQnCQkccI2L2p4Y9o9ofUaxpJKK/peSIOzBZnzrxGsVRjTofqz8Z88CMRR7KcwP9rtYwWaXnj2GqfHh5yenNJfuUAaxbz2qU8RphGurUFrkrQANOiIa/kEd60jCIVqdT5E7O6mxOHou55qXTIaj6RaRGnS6UyezEm4mcDaAKdBacX+wQH7B5dBBzgUtmlIwoBP/8xnGO1eENDgU0r7piVwjr3dPQKtuXv3LuvNioMrl9m7LAWJXdtijPNFhIpxOiNJFWiN6SXssphOINIcPrzPtWvXsKZjPj8hvXwFFWi0ceRFgbWGMI6pqorToxOcc+xfvEgxHnldFFy6ehXTdx7wWWzfo4OA8XRK0jSAOPHqWrrQlA5QzpGnKYmPO2ibGqU1y9WC2d4F9i5eotmsqVerbYp5qBTj2S6TyQxMj+s7XKDRcYQzEmimFBxcvsrBleuAgl5OCmMt9XJJMRqx3kinUYGU8NZlyc7OHrOdXfq2QQM/+5mf9X1NPW1Vb8eTUSzap76ut1Ze2zc0dYlxhqgTyrhpWrqRJVbivkmyDI3yoWwZOgoJtCKNYkb5iLps6OuOIBarc1QUUmpsDToo+NRnPws4LyjvubC/z2q9pN5s6JpGYiDShOeefxYLUjWw2dC1LUmS8NInXwcHm+WCpml47hMv05Vr+l76xS5euOSDTy1luQGkIzCKE/q2I88Kiqtj2q7BOkuciq7ONi3desOLL78iujMtqfWr42MOHz3imevXGY1zLl49AK1Zz+fgZGXYdB0nx8dMphMRdCiH6Vr6rqXxiwxjLat1+RO+Kny8t6dHcT/Idn7UIQ5Ptb0Jhz6wV3rcWubz+ZYxOv9cg13dOcdms+Ho6Ij79+9z9+5d3nvvPV566SX2L+yTJTl9Z3j33fd44403eO+998jznJs3b9J1LfOTJXGcsLOzw7Vr1/iX//Kb3Llzh4ODA/b2dphMJyhfLDywKM452rbl+PiYu3fvcu/eXV544XlGkxGnpye88e1v8+03vs18MZc+xix7gg0ZymSHm/QAIH4c21bQzJOMzPmv4d+0//eBdUmShCzLCMMYY85ed9d122vUwOIMx/M8mzSAwUHT9DS4gXPgyb9/sfYHT0QAnI3Etu9k+xl6Min93Ps+p7EbgNxQwqyUZr2ueHD/MZIzFW3ZPhn1+UoZ/zXoeiMvVJ9Mp1vQr8+NniUYMxRxexhKOwMywYiigGAY6W1ZqYF56uUeOrxHr3tSWSqaux/x9rECTTeu7TAaKdJE0ZuGcrVgZzpGdRalUz9mAGs1QRSjQ1itSh4/+hDT19R1y97eHqPRhHw0kTTvwKeII6nT1hjatqPrmm2nUhyExGlMnImlsywr7t+/z2g05tKly+TjqYi4N2viNEUlEYEOWC/mki01GhHGEddvXef927c5PD5CRVIxsilLjBVt0Hq14epVSxKGVFWNDhSTmTw21nD12jXRJkUp127cIoliYZKimL4sWayXZMZRlhVvvf0uAHGasxsmcuNrpShx98I+GEtdNVJdEIVo59DGFyArTYCirWq6uiFKxCWllKKra/quYzQakWY5SkHfNERpKuGYlYw337/9vuRYWcVkNCbJUvrO0jcNKMW6XHuh/VSKLB2gY5yzaGfJdy8RVxt29g9wON59922iKGB/dxdreo4fPsA5Q101JHHCeDIhDEM2yyWnp6diz5/N6Pueo6MjrHEyZsgy0iCktZYHHz5Aa8VisWaz3rAz2xcXiC/tdL4D0PQd86MTVquVf9w9sqyg3GxYLTdcvHSBx48fsr+3i05Tms2G1ekJOorZuXyBdrMh1IpvfvtbrNYrrly9yvXr1zGm57333qWsK565cYMojKhWG/JsxGg04cM7D7h75wP++t/4BZLRjLBtUcrh6OhbyVWJdIoxlnrZUtNydHTEcrlkOptx6coVsjTm9PABdz74gLZu6I1hb2+PZ194gSgMuffBBzR1zfMvvshmtcJWckNdrdc0bUOW5+x5/d/Ro0fc//Ahe3t77O3uU9U1pu+YzKaMJhOa5ZKjw8d/gVeIj8f2UUzSDwKenr7BDUBkqCc5bzyJomh7cx5+93yH2SBgHrq5Jv78OTk54Stf+Qp/8Ad/QBRJPUaSZNRVzXK5omkaVqsVv/mbv8l3vvMdnnvuOV577TWcc9y+fZvFYs5yueSf/tNfQynFF7/4RQIto5myLDk8PKSuG3Z2Zty8eZNXX/0kL7zwIrPZVBYKPrSw83ErSqkt0Npa7r2reLghN03zYwNNH3UMhuN3/jmf3r+y4Eq4cOESzz//ApcvX9kygIvFguVyKeP5sqSqqi3IHY7JE2BM6+1jDl/n3ZRPgjlL37snXtdgCBoA0PnPDbAFPIPJYMhtOmOmhHSQEZvkPSkVkCbbZ8E56PuG1bpmuTpjr85/XrXWpGlKlmUC8AMpQE/TlDzLyHIvQPfZU3EsVVt5Jv+u9UePYJ1z9MN7Mka0xNbSdw3VT7um6Z/89/81zz17k7/9C5/llZefIcki6rYmz3N0JHUWUgLZEwUa04jL5/r169TVmmw0ZjVfMh6PRQvUnxXS9r1QojqKCJHFdJKIq0t7keFQ+RgEAXt7+6S55AU1ZUVZlezs7YO1lJsNzsmHUbKHZMUSxzG3nn8elCMIROk/Gk9ECO1LVyMdeoeIAInjw2PiNGZ3dwfTdagw5IPb73Fh9wLJaCxApqzoO8N0NJMVTpjwuc9/AR2GxGFIU9cAjIqCrmk4OTxiZ3ePMIyoyzWxExu/dhBlGaCIonC7InTWYvyq9c6dO2it+dSnPkVdVzhnefToEaEO2NvbI00SwjjmuVc+AU0rie1Ko9IUuhZbVyRJxoXxZfC6AJSWvrq6RinlAzxbotEYlAUV8OzzL4LrCLVoAzq/wt67kAhb5hQ6CNnZ3WUyncqqLQho61ouSlVNXW1o24Y0yykmY3GVZCmXroa4riMIU/q2p2sa5ssV69USlKOta7I05fLBFZxSlMsNlS0ZzWZkY3FXXtjdRSuH6luSOCQ6uIQxltXhobgcixGf+/wXOHz8kKoqaaqSnd0dPvPZz3q3iWM9n2N6x3q1oW969nZ22d3ZRSvF6uSUvuvI8oRkVKCMoatrgjAky6R1va5KLly6xGx/n77tQAfoMCLLCg4uHzC7dAllDK5pefzhhzLum83ouo6ToyNCf+H84M4HcoHzujy0ZLvESUyWZ+RFTpplJEXB6ekxm8UC7U0IO7s7P/kLw1/y7TxY2lq0n/p+uCG2bbu9QZ4f5Q2VTAP4OP/nMP4CpDnA38TFth2gnJKk/VwcV13Xc3JyzHx+yv37H/LgwQPCMGQ+P/XVFzVvvvltnDPcvXuXoii4c+cDTk6OiWMpvc6ylNFoxN7eLvv7e/LfFezt7/Hqq6/y6NEj3nrrLcqyfIJNGfRRSZLgnOQZ/aQA0/nt6WiI84DmbJ9LYnccJ+R5zmg0Yjwec3Bw4MMZ2+3+Hr6apqHruu1/H0pnhz8H1udptutsLHhe+wTiXhtGcoqh886/C5wbnJpDXpQXU9th8nhWv3LGTnnWcmB+ENONNR5sDVU0bsirOhtjNs2G1drvoyGTaZs9FRP7UV4cRYRepzQqCqbjsYxotZbC4EjKg+Mo9gXyktcXhWKfGaZ00Y8B4XysQJNWlu/+8RvcvXsb4zSjSUHXlCyWxyRpRt10WAKKyYwg1vRNS11tKG3L3v4epyfHfO+7b/Hapz5NrkLqsiJNU6wTEXa5WmE2G7I8J05Tural9EW6cRxJMXDTMNvZYTpO5EZ3Ome1WgvLcbokiSOSJMP2HaoXKz4K6motls4goqlLqr6V0EggyQsgoFptmE1nNJuKAE0UBkTTHdq24ejDR1hjmEynbE6WXBjvcv/d25RVRYAiiTN2L10kTvKzHdYbnHXEKhbxdWdJwoS2aqDtoJXX2K8rVCIXzeXRCVGSCPGjtWR+aI2zlslswgvZCwRRRFJkhEmEM5ZLVw9wvSEK5ON0+PgR9+7e48Yzz7Azm6KUZX7yUDKo8pyur5gvj7HWkkQRSmkBc8fHBEHAwdUrZMWIk8P7fPVrv8+zz7/A1evX0AoaIwL25XLJarVCoUiTlN2dXcn58PUReV5IX5Lr0XHApd0DsiynaRrWmyXH88eEoWbnwh73P7zDpd19ir0xh/ce0DY1ly5fJgoDmnKN8g5I0zQiQu0Ni+Upm/WGPC/AGR49vs9olLFel+zu7ZFkKSfzBcp2lJsl+5cvk06nXJ+O+d6b3+a92+9yvZc05t6vMPuuQ6Ep8ox0NMIai+ukYPje3XeZTqd8eP+UZ248w3g6RStHXVeUm6W87+mUfDri6PFjHj34kPTkkDzLwRomkwmpDlmvNgTAdDzl+OSEuhT3X1WWPHz0IRcvXSQIJFguSlI2m4r5QhYao+kMHcY8fPSYTbkhDEMmkxHWGA4fP2J3d5fjo8Of8FXh47/9sCzTeeC01b2cYxHO/855V1nbtk/c7Icb/HnmamCrkiQh9CXXq+WKKAyZzaZbRsSYkKZpODk55tvf3myBTZIkhKHm5OSYr3zly3z729+SwFvPGE0mY7quI44jX4mx4uTkGHDs7e2iFLz00ks8ePCAe/fuSeadf0/DeG9gI8qyZLPZ/LmF4H+e7YkuOs/0DGDiX3U4Wj+V+JC6bsjznKIoGI/HzGYz6RD1AHeo/hiO4Wq1Yj6fb8HU+b61YQJy/ufPM4rDdja+GgDUk4zV8Fq7rnuC3ToDS+opsPTUKG9IaXYOHQhgC1wADEyYOvc4nhHqzZYBNdZg2o6mkTDepwlY+R1Ik5Qiz0niiDiS8XOeC4jP80Ku90VOlksuYhLFhFFIkmSQ/ug/Fx8z0DRYE4cALAvOECWaMFEkQYTSEVGs6PqKsi4xtiOJYgI/ex2PJ+BP8MBXmzgnVtAoTdFdB85JEWbXS/N0JM47i4hvy7IUdktJaOVoPCJUIqxbr9ckeY5y1gszxamA6WmMIUmga1rKaoPFEsYReLdEmqaoICDESuFuoAi1RlmLCQKifETohdJxlpO3higQgWScZIQ6YHl0QlM3pHlG4ssTsQ7tFH3dsKpK2rZlPJ6wWS5Js0xqPbTCOkeolRTbJumZrVNLfUdSeL0YYI2AvjAIKZKczXqDUZBkKVmdsre/TxyHnJ4eoQOF8j1Jxnb+RG+J4oQgkiyNrm8ZT0YkSUxvOk5PjgjDgJvP3WJnb4rDoMOIMEzoTcB0d4fElxZFcSIuxq6T9Pf1krIuhU4OAnYu7FNkI7yxF6cKEjK0s8RaM8lH9H3P5vhEVjxxTKAhm07IRznjusY5CGOx6Y6CGPSSdl3S64owlHyqNM0I0wwFNOsVkYZsukO1WmGahsePHpJkKdPplMQLQDerNSoMyIpCIhSCkFCH6FBjlMWgifKEC5cvEAB7asrjh/fpu4bJdEKaJ3Q2ojwuqao1UaSJlGW2MyIKY7GToiiKMaBlRNe0KK0ZF2OUVpSbkqqpme3uMt7bYzybyfkWhfTO0bUNoReOFvkIdh1ZMcJhCbXGBT1BkBBpRXyuqPqvto/ens5k+qiYgT9tGwDS+bHH+Zvb8P3wc4NwevjZwb20LWr1eqinHVpRFGEDh8vl9S0WC05PT4miiCtXrlKWm+3iZTQasbOzQ+3NNKPRiMHhVlXVEyCt7w3r9Ya7d+/wO7/zO7zxxhus1xt2dqZcunSRppEk8NPT0+3vhmG4fc0DwBhe4wAcf9zb+eN1fr+f//uTGiDji5JPEPY+2up9hjHjoAMawNR4PKYoCqbTKXt7e9vHPP/YAxNVluW54tpKqrLOjTebppEIgral9cG/A9gbjsUA2s6P0rYidO2Tkj6C0XROdEbb4AX3tOZKYcx5wDSM7CxaC8CS4uFB5O6nOUNO0zmwhbOUm+W2FDnYslRyX5b0cGGeZOwn9/Wd6ZQwjH7kn4OPFWgyXcPFvSl7O1O07lmvF1TlmrzI0IEiTxLQkbRe44jCgDjIScKI5XyFs3Dz5i3yrPB9dOJkMp3UhKRZhkskfqCqK9CQpClRHG0/FEEkKyzjwUSUpaRFISnT1hL4dmkVaPBp0ChQcURb1QQ2JMoSslCBxp/0EUqFRKnGOUWYxRj/WQqR8t0imGC1IkgSpsEOTmmmuztiA9fSdddUFW3f+Togi1UOFSiMNWgNQRhhKsemWtPaTrI1lGO5WYITa7+xUqS5XFekHiQ4rdChjAz7XpKFJ9MJR8cn5FlGqLXowRzUZUmkNQcHl3GmozWScp4nKQ7JEZFMEe9uc4qybOjbjt29PUaTCavFnAcPHnLx4kX29y+RZAmr5RKtNFmSoBAnSKIjur5HGTlxkyAicEjHXl/6iIeQ2WyHxfExkQ6E8raGbCyg94M7d7h+7RZHh4+JIsVkMmG1XPLhnbtMZjOSJIZeUtN7K6WV1khcwmh/H9NbtLZc2t+X0WSSYnCYpmUEBElMZBymbtgcn2DznL3LlwkmM8r1CmcsWZaDhfVihXESepkmMWGaECQBVblhZ3fGaj7n4qWLPHj4gLv37hA+CpnOZmRFQZLGWNOyXi/QgWYyHRNFCcpID2NVlWw2JXEYYAONShKUc0J3pykuksqFzWZDU1dMp1OctdumcWcdWZySJRlBELF3YUwQatqqRCmI44BqsSD6cfDhf8m2P80d9yf93NOr/AE0PW1df1rwPfzuYPcGtpqmoX5kyFN6miUJfK5PkiU+fb4FnGcQ5AYnWpdARilaNDHAlvkZnkNcoj4nLZDryMnJCVVVonVAWW5IkoTZbEqapk9oe4YbeuyvR4CUlcNWRP3j3J4Wf5/fp+dBzbANY9K+73yJbEfb9mdi53OAZRDuD0BpOp0yGo0YjUYURbEFVufBVprKaHN4nvNaqAE8DWB1syklT66ut/+96zqsEd2PHbL9zLkRp5P3IK68Yazngek50b07S2DYslpP7rfz3z3JcA6jwbP9Kc81/Mz5L2sNxvUY652Czj2RVTWI8IdQ58A7/EaFz5H6EW8fryucNVy/eoXZeIw2QzGsBBr2na+b8EFaaRJ7+7mh2pSs1muiOGZvMiVKi20YFli0T+lGaVSgUEFEHgQeFUuytvP2+ziL5fexvsxXo1SIC4C+YTTbZXhkoS8tYRLjXA9KqlB0GBKrMQxRkR48GNODk8TYviqJ4hjiCEVEEDR0dYXSKXZYDSK1AUoHWJ/MO97bkV46a8AYVBSgW2GKgiRjlkW4QEpsJ3tTetOxWa7Eolmk2NawXCx5cP8R+aggLyRwDi0Xh7bvyfKM0URYCuf3f5ymtFXF4vTYu10QSn//Ik1VgbPedhwRhTF9X9E2ljAwWKNxhBijwGgCnZAkOVXdstlsmO7sUJUtzvTYpCXQelt0XNe1XGynU+IwQmlHmmQkkcUaS9/01EtZieV5Tt+0rDdr2qoiziIeP37MqJiy3qy5dLCLs5bl/JTlYiGr4jii8dqtyIc41nXDaDzl4uUrJHmKbRuKNKVrG5qqwaYppu8wGymSdjgCHbC/Kzq4RGmq9YZ6tZHYhEZWjsePD1mXa/b294WxG+WEccJqfsru/j5HJ8dEUcTBlSscHR1x985dlqslo9GICxcuEgSaum22TsG6bSnSgiTPOT58zKOHhzz73LPUdc3+hX1uv/0OCsezn/wUYwxHH96jr0ToG/mx6en8lNV6jVaanckOWCmKLtqcYpQTBSG9s1RNg7FmezP7q+0H2z6KIfmTQNV5we/AEm0dVn78NQh8z9vaxcEVslwuPUMQbd1LZ0G4Z7qngc3oTS+akWAIkI0Yjy8BsFrNKcsSYwxFkQOOk5Pj7WN1XYtzVqzi2CcAVuTz2bqupapK0jTl4OCAtm04Pj7eaqyyLPOP1W3daEEgC5+yLBlyqv4itvPgdDge51k60YOe9eENr/3pkWjf975sdsmDBw+eYAuHfKOBjcqyjKIotqzUAK6GsV+SJBR5QZwkmN7QG+kfbbuGuqrZbMrtmG+9XlNWJXVVbwXpA4MnwKql66TGZaiPGYTggQ5QWoqS8QvYj9LaPQ1mlfLap62Tz53bf2eZTU9uHrT5qhjl9VHODgwZnAdkMhkBpywnc5m6/Ki3jxVoaowjiWJJHQ5C8tGEUbRLoJxoiPCCNh9g6ayvVVGKvCiY7uzRNxWLo0fY3mD6HpB8nCAI6NpWgIo6oyF7v9oJwxC8dTaORNUvFxaDVhrrk7gditSzVcaarYWyNx227wnCUBT+OJQeLpACvLQShGyM3b6+tqqwWJwPhNwsllKg2vf0bUfoBXE4qVmx1sJyISBQKxGP4zwoXGFsh9bQtA21szjTMcpzVBjI6w00lw8OuHj5qlj2fQqrz7vEKQmrVEoJfQyAItSKUZ4x292Rz20Q+soAcUygtaTwGovSAdlo4sd/isnunszFdUCgNUUccWM8BsXWGWStkZ9RPcr0WC0g11knPXkegEZKkynFUC3ulEZhmTh87pVlp20Agw4ckwsXmB8dM53tEIUBysLe/j77F/ZJ8oxAwfLkhKapybIUPVasVxvKqub05JhLlw/QUYjpDXGW0bUdy9MTTuen4GAyndC2LRf295ns7sjJby26a0FrqqaVHKjxmIMwZrlZkRU5ZV2xmK8oRo4kLYjyMVXd0bSGfBxx67kX2L94mbZtWa83rMuGNE0oy5owjkmzjMXpKWVUc+3KNfLRiPFOSzYZc/f+fXa7jjt3P2C9WrF3cEBv5fPVWct0NuP09JQwDNnf3+PylSs0dYOyECpNF4kwtSxLwkBTlmu6tvJ0+I9+ZffTtP0gLNR5tmLQtQxjtgG0DI9zxvScjV/63vcdelA1aGrOi8IHNkokCto3GNgnQPHATg2jvyfHOHb734efC4LgCddWURSAGGvattlqrp4Gh8NI7nyi9EfZ5H8c20exfcO+G/59YN+H/z6Ao3P6BsRxdvaehsc7vw1g9fy/n48lGCIlhr+ft/UPjNSokPHeUPSeJIn0s/pYiJ2dne1nYisw7zsBWX2/HeltNpvt2E9Ys15K4DsZ9fV967Vt/b/CFg2hk+d1Uef3odZnzJ0AqLPiePnz/F6xWCcEwRaUcpbB5cB38w2SHbzrz4IS5ulHvX2sQNPBhQsUxVis7ga61pBGCcpZ+tagQ4VpOrq2x1mo6wbjYHfvItNZjjUdb775Jv/sn/0z5vM5dS3N9a+99kkuXrzEfH7KrVu3mM8XrNdrlqslDx88oG4aLl64QJqkbDYbrl29yuVLlyg3JYdHj9lsVnRtRxhF3Llzl5/7/Oew1vHg/gNGkzEXL1zg5PSENBbrfllucE5E4nVd07WGIIwZFxOmsxknx8eMxmMBQx70tV3Der1id3eXRw8f8+H9+5RlyZWDK+zv7Unnj7XUpTgwUu902axXAOR5hjU9R8fH5HnCwcEV6qbCmp4iT0myRATxKiDNp+TFlJu3brFYLDg8PgItQWNOQ5pnKAcWR9919G3HdDJhZzojjWPSLCWKE7/6EnBDoLeOQPrOnxleJQ+gxUHXeXCotMY0rYC9cChqlEJH660dgfarbGuFSXYgTYLnSiOd9CHpLbVsJMZLawzy38c7O9LrrQMIA7I4FKbOyYoon4zJ3MhLuxRxUTC1CoXcBFSgsYmAap0kjJKIbGcmvXdRiGmkrkZFEU4pHAFROGUnz8E6zz4G6Kwg299D6YC8bsA5Aa7OoYl47oVXiBPp/BpP9ymm+9RliXr8mN4Y2g6MDSiXJScnK6qyBGepu54L+7u4AHQScvPlFwnTMV/9+tf4zd/8Tb76h38IOKbTKV/6vd/jk6++SpKmXD444BOvvMLVK9ek29BYulayw3Z3dymrNUpBkkSEgaKqa8LoL2bl/5dlO29p/5O2gXVR6iz/6LxbLgzD7Wp+2412buXfdd32BjRk+pzXPA1MVNM0mL73lS3yXG3boIOAPMuJ43h7o5WbtOhH5AbmV/0KQp/Y7Kzbio6HzJ4B+K3Xa6y1W63P8N/Pi5wH8fOwiH1a7/OTOB5PsydPO9jOM0Tyu+d/lidA03mQOxzL4ViI81o0jsM49DxUGJx3AzvknOXpzKQ8zxmPx9ux33Q63YacDg7EAVwNoO+8q2+IQxiYveGrbmoxoPjsuAE4nhfJw5l7T9go64uInX+vT+qfZF+Le+88qJLr+ODCkwT084J7AU12u+/csFgGokiceD/q7WMFmv7h//H/xNXLF3jppRcIw5i+azl88BBrekm9TTKMaXHWEcUpSZr7D6rl7e+/xZe//BW+993vcvPmTX7+538e5RmXGzdukIQxnenpu46DgwNZgQHL5YrlYk5RiFj49PCI55+Tdu7lcsmNG9dJkphyvQGlePDsA159+ROsN2uuXLpMMSpIkpTZZMrB1atkSSKgyRoc1n/wLXGcMh6NCaOYe3c/JPQ5HzLLlxHR/fv3uXZwhUt7+7z+yiuoOGE0maCMYXFyzGazovUFwKNRgVKarmuJk5goDKkrSXyOIln5dX2Dc8bbfwuUUqzWJU7F7MymtE3Dl7/8ZX7vy7/Har0WkGRlJHp6siDLY65evUrbNCwXC9q6I9AhaSrsWtNanFY4X6Qo58/QYyehbUGgCUKxwfbGgyOkw0k73z806Di8vVVpCENJivUk2NlqZrDBurN1nlzMRHiolUNioYaLl0KhuXx5l5PjJS++9DIvPHeLcrNiOZ+TxPG2MTsIZcWdpQXOQbkpicOYdFzQ9A3Pvfg8ZVWKgygIMdZQjEaU6zVJmrK3u0dVVSxXK+I4YX9/n9lsRmwttrG0dcvuzh7lpgTrKIoCrTVNVZNFijTJ0HEKCFAMtEaPEw7iHOv7Bwf625oeazocliAKiaOQfDJBhxFFGOKU5t/7n/99vvDFL7K7twdWNCef/fzn2dnb24pWx8WIzAdnKjS2FeYyyTPSLqEqNzRNSRCFTHd32Gz+Ktzyz7MN7EkUReeqND5aq2Ot3QIfYCuSPp/F87QGZ/j3p0XoT49Szocgaq0l/kR+A6UdYTTc4I13BluC0Bel/iuMgRtumXKTtLLwCUKxwbdtsw2+dFg5z556DwOIG17r+df89Hv5Ybfzv/80G/KnbU+DqKeF6OfB6BmAGMAA2+/Pvvz+VWz/vr+/y61bzzL2wcsbX4WVZeIaE3au5eTkhMePH9E0LWEYUZYVddXg3Fm46Z07dwiCgCzLniiQT5JkO+47L04fgk8HMTqwZTOtleNkjWG1XvmuU8Pp6YLVaklV1X7sV22ZrLbrqKuKpqmxVuJ2BpAV++vrEHmglCYMgy2oks+lxTipxRKRzFkHquia/PdKelMZtFbe9f2j3j5WoGkyGjOZTAmDCGssTdvTNC04R9cbKQ1VmizNCNMcVMjRo/v8wR98lW+/8W02Zc1sNuETn3iZl19+kbIqWSzm5HlKsxFk/e4773D54IAg0JRlJReotsNEiTRbhxodaqq25nR5SpImFJOCsItYrpY8c/MZsknByXJOmEZkhazIZuwShAF121LWNWW1ASxZljGajjG95fBEohPSIuXSxcvc+/Aujx5JaOLB5cvcyn0Evfa5IFECWrNazNER7F/cF4eaOWuPtjbbUuXZOGfv8gW6rvV9RWCdVMj0pmO9WjJfrEmSEc9fusw7b32f4+NjJpMJeZFjfAZHnKRcv2qYL44JveD98qVLRGFI08h4arVccPHyFVSgcNpJXhWKuqoZgtadv/E7JQJxazsR+yECZdN3hDrcRhlMZjNWm5JHjx9RFDldJ8nl1hrfjyej0t2dHfZ293AO2q6lyAuOjg5J04TJeMxyccrpYs7O/h6npydUm5rxKObrX/8qFy/sc/OXfoHVasXtvkfhCHzpZOhvINlQFZAk1FUtjF2oufvgAY8fP+LC/gWCMODDDz/0F1PlO8I0aZbS94ZyU+JgK/Ysyw3r9Ybd2Q6LkzmmM+zs7hAnMcvFitl0hsWS5yOSJGK1XHPnzh3KsuLa9eu+hDPi5OSYqqrJ0oSqLgmTkEsHl+jahrfeegtnHUkcU1Utk9GImzduULc1pu1Yr9d0vWGxXEq9gxdVgqzyinxEkeYEWnPj2VuMd3cJ4oT2UKo5VpsN053ZT/qy8JdqG3RF5wMrnwY4A2japiqfA1cD4Bi2jwIAT9vSzwOQpzOP9KAjQZYXAyNhndnWdwTBMKo6+90tsNhKHYbn5lwFTMdQoyGA4V99DWfC4T95+/OAph8GJP1pz/U0gPuI3/hXANJ58Hf2mGdBkuDIsowLF/Z9mHJJ6UvDkyRmOp1sWbq9vV0mkwnG9GRZTllKCOnG65ZWqxVlKQuZvuupogrnJElea70FYQPrNGjgBoZq0EoNI8cojIhCcVfHSbwF4pPJjNq7jLu2pW7qbTxF0zRUZSVaUj/qk4iVHtQZu9W2Hc5Zf38YFn8Oi8E6mQrY0HfNBUOkgaAkpRXKui0ZN1x3LT/4sf1Bt48VaPrv/tv/lqsHl/h3//1/lxc+8QmiMGayI7bMKBTBdhDFOOe4f+ceb37nO7zz9vc5mZ8ymU35mZ/9LAdXDphNp2R5St+3KCt6oSiJOHz/EVmSMC5yTk9PWa+WTGc7jEZj4jDCOcNsd0oYB0RJxMzM5KKjFWmR0fQtu/t7qFATxSFpISWtCuiMYblaEgQRFuhMTxBqkjwljCOadoPUqVriLCOIQxGfBwoVKlQItnMsyyWXL1/GAaulCAjLzYog0OSjDB0FtLXosJI4ISSUk66tRBCaxszXC+I0JQzEwdY2LeuNdM+NZzMm413u3r3DG298m81mw6VLF8lyaa+3VkDTaFTw8OH9rWW0yHLyIqfvelbLJQDFuMCqHuN60jQhDhNM29N1PpW2FW1FmueoIKCsKspKRNdaK/q2Y393h6qUEtliNObo+IT51QNu3LzJ6ckJp/M5WZYSBiF916F1wO7OLvt7+/Sm5/j4hNFoRLkpsc6SJil1XVJ3NbO9He4/+BDlNAdXLtPUDa+//kle/eQnWa9WJGlKGkdo7URwHsfkWU6gw62VdbXa0NmeRVVilDgop7Md8jwjimIBIl2/DXgTXUK0pdWHVbokMzuOjg/JEwG6q/WS3OZ0fcPR8WPG4zHHZUmgFXXdsF4tpTLn+BF937O/t0e9WbFerXCmoKordKfZ63aEDS03LBcLdnZ26NqOti6ZTSeEgeLuB3dIkpSj4xMZLXB2w3Jey5ImGdPxRBYTSnPwzDX+zb/zt9nbvUTZ1Dx6/Jhvv/nGT/y68Jdpe9K2Ldufxaz8oIzLnwtcwBbY/Lk2eYAnvhf2Vxiosx96cvthXuu/jrbpB/3dP+31/GmA6fyf7pzb7KMf80l3nrWW+XzO4eEh9+/fB2A6nXBycsLOzg6z2YwoitjZmRGGIePxeMs2fXjvQx4+eAScVebszHYIo5C2bTk9Pd0Cp0EDN2RADXEOg/h8m+Dtx47DSC/PM9I02/YCFkXObLZDlqUopX3cTi9/9iJK32w2rDdrqZRpWlar1fZrvVnT1M22nFfaORrpwRs0wAYxWGlgCN7Ef0aDM0eeEHqeefoRbx8r0ORMS9c19H3n3XKK0cSLa/3PnByfcvvdd3nnu9/h9nvv0lnDsy88ywsvvcDOdJfRdExdlZyeHNF3PabrmJ+csLO3j3KGa1euMC4K2rpBoRiNxxJW2XX0vWGUZ6KR0VBkqVDpvjE61Jq2qUVHZA2BimkqqR1p6xrbG9I4JYtHpFGAjiQ8Egs2TSiKEVEQgQ5o64pJUZCn19jf3yVNE+anEv44LjLqpgHb4fqOOPSjPBx924A1ognqey/gtoT+g661put68lFE3zcorUVopwMmsx2uXrlGGo/53X/xe5yenGyzMDxrDNbRNw2ls4QogjDasjCtP+GyLGVnZ4eT+QmGns40dG1DHMRglawwUDgjY7gw7BmNc1wK1lhJlHWGOA2ZTmdoB23T0rUtSsH+/i4XL+wTaojjiKLIUEpjjZC2YRAIMGpq1qs5dbUhSVKcs8znJYEOmExGmF5C9tIoZXdnhysHB+Acp6enOGuJwoBilIOztG1LFAtz2NQNm3JDFMeEcYjrHK7rKYqccDqjyAtGo4JJlosNuKxwQBxFNG2DMT27Pi8sz8XZVO/vE2jNcrXk4OCAzXpD1dRMJxM2mw1t27Iz3aGppRrAWsdzz90kjqTm5fT0mL29PW9o6IjTRMLjcEx3pOH+2pUrPHjwgMSvJterDXmayaj5dM6Vq1dJk5TGCzzbVj7X2/Fo31OVGx49esz79+6x/B9+k+OTx/x7f/9/weUrN2h6w9e/8Ud/AVeGj/d2/uZ5/qZ5fjuv4xi2j4oW+NMe+8/56p7684f/9Y/CFNtp+tb9pM7m6X+O7UcBfv51tj/tcc/A0p/+BoeHGMZSWutt6fKjR49wzrFarTg6OqYoCnZ3d58Q0+d5zs7OLkpp6qbGWEOSxEymE3Z2dtjf3ycMA5q6ZTQq6Lp+G3ngnGW+WHD4+HArArfWSEF7mgi7FAZodTYSLoqCKIr9Y2RMJmN2d3cZj0e+NFjYoDCMiCL5mSiKGY8nEuWCkkV/WW4DO4cvYVNbqZvZrKmryovOve7LG616023zp5waXOuiZ8WPrX/U28cKNP3P/u2/S5qkzKZjurrEokhDCbO6d/ce6/WGd976Pm9+81tgDTduPMe1G1fZv7RLXZe8f/tdZjszNhsRHSpgtVjStC03b93Cmp7l4pTVaoE1jqZuWZwuxK7vZ6tBKBUj4j7ot2K3tmlwwGlRyEFtO3GQKQ0OwkDmyM3ahyG6HrSjXC0FNFmHDiK0U4RJTKgCKZa1HaYdYQNNuVwQJwkP7t7l+OSYMAhJ4wSdSDWBaVvRAnhty7ppvUNQxkPaQlPWmM6QBBFdXYNyuN4RqojpaEaeFrz1/Xe4d/9Dmr6V7KEwFNDSiN1fByGb1QqcjNfQkOU5fd+xWq29e0NzcnKC0w5je/k9V2E6i7JiWw2DmN5a1uua0/ka4wQwKQV93xKHESdHp9i2ZbXeYBX0ztCajvsP76MVhHEo4yVjsb31Ym8Rx3ZtR9uJxq3tevb2dqnrWubzdDRtTWcM63bFwaVLmL7n0YOHPHzwgJ2dHfq+Y+MFqmKPxgtfW5qmJYlT0jTn9OSEdlOhmo6uaxnpiKrtWC5XMsoLA3DgTANNTRwEOAvr9ZqT+Zw0EfdLkMRkQcB6ccK63BCEIc511NUK6xyr5Qnj8QRjNKv1mkBb0lFKq0Bry3J5QpEXUhFkWxFQGsvi5ATX90xGI/TlAw4fP0JZSxRoqnJDZQXQhYGmKHJG42Krmajqhs4X8uZJShyGZHnKxcsX+P2vf43/36/8Ck4H/IN/8B9y7fotPve5L/4FXiE+/tvTF/nz+qLzjq2fhHMMnr7N/3gAx9N01I8L2PxFbk8L1p8ekT79s9JdJ+y+sNIG7fVfq9VatLbLJUmSbB13YRiys7NHmmbM56d0bUcUR+RFThBorDX0vcM6Q5zE5EW+jS1wzpHlUuBsbO/jQ2qRbvQyNlO9ZzStvP7NRsrCnWPLRA0atNiH96ZJQl5Icvfg9CuKgp2dHR+VMGI8nmxZLIlqaL24vGOxkK6+9WpJux331ZRVTVlXkk3YtvS9wbgzJ54MkxXG/pQX9r7w3C2WqzUKRxgnOK1YzOfUTcuv/Mqv8OHde/RtR5YkfPLVV3juuVskaUBnGx7c/5Cjo2PCEJbLJVmeoxycnh6jlBTRxmFIVW4kODHJ0Dj6psYaS5qlOCxNWaMDTdM0OCduD5xocmazGbZtybKcVSmitzwr/OrfcHp6TJZkOAVtWxFEijRLsMb5+W5AW7dMZlOyRFx2dVPS1hVZnnJ8+Ii6abijNPP5gmefvcXlgysoZT0Lo7390luHg5C684JKpXHGsmnWYC3LxQpjBNg1dUeaJUynO5yczPl//1f/FetNhUaxu7NDkWU0VUXftiSxRP9X5YYgkKqXsioZjUdYBavVCmuNsDzrNU3foENNnmRopzCtIVQBxhlMYLEOyqahX65wiFssjALapqYNQ7R1OJ9l0jlDj6EzhsViTlHkklrsnFSRNHKyRaG4fxyOJJHjs96ssKZHK0djWjYbcZcFYchivpTspWJE2zY0tfQZmt6wqERLoBTUdSXN62Es9LM1khPjwHY9q81aglGjiLZtefjoEXmasbu3y3w+p+vbrdDS+aqEqiyJk5jRqECrABcqjlcnqDBgd2+Pk5OOo8PHtG1PFIY8c+MmfdexXi44qipOvLhzU64pNxu5MCWJB88SuWEdpJkwSuPRiPt378nIuJPmcmOsB5M98/mSJJOLm7FGwvGqmkApbNaRRlJlUIz3ef31T/LNN7/DP/v1X6fpOv7j/9X/mivXbv4FXyU+vtvT4Ynntz8puPKHffyPeuyP/Nmtjkka5dUPCJiefo4f7OWesVlukPh6wfSfRMz8MO/lf2zb+dd+XjT/9L9J+K9vrwgkumU8HtM07XbxNwRbDgvFw8PHgKKua7TWW9G4nMvSiWqMSAOSJPGBpWdhpIOmaTRqCYIn4yTOxwTg8Dlh4Tb+wBjDfH7qA1OHfKmEzDstg0AThvKzk8lYroNOOl6LIifPR1707wjDyOctphSXM+Jr1wAfA9TL+278tbquajZ1tc2dalrp79uUpZhqfsTbxwo03bt7jzQvCKOQMEk4Oj7kt3/7d7hz70P+8OtfJQg0r3/yk/zMp17nxjPXaeqKB3ce4TCs1kvGk7F0asURo6JAAVmWEkUxk0nBerkhjWK63hAGEYQRodaEQUgUR5T1hkp7nZGSC9loVOCsIwgU165dYbPecPHSZQ6jCGth6qPcl/Ml88Wc3d1dHI7VWhMlAXkhLfVt02P9BzEMA6q6RGkkkNF0lJuOPM9YrRe4ICTJYsaTMUkSCYBDBJmST6HI84IwillvNlR1LUWzdU1ZVmRFzqPDI/IsRYfibJnt7hLGCd/40lf49htviqsiHxEGoYzdmkZunJ2vhbEyY26DgLoq6doGpxWNX/UEYYDBsFovZITVtzgDrrNkcSbBk71Fh4FovHoRmRvboltFXVfY3tDlmR/P1Wyait45dBhQb9acHD4my3OSOKbzKxCQLrokTui6nrquJOS065ifHuOMpbeWRosoPYwiVssV9+7eZTGfc7pY8vwLL/Diiy/SNg1NU0sxZBzhrKOualSqCIOQuitpshoUdKajbmvSJGG1WVPVFVVTU9UVxXTEqlzRtA15nlE1FThHEAbEeYKzhrqrscahIxmvjfKMQGvqTYntDE1ZY5OYalPRt73/6qisI0CThDHrfklT1rhOQiaHmiCFot2UHK9W1LMZtu3Rgaata5q6put7VpsNQRBydHRElMgqsbeWqq5pPcOYxwmRDoTNCzRJkfH666/xR996k//6v/7/UHctz9y48Rd7kfhLsJ3PKRq2QXsCH617+vFsDhgcrX+e3/3hAI1D3LLbrIIf5Hf+R81KnWfQ3EdomrZ/89+fCcEls6l5Ql+0t7fLhQsXBRBsNh5IS7VXEGjiOPFJ4M02Ly0Igm1H3/Acwh6ZraNub2+XohhhTM9mIwXJA4ganJxDN9w2I8sfI9Exzbh48TJ933P33h3KUiQQSSrGpSECY7PZsCk31FXtoxmCrVtOnHxjfz+W3rgsyyiylP3dHS7s7pFliYC0cGC15BzpOkPVVGyqDeu1xCTUbc38dM7p6Zwv/f6/+JEe1Y8VaOp7+NRnfganFd/7zrf4jd/+Hf6f/+gfkWQ5X/z5z/FzP/d5nnv2JnEQcHRyn/nJKcb0foUGk8mIxWKOVprNekWgpUF5tVpSbjas1yvSSCL8a1fRdzLiydIMHQZsyjW9685GQl7IG0cRm3LNarNivV6T5jld31HXLU3bEkcx1jhGozGTmYjHsyLFaUkVz4KQaDehblrpn9Ma6yx5lpIXGWDQCrLsgN0Lu14noNnf25NLWtdK8rRS9P7D3VmD6TrqtqXte6ztadoWB0ySlKbv0V2H6wyXLl/iwsFVvvu9t/hv/vF/z61nn2UynhBHMXEYYdqWHiVON9NL5UyS0LUdm3VJFAeEYYxxFq0MTduyPlkx25tKOGbfUVqDaXviIGY2nRDogOVyhVKGNE7oTEPft/RGVkJt09DXDZtTy+5kCtbi2oY4iYmjkPVqxWq5wlQVCyO2Vpwkx260MG7rTclisSDPc/b39lidnkr9hA5wWmOd1KGslyu+88abfPDBHd65fY+93T2+8PnP0zaSpNt3HdbIydx3PRtT4qylrhuWS6kZ6ZqG1rQQaTalNJTXtsNaQ9nVTPd3Wa2FTWv8yjAJE3SoqKqG5WpDoAKoNToMKNcVddmAdbSNRDm0dcvp8SlVWVHXFXEUEQUxXdvLBaRX9BhsL3ldcRwTFuLAasoK23Y8/PA+QRD6mAFHW9dopZmMx1RVSdPUGOdXsJ24GcMoIvYXvtKv5KzWmMUpo+mM55+7hXGW/8f/7b+kmI7+gq8SH+/t6eDDARQM7qVBIPyvUx3ygwONcyOzHxo4Pa0A/7N+9qOf9Qd9iB8WPP3ExptbzdZHAabzr/m8dsxt09KN6YnjkMlkzKVLl7h+/TrGGJbLpY8BELlAkiRMp1PiOEGYpoq2laaB09NTVqsVdV17957cD43p2GxEbD10Dw7gTj5jiiAQh5pzA3MY+O9ljKa1ohgVXL58EWsdZbUmS1OSJOH/z95/xdq25el92G+EmVba+eR7bq5cXV3N7uogiYJpmSIIybYYYEAwn2QTDoDhR/rBfrQpioBhG2CQDJACKUuMlk2RMpsEKXWs7uru6gq3boUbTj47rjzzCH4Yc669z65zY90Kp+uOi3N3WmGuGcb8xvf//t83Gge3cq0jrLWsVksOHx+y6rY9SWKSJKWqWhbLOersrCvfDTowFTEZDXl8XyPxRL1TfJqRDcK/QTYgzQYBoCUxw739cJ2oIHlZFwV/+//1tz/SY/pMgaZPfOFnOT6b8tU/+H3+4T/+x3z5K1/l6rV9/o//p/8Dt2/fIl+vMG1NVTSUVUnT1uTrNdY4tnd2Qqv2aslkMgl+ILYLzPWC2XRKkVes/Zo4TnDG0tRNKP20Bh1HeOGxxtKahkhryqZhNpuyvbWNUpqz0zMGgwGPHwU7/DgJnQXGhjKHd5633n4rdGJlMV448vWKOImZTCYYY3Hes1wseP7552nqhun0FCE9WZYQRxHOOU7PTrvHG6yHumkQUjEcDDHWka9zZstVF0YMkY4QSpCkKXGScXh0HFayBJH33t4BZ6cz/tk/+W/58m//Dn/qT/9ppI7J8xIbGyIZfDOqYo3r6NqyaxGtW0PRVOwNBoDCtg1bkwnHJydss0usk86KXzHZ2mZvb5+9vf1QO3/4kGKdEycpGQ4KhxIQiQSjIwY7u5wdHqGqYOsvbBvq63VF5D17kzECQV7XZFIHk0ilQ04egiwbsr93gPWO8WjEcrnENZZIKryQtLbFC8l4ewulFJ/+zKeYzhbUbYuXAqTospzqIHwcDCnKclN/B9npCSzz5YLWBBCZZgnGWtrWkqQxDw8fsb29TdNUKCVIh0ko99U11bqiqmucteG8sx7RlV+FEOx0IFsJiWkbZtPpE6Gdy9UqMGJVxWgwJIrO9WfOLrtYoeCoLqWisYb1esViscCL0IVnrMV1E3prTHC+x2B6vx/AGcNqtqBcF2zv7FC1JU4GbsBLya2b13nuuef4B3/vH/2YZoc/GuOiSeVl08SPgml6/3oov/n/B+WLgs7lw+muvP+g7/eTPp4ERufg49KjLpTnejbIdfYrPSuTJKHc1YMVa8M+1lp1zTfbTCZbRFFMVYV5a71eB21mHr6Gx2ugN5xsKctio6PqtUk9OH/CZFL0+YaK3tAS4RHCk6RJ54kUSqxSSeI4YjBISZKM4DeoWK2WaK0YDDP29/c52D9gtV7x6NFDyqJECBiNR4ExywvqfE1dFlSrNcI7Eh384wJoGjAaDRmNJ0wmY0ZbE0YdWxWnwYMqi37KA3t/76tf5a/8x/8Jr33rW3zus5/mL/7F/4jPf/7zPP/cLXCO1WLFerVEdczSeLTFaDAiz9c0TUWWDRiPmy77qGG5WLJe5bRVTVPWxElKHMVkaUpTtwipiZOEtm2wrmMoXKCsh9mgiy9xTMYjxuMhb7/9Nuv1migKAZf7+xlCaeqyRuoIEMEsUUsiLdFKECuF1BodxSAMo+GQ3U4kZ6IoiPfahiwakKUpx2fH1HXLzZvPUbcNp2dTyqomTjO8VHjhww0fgXU+lP6MDcaECObLnPl8jpaKG6/e5NatWzw6fMzf+Bt/k//y7/9Ddra2eON7d0L3RNsyGQ5QsAlmHQ47R9koZjCZYJRkua45fusOdRO6vbZ3D2idIElGCCKqqmFv9wrGNPz6b36FN+89JC8qROe4ffPGNf6tX/5j3Lh6jXKxRLWW4XjAZDiiPZ6ToLh68waFrZlXa6q2pa5qhumAIs8Zb4fS4jIvaK1lMtoiG45Y5zmL5RqsJdIZUWzY2jpgPNmiaGuW+ZqqqmldBb7l5774M9x9eIh1jqYNZTepNFkWIYUCNHhNXRkKW1N3IChJ4o0bbdO2RDpinefMZwv29nbwwmHMcXAjF44o0ptIibpqsNYFM0xTBU1WFJMlWYjqaVqOjw+ZjEakWcZ8frbZLmvdRpRZG4OsKlQTym7CB3aibgweS5YNmM6mKCVZzNdY55E6CtmFWlOVZbC+MC21bfEQ2D5j0FIhEJjG0TrBo8MThAKdaoqqYpWvsd5z/eYNfvFLX+Bf/6sv/ziniWdyPGEI+ZTy3EUDyo/qvd55XPAw6m6CH6xE94MI1c/b89/PuOh59JNcqutNLp+6b/zFx4SvWutNq39f2jo9PUVKyWKx4PT0tJsroo3552QyoWnaMC90r9V3o5VlyWoV0iF6wfZF8H3R3LQPSu4NRnvrEaU0QnSpCkAcB/CV52vWqwVN2/D48cOgGc4yFoutzlMqAJkgaG9QSnDz5g1+5mc+z6uvvsrx8TFf/epXOTk5ZXt7m8985jMURc6D+w+499bb2KZmlKZoIVDdtWGbmrxtyJcLjni0ieySQiKkRChJNhyE6KyPeDxToOl/9b/535IlKb/8y7/EZz/zaax13Llzh7apuXf3Dq998+toLXnxxReIIs39e/dYL5dcu3qFsi544423uHf/Hs/dfp7WNFRlxWQ4wVtPGifcunGTRw8fce/OPdbrHB3FjMbj4COxXpKmCWkaykOPHx0iVBBYx9EDpJacnc0YDEfs716hNYaiLJktVhwdHZOmGev1kr29vcCadCJl5z29O3XdNIzGIxbzkBYfSYXwnpdefAlvPX/rb/0thpMR/+Yf/7c4PDqhKEuEEOzu7aOiiMeHh6xXa+Ik6ez3Q53bGkfb1mgdIQgX0nB7h8Viwd7eHr/167/F7/727zBOU65fvc709DSUGeOYYrmgWK8xTY1zHkNw5lYIGgFaSYxzWNetnjz85m98GeE9v/vl3wcPjfd8/pOfwluLijX/8V/+P/PHvvTHaMsKWxZ8+Td+k7/zd/8LvG3YHw1Yn83Z3d7m+rVr5GXJummZzqY0viUZDYiSlKYoGUQxTVFycnjEfLnkdLZiUYRSopAS4zxNf4FHCmMcnnBBBf1U123nPdY73v7udzk8OmaQJnzr618POYCtwbYhF6/qGC/nQhimJyyL54tFmA6tY7FYhdftTP2KomC8NWI2mwEW5wxRHBHHOmyP80RROEZ1XYcA3MaRr3JGo0Ew8WxqYERV5qxWa0bjUbCUaE1n7umIowQnBM5Y6ralqWqkEBjnMK0jivJQmhSBQVuu1iAk1jlaGxLEnQir2rPpafBfiaIgFDcGKRSDdMRgMGS+mAc2yzQYZ7He4pzl7ptvMBlPfmzzw7M4Lt/oL+eqvVP33I9o635E7/Ph3/Uio/VuQO0nAVBd3Nbe3w6eZNYulmPjOEF3TS19abZtW+bzOYeHhzjnNpEoRVGwWq26spnpzE/DObNarbqOcdt1tkUbb6ZwvgUmS8reTNV1+YQBLPdaLCnpDCVDd3ZvAVBVZRdeHz5ZL1Bfr5c0TcViMd8Au9VqRZxotrbG7O3tcuVKMFtOkhhrW9q2RkpI0pg4iYJko20YpCmJ0kRShGguKcAHX7zGtLQm5MBa11nWIJBaYf2TZq0fxXimQNNnP/05Dvb3GWQpaZKSxgmutTRlzdZ4wi9+6UvojYOxY/ypIXGsGWQpJ6fH7O7t8fzzt0NGmFbEOkKhaKoapTSDOKXtdEVt25sPRgxGA67Ja4xGw84TyPP40SOE9Ozt7ZImCXXbsr17JbispgPKMuTaKaXY398DPFtb400nxHCY4bwNZRIfHLMHHYtTFkUw9jKWSCmOj49o25aXXnmVn/+FXwABi/WaphP0npycYDudQ0gcF8Rd5k5dBm2WMRbSlP39A/AO4SXr1Yr5dMoXfuZz/Cd/9a+ye+UKcRyHkGAh0F1umm0bhAqBu6ZtQ8SDVpigPg/tqAIipZFOgLNo4bFOgg45blf2r/CvfvVX+Rt//a/xt//mf8bu9h6/8gs/z+H3vsPp3fs08zWf/dyniITj937r9zFe4CNJ04ZummD4AJmC8TBBCIuWGu8de3t7fOpTn+KVT3yS7f0DnJBB24WgtgG8Pnr0gK3JDsPhkKIoaaqWQTJiPpvRmIoHD+6zzmccPzzh5OEh3/3ma1y7eo08r6jqFqE0zofyldQilGq9QSqBqRu8daRxgo4UrekM2XwAJEdHR4wnI6xtMDaYutnW4hxYY2mkJY4STGuxOMqixLYhV7EWNaa2NFVDXuYY56m7smjTWqwLk5qtG+qqJlKKYZYSJylt3eLKYHEQ6YhlVYXOlfEWRdfd6LxHaY0xNUVZhFUaobwplMY5y7qsEV6Q6AFxLIiFQuqYslwjteyYW4e3jmK1/rHMDc/SuNxu/k439MtC8B+0e+79jydURR+sXPZ9T3i3Z78DkPnx45t3H32H2wd70pPP8Oe/853b52VmUalQMek753rnbuccZVkCdL57LXVdM51ONwz2YjHvOtOSTZZgWPSnnV9fHx58kUkKWaD9PUqIc42d1nLzu16P1Zf7rDWdeDsmy1LqOkFK1QnLa1YrH/LhfEh6GI6GCHyXX5cH3Vak8c6yXq84PHzcdZzXCOGR0iO8QwtPrDVxFKG6EqJ1lsZENF1ki3Gha9gLaEyLqdsPd4zfZTxToOnmjVtIEbqezk6meOuZjAeMhgO8s1y9cpU4jpjPpjhniUfDDm3D1atXQQhu3LhOY1pUpMFBXQYRnGkseVkQJTFXr13FO5jOZjQm3FzTLENKifOhRTsdDEJQaZSwLkra1nDz5i1WqzXT6YwsHXQnju9ozdAZNxoNQ2K3EsQ6mE1WdR0o2SgYiQ3Hw1BWXK2ZjEYoKVnMF7z66ieQSjGdzjAutHsKJSmqgrquGQwGXahmg4ljpFBBc2NDa3/bNKxXIS/INoHCnc1mm9WGMw0q1sTdaqIp8w6seMoip+1WKulwSDocUna5ecGB9j5HDx9Sr9c0ZUVblCSDAQZH1Vqu7u1x/43vcbxc8Ma//Ffspynf/PSnufu9N/jyV79KqxS/+LM/ywvP3eDG/j5Xn38esTWhwqMFCGtQQvPmt77NN37nt2nbmrqoiSPN7du3UUJw9PgR6/UKFUdMtrfY3d/DoLl2sM3uJHRJNnXDIBZsD7fZHR9w9Pgx48mQ6nOf5dqNKzw8fMzx8TFJGlOuV8G13IWOSaUUZdOwnq/xXRdl09TsbW+zmC9woxFCQmvCPpdadp0oCVIoVkUdmCprwmRF7+EVIzKNtWFFl6ZD9EjjgOPTKWXZkGYhfDqOA7jKixKkJM2GVHVLvl5hjSGJok23G87RFBVxFLM1GPD46JCmbdFK4/DkZYk1ljRLscZijSVJNINsSGsNeVEGAznvsa1htlhQqIK2aYiTTivgPd4E89QkjpnO5z++CeKPyHinWI6L7NMPnzl50gbggwzh6Wjn9zIqeBczzg9QnvtB9sUHeu47Ybz30G+d2wiErEvv5ZOaJh/28kVAEgTeNXmxpuxSEpRSoXM8CqW74WgY0hqyAc4G8fje3h5aKe7dv0++XiMuZPfVVU3TBlfuXtukdXAI70OCkzhGqqB1xfuNrrFnocLnAQQMhkGrlGYJ2SDtGPhgBdBLEJSSRJFE0oM0h2lqHj28j9aK7a0Je3v7/Mqv/Ao7uztBKL5YBCBYVUzGE4xUiLoh5Ks4wOGd7WQULXXb0nYu4kBXqqPzx/spD+ydns3Ae2KtiLVmejalKQv29/YQ3vHo4WOkAGsapBJ4XFdDlSRpTGuCA2oyyILrcduv8FvKvGTYefMEus+yLkuMtQitibyjrhtWy2VnJKawVrDOS/Iix7QtV62jqhvOzqZcuRIziBJME/LwoihmNp+xvT0JWhXTdheAwvlgYU8jiJOQzCyEIh1kDEdjdEelWuu4e/cuUmriLOmCiQON29Qh4wcCOAv+FII0zUJenoMkjsGDbS29WLPvxDg8PEQIwf7eXmCabGCU0kFG3TYsVmuk1iTZACMEUZJ2IZHB58h3OUHFKqdarajzgrRYU/mGw5NTTPUcg1HCz7/8PMXBDq5Y8ta3v8mjt+9Sn51w+8UXiF3Ddqz5lZ/9Asl4wODqPp/+4ueZnpyhhWT7uRdozub8we/8LsvjE1aLJbPplMlkzP179zg6OWZZ1UglmR8f8/Dtt1GRYqtzqVXCkykRxIlCsJyeUq6X7G6NUEmMxrEzzFglitl8SmMtUZzipaKpDUIrhHIgg+Gb8AIkJIOYrErZ2t1iMV1gWo+IFLQOYw2DJKOsKpwHFQUn3LatQ6lORyilqZsQM6BURJwmIUbAWVoBxJrWBvDSVgEEBWldEIOa1oCHSEYIJ6jqCm+DqL7M12FCNS33H9zn4OCA4XiE7cpAweVXBCNW69FSB7fdxlKbFidASB3sJawNtgt1CVqCVrTOBBd6CbFIaH8IdPhPw3ga+3QRPD3NEfxHtGUf+BnhnhpCVN/z2f0D/CUd0zuAmcvg5LJp5Ls99glfpL489n4/Xvdw3ynVRc8Q8STIPY/wEJt/4T0DGx/e/WJG3zlgklJ1i1SL9QbT+Sk1psYZj8cSRQqtFWmWcO3aFV586SWSON34E73yyivs7+9z4+690ATlLE3ThvLdcklVV5RFyWw+33RTGxNiTvp90wvBff/ZpNgYtm8+b5f5JqUImZw+sElSBBfwIF4Pwe1KSCRBElE3NaY1VEXJ/OyMe3fuMUgybt68ydZowttvv8W3vvkaWE86HLM3mkDb4MoS6QJoMtZuGLJ4kBG3hrwsA/u26SQI+qfoPXILP8x4pkBT27Qknbo/jqKQB2cdbVUxGo1YzKfUTUWk1QaVqkiB9KzWK/b294NJovOs12uEUOxsJaxWa/LVmq3xhOPjUxpjQqdR24YW9qoMJT2pMT68tLCevKjQWnaW8Yqz6YyyLBBKUVYVUZyE9n/r0HGIKlkXBb1rad3UlGVBawxCCmzjUWWJD8I7WgABAABJREFUVpr5YolWkuPqNLhpSxnEfOucrZ1dnIeyqYlc0CvlRRlWIlrjOgOwwKwqnLEbozBc1yoqFFJK1us1UkquXLnSrSgEzjuSLCUbDkJgZFkyGE2Y7OwQZRmnZzOs86TJgN3dXQ72al59+SVefekltLWYsmB+dsLO7gSnPY8PH3Lt1i3y+Yp/+d/8E954/XVuXr3Kjb19Hr55i1fv3+bGzZu4suC1r36V1WzOG2+9STYZ8+nPfZrjoxOG4wnPPf8i1196hb3dKwxR3LpyndV8zmQ85tXnnicvC+IkQSjBg0ePuPfgLtY7muUCKz061nz+859jOBjz+utv8PD0BGtaqqogz5f8+m/+a5brJaPRiPH2FnGWUpYFq3WBAeIsDQylDUDB4YjSmIcPHzE/W5FmKYvlnLpuEVLSNG3XJVd2K7KgPyidx7pQgsMJbFvirMdb0IkjbyoaZxiNhgyGGa1rsRDyAaua0XiM0jHWedbrfKNp8sYinEf6oDNrTBv0DMWaZbHk+OyYrb1tEuFoTAhtTuMUiaDIV5SrNRLJOi9wwuOkwBDiCaRWodPUeQyO2jYh10mCseGcaoWH6KOfpP6oj6eFx1prnyjH9df2jw449V5J/sK/9/vU9wmYLoxNweviW73PF/hQnXoevHj/qMnjnwSzUpxjLtmBIwHWu83fhRTgglQEPF4Efz/vzsX1IVct6IVUpJBKYisTAmqVIEoUQkLZlBhrUFqBCB5vu/u7XLt+NSQUNBVSCrZ3tnnp5ZfZ29/feCydnJxwcnLCarXaRLN84xvfYD6fI4xAKolC4QXUbUNj2k5Q3XXLCbnppZSiL9t5qrJGK4XNLGVe4D2kccYoG3WsWIQUAtyF8F0TGPbJeIsszrj75ttor5hkY1555RV2hhPmR1PKqiJNYmKliBVEOLwz5Os1J6dnVE3DYDhiMBzhEJydnXF0dETbNFjTdl3eP5wF3DMFmnQUTCajKMJZhzWGLE45PjqhqWviOCJLQxCv6XLK+pMbIRFCIbvOObzAGMt8Me9oPsPpdMbjw8fMZnO2trbDit8aVkXBbLFECo01FuE9Sks8IexXiLBty8USa0Opa7lcE0VZKCWenbG1tYWUkpPjU9Z5YKucc9RtTRTHDIYjEFDULaPhmLKogk9QEZy5t8ZjjLW8decuz78oiOII4ywISV1X1E0IlM3XOVVZMhgOiXTEbL6gaRqyJGW9KijLiqoKJcmDq1eYnU2xzjEYDJBCsM7XrFYLdvd2yYoiXGTrnLyoSLJDssEQLyRRlDA9m1GuKuJI8y//+b/gv5fw3I3rYFve/O7rvPqJV/CRQCrBcrnCt4aibVlZw6xpGHvP+MY1PnXtKrv7eyRJwvp0ipcRv3TreXQa8av//Fe5fvUqi2XJ7//BN1h3+jMl4MXbz/OJl1/h2sE+3jnyfE1rWqQSCGu4trtP4xoeHx0xPznlu996HWFa9vf3ee3r36KpPGXVIrFcvXqFva1tfuYzn+HzP/MzCCl4+84dXnv9O1TOMRkOEUoym85Y5znIoPlaTmfcu3efew9mrE6P0UpTN23XvhsTaY1L0+Aiv85RWlHkOUVeBJv/1tE2hsl4i/29A3Qc42QAxNJYTF5SL9dE4yF1Xga/sOUq6I0QGOuJk5iGBiUCQC7KvNOetcyrHKE8+/vbJLNj5tWKo/tnZNkQ6UAhwMDybEYaRWQq4vBsRustUZLQekvrPUmWYqMG10UHTfM1cRqRjTJEt9rUUlG75sc9TTwz453YkV5H0v/8tO9/dENc+vo+H06vLf0gwz/pX9QLpd+FMdo889LPT+s+fPKdnvjfBx4X2aUeHl4sm56/t++Yl04zJEK+5jnDFmKvemYGBEpqpFABnFiP7eKhpAxApo91ajpn8NVyzeHhIdPpFCEl0+ms82QKb9L7NNV1TZZlXeyIecf9473H4ZBehq5gQPYdncJvzlFrLbLbzz5UHs8ZN+vw0uGl3LCGsjOKxkMap6RJRtO0rBYL7rz1NpLQGXh8eMyjR48wpmV3Z5tRFhPrUG6z1lA2NdaBFxKdJCRphk5SBsMxxoRmnbosWEzPWHTh8R/leKZAUxwnLBcLykjB1phYj5FdZ0FZ5CTRmCJfU6zX7O/vkaQJs8Wc1luywYAHDx6SDgZB3d+2m3bE0Ang0VHMy698giwb0BrLfLmkKEt8kIfQVKEkk8YxxraByVKSKFIo67BtS5YNaVqLAGbzRSea1rStwTnLYDhgZ2cfj6eoCkSliZO0c+9eY51gPl0ynU65ce0aQuoQ+tq0mNayu7NPVdSs85KqqRmORkilcF5iHGSDMVIlob7sBNlgxGAoaeuGxlh29w+oqoqqrCnKhtFkm6LMWaxytFToJGX32gAdR6yrkjRJGG/vMhha6qqlLS1xHDLvhIHVbImUgtFgxHic4XGoRPP8J15i1VYcH81QWnM6XyO8wEQZz3/uZ0iShJXzuI7jXlcNKYJ0a4ud/WtsjcZ8983v8Ut/+t9jezRmMZ2zu5jhnMU2LavZnDSJOTw+5stf/m1wjuEgA+HDiizSXL9xndsv3EYIiVQQac3dO/e4f+8B3iuG2ZD7d++TaEG9XvM7//1vcOu5W8jWUVc1b7zxBnv7e/z5f/dPA4Kv/N5XEDLnpVsvsLe/h9IRRVlwduN5vviLXyJOYr7zxludnYLYdLQ8eHCfnf09iqJAe8HBaIxNBxRFRSUqVJays7XFMInRSlO2DdJ7EilRQmCEQLcWbSxZlm66Jsu6pWhqdgdXQntyHNO2hlZ7MA2alO1RaBP2g4gmlpTC4hPJ9pVdFmczpmdzJnHGrWtX+eTLr/D8jdv8mnGczaZY4amtpSVMivUqaCGyQcYqL4mVpFisaNsG1ZmEnpyc/PgmiD8i42nxGh/167/3uCgG78HBezzjiU39wSwHflij587olVof8O16sPQ0Wwh48ngFn6NgCNl3wIUH2fPHCEJjDuesovcBZGkdE8CnRKkIKTV4iXNgTDDXXa8KlssVp6en3L17j/l8ydtv39kwlUop6s46xHvPaBTMZ6uqekfQdFk7J3zPkp13KV70cHLe4QKlFkqKbRvKlz3Qu9QJqjp9lNKaSZLiPTx8+JCiCOkDZVFy+PiQk5MTrl7ZZzxMiLRnOEiJ4hjnPEiFLxREmhESpTWj7S2sMTR1iVSC9eqjB0zwjIEmaw37+3toJVHCU5Y5RbkijSP2dndQ0rNeLXDOduVkgXMheHe1Ktna3eHsbI7WQbfhnafRFmctWmlWq5zVMmdvfw/jHat1jnE26D2kIIqTrrOI4LGU6K7lsiIRivHWDk3Tdq3qDoEKJmRtTStMlxy/oGnrQOV2ac2ZEwyERAoNPmzzc8/dRiAxrqUoVqzXJbdu3KIxntPpnMn2FtZJFougNXI48rIOF6fo/ukYJWOiKEaKEDXSWo9HIqOUxoT6bxQPSX1wfFZxTGNqvPUInVC1nvV6iRKayXiCcILFbEFTW7YmO0RdWzreBV0TgsVigaNha2ubUW3IsiGj4QhnLRZJVeZkccbe1jbVOmc6n6GiiFoqqtowW055NJsTb+8wm81YTuekWiGHA3YmY5JI893XXidKMyajIc9LuoulpixzEpkSJRGz9Zqz174ZDEq73LednR3iOMY7ifAh3wgv2Nra5j/8D//nHB8d8fjBY2xrOH78mN/5zS/zm//613nxxReYbG2H55YN3/v6t3j86DHj8Rbbe9t868tf5c233+bVVz/B7vY2Z9MpN3b3+Pk/8Sf56h/+Acenp6hritFowK1bN8nznD/4gz9gPofxaMxbd+7x+197jTSOeOn2c4huRbe9s01rGtI0YStNwFkmOzuUrUEAUSSpV0vOjg/RccTWzg5gWK9mOB88oaJIYU3D9RtXKaqSqigo8jXWtHhnKMuCWdPy5lvf43vf+jbT6RQnQHdhzdYajLfEkQLnaauS3e0JSktaa0BplJREUcYgHbBcf/R5T38Ux0+6t1AYH6zMdlFAHp77/stfFyHNBx3vFIb7wxiXAcbTAO6TxzWIv50P5UcXkFKnKfWdSaUFgui6LBtAdbEiI9J0gFIRWidAVzGRGlB4LxBC4Vwwqj09Pd2wTD1g6QGOEKIzxpQbL6be4uJpn+2dPw8b7ZDvFqlN2+DxVE1F1dQIIYg2zvWdQa5zOOtQUlOUFVtty9bWFgKoyorT6SmDbMjLn3iVdVlQNyHpwliDsxalRCh0emhtxTIvOJvPSdIBURQyUU3bUFclVb7GNvVGd/ZRjmcMNLWUpUVJSRJHZEkwnxykKdkgDYnOoyHWWtZFzuHJMcZ6rt24yWy+IM8L1nnJtWtXg1i6KLE+1FgjnZAmGcfNjOPTs+C6bA1RkmBcCwiE0AjniXVEXTXUbYvWktY4WlMCgiIPQjdjLQciIOq8rLDO09Q1CDC2JU5idBwjncMax2qZ07QtcZLgrWe5WGGMwzYGiWA0HOO8wLQOJRNWqwqpFIPRiKZtqcqC8Xi8Me40ztEWLetVvWkxdc4jlAk5bU2L9Z6DvX1a01LUoatKoXAiYrkONvxJFBPrGK8UVWPx1mMQeOOo6hWDLEPgODp+zO7+Fq9+6iVW1QK8IC9yYi0xVU1pHJHWJB6s0Ix1gqwtdl2TeI0ips5rnGlJdRRiTqSgEp7dnW3y9ZK1qakKRxZFJPtb4AU+S7jy8os0VdCHResli+WCZROaAVznQ+QRWKUoF0u8tUHL4xRECpnEPD45ZjQYhkBo79nb3eNX/o0/zr/zP0xZLpZdoOSAx48eoZVmcCVFWoH3gv3xLvnZim/91u/xva98HSchThI+8+nPUh5PefveHabzKcPxEKng7MF90jRheXzEdDrlc/+jP8n/+n//v+O7X/0ab775FlevXOUb3/g6xyfH3Lx5k4f37yKkYH97m6OTY6q2ZTSesLO/T2Mt67JifziibhrK2YKqKlFNjZYC2RqMN9R4RllGEiUsREWT5yRKIYdDqnXIbbLWkMYpo/0trIDG2S6qQeBdyJ6LpabI10jVdX12ItbQpqzZ3t7m8PRj0PTTOX7SAeD5+LBberlz8bJB5EU2RsrACvXhuwIVWF/RC8qDrumcudIIQre1c9A0ocFDqcA4CSGxFpSKSJMBcZSiVYzWCUkyIE2GeBzWuk3Y7uVYnj7It//dReD0NBD/bsBe6yBKRxC82mwo2TnvQiRUJ4vzhJKkc703ng1hu21LFMeb0Pv5YkmaDnj+hRcoqpDbuZhPqZoKrTzaKGSXV2espagajHVIHVI8dBThrME0TbDJ8Rb7QzgnnynQpFTwhPDOI4RBqYThKENHinW+oqqLkJcjRZfebvEI5osleVmBFCRJxmpVIIWkrluaOmcyGtM0hvX6hKqsGA5HRBG4qsQYi/WdAaKpwQpGXQkuFrpbyQdDsLpqg34oG2GLgtaENk4VxZ2mSjMaDbHeYF3wTqrbpjuhQ9dT07oQcpjE2LrEuuAQ7Rycns66NGvLeGcb6xxV1VKbButgnZd4l2OMIYpi0jhGqiicZKa/AVoQQRwfSck6D8GydV0TJx5tE7RIQqdZEszVvPMY47G2oalqTGPY3t6mynOEDi7pcRIzHA3RkQrZZK5FDwdUVUFTGXLnSaPgGZJIiL3HrHPqVbgBG2soihwlBelI0rYNRjmKKmc8GVJjKH1LkVeksSZNYlrreDA7ZTIaIb3ASo8YDoiloMlzWtsGj5IuAdzHoWxpCZYAtqpwWuK04tHRIdevRbQCTGtQRQ7edQC5JG9KsqqkMpZYhLiB5156mRdfeJHbz91iECV88Ys/x+OHj4izFC8E0+mUfLUmlopBFDNOUu7ceYt8NuOll19gbzRiPZ3y8Dvf4Wu//mt877Vv8fjwmD8sf5+T4yMGacw40igTWKWTu/eYzqeUdU2cZBgvab1g52CH3b1dhoOMbG+XJAkxCuvVmqJYU9UlVV2Tz5bsXTmgmi1xVY3rzOLi8RDTBvq+xVI3ObWziDhCxKFL1FuBTIIPV9s0xGmY7KomxMBYPFKG8N+Pxx+F8ePQT/0oRsdm+U2B7gONiyCkBxQXOx0vap2CgaQAH7rNJBKEQMpes+bBy00DTqhkQJa2tI3l+OiU1arAGof3grax1FUNXpIk2aZ8J4QkidMQrWUNWgf9lOsWPRdHH39y0Sj1MlP2XpYX/fe6s2EJbD1PgEcl1eZv3vtzmVr30YMRpQUZGmSUlJyezTDWsLW7zYsvvchyueBbr61YlisQwZPJWIuOFHTZoVVT45qGqG2J44SOi0IohTOG1n3089EzBZriRId2bk/QI0lBkiaMxsH7aDBISOKOaWlbBlGMR2K94ODgCqdnU7RW1E3DMBuSpUNMGzLCTB18MfCSsmo2Dt3BFaI7uZwgkiEOpW1aqjKnqjRRFOG9p6lDnThNhzjnOTubdaxOQ5YkVFXJ1tYE6yyyMworqwpPYCasdeRFw9aWwBlHvl6jhEQJRVEuKYuKydYW0+kMnSZ476lMS1nXwX+naxWVInTq1W1gqYJhGTSNxedl6K6zjiiKWNYrhJIdGxVcr4UXFHmBVIo4CasIY4KJkHMWpSVOC/KmZOAy2rogHiQIBcdHx3hruX79KrPZlJaWeKDwxmJdjXcN3oMxEViLogEXnKtHWXBAb4slToIXNU25oMwj2qbG49BxDFqxKkvqqkJ6UG2DQpKv11R1HdzApaRuPR6BjCI8YKUKeYCpp61rnPWk21tUeCo8q6amtS1VU7JsgxA/jWO8d+hS4adnrJZL4jjBIyjLhtcePmBvZ4fJeMx6nbNezLl54zpJmlJFkvG1fdQoYd8fcPP6NZxr2b92hRdeeoHXv/aHrI6PObWerzUtR48fc+vWberZmqhuoW05fesOzgWTVA/sZQOO1jn4mrpuOZwtOXz8mHSQEXXXQpJlYbK0hkRJBnGM8prZdIUYThg4yc6VfYSS1MYQpwlewGK1pm4bVnmBkWB8g2mCniGSisa12MbSOoOtPTqJiNIYnSYYH/y6iqb+Mc4QH4+Pxw9v9KzM07r1Lv4sZWCWvBcoabE2sD90uqVQUgv/QiBuMLCMo6RjiByz2YJvfOObVFVF07RoHRpL6rohzwva1lAUJdPpjOVyxWq1xlq/2UbwG9DUM2BCiM7pm01Abw8CLwKmy9EqlwFh/xhjzca2prda8NClLViEEZ1ZLpvX852sK4o1UstwL/aeKIlpTcN8NWc6myKUYGd/l/H2hLJc0rYVTSuI4hjRNZ2opoVGYFoTMjBFZ3QhQEGn6/0pdwQ3pnP9NAatJNYZTs7OqOqKJNYY16JkODGkgDhOg8ZH68CYeMiSAc4LBqMxo8GQQTagrmqkh/39fRyCdVFQNzXD4QihVRd1IvDG460niiMmaoKQbmMrr7WkbZLQWp6kWBvYDSUVg+GQQZYxX8xIB9nG4j6OgyePR5CmKQ4oi4ooirGdUF0rTZIGhqYqq9AmniYYEwy9jAtaInzII4u0RgBaRxubAudCllHvJOts8JmqKkNZ1aRZSpIEM8zVao20grasg8heSKxw1G2DF8HXxwnH6WLKMl+RZhHr5YLxIGG1hpOTFbGCqQqt9vP1nJ2dbVCOqqk7PyeH8w2+anFVg9YRRVuTDTKUEOSrJRZopEMLS7mc0RiDQ+KVpmna0I1nHVpIlKoRCNZVRVEUAThqjROSyliEc5jWUjUlgxSE8BuPKa8cy7JAJAk1jsZZKu/x3mJwoIJVvxYSIyynbUFE8Lk6nc/4ztERSkum8xll43HWsDvK2BqNmQyH3L51k2K1YhBHGBmAWdYaTmczTk9PsE3NQGu0dcQOnrtyBdFY7GpNXRdUixXr1RyEYnt7i1RpXGlQqeL5mzf41Kc/g9cRPs1oTYPA0RjDo8ePOXzwEO3g1sE+W1tb7I+3WZ/O+OIXvsB0teb+44esqhIRx6Aky2LNlesHRMKihMA3IWoIAVoHwCwI57/tOnh0ErIUccGSITYt8HEH3cfjj+boS1mXtUBPsjRBx6PUuW4pMDqd5+cFp4Pw2PA1gBuFUhF5XjCfzzdWE95F5HmYy5umxTmoqpr5fBE646oarVWwmRHn9ggXmbBzQMX36Z0uj4tlxv5f/5xQdrRhEa4CWJJKhftFF5ZurMO6Nvg2bWJaeruMMIdEcUTdLfjTJMELzzpfc+/hfba2ttjamfDiyy/iXcPD+3coqxodx6g4wXU2CJ6gDfbCBsNnKZCekAjhHOanPUZFiM5TwiuiOELgmS8W1FVNmsY4QrhtpDWDLAOWeCTZYMTDR4+p6prRaEieF1jnaZsWU4fymHcuMDUdE2VtoA6tD54pSmu8ILizWsd4PGQ8GQaPpDzUVOM47roGII5D8nSWDcg6N3EVSQ4O9qmrqus46uMqwgnXtk0IC25aJD7QrEIyGY8QQrFcLEjTmNGt61g8i+UKbyHVCTHBhkDJYDSWpSlJkiB6wzIlkEm0ySMbDQehdV54pAorHq0VzrZEImacZSBFEHorEFpiCflkzoeVxXCUobRikGWkaRS8s+QIZ2oOHz7m1U+9wmw5Q3R0rYg9SiqsCUxXVVXEQpEkEaYpqdqWyXDIYDBktlzgJNy6eYPp6RmpjollhOkEj8IJdrd3Wc4XVEUQwCdJhlKaoiyxzqG6rkWlFFJpTBNadJ11eG9IhkOKoqapgsN1OhzipaTtQpnRCpFECGtwMrQCTw72UCoE3SaTCR6Fl4Lm7e8SG8dyseKsKlnUNclixtFySpXXxFLwvbffCI7sb95hpBWRqdBpRilgXpc8Xi75zp27PH7wiLOTM8bjAcPxuANDChVFVHUNQtEay2g04vNf+Bl2rl4j2tkhX69o65K6qXnzzTf5fdPy4M597h+e0DjBYJjx+OEhf/6zn+cf/Nf/hOm84OqtG1TGYIVnHCXUjaHtLBMyqVFRTNu0+LLFCodwAokgimJWxTpoFKSgdZbBYESWJMDqxzhLfDw+Hu8yNtU5/wTQeT/jacLodzPW7Js5hAzXTAAdstMyhREwTJhvQgC32MgovIMoDdEkWoVFcMiPlESdbKLpuqqV0ozG46DpsQ3GtJ3AXHwfcOrZISHEE6G8F4HSO5UdLz6m/6q1Jk1TAJI4QUkVFrndIit0y3VMkw/aJqWDxKFpw5yslELHMVW+5uGjhyAlV67s84W9n0VJz/HRY+oyVBKQGi8Uxli8F3h6H6kOvBHiqdyHKsC+93imQFPWxYSkScJgEKJTnDGMRiOcC34XbVNj2hDeN1/MKaua7e1d8rwgz3OiJObo6JiqqcBBpiNu3LrF6dExbdMwGo+ZbG1jvWdd5BR1hXWONBuQxDG2alFCsphHIFzQ7DShu0lpjZKKOE4xrSGKY7QOq/K6bljMZzz/wu3QjWRa4ihGSIkX4LynbVu0Co/3JnQARkoTS0lrLCcnR+xd2aNpKlQSAw6lRHCKbhrA4VVfSw415/FoGPaTd2gVMRwMmC+XpLHGtBFpGtGYFu+D2B0X8vaSKKa1LdbWqCgmS2Ia11I3NVkSc+XgCm1RgrXs3byO9JYs0ezt7NA2BVWV8/KrL2GMJRsOkF6E/LI0IS9LqrLEjFtSFaHjCDmaIJVib3sbaR3i4UOiQcYLLz5PW75GkqSgYsq6JZc5sVdc379KqhKqskQoFfyK2hbrPHmR09Rh8lBJymg0Qg2HeOuxpgUckY4woqaoa2Kl0EqHPDsXnLylEsRxhJAxeZnTNg06ivAIiqrEGE+WxahY8+qrr7C9vcVisQz6r8awWi7Bu+CG27YkGtoIjo6nnM4rBpEmSyLy+ZQoX3K6XpM/uEMxX1I3FWmyw+0XbrP33A2SKKGsCqRUbN9wlHXDYpXzm7/9O8SjMel4hHWhQ3M0GrOzu8u/+cf/bd648Sbf+d5bHLctg1JCmvHm/Xu8eece//af+pP8uf/Zn6GuK1QScffeXf7v/4+/yXx2ymQYk0UxkRBo4yjzirIJWggdRUz298iXa5QKzGpjLJFQlPbdr+GPx8fjWRyX9T496/I0IHX5q5IKoVQXpyUvmFsG/6Xg09S/NgSWSiKl2sSRSCWxxmDahqYJc3VRFqgmlOmF8CQqJsoSnEvI8zXWytBp1rZPWgR0HksQ3MAvf57LlgP953CbbrjQwaa1Iu6E3H25L4ljojjqynfByLN3CH9SKB8ApCekFdRNjY40CMFssWC8vc3B1QP2rxzw0suvcHp8xHx6Cj7Y1DgkSZqRZpairqnbZvMZXR84/EPqTH2mQNNrr30LCCGEcRQjBGRJHPJ2opDvZkyLNQYlJDpKuDLZ5eDggLzIOTo+Drk4cUSUREgEsdIY0zKdnrFYLHjuuef45Kc/xWAwYLpYsC4LkJIsG5BEMU1RsFwscNayWq+JIs3+/i5lVQT37HWBcy2LxZzRZIIybWjBd46mqTk8fISxpkPoCmu7IMOuKyCclIJitaZtQrpzXYRsuePjx2zvbTGbTtFZQlFUDMdjpNKUdU2WdloW7ykGA7xzKCHJ0oSQXaTRwnP46D6j4QghJdYb8irol5I4oVjmLBvJ9mSbqilpvUElEUSS1hkaE6hUYRpmp2e41tBcu0GkBFka01QVZbkmShRf//pr3LvzgEGa0dbBuyPOMhbFCus8WZwiPVg8XkmMs0znazKlWS5zfFExHE84PZmhVYxUumNZDPk8Z5EtqYqwSmnalqIsaa2hbtugRXMhbTvP1+T5it2dHSKlMbbFmZZ8uaAua85OTzkD6ibkr82Xc6wLLfZ5vkJqSVmXtNYyGAywDlbrnKo2JHFGNhwQRYJSBzZLSkWUxggmIGF7MmE2neK9ZSuL2LqyT6yT4FzflXadc6jdOaM0YefmDZZnM5zSrKVHCMm8WPH46IxYKUbjATt7e8QOyqZlsrPH1etXg2u31uzu73Pz9nNMtiaczWfcOzrGComwlna95p/8f/5rThcL7j14yK//xm/z/PPX+MLPfZF/+s/+OadHpwjf0jiHcTWRhwRBiiBKB4hYo6KIF2/fJo5iZos5GEckFViPbz9GTe81fhJtBt5rm348xpo/mvFBjkd/87/szH65HAZdi73zCKlCnIhS9DY43ashRcgFVQriSCOUIjKaoEnqopTaZsMIGdMG0OQ9dROCeJ11SBU687LBGK1DI0gPgC6Dpov/LgIheIrJ5SXh+8WyZF9aDFYCduN7qJWCJEK70CQkEJvP3r9X0NUasmwAMqRjiM5FvDEtZVWyznPyomBvb5cvfekXOD064vT0LFRYhCDLRiAV6zznZHrGbDYL+8labAcGgxHWRzueKdD0rW+9/sTPUshQrhqNyAYZ3geKUytFmqRsbU24efMG6SDj8eNDWtOyzlc4B0oHgXXp4OTkmKIIQaTL5ZL5bEpVl+RlcGAejoYMh0OkEChnGKT7ZFnKo0ePiGLNjZvXOT4+ZHd3j4cPHrG1tc14PGI4miCVZDweEkURcRyztTXh+OQ4tIXLkEHnBUSdEHswGFAWBXZrQlPXbI0neOtYLucMhimPHz+grmrmD1dY73nxpZdw3vPg4SOGwyFKyuCOHsWUec5ysWJ3Z4vhcIh3jjRJeeOtN3nllWC1v1zNKaoKHUUMBgMioVnMlkzGA9b5krxao9MYFSusd2EfF5pmuaTMcxZnc47v3qMscuqmQShBa1pu3rrGcDRgdnLK9avXKPMC01qy8YhVHTL9vIX1ahnKhTtbQTzYWraGQ5QU3HvwgAcPH3B4/yHj0ZgsHaKjkPK9XCypmxA+mQ1TrHcbYCO0JM1S4mGGTDSmbamKgrP5WSesDx2WAhhtT7CEuniUJEQxbMngfBvHiqJYo2OFTjRVXREnUfC50pqssRgTLBFOTw45LguccwwGQ7wXVFVLNsg4PT0L7u5VibEtuztbTHYmRFptJiFrHPtXdBDg64hSK1ZNRTU95fDxCflyRbEuONjbwRw5tmdLtra2GWQDhkDpwbcW6aCZLznJvx3YS2cROiIdZGgPaZaxXhcMBwMe3bvLP3z7LW69eJ15nvN3/4t/QKQdf/bP/E+4eXCV2aMTTu49oF4uQ7k4yyicYb5aUZYVCEdZFTTWEmdZsMj4ycMDP9Hjvbya3sk35+MRxocGc/79l+cus0xPdIU95bE90Oh9iXxnDtt3yQVtU8hoE1KgEGgtg7xCK5xNgBBN1DQ1xrREUbwBTQFEOYwN2ljv++QLQzZIECKiqqquizraOIBfFH4DG/PLp3XLXTbpvLifNkJw09LFnTKbzVBSsV6taZoG3SUieB90wM55jA8WCiE1oaEqS3Z2gtdfXdcIKUjSFO1ihBAsl0vumIabV/Z46aWXGSQpTd2QFyVxmnHj5i1Gky2WqxXqrbcpipK2qnDChYqAtT+U+eiZAk3//r//75EkCSA2FKcI6YmbLoEeyVZlyeHjx3zt61/De8/R0Qn7+7tUdYFAYp2lrhqEDTdHLRVN1XL/4UMeHT7GekdR1WSDAVs72wE0AcrD1SsHbG9v8fDhA5qm5v6De8zmcw72dzk+PuHFl15Gq4j1cRmCEJ3rNE8R125cwxxaVqsVWZqFGrsUlGWJEJJXX32Z9TpnazShLiuuHBxQFSVplpImKYfHh2xvb9N6h/OOF196meksGBleuXKFKwcHFHkeXFzxHB8eh+DCKORWRVHMdHbCF77wuVCWPD3BOktrDPPFApxgf3ePl196lYePHrLMl+gk6sSFwf9DenDGM7l1i7tv3aUsCuoipzUtpnEUVcUtHZMNx3jreeHFl8jXOWVZMdreomwbdByjpSZfr3F40tEw+EUtFozSAVevHIBWXL12g53JNmmSBtM3L8myjN29vdCG6n14/mDAZHeH2raUdUmSJcRJzHq14vrVA8qi4PjwMWW5QgmFFMF5t2wBJRmMhkihiaOYdDSgLEvKqsQBtosr6XOZhNDEcUaaRTSNYTQa8PhRw3g0DOXIrhHAGc8wG3D4+DGTyYjtra2QTVc1HD0+xhqzEeyb1rC7vUNrPekwZbi/h1KKg709hvs3WM0XbI8TXnj+Nocnc/LGYJqWvG54MD3jcL3Eeo/psqy8d9R1TVlWrNdrynXO1nDAtYM9fu6LP8etW7co6pKHR8ccHZ3w5S//HnVrSJKIV159lVdfeImjyV2GeJYnIRB61TTUrcEIz5v33iZKYmQSk0nJYDJhuVqTCQlHsx/bHPGsjXe68V4eF29sPyh4+jDPf1rH2LuND7et7++xl0tiH0aT9GHHxRLT017/YmnOOYfztjOHNbRGbuJIeofwHkA5p4DgDB5uYW4j4wi6KLrct+Dv5Am6UOdDd5q3Hu8NRZHiXEJVVRttVFVVQfrRdVf3Xy+W5d7p/LoMLC8yTnXdvYcxPHzwACEEs7Mz2qYhTVNEkmzYKN/FhsmumahtmnDPI9wCe41VkgSZSxRFrPM1JyeHxNJz/eCANE03ui5rbWiyGgxASEbDEXGUYFsT7Bt0qHqYH4IFyjMFmv7yX/m/BJpTarJ0gBCSqgylpWAsuUYgGE3GnJ6c8I/+wT/gP/87f5vDoyPiOKZpGlarPAiWkyRolNoW5zzLdWjZtEC8WGKtZbUusP4M+eBRONk7P/A41mgtQxSLEGglqRrLeBxuLt/53ttY6yiKGo9gMEhJ4piqrrl16yZnZ1OqsmIwSIm6Om5Z18RR0jnGekbDEYv5nP29XXa2d8jGIwaDIeOmYv/aNV75xCfIy4IkicmGA3Skef7281y5ciVk4JmWydYWxhhmZ2fMpzMGoyG7e/sMxgM+87nPsbO/x96VA/Cek9NTjk7PWMwWfPITn8SKEEY7ibdo2pblaon3jkGWUtUtq/kSb8ELwWA8Zmt3h2w4xDjLg0cP2dk/oKprVJzipKQFWqBqDY117O92q4csQwhweBbLJQ/v30cj+dSnPslwe5vtrS1irTk9PuH1b30bpTTXn7uFMW1gm9oGISXD8ZC9gz2s8JycndAaQ1VXPD48Ynt7m/2Dq6FrUEkirZnPZsymSx4fz0liz8g6trd20WlY5ZmyZLEu0ZGiLINJXJKNQmSAA6EjkiRlMAoGq2k24Pbt25yeThFCsLe3HxoAlGJ/f5uyyHnxpZdo24blcsFqsaDI8+DWLiVVUbK/t8/+9avsXrmK7zL/JqMxv/grV7j39h12drcZDoa8hERFMXlesFqtGYxGXTi0xzobgJ5pWS6XzE7PmJ2dkS+XjLIMrSU/+/M/x+HRIb/6L/4lDx49BiU4m32Zg71tyjzn7/7nfyeYnJYVsq1JhCSKE/K2pRGSnav7eGcZ7uwisxKkZGtnBxklVPXHnXPvNZ52c+p/d5EJuNjifjGG48Pe/N8vo/VurMNH8V7v8qQP9VofdnvfL9i6vG8uHpeLfw9MUgAJxpgNi6wUm25q7z1SiMCEeB0qDV2JzHtLFEXkXdXDeYfSgRF3roUOgCgVLGScM1gbch9DBpwLuatGkSTJptN8488HwTdwUyZ0G7KhB1KXS4yX98FFcBXHMbLT3Z6enGCMoW1rvO9YsrZFK0USJaHbve9k63w926bh7PSUOE1Dh/qGxQPnLU3dMJvNOB4NOLt6BecsaZqFQPvZnO8032H0+DFSaaZdac46hxKhe9xpi5Qf/Xz0TIGmuip488F9tsZbvPzyJ/A4zo4PGY1GTLZ3ODs+pi5KPvGpTzEajqjLmkcPD/FC0LaW5XJNUxvaxlHXLRACEYUPgrVejbFYrroOAxc8mlqHEwETey2RFprWYJ0nSROsF+RljcNS147F+gS6PBwIppPBmgCWq7xzavXIefBUQgZNCx7+6T/7VaRU4AXWhtVBlqboKAIJTR00WeOtyQUxX4uzhsFgSKRDt5PzoRvQ42mqpltphJpxnq/45//qv+ucZ4NtQlXWzObBXOwrf/BV6LoGEZ62tbRtC10nhHMO25rQUditEJRSxEmCkJDnBUIq8MEBdjTMaNtw8SqlsN6TpRmj0YhIR0F/3gnh66oED8PhiLqpSRKNdwJjWhazGYNswJWDfdZ5EHqDJ8syJpMJ2zvbCCVCNpsMER/5Omcy2eH5524zTIbcvv0ct59/IZRj65baOCQeFUVEKjBfy9UKISTD8Sh0ZfruxqWD6WZeFJR9h57UaKV4+eWXuXn9GicnJwgh2dneBqAsc9q24e233ySKAv08GsWU2yOmZ2eUVUXTNAjtkRFM5zNOpjPQCScnU+bTGVeuXqHIc6TwSB0HLZQK3SPgSJKE1vSlgG50+xNj2d0eceX6AY8fPGBnb5fX3/4ub7/9NpWruXn7Kk3T4PB86hOf4N69+yxmZwghUQhULEFpShwrZ8irhpO7JXXbcv9oivOeJI3IzhZBU9Z8bG75g4zLXVkXu5h60PRObeJPG5cB2tOA2kc1LguIP9RrcCG4RXw/AHq/zNz7eowAxHsHvTwNHG0AUAc4Lh4nOD9G58xUJ1J2BuM8UoGTIeWibZvwXAmRicjzNUHrGKF1YKOMMXgRmH6lQixL6B7zSHXe5Ra0T5YsC9Y2ZVl2QcCBJYvjuPPkO49Sufg5nwaaLgvEIfhKDQdjxqMBgzQN86G1XNnfxbSOIi9Yrla0dUMkFDqOghO683jrcMaSr/OQx9mRGLoTkSdJQhTHWGsoyoLj42PeeitjlKT4LkB4Np9z9/59hNIMhiOs811JskUohRAKKeTmvvtRjmcKNP2n/+lf5+zslCt7Bzz33HMUecnp0Qm3bt5ikA6IdERT1UxPz7h64yZ5WeIQaKG6G7/pLDKCNgcE1vkNRdjbZ7QmgBEBRFJ0J2oQ2SqtQseaUTStQUcxkY6Jq4Y0G6KiwErEScLWZBuAIi82XRC9wdj5SSnQKoCbsixxXpKlGYv5kjiOsA6qpkVYy2K9Jo41xekZ4ugQJWXwxfCB9hxmWed0HlC+NYambTflQdH5Nm3tbHNyOsUTDDbTNMVZx3q97uJlquB9lQSH2rYN7avOWbwP3SBSCuo65BhZ0yKlwvc3akK4ctvlEQklwiryfKERVhvi/Of+1O53Sx/7Ak/qPoQQRG+8GY6Rh7jzrdJR1InpRRBWdo0B3kNdG7a3tijWJbs7wT3bdUHjSsvgo0KwnSirgrKqUFozHo064OufiANomobGNDgXtsFbh6lrrh3scXI2Q0eSuGsJvnZ1B+9tSO8+PmY4GnDl4IAsy7h+4yaTyZjBYEBd1+zv7XNytuBsumI42ubBo0c8PjziT/yJ/wFta7hz922uXb9BksR4D16EyUsiwnH2/eTZlQEALRWT8RDvWn7rN36Dz3/hc1R1w+9+7Q+ZrlfEWpOXBbFSzKdTvLUMByMQ0BiD8R4jJUXVkDc1TigkEVG3DRJP00JjAtXu/JMC2Y/HDzYut3h/PJ4+PjBIE+F/H3SvXtQswZPMy0Xjx4uP950WVOvgw9SHt4cyv0bKrotNBg85pQLQihONdwFoOWc7ACC6zjrdJWS4bp5z9FYCxpwbWkadVrVnlJIkWOMAaK03/n1P7JqnaJye6LBzDiUl4+GQm9evc/PGdSajMYNBRpIkLOYrHtx/wBtvvMHx8THOWNrOMDqIxSVVWdKuWoqqCmbPWodmpDRlOBoyGo+I4oiqKDhpG1xdMc4GWGNYLoOetTEWbzzOF1jnMG27AZauk6d49+HA+7uNZwo0/aN//P8GB6PRiEQrinWJtZ6rV65gmpbr164jhQQp2b9yhbfu3NmsXKQQOO835o+e3uSroykvnDdShptt1xQaHMhdMDnEgDW+izNxuKKiUYYg3AtfBd1rmjZgBRfyu4QIHQfehY4578NjLd1F7z1a6VAGlOcXYhAMSpIo6k78uAsL7gCFswgPSgXLA+/DqsL1F4XtPKikxLtQG4+SNNTDBUQ6xgpLpGOkNKRJhBAQ6+DrpETQ3IS6u0fK0AKrRBBDWhlKXs55vAk7ctj5D1lnEB3FHFYJ3SrGh+3okY3sVpWtDa8/yBLyvMJaF6wEmgatQinTtJZIxcHuxDlsK2jqOhwfQsOE79p2EZKjo9/f7MsgAIV3ny0vr24v/CC+78+bX2WRpDIeSWDOAG7f2A8TYBzo79EwY3dnhzQbkiYpB1f22d6a0BpDlh1TVi1l2TAYrZjNFjRty2q9pq5r8rIMRnZN5wfG+c3CP9HlEkBTv/Fn0zNm0zPuPzpCJRlV0/D4aMrjozPwHlM1xFJycrIIuXpZSuscZeeIjxDUbeh0+bmf+yK/8sv/BnEcbY6ls30ZQpDnJX/5L/+Vd9u5H493GU8rk13+3U86gPog2+e/75v391oflSj+fOn6Ph9/iYF5mnfR5cf23kS9pUBwAfdYG8Ldo1ijdCjPhXuI6+xrJK1tO+dth/cSCMBDCbXJlwuVi14bFaoCF928oygKC/5+4d8xlhe382ks4Tsxht478A7hIVaK8WDItatX2dvdZZBlzKZzvDFMT09ZLxaUNnTyKSFJ05TBcIyOuiqIsV3Hc8M6z6nKkqIsyMuCLMvw1pA7y/z0hGGaIqViXZSs84KmtSAVWEt3QwmLd6CqK7wxH/s0nRwtQgbXSfiqBTgPh8enAHzzu9/jImZWCGIhMa4rJRmPFAqlQ9qz954oklx2Wg8nYPjeOI9xNkCbpu1W793jANsYPDVSQFO3AYT54MKxmM3CDfUDHLlgv/90F9OeOQHQicK0Fp7y0HeaADwBTE3Pphfu/RdRgCSOFJUMHXA4UF2XiOvYLAji60gpQHYlN0ndtHjniYRGSkG5ztGRwliHMX6zUaIDLFJAFHcXvgtGbs45pPco0V/I4TNrpWkJE4yUCuENkYwDELYWrTRSeyK6qADvsM6gdKDOtQoHLIljnLHYtkUqQAaPK3XBo8S6foLrtqHzyDKtCYyjVh09b/EuUOiR1piqoW4t43QAEPafh6PDGa1zJJHCOxdiYrzHdkxRv+t1N/majeme2Bzv//K//Hsb5k0qeaEJ4slj++Q3738oCO7njUPQslyVhMJfFyFEOC8n44Rf+Pkv8Zf+0l9iOj3pukF3ESI4BDvXUhT1x6DpIxw/6QDp4vhoSn+XVyg/ueMyaHqncb5wtkgVQIwxlrIqQxZnx/6s12vqjuEHhfMWY9uNSFpIuoVKMJSVQuAJzHowz1Q45zcVgBDn8v3eS32loy8h6k54fVGv1T+2f0wf/ut9mMuddVR5wfxsxlGWgXE0RcV4PKbMC5qyRFhLJEJQOh6U1mxPtrh2/Qa7+3ts7WyHIHopmK8WvPXW27x1520eHT7u9kWNkgJvWnzTsOr0WE3bUrUhIUJLiZZBw3QuKm/Jy4JI9O7sH+14pkATQJyoLsDQkyQR3juaxqIIrqvOdbbuItwIg8V6p/jHU7c1yoogNnOeqj73ldmUiAg3aynCTb6fC0QAs9ggP0JJgZYBWAlBiCJRiqIoiaNwgwc6YVy7ced+Ylro6uodbYD3DiXpLqwQThxphVCCqjLoSGCMx3Xml1L32yaIOo1SyIm79EYE8WE46Rq8tahIY1ob9FIiAMm2dQGYdPvR2ZA51+9PvEeGWaAz5bQoHzRZAkEcBZaqMTXSBhdte9lorGPxrLEI7zuHcN+V/gTCOdarEk9wVm+aGuccVVWCECihAvARXVZdW4dt68SNjtBR5+uefZHgw/Pl5mAKpA4lVoBIS6I47sSM4XdSSuIkuNkKCTrSJEnSTSphm9I0ZTgYsDJLhDGAIk0TYhcmpjSNWec5w2GG856mqtBtuwm0bFuD847JeMRgOKSsK/KiQGsd8uuamuVyye7uLovVkmvXrjGdTmlNl0fnw34P6eHhswWWMwBRiUbQmchZAz0oD9Rfty9C4KZWwVel9TYsLjp3edtNlicnp7z22mv82q/9d/y9v/9fsbOzwy/+4i+ytbXF/fv3mc1mvPjiix/oev54vP/xfm/SPwnjwzBBHxpmfRjWqStff9D37Pf9ZW+jd+t4DH/vdE4IlAxzXlMbTBsWX2nWd7MFPU7QMYVruG8CEAQNrrUtQgSjWe+COSQIIt2bYzqcOy+n9dviuuByeJJFappms50Xfah6a4VeC7X5XHiEc2As89kM0zQcPnjEIMsYD0c451gtl5wdH2+sZrwHpQx20KKFYGs04ta1G9y4eYt0MGBVrBkOhjhCCG9IqwjvNxiNmAwy0jR049VtQ1W3NMaFuBbrMN39J9IKJSKcsyG65YdANT1ToClOZNdmCXioWssgjtHKY4wLrqudUaDrSkGorrtAKSbZhPU6iOw2VKkAHUnqxqIjGV6HwGBdZKA22Maf/+ycD54/3eMB8rwEoGkMSgZA5n0APrJDC9YFQY/ogJt3oTylZHgd4wJr1Q/fd2EIQZbEVD64p1pnMa3vGCEfOgE74CFluMC6hUG48XuDtw4henv5sCNV50brOqRlXdjuJNJBBO3YlDe99yRdjl9twwXY2hYtFVIqahNq10pGNKYJ7FC33xQgZIgXcC6sVlp7XhpVErSWXYHTEGw+gxC9H1qHz9rWFustkQz2E6Y7J/qD00+GQnbAD4iU6BitbsXVnouWW+OwrumOV/idtY6qrDeTS1021OWT3RjrNme9WncTsKSqG+qmxrhQqtVlMONcFcX5Z5ASLT20nXcKcHQ2R02XWB8udutK5HKNjoLmarVaYZqWBw8ehPNOyU1Jrme+nA2smVIhtgYp8E4E/YOOMGXRHfVQLrQmOIgH2waBtabbb4Gt85d8Kr33fO0Pv8Zf/2t/jfl8xmKx4He//Dvkec7bd+6wvbPNZz7zGT4eP/h4mlbmJ338oJYI51Xzn9zP+rTjcLmkdfFxURQxHE42i2BrQ+NN07Xc13VNEEwKkjgj0klgpcy5GaWUEiV11+3WbjyXVNc17h0bYMYFpvq8JPj0smE/ehapf81e69R3bfYAcZM9J0D1ppjOspovmLUtEkGWJMFw0znqogxkhtabudm2hny9ZrVYst5akm+tOl0mZEnKeDBkNBh294ZguTAej7n9/HNsbW8RxTEIQWssRdkwmy85Oz1lPl9SVRU4gRQelERGGvVDOJWeKdB069ZtVssVw9EAJQVNVRMnCWVZIqUkTYKgWUnFaDRCa81oMmEwGvOVr3yFv/gX/yL/6B/9Q15//dvs7u4Sx5qqLtnb3eeNN99mPB6R5zlpGnx26qrpXFnPQZZ34cbkvA8H13uW+RrpoazOb6iCDhx1o23tZlLYAJYLNyXrob+FS+gYrJ5u9BjriWNBsa7xEJiGkAuJIzwWAgjsqlHB2EtArEKJrTF9wSW8runASNsBNA9EOjBZgiAEvqij67815hzQKdVTXeICOgsarFQPupw6tylZSiGQSIRwCCXCTd51WUj0NHZ4qUgL2iZsbydPAjzeWXRXypYShBII64MrrhS0JoCxzVzR7ROh/GYfhf0jSNIMrTWr1aq72IOSrV9hbahx0WmIQiJmKMtFgemsqgoQJJHqonAUkVRdcLQg7py667oJ8yOexppOa8dGD2St7dLMI3y/+iOUj4uiIYkltXF4D7Y/lr0szPYfNfy+B/weSRYnpFnCOl90O9FtohiC4k8QaU1jmu41+uJfX0tVIMI5ee3KFX75S7+EihV7+3vsbu9QFAVREvMLX/oS0+mU//af/Qs+Hh/9eFbAE3xYpukn/7Nd9Cl6mvbnImiy1jIcDvnc5z7Hyy+/DEBRFFRVRVmWrNdrjo4Omc/n1HVDkobGm9PTs43PkpDh2uyZqtAxF3RRUqoLnksB0AT3b7v53fspmV70bboo+O71Ur2UZeMqbi1ZHPHC9RuMshScD47gQKSioLe0DrNlNgtWgcA4S1U3nBwdM5/OeOO73yMepMF4N45ZFyWzxYJynQf9qwvWBQCT7S1uv/gCV69fY//KFbLBkLJquXfvAV//xjd5/Vuvc//efZZlToiuIsRgafXUz/yDjGcKND14+JDdnV2Gg1F3MC2j8Ra/8KVf4rvf/U44yTwsFgtW+ZokTXFrwboocc6R52vKsuD27Vtcu3qN07MT5ss5Hsdka8j+/j6j0YjhYIAXsusaC6UO2RloDbIBdd1QVSXXr11DacXxyRFJHNO0NS+99BLHxyfM5zOuX7+Oc47vfOc7eGP4Yz//8xhjuHr1CnXdUJQVt557jr29Pe7eucM3vv41pBSkacp6ueKLX/xiWAXYUMN+/duv89nPfpb/7z/5bzi4cpUrVw64cvUKpycn3L9zl2tXr4afz055+PAhddXQmroTj6cYY1guV+zt7bJarZhMJhRFwWQywTnHw8en1MaTpRGDzuBRyFC6CV1xYTXTdgZifUddL3oOZR4BWLgArHSXSG0BY905Le48Arcphbr+Qd33tTnnw5w7ryYpGYCPJwDKHthZHyYU121KX4GKohBY3DThj7EOjJe3dKaidKsqgZIgCWJCZ2zXZBPeTMpQp3OdH4rrmBmBRwhP6xosECmFjiSuFlRdN2FtbGefEOhkOgAYMFqgA4UIlLzpujelFBjjiBWoOKaqmg0b2ZOgsisZi0jgrN902ATnKwBH3eRYVwMepbuJD7rOksBGNhc6RoUAqSTOg3cWXM8owte/+Q3evvM2EEzqhJKbbU2SFOc+thz4KMezVJL7QcfGrBh+Ysmmp4m+L4OmiyW8UObS3eJSsLu7z+7uLtZYzs5O2dnZY71aA7Ber1iuVkQ62XTDeu/xztFe8FzaADbvu3nRYWzTdchdiCrZ+D89CZwub/vl0QOki92A/uJn79j9Yp0jjSXWGi0VSgbj415yonSEjGNUFw/mgbptKco6BPUWBYvlAqEVKo6wPvxdIkjjhFQQUhO85/D0GB8rCtuwbmpG4y1a6zlbzVlVOaVtMQJEEhFJSRpHvPD8c+xOJvzWb33lozr8wDMGmtrGMJ8vWK9yQveBYz5fcHp6xmq1CmUhJTu0HUzEwsEKbf5//+//febz0Gq/XK6o65q6aXjw8BHGOE7PzqirlrOzeSf4DaUdF/hPZvNFCNQ1BmcN09m0u6E6xuMR8/mM6XRBXVe0bUNR5EFDVFfgHG+88V2qquF733sDax2ttbz+7e+QJDF5njOfLwLLpBXOGMqyCiCjK+UtFmuiOKGuDY8fH+GcYzabsZjPqYqSJImpmorZfE5ZNwyyFF862rbpjMEStA6Os9YatFakSRK8QJzgYG/CcDxmtpgznowYjofEcXBgr+v6vOaOJBsMyPMiCMvnc3Z3d9nZ2WW5XGGM4ZWXX+H1117j7OSY8XBIEscXauZhQjHGUJShsyKOdeiQiyLiKOLw+BghRAfw1hjThs6LQYa1LfPpvCtTCbSA7Z0dPDBfrKmaBq0IN/ReS9Dpzjx0GUeKtgwTU5alobzl/CazSErJRibUoTVv2XStAUEcz0aOtvm+qluU4ZykEWBMYAcvlnfDa9AJOc+B46bK6AOgMTY0HGxsGDrmzPuwadaC7B4b/h6+lx2GtdZhbYPSASj1fxOKziSPTfC6js/BsJASoc5p/6BVazg5m/7wLvKPx3uOP9Ig6kN8rA/Fannxgd/raZYC/e+fNsIc4lkul9y79wBrHdeuXuPTn/psp3O17DeWa1dvcOXKFb797W9z7949bt54Dq011gabnKLIWS6XG1+l3jyyZ62apsW0LU3TkKQxSRI9sW8uArqL5baL7NHlLL2Ln7e3yul1YNJ7hLWcnZ6ykophmjJIUqIo2jTVCCRKytAwpCNEZ9czHo7YGm9Rd35367KgJeiOXWuQBDlNlMToJCbzAypTc+fBfR6cHjF4K2g/k3SAUJqyrDmbzpiul1gtGSQpWRwzSBOef+lFbl6/9sEO8vsYzxRo8h6qsgaqJ36/6pD6uw0hBA8fPtr8XFz4W9s0GAOL2RJ7Mazmkpi6pAT6C8exys/fd52vqKqW2Xx1vl3rPJRgZCjtLd66H05Uvr+j7qJmqucKzhbrJ7TcWsDXv/ktnPfUZcXDh4/w3mFs0FbV7VH4PNaglaJtQrijc5bWWHTUgAdjHE3rmS+WOOcp6wbd0ZjL1ZqiqHEsAYGUZdAzOdvdrDVSaoqqxNlgqGjalvVyiTOWuvOHenD/PlVVhbJU09CasE266z4LEQOOpg1lKKUVrbXUraGNmk77BHlR0Xbuuk3TQldTz7IB41HEep0zHAwYZCPyskJKyWQ0Yjga4XDUTYN3jnydkyQa19HbIImThGwguXHzJuv1mvVyRVV5kjhhazLuSmr9PgyhlNY5ojiYsdV1CAO+efMG88WU1pR4L4g7E7amrTrgfm7Q2bQBrMadj5S1QY9njAnaLgmRkmjVdbM4ixehe5HuHOm7O5UIxFxf9nQdaNqAukv9APai7kv2gPDyae43YMz3Zlb9X0QAV0LQNVKw+Xt/3n5IOcvHoxuXNTH9988a43QRXLzbNnuenRPmvcXe/vt+llKRpgOSOGU2m7NcrlguVxweHvI7v/NllssVN2/eYHd3j6YJfndf/OLPdferh1RV6KbrS3HD4XCT1bZYLDojzKBHrOqS1jRdaHuYU6qq2pT64FyX1JffAJI0IU3Sc1bp0r/LZqrCO5R1KA+2i3KqdRGYJkLTUN/oopTepDDoKEZpjVCd36FWTEZjnAoayuVqTdPUNHWF8w7rPRaLTjXJYIjXgtpZ1rMZzs8QUgfjamORUUSUOoz3LPIV0/kU+Q24e+/OR34ePFOgCX8uFXzfQ/STeTjocRwxHKTdywm8txRFvulMePL93uk15fnSHMB5qqq98PeNtKfrbnhCIdJ9fXJSDNsTSlfywmMvflrrwVRNl19E8KnohvEeU7UbgFYZQ1ldFDpbqOwTGqWiaGlN0FrFWiKkpGktQkHdgz//5G4QiCcuLkEIncyXq27XBOr25PAI74OYvG9Z7Vc0vcN2/+GkkFRV3YVQ2sBEdQxHWS7pmrioa4vIKyKtGSQpgTpROCNYzAvKpsY5CSLC2gAirAGtY7wvUTJBKEtdV7hOcI2Auq5pmiaYrHW1QGssUWdsGSwgAtOiOvo7gD9NVamunCVwbdg3vY1FVYGXwR5AK4Hoy4sqRK/0HmHGWIoidIIoJdna2urocUkURcxmM7SM2NvbY7VYkCUZrsuW20QkeI9QoQPw2vVrDEdDjk+O8Xh293ZZLJY0TcWnXv0kf/i1rxNpj3PhOA7HGVIltE3LfDEjjmFnZw/ZmaE2TShTN3UTJvaOxfPOnQOlD3FpPuvjnXxs3s+4LJq+DJLez3N/1OODAZwnYPhHvzEf0fi+bXuPTe3LVsC5/kcKhOt/d+GlfJgvldQMBkOkVOzt7TEajVHqFBDEcUyaZsQdE6+UZm9vbxPBVHXebIPBgOEwSEiyLCPPcx4/fsx8Pt/YBTRNTVkW1E2Jc5aiKFiulriZoyzKzbl2MaPVex/88i7ugh4oOX8+V1/kEjqGPY4ioMXUDbbPles8pKQw58/xoQOwdyJHyhCOnoacORFrnKCzLbGdrtXgG0FtW2KVEusBMosQ3uHrhtZ0tgtSk0YROo6ImxhnDG0lqLzldHbGcjn7AGfD+xvPFmjajP7MfB8XY3e8hRAMhxlaB5E4eJI4RinNyckxAEVZIFWM1hpjQoit1poojnHOd+2aAalrLTf6DaUktg1rfqkkUoQOtCRK8M5TVfWmhbxpwskUvG385oSSSuCswRnbuVSbzc1qPB6QZhmrdYGKFGVZobVCSIXb6IoC+EiSsK1VVeJ8EK3jxXm5UimapiSOUsB3HYOm0xUJIh0AQOj+67x63LknkPMBPGgV/IsiHToa8d1jbO+wHhi1pm2IulBa6+ymM0/KIDoKxqIO2wRtk+gYEGtDMjgdILHm/Fi3xrAya2RRoYRklecYHFKEDr48z1mtVuHJIgjErXU0xuBtCLrEdtEwFaxXD7oJxaN8KAPnRRUEn/78swetu8Cuc6SSRJ0P1dt37iOk37CUTR2OcdMGbwoh7cYw1TmQ0qFU0EO5zUqufw+BVJq6LlHKk6YZTduSJgmj0Zi6KJlsTbrk8k7wSfDVkkqSZinXrl9jd2+PtjPFu/3cbY6TQ/LVis9/9jO89cb3oOvAs94zHk4YjkeUVcl6vUJIyd7eHlor5rMFVVkHMKskXjqUUoxHI8qyxhpDHAfxZ1kUKB1TFOUHu5x/yse7iXU/KoD0Tm3xH9Xjf9Tjw3TqPbGfPyTK37yGOF9ASiFx/YqI8+3yPnj81U2Y/2/fvs1zzz3H7dvPdfN0mAvH4zGTyYTRaMR6vWa5XDIYDHn+9m3atmU+n28ec+3aNZIk4fT0dMMgnbt8e6w1tKamaWrm8zlKKeqqpqmbcP8SYWGrld6wTqY1tE0bAGS/ELiwe2SnSRJSdveJEHCfxhGuUXiC/2GoJARDZG8dbdN081TXEGNbhAxaSYo8fC9CeoOKNBawznb3Qx00UHVNQ4uLPaN4i2w4ZDga4QmyG08fwdV56Qnw3mLbBtM0NEX9gY/xe41nDDRdPskvgKfLf7p0LWmt2N7aYrlcMj+b0raG4TDj4OCAg70DhICHjx8xGI3YmmyzLgryPCcbDtne3qZpW/I8R4ogTNMqABspIEtTvLVY0xInMWWn04lU1DEWEWmaMZls0bQmdCbZcEJYE0760XiIkpLlaonWgrIM72+t5erVqxxcucLZbA5CcHp6RpwmZFmKs571uqCuW+I4IY4ijLXEyQBrDUpqEBLThEiTLM1YLOZsb+8yn0+ZTLa7Vtagt1ktZwgRqFflBRrwsncz7zUwHmEdGoikAtt1mwmxyT9TUuJF6Jhwpg0lJCkQSgWQ5M7LTZsSz+Xjtin9+A179+ThttCxHVr2Am3TiaT95rH95FRXFfiL+Wx0oDq8Rud3ufljH+OiVZgsjA30ryeIxHs/LCEDo9WPvuNvswmu0yp1P1tjKfKO+u5YrP6zW+s4PT3rWLyQVRhYq5o33nwL6R3zLhuxF4P24MsbS1k3zH/3K0ipQqwA8MYbd8I+cI77D/4r6iLvOjTDpPP4eNq5xQcvJwG8/q3vAIEJwwdheJImZNkABFy9eo31ekVd1ezsbCOF4sH9+2xt7XD33n1+msZ76Vwu3uCf1nHV+/D0Pjq97iT4cJ2bDW5iLC6VTz7suMhyXTZA3Pzs4ckL9P2++Dt8/y7j/PL+cEDo+17vaQDUd/Vl+qaVy+//dCAluhKBs6HpI2h8Ok1hvxAUstNOdlpa7yiLgnt377Jerdna2ubu3Tt85zvfwZiW9WqNjjSLxZzjo6Ow0POeO3fusF6tGI1Dw1PbtEFbay3r1YrFYs7p6Rnr9QrwjMcjrl+7xsHBPkLAfLHg7t07LBcrqrLCtDZEX8VBgmEbixeuA0EB9PluXvSdbokLwEl4AQ420WM4jPQ0vdYp1ngkViq87BbM0mO1xKFBhZ7epptHgzzFdf9AWo/qTBD7bXDd9kSxxklPXZeYM4POV8RxitJRaFIKKC5EsETBj05KRTQchZLloPq+Y/mDjmcMNMGTCowL37/DtSMkRDqg6qPjE+jMAGOtMNZyeHgcIktEyOmqypKy8+ZxwHq1oshzvBAbEW3PZvqudpyr0DnQti3jQcY6r4giReHKzvfJU+YVi/m6sy/ouhus2zhtn5xKBC6kXuMR3SrFeU959z4PHx0itKSqOuSuVAciBGEzwuv2F7MQil4ETbfteLsBYtPplKZtaM9CuHBvZqikxrQlSni0kMRJ3FksBL+mADLChVVWFcJ5NCJ0TyhF69vgoO5DZE2WpTRtGz6XAKUVsUoCk9fd1FVnieCs7dimvgQWUdfVxvrfNAacRW1WNRpTt2ilieKExhjyujrvauv8uZyzCHzXpn9e9rx4ygTvpPNWYomkj9qxHRO0KZl295LeL2zT7r+5oV0+X3vAFN7VuYurZP/kaez9E1YVvS+F8+DadlO6jVQoC/SfyXfPVUribEgI71ky03XwCQSmXXX+V9BYuslG05h200GntA4RQX32H0H0bgu7AbBvfPetTefkbLoAIXCmJf8hrOx+UsdlL6Wn/f1pP18EOxfbvIPflumYbL15bN8F9bQk+h+UDboImJ6q2UGEm+YHHBeq7x/kWR+qlPd+2uo3f++6S/v/Li/Sevbo+56HCNrCzuy3f51wvQbvJalUmAsudODWZc3D/CGPHj4iiiLSNCXLUtI05LTFSVjo9sw7HpbLOY8fDRmPxwFUKMUgSymKNWdnpxs903q9Dn5KzjEcDrh27eoGbD969BBrLFUZFuZChNgr6yzWWKwPDLzWehM91sH5bqI7v6+eM1D992HxWJk2OHYrifehRca7IBXwdPfJSOOjcC+yFyj73kdRiNAKHYySw/F3PizKhZSBnFCOtm0pyqDxjbpAX6lDZUFFiiRJ0XFEa0xoABpkIeor+dgRvBvv/3L0LvgRCdFrahR/6t/9d/jzf/bP8sJLL2NsuKH2dzrftccjAv0HG0IivLOHuq1ZrVZ4Z0njOJz0eNbrnLauUDomjRPqpsHYADbmpzNOpzOKokKpIAJ2NvhYGGtpmpqjw8f85m/9BifTkxCn0X1SW9ehe+1CiC3m3Vq7xYV/0DXhA12/Oo66aQF/4WU6ASPB8bv10HpLVdUdRDoHAz0rg3c0JpCkznpkByI9BHPEtqHuzBODP1Oggmm7Y9FNTtbYzWQpZUinNq3B2QZvoaktUgTDUgkg+rZbt6mBN6YFEbyd4iRFSInszE77G1HbNPSu5aa1lEXJIBsQxynGtLRNjfeOWMfEURz8qETfft+BWBe21TpQQhKliqpsUDIiUvHmxocI2iXn6cCg6KwEzpPPnQ/mk/2h8pfdJJ92ZDvUZHpthTwPnYbAVIVz1m9Kd8770JUiIyQO5wwOj+6ch6smACwpJVKrDZiHc2ree4e3riupempbbc4H05rNSlyI9/4Mf1TH5Zv3O7VzXxw9u9SD9d6Yt9eqXWR9ftRi8A2wu/D/9zs+uOdSmO1+mPqnJ/dbeD9x4T2fts2b57zHZm0O6/e9RABa4Ti3NE3Fei03BpKBVQwL4D4TzlrHaDRka2urO/Yhs204HJDESadnbIInnNZEkSKKFHUdWJXFYsFivsB7T5qkrHUeynOuCeBJyo2OKcyp/dz/xAfffB554S+BEQ+VB9Mv8Dqwc85f+Ev765xJ7xQTGyYpEBbhl7ZbOIfFngARDIzR4LUHEzrZbaD8umahMGeZpkYoSVnVWO/QcdA5/TBKzM8oaHrHM/T7Rg97rAOtIIkV/9P/4D/gf/xn/hzj8ZjVfMqjR494/vnng2VBFKOiqGODzieoTRHFB9+Mqgqp7rHWAXS5oPZvypJIpyipNi7N1nrqsqYoKowLERV9B563nfjNO/Ii58/f/XOsVguc8CitMKbBAzpNED4YS3oCGyA2DNA5y9QzHT21Lrq7bDgZ6cTFelNnDp5LsmuP7YCJ97SuRSoVmA3vO78mT1M31GVJvlrx1ptvsM5z2rphuVxSlgUSQWsNVVXjrGEwGnF0fMxsNqO2YYWzOXpP0SU463D0JmodkHU+XCgE3sV1V6BxfQt/AIKBYRMI67pjx0aYjqALlmVjL+AAYzyNzcMx6LRZtTG03iGcC0ac1uHEuSAcESZa5xw0HUXvLZI+08l2ALe7Abq+ESCwnL6bZAIl7roJgg0673VdUkikDswRPkyq1thuIg3gRWtFmsbUTQM+nDMB5NBpniSx0njrcMaGAGfoApZDeVFKQV7X4dy2HQ3fHRLX1RY3rFh/7DrwdrExwPV03E/JuMz0vB8wc5mZuiws7m+kF7ub4Fy8ezEa46McF8tz3/83+LAH9odpWPlR6K6eBtIulypFd71/KGsDIYiic9awB8Z1XXemuOfjoji7qkrKstzoHaMuwilN0ydYSCkVZVmwWCy4fz+Uxds2uIbXdR261DoW0xqL1kGYfrEcbMV57NjF7e6/ngOqPqala4QSF+atsEH9o57YV75vu4VQEgwvHKo+Um7Kcra7F0LP5HWVEhPmKOm73E3jsd3CL8xfirZu8VJQ1UF/KbRExxEfhu98r/GMgaYPeeF2xzaOU/6X/4v/iD/5p/40pql5/GjJ/v4e+3t7xOkgnIRKbWqlm1Nlc9GEbRBApKNO5xMOfr8iyLJRACodueMN2LplsrVFWxqc9ZxNp1RVhVLhhBHeY5ogXLt94xbjyado2pJVvsJ5x2gywgk4Oztl/+CAo+MTkiSlKEvKugxt6T7ohMJNNGNrss2w81IKF1uCtRZjHCcnJ3gPN65e48aNG+FEM4adnd3gfyG6yVz29W02Kw9rDKZtqaqShw/uUeYFbdOyWgUb+wC++vJQAJXzxYLlakXZiwJlJ2L3oUNMXSg79Cvuft8KH2waQyZf8F2yHXWLD2XMAGLOwQiwkSv4nq7rXrupW+q6pigKFrMFtrU0bZjA5tMZWZoyHA1Zr1YcPn4UdARVwarIu5MJeh8k78G1AdB472h8s6GzhRWbzLbzMzfY953nMXe/7WsZ/sI2d6/jrOtMKMPn3hiJdvjEWkfTthvdC6YDeAQfMyE6ANpZPHgb6HVrwz4N03Ro/7167QrT+ZSmaUInqe/2W/eeWsmuSaFzHd9cjv7S5/zpGk9rNe9//7RxuazXl1SAJ76/7KGz8cu59Frvdxsvb9NlW4D3YxHwgccHfCnxLj991OP97osPC8z686LXrPVdxP2/jcP2pfKolPKJztgeWPfBulEUhbLaBfuAi1lxWmvSNN2w3pePff/48w19Ejx+H2DqQdPmZxAdaOolIhd30WYxIc4BVD8N99vYmlBlUd391ncLzD4iqj/XHS6I7J1AOIVw4TUDGRFyRoXyyM6wGCOQKLAS34rN3PlRjmcMNMGG4/tAF5QgUpI4SfkLf+EvcHJ0yOzkmFc/+UnW6zW/9mu/htTJBhhs1tR9ydufE9QBIwXGIk3ijQB6e2sLCF0JsY5wxiMijW1CoO2V69cwtWNrZ4J1QdSnZMc2Cajbhvl0ys7uNqPJiLKAqqmQSjLKBtSmQUpFlqTEcYTSsmMpHEL6DqEbhPBo5UhiSFNF23qSVDAcxjRtQ1kYpDR4B2mm2Nkbk68ETdOytT0mTga47oZ7enpKmmVMxmOECJ1wNAIZKWSsufX885sLQXVMVVHmLOeLIIBPEq4fHHB6cspitSKKYiKtqZtgddCalqqqKaqauqzQkSaOkw3zNhwMqKqKnd2d0BXShKDbKAr73bQtWE+UJJ1VgQpRBJ1GKmgPJHESEeuIuHOsbpuWMq+Znp1RFiXgqcqSo8dHjIZDtre3mM9n3L1zh6ZtmOdLFqtFaMHty6bi4g1GBAB3YUUq6F1wO/OlzZnFBtxtGKYOMIUulHO/pQ0Yp9dP9D/TsTqds/elSfecbeScZge8s0GQOplw5+23ODs94+zsjKJu0LEiScJ5hWFD3kG/6BBoHcqdAo/vQ5o3r99149kNIvypGU/TKvXjnQTiTyvTXbyZbcq8sNE72a7jMyw8PrhW48Pc/Pvz7YcNYH6U40dU3bzwfucg+aJD92Vwdrn8evF6DgvewB6dl/W6Mn9XzvXebwBTciEDrr2gk7t4fvXvy4X3eQIkdoBIAD4gpQvnwkWK6emf+YnRz8dCbLz6rPcXwNH5PuiBn5QSh8O6EGoc4re6qoj3QAsusG3Shb9pHYGSCB0isBr70ScUPIOgqR/+0td3eZQQG3HxtasH/PP/3z9lNBzzy7u/xG/99u/yV/+v/zfOT44nihDn319gI50PAbtJEhPJoJG5evVa17IfM8wGeOuROmhnsmzAyy+/gkDxwideCN0LNpRWnLMgoSoLzs5OcMphI2jrMpR8hKSp6g1LcHZ0gu88eoSAJE5CjAZB4yOEQAkVuiaMoW4aqqZlnRe0TUNd1UipcM5wNp3Svv46bd1gvefR0QnWeGrToiLN3Tt3SbOMg729UBqywca/LEvmszlSiBBV4yw72ztEkWY6PaPIC8qyxDQtP//Hfo6333yT2WzG9vYuw9GIpmnpp+KirJhOZyzmCwaDAVmWURbBpG08njBbTHnppRcBwfR0ShwFvyJrLFVVUpVVEIE3LZPJhKvXrxEnaRCjd6BBKY2TMB5PSNM4xBI0lv39vWBs6lrqouKTn/wkTdvgveP6zWvcev4WAFVTkxc5UaRJBwOc9zRt2P9SSoq8QOCpqoo0y4giTdsY6roJIbj9ROU8OtIbcBV+GSB6mAglxlpUlKB0RJLEJFEc2D1jGI1GtE1D2glIpejBVAfmucQoXKTbEThnWMym7O7u84d/8HvcvXOXO/fv8Qd/+DXevvM2pydT6rbpfCbogn8FvtM9SAl0fiwblZzsgV/QSvw0gabL+qL3A0qeBqp6lrBnJS7eCPubrHNBDBuinex7slnvZ7wfZkVsWIQP9j4/Kt3VhxkXP89lgPJu++LDgs4nok94koHqRw9mLoOn/jx42teL2/6E/o1OvylCN5nr36/zc4KOxaQHceeWJxtivvsnOY+puvj7nl4QPbvvL3BVF7/vKPH+M1vnEMZt3j9IKmTnLyjpOxN7LaUUahMIpUSwPVAd6JKEcp3UKjRqEYCh0CGDtLYNxn30GstnFDR9gJNXCDznXkJNXXD1yi7Xr99AK8m3vvVtvvL7XwUfsoEEBu8tfedUOEkEvfitnwv6iU6JEHg4Go0692dFliSb9lPjPEmSsLe3h5Sag6sHaKXCiSIEzgfdiXOWpqrIRgNUFIVQ2kiHUrCxaBHa9VvTIpPovGW/K6FJzi/yvkjV30QdhDJLd7ElSQbedq3wblN+c94HF3Fj2dndYrVcUzdNoIOVZjDIGI1GLJZLHj54yNbWFmmacnpywv7BAcPhiKqquHr1KsvFgrt33+b173yH46PHSCkZDEYIBOPxhP39A4bDMTqKkFIzHI1JOpft6WyOc5Z1XmK94bvfe5Od3V1Oj46RQFOFqJi6qWmbhiQdsFwsGU/GrIucbDgI9K8PVgppmmGt4drVa6RpSt2EwMxindO2Da1pODs748b1mxwdH7FcBgBXlhXj8RjnHKvVmslkzLXr17HWsl7nG9r9wb27ZIOEo6Njrl+/zvbONsvFkpOTk24y67y8rA0RAJ2ninMO2xqcs4yGQwbDIcv1msnOLpOdHQbDAZOtbYoip1kuOTi4SpGv2d7eYm93N3SQSEXcCR61jjadiJuJTQQHde87TyxjUDri4GCX+w8eUtU1r/zWb/D//M/+Fscnx8HbSwlwMBwOSJOIqixw3Q3cWht8WC7otMJ5D64LEOanBDc9jRW4/Pf3ev7FLriL7FF/U+1X3X1w6YcVgr8TyLrYPdf//ScZ8PwwxmWQ9FHti/71epDT/+4yo/Q0cNYLxfvzou+g7F/ncim1B2BKKQShY9Y5jxeh+y50N8vzsn/XRR4mh2DvvwFMFwo55z514a++e9DFMl2PQTfPv/A9F8pzeHCtoWnNpgsRqYg0Qesk1QaYWe87c00QQoUFog+PU0JD97lQILs0CWNt0Ef5oDEVThD9ECDOMwqa3v8QXYO1604Waxt+9gufJRsMsQiywQgQZOl+4JV8gXMV3jc4b7EenA/uz0oJoigc1sD0nLcBV01DXdc45xgPR/9/9v470LKsLPPHP2vtcPLNdW/l0FUdoeluukEBxQTomHAQA2MAGQM6BgyjP2dGxziOg2lglO/oqKgIKKZByXSDNNA5h+rK4VbVrZvPPXGntdbvj7X2PufeulVd3XQ7Ns0L1fecfXbea6/1rud93ufF9y1Brx/bLLWTZ0/ZRpDPEESBWBZgp/Q9wlLJysK7OK8nJL6QeO5lSVw2mhUlH4RqUmUJx0HgI4SNGSsXRszPMUnTAjEzQ38H92rgjjYadQLfp9vtWZKxWzY2Nka/b52WkdExfE+yuLBIrV6jUqnieT579u6l3+9z+vRJjh47jjaK6elpPOmxttZifHScmZkZ6vUG1VqNWrXmyI2CdrtNs9l0nYHkhddfxz333MOVVx4g6feJ+32WlxY4dWYWpRW7d+1iz8QEZ8+e4fjJYxw5foRSqUTmyp6kLpQRRX1qtRoIQb/fJ476KKVot9vs2r2Ds2fPMTo6Yh2bao1ypcy5uTmmp2colUr0uj1KYYnasaP4no/RFuXzfZ/Z0yfZtWtnodyrtabfjzDGEIShI2A67R2XuWJBMCsBaknoxjrmNvcWkynifkzH75DEMd1Oj1arRRJZtd+obwUkhRCUStYJqzcahGHoMlNsiFUKQSkIMUAUR9Tro3SXFzl1epYHH3qYer3Ol37py+n3Y975znfS60bkUcOxkRHq9SrLS9aZ9jxJv9enYPsJi57lM0b1LPAH/jXbsG4SXOhEXe6Au5GAOzwQbgypbFRzfirox6XW3egs5OcyCA8/NXuuOF+b3ZOL3Yuna3lCyPB+hgn9mzlQm4XRhgngw3/XJQgYg1EaoyVGWOqHdJEWxICbm4sUCwbaSJDPn4VjpBgG/8Xu283G7flIEBfk1llgYd19HWT3Zg41t3ViPTxPFSE3m4ctUBqk0uDCbUaDUZazq5W9PnupAs/3CbwAo1OUNkhjZR90qhDSwzdfdJqeotmHYIlstkFF/T712ggApdoo+/bvRwiDH5bsYKVTEKlNBdcDrzn32MPQR6mMUqnE3r172b59G2EY0hgZ4Y7PfY7z589RH6kBglRneFmG1oJyqUSjUWdtzRZeLJXLRXFhk1klU78UUh0ZoR9bWQKjNKH0CYTn0Cbox3awL4UCoRSBFKTa0OwnUA7Ys2c39UaDM2fnOD8/j+f77Nq9g1K5zOzpM6w1W0gfVAqeL5yiuA2zeEI4HSlBt9tzDdO4si2CXq9fqD1rbWiurrpwgqTb6dBzBXwXF+bBwbqttSaAPa4n6XV6HE4Pr3tG+bykkDbAcoLK5Tqrq8t84tZPsHfvHsphQNLvk8QxZxdsnb2bb7qJOEl4/PGDPPrYY3RywvbwETbM5HIbGWnQbrfJ1C20Wi1OnjzB9u07GJ+YQBntnJ8+lUqF0bEx+r0+c3Nz1Co1GvWGQ+8CxsYn2LJlmp27dljETmkajRFGxywSp5VaF6LLeSIF/J46ZfkgYMYYtIOae+0O7bU1fM8njmOOHD6MFNb56XV7xHFEktpinVIK6o06fhAUg2nidJ0q5Yp9Rr0uE1MTLC+u0Gw2efDBh9Ba84M/9AP8zM/8NH/6J++i340KZfNSKaBcKuF5AbJkNVN63X6BtuYJ27hO2Liw3vPFhlGijUhR/vvFttuMeD0cOhkOvyilbHFxRw7+fM85tycjgpscVngaYannmsN0effiqdswETzfz8X2NczpAS4gjOfO0XBY74Jwr5MA8H0fXwgrTDy07vDxhz+ve16XeHRm0F1bWZd8OwdKXOy88mU5IT538mw4P12XvW2ReUPoeVSrNVSm6Xd7JHFk0SRjgQgpJWXp4Usf7Rm0Z6iEVp8p8UPSLCVRycUv5mnaF7jTNGiIWllm/tLSCtdee5VTJVXkRT9a7TZGGQxdBJElWIPjw9qOKjWKLO0XKt3f8R2v58abbqTT6fCSl7yEH/zBH+Tc3FnmF+YtNJoLQxs7a3zNa17NI488wuOPP8GePXu4Yv8+FhcWOHLoEJ12m7Qf02zNDRqtwGULgCcEgZSkShFKiUhApJqx8VHSQBLLgL3XXcVLXvoSdu3YycHHD/Kxj3yc7dt3csMN1zM5Ocmjjz/Ohz74YfJEKyE9UMrGkd29qpRDrr72WkBQLpdc0ciYrdtmGB0ZJVOalZUVDh+2jk+5FNI3VkBtass0+/fvdwRExV333otAsWPnDrZumaHf73H23DmbtaZs2ROBdFl0uTKtQRtFvd7g+uuv5447P4cxmnIpZGpyCk/A4uICy6vLXHPNNfiBz6FDhzg3d440TTelJw53QsPtot3uIKXg7rvuLZafmZ1b9/2pmOcJlLpwlvpkVnRqbKSMc8H1bHZ9XGTZBcdxvCcpJD/0Qz/IG9/4vfi+z4H9B+h1++vQRwkuSzIv9mlF/HJkyZiC1WCdJn2ZJ/EFZPmgo5wOVz7gbXSmNm4zHJ4ZRgvy7fOMp2EnKucyeZ5H4MQQnwoK8tQH/WF2y79+B+ip2ecX3rxch3Djs82XbWwX+XPfGLobzrS7FKI5/FcrK6Yc+LYcGHLAixrO5s1f9CIhZd1tybPiBue3ztY5Yu5vHp/LP9sPg3vheE/S8yh5nkO7dFFvVCgrNZCXldLGhicbjQaBXyIdSel2e7Y/yiy31mAIwxJBEJIWCJbB823B4Ay1rkTsM2Vf4E6TwegMIyzMrbXm/NwcUb9PmmVURgYZcwgPpAYjEW7QyMMUuJm/ffC2jlilUmbfvis4ffo0v//7f8C/f/P3cccdd5Cl9imFoURjS3EYYGR0lDe96U187OMfo7nW4ku+9Et485u/j9XlVf7X7/8+H//Yxwg891JIMNqJKuYzeSHoKYUHVKtlRD8iBG6+8hrSwKcZSL7tzd/DfQ88yNLiEl/z1V/N4sIS3/u938tnPvNZpPR4yw+9hQ998MOgHRtFa6RnU/eN1iSZYvuOHbzrz95Fu9WiWq3zd3/3t8zPn+c1X/t13HTTi1Eq45Of/BQ/9ZM/RaYyut1uMXN41au+hl/91V9zWT4xr3v9t5LFCd/6Ha/nzd/3fZw8cYrf/Z3f4TOf+Sz9Xr8IyaHsQK6NRhuF5wuuueYq/uH//j2vetXX0Gyu8gu/8AvcfMvNTE1O8tnP3M5PvfWtvPOd/x9TM9P4QcCf//mf8973vpdjR48Rx7F7+tZyQu16xMkQBiGlcpmJiQlmZ2eZnJyk0WggXAhveXWZcqnMSKNBEAREUczq6qrLuLOhX4FBOQc7R2h8384E09SKXPqe1ewa8oUHRZmFsEgeFCKgATZTDgMeksTtXzqnJ+cZ2PCelWTQOSI6tN+caCCkzSL0rI+MUprTp09x4sROrrvuOqanp+l0Oq4zhx07ttPtdvG8wGajYOj3+3hSOA6c/ZtzJp6vNpzVJKUtrpwvz2fNF1SI3zBA5uvmBPB8n0qpos/Kj2GMoVwuUy6XP2/EKT/PSxGh4emF5p7aOQw+/0uCUxt9yCe/F08fbdps+2H+1GZO2LBjlaMzMJCh2NRpMu46pMDHWN6Pc66L9d1/BTjnx016ihBazlVaHzIs2isDJ0ppNbiPZj35e/hW2XMzRYgxKIWAcwqjyEmjCJtUQl6j1IbnKpUKW2e2Mzo6Rppm9Lp9kshKxPSiPpnKSOKETq9Hp9PDGFsEOeeASfnF8NxTNs+3GjRpHCPCgKWlZfr9iLGJCWRYI0pTQNqsOdfwdN588oHIGAwKz5P4nk/qUoGN0aRxxpHDx/jbv/1bVGoHRs+DNNV4UuB7kGTQ6/X4+K2f4Oixo9z04hs4deokv/3bv8P+/fvdC2RrmwnhiuFqjU4tzGSwA6rzL+h2+1SMYQQIeglaplRHqzR8n0BpBCmBhHLoUynb6vVJmFFx2WnaGOIoQqWaIHCwsbtffuCjlObt//MdfOb223nDd76Ba6++jv/zh/+HU6dO8fKXv5yrrrqKm2++mc9+7rMWcpUeCNixfRvtVpPXfN3XMjE+ztve9jZe/OKb+Jmf/mn+4PffybZt22h3OmitkFKgdFZkWzk5SqsQ68I+nU6L48ePceutt/KOd7yDt/zQD/IL/+U/85KXvsQWhgwCfuqnfhpjFN/13d/Nm9/8Zv7sXX/Gw488ckE7GOaSuUkPvu/zzd/0DfzXX/olvuqrvpq3vvWtfNM3fxPVaoU777iD//G23+KVX/Hl/Nt/+y1snd7K448f5J1/8E7++Z//mX4/skV4tSbwfBDaVt4GlHKFMb0hbR01PCuz/3xPEPgBvoPbhRCEwkckKSa1EDRCUPVDhDbIwBJDbakdQVgKAWhH/QKt0xiE9Ah8j8zYWVxQKlGplun2umyZmuT8/BLf//3fz+joKA888ABbtmxh27ZtRaf4W7/1m3ziE5/g1k98gvnFBWJXmsZVSHCzwUFWjX0XPt839blnw4Nqr9fj/PnzpE4zaxglyAfA3GEaThfPl4VhOOAeJklRiNXzPJIkYXV1lcXFRTqdTiFFMEwkvxy7WMgk39eltivkMS7jGE83NHfh6T1bjWojdpsf//KRu8ux4faxGUq0kce0mXO90RHP29TwOkVIz7h3UkCqlC1b5fo9O3ESjoJiHRwvFzV2mdk50jx8nsP7t/3nwNnPUqv1VjhTxX0c7Me4cQvHnVJakyRpIYcwfF+CIEApRRRFxHFCFMcobQhLJUbHxmzWt7IEdqU0/ahPP45ot1ooY/Xq6vU61WrNZiJn2QUCos+EfcE7TdoohEMAtBCMjIxRqtRIUk05zJxaa542mSdSD16anGUjsBlDysXcut0u995zL71eH6U0jz7yWBH2AixK4AoUikw57aOM1dVVXve6b+X06TPMzp5h+7btrCyvFgOp5wuyTF/YX0iBKPkEAtJeSgOo4ZGsdkiFYnS0RlUGBDY5PD9jDIZUZzz86KO84IUv4uu/4Rv527/5G8p+CYQhUzaeLKTNoGi1O3zsIx/jsccft6reDp07d+4c7U6HM2fOsLKywiOPPIxSqmjoRhvuvPMuVlZW6LTadFptjh07xle+8itQSvHuv/gLvvEbv5Hv/M7vRCvNXXfeafU18hi3i4lrbcUiDx06xPd8z/cQRRHXXHMtS0tW6+nc3DkC32NhZZkf+v5/z/FTp/jP//k/c/bsWf7xAx/g+InjF7SB4ayTYev2etxz973s2bOP8fEJ/viP/5iPfvRjBIHP8soCnU6ba6+5invvuYc777iLqckptm3bys4dOzly9CjC0t7IdDZ0LAbhKgGZUevg7zyTzbZN2ya0UpbEKARKJHhKExqbwmCEQQKxSvFMhic8B18b0sSiSPksT+ezPa0xmSmacZokZCojTTOaq6sArpSPYNu2bVx77bWsrbXyZsb//L13cH5+npXVpkXScji8eEY26cC4dr7hZXneWLfbLdpVu93mzJkzhYpzPrjlCNEwJyV3nIqZdxBQKpUIXEp4mqYFqpQ7TWtrazSbTdbW1uj1esX68OQIyMU4fZutV6w/FG55qvbUSPD5NmARiqd+0GeDP7XRodns8+XuZ+M1bRZey/mNw20l50IN1yXMOU45OjncphAuedVotNIkqcbzBuvljooeOh8p5aDywobntjFcWDhNUtokJOnhiXy83ByR2/g5rzmah5iHOV95tqAxhiRJ6PV6dDpdVptrduJXLhOWylQrNWphSF3XiZOEcrlMu9MmzRJGRsYYHR0hDAP6vYiW69eeSfuCd5qMNhhp8AOJVpqTJ0/S60fUGyGZytWOrdaD1sZVFXTVk9fBjBpbcU2CUaytrfFP//TBAinpdCOX9m/XlhJbrFdYuDOJIm679ZOcO3eOfxD/l16vT3NtjXNnz7G4uGjPVVBwoaTMtXwcJ0uC0QolBJVSiIgTUqPYOjPDxMgIJ5fn+MS7/pJTCws0pqe4p7XGwccO8pfveS+lWo1tO3dw1913c/TIMcCA0oSBT+hLqyDtCTQRxlhH6N++9t8yPTXFyZMnWVld5Qd+4AfYf+AACwsL3PqJT1inUIpCxkAIge9JXvnKL+flL/9SPN/n7W9/B54neNWrX8UrX/lKtm3fztLSEkkSuxcGYFBkFpf1JaXHlpktfOd3fgd33XUXWmd8+Zd/GR/72MeQnk9YLhP6Pt/02m+k0+1x6tRJPvWpT9vq4SpDDtXoy2dJG8MkuZ2ZPUOWZnhScvzUKY4eOYoBSiWf6elparUGDz34KO9//99x04038JJbbqYx0rBObq4FhSVFWsVuJ/k/5ByJnD9gFdwsnUDb0j5aaxdqc4iiggCLKEos2kmqiTGUhCBTKZn7LVY2YUH6lktlnW4PbSBWLp7vkCEyW1S007PieEplHDt2nDS1kgyf/ewddoAH7rzrbgB8X+LJXGjTdo15FurziO99Ubv//vsZHR2lUqkwMjLCtdde6/TMBgTe4XRzuJDsOzyAbiT55hIPSZIwPT1NkiTcf//9pOmFM/VL2aXWuejgby5jnaeyv8u0p+o0bURtLn/f6yfHl7OPjejLZue7ka+2sUzOxc4fBtzL3HKUcnid3AHKty0mvORojQ9iUMLJ5BOqXCInL7/lOE5ZvL7Atq1TyToHLnfoNobp7KXn5zyEoA1YTUUf79ZAGpec5USK/aDk3pOMOE5RWhGWyvbeIjg/P0+z1aJ8vExQstnKvmcnGZVKGd/zSJKU1tqadRo9AR54gcfo2AiVSvkST/Tp2Re80ySljZUarTFZyr333sfr5hcpVxqQ5TETO1BrqVxRecEgTSA3UQxsxlgF7RMnT687lhCelSmwchEAxayh1+1zzz33A3D27PkN29mHrbVz0oRrbHrotbajNCqDoF4mEJI4iqlPjHLlrn0szc3x8Mc/SQuFatRpSsGplVWOnznPlddcxczWbZw4doKjh56gJELKQUDoeaQ6sbFpLIcqyWLuf+B+xkfGEcDd995NHCfs2rWLlaUlDj3xBGfOnCUMAxtic85Xmimaa6ucPXeWmekZVKZ49JFH+d3f+T2+8Zu+gRtvvIknnniCj330oxw9cmTdwGFrsKnBdTPUAWD47d/+bVZXV7jpphuYm5vjT/7kXVSqFb7qq7+GBx+8n7vvvodSOWRsfIy1taabheWpu95g5iwsf6jgCggf4Xl0Ox26nS7f931vZu/e3QSBz8LiIrd+/DY8EXDjjTfxfW/6Hl7wguvIsox7773PCZm6+k1GIyUo56TkyKQYfnhD6NJw32lccxOORKcxZEDf/R5i0SaNbZLKw6bk5m1RGFsJPA+feHkIzbV/4/4NEYfBykfMzy8QRRFaGVZXV10HCWFoEQytMtI0F+Jzg8CFr9iGi3n+WE7KFkJQKpWYmZmhUqlc4DQ9GdKQrwcDeZD89zRNSZIEIQQLCwtUKpXP65wvNdhvvkH+56k6Jf8ydile0Mb1NprIw0aXey+4eHht4zpPJ9S3EeXJz3tjaO+CLE2d0zhEUXjb6ve58iQuDXzdZTpJm0td80anb926QiCEt27ZelQJhjuEfB/D6Jjtj+21ZBkkDmEdFnhdaTbRqysIh0TZkitWWqdWq1EulxHGlvcC20emyiJZ1XJl3TU+U/YF7zSBfYBZZgg9w1pzjSS2yrr4qS0bASiV2phK7hkxcFgsAimxBFhja9sAwhfsv+IAL77xRubnz3PbP3+y6Fy0oeCqpCqjXC5zww03sLq6yuzsLNPTM+zZs4dut8OJE8dZWVm1SpkekJkiZCWxfCctwAgJUpNkGcb3SQSs9vt4pZD9+/aysrrMyeUFVvoxKvRo1KqsRn0evv9B4EEAAr+C1FDyQ4RRtOKIcmgRtCzWrGUd/uLd7y7uXe463nXvPZveW9+TtqyLJ3nwwYd58MGH1/1+8sRJ/tfbf3/dshx9EyK/9xkwQIdAceTIUd7ylh8B4Nd+7depVMrccsuLuffeezhy9DhjozV+4zf+Gw888CDf+I3fyPUvfCGdTpdWq2WLUjrytX1urtSLE4xLU1tbLQxLZGnKe97zHlabq0xNTvGSW76EWr3Cgw8+SKYUDzxwL7t37+alL30pk5OT3H333bTW1vCELXqbZCkmo+AzYbBFbzc2ovV3wF3/gEdned/WKQkrITdcfz1JnPLwQw8jfOcwOc4buYPm9p8qVdQ61JlyhKmB450jX2lqcBx1wrDM1NQUcZwgPElzrYkxxgkp2k5HSA9RoGlWmuL5yF26mL3iFa+wYqlxXGS95YkRwzP1jYPyRqRiWMdnGFkYDtkMc5/CMLxg4Pyi/cvbxZzEYR7T081wzMNWG3lReXu54DjCbpMZjfC8wiEEitC9dtmv5OG/YckAhhzQC84r38sGbpa8tLjrZteWO015sWBgXTjbYIqxQSlFnMQW+Zd5YV9DolLiLCFRGWGvh4BC+2mt3SYM7UQm9IMvOk2wvvFdzszG6hDZjLfMrR6GoZVixxTK2iZL7AArBWgPSTaUrQTWj/esMjc2blytVvmWb/5mfv03fp377r6Pr//Gb6LZXBk6+kD4a9v27fzBO9/JJz5+K+985x/w6le9irf+xE8wv7DA2972Nj78sQ8jQuuBK21T5z0h8JQrASMgzjLwoRvHKBNT0nCu16JtMirjo/jVMjtH93HV9hkW04jT7TUeO36SXpwwMjlBJSgxf3YBaQztXhtDxu6ZbYxMjJJoxWKzycryMlIKfOkPeDLGygBIIahUKpRKJZTOmJiY4PTpWTrdPmFgyccCqNZqqCxlbGKSpeVlW1DWKBeGsmFSBGQqKxzUMAwYHR23Ct1pwu7du4mjmKPHjjM5MUa73eX22z9n24CEldU2H/norQD88Z/82YY2Yh0Fnc+ujBUOFdIWwsyyDJ0p0jijUq3wEz/5VsIw5Jd+5Zf5pV/55XX7uve++9Z9l0Dg2VBkrx/h+bYEgFaX4U3kiA/WQZKeZwsSZ6ogZmqlqI82+NGf+DFWV1b5yR/7yQJ9xGVVbtynvWj72c7GhA21OXKnrWlnP3uejxPHp9vrkSrFyVMn+fN3v5soiZ2OFOsInkoZHBD4RRuy/F0A1jlK+ffNHKYnC9EMWz6I5bXG8ir3OZfpi/b/3oZDWM+0beRVbeZA5ci2nVhbpW+0Qnge0pPY8kkDfl0OQfueB76PEL4thZLvNz/28HkM/3cdTK7RJg8PiPz/QxOrjeHQPGtYkCn7roRhiPQAIYiTxPJ5kQSBRHp5KE8gAx8v8C1pPdUkaUKSxvh+QOgSZDAglcAkFt3taWPrkz7D9pxymsrlcqFXctlWhD5Aorj7oUdZXFxkZGwE7Qd4Mo9jAGkGJsUTbm3HsdEuNCZReHiW3VTy2TGzk1s//kn+9I93Mj0zzZ/+6Z/y7d/27URpn1qlDhiyLLWZU9o6cEeOHGGtucaHP/ghUJrrXngdYcl1glqjUmVRAg0ohz1oMeBFOGqVKAVkKD714H089tBDhEYxPjXJD/zcz9GMIt75//0BlalJwlJAtT7CW3/mZ7jhhhv5hte8Gj+skyU9YuDv//GD/Mmf/SnnF8/xute/jh//sZ9kYX4BT3g25m0UnuMfVaoVbrjhBm684UYqlQo/+JYf4Ou+9ms5dvwkaWoH/muu3s/rX/+t3HXX3fzGb/4Gv/gLv8onP/VJ+v3MRjy1crFniXHAnhCSa665jp/+mZ/izOwZbrvtk3zoQx9iaXGeF7zgepZXVvE96zi6OraFlUohWiub0u8UZHMUR/pWeDNXoY3jGIxx9Yt8VKbp9XrU6yN0u51in35RlTwXWrPeSpZmtnPSeYFkizBd0l3KHR7XFrXWxTKlHCqEQRhRtOvl+RV++C0/jPQEInTXW5CcIPAlKrNcMN/3ieMUrTRh6PGC666hVKrwxBOHWWut4bm2nynwZN7p2Q743Lk5lBPwnDt3Dp1ZHSBtmyLGyQrAFx2mzSxHmPKBZGM9sI32ZGGfiy3PSbJBEBRZdpvxXS7HLjuE9q/ggV/q2p4u9+nzjSFvdJIu5jRtPL+ner7DobGNGXjrzE2WCuqIW8+Kzg4qRiitnc6a41BKgTQGLYzNEt9Q5gUHDgjHVZJDy3DnlOq8RJc9D2EoJEguQNmcQ6W0cdIyNgxXrdXwXJmgbj8iVYkFyoMQ6fo34VmEtVypgDCWbiOMzVQ1Co3n5AqsFEqq7LVrxwd8pu055TTl6YO5nHwO713S7P2zPpGGIADteezYuw+tMh55/HG7njCUayXSCEoegBXEVEaBFgRhCSk8TGYohSVGJxv83M/9HNu3b+Pt//MdXPeC6/ja17yK0A+I0j69fo/Az2+voFwq26yvhXl+4Rd+gU5rjYX5ecZGx9i5facL65iCV5dllrtSwhLrJMLWL7OxH/pxSsX36ACRSpEGOknCY0cOsRJFXP3SL+Hb3vAd/NEf/zE33PRi3vzGN/KZT38WgF7SI3BFwqI0op9EjIyOsXXrtqGQYgq4DC6dZwz2eOCB++n1Otx00028+lWvYnb2bH6bAQhKFaa2bOHxJw7yoQ992NYzU5owCPE8YQd4rVGptqFRIciywQAQlkJuu+02piYnmd4yzWOPPc5Xf9VXcfDQEwhhic4Ig8o00vNsgVnXW3z7t38bP/TDbwHgH/7hH/jkJz/J2NgY//W//lfe/H1v5vSpU0jfZWloq6+EEPR6A4cJ4N9//5v56q/+Gn7nd36XHdu38uY3v4n77nuA//n2d7C6sjqI4LqO4EkR4CHHSQ4pA7sPxb4Giwxp3xE0M/B82+bjJLMdj3LCq0qRpZqykGQIskTxla/4Svbt3s0f/skfkyYJu7ZtRYQBhw5ZMdI4TvGDwGactDtkmSLqJ6wsLTvSqH0mSisuNr58MUJnrdfr2bqM7j3frFxGbhv5L5caRPPfhh2yPMNoOLzxRfuXtc3CbjBAgjZDFvP1n45zm2de5sfY+Hcdl0tYCocQoKVA54W0Ddi6qcJmKjsHSPq+C+NZB8Q4p6e4JuxkNq8G4QlROFA4FD8zyhXsdTGV3GHapIfIr962W+1qkVaZmpqiXC7T7/dpra3RabfRSlEKSwSByyj1JEEppN6oUyqXiJOESqdDv9slyyd6WpOlKX2X8S0lhF6I5z3zqOxzymkCCmG39HJht7xNSVGUeJDasHBmFiOgHFqtG9IuWoLWKf0sQ5Di+9bjRduSFB4aaSS9qE80H/Pf/ttvEAYB8/PzPPLoI9z2ydvoRd3iwFYdGjzpc/r0LD/8I/+BI0eOMHv6NGmS4AEnTxxnbmHexQGxTpP7Zx2WQRM0ysVIwgDSlChVSI1L14dAQlar8NKX3ExYqjCzZYa11TZ//b6/5rprrmPf3isACIOQJLUeeKYzrrv2GsJySLVaJmcpW7GzC2dLQkiuvvoavvu7vou/fM97i/Bmbrt37+b1r/82/uPP/ieuuuoa/sOP7OYd7/h9Hn30MaIoGXQCEqcPgpspCcqVMtu3b2f7thnmFxY4d36OyanJYt+e9K2j7DoDrXL4xZ7DRz76EY4eO8rk5CRnzp7lyKHDfMVXfgUz0zMWLcI6gGma2th+zj/TmjAMyTJb5DJTmr1793DDi67nE5/4BO9//984HpTNLsvpS1IIfN8Sh9TlOE7YUJyQ0iJtxhQVvDd2tlk2UKQ3GcRkVALfwu15aQW7IalR1Eplfv7nf543ff/388C99/GC/Vcy2RgB3+PkmVmMKR5t0SEf2L8fpc0QWVkWhMqLnD5C8EWvydmwA7SRCHsxe7J1Nu5n4yD5VPf3RXv2bfj5b3SQLokSXcSGM++Gl+V/12VeOmfF9mYCzxeIvGh3DskLgfA9AjybwZbLXggnzysMxskISPIC8IPyY8OIUx6yz38pUCkpi/7KxvcHsjl2u/x3G0oMgoBavc7k1CTVSpVWu01YsqKUmcoopSW8wLPZ6S7CEJbLjE+M40mPKIrodrukSQJGkKQJvV6XbquLyjJKoc2yM89CPcznnNOUJEnx0C97tiWsOGTelKUBk2XgS0xem0bHJFFShEsMVpRSOL4JQObKqdgfBcdOHnYhF4Nowdnzs85btwOpVU+2pLskTbj33nvIsoy1tSbCaELf5+zZM8R5+Q8JMpBW1NLSSIoG7OGR5gqJ2i41xhZ4yRP+VuM+f/fRjzBx190EXkAQ+Bw5cphOt8vv/fZvM7VlGokkUyl+4JGmmv/8C79AP+rhSY+/+7u/Z2V5xRGU3YsoKMiC2hiiOOauu+5itdkkSazj6ns2WyNNFQ89+DA/8zP/EYHhfe99H+12h1OnTg9mTcKGu4wxaGULDyNgfn6e97///WiluebaF/Arv/Jr3P7pz/Kbv/GbrK6s2vuvNAivSOuWQjiUx0JNUT8mimKuvPJqXve613P48GFOnDy+vvaTdsWOhX3KXuCjEkWSJvi+Lf/ie5Kx0REOHLiC2dmruOmmm9ixYztKGT760Y/TXFuzMglao7V4cqQpb4Yuq8UY66igbduYnJripS99KXt27eY973kPL3v5yzh9+jSLi0vceMONRHGP8clJPvyhD6HSFHcbCwDL8ySxyvjHD3+YRw8e5PzcHCePH6cfx2Ra0e52EMJ2VMZYqYMsyzh79hzS87j22msBXMgyd5ouFs74oteUW66jtFGkb7MsqNwuB424mOM0vP3F9nMpu9j+v2iXts3u18YQ3UZ7Os9neNvL/b3QEHSfjdUCH/B1h1BL3/PwpIeXO105D8ihTI6Z5H5z3836ZcaRRI3J27rt14rEhHzgMLI49jpHz+0rCENK5RLVWo1qtUqiMjzfR2NIs4w0S/FS307ulSBNy3jCo1Fr0Gg00FoT9WxlDwwWqWqtsRbY7OlarYoUkjR+nofn8sHyKdlQS3A+P2kcW9gvMFxz9QH+48/+BFACIxHaedBSg4MttTFoI0iimHZzjV67S5Zm9KIerXab8/PnOX36FONj48wtzg0dy5rSlpyXL00dypNmGb04dmERILO1g4Y3zgDrtmWDxXE6iOXkZiDuJxw9fBwpTxWx5yRJ0Frz6MGDeIcPg3uZjNMzuPOOO4sXTwhBHNtsBSvMuR52FgK0VpyePcPc+fNorfF9l92nNJ4nmZub4x//8YNobbj99tvp9SJX1sQU+knGuBdN6ILAt7K6zKc+9Sl0ppmcnOLRRx/l7NmzfPKTn6Td6uDhgwtX2LR9z7m2OTxn+Udra20OHz6CEJKV1WWOHzvGH/zB77PabA61B7H+/kmxTh/n3nvv5R1vfwfHjx9jYX6eO++8i7Gxcc6dm8Og3X2w16H0Jvu7iA0LBg67HkopmqurbJ2e4fu///t57b/9Fj7wgQ9w7733s3P3bhYW5/n5//Sf+OhHP2odVTnQespNo7nvoQd54NFHUJkVzLSnZUgz5cJ/xqKfLgvLEvrXZ24NbHiwH1r+xcF2nW2sRv90bLNQz2b73JiV90X7f2MbkaXNfh/+7ekSxYf7pM32P7TEJmxIgUoUmYRMKTtpxyYU+Z5fyBAIM8iiE9qiSzJHrcidJVEMnRv/mfw3QVGiSbiSLBsz/GCTUCI2ypDzTD3PI02SAhkTMk9AUaSp5Tgl/YQsVWAkgV+yKL/w7cTT84irMaWgROiF+J7Pli1bKIUhaZzwiY998Cnf+0vZc8pputz0xvUbYZUmdV5HSYDnsbi4wMnTx7j2+hfwYz/6k3YQNhQNx8aMik8ApElGp9mhs7ZGHMfOu20xe2aWo8eOMDW9haNHDxOnqeXw+HnqpyWT55yHIPAJPCu4o5WyWV7CavQoo1FSo4xBK42vJUJb2MnzA7IixmgH6tzzh4EukFYalSnSJCGKIrIsJY4iFhcXWVhYICiFtFot62htEDfLLa8HtG5W41A4rbOCT6bU8Ats/3a7PQBWV5tDMgID00oXodL85sZxwuLCEhhotdr8xbvfTb8f0ev1nNaIcc8Q6/DhA9o+NwfTpVnG/Pk51porPPjAfZbkHPU5M3uGZmvNcpA2tKGi5IkQqMwKYz7++EFOnDhFmiZkWcaxE6fxPI84joijyDpuOnfBNdLlhDxZq/Scw4fJOxt7f1utFg8+9BBLy8u89pu+mdlTszRX14j6EU8cfIK11hprzTWGu64C0DQG40TBoiRe5+p4vlfMCD1PFqreGMH58+fp9rrUavViIjJ4Vpt18KJwwr5o1jYiTE9mF7t3G0Nww2GezRCNLz6D/ze22bPejBh+8fDc5aO06/SYzKB32YxXlU+6Mw2ZMaQin2jn00nhqg9YdFsJF+/SxobkhJMfQBQyN0VPIwYOlZ2IY8NxDt8ahO3yrfKw4Hq0NV8HcHIDil6vz9paizi2obU4ThnOFi0mJNqNaakijVLifoInZSFKHPgevufjez6BZ/mak2MTVKtVsuc7ETzLsqcHLRdt1T5IPwxpdVo89OADyEDy6n9zHUrB2dk54jhlcnySkdE60pcYITDWpbahtxmFSlNHEtcondGP7OAehiHdbsfyYtKYleYKUT8CDEEYUi6X8QOPXrdHEsUIDaHvUalUKVUrCM+lVvqlApUQZoBiCCEvY2gevFBpktJpd+j1uqyurnLk8CGeeOIgIyMNzpw5y+pq0xLLhUWL7ABLESPP6xnBQBTSuJVyJENrCAK/0NmwoVNL7r7Q8XLCju549poGxzVmkF0mpSVrd7tdKqWKrb/VatGPeiitSTNFliQuszEDl+mW6ow0TWi328X96DsnLr836zouR0ZSOis6n34/ot8f1Cwa/nzhzbYoz+WYHWTdZkMSBXl234kTJ/ir9/81//ShD9JcbdJxZTq00vzyL/0KiUPs0GZd/2uMthQCbYoMO6PBpHmNKivBkKY2Y2XHzp0cPHiQw4cPc9NNL14PrQv38Ne1s9xRG8LqvzhuD9K4uZBrMmwb+6uL9WEbM7I2226z/T+TWWZftMuzzdCjSyGEn48kQdF7boZs5ccceiGlNyitIqUYTOqULamFtggT2NfZFxKjJFKwjtMkHf+JgstkkSouuJZ8ed49DAf6bBwD8gmBXWqpHV2Wl5cJgpAoilwJIvC90KqF+z6+H7plAUYZon5ES7QwRqOyjDAIqVYtItVtd2mvtdGJIokSAi+wRdOfYXtOOU15qu1TcpwMFoKUAqM0YIiTiAMHrmL/VVchPKs6ncQxt912K0tLTV72slfwwhe9gFqjhpKWM2Qct0eGHvgenoFQWMBnxB0HDMJl3t35udu59557OHNmliRJmZiYZO++vWijuPeeezkze5aSHzA9McHOnTuZ3DJFUC4xMjnGrt17qY80KJXKlIKQkhcghY/Bok/SH3B6hLCOVe6NC2lPyobAQG3JSBKrfXT11Vfx5V/x5QRBQBTFxP2+IyILiksQeTMfdOByiOg3GBwEaIX0fKQUeJ6PcZBrDisLK8AxeBC5bTYwDzll2hh8acvanD0zx9aZrSwvL/Pgww9z5twZev3I1eFq0lxrEfe7aOVeDgGFJoEowJjCpLTnprW9V0ZrkBLfaS0NZz/lpOcifLpp8zLFLGxwHLv+oI0OPJwi00oNakDlUDcIkighjmK0UoR+YJ1zqTl65AhjY2MEgU+WZqSZrZvoeR5xkjin0abhKq3c+Zhc4aEosBv6Pm/83u+l1WqxtrbK1q3TQ0KJDrkzeujZ5OduGGoqT2XS/AVs+ex5c3ThgrWH+q3hsN5m6MTGAXhjiONyHKXNBvVnXE/oIm3gchTEn/bO1x3nMrfdOAe4rL1fZM9DKODwsuHf8s/rUfgL+4OLzUKMwYrzDsfzN16LHJrhom11gFJAqRzihSF+GNhKF1qTxjFxP7aZyig8ARJZTIa10UWBeoywRevzPhkLHtgwnpXB0cKW9kLIwQTYDN0TYwYyBvlkyzh+rPDAZCRxQrvdwZOSKI7odXt2/JDSoV8+gR8ihKBcKlntQA0qzbMKbcHzSrmC53l0Oz20MmRphs4MWaIK3u0zac8pp2lqaor5+XnbEBw0sh72zKESBuiMAemKl0ph1xWej1cqE2Uguwngo03EtdddS6fXZ3LbFKkHHZ2RYkhc+ExIK0boCUlgwFeGEEHoSYyyxG9fCIQwTM9s4boXXEOjXqfX7TE5Oc3uPXv4oz/+Iz7ykY/QbDbxhKRWrTI2OkqlWsULfMrVCuOTE5QrFQI/IPB8fOnhCen0c7SrTO2UU8EN7LagcI4UwfoQkEBYsbO88W6mKJy/f84R2NhhB75HpVymVq3RaDQIHIxaqVbYtnWGLFPMnZ8jCAIajRG+7uu+jpOzZ5ienqZeryMcdygHNETOAicvLzIIT+S125qrLRqNBlHc54U3voC1tRaZUvR7Pavt0e0RxzFC5Bo5BhhoIeWzG6utZBWvBZbEbUvXKKRn72+WZu7e5DdxUPZ4XedWzLwsx8uT3sBpMrbRrfMJhzq8PI18Y90w4RxJgdUa6Xd7NJtNlpdXWVpe5AUvfAGd9hr7rtjP6VMnWFxcYMf2HYyMjXLi6DHSLKHf7XHmzBkWF+dJ0sTVtrPn1Y9jzs3NEfg+r/2W1/Knf/qnnDx5AjCFlpOF1lXxef3f/EIGDtYFI9bzzIkSQrrCx7YvGi6ECps7L0/m9AwPvht1evLtngzJeirLNzoAm09ILwffvnCL4T7k8ra59PfNt8knLM9c47v0vXhyGw6rbkTah47i/smh70PojMnLntixyxgx6NvzDwJMXvBU+PiBpNKoUWrU8EvWcZK+R6aVjXh0u+hOl6QfoTNFgC3BJDW2qL3ecIquT9LGuGxq8HIUSdq+wAgDwgUDhesr875tnQPl+uV88o1BKU3U74MxRFGfKIpsoovvg7FAR66ZVwpL+H5gUTPtxnUpCf2QSqmC7/n0yj3CoITvh0XZKPOkac1P3Z5TTpMtaOkP8UnAjYzFw5bOUdBOcdTGaF3n4PbjBSGlygg7d1/B2Og4a2tr/N7vvR0pBGGlwt0P3E+pUWNkaor6xDhho4ZfKuGHJWq1OmP1BjXhM16uUPYD264dDJr3Z1fsv5Jdu/Zy9NAhzs6e47prr2dsYoLf+72302p1MEaQak2z3aY5FEraaJ5rZHlqvxQCZQxhGLiQWN4oBgM9Q3fnAnOrWJRlfad+qQ5CCgj9gFIYEIYlatUqpTDE9zyCUsj01CSZUswvLhIEPtVqlfvvv58z56xsQK1WAyEKnaKcaK4dMVm4Rm5J+AKVpm6mI4oXSOTijLlPlMO9rgkIBk6jwb6oRcKsEGRKIfGKMKcQNv6eQ85GazzhEbjrEkJQq9UIfB/pBAY936dcqVCtVghLJfzAt7M1DbVKlVqtTr1Rw/cDMNZJq1QrtkaSHLSPJLZ8M9/3KZVLRP2IJEkZHxvDYIiiiHarTbPZpN1a48CBAzRXl9m9dx/nz82y1mwyNTVFrVZl7uw5MpUR9fosLsyzurpiy7sgMVKCkCRpxvz8PEIIDhw4wMTEODe/+CauPHCAfr8/eNBPNkgYMei4nePuJrrPOxtM2CQIDUaxmaOU22Yk74s5QZv9y1HeS5GQLxyo1x/7Yryc4XUuQL027UsGHMr1qKMo2sRTcrVMvv5TR6iMuUhvJ4Y/igt2vbHHvHC/Fz//jY5w/i9/PsX91oWoyGC/rBuy3D7Wn7SdDNuzy5OR8h+MACWsTpLByuKUqmXqIw3GJieojDRQaJQnML7AK4Vo16cszp1n+fwC/dU+yggCWUEKiTQGoTU4wUopfISQrnanJFMKo2x/GgQB0vPQOnNJJwYlvGLsxYDIJ/UFYq9RWe4AglEKozVxv0+mMuKoj0pTJ3tgJ83SrBfpFMags4w4y4oEpziKSJMEISCNYsLAp1QKbf9vT+Siz/Dp2nPKaVpbW0NKD6VScrKZ9LyiM/C8gf5NbgYn+sWgkSZxQq1a40te8iW0Wh3+z//+Q37rN/8HmVaUq1WE71MfHWFy6zYmtk4zMj5GUC7jl0qMNEYZb4xQET5TtRqNILSZBsLWxCmFgUvNj/iKV76SLEk5deIkWyZnqFZqBF6AJz1UlgsKXPqhCiDwfDzPoNLUIjBC4EmJUprrX3Q9tWqNw0eOsrK8csl9gXWWPOegWBkGq7qah9YuZtKhI91+n063R3N1FYTAkx7aaI54Ag0kmbbQr4B773/IFZC08fC8Ey1QL7BOi5QFjJxfdZaltqPLISfn6NlYk+tPnNNnnzngMttshh3rsss8TzpH2t3VAhpab570CPwA3w/wPI9arWbL7rjirH7gU6mWqdVqVKoVSuUSRtsw2EijwejICKNjY4Rh6PSfAmq1OtVqxQpzYs8vjmI6nQ5hKWTrzAxZlnHq1CwzM9MgBC+6/kU06nXm5+fZtnUr7bUW3W6XwwefKMrszM+dJ00TfOlZna1KhR3bd7Jt6zaUzmh1uqy2WrTaXUBw1VVXsXXbNh544AH6/T4vfOH1bNu2lUceeeQyZ9S2E7KlYIRzPAcdu20jzx+4SQiram+EFezbiCJdzHHaaJdCpIa/Dw/Kw9tt3OZSz3Kz0NLG423GyVm/LhSj98YurNjNxTyZS9lGZbjLM6s6ssn1bHpem37F7uLJj77RyR12mDYLoQopEUOZ0/l/RYEYDZ3vUOxbCBBGgrAoj1FO000KjOcKe2uDluBXAirjI0zMbKHeaCB9n14a0c9iEqWp+D7VRo1Rf4Qk0PR1TBR3SXoJsYrxReAiJKCMtt2qdMQN6SFcORaDrRAgPB/he5DYem/KWOEbC2rIwfkPiWNabmae0efujzGkma0Dq9IUTwo86TtpBNuwtLIcLJVmqCwlSxO0UvR6NsLQabfo9ToEvpVTyTJFlqW02y1KpdLlCWA/RXtOOU02DJMXMRx0IrnQn3SFCrMsKwZMU3B8tJscG7rtNlnUp1StceTQ/fz33/hN4jimXCqj4gTdj1lu91idX0Y+7pCRIMAPQusEGEFoJLUgQCiNzhQT4xPM7N7KA488iDKGbq/L7/7Ob/Oar3kV5VKZ2VOn2Dqz1RHsrPeOGSrBMfwWD727Smt8aQh8D5267AJhxbyUMrz61a9i565d/Nmf/hkry0tPfhONcSmemQMMJEJItDZMTExy7bXXsrq6wsjICMeOHaVerxOGIQvzC2it2L59B5Pj46wsLrLaarFt+3bOnz/P8sI823ftYuvWbRw6+Dj1eo0DV13NQw8/TGOkwc6dO2k21zh9+jSVSplSqcxac412p+Wey2AmqPWgw63VqnS7Paanp4njiLVWq5jNFEhZZjWrCtAxR76HbDiyLTZU+F6/XmrbmYPNV5urAI7Xlc/EKdJsbZjTYPSgGKXneUghC95UPtDlkHEeGk3TlFIYcvMtN3PVVVfx6U/fTqlUYnW5yQ//8I9wyy03c/ToMcIgoFap0u60ydIMpVMwkCQxcRxR8gNqVRvXT9MUlWVEUcTc/HlOnZll9tx5Ot0eBw5cwVVXXc3c+fOcPHmSNE2Zm5uj1+ttei82mg2vyoL3UAxyxnb6pVLIy172cm677ZOXtb/nuvmeLdmTZhlgCAKPjRIEwyjRZoNtnkCRL8vb0GZO0sZBeTO7FJ/pctdZ50BdHIdZ9+zththQzEXP7nJs2Hm4lD2dowx7eOufxVMNww07nxvRuWHCs8O83dFyZHYQ8jfC7Qf3Njl0CWOdlDzsq9FWnk9gQ2JS4AcBMgyojtapT4xSadTpRj06nTaRylhLerRVhF+tMDY1zpbpLZQbVaZ3bkUqRevcIt21LlIqPL+ML60wpkWbjEWXRGYL1ItBNCDOMlAKjbZUBuMNAot55KNAnQYhOozlx0onlKedM2U8D5yGHAg7SXVim0ZrFAaVpfafyhAMSgtJV68zVgMJg25X0u/3KZfLTyOw/OT2nHKaBnFgCk9da2VDOgKL3ogh2FqA8K33mya6EOeaO3uWqNOjXG2Q9SPGx8fZvnMnu2ZmHCktRWmD8CRKWyZ+uVKhVK7Q6/dAacoyZLTWIO716fQ63HjTTXz9676Zt/7/fppOr8vCyjLNTpup7Vt50U03MD+/CL4k0RkqL7oGg3f4Es82yzJb1kLYys2ZVmQu+2psfJzp6S2UyqXBBhv7nOEQlhQkhSK2VWNQmULplKuvupZf+C+/yP/9p39g+/btfPQjH+HGG25gx/bt3HbbJ5mdPc1XfPkreNmXvoxHH3qYO++6k9e+9lu4/bOf5hOf+Dhf9cov4+u+7t/w8z//8+zdt5f/9t9+nZ/8ybcyNj7Om978fRw/foK/et9fcc01VzE2NsZnP3sHjz32uA1jSRt2NAZ8KShXauzffwXVaoU777yLV7/6azh95jSfuf1zVMplduzYyamTp5mYmKC5tsZIowFCkCYJ5XKJTGUsLi5Sq9bYs2c3q80mpVKJ5mqTdqdzgZL5+ltnyY6AE7Z131xlkSAIGB8bJQxDZmfPWG7U5RTs3cTq9Rpbtkzxmte8mn6/h5Q+C+cX2bt3L1dffTUG62T1+hG+H+L5IVpnGK0plcs0GMMTEHiSIAzRLkEgCHx27t7NTbfcAp5PlMS01tbYvXsPx44fZ35+ni1TUyilmJ+fv6xBw/M8pOeRZZnrrBzZUwgq5Qo33/xifuRH/sPzxmnKRVozR75P03zwXa+nNOyMDDtE+fdhZ3rYQcp/Ay7YV77scojhn5c9WdzrSdd9qgcbcj6eZM11ZzTk+FzsXlwK178cR3T4+6U4aTkXyapf5xpoQxM6MTSUu7BkLlAphgYEYSzyU4RIhSh8SulJ/MAnrFWp1etUq1WEJ1ltrrK4uEBYrZDohF7UpbeywlqzSRbFzGyZZnxkFG8mg17CYqtLpFJCLwBh67wJYTDallJSqUEqz/GEbFSn79CesOTj+xI/j1wMTwCMKK4pv/C8zXuO+qCFQHue1Vxywptgk2VkXjDWtXml7GQ2DErF9rnTJIQsqoREUYRSqnCano334jnlNGk9qFE2SO/Nl7vQixPbyoloRhuU1Fb9O7M6SY889DCnX/qloDQiTXnLG9/EgeuuZXx0FOlJlDJoIWwj0RqdZvhBiFcKyLIUozSB8SiHFVJl46tjk+NM7JzmV3711zh+5iS//hu/wUJzhbVWhy07ttFLUys6JsB4luh2uSQ1geVm4Rqd7/ukub6TMg6ZMeTZoPk/M3j/3L2y9y1LBwRTIWXBR5kYn+BFL7qeP3/3u5ibO8cVV1zBtddey3XXXsPExDgf+chHWV5dRUjBgauu4jOf+xxBKeSKK65gz5497N6zh207trOwuIQfhFx/402UyiUee/wgwgj27NrFzh07efWrXoPWGQ/c/8BghuKI2gBB6HPVVfv5T//p5zl16hT3338f3/zN38hdd9/FnXfczYtedANvfNOb+KX/+iu87nXfyqdvv52bX3wzAHPnz7F37x6aqyv8zd/8Dfv27+NNb3ojd999Fy960Q188pOf4taP33rRez3otgcd5UD2zXZw4+NjvOY1r2Hr1q287W2/NYjbY4rOEUGRxLdRLLfwlw34fsDk5CRbt24FoFKp8O++67u55SW38Pjjj/Oe9/4VC4sLhJ7ndL8cauo600xl6CxDCBgdHUVnitGREfbu2Uu5VmZm+3aue+GLyFaW+Ku/ej+vf/3r2LF9O7VqlVqtRrfT4djRo5elrq+K0i159qGF7avVKjfc8CL+w4/8CHfffceT7ucLxaz8Rd5Z9+n1uiiVFQ5Trss2nDE3jDoKIWg0GoyMjBTr+b5/AcqUO0tJkqDyEjpDSNZmDsPlZMptzPTKt79UeG4zYvqTOSv/quxCoOmpbb7JPcs/X0D4NhduS+E4ieJ7bjljQIDlGbmi5hgX+pU2i80G+6zzEfpO808b4igi7kcIBNOTkwSVMu2ox+y5s7TbXVZOnaOSQjgxSTWsMDo2RtTukXR69FWGwhBIH89Fc3Tm2rAxSM9H4GE8QRorMpUSCMv1lA4xtfpLNhw2XCNxmCuXvxfD9yx3pHJbv23ukFn+kkAUzlueFJW/M1rrdWg/UBRWfybtOeU0bZwr5HDecKdETnAuGqgDQn0Pk9nOpt1qIzzJvffey/333suP/PhPUJsYZeHcOaTnMzE1hV8qF8cYPNyhODXCIg/CEHf6fOr2T/O+f/x7fumXfpH5Tot3v/e9fOTDH+GGG27gy17+CjpRl1RleK4Aock2iSFtEqKzXrjViNJuNlqpVEjTjExnBUEu3zx3mHCfhxFfyAdyUzRGz5GwM5UQRX3iOOKlL3kpC4vzvOxlL+PI4SM8fvAJbrzxBiYnJzl06BCTU1OM7Bvh33z9v2Hv3r0cuHI/11xzDQeuvJIkTak3Gpw8eYo07hNHMWfPnefosSNMTk5SLpfYtm073W7bpZbmyG2u5mF1n2a2TLH/iiu48YYb+LVf/XX27dvHY48/xsTkGN/9Pd/DN33TN/M7v/W7vPZbvgUDfMVXfAWrK6s8+uij3HTTDSwuLvCpT32KrVu3MX9+njiK+Yav/3oWFha49eO34vkSlV3oKAyco8H3XNohd+pqtRo33ngD+/fv521vs4iD7wtXNFJd2CEX1zg4gJR5oV/78ne7XWbPnGH7th286EXX8+CDD/K233obhw8d46YX38gV11yN1gqVKYeqpqwsLzN/fp6zZ8+wtLJCmiRoY2jU60yMjhKEAeOTkxw4sJ/l5iof+MAHOXPmDL/7e79DvV5nZGSEtbU1HnvssctymoyyldKFEBhlis7pxhtv4Du+49s5+MRB/vf//sMn3c8Xis3OzlKv14G8nxgQwTeSuIeXDTsXSZLQarUKYms+ix4O52qtKZVKLC8vs7q6SrfbvWCguRy7GBH8ybZ5di1HZgbn9K/N97qYI3mxsN5652mAPtolrhPIvSMLxRRoUxHeFE7iRgjA9tFIm+avjbYIbw4KJCm9TofMaKJ+hC99GrUGo6MjjCV1eitrRCstuu0VWiKgRkC9Yd9/NZ2ybBbpt7pk2lCRktDzbKgOAZkLjxkFeBgpMb6wSJIcOI8qy0jTtNBS3Mxhyp9vvnz4PuZgwDAoIqW0ITYDQRBisAXHwXKa88nDsKM0/N7EUUycbC7e/PnYc8ppGk7Xzm+YHHoYwqWKm6FigbZhCoQyGAUyEDRGGsxs387tt9/Ox267lbf+4i/Saa6ytLwEvke5XqEsDFmSEscJ/X6EShJ8X1CqlinX6wReGaGsHP3qyhInTh7n8UMHOXzqFPc++BD9Tp/DjzzCh//xg3zJLbewtraKQFOvVvGRJMMaHPl5bmLFID7U4YWhTam075wp0vWl8NyLt8k+hvZvG6e9n0ZrjLAe++yZWW775Ce54op97N9/BZ4n+dwdn2P29Cxzc6/mhS98AXv27qXb67HabLLviv3EaYoyim3bdtDt9zl65Ajbtm5ltbnGqdOniKIYjOHQ4cNUq1WWV5ZYXJxjfGyM8fEJq1arbdq7J0AbgfRs5tmhQ4d57WtfRxiWSJOMteYae3bv4ZVf/mXoTLF163amt0wyMTGJJyVpkjjFbsm+ffv5yq/8akZHGxw+dJix8XFGGg284bYiLM889xc8x0/aSGYWwjqeRuRp1FZJXCnbQVQqJZI0QwhDlg3xVVw71EOpvMMoU947CiHd8zFMTEwgpOHdf/lu7rvvHn78x97Kd33Xd7Pvit0XtI2VhUWOnzjBkSOHOfj4QRYWFshUWsz4tLZOztz8PMsrK2it+NwddzI6MkKpXGJkZITV1VUefviRy6/jyGBQyJGRXbt20uv1+OVf/tXL3scXgp04cYKpqSnCMGRkZISpqQl83y/6p40lT4YHg7wv6/WstEQcxxeUiBpGmcrlMktLSywvL9PpdKhUKoSu2PhmKNGTIUaXQqiGv2+0S6FMG397qlagtP8SVoBBT98pXD+hvpAfVkyUHCleD4SQXKcymNAKN8EXQ+dljLFlQnB9kyfJyLPcNCZVZHFCH0On0yXOEpI0pVIuo+IUEkWoBHV8KgqiTp/u4iptv0xJ2oQWb8sUcRzT60dWkVsrhPCQ0kOGAinBZJrMaLTOkF6AHwaYwEP6vuVbpSlRHJOmqSut5RfO/kY0KV8+rPgNF7YXIVyWnpT4vs3YNkagM21FZTXFvhKn+j3MczLG0O/3Li1M/DTtOeU0CQSetESyWGWELrTR7nbo9XpuFi7WdwDaQKKKjOjA9zh24gRnz8wipGCltcbhQwdJ4h6eH5KkMYdPHKPT6dJsrrKyvMLi/AJxHNMYrTM+Oc7E2DiToxOUSjUkgjD0+bbv/DZ+8Iffwm/+3tv5H7/127RWVzHGEAhJr7lG1usRGEOtVCLwQ6wQwkWgw6H3WBmDUNrqYxhDr9ejUqkUse8sy2w4Zt8+4iQhDIJ8zlLsbLg5CgG+L4ljhZQ2k3Ct1XKFhJf5oR/6wU0zDu666+5NnsfFOQJSCP7kT/6U8+fPY4zhr9771xiTMT4xwSduvY3rrrvOhhi17baU6zg0kKkMZQz1eoOHH36QJIl59NHHmDu/wJnZs3z4Qx/m+utfxK6duzh18hS3f/p21laaSE8yN7fA/PklvuIrr+PVr34Nq6tLvPWtP47nC9rtFqdnT9n7mrnYu01QsS9paGFeS6Y2AxV014HlCrtwIVclSdIC1QwCH8RAJmIoKcZlDVotJoOx5HIBSZqCkGzdvo3FpSXa7Q7/7t99N7/0y7/I/PwC9913PwhjFeLdMQPfpzHS4GWveBlf+2++jkZjxNHXjdVocfF+YzRHjjzBt73+9Zw8dYZur8vC/Bw7tm9ndbXJ0mVkXdp7Zd89G9a1cgxpmvCe97yP97znfZe1jy8k279/H2Nj43ieR7lcolqtrivgu5kjujHLqlSy28VxfNFMH2MMlUqFKIrYsWMHCwsLpGl6ybDZZsfN9/Vk62/8beN6T7avp0Os/hezotMaIFxP1TY6SZuhikNruwmS9ZByFMmSTIfQZ+Mmuw6A0pkhJRsgL96A0lRUclcGnSlSA7FK6UU9u/9Mcf7MWaKVJiXhQyeipjziTKLW+rS9ZUI/YMyfolStUp+YoBMl9Fod0lSBzjAGSl5ow29IlHYCmCj8ko8vJJ6BLI5J0rRw+od5ecOTh0vx9fJ3YeMkw/d9wrDEyMgoExOTCCHtZFANqAlxHBP1bX3TNLNIV5rExEmC7/nU6pWn95AvYc8pp2lkZASbGp+iNdx40418/dd/Pe997/s4dPjQJYQDLYIAhl4/5lOfu4MTZ88wtnUrDz9xmBtvunnwIogNs61hDchCJ8jNItzAunV6nO/63jfxwz/64/zCz/4c1990C7NKsLLYp9Nc467Pfpa77rqD666+FpWkWD0hp0G0Dm0aCg7lyASDMdcTEoXV20gzq9GxsrLM1Vdfyx/90R+Rkw6L3W3GMnfTOeMg/zNnznDvvffQbrdZXFzkfe/7a86dn7czMBd7F9iw3vA9tcJjksyFPPOKK8Y5IJ4neNe7/oIwlExPT1jUQ0GSaD74wQ/xgQ98ECEkk5PjhZhkjvAIIbj33gf4ru/6LrevgJ/5mZ+lVPaRnuRXf/VXUZmmVKryd3//DwRBmbvuup8sSwHDhz/yQbQe1Bz62Z/92XX3xQ8EUgqUyrM/1s+uBcJW/pYSzMCxE2I9zK6dcrZSinq9RhwnxFFiRTJdu7F8RtuuLGHRzhKFExq190tSqza46qqrefHNL+Y9f/kePvWpT/Ld3/3dxHHCBz7wf/npn/6PGKMK6YdarUa1UqFWr7Jz504OHDjAS26+hcD3AMOevXuZ2bqV8YlxqtUq27dv40Mf+hAveMENfPYzn2XHjp3cdNNNrK1dXCNso0mn1JumKaOjo3zJl7yEpaVlHnnkUbIsxfc92+k+T+yKK66gVqtdMFgOk7cvhbgIISiXy4yOjhZ8pXx5bvl+c6dp27ZtnDx5ktXV1XXq9ZfDYRq2y+U8/cvYcFjg2bZLTfUubZcKs26+Xo4om2JCZUTxnyHYmWL9QWSBIgSHtHxc4zSU0FZjDmUg1WTKkKYJaZQgpSTVhuX+IpG/ykhQxhceI14J7ZfpRH3aSyvgecgwYCwMqNbrjG+ZQngevbU2ST+xmkzCI5QB0peoLEOhQWf40rfyKalNSMkcwiSlLQQ+zDHKuU55W80RpjyctrGN5Sh5EAROsiWkUqlQq9WKkN8wZylNU3q9Hv2+FceMY1sCK4oiwlLoogfPrD2nnKZmcxWl7Yz/Dd/5nfzWb/82SRzxsY9/DA5ToAG+71vZAWMxFw9hU/cD26mXKwHlWoWOE/Xbc8UOTp86y4Er91Gt1Tg5e5pypczE1BSZ0cRpShD6lKslOygpTdZLOHviLKvLa8wvrfJ7v/t7/K//9QdonfKm73oj7333X3Lv0jmyKKYWhhzYt5dGo8bMtq00RkZtBpfKQ3T2BZKeBOleysQ2Mt/pGeVq1vbN0njSVqd+5zv/N3/0R39s1b8d3wryF87eNyGsTIHvWWfB9322bdvGTS++kR07dvDCF7yQl7zkZqZnZvjxH/9xPD8glyHItx8OFeRxZBg08iIjQin8ICg6CrCDSKbUuu3s/vIna0NiOc8qh3iTJLH6Slrg+x6eL+h1u8yfn+fc2fMsLCyjleHBBx/m5MnTdLsdOh3r/M0vzLHWWsEUGKOypWYMZOnA+dFuamfh3Lg4N/uj5Q/l/awlJA7uaZ7hoZSi24swbiZptCnUdXMXy/43nyIKcuFtow1RL2Zlpcny0iqL88ucPXuGXq/Drbd+gve97738wA/8AG94wxsQwqJYcRIjkfiBFZTLz18Iz51TPn2VLC4u8Jd/9Ef84wf/iT9/15+jdcYv/eIv8uVf9gpe/iUvo9fu2OspztKFLF34Vkg3CdHGZlk6dLTdXuOhhx7hV37ll+h0uvz0T//MBeGlL3RLU6sxA4XfPcQnXM/ZuFiozBhT1B5cn+CyngOSr5sPRJsN3MPfN/tto21cPz/upcJ6l0MEf7rOlngWhAifTdt4zzby14r11m2ElQwYolGIIUaFBqRznoS0XFYJCJ2Tsi1CjScQ2mAyZdFyZfCRBEGI50kyV4rJoKjXqzSCEp4yJL0+K6ttIp2hPNCepDExzsTUpOX9GkMnXSPq20mnCKHslyzapJSjJSi09DDOIQIolUoEQUCpZLO4c37TMHqUj8lBENhrdW0ZKBwipVQRbsvfh+XlZdbW1jAGh+qWqVarLkQdUKvZpJZhBFcpRau1RrPZfAafuLXnlNOktOHGG67nR3/0x3jDv/su5s+f4y1veQv3338/ntPQGS6gKcBqPRhbiytzs+Con7LUbNp0TgylsMSVV+xFpRn9dodqECKEJI4iMgz9JKYbaUTHEEd9TKKYaEwwPjHG/n27mZycYG5unuXlJuWwwcTIKDlFs14tc+CKvTRGq4RhwNTUFKPjE8wvLKB1ivR9G5ZSCp2pCyZCmdL4zjOPswzf9+hHEV7g86M/8P3c8pKX0up0OXr0OAefOMSJY8dZWlyi2+kQx30uZvPLTR589CC+gEqlTLlaQno+YAdeMeShD5PzrGlXnsMKinrekDdv8smRtHXSMrWu0/CcLP5wpz/ovEXxN1+Wv2BxHDE1Nc7VV13Frl27qFbqjI9v4Yp9B3jd676VNE0Jw5Dx8XFqtTK+L8lL3x09ephv+qZvYGm5ecF9MMV/hpatC01ceO/ydhaGlhuWJVbCweQ727TvtwOKF1gFeVsrT1KpNvCDElGUMDU1je97zJ45Q5ZlNJvLLC7Ok6YRN9xwPZ4XWOTNH5rFkpfWgTCwtZnSodpMSZrS6bS54YYbkV6IMYIz5+bpdnucPHmS06dPFTvKKzLkLpfWGk96lEolm9I7VMdJa1haWuK//JdfXDezDoLnD9okRF5yZ7jNXKjHlP/dzBnJZ+Ibw3bD+9zoQG1E0i/m/Az/ttlgfjGH51II1JM5ahuv+6mYGQJfnl0bIPo5tvxkINdmz27YWdrcQbWTMws2DYfm3JFN3t0XZ1G80xKBJySezENXBjOUPS6MRGhsiE4Z0BpfeAQy71+tLECqIlK/RDko0ajU6NVHSLSmm2rWlpsYP0AEAfXxMeqNBmkUk8UpfaVJlCI0GaEIcsIpJlPWWRIZZFaDpVQuU6tWqVarrq5pRKvVugCBzR2knPOU96M5cpT39Tk3KpcQSNMMpWwmXE4YL5VKlEqlQnh4eB/5siRJXf2+Z9aeU04TwI/+6I/yPd/7Rh649y5+8id/igcefIiXf+nLOHn6FKdPn7bp4ULYdHzH7vEdqpGgEd4AEpWeJU6rLKPiB0Rxgo4VFQJUBlm7j/EkgYTMgC8DxqdG8BBk/ZS0EyENhJ6HyVKyKGJxpcepYyfInKc+P3eOBx+8j1K9jF/2KTdqSKfAbSs166KjkJ4AT9p3TWmM0/5RWtuCidJ20v3YoLOU+x94iJOzs2RK47nK0DfffAsT4+MEfoDWisCTTE1MMr1lilq1jOdL6vUGE9NbkFIwPzeHENoiOsZQqzfodHqcPTtnC+cGAUEQEEcRi0uLxFHC2NgY5XKFqalJ+v0+R44cod1u2xmwNoRhyPLyEuVymX6UsLK6SpzERFHE6soqa2ston5Ep2O3kZ5HHF3cwcvt6FG46867aTQaNOoNPK9EEISMNEaQ0ndOlo/nHCbPIWszM9O8+93vZmpqqhigpJT2hTJDM/t8AIQLBj6braGKEOjCwjz/8Pf/F8/z2L17J8dPnNrUAXN7cH8lOhuEZmu1Ggf2H2DXzl0sLy3TaXXYMjWDJzzC0MpdvOPt/4t/+sd/5MSJUwUC5Lm3Nq9RXKuViZO0mBTYcx5GPwT3P3A/3/kd34rWKQZIM8XcubPMnz+PAcqBR5wqi3QO9TNaKaJM0Rhp8B3f8e3s3LGD//7f/wdC2KzBJLFEy3I5JM0S0mehk/rXatKJfUIesoGNThOsd0g2Oin57HvjusN8qI08kctFcjZzlL5oz7xdOlTn/q1bOgQtDcNMQ9sIYast+E67CJ2Rx/ekkAjhIYwrMo52NVEFwtWRC/0AkSWkaUan1cVUbDLBxOgEIghY7LVp9WKWzi8QlEsEpZCwXGZkbAydWeXvqN0lQVMizxwHYRRkGcZIdJIigVq1xszMDKOjowAsLy/T7XaLdpdPkpMkuaAdK6WKyXGpVCrafT6ZSNMUpew0LkdkhxHtnPSdW76fcrmM53ko9TxXBAfo93r87V+9l3e+853c/9DDvOiFL2B0dKTQk4BB55T7+YXAdM5XEriyHS5Ek2k6/Q69tTblUpnRkVESlZHqDKU1nV4XPIFfq9CJW1TLZUQGtVIZk6acPzNLr9Ui9D0SFMuLi/TjHgZDpRSC0Bw7dog77/gsR48dpdPr2SrS9mRtZ+v4Uja0Y98020jzKzcYI8i0QngCieG+++8nTTMQglq9QbVaY3xsjLGREcIgxPcCpDDUqhVGGyOUwsCeU6VCY7QOSLqdFkIOitQGYYUoTmi3m2gXTpOeR5JktNsdjNFUK1WEkNTrdbIsZX5+fvBCuBDp2lqLcrnE6Ng47XaLcqXC5OQkAEJ6jO0bt7XrXPx7ZWWFLVNTSClpNpvs3LkLgVV2HR8fRwhDkvQBjecFqEzT70c0GqOF+rYdsCyPyA8kKks5eeokjzzyCLfddiu1Wp04jl2nZNE7VXCynEqt0z8aYpcVBHCt7cwqjiMOHz7EfffeSyn0XO02UcwX12fkiHWfzVAPGvg+E+P2PpyfmyNNU2rVCpVyiZe+5CXMzMzwwQ99mNkzZ9gyNUaSxASlCkIYPN8jTVK0NoyM1On3e0RR4siTAUpl9HoRICiXK6gs4zOf+SzGQL1SIYkjnnjiCc6cmQWc8rwnUEN1HXPHyGY3euzdu483vfGNRFHE29/+dgBmZmZI04xTp0/heYMCwc8Hs+GEYafHc87q5Ts1+cx5OKtOKVU4U8MO00axy437gs0RnosRs/+18Joude7PhD35fnNs9ZmxJw/RWUhtgEMJNjJvlNagU2Tgwk1aY5RGCMuj9Zxmkza6SP8VUhD1+vhSMjHSIChX0f2EtB/R7rTQWuG56gIdldDrpUT9mLXlVSq1Glu2TjMy0sBoQ5KkpHGC1gYlTFFmRRiDMBq0Io0TyqUS9Xqd7du3MzExQZIkxHHM4uJicf156C13iIIgwPf9gjju+37B7cvbfL/fL2gaOd81XzeneAxzmPNQeeakD3KnLHm+Sw4A/OW73835uTnOzM2zY9sM1UqFublzdFptFxkZdPrAEOoEoR+SZMl6th0wOjLC7MmTbN+6FaMM1VoN0evTa/XQwHjNat6stVu0Wm380YbVwxCSNWFI4h4jo3V2bNvJ4fZJPDG4saO1KlOjI5w40SPrdal4PmUvRCLW586ZXAdokPvmJha24LCw4UmJQfq2llu/HxXhvLWVFdZWVjh/9swFM5cLZr524eAeiI2/PvXMEuOcvLz/ybcfGRmh1+tRH6nTaDQcUdq+cOPj45RKIVmakcYxSRw7bhTkxXeFcAW1sW5Jjnz1dUIYhlSrlXUhvjxkEoQBKssYHR1ldvY0v/M7b3+a9M8LzZOCUikgDH08T1q19/zCnaO7vrsU6z9baXo836NWr6G1ZmFh0aWrW1X7l7/85WzbupWPfuxj+L7H2Ng4CwsLjnyfkSaqCGdEUULivhtjrIaXG3R93yMIPCrlEmHgs7jUZGpynLW1Jk8cOsT8/IKd8Wl7n4coXBjXJoUnGB0bYWxsjCuu2M+P/uiPAfC+972PXq9XqIPDphPnL1izfL5cT0Ig5eaOzEYbfh+HOUzD213Mydn4/XIcjafq+DyZE7NZiG/jupc65rPhHF0sXDn8faMT8/mexWahzWICts4s4m3WL7K0DGy/JtxJ5e8fxqC0JpMKT1rKhCc9BMYiTU7JWGgbkZBCIH2fJIrJUAikrcvqe+BJl0SuwBhSo6gGZZKSRsc9knaPtcVlwiBgbHLCToanMuIoImp36KYRjaCM73sgbdYcSpGr92qtabfbxT3OVbnze5Pfj2GZgTz7MwgCRkdHmZ6eZszV7EyShG63S6vVYm1tzXGcjJvE5Q7TevHXfEKhtXLil7qYfDzT9pxymgRwz30PoI1hamyMWqXMyRPHGRkds/XMeiFJmrkOf0hFVWtbW1IOenSVZag0Iww9JienOHr4KN/zpjfyqdv+mTRJUQbW5ucZHWnw6q95FSrN+PSnPsmBa3Zx8803c/DgQWZPnSIsBdTqZYIQGmMT3PfAo9SqFcp+YF8GYyhLyfToKNdfcw1JXOX2xn2cQK6rh2ZDcoBvz9MoUyhKK/e7wKXmK+OKKgLeYNuiJpuyo7YUksBzquZa4QtbDDHOYkqejxGGVBtKYQAI4iQBYflKUm4kA1tV1kHIxzbgImYsineoIC9qbJFlgNWVVZorTSfWJmiurqzrtQZOj1jX8VxyYDAbOuBNVis4INrVpkOAGWRv5C9VGIZ4gU9eXNhyhzw8TxLHEVmaEvi2gna1UqZUCun3+/T7yyi10VXf5ETt2WDRJrusUimzffs2xifGad5/P2HJ1jY0Bvbu3ku1WqXd7hH4kna76zTDYsIgIM0y0iwjLAWUS9CPErIsJQxzXoDVjRrORvGlrRHl+wGrq6tkaUaWWYhdGhuC9jbcQ60NE+PjvOzlL2N6yxSPPPwQL7z+et785u/n7rvu4Z5776IxUgcBz4L47r9ys8VJNwvJXa7lA++w4B8MZtH5+zAclhvm0gyvm+9vo9M1zInaDHW61G8b17vY75fjIF5qv5e77uUc+0nJ6WJYkuXyj3kprhlwwTPKf/M8D+HCbHmkAxyKnU8Ic4cph3gzBUqRZRoZeIRBCe1p0syW+Mo1WoTTLpK+R+j5mCAky1L6vYjYIf/10VFKYQgGVlZW6Pd7NMZGGfd9DIJukrC2tEqSZmRKs23XDqamtxBFPc5HPTrNFhXfo+SX8AIfTwPa4Hu202+1WiwtLRUFzpMkIYps2H4YQQ3DEGMM/X7fJfl4jI+PMzU1xc6dO6lWq4yOjjIxMUGn0+HUqVM8/PDDtNsdN5m2SS6D+X4eXbCyP6GTjMkjDjb098xP4Z5TThO4hgkYY8uI5GJyo6OjJGnGymqT1MUxDcbWbLNfiGML1RljiLo9m7GCtCiOELz227+NO+65l/5qk8QolNGMTUzwki95GfPn5rjrjru5+aVfyute9zqOHDzEww/fz8hYlfEtDR55/FGWmx28kodfK5N4llMVlAKq9Qrlsk+v2yaLU6SryyPcdGPdC61AK5e+WXGPx+TKsMrK6CtDuRKgMls12kKWkizJinCezgy+tKmiUkoC6bl6c1YqXwg7gJYCQZqkSM+jXK6QpgkCS7wLfEHmwjW+425k2aCWUubUqS23w0MpbdP4BXjCQ2qrJuuEr/F8STkMQVgl5DRVxYwvvwMSS4LMOxFjDJnzxnxpw0cK8KXnSt5ojFZoctja7U24Ts5oBFaDKUsVOu+cxIB0m8fWM6UsDF54f7Z6Ui4bEEWx1e9yPZ02VpTT9yG9rAmNcf/svaxUKuzetZOtM9P0ox5hGOB7Es/xGYTLztMaFheXqFbLxJEVsPN9H4NNbhBlQaVcIss8qtUKUko63S4qS4vstzRJKdVChMBqmqmUUliyz86dXTUISLK0eBaetCjtzMwMX/ZlX87i4iJve9vb+NVf+1UOPfE4q6urpJmm03EFf4cf5PPA8tDasGjf0yFDDzs5FxuQn6ptdLCG2/rlIEiX8/1iy/LjXmr55mHETTe55H42cxA3XvsF51Qsu/zns9myy3FEi++DH9hIadrMgiCwuoFxbDWQUsvlSbIUIyAIQ0p+hVJYwghBmiVE3Z5Flktl0sT1E4HP3gP7uWr/AaYmJzn8xCEOHjxIp9Mm1ZqyH5ColH4vJhKCzuoarZE61UadkfExsjRj1Qj6cUzWi6goQTUo40tJts5hNCiVkJO7c7rGcGhuPU9JFevmxPFOp1OE5XLHy2anDqIIm93b4fa9kYsq5VNTzb8ce045TUEg0UqjtEWKhDE0GnWiOMH3AzzPhj6EGFRjV9rGHYQQmMxCndZp6lpOiNF0owjjSe576CFOz53j0MFDKKWIo5jmWot777uP5mqTs4sL/NXf/h3/8KEPEccRUdzHCwRBxaNU8/nSV7yCSCsi35A4hHax3eb0+bOsttfo9LrFgGxdKksAERJLUDcCreyMc2x8lG3btrniqxLfD2m3WoTlEt2oT61eJ8lSlNbU61XCwKfTapElKVmaoTJNOSjRbXXRyirNCmP5LXnJBqUU7W6HMPAISyWEZ9VX46iP0SlCSDwJmbaFJz1PIkNZOEtaW7K1DbEppJD4YYDKUoIgIEkTAiHJtCYMfXzpFSKQQeDhCUtktDWWXFFKY5wOlkYD1XLZZgtK+zy10ATk6t2CSrVssxxVVmwrpVPtzuw5l8uhDV1lVoYh8ALL08qSdTP9IrbkYqJCCoy2iJ3v22raNvSlLclcCJJUDyD1TWxjd+s7eFlpgxS4/aZEvS5bJicsKmqsMnriiN0TY+Mkccq3fdu38sEPfYT5hfM06iMIIWiurbG21i54Mc1m24b4cI62tB1WGPiMjDRoNpt0ex3SOLEVxJVFMhvVMnv37uHkqVP0osgmxjj+Vb/fZ+7cObZt3cro6Ah/8Pu/z/vf/7ecO3fOtg+VWWQyV1R43jhOYh3Ks+6XS3gAFxvUh397ymfyJCjRk53XxRyPi32/3HO9mJO2WQgwDzf/i9pTdNQ2u+bLvfcmv29QaFmuOw83gxTunvmej9G2f6mUymijiVxpEKO0jZT4PpVyGdPXtNptatUKfmA5jWmaojGUa1X2X3M1L3nxzezZu5dKrcqD9z9As9nE93wSzyLXJIpeu8PSwiJjRlFrNJia3oJnYP70LFG3g5ABVT/ADwJkkiOtUCqF5ITuXOgy5zDl5O3833CSQ5IkNJvNwpkql8ssLy8jhCjKC+Vo0mYO94WTjcHNFGLzd/PzteeU05RluohItXt9gtU1du7YRpokjI6O0Y9j5OoqQmGzCoyb2QuKqqk5+WxhcYFMadJM0et22H9gPzfc9GKuv/kmZnbu4Gtf83VkUcKD9z3Ann37uOa6Co2xEa68+kquuvpqPv3Pt/OxT3yUTtSiOlrCSMOuK/aBD3OL88QO7WqMjrF1xy4OHzvCqdmzTM8cYHp6ykKbeVCnEI405Lys1lqbqB/biYmwpTYyleF5ApVpwmrZhhi1ttliQqAyhXFqqYHvU6tVUakmjRI84VEOy3T6vQICrtZqjIyPEQQBlWoVpRVra6uMjo6QJJF9GYVVqxaeR6VcJghLDmWys4VqtVqEgDA5QiLodjporbjqyis5cfIExmhq1ZoteGwMnY4VVSyXKy7+bOj3+iilbAgxjp24pG3858/Pk6Qpo6MjTqsoLeLmmdHU6nWbWdjvY7Qi9H0I7PPuRzECgY3QCacZpVE6AwSSAe+t6MqMKAoZ2yKsrikJ8DwfMGitKJc8wtCG6vJO35jBP/dIAdsMtVLr0MVcm2ppaZF+316/53ksLs5TCm39Q+06mp/9uf8fiwvLfPAjH+INb3gDBw4c4M/+7M944vAh9u3bx5VXXsns7CyPPfYYMzPTfNVXfRXbt2/jjjvv4J6772ZiYpITJ2bpdSN8L5dzsHIc/Sjm1OnTxHGC0gNOE8C5s2f5P3/0x1QrFV75yi/nK7/yq/jLv3yvbRdyEJZdR4h6HpjWmtiRZaWk0Cp7MrtcJ+tyOUuXe4ynsv6Tfd9s2dMnjT815Ofzts/z3jz7JkjSDE1Go95g68xW9u7eg+95rKyuMDs7y7nz5+l0ugRZSmN0lFJQwpc9+r0+WmU0RutoIegmMYvLK3SjiJ2797Bz9x62TM9QCss8/sijrKyukjmENDaaqNUhMRlaGEpOD0lMjtNeWWGt0yVOMzI0IvAIyiEqNbavFoLR0VFmZmbIsozV1VVWVlaKLLocdcoR2Xyykdde7Pf7xHGM7/t0u10nGZC46JApIhqDJjZwljYuy5f7Ts7nmbbnlNOUZ8EJwGhDp9tlZWUVsKGOMLA6NPYmDnpwYcCkyoWILLoxv7BAqlLGJ0e56cU3ceLkKT74gQ9w4sgxlNKcOXnaQpMrq5w+NcuWLVOMjI4QlgLOn5/j8JHDaAO7d+9FBNBsrVIPG6AEgZEErtxFGFSo17dQrU0yNbWNnbt20xirIjyHUOR1PBhk0RmN5ZuojHUhWQukYYwgiZPCKVy3ivuSeKktbmhcuA9JP47IARUhBN1+18aoPQ/f99BGEccx7a6VArCZC9iMPbE+c0G5opF+EFhZe62LcKNEEKcJQhv6UURrbQ1tDEHgO6TNkMRWksEPfOdgWM6ULz0C36JUUkhGRhqUS1XC0hrlap3xyUkXfuqQpAlhYImOY6OjVGsVer0eUa9H4Pv4gU+r1abZbGEweCKwKA4WkbJmU3YvMAO5lsrGe2wrNbtdZBqlU4aoXcVfMdxe3b9c1iAH1QwaZRRR1C8kDSSQRlHxLDudNgbYe8U+Wq0mBkO9VuHA/n3s27eX4ydPcGZ2liSOaHc6pGlKc7XJfffexxO1GmfPncUY6PWss1TyA/csDGmWFVfvl0tk/Zjhsa9SrbBzx3a2b99Go97ghde/gJmtWwlLru6Zdu1yg1TB88Gk9PCkz8ZySP/yg+wzb89kRttmyNrmaJV5Vga5fx0mBp0BbO4nDqFNWZaRKcXoyCijI6Ns3baNUimkXK0QxTHdqI+Rgkqtyo6du/B9n+XVZdqtNZTKaIzUSbRmtdPh3NwcZ8+dox/HbJ2aYueuXUxvmeZUtU671aZaKiN9DxP1iOI+qUmRgUe1WsWfmqBSrzI5MwVZSn+pSYwi1Clppq0On+tAarUa119/PZOTEzSbTR5//HGOHj3K6mqzmAwO87xyTqnWmiiKSJKkyKzLs0jTNC3GnBz5Hp7cXnQZAz7rM23PKacJ8nRu+zdKUlbXWmzdMk0cxfR6EQLh2PNWSRtph6vAlzYs5eCAdrdF5AQTX/jCF3Di+Alu/+d/ptNsUylXOXzwECrTxFHC4uKi1RQyGUvLC3S6Xc6ePYdSkKUglMY3IVOjUwgtGClVCBwfRYoACNE6BBkytXUL9ZGK5ck4v67oOIRDKnIWNRsm7Xrw3uXZSuvfOqtMBaAyg1bp0E4UyQbejUjiAs60kalilF73gg/I34O48hBYD5vMMI2xZ7Ky1rQhJ4YmeA7GMMIiP/lxbW1BiS+9Qj4iyRTVUkwUJwhPIVttEBD1IzKdkoY+aZbS7fXIlCJJYpI4sdpVmSKOU+ewCCfSbS9MCsulvKRdZOY8XNBXKbNOC2S4X9wYtlvn4hY/WsQqCH0mJsedFpMiDAJXEgVGanW+/Cteydv++3/nzJlZvvs7vxOtNO9973s5cuQwlXKJq6++mvPnz1MKA3bt3MHS0hJHjx4lCEKk5zM+Pkmr1QYjOHBgP9NbtnB+4TzHTpxCCCiXS8RJapudtAkMIBhp1JmZmWZifBzPkxw9eoSjR4+uU9q1s0AJ0kNl69IbvqAtCEJKpQrGKJRWLiyaix3maz23vYCnix5dKpy3mfP0hWt5J8+GpjAI1RU21DVoBkkcURzR6XTIshJplhKWy0xMTlKp1xifmODAVVdRrVVZba6ytDDPWnOVVKesuPDW0uISp06d4sixI3S7Hc6fm7PaeW5iVglLeCYgUiki7WOSjG6rzdLCPIEvmZqaYmrrFgSac/0+kVaQRKg4QxqPQATESUwYhlx99dXc9OIbMFpz++0zGANPPPEErVZrHQcw5wHm2kuw3tnJM99yHSelVVGBYdM2aS6c2mqXAPVM23PMaQKQGDTXXXst+/bu4Z6777H1drIMozSetGnbyoVbhEMKimwsGx2h3+vR6/cIvJCJ0TEatRpJFLFz+06ElPjSpxz61Ko1jNHEaYIUhqifYDLD9JYZWu0ugRfgeVDyS1x31bV4wiN1xD0AhEeqDGutDqfOnuXKF/ZAuuE7j8fmUwwDwpOg9GBQtW/QOrt0Ze5hiBIuFS/JvfRn09Rlevo5aVspRYotKaGNIWut0ZU9yz8C+nGMzdSyFPN+3wpPZmnmeFbalWSxnJ5MWaKN5wiJBoMU3oC/9CyY2eTzukMVj8QS0PMU3dGR0QLCzpVywZJCrzxwgP/9h3+IyTJ27dzBwYNP8Jk77qDTabNlyxa2b99GksSkaZ0wtPyCMCgxMTFFkqbML8zTajYJAp+9e/cyOTHB8upywRkAiCPH8XL8gZorVdBaa5EmCbt27aS5usan/vmfabvyK/l4YIxZl536fLA8o1Rrm5DxhTT+Xy5P53LsycjnX/COk3OYBK5Pvujl5tqC4PmWA5pmGYtLS0W4SSmFEVCt16mNjDAxNcnU1BSVWpUgDKiUQmqVMrNnzmBS5Xijivnz89x5x11UK1Xm5+Z46MGHWF1cwvMkvhfiCSj5IdWwhNAJSZzSWlmlUi5Tq9dojNRpTI4TLi0RtbuoLKFSCql6VUIZ0DnfYXV1lWazSaVcYcv0FItLyxw7dpzZ2VnW1tYKnSVjjEscUkW4zkY7fEuAd7+Ba4eunzRyUBoMN+HOPxcm3FIHRuhnAf5+jjlNud4GbNu6lRdcex2PPfooY2NjeL5PENqwQ5Jl5AVnbWdmUGnqwmH2vkf9Pkkco1JlBRRHx0gzxfj4FpZXVjHaqmEbY0gym3UX+D7VUg1/tMRKs4XREiM0vX6HaqXEnj17kEJy9tw5otimXPaiPivNVXpJn1a3RbffIdUp6xg00nKWCncoP1EpCuEy8uXFRpd/z55xe1rn8SS7LAZfQFhSuAFbudqFP4wAlcVD0JdF3Dx/APNuvm9HIhSWeG9sgA4ulY76eV5b3vltcjbrftHGcgI6nQ5xEhedwdraGmFoQ2CZVhw+fJgzZ84wMzXFp/75nzl37hyddhtjNEmacPjwYbrdLkEQuBIvARPj48zMTNNsrnHy+DHSNKUU2lIsVgelTZpZ4nmcpEVJkNyq1arjky1QLods27YdYxQLC0uDa8j/GJ53mgNZmrmkClv+yPedxthzrIbak9nn6zxdygH7Qkedhue+GNuHWdtk6juYO1tERkiUMayuNd27bp2N0fExGqOjtkB2lrHWbtGPI/r9HkbbSWG/10Nnmka5Soah1Wxyz513k6YZq0tLLMydRwoYGx0FT6HQhEFAQ9YQiUTFPeJun7Vmk2qjigw9TOhR3jJKLDVZJ0bqgJJfpoTveJlLPPLII+zdu4fruIZKucK2bduYmppytePWZ7ldmOkmL5DPkDKvvZqto90UXN+h73Y/zpnKb/izYM8xp2lgS0tLHDlyBJUptm7bRqu15rgVAxExi15YtVWlLOHVldBBpRlCG5Io4tzsWcqlMmFoFZB7vS7dbo9SuYzn+zbTQSl6CEy9Qbkk6HX6CCBJMlZWVqju2E7oiLvNtWYhwLGwvMDs2VlqjRpbZrbQaDTwfH99DAdscUZBEd4QvnQSARtsY7znX9qcQ2dnTc/cCeQk6oJIvclxixBiMV0bhPUuFQWxMw6L7tkBTRck7X9xGzpPIaWdUfo+CGFJj65G1erqKpVyBQ+bKfj4448TeB5plnLffffh+T6VcsmR6jscOnQYKQWlUqkYiKIoZrXZJIkTBIbRkQa9bpuFhQVKpZBOu2u1uLDyEb4/qDqec8wCP6Ber5FmKUeOHnWSEqIInwIXdIDPF8ud9FzVfpib+Dy8HRe1zXhMm3GavtCczeGXfSPenztPZkN/LkT+1yLlNjRlM8uyNMWTHl7gW+kZo1laWWFpZdlyFF2d0rjXY63ZwheCidEx+nFE1O5y9NARMpWRugxmPwjIjEGlCUJCWAqQeKQ6oxcDKqPbbjG/4JEIRXWsTnVmHBV69OabxCsp3X4PI8KiPuWhQ4cYGWmwutpkamqS6elprr76arIsY2lpiV6vd0E4LrfhsFwuzjsoDD+gQGx0vAaO04XI0/OeCG7NVVUul2iMjoAQJElMnCQkSYLSyo3rxa10AT1r+XtptEEa6LbbnDx+nKmpaTSWzDbSaIAU+GGA5wcIaVn+aS+iF0V4XsmWLKlW0TpldKzBC194LWEQotFs2bKF5koTMStYWlzk/Pwc09Nb2DozQ73eKNLOi8sRTn5AYDmlEuvZ6SES00Z7VvsXyfoDDwWZjL2rZriDu3gE8NI2tPtBWIj1XtPQjMGYDefiTGs9NAsZtvXrIGwSpbFs8GfWLgd9E0MdhQGE1beqVmpMTk4xMT6JkNZZzjkA1VIJKSW9fo9azYq/RXFsSfi+j5CCtN+37V74pGlaZFu2WmssLy0jgJFGg4nxMVqrK6RJTL/XoR9Zx99WkJGFUKlwgrBRFDExMc7u3bs4fvw4Dz/8qA0bep4NIWMnKSL3dIGiQN7zwqQrTB0ghCFT6YAGV7TfjS/wuob97Nlmh/lX/liGx9B/kVMdPsiz9DgG9R0YeET5bxeJZgthEejM2MQYW2vO9o+pVnSjPqlW9OOYKEnwfI8gDAj9AB+BShM67TblcplGpY400Ol26XVbVlamVKLsso0zrTFZRhDYyZsUEEiPUHhkwiOOIlaWl0ikYqrqsWX7dkSlhE4NvZVF2t0uWqTUalWElCwuLvK5z93B/PwCX/qlX8ru3bu55ZZbGBkZ4dFHH+XkyZM0m811St3DWnmwvp5ckSiluMDJuhiC+UxpnV3MnltO01AjGxkdZfuOHQhPcn7hPCqzJOA0saEvm66ODce4jXxp3SdjDHEUIQXEUcS5s2fZsX0H5UqV/VdezZVXXYWWtiK9diG+NE1JehGekdSrdXy/4uLIHn7gsX//HkrlEgBTW6Y4fuwEAmi11jh79izjEyOoNKW5ski/13OzCYlNOcdKBbiRXCAwTgnWaG0Rp818mM/zXm6+r1wEbwgKK7DloVjMM3EOGzqtdSHLPEQJF0XcbHkZMUTMHpyz7Z5yZGnwIj1rnfFlPptcsNII+1xVZgVBK6UKU5NTdDodkiRhTWWUKxUajQZZllEp2SriExMTxHFMa20NpRShX6JSqVDSmkqlUihLl8tl0jQl7sf40mOk0bD3y8DUxCTzi3NkqVMDl7bMQiF+aESBfFWrNa668mparQ5JmuIVStXWsc51z7R2IqD/ygfmZ9IMZsChcwKog9+sbYbICnHhss//XC48N3v89ee0vpkOoTyIdcsuNdSsSwYxG5cPBrKN2jm5GrbgMkJy63a+2ct1aXT5krbZrb+kw/b0HV1xwTdxwWVsPKDBOk1aKbRwWcnGkCnb3qI4IdOGNEtJ04w4TQm1RpQFSZaRRTE6y5AaPC0oiQDllZChB1Igg4AgDFFoVGJDy0YKJ0Js8IWkHIQ2i1v1iZKEOEvQvqA8NoIOfIJmh7BWRvT6mMzKE0gpCyXvVqvFzMwMBw4c4EUvup6tW2fwPFuns9VqFZzNnLdpC6hnRe24PFxn3c5BBGIzjab880YnaWMx32fKnltO09D133rrrdx2222Ug4DxkQatVotut2tFuoQtkaGVJjMa6W68LcoKJjMsLy0yPjbC9q1b6HXanDtzlm07ttOPekzPbCWsVJCeRDj5doMdpAMhMUaQKYEQToRRG+bmzvHY449itObRxx7n9JnTaGPodLocO3KUkZEa27ZuY3FskZkt07z0lpewvNKk0+0RRRG9bpeV1WWEsBoWmVZ40sa0C6fhYqjT53kvL/hpYyfvXvaiUy16zWewQQ75Y+6ghZrs0ImsW9+I9UE6e6aDAR2s0KUQBuVi4nmK/NNGxy51/k/ybGxmp3S+vHVCtVL0+z263S5JknDPXXfT7XWpVioIKdm+fTuLS4uAVbQ/P79AvVGnXCkXxYfzmk7lcpkoiopCy3ka79TWbezZtYNHHnkYIwTSg2qlQhiG9PqJBTi13SbvmIWw9aGSOMH3fKse7sTisixDuHcsv4m23I4cZGw+D2xpaRmBZHQ0IPB9jFFusuEalthssIdB4zMbvm/8/ORWrGkGHJnc+d3oJImNqNeT+QLP0PtxYbhy47WKC9bPVxMbzm8jAG2Gf9hkd2KTixMbvzzbjn5xwOGTvPC2Cwa9mF1gQ3MCge+I0krbShhKKYT08ANTVMAQUuSBAHzPBwNpnEKqCD2fwA+JlRXnjZMELUEJ+95rKYjSBGM0wpNUKhWE8UkSg5KKWr1BvTFCWK3Qi2OMJwlKIUEFyqksJgd5v7O4uMjs7Cyrq6vs338FW7ZsYXR0lEqlQhAE61Cl/F9ej26z8K1ArqvtKIberXzSvJn4pX4WUO/nltM0ZL4nKQUeSmU8fvAJtHY4jbCtRicDqFw7pylxkGDJk6w119i5fTsvvvHFdDs9er0+h584xKc/c4fN0JIWxci0lXnv9yOSNEGniiwzGOGjMgVGWUTDWPK5MYb777+PqGeJ4Mpl9JlU01xeZfSGEX7o3/8AP/B9P8jxE6d58OGHWFxaYm7uHHfffSf5C5Xk3KZ1V32pBiCK+Pilxm5ziV/zhm9T9pWrHJ2RpClpktrwZ55eP9zZPFPtckMnboy6YFYm3Xp5aZbBjFoWKI7l21gOk1IgPC7oHP9f6ArlJX+0HlyX53nUa3W2b9+G73u8//1/Ta/XZWZmmnqjzjXXXcfY+TnOnz+PUopeFLFz907q9TrLy8t0Op0i26TRaFCr2QLA5UqFSqWCGVPs27uXrdPTfPazn6HkS44cOUKv1yNNIoSworG2zp6HUmkhBIsxjI2PsWv3Lk6cOonvBwRBSBRH5HWf8usRQmDU84sIfvDgYZL9Gs8vMeLVnSK6QAiPnHu3KajxpJONJ0Fh1q3lPjluR4F+5VMcYwfTvP1rowupD8sFdBlbQwqlw9HWTY87/B65l9W4wSuPQG28cuN4iDIPwTM0KBpjFfgZLF8/PzLr5mrr9m8vYd39GPgo607UrTe0AU81Q3Dzda2Dui4QxyAPbsiLc0TK4Tnb8HXmpaPA4EmB9DxUBlLYqgtAUTpLCIHwJNITmEQX6KX0PPxyCJnnEjwisjQDTxKEAb0kohP1UbHAK4X4lRAv8EgFdOIexmjK5RJeWKWEoq4CyiWPbdt3MTmxBakEaTcmacd4ma2fWS+HtFotWwszDKnX62itWVxc5LHHHiMMQ7Is48SJE6ytreH7PqVSaZ2DlGfPDUQp7Z2wQpgW+ZYuW2iDMqG928a24YITa0QBaDzT9px1mhCWtO37PpmKUeug5iEP1RGstVJFFdm+0uzbsoU9e/aw74oreOF113PgwJUgBCkCISXSs7HUTr/LyRMnuP+hhzh+4jhrzRYrK6ukmWBleYUs6gIUL83oaIO9e67k4MEnSJIeM5NTXHfdC7jm2msw2nDw8Se4+qpr2bl7L7fccgsvuuEGlMpQKqPf79k3RwpUmtqOZN1sdbgB5C+pffVMAVdvBqMM5jDrmtBGR0IIsjRj9vQ5lpdXWFxaZG7uHGfOnuXMmVnOnZuj2VxBZcllPKBnCha79H5yYTSt7QuYZQrrG1vHSWsrRbFxl8/K5PIiO81vc5ambsZvn1XuoK6urHD65Cl6rQ6PPfYY3V6PbrfHw488wtTUFt7whu/goQcfola1mTDn589TLVs5gGazSb1ep922pVQajQblcoV6vcbExATbt23jqv0H6HQ6/J93vYvtWybptDvU6jV6/YROv+MSEJQrj2NDwtKzsgzHjh3ntltv49iJ49aJzqwcRBCEtkyDVoXuipUveP7Y6dk5ur2MufNLTG+ZZMvUGFNbJqjXKi7M4GbuRUV2XCjT6jl5Xl5bMHeScOi1DSsLkb+3OVqzXmBWk4duBD6CUqVMmiZ0mx38ICAoha6Qt0UBsySxYqZau8QXW14oSRPSNLUFvsPAIutOGDAvjQG2DxUIMm3JumEQEgR2CEkSO7HKBQxt5QU7GJp8WyFA+oRBqahIH0UxWZbheRIv8PDcPRLYWpVpptCZss6W51n0fah6vXTLBEU5bPt5GIJyg63YyCla3/kNPg494+EkE4Nx+ngDL026PtVs+F+xPgahFTm3zeR9dO4ErkPF3PGMAWkzMaXjDBaOkgThZozCGDwhCFxiUa5phBBoz5aA0gZkGOB5EuFbdFj6HpnQpEahdIYyEi0gEprUZPSNpOx7lKtVJupT1EYbjG2ZRIqQ9twy7TNL9BeblNoxijLKHwgUa60pl8sYY2g2m9x3333Mzs4CMDc3x8rKCkopyuUyvu8Tu9p6FjU360J2w5QKwbAPbN+lgcafjfhgcLXmrK6T1dB7vksOyEEWvlKaNM0KkSxPSgJXRTnLFEZ6jmg3qHA/LGY4O7dA59Of4d77HmB8bILxsTGklESZnZl/yUtfyotuuoHdu3dy9YGreMUrvox2p4XSiiRRaCNIk7xjVBidEfi2Q6iUR+l02sRRTCkIGBurMzpSQxmDJ0PCcgUZeGQmRQYgQ58An1KtNHCShjBt+/J4RTqmrQPnYbsJWayeNy8pRbH+07FdO3cTxwlJlhDHEb1e3yptR32iKCbNUqwkjy2SmxcYtllVClzmlc0Ks6VQrPaUACkRxgpbGjd4CM8rZsG2T7FhOU9a5XFPSPIpRAHdCokwktXVFp/73J10uz26/T6nTp3m/Nx5MpVhjKLdadHrrl3WdW8WdVzvtMLGOfD6ZZfnKOYzrEq5zML8PP1ul3K5zMLCgm1TSnFubg6lNcvLq/z5X/wlvuczOTnJxOQka6stdMNQKpUYHR1ndGyUdqfH2Pgk9XqNc+fmePzgQaJen1IY0KjXbZhNemzbto0nDh2hUq0Shvb1V8NOpACMZueO3ajMnsf8+fOkaULeORkjC4cJuKCe1PPFOu0ea2unOXnyDNPTE1x91RUEQUC1UgEcV9GhBznfwobw3MCbS8ZjisEuR0itDYrswpDT5A90xowwKKMRRuAFHpm25ZbGpyaYmtpC1I+Ioph+FCE8D2ksIq+xvEVlNKnSpKnC9wWh9BDSQ6sUpQ0CQxA4VMq4d1iDzjQylAR+aJ+9SVGZFSAc1JK0aAjGDJTjja2pZtWgBZi42C6/FyZ3ELDogVKZyzT1kNJWdDAqs+iaqwEp837GwmYufDf0fg5IMesmTZui80KsX5D7WsaNJ0NOkyOzua+6+GfTj/IKo7kmX57p5Xbr+AXFkiGnzhjLecwdvSxzemrF+q6vxWrQWV9LFe+yNrZvDj2J5weFTp1fCqhIjTEZiTAkRhGrBC0EWSjJREAaSEzZpzI+yvj0FibGx8AY1potls6dpbOwTLbcQbQzekIjSyA9iY/Vkcq1lqIo4oknnuDhhx9GCEEYhsUYVq1ajqaUVnS63++vc5iKe+qSdoonmd8j94wsL9j9LkBK16cpbevUZs97RfAhX8JAmloHRngeV+zbx6tf9SpWmy1uve025hcWgYCic4IhlxVa3R6trqvOzvF1w95oqcSnP/UpduzexZbpKeojI67qtEB6Eq2lnd25MiNKKzxfYhROngA7Yy9eWk2aRlg4WGJ0HtawPYktxGyVpa3YmJ1hVcoBoCmVymzftpWZbdPU63XW1prs2m2l8623LUkSyz+RUjI5OfH/Z+/PYy3LsvM+8LfPfOc3R7wYMyLnOatYWVWsIqtYlGyV2BJEC7Balu2WgPYIGA2hW2i40XaPQPcfsmHJ8NDobqvRtiFbNE2ZIpuiOKiKQxbJGnKuyiEiMuaIN9753jPuvfuPvfe59714kZlFJkWlShuIzPfuu/ecc8+w91rf+tb3cfbsBVZW1/nDoD3Cg6QRkhBCp/VDf/7Ito5ncMu/P+Rv7sE4gpw7jobN1Oqat4QXX3yRNDWq4Yf9AdPJBFAITzObzxgOR3iGA2nJqHYRUib4KoqcMIyM4TOasjTZeBSFBH5Qu9ebTE/YBjGFsj6GnhegEciqxEPY7FsgpTLmm0FAVRTgCXw/ZDKZkKYps+mU9957l/fefY8sTdnb3aWqCoMKaNPMkFcl33/nHcCYdjYaDfI8J45jwshMTlEcMxqNaLfbeL7P4cEBa2tr/M9+5md4+skn+NYrr/ALv/ALaM9nnqZWGd+YApt5XxzlzQE/+7N/ga997WsMDg75/d9/hW9885tcufaBuxgPKO3+cRAu/2kfK6trxEnLoh6au3fvIYQiS2dsbG4YXohnzpWUJlQBaj0sF2geQRzqFeBDhjDzUOiHBJHpmMyLnGowwPd9trZP89M//dM899xz7Nzf5dVXX+Ptt99m89QpE0AXBbs7exwe9lFaWe/JBnGSWFRJUHpGGNahAGFgngVZVQjPNBoIIUxyVRgdu3a7bQ9P1BpjDn2oLZmqitFoZM+DWSQbjYZZLLWRgHH3Uu6+rn2v255nty/cImv3UZdDtX7gfjSJ2mK+qRHfh43lfOjY9XCBz8frzFou0T14YY/HZouPKXO8evEXYYNsjePT2qNZumU0FhwQoJW24r7GdByljPSK5+ETEPigfUhVRaUgaDTorXRp9Tq0Ox1azSae7zPoDxn1B4z6A+bDMbGCJGxSMDMItWekENw1cqKVQB0ouTnU3QtlWRqKQFnWqKHzhV3mNDn9Pc/zahRTLaGgvh/U1itu+25fcDQB/qTGpytoslF0HAdGWyaXrK/3+Kmv/hRf+1N/in/hz/xZXnnl93n9jTfZ3TXkWcPPWLphl+5fgTH/9QDftUoLIy9w9do13rl27cTOdA9z4ko+en77RIaAM2c2OXV6i16vR78/4Mz2GZqtBBBUhaTIC8qixPN81tZWOXfuHBubmwRhZCcnv057jR6GNhOkEAax8zyC0EdWBoI3kPBi/4tx9DwqiwQZHSxN6BuEL/B9POGz0utR5SVBGBInCUHg44cRHrC2tkqjkZDlOe12izhJSBoNWq0GUlY2A5nj+QGz2dTww/yQZqvN6soaSkGr3eTipfOfyGk+SU+mqgxR2vcNalaWBdPpiKLIjQu5tuKUszm9TgeBYDIeo6QpgZRFwWAyIfA9VlbXKKuFp9+V968wnc6YTefEUcxjlx8jLwrmufGV2lhfp9FocPfuXdI0RQhBt9slTVMmkwlZnpPO5xwcHADQ7/cJgoCVlRW+9rWv8W//O/8WZ8+cY3N9nclkRJaljEdjNHD50cvcubvL/YMBgaDmAzoe887OfSbjEY2mCdSOLhA/egHSSaNSkl4cIRoNijxlMBhSFHPybI7wBNtntonCwGANcmGv4rqCpDIq9rUfYT05nbzALgo/2Kzaq/27KinJ8pxWq8XW5ik+//Ln+cpXvsr7V66wu7vHe++9x+dffpnHHn+cQb/P7/3e7zNPU9qdDqdPnaLb7TKdzdjb3WU+n4MtlURRZH0lLfphlux6gVRK1gRe4IhNhtPd8X2f0AoFV1Vly+lm8XQq0AZ8kygrd6KUrhdGt00nbbHorAKsRZNeCpQeCOCXuWUu63YJ2g95zesOU7FI3B4u3rlUyqv3ZREnsfj55HHUiHa5VOv+VieXAPqowGP9u7Y8KOuX6aoDgeeZBF9AgIcXBXRWeqyfPkV3bYUojqnKivl4xPjgkOHBIdlohlcp2nGDxAvAD9FVRZ4bCxV3naSU9XWPY9PdW9n3uXPngnHnL7csO+A+q7WmKIqa5+SCI2V1q0xwpCwy6e7HhXbaHxf6/akLmgAQHlEUsLXZ5Wtf/Sn+/X//f8cTTz8LwuOb3/ydhVaNbz3MtAdiyVF1ad2XWEBcmZup1JoAiHyfUGsq20kUBiaIqKSmHQe0Wi0mk5RZURB6ZjupcoiJ6dfzwgiQqLJYzIVeYI6nLqepBSwpNFVhhMbiOKQoSpSNzO7d3efe3f36+N98/Z1/Iqf8jzq2VlbQlfFSa3c7JM2mgfS14tz583Q7ZrLe2tqg1W7S6XRYXV0lz3OyLKM/6NNqt9nd36MoKqI4YXXFBIVRlLDSW6HdaeN5gYHpLUqvlHDVQITwAWOAjMBwPDCdaw5OllKxu7trlLTX1tBas7u7x879+yilaLUaRFFIlmUcHh4ihGBjYwMhBMPhkL29PV544XkCL6B/cABK0+uskOU5N2/d5NJjl7l06VH6wxGnt7c5f/4c22e2+at/7a/RarSYTaeEYchoNCLNU0TgcfHiRVZXV3nrrbeYTqd4nsf58+fp9/s1CXw0GjEYDCiKos7az549y5NPPEESJ7z+2qsMR0P+0v/8L3H69Cl+4zd+A+EJnnjqKaapyeVlzamxBGbh83M/93P8wi/8AlEQUFYVsqqOSjCdvK7/SI3bt24iS836+jqeJyirkr29IVk2o9Np0+60WfG7RxZYWJTqZGXLQx+zil4vmdp01LoM3G1TKYUnPNrtDufOXuDxy4+R5wXdbpdms8nP/MzP8LWvfY2rV6+S5wV5XvDMM8/ywgsvkCQx3/ve97h75w6Hh4e0223WVlfpdboEvk+eGv8zFwi5ORYMmjCbzZhOp3S7XVqtllVKV7V1hksU3M+Oa+L4We5np1jvfNdggUgBRxALh2i47lF3bn6YcWJe+FHX4SFo+YfuRBzfyw9fATgpqTNbsrp52txXnvsHCKUNp9Nq/pm/g4+HVApPaJpxTNLtsLFxivWNU8TNBlmRMxqM6O/tMen3kWlOiCAOInxtXOXjMKKwCbs7rtrFwF47R/IG6uDadcq5gMjdJ+77OGFLF2QDddDUbDYBKIqCPM8fQJbA3EeRpeoYWsEnOz5dQZMdeVogNLz00ov8G//mv8H5C+fZ29vl/fev8Z1vf5fhcIrwYnzPRxU5tcCcXoqWWPyosfOWEDz5yCWiOEJWFUVVopUi9L36zZFvYM6qkqyvbuEFAXmZgdAo3+fWvZtkBZSFRBXp0R35QFUtF2jBA1XJI4uRriBT5VEBxqW6Otht2a/lAaiFJqbbtNZGAV0vfVRZhC0MwrpDb3kXkVV2NVG/CeqOP6QPSAHY4XF0/tfAwXBIIwiZzabcH/RRQIDRBHn17bce2MY/7SMKA1ZXu2xsbLK2tsbKSo8gCMiylDCM2Fhdpd1ukYQJwvNIGg0uPfooz7/wAqe3zzIYDk0HE9BbWTEEdqXprPTwMOKq+KZrsyxLDg8PuXDhQu3PpLWm2Wxy7ty5WlgRqIMmt3gcHBzwjW/+Fu+99w7ddptnn34KJTVbW6f4F/7Mv8hbb7/N3Xt3CQOP0tb9taqwUSVCGARDVpUhm/qe0Y75ESzDPWzs7u7iezGNJKHXbdLrdOiXKZPxmPv379Fb6RIGPu12m8D3cZ2rjuxdjxMX1ZPG4jksqxKpJNq2jMdxDHqhd1M/n3qxmHW7XVpxQqdtSu5BEPDCCy/wZ//s14mikDzPeP3114miiOeee44Xn3+Bxx69TKvVRJYV9+/t8t3vfoc333yTO3fuIKUkjmOazSZJktRIglsM3bzhvBUdutRut2k0GrY7t6wDsKqUZMokMVEYEsdxjSoUhZXXsOjacqB0/Aw97A51pX1t+TD6h49bTgySPjRu0h+GJn38sRwwPSx4cl2RBknyjVc9NsDWZkHQYmmetsFU0mrQa7RohzGiksyGIwbDIXv37zMeDtF5TiOMaSYRkRaQV+hKEYchBBVFmtXzTxiGD5TSHEHdSQ2415aDXiMQuyhbu/un0WjUxP+iKOqgexlldWW95WaUPM/rQP2THp/KoMnzIcsKRqMxm6fPUVWKf/ff+Xf5B//gV0D4aC0MzOuBHybIcv7hSCgg/IBcVvyn/+V/wU985SvosmD3nsm6kjgmCgP8wCgrZ7OU8Thl49QZNjY3mAwPGc5GrJ02+kvvXruFF0SEkb+AYQ2GbJ0WFvIAeKIWBdTSwPdBJKhyVR+vMwYFTRwZc2KlTXeEVAqlFhre/jJs7BBsS5AUWCI9UMmqPh2+7+F7PlIqY+XhDg1RQ99HhlaEgdl4VakFHcP+P/R9tJRG3FyYDNwFpkuAsv3dcYQ0yywZ22XqhKfr76I1HykXcJJFgxthaFpxnfq1G74vbAfeYv+etxDONHwnkLKi3x9weNinki7YNojWP/zV3zT7EBAGPrJSFFoTeIJKH+NosSjLnCje+QmMOuv0bEnWnsswWMTuwhN4gbC2ccvwvzmm0PeIwoBKGlPkfz4WYzqdMhoOGPU6xJGg2TAB1GRacHh4yN7uHqsrq7RaLUPexqvLCdqSqj8WLebY0FohK4Oc+oF/pGyhtCbLMrI8o5AVZVXWi1BZuiRJ1EjQmTPbbJ/aotFocP7cec6ePcvFixf5+p/5Ol/+8R/niccfMxwsDbPpnO3t03iex3g8ZjKZ0Ov1eO6552g2m4xGI27cuMH+/n6NKLiSi5ErkUdKbstlFld60xobRHmEUYyUyi6YVY1aCbsgO5uRmuOC5ftwnNO08CpjiXn9wz91y3IMD4opnhRQfZjEyyc9hCv/4uG7DjOhkNZKzBU3fJu0etr8C/EIFJSzlDzLmWZzBv0Bg/4hMi+IA58k8mmIAF9pSqmgUkSRj7ZlOXedq6qqieAORXKBUc0BXjLjBRO8N6xunNaa1PIuhRA1ajmdTknTFK11vf3lbUSRsXIJgoA0Tcnz/AjK9UmOT1/QJBYetp4IUVXJ3/7bf5vvfPd7/O//w/+A8+ce4e/8f/5rvv2dVy3B3i3GYnHXnJCOlMo8lH/5r/yroKTJsAHfs+iN0HXnnu9DWZmyXxwGgKaoSpJWwMHhBBT8b/7GX+ff/Lf+l5w9t01Zpqgqx9cCIT28IMYLQ/A9lIW6AdCa6XhEu9Mmnc3I8xzP87h39y6/9/u/z+3bd1lZXeE3f/M3eebZZxn0B3zve69x+849yrKkEcd85jMvcv78OdI0pdvtEscx169fpz8Ysr6+ThTFjMcT5umc8XRCv9+vIXxP+NzbuY+S0hYOT55WlAZdLf6qgU6jge95TGYzUuvz5wHl0uTlXlvuDwJ9pKtxeR8cnfdM4GSksRA+JEmCrArKYulI9VKwJBb0Bcdb1kvBi+8vJt2FRIWdU+17osinqjRlpQgDnzgJKMsKrSRxaAjUhbTbF5DEAapUZKU5B5HvUcil43P8lWWyquebHdeTiQvYzI0ehmGdQTlD3uVFwS1Oy0igQwPlUrDkvt8y2Lk4Hyby84PQmCJbmyGBBiXx0PgCqn8ONNWj2YhBK+7fv0dRzDl3dpsgDGk0jBq7U3fXLmmxdjVSLjUWOIRJY+FhwckeP7p+fXnRDvyAMIwQCGRl3Odd4hP7Qd2hZIROjcBrt7dKs9k6UhIpygrh+Zw+fYZz587xxR//cS4+crFueddoWp0mf/7P/TlWVlZI05QPPviAS5cu8e/9e/8eKysrvP766/zyL/8yeZ7X37uqKoLAHIe7n5VSpGlaL7BxHJPnOa1Wl06nzcHBIVmWgZUjqKTCD0MajSZtq0OWZhmV5VGa8+hKy3qJ+mCHO8d2MlhM/26S+bCy2YNBkSNhH+cZmeFwHMFymviHHQ/nSx0b9ZTn/Ni0TcRNw4paEpMUeITCg0DgodBlxXQwYjQekVcVaWGunycVsR8a/lJWkekSv1T4CALb6BQGAd1ut0YTHdoYRdGRgNbxktxwSKGjRzgk1L2/stImzWaz9rWbTqcMh0PTuNBoEMcxcRzTarVMZ/HaGmmaMhwO6ff7Nc/qkx6fqqCp2UjIy9KUN6Rmf3+Pt996k7/8r/wV/rX/xV9jbW2dX/zFX2Y43EdWc4SXgG3/FDiTEv3AYuyGUoq8LDi1tUlZFMiyIAx85jMz+UkNk3lhaqZhSJrN8YSgnUR0um3u7x7S6YWoWcl/+9/8f/mVX/5llCr4qa98kf/w//R/Rk7H3Ltxm+98+1XeePv7zPIcPwpJ0xRZSU5tbbG6ssL6xhqhH9BoJJw/d56yKAgVvPD0s5w7e5af+NwX2Tp9iv/3f/V3qNIcISWxL1hf6bGxusqFc+dM7VfDNJ3z3HPP0e12SeKEg4M+3ZUe49GIL3zxi/zg3XdJMxNgVZVEKcn2qVM88/RTjCZjbt2+zeFhn/l8Rp4XzOcz3n//fX7nt3/XlvkCXnrxJV54/gU2NzZpNRq0Wi08ITh/8SLj6YSyqvB8j0azgdZwcNjH9wMmkwn7B/toNKPRiMPDAzxfEEYhBweH+L7HeDIi8H0GgwF3795jOBwjlXkon3/+WSbTcZ2lZmlOVZVEcUyn00FoSOdz2p0O7793pW5vzfOCRhJRVkYd140kjqmqkspGGUoudZ8YggBFkRsuCgsuUD00ZNlRNEYe14g6CVXSEpZwNq3B80M836csCsqqRHgeQRhQFSbAdpkcmEBqmS+wPFmBCTDddjXUnYTSRa8a8AOQFbLUdVaigGarzeb6CtPZjN39/oPH/iM8fN9Dqop+f0yWTRFC0m41CMIH+TVSSSNeKBbo5UJeBOoWqI8BTBgitG8TCfM5pe31tdm9G4m1t1hIRGiadsFxWknOqqnd7nD58iUeffQyZ7ZPM5+nXLtylTu3b5MkCRcunOfyo5d56aUX+epXv4rv+3S7XR599FE2T20QJzGvvvYqb731Vn0v9no9Ll26xBNPPEGz2WQwGHDlyhVu3bpFlmV0Oh0ajQb7+/tmsZSKeZoyn83q7+EkGszfpf2O2rbVe0sIj1PuOf6MWXRpCRnS9TpwHCd/+M91gFsHQt4JSKFY+vdHHz8MR6tG65SiEgKkRBUVsiwNr8kifIH2CTzPaFyhKaQim8wptSSXFaU0Qs2hH5CIgEh4CKkQlYJS4QUBgRCoShIEYX0vpWlKZrsfXTnOBUbLqJNDgJaRueWS4zIp3JXZXNdwURhNMRdsO57T6dOnuXDhAgDz+ZzxeFxLGbzyrU/kUtTjUxU0ZUVu9T/MuHr1Kn/nv/o7/N//5jP86i/8fR599AnS1ES6CAVCLxKP+oE6NpYTBU+Q5Tn7h4ek8zmyqmgkMUhJWVU8/vSTnHvkAr/zO68gKvixz32OdD7jnR+8wzwv0ECamvrt/t5e3cG3udIjn41ZbzWYNBu8/4Mf8Ou/9puM0zlRaJTFy0qSNBLarSZRGAAecRzSbrYQGvKiIIwTms0GjUYTfJ+rV64yHozw7QSVpSlvvvkWH3zwAWFoVFillmgEUWgi/yzLaTabzLKMN958k9F4fKRbIQh84ihmZWUFrRWT2ZTZfG7r0gY67/cP8ZyPotQcHh7y6muvmc97Ho2kgdCapNmkkCYI8QKfMAyQSlHklZHTj0LKsgJhgqBer8vKao+qqoytSFnQaCTs7OyYoC9pIOWIuBGx0uvgtI6arQYokF1FHMeWM5GQzlMG/QHdTof5uTO0WkYxO0kabG5u0my1uHPnLuvra+zs7LK1uWEftAylJINBHyE87t27y+3bd1hZWePxxy7RiBOK3MDHYRzR7nbpDw65fPkxRqMx89mMVqtNHIUUZcnp09tcvXaVzVOnabRapGnKdDY1KvMWFVjpdjl75gyNZpN/8D/9Iu+//96C3KpMBl0WZc3LW4a4HQrlJpvjfDPtTKBtgKRc0u1ZsVOpQSv8KEIW5YIQh1EUPhz0yfPqKCL4zxEntFYoWVn9silSZmxubbC+vkYQNOt2aOEJpFy4ti9KFguNtR9mifU8QRD4VNbE1CxK5m9uMVruZltwEEX9Hqe6LKUp6QVCsLq6wuVLlzl/7jy+73Pt6jW++Y1v8srv/g7tVosvf+lL/IW/8Bc4dfoUX/7yl9nf32cwGNSlkHPnzhq0wHp1mme6x+c//3l+9md/llOnTnHlyhX+3t/7e+zu7jKbzdja2uLcuXO02y0OD/umqSHPazTBLbrFEnHYSWTYOMiO5Tb8k0tlf/SxWNCPoIQsL/gfwXH6w+75Q8qAy0Gj0oALUsqq/scSx8zzfQINgfAQvqAsM2RZoDFK5EYDR+Mj8DUECnzt4SHAE0TCwwOKsqyDJlf+df9395hDvRf36aL7UQhR8+/yPCdN0xqBcnyl8Xhco0pJktT3wHIw6Xhyp06dotVqEYYhSin29va4f//eJ34tPlVBk9Kuag2BL8iKilde+T3+y//0b3HuwgUunj/LZDIBK5JovJaWVVpPyvIxlYnAetWVBZNRyU/85JdYW13j9dff4O6dO2xtbvH1r/8Znn7uWV5/9Q2Q8PWv/xm2Njf5pV/6Byip+Mff+CZlaZVJfYgCyAqNtHpArc1TXBARfpwwyXMmWUZYmBu5UppJWTKYTLFVKJzgnctbFAueihICUVVL+jqKNEuZpilyx2iSOAjaTJdmO8retBq4dv265Ttp+6RblXG7D7CZa92lYbdnOzG0NqrE9/d2uXv/Pnkl7Xm3D7HLfNxnl7bpCaMq7fwAG42E1ZUejWaTqqqYTifGNR4YDgdsndoinWdG6EwpsjTj5o1bAMRRhCtXhFFUa4OUZUE6TwnCiHSWMp9nTKczojhhMp3RaDQYDYfMZzOG9v+uLm86hAouXDjHZDIy6I+Vr5jN53W3W0M36a2sMBqO2N3dIfSjuvSQJAmTyQTP89jfP+D09jZbmxt4nmc6P6Qk9I3Mw3g8RhYFpRCsr63SajWZTedWCiIgLwqSODSLhWe8qJRUlLIyi/MSOVJ4RljRTVxZljEYDEzC7Zsgyl1yB3CgJFqaO80PTNkYrcmKEqU1rVaLII4ZjcaL0oQ4ysf6oxUiPn2jKnMCP6TZalCVObPZBO9AIdAE/sK92D3Bjt/o1OuDwAVNi4DmKJJx8hDCSIgIpVBSUxgfg5ok6wLpQkmUMu8XYsHtmM/n9h6V9m+CMPTpdjusb2zQbLWYz41Q7Ouvvc7v/u4ryKpkf2+fy5cvs7a+zqXLl3j88ce5efPmEoHXknqDRXNCVVW02222tjZZXVthZXWFLMuYzWZ0u12+9rWv8fWvf500nfOLv/QP+Lmf+x+4dPky58+fZ3Nzg7KsODjY5/p1w5WaTCe2nLOw4Ygio0xukNTlBXUB3R0B9f5ol32pTLe4Zv8kdMoeHiy5hAmw87WUEiUrS7Uwk7W3zC9D46kAHVghYQRBFBq1cCmpypKqKNGeQoQhcRgTBZ6Bp6UxGUc7+RqTwAVBQBRFNXfNIULHj3mZ8O1QcRd4N5tN2u02cRwzm804PDw8QkdwZHIn9eK66waDAf1+n7IsWV1dZXV11VpM/YiX56yvSE30C31Bmqb8vZ//Bf6P/4f/gEcuP8I777+HEFa1VVcYKdoFH+CBsSB31Al2q9Xkz/+5P88Xv/gF/u7f/e/4+3//f2Jza4ONtTVmozFIRRw3eO65Z/nTf+pP8/TTT/Hmm2/yj7/xzXpBUWpBbM6LnBtXr3Ows4dQkuF8isR0MijPTaqmHFRqjW3Qc9/46BQqBEp4CITtZnKlJFBaWh6XETMLAp+8lOhKmq4JN3HYbTuez/HMyCWljqDurATM+dFGEbjO6DSVLK1XncLzDFeoLDW+b/bxsOkkz/N659PJhMPDfp2VmMxY2+NR9L1D0tQgKrKsmEnJdJraYzg6eTiugbkHjPOgk5QQNsNxCubStrsqKRFWP8dxRkynjsd4PEUqTZ6l3L9/j/F4ymw2JwgEfhBx9959dnZ2ufbBTVa7PSPaVpWEUUiW5YRhQJ4XTCcTtq5eJUkM/ysOA1qNBkVRcv/ePYb9ARqotMK336nT6XDh/Fnefe9dNtfXmM/m+IE176yk4XX4RsLCdKFE5h5RhjAZxUZl3gVN5iRheVP6SLVQK6fO6sJjc68EYcz22TNUZcVoNK7PsUAgWbSOO97Ej8qoqhIpC8Iwxvcgzyumkyloal4TWGTIckAWZYqTAqXjPz98GG8tExwrWwJ2QbJDmxxx2uiLLbSUCuspacqEXs2DiaKIZrNhLHIqyXyeghCcP3+O4WDAZDLh1q3bHB4esn3mNGfOnKl5l2COI/CDmpC7ECYUBGFQW8k4ou/KygqPPfYYzz7/LForbt+5wze++Vu8/PLLfOYzn2Fzc5Msy7h37x6nTr/D3bt3ydKURqOBHwT0+312d3YpKzP/LCOtSqo6gKobIgLflvPcef4442iYVWsjLSFay52CH3dbxzvgjo8PI5k/NHgSi/Ks4zMqvZhHAZSsKJVEVAUi8CHwkZ5G+8aKxxeeWZLZQJ0AAP4wSURBVCeEqVAIqUzVxpOmacl2+2KTSFlJZrNZXXJzmkwu+XRcpuPE7eXOueVmgWazydraGmEY1s0Gxvs1rbvz3Dl3CLvWmv39feI4xvd91tbWuHjxIv1+34Aon/D4dAVNtuNEa5BaIzzotFucP3+et994k19b/xWuvH+FNE3NYqkqW2Y4tmwfT+Q0dQavteb06dPMZlOmsymXLj3CqdOnmEwnvPKtbzEcDJhOxmgteO/dd3n6qaeQUnJ/Z6cW9vIsp7ey9YzJdMLvffs7hL5g+/RpJvOZCXgAaR9oVyZxRrRuSl3OkNxPzmdnudhYa01Zvzrf9/GDkEBZwTAbzAVWg6A2uj0pohFmO0a5eAk1wgQwysLgruNXO0TM7lt4VhfLQmYfOp3oRTWoKqvFcS19Z09AWVT14mA+Z7Inx9NZftmgyJ7tTlsEYApRE3nMRGdFnZzSrq6rUmiMKfTu7g5pmiMwiu3GTiajLAuk9KjmBfN5SlVWiLygzAqkVEi9mMTduHnzJnfv3sPzPLqdNpvra5RFyf3dPbSUREIwTVOUZ6wwwOh19VZWiIKAwPfxPYHQGmnZ3GHgU9muR98q7Zp2Xk0QRnS6XUs0p/6OCPB8Y7ejl0jqDqEymePiuMMwZHV1jTw3Vg61wu9SWQIeeMr+mR9SVbUOTBiaEsFwOOCwf2iaOdLUBk6CIAhxliDLJdSTF80PX4CXk4plfSL3z3Wn+ZjuI6OLVNVK5J4QpmPWLmSmomMEb4UwiUIYBnS7HcNf+spPcHhwyOHBAUopdu7vcGpri263y+bmpiXbCoLQt11MCXGc4HmF3YeoUSBYiHK68k0cm66pXq/HM888w1e+8hU+//nP0+l0KMuSwWDA888/X2uVtdttiiLnjTfe5JVXXuGDDz5gPB4f0XJSWh05H3geQmmUUMfKaA8PRk4aZq5a7sbjgef8kxofhSw9cGA2Edc2eNJgAh08w2u0UjlVKZFammqF7+GFPkEUmLXVKW3j4YUhQmmwyH6hNL4WhL6PF4QmGSwqZvNZ3b3mut1cYLysx7VcOnYBtRM9ddep2Wxy5swZer0eo9GIg4MDdnZ2KMuy5jAB9TyntekWPTg4qBH1ZrPJ5cuX8X3f+Ll+wuPTFTRZ9ARbOqiUpt3t8b/663+d3/i1X+M/+Vt/mzTPOTg8RHjuYbWR9vFAyd1zS4uIkiYQ293d5b/9b/8uv/zLv0Kapty7f488zxkPR0hVUlaS8XjEz/3c/8Bv/MZvkmU5w/EAjbG+wBMEnglQirJkPJnxwfXrPPnE42xtb9Nsd8HzTJZfKeTSDS9spOQpOzkeO3CHgpgc0eaJwtXzNZUCLTRVWVEqhaxkLauvBBRLEgHHT4UbTsRQyUUU4cp9plxn3ucQKVmZCVcrkEKb7ER/OMpU73s5wV7K3hwK5P5WllX9AC7PT8tzh0PQXNkQKeoXNNTlWvceIdQi+VsKvtwmldZMJlPjZaiptUKchUgtwFZWeMK0kxdViXOQ10vopfCNqamZLDRntk/z/PPPs3P/PleuXKXTbtLtdZmkc/KiqI9lOBrz/R+8i1Sa3b1DUJX9fuZIHfqnpAJR1Nmm1ppIKbp0abdaBGFIVS50uZRy5qMGUdTKOYpbJV3hgipT2paWROwCpuMZ9rIFwj+JUsU/DcP3PcLApyxy0BVRu0kYBpRVznA04N69u9zfucfa+iqrqysIYYLahYekQ/z8xXO9dP5cOc+4tUsbDFc1oXtxHCb4cXyPNM3Y39/ncDCk0+3yzNNP8aUvfZn19XWE5xHFMe12x5JrM/KipNFqITDob5GXrK0mXLp0mfW1Nc5sn2Y2nbBzf7fuhAPPutUnZjFk0cWplu4VJ4dQd6kqIwZ67tw5VldXaTQa9fdotzo89uhjXLr0KKdOnUYDsVI0Gk02NjfQaqFJJaXk9PYZWq0Wv/Zrv84bb75BluUEYUCv1+PUqVOs9FZst2nJYDDkxo0bpFlGFEVEFvlabpV3pWaNrgUSHRJSlZX1ejM+buZBs3ZKlq7g+55Fspx+lPv5ZCTqwTLf0deXuYnHn7flnx36ZV43gIISmqTVpNNu02w0kGXJ4HDAeDymlLmpcmiDxHsotLCUDq3w/cB4CAofzyxEqEpaHplAWXFgrc28sNx97OQHHAq0LHK5aD5YJAzLIqUrKys8+eSTfPWrX6XX63Hjxg2uXbtWm/w6fpTjOy0ji86SxSW1Sqmaa/VJj09X0KQkp09tM5/NmAyHljsQ8Oijl/hGIHjhhecIo4R/+I9+jensHkdW4mPI0vGhLeLgCY/p1JRfXNu3qdULtjbXefTyZfYODtjb2+envvZVprMZv/Ir/5Af//KPc/XKNcI4ptPrkFjjztFozGg84bd++xXefPsdXvnW7/P++1epZEkcBWiHhAiPIDAZYhgENhCRlr9gyi3uBlPSmHS6r7JoNTd9gq5spyqFXjI7dMjMSWM5dvHcInoMwlkGel0w4OrariPI/L54gD9qaLdy6EXXS90yX2fPUJQLP6NlRGiRMdpzsTjc+jdXsjx+PEeoDycfnM1oDBpTVhKR2xZyhDE0FaaDUCvqrjswqKXAdkq5wMlb7Mw97EJAu2X4XK1WG8/bZ0kkg6qsmIxHSKUQVEYKQCwO2UzUAnkMUfU841Q/HA4Q4xFKSXzfmIgqG+DWO8FOykqDD8IGgA6FS9OUu3fv0ml32NjYZH9/z+7bLNZu0vxRCpiAGpExXT0VUeEThD6xjJjPZty+c4urVzZZWemRJIbIaoIJd960fb69WoftaELgEAx3bk13K8cWT4fYGBX9nPF4wvvvX+Xd997js5/5DJ956TMI4XH2rBFETZIGZ8+e49KlS6bdv8ityGRBlmXkeUHgB2xvb7O1uc7aag9ZSbZPn6EoCkvQdpIYuuYlgtExc35iDvH2fFMScwKtjz76KGCEXDc2NtBaM5uaEs/pU9usrKxQScnu7i5FURBHMRuba8RxTFWZ0lsQRjz66KMoKbl79x43bxmj7larzRNPPMlLL32GRy5eJGk0rDzCdaTS3L171zS+eNY02AahztjVPDfGZNc8wR7OZFlJjfaxPAfzjChpNPQ8+5rwjH2U1rpOnpZRqeXxYUHTcTTSBUkPKwU6IQWlFVIrFJA0G2yePsXayip5mlKUFdN0jspNQuQLgcRWN6REVKWdKwWelWNxlQhDZzAUB+2ZhFaVC7XuZd8582z4tdCqEy+ty8VLf18Wv+x2u1y8eJEXXniBbreL7/usrq7SbDaPWK24z9eBuF5okDnF8Ol0amQr/hjGpytoQrO5sUFfCMajEe1mk/W1Ff6b//q/5rXX3uCnv/ZTJM0WQRjAEc7AR27WLNoWpdEavvKVn+D8+XN873vf4/33ryCE4IUXX+Tf/Xf+bd599wp/97/7u3zuc5/jySef5NTp01y89Ai/9Pd/yfI6BEVVolKj31TKkns7O1y7fsuYJGqNHwSEUXBEuNITHmEQGghSGwJuGIaG9GtvLtNAJUEatMFzqJMQBFFEFIZogc1KFxoZtZrzcSXvYyiLTaAW7t+21GViEL0IjpaGQhuVNMeVcujNcXTvYadfH5049DGp3uUAyHPB2pHPmn2dhK4Lh5hot7gvvlf987HDXA68itKR200JI9eVLafaCQwjMOqcuE3b72IfRzYqFqdfKoXU2vjora6B8Eiz/KjitnAegIZH4PkeroDpRO7Pnj3LT/zEl/nW7/0BN27cwBeCRy49wle++lOcOrXNK9/6Ft/4xjcMmuo5sih2UjQLt8Z01EhVoZeUpB1SVpYlw8GYp596lueefYGf//mfr7PDZWj9RylgApCVCTTCMEApyXw+M2KhvpG2ODg44N333iFpJERRZK1/IrQ2nKIgWPA5FufRbHvRnbW4h33fJ4yMFIVruw89D99ydUzwZQxx33nnHV579VXOnTvH2bNnWVnpsbKygicESRTy2c8aztDFCxfwPZ+93X0ODvrIStUda3megzaq3zrSOJNys/CZEktVLcp+ixL3UR84p/4shEe32+Pllz/HU089SZIkXLx4kelkxrvvvsvO3i6tdptmq81sNufNN9/i3r17tFotXn75c2xubjKZTNnf3wetOXNmm/WNDR5/4gkeu3IVKRWPPfYY/+K/+Gd46aUXOHv2LGEYks4zHnnkEr1el29/+9t8//vfN8iJRZCk9YNstVrEcYKUxietLEtbkjaBQZJE9rwvURO0S1oNeigeWHMW6ahDhRfX9+OVBj8qYLKnHo2Zi10iqtCUVUVRllTKdDD7UYgfmmVfC/CWktX6utX/lr6FMNUGZas3JgkVdQOCC2Bc8OJKxAsOn64DH9eoA0cTLtctt7e3V5+fXs/ctzOrW+i4ectdoU4gVUpZazO5IO24sfgnMT5lQROMhkPSdI7SmlOnNvnqV7/K/+9X/iHXrl83TvVBQP+wz4mr9ocs4nWibu/Jl19+mZ/5mT/L6dPbjMf/Pffv32c0GrN/eGil3RXf+ta36K30OH/+PG+8/gZgFsPpZIJU5kHUddDidF2MZYope8h6MUeDQqJ1aerOlbRdCJrKdyRSVbcI18e9/N1cPZvFPyGcCas2reUfQaHQWJfs5cCktqFZ2ufy35dkID7eOGliWdRLzUT7IGpx/FP6+GsnXF+HMn1Y/PxQoAksT8xwQAzapxeMevueZS0m05nojl0cPbdLO0rnKYeHfeIopN3pMhoOGc3mBmnUgtLC2poFcqe05l//1/5Vtre3+cY//se88dbbbG+f5V/+l/8S93b2ufbBdS498ghPPvEErVaTRiOm2WzUB6qkxvOpy8CW6oHJml0rtzJkeQvBC88kEVKagMoF8FpbTysnfvgjFjCBOSdZltJqtdBaMJ3Ojc5YaMo6WZZx69Ztoiih2+nSarY4c/ZMPdE7JNIYjjrUAU66WV2WHi2VuxwqYAJfVaNWZVly48YNvv2d77B95gzPPPMMnXabe3fvG2Sz3eaRixe5eP48jUbMcDThxo2b3L+/U4thVlVlTZvHKGm4obOpuT97vR6NRqNe7NwCaMjli2fBfc/Kks59z6Pb7fLEE0/i+aK2z7h/b4fvvfoqw9GY8+cvEMUh89mcq1ev8vb3v0+73ebCxQs0mk12dnd56803mU6nvPTSSzz19FOcO3+ep556iqIoeOmll/ja136Ks2dP4wchk/GEpJFw+fIlTp3apNlsMBqNuHfnLtPpFCklKys9zp49y5kzZ+h2O1RVxWAwYm9vl93dPYbDwZGStBknBz8njeUy/fHX3TaWn5+Pep4eLM2Z46nnZTvpZUVBf9A36tqARBPGIWEZUzrRVXcfublTabD3k7aSJI4ygku8wcwNDkETi+u9fP2dcrcrf1ZVVSNE7ridKbRSiuFwyDvvvEO73eaJJ54gDEMuXLhgOMP377Ozs8N4PK4Ry2U+lOM4ObkDd3+Wx6zCPonxqQqaPCG4d/cOUikuXjjPl7/8JbrdDlJptk+f4vDwkIP+gCzPOQI96OXV/uThyhWuPnzr1m1u3DDttI1GgtaK1159jf/ob/5HrK+vcfv2LXZ2d7h1+xZC+Pz+H/yBIdJ6giLLwTfkT4GgkhqlSiN/YW8+cyMtFKDdMVRSLpWrHNzpUk5dZ3P1EObGlkqTFyaDNSU/817hAimtF+zy4+fh+MNs/+MQDW2zkaP7XXzuoc/3Q8/3RwVN6ijatBR7HL+UJqhYfPz4sZggx2VRS4f2sGMTJ52gZfTpWDB37B7TNkPT7mfv6AacdcpgOOD7P/g+rWYLVVXMs4zpbEYYRUZ88giKgw2uNX/hL/x5/uyf/Rl8D96/cpW8yBmNR6z0OnTaLdrtFu12h/5hnzt37nD9g2v2MJczVW0XYa9GOmUp8QMfadE4w9swPws0RVnw9ttvsbW5xSOPXOTunbuGn+OQMG10hx4Q8/xneAhhOujAoICuDCWEIYYXhVEFv3fvLtev32Br6xRr6xvEcUQcW/f5JXTJBecLBG9xv2k0CCs3EfgEYVCjx3lW4JIN04oP/X6f1157jSiKuXbtOuvra/T7A8Iw5NKlR3juuec5t71NFIWMRnd57733uHfvPufPn+fixYtUsmJnZ5f3332Xa1eu1ujCmTNnePzxx9nY2Fh06tmgyXFJHBq1TNZ1mX8YBiitjB6c5ZxMp1OuX7/OfJ6ysbGJEGbRnqVz+oNBrSMXxRFZnnHj5k12dnfZ2j7N8y+9yObWFucvXGA0GXP5sctcuHAeIeDWnbu8+eabNBpNzp87x6XLj/CZz36Gm7duMRwMOTw8pNVq8eyzz/KTP/mTPPHE43Q6LSqpGA3H3Lx5k9dff53vfOc7HBwcUBQFSZIQhguT5IeV3uqxNC8Id9O435d+dts7+tGjQdVJaFO9XYtkYY9HeIKyLBlNJuR5QWzRnSiOiUvjKSkd34jFfKWVMTKvS5M2YKqnOXufCa3wPTOfLXOwXKntuG2KUsrIYBRFLXp5HHWbTqdcvXqV2WzGeDzm2Wef5aWXXuLChQv84Ac/AKj1nJbPl/u8QzVdUFUtfcdPcnyqgqbAZlFSw1d+4sv8S3/xX+Jv/+3/lLt37/Lii88znk7YOzx84Cb8uEnwcoT/W7/1W/zgB99nOp2xt7cHCGazOe/84D2iOCTPC5TSjEZjfM9nMp7ghQFayjo4KfN8sdoKgbQk7Fpm8zhU4sAcu/C65q4PC/qWX1JS1VpOpqxjKvPHodaPPhGLDS/vfulE2Tjn2HaPxBtOr+Ckg3ZffBEsmX0so4Nq6WlevHzS1ziOfB3flTqOhNXHeTQwM19NHPndlTaXlR3V8vFrYbsuF+TwB/az9LXMJANFWbC3t0/gDwwnyWVjSqHwjIwAoKrKIj6Gc/SNb3yD7e2zzOdztNbs7e7yO7/9W4SBYHWlx9Vr19jc3OKnvvpVDgd93nnnPXMoltTtgukkiYjjiDQrKKyxdK0Ipl0ZU9c6XlKW7O7u0Wy0ePyxR9nf2yUvCqPvpKl9wH6Uhu95huenNb7wSJIEpSpMULrIhKfTGTdu3GRtbZ3V1TXOnz9PEjeY5BMUijAMLOonMMRh50+3aB/Xapk3xlLCoi05XNkSkukwmkwm3L17lz/4gz/g+vUbtNttRqMRSZLw6KOPsr9/yBOPP0Gz2eCdd97lO9/5HuPxiCAwSv2rqyuMRyPee/99du7tUJY5a6srfP7zn+f06dMIT9Qkb98S24WAtbU1zp07RxwnTKcTIw+wxD+ZTqfcunWLqqrY3NpkbW2dOI6MGG2akTl7FN8jCE0p0g98Gs2EOImRSnLQP+Te/XuMx2NTVmu3Wd/c4MLFi2yf2SYIfK5+cJ1v/tZv849/8zfp9Hp8/uWX2dza4JFHLvGlL32Jd995l0G/z6OPPcaXv/wlfvqnv8alSxfNNQOKvGR3d5fNzU2klPze7/0e9+/fJwzDI9pDbq49KeDRWtfzubDJ7cd5Qh6GMp0YMAlRz23aHo9nAx0ppXW2MChxo9kkCCOiuDLAgqxslUDbcpuphIjKbswPwLcef660b9u7l3Cxet10QZMjay8biLsgqapM6XMZcXLBt9aawWDAcDi0vLuzfPnLX8LzBEkSMxqNGA6HdTnOSRosB+XL5eBlOsYnOT5dQZPv1USzTqdDnCTc292j0prhaMR0SXrfxsSIj3nWXH3U83yEUOzu7rK7uwssB18aWUnmS2av86ltaRSu20HUC5zSGmGhITfp1RmB/YzZbn3IR0CXj3XoBgowgYyzZ/Bs279Si4DJBTofwxxWfxRY8LADO4aALcScTvj7gx9Y+r89EUtB2Unv/LjPQ324D3zwWHR43H3+oRs8eiRHDHeX4zBHqF76m0OilHLBSGEWTCGsbISy335BEq4nX+Dv//1f5MrVq9y+dZs0S1ntdVBKMplOSRoNvPGUDz64bsTh5nMO9g+OHONyHGeECAN83zPBmloEqlo9iJJVlWR/fx/PE0RxTJAZ8UJnruw0eU7ivf2zODzf2ZSYBKfRSMiy1HYOCXzr/VaWFXfv3q1VrhuNBmfOnKkTp2UezDK65NCj4/9cF2clFXES1qTsIAjr1u+qKknTjNFoxGw2R9hFNIoi9vf3uXLlKqsrq0RRxMHBIXfu3KbRaLC+vs7t27cQAkbjMf3+gFdfe5X5bMqF8+c4f/4cs9msNgJe1uNpd9o88+wzZFnG/fv32d/fJwgCVlZWAFO+29nZ5Vvf+j329/c5f/48X/nqT3Lq9Cl+7Md+jA8+uGGP03RvJklC0kgIoxBh52e36KZZtiCbBz6dToezZ8+wurrOPCv49ne+yz/61V/lW6+8QhCGjIZDnn7qKT772c/y2R/7MX77t3+bsiz5yZ/4Sb7wxS+wvX2KLMs5OBjg+x6NZszW1jovv/xjgBFPPDgwz5JDMpaDpeWfj84hS0jTUlL+w5TkTkKiHkSezL3kIWxZTbubCakVlZIotFEDj0LCKKpFLLV97pVVvXX5oQA84VvJkoUHoWtCdonV8qR30jlwAY0r1Tluk0OcXNDUsr6C4/GY3d0d7t27R1kWbG+f5vLly1y5cpWbN28ymUxqSQOXVLjtuO2naWqC+vBHXNyy2+mwc9in1Wzy5ptvcnDYZ9AfcOb0KQ77ffqDobHlcEP/8JHmR964fMiibTulFpARi1KbC1xYcFSOpB36If8/qZJ1LII4/sBpZTptjrxP60Up6eMMG2M5UKWGQJUjCx47pvqYjfZH7e79sXd5/Itz7PixYpvmfHpLpaCT+AIP3YU4/sJDfn3Y+T82zLlf1o96yOeWkIGTjnU5cwMbPAlH3rYTogfT6YzvfvdV5vM5VVmxubnFF7/44/zcz/08n3nxRfr9AYPhkM3NTdjfr4/t+H3teHSetwjiTdBv2qnrcu6RLwuj0YgwDHnxhRf4/vffsSgstkRHXfL7URhhaDS0ZrOZ7SqLLSl6wfXxPI8iLxhOTMt7u91mY2ODZrNFHEdLLdiGYCzEsr7QYsFxxrnmWtntC69W1V9WYTaPuhE6LcuKyWRKnue2tBQxHo/Z39s39AEhrB9jTrfb5f3336eqKlZWVjnY3+PevXtWQXxudXMO2dnZ4e7tu9y6dYvd3V1OnTrF5fASjWaDJ594gm6nS7/ft2rOJZcvP0ocx4Apr9y+fZsf/OAH3Llzh8cff4xHH3uU559/Ac8P2dvbr9FhP/BrnSfzrBsScaPRMIhaFBjivZU2aLc7JI0GeVmabrpdQyyfTCbcvnOHGzdv8dTTT7O1tcHFRy6ipOKF55/n0iOX8Dyft956mzfeeBOt4ZFHLvDMM0+xurrGCy+8wFtvvcX169eZTCZ1udCN48HM8bGYdh7sjnzYOE4DOKmUt0CazPyghXVZ1dr4YiqNh8D3LNInJZ7nE1hfUwFkQFma9cI0zKrak1JgGhCEsveeZ2xXHPrv1oLlY3JyGrAQsHRBjbM3yfO8BjfiOK4DH+cxV5albWb4Ab/7u7/L888/jxCmnBtZxweXOCyX45bLgQ7NckjXJzk+VUFTs9nCHwz58hc/T1VV/INf+mW6nTadXpfheHI0YIKPD0Ww0MSo0aKlhfhIQMIRo/pj+7I3r1tt6mjclT0WnQr1Z056bj7quJc+4/lm23qpfPRAUFN/yY95QixycLQ8Zl48rhJuFsijEZRwPX2WF/OHw0gf/A52OXrgq7lM68TdfCgsdTwC/ZCXxdKG7DX1XEaHePA7HrsPPuyaOnTBvNW82emMKblQB9XaWPecOXOa3/2d3+HNt39At9flueee49d+7df5whe+wN279xgMh/zVv/qvc/XqFW7fucPewWF9GhwQ6wufJEnQeXaMh7REfDvh/AsBYRSyvr7B+fPnGNusr/708e7Mf4bHxsYGp0+f4fr1G4xGQ5xWmu8HeF5Qo0JuYp9MJrZM931arTbPP/88cRwxmUxwnCFsUAROegMbgLnW/UUnHXh1ti6EcUdI0wwQdDodut1evf88z01nbhAuBVbWsFebbqQ0TXnvvfd49913kdI4E4RBwOrqKt12C8/z6Pf7vPfee0gpee+99wz6oqHICy5dvkSvt8Lq6hq+75FlGfN5SmTLb2BawoUQ9Pt9RqMRe3vGn/P8uXOMJhNmeY7wfao8t3w74/tZVhLh+bTaHVZWV+kNBiRJo+7cs5h+HawKITh/4QJf/PEvcuPGTYNgTMYcHB7QbjfYWN9AV4rTp0/TbDYYDAZ861u/z//4P/48aZry/PPP8/Wvf50vfOELrK2t8tRTT/Hee+/x1ltvMZ/PiaII+PDynBmLyeSBKeJDUKeTxkloFe52EdiqisDTmEYSrQnCwIhRIpBFiQiMBlWr0ST0jWDufDYlz7O6JOxyfiUE0vKWPM8zt6Yn8PCNK4TStZae4zEti626Mpm7l53FSp7n9etOc8t5FcZxzNraGqPRiLfeepvJZMrLL1/jqaeeqi1VHHEcOLIP31+o7rv7/uMn7R9/fKqCpvF0wiMXzvMX/+JfZGd3j7e//w6ra6tcv3GTyXiy4PK4C2k/t3yvep5nHy518npZRwuCD7NeWY53Hp5bHIeSFmhV/bkfJp44VsESHni+Z9vdT7QjfuCjH3eXy7wJ00q7eGgbjQbdbhcwyEM6n9to0hyLsmrWIgxtMPeQIO7Eo3n40ZlMXFjF9eUuHW05IS4AOWFzJz48H30mhF3JHNnVIDRioT1TSYQv0Cd1tj4sKP7ovSJ83wTZCvwgMLwEDX/lX/0rfOYzLzGbznj3/auMRiPefvttsizl29/5Nnu7e5RVxW/8+m8wGA6Yz2f1Vj1rHeUJQbvTZnVlFUbDj30Pam2en4P9ff7hP/xVnn32aR5//DGuXr3GfD635+cP830/nePLP/FlVlfXAXjvvfeYTMaWV5QQhiFlKUnTFKUUnU4bz/OZz017fZIkXL58mSSJ7ZxjuhVdCXexGIPNUo7c3BrDM8uyDN8PTNAlTGt8GEZcuHCB7e1t8jxnNpsxn88J/JDZzHA0lTTq5GD253ziTNBn1J6VBOVJPBFSKk2Wzfjggw/q7rwbN24wGo0YjUbs7u3y+OOPs7W1xebmplUKt6KbeU4a+DQazfrcGH/JKWmamjTLljqddVEQeHWA4AyJPc+3AalnicUleV7azqyyLtkEQcClS5dYW1vlc5/7LLdv3eHuvfusrq7WhtxSKsIostfFYzgcc3Cwz8HBAcPhsO4SXFtb5Stf+QnOnj3D9vY277zzzhFy+8OCJurLpW2ybMbDEKPjf6uTqGOB0vL7jmzHlQbs7aIdCu+2oxSlDS4iG8DEUWT4WUqhqgqpjXq9e7+WEuWB8qygpdWiwhr9GiVlfaRcuWyD5c6TK525f3EcM51OawRqWa/JST84tOm9996nLMv6ugwGA8AE38782iUlyyR0sN18xSc/IX2qgqaqyPm5n//v+fVf+3X+47/1t4wp6rxBVVZgCXF6adauA/Glxd/8wEMXiqPw6nG4xYzjClB1LGMX7oWnm5vgOBI3iA8LcE6KxI7szKIaGqPAXZ2gkn3CJsBMy6bbCVhyBD8JvTGmrebhDcKg9pJCw/b2Np/73OcYj8dcuXKFW7duUWS5+f6htdiQGm0tJh4+fjgEKo4jut0OSikGg2F9fl3D4ZHz8CHXeDE+IpQUTv7fPCZKKQQCqS0x101yH9UxdvxafsRxeb6PqiRoeUTeAOBnf/Yv1psRAt56+23+13/jf8s8zSjLqv5Gv/3K79Wf8d03tX9stRJ6va5BmvoPC3QePFBDU1PIUlGUFW++9Taf+cxLvPjiC/zBH3zbTp4CKX+46/ppHS+//DlazQ55lpFlBqVxrelG5busW6BbrRa+b9Ccmzdv4vs+L774oumECwI7PywkHoSwXZPY9m+9/M/d98oQfX0nrigIAo84jjlz5gxPPPEEWZbVfzs87LNzfwcwSs2OS5JlGePxmPl8jtFS6tJqtUjnM7I0tRYrkqLI2d3dYTqdcvPmTfb390nTtA7Erl27xubmJtvb25w9e5YkSepS1traGpubmxweHjKZTGrNHddxh6C2mDFB00IwEcy8aqY+VQsmOisOKVVdrjHnO+DixYv4/iW+8NnP8tSTT3P9piklemJR1g98jzD00VpRFPmRsttgMOC9997lueeeQYifZHNzg9OnTz+gaH38/0eDJ7NuHEmwtT4xwDppHA+cTvq7WJr4nMClQNsQXBv5ACtTozHrhfSlRaFCWp5HmaWUWUhpt+m+n+ukU55Aeka6xjSyaEN3FIvjOB74OZ7e8nrqDH2dAKVLKNznnQaTK8HFcWz9Dm/V8hBaa9rtdo30udKyQ5fcPe2QMfnHkMV9qoKm4Szll37xF3njze/XJ/n2nbtobQxFjDedWWdcYnbEkHS5TvzQ+3Y5UFp+jSOvnbQsaE0tclgT5k5470eqGtVPmTi2rh+DUWzgdHxzx7e+wLxsGzPLAdPJ+5elqonAZW4cr/3AlAhu3rzJjevXwdbJhTAO2YYoX+HKVx/9JX+4v33mM5/h2Wef4cqVK7z66qvMZgtfoQdKcw8LOh+6L33iS1IqXnzpOc6fP8+Nmze5dvUak9FkMRV+FMH+4wNp9ljNgoun8IXRS9JS4UfRiVmn1prZ3CjfOjG55b8B9QQnNcSRxzwteP/KDT7z0gtsbW1y6/bdjzgoM1zZo6qMMXRZFnz7299ha2uLRx+9zNWr135kAiaAtdVVer01Pve5l+uk4t69u5RlgdNU833DIXE1DyUVs5lpsf+d3/ldtFY899xz9YTleY5T5vTdJEoZt3pjwGqeXWPIHRBYLRyzaJSUpcTIdRjfOYCtrS3W19d59dXXmEwmPPHEE7RaLatQbtSTDw/7DAamM6nVahGGIbs790nnc6qqxPMEjUYTKY2Ip+OTeJ6g3z9kPp9z+/Ztms0GrVa7LsNprWk2Gpze3ubcuXMURcHrr7/G/v4+nU6H4XDIrVu3ieOY8diYqwZ+UC9+ziHBfEdlVc8zpDTioEkSkaZBTTQWwtyjcZwAikxq2q0GZ7a3mU2nJI2EVqvJ6moPXwiUFoBBwVqtVi28qJQiy3Lm84yyVKyurnP27NmabPxh4ygf6cHkYznYORr48MB7P2wfDwRSdsoWWpAEIZVWhuytC3w/MGa3WlMVJak0NiNhEBL5IUmc4AuvLndJKa2rhKDECC17gC+AwLfGvl6NjDpaSz2EUxH3bGPS4m+Ou6SUYjab1UiRSzjyLCO0wVWz2STPcw4PD+tGhmUSvukW9CirkkIXCGHFM4UgsN2rn/T4VAVNAvjP/7P/gr/5H/9NfvpP/TT/l//r/82owx7jYTj9QSHMTVQrVB9ZSC1hWamTF3m99IMQi0DG4+g6e+yaKG3IsCbSXYJRMQGPJ8yDrbXxCHJdMuoIqc4FR/rBZ+gkpOJjoSomB/k4H3CNd8vVSa2hKiVUC4XshXi0PmKmu1B7fdg+/nA38muvvc7Vq1eQUjKbpUfjR41RvPU8kx254zke7x65zCcETCcc6mgwpCoKbt++Qzqb2RZxQVGY7p0kjo/wej7+OCGa04ZL8tNf+2n+0l/6S5zd3mZ3b5eN9Q329nZ57tmnuXvXCPPN53N2d3fJ85yVlVUGgwHj8dQaVabs7+2RZSmJ5QBIKfnyl7/E3uFBPRm98sorJxDpT8YqpTIlIOGbYL2qzCQ3mY5ZXVvhhRee54033vxDnIdP53BUgFOntnjxxReQUvIH3/4Drrz/PkWRkSQJq6s9lII8z+pMudlqobTie69+l1Ontnj++eeRSlFJReTVuKBJcJSuPRw9IazZrtFzQph2bYRA2xIWmOev1WrRarUYDAY1X6TVanH61GlWV1cJgoD5fM6tW7dI04z19XUeeeQRAKbTOcPhkFOnT7HS6yLLEt+S0cfjMdPplLzI6XS6IIxQ63yeMplOAfCt3ILL+KMo4uat27x/5Qqe57G7u29KbUHA97//AzT/I2EUMppMKKqKjc11pJRcv36dw8NDNjc3CMOwDmhcM4Xxh4uJo5goDE1gEPiAZGfnPsPhkDhOWFlZqcuUvW7HniMfzx6fb/+5pGBhJuzb3xf0hONr8MNKc4vxIcnUAx/7YUWCT96o7/lEYUReFsiqopLSNg2EKG30kmZ5QWhLyVot/u4Jg9BopVASKipbGQHfqt1r6aGt/Z6wk4chn1cYKxl7Hm15lpo2I00A7vtsbmywurLCnbt3mc1maK1pJA0836u9N5NGQhgYeQeXFIBmPp9bVNAYUSvbHYg03DylFAQ++sMuyx9hfKqCpsCHg8GYv/Wf/Gfge4zHYwuj+HaGscRj3CRjF3fJg4iA1kZTyb35wVreYsfCbgxRo1lHFuFja0wllTGtPaH846GtirLLKBeZyceOik+KfU5Y58Sx/7vvVWvqLCU4CpstaHPjVaWsyeAu+Dlac1/U45fFxVzW4erbnyQxOM9zY+2wNPwlXytlM5Yj44dFeh4Ymjt37gDUJpRGR2dxLo4f08cbJz3R5rWyLHnjjdcZHB7W+jXtdos8z2kkxkvLlCYkZZkTBAHdlR5aYcsxPu12u144kzih1+3S7xtewHg85Nq1q1y/foObN28fWwjsbHjCiRIWnnelJGwAlWYZH3xwnc2NDf7Un/ppfvM3//Ef4nx8+kZeFHjzKd1ul7NnzxLHCbPZjMODQw4PD8116XYpipLJZEya5TSaDVbXVynLihs3b3Ll2lUOBwPj5+b7xs9waR8ONTLyEz5BGBEEhvytYCmIMEmW55ksPo4N6XZ/f588z2s18e3tbTY2NuqA+80330QpxRe+8AWeeeZpWq0W1659QJrOWV+/QKvZYD6dElgF5v39Aw4PDxhPJga5EILhcMjUNgSUZelEpa0PJsyznOl8h529fULLa2m1O3hewKuvvc73XnvdzB8CWp02lSwRnuCN115jd3+XKDIuDwf75v4t8hy0psgLJhPDi6qkMQk2tALJtWtXef/9Kwz6fbq9ngkYpGSl16MsJUVZImVlhGQDYcuD8kinVxgGRFFAGPqUVWEXd3UkSHLEaZdb21cfes88NMAS7lPH1qF6i678tkCZjqDKSwGMJyAIA+NPWojasSHwfcrKyjZkGcSxQRXjiCgKTXlTFGglqYQHQtrynkYKCZWH71VGhoDFOik8Bz5oKiXxhcb3fPAtF9SuIUqVlHkFnseZM9skccJwNGY8nlCWBVEcE4YBlZKUSiKKAgXEUcRqs0mRZ6TZ3Ho9GpFMhLBeewaZlVpRaYWnBULJj6ZO/CHGpypoQkPgwetvvQWAH/o2oLHQiADfCxAW+ThSKlguyx1HiDxhg6ljJRpX46uN2PTDSz0ueHLoy0OgVonC+LHah+/DymTL2/44r7nXH5rcOLkAUbcwH33dfFBpbaxThMukqIMkp8MjpVpwtrRTNT66Y2ca+0kApCcAhab8agO0GtT6oeDYj36v1pBli6CopsXZc2HQg082pVFKs3t/l937ux/7M75VpDacEMOhaTSaBqb2fVpJgzSd8dprrzJPcybTMaPRDAWEgUdZucnl4d/F3NJL2bYw94oqJVUh2dndoWlLQj8Kw/MEQehRlgXCilt+/vNfoNPp8s1vfoMbN29QydKgIUmE8E2mXskSBHR6Xd6/epX/7r//Of6Vf+Uv8+hjj3FwcGgWNDRxHJprSomS5rlN4oQgiMhz01EWBDFCeEipKYqKwHpX3rUZ/M2bN2m320gp2d7eptNtU1YF09mU6WxKlpmAYzqborWi2+uwtbXBcDig0+kQRVHtJxY1Gmyc2mJlfQ2ALDXoUtJsGm+wLCNpNAiCgDzLGI3HjMdjirK0+apGeB5RGKKAkeWpAMaTD81kPuFXf/VXEAgmwwl5lfPBB1f5f/0//x90uz3G4zE7O7sEgc+N6x/wG7/+j6jKivl8RrPVRmiB85K6fesWP3jrbUbjEatra/wLf/pP88xTTxF6vmUeaHwh8D2BH3hoFEWZU1YFZVUwn0/JsjkIaLVbrKz2jMyBVSw3ht2inoMMoqJtIGH4aY5Pthwr1UvDiY+a+YNLSpZfdZSP4wmqCRatAbGSVGVFXhUgIAh9lNBoVZGmM6RWBL7H6dNbPPPMM7zw4ovM53M+uHaVGzduGJNkyxXzgwAPH2Wtv1QpKRQIJfCjGG115IwAqbEBq1SFFAolNIEf1A0zWmnDhSxKsrKkdXjIam+FbrdHnpcMBn3yoqBSFUEUugwNhaKSFUKb0odpLogsF0oZBwMUfhgYhCsKDLoSeBRq0ezwSY5PVdBkNc/wfWwEK/CDwBCWOQr6ODqQqFV7zSJuzAqFgSQ941HTbLVQgKpKlFSmc8nydarSvGYyCvOwIAw3oaqM0qoQAuELZFUZtr4r4Z1QvluMHwJZOultJ712DIGq44gHAimr4/GQ3VeltXex7c4OvFno+Ni/LQXxJxr5/hEDpuNzyknxoEPqHvjAxyxZ/rBjGYj8YUidR4dDGpe3e/SG8SwSuCCDuk+ac19vSSyCWEdwlZWiKEb0DwcPPQK3icA/fvwP89oxf9Lug8eTDwFFUfGuVSD/URim8cMJ7kp83+Ps2TMEgc/u/i7TdEa/f4gqFEkzISahLCsqWaERNNst+sMB3/nud/nxL32Z7TPnAR9EgFY2oxdO2R+E8Gg0WjQaDQQYs1mrwOy4cI7ftLOzw87ODru7u/R6PQKr3u08wNI0RaNYXVsFoNlsoDClm3LJGkZ4xgdRKkUQRbQ7HRrNphXJPEAEAZ3eSs2D0cosclmW0bFlsTwvSNO5CaCKAu3QiaVkp6yk8UVUFffu3kUrRRIlCF8wm0x49bvfNTQAbfSpup0u7/zgBxzs75nzr2FtYx1VSfZ29rh14ya793e4c+sWu7v3Wd86xRc+9zKB55PEPkpVZOmcvChQ2qiyJ0lMFAX4vkCpikqWFsGCIPCJotCWRcHd+I575jogtSXw10CM5fIsJ6jOytONBxNsE1waTpB540Lh2uFRZn5wQsrLHWTGL9N0IXqeTxQFaKCsjK1Np9Ph8Sce50tf/hJf/OIXKcuCt79/jsYffNuS6iVZliJq/a8lVElgnAUqo9ruC98eK4CRzZB2jRS4QFKghEGBiqpElTnj8YTIj/D9gCiM0VqQphl+6BMlkbk/UAgNUlUILSw+YmyElBKU9rs61XIv8A3lxrcCn1rzydv1fsqCJqvgbpAaeYx0Y/9UHUfjjrSmmwXHBAPmJg+CiCRp4AcBVVUaUTXfs10WFXleUFbOhVzadnqF8s2DUknDaxEIRCAotemWceJiLjP3PNMKHMcxeZrjW5ExJU3k73kelayI44j5fGa4C4HJYqVVWq4qiSNxR3FgyZguMDkSMtoTsqQTVI+lJ1ofe3n5Ybcn/MijrG2rvzi6cD8w7INeK0w/cAzH3rs8EeCg6PooOLIyu8nKTRoY0qwfLALdI8Hose910rG4P590qA87fNNJZ47f7O5hUdryARy9PouypnpgG2ZxMJOU4+S5U+W+Xhga8muRm66jMPTtBLoQRlwWSUQbHl3gL/z6FBpVHT85J5yg5T8fD06X//YjNhZX1mTcUkriOObFF15Ea8V3v/sdBsMBoVVhzvOC2SylKivCMDb8szzje6++StJo8cRjTxBFMZUUpoyPJdNiFJQ7nTbdjul8nGc5ZVnVViZ+kuD7PmVZMhyNKPKcvb09qqqi11shz3IrChnRabeJwohHHnmEJEk4e/YsaNjd3eXW7dvc37mPVJJWs0WWZWRFQaU0rXaHpGHavf0gIGk0WV9fp9PpIKXk2rVr7O7u4vsBjz72OKdOnUJKya1bt3j33Xe5d+8eGkEUJySNJnmWkaYpsqrwPJ8kaRMGvjGHth10Ds0pyxJZFKR5Sjo3YptRFNWdeEEQ8Obrb3Du3AXef/999vf2SBoNeitrdNsdyqIkL0yn1WyJ/H769LYNSBt0uz0mkylgFN1d56zr2DuaKC2Vx9xzeewZOLGr7EgydPJD4/tGR809w8vdgU5XK47iOmh3nmzub2EUWo2rkna7jRCC2WxmOs6ThCeffJKnn36KixfP02q1OHVqy2g2zecUZcHh4SFlabo/6/Ju5BNahW1ZKYSWSJwUgaCWCndzmNGfQSBQUtcBnVbGWaMoKsqyIk0z8rwwfKfSJIF+6OMFHp4f1qtBvQq43XmGBK7BdDLbEq/WlhvmGTDjkx6fqqDp4sXLBH5oFxlpSkiYGrSrq2pXZtPCnmxvsYgLQVUUJviw7yurip37e3VEq5WiqvJ6ZXrYTe1KFUdGKPC9ELuU4/lBLQvg+TGd7hpbm6fYP9gnCSNaSZM8S2kmMY0kYTafsblpbAyiKCBJIsaTMUWRE0chszRFapNNdHodBv0BeV6A55n4DM/ASnYmV5WyC69nrAiERzqfW5NNrHy+RiNRviHuBVpApfC0QTuqWpnanEdp6+PCMw+OK2HW0vvuhtWAD34o0HIhmobG1JmV1ZgSwrTX20VHYFA8jZHwV1rhCQ+lDTJogkZFGJpJQdrV35A27QMjF0iZE2ZzGh5GGdc8uJ4NQkLPR2rTqiuEZ/ld9rbRug4QnR5UfQ9oV7b0zcQhHCG+bvg1x2xJ/gLzEBuNMG2rCAZexnP8FG/BWcAGkAKr1+RTVQrPtqUrrREWBnRWGlobN3lzzjRKqDpuC8IAVZQm1bBoqNIsYKfjN7i259AFbOqYPIE49n+HsP5xpHf/NA9331nUudPpcPnyZYoy5+Bgj+KKXdStjUqa5VRZju87Re6SN998k067y8XzF2m1m2h8o9CMa8k2XUEb6+usr60ZQ9yiMNc5COoyUOUJpCV+SylBSdLZjL3d+zXyG1prk06nC1rXZOCDvV0yG2jN5nPiOCK0CtxlVTEZDenHEVpJYwysFL1ul4sXLtDrdRmNxrz15hvcvnWLJIk5d/YMlx55hCgKSeII3xN0Ox2m0wnNZpNTp04zm02ZzWYkScx0OmZvb5eyrGxJSBOGMZ5nkKiyMublUdJAK02aZUxmU4IgxPd8yqri3v0dxtM5g0EfhKDb64FdWN94802UVvzet17hnXfeoaoqWq0We3t7TKdTPvjgA6bTKVVV4ftGgbrRSIiiCCEEKysrNZJXX3jg42QLP2wXl9OnOq5o7biizWaT1dVV2u02YDogR6NRrXvkgp3lwGc0GgFY0dF5LXJqGgfatNumHBuGIXEcH7E9WebdVmWFkoUhbnseSlrk3MPOqWb+1C5I0nbOlQofQRTGRGFE4PsWNDDdeBVmTpNKgRL4+HjCoG1m/lw+BwJfmPVOaZaMhReK5K6z7pMen6qg6dd/7R/R7XZqrKEeDu88AnM+mA5nUnHv7j36hwMjhJZn3L13n3fe+QG+8BgOB8ynM65eu8LOzn2qqrQLsQnGXGeF8AzCVG9aG1SoKtxrxh06iCIQHlVegvCZjifMRlMQPnMx51AdoKQ1ZMWswHfv3KKSxj1eqhJp/YB8u4hKrVAamzWYGq9vlX1VVdQImvB94laTTqdLM2lyZnubQPh857vf5dSpDcIgoChzvNBnnqXksmB1tcd6b43712+y1lkxImTzOQpodzuE1mW8qCo836ff75vypzZqx0pKQt+4m8dxwjydoVRJ0ogpZWW7aiKyNKXMC4tECfwwIPCM4F0lJb4IzM1ur2cYRsjK1MuFgNAL0dIs4lEYma6M0rQ/myDRr923TTnCalnZSMCIrbkgw+wmCAKKsjJdOMKjKHKUQ+qUNno6vrDbsqrDNiCXVVWLWwphuEWe59eKxotyqVq8b1lW3h6EWoo2tNK2C8VHKwkBpkPEF9aU2aBGsjTBoFYaT5vOJYfYeb6H7xlXeWNyWRIEPpU0lgmeX8e8dWC0OAB7WCdpwDq06fj/jz2W/6yP5ZhRaSM2GYYmyGh12ly4cIHHH3+C0XjMvfv3TMmq07WdZYUt65lF8MbNm2xsbPGTwyFxHNUm0Eax2zQexFHExvoGGxvrJI0IUChZWa8xExggTdIQ+B6+ECRRhEAzm0zI8pyyMF6H09GITrdL4PsEYcjertFvKoqcdJ7iB8Y3LwwiPM8snrPplEPPI88y4jgmSRLarTWajYQoDEErsjRlOh5RlQ2UlCRxSBTFtFstLl28SK/b5datWzQaDR577DHG4zGz2ZS19TVu3brJ3t4uUkmqSpoigfDxfY+8MKWyIAhoNZsEYUg6n5POU6I4JopigzQoxWA4IMtzkjgmCCOarRZpmnL1gw+4decOzSRGKUmv26UsS959913yPOfatWtMp9N6njcq7lPu39/F8zxmS96mH1WWP0n644cZTlJnucnG7Tuw12ZtbY2trS3C0IiWHh4e1u8djUa1YKmzLXGWJXmec+PGDd577302N09x9uxZDg+NPc54PK7fG8dxjXItNLFkXeFYaCIJhE0yPbuaoRW6snQWbR0rpCYQHkmYkIQRoR+gA+ogTSqJ1PZ+Vy65NHOhtZ6vz0dNRtcaoTSqDprE0oP5x2Mi/qkKmt743utsbG7S7XYIwsDI4W+sI6VRl+10OnQ6XYvzCKRSZHmOVIooToiKkq0XngVtMhctzaQzz1K76EKRFRRlYQm+ZtICHOwA2pgeAnZR1XZ/VvhQKQPpalfDMktYWUp27u+zt3eI5wWgNYPDQ/Z27zMcDJhNp0wmQ7IsIwx9dnd2uHnzOoNhnwqNsQBdrEmqlGigKCtEtVySMu/SVUU+KynTGUPhsbdzG4Egz+fcuzczi7EWIBRKVWg088GQHf8OqiiZDkbgiVoUjR2xpONgvmdVFDUiIg30QeWZ34s8RdlVN0fjByGy1JTpzARXYUgcJUZPRRtX8UbSoml/n85nVGWF1gGdVpdut0O/f4hQmlarxWgyptlssLKywnA4Is9zhGfUyPMip6xKNKaF2/OMjocnPPwwQmq56ITzoNQSUZosvLKBmkFjPPwoQEtpyrCeb1pZlTTf1xO2QuwiDoXGaFVJUbn0xwRQVCyeZrUkwro0oS5zsW3AUpsWu4/b37ULVAQIzyMMjWp4KVUdNJltGa2gRqNlJ9ES4dtAzPNM42mpTHB0vNT2YQn18TLdj1CwVI+l76+Upqokvm95ZVLSbrd54oknGY6GHPYPGY3H+H5gfdg8hoMxnhcQxwnj8cSUxm7dotlq0Ot1UdLwKRGGyxOEIatra2xsbNQlujCMSOKkRl6XRR993zfO9laFOQgCqtgEDJ7vG80juxgvKysrZZBc79Yto22UGCPVqqqM56E1yl1bW6PZbHLv3j3CMGQ6nRLHMWfPnTMaQGHIwcEhWZaxt7dHt9vl/Pnz9fY3Njbq9v5et0ez2SSOY8qyQVkWVJXVqpKqRkwEhjS+vr6G520xmUwYDkdUVUW706Gw0gLOUFgfHtpOuLBGb8aTCVop8ixjNjMSHS648H2/RpQmkwmvv/66NV0uuXbtGrPZ7IiFynL3MDwYLB3vij5pET9J4dtdD/d31yXphivfOSHSdrtdt9trrWk2m3UnoAt6qqqq7Wu+//3vU1UVk8mEjY0NDg8PefPNN7lx4wZ5ntdBsUOb3HUvysJKC5gKSuWXNRrtW90mD2ERIIuSKlNd0Ipa4sGR9X3PJwpiosjw/VSVoaVEoilt/BOFAaFNAnC0DLcfIew0qFHiGOH+2Ln/pManKmjqttpcuHCeOAw52Ntl/96QdDJGa02apYwaTcIgNARJ64sTRCHdlS7zPCOdp/S6PcaTMZ7n0Wl3kGXFeDymt7KCEJoyz+l1ewSBgWGNt5smabaMP9NsYiHMyJLPvNryAhEBmnQ2ptFsgTCkc9c9kD9dkKUFQntoZVpmszQ1HRtlYVybdQVSkRU5s8mEPEsNquUJ/Dg0woKeQFeV+b8NzNC2dKatCau1GqkqabWLFGVlkKB+/5BbN24xmU/rFs7+cMD+4T6T4ZBbN27SarcYz6bMcyOcWN96wunO6hMXygqORneALCSqtFmDvYmrokJVc7LUbt+WNuZzk0FIaVAggWA8GZlApKjQSjKuKvIyRypDhM2yrEb+tFJIJamrgVqbv2lDRFRlsYgFfPDDgKqwvDDPt/wxjyhpmOCwKowzeBKZbXqe8UD0AubzlDzLUZn5DtiFbaW3gu8HzKZzSlmRRDHdlR7ZPGWeGrFAZTMqIajLkrIy7drSEfMcDGbROFmVC572Iv5CKYl05bp68rZlS8zEnufpQlfL8tJUZWEkp9Txh5lfjgdSP0pjGey26EStbAxEUcz29jaPPvood+7dZX7lCpPJhLX1TXq9JpPxHMOrNAKOs9mMq1evsr6xTqfdqjNljVucTePK5uYmZ8+eZX//kPncyE+AWUhDa1qKDRB8a2Hhex46CMx9pkzJQlp/rhpJ0Lpe1LRSRvdrMsUPIqI4JonjOtlwTvJO6wnMQj+fz+l2u3Q6HbTW3L17l+FwyGRihCubTcOBAiPhEQYha2trrG+ss7u3U3NVtMYo3AuTEPi+b5IwBJ7v02y26HS7bG2dMomjEDSbTSaTCYcHByht5tfpdGYeF89bkJzzwtzyFjUpigUnallt2jkeuIBtOBwuWt2XrvuHLcwftXCf9Pnl8tIy+dsNN+fNZrMaGTvOd1pZWam3kSRJzY+6fv0677//Pnfu3OHdd9+lKIr6vO3s7DCZTGxZslEH2y4Ay7LMli5N0GNQKLMe+b6Hp6xyO8J0MSqryK0MyiS0rQRoUUsZeMLoeDWtcKoooVIl2vo4KqlQnjScJc88dELbqEmAZ0tzAmq6hVtnBO6FT3Z8qoKm9dUVWnFsvLNaDeLIs+Z9pnYfeR4eUCkFUhF4HpHvocqCWAga3Q55nlKlc5N5FDHpbM7uvbuoIsPzfbL5HGHJ3Vor/NAQK5OkQRzHjEYDEBCFMUVpBMKCKGDQ79PprjIejem0mzQbLQ76fWSlWFtbIwoTirKg0WiS55XhVGkDh2s0cRgQ+gKlAwJPsBGuIs5sk8QxUpYcHB5QKsnjTzzBeDy23lCSRruN5wdkaYHvBURJVJcvDc9roenhsqIiLzjc3Wdugw2lFdN0xnA0ZD6Zcri7T7vVYjafkeXZUucHOLKdM5J1iq9aWTdrYRA2rZTF36g7Qdx7lDICZe618XBc8xiEMAHe/v4+1658wGAwYDIeM5z2wcK+ttBlgs4yM4ifQ8TqYX9bFuhELzxuMJuShVxod9lropSgstmUsoKOqqpc/MKskgg/ROalRSIXNSxlO4cMx8KQ+ItSM5/5JpivSpQ95xptz5cJ2pTSdmGARWSkrFFpm/FwAIDnWwE3aQjwzUbTAj+CKIzMOXToQWVUdFvtFnfu3qXT6RqfOAFBEKJQlOVSKbIOyOwPy8zz4+NHLUg6PmrepJF8MGJ7hpTr+R5hFNFumzLdU089yWAw4MaNmwiBVc1uUNiSfrPZRFYVV65c4cLFC5w9e5okDvEC3y4EhsfpeT4bGxs8+eSTTMZTPvjgOoPBgCzLaFhft8CiA4ZcbZ7FIzpmFl3SS/8XQhBaLSanLJ9Z2YBKmuPrdDq1gnOWZUynU/b392sfsSiK6HQ6rK6u4nmeVRo/rIMNt59Wq0VVVdy9e5dms8nm5iZra6tGrsIh10pRlKUhm/shnW6jFgcNwwjPD2k0WibgWt+g026jgcGgz8HhIVEQMp3PuXXztkGebOADxq4Frez3NRICzsLFnRMXJOR5XgelZVnWPKMfpjx30u9wFKk6PpwwqPMDdNwmR0ifTCbs7++jtUHepZRMp9Na2HRtbQ0hjCRKs9ms0bP9/f0aeZtMJly5cqUOuoqiqMVIXVnOqKvHdWnQiPga0U9dWd6VNDQITwRWgBUb6LPgQSrb1KCFKadVJlgPQsOhajabKBR+LiiqzGhMofE9Ydg3SqHxrM2QXnDOl+FemyjX87n3xzNBfaqCpiSJqPKMRrNBt9shzxaKrnEU1Uq0CA/h+XhBAFpSlQVB4KERVGnKqa0tkjhGKI0v4PyFc7RbLXzPZ9g/JA4ND0QrQRwGxIFv9CG0NOaGSlEqUKpCBL4h4yrrFSZLqCpkmaPyjDIvyKIQHRn4l7IizXLDuxGCvCgIwgBPRBTZ3EDnScSBFbyrGomFlSvyNKXMc0bDgQkWtUGPPOGbdk0vIGk2jDeaNnVpg2pooiA0gmeVZDIa0223WF1dAU8wmk5oF23OnzsHUqPKEilLZFVZ6xRD2KuUI9+bluE4jomjiEpKZFmSJAlRGDFP57Q6XabDEVGjgdKaPMvxPI9Gs4Hv++RZgScEZVkxHU/tBGJb7TXs7uzy/rtXGE3GTKZj+oNDIzPhLR5mRwI3uiksUdqcX5dxh0eAb8ngWpgMxATa2gpw2sxEGNXjLMvJspwoiSmrkr3dXabTMXEU0e52iZKEw8MBd2/fpZKK+WRMWeRk2dxm2zMTaNmHV1VYaw1DFjbdHTa0Fbbjj5MmVvOttJJks3kdACrbYm1+VrX3khDCNB1YPpiqUTZIkgZaG+6XUZ9XpkwsQCths/tl9FCcUK/75+PoqMlqCM8zQagrqywtiKurqzzx+OPs7Oyyv39YBzDtdpvJZEaaZkR2Ybp16xZ37tzh8ccvk2xuEEYRSpVUZU5RVPi+ZH19g8+//DJxlBAEIVeuXOHw8JA0TU2Z2i6WQO3/dtwcddm2wvgr+rWxquPzgHluHDl4GZVypayqququLCdpUOQF/X6/3pdztR+PxwwGA7a2tijLkv39fdbX11FKsbLaI45jNje2LLlYk6YFURSzsrrKmiW/F5YS4LoE0zRlMBjU3mSz2Yy8yOludmi0Wgi8+rtrbZK14aDPzr17KFmxuWGCi8lkzHQ6q7+jQ26WTWCFMArsy0jQhwVHHyZa/Ifh2jgk05gVF7Xa+zJXKYoisiyrz5MLfobDIVJKdnd3awJ8ZjsXXeDsrr8jkrsSr7M9cSiqu7fcnGXOmVFOD7SzsgEtPMtpWiTcAlN1KcsSvyjMWi1Mk0NTN/B8CEqPSlUmSRCLQMikyotnTyPMeqRN242Zruw7tAGmfuSJ4EkUomSJ0AlRGJCnGlmWhHFMZC926PmIIMQPI/zAR8uCEiPEmKYZcRiwurJC4PtMRyMaccTGqU2qsjToiKxoJHGNXsRxZNcPQwJuN1u2k8wQ2JJGghDQTBokzSZr7RZ5ltNoNmnGMXmaIYRPGMXE9hg9NGEUIzyfoqqIIhP5z6amxttsNSiynGazUZeWVldXiaOI2WRCmedUZUUQBqRqjpKasjDQ7GQ8RqPwAhNMzqzLfbPRIAxDijxneDhg69RpiqIiaiQcHPbJy8J0migNsmQ+myzaXn0PaRcDqcwCH8UxAk3gwWQ8oSoKgtU1hFQMDw/odXpk8zme7zFPM+azOXESE4UBHhgiYBhQ+AU+gjhJTNlIQxRHbK5v8NTjTzCbzymlKWkqFKWu6jJUYR/sRjM2pGsnF+GbjgthOVlFWZhALTeoXqOREAiBzAvCpFF3OPqBmRiLrGA8mdJsG8j41u2bDPqHJI0Ga2vrNFtt7ty5yzvvvEuRlwz6fbJ0ytQGT5W0kLXwjM5Lzf9ZJgIthBUckuP7zn7GLMJmgpSm4066jrpFl5ZZCKhlFuxUTlGUpLO05rcUeUEpKzbXNgwaUpUEYYLvB8zTueVXLQVMdfS5TGz65+OjxrIflrCkZC1Md9r29hkeffRR7t/fYTKdMU/nNJoGdZ6MZySNCKkkg8GAe/fusbOzw+pKl0YjAjy0JfNLaTqnLl9+FM8zSVCn0+H+vXuMxmOjjl2WRn/IEnodedct+Kr2FlsIEDr0RLAIBEwAFRAniSldV9IkSDaQUFY80JncAvVCnOcmSXJlobIsGQwGDAZ9BoMBUkoODg84PDw0yaTQFEVBt9tjns4Yjyc16tHr9Th37hydTqfmVTlUxOwvqxFsV6JKYsOX9H2fbrdLu90mDEPKouDw4IAiy6iKnAsXLtLtdkjTlH6/z3A4rLvLXGB23NPxuHXKh/GZll//sFLc8eECGReguIDXHYu7puPxuA6mnNxFURS1Ir2pjoyQUprmJ3tfODRp2b3BBUhAjcw5VCoMQ5IkMfITmbUEsu2yUik8KZHSQwU+PqYMp4VGC98g6QITTGHuvyLP0dpQSDz7nepjUPYz2jRIaZvkaenkXUBrgRauw9gFTYLjhq8/8kTwKAyQZUk6nxLHkblkArPKaNO6qJUknxV4XoYfBIBEa9N1VRU5nWabUECezsnTOX4YEmjTXVKWJb7wiMPQmGMq47Pmurq0hsjeQFJKojAi9HwqWRH6PgEaaSHFMPCIghhPG1+5OArwbE3e9yITbPkBeWk4AkJLAs8jSWJjvZDltLodxuORyUo7HdulJUlCI/4VxrGRxi+liao9oy0lhE+cxEhZUUWBJdP5RlVVVqyvrdBMIvZ2dlhZXaMZm7ba0A/I8hRZZjYgMpO1Vhphb2DrSYOSJdlMEniCyXBoHsQgoMxL9vf3aEQJB3t7NNMpk+kMraFZGil8IQQeppzh4ObADyiLEg9jtukJn+3t0xwcTJFKsrV5irwq0NkcPwprnv1sNqPdXrOaJJIoNsGysUCISDOjfh1FEVJXyMpk+J1Gg/loxOraBr7wmM8z2u2O8T5KC9Ispd1tE8cRly+eZ57O7SSU4HsBT1x+lB976SWqUpHO50hVIKuSIsvQGIVege3eKyuCOAS9ZH+gdS3QqqwlihcYRM/3IzxrUKqUBqHqTF54PrI0C4ayyJWRNXCqxB6T0Zj9nT2yLKesKg72D9jZ3aHRaOD7HteufUCr20YI2NndZTKd4PmecUC3ZUgzlkuFi9cWry4hLfY1qRTXPvjgj2kG+KdxLBZMF0zUfnCWYOH5Hs1WkwsXLvDkU4e88867DIdj1tY65l73fTwhUDYQ2dnZ4YMPrnPu7DbtVsNuf7EQO9+tjfV1PvvZz/Loo48xGY9NEHJwSL/f57B/WJfNRqMx/f6h6VjVi7KcsJy3IPAJfXvMyjYp2LKMxgorVtKIY2jTpeoHRvqltCVAzzMLl/CE4aEo283qG+RYKssfKhflrrkNrubpnOFoYMp2UjKbzknnWY3UBUFAq9ViY2Oj9swz6JY8go4JYVDi2WzOdDrn4OCQNM1otVqsrKyYMpYNugCj4dTrsb19Gt/3mU6njG3gOZ/PmU6nNY/HBVLj8fiBEt1HcZqWfz6+iD+MJH48aKq7xqzFi0POHIrotuVQstFoZL0PVxmNRjX/qdFo0Gq16n05lNEFvm5+Oi6m6RoJwjA0AXklLYkIKq0RWhHYwFX7lrwuPLBin9oTOIsmKSWlzMiKkjAvajXvSpWUVUFVGW6vtp3jQpvwTFlFZWHbh+u0zv0sloyB9dHA9ZMcn6qgyUrxURY5gQeBJ/A8o9JalIagprUAVdSEYCkrI2cfCFvTryizjCLLEIAvBPl8ZrhMQhBYF+XSZujowEwGtnyj7Y1b18i1NjwVJdFVxXgyotvpglJUlUFnDDEO0EZW37RSKkCiZElZmrp1UZiMqydWqGSJVkZh17NcA7QmDA3RPUriBeFNa6LAp2WdxT3fM2U6WdHOWzVs63seU2B1bQ3fD+i0WvQ6HZJmC6kxWiwCAhETBr5BXqwGllKKNMsoS5ux+h4egtX1dULPGOVGYchoMGS110OWBb6AbrtthDxt1iPtA15VFUqa8+dhSoIeijhO0FpyeHhIEgfcv3vHKB3jkZUZpSyJGzFFZdpxD/b2SCLD8Smq0hI+K4QvaDQaFHlhyoXtFpPplCzNKNM5kyRhdHDIbDSmETfoD4a0Wh2UMhmv53uM+jHdXpcsz5jNp0RRzBTBZGy4Aytra4jIp53EIBSeL5iOx7RaLeIkobTeSZPRmE6va6+/VwdDQWCy/CLP7bUXZFlGt7ti/MXstcyytEaqNja2GA37KKmYp2byP3fhEdLp1LSAd3uMBkMO9/bxgxDPD+jvH9Af9On0enhCsHP/Po1uqyb7pla8rioKw3EzxmEIowJbkyqdrpRpUTFkTjMVLib7dD7nr/+Nv/FPclL4ExsGkDsaNNX8E29hP+TKQ1tbWzz5xBPs7OzS7w8RCOLYuLmjjdNBs9nk8PCQa9eu8uLzz7CxvmK3o+tFU0qJLCWBH3Du7DkuXzKowWg04uDAoDd7e3uMhiOkkhwcHHD//n2m02mt+J1neV02q1vaLXpR2UAoThKEZ5SaTZIW1cHColPLsH9dIBYEAUEQGZ6uQ908QRSHIIyoYZqltWK10orRaMhsNqXT7dJsNhmNxnX5T0rJZDJhb29vib+jahkRF1C12ybBKYqKwWDIvXv3yLK8DnhGo5HRJioKZtZYOLYWMfP5jDA05O9Go3GEm+Xm+SzL6Pf7dRdZfQ88ZFE+CWk6aTwMCTleAnTXCXgAcVq+79x1cHIFrvzmym2uw9Jt1wWbwJHtuNLqcpnO8ZuMb59CYSgSSis8pZHKN4KW9tr7nvWe0w4h8gHPFjMqqqIkL0pjtB6YkpyUJVIWpqPb0mF8zyMKQnwPg25pZQIlsQiaFBZl0qJORh1C9UmPT1XQJLWk2Uhqk9ayLA0jP/DQhUJpSRhEtqWXWkHb9zyjSyRKBv0BrVanLm3IuhXSrzs+tIbZ1BhBdrsdg174RkpcaSwMbKwFDEPfihjaz8ZRzDydUxbVUo24QGAEFSttMjFdSPJ0jheENFpNlJZMpmPWynXG4xFCGJPaMIrN4mS5Bp7vkVhz0DzPEUAcRHgIwiBABCZz9cIApQOydIasjHeU5xkydZ7lrK2u0G63EJ6PKiq0rAg8QRLFBIG1SxGmW0UK43JNENBqNQmjGC2lkQjodgiDwEg4tFt0ux3KoqAscrZPnWI6N0hT4Pt1SU1rS3o2Eq4IBI1mk063x3w2N5IFquL82W0rLucTxx6NZoIf+AyGQ9bW1mhHIevr6xRlQZbnFGXBZJqD1PhRRIQm6bSJopheo8lkOjUoGhD6Pvl8RhwGzKdjstncBCsW7p5Px2TzKUWZM51OieIYEOzu7OF5gjw9Y1rwNMznU0TgMez3WVlZIbLK7mEUM+z3CeOIwDMIoEHXKoLQcOWK3BDGwyBgPBoT2Y6VMAhBQGVV4WezKU88/jg3b96kqkpLyhRk85SDvT3arRYbG1tkeUFVFDSaLSqpUFXJ2dOnycqSlV6H0Bc02qYseXpr0xh2hglVmRGEIePJxPAYmk2qsjSCqDZQzquCMI7xhUc2Twl8nziMEJhyZL/f/xOaHf6kh2bZ2NrzjdhqvagJTbvd5uzZs5w+fZrd3X2kNIHPysoK08kUrSra7TbjyYg7d+4wGA5I09PGBqPmnPh4XmDQX9t95MoazitufX2d8+fPU1qRS0fYnozHzGYzZvM5h4eH7O/v14FAURQ18bksSsMe8YRdnBVhEFo5gAKlFgGctmzfsiwpq9IGTQb5cbyg2HbeBb5PlhnNIITxRIzjGKdWPZlMjAxCltWlPcfDGY1GNJtNYzAbhrV0g3ut0+nQ6/Xw/cDMNUFgNOqajXo9yK1OFWh6vR5xGDAcDhgMDmsLIt8mqGYf5mdnfK215tatWyciRsfH8RLeDztc4OICR/dvedsuQHKokEOMzHfv1HyzJElqOQcXNC2rrS94SVVdEnUlv2W1dafd5Eq7WVWhSmU0ygBpWUcaTNKlTcDsaYHWFtf0AhSCUipUWSGVRCgPX3omYFKl4QoLhUAZnrAODMigjcesEovuOa1BWf4U3qI857hTJ9l7/VHHpypoQkCn20EAeWEWyFIWJF5sFzrzNt+3Cs+eIggDw6WJzGSTF8YWJQx9irKCsqLZbNnMP7Q8EElRlNa41qMsTdkiCCMUUJQleV6YjpkwJApDyjwn8gMaSQuET5YX5FlOp91GeII8zQmtO7kxOfTI8ow0ndFZXaO3toKsCoaDPqoqTfnRM4EGSpGmKZ4QJogKQ+I4YjIZEwQ+jaSBr0CVJXlVgicocs/44amKqsyRZUEgqF3AZ5MpzWaL0PfMsaY5YPQ9ZJkyT82CqJR5KHzfJwoD/CQiSWKarRaz8YTJaMB8Pifw/LrTY2tri/FoxOpKj0bDiFxigy6tDGkcwAsCiw4CFvovrJSA7wui0OfCmYtorRmPx2htbCRAo/Oc9W6HThSRNBpMphN8pYgFxB0jxBnFMel8TrPVIs8zVno9mr5PKSt6vS7dVoMiK+j1OsxmU6Ig5tT2GZTGelpV+IFHb6XDmTOn0WBRO7MwrK6uUZWKIAjZ3ZUEcUQraZAkMUa3y5TpehfOMxqP6HQ6BEFgIX9JGBolZxknJlBqNBh1RkwmU7rNJn4QkOc5K2trDIdDSCKKdGZ4d2FIt9WirCpGBwfMxxPm4zEHe/sY89gYKRUHe/tURcmp06e5fuMG22fOcP/+fXor3RrZCuIQ3wuZTEdsbmyyu7dLHMX0el3SdM5wOCQII/w4YpKZACqJYmbTKa04oREnpkvPj+gPBn8CE8OfzFgWQVR2gvY8o8tl7B1sAUGbEoVrBX/88SeYTlM+uHYd34/Y2tpkNp1RFAWddgII0nnK3u4eZ7ZPs7m5USd5rihqUG+T5RdFYflwhoPZaCT0ej3TeNFoUFUl0+kFsiytS0/9/oDDwwP6/T6TyYQ0TWtV6dFoTJrOQQgqKU2i4fmGClCYphTXPStwkiqq7nJCK8tFMT50ge9DGNbm0U7dQtjuUVfeqyqHdEoj2msDBocUuQU9tB1XSRKTJI06KHByAVpjrKi8gDC0QZwtLbVaLeI4YqXTQSvJ3p4RdASD8jm5gSwz5cHQqrg75MVd9+XOt4cRvZdLasf/tvye5dfgKAfKBabHO/uO7/d4qdAFo04iIgxDQ1dI07qLcHl/y/9c+XEZKXP0gCAITDCLJp1OUEKD9SjEavQprZDLDlqewHWaaqs5JrWksvxUT/gIrSnLHHRFGAb0Om0ajQglK8qioiolZSGNSbLlrSqswK/WeFYgWfg+wvNrFEqdcG3+qONTFTSlVtPHEZKjyAiW+UFg2lVrBWaPODFcACVLoiikksbGYHt723ShBAFRHFmym/WEy3PSLKPX7bK+voESEDUSqjxHViV+aLyjlFLEScOKxxlYeJTvG1PMMCIIIjbWN00Xm1KoqqJhidhxnOCXhbnwpTEZlKpiPOwzGQ6MuaaFnJMkIYqMM7RSim63a8t4BVorfE8QJU2aSZMqzZBVSVUaTadA+LXcQGh9ohqNhCiQCAFh4JPEMaHvUwUK3wN8n3avSzrRVGWBb9WtS2mUslvtNlpr43GkzIOUpSlJnOAJU2YKAhNEVkVJr9NFA6X1ewpbLdZWe4RhyHA0QthMdjadU1WK0Wgf3ztEIJjPJly6cJ7JdEhZGr6XUhXz2YRmkrC2soLvCQLfA1kRCEHgCeJGwtbWBlEjIU1TQqFptRsMZYUvIA4DkihktdshSSL2dvfxgoC1jTWajTat1R6TyQSpJV4gWF1fBaGs71eXLMtpdZoIbVS/5/PMwPndDgBBs2lJtyFeu8PBwT5nH3+cVqtJr2e27QlIbKBUWkPopGECj0aSkEQx7XYDpTR5FNJqN0FWNE9vorXm/LmztFotGo2ESkrDIdGQzo3rvFKKOEoQnsf6Sg9VScIootfrMhwPef75Z4miyHTYWNR2b3cPQYvNtTW2t7Zsdq7JiwxZSbPQBQFpWTAcDul0ukSBT6/ToyoKhoMh3Wab1dXVP5G54U9iCKsxYxSSqzrxCgKj9yWVbRd3dkHK3EfPPPM0WsOVK9fwPM3GxgY3b94my3PiqKDRSIiTmHv37nH27Gm2tjYIQ+ON6fy7hDPONq1FSKnqRUpJQzIzooSm5BIEHu12i1ariZSrnD59ypRYLJm7KApmsxmj4ZDDfp/+YZ/RZEx/MGB3b5/ZbG45nywIxVLVdIXQ94gCH7SiyIyEhS8MOo6SZLaUrGv+m0ZVFenMlMqkdp1RCj8wBd+yNJK+rrMPsF1iGVmWMp06AnNwpMxkmXaEYWT0nDodeisrrK6s0Go26XbarK+tkmcp/f4BSWKCzLNnz9Lr9ZjP5xwcHNSkcIfkjEajRfmVRenMBU/LQdIysfnjIFPHh0PjXeCyHHi54M0dgyNvm/Or6+TVNfI4xMgN9/5lLtRyh2Bl9bvc9sPQ2P241/wgIExi5HSCstszXdZGlb5S0tIuDNpk63UmWKpKiqoil1VdyvN8rB6iIhCClXaLyxfOsrW1YTv/xuzvHdLvD5mXGZ7262sspaTUmsjz8GPjOIHnUVqtPvWjTgQvipLheGzFG03NtCgKKnszK6XJ8wKBZ+wKohglS4oiMwhRUdUk7k6zQeh5zGZzZmmKH0bc390xD1q7xdxqFLVabUuS06jCePV0ez2ana7pAKhK8jQjLyWIwnrXGTXxZq9HUeRMR2N6HdOqbmc5yjzHF4J2I2E2m7J77x5VWXHx0mUz0QioqpI4btQ8m0aS2DpzTlka1VYtK0b9A6hMR0SSJHiB6TTLS8OVMcJ1gvF4QpHleAiiKGF/b492p4MIQpSSjIYD9nbuc/7sGZpxw9Sq/aU25MoEq1VRIkujOlzmJc3EBIRREBpNJmE698bjMZWqGA365kELfILQI9CeIYR7Jjsuy5zAD1ld6ZLEiekGUxWdXocb167hiYAkNnpcaTZjMOwjhEev10UIj1argcLw1YwGFGTZjEoauDmdTxkNDvGFAY/n8xQlJM1eh3E2o7Xaoz8e0x9PaM+nzGdzRpMRW1vrdFa73L9/l2KYozDaNWVZ4ns+WgkqrfB1QBT6pvOx2TTdQZVkpdtldaXLqe1t7t+9jVaSRhLRSAw3JM9zsjQlCEJCP6AqTNBrgkEf5WnTLRVHpIGHwGhGFXmGLAvKvGEtfUxTAlriYXgmVWFazwNrp2OCrBavvv49up0OSSO23C2f4WhIOp/x9LNPMxwMiT1zvoWAOApQyggCFlqTxDHra2skUQJas2ad7LM0Y+vUltHA+REZhpQb1R5eNQLBQo+s7kgTi8W13W6zsbHB+toaRWEWorXVVcq8pCwqPN9DCM/ykwZUlSQKHVGXRVSw+AFO+knzgKo0GBkBorBe1F2nX2mToNlsxnQ2ZTqbMZ3Nmc5mllxtkKjhcMhoOGQ2mzFP07pbLs8ysjQ1tAW7P9dlZ4jBDlXy7bO/+AaBF1huiq7nci0tsZgljXv7gxBATXwua94dwihTC9+z4sGZPe4h/cMOjSShkcSsra5QFgV379yqgwdnPLzcYg/UdiK+73P79u0jBOnj3KWTyMcLNOnD76flj/0wJOaTUKrF/h4kq3ue4PimT0K1zHsX19BxnrRYCtysKa6wbhEO3VG4xhRwXbgaY+KrPUGjadbCsixAge8LOs0Gjchnrddhvdvh1MqKKS/3Vmk32oRhxP5Bv+4Q1e54lbINNyWhNTH2MRzkwPP5pMenKmhaW1uj1WyjdGVbLnPKSuILo0eSNJpUlaLIioXRn+eRlxmBrck3Ggmz2ZzC1u6NumyT6WzGqVOnkUoRR0b3QtpW7MCWSfI8M+RI4SG0Jp1OKYuCZqtJaOvI0+msNnzM5ymyKolCo2WUjUa2A8oYDpqs1Ey87W4PLwjxEIgwpNvtUeY5hSXtuVZSgG6ng+8HhI2EssjJJzN0abtYrOqXWzBFlNBuR8RNAzGXs4x0PmdlZYXJZEIlK2IRsrG5QdxpcbhznzgK0KqiqjRRaDg+ykLDaGNjEPg+SRQznU5p9VZIbDvvbDYjzzLCyIhstpoNNtfXCZLYoCllyXQ6Ic+NZ1QQNFldW8VDMBqNyLI5vW6P1dUV5tMpUWgI7GVZIDxodw2xvCgKFJpOu0EYhRRViR8Yg98sz9BCEwYhQWh4E92VnFKWtFot8AR5kdMQHZJmg5W1NaZpilZG8DHWmq1mg+3TWzRaTZJmA6UlSlckjdhw3DxjqDqbpaYEHCc0koTuygqdbpfhYEBR5DSTBulswnRiCOLLLb3G2btCScV4PLZCp4o4NhO2qkpC37PlHWOUubrSo7Kf8xy8b9E/3w/oWrK3skggmHsty3MExvoiikOroK7pdju0W0ZY8alnn6UoSw7+/+z9eZNlyZneB/58OevdIiIjM2tDdTUAdpMipRHHbEya0WiM+pajryMbmShNUzMaiqQZ1VgKqCW32O5yNt/mj9fPiZuJLGxdALpY5bBEVkbc/fpxf/15n+XrFxRGE5wjBI8tGtbrhCoqiqZhHCfqSojul1eXTONE8mKid/306V9odfhLDNnyU5Kg0HnT8f5ReSRFztKokILWOcqy4rPP/pqXL8Uc8tnz5xhd8Iuf/0Jabylxe3fH/f2drDulzBttchF27uB+tjn+xst7fObHTe8ddGSmDCitaFcr2tWK6/RUbFeUxljLNDmOpwP3d/e5tScqvbu7u/z37aLSGoYx815kvZgVe7OrRUwRnXIbMxPCVQ7VjSn7wUXP7LUW833ngmv25lNKojq8E/TN5JaeyYIdIbd7+tPE6XTg5vWbjAAptps1xMjDw92yvgMcDoclcsXkvWG1Wi2twhlxmYuG9xVQ77bS4G2k6Btnk+I3ipnfZ/x2ovlv//lc7J/Ph3cf8xz1UkrhY2DK1jNovajXYJ6SM+fo8Ulmj6WkxVZnvdlQFJbj4cA0jqiY2K1bNm3Fui7R3pH6jqauaHY7yqIiJYVz0q6dxkl8GJc0Bc80jlRNi1EivDJaC/f8Wx7fqaKpKArGYcAaTdPUIpFXPTF7GUlyvXhBjGOGiAvZSLv+xGZ7Id5NKTEOA6CoSiEJDl3HarUSDyBjWa3WuHGkO3aQyd1ay4U9ZYQgeJ9NL+UCG4YB5ye++vWXVGXJdrNePJh8VoGs12uM1Yy5ryzETo1B5PP7+wd8jAxZ4k5S2W9EczjsaduWaQpcXq5FPZWkfu+HjqZuFlWLCx4fI5uL3WKDP/Yj0yQeLjOEu95uJWTzcEdRVaw3K5x3i2O31hrcozfSvBDoUv7bO8frFy+4vLigrusF5q9yMri1BVVZSfacjxyPHZObSEmQw0MmHc/EPVNK0K01kvlW1RXjMLC7vOB+f8/gJq6vr/HHhJsmBj/x5cuv+eqrryjLimfPnoq0PqNj4zShC8tqs+F06vD55Do4z93DnnFyvLm9JcZE26wAxeyavN/v2T/cEaOX9mOOUPjVL3+FnzwffPAh3icqWzH0PRrFrz//nL7vRcZtxSleK1FurtoWYwyn0xFSZLNuc0Eo6NLQdwvR1fuwhCCXRYH3kZcvv2K9as94AxJubLSRCB0eNxbnR7pB7B3quib6gHMTu4udxFNoRUiJ07HDGE3d1Nzd3jKNE2PfU6xENGCMJNRPzksWX1L0hxMGlX1vHClEmZduoiiL9126/8mO+/sHdFaHqhycPRcnjzZceVOVf6AUrNcrfvrTn+Kc59e//pJPPv4Uoyw//9nP830it7e3vHnzhuPxSNtUGHtGLlbz2f0PGLmAU4/NPebtLAaPinOLSS9k9plwawvDei18oKsnl/zVZ5/i3MQ4TgxjT98J4uS9qNd+9etfcfNGDDdPpxPHw4HD8UjXdYy5qAJp381cmUXFlaS4VEm8y6wSjtiCXCRIweOT8MWslvZQmiM88ucnjxHzdaGw+SAtYduChszF0GwU+cUXXyzFQ13XCwl8tiSYzUPfVca9D3H6Q8efgH7zR4930au5CATAK6Y0G6Pqt1qTj+jSjC+d1e5KEfP/lAZbWAkQHxMpOlQwpHFknAbuh564v+fwZsvm6inFaktdlOJ955xwe1PCqvn7liczSua2c26htXzb4ztVNI3jhEngMqnSKI1GY7Nc1I0jCrkwjBaCpgpyoklIwnvXd5Jgn92gPZ7hdliKIK0l/qCuSuqqIfpI8gFba6ytOQ0dSsE4DpQ5cmBGgKZJEK5msxFUa7OmsoaYlQnNqqVZrQjBMY4DpERRSCtu6geGLIVVOSCxrmsKWwIyGSc3sbEbDg971utNRmJE3eKco7CF9HSVIgbpY5fjBMhEC5NMpFlVIQVW4Hh8IKXIRm/ROuL9Y4bR4t8xQ6D5MzKza2wSjyuT2w4ikR/wzrHZbiUjqeswoyzA3vkcSVNz6jq6rhO+k7FcPbmirlr6ozhrt5utIG2IfcBx6PAEggYXPZ7E19nhtm5bNtst9Xotqp2+l1al0XTDyJiT270PhBCFBF+3jGOgP3W40REKKWwKazgdDuzv7lEkytLK9uIjLoy4YSDlHKXSalZ1jZ+kCB77ntPxwMWFKBPv7u4IBFZtg9FqKfDJ7YT1qpUWREoodZVjgWJ2PTfS78+ZWPf3Nc4H6qamKEuMEZg9omiakmEcOXW9mBEaccQvy1Lc892ImwLtekXfD9RNjZ88wSdsYdhdXPD6xQtSjFhjsUUhJ7gsjIg+oJKW998PDEqT6go3CL+O3Fbph/7PvCr85YaQlGP2vzJIwR2yz9EZOfesXZaye3Fd1/zoR5/w4sUrfvnLX1GUBU3bZFWTytfokfv7e+7u7rjYbWitkMRTSo8n+rf+/u0jprOk+DOfoVnscR6porUWWRJvF2dzIv3MM5rRnjlUN8bI8Xjks6/+ittbaaV0J/E3mlt7h8OBYy6gJPfNL/J270WVPPvDkVv9xtjHQ9zc08sqvMUPivkzV+iElIZGNlGlNSY/1qwglvX30en87u5uKQ601ksG26y+k414WhDFc1Tm2zBRfHyI355n97sf5+227e8a39ROPP/dWygaMjVMRs+0zq3X+c/CWsuyhaWnPM+38OhYHgMmeznp4FCTFwsaPzGMA+54ZBoc5cVEdzxlXmBJ06TF+kbWpvmKyP/LiKB33z5d4DtVNLlhpF6vxS7/2AEy2TdzwnNIovKwEpQ7uQkf/SIdFd7OvagvskO0dx6FXCBi3x+Y+h4/jZS2YLfZ5NtJgTANA6vNiroqCT7QnU6kJItgURZM48Q//xf/nKQUY9eR3EQMnpQiq7oV1EBDqsVJ3FqLD5GqadhdPZFToNaEJHDzOIghnS2q5YIWVKvPcK5Mnk0maXvvhFmXF+39/T0gLcC2adhuttmzCTabDV9++QXD0PHk6TV1WzOMAz64ZSMoy5Iy+5ns93tSjKw2G5q2pTud2G63C1Fw9v/QxrA/Hrm4vGTsh2xCJ/L6tihoVqvsFyOu5hLKWWDLKmd3DcQktg6nrqMtKrrjkcJYqqrBO08MwnnwzvPZj3/MbrslIiGbLx9ecDoeSCGy3e7QSnPf9ey2WyE1x8TlxY7VbgdBYbTmdrzD9SPFZkdb1oR+Qlew3qwY+hN+GqmyieQ/+ezHNM0KkmZ/OKBSYLtecXt7x9On1zx79lRQtsJwsdvSDx2rtuVw2It5pTW4aaI7ndhsd1hbcNjvefbsGX2OQCjKksvLKyHLFwVNJe1n54QoDHnxGR0+jLTrNT5GRucxtmC1XrG92GHzST71GrSgRXX2olHKULdSkI3TyKuvv+b6+gkxBKqyXMKiTTKUZUVVNWhVUGAoqoJj36Eyd0D4byaHan4/hnNyncBvbpyPCqtcUOafoaTQtdZydXXF1dWlXF/ThHcuI8mPhcRsF/D8+VNW6yxuma1Oslng77dB5ggfHnkq54TmeZN5t600I03M7yK3c/R7Cq+5+Kjrmn/6t/9UCsmMyoTsl3Q8Hrm7u+PVq1e8ePGCly9f8ubNm6WoOhwOTNPANDmm4AlhWlpls7miMcLHmsOJRakX8vvIhHCtF+jmsShISzzW7I0nmYFyMJxl/eek8z5ztt515j7/nt/9zt+njpPX8bt5TW99Y2/xk775jr/v7X6f5/pdVgnp7LYmf9ZKgVGIEziP5dFjmTS/NrHcgUTfd/hhQBGxCiqrsClQotlWFbvCUgKD8+xfvub05oFjkBbzk6sr0FKsd8OAPxyJPCYmKKWoq4pYJELxPS+a5l5yYSxKgc9oUfCSI6aUhpjweVEJ3hMQN/Bh6MVYrTvx5JMnlGXJYb8n+kBdNbKBW8vtzY3I5LueN/f3NHVFXVeAZpxG/DSSQpVdSnO2F6LKOh4OghCVBd3pxP6wX/yfmrrkeNhTVVX2wbCEKO603TBSti2XqxVjP/D11y9Q2rCrVxxdhzYOW+ww2tB1HVUh3ks+BDrvqKylKRti7rm7IIhQ07SEGMWrJ0rvv6kq+vueqizExsAaqrqW01eWAHenju1qS98J92sO2Bz7XkibRQFGs3944OrqSnywqhJrC0Y3URuDLcXvpywKyqIStVhZMnrHMIwZnRMvFiFdQrcX+DuGiJscYQrc3d7TPHueDe38YlcQo/hrfXD9AW274fWrG+7evEEBm3bFs524hPennsvLCz5+9iEKeP3mDeMwUuiC7jjw6s1rtpstcfLouuB4fyB6zzj0GKupqwYDdM6DF3WHUYbpJCfow/HEw5sbirrh1evXPH/+nIvLS9w40p/E32ToOq6vrvBuFFRUafw0YbSckrp+ZDh1HPd7+mnC2lJiZZIIAIZxwhtLipHJe9xBsvqMNfK5xIQ2lna9xhTSCvVZ/t6PI7Hv0MZy6nppswHp1EmwsDGchh5ikrziKHlRx67juBeZ98VuBSHhBkdVGlRIFNqSXCCoiCkl+DelJKar35NhjFkIwyLJLhBUOKMZM69mRp6MIB0+CEetKCuurq744IMP6Lqe4+HEbrfj9u6G4dRRZrXu119/zV9/9inPnj6RAmHefR9TS3+vMaMxkAhRxCHv/l7PyLLWc6PrrfbMvD/L3jdvtGInojBilaNE1WqyV5VSCgpLXZas2pbLix0fPH/OT3/yY+FA9QOTmzgeT+z3D3Tdkf3hwMP9nv3DnsPhwKk7nRHOJ/qul9ZwRu3KsswIg/Cj5u1aKbI45Yy7BUx+Dtsu34qBmRG0R36aX6gMs5Ls3fFuofFNRdM/ZPxDkKc/5rHf5Tct8zije0mBReXoEiV2NuS8t7l4ytM0AUnJnlzVJbU2QuR2Dp0SJkUKBWtrebpq+XC340ldUirFMHl+fX/g5nZPnxTVZsvl9ROatuHU97x6/ZrTqcc5McIMzqNrxcXFBVdXTyiLkv/tf/1fvtXP6ztVND199oy6LNnvH6iqmsvdlfCOTif6caQqK1KQcNiUZJGylZhylVXBZnPBqxcvFkWZMYZ1K6ZlL15+zaeffioQ+aqhNoqylrydcZBe/Xq9pl2vJErAKFSSttEs05zN1qZJQnnXazHRHIZBEJpVg3MelAThkjSzUPj0cODzQw9GCNHr9ZZpHMSsMiX8OPLs6XOKwvDy5Svq0uTYkhEXPU2JcJ/KkjJKQrc2htIYcc9NInPv+4F/82/+DkXiX/7L/xNXT54wTj3H05FTdxI0bF1K+8/0S3yAhO22csrSGp8iXXfEFhofI9u6pB8Hbm7esF5vaJpmgbmHocP5kbKqCClSlCVNU6FGlnBJk43uyrKk73qRyxYlZWl58vQpDw8PFEZhUsQFx7Ora6qq5Ne//jWHV6/pjkcqW7DLad4pBqq6ZLVqcc7x8uXLRXKtjcG5CaPEMXfsBz7+6BM++Ku/4usvvhQjvaZBxcjh5obLyx311SXd8YAtCup6xel4pCxL/tk//ZTXr19jq5qLy0suLy8p12u6w57+sKddtdhC8sKcmyjLAuc82+2WqpIsp1BaPvnRRzw83BNTwFYFIThAEuGXJowSjkeITvh6RlzbXQyM0yTWGc2KPqfTF4WhLAuaVSOmiflns7CgbRuMsdRVzTQM1GVJSpG+75jGCe8jWkfevH7D/mGPSprL3QXT5AivA0nB6EZsVXBx/QRjNC9evfmLrQ9/7jG3qmbX5RhFDXreRkppJl9nXEhpwBNCxOjI5eUVf/3ZX/O//W//P7pOiqb9YY/3njYrHF+8eMH+cFhO0fAOV+T3HEvEBML5Oad7PJLCHwna+YbirBzjkgyQ8pOL8w6Ltm22GEiZ6+jP/ICEqy7P0dRV5vc9eysseBzHXBxJ9tz93X3Oq7vj7u6e+7s77u8fuL9/4GG/53Q6Luj2TJ4Xv6yYzSofmfJLIDFxue1cK5zzdWZFodZaRBozh/NMafiNn+873J73+Sl9V8a7hd/StlWakEN4tdLMjgKW86LpEQOdJ6ogrVBaQ1FVeKOYiDBNEBwKWDc1zzZrPn3yhKtSkCaPxiXDi8NAN4nn13azZbPdYOye+4d95jEJ+Xuec7vtlk8//ZTNavOtfzbfqaLp65dfi8/GdktKgV9//kuMldaOMlDUBcEFTLS0zQpbWMZpZBzEdsCNXjyGUsIY4fOEmF27C4npUFZzerin7455woiBnAsBvKOoCojSNvPZV0drxf3hyEfPP2AYBtZty3F/oCpr2tWKw+EgPXo0TbNidI6+6xfFRtOuMc5ze3vHkydP+OijjSi5jKVsS4iR7nDM6dQN/fHI68OBjz76iHW7wTnHqesZhju0EWWCMVLNl0px2j9QVDVXz54R3rzh6uoJP/3pjylLy939Pd6PdENHIhPeEe8khaKpKomjiUEMJ8sCP5zo7jqePr+mbVtub++wRjFNEiOjjKIfOp5/8JxpGBn8SKkNUzfQDz1t2+JD4P7uQF2tePr0OUQYukG4ZLWYsD3/+AOKuqAbOj7/1a+IKfLZZ59xeXHBy5dfM+TPb3IOozX1aoUpCw4P9wyDcJpISVqMlWWzyyamVYEtS/anjtFNDF3Hervh9f/331Bow4cff0jbtrz4+ite373h0D1QFeLKbLTmdDpRFmLwuT/sCTFA8PTTyP5XezbrDV134ng88Pz5c+4f7s9Oq5q+H3DTSHAS+zINPa9efJ2NEj2FhZvbN8So+PijTzkeTxwOe548eUI/9gsh2MeJlBRlVVLWFSEm3DiiteLqySVKRbruwM3re477ks2mEh5CShgV8dGDUkxjx1dffom1lpubW5qm5uLiYmm9ygHDUhrLdr2i63tCiozesbJyndnSULct20P7F1wh/rzDGHO2sSYJH838GpWR2yWkNMVFqSStsMg4OjabDZ/99V/z7//9fyCEyO56x8uXL4ghUhYF0zTx6tVL9vuH3PoqRTL+R7zeuRASmfY7CENKWXEm2XFLGyYXTedjtlHI/zj/hXD1tKZYWlgpU/geydHC4Qtn3D2V2/uSOyaty0uuri5J6a/z7YL46HViwnk6nXLI7h0//8UvePXylfjGZZFO1/WLPYgorc8LQ724W8+xNHN78dxj6Tx/bSarn/O+vonw/bvUaP/YxjcVdu9Fz4zICHRSBB6L97lgMsRHlAmWyl5nblOMkegD1hhMWTBOI2EccDHS7LZcNg3b0lKFiAmBsmp4ttvxsU8MN3d0bmK/fyBEzyGr4JVSlEWBycamMUaM0kJNKL7nlgNKKdq2ISWxG4hKctcSQkqbSWdJgQuekCKTm8REi8devlKKaRpzqvaWupU0+9PptNjpF1UlwYjWSuVsLSEmxtEha4hAtUVZoiDzfuQEOUYhpBftCl3X2L5fLrjRezCGSjdC5rWWumkgQtuuljDFsRffk6ptSUn8gTQQXMRqw6effsp6tZIcurKiblppQ/Y9MQZ88ExuoKoKKXTcxJuvvqQfJ/6rf/Wv0GPPy1dfY4zm4uKaRMIFT11UHB46HvZ7KUTaWjbFwuYTmCKERFmKK7n3jrqtmZwjEfnRp5/y5MmTbNEwgoYPP/qAdrsi4jnc3TEMA0UqWbVbqnKFtQV9J+3Tkz+htQR/fv75L1i1a5SG7cWOQhs2uwtUdvn96OqKcRyp6lo8bd68YbVqxQG+KhnHkYf9Hmstzo+00wprDKObiAp2uw0fffIhHzz7gMJa/o//+PccuxOvX78SZZgb2V3saJta3IonJw7muQVsonCqiqqkbldgDYeHfeaYbSjrit3Tp7y+eQWwcOge7u8Zx4Hr6ydstyK9FTKuzOunFx/xUVXw8qtX3N7eslqtaaqVRD/UGqUjITqGcWAcpX0w7T3R5w0qa6KsBVJAJTFYdT7gtSH4mB3d9yglGWPExCc/+oTLi8uMhrlcCEhy+uVuSwqe/nTCWp35Ymt0UeD8tHgTfZ/G+ab6SCA2qDMir84xEsHPOVhpsYqIIVDXjViprNdU1R11LXwzk1W1zo3s957D/sA4jDkpQOKc/tB9OKRHM8NZSTer0t7i3ohPQBYsyM/Of3++jn7TeN9ceFd19r6NelbdCjdSY2ZPpyhh0D5LyyX813E6nvirTz/l/v4B5yUj7nQ8cjgeOR4O7A8HDvsjXdfT912Oe5EYozlDbi6YQNDDOWJkLp7gMdh2UUaevZ8ZiXpXZfZNvk2/7/hj7vfHtgPPkaXz5363cJqnjlaPjCUN6ASGJK25dOasNbeQlRwigpsIzlEYg04RFQLaB0wI2BAoQqAKgdJ5TIgUZWJdluzaFnu3pz8ceBVfUe1rRidK7ZQSthC+W1SyNnVdx4sXL7i9+fZjnb5TRdMm+xOF4BYPjXMI1IfswAskL3wLsmxVabtcDG12bZ5PFlqJvLTO5pFzkjQ8KibmvJ2EQmlDaYT8TJSQylXTCoSN2AgkIiojRA8PD3z4wQeSjTaOKKMXW3utlKSPhyD/PTlxEQ9RFgwlF7UtCjarFZObeHItrak5psCHzFtCQn+NEafiStdMwWOKgqQU/TjgnOf1119RWU1Z1Uxu4tR1aPtoYlaWVoolbSRvTclnG2IQM0XvCF5c1d+8eYMpLNoYhn5E5QiJYRyWhWYcAw+He7SBsigEBTQl4xBw00QK4F1kGh1NW7Ja1fRDh06KoT+hlFsWUjcMjMeRV69ecXF5yTCO+CCRJHOWmwQoy/coeVSGrjtx8+aW6+trirLk5vaGqmnwk6eqSqp6xbNnz+lOJ1G5xUjSmsJYyqKAFIlassIiScQEVUk/jhAU3d0NSYGPfpHU+uA57h+oK+Ft9X3P6mJHu9twf3PDOI3c3vnFuZfMv3h480bmgZOkeKsNQ9fz6uVLXr35irqybLZttr/wxKgZp57CVpRFScgtIaPFS8tmg1Y3TgRA2URVVISmpaxKiqLi9HAUJapSNGVNXVRSSLuJ/nSkPx7QKmGUEhds7zLCqxmniI+RqIQY/30Z0+Romjla43yTZDF4tFajk6h35yJLaY01CpdvU1Ulu92Oh+0D1so1V9dC9hdFGZxOHX3fs163KC2tiEgivWePPKeGP7KYENPYM3Xckm+WhSPziElEFinzppR6bMHJe33cYL/Jz+e8aDonnr+7KZ//ba3JHnm5GI1JCvL0iFaBrOd1VVKWBdvNmk8//XRB/SS7rud4PC1qvTdvbri9ueP27o7ueOTUdbx4+Upanvn1zsXRee7cu+/vPP/t/P2fezadv+93/Y3+3OObaq75paT3FMSP9/0trzsJOGF0bkWmuXCKGWFSb6FM+U6QxE/Je09A2rlFkkBeGwKu6xgOB1zTyL4dQTmHSpJuEbxwgA+nE8ZaAgkXAyFBOQtelLSQ9/s9x1OHP/u+vq3xnSqaQAjd1poF5RA7f/lgXCYXixp1dqVVuf+u3prUc9ESQuC43y+cpNevX7PZbBbPEckyypb0ZYnzgbKs0IjkfhxGghEZfnAe7x06xysUZUn38MDh/oHr62t5Xmul3TeNODctZEMFQvwcJFRyVpSdboQc3bYtGM1p33F1sWOaJu4f7sX2oChxXtR9RuAFur6j7zu0NlxePUEZw+g9l8+e8/rrL7n45GNCyCGbRYGxOpOOBwyWpm0w2uQi1ROdqPgKLQuKc074Sv1ApWq2qxWFLYiBJb181awoi5Lj8YAbPUon3OiySeMTDocT3XHko48+ycGie6yXzK7j8cjz62tubm5Zr1pGJeZ7Klvtb7bbRd2y3W4pq5IQnNhNFKKucV7Ud03TsN7t+PLXv84uzRB9pKlrgjvy5tUryrJh6Lvsi1SxXq2IsSV4kca2TUtVluI4fpKAU3EfH8Q/bBpIKs+rStqLwzBQ1zPnRbLAbm9u0FqJbcBuS386iSkbsqn1XSfRMHk+T+OesqworGUcBrarDVUt7uo+eJqqoapaTl3PanNBWUhQcAgTxohviVGa0paE7D6NSgy9eDg1VU0/DLx88VKIzXVDWY6AIgSxqCDJwcBoieQpipKHw5F+GIgkur4nKUWN/sPhj+/wqMqaaXILB0ZS1aX9JO1sOTSkhUwdiDFgC4PWhphkEx6HkSbzIckWF3VdEYOY+BqjmNyYla0eqwopls5J2fxunpNReiZWLRvbkmXGWdGXHydnn6J5bOud/w2/WTgJL/hRkTa/snenxSyLny0C5jHlyKW339nj4y4/N2Q52iMHaT4Ul2XJarXm6upK1rR+YBjGjPQGxmnk9Zsb7u/v6fqeh4cHbm9vuc8u59M0vR1gnNc6KaQkNPldy4H570eu1GPxOL/HP6Ruevfzmh/znU9x+Wx+c/w+T/abirm3/abm2zx+H7kDi9bij6XQpJw1qFBoZYCY73x2v/xCq6qmqliMoeuqQpclwTlevHlD4TwrW2A2W9ZlhQ+RN13Hq7sHjkMnocBaoeYMuyyOiVqhrKWsapSSTM2YwIW3c/a+jfGdKprGYaQsBU2Y+9whBOmlGpO/tOzLFPNk0EpO3jxWzjOxeUZWTuP4lsvrTOyeHWBnuau1hSySWkmkSEahJsSRXNKco0SFWIPPdgNlU6MLKwWK0RBy+rIQBx77v+/A4TMpcz75DDnz7TD7OSmB6e8fHjicOna7HbvdlvmCnb1jfPAEHyiKgu1uy93rV6DEAFRkvCWSU01GUSzBzzl+Em0gVwugEPuA3JpcbdZUTS1t0yyuKMty8fbRlXiszEiKm0QdF7zHaIu1kXEYRT10PFJWliJ7cM32+NIKlE1EolEU7aplHCVUsyyz18pM/s/+ROF04nQ8sllvltfgndgcbDcbmqrGjQ5iZOxOKMQB3SqTN43IcX/kdDhSNxXF3KJUEp4avIT1brdbTieNi566qel78aXx3om5X3fK7vLwcHeHVorrp9cSwWEnCbvNLvSr1ZrCSmRM2zYYVUv8zGZD8CPt6gJbKUY3oCYlzvC24JR6UYyGgJ8mlBKVi/MTLoFK4hJeFCUuOFIQXoE452uauqI/nQhu4nTas1qvsUa+A2MMVVmKXDjJ6d8Hj5qz6/J8tcawatd/ptXgLz9u7+55+vSaoqgAxZQjnLTW6Lz5COcpLA7W4hEUSCmiFaTkUSnw4QdPcdPAi5evMUbR5nilmQc39IO4Jz+5wtgyb0LpreJmhpjmU37Kv8tsgqUtmOZdNqZsI/C2B4/cXOVKhXc2UZiLoHkDnYfKT7yoB+db5w30fYXTfB9QpHSG5OTfKR5//1gcKEgqk5CToAkhnK1POtulVKxWwJV8LimjICEEUYceT5xOJ+7v77m5ucmu5ntOJ/n5I3eqw/uQI1vGt5AlQZ9ipi2EXOSaZW8RFEuh9YyWyef37ni7sFKQ45AeC6PfbIeeF2fn340oCFnuu3xm/Ob3OP/uET3kLVQvvVX8zLeV+/ocJyXtaEnQCEminoxSuYU3v9Y8L4z4kNmyIqKwZY3W4p93P46wP1C8ueUUYde2TAm+PHZ8vT9wCgFK8Y+LuQOTpFcoQEme/DEl8c3Tipi+50iTUixGY3PLbF6sS1tIfz7JBxiSZNygIHhJVJ4Tq8dxZA7PLDIqUVWVkLiz51PTNL8ROzCOI9MwUtiCFAKK9NjKSRFtLYW2wnshMnQnCiOkXF1YpvsRsmOy1RKiO5/2fA7iFD+cannO3W5LUVVMw8DhcGB3sePu9oakFav1ihgixxvxOVmtVrktpbJSTKwUxNcoUFcVfhgoy4Lj4SjxDFWZfZF8XgiFS3A8HklResUyEWMuQOUUYQuLj0HUatZIq8g5gs/ESWPohwGjLX7yUvyVlqasSe2G4CNVYTFtweFhT9f1WK2pyxJrLOvVSiz6g2wu1mgGJ6afxggK4jWQQx99cAupVFAgac1WZU1pSw43d1hl8+lPitzueEKlSFOXjMNEXbUUthBPma5nGHrGyeH9yKk75sT0gsKWghhGeT5bFJKFqDQheLruxDhOOZ5hYHd5ic1xNLuLCzHYU5rh1JFiom1X1HWTFZrgJgkJ3W22bNdX+BiZnMcp4VBNwRGCy5thktfa9YCGhETUlBZVWlJMYnQ6TML9Ky3H4SRFvtLCydIFT58+pV6tJKfv4Z66qlg1G1KSwrXvu3wgUaSYKKy0LYtKENXJS+6fendn/E94/Lt/9+/5V//df4ctKoIPhAimsBRliQ6GcerpugGtE2VhKcoin4JFvVtUNVoZ6rrgb/7mJxij+PKLX+c5WTP0p7zhKg7HAzd3dzz/8EPa1Vp0azG9vS2mb/hvWPa9NMNIb6FUjxt2euf+C4p19ngpqeVnv4lyZF+m32Ovks35HCeDxehrfr1vbbzve5B3XsBcJOb/e6wd8msmoYymbRuqquTycseHHz7PB3C/tPj6vhcF3/09Dw+CQN3d3fHFF19wOBxwzmWyeZ/vJ6bAzo0Z+TJLoSMFpHzo7+M4yQH+saiWkN+4oFSz67bW54gQpCwumBWDkIhRCrS5EHpE59J7nv+xACJnxc3F21w4yZ+4fNdKkbsPkaGT1Ioq54UG7wQJNZaqLFBzZmuacxgTyQslpqpbbLsSwDAE0JqA4s04cby944vRsWobIpq9c9xOE84WGF3gYhCCv6ilJNbKCkUkEuWwHYOYPf8JSpzvVNFUlKW0QsYRcsGiMLmNZCDO1X/+knKrzgfPOEmxIxC4kHLnfn6ygkC5LFcUye2j4du5YqLMm3pSClNIG0pCgue+vGW/30vro6mWx+0e7hn6jrKUNo3P7rnzCfQRVamW9yjE9MA6ExBloa3YXuy4vbtjGEdWbcsnn/yIZ8+lvZiy26rPES9aqfyY8tnsHx6yakRlN3OH0glj5YjhXMCYia7vhdcVxOVYCks54YYYMIXleDzRNg3jMGIL8RHquk5eZ1GwXm0IQbK5CiN8BaVg6kdRkLmIwlKUFc+3O4rCYEzKYaBHrp5eY+xrAEJw+QIXFeJqtUJrzddff82TTHCdHYWFS2Ipi4rddpd5VdlvKgT6cSAEOc1BpDse2O+PrFYbyqKCpAg+oJTm+QcfAOLv1fWdZFOZgrKs83cSc16Wk9PmGKjqmvV6Izw1a8W3aZposms5SXLkUpKQ5VnePI4jh8MDpChFWYh0fQ9aM41iM6BV5NQfCUHyy6wRFHK33WKLkhAjxiiMEa8ctCUQ8T4JF1BbHvIcGMaBNAga8uLFC/6r/+5fQQx8/fkvIUbh5WkJnAjBS6J9khbIqm0gz11jDMpLQXt7++0TL/+xjn/37/4Dz59/yI9+9Anb7QV13coc8LLJ6HziV0rl+KYoporRi9O7SigSGFitWnbbDWVZYMwjqpIShBDp+4HTSYjMLLyRXPK8gzRBpp2kx6IHNSMt8FhZzIjS2ZvK9zvv+f1GU+gbOj9v1wK/b/H8NorBssmfFz9/zEjf8Bjyj7l7IOu6PPe5VcB58XQ8Hpc4mLu7W45HUe4dDgcOhz2HgyBSXdctbfm5ree9XxR8M69rbtsu7dskBQXIZzivYcKzmlt86TcsIua/zw033+WNPSoW41vvdd7fZnL3uy3Xx7bc/PPHSZEymjTbl0jiLuik0Rjh1gZPUAmd0WwFpAg+RhH6FIaIYcjpE5gS1jtiNXH0AZcUpp8YnWOKAa8EVVdK4QGfeyNLWza7g881tJrf3x87fX7L+E4VTW5yuIzUaKUhf3kp5QTmIBtA8GKkJdb7BqsMQ0gEFyBCaUsIsmDMJ3GxIIhLJIpSIv8WrxqNSHEDMUS8m8RIi/ToIlsUxCgkcGvNUnChckjnNDInAMzZZgI/p/waAj4EvJ8WaHtW3B0PB4GgtSSf73Ybthc7CSmc/aE2W7xzdCcJ4S1MQSJyOJ7kPZcVznn2+yNt29C2NdMkhWNVlAJzeydeP1XN9oIll0/iYYSvk2IkefnMTqeOwpbc396z2W5omxatesmE6zp2OyNFhY8kqxf0JoSESprNakOz2lE0K5SGqTuiCIxh5P7ujk9+9AnbzZabm1uqql6k3Lawi6XAOI64GNleXrJarzns99ze3qFQ7HYa70MOs52/E3k9hSlomhLnPQRBxozSjIOQ5efsOoGgFZvtjna9pizvKW3Fql3jQ2SaHFVVUjQlPpNW5wK9qSqGaeJ0PIoxZUbARP0kczY6B0kM+rbbDcZoxv4EKGKCU9cRgiycVVPgos+nWDn6LdyZ7IKvjaFUReYYJEQMKsTjaZxIEYqiWopDrTRUUDcN06nj4eGOrjtRWIOZZMJ6J6iWmb1wgidFRSCBn8nEijqrFr8v4+9/9jOqquKf//N/wY9//BOePntGWZRAQCmR3ksU0XwA8zgnMRyFtdK+yJVMStIW32w23N7eCo8w841SSiKlzwjVvMmq96As81DvFBuP/363ClK/9X7vG+dF0/uBxW96kHe3sPfd7gzx+r0Lpvdtjb/9zueu3oLSRM4Lp7kL0bYtFxcXi7XETN3ouo7j8Sj5lPv9Endzd3eX0akHuq5b2nzDMCyUiXMSPrBwn94mxdtMQ0lvzQPgreJrMezMr3vmdZ3bKLyNbL0tApiLxPl274uIOX9d8+1nAIEkX5Y2BqUNKgScd0whoGPE2rTQS0iy14lCPeCVqIVJYKzFVjW2rnNnKNF7TxdHQg5jlrJNBE9pfs9aS2iv0fJnLh7znAi/D+T5B47vVNE0T5Y5FFMhNvhz9IgQhRUC9GmMLSjqksq0kj1kDC4TN5f+cEKya5SibloOhyNl6Uhpyq0gw3q9Fr+ayS2QeIhhOVkURcHl5eUZGlXkk8pwVumHJevNOY8iYXXB7LpbWrlIhWs0EANYY7JiD5x3vHr9mqurK4rS8vrNa6psi9D1PdZ2wjvJEl054UgitDUFjTEYUyxS25TSWzEQolhJ+DBRekF1+qGXU3J+TOd89lKRk1Pf9VzsLkjA2A9UZUWbH09OxNJO6k491q4oS0MMieDFg0ZrUXT1wwREpqFDq4gpFNvthq+/+opVu15QMR8E1Xl4uOfly5f81V99xma9lsDh/R6S+Cg1TUtVN9RNw+3r16SUWK1WSwL7NE0Ebylsy+3tG55cXtM24v/lJs/Dfo9WYizp3EQ3jLSqkf0lKgpbUtUN2vmF65DGxMPhgXbVEGKg3/ccERVVCGEJM9ZaPyqjvCd4yaKbW7w+E/rrqsbokr6Xk1hZVQQECTXWCg8gBCbnqJRhHD0oKZ7mLCexDPAUpmK1WhOC5DY9uXqSX3egKHTesAvu37ziYf8g3mdlCfm0VtYlpERlSwyG4+lISJGyrlHWMnlHCKLw3G23f7b14C89Xr58yel45P7hgVevXvPpp5/x4x//Nc+fP0UpOQTJWvPYPrH2cclNKWU+yKyQtVxcXPDVV18tOW6yocp1LOIU2QTmhs8fPn6/e808qN916/ejTr/vKzuDxt65729Qcn5j/K6K6v13nuuHd5Vh4h31yGl9txB5RKWk/VjXNZvNhidPniyk8b7vl2Kq67qlWzCTzGcEai6muq5bfnZOBQGW+Cp5bb+pRJzfw7sF1VxwnYue3n0/54+xfFrv+SLn277bUpxvG2NEKyNoUr6dwGFSSEmGoEfI88XCPR2G8fEwMMhBsygTeglezt5gxlLUDTYl6f4YDZnna0x+b7lbNO//82ewtC3/BGyB71TR9KhEeIQIJaIgw4w5+FQbI8VQnIl5j4ova+2ySXnvKcqS9XbL0PdURbl4L3nvl8c9n1BS4Vs4a61YaygKm2FYabvIRm/O8oxMVutl9Z4tHl9LQnxZjCEhJ5nBTXR9x2qzYXN9jXq4J758yXqzwflJ5OvrFVVdMQ6TnGCVSJwBnPNoLdydpmpYtStCFKfwMrc5t9udZOGFCYJI1FPwHA576nr1iJblEfJpq6hqSi18mLqqefbhh/TZbK6qKpq65hSz4rCW1pjNMLQgecI7Oh16nIuURcl6vaIsLD6MBOfZbNYc9gf2zrPdXGELS0w+b+TZwMxoUdGFkInlmrIoKcqKpmlRxtDkOBKTT19ayUVmC0tZVUzjKC7wk2xqKhdoKUnRWpYtMUrhpzKfx3lP1/VSUHW9+Cy5SN91GJMjGGLEpdk2Qi2nOJ3NJVNmhRpjKHJxfDyeSElsBlTOVvIhoE0hShESTVUxuiCLB5LPF+ZWmpJwaElvl3gh5zwkIxu4tcSUsEXBOE0zn5aYEu16hbbCuVJGUeUFLqYg7ykljLYYhI8V8wZirSVkzoK1VpCr78k4nY7c3t4wuYm723u++uoFXXckhL/l+vqKqragIiFMOX+RZfMVBMEDGq2EI2KMZbfbLYaiwFL8juOUuZhhqSjeje397o8/3zs530uAtwqR9/1+/t1vokGWtm2Xa/yRHC5IzzAMi/XB3Op7eHhYUKk5b28uss6L5flwe14Unbf35uLnnHcrfLK3KSVzC/13IVznLb93bSXOvaeA3B2JFIXB2mKJmwFYrzfUVbXsN+M40HU9VRWXx2qalVBdbCGpCCG89d611lht0IVe5nlUwMz7sjy2p89e73mxGNOfxo39O1U0+fD4xdiiwCiyjb5Mgkf1gn389zBgghe4OwTUWRXuctuhWa1kwy9LyqKgqmuKLMOdW3YxBGbbguDFxLKua6wVspsxmnGcC4w5p0gg1mnKzs1ZVWaMEe+fGHFuyl+4EZ4Lsllaa/EnicfQ3YnROXa7nRQ8WJ48ecJqLUqlokhZuRaZnLwGa4WXNVfdwzDQ5Yv2k08+YRz7xazSREXIeHhZFAyjXy7OyU2C6Kf5ItSUVYUtSo6393THjs2TS7z37G9vSSlRlRXBe5GfKs04TnhfY6xc0NIaElKmzQQ+5x0JzTD2xOh4+vSKuiiF6GeLjNhEqrJgt9vhfWC9XnM4HAFBhaqqWvLmuhBp2lY29XxBFoVIZLVSbLdr1rvtkh2WSFlq7KWNq6QtKaR8mWdz6nkIEjg5D2stLjjqqqIwYrJW2mJZwOZE79m+wjnxuZoXXm0MGsSuQucCPCt1XJC2WIgBVCSW8p1YpCAfxpGyqKmamhjE5ytGT4xy2iuKEq0th6M4yhe5PWjLErHATGhradqWchYgKCkeU/C4IATZOTIhItEtMQhSRl6Mq0o8otz07Ut8/7GOkNeVu7tbmf/7Pff3t7x8+YL/+v/6f+EnP/ksI6QTMQbhOWVid0oinFBo4Z6lgFaGzWYrYoQwG/KaZSN2zr2DkMzEpG+/2PjzlS9/7DP9kTibUst1Of9bvedwDG8XUPNtZ6TwvID6bQhOCIGLiws+/PDDpbCY0aa5gJpbe/f39xwOBzHlfXjg4eFBgs2dW/hRc5twVum9+5yz6vu8aJp5P3N78N0i632f0btF5Pl7Fg8tRYpZSW0tXSc2LE3T8OGHH/Lxxx9DghcvX/D5559zc3PD6XRa8hovLy+5vr4W65z7e25vb5eiad57QEjis6ozBlHLGSM85ne/p+W/zzl5f4LxnSqajJa+qFZCStZINW5yH9fkynSGvFNiIYbZLM221jLmrLM6oxBD1zH0PW69Jqa0BN9aa5mcozudqMuKzWrFmAszIZLXgFn8i0BCL1X2cDon+83uykVRMvtYOC8Xgy0KDGIkGIIUae16w2q9kmLn668xRcF6s6XrOna7DU3dSLCtd6JmMmVGzxxlVeTiR9K5XZTnGYYBBdKqOu6Z3Ehbi6Td5MJFK8t2s2Vy/pGDozQsfJbI1A/ESQiO9/f3KJMRqUxinKYx974Dp67j5uYNdVNy/eySixhIEabRcXF5wbrd0fU9h8N+yairVyXNasU49GyfXDPsB8ZpwkdPhZxO2rbFGJGqNk3D7MKcYqLve2afq/PTX103BDcuSOBMcO+6jqGf2Gx2xBjph16oJkoxjCPjOCD+LxLIOk0DIFl/bVsgRTILRD0/b991koWU51Sa+Uz59czo6DRNcnKLScjw1uaYnUBVi12CSpq+PzI6jQ/i9ROSqER0JrUObm7fSf6XUlDVNWVRLYtx27bS0jMG52RRL7OXUEQxTkNGSFPm2InNg1GaaAoKFUFpjAWfIjEGMTTNCp8/BYfgH+vQWrNer1EKadV3A69eveTVq5cYq9AGPv7kQ4rC4oMnBo/zPs+T+VFkrfJ5bqzWK0xRLHJqbTQG2ehCDEtL7z8ddOkfx3i3aDpvxc3IzCM5m7eQl/M/72tznUe2nPOInHMLwXwukvb7/VI03d/fv8WLmgnps9DnvIB6n7HmeXtxfv3nvKV5bXy36JsLl/Mxr1tSoGSFeng8SHsviHdV13z88Sf85//Ff05RFPzq81/Jc2stB0WgbRqePn3Gp59+itaa/X7Pq9evOOz3ch0NwxLdY4sir0NObFvSY97hck3Ex+KQ9Pb3FL7v7bnCWmkhxUAKQYyu8mTxXk7D2gj51ztRPxmbK9Izgl/f9xS5ZzxX8FVVUa/XHDsxG0wx4vKXHMPjhNRKJoxAp2R1XqDNiEGIonwKMRAzGTciLcNhmChLURvM6VFzEniInkTMqgYpoGyRw4aVot2sMcZy8+a1ZI/pRx8Qay11VRFDpE/iv1MWBW4awVjKoswngoLgPNPkSEHk4bNvCSRBNLynXZUcs4GjNYIWzcb4KSZccOhasV6vCd7jBmn1lGWVX7twoowR9Ye1FluWlHVDimLg6IKj1op23VJUFdZoxrEHC822xWQZf4qRU38iRYWtxKJBCo6c4eYcq/UKN7rFhK4oH7li8+fTNE1WC45ScITAw90tKHGDv7vbs1pvQSkm78SRWwl3YZ438+IwS3C1VrRty+Gwl2I7RqZxXNpt4ziK260xOC9IZ1lV1FWFqqql7z9bQiglB4MYhN9kbeYoaYNWhtMxYqyh1rUYJBZiPhpJGdoP1FVNXZeMYy9hxNMguXF1xTiNrDZrwl7y8pQRAndCWnQmL+oxRYHEFaAVxMfbhBTz4m9QOb8sBi/3z+qW78uwxrBarQAYxonj8SQZad7B//A/cOpP/Df/9/+aH//4s9zyV7hpIkaf52mJUpYU5LNVRlM3LbYoRLeUkrRhlTh2x/i2XPyH8YeP8zbX+xCmb0KQ5k14LlTOx29DZc5/9j7S9mq1omkaLi4uFkTpvCA7nU7s93tub2/faunNBPTzHL7FauXsNZ2/npk3Nf+9GJuete/eLaDebdnJ5yBr6xQnfE7KUMpQFJq6bmhXKzabLev1ihQT4zRxcXklB8n8One7HevNhrZtefb8OZ/+1V+x3+95+eIFv/r1rzmdTigle0xMMZs1D4RsJaCUkoi0IeBd7owUj/xmk70cY/j2r5PvVNGUvMc76e2rHgnjS2kh0xolRlvTODJOjqqqKXQhlSti1FjkEEyQAmj2bEIp0uLILDycmBLb3U74Hkl4Qn0/iF+NEnfYGDJ07sVd23uPalv2hwMXVYXSmnF0tG2JDyPd/T2Tm1g1NetWNj+lVH5tBq2toO1GyNdlqYnzApk3rHEcubi4pBu6PJl1JpcLohaD9JKDD8va6pwj+kBZlLz86ivZC4lyijk8oHSiqRucD4SkOR1PslmbYuGHzReMUgprLL3vskJRionCFhSFKM5mhdt2u8WaQgqbScj1ITutG63F7DKvQSEElIYUPA8PHa/fvKFdbQBxOi9ym85NQjwexynzlQxTmpiyWWdZyEU1S4bbtqUsS4ZeiO3iZ+V5ONxTlCXtakVZSXagNprtdofLHKm6rnJhIM7N8+cZgsDsVSVopTxXJyaRfcc4OiDRbLe4yVEWBUFrIUqmuVhWWOQwMCsZy1KK+rpqqNs10wT7/YG6aeT3RUFUMIxdbr9JTpx3geAi6/Waum1wfsC5USwl0BSlRHP44IkpMgy92F0Yg8vXVFmWGFMw+ZEQPQol36mVyB6rNIUyC9diPnX64InOYWee3vdkKMjqN9lEFIqiFMT357/4OZMbiQRi9Pzkpz+mqmrhqC2b0uwMjlzbRlHXuZjSOv88K72UksI85UPYH6/F/2HwdtF0zlP6pmJn/vdccLxbcH3TYeHdguwc2Zkfe47tmu1wgIUvNU3TgkadG26K3cHjn7lwmgupubNwjkydt3fPX9f7ImPOUbPfbF3mDLozWoyxlpQi4zjyxRdfLFyvYRg4nU5y+C9LmqbhdDrhvOfu7g7nHFdXVzx99ozLy0usLTgcT+KHV1ieP3+OLazwwU4nORgrjQ+ew+HAw/39W4XiUuQq4ULpP8Hh4ju3wgUvbQ5TFdiqEufdrKpSSvgqktYtXg1aa0F7ovCEgvOUc/uk75eMqBQjU9cJMmCt+BwZw2qzQQFT1xFcQBtBO8rK5BBfUSF1XU9SSbykmprJO1kgQ6QfRowtRZZZWFTOcAOyXFNjCyH7ajVLxy0hSctm7Ae6vsstKYMxiqqs6DKvxntHN5ykBVlI7Il3brEu8M4zDCNKaa4unnDY73lyfcXkPJHsC6IVyhhMEiPLuq4pjc0mjrPDuPy3MpqyaRlfv6GuqyXrzlp5bQBlabm7u+fZ848xxpKi57Q/cDoeqOuG1apFK0Pfd2KU1vcMYw+jJ7keHz03Nzf8k3/ytxLEmDk6bvKMg2TcWZvkFJ7U8r1ZaxeIOcZHjtk49BmJehQNKCVcnG7oxVwySqvqydNr0IbTQfLWUErQMS+8pTnnSha0TiwmELl2UzcopbLflVkKyt1uJ1LcccR7h88ROlVVSfyMtYLqkRfMomS13mCnxKtXrymKkrppKIsSHyNjP4oH17okKLFMIMo8dtOEdz6TzEus1gTnWDctp+MB7wJl8ZjYTmLxASvzpu+dqD2tKRHvIU9pLaW2TG5aUFZjDAYIYbZT+P4gICGIghYlBg9NU1PVDcpo7h9u+fnPf8409aTkadqaTz/9JAfS6nziFwK9MQUpBYQUDmVRYUxBDHJt2GWDZymcxLf3+/NZf9vjvFA6/9n7/pzzn85Rp/ln7yNSn//+3ec9L8rezbKb/8wWBTMitdvtuLy8XG47cyWHYWAYBqZ8IP36669F1ZmdzmeEaiabnyNl50XGOYH8XJF3fptHNEov9AdQy7oaAuz3e/7tv/23/If/8B+W9XhWGs7G0c45TqcTr1+/pixLPvroIz755JPcDTCsVi1FYRd+1Hq9JgGn7sQwCDWi73tevXolHNH8nhYe1Nnnm/4Eh4vvVNHUNC1lYbDbDdV6RVSJ8XDA4ZbJZbRFG9l0ZtPL3XbLNI1UZck0TSItzz3hOdW6z6qyZ0+fYQqBTWcQ1lpLsJYU4frpM2KK9KeT8Dg0KKdR2rDbrSV3brvh8njAoEg+oLSmHwY22y1Xzz8A7/FjR3SjeC1NIk/WOTfOx8g4nHJl3qKNpRv6zNmJuKi4eXNLwKOMyqbY4iPlFZlnIr45wQeCD6zqlma1omlWXF5ccnlxyTCcaDdbPn7+jJTEa8p7h65acOL4Ol8Q1liMkQnvU8KWFdfXT4HEMPQLKX3/cEvTtHjvs7Q+sd8/EKPDFjmTzxpCEI8ppQtRz203bC+2GJOIceJw2PPkyROAhTyt1MwF0WIwmt/32I+o3CqbL0rvPW3bstls8F4KtjmQuTse0FazXm9Bw+s3bzC6oG3XjKPDn07YsuT+QQJUt7st9aqVdqfWGKWRdyYgQ1mWbFZrdustVS2Gpm2zku8qc5rmYGaVkTrnpXgiRVIh4oaQIsM0sVmv6YeR29s71rsnXGx3aFtQlHoGNSiz90ldVlidIEAoItF7Tg89KUa2K+G+TZNwy9brJwxjRwC2m634f8XAKmeebdZbnJsgigVGVcl7mcYJUkRZsTkIQQ4PKGlRnfMdnPN/vgXhLzxmg1xtLXXd0DQtMUTGSRDIECbevHnD//I//y8olfhv/pv/G//iX/yL3Fru8S5izGzSKy1P9HwAsTncNGDsI8IL/G61/Q/jG8c5EXzm7syo07uk6vfxleZD2bny7FxdNj/HN7X93odkvU+t9u7939cmM8ZQVZW0sDJF5enTp0sRNav19vv9wpE6b+vNBdd+v+dwOCwF1XnBJI7fglLPQpiyrEgR9vujmLWeoVIxSh7obBRdFAV1XXN/f794SMHbPlkvXrzgF7/4BdvtNvN/44JMHQ4H2cMvL9BGL2tSWZaculNGx83ymPP38K6Fw7c5vlNF0ziOpGSYvKPrO4ZRvpztdisfptYSXBvSMrmdd9hcSXfZrdVM0noT12yZFHNA7zBNTG5WTQlJd5aUjv1If+yWRe2wP5BSYLPbUDcNITq+evkVz8NzDLKZVlWdFTDCVznc3nE87FFE2rqA3LcWcrolIj4UXdfz/PlzTvHEOIllfN/3rDcbhr7n9ZsDdVNRNxWlKRb7A2kTZEl9YQVlShJnopXi+LDPEzuHzvo7vB8hRYyV28d4S0gq+z/JhuBNIEW5yNw04fqecZqyl5GQ5WMM9L1waPb7e0KQ4N0kQYCiXtOFSP5tIbEgYb7QBryfqOuCqpLW5Ha7JcaUYd6HhWRtbUHT1JxOJ9brtYQzxsA4inllCH4xp5uLluNJlCrXT65Zr9fsuwPBpWxlkPDjyPPnH7M/fMUXP/uFeE61Dav1ipDd1ZUSdVphDbYsFouLGCPH41GUeEPK7UwhYiv16MMzDAOFLSR41U/stluaXMgpYBhHiJHueCSi6E73HG+PbHYX4mk1jbx5/QZdROq2Eo+g2zs++OAjQAm6qKSoUxqi9xweDvTDgA9J8vAmR3foqCsxBh0nKeR0UnTHjhAcwzAtSODkJLxUo3AmYLTCliWmsLld4BZSqFIiJvi+jMeNTfpoSsGYW53NqqGqCrruyOef/4oQBcW7urriyZNrJC7kMQZjPgyIerIUkcOM/qUi+9CIT80P448f7yJM8/V7bur4bsFyzv35JpTp3cc8l/Wf/31+3/f9/Lx9Nz/+Odr1vmJrFhsVRcFms3krLHlW683cqLu7uwV9mlt+r1+/5vXr10srb7Y+mMUqRVEshVPbttR1Q4wwjlP2TDKLlYoPYTm4zoXLXMDN7UF9djCQw7Diiy++WA65FxcXbLfbxV19u9vyrH++HBbqus6UBL9QYuDtNupjcfr7z43fd3yniibvPRcXW1ExKWh9kwupRNd1KDQxCs+iLFj4TtYayMZl282GpMVJdO71nn/YVVXy5IPnGCUtieB93vgTV1eXFFo8Kbq+y0hQBRoeHh7wYWRyjpvbW5ybiElx7Drc5NhtLnj9+g113UC2lUdp6rpge3khi+TkCSkxDANPrq+pqppxdDSt+DEdTkd2W1GbffZP/gnBT6gYmMaRYZho6hqrxdXae0dQ0DQ1bd1kPo5i9+RKpNLOU2jJL1cJqqalaVv2Dw9YU1AUNdv1hnK1IqHpjwfC5NiuNzlbT1xgm6ZhGMWfaSZ+N03D119/yScf/4hpHHDBid9PUDRlRSIxTCN9N2F1SVVWS9tMG5HXD8PI5dWOm9sbgovLCcN78SDqczr59fX10lKS789QlgUpRU6HI36Uovezzz5lmibq5nHONG1DIrHZbHAucjgeuX76jA+efyQ+UqcjTSNE8L7vSClSN6W0wpJ457hRWm9DN2K1Zd2u8T5wc/sGiDx9+pT9/l5OhG1L26wyiV/4U/JdeVZtKyrMouXV69dcbC+oyxX7h46mbjmejtjK8vEnH3Pq9lS1ZdOuGfoBYmLoBrQylIXFatmQffCY0nJ1ccHoHG/evOYnP/kp//u//Xes1y1tW7N/OLLdbNBG8+bmZoHRyZtJYQpMLQWbcw5dFhSZ2H86nYgpUZYlu13L1dU1tij+3MvCX2zMOYLDODKMI10/yJHFGBEcWJN5ZI6bmxv+9b/+n3HO89/+t/8P/vZv/xajLSEkgpu95uZYGmlrxJAk+Dqrg+bMQiXyIeARdPqhlPrDRlmWv4EYzZY17yto5vVltg95lzQ9j/chG+8jlp8/x7t/3kWt/pC/57Xt/HmstWy3W9q25fr6ekGlzg98MwL18PDAzc0Nb9684fb2ln1WtLkcTu+c43g8YYyADRcXq4WLNT/uPM7Rrr7vmZ3K5+ediygz29tMEzc3N9ze3lJVVc7uFK5lu5Z1s65qnlxdYQrL7e0ttze3EvSehV3ne/m5B9W3Ob5TRVNT1wTv6fJp1hYyGfpOkJS6rjFG8tTGcRIJZFvTNjWHw0FaCYVAhNOZUiHGSNs0DOPI7cM9h+ORyTlZ+LSRSIym5WF4YBhGNpsNSsPxeOT+4Z6ysIglkUGhISm6Yw9Js1lv0RiGYeDi4oKmqmVTjx5UEnm39+zv90zjxHq3ZbPZ0PVDRlAitSqYxpGbl69ZrVq8c3z9q88Zhl7g0hwyG53LXlLZZ8jJCUMjTuiTczStOGc3ZUWIYldwd3dL1Vc4N/DyxSuuLq/pGMSDabfDWsupO2GNoapks3z5+iUXl5fEFCgy+XfKbut1XfPs2XPKuqJ/OGL03KfW2fwzMS1+PkKiXuDdEHIUjWTHhXBEacvV5TX7w4EEbDYbHh7uBJYOkdWq5Xg80J16jDWEUCwnrznJ/fb2lnEcsEY2/BQTd7e37C4v6E4dYChtYOh6UpDCrK1r6aWfTigldgDDIPYDhbFLrM7d3b1E11Ql//E//kfW6xVVXdJ1jyTIu7s7ClvwcHgQgqKWWA1po9b008A4DNhMeNfa5HYiTMOR0+mBN796w1//+EdYYzjtj/jgsKaAwGIq6b0nIJE4KYkE+DR0HA4nlNY87B/YXV7w8PCAsYb1usH5iZsXtzSrmuPxkDduLShc8PRDL35eynA8dthC02abijJHuAz9wFdffsXPf/mLP++i8BceYiyaF+m8PksbXf5ba4NOEuXzy19+jnMeayRc/Kc//RtpMxwGZpRp5sCN44TORqkiwlCYrBr6jY3gT+xL830b736+5+2n9yFLv+v+v+33v8/jvfu7d8nr7/57iWp6B7WSzkf1G8XezM+cSdszuXx2Np9/fnNzw5dffsmvf/0F9/cPtO2atu2Ei2RsjnMyS35qeYZOTdNEiIFpckuxBG8jaz6rll3uJKUkfF2QkGDZWyru7+4pq5JpnBiHccm2C86Tzcgz+Dtrvr/d8Z0qmqqmoigLur4TqDpGDscOW5QM/cipf5DE+ZkI7APDOHB/d4sPjpTky6mqSlCofLLYHw5cXl5RlhVt3dI2Db4MxHZN27YUxmKKkuAD9uGBOqsAttsLjt2JFB27qwuwBdP+gDaaJ1fPxNrAaDabLafDEWM0hTV0xyMuZ8zFmP17lMaWFW50RB8hRqxWlHUNKXJ82DN2Hfc3b8SA0U/UpaWpSlHC5Yw3ayVhWgrDKZO8oawLyrqgrioppKzGT7BatyhznYN9S7bbCy4+/ADtI4W5xwdHVIl12wAJ7wYKq/n0k48xZUE/jCQSXX8iAXVV41zg1A1UtRhOrjctxoLzI9pKjMg0ekgWrUTNGHKhInJ+nVuiWmwaCktZ1Uy3d+wPD0ASHyEf6LoHVtsdzXorwrRswjlHE1hjqOqavo9cPrnOeUlw6juON695uH9ARyVhxFEUlKSczB0c2hjWq0Y+y+Bo2iYTNUecD1RNzWqz4f7+DmUU24sNF7udyMZTQmtoVi1GG8ZppFmL07rKJocuL1YpJZSxjJOjrCsmP8EIVS32CbtNQ0xr2kZczAsrRqxaGawpsneYeJkoDTobq9pS5OvVuhZvFxWJKlIYSARGNwhi2NYYq7F50UkqkVQkJE8kUNiKzWYrthZ5QbPWkkgE70gJqqKgrau/5BLxZx3DWfyF1oaqsksbZVHwJGnXpgjHw4mf/+wXkKRle3X1lKdPn4vvlREX+nEc88m8E7Vo4ZdNsCgsi2GtSuJBl36ol/6Qcc5pene86030u8Yfypn5fW//7ut4XxH3Tf+eOVrvtql+25hbfOv1mvV6zfPnz99q/c2dnM8//5y/+7u/43A48vXXLxnGibv7+9xSLhZ1XNu21E1DU9cURcFqs2aTX9OMWM1gxSxAmaYph4MDWmVAQxTvMURimAhagu6HflgCl40xWVyhFh7VWwT2P8HF8Z0qmm7vbkWGmRfpuR+6u7hkyoaV8OgMHoJHa0XTrJB5KBYBVc70mU8PIQTazY7ueGSzWrO5upKNc5LHnzffpm5obCFy4BAorWXTtKIAGzwBL4HPhXByrNZEH/DjhNEK70YmlXBuQCjMWja9ql4KueTFk6Kuy2xtIO+9rWueP7tm1a5whWNyktOjtexydVHmCSiTrSgKCmOytYCoCKw2lEXJYAYmN2Ksloo+RrSVk8l6tWJ/e0uh5L5GK4xWVHWFsRJF46Zp8V5KKXHqTrjgKaz4QSXARyjKChM8JEFwRtdLJIdSeCffwyyrneNolJaLXC4kJ3l304k3Nzf0fZ89s2b/pIaHhwdu7+4IbkKRMuwOSslRf5ykOP36xdf86Ec/EuuBaSTEwCarMp5+8DH3r1+jYqQqCpSS4ssFBylgbUnCEKJjjh5BiWGb0poQI6e+o6hLJjehjaFtGk7Hg8zRpCnbktN9x8XFBSklDscDxhpMlpMX1hJ9yGhgk6MHjlRVxRzauWoqHu7uePXqFZv1mvVmzeQdRRkzsmYXJVdMEZ8S3fGEjz77TEHMvmFlXeGdEy8pY9BWduCYktBtgKiEhFwmCX+eppEUFePoREFHYrNaSbu065nG4Xu1gS9IdT7aaiPeZPLHQPZAC0F8bLS2OBd48eIl/+bf/H/QuuC//C//z/zNT/8ZZVNxPBz5+5/9jFevXgvvo6qkaI2BWSsnGVzy2GkxI/k+feo/jG97vGsvMP9sLmxmx/8PP/yQf/kv/yW73QV/+7f/jL//2c+5v7/PVIaOw+HAzc3NwlsqyoKqFGPPIrf1m7alqSWf9NycMyVJIajqiipVy5RWiE/irJqPMQq6G4W7KlxC9VaerBTFfjmIfNvjO1U0KaVAC1KkjRF0IpsHqqoSVCiTy2YFWWFtjq+QxSulRJFtA2YDOe+9HAVTZHJi664Bg0Kl7DCtNU1do1AYazkdj4tyaBgH3DiiM9fHeccwDoJYeYlZaNomx5I4yqqS/K5MADXGZpNDh4qP5mMoxeF4xGhBXow1YivPY1r2OIzi8t20BGPwfViKSZHqy/uUIF+DC56YxJW7rFrGYRJ+jvP0w0ihC8LkSLrAGiFsxxTpuhNFRi26/kQIgcur6+Vk0xQNIOTxqqooqwpTFHTdidKKpYLS4js9jtk/SFmcS3Rdv/SjjRb5+n6/hyScrBQ0yhZcPnkCURCpOeKkbYWv5aaRuq6oyuzMnu0mQowUTUuIUeThkPPcpMV2c3/HBx8ZDg97ri+uqcoim2aOxBwnMgwDyojkuzudFm8qYwzBBU5RlJSr1Yo3L17x5s0bTscjx+MDVd1gkllOVzM6MTlHqZCFIEWaqkFZIWuTzVBNUdCs1vT7/RLb451nu92xubigtJbueHwkiaIWE1cQC6H53/NcTykJCT8E/JmRnvee1WrFFCastpn3JuKAqq7w3tMPA0SwVtO2jSxMSuHHCTdNlGVFab8/nCaVW8ApJUKcs0pVzpDMuYDBo42SE/hK1HUxJb7++iV/93f/KxcX1/yzf/pfoNC8fPma/+l/+tf87Ge/oOtHZqJ4yg1AlQ8WxhiIv+PF/TC+k+Ndi4I/13i3lXfOe5q7M5vNhh//+Mc8ffqMzz77a378k59wd3dH1/VZnffAfn/P8dgtCj7vheYyDD2QKA6H7G33Nt9I1mub479mewMxfjZa0j5k/XpEqGaqx/w6bVaZzrfROocHf8vjO1U0XVxcSMFhpMVjrPBK+l5IqlVdo+eKWZe5aBAvlVnZVVUVTksYYIK8gQtKtd5uOZ5O9Pf3GKUlj6ssF3VV8J4UBMUZpkmy55RdcoCqVb0YSwrrf4cpCo7HI1pbiqLMMmVAs1TMgg7lU2OIC9QYjeTRBa2wzi2hjtpKy0lbQyRJcadUDuw1ufVjAE2MfiGXyulXTsartkUZzWa7w7mJfugJzqFV5MnTa1Ik+8hIO2YYe5wvsjcHoKQgmYONN5sN4+S4ubtHWwknHoaRaRxZry7RJhGTRSmBelOc6PsBks8nEyXIi4HoBOooq4rruiVMsLq4hOiJ04TzTjZwxGRys13jXLVcHtIzjwIZG0O73fDxxx+LwWQ+OTVtSz+IGdz9y5d4N2WDNoWbFWPWUBiLdxM66eWxldIUpkBpcSaPkFUlNU1T40NkcBNos7hspyhZeFppohYTylXbCFmy62TeKvE9CiESQ6BdrVjvdkzZHVcg9A3Xz5+jypKx7wleWogpJYZxxIcsozYGU1jKWvyzUko8PDwsn7WbpFg+J8SWZUmYRP0y+z0pFGVbYo3FKUcKUBSW7WYjbaSuywuizIFhHP8sa8E/hiG+WRVKK5wP9N0gggYtBycplAOrdcvu8kII9sDx2DEOI4fjAec9ZVXTDT3/x8/+nv/X//g/cpcR9aoqJN+P7BOkHwNXY/qhavpjx7dVmPy5Cpw/5nl+3/ucE8/nNeacqD4Tu+f1YbYduLq65J/9Z//ZQugehnHhQ93f33Nzc8vr16/ynxvu7u7Y7yWc+O7u7gwl0szZeDMXajZ51lnQkhZxV0FRFpRlQYgSAyUqPU+InhAVRkvGXiIS4qL/+lbHd6poOp1OzCGv534S3stGMZ+cZ5dvY3ROpZ+wVudq1BKdx6fHxOiZqFat1+xCpKlrkY0HUbUYY7DZz2k4HVEomhips9nkyVrGaUQhbtM6K/NcSqxWKz7KPCufzfCGYaDO/V5ypSwhviUqR3+4rFQoa4E3Y/YVWeU0eq01LgSJWrHiIOzGEe98jsfQeQF/zAQLwWNLK1l3tmAcOmzTELyEEa/WG0CCRFMUd/F5QrdmRdM2UnRmawYSTNOIz8WTtZZdVmkcDifxfMob+t3dLZMb2G7XC6dnmibKQnrpZVkiF1EEIhcXF1xeXuLGkbvuyOnmDfv9HqUV24sddb5437w5Lpv10PeCFk4jRSHBvlPXkYK4pRdFge97hnGkqEqePnsmfB9bcXl1JQtAdDTZfDPGAFrRTz0qJWxuo1VlKVls3lNVJSFGjt2J2zc3xBB59uGHXH38EeP+gddffAExLTElZVmCB+cnTl2isJb1agMxcnd3R9uuKMuCcRw43d8Tu55Xr19zdXXFzc0Nm43YO2itl2iT7cWF5O9pLRE5ZG+rcWIYAnVdiSXGwr/J0nYriKv3nqYRJPTRNRiMtQunQIpyTUiBoe/pTqK4Wa1W7Ha7ha9w2B/+UsvDn30cTz11mbKx7iz91tIAzetLVVc8ubrmo48+5vLiAqUV+/2BvpM14OrqCdZavvzqK/7+7/+eFy9eoBTi7F5XFGFCKbnuZ58x5xyaP4x/88P4YXzTeJdIPqNLM93lPOBea53XCGhaMWoVU1xZj7ebNVdXV3z00YecTp9xOp445YOVeEIdeHh4oMvu5g/7PYfDXmx3RuHgDv1EiLOqUaoea0ToVcdKuh8xLvYIs+Jxtl4AiSUri+oP5qn9PuM7VTQFL6GwReYkzd5KoLBWCMXLF640Vhl0qTF2loea3PaCc88IpcQYcf/6tRhrFZbT8UR36imyiaDWSrLuVGKz3ZCiuFiPU0AXGov41vicVxdi5HB7y5g3OJFo1ux2l0LOVjIZp4x0zUVeXVdL2K+2hrv7e8qy5Pr6mjqJtHPMEn83jkxhQjdQGIO8NUVpS1GexSDFk5GstWGcwDlCEFuEshIvmMIY6rxpjv3A5CaGfqTKhmZ126AUODdxc3sixCg+G0VBu2qYnKiDbFFQluKhcXt7y8XFBff391Lua5b8trKU28WgCV7iaYIPGA2oiPcjk5u4efWK0/EIqeTyyVraGzEQvON03FOVFpUiuqqoUiQFl32cdgzDwO3tDcZoNts19/d3TNMk6FNGy5xzlGXJNAinqh86rC5o6poUPIfjnnKOUUEuTI2CmJjcwDiOWFMw5iKtbRrc2KOSJ+7v6e8fpNBdreiGHrRGebdw7pxzHJzDGsOmXbFei/AABU0ucqKCq6fXfHD9jFf6FZc5w2mJMgmB4BwDUuwqJZ9zArEcSFJcdaeTRKEYI7En3nOfjePquib4kL8ftRTxxphs/9ALchnAT04KvfWGpqypGrGa6LNXVvoe9Y2MtZiigOCxRhyMpS035pO0kFTruqau6hzWDW2zom3WXF8/5fnzDwD45S9+yS9/+UuUktbner3GGIXxoA1s1htW7WrZpKyxPziC/zD+QWMulM5DiJeM1TPDz/l256TtEGPOnNOZ1ydhurYo2JYlu91WrDHO5mhMkWmcOJ6OC+L06tUrXr96zc3NDQ8PD1mxd6LvB6ZJLHyC9znBgVy8SdtvGMYciZbyAV8RQlpuQxJ1+rc9vlNFk7UWFzyp6xbil6hKBNYLIYh6CPEeWqrP0jCN02J2uFmtF1SqruvFGfxwOHB5dQVIa6qqK9pGUKjudMr+OsIFmYuzub1B9mepq4pTL5lr88lwtg4IIdA0zcKZmaviqqryZi/FWVmWC3JzcXGxxLocDgc2m83iAhtiRHmoiswHMpqUZsm6ztD+o/2/LQqqWmSg49jLpujlvfkQKCtJlL5Y72gbCSI9HI90fUdhRe5ZVRW2kMc/nMQfo64rCUkOEaXFMmC92fD02TOeXF7gvZyWUZLDNU2OvhspbENR1GgtGUFKSUFblCV1U9NWtRhIRkthNKfoCSlgMQvKWFYVeE+hFC4jkCq3C601HA7idjs702qt8UG8rU59R9d1XGwvWK/WuMnjx8A0TiiVqMoKoyXPUGn5bgorWXcuJFRRURclY9djS3jYP+BDoD+dON3vceNIu5a5tmpX9H0PMaIzR0B8tcRyQgog4V05P1HksGaQbDpx8T2x2WzFt2Q2XM2t5RAjwft8OJDcpRgkL9BqQ9Sa9XotUHo38Pz5c+pacvOMMbx69YrVZrWYlarcKnTeLQaY7arFW8fxeMzOvYH+vmeaZC5vdzvK8vujnvvp3/wNq2ZNDCJ1NkrTdSfuH+4yl0OiLu7v75cWfoyRoih59uw5V5dPuLy4IsbEz372c77++iXPnj2naWqKwjIMJ6ZppG4qmrZls92KWa4yP7iC/zC+tTEXSPNh7F2n9Pf5UXkvfL2kADTaqJxVaiXTUinh4qVH3qaOEqRbNSW7ix3PP3jOj3/8Y8ZpZBgH+q7nmBHs+/t7Dg8Huv7E/e0dtzc3TOOEm7KxtXP4SZIKyrKkqeuFKzWOo3RqxgnvvuecJmMVVkk0SeLRbl5r4bZYrdEoaTdlozjvPcGFxZp9JqgdD4e3ZKfLhMmZPs5P2biuwzmfJZWWYQxEpM1htJFQzhQpywpNwSEjS1KhA+gcP+KFUJ4J2taKSV1MAQUYKzyF06mnrZslVDglIQVXTUPVixR5Vi6N07iQ0YGlDTeOI2Wmj858Gu/9UvN3XYfW5MDZAlPIRC+rCpc9rgSFkfestfCfUAqXkbSyKs/afkLQmybHME65RanouiMaJR5PNkEKZ3L1MkvxS9x4wLkJraGs7ULmTiFSGAu6wGakLqlIWRSsVg1KZb8jJwq5FCPjOJBSZLNZL2RlIJPGczu16ximAVsKzBt84H64py4k+yjFRIxeCIUBabfEhMpqxBQDiiTFKbKYNHWN7zrWqxUKOA0jSks+08Phgcu6FMPR5EErVNTiM5IdwEMIgiRAdtktCF5ajaEs8T6w2mzwMYrHWOaSBe8x1i7I1cxFMFbabzElhozIzkVzl0+UM7QtRfS4uLcH7zGFxRZSrI/9wISiMAV1WeHsSAqSd1evKxKrt2Tx35fRNi1FVeG9ztYhSYrMfP3Mp3Y5JLmloH/69BlXV1f8zd/8DU+ePOHu7o6vvvyaoRv4yU9/IqrW7sTh8MA0Oaq6pK6b3LITRHi2xfhBPffD+IeOdx3Q4W0n83fVdUu+Z67c04woKZW5ROcVveSnzl5mSoMxBcYK33K1WWXENC1h7F3fCeJ0lLbeMQfzDl1P33UcDoJUDVmtK1zSlqqqpKMy9BwOB/Z78Zr66uuvvtXP6zu1wvngqatmybCZfOZvZDO/ol2hPIxTl/1rLIGzcFzkA7bWLq2IcwhSK7XI35tG/Hhsbl8UhRWUKZqF0F1WNUarTO5WpKiWx31soQCoR5VNCCK71wUhegiywVlrGHOhNBM+50VXHkvQrbKuaC+21E2Du70lZpt8CestUDn/3BhDSI8us2VZYssSH6R4W60aURKWBWUQdEAbg1LClRJyfFa9GSF9z22l+QQSQmC1WmVJtaAxwyh978IatFL0fU/fD7RNSVmKUajPxeo0jmgVMdk5mRTQWgqwoR+yU2zBNEUxoEwxS+NZcubEAVe8nay1pFRhrVmQPaUU682G7nRa5oACqrISzyQUdVEz9ANN05JCYhoG8dSqGikG3Yj34t9kjc199oQ1CkXCFobgA3Xmns0mn7YwuGli1TS5TVOJJYGWJHttNG4U4nXTtKxqMWFVWpNCEu8uWwhPKYlz+Sy5neetyYGYs8Ll3AVXa4NKiWQeF73Nak137JeFcIbkV6sVhZXg2LoSya9PYZkjhS4otCGGAEna32b2U0lx4Rmczj7n/9THqe/p8+k3RBGJ7B/u2T/c4yaxnhD0V3zDvA8ZmW14ev2MTz/9FGstX3zxC+4f7rGF5cMPPySmyOvXcu3HjATKxtCIN5fJ38MP44fxLYxzL6d3c/jOLQFmTjBKYQpLoQXRDjEIBykEYgiy5yUFKknXR+nMrc2K3hRFwZ0iJEHGtWEx9m3amsvLC8lNDV5cYxN4NzH0A91JOgSTk/2yLMrcrbGEIIfB4+nIV19+xdcvXvD//jd/961+Xt+toim3Deb21ZzKDgqVyAhBXNAcbTTBPSZUj6NwT7z3QujNifXz5rq5uFhaGsboHJjpJDwzeMZJUAyl0tKS0GWxRLGQJ5X3MxQvuVJz8vtjNS8oVAzSBzZGfJ2WpGaFIDuQDfIibhzJCnLhMzmH0ZqirLDaoJJMcJ0X6fmxXEYlKmOwRgwW59OCOHPLqVgb8bDSSi/EdznApkWZp7KZWEqJvhP0wlhL35/EJ8hKLIRWEvnQVDWHh4PIpLWRYqWuKBNEn/CTIwYvWVtaSyQInrou8H4iKiPoUiQT3IEYcdOwFC8pBryblm5FVZXL5z27guuMOPZ9L5YQ+f3OMSAxCZI0TZNE2XhPXYt5pA+ONMnCYQtL2dQE53HjJD36acRktWO7XrF/uBdUsyyJPuD8RNPUotozGhOlMIW0fP/WWuo8F5XKJzaFQN6ITUVK8S3jxHk+pfTYhpO2tIQ6zyTxiBTlbdlyPB6zIk6c18/JnlUlpMl5gRT4PQDprWxDP2a1o9GgkmQVpricPmeE9PswulMnIdvZuXjoxatmjowos92JUoKEe+8IIXJ/f8+vv/iS//1//3es2i1ffPkVJMXFxRWbzYaErC8PD/dUlSTcr/IhYVasqh/acz+Mb2mcx7Y8yv0feU7zYX+JecliJWMM6Lx+phmdmhGlua0ne5gUTWnxnIO4cJ6UEr5SFhDn/xY/P2tEgGK1FZHURcT7mPf/zEtWOrvnK2IQcGWaBp49fconN5/w//zv//tv9fP6ThVNxhYLh8g5R3AeU9VCeFaamCHwqqwoypKUIjbqxYRxmiZM5pCoPDnmCVEUBavVirHv82k6GxzOnk+5R1tUQuaMMRCCo0h2Qaqqqlps3+f4A0ElzFK168yNSZk4Nxck59LPReVkzONk0moh5k45/FCjsgJKEVwQMrUx4hWjFdaWj5LRmPA5RLEsH710vJfQQx2E+6KNIDbRS0szRSkglVKLXD8huUIxBKLz4llVlKxWTW6VOvpTx2a1Fg+ptsVoRYiTFFTG4KPHJ0FNcLKR931HWRqMreU50SilsUaLdFRHQvT4rIYjpGyWmT23UIs/03zhhxAYs0Bg/mxjlOiLbpJw4dJI4T0rDKVQlEXAWIu1wpez1lCUhRRumTOktMT5BO8w2cupKAoKaxgz4T+lxDgMS3GTgJAz9LQyRCUHApKYdkYARSZbRokziZHRTQt/bZ5jIaOCirchdeYCPRdlZVmKPUfXY7TJ8zst6si5IJsPACEE4SHM8zOKCELlolopAdWtNSLtTVGuKfv9UXX140SZFHVdYY14qFWlmIbK9WdzYRvzhgMxeu7u7vn7v/+PxJC4vLwG1BKEejr10rbPDvlVVbNatZRVJdma2Y5CiuQf2nI/jH/4eLdomjsw587a85j3SymUJH0g5XVBaZXXbCVzcymE1KNZpSLvf/M6oZeiiRSz8lcMjIXPmqk4Sq4HoxVVVeQosjIXdeT2XwKjsEnc+derNZvt9lv/vL5TRVNRyOLvls2/XBb4JaUaceWNwWdEwVBXNZObsueScJu8c1nF9phU7qYRo0VKHmLIflDVUsQUVUVZl8Qs0Z4VZCVkJ3CLc9OiKFBKAmhBJs65hb9U1GqRcpL/22jJZ5vfkxRaijK3BJWStlZZloTcmiMJ8Z0kBnvSjrQZCyVD+kpaQ8jvZyShbkp8zuELXnKx5oslhLR4B0XvBenL7rCqrulPJ9yiWCS/TyH+dV3HKRtBFqYkJS/OrimgCFnGLm3JGKSgs9aKWi1/FkLSn0he0BDvHC6MoCBG+d4LaxfIeM4TFHt9zTT5rEQUblHTNIzTSJrl9H524y6hgNKWOahZIGbnHLY0ORDYC4ITsjNtCBhTUNcVymimIL48ZVlSWJstCcJiRTErIqNUN9kSwsjCM7dRM+/OhcDkJuGP5dbqXPRbLciomVG/GJfveXYXP+cgLIX4bFSXW8XmTBUzXztzMbQo54LP95cC0UeHVQVlURKCQ9vHQt65CWXA6m9frfKPdTgfsYWiKGtWdcPV5RXrds3d7Q3jKGpCn1sMEhM0F/KeV69eczz2fPjBR3zwwUdcXF4w9gO//vWviCngppG7uzvKcl6iRWwiXM0c+L3gqz8UTz+MP3ycWw3AY5vunBYCj9Es82EtwXJglaJpzr4UYCE3KB4RKMTRfvZlMkqR5v06SfftkZ+X8mPkfVnL45IJ5SEp2eMUzHYE5APcfFjUChQGrH4LIPi2xneqaJqGER8EabCFhMfGmBbl3Lz4ey8tFh8CpVLE7I80K+SePX2GSHcHilLiJUIMjMeBKqeShxCoypKinLPRAiSxGZjVT0UhcvMxk5+Px6O4p8QkHJe2lZbGNEGWqj+6l5oFtp/ciCyIYPJ7IVfyi6QyG3nWVc0YBLHpxolpHMWzqG6WU2jMssuQN3db2FzU5IDgKL1nhUSrWGMIzmGURKaEacqE3syryI+jyAkbSYqHkB2u67pitVoBChcdVSkKvZSVYtMwEdKI8wMpBkG0tKEuS7QuGEcHMWGMxWrN0PWcTh1lKRwgo6VF13UdITrqRny0gvfy+LkVGbIcNsUIKkmIcYoYJQTlGfVRStAhbTQhROFxeS+FTvYGiVERvGMYgzwf4kk1F9ExJerCUpQlzjt88LiTo67E4JIIY9fjJicFGoJATd4tnIHgPS6jn1VZLp5JMct558VrhsiFuyau7korQlTElBbHnrnoSblgjykRQsQlaZkZJeanwyCS4XmBnMUR52GeSkGM6rFdmITkjJbHGIeJEosujJi/Onmf36ft2xYtIK72ZlXx9MkTmnKFVpbjYU8ILjvg9xyO+0WUoLUIFo7HI6u25dNPP+Xq6ppxGBndQN8JR2qmCcimIUiVnOhNdu2bNzwx//th/DD+mHGOKp0XTt9UNBkteRlaGSmaBOOXtj/CP5qdR5aiLEleotGyiyxFToRlHs+3zT06ndce2bpUFj8kiDOfL897NZdr891zAkKU9e/bHt+pomn+VIzVmKSz39JEUVRLK2WGF43W2LzpzDwLk8nYKW+yQig/a43puc2gUPGRJHvuwuvGgSKTque2iwsBW5VM9w+sVit8ECTFaDFxnJ8PbdCFtO2MNRgrvVjnpsVDpypF4WcLS1kWuQ0oPkhDP8imnK0MSFIwNnVNXZWMw7gQ68jtRLKlfJnVUKKykqJIYUg+QBQ0oyxLoo9M00BZbnJRF5aJb3K70DtJo9bIxG9yOGPfDwQfqNYVuwUWlZO1SiJJVVrnoFlNYS0gikel5YSQQiTmtpZwQBLNdo3vpaBVqszBpXqxfpjbFYW1wr3K7cQ51gYABeMkRbfSj1b7c3s0BI/PbShprWgIghbNJGt5TVLUWCOtFzdNDE6y+IahZ2JA73ZE4iO6FCJ1I2amPgThSJ6d8myeizPiJwuGzghXXFqKRVGIR5lS+b1qUohLJpM1Ztk6Vf5TWJEBa21ymy0u87G0xaP4QM2iCPMWXD8XUKAIKeJjwKpi4SZYXWTeGEur+Psynl5/wOR8LtZB65qqDDT1DpLBFoq6Ljie9vgvPYfDQ27jS0ajcyOJQFUZNpuGq8sd603L/f09+/2e0+lIjIG6btDaZi5c7kWjWBJ7fxg/jD9inKNM77bgzn8HvFVEkTKylCAFKYaY0fKZbJcDOpeCh1zYxKxGVmdUghmFmteO3DV5/M38upZX+84bYca3Hu+hcviv+vbXo+9U0SSGb5boBQ2QNspjZRxTwmbXZYUgEZMThV3diJv1breTIsIWmLYVy4C8edZVvbRDghY0Zjb6K3LBEVO5ENFjFNWQLUqc87RZJZWCKLpi8IRRCp6mlpZQUZS4aSIS0ErQssKaM2KuqAwKWywn/xCDCAgyCbSpa0m6L8SdurBWjD+zxUBZymuanGPoe5J3GGsoM1rlvaMsLIURdZVRCluUlLbg1B8luiUExkFUgXOY8KLYyUhN0zQ5GkSKB+99hmmz/YH3gqgZiy40KAtJWpBucgzDAMngfcynFznVtG2LLS3BeaIKFGXJcDxlZ+uKlIJwbbQWMiCivLA5sDYmgXiFo5TjVuIjF8QUIsXv+15iebTGLjEz0pKcSdhlJUaRtigx1ohHkTaUlSgRvfe44KnXrcSvDKM40nuxANhsN3gfFhf7EETNqQtLWclnbrQWX7Fc6iSkhWpz4OQwDKAVu90F1sxeVGeRGoA+s5cIISzKyaIQc0VpCctnpjPnYFaHFrk4k9fnF6GAUjqjnjn7bC7Gk4QxK5OWhW1GRbT5/hRN//Sf/XPGydOfOpIPjL3neBzxLmFNSdvUbLaNxKG4kfV6RV2XXF1dkkh8/vmvGMeOX3/xOYnE9ZNnbDZriqJgvV6LD1bfkYjUdUNhi3cKpzy+Px/5D+NPMN4tmN433ldIJem95Ronyvp1VunIWUu9XbjM/KOzAkgpQarmiazeuu3ZP5bfvjvh59LqrGmd5sf+vjuCh0DfnZi8oyykqLC2wJpCvqQoKrOQN+8l6DJv6gpyvMlEqmvKslhI5Y/eSo9huD5nzUEUgnMmhCutOWQ35e1mR/CBw80dm81mIZDPNgcppswHmgjeoUgMw0DwjlDaRb0nPBXzSE5KamkLVk2NsQW+H0gp0Z1OdF1Hnd2q3STqqsf2TmBlpR3lQuYrpUjrpXBUInVY+F8m6EV1JcTh4rFVFSHasBB/F8jWiNllXZcc98dsxlnQ1jUpBNwoPlduzO3FQgGZ3wEMw0iMmqpspdBMaSHOa1PQli2nwxGIhK7j4e6OupHQ39nQrMgcMmOtFJaTNLtDDBIlk79HnU0F5ygTrSTSJoTA0A/0ul9OUoJuebyX4hVVCi9pkkyxyTs2Kym0xkmI2SHzjWYDwzCJDB2VoWyrHnlrWbIfckZcyEabpS0obMk0eZSZidoz307lwq1g7KWto5XKFgECedtFvSiF/szZmyZRbKWsgmzqmtOxY7/fE6IIIGJ4jB7qutMi350/E5XncS6PGCYJp47JM01jdtf3GKMeeQbfg/Hjn/yUoqzx08Td6xt+/Ytf4V1EK4O2OrdgBSn6yY9/QtNWXF5e8OTJJf3QU5UFv/zl5/z7f//v6bqe4UcjT58+kza1NdR1S1FYpmmirhvKssYYC7kR8sP4YfxjGG+jy+qtguhtIPQbqvv0viLod9znt97u/L+/fST2O1U0xRSJUVCG2ZV6mkRFN5+i5cScSAhZtq4KlLUwiKx6miaSF45TCCVNK+jSMAyLis5ai4vurNU6a7MUaAVGvIzK7FwdnJziZz6VUgqrNaassEUldzF2IZOHGKHIJo6Kt+JfztVMQz/gvFgLaGSjbdsKN405r63IxUBA52JbFH+BcRwoq5rLiwv6YSAhG7giE8MhO0gXS4zLul1RVTUwZiNEKYCEfMpiUxBy+zGmyND1jPnxzzOBtFI53FcsE4xJxARKiSVBBRANVd4MUpS24OQGuu6UM7wSbdtQlAVPnz7Flhqt4XQ8MAyDcNm8p9lu5f7jkC0WhF9jGiFx+2zgOFtLzKqxZ8+eZUPHRFGWFLaASG7RacpSOEvKKLqhw09OfKsykV6KwQJTFoTcjmuqekFqNut1NlWt8TEsruxjLqqskZZZhEzol/uZwuJjeOtkNz+nNZLVl6IUM1LYym1jVp9YayWcOGRFZZ5TD4d9Rg9NDi+WUN+ZC7VarYiZS+N9YJokfLfMwdfWVNR1y8PDAzmVCmt1di63b/mhfR/GzZs3XF4+oapKnjx5Qm1Lxv5jpvHE8fTA4XDPNJ3YbNZ89tef8tHHH/D06RMuLnb0fUfbNDRNyy9/+SsuLy+z51ng7u6OvuuWKB3hM8YlMeCHltwP49sY7wv1/X2CftP5/EvvFk3f3pi70Gd//TGP8K2P71TRZIwYVo7TJLk31uK7jqHvWbUrTJblU0l4bYwBkqTUAxyPR5q2IfkogI4SZEFrzWq1ZprGJc9OK0VZFdiiIAa/ZMQpbbC6ZFW31GWFn6RFM29qSitWTcsw9MRJlG1iW5CJdl3PNAxZiin8GYVk581uzPPmXmZPqXEc5fWU4ug8Dj1VVS5qLp3begDtjNoEabUUhcV7K0TnLMU3uS03t17cJAhYDH5p+4QQ0IApy+w9JYhVQgjNZV1zOhwYuj4jOmLwGFzI6izxzpKcooC1YIucoq0UWhlSEvftvu+zYtGjDdnhWkizPvO5MotbyOX59c0+VGTi/ExWrwrh/QQnYoBhGBbUKSYJQ54NA6dpIjmkpamtcKqitDpTkuf3+fuossmpnPYDOgb6caCsKla7rWS6hUA/TeIDVhQMk7iVV0UpJpWZHJ/mxSmjX3P0zenUkSYgF1jWFpiiwCgJ9G2ahnEYKK2VdnKMBJ+9mxKi7gsRZQwqpqw0TaxWUubMRPEYxMogId+RtZYUE2VZZb5XyO1nKcC8MoSQGKYRUxjqqiRGh/cTIYkbv9Watmn+fAvCX3j84uc/Y3+9Z7vZcLHd8fz5MzarBmLg9ZsXfPnF59zevebp9RV/89Of8ulnP2K73bDZrAjB0zZtNuEt+au/+ozd9pLD4civ7+/56quvFp7h5eUlzJw1pbJI5Ieu3A/jLzPeKpL+hJPwH/7Q72vl/cPHd6poSklS4UMMoo6bJiRsd4XSmmEcMZkjQkpiLplXl7ISHlJZlKhCgk7HrD7TRk7LdZ2dsbUWFVpKaGvRpkCbSFFkx+8QIURCiItia/YGCiGf1J1HZ5fu89BA76RAEc6PIyDmg8Zo8buJUYqV3FYiFzfeOWkrTVPOMmtZvJ6yN5G0YBBUNDwSkefW2mxp4KYJzj2qjJXIjRA5no7srp/k5/XLY0qrZyTGtCAXVmt0UVCWVS7uJowtqNsWN44SOGo0KfnMGzLElH2FVP53CLm9KJEytiho6mqxyCc7nHvvIGgKqxeZu/hieYITb6zZzHJJuiahQ0Dn4qquaynQgtgZjMOQ0blA9AqXxE7AiGZVkE0vbZCyqjBGfu5D9t9SooAb+oAqjLh3uyG3GpOETeairahKMc88K8KKbGgZQxADt8Ky2W7ohwGf7QGcE7Vo3dSM454nV9eIX0/Inxn4ELGASz7nMT0SM6XVU+CzW3fTtthC0Cy0WlRvl9sNdze3QlI2lkTKLvYIW8FYgpsYp1O27sjhnc5hrDgEh5Dox/HPsxj8IxgxOB4e7jk83PNSfUlb1fzVj37EJx99yAfPn7FeV5y6D7m62vLxJx9TlgX7/T3T1NO2Dc+fP+fy8hKloGlqdrstKQl3U9So4nx/ff2Ei8sLUXEqQZwNj0yOH7RzP4wfxp9vfKeKJpX7+IW1+BBw44DRhqZt6bse1FxoyGbipomoFNWqwXlp1x2PBwoz83MiWkHwEyc/SduDiLYFzkt7bHKRoqhwU4AYqUuRWWZziaVVVRUSixG8oFIzH2RWIs15UcoojH7ctFLwi+WBd47m7KSeYhQfIm1ISPmQIliVOTrZZRiV8EGkyIq4tCgFkXl0a17AjRhJs6t0krZlXVUMgwQnNuMoXjHBYZHPKOUw3ZSEjD71PT77JM3KK2utEO4zUV0bhdJRMgMzEhZcwDsvnkhGeEJ1WRATjFPMaBL4KUuurcJNY+YqJYKTjXoaR+q6YRwFRfRZ5WcLKXAB0Cq7kxtR+ymd3WUNpba4wZFcwGoxRnU+UBRW+DwpERBvKaW1OIafRVeorCDRVmwnVOYWKQVt20AmcCutcN6hjKgm4+zBlV3BtVEyV/1E0omqqrHeMvYSRmmMwbtAqhRVUcs8NgXeR4ZpQqUkQclKkZBCn8Qjj0lriroErRiGAR8DPkVUihDAOQ8oVBBLD6sNRIMyCq3zdeYDWlu5tkIADVOYxEMrJQpbUlU1x+ORh/3hT3X5/6MbXXfM36VCRegOeyyJ4EdWq4ayMux2W66vr1m1LWi5doYh5pa15Xg8cn9/z/FwZNWuc9u5ZLPZEELI93/CatWeBaFqHpO/fhg/jB/Gn3N8p4omISoHjBaDqxDCYkaZYhQ1UlVLayomphAJKaG15XA80K5W9IcuE8kNdVVhC0M3CGl3HAdOpxPrwmKqEpsUKENICh8SwTlKM9C2zVLciGRfNi2VxIFiGEdBqJRiHAamaWSzWkvLxdociSIqPU1B8BNTP2U38XoxvCzqGsWjV8voR8qqhKaBEIVcHgNJKVJIUpxYabnElLJpojhZGytuwiFnqM3O3OLtFIlBPJuqqsSNPRgom1IcNZSi0IKYeR8eI1u0QhtDd5IYlbKsQElczd3dHc+ePsX7CRDr++SkiCxMidGijgOVjSYDSokL7DhMjIMUEat2JVEl1kDSYAwBhfdhIXNrW+CGEe8cVUqgJSsPRDVXZr7S0A/iEWUVQzdQlzU+TLTtShRQJPF0sgWRSHRhsaiYydjnqhBxUBe7gu12x+FwT4yeqt6gjJGWaCb6z15i5w67KUlIZSIRVaSfeoZpRCWbEUklMTJdh/eR1XrFqetpaollSTHQzPEryHcqbbb5O/WE4AhUtOsN9bqlqmuG/QGTItZIjtk4TPjhhtIIp2vyTlA6LRy+rEVEWUVlykcTzSjGnNoYaU2m+CeR+P5jHTevX3B19ZTLywuaqkElxd3DPS9efg0pcP30kn/yNz/myZMLvPdUTclms87eWD2vXr3m5z//OV9//TXbzY5xnOh7ERSs12ucG9luN1xeXi48x9IWqMKSgifFH/Lnfhg/jD/3+E4VTd2pY7fdLh4RAvkIaTUlMTMcY4+xBc1mTbVaEaJD6UhbF6zahk1T0N8/EPyERpLsDZ7NZoOPil6Joqxt12zaOR9OozdrgpuYhg7n3eK1I8bxhtPQUdUVurT0pxNlUWJKS1QJnXx2TVWYqqCfeqZTJzEsVSk5aylQtxVJJYZpoO8GLtZrVFIQxNTr4WGPNrKJNU2NLo2gFykRFkWmkjwspIUprayMPvkgfzIRuqoqhl7S7b0TRaIxhosnVyA2WAzjhJsmrFKUTUuKEJynrEpWqzUhRo7HEyhFmb2TpuORN2/ecPXkQswUpwkXHAohIJdFxThNjLlNp1224LcGZQwRsHVNjI7VboupCt7c3lIUBVdPLgFFV0gb0GhDs9lIoZRd3+fPwcfI/cM9q/WKsii5u7mhbFsUicP+RF3W0tIzFj+T44PHe4UtHx3Vi7KkqiuU1jgnETHRe4rKIA7bJa4fcJNjHEZ8ShQ5O6ksSnov0SU2m1yauchAjNpExWgJcSLGQGEL6tpS1SVKR/pxT7UqMKbk6fUlRWmZXp7wzlO3G4hRkD2dcpyHZPspo0kuEJNHG8Xl1QV1UxJSKynjtqT2Jd2pYxoc280GABfFcDUqBCHLcULWaCpbAGJlEJ0XBV+M7O8OVFXJ1e7iz7MY/CMYQ9/xcH/LNA5s1xt26x3BO8axx4eJrVtRlBJAagtLURhSEsTPOXGQ915UjsM40HUdzgmvUJSIlu12ywcfPGe9Xi1ZhyolKZL/0h/AD+OH8T0c34miaVYRnbqBw6lntle3tuR46hmGQVpkdU3iUc49jQNdf6CqDIXR3L96QWENVWkxCrqjFBMoxee//CWb7Y6h63DTxI3/isk5SCLtNqZYQnVD8pJgbqWNkaKQqStXUVc1x9MRYy37k7QqlFI83D9IWG1Z0vcDzk0cxx5IuDBhjGW9WWNCIKrEFCNffP45McFmu6Ouavo0Md7coFSinVrGaQBy6O7ouNhdUpYlXdcvyFulDVO2Kkgp0Q8jV9fXnA4HXt/cUlrJ8uuGkWPfM4w9XqX/f3v3GxtVlf9x/NNOOzNlpQW3y7TFwQYM4h/+rEW6FQnxl65NNLg82MiKgS5RWZU1SrMrIMKoKGWVZUmkSmRVfKALasQYaepqV2LUGhKgiS5/DIJS3Z2BrtIpUzpDZ87vQWG0tsVzCzMF5v1K7oMezun9DnC/+cydO/d2h0B3nnJcLsWise4H6+Z0P7Ymerz73kc/+9nP1N7erosuyleuZ4gikQ51HO8OCMUlJfr2u6PqSpzoPvPm6j6zFzUJHYtGZYzk8eYpcfKbi263u/vv/uSNIhPG6FhHRJ1GCrcfTV6UHO6InDyjJ3UcP65oR4d+duyYuuIJRWOdyefjefOGyJPnkTcvT8ciEZnEMYXDYbWFw/J6uu/wHWoNyeVy6VjkmHJzcxVpb+++OWeeV1mu7O6P1bKzlTjW/TFMlqv7Oqzuj0fjam/v/up/Xl6e2o8dU/eDVrsUaY8okWjX0f99J6/Xq87OToXD7Sfv/9T9UeuJkxe+S1LWyY9Ku05+dNd9q6O4huYPVU6uW5GO75RwnVC8K64hed7ue0x1HFcsFlXr//6j3Fy33Dk50sl7cGWp++aXiUT3fnLa3TrWHlb4WIda//e/7o/ksrLl8eQqS923QXBlueTKzVIs1imT1f3Ina54XHJ1X7RuEgm5s3PUlZXb/bDik/eDyh96kbKypWjkuLo6Y4p1RHscsxei5HMMT0R1rD2uY+1hxY53Kvfknbq7H5vS/e/b/SiVmDqjHTJyy5i4IpGI2tvD6uzsSD6qJhaN6fjx44p3dV+Dd+rLDllZUo6r+81I5NjJ52Cq+7rCU/c2A9C3SCQi6ez2oyxzHnS3r7/+Wn6/f7DLAGCppaVFl1xyyWCXkRL0I+D8cjb70XkRmhKJhP7zn/9o6NChVveRADA4jDFqb29XSUlJ8llVFxr6EXB+SEU/Oi9CEwAAwGC7MN8KAgAAnGWEJgAAAAuEJgAAAAuEJgAAAAuEJgAAAAuEJgAAAAuEJgAAAAuEJgAAAAuEJgAAAAuEJgAAAAuEJgAAAAuEJgAAAAuEJgAAAAuEJgAAAAuEJgAAAAuEJgAAAAuEJgAAAAuOQ9MHH3ygGTNmqKSkRFlZWXrzzTd/cs22bdt0zTXXyOPx6LLLLtPGjRsHUCoA9EQ/ApBOjkNTJBLRxIkTVVdXZzX/4MGDuvnmm3XDDTeoublZDzzwgO6880698847josFgB+iHwFIpyxjjBnw4qwsbdmyRTNnzux3zqJFi7R161Z99tlnybHf/e53Onr0qBoaGga6awDogX4EINVyUr2DpqYmVVZW9hirqqrSAw880O+aaDSqaDSa/DmRSOjbb7/Vz3/+c2VlZaWqVABnyBij9vZ2lZSUKDv73Ltkkn4EZI5U9KOUh6ZgMCifz9djzOfzKRwO6/jx48rLy+u1pra2Vo8++miqSwOQIi0tLbrkkksGu4xe6EdA5jmb/SjloWkglixZopqamuTPbW1tGjVqlFpaWpSfnz+IlQE4nXA4LL/fr6FDhw52KWcN/Qg4P6WiH6U8NBUVFSkUCvUYC4VCys/P7/NdnSR5PB55PJ5e4/n5+TQp4Dxwrn5sRT8CMs/Z7Ecpv+igoqJCjY2NPcbeffddVVRUpHrXANAD/QjAmXAcmo4dO6bm5mY1NzdL6v4Kb3Nzsw4dOiSp+1T23Llzk/PvvvtuHThwQA8++KD27t2rZ555Rq+++qoWLlx4dl4BgIxFPwKQVsah999/30jqtVVXVxtjjKmurjbTp0/vtWbSpEnG7Xab0aNHmxdffNHRPtva2owk09bW5rRcAGmU7mOVfgSgP6k4Vs/oPk3pEg6HVVBQoLa2Nq4hAM5hmXCsZsJrBC4EqThWz70bqQAAAJyDCE0AAAAWCE0AAAAWCE0AAAAWCE0AAAAWCE0AAAAWCE0AAAAWCE0AAAAWCE0AAAAWCE0AAAAWCE0AAAAWCE0AAAAWCE0AAAAWCE0AAAAWCE0AAAAWCE0AAAAWCE0AAAAWCE0AAAAWCE0AAAAWCE0AAAAWCE0AAAAWCE0AAAAWCE0AAAAWCE0AAAAWCE0AAAAWCE0AAAAWCE0AAAAWCE0AAAAWCE0AAAAWCE0AAAAWCE0AAAAWCE0AAAAWCE0AAAAWCE0AAAAWCE0AAAAWCE0AAAAWCE0AAAAWBhSa6urqVFpaKq/Xq/Lycm3fvv2089euXavLL79ceXl58vv9WrhwoTo7OwdUMAD8EP0IQLo4Dk2bN29WTU2NAoGAdu7cqYkTJ6qqqkqHDx/uc/4rr7yixYsXKxAIaM+ePXr++ee1efNmPfTQQ2dcPIDMRj8CkE6OQ9OaNWt01113ad68ebryyiu1fv16DRkyRC+88EKf8z/++GNNnTpVs2fPVmlpqW688UbddtttP/luEAB+Cv0IQDo5Ck2xWEw7duxQZWXl978gO1uVlZVqamrqc811112nHTt2JJvSgQMHVF9fr5tuuqnf/USjUYXD4R4bAPwQ/QhAuuU4mdza2qp4PC6fz9dj3Ofzae/evX2umT17tlpbW3X99dfLGKOuri7dfffdpz0dXltbq0cffdRJaQAyDP0IQLql/Ntz27Zt08qVK/XMM89o586deuONN7R161atWLGi3zVLlixRW1tbcmtpaUl1mQAyAP0IwJlwdKapsLBQLpdLoVCox3goFFJRUVGfa5YtW6Y5c+bozjvvlCSNHz9ekUhE8+fP19KlS5Wd3Tu3eTweeTweJ6UByDD0IwDp5uhMk9vtVllZmRobG5NjiURCjY2Nqqio6HNNR0dHr0bkcrkkScYYp/UCgCT6EYD0c3SmSZJqampUXV2tyZMna8qUKVq7dq0ikYjmzZsnSZo7d65Gjhyp2tpaSdKMGTO0Zs0a/fKXv1R5ebn279+vZcuWacaMGclmBQADQT8CkE6OQ9OsWbN05MgRLV++XMFgUJMmTVJDQ0PyYsxDhw71eCf38MMPKysrSw8//LC++eYb/eIXv9CMGTP0xBNPnL1XASAj0Y8ApFOWOQ/OSYfDYRUUFKitrU35+fmDXQ6AfmTCsZoJrxG4EKTiWOXZcwAAABYITQAAABYITQAAABYITQAAABYITQAAABYITQAAABYITQAAABYITQAAABYITQAAABYITQAAABYITQAAABYITQAAABYITQAAABYITQAAABYITQAAABYITQAAABYITQAAABYITQAAABYITQAAABYITQAAABYITQAAABYITQAAABYITQAAABYITQAAABYITQAAABYITQAAABYITQAAABYITQAAABYITQAAABYITQAAABYITQAAABYITQAAABYITQAAABYITQAAABYITQAAABYITQAAABYITQAAABYITQAAABYGFJrq6upUWloqr9er8vJybd++/bTzjx49qgULFqi4uFgej0djx45VfX39gAoGgB+iHwFIlxynCzZv3qyamhqtX79e5eXlWrt2raqqqrRv3z6NGDGi1/xYLKZf//rXGjFihF5//XWNHDlSX331lYYNG3Y26geQwehHANIpyxhjnCwoLy/Xtddeq3Xr1kmSEomE/H6/7rvvPi1evLjX/PXr1+upp57S3r17lZuba7WPaDSqaDSa/DkcDsvv96utrU35+flOygWQRuFwWAUFBWk7VulHAPqTin7k6OO5WCymHTt2qLKy8vtfkJ2tyspKNTU19bnmrbfeUkVFhRYsWCCfz6err75aK1euVDwe73c/tbW1KigoSG5+v99JmQAyAP0IQLo5Ck2tra2Kx+Py+Xw9xn0+n4LBYJ9rDhw4oNdff13xeFz19fVatmyZ/vrXv+rxxx/vdz9LlixRW1tbcmtpaXFSJoAMQD8CkG6Or2lyKpFIaMSIEXruuefkcrlUVlamb775Rk899ZQCgUCfazwejzweT6pLA5Bh6EcAzoSj0FRYWCiXy6VQKNRjPBQKqaioqM81xcXFys3NlcvlSo5dccUVCgaDisVicrvdAygbQKajHwFIN0cfz7ndbpWVlamxsTE5lkgk1NjYqIqKij7XTJ06Vfv371cikUiOff755youLqZBARgw+hGAdHN8n6aamhpt2LBBL730kvbs2aN77rlHkUhE8+bNkyTNnTtXS5YsSc6/55579O233+r+++/X559/rq1bt2rlypVasGDB2XsVADIS/QhAOjm+pmnWrFk6cuSIli9frmAwqEmTJqmhoSF5MeahQ4eUnf19FvP7/XrnnXe0cOFCTZgwQSNHjtT999+vRYsWnb1XASAj0Y8ApJPj+zQNhnTf+wXAwGTCsZoJrxG4EAz6fZoAAAAyFaEJAADAAqEJAADAAqEJAADAAqEJAADAAqEJAADAAqEJAADAAqEJAADAAqEJAADAAqEJAADAAqEJAADAAqEJAADAAqEJAADAAqEJAADAAqEJAADAAqEJAADAAqEJAADAAqEJAADAAqEJAADAAqEJAADAAqEJAADAAqEJAADAAqEJAADAAqEJAADAAqEJAADAAqEJAADAAqEJAADAAqEJAADAAqEJAADAAqEJAADAAqEJAADAAqEJAADAAqEJAADAAqEJAADAAqEJAADAAqEJAADAAqEJAADAwoBCU11dnUpLS+X1elVeXq7t27dbrdu0aZOysrI0c+bMgewWAHqhHwFIF8ehafPmzaqpqVEgENDOnTs1ceJEVVVV6fDhw6dd9+WXX+pPf/qTpk2bNuBiAeCH6EcA0slxaFqzZo3uuusuzZs3T1deeaXWr1+vIUOG6IUXXuh3TTwe1+23365HH31Uo0ePPqOCAeAU+hGAdHIUmmKxmHbs2KHKysrvf0F2tiorK9XU1NTvuscee0wjRozQHXfcYbWfaDSqcDjcYwOAH6IfAUg3R6GptbVV8XhcPp+vx7jP51MwGOxzzYcffqjnn39eGzZssN5PbW2tCgoKkpvf73dSJoAMQD8CkG4p/fZce3u75syZow0bNqiwsNB63ZIlS9TW1pbcWlpaUlglgExAPwJwpnKcTC4sLJTL5VIoFOoxHgqFVFRU1Gv+F198oS+//FIzZsxIjiUSie4d5+Ro3759GjNmTK91Ho9HHo/HSWkAMgz9CEC6OTrT5Ha7VVZWpsbGxuRYIpFQY2OjKioqes0fN26cPv30UzU3Nye3W265RTfccIOam5s5zQ1gwOhHANLN0ZkmSaqpqVF1dbUmT56sKVOmaO3atYpEIpo3b54kae7cuRo5cqRqa2vl9Xp19dVX91g/bNgwSeo1DgBO0Y8ApJPj0DRr1iwdOXJEy5cvVzAY1KRJk9TQ0JC8GPPQoUPKzuZG4wBSj34EIJ2yjDFmsIv4KeFwWAUFBWpra1N+fv5glwOgH5lwrGbCawQuBKk4VnkLBgAAYIHQBAAAYIHQBAAAYIHQBAAAYIHQBAAAYIHQBAAAYIHQBAAAYIHQBAAAYIHQBAAAYIHQBAAAYIHQBAAAYIHQBAAAYIHQBAAAYIHQBAAAYIHQBAAAYIHQBAAAYIHQBAAAYIHQBAAAYIHQBAAAYIHQBAAAYIHQBAAAYIHQBAAAYIHQBAAAYIHQBAAAYIHQBAAAYIHQBAAAYIHQBAAAYIHQBAAAYIHQBAAAYIHQBAAAYIHQBAAAYIHQBAAAYIHQBAAAYIHQBAAAYIHQBAAAYIHQBAAAYIHQBAAAYGFAoamurk6lpaXyer0qLy/X9u3b+527YcMGTZs2TcOHD9fw4cNVWVl52vkA4AT9CEC6OA5NmzdvVk1NjQKBgHbu3KmJEyeqqqpKhw8f7nP+tm3bdNttt+n9999XU1OT/H6/brzxRn3zzTdnXDyAzEY/ApBOWcYY42RBeXm5rr32Wq1bt06SlEgk5Pf7dd9992nx4sU/uT4ej2v48OFat26d5s6da7XPcDisgoICtbW1KT8/30m5ANIo3ccq/QhAf1JxrDo60xSLxbRjxw5VVlZ+/wuys1VZWammpiar39HR0aETJ07o4osv7ndONBpVOBzusQHAD9GPAKSbo9DU2tqqeDwun8/XY9zn8ykYDFr9jkWLFqmkpKRHo/ux2tpaFRQUJDe/3++kTAAZgH4EIN3S+u25VatWadOmTdqyZYu8Xm+/85YsWaK2trbk1tLSksYqAWQC+hEAp3KcTC4sLJTL5VIoFOoxHgqFVFRUdNq1q1ev1qpVq/Tee+9pwoQJp53r8Xjk8XiclAYgw9CPAKSbozNNbrdbZWVlamxsTI4lEgk1NjaqoqKi33VPPvmkVqxYoYaGBk2ePHng1QLASfQjAOnm6EyTJNXU1Ki6ulqTJ0/WlClTtHbtWkUiEc2bN0+SNHfuXI0cOVK1tbWSpL/85S9avny5XnnlFZWWliavNbjooot00UUXncWXAiDT0I8ApJPj0DRr1iwdOXJEy5cvVzAY1KRJk9TQ0JC8GPPQoUPKzv7+BNazzz6rWCym3/72tz1+TyAQ0COPPHJm1QPIaPQjAOnk+D5Ng4H7ogDnh0w4VjPhNQIXgkG/TxMAAECmIjQBAABYIDQBAABYIDQBAABYIDQBAABYIDQBAABYIDQBAABYIDQBAABYIDQBAABYIDQBAABYIDQBAABYIDQBAABYIDQBAABYIDQBAABYIDQBAABYIDQBAABYIDQBAABYIDQBAABYIDQBAABYIDQBAABYIDQBAABYIDQBAABYIDQBAABYIDQBAABYIDQBAABYIDQBAABYIDQBAABYIDQBAABYIDQBAABYIDQBAABYIDQBAABYIDQBAABYIDQBAABYIDQBAABYIDQBAABYIDQBAABYIDQBAABYIDQBAABYGFBoqqurU2lpqbxer8rLy7V9+/bTzn/ttdc0btw4eb1ejR8/XvX19QMqFgB+jH4EIF0ch6bNmzerpqZGgUBAO3fu1MSJE1VVVaXDhw/3Of/jjz/WbbfdpjvuuEO7du3SzJkzNXPmTH322WdnXDyAzEY/ApBOWcYY42RBeXm5rr32Wq1bt06SlEgk5Pf7dd9992nx4sW95s+aNUuRSERvv/12cuxXv/qVJk2apPXr1/e5j2g0qmg0mvy5ra1No0aNUktLi/Lz852UCyCNwuGw/H6/jh49qoKCgpTvj34EoD8p6UfGgWg0alwul9myZUuP8blz55pbbrmlzzV+v9/87W9/6zG2fPlyM2HChH73EwgEjCQ2NrbzdPviiy+ctJYBoR+xsbHZbGezH+XIgdbWVsXjcfl8vh7jPp9Pe/fu7XNNMBjsc34wGOx3P0uWLFFNTU3y56NHj+rSSy/VoUOH0vLu9Ww6lXTPx3el1D44zufaT52Fufjii1O+L/qRc+fz/y1qHxznc+2p6EeOQlO6eDweeTyeXuMFBQXn3T/aKfn5+dQ+CKh9cGRnXzhfzKUfnVuofXCcz7WfzX7k6DcVFhbK5XIpFAr1GA+FQioqKupzTVFRkaP5AGCDfgQg3RyFJrfbrbKyMjU2NibHEomEGhsbVVFR0eeaioqKHvMl6d133+13PgDYoB8BSDunF0Ft2rTJeDwes3HjRrN7924zf/58M2zYMBMMBo0xxsyZM8csXrw4Of+jjz4yOTk5ZvXq1WbPnj0mEAiY3Nxc8+mnn1rvs7Oz0wQCAdPZ2em03EFH7YOD2gdHumunHzlD7YOD2gdHKmp3HJqMMebpp582o0aNMm6320yZMsV88sknyT+bPn26qa6u7jH/1VdfNWPHjjVut9tcddVVZuvWrWdUNACcQj8CkC6O79MEAACQiS6cr7gAAACkEKEJAADAAqEJAADAAqEJAADAwjkTmurq6lRaWiqv16vy8nJt3779tPNfe+01jRs3Tl6vV+PHj1d9fX2aKu3NSe0bNmzQtGnTNHz4cA0fPlyVlZU/+VpTyenf+ymbNm1SVlaWZs6cmdoC++G07qNHj2rBggUqLi6Wx+PR2LFjB+3/jNPa165dq8svv1x5eXny+/1auHChOjs701Tt9z744APNmDFDJSUlysrK0ptvvvmTa7Zt26ZrrrlGHo9Hl112mTZu3JjyOs8G+tHgOF/7kURPypieNNhf3zOm+14rbrfbvPDCC+bf//63ueuuu8ywYcNMKBTqc/5HH31kXC6XefLJJ83u3bvNww8/7PheK2eL09pnz55t6urqzK5du8yePXvM73//e1NQUGC+/vrrNFfuvPZTDh48aEaOHGmmTZtmfvOb36Sn2B9wWnc0GjWTJ082N910k/nwww/NwYMHzbZt20xzc3OaK3de+8svv2w8Ho95+eWXzcGDB80777xjiouLzcKFC9NcuTH19fVm6dKl5o033jCSej0o98cOHDhghgwZYmpqaszu3bvN008/bVwul2loaEhPwQNEP6IfOUVPypyedE6EpilTppgFCxYkf47H46akpMTU1tb2Of/WW281N998c4+x8vJy84c//CGldfbFae0/1tXVZYYOHWpeeumlVJXYr4HU3tXVZa677jrz97//3VRXVw9Kk3Ja97PPPmtGjx5tYrFYukrsl9PaFyxYYP7v//6vx1hNTY2ZOnVqSuv8KTYN6sEHHzRXXXVVj7FZs2aZqqqqFFZ25uhH9COn6EmZ05MG/eO5WCymHTt2qLKyMjmWnZ2tyspKNTU19bmmqampx3xJqqqq6nd+qgyk9h/r6OjQiRMn0vJU+B8aaO2PPfaYRowYoTvuuCMdZfYykLrfeustVVRUaMGCBfL5fLr66qu1cuVKxePxdJUtaWC1X3fdddqxY0fydPmBAwdUX1+vm266KS01n4lz5Th1gn5EP3KKnpRZPSnnbBflVGtrq+LxuHw+X49xn8+nvXv39rkmGAz2OT8YDKaszr4MpPYfW7RokUpKSnr9Q6baQGr/8MMP9fzzz6u5uTkNFfZtIHUfOHBA//rXv3T77bervr5e+/fv17333qsTJ04oEAiko2xJA6t99uzZam1t1fXXXy9jjLq6unT33XfroYceSkfJZ6S/4zQcDuv48ePKy8sbpMr6Rz+iHzlFT8qsnjToZ5oy2apVq7Rp0yZt2bJFXq93sMs5rfb2ds2ZM0cbNmxQYWHhYJfjSCKR0IgRI/Tcc8+prKxMs2bN0tKlS7V+/frBLu0nbdu2TStXrtQzzzyjnTt36o033tDWrVu1YsWKwS4NFxj6UfrQk85fg36mqbCwUC6XS6FQqMd4KBRSUVFRn2uKiooczU+VgdR+yurVq7Vq1Sq99957mjBhQirL7JPT2r/44gt9+eWXmjFjRnIskUhIknJycrRv3z6NGTMmtUVrYH/nxcXFys3NlcvlSo5dccUVCgaDisVicrvdKa35lIHUvmzZMs2ZM0d33nmnJGn8+PGKRCKaP3++li5dquzsc/d9T3/HaX5+/jl5lkmiH9GPnKMnZVZPGvRX53a7VVZWpsbGxuRYIpFQY2OjKioq+lxTUVHRY74kvfvuu/3OT5WB1C5JTz75pFasWKGGhgZNnjw5HaX24rT2cePG6dNPP1Vzc3Nyu+WWW3TDDTeoublZfr//nKxbkqZOnar9+/cnm6okff755youLk5bc5IGVntHR0evJnSq0Zpz/LGR58px6gT9iH6U6toletJgOSvHqtMr1FNh06ZNxuPxmI0bN5rdu3eb+fPnm2HDhplgMGiMMWbOnDlm8eLFyfkfffSRycnJMatXrzZ79uwxgUBgUL/i66T2VatWGbfbbV5//XXz3//+N7m1t7ef87X/2GB9W8Vp3YcOHTJDhw41f/zjH82+ffvM22+/bUaMGGEef/zxc772QCBghg4dav7xj3+YAwcOmH/+859mzJgx5tZbb0177e3t7WbXrl1m165dRpJZs2aN2bVrl/nqq6+MMcYsXrzYzJkzJzn/1Nd7//znP5s9e/aYurq68+aWA/Qj+pET9KTM6UnnRGgyxpinn37ajBo1yrjdbjNlyhTzySefJP9s+vTpprq6usf8V1991YwdO9a43W5z1VVXma1bt6a54u85qf3SSy81knptgUAg/YUb53/vPzSYTcpp3R9//LEpLy83Ho/HjB492jzxxBOmq6srzVV3c1L7iRMnzCOPPGLGjBljvF6v8fv95t577zXfffdd2ut+//33+/y/e6re6upqM3369F5rJk2aZNxutxk9erR58cUX0173QNCPAukv3Jy//cgYelKm9KQsY87x82kAAADngEG/pgkAAOB8QGgCAACwQGgCAACwQGgCAACwQGgCAACwQGgCAACwQGgCAACwQGgCAACwQGgCAACwQGgCAACwQGgCAACw8P+J+/AF07Vw5wAAAABJRU5ErkJggg==\",\n \"text/plain\": [\n \"
\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"# The mistakes look all quite reasonable.\\n\",\n \"show_img_grid(\\n\",\n \" [batch['image'][idx] for idx in error_idxs[:9]],\\n\",\n \" [f'pred: {imagenet2012_label(preds_labels[idx][0])}\\\\n'\\n\",\n \" f'label: {imagenet2012_label(preds_labels[idx][1])}'\\n\",\n \" for idx in error_idxs[:9]],\\n\",\n \")\\n\",\n \"plt.tight_layout()\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 53,\n \"metadata\": {\n \"outputId\": \"4fae2533-5598-4f2e-c133-50bfba463311\"\n },\n \"outputs\": [\n {\n \"name\": \"stderr\",\n \"output_type\": \"stream\",\n \"text\": [\n \"INFO:absl:Constructing tf.data.Dataset imagenette for split validation[0:3925], from /root/tensorflow_datasets/imagenette/full-size-v2/1.0.0\\n\",\n \"WARNING:absl:options.experimental_threading is deprecated. Use options.threading instead.\\n\"\n ]\n }\n ],\n \"source\": [\n \"# Define parallelized inference function in separate cell so the cached\\n\",\n \"# compilation can be used if below cell is executed multiple times.\\n\",\n \"@jax.pmap\\n\",\n \"def p_get_logits(images):\\n\",\n \" return model.apply({'params': state.params, 'batch_stats': state.batch_stats},\\n\",\n \" images, train=False)\\n\",\n \"\\n\",\n \"eval_iter = train.create_input_iter(dataset_builder, config.batch_size,\\n\",\n \" input_pipeline.IMAGE_SIZE, tf.float32,\\n\",\n \" train=False, cache=False, shuffle_buffer_size=None,\\n\",\n \" prefetch=1)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 54,\n \"metadata\": {\n \"outputId\": \"d01e1993-28ab-4a4a-ac58-01c83b80e6c9\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"Step 1/7...\\n\",\n \"Step 2/7...\\n\",\n \"Step 3/7...\\n\",\n \"Step 4/7...\\n\",\n \"Step 5/7...\\n\",\n \"Step 6/7...\\n\",\n \"Step 7/7...\\n\"\n ]\n },\n {\n \"data\": {\n \"text/plain\": [\n \"Array(0.9118304, dtype=float32)\"\n ]\n },\n \"execution_count\": 54,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"# Compute accuracy.\\n\",\n \"eval_steps = dataset_builder.info.splits['validation'].num_examples // config.batch_size\\n\",\n \"count = correct = 0\\n\",\n \"for step, batch in zip(range(eval_steps), eval_iter):\\n\",\n \" labels = [imagenette_imagenet2012(label) for label in batch['label'].flatten()]\\n\",\n \" logits = p_get_logits(batch['image'])\\n\",\n \" logits = logits.reshape([-1, logits.shape[-1]])\\n\",\n \" print(f'Step {step+1}/{eval_steps}...')\\n\",\n \" count += len(labels)\\n\",\n \" correct += (logits.argmax(axis=-1) == jnp.array(labels)).sum()\\n\",\n \"\\n\",\n \"correct / count\"\n ]\n }\n ],\n \"metadata\": {\n \"accelerator\": \"GPU\",\n \"gpuClass\": \"standard\",\n \"language_info\": {\n \"name\": \"python\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 0\n}\n" }, { "path": "examples/imagenet/imagenet_benchmark.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Benchmark for the ImageNet example.\"\"\"\n\nimport time\n\nfrom absl import flags\nfrom absl.testing import absltest\nfrom absl.testing.flagsaver import flagsaver\nfrom flax.testing import Benchmark\nimport jax\nimport numpy as np\n\n# Local imports.\nimport main\nfrom configs import v100_x8_mixed_precision as config_lib\n\n\n# Parse absl flags test_srcdir and test_tmpdir.\njax.config.parse_flags_with_absl()\n\nFLAGS = flags.FLAGS\n\n\nclass ImagenetBenchmark(Benchmark):\n \"\"\"Benchmarks for the ImageNet Flax example.\"\"\"\n\n @flagsaver\n def _test_8x_v100_half_precision(\n self, num_epochs: int, min_accuracy, max_accuracy\n ):\n \"\"\"Utility to benchmark ImageNet on 8xV100 GPUs. Use in your test func.\"\"\"\n # Prepare and set flags defined in main.py.\n config = config_lib.get_config()\n config.num_epochs = num_epochs\n workdir = self.get_tmp_model_dir()\n\n FLAGS.workdir = workdir\n FLAGS.config = config\n\n start_time = time.time()\n main.main([])\n benchmark_time = time.time() - start_time\n summaries = self.read_summaries(workdir)\n\n # Summaries contain all the information necessary for the regression\n # metrics.\n wall_time, _, eval_accuracy = zip(*summaries['eval_accuracy'])\n wall_time = np.array(wall_time)\n sec_per_epoch = np.mean(wall_time[1:] - wall_time[:-1])\n end_accuracy = eval_accuracy[-1]\n\n # Assertions are deferred until the test finishes, so the metrics are\n # always reported and benchmark success is determined based on *all*\n # assertions.\n self.assertBetween(end_accuracy, min_accuracy, max_accuracy)\n\n # Use the reporting API to report single or multiple metrics/extras.\n self.report_wall_time(benchmark_time)\n self.report_metrics(\n {'sec_per_epoch': sec_per_epoch, 'accuracy': end_accuracy}\n )\n\n def test_8x_v100_half_precision_short(self):\n \"\"\"Run ImageNet on 8x V100 GPUs in half precision for 2 epochs.\"\"\"\n self._test_8x_v100_half_precision(\n num_epochs=2, min_accuracy=0.06, max_accuracy=0.09\n )\n self.report_extras({\n 'description': 'Short (2 epochs) 8 x V100 test for ImageNet ResNet50.',\n 'model_name': 'resnet50',\n 'parameters': 'hp=true,bs=2048,num_epochs=2',\n 'implementation': 'linen',\n })\n\n def test_8x_v100_half_precision_full(self):\n \"\"\"Run ImageNet on 8x V100 GPUs in half precision for full 90 epochs.\"\"\"\n self._test_8x_v100_half_precision(\n num_epochs=90, min_accuracy=0.76, max_accuracy=0.77\n )\n self.report_extras({\n 'description': 'Full (90 epochs) 8 x V100 test for ImageNet ResNet50.',\n 'model_name': 'resnet50',\n 'parameters': 'hp=true,bs=2048,num_epochs=90',\n 'implementation': 'linen',\n })\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "examples/imagenet/imagenet_fake_data_benchmark.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Benchmark for the ImageNet example using fake data for quick perf results.\n\nThis script doesn't need the dataset, but it needs the dataset metadata.\nThat can be fetched with the script `flax/tests/download_dataset_metadata.sh`.\n\"\"\"\n\nimport pathlib\nimport time\n\nfrom absl.testing import absltest\nfrom flax.testing import Benchmark\nimport jax\nimport tensorflow_datasets as tfds\n\n# Local imports.\nfrom configs import fake_data_benchmark as config_lib\nimport train\n\n# Parse absl flags test_srcdir and test_tmpdir.\njax.config.parse_flags_with_absl()\n\n\nclass ImagenetBenchmarkFakeData(Benchmark):\n \"\"\"Runs ImageNet using fake data for quickly measuring performance.\"\"\"\n\n def test_fake_data(self):\n workdir = self.get_tmp_model_dir()\n config = config_lib.get_config()\n # Go two directories up to the root of the flax directory.\n flax_root_dir = pathlib.Path(__file__).absolute().parents[2]\n data_dir = str(flax_root_dir) + '/.tfds/metadata'\n\n # Warm-up first so that we are not measuring just compilation.\n with tfds.testing.mock_data(num_examples=1024, data_dir=data_dir):\n train.train_and_evaluate(config, workdir)\n\n start_time = time.time()\n with tfds.testing.mock_data(num_examples=1024, data_dir=data_dir):\n train.train_and_evaluate(config, workdir)\n benchmark_time = time.time() - start_time\n\n self.report_wall_time(benchmark_time)\n self.report_extras({\n 'description': 'ImageNet ResNet50 with fake data',\n 'model_name': 'resnet50',\n 'parameters': f'hp=true,bs={config.batch_size}',\n })\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "examples/imagenet/input_pipeline.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"ImageNet input pipeline.\"\"\"\n\nimport jax\nimport tensorflow as tf\nimport tensorflow_datasets as tfds\n\n\nIMAGE_SIZE = 224\nCROP_PADDING = 32\nMEAN_RGB = [0.485 * 255, 0.456 * 255, 0.406 * 255]\nSTDDEV_RGB = [0.229 * 255, 0.224 * 255, 0.225 * 255]\n\n\ndef distorted_bounding_box_crop(\n image_bytes,\n bbox,\n min_object_covered=0.1,\n aspect_ratio_range=(0.75, 1.33),\n area_range=(0.05, 1.0),\n max_attempts=100,\n):\n \"\"\"Generates cropped_image using one of the bboxes randomly distorted.\n\n See `tf.image.sample_distorted_bounding_box` for more documentation.\n\n Args:\n image_bytes: `Tensor` of binary image data.\n bbox: `Tensor` of bounding boxes arranged `[1, num_boxes, coords]`\n where each coordinate is [0, 1) and the coordinates are arranged\n as `[ymin, xmin, ymax, xmax]`. If num_boxes is 0 then use the whole\n image.\n min_object_covered: An optional `float`. Defaults to `0.1`. The cropped\n area of the image must contain at least this fraction of any bounding\n box supplied.\n aspect_ratio_range: An optional list of `float`s. The cropped area of the\n image must have an aspect ratio = width / height within this range.\n area_range: An optional list of `float`s. The cropped area of the image\n must contain a fraction of the supplied image within this range.\n max_attempts: An optional `int`. Number of attempts at generating a cropped\n region of the image of the specified constraints. After `max_attempts`\n failures, return the entire image.\n Returns:\n cropped image `Tensor`\n \"\"\"\n shape = tf.io.extract_jpeg_shape(image_bytes)\n sample_distorted_bounding_box = tf.image.sample_distorted_bounding_box(\n shape,\n bounding_boxes=bbox,\n min_object_covered=min_object_covered,\n aspect_ratio_range=aspect_ratio_range,\n area_range=area_range,\n max_attempts=max_attempts,\n use_image_if_no_bounding_boxes=True,\n )\n bbox_begin, bbox_size, _ = sample_distorted_bounding_box\n\n # Crop the image to the specified bounding box.\n offset_y, offset_x, _ = tf.unstack(bbox_begin)\n target_height, target_width, _ = tf.unstack(bbox_size)\n crop_window = tf.stack([offset_y, offset_x, target_height, target_width])\n image = tf.io.decode_and_crop_jpeg(image_bytes, crop_window, channels=3)\n\n return image\n\n\ndef _resize(image, image_size):\n return tf.image.resize(\n [image], [image_size, image_size], method=tf.image.ResizeMethod.BICUBIC\n )[0]\n\n\ndef _at_least_x_are_equal(a, b, x):\n \"\"\"At least `x` of `a` and `b` `Tensors` are equal.\"\"\"\n match = tf.equal(a, b)\n match = tf.cast(match, tf.int32)\n return tf.greater_equal(tf.reduce_sum(match), x)\n\n\ndef _decode_and_random_crop(image_bytes, image_size):\n \"\"\"Make a random crop of image_size.\"\"\"\n bbox = tf.constant([0.0, 0.0, 1.0, 1.0], dtype=tf.float32, shape=[1, 1, 4])\n image = distorted_bounding_box_crop(\n image_bytes,\n bbox,\n min_object_covered=0.1,\n aspect_ratio_range=(3.0 / 4, 4.0 / 3.0),\n area_range=(0.08, 1.0),\n max_attempts=10,\n )\n original_shape = tf.io.extract_jpeg_shape(image_bytes)\n bad = _at_least_x_are_equal(original_shape, tf.shape(image), 3)\n\n image = tf.cond(\n bad,\n lambda: _decode_and_center_crop(image_bytes, image_size),\n lambda: _resize(image, image_size),\n )\n\n return image\n\n\ndef _decode_and_center_crop(image_bytes, image_size):\n \"\"\"Crops to center of image with padding then scales image_size.\"\"\"\n shape = tf.io.extract_jpeg_shape(image_bytes)\n image_height = shape[0]\n image_width = shape[1]\n\n padded_center_crop_size = tf.cast(\n (\n (image_size / (image_size + CROP_PADDING))\n * tf.cast(tf.minimum(image_height, image_width), tf.float32)\n ),\n tf.int32,\n )\n\n offset_height = ((image_height - padded_center_crop_size) + 1) // 2\n offset_width = ((image_width - padded_center_crop_size) + 1) // 2\n crop_window = tf.stack([\n offset_height,\n offset_width,\n padded_center_crop_size,\n padded_center_crop_size,\n ])\n image = tf.io.decode_and_crop_jpeg(image_bytes, crop_window, channels=3)\n image = _resize(image, image_size)\n\n return image\n\n\ndef normalize_image(image):\n image -= tf.constant(MEAN_RGB, shape=[1, 1, 3], dtype=image.dtype)\n image /= tf.constant(STDDEV_RGB, shape=[1, 1, 3], dtype=image.dtype)\n return image\n\n\ndef preprocess_for_train(image_bytes, dtype=tf.float32, image_size=IMAGE_SIZE):\n \"\"\"Preprocesses the given image for training.\n\n Args:\n image_bytes: `Tensor` representing an image binary of arbitrary size.\n dtype: data type of the image.\n image_size: image size.\n\n Returns:\n A preprocessed image `Tensor`.\n \"\"\"\n image = _decode_and_random_crop(image_bytes, image_size)\n image = tf.reshape(image, [image_size, image_size, 3])\n image = tf.image.random_flip_left_right(image)\n image = normalize_image(image)\n image = tf.image.convert_image_dtype(image, dtype=dtype)\n return image\n\n\ndef preprocess_for_eval(image_bytes, dtype=tf.float32, image_size=IMAGE_SIZE):\n \"\"\"Preprocesses the given image for evaluation.\n\n Args:\n image_bytes: `Tensor` representing an image binary of arbitrary size.\n dtype: data type of the image.\n image_size: image size.\n\n Returns:\n A preprocessed image `Tensor`.\n \"\"\"\n image = _decode_and_center_crop(image_bytes, image_size)\n image = tf.reshape(image, [image_size, image_size, 3])\n image = normalize_image(image)\n image = tf.image.convert_image_dtype(image, dtype=dtype)\n return image\n\n\ndef create_split(\n dataset_builder,\n batch_size,\n train,\n dtype=tf.float32,\n image_size=IMAGE_SIZE,\n cache=False,\n shuffle_buffer_size=2_000,\n prefetch=10,\n):\n \"\"\"Creates a split from the ImageNet dataset using TensorFlow Datasets.\n\n Args:\n dataset_builder: TFDS dataset builder for ImageNet.\n batch_size: the batch size returned by the data pipeline.\n train: Whether to load the train or evaluation split.\n dtype: data type of the image.\n image_size: The target size of the images.\n cache: Whether to cache the dataset.\n shuffle_buffer_size: Size of the shuffle buffer.\n prefetch: Number of items to prefetch in the dataset.\n Returns:\n A `tf.data.Dataset`.\n \"\"\"\n if train:\n train_examples = dataset_builder.info.splits['train'].num_examples\n split_size = train_examples // jax.process_count()\n start = jax.process_index() * split_size\n split = f'train[{start}:{start + split_size}]'\n else:\n validate_examples = dataset_builder.info.splits['validation'].num_examples\n split_size = validate_examples // jax.process_count()\n start = jax.process_index() * split_size\n split = f'validation[{start}:{start + split_size}]'\n\n def decode_example(example):\n if train:\n image = preprocess_for_train(example['image'], dtype, image_size)\n else:\n image = preprocess_for_eval(example['image'], dtype, image_size)\n return {'image': image, 'label': example['label']}\n\n ds = dataset_builder.as_dataset(\n split=split,\n decoders={\n 'image': tfds.decode.SkipDecoding(),\n },\n )\n options = tf.data.Options()\n options.experimental_threading.private_threadpool_size = 48\n ds = ds.with_options(options)\n\n if cache:\n ds = ds.cache()\n\n if train:\n ds = ds.repeat()\n ds = ds.shuffle(shuffle_buffer_size, seed=0)\n\n ds = ds.map(decode_example, num_parallel_calls=tf.data.experimental.AUTOTUNE)\n ds = ds.batch(batch_size, drop_remainder=True)\n\n if not train:\n ds = ds.repeat()\n\n ds = ds.prefetch(prefetch)\n\n return ds\n" }, { "path": "examples/imagenet/main.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Main file for running the ImageNet example.\n\nThis file is intentionally kept short. The majority for logic is in libraries\nthat can be easily tested and imported in Colab.\n\"\"\"\n\nfrom absl import app\nfrom absl import flags\nfrom absl import logging\nfrom clu import platform\nimport jax\nfrom ml_collections import config_flags\nimport tensorflow as tf\n\nimport train\n\n\nFLAGS = flags.FLAGS\n\nflags.DEFINE_string('workdir', None, 'Directory to store model data.')\nconfig_flags.DEFINE_config_file(\n 'config',\n None,\n 'File path to the training hyperparameter configuration.',\n lock_config=True,\n)\n\n\ndef main(argv):\n if len(argv) > 1:\n raise app.UsageError('Too many command-line arguments.')\n\n # Hide any GPUs from TensorFlow. Otherwise TF might reserve memory and make\n # it unavailable to JAX.\n tf.config.experimental.set_visible_devices([], 'GPU')\n\n logging.info('JAX process: %d / %d', jax.process_index(), jax.process_count())\n logging.info('JAX local devices: %r', jax.local_devices())\n\n # Add a note so that we can tell which task is which JAX host.\n # (Depending on the platform task 0 is not guaranteed to be host 0)\n platform.work_unit().set_task_status(\n f'process_index: {jax.process_index()}, '\n f'process_count: {jax.process_count()}'\n )\n platform.work_unit().create_artifact(\n platform.ArtifactType.DIRECTORY, FLAGS.workdir, 'workdir'\n )\n\n train.train_and_evaluate(FLAGS.config, FLAGS.workdir)\n\n\nif __name__ == '__main__':\n flags.mark_flags_as_required(['config', 'workdir'])\n app.run(main)\n" }, { "path": "examples/imagenet/models.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Flax implementation of ResNet V1.5.\"\"\"\n\n# See issue #620.\n# pytype: disable=wrong-arg-count\n\nfrom functools import partial\nfrom typing import Any, Tuple\nfrom collections.abc import Callable, Sequence\n\nfrom flax import linen as nn\nimport jax.numpy as jnp\n\nModuleDef = Any\n\n\nclass ResNetBlock(nn.Module):\n \"\"\"ResNet block.\"\"\"\n\n filters: int\n conv: ModuleDef\n norm: ModuleDef\n act: Callable\n strides: tuple[int, int] = (1, 1)\n\n @nn.compact\n def __call__(\n self,\n x,\n ):\n residual = x\n y = self.conv(self.filters, (3, 3), self.strides)(x)\n y = self.norm()(y)\n y = self.act(y)\n y = self.conv(self.filters, (3, 3))(y)\n y = self.norm(scale_init=nn.initializers.zeros_init())(y)\n\n if residual.shape != y.shape:\n residual = self.conv(\n self.filters, (1, 1), self.strides, name='conv_proj'\n )(residual)\n residual = self.norm(name='norm_proj')(residual)\n\n return self.act(residual + y)\n\n\nclass BottleneckResNetBlock(nn.Module):\n \"\"\"Bottleneck ResNet block.\"\"\"\n\n filters: int\n conv: ModuleDef\n norm: ModuleDef\n act: Callable\n strides: tuple[int, int] = (1, 1)\n\n @nn.compact\n def __call__(self, x):\n residual = x\n y = self.conv(self.filters, (1, 1))(x)\n y = self.norm()(y)\n y = self.act(y)\n y = self.conv(self.filters, (3, 3), self.strides)(y)\n y = self.norm()(y)\n y = self.act(y)\n y = self.conv(self.filters * 4, (1, 1))(y)\n y = self.norm(scale_init=nn.initializers.zeros_init())(y)\n\n if residual.shape != y.shape:\n residual = self.conv(\n self.filters * 4, (1, 1), self.strides, name='conv_proj'\n )(residual)\n residual = self.norm(name='norm_proj')(residual)\n\n return self.act(residual + y)\n\n\nclass ResNet(nn.Module):\n \"\"\"ResNetV1.5.\"\"\"\n\n stage_sizes: Sequence[int]\n block_cls: ModuleDef\n num_classes: int\n num_filters: int = 64\n dtype: Any = jnp.float32\n act: Callable = nn.relu\n conv: ModuleDef = nn.Conv\n\n @nn.compact\n def __call__(self, x, train: bool = True):\n conv = partial(self.conv, use_bias=False, dtype=self.dtype)\n norm = partial(\n nn.BatchNorm,\n use_running_average=not train,\n momentum=0.9,\n epsilon=1e-5,\n dtype=self.dtype,\n axis_name='batch',\n )\n\n x = conv(\n self.num_filters,\n (7, 7),\n (2, 2),\n padding=[(3, 3), (3, 3)],\n name='conv_init',\n )(x)\n x = norm(name='bn_init')(x)\n x = nn.relu(x)\n x = nn.max_pool(x, (3, 3), strides=(2, 2), padding='SAME')\n for i, block_size in enumerate(self.stage_sizes):\n for j in range(block_size):\n strides = (2, 2) if i > 0 and j == 0 else (1, 1)\n x = self.block_cls(\n self.num_filters * 2**i,\n strides=strides,\n conv=conv,\n norm=norm,\n act=self.act,\n )(x)\n x = jnp.mean(x, axis=(1, 2))\n x = nn.Dense(self.num_classes, dtype=self.dtype)(x)\n x = jnp.asarray(x, self.dtype)\n return x\n\n\nResNet18 = partial(ResNet, stage_sizes=[2, 2, 2, 2], block_cls=ResNetBlock)\nResNet34 = partial(ResNet, stage_sizes=[3, 4, 6, 3], block_cls=ResNetBlock)\nResNet50 = partial(\n ResNet, stage_sizes=[3, 4, 6, 3], block_cls=BottleneckResNetBlock\n)\nResNet101 = partial(\n ResNet, stage_sizes=[3, 4, 23, 3], block_cls=BottleneckResNetBlock\n)\nResNet152 = partial(\n ResNet, stage_sizes=[3, 8, 36, 3], block_cls=BottleneckResNetBlock\n)\nResNet200 = partial(\n ResNet, stage_sizes=[3, 24, 36, 3], block_cls=BottleneckResNetBlock\n)\n\n\nResNet18Local = partial(\n ResNet, stage_sizes=[2, 2, 2, 2], block_cls=ResNetBlock, conv=nn.ConvLocal\n)\n\n\n# Used for testing only.\n_ResNet1 = partial(ResNet, stage_sizes=[1], block_cls=ResNetBlock)\n_ResNet1Local = partial(\n ResNet, stage_sizes=[1], block_cls=ResNetBlock, conv=nn.ConvLocal\n)\n" }, { "path": "examples/imagenet/models_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for flax.examples.imagenet.models.\"\"\"\n\nfrom absl.testing import absltest\nfrom absl.testing import parameterized\n\nimport jax\nfrom jax import numpy as jnp\n\nimport models\n\n\njax.config.update('jax_disable_most_optimizations', True)\n\n\nclass ResNetTest(parameterized.TestCase):\n \"\"\"Test cases for ResNet v1.5 model definition.\"\"\"\n\n def test_resnet_model(self):\n \"\"\"Tests ResNet V1.5 model definition and output (variables).\"\"\"\n rng = jax.random.key(0)\n model_def = models.ResNet50(num_classes=10, dtype=jnp.float32)\n variables = model_def.init(rng, jnp.ones((8, 224, 224, 3), jnp.float32))\n\n self.assertLen(variables, 2)\n # Resnet50 model will create parameters for the following layers:\n # conv + batch_norm = 2\n # BottleneckResNetBlock in stages: [3, 4, 6, 3] = 16\n # Followed by a Dense layer = 1\n self.assertLen(variables['params'], 19)\n\n @parameterized.product(model=(models.ResNet18, models.ResNet18Local))\n def test_resnet_18_model(self, model):\n \"\"\"Tests ResNet18 V1.5 model definition and output (variables).\"\"\"\n rng = jax.random.key(0)\n model_def = model(num_classes=2, dtype=jnp.float32)\n variables = model_def.init(rng, jnp.ones((1, 64, 64, 3), jnp.float32))\n\n self.assertLen(variables, 2)\n self.assertLen(variables['params'], 11)\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "examples/imagenet/requirements.txt", "content": "absl-py==1.0.0\nclu==0.0.6\nflax==0.6.5\n-f https://storage.googleapis.com/jax-releases/libtpu_releases.html\n-f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\njax[cuda11_cudnn805]>=0.3.16 # change to jax[tpu] if running on tpus\nml-collections==0.1.0\nnumpy==1.22.0\noptax==0.1.3\ntensorflow==2.11.1\ntensorflow-datasets==4.4.0\n" }, { "path": "examples/imagenet/train.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"ImageNet example.\n\nThis script trains a ResNet-50 on the ImageNet dataset.\nThe data is loaded using tensorflow_datasets.\n\"\"\"\n\nimport functools\nimport time\nfrom typing import Any\n\nfrom absl import logging\nfrom clu import metric_writers\nfrom clu import periodic_actions\nfrom flax import jax_utils\nfrom flax.training import checkpoints\nfrom flax.training import common_utils\nfrom flax.training import dynamic_scale as dynamic_scale_lib\nfrom flax.training import train_state\nimport jax\nfrom jax import lax\nimport jax.numpy as jnp\nfrom jax import random\nimport ml_collections\nimport optax\nimport orbax.checkpoint as ocp\nimport tensorflow as tf\nimport tensorflow_datasets as tfds\n\nimport input_pipeline\nimport models\n\n\nNUM_CLASSES = 1000\n\n\ndef create_model(*, model_cls, half_precision, **kwargs):\n platform = jax.local_devices()[0].platform\n if half_precision:\n if platform == 'tpu':\n model_dtype = jnp.bfloat16\n else:\n model_dtype = jnp.float16\n else:\n model_dtype = jnp.float32\n return model_cls(num_classes=NUM_CLASSES, dtype=model_dtype, **kwargs)\n\n\ndef initialized(key, image_size, model):\n input_shape = (1, image_size, image_size, 3)\n\n @jax.jit\n def init(*args):\n return model.init(*args)\n\n variables = init({'params': key}, jnp.ones(input_shape, model.dtype))\n return variables['params'], variables['batch_stats']\n\n\ndef cross_entropy_loss(logits, labels):\n one_hot_labels = common_utils.onehot(labels, num_classes=NUM_CLASSES)\n xentropy = optax.softmax_cross_entropy(logits=logits, labels=one_hot_labels)\n return jnp.mean(xentropy)\n\n\ndef compute_metrics(logits, labels):\n loss = cross_entropy_loss(logits, labels)\n accuracy = jnp.mean(jnp.argmax(logits, -1) == labels)\n metrics = {\n 'loss': loss,\n 'accuracy': accuracy,\n }\n metrics = lax.pmean(metrics, axis_name='batch')\n return metrics\n\n\ndef create_learning_rate_fn(\n config: ml_collections.ConfigDict,\n base_learning_rate: float,\n steps_per_epoch: int,\n):\n \"\"\"Create learning rate schedule.\"\"\"\n warmup_fn = optax.linear_schedule(\n init_value=0.0,\n end_value=base_learning_rate,\n transition_steps=config.warmup_epochs * steps_per_epoch,\n )\n cosine_epochs = max(config.num_epochs - config.warmup_epochs, 1)\n cosine_fn = optax.cosine_decay_schedule(\n init_value=base_learning_rate, decay_steps=cosine_epochs * steps_per_epoch\n )\n schedule_fn = optax.join_schedules(\n schedules=[warmup_fn, cosine_fn],\n boundaries=[config.warmup_epochs * steps_per_epoch],\n )\n return schedule_fn\n\n\ndef train_step(state, batch, learning_rate_fn):\n \"\"\"Perform a single training step.\"\"\"\n\n def loss_fn(params):\n \"\"\"loss function used for training.\"\"\"\n logits, new_model_state = state.apply_fn(\n {'params': params, 'batch_stats': state.batch_stats},\n batch['image'],\n mutable=['batch_stats'],\n )\n loss = cross_entropy_loss(logits, batch['label'])\n weight_penalty_params = jax.tree_util.tree_leaves(params)\n weight_decay = 0.0001\n weight_l2 = sum(\n jnp.sum(x**2) for x in weight_penalty_params if x.ndim > 1\n )\n weight_penalty = weight_decay * 0.5 * weight_l2\n loss = loss + weight_penalty\n return loss, (new_model_state, logits)\n\n step = state.step\n dynamic_scale = state.dynamic_scale\n lr = learning_rate_fn(step)\n\n if dynamic_scale:\n grad_fn = dynamic_scale.value_and_grad(\n loss_fn, has_aux=True, axis_name='batch'\n )\n dynamic_scale, is_fin, aux, grads = grad_fn(state.params)\n # dynamic loss takes care of averaging gradients across replicas\n else:\n grad_fn = jax.value_and_grad(loss_fn, has_aux=True)\n aux, grads = grad_fn(state.params)\n # Re-use same axis_name as in the call to `pmap(...train_step...)` below.\n grads = lax.pmean(grads, axis_name='batch')\n new_model_state, logits = aux[1]\n metrics = compute_metrics(logits, batch['label'])\n metrics['learning_rate'] = lr\n\n new_state = state.apply_gradients(\n grads=grads,\n batch_stats=lax.pmean(new_model_state['batch_stats'], 'batch'),\n )\n if dynamic_scale:\n # if is_fin == False the gradients contain Inf/NaNs and optimizer state and\n # params should be restored (= skip this step).\n new_state = new_state.replace(\n opt_state=jax.tree_util.tree_map(\n functools.partial(jnp.where, is_fin),\n new_state.opt_state,\n state.opt_state,\n ),\n params=jax.tree_util.tree_map(\n functools.partial(jnp.where, is_fin), new_state.params, state.params\n ),\n dynamic_scale=dynamic_scale,\n )\n metrics['scale'] = dynamic_scale.scale\n\n return new_state, metrics\n\n\ndef eval_step(state, batch):\n variables = {'params': state.params, 'batch_stats': state.batch_stats}\n logits = state.apply_fn(variables, batch['image'], train=False, mutable=False)\n return compute_metrics(logits, batch['label'])\n\n\ndef prepare_tf_data(xs):\n \"\"\"Convert a input batch from tf Tensors to numpy arrays.\"\"\"\n local_device_count = jax.local_device_count()\n\n def _prepare(x):\n # Use _numpy() for zero-copy conversion between TF and NumPy.\n x = x._numpy() # pylint: disable=protected-access\n\n # reshape (host_batch_size, height, width, 3) to\n # (local_devices, device_batch_size, height, width, 3)\n return x.reshape((local_device_count, -1) + x.shape[1:])\n\n return jax.tree_util.tree_map(_prepare, xs)\n\n\ndef create_input_iter(\n dataset_builder,\n batch_size,\n image_size,\n dtype,\n train,\n cache,\n shuffle_buffer_size,\n prefetch,\n):\n ds = input_pipeline.create_split(\n dataset_builder,\n batch_size,\n image_size=image_size,\n dtype=dtype,\n train=train,\n cache=cache,\n shuffle_buffer_size=shuffle_buffer_size,\n prefetch=prefetch,\n )\n it = map(prepare_tf_data, ds)\n it = jax_utils.prefetch_to_device(it, 2)\n return it\n\n\nclass TrainState(train_state.TrainState):\n batch_stats: Any\n dynamic_scale: dynamic_scale_lib.DynamicScale\n\n\ndef restore_checkpoint(state, workdir):\n return checkpoints.restore_checkpoint(workdir, state)\n\n\ndef save_checkpoint(state, workdir):\n step = int(state.step)\n logging.info('Saving checkpoint step %d.', step)\n\n # Orbax can not handle host local arrays from pmap. Convert to global arrays.\n replicated_state = jax.tree_util.tree_map(\n ocp.utils.fully_replicated_host_local_array_to_global_array,\n state,\n )\n checkpoints.save_checkpoint_multiprocess(\n workdir, replicated_state, step, keep=3\n )\n\n\ndef create_train_state(\n rng, config: ml_collections.ConfigDict, model, image_size, learning_rate_fn\n):\n \"\"\"Create initial training state.\"\"\"\n dynamic_scale = None\n platform = jax.local_devices()[0].platform\n if config.half_precision and platform == 'gpu':\n dynamic_scale = dynamic_scale_lib.DynamicScale()\n\n params, batch_stats = initialized(rng, image_size, model)\n tx = optax.sgd(\n learning_rate=learning_rate_fn,\n momentum=config.momentum,\n nesterov=True,\n )\n state = TrainState.create(\n apply_fn=model.apply,\n params=params,\n tx=tx,\n batch_stats=batch_stats,\n dynamic_scale=dynamic_scale,\n )\n return state\n\n\ndef train_and_evaluate(\n config: ml_collections.ConfigDict, workdir: str\n) -> TrainState:\n \"\"\"Execute model training and evaluation loop.\n\n Args:\n config: Hyperparameter configuration for training and evaluation.\n workdir: Directory where the tensorboard summaries are written to.\n\n Returns:\n Final TrainState.\n \"\"\"\n\n writer = metric_writers.create_default_writer(\n logdir=workdir, just_logging=jax.process_index() != 0\n )\n\n rng = random.key(0)\n\n image_size = 224\n\n if config.batch_size % jax.device_count() > 0:\n raise ValueError('Batch size must be divisible by the number of devices')\n local_batch_size = config.batch_size // jax.process_count()\n\n platform = jax.local_devices()[0].platform\n\n if config.half_precision:\n if platform == 'tpu':\n input_dtype = tf.bfloat16\n else:\n input_dtype = tf.float16\n else:\n input_dtype = tf.float32\n\n dataset_builder = tfds.builder(config.dataset)\n train_iter = create_input_iter(\n dataset_builder,\n local_batch_size,\n image_size,\n input_dtype,\n train=True,\n cache=config.cache,\n shuffle_buffer_size=config.shuffle_buffer_size,\n prefetch=config.prefetch,\n )\n eval_iter = create_input_iter(\n dataset_builder,\n local_batch_size,\n image_size,\n input_dtype,\n train=False,\n cache=config.cache,\n shuffle_buffer_size=None,\n prefetch=config.prefetch,\n )\n\n steps_per_epoch = (\n dataset_builder.info.splits['train'].num_examples // config.batch_size\n )\n\n if config.num_train_steps <= 0:\n num_steps = int(steps_per_epoch * config.num_epochs)\n else:\n num_steps = config.num_train_steps\n\n if config.steps_per_eval == -1:\n num_validation_examples = dataset_builder.info.splits[\n 'validation'\n ].num_examples\n steps_per_eval = num_validation_examples // config.batch_size\n else:\n steps_per_eval = config.steps_per_eval\n\n steps_per_checkpoint = steps_per_epoch * 10\n\n base_learning_rate = config.learning_rate * config.batch_size / 256.0\n\n model_cls = getattr(models, config.model)\n model = create_model(\n model_cls=model_cls, half_precision=config.half_precision\n )\n\n learning_rate_fn = create_learning_rate_fn(\n config, base_learning_rate, steps_per_epoch\n )\n\n state = create_train_state(rng, config, model, image_size, learning_rate_fn)\n state = restore_checkpoint(state, workdir)\n # step_offset > 0 if restarting from checkpoint\n step_offset = int(state.step)\n\n p_train_step = jax.pmap(\n functools.partial(train_step, learning_rate_fn=learning_rate_fn),\n in_axes=(None, 0),\n out_axes=(None, 0),\n axis_name='batch',\n )\n p_eval_step = jax.pmap(eval_step, in_axes=(None, 0), axis_name='batch')\n\n train_metrics = []\n hooks = []\n if jax.process_index() == 0 and config.profile:\n hooks += [\n periodic_actions.Profile(\n num_profile_steps=3, profile_duration_ms=None, logdir=workdir\n )\n ]\n train_metrics_last_t = time.time()\n logging.info('Initial compilation, this might take some minutes...')\n for step, batch in zip(range(step_offset, num_steps), train_iter):\n state, metrics = p_train_step(state, batch)\n for h in hooks:\n h(step)\n if step == step_offset:\n logging.info('Initial compilation completed.')\n\n if config.get('log_every_steps'):\n train_metrics.append(metrics)\n if (step + 1) % config.log_every_steps == 0:\n train_metrics = common_utils.get_metrics(train_metrics)\n summary = {\n f'train_{k}': v\n for k, v in jax.tree_util.tree_map(\n lambda x: x.mean(), train_metrics\n ).items()\n }\n summary['steps_per_second'] = config.log_every_steps / (\n time.time() - train_metrics_last_t\n )\n writer.write_scalars(step + 1, summary)\n train_metrics = []\n train_metrics_last_t = time.time()\n\n if (step + 1) % steps_per_epoch == 0:\n epoch = step // steps_per_epoch\n eval_metrics = []\n\n for _ in range(steps_per_eval):\n eval_batch = next(eval_iter)\n metrics = p_eval_step(state, eval_batch)\n eval_metrics.append(metrics)\n eval_metrics = common_utils.get_metrics(eval_metrics)\n summary = jax.tree_util.tree_map(lambda x: x.mean(), eval_metrics)\n logging.info(\n 'eval epoch: %d, loss: %.4f, accuracy: %.2f',\n epoch,\n summary['loss'],\n summary['accuracy'] * 100,\n )\n writer.write_scalars(\n step + 1, {f'eval_{key}': val for key, val in summary.items()}\n )\n writer.flush()\n if (step + 1) % steps_per_checkpoint == 0 or step + 1 == num_steps:\n save_checkpoint(state, workdir)\n\n # Wait until computations are done before exiting\n jax.random.normal(jax.random.key(0), ()).block_until_ready()\n\n return state\n" }, { "path": "examples/imagenet/train_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for flax.examples.imagenet.train.\"\"\"\n\nimport pathlib\nimport tempfile\n\nfrom absl.testing import absltest\nfrom absl.testing import parameterized\n\nimport jax\nfrom jax import random\nimport tensorflow as tf\nimport tensorflow_datasets as tfds\n\n# Local imports.\nimport models\nimport train\nfrom configs import default as default_lib\n\n\njax.config.update('jax_disable_most_optimizations', True)\n\n\nclass TrainTest(parameterized.TestCase):\n\n def setUp(self):\n super().setUp()\n # Make sure tf does not allocate gpu memory.\n tf.config.experimental.set_visible_devices([], 'GPU')\n\n def test_create_model(self):\n \"\"\"Tests creating model.\"\"\"\n model = train.create_model(model_cls=models._ResNet1, half_precision=False) # pylint: disable=protected-access\n params, batch_stats = train.initialized(random.key(0), 224, model)\n variables = {'params': params, 'batch_stats': batch_stats}\n x = random.normal(random.key(1), (8, 224, 224, 3))\n y = model.apply(variables, x, train=False)\n self.assertEqual(y.shape, (8, 1000))\n\n def test_create_model_local(self):\n \"\"\"Tests creating an unshared convolution model.\n\n Uses smaller inputs than `test_create_model` to due to higher compute.\n \"\"\"\n model = train.create_model(\n model_cls=models._ResNet1Local, half_precision=False\n ) # pylint: disable=protected-access\n params, batch_stats = train.initialized(random.key(0), 64, model)\n variables = {'params': params, 'batch_stats': batch_stats}\n x = random.normal(random.key(1), (1, 64, 64, 3))\n y = model.apply(variables, x, train=False)\n self.assertEqual(y.shape, (1, 1000))\n\n @parameterized.product(model=('_ResNet1', '_ResNet1Local'))\n def test_train_and_evaluate(self, model):\n \"\"\"Tests training and evaluation loop using mocked data.\"\"\"\n # Create a temporary directory where tensorboard metrics are written.\n workdir = tempfile.mkdtemp()\n\n # Go two directories up to the root of the flax directory.\n flax_root_dir = pathlib.Path(__file__).parents[2]\n data_dir = str(flax_root_dir) + '/.tfds/metadata'\n\n # Define training configuration\n config = default_lib.get_config()\n config.model = model\n config.batch_size = 1\n config.num_epochs = 1\n config.num_train_steps = 1\n config.steps_per_eval = 1\n\n with tfds.testing.mock_data(num_examples=1, data_dir=data_dir):\n train.train_and_evaluate(workdir=workdir, config=config)\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "examples/linen_design_test/attention_simple.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport functools\nfrom pprint import pprint\nfrom typing import Any, Optional\nfrom collections.abc import Callable, Sequence\nfrom flax.core.frozen_dict import unfreeze\nfrom flax.linen import initializers\nfrom flax.linen import Module, compact, vmap\nfrom flax.linen.linear import PrecisionLike\nimport jax\nfrom jax import lax, numpy as jnp, random\n\n\nclass Dense(Module):\n features: int\n use_bias: bool = True\n kernel_init: Callable = initializers.lecun_normal()\n bias_init: Callable = initializers.zeros_init()\n dtype: Any = jnp.float32\n precision: PrecisionLike = None\n\n @compact\n def __call__(self, inputs):\n inputs = jnp.asarray(inputs, self.dtype)\n kernel = self.param(\n 'kernel', self.kernel_init, (inputs.shape[-1], self.features)\n )\n kernel = jnp.asarray(kernel, self.dtype)\n y = lax.dot_general(\n inputs,\n kernel,\n (((inputs.ndim - 1,), (0,)), ((), ())),\n precision=self.precision,\n )\n if self.use_bias:\n bias = self.param('bias', self.bias_init, (self.features,))\n bias = jnp.asarray(bias, self.dtype)\n y = y + bias\n return y\n\n\nclass SoftmaxAttn(Module):\n\n @compact\n def __call__(self, weights):\n norm_dims = tuple(range(weights.ndim // 2, weights.ndim))\n return jax.nn.softmax(weights, axis=norm_dims)\n\n\nclass Dropout(Module):\n rate: float\n\n @compact\n def __call__(self, x, deterministic=False, rng=None):\n if self.rate == 0.0:\n return x\n keep_prob = 1.0 - self.rate\n\n if deterministic:\n return x\n else:\n if rng is None:\n rng = self.scope.make_rng('dropout')\n mask = random.bernoulli(rng, p=keep_prob, shape=x.shape)\n return lax.select(mask, x / keep_prob, jnp.zeros_like(x))\n\n\nclass SoftmaxAttnWDropout(Module):\n rate: float = 0.0\n deterministic: bool = False\n\n @compact\n def __call__(self, x):\n x = SoftmaxAttn()(x)\n x = Dropout(self.rate)(x, deterministic=self.deterministic)\n return x\n\n\nclass RawDotProductAttention(Module):\n attn_module: Callable = SoftmaxAttn\n\n @compact\n def __call__(self, query, key, value, bias=None, dtype=jnp.float32):\n assert key.ndim == query.ndim\n assert key.ndim == value.ndim\n\n n = query.ndim\n attn_weights = lax.dot_general(query, key, (((n - 1,), (n - 1,)), ((), ())))\n if bias is not None:\n attn_weights += bias\n attn_weights = self.attn_module()(attn_weights)\n attn_weights = attn_weights.astype(dtype)\n\n contract_dims = (\n tuple(range(n - 1, attn_weights.ndim)),\n tuple(range(0, n - 1)),\n )\n y = lax.dot_general(attn_weights, value, (contract_dims, ((), ())))\n return y\n\n\nclass DotProductAttention(Module):\n qkv_features: int | None = None\n out_features: int | None = None\n attn_module: Callable = SoftmaxAttn\n\n @compact\n def __call__(self, inputs_q, inputs_kv, bias=None, dtype=jnp.float32):\n qkv_features = self.qkv_features or inputs_q.shape[-1]\n out_features = self.out_features or inputs_q.shape[-1]\n\n QKVDense = functools.partial(\n Dense, features=qkv_features, use_bias=False, dtype=dtype\n )\n query = QKVDense(name='query')(inputs_q)\n key = QKVDense(name='key')(inputs_kv)\n value = QKVDense(name='value')(inputs_kv)\n\n y = RawDotProductAttention(attn_module=self.attn_module)(\n query, key, value, bias=bias, dtype=dtype\n )\n y = Dense(features=out_features, dtype=dtype, name='out')(y)\n return y\n\n\n# Trying out a slightly more compact vmap notation:\n\n\ndef concise_vmap(module, in_axes, out_axes, axis_size=None, **var_specs):\n variable_axes = {\n k: v[0] for k, v in var_specs.items() if isinstance(v, Sequence)\n }\n splits = {k: v[1] for k, v in var_specs.items() if isinstance(v, Sequence)}\n return vmap(\n module,\n in_axes=in_axes,\n out_axes=out_axes,\n variable_axes=variable_axes,\n split_rngs=splits,\n axis_size=axis_size,\n )\n\n\nclass MultiHeadDotProductAttention(Module):\n qkv_features: int | None = None\n out_features: int | None = None\n attn_module: Callable = SoftmaxAttn\n batch_axes: Sequence[int] = (0,)\n num_heads: int = 1\n broadcast_dropout: bool = False\n\n @compact\n def __call__(self, inputs_q, inputs_kv, bias=None, dtype=jnp.float32):\n qkv_features = self.qkv_features or inputs_q.shape[-1]\n out_features = self.out_features or inputs_q.shape[-1]\n\n # Now, vmap attn.__call__ along heads and spatial dims.\n Attn = concise_vmap(\n DotProductAttention,\n (None, None, None),\n -2,\n param=(0, True),\n dropout=(None, not self.broadcast_dropout),\n axis_size=self.num_heads,\n )\n for axis in reversed(sorted(self.batch_axes)):\n Attn = concise_vmap(\n Attn,\n (axis, axis, axis),\n axis,\n param=(None, False),\n dropout=(None, not self.broadcast_dropout),\n )\n\n attn = Attn(\n attn_module=self.attn_module,\n qkv_features=qkv_features // self.num_heads,\n out_features=out_features,\n )\n\n # evaluate multi-headed-attention.\n y = attn(inputs_q, inputs_kv, bias)\n return y.mean(axis=-2)\n\n\n# run it.\n\n\nif __name__ == '__main__':\n inputs = jnp.ones((8, 97, 256))\n rngs = {'params': random.key(0), 'dropout': random.key(1)}\n model = MultiHeadDotProductAttention(\n broadcast_dropout=False,\n qkv_features=256,\n out_features=256,\n attn_module=functools.partial(SoftmaxAttnWDropout, rate=0.1),\n num_heads=8,\n batch_axes=(0,),\n )\n\n y, params = model.init_with_output(rngs, inputs, inputs)\n\n print('input shape: ', inputs.shape)\n print('parameter shapes:')\n pprint(jax.tree_util.tree_map(jnp.shape, unfreeze(params)))\n print('output shape: ', y.shape)\n" }, { "path": "examples/linen_design_test/autoencoder.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom typing import Tuple\nfrom collections.abc import Iterable\n\nimport jax\nfrom jax import numpy as jnp, random\n\nfrom flax import linen as nn\nfrom flax.linen import Module, Dense, compact\n\n\n# A concise MLP defined via lazy submodule initialization\nclass MLP(Module):\n widths: Iterable\n\n @compact\n def __call__(self, x):\n for width in self.widths[:-1]:\n x = nn.relu(Dense(width)(x))\n return Dense(self.widths[-1])(x)\n\n\n# An autoencoder exposes multiple methods, so we define all\n# submodules in setup().\nclass AutoEncoder(Module):\n encoder_widths: Iterable\n decoder_widths: Iterable\n input_shape: tuple = None\n\n def setup(self):\n # Submodules attached in `setup` get names via attribute assignment\n self.encoder = MLP(self.encoder_widths)\n self.decoder = MLP(self.decoder_widths + (jnp.prod(self.input_shape),))\n\n def __call__(self, x):\n return self.decode(self.encode(x))\n\n def encode(self, x):\n assert x.shape[-len(self.input_shape) :] == self.input_shape\n return self.encoder(jnp.reshape(x, (x.shape[0], -1)))\n\n def decode(self, z):\n z = self.decoder(z)\n x = nn.sigmoid(z)\n x = jnp.reshape(x, (x.shape[0],) + self.input_shape)\n return x\n\n\n# `ae` is a detached module, which has no variables.\nae = AutoEncoder(\n encoder_widths=(32, 32, 32),\n decoder_widths=(32, 32, 32),\n input_shape=(28, 28, 1),\n)\n\n\n# `ae.initialized` returns a materialized copy of `ae` by\n# running through an input to create submodules defined lazily.\nparams = ae.init({\"params\": random.key(42)}, jnp.ones((1, 28, 28, 1)))\n\n\n# Now you can use `ae` as a normal object, calling any methods defined on AutoEncoder\nprint(\"reconstruct\", jnp.shape(ae.apply(params, jnp.ones((1, 28, 28, 1)))))\nprint(\n \"encoder\",\n jnp.shape(ae.apply(params, jnp.ones((1, 28, 28, 1)), method=ae.encode)),\n)\n\n\n# `ae.variables` is a frozen dict that looks like\n# {'params': {\"decoder\": {\"Dense_0\": {\"bias\": ..., \"kernel\": ...}, ...}}\nprint(\"var shapes\", jax.tree_util.tree_map(jnp.shape, params))\n\n\n# TODO(avital, levskaya): resurrect this example once interactive api is restored.\n\n\n# You can access submodules defined in setup(), they are just references on\n# the autoencoder instance\n# encoder = ae.encoder\n# print(\"encoder var shapes\", jax.tree_util.tree_map(jnp.shape, encoder.variables))\n\n\n# # You can also access submodules that were defined in-line.\n# # (We may add syntactic sugar here, e.g. to allow `ae.encoder.Dense_0`)\n# encoder_dense0 = ae.encoder.children['Dense_0']\n# print(\"encoder dense0 var shapes\", jax.tree_util.tree_map(jnp.shape, encoder_dense0.variables))\n" }, { "path": "examples/linen_design_test/dense.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom jax import lax\nfrom flax.linen import initializers\nfrom collections.abc import Callable\nfrom flax.linen import Module, compact\n\n\nclass Dense(Module):\n features: int\n kernel_init: Callable = initializers.lecun_normal()\n bias_init: Callable = initializers.zeros_init()\n use_bias: bool = True\n\n @compact\n def __call__(self, inputs):\n kernel = self.param(\n 'kernel', self.kernel_init, (inputs.shape[-1], self.features)\n )\n y = lax.dot_general(\n inputs,\n kernel,\n (((inputs.ndim - 1,), (0,)), ((), ())),\n )\n if self.use_bias:\n bias = self.param('bias', self.bias_init, (self.features,))\n y = y + bias\n return y\n" }, { "path": "examples/linen_design_test/linear_regression.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport jax\nfrom jax import numpy as jnp, jit\nfrom dense import Dense\n\n\nX = jnp.ones((1, 10))\nY = jnp.ones((5,))\n\nmodel = Dense(features=5)\n\n\n@jit\ndef predict(params):\n return model.apply({\"params\": params}, X)\n\n\n@jit\ndef loss_fn(params):\n return jnp.mean(jnp.abs(Y - predict(params)))\n\n\n@jit\ndef init_params(rng):\n mlp_variables = model.init({\"params\": rng}, X)\n return mlp_variables[\"params\"]\n\n\n# Get initial parameters\nparams = init_params(jax.random.key(42))\nprint(\"initial params\", params)\n\n# Run SGD.\nfor i in range(50):\n loss, grad = jax.value_and_grad(loss_fn)(params)\n print(i, \"loss = \", loss, \"Yhat = \", predict(params))\n lr = 0.03\n params = jax.tree_util.tree_map(lambda x, d: x - lr * d, params, grad)\n" }, { "path": "examples/linen_design_test/mlp_explicit.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom pprint import pprint\nfrom typing import Optional\nfrom flax.deprecated import nn\nfrom dense import Dense\nfrom flax.linen import Module\nimport jax\nfrom jax import numpy as jnp\n\n\n# Add `in_features` to the built-in Dense layer that normally works\n# via shape inference.\nclass DenseExplicit(Dense):\n in_features: int | None = None\n\n def setup(self):\n # We feed a fake batch through the module, which initialized parameters.\n # Assuming we're in a jit, should use no FLOPs -- \"just shape inference\".\n self.__call__(\n jnp.zeros((\n 1,\n self.in_features,\n ))\n )\n\n\nclass MLP(Module):\n\n def setup(self):\n self.dense1 = DenseExplicit(in_features=3, features=2)\n self.dense2 = DenseExplicit(in_features=2, features=1)\n\n # explicit instances are materialized immediately at init\n pprint(self.dense2.variables)\n # {'params': {'bias': DeviceArray([0.], dtype=float32),\n # 'kernel': DeviceArray([[ 0.6704609 ],\n # [-0.90477365]], dtype=float32)}}\n\n def __call__(self, x):\n return self.dense2(nn.relu(self.dense1(x)))\n\n\n# Return an initialized instance of MLP by only calling `setup`.\nrngkey = jax.random.key(10)\ninit_variables = MLP().init({'params': rngkey}, jnp.ones((1, 3)))\n\npprint(init_variables)\n# {'params': {'dense1': {'bias': DeviceArray([0., 0.], dtype=float32),\n# 'kernel': DeviceArray([[ 0.18307537, -0.38739476],\n# [-0.902451 , -0.5190721 ],\n# [ 0.51552075, 1.1169153 ]], dtype=float32)},\n# 'dense2': {'bias': DeviceArray([0.], dtype=float32),\n# 'kernel': DeviceArray([[ 0.6704609 ],\n# [-0.90477365]], dtype=float32)}}}\n" }, { "path": "examples/linen_design_test/mlp_inline.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport jax\nfrom jax import numpy as jnp\nfrom flax import linen as nn\nfrom collections.abc import Iterable\nfrom flax.linen import Module, compact\nfrom dense import Dense\n\n\n# Many NN layers and blocks are best described by a single function with inline variables.\n# In this case, variables are initialized during the first call.\nclass MLP(Module):\n sizes: Iterable[int]\n\n @compact\n def __call__(self, x):\n for size in self.sizes[:-1]:\n x = Dense(size)(x)\n x = nn.relu(x)\n return Dense(self.sizes[-1])(x)\n\n\n# Return an initialized instance of MLP by calling `__call__` with an input batch,\n# initializing all variables.\n#\n# Variable shapes depend on the input shape passed in.\nrngkey = jax.random.key(10)\nmodel = MLP((2, 1))\nx = jnp.ones((1, 3))\nmlp_variables = model.init(rngkey, x)\nprint(mlp_variables)\n# {'params': {'Dense_0': {'bias': DeviceArray([0.], dtype=float32),\n# 'kernel': DeviceArray([[-0.04267037],\n# [-0.51097125]], dtype=float32)},\n# 'Dense_1': {'bias': DeviceArray([0., 0.], dtype=float32),\n# 'kernel': DeviceArray([[-6.3845289e-01, 6.0373604e-01],\n# [-5.9814966e-01, 5.1718324e-01],\n# [-6.2220657e-01, 5.8988278e-04]], dtype=float32)}}}\nprint(model.apply(mlp_variables, x))\n" }, { "path": "examples/linen_design_test/mlp_lazy.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport jax\nfrom jax import numpy as jnp\nfrom flax import linen as nn\nfrom flax.linen import Module\nfrom pprint import pprint\nfrom dense import Dense\n\n\n# Here submodules are explicitly defined during init, but still materialized\n# lazily only once a first input is passed through and shapes are known.\nclass MLP(Module):\n\n def setup(self):\n self.dense1 = Dense(features=2)\n self.dense2 = Dense(features=1)\n\n # shapes aren't yet known, so variables aren't materialized\n print(self.dense2.variables)\n # FrozenDict({})\n\n def __call__(self, x):\n return self.dense2(nn.relu(self.dense1(x)))\n\n\n# Return an initialized instance of MLP by calling `__call__` with an input batch,\n# initializing all variables.\n#\n# Variable shapes depend on the input shape passed in.\nrngkey = jax.random.key(10)\nmlp_variables = MLP().init(rngkey, jnp.zeros((1, 3)))\n\npprint(mlp_variables)\n# {'params': {'dense1': {'bias': DeviceArray([0., 0.], dtype=float32),\n# 'kernel': DeviceArray([[ 0.18307537, -0.38739476],\n# [-0.902451 , -0.5190721 ],\n# [ 0.51552075, 1.1169153 ]], dtype=float32)},\n# 'dense2': {'bias': DeviceArray([0.], dtype=float32),\n# 'kernel': DeviceArray([[ 0.6704609 ],\n# [-0.90477365]], dtype=float32)}}}\n" }, { "path": "examples/lm1b/README.md", "content": "\n## Language modeling\nTrains a Transformer-based model (Vaswani *et al.*, 2017) on the One Billion\nWord Benchmark (lm1b; Chelba *et al.*, 2013).\n\nThis example uses linear learning rate warmup and inverse square root learning\nrate schedule. Based off of Machine Translation `wmt` example.\n\n\n### Requirements\n\n* TensorFlow datasets `lm1b` need to be downloaded and prepared.\n A sentencepiece tokenizer vocabulary will be automatically generated\n and saved on each training run.\n* This example additionally depends on the `sentencepiece` and\n `tensorflow-text` packages.\n\n\n### How to run on Cloud TPUs\n\nSetup the TPU VM and install the Flax dependencies on it as described\n[here](https://cloud.google.com/tpu/docs/jax-pods) for creating pod slices, or\n[here](https://cloud.google.com/tpu/docs/jax-quickstart-tpu-vm) for a single\nv3-8 TPU.\n\n\nFirst create a single TPUv3-8 VM and connect to it (you can find more detailed\ninstructions [here](https://cloud.google.com/tpu/docs/jax-quickstart-tpu-vm)):\n\n```\nZONE=us-central1-a\nTPU_TYPE=v3-8\nTPU_NAME=$USER-flax-lm1b\ngcloud alpha compute tpus tpu-vm create $TPU_NAME \\\n --zone $ZONE \\\n --accelerator-type $TPU_TYPE \\\n --version v2-alpha\ngcloud alpha compute tpus tpu-vm ssh $TPU_NAME --zone $ZONE -- \\\n -L 6006:localhost:6006\n```\n\nWhen connected install JAX:\n\n```\npip install \"jax[tpu]>=0.2.16\" \\\n -f https://storage.googleapis.com/jax-releases/libtpu_releases.html\n```\n\nThen install Flax + the example dependencies:\n\n```\ngit clone --depth=1 --branch=main https://github.com/google/flax\ncd flax\npip install -e .\ncd examples/lm1b\npip install -r requirements.txt\n```\n\nAnd finally start the training:\n\n```\npython3 main.py --workdir=$HOME/logs/lm1b_256 \\\n --config.per_device_batch_size=32 \\\n --jax_backend_target=\"grpc://192.168.0.2:8470\"\n```\n\nNote that you might want to set `TFDS_DATA_DIR` as explained below. You probably\nalso want to start the long-running command above in a `tmux` session and start\nsome monitoring in a separate pane (note that we forwarded port 6006 locally\nabove):\n\n```\ntensorboard --logdir=$HOME/logs\n```\n\nYou should expect to get numbers similar to these:\n\n\nHardware | config | Training time | Loss | TensorBoard.dev | Workdir\n-------- | ------- | ------------- | -------------- | ------------------------------------------------------------------------ | --------------------------------------------------------------------------------------------------------------------------\nTPU v3-8 | default | 13h18m | 3.127 | [2021-08-08](https://tensorboard.dev/experiment/n30WkNOZTJq3RHWD7wNslg/) | [gs://flax_public/examples/lm1b/default](https://console.cloud.google.com/storage/browser/flax_public/examples/lm1b/default)\n\n### Downloading the LM1B Datasets\n\nWe recommend downloading and preparing the TFDS datasets beforehand. For Cloud\nTPUs, we recommend using a cheap standard instance and saving the prepared TFDS\ndata on a storage bucket, from where it can be loaded directly. Set the\n`TFDS_DATA_DIR` to your storage bucket path (`gs://`).\n\nYou can download and prepare LM1B datasets using TFDS directly:\n`python -m tensorflow_datasets.scripts.download_and_prepare\n--datasets=lm1b`\n\n" }, { "path": "examples/lm1b/configs/default.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Default Hyperparameter configuration.\"\"\"\n\nimport ml_collections\n\n\ndef get_config():\n \"\"\"Get the default hyperparameter configuration.\"\"\"\n config = ml_collections.ConfigDict()\n\n # Path to load or store sentencepiece vocab file.\n config.vocab_path = None\n\n # Vocabulary size if `vocab_path` is not given.\n config.vocab_size = 30_000\n\n config.max_corpus_chars = 10**7\n\n # Name of TFDS translation dataset to use.\n config.dataset_name = 'lm1b'\n\n # Optional name of TFDS translation dataset to use for evaluation.\n config.eval_dataset_name = 'lm1b'\n config.eval_split = 'test'\n\n # Per device batch size for training.\n config.per_device_batch_size = 32\n\n # Per device batch size for training.\n config.eval_per_device_batch_size = 32\n\n # Sampling temperature for language model inference.\n config.sampling_temperature = 0.6\n\n # Top k cutoff for logit sampling. If 0 then no top-k cutoff is used.\n config.sampling_top_k = 20\n\n config.num_train_steps = 500_000\n\n # Number of steps to take during evaluation. Large enough to evaluate all.\n # Large enough to evaluate all samples: 306_688 / (32 * 8) = 1198\n config.num_eval_steps = 2_000\n # Number of steps to generate predictions.\n # -1 will use the whole eval dataset.\n config.num_predict_steps = -1\n\n # Base learning rate.\n config.learning_rate = 0.0016\n\n # Linear learning rate warmup.\n config.warmup_steps = 1000\n\n # Cross entropy loss label smoothing.\n config.label_smoothing = 0.0\n\n # Decay factor for AdamW style weight decay.\n config.weight_decay = 0.1\n\n # Maximum length cutoff for training examples.\n config.max_target_length = 128\n # Maximum length cutoff for eval examples.\n config.max_eval_target_length = 512\n # Maximum length cutoff for predicted tokens.\n config.max_predict_length = 50\n\n # Final logit transform uses embedding matrix transpose.\n config.logits_via_embedding = False\n\n # Number of transformer layers.\n config.num_layers = 6\n\n # Size of query/key/value for attention.\n config.qkv_dim = 512\n # Size of embeddings.\n config.emb_dim = 512\n # Size of the MLP.\n config.mlp_dim = 2048\n\n # Number of attention heads.\n config.num_heads = 8\n\n # Dropout rate.\n config.dropout_rate = 0.1\n\n # Attention dropout rate.\n config.attention_dropout_rate = 0.1\n\n # Whether to save model checkpoints.\n config.save_checkpoints = True\n # Whether to restore from existing model checkpoints.\n config.restore_checkpoints = True\n\n # Save a checkpoint every these number of steps.\n config.checkpoint_every_steps = 10_000\n # Frequency of eval during training, e.g. every 1_000 steps.\n config.eval_every_steps = 1_000\n\n # Use bfloat16 mixed precision training instead of float32.\n config.use_bfloat16 = True\n\n # Integer for PRNG random seed.\n config.seed = 0\n\n # Prompt for language model sampling,\n # taken from MaxText (https://github.com/google/maxtext/blob/main/MaxText/configs/base.yml).\n config.prompts = 'I love to '\n\n # Parallelism\n config.mesh_axes = ['data', 'fsdp', 'tensor']\n config.logical_axis_rules = [\n ['mlp', 'tensor'],\n ['vocab', 'tensor'],\n ['embed', 'fsdp'],\n ['heads', 'tensor'],\n ]\n config.data_sharding = ['data']\n\n # One axis for each parallelism type may hold a placeholder (-1)\n # value to auto-shard based on available slices and devices.\n # By default, product of the DCN axes should equal number of slices\n # and product of the ICI axes should equal number of devices per slice.\n # ICI (Inter-Chip Interconnection): A high-speed connection between\n # sets of TPU chips, which form the TPU network.\n # DCN (Data Center Network): A connection between the TPU networks;\n # not as fast as ICI.\n # ICI has around 100x the bandwidth of DCN, but it is not a general\n # purpose connection, which is why DCN is necessary for scaling to\n # extremely large ML models.\n config.dcn_data_parallelism = -1 # recommended DCN axis to be auto-sharded\n config.dcn_fsdp_parallelism = 1\n config.dcn_tensor_parallelism = 1\n config.ici_data_parallelism = 1\n config.ici_fsdp_parallelism = -1 # recommended ICI axis to be auto-sharded\n config.ici_tensor_parallelism = 1\n\n return config\n" }, { "path": "examples/lm1b/input_pipeline.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Input pipeline for a LM1B dataset.\"\"\"\n\nimport os\nfrom typing import Dict, Optional, List, Union\n\nfrom clu import deterministic_data\nimport ml_collections\nimport tensorflow as tf\nimport tensorflow_datasets as tfds\n\nimport tokenizer\n\nAUTOTUNE = tf.data.experimental.AUTOTUNE\nFeatures = dict[str, tf.Tensor]\n\n\nclass NormalizeFeatureNamesOp:\n \"\"\"Normalizes feature names to 'inputs' and 'targets'.\"\"\"\n\n def __init__(self, ds_info: tfds.core.DatasetInfo):\n self.ds_info = ds_info\n\n def __call__(self, features: Features) -> Features:\n features['inputs'] = features.pop('text')\n # Unnecessary step used for uniformizing with examples/wmt.\n features['targets'] = features['inputs']\n return features\n\n\ndef get_raw_dataset(\n dataset_builder: tfds.core.DatasetBuilder, split: str\n) -> tf.data.Dataset:\n \"\"\"Loads a raw text dataset and normalizes feature keys.\n\n Args:\n dataset_builder: TFDS dataset builder that can build `split`.\n split: Split to use. This must be the full split. We shard the split across\n multiple hosts and currently don't support sharding subsplits.\n\n Returns:\n Dataset with source and target language features mapped to 'inputs' and\n 'targets'.\n \"\"\"\n num_examples = dataset_builder.info.splits[split].num_examples\n per_host_split = deterministic_data.get_read_instruction_for_host(\n split, num_examples, drop_remainder=False\n )\n ds = dataset_builder.as_dataset(split=per_host_split, shuffle_files=False)\n ds = ds.map(\n NormalizeFeatureNamesOp(dataset_builder.info), num_parallel_calls=AUTOTUNE\n )\n return ds\n\n\ndef pack_dataset(\n dataset: tf.data.Dataset,\n key2length: int | dict[str, int],\n keys: list[str] | None = None,\n) -> tf.data.Dataset:\n \"\"\"Creates a 'packed' version of a dataset on-the-fly.\n\n Adapted from the mesh-tf implementation.\n\n This is meant to replace the irritation of having to create a separate\n \"packed\" version of a dataset to train efficiently on TPU.\n Each example in the output dataset represents several examples in the\n input dataset.\n For each key in the input dataset, two additional keys are created:\n _segmentation: an int32 tensor identifying the parts\n representing the original example.\n _position: an int32 tensor identifying the position within the original\n example.\n Example:\n Two input examples get combined to form an output example.\n The input examples are:\n {\"inputs\": [8, 7, 1, 0], \"targets\":[4, 1, 0]}\n {\"inputs\": [2, 3, 4, 1], \"targets\":[5, 6, 1]}\n The output example is:\n {\n \"inputs\": [8, 7, 1, 2, 3, 4, 1, 0, 0, 0]\n \"inputs_segmentation\": [1, 1, 1, 2, 2, 2, 2, 0, 0, 0]\n \"inputs_position\": [0, 1, 2, 0, 1, 2, 3, 0, 0, 0]\n \"targets\": [4, 1, 5, 6, 1, 0, 0, 0, 0, 0]\n \"targets_segmentation\": [1, 1, 2, 2, 2, 0, 0, 0, 0, 0]\n \"targets_position\": [0, 1, 0, 1, 2, 0, 0, 0, 0, 0]\n }\n 0 represents padding in both the inputs and the outputs.\n Sequences in the incoming examples are truncated to length \"length\", and the\n sequences in the output examples all have fixed (padded) length \"length\".\n\n Args:\n dataset: a tf.data.Dataset\n key2length: an integer, or a dict from feature-key to integer\n keys: a list of strings (e.g. [\"inputs\", \"targets\"])\n\n Returns:\n a tf.data.Dataset\n \"\"\"\n shapes = tf.nest.map_structure(lambda spec: spec.shape, dataset.element_spec)\n if keys is None:\n keys = list(shapes.keys())\n for k in keys:\n if k not in shapes:\n raise ValueError(\n 'Key %s not found in dataset. Available keys are %s'\n % (k, shapes.keys())\n )\n if not shapes[k].is_compatible_with(tf.TensorShape([None])): # type: ignore[wrong-arg-types]\n raise ValueError('Tensors to be packed must be one-dimensional.')\n # make sure that the length dictionary contains all keys as well as the\n # keys suffixed by \"_segmentation\" and \"_position\"\n if isinstance(key2length, int):\n key2length = {k: key2length for k in keys}\n for k in keys:\n for suffix in ['_segmentation', '_position']:\n key2length[k + suffix] = key2length[k]\n\n # trim to length\n dataset = dataset.map(\n lambda x: {k: x[k][: key2length[k]] for k in keys},\n num_parallel_calls=AUTOTUNE,\n )\n # Setting batch_size=length ensures that the concatenated sequences (if they\n # have length >=1) are sufficient to fill at least one packed example.\n batch_size = max(key2length.values())\n dataset = dataset.padded_batch(\n batch_size, padded_shapes={k: [-1] for k in keys}\n )\n dataset = _pack_with_tf_ops(dataset, keys, key2length)\n\n # Set the Tensor shapes correctly since they get lost in the process.\n def my_fn(x):\n return {k: tf.reshape(v, [key2length[k]]) for k, v in x.items()}\n\n return dataset.map(my_fn, num_parallel_calls=AUTOTUNE)\n\n\ndef _pack_with_tf_ops(\n dataset: tf.data.Dataset, keys: list[str], key2length: dict[str, int]\n) -> tf.data.Dataset:\n \"\"\"Helper-function for packing a dataset which has already been batched.\n\n Helper for pack_dataset() Uses tf.while_loop.\n\n Args:\n dataset: a dataset containing padded batches of examples.\n keys: a list of strings\n key2length: a dict from feature-key to integer\n\n Returns:\n a dataset.\n \"\"\"\n empty_example = {}\n for k in keys:\n empty_example[k] = tf.zeros([0], dtype=tf.int32)\n empty_example[k + '_position'] = tf.zeros([0], dtype=tf.int32)\n keys_etc = empty_example.keys()\n\n def write_packed_example(partial, outputs):\n new_partial = empty_example.copy()\n new_outputs = {}\n for k in keys_etc:\n new_outputs[k] = outputs[k].write(\n outputs[k].size(),\n tf.pad(partial[k], [[0, key2length[k] - tf.size(partial[k])]]),\n )\n return new_partial, new_outputs\n\n def map_fn(x):\n \"\"\"Internal function to flat_map over.\n\n Consumes a batch of input examples and produces a variable number of output\n examples.\n Args:\n x: a single example\n\n Returns:\n a tf.data.Dataset\n \"\"\"\n partial = empty_example.copy()\n i = tf.zeros([], dtype=tf.int32)\n dynamic_batch_size = tf.shape(x[keys[0]])[0]\n outputs = {}\n for k in keys:\n outputs[k] = tf.TensorArray(\n tf.int32, size=0, dynamic_size=True, element_shape=[key2length[k]]\n )\n outputs[k + '_position'] = tf.TensorArray(\n tf.int32, size=0, dynamic_size=True, element_shape=[key2length[k]]\n )\n\n def body_fn(i, partial, outputs):\n \"\"\"Body function for while_loop.\n\n Args:\n i: integer scalar\n partial: dictionary of Tensor (partially-constructed example)\n outputs: dictionary of TensorArray\n\n Returns:\n A triple containing the new values of the inputs.\n \"\"\"\n can_append = True\n one_example = {}\n for k in keys:\n val = tf.cast(x[k][i], tf.int32)\n val = val[: tf.reduce_sum(tf.cast(tf.not_equal(val, 0), tf.int32))]\n one_example[k] = val\n for k in keys:\n can_append = tf.logical_and(\n can_append,\n tf.less_equal(\n tf.size(partial[k]) + tf.size(one_example[k]), key2length[k]\n ),\n )\n\n def false_fn():\n return write_packed_example(partial, outputs)\n\n def true_fn():\n return partial, outputs\n\n partial, outputs = tf.cond(can_append, true_fn, false_fn)\n new_partial = {}\n for k in keys:\n new_seq = one_example[k][: key2length[k]]\n new_seq_len = tf.size(new_seq)\n new_partial[k] = tf.concat([partial[k], new_seq], 0)\n new_partial[k + '_position'] = tf.concat(\n [partial[k + '_position'], tf.range(new_seq_len)], 0\n )\n partial = new_partial\n return i + 1, partial, outputs\n\n # For loop over all examples in the batch.\n i, partial, outputs = tf.while_loop(\n cond=lambda *_: True,\n body=body_fn,\n loop_vars=(i, partial, outputs),\n shape_invariants=(\n tf.TensorShape([]),\n {k: tf.TensorShape([None]) for k in keys_etc}, # type: ignore[wrong-arg-types]\n {k: tf.TensorShape(None) for k in keys_etc}, # type: ignore[wrong-arg-types]\n ),\n maximum_iterations=dynamic_batch_size,\n )\n _, outputs = write_packed_example(partial, outputs)\n packed = {k: outputs[k].stack() for k in keys_etc}\n for k in keys:\n packed[k + '_segmentation'] = tf.cumsum(\n tf.cast(tf.equal(packed[k + '_position'], 0), tf.int32), axis=1\n ) * tf.cast(tf.not_equal(packed[k], 0), tf.int32)\n return packed\n\n dataset = dataset.map(map_fn, num_parallel_calls=AUTOTUNE)\n return dataset.unbatch()\n\n\n# -----------------------------------------------------------------------------\n# Main dataset prep routines.\n# -----------------------------------------------------------------------------\ndef preprocess_data(\n dataset,\n shuffle: bool,\n num_epochs: int | None = 1,\n pack_examples: bool = True,\n shuffle_buffer_size: int = 1024,\n max_length: int = 512,\n batch_size: int = 256,\n drop_remainder: bool = True,\n prefetch_size: int = AUTOTUNE,\n):\n \"\"\"Shuffle and batch/pack the given dataset.\"\"\"\n\n def length_filter(max_len):\n def filter_fn(x):\n source, target = x['inputs'], x['targets']\n l = tf.maximum(tf.shape(source)[0], tf.shape(target)[0])\n return tf.less(l, max_len + 1)\n\n return filter_fn\n\n if max_length > 0:\n dataset = dataset.filter(length_filter(max_length))\n\n if shuffle:\n dataset = dataset.shuffle(shuffle_buffer_size)\n dataset = dataset.repeat(num_epochs)\n\n if pack_examples:\n dataset = pack_dataset(dataset, max_length)\n dataset = dataset.batch(batch_size, drop_remainder=drop_remainder)\n else: # simple (static-shape) padded batching\n dataset = dataset.padded_batch(\n batch_size,\n padded_shapes={'inputs': max_length, 'targets': max_length},\n padding_values={'inputs': 0, 'targets': 0},\n drop_remainder=drop_remainder,\n )\n\n if prefetch_size:\n dataset = dataset.prefetch(prefetch_size)\n\n return dataset\n\n\ndef get_datasets(\n config: ml_collections.ConfigDict,\n *,\n n_devices: int,\n vocab_path: str | None = None,\n):\n \"\"\"Load and return dataset of batched examples for use during training.\"\"\"\n if vocab_path is None:\n vocab_path = os.path.expanduser('~/lm1b_sentencepiece_model')\n\n train_ds_builder = tfds.builder(config.dataset_name)\n train_data = get_raw_dataset(train_ds_builder, 'train')\n\n if config.eval_dataset_name:\n eval_ds_builder = tfds.builder(config.eval_dataset_name)\n else:\n eval_ds_builder = train_ds_builder\n eval_data = get_raw_dataset(eval_ds_builder, config.eval_split)\n\n # Tokenize data.\n sp_tokenizer = tokenizer.load_or_train_tokenizer(\n train_data,\n vocab_path=vocab_path,\n vocab_size=config.vocab_size,\n max_corpus_chars=config.max_corpus_chars,\n )\n train_data = train_data.map(\n tokenizer.TokenizeOp(sp_tokenizer), num_parallel_calls=AUTOTUNE\n )\n eval_data = eval_data.map(\n tokenizer.TokenizeOp(sp_tokenizer), num_parallel_calls=AUTOTUNE\n )\n\n batch_size = config.per_device_batch_size * n_devices\n if config.eval_per_device_batch_size > 0:\n eval_batch_size = config.eval_per_device_batch_size * n_devices\n else:\n eval_batch_size = batch_size\n\n train_ds = preprocess_data(\n train_data,\n shuffle=True,\n num_epochs=None,\n pack_examples=True,\n batch_size=batch_size,\n max_length=config.max_target_length,\n )\n\n eval_ds = preprocess_data(\n eval_data,\n shuffle=False,\n pack_examples=False,\n batch_size=eval_batch_size,\n max_length=config.max_eval_target_length,\n )\n\n predict_ds = preprocess_data(\n eval_data,\n shuffle=False,\n pack_examples=False,\n batch_size=eval_batch_size,\n max_length=config.max_predict_length,\n drop_remainder=False,\n )\n\n return train_ds, eval_ds, predict_ds, sp_tokenizer\n" }, { "path": "examples/lm1b/input_pipeline_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nimport pathlib\nimport sys\nimport tempfile\n\nfrom absl.testing import absltest\nimport tensorflow_datasets as tfds\n\nfrom configs import default\nimport input_pipeline\n\n# We just use different values here to verify that the input pipeline uses the\n# the correct value for the 3 different datasets.\n_TARGET_LENGTH = 32\n_EVAL_TARGET_LENGTH = 48\n_PREDICT_TARGET_LENGTH = 64\n\n\nclass InputPipelineTest(absltest.TestCase):\n\n def setUp(self):\n super().setUp()\n if sys.version_info >= (3, 13):\n self.skipTest('Test (and tensorflow-text) does not suport Python 3.13+')\n self.train_ds, self.eval_ds, self.predict_ds = self._get_datasets()\n\n def _get_datasets(self):\n config = default.get_config()\n config.per_device_batch_size = 1\n config.eval_per_device_batch_size = 2\n config.vocab_size = 32\n config.max_corpus_chars = 1000\n config.max_target_length = _TARGET_LENGTH\n config.max_eval_target_length = _EVAL_TARGET_LENGTH\n config.max_predict_length = _PREDICT_TARGET_LENGTH\n\n vocab_path = os.path.join(tempfile.mkdtemp(), 'sentencepiece_model')\n\n # Go two directories up to the root of the flax directory.\n flax_root_dir = pathlib.Path(__file__).parents[2]\n data_dir = str(flax_root_dir) + '/.tfds/metadata' # pylint: disable=unused-variable\n\n with tfds.testing.mock_data(num_examples=128, data_dir=data_dir):\n train_ds, eval_ds, predict_ds, _ = input_pipeline.get_datasets(\n n_devices=2, config=config, vocab_path=vocab_path\n )\n return train_ds, eval_ds, predict_ds\n\n def test_train_ds(self):\n expected_shape = [2, _TARGET_LENGTH] # 2 devices.\n # For training we pack multiple short examples in one example.\n # *_position and *_segmentation indicate the boundaries.\n for batch in self.train_ds.take(3):\n self.assertEqual(\n {k: v.shape.as_list() for k, v in batch.items()},\n {\n 'inputs': expected_shape,\n 'inputs_position': expected_shape,\n 'inputs_segmentation': expected_shape,\n 'targets': expected_shape,\n 'targets_position': expected_shape,\n 'targets_segmentation': expected_shape,\n },\n )\n\n def test_eval_ds(self):\n expected_shape = [4, _EVAL_TARGET_LENGTH] # 2 devices.\n for batch in self.eval_ds.take(3):\n self.assertEqual(\n {k: v.shape.as_list() for k, v in batch.items()},\n {\n 'inputs': expected_shape,\n 'targets': expected_shape,\n },\n )\n\n def test_predict_ds(self):\n expected_shape = [4, _PREDICT_TARGET_LENGTH] # 2 devices.\n for batch in self.predict_ds.take(3):\n self.assertEqual(\n {k: v.shape.as_list() for k, v in batch.items()},\n {\n 'inputs': expected_shape,\n 'targets': expected_shape,\n },\n )\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "examples/lm1b/main.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Main file for running the Language Modelling example with LM1B.\n\nThis file is intentionally kept short. The majority for logic is in libraries\nthat can be easily tested and imported in Colab.\n\"\"\"\n\nfrom absl import app\nfrom absl import flags\nfrom absl import logging\nfrom clu import platform\nimport jax\nfrom ml_collections import config_flags\nimport tensorflow as tf\n\nimport train\n\nFLAGS = flags.FLAGS\n\nflags.DEFINE_string('workdir', None, 'Directory to store model data.')\nconfig_flags.DEFINE_config_file(\n 'config',\n 'configs/default.py',\n 'File path to the training hyperparameter configuration.',\n lock_config=True,\n)\nflags.mark_flags_as_required(['config', 'workdir'])\n\n\ndef main(argv):\n if len(argv) > 1:\n raise app.UsageError('Too many command-line arguments.')\n\n # Hide any GPUs from TensorFlow. Otherwise TF might reserve memory and make\n # it unavailable to JAX.\n tf.config.experimental.set_visible_devices([], 'GPU')\n\n logging.info('JAX process: %d / %d', jax.process_index(), jax.process_count())\n logging.info('JAX local devices: %r', jax.local_devices())\n\n # Add a note so that we can tell which task is which JAX host.\n # (Depending on the platform task 0 is not guaranteed to be host 0)\n platform.work_unit().set_task_status(\n f'process_index: {jax.process_index()}, '\n f'process_count: {jax.process_count()}'\n )\n platform.work_unit().create_artifact(\n platform.ArtifactType.DIRECTORY, FLAGS.workdir, 'workdir'\n )\n\n train.train_and_evaluate(FLAGS.config, FLAGS.workdir)\n\n\nif __name__ == '__main__':\n jax.config.config_with_absl()\n app.run(main)\n" }, { "path": "examples/lm1b/models.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Transformer-based language model.\n\nReusing decoder only model from examples/wmt.\n\"\"\"\n\n# pylint: disable=attribute-defined-outside-init\n# See issue #620.\n# pytype: disable=wrong-arg-count\n# pytype: disable=wrong-keyword-args\n# pytype: disable=attribute-error\n\nfrom typing import Any, Optional\nfrom collections.abc import Callable\n\nfrom flax import linen as nn\nfrom flax import struct\nfrom jax import lax\nimport jax.numpy as jnp\nimport numpy as np\n\n\n@struct.dataclass\nclass TransformerConfig:\n \"\"\"Global hyperparameters used to minimize obnoxious kwarg plumbing.\"\"\"\n\n vocab_size: int\n output_vocab_size: int\n share_embeddings: bool = False\n logits_via_embedding: bool = False\n dtype: Any = jnp.float32\n emb_dim: int = 512\n num_heads: int = 8\n num_layers: int = 6\n qkv_dim: int = 512\n mlp_dim: int = 2048\n max_len: int = 2048\n dropout_rate: float = 0.1\n attention_dropout_rate: float = 0.1\n deterministic: bool = False\n decode: bool = False\n kernel_init: Callable = nn.initializers.xavier_uniform()\n bias_init: Callable = nn.initializers.normal(stddev=1e-6)\n posemb_init: Callable | None = None\n\n\ndef shift_right(x, axis=1):\n \"\"\"Shift the input to the right by padding and slicing on axis.\"\"\"\n pad_widths = [(0, 0)] * len(x.shape)\n pad_widths[axis] = (1, 0)\n padded = jnp.pad(\n x, pad_widths, mode='constant', constant_values=x.dtype.type(0)\n )\n return lax.dynamic_slice_in_dim(padded, 0, padded.shape[axis] - 1, axis)\n\n\ndef shift_inputs(x, segment_ids=None, axis=1):\n \"\"\"Shift inputs and replace EOS by 0 for packed inputs.\"\"\"\n shifted = shift_right(x, axis=axis)\n # For packed targets, the first shifted token of a new sequence is made\n # 0, rather than being the EOS token for the last sequence.\n if segment_ids is not None:\n shifted *= segment_ids == shift_right(segment_ids, axis=axis)\n return shifted\n\n\ndef sinusoidal_init(max_len=2048, min_scale=1.0, max_scale=10000.0):\n \"\"\"1D Sinusoidal Position Embedding Initializer.\n\n Args:\n max_len: maximum possible length for the input.\n min_scale: float: minimum frequency-scale in sine grating.\n max_scale: float: maximum frequency-scale in sine grating.\n\n Returns:\n output: init function returning `(1, max_len, d_feature)`\n \"\"\"\n\n def init(key, shape, dtype=np.float32):\n \"\"\"Sinusoidal init.\"\"\"\n del key, dtype\n d_feature = shape[-1]\n pe = np.zeros((max_len, d_feature), dtype=np.float32)\n position = np.arange(0, max_len)[:, np.newaxis]\n scale_factor = -np.log(max_scale / min_scale) / (d_feature // 2 - 1)\n div_term = min_scale * np.exp(np.arange(0, d_feature // 2) * scale_factor)\n pe[:, : d_feature // 2] = np.sin(position * div_term)\n pe[:, d_feature // 2 : 2 * (d_feature // 2)] = np.cos(position * div_term)\n pe = pe[np.newaxis, :, :] # [1, max_len, d_feature]\n return jnp.array(pe)\n\n return init\n\n\nclass AddPositionEmbs(nn.Module):\n \"\"\"Adds (optionally learned) positional embeddings to the inputs.\n\n Args:\n config: TransformerConfig dataclass containing hyperparameters.\n decode: whether to run in single-position autoregressive mode.\n \"\"\"\n\n config: TransformerConfig\n decode: bool = False\n\n @nn.compact\n def __call__(self, inputs, inputs_positions=None):\n \"\"\"Applies AddPositionEmbs module.\n\n By default this layer uses a fixed sinusoidal embedding table. If a\n learned position embedding is desired, pass an initializer to\n posemb_init in the configuration.\n\n Args:\n inputs: input data.\n inputs_positions: input position indices for packed sequences.\n\n Returns:\n output: `(bs, timesteps, in_dim)`\n \"\"\"\n config = self.config\n # inputs.shape is (batch_size, seq_len, emb_dim)\n assert inputs.ndim == 3, (\n 'Number of dimensions should be 3, but it is: %d' % inputs.ndim\n )\n length = inputs.shape[1]\n pos_emb_shape = (1, config.max_len, inputs.shape[-1])\n if config.posemb_init is None:\n # Use a fixed (non-learned) sinusoidal position embedding.\n pos_embedding = sinusoidal_init(max_len=config.max_len)(\n None, pos_emb_shape, None\n )\n else:\n pos_embedding = self.param(\n 'pos_embedding', config.posemb_init, pos_emb_shape\n )\n pe = pos_embedding[:, :length, :]\n\n # We use a cache position index for tracking decoding position.\n if self.decode:\n is_initialized = self.has_variable('cache', 'cache_index')\n cache_index = self.variable(\n 'cache', 'cache_index', lambda: jnp.array(0, dtype=jnp.uint32)\n )\n if is_initialized:\n i = cache_index.value\n cache_index.value = i + 1\n _, _, df = pos_embedding.shape\n pe = lax.dynamic_slice(pos_embedding, jnp.array((0, i, 0)), (1, 1, df))\n if inputs_positions is None:\n # normal unpacked case:\n return inputs + pe\n else:\n # for packed data we need to use known position indices:\n return inputs + jnp.take(pe[0], inputs_positions, axis=0)\n\n\nclass MlpBlock(nn.Module):\n \"\"\"Transformer MLP / feed-forward block.\n\n Args:\n config: TransformerConfig dataclass containing hyperparameters.\n out_dim: optionally specify out dimension.\n \"\"\"\n\n config: TransformerConfig\n out_dim: int | None = None\n\n @nn.compact\n def __call__(self, inputs):\n \"\"\"Applies Transformer MlpBlock module.\"\"\"\n config = self.config\n actual_out_dim = inputs.shape[-1] if self.out_dim is None else self.out_dim\n x = nn.Dense(\n config.mlp_dim,\n dtype=config.dtype,\n kernel_init=nn.with_logical_partitioning(\n config.kernel_init, ('embed', 'mlp')\n ),\n bias_init=nn.with_logical_partitioning(config.bias_init, ('mlp',)),\n )(inputs)\n x = nn.relu(x)\n x = nn.Dropout(rate=config.dropout_rate)(\n x, deterministic=config.deterministic\n )\n output = nn.Dense(\n actual_out_dim,\n dtype=config.dtype,\n kernel_init=nn.with_logical_partitioning(\n config.kernel_init, ('mlp', 'embed')\n ),\n bias_init=nn.with_logical_partitioning(config.bias_init, ('embed',)),\n )(x)\n output = nn.Dropout(rate=config.dropout_rate)(\n output, deterministic=config.deterministic\n )\n return output\n\n\nclass EncoderDecoder1DBlock(nn.Module):\n \"\"\"Transformer encoder-decoder layer.\n\n Args:\n config: TransformerConfig dataclass containing hyperparameters.\n \"\"\"\n\n config: TransformerConfig\n\n @nn.compact\n def __call__(self, inputs, decoder_mask=None, encoder_decoder_mask=None):\n \"\"\"Applies EncoderDecoder1DBlock module.\n\n Args:\n inputs: input data for decoder\n decoder_mask: decoder self-attention mask.\n encoder_decoder_mask: encoder-decoder attention mask.\n\n Returns:\n output after transformer encoder-decoder block.\n \"\"\"\n config = self.config\n\n # Decoder block.\n assert inputs.ndim == 3\n x = nn.LayerNorm(\n dtype=config.dtype,\n bias_init=nn.with_logical_partitioning(\n nn.initializers.zeros, ('embed',)\n ),\n scale_init=nn.with_logical_partitioning(\n nn.initializers.ones, ('embed',)\n ),\n )(inputs)\n x = nn.MultiHeadDotProductAttention(\n num_heads=config.num_heads,\n dtype=config.dtype,\n qkv_features=config.qkv_dim,\n kernel_init=nn.with_logical_partitioning(\n config.kernel_init, ('embed', 'kv')\n ),\n bias_init=nn.with_logical_partitioning(config.bias_init, ('embed',)),\n use_bias=False,\n broadcast_dropout=False,\n dropout_rate=config.attention_dropout_rate,\n deterministic=config.deterministic,\n decode=config.decode,\n )(x, mask=decoder_mask)\n x = nn.Dropout(rate=config.dropout_rate)(\n x, deterministic=config.deterministic\n )\n x = x + inputs\n\n # MLP block.\n z = nn.LayerNorm(\n dtype=config.dtype,\n bias_init=nn.with_logical_partitioning(\n nn.initializers.zeros, ('embed',)\n ),\n scale_init=nn.with_logical_partitioning(\n nn.initializers.ones, ('embed',)\n ),\n )(x)\n z = MlpBlock(config=config)(z)\n\n return x + z\n\n\nclass Decoder(nn.Module):\n \"\"\"Transformer Model Decoder for sequence to sequence translation.\n\n Args:\n config: TransformerConfig dataclass containing hyperparameters.\n shared_embedding: a shared embedding layer to use.\n \"\"\"\n\n config: TransformerConfig\n shared_embedding: Any = None\n\n @nn.compact\n def __call__(\n self,\n inputs,\n inputs_positions=None,\n inputs_segmentation=None,\n decoder_mask=None,\n encoder_decoder_mask=None,\n ):\n \"\"\"Applies Transformer model on the inputs.\n\n Args:\n encoded: encoded input data from encoder.\n inputs: input data.\n inputs_positions: input subsequence positions for packed examples.\n inputs_segmentation: input segmentation info for packed examples.\n decoder_mask: decoder self-attention mask.\n encoder_decoder_mask: encoder-decoder attention mask.\n\n Returns:\n output of a transformer decoder.\n \"\"\"\n config = self.config\n assert inputs.ndim == 2 # (batch, len)\n\n # Target Embedding\n if self.shared_embedding is None:\n output_embed = nn.Embed(\n num_embeddings=config.output_vocab_size,\n features=config.emb_dim,\n embedding_init=nn.with_logical_partitioning(\n nn.initializers.normal(stddev=1.0), ('vocab', 'embed')\n ),\n )\n else:\n output_embed = self.shared_embedding\n\n y = inputs.astype('int32')\n if not config.decode:\n y = shift_inputs(y, segment_ids=inputs_segmentation)\n y = output_embed(y)\n y = AddPositionEmbs(\n config=config, decode=config.decode, name='posembed_output'\n )(y, inputs_positions=inputs_positions)\n y = nn.Dropout(rate=config.dropout_rate)(\n y, deterministic=config.deterministic\n )\n\n y = y.astype(config.dtype)\n\n # Target-Input Decoder\n for lyr in range(config.num_layers):\n y = EncoderDecoder1DBlock(\n config=config, name=f'encoderdecoderblock_{lyr}'\n )(y, decoder_mask=decoder_mask, encoder_decoder_mask=encoder_decoder_mask)\n y = nn.LayerNorm(\n dtype=config.dtype,\n name='encoderdecoder_norm',\n bias_init=nn.with_logical_partitioning(\n nn.initializers.zeros, ('embed',)\n ),\n scale_init=nn.with_logical_partitioning(\n nn.initializers.ones, ('embed',)\n ),\n )(y)\n\n # Decoded Logits\n if config.logits_via_embedding:\n # Use the transpose of embedding matrix for logit transform.\n logits = output_embed.attend(y.astype(jnp.float32))\n # Correctly normalize pre-softmax logits for this shared case.\n logits = logits / jnp.sqrt(y.shape[-1])\n else:\n logits = nn.Dense(\n config.output_vocab_size,\n dtype=config.dtype,\n kernel_init=nn.with_logical_partitioning(\n config.kernel_init, ('embed', 'vocab')\n ),\n bias_init=nn.with_logical_partitioning(config.bias_init, ('vocab',)),\n name='logitdense',\n )(y)\n return logits\n\n\nclass TransformerLM(nn.Module):\n \"\"\"Transformer pure decoder stack for language modelling.\n\n Args:\n config: TransformerConfig dataclass containing hyperparameters.\n \"\"\"\n\n config: TransformerConfig\n\n @nn.compact\n def __call__(self, inputs, inputs_positions=None, inputs_segmentation=None):\n \"\"\"Applies TransformerLM on the inputs.\n\n Args:\n inputs: target data.\n inputs_positions: input subsequence positions for packed examples.\n inputs_segmentation: input segmentation info for packed examples.\n\n Returns:\n logits array from transformer decoder.\n \"\"\"\n config = self.config\n\n # Make padding attention masks.\n if config.decode:\n # for fast autoregressive decoding we use no decoder mask\n decoder_mask = None\n else:\n decoder_mask = nn.combine_masks(\n nn.make_attention_mask(inputs > 0, inputs > 0, dtype=config.dtype),\n nn.make_causal_mask(inputs, dtype=config.dtype),\n )\n\n # Add segmentation block-diagonal attention masks if using segmented data.\n if inputs_segmentation is not None:\n decoder_mask = nn.combine_masks(\n decoder_mask,\n nn.make_attention_mask(\n inputs_segmentation,\n inputs_segmentation,\n jnp.equal,\n dtype=config.dtype,\n ),\n )\n\n logits = Decoder(config=config, shared_embedding=None, name='decoder')(\n inputs,\n inputs_positions=inputs_positions,\n inputs_segmentation=inputs_segmentation,\n decoder_mask=decoder_mask,\n encoder_decoder_mask=None,\n )\n return logits.astype(self.config.dtype)\n" }, { "path": "examples/lm1b/requirements.txt", "content": "absl-py==1.4.0\nclu==0.0.9\nflax==0.6.11\njax==0.4.13\n--find-links https://storage.googleapis.com/jax-releases/jax_releases.html\njaxlib==0.4.13+cuda11.cudnn82 # Make sure CUDA version matches the base image.\nml-collections==0.1.1\nnumpy==1.24.3\noptax==0.1.5\nsentencepiece==0.1.99\ntensorflow==2.13.0\ntensorflow-datasets==4.9.2\ntensorflow-text==2.13.0\n" }, { "path": "examples/lm1b/temperature_sampler.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Fast decoding routines for inference from a trained language model.\"\"\"\n\nfrom jax import lax\nfrom jax import random\nimport jax.numpy as jnp\n\n\n# Constants\n# The default End-of-Sentence token id is 2 (SentencePiece).\nEOS_ID = 2\n\n\ndef temperature_sample(\n prompt_inputs,\n init_cache,\n tokens_to_logits,\n prng_key,\n temperature=1.0,\n topk=20,\n eos_token=EOS_ID,\n):\n \"\"\"Temperature sampling for language model generation.\n\n Args:\n prompt_inputs: array: [batch_size, max_decode_len] int32 sequence of tokens.\n init_cache: flax attention cache.\n tokens_to_logits: fast autoregressive decoder function taking single token\n slices and cache and returning next-token logits and updated cache.\n prng_key: JAX PRNGKey.\n temperature: float: sampling temperature factor. As it approaches\n zero this becomes equivalent to greedy sampling.\n topk: integer: if nonzero only use the top-k logits to sample next token,\n if zero don't use any cutoff and sample from full logits over vocabulary.\n eos_token: int: end-of-sentence token for target vocabulary.\n\n Returns:\n Array of sampled sequences: [batch_size, max_decode_len]\n \"\"\"\n batch_size = prompt_inputs.shape[0]\n max_decode_len = prompt_inputs.shape[1]\n end_marker = jnp.array(eos_token)\n temperature = jnp.array(temperature)\n\n # Initialize sampling loop state.\n # initial loop PRNGKey\n rng0 = prng_key\n # loop position counter.\n i0 = jnp.array(-1)\n # per batch-item holding current token in loop.\n token0 = jnp.zeros((batch_size, 1), dtype=jnp.int32)\n # per batch-item state bit indicating if sentence has finished.\n ended0 = jnp.zeros((batch_size, 1), dtype=jnp.bool_)\n # (batch, length) array containing prefix prompt tokens for sampling loop\n # as well as the generated output of newly sampled tokens.\n sequences0 = prompt_inputs\n # Sampling loop state is stored in a simple tuple.\n sampling_loop_init_state = (i0, sequences0, init_cache, token0, ended0, rng0)\n\n def sampling_loop_cond_fn(state):\n \"\"\"Sampling loop termination condition.\"\"\"\n (i, _, _, _, ended, _) = state\n # Have we reached max decoding length?\n not_at_end = i < max_decode_len - 1\n # Have all sampled sequences reached an end marker?\n all_sequences_ended = jnp.all(ended)\n return not_at_end & (~all_sequences_ended)\n\n def sampling_loop_body_fn(state):\n \"\"\"Sampling loop state update.\"\"\"\n i, sequences, cache, cur_token, ended, rng = state\n # Split RNG for sampling.\n rng1, rng2 = random.split(rng)\n # Call fast-decoder model on current tokens to get next-position logits.\n logits, new_cache = tokens_to_logits(cur_token, cache)\n # Sample next token from logits.\n # TODO(levskaya): add top-p \"nucleus\" sampling option.\n if topk:\n # Get top-k logits and their indices, sample within these top-k tokens.\n topk_logits, topk_idxs = lax.top_k(logits, topk)\n topk_token = jnp.expand_dims(\n random.categorical(rng1, topk_logits / temperature).astype(jnp.int32),\n axis=-1,\n )\n # Return the original indices corresponding to the sampled top-k tokens.\n next_token = jnp.squeeze(\n jnp.take_along_axis(topk_idxs, topk_token, axis=-1), axis=-1\n )\n else:\n next_token = random.categorical(rng1, logits / temperature).astype(\n jnp.int32\n )\n # Only use sampled tokens if we're past provided prefix tokens.\n out_of_prompt = sequences[:, i + 1] == 0\n next_token = (\n next_token * out_of_prompt + sequences[:, i + 1] * ~out_of_prompt\n )\n # If end-marker reached for batch item, only emit padding tokens.\n next_token_or_endpad = next_token[None] * ~ended\n ended |= next_token_or_endpad == end_marker\n # Add current sampled tokens to recorded sequences.\n new_sequences = lax.dynamic_update_slice(\n sequences, next_token_or_endpad, (0, i + 1)\n )\n return (i + 1, new_sequences, new_cache, next_token_or_endpad, ended, rng2)\n\n # Run sampling loop and collect final state.\n final_state = lax.while_loop(\n sampling_loop_cond_fn, sampling_loop_body_fn, sampling_loop_init_state\n )\n\n # Pick part of the state corresponding to the sampled sequences.\n final_sequences = final_state[1]\n return final_sequences\n" }, { "path": "examples/lm1b/temperature_sampler_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom absl.testing import absltest\nimport jax\nimport jax.numpy as jnp\nimport numpy as np\n\nfrom temperature_sampler import temperature_sample\n\n\njax.config.update('jax_disable_most_optimizations', True)\n\n\nclass TestTemperatureSampler(absltest.TestCase):\n\n def test_temperature_sampler(self):\n tokens = jnp.array([[5, 0, 0, 0]], dtype=jnp.int32)\n cache = None\n key = jax.random.PRNGKey(0)\n\n def tokens_to_logits(tokens, cache):\n jax.debug.print('tokens: {}', tokens)\n logits = jax.nn.one_hot(tokens[..., -1:] + 1, 10)\n logits = jnp.where(logits < 0.5, float('-inf'), logits)\n logits = logits.squeeze(axis=1)\n return logits, cache\n\n new_tokens = temperature_sample(\n tokens, cache, tokens_to_logits, key, topk=5\n )\n\n np.testing.assert_array_equal(new_tokens, [[5, 6, 7, 8]])\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "examples/lm1b/tokenizer.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Provides op for tokenizing a dataset.\"\"\"\n\nimport os\nimport sys\nimport tempfile\nimport time\nfrom typing import Any, Dict, Tuple\nfrom collections.abc import Iterable\n\nfrom absl import logging\nimport dataclasses\nimport jax\nfrom sentencepiece import SentencePieceTrainer\nimport tensorflow as tf\nif sys.version_info < (3, 13):\n import tensorflow_text as tftxt\n\nFeatures = dict[str, tf.Tensor]\n\n\ndef _dump_chars_to_textfile(\n dataset: tf.data.Dataset,\n maxchars: int = int(1e7),\n data_keys=('inputs', 'targets'),\n) -> tuple[str, int]:\n \"\"\"Write part of a TFDS sentence dataset to lines in a text file.\n\n Args:\n dataset: tf.dataset containing string-data.\n maxchars: int: approximate number of characters to save from dataset.\n data_keys: Tuple[str]: what keys in dataset to dump from.\n\n Returns:\n name of temp file with dataset bytes, exact number of characters dumped.\n \"\"\"\n char_count = 0\n ds_iter = dataset.as_numpy_iterator()\n with tempfile.NamedTemporaryFile(\n delete=False, prefix='/tmp/ds_chars'\n ) as outfp:\n while char_count < maxchars:\n example = next(ds_iter)\n for k in data_keys:\n line = example[k] + b'\\n'\n char_count += len(line)\n outfp.write(line)\n return outfp.name, char_count\n\n\ndef _train_sentencepiece(\n dataset: tf.data.Dataset,\n *,\n vocab_size: int,\n maxchars: int = int(1e7),\n model_path: str,\n model_type: str = 'unigram',\n character_coverage: float = 1.0,\n data_keys=('inputs', 'targets'),\n):\n \"\"\"Train SentencePiece tokenizer from subset of tf dataset.\n\n Args:\n dataset: tf.dataset\n vocab_size: int: size of vocab tokens to train.\n maxchars: int: number of characters to use for sentencepiece training.\n model_path: str: path of model file to save vocab model to.\n model_type: str: type of sentencepiece vocab to train.\n character_coverage: amount of characters covered by the model, good defaults\n are 0.9995 for languages with rich character set like Japanese or Chinese\n and 1.0 for other languages with small character set.\n data_keys: Tuple[str]: keys of dataset to use for training.\n\n Returns:\n path to the trained sentencepiece vocabulary model.\n \"\"\"\n if model_path.startswith('gs://'):\n abs_model_path = model_path\n else:\n abs_model_path = os.path.abspath(os.path.expanduser(model_path))\n fname, _ = _dump_chars_to_textfile(\n dataset, maxchars=maxchars, data_keys=data_keys\n )\n with tempfile.NamedTemporaryFile(\n delete=False, prefix='/tmp/sp_tmp'\n ) as model_fp:\n pass # we just want a prefix'd tmp-filename\n argstr = ' '.join([\n f'--input={fname}',\n f'--vocab_size={vocab_size}',\n f'--character_coverage={character_coverage}',\n f'--model_prefix={model_fp.name}',\n f'--model_type={model_type}',\n ])\n SentencePieceTrainer.Train(argstr)\n if jax.process_index() == 0:\n # Use an intermediate filename that is renamed to the target name to address\n # create and fill delays.\n copy_rename_path = abs_model_path + '.rntmp'\n tf.io.gfile.copy(model_fp.name + '.model', copy_rename_path, overwrite=True)\n tf.io.gfile.rename(copy_rename_path, abs_model_path, overwrite=True)\n logging.info('copied %s to %s', model_fp.name + '.model', abs_model_path)\n else:\n while not tf.io.gfile.exists(abs_model_path):\n time.sleep(1)\n time.sleep(1)\n return abs_model_path\n\n\ndef _load_sentencepiece_tokenizer(\n model_path: str,\n add_bos: bool = False,\n add_eos: bool = True,\n reverse: bool = False,\n):\n \"\"\"Load a tf-text SentencePiece tokenizer from given model filepath.\"\"\"\n with tf.io.gfile.GFile(model_path, 'rb') as model_fp:\n sp_model = model_fp.read()\n sp_tokenizer = tftxt.SentencepieceTokenizer(\n model=sp_model, add_bos=add_bos, add_eos=add_eos, reverse=reverse\n )\n return sp_tokenizer\n\n\ndef load_or_train_tokenizer(\n dataset: tf.data.Dataset,\n *,\n vocab_path: str,\n vocab_size: int,\n max_corpus_chars: int,\n data_keys: tuple[str, str] = ('inputs', 'targets'),\n):\n \"\"\"Loads the tokenizer at `vocab_path` or trains a one from `dataset`.\"\"\"\n try:\n return _load_sentencepiece_tokenizer(vocab_path)\n except tf.errors.NotFoundError:\n logging.info('SentencePiece vocab not found, building one from data.')\n vocab_path = _train_sentencepiece(\n dataset,\n vocab_size=vocab_size,\n maxchars=max_corpus_chars,\n model_path=vocab_path,\n data_keys=data_keys,\n )\n return _load_sentencepiece_tokenizer(vocab_path)\n\n\n@dataclasses.dataclass\nclass TokenizeOp:\n sp_tokenizer: Any\n data_keys: Iterable[str] = ('inputs', 'targets')\n\n def __call__(self, features: Features) -> Features:\n for k in self.data_keys:\n features[k] = self.sp_tokenizer.tokenize(features[k])\n return features\n" }, { "path": "examples/lm1b/train.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Language Modeling example.\n\nThis script trains a Transformer on a LM1B dataset.\n\"\"\"\n\n# pytype: disable=wrong-arg-count\n# pytype: disable=attribute-error\n\nimport collections\nimport os\n\nfrom absl import logging\nfrom clu import metric_writers\nfrom clu import periodic_actions\nfrom flax import linen as nn\nfrom flax.training import checkpoints\nfrom flax.training import common_utils\nimport jax\nfrom jax import random\nimport jax.numpy as jnp\nfrom jax.sharding import PartitionSpec as P, Mesh, NamedSharding\nimport ml_collections\nimport numpy as np\nimport optax\nimport tensorflow as tf\n\nimport input_pipeline\nimport models\nimport temperature_sampler\nimport utils\n\n\ndef rsqrt_schedule(\n init_value: float,\n shift: int = 0,\n):\n \"\"\"Applies a reverse square-root schedule.\n\n The reverse square root schedule is simply `lr = init_value / sqrt(step)`.\n\n Args:\n init_value: Base learning rate (before applying the rsqrt schedule).\n shift: How many steps the rsqrt should be shifted. Shifting the rsqrt\n schedule makes it less steep in the beginning (close to 0).\n\n Returns:\n A schedule that applies the reverse square root.\n \"\"\"\n\n def schedule(count):\n return init_value * (count + shift) ** -0.5 * shift**0.5\n\n return schedule\n\n\ndef create_learning_rate_schedule(learning_rate: float, warmup_steps: int):\n \"\"\"Creates a rsqrt schedule with linear warmup.\"\"\"\n return optax.join_schedules(\n [\n optax.linear_schedule(\n init_value=0,\n end_value=learning_rate,\n transition_steps=warmup_steps,\n ),\n rsqrt_schedule(init_value=learning_rate, shift=warmup_steps),\n ],\n boundaries=[warmup_steps],\n )\n\n\ndef compute_weighted_cross_entropy(\n logits, targets, weights=None, label_smoothing=0.0\n):\n \"\"\"Compute weighted cross entropy and entropy for log probs and targets.\n\n Args:\n logits: [batch, length, num_classes] float array.\n targets: categorical targets [batch, length] int array.\n weights: None or array of shape [batch, length].\n label_smoothing: label smoothing constant, used to determine the on and off\n values.\n\n Returns:\n Tuple of scalar loss and batch normalizing factor.\n \"\"\"\n if logits.ndim != targets.ndim + 1:\n raise ValueError(\n \"Incorrect shapes. Got shape %s logits and %s targets\"\n % (str(logits.shape), str(targets.shape))\n )\n vocab_size = logits.shape[-1]\n confidence = 1.0 - label_smoothing\n low_confidence = (1.0 - confidence) / (vocab_size - 1)\n normalizing_constant = -(\n confidence * jnp.log(confidence)\n + (vocab_size - 1) * low_confidence * jnp.log(low_confidence + 1e-20)\n )\n soft_targets = common_utils.onehot(\n targets, vocab_size, on_value=confidence, off_value=low_confidence\n )\n\n loss = -jnp.sum(soft_targets * nn.log_softmax(logits), axis=-1)\n loss = loss - normalizing_constant\n\n normalizing_factor = np.prod(targets.shape)\n if weights is not None:\n loss = loss * weights\n normalizing_factor = weights.sum()\n\n return loss.sum(), normalizing_factor\n\n\ndef compute_weighted_accuracy(logits, targets, weights=None):\n \"\"\"Compute weighted accuracy for log probs and targets.\n\n Args:\n logits: [batch, length, num_classes] float array.\n targets: categorical targets [batch, length] int array.\n weights: None or array of shape [batch, length]\n\n Returns:\n Tuple of scalar loss and batch normalizing factor.\n \"\"\"\n if logits.ndim != targets.ndim + 1:\n raise ValueError(\n \"Incorrect shapes. Got shape %s logits and %s targets\"\n % (str(logits.shape), str(targets.shape))\n )\n loss = jnp.equal(jnp.argmax(logits, axis=-1), targets)\n normalizing_factor = np.prod(logits.shape[:-1])\n if weights is not None:\n loss = loss * weights\n normalizing_factor = weights.sum()\n\n return loss.sum(), normalizing_factor\n\n\ndef compute_metrics(logits, labels, weights, label_smoothing=0.0):\n \"\"\"Compute summary metrics.\"\"\"\n loss, weight_sum = compute_weighted_cross_entropy(\n logits, labels, weights, label_smoothing\n )\n acc, _ = compute_weighted_accuracy(logits, labels, weights)\n metrics = {\n \"loss\": loss,\n \"accuracy\": acc,\n \"denominator\": weight_sum,\n }\n return metrics\n\n\n# Primary training / eval / decode step functions.\n# -----------------------------------------------------------------------------\n\n\ndef train_step(\n state,\n batch,\n config,\n learning_rate_fn,\n label_smoothing=0.0,\n dropout_rng=None,\n):\n \"\"\"Perform a single training step.\"\"\"\n # X_position and X_segmentation are needed only when using \"packed examples\"\n # where multiple sequences are packed into the same example with this\n # metadata.\n # if such features are not present they are ignored and the example is treated\n # like a normal, unpacked sequence example.\n train_keys = [\"inputs\", \"inputs_position\", \"inputs_segmentation\"]\n (inputs, inputs_positions, inputs_segmentation) = (\n batch.get(k, None) for k in train_keys\n )\n\n weights = jnp.where(inputs > 0, 1, 0).astype(jnp.float32)\n\n dropout_rng = jax.random.fold_in(dropout_rng, state.step)\n\n def loss_fn(params):\n \"\"\"loss function used for training.\"\"\"\n logits = models.TransformerLM(config).apply(\n {\"params\": params},\n inputs,\n inputs_positions=inputs_positions,\n inputs_segmentation=inputs_segmentation,\n rngs={\"dropout\": dropout_rng},\n )\n\n loss, weight_sum = compute_weighted_cross_entropy(\n logits, inputs, weights, label_smoothing\n )\n mean_loss = loss / weight_sum\n return mean_loss, logits\n\n step = state.step\n lr = learning_rate_fn(step)\n grad_fn = jax.value_and_grad(loss_fn, has_aux=True)\n (_, logits), grads = grad_fn(state.params)\n new_state = state.apply_gradients(grads=grads)\n metrics = compute_metrics(logits, inputs, weights)\n metrics[\"learning_rate\"] = lr\n\n return new_state, metrics\n\n\ndef eval_step(params, batch, config, label_smoothing=0.0):\n \"\"\"Calculate evaluation metrics on a batch.\"\"\"\n inputs = batch[\"inputs\"]\n weights = jnp.where(inputs > 0, 1.0, 0.0)\n logits = models.TransformerLM(config).apply({\"params\": params}, inputs)\n\n return compute_metrics(logits, inputs, weights, label_smoothing)\n\n\ndef predict_step(\n inputs, params, rngkey, eos_id, max_decode_len, config, temperature, top_k\n):\n \"\"\"Predict language model on a batch.\"\"\"\n target_shape = (inputs.shape[0], max_decode_len) + inputs.shape[2:]\n initial_variables = models.TransformerLM(config).init(\n jax.random.PRNGKey(0), jnp.ones(target_shape, config.dtype)\n )\n cache = initial_variables[\"cache\"]\n\n def tokens_ids_to_logits(flat_ids, flat_cache):\n \"\"\"Token slice to logits from decoder model.\"\"\"\n # --> [batch * beam, 1, vocab]\n flat_logits, new_vars = models.TransformerLM(config).apply(\n {\"params\": params, \"cache\": flat_cache}, flat_ids, mutable=[\"cache\"]\n )\n new_flat_cache = new_vars[\"cache\"]\n # Remove singleton sequence-length dimension:\n # [batch, 1, vocab] --> [batch, vocab]\n flat_logits = flat_logits.squeeze(axis=1)\n return flat_logits, new_flat_cache\n\n # Using the above-defined single-step decoder function, run a\n # beam search over possible sequences given input encoding.\n seqs = temperature_sampler.temperature_sample(\n inputs,\n cache,\n tokens_ids_to_logits,\n rngkey,\n temperature=temperature,\n topk=top_k,\n eos_token=eos_id,\n )\n\n return seqs\n\n\n# Utils for prediction\n# -----------------------------------------------------------------------------\n\n\ndef pad_examples(x, desired_batch_size):\n \"\"\"Expand batch to desired size by repeating last slice.\"\"\"\n batch_pad = desired_batch_size - x.shape[0]\n return np.concatenate([x, np.tile(x[-1], (batch_pad, 1))], axis=0)\n\n\ndef tohost(x):\n \"\"\"Collect batches from all devices to host and flatten batch dimensions.\"\"\"\n n_device, n_batch, *remaining_dims = x.shape\n return np.array(x).reshape((n_device * n_batch,) + tuple(remaining_dims))\n\n\ndef evaluate(\n *,\n jit_eval_step,\n params,\n eval_ds: tf.data.Dataset,\n num_eval_steps: int,\n config,\n):\n \"\"\"Evaluate the target an return a dictionary with the metrics.\"\"\"\n logging.info(\"Gathering evaluation metrics.\")\n eval_metrics = []\n eval_iter = iter(eval_ds) # pytype: disable=wrong-arg-types\n for _, eval_batch in zip(range(num_eval_steps), eval_iter):\n eval_batch = jax.tree_util.tree_map(lambda x: x._numpy(), eval_batch) # pylint: disable=protected-access\n metrics = jit_eval_step(params, eval_batch, config)\n eval_metrics.append(metrics)\n eval_metrics = common_utils.stack_forest(eval_metrics)\n eval_metrics_sums = jax.tree_util.tree_map(jnp.sum, eval_metrics)\n eval_denominator = eval_metrics_sums.pop(\"denominator\")\n eval_summary = jax.tree_util.tree_map(\n lambda x: x / eval_denominator, # pylint: disable=cell-var-from-loop\n eval_metrics_sums,\n )\n return eval_summary\n\n\ndef generate_prediction(\n *,\n jit_pred_step,\n params,\n tokenized_prompts,\n eos_id,\n inference_rng,\n decode_tokens,\n config,\n predict_config,\n):\n \"\"\"Generate text from the prompt.\"\"\"\n n_devices = jax.local_device_count()\n\n logging.info(\"Generating text.\")\n predictions = []\n # Use batch of prompts provided by user.\n for pred_batch in jnp.array_split(\n tokenized_prompts, int(np.ceil(len(tokenized_prompts) / n_devices))\n ):\n cur_pred_batch_size = pred_batch.shape[0]\n if cur_pred_batch_size % n_devices:\n padded_size = int(np.ceil(cur_pred_batch_size / n_devices) * n_devices)\n pred_batch = jax.tree_util.tree_map(\n lambda x: pad_examples(x, padded_size), pred_batch\n ) # pylint: disable=cell-var-from-loop\n pred_batch = common_utils.shard(pred_batch)\n inference_rng, sub_rng = random.split(inference_rng)\n inference_rngs = random.split(sub_rng, n_devices)\n\n predicted = jit_pred_step(\n pred_batch,\n params,\n inference_rngs,\n eos_id,\n config.max_predict_length,\n predict_config,\n config.sampling_temperature,\n config.sampling_top_k,\n )\n predicted = tohost(predicted)\n # Iterate through non-padding examples of batch.\n for s in predicted[:cur_pred_batch_size]:\n prediction = decode_tokens(s)\n logging.info(\"Sample: %s\", str(prediction))\n predictions.append(prediction)\n\n # Save generated texts for tensorboard.\n exemplars = \"\"\n for prediction in predictions:\n exemplars += f\"{prediction}\\n\\n\"\n return exemplars\n\n\ndef train_and_evaluate(config: ml_collections.ConfigDict, workdir: str):\n \"\"\"Runs a training and evaluation loop.\n\n Args:\n config: Configuration to use.\n workdir: Working directory for checkpoints and TF summaries. If this\n contains checkpoint training will be resumed from the latest checkpoint.\n \"\"\"\n tf.io.gfile.makedirs(workdir)\n\n vocab_path = config.vocab_path\n if vocab_path is None:\n vocab_path = os.path.join(workdir, \"sentencepiece_model\")\n config.vocab_path = vocab_path\n tf.io.gfile.makedirs(os.path.split(vocab_path)[0])\n\n # Load Dataset\n # ---------------------------------------------------------------------------\n logging.info(\"Initializing dataset.\")\n train_ds, eval_ds, _, encoder = input_pipeline.get_datasets(\n n_devices=jax.local_device_count(), config=config, vocab_path=vocab_path\n )\n\n train_iter = iter(train_ds)\n vocab_size = int(encoder.vocab_size())\n eos_id = temperature_sampler.EOS_ID # Default Sentencepiece EOS token.\n\n def decode_tokens(toks):\n valid_toks = toks[: np.argmax(toks == eos_id) + 1].astype(np.int32)\n return encoder.detokenize(valid_toks).numpy().decode(\"utf-8\")\n\n def encode_strings(strs, max_len):\n tokenized_batch = np.zeros((len(strs), max_len), np.int32)\n for i, s in enumerate(strs):\n toks = encoder.tokenize(s).numpy()\n # Remove EOS token in prompt.\n tokenized_batch[i, : toks.shape[0] - 1] = toks[:-1]\n return tokenized_batch\n\n tokenized_prompts = encode_strings(\n [config.prompts], config.max_predict_length\n )\n\n logging.info(\"Initializing model, optimizer, and step functions.\")\n # Build Model and Optimizer\n # ---------------------------------------------------------------------------\n train_config = models.TransformerConfig(\n vocab_size=vocab_size,\n output_vocab_size=vocab_size,\n logits_via_embedding=config.logits_via_embedding,\n dtype=jnp.bfloat16 if config.use_bfloat16 else jnp.float32,\n emb_dim=config.emb_dim,\n num_heads=config.num_heads,\n num_layers=config.num_layers,\n qkv_dim=config.qkv_dim,\n mlp_dim=config.mlp_dim,\n max_len=max(config.max_target_length, config.max_eval_target_length),\n dropout_rate=config.dropout_rate,\n attention_dropout_rate=config.attention_dropout_rate,\n deterministic=False,\n decode=False,\n kernel_init=nn.initializers.xavier_uniform(),\n bias_init=nn.initializers.normal(stddev=1e-6),\n )\n eval_config = train_config.replace(deterministic=True)\n predict_config = train_config.replace(deterministic=True, decode=True)\n\n # Mesh definition\n devices_array = utils.create_device_mesh(config)\n mesh = Mesh(devices_array, config.mesh_axes)\n\n start_step = 0\n rng = jax.random.PRNGKey(config.seed)\n rng, init_rng = jax.random.split(rng)\n rng, inference_rng = random.split(rng)\n\n m = models.TransformerLM(eval_config)\n\n learning_rate_fn = create_learning_rate_schedule(\n learning_rate=config.learning_rate, warmup_steps=config.warmup_steps\n )\n\n optimizer = optax.adamw(\n learning_rate_fn,\n b1=0.9,\n b2=0.98,\n eps=1e-9,\n weight_decay=config.weight_decay,\n )\n\n state, state_mesh_annotations = utils.setup_initial_state(\n m, optimizer, config, init_rng, mesh\n )\n data_sharding = NamedSharding(mesh, P(config.data_sharding))\n\n if config.restore_checkpoints:\n # Restore unreplicated optimizer + model state from last checkpoint.\n state = checkpoints.restore_checkpoint(workdir, state)\n # Grab last step.\n start_step = int(state.step)\n\n writer = metric_writers.create_default_writer(\n workdir, just_logging=jax.process_index() > 0\n )\n if start_step == 0:\n writer.write_hparams(dict(config))\n\n # compile multidevice versions of train/eval/predict step fn.\n jit_train_step = jax.jit(\n train_step,\n in_shardings=(\n state_mesh_annotations,\n data_sharding,\n None,\n ), # type: ignore\n out_shardings=(state_mesh_annotations, None), # type: ignore\n static_argnums=(2, 3, 4),\n donate_argnums=0,\n )\n\n jit_eval_step = jax.jit(\n eval_step,\n in_shardings=(\n state_mesh_annotations.params,\n data_sharding,\n ), # type: ignore\n out_shardings=None, # type: ignore\n static_argnums=(2, 3),\n )\n\n # Since the inputs and rngkey args for predict_step will be batched,\n # we must vmap them, otherwise the global arrays will be seen in each device\n jit_pred_step = jax.jit(\n jax.vmap(\n predict_step,\n in_axes=(\n 0,\n jax.tree_util.tree_map(lambda x: None, state.params),\n 0,\n None,\n None,\n jax.tree_util.tree_map(lambda x: None, predict_config),\n None,\n None,\n ),\n ),\n in_shardings=(\n data_sharding,\n state_mesh_annotations.params,\n data_sharding,\n ), # type: ignore\n out_shardings=data_sharding, # type: ignore\n static_argnums=(3, 4, 5, 6, 7),\n )\n\n # Main Train Loop\n # ---------------------------------------------------------------------------\n\n # We init the first set of dropout PRNG keys, but update it afterwards inside\n # the main pmap\"d training update for performance.\n dropout_rngs = rng\n\n logging.info(\"Starting training loop.\")\n hooks = []\n report_progress = periodic_actions.ReportProgress(\n num_train_steps=config.num_train_steps, writer=writer\n )\n if jax.process_index() == 0:\n hooks += [\n report_progress,\n periodic_actions.Profile(logdir=workdir, num_profile_steps=5),\n ]\n train_metrics = []\n with metric_writers.ensure_flushes(writer):\n for step in range(start_step, config.num_train_steps):\n is_last_step = step == config.num_train_steps - 1\n\n # Shard data to devices and do a training step.\n with jax.profiler.StepTraceAnnotation(\"train\", step_num=step):\n batch = next(train_iter)\n batch = jax.tree_util.tree_map(lambda x: jnp.array(x), batch)\n state, metrics = jit_train_step(\n state, batch, train_config, learning_rate_fn, 0.0, dropout_rngs\n )\n train_metrics.append(metrics)\n\n # Quick indication that training is happening.\n logging.log_first_n(logging.INFO, \"Finished training step %d.\", 5, step)\n for h in hooks:\n h(step)\n\n # Periodic metric handling.\n if step % config.eval_every_steps == 0 or is_last_step:\n with report_progress.timed(\"training_metrics\"):\n logging.info(\"Gathering training metrics.\")\n train_metrics = common_utils.stack_forest(train_metrics)\n lr = train_metrics.pop(\"learning_rate\").mean()\n metrics_sums = jax.tree_util.tree_map(jnp.sum, train_metrics)\n denominator = metrics_sums.pop(\"denominator\")\n summary = jax.tree_util.tree_map(\n lambda x: x / denominator, metrics_sums\n ) # pylint: disable=cell-var-from-loop\n summary[\"learning_rate\"] = lr\n summary[\"perplexity\"] = jnp.clip(\n jnp.exp(summary[\"loss\"]), max=1.0e4\n )\n summary = {\"train_\" + k: v for k, v in summary.items()}\n writer.write_scalars(step, summary)\n train_metrics = []\n\n with report_progress.timed(\"eval\"):\n eval_results = evaluate(\n jit_eval_step=jit_eval_step,\n params=state.params,\n eval_ds=eval_ds,\n num_eval_steps=config.num_eval_steps,\n config=eval_config,\n )\n # (clipped) perplexity after averaging log-perplexitie\n eval_results[\"perplexity\"] = jnp.clip(\n jnp.exp(eval_results[\"loss\"]), max=1.0e4\n )\n writer.write_scalars(\n step, {\"eval_\" + k: v for k, v in eval_results.items()}\n )\n\n with report_progress.timed(\"generate_text\"):\n exemplars = generate_prediction(\n jit_pred_step=jit_pred_step,\n params=state.params,\n tokenized_prompts=tokenized_prompts,\n eos_id=eos_id,\n inference_rng=inference_rng,\n decode_tokens=decode_tokens,\n config=config,\n predict_config=predict_config,\n )\n writer.write_texts(step, {\"samples\": exemplars})\n\n # Save a checkpoint on one host after every checkpoint_freq steps.\n save_checkpoint = (\n step % config.checkpoint_every_steps == 0 or is_last_step\n )\n if config.save_checkpoints and save_checkpoint:\n logging.info(\"Saving checkpoint step %d.\", step)\n with report_progress.timed(\"checkpoint\"):\n checkpoints.save_checkpoint_multiprocess(workdir, state, step)\n" }, { "path": "examples/lm1b/train_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport pathlib\nimport sys\nimport tempfile\n\nfrom absl import logging\nfrom absl.testing import absltest\nimport jax\nimport tensorflow as tf\nimport tensorflow_datasets as tfds\n\nfrom configs import default\nimport train\n\n\njax.config.update('jax_disable_most_optimizations', True)\n\n\nclass TrainTest(absltest.TestCase):\n \"\"\"Test cases for LM library.\"\"\"\n\n def setUp(self):\n super().setUp()\n if sys.version_info >= (3, 13):\n self.skipTest('Test (and tensorflow-text) does not suport Python 3.13+')\n tf.config.experimental.set_visible_devices([], 'GPU')\n\n def test_train_and_evaluate(self):\n config = default.get_config()\n config.max_corpus_chars = 1000\n config.vocab_size = 32\n config.per_device_batch_size = 2\n config.num_train_steps = 1\n config.num_eval_steps = 1\n config.num_predict_steps = 1\n\n config.num_layers = 1\n config.qkv_dim = 128\n config.emb_dim = 128\n config.mlp_dim = 512\n config.num_heads = 2\n\n config.max_target_length = 32\n config.max_eval_target_length = 32\n config.max_predict_length = 32\n\n workdir = tempfile.mkdtemp()\n\n # Go two directories up to the root of the flax directory.\n flax_root_dir = pathlib.Path(__file__).parents[2]\n data_dir = str(flax_root_dir) + '/.tfds/metadata' # pylint: disable=unused-variable\n\n with tfds.testing.mock_data(num_examples=128, data_dir=data_dir):\n train.train_and_evaluate(config, workdir)\n logging.info('workdir content: %s', tf.io.gfile.listdir(workdir))\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "examples/lm1b/utils.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# Copied over from MaxText (https://github.com/google/maxtext/blob/main/MaxText/max_utils.py).\n\nimport functools\nimport logging\n\nimport numpy as np\nimport flax.linen as nn\nfrom flax.linen import partitioning as nn_partitioning\nfrom flax.training import train_state\nimport jax\nimport jax.numpy as jnp\nfrom jax.experimental import mesh_utils\n\n\n# Mesh utils.\n# -----------------------------------------------------------------------------\n\n\ndef create_device_mesh(config):\n \"\"\"Creates a device mesh with each slice in its own data parallel group. If there is only one slice, uses two replicas.\"\"\"\n devices = jax.devices()\n num_devices = len(devices)\n try:\n num_slices = 1 + max([d.slice_index for d in devices])\n except:\n num_slices = 1\n num_devices_per_slice = num_devices // num_slices\n logging.info(f\"Devices: {devices}\")\n logging.info(f\"Number of devices: {num_devices}\")\n\n multi_slice_env = hasattr(jax.devices()[0], \"slice_index\")\n\n dcn_parallelism = [\n config.dcn_data_parallelism,\n config.dcn_fsdp_parallelism,\n config.dcn_tensor_parallelism,\n ]\n ici_parallelism = [\n config.ici_data_parallelism,\n config.ici_fsdp_parallelism,\n config.ici_tensor_parallelism,\n ]\n\n # Find possible unspecified parallelisms\n dcn_parallelism = fill_unspecified_mesh_axes(\n dcn_parallelism, num_slices, \"DCN\"\n )\n ici_parallelism = fill_unspecified_mesh_axes(\n ici_parallelism, num_devices_per_slice, \"ICI\"\n )\n\n if multi_slice_env:\n mesh = mesh_utils.create_hybrid_device_mesh(\n ici_parallelism, dcn_parallelism\n )\n else:\n mesh = mesh_utils.create_device_mesh(ici_parallelism)\n\n logging.info(f\"Decided on mesh: {mesh}\")\n logging.info(f\"Mesh shape: {mesh.shape}\")\n\n return mesh\n\n\ndef fill_unspecified_mesh_axes(\n parallelism_vals, target_product, parallelism_type\n):\n \"\"\"Evaluates unspecified DCN/ICI parallelism values\"\"\"\n if -1 in parallelism_vals:\n assert parallelism_vals.count(-1) == 1, (\n f\"Found unspecified values (-1) for more than one {parallelism_type} \"\n \" parallelism axis. At most one axis can be unspecified.\"\n )\n\n determined_val = target_product / np.prod(parallelism_vals) * -1\n\n assert determined_val >= 1 and determined_val.is_integer, (\n \"Unspecified value unable to be determined with the given \"\n f\" {parallelism_type} parallelism values\"\n )\n\n parallelism_vals[parallelism_vals.index(-1)] = int(determined_val)\n\n target_type = \"slices\" if parallelism_type == \"DCN\" else \"devices per slice\"\n\n assert np.prod(parallelism_vals) == target_product, (\n f\"Number of {target_type} {target_product} does not match the product\"\n f\" of the {parallelism_type} parallelism {np.prod(parallelism_vals)}\"\n )\n\n return parallelism_vals\n\n\n# State initialization utils.\n# -----------------------------------------------------------------------------\n\n\ndef unbox_logicallypartioned_trainstate(\n boxed_train_state: train_state.TrainState,\n):\n \"\"\"Unboxes the flax.LogicallyPartitioned pieces in a train state.\n\n Args:\n boxed_train_state: a train state that includes LogicallyPartitioned\n leaves.\n Returns:\n a TrainState where all all LogicallyPartitioned leaves have been unboxed.\n \"\"\"\n return jax.tree_util.tree_map(\n lambda x: x.unbox() if isinstance(x, nn.spmd.LogicallyPartitioned) else x,\n boxed_train_state,\n is_leaf=lambda k: isinstance(k, nn.spmd.LogicallyPartitioned),\n )\n\n\ndef init_train_state(model, tx, config, key):\n \"\"\"\n We pass in \"static\" objects like model, tx, config as JAX compares them by\n object hash, and instantiating them inside causes pjit top-level annotations\n to fail to match as pytree prefixes if we re-instantiate.\n\n Args: model, tx, config, key\n \"\"\"\n input_shape = (config.per_device_batch_size, config.max_target_length)\n initial_variables = jax.jit(model.init)(\n key, jnp.ones(input_shape, jnp.float32)\n )\n\n state = train_state.TrainState.create(\n apply_fn=model.apply, params=initial_variables[\"params\"], tx=tx\n )\n return state\n\n\ndef setup_initial_state(model, tx, config, rng, mesh):\n \"\"\"We initialize the model and optimizer state, and optionally load from a\n checkpoint as necessary.\n\n Args:\n model: the flax model to initialize\n tx: the optax.GradientTransformation\n config: config object\n rng: jax.prng key\n mesh: jax.devices() mesh\n\n Returns:\n state: the initialized train state\n state_mesh_annotations: the mesh annotations for the train state\n \"\"\"\n init_train_state_partial = functools.partial(\n init_train_state, model, tx, config\n )\n abstract_state = jax.eval_shape(init_train_state_partial, rng)\n state_logical_annotations = nn.get_partition_spec(abstract_state)\n\n # Initialization\n with jax.set_mesh(mesh), nn_partitioning.axis_rules(config.logical_axis_rules):\n state_mesh_annotations = nn.logical_to_mesh_sharding(\n state_logical_annotations, mesh, config.logical_axis_rules\n )\n state = jax.jit(\n init_train_state_partial,\n in_shardings=None, # type: ignore\n out_shardings=state_mesh_annotations,\n )(rng)\n\n state = unbox_logicallypartioned_trainstate(state)\n return state, state_mesh_annotations\n" }, { "path": "examples/mnist/README.md", "content": "## MNIST classification\n\nTrains a simple convolutional network on the MNIST dataset.\n\nYou can run this code and even modify it directly in Google Colab, no\ninstallation required:\n\nhttps://colab.research.google.com/github/google/flax/blob/main/examples/mnist/mnist.ipynb\n\n### Requirements\n* TensorFlow dataset `mnist` will be downloaded and prepared automatically, if necessary\n\n### Example output\n\n| Name | Epochs | Walltime | Top-1 accuracy | Metrics | Workdir |\n| :------ | -----: | :------- | :------------- | :---------- | :---------------------------------------- |\n| default | 10 | 7.7m | 99.17% | [tfhub.dev] | [gs://flax_public/examples/mnist/default] |\n\n[tfhub.dev]: https://tensorboard.dev/experiment/1G9SvrW5RQyojRtMKNmMuQ/#scalars&_smoothingWeight=0®exInput=default\n[gs://flax_public/examples/mnist/default]: https://console.cloud.google.com/storage/browser/flax_public/examples/mnist/default\n\n```\nI1009 17:56:42.674334 3280981 train.py:175] epoch: 10, train_loss: 0.0073, train_accuracy: 99.75, test_loss: 0.0294, test_accuracy: 99.25\n```\n\n### How to run\n\n`python main.py --workdir=/tmp/mnist --config=configs/default.py`\n\n#### Overriding Hyperparameter configurations\n\nMNIST example allows specifying a hyperparameter configuration by the means of\nsetting `--config` flag. Configuration flag is defined using\n[config_flags](https://github.com/google/ml_collections/tree/master#config-flags).\n`config_flags` allows overriding configuration fields. This can be done as\nfollows:\n\n```shell\npython main.py \\\n--workdir=/tmp/mnist --config=configs/default.py \\\n--config.learning_rate=0.05 --config.num_epochs=5\n```\n" }, { "path": "examples/mnist/configs/default.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Default Hyperparameter configuration.\"\"\"\n\nimport ml_collections\n\n\ndef get_config():\n \"\"\"Get the default hyperparameter configuration.\"\"\"\n config = ml_collections.ConfigDict()\n\n config.learning_rate = 0.1\n config.momentum = 0.9\n config.batch_size = 128\n config.num_epochs = 10\n return config\n\n\ndef metrics():\n return []\n" }, { "path": "examples/mnist/main.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Main file for running the MNIST example.\n\nThis file is intentionally kept short. The majority of logic is in libraries\nthan can be easily tested and imported in Colab.\n\"\"\"\n\nfrom absl import app\nfrom absl import flags\nfrom absl import logging\nfrom clu import platform\nimport jax\nfrom ml_collections import config_flags\nimport tensorflow as tf\n\nimport train # pylint: disable=g-bad-import-order\n\n\nFLAGS = flags.FLAGS\n\nflags.DEFINE_string('workdir', None, 'Directory to store model data.')\nconfig_flags.DEFINE_config_file(\n 'config',\n None,\n 'File path to the training hyperparameter configuration.',\n lock_config=True,\n)\n\n\ndef main(argv):\n if len(argv) > 1:\n raise app.UsageError('Too many command-line arguments.')\n\n # Hide any GPUs from TensorFlow. Otherwise TF might reserve memory and make\n # it unavailable to JAX.\n tf.config.experimental.set_visible_devices([], 'GPU')\n\n logging.info('JAX process: %d / %d', jax.process_index(), jax.process_count())\n logging.info('JAX local devices: %r', jax.local_devices())\n\n # Add a note so that we can tell which task is which JAX host.\n # (Depending on the platform task 0 is not guaranteed to be host 0)\n platform.work_unit().set_task_status(\n f'process_index: {jax.process_index()}, '\n f'process_count: {jax.process_count()}'\n )\n platform.work_unit().create_artifact(\n platform.ArtifactType.DIRECTORY, FLAGS.workdir, 'workdir'\n )\n\n train.train_and_evaluate(FLAGS.config, FLAGS.workdir)\n\n\nif __name__ == '__main__':\n flags.mark_flags_as_required(['config', 'workdir'])\n app.run(main)\n" }, { "path": "examples/mnist/mnist.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"# Flax MNIST Example\\n\",\n \"\\n\",\n \"\\\"Open\\n\",\n \"\\n\",\n \"Demonstration notebook for\\n\",\n \"https://github.com/google/flax/tree/main/examples/mnist\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"The **Flax Notebook Workflow**:\\n\",\n \"\\n\",\n \"1. Run the entire notebook end-to-end and check out the outputs.\\n\",\n \" - This will open Python files in the right-hand editor!\\n\",\n \" - You'll be able to interactively explore metrics in TensorBoard.\\n\",\n \"2. Change `config` and train for different hyperparameters. Check out the\\n\",\n \" updated TensorBoard plots.\\n\",\n \"3. Update the code in `train.py`. Thanks to `%autoreload`, any changes you\\n\",\n \" make in the file will automatically appear in the notebook. Some ideas to\\n\",\n \" get you started:\\n\",\n \" - Change the model.\\n\",\n \" - Log some per-batch metrics during training.\\n\",\n \" - Add new hyperparameters to `configs/default.py` and use them in\\n\",\n \" `train.py`.\\n\",\n \"4. At any time, feel free to paste code from `train.py` into the notebook\\n\",\n \" and modify it directly there!\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Setup\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 1,\n \"metadata\": {\n \"outputId\": \"8520b2f8-2b9d-4216-ba1f-d96175455bbc\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"\\u001b[?25l\\r\\n\",\n \"\\u001b[K |███▊ | 10kB 21.5MB/s eta 0:00:01\\r\\n\",\n \"\\u001b[K |███████▍ | 20kB 12.7MB/s eta 0:00:01\\r\\n\",\n \"\\u001b[K |███████████ | 30kB 9.4MB/s eta 0:00:01\\r\\n\",\n \"\\u001b[K |██████████████▉ | 40kB 8.4MB/s eta 0:00:01\\r\\n\",\n \"\\u001b[K |██████████████████▌ | 51kB 5.2MB/s eta 0:00:01\\r\\n\",\n \"\\u001b[K |██████████████████████▏ | 61kB 5.2MB/s eta 0:00:01\\r\\n\",\n \"\\u001b[K |█████████████████████████▉ | 71kB 5.7MB/s eta 0:00:01\\r\\n\",\n \"\\u001b[K |█████████████████████████████▋ | 81kB 6.3MB/s eta 0:00:01\\r\\n\",\n \"\\u001b[K |████████████████████████████████| 92kB 4.4MB/s \\n\",\n \"\\u001b[K |████████████████████████████████| 634kB 8.5MB/s \\n\",\n \"\\u001b[K |████████████████████████████████| 102kB 6.0MB/s \\n\",\n \"\\u001b[K |████████████████████████████████| 61kB 6.3MB/s \\n\",\n \"\\u001b[?25h Building wheel for flax (setup.py) ... \\u001b[?25l\\u001b[?25hdone\\n\",\n \" Building wheel for jax (setup.py) ... \\u001b[?25l\\u001b[?25hdone\\n\"\n ]\n }\n ],\n \"source\": [\n \"# Install ml-collections & latest Flax version from Github.\\n\",\n \"!pip install -q ml-collections git+https://github.com/google/flax\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 2,\n \"metadata\": {\n \"tags\": []\n },\n \"outputs\": [],\n \"source\": [\n \"example_directory = 'examples/mnist'\\n\",\n \"editor_relpaths = ('configs/default.py', 'train.py')\\n\",\n \"\\n\",\n \"repo, branch = 'https://github.com/google/flax', 'main'\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 3,\n \"metadata\": {\n \"cellView\": \"form\",\n \"outputId\": \"2dfbdfa6-d213-4b5b-dc82-ee1765705255\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"Cloning into 'flaxrepo'...\\n\",\n \"remote: Enumerating objects: 330, done.\\u001b[K\\n\",\n \"remote: Counting objects: 100% (330/330), done.\\u001b[K\\n\",\n \"remote: Compressing objects: 100% (298/298), done.\\u001b[K\\n\",\n \"remote: Total 330 (delta 58), reused 126 (delta 14), pack-reused 0\\u001b[K\\n\",\n \"Receiving objects: 100% (330/330), 1.81 MiB | 6.59 MiB/s, done.\\n\",\n \"Resolving deltas: 100% (58/58), done.\\n\"\n ]\n },\n {\n \"data\": {\n \"text/html\": [\n \"

WARNING : Editing in VM - changes lost after reboot!!

\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {\n \"tags\": []\n },\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"application/javascript\": [\n \"\\n\",\n \" ((filepath) => {{\\n\",\n \" if (!google.colab.kernel.accessAllowed) {{\\n\",\n \" return;\\n\",\n \" }}\\n\",\n \" google.colab.files.view(filepath);\\n\",\n \" }})(\\\"/content/examples/mnist/configs/default.py\\\")\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {\n \"tags\": []\n },\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"application/javascript\": [\n \"\\n\",\n \" ((filepath) => {{\\n\",\n \" if (!google.colab.kernel.accessAllowed) {{\\n\",\n \" return;\\n\",\n \" }}\\n\",\n \" google.colab.files.view(filepath);\\n\",\n \" }})(\\\"/content/examples/mnist/train.py\\\")\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {\n \"tags\": []\n },\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"# (If you run this code in Jupyter[lab], then you're already in the\\n\",\n \"# example directory and nothing needs to be done.)\\n\",\n \"\\n\",\n \"#@markdown **Fetch newest Flax, copy example code**\\n\",\n \"#@markdown\\n\",\n \"#@markdown **If you select no** below, then the files will be stored on the\\n\",\n \"#@markdown *ephemeral* Colab VM. **After some time of inactivity, this VM will\\n\",\n \"#@markdown be restarted an any changes are lost**.\\n\",\n \"#@markdown\\n\",\n \"#@markdown **If you select yes** below, then you will be asked for your\\n\",\n \"#@markdown credentials to mount your personal Google Drive. In this case, all\\n\",\n \"#@markdown changes you make will be *persisted*, and even if you re-run the\\n\",\n \"#@markdown Colab later on, the files will still be the same (you can of course\\n\",\n \"#@markdown remove directories inside your Drive's `flax/` root if you want to\\n\",\n \"#@markdown manually revert these files).\\n\",\n \"\\n\",\n \"if 'google.colab' in str(get_ipython()):\\n\",\n \" import os\\n\",\n \" os.chdir('/content')\\n\",\n \" # Download Flax repo from Github.\\n\",\n \" if not os.path.isdir('flaxrepo'):\\n\",\n \" !git clone --depth=1 -b $branch $repo flaxrepo\\n\",\n \" # Copy example files & change directory.\\n\",\n \" mount_gdrive = 'no' #@param ['yes', 'no']\\n\",\n \" if mount_gdrive == 'yes':\\n\",\n \" DISCLAIMER = 'Note : Editing in your Google Drive, changes will persist.'\\n\",\n \" from google.colab import drive\\n\",\n \" drive.mount('/content/gdrive')\\n\",\n \" example_root_path = f'/content/gdrive/My Drive/flax/{example_directory}'\\n\",\n \" else:\\n\",\n \" DISCLAIMER = 'WARNING : Editing in VM - changes lost after reboot!!'\\n\",\n \" example_root_path = f'/content/{example_directory}'\\n\",\n \" from IPython import display\\n\",\n \" display.display(display.HTML(\\n\",\n \" f'

{DISCLAIMER}

'))\\n\",\n \" if not os.path.isdir(example_root_path):\\n\",\n \" os.makedirs(example_root_path)\\n\",\n \" !cp -r flaxrepo/$example_directory/* \\\"$example_root_path\\\"\\n\",\n \" os.chdir(example_root_path)\\n\",\n \" from google.colab import files\\n\",\n \" for relpath in editor_relpaths:\\n\",\n \" s = open(f'{example_root_path}/{relpath}').read()\\n\",\n \" open(f'{example_root_path}/{relpath}', 'w').write(\\n\",\n \" f'## {DISCLAIMER}\\\\n' + '#' * (len(DISCLAIMER) + 3) + '\\\\n\\\\n' + s)\\n\",\n \" files.view(f'{example_root_path}/{relpath}')\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 4,\n \"metadata\": {\n \"outputId\": \"e9061488-ac3e-4d23-f24f-06e1988e7541\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"/content/examples/mnist\\n\"\n ]\n }\n ],\n \"source\": [\n \"# Note : In Colab, above cell changed the working directory.\\n\",\n \"!pwd\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Imports / Helpers\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 5,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"from absl import logging\\n\",\n \"import flax\\n\",\n \"import jax.numpy as jnp\\n\",\n \"from matplotlib import pyplot as plt\\n\",\n \"import numpy as np\\n\",\n \"import tensorflow_datasets as tfds\\n\",\n \"\\n\",\n \"logging.set_verbosity(logging.INFO)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 6,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# Helper functions for images.\\n\",\n \"\\n\",\n \"def show_img(img, ax=None, title=None):\\n\",\n \" \\\"\\\"\\\"Shows a single image.\\\"\\\"\\\"\\n\",\n \" if ax is None:\\n\",\n \" ax = plt.gca()\\n\",\n \" ax.imshow(img[..., 0], cmap='gray')\\n\",\n \" ax.set_xticks([])\\n\",\n \" ax.set_yticks([])\\n\",\n \" if title:\\n\",\n \" ax.set_title(title)\\n\",\n \"\\n\",\n \"def show_img_grid(imgs, titles):\\n\",\n \" \\\"\\\"\\\"Shows a grid of images.\\\"\\\"\\\"\\n\",\n \" n = int(np.ceil(len(imgs)**.5))\\n\",\n \" _, axs = plt.subplots(n, n, figsize=(3 * n, 3 * n))\\n\",\n \" for i, (img, title) in enumerate(zip(imgs, titles)):\\n\",\n \" show_img(img, axs[i // n][i % n], title)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 7,\n \"metadata\": {\n \"tags\": []\n },\n \"outputs\": [],\n \"source\": [\n \"# Local imports from current directory - auto reload.\\n\",\n \"# Any changes you make to train.py will appear automatically.\\n\",\n \"%load_ext autoreload\\n\",\n \"%autoreload 2\\n\",\n \"import train\\n\",\n \"from configs import default as config_lib\\n\",\n \"config = config_lib.get_config()\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Dataset\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 8,\n \"metadata\": {\n \"outputId\": \"bb4525f4-8ca4-4e9d-d1cc-48a3e0533645\",\n \"tags\": []\n },\n \"outputs\": [\n {\n \"name\": \"stderr\",\n \"output_type\": \"stream\",\n \"text\": [\n \"INFO:absl:Load pre-computed DatasetInfo (eg: splits, num examples,...) from GCS: mnist/3.0.1\\n\",\n \"INFO:absl:Load dataset info from /tmp/tmpqyu1t56xtfds\\n\",\n \"INFO:absl:Field info.citation from disk and from code do not match. Keeping the one from code.\\n\",\n \"INFO:absl:Generating dataset mnist (/root/tensorflow_datasets/mnist/3.0.1)\\n\"\n ]\n },\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"\\u001b[1mDownloading and preparing dataset mnist/3.0.1 (download: 11.06 MiB, generated: 21.00 MiB, total: 32.06 MiB) to /root/tensorflow_datasets/mnist/3.0.1...\\u001b[0m\\n\"\n ]\n },\n {\n \"name\": \"stderr\",\n \"output_type\": \"stream\",\n \"text\": [\n \"WARNING:absl:Dataset mnist is hosted on GCS. It will automatically be downloaded to your\\n\",\n \"local data directory. If you'd instead prefer to read directly from our public\\n\",\n \"GCS bucket (recommended if you're running on GCP), you can instead pass\\n\",\n \"`try_gcs=True` to `tfds.load` or set `data_dir=gs://tfds-data/datasets`.\\n\",\n \"\\n\"\n ]\n },\n {\n \"data\": {\n \"application/vnd.jupyter.widget-view+json\": {\n \"model_id\": \"1c176027dbbf459b8ed946ccc58de845\",\n \"version_major\": 2,\n \"version_minor\": 0\n },\n \"text/plain\": [\n \"HBox(children=(FloatProgress(value=0.0, description='Dl Completed...', max=4.0, style=ProgressStyle(descriptio…\"\n ]\n },\n \"metadata\": {\n \"tags\": []\n },\n \"output_type\": \"display_data\"\n },\n {\n \"name\": \"stderr\",\n \"output_type\": \"stream\",\n \"text\": [\n \"INFO:absl:Load dataset info from /root/tensorflow_datasets/mnist/3.0.1.incompleteJ6TJES\\n\",\n \"INFO:absl:Field info.citation from disk and from code do not match. Keeping the one from code.\\n\",\n \"INFO:absl:Constructing tf.data.Dataset for split train, from /root/tensorflow_datasets/mnist/3.0.1\\n\"\n ]\n },\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"\\n\",\n \"\\n\",\n \"\\u001b[1mDataset mnist downloaded and prepared to /root/tensorflow_datasets/mnist/3.0.1. Subsequent calls will reuse this data.\\u001b[0m\\n\"\n ]\n },\n {\n \"name\": \"stderr\",\n \"output_type\": \"stream\",\n \"text\": [\n \"INFO:absl:Constructing tf.data.Dataset for split test, from /root/tensorflow_datasets/mnist/3.0.1\\n\"\n ]\n }\n ],\n \"source\": [\n \"# Get datasets as dict of JAX arrays.\\n\",\n \"train_ds, test_ds = train.get_datasets()\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 9,\n \"metadata\": {\n \"outputId\": \"89de05b0-aede-414f-cf43-5e7c71871140\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"image/png\": \"iVBORw0KGgoAAAANSUhEUgAAA1MAAANRCAYAAAAGcOaXAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjIsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+WH4yJAAAgAElEQVR4nOzdebzOdf7/8dfLcpBdSTE4FTIVGdF8lVC0qGnKpMVSGi2GNMUwE9+0KISMqUgLRX1Li1TTIC1M2qcUqdMyKrJn32V7//44x2/OeL2PPud9Xde5tsf9drtut+PZ5/P+vM7p7bhePj6vo845AQAAAAAUT6lkFwAAAAAA6YhmCgAAAAAC0EwBAAAAQACaKQAAAAAIQDMFAAAAAAFopgAAAAAgQFY1U6q6RFU7RDjOqWqDwGsEnwscjD2LdMOeRbphzyLdsGdTS1Y1U6lOVXNU9UtVXZ7sWoBDUdUzVXWuqm5W1SXJrgf4OZpvpKquL3iNVFVNdl1AUVR1lqpuK/TaraqLkl0XUBRVLaeqD6nqGlXdoKqvqGqdZNeVaDRTqWWgiKxNdhFABNtF5DHJ37NAOrheRC4WkZNFpKmIXCgivZJaEXAIzrmOzrlKB14i8p6IPJ/suoBDuElEWkn+99jaIrJRRB5IakUlICubKVU9VVXfV9VNqrpKVcepas5Bh52vqt+p6jpVHa2qpQqd37PgDtJGVZ2tqvXjUNMxItJdREbEuhYyT6rtWefcv5xzT4rId7Gsg8yVantWRHqIyBjn3HLn3AoRGSMiV8e4JjJICu7ZwrXlisgZIvJEvNZE+kvBPXuMiMx2zq1xzu0SkWdF5MQY10x5WdlMicg+EeknIkdIfgfdXkT6HHRMJxFpISLNReQiEekpIqKqF4nIYBH5nYjUFJG3RWSq7yKqekvBBve+Djr8gYJ1d8bjE0TGScU9CxxKqu3ZE0VkYaFfL5Qs+EMexZJqe7awq0Tkbefcklg+QWScVNuzk0TkdFWtraqHiUg3EZkVp881dTnnsuYlIktEpIMnv1lEXiz0ayci5xX6dR8RebPg41kick2h/1ZKRHaISP1C5zYoZl2dRGRWwcftRGR5sr9WvFLjlap7ttBaHURkSbK/TrxS55Wqe1by33Q0LvTrhgXraLK/ZryS+0rVPXtQLYtF5Opkf614pcYrVfesiFQVkWcKzt0rIp+KSI1kf70S/crKO1Oq2khV/6Gqq1V1i4gMl/yuvrBlhT5eKvn/9lNEpL6I3FeoG98gIioiQQ/YqWpFERklIn8MOR/ZIZX2LBBFCu7ZbSJSpdCvq4jINlfwDgBIwT17oK7WInKUiEyLdS1klhTcs+NFpJyIHC4iFUVkumTBnamsbKZEZIKIfCUiDZ1zVST/NufBU53qFvq4noisLPh4mYj0cs5VK/Sq4Jx77+CLqOpg/e9JPP/1KjisoYjkisjbqrpa8jfe0QW/MXLj9Pki/aXSngWiSLU9+4XkD5844OSCDDgg1fbsAT1EZLpzju/BOFiq7dlmIjLZObfBOfeT5D/CcqqqHtzgZZRsbaYqi8gWEdmmqo1FpLfnmIGqWl1V60r+dJJnC/KHRGSQqp4oIqKqVVX1Ut9FnHPDXaFJPAe/Cg77XPI3erOC17Uisqbg42W+dZGVUmnPiqqWUtXyIlI2/5daXu1Dr8huKbVnJf/B/f6qWkdVa4vIn0Rkclw+U2SKVNuzoqoVROQyYa/CL9X27EciclXBWmUl/58VrnTOrYvPp5uasrWZGiAiXUVkq4g8Kv/ZWIW9LCLzRWSBiMyQ/IfqxDn3ooiMFJFnCm6pfi4iHUMLcc7tdc6tPvCS/Nus+wt+vS90XWSclNmzBdpI/rCUmZL/N107ReS1GNdEZkm1PfuwiLwiIosK1ptRkAEHpNqeFckf579JRObGYS1knlTbswNEZJeI/Fvyf9TP+ZI/FyCjKf9cHAAAAACKL1vvTAEAAABATGimAAAAACAAzRQAAAAABKCZAgAAAIAANFMAAAAAEKBMcQ5WVUb/oViccwf/8LgSxZ5FcbFnkYbWOedqJuvi7FkEYM8i3RS5Z7kzBQBAelua7AKAYmLPIt0UuWdppgAAAAAgAM0UAAAAAASgmQIAAACAADRTAAAAABCAZgoAAAAAAtBMAQAAAEAAmikAAAAACEAzBQAAAAABaKYAAAAAIADNFAAAAAAEoJkCAAAAgAA0UwAAAAAQoEyyCwAAIFSjRo1MNn78eJO1b9/eZJMnTzZZnz59TLZr166w4gAAGY87UwAAAAAQgGYKAAAAAALQTAEAAABAAJopAAAAAAjAAAoAQNo67bTTTHbWWWeZzDlnsh49ephs3759JrvhhhtMtnv37qglAkBWqFy5ssl83z+HDx9uslWrVpnshBNOMNnmzZsDq0sc7kwBAAAAQACaKQAAAAAIQDMFAAAAAAFopgAAAAAgAAMoUkDnzp1N9txzz5msV69eJnv00UcTUhNwQIUKFUz24IMPmuywww4zWZcuXUy2f//++BSGrHLeeed587/97W9xvU7Pnj1NlpeXZ7KxY8fG9boAEA/HHnusyXyDIS655BKTlS9fPvKavmzhwoUmu+qqq0zmGwh09NFHR6qHARQAAAAAkCFopgAAAAAgAM0UAAAAAASgmQIAAACAAAygSAFdu3Y1me/hvBo1apREOchiqmqyhx9+2GTdu3ePtN6IESNMtmDBguIXhqziG2YydOhQ77G+B6vjbciQISZjAEV6mzt3rsnatWtnspEjR5rslltuSURJwP9Xrlw5kx1zzDEmmzBhgsl+9atfmaxKlSom873PLA7f+4WTTz45pjXTFXemAAAAACAAzRQAAAAABKCZAgAAAIAANFMAAAAAEIABFCWsfv36JuvYsaPJ5s+fb7Knn346ITUBB5xwwgkmizpsYsuWLSZbv359zDUh+7zwwgsma9GihffYqA9R+wafNGvWLNK5ZcrwR2W68D0Uf/zxx5vM95D+/v37TXbTTTeZbN++fSabPn26yYram19//bU3P9hZZ51lsmOPPdZkS5YsMdnMmTNNtmfPnkjXRcny/T997rnnTObbs1G9++67Jvv2229NNmPGDO/5mzZtMtns2bOD6/FZsWKFyXbt2hXXayQKd6YAAAAAIADNFAAAAAAEoJkCAAAAgAA0UwAAAAAQIOOfqvU9jOoT60+CjuqPf/yjyXJyckz23XffmWzZsmUJqQk44NJLLw0+94cffjAZexY/59prrzVZu3btYlrT9/2zbdu2JvMNuujQoYPJfAMojjvuOJP5HuhGyWrSpInJPv300+D1fH8+33LLLZGyZHr77bdN1qlTJ5Nt3LixJMpBAd/AsaKGPhxs69atJps7d67JRo8ebTLfAIriuPLKKyMdt23btkjHVa5c2WRvvvmmyTZv3hxpvWTjzhQAAAAABKCZAgAAAIAANFMAAAAAEIBmCgAAAAACZPwACt+DzGPHjjXZH/7wB5N98MEHca/H93Csz4IFC+J+beDn3HTTTZGO27t3r8lGjBgR73KQYa666iqTjRs3zmRly5aNvObixYtNdu6555rM92D0+vXrI12jXLlyJvP92cIAipJVv359k7300kvB623ZssVk+/fvN1n16tVNVpwhVr7BWFHP9z2QX7VqVZO1adPGZMOGDTNZnz59Il0XxXfiiSeazLc/ff/v//Wvf5msc+fOJluxYkVgdcUzf/58k40fP95ky5cvN1m/fv1MVqlSJZP17t07sLrk484UAAAAAASgmQIAAACAADRTAAAAABCAZgoAAAAAAmT8AIqdO3eazDcEom3btiaLdQDFL37xi0jX8f1U6ylTpsR0beBQqlWr5s19DzL7rF271mRTp06NqSZkljp16phs0KBBJos6bGLVqlXevFevXiZbsmRJpDVj0b59e5NNmjQp4dfFf1x//fUm8w2l8Bk5cqTJ/va3v5nM9x7irLPOinSNRPj8889N9s0330Q6t3LlyvEuB4fQtGlTk5UpE+1t9/nnn2+yjRs3xlxTqLy8PJPdeOONJuvSpYvJatasabIdO3aYzPd7LV1wZwoAAAAAAtBMAQAAAEAAmikAAAAACEAzBQAAAAABMn4AxY8//pi0a3fq1MlkvoetP/74Y5MV9bA1EA9Dhw6N6fxFixbFqRJkAt+wnZkzZ5qsUaNGwdcYNWqUN//nP/8ZvGYsTjzxxKRcN1u1bt3aZDfffHPwevfff7/Jor5fePnll4OvG6sGDRpEOs45Z7Jzzz3XZOXLlzfZrl27il8YjF/96lfB555yyikme+ONN2Ipp0QMHDgw0nFjxoxJcCUliztTAAAAABCAZgoAAAAAAtBMAQAAAEAAmikAAAAACJDxAyhq1KiRtGvXrl070nHJeoAa2evaa6+N6fz77rsvTpUgE0yaNMlksQxoWLBggckmT54cvF4ipFo9mc43HMI3PGH37t0mGzdunMk2btwYn8JKWNeuXSMdp6ommz17tskYNpE4Tz31lMkGDBgQ6dzXXnst0nH/+Mc/TObb276hZi+99JLJPvjgg0jXFRHp0aOHyZo1a2ay1atXm+yOO+6IfJ10wJ0pAAAAAAhAMwUAAAAAAWimAAAAACAAzRQAAAAABMj4ARSdOnUyme/BzFjUqVPHm/fu3TvStR977LG41gPE06ZNm0z2+uuvJ6ESpIJzzz3XZGeffXbwetu3bzfZxRdfbLLNmzcHX6Movu/HUf982Lp1a7zLwSH8+9//NplvyInv/8uKFSsSUlMyVKlSJdJxzrkEV4Kfk5eXZ7ILLrjAZMOGDTOZ7//zMcccE2k9H9/3tX79+pls/fr1kdYTEalatarJfPvuhx9+MNnJJ59ssoULF0a+dqrhzhQAAAAABKCZAgAAAIAANFMAAAAAEIBmCgAAAAACZNQAinLlypns+uuvN5nvAbkuXbqYLDc312Q1atQwWdOmTb31VK5c2WSffvqpyb7//nvv+UA8+H4iedmyZSOfP378eJPt3bs3ppqQHqpVq2ayiRMnmizqw+6+YRM9evQw2bJlyyKtVxw5OTkmO/LII03m+1z27dtnskwaapAOfP9fvvrqqyRUUnKGDh1qshtuuCHSub5BHJMmTYq5JkS3Z88ek82aNStS5nv/GHUAhe/7tm8Ahe/3lO/7sYhIzZo1g9ds2bKlyT755BOTLVq0yGQDBw40WSoOwOLOFAAAAAAEoJkCAAAAgAA0UwAAAAAQgGYKAAAAAALQTAEAAABAgIya5te1a1eT+abv+TRp0sRkvil9UadWFeWee+4x2f79+2NaEziUUaNGmaxMGf9vfd/0Id80P2QH34TU2rVrB6/3yiuvmOzFF18MXq84brzxRpO1a9cu0rm7du0ymW8CFxDqrrvuMtmgQYNM5pug5uObuvnPf/6z2HUhOXzTGD/77LNIWVQdOnQwWa9evSKfP3/+fJONHj3aZOeff77J2rdvbzLfe+7nn3/eZM2bNzfZd999V2SdJYE7UwAAAAAQgGYKAAAAAALQTAEAAABAAJopAAAAAAiQUQMoWrZsabIdO3aY7LHHHjPZypUrTbZhwwaTrVu3zmTTpk2LWqK8+uqrkY8Fiqt+/foma9WqlcmKGqSyePFik61evTr2wpDy2rRpY7K///3vwev59tjMmTOD14vVb37zm+Bzc3JyTNaiRQuTffzxx8HXQOYpalhEt27dTPanP/0p8vkHmzNnjsluueWWSOciO9xxxx0mGzhwoMkqVKjgPf/dd981WY8ePUzmGwTx3HPPmax169YmmzdvnsmqVKliskqVKnlrTCbuTAEAAABAAJopAAAAAAhAMwUAAAAAAWimAAAAACBARg2g6NOnT6QsFp07dzZZUQ+JTp8+3WRbtmyJaz1AYQMGDDBZxYoVI58/atSoeJaDNDJu3DiTVa5cOXg934PITz31VPB6xXHmmWea7PTTTw9eb//+/SbbuHFj8HrIPLm5uSa78847vcdeeeWVJitqKNDBvv76a5P9/ve/N9nevXsjrYf0VrZsWZO99NJLJuvYsaPJfHuuqO/Rffv2NdnmzZujlOjVvHnzSMd9/vnnJsvLywu+bqJwZwoAAAAAAtBMAQAAAEAAmikAAAAACEAzBQAAAAABMmoARUno2rWryYp6cPSjjz5KdDnAf2nXrl1M50+ePDkudSD9+H5KfVEP0Efx7LPPxlJOZN27dzfZHXfcYbLSpUsHX+P222832bfffhu8HtLbSSedZLKRI0ea7LzzzvOeH3XYxIsvvmgy35Ch5cuXR1oP6eOoo44ymW8A2uWXXx7p3J9++slkvj3ry0REdu7c6c2j8A3B6t27d6RzR4wYYbJUHK7CnSkAAAAACEAzBQAAAAABaKYAAAAAIADNFAAAAAAEYABFMbVt29ZkRT1M+tZbbyW6HGSxk08+2WSNGjWKdK7vJ6Qju61evTqu6+Xk5JjsmmuuMdkpp5xismXLlpmsqOEqbdq0iXRtn/3795vMN4hjzJgxkdZD5qlTp47JJk2aZLIWLVrEdJ2+ffuabMKECTGtidRToUIFkz344IMm69Gjh8miDi554403TDZo0CCTTZs2LdJ6sWrSpInJfO9VVqxYYbKZM2cmpKZ4484UAAAAAASgmQIAAACAADRTAAAAABCAZgoAAAAAAjCA4hCaN29usjJl7Jfstdde857/wQcfxL0m4IBx48aZrGzZspHOHTp0aLzLAf7LgAED4rpeqVL+v/vzDZHwWbNmjcn++te/muzee+8tXmHIaDfddJPJWrZsaTLfcIBt27Z517zllltMNnHixIDqkMp+/etfm8z357ZvCI+qmsz3/WrYsGEm27hxY9QS465evXommzFjhsl8n99dd91lss2bN8ensATjzhQAAAAABKCZAgAAAIAANFMAAAAAEIBmCgAAAAACMIDiEEaOHGmyypUrm6x9+/be83v37m0yfqI5QlSqVMlkxx57bKRzfQ+j5uXlxVwTMovvJ8379skJJ5xQEuUYvgf8RUTWrVtnskceecRkkyZNMtmSJUtirguZw/cAvG8AhW8v+h6UHzRokPc6Dz/8cEB1SDeXXHKJyXyDzYr63nawL7/80mS+96S+IRCJcNppp5nMt+erVatmsm+//dZkvu/b6YI7UwAAAAAQgGYKAAAAAALQTAEAAABAAJopAAAAAAjAAIpD8D0U6Mu++OIL7/nTpk2Le03ITo0aNTLZ0UcfHenc9957z2S7d++OuSZklpUrV5qsTZs2JrviiitMNmTIEJPVqlUruJbJkyeb7B//+If32Pfff99kq1evDr42soPvofiuXbuarEwZ+zZJVU32zDPPmIxBE9nN933swgsvNJnvz3cf34AG34Cp6tWrm8y3Z6MOviiKb03fewvfcCPf77V0xp0pAAAAAAhAMwUAAAAAAWimAAAAACAAzRQAAAAABGAAxSH88pe/NNn27dtN9rvf/c57/tq1a+NeE7KT76HVqCZOnBjHSpBNfA83T5gwIVIGpLIuXbqYLDc3N9K53333ncmGDx8ea0nIMHl5eSZr1qyZyXyDfk4//XST+fZnhQoVTNa5c+eIFVq+mkVE5s+fbzLfoJ+XXnrJZB988EFwPemCO1MAAAAAEIBmCgAAAAAC0EwBAAAAQACaKQAAAAAIoMX5CciqGtuPS04z69atM5nvgeyGDRuWRDlpyTlnf0R2CcqUPXvEEUeY7IsvvjCZ7/fzcccdZzLfIBXkY88iDc13zrVI1sXTcc927NjRZDNmzDCZ73tq7969TfbII4/Ep7DswZ5Fuilyz3JnCgAAAAAC0EwBAAAAQACaKQAAAAAIQDMFAAAAAAHKJLuAVOZ76B9IBt8wlFq1aiWhEgBIf3PmzDHZhx9+aLLjjz8+0rkAshd3pgAAAAAgAM0UAAAAAASgmQIAAACAADRTAAAAABCAARQAACCr/PTTTyZr1apVEioBkO64MwUAAAAAAWimAAAAACAAzRQAAAAABKCZAgAAAIAANFMAAAAAEIBmCgAAAAAC0EwBAAAAQACaKQAAAAAIQDMFAAAAAAHKFPP4dSKyNBGFICPVT3YBwp5F8bBnkY6SvW/Zsygu9izSTZF7Vp1zJVkIAAAAAGQE/pkfAAAAAASgmQIAAACAADRTAAAAABCAZgoAAAAAAtBMAQAAAEAAmikAAAAACEAzBQAAAAABaKYAAAAAIADNFAAAAAAEoJkCAAAAgAA0UwAAAAAQIKuaKVVdoqodIhznVLVB4DWCzwUOxp5FumHPIt2wZ5Fu2LOpJauaqVSlqmeq6lxV3ayqS5JdD/BzVHWgqn6uqltV9XtVHZjsmoBDUdVZqrqt0Gu3qi5Kdl1AUXhvgHSkqs1VdV7B99k1qnpTsmtKNJqp1LBdRB4TEd6QIl2oiFwlItVF5DwR6auqVyS3JKBozrmOzrlKB14i8p6IPJ/suoBD4L0B0oqqHiEir4rIwyJyuIg0EJHXklpUCcjKZkpVT1XV91V1k6quUtVxqppz0GHnq+p3qrpOVUeraqlC5/dU1S9VdaOqzlbV+rHU45z7l3PuSRH5LpZ1kLlScM+Ocs594pzb65z7WkReFpHTY1kTmSXV9uxBteWKyBki8kS81kT6S7U9y3sD/JxU27Mi0l9EZjvnnnLO/eSc2+qc+zLGNVNeVjZTIrJPRPqJyBEi0kpE2otIn4OO6SQiLUSkuYhcJCI9RURU9SIRGSwivxORmiLytohM9V1EVW8p2ODeVwI+L2SulN2zqqqS/8b0ixg/R2SWlN2zkn9X9W3n3JJYPkFknFTes4BPqu3Z/xGRDar6nqr+qKqvqGq9uH22qco5lzUvEVkiIh08+c0i8mKhXzsROa/Qr/uIyJsFH88SkWsK/bdSIrJDROoXOrdBYH0dRGRJsr9OvFLnlep7tuD8O0VkoYiUS/bXi1fyX2myZxeLyNXJ/lrxSo1Xqu9Z3hvwOviVqntWRL4RkU0i0lJEyovI/SLybrK/Xol+ZeWdKVVtpKr/UNXVqrpFRIZLfldf2LJCHy8VkdoFH9cXkfsKdeMbJP/5kTqJrhvZK1X3rKr2lfy/5b/AOfdTrOshc6Twnm0tIkeJyLRY10JmSdU9CxQlBffsTslv5j5yzu2S/L9sPU1Vq8awZsrLymZKRCaIyFci0tA5V0Xyb3PqQcfULfRxPRFZWfDxMhHp5ZyrVuhVwTn33sEXUdXB+t/To/7rlYDPC5kr5fasqvYUkVtEpL1zbnmcPk9kjpTbswV6iMh05xzfg3GwVN2zQFFSbc9+Jvl3tA5wkgWytZmqLCJbRGSbqjYWkd6eYwaqanVVrSsiN4nIswX5QyIySFVPFBFR1aqqeqnvIs654a7Q9KiDXweOU9VSqlpeRMrm/1LLq32AENkt1fZsN8n/G7CznXM8HA2flNqzBetUEJHLRGRyXD5DZJqU2rO8N0AEKbVnReRxEemkqs1UtayIDBGRd5xzm+Pz6aambG2mBohIVxHZKiKPyn82VmEvi8h8EVkgIjNEZJKIiHPuRREZKSLPFNxS/VxEOsZYTxvJvzU6U/L/1mCnZMEoSRRLqu3ZuyV/7OlHhf526qEY10RmSbU9KyJyseT/e/65cVgLmSfV9izvDfBzUmrPOufmSP7dsRki8qPkj0bvGsua6UCdy4o7cAAAAAAQV9l6ZwoAAAAAYkIzBQAAAAABaKYAAAAAIADNFAAAAAAEoJkCAAAAgABlinOwqjL6D8XinDv4h8eVKPYsios9izS0zjlXM1kXZ88iAHsW6abIPcudKQAA0tvSZBcAFBN7FummyD1LMwUAAAAAAWimAAAAACAAzRQAAAAABKCZAgAAAIAANFMAAAAAEIBmCgAAAAAC0EwBAAAAQIBi/dBeAAAAANmjfPnyJjvnnHNM1r9/f5P99a9/NdlHH31kslWrVgVWl3zcmQIAAACAADRTAAAAABCAZgoAAAAAAtBMAQAAAEAABlAAAAAA8Bo/frzJevToEenc1q1bm2zChAkmu/HGG4tfWIrgzhQAAAAABKCZAgAAAIAANFMAAAAAEIBmCgAAAAACMIAigcqVK2eyd99912THHnusyTp06GCyTz75JD6FIWU88MADJjvllFMinfvqq6+abOnSpSZbvXq1yWbPnh3pGgCA1NG4cWOTLViwwGQfffSRyc4444yE1IT0lJOTYzLfexIRkauvvtpkzrlI19mzZ4/JPvjgg0jnpgvuTAEAAABAAJopAAAAAAhAMwUAAAAAAWimAAAAACAAAygSqHr16iZr3rx5pHMnT55sspYtW5rsp59+KnZdSDzf8BHfTxDv2bNn8DVatWplMt8Dofv37zfZxx9/7F3ztttuM9lrr70WUB0AIN5at25tstKlS5vspJNOMtlxxx1nsm+//TY+hSHtXH/99Sa75pprYlrTNwTr7rvvNtlTTz0V03VSDXemAAAAACAAzRQAAAAABKCZAgAAAIAANFMAAAAAEIABFAl0xx13BJ9bpUoVk9WsWdNky5cvD74GEufPf/6zyWIZNuET9aePlypl/87k1FNP9R7rG5LRpUsXkxU1wAJIlDZt2njz+++/32THH3+8yfr372+yCRMmxF4YkCAdO3Y0mW9IUJky9q3cjh07TLZr1674FIa04/v+2a9fv5jW9A0v6dChg8mWLVsW03XSAXemAAAAACAAzRQAAAAABKCZAgAAAIAANFMAAAAAEIABFHHQqVMnb96rVy+TRR0akJeXZzKGTaSP2rVrRzpu+vTpJlu4cKHJtm3bZrInn3zSZOXKlTOZ7yeNn3baad56jjvuOJM98sgjJmvZsqXJ9u3b510Tma9SpUom27t3r8l8++ukk04ymW9/FjWAokmTJlFKlFatWpmMARRIFaVLlzZZnz59TFa3bl2T+b73vvnmmyZbsWJFYHVIJ1WrVjXZ0KFDTZabmxt5zTVr1pisW7duJsuGYRM+3JkCAAAAgAA0UwAAAAAQgGYKAAAAAALQTAEAAABAAAZQxEHjxo1jOt/3UGjPnj1jWhPJ5Xuw/YcffjDZqFGjTBbvQQ7t2rUz2auvvuo99pxzzjFZs2bNTPaHP/zBZOPHjy9+cc5TUHEAACAASURBVEgZhx12mMlmzpwZ6dzdu3ebrEGDBiarVauWycqXL28yVTVZ1OE9Rdm6dWtM5wOJ5BsQ8Jvf/CbSuR999JHJrrrqqphrQnoaN26cyVq3bm2y4nxP9Q2dYqDJf3BnCgAAAAAC0EwBAAAAQACaKQAAAAAIQDMFAAAAAAEYQBEHsT7o+cgjj5hs9erVMa2J5Prss88iZcly9913e3PfsIqcnByTDRkyxGSvvPKKyXxDN5CaKlSoYLIzzjjDZLEMh9i1a5fJtm3bZrLHH3/cZIcffrh3zcsvv9xkpUuXNplvSAaQDL6hVTfffHOkc30DinzDK5KlRYsWJvv444+TUEl2uO6660zWuXPnSOfu2bPHZH379vUey7CJQ+POFAAAAAAEoJkCAAAAgAA0UwAAAAAQgGYKAAAAAAIwgKKYunTpYrKGDRvGtOayZctiOh8ornfeecebjx492mT/+7//a7IjjzzSZLm5uSZjAEX62Lp1q8kuuOCCuF5jyZIlJtuyZYvJVq5cGXnNU0891WQNGjSIdB0gkQ477DBvfvvtt0c+9mBTp0412axZs4pXWALt2LEj2SVklYceeshkUQcCPf/88yabNGlSzDVlI+5MAQAAAEAAmikAAAAACEAzBQAAAAABaKYAAAAAIAADKIppyJAhJitVKnpPunbtWpNNnz49ppqAeHn55ZdN5htA4dOkSROTzZs3L+aaUDJ2795tsldffTUJlfhVq1bNm/se3FdVk/mGXwCJdOGFF3rzK664ItL5GzZsMNnDDz8cU02JlpeXl+wSMpZvAFpUvj+L+/btG0s5KIQ7UwAAAAAQgGYKAAAAAALQTAEAAABAAJopAAAAAAjAAIpiql69ekznjx071mRbtmyJaU0gFfgetvb9dPZ9+/aVRDnIMMcff7w3r127tsmccyY788wzTfb444/HXhggIu3atTPZlClTIp/v27P9+/c32TvvvFOsupA5Bg8ebDLfsB2ft956y2S894wf7kwBAAAAQACaKQAAAAAIQDMFAAAAAAFopgAAAAAgAAMoDuHKK6802ZFHHhn5/G3btplszJgxMdUEJNLatWtNtm7dOpMdccQRJmvQoIHJcnJyTLZz587A6pDNmjRpEtP5ixYtilMlgHXbbbeZrFy5cpHPHzdunMmKM8ACmaVZs2Ymq1+/vsl8g0t82UsvvRRTPVWrVjXZOeecY7Lf//73JitfvrzJnnnmGZM98sgjgdUlH3emAAAAACAAzRQAAAAABKCZAgAAAIAANFMAAAAAEIABFIdw9tlnm6xUqej95969e022Z8+emGpC+qpWrZrJateuHelc31765ptvYq7pYDVr1jSZb9iEz9ixY03GsAnES6wDKBLx+wXZqXfv3iZr3bp15POXLl1qsltvvTWmmpBZfO8XDjvssEjn5uXlmWzx4sUm8w2IOv30071rTps2zWS+oRRRtWnTxmS+oRt9+vQJvkZJ4s4UAAAAAASgmQIAAACAADRTAAAAABCAZgoAAAAAAtBMAQAAAEAApvkV8E0RufDCC03mnIu85qhRo2KqCemrY8eOJvNNu2vUqFGk9Xbv3m2yO++802QzZ8402cKFCyNdQ0TkoosuinzswRYtWhR8LjKPby/5JvJ9//33JuvWrZvJGjduHFM948aNM9kpp5xisttuuy2m6yCz1KpVy2R/+ctfTFa2bFmT+aawioiMHj3aZFu2bAmoDrB80/x8k6R9k/Luvfde75qqarLivB+OokuXLibzfd/2fX7Jxp0pAAAAAAhAMwUAAAAAAWimAAAAACAAzRQAAAAABGAARYGGDRuarGrVqjGtOWPGjJjOR/p6+eWXTVamTPhvt5ycHJMNGzbMZLfffrvJXnnlFZMVtTf//Oc/R6rH9zDrTz/9FOlcZJ6JEyea7PLLLzdZxYoVI60X68POvoEtvt9DQGG+79FTpkwxWf369SOtV9RQnvHjxxevMCBGDzzwgMmuueaamNb0vc95++23TVbUUIuDValSxWQ1a9YsfmFJwJ0pAAAAAAhAMwUAAAAAAWimAAAAACAAzRQAAAAABGAARQKdfvrpJvvss8+SUAlK2ooVK0wW9aHlVatWmcz3IPM555xjMt9D9pdcckmkrDgWL15ssg8//DCmNUM1b97cZHXr1jWZ72FZxMfdd99tsjp16pjsuOOOM9m6detM5htAUa9ePe+1jzrqKJPNmTPHZL6BGFu3bvWuiex00kknmezcc8+NdO7evXtNdtddd8VcE3CA7/uiL7v00kvjfu3u3bubbObMmSbz/Tnrq9GnVKn0vb+TvpUDAAAAQBLRTAEAAABAAJopAAAAAAhAMwUAAAAAARhAUeCKK66I+5qjRo0y2YQJE+J+HaSeoUOHmuzhhx82WZky9rfg/PnzTXb99debrHz58ibz/fRx3yCAWDVs2NBkvqEbeXl5JjvhhBPiWku1atVM5nvg9bDDDovrdfEfS5YsMVnHjh1NVrlyZZNFHQLhGyoh4h9A0bhx4+DrIHsNGTIk+Nz77rvPZC+++GIs5SCLrV271mRbtmwxWZUqVUzmnAu+Ru/evb3H+oZbPffccyZr3bp1cD2vv/66yRYuXBjp3GTjzhQAAAAABKCZAgAAAIAANFMAAAAAEIBmCgAAAAACaNQHw0REVDX6wWnm008/NdnJJ58c05rbt283me8B7EzmnIv2o68TJJX27JlnnmmyRx991GTHHntspPXef/99kzVr1sxkFSpUiLReOvjhhx9M5hu6MWbMGJMtWLAg0jXYs8mXm5trskWLFnmP9Q0WWblypcnq1q0bc10pbL5zrkWyLp6Oe7ZFC/vlmjdvnsmifv8844wzTPbOO+8Uv7DswZ4tJt9wiHHjxpks6vv6PXv2mGzz5s3eY2vWrBl8HZ+NGzearF69eibbuXNn8DUSoMg9y50pAAAAAAhAMwUAAAAAAWimAAAAACAAzRQAAAAABCiT7AIyGT/9HIXNnTvXZP379zfZ6NGjTeZ7IL9Vq1aRrrt7926T+QauiIgMGzbMZF999VWk6/j07NnTZL6fpD5//nyTffTRRybbtGmTydatWxdYHVLVL3/5S5P5Bk0U5YUXXohnOchAAwYMMFnUYRNvvPGGyT788MOYawIOZcKECSbzDaCIqmzZsiY74ogjgtcT8Q+1eP75503Wt29fk6XYsIli4c4UAAAAAASgmQIAAACAADRTAAAAABCAZgoAAAAAAmTlAIq2bduazPfAc1SfffaZN7/qqquC10R2+Pvf/x4pa9asmcmaNm0a6Rrz5s0z2ZIlSyKdG6vBgweXyHWQWXwDV1Q18vmrVq2KYzVId0ceeaTJog7w8bnnnntM5nvwHki0W2+91WR33XVX3K+zceNGk02dOtVkEydONNnChQvjXk+q4c4UAAAAAASgmQIAAACAADRTAAAAABCAZgoAAAAAAmTlAIrDDjvMZDk5OcHrzZgxI5ZygJ+1YMGCSBmQCY444giTOecinz937tx4loM0V716dZPVq1cveL39+/fHUg4QNyNGjIiUIbG4MwUAAAAAAWimAAAAACAAzRQAAAAABKCZAgAAAIAAWTmA4vXXXzfZzTffbLKzzz7bZN9++63J3nrrrfgUBgCQRo0aRT52yZIlJvvss8/iWA3S3ffff2+yBx980GR9+vQx2YYNG0y2bNmy+BQGICNwZwoAAAAAAtBMAQAAAEAAmikAAAAACEAzBQAAAAABtDg/VV5Vox8MiIhzTpN5ffYsios9m3xPPvmkybp27eo99osvvjBZ06ZN415TipvvnGuRrIuzZxGAPYt0U+Se5c4UAAAAAASgmQIAAACAADRTAAAAABCAZgoAAAAAApRJdgEAAIR64YUXkl0CACCLcWcKAAAAAALQTAEAAABAAJopAAAAAAhAMwUAAAAAAdS56D8Emp8YjeJyzmkyr8+eRXGxZ5GG5jvnWiTr4uxZBGDPIt0UuWe5MwUAAAAAAWimAAAAACAAzRQAAAAABKCZAgAAAIAANFMAAAAAEIBmCgAAAAAC0EwBAAAAQACaKQAAAAAIQDMFAAAAAAHKFPP4dSKyNBGFICPVT3YBwp5F8bBnkY6SvW/Zsygu9izSTZF7Vp1zJVkIAAAAAGQE/pkfAAAAAASgmQIAAACAADRTAAAAABCAZgoAAAAAAtBMAQAAAEAAmikAAAAACEAzBQAAAAABaKYAAAAAIADNFAAAAAAEoJkCAAAAgAA0UwAAAAAQIKuaKVVdoqodIhznVLVB4DWCzwUOxp5FumHPIt2wZ5Fu2LOpJauaqVSmqs1VdZ6qblPVNap6U7JrAn6Oquao6pequjzZtQCHoqqzCr6/HnjtVtVFya4LKAp7FulGVfup6nequkVVV6rqWFUtk+y6Eo1mKgWo6hEi8qqIPCwih4tIAxF5LalFAdEMFJG1yS4C+DnOuY7OuUoHXiLynog8n+y6gKKwZ5GG/i4izZ1zVUTkJBE5WUT+mNySEi8rmylVPVVV31fVTaq6SlXHqWrOQYedX9Bdr1PV0apaqtD5PQv+Nn6jqs5W1foxltRfRGY7555yzv3knNvqnPsyxjWRQVJwz4qqHiMi3UVkRKxrIfOk4p4ttHauiJwhIk/Ea02kP/Ys0k2q7Vnn3LfOuU0HlheR/ZJ/gyCjZWUzJSL7RKSfiBwhIq1EpL2I9DnomE4i0kJEmovIRSLSU0REVS8SkcEi8jsRqSkib4vIVN9FVPWWgg3ufRU69H9EZIOqvqeqP6rqK6paL26fLTJBqu1ZEZEHCtbdGY9PEBknFffsAVeJyNvOuSWxfILIOOxZpJuU27Oq2lVVt4jIOsm/M/VwfD7VFOacy5qXiCwRkQ6e/GYRebHQr52InFfo131E5M2Cj2eJyDWF/lspEdkhIvULndugmHV9IyKbRKSliJQXkftF5N1kf714Jf+Vwnu2k4jMKvi4nYgsT/bXildqvFJ1zx5Uy2IRuTrZXyteqfFiz/JKt1ea7NmGInKXiByV7K9Xol9ZeWdKVRup6j9UdXVB9zxc8rv6wpYV+nipiNQu+Li+iNxXqBvfIPm3MuvEUNJOyd/8HznndonInSJymqpWjWFNZJBU2rOqWlFERkkW/DtohEulPXtQXa1F5CgRmRbrWsgs7Fmkm1TdsyIizrl/i8gXIvJgPNZLZVnZTInIBBH5SkQauvyH5AZL/gYqrG6hj+uJyMqCj5eJSC/nXLVCrwrOufcOvoiqDtb/nsTzX69Ch34m+X8DcIAT4L+l0p5tKCK5IvK2qq4WkekicnTBN/PcOH2+SH+ptGcL6yEi051zvv+G7MaeRbpJ1T17QBkROS74s0sT2dpMVRaRLSKyTVUbi0hvzzEDVbW6qtYVkZtE5NmC/CERGaSqJ4qIqGpVVb3UdxHn3HBXaBLPwa9Chz4uIp1UtZmqlhWRISLyjnNuc3w+XWSAVNqzn0v+N+dmBa9rRWRNwcfLfOsiK6XSnpWCdSqIyGUiMjkunyEyDXsW6Sal9qyqXquqRxZ8fIKIDBKRN+P1yaaqbG2mBohIVxHZKiKPyn82VmEvi8h8EVkgIjNEZJKIiHPuRREZKSLPFNxS/VxEOsZSjHNujuT/bcIMEflR8iefdI1lTWSclNmzzrm9zrnVB16S/08D9hf8el/ousg4KbNnC7lY8p9PnRuHtZB52LNIN6m2Z08XkUWqul1EZha8Bse4ZspT5/gXZQAAAABQXNl6ZwoAAAAAYkIzBQAAAAABaKYAAAAAIADNFAAAAAAEKFOcg1WVaRUoFufcwT/voESxZ1Fc7FmkoXXOuZrJujh7FgHYs0g3Re5Z7kwBAJDelia7AKCY2LNIN0XuWZopAAAAAAhAMwUAAAAAAWimAAAAACAAzRQAAAAABKCZAgAAAIAANFMAAAAAEIBmCgAAAAAC0EwBAAAAQIAyyS4AIp9++qnJnnjiCZONHTu2JMoBAAAAEAF3pgAAAAAgAM0UAAAAAASgmQIAAACAADRTAAAAABCAARQl7IEHHjBZ3bp1TfbUU0+VRDkAAAAAAnFnCgAAAAAC0EwBAAAAQACaKQAAAAAIQDMFAAAAAAEYQJFAV199tcn69OljsrFjx5rsxx9/TERJSANnnnmmyfbs2WOyefPmmezKK680GcNMAADIXr73o0ceeaTJqlat6j1/8ODBka4zZMgQk+3evdtkTz75pMlWrVoV6RqpiDtTAAAAABCAZgoAAAAAAtBMAQAAAEAAmikAAAAACKDOuegHq0Y/GPLll1+a7KijjjLZ2WefbbKPP/44ITWVNOecJvP6qb5np0yZYrLLLrvMZL7fpxUqVDDZ008/bbJu3boFVped2LNIQ/Odcy2SdXH2LAKwZw8hJyfHZL4/y0eMGGGy0qVLm6xatWqRjispmzdvNtkjjzxiskcffdRkixcvTkhNERS5Z7kzBQAAAAABaKYAAAAAIADNFAAAAAAEoJkCAAAAgABlkl1AJhg0aJA3P/744032xz/+0WSZMmwCxffVV1+Z7McffzRZvXr1SqIcIKEqVqxoslq1apnM932yKN27dzdZjRo1Ip377LPPmuy6664z2bZt2yLXA/wc3/CgDh06mGzatGkmK1PGvm275ZZbTDZ69OjA6pAKypcvb7JJkyYl/LqffvqpN1+xYoXJGjZsaDLf+16fqlWrmmzgwIEmO//8803WpEmTSNcoSdyZAgAAAIAANFMAAAAAEIBmCgAAAAAC0EwBAAAAQAAGUBRT27ZtTfaXv/zFe+wnn3xisqlTp8a9JqQv308vL1XK/h3H3XffHWm9H374IeaagOJq06aNyc477zyTnXXWWSZr2bKlyVTVZM65yPVEPfayyy4z2RtvvGGyknjwG5mnUaNG3nzIkCEm69KlS6Q1fXu7adOmxSsMKW/fvn0m+/DDD0126qmnmsz3/XPz5s0mW7x4sckuuOACbz1RB2P5Mt8wlF//+tfe6xysTp06Jqtdu7bJVq5cGWm9ROHOFAAAAAAEoJkCAAAAgAA0UwAAAAAQgGYKAAAAAAIwgOIQatSoYbKHHnrIZKVLl/aef+2115ps/fr1sReGjLZly5bgc8eNGxfHSpDtjjzySJP93//9n8natWtnsqK+L4Z6/fXXvblv6Er58uVN1q1bt0jXyc3NLVZdgIjISSedZLK5c+d6j61evXpcr52XlxfX9ZB827dvN1mrVq1Mdu+995rs888/N5lv2MSmTZtM5hs0URTf915f5qvx+eefj3SNatWqmaxr166RrlGSuDMFAAAAAAFopgAAAAAgAM0UAAAAAASgmQIAAACAAAygKFCqlO0rH3vsMZP5fqL573//e++aCxYsiL0wZB3fg8w+u3fvNtn+/fvjXQ6yxMUXX2yy22+/3WRNmzYNvsYnn3xisunTp5vs/vvvN9lPP/3kXXPv3r0mq1WrlsmiDqBAevP9WT5w4ECTvfnmmyb7+OOPTVamjH2bdN1115nsjjvuMJlviJWIyKJFi0x22WWXmWz+/PkmW7ZsmcmmTJnivQ4y34ABA5Jdws/yDQSKyjlnMt97n2TjzhQAAAAABKCZAgAAAIAANFMAAAAAEIBmCgAAAAACMICiwODBg03229/+1mTjx4832RNPPJGQmoBDeeihh0y2atWqJFRStD/96U8m27x5s8kmTpxYEuXgEEaPHm2yY4891mTr16832dSpU032/vvvm2zu3LkmW7NmTdQSgZ9Vs2ZNkw0bNsxkP/74o8l8D8r/5S9/MVnHjh0j1TJq1ChvPm7cOJP5BllVqFDBZEOHDjXZypUrI9UDJNqQIUNM1q9fv+D11q1bZzLfgKJk484UAAAAAASgmQIAAACAADRTAAAAABCAZgoAAAAAAmTlAIqjjz7aZP379zfZF198YbK77rorITUd7OKLLzaZ7+HYWbNmmcz3gD9SU7Vq1UzWpEmTJFSSGDfccIPJcnNzTbZixQqT+fY2EufWW281WcuWLU3mG3yyePHihNQU6vLLL092CUhxjz76aPC5vkEqvmETr7/+uvf8Fi1amMw36GLjxo2R1wQSqW3btibzDZj6zW9+E9frvvvuu3FdL1G4MwUAAAAAAWimAAAAACAAzRQAAAAABKCZAgAAAIAAWTmAYuLEiSbzDXfwPcTs+6npRfE9ZNqpUyeTDRgwwGQ5OTkmc86ZbNu2bSbzPQA4b968IutE8hx++OEma9WqVaRzS2oYSlQNGjQwWaVKlUy2du1ak61atSohNSG6Z599NlKWDurUqRN87r59++JYCUqa7/tLz549TeYbuFKjRg2TDRs2zGTjx4832e7du6OWKB06dDBZhQoVTOYbarF+/frI1wEOpXv37iYbMWKE91jfsKyKFSvGtZ4ePXqY7IUXXojrNRKFO1MAAAAAEIBmCgAAAAAC0EwBAAAAQACaKQAAAAAIkPEDKNq3b2+yc88912S+B/Hy8vJMVrp0aZPdcccd3msPGjTIZBs2bDCZ7yHTb7/91mS+wQTXXXedyS688EKTMYAi+XwPcBb1sGcUW7ZsiaWcuPM9VF2zZk2T+fbiggULElITMp9vaEDfvn2D13vwwQdjKQdJtn//fpM98cQTkbJ4833/ExHp1atXpPPHjBkTz3KQxXzDhDp27Ggy39CoWPkGt/neu7711lsm27FjR9zrSQTuTAEAAABAAJopAAAAAAhAMwUAAAAAAWimAAAAACBARg2gqFy5sskeffRRk3355ZcmmzNnjslq165tsieffNJkZ555preeN99802SdO3c22ebNm73nR+F7iG/dunXB6yFxGjdubLJLL700CZUULScnx2QnnHCCyc4++2yT9ezZMyE1AYfyu9/9zmTly5ePdO7cuXNNtmnTpphrAkREevTo4c3r1q1bwpUg2/kGSyRi2ITP22+/bbJXXnmlRK5dUrgzBQAAAAABaKYAAAAAIADNFAAAAAAEoJkCAAAAgAA0UwAAAAAQIKOm+XXp0sVkubm5Jjv11FNNVq9ePZM9/fTTJmvQoIHJHnvsMW89N954o8l27tzpPTYKX92+yX1TpkwJvgYSp2vXrsHnzp4922RHHXWUyfbv32+yfv36mcw3qVJEpEKFCia76KKLopQIpB3f76vdu3cnoRJkogsuuCDysc8995zJdu3aFc9ykMWuueYak3Xr1s1kN998s/f8qlWrmqykpgGmA+5MAQAAAEAAmikAAAAACEAzBQAAAAABaKYAAAAAIEBGDaDo37+/ydasWWOyrVu3mmzOnDkmO/roo012++23m+yvf/2rt55Yhk10797dZL4HCEeMGGGy1atXB18X8XHSSSeZ7JJLLgle79xzzzXZ0qVLI53rnDPZnj17vMdu3rzZZDfddJPJOnToYLILL7wwUj1APLVt29Zkqhrp3Hnz5sW7HGSp3/72tyZr166d99ivv/7aZL6BVb6BQkAI3/vCMWPGRMpERO677z6T+fasT7NmzUxWv359k0V9T5OKuDMFAAAAAAFopgAAAAAgAM0UAAAAAASgmQIAAACAABk1gKJixYomW7t2rcn27dtnMt+widtuu81kf/vb30y2ffv2qCV6dezY0WTDhw832cyZM012zz33xHRtJMYxxxxjstq1a8f1Gr494ntgefHixSabMmVKTNe+//77TeYb7MJPSEe81KhRw5ufdtppJvMNXfnmm28iZUAI3yCqogZIvPvuuyZbv3593GtCYvi+5yxYsMBkO3bsKIlyUt5xxx1nslq1apmMARQAAAAAkGVopgAAAAAgAM0UAAAAAASgmQIAAACAABk1gCKqTZs2maxXr14me+aZZ0y2bdu2mK7dqVMnk/l+srRvaMDAgQNNtmvXrpjqQWLMmjXLZO+//77JWrVqZbJ58+aZbOzYsSZ7+eWXTeZ78B7IBPfee683z83NjXT+xo0bI2XAz6lfv77JKleubDLfUB4R/wAfpI+5c+eazDdUZPTo0SbzvTdIB0OGDDHZjTfemIRKUhN3pgAAAAAgAM0UAAAAAASgmQIAAACAADRTAAAAABAgowZQ7Nmzx2THHHOMyX7xi1+Y7PHHHzfZ3r17TaaqJqtevbq3Ht9gic6dO5ts2rRpJhs0aJDJli9f7r0OUo9v7/j+35cvX95k69atM9mWLVviUxiQBmrUqGGy008/PfL527dvN1nv3r1jqgnZqUwZ+zbJ9+D94YcfbrKnn37au+bChQtjLwxJ880335isXbt2JmvZsqXJ7rzzTpNNnTrVZCtWrAgrLkH69euX7BJSGnemAAAAACAAzRQAAAAABKCZAgAAAIAANFMAAAAAECCjBlB07drVZLNnzzbZJ598YrIPPvjAZJ9//rnJfAMt2rdv763n+++/N9mVV15pMt8ACmSelStXJrsEIC3cfPPNJmvQoEHk830DKHjoHyEaN25sMt/+9JkxY0a8y0EKGD58uMkmT55ssooVK5ps1KhRJrv++utNNnHiRJN9+OGHESsMV7VqVW/+hz/8IXjNuXPnmiwvLy94vVTEnSkAAAAACEAzBQAAAAABaKYAAAAAIADNFAAAAAAEyKgBFL4hEj179jSZ78G+//mf/4mU7du3z2QjR4701jN+/HiTLV++3HsskO7WrFljskqVKiWhEqS7W2+91WTOucjnP/vss/EsB1miTBn7lmjQoEGRzn3uuedMxj7MTFOnTjVZbm6uyXxDG+rWrWsy33Cde+65J6y4FPTVV1+ZbNu2bUmoJHG4MwUAAAAAAWimAAAAACAAzRQAAAAABKCZAgAAAIAAWpyHelU1+sGAiDjnNJnXZ8+WHN/DthMmTDDZvHnzTNa2bduE1BSCPVuybrjhBpM98MADJivqz6q9e/earHbt2iZbv359QHVpY75zrkWyLp4pe/bPf/6zyYYPHx7p3KZNm5osLy8v5poyWMbvWd+wPHj1NQAAIABJREFUifbt25vsyiuvNFmrVq1MVr58+fgUFidr1641mW8QXPfu3U22devWhNSUYEXuWe5MAQAAAEAAmikAAAAACEAzBQAAAAABaKYAAAAAIID9cd8AEMD34OmGDRuSUAlSVbNmzUzme8Bf1c4AKWoAxZQpU0yW4cMmEAennnqqyYYMGRLp3IULF5ps2bJlMdeEzOLbE5MnT46UnXPOOSbzDTnxufXWW01WpUoVk61YscJk9913n3fNffv2mWzs2LGR6skG3JkCAAAAgAA0UwAAAAAQgGYKAAAAAALQTAEAAABAAAZQAIiLBQsWmOzwww9PQiVIVfXq1TNZpUqVTFbUsAmf6dOnx1QTstOAAQNMVqFCBZNt3brVZL4H/H3HAaFee+21SJnPvffeG+9y8DO4MwUAAAAAAWimAAAAACAAzRQAAAAABKCZAgAAAIAADKAAAJSI7777zmRbtmwxWdWqVU12/fXXe9ecM2dO7IUh67z33nsm69ixo8mWLl1qslmzZiWkJgDpiTtTAAAAABCAZgoAAAAAAtBMAQAAAEAAmikAAAAACKDF+Unzqhr9YEBEnHOazOuzZ1Fc7FmkofnOuRbJujh7FgHYs0g3Re5Z7kwBAAAAQACaKQAAAAAIQDMFAAAAAAFopgAAAAAgAM0UAAAAAASgmQIAAACAADRTAAAAABCAZgoAAAAAAtBMAQAAAECAMsU8fp2ILE1EIchI9ZNdgLBnUTzsWfy/9u47zKrq+v/42lQF6SK9JIIYEEHASp2ICqgUg2AgdLDgA4IFUSQUEQLql2ASpAmKGiE0RZCiBFQkkRaRL4oE+A1SpUmHoe3fHzPmO2HtwTP73pl7z53363nu8wwfzzl7DW7uzOJw1oRRrPctexaZxZ5F2GS4Z421NjsLAQAAAICEwD/zAwAAAAAPNFMAAAAA4IFmCgAAAAA80EwBAAAAgAeaKQAAAADwQDMFAAAAAB5opgAAAADAA80UAAAAAHigmQIAAAAADzRTAAAAAOAhRzVTxphkY0zTAMdZY0wVzzW8zwUuxZ5F2LBnETbsWYQNeza+5KhmKl4ZY/obY7YbY44ZY/YYY8YaY/LEui4gI8aYRcaYE+leZ40xG2NdF5ARY0xRY8xbxpj9aa+hsa4JuBz2LMLGGJNkjFlujDlqjEmOdT3ZhWYqPswXkTrW2sIicoOI1BKRvrEtCciYtba5tfaqn14iskpEZsW6LuAyxopIARGpLCK3iEgnY0y3mFYEXB57FmFzUkSmisgzsS4kO+XIZsoYc4sx5h/GmCPGmL3GmD8bY/JdcliLtLtFB40xLxtjcqU7v7sx5ltjzI/GmCXGmEqR1GOt3WatPfLT5UXkoohwaxX/EW979pLaKotIQxGZHq1rIvzicM/eLyJjrLWnrLXJIvKGiHSP8JpIIOxZhE287Vlr7Wpr7dsisj2S64RNjmymROSCiPQXkatF5HYRuVNEel9yTBsRqScidUSklaS9gRljWonI8yLygIiUFJHPReQ91yLGmIFpG9z5uuTYDsaYYyJyUFLvTE2MzqeKBBF3ezadziLyedoXe+An8bhnzSUf3xDJJ4iEw55F2MTjns15rLU55iUiySLS1JH3E5F56X5tRaRZul/3FpFlaR8vEpEe6f5bLhE5JSKV0p1bJYIaq4rIiyJSOta/X7xi/wrJnt0qIl1j/XvFKz5e8bpnReQdEZkrIoUk9c7/NhFJifXvF6/Yv9izvML2itc9m+5aTUUkOda/T9n1ypF3powx1xljFhhj9qXdDRopqV19ejvTfbxDRMqmfVxJRMal68YPS+rfFpWLRm3W2n+LyCYRGR+N6yExxOueNcY0EJHSIjI70mshscThnu0rIqdF5N8i8oGk/g3srgiuhwTDnkXYxOGezZFyZDMlIq+LyGYRqWpThz48L/99K11EpEK6jyuKyJ60j3eKyCPW2qLpXldaa1dduogx5nnz3xPP/ut1mfryiMi13p8dElG87tkuIjLXWnu5/YycKa72rLX2sLW2o7W2tLW2hqR+/Vsdxc8X4ceeRdjE1Z7NqXJqM1VIRI6JyAljzPUi8pjjmGeMMcWMMRVE5AkRmZmWTxCR54wxNUREjDFFjDEPuhax1o606SaeXfr66ThjTE9jzDVpH1cXkedEZFm0PlkkhLjas2nXuVJE2onIm1H5DJFo4mrPGmOuNcaUMMbkNsY0F5GHRWRE9D5dJAD2LMIm3vZsLmPMFSKSN/WX5gqjB2IknJzaTD0tIh1E5LiITJb/21jpfSAi60TkKxFZKKlTdMRaO09ERovIjLRbqv8rIs0jrKe+iGw0xpwUkY/SXs9HeE0klnjbsyIirUXkiIgsj8K1kHjibc/WFZGNafWMEpGO1tpNEV4TiYU9i7CJtz3bSFL/aepHknoX7LSILI3wmnHP2NQHxQAAAAAAmZBT70wBAAAAQERopgAAAADAA80UAAAAAHigmQIAAAAADzRTAAAAAOAhT2YONsYw+g+ZYq299IfHZSv2LDKLPYsQOmitLRmrxdmz8MCeRdhkuGe5MwUAQLjtiHUBQCaxZxE2Ge5ZmikAAAAA8EAzBQAAAAAeaKYAAAAAwAPNFAAAAAB4oJkCAAAAAA80UwAAAADggWYKAAAAADzQTAEAAACAB5opAAAAAPBAMwUAAAAAHmimAAAAAMADzRQAAAAAeKCZAgAAAAAPNFMAAAAA4IFmCgAAAAA80EwBAAAAgAeaKQAAAADwQDMFAAAAAB7yxLoAAAAAAIln6NChKhsyZIjKkpKSVLZixYosqCj6uDMFAAAAAB5opgAAAADAA80UAAAAAHigmQIAAAAAD6EdQJE/f36VVa5cWWVdu3ZV2VVXXaWy3/zmNyorUaKEyiZNmhSswAwsXrxYZevWrVPZvn37IloHAAAAiKXGjRsHOm758uUqC8tQCu5MAQAAAIAHmikAAAAA8EAzBQAAAAAeaKYAAAAAwEMoBlBce+21KhswYIDKevXqleW1PP7441E/f/fu3Sq76667VLZ58+aI1gaAWCpQoIDKypcvr7Lu3burrEqVKs5ruoYHrVmzRmWu4T8jR45U2ZkzZ5zrAAAuzzVEokmTJt7Xc53LAAoAAAAASBA0UwAAAADggWYKAAAAADzQTAEAAACAh7gbQOEaNrFs2TKVVaxY0XuNCxcuqOzixYve18tI7ty5VZYrl+5fy5Urp7K5c+eqrHr16tEpDHGjYMGCKqtcubLKNm3apLIiRYqorE2bNiobNGiQylx/zkREjhw5orJ3331XZa6H9N98802VHT58WGU//vijyp566imV3X333SobPny4ylzvD4i9kiVLqsz1/6pGjRoRreN6765bt6539tBDD6ns+PHjntUhTFwDUlxfdx944AGV9enTR2VXXXWVc51//vOfKrvhhhtU1qxZM5WtXr1aZefOnXOuA2Ql13CISIZNhBl3pgAAAADAA80UAAAAAHigmQIAAAAADzRTAAAAAOAh7gZQVKhQQWVlypRRWUpKisq2bt2qsqlTp6rso48+Utl3330XtMTAkpKSVBb0YXnX5+waGrBt27bMF4aYuOKKK1TmGtrgerj5888/V1mJEiVUFnRIibXWmRctWlRlvXv3DnTNJ598UmV79uxR2d69e1XmGgTgMmvWLJXVqVNHZcnJyYGuh+goVaqUyhYvXqyySIZNfPPNN8785MmTKrvppptUlieP/nLnesD/3nvvVdmMGTOClIg4lTdvXpVdffXVKps3b57KbrnlFu91M3qfve222wId63rff+utt1T26KOPqsz1PRLgyzVYYsiQIdlfSJzizhQAAAAAeKCZAgAAAAAPNFMAAAAA4IFmCgAAAAA8xN0AihUrVqhs0aJFKnMNY6hZs2ZWlOStWrVq3udu2LBBZQybCLd+/fqpzDVswqVRo0Yqcz0U/+tf/zrzhaUTdNiE66Fs18P8ZcuWVVm5cuVU5nr42pX16dNHZbt27cqwTkSfazjO/PnzVXbjjTcGup5rgIRrbz/zzDPO848ePaqy+vXrq2z58uUqy507d5ASESLXX3+9ykaPHq2y+++/33uNHTt2qOzcuXOBzx81apTK7rrrLpXdfvvtKuvSpYvKli5dqrL33nsvcD3Az3ENm3ANpYi2oUOHZvka0cCdKQAAAADwQDMFAAAAAB5opgAAAADAA80UAAAAAHgwGf2EbufBxgQ/OIpcP5H+tddeU1njxo1Vlh0Pp+fJ457j4XpI1fXwtkupUqVUduDAgcwVFgestSaW68dqz7qsWrVKZbfeemugc1evXq2yli1bqiy79kjt2rVVtm7dukDnnjp1SmUvvPCCyrZv366yDz/8MNAakWDPXt6SJUtU1rRp00DnnjlzRmWdO3dW2Zw5czJf2M9ISUlRmeu9u2PHjipzDcSIM+ustfVitXis9myxYsVUtn79epVVqlRJZa6BTuPHj1fZzp07VbZ48WKVnThxIsM6fVWsWFFlX331lcrOnz+vMtfQos2bN0ensOjIkXs23rkG9Yhkz7CJpKQklbmG0sVQhnuWO1MAAAAA4IFmCgAAAAA80EwBAAAAgAeaKQAAAADw4J6cEGcWLlwYKMsOBQsWVNm0adOcxwYdNjFv3jyV/fjjj5krDDFTvHhxlQ0aNEhl9eoFe9bW9bB7nz59VHb48OFA18sKkTyQ36lTJ5W9//77kZSDLNK/f3+VuQb9uJw7d05lVatWVdmePXsyX9jPcA1IyZVL/93hsWPHVFahQgWVuQYBfP/9957VIVquuOIKlbmGTbgeYm/fvr3K4m3Ik2uPdejQQWV/+MMfVOb6vWnQoIHKXIODTp8+HbREhJhrqER2DJoQcf+ZjLNhE5nCnSkAAAAA8EAzBQAAAAAeaKYAAAAAwAPNFAAAAAB4CMUAilipU6eOyqZOnaqyG2+8MfA19+7dq7KPPvpIZS1atFDZ/PnzA6+D7ON6kNn14L7Ll19+qbJYDptwPbQ8fvx4lVWrVk1l1uofKD9ixAiVMWwiPrkGpIwZM0ZlrkEOFy9eVFnbtm1VlhXDJlyaNm2qMlfdhQsXVpnrYf5hw4apbM6cOSpzDVdB9nINb+rbt6/K4m3YRFCrVq1SmWvYi+v7hfLly6usefPmKluyZIlndYhXQ4cOVdmQIUOyZW3XYImkpKRsWTu7cGcKAAAAADzQTAEAAACAB5opAAAAAPBAMwUAAAAAHhJ+AIUxRmVFixZV2cCBA1XWu3dvlRUsWDCiesqUKaOyyZMnq8z1QPfXX3+tshdffFFlrgf8XcMBEB2uoQ2u3+8333xTZa59l13DJlxatWqlsi5duqjMtT9dD+m/+uqr0SkMWa5QoUIqcw1tcJkwYYLKFixYEHFNQbj+DA0aNCiqa+TPnz9QhtgrVqyYyqZPn66ywYMHq2zhwoVRreX666935q6hVa6vGcWLF1eZa7jRdddd51Fdqpo1a6qMARTh1qRJE5Vl17AJF9f3BomGO1MAAAAA4IFmCgAAAAA80EwBAAAAgAeaKQAAAADwkPADKB599FGV/eUvf4lBJZnjevC7du3aKpszZ47K7r//fpVF+8Fa/J8//elPKnvnnXdUdurUKZWdPHkyS2r6Obfccoszdw1dcXH9RPPZs2erLFafHzIv6P97l48++ijQcXny6C85rofny5cvr7IOHTo4r9myZUuVFShQIFA9kTh06FCWr4HLc72nTps2TWUpKSkqy5cvn8o+//xzlZUuXTpQLa7vNXr16uU8tn379irLjiFRo0ePVtm4ceOyfF1kr+XLl8ds7aSkJJW5vl9INNyZAgAAAAAPNFMAAAAA4IFmCgAAAAA80EwBAAAAgAeaKQAAAADwkPDT/G644YZAx+3YsUNlrsk+WeG9995TWZcuXVTWrl27QNd74YUXVMY0v6xz/vx5lR04cCAGlbjVqFFDZS+99JLz2IYNG6ps//79KhswYIDKvvnmG4/qEAvFihVTWbVq1byv17hxY5W5Jq09//zzKmvatKn3uiIi69evV9lnn32msp07d6rs1VdfDbTGxo0bVfbcc88FOhdZ5+jRoyrr0aOH9/Vce2T48OEqq1ChgsqKFi2qspkzZzrX6devn8pat26tsvHjxzvPv9Tu3btV1qZNG5Vt2LBBZefOnQu0BuJTLCf3uab05YTJfS7cmQIAAAAADzRTAAAAAOCBZgoAAAAAPNBMAQAAAIAHY60NfrAxwQ+OE8YYlRUpUkRlrocwT548mSU1BVGlShWVbdmyJdC5mzZtUlnNmjUjrsmHtVb/D8hGYdyzkbjyyitVtnLlSpXVrl3beb7rz8vf/vY3lT300EMe1YVDTt2zrgfWZ8+eneXrnjhxQmXvvPOOyv761786z//Xv/6lsjJlyqhsxIgRKgs61Kdnz54qmzZtWqBzs8k6a229WC2eyO+zefLoOV25c+dWWUpKSkTruIah9O/fP9C5rVq1UtmHH34YUT3ZgD2bSUOHDlXZkCFDsmXtYcOGqcxVT4LLcM9yZwoAAAAAPNBMAQAAAIAHmikAAAAA8EAzBQAAAAAe9JOVCcY1YOPIkSMxqCRzInnwe+rUqVGsBGFyzTXXqKxw4cIqcw2aEBH54YcfVPbUU09FXhji3oYNG1T2ySefqKxp06aBrnfx4kWVTZgwQWVjx45V2fbt2wOtIeIe1rN48WKV/eIXvwh0Pdd776xZswLXg8Ry/vz5QFlm5M+fX2WNGzdWmet9eseOHSr75ptvIqoH8adJkyYqy45hEytWrHDmOXDYRKZwZwoAAAAAPNBMAQAAAIAHmikAAAAA8EAzBQAAAAAeEn4ARbwpXbq0ylwP+NesWdN7jRkzZnifi3Dr1q2byn75y1+qzDWYRUTkyy+/VNnu3bsjLwxxzzX0oWXLlirr37+/ynbt2qWyjz/+WGWuASdBXXvttc48kmETx48fV9moUaNUduLEiUDXA4IYN26cyurUqaOy/fv3q8w1AGbbtm3RKQxxY/ny5TFZNykpKSbrhh13pgAAAADAA80UAAAAAHigmQIAAAAADzRTAAAAAOCBARRZyPXAtOunSHfs2DHQ9c6ePauyP/3pTyrbt29foOsh3KpUqaKyxx57LNC5GQ0p6dmzZ0Q1IbGkpKSo7A9/+EOWr1upUiWVffLJJ85jK1asGOiarvfPu+++W2VfffVVoOsBP6d48eLO/JFHHlHZ+fPnVTZhwgSVbd26NfLCEFeaNGmS5WusWLFCZcOGDcvydXMK7kwBAAAAgAeaKQAAAADwQDMFAAAAAB5opgAAAADAQ8IPoChfvrzKdu3aFejcvHnzqqxBgwYqa9eunfP8Hj16BLqmtVZl27dvV9mLL76osrfeesu5NhJLwYIFVfb444+rrHDhwoGut2bNGmd++vTpzBUGZIEKFSqoLOigCRGR48ePq2zu3LkqW716deYKAzJQsmRJlS1evDjw+ZMmTVLZkCFDIqoJ8cc1bGL58uVZvq5r2IRrKAX8cGcKAAAAADzQTAEAAACAB5opAAAAAPBAMwUAAAAAHkI7gML1k8UHDhyosqpVq6qsTZs2KmvdurXKfv/736usdu3aQUt0cg2b+O6771SWlJSksn379kW0NsLLNViibdu2KsuXL5/K9uzZo7KZM2dGpzAgC9x3330RnT9q1CiVjR49OqJrAj8pXbq0yubPn6+ym266yXn+iBEjVDZlypTIC0Pccw2giDZjTJavgf/GnSkAAAAA8EAzBQAAAAAeaKYAAAAAwAPNFAAAAAB4CO0Aivbt26vs6aefDnTu2bNnVZY7d26VRfoQ3+bNm1U2fPhwlc2aNUtlFy5ciGhtJJalS5eqrFy5cio7c+aMylq1aqWyvXv3RqcwIEINGjRQWdeuXQOf7xrMM2nSpEhKAv6jTJkyKnv//fdVVq9ePZW9/fbbzmsOHjw48sIAEVmxYkWsS4BwZwoAAAAAvNBMAQAAAIAHmikAAAAA8EAzBQAAAAAeQjuAolatWt7n5snj/2lv3LhRZa+99prz2OnTp6vs3Llz3msjZ2jTpo3KatSooTJrrcpeeuklle3cuTM6hQFZoGfPniorWbJk4PN/97vfqezHH3+MqCbgJ7169VLZzTffrLKtW7eqjEETuNTQoUNV1rhxY5U1adJEZa5hE0lJSVGoCpHizhQAAAAAeKCZAgAAAAAPNFMAAAAA4IFmCgAAAAA8hHYAxR//+EeVnT17VmV169ZV2alTp1Q2f/58lS1ZskRl+/btU9mxY8cyrBO4nHz58qnsgQceCHRuSkqKylauXKmyAwcOZL4wIM58//33znzz5s3ZXAlyklKlSgU6zvU+6/p+AbgUQyTCjztTAAAAAOCBZgoAAAAAPNBMAQAAAIAHmikAAAAA8BDaARSuh4779u0bg0oAf3fccYfKOnToEOjcCRMmqOzTTz+NuCYg1nbv3q2yZs2aOY/du3dvVpeDHOLOO+9UWa9evQKdu3btWpW5hmIBSDzcmQIAAAAADzRTAAAAAOCBZgoAAAAAPNBMAQAAAICH0A6gABJBzZo1Ax23bds2lQ0ePDja5QDZrmvXrrEuARARkUGDBqksTx79bdLXX3+tso8//jhLagIQ/7gzBQAAAAAeaKYAAAAAwAPNFAAAAAB4oJkCAAAAAA/GWhv8YGOCHwyIiLXWxHJ99iwyiz2LEFpnra0Xq8XZs/DAnkXYZLhnuTMFAAAAAB5opgAAAADAA80UAAAAAHigmQIAAAAAD/pHe1/eQRHZkRWFICFVinUBwp5F5rBnEUax3rfsWWQWexZhk+GezdQ0PwAAAABAKv6ZHwAAAAB4oJkCAAAAAA80UwAAAADggWYKAAAAADzQTAEAAACAB5opAAAAAPBAMwUAAAAAHmimAAAAAMADzRQAAAAAeKCZAgAAAAAPOaqZMsYkG2OaBjjOGmOqeK7hfS5wKfYswoY9i7BhzyJs2LPxJUc1U/HKpBptjDmU9hptjDGxrgu4HGNMHWPMZ8aYE8aYH4wxT8S6JuBy2LMIE2NMUWPMW8aY/WmvobGuCbicnLpn88S6AIiIyMMi0lpEaomIFZGPReT/iciEWBYFZMQYc7WILBaR/iIyW0TyiUj5mBYFXAZ7FiE0VkQKiEhlEblGRJYZY3ZYa6fFtCogYzlyz+bIO1PGmFuMMf8wxhwxxuw1xvzZGJPvksNaGGO2G2MOGmNeNsbkSnd+d2PMt8aYH40xS4wxlSIsqYuIvGqt3WWt3S0ir4pI1wiviQQSh3v2SRFZYq1911qbYq09bq39NsJrIoGwZxE2cbhn7xeRMdbaU9baZBF5Q0S6R3hNJBD2bHzIkc2UiFyQ1L+dvFpEbheRO0Wk9yXHtBGReiJSR0RaSdpmMMa0EpHnReQBESkpIp+LyHuuRYwxA9M2uPOV7tAaIrIh3a83pGXAT+Jtz94mIoeNMavSbuV/aIypGLXPFomAPYuwibc9KyJiLvn4hkg+QSQc9mw8sNbmmJeIJItIU0feT0Tmpfu1FZFm6X7dW0SWpX28SER6pPtvuUTklIhUSndulUzWdUFErk/366pp1zGx/j3jFdtXHO/ZLSJyRERuFpErROQ1Efki1r9fvGL/Ys/yCtsrjvfsOyIyV0QKiUgVEdkmIimx/v3iFfsXeza+XjnyzpQx5jpjzAJjzD5jzDERGSmpXX16O9N9vENEyqZ9XElExqXrxg9LauddLoKSTohI4XS/LiwiJ2zazgTicM+eltQ37DXW2jMiMkxE7jDGFIngmkgg7FmETRzu2b6Sum//LSIfSOpdg10RXA8Jhj0bH3JkMyUir4vIZhGpaq0tLKm3OS+dnlch3ccVRWRP2sc7ReQRa23RdK8rrbWrLl3EGPO8SZ0a5XylO3STpA6f+EmttAz4Sbzt2a8l9W+tfkLjj0uxZxE2cbVnrbWHrbUdrbWlrbU1JPV7ttVR/HwRfuzZOJBTm6lCInJMRE4YY64XkcccxzxjjClmjKkgIk+IyMy0fIKIPGeMqSEiYowpYox50LWItXaktfaqjF7pDp0uIk8aY8oZY8qKyFMi8mZUPlMkinjbs9NEpI0xprYxJq+IDBaRldbao9H5dJEA2LMIm7jas8aYa40xJYwxuY0xzSV18u+I6H26SADs2TiQU5upp0Wkg4gcF5HJ8n8bK70PRGSdiHwlIgsldSKJWGvnichoEZmRdkv1f0WkeYT1TBSRD0VkY9r1FqZlwE/ias9aa/8uqX8DtlBE9kvqv43uEMk1kXDYswibuNqzIlJXUr8vOC4io0Sko7WWf7WC9NizccDwWA4AAAAAZF5OvTMFAAAAABGhmQIAAAAADzRTAAAAAOCBZgoAAAAAPNBMAQAAAICHPJk52BjD6D9kirX20h8el63Ys8gs9ixC6KC1tmSsFmfPwgN7FmGT4Z7lzhQAAOG2I9YFAJnEnkXYZLhnaaYAAAAAwAPNFAAAAAB4oJkCAAAAAA80UwAAAADggWYKAAAAADzQTAEAAACAB5opAAAAAPCQqR/aCwBAvGvbtq3KZs2apbKBAweqbPTo0VlSEwAgMXFnCgAAAAA80EwBAAAAgAeaKQAAAADwQDMFAAAAAB4YQAHgPypXrqyyPn36qKxu3boqW758ucqGDRsWlbqAzBg0aJDKrLUqa9CggcoYQAEAyAzuTAEAAACAB5opAAAAAPBAMwUAAAAAHmimAAAAAMADAyiAHCh//vzO/Nlnn1XZww8/rLLJkyerbO3atZEXBkSBMSZQtnTp0uwoBwCQwLgzBQAAAAAeaKYAAAAAwAPNFAAAAAB4oJkCAAAAAA8MoLiM2267TWWtW7dWmeuhfRERa21U65k9e7bKJk6cqLLly5er7OLFi1GtBeHWoEEDZ+4aNjFhwgSV9e/fX2Vnz56NvDAgk0qWLKmyEiVKqCza78cAAIhwZwoQ7kTBAAAQRUlEQVQAAAAAvNBMAQAAAIAHmikAAAAA8EAzBQAAAAAeGECRplWrViqbMmWKyooXL66y7Bru8Jvf/CZQ1rdvX5X95S9/yZKaEE4tWrRw5nv27FHZmDFjVMawCcSLIkWKqKxs2bIqO3funMr+8Y9/ZElNyB6NGjVSWefOnVXmGj7SsWNHlV155ZXetRw6dEhlnTp1ch67aNEi73UAxB/uTAEAAACAB5opAAAAAPBAMwUAAAAAHmimAAAAAMBDjhxAkSeP/rTvvfdelbmGTWSGazDF+fPnVZYvX76I1rnUSy+9FGjdiRMnRnVdxKf77rtPZU888YTz2OHDh6tsx44dUa8JiJbBgwcHOs710P/atWujXQ6yiGs4xAcffKAy10CSoFxfJ13DK1zfQ7i+X1i4cKFzHddAKNd774EDB5znA4gv3JkCAAAAAA80UwAAAADggWYKAAAAADzQTAEAAACAh4QfQOEa7jBs2DCV9ejRI+prv/322yqbNWuWyh599FGVFS5cWGWun/buUqhQIZWVK1cu0LlIPOPHj1fZunXrnMe6HoIG4kG9evWceefOnVWWnJyssueffz7aJSEb5cql/+43kmETr732mspmz56tslOnTqmsWrVqKuvXr5/Kbr75Zufajz/+uMoaNmyosq5du6rsq6++cl4TyCzXnh0yZIjz2KJFi6rsww8/VFn37t1VdvDgQY/qwoU7UwAAAADggWYKAAAAADzQTAEAAACAB5opAAAAAPCQ8AMonn32WZUNGDDA+3pr1qxR2YwZM5zHTpkyRWUnTpxQ2aJFi1RWsmRJlU2cOFFlrVq1cq6NnOmxxx5TmWsvjR07NjvKAbwULFhQZfPmzXMea61V2YQJE1T2zTffRF4YYqZWrVre5y5evFhlTz/9tMrOnz8f6Hrr169X2dy5c1V21113Oc+/4447VOb6/L744guVbd++XWUjR45U2cyZM1V28eJFZz1ILG3btlWZ6/tH11AJY4zzmq732fvuu09lrqEpr7zyivOaiYQ7UwAAAADggWYKAAAAADzQTAEAAACAB5opAAAAAPCQUAMoXA8tt2vXzvt6W7ZsUdmDDz6osp07d3qvkZEDBw6o7PPPP1cZAyhyrrJly6rMNVwlX758KnMNRwFiwfW+vXr1apWVK1fOef78+fNVNnr06MgLQ1wpVaqU97muh++DDpsIKiUlRWULFixwHuvK8+bNq7Lp06errH379ip79913VVahQgWVjRkzxlkPwqFYsWIqc73/uQacZDRY4lJ79+515mXKlAl0vmv4BQMoAAAAAABONFMAAAAA4IFmCgAAAAA80EwBAAAAgIfQDqBwPXzvesi0evXqga63bt06lbVs2VJl+/btC3Q9IKtVrFhRZZUqVVLZ8OHDVXb8+PEsqQnIrD/+8Y8q+9WvfqUya63z/I4dO0a9JsSfpUuXep/bs2dPlbmGQFy4cMF7jUidO3dOZdOmTVOZawCFS69evVTGAIpwmzt3rsrq16+vMtfX9yFDhqhs1qxZKjt06JBz7VtvvVVl77//vsquvPJK5/mJjjtTAAAAAOCBZgoAAAAAPNBMAQAAAIAHmikAAAAA8BDaARR16tRRWYsWLbyvN2XKFJUxbALx7N5771WZ6yF9194GYqFp06Yq69Gjh8qMMSpzPVAvInLy5MnIC0PcO3PmjMqmTp2qsu7du6vM9V7peqB+1apVntUB0eUaLFGvXj2VrVmzRmWtW7dW2d69e1VWoEABld1+++3Oelzv3a4/k59++qnz/ETHnSkAAAAA8EAzBQAAAAAeaKYAAAAAwAPNFAAAAAB4CMUACtdPVH799dcDnbtlyxaV3X///So7cOBA5gsDskmpUqVU9uijj3pfz/XgqYj7Qe3rrrsu0HGuB08/+eQTlS1btixIiUhA48aNU5lraMr06dMDZRkpWbKkylwPb7ssWrQo8DrIPhcuXFCZ6z2wbt26KqtVq5bKRo8erbKnnnpKZatXrw5aYkTKlCmjsldffdX7ekePHo2kHMTYe++9p7KCBQuqbNu2bSp77rnnVFakSBGVub6OFy9ePGiJTq1atVJZ3759I7pmGHBnCgAAAAA80EwBAAAAgAeaKQAAAADwQDMFAAAAAB5CMYDinnvuUVnZsmUDndu+fXuVbd26NeKagOzUu3dvlbkeFF2wYIHKXANXWrRo4VwnozyI2267TWX9+vVTmeth1MmTJ3uvi/jUvHlzlVWvXl1lycnJKuvatavKMhog8corr6isUaNGP19gBj777DOVtW7dWmVHjhzxXgPRcf78eZU9+OCDKnMNvalfv77KFi5cqLK33npLZRs3bgxaYmCu98oaNWqozBijMtcQl9mzZ0enMMRE/vz5Ax330EMPqez7779X2e7du1X23XffqWzevHnOdVyDLubMmaOyoHUnGu5MAQAAAIAHmikAAAAA8EAzBQAAAAAeaKYAAAAAwAPNFAAAAAB4CMU0vwEDBsS6hGxXsWJFlbmmtiBnKFSokMpcU51ck/tatmypsl27djnX+dWvfqWyLVu2BClRmjRpojLXFLSJEyeqbNGiRSrLqEaEwyOPPKIy19SxCRMmqGzKlCkq+93vfudcJ1++fCr79NNPVbZjxw6Vufa7axLgoEGDVPbMM88460Fsuab1Nm7cWGWuCX+/+MUvVPbkk09Gp7Aocf0ZcnnjjTeyuBJkpTp16qgsd+7cgc51TRo9duxYRPVUqlQp0HGHDx+OaJ2w4s4UAAAAAHigmQIAAAAADzRTAAAAAOCBZgoAAAAAPMTdAIpnn31WZTfddFOgc4cPH66yb7/9NuKaYmHp0qUqq1q1aqBzT506pbJDhw5FXBNix/XQcdAHkSdNmqSy//mf/3EeG3TYhMuKFStUdvDgQZV169ZNZa6HbRlAER5VqlRRmWvwiWtoyqhRowIdd+LECefa99xzj8pcAyhc2rZtq7K//e1vKnPtY4RHcnKyyho2bKiy22+/XWW//vWvVeYa9FO6dGmV5crl/vvq5cuXq+yKK65QWf369Z3nBxH06wPi0+7du2Ndwn+pXLlyoOPWrl2btYXEKe5MAQAAAIAHmikAAAAA8EAzBQAAAAAeaKYAAAAAwENMB1CUL19eZa6fPu/6Cfeun2b/2WefqezcuXOe1WWfu+++W2XlypXzvt6IESNUNm7cOO/rITwefvhhlb399tsqO3v2bHaUE9j69etjXQIi0LhxY5UFfQDeddyBAwdUltHD+Fu3bg20TvXq1VU2duxYlZ08eVJlrj9DCLc9e/aobM6cOYGyxx9/XGUlSpRQWUZDo/75z3+qzPV9zsqVK1V28803q4xhE8hq9erVC3Tchg0bsriS+MSdKQAAAADwQDMFAAAAAB5opgAAAADAA80UAAAAAHiI6QCKatWqqaxZs2YqS0lJUVmnTp1U9sUXX0SnsCzkGjYxefJklRUoUCDQ9UaOHKmyV155JfOFISFs2rRJZdk1bKJQoUIqe/3111XmevD7xIkTWVITos8Yo7IHHngg0HGuzDXwISkpSWVBB02IuIdNzJo1S2VFixZVWbdu3VTm2rNAeocOHQqUZcT1Pp2cnKyyoIMAgGjatWtXoOO+/PLLLK4kPnFnCgAAAAA80EwBAAAAgAeaKQAAAADwQDMFAAAAAB5iOoAiqAsXLqgsDMMm7rnnHpXNnj1bZUGHTezfv19lS5YsUZnr9wvh5npwf9++fSrbu3dvROvkz59fZddcc43KevbsqbJBgwapbMGCBSpr2LChZ3WIB9ZalR04cCDQcS5z5sxRWaVKlVTWvHlz5/mu/dSoUSOVHT16VGWtW7dW2bJly5zrAEBO9fTTT6vM9T1I0EEViYY7UwAAAADggWYKAAAAADzQTAEAAACAB5opAAAAAPAQigEU8a5Zs2bO/K9//avKIhk28cgjj6hs5cqVga6HcHM9zF+6dGmV1alTR2U7duxQWYkSJZzrzJw5U2VJSUlBSpRhw4apbPjw4YHORbj9+c9/Vtlvf/tbleXLl09lnTt3VlmXLl1UltFAC9dgiddff11lY8eOVZlrcAYQL06fPq0y1zAi15+Nm266SWUff/xxdAoDROSNN95QWXJycvYXEge4MwUAAAAAHmimAAAAAMADzRQAAAAAeKCZAgAAAAAPoRhAkT9/fpX9/ve/V1l2POx+5513qmzSpEnOY4sUKeK9Trdu3VS2ePFi7+sh3P79738HOs61F5955hmVuf5MiYjUrl1bZX//+99VNnr0aJV98sknQUpEAlq7dq3KevXqpbKXX35ZZVdffbXKPv30U5XNnTvXufaMGTNUxmAJJIJ3331XZZ06dQp0bt26dVXGAAr8nIwGTl1//fUq69OnT1aXExrcmQIAAAAADzRTAAAAAOCBZgoAAAAAPNBMAQAAAICHmA6gcP3k+t27d6usXLlyKnvhhRdU9uyzz6rs4YcfVtm2bdtU1r9/f5Xdd999KsudO7fK8ubNqzIRkVOnTqlsxIgRKhs3bpzKUlJSnNdEzjRlyhSVlS5dWmWDBw9W2a233qqykydPOtdp2rSpyr744guVnT171nk+8JPp06cHygAEZ4xRmbVWZYUKFcqOchBiefLoFsA13E1E5ODBgyrbtGlT1GsKK+5MAQAAAIAHmikAAAAA8EAzBQAAAAAeaKYAAAAAwINxPbiY4cHGBD/Y01133aUy109kdg2biJX58+c78xUrVqjMNWwikVlr9dOy2Sg79iwSC3sWIbTOWlsvVouzZ7PGjTfeqDLXQKCCBQuqbP/+/SqrVauWyn744QfP6iLGno2xNm3aqGzOnDnOY++9916VLVq0KOo1xbkM9yx3pgAAAADAA80UAAAAAHigmQIAAAAADzRTAAAAAOAh7gZQuOTOnVtlAwYMUNmIESOiuq7rAc4OHTqo7Ntvv3Wev2/fvqjWE0Y8zI+wYc8ihHiYP4d44oknVDZ27FiVGaPfxl5++WWVDRw4UGUXL170rC5T2LMxtnHjRpUVKlTIeaxrGMqxY8eiXlOcYwAFAAAAAEQTzRQAAAAAeKCZAgAAAAAPNFMAAAAA4CEUAygQXjzMj7BhzyKEeJg/hyhatKjKxowZo7JOnTqpLH/+/CorWLCgyk6fPu1ZXaawZ7PRL3/5S5WtX79eZe3atXOev3Tp0qjXFEIMoAAAAACAaKKZAgAAAAAPNFMAAAAA4IFmCgAAAAA8MIACWYqH+RE27FmEEA/zI2zYs9lo2bJlKluxYoXKXnzxxWyoJrQYQAEAAAAA0UQzBQAAAAAeaKYAAAAAwAPNFAAAAAB4yBPrAgAAAABErnLlyiq75pprVDZu3LhsqCZn4M4UAAAAAHigmQIAAAAADzRTAAAAAOCBZgoAAAAAPBhrg/8QaGPMARHZkXXlIMFUstaWjGUB7FlkEnsWYRTTfcuehQf2LMImwz2bqWYKAAAAAJCKf+YHAAAAAB5opgAAAADAA80UAAAAAHigmQIAAAAADzRTAAAAAOCBZgoAAAAAPNBMAQAAAIAHmikAAAAA8EAzBQAAAAAe/j+4x3N5ZeguRgAAAABJRU5ErkJggg==\",\n \"text/plain\": [\n \"
\"\n ]\n },\n \"metadata\": {\n \"tags\": []\n },\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"show_img_grid(\\n\",\n \" [train_ds['image'][idx] for idx in range(25)],\\n\",\n \" [f'label={train_ds[\\\"label\\\"][idx]}' for idx in range(25)],\\n\",\n \")\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Training\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# Get a live update during training - use the \\\"refresh\\\" button!\\n\",\n \"# (In Jupyter[lab] start \\\"tensorboard\\\" in the local directory instead.)\\n\",\n \"if 'google.colab' in str(get_ipython()):\\n\",\n \" %load_ext tensorboard\\n\",\n \" %tensorboard --logdir=.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 11,\n \"metadata\": {\n \"outputId\": \"a0eb78b5-ee73-4f4f-8400-41b521f42b75\",\n \"tags\": []\n },\n \"outputs\": [\n {\n \"name\": \"stderr\",\n \"output_type\": \"stream\",\n \"text\": [\n \"INFO:absl:Load dataset info from /root/tensorflow_datasets/mnist/3.0.1\\n\",\n \"INFO:absl:Field info.citation from disk and from code do not match. Keeping the one from code.\\n\",\n \"INFO:absl:Reusing dataset mnist (/root/tensorflow_datasets/mnist/3.0.1)\\n\",\n \"INFO:absl:Constructing tf.data.Dataset for split train, from /root/tensorflow_datasets/mnist/3.0.1\\n\",\n \"INFO:absl:Constructing tf.data.Dataset for split test, from /root/tensorflow_datasets/mnist/3.0.1\\n\",\n \"INFO:absl:train epoch: 1, loss: 0.2352, accuracy: 92.94\\n\",\n \"INFO:absl:eval epoch: 1, loss: 0.0592, accuracy: 98.00\\n\",\n \"INFO:absl:train epoch: 2, loss: 0.0584, accuracy: 98.15\\n\",\n \"INFO:absl:eval epoch: 2, loss: 0.0575, accuracy: 98.14\\n\",\n \"INFO:absl:train epoch: 3, loss: 0.0423, accuracy: 98.66\\n\",\n \"INFO:absl:eval epoch: 3, loss: 0.0357, accuracy: 98.78\\n\",\n \"INFO:absl:Load dataset info from /root/tensorflow_datasets/mnist/3.0.1\\n\",\n \"INFO:absl:Field info.citation from disk and from code do not match. Keeping the one from code.\\n\",\n \"INFO:absl:Reusing dataset mnist (/root/tensorflow_datasets/mnist/3.0.1)\\n\",\n \"INFO:absl:Constructing tf.data.Dataset for split train, from /root/tensorflow_datasets/mnist/3.0.1\\n\",\n \"INFO:absl:Constructing tf.data.Dataset for split test, from /root/tensorflow_datasets/mnist/3.0.1\\n\",\n \"INFO:absl:train epoch: 1, loss: 0.2745, accuracy: 91.83\\n\",\n \"INFO:absl:eval epoch: 1, loss: 0.0478, accuracy: 98.36\\n\",\n \"INFO:absl:train epoch: 2, loss: 0.0508, accuracy: 98.42\\n\",\n \"INFO:absl:eval epoch: 2, loss: 0.0382, accuracy: 98.81\\n\",\n \"INFO:absl:train epoch: 3, loss: 0.0374, accuracy: 98.85\\n\",\n \"INFO:absl:eval epoch: 3, loss: 0.0264, accuracy: 99.09\\n\",\n \"INFO:absl:Load dataset info from /root/tensorflow_datasets/mnist/3.0.1\\n\",\n \"INFO:absl:Field info.citation from disk and from code do not match. Keeping the one from code.\\n\",\n \"INFO:absl:Reusing dataset mnist (/root/tensorflow_datasets/mnist/3.0.1)\\n\",\n \"INFO:absl:Constructing tf.data.Dataset for split train, from /root/tensorflow_datasets/mnist/3.0.1\\n\",\n \"INFO:absl:Constructing tf.data.Dataset for split test, from /root/tensorflow_datasets/mnist/3.0.1\\n\",\n \"INFO:absl:train epoch: 1, loss: 0.2676, accuracy: 91.85\\n\",\n \"INFO:absl:eval epoch: 1, loss: 0.0485, accuracy: 98.57\\n\",\n \"INFO:absl:train epoch: 2, loss: 0.0483, accuracy: 98.54\\n\",\n \"INFO:absl:eval epoch: 2, loss: 0.0461, accuracy: 98.64\\n\",\n \"INFO:absl:train epoch: 3, loss: 0.0341, accuracy: 98.91\\n\",\n \"INFO:absl:eval epoch: 3, loss: 0.0396, accuracy: 98.74\\n\"\n ]\n }\n ],\n \"source\": [\n \"# 3x 3 epochs trains in ~1 minute in the GPU Colab...\\n\",\n \"\\n\",\n \"# We don't use TPUs in this Colab because we do not distribute our\\n\",\n \"# training using pmap() - if you're looking for an example using TPUs\\n\",\n \"# checkout below Colab:\\n\",\n \"# https://colab.research.google.com/github/google/flax/blob/main/examples/imagenet/imagenet.ipynb\\n\",\n \"\\n\",\n \"config.num_epochs = 3\\n\",\n \"models = {}\\n\",\n \"for momentum in (0.8, 0.9, 0.95):\\n\",\n \" name = f'momentum={momentum}'\\n\",\n \" config.momentum = momentum\\n\",\n \" state = train.train_and_evaluate(config, workdir=f'./models/{name}')\\n\",\n \" models[name] = state.params\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 12,\n \"metadata\": {\n \"cellView\": \"form\",\n \"tags\": []\n },\n \"outputs\": [],\n \"source\": [\n \"if 'google.colab' in str(get_ipython()):\\n\",\n \" #@markdown You can upload the training results directly to https://tensorboard.dev\\n\",\n \" #@markdown\\n\",\n \" #@markdown Note that everbody with the link will be able to see the data.\\n\",\n \" upload_data = 'no' #@param ['yes', 'no']\\n\",\n \" if upload_data == 'yes':\\n\",\n \" !tensorboard dev upload --one_shot --logdir ./models --name 'Flax examples/mnist'\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Inference\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 13,\n \"metadata\": {\n \"outputId\": \"3af424f7-4433-475d-817c-5c0bbc4599ae\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"0.0126\"\n ]\n },\n \"execution_count\": 13,\n \"metadata\": {\n \"tags\": []\n },\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"# Find all mistakes in testset.\\n\",\n \"logits = train.CNN().apply({'params': state.params}, test_ds['image'])\\n\",\n \"error_idxs, = jnp.where(test_ds['label'] != logits.argmax(axis=1))\\n\",\n \"len(error_idxs) / len(logits)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 14,\n \"metadata\": {\n \"outputId\": \"949487f5-8aa2-45c8-9b54-efbf34ab58f1\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"image/png\": \"iVBORw0KGgoAAAANSUhEUgAAA1MAAANRCAYAAAAGcOaXAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjIsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+WH4yJAAAgAElEQVR4nOzde7zVc9r/8etKqelEaiuFUigR6TDMOFQynYZu/RJjKodRQjRKdJIGHSZTFCoGk+PcIjWYSEqoUJSkExWddNBOEZUO+/P7Y+/u2XfXZ3d/12evvdfp9Xw81uNR7/09fNZ2WXtdfff3WuqcEwAAAABAbEokegEAAAAAkIpopgAAAAAgAM0UAAAAAASgmQIAAACAADRTAAAAABCAZgoAAAAAAtBMxZGqrlXVSxO9DiAqahaphppFqqFmkWqo2djQTCUBVe2tqltU9UdV/Yeqlk70moCCqOpZqvq2qmarKh9Uh6RHzSKVqeosVXWqWjLRawEKksmvszRTBSiuFy1VbS0i/UWkpYjUFJHaInJfcZwb6aUYf9DuF5GXReTGYjof0hQ1i1RT3A2NqnYWkVLFeU6kF15ni17GNVN5ly4HqOpyVd2hqhNVtYyqNlfVjaraT1W3iMhEVS2hqv1VdY2qblfVl1X1uHzH6qqq6/K+NihwSdeJyNPOuWXOuR0i8oCIXF/4Z4p0kWw165z70jn3tIgsi9dzRHqhZpFqkq1m845zjIgMEZG74/AUkWaSrWYz+XU245qpPJ1FpLWI1BGR00Xknry8mogcJ7lXiG4SkdtF5AoRaSYi1UVkh4iMExFR1foiMkFEuuZ9rbKInHjoBKr6R1XdeYTHyXmbnikin+db2+ciUlVVKxfB80bqSqaaBaKgZpFqkq1mh+cda0tRPWGkvGSr2czknMuoh4isFZGb8/29nYisEZHmIrJPRMrk+9oKEWmZ7+8nSO5lzJIicq+IvJTva+Xy9r80xvWsEZE2+f5eSkSciNRK9PeKR3I8kq1m8+1/au5LSOK/RzyS60HN8ki1R7LVrIg0EZHFeceslfe+oGSiv088kueRbDWbb/+Me53N1JsZN+T78zrJ7cRFRLY55/bm+1pNEZmqqjn5soMiUjVvn/85jnPuZ1XdHrCWn0SkYr6/H/rzroBjIX0lU80CUVCzSDVJUbOqWkJExovIn51zB1Q1lt2RWZKiZjNdpv6a30n5/nyyiGzK+/Ph00c2iEhb59yx+R5lnHPfisjm/MdR1bKSe2n00N87q+pPR3gcuiy6TETOyXfOc0Rkq3OOQkZ+yVSzQBTULFJNstRsRcm9MjUp756XT/J236iqF8X1GSPVJUvNZrRMbaZ6quqJeTffDRKRSQVs97iIDFPVmiIiqpqlqv+V97XJInKZql6oqkeLyP2S7/vpnHvROVf+CI/1eZs+JyI3qmp9VT1Wcn/f9Zm4P2OkuqSpWc1VRkSOzvt7GWWcPyxqFqkmWWr2B8m9WtAw79Eub/fGIjI/zs8ZqS1ZajajX2cztZn6p4jMEJGvJff3S4cWsN1YEXldRGao6i4R+VhEzhMRcc4tE5GeecfaLLk3822MdSHOueki8qCIzBaR9ZJ7mXZIrMdB2kuampXcXxfYI/+Z2LNHRL4MOA7SGzWLVJMUNetybTn0EJFteV/a6pzbF9tTQppLiprNk7Gvs+pcRn2ulqjqWhHp5pybmei1AFFQs0g11CxSDTWLVEPNJo9MvTIFAAAAAIVCMwUAAAAAATLu1/wAAAAAIB64MgUAAAAAAWimAAAAACBAyVg2VlV+JxAxcc4l9KPbqVnEippFCsp2zmUl6uTULAJQs0g1BdYsV6YAAEht6xK9ACBG1CxSTYE1SzMFAAAAAAFopgAAAAAgAM0UAAAAAASgmQIAAACAADRTAAAAABCAZgoAAAAAAtBMAQAAAEAAmikAAAAACEAzBQAAAAABaKYAAAAAIADNFAAAAAAEoJkCAAAAgAA0UwAAAAAQgGYKAAAAAALQTAEAAABAAJopAAAAAAhAMwUAAAAAAUomegEAAABHkpWVZbKbb77ZZB07djRZhQoVTDZ79myTdevWLXB1ADIZV6YAAAAAIADNFAAAAAAEoJkCAAAAgAA0UwAAAAAQgAEUAIBCqVKlislmzJhhsjPOOMNkH3/8caR9R48e7T33vn37oiwRKW7nzp0m69Kli8lOPfXUSMe74YYbTLZ69WqTvfHGGyZbtmxZpHMARa13794mu+KKK0yWnZ0daTtV9Z7HOWeyuXPnmqxPnz4mW7hwofeY6YQrUwAAAAAQgGYKAAAAAALQTAEAAABAAJopAAAAAAigvpvKCtxYNfrG8N7IV758eZNVqlTJZD169DDZ1VdfbbLatWub7O677zbZY489ZrK9e/eaLN6cc/67GYsJNetXr149b/7ZZ5+Z7IsvvjDZ+eefb7KcnJzCLywJULOx89VTYW7S9712FjSA4q677go+TxpZ6JxrkqiTJ6pmfXV34403muzWW281WZkyZUzmq7s9e/ZEOt6kSZNMVhw/Y1NYRtZsvPleZ+vWrWsyX2373v/HMoDCt+22bdtMdsstt5hs6tSp3vMkuQJrlitTAAAAABCAZgoAAAAAAtBMAQAAAEAAmikAAAAACMAAijioVq2aNx8yZIjJbrrpJpP98MMPJvPd9Oq7wfWbb74xme9mv4MHD5rsggsuMNmPP/5ossLgZv7EO+GEE0w2ffp077a1atUyme+G7smTJxd6XcmKmo1dqVKlTFanTh2T9evXL9LxrrvuOpPNnDnTu+0VV1xhst27d0c6TxrhZv4jaNiwocmeffZZkzVo0MBkUd8jTZkyxWR9+/b1brtu3bpIx0xz1GwcdOjQwWQPPfSQyapUqWKylStXFurcjRs3Npnv/5dFixaZrGnTpoU6d4IwgAIAAAAA4olmCgAAAAAC0EwBAAAAQACaKQAAAAAIwACKI6hYsaLJhg4darIePXp49y9ZsqTJFi9ebLKrr77aZKtXr46yRK+rrrrKZP/93/9tsho1aphsy5Ytwef14Wb+4tWsWTOT9e/f32StW7f27t+pUyeTvfrqq4VfWAqhZhMvJyfHZAX9rGrTpo3J3nnnnbivKclxM38cPPXUUya74YYbgo9X0KAe38/8DETNFhHfsImiGEDhq2/fQCDfQLVrr73WZFOnTi3UeooBAygAAAAAIJ5opgAAAAAgAM0UAAAAAASgmQIAAACAAHZCQoY6//zzTfbII4+YzPeJz1u3bvUe84UXXjDZ3XffHbC6wvv+++9Ntn///gSsBPkde+yxJtu5c6fJVO1MhF69eplsyJAhJvPVcd26db3rqV69ujcHitO4ceNMduutt3q3/eMf/2iyDBxAgTjo1q2byY4//niTXXbZZZGO5xvoI+J/Pe/cubPJ+BmNENnZ2ZGyWGRlZZns//2//2cy36CgcuXKmaxVq1YmS4EBFAXiyhQAAAAABKCZAgAAAIAANFMAAAAAEIBmCgAAAAACZOQAinPOOcdkvhv3fcMmfJ5++mlvPnjw4NgWFidXXXWVyV577TWTbd++vTiWgzy+TyAfNWqUyXyDS3y1dMopp5isa9euJps2bZrJWrZsWeA6gUQbPXq0ya6//nrvtr4bmU866SSTbdiwodDrQua57bbbTFanTh2T1atXL/Ixr7zySpPt2LHDZD169Ih8TKAoPffccybzDZvwZT6pPGzChytTAAAAABCAZgoAAAAAAtBMAQAAAEAAmikAAAAACJBWAygqVqxosj/84Q8mGzlypMl2794dabubbropcHVF4/zzzzeZ74bs9u3bF8dycAQ9e/Y0WZcuXUx23XXXmezf//63yS655BKTrV69OtJaPvjgg0jbAYmwdu1ak/lugBYRufnmm01WtmzZeC8JGWr9+vUm69+/v8mefPJJk2VlZUU+T4MGDWJbGFAEChq81qhRI5OpaqRjLlq0yGQzZsyIbWFJjitTAAAAABCAZgoAAAAAAtBMAQAAAEAAmikAAAAACJBWAyhq165tsgkTJphs9uzZJhs0aJDJtm7dajLfzc6JdOedd5rM90nqK1euLI7l4Ah8n/j9008/mWzatGkmW7VqlckOHjwYn4XlU6IE/76C5LRixYpELwEQEZE33njDZMOGDTPZmDFjIh/zzDPPNFnbtm1N9tZbb0U+JhCrhx56yJtXrlzZZM65SFkm4J0TAAAAAASgmQIAAACAADRTAAAAABCAZgoAAAAAAqTsAIqKFSuabOTIkZH2/ctf/mKy4447zmSvvfaayb799luTjRs3LtJ5C+vGG280Wbt27Uzmuzl2y5YtRbImRLdkyZJIWXGYP3++N+/bt6/JHnnkEZNl6k2mAODz+uuvm6xLly7ebZs2bWqy8uXLm+zhhx82GQMoEC9ZWVkmu+iii7zb+n7mq2qk83Tt2jW2haUgrkwBAAAAQACaKQAAAAAIQDMFAAAAAAFopgAAAAAgQMoOoBg6dKjJLr30UpP5hkP4hjE8+eSTJtu2bZvJWrVqFel4heW7MfCOO+4wWZkyZUz26KOPxn09SC/vvfeeN3/qqadMdswxx5hs586d8V4SAKSEU0891WSnnXaaydatW+fd/9e//nWk85QrV85kxx57rMl4Pcb/xfee8s033zRZQcOlog6dmjJlislWrlwZad9UxpUpAAAAAAhAMwUAAAAAAWimAAAAACAAzRQAAAAABKCZAgAAAIAAKTHNr1q1aibr0aOHyZYsWWKy+fPnm2zatGkmO3DggMl+97vfmawoJvf51KhRw2T169c32X333Weyjz/+uEjWhPTx448/evNffvnFZC1atDDZ1KlT474m4EiGDx/uzVU1UobMVbKkfavz2GOPmaxJkyYmq1mzpsl8U3TLli0beT2+yWi+rHr16ib78ssvTfbSSy+Z7KOPPjLZF198YbJly5YVuE6kj169epmsUaNGJovltXP37t0mGzx4cGwLSxNcmQIAAACAADRTAAAAABCAZgoAAAAAAtBMAQAAAEAA9d30WODGqtE3jqMTTjjBZN9++63JfDfO+Z7fnj17THbnnXea7PHHH4+6xELxPb9Zs2aZ7KijjjLZpZdearINGzbEZ2Fx4JxL6J3giarZVPXQQw+ZrFmzZiZr2rSpyXJycopkTcWNmk1OP/zwgzcvX768yXzDenw37qeRhc45Oz2hmCR7zTZs2NBkH3zwgcl8tRTLe6Soor5Xifc5fO99+vXrZzLfgCHfe65ComaLSIcOHUw2efJkk/lqrqABFL5tfcMmRowYEWWJqarAmuXKFAAAAAAEoJkCAAAAgAA0UwAAAAAQgGYKAAAAAALYjwVPQgcPHjTZzz//bLKjjz7aZNnZ2Sbz3Zy3YMGCwNXFpmLFiibz3fRft25dk/3+9783WTINm0Di+W609tV7pUqVvPvXqFHDZOeee67J1q1bZ7Jrr73WZLNnz/aeBziSs846y2SlSpXybrtmzRqT/fjjj3FfE1LX4sWLTeYbouN7/TzuuONMVrp0aZO1bt3aZC1atPCux/deJaq9e/ea7MCBAyarUKGCycqUKWOysWPHmuzqq6822UUXXRR1iShG5cqVM9nQoUNNVtBgiajb+YaSpPmwiZhwZQoAAAAAAtBMAQAAAEAAmikAAAAACEAzBQAAAAABUmIAxXfffWeytm3bmsx3o+jrr79eJGuKwnez57x580xWv359k61fv95kS5cujc/CkBa6dOlisokTJ5rsqKOOinzMmTNnmsx3c/Pxxx9vslWrVkU+D3Ak7777rsl8N/2L+F9TN2/eHPc1Ib18+eWXkbKoxowZY7JHH33Uu227du1MVqtWLZPt2rXLZJdddpnJvv76a5OdffbZ3nNH2S7qsAIkXv/+/U3mG2DmnIuU+Ya2iYj06dMnYHWZgytTAAAAABCAZgoAAAAAAtBMAQAAAEAAmikAAAAACJASAyh85s6dm+gl/J9uv/12k/mGTfzyyy8mGz16tMk2btwYn4UhLbzyyism89WSzzvvvOPNf/jhB5MNGTLEZP369TMZN/0jXqpUqWIy383SQDKbP3++N7/ttttM5qtv3/Cgbdu2mWzTpk2RMp/p06dH2g7JqWPHjibzDRCJOlSkWbNm3tw3FA3/wZUpAAAAAAhAMwUAAAAAAWimAAAAACAAzRQAAAAABNBYbupVVe4A9rjrrru8+ciRI022evVqkz3yyCMme+yxxwq/sCTgnEvoR6lTs0Vj//79JrvyyitN9tprrxXHcuKKmk28nJwckxX0s8p3w3QqDCiKs4XOuSaJOjk169emTRtv/uabb5rMV9+rVq0yWb169Qq/sORAzR5Bhw4dTDZw4ECTNWrUyGS+WvINoFi+fLnJGjRoEHWJmajAmuXKFAAAAAAEoJkCAAAAgAA0UwAAAAAQgGYKAAAAAAKUTPQCUk2/fv1MNnToUO+2Bw8eNFnfvn1N9vrrrxd+YUAx+sc//mGyv/71ryZLxQEUSDzfzdIFDaA47bTTTJaBAyiQhKZPn+7NlyxZYjJu/Ed+vuElvmETJUrYayK+AT67d+822b333hu4OhyOK1MAAAAAEIBmCgAAAAAC0EwBAAAAQACaKQAAAAAIwACKI/DdPO+7KdB3A6CISM+ePU3GsAmkgyFDhphs1apVJjvrrLNMtnTp0iJZE9KH78b9Vq1aebd9++23i3o5QFzt27cv0naffPKJyapUqWKy7OzsQq8JycU3cMeX+YZN+LZbuXKlyaZOnRq4OhyOK1MAAAAAEIBmCgAAAAAC0EwBAAAAQACaKQAAAAAIkJEDKE488UST3X777SZr3bq1yb766iuTderUyXse3w1/QDrw3fC8ZMkSk40ePdpkvv+vgPyaN2+e6CUARcb3Wtm4cWOT/fTTTyZj2ERmmDt3rsm6d+9uMt8ANN9Qin/961/xWRi8uDIFAAAAAAFopgAAAAAgAM0UAAAAAASgmQIAAACAAOr7pOQCN1aNvnGSqFGjhsmmTZtmsgYNGphs+fLlJmvVqpXJNm/eHLi69Oec00SePxVrNlVVrVrVZAsWLDDZhRdeaLINGzYUyZpCULOJN2DAAJMNHTrUu+1JJ51ksk2bNsV9TUluoXOuSaJOTs0iADV7BGXLljWZ73Vx4MCBJpsyZYrJChqUhpgUWLNcmQIAAACAADRTAAAAABCAZgoAAAAAAtBMAQAAAECAtB9AUatWLZO9//77JpsxY4bJBg8ebLItW7bEZV2Zgpv5M9ujjz5qsiFDhpjs+++/L47lRELNIgVxMz9SDTWLVMMACgAAAACIJ5opAAAAAAhAMwUAAAAAAWimAAAAACBA2g+gQGJxMz9SDTWLFMTN/Eg11CxSDQMoAAAAACCeaKYAAAAAIADNFAAAAAAEoJkCAAAAgAA0UwAAAAAQgGYKAAAAAALQTAEAAABAAJopAAAAAAhAMwUAAAAAAUrGuH22iKwrioUgLdVM9AKEmkVsqFmkokTXLTWLWFGzSDUF1qw654pzIQAAAACQFvg1PwAAAAAIQDMFAAAAAAFopgAAAAAgAM0UAAAAAASgmQIAAACAADRTAAAAABCAZgoAAAAAAtBMAQAAAEAAmikAAAAACEAzBQAAAAABaKYAAAAAIADNVByp6lpVvTTR6wCiomaRaqhZpBpqFqmGmo0NzVSCqerjqvpTvscvqror0esCCqK5hqrqt6r6g6q+p6pnJnpdQEFUtbSqPqyqm1R1h6qOV9VSiV4XEIWqzlJVp6olE70WoCCZ/H6WZqoAxfWi5Zy72TlX/tBDRP5bRF4pjnMjvRTjD9pOIvInEblIRI4TkY9E5PliOjfSSDHWbH8RaSIiZ4nI6SLSSETuKaZzI40Ud0Ojqp1FhMYfwXg/W/QyrpnKu3Q5QFWX5/0L5URVLaOqzVV1o6r2U9UtIjJRVUuoan9VXaOq21X1ZVU9Lt+xuqrquryvDYrD2sqJSEcRebawx0L6SMKaPUVE5jrnvnbOHRSRF0SkfhyeKtJEEtbs5SLyiHPue+fcNhF5RHL/QQAQkaSsWVHVY0RkiIjcHYeniDSTjDWb73gZ9X4245qpPJ1FpLWI1JHcf6U89C+U1ST3X9prishNInK7iFwhIs1EpLqI7BCRcSIiqlpfRCaISNe8r1UWkRMPnUBV/6iqO4/wONmzro4isk1EPoj3E0bKS6aafUlE6qjq6Zr7q1LXicj0InzuSE3JVLMiInrYn0/Me7MKHJJsNTs871hbiuoJI+UlW80eklnvZ51zGfUQkbUicnO+v7cTkTUi0lxE9olImXxfWyEiLfP9/QQR2S8iJUXkXhF5Kd/XyuXtf2kh1jZLRP6S6O8Rj+R6JFvNisjRIjJWRJyIHBCRb0TklER/n3gkzyMJa3aoiMwTkSzJfZMxP69+T0j094pHcjySsGabiMjivGPWyqvXkon+PvFInkey1exha8uo97OZejPjhnx/Xie5nbiIyDbn3N58X6spIlNVNSdfdlBEqubt8z/Hcc79rKrbQxeU19k3F5HuocdAWkummr1XRJqKyEmS+y+mXUTkXVU90zm3O+B4SE/JVLPDRORYyX1z+ouIPCki54rI1oBjIX0lRc2qagkRGS8if3bOHVDV/2sXZK6kqNn8MvH9bKb+mt9J+f58sohsyvuzO2y7DSLS1jl3bL5HGefctyKyOf9xVLWs5F4aPfT3zvq/p5oc/jj8smhXEZnnnPs6Xk8SaSWZarahiExyzm10zh1wzj0jIpWE+6bwvyVNzTrn9jjnbnPO1XDO1RaR7SKy0DmXI8B/JEvNVpTcK1OT8u55+SRv942qelFcnzFSXbLUbH4Z9342U5upnqp6Yt7Nd4NEZFIB2z0uIsNUtaaIiKpmqep/5X1tsohcpqoXqurRInK/5Pt+OudedPmmmnge6w8717Ui8kwcnyPSSzLV7Cci0klVq2ruTa1dJXfa1Or4P22ksKSpWVWtoarVNdf5IjJYcm/sB/JLlpr9QXKvFjTMe7TL272x5P6KKnBIstRsfhn3fjZTm6l/isgMEflacn+/dGgB240VkddFZIbmzsr/WETOExFxzi0TkZ55x9osuTfzbQxZjKr+RnJv9suIEZIIkkw1O1JEPpfcX5naKSK9RaSjc25nwLGQvpKpZuuIyIci8rPkTpfq75ybEXAcpLekqFmXa8uhh+TeyC8istU5ty+2p4Q0lxQ1e0imvp9V5w6/EpjeVHWtiHRzzs1M9FqAKKhZpBpqFqmGmkWqoWaTR6ZemQIAAACAQqGZAgAAAIAAGfdrfgAAAAAQD1yZAgAAAIAANFMAAAAAEKBkLBurKr8TiJg45xL60e3ULGJFzSIFZTvnshJ1cmoWAahZpJoCa5YrUwAApLZ1iV4AECNqFqmmwJqlmQIAAACAADRTAAAAABCAZgoAAAAAAtBMAQAAAEAAmikAAAAACEAzBQAAAAABaKYAAAAAIADNFAAAAAAEoJkCAAAAgAA0UwAAAAAQgGYKAAAAAALQTAEAAABAAJopAAAAAAhAMwUAAAAAAWimAAAAACAAzRQAAAAABKCZAgAAAIAAJRO9AADxU7p0aZP99re/Ndm9997r3b958+Ymy8nJMdmoUaNMds8995hs//793vMAAACkA65MAQAAAEAAmikAAAAACEAzBQAAAAABaKYAAAAAIIA656JvrBp9Y0BEnHOayPOnS81Wr17dZNdee63JLrnkkkhZQVTtf66orxENGzY02dKlSyOfO1lQs0WnQoUKJrvtttsi7duqVStvfv7555vsoYceipRt37490rlTwELnXJNEnTydazaqJk3st3/kyJHebVu0aGGyefPmmezyyy832c6dOwNWl5SoWaSaAmuWK1MAAAAAEIBmCgAAAAAC0EwBAAAAQACaKQAAAAAIkNABFJUrVzZZlSpVTPbzzz+bbNOmTSbLycmJz8IQN9zMH7vmzZub7B//+IfJTj755LifuzADKIYNG2ayIUOGFHpNxY2ajY+6deuabMGCBSYrV65cpOP5alMken3u3r3bZAMGDDDZuHHjIh0vyXAzfzHyvfa+9957JqtVq5Z3/6g16xvOMmHChEj7pgBqFv9L7dq1TdaxY0eTnXvuuSbz/WwZM2ZMfBb2HwygAAAAAIB4opkCAAAAgAA0UwAAAAAQgGYKAAAAAAKUTOTJJ0+ebLJKlSqZbM6cOSZbtGiRySZOnBifhQFFoHr16ibz3WDcp08fk5UqVcpksQyPOZzv/z0RkTPPPNNkZ5xxRqRj1qxZM3g9SG2+wUHjx483WdRhE0WhbNmyJnvwwQdN1rp1a5O1b9++SNaE5HfWWWeZbNKkSSaL5fVv165dJqtQoYLJ9u3bF/mYQLKqU6eOyfr27Wuy7t27m6xECXvNZ+3atSabOnVq2OLihCtTAAAAABCAZgoAAAAAAtBMAQAAAEAAmikAAAAACJDQARRlypQxWYMGDUy2ZMkSk/k+Bdx3E/Tf/va3wNUB8fXYY4+ZLN43to8aNcpkDz30kMmys7O9+2dlZZls5syZJvMNpfB9Uvldd91lsm3btnnPjdRw/PHHm+zFF180WbNmzYpjOYVSunRpk/l+jiAznHrqqSabNWuWyXw18ssvv5jshhtu8J7n008/NVmTJk0infvKK680mW+40TPPPGOyH3/80bseIMQ555xjsieeeMJkjRo1MlnJkrb92L9/v8l87/VHjx5tMt9QiuLElSkAAAAACEAzBQAAAAABaKYAAAAAIADNFAAAAAAESOgAiokTJ5rsvPPOM1mXLl0iHe8vf/mLyS6//HKT+W4KXb9+vclq165tsp9//tlkGzdujLQ+ZA7fjcwNGzY0mapGOt53331nspYtW5ps+fLlkY5XkK1bt5ps7NixJvPdZFq2bFmTtWjRwmQvv/xy4OqQDDp06GAy33/nqA4cOGCyQYMGmez999/37u+7Ib9v377B60FmqFGjhskWLVpksnLlykU6Xs+ePU02adKkyOtZs2aNyW6//XaTjRkzxmS+nyM5OTkm8w1BAvLzvXcR8ddO8+bNTXb00UebzFefq1atMln37t1NVtDrfrLhyhQAAAAABDZduXYAACAASURBVKCZAgAAAIAANFMAAAAAEIBmCgAAAAACJHQAxccff2yyDRs2mOykk06KdLxf/epXJrvwwgtNNm3aNJP5bnLz3Qznu0n0X//6l3c9vk8vnz59unfbePI958aNG5vMN1gAsTvmmGNM9tZbb5ns5JNPNplzzmTbtm0zWVEMm4jq1VdfNZnvxugzzzzTZL66YwBFavMN8CmMr776ymS+T7gviG9oEfB/uf/++01Wvnz5SPtec801Jotl2ISP78b/W2+91WS+m/l92Ycfflio9SD91alTx2RPPvmkd9tmzZpFOuaLL75ospkzZ5rslVdeMdnu3bsjnSMZcWUKAAAAAALQTAEAAABAAJopAAAAAAhAMwUAAAAAARI6gGLJkiUma926tcnuu+8+k3Xs2NFkJUpE6w1PP/30SJmP74a9O++807vtnj17TLZ//36TTZkyxWQrV640WfXq1U1WsWJFk5UqVcpkTz/9tHeNKDzf4JNTTjkl+Hi+G6OLa9iEz4EDB0w2b948k/kGUAD5rVixwmTt27cv1DE7d+5cqP0P5xsAg9TWpEkTk/kGqfgGAo0fP95kvqE8sfANuvCdx/e+xLdG3w3+ixcvDlwd0lHt2rVN5ntfePHFF3v3972fvfbaa032+uuvm8z3vjfdcGUKAAAAAALQTAEAAABAAJopAAAAAAhAMwUAAAAAARI6gMLHN3jh6quvNlnLli1N1q9fP5OddtppJitXrpzJqlSpEnWJkR1zzDEm27Fjh8neffddk/luRn3iiSdMVqFChUhrKehTrVF4f/3rX4P3/frrr03m+wTxeCtoWMQdd9xhsgsvvNBkvv+vfOrXr2+yTp06mcz3aehIP2vWrDHZ2rVrI+17+eWXe/Nzzz23MEsyxowZE9fjoXiVLVvWZI8//nikfffu3WuyUaNGmcw3lMenWrVq3tz389j3niaqhQsXmiwnJyf4eEhtHTp0MNmgQYNMlpWVZbIhQ4Z4jzl37lyTzZ49O2B16YkrUwAAAAAQgGYKAAAAAALQTAEAAABAAJopAAAAAAiQdAMoopo1a1akzMd3U7zvhr1LL73UZFWrVjVZLAMDbr/9dpOVLGn/M7z99tuRjqeqJvN9QjqKzhlnnBG874cffmiyH3/8sTDLMXzDJt555x3vtr76Lkw9tW3b1mRt2rSJlD3wwAMmizqsAPFxzjnnmOzkk08OPt6GDRuC923cuLE3L1WqVPAxv/zyS5OtWrUq+HhIvIYNG5os6pCSoUOHmmzdunWR9j377LNN9tRTT3m3LaiWQy1dujSux0Pq6N69u8lGjx5tsqlTp5rsb3/7m8mopTBcmQIAAACAADRTAAAAABCAZgoAAAAAAtBMAQAAAECAlB1AURjLly+PlI0dO9ZkvmERO3fujHzuVq1amWz69OmR9z+cbzjA008/bTLf80N8fPDBByZr0qRJpH19A0QKw3fztW+YSZUqVbz7lyhh/30lJyen8AvL56ijjjLZ9ddfHyl74YUXTHbdddfFY1nwqF27tsmOP/744OMdc8wxJitTpozJBg8ebLL+/ft7j1mYASkbN26MlCF1dOzYMXjfJ554wmQVK1aMdI4RI0aYrKD/VwpTs998843JZs+eHXw8pLYBAwaY7P333zfZ+PHjTcawifjhyhQAAAAABKCZAgAAAIAANFMAAAAAEIBmCgAAAAAC0EwBAAAAQICMnOYX1U8//RRpu0qVKnnz0qVLR9rWN9nnu+++M5lvAptvkqBvCuGBAwe8a0Th+f77RZ3WVJipTr/61a9MNm7cOJNVrlw58nn37dtnsvfee89kzz//fIQVijRu3Nhk11xzjckKmi54uN/+9rcmO/XUU022evXqSMfDkU2dOtVkCxcuNFnU6ZWdO3eOlPn4Jk2KFG7a5KBBg4L3RXK64IILTOabmurLJkyYYLI2bdqYrEKFCpHWsmTJEm/um2pZs2bNSMf86KOPTLZp06ZI+yK1+aarZmVlmcz3HnD+/PlFsibk4soUAAAAAASgmQIAAACAADRTAAAAABCAZgoAAAAAAjCAIkann366yV577TXvtnXr1o10zF9++cVk48ePN1mvXr1M1rt3b5MtXbo00nkRH6+++qrJ+vTpE2nfli1bmqxhw4YmW7x4scl8tXjeeedFOm9BBg4caLLRo0cHH++FF14w2bPPPmuyRx991GS/+c1vTHbKKaeY7O233zZZnTp1oi4RMSrMwJXCKGjQRHGcG6mjf//+Jps1a1akfa+88spI223evNlkzzzzjMnuv/9+7/4PPvigyW677bZI5/70008jbYf0s3XrVpN99tlnJtuwYUNxLAf5cGUKAAAAAALQTAEAAABAAJopAAAAAAhAMwUAAAAAARhAEaOOHTuaLOqgiYL4hkjUqFHDZJUrVw4+x29/+1uTlS9f3mQzZswIPkemWrlypcm++eYbk/mGJ5xwwgkmmz59usmqVasWuLrYFGbYRFS+YRqdO3c22bx580zm+37VrFkzPgtDJCNHjjTZK6+8koCVANacOXNMdvPNN5usb9++JsvOzjaZ72fipEmTTOb7OVCQBg0aRNpu9+7dJps2bVrk8yC9/Pzzzyb7/PPPTda2bVuTffnll0WyJuTiyhQAAAAABKCZAgAAAIAANFMAAAAAEIBmCgAAAAACMIDiCK655hqT3XvvvQlYSS7fp6736NHDZL6bFCdOnGiyUqVKmaxcuXKBq8tcP/zwg8kGDx5ssmHDhpnMNzyhSpUqJluzZo3JBg0aFHWJkd1yyy0mmzBhQtzPc7gWLVqYLOrAleJYH/7Dd0N+69atTearJZ+zzz7bZLVr1459YQF8w398w1CQOg4ePGiyJ598MlIWb7Vq1fLmTZs2jbS/b1jP6tWrC7MkpJkNGzaYrLBD0RA7rkwBAAAAQACaKQAAAAAIQDMFAAAAAAFopgAAAAAgAAMo8rRp08Zk9evXN1np0qXjfu7x48eb7P333zeZb2DEhRdeaLILLrjAZFu2bDFZhw4doi4RMXrppZdMdsUVV5jsxBNPNFnJkvZ/S98N+S+++GLg6gr24IMPmmzlypUmW7FiRaTjNWnSxGS+T2e/+eabIx1v7969JnvzzTcj7Yv48A24mTVrVqTMp169eiZbunRp7AsL4Bv2AsRLQUNYypYtG2n/hQsXxnM5SENfffWVyXzvXVG0uDIFAAAAAAFopgAAAAAgAM0UAAAAAASgmQIAAACAAAygyOMbQNG1a9diObfvZv4ZM2aYbO7cuSZ7/vnnI53j6quvNtn8+fMj7Yv4+MMf/mCyYcOGmaxfv34my8nJMZlzLj4Ly8d3Y/Q777wT13OoqsmiPhffsJa33nqr0GtC4mzatClh527UqJHJGjZsaLLFixcXx3KQZqpXr+7Nfa+BPps3b47ncpCGeL1KDlyZAgAAAIAANFMAAAAAEIBmCgAAAAAC0EwBAAAAQAAGUOR5+OGHTdazZ0+T+QYB7N+/33vMf//73yZ77bXXTDZnzhyTrVu3znvMUA888EBcj4f4eOyxx0z2008/max79+4mO/nkk4tkTcli+/btJhs3blwCVoJ0ValSpUgZEE++gTt79+41me89BDJXrVq1TOYbbHXnnXcWw2qQH1emAAAAACAAzRQAAAAABKCZAgAAAIAANFMAAAAAEIABFHl8Ax86dOhgMt8nmv/9738vkjUh/fk+4X7EiBEme+aZZ0x27bXXmuySSy4xWcuWLcMWV4x27NhhslGjRpls7dq1xbAaACi89u3bR952w4YNJlu2bFk8l4MU169fP5N98sknJps2bVpxLAf5cGUKAAAAAALQTAEAAABAAJopAAAAAAhAMwUAAAAAARhAcQR8+jiShW9QxciRI002ZswYkx1zzDGRz9O8eXOTNW7cOPL+UXzwwQcm++ijj0z2/fffx/W8SE7OOZPt27fPZKVLl477ub/88kuTrVq1Ku7nQWaqUKGCN49a80htRx11lMlKlrRvu0uUsNc1Lr/8cpO1a9fOZF9//XXg6hBPXJkCAAAAgAA0UwAAAAAQgGYKAAAAAALQTAEAAABAAAZQAGnkl19+Mdl3330Xef+XX345UgbEy65du0zWtm1bk7377ruFOs+yZctMNmLECJNt3LixUOcBDpk9e7Y3P+uss0w2ZMiQol4Oiln79u1NNn78eJNlZWWZTFVN9vnnn5vshhtuCFwd4okrUwAAAAAQgGYKAAAAAALQTAEAAABAAJopAAAAAAigvk/iLnBj1egbAyLinLN3URYjahaxomaRghY655ok6uTULAJkZM3OmTPHZDVq1DDZTTfdZLKZM2cWyZoQWYE1y5UpAAAAAAhAMwUAAAAAAWimAAAAACAAzRQAAAAABCiZ6AUAAAAA6e6iiy5K9BJQBLgyBQAAAAABaKYAAAAAIADNFAAAAAAEoJkCAAAAgAA0UwAAAAAQgGYKAAAAAALQTAEAAABAAJopAAAAAAhAMwUAAAAAAUrGuH22iKwrioUgLdVM9AKEmkVsqFmkokTXLTWLWFGzSDUF1qw654pzIQAAAACQFvg1PwAAAAAIQDMFAAAAAAFopgAAAAAgAM0UAAAAAASgmQIAAACAADRTAAAAABCAZgoAAAAAAtBMAQAAAEAAmikAAAAACEAzBQAAAAABaKYAAAAAIADNVByp6lpVvTTR6wCiomaRaqhZpBpqFqmGmo0NzVSCqer1qnpQVX/K92ie6HUBBVHV61R1oar+qKobVfVBVS2Z6HUBUajqLFV11CySnar2VtUtea+1/1DV0oleE1AQzTVUVb9V1R9U9T1VPTPR6yoONFMFKOYftB8558rne7xXjOdGmijGmi0rIneISBUROU9EWopI32I6N9JIcTc0qtpZREoV5zmRXoqrZlW1tYj0l9zX15oiUltE7iuOcyO9FOPrbCcR+ZOIXCQix4nIRyLyfDGdO6EyrpnKu3Q5QFWXq+oOVZ2oqmVUtXnev7L3U9UtIjJRVUuoan9VXaOq21X1ZVU9Lt+xuqrquryvDUrg00IaS7aadc5NcM7Ncc7tc859KyIvisgFcXq6SAPJVrN5xzlGRIaIyN1xeIpIM0lYs9eJyNPOuWXOuR0i8oCIXF/4Z4p0kYQ1e4qIzHXOfe2cOygiL4hI/Tg81aSXcc1Uns4i0lpE6ojI6SJyT15eTXK76ZoicpOI3C4iV4hIMxGpLiI7RGSciIiq1heRCSLSNe9rlUXkxEMnUNU/qurOIzxOzreec1U1W1W/UtXBxf2vtUgJyVaz+V0sIsvi+myRDpKtZofnHWtLUT1hpLxkqtkzReTzfGv7XESqqmrlInjeSF3JVLMviUgdVT1dVUtJ7j8ITC/C5548nHMZ9RCRtSJyc76/txORNSLSXET2iUiZfF9bISIt8/39BBHZLyIlReReEXkp39fK5e1/aYzrqS253XwJEWkgIstFZECiv088kueRbDV72Nr+JCIbRaRKor9PPJLnkWw1KyJNRGRx3jFriYgTkZKJ/j7xSJ5HEtbsGhFpk+/vpfLqtlaiv1c8kuORhDV7tIiMzavTAyLyjYickujvU3E8MvUKyIZ8f14nuZ24iMg259zefF+rKSJTVTUnX3ZQRKrm7fM/x3HO/ayq22NdiHPu63x//UJV7xeRu0RkRKzHQlpLmpo9RFWvkNw6vdQ5lx16HKStpKhZVS0hIuNF5M/OuQOqGsvuyCxJUbN5fhKRivn+fujPuwKOhfSVTDV7r4g0FZGTJPc3ALqIyLuqeqZzbnfA8VJGpv6a30n5/nyyiGzK+7M7bLsNItLWOXdsvkcZl3ufyOb8x1HVspJ7afTQ3zvr/57Qd/ijoF+ZciLCT3scLqlqVlXbiMiTInK5c+6L+D5VpIlkqdmKkntlalLe/QOf5O2+UVUviuszRqpLlpoVyf3V6XPynfMcEdnqnAv+BzCkpWSq2YYiMsk5t9E5d8A594yIVJIMuG8qU5upnqp6oubefDdIRCYVsN3jIjJMVWuKiKhqlqr+V97XJovIZap6oaoeLSL3S77vp3PuRfe/J/Qd/lifd8y2qlo178/1RGSwiLxWFE8aKS2ZavYSyR060dE5t6Boni7SQLLU7A+S+y+vDfMe7fJ2bywi8+P8nJHakqVmRUSeE5EbVbW+qh4ruffCPBP3Z4xUl0w1+4mIdFLVqpo78KKr5P566ur4P+3kkqnN1D9FZIaIfC25v186tIDtxorI6yIyQ1V3icjHkjsKWpxzy0SkZ96xNkvuzXwbA9bSUkSWqOrPIvKmiEyR3BulgfySqWYHi8gxIvJmvn+ZeivgOEhvSVGzLteWQw8R2Zb3pa3OuX2xPSWkuaSo2bzjTBeRB0Vktoisl9xf4RoS63GQ9pKmZkVkpOQOSlksIjtFpLfk/qPrzoBjpRR17vArgelNVdeKSDfn3MxErwWIgppFqqFmkWqoWaQaajZ5ZOqVKQAAAAAoFJopAAAAAAiQcb/mBwAAAADxwJUpAAAAAAgQ04f2qiqXsRAT51xCPzOLmkWsqFmkoGznXFaiTk7NIgA1i1RTYM1yZQoAgNS2LtELAGJEzSLVFFizNFMAAAAAEIBmCgAAAAAC0EwBAAAAQACaKQAAAAAIQDMFAAAAAAFopgAAAAAgAM0UAAAAAASgmQIAAACAADRTAAAAABCAZgoAAAAAAtBMAQAAAEAAmikAAAAACEAzBQAAAAABaKYAAAAAIADNFAAAAAAEoJkCAAAAgAA0UwAAAAAQgGYKAAAAAAKUTPQCABS/nJwcbz5lyhSTqarJli9fbrLBgwcXfmEAAAAphCtTAAAAABCAZgoAAAAAAtBMAQAAAEAAmikAAAAACKDOuegbq0bfOIldf/31JuvZs6fJ9uzZY7Lhw4ebbPr06XFZVzpyztnpBcUoXWo23g4ePOjNfa8HvgEUvu06depksqlTpwasLrGoWaSghc65Jok6ebrUbNWqVU12yy23mKx58+Ym27Rpk8kKen+1ZcsWk915550RVphWqNk4KFeunMnq1atnsu7du0c6XlZWljf31XJ2drbJfD/z33777UjnTgEF1ixXpgAAAAAgAM0UAAAAAASgmQIAAACAADRTAAAAABCgZKIXUNQmTJhgsqOOOspkv//970123HHHmWzKlCkmO/HEE73nfuqpp6IsESh2vpuqCzJ06FCTVa5c2WQDBw40WSoOoEDsfK+LHTp0iLTv3LlzTdauXTuT7dq1K/aFASJSq1Ytk914440m69Gjh8l87wOiDuWJxbHHHmuyu+++22Tbt28v1HmQGnyDIHyvqX/+859NVrduXZNFrVnfdrFs261bN5Oly3CqI+HKFAAAAAAEoJkCAAAAgAA0UwAAAAAQgGYKAAAAAAKk1QAK341vxx9/vMn+8Ic/mGz//v0m++6770zWpIn98ONTTz016hKBpPD3v/898raNGjUyme//NWQu383JUW/Iv+CCC0y2du1ak40aNcq7/7PPPmuyTZs2RTo3MsO8efNMVrVq1eDj7d6922S+wRAFDafyuf766022evVqk40YMSLyMZG6evXqZTLfkKfCDpYI3S6WbZ944gmTrV+/3mQLFy6MfO5kw5UpAAAAAAhAMwUAAAAAAWimAAAAACAAzRQAAAAABEirARR/+tOfTHbppZeazDdsIqozzjjDZPXr1/duu2TJkuDzAMnMd+PpnDlzErASJIO//e1vJuvQoUPw8SpVqmSyYcOGebdt3769yQYPHmyymTNnBq8HqaNx48Ymq1atmsl8N+l/9dVXJnvggQdM9t5775ls3759Jjv77LO9a5w6darJKlSoYDJfHZ900kkmu/XWW73nQeryDZuIOtQn3tsV9piVK1c2mW+IFQMoAAAAACDD0EwBAAAAQACaKQAAAAAIQDMFAAAAAAFSdgBFixYtTPbZZ5+ZzPdJ5YXhu/mzb9++3m2ff/75uJ4bSATfIAHfTaZTpkwpjuUgCX3++ecme+ONN0x2+eWXRzpex44dTTZy5Ejvtuedd57J/v3vf5usc+fOJnv11VcjrQepwze0oUQJ++/GOTk5Jtu7d6/JfLW0a9euSGuZPXu2Nz/22GNNdt9995nsnnvuMVmPHj1MdtRRR0XaDqkjas36hkFlZ2ebbP369YVaT7169UxWrly5SPv61phuuDIFAAAAAAFopgAAAAAgAM0UAAAAAASgmQIAAACAACk7gOJ3v/udydauXVvk5/XdeHrKKad4t61fv77Jli9fHvc1AUVp6tSpJuvevbvJ5s6dWxzLQRLas2ePyTZv3hx8vAULFpjslltu8W47evRok5199tkmGzBggMl8AwK+//77KEtEkjr33HNN5rtx3/ez/IEHHoi0XVHwDVi55JJLTPab3/zGZDfeeKPJfANgfMM0kHi+IU++mvUNfpozZ47J+vTpY7JFixYFri7Xc889ZzLfUB/fGjMBV6YAAAAAIADNFAAAAAAEoJkCAAAAgAA0UwAAAAAQIGUHUDRt2tRky5Yti+s5srKyTDZ8+HCTTZs2zbu/75PYr7nmmsIvDCgkX22L+G/S990cyyAVFLdZs2Z588cff9xk48ePN1mjRo1Mdscdd5js3nvvDVgdEqFly5YmK+i17XBffPGFyXzDdorL7t27TfbQQw+Z7JVXXol0vNNOO63Qa0LxWLFihclmzJhhsq5du5osOzs7+Ly+/1d69erl3dY3bEJVI53Ht11h1p2MuDIFAAAAAAFopgAAAAAgAM0UAAAAAASgmQIAAACAACk7gMJ3Y+a4ceNMFvUTzX2fmv7oo4+a7MUXXzTZvHnzvGv03TB91llnmWzp0qXe/YFY1axZ02Tbtm0zWZcuXbz7//nPfzaZ78boTp06BawOiKZVq1Ymmzhxonfbjz/+uKiXgyR19913m+zoo49OwEqKxubNm4P37datm8kefvjhwiwHRWTlypUma9u2bVzPcdNNN5mse/fuJvMN6hERcc5FOo9vO99glxEjRkQ6XqrgyhQAAAAABKCZAgAAAIAANFMAAAAAEIBmCgAAAAACpOwAirfeestkb7/9tsm++eYbk23dutVkO3bsMFnv3r1NVtCwCZ8JEyaYbMiQISa76qqrTBb1Zj8gvwULFpisT58+Juvfv793f1/dDR8+3GS+G2aBeHnyySdNVqZMGe+2H330UVEvBylEVSNlY8eOLY7lxJ3vufhUqFChiFeCZNC4cWOTvfnmmybLysoyme/nfdT6imXbW265xWS+wVapjCtTAAAAABCAZgoAAAAAAtBMAQAAAEAAmikAAAAACJCyAyh8fDe5+W5kLlHC9pCffvpp3Nfju3F/7ty5JuvVq5fJUvXmWBSfDh06mMx3k+nAgQMjbScismLFCpOl2yeVo3h89tlnwfv6XqNr1arl3ZYBFJnrqaeeMlnLli0j7fv73//eZJMnTy70mopa1OFUDLHKDL5hE5UrVzaZrx5iqZHC1J3vvcrf//73yOdOBVyZAgAAAIAANFMAAAAAEIBmCgAAAAAC0EwBAAAAQACaKQAAAAAIkFbT/HwWLVqUsHP/+OOPJrvhhhtM9sYbb5js4MGDJnvsscfiszAktXr16pmsY8eOJuvfv7/JfJN0Xn31VZPVr1/fe+4rrrjCZIMGDTLZsGHDvPsDh/gmrX388cfBx/vqq6+8ed26dYOPidT27rvvmmzbtm0mO/7440126qmnFsmagHioWbOmyRYsWGAy32Re3/sAVY103qjbxbLt0KFDTbZw4cJIWargyhQAAAAABKCZAgAAAIAANFMAAAAAEIBmCgAAAAACpP0AimTju8GuZ8+eJnvhhRdM5hum8eGHH8ZnYUgI302mvuEOHTp0MNkHH3xgslq1apnsn//8p8nKlSvnXc+KFStM9sADD5hs7dq1JnvxxRe9x0Rm8g3R+fzzz+N+nho1akTa7sCBAyabMWNGvJeDYrR9+3aT7du3L9K+jRs3NlmjRo1MlsghVshcVapUMVnlypVN5hs24ct8om5X2GP61t2tWzeTMYACAAAAADIMzRQAAAAABKCZAgAAAIAANFMAAAAAEIABFEngX//6l8kef/xxk40fP95kzZo1M9kPP/wQn4WhyD333HMmu+CCC0y2bds2k/Xp08dk69evN1l2drbJypYt612PbwDF1KlTTTZw4ECT7d69O9K+QH7ly5c3WYsWLUx2zTXXePf3betTsqT9cTdhwgSTffHFFya79957TbZ69epI50XxevPNN03Wo0cPk5UuXdpkn3zyicnOO+88k3366aeBqyvYSSedZLJevXqZTFWDj3fllVeabPLkyZGOh+Ll++8cNVu5cqXJxo4da7Lly5ebrH79+t71+IZgtW7d2rvt4aLWbCrjyhQAAAAABKCZAgAAAIAANFMAAAAAEIBmCgAAAAACaCyfgKyq0TdGofzqV78y2bJly0zWu3dvk7322mtFsqYQzrmE3nmYTDWblZVlsq1bt5rsgw8+MFnz5s2LYklx9d1335msTZs2Jlu0aFFxLCcYNVt0fEMgZsyYYbJE1vvBgwdNdsMNN5jshRdeKI7lRLXQOdckUSdP9pp9/fXXTdauXbvg491xxx0me++990y2efNm7/6nnXaayXzDJq666qrYF5fHd9P/1VdfbbIEDqCgZo+gcePGJps/f77JfP+dmzZtarKi+Lnre6309RS+NfqGalWrVi0+Cys6BdYsV6YAAAAAIADNFAAAAAAEoJkCAAAAgAA0UwAAAAAQwN4NjKSwZ88ek02dOtVkv/71r02WTAMo8B++TxD33azp+++cCrp27Woy36epJ/sACsRH6dKlTTZ9+nSTNWvWrDiW4zVnzhyTjRo1ymRvvPFGcSwHRWTgwIEm8w15atGiRaTjjR071mS7du0yWSwDKHw36ccyIAzpZf369SbbsGGDyWrW09aY7wAAIABJREFUrGky33uNwv7c9Q3Q8tWsj2+7qPumCq5MAQAAAEAAmikAAAAACEAzBQAAAAABaKYAAAAAIAADKFJIvXr1TLZ48eIErAQhfJ/4vX37dpPddNNNJvPdjJrIQRW+G1ynTJlispycHJO98MILRbImJJd9+/aZbMGCBSbzDaB4/vnnTdauXTvveSpXrmyyn3/+2WSdOnUy2ezZs032yy+/eM+D1LV06VKTXXbZZSa7/fbbTXbPPfeYrEKFCiYrX768yXyDJmIxbdo0kzVq1MhkJ5xwQqHOg+Tje7/Qp08fk02ePNlkAwYMMFm5cuVMtnLlSpP5Bk2IiHTr1s1kvgEpUYem+N4vpDKuTAEAAABAAJopAAAAAAhAMwUAAAAAAWimAAAAACAAAyiSlO9mwQYNGphs4sSJxbEcxIFvYITv08tvvPFGkz333HMmGz58uMlGjBgRuLqCDRo0yGT9+/c3mW/YxLBhw+K+HqQG343II0eONNkTTzxhMt/AlU8++cR7Ht8Aip07d5ps+vTp3v2RmXyDRkaNGmUy35CSiy++OO7refjhhyNt9+GHH5qsevXqJlNVk/kGZyB1+N5DlChhr4n4fhb37t070na+48WybdTtsrOzvedJVVyZAgAAAIAANFMAAAAAEIBmCgAAAAAC0EwBAAAAQAAGUCQB3815d911l8mWL19uMt+nXyN1jBkzxmS+G+V92ZNPPhn39Tz//PMmq1evnsl8N4/6Pp3dd8MsMtf3338fKSusRx99NO7HRGZauHBhpKy4bNq0yWS+YS8+f/rTn0zGEKvUNnToUJMNGDDAZL7BEL668W0Xy7a+7V599VWTFcWwrETiyhQAAAAABKCZAgAAAIAANFMAAAAAEIBmCgAAAAACMICiCFWuXNlk7du3N5nvxv1q1aqZrHv37vFZGJLaypUrTdalS5e4nsM3VEJE5IorrjDZ8OHDTeYbfpFun2iO1LB3716TffHFFwlYCVD0Jk2aZDLf6zYyw+DBg03me030DarwDZDwDUQraNtFixaZbMqUKSZLt2ETPlyZAgAAAIAANFMAAAAAEIBmCgAAAAAC0EwBAAAAQAAGUOSpVauWyS6++GKTtWvXzmRnnnmm95hnnXWWyXbs2GGy3r17m8x3k6nvpkJkhrlz58b1eL4hFyIiFSpUiOt5gKLmq+W33norASsBit60adNMtnnzZpNVrFjRZC+//HKRrAmJk5WVZbILL7zQZL4BEs45k73//vve8/gGUc2YMSPKEjMCV6YAAAAAIADNFAAAAAAEoJkCAAAAgAA0UwAAAAAQQH03oBW4sWr0jQERcc5pIs9PzSJW1GxyeuD/t3fv8TbWaR/Hr1+2QzaRQxEhppNDsTORLSGnSEzimTSmRtOknFUqkRGlV3TSSDPVYColNaJ5MuNJitHjKUqTw8Pg2aJsbaKdc/g9f+xtZo/rt5t7/fZa61732p/363W/Xnzdh2vvflbrcu/7WhMnOvMePXqoLCsrK9HlpJrV1tqWYV2cNQsPrFlETbFrljtTAAAAAOCBZgoAAAAAPNBMAQAAAIAHmikAAAAA8MAACiQUD/MjaliziCAe5kfUsGYRNQygAAAAAIB4opkCAAAAAA80UwAAAADggWYKAAAAADzQTAEAAACAB5opAAAAAPBAMwUAAAAAHmimAAAAAMADzRQAAAAAeMiIcf/dIrItEYUgLdUPuwBhzSI2rFlEUdjrljWLWLFmETXFrlljrU1mIQAAAACQFvgxPwAAAADwQDMFAAAAAB5opgAAAADAA80UAAAAAHigmQIAAAAADzRTAAAAAOCBZgoAAAAAPNBMAQAAAIAHmikAAAAA8EAzBQAAAAAeaKbiyBiTY4zpFHYdQFCsWUQNaxZRw5pF1LBmY0MzlUKMMUuMMdYYkxF2LUAQrFlEDWsWUWCMudkYs9oYk2+M2WGMeYw1i1RmjHnOGLO/yHbEGPNd2HUlA81UMZL9omWMuUlEyibzmkgvrFlEDWsWUZPENVtRREaISA0RaSUiV4vI3Um6NtJIstastXaQtbbSyU1EXhWRecm4dthKXTNVeOvyfmPMemPMXmPMTGNMBWNM+8J//bnXGJMrIjONMacZY+4zxmwxxuwxxrxujKlW5FwDjDHbCv/sgRLUVEVExovI6Dh8iUgzrFlEDWsWUZNqa9ZaO8Nau9xae9Ra+6WIvCIi2XH6cpEGUm3NnlJbpoj0EZHZJT1XFJS6ZqrQTSLSVUQaicgFIjK2MK8lItVEpL6I/EpEhopIbxG5SkTOEZG9IjJdRMQY01hEZojIgMI/qy4idU9ewBjT3xiz7we2ekXqeaTwXLmJ+oIReaxZRA1rFlGTamu2qHYisi6uXy3SQaqu2T4ikiciy+L9Backa22p2kQkR0QGFfl9dxHZIiLtReSoiFQo8mcbROTqIr+vLSLfi0iGiDwoIq8V+bPMwuM7xVhPSxFZU3jOBiJiRSQj7O8TW+psrFm2qG2sWbaobam2Zk+pbaCI7BCRGmF/n9hSZ0vxNbtERH4d9vcoWVtpfZhxe5Ffb5OCTlxEJM9ae7jIn9UXkfnGmBNFsuMicnbhMf84j7X2gDFmTyxFGGNOE5FnRWS4tfaYMSaWw1G6sGYRNaxZRE1KrNmijDG9RWSyFLyx3e17HqStVFyz9aSgobvN9xxRU1p/zO/cIr+uJyJfFf7anrLfdhG5xlpbtchWwRb8/PLOoucxxlSUglujJ39/k/nXqSanbvVE5Awp+BfTuYU/1/px4eE7jDFXxvUrRtSxZhE1rFlETaqs2ZP7dhOR50Wkp7X28/h+qUgTKbVmCw0QkRXW2q3x+iJTXti3xpK9ScFt0c+l4OdBq4nIX6XgZ+nbi8iOU/YdKSLvi0j9wt/XFJFehb9uIiL7RaStiJQTkakickxiuC0qIkYKfq715PZjKfgLUEdEyoX9vWJLjY01yxa1jTXLFrUtldZs4Xk6isgeEWkX9veGLTW3VFuzRa61UUQGhv39SeZWWu9MzRGRxSKyVQp+vnRSMfs9LSILRWSxKZiVv1IKRpSKtXadiAwuPNdOKXiYb0csRdgCuSc3KXhYT0Rkl7X2aGxfEtIcaxZRw5pF1KTEmi00TkSqiMg7Re4ALPI4D9JbKq1ZMcZcIQXNXakYiX6SKewiSw1jTI6I/NJa+27YtQBBsGYRNaxZRA1rFlHDmk0dpfXOFAAAAACUCM0UAAAAAHgodT/mBwAAAADxwJ0pAAAAAPBAMwUAAAAAHjJi2dkYw88EIibWWhPm9VmziBVrFhG021pbM6yLs2bhgTWLqCl2zXJnCgCAaNsWdgFAjFiziJpi1yzNFAAAAAB4oJkCAAAAAA80UwAAAADggWYKAAAAADzQTAEAAACAB5opAAAAAPBAMwUAAAAAHmimAAAAAMADzRQAAAAAeKCZAgAAAAAPNFMAAAAA4IFmCgAAAAA80EwBAAAAgAeaKQAAAADwQDMFAAAAAB5opgAAAADAA80UAAAAAHigmQIAAAAADxlhFwC3M888U2V79+4NoRIAAAAALtyZAgAAAAAPNFMAAAAA4IFmCgAAAAA80EwBAAAAgAcGUKSAF198UWXNmzdXWffu3VW2a9euhNQEAMkwduxYlY0ZM0ZlOTk5zuN79+6tsk2bNpW4LgAAguDOFAAAAAB4oJkCAAAAAA80UwAAAADggWYKAAAAADwwgCLJunXrprKf/exnKitbtqzKnnnmGZUNHz5cZTt37vSsDlGSkaH/+nbo0EFlr776qsqqVavmPOfmzZtVlpWVpbL9+/cHKRH4F08//bTKhg0bFujY6tWrO/M1a9ao7NZbb1WZ6+8BSofMzEyVPfbYYyobNGiQyrZt26ay7OxslfH/XaD04s4UAAAAAHigmQIAAAAADzRTAAAAAOCBZgoAAAAAPDCAIslq1aqlMtewCZcWLVqorH379irjQevSYdKkSSq7++67Ax1rrXXmDRs2VNnpp5+uMgZQ4N958MEHVRZ02MTo0aNV9sILLzj3vfTSS1XWvXt3la1YsUJlX3zxRaB6EG133HGHylzDJg4ePKgy14CTw4cPx6cwIMUU9370wgsvVNmUKVNU1rVrV5VNnz5dZUOHDvWoLnVxZwoAAAAAPNBMAQAAAIAHmikAAAAA8EAzBQAAAAAeTHEPojt3Nib4znByfRL7//3f/6msRo0agc73/fffq6xLly4q++CDDwKdL96stSaUCxdKlzXbsmVLlb333nsqq1ixosoOHTqksiNHjjivU7VqVZWNGTNGZY899pjz+HTAmo1dzZo1VbZ161aVVapUSWUvvfSSym699VaVuV7rYjFy5EiV3X777YGysF4/Y7DaWqtfJJIk1dfstGnTVDZ48GCVffbZZyrLyspKSE1gzYatSZMmKnviiSec+3bu3FllGzduVNnOnTtV9umnn6rsrrvuClJiqil2zXJnCgAAAAA80EwBAAAAgAeaKQAAAADwQDMFAAAAAB4ywi6gtDlw4IDKTpw44X0+16dVV6hQwft8SK46deqozPXw/Z133qky17CJffv2qWzixIkqe/fdd531uB7AbtiwoXNf4CTXJ9y7hk24Hlh+4IEHVFbSYROua48bN05lZ555pspatGihsggMoMAPcP1/1+X8889XWevWrVW2cuXKQOcbMGCAM7/ssstUtnr16kDnfOONN1TmGjKE0qt69eoq69Onj8pc7w1c7wFERC655BKV5eXlqez0009XWU5OjvOc6YQ7UwAAAADggWYKAAAAADzQTAEAAACAB5opAAAAAPDAAAogRNddd53Kxo8frzJr9Ye1f/zxxyobM2aMypYuXaoy1yefF6dRo0aB90X6cz1g3KxZs0DHuh543r59e4lrOpXrYWvXsAmXLVu2xLschGzPnj2B9nOt7T/96U8qW7Nmjcp2796tsr59+zqv43o9D2rKlCkqu/baa1W2atUq72sgOlxr7LbbblOZa+jJHXfcobI333zTeZ3jx4+r7IwzzlBZuXLlnMenO+5MAQAAAIAHmikAAAAA8EAzBQAAAAAeaKYAAAAAwEPaD6CoUqWKyurUqaOy7OzsQNk777yjsrVr1zqvvX//fpVlZWWpzPUQH9JPZmamykaNGqUyY4zK8vPzVda9e3eV7d2717O64q+9cOFC7/NVq1ZNZd988433+RA+1+vVRRddpLIdO3ao7LXXXktITaeqV69eoP2+/vprlS1btize5SBkU6dOVZnr9Xjw4MEqq169uso6dOgQ6LoHDhxw5q73Bm+88YbKXINUatWqpbL27durjAEU6ce1Rpo3b66ybt26qWzz5s1xr+eSSy5RWd26dVX285//XGUjRoxQ2aZNm+JTWAi4MwUAAAAAHmimAAAAAMADzRQAAAAAeKCZAgAAAAAPkR1A4XrguXbt2ipzPXjaokUL7+u6HqRzPbQvIpKbm6sy10OvFSpUCHRt13W2bt2qsk8//TTQ+ZBcP/3pT1V23nnnqezbb79VmetTzksybMI10EJEZN68eSqbPXu2ys466yyV3XnnnSq74oorVNa1a9cgJSJFudaxy7Bhw1R2/PjxeJfj1LZt20D7ff755ypz/f1D+pkwYYLKpk2bprIePXp4X+OTTz5x5uvXr1dZ5cqVVda5c2eVnX322d71IDqGDh2qso8++khlAwYMUNmhQ4fiXo/r/bXr/+WuARSugRiuAW8MoAAAAACAUoZmCgAAAAA80EwBAAAAgAeaKQAAAADwEIkBFK5PeJ4+fbrKXJ96X6dOnYTUVJS11pnH+0HR4q5zqtatW3tf489//rPKjh496n0+/FOXLl0C7ff3v/9dZe+++25ca9m+fbszdw0XcA1Nefrpp1XmGpJx4sQJlbVp00ZlH374obMepJ5+/foF2i/o61VJlS1bVmVVq1YNdOyqVaviXQ4izDXU5+WXX477dZo2baqya6+9VmXnn39+oPO98cYbJa4JqWXmzJkq279/fwiVFNi5c6fKxo0bp7JevXqp7K9//avK0u3/+dyZAgAAAAAPNFMAAAAA4IFmCgAAAAA80EwBAAAAgIeUG0DhenB46dKlKqtSpYr3NVwPa/bu3VtlGRnBvj3GGGfuegDbtW9JHtQ+77zzVPbWW295n++hhx5S2a9//Wvv85VW5cqVU1nNmjUDHfvII4/EuxzFNVRCxD2AYsSIESq7+OKLA12nTJkyKqtfv77K0u1h1HTmWtthatu2rcouv/zyQMfOmzcv3uUA/3D66ac782effVZlrsE8Lvfdd5/KcnJyYqoLqS8ZwyYqV66sso4dOzr3rVWrlspyc3MD7bdu3TqVbdy4MUiJkcGdKQAAAADwQDMFAAAAAB5opgAAAADAA80UAAAAAHigmQIAAAAADyk3za9Vq1YqCzq5zzX1r3///irbvXu3ylzT/IJOeoplGl9JJvclw9NPPx12CWnhxIkTKjt8+HCgYzdv3ux93a5du6qsYcOGKhs1apTzeNd0yHhr1KhRwq+B9FOnTh1nXpLXrPXr13sfC/w7xU3Zq169eqDjXZP7nnjiiZKUhDRTvnx5lR05ckRlAwYMUNmGDRtUtmfPHud1FixYoLKsrCyVzZo1S2UdOnRwnjOdcGcKAAAAADzQTAEAAACAB5opAAAAAPBAMwUAAAAAHlJuAIXrAXqXnTt3quz2229X2a5du1SWmZmpsttuuy3QdcPkGl5x/PhxlWVk6P+sc+fOVdmKFStUFnRIAn7YsWPHVPbdd9+pzBijsg8++EBlBw8eVNk555zjWZ3Iaae5/x3FNTjDtcYWLVqksmuvvTbQdZYvXx6kRKSooK8RlStXVplrvbter1xr6amnnnJe5+yzz1aZa7BE48aNVda5c2eVLVy40HkdlE6uB/zbtGmjMte6qVSpkvOcrtfZRx55RGVTp04NUiLSkOv1s1u3bipbt26dylyvf4sXL1bZoUOHVJafnx+0RLnllltU5hrctmbNmsDnjCruTAEAAACAB5opAAAAAPBAMwUAAAAAHmimAAAAAMBDyg2gaNq0aaD9ateurbIpU6aobMeOHSpzPTTnGkoRlOuhahH30ICyZcuq7E9/+pPKlixZojLXAIPc3FyVnX/++SpzPRT49ddfqwyJs3nzZpW5hopUrVpVZVWqVAl0bFCuB6BFRN5//32VTZo0KdC1e/ToobJ9+/apLC8vL0CFSFWuh+Kzs7NV9oc//EFlAwYMUJlrbV9++eUqK27N3n333Sr75ptvVDZr1iyVVa9e3XlO4KTJkyerbNiwYYGOLW7Nul4/Xe9LXK/Hy5YtU9n3338fqB5Eh+v9nut9XFCuYWyxaNKkicrq16+vsieeeKJE14kq7kwBAAAAgAeaKQAAAADwQDMFAAAAAB5opgAAAADAQ8oNoPjLX/6isk6dOgU6tlevXvEuJ5DiBgEMHz5cZTt37lSZa9jE4cOHvetxfdI1wjdhwgSVuR6A7927t8pca8z16eWuh+xd9u7d68yfe+45lbkebg56HdcAGNensyM6FixYoLKJEyeq7LbbblNZ586dA10jPz9fZYMHD3bu+/LLL6vs5ptvDnQd4N9Zvny5yi6++OJAxxY3nMr1ML9rqJbr/VCbNm1U9tFHHwWqB/A1ZMgQldWtW1dlH3zwQTLKSTncmQIAAAAADzRTAAAAAOCBZgoAAAAAPNBMAQAAAIAHU9zwBOfOxgTf2dMZZ5yhsg0bNqjM9bBmMrhqWbVqlXPfESNGqKy4B//TlbXW/QRukiRjzZZGc+fOVVmfPn1U5vr70qxZs4TUFC+s2fgoV66cyrKzs1VWuXJllbkevD9y5Ejga7sGULiGpriG9XTt2jXwdVLIamtty7Auni5rNlnOPfdclY0bN05lAwcOVNmiRYtU1rNnz/gUllys2RR04403OvM5c+aorHv37ipzrc80Uuya5c4UAAAAAHigmQIAAAAADzRTAAAAAOCBZgoAAAAAPKTcAAqXBg0aqGz06NEqGzRokMo2bdqksjJlyqhszZo1KnvvvfdUNm/ePJXt3r1bZSjAw/zpiQEUicOaLTnXUIv8/HyVLVmyRGWdOnVKSE0JxsP8EXfBBReobNmyZSqrUaOGyoYOHaqyGTNmxKewxGHNhqx+/foqe+ONN5z7Hjx4UGVXXXVV3GtKcQygAAAAAIB4opkCAAAAAA80UwAAAADggWYKAAAAADxkhF1AEDk5OSobNWqUyh5//HGV7du3T2WnnaZ7yLy8PL/iAIiIiDF6bsPevXtDqASl3eHDh8MuAYiJa1jWPffco7JZs2apLDMzMxElIc1lZWWprHr16s59+/btm+hyIo07UwAAAADggWYKAAAAADzQTAEAAACAB5opAAAAAPAQiQEULq4HjLds2RJCJQBERKzVHyj/6quvhlAJAKQn1+tsjRo1QqgEUVK+fHmVjR8/XmWrV692Hu8aBId/4s4UAAAAAHigmQIAAAAADzRTAAAAAOCBZgoAAAAAPER2AAUAAC4nTpxQ2datW0OoBEi8q666KuwSkOJq1aqlsvvvv19ln376aTLKSTvcmQIAAAAADzRTAAAAAOCBZgoAAAAAPNBMAQAAAIAHBlAAANLK8ePHVfbJJ5+o7Mwzz0xGOUBCvfnmm2GXgBTStm1blVWvXl1lCxYsSEY5pQJ3pgAAAADAA80UAAAAAHigmQIAAAAADzRTAAAAAOCBARQAYrZu3TqVXXjhhYH2A8LQt2/fsEsAYnLZZZeFXQJSXIMGDVQ2cuRIlT355JNJqKb04s4UAAAAAHigmQIAAAAADzRTAAAAAOCBZgoAAAAAPBhrbfCdjQm+MyAi1loT5vVZs4gVaxYRtNpa2zKsi7Nm4YE1GwfZ2dkqe/3111VWp06dZJST7opds9yZAgAAAAAPNFMAAAAA4IFmCgAAAAA80EwBAAAAgIeMsAsAAAAAEJsVK1aojGETycedKQAAAADwQDMFAAAAAB5opgAAAADAA80UAAAAAHiIdQDFbhHZlohCkJbqh12AsGYRG9YsoijsdcuaRaxYs4iaYtessdYmsxAAAAAASAv8mB8AAAAAeKCZAgAAAAAPNFMAAAAA4IFmCgAAAAA80EwBAAAAgAeaKQAAAADwQDMFAAAAAB5opgAAAADAA80UAAAAAHigmQIAAAAADzRTcWSMyTHGdAq7DiAo1iyihjWLqGHNImpYs7GhmUohxpglxhhrjMkIuxagOMaYm40xq40x+caYHcaYx1izSGWsWUSNMeY5Y8z+ItsRY8x3YdcF/BBjzEhjTG7ha+3vjTHlw64pGWimipHs/9EaY24SkbLJvCbSSxLXbEURGSEiNUSklYhcLSJ3J+naSCOsWURNstastXaQtbbSyU1EXhWRecm4NtJLstasMaariNwnBa+v9UWkoYhMSMa1w1bqmqnCW5f3G2PWG2P2GmNmGmMqGGPaF/6L5b3GmFwRmWmMOc0Yc58xZosxZo8x5nVjTLUi5xpgjNlW+GcPlKCmKiIyXkRGx+FLRJpJtTVrrZ1hrV1urT1qrf1SRF4Rkew4fblIA6xZRE2qrdlTassUkT4iMruk50L6SME1e7OIvGitXWet3SsiE0XklpJ/pamv1DVThW4Ska4i0khELhCRsYV5LRGpJgUd9a9EZKiI9BaRq0TkHBHZKyLTRUSMMY1FZIaIDCj8s+oiUvfkBYwx/Y0x+35gq1eknkcKz5WbqC8YkZdqa7aodiKyLq5fLdIBaxZRk6prto+I5InIsnh/wYi8VFqzTUTksyK1fSYiZxtjqifg604t1tpStYlIjogMKvL77iKyRUTai8hREalQ5M82iMjVRX5fW0S+F5EMEXlQRF4r8meZhcd3irGeliKypvCcDUTEikhG2N8nttTZUm3NnlLbQBHZISI1wv4+saXOxppli9qW4mt2iYj8OuzvEVtqbam2Zguv3a3I78sWvqdtEPb3KtFbaX0Ad3uRX2+Tgk5cRCTPWnu4yJ/VF5H5xpgTRbLjInJ24TH/OI+19oAxZk8sRRhjThORZ0VkuLX2mDEmlsNRuqTEmi3KGNNbRCZLwQvubt/zIG2xZhE1qbhm60nBm+PbfM+BtJZKa3a/iJxR5Pcnf532g1NK64/5nVvk1/VE5KvCX9tT9tsuItdYa6sW2SrYgp+531n0PMaYilJwa/Tk728y/zqJ59StnhQstJYiMrfw51o/Ljx8hzHmyrh+xYi6VFmzJ/ftJiLPi0hPa+3n8f1SkSZYs4ialFqzhQaIyApr7dZ4fZFIK6m0ZteJyKVFrnmpiOyy1nr/Y0JkhH1rLNmbFNwW/VwKfh60moj8VQqeWWovIjtO2XekiLwvIvULf19TRHoV/rqJFHThbUWknIhMFZFjEsNtURExUvBzrSe3H0vBX4A6IlIu7O8VW2psqbRmC8/TUUT2iEi7sL83bKm5sWbZoral2potcq2NIjIw7O8PW+ptqbZmRaSbFDz731hEqorIeyLyaNjfp2RspfXO1BwRWSwiW6XgZzwnFbPf0yKyUEQWm4LPd1gpBWN1xVq7TkQGF55rpxQ8zLcjliJsgdyTmxQ8YCpS0Mkfje1LQppLiTVbaJyIVBGRd4r8y9Qij/MgvbFmETWptGbFGHOFFLxRZiQ6ipMya9Za+2cReUxElorIF1LwY4fjYz1PFJnCbrLUMMbkiMgvrbXvhl0LEARrFlHDmkXUsGYRNazZ1FFa70wBAAAAQInQTAEAAACAh1L3Y34AAAAAEA/cmQIAAAAADzRTAAAAAOAhI5adjTH8TCBiYq01YV6fNYtYsWYRQbuttTXDujhrFh5Ys4iaYtcsd6YAAIi2bWEXAMSINYuoKXbN0kwBAAAAgAeaKQAAAADwQDMFAAAAAB5opgAAAADAA80UAAAAAHigmQIAAAAADzRTAAAAAOCBZgoAAABULNc6AAAS/UlEQVQAPNBMAQAAAIAHmikAAAAA8EAzBQAAAAAeaKYAAAAAwAPNFAAAAAB4oJkCAAAAAA80UwAAAADggWYKAAAAADzQTAEAAACAh4ywCwAAAAAQm9q1a6ts4MCBKhszZozKKlSooLK1a9c6rzN+/HiVvfXWW0FKLBW4MwUAAAAAHmimAAAAAMADzRQAAAAAeKCZAgAAAAAPxlobfGdjgu+MEqlWrZrKfvOb36isRYsWKrv44osTUpMPa60J8/qsWcSKNYsIWm2tbRnWxVmz8MCajdGPfvQjlT3zzDMq69y5c9yvfejQIZXdcsstKnvzzTfjfu0UUuya5c4UAAAAAHigmQIAAAAADzRTAAAAAOCBZgoAAAAAPGSEXQBEqlatqrL//M//VFnr1q1V9u233yakJoSnXr16Krv66qtV1q5dO+9rGOOesdCnTx+V5ebmqqxRo0Yq+/jjj1W2fv16j+oKXHLJJSrLy8tTWbdu3byvgR9WpkwZlT344IMqcw3CufbaawNdw/XftEKFCir77//+b+fxa9asUdlzzz2nspycnED1AD5q1Kihsv79+zv3HTJkiMpcwwW++OILlc2bN09ljz32mMpcf6+Qmk47Td/XcK2diRMnqsz1fiGWwXJBVaxYUWV9+/ZVWZoPoCgWd6YAAAAAwAPNFAAAAAB4oJkCAAAAAA80UwAAAADgwcTyoFoUPzE6qMaNG6ts9OjRKrvzzjtVdvDgwcDXueyyy1Q2bdo0lbVp00Zl3333ncpcD3kvW7YscD2JZq11TzpIklRfs7Vr11bZrFmzVNapUyfva+Tn56ts165dzn0T8eDqqerUqaOyzMzMQMfu2LFDZfXr1y9xTUWxZv8pOztbZcuXLw+hkti4BvP8z//8j8puvfVWlX355ZcJqSnBVltrW4Z18VRas/F21llnqcw1HGDs2LEqO/PMMwNfxzXoZ+DAgSqrVauWylzDVT744IPA1w4Ja7bQvffeq7KHH3440LGbN29W2bPPPquyr776KtD5Hn/8cWdet25dlW3fvl1lrsFRrvcgEVXsmuXOFAAAAAB4oJkCAAAAAA80UwAAAADggWYKAAAAADxkhF1AqmjZUj9TdvPNN6use/fuKluwYEHg6/zsZz9TWYUKFQIdO2rUKJWl0rAJxK5MmTIqcz1Q6srmzJmjsuPHj6ssLy9PZVu2bAlaYolUqVJFZa4127RpU5W5hgh06dIlPoUhkBYtWgTa78SJEyo7dOiQyt577z2VuQY+NGvWTGU//vGPndcuX768ylzrzrV23nnnHZX94he/UNknn3zivDaiq2zZsiq76667VHb77berrF69eio7duyYylasWOG8tuvvlWuwxAsvvKCyefPmqezdd991Xgepp2vXrip76KGHAh27du1alXXs2FFl33zzTeyFFRo/fnzgfV1DKQYNGqQy1/C0devWqSzK72e5MwUAAAAAHmimAAAAAMADzRQAAAAAeKCZAgAAAAAPDKAo9Morr6isSZMmKhsyZIjKfvnLX8a9HtdDfLNmzYr7dRCuHTt2qGzw4MEhVFJybdu2VZlrOEvVqlVV5hpMMHLkSJVt3LjRszr46N+/f6D9/va3v6ksKysrrrU0aNDAmbsGU4wZM0Zll156qcpcgy5cD/3H+2tBctWuXVtlL730kso6dOgQ6HwzZ85U2ejRo1VW3CAA12ugawjWjTfeqLLhw4erLDs7W2U33XSTylyDjJA4ruFi/fr1U1lGRrC34suXL1dZSYZNZGZmqqy4gWjGmEDnnDx5snc9CxcuVJlraNuBAwe8r5Eo3JkCAAAAAA80UwAAAADggWYKAAAAADzQTAEAAACAB2OtDb6zMcF3TlMXXXSRyqZPn+7c1/XJ1C5PPfWUylwP30eRtTbYU4sJwppNDNewCdfDo1WqVFHZ0qVLVdarVy+VhfWQKWv2nz788EOVtW7dWmWuoQ2/+tWvElJTEBUrVlTZlClTVHbHHXeo7NixYypzfS0pNhBotbVWTzBIklRasy4TJ05U2dixY1Xmej/09ttvq6xv374qO3r0qGd1sWnevLnKZs+erTLXcKMePXokpCZPab9mr7vuOpX98Y9/VNnixYtV1qVLl0D7de/ePVAtrsESriEsP/nJT5zHuwZQxNI/BOG6xptvvqmyAQMGqOzIkSNxraUYxa5Z7kwBAAAAgAeaKQAAAADwQDMFAAAAAB5opgAAAADAQ7CPXcY/uB6Kb9q0aeDjDx06pLLnn3++RDUBiTR+/HiVDRs2TGWuT1P/7W9/q7L77rtPZan4ieYIbsOGDWGX8C8OHjyoMtcQgmuuuUZlDRo0UFksr/FIPa4BDVlZWSrr1q2byubMmaOyZA2bcFmzZo3KXAN8li1bprL7779fZZMnT45PYVC+/PJLlbkGNbkG4biOdQ2lmDdvnspc7zNda7tatWoqi8Xq1atVdsMNN6isXbt2Kps0aZLK6tWrp7Lrr78+UC39+vULtF+icGcKAAAAADzQTAEAAACAB5opAAAAAPBAMwUAAAAAHhhAEaP27dur7Kyzzgp8/MCBA1W2fv36kpQExM24ceNUNnToUJVVrVpVZQsWLFDZnXfeGZ/CEIolS5aorHXr1irr2LGjyp588smE1OTL9bC1a9gE0s/mzZtVtnTpUpXVrFlTZfPnz09ITfGUk5OjssWLF6usUaNGSagGJ+3Zs0dlrgEi1lrva7gGNJTkfMX56KOPVHbdddepLC8vT2WvvPKKylauXKmy999/X2W1a9dW2ZVXXllcmaHhzhQAAAAAeKCZAgAAAAAPNFMAAAAA4IFmCgAAAAA8MIDiB9SpU0dlY8eOLdE5y5cvr7IHHnhAZa6HvPPz81U2bNgwlbkeegSKmjp1qjMfMWKEylwPsw4ZMkRlzz//fMkLQ0qZMGGCyh555BGVnThxIhnllEi7du28j929e3ccK0EquOGGG1T23nvvqezo0aPJKAdpyDUY5KGHHlJZRoZ+K/7EE0+obNSoUXGp66SvvvpKZa5BUiIi48ePV9k333zjfe0tW7ao7OGHH1bZb37zG5W5hgndddddKnv88cc9q4sdd6YAAAAAwAPNFAAAAAB4oJkCAAAAAA80UwAAAADgIe0HULg+8bty5coqa9asmcpcwyYuuOCCEtUza9Yslbke8P/ss89U9tprr6msXLlyJaoH4apatarKXOvT9QD8oUOHVOZ6kNU1RKC4B1ldD1v/7ne/U9mMGTOcxyO9HDt2LFAWBZmZmYH2c/0dWLhwYbzLQRLVr19fZa1atVLZlClTklFOUvTq1UtlP//5z0OoBP+O6zV19OjRKrvmmmtU1qRJk0DXcA2IGjRoUKBjk2X58uUqM8aozPU+p0ePHipjAAUAAAAApDiaKQAAAADwQDMFAAAAAB5opgAAAADAA80UAAAAAHiIxDQ/18S6cePGqaxnz54qa9iwocpc09ISIT8/X2XPPvusyl599VWV/e1vf0tITUgtruk1jRs3VplrPUybNk1lN9xwg8q6desWuJ7/+q//UtmwYcMCHw+kqj59+gTab/369YEyREfz5s1V5pqgtmXLlmSUE3eXXHKJynJzc1W2aNGiZJSDJHJNg37//fdVNnTo0CRUUzKu6duur8+lRo0aKjvjjDNU5npfHg/cmQIAAAAADzRTAAAAAOCBZgoAAAAAPNBMAQAAAICHSAygaNOmjcoeeOABlR09elRlGzZsUNn27dtV5nrwzTUIwPXwWnEP6M+ePduZAye99NJLKnv00UdVdumll6rsxRdf9L6uMcaZ79q1S2UXXnihyjZu3Oh9bSAMxa35U61YsSLBlSAV7N27V2Vr1qwJoZKS+8UvfqGyt956K4RKkAry8vJU9v3334dQSWw6d+7sfWxmZqbKypcvX5JyYsKdKQAAAADwQDMFAAAAAB5opgAAAADAA80UAAAAAHiIxAAK16c516xZU2UnTpxQ2cGDB1V25MgRlT388MMqcw2g2Lp1q8oYNAFfTz31lMrmzp2rsjlz5qisdevWca/H9SBz3759VTZ9+nSV/fa3v1XZtm3b4lMYIqdMmTIqcw36cb2Wd+vWrUTXzsrKUlmLFi0CHesauNKqVSuVrVu3TmX79+8PdA0kV/PmzVWWm5sbQiUl96Mf/UhlN954o8quvvrqZJQDxE2dOnW8j83JyVGZaxBHonBnCgAAAAA80EwBAAAAgAeaKQAAAADwQDMFAAAAAB4iMYDCZc+ePd7Huh6Mvv766wMd+/bbb3tfFzjVsWPHVNaxY0eVuYZNuAauvPbaayrLz89XmTHGWY/rIf3LL79cZffee6/KhgwZorJVq1apbNCgQSr7+9//rjJrrbNGhMv1+jlixAiV/cd//IfKWrZsmZCa4qlTp06Bsv/93/9V2e9//3uVTZ06NT6Fwdu5556rso8++iiESmJz2mn637td62nJkiUqcw1IAVKFa2hKr169VBb0fYDrvU8ycWcKAAAAADzQTAEAAACAB5opAAAAAPBAMwUAAAAAHiI7gKIk7rrrLpVddNFFKtu1a5fK5s6dm5CaUDp1795dZS+88EKgY4cNG6ayGTNmlKieypUrq6xHjx4qcz082rhxY5VdddVVKtuwYYPK5s+fr7JXXnkl0H5IrsmTJ6vs7rvvDqGSxHANbFm2bFmgYxcuXBjvchAH+/fvV1mlSpVCqCQ29erVU1l2drbK2rZtm4xyELLnnntOZc8884zK+vbtq7KnnnpKZStXroxPYUVkZOi2okOHDipz1e0auOIatJWXl6ey559/PmiJCcGdKQAAAADwQDMFAAAAAB5opgAAAADAA80UAAAAAHgolQMoGjRoEGi/FStWqIxPFYcv14CGP/zhD4GOXbp0qcoS8cDld999pzLXJ4u7srp166ps9OjRKtu6davK+vfvr7KePXuqzPXA7M6dO1WGxAn6+nns2DGVrV27NtCxCxYsUJlrMISI+3W6X79+Khs1apTK/vKXv6jspz/9qcq+/fZb57URDa7XT9frS6pxDbxyreONGzcmoxyEbPPmzSqz1gbKXH8HJkyYoLItW7Y4r22MUVmjRo1U5hqqdcUVVzjPeSrXsAnX1/L4448HOl8ycWcKAAAAADzQTAEAAACAB5opAAAAAPBAMwUAAAAAHkrlAIqgNm3aFHYJiKhKlSqpzPWJ31WqVFHZrl27VHb99derzPWAf5h27NihsmHDhgU61vXp7EhN99xzj8oefvhhlbnWZ7IG+PTt2zfQfvPnz1cZwyYQBtdQn9NPP11lL730UjLKQQpyDcxZtmyZyq688kqVlS1bVmWTJk0KfG3XAArXcIiS2L9/v8pcNTKAAgAAAADSBM0UAAAAAHigmQIAAAAADzRTAAAAAOChVA6gcD2c5/Lhhx8muBKkqzFjxqisffv2Ktu3b5/KXA/P5+fnx6UuoKS2bdsWKAtTy5YtVXbkyBGVvf3228koByE7evSoylq1ahVCJQV69uypsiZNmqisS5cuySgHEdarVy+Vvf766yrr1KlTMsoJbObMmSp79NFHVbZly5ZklFNi3JkCAAAAAA80UwAAAADggWYKAAAAADzQTAEAAACAh7QfQHHdddeprGnTpio7dOiQyjZu3JiQmpBeRo4cqbJ77rlHZZ9//rnK2rVrpzKGTQDxl5ubq7KdO3eGUAmSbdGiRSp7+eWXVfaTn/xEZfPnz/e+rms4gIjICy+8EGhf1if+Hdf7hX79+qnMtbZdg7Jc74VFRJo1a6Yy13CIBQsWqGzWrFmBjnUNCYoK7kwBAAAAgAeaKQAAAADwQDMFAAAAAB5opgAAAADAQ9oPoFi5cqXK1q5dq7IDBw6obNOmTQmpCdHlGmjy4IMPqmzVqlUq69Onj8oYNgGUTKdOnVSWnZ2tMtdD/yi9hg8frrLRo0er7ODBgyqrVKmSyoYMGaKyc845x3ntW2+9VWUffvihc18gVq73FbNnzw6UwQ93pgAAAADAA80UAAAAAHigmQIAAAAADzRTAAAAAOAh7QdQfP311ypzfZIzEMSMGTNU5vok73vuuUdlX331VUJqAkqz8847T2WuwUNz5sxJRjmICNd6qFWrlsqmTZumMtcAiieffFJlrv9fiLgHXgGILu5MAQAAAIAHmikAAAAA8EAzBQAAAAAeaKYAAAAAwIOx1gbf2ZjgOwMiYq01YV6fNYtYsWYRQauttS3DujhrFh5Ys4iaYtcsd6YAAAAAwAPNFAAAAAB4oJkCAAAAAA80UwAAAADggWYKAAAAADzQTAEAAACAB5opAAAAAPBAMwUAAAAAHmimAAAAAMBDRoz77xaRbYkoBGmpftgFCGsWsWHNIorCXresWcSKNYuoKXbNGmttMgsBAAAAgLTAj/kBAAAAgAeaKQAAAADwQDMFAAAAAB5opgAAAADAA80UAAAAAHigmQIAAAAADzRTAAAAAOCBZgoAAAAAPNBMAQAAAICH/wc2hN/DNrkeJgAAAABJRU5ErkJggg==\",\n \"text/plain\": [\n \"
\"\n ]\n },\n \"metadata\": {\n \"tags\": []\n },\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"# Show some of them.\\n\",\n \"show_img_grid(\\n\",\n \" [test_ds['image'][idx] for idx in error_idxs[:25]],\\n\",\n \" [f'pred={logits[idx].argmax()}' for idx in error_idxs[:25]],\\n\",\n \")\"\n ]\n }\n ],\n \"metadata\": {\n \"accelerator\": \"GPU\",\n \"language_info\": {\n \"name\": \"python\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 0\n}\n" }, { "path": "examples/mnist/mnist_benchmark.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Benchmark for the MNIST example.\"\"\"\nimport time\n\nfrom absl import flags\nfrom absl.testing import absltest\nfrom absl.testing.flagsaver import flagsaver\nfrom flax.testing import Benchmark\nimport jax\nimport numpy as np\n\nimport main\nfrom configs import default\n\n\n# Parse absl flags test_srcdir and test_tmpdir.\njax.config.parse_flags_with_absl()\n\nFLAGS = flags.FLAGS\n\n\nclass MnistBenchmark(Benchmark):\n \"\"\"Benchmarks for the MNIST Flax example.\"\"\"\n\n @flagsaver\n def test_cpu(self):\n \"\"\"Run full training for MNIST CPU training.\"\"\"\n # Prepare and set flags defined in main.py.\n workdir = self.get_tmp_model_dir()\n config = default.get_config()\n\n FLAGS.workdir = workdir\n FLAGS.config = config\n\n start_time = time.time()\n main.main([])\n benchmark_time = time.time() - start_time\n\n summaries = self.read_summaries(workdir)\n\n # Summaries contain all the information necessary for the regression\n # metrics.\n wall_time, _, eval_accuracy = zip(*summaries['eval_accuracy'])\n wall_time = np.array(wall_time)\n sec_per_epoch = np.mean(wall_time[1:] - wall_time[:-1])\n end_eval_accuracy = eval_accuracy[-1]\n\n # Assertions are deferred until the test finishes, so the metrics are\n # always reported and benchmark success is determined based on *all*\n # assertions.\n self.assertBetween(end_eval_accuracy, 0.98, 1.0)\n\n # Use the reporting API to report single or multiple metrics/extras.\n self.report_wall_time(benchmark_time)\n self.report_metrics({\n 'sec_per_epoch': sec_per_epoch,\n 'accuracy': end_eval_accuracy,\n })\n self.report_extras({\n 'model_name': 'MNIST',\n 'description': 'CPU test for MNIST.',\n 'implementation': 'linen',\n })\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "examples/mnist/requirements.txt", "content": "absl-py==1.0.0\nclu==0.0.6\nflax==0.4.1\njax==0.3.4\n--find-links https://storage.googleapis.com/jax-releases/jax_releases.html\njaxlib==0.3.2+cuda11.cudnn82 # Make sure CUDA version matches the base image.\nml-collections==0.1.0\nnumpy==1.22.0\noptax==0.1.0\ntensorflow==2.11.1\ntensorflow-datasets==4.4.0\n" }, { "path": "examples/mnist/train.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"MNIST example.\n\nLibrary file which executes the training and evaluation loop for MNIST.\nThe data is loaded using tensorflow_datasets.\n\"\"\"\n\n# See issue #620.\n# pytype: disable=wrong-keyword-args\nfrom functools import partial\nfrom typing import Any\nfrom pathlib import Path\n\nfrom absl import logging\nfrom flax import nnx\nfrom flax.metrics import tensorboard\nimport jax\nimport ml_collections\nimport optax\nimport tensorflow as tf\nimport tensorflow_datasets as tfds\n\ntf.random.set_seed(0) # Set the random seed for reproducibility.\n\n\nclass CNN(nnx.Module):\n \"\"\"A simple CNN model.\"\"\"\n\n def __init__(self, rngs: nnx.Rngs):\n self.conv1 = nnx.Conv(1, 32, kernel_size=(3, 3), rngs=rngs)\n self.batch_norm1 = nnx.BatchNorm(32, rngs=rngs)\n self.dropout1 = nnx.Dropout(rate=0.025)\n self.conv2 = nnx.Conv(32, 64, kernel_size=(3, 3), rngs=rngs)\n self.batch_norm2 = nnx.BatchNorm(64, rngs=rngs)\n self.avg_pool = partial(nnx.avg_pool, window_shape=(2, 2), strides=(2, 2))\n self.linear1 = nnx.Linear(3136, 256, rngs=rngs)\n self.dropout2 = nnx.Dropout(rate=0.025)\n self.linear2 = nnx.Linear(256, 10, rngs=rngs)\n\n def __call__(self, x, rngs: nnx.Rngs):\n x = self.avg_pool(nnx.relu(self.batch_norm1(self.dropout1(self.conv1(x), rngs=rngs))))\n x = self.avg_pool(nnx.relu(self.batch_norm2(self.conv2(x))))\n x = x.reshape(x.shape[0], -1) # flatten\n x = nnx.relu(self.dropout2(self.linear1(x), rngs=rngs))\n x = self.linear2(x)\n return x\n\n\n\ndef loss_fn(model: CNN, batch, rngs):\n logits = model(batch['image'], rngs)\n loss = optax.softmax_cross_entropy_with_integer_labels(\n logits=logits, labels=batch['label']\n ).mean()\n return loss, logits\n\n\n@nnx.jit\ndef train_step(model: CNN, optimizer: nnx.Optimizer, metrics: nnx.MultiMetric, batch, rngs):\n \"\"\"Train for a single step.\"\"\"\n grad_fn = nnx.value_and_grad(loss_fn, has_aux=True)\n (loss, logits), grads = grad_fn(model, batch, rngs)\n metrics.update(loss=loss, logits=logits, labels=batch['label']) # In-place updates.\n optimizer.update(model, grads) # In-place updates.\n\n\n\n\n@nnx.jit\ndef eval_step(model: CNN, metrics: nnx.MultiMetric, batch):\n loss, logits = loss_fn(model, batch, None)\n metrics.update(loss=loss, logits=logits, labels=batch['label']) # In-place updates.\n\n\ndef get_datasets(\n config: ml_collections.ConfigDict,\n) -> tuple[tf.data.Dataset, tf.data.Dataset]:\n \"\"\"Load MNIST train and test datasets into memory.\"\"\"\n batch_size = config.batch_size\n train_ds: tf.data.Dataset = tfds.load('mnist', split='train')\n test_ds: tf.data.Dataset = tfds.load('mnist', split='test')\n\n train_ds = train_ds.map(\n lambda sample: {\n 'image': tf.cast(sample['image'], tf.float32) / 255,\n 'label': sample['label'],\n }\n ) # normalize train set\n test_ds = test_ds.map(\n lambda sample: {\n 'image': tf.cast(sample['image'], tf.float32) / 255,\n 'label': sample['label'],\n }\n ) # normalize the test set.\n\n # Create a shuffled dataset by allocating a buffer size of 1024 to randomly\n # draw elements from.\n train_ds = train_ds.shuffle(1024)\n # Group into batches of `batch_size` and skip incomplete batches, prefetch the\n # next sample to improve latency.\n train_ds = train_ds.batch(batch_size, drop_remainder=True).prefetch(1)\n # Group into batches of `batch_size` and skip incomplete batches, prefetch the\n # next sample to improve latency.\n test_ds = test_ds.batch(batch_size, drop_remainder=True).prefetch(1)\n\n return train_ds, test_ds\n\n\ndef train_and_evaluate(config: ml_collections.ConfigDict, workdir: str) -> None:\n \"\"\"Execute model training and evaluation loop.\n\n Args:\n config: Hyperparameter configuration for training and evaluation.\n workdir: Directory path to store metrics.\n \"\"\"\n train_ds, test_ds = get_datasets(config)\n\n # Instantiate the model.\n model = CNN(rngs=nnx.Rngs(0))\n\n learning_rate = config.learning_rate\n momentum = config.momentum\n\n summary_writer = tensorboard.SummaryWriter(workdir)\n summary_writer.hparams(dict(config))\n\n optimizer = nnx.Optimizer(\n model, optax.sgd(learning_rate, momentum), wrt=nnx.Param\n )\n metrics = nnx.MultiMetric(\n accuracy=nnx.metrics.Accuracy(),\n loss=nnx.metrics.Average('loss'),\n )\n rngs = nnx.Rngs(0)\n\n for epoch in range(1, config.num_epochs + 1):\n # Run the optimization for one step and make a stateful update to the\n # following:\n # - The train state's model parameters\n # - The optimizer state\n # - The training loss and accuracy batch metrics\n model.train() # Switch to train mode\n\n for batch in train_ds.as_numpy_iterator():\n train_step(model, optimizer, metrics, batch, rngs)\n # Compute the training metrics.\n train_metrics = metrics.compute()\n metrics.reset() # Reset the metrics for the test set.\n\n # Compute the metrics on the test set after each training epoch.\n model.eval() # Switch to eval mode\n for batch in test_ds.as_numpy_iterator():\n eval_step(model, metrics, batch)\n\n # Compute the eval metrics.\n eval_metrics = metrics.compute()\n metrics.reset() # Reset the metrics for the next training epoch.\n\n logging.info( # pylint: disable=logging-not-lazy\n 'epoch:% 3d, train_loss: %.4f, train_accuracy: %.2f, test_loss: %.4f,'\n ' test_accuracy: %.2f'\n % (\n epoch,\n train_metrics['loss'],\n train_metrics['accuracy'] * 100,\n eval_metrics['loss'],\n eval_metrics['accuracy'] * 100,\n )\n )\n\n # Write the metrics to TensorBoard.\n summary_writer.scalar('train_loss', train_metrics['loss'], epoch)\n summary_writer.scalar('train_accuracy', train_metrics['accuracy'], epoch)\n summary_writer.scalar('test_loss', eval_metrics['loss'], epoch)\n summary_writer.scalar('test_accuracy', eval_metrics['accuracy'], epoch)\n\n summary_writer.flush()\n\n # Export the model to a SavedModel directory.\n from orbax.export import JaxModule, ExportManager, ServingConfig\n\n def exported_predict(model, y):\n return model(y, None)\n\n model.eval()\n jax_module = JaxModule(model, exported_predict)\n sig = [tf.TensorSpec(shape=(1, 28, 28, 1), dtype=tf.float32)]\n export_mgr = ExportManager(jax_module, [\n ServingConfig('mnist_server', input_signature=sig)\n ])\n\n output_dir= Path(workdir) / 'mnist_export'\n export_mgr.save(str(output_dir))\n" }, { "path": "examples/mnist/train_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for flax.examples.mnist.mnist_lib.\"\"\"\n\nimport pathlib\nimport tempfile\nimport sys\n\nfrom absl.testing import absltest\nimport jax\nimport flax.nnx as nnx\nfrom jax import numpy as jnp\nimport numpy as np\nimport tensorflow as tf\nimport tensorflow_datasets as tfds\n\nfrom configs import default\nimport train\n\n\nCNN_PARAMS = 825_034\n\n\nclass TrainTest(absltest.TestCase):\n \"\"\"Test cases for train.\"\"\"\n\n def setUp(self):\n super().setUp()\n if sys.version_info < (3, 13):\n self.skipTest('Tensorflow 2.20 required for this test, which conflicts with tensorflow_text.')\n # Make sure tf does not allocate gpu memory.\n tf.config.experimental.set_visible_devices([], \"GPU\")\n\n def test_cnn(self):\n \"\"\"Tests CNN module used as the trainable model.\"\"\"\n inputs = jnp.ones((1, 28, 28, 1), jnp.float32)\n cnn = train.CNN(nnx.Rngs(0))\n cnn.eval()\n output = cnn(inputs, None)\n\n self.assertEqual((1, 10), output.shape)\n\n def test_train_and_evaluate(self):\n \"\"\"Tests training and evaluation code by running a single step.\"\"\"\n # Create a temporary directory where tensorboard metrics are written.\n workdir = tempfile.mkdtemp()\n\n # Go two directories up to the root of the flax directory.\n flax_root_dir = pathlib.Path(__file__).parents[2]\n data_dir = str(flax_root_dir) + \"/.tfds/metadata\" # pylint: disable=unused-variable\n\n # Define training configuration.\n config = default.get_config()\n config.num_epochs = 1\n config.batch_size = 8\n\n with tfds.testing.mock_data(num_examples=8, data_dir=data_dir):\n train.train_and_evaluate(config=config, workdir=workdir)\n\n\nif __name__ == \"__main__\":\n absltest.main()\n" }, { "path": "examples/nlp_seq/README.md", "content": "## Part-of-Speech Tagging\nTrains a simple sequence-based part-of-speech tagger. The following sentence\nshows an example.\n\n```\nFrom|ADP the|DT AP|PROPN comes|VBZ this|DT story|NN :|:\n```\n\n### Requirements\n* Universal Dependency data sets: https://universaldependencies.org/#download.\n\n Download via command line:\n\n ```\n curl -# -o ud-treebanks-v2.0.tgz https://lindat.mff.cuni.cz/repository/xmlui/bitstream/handle/11234/1-1976/ud-treebanks-v2.0.tgz\n tar xzf ud-treebanks-v2.0.tgz\n ```\n\n### Supported setups\nThe model should run with other configurations and hardware, but explicitly tested on the following.\n\n| Hardware | Batch size | Learning rate | Training time | Accuracy | TensorBoard.dev |\n|:---:|:---:|:---:|:---:|:---:|:---:|\n| Nvidia Titan V (12GB) | 64 | 0.05 | 5:58h | 68.6% | [2022-05-01](https://tensorboard.dev/experiment/F5ULHlyzQlieVJn5PG8mRQ/) |\n\n### Running\n```\npython train.py --batch_size=64 --model_dir=./ancient_greek \\\n --dev=ud-treebanks-v2.0/UD_Ancient_Greek/grc-ud-dev.conllu \\\n --train=ud-treebanks-v2.0/UD_Ancient_Greek/grc-ud-train.conllu\n```\n" }, { "path": "examples/nlp_seq/configs/default.py", "content": "# Copyright 2025 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Default hyperparameters for NLP sequence tagging.\"\"\"\n\nimport ml_collections\n\n\ndef get_config():\n \"\"\"Get the default hyperparameter configuration.\"\"\"\n config = ml_collections.ConfigDict()\n\n # Model directory for checkpoints and logs\n config.model_dir = ''\n\n # Experiment name\n config.experiment = 'xpos'\n\n # Training hyperparameters\n config.batch_size = 64\n config.num_train_steps = 75000\n config.eval_frequency = 100\n\n # Optimizer hyperparameters\n config.learning_rate = 0.05\n config.weight_decay = 1e-1\n\n # Model hyperparameters\n config.max_length = 256\n\n # Random seed\n config.random_seed = 0\n\n # Data paths\n config.train = ''\n config.dev = ''\n\n return config\n" }, { "path": "examples/nlp_seq/input_pipeline.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Input pipeline for the sequence tagging dataset.\"\"\"\n\nimport codecs\nimport collections\nimport enum\n\nimport tensorflow as tf # pytype: disable=import-error\n\n\n# Values for padding, unknown words and a root.\nPAD = '

'\nPAD_ID = 0\n\nUNKNOWN = ''\nUNKNOWN_ID = 1\n\nROOT = ''\nROOT_ID = 2\n\n\nclass CoNLLAttributes(enum.Enum):\n \"\"\"CoNLL attributre names and indices.\n\n A UD CoNLL file looks like:\n 1 They they PRON PRP Case=Nom|Number=Plur 2 nsubj\n 2 buy buy VERB VBP Number=Plur|PTense=Pres 0 root\n 3 books book NOUN NNS Number=Plur 2 obj\n 4 . . PUNCT . _ 2 punct\n\n For details, please see: http://universaldependencies.org/format.html.\n \"\"\"\n\n ID = 0\n FORM = 1\n LEMMA = 2\n UPOS = 3\n XPOS = 4\n FEATS = 5\n HEAD = 6\n DEPREL = 7\n\n\ndef create_vocabs(filename, max_num_forms=100000):\n \"\"\"Loads corpus and create vocabulary lists.\n\n Args:\n filename: file name of a corpus.\n max_num_forms: maximum number of tokens included.\n\n Returns:\n Dictionary containing named vocab dictionaries.\n\n \"\"\"\n form_counter = collections.Counter()\n xpos_counter = collections.Counter()\n with tf.io.gfile.GFile(filename, 'rb') as f:\n for line in codecs.getreader('utf-8')(f):\n line = line.strip()\n split = line.split('\\t')\n if not line.startswith('#') and split[0]:\n form_counter[split[CoNLLAttributes.FORM.value]] += 1\n xpos_counter[split[CoNLLAttributes.XPOS.value]] += 1\n\n special_tokens = {PAD: PAD_ID, UNKNOWN: UNKNOWN_ID, ROOT: ROOT_ID}\n\n # create word form vocab\n vocabs = {'forms': {}, 'xpos': {}}\n vocabs['forms'].update(special_tokens)\n vocabs['forms'].update(\n {\n form[0]: id\n for id, form in enumerate(\n form_counter.most_common(max_num_forms), start=ROOT_ID + 1\n )\n }\n )\n\n # create xpos vocab\n vocabs['xpos'].update(special_tokens)\n vocabs['xpos'].update(\n {\n tag[0]: id\n for id, tag in enumerate(\n xpos_counter.most_common(), start=ROOT_ID + 1\n )\n }\n )\n\n return vocabs\n\n\ndef create_token(token, attributes, vocabs):\n \"\"\"Map for a token a selected subset of attributes to indices.\n\n Input example: CoNLL 09 representation for a token.\n ['Ms.', 'ms.', 'ms.', 'NNP', '_', '2', 'TITLE]\n Output example: Indices as defined in self._attributes, e.g., [word form,\n part-of-speech tag, and head].\n [1025, 3, 1]\n\n Args:\n token: CoNLL token attributes.\n attributes: selected attributes.\n vocabs: dictionary of vocabs.\n\n Returns:\n List of attribute ids for a token, e.g. [1025, 3] with word id and pos id.\n\n Raises:\n ValueError: CoNLL attribute requested but not covered by mapping.\n \"\"\"\n selected_attributes = []\n for attribute in attributes:\n index = attribute.value\n if attribute == CoNLLAttributes.FORM:\n selected_attributes.append(vocabs['forms'].get(token[index], UNKNOWN_ID))\n elif attribute == CoNLLAttributes.XPOS:\n selected_attributes.append(vocabs['xpos'].get(token[index], UNKNOWN_ID))\n elif attribute == CoNLLAttributes.HEAD:\n selected_attributes.append(int(token[index]))\n else:\n raise ValueError(\n 'CoNLL index %s not covered by mapping.' % str(attribute.name)\n )\n return selected_attributes\n\n\ndef create_sentence_with_root(attributes, vocabs):\n \"\"\"Create a sentence containing a root.\n\n Args:\n attributes: attributes extracted from token.\n vocabs: dictionary of vocabs.\n\n Returns:\n A list representing a sentence containing the root only,\n e.g., [[2, 1, 0]] for root word, unknown xpos, and head 0.\n \"\"\"\n # Create the token properties of an artificial root node.\n token_properties = [ROOT for _ in range(12)] # CoNLL 09 has 12 columns.\n token_properties[CoNLLAttributes.ID.value] = '0'\n token_properties[CoNLLAttributes.HEAD.value] = '0'\n token = create_token(token_properties, attributes, vocabs)\n if len(token) == 1:\n token = token[0]\n return [token]\n\n\ndef sentences_from_conll_data(\n corpus_filename, vocabs, attributes, max_sentence_length=1000\n):\n \"\"\"Load and returns conll data in list format.\n\n Args:\n corpus_filename: filename of corpus.\n vocabs: dictionary of vocabs\n attributes: list of conll attributes to include into the batch\n max_sentence_length: cut off sentences longer as max tokens\n\n Yields:\n A sentence as a list of tokens while tokens are lists of attributes.\n \"\"\"\n with tf.io.gfile.GFile(corpus_filename, 'rb') as f:\n sentence = create_sentence_with_root(attributes, vocabs)\n for line in codecs.getreader('utf-8')(f):\n line = line.strip()\n if line.startswith('#'):\n continue\n split = line.split('\\t')\n if split[0]: # Not an empty line, process next token:\n if len(sentence) < max_sentence_length:\n if len(attributes) == 1:\n sentence.append(create_token(split, attributes, vocabs)[0])\n else:\n sentence.append(create_token(split, attributes, vocabs))\n else: # Sentences start with an empty line, yield sentence:\n yield sentence\n\n # Reset sentence.\n sentence = create_sentence_with_root(attributes, vocabs)\n if len(sentence) > 1: # sentences does not only contain a root.\n yield sentence\n\n\ndef sentence_dataset_dict(\n filename,\n vocabs,\n attributes_input,\n attributes_target,\n batch_size,\n bucket_size,\n repeat=None,\n prefetch_size=tf.data.experimental.AUTOTUNE,\n):\n \"\"\"Combines sentences into a dataset of padded batches.\n\n Args:\n filename: file name of a corpus.\n vocabs: dictionary of dictionaries to map from strings to ids.\n attributes_input: attributes for the input.\n attributes_target: target attributes empty targets is not included.\n batch_size: the size of a batch.\n bucket_size: the size of a bucket.\n repeat: number of times the dataset is repeated.\n prefetch_size: prefetch size of the data.\n\n Returns:\n Returns dataset as dictionary containing the data as key value pairs.\n \"\"\"\n data_keys = ['inputs']\n if attributes_target:\n data_keys.append('targets')\n\n def generator():\n \"\"\"Generator to create the data.\"\"\"\n input_generator = sentences_from_conll_data(\n filename, vocabs, attributes_input, max_sentence_length=bucket_size\n )\n\n if attributes_target:\n target_generator = sentences_from_conll_data(\n filename, vocabs, attributes_target, max_sentence_length=bucket_size\n )\n\n for inputs in input_generator:\n data = {'inputs': inputs}\n if attributes_target:\n data['targets'] = next(target_generator)\n yield data\n\n output_types = {k: tf.float32 for k in data_keys}\n output_shapes = {k: (None,) for k in data_keys}\n dataset = tf.data.Dataset.from_generator(\n generator, output_types=output_types, output_shapes=output_shapes\n )\n\n # cache the dataset in memory and repeat.\n dataset = dataset.cache()\n dataset = dataset.repeat(repeat)\n\n # static padding up to bucket size.\n padded_shapes = {k: [bucket_size] for k in data_keys}\n dataset = dataset.padded_batch(\n batch_size=batch_size, padded_shapes=(padded_shapes)\n )\n\n dataset = dataset.prefetch(prefetch_size)\n return dataset\n" }, { "path": "examples/nlp_seq/input_pipeline_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for flax.examples.nlp.input_pipeline.\"\"\"\n\nimport os\n\nfrom absl.testing import absltest\nimport jax\nimport tensorflow as tf\n\nimport input_pipeline\n\n\n# Parse absl flags test_srcdir and test_tmpdir.\njax.config.parse_flags_with_absl()\n\n\nCONLL_DATA = \"\"\"1\\tThey\\tthey\\tPRON\\tPRP\\tCase=Nom|Number=Plur\\t2\\tnsubj\n2\\tbuy\\tbuy\\t VERB\\tVBP\\tNumber=Plur|PTense=Pres\\t0\\troot\n3\\tbooks\\tbook\\tNOUN\\tNNS\\tNumber=Plur\\t2\\tobj\n4\\t.\\t.\\tPUNCT\\t.\\t_\\t2\\tpunct\n\n1\\tThey\\tthey\\tPRON\\tPRP\\tCase=Nom|Number=Plur\\t2\\tnsubj\n2\\tbuy\\tbuy\\t VERB\\tVBP\\tNumber=Plur|PTense=Pres\\t0\\troot\n3\\tbooks\\tbook\\tNOUN\\tNNS\\tNumber=Plur\\t2\\tobj\n4\\t.\\t.\\tPUNCT\\t.\\t_\\t2\\tpunct\n\n1\\tNY\\tNY\\tNOUN\\tNNS\\tNumber=Singular\\t0\\troot\n\"\"\"\n\n\nclass InputPipelineTest(absltest.TestCase):\n\n def setUp(self):\n super().setUp()\n self.test_tmpdir = self.create_tempdir()\n\n # Write a sample corpus.\n self._filename = os.path.join(self.test_tmpdir.full_path, 'data.conll')\n with tf.io.gfile.GFile(self._filename, 'w') as f:\n # The CoNLL data has to end with an empty line.\n f.write(CONLL_DATA)\n f.write('\\n')\n\n def test_vocab_creation(self):\n \"\"\"Tests the creation of the vocab.\"\"\"\n vocabs = input_pipeline.create_vocabs(self._filename)\n self.assertEqual(\n vocabs['forms'],\n {\n '

': 0,\n '': 1,\n '': 2,\n 'They': 3,\n 'buy': 4,\n 'books': 5,\n '.': 6,\n 'NY': 7,\n },\n )\n\n def testInputBatch(self):\n \"\"\"Test the batching of the dataset.\"\"\"\n vocabs = input_pipeline.create_vocabs(self._filename)\n\n attributes_input = [input_pipeline.CoNLLAttributes.FORM]\n attributes_target = [] # empty target for tagging of unlabeled data.\n sentence_dataset = input_pipeline.sentence_dataset_dict(\n self._filename,\n vocabs,\n attributes_input,\n attributes_target,\n batch_size=2,\n bucket_size=10,\n repeat=1,\n )\n\n sentence_dataset_iter = iter(sentence_dataset)\n\n batch = next(sentence_dataset_iter)\n inputs = batch['inputs'].numpy().tolist()\n self.assertSameStructure(\n inputs,\n [\n [2.0, 3.0, 4.0, 5.0, 6.0, 0.0, 0.0, 0.0, 0.0, 0.0],\n [2.0, 3.0, 4.0, 5.0, 6.0, 0.0, 0.0, 0.0, 0.0, 0.0],\n ],\n )\n self.assertLen(batch, 1) # make sure target is not included.\n\n def testInputTargetBatch(self):\n \"\"\"Test the batching of the dataset.\"\"\"\n vocabs = input_pipeline.create_vocabs(self._filename)\n\n attributes_input = [input_pipeline.CoNLLAttributes.FORM]\n attributes_target = [input_pipeline.CoNLLAttributes.XPOS]\n sentence_dataset = input_pipeline.sentence_dataset_dict(\n self._filename,\n vocabs,\n attributes_input,\n attributes_target,\n batch_size=2,\n bucket_size=10,\n repeat=1,\n )\n\n sentence_dataset_iter = iter(sentence_dataset)\n\n batch = next(sentence_dataset_iter)\n inputs = batch['inputs'].numpy().tolist()\n self.assertSameStructure(\n inputs,\n [\n [2.0, 3.0, 4.0, 5.0, 6.0, 0.0, 0.0, 0.0, 0.0, 0.0],\n [2.0, 3.0, 4.0, 5.0, 6.0, 0.0, 0.0, 0.0, 0.0, 0.0],\n ],\n )\n targets = batch['targets'].numpy().tolist()\n self.assertSameStructure(\n targets,\n [\n [2.0, 4.0, 5.0, 3.0, 6.0, 0.0, 0.0, 0.0, 0.0, 0.0],\n [2.0, 4.0, 5.0, 3.0, 6.0, 0.0, 0.0, 0.0, 0.0, 0.0],\n ],\n )\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "examples/nlp_seq/main.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Main file for running the NLP sequence tagging example.\n\nThis file is intentionally kept short to allow config-based execution.\n\"\"\"\n\nfrom absl import app\nfrom absl import flags\nimport train\nfrom ml_collections import config_flags\n\n\nFLAGS = flags.FLAGS\n\nconfig_flags.DEFINE_config_file(\n 'config',\n 'configs/default.py',\n 'File path to the training hyperparameter configuration.',\n lock_config=True,\n)\n\n\ndef main(argv):\n if len(argv) > 1:\n raise app.UsageError('Too many command-line arguments.')\n\n # Convert config to FLAGS for train.py compatibility\n config = FLAGS.config\n\n # Override FLAGS with config values\n FLAGS.model_dir = config.model_dir\n FLAGS.experiment = config.experiment\n FLAGS.batch_size = config.batch_size\n FLAGS.eval_frequency = config.eval_frequency\n FLAGS.num_train_steps = config.num_train_steps\n FLAGS.learning_rate = config.learning_rate\n FLAGS.weight_decay = config.weight_decay\n FLAGS.max_length = config.max_length\n FLAGS.random_seed = config.random_seed\n FLAGS.train = config.train\n FLAGS.dev = config.dev\n\n # Run the training\n train.main(argv)\n\n\nif __name__ == '__main__':\n app.run(main)\n" }, { "path": "examples/nlp_seq/models.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Transformer-based language models.\"\"\"\n\nfrom typing import Any, Optional\nfrom collections.abc import Callable\n\nfrom flax import linen as nn\nfrom flax import struct\nimport jax.numpy as jnp\nimport numpy as np\n\n\n@struct.dataclass\nclass TransformerConfig:\n \"\"\"Global hyperparameters used to minimize obnoxious kwarg plumbing.\"\"\"\n\n vocab_size: int\n output_vocab_size: int\n dtype: Any = jnp.float32\n emb_dim: int = 512\n num_heads: int = 8\n num_layers: int = 6\n qkv_dim: int = 512\n mlp_dim: int = 2048\n max_len: int = 2048\n dropout_rate: float = 0.3\n attention_dropout_rate: float = 0.3\n kernel_init: Callable = nn.initializers.xavier_uniform()\n bias_init: Callable = nn.initializers.normal(stddev=1e-6)\n posemb_init: Callable | None = None\n\n\ndef sinusoidal_init(max_len=2048):\n \"\"\"1D Sinusoidal Position Embedding Initializer.\n\n Args:\n max_len: maximum possible length for the input\n\n Returns:\n output: init function returning `(1, max_len, d_feature)`\n \"\"\"\n\n def init(key, shape, dtype=np.float32):\n \"\"\"Sinusoidal init.\"\"\"\n del key, dtype\n d_feature = shape[-1]\n pe = np.zeros((max_len, d_feature), dtype=np.float32)\n position = np.arange(0, max_len)[:, np.newaxis]\n div_term = np.exp(\n np.arange(0, d_feature, 2) * -(np.log(10000.0) / d_feature)\n )\n pe[:, 0::2] = np.sin(position * div_term)\n pe[:, 1::2] = np.cos(position * div_term)\n pe = pe[np.newaxis, :, :] # [1, max_len, d_feature]\n return jnp.array(pe)\n\n return init\n\n\nclass AddPositionEmbs(nn.Module):\n \"\"\"Adds (optionally learned) positional embeddings to the inputs.\n\n Attributes:\n config: TransformerConfig dataclass containing hyperparameters.\n \"\"\"\n\n config: TransformerConfig\n\n @nn.compact\n def __call__(self, inputs):\n \"\"\"Applies AddPositionEmbs module.\n\n By default this layer uses a fixed sinusoidal embedding table. If a\n learned position embedding is desired, pass an initializer to\n posemb_init in the configuration.\n\n Args:\n inputs: input data.\n\n Returns:\n output: `(bs, timesteps, in_dim)`\n \"\"\"\n config = self.config\n # inputs.shape is (batch_size, seq_len, emb_dim)\n assert inputs.ndim == 3, (\n 'Number of dimensions should be 3, but it is: %d' % inputs.ndim\n )\n length = inputs.shape[1]\n pos_emb_shape = (1, config.max_len, inputs.shape[-1])\n if config.posemb_init is None:\n # Use a fixed (non-learned) sinusoidal position embedding.\n pos_embedding = sinusoidal_init(max_len=config.max_len)(\n None, pos_emb_shape, None\n )\n else:\n pos_embedding = self.param(\n 'pos_embedding', config.posemb_init, pos_emb_shape\n )\n pe = pos_embedding[:, :length, :]\n return inputs + pe\n\n\nclass MlpBlock(nn.Module):\n \"\"\"Transformer MLP / feed-forward block.\n\n Attributes:\n config: TransformerConfig dataclass containing hyperparameters.\n out_dim: optionally specify out dimension.\n \"\"\"\n\n config: TransformerConfig\n out_dim: int | None = None\n\n @nn.compact\n def __call__(self, inputs, deterministic=True):\n \"\"\"Applies Transformer MlpBlock module.\"\"\"\n config = self.config\n actual_out_dim = inputs.shape[-1] if self.out_dim is None else self.out_dim\n x = nn.Dense(\n config.mlp_dim,\n dtype=config.dtype,\n kernel_init=config.kernel_init,\n bias_init=config.bias_init,\n )(inputs)\n x = nn.elu(x)\n x = nn.Dropout(rate=config.dropout_rate)(x, deterministic=deterministic)\n output = nn.Dense(\n actual_out_dim,\n dtype=config.dtype,\n kernel_init=config.kernel_init,\n bias_init=config.bias_init,\n )(x)\n output = nn.Dropout(rate=config.dropout_rate)(\n output, deterministic=deterministic\n )\n return output\n\n\nclass Encoder1DBlock(nn.Module):\n \"\"\"Transformer encoder layer.\n\n Attributes:\n config: TransformerConfig dataclass containing hyperparameters.\n \"\"\"\n\n config: TransformerConfig\n\n @nn.compact\n def __call__(self, inputs, deterministic):\n \"\"\"Applies Encoder1DBlock module.\n\n Args:\n inputs: input data.\n deterministic: if true dropout is applied otherwise not.\n\n Returns:\n output after transformer encoder block.\n \"\"\"\n config = self.config\n\n # Attention block.\n assert inputs.ndim == 3\n x = nn.LayerNorm(dtype=config.dtype)(inputs)\n x = nn.MultiHeadDotProductAttention(\n num_heads=config.num_heads,\n dtype=config.dtype,\n qkv_features=config.qkv_dim,\n kernel_init=config.kernel_init,\n bias_init=config.bias_init,\n use_bias=False,\n broadcast_dropout=False,\n dropout_rate=config.attention_dropout_rate,\n deterministic=deterministic,\n )(x)\n\n x = nn.Dropout(rate=config.dropout_rate)(x, deterministic=deterministic)\n x = x + inputs\n\n # MLP block.\n y = nn.LayerNorm(dtype=config.dtype)(x)\n y = MlpBlock(config=config)(y, deterministic=deterministic)\n return x + y\n\n\nclass Transformer(nn.Module):\n \"\"\"Transformer Model for sequence tagging.\"\"\"\n\n config: TransformerConfig\n\n @nn.compact\n def __call__(self, *, inputs, train):\n \"\"\"Applies Transformer model on the inputs.\n\n Args:\n inputs: input data\n train: if it is training.\n\n Returns:\n output of a transformer encoder.\n\n \"\"\"\n assert inputs.ndim == 2 # (batch, len)\n\n config = self.config\n\n x = inputs.astype('int32')\n x = nn.Embed(\n num_embeddings=config.vocab_size, features=config.emb_dim, name='embed'\n )(x)\n x = nn.Dropout(rate=config.dropout_rate)(x, deterministic=not train)\n x = AddPositionEmbs(config)(x)\n\n for _ in range(config.num_layers):\n x = Encoder1DBlock(config)(x, deterministic=not train)\n\n x = nn.LayerNorm(dtype=config.dtype)(x)\n logits = nn.Dense(\n config.output_vocab_size,\n kernel_init=config.kernel_init,\n bias_init=config.bias_init,\n )(x)\n return logits\n" }, { "path": "examples/nlp_seq/requirements.txt", "content": "absl-py==1.0.0\nflax==0.3.6\nnumpy==1.22.0\ntensorflow==2.11.1\n" }, { "path": "examples/nlp_seq/train.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Sequence Tagging example.\n\nThis script trains a Transformer on the Universal dependency dataset.\n\"\"\"\n\nimport functools\nimport os\nimport time\n\nfrom absl import app\nfrom absl import flags\nfrom absl import logging\nfrom flax import jax_utils\nfrom flax import linen as nn\nfrom flax.metrics import tensorboard\nfrom flax.training import common_utils\nfrom flax.training import train_state\nimport jax\nimport jax.numpy as jnp\nfrom jax import random\nimport numpy as np\nimport optax\nimport tensorflow as tf\n\nimport input_pipeline\nimport models\n\n\nFLAGS = flags.FLAGS\n\nflags.DEFINE_string('model_dir', default='', help='Directory for model data.')\n\nflags.DEFINE_string('experiment', default='xpos', help='Experiment name.')\n\nflags.DEFINE_integer('batch_size', default=64, help='Batch size for training.')\n\nflags.DEFINE_integer(\n 'eval_frequency',\n default=100,\n help='Frequency of eval during training, e.g. every 1000 steps.',\n)\n\nflags.DEFINE_integer(\n 'num_train_steps', default=75000, help='Number of train steps.'\n)\n\nflags.DEFINE_float('learning_rate', default=0.05, help='Learning rate.')\n\nflags.DEFINE_float(\n 'weight_decay',\n default=1e-1,\n help='Decay factor for AdamW style weight decay.',\n)\n\nflags.DEFINE_integer(\n 'max_length', default=256, help='Maximum length of examples.'\n)\n\nflags.DEFINE_integer(\n 'random_seed', default=0, help='Integer for PRNG random seed.'\n)\n\nflags.DEFINE_string('train', default='', help='Path to training data.')\n\nflags.DEFINE_string('dev', default='', help='Path to development data.')\n\n\ndef create_learning_rate_scheduler(\n factors='constant * linear_warmup * rsqrt_decay',\n base_learning_rate=0.5,\n warmup_steps=8000,\n decay_factor=0.5,\n steps_per_decay=20000,\n steps_per_cycle=100000,\n):\n \"\"\"creates learning rate schedule.\n\n Interprets factors in the factors string which can consist of:\n * constant: interpreted as the constant value,\n * linear_warmup: interpreted as linear warmup until warmup_steps,\n * rsqrt_decay: divide by square root of max(step, warmup_steps)\n * decay_every: Every k steps decay the learning rate by decay_factor.\n * cosine_decay: Cyclic cosine decay, uses steps_per_cycle parameter.\n\n Args:\n factors: a string with factors separated by '*' that defines the schedule.\n base_learning_rate: float, the starting constant for the lr schedule.\n warmup_steps: how many steps to warm up for in the warmup schedule.\n decay_factor: The amount to decay the learning rate by.\n steps_per_decay: How often to decay the learning rate.\n steps_per_cycle: Steps per cycle when using cosine decay.\n\n Returns:\n a function learning_rate(step): float -> {'learning_rate': float}, the\n step-dependent lr.\n \"\"\"\n factors = [n.strip() for n in factors.split('*')]\n\n def step_fn(step):\n \"\"\"Step to learning rate function.\"\"\"\n ret = 1.0\n for name in factors:\n if name == 'constant':\n ret *= base_learning_rate\n elif name == 'linear_warmup':\n ret *= jnp.minimum(1.0, step / warmup_steps)\n elif name == 'rsqrt_decay':\n ret /= jnp.sqrt(jnp.maximum(step, warmup_steps))\n elif name == 'rsqrt_normalized_decay':\n ret *= jnp.sqrt(warmup_steps)\n ret /= jnp.sqrt(jnp.maximum(step, warmup_steps))\n elif name == 'decay_every':\n ret *= decay_factor ** (step // steps_per_decay)\n elif name == 'cosine_decay':\n progress = jnp.maximum(\n 0.0, (step - warmup_steps) / float(steps_per_cycle)\n )\n ret *= jnp.maximum(\n 0.0, 0.5 * (1.0 + jnp.cos(jnp.pi * (progress % 1.0)))\n )\n else:\n raise ValueError('Unknown factor %s.' % name)\n return jnp.asarray(ret, dtype=jnp.float32)\n\n return step_fn\n\n\ndef compute_weighted_cross_entropy(logits, targets, weights=None):\n \"\"\"Compute weighted cross entropy and entropy for log probs and targets.\n\n Args:\n logits: [batch, length, num_classes] float array.\n targets: categorical targets [batch, length] int array.\n weights: None or array of shape [batch x length]\n\n Returns:\n Tuple of scalar loss and batch normalizing factor.\n \"\"\"\n if logits.ndim != targets.ndim + 1:\n raise ValueError(\n 'Incorrect shapes. Got shape %s logits and %s targets'\n % (str(logits.shape), str(targets.shape))\n )\n onehot_targets = common_utils.onehot(targets, logits.shape[-1])\n loss = -jnp.sum(onehot_targets * nn.log_softmax(logits), axis=-1)\n normalizing_factor = onehot_targets.sum()\n if weights is not None:\n loss = loss * weights\n normalizing_factor = weights.sum()\n\n return loss.sum(), normalizing_factor\n\n\ndef compute_weighted_accuracy(logits, targets, weights=None):\n \"\"\"Compute weighted accuracy for log probs and targets.\n\n Args:\n logits: [batch, length, num_classes] float array.\n targets: categorical targets [batch, length] int array.\n weights: None or array of shape [batch x length]\n\n Returns:\n Tuple of scalar accuracy and batch normalizing factor.\n \"\"\"\n if logits.ndim != targets.ndim + 1:\n raise ValueError(\n 'Incorrect shapes. Got shape %s logits and %s targets'\n % (str(logits.shape), str(targets.shape))\n )\n loss = jnp.equal(jnp.argmax(logits, axis=-1), targets)\n normalizing_factor = np.prod(logits.shape[:-1])\n if weights is not None:\n loss = loss * weights\n normalizing_factor = weights.sum()\n\n return loss.sum(), normalizing_factor\n\n\ndef compute_metrics(logits, labels, weights):\n \"\"\"Compute summary metrics.\"\"\"\n loss, weight_sum = compute_weighted_cross_entropy(logits, labels, weights)\n acc, _ = compute_weighted_accuracy(logits, labels, weights)\n metrics = {\n 'loss': loss,\n 'accuracy': acc,\n 'denominator': weight_sum,\n }\n metrics = np.sum(metrics, -1)\n return metrics\n\n\ndef train_step(state, batch, model, learning_rate_fn, dropout_rng=None):\n \"\"\"Perform a single training step.\"\"\"\n train_keys = ['inputs', 'targets']\n (inputs, targets) = (batch.get(k, None) for k in train_keys)\n\n weights = jnp.where(targets > 0, 1, 0).astype(jnp.float32)\n\n dropout_rng = jax.random.fold_in(dropout_rng, state.step)\n\n def loss_fn(params):\n \"\"\"loss function used for training.\"\"\"\n logits = model.apply(\n {'params': params},\n inputs=inputs,\n train=True,\n rngs={'dropout': dropout_rng},\n )\n loss, weight_sum = compute_weighted_cross_entropy(logits, targets, weights)\n\n mean_loss = loss / weight_sum\n return mean_loss, logits\n\n lr = learning_rate_fn(state.step)\n grad_fn = jax.value_and_grad(loss_fn, has_aux=True)\n (_, logits), grads = grad_fn(state.params)\n grads = jax.lax.pmean(grads, 'batch')\n new_state = state.apply_gradients(grads=grads)\n metrics = compute_metrics(logits, targets, weights)\n metrics['learning_rate'] = lr\n\n return new_state, metrics\n\n\ndef pad_examples(x, desired_batch_size):\n \"\"\"Expand batch to desired size by zeros with the shape of last slice.\"\"\"\n batch_pad = desired_batch_size - x.shape[0]\n # Padding with zeros to avoid that they get counted in compute_metrics.\n return np.concatenate([x, np.tile(np.zeros_like(x[-1]), (batch_pad, 1))])\n\n\ndef main(argv):\n if len(argv) > 1:\n raise app.UsageError('Too many command-line arguments.')\n\n # Make sure tf does not allocate gpu memory.\n tf.config.experimental.set_visible_devices([], 'GPU')\n\n batch_size = FLAGS.batch_size\n learning_rate = FLAGS.learning_rate\n num_train_steps = FLAGS.num_train_steps\n eval_freq = FLAGS.eval_frequency\n random_seed = FLAGS.random_seed\n\n if not FLAGS.dev:\n raise app.UsageError('Please provide path to dev set.')\n if not FLAGS.train:\n raise app.UsageError('Please provide path to training set.')\n if batch_size % jax.device_count() > 0:\n raise ValueError('Batch size must be divisible by the number of devices')\n device_batch_size = batch_size // jax.device_count()\n\n if jax.process_index() == 0:\n train_summary_writer = tensorboard.SummaryWriter(\n os.path.join(FLAGS.model_dir, FLAGS.experiment + '_train')\n )\n eval_summary_writer = tensorboard.SummaryWriter(\n os.path.join(FLAGS.model_dir, FLAGS.experiment + '_eval')\n )\n\n # create the training and development dataset\n vocabs = input_pipeline.create_vocabs(FLAGS.train)\n config = models.TransformerConfig(\n vocab_size=len(vocabs['forms']),\n output_vocab_size=len(vocabs['xpos']),\n max_len=FLAGS.max_length,\n )\n\n attributes_input = [input_pipeline.CoNLLAttributes.FORM]\n attributes_target = [input_pipeline.CoNLLAttributes.XPOS]\n train_ds = input_pipeline.sentence_dataset_dict(\n FLAGS.train,\n vocabs,\n attributes_input,\n attributes_target,\n batch_size=batch_size,\n bucket_size=config.max_len,\n )\n train_iter = iter(train_ds)\n\n eval_ds = input_pipeline.sentence_dataset_dict(\n FLAGS.dev,\n vocabs,\n attributes_input,\n attributes_target,\n batch_size=batch_size,\n bucket_size=config.max_len,\n repeat=1,\n )\n\n model = models.Transformer(config)\n\n rng = random.key(random_seed)\n rng, init_rng = random.split(rng)\n\n # call a jitted initialization function to get the initial parameter tree\n @jax.jit\n def initialize_variables(init_rng):\n init_batch = jnp.ones((config.max_len, 1), jnp.float32)\n init_variables = model.init(init_rng, inputs=init_batch, train=False)\n return init_variables\n\n init_variables = initialize_variables(init_rng)\n\n learning_rate_fn = create_learning_rate_scheduler(\n base_learning_rate=learning_rate\n )\n\n optimizer = optax.adamw(\n learning_rate_fn, b1=0.9, b2=0.98, eps=1e-9, weight_decay=1e-1\n )\n state = train_state.TrainState.create(\n apply_fn=model.apply, params=init_variables['params'], tx=optimizer\n )\n\n # Replicate optimizer.\n state = jax_utils.replicate(state)\n\n p_train_step = jax.pmap(\n functools.partial(\n train_step, model=model, learning_rate_fn=learning_rate_fn\n ),\n axis_name='batch',\n donate_argnums=(0,),\n ) # pytype: disable=wrong-arg-types\n\n def eval_step(params, batch):\n \"\"\"Calculate evaluation metrics on a batch.\"\"\"\n inputs, targets = batch['inputs'], batch['targets']\n weights = jnp.where(targets > 0, 1.0, 0.0)\n logits = model.apply({'params': params}, inputs=inputs, train=False)\n return compute_metrics(logits, targets, weights)\n\n p_eval_step = jax.pmap(eval_step, axis_name='batch')\n\n # We init the first set of dropout PRNG keys, but update it afterwards inside\n # the main pmap'd training update for performance.\n dropout_rngs = random.split(rng, jax.local_device_count())\n metrics_all = []\n tick = time.time()\n best_dev_score = 0\n for step, batch in zip(range(num_train_steps), train_iter):\n batch = common_utils.shard(\n jax.tree_util.tree_map(lambda x: x._numpy(), batch)\n ) # pylint: disable=protected-access\n\n state, metrics = p_train_step(state, batch, dropout_rng=dropout_rngs)\n metrics_all.append(metrics)\n if (step + 1) % eval_freq == 0:\n metrics_all = common_utils.get_metrics(metrics_all)\n lr = metrics_all.pop('learning_rate').mean()\n metrics_sums = jax.tree_util.tree_map(jnp.sum, metrics_all)\n denominator = metrics_sums.pop('denominator')\n summary = jax.tree_util.tree_map(lambda x: x / denominator, metrics_sums) # pylint: disable=cell-var-from-loop\n summary['learning_rate'] = lr\n logging.info('train in step: %d, loss: %.4f', step, summary['loss'])\n if jax.process_index() == 0:\n tock = time.time()\n steps_per_sec = eval_freq / (tock - tick)\n tick = tock\n train_summary_writer.scalar('steps per second', steps_per_sec, step)\n for key, val in summary.items():\n train_summary_writer.scalar(key, val, step)\n train_summary_writer.flush()\n\n metrics_all = [] # reset metric accumulation for next evaluation cycle.\n\n eval_metrics = []\n eval_iter = iter(eval_ds)\n\n for eval_batch in eval_iter:\n eval_batch = jax.tree_util.tree_map(lambda x: x._numpy(), eval_batch) # pylint: disable=protected-access\n # Handle final odd-sized batch by padding instead of dropping it.\n cur_pred_batch_size = eval_batch['inputs'].shape[0]\n if cur_pred_batch_size != batch_size:\n # pad up to batch size\n eval_batch = jax.tree_util.tree_map(\n lambda x: pad_examples(x, batch_size), eval_batch\n )\n eval_batch = common_utils.shard(eval_batch)\n\n metrics = p_eval_step(state.params, eval_batch)\n\n eval_metrics.append(metrics)\n eval_metrics = common_utils.get_metrics(eval_metrics)\n eval_metrics_sums = jax.tree_util.tree_map(jnp.sum, eval_metrics)\n eval_denominator = eval_metrics_sums.pop('denominator')\n eval_summary = jax.tree_util.tree_map(\n lambda x: x / eval_denominator, # pylint: disable=cell-var-from-loop\n eval_metrics_sums,\n )\n\n logging.info(\n 'eval in step: %d, loss: %.4f, accuracy: %.4f',\n step,\n eval_summary['loss'],\n eval_summary['accuracy'],\n )\n\n if best_dev_score < eval_summary['accuracy']:\n best_dev_score = eval_summary['accuracy']\n # TODO: save model.\n eval_summary['best_dev_score'] = best_dev_score\n logging.info('best development model score %.4f', best_dev_score)\n if jax.process_index() == 0:\n for key, val in eval_summary.items():\n eval_summary_writer.scalar(key, val, step)\n eval_summary_writer.flush()\n\n\nif __name__ == '__main__':\n app.run(main)\n" }, { "path": "examples/nnx_toy_examples/01_functional_api.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# %%\nimport jax\nimport jax.numpy as jnp\nimport matplotlib.pyplot as plt\nimport numpy as np\n\nfrom flax import nnx\n\nX = np.linspace(-jnp.pi, jnp.pi, 100)[:, None]\nY = 0.8 * jnp.sin(X) + 0.1 + np.random.normal(0, 0.1, size=X.shape)\n\n\ndef dataset(batch_size):\n while True:\n idx = np.random.choice(len(X), size=batch_size)\n yield X[idx], Y[idx]\n\n\nclass Linear(nnx.Module):\n def __init__(self, din: int, dout: int, *, rngs: nnx.Rngs):\n self.w = nnx.Param(jax.random.uniform(rngs.params(), (din, dout)))\n self.b = nnx.Param(jnp.zeros((dout,)))\n\n def __call__(self, x):\n return x @ self.w.value + self.b.value\n\n\nclass Count(nnx.Variable[nnx.A]):\n pass\n\n\nclass MLP(nnx.Module):\n def __init__(self, din, dhidden, dout, *, rngs: nnx.Rngs):\n self.count = Count(jnp.array(0))\n self.linear1 = Linear(din, dhidden, rngs=rngs)\n self.linear2 = Linear(dhidden, dout, rngs=rngs)\n\n def __call__(self, x):\n self.count[...] += 1\n return self.linear2(jax.nn.relu(self.linear1(x) * 0.5))\n\n\ngraphdef, params, counts = nnx.split(\n MLP(din=1, dhidden=32, dout=1, rngs=nnx.Rngs(0)), nnx.Param, Count\n)\n\n\n@jax.jit\ndef train_step(params, counts, batch):\n x, y = batch\n\n def loss_fn(params):\n model = nnx.merge(graphdef, params, counts)\n y_pred = model(x)\n new_counts = nnx.state(model, Count)\n loss = jnp.mean((y - y_pred) ** 2)\n return loss, new_counts\n\n grad, counts = jax.grad(loss_fn, has_aux=True)(params)\n # |-------- sgd ---------|\n params = jax.tree.map(lambda w, g: w - 0.1 * g, params, grad)\n\n return params, counts\n\n\n@jax.jit\ndef test_step(params: nnx.State, counts: nnx.State, batch):\n x, y = batch\n model = nnx.merge(graphdef, params, counts)\n y_pred = model(x)\n loss = jnp.mean((y - y_pred) ** 2)\n return {'loss': loss}\n\n\ntotal_steps = 10_000\nfor step, batch in enumerate(dataset(32)):\n params, counts = train_step(params, counts, batch)\n\n if step % 1000 == 0:\n logs = test_step(params, counts, (X, Y))\n print(f\"step: {step}, loss: {logs['loss']}\")\n\n if step >= total_steps - 1:\n break\n\nmodel = nnx.merge(graphdef, params, counts)\nprint('times called:', model.count.value)\n\ny_pred = model(X)\n\nplt.scatter(X, Y, color='blue')\nplt.plot(X, y_pred, color='black')\nplt.show()\n" }, { "path": "examples/nnx_toy_examples/02_lifted_transforms.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# %%\nimport jax\nimport jax.numpy as jnp\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport optax\n\nfrom flax import nnx\n\nX = np.linspace(0, 1, 100)[:, None]\nY = 0.8 * X**2 + 0.1 + np.random.normal(0, 0.1, size=X.shape)\n\n\ndef dataset(batch_size):\n while True:\n idx = np.random.choice(len(X), size=batch_size)\n yield X[idx], Y[idx]\n\n\nclass Linear(nnx.Module):\n def __init__(self, din: int, dout: int, *, rngs: nnx.Rngs):\n self.w = nnx.Param(jax.random.uniform(rngs.params(), (din, dout)))\n self.b = nnx.Param(jnp.zeros((dout,)))\n\n def __call__(self, x):\n return x @ self.w.value + self.b.value\n\n\nclass Count(nnx.Variable):\n pass\n\n\nclass MLP(nnx.Module):\n def __init__(self, din, dhidden, dout, *, rngs: nnx.Rngs):\n self.count = Count(jnp.array(0))\n self.linear1 = Linear(din, dhidden, rngs=rngs)\n self.linear2 = Linear(dhidden, dout, rngs=rngs)\n\n def __call__(self, x):\n self.count.value += 1\n x = self.linear1(x)\n x = jax.nn.relu(x)\n x = self.linear2(x)\n return x\n\n\nmodel = MLP(din=1, dhidden=32, dout=1, rngs=nnx.Rngs(0))\ntx = optax.sgd(1e-3)\noptimizer = nnx.Optimizer(model, tx, wrt=nnx.Param)\n\n\n@nnx.jit\ndef train_step(model: MLP, optimizer: nnx.Optimizer, batch):\n x, y = batch\n\n def loss_fn(model: MLP):\n y_pred = model(x)\n return jnp.mean((y - y_pred) ** 2)\n\n grads: nnx.State = nnx.grad(loss_fn)(model)\n optimizer.update(model, grads)\n\n\n@nnx.jit\ndef test_step(model: MLP, batch):\n x, y = batch\n y_pred = model(x)\n loss = jnp.mean((y - y_pred) ** 2)\n return {'loss': loss}\n\ncached_train_step = nnx.cached_partial(train_step, model, optimizer)\ncached_test_step = nnx.cached_partial(test_step, model)\n\ntotal_steps = 10_000\nfor step, batch in enumerate(dataset(32)):\n cached_train_step(batch)\n\n if step % 1000 == 0:\n logs = cached_test_step((X, Y))\n print(f\"step: {step}, loss: {logs['loss']}\")\n\n if step >= total_steps - 1:\n break\n\nprint('times called:', model.count.value)\n\ny_pred = model(X)\n\nplt.scatter(X, Y, color='blue')\nplt.plot(X, y_pred, color='black')\nplt.show()\n" }, { "path": "examples/nnx_toy_examples/03_train_state.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# %%\nimport jax\nimport jax.numpy as jnp\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport optax\n\nfrom flax import nnx\nfrom flax.training import train_state\n\nX = np.linspace(0, 1, 100)[:, None]\nY = 0.8 * X**2 + 0.1 + np.random.normal(0, 0.1, size=X.shape)\n\n\ndef dataset(batch_size):\n while True:\n idx = np.random.choice(len(X), size=batch_size)\n yield X[idx], Y[idx]\n\n\nclass Linear(nnx.Module):\n def __init__(self, din: int, dout: int, *, rngs: nnx.Rngs):\n self.w = nnx.Param(jax.random.uniform(rngs.params(), (din, dout)))\n self.b = nnx.Param(jnp.zeros((dout,)))\n\n def __call__(self, x):\n return x @ self.w.value + self.b.value\n\n\nclass Count(nnx.Variable[nnx.A]):\n pass\n\n\nclass MLP(nnx.Module):\n def __init__(self, din, dhidden, dout, *, rngs: nnx.Rngs):\n self.count = Count(jnp.array(0))\n self.linear1 = Linear(din, dhidden, rngs=rngs)\n self.linear2 = Linear(dhidden, dout, rngs=rngs)\n\n def __call__(self, x):\n self.count.value += 1\n x = self.linear1(x)\n x = jax.nn.relu(x)\n x = self.linear2(x)\n return x\n\nclass TrainState(train_state.TrainState):\n counts: nnx.State\n graphdef: nnx.GraphDef\n\nmodel = MLP(din=1, dhidden=32, dout=1, rngs=nnx.Rngs(0))\ngraphdef, params, counts = nnx.split(model, nnx.Param, Count)\n\nstate = TrainState.create(\n apply_fn=None,\n graphdef=graphdef,\n params=params,\n tx=optax.sgd(0.1),\n counts=counts,\n)\ndel params, counts\n\n\n@jax.jit\ndef train_step(state: TrainState, batch):\n x, y = batch\n\n def loss_fn(params):\n model = nnx.merge(state.graphdef, params, state.counts, copy=True)\n y_pred = model(x)\n loss = jnp.mean((y - y_pred) ** 2)\n counts = nnx.state(model, Count)\n return loss, counts\n\n grads, counts = jax.grad(loss_fn, has_aux=True)(state.params)\n # sdg update\n state = state.apply_gradients(grads=grads, counts=counts)\n\n return state\n\n\n@jax.jit\ndef test_step(state: nnx.TrainState[MLP], batch):\n x, y = batch\n model = nnx.merge(state.graphdef, state.params, state.counts)\n y_pred = model(x)\n loss = jnp.mean((y - y_pred) ** 2)\n return {'loss': loss}\n\n\ntotal_steps = 10_000\nfor step, batch in enumerate(dataset(32)):\n state = train_step(state, batch)\n\n if step % 1000 == 0:\n logs = test_step(state, (X, Y))\n print(f\"step: {step}, loss: {logs['loss']}\")\n\n if step >= total_steps - 1:\n break\n\nmodel = nnx.merge(state.graphdef, state.params, state.counts)\nprint('times called:', model.count.value)\n\ny_pred = model(X)\n\nplt.scatter(X, Y, color='blue')\nplt.plot(X, y_pred, color='black')\nplt.show()\n" }, { "path": "examples/nnx_toy_examples/04_data_parallel_with_jit.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nos.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=8'\n\nimport jax\nimport jax.numpy as jnp\nimport numpy as np\nimport optax\nfrom flax import nnx\nfrom jax.experimental import mesh_utils\nimport matplotlib.pyplot as plt\n\n# create a mesh + shardings\nnum_devices = jax.local_device_count()\nmesh = jax.sharding.Mesh(\n mesh_utils.create_device_mesh((num_devices,)), ('data',)\n)\nmodel_sharding = jax.NamedSharding(mesh, jax.sharding.PartitionSpec())\ndata_sharding = jax.NamedSharding(mesh, jax.sharding.PartitionSpec('data'))\n\n\n# create model\nclass MLP(nnx.Module):\n def __init__(self, din, dmid, dout, *, rngs: nnx.Rngs):\n self.linear1 = nnx.Linear(din, dmid, rngs=rngs)\n self.linear2 = nnx.Linear(dmid, dout, rngs=rngs)\n\n def __call__(self, x):\n return self.linear2(nnx.relu(self.linear1(x)))\n\n\nmodel = MLP(1, 64, 1, rngs=nnx.Rngs(0))\noptimizer = nnx.Optimizer(model, optax.adamw(1e-2), wrt=nnx.Param)\n\n# replicate state\nstate = nnx.state((model, optimizer))\nstate = jax.device_put(state, model_sharding)\nnnx.update((model, optimizer), state)\n\n# visualize model sharding\nprint('model sharding')\njax.debug.visualize_array_sharding(model.linear1.kernel.value)\n\n\n@nnx.jit\ndef train_step(model: MLP, optimizer: nnx.Optimizer, x, y):\n def loss_fn(model: MLP):\n y_pred = model(x)\n return jnp.mean((y - y_pred) ** 2)\n\n loss, grads = nnx.value_and_grad(loss_fn)(model)\n optimizer.update(model, grads)\n return loss\n\n\ndef dataset(steps, batch_size):\n for _ in range(steps):\n x = np.random.uniform(-2, 2, size=(batch_size, 1))\n y = 0.8 * x**2 + 0.1 + np.random.normal(0, 0.1, size=x.shape)\n yield x, y\n\n\nfor step, (x, y) in enumerate(dataset(1000, 16)):\n # shard data\n x, y = jax.device_put((x, y), data_sharding)\n # train\n loss = train_step(model, optimizer, x, y)\n\n if step == 0:\n print('data sharding')\n jax.debug.visualize_array_sharding(x)\n\n if step % 100 == 0:\n print(f'step={step}, loss={loss}')\n\n# dereplicate state\nstate = nnx.state((model, optimizer))\nstate = jax.device_get(state)\nnnx.update((model, optimizer), state)\n\nX, Y = next(dataset(1, 1000))\nx_range = np.linspace(X.min(), X.max(), 100)[:, None]\ny_pred = model(x_range)\n\n# plot\nplt.scatter(X, Y, label='data')\nplt.plot(x_range, y_pred, color='black', label='model')\nplt.legend()\nplt.show()" }, { "path": "examples/nnx_toy_examples/05_vae.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# %%\nimport typing as tp\n\nimport jax\nimport jax.numpy as jnp\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport optax\nfrom datasets import load_dataset\n\nfrom flax import nnx\n\nnp.random.seed(42)\nlatent_size = 32\nimage_shape: tp.Sequence[int] = (28, 28)\nsteps_per_epoch: int = 200\nbatch_size: int = 64\nepochs: int = 20\n\n\ndataset = load_dataset('mnist')\nX_train = np.array(np.stack(dataset['train']['image']), dtype=np.uint8)\nX_test = np.array(np.stack(dataset['test']['image']), dtype=np.uint8)\n# Now binarize data\nX_train = (X_train > 0).astype(jnp.float32)\nX_test = (X_test > 0).astype(jnp.float32)\n\nprint('X_train:', X_train.shape, X_train.dtype)\nprint('X_test:', X_test.shape, X_test.dtype)\n\n\nclass Loss(nnx.Variable):\n pass\n\n\n# %%\nclass Encoder(nnx.Module):\n def __init__(self, din: int, dmid: int, dout: int, *, rngs: nnx.Rngs):\n self.linear1 = nnx.Linear(din, dmid, rngs=rngs)\n self.linear_mean = nnx.Linear(dmid, dout, rngs=rngs)\n self.linear_std = nnx.Linear(dmid, dout, rngs=rngs)\n self.rngs = rngs\n\n def __call__(self, x: jax.Array) -> jax.Array:\n x = x.reshape((x.shape[0], -1)) # flatten\n x = self.linear1(x)\n x = jax.nn.relu(x)\n\n mean = self.linear_mean(x)\n std = jnp.exp(self.linear_std(x))\n\n self.kl_loss = Loss(\n jnp.mean(\n 0.5 * jnp.mean(-jnp.log(std**2) - 1.0 + std**2 + mean**2, axis=-1)\n )\n )\n key = self.rngs.noise()\n z = mean + std * jax.random.normal(key, mean.shape)\n return z\n\n\nclass Decoder(nnx.Module):\n def __init__(self, din: int, dmid: int, dout: int, *, rngs: nnx.Rngs):\n self.linear1 = nnx.Linear(din, dmid, rngs=rngs)\n self.linear2 = nnx.Linear(dmid, dout, rngs=rngs)\n\n def __call__(self, z: jax.Array) -> jax.Array:\n z = self.linear1(z)\n z = jax.nn.relu(z)\n logits = self.linear2(z)\n return logits\n\n\nclass VAE(nnx.Module):\n def __init__(\n self,\n din: int,\n hidden_size: int,\n latent_size: int,\n output_shape: tp.Sequence[int],\n *,\n rngs: nnx.Rngs,\n ):\n self.output_shape = output_shape\n self.encoder = Encoder(din, hidden_size, latent_size, rngs=rngs)\n self.decoder = Decoder(\n latent_size, hidden_size, int(np.prod(output_shape)), rngs=rngs\n )\n\n def __call__(self, x: jax.Array) -> jax.Array:\n z = self.encoder(x)\n logits = self.decoder(z)\n logits = jnp.reshape(logits, (-1, *self.output_shape))\n return logits\n\n def generate(self, z):\n logits = self.decoder(z)\n logits = jnp.reshape(logits, (-1, *self.output_shape))\n return nnx.sigmoid(logits)\n\n\nmodel = VAE(\n din=int(np.prod(image_shape)),\n hidden_size=256,\n latent_size=latent_size,\n output_shape=image_shape,\n rngs=nnx.Rngs(0, noise=1),\n)\n\noptimizer = nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param)\n\n\n# %%\n@nnx.jit\ndef train_step(model: VAE, optimizer: nnx.Optimizer, x: jax.Array):\n def loss_fn(model: VAE):\n logits = model(x)\n losses = nnx.pop(model, Loss)\n kl_loss = sum(jax.tree_util.tree_leaves(losses), 0.0)\n reconstruction_loss = jnp.mean(\n optax.sigmoid_binary_cross_entropy(logits, x)\n )\n loss = reconstruction_loss + 0.1 * kl_loss\n return loss\n\n loss, grads = nnx.value_and_grad(loss_fn)(model)\n optimizer.update(model, grads)\n\n return loss\n\n\n@nnx.jit\ndef forward(model: VAE, x: jax.Array) -> jax.Array:\n y_pred = model(x)\n return jax.nn.sigmoid(y_pred)\n\n\n@nnx.jit\ndef sample(model: VAE, z: jax.Array) -> jax.Array:\n return model.generate(z)\n\n\n# %%\n\nfor epoch in range(epochs):\n losses = []\n for step in range(steps_per_epoch):\n idxs = np.random.randint(0, len(X_train), size=(batch_size,))\n x_batch = X_train[idxs]\n\n loss = train_step(model, optimizer, x_batch)\n losses.append(np.asarray(loss))\n\n print(f'Epoch {epoch} loss: {np.mean(losses)}')\n\n# exit()\n# %%\n# get random samples\nidxs = np.random.randint(0, len(X_test), size=(5,))\nx_sample = X_test[idxs]\n\n# get predictions\ny_pred = forward(model, x_sample)\n\n# plot reconstruction\nfigure = plt.figure(figsize=(3 * 5, 3 * 2))\nplt.title('Reconstruction Samples')\nfor i in range(5):\n plt.subplot(2, 5, i + 1)\n plt.imshow(x_sample[i], cmap='gray')\n plt.subplot(2, 5, 5 + i + 1)\n plt.imshow(y_pred[i], cmap='gray')\n # # tbwriter.add_figure(\"VAE Example\", figure, epochs)\n\nplt.show()\n\n# %%\n# plot generative samples\nz_samples = np.random.normal(scale=1.5, size=(12, latent_size))\nsamples = sample(model, z_samples)\n\nfigure = plt.figure(figsize=(3 * 5, 3 * 2))\nplt.title('Generative Samples')\nfor i in range(5):\n plt.subplot(2, 5, 2 * i + 1)\n plt.imshow(samples[i], cmap='gray')\n plt.subplot(2, 5, 2 * i + 2)\n plt.imshow(samples[i + 1], cmap='gray')\n\nplt.show()\n\n# %%\n" }, { "path": "examples/nnx_toy_examples/06_scan_over_layers.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nimport jax\nimport jax.numpy as jnp\n\nfrom flax import nnx\n\n\nclass Block(nnx.Module):\n def __init__(self, dim: int, *, rngs: nnx.Rngs):\n self.linear = nnx.Linear(dim, dim, rngs=rngs)\n self.bn = nnx.BatchNorm(dim, rngs=rngs)\n self.dropout = nnx.Dropout(0.5, rngs=rngs)\n\n def __call__(self, x: jax.Array):\n return jax.nn.gelu(self.dropout(self.bn(self.linear(x))))\n\n\nclass ScanMLP(nnx.Module):\n \"\"\"\n An MLP that uses `vmap` during `__init__` to create a Block instance\n with an additional `layer` axis, and `scan` during `__call__` to apply\n the sequence of layers iteratively over the input / output `x`.\n \"\"\"\n\n def __init__(self, dim: int, *, n_layers: int, rngs: nnx.Rngs):\n self.n_layers = n_layers\n\n @nnx.split_rngs(splits=n_layers)\n @nnx.vmap(axis_size=n_layers)\n def create_block(rngs: nnx.Rngs):\n return Block(dim, rngs=rngs)\n\n self.layers = create_block(rngs)\n\n def __call__(self, x: jax.Array) -> jax.Array:\n @nnx.scan\n def scan_fn(x: jax.Array, block: Block):\n x = block(x)\n return x, None\n\n x, _ = scan_fn(x, self.layers)\n\n return x\n\n\nmodel = ScanMLP(10, n_layers=5, rngs=nnx.Rngs(0))\n\nx = jnp.ones((3, 10))\ny = model(x)\n\nprint(jax.tree.map(jnp.shape, nnx.state(model)))\nprint(y.shape)\n" }, { "path": "examples/nnx_toy_examples/07_array_leaves.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# %%\nimport jax\nimport jax.numpy as jnp\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport optax\n\nfrom flax import nnx, struct\n\nX = np.linspace(0, 1, 100)[:, None]\nY = 0.8 * X**2 + 0.1 + np.random.normal(0, 0.1, size=X.shape)\n\n\ndef dataset(batch_size):\n while True:\n idx = np.random.choice(len(X), size=batch_size)\n yield X[idx], Y[idx]\n\nclass Linear(nnx.Module):\n def __init__(self, din: int, dout: int, *, rngs: nnx.Rngs):\n self.w = jax.random.normal(rngs.params(), (din, dout))\n self.b = jnp.zeros((dout,))\n\n def __call__(self, x):\n return x @ self.w + self.b\n\n\nclass MLP(nnx.Module):\n def __init__(self, din, dhidden, dout, *, rngs: nnx.Rngs):\n self.count = jnp.array(0)\n self.linear1 = Linear(din, dhidden, rngs=rngs)\n self.linear2 = Linear(dhidden, dout, rngs=rngs)\n\n def __call__(self, x):\n self.count += 1\n return self.linear2(nnx.relu(self.linear1(x)))\n\ndef is_param(path, value):\n key = path[-1]\n return key == 'w' or key == 'b'\n\nmodel = MLP(din=1, dhidden=32, dout=1, rngs=nnx.Rngs(0))\ntx = optax.sgd(1e-3)\noptimizer = nnx.Optimizer(model, tx, wrt=is_param)\n\n\n@nnx.jit\ndef train_step(model: MLP, optimizer: nnx.Optimizer, batch):\n x, y = batch\n\n def loss_fn(model: MLP):\n y_pred = model(x)\n return jnp.mean((y - y_pred) ** 2)\n\n diff_state = nnx.DiffState(0, is_param)\n grads: nnx.State = nnx.grad(loss_fn, argnums=diff_state)(model)\n optimizer.update(model, grads)\n\n\n@nnx.jit\ndef test_step(model: MLP, batch):\n x, y = batch\n y_pred = model(x)\n loss = jnp.mean((y - y_pred) ** 2)\n return {'loss': loss}\n\n\ntotal_steps = 10_000\nfor step, batch in enumerate(dataset(32)):\n train_step(model, optimizer, batch)\n\n if step % 1000 == 0:\n logs = test_step(model, (X, Y))\n print(f\"step: {step}, loss: {logs['loss']}\")\n\n if step >= total_steps - 1:\n break\n\nprint('times called:', model.count)\n\ny_pred = model(X)\n\nplt.scatter(X, Y, color='blue')\nplt.plot(X, y_pred, color='black')\nplt.show()\n" }, { "path": "examples/nnx_toy_examples/08_save_load_checkpoints.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom tempfile import TemporaryDirectory\n\nimport jax\nimport jax.numpy as jnp\nimport orbax.checkpoint as orbax\n\nfrom flax import nnx\n\n\nclass MLP(nnx.Module):\n def __init__(self, din: int, dmid: int, dout: int, *, rngs: nnx.Rngs):\n self.dense1 = nnx.Linear(din, dmid, rngs=rngs)\n self.dense2 = nnx.Linear(dmid, dout, rngs=rngs)\n\n def __call__(self, x: jax.Array) -> jax.Array:\n x = self.dense1(x)\n x = jax.nn.relu(x)\n x = self.dense2(x)\n return x\n\n\ndef create_model(seed: int):\n return MLP(10, 20, 30, rngs=nnx.Rngs(seed))\n\n\ndef create_and_save(seed: int, path: str):\n model = create_model(seed)\n state = nnx.state(model)\n # Save the parameters\n checkpointer = orbax.PyTreeCheckpointer()\n checkpointer.save(f'{path}/state', state)\n\n\ndef load_model(path: str) -> MLP:\n # create that model with abstract shapes\n model = nnx.eval_shape(lambda: create_model(0))\n state = nnx.state(model)\n # Load the parameters\n checkpointer = orbax.PyTreeCheckpointer()\n state = checkpointer.restore(f'{path}/state', item=state)\n # update the model with the loaded state\n nnx.update(model, state)\n return model\n\n\nwith TemporaryDirectory() as tmpdir:\n # create a checkpoint\n create_and_save(42, tmpdir)\n # load model from checkpoint\n model = load_model(tmpdir)\n # run the model\n y = model(jnp.ones((1, 10)))\n print(model)\n print(y)\n" }, { "path": "examples/nnx_toy_examples/09_parameter_surgery.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nimport jax\n\nfrom flax import nnx\n\n\n# lets pretend this function loads a pretrained model from a checkpoint\ndef load_pretrained():\n return nnx.Linear(784, 128, rngs=nnx.Rngs(0))\n\n\n# create a simple linear classifier using a pretrained backbone\nclass Classifier(nnx.Module):\n def __init__(self, *, rngs: nnx.Rngs):\n self.backbone = nnx.Linear(784, 128, rngs=nnx.Rngs(0))\n self.head = nnx.Linear(128, 10, rngs=rngs)\n\n def __call__(self, x):\n x = self.backbone(x)\n x = nnx.relu(x)\n x = self.head(x)\n return x\n\n\n# create the classifier using the pretrained backbone, here we are technically\n# doing \"parameter surgery\", however, compared to Haiku/Flax where you must manually\n# construct the parameter structure, in NNX this is done automatically\nmodel = Classifier(rngs=nnx.Rngs(42))\nmodel.backbone = load_pretrained()\n\n\n# create a filter to select all the parameters that are not part of the\n# backbone, i.e. the classifier parameters\nis_trainable = lambda path, node: (\n 'backbone' in path and isinstance(node, nnx.Param)\n)\n\n# split the parameters into trainable and non-trainable parameters\ngraphdef, trainable_params, non_trainable = nnx.split(model, is_trainable, ...)\n\nprint(\n 'trainable_params =',\n jax.tree.map(jax.numpy.shape, trainable_params),\n)\nprint('non_trainable = ', jax.tree.map(jax.numpy.shape, non_trainable))\n" }, { "path": "examples/nnx_toy_examples/10_fsdp_and_optimizer.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport dataclasses\nimport os\nos.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=8'\n\nfrom matplotlib import pyplot as plt\nfrom jax.experimental import mesh_utils\nfrom jax.sharding import Mesh, PartitionSpec as P, NamedSharding\nimport jax\nimport jax.numpy as jnp\nimport numpy as np\nfrom flax import nnx\nimport typing as tp\n\nmesh = jax.sharding.Mesh(\n mesh_utils.create_device_mesh((2, 4)),\n ('data', 'model'),\n)\n\n\ndef named_sharding(*names: str | None) -> NamedSharding:\n return NamedSharding(mesh, P(*names))\n\n\n@dataclasses.dataclass(unsafe_hash=True)\nclass MeshRules:\n embed: str | None = None\n mlp: str | None = None\n data: str | None = None\n\n def __call__(self, *keys: str) -> tuple[str, ...]:\n return tuple(getattr(self, key) for key in keys)\n\n\nmesh_rules = MeshRules(\n embed=None,\n mlp='model',\n data='data',\n)\n\n\nclass MLP(nnx.Module):\n def __init__(self, din, dmid, dout, rngs: nnx.Rngs):\n self.w1 = nnx.Param(\n nnx.initializers.lecun_normal()(rngs.params(), (din, dmid)),\n out_sharding=mesh_rules('embed', 'mlp'),\n )\n self.b1 = nnx.Param(\n jnp.zeros((dmid,)),\n out_sharding=mesh_rules('mlp'),\n )\n self.w2 = nnx.Param(\n nnx.initializers.lecun_normal()(rngs.params(), (dmid, dout)),\n out_sharding=mesh_rules('embed', 'mlp'),\n )\n\n def __call__(self, x: jax.Array):\n return nnx.relu(x @ self.w1 + self.b1) @ self.w2\n\n\nclass SGDState(nnx.Variable):\n pass\n\n\nclass SGD(nnx.Pytree):\n def __init__(self, params: nnx.State, lr, decay=0.9):\n def init_optimizer_state(variable: nnx.Variable):\n return SGDState(\n jnp.zeros_like(variable.value), **variable.get_metadata()\n )\n\n self.lr = lr\n self.params = nnx.data(params)\n self.momentum: nnx.State = nnx.data(jax.tree.map(\n init_optimizer_state,\n self.params,\n is_leaf=lambda x: isinstance(x, nnx.Variable),\n ))\n self.decay = decay\n\n def update(self, grads: nnx.State):\n def update_fn(\n params: nnx.Variable, momentum: SGDState, grad: nnx.Variable\n ):\n # v_t = β * v_{t-1} + (1 - β) * ∇J(θ_t)\n momentum[...] = self.decay * momentum[...] + (1 - self.decay) * grad[...]\n # θ_{t+1} = θ_t - α * v_t\n params[...] -= self.lr * momentum[...]\n\n jax.tree.map(\n update_fn,\n self.params,\n self.momentum,\n grads,\n is_leaf=lambda x: isinstance(x, nnx.Variable),\n )\n\n\n@nnx.jit\ndef create_model():\n model = MLP(1, 32, 1, rngs=nnx.Rngs(0))\n optimizer = SGD(nnx.variables(model, nnx.Param), 0.01, decay=0.9)\n return model, optimizer\n\nwith jax.set_mesh(mesh):\n model, optimizer = create_model()\n\nprint('Model parameters sharding:')\njax.debug.visualize_array_sharding(model.w1.value)\nprint('Optimizer momentum sharding:')\njax.debug.visualize_array_sharding(optimizer.momentum['w1'].value)\n\n\n@nnx.jit\ndef train_step(model: MLP, optimizer: SGD, x, y):\n def loss_fn(model):\n y_pred = model(x)\n loss = jnp.mean((y - y_pred) ** 2)\n return loss\n\n loss, grad = nnx.value_and_grad(loss_fn)(model)\n optimizer.update(grad)\n return loss\n\n\nX = np.linspace(-2, 2, 100)[:, None]\nY = 0.8 * X**2 + 0.1 + np.random.normal(0, 0.1, size=X.shape)\n\n\ndef dataset(batch_size, num_steps):\n for _ in range(num_steps):\n idx = np.random.choice(len(X), size=batch_size)\n yield X[idx], Y[idx]\n\n\nlosses = []\nfor step, (x_batch, y_batch) in enumerate(\n dataset(batch_size=32, num_steps=10_000)\n):\n x_batch, y_batch = jax.device_put((x_batch, y_batch), named_sharding('data'))\n loss = train_step(model, optimizer, x_batch, y_batch)\n losses.append(float(loss))\n if step % 1000 == 0:\n print(f'Step {step}: Loss = {loss}')\n\nplt.figure()\nplt.plot(losses[20:])\n\ny_pred = model(X)\nplt.figure()\nplt.scatter(X, Y, color='blue')\nplt.plot(X, y_pred, color='black')\nplt.show()\n" }, { "path": "examples/nnx_toy_examples/hijax_basic.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# %%\nimport jax\nimport jax.numpy as jnp\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport optax\n\nfrom flax import nnx\n\n\nX = np.linspace(-jnp.pi, jnp.pi, 100)[:, None]\nY = 0.8 * jnp.sin(X) + 0.1 + np.random.normal(0, 0.1, size=X.shape)\n\ndef dataset(batch_size):\n while True:\n idx = np.random.choice(len(X), size=batch_size)\n yield X[idx], Y[idx]\n\n\nclass Linear(nnx.Module):\n def __init__(self, din: int, dout: int, *, rngs: nnx.Rngs):\n self.w = nnx.Param(rngs.params.uniform((din, dout)))\n self.b = nnx.Param(jnp.zeros((dout,)))\n\n def __call__(self, x):\n return x @ self.w + self.b[None]\n\n\nclass Count(nnx.Variable[nnx.A]):\n pass\n\n\nclass MLP(nnx.Module):\n def __init__(self, din, dhidden, dout, *, rngs: nnx.Rngs):\n self.count = Count(jnp.array(0))\n self.linear1 = Linear(din, dhidden, rngs=rngs)\n self.linear2 = Linear(dhidden, dout, rngs=rngs)\n\n def __call__(self, x):\n self.count[...] += 1\n return self.linear2(jax.nn.relu(self.linear1(x)) * 0.5)\n\nnnx.var_defaults(hijax=True)\n\nmodel = MLP(din=1, dhidden=32, dout=1, rngs=nnx.Rngs(0))\noptimizer = nnx.Optimizer(model, optax.sgd(learning_rate=0.1), wrt=nnx.Param)\n\n\n@jax.jit\ndef train_step(model, optimizer, x, y):\n graphdef, params, nondiff = nnx.split(model, nnx.Param, ...)\n\n def loss_fn(params):\n model = nnx.merge(graphdef, params, nondiff)\n return jnp.mean((y - model(x)) ** 2)\n\n grads = jax.grad(loss_fn)(nnx.vars_as(params, is_mutable=False))\n optimizer.update(model, grads)\n\n@jax.jit\ndef test_step(model: MLP, x, y):\n return {'loss': jnp.mean((y - model(x)) ** 2)}\n\n\ntotal_steps = 10_000\nfor step, (x, y) in enumerate(dataset(32)):\n train_step(model, optimizer, x, y)\n\n if step % 1000 == 0:\n logs = test_step(model, X, Y)\n print(f\"step: {step}, loss: {logs['loss']}\")\n\n if step >= total_steps - 1:\n break\n\nprint('times called:', model.count[...])\n\ny_pred = model(X)\n\nplt.scatter(X, Y, color='blue')\nplt.plot(X, y_pred, color='black')\nplt.show()\n" }, { "path": "examples/nnx_toy_examples/hijax_demo.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# %%\nimport jax\nimport jax.numpy as jnp\nimport matplotlib.pyplot as plt\nimport numpy as np\n\nfrom flax import nnx\n\n# ## Data\n# We create a simple dataset of points sampled from a parabola with some noise.\nX = np.linspace(-jnp.pi, jnp.pi, 100)[:, None]\nY = 0.8 * jnp.sin(X) + 0.1 + np.random.normal(0, 0.1, size=X.shape)\n\n\ndef dataset(batch_size):\n while True:\n idx = np.random.choice(len(X), size=batch_size)\n yield X[idx], Y[idx]\n\n\n# ## Model\n# Here we define a MLP made of a stack of blocks. Each block contains a linear layer,\n# batch normalization, and a dropout layer.\n#\n# In this version we want the Modules to be pytrees so they can be used with JAX transforms\n# so we use a new Pytree type as the base. The main difference with current NNX is that\n# attributes that contain arrays or other pytrees now need to be explicitly marked as\n# using `nnx.data` to be included in the pytree.\nclass Linear(nnx.Module):\n def __init__(self, din: int, dout: int, *, rngs: nnx.Rngs):\n self.din, self.dout = din, dout\n initializer = jax.nn.initializers.lecun_normal()\n # nnx.data is used mark attributes as pytree data\n # Param, BatchState, and Cache are built-in Variable subtypes\n self.w = nnx.Param(initializer(rngs.params(), (din, dout)))\n self.b = nnx.Param(jnp.zeros((dout,)))\n\n def __call__(self, x: jax.Array):\n return x @ self.w + self.b[None]\n\n\n# Block implements linear, batch norm, and dropout. Its behavior\n# is controlled by the 'use_stats' and 'deterministic' flags.\nclass Block(nnx.Module):\n def __init__(\n self,\n din: int,\n dout: int,\n *,\n dropout_rate: float = 0.05,\n moumentum: float = 0.95,\n use_stats: bool = False,\n deterministic: bool = False,\n rngs: nnx.Rngs,\n ):\n # ----------- linear -------------------\n self.din, self.dout = din, dout\n initializer = jax.nn.initializers.lecun_normal()\n self.w = nnx.Param(initializer(rngs.params(), (din, dout)))\n self.b = nnx.Param(jnp.zeros((dout,)))\n # ----------- batch norm ---------------\n self.mu = moumentum # momentum\n self.use_stats = use_stats\n self.mean = nnx.BatchStat(jnp.zeros((dout,)))\n self.var = nnx.BatchStat(jnp.ones((dout,)))\n self.scale = nnx.Param(jnp.ones((dout,)))\n self.bias = nnx.Param(jnp.zeros((dout,)))\n # ----------- dropout ------------------\n self.dropout_rate = dropout_rate\n self.deterministic = deterministic\n\n def __call__(\n self, x: jax.Array, *, rngs: nnx.Rngs | None = None\n ) -> jax.Array:\n # ----------- linear --------------------\n x = x @ self.w + self.b[None]\n # ----------- batch norm ----------------\n if self.use_stats:\n mean = self.mean\n var = self.var\n else:\n mean = jnp.mean(x, axis=0)\n var = jnp.var(x, axis=0)\n # ema updates\n # stop gradient is used until a Hijax supports updates from grad tracers\n sg = jax.lax.stop_gradient\n self.mean[...] = sg(self.mu * self.mean + (1 - self.mu) * mean)\n self.var[...] = sg(self.mu * self.var + (1 - self.mu) * var)\n x = (x - mean[None]) / jnp.sqrt(var[None] + 1e-5)\n x = x * self.scale + self.bias\n # ----------- dropout -------------------\n if not self.deterministic and self.dropout_rate > 0.0:\n assert rngs is not None\n keep_prob = 1.0 - self.dropout_rate\n mask = jax.random.bernoulli(rngs.dropout(), keep_prob, x.shape)\n x = jnp.where(mask, x / keep_prob, jnp.zeros_like(x))\n # ----------- activation ---------------\n x = jax.nn.gelu(x)\n return x\n\n\nclass Model(nnx.Module):\n def __init__(\n self,\n num_blocks: int,\n din: int,\n dhidden: int,\n dout: int,\n *,\n use_scan: bool = True,\n rngs: nnx.Rngs,\n ):\n self.count = nnx.Variable(jnp.array(0))\n self.block_in = Block(din, dhidden, rngs=rngs)\n self.linear_out = Linear(dhidden, dout, rngs=rngs)\n\n # 'blocks' is either a list of blocks or single block\n # whose parameters contain an additional 'layer' dimension,\n # here created using jax.vmap\n if use_scan:\n\n @jax.vmap\n def create_block(rngs, /):\n # return nnx.stateless(Block(dhidden, dhidden, rngs=rngs))\n return Block(dhidden, dhidden, rngs=rngs)\n\n # self.blocks = nnx.stateful(create_block(rngs.fork(split=num_blocks)))\n self.blocks = create_block(rngs.fork(split=num_blocks))\n else:\n self.blocks = nnx.List(\n [Block(dhidden, dhidden, rngs=rngs) for i in range(num_blocks)]\n )\n\n def __call__(self, x: jax.Array, *, rngs: nnx.Rngs | None = None):\n self.count[...] += 1\n x = self.block_in(x, rngs=rngs)\n\n # on the forward pass we either iterate over the block\n # list or use jax.lax.scan to apply the blocks, if we\n # had shared state we would use split and merge to\n # pass the shared state as a capture\n if isinstance(self.blocks, nnx.List):\n for block in self.blocks:\n x = block(x, rngs=rngs)\n else:\n\n def block_fw(x, block: Block):\n x = block(x, rngs=rngs)\n return x, None\n\n x, _ = jax.lax.scan(block_fw, x, self.blocks)\n x = self.linear_out(x)\n return x\n\n\n# ## Optimizer\nclass OptState(nnx.Variable): ...\n\n\n# Optimizer are an interesting case as they are inherently stateful and\n# pose a good use case for MutableHijax. Here we implement SGD with\n# momentum. The optimizer receives the params as constructor arguments but doesn't\n# hold a reference to them, it only uses the params to initialize its state\n# by creating new OptState Variables that reuse the param's metadata.\nclass SGD(nnx.Pytree):\n def __init__(self, params, lr: float, decay: float = 0.9):\n self.lr = lr\n self.decay = decay\n\n def make_opt_state(x):\n if isinstance(x, nnx.Variable):\n return OptState(jnp.zeros_like(x[...]), **x.get_metadata())\n else:\n return OptState(jnp.zeros_like(x))\n\n self.momentum = nnx.data(jax.tree.map(make_opt_state, params))\n\n # during the update we simply map over (params, momentum, grads),\n # for each triplet we implement the SGD update rule which updates\n # both the optimizer's state (momentum) and the params in place.\n def update(self, params, grads):\n def update_fn(\n param: nnx.Variable[jax.Array],\n momentum: nnx.Variable[jax.Array],\n grad: nnx.Variable[jax.Array],\n ):\n momentum[...] = self.decay * momentum + (1 - self.decay) * grad\n param[...] -= self.lr * momentum\n\n # is_leaf might not be necesarry as MutableHijaxVariable are not pytreees\n jax.tree.map(update_fn, params, self.momentum, grads)\n\n\n# ## Training\nnnx.var_defaults(hijax=True)\n\nrngs = nnx.Rngs(params=0, dropout=1)\nmodel = Model(\n num_blocks=3, din=1, dhidden=256, dout=1, use_scan=False, rngs=rngs\n)\noptimizer = SGD(params=nnx.state(model, nnx.Param), lr=3e-3, decay=0.99)\n\n# Create a copy of the model structure and set its attributes to eval model.\n# This works because they share the underlying ArrayRefs so both models\n# will always be in sync.\neval_model = nnx.merge(*nnx.split(model))\neval_model.set_attributes(use_stats=True, deterministic=True)\n\n\n# The training step uses 'jax.jit' and receives the model and optimizer as arguments,\n# this is supported as they are now pytrees. The first thing we do is group the model\n# state into the params and the non-differentiable state using 'split'. We differentiate\n# the loss function using 'jax.grad' with respect to the params-only. Inside the loss\n# function we merge the params and non-diff state back into a single model and then\n# compute the loss by calling the model with the inputs.\n@jax.jit\ndef train_step(model: Model, optimizer: SGD, rngs: nnx.Rngs, x, y):\n graphdef, params, nondiff = nnx.split(model, nnx.Param, ...)\n\n def loss_fn(params):\n model = nnx.merge(graphdef, params, nondiff)\n loss = jnp.mean((model(x, rngs=rngs) - y) ** 2)\n return loss\n\n # For the time being we have to use 'immutable'\n # as 'jax.grad' doesn't support QDD types yet.\n grads = jax.grad(loss_fn)(nnx.vars_as(params, is_mutable=False))\n # 'update' mutates the optimizer's state and the params in place\n # so we don't need to return anything 🚀\n optimizer.update(params, grads)\n\n\n# simple test step that computes the loss\n@jax.jit\ndef test_step(model: Model, x, y):\n return {'loss': jnp.mean((model(x) - y) ** 2)}\n\n\n# minimalistic training loop\ntotal_steps = 2_000\nfor step, (x, y) in enumerate(dataset(32)):\n train_step(model, optimizer, rngs, x, y)\n\n if step % 200 == 0:\n logs = test_step(eval_model, X, Y)\n print(f'step: {step}, loss: {logs[\"loss\"]}')\n\n if step >= total_steps - 1:\n break\n\n# ## Sample\n# Sampling is trivial, just use 'model_eval'\nprint('times called:', eval_model.count[...])\n\ny_pred = eval_model(X)\n\nplt.scatter(X, Y, color='blue')\nplt.plot(X, y_pred, color='black')\nplt.show()\n" }, { "path": "examples/nnx_toy_examples/requirements.txt", "content": "matplotlib>=3.7.1\ndatasets>=2.12.0" }, { "path": "examples/ogbg_molpcba/README.md", "content": "## Predicting Biological Activities of Molecules with Graph Neural Networks\n[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google/flax/blob/main/examples/ogbg_molpcba/ogbg_molpcba.ipynb)\n\nThis example trains a Graph Neural Network to classify molecules\non the basis of their biological activities.\n\n![Prediction on a caramboxin molecule](https://www.gstatic.com/flax_examples/ogbg_molpcba.png \"Prediction on a caramboxin molecule\")\n\nWe use [Jraph](https://github.com/deepmind/jraph/),\na JAX library for Graph Neural Networks, to define models\nwhich are trained on the\n[ogbg-molpcba](https://ogb.stanford.edu/docs/graphprop/)\ndataset, part of the [Open Graph Benchmark](https://ogb.stanford.edu/).\n\nYou can run this code and even modify it directly in Google Colab,\nno installation required!\nThe [Colab notebook](https://colab.research.google.com/github/google/flax/blob/main/examples/ogbg_molpcba/ogbg_molpcba.ipynb)\ncan even create visualizations of model predictions:\n\n![Visualizing predictions of a trained model](https://www.gstatic.com/flax_examples/ogbg_molpcba_predictions.svg? \"Visualizing predictions of a trained model\")\n\n\n### Requirements\n\nWe depend on\n[TensorFlow Datasets](https://www.tensorflow.org/datasets/catalog/ogbg_molpcba)\nfor ogbg-molpcba.\n\n### How to Run\n\nTo run with the default configuration:\n\n```shell\npython main.py --workdir=./ogbg_molpcba --config=configs/default.py\n```\n\nSince the configuration is defined using\n[config_flags](https://github.com/google/ml_collections/tree/master#config-flags),\nyou can override hyperparameters. For example, to change the number of epochs\nand the batch size:\n\n```shell\npython main.py --workdir=./ogbg_molpcba --config=configs/default.py \\\n--config.num_training_epochs=10 --config.batch_size=50\n```\n\nFor more extensive changes, you can directly edit the default\nconfiguration file or even add your own.\n\n### Supported Setups\n\nThis example supports only single device training.\nThe model should run with other configurations and hardware, but was explicitly\ntested on the following.\n\nHardware | Batch size | Training time | Test mean AP | Validation mean AP | Metrics\n-------- | ---------- | ------------- | ------- | ------- | ---------------\n1x V100 | 256 | 3h20m | 0.244 | 0.252 |[2021-08-03](https://tensorboard.dev/experiment/AAJqfvgSRJaA1MBkc0jMWQ/)\n\nThese metrics reported above are obtained at the end of training.\nWe observed that slightly higher metrics can be obtained with\nearly-stopping based on the validation mean AP:\n\nHardware | Batch size | Training time | Test mean AP | Validation mean AP | Metrics\n-------- | ---------- | ------------- | ------- | ------- | ---------------\n1x V100 | 256 | 2h55m | 0.249 | 0.257 |[2021-08-03](https://tensorboard.dev/experiment/AAJqfvgSRJaA1MBkc0jMWQ/)\n\n\n### Model Description\n\nThe default configuration corresponds to a\n[Graph Convolutional Network](https://arxiv.org/abs/1609.02907)\nmodel with 695,936 parameters.\n\nWe noticed diminishing gains when training for longer.\nFurther, the addition of self-loops and undirected edges significantly\nhelped performance.\nMinor improvements were seen with skip-connections across message-passing\nsteps, together with [LayerNorm](https://arxiv.org/abs/1607.06450).\nOn the contrary, we found that the addition of\n[virtual nodes](https://arxiv.org/abs/1709.03741),\nwhich are connected to all nodes in each graph,\ndid not improve performance.\n\n### References\n\n- Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong, Hongyu Ren,\n Bowen Liu, Michele Catasta and Jure Leskovec (2020).\n *Open Graph Benchmark: Datasets for Machine Learning on Graphs.*\n In Advances in Neural Information Processing Systems 33: Annual\n Conference on Neural Information Processing Systems 2020,\n NeurIPS 2020, December 6-12,\n 2020, virtual.\n\n- Thomas N. Kipf and Max Welling (2016). *Semi-supervised classification\n with graph convolutional networks.* arXiv preprint arXiv:1609.02907.\n\n- Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton (2016). *Layer\n normalization.* arXiv preprint arXiv:1607.06450.\n\n- Junying Li, Deng Cai and Xiaofei He (2017). *Learning graph-level\n representation for drug discovery.* arXiv preprint arXiv:1709.03741.\n\nThe caramboxin molecule diagram depicted above was obtained and modified from\n[Wikimedia Commons](https://commons.wikimedia.org/wiki/File:Caramboxin.svg),\navailable in the public domain.\n" }, { "path": "examples/ogbg_molpcba/configs/default.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Defines the default hyperparameters and training configuration.\n\nUses a Graph Convolutional Network model (https://arxiv.org/abs/1609.02907).\n\"\"\"\n\nimport ml_collections\n\n\ndef get_config():\n \"\"\"Get the default hyperparameter configuration.\"\"\"\n config = ml_collections.ConfigDict()\n\n # Optimizer.\n config.optimizer = 'adam'\n config.learning_rate = 1e-3\n\n # Training hyperparameters.\n config.batch_size = 256\n config.num_train_steps = 100_000\n config.log_every_steps = 100\n config.eval_every_steps = 1_000\n config.checkpoint_every_steps = 10_000\n config.add_virtual_node = False\n config.add_undirected_edges = True\n config.add_self_loops = True\n\n # GNN hyperparameters.\n config.model = 'GraphConvNet'\n config.message_passing_steps = 5\n config.latent_size = 256\n config.dropout_rate = 0.1\n config.num_mlp_layers = 2\n config.num_classes = 128\n config.skip_connections = True\n config.layer_norm = True\n return config\n" }, { "path": "examples/ogbg_molpcba/configs/default_graph_net.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Defines the default hyperparameters and training configuration.\n\nUses a GraphNetwork model (https://arxiv.org/abs/1806.01261).\n\"\"\"\n\nimport ml_collections\n\n\ndef get_config():\n \"\"\"Get the hyperparameter configuration for the GraphNetwork model.\"\"\"\n config = ml_collections.ConfigDict()\n\n # Optimizer.\n config.optimizer = 'adam'\n config.learning_rate = 1e-3\n\n # Training hyperparameters.\n config.batch_size = 256\n config.num_train_steps = 100_000\n config.log_every_steps = 100\n config.eval_every_steps = 10_000\n config.checkpoint_every_steps = 10_000\n config.add_virtual_node = True\n config.add_undirected_edges = True\n config.add_self_loops = True\n\n # GNN hyperparameters.\n config.model = 'GraphNet'\n config.message_passing_steps = 5\n config.latent_size = 256\n config.dropout_rate = 0.1\n config.num_mlp_layers = 1\n config.num_classes = 128\n config.use_edge_model = True\n config.skip_connections = True\n config.layer_norm = True\n return config\n" }, { "path": "examples/ogbg_molpcba/configs/hparam_sweep.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Defines a sweep for the hyperparameters for the GNN.\"\"\"\n\nimport ml_collections\n\n\ndef get_config():\n \"\"\"Get the default hyperparameter configuration.\"\"\"\n config = ml_collections.ConfigDict()\n\n # Optimizer.\n config.optimizer = 'adam'\n config.learning_rate = 1e-3\n\n # Training hyperparameters.\n config.batch_size = 256\n config.num_train_steps = 500_000\n config.log_every_steps = 50\n config.eval_every_steps = 1_000\n config.checkpoint_every_steps = 10_000\n config.add_virtual_node = True\n config.add_undirected_edges = True\n config.add_self_loops = True\n\n # GNN hyperparameters.\n config.model = 'GraphConvNet'\n config.message_passing_steps = 5\n config.latent_size = 256\n config.dropout_rate = 0.1\n config.num_mlp_layers = 2\n config.num_classes = 128\n config.skip_connections = True\n config.layer_norm = True\n\n return config\n\n\ndef sweep(add):\n for add_virtual_node in (True, False):\n for add_undirected_edges in (True, False):\n for add_self_loops in (True, False):\n for layer_norm in (True, False):\n for skip_connections in (True, False):\n add(\n add_virtual_node=add_virtual_node,\n add_undirected_edges=add_undirected_edges,\n add_self_loops=add_self_loops,\n layer_norm=layer_norm,\n skip_connections=skip_connections,\n )\n" }, { "path": "examples/ogbg_molpcba/configs/test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Defines a CPU-friendly test configuration.\"\"\"\n\nimport ml_collections\n\n\ndef get_config():\n \"\"\"Get the default hyperparameter configuration.\"\"\"\n config = ml_collections.ConfigDict()\n\n # Optimizer.\n config.optimizer = 'adam'\n config.learning_rate = 1e-3\n\n # Training hyperparameters.\n config.batch_size = 32\n config.num_train_steps = 10\n config.log_every_steps = 5\n config.eval_every_steps = 5\n config.checkpoint_every_steps = 5\n config.add_virtual_node = True\n config.add_undirected_edges = True\n config.add_self_loops = True\n\n # GNN hyperparameters.\n config.model = 'GraphConvNet'\n config.message_passing_steps = 5\n config.latent_size = 256\n config.dropout_rate = 0.1\n config.num_mlp_layers = 1\n config.num_classes = 128\n config.skip_connections = False\n config.layer_norm = False\n\n return config\n" }, { "path": "examples/ogbg_molpcba/input_pipeline.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Exposes the ogbg-molpcba dataset in a convenient format.\"\"\"\n\nimport functools\nfrom typing import Dict, NamedTuple\n\nimport jraph\nimport numpy as np\nimport tensorflow as tf\nimport tensorflow_datasets as tfds\n\n\nclass GraphsTupleSize(NamedTuple):\n \"\"\"Helper class to represent padding and graph sizes.\"\"\"\n\n n_node: int\n n_edge: int\n n_graph: int\n\n\ndef get_raw_datasets() -> dict[str, tf.data.Dataset]:\n \"\"\"Returns datasets as tf.data.Dataset, organized by split.\"\"\"\n ds_builder = tfds.builder('ogbg_molpcba')\n ds_builder.download_and_prepare()\n ds_splits = ['train', 'validation', 'test']\n datasets = {split: ds_builder.as_dataset(split=split) for split in ds_splits}\n return datasets\n\n\ndef get_datasets(\n batch_size: int,\n add_virtual_node: bool = True,\n add_undirected_edges: bool = True,\n add_self_loops: bool = True,\n) -> dict[str, tf.data.Dataset]:\n \"\"\"Returns datasets of batched GraphsTuples, organized by split.\"\"\"\n if batch_size <= 1:\n raise ValueError('Batch size must be > 1 to account for padding graphs.')\n\n # Obtain the original datasets.\n datasets = get_raw_datasets()\n\n # Construct the GraphsTuple converter function.\n convert_to_graphs_tuple_fn = functools.partial(\n convert_to_graphs_tuple,\n add_virtual_node=add_virtual_node,\n add_undirected_edges=add_undirected_edges,\n add_self_loops=add_self_loops,\n )\n\n # Process each split separately.\n for split_name in datasets:\n # Convert to GraphsTuple.\n datasets[split_name] = datasets[split_name].map(\n convert_to_graphs_tuple_fn,\n num_parallel_calls=tf.data.AUTOTUNE,\n deterministic=True,\n )\n\n # Compute the padding budget for the requested batch size.\n budget = estimate_padding_budget_for_batch_size(\n datasets['train'], batch_size, num_estimation_graphs=100\n )\n\n # Pad an example graph to see what the output shapes will be.\n # We will use this shape information when creating the tf.data.Dataset.\n example_graph = next(datasets['train'].as_numpy_iterator())\n example_padded_graph = jraph.pad_with_graphs(example_graph, *budget)\n padded_graphs_spec = specs_from_graphs_tuple(example_padded_graph)\n\n # Process each split separately.\n for split_name, dataset_split in datasets.items():\n # Repeat and shuffle the training split.\n if split_name == 'train':\n dataset_split = dataset_split.shuffle(100, reshuffle_each_iteration=True)\n dataset_split = dataset_split.repeat()\n\n # Batch and pad each split.\n batching_fn = functools.partial(\n jraph.dynamically_batch,\n graphs_tuple_iterator=iter(dataset_split),\n n_node=budget.n_node,\n n_edge=budget.n_edge,\n n_graph=budget.n_graph,\n )\n dataset_split = tf.data.Dataset.from_generator(\n batching_fn, output_signature=padded_graphs_spec\n )\n\n # We cache the validation and test sets, since these are small.\n if split_name in ['validation', 'test']:\n dataset_split = dataset_split.cache()\n\n datasets[split_name] = dataset_split\n return datasets\n\n\ndef convert_to_graphs_tuple(\n graph: dict[str, tf.Tensor],\n add_virtual_node: bool,\n add_undirected_edges: bool,\n add_self_loops: bool,\n) -> jraph.GraphsTuple:\n \"\"\"Converts a dictionary of tf.Tensors to a GraphsTuple.\"\"\"\n num_nodes = tf.squeeze(graph['num_nodes'])\n num_edges = tf.squeeze(graph['num_edges'])\n nodes = graph['node_feat']\n edges = graph['edge_feat']\n edge_feature_dim = edges.shape[-1]\n labels = graph['labels']\n senders = graph['edge_index'][:, 0]\n receivers = graph['edge_index'][:, 1]\n\n # Add a virtual node connected to all other nodes.\n # The feature vectors for the virtual node\n # and the new edges are set to all zeros.\n if add_virtual_node:\n nodes = tf.concat([nodes, tf.zeros_like(nodes[0, None])], axis=0)\n senders = tf.concat([senders, tf.range(num_nodes)], axis=0)\n receivers = tf.concat(\n [receivers, tf.fill((num_nodes,), num_nodes + 1)], axis=0\n )\n edges = tf.concat([edges, tf.zeros((num_nodes, edge_feature_dim))], axis=0)\n num_edges += num_nodes\n num_nodes += 1\n\n # Make edges undirected, by adding edges with senders and receivers flipped.\n # The feature vector for the flipped edge is the same as the original edge.\n if add_undirected_edges:\n new_senders = tf.concat([senders, receivers], axis=0)\n new_receivers = tf.concat([receivers, senders], axis=0)\n edges = tf.concat([edges, edges], axis=0)\n senders, receivers = new_senders, new_receivers\n num_edges *= 2\n\n # Add self-loops for each node.\n # The feature vectors for the self-loops are set to all zeros.\n if add_self_loops:\n senders = tf.concat([senders, tf.range(num_nodes)], axis=0)\n receivers = tf.concat([receivers, tf.range(num_nodes)], axis=0)\n edges = tf.concat([edges, tf.zeros((num_nodes, edge_feature_dim))], axis=0)\n num_edges += num_nodes\n\n return jraph.GraphsTuple(\n n_node=tf.expand_dims(num_nodes, 0),\n n_edge=tf.expand_dims(num_edges, 0),\n nodes=nodes,\n edges=edges,\n senders=senders,\n receivers=receivers,\n globals=tf.expand_dims(labels, axis=0),\n )\n\n\ndef estimate_padding_budget_for_batch_size(\n dataset: tf.data.Dataset, batch_size: int, num_estimation_graphs: int\n) -> GraphsTupleSize:\n \"\"\"Estimates the padding budget for a dataset of unbatched GraphsTuples.\n\n Args:\n dataset: A dataset of unbatched GraphsTuples.\n batch_size: The intended batch size. Note that no batching is performed by\n this function.\n num_estimation_graphs: How many graphs to take from the dataset to estimate\n the distribution of number of nodes and edges per graph.\n\n Returns:\n padding_budget: The padding budget for batching and padding the graphs\n in this dataset to the given batch size.\n \"\"\"\n\n def next_multiple_of_64(val: float):\n \"\"\"Returns the next multiple of 64 after val.\"\"\"\n return 64 * (1 + int(val // 64))\n\n if batch_size <= 1:\n raise ValueError('Batch size must be > 1 to account for padding graphs.')\n\n total_num_nodes = 0\n total_num_edges = 0\n for graph in dataset.take(num_estimation_graphs).as_numpy_iterator():\n graph_size = get_graphs_tuple_size(graph)\n if graph_size.n_graph != 1:\n raise ValueError('Dataset contains batched GraphTuples.')\n\n total_num_nodes += graph_size.n_node\n total_num_edges += graph_size.n_edge\n\n num_nodes_per_graph_estimate = total_num_nodes / num_estimation_graphs\n num_edges_per_graph_estimate = total_num_edges / num_estimation_graphs\n\n padding_budget = GraphsTupleSize(\n n_node=next_multiple_of_64(num_nodes_per_graph_estimate * batch_size),\n n_edge=next_multiple_of_64(num_edges_per_graph_estimate * batch_size),\n n_graph=batch_size,\n )\n return padding_budget\n\n\ndef specs_from_graphs_tuple(graph: jraph.GraphsTuple):\n \"\"\"Returns a tf.TensorSpec corresponding to this graph.\"\"\"\n\n def get_tensor_spec(array: np.ndarray):\n shape = list(array.shape)\n dtype = array.dtype\n return tf.TensorSpec(shape=shape, dtype=dtype)\n\n specs = {}\n for field in [\n 'nodes',\n 'edges',\n 'senders',\n 'receivers',\n 'globals',\n 'n_node',\n 'n_edge',\n ]:\n field_sample = getattr(graph, field)\n specs[field] = get_tensor_spec(field_sample)\n return jraph.GraphsTuple(**specs)\n\n\ndef get_graphs_tuple_size(graph: jraph.GraphsTuple):\n \"\"\"Returns the number of nodes, edges and graphs in a GraphsTuple.\"\"\"\n return GraphsTupleSize(\n n_node=np.sum(graph.n_node),\n n_edge=np.sum(graph.n_edge),\n n_graph=np.shape(graph.n_node)[0],\n )\n" }, { "path": "examples/ogbg_molpcba/input_pipeline_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for flax.examples.ogbg_molpcba.input_pipeline.\"\"\"\n\nfrom absl.testing import absltest\nfrom absl.testing import parameterized\nimport input_pipeline\nimport jraph\nimport tensorflow as tf\n\n\ndef get_dummy_datasets(dataset_length: int):\n \"\"\"Returns a set of datasets of unbatched GraphsTuples.\"\"\"\n # The dummy graph.\n num_nodes = 3\n num_edges = 4\n dummy_graph = jraph.GraphsTuple(\n n_node=tf.expand_dims(num_nodes, 0),\n n_edge=tf.expand_dims(num_edges, 0),\n senders=tf.zeros(num_edges, dtype=tf.int32),\n receivers=tf.ones(num_edges, dtype=tf.int32),\n nodes=tf.zeros((num_nodes, 9)),\n edges=tf.ones((num_edges, 3)),\n globals=tf.ones((1, 128), dtype=tf.int64),\n )\n graphs_spec = input_pipeline.specs_from_graphs_tuple(dummy_graph)\n\n # Yields a set of graphs for the current split.\n def get_dummy_graphs():\n for _ in range(dataset_length):\n yield dummy_graph\n\n datasets = {}\n for split in ['train', 'validation', 'test']:\n datasets[split] = tf.data.Dataset.from_generator(\n get_dummy_graphs, output_signature=graphs_spec\n )\n return datasets\n\n\nclass InputPipelineTest(parameterized.TestCase):\n\n def setUp(self):\n super().setUp()\n dataset_length = 20\n self.datasets = get_dummy_datasets(dataset_length)\n\n @parameterized.product(\n valid_batch_size=[2, 5, 12, 15],\n )\n def test_estimate_padding_budget_valid(self, valid_batch_size):\n budget = input_pipeline.estimate_padding_budget_for_batch_size(\n self.datasets['train'], valid_batch_size, num_estimation_graphs=1\n )\n self.assertEqual(budget.n_graph, valid_batch_size)\n\n @parameterized.product(\n invalid_batch_size=[-1, 0, 1],\n )\n def test_estimate_padding_budget_invalid(self, invalid_batch_size):\n with self.assertRaises(ValueError):\n input_pipeline.estimate_padding_budget_for_batch_size(\n self.datasets['train'], invalid_batch_size, num_estimation_graphs=1\n )\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "examples/ogbg_molpcba/main.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Main file for running the ogbg-molpcba example.\n\nThis file is intentionally kept short. The majority for logic is in libraries\nthat can be easily tested and imported in Colab.\n\"\"\"\n\nfrom absl import app\nfrom absl import flags\nfrom absl import logging\nfrom clu import platform\nimport jax\nfrom ml_collections import config_flags\nimport tensorflow as tf\n\nimport train\n\n\nFLAGS = flags.FLAGS\n\nflags.DEFINE_string('workdir', None, 'Directory to store model data.')\nconfig_flags.DEFINE_config_file(\n 'config',\n None,\n 'File path to the training hyperparameter configuration.',\n lock_config=True,\n)\n\n\ndef main(argv):\n if len(argv) > 1:\n raise app.UsageError('Too many command-line arguments.')\n\n # Hide any GPUs from TensorFlow. Otherwise TF might reserve memory and make\n # it unavailable to JAX.\n tf.config.experimental.set_visible_devices([], 'GPU')\n\n # This example only supports single-host training on a single device.\n logging.info('JAX host: %d / %d', jax.process_index(), jax.process_count())\n logging.info('JAX local devices: %r', jax.local_devices())\n\n # Add a note so that we can tell which task is which JAX host.\n # (Depending on the platform task 0 is not guaranteed to be host 0)\n platform.work_unit().set_task_status(\n f'process_index: {jax.process_index()}, '\n f'process_count: {jax.process_count()}'\n )\n platform.work_unit().create_artifact(\n platform.ArtifactType.DIRECTORY, FLAGS.workdir, 'workdir'\n )\n\n train.train_and_evaluate(FLAGS.config, FLAGS.workdir)\n\n\nif __name__ == '__main__':\n flags.mark_flags_as_required(['config', 'workdir'])\n app.run(main)\n" }, { "path": "examples/ogbg_molpcba/models.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Definition of the GNN model.\"\"\"\n\nfrom collections.abc import Callable, Sequence\n\nfrom flax import linen as nn\nimport jax.numpy as jnp\nimport jraph\n\n\ndef add_graphs_tuples(\n graphs: jraph.GraphsTuple, other_graphs: jraph.GraphsTuple\n) -> jraph.GraphsTuple:\n \"\"\"Adds the nodes, edges and global features from other_graphs to graphs.\"\"\"\n return graphs._replace(\n nodes=graphs.nodes + other_graphs.nodes,\n edges=graphs.edges + other_graphs.edges,\n globals=graphs.globals + other_graphs.globals,\n )\n\n\nclass MLP(nn.Module):\n \"\"\"A multi-layer perceptron.\"\"\"\n\n feature_sizes: Sequence[int]\n dropout_rate: float = 0\n deterministic: bool = True\n activation: Callable[[jnp.ndarray], jnp.ndarray] = nn.relu\n\n @nn.compact\n def __call__(self, inputs):\n x = inputs\n for size in self.feature_sizes:\n x = nn.Dense(features=size)(x)\n x = self.activation(x)\n x = nn.Dropout(rate=self.dropout_rate, deterministic=self.deterministic)(\n x\n )\n return x\n\n\nclass GraphNet(nn.Module):\n \"\"\"A complete Graph Network model defined with Jraph.\"\"\"\n\n latent_size: int\n num_mlp_layers: int\n message_passing_steps: int\n output_globals_size: int\n dropout_rate: float = 0\n skip_connections: bool = True\n use_edge_model: bool = True\n layer_norm: bool = True\n deterministic: bool = True\n\n @nn.compact\n def __call__(self, graphs: jraph.GraphsTuple) -> jraph.GraphsTuple:\n # We will first linearly project the original features as 'embeddings'.\n embedder = jraph.GraphMapFeatures(\n embed_node_fn=nn.Dense(self.latent_size),\n embed_edge_fn=nn.Dense(self.latent_size),\n embed_global_fn=nn.Dense(self.latent_size),\n )\n processed_graphs = embedder(graphs)\n\n # Now, we will apply a Graph Network once for each message-passing round.\n mlp_feature_sizes = [self.latent_size] * self.num_mlp_layers\n for _ in range(self.message_passing_steps):\n if self.use_edge_model:\n update_edge_fn = jraph.concatenated_args(\n MLP(\n mlp_feature_sizes,\n dropout_rate=self.dropout_rate,\n deterministic=self.deterministic,\n )\n )\n else:\n update_edge_fn = None\n\n update_node_fn = jraph.concatenated_args(\n MLP(\n mlp_feature_sizes,\n dropout_rate=self.dropout_rate,\n deterministic=self.deterministic,\n )\n )\n update_global_fn = jraph.concatenated_args(\n MLP(\n mlp_feature_sizes,\n dropout_rate=self.dropout_rate,\n deterministic=self.deterministic,\n )\n )\n\n graph_net = jraph.GraphNetwork(\n update_node_fn=update_node_fn,\n update_edge_fn=update_edge_fn,\n update_global_fn=update_global_fn,\n )\n\n if self.skip_connections:\n processed_graphs = add_graphs_tuples(\n graph_net(processed_graphs), processed_graphs\n )\n else:\n processed_graphs = graph_net(processed_graphs)\n\n if self.layer_norm:\n processed_graphs = processed_graphs._replace(\n nodes=nn.LayerNorm()(processed_graphs.nodes),\n edges=nn.LayerNorm()(processed_graphs.edges),\n globals=nn.LayerNorm()(processed_graphs.globals),\n )\n\n # Since our graph-level predictions will be at globals, we will\n # decode to get the required output logits.\n decoder = jraph.GraphMapFeatures(\n embed_global_fn=nn.Dense(self.output_globals_size)\n )\n processed_graphs = decoder(processed_graphs)\n\n return processed_graphs\n\n\nclass GraphConvNet(nn.Module):\n \"\"\"A Graph Convolution Network + Pooling model defined with Jraph.\"\"\"\n\n latent_size: int\n num_mlp_layers: int\n message_passing_steps: int\n output_globals_size: int\n dropout_rate: float = 0\n skip_connections: bool = True\n layer_norm: bool = True\n deterministic: bool = True\n pooling_fn: Callable[\n [jnp.ndarray, jnp.ndarray, jnp.ndarray], # pytype: disable=annotation-type-mismatch # jax-ndarray\n jnp.ndarray,\n ] = jraph.segment_mean\n\n def pool(self, graphs: jraph.GraphsTuple) -> jraph.GraphsTuple:\n \"\"\"Pooling operation, taken from Jraph.\"\"\"\n\n # Equivalent to jnp.sum(n_node), but JIT-able.\n sum_n_node = graphs.nodes.shape[0] # pytype: disable=attribute-error # jax-ndarray\n # To aggregate nodes from each graph to global features,\n # we first construct tensors that map the node to the corresponding graph.\n # Example: if you have `n_node=[1,2]`, we construct the tensor [0, 1, 1].\n n_graph = graphs.n_node.shape[0]\n node_graph_indices = jnp.repeat(\n jnp.arange(n_graph),\n graphs.n_node,\n axis=0,\n total_repeat_length=sum_n_node,\n )\n # We use the aggregation function to pool the nodes per graph.\n pooled = self.pooling_fn(graphs.nodes, node_graph_indices, n_graph) # pytype: disable=wrong-arg-types # jax-ndarray\n return graphs._replace(globals=pooled)\n\n @nn.compact\n def __call__(self, graphs: jraph.GraphsTuple) -> jraph.GraphsTuple:\n # We will first linearly project the original node features as 'embeddings'.\n embedder = jraph.GraphMapFeatures(embed_node_fn=nn.Dense(self.latent_size))\n processed_graphs = embedder(graphs)\n\n # Now, we will apply the GCN once for each message-passing round.\n for _ in range(self.message_passing_steps):\n mlp_feature_sizes = [self.latent_size] * self.num_mlp_layers\n update_node_fn = jraph.concatenated_args(\n MLP(\n mlp_feature_sizes,\n dropout_rate=self.dropout_rate,\n deterministic=self.deterministic,\n )\n )\n graph_conv = jraph.GraphConvolution(\n update_node_fn=update_node_fn, add_self_edges=True\n )\n\n if self.skip_connections:\n processed_graphs = add_graphs_tuples(\n graph_conv(processed_graphs), processed_graphs\n )\n else:\n processed_graphs = graph_conv(processed_graphs)\n\n if self.layer_norm:\n processed_graphs = processed_graphs._replace(\n nodes=nn.LayerNorm()(processed_graphs.nodes),\n )\n\n # We apply the pooling operation to get a 'global' embedding.\n processed_graphs = self.pool(processed_graphs)\n\n # Now, we decode this to get the required output logits.\n decoder = jraph.GraphMapFeatures(\n embed_global_fn=nn.Dense(self.output_globals_size)\n )\n processed_graphs = decoder(processed_graphs)\n\n return processed_graphs\n" }, { "path": "examples/ogbg_molpcba/models_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for flax.examples.ogbg_molpcba.models.\"\"\"\n\nfrom absl.testing import absltest\nfrom absl.testing import parameterized\nimport jax\nimport jax.numpy as jnp\nimport jraph\n\nimport models\n\n\nclass ModelsTest(parameterized.TestCase):\n\n def setUp(self):\n super().setUp()\n self.rngs = {\n 'params': jax.random.key(0),\n 'dropout': jax.random.key(1),\n }\n n_node = jnp.arange(3, 11)\n n_edge = jnp.arange(4, 12)\n total_n_node = jnp.sum(n_node)\n total_n_edge = jnp.sum(n_edge)\n n_graph = n_node.shape[0]\n feature_dim = 10\n self.graphs = jraph.GraphsTuple(\n n_node=n_node,\n n_edge=n_edge,\n senders=jnp.zeros(total_n_edge, dtype=jnp.int32),\n receivers=jnp.ones(total_n_edge, dtype=jnp.int32),\n nodes=jnp.ones((total_n_node, feature_dim)),\n edges=jnp.zeros((total_n_edge, feature_dim)),\n globals=jnp.zeros((n_graph, feature_dim)),\n )\n\n @parameterized.product(\n dropout_rate=[0.0, 0.5, 1.0], output_size=[50, 100], num_layers=[2]\n )\n def test_mlp(self, dropout_rate, output_size, num_layers):\n # Input definition.\n nodes = self.graphs.nodes\n\n # Model definition.\n mlp = models.MLP(\n feature_sizes=[output_size] * num_layers,\n dropout_rate=dropout_rate,\n activation=lambda x: x,\n deterministic=False,\n )\n nodes_after_mlp, _ = mlp.init_with_output(self.rngs, nodes)\n\n # Test that dropout actually worked.\n num_masked_entries = jnp.sum(nodes_after_mlp == 0)\n num_total_entries = jnp.size(nodes_after_mlp)\n self.assertLessEqual(\n num_masked_entries, (dropout_rate + 0.05) * num_total_entries\n )\n self.assertLessEqual(\n (dropout_rate - 0.05) * num_total_entries, num_masked_entries\n )\n\n # Test the shape of the output.\n self.assertEqual(nodes_after_mlp.shape[-1], output_size)\n\n @parameterized.parameters(\n {\n 'latent_size': 5,\n 'output_globals_size': 15,\n 'use_edge_model': True,\n },\n {\n 'latent_size': 5,\n 'output_globals_size': 15,\n 'use_edge_model': False,\n },\n )\n def test_graph_net(\n self, latent_size: int, output_globals_size: int, use_edge_model: bool\n ):\n # Input definition.\n graphs = self.graphs\n num_nodes = jnp.sum(graphs.n_node)\n num_edges = jnp.sum(graphs.n_edge)\n num_graphs = graphs.n_node.shape[0]\n\n # Model definition.\n net = models.GraphNet(\n latent_size=latent_size,\n num_mlp_layers=2,\n message_passing_steps=2,\n output_globals_size=output_globals_size,\n use_edge_model=use_edge_model,\n )\n output, _ = net.init_with_output(self.rngs, graphs)\n\n # Output should be graph with the same topology, but a\n # different number of features.\n self.assertIsInstance(output, jraph.GraphsTuple)\n self.assertSequenceAlmostEqual(output.n_node, graphs.n_node)\n self.assertSequenceAlmostEqual(output.n_edge, graphs.n_edge)\n self.assertSequenceAlmostEqual(output.senders, graphs.senders)\n self.assertSequenceAlmostEqual(output.receivers, graphs.receivers)\n self.assertEqual(output.nodes.shape, (num_nodes, latent_size))\n self.assertEqual(output.edges.shape, (num_edges, latent_size))\n self.assertEqual(output.globals.shape, (num_graphs, output_globals_size))\n\n @parameterized.parameters(\n {'latent_size': 15, 'output_globals_size': 15},\n {'latent_size': 5, 'output_globals_size': 5},\n )\n def test_graph_conv_net(self, latent_size: int, output_globals_size: int):\n graphs = self.graphs\n num_nodes = jnp.sum(graphs.n_node)\n num_graphs = graphs.n_node.shape[0]\n\n # Model definition.\n net = models.GraphConvNet(\n latent_size=latent_size,\n num_mlp_layers=2,\n message_passing_steps=2,\n output_globals_size=output_globals_size,\n )\n output, _ = net.init_with_output(self.rngs, graphs)\n\n # Output should be graph with the same topology, but a\n # different number of features.\n self.assertIsInstance(output, jraph.GraphsTuple)\n self.assertSequenceAlmostEqual(output.n_node, graphs.n_node)\n self.assertSequenceAlmostEqual(output.n_edge, graphs.n_edge)\n self.assertSequenceAlmostEqual(\n output.edges.flatten(), graphs.edges.flatten()\n )\n self.assertSequenceAlmostEqual(output.senders, graphs.senders)\n self.assertSequenceAlmostEqual(output.receivers, graphs.receivers)\n self.assertEqual(output.nodes.shape, (num_nodes, latent_size))\n self.assertEqual(output.globals.shape, (num_graphs, output_globals_size))\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "examples/ogbg_molpcba/ogbg_molpcba.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"# Flax ogbg-molpcba Example\\n\",\n \"\\n\",\n \"\\\"Open\\n\",\n \"\\n\",\n \"Demonstration notebook for\\n\",\n \"https://github.com/google/flax/tree/main/examples/ogbg_molpcba.\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Setup\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 1,\n \"metadata\": {\n \"outputId\": \"6508ab2f-b0e5-4693-f6a0-7bc495ec1344\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"\\u001b[?25l\\r\",\n \"\\u001b[K |████▏ | 10 kB 20.9 MB/s eta 0:00:01\\r\",\n \"\\u001b[K |████████▍ | 20 kB 10.6 MB/s eta 0:00:01\\r\",\n \"\\u001b[K |████████████▋ | 30 kB 8.8 MB/s eta 0:00:01\\r\",\n \"\\u001b[K |████████████████▉ | 40 kB 8.0 MB/s eta 0:00:01\\r\",\n \"\\u001b[K |█████████████████████ | 51 kB 5.1 MB/s eta 0:00:01\\r\",\n \"\\u001b[K |█████████████████████████▎ | 61 kB 5.3 MB/s eta 0:00:01\\r\",\n \"\\u001b[K |█████████████████████████████▌ | 71 kB 5.5 MB/s eta 0:00:01\\r\",\n \"\\u001b[K |████████████████████████████████| 77 kB 2.7 MB/s \\n\",\n \"\\u001b[K |████████████████████████████████| 88 kB 5.5 MB/s \\n\",\n \"\\u001b[K |████████████████████████████████| 4.0 MB 38.2 MB/s \\n\",\n \"\\u001b[K |████████████████████████████████| 73 kB 1.3 MB/s \\n\",\n \"\\u001b[K |████████████████████████████████| 118 kB 38.3 MB/s \\n\",\n \"\\u001b[K |████████████████████████████████| 57 kB 3.6 MB/s \\n\",\n \"\\u001b[?25h Building wheel for flax (setup.py) ... \\u001b[?25l\\u001b[?25hdone\\n\"\n ]\n }\n ],\n \"source\": [\n \"# Install clu, ml-collections, latest Flax version, and tensorflow_datasets.\\n\",\n \"!pip install -U -q clu ml-collections git+https://github.com/google/flax tfds_nightly jraph\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 2,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"example_directory = 'examples/ogbg_molpcba'\\n\",\n \"editor_relpaths = ('configs/default.py', 'input_pipeline.py', 'models.py', 'train.py')\\n\",\n \"\\n\",\n \"repo, branch = 'https://github.com/google/flax', 'main'\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 3,\n \"metadata\": {\n \"cellView\": \"form\",\n \"outputId\": \"8261a349-b41e-4e1b-a2ca-8d23412155be\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"Cloning into 'flaxrepo'...\\n\",\n \"remote: Enumerating objects: 350, done.\\u001b[K\\n\",\n \"remote: Counting objects: 100% (350/350), done.\\u001b[K\\n\",\n \"remote: Compressing objects: 100% (312/312), done.\\u001b[K\\n\",\n \"remote: Total 350 (delta 65), reused 140 (delta 20), pack-reused 0\\u001b[K\\n\",\n \"Receiving objects: 100% (350/350), 2.10 MiB | 20.10 MiB/s, done.\\n\",\n \"Resolving deltas: 100% (65/65), done.\\n\"\n ]\n },\n {\n \"data\": {\n \"text/html\": [\n \"

WARNING : Editing in VM - changes lost after reboot!!

\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"application/javascript\": [\n \"\\n\",\n \" ((filepath) => {{\\n\",\n \" if (!google.colab.kernel.accessAllowed) {{\\n\",\n \" return;\\n\",\n \" }}\\n\",\n \" google.colab.files.view(filepath);\\n\",\n \" }})(\\\"/content/examples/ogbg_molpcba/configs/default.py\\\")\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"application/javascript\": [\n \"\\n\",\n \" ((filepath) => {{\\n\",\n \" if (!google.colab.kernel.accessAllowed) {{\\n\",\n \" return;\\n\",\n \" }}\\n\",\n \" google.colab.files.view(filepath);\\n\",\n \" }})(\\\"/content/examples/ogbg_molpcba/input_pipeline.py\\\")\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"application/javascript\": [\n \"\\n\",\n \" ((filepath) => {{\\n\",\n \" if (!google.colab.kernel.accessAllowed) {{\\n\",\n \" return;\\n\",\n \" }}\\n\",\n \" google.colab.files.view(filepath);\\n\",\n \" }})(\\\"/content/examples/ogbg_molpcba/models.py\\\")\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"application/javascript\": [\n \"\\n\",\n \" ((filepath) => {{\\n\",\n \" if (!google.colab.kernel.accessAllowed) {{\\n\",\n \" return;\\n\",\n \" }}\\n\",\n \" google.colab.files.view(filepath);\\n\",\n \" }})(\\\"/content/examples/ogbg_molpcba/train.py\\\")\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"# (If you run this code in Jupyter[lab], then you're already in the\\n\",\n \"# example directory and nothing needs to be done.)\\n\",\n \"\\n\",\n \"#@markdown **Fetch newest Flax version and copy of example code.**\\n\",\n \"#@markdown\\n\",\n \"#@markdown **If you select no** below, then the files will be stored on the\\n\",\n \"#@markdown *ephemeral* Colab VM. **After some time of inactivity, this VM will\\n\",\n \"#@markdown be restarted and any changes will be lost**.\\n\",\n \"#@markdown\\n\",\n \"#@markdown **If you select yes** below, then you will be asked for your\\n\",\n \"#@markdown credentials to mount your personal Google Drive. In this case, all\\n\",\n \"#@markdown changes you make will *persist*. Even if you re-run this\\n\",\n \"#@markdown Colab notebook later on, the files will still exist. You can\\n\",\n \"#@markdown remove directories inside your Drive's `flax/` root if you want to\\n\",\n \"#@markdown manually revert these files.\\n\",\n \"\\n\",\n \"if 'google.colab' in str(get_ipython()):\\n\",\n \" import os\\n\",\n \" os.chdir('/content')\\n\",\n \" # Download Flax repo from Github.\\n\",\n \" if not os.path.isdir('flaxrepo'):\\n\",\n \" !git clone --depth=1 -b $branch $repo flaxrepo\\n\",\n \" # Copy example files & change directory.\\n\",\n \" mount_gdrive = 'no' #@param ['yes', 'no']\\n\",\n \" if mount_gdrive == 'yes':\\n\",\n \" DISCLAIMER = 'Note : Editing in your Google Drive, changes will persist.'\\n\",\n \" from google.colab import drive\\n\",\n \" drive.mount('/content/gdrive')\\n\",\n \" example_root_path = f'/content/gdrive/My Drive/flax/{example_directory}'\\n\",\n \" else:\\n\",\n \" DISCLAIMER = 'WARNING : Editing in VM - changes lost after reboot!!'\\n\",\n \" example_root_path = f'/content/{example_directory}'\\n\",\n \" from IPython import display\\n\",\n \" display.display(display.HTML(\\n\",\n \" f'

{DISCLAIMER}

'))\\n\",\n \" if not os.path.isdir(example_root_path):\\n\",\n \" os.makedirs(example_root_path)\\n\",\n \" !cp -r flaxrepo/$example_directory/* \\\"$example_root_path\\\"\\n\",\n \" os.chdir(example_root_path)\\n\",\n \" from google.colab import files\\n\",\n \" for relpath in editor_relpaths:\\n\",\n \" s = open(f'{example_root_path}/{relpath}').read()\\n\",\n \" open(f'{example_root_path}/{relpath}', 'w').write(\\n\",\n \" f'## {DISCLAIMER}\\\\n' + '#' * (len(DISCLAIMER) + 3) + '\\\\n\\\\n' + s)\\n\",\n \" files.view(f'{example_root_path}/{relpath}')\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 4,\n \"metadata\": {\n \"outputId\": \"14b17380-5077-4354-e651-027f3d933cfe\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"/content/examples/ogbg_molpcba\\n\"\n ]\n }\n ],\n \"source\": [\n \"# Display current working directory.\\n\",\n \"# Note: In Colab, running the above cell changes the working directory.\\n\",\n \"!pwd\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Imports\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 5,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# Base imports\\n\",\n \"from absl import logging\\n\",\n \"import flax\\n\",\n \"import jax.numpy as jnp\\n\",\n \"import jraph\\n\",\n \"from matplotlib import pyplot as plt\\n\",\n \"import numpy as np\\n\",\n \"import pprint\\n\",\n \"import tensorflow_datasets as tfds\\n\",\n \"logging.set_verbosity(logging.INFO)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 6,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# Local imports from current directory - auto reload.\\n\",\n \"# Any changes you make to train.py and other modules below will appear automatically.\\n\",\n \"%load_ext autoreload\\n\",\n \"%autoreload 2\\n\",\n \"import train\\n\",\n \"import input_pipeline\\n\",\n \"from configs import default as config_lib\\n\",\n \"config = config_lib.get_config()\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Dataset\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"TensorFlow Datasets supports customizable visualization of the ogbg_molpcba dataset.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 7,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# Visualization helpers\\n\",\n \"# Dictionaries used to map nodes and edges to colors.\\n\",\n \"atomic_numbers_to_elements = {\\n\",\n \" 6: 'C', 7: 'N', 8: 'O', 9: 'F', 14: 'Si',\\n\",\n \" 15: 'P', 16: 'S', 17: 'Cl', 35: 'Br',\\n\",\n \"}\\n\",\n \"elements_to_colors = {\\n\",\n \" element: f'C{index}'\\n\",\n \" for index, element in enumerate(atomic_numbers_to_elements.values())\\n\",\n \"}\\n\",\n \"bond_types_to_colors = {num: f'C{num}' for num in range(4)}\\n\",\n \"\\n\",\n \"# Node colors are atomic numbers.\\n\",\n \"def node_color_fn(graph):\\n\",\n \" atomic_numbers = 1 + graph['node_feat'][:, 0].numpy()\\n\",\n \" return {\\n\",\n \" index: elements_to_colors[atomic_numbers_to_elements[atomic_number]]\\n\",\n \" for index, atomic_number in enumerate(atomic_numbers)\\n\",\n \" }\\n\",\n \"\\n\",\n \"# Node labels are element names.\\n\",\n \"def node_label_fn(graph):\\n\",\n \" atomic_numbers = 1 + graph['node_feat'][:, 0].numpy()\\n\",\n \" return {\\n\",\n \" index: atomic_numbers_to_elements[atomic_number]\\n\",\n \" for index, atomic_number in enumerate(atomic_numbers)\\n\",\n \" }\\n\",\n \"\\n\",\n \"# Edge colors are bond types.\\n\",\n \"def edge_color_fn(graph):\\n\",\n \" bonds = graph['edge_index'].numpy()\\n\",\n \" bond_types = graph['edge_feat'][:, 0].numpy()\\n\",\n \" return {\\n\",\n \" tuple(bond): bond_types_to_colors[bond_type]\\n\",\n \" for bond, bond_type in zip(bonds, bond_types)\\n\",\n \" }\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 8,\n \"metadata\": {\n \"outputId\": \"d9336190-e685-43e8-e3f1-73e3f1ce1cd2\"\n },\n \"outputs\": [\n {\n \"name\": \"stderr\",\n \"output_type\": \"stream\",\n \"text\": [\n \"INFO:absl:Load pre-computed DatasetInfo (eg: splits, num examples,...) from GCS: ogbg_molpcba/0.1.2\\n\",\n \"INFO:absl:Load dataset info from /tmp/tmpmdoxgxq7tfds\\n\",\n \"INFO:absl:Generating dataset ogbg_molpcba (/root/tensorflow_datasets/ogbg_molpcba/0.1.2)\\n\"\n ]\n },\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"\\u001b[1mDownloading and preparing dataset 37.70 MiB (download: 37.70 MiB, generated: 822.53 MiB, total: 860.23 MiB) to /root/tensorflow_datasets/ogbg_molpcba/0.1.2...\\u001b[0m\\n\"\n ]\n },\n {\n \"data\": {\n \"application/vnd.jupyter.widget-view+json\": {\n \"model_id\": \"78bb0e82e40b4c3c937ae1f248ba1d54\",\n \"version_major\": 2,\n \"version_minor\": 0\n },\n \"text/plain\": [\n \"Dl Completed...: 0 url [00:00, ? url/s]\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"application/vnd.jupyter.widget-view+json\": {\n \"model_id\": \"f933ef3ecafb4973b5c54cc4ff039ca1\",\n \"version_major\": 2,\n \"version_minor\": 0\n },\n \"text/plain\": [\n \"Dl Size...: 0 MiB [00:00, ? MiB/s]\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"application/vnd.jupyter.widget-view+json\": {\n \"model_id\": \"0607a3a123334ae882e26c39c543c9e7\",\n \"version_major\": 2,\n \"version_minor\": 0\n },\n \"text/plain\": [\n \"Extraction completed...: 0 file [00:00, ? file/s]\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"name\": \"stderr\",\n \"output_type\": \"stream\",\n \"text\": [\n \"INFO:absl:Downloading https://snap.stanford.edu/ogb/data/graphproppred/csv_mol_download/pcba.zip into /root/tensorflow_datasets/downloads/snap.stan.edu_ogb_grap_csv_mol_down_pcbapc4I82Cv1THcU-IggPHK8IHZ8qM-BJ3VDk-q_rtqrf4.zip.tmp.495be609cf3840d3a942e7f9ba5d705e...\\n\"\n ]\n },\n {\n \"data\": {\n \"application/vnd.jupyter.widget-view+json\": {\n \"model_id\": \"83b57edf962a4faa9e875b4fecbb1eb1\",\n \"version_major\": 2,\n \"version_minor\": 0\n },\n \"text/plain\": [\n \"Generating splits...: 0%| | 0/3 [00:00\"\n ]\n },\n \"execution_count\": 8,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n },\n {\n \"data\": {\n \"image/png\": \"iVBORw0KGgoAAAANSUhEUgAAA1MAAAM9CAYAAAB5Rim2AAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjIsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+WH4yJAAAgAElEQVR4nOzdd3hUdfb48fe9d2p6IaGE0BJ6B6mCgFhWXOwitrUgYC+rruv2dcvPVXe/axfsYK/rWtaCSu9VekkoSSC9J9Pv/f0RQSAzyUzIzAQ4r+fh8XHaPUlm5t7zKecohmEghBBCCCGEECI0arQDEEIIIYQQQoiTkSRTQgghhBBCCNECkkwJIYQQQgghRAtIMiWEEEIIIYQQLSDJlBBCCCGEEEK0gKmpO9u1a2d069YtQqEIISJh3bp1pYZhpEU7jhMl309CnHpOhe8n+W4S4tTT1HdTk8lUt27dWLt2bXiiEkJEhaIo+6MdQ2uQ7ychTj2nwveTfDcJcepp6rtJlvkJIYQQQgghRAtIMiWEEEIIIYQQLSDJlBBCCCGEEEK0gCRTQgghhBBCCNECkkwJIYQQQgghRAtIMiWEEEIIIYQQLSDJlBBCCCGEEEK0gCRTQgghhBBCCNECkkwJIYQQQgghRAtIMiWEEEIIIYQQLWCK6tENA/YugmVPwaGN4K4HzQxx6XDGLTDkarAnRTVEIYQQorjayRur9vPhugIqHW50HWKtGmOyUpk1PouBnROjHaIQfum6waLdJcxdnMuOQ9XUu31YTSrpCVZuHNuNS4Z2Js4a3ctBIU5m0fv0/PAefPNHcFWBu+6n270OcFXDd4/At3+C/pfCBY+BLSFqoQohhDg9FVY5+c1Hm1maU4oCuLz6kfscHh+f/3CIBduKyUiy8cjFAxib3S56wQpxnLdW7edf3+zC4fZR5/Ydud3l1al2evn7Fzv46+fbuXxYZ353YT/sFi2K0Qpxcor8Mj/DgAV/gk/vhpqDxyZSR/PUg9cJWz6CuROgpjCiYQohhDi97Sqq4YInF7NwVzFur35MInWYbjQkVXtK6rj59TW8uyYvCpEKcSzDMPjdx5v5y2fbKa11H5NIHa3e7cPp0flgXT6XPLuMijp3hCMV4uQX+WRq2b9h1RzwOIJ7vM8FFQfgtQvBVRPe2IQQQgjgUJWDaXNWUFHvQTeCe47To/PH/27hq60y+Cei67GvdvLh+gIcHv9J1PFcXp3c0lqueWkVziCfI4RoENlkqnQPLHy0YdYpFIYXKvPg20fCE5cQQghxlPvf20SN0xvy85wenXvf2UitK/TnCtEaNudX8eqyvUEnUod5fAa5JbU8+/2eMEUmxKkpsnumVj4H+rEf7m7/rqHeA3vviSPWogDw0no3b/zgYeGNsT890OeCjW/CuY+A2R7JqIUQQpxG8srrWbe/Al+AKam6bQupXvMfPGX5qBY75vQeJI6dhq1zfwAUBT5en8/1Y7pFMOrw21dax9urD7C7uJZal5dEm4mBnZOYPjKT9HhbtMMTP3pxSS5uP0tSofn3rsurM2/Ffu6e3BOzJgWfhQhG5JIpdx1seht0T6O7fAY8ucrNb8Zbm3+dLR/B0GvDEKAQQggB81fsRzf8J1LVqz+matUHpJ53B7buw1A0E46963DsXnXkgrTe7WPOolyuG90VRVEavYZhGH5vb6sW7izm6e/2sKWgCl038ByVZC7eXcoz3+9hXHY77p7ckyGZUoE3mqrqPXy1tdDv0tRg3rsAXl3n2+1F/GxAxwhGLsTJK3LJ1J5vQfE/yvHgWAuPLXNx+wgLSbYmTjDuOlj7iiRTQgghwuaD9fl4fI2vRnVXHZVL3yR1yr3E9B575PaY7FHEZI865rFl9W52FdXSu0M8Hp/ON9uKmLMoh51FNTg9OiZVISXWwrWjunDNqK6kxQcxmBhhhmHwjy938Pry/QGXjB0uyvH9jmKW55Typ6n9mT6ySyTDFEf5elshmtr4OiqU926dy8dbqw5IMiVEkCKXTNUWge5/DfkZnTQmdjPxxHIXfz27maUCtcVhCE4IIYRoUO1ovIICwFWwA8PrJqbXmGZfw6QqlNQ4WbanlH8v2IXPMKhz/ZSQeHWD4hoXzy3M4bmFOUzoncZjlw8iKcbSaj/Hifp//9vB/BWBE6mjGTTsF/vzp1vRVIUrz8gMf4CikZJaFy4/f69Q3rsARdWu1g7Nr4JKB19vLaS01o1P10mOsXBmdjsGZEjfNnHyiFwy5fOA4X8NL8Ajk6yc+Uod94xq5kSiS9lOIYQQ4WEYRsC9Uj5HNWpMAorafC8eQzeYsziXtfsqmkxGjp7ZmfLkEj64bSydkqK/L/i7HUVBJ1JHc3h0fv/JFoZkJtGzfXyYohOBuL263yV+obx3Ady+wNdrJ8owDJbsLmXOohzW7q8AfvocmFSFfy/YTUayndsmZPHzwR2xmqT3lWjbIpdM2RJBNYHPfzI0IF3j571MPLrUTd+0JjY9WuXLWQghRMt4fDoLthUxZ3Euu4pqcHp8mDWVtHgrN4zpxrQzMrGZNb9JhGZPQK+vxtB9zV6Uurw6q/eW++1N5T8ug6JqF9PmrODzu8eTaDe36OdrLU8u2B0wkWquiIHHq/Piklweu2JwJEMWQILNjNmkNipAEcp7FyDWVYenuBhzenqrxufx6dz37ka+21FMvZ/eV17dwKv72FNcy+8/2cKcxTm8NXM07eLa3jJYIQ6LXDKVOarJmSmAP0+0MWxOLfePCfChUc3QfWLrxyaEEOKUZhgGLy3N5elv9zRacufy6uRXOPjnNzt54uudxNtMfhMJa0YfFJOZ+l0riO0zrsnjeXQDf1METSUiPsOgqNrJP7/eySMXDzjxH7qF9hTXsrPQf1/HYIoY+Az476aD/GFqf+KskS0afLob0S0FzU9xk1DeuxZ0hhXvJnfqX9Di4rAPGYJ96FDsQ4Zg690LxdyyRN+nG8yct5aVuWU4Pc0PMtS7feSW1DH16aV8cfd4kmPbzhJYIY4WuW+5dtnQfgAUrA34kOwUlav6m3lqtZuB6X5mp1QNRt8WxiCFEEKcanTd4IEPNvG/zYVNLls7fIHnq/eg0LAP6GiqNZakcddS/s0LKKqGrftQFNWEc99GnAd+IHnSzQ0PNHQURcHg2IvaYBIRj8/g/bX5PHxBX+yW6CxvenPVfrx+EsFQihioisJnmw5KMYoIG9g5kU5JNnJK6o65Pej3LoDJxB1/uZ30uHtx79uHY8MGHBs3UvnuO3gKDmLr3/+nBGvoEEzJyUHF9vhXO1mVWx5UInWYVzcorXVx82tr+PiOM4N+nhCRFNkho3H3wsezG6ryBfCHCVbm/+B/8y8dB0NqVpiCE0IIcSp65LOtzSZSR/OXSByWMPIy1Nhkqla8S+lnT6BY7FjbZ5Mw5qqfHqSojRKxUBIRRYFPNx1k2ojoFHHYVVTr93cQShGDerePfWWBz/UifG6bmM0fPtnSaBldMO9dBTgzK5X2CQ3FwKw9emDt0YOkyy8HwFddjWPTDzg2bKDijTc4+KtfYUpNPSa5smZno2jHDgTUu728vnxfi5aOenwGOwpr2JRXyWApvS/aoMgmU70ugKSuULr7SL+pffceuwcqM1HF+buExs812eHcv0QiSiGEEKeI9QcqeHdNvt+LuOb2/mAYDZnNceL6TyKu/yS/xzOrCjoGx+/fDzUReWfNgaglU3Uu/5V3Qy1iUFkfYGBUhNXPB3XkqW92ku/yoB/Xkqap9y6AzazxwPm9A96vJSQQN34cceMblgoaPh+uPTk4Nm7EsWED5a+9hre0FPuggQ0J1pAh2AcP5pOdVf4+SkBwM7Yur48Xl+TyzDXDQvxtCBF+kU2mNBPc8CnMGQ+1pcFX5jPZ4KKnoMuo5h8rhBBC/OjFxbk4vY0TqaAamIbYWNdmUkmOtVBW68J33NxUqIlIaW30KtfGB9jnFGoRg+Q2VOb9dKJVlPH46pe5tfc0qhUtYHXK49nMKk9fPZT+nYIvS65oGrbevbD17kXyVdMA8FZUNCRXGzdR9tLLOLds4dkJv6Te2nhWKdgZW92Ab7YVUVXvITEmusVZhDhe5HeGxraD2Uth/sVQvhfctYEfa45p+O8Vr0DvCyITnxBCiCYVVDqYv2IfX20tosbpQVEUkuxmLhuWwfQRXdrMRvHyOjff7SjGOO5aMpQld2ZVAaVhD1CgynyqAlaTxpisVIZ3Teb/vtnV6DGhJiJePXylqZszICORlXvLGjUuDqWIQaxFo1eHuHCGKfxwHzjAgRm3kHX5Zfzv6vO59qVVFFY5qfNTOe+wGIuGosDc68/gzOx2JxyDKTmZ+EmTiJ/UMAOmezwc/MPXjTchEtqMrUVT2V9ex6AYWeon2pbolNmJTYVZi2HvIlj2JOxfDiYrGDounwGGgTU2EcbeBUOuBbt8cIQQItp2FtbwyKdbWbu/AsMwcB91sV1S4+LJb3fz7wW7Oadve3738750TIxuv6RvthWiqY1nl0K5gPPoBiO7JzOyWyrzV+4/MspvYKAqCm6vznn92nPL+B4Mzkzi3TUHMGsqXv3Yi9dQEhFoKHEdDj7d4Lsdxbyz+gCHqhx4fAaJdjPje6ZxzagupMVbuXZ0F15Ztpfjr35DKWJgABcM6BiWn+F0UufyUlLjwuHxEWc10T7BhsXkv32Mc8cO8mbNpt3tt5E8fToAX983gcW7SnhhUQ4b8yqxmFR0w0BBQTcMUmIszJ7Qg0uHdQ5b5UU3gdvdhDRjq0Ct0/8SVCGiKXo1S1UVsiY1/KsqgKIt4Kxi60EH/9uv8ttZ14W8xEIIIUR4LNtTysx5a/32hjnscJWu/205xPKcUt6eNZo+HfzsgY2Q0lo3Tj97pUJdclft8PLA+b2555yerNtfQVmtG6+uk2A3MzQziaSjlrON6JaCfvxUGKElIhaTyqTerdvfx+nxMWdRLq8t34vbpx9TGh5gc0EVzy3cw1m90njgvN4M65LMityyRq8TTBEDs6oc6dclQmcYBhvyKnlxcS7f7ijGpCqoioJPN1AUuGpEJjeN7U6X1Jgjz6lft478u++hw+9+S8IFP63k0VSFSX3SmdQnnfyKenYW1lDj9GIza3RKsjEwIxElzNdaVlPjgixH4gtlxtYgahUuhWhK22gAkZjR8A+Ib1/Dgi3r+K0kUkII0SZsyqvkltfXBl0NTzegot7DVT82oO2cHNP8k8LA49P9tXoKecmd58dqEmZNZXSP1CYf2yMtjn4dE9iQV9novqAqAdJQUe0XY7s2G1ewKuvdXPfSKvYU1+IMsFTx8BLGBduLWLanlLsmZ7N2f3mjpX7QfBEDk6YyY1z31gn+NJNXXs/Nr62hoNKB0+NDN+D43XNvrNzPW6sOcFavNJ6cPgTf8qUcevg3dHr8ceLGBS4f3jk5JiqfRUVRSImxUFbXeB9gKDO2bp8e9dluIfxpG8nUUbqmxlJQ6cDl9WE1yQiEEEJEk083mPH6mqATqaPVOn3c8eZ6Prmz+WVt4ZBgM2M1qY32OoW65C7RHtqSu1snZvHLdzf63afSXCICMLJ7SqtdNDo9PqbPXUlOSa3fxOh4htFQTfCJr3YRY1ZRFSPgXjF/bGaVZ64ZSmZKdBLok9muohqueGE5tU6v30GAwxr+jgaLd5Uw9f99yRPf/Zs+zz+HfciQiMUaqutGd+WFRTmN3kuhzNgOzEikQ6It0qEL0aw2l0xZTCqdk+zsK62nd4f45p8ghBAibL7bUdyi3jAAPsNgZ2ENu4tq6Nk+8t/no3qkoPpZ5RDKBZzVpDIxxCV3k/uk0z0tlp2FNUElMEezmzUevqBvSM9pyl8/28be0rqQ4/DpBroB953bi38v2IXb63+W7zCTqmDWGhKpyX3bn2DUp5/iGifT566kxuENuCTueC6vzgGXjz9f/DAfDBwU1vhO1LWju/DCohy/9wUzYxtr0bh1gvQZFW1Tm0umALLS48gpqZVkSgghouyFRTmN9tdAkKXFAY+u8/LSvTx6eeQv9vp3SiQzxc6uosZVY4NdcgdwzaguIR3XpKm8MWMUU59ZSlGVC/fxTacCsJlVnrt2GBlJdt5ctZ+88npqnF5SYi3065jAOf3aY9YCb+Y/Xp3Lywfr8wPOLDWXDBtAvM3E+7PHMmdxDl9vK0JVftobBw2V4AwDLh+ewS3jetCtXWzQ8YmfPPXtbqodHr+JVJMNbTUTuyo9fL2tiCkD227Bj/R4GxN7p/H9jhK/n4emZmwVINZqYlKf1t1HKERraZPJVHZ6HHuKmyiZLoQQIuxKa11szq9qdHsopcV9OvxnQwH/77KBYd/o7s9tE7P47cdb/BbOaG7JnaLAhF5ptIuzhnzcpBgLn901nptfW8P2Q9VH9r/4E2PRUBz1PNBN4eMNBdz6xjpURTlmRjDWqqEpCteN7soNY7vRPqH55U7/2VDgd2YOgkuG690+XliUw+IHJ/HMNcOoqHPz300HyS2ppdrpITnGSp+O8fx8UEdiLG3ycuKkUO/28uG6Arx+3iDB/p2eX5jTppMpgMevHMyFTy7hYJWjUVPrptgtGvNnjPJbmVOItqBNfvtlp8WxaFdJtMMQQojTWnG1C4tJbTSSHEppcWgoL+706FGpxDVlYEee+W4P+8vq/V6sNsVu1njw/N4tPnai3cwHt45hY14lLy7JZcH2YqwmtWHKR2kobJGZHMOtE7I4mJvPo6sL8ZgO+k26Ds8OvrR0L/NW7OfVm0YwoltKk8d/dflev0lkKMlwWa2brQerGZCRSHKshRvGdgv59yCa9ummg36LF4fyd9pdVMOe4hqy09vuip4Em5kPbx/L9LkrKahwNLsXTwVibSbm3TxSViqJNq1tJlPpcby8dG+0wxBCiNOay+vze5EXamlxTVFwenxRSaasJo23Z41m6tNLfyxpHlxCZTOrvHDd8BPe66UoCkO7JPPctcOpqHOT8+Osjs2k0T7RRlZaHM8t3MNzP1Ti0sx+G5seze3VcXt1fvHyaubPGMkZTSRURVUuv7eHkgyrisLBSgcDMhKbfaxomU83HfKb9Ibyd/IZBgt3lrTpZAoalvt9euc4nlu4h/krGvq2HV+oxW7R0L0+Jpbv5Hd/m0VmqjR/Fm1bm0ymeqTFsre0Dl03UGVaVwghoiLeZkb3k3yEXFpc14m3NZxuquo9fL75EAWVDmqdHlJirfTrlMCk3mmYQtgPFIr0eBuf3z2e619exf6yehxuX8CcJcaiYVIVXr1pBMO7Nj3zE6rkWAtnxB77mot3lfD0t3tCrpbo8Pi48dU1LHxwYsBliIH2aoWSDBuG0aJKjiJ45X5KhkNofyePz6Cs1v/rtDWxVhMPnt+He8/pxYJtRby/Lp/SGhde3SApxsx5/dpz2bAMym+cT8LyXjD159EOWYgmtclkKt5mJtFupqDSIeVVhRAiSjJT/JfnDrW0eNeUGHYU1vDikly+3FKIqio4fhyNVoAYq4ZZVblxbDeuG9O1RXuUmtMuzsoXd49nZW45z76zhNXVGla7BcNo2Bvl9Rl0TrZz64QsLhzUMWINZ//59c4WV0v0+HTeWnWAuyf3bPTcercXs6bg8jZ+3VCSYVVRiLeayCuvp7zOjQEk2c10SYmRwc5WEmgrYaiDFifb38OsqVwwsCMXBNjrpd13H4f+8AcSfnY+ijm09gRCRFKbTKagYanftkPVJNjMxNlMsvFQCCEizGrSuGpEJvNX7j+mtHYopcVjLBp9O8ZzxQvL8XgNfMaxc0IGh/cD+Xh+UQ6vLNvLvBmjGJKZdMLx+3SD73YU8+H6fIqrnfh0g+QYC6MKt/LwpCFUDR5OjdOL3azRKckW8SVSe4pr2VlY4/e+YAoPuLw6ry7by+0TszhQXs+GA5VsyKtgw4FKckpqA543Q0mG691efvPxFiod7iOVBL0+gzibiZnjuzPtjEySYiwn8FsQqXH+f3+h/J0sJpWU2FPr7xA7ehSWzhlUfvgRydMbV9kUoq1oc8lUTkktryzdy8rcMpbtKcWkKXh1gy7JMcya0INLhmQQa21zYQshxCnppjO78+aqAxy/mSfY0uJur87CnSXHlNMOxOXVcXl1rp67kndmjWZwCxMqh9vH3MW5vLZ8L26v3mhPxkprX55dr3Iph7jr7Gw6JbVOg9xQvb58n989XKEUHqh2eBn8yNck2S0M6ZLE0MwkLh3amf6dElieU8pdb29oVNo+lGTYAAqrncCxJdEdHh//+mYX//x6F/ed24vZZ/WISrXGU8ElQzJYt6+i0fs0lL+TQkN/s1NN2n33kX/nXSRecjGqTRr2irapzWQlh6oc3PnWBrYUVOHTjSMnmMOjofvL6/nb59v5y2fbuGVcd355bu+TbkpbCCFONpkpMZzXrz3fbC9qlBA1V1pcVcB71Pd5sBweH9e/vIrFv5oU8qxHaa2Lq+eu5EB5fcBqYU6TFbw67609wOebD/LWLaOjUmBhc0GV399NKIUHDAxuPrM795/XuOrghF7pWE2a3z5hwSbDTf3pDr8fnlywm4OVDv58UX9JqFpgysCO/P4/W/zeF+zfaUBG4inZ48s+cCD2QQOpePMtUmfc3PwThIiCNpFM7Smu5coXllPt9OJr4pv7cLWbl5fuY09xHc9eO0yW/wkhRJg9MW0wV81ZwfZDNc2WMz7MbtYwawrVTj+bdmh+P5Dbp/PumjxmT8gKOs4ap4crnl9OfoUjqATOpzfM7Fw1dwWf3HFmxJf51To9/uMKofCAbhDwvKmpCved05O/f7HD776s5pLhYDk8Pt5fm0/nZDuzzgr+7yUa2Mz+l9Me1tzfKcaiMfusHuEMMarS7rmH/b+4gaRpV7KnrqGU/MFKBx6fTrt4K2dmtWNSn3S5HhRRE/VkqrjGyVVzV1BZ77/ztz8Oj49Fu0r47cebefTyQWGNTwghTndWk8Y7s8Zw+5vrWZFTiqOJJXuaqmDWFM7pl84324r8PiaY/UBOj85LS/cyc3yPoFchPPTBDxyscoY8E1bv8nH9y6tZ+tDZEb0gswdodBtK4QFNgThb4FP5daO7svVgNZ9sPHhCVfmaS34dHh//+noXV4/sQrxNigWE6u7JPflicyHFNc4mZwOPZ9UUhndN5py+7cMXXJSZe2Sxcvyl3PWPb9iv2PF4dY7OOd9bk4fVpHLTuO5cP7qr7OETERf1ZOrRL3ZQFSCRaurL2+Hx8cnGg1w9skuL19ULIYRoXkGlg9eX7WPtvnKcHh2FI31n0RQwaQomTcPj0/n5oI7MGNeDeSv24fEzixXKfqB6l5flOWWM69mu2RiLa5ws2FGMO8DMWVPnEwOodnpYvKuESRHcd9IjLZatB6saXTyHUnjAZtbonBy46q2iKPz90oHE2Uy8ufIALq8vpIt1CC75hYZqch+uy+fGM7uHdgBBUoyF92aP4dLnllHpcBOgqv0xrIpB9+oinps68pTd9uD0+BoGccz9cXh1oPEvps7to87t45nvGnpXvTNrND3SpDeViJzwNPUIUrWzod+Iv1HE6tUfU/7tiySOnkbnO98g47ZXiR82BcfuVUce4/L6eHFJbiRDFkKI00ZFnZsbXlnN2U8s5LXle6l2ejH4qRSFASiqgs+A0T1SWPmbyfxz2hD6dUpgc0EVflYshbQfyOMz2FXkv9rd8d5edYBAl5PBnE/qXD5eWJQT1LFayw1ju2H1U4L96MID9btWoHucGD4vjpy1VHz/yjGPNYDz+jU9K6GqCr+7sB9vzRzFz/p3wGJSsZu1I78vs6ZgN2tYtMa/wcPJb8q5txHTeyyqxYaimYjJHnVMAQRoWIo/d0kuhhFitiYA6JIaw//uHc+QzGRsJpVAbdesJhWLSWXqsExe6lxG2T13ozudkQ02Arw+nZteXcPyPaU/JlJNc3l1SmpdXPLsMvLK6yMQoRANojoz9eHafFQ/m1WDHbnUDfhmWxEVdW6ST7GSoEIIEU0HKx1c9txyyupcfvdxHHb4vmV7yrj77Q28fMMILCaVWn8NjghtP5Dbp7Mxr4KVuQkk2hv6DybYzcRatEaFDuav3O93P1coM2Eb8yopqnbSPiEyVcOGZiaRHm9lf1njC79gCg+YVYVpZ2QG3RNraJdknrtuOOV1br7eWkhJjQunVyc5xszAzolcPXdlo+eEkvwClNW6Kap20SFRKq+1RHq8jQ9vG8ue4hpeWbqPjzbko+sNy2e9uk6c1cRNZ3bnmlFdaBdnxdAHcvChX1Nw/wN0fvLfKKaoLzhqNY9/tZONeZU4g9ynCWAYUOvycs2LK1n04KRTdsZOtC3RTabW5/tdwx3Kl7dJVVi8u4SLh2SEI0QhhDjtVDk8TJuzgpIaV6O+UIE4PD7W7Cvnvvc28szVQ7EHuMAPZT+QAmw/VMO/vt5FlcNz5J/Hp5NwOLmymYi3mSitdft9jVDOJxaTysFKR8SSKUVRuOvsbH7/n60tKhChaQo3ndkt5OOmxFqYPrLLMbcV1zixaGqjC9dQkl9oaMRa5fBIMnWCstPj+ftlA/nLJQOocXqod/uIt5mIs5qOGUhQVJVOf/srebfeSuEjf6HDn/90SlRUdLh9zFuxv0UNrXUDyuvcLNpdwqTep165eNH2RDWZqqg/8UpGXt2gvM7/SVQIIUTonv5ud0OTWz+JVFMXMU6Pznc7ilm6p5Tu7WLZWVjTaD9sKPuBYqwavzy3FxcM7HjM7W6vTrXzp+SqtMbFspwy/OV9oSYDDnfLizS0xOXDOvP9V6tZ4NRwacEXbrCZVR67fBBdU1unHLaqKH73LoeS/EJDqXaZDGg9mqqQFGMhKfC2OBSLhYynnmL/L35B6XPPkXbHHZELMEw+3XSQQDlhMHv46tw+5izKkWRKRERUkyk9wIhnKF/ehmGEvJlWCCGEfy6vj7dX5+H2s7QvmIsYh9vHI59uw6z5vzgPpREpKJzdt/HFkMWk0i7OSrs4K9BwHjhcFON4ISUDRtOV8cKh7KWXuHvRh1iu/iNf5lQ1W3FPAaxmlb9ePICLWnFFRqLdjNfP3zyU5BfA4zVk2X0UaHFxdJkzh31XX4MpPZ3kK6+MdkgnZO6S3CPtcFU8eJ0AACAASURBVI4WyrLdDQcqOVTloGNidJpyi9NHVAtQJNr9j8Id/eXdHJOmBnwdIYQQofnf5kL8TfGEUoggp6SW60Z3oWOApV4JIy8j+ewZVK14l/ynryX/+RupWf8Z9p4/LcUzawrXjMzEamp+NkRRFDIDVLQL5Xzi9ul0SWliCqCVlTz3HFUffUyPefP41w2j+ccVg+jVPg67WUM7bljealKxmlQm9k7jnVljuOKMzFaNxaypjO6R0uj2UIphAHRPiz2S5IrIMqWlkfniXEqeeoqa77/3+5hDVQ62HqxiS0EVeeX1bbZYSKACEqEu291XKoUoRPhFdWbqvP7t2Vta12jTcCgjl17dYGxWaqRDF0KIU9IH6/Kp8zMiHNJFjKbi0+HOSdn89fPtLdsPpCrcMLZb0HHPPKs7f/9iR6PR7GDPJ4quM85cR1x9NcQ0X4r9eOV1bt5dc4C3Vh2gvM6NVzewWzSGZCYx66wejOmRemQvi2EYlDz5JLXffkvX+fMwtWs43kWDO3HR4E5sKajirVUH2FtaR73bS4LdzNAuSVw7qmtY93PNnpDFxrzKRn//YIphAMRaNG4LocmyaH3W7t3JfOYZ8m69DdPzz2EfMgSH28enPxzkhYU5FFQ6MP9YJtCn6yTaLcwc350rzshsUwPTgVochLpstyZAY2whWlNUk6nrRnXlhUX+S5sH++U9slsKnZJkClcIIVpDaa3L7+2hXMS4vDrldW7untyTpTmlfL+jGGcTjX6PZzOrPHHl4Cb7Jx3vkqGd+evn2/3eF8z5xGbRuNq9l5wpFxJ/9tmk/OJ6bP36NXvcijo3v/vPFr7ZXoSqcMzP6fLqLNpZwuq95STazfz2wr5cOLAjxY89Tt3KlXSZNw9TcnKj1xyQkcjfLxsY9M/eWsZltyPWavKbTDeX/AKgwAUDO4QpOhEs++DBdHr0/5F3511s/sNT/H5pQ/PswwMNRw9gOzxOnvh6F499tZNfntuLWWf1aBMFLCwm1W91zlD38EV62a44PUX1XZaeYGNcVju+31Xsd+Nwc1/eMRaN2RN6hDFCIYQ4vbTKXlYaVg0oisJT04dy77sb+W57cdD7gf5+yUB+PqhTSHHHWU1cP7orb6zcj8NP4tbU+cSsKvTskMDkO36F74FZVL7/AXm33Y4lM5PkX1xP/OTJKFrjnzm/op4rX1hBaW3g8vEGDRex9W4fD7y/iY3vf8G129fQ9dVX0JLaVsN5VVV4cvpQbnptdUjJLzQkwP+8cnBQyzJF+MVNmMBXV9/PM9/sb7awyeHP5b8X7Kag0sGfL+of9YSqY6KNfX5aBoSyh8/tjeyyXXH6iuqeKYDfXNg3YAndplg1haGZSZyZFfpyDCGEEP4ltcJeVoumkhTTUITArKk8c/VQ/nLJALq1iyHGojWq9nZ4P9CEXg37gS4b3rlFsT/0sz4M65KMzRz8qc2kKqTEWXjtppEoioIpOZl2s2aSveAbkq+5mvJXXiXn3PMoe/kVfFVVR55XWe9m2gsrKKp2NtmH62hOj858Vxrf3/6XNpdIHTYmK5V/Xjk4pN+hzazy2yl9+dmAjs0/WETEF5sP8myhNaQKkQ6Pj/fX5jNncWSbV/tzy/juxFhOrKF1/04JIc1uC9FSUZ//zE6P45UbR3Dza2v8Vm7xx6oYZFYX8ew5/aUhmxBCtKKfDejI5oLqRrNIoexlVdWGJWOHKYrCFcM7c/mwDDblV/Hmyv3sK6vD4fGRaDMztEsy143uesK9iUyayis3jeDutzewZHdps+cUu1mjY5KNd2aNJuW4CnSK2UzClCkkTJmCY/NmyufNZ8+555Fw4RRSrr+ev6ytpqTWFXI1WZdq4h/f7eX8YV3JaKNL1C8c1InUOCu/fHcjlQ4PDrfPb6XEWIuGzaLxj8sGcU6/9hGPU/jn9en85uMtAWcXm2pv4PD4+L9vdnP1yK5R3UN1ost2Yy0at8r+PREhUU+mAEb3SOW92WO46bU11Dvc1Hn9n53MqoKqKkzqncYfjCpKZ88i9o35mNOlj4AQQrSGy4d35h9f7vB7X7B7WbPS4ujdIb7R8xVFYUhmEkMywzcrYzVpvHDdcL7bUcycxblsyqtEN4wjs0eqAjazRnq8ldsmZnHxkAxszayOsA8cSMbjj+EpKqbinbfZduNMPh1zDx7F//OauliFhgIU85bv4+EpfVv3h29Fo3uksuzXZ7Myt5y5i3NYvKv0SN8f3TAY2T2F2ROymNAzTQY125jvdhTj8flPpIJpb6AqCh+szWPG+Ohto4izmph2Ribvrclr1Egaml62qygQbzNxdh+5NhSR0SaSKWjYcLvy4cl89NCjvBXXl61uKxaTiqI0fHEbBkw7I5Mbx3ajW7tY4AxK62vIm3ELXefPa7NLJoQQ4mSSaDczZWBHPtlY4HfWJZi9rNEeEVYUhcl92zO5b3sOlNXz1dZCimuceH0GqXEWxma3Y2hmUsj7Qszt00m/5x4+GzwF9eud4Od6NZiLVbfP4M1VB7j/vN5YTFFfbR+QoiiMyUplTFYqhmFQ7/ahGwZxVlPU99SIwF5YlEOdq+U9mhweH3OX5HLzuO5R/Tv/7sJ+bM6vYuuh6oDV/fyJsWi8OXM0Jq3tfrbEqaXNJFMARlUlgxZ+zGULbqdCs1Fc7cLp9ZFgM9M52d5o9DB19mx81TUcmDWbLq+8ghbXOp3ghRDidHb/eb1YsL2IGqc3pOeZNYWstDh+NqDtVHTrkhrDzLNad4T97fUHceqNLzJDaSgKsCK3jAm90lo1tnBRFIVYa5u6ZBB++HSDjXmVfu8Lpb1BlcNDfoWDzCgWcLCYVN64ZRQzXlvDpvzmG1qbNYU4q4k3bxlNVlpchKIUog0UoDha1UcfEz95MlpiIu3irPTrlMCwLslkp8f5XYahKArpDz6ArXcv8u+8E93lv6SvEEKI4HVOjmH+jFHEWDSCHZe2mFQ6JdmZP2PkkT42p6qyOrff20O5WNUNg9IaOWeJ1lXj9KAFWHYZSnsDk6pS5Yh+j6ZYq4k3bhnFH6b2o2tKDDafu9F3UqxFI85q4sax3fjqvrPo1ykhKrGK01ebOeMZuk7Fe++SdNW0kJ6nKAod/vQntOQkCn55P4bH/4ffpxtU1ruprHfjC3XHsBBCnGaGZCbxnzvOpEOijVhr4IsvTWmo5nZG12Q+vWvckSp+p7JA55BQLlYNA9wB9rUI0VJNLcs7ur1BcK/VWlGdGJOmcvXILnw5tQOP7fqIa0d14Zy+6ZzVsx2XDOnE3y4dyLrfn8NvL+xHenz4mloLEUibmbOvX7kS1WbHPmRIyM9VNI2Mf/yDvDvv5OBvf0unRx9FUVV8usHCncW8sCiHdfsrMKkqBgY+XTbPCiFEc3q1j2fZQ2ezeHcJcxblsu5ABRatYS+rYTRUDZs6uBMzxnenT4fTZzQ41qr5HbUPpReXpipRrZYmTk3xVlPAZD+UHk0en05yGxsYqf32O0aN7s9Fl0a+obUQTWkzyVTFu++RPP2qFm92VCwWOj/5JAdmzqTor39j8xUzeejDzbg8viPd3I8eBVyZW87m/CrsFo1/Thty0qxbF0KISFJVhYm905nYO53SWheFVU6cHh/xP+5lPR330YzNasfHGwoaXbSG1FDUpzO0ixROEq1LVRXGZbdj8e7SxveF0N6gY7yFjifYqqC11Xz7LR1+83C0wxCikTaxzM9TXEzdihUkTJ16Qq+j2u1kPv887+51cOe8NZTXuY8kUv7UuX2U1rqZPX8t76/NO6FjCyHEqa5dnJUBGYmc0S2F3h3iT8tECmDGuO6YtcYDf6E0FB3VPYWOiW2zz5Q4uc2ekOW34S00tDdIPnsGVSveJf/pa8l//kZq1n+GvedP+/zs+Lh0+XscfOghHJs2YRjR3xrhzs/HW1SEfdiwaIciRCNt4kxY9dFHJJx/PlrciVdf+fZAHc9njMcVQhlNp0fn959sIS3eysTe0pdACCFEYH07JtA9NZbthTWN7gumF1eMRWP2WdJQVITH2KxUEu3mgE2rm2tvoFgszHjp77g+/S8FDzyIlpBA8rXXkjDlAlRbdGarar/9lrhJE1G05vcjChFpUU+mDJ+PyvfeJ+Opp074tdxenfvf3+S3wVtzTRSdHp373t3I2t+dG7ASjhBCCAHw+6n9uPm1NTg9oTUUtWgKvdvHMzYrNdwhitOUoij831VDuPHV1X7fn02xmVX+fukA4tJSibv5JlJu+AW1S5ZQ8dZbFD/+OEmXX0bS9KuxdM4IU/T+1Sz4lpSbboroMYUIVtSX+dUtXYqWkoJ9QP8Tfq2vthai+9l4Wb36Y8q/fZHE0dPofOcbZNz2KvHDpuDYveqYx7m9Ot/tKD7hOIQQQpzaxma1488X9cdmDv40atEayse/dvNIKXwkwmp0j1T+eeXgkN6fNrPKL8/txSVDOx+5TdE04idOpMvcuXR7520Mr499V1xB3m23U7t0GYYe/oqU3vJynNu3Ezu2+ZYDQkRD1JOpw4UnWsMLi3Ia7ZE63EQx5dzbiOk9FtViQ9FMxGSPOmazJTTsoXphUU6rxCKEEOLUdtWILvxrWsMFq72Zi1ZNVRjYOZH/3jVOqviJiLhwUCdeu2kkHRJsxAbYQwUN1SmT7GaeuHIws5pYfmrp2pX2v36I7O8altwVP/EEuVMupHzePHw1jZe8tpba7xcSO3Zs1JYYCtGciC7z8+kG9W4vdrOGSVPxHDpE/bp1ZDzx+Am/dlW9h11FjT/MoTRRBNiYV4nD7cPexBePEEIIATBlYCfOzErj/XV5zF2cS53bi4qCQUOfHrdXZ2xWOzbmVfDHn/cjwSaJlIic0T1SWfHw2SzPKWPu4hyW7C5FUxUUFDy6zpDMJG6dkMXkPumYgmy2rcbEkDxtGklXXolj/Xoq3nyLkmeeJeGCC0i+5hpsvXu1KNb8inrmLd/PfzYWUO30YBgN+wuHV5Zyy+jJdG7+JYSIirAnU1UOD++vzeOlpXspqnKiqQo+wyA5xsJlSiGXTrkYNSbmhI9TUe/GrKl4fMfOTIXSRBEalmFUOtzYLVJlSQghRPMSY8zcMr4HN5/ZnR8KqiiqduL26iTYzfTrmEBavJX5K/fz+Nc7mT9jVLTDFacZRVE4M7sdZ2a3Q9cNalxeDMMg3mY+oT3iiqIQM3w4McOH4ykupvK998mbORNL164kX3sN8ZMno5ibHzw4UFbPQx/+wPoDFeiGgcf303YNl1fnO1MHlm03k/GvRfz1kgGM7iH7DUXbErZlfm6vzm8/3szIvy3gn1/vorDKiQF4dQPDgPI6N/Oq4rnUNZhb31hHrct7QsfTDQN/XwmhdvxueK0TCkUIIcRpSFUVhmQmcX7/Dkwd3IkJvdJIi7cCcNUZmewvq2dFTlmUoxSnM/XHZtFJMZZWLbZlTk8n7c47yP52AcnXXE3FG2+y55xzKXnuObwlJQGf90N+JRc+vYRVe8twefVjEqnDdFXD4dHZU1zLja+u5qP1+a0WtxCtISzJlMPtY/rcFXy4Ph+XV8fh8Z/IuDUzbt3g+x3FTH16CWW1rhYfMynGckxT3sOObqIYDI9Pl/XsQgghWpXF1LC5/7GvdrSJvj1ChINiNpNwwQV0fWM+mXPn4C0sIufCn1Nw/wPUr99wzHt/X2kd1764ihqnN+hBbKdH5zcfb+Z7KRYm2pBWT6Z8usGs+WvZerA66JKcLq9OfrmDa15chaOJJrtNSY4xk5HUeGleKE0UAbLT44g7TRtRCiGECJ+LBnfC4faxYLtcCIpTn613bzo+8meyF3yDfdBADj78a/ZedjmVH3yA7nBw9zsbqHP7X5VUt20hh16/lwP/uoL8Z66n6L0/4szfCjQkVHe+tR5ngIF6ISKt1bOGL7cUsm5/hd+muU31evLoBvvL6nht+V5um5gd8nEVReHWCVk88tm2Ro3qgmmiCBBr0bh1gjRSFEII0fpUVeGB83rz+Fc7ObtPuvQ0FKcFLSGBlBtuIPn666lbtpyKN99kzfPz2DliJrqfMf3q1R9TteoDUs+7A1v3YSiaCcfedTh2rzrSGxTg000HufKMzEj+KEL41erJ1POL9vjtuh3Mh8Pp1Xl56V5mn5XVoh4cFw3pxJ8/3eb3vuY6fkNDQnbBwA4hH1cIIYQIxuS+6Ty3cA//3VTApUOlPpk4fSiqStz4ccSNH8fT81bg3dZ4/+DhdjapU+4lpvfYI7fHZI8iJvun4i2HW9lIMiXaglZd5re7qIY9xbWNbg+l15PD7WPx7sCbFZsSYzHx2yl9sJtDL2tuM6v8cWo/rCYpiS6EECI8FEXhVz/rw7++2YXb44U938Ind8Abl8P8S+Gj2bD9U/CdWFEmIdqyz/dU4fNTNiyUdjYFlQ4OlNWHIzwhQtKqM1P/21KIx08RiFA+HHVuHx9vKGBi7/QWxXDdmG4crHLy6rJ9AQtfHM9mVrl9QraMcAghhAi70Z3t3Gb5Au/jM7HgAHfdsQ/Y8RloZhg5C0bfBvbk6AQqRBh4fTr1Aa7PQmlnY9ZUSmpddEk98fY6QpyIVp2ZKqx24ieXCrnXU3F1y6v6AfzqZ334zZQ+P3alD3zMGIuG3azx54v6c/c5PU/omEIIIUSzakvgxbOZXjOfGHdp40QKwF0LjgpY+m94/kyo2BfxMIUIF6/uv5UNhN7Oxt8AvhCR1qozU3qA2pZHfziCSai8+ol/OK4f041Lh3Xm4/X5vLAol7Ja15Hu3h6nizSrwu0XDuDiIRnESvU+IYQQ4easgpfPgaoCVN3T/ON9Lqg5BC9OhluXQkLH8McoRJjZzBqKooCfFgFHt7OJ7TOuydcxDKSVjWgTWjWLSIuzogDHfzxC+XAApMRaWiWeOKuJ68d047rRXTlY5aSy3g2A+s3/SNixmYxRU1vlOEIIIUSzPpoF1YcgmETqMEMHZyW8dRXcujh8sQkRQf07JfBDflWj249uZ6OoGrbuQ1FUE859G3Ee+OGYffaGYdC9XWwkwxbCr1Zd5je+Vxp2S+OZp1B6PcVYNH42oHUr6imKQkaSnf6dEunfKZHuo4bg3LihVY8hhBBCBFSZB7nfN8w2HaXbv2tIf7yGOvdPw5AvrXcz8bWjlv/pXijbDQXrIhWtEGF124QsYq3+VyoljLyM5LNnULXiXfKfvpb852+kZv1n2Hv+tO/erClcNSITWwsKjgnR2lp1ZmpEt2SSYyzUux2N7gu215MCXDAgvEsZrD174i0uxltRgSlZNvYKIYQIs9Uv+l3WBOAz4MlVbn4z3hr4+V4nLH8WrmzcbF6Ik825/dpjUlXA/96o5trZqIrCjWO7hyk6IULTqslUQ+PcHvz9i+04PI33PTX34TCrCtNHdgn7SIOiadgGDcSxaRPxEyeG9VhCCCFOc4YB614Fn9vv3Q+OtfDYMhe3j7CQZAuwNd/QYedn4KoFa1wYgxUi/EyaysMX9OHPn24LuvLyYXazypSBHaWKn2gzWnWZH8AVwzPpkGgPubO7AsTbzdw6Iau1Q/LLPmQIjo0bI3IsIYQQpzF3HXgC98M5o5PGxG4mnljeTCVb1Qw1ha0cnBDRMX1kF24Y2zXk3qDZ6fE8evmgMEUlROhaPZmyWzTemTWadrEWzEEmVKoC8TYT78waTVp8E8scWlHM0KE4NkgyJYQQIszcdaA2vRDkkUlWnl7tpqSuiWq2igLumlYOTojo+fUFfbnv3J5YTSoWU+BLUk1p6Ak6rEsSZbUuKur8z/IKEQ2tnkwBtE+w8cU94+nbKQG7WaOpnCrWopGZEsNnd42nV/v4cITjl33wYJybN2N4pcu8EEKIMLLGNRSRaMKAdI2f9zLx6NImLhINA6wJrRycENE166wsFj04iVnje5BgMxFn/elfvNWEzaxy+fDOfHLHOD66/Uymj+zCzHlrcbhDWx4oRLiErcFSapyV/945jk15lby4JJdvthVh1lQOtxZwuzwMN9Vy1w3nMKZHakPPgQjSEhMxdeyIa9cubP36RfTYQgghTiPmGLDENZQ4b8KfJ9oYNqeW+8cEWKGheyG+davdCtEWdEi08cD5vbnnnJ5szKukrNaNTzdIijEzqHMi8baf+knddXY2e0vr+OV7G3n2mmGoIW4rEaK1hb1b7eDMJJ65ZhjVTg955fXUOr3EWk2kHNyH4w8Pk/XX6eEOISD7kMHUb9woyZQQQojwURQYOQuWP9VQlS+A7BSVq/qbeWq1m4Hpxy0cUTTofylYpK+OOHWZNZUR3VKafIyiKDx6+UCue2kVj3+9k4d+1idC0QnhX1iW+fmTYDPTv1Mio3qkMiAjkY7DB+KrqcF94ECkQmhE9k0JIYSIiBEzApZGP9ofJliP6Tl1hMkCY+4MQ2BCnHysJo0515/BF5sP8d7avGiHI05zEUumjqeoKnFnnUXtwkXRCkEq+gkhhIiM+A7Q50Iw2Y65ed+98ZzT46dFIpmJKs7fJbDwxqNmoDQLdBwMHQZEKloh2ryUWAsv3zCCx77cwfKc0miHI05jUUumAOImTKB2UfSSKUuPHviqqvCWyodQCCFEmF38LKT0aEiOgqWaILYdTH87fHEJcZLKTo/jqelDufvtDeSU1EY7HHGaimoyFXvmWBwbNqDX1UXl+IqqYh88WGanhBBChJ8lBm7+smGWyRxEw1GTHZK7w8zvIabpfSRCnK7GZrfjwfN7M+O1NVIyXURF2AtQNEWLi8M2aBB1K1cSP3lyVGIoH3gGHy/fj9u3E5+ukxJrYXzPNPp2lPKzQgghWpktEW78An54F5Y9CdUF4HECh/tLKbhVGw4tnsTJ98Ow66XohBDNuGpEF3JL65g9fx3zbxmJ1RRaI2AhTkRUkymAuIkTqF24MKLJlK4bLNpVwguLcthQmIHh9eL5fg8AZk3hX9/soltqLLdNzOKCAR2bbCQnhBBChMRkaUiShl0PBetg26dQewh0HeI7UJo+nos+VVgxYjJmTc4/QgTjofP7cPub63n4o83888rBEW+5I05f0U+mJkzgwCuvYhhGRN74To+PO99az/KcMuoPN3w7qjO9x2fg8RnsKKzh4Y828+KSXN6YMYqkmBDWuAshhBDByBje8O8onYAuy5exeFcJk/u2j05cQpxkVFXh/64awrQ5K3j2+z3ceXbPaIckThNRH/Kydu+OYrfh2r497Mfy+nRufHU1S3eX/pRINaHe7WNnYQ0XP7uMGqcn7PEJIYQQAJcN68xH6wuiHYYQJxW7ReOlG87grVUH+OyHg9EOR5wmop5MQeSq+v3ls21syqvC6dWbf/CPPD6DQ1VObp2/LoyRCSGEED+ZOqgTi3eXUFUvA3lChKJ9go2XbhjBHz/ZyvoDFdEOR5wG2k4yFeZ+U1UOD++sycPh8T8jVbdtIYdev5cD/7qC/Geup+i9P+LM3wqA26uz7kAFu4pqwhqjEEIIAZAYY+asnml8KqPrQoSsX6cEHr9yELfOX0deeX20wxGnuDaRTMWMGIErJwdveXnYjvH+2jzUAHuyqld/TPm3L5I4ehqd73yDjNteJX7YFBy7Vx15jMer8/LSvWGLTwghhDjaZcMy+Gh9frTDEOKkdHaf9tw6IYsZr6+h2s9Wjbzyetbtr2BFThlbD1bhDDDYLkRzol6AAkC1WIgZPYq6JUtIvPjisBzjpaV7/c5K6a46Kpe+SeqUe4npPfbI7THZo4jJHnXk/30GfLKxgD9O7UeMpU382oQQQpzCzuqVxkMf/kBuSS090uKiHY4QJ52bzuzG3tI67nhzPa/eOAKvbvDF5kM8vzCHvPJ6zD9WazYMA92AaWdkctOZ3eiaKu0IRPDaTFZweN9UOJIpj0+nqNrp9z5XwQ4Mr5uYXmOafR1NVSiocNCzfXxrhyiEEEIcw6ypXDQ4g5eW5BJjMbG/rA6HRycpxsyo7ilcMjSDeJs52mEK0WYpisIfp/ZjxutrmTV/Lav3VmAYBnU/FiE7fg/9myv38/bqA/xsQAcev2KwtMYRQWkzyVTs+LNYNPcdXvvvFgpr3fgMg7Q4K2f3bc/47HaoasvLptc6vZhVFbevceEJn6MaNSYBRW2+wZuqKFQ7vS2OQwghhAiGYRj8b0shC3cVs7ekDlVpWCFx2Hc7ivnbF9uZOqgTt0/Kpns7GUkXwh+TpjJ1UEce/OAHjGYe69EN0A2+2lpIQYWDN2eOkgbAollRT6bcXp2PN+Tz/MIcCodej2v5/mPe7B+syyfGYuKW8d25ZlSXFo3C2S0aXt1/BT/NnoBeX42h+5pNqAwDYizyoRJCCBE+Xp/Ogx/8wFdbC4+08fAddxV4+PaP1ufz+eZDzLl+OON7pkU6VCHavLX7yvndJ1uaTaSO5vTobDlYxb3vbOT564Y3/wRxWotqMlXt9HDDy6vZUVjTsJ9JbZwo1bl91Ll9/N+CXcxfuZ/3Zo+hU5I9pOPYzBoxFhO1rsazStaMPigmM/W7VhDbZ1yTr+P26nRIsIV0bCGEECJYhmFwz7sb+W57ccDqs0fzGQ2J1cx5a3ntppGM7pEagSiFOHn84ZOtOD3+B9Trti2kes1/8JTlo1rsmNN7kDh2GrbO/XF6dBbuLGFLQRUDMhIjHLU4mURtMajD7ePKF1aw9WB1UCcMp0fnUKWTi55ZSkmNK+TjTR+ZiVlrvFRQtcaSNO5ayr95gfpdK9A9TgyfF0fOWiq+f+WYx47JSiE51hLysYUQQohgzF+xP+hE6mhOj86M19dSWe8OU2RCnHx2FFaTW1rr975gKjm7vTovLsmNVLjiJBW1malff/QD+0rr/O5jCsRnGFTWe7jp1dV8dvf4kI53w5huzF+xH/xM9CaMvAw1NpmqFe9S+tkTKBY71vbZJIy56shjYi0as8/KCumYQgghRLB03eCZ7/c02Q8x0Cg6gE/XeXdNHrMnyLlK//9HjAAAIABJREFUCICXl+zFc/waWUKp5Gzw5ZZCquo9JMZIsRfhX1SSqdJaF19uKcTlDX3a1asb5JTUsSmvksGZSUEfMzMlhpHdU1iZW+b3gxXXfxJx/Sf5fa6iQEqshTFZsnxCCCFEeCzLKaXOz3J0aBhFr1r1Aann3YGt+zAUzYRj7zocu1cdSaacHp2Xlu5l5vgeJ1S0SYhTxcJdJfj0xtd8oVRyNmsq6/MqmNQ7PRwhilNAVJKpt1YdCHhfMCcMl9fHS0tyefqaYSEd98npQ5ny5BJKalz4jOC3IsZaTLx+80iUAE1/hRBCiBP1ytJ9R0o2Hy3YUXSAereXlXvLGJvVrtHr5NXksbF4I9XuasyqmVR7KmM7jcVuCm0fshAni0CDE6FUcjYMg6r6xk1/hTgsKsnUa8v3+Z2VCvaEoRvw9bYiapyekKr7pcRa+Oj2sUybs4LiGhfuADNjh2kqxFnNvDVzlDRMFEIIEVa5Jf73doQyiq7rsL+snrE/rvTz6T6WFizllS2vsLVsK5qi4dW9qIqKpmrohs7FWRdzbd9r6ZbYrRV/GiGiTw0wCB5KJWdFUdBkplc0IeIFKNxePeAG2VBOGCZNobDKfyPepnRKsvPFPeO5aWw34m0mYq2NP0QxFg2bWeXK4Zl8ee94+neSKi5CCCHCK9BeqVBG0b26fmQ0vspVxdWfX82vFv+K9cXrcflc1HvrcetunD4ndZ46HF4HH+z6gCs+vYKXfngJI4RVG0K0dQl2/3MGR1dyDkaqFB8TTYj4zFSdy4tJU/3OCkWqgW6CzczDU/py/3m9+WprIR+uz6do/yF8LjdpWZlMGdCRS4ZmEGuNehsuIYQQpwl7gD6GoYyim1SVOKuJKlcVV312FcX1xXj0ppcoeQ0vXp+XuZvnUu2u5pdn/LLFP4MQbcmlQzN4ccneRtecR1dyVlQNW/ehKKoJ576NOA/8QPKkm495/PBuyZEMW5xkIp4txFg1vAEq+IXaQDfedmLhW0wqUwd3YurgTpS/+SbunBw63HLlCb2mEEII0RK92sdzoKy+Uc3ZUPohqgp0bxfLbQtuCyqROprD6+DtHW+TlZTFxdkXt+AnEKJtuX50N15astfvfcFUcjZrCteO6oLV1Pwgvzh9RTyZspo0EuxmKv1s5gupga5Pp0Ni6zbQleUNQgghomXGuO4s21NK/XFFKEIZRU8wqtEqPiKnco/fRKpiSQWlX5XiLnaj2f4/e/cdHlWVPnD8e++dnp6QQkiAJBBKpPfQFTvYQUVRRAWxoru6u7qiq+vqb2Vdu4BdcREVewcVAQm9NylJgBAICSF1+r3390ckEDJJZoD083kenl1n7sycIeTe+57zvu9RCO0XSuw1sShBFTeLTtXJ8+ufZ1zKOGSp0baiFISzIi7MwuDkSJbvPuqz8VhtnZyhIgtq0pAO9TlEoQVolDPlTUM6YDZU/2h/N9CVJTivawyhATSfqJPo1CcIgiA0okFJkYTXsJdN6MCriDj3VoozFpDz0g3kvDaZ0vVfY+18osbYapSZ1l3l3Y2v4PA4qr1HwXcFHP74MHET4uj+aneSH03GfdRN9qxstJPSoOweOytzV579LygIjeCpK3sQZAl8ZclskLnnvE4kRNjqYVRCS9IowdSNgzr42Dq3gj8XDLNB4fYRyQ0zWEEQBEFoAJIkMeO8VKxG3zd+wWmjaXvz87R/YCGJd88jZvzjWBK6VT5vVGRGXTyaFSYF/ZT5QdWhcuTzI8TfGE9IzxAkg4Qp2kTinYm4C9wUryiuPNbutfPW1rcQhJYgIcLG/NsHE2ox4G9TPpMiI0lwrthbSvBDo3RYiAm1MKZbLD/tyPPZIr22ZVdF12gfEUSfADbsFQRBEITmYHz/BNbvP8YXG3Nr7O7ni9Wo8N6tg8gt345JMeHWqnbNte+2o3k0QvuFVnlcsSiE9AyhbFsZESNOFNnvKNxxZl9EEJqQtPgwvr5nOPd+uIGdh0rwajpeH5v5BpkUFFnioYu6Em4zMvntNSyYNoSkNkGNMGqhuWi0hOhnr+lJQoQVo+J/ep0sQYik8uSy2XgO5p79QYmaKUEQBKERSZLEv67swbUDEmpcoTqZUZEINhuYd9tAeieGU+opRfeR+6GWqRiCDUg+rrmGMAPesqrdcR3e6mmCgtCctY+y8fldQ/n63uFcOyCRILOCJIEiSSiSRI92oTw7vhfrHj2fGwd3YGzPeB44P5Ub31hFbpH4fRBq1mi9v4PMBhZOT+eGN1aRmV9e5wyc2SATGWTiw6mjCPmmhH3XX0+7F17A1rfPWRmPJGqmBEEQhCZAliUev+wczu0ay50frMPp0ZAk8KgngqQgkwKSxPUDE7l1WBJtw6wAmBUzEtWvZ0qwgrfMi67q1QIqb7EXQ3DV2wGDLLYGEVqmTjHBPHVlD566sgdeVcOr6ZgNss/7wOsGtqfM5eXGN1bx0R1DaBNsboQRC01do54tw20mFk5P58PV+5m7LJMiuweHW60ypxZkVjAqMrekJzE5vSNhNiPcdBOmDh3IuftuYv/2N8LGjW207yAIgiAI9SHEYiDcZuKD2wbx5aZcDhQ6sLu9RAWZ6NM+got7xFVr2Rxji0HTq6fP2zrZkAwSJetKCBt4YiN61alSurmU2GtiqxwfaYmsny8lCE2IQZGpq+v5bcOTKXF4uOnN1cyfOpgw61lsfia0CI0+9WQxKkwemsTN6R1ZnVXIV5tyyStxoeo6bYJNnN89jnO7xqCcUjUYPHIk7d9+m5zp03FnZdLm7ruRZNHGVRAEQWgZ3votm8npHekQFcQ953b26zXdIrsRag7F7rVXeVyxKcRcEUPuvFxki0xw92A8xzzkvp+LMdJIePqJOmSLYmFC6oSz+l0EoTm7//xUSpxebn1nDe/dOhCbqdFvn4UmpMn8a5AkiUHJUQxKjvL7NZYuqXT8aAE5d92NKyuL+KefRracwd5TomRKEARBaAIOFTtYuiufp648J6DXSZLELWm38Pz656vVPUVfEo0SpHB4wWHcR9zIVpnQvqEkTktENp6YjNR0jas6X3VWvocgtASSJDFzbHceWriZae+v442b+4uNfIVKzX4px9CmDe3fexdJVth308148/NP741EzZQgCILQRLyXsY8r+7Q7rf0UL0u5rMbnIkdG0vmpzqS9nka3F7vRbnK7yg17AYyykdHtRxNuER1zBeFksizxzFU9CDYbmPHhRrxq9XRaoXVq9sEUgGw2Ez/rWYJHjiDr2mtx7tzZ2EMSBEEQhFqVubx8sGoff/lkM1PfW8sDCzby0s+7ycwvY8GaA9wytONpvW+wKZhZI2dhUQLL1FAkhWhrNI8OfvS0PlcQWjqDIvP8db0pc3n566db0Hy0VxdanyaT5nemJEki+q67MCclsf+WKbR96ilCzvW9V1WNRGt0QRAEoZ5l5pcxZ2kmX2w8iCxJ2N0nutmaFJkXFu8m1Gokt8hJh6jT299mRMIInkh/gpkrZuJUnXUeb5SNxNhieOeidwgzh9V5vCC0VmaDwpxJ/bjpzdU88fV2HhvXXXSEbuVaxMrUyUIvuYTE2a9x+PHHOfrW2+h+B0jiF0EQBEGoXz/tyOPSF5fzyboDOD1alUAKwP1Hq+bCcjdT3lnDrB9+D+A6VtXFyRfz+gWv0yu6F2bFjEGqPn9qNVixKBau6HQFH4/7mLiguNP6LEFoTWwmA29OHsDqrEL+u3h3Yw9HaGQtZmXqZNZevej44XwOTL8Td1YmcY8+imQy1Xh8mcvLQZeEXTdjtbsJsxrFLIMgCIJwVv3y+xHu+t96nB7/ai0cHpU3l2fh1TT+enG30/rM3jG9mXfJPPaV7GPe9nlkHMqgzF2GQTYQZY3ims7XcGnypdiMttN6f0ForcKsRt67dSAT5mQQajFw2/Dkxh6S0EhaZDAFYIyPp8MHH5D75z+z//apJLzwPEr4iYJar6rx884jzP51L5tzijFiA+0cvE8tpm2YlWkjk7midzuCzC32r0gQBEFoIIeKHdz1gf+B1HEOj8q7K7Lp2z6CC9JOf9WoQ2gHHhn8yGm/XhCE6toEm5l36yDGz84gxGLg2gHtG3tIQiNo0ZGCEhxEwisvc2TWf8i+9joSZr+GOSmJ5bsLuGf+BtxelfI/Uiy8SIACqs7+QjtPfbODJ7/ezoMXduXWYUmN+0UEQRCEZu29Ffvwqr7T9cq3L6Fkzed4juYgm6wYY5IJS5+AJSENAIdH4/nFu88omBIEoX7Eh1uZd9sgrpubQZDZwNie8Y09JKGBtehgCkBSFGL/8hCmpI7su3ESG2c8xczN9jpnB4/nsc/64XcOFNpFgaEgCIJQxa68Ut5dkc223BLKXV5sJoXU2BBuTu/IOe1ONHFwezXmrdqH20cr5ZLVn1G86hOiLrgLS1JfJMWAI2sdjt2rKoMpgMyCMnYeLqFrXGiDfDdBEPyX1CaId24ZyKQ3VxFkNjC6S0xjD0loQC0+mDouYsIENlrb8ujSo7gU//ftcHhUFqw5QNswC9NGptTjCAVBEITm4KcdeTy3aBd788vweDVOXnDacrCYrzbnkhBhY8aYzoztGc9PO/LQfDSR0FzlFC3/gKhLZmDrkl75uK3TIGydBlU51u3VeHdFNk9f1bPevpcgCKevW9tQ5t7Un9vfXcurN/RlUHJU5XOHi528vzKbX3flU+LwYlAkYkIsXD8wkYvOiRMbADdzrSaY0nWdR7d7agykakuzcHhUnlu0iwn9E4kIqrmRhSAIgtBy6brOf37cxZvLM3HUkN2g6eD0aOw5UsaDH29mZeZR4kItOE7p2gfgOrgT3evGljqkzs/WdNh5qPSMv4MgCPWnb/sIXry+D3d+sJ53bhmIIkv83/c7WZl5FJ2KSZHjMvPL2ZJTxCOfbWXiwPbcN6azqNNvplrNT211ViGF5W6fz/mTZiFJ8OGa/Uwf1akhhy0IgiA0ES/9vIc3l2fVGEidyuFRWbjuICnRQfja21N1lCDbQpFk/2aly1zeQIYrCEIjGNqpDU9f1YMb3lyJy6Ph8tZ8vjhet/9uRjaLd+bx4dTBxIQEttm20Pha3D5TNZm7NNPnzODxNIvI86dj65KObLIgKQZsnQYRMXpK5XFOj8Yby7PEbteCIAit0Lp9hby2ZC8OT/XrSG0cHpXf80p97mSoWEPR7CXomn/vGSxmrQWhWQixGHG41VoDqZO5vBr7j9qZMDtDTJo0Q63mzLxsdwG+wqBA0iwcbpU9+WWkxoac/QEKgiAIDUbTdEqcHsrdKsEmAyEWA7Jcc5OhV5fsxVlDIFVXNz6PqiNLcGrZlLldVySDEfuuDIK6Dqt1vIosiWuPIDQDDrfK7e+txVND986aeDWd3CIHj36+lf9e27ueRifUh1YRTLm9Gl7N9+xAIGkWiixRZPec7eEJgiAIDeRAoZ13V2Qzf81+3F4Ngyzj1TTMBoWJA9tzU3oHEiKqbmB7pNRZ44Scv934NE2vyBc/iWwOInzYDRQumo0kK1iS+iDJBpzZG3Hu31wlO8KoSNyc3vFs/lUIglAPvtx00GfDGah74sWt6ny75RCPj0sjzOZ/szShcbWKYKo2J6dZ1BVQicbogiAIzVOJ08O98zeQsfcomq5Xzhp7VPWP//Xyzoos3s3IZlinNrxwfZ/KtLovNhz0ef4PpBvf8a01Tr3FCh14FXJQBMUZCyj4ehaSyYo5thOhQ66tclzHqCC6x4u26ILQlOm6zmtL9lZur3MyfydeZAk+WnuA20ckN+TQhTPQKoIpk0HGIMs+9/gIJM3Cq+mEi5kCQRCEZqWgzMWVr/5GXrHL53XgOLeqAzrL9hQw7qVlLJw+lMggE5kFdp+1D4GkievgM9UPIDhtNMFpo2t8rdWoMGNM5zo/QxCExvV7Xil5Ja5qjwcy8eLwaLyXkS2CqWak1TSgGJHaxufM4slpFvZdGWgeJ7rqxbF3Lcd+eavKsTaTQqfo4IYZsCAIgnDGHG6Via+v5FCRs9ZA6mRur0bOMQc3vLEKp0elvIaC8EC78aXGhmAxBnbZtRqVP/aiaRvQ6wRBaHiHipwYlOp3m4FMvAAUlPnuPi00Ta1iZQpg6ogUVuw96nPp1Z80C4tR5rZhSbUWKAsth6rplDg8qLpOqMWIydBq5h0EoUX536p97C+04/XRibW2+gWPqpNdUM7Haw/UmJEQSJo4QFKbIO4/P5UZH27E5VV9tks/mdWoMGlIB/52cVe/vqsgCI3L6VF9rj4HOvHi8XPiR2gaWk0wNaBjBJFBJuxuh8/n60qz0HW4dkD7+hqe0ATouk7G3qPMXZrJst0FKIqERMVJrVvbUO4YmcKFaXEisBKEZkLXdeYuy8TpY18of+oXHB6V2Uv2Mq2tG6vuxSFVvWQGkiZuMcr0SgznwrQ4vrx7KLOemccSSyKywVCl3bpRkZAliT7tw7lzVCdGpEafhb8JQRAaQojFeGqfGSDwiReL0b+gS2gaWk0wJUkSs8b3YvLbq31eWGtjNSrcf35nIoJM9TQ6obGtzS7knvkbKHZ4KlcvVe+J6aVtuSX8deFmHv5sCzPHdmd8/8TGGqogCH5asfcopc7qKXqB1C8cPVpMUPZqiBoGp1w6AunG51F1bEaFIyVOEgtz+PPa+fzf19/x+dYjbDxQRJHdQ5DFQHKbICb0TyQxsmpHQUEQmr7UuGCf9ZWBTLwAdIkWv//NSasJpgAGJ0fx7NW9eHDhJr8DKoMsMXFQe6aOSKnn0QmN5cdth7n3ww11/ps4vlP5zC+2cvCYgxnnpzbE8ARBOE3fbTnkM7U7kPoFl8HMtrGTGK/pzF+1H88puXn+duPTNJ1nvt/JP7/dQT/XESZPuJ1OoTYmD006sy8pCEKTERNiIT05iiW78qs8HsjEi1X3Mu6bORzM/JywKy4nKD0dydCqbtebnVb30xnXO57IYBP3zN+Ay6tS7vK9CaPNpKDpOookcUmPuAYepdBQ1u0r9CuQOpnDozFnaSbRoWZuGNQh4M/cc6SUd1fsY8ehEspcXoLNBrrGhXBTekexKacgnEV5JU6fjwdSv6ADeaUu/npRVz5Zl4PHR3BWV5r48fc5HtitlCPZdMRA+ntreXliX5HSIwgtyNSRyazOLqw2kePvxIslyMr1816k/IfvyX/lFXIfeYSwS8cSdsXlWLqK+smmqNUFUwBDO7VhzSNjWPL7EWb/upcN+4swKjKSVFEf0y7cyvRRKVzWqx3L9xRw34cb+ebe4YRZRVv0lkTXdR5YUPMqZW3F6Q6PyhNfbWdcr3hCLf79u/h5Zx7PL9rNrrxSPJrGyfWlG/Yf45P1OXSKDmbGmFTGdI89G19REFq0w8VO9uaXUer0YjMpJEbaSGoTVOfrAq1f0HWdxEgbcyb14/b31gacKl7t/SQZu0dj+Z4CbnhjFfNvHyxqMQWhhRiSHEVSmyB2HS6ttpLt3zYIqZijIjFPnEjkxIm4srIo/vJLcu68Czk4mLDLLyd07FiMsTH1/VUEP7XKYApAkSXO6xbLed1iKXd5KXJ4UFWdMKuRUKuhcoPF87vH8tueAv66cDOv3tC38nGh+Vu//xj5ZdX3gwD/itNlSWLh2hxuGVZ7mo6u68z64Xfe+i27SqH5yVQdVI/G1twS7pm/gZvTO/CXi7qKf2+CcApN01m+p4A5v+5l7b5jmAwyuk7lZFhym2DuGJXCRWlxtAmx+HyPQOsXokPMAAzvHM2bNw/g9vfWomq6z9qIQDg9Gttyi/nbp5v5z4TeZ/RegiA0DZIk8d6UgVzy4jKOlrl9dhL1xWqUuaJPO24aUjXjxZyURMx99xF9zz3Y166l+IsvyLzsMqw9ehB2+WWEnHcesk3UWDUmSffVw/EP/fv319euXduAw2manB6VK19dwY2D259WWpfQNE17fy0/bs+r1sZUc5WT88rNRF0yo84brbgwCxl/PbfWoOe/i3Yxd2lmjYGUL1ajwq3DkvjzhV38fo2/JElap+t6/7P+xg1MnJ9an8PFTia+sZK8YmdlDaMvQWYFs0Hh3vM68ez3v/s8tmT1pxSv+pSoC++qtX4hyKQw96b+DO3UpvKxvBIn72fs472MbFRdx+Guuc15bSvcx5kNMksfGk1sqO/grzVpCecncW4SAPJLXUx8fSUHixw+azePUyQJo0FicnpHvydRNYeD0p9/pviLL3Bs3ETIeecRdvnl2AYOQJLPbJW72O7ho7UH+HDN/spg0GZSGJQcydThKfRICDuj92+uajs3tdqVqUBYjAovT+zD+NkZ9O8QSZc4UdfSEvyyM9/nfhCBFKcX2z1kFpSTUsNmziszjwYcSEFFS+Y3l2eRnhJF+kk3cYLQWh0otHP5K79R7PCg1jHTW+6qqIf99/e/+9xAE/yvXwixGElPiaryWGyohT9f2IX7xnTm551HmL9qP8v3FFSbgfZnhfu49zP21cvkiSAIjSM6xMw39w7nh22Hee3XvWTml6Hr4PGqyIDJZEDVdC46J47bhydzTjv/gxTZaiXs0ksJu/RSvPn5FH/9DXnPPINaXEzYuHGEXXE55uTkgMZ7rNzN419u4/tth5Glivrw48pcXr7ZfIjF248QH25h5rg0RoptGyqJYMpPKdHB/PXirtwzfz1f3DUMq0kUDDdnbq+GV/OdohNIcbpBkSiy17xT+au/7KkxkKprxtrhUXn5lz0imBJavXKXl+vmZlBkd9e50e3J7G4VkwRG1YNHqV7bWFf9gkGWGN0lGlXTfQZlRkXmgu6x/P3zrdUCqUDar7u8Gu9lZDNjTGcMiqidEoSWwmSQGdcrnnG94tl5uIQ12cc4+O2PWMLDaD8qnfO6xZ5xPb4hOpqoWyYTdctknL//TvEXX7L/5skY4uIq6qsuvQRDRESt73Gg0M6EORkUlLqq1Xkdp+kV9yV788uZ9v5aHr64GzeldzyjsbcUIpgKwPh+CSzfXcCT32znX1f2aOzhCGdAqyW9NZDidN3rxbnvAFqMFdlqrfLcoWIHq7IKfb7O3xnrdfuOcbDIQbtwq8/3EYTW4OO1Bygs9x1I1TUp4dU0LGYjmibVuaJ1KlXT+XJzLt9uPcxNQzowaXAHYk5JxSt1eX1OqASywg0V9V4FZW7iwkSqnyC0RF3jQukaF0reLwUYIoxE9U04659h6dIFy0MPEvPA/ZRnrKT4yy/Jf+EFbAMGEHb55QSPHoVsqrpn6rFyN+PnZHCkxOn3ZJXTo/H0dzsJsxq5vE+7s/49mhsRTAVAkiSeuvIcLn1xOd9tOcTFPdo29pCE02QxKsiS5DOoCqQ4XfV4cT73b3ZlbkUJC8PYoT2mxPaY2icyj/boPha/Apmx1nWdj9Yc4H6xp5XQSum6ztxlmVVSTo7zZ1JCk2RUZCJtBorsnhpnXX1+NvyxfYbK3KWZvLEsi+cm9Kpy7i91ejHIMh616gp0ICvcALIsUeL0iGBKEFo6xQBaYKn/gZIMBoKHDyN4+DDUsnJKf/yRY//7H4dnziTk4osIu/xyrL17I0kSM7/YytEyV0Cr/lCxSvWXTzczPDWayCBT3S9owUQwFaAQi5GXru/DlHfW0CMhjIQI0UGluRqcHMXyPQXVHg9oc71gG6M+nYesa3jz8nDvP4D7wH48+w+w59Ah3IbqAXcgM9ZuVSezoPzMvqggNGOrsgopsnuqPR7IpIQsSUwe2pGfd+azLbcYj6oHvEp1vHPf/R9txOFRueqPWWWzQfY5KRN4+3WwGET6uCC0dJIio3vrN5g6mRIcRPhVVxJ+1ZV4Dh6k+KuvOfTwI+iain7pFfyY1wGP6vt8WGcDHR0+XL2fO0d3arDv0xSJYOo09EoMZ9rIZO77cCMLpg4WOe7N1B0jU9iw/5jPTl/+FKdbjDJThnVEkSVAwRgfjzE+nqDBFTdy2rtrYUdetfcOdMa6zFn9RlIQWovVWYU4fPyOBjIp4fCorMk6xsLp6ew8XMIby7L4auNBXDXcQNTG6dF4+LMtdGwTRN/2EYRZjT4b2QTaft2jakQFt+7ZXUFoFWSl3lemamJs1442d0wjatpUnFu38tJHGeB2gVL93OPPyr/Tq/Hm8iymjUz5416odRLB1Gm6bVgyy/cc5fnFu0UHpmYqPSWKYIuhxhbLdRWn6zpcP6B9jc+H23wXlQY6Yx1uEzdYQsvgUT2UuEswKSaCjEHIUt0TUQVlLnyFPIFOShz7o66pa1wos8b3IjUmmOcW7cLpY6+oumZjnR6N//zwO+9N6EbZ9z8wvCSfX6yJaCeNJZAVbgkY3imaILO4JAtCSycpMrp6ZnvUnfEYJAlrjx58810BLo+92vOBrPw7vSob9h+jf8fIeh93UyXO3KdJliX+M74XY19aJtpXN1OyLPHM1T2ZPm8dTh/1GLWxGhXuGJlMVLC5xmN6JoTx9ebcarUegcxYW40KvVrpng5Cy+DwOvg+63ve2voW+0r2YVSMHN/fcEyHMdycdjNpUWk1vt5Uw8p/oJMSxpMmTXVd5+0V2T4DKX+bw6zZc4QVl/ydpAE9mHrepaxY56n2u+5v+3WrSWHqyMDaGAuC0Ew14srUqQrLfXcjDmTlX0Iiv9R1tofWrIhg6gxEh5iZNb4XD3y0iW/uHVbrjbXQNI3uEsPMsd154uvtfgdUVqPC5b3jufe8zrUed0Wfdvzzmx3VHg9kxlrTda6sh44/glDfdF3n7W1vM3vTbCQk7N6K2U+3euLi/UP2D/yy/xfahbTjuVHPkRxWPaCIC7NgMsi4vac/KYGuE7R6GXu+fBJLl65s7diT4rK4aocF1BxGlln24H8YfkUvEoAO2UvZlVdarYi7rhVuCWgTbGZQUuud1RWE1kRSZDRP00jf95yFLWJ0XcfdyCttjU0U+5yh4Z2juaJPOx78ZHPlbKvQvEwc1IEXr+tDsNlAUC37h1l9I9lNAAAgAElEQVSNMmaDzJ2jUnj6qh517lIeYjEyrle8zzzi0IFXEXHurRRnLCDnpRvIeW0ypeu/xtr5xCyQIklc0qPtGe9BIQgNTdd1Hl/xOK9tfA2H11EZSJ1K0zWcqpPMokyu//p6NuVvqnbMRefE4aso6eRJCfuuDDSPE1314ti7lmO/vFXl2CCzgSl/vYXEV14h5IIL2F4uVzaUOFkgs7EeXWJt7onmMHMn9T+tND2bSeHNm/vXeT4RBKEF0DQkyYPkLQa18QMqm8n3Oevklf+6yLJEqKV136eIlamz4E8XpDJ+dgZv/ZbNrcOSGns4wmm4IC2OdY9G8+3Gg7zwzmJyw+Iw/dFZy6tphFqMTB2RzPh+iYTVUAvlyx0jU/hm8yEcPk5Idc1Ymwwy00elBP5lBKGRvbjhRb7L+g6n6vTreB0du9fOtEXT+PDSD+kY1hFd11mVVciby7NqbNnrbxqdzWxgeJdYZDkOc6dOaMbfUX/ZU+39Aq3DKnGcWGVrH2VjwdQhTHx9JSVOT51thmWpIsh7b8pAOseG+PV5giA0Q7oOB1bDihdh1/dEHj85/PN5iO4Gw+6H7peBoeGzm/p1iGDxjrxq81WBrPy7vBpp7ULrcZRNnwimzgKjIvPS9X244pXfGJQUyTntRI1Lc2Q2KFzgPUTvw98R9PT7HCv34NU0wm0m2oZakE+jU02nmGBeuK439364IaC6LItR5j8TepIqbrKEZuZA6QHe3/4+LrV6Dv2xZcco+KEA9xE3ikUhtF8osdfEogRVBC92j51/ZDzBZbFP8ObyLOwulSnDkpjQL4H7FmzE7qNZTF2TEhajzO3Dk6r8/lqMMrJEtYAn0Dosi7HqMd3jQ/n2vuE8/e0OftyehyRR7ffebJDRgXO7RPPwJd1pHyW21xCEFit3I3xyC5TmgccO6FSeiXTgyDb4ekbFn/OfgAG3Nujwpo5I5rc9BdXOrf6WI1Q0z2lDTEjr3h9PBFNnSWKkjccvS+Oe+Rv46p5hBIuuTM1S2ZIlBI8eRXSI5aydHC5Ii+Pl6/tyz/wNeDWtxv0cAIyKhEGWeeG63lyQVr2mQxCaug92fIDmY7fqgu8KyP8un4TbEgjuHoznmIfc93PJnpVN0iNJyAYZHZ21hzdQemATD5zfn1GpMciyhK7rXJgWx3dbDwU0KWFUJLrEhnBzescqj8eFWbEYlWo3EIG2M28XYa32WHy4lZcm9qXI7uajtQf4YmNuZSfBMKuRS3u25foB7UWNrSC0dJlLYP71fwRRtXCXVfzvj4/AsWy44Mn6Hlml/h0iiAwyYXc7qj3nz8q/1aQwdYRoniPu+M+icb3iWb67gJlfbOW5Cb0bezjCaSj75Rfin/vPWX/fMd1jWfTACN5Zkc38VfsBsLtVdCpmdixeF7LNxrUDErllaBKJkWK2Wmh+nF4nn+3+DI9WtRZAdagc+fwI7W5tR0jPitVWU7SJxDsT2fXgLopXFBMxIgIAgywxpM9Ozu16aeXrJUni2Wt6Uu7ysmx3AQ5P3Xn8ZoNMcnQQ700ZhPmUzXAvSIvlkc+2VHtNIM1hgkwKNw7uUOPnh9tMTB2RwtQRIlVXEFqdQ5v9C6RO5nHAmjcgOAbS76m/sZ1EkiQeHdud+2rInqlt5d+oeki1yQwUzXNEA4qz7bHLurM5p5jPNuQ09lCEALkys9CcTizdu9fL+ydE2Pj7pd1ZP/N8nrm6J/efn8qtw5K4//xUHji0lN+ubsfMcWkikBKarQ1HNvjcO8q+247m0QjtVzWvXrEohPQMoWxbWeVjXt3D91nfV3sPgyIz+8Z+TB+VQrDZgK2GZjHHG8Vc2acdn9051GeNY6jFyNiebU+7OQxU1DsNTRFbYgiC4MMXd1ULpDo+X0rMs6WUu09kp7yx3s2od040ssFjh5+fhLL8hhopF6bFcf+YVCxG/0MCkyITH2bh8R+eo/Dtd1p9AzaxMnWW2UwGXrq+Dze8sYreiREktQlq7CEJfipbsoTgUaPqvauW2aAwrld8lccOLQ1H27geevWo188WhPpU5CpC97HFrlqmYgg2ICnVf7cMYQYc+6qmmJR6Sn2+vyxL3HteZ6aNTOb7rYeZszST7IJynB4Vk0EmNtTCLekdubpfAiF1dJe6fUQy32w5hOqjU4Q/dVhTRySfVh2lIAgtXN42KNjt8ylVhxdWuXl4eG1pvhKsextGPlQ/4/Nh2sgUwqxGHvtyG4DPbqd/jAyrSaFb21DevmUA1ik9OTD9TlyZe2k7cyaSydRgY25KxMpUPejWNpQZYzpz7/wN1fZGweuC0sNQdABcZb7fQGgUZb/8QvCokY3y2bb+/bGvXdsony0IZ4umaz5nKJVgBW+ZF91HvaC32IshuOq8Xl2znGaDwuW92/HtvcPZ/sRFZD59KTufvJhfHxzN5KFJdQZSAF3jQrl/TCpWo3+d+058tky/9hHcMlR0bhUEwYeMV0H1vRnug+kmZq1wUeSs5RzndcKq1xp8Y9/rBrZn6UOjmToimTCrkWCzQpBZwWZSCDYbMBtkRnWJ5s2bB/DJHUMItRgxxsfT4YMPUI8Wsv+22/EeO9agY24qxMpUPZk0uAPLdxfw7+938vdLu0HWUvjtBcj6FWQjSFLFL1t4exg6A3pcAyaxitVY1OJinNu3EzSk7v1l6oOtfz/ynn4aXdfFfjNCsxVmDvOZ5mfrZEMySJSsKyFs4Ilup6pTpXRzKbHXxFY5PsjYMOfCqSOSsbtV5izd61djC6tRpndiBK/f3N9niqAgCAI7vwLddyDUP15hVEcDs1a4+Oe5tTS58rrh0CZo17eeBulbbKiFP13QhfvO68zqrELyy1y4vRqhViO9E8OJDa0+ZiU4iISXX+LIf54j+7rrSHxtNubk1jXZJIKpeiJJEv++pid/+u9bOLdfhcVbAu4/8mI174kDCzPh+7/B93+FUX+rKDoUN9MNwulRcXk0gi0GypYtxzZwILKlcdp7Gtu2RbZYcGdlYU4WnXGE5ql3dO9qzScAFJtCzBUx5M7LRbbIVbr5GSONhKeHVx5rkAyMShzVIOOVJIn7z0+le3woz36/k4NFTlxetVrL9CCTgtEgc9uwJO4YmYJBEUkdgiD4oOt1Zh09MdrM0LfKuW9QLSlxkgyOwrM8OP8ZFJn0Tv7XhEqKQuxDD2JOTmLfpEm0m/Vso01ONwYRTNWj8NylzNUeQ7HXsXGl548ga8nTFcHV2P+KgKqe7M0v463lWXy6/iAur4oiS3g1nXjNzuS0MUx0eQlqpLb21v79sK9dK4KpZmTTgSI+3XCQ3CIHbq9GVJCJEanRXNwjrloHudYg2BTMRR0v4uvMr1FPmZmNviQaJUjh8ILDuI+4ka0yoX1DSZyWiHxS4bMiK9zY/cYGHfeFaXFcmBbH5pwi3lyWxZaDxZS5vFiMCgkRViand+TcrjEiiBIEwQ+1pymfE6MwNtXAM8vddIuu5ZzSDJs6hF9zDcbE9hz805+IvvtuIq67tu4XtQAimKovuRtgwSQUtY5A6mQeO2xeACFxMOqv9Te2VuhQsYO7/7eBrQeLUTUd7x9Tz9ofNRwHJRv/zZP5zz8XcduwJB44v0uDF5fb+vXHsW4dERMmNOjnCoFRNZ2F63OYvWQvh4qrr2T8sO0wj3y+lesHJjJ1eDIxPtIiWrKb0m7ih+wfUNXqaS6RIyOJHFl7G93O4Z1JDmucCYWeCeG8cH2fRvlsQRBaAEkCo+3E3lE1+McoC33nlPGnITU1otDBGnH2x9cAggYNpOMH8zhwx3RcmXuJ/ctfkJSaJxe9qkaxw4NX0wmzGqttht4ciGCqvnxxj8+2mHYPZN0XTJCp4kb9jfVu5m32sGTyHzUCHjssfw763gSh8ae+q3Aa9hwpY/zsFZQ4vT47dx1n/6Nm4s3l2ew5Us4rN/Rt0LoI24D+HJ0zp8E+Twicw61y+3trWbfvWI17HZX/sRHsuyuy+XhtDh/cNohz2oX5PLY5KXd5WbwjryKA9GiEWAz0SAijf4eIKnV+qRGpnJc4hm+zf0THdxF2TSyKhUcGP3K2hy4IgtBwUs6FnV+Dj83Lj+sUKXNtmpEXV7vpEVPD6lRc8+3ua+rQgY4fzidnxgwO3Hkn7f7zH5Tg4MrndV0nI/Moc3/NZNmeAgyyhCSB26uREh3M9FEpXNKjbbMJrETOQn04vBWO7vH51PG2mLXSgTVvnv1xtUJHSp1cOzeDIrun1kDqZA6Pyq+78vn751vreXRVmZKS0JxOPIcONejnCv7xqBqT3lzFmuxCvzaN9ag6xQ4P187NYHee71bfzcGeI2X87dMt9P/nYh7+dAuzfvid/y7exTPf7eTmt1Yz/N+/8H5GNmWuilrQcpeXfbsuJVxKxaL4vypnUSzMGjmLc9qcU0/fRBAEoQGk3wuGus99M0eaq+w5VUkxQb9bwFBb+/SmTwkLo/3cuRjbtmXf9RNx5xwEYP3+Ywz9v5+57d21LNmVj6rpuLwaTo+GpsPuI2U8+vlW+j25iPczshv1O/hLBFP1YeUZtsVUXbDmdVCrF3ILgXnm250U2z0+M5jLty/h0Lsz2P/cNeS8PIm8jx7DmVOxx4LDo/L5hoNsOlDUYGOVJAlbv77Y165rsM8U/Pd/3+1kW25xjftv1MTuUpn4xio8amCvawre/i2LsS8t46O1+3F4VMrdamWKrFvVsLtVco45+Nd3Oxk9awmb9hdx45ur6BAZyuKJ73NZymWYFBMmueZCa5vBRrg5nDnnz2FkYuNsTSAIgnDWJPSvKNc4RfaMEMYkn0gISwyTcf499ERm0nGSDANvr+9RNgjJaCTusccIHz+e7Ouv47tvMpj4+kpyi5zY3TVPSpa7K643//p2J099s6MBR3x6RJpffdj5zZm3xdQ0yN0IiQPqaZAtX4nTwzdbDlXe/FV5bvVnFK/6hKgL7sKS1BdJMeDIWodj9yosCWkAuLwqbyzL5KWJDdea1NqvH/Z1awkbN7bBPlOom8Ot8sGq/ThqaJ9dvn0JJWs+x3M0B9lkxRiTTFj6BCwJaeiA3e1l8fY8Lu7RtmEHfgZe+WUPL/+8x6+W4Q63itOtcuVrv3FV3wSevqoHkiTx6JBHmdZrGh/9/hH/2/k/vJoXRVLQ0XGrbrpFdmNKjymMTBiJQRaXI0EQWgBJgnEvwgfjweuo+/iT6LIZqf+Uim1zWghJkoi8aRK/hycwY0keLqXufQCPc3hU5q3cR0yImdtHNN3mXOLqVR/qKDz0ry2m1KhtMVuChWtzkH10RdRc5RQt/4CoS2Zg65Je+bit0yBsnQadOE6HH7fnUWR3E25rmF29bf36U7xwYYN8llDBq2rsPFxKsaNiJTjcZqRrXGiVermvNuXW2GDTn8C83KXy2q97m00w9dOOPF76ebdfgdRxOhXNp5buysfp0bCaKnLdY2wx3N3nbu7odQeHyg5R4inBJJuItEQSZY2qp28gCILQiJKGw2Uvwpf3+h1Q6YqFslwzOiMJrefhNYaZe+QaA6naJiQdHpVnf/ydq/q2Iyq4aaY+imCqXrTetphNycL1OT5rW1wHd6J73dhS694DwSBL/Lorn8t7t6uPIVZj6dYVT+4hvMeOYYhonp18mosjJU7eX7mP9zL24dW0ysBb03TMRoUpQzty3cD2tAk28/qyTJ8pCf4G5gC78krZd7ScDlFNf3Pu//t+Z42BVG0XPYAyl5evNucyoX9ildcZZAOJoYm+3lIQBKHl6TkBbJHw8S0VzShqmmg3WAAdaeA0jAnXsn/qHSDJhF54QYMOtz5tySkm55jvoNKfCUkZmL9mP3eP7tyAo/afCKbqgzEIXCW1HtKS22I2FcfsvmvOVEcJsi0USa67S4xX0yksD6wj2ZmQDAak3v34adEanKkVJ5Fwm4n+HSIabf+rlkbXdf67aBezl2Yigc8aqHK3ysu/7OGln/fwwPmpHCzyfREIJDA3KjLZR+1Vgild1yl3qxQ7PBgViXCrCZOhcUtZtx4s5kDh6V/07G6V2Uv2VgumBEEQWp1OY+DBvbDzK1j+POT/DgYTIIHmBaMVBt8F/W6GoDZYgPavz2X/7VMBWkxA9cayTFze05+QdHo13l6ezfSRnRq0y7K/xN1ZfUi9ELZ+WmPdFPjZFjO+dz0NsHXQaljZU6yhaPYSdE2tM6DSdR0/mwCescz8Mt76LYtPYsairHfCti0AyEh4NI0r+yRw67COdIoJaZgBtUC6rvOXhZv5atMh3HU0kji+MvP84t04a+jeF0hgrusVne6gYvXm8w05zP41k8PFToyKjKbrqJrOyC7RTBuRwoCOVVuON5S3lmf5/LsJZBXuULGTzTlF9EwIr/fxCoIgNGkGE5xzdcWf0jywF1Q0GLOGQ1ginHL9sHTr1uICqkU78nzeSwUyIen0qPx+uJTu8U0vCVIEU/VhyN0VTShO2WfqVDNHmnl/s4/VkxbSFrOxhVmNHCquvmmyuV1XJIMR+64MgroOq/U9DIpMmNX/YsnToWk6T369nfmr9+PVdLz6H8G1q+oN/MdrD/DZ+hyu7pfAE5ef0yRnZ5q6l3/ew1ebDvnV2vy42o4NJDCXJAgyKbz0025eWbIHWZIqUwe92onP+HnHETL2HiUq2MTsG/uRFt+we1RtyilC9TEREchFD2B7bokIpgRBEE4WElvxpw4tKaDSNL3G62ggE5KyLFFkb7hMoUCIYKo+xPeu6MSSv7PKw9kzqq4oHG+LWY0kwcCp9TnCVuGCtFiyCsqrpXHJ5iDCh91A4aLZSLKCJakPkmzAmb0R5/7NRIyeUnmsqukM7VR/RfK6rnPfhxtYvOMIzjpWSryajlfT+XT9QQrKXLx2Qz9kEVD5rdjh4eVf9vhO66ujDqgmgQTmLo/K/1bvZ+muglobO1R0/1OxFzoYPzuDtyYPYHBywzVqOL5f1KkCS4/VKHX6fh9BEAShbi0poKqplUAgE5JAg2UKBUrsM1VfLnsJDNaAX2bHTG6XmyBc1BucqRsHdaixFUjowKuIOPdWijMWkPPSDeS8NpnS9V9j7Vx11n1gUiRtwwL/Ofrrv4t2sXjHkYBXSpbuKuCZ73fWfbBQaeE6390dS1Z/RuFPrxM2eAIJd8+j3fS3Cel7CY7dq6oe6GO15uTA3L4rA83jRFe9OPau5dgvb1U5NtxmYumugoB+1na3ypR31vD74Ybb9Nds8H1BO/miVxdFkrCYmsfO9YIgCE3V8YDq8JNPUvLDj409nNMiyxIWo+/rwckTknXRdJ1wW/1mCp0usTJVXxIHwtWvw8Lb/d9nwGijLOF8Ltt5ATM35XJZr/j6HWMLFxNqYVhKG37ZdcRnY8TgtNEEp42u8fU2k8K0ESn1Nr4Sp4c5SzNPa6XE4VF5d0U200emEBHUMG3bmzNd15m7NLNaIBNIHVBNvdFDB16FHBRBccYCCr6ehWSyYo7tROiQayuPMXtdFJZqeH3MX9X1s7a7VR7+bAsLp6dXe219SIiwsr+weopyoOmx8WG17KMnCIIg+MWvFSpdh6L9YD9a8d/WCIjoWON1q6ENT23Dou151e7FAskUUiSJLnFNs2ZcBFP1qds4uPET+Ohm8DpraYtpBXQYcjcxox9mXl4pU95eQ26Rg2kjkhulCL2lePjSbqzMOlrrTtu+mA0yfdtHkJ5Sf+lVn9ayUlJXxzSoOEcuWHuAO0bWX8DXUuSXujjmI9c60DqgmtQWmMsSKFYrqqd62pu/P+utB4sbrK36zekd2ZRTRPkpNXuBXPQkCUakRtf7WAVBEFqDGgMqVyls/gh+ex7K8uH4Pk6aBywRkH4v9L6+otlFI5o2Ipnluwt83ov5NSFpkLk5vSNGpWkm1Ilgqr51HAZ/3gW7f6xoi3lwXdW2mNYISL8Hek+sbIXeNS6UT+8cyuS3V5NzzM7j49IwNNF/QE1dp5hg3po8gCnvrPE7oDIbZDrFBDP3pvqrSdJ1nTlnuFLi9Gi8uSyLqcOTRe1UHYodHgyKxKnlQIHUAUFFYGRUZJ+rib5IVKxwqhrVVqUC+Vlrms7bv2Xz+GW113CdDed1jcEoy8DpXfRMBplJgzs02YueIAhCc1QtoIo+DN8+VDF7dbzh2cmZUB4H/PwE/PQ4nDcThtzV8IP+Q9/2EbQJNvvMeoC6M4UAbhzcoT6GdlaIYKohyAp0ubjij6OoYhn2eFvM4Fify7BxYRY+vmMId36wnjvmrePF6/tgM4kf1+kYnBzFR9OGcMs7a7C7vJTXEFQZZQlZlhjdJZrnr+tTY47v2VDi8FJQ5qr2eKArJSVOD0fL3USHiM6PtVFkyWcBbKDFrwZZYlBSJGv3HaszODcbZILNBv5yURee/HpHtecD+Vl7NJ1vthxqkGDKoMjcOiyJV5bs8dkoo66LnizBpCFN96InCILQXB0PqMr+dTl652NIWh3d7Y4HWT8/CcUH4aJ/1f8gfZAkiWeu7sGUd9bU2oDJF6tRYcrQjsSGNt3UcTF12NCs4RCVAjFdISSu1nzWEIuRtyYPIMJm4vq5K8kvrX7zLfjnnHZhrPzbebx4XW96FB/AJEsEmw2EWAwEmRVsJoUbBnfghxkjmD2pf70GUlARBPmauQ90pcQgS5Q4fW9OLJwQFWTGrVY/gQdS/Aqg6jBnUj/+fU1P0uJDsRhllFN+hYPMCuFWI9NHprDogZEEmY3oPiK5QH/WZQ3YHW/6qBR6JoRjDnADYYtR4dlretZr0xZBEITWzOLaQFRqUd2B1Mk8Dlj3Nqx8rf4GVof0lDY8c1VPLEb/rytWo8LF58Tx5wu71OPIzpxY6mjijIrMv6/pyYs/7eHq11bw9i0DSIkObuxhNUuKLJHuyaPjvq8Ie/ZTjpS4cHpVQi1GEiKs9R5AncxkkH1uKhzoSokOmEQ6VZ3CbEa6tw1lU05xlccDqQMCGJoShdVkYGzPeMb2jGdXXilfbDhIbrEDl1ejTbCZ9JQ2jOkWU2dqbqA/a18yizP5cOeH7D62G7vHTpAxiK6RXbm267V0CD391SGDIvP25AFcOyeDrbklfr3GYpT5x2XdGder3Wl/riAIglALrxu+ewhJrTq53vH5UuweyLovmCBTxQzfG+vdzNvsYcnkP2ptPXZY/A/oMwnMjXMfeUWfdkQEmbh3/ga8qlZjppDVKKPpFbVW943p3OR7B4hgqhmQJIn7xnSmXYSVa+es5LUb+zKgY2RjD6tZKl28mJDzz6dNsJk2wY2XGhdmNeJVqwdTgXRMA/CoGpGim59fpo9K4U8fb6rWWMGfOiCo2HB32inNPlJjQ3jwoq61fm5FK9fqF4JAf9ZB5hMB19Kcpby28TV2F+1G1VS8+olVqw1HNvDRro/oGtmVO3vdSXq7mrsAappOmduLIknYTEq1C1ap08MVveNZtqcAp1utduEzKhKyJNEzIYyHLuoqzkuCIAj1aceXPrfpgIrMiRdWuXl4eC33NpIEmxfAgFvraYB1G5kazdq/j2HR9jxmL9nLjgNHMZlNSLKMV9MItRiZOiKZ8f0SCWuirdBPJYKpZuSafgnEhpq54/11/OPyNMb2FK3TA6HrOqWLFhP/7LONPRQsRoURqdH8svNIlQSwQFdKBiVFEWQWv8b+GNMttsbGCv4Uv4ZYjKfV3bFP+3BUHzsNBvKzNsgSF3SPQ9d1Xlj/Ah/s+ACn6vT5eV7di1f1sil/E/f9ch+39riVaT2nVQZKXlVj8Y4jzP51L5tyijDIUuW1+dyuMUwdkUy/DhH846ttDEyK4tnxvdA0nV935zMvYx8Hixw4PSqhViP9O0QwOT2J9lG2gP9eBEEQhAAt/2+NnaEfTDfx799c3DnARLilhpUcj72i81//KY3aNt2oyFzSoy0Xdgxm7ZhLiPziG1QqJppjQyzNrqmWuAtrZoZ3jmbebYO49Z2K1um3Dxet0/3l3rsXzeXCck79F/H7Y9qIZFZmVm/bHthKSXJDDrlZMygyL03sw+3vrQ24ANZilHnlhj6n9btmMxm4qm87Fqw5gPeUoMrfn7VBlpgyLImXN77MBztrDqRO5VSdvLnlTQySgdt63sYXGw7y2Jfb8JyUXuE5aYV00Y48lu8pwGZSMMgSi/80CuCPxiwxjO4SE/D3FwRBEM4Ctx2OVG9mdFz/eIVRHQ3MWuHin+fW0qyhNA/KjkBIbD0MMjCuXbuI6RhPUlxoYw/ljIhgqhnq1jaUhXemc8vbazh4zMHMcWkV3cqEWpUuWkTIeec1meBzYFIkUUEm7O7qmzr7u1IyNKVNfQ2vRRreOZp/X92ThxZu9jugshhlXpnYl34dTj+FbcqwJD5Zl1MtmAL/ftapcSEUqFt5b9t7fgdSxzlVJ3M2zyE7N5ZPfjPU+r11vWKTYLtbxWSQWZNdKAIoQRCEpsBZBIqpavvzUzwx2szQt8q5b1At6f+KseK9mkgwZUlNbexhnDFRud5MtQ2z8tEdQ9ibX84d89bhCHBT2taodNFiQsaMaexhVJIkiTmT+mOVfec/18ZqVOp1H6yW7LLe7XhgTKqPKqbqJKBtmIU+7SPO6DNTooMZ3y8B62k0ObEaFf51ZQ/mbp7rM5A6tuwYu/++m21Tt7Hz3p3kvpuLWl71fOBUXXyW+X5AK3Jur8ad89az8UBRwGMWBEEQzjJJxuceHyc5J0ZhbKqBZ5bX0elPahq3/87ff8ec2rQ79fmjafxtCqcl9I/W6aEWI9e9vtLnvkVCBc/Bg3gOHcLWv19jD6WKhI3LeXLTfGwGya+be6i4uX79pv70TGjcHc2bq++2HOK5xbvquCRV0IGcQgeXvLjsjH+//nH5OQzv3CaggMpilHn1xr5EhJWxKX9TtecLvivg8MeHiZsQR/dXu5P8aDLuo26yZ2WjVdlYWEey7kYyVO3MV759CYfencH+564h5+VJ5H30GKHHL3UAACAASURBVM6cbZXPOzwq9y/YiF5DwbMgCILQQKwRFXuU1uEfoyy8vt7NwRLf523d7cB1pLRRzusOt8rPO/P4eO0BFqzZz+L9dso7dGrwcZxtIphq5kwGmVnjezIyNZqrXl1BZr7vwsTWrvSnnwgePRrJ0HQyW0u+/57DTz/N2Oce4/N7hjMgKRKzQcZ46sZFVGwobDbI9OsQwcLp6QzrLNL7TsemA0U88NHGgFZoPJpOQamLia+vxONjryp/KbLE7Bv7MWlIB0wGuda9NoJMCm2CTXxw22BGd4nh892fV7vwqQ6VI58fIf7GeEJ6hiAZJEzRJhLvTMRd4KZ4RXG19zWErq/8/yWrP6Pwp9cJGzyBhLvn0W7624T0vQTH7lVVXnO42ClWpwRBEBqbwQztB9d5WKdImWvTjLy42vfqlFcLZ/9df2bPyFEcfOghihYuxJ1z8GyPtorM/DL+/vkW+j65iHvnb+SxL7fx+JfbeTpyEOf9UMi099eybl9hs524azp3lsJpkySJB85PJSHcyoQ5K5l9Y1/6ixbFVZT+uIjIKVPqPrCBlPz4I4f/+RTt33gdS5dUUoGPpg1h/1E776zI4tsthyl1VcxABZsNXJQWxy1Dk+jYJqhxB97MPf3dDhw1BFLl25dQsuZzPEdzkE1WjDHJhKVPwJKQhlfTyTnmYNH2PC7p0fa0P1+WJR6+pBt3jExhwZr9vLE8i3KnF0WR0DUdt8tN/5QYpo1MYURqdGUtZGZxJh6t6oykfbcdzaMR2q9q4a5iUQjpGULZtjIiRpxIT5RkL7I5HwDNVU7R8g+IumQGti4nWqfbOg3C1mlQlfdzeVVeX5bJqzc0rVVdQRCEVmfoDDi0qcaOfsfNHGnm/c0+VrFMwRjH/ZtO/7gaz4EDlK9cSflvKzjy3H+RrVZsgwcRNGgwtkEDMcaceb2sruvM+uF33lieharp1euGDRZQdX7cnsfS3QWkJ0fxyg19G3Tfz7NBBFMtyIQBicSGWZj2/jqevOKcM7rpa0m8hYU4d+4kaGjN++00pNKffuLwE0/Sfu4cLF2r7lHUPsrGzHFpzBzXNDoOtiQHCu1s2O97haVk9WcUr/qEqAvuwpLUF0kx4Mhah2P3KiwJFT8Lu1tl9pK9Z+X3KjLIxPRRnZg2IoX8MhclDg9GRaZo/BV0v/dtTAlVL2LlnvJq76GWqRiCDUg+VjINYQYc+6oXKUtyRc2V6+BOdK8bW+qQOseq6bB4xxF/v5ogCIJQXzqdB0ZrtWAqe0ZIlf9ODJNx/t1HhzxJhm7jkCQJU/v2mNq3J2LCBHRdx71nD+UrV1H64w8c/uc/MURFETR4ELZBg7ENHIAhIrDaYV3X+dunW/hiYy4ub+1ZHbpekQL4254CrpubwYJpQzAbmk9AJYKpFmZkajTv3TqQ295dS26Rg1uHJfnsXqdqOmuyC8krceLyaoRZjfRoF0Z8uLURRn12ubwqxQ4PmgahVgPun38maNgwZHPjbdJ7XOnPv3Bo5mMkzpmDpXv3xh5Oq/JuRjaajxSCQFZpdh0pZc+RUjrFhJz6NqdFliViQy3EhJhZt+8Yr/SZwL63NuE0/Y7NqNApNoQpQ5MIMVa/KCrBCt4yL7qqVwuovMVeDMHVT++6WrGyqTpKkG2hSLJ/FyuvquH2apgMIjNcEASh0cgKXP0G/O+6Wrv6+WSwwhWvVaQLnkKSJMydO2Pu3JnISTeiqyrOnTuxr1xF0cJPOPTwwxjbtydo0CBsgwdh698fJTi41o97c3kWX2zMxeHxv0Ga06ux81Apf/poEy9P7BvY92tEIphqgdLiw1g4vaJ1es4xB4+O7V6ZLpRf6uJ/q/bzzoosPKqGDmhaRT2HR9Xo3yGCaSNTGNapTbPqFKdpOr/tLWDOr5lkZB7FKEtIErhVnUSPh9t6juEaj9qoS8dlv/7KoUcfJXH2a1ibyF5XrcmyXQVV9lQ6LpBVGglYt+9YZTDl8qp8v/UwX206xNEyF0jQJsjMuN7xXJQWV2fwoes6n64/yPOLd3G03I3D1BbdDtgrVpD2FpSzdFc+JtMw9LBCCFlVuc+irZMNySBRsq6EsIFhle+pOlVKN5cSe03Vtre6ZkRzVmz0rVhD0ewl6Jrqd0DlKxAVBEEQGljyKLjsJfjyHv8DKoMVLvwXdBvr1+GSomBNS8OalkbUrVPQPR4cW7diX7WKwnfe5eADf8LcuRNBgwYTNHgQ1j59kK0nJuNdXpXnF++uMZCqLa3e6dVYtD2P7ILyZlPaIIKpFio+vKJ1+vR567jzg3U8f20fvtt6iIc/3YIONS65/rb3KBsPFJEUHcT7UwYREVTLXgVNxJrsQu7+33rKnN7KjUjVk/Jys5VQns5WePrJRfztkm7cOLhDg4+xbNkycv/2MImvvYq1R48G/3yByhq0UwWySuNRdYodHo6Vu3l1yR7+t/oA6Hrlv7vjVmQW8PCnW7hhUHumj0r5f/buO06K+nzg+Gdmtl+vdO7g6FV6b3YNKhKjMYpGo6LGWGOJppmYYjS/WJLYNZbYEbsmakSlcyAdFO6Og+No18v2mfn9sYAcu3u3u3dHOZ7368UL2J2dnbnbnZlnvs/3eUh3hX+PDMPknrfX8/bXOw854TS9gfFd3ycNGr+HNbU79i5voSgmmksjd1Yu5S+VozpUkgclE6gOUP5iOdZMK+kTD6/2aBKoOwkAe7cBKBYr7m+XkDRgcov7rSrKcZfDLoQQHdawH0ByDsy/Fnx14A9PBQfAlhxKCzz379D/zITfTrFacY0YgWvECLKvvRbD58Pz9Woaly1l36N/x/vNNzgHD8Y1bhxJ48fxHzMnajGJWNLqDcPkucXbuPfc4+PGswRTHVia08q/rhjLXfPWctrfFlBR78fbQt4qQKNf55vd9cx8dCHv/2zyMR1Q/WfDbm569esWq7MduNj9wwcb2VHl5hdnDzwSmwdAw8JFlN95F93/8Xecw4cfsfcVTVnUyKNE8YzSqKpCvTfImQ9/SVWjP+JIF0CjL/R5e25RCe+uKef1uRPokelqssxv39vA21+XRS2IEca0EagbDkoQR5d3AMg5OwctSWP3a7vx7/WjOlVSR6bSY24P1EOqBZqmEnqt4Qjthz2J9MmXUPXJ4yiqhqPXCBTVgnfbarzb15Ixo2mxlrG9pKCNEEIcU3pPh1s3QckXsOhhKF4Aiha6J2fokDcxVLCi4BSIcv5LlGq3kzQ+FDhxExiNjbhXraJx6VL2/OnPPNL5LBpTOoe9Lta0+oBh8kbhDn5x1oDj4kaeBFMdnM2ict5JXXl3TXl4FZVmBHSTvfVe5jyzjHdvmHxMpvytLK2KKZA6lCdg8MKSUnJT7fxkcu923LqQxsWLKb/9drr//VFcI0a0+/uJ6HJT7Gyvcoc9Hs8ojUVVeHZRCR6/TixfJ79usqfOy6x/LOKjm6eQmxIKZhZ8s5c3CuMIpA4wbQRqR2JJ/gZLymYAMqdlkjmthWDHtBConNbkodSxs1GTMqhd8hoV7z+IYnNi79SH1AkXNVkuyaYxd1pBfNsphBCi/SlKKKjqPT2UyuCrC/1tT23zAKo5alISyVOmkDxlCqZpUnr3hxH7C8eTVq8qCtsqGxnQOUIhjWOMBFMngPs+2BQ1kGoubzWgm5RUNPLlln1M79/6EpltyTRNbnt9TdRAqrn98gR0/vLxN8we0b1dR90aly5j520/p/sjD+MaJWWlj7YfjevJpl11YSl58YzSePw6ikJMgdQBhgm1ngBzX1zJ/OsnAfDPBUUJ5ZIDYNrxVU4/GEy1xK7ZaSi/CMMf/h1OHjyD5MEzmn29y25hSh/payaEEMc0RQFHWsvLtTNPQEdVlIjzbONJq1cUqPME22MT25wEUx3c+p217KiOPEExlrzVRr/OE18UH3PB1KrtNeyp90V8Lpb9UhR4rXAH18Zxx93j13l/bTmbd9dT4/aT6rDSOzeZc4d3Jc1pbbJs4/Ll7Lz1Vro99BCuMWMS31HRZs4e2oVfvb0+4nOxjNKoZihwNyL0Om8pAAoaJpt21fHtnnocFo01UZrgxvLZBTC83VACuZjW6CXLraoVi2rh/in3s3JTV55ZWBJXVSUAh1Xlj+cPPSZHpoUQQhx7rJqKHmW+VLzFj46XCrISTHVwT31VTCDCPKl4ykGv2l5NWbWb7hmuw1dz1Dz1ZTHeCBeGse6XN2DwzFclXDOld4sXiqWVjTz1VTHzVu5EUUIFAQ5wWjXue38jZw3pzDVTCxjUNRV3YSE7b76Fbv/3V5LGjW2DvRVtwWHVuGhMD15auh2/Hv6daGmURlHViCNSsQZAgaDBMwtL6JzqaHWJdotipbf6I8otT6Cg4A5+l77osoS+p7P7zuaSgZfQPaU703uYlNd4+Gj97pgDKodV5RdnDeS0QZ1aXlgIIYQgFEy5rFpYFgjEl1bvDxrkphz9ljaxkGCqg1uxrSriHYJ48latmsKaHbXHVDC14Ju9RLrxEc9+NfqDbN3XQL9O0XsGfbJxDze+8jUB3YiYKnngwvS9Nbv4eMNu7hyWzKS/3UW3Bx8gafz42HdIHBG3nd6fBd/uo7SykQjxVFROq4Y/qIelgMcTAOkmvLN6J6cP6tTqEu1BA7o5RvDKhV/waemnlNaVUuurJd2eTn5aPqf0PAWHxXFweUVR+OuFw+me4eKJL4tQIGoxGpctdLfwgQuG8b1hXVvcFiGEEOJQs0Z047UVO8Kum+JJq++Tm3zc9D6VYKqDa/BFzjeNJ29VN6DOG7ms9NEQ1A18Ua6E49kvTVWobvRHff6TjXv42SurYipwoZsmesDkz8sr+Pl1v+PqiRNbfI048pLsFl69ZjwXPr6E8hoP/ijV+A7ltGqcO7wL76wuRz8sAIknAAIwA0H2rFoHWvgcpHgb6Xr8Ok6Lk3MKzolpeUVRuPX0flw+MY8nvyziyS9LcFg1LJqCaUJAN+ie4eTaaQWcM7zrcVFBSQghxLHnysm9eHNlWcSb0LGk1SfZtbimYRxtEkx1cDat9eWgFSX6eo6G5i5/483HjbaubRWN3PhKfJUCAXyajb9+62dkaRWj8qSc9LEoN8XBez+bzG/e2cAH63ahKkrE1Lckm4bTpnH7GQMYnZ/B+2t3hS0TbwCkqgqpmVlQFf5cvJ/ddJe1xWUOatgH9bsg6CPLkUpeqsr5I7rxs1P6UusJYFEVMpJsdDtO7gIKIYQ4dhXkJDO4ayprdtRGzI5qKa1eUxXOGBxeWv1YJcFUB5edbKeiIXz0JZ68VVVRyE09dvJWrZqKVVPxR0hTime/gj4/ts0b0HNHoiUnN3nu6YXFBKKMfrVUbMAbMHjksy08f+W4iK8XR1+Kw8r/XXQSvzl3MG8U7uDFpaVUNPgI6CYuq8bArqnMndqbqX1zUFWFvXXeiCeEeAMgQ9WYMHEAi/77TZO5dxDfZ9dl0xiVl9HCm+mw5RNY9BDsXAUWG6CAafD9QICJ+bPJ5+fQo0+L2y2EEELE45GLR/C9RxZS64kvs0kBTh3QCat2/BQ+kmCqg7tkXE/++OHmsDvv8eStKj4vwz17MM1sFOXY+HBPLsjm82/2ho0sxbNfNgzSX36aLb/cgC0/D9eo0bhGjcQcOpx5K3dGHJ6OtdjA0uIqdtd66ZzmCFuHOHakOa1cNaU3V01pvudYZpJtf9PfpgF2PAEQhCoTXTi6Bw/855uw5+L57JomzDqpW/Q32rUG/n0h+BtCfwD076pf2oG80jfh8flQcDJ8/xmwHTtzIoUQQhzfume4eG3ueC5+cim1nkBMLUUcVo0/nT+EF5aWcvNrq/nLBcOwW479lPNjJ3dLtIvzR3bHjJLMljp2Nhkn/4TaJa9R9ugllD32Y+pXvY+z73fzP2yawgWpDey95WZKzp9N1YsvoddELut8pJh+Pxeb27Hrkec7xbJfdovKVacNovdLL9Bv6RI6//JXWDvlUvv2O7xw/W8wveHl5A8UG8g87Tpc/Sei2hwomgVXn3FNLnQhlD748vLSNt1vcfRYNJUfju0Rdqfs0ADI/e0SjIAXUw/iKSqk+vNnmyxr01TmjM8j2WHh+6O6YYlQRTKWz66mhib3Jtmj3AsrXQzPngkNu78LpCJQjAAEvVD0P3jmNPBFX1YIIYSI14DOqXx40xROHdQJu0XFHqHUuVVVsFtURvZM5/W54zl/ZHdeuXo8/qDBnKeXNzu3/VghI1MdXLLdwnnDu/HW12URK4i12LQzGOSqy06j610X4162jJo357HvkUdInjKZtO9/n6QJE1COUJdtw+Oh5o03qXz2WfoVFJCZ933K3ZFT8WJpRnrx2J4AqDYbrpEjcI0cQdZVV1H73ga8i7aFLR9PsQF/0GD9zrqWd0ocNy6fkM+LS0o5fKZdLJNpAVBgzoQ8AOZOLeCtVTsJRigd29Jn127RmDs1ykhaxVb49w8g4I78fCRBL1RuhVd+CJe9C0fo+yyEEKLj65Lm5Mk5o6lo8PHq8u3MW7WTWncA3TRJtls4ZWAuV0zqRa/spIOvcVg1/vGjkdz/8Wa+/9hinrtiDHlZSc28y9ElwdQJ4J6ZA1lcXEF5jRc9lnHW/RxWlZuV7bgv+yGe++8nacIEkiZMQK+pofb9D9j74F8xamtJmz2b9NnnY+3aPmWU9YYGql9+haoXXsB50nC6P/oIzqFDeWBrBT95fkXcRSKcVo1rp/UmOznyPLAab+srIALUxZknLI5tPTJdTOuXwxff7sN32Hy9lgMglVMGdqJLmvPgup6+bDRXxvn5dVhVnpgzivzsKCeVT34F/sYmD+U/VI87ACU3JZNkC42GPb3Kz0trAyz48f71BL1QvgqK/wd9To15e4QQQohYZCfbueHkvtxwct+YlldVhV+cPZAemS4ueHwJj186quW5wkeJBFMngFSHlTfmTuTCJ5awu9YbsWHp4RxWlVtO7cfcaWdRv2AgZbfeSvoFF5Bz/fVo6elkXnoJmZdegnfjRmrenEfJ+bNxDBlE9sn5ONVvUNx7QxPgXVkwYCYMvQBsMd5VaKyEVc9jrngWs24viu4nQ7GTcflYtNN/BD2HADCpTzZ/PH8od89fF/MFqdOqMWtEV248JfqXOdUR+WsRb7GB5GhpWOK49dAPT+K8vy+itNId0/cIQvOk8rOT+L8Lhzd5fGKfbP51xViuer6QoGE0+xm2agoK0C83mQf+8w0Pf7qFbulOLhrTgwkFWaG5jA17YetnRKpRqZvw8DI/d09pppCMvxEWPSzBlBBCiGPGpePz6Jbh5JoXCvndeUP43rAuYct4vV7WrFnDli1bcLvdaJpGamoqJ510EgUFBajtnHEhV3sniM5pDt6/cTK/f28j760tRyG8HLSqhIZWc1Ps/PJ7gzh1UCcAUqZPx/nWW5TffQ/bLr2Ubg8+iK1HDwAcgwbR+a48cicqsOQx+MaLoh42slO6CD6+C4b/EKbeDqlRRrAaK+GDWzC/+QhTN1AJogBoAF7Y+RW8tDIUoJ35Rxh4DrNHdiczycaNr3yNbpgRO25DKIgyTJMbTu7D9dMLIhbSCOzZg3vZMrIWb8Ou98SnNS09HU+xAauq0LdTcrPLiOOPy2Zh3vUTufzZ5Xyzuz6sIl/48hqDuqTy3BVjIvZtGt87iy/vmMFrK7bz9MISvAGdoG4S1A0smoqmKviDBgHdxGFRWHtI6ujK0mo+2bSHNKeVuVN7c6nvdSxRCsTcPtHGXxb5uH6MjXRHM0VkdiyDmu2Q3jO2H4gQQgjRzmb0z+WFn4RuPu6odjN3am8URaG6upoFCxawYcMGFEUhEGiaEbRlyxasVisTJkxg/PjxWCztE/YoZoRyvweMHj3aLCwsbJc3FkdPvTfAvJVlvLRsOxUNPoK6icumMTovg6un9uakHukRgw3TMKh+8UUqHn+CTnfdSeq556I07oN/zYSa0lCqUHNUC9hTQvMyugxr+lz1NsynTgN3BQox3PG3OGH6XTD5ZiDUcPSTjXt4bEERm3fXhfpiKRDUTTKTbMyd2pvZo7qT6vguQApWV+NetpzGZUtxL12GXlWFa9w49DHjOH1zOr4Ic8zqlr9F7bK3yDrjp81WW7NbVP5z89To6VhHkaIoK03THH20t6O1jubxKaAbfLhuF49/UURJRSMB3eDAQJWmKlg1hYKcZOZOK+DsIZ2xxNCnTTdMFhdVsK2ikXpfkN21Xl5bsQN/0Gi2txqEbhYMV4t5ht+SpPiaPJf/UD1Pn+vknyv8DMpRue9kR3ia3wFWF5z9AIy4NI6fhhBtpyMcn+TaSYj2savWwxXPrWBkXgZXj0jl1Vdexu/301wsA2CxWMjNzWXOnDk4nYn1U2zu2CTBlIibd/Nmdv785zj796ZL989R6svBiDzP6HAmgD0F5eoFkB3qb+PbuArLazNRzUaUeEZirU44834YdXmTh6sb/VS5/eiGSbrTSk6KHUVR0BsacBcW4l66jMZlywjs2IFz1EiSxo0nafw47AMGHCymcfOrX/Peml0Rews1bPic+sJ3CFTuaFJswNF94MFlRuVlMO+6iXHszJHTES5W4Ng5Pm3aVcfn3+ylot6PokBWso1TBnSif+eUhNe5vKSKy59dhieO+VR2AgxVinnFdh9W5bsRswPBVOdkhUnPNrL1Z8m8800wcjClWuHU38LEGxLediFaoyMcn46VY5MQHVG9N8Ct/1pAlz1LUc3ms0MOpWka2dnZXHXVVVitcTS936+5Y5Ok+Ym4OQYMoNcbb+C7fzJm1Q4ULfaiFgpgeuvR/3kqnhkvUjPvLVI987B18RJ3C6uABz66A/qdASnfdcrOSLKRkWTD8HrxrF7NviVLcS9dinfLFpxDh5I0YTydf/0rnEOGoET5Qs2dVsDHG3ajB+KvgOi0avzsZGmEeqIY2CWVgV1S22x9Db4gV/5rRVyBFIAPK+vNfP4veAF3Wl8Le35IrsbMfhb+vNDPwJwody0UhfjuaAghhBBHjtOi0LfuaxpjDKQWLFhAVVUVs2fPprKykg8//JDzzjuvTbdJgimREDVQg1PbweGT3WOpHKYooPhrqfrV5egpfUgZ7UY57EsRUwUyCL1/4bMw427MQADPuvW4ly2lcekyPOvW4ejXD9f4ceTccjPOk05CdcTWRHdgl1R+NXMQ972/KWxuWXOcVo0fT8pnev/cmF8jxKHe/roMI0rGQOPGBdSteJtAZRmqzYk1tzdpEy882DDai50X9NO52TIPuxI+WnzvdAcjn2jgtglRClFo1tCcRCGEEOIYtHnzZoLB/XPqD7Fu3TqWLFlCRUUFdrudzp07M2XKlCbLBINB1q1bxxlnnIEjxuvBWEgwJRJT+GzUp2KpHKZYDLrPGYAv0B2zfEXEm+ExVSAL+jC+eoSyV0vwrFyNtUcPksaNI/OKH+MaPRotOfEiEJeMy8MwTP7w4aaYqgUeCKTuOKN/wu8pTmymafL4F8URC1vULZ9P7bI3yTr9pzh6jUTRLHhKVuLZsuxgMHXAR8Y4ZmmLwtbRJ1PlosFWHlnuZ2hupC9dUKr5CSGEOGYtXLgQv79pI98lS5awcOFCZs6cSUFBAZqmsXXrVjZv3ozNZmuyrKIorF69mvHjx7fZNkk+h0hM4TOg+yI+dftEGw8u9lHjjZ7+pwDqrkKc3qWoauRAJZb1AGDoZJ0ykIJP/kvv+W/R6a47SZk+vVWB1AFzJuTz76vGMa1fNjaLiu2w7t2Hdu7+56UjufPMARGLdwgRi3U7a6mK0O3d8DVSs/DfZJ52Ha7+E1FtDhTNgqvPuCaFTwAacfJ08Kyo7/HraXYa/RG+U4oK/c+EJBmZEkIIceyprKykoqKiyWNer5fPP/+cs88+m4EDB2Kz2dA0jf79+3P66aeHrSMQCLBs2bI23S4ZmRLxM3RwV0V9enRXjen5Fh5c7OO+k5sZRtXs4K5s9XpUm42kgT0ho32auY3Ky+T5K8exu9bLy8tLWb+zjjpPgGSHhb65yfxoXF6Tzt1CJGpHlQc1Qizu27kZM+jH1W9CTOspN7MP/nvbzU0LYfRIU/H+MsIcL4sDJvwsru0VQgghjpSamho0TSMY/C6NvaysjGAwyMCBA5t5ZVP19fVtul0STIn4BTygas1W8PvdDDuTnm3kpnG2qMsAocCsGTGtxzRBD7+b39Y6pzm49TRJ4RPtx+0PYkQYNNI9daiu1JiaRUOoGEVcrC4Y8n3oMSa+1wkhhBBHyOF9pADcbjculyuuxry6Hvtc+FhImp+In9XVYhB0aOWw6MxQefPWrkdVwZHW7HqEOB6kOCyoEdJENWcqhrsOs4Xv3QFOJRDq6xYLqwt6z4CZD8WzqUIIIcQRZbeHz6F3uVy43W4MI/YKuG3dvFeCKRE/VYW07i0udu90B0+t8rOzLsqcJz0AeZNaLMXc8nr80G1Ui9sjxLGuT24KQT38hGDvNgDFYsX97ZKY1tO3Z1foPiaUuqdGGaWyJoWaX4+/Hi56CTRJVBBCCHHsysrKapLiB9C9e3csFgubN2+OeT2ZmZltul0STInETLghdEe7GYdWDgujqDBgJky5NXTBl+h6UKDXtCZ9poQ4XvXJTaZPbnjhFNWeRPrkS6j65HHc3y7BCHgx9SCeokKqP29aWTPJpnHNyQPhyo/husWhptbWpNB3TrMBCmTkwxl/gNu3wim/Ct0gEUIIIY5hqamp5OXlNXnM4XAwffp0PvzwQzZv3kwgEEDXdbZs2cInn3wStg6bzcakSZPadLvkVqRIzEkXw6e/aXGxX0+z8+La8BxXLA6YeAN0HREa5ar4NrH12Fww6aZYt1qIY9610wu48821NB5WHj117GzUpAxql7xGxfsPotic2Dv1IXXCRU2Wc1g1pvXNCf0nqwC+99fQn4AXgh6wp4bmPAohhBDHmUmTJlFWVtakPPrEiRNJ/RYJxgAAIABJREFUTk7myy+/5K233sJms9G1a1emTJlCUVFR2DoGDRrUptskwZRIjCMNRl4OX78QKkixX0yVwzQbdBocCqQAzvwTvHoJBL1xrscOnYaGUgWF6CBOH9SZPzhDzaIPL0aRPHgGyYNnRH2t06px86l9USOVBLQ6Qn+EEEKI41SvXr1ITU2lqqqqyTypYcOGMWzYsLDle/TocfDfVquViRMnypwpcQw544+huUotpOk1oVogKQd+9Pp3j/U5FU6/r8ViFE1odkjvAZe8AdLXSXQgNovKq9eMJ8luieuj7bRqzBzehUvH57W8sBBCCHEcUlWVyy67DIfDEVdfT6vVSu/evZk6dWrbb1Obr1GcODQLXDIPCmaE5mS0xOqCzAK4ZgG4Dpv8N/ZqOOfRUGBmaS6oUkLv1W0kXP05OCL0yxHiOJeXlcTbP51EdrIdh7X5w7RCKJC6aEx37p89TJpGCyGE6NBSU1O55pprSE1NxWptuRWI1Wpl4MCBXHjhhXGVUI+VpPmJ1rE64KKX4duPYdFDsGsNGAYY+3NZFQ0s9lCBiMm3wNAfRB+BGvYDKDgZVr0AS/8RSh9UlFAfKUUF3QcFp8LEn0HP8TIiJTq0gpxkPrttGq8t38FTXxXT6Avi9uscyPyzW1RMYFJBFtdMLWBCQdbR3FwhhBDiiElPT+f6669nzZo1LFq0CI/H02QelaZpKIpCjx49mDRpEgUFBe12s1ExzSjlpoHRo0ebhYWF7fLGooOqLILNH0DDHjACoZS+3jNC6YDxfIgNA8pWhNaj+0NztLoMh+Tc9tv2E4SiKCtN0xx9tLejtU6k45NhmCwqquDr7TVUNvhw2jQ6pTo4e2gXOqXKPCjRcXSE49OJdGwS4lhgmialpaWUlpbS2NiIxWIhJSWFgQMHkp6e3ibv0dyxSUamRNvKKoBJN7Z+PaoKPce1fj1CdACqqjClbw5TDlTpE0IIIQQAiqKQn59Pfn7+UXl/mTMlhBBCCCGEEAmQYEoIIYQQQgghEiDBlBBCCCGEEEIkQIIpIYQQQgghhEiABFNCCCGEEEIIkQAJpoQQQgghhBAiARJMCSGEEEIIIUQCJJgSQgghhBBCiARIMCWEEEIIIYQQCZBgSgghhBBCCCESoJimGf1JRdkHlB65zRFCHAF5pmnmHO2NaC05PgnRIR33xyc5NgnRIUU9NjUbTAkhhBBCCCGEiEzS/IQQQgghhBAiARJMCSGEEEIIIUQCJJgSQgghhBBCiARIMCWEEEIIIYQQCZBgSgghhBBCCCESIMGUEEIIIYQQQiRAgikhhBBCCCGESIAEU0IIIYQQQgiRAAmmhBBCCCGEECIBEkwJIYQQQgghRAIkmBJCCCGEEEKIBEgwJYQQQgghhBAJkGBKCCGEEEIIIRIgwZQQQgghhBBCJECCKSGEEEIIIYRIgARTQgghhBBCCJEACaaEEEIIIYQQIgESTAkhhBBCCCFEAiSYEkIIIYQQQogESDAlhBBCCCGEEAmQYEoIIYQQQgghEiDBlBBCCCGEEEIkQIIpIYQQQgghhEiABFNCCCGEEEIIkQAJpoQQQgghhBAiARJMCSGEEEIIIUQCJJgSQgghhBBCiARIMCWEEEIIIYQQCZBgSgghhBBCCCESIMGUEEIIIYQQQiRAgikhhBBCCCGESIAEU0IIIYQQQgiRAAmmhBBCCCGEECIBEkwJIYQQQgghRAIkmBJCCCGEEEKIBEgwJYQQQgghhBAJkGBKCCGEEEIIIRIgwZQQQgghhBBCJECCKSGEEEIIIYRIgARTQgghhBBCCJEACaaEEEIIIYQQIgESTAkhhBBCCCFEAiSYEkIIIYQQQogESDAlhBBCCCGEEAmQYEoIIYQQQgghEiDBlBBCCCGEEEIkwNLck9nZ2WZ+fv4R2hQhxJGwcuXKCtM0c472drSWHJ+E6Hg6wvFJjk1CdDzNHZuaDaby8/MpLCxsn60SQhwViqKUHu1taAtyfBKi4+kIxyc5NgnR8TR3bJI0PyGEEEIIIYRIgARTQgghhBBCCJEACaaEEEIIIYQQIgESTAkhhBBCCCFEAiSYEkIIIYQQQogESDAlhBBCCCGEEAmQYEoIIYQQQgghEiDBlBBCCCGEEEIkQIIpIYQQQgghhEiA5WhvgBAiXJ03QE1jAIA0l5U0p/Uob5FoLY9fp8rtxx80SHFYyHTZUFXlaG+WEKKDMwyTKrefem8Qm0Ul02XDadOO9mYJ0WFIMCXEMcIfNPh4w24eX1DEt3vqsVlCA8cB3aAgJ5nrphdw5pDO2C1yEjxemKbJim3VPPFlEV9+uw+LqqIqEDBMUh0WfjK5FxeN6Ulmku1ob6oQooOpavTz2ortPLOwhDpvEKuqYJgQNAym9sth7tQCxuRnoChyU0eI1pBgSohjwAdry7lz3jpMTBp9OgBBv37w+c2767n7rXXcM389980awqwR3Y7WpooYfbunnqueL6SiwYfHr2MCAf2732lFg5+HP9vCQ59u4eKxPfnVzEFoMlIlhGgl3TD5/fsbeWX5dhQFvAEDAP8hy/xv016WFFWSnWznqctG079zytHZWCE6AJkzJcRR9uzCEm57Yw0NvuDBQCqSRr9Ogy/IXW+t5bEFW4/gFop4rSytZtY/FrGjyo17fyAViTdg4AsavLZiB1f+awVB3Tii2ymE6FiCusFPnl/Bayt24AsaBwOpw5mA26+zo8rN+f9cxMrSqiO7oUJ0IBJMCXEUfbC2nL/8Z3PUE14k3oDBI59tYf6qsnbcMpGokopGLn92ebNB1OE8AZ3lJVXcMW9tu26bEKJju3PeWpYVV+EJRL8xd6gDQdXlz66geF9D+26cEB2UpPkJcZQEdINfvLUuYiDVuHEBdSveJlBZhmpzYs3tTdrEC3F0HwyAJ2Dwy3fWc9bQLjisMoeqPfiCoQCnssFP0DBJc1oZ0TOd7GR7s6+7990NNPqDEZ9r7vfqCeh8tG43V0ysZWj3tPbYJSFEB7aurJYP1+2OGEi1dE5p9Ae5972NPH/l2CO92UIc9ySYEuIo+e+GPehm+NhF3fL51C57k6zTf4qj10gUzYKnZCWeLcsOnvgAMOGj9bs4f0T3I7jVHd/OGg8vLN7Gv5dtB0JFJExAVRQCusH0/jlcM7U3I3uGT9zeU+dlSXElEX6tMf1efUGdp78q5uGLR7T3bgohOpinvyrGFwwPpGI59pgmLC2uZE+dl06pjiO96UIc1ySYEuIoeeyLrWFzpAxfIzUL/03W2Tfj6j/x4OOuPuNw9RnXZNlGv85jC4okmGpDT3xRxP998i2maeLXIyfp/XfjHr7aUsGovAyenDO6SYnhF5eURnxNrL9Xw4SPN+ym1h0gzSXl8IUQsal1B/h4w26Mww5b8ZxTIHQM+/kZ/dt7c4XoUGTOlBBHQaMvyOZd9WGP+3Zuxgz6cfWbENN6SioaqXH7W15QtOj+jzbz0Kdb8AWNqIEUhO7guv2hFMDvP7YY7yEpNR+sK8cXDE/bjOf3atVUFhdVJLYTQogT0pLiCqxa+CVdPMceX9Dg/bXl7bF5QnRoEkwJcRTUeAIRT3y6pw7VlYqixjYPyqqp1LgDbb15J5w3Cnfw3OKSmCdtQ+jCo2hfA9f/e9XBx2o9kedKxfN71Q2TGo/8ToUQsat2B9APH5Yi/nNKnTfyMUwIEZ2k+QlxDNGcqRjuOkxDj/nkJ1pHN0z+9FH0iorNTdz2BQ0WF1WwaVcdA7ukRn2PeH6vhmny5sodvLxsO42+IC67Rt/cFH48MZ/hPdJbta9CiBOLnFOEaH8STAlxFKQ7rQQi9BSydxuAYrHi/nYJSQMmt7iegG6QLnNrWmXBN3sjTtqG2CZuB3STZxaWcNdZA6I23Y3n9+oLGny9vabJ3IeN5XV8vH43XdIc3HRqX847SZo2CyG+k+GyRjz+xHtOSXXIZaEQ8ZI0PyGOgiS7hQFdwjvOq/Yk0idfQtUnj+P+dglGwIupB/EUFVL9+bNhy/fKTiLdZTsSm9xhPf5FUcRmyQcmbmeedh2u/hNRbQ4UzYKrzzgyZlx5cDndMJm3sozpD3xOkk2LeEET7+81bBK5GepFVVzRyF3z1nHP/HUYEVJ6hBAnpgm9syPeoIvn2GO3qMwc1vVIbbIQHYbcghAiTkHdIGiYre7vdGU3uGdHAI/adGQpdexs1KQMape8RsX7D6LYnNg79SF1wkVNlkuyaVw3vaBV2yBg3c7aiI/HM3HbadN49sdjyctyMfUvn0ecuxDr77UlnoDOW6t2YtUUfnvukLheK4TomNJcVs4c3Jn31paH3YyJ59gzZ0Je2GPrymr597JSiisacfuDoZ57PTK4ZHxPuqQ522uXhDhuSDAlRAy+2V3PMwuLeX/tLjwBHQVQFIXh3dO4dloBJw/IxRKhoEQk7lWr2PfIowwq3402+qcQYapO8uAZJA+e0fyKFDhrSJf4d0YcZJomvihzpeKZuK2pCo2+IJ1SHUzoncUXW/ZF7DUV0+81Bp6AzmsrypjaL4eTB3Rq9fqEEMe/q6f25r8b90QspBPLsWdsr8yDPaZM0+TdNeU88tkWymu8+IJ6kyBtxbZqnvqqmHG9Mrn5tH6M7JnRpvsixPFEgikhmrF1bwM3vfo1RfsaCOjmwREHE8A0WbW9hlteX41VU/n1zEHMHhm955Nn9Wr2Pfp3/Nu2kX39dazKG41//kYiRlMtcFhV/jBrSKtHx050iqKgqkrEkaR4J27bLKFg+jfnDuacRxfS4GtdVazmCl9AKKD6+/+2SjAlxAnK49dZVlJJVaMfwwzNxT11YC6fbtobV2VSAIuqsKvWQ0lFIz0ynNz+5lo+Xr876nr8+1tAfLmlguXbqrj33MFcNKZnq/dJiOORBFNCRLFqezWXPbOcRl+Q5manhObb6Nwzfz2llW5uOa1fk+c969ax79FH8W3dSvbca0k/fxbPLCvjwXc24o+Q494Sh1Xl5lP7MUua9baJNKeVqsbwXl3xTNwO6ia5KXYgNI/t+SvHctkzy3D79WY/O9HEUvgCYEN5HSUVjfTKTkrgXYQQx6OSikaeW1TCG4VlaKqCaZqYgKooBHSDNKcVwzQj9rw7nAK4bBrPXzmWjbvqmP3PRRTkJLOhvBZPlFH7w3kDBr95dwM2TeX8Zm4oCtFRSQEKISIo2tfAZc8sp6GFQOpQnoDOk18W8/ziEgC8Gzey47rrKfvZjaTMmEHBxx+TcdGFfPxNJQ/+95uopbibY9UU/jx7GNdOk7lSbeUHo7pj01pXNCI3xU6f3OSD/x+Vl8H8n06iR6YLl00jco2/yGItfAGh4hevrdge1/4KIY5Ppmly/0ebOPOhL3l52XY8AZ0GX5BGv47bH/q3L2hQ0eAjaJgoCtgtkY8+B4KoHpku5v90EqPzM7lsQj6XT8hnZWl1zIHUAd6AwS/mr6O0srEN9lSI44uMTAkRwR1vrqXRHzlNq7n0K09A548fbGLkK4/iWLOSrKuvpttDf0O1h0YtdMPknrfXJ9TTCEBBYUrf7PbZ6RPUZRPzeW5RScTnYpm47bJpzJ3WG0VpetHSr1MKX9w+ncLSap74oogvvt2HRVVRlVA5dUUh4p3jeApfBA2Tkgq5eBGiozNNkzvnreW9NbtaHHEyTMA0sVtUclPsuP1B6r06Vk3BMCFoGEzrl8PcaQWMzss4eOwyTZP5X++MegOxpfOTrps8u2gb9547OMoa2pC7CjzVoX87M8CV2f7vKUQUEkwdx0zdILDHjeEJoiigJlmx5LrCLupEfLZVNLJ+Z23EAgKxpF+Zfj8f54/ntgf+iOpwNHl9a3saqQq8smI7N8zo27Y7fQLSDZPPNu3h8S+K8OvRxx9jmbg9a0Tkvk+KojAmP5Mx+Zl4AzrVbj++gEGq08o989fy0fo94dsVR+ELIGJZdyFEx/LEl8W8t2ZXXHOhfEGDfQ0+vjekC/fMHESdJ4DdqpLhskWcb7tqew37GnwR1xVTzz3D5I3CHfzirAHtM5834IWN78DCv0HVVtBCNynRfZBZAJNvhUHngdXR/HqEaGMSTB2HgrU+GpfsomFpeagSwoHYyTBRHBZSpnQjaXRnVKf8ehPx7KKSiD18DqRfZZ19M67+Ew8+7uozDlefcQf/79esvOJO4labPWwdLfU0amnd3qDBcwu3cd20PlEbxIqW/WfDbu6atxa/brQqGHFYVf7+oxG4bC1/1xxWrUkZ4azk8M8HxF/4Qpo2C9GxeQM6j3y2JWog1dyIkTdg8P66Xdx6Rn/yW5hb+cKSbRHfI9bzE4QuR/67cQ/nDm/jflVrX4f3bwn9298Q+lsPfPf8vs3wwa2hPzP/BsMubNv3F6IZMmfqOGIaJjUfFLP7gRXULyzD9OqYPj30t1fH9BsYdX7q/ltK+R+W0bB819He5GOaaZo0+oLsrvVS3eg/WNHtvTXlBCIEU/GkX/mCOusP619kmiYrS6sjLx/Huj0BXfLSW+H5xdu46dWvqXYHWh1I/fH8oQlX0xvePR2XLTxYOrTwRUucVo0RUpJYiA7t/bXRz+V1y+dT9dlTpI2/kO43vES3654jZeTZeLYsO7iMaZq8sHhbi+9TvK8xYkZGvOensmp3i8vFZeFD8O6NoSDqQCAVyYHn37sx9BohjhAZujhOmIZJ1aub8W6qgmDzJRHM/fNxat8rxqj3k3pKeBO+E1mdN8C8wjKe/KqYffU+rJqCbgImnD20M/XeQMTXxZN+pSoKVe6mFeIa/TqqomBEOFvF29OoxhN5G0XzPlxXzp8+2pRQ8Q8I/eytmkKv7CR+e85gxvXOSnhbZg7ryq/f2RD2+KGFLxRVw9FrBIpqwbttNd7ta5sUoTBMkwtGSfUsITqyxxZsxe1PfMTIr5v8e9l2bju9/8EWDgd4AzoVDT4qGvxUREnxi+f8ZJjQ4G1dW4gm1r4OC/4MQU/srwl4Qq9J6QLD42uKLkQiJJhKwK5aDy8tLWXNjlrqvAFcNgsFOUlcOj6PgV1S2+U9az/ehndT1cFAKRZmwKB+QRlappOkEbntsl3Nvr9psmp7NfO/3kl5jZegbpCVbGNav1zOGtoZu+XI9kgyDJO//Gczzy3ahqooB9MZgoeMQr23ZhfRps/ElX5lhibjHkpViDqxN97ULk3mxcXNF9S54811EQOpliZWAygKnDe8C1dPLWjV93xjeR1PLyzm883R58/FUvhCUxTOHNKZNKek+QnRUemGSfG+yJkI8WZL3PbGagJBk4oGH5WNfirqffiCofNydrI9am+8eM9P/9u8F01VKMhJpiAnmd45SSTZE7jcDPrg/VvDAqn8h+pxB6DkpmSSbKFz4dOr/Ly0NsCCH+9PZQx64IPbYPAssEROqRairUgwFYeVpVU89OkWlpdUYZo06RG0YlsV81aVkZ+VxI2n9OXsoV3a7H31ej8Ni3dGHJFaXraWP37+GN9WbENVVfpm5fGbU37GSV0GAqGAqvbdIlzDclAilH9uD0Hd4PXCHTz+RTEVDT48Ab1J6sB/N+zhnrfXc/HYHlwztTe5Ke0/WTSoG1z30koWbq1sthKSHinHYb94+g4BpO2fy2IaBv6iIrwrv0bRnaCEZ9fGs+6AbpCZZGvx/UVTH63bjRnh9xtrTyeXTeN7w7omHEit2l7N3W+tY1tlY5MG0NG0VPjCalH56Yw+CW2LEOL4UO8NYNEUAhHu8sUzYqSg0DnVwUk9MshOtpGdYic7yU6q03KwaNWv31nPy8u2N7nBCPGdn5xWlWn9clAVhU837eGJL4spqWggw2XbH1wlUZCbfDDQ6pRqj140a+O7RLsFqZvw8DI/d09pLlAyQwUrZP6UaGcSTMXopaWl3PfBxqjpQboRujjavLue215fw8ItFfx+1pA2KRLQuCxyvnS9r5Er3ryLP5x+K+cMmIFfD7K8bA12remFtmmYeDdX4hzc/iW1G31BrvzXCtaW1UafLLs/XeH5xdt4c2UZL181nkFd22dE74C7569j4daKuHtnHCqe9KtgMEjnD99g++qVeFavQUtLwzVyBFPSJ7KgVg07PcSz7k4pDrpnOBHxefyLooOfvQPimVjd6NN54stiThkY/xypj9fv4pbXVrfq83coh1Xl4R+eRL9OKW2yPiHEsclmUYl23yWeESObRWXWiG4M7poWdZnLJ+bz+oodYcFUPOcnh1XjjjMHNLn2MQyTnTUetu5roGhvA5t21fP+2l0U72vAGzBCAVZO8v4gK/TvvKwkbIseijpH6vaJNv6yyMf1Y2ykO6JcZ/kbQnOnJJgS7UyCqRi8tmJHs4HU4TwB/WCvhj+eP6RVpcpN3aRhcXnEUaniqh0AzBp0KgBOVWNar7Hh6/Dp1C0oa/dgyh80uOTppWzaVR9T5/WAblLjDvCDJ5bwzk8nNWl62pbWldXuLymbWHrXoWJKvzKCnF67FVtKEOcFF9D1D3/AkpMDwA2lVSx7ZnnE/PfW9DQSzatq9FO0L/ykHE+aDMDK0mq8AT2usr9Liiq5+bXVCc/TOpRNU9FUhX9cMiLhwhdCiOOH06phUZWII9nxZjTkpDSf7laQk0z/zimsKasNey6W85PdonLl5F5hN5FVVaFHposemS5m9G865aDG7adoXyNF+xoo2tfAmyt3UryvgZqaSpZZNhEtiXl0V43p+RYeXOzjvpObyW6p+Ba8teCIHkQK0VoSTLVg6956fvNu9Car0XgCOm9/vZOJBVmc04oSocEqD2aUwKR3Zg9UReWWD/7AuQNOYUS3waQ7It+pDpTVY5pmu16E3/fBRjbvji2QOpTbH+TSp5ey8M6TsWhtX2Dyya+KIs5NiTW963AtpV9ZbDZu/N1cOuWEB4cje2aQk2yntCpytaOW1m2a0Xsaieiq3X5smkpAb/o5iLenk1VTqPMEYg6mgrrBdf9emfA8LQCLqmC3aqgKzBmfx5wJeU1KrAshOi5FUThneFfmf70zLKCKZ8RoQOeUmFLqf35Gf65+oTDiMavF1GNN5eKxPePYO0h32RiVZ2NUXtOqpP59JWhP2CEYvTLg72bYmfRsIzeNaybtXbOFmvu2QzDlC+p8vH43T35ZzI4qN96ggd2i0iPTxdVTenH20C5HfG64ODokmGrB01+VENAjBwctXQx5AjqP/m9Lq4Ipwx0MVS6IIMWexFuX/J1/LnuZOz5+gH2NVcwoGMdfzryDnKTDuoErCqZfR0lkEmgMGnxBXi/cETXobO5nZZpQ7w3yv817OX1w5zbdrlp3gP9u2BOWJhFPelc8nFaNmcO6UBAhkILQifFvPzyJHz21NO4APVSKe0hMPY1EU5HmSkH8E6uBqCk3kXy6aU/E40esgbymwI/G9WRiQRYnD+gUVolLCNHx/WRyL95fWx5xdCqWEaMku8Z10wtieq8pfXO46ZS+PPLZ1rgaBDutGs9fOYbsKP3z4mWzKN/10IxiSK7GzH4W/rzQz8CcaMfGOA7YMdINk4c/3cKzi0pCLVYOyTTxBw02ltfxy/nr+eXbG7hiYh63nNZf+kJ2cHJV1oxGX5C3V+8kUiwV68XQ9io363fWMqRb83dFdMOkssHH7jovu2u97KnzsrvOi1HeyGy/TrT70H2z8/nb9+4GYGtlKTe+fx+//exR/nHubw5bv8GFTy4hNcVBdrKNrGQ7WUk2clLsZCXZD1bzyXBZExodevvrMtQoo16x/Kwa/TqPf1HU5sHUwq0VWDSFw4sUxZveFQunVWVUXgZ/mj202eVG9szgn5eM5Pp/r4o5oHJYVW49rR/nj5Qy2IlIc9oiTuCOt6hIIGjGVT3vsQXhTZrjCeRNoKLex5lD2q6gjRDi+DKwSyoFOcls2lUX8WZOLCNGp8Yx1/O66X2wWzT+8p/NBIJms4WZbBYVm6by/JVjGJWXGXW5uDkzQfe3uNi90x2MfKKB2yZEDuL0oJ9Pi/2c1NdLp9TWF7vyBXWuer6Qwm3VzQabBwKsZxaWsHpHLU9fPjqu9HBxfJFgqhkfrtsVMUCI52IoEDT51+Jt3DCjD7sOCZIODZj21HrZ1+Aj1WGlU6qDzmn7/6Q6yOuViX2rh1jurvTJyuPCIWfy0up3w55TNZV7zhtMRYOfyv1lUcuqPawpq6WywRcqldrgp9YTINVpDQVchwRZhwZgB6oAZSXbDpY7ffqrklb1wQDYUF5HWbWb7hmuFve1OaauY3q9GF4v+8r3okdIO4w3vas5Vk1BVRTOO6kb980aElMwevKATrx89XhufnV1xIqHByTZNOxWjftmDebsoW3cUf4Ekp1so3Oag+2HpVfGkyYDMKBLCs4IjXYjqWzwsWlXfdjj8QTyhgn/3bgHwzBR5c6mECesf/xoJDMfXRi1fHk0DqvKU5eNjvsm6ZWTezG+dxZPfVUcuhZSFTyHnOOT7BoWVeWyCaHU4zavyutIhdxBsHtts4v1yVS5aLCVR5b7GZobvo9VSb15Y30td33wJUl2C2PyMxmVl8Ho/Az65abEdVw1DJMbXv6aFSVVeGOczuAJGBRuq+KnL6/iqTmj5TjeQUkw1YwdVe6IAUI8F0O6aTJvZRnLSirpnOoIBUupoWpso/Iy6JIWeiw31R41t3Z3YSXBivCGdVsrS/msaAnnDjiZLqm5lNft4Z1NnzGy62HzfRRwDcpiRM+MsHWEba9hUu32HwyuDjTzq2zwsWN7Tej/jf6DAZiCQlayjZ3VkRvqxfOzspoGm97+Dy67B9PjxfB5m//b6z0YNB36txkMojgcqA4HlT3HYfScBlrT0YR407tS7Baykm3srvNi3X9SMkwTBYWLx/bgsgn59MiMLwgc2TODL26fzqrt1TzxZTH/27QXk1Bmg2GajMrL4NppBUzvnyspAq2kKApzp/XmDx9sCvtOx5ImA6GLh2unxZYqA1DZ6MdqUTj8EBJvIK8ooTTXAWIcAAAgAElEQVTYA6X2hRAnnvzsJF65ejyXPL2UBl8wpnRjp1Xj7z8awZj8xEaMBnVN5W8XncRvzx3Mh+t2saPKTYMvSKbLxoAuqZwyMPfg+bBdTL4Z3r0xakW/A349zc6LayM0srclk3PGnTw9dAyGYVJc0UDhtmoKS6t5ZmEJlQ0+RuZlMDovg1F5mZzUI73Zm2UfrNvFoq0VMQdSB3iDBkuKKnl3TbnMee6gJJhqRn2ULt7xXgwN7JLKhzdNSXg7UqZ3p+bdIkx/0y9wks3F6vJNPLXidep8DaTakzm1YAL3zLi+yXKKRSVlamzpYZqq7B+Jajnv2TRN3H6dygY/0x74POIy8fyszGCAqm+24nO6Ue0OFKcDNSkJS3YWit2B6nSE/70/aDr0b8VmO1hoo++acqzz1uI/7Io23vSu3FQ7n946jdJKN5WNPgJ6KN2rd05SqyaYKorCqLxMnpyTiWmaeAMGummSZNOkYl8bm3VSN37//saIz7WUJgOgKgpnxJGGGtANlAhJ//EG8qqiNOlpJ4Q4MQ3tnsYHN07hV2+vZ3FxJRCao3Moi6qgqQqDuqTyu/OGMLR76wsvpDmtcReWaBMDzoH3bg57eNvNTQtt9UhT8f4yQnsVRYGB5wKhioJ9clPok5vCD/fvy756HytLq1lZWsVf/rOZzbvq6dc5hdEHAqz8jCYjbo8tKIp4gx1ankPv9us89kWRBFMdlARTzUh3WVEIT7CL92Io1dm6H7NzWA417xaFPd4lJYfHZt3b4uu1NDvW7m1fdlxRFJLsFpLsFmwWNWIVv3h+VqrTSc85V9KlIKvNtnFiQVZYzwyIL73LblE5e2gXFEUhPzuJ/OykNtu+QymKEnMKmYhfkt3CPWcP5I8fboq735PDqvL784Zg1RRWllbx5JfFLCuuwh3Q0RSFNJeV74/sxmUT8g/m5ac6rK0uZwyhoKy1xxAhRMfQI9PFv64cy546Ly8tLeXt1Tup8wQxTJMUu4VTBnbiikn59I5SBOm4YrHBuY/C/GshGDn7JfprnXDOo6F1RJGTYufMIZ05c0joJpk3oLNmRw2FpdW8sbKMu95aR5rTyuj8DLqnuyK214DY59CXVjayoby22V5f4vgkZ+hm9OuUgsumhTX6jOdiyKopDG2h+ERLVJtGxg/6UfX6txDnRaBiVcm8eEC7j3L0yHSxdW/4gSaen5U/aJCf3br5UofLTrYztV8On27aEzYnKdb0LoBLx+e16XaJo2POhHzKa7z8a/G2mCtVKcCIHumkOa1Mf2AB+yLMcfPU6jz9VQlPf1XChIIs/jR7KF3TnTisatj7xDtPq1d260Y/hRAdT6dUB7ed3p/bTu9/tDelfQ2eBfW74NN7Yw+oLE445dcw5Py43sph1RjXO4txvUM3dA3DZOu+UGrgc4tLIt4wjmdeuD9oMH/VTgmmOiAJpppx6qBOEScLxnMxpCoKl03Ib/W2uIbmYDQGqf2gGDPGgEqxqWTNGYStW/vfobpmSm9++96GsCHweH5WI3qmx9U/xzBM6r1BGvxBXFaNVKc14tyiuVN7s3BLRcSL55bSuxQlNLrVFlWAxLHhzrMG0DnNzp8+2oyqKFHTNlw2DdOEO87sz6P/28LVLxRGHOU84MCJ9qtv93H2w1/x+twJ/HhiL/65YGvYSTjWQN5l07huep9W7rEQQhzHxl8HKV3gvRvB0KPPobIlg6rBOY+EgrBWUlWFfp1S6NcphSVFFWzZ07rG74YJ5bVxjrCJ44IEU82waipzxufx9MKSsLzkWC+GTuqRHndhgmiSx3dBS7dT/dYWTG8wbA4VAEpojpSW6SDzwv5HJJACOGd4V37z7oaIz8Xys3IGfVwSLMUMjEaxNj/RflethxeXlPLS0lI8AR2LqqIbJpqqcMGo7lw5uRe9DknFG5WXwZS+2XyxeQ++OKeeuGwav5w5KL4XiWPe5RN78f1RPZi/qozHvyimosF3cCJ1QDfokubguukFnDO8K59u3EOjT282kDqUbkK1O8CFTyzhvllDwo4dB8QyTwtg5jApiy6EOMENngUDvgeb34eFD8Ge9aGGvBAqod5pMEy+BQbMDCs41RYijUpB/HPoo50PxPFNgqkWXD4xn+eXbMMfoRZFSxdDDqvKzaf2a9PtcQ7IxPGLsfiKaqn/sgxfcU3o6s0MpfQ5hmSRMrn7EQuiDm6XTWPOhDxeXFIa9wiQqkBWehJDV35I8bvPk3vH7STPmBGWmujx69z2xho+27QHk+8OSgF9//vp8Ory7bxeuONgL6eMpNDB9l7fWq6u9rM1owfeGI9lLpvGcz8eG7UBrzi+JdstzJmQz6Xj89hV66XGHUBRIMNlo1OqHUVRqHUHuGPe2ogn0pYmHNe4A9z86momFGTx9fbquOdpOa0at5zWV3qTCCEEhIKkweeH/vgbwVMdetyRDvb2PU9nRSnKFe8c+gPXJKJjkWCqBZ1SHTx92Riu+NfymBusQuhC6Odn9GNCGxZTOEBRFBx90nH0SQfA1A1QFJSjXD77jjP6s66sllXbq6PexTmcokCKw8or10+mR+YZNHz1FXvuv5+q51+g05134BgUGhWq8wb4wWNL2FbZ2Oy6A4YJhklhaRVnPvwl8y4bjvnnewnu3cfrf/0/fr20gvfWlDcJxg6XZAulDD5z+RgGdY1QIUh0KIqi0DXdSdf08BTT1wt3RKjHF9uEY3P/uh/4wTB++84GFm6tjHmeltOq8f2R3bhqcu9W7JkQQnRQtqTQnyNkUp8s3l2zM6wJezzzwl02jcl9sttzM8VR0o4NAjqOCQVZPPvjMSTZNGwt9FTQVAWHVeXuswfwkyN0IaRo6lEPpAAsmspzV4xhUp9sXDFUpbNbVHKS7cy/fuLBVMjkKVPo/fbbpJ51FtvnzqX8F3fjLt/Fj59dTklF84HUoQK6SUW9jx/89b/4O3Uj/6UXcfXoxoM/GM4Xt8/g6im9SHVYcFo1ku0WkmwaVk1hSt9snpgzmkV3niyB1AnOMEye+qo4bETpwITjzNOuw9V/IqrNgaJZcPUZF1Y8AgX+vXQ7j88ZzeyR3XBYVCzNfFdtmoLdojJ3Wm9+P2uIlMcXQohjwOmDOqNGOB4fOi/c/e0SjIAXUw/iKSqk+vNnmyzr9uv8Z8MelhVXYh5eEUsc15TmfqGjR482CwsLj+DmHNt21Xp4fvE2Xlq6HROzyR0Kp1XDME2+N6wLV0/pzcAuJ+6FuGGYfLh+F48tKKJoXwOBoIF+yMcsyaZht2pcOSmUYpXuijzsrTc0UPnkU7y5YCOPDjoXT4TYv6VUK5sK18/oy82nhadbBnWDPfU+6jwB7BaV7BQ7qY6O3xhVUZSVpmmOPtrb0VrtfXxaV1bLD59cElbN01O8kr1v3kvPn8+PKa0j9//Zu+/wqMq08ePfc860TCaTQhohoYaWAErvAta1YwFfRdeGBSvr7v62u++6677qurvq6tqxACo2xLKsHaX3Jr0FCC0J6Zk+5/z+GIGEmSQzIQkB7s91caFnzsx5opmZcz/P/dx3gpVlvzsfgJ3F1by2sID3VxZiUpU6bRcUBW4c2ombhneKuEomxJngdPh8knun09MT/93MK/N3Rez7V73hW6pWzMF/eG+dfeG27N5AqLLzDUM70SnFzsylu9FUhUlDO3HVgA5Nvu8IBHW+2lTE64t2sbfUhcevY7do5GU5uWN0VwZ2SpYJuWbU0GeTBFNN4AvofLHxIFsOVlHm8uG0mclJsXNpv/ZnxM14LLYcrOKTtfs5UOHGF9BJdVgZmZvKuF7pESvvRXLBE1+zrdQTdry+VCvv3g11VgiS7GZW/O58TC3Zqf0UcjrcrEDLfz59s/kQD76zJqx5d/WGbyn79lVy7psR1etYTSpb/nJxnWNuX5A1e8upcPtQFIVku4Wzc5KwmOR3VJzZTofPJ7l3Oj0VV3m54J/fUe7yx/zcxDgzX/7sHNKdNgzDYPHOw8xcsof524q5tF97Jg3tRJ8o2+gEgjrPz9vBqwt24Q/qYRN+ihKa4G8Xb+GXF/XkirOlUXBzaOizSfZMNYHFpHJZvywu63eyR9L29cxMoGdm0/tg/LCvgsKq8A+uWHo7+AM632wu4sL8zCaPQ5x5/EEjrDcZxL7hOFIVwDiL1iL7KYUQQrSMtAQrb00exoQXF+HyBokmUU8htFdq5uShpP/YYkVRFEZ0S2VEt1SKKj3MWr6XO99cQbrTxo3DOnFZv/b1Fh5y+4JMfmM5qxooamQYoZRCl8/Nrz5Yz7rCCn53aW9ZpWpBMg0q2rSvNx3CGwjftB9Lb4caX5BP1u1vieGJ05jTZibSd0/tDcfRsEs1PiGEOC3kZTmZc+9IUhOsxFsb/myPt2i0c1iYfe/Ieled0p027j+vO/N/dS73jsvlk7X7GfHYN/zl043sKqmpc25QN7hz+gpW7I6+OqzbH2Tm0j089dW26H5A0SSyMiXatIOVXiK194m1t0Nxla+ZRyZOd/kdnBErPsbSiFoh1OdMCCHE6SE3PYGFvzqXLzYePLo33Kyp6IaBqij4gzpd0xxMGdONi/Izo0rf1lSFC/IyuCAvgz2HXcxctptrn19E7/ZObhzWkfN7Z/DW0j2sKIi+WvIRbn+Ql77fybhe6Zydk9TUH1s0QIIpcUqKNdVKKueIWDltZi7p256P1+wneNzvT7RNu+MsGneOkfLmQghxOjm23SOLHcXV7D5cQ5UngMNqolO7eHLTm973qmM7O7+5uDcPXdCDuesP8sr8XTw85wfcPr3e9hqNFePyBoK8/P1Onps0oMnjEvWTYEq0aWkOC6pC2OpULL0dAFLrabgnREMmj+7C3B8OEPSHB+ONNe2G0Kbj4V1lb5QQQpyuuqU56JbW/E2DrSaN8f07ML5/B95dsYfffvhDxPOi6XuoG/DVpkOU1fikcXALkD1Tok07t3cGVlP4ylMsvR3iLRqX9G3fWkMWp5H8rERG5qZibUKVPZtZ5Y+X58mmXyGEECdkwbbDBCPseYil76GqwGfrD7TWkM8osjIl2rSzshPJTLSFbcSE6FOtNFXhwvyM1hqyOM08d8MAJrywiG2HqvFEmaseZ1Z58Pzu/KSPBPFCCCFOTGGZK2L1wFiKcbn9OocqwtvMiBMnwZRo0xRFYcqYbvzx4w0Rc4UbS7WyaAo3DuuEWXpMiSaymTXeu3sE98xcxeIdh/EGghGLokBo5s+sqfzx8nz+Z0jH1h2oEEKI05Knnup9sRbjqvEFGj9JxEzuMEWbd2X/LLqkxmPWYkuXUhVIcVi58xwpACBOjM2sMe2Wwbx713Au75eF1aSSYDXhsGokWE3EWzWS7WZsJpWZk4eEAinDAG81+FxEbFglhBBCRCExzhzxeO1iXI1RFUiR/VItQlamRJtnNYUa3o1/biF7SiMvdR/PpCo448y8e+dwkuzy4SGaR9/sRJ6+vj8VLj/r91VQ4fZj1hTaOaycnZPEs19vZe282QxSPoaCBaFW9BiAArnnw4gHoNMIIjawEkIIISIY1rUdq/aEl0WPpRhXnFmjX7aURm8JEkyJU0JyvIVJwzryzNfb8AcNgrpBIEKulaYoWEwquenxvHrz4KMdx4VoTol2M6O6p9Y9uPUL7l9zH+6aCgw8KECdyH/r57BrPthT4OqXQkGVEEII0Yjrh+bw3LztYcdj6Xtot5oYlZsa9hrixEkwJU4JWw5W8cJ3O/nk/tEYhsG0hbv4YOU+IFRgQjdCwdXFfTK5Y3TXeruNC9EiVrwO//01asBNfL0nGeCvgYoamH41XPUC5I9vvTEKIYQ4JaUn2BjdPZVvNhWFZedEU4zLZla5Y3QXVFWyIlqCBFOizfP4gzzw9mp+fXEvuqSGblX/Mr4vf7gsj+IqL9XeAPEWE6kOK3GW6DZhCtFsNn8G//01BNzRPyfghtl3Q3wqdG68T5oQQogz20MX9GDh9pKIxSgaK8YV1A1pEdOCpACFaPMem7uZ3AwHEwZm1zluNWlkJ9vplekkJ8UugZRofQEvzL4rLJDq/FQV6X+rosZ3bA7xlVU+xr5eq8R/wA0f3A56dOXWhRBCnLnysxJ56rqzsZmjv3VXCPXavLhPJje9uoxth6paboBnMAmmRJv2zeZDfLnxEH8d31ean4q2Z+PH9VbqCxrw9FJfw8/3VsOuec0/LiGEEKedn/Rpz/M3DiTOrGFrpJm83aLRzmFhzn0jeeb6AdwzthvXvbSE//4gjXubmwRTos0qqvLwqw/W88/rzibRHrksqBAn1cJ/gq864kO/HGHhyUVeyj0N1J/0VcPCZ1pocEIIIU4343qms+BX47j/vO60i7cQb9WwWzSsJhW7RSPOotEtLZ4/XZHPgl+dS256AgATBuXw+q2D+fOnm/jb55sJ1tcwUcRM9kyJNknXDX7+7lquH5zDkC4pJ3s4QoSrLoaSbfU+PChLY2xnE08u8vKXcxuoKlnwPQT9oMmEgRBCiMa1c1i5d1wud4/pxtKdhyksc+PyBXDYzPTMSKBvduQiXP2yk5hz30jue2sVt72+nGf+p79MVjcDCaZEmzRt4S5qvAEeOK/7yR6KEJG5DoNmgWD9qXyPjLMycloNDw5toNeZagJ3OTjSWmCQQgghTleaqjAixnLnqQ4rM24fyv/N3cwVzy3gxZsG0ivTGfHckmov763Yy8YDVVS6/ThtJnpmJjBxUI60nqlFginR+kp3QsFC8JSHbiTj00INTeNCzeQ27K/g3/N2MOfekZg0yUQVbZQRbLT5bp90jct6mHhsgY/eafX9Liuh1xJCCCFagUlT+cNlefTtkMgNLy/lkSvzuaxf1tHH1+4t57lvt/Pd1mKAOs2Cv9h4iH99s52R3VK599xcBnZKbvXxtzUSTInWoQdh2xew4Ck4sAYUFfQAoITSm/QA5I3HM3gKD7xbzsOX5ZGTYj/ZoxaifnHJofS8RvxprI0BL1bz8+HWyCcE/WCTrvRCCCFa1/j+Heie4eDuGStZX1jBLy/qyawVe/nzpxvxBvSI9ZWOBFbfbili8c7D/OKiHtw+qmsrj7xtkWBKtDx3Ocy4Goo3g68m/PGgN/T3+vdQ18/md8njOfesF1p3jELEKqE92FOgcn+Dp+WmqFyXb+aZZT76pkdYncrIA7OkSwghhGh9+VmJfHzvKO5/ezUXPfU9+8rceAKNt+wwALc/yJOfb8UwYPLoMzegkhwq0bI8lfDyuXBwfeRAqjYjiMXwMq7qE/j4/npLTgvRJigKjHgQzI2voD48xlqn59RRFgeMnNoCgxNCCCGikxxv4ZcX9WD3YVdUgVRtbn+Qv3+xlZW7S1todG2frEyJlvX29VBR2OAm/eMpfhds+BDSe8OI+1pwcEKcoLOvh6/+GHa4YGpCnX/PSVTx/D7CBl9Fgd5XtNTohBBCiKi89P0u9HomsWs2zqNy+Uf4DxeiWuIwp3clccREbNn5AHj8QZ79Zjuv3TqkNYfcZsjKlGg5+1bB/lXH0vh+1PmpKtL/VlVnpv6VVT7Gvl5r5crvgu8eh0D0QZgQrc6WCGN/E9Xq1PEMUxz85HEwNVDpTwghhGhhZTU+vtp0iEitpyqXzab065dJHDaR7Ptm0GHKayQMuAT3tqVHzzGAhTsOc6jS03qDbkMkmBItZ/FzEIj8xgoa8PTSRgIlIwibP22BgQnRjEY+CGddH1NA5VNtvG2+isreE1twYEIIIUTjZq8ujFicVvfWUL5gJikXTMHecwSqxYaimbDnDiV53G11zlWAWcv3ts6A2xhJ8xMtw10eCoSMyLm3vxxh4YmFXu4ZbCHJVk95aV9NqPpfn6tbcKBCnCBFgUv/DonZ8N1jgAoBd+RzzXYwdMw/eYxNhYP58LXlvHHbEOKtxz6KDcNg1Z5yFu8ooaTah0lVSHda+Ul+ezq2kwqXbVWFt4KPd3zM5tLNVPmqcJgd5CblcmXulbSLa3eyhyeEEPXacrAajz/8fs27bzNGwIe9x/BGX8Mb0NlysKolhtfmSTAlWkbJ1lBD03pWpgZlaYztbOLJRV7+cm4DlcxKNrfQAIVoRooCox+CgbfAqjdDq7K+6lAfNQDdD/Z2MOIBOOt6FJuTPw0w+PWH65j8xgpeu3UwAB+uKuTF73ZSXO3F4w8eTbkwawp//2Ir/bITuXtMN87tlY7SSI8r0To2l27mlfWvMG/vPAC8tdKarZqV59Y8x8gOI5ncdzL90vqdpFEKIUT9KjyR23wE3ZWodieKqkX1OpX1vM7pToIp0TI8lY2e8sg4KyOn1fDg0Ab2jAS8oR5VUb6RhTip7CkwamooaCrdCe6yUE81ezIkd6nT5FdVFf7v6n489O4abnltGYcqvRys8OD2hzfw9QcNwGB5QRkb9q/mnO5pPHN9fywmydQ+meZsn8NflvwFX9CHToRZ3R8Dq3l757F4/2IeHPggN/a+sbWHKYQQDUqMM0c8rsU50V2VGHowqoCqvtc53ck3sWgZpnoalNbSJ13jsh4mHlvQwN4pVQvdjApxKlFVSM2FnMGQPRBSuhIpIV1TFX57SS9W7ymnoKQmYiB1PJcvyLytRdz2+nKCkXYLi1ZxJJDyBD0RA6naDAw8QQ9Pr3ya6Runt9IIhRAiOr0yE4gzh99rWTv0QjGZcW1d3OhrWE0qvTITGj3vdCR3qaJlOLMg2Phy75/G2nh5lY99lfXcFMalRLwJFeJ0YBgGt7+xAl03iCUs8vh1Vu4u5bG5m1psbKJ+W8u2Hg2kYuEJenhm1TOsKVrTQiMTQojYXdW/Q8RKfqo1nqRRkyj98gVcWxej+z0YwQDuHSso+3ZanXMNYOLgnNYZcBsjaX6iZbTrBildoGhjg6flpqhcl2/mmWU++qYfF9trVhh4awsOUoiTa9WecnYW1+CP8C3WWF8Pt19n+pLdTD2/R50CFs0iGIDKfeCtDL0PHWkQl9y81ziFTVs/DZ8evqJeNr+Mks9L8BX50GwazoFOMq7NQIs/lh7jDXp5cd2LPH/+8605ZCGEqFeS3cKF+Rl8tu5AWFDlHHI1anwyFYtnUfLpkyiWOKwZuTiHX3f0HAUY3T2V9IQG9sCfxiSYEi1n1M/g05+FNuI34OExVqavq2cVa9BtkY8LcRp4+fudeCKk9lUum03F0vdpd+G92LoMQNFMuHetxL1t6dFgCkBVFD5as49JQzs1z4AqD8DyV2D5y6GVZVUDwwg13c4ZCiOnQrdzQ2mMZ6hKXyVf7fkK/bhKpSVzSyieW0z25GwceQ78ZX72T99PwZMFdPldF9Qf97cZGCw7uIwiVxHp9vST8SMIIUSYKWNy+XLjoYhV/Rz543Dkj6v3uVazyn3jcltyeG3amfuNKFpe3pURC0cUTE3g/K7H4vicRBXP753MuyX+2EmqGbqNA2f71hipEK2u3OXjmy1FYbOAsfT1cPmCvPT9zhMfTDAAHz8Az5wFi/4FnopQ42xvVWgyJOiDgvnw3s3wzzzYf+amqX224zPU4746g+4gRR8VkXVjFgn9ElBMCpY0Czn35OAr8VGxqKLuixjw4bYPW3HUQgjRsLwsJ49c0QdbhL1TDbGZVX5zcS/6dzxzsxckmBItx2SF62ehm+Jie56iQnwqjJc0GHFqCZR58Gwtw7W+GPeWUvzFrnrP3X3YhTVCNb5Y+noAFJa6MYwTKEQR8MH0K2H9u6HqmbVKe4fxVUPVAXjtYtg1v+nXjHV8rtKo9mC2hm3l23AH6/YRc21zoft1nAOddY5rNo2EfglUb6i7Ou/TfWwt29riYxVCiFhMHJzDn68MBVRqFNvVbSaV313Sm5tHdGn5wbVhkuYnWlRF+iCesPw//ld/HLMexWZtzQKODLj1P6Ey00K0cYZu4NlSStV3hfgKq1FMSmgnrgIEDUypcSSMzSYuPxWlVvBU7Q1EfL1Y+3oYGPiCOlZTE9oHGAZ8eAcUrqy/0XAkfhe8fR1M/hrSe8d+3cZUHYIV02D5S6EG4KoJ9ECoV9fQu2HAzaF9XCdBlS+8KWWwOojJYULRwu8+TIkm3LvD/9tWehtvHyGEEK1twqAc8rKcvDBvB19sPISiUCf1z2ZSMYDMRBu9MxO4aXjnkzbWtkKCKdFiPP4gd765grzeF2EadD58+TDsXgT8uAejNnN86Hi/6+D8P8pmd3FKCJR6KH55PXqND8MX+rIxjouR/AdqKPtwO+Uf7yT19j5YshwAxFkiBz+x9vUwDPh+SzE9M51kJ8ehRjOdeMS+lbDtizqBVOenqnD5YdeDDuItodd6ZZWPGev8dVNxfS6Y+yu4+ePor9cYXw3MuRc2/ydUxfNI0+8jnxc1xfD93+D7J6D3lXDFM2COceX7BCVYwkv/ag6NQHUAI2iEBVSBigAmR/hXbaTXEUKItiA/K5F/3TCAcpePD1ftY9PBSipcfhLjzHTPcHDNgGw0VeG8v3/H1kNV9Mg4sz/PJJgSLULXDX7+7lpSHVb+cGkeiqrATz+CikJY9kroBs5bAYoWSunr/1PoNwEs8Y2/uBBtQKDEzaHn1mB4AjRW19zwBjG8QYpfWEvq5L5YOzrpkBSHLxC+0bd2X4/4XqMaHYfNrDF96R62Haqiwu2nW5qD7hkOemQk0DMjge4ZDjokxaFEajGw6F/HApZaggY8vdTHb0c31C/OgL1LQu/pxOxGx9koVylMuwjKdzecanhkvJs+DlULvfU/YEs88etHKTcpF5tmq1MW3Z5rRzEpVK6sJHHIsbEEPUGq1lWRcW1GndewqBZ6JPdotTELIURTJNkt3Daq/hS+e8bl8tf/bOL1W4e04qjaHgmmRLMzDIM/f7aRkmovb9w2pO5MeWI2XPC/oT9CnKJ0d4CiF9dFFUjVZvh0Sqb9QMbUAWQk2eiXbmL5/rqrtLX7eiiqhq1LfxTVhKdgDZ496+oUobCaVO4Z1437z+Zh4X8AACAASURBVO0OQKXHz7ZD1Ww7VMXWQ9Us2FbC1kNV1HgD5GYk0CM9FGT1yEygp9NHxta5KEZ4QPfLERaeWOjlnsEWkmwNrHQZBix7GS74U/T/ESLxe+DN8VBWEL5qXZ+AB0q2wsxr4Zb/gGY+sTHUwzAMlu0q5fVFBewsrqEmkIgnLRBK4/yRZtdIH5/O/hn7UW1qnWp+5hQzSSOSwl736u5Xt8h4hRCitdw0rBPTFxfw/dZizulxclKv2wIJpkSze3n+ThZuL+G9u0dgMzdhH4cQbVz1sgMY7siB1LLCdfz12+fZWlKAqqp0b9eJP553P2e3D+0tMnxBqmYvJJl/cHdNHBu1m6kJ1i1EEU1fjyOuH9Lx2PNsZgZ2SmZgp7ppshUuP1uLqth6qIpth6r5ZnMROQe/5GFdxRHh5xuUpTG2s4knF3n5y7kN9A0J+mDjnBMPpla+Doe31gmkoko3DPrg4A+w9h0YcNOJjeE4hmHwzrK9PPvtdspcPty+4NH/3TZrPibnOhTl2C9A2iVpaPEaB2cdxFfkQ41TcQ5wknNXDupx1bEGZw4mI77uapUQQpxqLCaV31zSm0c/28TI3FS0WNLMTyMSTIlmNWfNPl5fWMAH94wgMa5lZoqFOJkM3aB6/j6MCCl6Vd4abn3/1zx64UNc3mscvmCAZYVrsWqWYyfp4NoKiRNuZeyNV5Dw5HxcFZ6wuKzRvh4mlfN7Z5DqaCgVLyTRbmZw5xQGd65V1GXlZoy5QOQ6GDwyzsrIaTU8ONQS+YQjTrSQgmHAomfAH16kIap0Q78LFj7VrMGUP6jzs1lr+HpTEe4IfcB8h8dgStgISt0KgyljUkgZ03DhHJtm485+dzbbWIUQ4mS6MC+DVxfs4t0Ve+tM7rUWwzBYvPMwL323k+UFpbj9QTRVwWkzc1X/Dtw8ojM5KfYWHYMEU6LZLNxewp8/3cjMycNon9i6m8KFaC3e7eUYvvAbbICdpXsBGJ93PgBxqsaYLhFyyc02XP6ROMwWpt82hPH/XkiNN/JrRmJSFbKS4nj82n6x/wBHKZH3Uf2oT7rGZT1MPLbAR++0+rtoVHiCPPzOatIcVlITrKQ5rKQlWEn98e+UeEvDs5UFC8BTHvGhqNMNK/eHiml0GFj/OVEyDINfvLeWrzZFbl4JoHuz8By8Alvmxyhq9CXbbZqNe8++lwEZA054nEII0RYoisIfLs3jtjeWc/lZWTisrRdafLHhIH/8eAMVbn+d7AE9aHC4xscbiwuYvmQ3Z+ck8eSEs1osqJJgSjSLDfsreODt1fx70gB6Zp7ZVV3E6c2zo/xo5b7jdU3JQVVUfvbZo1zR6zz6d8gnyRb+fjB8Op6tZTiGtqd7RgLv3DGcG19dQo03SOD4Lr7HsZpUOrWz89Ydw07sS8veLmJT7dr+NNbGgBer+fnw+leGzPHJjOmRRkm1l+IqL5sPVFJS7aO4yktxtZdKt58ku4W0hCNB1o///GOwNXTDTDJ8LiKFSlGnG/rdsOGjZgmmPl67ny821B9IHRGoGIwHsGV+DEqgTsrf8RQUrJqV+/vfz0/zf3rCYxRCiLakb3Yio7un8sK8HVw/tCPTFxfw9aYiKj1+VEUhJd7ChIHZXDMwmwRb82QtvTJ/J09+saXBz2p/0AAMlheUcukz83n7zmHkZzV/wSIJpkS9fAGdbzYfYmdJDVWeAA6ria6p8ZzXOwNLrX45e0td3P76Ch65sg9Du7Y7iSMWouXpNfWvRCRY4/lw0rP8e+lb/L///o3imlLGdRvKEz/5f6TF103/0l3HXqdvdiJzHzyHf32zndmrC1EVBddxq1/xFg2LSeXWkV24Y3TXekurR63L6EYb4eamqFyXb+aZZT76pkdYndKs2Af+D1cPqL+anz+oU1pzLLgqrvJSUu1lX7mbNXvL6bBnF5kNVPGILt3QCDUTbgbPfrM9YmofQM3GeVQu/wj/4UJUSxzm9K4kjz0fZ/+9mBybAVDUY3mTVs2KYRgMyxrGHX3v4Oz0s5tljEII0daMPzuL215fwUvf7wQMfMFjn+sHKjzsKtnCY//dzGX9svjtJb1JiW8khbwBH64sbDSQqk03oNIT4PqXlvDZA6ObfYVKgikR5kCFm+mLdzN9yW50Azy+AEEDNOVIbxyFSUM7cvOIzsSZNW5+bRl3jenKpf3an+yhC9HiIjVmra17amf+eelvAdh+eDcPfPoX/vfrf/HcFX887nXqBidZSXH839V9+f2lvflo9T6+3HSIMpcfk6qQnmDlmgHZjOuV3nwbfG2JkH8VrHsXjPpTDB8eY2X6ugaCroG3NngZs6aS4bSR4axnZWlmImyr//nRphuu21vKt19tO7oCVnslLNqGxusLKygsi9y8uHLZbCqWvk+7C+/F1mUAimbCvWslNRs2YEm/DUWrxpS4EnPcQUb2sJMSl0i3pG5clXsVafYzt8qVEOL0N29LEVNmrPoxsyLy5NiRCcKPVu9j/tZi3r17OJ3axd4Op9zl4zez1+ONsG850oRX4oiJ2LLzAaj2Bvjl+2t5587hMV+3IRJMiTrmbSninpmrCAQNfMG6v6hBA6p/3Nfx2sJdvLl4Nx2S4rggL4NbR9bfh0CI04mWZA3NLAQbr4me264TE/v8hBlrwhvbaomRU+firSYmDevEpGGdTnisjRp+Xyg9rlbT3oKpddMScxJVPL93hj9XUaHbOEg4wap0jsYDjcbSDQ3A5MwgoOusKywPWwWzW0x1UguP/jnu32cu3R2x95furaF8wUzaXTIVe88RR4/bc4dizx0aGkPQgb90DJpZZfSg3vx0eOcm/ecQQohTyYqCUqbMWIk7ylWigG5QXO3lmucXMffBc0hLaLyIUm3vrthLpO2+9U14ubctPRpM6Qas3lPO3lJXs65OSTAljvp2cxFTZq6MatnUFzQgGGRXSTUDcnq2wuiEaBvs/dKo/HpPxMe2H97N1zsWc0Wvc2nvTGd/5SHmbPqaAVn5dc5TLCrxA9tAaezMPnDW9bDunVBVvFhYE+Dix098DD0uDgV0vup6T2ks3VCxOMgbM4G8buGfRbpuUOH2Hw2ujv6p9rL1YFWd46U1vohzqt59mzECPuw9Gp/N9Pj1ele3hBDidOIL6Nz+xoqoA6kjdAPKXX5+NmsNMyYPjf55usHL83eF3adGM+F17NoGbywq4PeX5cU05oZIMCUA2FVSwz0zV0Wdf3pE0ICps9bwSfpIctOl8IQ4/ZnaxWHpkIBvd3hJ8HiLnTX7N/Hy8nep9FbjtDo4v9twfjfunjrnqXYzli4RVntOhkufBFcJbP8quoBKUcESDz+dA8mdT/z6PX4SVcPdBtMNLQ7oMjbiQ6qqkBxvITneQo+Mhj+jrnx2AWsLK8KOB92VqHYnSiMFO46o9tRTb14IIU4j/91wkEAw8n1jYyl3AT1UGKKwzEV2cnSrRBsPVFLjDf98jWXCyx80mL16nwRTovm9+N0OfMHoN13XfkP4gjrPz9vB3yfK5mpxZkgYm03p25vDqvq1T0jj+fENN7BVzCqOc7IbLEveqlQNJrwB3z4Ki58NBUuRgipFA5M1FEBNnA6puc1zfc0EQ+6Chf+EgPfo4ajTDU1xMOI+UOvfTxWtRHvkDdFanBPdVYmhB6MKqNqdwMZqIYQ4Vbwwbwc1EVqFRJNyB6FWFNMX7+Y3l/SO6nol1d6I+4ZjnfCq9ETf0iIaEkwJarwB5qzZT6TJhWjeEEHd4NN1B/jjFfk4m6nkpRBtma1XCrZeKbg3lUIsq7kmBXMHB46hmS03uKZQVTjvDzDyQVg3CxY+HerdpJnB+PHny7sSht8LWf2b//rD7obVb0LVwWPXi4aigSMDBt7SLMMY1CmZpTsPh21stnbohWIy49q6mPheoxp8DbtFIy+r+VYdD1Z42F/hxu0L4rCa6JwaLw3RhRAn3d5SFzuKw9OzY0m58wUNZi3fy28u6Y2uG1R5ApS5fJS7/aG/XT7Kavyhv11+Nh+sDKt0C7FPeAUbaUESKwmmBB+t2RdxM18sbwhVUfhwZSG3SCEKcQZQFIWU63pyeMamUBPfaAIqs4o5M57UW/PDKvm1GTYnDLkDBk8O9W7yVIRWo2yJjfakOiFxyXDLZ/DK+aFr6lGkyakmiEuBWz8L7d9qBtcP6chz324Pv5Q1nqRRkyj98gUUVcPWpT+KasJTsAbPnnUkj7vt6LkuX5B3V+xFVRXO7ZWOuQn/r4O6wbebi3jhux2s31dxrBWFEcoEuCAvgztGd+WsnKQm/6xCCHEiDlV6sJjUsMmnWFLuAMrdfvo/8gWVngB2s0ZSvJlku4Uku4Vk+5F/NtMtLZ5ku5kf9lXi1usGVLFMeAEn3lrkOBJMCdbtrYgY6cfyhnD7gxH3GghxulI0lXY35VH1zR6q5u8Dw4jYzFexqGCAfXAGSZd0RTG10UCqNkUBiz30p7WkdIW7F8CMa6CsIBTMRSoHoahgskFqd5j0PjjSm20IaQlWzumRxlcbD4Vd2TnkatT4ZCoWz6Lk0ydRLHFYM3JxDr/u6DkWk8ptIzvTLc3Bq/N38bvZP3D1gA5MHJRDbrojqjFsOVjFT19dSrUvQM2P1VOPv1n5z/oDfL2piN7tE5h2y2CS6klPFEKIlhKpNDnEnnKnKPCfB0eT6rA2OvlU7Q3w0vydYcdjmfBSgMGdU8Je40RIMCUoc/kiHo/1DVFez+sIcbpSVAXn+Z1IGJuD+4cSqr4vJHDYg+HXUUwKWqIVx6gO2PunozbzTNhpyZkFUxbB3mWw6BnY9mUo1VBRQ+l/uh96XhraI9VhYIsM4cHzujN/W3HEYjyO/HE48sfV+1yLpnLbyC6kO21MGJTDzuJq3l1RyA0vLyEnxc7EQdlc2i8LhzXyV++aveXc8PKSiJNbtelGaAJr/b4KLnl6Ph/fP4pUR2zlhYUQ4kQk2EwR57tiTbnTFIX2iXFRXdNhNXHl2R14f+XesK0p0Ux4QWhV6q5zukV1vWhJMCWIr+eLPdY3RH2vI8TpTjGp2M9Ox352862SnLEUBToOhY4zwV0e2rvlqw6l8jk7hFIRW1CfDok8cU0//t8H62Kqbhpn1njjtiGk12pO3DXNwa8v7sUvLuzBvC3FvLtiL49+tomL8jO5bnAOAzslHy1EsrfUxU2vLm00kKrNHzQoqvJyw8tL+OT+UVE3JxZCiBPVNc2BP8Jm+1hT7qJdtT/i9lFdmLNmH0E99gkvgCS7mWFdZWVKNLMuqfFYNDWsSW8sbwizptAlNfZO1kIIUa+4pNCfVnbF2R2wmDR+NmsNQV0P9dWrh82sYtFU3rhtCP07Jkc8x6SpnJ+Xwfl5GRRXeflwVSH/74N1AEwclMPVAzrw1FdbI5b8hYYrqgZ0g8IyN5+tO8DVA7JP/IcXQogoOKwmLjsri9mr99Up6BBLyl28RWPK2NhWiXpkJHBe7wy+3nQo5nY+NrPK/16e3+zVdCWYElw7MJtnT3DTtaooTByU05rDFkKIFvOTPpn0yx7Dm4t3M3PpbnQDfIEggaCBWVMxmxTsZhO3j+7C/wzOiXrfUlqClbvGdOPOc7qyak8Zs5bv5bwn51HjCxKpwFQ0FVVdviDPz9shwZQQolVNHt2FT9ftD6uOF23KnaIo/KRP7NVt/znxbK5/eQkb9ldEHVDZzCq/uLAnF+Y3fzVdCaYEDpuJ9AQrhWXusMeifUOcnZNETkorblYXQogWlpUUx68v7sVDF/Tgm81FFJa5cP1Yorx7hoOR3VJRI/Q8iYaiKAzslMLATil0S3Pw5Bdb0I9bAYulomphmZsf9lXQp0Nik8YjhBCx6pXpZGiXdiyJ0FKisZS7OLPG/eflNik92WJSeeuOoTz07lq+3ngIv65HbO8Tuo6KYcBfx/fl6oEtM+EkwdQZzBfQeWvpbp79djv5WU5Kqr1N2nQdZ9a4d1wzNfAUQog2xmJSmzR7Gq3vthbjj5BKGEtF1YCus3B7SZODqaIqD0WVXryBIAk2Mx1T7NjMsgdLCNGw528cwOX/WsDO4ppI9SgiijNrXNI3kztHd23yda0mjeduGMDWQ1W8umAXc9bsw3ykebsSajHhsJq4Y3RXJgzKbtGqpxJMnYEMw+C/Pxzk8f9upmO7eKbfPpTe7Z387b+bmbawALc/+g3QcWaNm4Z14pweaS04YiGEOH01R0VVf9Co93XqE9QNvqndz0pTURTQDQNdD6WA3zqyM13TYtsgLoQ4c8SZNXplJlDm8uPxBxssoqOpCmZNYdLQjvz2kt7NsnepR0YCj1/Tjz9clsemA5VUuPyYNIVUh5W89s4mZw/EQoKpU4Q3EOSrjUXsKK6mwu3HaTPTOdXORfmZMc0ertxdxl//swmXL8ifx/dhdPdjQdAvLuqJXzeYvnh3VAFVnFnj+iE5/PriXk36mYQQQoRuMCIej7Gi6tFZ2SisL6zg1teX4fYHj/az8h2XpvP2sj28u2Iv5/RI41/X95eVKiFEmJe+38nuUhfzfzmOFXvKePG7HazcXYamKgR0AxXQNIVA0OCyfu25fVRX8rKavyqrw2pq9v5R0ZJgqo3bX+7mjUUFzFy6BwMDlzeIQajpmN2i8ZsP1zNxUA63jexCx3b171kqKKnhic83s3pPOT+/sCdX9e8Q9gWuKAq/vaQ3Azsl89SXW9l1uAZ/0KizsVBTFMwmhU4p8Uw9vzsX923fQj+5EEKcGdITbEBl2PFYKqpaNYV2jujSWBbvOMxtry9vdNIsoBsEdIPvtxZz1b8X8sGUEdgtctsgWke5y0e5y4+iQJLdQmKc+WQPSRzn2y1FvLpgFx/dO5J4m4kxPdIY0yONwjIXy3aVUuH2Y1IVUuKtjO6RitN2ev4/lE/FNuzbLUXcM3MVgaAelk9vADU/LqXOXLKbWcv38veJ/bikb1ad80prfDzz9TbmrNnH5NFd+fuEs4lrpHnoRfmZXJSfyaYDlbyxqIDNB6uo9gZwHP6BHn0GcfOYXuRnySZnIYRoDtcOzGbpzsNHP9OPiKWiqu730+sfv6d45VAc552HLS8vYgrN9qIqJr/ReCBVmzegs7O4hslvrGDG7UNbJW1GnJk8/iCfrN3PC9/tYE+pC4umYhBaNe2ZmcDdY7pxUX4mFlP0q7CiZeworuYX767lxZsGkpVUt+ludrKd7OQzpyiZBFNt1DebD3HPzFVRlXz06wZ+PchD764lqMPlZ2Xh8Qd5bWEBL8/fyWX92vPlQ2NIdVhjGkPv9k4eu6bfsQMvPgyjh4EEUkII0WwuyMuoN9Uv2oqqQ7pn0O+6h6j6+hv2PfQQhs9PwrhxJJx/HvbBg1HMoRnhRz/bVO+ehob6WXkDOmv2lrNwR0md9HAhmsus5Xv50ycbAI7+jvqDx35XN+yv5NcfruO3s9fzt2vPatGiMKJhFW4/d7yxgl9e1JNBJym1ri2RYKoN2lVSw70zV8fcjMzj1/nl+2vZU1rDW0v30qeDk/fvHn7Cm4e3lG5h+sbpLLRV4vr2ThTNjNPi5PJul3Ndz+tIt6ef0OsLIcSZzKyp3DS8Ey/P3xW2bwmiqaiq0i87ife8NpTRE0i+6AYGa9WYFn5H0dNP4yvYjWPUKNyjz2XhdiJW3Iq2n9UL3+2QYEo0u6e+3MoL3+9o9L7nyP6+qbNW89uq3vx0eOdWGJ2oLagbPPjOakZ3T+V/hnQ82cNpEySYaoNe+n4nvmDsM4cQCqhe/n4XL9886IQ34q06tIpHlz7Knso9+HU/QYKgA7qPGn8Nr//wOq9veJ0hmUN4eNjDtHfI/ikhhGiKe8bmMveHg+w57CIQqXtvPVQllIb35uKCo88zaQr+oMEFvQdxx5MTybf4qP72W176ZhNYu4FWd99CLP2sVhSUsb/cHZbWI0RTvbNsT1SBVG0ev85f/7OJDKeNi1qgCauo3xOfb8br1/n9ZXkneyhthgRTbUyNN8Ds1YURm49FM3MI4PYHyT3B1ai5O+fyh0V/wBv01nuOTw+V4V20fxHXfnIt0y6aRs+Unid0XSGEOBPFW028c8cwrnl+EYcqvfjq60BZh4FuhNIDa++38gZCf8/94QDfbC7i4r6ZPDFhAksOLMB3qCrsVWLpZ6WpCot2HObaFmp+Kc4sLl+AP32yMWIgFc3k8a8/WMd5vdIxabKHqjV8tHof/1l/gI/vHYVZ/psfJcFUG/Ppuv2oETYNxzJzqCjw/spC7jinac3QFu5b2GggVWdshk6lr5JbP7+V9y5/jw6ODk26rhBCnMnSnTY+e3A0v3xvLd9uKUYhtOp0PFWB0CJUw4UgdCM0uTZ3/UHKanxUuE+8n1UgqFMeYz8rIeozZ81+IrUainby2BfQ+WZzERfK6tSJKd0JW/4LNSVg6BDfDrqdBxnHVp/WFZbzyKcbeeuOoSTHt1wD3FORBFNtzLrCioibg2OZOfT4ddYUljfp+p6Ah5/P+3nUgVRtNf4afvX9r5hxyYwmXVsIIc50TpuZF28aRFGVh7eW7GH6kt2UunwogIKCzaLiC+jowehTAd3+IEt2Hq63Cl8s/awURam3WIYQsTAMgxfm7Qi754ll8rjmx318Ekw1ga7Dts9hwVNwYE0oiAr+OFGiWuCbRyG1O4yaSlGHC7lr+kr+elVfemU2f4+oU50EU21MfR3sY5k5BKhw+Zt0/c8LPseIuD0ZyuaXUfJ5Cb4iH5pNwznQSca1GWjxoTHphs7m0s0UVBTQObFzk64vhBAi1Htq6gU9mHpBDwzDwBvQKavxMebJeWGtMqDxlCi3X693HSuWflahnjEyKy1OXHGVl4OVnrDjsUweA6zZW44voEu59Fj43TDrJti9CPw14Y/rvtAe+YPrMObcT7mRzU8HPS8VFOshv3ltTII1ckOz2jOH0XBYmxYnv/rDq7gCrrDjJXNLOPjeQTInZpL37zy6/qErvsM+Cp4sQK+VhhLUg8zYJCtTQgjRXBRFwWbWmLF0d8SAqHLZbEq/fpnEYRPJvm8GHaa8RsKAS3BvW1rnPJOmYIqwqlS7n5Vr62J0vwcjGMC9YwVl306rc25ANxjbQyq4ihNX5vJH3HcT6+SxWVOp9DRtAvmMFPDBm1dCwfzIgdRxFH8NXQI7uXvbXeCtboUBnnokmGpjuqbFY40wu1J75rAxZk2ha1p8zNfeW7WX/dX7w44H3UGKPioi68YsEvoloJgULGkWcu7JwVfio2JRxdFzA0aAz3Z+FvO1hRBC1C8Q1Jm+eHfYHqojKVEpF0zB3nMEqsWGopmw5w6t09QXwB806q0U6BxyNcnn3k7F4lkU/msShc/fQtWqT4nrfmx1QFXgwvwMEu2RJ/2EaA6xTh6LGP3nF3BgHQTCVwXrY8aPUr4H3rul5cZ1CpM0vzbmqv4d+PuXW8OO1545VFQNW5f+KKoJT8EaPHvW1fnSVBWF6wbnxHztEncJZtUctl/Ktc2F7tdxDqybJ6vZNBL6JVC9oZrkc5KPHq/x16AbOqoisboQQjSHoqrIFf5iTYlSCAVFkbZcNdbPymrSuGN00wobCXG8ZLs54u90LGmnAP6gjtMmAX5Uag7D2nfguPu8zk9V4fLDrgcdxFtCq9evrPIxY52febf8ODkf9IZWs0q2Q2pua4+8TZO73TYm3WljVLfUiKkc0cwcAvTLTqRTu9hXpvzByMvkweogJocJRQsflSnRRKA6UOeYoigE9EDYuUIIIZqmyhPApJ54SpTVrOKMM0esoNaQOLPGVf070C87KbYnClGPtAQrHSL0K4sl7RRgQMdk2S8VrVVvUN+bP2jA00sbqdSpB2Hp8y0wsFObrEy1QfeM68ainSUR+y40NnMYZ9a479zuTbpugiUhYvEJzaERqA5gBI2wgCpQEcDkqPtrpKBg0WSDshBCNBerSUU3Inw+x1CJD8Aw4NWbBzP5jRVUePwEo2gQHGdWGdszjT+P79OksQsRiaIoTBnTjf/9ZENYRT/nkKtR45OpWDyLkk+fRLHEYc3IxTn8ujrnxVs07h7TrTWHfeoyDFjy73rT+345wsITC73cM9hCkq2e2RbdD2veggsfBbOtBQd7apFQvg0a1DmFe8fmEmeObqbxiDizxk3DOzGmR1qTrtslsQu6ER7A2XPtKCaFypWVdY4HPUGq1lURn1d3Fax3Su8mXV8IIURk7RwW/I2kREXDMELZC3OnjmZAxySsJjViUQoAu0XDZgS4MaGSf08agKrA8oJSPlxVyPQlu5mzZh8b9ldEfK4Q0bj8rKx6H3Pkj6P9zU/R8aEPyLlvBukT/hdbdt37C4tJZVwvKYgSFb8L3KX1PjwoS2NsZxNPLmqsNY4CVeH7689ksjLVRt13bi6VHj8vz98V1flxZo0bhnbkNxf3avI1bSYb47uN572t7xEwjqXpaXaN9PHp7J+xH9Wm4shz4C/zs3/6fswpZpJGHEv7iDfFc3vf25s8BiGEEOESbGZGdGvHd1tL6hyPZT+tpihc3DcTk6aS4bTx3t0j2FVSw2sLd/HBykLc/iCaqhDQDTql2JkythsXJQXYeettvNy3F9NWl1D1Y9W0oGGgKQq6AR2SbNw9NpfL+rXHFuMkoDizxVk0Hrkin9/P+SFiNk5DbGaVJyecJX3PouWtAtUcStWrxyPjrIycVsODQxvILlJV8FTW//gZSIKpNqzgsIuL+2Syv9zNloNVBPS6lZg0Bcwmla6pDh44r3uz1P+flDeJD7d/SCBYd89T2iVpaPEaB2cdxFfkQ41TcQ5wknNXDqr52AKnpmqMzRl7wuMQQghR113ndGNFQRk1TUyJspjUsAISXVLjeeTKPjxyZR+8gSAev06C1XS0we+iHSVMHv1z9HkFeJTItwzbi2v445wfeGzuJt6+YxjdMxKa8acWp7trB+VwqMrLv77ZFnVAZTOrPHxZHuf0TGFt8VoqBfbwbQAAIABJREFUvKEV0kRrInkpeZg1KUgRxhzXYCAF0Cdd47IeJh5b4KN3Wj3Ja4YBltj35Z/OJJhqo95etpf95W5m3zMSi0lle1E1by4uYNOBSqo8ARxWEz0yEvjpiE7N2o26k7MT43LGMW/vPDzBunm1KWNSSBmTUu9zbZqNqQOmYlLl10oIIZrb8G7tSLJbcPncYbtbG9tPqyrQsZ2dPh0S6z3HatKwmo6tLM3bUsTdM1biMVRopDprjS+Iyxdk/L8X8sGUEc36vSROf/eOyyXTaeOPH2/AMIywCYMj4i0abn+QiUNTKDHPYcw776Cjo/xYtsvAQEFhYs+JXN/rejLjpcnsUVYnaObQvqcG/GmsjQEvVvPz4dbIJwT94MhogQGeuuSutw3aUVzNk19s4d27hh2tUJOb7uCRK1tn8+9fR/2VWz+/lc2lm8PKpNfHZrIxscdEJvSc0MKjE0KIM5OiKLx6yyCu/veisA37jbGZNUZ2a8fvZq9HNyDNYWFMz3QGdExCiVDda3tRFffMXBVT6pUB1HiD3PDSUr7++RiS46UQkYjeNQOzueys9sxdf5AXvtvB9qLqo019/UGd/Cwnd4/pxndFM5m9dybmMhW/Hrn63IyNM5ixaQY39b6JBwc8GPF3/IyjKHD2DbDyjQYDqtwUlevyzTyzzEff9OMnURTIPQ9sMllSmwRTbYwvoDP1nTU8dEEPctNPTqqEWTMz7aJp/Or7X7Fg3wJ8ui9iYQoAs2pGVVTu6ncXt/eRvVJCCNGSemU6mX77EG6etpwaX4AIBf7qUAwdRVUJ6gZvLCo42l9KAV5ZsIu0BCt3ndOVqwdk19nv9PRX2/D4IwdsNRvnUbn8I/yHC1EtcZjTu5I4YiK27HwAXL4AM5ft5r5xTassK85cVpPG+P4dGN+/A1UeP+UuP6qqkBRnxm7ReGTxI3xz4DNQAjQU5/t+DLJmbp5JkauIR0c9KgEVwNApsHpGo6tTD4+xMn1dhHPMdhjxQAsN7tQlwVQr8gV9fLX7K344/ANlnjLizfHkJORwaddLSY1LBeCfX20lw2ll0tCOJ3WsFs3CP8f9kw2HN/Dmhjf5es/XmFUzuvHjcroSKoE+occEru91Pe0d7U/qeIUQ4kwxsFMKn9w/ij9/upEF20tQAG+g7p2l1aTgCxgogG6EP24ALl+Q3Ydd/PnTTby+qIC37hhGqsNKucvHFxsPEalqeuWy2VQsfZ92F96LrcsAFM2Ee9dK3NuWHg2mPAGd1xYUMGVMrhQHEE2WYDOTUKsZ76vrX+XTXZ+GbUFoiCfg4avdX5GTkMOUs6e0xDBPLam5+DL7oxQuw8yxyZKCqXUn73MSVTy/P271SVHB2R46DmuNkZ5SJJhqBQdrDjJj4wze3/Y+GFATqDn6mFWz8syqZxieNZyhKdfw4aognz0wus3MoOS3y+fxcx6nwlvBuuJ1VPoq0VSNZGsyA9IHyCZPIYQ4CbqkxjPtlsEUVXl4a8kePlt/gMofK+05bWaqPH5Ka3z4go13QHH7g+wsruHKZxfwnwfOYdbyvRH7eureGsoXzKTdJVOx9xxx9Lg9dyj23KF1zvX4g3y3tYhze8neCnHiXH4XL6x9IWIgVTa/jJLPS/AV+dBsGs6BTjKuzUCLD620uoNuXv3hVW7Mu5EEy5ldHGX1njJ+V3QH75m3YgqUodSTdRSRxQGT3qu36e+ZTIKpFrby0Eru/fpefEEf/gjLqkf2JH1f+D3f7VnERYMn0i7+vNYeZqMSrYmMzh59sochhBCilvQEG1Mv6MHUC3ocPfaL99by6dr9+IKNN+Q9IqAbFFV5uWP6CuxmLeJeKe++zRgBH/Yewxt9vRpfkNW7yyWYEs3is52fRZxkLplbQvHcYrInZ9dp21LwZAFdftcF9cd956qiMmf7HG7Mu7G1h95mzFq+hyf+u4XHrhlFfPtv4fVLoLoIgpH3nR2lmsCWCDd/CildGz73DCVNe1vQmqI13P3l3dT4ayIGUrUZGKD6+a7oA/658p+tNEIhhBCnk8PVXj5Zux9PIDwYqtk4jwNvTGXPP66l8NmbOPTuH/EUbjj6uD9osK6wnIMVkdOogu5KVLsTRY2ul1RJTSM3aUJEwTAMpv0wDXfAXed40B2k6KMism7MIqFfAopJwZJmIeeeHHwlPioWHWso7Q64eX3D6xiNbTI8DfkCOr//aD0vfr+TWXcN54K8DEjuBHcvgCF3hlacLI7wJ5rjQ+XU+/8UpiyCjLzWH/wpQlamWki5p5wpX02JKbcXQvm9b29+m75pfbmg0wUtNDohhBCno7eX7yFSEk40e50gdON1uCZyFVctzonuqsTQg1EFVHaLNPAVJ67aX83BmoNhx13bXOh+HefAunt7NJtGQr8EqjdUk3xO8tHjZZ4ySj2ltItr1+JjbiuKqjzcO3MViXFmPrp3JM5ae9CIS4aLHoXzHoaNH8O6d6CmBAwd7CmQfzX0vVZ6SkVBgqkW8sG2D+pdjWosv9cT9PDs6mclmBJCCBGT1xYUhK1KxbLXSTegtJ4VJWuHXigmM66ti4nvNarBcVhNKllJcU38KYQ4ptJXiVkzEwgE6hwPVgcxOUwoWvj0gSnRhHt33ZUsk2qi0ld5xgRTq/eUcc/MVUwclMOD53U/2og7jMkK/SaE/ogmkTS/FqAbOm9ufDNij6aSuSUcfO8gmRMzyft3Hl3/0BXfYR8FTxag1/oC3F+9nw2HN4Q9XwghhIjE4w9S5goPhGLZ6wRgNqnYzOG3B6o1nqRRkyj98gVcWxej+z0YwQDuHSso+3ZanXMN4PJ+UuVVnDiTYoqYnqc5NALVAYwIewMDFQFMjrrrBQYGZvXMKJr17vK9TH5jBX+6Ip+fXdCj/kBKNAsJplrAkv1L8ATC0/tiye/16T5mbpzZmsMWQghxCqv2Bo42Oa0t1r1OZlWtmw5Ui3PI1SSfezsVi2dR+K9JFD5/C1WrPiWu+7FATQFGdmtHutPWpJ9DiNoSrYkE9EDYcXuuHcWkULmyss7xoCdI1boq4vPqpqf5dT9J1qQWHevJ5gvo/OGjH3jhux3Mums4F+ZnnuwhnREkza8F7KzYGfGNH0t+r27obCnb0irjFUIIceqzWzQCEZpDxbrXScfghiE5vPj9LtwRGvc68sfhyB9X7/NtZo27xnSLbfBC1MNmsjE4czCLDyyuc1yza6SPT2f/jP2oNrVONT9zipmkEXUDp7PSzsIRqdDCaaK4yss9M1fitJn56L6R9U6IiOZ3xgRTQT1ITaCGOFNciy/z1le9L9b83hp/Tdh5QgghRCRxZo04s0a1t+5kXix7nQD8AYNbRnZh9d4Kluw8HNbwt7Ex3DC0I8O6nhn7UkTruLXPrawtXosr4KpzPO2SNLR4jYOzDuIr8qHGqTgHOMm5Kwe1Vqqq3WTntj63tfawm8TtCzJvSxHF1V58AR2nzcyATsnkptcfCK7ZW86UGSuZMCiHqQ3tjxIt4rQOpiq8FczZPoc3N75JkasIk2oioAdIsCRwbY9rub7X9WTGN/8SqN1sx6SawgKq2vm9xwdUkfJ77SZ7s49NCCHE6UlRFCYN7ci0hbvw19pHUnuvk6Jq2Lr0R1FNeArW4NmzjuRxt9V6DRjbK40ku4UXbxrI5DdWsKLgMJ5A4yWl48wa4/t34HeX9G6Rn0+cuYa2H4rD4ggLpgBSxqSQMialwefbTDZGZo1sqeE1ix3F1by2cBcfrNyHpoZaFeiGgUlVMQyDHpkJTBnTjfPzMuqk8767Yi+Pzd3M/13dl4skre+kOC2DKb/u5/FljzN7+2w0NNxB99HjEKoMM2PjDGZsnMGwrGH83+j/w2lxNvSSMclJyMGiWsKCqdr5vYlDEo8eP5Lfm3HtseaGCgqdEzs325iEEEKc/m4a3onXFxUQKgFxjHPI1ajxyVQsnkXJp0+iWOKwZuTiHH5dnfOsJpVr+mcTCOrYzBpv3DaEv079B+85ehA0manxhaf9xVs0EmxmfnZBdyYOyonYXFWIE6EqKv8Y+w8mfz455pYzNs3GP8b+Ay3KPYOtzTAMXvhuB09/tY2AboSl6vqDoffcusIKfvHeWjokx/HWHcNIjDPz5083smBbCe/eNYzc9ISTMXzBaRhMeQIe7vryLjYe3oivga7OPj302OL9i5nw8QRmXDKDNHtas4xhZIeRqGr4JuBY8nttJhuTek9qlvEIIYQ4M2Qn2xnerR2LtpfgO67KWWN7nQC8AZ0H31mN2aRy07BOXMV+btj6Fb+afS/f7Cjjlfm72FvmwuvXibNo9MxM4M7RXRnerZ0EUaJFnZV2Fv8Y+w8emvdQ1AGVTbPx+DmPMzBjYAuPrume+HwLry/8/+zdd3hUVfrA8e+9d3p6AgkthISAQOiE3uyiIpYFUbE37GV33f3trlvcblnXsiooIirFXhAFsQDSkd5bQhqkkF6m33t/fwwgYSbJDEwSgufzPDzo1JOQ3DnvOe95X/+WBoHUuVWyj9Zx+Ysr6RJrJT7CJM5HnQXOqWBK0zV+ufyX7CrbFbAseSAezUOJvYQ7vr6DDyZ+gM145ql1RtnIjb1uZM7OOSeCtuOCze+Nt8QzOHHwGY9FEARB+Hl5ceogrnhpJcXVzoAFKRqj6+D0aji9GrNWHuINj5vLrnyM/xiNTOjbkQl9RblzofWM7TKWORPm8Kc1fyKvOg+P5kHV6++WypKMUTbhcsTw1Pi/c2HXEa002qZ9urmAOasP4fAEfy7Rq+kcrXGh6zofTh+JwSAKc7e2cyqYWp6/nI3FGwMGUo01yvXqXgrrCnlr51s8OOjBsIxl6nlTeXf3uxDg96Op/F6LwcL0/tPFKp8gCIIQshibkU8fGMUNb6zjcIUjpAISJ3OrGsgGvi3RmDpzLe/dOxKL8exMlRJ+PjLaZfDxpI/ZV76Pd3a/w/L85ScKdtkMNsZ1GcetGbeyepeF91aVcnkP/aycT+m6ztNL9jUYSNXtXk71j5/hKStANlkxJqYRM+p6LF0yALC7VdYeKmNsj/BkVQmn75wKpmbvnI3D6/C7vXRxKUcXH6XL3V3qpdblPJdD6h9SkQ0ybtXNgr0LmD5gOgb5zL8tibZEnhv/HL9a/quQ8nstioVLul7CNenXnPEYBEEQhJ+nxGgLXzw0htmrD/HW6hxcHjXgeadgOD0aewprmP7uJt66faioFCacFc6LP49/jPlHg/f3GK3x/sZ8vtldfFb2W1qbVUaN07/yM0D1hk+pWv8RCZc+iCV1MJJiwHFoE44D6+sFUzNWZIlg6ixwzuwN5lbnsrd8r9/toTTK9WgeVhSsCNuYxnUZx7/H/huLYkGWmv5WWw1WJqRO4K+j/3pWrqIIgiAIbUeE2cDDF/bgxz9czIs3DGJi/44kRJgafHzd7uUUvv0Yec9PpuB/t1D8wZ9xFuwCfGepfswpZ3VWaUsNXxDOiFGR+ctVGfzty904j/VL03Wdijo3h0rryC2ro8oROJhpCTN/yA64wKG56qhcNY/4S+7Hdt4oZJMFSTFgSx9er/ImwMacCgqr/DcRhJZ1zuxMfZ/3PZruv1UaSqNcu9fOF1lfcFHXi8I2rotSLmJe9Dxe3/46y/KWIUlSvTRERVIwyka6xXTj7n53c2nKpSKQEgRBEMJGkSUu7pNEZrc4hv/zu4CPCXYlfOaKbLESLrQZY3q0I6NjDP/7/iBd4qzMWJHF4UrHidLibq9G747R3De+O5dm1C853tx2FFQFvN11eC+6142t58gmX8NkkNlTWE3HGGu4hyeE4JwJpo46joalUW6pI/yrbj3jevLc+OeodFby4b6P+e/KbxiWbsVmsNI1uiu/6PELesT1CPv7CoIgCMJx7/+YT6AMveMr4QlXPIbtvFEnbrelD8eWPrzeY3/MKedIpYNOsWLyJrQNfTpF8fw3B7AaFRzHdqiOlxsH2HG4it98vI3ffSLx/PUDubhPUkMvFVZ2jzfg7aqjGtkWjRREKXdN11t1d03wOWeCKVULnAseaqPcU6vChFOsJZYLOk5lvrM7cyac32zvIwiCIAinWrS9MOBh91BWwhVZYuWBo0wd2rU5hij8TOSW1fHW6hy+3FFIrdOLjk6k2cDFvZO4a0wqPZLC0zPp6cV7j/Vd40QgFUidy3ffQws28+eJfbhxeEpY3r8xRkXGGeD3UbFGo9mr0TW1yYBKQsIqisK0unMmmEqwJiAjo51SPi+URrkAcea4Zh1nfrmd5PgzL78uCIIgCKFoaAU7lJVwt1ej0i5WwoXTc6C4ht99soMdh6vQNB3PSaX7nR43H27K57Mth+mRFMk/ru1H/y6xjbxa495Zk8OcNTmNBlGncno0nlq0mw4xVi7olXja7x2M9pEmapz+u1Pmzr2QDEbs+9cS0WtMo6+h6zpJ0ZbmGqIQpHMmmBrVaRSzdszyq+YXSqNcq2zmkpRLmnWceeV2usaL9AhBEATh7BDKSjiAONbbdtU4PSzZWURhlRO720uM1URGp2jGpLdr9iqN67PLuHPOj9jdKg11P1M1UDWNHYermTpzHf+7aRAX9Q497a7a6eGfi/cE3PlpquS406PxxEfb2PD7i8P+PdFVFfumTdQsWcKEHWXMTr0Ap1y/4a5sjiB2zDTKv5mBJCtYUgchyQacOVtx5m2vV4Qi0mJgYPLpB5xCeJwzwVTfdn1JsiWRU53jd1+wjXJRXUxY9gLU1ELfX4Ap/DtIeeV2kuPEzpQgCILQsuJsRvLK/W8PZSXcZJCJtTVcEVA4O+0rquGNldl8se0IiizhOBbQKDJYDAo2k4G7xnTjhmFdm+Xfd09hNXccC6SC5fCoPDh/M3PvGk5mt4Z7cwbyyaaCgMW8gim0AuBwq/xw4Cjnn3fmu1O6quLYvJnqxUuo/mYphnbtib7sMu76+zRmz90PAQK+6GHXIUfEUbX2fUoXPYdksmJOSid65NQTj7EaFe4ZmyaKlp0FzplgCuDOvnfyrw3/CthrqqlGuYqkMKnnL7AmDIMf34Rv/ggDboTMO6Hd6RWH0HWdbQVVzFqZzbb8SurcKnUuL51jrXSMtTIhowMm0blaEARBaAFXD+zM/uJav7SnUFbCvZrO+J6iml9b8uaqbJ79eh8eVUfV6u8JqRq+uYlb5YXvDjDjh2zm3T2cjE4xDbxa6HRd5/65m0IKpI5zejSmv7uJDX+4GCXIXSJd15n5QzaOU94vlEIrdccqV55uMKVr2k8B1NKvMSS0I3rCZaS88w7m1NQTj7uqfy2fbzuCO0Bj7ciMC4jMuKDR95mSmXxa4xPC65wKpq5Iu4J3dr9DTlUOXj1wlZSGRJmiuHfAdLAlQs/LoCIHNs2Bty6HxD4w9C447wpQjE29FABLdhbx9JK9FFU5cXlVTr5+ZZfW8buPt/P7T3Zw68gUHru4pwiqBEEQhGY1ObMLTy/x78cIwa2ES8Do7gnijEYb8uryg7z83cGA6W6ncno0nB6NKTPW8tF9o+jTKbrJ5wRjc14FJTWugPc1lXLnG5fKsr0lQVfZK6hwUGF3+90eSqEVgPWHytA0PehUP13TcGzZQvXiJdQsXYoSF+cLoN5+B3NaasDn/HlSBhtzK8gvt+PVGkp+9GcxyvzvpkHEWIObkwrN65wKpsyKmTcve5MbFt1AqaM0YKn0U8nI2Iw2Zl06i0TbSSsQcd3g4r/A+b+DPV/Autdg8W9h8K0w+DaI6dzga77w7X5mrMhq9OJ1vFHb7NWHWJtVxjt3DSPKIn4pBEEQhOYRbTFyZb+OfL7tMGqAj6emVsKtJoV7x3VvxhEK4bRsXwkvfXfwRMPaYNndKjfNWscPv7mA6DDMS17/ITtgEYhgU+7q3CozVmQFHUxV2j0YZRnnKQXJQim0Ar7KlbVub6PfA13TcGzd6gugvv4aJTaWqAmX0XXOW5jT0pp8j0izgfenj+CG19dxuMKBK8AO1aksRpmnr+t/WmfJhOZxTgVTAPGWeD686kMe+f4RdpftxqN5Gix3bjPYiLfEM/OSmXSNbqDMq8EM/Sb7/hTvgo2z4bVR0G2Mb7cq9XyQf9pVeuOHbGauyA5qFQh8K0G7jlRx2+wNvD99ZIs2jBMEQRB+Xn57eS+W7ztKhd3dYBGAQCxGmfE92zMiLbSzK0LreXbJvgYDqaZ2hFxejU82FXD76MA7KsHyqhrf7SlBP+WHLZSUO4BtBZWUV9YR5apDraxArahAraw88be3ogK1ohK1spLDdhm146W++dtJQi200hBfALWN6iWLqfl6KUp0tC+Aems25u6hLzYkRln44qExPLd0H+//mI/ETwvuxxk0FcWo0LdLHL+/ohdDUsTv4dnknAumAGLMMbx9+dvsLd/LO7veYWnuUgyyAQnfVq1bdTO0w1Du6HsHwzoMC/7wXlIGXPkf347V9g9g6R/B4/Cdqxp4Ewdrjfxn6T6cQawsnMyt6uwurGbG8iwevkg07xUEQRCaR1K0hQX3juD6mWuocXoJJrPIalQYmBzLizcMEofd24i9RdVkl9YGvC+YHSGHW2XmD9ncNqrbGf2bVzk8KLLkl8IWasqdweVk02UTSTF5McTGocTGosT99LexUycsffpgiIsj3RiF+4tCTl0tCKXQCoCq6USafNNkXdNwbNtGzZIlVH+9FDkygugJl9P1zVmY09OD+2Y0IsJs4M9XZfDbCb34cnsh89bncrTGhUfViTQrdKgsokw3sPuIwuQZazHIEnE2E1OHJnPLiBQSReptq5L0U5cLTpKZmalv3LixBYfTPOo8dRypPUKdpw6rwUqiLZE4Sxj6Sek65K/3FazY/zX/Z/4/PjyaghrgWxpMXnCczcjGJy8J+pClIJwOSZI26bqe2drjOFPnyvVJEFpDfrmdK19a6StTrYMaYC5gMcpoOlw/pAt/mZSBoQUyJ86F69PZcG367cfb+WhTgV/BCc1VR8Ert5FwxWNNBhQRJoU3bx/KiLSEgPfrXi9qTQ1qZSVaVRXq8T+VVajV1ahVVRRV2ZkqDcUl1V+7r921jIplb5L80Nygvp5Is8IH00fSJ4jCGFUONwOe+ibgfUe/eBb73lUgycjmCIyJaVi7DUCzV9UrtAIwqnsCbw61UL3ka6q//hrZZiN6wgSiJ1yGuUfLLHy//2MeTy/Zh8vloc7r/ztqPnbefnR6O56Z3J92kWa/xwjh0di16ZzcmTpVhDGCHnHN8IMvSdB1BHQdQV1FMZ89uyFgIBVsXrBb1UI6ZCkIgiAIpyO/wk6UxcicO4bx7rpcvtxRiATIkoRX04i2Grl7TCpTh3YlPkKUQm9rtuVX+gVSENqOkMfjZeP8haR4clGrq3wBU+VPQZNmt6NERSHHxqDExKLExKBER/v+jo3B1KUzSb1i8G40+O0ShZpy57E7qXvydxSldsKc3h1zejqm9HQMcf4L459sOowhwG5Y9YZPceZsJWrQlTjzd+ItP4y7cB9qxWESJv663mONaFz15UyOfFpC9IQJJM+cgblHjxbbmdV1nX98tYd563JxNHJs5PgZqx/2H+WKF1fy4X0jSUmIaJExCj/5WQRTLeHbXC+KwQhnUorTpfLuulwRTAmCIAjNxqtqPLVwN09e2ZvBKXEMTonj2cn9qXJ4cHhUoq1GoswGkdLXhtW6Alc0DqUIg0eHOoMJc1pPX4B0PGCK9QVNclQUktz0bmWXrGXkltvr3RZqyp0t0krPSTfhzTqIc/duqhZ+gevgQSSjEXN6Oub07pjS0zGldWfmimq/QOrkuVhQKX5ITPrXb7H1Oq9Vfg9eXZ7FvHV5jQZSJ/NqOqW1LqbMWMviR8eSIHaoWpQIpsKkuNoZsApLqHnBRyr9e2QJgiAIQkg0DVxV4KwGUyRYY+HYBHre+jwSIk1M6NvhxMMNiiwmYGeZ3LI63tuQT9bRWuxulWirgUHJcUzJ7NJkY12rMXCwFMqOkFGRaTdmFPFjzqwIxfTxafz9yz31+kyF0tvMYpC5a2waMWN7wNifAiFd1/GWHMV18ADurCxce/eR99W3lHW6GuT609tQ52Kg4+7YhYhWCKRyy+p46bsDAeeUjR0Z0XQor3Pz9y/38N+pA1t83D9nIpgKE5dHC7ilHmopzmDKYgqCIAhCQNVHfOd4f3zDVyBJNoCmgiTDgBuoHHAPL313hPn3jBA7T2epZftK+N/3B9l5uApV0+vtsny/t4Tnlu7jkj5JPHRhOr06BO4F1TXexoES/wIUoewImQwynWOtZ/bFANcM6szfFu3xuz2Y3mYAGnDDMP+Ky5IkYUxKxJiUCKNHA+AsrcP00ko8p2QJhToXM+kau+9/hGH//SfGDh2afkIYvbU6By3AfDKYIyNeTeerHYX8ZVKG6EHVgkQwFSbRViMmg+wXDIWaFxxlEf8kgiAIQog8DvjsAdj7pe//1WNNUtWTmpdueZeIzfP4xNablMiPgagWH6bQME3znZOZvz4vYF8m4ETbla92FPLdnmKenTKAif07+T3u1pEprD1Qgv2UlwllRwgkLujV/oy/LpvJwEMXpvO/7w/6fV1N9jYzKtw4LDnowgoGWQpY8j/UuZhkMhE3fgw5U66n8/P/wTZ0aFDvf6acHpUPNubjaSRNsakjI7Ik8fGmAu48wx1FIXiiqVGY9O8Sgxxgle/kVaCmGBWJ4amid4AgCIIQAlcNzLoE9n3lC6KOB1Kn0rwYdTddHbthxmiozGvZcQqN+tuXu5m/PrfBQOpkmg4Oj8avP9zGkp1FJ27XNY2ab7+ly58exuqsC/jc6GHXEXfhXVStfZ+Cl6dR8Nrt1GxehLXHTylwRkXipmFdMRtOvx/TyR44vzsT+3dsMP0wEKtRYXR6Ak9e2Sfo58RHmHAHyPAJZS4G4FF1Uu+8jY7/+hcFjz1O+dx5NFb9OlxWHywNOJcMJU3R4VF570fxu92dazZ5AAAgAElEQVSSxDZImAxMjiUp2kxOWf1DlqGsAsmSxO2jxEqCIAiCECTVC/OmQOn+hoOoU0iaB+pK4a0r4L5VvvNUQqtasrOQ9zbkB11w4DinR+Px97fSJ3Ek0auXUTZrFrLFQvt77uHBiHSeXbo/4Gs2tSOkyBK3jkoJ+etoiOR18cx4C+10M7O32dGQ8AQqf4xvd8kgS1w7qDN/u6YvcpDtYnRNgy0b6adVsoX6JdRD25GDEWkJWE0KjBlNt/cWUPDQwzh37qTDX/6MbGm+nk5ltW60AEFbqGmK5XXuph8khI0IpsJEkiTuG9+dvy7aXe+QJQSfFzwgOZauCbaWHLYgCILQlu3+DAq31wukur1Qg90Dhx6NJMLkm4jO2uxm7nYPy28/VjZZV6G2GFa/BBf/qTVGLpzkhW8PNLgj1VSfSo/HywtPvMCjUg4d/vB7bCNHIkkSt2s6q7LKWXOwFGcI57EtRoVnftGPLnFhmI+UZcG612DrPCRJ5reSzE22eOY4xvGedCGSwXSiMApIqJrOdYM7c8foVNITI4N6C9eBA1QtXEjVF4tQ4uK448LJ7C+RqTsliAx2LhZhUpg+Pu3E/5uSk+m2YD6FT/6R3Gk30+XllzB28k+tDAePpgXcAQs1TTHQGX6h+YhgKoyuHtiZF747gNOj+nWVb2oVSJbgwQvOvIu2IAiC8DOy+gXw+KdzqTq8uN7N78c2ctZEdcOPs+CC34EiDqu3lt1HqsktC5ySF1TRASSWJA/ln39+EstJaXSyLPHazYN5YN5m1hwsCyp90GKU+eukDCYN7HxmX5THAR/fDQe/9RVA0Twn7kqmmj8acnhCWcBmpT+VhgS00b8irnM6g7rGYjM1PTX1Hj1K1ZdfUrVwIWppGTGTriL59ZlYevakm6bzt399R53Hf6e2qbkYQJTFyOju7erdJttsdPrPc5S/NYdDU6fS+bn/EDF8WJDfjODFWI0oAXbiQi0lH2kW0/uWJL7bYWQ1KXxw70iu+t8qapwev4CqIRajTGZKHE99sYs3bxtKajvRcE0QBEFoQvEuKD0Y8K4nRpl4ZrWLB4aaiLU0kialq7B3EWRc20yDFJryztoc3AFS3kIpOiDJMkt3FzNpQP0dE7NB4Y1bMpm/Po/XVmRRYXfjcKv1ijSYFAlJkhiSEsevLu3JkJQzPLvtroPZl0HpAfA6G3yYRfIwStsEbmDVarj5YzA1fCZIczio+fY7qhYuxLFtG1EXXkjSr3+NbfhwJOWUIHLaYG5+c/2Jgh3BshoVXr15cMDUQkmSSLjzDiy9zuPwr35Fu3vvIe6WW/yqYmqazooDR5m1Mpv9xbU43CoWo0yXOBt3jkllQkYHTIbAJQsyU+IDpj+GkqZokCXG9jjzwiFC8EQwFWZdE2wsengMU2eupcrhoc7d8EqQ2SAjSfD89QO4ol8nFmzIY8qMNfx36kDxiyAIgiA0bv+Seiv+J8vspHB+NwPPrXHx9wsbOePhroWdH4tgqhXtL64JmJYVatGBvAZ2t2RZ4uaRKUwb0ZUNh8qZuy6XvHK7r0GzxciQlDhuGZkSnrQ+TYP5N/jO8HmDO8MH+HZX502Ge1dAu5+ydHRVxb5hA1WfL6Tm+++xDhhAzKRJdHnxBWRbw+PN7BbPKzcN5sH5W3AGsSMHPwVSg7vGNfq4iFGjfOeoHn4Ex86ddHzqKWSrFV3Xmbsulxe/O4DDrdab/9W6oLTWze8+3s7vP9nBHaO78ehFPTAo9YOqDjEWhqXEsTKrFKgfpAWbpqjIEneNFefvW5IIpppBcryNZU+cz5KdRby2PIucsjoUScKr6SiyhCSBUZa5fVQ3bhrRlcQo3wfdjcO6ktYugocWbOGB87tz+6huog+IIAiCEFhNMWjeBu/+6wVmRs+u49HhjTd4pbYkzAMTQlHnCjzZD6XogKZDpSNwYA2wvaCS9zbkk1dux+lRSYwyMzwtgSlDkomxhTHFM/t7OLypXiAV1Bk+8O1offNHuHEBzv37qT5+DiohnphJk0j81S8xtA9+ofmi3km8d+8I/u/j7eSU1eFRdb+gVZEljIpEartInv5FP/p3Ca4Yi6lLF7rNn0fhH/9EzrRpdHrxJZ5cV8aX2wsbTac8HmC9sTKbH3PKeev2Yb5CF/iaENcsXsxViz9mY49JOCT/KXowaYq9OkTRvX1w582E8BDBVDMxGxSuHtiZqwd2Zk9hNXsKq6lxerEaFTrEWBjVPcFvRQJgeFoCn9w/irvf3sj+4hqemtS3we1gQRAE4ees8VzyvokKE3sa+PcqN73bN/Y5Ig6rt6aGzreEUnRAliDOVj9o1nWdTzYf5tXlBzlS6cTlrX+ee9XBUp79eh+XZXTgkYvSSU8MQ9+x1S+e/hk+dPR9S8mbciXuo3ZirppI11lvYO7R47SHMzA5liWPjWP3kWreXJXNt3tKsLt9CxA2k4FL+yRx19jUBpsfN0a2Wun07DOUv/02v/nDbJYmD8EZ3CYYTo/GlrxK7nlnI2/fOQzX5k0UP/MseL1c8cSv+XCbzraCqoBl3htjMcr86argS8kL4SGCqRbQu2M0vTsG/4uaHG/j4wdG8dh7W7l51npeu3kwCUE2rBMEQRB+JiITQVJ8554a8NT5FgbPrOVXIxv5DLG1a/g+odn16RTN1oIK1FPmzaEUHbCalHrnrV1elUcXbGXFgaM4GjhucLxk+pfbC/lmdzGv3jyYC85LPP0vpDIP8tcHvCvYM3y6rtPx6jSMN75Q7xzUmerTKZr/XD8wbK93nCRJ7Bs7kW8OrA86kDrO5dXYlFPGq798hsu3L6X9448RfeWVSLLM7EEern11DfnldlxBBlQWo8yzk/uf+Zk3IWRiy+MsFWk28PotQxiaGsfVr6xmT2F1aw9JEARBOJukXwKGxlP40uNlpmYYeWlDA31nTJGQcV0zDE4I1m2jUjDK/tOxk4sO2PevRfM40VUvjqyNVCybXe+xkqZxce8kwFcA4YF5m1m+v6TBQOpkqq7j8KjcP3cTaw6Wnv4XcugHX3AfwMln+BojSyqmms1hDaSa28wVWTga+DbX7V5O4duPkff8ZAr+dwvFH/wZZ8GuE/c7vDoLInuR+tWXxFx1FdKxn4Moi5HPHxzNkJQ4rEaFxlpt2UwKNpPCjJuHcNWAM6zCKJwWsTN1FpNliScu60XPpCimzVrPv6/rx6UZHVp7WIIgCMLZoNNAiE2Bo3sbfdifxpt5d3vD52noMynMAxNCkZ4YRc+kKLYfrvK7L5iiA0ZJZ2LeBo7cPp9206czX+/EmoNlIVeyc3o07nlnI6v/70JibU2cswvEURGeM3xO/+/D2aqoysm67PKA9wVT1h6gUjazpchOZrf6hWIizAbm3zOC7QWVvLEym6W7in3HPnRAAq+q0z7KzH3j07hmUOegSsoLzUN859uAqwd2pltCBNPf3cT+4hoevCBdFKYQBEEQYMzjsOhx8NhP3JTzWP2zL8kxMs4nA6SaK0bIvAMMIo28tT1+SU8emLc5YPGCpooOGI0GHnn+CSLWLKfwmed49bxbcBisAR/bVANgVdd5/8d8po/vHvTYdU3DW1SElpePSdVoaHYS9Bk+qe0kTS3afiTg7aGUtXd4VN7/MZ/MboHT8/p3ieXlGwdTZfewr7iGaocHs1EmKdpCj8RIMR88C4hgqo0YkBzL5w+N5t53NrKvuJZnJ/ev15xPEARB+BnKuA7Wz0Av3oWkNpDKF5AEtgQY/XizDU0I3gW9ErlnbCpvrDwUVHPd4yxGmdduHkKnhEi4aiIbu2finrcJArxEMDslTo/GrFWHuGdsWr1eS7qm4S0pwZ2TizsvF3eu748nNxd3fgFKdDRxfSTik3wVixsS3Bm+tnPm50ilI+CZplDK2us6FFQ6mnxcjM3IsNS28735ORHBVBuSFG3h/ekj+e3H27l+5lpevyWTDjH+/UM0TWdNVhlLdhZSUuNC13XaR1uYkNGBMentAjajEwRBENoggwlt2qeUvDCWBK0Io95IOt9xsgHM0XD7VxCR0PxjFILy+CU9MSgyry3PwulRG62xaJAlTAaZV6YNZnzPn8qFv7UmF3uAQCqUnRK7w8V3cz5hQPkh3Ll5vsApPx85MgJTSgqmrimYUlKImXgVpm4pmLp29fV8qiuD53sHDOSOO/kMX7/EADtQRhsMuLGRr/zs4mgglTKUsvYArhACaOHsI4KpNsZiVHhh6kBeW5HFNa+sZsYtQxiY7OuL4HCrzFufyxsrs6lxerGfcvD08y2HsZkN3DM2lWnDU4hooByrIAiC0Hb8a3khe+NfYk7Ua5DzA2hq4Ga+kgwGCySkw03vQ3Snlh+s0CBJknjkoh6M6p7Aq8uyWJVVigT1dj5sJgVdh+sGd+aesWl0O6mCH0BOA417Q9kpUd0esvfnM6hbFNFXXI4pJQVj1xSUyIjGnxiRAD0vg72LQG/4vFajZ/h0rU0FU+0iAp//CqWsPUCMNYy9voQWJ2bTbZAkSTxwfjo9EqO4c86P/GliH0ant+PGN9ZRUG7H2UAZzbpjHbmfX7qfBRvyee/eESRF++9sCYIgCG3Du2tz+G5vCZ/cPwbFdiGUZ8P6GbBlri9/6HjpdE31FZoY+SB0GtTawxYakdktntl3xFNS7eTTLYc5VFpHjdNLXISR/p1jmTigY4PFBpwN7HCE1ADYZMY44WrajUkNffCjHoGD357eGT7Z4EtbtYTe86m1DE6JI8Ks+DVeDqmsvVFhdLpoT9CWiWCqDbukTxLJ8cO5a86P/PEzLw6Pildruvmi06uRX27nmldWs/jRsadXtUcQBEFoFna3l0XbC9lXVEOF3U20xUj3xEgm9e9EjO2nFezv9xbz0vcH+ei+kT9dx+PT4PJn4JK/Q10JuGp8qVMR7cFka6WvSDgdidGWkApBgG9iHkgoOyUGWSLScprTw+ShvoBo1yf1AqqmSWCNh0ueOr33bSXjerbHYvAPpk4uay/JCpbUQUiyAWfOVpx524m74M4Tj9V0nSmZyS09dCGMRDDVxvXqEE2XOBsbKstD6mHv1XRKa11Mn7uJ9+9tettfEARBaF45pXW8sTKbTzYfRpKol6ptNcr8fdFuJmR0YPr47mi6zq8/3M4bt2aSkhAg/cpggpguLTh64WzQq2M0BRUOv/lAKDslAD0SI09/EFe9CI5yyF4eXEB1/AzfHV/5GlG3IYoscdeYVF767oBfVlAwZe0VSeLyvh1Eml8bJ4KpNi6ntI6t+ZUBA6mmSqB6VJ1teZUcLKkhPTEqwCsIgiAILWHpriIefW8rHlULmGFw/KD7F9uPsGRXESZF5unJ/RmSEtfSQxXOYneNSWX1wVK/M9Oh7JS0izSfOIt9WhQDTJ0H3/8d1r0KkhQ4qJKNICvQcQBMmdNmz/DdNLwrs1YdwuV1+83FmiprbzbKPHxRj+YdoNDsRDDVxr21+hBagA/eYJvFeTSN2aty+Od1/Vpy2IIgCMIx3+wu5pH3tgTVZFXTfQUJVE3naI2zBUYntCXDU+OJtRn9gikIbqfEZlKYPj7tzHsXyTJc/CcY+zhs/xBWvwCVeb7eZprq62024EYYcT+0a9vBRKzNxPv3juDaV9dQ5/aiB5kmZDXKvH5LJt3bn8EuoHBWEMFUG+ZVNT7cVIDnlGAqlBKoqgafbCngL5MyfJ21BUEQhBaTU1rHIwuCC6RO5tV0/rV4L307xzAkRfSeEXwkSeLxi3vyp893nVYDYJNB5pqBncM3IHMUDL3T90f1gqsajFbfn3NIj6QoPn9oNDe+vo46l5e6AMHscWaDjFGReeuOoQxtoFGv0LaIYKoNq7B7UAPsSoVSAhVAAsrr3AF7VgmCIAjNZ9aqbDxqAxVYm0jVdno0XvruAG/fOTzg84WfpymZyWzNr+STzYdDagBsMynMu3t487VNUQxtqiFvqLq3j2Tlby9gyc4iZizPIqfMjiL7Fq0VWULTdYyKTNd4KwvuHUmkaE9zzhD/km2Y3e1FCdCAN9RmcbIsUevyhnt4giAIQiPsbi8fbzoc8IxUsKna67LLKapyisUwoZ6/Xd0Xq0lh3rq8JgMqs0HGYpSZd/cIMjrFtNAIz01mg8LVAztz9cDO7Cuq4WBJLbUuD1aTgeQ4K6kJEYx7dhl2t1cEU+cQ8S/ZhkWYDQF3pkJtFqeq+lnxS737SDUbDpVR5fBiUCTaRZq4uHcSCZHm1h6aIAhC2H25vZBAR1NCSdXWgfkbcvnlJec182iFtkSWJZ68sg8X9Upi5oos1maXoes6bvWnOUOEWcEgy9w6MoVbR3ajfZT4rA2n8zpEcV4H/+JeEwd0Yv76PB67uGcrjEpoDq0/gxZOW6zVGHBnKtQSqLrLScUvJqEO6I910CCsgwZiOe88JEPz/3i4vRqLdxby2vIscsrq0HXfbZIEFqPCHz/fxYXnJXLPuDRRtUoQhHPK3qKagIUCQknVdns1dh6ubo7hCeeAkd0TGNk9gaIqJ19sO0JBhR27WyUh0sTA5Dgu7p2IQRHnpVvSbSO7ccub63ng/HRxVv0cIYKpNsygyEwdmszcdbl4TlptCqUEqkGWmDI8nbT7Z+LYsgXH1i1UvLcA75FCLP36YR00ENugQVgHDECJPYNSqQGU1rq46Y11FFQ4/CYUuv5Tj5Wlu4tYsf8oUzK78JerMpADBJCCIAhtTaXdHfD2UFO1qx2ecA5LOAd1iLFwz7i01h6GgG/Hqnv7SBbvLOTqcBb7EFqNCKbauNtHdWP++jw4pbtBMCVQwXco8o4xqZjbR2JOSyX2F9cBoFZV4di2DfuWLZS9NQfn9u0YOnb8KbgaNAhTauppl08tr3Mz8eVVlNa4Ap4XOJmmg8Oj8uHGAqocHl6YOvDMy7YKgiC0smhL4EadoaZqR1rER7kgtCW3jerGzB+yuHpgZ+weO+XOctyqmyhTFPGWeJQgF1KEs4O4ArdxKQkRDE9NYF1WKe5TgpKmSqAaFYkhKXEBexwoMTFEjhtH5LhxAOheL679+7Fv2ULd2nWUvvoaWm0t1oEDj6UGDsLary+yzdbkmHVd57bZGyirbTqQOpnDo7J0VzGv/5DN9PHdg37e6ahxevhk82FWHjhKhd2DSZHpFGthSmYyw1PjRTAnCD9DOwqqeHNVNjuPVFPn8mI1KqQkRHD76G6MTW8X8q55WmIkVqPiVyAglFRtoyzRI1H0qRGEtuSiXu3589KF3LxoDrsqNmKUjcjIeHUvFoOFab2mMeW8KbSztmvtoQpBEMHUOeD5EdFM3HGAUksMXoL7MDfIEolRZl6bNiSox0sGA5Y+fbD06QPTpgHgKS7BsXUrji1bOPr88zj378eclnbi3JVt0CAMHTv6BR4bcyvIOlpbLzXxuKZKATs8Kv9bdpA7Rqc2S65xXpmdl78/wBfbjyAh1ZvkSMDinUXE2UzcOy6NacO7ilxzQfgZWLKzkOe+3sfhSicur8rJa0DZpXVsOFSGzWTg/vO7c/uobkEFVbqukxxnxeX1PzMVSqq2LEPf7ke5/9v7OVBxAIfXgUkx0TmyMzf3uZmLki/CqATeARMEoeUdqDjAw98/jKddGdvKfI23vdpPFZVdqos3d77JrB2zuLr71fx+xO8xyGK6fjaT9EZaNWdmZuobN2484zdxelS+2lHI3qIaquweoiwGUttHMLF/J2Ks4iJ/Jhw7d5F/330YHnuC+/NjyS2va7L5o8Uo0yXOxoJ7RoS1eo/mcuHctevY2aut2DdvQTIYsA4ahG2QbwfL0qsX9763nW/3FPt1CW+oFLArf1e9yUOEWeHpX/RnYv9OYRs7wIZD5dwxZwNOt0qAOK8eq1FhQHIMb942tPl6cjQTSZI26bqe2drjOFPhuj4JQkN0Xee5r/cxe3VOUP16rEaF0ekJvDJtMGZD4DSdWpeXz7YcZu66XNxeDZtZYfeRagJt0tfuWkbNxs/xlOXXS9W2dOkNgCFmI1EdvsVkcmP32v2eH2GIQJIkpvWexn0D7msTE7Jz4fokrk1CQ7aWbGX6N9NxeB3oNJ2ZY1Es9G/fnxkXzxCLIq2ssWtTswZTeWV2Zq3K5qNNBQD1igxYjQqarnNFv47cOy6N3h2jT/t9fq7smzZR8PAjdPzrU0RdfDFOj8p7G/J4fWU2lXYPDrda71c1wqQQZTUyfWwaNwzritXUvDm5uq7jKSjAsWUL9i1bcGzZSklhKbee/wRuqf57a646Cl65jYQrHguqAmHfztEsenhs2Ma6Lb+SG15fF1KDQ7NBJqNTNO/dO7JNVeQ5FyYrICYsQvN76bsDvLY8K6TrgsUgM/689sy4eUi9Xfn9xTXMXZfL51uPMDItgVtGpjCqewJ7i2q49tXVTS6C1adjTvocY+wmJLnp4hMWxULfdn155aJXsBmbTsVuTefC9Ulcm4RAcqtzmbpoKnWeupCeZ1EsjE8ez7PjnhVHDFpRY9emZlum+m5PMQ/N34JX1fAEWHI7/uG0cNsRFu8s5M9XZXDjsK7NNZxzTu3q1Rx54jd0evYZIkePBnylxG8fncpto7qx/lA5S3YWUVLj20JuH2nmsr4dGJmW0GK/jJIkYUpOxpScTMykSQAU7yrAtGAb7lN6BIdSChjgYElt2Mbp9KjcNntDSBMmAJdXY3dhNc98vZcnr+wTtvEIgtD6thdU8trygzhCCnLA6dX4YX8pH24q4NpBnVm6q5h31+WQdbSOG4cms+SxsXSMsZ54fO+O0fxpYgZ/W7Q76GuQqf3ioAMpAKfqZPvR7Tzy/SPMvGSmONwuCK3gmR+fwe7x30GuWFlB6deluEvcKBaF6CHRJE1OQonw/Z46VSc/FPzA9tLtDGg/oKWHLQShWYKp7/cW8+D8zUGttKmajqrp/PWL3XhVjVtGdmuOIaFpOpUODzVOD1ajQqzNdFbtJjg9KqW1LhxulQizgXaR5gbHV/3NNxT9+S90efklbEP8zzxJksSItARGpCU097BDVicZkRQDeOtHU6GWAnZ5NHRdD0tguGh7IR418M9qU2e4nB6N+evz+PWl52ExigmKIJwrZq7IxuU9veuCw6Pyjy/38OySvaS1j+SWkSlc2qdDg9f0m4Z3RdU0/vHVniY/N2VrLqb4tUEHUse5NTfbjm7j/X3vc1Pvm0J6riAIZ+ao/SjrjqzzS+0rXVzK0cVH6XJ3FyL7ROKp8HDk3SPkPJdD6h9SkY9dM1xeF3N2zuG/F/y3NYYvNCHswVRemZ2H5m8JMWXh2IfPV3vo0ykmrM1ZS6qdzF2fyztrcnF4VBRZQtN8P85X9OvI3WNTyegUE7b3C9WuI1XMWnmIr3YUIksSsgyaBpIEUzKTuWNUN7q1izjx+KqFCyl+5lmSX38da9+MVhv36bKaFALVyAi1FLAOXPT8CjpEW+gQbSEp5tjf0RY6HPvvdpGmoApEvLb8IHUBGmc2dIbLcWD9iUnTcV9sO8KUzOQm30sQhLNfRZ2bb/cUBzzHFOx1oc7l5a/XDwi6j8wtI7vRp1MML393gLXZZQD1gjmjLCHLEvFd1lAnewOetghmhfutnW9xY68bRbqQILSgD/Z94Heb6lAp+ayEznd1Jqp/FACm9iaSH0hm/xP7qVpTRdw433xYQ+OHgh+ocFYQZwnfHFkIj7AHU7NWZeM+zdU8p0fjpe/28/adw894HG6vxu8+2c4X2wuRIOAK48KtR1iys4i09hG8cWsmnWKt/i/UTEqqndz9zkYOFNfg9uqoAc6uzV+Xy3sb8hiRlsD/bhqE97OPKX1tBilz3sKcnt5iYw2njjGWgLtAoZQCBkiKNjPj5iEUVTkpqnZSXOVkf3ENKw+UUlztu63S7iY+wlQvyEo6Fnwd/+9qp4cjlU6/19dcdVSumkfCFY9hO2/Uidtt6cOxpdf/+bS7VWatOiSCKUE4RyzaUUiggnyhXBc0XWflgdKQmnIOSYljzp3DKKpysmBDHjsPV1Ht9BBpNtAzKYrLB0Ry1/d70ANEecGucFe5q9hYvJGhHYYGPS5BEM7M0tyluLX6TbrtB+xoHo3oIfVrBigWhaj+UdTuqj0RTAEYZSPrC9czIXVCi4xZCF5YgymH29dYNVDvoGBX89Zll1NU5aRDjOWMxnHTG+vYU1TdYGAHoOo6Do/K3qIarnhpJR/dN4r0FujXkVdm59pXV1Pl8DTaZ8mj6aDprM0u44p/fMWLG+bSd+67mJLb7qQ9rX0kXeNt7C+uf+YplFLAFoPMbSO70TMpip5JUQ2+l0fVOFrjOhFsFR0Lsg4U1/huq3ZxuMKBO0BwF+oZrsJKR5DfAUEQmpvLq1Lr9GI1KViNSsi7MEcqHAHPSoVyXdB0yC0L7aD5cR1iLDx+SU+/2xfsXYAs+e+2h7LC7fQ6+Wj/RyKYEoQWVO2u9rtNrVUxRBqQFP/rkyHGgCO3/rzCq3upclU12xiF0xfWYOqrHYUE+swKZTVPB+atz+VXl553WmPQNJ3pczexu7C6wXz3U6maTpXdww2vr2Xxo+PCWi78VFV2D1NfX0uF3R0whSQQt1ejUIMnL/0ln3YMbznw1vDA+en84dMdfql10cOuQ46Io2rt+5Queq5eKeCT6cANQRQrMSoynWKtje44fralgN9/urNepUk4jTNcQf6sCYLQPCrtbj7YmM+slYcorXVhUGRUTcdskJmamczto7uRkhDR9AsBdadWyDkm1OtCqOnuTSmuK8ap+u+kh7LCraNTWFsY1nEJghA6JVLBW+tFV3W/gMpb5cUQWX+KLgXZR1RoeWENpvYWVftNSiG01Ty3V2Pn4dOPvJfvL2FjTnnAyW1jaYY6UGn38MK3+/nHtf1O+/2bMvOHLMrqAgdSjY3PKxs4VOVm4da2fzbn8n4d+PPCXYD/z0pkxgVEZlzQ4HNNBplL+yQRH2EKy1iiLEaUAPk8oYePjo4AACAASURBVJ7hEsUnBKF1eFSNpxbu5sNN+cgSJ3aUjmcl2N0qc9fnMn9DHkNS4nj5xkEkRDa+YBZvMyGB37mkUK8L0WHuo+jwBt4BD3WF26W6wjouQRAaF22KptRRWu82W7oNySBRvamamGE/nd1XnSo122tImpxU7/GKrBBjab0z/kLDwlrOrtIeuLpQqKt51Y7QqhSdbMaK7IABXfWGTyn/7g1iRlxPl4fm0vn+t4gafAWOA+tPPMar6Xyy+TD2BlYlz5RH1U40ajyd8TncKjNWZDXL2FqS2aAw+/ahIQcgiuw7c/XP68IX7KYnRgb89zj5DFcw0toHt+ItCEL4OD0q095Yz8eb83F5tQbLmHtUHZdX48dD5Vz+4koKKvzLE58so3M0ZqP/x2Mo1wWzQWZYt/AeFI+3xAdcnT55hftUgVa4o82ir6MgtKQrUq/ArNRfxFFsConXJHJk7hFqttege3XcR93kv5qPMd5I7KjYeo/3al5Gdgzu6IHQssK6MxVlCfxyoa7mbc6rZMzT39PxWJGAn/620iHGTIcYK4lRZoynVGrLK7OzLb/S7/VCSTOUJPhsy2FuGp4SzJcckqW7ilEDbEmFMr4jlU625VcyIDn21JdpU4akxPHmbZnc885GHB6VRnpHA76JSec4K+/dO4IoS/hWe1MSIujdMZqtp/zchHKGK8KkMH1cWtjGJAhC0zRN58H5m9leUIkzyDRbj6ZTVutm6sy1fPXoOGJO2TnaV1TDJ5sL+GzLYTwBApNQrgsANw4Pb+/Efu37YTVYsXvrB4OhrHCbFTPDOg4L67gEQWjc5J6TeX376363t7+iPUqEQtH7RbhL3MhWmejB0SRPT0Y+aUFHlmQuSL6AGLPYmTobhTWYSmsfidWo+DUeDKVSm0GWuHlEV+4YnUphlZPiaieFVU4KKhxszKk4VjjASWmti1ib6URltg7RFo5UOdACzMpDSTO0u1U+33qkWYKpr3YcCViCO5Txubwqy/aWtPlgCmB0ejsWPjSaZ5bsY/n+o0heL65TNksjTL7D4zeP6MrDF/Ygwhz+1mj3je/Orz7cSp3r9M5wGRSZi3vX344XBKF5rdh/lLVZZQEDqcZSplVd52iNi1e+P8jvr+xNSY2ThVuP8Mnmw5TXublmUGfevXs4S3YW8cqyg34p48FcFyRgTI92JEadfiGlQEZ0HEGEMcIvmDp5hVu2yPWq+Z26wq2jM7nH5LCOSxCExiVYExjdeTQrClag6fWvKfHj44kfH9/o802yidszbm/GEQpnIqwz06sGdOJvi3b73R7Kap4iS9w60ndQuLHDwqqmU1rrorDK6SuPXeVgT2F1wNXEUNMMD5bU8t9v9vvdHnDzJEDwFuhxug5bAuyahTo+TYeSmnMn3z09MYrXb82ktNbFq0/8h40pg6k1WDAqMu2jzFw/NJkJGQ03uwyHi3snEms14nCrfmfZmjrDZTUq3H9+WlD9rARBCJ8ZK7IaTOluqnKsW9V5Z20Oe4uq2ZJfyaV9OvCHK3szIi3hxBnKhAgTb646FPD8bVPXBbNR5vGL/avxnSlZkrm1z628svUVv0IUwaxwy8iM6zxO9KkRhFbwm6G/YWPRRmo8NSE9z6JYuCTlEjLatb3eoj8XYQ2mYqxGJvTtwKJthX59k4Jd5e/TMZq09k2XJ1dkiaRj/YM4Vo+hyuFhU27FGR8a1o/9OTUzXQL/aoWSFDCHPVBVQ7mB8ryhju9cKOhS7a5mac5SDtccptZTS4whipTSD/nN0w9iiW7ZbWyDIrPgnpFMfHklNS5vkymHx1mNMhf0as/0cd2bd4CCINSTX273S82F0FKmPapOSkIEM24Zgs3k/1GYEGlm3t3DuX7m2oBBW0MsRplnJw+gb+fmuY5d1/M65uyag0t1oZ/yadfUCrdJMfHAwAeaZVyCIDSuS1QX3rj0De5aehd2j93v9zcQi2IhMymTv47+awuMUDhdYc+Zmj6uO1/vKkL1+P+QBLPK//BFp9+MNtZmwmSQ/VYSQ20I2yMxkl8G6PFxpvYX11BQ4V+NKZTxyRIkNWPp9ua2r3wfc3bN4Zvcb5AluV51KutlEjO/vJwbzruBG3vdSHtb+xYbV9cEG589OJqpM9dR4/Q0eQbDalSYOKAj/7q2X8g9bARBODPf7ikOeHsoKdOqrpNXbg8YSB3Xt3MMH943kmmz1vsKXDQSVJkMMook8dKNg7ikT/Ol/Uabopl92WymfTWNOk9dUBMy8E3Knhv/HD3iejTb2ARBaFxGuwwWXLmAR5c9SlFdES6vCw3/+YZFsfhScntO5teZv0YJMrNKaB1hD6b6dIrmySt684+v9jRYWSkQq1HhlpEpXNjr9D+Ezj+vPf/8ao/f7aGkGdpMCpMGNE8vp4n9O7FsX4nf2ZxQxmfSVcZQhq7rZz6Jr8yDDW/Ajg/BWQW6BqZI6H4hjHoIOg44s9c/xdw9c3lh0wt4NI9fzjCAw6iDu4a3d73N/L3zefWiVxmcNDisY2hMWvtIvv3VeBZsyOPNlYewu731zrgZFQlZkhicEsd947szrkc7EUgJQjgU7YA1r0DWt+Cq9W3tm6Oh3y9g2L0Q163ew8tr3QHT70JN6S6tbTplOqNTDCt/cwGfbD7MzB+yqLJ7UHUdr6qjyBJGRUKRZW4dmcLNI1J82RLNLC02jQVXLuDOr++kzlPnd4bqZBbFgizJvHjhi4zoOKLZxyYIQuNSY1L5/OrP2VG6g7d3vc2y/GXIkowsyXhULx6PBXf5+bgqh/DWPhsfLV3GpRlJ3Dkmle5BZG4JLU/SG8lpyszM1Ddu3HhaL/z2mkP8e/HeoAKq44HU7y7vdcaT08mvrWFjbkXA+2p3LaNm4+d4yvLrpRlauvQ+8RiLUWbTk5c0S6EDr6qR+fdvqWyg9Hsw4+tu9PD6llnoThcxkyYRc/UkTCkhFssoz4YvHoP8db7DXKq7/v2SAgYTxHaDK/8D3UaH+JX6m7NzTsA8/8ZYFAuvX/o6gxIHnfH7h0rTdFYeLOXHQ+WU1rqwGBU6xli4sn9HusTZWnw84SRJ0iZd1zNbexxn6kyuT8JZIm8dLPolVGSD1w36KTs/igkkGTpnwqSXIMGXUvv04r28FqBNhCN7EyUfPUXXX38aVEDVu2MUix8dF/RwdV1nY24F+4pqqHV5sZkUusRZGdejfaucm3SrbpbmLmX2jtnk1+SjyAqqpiJLvrFEGiO5LeM2rk6/us1UATsXrk/i2iSEwqN62JCXzz++2k5WsYrbbUY/5TyHQZZQZIleHaP417X96dNJtDdoaY1dm5otmALYmFPOS98dYP2hcnQd3OpPgdXxH4zeHaN55KL0M9qROtm3u4t55L0tIeW4nzymyUO68O9f9A/LWAJ5/pt9zFyRHXBVtSlWk8I/r+3LNQM749y1m6qFn1P95VeYkpOJuXoS0ZdfjhLbRJW/w5vhnavBXevbiWqKwQqTXob+U0Ie73HrCtfx8HcPhxRIHRdhjODLa78kwZpw2u8v1HcuTFZATFjavF2fwaf3QQONaOuRZDBFwM2fQvJQZq3M5ukle/0KDmmuOgpeuZWEKx4PKqV7VPcE5t9zbuzW7K/YT1ZlFjXuGmxGG50iOjEocVCb2z0/F65P4tokhOKH/Ue5b+6moOetNpPC67dkMqZHu2YemXCyxq5N4d9+OUlmt3jeuWs4hVUO5q3LY+fhKqqdHiLNBtITI5k2IiXsW5YX9EpkYHIsm3IrQgpYJHzd6pvjrNTJ7h6bxqdbDnOk0hmw51RDTIpMrw5RTOzfCUmSsPbNwNo3g6QnnqB29WqqPv+ckv88T8TIkcRcczWRY8cimUz1X6QsyxdIuaqDH7DXAQsfBlscpF8c/PNO8vLmlwMGUhUrKyj9uhR3iRvFohA9JJqkyUkoET+tKHs1Lx/s+4D7B95/Wu8tCMJZKHt58IEU+BZ+XDUw91q4ZxljeiTx3NcSnlPOC4Wa0n1lv45h/KJaV8+4nvSMa97PL0EQwmtbfiXT393k11KoMXa3yj3vbOT96SPo36Xtt8k5FzTrzlRrsbu9XD9zLQdLanEGkWYoSxBpMfDRfaPomRTV7OM7Uung2ldWU17nxhNEQGU2yKQk2Pjo/lFEN9KwVq2upvrrr6n6/HPcWdlEX345MddcjaXfsSIJr58PhduC25E6lSkSnsgCY2jnAXKqcpj8xWRcav2zCaWLSzm6+Chd7u5SryeKWqOS+odU5JNKoceYY1h+/XIMcrPG/j8b58LKL7Td69PPntcNz6aDq+o0niyhJ/bhh6v+xQNz8qmrC1ziu7VTuoXTdy5cn8S1SQiGpumM+vf3FFWHnrUD0DHGwpr/u7DN7T63Va22M9VabCZfYPTER9tYuqvYL8XwOFkCs0Gha7yNWbdlkhzfMmdhOsVaWfzYOO6ft4mteZWomo43QFBlUmQkybfb9t/rB2I1NX4GQImOJm7KFOKmTMFdUEDVwoUcfuIJJEkm/srhxFbuRjolkOr2Qg12Dxx6NJIIk+8XctZmN3O3e1h++yl9vnZ/BgNuCOlrXbB3AapWf8VFdaiUfFZC57s6E9XfF7ya2ptIfiCZ/U/sp2pNFXHjfpokeVUvqw6v4vzk80N6b0EQzkJ7FvqfjSLYa5GO++geFq35F1NHPMJ7K+WAqTFNVY5VZLhmYGcRSAmC0GpWZ5VS4wx8hr6xxuPHVTk8rM0qY1S6SPdrbefsJ4nFqPDyjYM5Uulg7rpc3l2Xi8erocgSmu5r+ntpRhL3jE1jQHLLb5PGR5h4/96RZB2t5a3Vh/h402E03VcdyqvqmI0yt4xI4ZaRKXSMsYb8+qYuXWj/wAO0u/9+HFu3on/+EGguCHBGWtXhxfVufj+2kZLr7lpY9d+Qg6ldZbvw6t56t9kP2NE8GtFD6h+gVCwKUf2jqN1VWy+YcqpOsiqzRDAlCOeCVS/4ricBBHMtMgLPGFPwXjiBjXvXsKcocLP2xkSZjTzWDE11BUEQgjVzRXa9isHHBdN4HHzpfjN/yBLB1FngnA2mjusUa+U3E3rxy0t6Ul7nptrpxWpSSIgwYTG2ft3+7u0j+fs1/fjLVRlUOTzY3SqRZgMxViOyfOZbt5IkYRs0CL46BIEXQHhilIlnVrt4YKiJWEsj71mR4/tzUpli3etFq6tDq61Fra3z/XddLVptLVpdHVW1BX5NhtVaFUOkAUnxfy9DjAFHbv1zFKquUuU+nZQgQRDOKtVHoGx/g3cHcy2SdQ12fYrxupm8e9dwrnl1NUcqHUGdkZUlX+bC/HtG0CGm+UuYC4IgBGJ3e1mXXeZ3eyiNxwFWHyzD4VabzFwSmtc5H0wdZ1BkEqMtJJ6l1SQNikxCpJlmqVmnesDTcB+SzE4K53cz8NwaF3+/sOEJhubyUvTw3diPKmjHAifd7UaOjESOiECJjECO8P23HBmJHBmBuZsOp9TBUCIVvLVedFX3C6i8VV4MkfV/LCUkoozNf5ZNEIRmVlsMihm8gfs7BXstQvOCx0GMzcoXD4/h/rmb+PFQOd4GUqYlfNVQ20eamXPnMFLbRfi/piAIQgspr3NjVGS8pxyDCKXxOIBRkamwu7GaQs9gEsLnZxNM/aypHpBl0BquFvPXC8yMnl3Ho8NNDT5GMptJuOt2ErqNRomMRI6MRLJYGj382HPl7zmQ/WW9Dt+2dBuSQaJ6UzUxw37qfaI6Vf6fvfsOj6pMGz/+PWVKeoGEGgihVxGQ3sSy9oJtbaigYvnt6uu+um51d93i67rNdQUUkbWs4tpFsaEgJVSl94QAIUASQur0c87vj5ESZpLMhHTuz3W9716eOXPmmYScc+7z3M99V2yqoMP11cvkx+gxdI5vnEbKQogmZNQwPX6KSM5FqGowILPFEO/QeXXGKPYUVjBveR7vfXcQVQFVVTAtC79hMb5ne+6dlMWoHqmyWFsI0ez8hkW45KNoG48rCvjD1AQQTUuCqbOBLQbqWFIwKF3jij46Ty330T8tfPNJRVVw9B8KHXtE/NE/7PdDvtz/Je5TSiBrsRrp16RT8FoBqlOtVs3PlmojeWz1NWwWFlO6TYn4M4UQLZQzuc5qopGcizAC4KieZtArPYE/Th3Mr68cwMFSNxWeYFPdDglOkmJrroIqhBBNLSnGFnatpxaTiOkqxzKNiAIqv2HWWuVZNA0Jps4GigKdhkDBd7Xu9tvJTobNqeQnY2opRNG+d1QfPbj9YNJi0thfsb/a9rTL0tDiNA4vOIyv0Icao5I4LJGMmRmotpM3ULqic3XPq4nRZQpbiFYvtUewAW8d6jwXdRgQnJ0Kw2nTGrx/oRBCNKSUWBspcTaOlFdPeXZ06Yei23Dtyo6o8Xj7eAfJ8rCo2dV9VRNtw7iHg72iatErVeWmgTaeXeMLfVGzw/A7Qa8l0ApDURTuO+e+sMFQ6qRUev+hNwNfHEj/Z/vT5c4u1Rr2Auiqzm0DbovqM4UQLZRmg/NmBM8ntaj1XGSPh/GPNNIAhRCi8SmKwj0Tsog5rRDaqY3HXbuyMf0eLCOAO2cdx76eV23fGJvGzIlZkrrcAkgwdbbodzlEMGX860kOqnxhcgIVBUbeW6+PviLrCi7sdiFOLbrqWU7Nya/H/Jruid3r9blCiBbovHuC55M61Hou6n9lIwxMCCGazg0jMjCt0HNc4sippEyZQVn2AvL/eSv5s+6k4tuFxPSuXpTCwmLq8K5NNVxRC0nzO1toNrjwd/DZ4+A/uX4p7+HqVfIyklQ8vzyt5KEeAwOvhZT6BTWKovC7cb9DVVQ+3/d5tfVTNXFqTh4f+ThX9pSbJiHalKQucM7NsGlB9OciWyxM+VXUM+RCCNHSJMXYuGdCD15anofbX71AWF2Nx2NsGvdM6CHrpVoImZk6m4y4E867O3hDEik9BroMh6uePaOP1lWdJ8c9yW/G/IaeST1xak7U09ZO2FU7dtXO2M5jmfuDuVzX57oz+kwhRAt12V8gY3SwOE6kbLFw7u0wambjjUsIIZrQTy7uywX900PS/WoTY9O4aEAH/uciaTzeUsjM1NnmoichNg2W/CFY4c8I3+8FVQfVFkwPvHZ2cGbrDCmKwmVZl3FZ1mVsO7qNt3a+RV55Hu6AmwRbAkPShnBj3xvpGNfxjD9LCNGCaTrc+l/48Eew9X0wfGDV0LpBsweLVkz4SfD/hBCijVAUhWd/eC6//3gbr6/ej2la+MP0ygOwaQqqonDLqG784rL+slaqBVGsMPmax40YMcJat25dEw5HNJnyAlj7Eqx5EY73gLIIrkcw/DD4Ohj9YLBqlmhTFEVZb1nWiOYex5mS81MbUbgdsv8Fm98OPrSxrGCXXQCUYMGK8+6GJFkbcDZoC+cnOTeJ+sgrruLlFXv57/p8NEU50dFGAQzL4sYRGdw1LpPu7aTpeHOo7dwkwdTZzvDDgTXgOgpmAGKSocsIcCbW/V7RKrWFmxWQ81Ob462A/HXgKQ3ORMW2g64jQa+98p9oW9rC+UnOTeJMePwG3+47Rqk72OQ8OcbGsO4pOKNIBRQNr7Zzk6T5ne00G2SOa+5RCCHOdo4E6FnzgmshhDgbOG0aY3u1b+5hiChIAQohhBBCCCGEqAcJpoQQQgghhBCiHiSYEkIIIYQQQoh6kGBKCCGEEEIIIepBgikhhBBCCCGEqAcJpoQQQgghhBCiHiSYEkIIIYQQQoh6kGBKCCGEEEIIIepBgikhhBBCCCGEqAcJpoQQQgghhBCiHhTLsmp+UVGKgH1NNxwhRBPobllWWnMP4kzJ+UmINqnVn5/k3CREm1TjuanWYEoIIYQQQgghRHiS5ieEEEIIIYQQ9SDBlBBCCCGEEELUgwRTQgghhBBCCFEPEkwJIYQQQgghRD1IMCWEEEIIIYQQ9SDBlBBCCCGEEELUgwRTQgghhBBCCFEPEkwJIYQQQgghRD1IMCWEEEIIIYQQ9SDBlBBCCCGEEELUgwRTQgghhBBCCFEPEkwJIYQQQgghRD1IMCWEEEIIIYQQ9SDBlBBCCCGEEELUgwRTQgghhBBCCFEPEkwJIYQQQgghRD1IMCWEEEIIIYQQ9SDBlBBCCCGEEELUgwRTQgghhBBCCFEPEkwJIYQQQgghRD1IMCWEEEIIIYQQ9SDBlBBCCCGEEELUgwRTQgghhBBCCFEPEkwJIYQQQgghRD1IMCWEEEIIIYQQ9SDBlBBCCCGEEELUgwRTQgghhBBCCFEPEkwJIYQQQgghRD1IMCWEEEIIIYQQ9SDBlBBCCCGEEELUgwRTQgghhBBCCFEPEkwJIYQQQgghRD1IMCWEEEIIIYQQ9SDBlBBCCCGEEELUgwRTQgghhBBCCFEPEkwJIYQQQgghRD1IMCWEEEIIIYQQ9SDBlBBCCCGEEELUgwRTQgghhBBCCFEPEkwJIYQQQgghRD1IMCWEEEIIIYQQ9SDBlBBCCCGEEELUgwRTQgghhBBCCFEPEkwJIYQQQgghRD1IMCWEEEIIIYQQ9SDBlBBCCCGEEELUg17bi+3bt7cyMzObaChCiKawfv36Ysuy0pp7HGdKzk9CtD1t4fwk5yYh2p7azk21BlOZmZmsW7eucUYlhGgWiqLsa+4xNAQ5PwnR9rSF85Ocm4Roe2o7N0manxBCCCGEEELUgwRTQgghhBBCCFEPEkwJIYQQQgghRD1IMCWEEEIIIYQQ9SDBlBBCCCGEEELUgwRTQgghhBBCCFEPEkwJIYQQQgghRD1IMCWEEEIIIYQQ9SDBlBBCCCGEEELUg97cAzgb7D/qYueRCiq9fmJsGl1TYhnYORFFUZp7aEKIFqDM7ee7/ccoc/vRVIXUODsjuqdi1+V5lxBCCFGXY1U+NhwopdzjR1dV2sfbGd49BV1r/OuoBFONxDAtvtpRyOylOWw5WIZdUzGxUFAwLYv28Q5mTsrimqFdiHPIr0GIs9HWgjLmLtvLJ5sPYddVTMsi+Igl+P9vHdWNaWMz6ZIc05zDFEIIIVocy7LYcKCUF5flsnh7Ych1VFMVpo3pzu2ju5Oe6Gy0cSiWZdX44ogRI6x169Y12oe3VQdL3dzy4iqKK7xU+Ywa94u1a2iKwkt3nsfIHqlNOEJxNlMUZb1lWSOaexxnqjWfn3wBk/95awNfbS/EFzAxajgP2zUFRVH48QW9eWByT5nNFm1eWzg/teZzkxCthdtncP9r61mTV4LHb2DWEM44vs/w+OXl/bl9TGa9P6+2c5PkkDSwAyUuLn92Gfkl7loDKQCXz6DCG2DavNV8s6uoiUYohGhOfsPktrmrWbz9CG6/UWMgBeAzLLwBk+e+2sOTC7c14SiFEEKIlsnjN7h+9kqyc4/i8tUcSAF4AybegMkfP9nBs4t3N8p4JJhqQFXeADfNyabc7a/1Bul0Hr/Jfa+tZ/eRikYcnRCiJXjs7U1sOliKx29G/B633+CNNQd4bdW+RhyZEEII0fI98Pq37CmsxBuI7jr6/JI9fLShoMHHI4t1GtC73+ZzzOUPGyFXbVtC+dr38R/NR7XHYEvPImnsjTi7DgSCUfZfv9jFrNuGN/GohRBNZd/RKj7ZfCjsBaCuc4Tbb/D0pzu46bwMbE2woFYIIYRoabYcLCM7p7he11GP3+TJj7dx+ZBOqGrDpc1LMNVALMtizje5uP2hqX3la96jbPXbtLv4QZw9hqFoOu6963HvXn3iF2xa8NWOQkqqfKTG2Zt6+EKIJjB/ZR5mmFnrSM4RECxs8/nWI1w+pFNTDlsIIYRoEV5avhefUf/raJU3wIqcYib0TmuwMcnjzQayNu8YJVW+kO2mt4rS5a+TetH9xPYdi2p3omg6sb1GkXL+9Gr7KsCba/Y30YiFEE3J4zdYsPYA/tMuAtGcI6p8BrOX7mnKYQshhBAtQrnHzyebD2GYZ3odzWnQcUkw1UDW7zuGL8yUo/fgDqyAj9g+Y+o8hidgSiEKIdqovcVVhCvGF805AmBrQTm1VWEVQggh2qLtBeVh+y9Gex3dsL+0QcclwVQDOebyEQizWMpwl6PGJqKoWkTHKXX7G3poQogWoNztRw0TTUV7jlAUJariFUIIIURbUO4JQJhnidFeR8MtyTkTEkw1EGeYSBlAi0nEdJVjmZH94hw1HEcI0brZdTXsRSDac4RpWdg06TclhBDi7GLX1eM97auJ9jqqqw17ry137g0kLdFJjC30x+no0g9Ft+HalR3RcTomNV6HZiFE82kf78BnhM4oRXuOiLFp6FLNTwghxFmmfbw9ZL0URH8dTYxp2Pp7ckVuIJcM7EiY4iKojjiSx99KyRezce3KxvR7sIwA7px1HPt6XrV94+waPxzZrYlGLIRoShmpsXRvFxuyPZpzhK4qXD20S1MNWQghhGgxBnRKJCnGFrI9muuoXVe5YXhGg45LSqM3kLQEB5P6pPHl9iOcvjY8ceRU1LgUyrIXULzwGRR7DI4OvUgcc1O1/Zx2jUkNWKpRCNGy3DepJ796fwtVvuqpCJGeI3RNYcb4zCYcsRBCCNEyKIrCvROyePrznbjreR1VgGljuzfouCSYakAzJ2axfHdx2IVt8QPPJ37g+TW+12lTuXt8jwZtIiaEaFkuG9yJX3+wNexrdZ0jVAX6dEigV3pCYw1PCCGEaNGuG9GV//tsR9jX6rqOagqM7JFKp6SYBh2TpPk1oBGZqfxwZAYxtsiqiRxn11T6d0xkxvisRhqZEKIlcNo0nr91GM4w6yvrEufQee7mYY0wKiGEEKJ1SHTa+OuN5+DQo5t8UICkGDvP3HBOg49JZqYa2K8uH0ClN8DCjYciLr3YrV0s86ePxK6r7Dq2iy/yvuCI6wiGZdDO2Y7RnUczutNoVEViXyFau4l90vjz9efw6NsbIypxrioQ79D5zz2j6RZmzZUQQghxNhnQKYkEpw1vpS+i/VUFkmPtnbVmjwAAIABJREFUvDVzNB0SG77QmwRTDUxVFZ6+bgiDOyfx98W78fqNkPURALF2DcuCc7snc6TcxZL8z3l1+8vklecRMAMY1sn3LNi5gHhbPNMGTuP6PtcTZ4tryq8khGhgV57TmU5JTn71wRb2FlfhN6yQCkV2XUUhmJLwx2sHk5EqgZQQQoiz29aCMn44ZxVVvkDE74mz67xz/xh6tI9vlDEp1unVEk4xYsQIa926dY3ywWcDw7RYuquQF5bmsvNIBW6/gV1X6ZQYw/TxmVx5TmdQfFzyxgzKrd0YeGs9nlNzkupM5eVLXqZzfOcm+haiMXkNL17DS5wehxZhs7kzpSjKesuyRjTJhzWitnJ+2n6onOeX7GHhxkPE2DVURSHRqXPNuV24fUz3Bs/tFqIlawvnp7ZybhKipSkodXPpP5ZR5vZH9T6HrjKsWwqv3z2q3rUJajs3ycxUI9JUhSn9OjClX4ewr/tNP9M/nYlL2Y1h1R5IAXgMD4ddh7n545t5+8q3SYuVyn+tUV5ZHq9vf50Pcj7Aa3jRFI2AGaBrQlfuGngXl2ddTqxNZiHOFv07JXL76EwOlLh5/8FxzT0cIYQQokX6y+c7qfSED6Sqti2hfO37+I/mo9pjsKVnkTT2RpxdB+INmGzML2XJrsIa78nPhARTzejpNU+zo2QHXrPuQOo40zIp95Yz84uZvHPVOyiKVP9rLY5UHeHRbx5l29FtGKZBwApOUZtWcN3MgYoD/Hndn3l67dNMGzCNB899UNbJtWEHKw+SV5ZHlb+K1bmVdEmTAFoIIYQIp9zjZ+GmQ2F7upaveY+y1W/T7uIHcfYYhqLpuPeux717Nc6uAwFw+QxmL82VYKotqfBV8N6e9/AaoYHUsWXHKP6sGF+hD82pkTg8kQ7Xd0CLC6aBBawA+ZX5fFv4LcM7DG/qobcIHr9BucePQ9dIcOgtvqT83rK9TFs0jQpfRbX1cKdzB9wAvLr9VXLLcnlm0jNNlv4nGp9hGnyT/w3ztsxje8l2bKoNCwuv3wTF4IYPX2L64Olc2O1CbFpoY0IhhBDibPTu+nzUMBMIpreK0uWv0+6yh4ntO/bE9theo4jtNaravhsPlLL/qKvBizlJMNVMPsz5MOysQ/GiYooWFdH17q7ED4jHf8xPwasF5D2TR49f9EDVg+/xBDzM3zL/rAqmqrwB3v/uIHO+ySX/mAubpmJ+v+bvskGduGdiFoO6JDXzKEMVu4u569O7KPOWYVHzGsVTuQNulh9czu9X/Z4nxj7RyCMUTeFAxQHu/uxuyrxlVAWqAKo/TLFgx7Ed/Gblb3hqzVO8cNEL9E3t20yjFUIIIVqOhZvCV8n2HtyBFfAR22dMncdQFFi6q5Dbx2Q26NgkmGomr2x95cQsxHGG26Dw/UK6zOhCwpBgY057mp2MBzLY9eguylaWkTIxBQALi5UFKznmOUaKM6XJx9+ULMvi2cW7mbU0B1VRcH1fHdEbOFlWeuGmQ3y27QjdU2OZfftwerRvORUP/7b+b5R6S0MCqbpmID2Gh4W5C5naeyqD0wY3x9BFA9lbtpdbP7mVKn/VibTOmrgCLlwBF9MWTWPuxXPldy+EEOKsd8wVvgy64S5HjU1EiSCLxxcwKXVFV7wiErIgoxlYlsXhqsMh2127XZh+k8ThidW2a06NhCEJVG6trLbdrtk5WHmwUcfa3EzT4qE3NzB7aS4ev3kikDqdYVl4/Aa7Ciu48p/L2ZRf2sQjDa/CV8FneZ+FpPYVLyrm8H8P0/HGjgx4fgBZv8rCd9RH3jN5mKcEiT7Tx7+3/buphy0aUJm3jOmfTqfSV1lnIHUqV8DFzC9mhj1XCCGEEAK0mERMVzmWGVlv18YoNSAzU83AZ/qCrZhPy/gyKg30eB1FC/1N60k67n3ukO2V/sqQbWfKMC02HDhGYbkXn2GS6LQxuGsS7eMdDf5ZdXly4Ta+2HYk4gbIlgWV3gC3zV3Nwh9NqDEvNmCYfLWjkA83FlBY4cWyLNrFO7h0UEcuHdQJu94wzxnCpXNGMwNpWiZf7/+aMm8ZSY6Wl8Io6vbWzreo8FdEPTMJwXTPlza/xC9G/6Kphy2EEEK0GO3i7OQUVYVsd3Tph6LbcO3KJq7f+FqPYddVkmPtDT42CaaagV21hwRSAFq8RqAygGVYIQFVoCyAHh/664rTGy6draTKxxtr9jNv+V48AQMFBQsLVVHwBUwm9knj3olZjOie0iRVBHcfqeCNNfvxBEKf5tdWAhOCAdWvP9zC/LtGVntfpTfAS8tyeXllHn7DpMpbPUhbtquIX7y3hVtHdWPmpJ6kxp3ZH92HOR+GpHNGMgN5PJgC0FWd5QeXc3nW5Wc0FtH0DNPg1e2vhhSaiXRtZMAK8EHOBzwy4hFidOk3JYQQ4ux09bld2FJQHpKhpDriSB5/KyVfzEZRNZw9zkVRdTx5G/Ds30TK+dNP7GtaMKVfeoOPTYKpZqAoCu1i2lHkLqq2PbZXLIquUL6+nKSRJ2chDI9BxaYKOlxfvZyjz/TRKb5Tg4zpvW/z+dl7m8EibPAC8OX2I6zYU8ygzom8dOd5JDgbt9rYS8v34jdDxxJJCUzTgpU5RzlS7qFDohOAw2UefvhCNofKPNXWW52q6vs/0nkr9vLedwd5897RZKVF3jH7aKWX3YWVVHgCxNg0iqpKQvaJdgYyYAYo9baMtEURnRUFK/AGqgdS0cxMAigofLr3U67tfW2Tjl0IIYRoKa4Z2oUnF24L+1riyKmocSmUZS+geOEzKPYYHB16kTjmpmr7jcxMpXNywz+YlGCqmdzS/xbmbJyDx/Cc2KbFaqRfk07BawWoTrXaE2tbqo3kscnVjjEsfRjtY9qf8Vjmr9zLU4t24PHXvp7DsoJ1+jfkl3H1cyv44P+Na7SAqtIb4P0NBzFOG1I0JTAV4LVV+/jJxX0pdfm49vkVFFZ4Mcy6K+r5DYuiSi9TZ61k0UMT6JRU8x+fZVmszTvGC9/ksGx3cTBF0AoOwOrqRjntRxTtDKSFhRFhLrBoWdYcXoMr4Kq2LdqZSVfAxfKDyyWYEkIIcdaKc+hce25X/rvuAIEw93HxA88nfuD5Nb4/1q5x78SsRhmbFKBoJtf1vg6T0OAl7bI0OlzXgcMLDrPt/m3kPJmDLdVGj8d6oNpO/rpi9VjuGnjXGY9jyc7CiAKpU/kCJvmlbqbPX4tlRVbqO1qrc4+iq6H/PKMpgekNmHy4sQCAH73xHcWVkQVSx1kWVHgC3PVyzd+zpMrHVc+t4M6X17B4eyHegEmFJ0CFN0CFJ4BphAZhp85Anur4DGTcgOqpm7qqy3qpVqrEHf3MZKAyELL9mOdYo4xPCCGEaC0euagPybG2qItIOHWVsT3bMaH3mU9AhCMzU80kxZnCxd0v5ot9X4Ssp0idlErqpNQa36uikuJMYXTn0Wc8jt99tK3GQKq2dUm+gMnWgnLW7C1hVFa7Go+/5WAZLy3fy+aDZVR5g6lv3VJjuWNcJpN6p9XYbLekyneih9SpoimBCVDu9rP/qIs1e0vwh2mbXdfaK8O02HfUxcb8MoZmVJ8ZLK70csU/l3O00hv22ACBioGo9mIU9eQNcrQzkIZpMKrTqNMPLVoBuxa65q4+ayPtesMvmBVCCCFak7QEBwtmjuH6WSsp9wQiekAeY1MZ3CWZf906rNHW+0sw1Yx+PebX7CjZwb7yffjNSOveK8TaYnnxohfDNv2NxsYDpRwq84R9LZJ1SW6fwQvf5IYNpj7fepg/f7aTA8fc+AMmximBUW5xFWvzSoixa8ycmMWM8VkhQZVlEb5IxyklMCMJqCwrmMYYLjCL5DsCeAMGL36Ty79uHXZimy9gcsuLqyiu8Iadbj7OXzoKe7ulIdvTLktDi9M4vOAwvkIfaoxK4rBEMmZmVJuBBBjUfhAdYjuEHEO0fB3jOqIpWrXS+NGujVRQ6BjbscnGLIQQQrQElt/Ek1uKWekH00KN0cnsnsgnD03gpjnZHDjmRgXCPc926ioWcO25Xfnt1QOxaY2XjCfBVDOK0WOYf8l8Zn4xk5zSnGrrp8LRFZ1AwMnDQ/9BRmLGGX/+i8ty8QZC1+JEui7JApbtKaawwkN6QrDIg2VZ/P3L3cz5JqfW1MEqn0GVz+CvX+xmZc5RZt02HKftZHCUFGsLO2sVTQlMgHjvERZku/Cb1RcuRbP2yrTgi+1HqPIGiHME/2QWbTlE/jF32EDq9NkuZzc7adfEEt+3evpeXTOQx20q2sTl713O9EHTuazHZcTawpd7Fy3PJZmXMHfzXAzj5N9ZtDOTTt3J1b2ubuqhCyGEEM0icNRNxcoCXGsPBxtDHX8grihYhomtVzK9fAoPXT+E1XtL+GBDARagKQoBM9jS5+4JPbjpvG5nXJU5EhJMNbMkRxKvXPoKC3Yu4N9b/02FryJkwXqsHouiKNzQ5wYGx1/Fz9/ex9gMFxmpZ3ZTvSr3KOEmVaJZl+TQVTYdKOPCAcFgavbSHF74JjfiNVhuv0F2zlEe/M+3vHj7iBMB1KgeqfjCVNyLpgSmTVM4f1A33txUcUbfEcCmKhRWeOnxfTA1a0lO2AbC4Wa7vAe+ouLbN0OCqUgFrAAHKg7w9Nqn+b81/8fjIx/nuj7X1etYomllJmXSN6Uvm4o3VdsezcxkqjOVoWlDm3LYQgghRJOzLIuKrw5Q/vWBYABVwxIK745j/FLRSdhQyrW3D+JPU4dQ7vbj9hskOHXiHXqTtPA5ToKpFsCu2bl9wO3c1v821h5ey8LchRS6CgmYAVKcKUzOmMxF3S86sf7iwGSd+15bzzv3j602mxOtcMEARLcuyTAtyj3BFMWtBWX8Y/HuqIpZQLAU+8o9R3lj7X5uPq8bS3cX8fqqfVhY4XobR1wCU1UUrh07hP9uW4X/tDFFu/YKBaq8we+5/VA5eUdDG8fVNNvlzLwYLbYXlvkyihppOmeo4/2qnlrzFIddh3lw6IP1PpZoOjMGz+DxZY+H9BuLZGbSqTmZPmh6k14UhBBCiOZQtjCXqjWHoYb2NcepgN0C794yimZtJP2Bc0iJs5NS67sajwRTLYiiKIzsNJKRnUbWut/0cZlsyi/l5+9u5i83nlPvGy29huIP0axLUhRO5KG+sDQXfyD8U4S6Cj24/Qb/t2gHs5bsITnWzm2junPvxCzumLcWtz806KurBCbAORnJ9O2YQCDMk41o115VeQ2mzlpJ1+TgbGC4WbPaZrsMVxauffcRkzEfRfGhaN6QfSLlMTzM3zKfjrEdZYaqFTg/43wmdJnAN/nf1JnKeyq7amdw+8FM7T21EUcnhBBCNL/KVQVUrTmMFc0D+YCFv8jN0f/soP0dA+vev5FIMNUKKYrCn6YOZurzK3klex93jM088drWgjKW7CyiqMKLAqQlOriwfwf6dEgIOU5qnJ1yT2gZ5mjWJVkuN/qHb5N/oC+fbjGqFZo4LtJCD1U+g59e2o9bRnY7ESBO6Z/O4u1Hop7tirVr/PaqgcTaNeKdOqWu6jNC0a69ctpUVj4+hZIqP89/vYfc4tCZqbpmu0xPF6p2/wwtfif2dkvRY/JBCf35Axxbdoziz4rxFfrQnBqJwxPpcH0HtLjgsT2Gh6fXPs3lWZfj1J11jl80H0VReGrCUzz09UOsPbw2ooDKqTnpk9qH5y54Dl2V07QQQoi2yzIsyj7bFzaQWpO/iT9+PYtdxXmoqkrvdt154oIfMbRT/+AOARPvnlL8R6qwdajfcoozJVfpVirWrjPn9uFcN2slvdPjOVLhYdaSHA6UuPEZxolmt7qq8Ozi3fRMi+f+yT25ZGBH9O9nkn44sht//3JXSKAS1bqkGAfnpsfw7mfrURz94bQSzlEVejAtVueWcOuo7ie2/e3GodwxbzUbDpTijjCgirFpzLl9OP07BZui3jU2k+eX5OA9ZTYpmu+oqcFqMKlxDlLjHPRMj0dVCFlvFtlsl4pR2R93ZX+6dPsOX+IHITfXxYuKKVpURNe7u1YrTpD3TB49ftEDVT+5pubTvE+5ptc1Ef1cRPOxaTaeu+A5Zm+czSvbXsHtNzAJDapi9Bgsy+L6PtfzyPBHsGmN0xRbCCGEaCk824+G3lQBFd4q7nr7cf5w8SNc2e98fEaANfkbcZzWdsQKmFQsP0jqdX2aasjVSNPeVqx7uzh+e9VAps1bw8/f3cyuI5W4/ScDKYCAaeHxB3tCPfb2Jm6du5oqb3A25KYRGVg1lPVOHDmVlCkzKMteQP4/byV/1p1UfLuQmN4nU9gcuspdE3rRceY9eK+6Dk+YXjjRFHqwgAPHqhffsOsqr84YxVVDu+DQVRx6zf9k4+wa7ePtvHnvaCb0Tjux/eZR3cJVWY/oOwLYVJUZ4zNP/HdavANHmLVqp852RcJn/yokkDLcBoXvF9L5ts4kDElA0RXsaXYyHsjAV+yjbGXZiX1dARfzNs+L6LNE81MVlQeGPsBHVy7GPHIdfZL7EWeLQ1M0YvVYspKyeHzk4yy9aSk/HflTCaSEEEKcFcqX5mN5Q5d05JYcAOCaAReiqRoxNgeTeoykf3rP6jta4N5QhOkNn+3T2GRmqhWr9Ab4y+e7sCwLt7/uxmUun8GGA6XcMDubd+4fg7rkC8Yc3sbytH4EwsTVkaxLumV0NyC4piicaAs9eMKsj9I1lf+7bgj/c2EfXl2Vx6vZ+wiYFqqiYGHhC5ic2y2F+yZlMalPOtppa8HSE5xcNqgjn245jOe0tU51fUebpjCsewq90k+mSV7QP51ffrAlZN9oZrvi7D58tuKQY7h2uzD9JonDE6tt15waCUMSqNxaScrEk0ss8yvzKfWUkuxMPv1QooX6asdRxnW6iNlXP97cQxFCCCGanb+gMuz2rNQMVEXlfz7+A1f1u4Bzuwwk2Rm6bAUAVSFQ6MaeUcPrjUiCqVbsR//5loOl7poqR4blDZjkFFbw/x6byy9yF/GHn/+KqV+VUlThDVsmvSYxNo3fXDXgRH+plDj7GaS+nZTgrPlpfMckJ4/+oB//c2EfDpV5KHP7cegq7eMdpNTRR+Cp64awu7CSPYWV1dL9amNTFTomOpl12/Bq29vFO5jSN53Ptx0O+b4RVxq0e7BrTjxG9QpvRqWBHq+jaKHFQfQkHfe+6vvbVBtlvjIJplqRDzYUcPvo7nXvKIQQQrRxllFzCfQERxzv3vocz6/+D499+meKqko4v+conr7kMdLiTquGq4DplpkpEYW84ipW5hwNGxjUVTnPa1gsi+lK3PzX6ZAaz9t9XFw/eyVHK31hm9CezmlTefjC3tx0XrcT2wZ2TiTGplF1Wrn1aAo9OHSV8zLrLmypayoZqbFE07bYadNYMHMM019ey5aCshrLwh8XY9Po3i6W/9wzmqSY0ADv3klZLN1VVK9Kgw5d5epzuvJZeZgqg/EagcoAlmGFBFSBsgB6fOifrCrZuq3GkXIPWw6WcX6/9OYeihBCCNH8VAjbB+d7vdtn8rfLfw7AnqP7+PHC3/Obxf/kX1c9EbKvojdPGxG5C2ulXl6Zhxkm8Clf8x4li18kafSNdP1/r9Hl/pdJGHYZ7t2rq+2naBqvrTsIQEZqLIsemsjFAzpg11WcYdYlKUqwQl5GSgzP/vBcZk6qnq86sXcaTnvozNOpqW+uXdmYfg+WEcCds45jX4eu97l9dGY0P4aoxDt0/nPPKP40dTADOiXitKmcGq+oSjCIymwfy2+vGsj7D46rsXP2sG4p3DCiKzFR9vnSVOiSEsPDU4biN0J7TsX2ikXRFcrXl1fbbngMKjZVEDegeqUan+kj0VE9JVC0XB9tLODigR3PqD+cEEII0VYoioISE9ncTq923blx0CXsLMoNfdGwUBNqz1JqLDIz1Qr5Aib/XXcA/2nBVDSV87wBk1ey9/E/F/ZBVRVS4+w8f9twSqp8LFi7n/+s3s8xl5+AaRJr1xnePYWZE7MY3j0lbF8rVVW4e3wP/v7l7pDZskhT30ZnpdIxqXHLfOuaytVDu3D10C7sOFzO51uPUFjhwbQgPd7B5H7pnNM1KaLeXb+5ciDlbj+fbT0SdobqdHZdpVOSkzfvHU37OCcD2g1gc/HmavtosRrp16RT8FoBqlOtVs3PlmojeWz1dL6eST1JciRF90MQzebDjQU8+oO+zT0MIYQQosWIG96BypUFIel+e47uY3FONlf1m0KnxHQKyo/wwfbFDOsc2lNKS3ZgS4ttqiFXI8FUK1RU6SVMO6eoKucBuHwBKryBamlsqXF27p/ci/sn94p6XDeP7Macb3LDph7WlfrmtKk8clHT3mT265hIv471n9VRVYW/3TSUF5bl8txXezBNKyTNEcCpq5jADwZ04I9TB59YFzZj0Ax+vvznuALVKximXZaGFqdxeMFhfIU+1BiVxGGJZMzMQLWdnDWM1WOZMXhGvccvGodpWizdXcR/Vu0nv9SFL2CS4LTRv2MCB0rcjMlq19xDFEIIIVqM+DGdqcwuCNkeZ49lQ8F2Xlz7FuXeShId8VzYcwy/OP+BavspdpWESdEs/mhYEky1QpWeAFqYBM1oK+fpqkrlacHUmUiOtfOfu0dz/eyVda5JOpXTpvLU1MGck9H6iigoisLMiT25a2wPPt92mDlLc9lTWInHb2DTVdrF2Zk2pjs3ndctJGVwUsYkbGr4n33qpFRSJ6WGfe04VVG5sNuFDfZdxJkJGCZzl+9l7rJc3D4jJLDenF8KKNwxbw2PXtKPoa3w37sQQgjR0PRUJ/ZuifjyyqtVMuuUkMasa35b9wEUhdhz2jfiCGsnwVQrFGvXwlbei7ZynmFZxDbw2o0BnRN55/6x3PLiKrwBs9agyqGrqIrC3246h0sGdWrQcTQ1u65yxZDOXDGkMwCWZdWZKqirOn8Y/wf+d+n/hvSbqotTc/K7cb+TXkQtRJU3wF0vr2XzwZqbSwezFyxW5Bxl/QvZPH3dEK4a2qVJxymEEEK0RO1+2Jcjf/8WwxUgmjISik2l3W39UZpxLbIUoGiF2sc7wlbdi7ZprKYoJDbQrNSp+ndKZNlPp/CzS/vRNTmGGAwcioWuKjh0lXiHTqJT5+4JPfj6fye3+kAqnEjWXEFwduqnI3+KQ3NEfGyH5uCREY9wUfeL6js80YD8hsm0eWvYmF9zIHU6j9/ksXc28fnWw408OiGEEKLl0xIdmLf0pQwLK8JoSrGppNzQB2fvuitBNyaZmWqFYuwaPxjQgY83H6o2QxVN01hdVZg6rEtIg9uGEu/QuX1MJreN7s5HN9/H0RvuwNu+A7F2jS7JMUzsk4YtXK7iWej6PteTFpPGL1f8Ep/hC1lDdVycHoeu6vx23G+5oNsFTTxKUZNnF+9ma0FZxP3LjvP4TR56cwPLfno+7eMjD6aFEOJsZ5oWlb4Adk2V6qhNoWQvHFwPnlLQHJDQCXpMBL3hque5fAHuXriFWy/szGVHfLi3lwRLSZ/+kFIFVBVbx1iSr+yJo3vzVzSWYKqVumdiFl9uLwypIhdp5TxdVZg+vkejj9Nyu+mzcy19rvwnqr15Sla2BpMyJrHkxiUsO7iMeVvmsbFo44n1VH7Tz6B2g5g+eDqTuk5CV+XPtqXwBUzmr8zDU8OMVF0930zL4o01+/nRlN5NOWwhhGh1PH6DTzYfYtaSHPYUVaKrCoZp4dA1rh7amRnje9C7Q0JzD7PtMA3Y/Tks/zsc2gCqHtymqKB+3xxqxAwYeTckda3xMIfK3LyavY9VuUcp9wRw6CpdU2K4dVR3xvdqj6oqmKbFT97ayKAuSdx2QU8URcGo8lO19jCudUcwXX4sC1SHhrNPCvHju2BLb57KfeHIXVkrNaRrMj3ax7HzSAXGaSl/dVXO01WFIV2T6ZkW39jDxLN9O45evSSQioCmakzOmMzkjMn4DT9lvjIAEu2J2DX5+bVEn287HLbfGwR7vpWtfpt2Fz+Is8cwFE3HvXc97t2rTzbQDpjMW76XByb3arRZYiGEaM0sy+KFZbk8++VugBPFffzfl9F2+w3+u/4A7393kD4dE/jXLcPISG05N9qtUlUx/PsqKN0Hvsqa91v1L1g9Cy55CkbcVe2ljQdKeebznazZW4Jlgc84+dBxa0E5y3cXE+vQmTkxiwp3gCPlHt64d/SJZRJanI3EyRkkTm6+Kn2RkmCqFXvpzhFc+o9llLn9YUulh6Mpx3tKDWvcwX3PvXkzMYMHNclntSU2zUb7mOarTCMiM39FXthy+NH0fPMZJtk5RxnfW37fQghxKsuyePydzXy4saDWfo6GCYZpsuVgGZc/u4w37h3NwM7Sg7FeqophzgSoLALTX/u+hi/4v5/9DDxlMP5hAD7ccJDH3tlUY9YGBIPiKp/B05/uBCy+eGQSDr11pmzKopVWrFNSDO89MI60eAc2re6n2nZdpWOSk3cfGNtkazQ8m7fgHDS4ST5LiKZ2sNQddns0Pd8sEwpqOI4QQpzNnv5sZ52B1KlMC8o9AW55cTX5x8KvPxa1MA145arIAqlT+d2w5E+wcxGfbz1cZyB1Kp9hYgGPvb2JgBHd2uOWQoKpVq5H+zg+fXgid4zNJM6uEWcPjepjrUCwet74Hnzy0AS6pjTe9LdlWVR4/BSUujlW5aNy8xZihkgwJdqmmopORNPzzbCsiG8UhBDibJFbVMm85XvDnh+rti3h0L8fZv9fryf/uds58tYTePK3nni90hPgiQ+3hrxP1GHPYji2r1oglfn3CtL/XEGV72QK1NxvfUyeX1X9vQEPgUU/56E3v4s4kDrOb1hszC/lua/3nNHwm4uk+bUBqXF2fnn5AB79QV8WbT7Mx5sOcbTKi4JCqm4w6uN/M+3t2dhtjffrPlblY8Ha/cxdvpdSlx+bpmKYJtY593PV+gpmxJYzoHPzV1wRoiHF2TUkIBVsAAAgAElEQVRKqkK3R9PzTVcVEpxyKhZCiFO9vCIvZE04RLYe1bAslu0qprDCQ3qCM+QYB0pcrNhTTJnbj/r98ofz+6WTGneWr09e+Y+wa6QMC/6x2sfPJ9Se1WSVFzDQ2sM6skJeq6sgk8dv8vKKPB48v1erq/YsV/A2xKFrXHNuF645t3oj0D3v/B4zZw/069fgnxkwTH770TYWrDuAqnDiaUTA/P5Jkqbz/sZDfLzlML3S4pkzbQRdkmMafBxCNIfBXZI4WOoOaaJ9as+3uH7jaz2GaUG/jvKgQQghjnP7DN5enx/SUzOa9aiKAm+s3s9DF/YJvte0WLqriNlLc9hwoBRVUfAbJopC8AHwexZT+qVzz8QshnVr3r5FDerQpmBZc2856M5gWfPeF4HttHuxY/sgf23YQzw61s7TK7w8cJ6dZGfNy0pU08udykes46Fq2yMJgCF4T7l4+5FW139UgqmzQNzYsVStWImzgYMpb8Dgznlr2XCgFF8tPXYM08IwLbYfquCyfyzj7fvGSPlS0SbMmJDF17uKcJ9WhCKanm8ZqTEyayuEEKdYlXs0bIXTaNajegMm7353kIcu7EOlN8D0l9eypaAMV5iiQX4juO2zrYdZsrOIK8/pxJ+mDmm9VVb9Htj2Piz/W7AiH0qwWISigWYHTBh6G4yaCe16Bt9T8B2odsAbcrgRnTUmZ+o8s9LL76eEzvQdp2ExXN1VbVs0AXCVz+DlFXmtLphqXfNool7ix42jasWKBj2mZVk8/OYGvjtwLOL1HoZlUe72c9MLqyis8DToeIRoDsO6JZOeED7tIXHkVFKmzKAsewH5/7yV/Fl3UvHtQmJ6n7wJiLNr3D+5Z1MNVwghWoWjVT7MMGWKo1mPClDm9uPyBZj6/Ao25JeGDaROZVrBUusfbSzggdfX19j6okUr3Q/PnQcLH4GiHcHiEH4XmAEwvOCrAF8VrJsHs8bCmheD7/OWg1Xzz+d35zv45xofRVW1r4eKOy0YiyYABsg/1voKMsnM1FkgdtQoCh77KabXi+pomCp+2blHWbqrKOwiw9ryYi2g3O3n6U938swN5zTIWIRoDKZpseNwBSVVPgzLIinGRr+OCThtJy/iiqLw00v68ciCDXjCzM7W1vNNVSDBqXNpK3sCJ4QQjc20LAgTx0SzHvX4cf7ff75j31FXrRk0p3P7Tb7ZVczfvtzFTy7uG83Qm1fpAZgzKVimvJbACAgWmTD98MWvwVsBiV2CDXlrMChd44o+Ok8t99E/reb9/FT/vUQbAHsDra8gkwRTZwEtIQFH3764168nbuzYut8QgRe+yQ1JbYLI8mIDpsXCTQU8ceUAEpy2BhmPEA2l1OXjrXUHeHHZXqq8gRNpHpYVnJH94chu3Dk280RTyIu7OrnuyLe8024wHiK7WCgKxDt0FswcUy04E6IlMAyD3NxcSktL8fv9OBwOOnbsSOfOnU801BSiMSXF2FDDpNhFsx4VwKGrrNxTHLbyal0FEdx+g7nL9nL/5J7E2lvB7XLAB/MvjyyQOpXfBUufhon/G7w41eK3k50Mm1PJT8bU/GC+yEqu9t/RBsBxreFnfZrWN2JRL3Fjx1K5YkWDBFNHyj2szDka8tAoqoWhKLz77UHuGJt5xuMRoqG8uWY/T3y4FVUJPpkM55XsPF5btY8bhnflV2M7cPDuGfx40mQ6D+/PP7/ag9+wwlagOi7GppHgDAZS3dvFNdI3ESJ6FRUVrF27ljVr1mBZFoZhYFkWqhp8Cp2QkMD48eMZNGgQdvtZXvVMNKqRman4w/QcimY9qmYaxFeUUqKEtoOJtCCCosAHGwq4eWS3xvmiDWnHR+AqrhZIZf69Apcf9j4UT5w9GCjN/dbHa5v8LLnzlOtPwI214Q0sv1Hr+p9eqSo3DbTx7Bofg9ND93RZDl41Lqy2LZoAWFVgUNfW12xZgqmzRNy4cRx+8skGOdbSnUXoqoLvtO3R5MW6/QbvfSfBlGg5nvtqN//6OqfG3lHH+Q0LsHh7fT57Pl/KPy69jPT77+NBRWFKvw7MXZbLwk2H0FQFj9/ANC1suopNU0mNtTNzUhbXDutKvENOv6Ll2L59O++++y6WZREIBKq9Zny/OL+kpIRFixaxePFi7rzzTtLS0ppjqKKNsyyLnKJKOiY52Xc0tPFu4sipqHEplGUvoHjhMyj2GBwdepE45qZq+2m6RoEVg3Has61oHvy6fAazl+a0jmBq+d+Da6FOE3FZ88Icqqy+xDtyUczT7/BO+vUkB69uCt/QV1Us3jeqB0zRBMAOXeOeCaFl1Vs6uZqfJWKGDMafn0+guBi9ffszOtYxly9s7nG0ebGlrpr/WIVoSh98d5Dnvt4TVaNBT8BkQ3J3nuvcgye/T43o3ymRv9w4lCeuGsjnW49QkFdA4UefkDljGudkJDOie4qkSYkWZ+PGjXz00UchQVQ4fr8fv9/P3LlzmTFjBunp6U0wQnE28AVM/rvuALOX5nC0yldrsYja1qMe1zM9gQPHXHi91Y8TbUGE/UddmKYVNu2wxSjcDsW7w74UaVlzxaaS0Lc37MmDUy6FeQ9Xr76ckaTi+WWYCrSag4IuV2DlxcNpv7tIA+COSU7OkZkp0VIpuk7sqJFUZa8i6corzuhYwefyoWlM0ebFtsIaOaINMkyL33y0tcZAqra8eo+l8ta6Azxwfk86JZ3s2ZHotHH98K54U/3kv7CKnhN+00TfRojo5OfnRxxILVmyhJKSEqZOnYrX62X+/Pn8+Mc/xumsuVSyEJEoc/u5Y95qdh6ujLhCcG1ibBp3juvB7z/eFvJatA9+dU2h0hcgsSWv8T60EdTwCXqRljVXLAOObIGr/gkf/ThYBTBSig5JXel841+J/8c63H6D04sx1hUAx9g0HvtB31b5wFGCqbNIsN/UijMOppJjbNg1DbdZ/YQX7cLQ5JgWfGISZ42vdxTWWOUpkrx6y4JXs/fx2CWhfdwUVcUKU95XiJbiq6++CgmkNm/eTHZ2NsXFxSeKT0yYMCHkvX6/n++++44xYyJ7wi9EOG6fwY1zssktqvw+jfrMOG0qz91yLh0SnSE39BD9g1/DtHDqLbxQkKc8WPq8Br8738G4eVU8NKqOtY6+ChhyY/B4n/8SAhEEVJoDkrrCXYtwxifz5r2jufpfK6j0BsL+/MOJsWnMGN+DSwe3zuq2EkydReLHjuXo7DlYlnVGkf/43u3D9n+IJi82xqZyxZDO9R6DEA1l9tIcqsKkk0SaV+8zTF5dtY+HL+yDXQ8+GdxTWMnrq/exc18Rx/rcQNq81QzuksSto7rTOTkm5LOEaA5lZWXs27ev2rbs7GyWL1/OFVdcQc+ePdE0jT179rBjx46QohN+v5+VK1cyevToVvk0WbQMP3t3E3nFVWccSMXZNWyaygvTRjCyRypFFd6wD8qiffAbY9dOnNtbLN0RbMhbg0jLmqN9v65q5N3QLisYUJXkBisFnl4h0BYHWHDOzXDRb8ERTAfMSovnvQfGceM/vsLlt/CoNYcamqpg0xR+dEEv7p/UensuSjB1FjmSmMacHlNY8fvPqQwE0+wSnDo/GNiRu8ZlRlxZrGtKLMO6pZCdezTktUjzYk0LbhyR0RBfS4gzsuFAadjt0eTVm98vmD54zM2zX+1m15EKAoZFwLQgvhPbdxWzKqeEucv2cl5mKg9f2JsRmakN/VWEiMratWur/bfH4+Hrr7/m6quvpn///ie29+3bl759+7JkyZKQY3i9Xvbu3UtWVutbNC6a39FKL4u2HK6x8E9d5csBbJpCv46J3D+5JxcN6IBNCwYLaQkOhnRNYt2+Y9WOGc2DX5umtI57laQuUMcsWyRlzUk85SF3zylw/0o4vBmy/wX7VoKvEjQbxHeAETNg8PVgD7137OY9xktL/sK3P/8bczcdo8ztx2+Y+A0LVQGnTcMwLS4f3Im7J2QxoHOYNVitiARTZ4E9hRX88v0tfLe/FKPTMAJVJ6eCK70BXl+1jzfW7GdwlyR+f+0g+nWs+x/1zElZbKyhm3hdebGqAj8Y2JGkWEnzE83L4zfCzrJCdHn1CvDPxbv5emdRjfn+vu/L/C7fU8z6fSX87NL+TJNqlqIZ7du370SlPgiunwoEAtUCqboEAgEOHTokwZSolzfW7K+xtVGk5cvHZLXjlRmjwh7jvkk9eejN70KyDyJ98KsqCne2hvN0j0m1NtyFusuaY4+HkTNDt3ccDNfOjngolmVx+Ikn6DZ9GudePozpl1msyi1hy8EySt1+Yu0a6QkOfjCoY8tehxYFCabauDV7S7jr5TW4fMb3BR9C/4D8pgWmxbp9x5j6/EpenDaCcb1qr/g3sXcaw7qlsDavpM5S0qeLd+g8dkkr6igu2ixNVWoshBJNXr03YLJ4R2HEfwtuv8mfFm1H0xRuHdU9ylEL0TA8Hk+1/3a5XMTGxp7oKxUJ0zRxuULLVwsRifkr88IW/4mmfPmqvSWUufxhH9Ce3y+dWIceNpW7rge/mmUytFNijVk7ZS4/H20qIK+4igpPgORYG/07JXLJoI5N34xds8F598LKZ8Hw1rhbbWXNARhwVeSfaRrBRcNa9VCi7P0PCJSWknrHHQAoisKYnu0Y07Nd5MduZSSYasO2Hyrnzu8DqUi5fAZ3/3sdC2aOZkjX5Br3U1WFF6eN4KYXstl1uAJPBDeRihLsbP363aPpmhLaRE+IpmbTVJy6FnY2KZq8+uO9p6Lh9ps8uXAbw7ql0L9T605xEK2Trle/BYiNjcXlcmGaZlQBlTTwFfURMEyOVoVvkRJNmrVdU8kvdZEUG1pSW1MVXr7zPG6YnR1VlUBFgUTF5CcfPYNnSnucp8zWbi0o44Vvcvl0y2FURal23Fi7xs/f28xN52UwfVwPMlLP7F7H7TP4aGMBC9YeoLjSi2FZJDp1pvRL5/YxmXRIPKU633kzsLKf49SJvojLmuvOYNqeXnsvKg6uh5XPwa5F4P/+YYxmh8wJMO7HBBIGUPjMM2S8MAdFP3tCjLPnm55lLMti5qvrowqkjnP7De59ZT0rH59Sa1+FGLvGf+8bw6P/3cRnWw9jBgz8hO6vfL9vWoKDeXeeR8+0+KjHJERjuXRQRz7YUIBxWrpfNHn1Nakr398fsHjxm1z+etPQRvluQtQmOTmZQ4cOnfjvrl27ous6O3bsYMCAAREdw2azkZgoDwNE9Fx+A11VwhaeiCrNWoEqb833OoO6JDH/rvOYPn8trjAlu0+nqwqpcXbemjmZlNU6+6fPoMPPf0bSlVfyyso8/rhoO/6AFXLNAE7cc722ah8L1h5g1m3DmdQn+ubWZW4/z3y2k7fX56MohNzL5RRV8eKyvYzJasfPLutP344JmI5USstHkuxchqpGkTGk2SGtH0z5Rc37HFwP782EsoMQ8IB1yvENL+R8CQeywQPpV15FzMCBNR+rDZJgqo1at+8YxZXhp3ojWdBZ4fGzIqeYCb1rPwk4dI1nbz6Xg6Vunv3fv/Jxx3MxFBVNVTAti4BhMaFPe2ZO7Ml5mdKwVLQ8Myb04JMthzD8oRfGSPPqw4kk39+wLD7efIjfXD2wzeSOi9ZjxIgR5OTk4PMFZwecTieTJ0/mk08+QVVVevbsiaqq5ObmkpeXh80W+m/Usqyo1lgJcVysTQsW6QkjmjRrywouH6jNqKx2fPSj8fxp0Q6+2VUEEJKWHWvXMC2Lq4d24bEf9KVdvAMuvxxHr17k/+jHzF9XwGyzW0TN3f2Ghd8wmPnqOubcPiKqgOpgqZsbZ2dTWOGpscLh8bEv3VXEmrwSnr+2L5nP/Ao9rSfJU6+AL38VDHrqojshvT/c/l7Ns1K7v4C3bq+775SvCl2FpMB7sH4MDL+j7s9vIySYaqPmLM3BHWZWKtIFnVU+gzlLc+sMpo5rX3yQu3d/wZOzfkpRlZ8Kj58Ym0a7eEedJzkhmtPAzklktotj5+GKsIl6deXVhxNNvr+qKny0sUDWTokm16NHDxwOx4lgCmDs2LHEx8fzzTff8O6772K32+ncuTMTJkwgJyen2vsVRWHAgAHStFfUi66ptI9zUBTmwW80adY+w6RLBC0nstLieXHaCIoqvLyxZj+fbT1MmduPpiikxNm5YXhXrjm3C3Gn3bM4+/al6C8v8vxr3+FVolsj7vGb3P/aer54ZFJEYzxW5eO6WSspKveGnfk6nUVw1mrmGxt5bsBoLvzZAyiqCh36w5dPwJGtwf5Tp/egsscHS6mPvBcmPVpzIJW/PrJA6hRKwAOLfgpxadDvsojf15rJXW4b5A0YLNlZFHJjGM0NHsDqvUep9AYiCobKPv6YxEsvxW7T6ZKsA9JLR7Qez958Ltf8a0XUabGKQtiUkWjy/d0+g7ziqqg+V4iGoKoq48aNY/Hixfj9JxelDxkyhCFDhoTsn5FRvUS0ruuMHTs2ZD8hInXHuO48t3hPyLrrSNOsFQXO75sWVXXgtAQHP76gNz++oHfE73k2uwCvEv5eqM50bsPk5RV7+eXldafOPvbOJo5WRhZIncqr2vhJoDdrTQuHCvSYAPd8BUW7YNUs2J8N3vJg0JTQOdhHqt8VwcIVNbEsePuukEAq8+8VuPyw96F44uzBbKO53/p4bZOfJXd+X6wj4IZ374FHc8DW9h+2/H/27js+qip9/Pjn3jt90iHFhFCSUENRegfFtRcWFVR0VSyIBd1d26qr636tq/7WsvaGggXXvip2qdKLdAlggBBCKinT5977+2MgMMwkmYE0kvN+vVxf3tyZObOvuTP3Oec5zyOCqTao0unDoEghy+fR3OBBYHN+hcPbYDCl6zpVX31NxtNPH/OYBaEl9UiNZdY1Q7n6rRW4aitf1s9ilEmOMbOnInTGLpp8f4Aqd92d6wWhKQ0dOpRt27axe/du/P7IP4dGo5Fx48aRlpbWhKMT2rrLhnTm+R+3h/1bJGnWVqPCDWObtiz/nnIna3ZXhP1bJNk+PlXn/RW7ufPMnpgNdf8mFFe5WbCtpM7UvoaCNlXV+WZjEReenHH4Qck94Px/H9sb370UHKVh/6Tq8OxyL/eOaaBgxebPYUDDafEnOhFMtUEev4YUphBEtDd4EqE5xeG4N25EkiQsfdvXhkOhbRnaLYnPbh7FXz/8lbz91fj8KupR19GhYioJNiOPTOzHrF/ywwZT0eT7AyTZRDU0oWXIssyll17K+++/T0FBQdAKVV2MRiMjRoxg1KhRzTBCoS3rEGPmnH4nMW/DvrBVgetLs1ZkicxEGwM7J1Lp8vF7qaN2i8FJCdaI0uoi8e7y3WH7EUaV7aPDt5v2c8GAdOoye9muMHduAZEEbQ6vykvzdwQHU8djyXPgC9/24M6RJv61xMNNQ0wkWOoYtbcGlvxbBFPCiSnOasSvhX4pRXuD59d04qwNf0SqvvyKuHPPFcUlhBNej9RY/nfraPL2V/Pcv97lZ0snnFqgF5XFoDA8K4np47IZ1i0JSZJYvauCpTvL8B51ExBNvr/dpNAjTVS4FFqOyWTiiiuuYMmSJSxduhRVVYP2UUFgf5TBYCA+Pp4JEyaIohNCo3lsUj+2FlWxo9hR29y8ITIa8VYL95zdk5nvr+W7zfsxKTKHohGvX6NnWiw3jsvmD31SMSqRl/o/2oa9lWFXi6LJ9nF4VfKKqmFA3ee8v2JP2AnsaIK2/DIHu8ocdfbGipimwfbvqavlx+B0hfFdDTz1i4eHT6snja98J1Ttg7iTjm88rZwIptqgOIuBjjFm9lUGV3KJ5gYPIMZioKO9/iVcXVWp+vprOr8963iGLAitSk6ynVsXz+L/ffE5hpQUdJ2wbQKmDu/Ma4t2hhyPpqy6Dpzdt23/0Aitn6IojB07llGjRpGXl8fy5cvZsLOQ1BgjFouZk046iREjRpCR0Uiz3oJwkMWoMHf6CK56cwVb91U32A/KbJCJdVaTYoeb3l2Lx6+i6aGZNOsLKrnzo18xKTJvTxtab+/M+lS7w6/WRpvt8/mve9lT4cRsUDAZZMwG+eC/FYyKRJkjfAXmaII2oyKzr9J9/MGUtzqwIa0e/zzVzKg3Hdw2rJ7MCsUEzlIRTAknHkmSmD42iye++S3oSymaGzyLUea60d3q7TMF4Fy5CqVjR8xZTZuzLAjNyb1lC0pSEsbUVKDu35ST4q0M7ZbEorzQvPJI8v2NisSUIZlYjJH9GAtCU1MUhV69etElqzuP/PM7frvjLJF1IDS5OIuRuTeM4JM1Bby0YAcl1Z6Q/at2k4LFqHDp0Ew+XZlPXqUHn1J/4OXwqDhQmfLKMt68eggjsjtEPTa7KfytcrTZPr3S4hjbIxmvX8Pj1w7+W8Xr13B46u5/FW3Q5o6iOXGdVD9Q/2pe3xSF83oYeHyxl97JdZ0rhVYSbINEMNVGTRrUice/2RpyPNK+OboOU4Z0rv3vCoeXuav28NPWYqpcPgyyREqchT/k/cJp57SP0pdC++Fctgz7iOERnXv76d1ZmV8etvdIQ2XVDbLMtFHdjnmcgtBUSms8JMeYRSAlNBuTQebSoZ2ZMiSTtXsO8OWv+9hf5UbVdDrGmDitdwojszsGSoe7NHxK5LewLp/KtW+v5H+3jiY7Obq06u6pMaz4vYyjM/2iyfaxmRQm9E5h0sBOdZ4z65ddYdMcow3aYhujZ6ElHlRvg6c9NN7CwFdq+OuIOrKYND9Yjm1F8EQigqk2Ks5i5PoxWby+6PeQJfOGbvAsup/LB3cjyW5iZ0kNT3+3je+37EeWCL5hLKxiqb8LjxfYuOr735gxLgerScywCyc+x7LlJFxycUTnDuqSxL1n9+axeVsbTE85ksUo8+IVA8lMsh3rMAWhyRRXe+gY20ClLkFoApIkMbBzIgM7J4b87av1+/i91BF2D1ND1e5cPpUnv9nKy1cOjmo8Vwzvwoer9qD6jq18O4Cm65zbv+7iExAI2jYVVoUcj6rnll8jJ8pgMSzFACf1h32/1ntaTpLMlFwjz63w0i8lzOqUyQ4Jbb+Hogim2rC//KEHO0sd/LSlOOKbPKtRZqDnAFfMfZzF3R5l+seBG8Q6mpTjMphxeQINfr/dtJ/3rhsW6BouCCco3evFtWYNGU/+K+LH/GlkVwyKzD+/3IRP1VHrumAIpPYZFJmXpg5kfM+UxhiyIBw3TdNZmFfCO0t3savMQYXTh8evcddHvzJtdDd6pcW19BAFgZcXbA/bDzCSane6Dj/9VkJZjSeq+5QeqbFkJ4cPdCLJ9lEkuGBAeoNtZqaPy+Zvn6zH4Ql+f5EGbbIEZ/RJjarnVr1G3Q5f3BqoylePB8aZmb0+zL4ygxWG3wzysRf/OFFIej2NwQYPHqyvWrWqGYcjNDZN03noy03MXbkHv6qH9J46RJHAaJCZeHIGD1+Yy+InX2J6eSc89TV0O4pRlsjsYON/t4wO6SAutB6SJK3WdT26qblWqDG+nyqdPj5ctYf/rS+k0ukDCeJ1L+O2LuKGFx+MOl1ia1EVry3cyZfr9iL5fbjlw4+3mxR0YMqQTKaN6iZWpIRWQdN03lzyO68s2InT68dx1I2qIksYFYluHe3cdWYvTu3VtBMAbeH7Sdw7NY28/dWc//zikBLqmsdBwQtX0eGc2xtcubEYZG45LYdbTou8YS/At5uKuP2DdVFlH9S+plHmi1tG0yM1tt7zvH6NQQ9/T3UdfQdrNv1M9arP8ZXtCQraLJ0ClTWtRoW504cfc6GNEH4vPJkdaPZ7LAxm+PMWsEe/T601qu+7SdzxtnGyLPHQBX25cnhX3lz8O5+u3YtBltB0HU3X8fg1TAaZC0/O4NrR3eiRGovHr3K7twcepeF+I0fyaTp7K1zc/9lG/j3l5CZ6R4Jw/AoqnDzxzVa+27Qf6ej0VWBbylBee+QHzuufzl1n9iQlLrIO7r3S4nh68sncvOtnvvcnUDZgOFUuH4k2Ez3TYjmrb5ooNiG0Gm6fyo1zVrN8Z3mdN4mqFlhp3bKvmhnvrubW07pz86k5zTxSQYClO8vCFuqOptqd26/xw5biqIOpM3PTuHRIJh+s3BNVQGU1yvzjgr4NBlIQ2DN2y2k5PPN9XtjXqG+LhkmGfhlxjRdIAbqkUC2dRoz6ObISSSv7IxhtMOLmNhNINUQEU+1ETkoMj07qx/3n9WbF7+UccPpw+VQe/HwTK+49nTjr4Rn0bzYW4auj10NDOckev8bXG/bxj/NzG2+pWRAa0YaCSqa+sYwat7/O9FU3Cvg0Pl1bwE9bivlg+vCIfgwPkZYu5sq/3YNtiOjFI7ROmqZz45zVLN1RFlFzdghMOvznp+1YjDLXjhYVXIXmVen0hfT0g+ir3VW6opsoPuTv5/UBCT5YvhtXBNeM1ajw9/P6MGVIZsSvccOYLH4rqmbehqKIgzajBB0d5TzXr+7ApajSTWmNB7+mk2A10inRiqGe3ltqZSV777gT3QP2KbfCutfAF9qgPvyAbND7fDj1vsjObwNEMNXO2EyGoH0az/ywjWqPPyiYenH+jpBUD4gsJxkCebsfrtrD9WPFj63QuuwoqeGy15ZR44msVKuqQYXTyyUvL+WrmaPplNhwap6/ogLvrl1YB9TTnVEQWtjsZbtYvrM84kDqEJdP5clvf2N0TjI90yKfYBCE46UoErJESFW9aKvdKQ20fKmLLEs8eH4uPed/wTu2HuzEhk/VOHLu2WwIBCgjsjow8/TuYYto1EeSJJ66eADxViPvr9iN16/VOekHgfTx7JQYXu6fSNUdfyH26aewjww09vX4VeZtKOKlBTvIL3VgVGQkKbDabDLIXDOyK5cN60xKbHDmhfu33yi4dSaxp55Kyp13IBkMkNIVvr0XkMBfR1BlMAcaJw67ESY80GCfqrZEBFPtXHZyDDuKa8hIsAKws6SGXWWOkPOi6cDt8njg6PoAACAASURBVGm8vTRfBFNCq6LrOtNmrcThja7nhQ7UuP3cOHs1X84cU+dzu3wqFoOCc9kybIMHI5nqaWQoCC1I13VeXrCjzpnvhjIQfH6N1xft5MlLxISB0Hw62E2YjUpIAYpoqt0BdIyw+IRP1ZAlKSj4cq5dy6BV3zJ53m3srFF5d/ludhTX4PCoxFoN9O8Uz+VDu5AWH1lqeDiHgraLBnbitYU7+WZTEQZZwqtq6HqgiJGmQ256HDeOy+a0XikYFBnnc89SMPM20v/1BCuTe3Lr+2vRdb12cvzIiROnV+XF+Tt4cf4OrhzehXvP6Y0sS1R+9RX7H36E1HvvJf788w4Pasi10HcSrH0Xlj4P7mo4FLjqGsgGGDYdBk+D2LRjfu8nKhFMtXPZyTHsKKlhbI9kAPYecGFU5JA9JNHkJAOUVIfv5C0ILWVlfgUl1Z6wjREbunlUdZ0dJTVsLqyiT3qgqlm128ena/by2qKdFBxwIUuBvYjxuo+Lu0/g+io3qRHutRKE5rRsZ3mdqU6RZCCoOvxvfSEPXpDbYIUyQWgsE3qn8sDnm0KOR1Oi3G5SmDw4fK8nXdf5ZUcZryzYwbKd5fg0DXSwGBXOzE3lutFdsT/+OMm3345ss5FjgwfPzw37XI2hb0Y8z152CpVOHz/9tp+yGi+qphNvNTK0WxJZR5VAtw0ZQqcXX2DWQy/xXJ8LcTew6HwouHp3+W72lDv4R+HPOH74ns5vvYmlV6/QB1gTYeQtMPwmKN0GzrJAIGVNgOTegXLq7VT7fecCAFnJdnaUHC576QqT3gfR5ySHazwnCC3plYU7wn6+I01f9ao6byzeyRMX9eexeVt5d/kuZEmqnSVVD0ZpBzDyToWJd/71M+N6JPPUJYGUDUFoLWYv2xX2WogmA0GWJL7ZWMTFg+puQioIjaljjJlxPZL5fsv+kEmxSEqUQyBgOqffSSHP/dPW/dz7yUaq3L6QlS+XT+V/v+7jm/V7SUs/n/8MGUdztqGNtxn54ymRXWfrYjN5Lnci7jB9uOri8qks2FjIk24LT370X5SEBt6dLENKmGCrHRPBVDuXnRzDd5v21/53jCX8RyLanGSLQVQsE1oPh8fPom2lIZWgorl5VDWdL34tZN8BN2v3VISs3h7Jq+qAzoLfSjj3uUV8MmNkxBUBBaGp7S53HH9VNJ9KUWWEG9IFoZFMH5fForzSqKvdARjROWvnL1Q+sxHjjBkoMYGVndlL83nk6y31fqeruo6qS+SbE5n86nJevXIwo7t3PO7305h0Xedvn26oM5CqLwPDLRv4Kq47f9ZMRF4uQzik7XfSEuqVnRLDztLDK1M5KTFhq+UcmZMc0fMm2xttjIJwvMpqvBiU0M2w0aavqprO6t0VuOr50T2SV9XYV+ni0teW4Yiw6IUgNDVPHZ/faDIQNJ2Q5qKC0NQGdUniwpPTsUbZYkKRJdKSbPz9iZtQKw6w8+xzOPDpZ3z5694GA6mjOb0qN8xexca9ldEOv0mt2X2gzi0WVSs+pfzH14gfPplOt8whY8ZbxA48B1fe8tpzNF3n7V/ym2m0bYtYmWrnToqzUOXyU+32EWsxkhJrYURWBxZsKwmauYwqJ9msMH1cdvO/GUGog8unEq6AU7Tpq5pO2Opn9c34qRrsrXDx8oId/PWMnsf7VgThuMU2QgaCQZZIEO0vhBbwyB/7UeX28fPWkojKhxsVieQYMx9OH0FivJXERx/BtX49+Y8+zh2dCWqufkhD+2idXpWZ76/lx7+OQ2olVeteXbgz7P8fkWZg+FSd91fs5o4ze4p+iFESK1PtnCxLZCXb2VlyuILfDeOysJpCL6S4oZNIPO1aKpfOpeD5qRS8dDXVa77E2j14Vl9C4szc9lfNRWi9Yi0Gwm3jO/Lm8VhFMuPn8Wu8s3QXfrGXUGgFhnRNwhhmpTaaDASzQSY3Pb4phicI9VJkiRcuH8iM8dnYTQr2MPcrEPiMmg0y43umMO+2sZwUb639m7V/f9b/5TFkQ+jEQiTf6QD7Kt38WtB6VqcWby8NW2ApmgwMSZLYvK+qCUbXtomVqXaswl3BJ3mfUBL/PtfNr0GSNWwGG6eknEJa8sns3teRoyfhG8pJthoVZozPxmQQcbrQeqTEmgNpfkcVMIu2pO7Rotlz5dc0fthSzFl9xUSD0LKuHNGFWb/kw1E7p6LLQDAwMrvuJqGC0JQkSWLmhO7cMDaLrzfs4+UFO9hZ4kDVdSQg0WZi6vDOXDGsS9j9qrqu89LCnTi14EmFaL7TPX6V1xbu5IWpA5vkPUZD13VcdbT9iCYDQ5KOvalxeyaCqXaowl3BI8se4ec9PyNLMm7coAEauPwuftr9E6bExVitMbj3n4+vOrLUJKtR4bTeKdw0XqT4Ca2LQZG5YlgX3lj8e1ClyWhuHsOJZsbP4VH5eE2BCKaEFtcp0cbAzoks3VkW8rdIqqJZjTLXj8lCPsbmp4LQWCxGhUkDOzFpYKDanU/VMMhSg6l3BRUuig64Q45H852u6fDd5qJjG3gTkJA4eoIEoi8gJreStMUTiQim2pmC6gKu+uYqyl3l+PXwsxgaGm7VDUY3low5mErPxlE6Muy5ABKBL7SJp2Tw8MS+rSZ/WBCOdOWILryx5PeQ45GW1A0n2j1Xov+a0FrccWYPpr6+POzG+4YyECxGhcmDRc0vofUxKpFlxZQ7vBgNMu6j0m+i3kerBSpbtvQeI0mSsJsVqtyh93XRZGBouk6STTScj5YIptqRCncFV31zFaXOUjQi3Lsh+TClfEucOY4DxQPwaxq+g2U3LUYZTYfROR2ZPjaLod2SRCAltFrpCVbOyk3ju81FITeQDd08KlKgUWnI8Shn/Pya2DMltA6DuiTx0AV9efCLjVFVMrOZFB6b1I+dpTXEWoykJ1iwmcSthHBi8Wvhy4dH+50uSXU/V3M7q28aH68pCNkfHE0GhkmRaxvTC5ET34DtyMPLHqbcVR55IHWQX/fgS/wvr5wzma17NR6ft5WbxmeTGm/lzNxUUmJF/xzhxPDkJf3Z9bKDrUXVYavyhWMxyqTEmik84A750Yx2z5WY8RNakylDMrGZZO78aD2qx4tPrvuWwKSr6AYDXr/GHf9djySBpumoms55A9K5bkw3eqWJmzDhxJBgM6KGCYKi/U7XdL3OAhjN7boxWXzxayFqmEm7SDIwLEaZaaO7oYj03aiJYKqdqHBXML9gftjUvopFFZR+W4q32ItiUYgbFEfqxako9uAviK2O77lu9PU8Nm8rfz2jp1iFEk44ZoPCBzeMYPqc1azKLw/pdH8kCbCaFMb1SOa6MVlc8fpy/EdV/Ytmxs9mVDgjN7Wp3pogHJPzB2TQPX8jb321ma9OGohOYN+JX9UwKTKSJOFTNTS/jv/g8mzNUT3TPl27ly/XF9K/UwKvXTmYeFEyXWjluiTZMCkyTo79Ox2gf6eEVnMv1CM1lpzkGDYWhq/G11AGhqbDZUM7N9Xw2jQRTLUTn+R9cnBzYrDSeaWUzCuh03WdiOkTg6/CR+HsQvKfyqfbfd2QD1bl86ge5myew9SeV2M2yK3my0MQomU1Kcy6egiLt5fyyoIdrNxVgXLwhhECPUn8Hi8js5KY/ofejMjqgCRJpCdY2HFEC4FDIt1zpaHzx1M6Nct7FIRI6ZqG6fUXuP+vf+H/Ro/l59+KKTzgwulVcftU3lj8O6qmo9ezaqUeXKFau7uCc55bxOe3jKJjjLkZ34UgRMegyFw9sisvLdgRkqUQ6Xe63RyoXtyaPHnJAC566Zd6JwrDsRoV/npGD3HdHiMRTLUTH237CI8avPlddakUf1ZMxrUZxPaPBcCUbCLzpky23bmNyl8qSRybWHu+V/Oyev9azIbWsaQtCMdKliXG9khmbI9k9h5wsWhbCQdcvtqSuj0/fp1Me3eSssfUPubGcdk8+MWmsD9SDe65kmHiyRnYzeIrV2hdqr/9FslqJWZcoPnooR6BpTUezvz3QlxeNUx9sPB8qs7+KjdTX1vO57eMavFN+YJQn8uHd+alBTvC/q2h73QIFLuY0CulKYZ2zHqfFMcbVw3h2rdXRhxQWY0KV43swnVjspp4dG2X+GVvJw54DoQcc+Y50XwacYOC89wVi0Js/1hqNtUEBVMA+2vKsBjFvg+h7chIsHLpUakN1VVjKX9nNklXXlF77PwB6by2aCc7SxxRbziOMRuYOaF7o4xXEBqLrqqU/OcFUu+5OyTb4MWft1Pl8oUNpByb51O18jN8ZQXIJivGlCziR07G0ikXv6azu9zJF+sKmTxEVPwTWq+UWAtXjezK7KW7cPmiW8mxGGX+cX4uhgirBzanEdkd+OjGkdzy3hp2lzvr/L2ymxRkWeLes3tx2bAuzTzKtkUEU+2EXwvdK6XWqBhiDEhKaMqeId6Aa5cr6Jiu6zj9bswGa8j5gtCW2EeNovDue/BXVGBIDEwoWIwK710/nAueX0xJjae2qmV9JMBmVphz7XDSE8R1IzQ/ze3HsbYY764qNKcfySRj7GDFNjgV57KfUGJjsY8O3mjv9qnMXbkHX5ibsKoVn1K5/CM6nHEzlm4DkRQDrt9X48pbjqVTLgAun8qL87dzyeBOIiVcaNXuOasXew+4+GlLccQBldUoc/OpOUw8JaOJR3fs+qTHMe+2MQx99EcGZMbzy/ZATzlZAp+m0+ekOGaMz+aMPmmYDK0vIDzRiGCqnbAarIHeUUdQYhT8NX50VQ8JqPyVfgwxwR8PWZIxSXYsRnHhCW2bbLFgHzWKmp9+IuGii2qPd4wx8/VtY5g2ayVbCqvweP1ocvjrwW5SiLcaeefaoeSkxDbX0AUBAF+Jk+r5e3CtLwUJdO/hfSFuGaqX7EWrrCZ+8s0hAc/XG/aFfU7N4+DA4nfpcM7t2Hoe7j1oyxmGLWdY0Ln7qzz8WlDJyZkJjfiuBKFxybLE85eewmPztvL20nwAvHVUepWlQEPbB8/PDclmaI2+27yfPifF8c60YYHJcK+KT9WItRhFxb5GJu6K24mBqQNDClDYcmxIBomq1cGVX1S3SvX6aux97EHHvaqXTrbuIg9eaBdi//AHqr/7PuR4gs3EJzeN4tWUIibIZZgNMjFmA7EWA7FmA2aDzIisDrwwdSCL7z5NBFJCs3P9Vk7xc2txri1G92lBgRQAGuDXke2dcK6VqPg0D/2IVah5G4twhNlv4dm7Fd3vxdZjRINj8PhVFvxWfLxvRRCanCxL3HdubxbeeSrXj+lGnMWA3awEvtMtBmwmhbQ4C9NGdSPRbuTCk1vvitSRZi/bxRXDA+l7gaa+BhJsJhFINQGxMtVOXJ17Nb8U/oLLfzh1T7EppExMoXBOIbJFDqrmZ0wykjDy8IyihMSojFGY5QTMhqKWeAuC0Kxixo2l6MEHUWtqUGJigv6m+/2kfTqb5/7zPL7sHuwuc1Lj8Qd+dOMtovea0GLceRWUz9mCHmEjXt2n4VhTjK7qJF7UHUmSKHd4w56ruqqQbXERNTPVdCip8TR4niC0FmnxFu48sxe3n96DvP01HHB5USSJJLuJnJQYJEliX6Wbt375nZvG57T0cOv1W1E1+aUO0Y6jmYhgqp0YkDyADpYOFNQUBB1PPicZxa5QNLcIb7EX2SoTNzCOzOmZyEek81kMFq7KvQpHpSpWpoR2wW+1sXjE+cz89wJ2exU8fhWjIpMaZ+Hyjm5Gn5SJNTcXK9A3I76lhysIqNVeymZvjjiQquXTcP1agrlbHPZBadQ1ca1Y49CcVeiaGlFAJWbAhRORUZHpkx6+AfVfzujBJS8vZerQLq26n9q7y3dx6dDOGFthgYy2SPy/3E5IksQ9Q+/BooTOmCeNS6L7I93JfS2X3s/1JuPqjKCGvSbZRO+k3gxMGYjHr2EWmxWFNkzTdJ77MY+B//cDT8cN4jdHYEO9poPHr7G73MkzW71c1PUS7vl4Pe4oq0AJQlNxLN8XlK53pBUF65k4ewZ9/n02fZ89lz/OuYl1+7bU/l33aVT9uAdd10mODd9rxpzRC8lgxLltaYNjUWRJrNAKbU52cgxn5qbWWVK9Najx+Pl8XSGXDRXVNJuLuCtuR8ZljuPWU24NG1DVxSSbSI9J54UJLyBJEm6filmsTAltlF/VuHHOal6avwOHx49TCz+z7pYNeDSJT9fuZeILS6h0+pp5pIIQTFd1apYUgj80mKr2OLjmo3u4etBFbLjtS1be9Am3j7oasxLc5kKr8eLdXX2wJ1ro97xstpMweirl37+Mc9tSNJ8bXfXj2rGKip/fDDrXKEucKVKMhDZo5oTufLByN0WV7oZPbgGfrd3L8KwkTooXFWSbi0jza2f+lPsn4sxxPLzsYYCQRr6HSEhYDBZ6J/XmhQkvEGMK7BkRK1NCW6XrOnf891cW5pXgjjBNyuPX2FFSw5VvLue/N44QDa2FFuPOq6hzVWpn+R4AJvY5HQCrrDCu29CQ83SfRs2SvZw2pefB9KDQVde4oZOQ7YlULp1L6ZdPIZmsmFNziBsxJei8HqmxoviK0CadFG9lyuBMnv0xj8cm9QN3JexaCq4KkGSwJUGXkWCyN/xkx6iwppD3trzHVzu/ospXhaZr2Aw2BqcOZv2mU/jnmec22WsLoUQw1Q5NzJnI6IzRfLTtI+ZsmVPbg0rXdWRJxuF10zN+MPeMvJGBKQODyua6fce3Z8qjeliwZwFFjiLcqpsYYwy9O/Tm5OSTRT8SoUUtzCvlu837wwZS9TUp9ak62/ZXM2tJPtPHZbfAyAUB/KUudDX8JEBWUiayJPPnrx7hgl4TOCUjlwRLmEBHB99+JwZF5pqRXXlx/g48YcpEx+SeSkzuqXWOxWZSxLUgtGkzxmcz/clZVHueIXbHl6AYQT94rUgyaH7ofymMuBk6Nl7D9r01e3lgyQOsK1mHruv4tMNZEVXeKn7e8zN6zGIe3TCX++z3MTpjdD3PJjQWEUy1Ux2tHblxwI1c1+861hWvo8xdhlf1EmeKY9ueeH7N1xiUOhAIBFneXVV4d1fTaVs5sV4Vx8oirLkdkCPcgLmneg/vbnmXT/M+RULCq3lRNRWjYkSWZJIsSVyTew3nZ5+PzWhryrcuCGG9PH8HzjDloCNpUur2aby+6HeuH5OFLDbdCy1A96pQRyPpWLOdT6b+hxeXv8dd3zxJiaOcU7OH8a+z7iLZnhT6PMB1Y7L4bF0hu8udqHWseIVjNsgM7JzIWX3Tjv3NCEJrpvpI+GYmc6RPkX/zAhr4w6T8rZ0Nv74Pg6+BMx6FOnoSRmpL2Rau/e5aHD4Hmh5+4kRHB9lLQU0Bf/75z9w55E4m95x8XK8rNEwEU+2cQTYwOG1w0LE+CR6e+24+bocX/4ZSqhcUoDl86KpOjqqjAQeKdlDx+XasfToSOzYDU6e60zk+yfuER5c/iqZrQbMocDjNcG/NXp5e/TQv/voib535FlkJWY3+XgWhLgUVTtbsrgg5Hk2TUqfXz+LtpYztkdzk4xWEo0lmBRQp7J4pgO4du/Lvc+8FYHvZLmZ++TD/+PF5XrjgwdDnAexmA3NvGM4fX1xCSbUHbx2B2pEsRpleabG8+qdBopKf0DZpKrw3GXYvxag1sGdK8wf+Wf021JTARa/DMWbgFFQXcO2311Ltq474MW7VzZMrnyTeHM+ZXc88ptcVIiM2vwghkmPNjEyKYd/Tq6n86nfUCk+g6ePBH1MZAv/t13FtKKHklfVUfpeProf+2L635T0eW/4YHtUTEkgdzeV3UeGu4PKvL2fngZ1N8dYEIaxvNhYR7lYxmialDq/Kf1fvafzBCUIEjB2tSBGWQc7p0IXJfc/it5KjvmclMKYe3ueREmfh65ljGZXTEZOuYpTCB1QWo4zZIHPhyRl8OH0kNpOYpxXaqHl3we5l4HM1fO4hPif89jUseOKYX/a+xffh8DtCjlcsqiDv/jw23bCJrTO3Uvh2IarjcIaFW3Vz/+L7cfhCHys0HhFMCSH85W7uLtFRnP6G+5XoBzctL9rLgf8F/zAv37ecf6/+N2418oo3OjpOn5Nrvr1GXPxCsymqdOMNszckmialh55HEFqCuXsikiH8rPf2sl28suID9lUVA1BYtZ/Pt/zIwPTcoPMkg0zs6IygY/E2Iy8MMPDWhllMG9WNGLMBWQqk88kSJMeYuX1Cd5b9bQJPXNQfkyhQJLRV1fthzexAcHSErs9Uk/JkNQ7v4cmG19d4GT/riHsYnxMWPwOemqhfdnfVbjaVbQpJ7SudV0rRf4tIm5xGnxf7kPX3LLxlXvKfykc74vdMQuKL7V9E/bpC5MT0kRBEVzVKXl2P0a8TzWK07tNwrizClBGDfVCgHO6za54NG0hVLKqg9NtSvMVeFItC3KA4Ui9Ore1tpaPj8rv4cseXTOk1JeTxgtDYvHVs3I+2SakvglQoQWgKkiwRMyqDqp93gy/4c2g32VhXuIXXVn5IlaeGOHMMp2eP4L5Tbwo6z2838NOBGlyllcSYjXRPiaFrRzulr7xCnz9NYdR5ufztvFzcPhWnV8VuVkQFS6H9WPVmnWl6qg7PLvdy75jwPdqAQGGK9R/CkGlRvey7W95F1YL386ouleLPism4NoPY/oFtFqZkE5k3ZbLtzm1U/lJJ4thEAFyqi7c2vcWlvS4Vhb6aiAimhCCuTWVoTh/hsjlWFKzn0Z9fYltpPrIs071DFx6ccCsnn9QbONj08bt8bANTyK/KZ1vFtpDnKJ1XSsm8Ejpd14mYPjH4KnwUzi4k/6l8ut3XDfngrKbLH7j4J/ecLC5+ocl1sJuRICTV78gmpfZeDVdFSoywIIsgNAX70DSqFxSgH1XS/KTYZF6a+FC9j3Wj8/+qq1j40Xo0XUeWJHyqRq8EIxP3q0y98MLacy1G5biqugrCCUdTYcUr4QtNAHeONPGvJR5uGmIiwVLHPYvPAUueiTqY+mrnV/h1f9AxZ54TzacRNygu6LhiUYjtH0vNppraYArggOcA2w9sp3ti41UWFA4T6/FCkOoFBYH9UAf50clDZZGnkis/upuzB/2Rdbf9r+6mjy4/np2V9c6kpF+RTmz/WCSDVDuT4i31UvlLZdD55e5y1havbbo3KwgHDctKwmo6vialVqPC+J6i+ITQcpQYEx2vykUyRvfT7kLnS7x8rXqp8fhxelVqPH48fo1fSz082ftCTnvuF/aUOxt+MkFoi6oK6wykAAanK4zvauCpX8L37jz8PAXgjfw60nU9bNEJtUbFEGNAUkIDN0O8AX9NcPClSArl7vKIX1eIjgimhFq+Yie+4sBFXorG67g5n2puxsF95dtxAx/1Gc4k2cV/jTr9uw2md0pwLxHdq1GzsIB1xeuOaSblSJqmsaV8S+O/UUE4yrBuSSTUsaoUN3QSiaddS+XSuRQ8P5WCl66mes2XWLsHF6XQdZ1JAzs1x3AFoU7mrHg6XNUHySQ3+AuvoeNC5zO8PEvdN4FOXabwgIvznl9MfqnYyyq0Q+5KkOtP5vrnqWaeX+GlxFHPXnPZFHiuCOnoYYt7KTEK/ho/epjUcn+lH0NM6FgbKgImHDuR5ifU8hU5QILXcfMeXoCD/wtaUgZIMnu++n/Ye43lzYxezLLEcD1mLiM4R9i7uwyHpSjk+RuaSXHtCq6O49W81Hij36wpCNGSJIkbxmTxxDdbcYUputJQk1JFgvMHpBNrEWl+Qsuz5CSSevsgqhcW4Fy9H83tQlIOf0/rsoRX09iAyhw8rCK0v9rRNB2q3T6mvLqUH/4yTnzWhfbFYIYwQc2R+qYonNfDwOOLvfROrmMmQ9cCzxUhWZIxKabaNjKH2HJsSAaJqtVVxA+Nrz2uulWq11eTenFqyHPFmeJCjgmNQ6xMCbU0l5/HfA4+wIuXw4EUgGy2kTb1X4BE2TfPs+O5yyn4+J+84ijiPwQvfeseH2Z/6AxItDMpBtmAxWBphHcmCA27ZHAmSXYzyjHs0bOaDMycIHLRhdbDkGQhcWIOHa5Mw7fra2xD07D07YBtYAqLUgxcQQ234wwJpByb57Pv7dvZ/f8upuA/V7L/wwdxF2wCAgFVlcvHf1cVtMRbEoSWY+8IqrfB0x4ab+G1NV72VtUReOkamKMLagYkDwg5ptgUUiamUDinkOr11eh+HW+Jlz0v7sGYZCRhZELQ+X7NT05CTlSvK0ROBFNCrVfz9vOj7qOurGBjx0w6nvtnOt38NunXvoBaU86+H1/jM7x8ckSKiGSLISN9cMjjj5xJOdKhmRR7H3vQcZNiItUeOrsiCE3BbjYwd/pw4m3GqBqOWo0Kb08bQmaSrQlHJwjHxrV6BZYsM0mTutPxij4o52XxUGkF+8J0Vqta8SnlP75G/PDJdLplDhkz3iJ24Dm48pYffj6fxqsLd4ZNPRKENsuaCJ2GNHhaTpLMlFwjz60IE3hJMvQ+D5ToksKu6XsNNkPo70vyOcmkXpRK0dwiNs/YzI7/24ExyUi3u7ohH7Fv0iAZuCD7AmxG8RvVVESanwBApcvHK1sKw2bNOzbPp2rlZ/jKCpBNVowpWcSPnIy97wRq1n2DG3gJD+diwoyEEm/m0p6XsqpoFU7/4Y2WR86kyBY5qJpfuJkUTdcY32l8k75vQThSp0QbX88cw+WvL2N/pRuHt+70J7tJwWxUeGfaUPpmxNd5niC0JMfyFcSMG1f73x+u2hO2urPmcXBg8bt0OOd2bD1H1h635QzDljMs6Nwqt4+lO8oYmdOxycYtCK3OqNth36/QwPaDB8aZmb0+zP4kgwVG3Br1y45MH4nVYA26nzokaVwSSeOS6n28Iitc0eeKqF9XiJwIpgQAPlq1B1mWAs0SjlC14lMql39E/Igp6D4Xtl7j8JXtwrHxZ3yluzCn96w99yd8nGOyEjMqnVEZyVgMlpCLP/mcHYkbmQAAIABJREFUZBS7QtHcIrzFXmSrTNzAODKnZ4bMpEzMmSjS/IRmlxZv4Yc/j2PJjlJeXrCDVfkVmAwyuh5oMeLw+Dkp3so9Z/fizNw00aRUaLV0TcO5YgWpd91Ze2zhthLcYfYFevZuRfd7sfUYEfK3o7m8Kqt3VYhgSmhfciaAyR4STOXfHhv035nxMu77j07lkyC+E2QMjPplZUnm7qF388CSB8L27qyPRbEwPnM83eK7Rf26QuREMCWg6zqvLfo9ZOP9kTOV5oxeVPz4OtVrvkLzOJDNdqzZQ0k8NdAvwQXMwcs5WLH1S0aSZK7OvZoX170YcvFHPJPSW8ykCC1DliXGdE9mTPdkiird7Cytodrtx24ysCq/jD0Vbs4fkN7SwxSEennytiPHxGBMP/xZPeAKX9FLdVUh2+Iiak6tA2WOhvePCEKbIitw6Xvw9nngczV8/pFMdpgyp86mvw05u9vZ7KnawyvrX8WrNVB+/SCLYqF3h948OvrRY3pNIXIimBLYX+Whwhn6w3jkTKUkKyRPvKfe5ylAQx6cUtvj5Mo+V7Jk7xLWFq/Fq0X+w2tRLNw95G46x3WO7o0IQhNIi7eQFn94hTQj0cqlry5F13XRUFpo1ZzLl2MfNjTomLmOlVTFGofmrELX1IgCKqto2iu0R50GBwKqD6aCr+F+URoSfsWG6crPILlng+fXZ0qPa3h9YTGG+E9RZCmkwt8hiqRglI2MyxzHY6Mfw6iIyptNTeSnCFS6fBiV0I9CNDOVAAZ08LwBamDms7CmEL/uj6q3gUWxMHPgTC7ueXHEjxGE5tS1gw2DLLO9WJTtF1o3x4rl2IYNDzp2UkL41GlzRi8kgxHntqUNPq/ZIJMSF3l5Z0FoU7JPg2u/g8zhgX1QcphgRTGBYsabOYZLeYw1+vFV0tN1nXs+Xs9ZnSfy4+Tvmd5/OonmROwGOzHGmNp/zIqZC7Iv4L1z3+OpcU+JQKqZiJUpAUWWwjeFi3KmEoOMwbEXZv+RDaf/jRsW3oHT50QPUzUqnE72Tjww8gFGpDecsy8ILUWSJMZ078iivFK6p8Y2/ABBaAG6quJcuYq0Bx4IOn7pkM78tKU4pLiKbLaTMHoq5d+/jCQrWLqdgiQbcOevw717fW1KNwTS/M7pd1JzvA1BaJ3S+sG130L5Tlj+Mmz5CjxVgTQ+SzzkXgRDr8MS34kbNhZx2wdr+WrmGOKOsT/b7GW72F3u5N9TTsZiVLi+//VM6zuNTWWbKHeX49f8xJniyO2Yi91ob/gJhUYlgimBDnYTXjV0Q/KRM5X2XqMbfB6/BvGXv0H+D3dy/Q834ogyA6rUXYqm19M5XBBaidHdO/LJmr1MGy029Qqtk3vrVgwdOmBMSQk6PjK7AzEWQ9hKlXFDJyHbE6lcOpfSL59CMlkxp+YQN2JK7TnSwedIjRPFgQSBpCw4+1+Bf+pwVt80Fm8v4b5PN/LcpSdHnR6+cW8lz/yQx8czRmI5Ir1WkRX6J/c/5qELjUcEUwIJNiNJNhP7q4Pzb6OZqQQYld0Bk8nInepenGG+KyoWVVD6bSneYi+KRSFuUBypF6ei2ANfDm7VzV8X/JWFUxZiUkxN9n4F4XiNyu7IPR9vwOvXRDU/ocU5vX4+W7uXRXmlHHD6MBlkOhbt4uwhp5J11N4+SZK4fkwWT32zFXeYBuoxuacSk3tqna9lMSrcMDarSd6HILRV95/bhwv+s5iPVhdwyeDM2uOqplPu8FLl9mE1KiTZTUEBU7Xbxy3vreEfF+TSraNYcWqtRDDVzmmazv99tRmDImMzKTiPmq2MZKYSAj13po/LZlvFNvIr80MS+0rnlVIyr4RO13UK6i+V/1Q+3e7rhnzwhlTXdb7b9R3nZZ3XlG9bEI5Lot1EVrKdNbsrGJ7VoaWHI7RTe8qdvDR/B5+u3YskEfT9LekWvlZ6k/rUfG4an8NFgzrVNqO+sGQ9n5YXkZeQiTeK3rtWo8LFgzoxMluURBeEaFiMCs9fNpDLXlvGwC6JxJgNzFm2i7eX5uPxaRhkCU0Hv6ZxRp80rh+bxYBO8dz76UZGZHfgAlE9tlUTwVQ75lc17v54A/llDj6/ZRQTnl4AhKZ+NDRTCRBrNTIyuwMPLHk2pOCE6lIp/qyYjGsziO0f2GNiSjaReVMm2+7cRuUvlSSOTQTA6Xfy5oY3RTAltHqjczqyOK9UBFNCi1iVX87Vb63E7VPxa6ERkS7JuDTIL3Py4Beb+HLDPl6+7GSqn3+W6h9+4J1/P8e180vYuq8at7/h9GqrUeGsvmk8dEFuU7wdQWjzeqbFMnNCDpNeXILLq4EE3oPX3pF5QfM27uOnrcXEWQzYzAbm3TamZQYsREzkp7RTbp/KTe+uoaTGw+xrh9Ixxszzl52CxRj9R8JilHnh8lOQJIlv8r9B1YMDMmeeE82nETcouImdYlGI7R9Lzabgqmi7q3dTWFMY/ZsShGY0untHFm0vbelhCO3Q+oIDXPnGCmo8/rCB1NFcPpXlO8uYev971GzaTNe5H9AhtycfTh/J1OFdsJkUbKbwRYbspkDq0d1n9eT/TR4QaO4uCELU3D6Vz9YWUu3241W12kDqaJoeuGb3V3sornaTX+Zo5pEK0RLBVDtU4/EzbdZKjIrM638ajM0UWKAc2yOZxyf1jyqgCgRSAxnUJQmv6g3b90CtUTHEGJCU0B9hQ7wBf40/6JhRNlLmKovyXQlC8xrUJZEdxTVUOiMv/S8Ix8vlVfnTGytw+UKzCOrj8WtsMSbxyZQ7MSQGMgFMBpm/n9eHNX//Aw9dkEuvtFjirUYsRplEm5GhXZN49tJTWHnf6Vw9qpvoqyYIx0jTdKbPXs2WfVVEMP9Ry+FRuezVZRRVuptucMJxE2l+bYHXCcWbwXUAFAPYkyGlT9hO2xUOL1fPWknvtFge+WO/2hz6QyaekkFKnJm7P15PWY0Xl0/l6KrpsgRmg0J6goWnLhnAKZ0DP8x+zY8iKfj14OBIiVHw1/jRVT0koPJX+jHEhH4Mo+lNJQgtwWxQGNDVwNPL3qJDvAu36ibBnEC/jv0Ynj4cWRJzVULj+9/6wrDVVwEcm+dTtfIzfGUFyCYrxpQs4kdOxtIpkJrnRmHW8j3MPKMnZsPhlSiLUeGSwZlBG+MFQWg8i7aXsjK/HE+Y1aiGrtsqt4+nv/uNJy8Z0NzDFiIkgqkTWel2WP4SrHsPZIVA0VpAU8GaACNnwsmXBXoeAPur3Fz5xnJO7ZnCPWf3qnOWcWR2RxbeeSqrd1Xw6sKdLMwrweMLfAFYjAqn907hujFZDMhMCHqc1WANW9rclmNDMkhUra4ifmh87XHVrVK9vprUi1ODztd0jThT3NFPIwitxq8lvzJr4yw2KQvYuEdC2+MFQEbGYrBgM9q4qs9VTOoxSXyWhUb18vwdIYWCAKpWfErl8o/ocMbNWLoNRFIMuH5fjStvee1NGQC6zjcbi7jw5IxmHLUgtG/Hc92qWmAS5YHz+xB7jH2qhKYlgqkTkeqDz2+BzZ8FAqdwqzg+B/z4EPzwD7jwP+xKP5sr31jBpUMzuWl8w524JUlicNckBndNCrykpiNBvfnykiTRM6knW8q3BB1XbAopE1MonFOIbJGDqvkZk4wkjEwIeZ4ucV0aHKMgNDdd13lu7XPM2TwHj+oJaUitoeH0O3H6nbyw7gXe2vQWb575JtkJ2S00YqEt2VxYxb4w6T6ax8GBxe/S4ZzbsfUcWXvcljMMW86woHMdXpXXF/0ugilBaCZ7yp2s2V0Rcjya61aWJD5eXcDVo0Rvw9ZI5KGcaFQfzP4jbPkc/O7wgdQhPif4XWif3cx/X7ifG8ZmRRRIhaPIUkQbj6f1m4bNYAs5nnxOMqkXpVI0t4jNMzaz4/92YEwy0u2ubshH7NEyykYu6XEJRkXMvgitzxMrnmDO5jm4VXdIIHU0t+qmwl3B1K+n8nvl7800QqEt213uCEnNBvDs3Yru92LrMSKi59lT4WzsoQmCUIeFeSXhdl1Edd06vSpf/LqvCUYnNAaxMnWi+fxm2LsKfK6IHyKrbm5X3sOQeAbQtCs+EzIn8JD0UNi/JY1LImlcUr2PlySJy3pd1hRDE4Tj8sWOL/g472PcauQbgXV0nD4n076dxrxJ87AYLE04QqGtc3hU9KM3sQKqqwrZFockh6/Id7RDaduCIDS9A05f2Mp90V63B5zexh6a0EjEytSJpGQbbP48JJDq+kw1KU9W4zii++Lra7yMn3W4nKZBdcPXdxBSTaKRGRUjtw+6HYsS/U2jRbFwbrdzSY8RzemE1kXXdZ5f+3zYQKpiUQV59+ex6YZNbJ25lcK3C1Edh3PjDwVU3+36rjmHLLRBMRZD2L2uijUOzVmFrkVW4e9YWmAIgtC4or1uhdZLfKOeSJa/BJo/7J9UHZ5d3sCshasCdi1pgoEFm9JzCpN7To5qFt6iWBiQPIAHRjzQhCMThGOzav8qqjxVIcdL55VS9N8i0ian0efFPmT9PQtvmZf8p/LRjpiJdPqdvLHhjeYcstAGZSfH4NdCZ7jNGb2QDEac25ZG9DxZHWMae2iCINQhwWbEbAi93Y72uk2wmRp7aEIjEcHUicLrgF8/qDOYunOkiad+8XDAXc/Kk9cJS55togEGu2PwHczoPwOTYsIk1/0FoEgKZsXM6V1O56U/vIRBFpmnQuvzzuZ3cPmDV4RVl0rxZ8WkX5FObP9YJIOEKdlE5k2ZeEu9VP5SGXR+YU0hW8u3NuewhTYmJyUmbCAkm+0kjJ5K+fcv49y2FM3nRlf9uHasouLnN4POtZsVrh+b1VxDFoR2b1yP5LC9paK5bq1GhYkni6yd1krcuZ4oijYeLH8e3uB0hfFdDTz1i4eHT6trRUhvlpUpCOx9mtZvGudmncvc3+bywdYP0NCQkNDRkZDwa37OyTqHK3tfSU7isRXGEITm8Fv5byEFJ5x5TjSfRtyg4NLnikUhtn8sNZtqSBybWHtclmS2H9hOr6RezTJmoW26cXw2f/t4PY6jyizHDZ2EbE+kculcSr98CslkxZyaQ9yIKUHnGWWZ03unNOeQBaFd65RoY1CXRH7ZURbyt0ivWx2dPw4UFThbKxFMnSjcB6jtI1WHf55qZtSbDm4bVs9SsM8V2DfVTJ3sU+2pzBw4kxknz2B9yXoOuA/g1/3EmeLon9wfu9HeLOMQhOPh9IdWP1NrVAwxhpBG1ACGeAOuXUetZOkqNd6aJhuj0D6clZvGw19uxhmmoXpM7qnE5J5a52OtRoUZ47MxKCIpRRCa0/Rx2azbcyBsr6mGrltFhvP7p4seU62YCKZOFBFUe+mbonBeDwOPL/bSO7mOH0tJbrZA6khG2cig1EHN/rqC0BjMijnkmBKj4K/xo6t6SEDlr/RjiAn+epUlGavB2qTjFNo+k0Hm/RuGc+F/luDw+Bso0H+Y1agwvmcyN4gUP0FodmNyOjK0WxJLd5ThCVPZrz7xFhN3nNmziUYmNAYxPXWisCeD3nDFl4fGW3htjZe9VXX8xJrjwh8XBKFOGTGh6RW2HBuSQaJqdXBhCtWtUr2+Gnuf4FVXCUlUqhQaRXZyDB/PGEmi3USYhdEQNpPC2X3TeP6yU8JWAxQEoWnJssTLVwyiT3pcxNU0ZQnirAbev2E4qXGirUZrJoKpE0VqPzA1XIEpJ0lmSq6R51aEqewnG6H/JU0wOEFo267ofUVIM2rFppAyMYXCOYVUr69G9+t4S7zseXEPxiQjCSMTgs63GqxidVZoND3TYnnjT4MxKjLxViN2U3D2glGRMBtkhnRN5PnLTuHpyQNEep8gtCCLUWHuDSM4r386JoMctsIfBIIoq1GhR2osX906hp5psc08UiFaIs3vRCHLMOIWmP9ogw17HxhnZvZ6X5jnUGDYjU00QEFou07tfCrK0tBU2+RzklHsCkVzi/AWe5GtMnED48icnol8xOyjWTHzp9w/IUviZlZoPK8v/p1bJ3Rn+tgsftpazMr8ckprPFiMBtITLFwwIJ0uHcS+VEFoLUwGmacuGcBdZ/XkvWW7mfVLPpUuH+aDvxeaBmf1TeP6MVn06xTfwqMVIiWCqRPJKVeg//xISBmK/NuDZy0y42Xc9x+VzifJkH4yJIl8eUGIllE2cmWfK3lzw5shjXuTxiWRNC6p3sfLksyknElNOUShnVm35wCrdpXz5CX9MSgyZ+SmcUZuWksPSxCECKTEWrj9Dz249bQc+v3jWz68cSTJseaDPaka3iMvtC4imDqB5FUb+Vi+nr9Ir2LSPdE92BQDF77YNAMThHbg+n7Xs3LfStaXrsejRn79WRQLz4x/hgRLQsMnC0IEdF3n8XlbuG1CD2wm8TMuRKbIUcTa4rVUeapQZIUkSxIj0keIwjgtqMajIssyfTPEKtSJTHwLnyC+3VTE3z7ZwN/Ono7Jkww/Pwb++tP9gMCKlCkG/vQ5dMhu+oEKQhtlkA28cPoLzPxpJuuK14WsUIXQJSwGM0+MfYKRGSObZ5BCu7BgWwnFVR4mD+7U0kMRWjlN11i2bxlvbXyLNfvXYJSN+HU/EhKKrKBqKhfmXMjU3lPpFt+tpYfb7hRWukhPEMUlTnQimGrlNE3nmR+28dHqAt66eggDMhOA2yCxK8y7CzzV4HWEPlA2BvZIpfWDP74iAilBaARWg5WXT3+Zub/N5a2Nb1HprcTtdwc19DUrZnR04vV+DI2/jNM6n9aCIxbaGk3TeXzeVu46q6coKCHUq8Zbw80/3szW8q21vfK82hHFqQ4WCP5428d8tv0zpvWdxowBM0TFx2ZUeMBFeoJYGTzRiWCqFat0+fjz3HXUuP18fstokmOP6HXT50LodT78Ph8WPwsFy8HnDvSQMsdBv4th2AzomNNi4xeEtkiRFS7vfTmX9bqMVftX8dG2jyisKcSjeogzxTEkbQgX97gY1WfnzGcWcsPwGrKSG67EKQiR+PzXvViMCmeK/VFCPZw+J1O/nkpBdUFwABWGX/fjV/3M2jSLSk8lfxv2t2YapVB4wMVJ8SKYOtGJYKqVyttfzQ2zVzO2e0fuP68PxnAzkLIM2acF/oFAGRhJapGmvILQ3kiSxJC0IQxJGxL+BCvMGJ/NQ//bzKxrhojZXuG4efwqT327jacnDxCfJ6Fet/18G3tr9jYYSB3J5XfxSd4nZCdkM7nn5CYcnXBIYaWbDJHmd8ITwVQTc3r9bC6s4oDTh6JIdLCbyE2PR5Hr/iE8vD+qF5cMzoz8xWSR8iEIrcnVI7sxd+UefthSzB/6pLb0cIQT3Jxlu+mZFsvwrA4tPRShFdtUtol1xevCFsqpWFRB6beleIu9KBaFuEFxpF6cimIPVJBzq26eW/sck7pPwiCLW8SmVnjAxdjuyS09DOE4iSuliewoqeHNxb/zyZq9GA4FTlIg391iVLh2dDcuHdqZJLup9jHh90cJgnCiMhlk/nFBLvd+uoEx3TtiMYqSt8KxqXL7eGn+duZcN6ylhyK0cu9seifsilTpvFL+f3v3HR9VlT5+/HPvnZnMTHpCCpAAgdB7V6qIBQuKXdcKiIqr6677dYu66rquu+v6W3fVFRERe8MV7AVQmtKV3gMpJJBC+vS59/7+GAgZ5k4yA6EEzvv1yuulM/dOTkLmzn3OeZ7nlH9VTtadWcT1isNX5aPkrRLyn80n55Ec5EObyPpUH0v3LRX1nifB/mq3qJk6A4iljBbmUzV+++F6Lv3PMj5YU4TLp1Ln8Qe+3H4cXpWDDi/PL9rFuX9bxDurCoDAB+W0N9eyck8ln9w3SgRSgnCGGN01jd5tE3ll6Z5TPRShFZu1dA9ju6XTIzOh+YOFs1att5aFhQvRdC3ocdWlUja/jHa3tCO+XzySScKSZiH73my8FV5qfqxpONbpd/La5tdO9tDPSsXVLtqLYKrVE8FUC/KrGne8tpovN+3H49fwa3rYY91+DY9f46nPt/LEp5uZ9OIPZCXbeGfa8OBGE4IgtHqPXt6T137Yy74q56keitAKldW6eWtlAQ9e1O1UD0U4zeVV52GRLSGPO3c50XwaCYODg3HFqhDfL576LfVBj2+v3H5CxymAqumU1bnJSBT3fK2dSPNrQX+ct4mfCqtw+bTmDz7E5dN4/ccCbhqWzZ+v7HMCRycIwqmSlWxn8ogc/vrFNmbcMhiHx88Xm/azt9xBjdtHks1M14w4LunTVqQCCiH+s2gX1w3OEjPYQrPqvHWGj6v1KqY4E5ISWq9tSjThKgjet9KjetB1XTQ6OYHK6zwk2S3EmMQ1v7UTwVQLKTzo5NP1JXj8oYGUY+tiatfMx3dwH7LFhjm9M4kjrsea1bvhmG+3lPLUJL3JxhSCILRed4/tzHnPfs8dc1azcs9BZEnC6VUbno+1KDwybzPXDclm6sgcOqTaT+FohVNB13RQdSTzkaSRPeX1fLlpP9/99rxTNzCh1bAooatSAEqcgr/ej67qIQGVv8aPKS74dtAkmUQgdYIViz2mzhgimGohr/+4F00PTeurXT2PmlUfkXrRL7HmDEJSTLj2rsO1a1VQMOX2qSzdWc64Huknc9iCIJwkX2/eT2W9j8U7yg2fdxwKrN5dWcCHa4p4/qaBogPgWcB/0EXdD8U4fypDd6sgARKYM2KJPy+Lf20sZNqYziTHGt8kC0Jj6fZ0/Jo/5HF7rh3JJFG7rpbEYYkNj6tulbqNdWRcG3ytSYxJPPolhONUVOnkjR/zmfdzMbVuH35NR5Yk7pizmrvGdObczqkigG2lRDDVAtw+lQ/WFOFTg4MpzeOgevk7pF76a+zdRzQ8bs8djj03uCOTw6syY0meCKYE4Qz0v3X7eGT+Jrxq8ynAPk3Hp6nc/95PPH/jQC4Sm7Oekfw1Hirf3463qB70wIoUAHrgy7ffQcVHu7jfp5LWIUOkXAkRyUnIISM2g4LagqDHFbtC+qR0St4uQbbKQd38zClmkkYcaXplkS1c0+2akz30M1ZRpZOHPtrAz4XVaLoedK+o6jpLdpSzem8liTYzT1zRW2zI3QqJBhQtYHdZveGHnKd4O7rfi73buRG9zs+FVS09NEEQTrHNxTU8Mn8T7ihqKQHcPo0H3l/PnvL65g8WWhVfmZOy//yEt6AW/NqRQOoosk8jFgnXwkKq5+1GN8h+EITGJEliap+p2E2hacJpl6aRcU0GBz44wNbpW8n7Sx7mFDM5v8tBbpRaigQ3dL/hJI76zLW1pJbLnl/G6r2VePxayKQ7BOZPnF6V/TVuHnj/Z15dJjq/tjZiZaoF1Lh8GE0Yqq5aZHsCkhxZcaFf0/GpGmZFxLiCcKZ44btdhrWU0Hw9pVdVeWXpHv5+Tb+TOWThBFJrvZTP3IjmDE3FCkf3aTh/LkOym0iakHMCRyecCSbkTOAfq/9h+FzK2BRSxqaEPdckmRieOZx0u8iSOV5FlU5unLWCWnfk73W3T+PZb3eQbLdwzeCsEzg6oSWJYOo41Dh9zF1XxGs/7KXO4M2i2BLQnLXomhpxQKWINA5BOGNU1HtYvKMcowWFSOopVQ3mry/mT5f3IjZGXK7PBNWf5aG5fCGPr963kae/n8HOinxkWaZrakceH38/A9r2BAIBVf3yEmIHZ2BOE81JhPBsJhv/GvcvHvjuAdyqO+LzZGSSrEn8ZeRfTuDozh4PfbSB+jCBVFMTaW6fxiPzNjG+ZzpJdlEr2RqIT+djUOXw8sRnW/h68wFkibCt0GPa90AymXHuXEFsj1HNvm6sxYQsuvkJwhlj7toijN7R0dRTypLEZxtKuHFYhxM8WuFEUx0+XNsOwlEfGXUeB5M/+gN/vehBJvYYh1f1s3rfBmKO7symadT/UELypNyTN2ihVRrRbgRPj3qah5c/HFFAZZJMpNhSeP3i10m1pZ6EEZ7Ziiqdh2qkQp+LZCJNkmDu2n1MG9P5JI9cOBYimIpSUaWT615ewUGHxzD3tTE5JpakUTdTueBlJFnBmjMQSTbhzl+Pu3AjyeOmNByryDCxf7sTPXxBEE6ijftqcBuk+EVTT+n0qmwpqT0RwxNOMseaA4G7JII/O/ZUFgEwqdcFANhkhbE5w0JfQAPnulISL81BtgSyHdx+NwsLF5Jfk0+tt5YESwKdEjtxQYcLsJqsJ/TnEU5vF3a6kMzYTJ5Z8wzbKrehaip+PXilxGayoekaF3W8iIeGPkSyNfkUjfbM8saP+YYdniOdSHP5NGYt28PUUTlikr0VEMFUFCodXq59+UfK6zyGsw1GEoZdjRybTM2KD6j4/Fkki42YjFwSzg0u7jTLMlNHdWr5QQuCcMrUGqRzQfT1lFVOb0sOSzhFnOtKwSCToXNKNrIk85sv/soVPcYzsH1vkqzxxi8iS3h2V1OR7eTtrW8zf/f8wGv7nQ2H2E12nlzxJFflXsUtPW8hOyH7hPw8wumvZ0pvftnj32w/uIflB+aRV/ol/th4zIqZpJgkrul2DVd0uYJ4S5i/N+GYzPu52HDCPZqJNIfHz9b9tfRpL9rUn+5EMBWFR+dvotLhjTiQOiyu9zjieo8L+7wkQbfMeHLTxcVMEM4k4eqcoq2nTLCaW3powikQrulEfEwsH9/8Ii+tepffff1Pyh2VjOsynGcm/I602KOaBWg63x34nid+/geqpuLTQgP2w4HVhzs+5ONdH/P30X9nfMfxLf7zCCdeQW0BxfXFuP1u4sxxdE7qTBtbm2bPK6118/bKAt5cUYCqBfrt6/poZN8gvEos5/dIZ9o5nRmYnSRa7p8Ate7jn0iTZYmKek9LD004AUQwFaFKh5dF28rCpvY115WrKbEWE/++YUBLD1kQhFMsNz2O73eEXjeiqaeMMcl0Tos9kcMUTpYmWpt3bdOJ5y57GIDdBwv41edP8cSiF/jvFY8HHbcsdh3P7nsTj96UDVKAAAAgAElEQVT8TZZf9+NX/fxh2R94Sn+KiztdfHzjF04Kr+plYcFCZm+eTWFtIWbZjI6OhIRH8zAscxiTe09maOZQw0DozR/z+euX2wAMOonawK/xzZYDLN5ZztCOycy8dQg2S2Sr5EJktDA7YUQ1kabTbDmJcHoQPbgj9P7qQsNCcggUE1YumkXiOdeTdd/btJ8+h/hBl+LatarJ15QkiIsx8fadw+mcFtfygxYE4ZS6cWgHZIObncb1lM6dK9B8bnTVjytvLVXfvxZ0rA5MGtj+JI1YOJGkCDsy5qZ25Po+E9hRHrzfTIFlP89mvBFRINWYW3Xz6PJH2V21O6rzhJNvS8UWxs8dz5MrnmRn1U7cqps6Xx31vnrqfHV4VS8/FP/A/d/dz7WfXUuFqyLo/OcX7eJvX23H49fCbskAoOng8qqs2lvJNTN+xOVVT/SPdlaxWoxvrxtPpEUi0SayEloDEUxF6N3VhYaF5IeLCVMunI69+whkixVJMWHPHR7UYKIxsywRY5IZ1CGZz+4fxYDsJMPjBEFo3Tqk2ukf5v2dMOxqks+fSs2KD9j3ws3sm3EHdT99jq3rkVx6SYJx3dNoExdzsoYsnEDWHslgUEy++2ABM1e/z/7aMgBKakv5ZNsiBrULzmyYm/oNPkJTBauWVbHr0V1suWsL23+1nZI3SlAdwTfHXs3L7M2zW/CnEVra2gNrmfzNZKo91Tj8jrDH6eg4/U721Ozhus+uo9RRCsBn60t4afFuXL7IAyOPXyOvvJ7p76w77vELRwzrlGI4AR/NRJpP0+jVLuHkDFg4LiLNL0JVDuMC8GiKCQFiLQrXDcnmjhGd6NRGpO4Iwpnu/vNzuevNdYY3OM3VU8aYZO4Z2+VEDk84ieJHtsexppSjC29jLXbWl2xj1poPqfXUkxATxwVdzuWRcfc2HOOQXSxL+BlNDp7Uq/iqgvKvysm6M4u4XnH4qnyUvFVC/rP55DySg2wKzJlqusaCggX8cfgfSbCIG7TTTUFtAb9c9EtcflfE5/g1P1XuKqZ8M4W5l3/EX77YijvMVi1NlSJ4/Bqr9lSyubhGNDtoIXeN6cKqvZU4DVb8ImlMpshw5YD2xIn9BVuFs/5fye1TqXH50HVIspuxmo1zWP1huk5EU0wYG6Pw96v7iRbognAWGd01jSkjO/HaD/lRzRhLwPSxXRjYQbQqPlOUmqDQrNPhqMWltvFpzJj05ybPXZS4CumouW7VpVI2v4z2U9sT3y/QwMiSZiH73mx2PrSTmh9rSB5z5O9HQuKzvM+4uefNLfMDCS3mv+v/a7gfVNWyKiq+qcBb5kWxKiQMTiDj2gyU2MA9h6qrlLvKeWH1XByedMPXjmRfI69fY/byvTwn6rdbxDmdU0iymw2DKWh+Is2syEwdlXOihie0sLMyzc+vBoovr/rvD/R+/BvG/vN7znv2e3o/9g0TX1jOl5v241ODZ3fCFWc2LiZsjixJJIj8V0E46/zfxd25Y0RHbGEma45mMysMy0nh262l1DiNu0IJrYdP1Zi1dA+XP7+MXX1TwBz9R+8ueyEeOThDwrnLiebTSBgcvNKkWBXi+8VTv6U+6HG36mZLxZbofwDhhKrx1PBd4Xdoeuiq44G5B8i8PpNeL/Wi85864z3oJf/ZfLRGZQcuv4sPd7+BwxuaAhppKYKq63y5aT81YbZzEKIjSRJPTOyN9Rje61azzPk90umWITo8txZnXTC1cGspQ55ayIMfrOfnompUTcft03D7NFRdZ1NxDQ/N3cDgvyzgi40lDef1TzEh6aHL59EUE3r8Gr1F/qsgnHUkSeL3l/TkpZsH0T87kRiTjOmo2hmzEqilHJaTwuzbh/D+XecwPCeV2+aspi5Mm13h9LeuoIqJLyxnyc5yPr53JLdf3YvUm3ogRXOTZZJx2EJXLdR6FVOcCUkJrc4wJZrw14feXNd6xQbQp5t5u+aFXXVsd0s74vvFI5mkhlVHb4WXmh9rgo736FXI1qKQ146mFMFiktm2X/x9HI+Cgw4+WFPIzCV5FFY6ubh3JjGmyFvPW80yvdslihXCVuasSvN7d1UBT34ePqf4MMehZdnfzt3A/ho315atZ+LnH7J60K24jjq1cTGhJCtYcwYiySbc+etxF25smPmRJBjbTRSSC8LZbFyPdMb1SGd3WT3vrCpgZ2k99W4f8VYzvdslcPPwjnRItTcc/6fLe/KnTzYzec4a3pgyLOy+VcIJoGmwbzXUloDfDdZEyOwHSZFtgFvj9PH3r7ezaFspj1zWkyv6t2toY23rlUqbyb2peHMraDq6N9xnkgYmE9ZuyaR0bAf564OeVeIU/PV+dFUPCaj8NX5McaF/L7FmUat7ulm6b2lIil8kq46NUziRVBT7XjR3h6Djo9ogXA+/0bgQnqrpfL+9jJeX5LGpuAZZkvCpGrIUSNdTtcDKhaL58MnG2UmKHDh2fI8MnrthABbTWbfW0aqdNZ/MC7eWRhRINeb2afzzi83oBUu49cWn+NdHBbiqQ2cHIykmtJkV7hrTuUV+FkEQWrfc9Dgen9j8HnSSJPHkFX34/f82cucba5kzeWjYuk5d1/H4NWJMstiE83g4K2HdG7Dyv+A71AxA10GWQfVC1jAY+Wvocn7gsaPous789cU8/eV2JvTOZMGDYw3bG8d0TqLdo+fg2lxB7eIi/BVupMMz2DqgqtjNy4m7+wHMmXHkbOiMRbbg1Y6k+tlz7Ugmidp1tSQOO9I4QHWr1G2sI+PajKDvaZbNdErsdNy/IqFl1XhrQh5rbtXRVRDcqEKSNCTFGXJsVPsaSYib+CjVuHzc8dpqdpbWNUzEN+ZVA4/JmoaumLGZZHSvF5NJQTIFrgs+TePKAe2ZOipHpPa1UmdFMOVXNf5v7gbDQKq5zXY9KPy76yXc1imHP10Wy28+XG/4Ok0VE1pMMr3bJTCkoygkFwQhOrIs8fdr+vHgh+u56611zLptMDGmwE1RWZ2bd1cW8vaqAg46vA2JQh1S7Nw1pjOTBrbHbjnzL/MFBx288WM+64uqqXP7sVkUctrEctu5HRnUITny4HLrpzDvrkAwE66rWv4yKPkJkjvBbZ9CbJuGp/LK63l03mZqXD5m3Tak2W0vJJOMfUA69gHp+CvdqHVedL+GbDVhSrMhz/4DuMYAo7ky90pe2fhK0PmKXSF9Ujolb5cgW+Wgbn7mFDNJI4K/vyRJXJV7VWS/C+GkMcmh79GoVx11QA8NlqLZIFzVdNLiRfZMpOo9fib99wf2VTmb3VxXk2U0HcxI9HPuZ+ol/ZE7dSbRZqZXuwTRta+VOyv+9RZtL8NnsB11JB1uDvtq0wEmDWxPXnk9//0+L+KuXBaTTPskG6/dYbxTuSAIQnMUWeL/Xdef+9/7mV++8xP/uKYfj8zbxHc7ypGgYXPOwx/n+QedPPXFNv7y+TZuH9GRhy7ugWKwv1FrtyLvIP/v2x1sKq5B0/WgG5pNxTUs2FpKm7gY7hvXheuGZDd9Df75XfjiwfBBVGNeB5TvhJdHwd3LcMek8NL3u3lrZQH3n9+V287tiEmJbobflGLFlGINeszT91q+WfMf3t3yX8qcZSENCgDSLk1DiVU48MEBvGVeZJtMwqAEsu/ORj6qLmtw+mAyYzOjGpdw4qXb0tnK1qDHol11VCQLsh5akx1pKQIENojt1VbUdUfqrjfXUlztajaQaszlU9kQk8FuOZ5f9spo/gShVTgrgqmXF+fh8AQHP4c73KRe+mvs3Uc0PG7PHY49d3jQsQ6vyowleUwa2J77zu9Kos3MU19sQ9V1/E28iewWhZ5tE5gzeSjxVtHFTxCEY2dSZP5z40CmvrGa0c98j1/V8arh05YPt+R948cCtu2vY9ZtQ86oFJ45P+zlH19vD5u6reuB30FhpZPHP93Ksl0V/OuGAZiNgpz8HyIPpA7TfOCowDHrUib6/kb3dsl89cAYMhOtzZ/bDI/q4fmfnuejgrngdeJ0NR0Ip4xNIWVsSpPHWBUr0/pNO+6xCS3vqq5XsfrAapz+I2l60a46yjLo9capw5GWItw9prOY9I3Q5uIafi6sxus/hownk4UZK4q584KeDVkGQut2xgdTLq/KxuLQfORoN9vdU15PlcOLzaJgNStkJlgprm70wavrIEnIUiCVYmSXVO45rwvndk4VFydBEFqET9UornKH3bvEiMunsmrvQR78cD0v3DTwjLgevbuqkGeaCKSO5vKpLNhWyoMfruf5Gw1+B98+GhJIdfp3HU4f7H0gjlhL4PhXf/Ly9kYfi+841MRB80FNIf8edYB+F1xw3D8XBNpk3/nNneyt3YtH9UALrChaFSuT+0xmaObQFhih0NLGZI3BrJjhqOaLka46SkgMbzsMp68zy3eXYzTH0ty+Rjo6Vw/OMn5O16l1+al1+zApEsl2S9jazbPF7OV7DQOpSDOedF1vyHgSWr8zPpiqdnmxKBKuozbdjarDDWBRZFbsqeDhjzfjVbXQm5lDH86aDnazzI7SOjITrGfEjYsgCKeHl5fkBU/iNNLUbKjbp/Hd9jKW7CznvO7GG3ueLG6fyhcb97N0VzmVDi8WRaZ9so2rB2XRPyux2Wvm7rI6nvx8S1TNhALfV2Ph1jI+WreP64Y06shXvgPKthqeo+rwn1VeHh4dvo4kFjf98l8HbolqPEY8qoc7v72TvJo8fFrLdFWzKlZu7XUr0/tPb5HXE1qeSTZxc4+bmb15diCAbiSiVUeTlSl9ptBhWF8u+c8yqhxeIk88C2wQPig7OWQfvDq3j49/KuaVpXsoq3NjVmT0Q+m0I3NTuXtMF87tcvZNGNe5fXy5aT+qHvxbjjbj6eVDGU9C63fGB1MBoW/0qDrcENjQ7sEPjZtYHM3pVXH5VK548Qc+vneE6M4iCMJx86kab64oaKiPaiyS2VCnV+WVpXtOWTB1oMbNzKV5fLCmCAmCOl/JEsxdu4/MRCvTx3bhmsFZYWu8Xl22N2RT9cOaS69x+VRe/G431w7OOnIDuHIGaKH7MQE8NMLCMz94uHeohSRrEzeMBzZBxW5ok9v8L6IJMzfMZG/N3uMOpEySCUVW6Jbcjen9pzM6a/RxvZ5w4k3uM5mFhQvZU70Hv27892jEqli5JOcShmQMQZIkPrz7XG6YuYJql9dwhSrkfLPMHy/pydKd5dz+2mpm3DyYBJuJF77bzUuLdyMhNdSI+9Qj79klOytYk19Fos3MzFsH0y+r6WYrZ5K8cgcWRQ65Fkeb8bS7rL75g4RW4bQPpnRNx72zCue6UtRaL7qmIdvN2HqmYB+YgRzTdCCUaDMb1hVE0+EGiHoWVNfB4fFz4ysrWfTgWJJjLVGdLwiC0NiibaX4DRrpRDMbuq6gin1VTrKS7Ue/zAm1aV8NN89eicurGhZra3og0Nlb4eDxT7fw6YYSZt02BJsl+Pru8PiZv77Y8CYx0vSa8noPPxdVM6jDoe6q+cvCBlND2imc18nEsz96eOr8JmqhZAX2rTmuYMqn+nhv+3shKxMAVcuqqPimAm+ZF8WqkDA4gYxrM1Big38/yTHJdEnqQrfkbtzY40ZyEnOOeTzCyWU1WXn1ole5/avbKa4vDmqB39Q5o9uP5rFzHmuYHMhNj+OrB0bzyPzNLNkZ3KDmMFmCGJNC+yQrT1zRh1Fd23DLOR35y+dbueql5XRNj2fpropm73ucXhWnV+WGmSt55bbBjO6adsw/f2tS5/YZzdFHnfGkajo+VTOu4xRaldM2mNJ9KnXLi6lfXoLu09CPSqvz7q2h5ou92AekET++A6Yk4w86u8VEj8x4tpQE7+odTYebpjQ1E6oT+PB/a2UBvxrf9Zh+D4IgnF3KnGVsqthEnbcOi2yhja0NgzIG8dG6fSGNdCD62dBvt5QyZdTJu8necaCOG15ZEXGdl8unsia/ktteW8W7084JutH4evMBZIOUomgCSpdP5c0fCxjUIRld19HdtTR1K/PkuBhGvubggeFNTIipfnCH1uZGY1HRIsNufRVfVVD+VTlZd2YFNSHIfzafnEdykBs1FXGrbl664CVsJttxjUU4NZKtybx/+fs8tfIpvi34FgkpZDNfALvJjizJTOkzhTv73hmSZpeeYGXWbUMor/Pw3upC5q4totrlQ9V07BaFczunMm1M56DVJEWWeOKK3tz8ah3fbC2Natwun8pdb65j7j3n0qd9YvMntHLh6sWizXjSgfveWcdzNw48K7awOJOdlv96qsNH+axN+CtcYJDSAjTsGO9YV4pzUwVpU/tiyQ5Np6tx+shJjQ0JpiCyDjdNiWQm1OPXmPPDXn45LveMbE0sCMLx03Wd1QdWM2fzHNaUrsEiW1B1FQkJSZIwSSYkz0gk0wB0f3Dr4mhmQz1+jYOO5me8W4rXr3Hzq4EVqWh4/Bqbimv45zc7ePjSng2PF1e7DF8rmoBS1+Hrzfs55+mDVDq9LFZU2jVxae6TrnB5NxN/X+6lZ1qYsEuSwXR82QcfbP8gqJsbgOpSKZtfRvup7YnvF/h8s6RZyL43m50P7aTmxxqSxxzZv1BGZnnxci7seOFxjUU4dexmO0+PfprfD/s983fP593t71LhqsCn+ohRYshJzGFyn8lc0OGCQNOKJqTFx/Cr8V0jnswtqnSyNr/K8LlIUmgfnreJT+9rPtOntctMsBo2n4g24wlg8c4KrnjxB/53zwgS7aLrc2t12gVTmkel/OUN+CvdgerfZk8A3a1SPmsT6b/sjzkj0GWprM7N7OV7+WBNEeO6pRNrUQx3p26uw034cUY+E+pVNb7fXsYFYk8BQRCOUuutZfqC6eyu3t1wM+1VDQIey7fEdvkW94Er8dcc6coW7WxouHqjE+GbLQdweVXDYvjmbs7cPo23Vxbw4IXdGmaC6z1+w9eKNr0mI8HK+3efE+hK9nonKD7Y5PF/Ps/KoJn1/PbcMI0oFBPEHd/1vdQZuhrg3OVE82kkDA4OoBWrQny/eOq31AcFUz7NR5mz7LjGIZweEmMSub337dze+/aT9j3f+DEfTQ99h0WaQrvzQB27y+rITT+z68SzU+zkpse1SMaTx69ReNDBba+tYu49I86o7SvOJqddMFU1bxf+qggDqUZ0r0r5q5tR7+7DzOV7+HR9CVcNbM/n948iK9nOZxuKeeijjVHXPslSIJ//aNHMhDo8Kqv2VopgShCEILXeWm76/Cb2O/Y333RA8iNJYM38FI/swlc1BohuNtQkS6QcS/1mbQmsmQ07vgJ3NcgmsLeBgbdAv+shJs7wtJeX5BlOYkWzYfpnG0q4ZlAWBZVO9le7kCAkoIo2oEyLj6Ft4qFUuCFToHx7YCPeMHJTZG7obeb51V76phvc7OgadBnf7Pdtik8N/fdX61VMcSYkJXTpzJRowlUQ3NlR0zXjQFwQmuH2qby3ujCkpjGaiWOfpjF7eT5/u7rvSRnziVDv8TPvp318v6OcKqcXsyKTmWDluiFZjOzSBvlQhtE9Y7vwh483hqReh8t4srTtyv43fm04eeRVdXaW1jP/52KuH5ptNCzhNHdaBVOqw4drcwX4Q6OX1fs28vT3M9hZkY8sy3RN7cjj4+9nQNsjKSBuh5e/P7+C3BFZLPrteaTFH5lFnNi/Pftr3Pxrwc6IAyqbWSYl1kJxdWjOcrQzoQfrQ4uKBUE4e+m6zr0L740skGpEkn3EpC9A87VBre8V1WyoIkuM6JIa+SBLtwb2YMpfHvj/xs0RqgsCQci3D0O/G2D842A/0sJ5T3k9eeWh3aqiuTlzelUe+2QzT362lQSbmYyEGMyKhPeoG75oAkqzIjEgu1Hnsd5Xw5cPNfeb4LGxMby10eDfSTbDwFvBfHyb9cZaYiE4yw8lTsFf70dX9ZCAyl/jxxQX/BFukk3EW87sVQHhxNhQVG1YjxjNxLGqBVajW2MwVVTp5IXvdvHphhJkSQqp8Vy0rZTYGBPTRnfmthEdubh3Jn+avxloPuMpkskjly/QKl0EU63TaRVMOdYcAIN5xzqPg8kf/YG/XvQgE3uMw6v6Wb1vAzFK8AyrRYM/Z6TSdkIPw9e/a0wX0uOtPDJvU+D7hcnjtx/qIPX4xF58sWm/YTAV7Uzo2b7BnSAIwdaWrmVn1U7DQKq57m2S7MOa8TmO+p6AFHH9p1/TeeyTLfxieAcu79e26aLnvO/g/ZvB5wx/jO/Qas7P78DOb2HKV5DcCYCCSidmRQ6ZvIq2YYYOLPndOFJiLaiazpCnFuB1Bv/OogkoZUni9hGdjpxsscPAW9DXvYHUKFjM/3VwUJKdKON+NDjdLvCCCgy/O6KfpSlDMoZQUFMQ1BbbnmtHMknUrqslcdiRwn7VrVK3sY6Ma0OzHfq2aX03ssKpV+X0tUgKrcMTeVv308W6girueG01Tp8/bDt5h1fF4VX5fwt28MWmEubc0Ie/HVjErxNG4JHCX0ejmTzaX+NmQ1E1/bPPnjbzZ4rTK5j6scSw4cSeyiIAJvUK7DBvkxXG5gwzfA11vxN/tQdTknFu+6SB7ZnQJ5MvN+1nxpI88svqMWt+JGsMPr9O+2Qb08/rwsR+7bBZFNYX1Rim+kU7E9o++fhmLQVBOLPM2TwHtz90oibS7m2SqR7ZVojm6gg0X/9ptyg8MbEXKbExvLe6kL9+sY0r+rfjF8M70LPtUUFC0Rp4/xfgM94gOITmg/oDMPsiuOcHiEvD6VHRDeovor0586t6Q2qiIktMGZXDi9/tDmn3HGlAOSA7ieyU4Nbw2vmPUbnha5LUYkwGM81hme0w/jFI6Rz5OWHc0vMW5u+ej189cjOq2BXSJ6VT8nYJslUO+nswp5hJGhF805Udn033lO7HPRZBOCzaieNo+VQfJtl0yjb+3VJSw62zV0XcbdTt09hSXMt1f/mEVzNimXXjcO55d33Y86OZPPL4Vb7dckAEU63QaRNM6bqOWm+c6905JRtZkvnNF3/lih7jGdi+N0lW41QGySShVrvDBlMQWCW6elAWVw/KIu+1tykvLKHNvfeSZDfTJi74vOuGZDF/fXFIB6loZ0In9hO7XAuCEFDuLGfV/lXoR80FR9W9TfJhSVmKu/jWZr+fSZZon2TjigHtsZoVLuiVQUm1iw/WFDHl9TVkJFiPrFbJGrx7XeSB1GG6Bs5KmHc33PoxsTGK4Q3S8a7q/2JYB15Zusdw8+LmAkqrWebBC7sFPeZTNX7/SR51CX/jZf9jULsvOJ0xHLMNRj4A50xv/tgIdErsRPfk7mys2Bj0eNqlaSixCgc+OIC3zItsk0kYlED23dnI5iP1W3aTnSl9I9vSQxCOlmQ3G22dFHWHutiYpm8rXX4XX+39ijmb51BUV4Sma0iSRHJMMjd0v4Hrul9HG1ubY/wpouP1a9w6e3XEgdRhPk2nyJ7KK/0G8nSPTObdO5IbX1lBldOg7jGKySNNh9I6URLSGp20YGpvhYPdZfU4PH5sFoUOKfbg2VDt0JeB+JhYPr75RV5a9S6/+/qflDsqGddlOM9M+B1psSkhxx9um94Uv6pR6/ajHqygU0YimenGBdQDs5NIj4+h4GBoqkukM6H9s5PokHpyN8kUBOH0teXgFiyKJWRjzmi6t0mSjmIvaPZ7WRSJ1LgY3p12TlBg0i7Jxm8u7Mb95+eyeEd5w2rVwx22ca3fy9Ef/Z3+XYfTB3sfiCPWErjtevUnL29v9LH4jkAXVTQf5C9Hry5C1SyGNynR3pwdfe1MjYvhzSnD+MWsVbh8kd8E2cwKf7ykB8M7H6kZc/tU7nv3Z/yaxoxpl6DoY+HrP8Kmj0CSDFIcpUBaoC0VLvoL9J4U8fePxINDHuSeBfeE7C2UMjaFlLGhn3WHKZJCsjWZizpe1KLjEc4eA7KTUA1WkqOryYQLwzTa0nSNF356gXe2v4OEFLQNgK7rHHQfZPbm2by66VXGZo/lyRFPEmcxvi9rKd9sOYAnzDWkuW6jHl3i45+L+eOlPemeGU/vdoks310R8jpR7z1l8G8gnP5OaDDlUzUWbC3l5SV57CytwyzLaLqOJEmomk7bJCv3jO3CFf3bBT7kFSlsF7+ubTrx3GUPA7D7YAG/+vwpnlj0Av+94vGQY6UY4z9YTdNZsrOcl5fksSa/EpMig68L/oMSA2f8yD1ju3B+j/Sg/aAkSWL62C78+bOthh/ckaTW3DP2+FNABEE4c9R56ww3aI22e5skh+/cJksQY1LonhnPnDuGkhymi59JkbmgV0bDapUy61EUn3FnO1WH/6zy8vDo8Cv/fk3jvRf+xCsxt5FiN1N+VMZBNDdnsRaFu8eEXj8Hdkjm3WnDue211fj8Gu4w+xEeZlYkHr28BzcP79TwWL3Hz7Q31pIaZ+Ff1w861JI4Hq58ES5+GjZ+AGtmQX0ZqF4wx0L7wTDyV9Dh3ECw1cIGZwzmj8P/yN9W/c1ws1YjiqSQYEng9QmvY1GOb68r4exlNSvcMDSbt1cWhHT0i3TiWNOgW3ocmqY3dL2DQMv+B757gDUH1jT5d+05tCK8pGgJ139+PW9e8uYJXaU63m6jsiTxv3X7uGNkTkhW02HRTB5JQHq8KAlpjU5YMJVf4eCmWSupdfsaWke6j1p62lPu4M+fbuGvX2zjzSnDyEix4i9vPrUkN7Uj1/eZwNvrPw15TvdrmNqE7v6+eEcZ/zd3A65DRYQQWOJFCqRJrCuo4tfv/4zFJPOPa/pxUe/MhnOvG5LNF5v2s3pvpWFqSTg2s8KEPpmM654e8TmCIJz5zIoZySCpJtrubXazhR4dk9lUXINFkRv693hUjQt7ZjBtTGf6ZyVGXI/QjnLwhF/temiEhWd+8HDvUAtJVuPXNOk+bjIt5paH5vDNllJ+O3d9xO2DQzZMlyQm9MnEyMAOySz+v/N4d1UBM5bsaTJVx6TI/O3LHYCiPDkAABWYSURBVByo8fDA+K7Uuf3cMWc1vdol8tSkPqEbqlsTYNi0wNdJdnXXq4k1xfLoD4+iozfcYBqxm+ykWFOYM2EOmbHGvydBiNTkETm8u6qQ0M0HItuTMzvFzicbSnh/TRH3nZ/L5f3aIUvwyLJHmg2kGvNqXvbX72fqN1N577L3sJtbPrMnv8LRIt1G5/yYzx0jc7ikbyYLth0IudZFM3lksyiM7ynuF1ujExJM7S6r56qXfsDh8Rvu0dRYILBRufGVlcwc0YUuKzwhaXq7DxawKG8FV/Q4n7YJ6ZTUlvLJtkUMatc75PWs3ZJRYoN3kZ67tog/fbK52Zboh7u1/Or9n3nk0p7cem4nIFD0/MqtQ7hjzmo27quJKLXEZlY4r3saz1zT75QVVgqCcHpKs6VhVKAQbfe2NHsKH90ygtJaNyXVLlxelTiriY4psSTazUe/fPNqS0CxgEFjDIAh7RTO62Ti2R89PHV++BlUk7cagAt6ppNgNePyqiGfBc3dnNnMClNH5hBjCp8ak2y3UFjlornMmMM1r68u28vyXRXUun1c2CuT30/oflpeny/OuZjBmYOZu2Mu72x/B7/qR0ND13VkSUbVVXKTcpnadyrnZZ+HWT6Gf2tBOEqHVDu3nNORd1cVRpVCC4H3639/MYg+7RNYuquCFxbt4t8Ld3HRkBoWFy82DKSa6lrq1/3sq9vHa5tf476B90U1Fr+qsXx3BSXVblw+lXiriZ6ZCfTNOnJNLapqmW6jpbWBn2t8j3TMsoxRq/RIJ49SYy0M7pgccr5w+mvxYKrG6ePGV1YEdqqPIvXT5VOZviKPV3w2OhK8KWKsxc76km3MWvMhtZ56EmLiuKDLuTwy7t6g4ySLTPyYrKDHvt9RFlEg1Zjbp/HXL7eRFm9tmBW1WRTeuXM4/1qwkzd+zEf3+3FqoR/CsRYFkyIz/bzO3D2my2n5QS0Iwqk1IG2A4Q1wNN3brIqV67pdB0BGgpWMhBZIDwkTRDX25LgYRr7m4IHhTaWUSaD6MJksvDvtHK54cTn1br9h62UjVpPMsJwUHriga5PHPfHZFj7fsD/iGz+XT+Xnomo6ptp56OLTM5A6rI2tDdMHTGdav2msLV1LqaMUt99NvCWeHqk96Jwo0seFlvfIpT0pq3WzcFtZxO8rq1nmpVsGNQQrY7ulMaZrG1buqeQ3S+/BJblCJo8i6Vrq1by8t/097ul/Dya5+dvVslo3b60s4M0VBfg1DVXT0TQdkyyjA+0OlZZM7N8Op1c1vEeNttuo91C2kkmRuX1EJ15ekndMzXFsZoV7xop7xtaqxYOpd1cXUOc2DqSaK+hzeFVesPl41m+DRsFP2/g0Zkz6c9PfWAZTqg1LxyOF235V48EP1hsGUs2Nxe3T+L+5Gzi/RzoWk0yN00eV08t1Q7K5fURHvnjs38xN6cN+YnD7VGxmhc5t4pg2pjMX9EwP1GMJgiAYUGSFW3rewqxNs0LSuCLt3qbpGld1vaplBxZjsI/SUfqkK1zezcTfl3vpmRbmOicpYAoEWzltYvl4+ghufGUl9R5/s6nSdrPCmG5p/OemAaHpd42syDvI3LX7op5BByir9fD+6kJuPqdj1OeebCbZxDltzznVwxDOErIs8fxNA/nXgp28snQPsiSFfY/FWhRiY0y8fOtgBnUIXlGRJImOGR585j0hizXRdC31a36WFC1hfMfxTY77q037+c2H69F1Qq4xXjUwgLxyB49/uoV/frOD303oblj2eDzdRqeMymHu2kIOVLvQpMjvAU2yRIcUO9cMzmr+YOG01KLBlKbpvLpsr+GHZaQFfev9Pvwd0zAV1AcFVE2SJeRYE22m9AmK6r/bXobXYAe2SMei6zp//WIrq/ZWkldej0UJzG74VI22/kzuH9uViSO6ig15BUGI2rXdrmXWplmGzzXXvc0smxnfcTyJMYlhjzkmqbmBZgvN+PN5VgbNrOe354ZpRNEmeEWpa0Y8Cx4cy1sr8pn9w15cHhWfQQ54XIzClJE5PDC+K0ozE1Izl+SFvclrbrLM5VOZsSSPXwzvIGaCBeEokiTx24u6c+fozsxdW8SsZXuocvgwK1LDPdCQTincM7YLo3PbBDWbaOy7wu8Ma0Oj6Vrq9Dv5JO+TJoOpj3/ax8PzNkWUgeT0qrh8Ko9/ugWfwb1qtN1G2yUdqdGP1308t+MjpmdeTK0cg7+5OhcC3VbTE6y8M224uJdsxVo0mFqyqxy3P/TDLZqCPlmSWNozgctsFtw7Kpttcy6ZZeQEC+l39UOJD047mbl0T0gxYDRjcXhV3lxR0JCa4lOPvFZhXDqPL8rn8UX5PHVlH64WMwqCIEQh1ZbK4+c+zpMrnoy4MBtAlmRSrak8MvyRlh+UNQF6XwUbPwQ9/IpPborMDb3NPL/aS9/04KBHNceijPpNyDlJNjP1Hn/ghidM/FLvUXl1+V4+XLuPWbcNCapxaKy01s2KPQcNn4t0sqzS4WVtQRVDO4UPWgXhbJZoM3Pn6M5MHZXDQYeXGpcPiyKTEmtpdj8pgIOug4YNVKLtWlruLA/7PX4qrIo4kDpM1wNBlckgCIymYQQEGln86r2f+celuZT+cjrZHTry9e8vZtpbP7GztA6vqqMaBFUmWcIkSwzqmMyMWwaTaBN1j61ZiwZTq/dUhgQvEF1Bn9Orsmz3QW69bTDuHVXULSnCW1QXaC5zuF2ndCiIirUQf14W9oHpyJbgiN7tU1lfVH1cYwGjnjbBYwV4eP4m9te6+eW43IheUxAEAWBil4nUemt5bt1zTXZtO8wsm0m1pvLGJW+0/KrUYef+ErbOb3bT3sfGxvDWxtBNKl0+jT9v6cSvspxkpwS6cKmazvS317FsV0VDjUE4Tq+K06ty/cwVvHr7EEbmhrZG/nzjfsNzo5ksc3lV3l9dKIIpQWiGJEm0iYsJ2/47HL/uN3w82q6lahMTO//8ZkfYQKqpFWpdB/x+zLqO76h6rIi7jQJ+TefbLQfYsWYTMzp2pu1fnkCSZT65bxRbSmqYvWwvX27ajyxLyIe2BQK4elB7Jo/MITfMHqdC69KiwdRBh/HNQLQFfVVOL5IkYeuRgq1HCv6DLlxbDqLWedH9GkqchZguiVg6JoRN0ahyerEoMi4t+E0Y7Vgi4fZpvPDdLjITrCLnVRCEqNzc82Y6xHfgmTXPcMBxAK/mDdmDyqpY0dC4sMOF/HH4H09cIAWQ2Reyh0PBCmgU4OX/Oj7osOxEGfejR9VYme2YRvyWtv4kJr64nKsGtue+cbk8t2Any3ZVRFXf5PKpTHtzLR/fO4IemcHfZ3+1yzCdPJrJMh0ormp+Kw5BEI5NUkwSiqSEBEPRdi0tqZT57/e76dUugd7tEhr2YtpX5eSngirD7x3JCrVPUlAUCQVCVo8iaQV/mNuvsdeSxGM5l/Jmo8TG3u0S+dcNA3j66r6U13lweP3ExZhIi49pskup0Pq0aDAV7o8j2oK+GFNw2ogp1RbSpa854ToJRjuWwyJpWPGnTzZzWb+2Iu9VEISojM4azeis0Wyu2MzrW17np9KfcPqdmCQTydZkru12LZNyJ53YIKqxG96GWedDVUFQQNUksx16XI71vN/yoCRx6zkdefG7XYx7djFOr2pYP9DcddXpVXl03mY+mj4i6DxnmKAs2smyaPYNFAQhOsPbDmfmxpm4/MGTFtF0LY1RbIzKGEOty8ery/awpaQWiyLTu10CNS6f4XUlmhVqsyJhlmUc3gi28mnieuVF4afCapbsLGdcj+C9oqxmpWGVXjgztWgw1T7ZhlmRQnbPjnYH6MYFfccqyW42bD4RbXEhRJ6DLwFfbNwvVqcEQTgmfdr04dmxz57qYUBMPNy5EN65Fkq3gNcR/tjDnfsG3goT/s7hFllp8TH8+co+VDq8hml5kV5XNxXXUHDQQcfUWAAO1nuoc4WmF0L0k2XHtBeXIAgR6dOmDxn2DPJr80Oei7RrKWg8OvY24iyBdDhd1ymudrGlpJbHPtlsWI8UzQq1quncOaoj834u5qDDGzZlMJLrldOr8vKSvJBgSjjztWgwdVnftjy3YCdHVxpFuwP09UOzj3ssdouJrulxbD9Qd8xjgegbVsxYkieCKUEQWj9rIkz+GnYvhOX/hpJ1gcBJ8wMSKObAf/e8IlBn1W5AyEs4PH4WbisLqT2N5rrq1zR+99EGMhNtrC+qptLhpUOK/bgn7mxmmVEG9ViCILScKX2m8LfVfwtZnYLmu5bKyEzoNKEhkIJA/VZWsp2sZDvPfL2dUgwaXESxQu1TdSQkvvnNWB6au4GvNh8IOSaa69X6omqKKp1iJeos06LBVHaKnYEdkli5pzLkuUgL+lJiLQxpoR2gp5/XhYc/3oTDG5wSEk1xYbQNKwoOOqh0eEmJbWpDS0EQhFZAVqDbxYGvyr1Q8AO4qkE2QWwbyB0PtvDX67UFVYYds6KbOYaN+2q5ZlA2943LpUta4MZq6F8XctAR3MY9mskyTYfrBh//xJ0gCOFd2vlSXt/yOoV1hfg144YU4djNdqYPmB72eYvJePuEaFaoZV3DPf8jahbWUBIzFAgNgqK5XknA4p3l3NoK9rATWk6Lb9p7z9gubNhXg8sbmtMe2Q7QnVts348JfTJ5ZN5mw+ciLS6MNgffoshUO0UwJQjCGSYlJ/AVhWqnF82ggDXa66oiSyEZC3eOzuHfC3eF1D1FMlkmSzChd6ZI8xOEEyxGiWH2xbO54fMbqHJX4dOMU3SPZjfZmXnhTNrHtQ97TFaynW3760Iej2aF2moxkTNuFLGWKmrWGR8TzfXKq2rUOJvfq084s0S+RXOExnZL47K+bbGZo3vpGJPMwA5J3Di0Q4uNJcak8OQVvbFGOZbGGs9wCIIgCFEymBtrievqL4Z1JC7GZLhlVVzvcbS9/d90ePB/ZN/3NunXPYE1q2fD81azwq8u6GpwpiAILa2NrQ0fTfyIbsndsJlsyE3cesaaYmlja8Nbl75Fv7R+Tb7uLed0JNYSGuA0XqF27lyB5nOjq35ceWup+v61oGM1YOIVI0iaNAklKSnktUDcBwrNa/FgSpIk/n51X87vkYEtwq52VrNM3/aJvHr7EEzN7HofrasHZ3HfuNyIx3K0xjMckfCqGsl2sSolCIKQZLcgG4Q70V5X4ww2CE20m3n/rnOIjTERTTKD1Swz89bBDemCgiCceMnWZN677D1mXTSL8zucj0W2EGeOa/iKUWLon9afp0c/zYJrF9AtuVuzrzk6tw12i3GCVcKwq0k+fyo1Kz5g3ws3s2/GHdT99Dm2rkdS9RRZYmK/dsRbAyvU4TKKorleWUwyieIe8KzT4ml+ACZF5sVfDOSVpXt4aXEefk0z3Mw31qKgE5hdeOji7phbOJA67L7zu5KRYOWxT7YgSUc22z2aJIW2VI+2YUVOm1iSRYqfIAgCQzslG7Yujua6alYkLu2bafj6XTPi+fS+kdw0ayX1bn9IfWxjNrOCIku8dsdQhuWIjXoF4WSTJIn+af15btxz1HhqKKwtpN5Xj81kIzM2k8xY4/d5OLIsMW1MDv9asNOwC19z5RxmRWLq6COpy1cOaMeWktqQe8Rorle6DuO6p0X1cwit3wkJpiDwprl7bBemjsph4bYyXlmax55yBy6fitWs0C7JytRRnbn8JO3LdN2QbC7r15ZP15cwY0keB2rcDcGbT9VIi49hVG4bPt1QjNMbfQ4+BILD6ed1OeE/iyAIQmtgt5i4elB7PlhTFBJURXpdlSWJO0aEr9XqnBbH0t+N4+vNB3h5cR57DzowyRKqFph51nSdZLuFe8Z25qpBWYarXIIgnFyJMYn0Tet73K8zZWQO328v46fC6qj2jbOZFX59QdegDcGvGpjFU19sMzw+0uvV4I7JZCWLTn5nmxP+qWJSZCb0yWRCn+hmHE4Eu8XEjcM6cMPQbEprPVQdKhJMspvJTLDi13S+2ryfQBZtsIgaVkhwSZ+2J2DkgiAIrdOUUTl8tG6f4QpVc9dVCeifnUSH1KZvTmJMClcOaM+VA9qzs7SOXaX11Ht82CwmspNtDMhOarHGRoIgnD5MiszsO4Yyec4aNu6rwRVmQ+/GbGaZu8d05u6xwZPfsTEmJg1of8zXK7tFCXlN4exwVk7RSZJEZqKVzERr0ONmReKZa/vzwPs/h924LRybWeHpSX1PyiqbIAhCa9ElLY6bhmXzwZp9Ed3oNGazKDw1qU9U53TLiKdbRnxU5wiC0HrZLSbeuXM4Ly3ezezl+YalJbIUmHRpl2TldxN6cHFv4wn+By/qxoKtpVQ6vCH74zUlxiQzPCeF0WLvurPSWRlMNeXi3pk8fGlPnv5yW8QBldUs85sLu3LlwPAtPAVBEM5Wj13em/J6L99tK4s4oLJbFF67Y6gIjARBaJZJkfnV+G7ce14uC7eV8fqPe9lX5cLj14i1KPTLSuLO0Tn0yzLu2HdYeryV9+86h2tfXkG9249qsLXD0axmmV5tE5hxy2Bkg331hDOfCKYM3HZuJ9LjrfzufxtQNd2weQYEaqRkWeLpq/owsb8IpARBEIzIssSLNw3k2W938OqyvciSZBhUSQRWo1JjLcy8dQi92iWEvpggCEIYLVFa0jUjni8fGM20N9ayt8KB168ZBlUxhzYNvrxfO56+qm/YTYSFM58IpsKY0CeT8T3TWbi1lBlL8thaUhvUsKJbZjzTx3bh4t6Z4g0kCILQDEmSeOjiHtw9tgv/W7uPV5btoaLeg1mR0XQdv6ozplsad4/pzLCcFFHjJAjCKdM+ycaXD4xmc3ENs5fv5ctN+5EITAz5NR27RWHyiBx+MbwDafExp3q4wikmgqkmmBWZS/q25ZK+bXF4/FS7Ajt3J9rMoiOUIAjCMUiwmpk8Koc7Rnai1uWn1u0L7M1iM4uaU0EQTit92ify3A0D+Mc1/ah2eXF7NeKtJhJtZpHSJzQQEUGEYmNMxIoAShAEoUVIkkSi3Uyi3XyqhyIIgtAki0kmPd7a/IHCWUnkpwmCIAiCIAiCIBwDEUwJgiAIgiAIgiAcAxFMCYIgCIIgCIIgHAMRTAmCIAiCIAiCIBwDEUwJgiAIgiAIgiAcAxFMCYIgCIIgCIIgHAMRTAmCIAiCIAiCIBwDEUwJgiAIgiAIgiAcAxFMCYIgCIIgCIIgHAMRTAmCIAiCIAiCIBwDSdf18E9KUjlQcPKGIwjCSdBR1/W0Uz2I4yWuT4JwRmr11ydxbRKEM1LYa1OTwZQgCIIgCIIgCIJgTKT5CYIgCIIgCIIgHAMRTAmCIAiCIAiCIBwDEUwJgiAIgiAIgiAcAxFMCYIgCIIgCIIgHAMRTAmCIAiCIAiCIByD/w84YL7pSAqAmgAAAABJRU5ErkJggg==\",\n \"text/plain\": [\n \"
\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"# Visualize examples from the training set.\\n\",\n \"# It takes a few minutes to download & prepare the dataset.\\n\",\n \"ds, ds_info = tfds.load('ogbg_molpcba', split='train', with_info=True)\\n\",\n \"tfds.visualization.show_examples(ds, ds_info,\\n\",\n \" node_color_fn=node_color_fn,\\n\",\n \" node_label_fn=node_label_fn,\\n\",\n \" edge_color_fn=edge_color_fn)\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Training\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# Start TensorBoard\\n\",\n \"# Get a live update during training - use the \\\"refresh\\\" button!\\n\",\n \"# (In Jupyter[lab] start \\\"tensorboard\\\" in the local directory instead.)\\n\",\n \"if 'google.colab' in str(get_ipython()):\\n\",\n \" %load_ext tensorboard\\n\",\n \" %tensorboard --logdir=.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 10,\n \"metadata\": {\n \"outputId\": \"7696aae4-beb5-4df2-b72a-7f391ac30c2e\"\n },\n \"outputs\": [\n {\n \"name\": \"stderr\",\n \"output_type\": \"stream\",\n \"text\": [\n \"INFO:absl:Unable to initialize backend 'tpu_driver': NOT_FOUND: Unable to find driver in registry given worker: \\n\",\n \"INFO:absl:Unable to initialize backend 'tpu': INVALID_ARGUMENT: TpuPlatform is not available.\\n\",\n \"INFO:absl:Hyperparameters: {'add_self_loops': True, 'add_undirected_edges': True, 'add_virtual_node': False, 'batch_size': 256, 'checkpoint_every_steps': 10000, 'dropout_rate': 0.1, 'eval_every_steps': 500, 'latent_size': 256, 'layer_norm': True, 'learning_rate': 0.001, 'log_every_steps': 500, 'message_passing_steps': 5, 'model': 'GraphConvNet', 'num_classes': 128, 'num_mlp_layers': 2, 'num_train_steps': 1000, 'optimizer': 'adam', 'skip_connections': True}\\n\",\n \"INFO:absl:Obtaining datasets.\\n\",\n \"INFO:absl:Load dataset info from /root/tensorflow_datasets/ogbg_molpcba/0.1.2\\n\",\n \"INFO:absl:Reusing dataset ogbg_molpcba (/root/tensorflow_datasets/ogbg_molpcba/0.1.2)\\n\",\n \"INFO:absl:Constructing tf.data.Dataset ogbg_molpcba for split train, from /root/tensorflow_datasets/ogbg_molpcba/0.1.2\\n\",\n \"INFO:absl:Constructing tf.data.Dataset ogbg_molpcba for split validation, from /root/tensorflow_datasets/ogbg_molpcba/0.1.2\\n\",\n \"INFO:absl:Constructing tf.data.Dataset ogbg_molpcba for split test, from /root/tensorflow_datasets/ogbg_molpcba/0.1.2\\n\",\n \"INFO:absl:Initializing network.\\n\",\n \"INFO:absl:\\n\",\n \"+-----------------------------+------------+--------+-----------+--------+\\n\",\n \"| Name | Shape | Size | Mean | Std |\\n\",\n \"+-----------------------------+------------+--------+-----------+--------+\\n\",\n \"| params/Dense_0/bias | (256,) | 256 | 0.0 | 0.0 |\\n\",\n \"| params/Dense_0/kernel | (9, 256) | 2,304 | -0.000941 | 0.335 |\\n\",\n \"| params/Dense_1/bias | (128,) | 128 | 0.0 | 0.0 |\\n\",\n \"| params/Dense_1/kernel | (256, 128) | 32,768 | -0.000737 | 0.0623 |\\n\",\n \"| params/LayerNorm_0/bias | (256,) | 256 | 0.0 | 0.0 |\\n\",\n \"| params/LayerNorm_0/scale | (256,) | 256 | 1.0 | 0.0 |\\n\",\n \"| params/LayerNorm_1/bias | (256,) | 256 | 0.0 | 0.0 |\\n\",\n \"| params/LayerNorm_1/scale | (256,) | 256 | 1.0 | 0.0 |\\n\",\n \"| params/LayerNorm_2/bias | (256,) | 256 | 0.0 | 0.0 |\\n\",\n \"| params/LayerNorm_2/scale | (256,) | 256 | 1.0 | 0.0 |\\n\",\n \"| params/LayerNorm_3/bias | (256,) | 256 | 0.0 | 0.0 |\\n\",\n \"| params/LayerNorm_3/scale | (256,) | 256 | 1.0 | 0.0 |\\n\",\n \"| params/LayerNorm_4/bias | (256,) | 256 | 0.0 | 0.0 |\\n\",\n \"| params/LayerNorm_4/scale | (256,) | 256 | 1.0 | 0.0 |\\n\",\n \"| params/MLP_0/Dense_0/bias | (256,) | 256 | 0.0 | 0.0 |\\n\",\n \"| params/MLP_0/Dense_0/kernel | (256, 256) | 65,536 | -0.000616 | 0.0625 |\\n\",\n \"| params/MLP_0/Dense_1/bias | (256,) | 256 | 0.0 | 0.0 |\\n\",\n \"| params/MLP_0/Dense_1/kernel | (256, 256) | 65,536 | -0.000287 | 0.0626 |\\n\",\n \"| params/MLP_1/Dense_0/bias | (256,) | 256 | 0.0 | 0.0 |\\n\",\n \"| params/MLP_1/Dense_0/kernel | (256, 256) | 65,536 | 0.000217 | 0.0628 |\\n\",\n \"| params/MLP_1/Dense_1/bias | (256,) | 256 | 0.0 | 0.0 |\\n\",\n \"| params/MLP_1/Dense_1/kernel | (256, 256) | 65,536 | -0.00055 | 0.0625 |\\n\",\n \"| params/MLP_2/Dense_0/bias | (256,) | 256 | 0.0 | 0.0 |\\n\",\n \"| params/MLP_2/Dense_0/kernel | (256, 256) | 65,536 | 0.000274 | 0.0625 |\\n\",\n \"| params/MLP_2/Dense_1/bias | (256,) | 256 | 0.0 | 0.0 |\\n\",\n \"| params/MLP_2/Dense_1/kernel | (256, 256) | 65,536 | -0.000227 | 0.0626 |\\n\",\n \"| params/MLP_3/Dense_0/bias | (256,) | 256 | 0.0 | 0.0 |\\n\",\n \"| params/MLP_3/Dense_0/kernel | (256, 256) | 65,536 | -0.000224 | 0.0625 |\\n\",\n \"| params/MLP_3/Dense_1/bias | (256,) | 256 | 0.0 | 0.0 |\\n\",\n \"| params/MLP_3/Dense_1/kernel | (256, 256) | 65,536 | 1.42e-05 | 0.0627 |\\n\",\n \"| params/MLP_4/Dense_0/bias | (256,) | 256 | 0.0 | 0.0 |\\n\",\n \"| params/MLP_4/Dense_0/kernel | (256, 256) | 65,536 | 0.000319 | 0.0623 |\\n\",\n \"| params/MLP_4/Dense_1/bias | (256,) | 256 | 0.0 | 0.0 |\\n\",\n \"| params/MLP_4/Dense_1/kernel | (256, 256) | 65,536 | -0.000191 | 0.0624 |\\n\",\n \"+-----------------------------+------------+--------+-----------+--------+\\n\",\n \"Total: 695,936\\n\",\n \"INFO:absl:Checkpoint.restore_or_initialize() ...\\n\",\n \"INFO:absl:No checkpoint specified. Restore the latest checkpoint.\\n\",\n \"INFO:absl:Checkpoint None does not exist.\\n\",\n \"INFO:absl:Checkpoint.save() ...\\n\",\n \"INFO:absl:Checkpoint.save() finished after 0.09s.\\n\",\n \"INFO:absl:Checkpoint.restore_or_initialize() finished after 0.10s.\\n\",\n \"INFO:absl:Starting training.\\n\",\n \"INFO:absl:Finished training step 1.\\n\",\n \"INFO:absl:Finished training step 2.\\n\",\n \"INFO:absl:Finished training step 3.\\n\",\n \"INFO:absl:Finished training step 4.\\n\",\n \"INFO:absl:Finished training step 5.\\n\",\n \"INFO:absl:Finished training step 6.\\n\",\n \"INFO:absl:Finished training step 7.\\n\",\n \"INFO:absl:Finished training step 8.\\n\",\n \"INFO:absl:Finished training step 9.\\n\",\n \"INFO:absl:Finished training step 10.\\n\",\n \"INFO:absl:Created artifact [10] Profile of type ArtifactType.URL and value None.\\n\",\n \"INFO:absl:Setting work unit notes: 2.0 steps/s, 12.4% (124/1000), ETA: 7m\\n\",\n \"INFO:absl:[124] steps_per_sec=2.048078\\n\",\n \"INFO:absl:Setting work unit notes: 2.1 steps/s, 25.2% (252/1000), ETA: 5m\\n\",\n \"INFO:absl:[252] steps_per_sec=2.127820\\n\",\n \"INFO:absl:Setting work unit notes: 2.1 steps/s, 38.0% (380/1000), ETA: 4m\\n\",\n \"INFO:absl:[380] steps_per_sec=2.116415\\n\",\n \"INFO:absl:[500] train_accuracy=0.9849439859390259, train_loss=0.06188433617353439\\n\",\n \"INFO:absl:[500] validation_accuracy=0.9839196801185608, validation_loss=0.06475622951984406, validation_mean_average_precision=0.036894\\n\",\n \"INFO:absl:[500] test_accuracy=0.9832062125205994, test_loss=0.06711383163928986, test_mean_average_precision=0.037820\\n\",\n \"INFO:absl:Setting work unit notes: 0.4 steps/s, 50.1% (501/1000), ETA: 23m (8m : 31.9% eval)\\n\",\n \"INFO:absl:[501] steps_per_sec=0.351547\\n\",\n \"INFO:absl:Setting work unit notes: 2.1 steps/s, 62.8% (628/1000), ETA: 2m (9m : 28.7% eval)\\n\",\n \"INFO:absl:[628] steps_per_sec=2.109877\\n\",\n \"INFO:absl:Setting work unit notes: 2.1 steps/s, 75.6% (756/1000), ETA: 1m (10m : 26.1% eval)\\n\",\n \"INFO:absl:[756] steps_per_sec=2.128443\\n\",\n \"INFO:absl:Setting work unit notes: 2.1 steps/s, 88.4% (884/1000), ETA: 0m (11m : 23.9% eval)\\n\",\n \"INFO:absl:[884] steps_per_sec=2.122990\\n\",\n \"INFO:absl:[999] train_accuracy=0.9868282079696655, train_loss=0.052991833537817\\n\",\n \"INFO:absl:[999] validation_accuracy=0.984052836894989, validation_loss=0.06297025829553604, validation_mean_average_precision=0.049514\\n\",\n \"INFO:absl:Checkpoint.save() ...\\n\",\n \"INFO:absl:[999] test_accuracy=0.9833315014839172, test_loss=0.06533924490213394, test_mean_average_precision=0.050113\\n\",\n \"INFO:absl:Checkpoint.save() finished after 0.06s.\\n\",\n \"INFO:absl:Setting work unit notes: 1.7 steps/s, 100.0% (1000/1000), ETA: 0m (13m : 0.0% checkpoint, 22.5% eval)\\n\",\n \"INFO:absl:[1000] steps_per_sec=1.697496\\n\",\n \"INFO:absl:[1000] train_accuracy=0.9897193908691406, train_loss=0.04319273680448532\\n\",\n \"INFO:absl:[1000] validation_accuracy=0.9840479493141174, validation_loss=0.0629250779747963, validation_mean_average_precision=0.049538\\n\",\n \"INFO:absl:[1000] test_accuracy=0.9833057522773743, test_loss=0.06529100239276886, test_mean_average_precision=0.050135\\n\"\n ]\n }\n ],\n \"source\": [\n \"# Training loop\\n\",\n \"\\n\",\n \"# Use a Colab GPU runtime to speed up training.\\n\",\n \"# We don't use TPUs in this Colab because we do not distribute our\\n\",\n \"# training using pmap() - if you're looking for an example using TPUs\\n\",\n \"# checkout the below Colab notebook:\\n\",\n \"# https://colab.research.google.com/github/google/flax/blob/main/examples/imagenet/imagenet.ipynb\\n\",\n \"\\n\",\n \"config.num_train_steps = 1000\\n\",\n \"config.log_every_steps = 500\\n\",\n \"config.eval_every_steps = 500\\n\",\n \"\\n\",\n \"# Construct the model and start the main training loop.\\n\",\n \"# The default config corresponds to a 5-layer Graph Convolutional Model\\n\",\n \"# with skip-connections and mean-pooling.\\n\",\n \"state = train.train_and_evaluate(config, workdir=f'./models')\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 11,\n \"metadata\": {\n \"cellView\": \"form\"\n },\n \"outputs\": [],\n \"source\": [\n \"#@title Upload to TensorBoard.dev\\n\",\n \"if 'google.colab' in str(get_ipython()):\\n\",\n \" #@markdown You can upload the training results directly to [TensorBoard.dev](https://tensorboard.dev).\\n\",\n \" #@markdown\\n\",\n \" #@markdown Note that anyone with the link will be able to see the data.\\n\",\n \" upload_data = 'no' #@param ['yes', 'no']\\n\",\n \" if upload_data == 'yes':\\n\",\n \" !tensorboard dev upload --one_shot --logdir ./models --name 'Flax examples/ogbg_molpcba'\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Inference\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 12,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# Create deterministic evaluation model.\\n\",\n \"eval_net = train.create_model(config, deterministic=True)\\n\",\n \"eval_state = state.replace(apply_fn=eval_net.apply)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 13,\n \"metadata\": {\n \"outputId\": \"51242dc2-5260-4279-b41a-9d7614b33c97\"\n },\n \"outputs\": [\n {\n \"name\": \"stderr\",\n \"output_type\": \"stream\",\n \"text\": [\n \"INFO:absl:Load dataset info from /root/tensorflow_datasets/ogbg_molpcba/0.1.2\\n\",\n \"INFO:absl:Reusing dataset ogbg_molpcba (/root/tensorflow_datasets/ogbg_molpcba/0.1.2)\\n\",\n \"INFO:absl:Constructing tf.data.Dataset ogbg_molpcba for split train, from /root/tensorflow_datasets/ogbg_molpcba/0.1.2\\n\",\n \"INFO:absl:Constructing tf.data.Dataset ogbg_molpcba for split validation, from /root/tensorflow_datasets/ogbg_molpcba/0.1.2\\n\",\n \"INFO:absl:Constructing tf.data.Dataset ogbg_molpcba for split test, from /root/tensorflow_datasets/ogbg_molpcba/0.1.2\\n\"\n ]\n }\n ],\n \"source\": [\n \"# Compute accuracy and mean average precision on validation and test sets.\\n\",\n \"datasets = input_pipeline.get_datasets(\\n\",\n \" config.batch_size,\\n\",\n \" add_virtual_node=config.add_virtual_node,\\n\",\n \" add_undirected_edges=config.add_undirected_edges,\\n\",\n \" add_self_loops=config.add_self_loops)\\n\",\n \"eval_metrics = train.evaluate_model(eval_state, datasets,\\n\",\n \" splits=['validation', 'test'])\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 14,\n \"metadata\": {\n \"outputId\": \"2a529c9f-fed6-4221-a738-a8ffd9049d7e\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"validation\\n\",\n \"- accuracy: 0.984\\n\",\n \"- loss: 0.063\\n\",\n \"- mean_average_precision: 0.050\\n\",\n \"test\\n\",\n \"- accuracy: 0.983\\n\",\n \"- loss: 0.065\\n\",\n \"- mean_average_precision: 0.050\\n\"\n ]\n }\n ],\n \"source\": [\n \"for split in ['validation', 'test']:\\n\",\n \" split_metrics = eval_metrics[split].compute()\\n\",\n \" print(split)\\n\",\n \" for metric_name, metric in split_metrics.items():\\n\",\n \" print(f'- {metric_name}: {metric:.3f}')\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 15,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# Helper functions for formatting labels and predictions.\\n\",\n \"def get_formatted_label_for_task(labels, task):\\n\",\n \" label_for_task = labels[task]\\n\",\n \" if np.isnan(label_for_task):\\n\",\n \" return 'Unknown'\\n\",\n \" elif label_for_task == 0:\\n\",\n \" return 'Inactive'\\n\",\n \" elif label_for_task == 1:\\n\",\n \" return 'Active'\\n\",\n \" raise ValueError('Invalid label.')\\n\",\n \"\\n\",\n \"# Predictions are computed with a threshold of 0 for the logits.\\n\",\n \"# This is the same threshold used to compute the accuracy.\\n\",\n \"def get_formatted_prediction_for_task(logits, task):\\n\",\n \" predictions = logits > 0\\n\",\n \" prediction_for_task = predictions[task]\\n\",\n \" if prediction_for_task == 0:\\n\",\n \" return 'Inactive'\\n\",\n \" elif prediction_for_task == 1:\\n\",\n \" return 'Active'\\n\",\n \" raise ValueError('Invalid prediction.')\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"We can choose one of the 128 different tasks and see how the model predictions\\n\",\n \"match up with the true labels.\\n\",\n \"\\n\",\n \"You can change this visualization to use any of the 128 tasks in this dataset, by changing the `task` variable below. See the appendix [here](https://arxiv.org/pdf/1502.02072.pdf) to understand what these tasks mean.\\n\",\n \"\\n\",\n \"The default is task PCBA-686978 (indexed at 93).\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 16,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# Define which task to plot labels for.\\n\",\n \"task = 93\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 17,\n \"metadata\": {\n \"outputId\": \"38934096-13a6-4701-823a-ac83f3b7eaac\"\n },\n \"outputs\": [\n {\n \"name\": \"stderr\",\n \"output_type\": \"stream\",\n \"text\": [\n \"INFO:absl:Load dataset info from /root/tensorflow_datasets/ogbg_molpcba/0.1.2\\n\",\n \"INFO:absl:Reusing dataset ogbg_molpcba (/root/tensorflow_datasets/ogbg_molpcba/0.1.2)\\n\",\n \"INFO:absl:Constructing tf.data.Dataset ogbg_molpcba for split test, from /root/tensorflow_datasets/ogbg_molpcba/0.1.2\\n\"\n ]\n },\n {\n \"data\": {\n \"image/png\": \"iVBORw0KGgoAAAANSUhEUgAAA1MAAArLCAYAAAAvFZZhAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjIsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+WH4yJAAAgAElEQVR4nOzdd5hU1fnA8e+5U3a2F3pZYCmCgIgUKaKALUbsBSuKFVsSU9TE2JLoLybRxESjWImKBSuWaOygSJPepC11WWB32b4zszNz5/z+mFmY3ZnZnVm2wvt5nnmYvfXc2eXMfe855z1Ka40QQgghhBBCiPgYrV0AIYQQQgghhGiPJJgSQgghhBBCiEaQYEoIIYQQQgghGkGCKSGEEEIIIYRoBAmmhBBCCCGEEKIRJJgSQgghhBBCiEaQYKqdU0pNUkrltfS+jTzfPKXUjS11PiFE65G6SQjRHrSnuqoxlFK9lFKVSilLa5flSCXBVB3BP7ial18p5Qr5+apmPO90pdSC5jr+4VJK9VFKaaWUtc7y/yilHm6tcglxtJC6KTKpm4RoW6Suii5YV/Vv5nPsUEqdXvOz1nqX1jpFa20253mPZtaGNzm6aK1Tat4rpXYAN2qtv6y7nVLKqrX2tWTZhBBHL6mbhBDtgdRV4mgjLVMxqmnKVUrdo5TaB8yK9BQk9KmDUipBKfWYUmqXUmq/UmqmUiqxEee+Tin1o1KqQim1TSk1I8I29yqlioJPJK4KWd4kZYixnNOVUguC5ytRSm1XSv00yrbdlFJrlFJ3BX+ep5T6k1Lq++B1fq6U6hiy/XlKqfVKqdLgtscGl1+nlPooZLstSqm3Q37erZQaHnyvlVK3BLcpVUr9WymlmuOzEKKlSN0UUzmlbhKilUldFXa+h5RSbymlXgmWa71SalTI+t8qpXKD6zYopS6ss/9NIde0QSk1Qin1KtAL+EgFWgLvViGt90qpy5RSy+oc55dKqQ+b81qPdBJMxacrkAX0Bm6OYftHgWOA4UB/oAfwQCPOWwCcA6QB1wH/UEqNqFOujsHjXws8p5QaGG8ZlFJPK6WebkT5Qo0BNgXL81fgxbo3BUqpHGA+8JTW+m8hq64kcH2dATvwm+D2xwBvAHcCnYBPCFQU9uBxTlZKGUqp7sH9xgX36wukAGtCznEOMBoYBkwFfnKY1ytEWyB1U8OkbhKi9UldVdt5wJtABvAh8FTIulzgZCAd+AMwWynVLXieS4GHgGuC13QecEBrPQ3YBZwb7Nr31zrn+wgYqJQaELLsSuD1eK9VhNBayyvKC9gBnB58PwnwAI6Q9dOBBXX20QT+ABVQBfQLWTcO2B7lXGHHqqdcc4FfhJTLBySHrH8LuL+hMgT3zYvxnH2C12ats/w/wMMh17A1ZF1ScJ+uwZ/nAX8Pfq5X1DnOPOC+kJ9vA/4XfH8/8FbIOgPYA0wK/rwbGAFcDjwHLAUGEagwP6zzu5lQ53P6bWv/nclLXvG+pG6qdU6pm+Qlrzb6kroq7Lwa6B98/xDwZci6wYCrnn1XAecH339WU/76PvPgz7XqSGA28EDw/QCgIlgnxvV5y+vQS8ZMxadQa+2OcdtOBP44l4c8/FRA3NlUgt1RHiTwtMAIHndtyCYlWuuqkJ93At2bsgwEKhoAW8j7mp+9IT/vq3mjtXYGz5sSsv4qYCvwToRz7At57wzZrzuBa6o5rl8ptZvAExMIPAGeRKDynQ+UAhMJVALzYzyHEO2Z1E1SNwnRHhzNdVUkdf/fO1RwLJlS6hrgVwSCIQjUCTVdjLMJtFw1xuvA48AfCbRKzQ3WiZ1p3ms9Ykk3v/joOj9XEfjDA0Ap1TVkXRHgAoZorTOCr3QdMjAzFkqpBOBd4DGgi9Y6g0BXktDuKZlKqeSQn3sB+U1VhqC9BG5M+tRZnkPIzUQMHgqW63UVe5rOfAJdAgAIds3JJvAEGA7dsJwcfD+fwA3LRMJvWIQ4EkndJHWTEO3B0VxXxVPm3sDzwB1Ah2CZ14WUeTfQL8rudT/jur4AOqnAmM0rONTFr1Wu9UggwdThWQ0MUUoNV0o5CHwZA4EnlAT+I/wjGO2jlOqhlKqvH7xSSjlCXwT62ScAhYAv+HTlzAj7/kEpZVdKnUygX/DbjSxDRDqQUvNd4BGlVAellE0pdQWBZulP4ziUF7gUSAZeUUrF8jf4FjBFKXWaUsoG/BqoBhYG188HJgOJWus84DvgLKADsDKOsglxpJC6SeomIdqDo6auilMygaCoMHjO64ChIetfAH6jlBqpAvoHAzCA/UDfaAfWWnuBt4G/ERi/9kVweWtda7snwdRh0FpvJtBM+iWwBag7v8E9BLqNLFZKlQe3G0h04wk8Faj7+jmBL+0SAk2yH9bZb19wXT7wGnCL1npjvGVQgawtM+sp321AMYFB0wUEnphM0Vrvr2efMFprD3AR0AV4qaGbFq31JuBq4EkCT07OJTC40hNcvxmoJHCjgta6HNgGfK9lXgVxFJK6SeomIdqDo7CuionWegOBrniLCARHxwHfh6x/G3iEQKtSBYExYFnB1X8G7lOBzKC/iXKK14HTCQSMod2j4/28BaC0bqg1UAghhBBCCCFEXdIyJYQQQgghhBCNIMGUEEIIIYQQQjSCBFNCCCGEEEII0QgSTAkhhBBCCCFEI0gw1QyUUv9RSj0cfH+yUmpTC51XK6X6N/ExD15LS+57OJRSvZRSlXHMFSPEEUnqosPft6Uope5VSr3Q2uUQoiVI3XT4+7YUqZsadtQGU0qpHUopV/Cme3/wD7rJJybTWn+ntW4wraRSarpSqm5K0CajlJqnlLqxuY7fVJRSk4KV3T1x7LNDKXV6zc9a611a6xRJPSzaA6mL2pbmvv7gOSYppfJCl2mt/09r3WY/F3H0kbqpbZG6qe06aoOpoHODMzuPAEYB99XdQCllbfFSHd2uJTBfzDWtXRAhWpDURUKItkjqJiEacLQHUwBorfcAnxKcXTrYMnK7UmoLgUnkUEqdo5RaFZwEbaFSaljN/kqpE5RSK5RSFUqpOYAjZF2tKF8pla2Uek8pVaiUOqCUekopdSwwExgXfAJUGtw2QSn1mFJqV/Cp0EylVGLIse5SSu1VSuUrpa5v7PUrpd5WSu1TSpUppb5VSg2ps0lHpdQXweubrw7Nso1SalBwXbFSapNSauphlCMZuAS4HRiglBpVZ/1NSqkfg+XYoJQaoZR6FegFfBT87O5WSvUJ/g6tSqnLlFLL6hznl0qpD4Pv6/2MhWhJUhe1jbqoTpl2KKV+o5RaEyzXHKWUI7guUyn1cfAzLAm+7xmyb5ZSalbwcylRSs0N1nOfAt2Dn3GlUqq7UuohpdTs4H6fKqXuqFOO1Uqpi5rzWoWIRuomqZuC+0ndFIEEUwT+4wJnAytDFl8AjAEGK6VOAF4CZgAdgGeBD4P/ie0EZp5+lcDs028DF0c5jwX4GNgJ9AF6AG9qrX8EbgEWBbunZQR3eRQ4BhgO9A9u/0DwWGcBvwHOAAYQmMm6sT4NHqMzsILA7N+hrgL+BHQEVtWsD/7H+4LATNqdgcuBp5VSg6Ncf6lSakI95bgIqCTwGX5GoJWqZt9LgYcItFilAecBB7TW04BdBJ+eaa3/WueYHwEDlVIDQpZdGSwz1PMZC9HSpC5qM3VRXVOBs4AcYBgwPbjcAGYBvQk81HEBT4Xs9yqQBAwJlusfWusq4KdAfvAzTtFa59c53xvAFSHlHRw8x3/jvVYhmoLUTVI3BUndFInW+qh8ATsI3LiXEvhP+zSQGFyngVNDtn0G+FOd/TcBE4FTgHxAhaxbCDwcfD8JyAu+HwcUAtYI5ZkOLAj5WQFVQL+QZeOA7cH3LwGPhqw7Jlju/lGudx5wYwyfS0bwOOnBn/9DoCKrWZ8CmEA2cBnwXZ39nwUeDNn34Th+J18CTwTfXxH8rGzBnz8DflHP7/L0kJ/7BK/BGvx5NvBA8P0AoIJAJVLvZywvebXES+qiqJ9Lq9RFEa5/B3B1yM9/BWZG2Xc4UBJ83w3wA5kRtjv4uwhZ9hAwO/g+NfiZ9w7+/AjwUvB9vdcqL3k11Uvqpqifi9RNUjfVeh3t/Vwv0Fp/GWXd7pD3vYFrlVI/C1lmB7oT+A+1Rwf/aoJ2RjlmNrBTa+2LoWydCNzwL1dK1SxTQE2Guu7A8hjOWa/gU6BHgEuD5/QHV3UEyoLvD34WWutKpVRx8Py9gTE1ze1BVgJPPOItRzYwGfhdcNEHwHPAFAJPtLKB3HiPG/Q68DjwRwKtUnO11k6lVGfq/4yFaClSF7WRuiiKfSHvncFzopRKAv5B4MlwZnB9avBasoFirXVJvCfTWlcopf5L4MnuXwg8XLopuLq5r1WIUFI3Sd10kNRNkR3twVR9Qv/T7wYe0Vo/UncjpdREoIdSSoVUFL2IfOO/G+illLJGqCh0nZ+LCDTLDtGBvsp17SXwH6JGr+iXUq8rgfMJNH/vANKBEgIVUo2D51GBTD5ZBJ4y7Qbma63PaOS5Q00j0Cz9UUil6CDQ1W9u8Fz9ouxb97Or6wugk1JqOIH/+L8MLm/oMxaiLZC66JCWqIvi8WtgIDBGa70vWMesJFDm3UCWUipDa11aZ7+G6iwIdKd5UCn1LYG68Jvg8ta6ViHqkrrpEKmbjuK6ScZMxeZ54Bal1BgVkKyUmqKUSgUWAT7g50opW3AQ3olRjrOUwH/uR4PHcCilTgqu2w/0DPYtRmvtD573H8EWFJRSPZRSPwlu/xYwXSk1OPgE4sEYrsMaPGfNy0agybYaOEDgCc//RdjvbKXUhGDZ/gQs1lrvJtCv+Ril1LTgtduUUqNVYKBovK4F/kCgKbrmdXHw3B2AF4DfKKVGBn8H/dWhAZ77gb7RDqy19hLoo/03AhXcF8HlDX3GQrQ1Uhc1f10Uj1QCN3OlSqksQq5da72XwDiLp1VgMLhNKXVKcPV+oINSKr2eY39C4EnvH4E5wd8DtN61ClEfqZukbjpq6yYJpmKgtV5GoBnzKQJPI7YSHOSntfYQSJwwnUBK78uA96IcxwTOJTBIcheQF9we4GtgPbBPKVUUXHZP8FyLlVLlBMYUDQwe61PgieB+W4P/NuQZAv+5al6zgFcINH3vATYAiyPs9zqB/4jFwEjg6mAZKoAzCTT35hNobv4LkBDp5CqQGebkCMvHEviP+W+t9b6Q14fBa7tCa/02gWb21wmMeZpLIDAC+DNwnwoM3PxNlGt/ncBTpbfrPO2K+hkL0dZIXdS8dVEjPAEkEnhCvhj4X5310wAvsBEoAO4Mlncjgae724L1Vve6B9ZaVxP4/Z3OoYQ5cV+rEC1B6iapm47muknV7sIqhBBCCCGEECIW0jIlhBBCCCGEEI0gwZQQQgghhBBCNIIEU0IIIYQQQgjRCBJMCSGEEEIIIUQjSDAlhBBCCCGEEI1Q76S9HTt21H369GmhogghWsLy5cuLtNadWrsch0vqJyGOPEdC/SR1kxBHnvrqpnqDqT59+rBs2bLmKZUQolUopXa2dhmagtRPQhx5joT6SeomIY489dVN0s1PCCGEEEIIIRpBgikhhBBCCCGEaAQJpoQQQgghhBCiESSYEkIIIYQQQohGkGBKCCGEEEIIIRpBgikhhBBCCCGEaAQJpoQQQgghhBCiESSYEkIIIYQQQohGqHfSXtG+bCnZwq6KXTi9TpJtyfRJ70Pf9L6tXSwhhBCtZE+pix/zy6mo9uKwWuiWkcjxPdNRSrV20YQQosmUVHlYtbuUcrcXq2HQMcXOyN6ZWC3N324kwVQ7V21W8/mOz3lx7YvsqdyD1bDi134MZeDz++iT3ocbht7Aab1Ow2axtXZxhRBCNDO/X/PtlkKenb+NFbtKsFsN/FqjUGitSUu0cfPJfbl4VE/SHPK9IIRon7TWrNpdyvPfbeOrHwtC6joAhcVQXDOuN9PG9qZzmqPZyqG01lFXjho1Si9btqzZTi4Oz8bijdz0+U14TA9OnzPqdknWJJJtybzwkxekpUqglFqutR7V2uU4XFI/CRGusKKaq19YQl6JkyqPGXW7RJsFpeDpq0YwaWDnFixh/Y6E+knqJiGan8tjcuvs5SzdUYzba+KPEs4kWAMtU/dNOZZp4/o0+nz11U0yZqqdWlO4hms+vYbS6tJ6AykAp89JkauIK/97JZtLNrdQCYUQQrSkgnI3U/71HbmFlfUGUgAur4nTY3LL7OV8sja/hUoohBCHz+01uWTmQhZtO4DTEz2QAqj2+an2+fm/Tzbyr6+2NEt5JJhqh/ZV7WPGFzNw+Vwx76PRVHmruOGzGyh2Fzdj6YQQQrQ0j8/P5c8vprjKg6++O4s63F4/v3prNat3lzZj6YQQounc9toKthZUUu3zx7yPy2vy9LytfLSq6R8eyZipduiltS/h9rnDlpd8V0LRZ0V4CjxYHBbSRqbR5ZIuWJItB7dxep28/uPr3HHCHS1ZZCGEEM3o03V72VfmjhhIVW2YR/kPc/EeyMOwJ2Lr3Jf08VNx9BwCBAKqRz/dyBs3j23pYgshRFzW7SljUW5RxEAqlrruT//dwJRh3TCMpkvCI8FUO+PyuZibOxef9tVaXvRpEYWfFtLzxp6kDE7BW+Il/9V8djy2g5zf52AE+4x6/B7e2PgGM46fgc2QgcdCCHEkmDk/F2eErn3lS9+nbMk7dDjzdhw5I1AWK67ty3FtWXLwBgNgxa4Sdhc7yc5KasliCyFEXF5csB2PGf7QKNa6rqrax/e5RZw8oFOTlUm6+bUz/9v+PxS1o2nTZVIwt4DuV3cndVgqyqqwd7KTfVs2niIPZQvLam3v8/uYv3t+SxZbCCFEM9m4r5ztRVVhy/3VVZQueI2sM24laeB4DLsDZbGS1H8MmZOvr72t1ry6aGdLFVkIIeJW6vTw37V7Meu0wMdT11V5TGbOz23Sckkw1c4s3rs4LOGEc4sTv9dP2si0WsstDgupw1KpXF9Ze3ufkx/2/dDsZRVCCNH8lu8sibi8es9GtM9D0jHjGjyG19R8u6WADfnlfLelkG82FbBiVwlub/2JLIQQorkVV3l4Zt5WJv7tGzwRuvfFU9cBrNrVtGNEpZtfO1NaHf4HYFaaWFOsKEt4/09ruhXXzvBEFZKEQgghjgzlLh9eM/wGw3SVYySloQxLhL3Cbd5fySUzF2IJjiXQOtBidcnInlx3Ug45HZObtNxCCFEfrTV/+2wTLy7YjlKBMU+RxFvXuZr4IZEEU+2M3bCHLbOkWPBV+tCmDguofGU+rCnhv+YES0KzlVEIIUTLsVsNDKUwqd31xZKYht9ZjvabMd1k+DURx129sWQXc37YzXnHd+fPFx2H1SKdWoQQzcvv1/z8zZV89WNBg1n74q3rrEbT1mFSI7Yz3VK6Yajav7ak/kkoq6J8eXmt5abbpGJNBcmDaz9NtCor3ZK7NXtZhRBCNL+OKXbsEQKchB6DUFYbzs2LDuv4Xr+m2ufn4zV7mT7rB3wRWsGEEKIpPfLJj3z1Y0FMrUjx1nVpiU3bliQtU+3M+f3O5/0t7+M2D6VGtyRZ6HxBZ/Jn52M4jFrZ/GxZNjLGZ9Q6hsWwcHbfs1u66EIIIZrBqYM6hw3IBjASksmYcBXFX8xEGRYcOSegDCvuHatw71oTNjC7IS6vyfKdJfz+/XX85ZJhTVV8IYSoZXexk9mLd8aV/jzWus5uNbh0ZHaTlleCqXZmSMchdEvuxvby7bWWdzq7E5ZkC/vm7MNT4MFINEgbkUb2jGwMW+0nlgMzB5KTntOSxRZCCNFMUh02zh3enfeW51E3Y3DaiRdhJGdStmgORR8/hrInktClP2njLot6vPrmanF5Teau2sPtk/vTq4OkURdCNL1XFu3Ar+NLf545+fqY6joFXDO+d5OWV4KpduiG427g4SUPh03cmzUxi6yJWfXum2hN5Pqh8T2NFEII0bbdOKEvH63Ox4wwQDtlyGRShkyO6TixzNXi15pZC7fz4LlDGjiaEELEp9pn8sbS3XjNyOnPO5x9J0kDxx9cntR/DEn9xwAN13UWBSfmZNEtPbFJyyxjptqhc/udy+guo+NOIuGwOJjYcyKn9jq1mUomhBCiNQzsmsotE/uRaIstm1Uksc7V4jU1c37YLWnThRBNbt2e8ojL401/XpcC0hPtPHbp8YdRusgkmGqHDGXw90l/5/hOx+OwOGLaJ9GSyNhuY/m/k/8PpcJTqAshhGjffnHaAC7p6ifB9DZq/3huVgyl2Ly/olHnEUKIaMpcHiLdpsab/jyUxYDMZDtv3TKWLmmx3TfHQ7r5tVMOq4Nnz3iWZ1Y/w2s/vobWOmwyX4AkaxIWZWH60OnceNyNYZkAhRBCtH1r88p4ZdEONu+vwOkxSU6wMqR7GteO78MxXVIB8OTmMm32wxx7z+M8sb6KqmofVRFSnUcTz82KUlDmalzQJoQQ0UR74B9v+nOABGvgnvfkAR155MLjmiWQAgmm2jWrYeVnJ/yMGcNm8MXOL3h5/ctsLNqJzWqSaHXQK60X04dMZ3KvydgMW2sXVwghRJw+Wp3PP7/czJ5SN9U+k9CkfWvzSnl3RR79O6Vw54Rscn53B53vuosrLxjH5edpvs8t4rlvt7F+TzlOrw+7xaBjip3tRU7Ch3a3/lwtQgiRmWRHR0g+EZr+PHnQhAaPo4DbJvXjijG96JzaPEFUDQmmjgB2i50pfacwpe8URj38BZ/8/GQ6N1P0LYQQovn5/Zr7P1jHeyv2RJ1nxdRgev2syy/njjmruXL0xdx/4QUAGIbi5AGdOHlAp7D9hj74GZXVvrDl8dys+ExNx5TwSeSFEOJwDO2eFnxQU7vei3eqh0tG9uQXpx/TImWWYOoIYvo1pU4vWcnyBSeEEO3ZQx+t570VebgiZOeLxK2svKF6kvn1Vn522oB6tz1nWDfeWZ6Hr87cVPHcrGQl2+nfOSX+CxNCiHpYLQbXjO/NU19vpe70eTVTPRR//QJmeQFoDRYb9s59ceetP5hxFIj4IKnZytxiZxLNrrjKQ1qiDatFul4IIUR79eWG/by9LHIgVf8cUH7+PW8r4/t3YGTv6NNk3DAhh7mr9oQFUxDbvFSJNgszTukryYyEEM3i3OO786+vtkZc568qRXvddDr/d1GnbwD4YsM+zhvevUXKK8FUO6a1Zsn2Yl5csJ3N+yood3uprPYxfdZSbpiQw0n9OmIY8mUnhBDtyZNfb4nYtS+WOaCqvX6enpfLi9dGD6YGdEllQOdU1uWXEWFoQoNztWg0F47oEf+FCSFEDL7+sQCrocIe+MQy11SNzzfsp8zpJT2p+XMGSDDVDmmteXtZHk98tZlSpxeXx6w1mHjepkKWbi8mJcHKHZP7M21cb3mCKIQQ7UBuYSWb9oWnHI/1JkID320uorCimk6p0ecifOLy4Zz31AKqquObK8phM/jrxcNIdUhSIyFE83h/ZeSW83imb7BaFN9uKeTc45u/dUr6g7Uzfr/md++t5cEP15Nf6sZZJ5Cq4fSYFFRU8+dPN3LHGyvxmrH1uxdCCNF65izdfdg3EUrB3JV59W7Tr1MKr94whmRDo3Rs3w8Om8FvfzqI84ZLq5QQovmUOiNPuxDP9A0+U1Pi9DR10SKSYKqdeeij9XywKj9qdqe6XF6Tr37cz11vr46YalIIIUTbkVtYGTGYiucmotrnZ8cBJ6bfpMRdQl5FHsXuYkx/7e+NIb4Snlw8k2Gdk0iwGlgjdAtXQJLdQrd0B/++cgTTx+c0+tqEECIWOmIzQe3pG2Lhj1CXNgfp5teOfLu5MDgoOfyPqL5ByW6vn8837OeTtfuYMqxbK5RcCCFELJxRJtmNZw4oZStmZeV8TnrzW7ymF4thwfSbWAwLFw24iKsGXUWPhM7s+eWvOOGWa/hg6mlsK6xk1vc7+HB1PpVuH340iTYLJ/bJYsbEfoztmyXdxYUQLSLNYWN/eXXY8nimb7AaiswWym59xAdTFW4vH63ey9aCCspcXjKS7BzTJYVzhnUnOaF9Xf4z83IbPSjZ6TF5et5WCaaEEKINS0+MPBYpppsIoxpH9zewJm8lzw/aH5xLqqYXnx/mbJzDO5vfYUhlBr8dMIiMSy8FoG+nFP50wVD+dMHQg70YJHgSQrSGs4Z2Zee32/D4andBjmf6Bq9fM65vhxYpb/uKJuKweX8Fz327jY/X5KOgVorZJLuFhz7cwPnDu3PTKX3p16ntz5Wxu9jJil0lYcvjyWxSM7B5YNfUZi+vEEKI+A3vlcG8zQW4vfHdRGSdfhmJfZ7GsJahDF+UTjLg0z58po81tv38YozidfcBOiZ2rLWNBFFCiNZ09djePD0vN+K6WKZvUMCEfh3pnOZokfIekcHUe8vzuHfuWrymxozQX7KmG8Xby3fzwap8Hrt0GFOGtUwu+sb6YNUe/BHGPMUzKNnr8/PWsl3cf86QBrcVQgjR8qaOyuYfX2yOuC7qTcT4i0ns9QKGtRRlxDaWwGeFQlcR1//vet48502SbElNeRlCCNFoH6/Jr3ecf0PTNyTaLdw8sW9zFC2iIy6YenvZbh74YH3YU71ITD+4/Ca/fns1fk2LpE9srJ3FTrzm4Q1KNjXsKnY1R/GEEEI0gaxkO6cd25n/rdtHpLHTkW4ibBmLMexFYYFUyXclFH1WhKfAg8VhIW1kGl0u6YIlOfB94dM+8qvyeWvTW0wfOr25LkkIIWK2ZNsBHvtsU8T6r778ADUSrAZj+3ZgTE70ufaa2hEVTK3bU8b9H6yLGEg1lKDh7nfWMKhrKgO6tM0ucK4mGJRc33GEEEK0DT87dQBfbwzv6heZxt5hPsqonUq46NMiCj8tpOeNPUkZnIK3xEv+q/nseGwHOb/PwbAGkvlWm9W8suEVrhlyDYaSBL9CiNb1z6+21BqaUyOW/AAAqQ4rT181okW7Kx9RwdRTX28NG6wGsf0CPKafmfNzeXzq8IP7eU0/n6/fz1vLdlNQ4cZnajKSbJw2qDOXje7VYllCIPC0MpJ4MpsAZCbLRItCCNGWHdstjb9dcjx3vbO6wYDKkrgTZamqtcx0mRTMLaDHDT1IHRZ4QO69RH4AACAASURBVGjvZCf7tmw237WZsoVlZJ6SeXD7Km8Vi/cuZnz38QghRGvJK3GyfOfh5QeoqvY1eznrOmIeQx2orObrTQVhzYI1v4CsM24laeB4DLsDZbGS1H9Mrawfpl/z8Zq9lLu9OD0+/vbZRkb+6Qvufmc18zcX8uPeCrYUVPLDjhKe+GoLY//8Fbe/toLtRVW0hJG9M0m2h7c8hQ5Kdm5ehN/rRps+XLnLKPnmpVrbJtosjO7dcs2eQgghGufc47vzxGXDSbQZ2CzRn7BaUteBUXtiSucWJ36vn7SRabW3dVhIHZZK5frK2tv7nHy+4/OmK7wQQjTC60t2RRwrFU9+AEMpPlm7tzmKF9UR0zL19vK8iJFhvL+A2Yt28u6KPPJKXFRHaOUCDj4p/HTdXuZtKuDF6aMZ28zpF88a2pV7318XcV0smU0AtNZcOEJmrhdCiPbgrKHdGNI9nWdf+pR3C60YCQm4vCZ+DRYDHDYLRkJlIHVVCLPSxJpiRUUIwqzpVlw7w8fOFrmKmusyhBCiflVFcCAXlfsDx/ur2UNH8jmUZTSe/ABVHpPN+yuas7Rhjphgat2eMtwRgp94fgEub2AuJrfXH3EG+rr8OvBLu27WD8yZMZZhPTMaVfZYJFgtXHFiNi8v3BExEUVDmU0sKvCkM9Uh3fyEEKK96JmZyPXzX+bXv/0d36f2ZlexkwqXj7REG307JfN54Zd8vrP2PpYUC75KH9rUYQGVr8yHNeWI+eoXQrRXWsPOhbDwX7DtG7AkcJvH5Ga7JgEv63UfnvWdw1f+EXHnByhxehvcpikdMTVquTvyBxfvL6Cq2ow4P0d9CSxcXpNrX1rKkntPx25tvp6TN07oy5tLd+M14+8PmmCzcNvk/s1QKiGEEM3FtWIF2uslY/xYzokwoHrV0k4oFDrkmyupfxLKqihfXk76iekHl5tuk4o1FXS5pEvYcToktszklkIIQWUBvHohFG8HrxPQ4KsmGQ62tI9UW/i77RmqcHB5jzspjCM/QLTJz5vLETNmKiUhclwYmqAhFpECqfKl71P81fOkj51Kzztm0+PWWaSOOBvXliUHt/H4/Hy2fl9jih6zrukOXr5+NIm2hoPCUA6bwcyrR5LTMbmZSiaEEKI5lLzxJhmXXxY1M9WpvU4l0ZpYa5klyULnCzqTPzufijUVaJ/GU+hh99O7sWXZyBhfuxdFkjWJ03ud3mzXIIQQB5Xnw8wJULgRvFVEvvMOSFFuOlHKx2l/5tgJZ8WcH6Cl73ePmJapYzqn8qWlAI8Z36zxoUkoAr/Q2l9YsWYQqfKYzJyf2+xzVY3sncUbN4/lmheX4DU1Lm/0VOcOm4HFMHjx2lHNPqZLCCFE0/IVF1M5fz5d7/t91G1GdRlFmj0Np89Za3mnszthSbawb84+PAUejESDtBFpZM/IxrDVfo7qsDo4qcdJzXINQghxkMcJs84G5wHwx9bLylCQTDULT17CqOSp7GooPwCac4Z1a47SR3XEBFNTR2fz9PzciOtiTdAQNoqX+BJY5BZUsuuAk14dmncm+eHZGSz47am8tzyPZ7/dRrnLi9trglLYLQaGAUl2KzdOyOGy0dlkJLVcCnchhBBNo/Tdd0k9/XQsGdHH4yqluHbAlTyx8gmqjdoPE7MmZpE1sf4MrgmWBK4ZLHNMCSHCVXgq+GDrB3yx8wvKqsswlEGHxA6c1+88zuxzJgmWhPgOuPpNqNxfK5Dq80QFTi9s/0UKyfbAffgLKzzMXuNl3vRDLUxJuPnD8fu5d8gTUQ9vMRTnDmv5/ABHTDDVPSORE/tksWBr5IxEDSVoiCaeBBY2q8GeUlezB1MAaQ4b00/K4drxffhhRwn3vLua8f06MrRHOr07JDE2pwOG0XITlgkhhGg62u+ndM5b9Pj74/Vu51q7lhH3vkGXC5LY43Bh6tgnZrcoC52TOnP5oMsPt7hCiCPI3sq9PLnyST7f+TkKhdt0H1y3pXQLawrX8PDih7l4wMXMOH4G6Qnp9RwtSGv4/ongGKnaTA3/XOLh3pOjB2c25edCy/c87JuGE0fEbewWg5tO6dtwWZrYERNMAdw+uT/Ld5bU2/UtXnElsNDg8rbsZGFKKU7MycJqGFwzrg8Du6a26PmFEEI0nunXfL+1iO1FVVR5fCTbrfTtlMyw/A1Y0tJwHHdcxP20aXLghRcpfvllet5/Hy9PGs2V/72SIlcRXn/DmaysykqmI5OXfvISyTYZTyuECFhftJ6bvriJKk8VfiJPEVTTrfjNTW/y5a4vmXXWLHqkNDD1zu4lgRToEdw13s5fv6/mttF2MhzRGwL8KC60fMdr5hlh6xw2gz+eP4RjurT8fXC7DqaKKqspdXrQGjKS7Iztm8Wtk/rxzLzcuAKqRJsFj88kQsbxWgksGswgogItRq3hQJWHrGTpzieEEO1BcZWH15fuZNaCHbh9Jj5T4zP9WC0GVovC7nJy9aSruc7pDavbvfn55N99DyhFzrvvYOsWGB/w9rlv86t5v2J14WpMv4lPhz/csygLVsPK4A6D+efkf5LpyGyR6xVCtH3bSrdxw2c3UOWriml7r9/Lfud+rv7kat49712yHPV0K97+HfjC57gDGNXdwqQ+Vh5bWM3Dp0ZudQJIVtWcZqysFUwpIMFm8NC5Q7h0VHZM5W5q7S6YcntNPlydz8x5ueSVuLBZAxGs16fpkZnILRP7csO4bJ6ft4Vqo+HAJtFm4ZdnDOD9FXv4cV/4JF/xJLDw+PytkjHP9GvKXV4yk2QOKSGEaOsW5R7gxld+wDR12PyIps9PtQ+qjASeLzB46a/f8NK1oxgTTCJU/skn7Hv4EbKum06H669HWQ71mEhPSOfFn7zItrJtzN4wm49yPwICAZSpTTSaKX2nMO3YafTPlKkyhBCH+LWfW768JSyZDUDJdyUUfVaEp8CDxWEhbWQaXS7pgiXZgl/7KXWXcs+39/D8mc9HP0FVAejILV0Af5ycwEkvVfGLMfU3DGSqwL16gtVAAxP6d+Tnpw1geHbzzfXakHYVTL27PI/7P1gHgNMTaHnyhDRAbS+q4g8fbUB7vFzo3c32QaNYtbsUv9a1Jrq1WxRKKUb2zuTnpw1gbN8OdElzcO97a6nyhLdoxZLAQik45ZhOdEiJczBeEyh1ekh1WLFaZACxEEK0Zd9tKeSmV5bh9ka/qajh9vnB5+faWUt5bupQ+s/+N67Vq8l+7jkShw6Jul/f9L48MO4Bfnfi7zjgPkCVt4pkWzJZjizsFunBIIQIt3jvYsqqy2rNWQdQ9GkRhZ8W0vPGnqQMTsFb4iX/1Xx2PLaDnN/nYFgNfNrHyoKV5FXk0TO1Z+QTNFD3DO1s4ZxjrDy6wMOxnaLfzyYmJHBm/y4M6Z7OFWOy6ZwavSWrpbSbYOrf32zlya+3NPgFFAiyDOYm9+f2/h157NLjeW3JLjbtK6fCHZg1/thuqVw5pjc9Mg7NzXHW0K7c+/66qMdtKIFFos3CjFYY9AaBLn6tEcQJIYSI3c4DVcx4dXlMgVQot9fPza8sY7YjjRHvvYuRFFuSI5vFRtfkro0pqhDiKDNr3aywVinTZVIwt4AeN/QgdVhgLJK9k53s27LZfNdmyhaWkXlKoKuwX/t5Y+Mb3DX6rsgnSO0aCKhMT9Qy/GGSgxHPVvLrcdHvaQcNOIbnpo6K8+qaV7sIpt5bkRdTIBXK7fPz9LxcuqU7+O1PBzW4fYLVwq0T+/Lvb+IbbwVgMxR9OyUzsnfr9D0/UCnjpYQQoq2bOX8b1b7I3y9VG+ZR/sNcvAfyMOyJ2Dr3JX38VBw9Ay1QPouduSPPZVSMgZQQQsSq1F3Kiv0rwpY7tzjxe/2kjUyrtdzisJA6LJXK9ZUHgymv38u7m9+NHkwdey58/ad6y9E/y+CyITb+tdTDcZ3DW6f8pgVj+JUxXlXLafP9wqp9Jg98sD5iIFW1YR57X76TXX+/hLynprH/rQdx560/uN7lNXnww/WBOZhicPvk/kwa1IlEW+wfixU/HVLsvHL9mKgz1De3A1XVdEyRYEoIIdqqqmof76/Mw4zwTLB86fsUf/U86WOn0vOO2fS4dRapI87GtWXJwW18wLvL9+D0tGzGWCHEkW+/cz82S/i4e7PSxJpiRVnC72+t6VZ8lbXrI5fpotqsjnySzD7Qo+EWpQcmJlDliZARDqhQiTyd14cDlVHO0UrafDD1v3X70Dr8Q43ly6fGJ2v3xnQupRRPXTGC84b3INFmiTCFb21JNoPuzmLmjEts1ZahYsnkJ4QQbdoHq/ZgRHjg5q+uonTBa2SdcStJA8dj2B0oi5Wk/mNqJTiCwNjcj1bnt1SRhRBHCZfPhYpw12tJseCr9KEjpLv2lfmwptTu4GZVVtw+d9i2B510J9hqt67vuDOV0/seOk52uoH7vrRaE/YCuLSdmeYU/vXNNsY/+jW3vLqMDfnlsVxes2vzwdQz83LDkkLE8+VT5TF5Zl5uzOezGIpHLzqOl68/kcmDOmO3GiTajIN/YjaLItFmYUCXFB6+8Dg+uOwYqh/4Ld79+w/3UhutqNJDh2QZMyWEEG3V8p0lBxMnharesxHt85B0zLgGj+H0mCzfWdIcxRNCHMXS7GkRGy6S+iehrIry5bWDFtNtUrGmguTBtQMer99b/7x1A86AAT9BWxOjbxOBR1vZrrvykvlT3F4/1T4/n23Yz8XPLOSTta3/gKlNj5mqcHvZWlAZtjyeLx+AHQeqKHN6SY8xdXjNRLgn5mRRUOHmiw37OVDpwevzk55kY2zfDgztUTPbc0/Mq65mzy9/Re+X/4OytXx68uKq6laZpEwIIURsSpyRJ9I1XeUYSWkNTwpfc5yqhifkFUKIeHRL6Yapwx/2WJIsdL6gM/mz8zEcRq1sfrYsGxnja6cj75LcBatRT2ihFFz0HNWPnIjdvwvDaHgYjlvbyNOduMpzL9Uc6oWldWA4z6/eWo3dYuH0wV1iv+Am1qaDqVKnF7vVwFfnaV68Xz42i0GJ0xNzMBWqc6qDq8b0rnebDjffhGvlSgoee5wuv/stWmuW7yzhue+2sS6vjCqPSYLVoFuGg+vG5/DT47qSYI2t7JH4TD9fbSxg7so9FFRUs7WgkuzMUhw2C+cd3x2HrfHHFkII0fSS7JHrZUtiGn5nOdpvxvSdlpQg9bsQomklWhOZ0ncKc7fODQuqOp3dCUuyhX1z9uEp8GAkGqSNSCN7RjZGSI4Bh8XBqIwLuff9tRRXeUiwGvTMSOSCE3owIOSBf8l7cyle1JWcn50Nq2YBCrzhkwQ7tR0Dzafmidyx5hiKfnggYoIet9fPz95Yyby7JtElrXXSpLfpYAoCkWdd8X75QCAYbi7KMOj+l0fZfvElzO9xPE/uTeBAlQeX16xV/oKKan7//lp+P3ct08b25ldnDMRujb2nZVW1jxe+28ashTvwmn6qqg/9wZe5vDz04Xoe+nA9l4zsyR2T+9O5lf6ohBBC1Na7QxJWQ+Hz1/5SS+gxCGW14dy8iORBE+o9hs2i6JUl2fyEEE1v2uBpfLztY0wzvLUoa2IWWROz6t3f7TX5cEE3qty7Di6zGPDigu0M6JLKrZP6cYpzF4VP/JPer83GyMmBM+6Dde+iv/8n3gPbsWkfPiwUkc5/fGcyx5zMrqVfUrbkP3Q483YcOSNQFiuu7ctxbVlyMNupX2teXbST3/xkYNN+KDFq08FUZrIdb4TUR/F8+QB4fH4ymzlBg5GezpzpD/LKmiKqLdFTuNeM//rP9ztYlHuAV28cQ5qj4RazgnI3lz+/mD0lLqp9kY9f0x//9SW7+Gj1Xt64eQyDuqZF3FYIIUTLuXRkNi98tz0smDISksmYcBXFX8xEGRYcOSegDCvuHatw71pTaxywoRSXjsxu6aILIY4C/TL6MbLLSJbtW4bHH30uqEi034an5EQ87tr32qYfTL+ftXvK+PWcVYzau4F//eUvJOTkBDawJ8GIaSxOP5sbXv4Bt8eLPySdQ02OhA5n30nSwPEHlyf1H0NS/zEHf672+Xll0Q5+cfoAbJaWTwfRphNQpCRYGdwtPBgI/fJxbl6E3+tGmz5cucso+ealsO0Hdk2NKWA5HE/Py2X25iqqY5xd3u3z8+O+Cq59aSmeKMFRjTKXl4ueWciuA86ogVQon19T4vRw6cxF7DwQ3nQqhBCiZfXpmBwy1ra2tBMvIvPUGyhbNIe8J68i75npVKz4mMQBtccFH5+dQa8O0jIlhGgef5/0d7qldMNmxH7PrP1WTGdvPAVT6t3O5fOztPNAfrXVjr/OQ6X/LNyOy2PWCqQgvhwJptbM31QYc7mbUptumQK4ZVI/7npnda0ubRD48jGSMylbNIeijx9D2RNJ6NKftHGX1douOcHCLRP7NWsZN+2riDqpcH0TMXp8fn7cW85z3+Zyx6kDoh7/V3NWsb/cHfZEs6HjV1X7uObFpcy7a1KrzYElhBAi4PbJ/bj9tZURJ4ZPGTKZlCGTo+6bZLdw66Tm/S4TQhzdkm3JvD7ldW754hZyS3Nx+pzRN9agtR1f5UDcey6jpn2mvvvSaq1YvK2Yp77Zys9PO3Tfu62wikgzS8WTI8FranaX1FPeZtTmg6kzBnfBakQOBBr68oFAt4ifDOkadb3WmsXbinlvRR77ygIBS2ayjdMGdWHKsG4xJXN44btteCPk4C9f+j5lS96pt5+n2+vnxQXbuXVSfywRrnNvmYsFW4sadXy/hsLKahZtO8D4fh0bvA4hhBDN59RBXbh4ZA/eWZ4X8eFbNIk2C5eM7MnkgZ2bsXRCCBFIk/7KT1/h611f89K6l9hauhWtNR6/B4XCbrFj+jVeZ29cRadgVvWH4ARCsdz3urwmz3+3jVsn9TvYJS/SAyaIL0eCz/RHnH6iJbT5YMpmMfjbJcfz8zdXxvXlA+CwGfz14mERkzx4fH5eX7KTZ7/dRpnLi8tj1oqK528q5IEP1jF1dDa3TuwXNZlDhdvLR2vyMeu0GsXaz7OmLN9sLIiY1vHVRTsjRuuxHt/lMXlu/jYJpoQQog146NwhzNtUyN4yd9j3RiSJNgvnD+/OQ+cOaYHSCSEEWA0rZ/Y5kzP7nEluaS6L8hdRVl2GYRhkJmTy0ueJbM6vHULEc9/r15rP1+9nyrBuACTbI4cj8SXoMUh1tE5Y06bHTNU4c0hXfn/2sThssRfXYTP43U+P5afHdQtbV+byMvXZhfzlfxvZW+bGWSeQgkCiiCqPyezFOznziW+jzrL85Y/7sUToQhdPP88qj8lrS3ZGXPfakl0Rx1TFenwNLNx2gJKq+AYTCiGEaHrPfruNrGQ7D58/hJ6ZiSTZLdT9BlEq0K0vOzORP10whD9fdBxGlB4aQgjRnPpl9OPqwVdz+wm3c+vxtzIq6xx2FYbnB4jrvrfa5Llvcw/+PLh7GpGquHhyJFgNRf9OKfFfYBNo8y1TNaaN60OXNAd3v7smLC14qOQEC1bD4C8XD+OsoeHd+9xek8ufW0RuQRWeCJkC6/KamlKnl6nPLuKDO06iX51f1L6y6ohJIeKdCyu/1B3h3H7K3Yc/0aPdYrC3zN3sGQ2FEEJE9/XG/by6aCcf3HESXdIcXH5iL5btLOHFBdvZur8Cp8ckyW5hQJdUbpiQw8jemTLeVQjRpmwrrMRqGEDte99473t3Hjg0vumGCTn8b92+iN39Ys+RYGVs3w7xX1ATaDfBFARaqE4d1JmvNxYwc34uq3aXYtN+UOBTFo7rkc4tE/tx+rGdsUZJjfjAB+vZVhg5kKo3mYPHx1UvLOH7e06tNbbJ4/NH7KoR71xY1b7wPyCnx8RqqIjjpeI5vlLg8voaLIMQQojmsbWgkrveXsNz14w6OLGkUorRfbIY3af++VuEEKKtqPL40BEmgY33vjc0cBraI50emYlsLaiMuG1DORISbQY3ndy31Vrw21UwBWC1GJw5pCtnDumK22uy7R9PYiSn0Ofm6xpMFlHm8vLBqj0RW5IaGjSnNVS4vMzbVMBpxx4a25SsTGwGeOocMt65sNISw9NQpiRYI2bwi/f4WkNKQvOmhhdCCBFZudvLza8u4+6zBjKyd2ZrF0cIIRot2W7FiNBiHu99b9179t+cOZBfzlmJK878CBAYLzV1VOvNwdfugqlQDpuFjoYPw6Zjyrr39rLdEf8AYh00V+UxeebLTYzetRrnsmU4ly+jY4mJMfp6MGp/lPFMxGi3KMZFaJq0GIouaQ72lYV3AYzn+F7TT4/MxAY/H9E0yt1eFuUeGqeWkWRjbN8OZCRF7ma54cAGPtvxGfur9mNqkyxHFif1OImTup+EJcbmciFE22T6NXe+uYoJ/Tty2eherV0cIYQ4LDkdk/H6wwOeeO5LAXpl1Z4z76yhXVm9uw//Wbgzana/SBJtFl65YQzpSa3XaNCugykIpDYPG70bxazvd0T8BcUzaG7NrmI2rf+GnJFD6Xr//fQZPJgn/rWQHQfCc9vH2s8TpZg2rnfE8904IYfHP9/c6H6kFqWYclw3UhLa/a+6zVufX8aL323nv2v3YrMYB7t/WgyF1/RzxuAu3HRyX47PzsDn9/HJ9k94ce2L5Ffm4zE9+EP6H8/dOheH1cG0wdOYOnAqafbwyauFEG3f37/YRFW1j/vPGdzaRRFCiMM2oEsqvbKS2Lw/vEtePHPA3nhyTtj+d581iASbhWfnb8PtM4nQm/CgBKuB3Wrw8vUnMjw747Cv63C0/ztsTcwDdPeXh7fwQHyD5hKSE9G3PUjHnEN93G+Z2I8/frwhYn77WObCUsCavDJ6ZCSGXculI7P522ebou7b0PHtVoMbIvzBiqZj+jW/f38tc1ftwWtqTL+O2JX0k7V7+erHAiYP6oAzayYbitfi8rkiHtPpc+L0OZm5eiZvbnyTWWfNIju19ZqwhRDx+++avcxdmc+Hd5x0cD4VIYRo726d1I/73l9HVaPve1XEJHFKKe48/RhOHtCJ577NZd6mQoBa91TJdgsWI9AIce24PlGnLmpJ7b921zqQYaEBfr+OOv4odNBcLKqqaydzOH94D9ITbRHTOjbEYTO49+xjefLrrUx9dhGrd5fWWp+eZOOSkT3jSgtfw2ZRDOmexpDu6fEXTMTE79fcOns5H6zKx+2NnIzk4LY6MODyf+vzWbDseJze6gaPX21WU+gq5Mr/Xsm+qn1NWXQhRDPakF/O/R+s49lpI+mQktDaxRFCiCbz06HdSIhheE0kiTaD607qQ4I1+v4je2fy7LRRfP/bU/ndTwdx8Z6lTB3akdsm9eNvlx7P8vvP4K6fDGoTgRQcKcFUDP38DENhs0TeLnTQXCxS6kwKlmi38NaMcaQ64guoaiYVvnZ8Hz7+2QQuHZnNTa8s45dzVpFfeqjF4qHzhjCkezoJESYfjkYpyEyy88K1o2IvkIjbXz/bxHdbiuLq3+v3W/E5s6ned97BZSXflbDlvi2sv3k9G3++kfyX8zGrAsf0az8VngpmfDEjYgYdIUTbUlzl4eZXl/HguYMZ2kMeZgkhjiwOm4XXbhxDkj2+gMphNRjZO4s7Tz8mpu07piRw7fg+3LjqfR6degJ3nzWIs4/r1uZa+ttWaRojxpYpgJ6ZSRGXxzMpmMfnp3eH8ONkZyXx8c8m0C09keQG/rgcVoNEm4UnrxjBecN7AIFxNVNHZ/P1bybRMzORs//1HY9/Huhrb7MYvHbjGMb16xDzH67WUOr0cMPLy/h64/6Y9hHxKXN5mfX99oiBVNWGeex9+U52/f0S8p6axv63HsSdt/7QBtqOt2wkfl8KRZ8Wse/tfXSd2pXBTw+m7/198RzwsOOxHfiDTdumNtlbtZfl+5e31OUJIRrBa/q5/bUVTBnWjfOD9bsQQhxpju2Wxps3jyXNYY3aWBEq0WbhpAEdeeHaUbWmGGqIv6ICw27HsLfduVLbZTBVWe1j9uIdXP3CEq6p6MfVOzO54T8/8PGafDwRxqrUuOnknKjBSNqJF5F56g2ULZpD3pNXkffMdCpWfEzigNpJKcb2zaJzauRmxeysJL75zST+cskwhnRPw2EzSE6w4LAZJNktJCdY6JBs52enDWDBPZM5Y3CXsGOkJFj59ZkD+eTnJ7O72Mmpj8/jrR92Y7MYvHTtaP4+dTjDszNi+sP1mJrlO0u4/bWV3Dd3bb1d0ET83l2eFzE7ZPnS9yn+6nnSx06l5x2z6XHrLFJHnI1ry5I6W2rce4dTMLeA7ld3J3VYKsqqsHeyk31bNp4iD2ULyw5u7fa5mbV+VjNflRDicDzy3x+xWw3u/smg1i6KEEI0q2E9M/j8lxOZNrY3yXZLWGOCoQK9sI7pksL/XTSU56eNiin7diizpARLVtuei69dJaDIL3Xxz6+28MGqPRhKBRM+JEE1/LixgMXbDvA7Yy3Txvbm1kn9SHXUTpN4/vHd+OPctUTrFtjQoLlku4UZp/Srt4x2q8E5w7pzzrDubNlfwY/7Kih3eUmyW+iWnsiJOVkxReTdMxJ54vITWLW7lIc/3sCshTu4f8qxnDW0K+mJNq6dVffGPDqX1+Td5Xvw+jSPXnxczAk7RHRaa577bltYq1SsafYDB7FTuS4Tv9dP2sja2fosDgupw1KpXF9J5imBeWk0msX5iyl2F5PlaNsVixBHo7eW7Wb+5kLm3n5SXE9ehRCiveqa7uCBc4dw91mD+GTtXr7fWsSeJStI6dGN3sf04sITeh5Wd2dfcTGWzLY9P1+7CabW7SnjqheWUFntxYzS+FSTVeTFBYH01HNuHkfX9EArknvTZgofeojz0ofzYYehuM34WmkshqJruoNx/cLng4pmQJdUBnRJjes8dQ3PzuDtW8bx6bp93PPeGnI6pLBsZzEeX3j5qzbMo/yHuXgP5GHYE7F17kv6+Kk4eg7BHWZJgwAAIABJREFU5TX5cHU+o/tkckkrTmx2pKio9lFUEZ5AIp40+wBmlRNLig0VoaXRmm7FtbN2tj+7xU5+Zb4EU0K0MSt3lfCXTzcyZ8ZY0iNMwi6EEEcyh83CRSN6ctGInuz+5jkyhl9C6qlDDvu4Zkkp1jYeTLWLbn65hZVc/txiylzRA6lQ1T4/eSUuLnr6e0qKyyl4/HF2TZ9O+vnn8ciTv+KE3plxJXMwFKQ5rMy+cUyrtOoopTj7uG588cuJKKUjpmCPpWuZy2vyz6+2SBKDJlDh9kUcABlPmn0AS1IKZqUPHSG495X5sKaEP++o9IbP7SCEaF711ZsF5W5ue20Fj148jP6dD+8BmhBCHBGa6F7TLCmWbn6Hy+/XXPPiUqo8vrB19bXEmH5NYbmbGfe+zD+S99L3ww+wduoEwKzrTuT211ewKPdAxMAklMNmkJVkZ86McXRLT2yWa4yVzWKwbk952PJ4upYdqPKwbGcJo/u07T/Mts5uMfBHqChC0+zHNG9Z94Eoq6J8eTnpJx5qBjfdJhVrKuhySfi4umRr8uEVXgjRIK01i3IP8Oy32/hhRzEur4lFKVIdVs4f3oPrTupD7w7JVPtMZsxezhUn9oo4DlYIIY46h9vwULYH1rwFpTtwbF2NLdkKa96GweeBte1NNdHmg6kFW4sodXrCAtzype9TtuQdOpx5O46cESiLFdf25bi2LMHRM9Cs6NWwKqsv6u4bsWYcCoQcNgvPTxvFZ+v38cz8XDbvq8Dr99dq9UqyW0hJsHLTyX25/MTssPFXreH7rUW4feHBXzxdy1xek5cWbJdg6jBlJNkiBlOhafaTB01o8DiGI43O53Ujf3Y+hsMgZXAK3hIv+a/mY8uykTG+9qzeHtND95TuTXYdQohwn63fx0Mfrqfc5a01KaVPa0qcXl5bspM3lu5iWM90Oqc66JLq4I7J/VuxxEIIcQTY/i0s+Afs+B7QYHo4mPLt4x/h4zth5HQYe9v/s3fe4VHU+R9/zcz29IQECAkk9N57tSsoFiwodlFEvTv11PPneXd6d57neZxdURGwYEHsIoiKSpcOobeQkEJIr9tn5vfH0sLOJjsQSALzeh4fH3Znd7+72Z35vj/l/YGYpuOW2uTF1FtL9gVNWNbV5C+IfLAqm8fH1nZWEkWBsb1aM7ZXa/YWVvFdxkEKKj14/TKJUVaGd2jByI4tEJtQE/H+4hr8GuVgekrLVBX2FBplYidDTqmTDQfKqHT7sUoifVJiWZ9dxvF/keNt9gVRwpbeD0E04c7ahPtABnHn33Xc0QqSI5PEK+KQolQK5hbgLfQi2kWi+0eTem8q4gnDmge2GkiCPfy+PQMDA328syyTaT/swu0LXVPuk1VAZV1WGYIAc6cMa1LXCgMDA4NGR0+Zn6rCj3+FtTPB59Q+xnt477rmbdjwHkyaB+3C608/3TRpMVVS7WFtVlnQ7XoyMV5Z4cM1wWLqeDomRfHgRU2/zr3G68ev0TSmt7TMqVEyaaCNrKgs2V3I9F/3kZFbgUkUkBX1aO+c1qkievAExIg4KlbNpXj+NASLHWvLjkQPm1j7QNGHJWEpAPFj4okfU3e20GFycGfPOxvibRkYGGjwxfrceoXU8agE9gB3vbeW734/irYaMwgNDAwMzjn0lvktfBw2fhBaSB2P7A38N+cauO0bSB18cmtsQJq0mDpY4cZiEvGeICD0NvlXuf34ZKXJTUzWS6TVhEkSkU+YpaW3tCzC0qT/7E2GshovN7+zmuySmqPZ0WD/Pm3qs9kHBUF0ITn2hfV8oiDSwt6CIa2G1H+wgYGBbiqcPp74cgsejVmFdfXnAtR4/DwybxPzpg4PeqyBgYHBuUCFy8e8dTn8squQIsdwrGu8tDm4nhsGpTCmc1LocREZ84KEVNpLVTh9sP/BSCIsgce9s8HLnAwfv95xuG/c54I518JDGWBvXLe/Jr2rdvlkTXGrNxNjEgVcPrnZi6n0FhGYJAHPCYklPaVlggCdkiL1vfDBzbDyVcj8Fbw1IIiBL26fG2HQ3RDV6tTfXBOjrMbLFa8up7DKfbikp4ER/NhT30MQwnhuVcAs2nlxzBvGjDADg9PEp+ty0LrWh9Ofq6iQkVtBdkkN7RIMgxiD5klxtYdP1hxgfXagnN1hlkhrEcHNQ9vStVV0/U9gcE6SVVzDSz/tZuHWAkQhsN9GioFKlS3bCli2twirSeKukWncPbJ97aG9qgq//EszIyWr8PJqL38eVYfhhOyDjXNg+O9PwzsLnyYtpqJsJs2SS72ZGL+snhXZmOEdWmAzSdR4gk0owi0ts6kyt3SwBT1ek/3LYMFjUJ4Ffi+ox72utxpWvgIrXoH2Y+CKFyEm5RTeXdNBUVRum7X6tAgpAbBbJNp3XcIhpQx33WaSqIqEqtipOXAf41/czvg+5Tx6SZej89MMDAxOHUVRmbEsE9cJ5X16+nMVVeXdlVk8Nf7U56oYGJxJtuZV8PLiPSzdXQRQKzu7Yl8x89bnkJ4Qwe8u6MS4Xq2MoJ7BUdbsL+XOd9fg9irIIXqkajwyNR6Z1xbvZeGWAj68ewixDkvgztx1UH1I83GPDbfw/AoP9w+yEGsL8Z3zu2DlazD0ARAbL2HSpBVGapwDn0aPkL4mf0iOtZ8V0+glUWDyyHReWbwHt0YpSv2lZRAvKSQ8MoWc3r2Jv+N2HENCzM7a9DHMfzjwRQ2F/3DR297F8OZIuGMBtOyu5y01SVbuK2FfUY2mkKqv3CcUFkkEAYakx/PU+O60S7iQL/Z8wextsylxleL2u+G4LJUqWwEBb+kwfGUjUeUIQOHLjbks3nGIj+4ZSrfWRqTQwKAh2FFQSfWJKX/09ef6ZJWvNuYZYsqgWfHNpjz+9HkGHr+iGbxWVHD7FHYUVPHovM0s2V3Ivyf0Drmn2ltYxczlWfy88xDVHj+iIBBjN3NV32RuHZpmBALPIjJyy7l91ppAJioM3H6F3YeqmPjWb3z1wAjsFgl+eyNQrqfBwGSJ89JMTFvp4ZkL6vjeeKsha1kgsN9INGkxFWE1cXmv1ny9OR9Zqf0rDzcTY7dI3DMq/Uwu+7Ry4+C2vP7rPtAQU/VhN0s8fHVvOv11MRXffEPBP59BMJuJv/12oi8fh2g5HCnYtbB+IXU8qgyuMpg9DqYug9hU3WtrSry1dF/Iwcj1lfsAmCWBFpHWo9G9KFvge3zL0HYkH2fRP7HrREa0HM+4t97Ha12HYK4EFJAj8Vd3wV/VA6hdxiorUOb0MfGtVcw3Gt4NDBqE4mqv5ubwZPpzDQwaCkVV2Fi4kfzqfFx+F5HmSDrHdaZjXMPY8C/Yks+fPs8I23DF5ZP5dvNBFBX+e13vWoHYTTnl/O2rrew+FDxqpsrtZ8ay/byzbD9D28fzzNW9SI03rl3NGY9fDimk6go6+2SVrJIanvpmG89f1xsObUXbyivAP863MmJWDQ8OsYRejOKHkj2GmKqLyaPSWbD1YJCYgvAyMYqiMmHA2VF+BhAfYWHW7QO5ffaasE+AEBBSV/VN5tr+KQiCQNwNNxB73XXUrFhB6bvvUfTCC8RNuonYCeMxfTY5SEiF1QzoqYQv74U7FzTY+z3TFFS4WbO/NOh2PeU+flmle3I0M28fVO/r3fvBemoq2iKrtQVo4GT0SMgMWLXHz13vreXHh0cbJRcGBqeIlksq6O/PVVQVVVWN36TBKVHhqeDLPV/y/vb3qfHVAAFhJYkSsiLTNrotk3tO5uJ2F2OWTm4GZm6Zkz9+ullzH1HXZtjlk/ku4yBD28dz3YDAdWvRtgIe/GRjnXsS7+Hg4rI9xVz+yjI+umcoPdvEhDzeoGnz/dYCTbOecILOHr/C15vy+MsV3Yj21tT5Oj2TJK7obOK55V66JYYo45P94K485fd0KjR5R4YeyTH0TonFIum/ONnNIjcNbkt0Exi425AMaZ/AzNsHYZYEwvlU7GaJGwal8Ow1vWpd5AVRJHLUKNrOfIfUme/gzcuj+P4LULxezec50gwYElWGvPVQul/nO2o6bMmrCJTknYCech8V2JxTXu9xW/MqyCyqCaozrlzzJaWLZxAz9AZSfjeHNvfNJqr/OFx7Vh89RlEhr8zFpjBex8DAoG5iHWbN4Ojx/bnhYDNLhpAyOCVW5q3kks8u4fVNr1PkKsLpd+L0O3HLbmp8NbhlN7vLdvP3VX9n7BdjyanMOanXeXdllmaQOpzrj8sn8+rivaiqysp9xfUKqeNRVKh0+7lpxm9kl9S9kTZoukz/NbiC50jQOf7i+3B0GY5osSFIJhwdhwS134iCwOfrcsFSf4by7+fZmLHBS15liAyWZAJr4443avJiCmDGbQNJirZh1iGobGaRfm3j+Mvl3U7jyhqPER1bMKFfG/q1jcVulnBYakdNTQJYVJnB6fG8cUt//n5lzzqHSto6dyb5n/+k5UgrouDTPOax4RamrfRQ7q7DlEFVAgPVmilVbh+KRuG43nIfrTLBE5m5fP/RaN0R9JyMPH6ZGcsyw1qPgYFBaLq1jsan1N2f69y9CsXnRpX9uPato+yXWbWOFYBBaXXPijMwqIufD/zMH375w1HxVBdOv5MiZxETv5tIVkWWrtfx+GU+XpMT1Bes5/pTWO1hbVYpU+esD5ndOvjeQxx44TpyX7uVQ58+hTt327H7PX4e+GiDrnUbNA0yi6rJ0hDCeoLOLp/M7JVZ0KIz1JMW6BgvMrGHmVfWhAjmiyaIa9x2niZf5gcQYzfzze9GHp35U99G1WGRGNWpBa/c1A9TM7dDrwtJEpnQP4UJ/dvw7eZ8NmSXU1rjJcIqkRJpYuBzf2TUX79CigzTqrd4D4KzOOTdYTUDyl7Y/DFc9u+TeEeNj8UkNowdfz3C3+tXWLDlYFBWSs/JSFHhp+2FuH1ybatRAwMDXTgsJq7pl8K8dTn4T6E/994x7c/ksg3OInaV7uLxpY/jkcOdZggKCtXeau5cdCffXv0tkZbwxp4s3lGIVipWz/XH7ZP5z/e7kDWMmsIdJ7C3sJpdBVV0adW4WQUDfeSWuTBLYpCI1ht0Lqxyw5D7jo3dqYO/jbHyQYZ2oB+TBTrU3fJzumkWYgoCvUJfPzCC77bk8/z3uzhU6UYShaORFaspIJr6tY3l3tEdOK9L4llfblHj8RNpNeGwmJg4qC0TB7Wtdf+BL9tRs3w50ZddGt4TVh8KKPw6CKsZ0F0RmB3QDD//VtE2zVZIvXb8LSLqmIsAlLu8mrEYvScjSRQoqfHS5jhjCwMDA/1MHpnGlxtyg8QUhNefG2M3M6x9wulansFZzmsbX9MUUmXLyiheVIy30Itkk4geEE3L61oiRQSuESoq1d5qvt73NTd3u1nzueUaH/4iJ4pbRjCLFOdW4tHIJum5/qgqbMurCHIW1tNf7PMrzFyeyfPX9an39QyaDk6vrOn8qDfo7PUrkDYSbLFBYirrodoCOzVGxP0XDQdjky0gyMLcM50umo2YgkDW4Jp+KSzYUsCNg1KxmSVKaryYRIGESCuXdG951jvEKIrKuuwy8std7C6oomWUlZ5tYuioMYg36qILqVq8OHwxFUZELKxmQFUNlPsJzS9b0r9tHDZz8CwvPXb8drPELUPb1fk6bq+iWXap92QkiuAKo6TQwMCgbjomRXFJj1b8sL1Al7kPBMrKn76yx1kfwDM4PRS7ilmZvxL1hFBe8cJiihYWkXJ3CpHdI/GV+cj/IJ+saVmkP5mOeDiI7JbdvLv1XSZ1nXT0O6iqKt7sSqqW5uLeXYZgOna9Hu2VeUtx8CEeluLniAel3uuP1ogWPdktWYVvNucbYqqJ4PHLLNhykFnLs8grd+HxydgtEp1aRjFlVHvGdE5EFAWibSbNWLneoLPNLAWC7qMfg0V/1hzcWy+CBAPu0P+4BqZZiSkI9LT8tq+E/93Q56wzlqiLshovc9flMHPZfpxePwiBTfT+khre/y2bTklRTB3TgUt6tMR8uLQx8oILKHzpZVSfD8Ecxmdli6Uui8oj/P08G/3fquaRYSGyLyZLo0cJThZRFLh7ZDovL94TtKEKt9xHUVWurcdBMspmwq9RHqH3ZCQrKtG2ZvczNjBokky7vg+3vONmc265plOVFjazyKOXdOHSHq1O8+oMzlbm7ZoXJMRll0zhV4W0mdyGqN6BKL0l0ULq/ansfmw3FSsriBsdd/T4Sm8l6w6tY1CrQcjVXopnb8Nf5ET1KaCC6j8WdDMBXZB4HDt/ROURnOxC0X390UJvdYXXr+Dxy1hNtY/3+GUqXX4clkBPuBGoOH34ZIUXftzN+6uyQVWpOS5AW+OVKa4uISOnHJtZ4sELO3FJj5ZB/d6gfwZseovDLSgD7oDslbDz25AzpzQx2WHiBxCZdDJvu0FpdruwH7cfYkj7hHNKSP2ys5D7P9wAqLhO2OAf2fBvyavgT59v5t8LLXwyZSgpcQ58cQn82PNC7vvPTxR4Bbx+BatZpF1CBFNGtWdsr1a1T2CJXQMWk/VwfDNgrySN7FTrfqfydhudiYPa8vLiPZr31VfuY5YELu/Vmhh77e/n9vxKPlydzb6iGpzeQHmmlh+I3pORRRKJj6ij5NLAwCBsLCaR9+8azMjnf0Z2Burztcr+IOAWq6rw7NW9zqrxGwZnntUFq4NK/Jx7nCg+hegBtUubJJtEVO8oqrdV1xJTHtlDRlEG/SP6cOjVjSjVvkBjUh1EIAACrxHBozjZrOP6IwrUarU4uj6d2S1JPLw3MUlUuHzMW5fDO8v2U1jlxiyJyIqKWRK5dkAb7hqRTvvE8PrCDMLD6fVz28w1bM2vqDMjX+OVqfHK/HvhTjLyKhiUFs/yvcE99uEGnSOsElPHdAj8QxDg6unwjQm2f1V/hkqQwGSF69+FjhfqfcunhWYnpr7LOMj4PsmNvYwzxvyMfB6dpz0L4kRqPDJur5vLX1nOJd2TmJ9RAElDcVUfe6zfI7M9v5Inv9zCk19t5e6R6Tx4YadAyZk1EnrfABs/BLVuURWyGdASCSMf0v0+mxLxERb+d30fHgnzcz+CJAq0irbx9FWBBltVVVm4tYBXFu8hq6QGn6zWsqINZa4Y7snIYhK5ZWi7s9pkxcDgTPPp+hxS4hzMmdyL91dl8+XGPCRROFrWIisqUTYT94xqz/UDUolxnDuBPYPTQ6UneEaOXC1jijQhaJgZmWJMuLJrR/BlVabcVU7RjAyUai/oqFS1I/A8DiZTQ26Y1x9REFA0xJre7JZfUbFKIk9/s42P1xxAFISjg2CPZIf9iswna3KYty6X3ikxvH5zf5KiQphgGYSNX1aY/O46tuRVhJ2JPzJnbFiHBMxSsJiG8HpMRUGonc2XTHD1G9D5Ulj2AmrBNkBGOL5ayuwItJB0vwpGPQKJXcJa85mgWYmpCpePNftLefmm5p35CJeM3PKwhdQRZFUNRHfW5x2+RXujfSSN+/bSTDLyynnrloFYTCIMvR8y5oK/tpgKvxnQCp0uCXu9TZXLeydT7fHz1Dfbwvr8LZJA61g7c6cMI9pmRlZU/vrVVr7cmKc5IRzqDhqGczISgFuH1d2bZWBgED5ZxTW8+ONuPrtvOB0SI3l2Qi/+ckU3dhyspMLlwySKJERa6N462ig7MmgwLFJwdYEUKeGv9qPKapCg8lf4MUUGb99MJQpyuUdTSK3JzeDZX6azuzgLURTplNCOpy78PX1bB8bH2IC7sfI0rrCuP4qqal7D9FZXdEqK5I5317LxQFmdG3q/ouJXVDYcKGPsS8v4/L7hpLUI06nYQJMPVx9gU472517f4OafdxZiN4soKprzyurCbg5kpSymE/anggA9roYeV5N7/QW0Gp+O2VQdyFTZ4qDtMOh7E9hjT+VtnxaalZj6YVsBwzsmEGltVss+af67aFfIjXxdX3Q9uHwyq/aV8PDcTbw2qR9CUteAGNrzI/h11K4CiiLB6L8iNtN+qROZOKgtqXEO/vXVOvaVePCpAjK135vDIqGqcE2/NjwxritRNjOqqvKXr7byVR1C6lSxmUQu7tGS1jGGi5+BQUMgKyqPztvM7y7oRIfjSokcFhMD2hnzowxOH8mRyWwr2VbrNkdHB4JJoHJ9JTGDY47eLrtlqjKqaHldy1rH2wUzkftUVG/wnqHKU8Odn/0f/7rkj4zvej5e2c+a3M1YjxNxEgKjMBEFVIWx5rr2z+FWVzjMEjazxIbsMk0zCy1kBcqcXia+tYqFD402ytxPElVVeWvpvqDWEQjP2h4C7tm7Cqopd/nCFlR2s8h5XRO5/7wOIY/x5eXhyvNguvltkJrHfrJZqZLvthxkQv9zozb9YIWLNftLNe8L94t+PHWJL7dP4Zddhfyw/VAg7TphBsweC4U7whZUqtmO09uTwue+IfXNsZhbNfNmbFWFjXMYvmwa33kK2WVOYKb/MlYp3anBjgmZFmIlN8fs4eorriSiS6+jD12wpYCvNmkLqYYQwVaTSKeWUUy73nBAMjBoKGYuz0QSBe4cntbYSzE4x7iu83WsyFuB03+sV0RySCRdnUT+nHxEm1jLzc8cbyZ2eO3ovKzIDD+oXfaUWZoDwNXdLwLALkqMSR8cdJwKXIGFjwkxHFUH4WS3nD6ZbXkVaFSK1XmtVFQorfHywo+7eObqXsEPDpeKPFg7A3bMD4x0EQSwx0OfG6H/beA4e4Moq/eXUu4MbtXQY22/PruceVOH8cCHGyiu9tYdPFZV7GaRK/ok8+9retWZ2a9asoTIUSMRmomQgmYkpspqvKzPKuP1Sf0beylnhA9WZWv66un5oh8hHPHl9Mq8uWRfQEyZbXDnQph3B+xfAn53oE5VC8kKAggX/JWIofcTM2sWWTfeROr0N7B163aKn0IjIfvgy6mwa8HRRsguopPnLTOCj60CPvsYzn8Shv8egFcW79G0Kz8ZEXwiDotE/7ZxzLhtYJD7kYGBwcmx51AV03/dx9cPjNQcWWBgcDoZ2nooEeaIWmIKIHFcIlKERMHcAryFXkS7SHT/aFLvTUU0HyuREhEZGTmCKBLQcuRtH5+KKIg8/N2/uLLrhfRr04NYW/CgXBsCIzCFLaZOJThoN0skRlk4UBocsA3nWulTVD5fn8eT47pjt+i8FhbugO+fgAMrA4FT+bj3W30Ifn0Ofn0WOo+FS5+FmDb6nr8ZMHdtjuY+RY+1PcC6rDIWPTyaLzfm8eav+yiu8eL1yUcF8pEZsP1NNdxUs4vx146tt0S6eskSYq+6St8bamSajZj6YXsBozq3IOIcKfFbua9E03pS7xddj/janl9JZlF1wC3HbINJn0Deelj5OuyaHxBOqkJAPQGCCAPvhkGTIaYNApAweTLmNikcmHw3yf9+lsgxYzTXVe708vPOQkprvMiKGhh42SGBdgmNXAOtqvDVfbDzu/DLHH0u+OVZMNnYlnIDB0qDnWj0imCbWTz697dIIirQNzWWqWM6HJ31YGBgcOr4ZYVH5m3m0Uu70Dbh7J5TaNA0EQWR23vczmsbX8Mtu2vdFz8mnvgx9WRIBBgq9UczxQNEWSP44ubXeGP1R/zp+/9SVFPK+R2G8PxlfyIxovZzxwhCLT0WymTgVIKDdrPEpT1asmDLwaD79FwrBQG+3ZzPDYNS63y9WmQugY9vOhwoDVGaduTav+Nb2L8U7pgPLfW1UDQ05e5yfs75mWJXMV7ZS7Qlmn5J/ejZoudJ9W/mlbs0370ea3uPX+FghQuHxcTNQ9oxaXBb1meXsSarlNJqLxaTSFKUlct6tibJJLP30pfx7L4OW5fQxhGKy4Vr3Xra/Pe/ut9TY9JslMn8jIPcOKhtYy/jjFHp0nDKQ/8MB10D9BSV+ZsP8oeLOh27sc0AuH4WOEshbwO4ywMzpBwtoO1QkIKdrKIvuxRTyyRy//AHWtx3H/GTJh29b3NOOTOWZfLj9kOHbVUVVDVwwlZU6NUmhqljOnB+1ySkxhAMGXODhFTaS1U4fbD/wUgiLIE1vbPBy5wMH7/ecVj8+Zzww1+Zk975lEWwzSwytmdrOiQGnjvGYeG8zoln/UBqA4PGYPqv+4ixm5k0+Ny5vhg0PW7udjO/5PzClqIteBV9ZXaKqvBC5Zt8mJ7AU7lTae1rEXRMpxZpvHj5nwHYW5LNH+Y/w9OLX+X1K5+qdVxClI2pfVtS5vQRaTPx7eZ8Cqtq27afTIUMBK7zoiBwz6h0kqKtLNpWECQA9VwrnV6Zzzfkhi+m8jbAxzeGPxxWlcFVCrPHwdRlEHvq5whZUdlxsJIypxdFhVi7mW6to4PNGA6ztXgrs7fO5tecX5FECbffjYqKWTRjEk0k2hOZ3GsyY9PHYjeF30PtCVGSp9fa/vjSPkEQGJgWz8A0bfHfYso9FL30MqnT3wj5fDWrV2Pr1g0pWsPgrAnTLMRUSbWHTTnlvH3rwMZeyhkj1A9L7xddj/jyKyoHK93adzriodNF9T7H0cP79SPto4/ImXIvvpxcWjzyCE/P38Fn63Px+OWg5tUjs1zWZZfx4Ccb6dwqivfuGnxm54mpKix5XvNEK6vw8movfx4VYlAxgOxlb1Y2spoQfJeOv4Pbp5ASZ+d3F3Sq91gDA4OTZ1t+Be+uzOLb34803PkMGhWTaOKNC9/g3h/vZWfpzqAMVX248XDAepDfp/+b57Mfpr0ndH95x4R23NDzMuZs+ibovj01bj74rZxOLaNoFW2jzBks7PRWyIhCIBs1cVAqdwxPp22Cg9d+3qPpIqc3YFxaE6bwVGT4aGLQ9T2sYKmnEj69Hab8Et5raVBS7eHjtQeYtTwLj08+Wl2iqoEM2y1D2nHrsHYkx9oP367y4voX+Xjnx3hlLwpKLZdGn+LDp/g4UHWA59Y8x9sZb/PuZe/SKiK8fvVYh7Zxhx5rewFoEVHHnujE15w4kZLZ7+IKvo3UAAAgAElEQVTcuBFHP21X7uolS4g8T7uiqSnTLAbULNp2iDGdE/XXxTZjWsdoz1A4/oseDseLr3AIFa04GSypqaR9/BGurVuZ8vgsPlufg8sXLKROpMYrsy2vkqteW0G1p/4hwg1G3gaoCi47AHhsuIVpKz2Uu+tYvCrjrNH2QdL7dwiVmTQwMGgYPH6ZRz7dzBPjuh3dwBgYNCYOs4NZl83ilu63EGGOQBL07XlUQaVGdPF/bV+myHTMwGpvSTZvrfmEg5WFAORXHuLrHYvpn1y7dE2wiFx8Yw9W/t+FPHl5t5BOeXoFjyQKXNazFXnlLh7/PIPLX1nG28syNfcCeq+VPjnM0TF7fgiZkToSLA2JqkDh9kCv1Ukwc3kmw5/7mdd+3ktpjZcar0yV20+V20+1J/D/mcv3c/60X/nXdzuQZYV//vZPPt71MW7ZHRBSdeDyuzhYc5Abvr2BQzWHwlrTyI4tsJmDJcDx1vbO3atQfG5U2Y9r3zrKfplV61iHRaJ/u7ig5wiFaLWS+MD9FL34Eqqqsr+4hndX7OelH3fz0k+7eW/lfvau2hiyPaQp0ywyU99tyefWoWmNvYwzyqQh7Vizv/ToPKgj6J3hoHeAXouo8KMM4SDFxjL/tj+z4qfduuZleWWF/HIX976/jg/vGdqgawrJ2ncCZhsaDEyWOC/NxLSVHp65IPSwwEi0+6z0/h0Mu1cDg9PLK4v3kBLn4Nr+Z19zuUHzxSyaebD/gwxtPZT7froPWQ0WFWXLyiheVIy30Itkk4geEE3L61oiRUggQI3kZFbSVzyeH9gPRFgcbMrfwYy1n1LpqSbaGslFHYbx5Pn3135iUcDePQGHJDIoLZ6cUidfbcrDJ9deg94KGVWFwenxRNvMRNvNRNvMLN5xiOlL9gVlp/ReKw+UOun7jx/okBhJh8QI2idG0iExkvaJEbSNd2A+MtR+xcvgrdZ8jseGW3h+hYf7B1mItYXIUMt++O0NuPLVetd0PP9esIP3V2XXOxTXe1gUzvktm/Wl35EjfKsrO6moCpXeSib/MJmvr/oa6fDfZVt+BdvyK6ly+7GbJVrH2hjZsQU3DExl2g+7NJ8rbGt7q4mRHYNLSusi8sqr+ObTn/ni+UXsqAqo6SOtEVZJQOl9J/1/LuZeJYYxnZpPf3iTF1NFVR4ycis4r0tiYy/ljHJB1yQsJjFITEH4X3TQJ74irBJD2zesFajbJzN9WRbuEBG2utyAPH6FDQfK2J5fSffk8Otnqz1+vt6Yx5qsUsqcXhxmE23jHVw/MIVOLYMdjI5Sui+0ayHwj/OtjJhVw4NDQgudPsJe1tMF/wlJX71/hzrXaWBwClR4KliZv5JSdymKqhBtiWZwq8G0jmzd2Es7Y2w8UMbctTkseHCUUd5n0CSZs30OfiW4MqN4YTFFC4tIuTulll161rQs0p9MRzSJKILKyqjNVIlOohQHraMSmX713+t+QZNA5LBkBOnYtcthMSFq/D70Cp5Im4mJJ/S8R9pMTF+yL+hYPddKu1ni0Us6c2XfNmQWVbOvqIZ9RdWszixhX1ENBZVuUuLs9In389+cdSE3vGEFS1U/ZHwK418J1OWFwZzfsnl/VbaueZMun4/d3s8RTMFCqk4RDciqTJGziF8OLKWytBPTf91HbpkLQQC/rCKJIIkiogC3DUtjTOdEftpxSDNDWJ+1vc0kcs+odF1ip9rjZ/K768joeDmusuDPxC2rIJr4bX8pGXkVDEmPZ/otA7CZm35VWpMSU16/wrb8CsqdPgQB4hwWNuaUc0HXpGbxYTYkkigweWQ6r/2yVzOjE84MhyOEK75sJokxnZMaZP1H0HLrOUI4bkBeWWXm8kz+d0Pfel8rq7iG6Uv28fWmPERBwHmcEJVEeH9VFp1bRnH/+R24tEer4E2Ut+6m1J5JEld0NvHcci/dErUrZG+VfuQ9dRxaQahw/w6iIHBRt5bBT2BgcApsL9nOe9veY/GBxZgEEz7Fd7SR2a/66ZfYj7t63sXQ5KGIQrOoAD8p3D6ZR+Zt5ukre5AUFTrLbGDQWBS7ilmZvxL1BL812SVT+FUhbSa3Iap3IOBmSbSQen8qux/bTcXKCuJGB8quBAR+jFnFhLIL639BEUxxNqLOr23k0D4xQnMYq94KmTQNl970FhH0SI5mw4HyoPvCvVYqqsp1A1OJsZtJjLIypH3tfmW3Tya7xEnh3vXIuRZMcujy+XCCpSj+QKmgpX7XYY9f5rmFO3XPm5QidoMQXHIYjogGcPqdPPLDy/hzpwYH42U40nj19rJMQMUkiZqmWXUhCNoCuS7cPpnr31xJZlENnjBezumVWbmvhJvfWc3H9wwN6SPQVGgSYiqv3MX7K7P4cPUB4JjoV9TAH2Bcz1YUVrnPuQvflNEd+GnHIbbnVx1NAZ8s9Ykvq0nkrpFpDe6g9+aSfZrZtXDdgGRFZX7GQZ6+sgdRdZhRLN9TzJQP1uE5br7B8cgKyIpCRl4FD8/dzLheh/jPtb0xHReFwxYT/MAT+Pt5Nvq/Vc0jw7TLIVPFEvrFuvmtVPtkW9/fwSKJ3Dq0XZM/cRg0HxRV4dnVz/L13q/xKl4UVcHDMXcunxLYYKwuWM2W4i30SOjBaxe+hsN8drpHTlu0i26to7mid3JjL8XAQJMVeSswiaYgVz/nHieKTyF6QO1KDckmEdU7iupt1UfFlEf08kvM2vrFlElAirGSOKU34gl96Z1bRpEab2f3oeDyuHAFT4RV4p5R7TVfeuqYDjw0d1OtwOcR6rtWigKM69WaGHvofYHNLNGlVRRd/FFgkg6LCW3CCZaqgojf48Ichpj6fmsBqqrfUt6SsBRBqv131yOiAWRzNk61CAhdaXREQJnxYxbAp4aXsBAFiLCamHvvsDo/+xN5eO6mgJDSIdw8hxMsf/16K/+5tnfYj2sMGlVMqarKvxfu5L2VWaiqijfEjIQfth/ih+2HePCiTtw3psM5U5ZhMYm8P3kIt76zml0FVbh1Rg/CRRAg2m7mliFpDfq8qqqyr7BG8z49bkBmSSSr2EmvFG2xszqzhLvfXxt2T5bLJ7Ngy0G8ssIrN/Y79n1KGYiauwZBDt2I2jFeZGIPM6+s8dIrSeOEa4ng4cER3L5Y1NUjdgSzJHDbsDTdjzMw0EJVVf609E8syVkSVv290+9kc/Fmbl14K3PGzdFltdscWLO/lG825/P9Q6MbeykGBiEp95QfDXIcj1wtY4o0IUjBeyBTjAlXdu2e3SpTDYJFRPVqXItMIqBi796CuAkdEW3a28H7zuvAX77cqhkUDadCRhQELumhXWlxYbeWdG8dTUZehe7sSKTVxCOXdA7vYFt0wOa8HuoLlqqyj97P/UZcpJ2UeAepcQ5S4uykxh/7f6toG5IoMP3X4EByOEFk0VoQ9Lp6RHRgoSZE6yFkX/1tGz4kLKKAXRQRBDSF7REirBJRVjOfTBlKWovwZ4LmlDr5eWehppCqb/Cz26fw5cY8Hru0Cy0iG7anvyFpNDGlqioPzd3ED9sO1atUj9z/6uK9FFZ6ePrKxh2ediaJtpn5dOow/rdoNx+uzgYI+oEKgN0iEeewMHFQCm8uyazzB3E8ohA4KX0yZSgxjoa1Iff4laAyhSPocQMSBKh0a6fny51eJr+3TlO41PUjdfkUftpeyIe/HeCWYe04WOFivut8bpNfp76f69/GWPkgI0S5gGRmyKhLedycw3++36lLUNnMIjNuH0irEE6OBgZ6eSvjLZbmLtXVyOyVvWRXZvPYksd47cLXTuPqziw1Hj+PztvMv67pZRi8GDRLpEgJf7UfVVaDBJW/wo8psvaWToq1EjegE5VLcvAXu8CvgiggRpqJHJZMxKCWSJF1/xbG9mzNM/N3aIqp+rCbRe4Z1f6YCcSJ70cUmH3nIK6dvpLsEmdYWQtBCLjIzbl7CClxYWbPY9sS2CnVTX3BUjEujYzfjaOgwk1OqZPcMhc5ZU5W7C0mp9RJTpmTshofLaOt5JQFm1GFE0QWxOBgrl4RDSqCeOycX59g8coq/dtEMaF/G95ckklpjRdRCFT0SKKAX1HomBjJ1PM6cEn3VrorZ95blYVyElm6I4jAJ2sONOlxMY0mpl74cTc/bDukszFPZu7aHNrGO7hrZPppXF3TwmqS+PPl3fjjJZ1ZuPUgL/64m4JKDzaTiMNiokdyNFNGt2dwejyCIHBx91bc8s5q3D65zhNghFUi1m7h43uG0jah4Ut6LJIYar64LjcgRVWxhfjxzl2bg18JPgGH8yN1+WT+9+MuluwuZE1WGVf3TcbfZjDWvBW1nivrodpmEKkxIu6/aBhimGww+F6QTNw5Ih2zKPDMgh34ZFWz7vwIFknALAWE1PAO+pxxDAxC4fK7mLV1Fi5/8EW9vkZmj+zht4O/kVmRSfsY7RKd5sa/F+5gUFo8F3c3+hENmjYx1hjMojkoO+Xo6EAwCVSuryRm8LFKDdktU5VRRcvran+3o63ROPol4egX6IVWFRVBZym/zSzx4T1DmPDGyrCDtBAwKBjSPoHfnd+xzuOibGa+emAEf/h4I8v2FKMoKr4Q18sIi0RchIV37xxMx6TI8N+EyQr9b4e1M6COyhOoI1hqjoARD2GWRFLjHaTGa++ZPH6ZrXkV3Pj2b/jkE3vewggiqyagtvGIXhEdeJ5AcDxcwbItv5IXJ/bl5iHtyMitILfMhdPrJ8pmomNSJB2TTs4UyycrfLImJ+iz0DP42e1XmLUiiwfO79hkK9MaRUyVO728vTTzpFJ+Lp/MtB92MWlI23POlMJmlrimXwr55W4q3T6eGNtN87huraNZ9cSFLNpWwJtL9rGvqBqzJB4dDuf1yXQsyeah+6/kgp7JDd4ndQRRFIixmyl3Bp+Y9LgB1Xhk7np3LX3axtG7TQy9UmLonRJDUqSVd5bvD8r+6PmRlrt8pMZH8NKN/YiwmuDA3+D9q0FjA1ovkgUGTT76z1uGpTEoPYEZyzKZn5GPgFAreBBhkRBFgVuGtOO24e1oHXN2lVQZNC7f7/8eQSMaG24js6zIzNk+h78N+9uZXnqDs2xPET/vKGShUd5n0AwYnjwcvxrs5Cc5JJKuTiJ/Tj6iTaz1+zXHm4kdHnv0WJtkY2za2FqP1yukjtC1VTRzpwzj5pm/4fLKQRvjE3FYJEZ3SuTlm/qG5fbmsJh45/ZB7C+uYfaK/Xy2PheBQImggorPrzKsQzz3ju7AsA4JJ7ehHjIF1r0TdHPYwVIU6H19vS9jNUl0TIzSzIOFE0RW/NFIUu1KAr0iGlQUX6yuvZCiqry7MounxvegT2osfVJjaQgKqzyawWS9g5+r3D6qPf46e+cbk0YRU3PX5mg6S4aroAHmZxzkugGhJ3yfzWSX1NA3te5BaRaTyPg+yYzvk0xmUTUHSp3UeGQibSbat4hA/v0UEoq7IImnd8bKTYPaMnN5ZlA/nB43oN5tYphx+0C25FaQkVvOx2sO8MQXFciKqjnUV8+PVFUhq6QmIKQA2g6Fsc/Bwv/TJ6jMDrjlC4is7YbYpVUU067vw1Pju/NdxkEOlDqpdPuIc1jo2iqai7u3NMwmDE4Ls7fOxumv7VCpp5HZr/r5dt+3PDbosabdO+VzwY5voXgPuMrAFgstOkH3K8Fsp9Lt4/HPMnju2t66GqYNDBqLJEcSg1sNZnne8qD7EsclIkVIFMwtwFvoRbSLRPePJvXeVMTjhrAqqsKEThMabE29UmL46eExzFy+nw9XH0BV1VqVL5IAZpNI+8RI7hvTgSt6t9YtetJbRPCPq3ry53HdOFjhpsrtw26WSIqynXobQlwadL0Cdi7QHyw1O2D478Ny8QOIspk0jbDCCSL7SocjJi2oZUKhR0QDqHIEirsNnrwNYe+FfLLKmv2l9R53PF6/wrqsUoprvMiKQozdTN/UuFpl1FVun2bAXu/gZ5MoUuk2xNRRFEU95WyC0ysz/de957CYcnJV3/BFUPvESNon1k6Jl42/gopvvyXqggsaenm1uHVYO2at2A8aBX/huAFFWCSmnteBltE2Wna3cdHhEh1VVZm+ZB8v/bg7yKBH7490X9EJTkUD7ghkmeb/MWCFqtEIfBSTLXDsLV9A6qCQh0XZzNw4OHwbUQODU0FVVbIqs4Ju19vILIkSOVU5dI4Ls9H7TFK6H36bDpvmAELtgZyWSJj/MPS7mTfKzue8rqmM7nx2zipUVZVV+0p4f1U2OWVO3D6ZKJuZgWlx3D4sLWQ5kkHT5s4ed7L+0HrNMt34MfHEjwltLiAgcF7qecTaGia7cISkaBtPjAu0HHy/tYCV+0oorfZgNUu0ibUzoX8KXVqd+oxEm1kiXYfBQdhcPR1mj4VD28MXVGYHdLoEznsi7JcRRYELuibx045DHN8qFE4Q2VfRD2vL74KeM1wRrSoWvCVjAEH3XqjKHRyc1iK/3MUHq7KZszobVSXQG68GMp9ev8KFXZO4Z3R7+qXGYjdLmv1Segc/y6qKowlXo51xMXWoKhBtOBG9Kb+s4sBF41wr9YPAxO+2p3iBjLr0Ugqn/Q+5uhopUkftsU6SY+0MbZ/Ayr3FmnXQ9VqFm0TNHgdBEEL2ZOn9kWrWgfedBKlDDm/WPgRBDETAVRlEU0BEmaww9AEYcDtEGL1OBk0Hl9+FJEhBpUJ6G5kFBKq8Vad1rSfFjm/hiykg+7SDHYeFlbJ2Ng+q78GEt4FeZ3aNpxlFUXl/VRZvLsmkyu0L6o/dll/BB6uy6Zsay2OXdmFgWsMOZDc4vQxqNYghrYaw6uAqPLKn/gccR6Q5kj8O/ONpWlmglO2qvm10BXWbBCYr3PEdfHobZK0IzIwK1dktSIFAaZ+bYNx/wx7Ue4Qpo9uzYm9x0P6i3iCyasVX3h9z7HoEsfb5uz4RHXi8gK+yH6B/L2QNo0rmnWWZ/HfRLlQI6cC4aFsBS3YXMTg9nv9d3we/RppO7+BngYDrdFPljIupCpcPkyhyZHDYEfQqaLNJoMLlO+fElNsnU1LjJTn21MpuTHFxOAYPpurHn4i95upTXpfXr7A2q5Ti6kB9bIzdTL+2gXTvCzf0YezLyyiu9mhO2g6F3Swx+87BIZ2AIqwmTKKIT659stL7I42whPgOJXSAy6fBxf+AXQugPBs8NQGb1aTu0PFCCPP7amBwJjFLZmQNK+CTaWS2SU3MXXLbV/Dl1LAiy6Lqww7wzVSQgB6nfq5rCrh9MlPnrGd1ZmlIE6dAX4vK6v2l3DJzNU+P72Fkx5sRgiAw7bxp3L3obnaW7gzLkVNAwGF2MOOSGbSJbGZC50xhtsOkT+HAKljxCuz7GSRzoAoFIRAsVXzQ/WoY9gC0Prn5RgPbxZEQacFZGnyeqi+I7Cm8AkvkAUy2Yvx1VcacgKqYceXeBkrAk1jvXqhNPfvK/y7ayazlWfW6LipqIEi9al8Jt89ew+jOLVi8s1B3lu4IkgBX9T19/f0NwRkXUyZR1LTL1qugVRVMTfiDPV3kljlpE2tvkC9VzPgrKJ/32SmJqVDpXlEQ8MrH0r2fTurB9a8uodwaha8eQXXE6n3GbQPpW0cTZPsWEWh9DHp+pALQqWU9ZQkWB/S6ru5jDAyaEGbRjMPsoMZXe86b3kZmr+wl0dGEyuMKd8JXtYVU2ktVOH2w/8FIIiyBE8I7G7zMyfDx6x2HS4X8roAAS+wCSdrGPc0FWVGZ8sF6VmeWhD0A0+1TePrbbdjNElf1MzbZzQWrZGXWpbP4+6q/szBrIagEDfKFgIiymWwk2hN5/cLXSYtJO/OLbU4IArQbHviv6hBkLQN3eaACxR4PHc4Hm/Zcy/BfQuD1Sf2Z+NZvulyrAWwmK88Nf4P3Mv9MZkVmvZlJAQFFMeHKuxnZ2eHo7Xr2QgDZpU7eWZbJFb2Tg0a0fL4+l5kaLTp14fEr7DlUjc0kYTdL+rN0hzGbRCaPbNqusmdcTCVEWDRTg3oVtF9Rm3TK73TRECV+R4g8/3wOPvU0vsJCzElJ9T/gBGYszWTaD7tQVfDK2j+w77cV8MuuIvrU5DEroZxPul/CFxvyNIfDWU2Bsr2RHVvwxNiu9YqcwenxRNnMmvbv4f5I7RaJyeeQzb7BucPVHa/m012f1rJX1tvI3DW+K0kO/eeG08byF8GvMYdFhZdXe/nzqDqmxMnewOMnvH0aF3j6mbU8k7X7S3W74bp9Co9/kcGAtLjw5/MYNDpmycwzI5/hwf4P8umuT/lo50d4ZA8mwYSCgl/xM7LNSO7seSd9E/s2WevoJktUy9MWLO2dEstbtw7g3g/Why2obGaRZ6/pxeU9Urio6wfM2jKLD3d+iE/2BRkKWSUrqqoyLHkYlI7le6fAiWWL4e6F4hxmnrqiO99tOcirPy+lS6soxvdJZlzPVsQ6LDy7YEdIIVXXecfjV9iSV05SlA23zxVUnVRflk7CTxfhEF0Kv4ekCU22GuiMi6m4CAvdW0ezObei1u16swljOiWGLP86m8kuaTgxJdpsRF1wASXfLWD3iMspqvLgkxWi7Wb6tY2t06r7+e93MntFeOlel09mg6kFT0R15rMruvPk5d34elM+n67LobTai6yqRNvMXNQ9iVuGtiMpKryyIkEQmDK6Pf9dtEvzRBXOdPZYh5kh6UYvgcHZx81db+az3Z8F3R5uI7PD5OCuXncFPb7RcFfA9q8CfYsn8NhwC8+v8HD/IAuxthCbSVWG7V/D2OfB3rCN+WcKRVF5a2mm5vkuHDdcWVH5YFU2T4xr3tm5c5FERyIP9HuAqX2mUuoupcpbhdVkJc4ah8NsiOOmyujOicybOoxH520mq6RGc+6kIATaGlpEWnluQi+Gdwz0YFslK/f1vY97et/D0tylfLb7MwqdhXgVL9GWaEYkj+D6LtfTwt6CvYXV/JSxDP9J9KbbzCKTR6VzXtckzuuaxDN+maW7i/l2cz7Pf7+TtvEOTedkCO+845dVeraJodzlC9vkAkBAIRonbwnPwrce2PwRTPwwUC3UxGgUa/T7zuvAI59uDsoo6MkmTBnTtFN+p4vsEiftGmjAbl65i1kdL+Tj7VWIWRtQVBVVDUy89skKg9LiuXdMe0Z0aFFrVsRn63OZvWI/Lh3pXq9oYk9hNb//eCMzbhvITYPbclMD1O9fOyCFF3/aDeGXFR/FZha4sI+HL/Z8gc1ko1VEK/ol9UMUzj2RbnD2kRqdSq8WvdhUuCnIiCKcRmarZGVMypjTuUR9bJ4bKMPRYGCyxHlpJqat9PDMBXUEYwQRNn8CQ6eepkWeXpbuKcKlkYkP1w3XJ6t8uPoAf7ykM1ZT04zwGtSNJEokOhKbVvmtQZ30bBPD9w+NZnt+JTOXZ7Jo2yGcXj8qARE1qlMLpoxuT/+2cZqZRZNo4oK2F3BB29Duyx2TIvnTpV1DBpdDYZEEeibHcO/oY+WBVpPExd1bcnH3lji9fsa/ulwzcB7ueUdW4ZedhXx4zxDufm8dlW4fIYqZjmLGRyzVzLX8k1ZCWWCPl70SPrgG7pgf6HNrQjSKmLqoW0vMJhE0LgrhZBPiHJZzNptwoNTJ8A4Jp/w8763cz7MLdgZK9CQLaEQdlu8tZsOBMrq0jOLduwYTYzcjKyrPLtgRUkjVl+5dtqeIXQVVDWKfChBjN/PeXYO5ecZqXScQUfShRmTwQ+kCFpUqiAR6+ewmO7f1uI0JHSc0uK2sgcGZ5j+j/8N131xHuadcs1c1FDbJxhsXvYFJNOH2yXy7OZ+vN+VTUuNBVSEh0sL43slc2TcZh+UMXUbyNx5239LmH+dbGTGrhgeHWEIeg88JBzedhsWdGd5bmaVZ1qzXDfeXnUVc1rNVQy/PwMCgDronR/O/G/ryPwIjDYAGLcu8a2Q6NR4/r/+6N6zeJptZpHvraN69K7TRl8NiIq9c2+xHz3kn8PwCCx8czYs/7uLrdfsRkXFyLPiVO/0u8LvpcO8bXG9fw8Omz/lyYwl3Hel/9buhYDMsehLGPV/va55JGiUEb5JEXp/UH5tZ/8vbzRJv3Nz/nK0Lzi6poV3Cqc1feGXxHp5buAuPXwnZ63QEp1dmW34FV722nEq3j193FeLxa4uWyjVfUrp4BjFDbyDld3Noc99sovqPw7Vn9dFjfH6FmcszT2n9J9K/bRzvTx5MpNWEWcPyOQjBgxSzFnOrz3D6a3D5XdT4a3D6nZS4S5i+aTqXfH4JS3KWNOg6DQzONEmOJN4f+z4JtgQkof5MhICAHYnX5DiSzWk89fVW+v/zR576ZhvL9xaz42AVOwuqWLG3hH/M386Af/7Ek19uobCqfqexU8ZdVufdPZMkruhs4rnlwT1VtXDV/TxNmQOl2mJSjxuuV1bID7E5MjAwODMIgnBa9rG/v7ATr93Un45JkdjNkqZJV4RFItpmYsqo9sy9dxiR1tABMVVV8YQQZrpcuAWodPloFWPjP/3LWBf1MH82fcgQYTudhBw6CzlY8WFXXUzaNIVnzO+SIGiM5fC5YMP74GlaIzsaJTMFMKJjC6Zd34dH520O2x3EbpZ489YB9KnD4e1swSN7WJ67nEPOQ3hlL5GWSLrFdye3zHVKPVPfZeTzRphRiyN4ZZX8cjd3zl6LgEqN5+TLTGQVvtmUz9/G96jzB6yXQWnx/PDwaN5ZlsnctTkAtSK4Jgn8ih/JfgBLwq+YIneHfK4jFrSPLnmUfwz/B2Pbj22wdRoYnGnSYtL47MrPeGH9CyzKWoQoiEGDQC2iBQQY3Gowj/R7COmbaVw2bRFlsl1zPhwcM5CZuzaHhVsLmDtlaP3OmCeJoqg4VTv1TcT7+3k2+r9VzSPD6jCisJyeNZ4JQvWo6nHD9cuKbncxAwOD5sNF3VtyUfeWbMmt4J3lmWzOKafa48dmlkiOsXPHiDQu7t4yLN8BQRAQRSGozwv0nQVV9UIAACAASURBVHcUVeFoZfGKl4n0lXKLaTG3mBYfPSZNqGLqcAvPr3Dy4OCo0P2vggAZn8KgyfWu/0zRaGIK4IreySRF2fi/zzPIL67Eg4hK7Q9PEgXMkkBaQgTTru9DzzanZlfZ1MmtyuWjHR/x+Z7PEQQBv+xHVmVMogkVsKbF8GOOj0vTLsUq1bFh0EBVVf713ck5snhlhR0HK/GFyGTpTffuPlRF/7ZxutZfH8mxdv42vgd/uqwrC7YcZH12GaU1XmxmgWWHvsTjWAbmkrCfzy27+dvKv9Emqg29E09u1oSBQVMgwZ7Av0b+iycGP8E3+77hq71fUe4pR1EVIi2RXNj2Qm7ofAMtI1pSUOFmXM71lPk8YZUG+hWVshov1725iu/+MPKUneJkRWV/cTVb8yrZklfB1rwKtudX8juTjbswY66jQbJjvMjEHmZeWeOlV5LGRkGyQItOp7S+xiRUAEqPG65FEom21X3pr3L7WJddRoXThyAESusHp8efc3MdDQyaM71SYnj5xn6n/DwxNhOlzuDzrp7zjtPn4aFlkxm2w85L+5eEFB9h9b/6nLDyVUNMHc/g9Hh++t1Qvhs/iYUTH2FpTjVOT+CPFmEzc2mPVkwemU631tGNvNLTz7xd8/jP2v8gqzJ+pXYPk18+/G+zm3/99i9e2/gasy+brWsw32+ZpZS7tDci4TiyuLxyyK2V3nRvRYh1NAQ2s8SE/ilM6J8CwNd7v2bF6kWaQz7LlpVRvKgYb6EXySYRPSCalte1RIoIvA+37OalDS8x69JZp229BgZnikhLJJO6TWJSt0khj5k6Zx0Vbj/qCVXgdQVbVKDa7Wfyu+tY9PDosNfjlxX2FgWE09bDwmnHwUoSIq30ahNDjzbR/O6CjvRIjiHe3xNe+RzqSar8bYyVDzJCnV8E6HdL2OtravRJiWXPoSrkE07EetxwRVGga4jr6Y6DlbyzLJP5GQcxS2Ktvg5FVbl+QAp3jEgnvcWplZobGBg0H64bkMrslfsPDwI/hp7zTkpMNN9O+oD8HV+ibFkCSugTeVj9rxW5p/y+GpJGF1MAVT/+RJ+28Yy/ewQA5V99TfWKFaQ83bQazE4n7297n1c3vlrvcDYAp9+JW3Yz8duJzB0/N2xBNWNp5ik5QdUVo9Y7dNl6Bm3tZ26dGVTWBFC8sJiihUWk3J1Sa95O1rQs0p9MRzQF1ri5aDN51XnGRHmDs56dBZXsLKgKKukIy3ZbVTlQ6mRzTrlmKbbXr7CnsOqwaApknXYVVNE6xkaPNjH0ahPNJd1b0iM5hhiHllNTG0gbAft+rnVr1kO1y/ZSY0TcfwkRfGs3HGKa7+/4tuHt+GJjLvKJago982QsDGxXuyrALys8/nkG3205eNS6Wauk8KPVB/hkbQ53jkjj8cu6nrO9ywYG5xK3DW/He6uy0NoFhnPecVgkpo7pQKwtltjIVBAtQOg+2+P7X7slhtgrKj6Q/SA1CRnTNMRU+WefEXfTjUf/LZhMCHL4XvTNneV5y3l146tHe3XCQVEVqnxV3Pn9ncy/Zj4WqQ4Ff5iNOWWagkiXE5SqBupVT0BPutcvqyRG6StRPFm2lWzjYPXBoNtll0zhV4W0mdyGqN6BzZgl0ULq/ansfmw3FSsriBsd2HCoqsrHOz7m0UGPnpE1Gxg0FjOX7cd3wiY63GALgMcvM2NZJv+7oQ+7CqqOiqZt+RXsPlRFSpwjkHFKjuaK3q3pnhxNlE2Hxe2oR+HAqkATsl7MDhj1iP7HNRFKqj088ulmzTkyR6jPDVcUwO3zs2BLAWN7tjraC3HXe2tZu7+s3l5an6KCovLeymyKqjxMu76PIagMDM5yUuIcDGgXx+r9pZq9U/Wdd1QVru53OIhlra/zNUC9/a+iqckIKWgCYsqbnY1n716iLjjmny+YTKi+c0dMvbj+RU0hVV8JmqIqVHgq+CH7B65ofwVVbh/lh+taYxxmok/YpDg1slKgr0TPbBJRFPWUykwSo6x0TArvB3WqZBRlaPZ9OPc4UXwK0QNqR7Alm0RU7yiqt1UfFVM+xceagjVnZL0GBo2FX1b4ZnN+0G9bT7BFUeG7jIP8uL2A9BaR9EgOZJyu7d+Gbq2jiThV05m0ETDiIVjxcp026UGYHTD8D5A+6tRev5Eoq/Ey/tXlFFUFrOlPBgGIdVj494RevPrzXqYv2cufLu3KD9sKWLu/VNfcQJdPZsGWAtonRvLA+R1PbkEGBgbNhhdu6MvYl5dS7vTpGLIRsF9/9aZ+x879se0gjAqsevtfo1rrWMXpp9HFVPlnnxNz5ZUIlmOZFcFsQpXPDbehXaW7OFB5IOj2cEvQnH4nL619i/9n77zDqyjTPnzPzOnpCUmAhJbQe5MqRQVFFEVXxbKooNi77ufu6u66urq71lVRQUVUsGAvWFF6772HEBIC6T2nzsz3x4GQw5kkZyANmPu6vLw4Z2byvuck78zzPs/z+73xfST7csqxHHvd41PolBjOXaNSGdezJVaThEkStOykdJXoWSQRt6qgdUcPPd2b0mi7mWWeMrxycP+EXC5jCjchaEipm6JMODMCd77LveUNNkYDg+ZATf2UuvohAYtJ5If7z6djAyn7Mepxf739qtdDy1CZHTDkHhj954YZTwOjqiqT31tDXrm7RmXFUHBYJT6ZNoQuLSMY2z2Rn7Yf5clvtpFZ6NR8OKqtRw78AdX0hfuZMrx943mNGRgYNAkto2x8dudQJs1cTYnTE7TppoXNLPLcxF6M6Z544sXIVpB8HhxcXuf5Nfa/mh0w5F4do294mnQFVL1eir/5mnbvvx/4hiSh+hpOoKA5MWfnHLxK4Fz1lKABHK08TGXhHhSlNb5q2addR8r469fbeOLr7fz76p7EOixUuIMfPvSU6CFAt1YR7DwS3FcBoZkuV6V7GwGLaEEURRQlcNdVCpfwlftQZTUooPKV+DCFB/5pmMXm5bZtYFDfOD0yosYmh95+SJMkEHqO4xQQBLjwCWgzCBb/B3K2+4Or6uuoaEIRJHYqbenyh39h7jquIUfUoKxKK+BAXkVQ8zfUHfAcJ9pu4tM7h1aZpQuCwPherdiZXcKMJQeCSgdD6ZHzXwe+35LNpPPaNsDMDQwMmhOdEiP46aERPPnNdpbuzQOC7RokUcAiCbSNC+OpCT0YmhoXfKHhD/lN2D0VAS+H3P+qKtD3htObTD3TaMFUfrmbjIIKylw+wqwmkmPshK9fiaVNW6ypqQHHCiYz+M6NMr/NeZuR1cAsnJ4SNABUEG3ZKO7WQdc/7gn1py+2MjQllrxyd1BdvJ4SPYfFxIzJA7j8teUUaUhl1sbxdG9j7mLGO+KxiJYgdURHRweCSaB0QylRg07I7csumbKtZSRekxhwfKIj8N8GBmcbETYTPiU4DNK12YLfEyqiDunteqHTWP9/+ftg3TuQs9Nv5GiNgITuiIOm8cJ3RVxS0pKatQubPzOXpmmWaIca8Fgkkeev6UPXloH3E5+sMGf1oaBASk+PXKVH5q3FaUYwZWBwjpAYaeOdmweSV+bm4zWH+GJDJsVOL7KiEmY1cX7HFtw+ogM9WtdiY5R6Edjj/JUFqs6tN5Mdel8H9ublN9ugdzxVVVmTXsjMJWmsSCvAajpR9+jxKXRx5nH7mOtoo6hI1WyaBZMJ1XdulPlVeCuCXtNbgoagIEi1l7u4vAqr0gpr9IkKpUTPbhaZNqIDSdGOU073XtStcYOS0W1G84+V/wh6XXJIJExMIHtuNqJNDCilNMeaiR524g/VYXJwXZfrGnPYBgaNTqTNTJjFFFTup2ezBUASRVqEN47ADOD3jbpUW/n1gYsKefDTzVw7MDkkg8rmRk6pi9UHCoNe1xPweGSFuaszuLhHy4DXj5S48Ggo9ukSJAIOFVbi9slYTYYHlYHBuUJ8hJUHx3TiwTGn4NsninDzN/D2aP8GWKhdWJIF4jvXuN43JQ0WTOWUupg8aw1ZRc4qf6KTF+6t5jj+clDiuf8uZO5tg0mND2NrVgn7j3g5am1L+61H6JQYTueGqr1vRI6WuNiQUUSpy4tJFGgRbmVoahwWMViFT28JGoigWOos+XD5FEyigElQ0biH1q3IAlw3sA3gT/f++OAI/vZt7ene44bLT13RgyEpGuneBibMHMb4DuP5Lu27oAxg/Ph4pDCJo/OO4sn1INpFIvtH0ubONojmEw9eJtHE6DajG3nkBgaNiygK3DKsPTOWpAX9LYcqu22WBG4a3LbZBC4D2sXSJsbBN5sOc+2xtetMYueRUswmMej70BvwbD1cEvRaidMbsIl5HL09cmZJpNTpIz7CCKYMDAxCJC4VbvsV3r8cxVWCqNRR6WR2QMtecNMXYK7BzLcJaZBg6nCxkyteX16V+quNCo9MhUfm0leXEuuwUOb2ISgKcswATF/6ZWDbx4Vx9+gTQgpNicsrk1FQSYnTi9UkkhBppVWUXfNYVVVZmVbAzCVprE4vxCKJ+BQFAQFJFFBVlYQukUB2wHl6S9BQRQqXbKJ46bK6Sz5MIhZJpNTlRU8vs80s8vJ1fYh2nAj+WkYFpnu/3JhJceXJ6d4UurduWsPlyd0n81P6T8gaoiaxo2KJHRVb47lWycqNXW/EJBoN1gZnN6qqkhxj1/QXgtD6IUVBYPLQdg0xvFPm/os68sTX27m6f7Jm8NCcKXP5qoxzq6M34NHyF7SaRE2lU709coqqVgkfGRgYGIRMQje4ZzUFC17CseUDwiwieKqLfQn+ICqshb/Pqv9kkJpn/3q9PyGWu31cN2OV/6Fah4arV1bJKasmlyia4Vi/z+6jZfz1q208+8MuPrljCKnxjSOrXZ2MggpmrzjIZ+szEQSqGrU9PoXUhHDuHpXKJT1aVt1Uylxepr6/jh3ZpVX17lolFdmH+mNOTEeQTsz9eAla5luZHJ59mC4vdUGukMmek41oFilaXhQQAMiVMkW//xZyjXvreBsOi0ROqSvkEr2nJvRgfK/gniw4zXRvI9ApphN39L6Dt7e+rcvLyyyaSY1KZVrvaQ04OgODpmfP0TKenr+DnFI3wzvGsSGjbs+hk7GaRC7smkByjKOBRnlqDE2JIy7Mwvyt2VzZ98wy7LWaRAROXxREK1MYF27VvCfp7ZEDiDhdyXsDA4Nzk/B4Coc+wc0HxvDTxcWw71eoyPcHTRGtoM/10Gawpr9pc6LeV8DP12VSWOHWDKRCVR7SosIjU+mVuXL6Cr66Z1ijlf65fTKPfbaFX3fmoCiqpjTtzuxS/vzlVp74ZhvvTB5Ij6QoJr6xgswip+bNKuD6Jd0xJ34V9Hr8+Hjyf8hHrpTZdd8upHCJyP6RhPcMp3RDadVxqiJRuiUZ1bc15JKPjMJKvrl3OFe/uRJJ9ZfjOb2BO5d+RRaRdnEO/n55d4Z1bBHStZsrt/e6nXJvOR/v+jikgMoqWekQ2YG3L347JENkA4MzkcIKDy8v2MNP245y/4UduWlIOxRV5boZq9h9tKzGLNXJWE0iHVqE8cqkvg08Yv0IgsD9F3Xi2R92MqF3a8QzKDvVMtKGonEv1RvwxIUHr2GxYRZ6J0ezIaMo4HU9PXKiAON6tDyjPlMDA4PmhUkScKsm6HWN/78zkHoNplRVZeayA5rmf6EqD9V+fahw+7j+7dUsfHRUQMlZQ+Dyytzwzmp2HSmt86Gi4lj26ZbZa2kXG0ZWCIGUHxOewqFY4pYhiIGKc6JNpMWlLcj/KZ/O/+6MFCZRuOTkZmQBT14b3TXus1ekc3X/JP4yvhtfrM/ik7WHKKjw4FMUwi0mhqTGcfv5TV+iV18IgsDDAx6mS0wX/rfxf5S4S3D6nEFlLg6TAxWVP3T6Aw8NeAir1IiN9AYGjYRXVpizKoPpi/YzoXcrfntkFDFhJ9bTT+8Yyh1z1rMho6hGs+/jOCwSPVtH8d6U87CZm2ffzMhOLXjJLPHLjqNc2qt5mT3WRu/kKCJt5qDvQE/AYzdLTB6iXXp516hUHpq3qUr19Tih9shZTRLTRqbU02wNDAzORUyiEKQqeqZRr8HUqgMFlGkYP+pRHoLaM1gqUOn28dGaQw3qvK6qKg98soldR0p1lbu4vAp7cso036tpXnARpvD9iNZsBDHwpmbvYCesaxj5P+eT+IfAXilVMeM6MhHRgu4a99925rLwsdFE2sxMPb8DU8/vEPIcz2TGp4zn0g6XsiFnA+/veJ9dhbuo9FZikSwkOBK4seuNXNrhUmym5tfgaGBQHyzak8u/5u+kdbSdT+8Yopnlt1skPpgyiIW7c5mxJI1txwQMjm8qWUwiItC1VSR3jUplTLcETM1EdEILQRC4/8JOvLJgL+N6tmw00/DTRRAE7hiZwgu/7AmqHgg14FFUtUbxjQu7JmAzSUHBFNTdIycASTF2eic3L4liAwODMwSfG/b8SHTmDu5x74bFayA2FbpdDmZtLYLmSr0GUyv2F1RlaKqjR3kolAyWy6fw3vJ07hqV2mANxZszi1m2L18zkDqVcsW65lV56DYcbd9FtB4NylAlXJXAgWcPEDf2hBqeqphx51yKr3QA1qQKXSUfbq/CI2M7Ext2bpavCYLAwJYDGdhyYFMPxcAgZFRVxaeop6yUtz+3nH/9sJOMgkqevKwbF3ZNqDWoEEWBMd0TGdM9kYP5FSzcnUtRhQcViAmzMLpLfJP0r54qY7ol8PKCvSzcndvoFg2nwx8GJPPCL3s036sr4LGYRMb3akWUXbtpWxIF3r55ADe9u0Z3j5zDIjHjj/11nWNgYGBAcSasnQnrZwMQ6angelRYPB8sYfD9g9DvJhhyN8SeGZnveg2m8svdmq+HqjykJ4Pl8sos2ZvLhV0b5qb4ztIDuDW8rk6lXDGkeSk2KjPuwhK3EHvscmzCib4eW7KNiD4RFPyQh6WVHRQrzsxbkCv9mTm9PjAWk8ht558Zv6AGBucy+eV+pcy5qzPIL3ejHutxbBfn4M5RqUzo3Rq7pfZ1taTSy/9+38u3m7O5Z3Qqb09ur1t9rX2LsDM+ey0IAvdd0JHXFu6vM5BsTkTZzUy/sR/3frxRV8AjCJAQYeXpK2svox/QLpa3/jiAe+ZuDMp+1XTdMIuJD28bRMeEM9+2xMDAoBHZtwA+uxkUH8gegGoSO+oJNb/1s2HTHJg4A3pMbIqR6qJe6zLMNWSJqisP1YaeDFaFR2ZDRvEpjbMuCis8/L47N0g6/HhQFDv2bhxdhiFabAiSCUfHwUEBS3VCnpdqwpN/Mb69/8czeQVYVZUEn49En4+Bl8dQsriQ/gfN4GpVFUgdJ3LQ1cRceBslq+aR9fpNZL11K2Ub52PvFPwzGzKjZ2BgcPqUubzc89FGhv9nIW8s2k9umRtF9Xu9+RSVtLwKnvpuB/2fWcALv+xG0ag398kKc1Yd5KKXF+P2KSx4eCS3j0g5p2WsL+3Zkgq3j2X78pt6KLq4qFsi/726Nzazju9OhdxSF499voUNGcHGv9W5oEsCn905lL5torGaREwa9wcBv9DI0JQ4vrtvOP3bxuichYGBwTnN3l9h3mTwVlYFUjWieMHrhK/vgm1fNM74ToN6zUy1jLJpNpKFqjyk1zujoIZM2Omy7mAhZun0jRKPo3deLuxcXOkkTpZ5Nr+QMZl+9b5p3U18tuYgcnw3zfNC8YEBGNjeuAkaGDRX8srcXDNjJUdKXLWK2BwXJXhv+UF2Hy1jxh8HVJUALt+XzzPzdxITZubDqYPPGiGZ00UU/dmp6Qv3M7JzfFMPRxdX9ksiOdbOsz/sYkd2KbKs4KulZ1sFPLLKrztzWLo3nxsGteHJy7rXqLzXKzmKb+4dTlpeObNXpLNgZw7lbh8iAmFWE/nlLn544HwjG2VgYKCfgjT4/BbwOfWd53PCd/dBfBe/aW8zpV63KC/t1Uoz41G9DK1y7yoUrwtV9uFMW0/Roveqjgs1g3WchlKOKnF6NXd69QZFx9E7Lwnt40aNi6LMC6IlH4Q6ovpamPbhehbuzjnl8w0MDBqGSo+P699exeGQ1UDB6ZVZsT+fP32+hfS8cm7/YD1/+XorD4/txCfThhiB1Elc3rsVuWUuVh8oaOqh6GZAu1i+umc4c24bjBRi75yq+n9HPlmbyeNfbtU0Aa5Oanw4/5rYizV/HcOOf45j2z8vYfVfL2JQhzjS8yvrYxoGBgbnGitfD8pGtf9fGQkvlFHhObEmvbvRw+j3KwLP9blh6YuNMcpTpl4zU6nx4XRrFcnmzODyu1CUh/R4Z1hMIknRDaP2IQoCGj6Juo0Sj6PXEyQGf83owYdO7ACut1n5X7t4erzbEgB3wS94ci8B9ItIuLwK9360iS/uHkqP1lG6zzcwMGgY3lqcRlaRU1MmtjbhG5dXYf7WI/y2K4d7LujI9Bv7NVuZ8qbGJIncM7ojry/cx5CUuLpPaGYoisqfv9qqW0rY6ZWZv/UIPZIiuXWY/v63C7smsGhPLmO7nzniHQYGBs0ATwVs/dTfJ3USsgqvrvHw1xG12NCoCuz5CSoKIKx5rtn1btp716hUHvlss6Y3SV1laHqEFATg8j4N4xcSF2bxB1QncSrO8KBvXlbcXC8tDDhfBf7aIg63eGIn0hq3AsUXTvESH6Vrv8FXfBTBasfRaSgxo27BeWB9rYqDTq/Msz/s4uNpQ3R+OgYGBg2BV1b4cFWGpqddKMI3PkWlR+so7hndcJYRZwtX9U/i1d/3sSGjiAHtzqyy5yX78sgpcSHrDLjBv+6/+ts+Jg9pr7tvdnSXBGYtX4OqqmeMeIeBgUEzYPuX1FQI96dhFp5f4eae8yxE22pZVwQRNs2F8x9smDGeJvUeTF3cPZHhHVuwbG8erhDLVKoTqnfGoA6xtIpqmMzUkJQ4zRuVXtW86oQ6LxWRm0y/B7y23malRKOko+j3NIqXL6bFZQ9ja9cHubyAgl/fJPu9+1B9HuIuua9WxcENGUUcLnY2WIbPwMAgdH7flYNPDl4z9aicbs4s5kiJs8HWxrMFsyRy9+hUpi/cx+wpg5p6OLqYsThN04IkVKVZj09h4W79GabU+DDMksjWrBISIq1UemQirCZiwyzN2mPMwMCgicneDN4KzbcGtpYY3d7Eiyvd/OvCWvw9fU7I3thAAzx96j2YEkWB6Tf2Y8rsdWw6VIRTp3cF1J3Bspsl7hqVejrDrBW7ReLagcl8vOZQUClFqEGRFnXNyyQKjEgyk1DsBRx+xRPg/agInCftBHpLIila9DNxlz6IPWWA//yoROLG3c/hN28lvO+ldT54qarKhysP8pfx2oIWBgYGjccXG7JO26cP4JftR7l1+JktY94YXDswmekL97Mtq4TurSPZdKiI/HI3Xlkl0m6mT3IU0Y7m5cWXVVSpWUavJ+Cu8MjMXJJWFUx5ZYUFO3P4dvNh8ss8qKjEhlkZ36sl43u1qioX3XWkDLMkcPVbKzFLApIgIKsqkiBww6C23DKsPW1iHQ04ewMDgzMSZ1Gtbz99gZXh71Xw4OA61ltnwyh41wf1HkwBWE0SH04dxH9/3s3c1YcQBDTL/kSBIPnxurCbRf7QP4nhHVvU02i1mTK8A/PWZWrWpYfiDC/onJsoQGyYhedvGQHmvbg3fEz+Ly/QylLBSrsd9aRgqmKn3f+AVe3GCeDNywAEZFd5nT/TI6ss3J1rBFMGBqdIdrGTLzdmkZ5XQYXHR4zDQu/kKK7sm0SYVd/ymlN6ej59AG6fQkH5qYvTnEtYTRI3Dm7LQ/M2UVjhwSufWLAFwZ/Bubh7ItNGptA7OboJR3qCfbnlWEynrzS7L7ecMpeXGUvSmLMqA1lVqXAH3qNXpeXzt2+2c0Wf1mzNKiYtvwKPT0FRCarc+GDVQeaszmBEpxa8en0/3b/7BgYGZzHW2hVAeyZIXN7ZxH+We+gWX0uW29p8TeIbbMUzSSJPXNadh8d25ptNh3lnWTqZhZX4FAUJSIiy88chbVm8J4/th0tDMgu0myXG9WzJP6/s2VDD5kiJkzmrMvh4zSHN3oVQiHKYefCiTvz3590hmSyaJYEou5n/Xt2bw8VOymxmtjom8HmbPswcJ6EuvN3fgFcNX5mi+YAlO0sRzBZUZ1lIYy11eUOfmIGBAQAr0/J5c1Eaaw8W+n0Gq5Xnfbclm6fn72Ri3yTuHJVKhxZhIV3Tp2ivFXqFb7w1XMcgkI/XHOKNRftrXed/2HaE33blMiTFb2zb1KIeZS4fioYan16l2UqPj/GvLiO3zF3j/I9nST9Zl1nn9fyBqMqyfflcMX05X90znCi7OaSxGBgYnOXEdQSTDXyuGg/552gb/WeW8+jQGoQoRDPEdWqgAZ4+Db595LCYuHFwO24c3A6A7BdfxhIeRou77gTgzpGpvPr7Pt5feRBFY3cMIMwqYTNJPHBRR24e2r5Bml9LKr08/NlmVuzP9/tznEIgJQoQbjPxybQhdGsVSc+kKJ7+fif7csrwKgont0PYRPD5fITZ7ZQ4vdz/6SYEQFZVPD6FEZ1akGNPQUAEAk+WwqyaD1iSPRLV60Gwh+YFYhKNWncDg1BRVZUXftnD7BUHa9wAOp6F/3xDJt9uzubNm/pzQdeEOq8dU0NJmR7hG5Mo1HgdgxO8sWg/0xfuq3PDTDkmK74yrYBrZqzki7uGNWlAZTdLiBpSs3oDbp+skl2DiMXp4PYpZBZWcvOsNXx+17Bz2iDawMDgGL0nwcJ/1XpIx1iRST3MvLbWQ68EjXVDlGDAzQ00wNOn0XPxljA7SuUJrwqTJPLoxV24/8JO/LLjKLOW+zNYbp+CXYKkw/t58OFrGNU5oUazwdMlp9TF1W+uILfMHVDqESomUcAkCnRrFclrN/Srqhs/r30s399/Pvtyypi1PJ3l+/Mpd/kwSQJWk0RhhRuzT6XY6c8OeeVA2cgV+wu45JV8rJ2CZaJZuwAAIABJREFUb/j2jmH+B6w9KwnrNqLqdXN8O0BFsoWWDo0LNx68DAxC5d8/7WbOqoyQMumyAk5F5u6PNjBz8kBG1WIS683NZVhFJhsUM24xcEdfj/CNWRIZmto8pWObCz9uy+b1hftCqho4jtunsD+nnHs/2sisW89rwNHVTlK0HVkjM6VXaVYluFSvLiXAUI/1yCr7csv5YVs2V/VLPuW5GhgYnCVEJELqBbD3F/yrjzZ/H2VlztYaqqVa9YXYlIYZXz3Q6MGUYLejFAY3o1lMIhP6tGZCn9YBr+8b/R/a2S9vsECq3O1j0sxVHC11n9IuncMicVW/JKYM70DHBO0AplNiBP/5Q++qfy/YmcP9n2z038zFmkshjpcOmSs7IDrSAt6ztjhI1PBJFP42E9HqCFDzk8LjqNy7EnuH/rU+eDksEpPOa6N7zgYG5yK/7DhaYyBVlwfU3XM3sPDR0bSMOqFW5MvLo3TBAsp++hnXnj2MveAiXjcNPzkJDYQufJMUbWs2/T3NEVVVefr7XTUGUrV+jz6FlWn57MguaTJ/vm6tIkiIsHKwINA893SUZiF0JcBQj630yLy1OM0IpgwMDPyc/wikLwGvs+ql6l6qAG2iRFxPapjMmx0w8rGGHuFp0ejBlGh3oDhDd1G39+tL5aZNWNq1a5DxvLloP0dO0bPDbhZ586b+jO5SdwnPcTZnFp8IpELElT8Ke3IWgniiQV2Q3MRdnIJo+yNFi2b5faYsDhydhtBiwp9wpq2rW4ZdhYl9k0Ieh4HBucwrC/ZqBlKhekB9uOogjwxsQdmCBZT+9DOuXbsIHz2a2Cm3EjZ8OKLVyvh5m/l282FN8Zq6hG8cFom7DI+pWlmVVkBZDX2ioXyPHlll1rJ0Xp7UtzGHXYUgCNw1KpWn5+8MEnU6VaVZPUqAeo7NLHSy/XAJPZMMY3gDg3OetoNhxGOw7KUqpeqQMDtg4FToNLbhxlYPNH4w5bCjVjrrPvAYjr59cW7eTPTEifU+Fq+sMHf1qZtkOr0KM5cc0BVMPfn1tlPYFe2GqpgDgikAS+xyIvreQ0SfcUHXquvByywJTOynX3HMwOBcZGd2KQcLgn0yQn249PgUPly4i8v+/l+iR44g9pabCTv/fERrYLPtI2M7s2BnDuXuYKf42jCJAq2j7Vzeu2GMzM8WZi7V9mgK9XuUFZUfth3hqSt7EGlrGoGFK/q25tkfd2m+V9e6r4UeJUBdx/pkFu3ONYIpAwMDPyMeBUWGFa8EZKhqxOyAAbfCxbX3WzUHGr07VLTbUZyhB1P2fv1wbtrcIGP5dUeOZkbq+I01duzdOLoMQ7TYECQTjo6Dg8olNh4qIrMwtCh7b04Z+/O0JctL135N4e/vEDXkOpLvm0vS3bOJ6D8e5741gIjryDWoSuDNW7IdxRy7EgR9UsiiAPHhVh4f10XXeQYG5yofrDyo2U+p5+FSMVs4NPMzkl58gYiLLgoKpADaxDr4YMp52JH9qeMQMEsCLSKsfDJtSJOrzTV3NmRo+5To+R4tJpGd2aX1PbSQcVhMzLrlPGxmfbfvmrQg9CgB6jlWUSGvXFvu38DA4BxEEGD04zBpLiQNRJFseDlpLREkvIKVo45OcPU7MO7f/vOaOY0eTAk6gylb1654MjORy+v2TdLLlxvrxyRzwc6ckI6btTwd78mSfoQWvMnlXXHnXBYUUFnjf8YctTnkgMosCcRHWJl359BmZ0hpYNBc2X20VHPjRZcHlAoHS+q2Iuiw9Ade2/8VLcIthFlrvq4g+Ev7urWK5McHRhAfUYOkrEEVTo31HvR9j6oKpc6mtZQY1CGWtycPxGGRQnrOsJlEUhPCsWpEVNWVAOtCz7EGBgYGmnQcA9N+Z/7QeayKuRLaDoPEntB2CPS/mYIbf2ac6zkqUoKrrporjVrj5ZMVclUrB30WKKggLtxKeB1lZoLFgq17N1xbtxI2bFitx+olr6yeTDIrQgtkFu7KDZJHh9CDN2/xEKxCNLakzxAFqPRVIghgbfkVovUo7rwxoIqg2oLOPX4TPb9TC164pg+xYUYgZWAQKjWV3emRpJaVun3dypcto+Cddxj1ySesbtWK33fnMmNxGjuPlAbITHt8Chd0TeCOkSn0axPdIHYRZyOSKCBrZBj1fI+CAOZmIPk9snM83947nJd+3cOiPXkIgOukkvUwi4TFJDJleAduGtyWwc/9HnQdPUqAeo6VBEgwAnwDA4Ma+OFIOONGPgMnCdW0BIZ02MCXG7O4eWj7JhmbXholmDpuhDt3dQYer4zYahy8ugyvrDCiYzx3jEphcIfYGh8IrH36sm71DtTYjsiKSrTDTM+kqNMuaalJvU+vZ4ccoklmTQ9kunZFK3vw1x7zsEbtZNb2WaSXpGMRLZC4BU+LjVhcA5GKx5FfYsHtUzCJAtEOCzcObsMfB7cjITI40DIwMKidmnoLdT1cihBhq3nJde/fT/bjfyZ5+utYkv3CMJf0aMklPVpyuNjJkWInlR6ZcJuJDnFhxBgbIrqJsps1S8/0fI+yohIf3jyChE6JEcyYPJCCcjefrDvEkj15lDi9mCWRVlE2rj+vLRd0TUASBVRVJcJmoqgyMKDXowSo51iLSeTCromN9lkYGBicOfhkhVVpBTwzsafm+1PP78DjX27lj4PbNZiad33SoMGUx6fw+Jdb+XHbkUAjXMkCx8otFu3JZXV6AS3Crbx360A6JpyQSswrc/PxmkPMLu2G2+3F9OkmUP0q9Yqqcu2AZKYM70D7FmGnNL6YMO0G4oYyyaxp81jXrihgkayMTxnP+JTxFLoKKXYVI6syUdYo4u3xVUGpoqhnxC+hgUFzp0vLCHYcLg3y+NHzcGkzQ/tjHnQn4ysqIvPue0j4vz/h6N8/6P2kaDtJ0fb6ndQ5yLUDk3l3WXqV7cRx9HyPkTYzPVpryPc2IXHhVu67oBP3XdCpxmMEQeDWYR14c/H+INElPUqAoR7bLi6M7s3sczIwMGgebMkqoXW0nYQI7Q3+89rHEGaVWLQnl4u6JSIrKmUuLwICETZTs3u2bbBgyuWVueGd1ezKLq3VZV7F70mRWVjJldNX8PG0IfRpE82s5Qd4/uc9ALh9+P2YXIGZnY/XHOLTdZn8oX8yz0zsiaTzwx3XoyUbM4qD5I71mmQO79gieF6qSl6Zm2KnF1GAGIeFSJspSM4W9BouCsQ4TgSBsbZYYm2xmkc2t182A4MzlVuGtmf+liOa0ug1PlwO+wOmyE1Y4hYjWvMQUPjLZpH/7orm2s7XMqnLJOId8ageD1n330/kuHENolpqcILJQ9vx7vJ0zfdCCRLMksDU89ufsWWVNw5uy5uL92u+p0cJMCSZ/lGppzRGAwODs59l+/IYWYuRvX/zpz0v/LKHNxbtZ3NmMSZRBFR8isr5HVtw56hUhqXGNYv1uEGCKVVVueejjeysI5AKOAeo8Mj8cdYaruzTmi83Hq7zXK+igqLy9aYsjpa6eOfmgboCqqv6J9coMRvq7lubWHuA9Gu528fXG7OYueQAeeVuTJJ/PF6fQpjVhCTAySX7eoI3FZX+7WJCnqOBgcHp0zMpiuRYO/tytIVwAh8uVSxxi7DEfQ6AIJ3oqVRUhUJXIbO3z2b29tkMbT2U+xbZCYuOJv7hhxp6Guc8raLsDOkQy4q0As0y77qCBFlRmbcuk55JUQxLDd5Ea+7ER1gZ36sVP207EtRfVV9YJJGuLSO4zJDpNzAwwL9ubjpURG6ZG49PIdJuYtHuXB67pGZF6R+3ZfOv+bsocXo5vlJXryhYui+f9RlFRNrMvDKpL0NT4xp4FrXTIMHUqrQCVh8o0AyG6jLCLXP5+GTtoaCAozacXn/t5T+/38HTV2rXX2oRbjVxRZ/WfLkhS/PnhbL7dvfoE7tv7y1P5/lfdiMKQlUGqnqblKey5ubzUHdFbxjU1pA/NjBoAh4e05lHP9uimZ06gYKt9aeYInYhiDX/vXsUf4C1InM5O5IEPnrwSwSx6UUNzgWev6YP419dRlGlBx23GWxmkRf+0BuzSeJPn2+lb9tonhjfjdZnWPnlv6/uxf7ccvbmlIW82RkqVpNIuzgH708dhFkyfp8NDM5ljvdyzl5+ELdPBvy9m4IA5W6Z95anYzVJnNc+JiC79O6yA7z4654aPVmPU+mRqfTITHl/LS9e24fLe7du4BnVTIMEUzOXpmmWs4VihAvBmZvj1BaIOb0y89Zlcu8FHUnUIbLw4JjO/Lz9KKUu/SaZsWEWlGOZsZ+3H2XJ3rw6v/zaqCt4EwWBW84QZRMDgzOF7YdLWLQ7t0qYIDHSxkXdEujaMrDfY3yvViydv5SvC624pRr6LRO/rTOQqo5PkCl0iExdeg+fT/icCEtE3ScZnBYto2zMu3MIE15fHnJ2xmYWeXJ8Nyb09QuDjOocz1tL0rjstWXcPiKF20d0wGo6Mza5bCaRT6f0ZdrH29icVaJ5rw443iwiIKCoCj5F1VSktUgCgiBwQdcEXrmuL3bLmfFZGBgYNAxfb8ziL19vAzVYZfQ4i/fksSa9kB6tI3nv1vOIsJn5fsvhkAKp6ri8Co99voW4MGuTZajqPZg6WuJi1YHCoNdDdZiviVADsTmrMmpNHZ5MUrSdubcP5oZ3VlPplkPaqRQFf9qyqMLD37/dgUdWNA096xO7WeL2ER1oU0MDu4GBQeh4fArfb8nmrSVpHC5y4vbJHK/6kkR4feE+UlqEc9foVMb3bIlJEsl7801uWzofx5Sn+WRbfnCvpe0Q5uiNmoFU0bIi8n/Jx5PrQbJJRA6IJPGaRKQwCQWF3Mpc3tz8Jo8Perwxpn/Ok55fQbTDQqeEcNYcLESAoCyNKIDVJJEYaeWpK3owuktC1Xt2i8QjYztzTf9knvlhJ5e8spR/TOjBBV0TaJZ4XbDja1jxKuTvJQyYo6j8bB7NDOv17HNG4FPAV630URSgQ4sw7hndkcv7tOJQQSWzlqfzzebDqOoxmXlFxWISuWlwO24e2u6My9IZGBjUPx+sPMi/f9pVZ0B0XDNhS1YJV0xfwbw7h/D4l9s0z6urqs3lVXho3iZW/+WiJumhElS15iBg4MCB6vr163Vd8IsNWfz92+1Bu13OAxvI/eKftH3s65AkwKujuCvIeuMW4sY/VKdAQ6TNxIa/jdVdYrA/t4zJs9ZS6vRqGvmCX0lPVdWQ3Zjr+vJDxW6WmNivNc9d1atZNNoZnNkIgrBBVdWBTT2O0+VU1ieA4koPk2etZX9ueR0le/5S3h6tI3nBtQH5t19oN3s2pvh4lu7Nq2qKVVQVr6xiS5qLKWIHghC4pub/lE/eT3kk355MePdwvEVesudkI5fJdHiiA+IxzyKHycHS65dilZqH7PbZSmZhJVe9uYK3bx5I/7Yx5Ja6mLsmgy/WZ1Hs9CIrKg6LxJCUOO4YmULfEHy8Fu3J5envd5LSIoy/T+hOu7jQFWZVVaW40kuJ04skCsSEWer0X9RxcVj2Eix/GRDAo9HzJ1nYpyaxIPxK8tpfzqpDTkqcXl6Z1JchKcG7vF5ZoajSg9MjE241Ee2w6BZ/qo2zYX061bXJwOBMZ/GeXO6au0F3lZbFJNI6ykZemTvoGbymZIo7c0eApkCYReLtmwdqisLVB7WtTfWemSqu9ODVqAPQ46V0MqGa2oI/Y5RRUBEgsR4KHRMiWPH4hSzdl8fMJQfYkFHkN7oV/PejSo8PAZAJ7aZRVyYtlEBLFMBmlnhoTCemjUgxAikDg9Ok3O3jqjdXklVUGVI2udIjs/lgAbe5Ivl2lj+QAr9h6sjO8RwqqOSz9YfYm5/Dau8e1JMCKdkpk/tNLkm3JRHR278mWeIttLmnDXv/tJeSlSXEjDwhKPPrwV+ZkDqhHmdsUB23T+bejzdy9+iO9G/r/9wTIm08MrYLj4wNvaLhZC7oksCw1DjeW36QiW+s4KbB7bjnglQclppvsRVuH99sOsyMJWkcLXVVbQB6fAp920Zz18jUKo+oU0JR4KtpsOdH8FbWfJzsoRPpdCx/g4qdc3g44r/Me2giUXbtUlazJNYoZ2xgYHBu88z8nTUGUrU993p8ChkFlUHVYXqq2io8MjOWpDVYMFUbjWLaC/qNcKujJxATRYESZ2j9Clrnju6SwOguCeSVuTla4qLS4+Pn7Uf4ZG1myPX1dX35oZYsSqLAp3cMoXdy9CnNx8DAIJD7Pt5IdrEzKJCqbZH3IpIVnsBfFmby+g2BUq5t4xw8dklXvt63i81rTTh9gWtP5b5KFK9C5IDA/ivJJhHRO4LyHeVVwVSlr5Kv9n1lBFMNyLM/7KJ1lJ2pw9vX+7WtJr8g0cR+rXnux92MfXkpT1zWjUt7tgzYCFNVlXeWHeCVBfsQBKqqOLzyid3Y9QeLeDB7E1azxPQb+52acuDPj9cdSFVDULzYlSJm+v6GqF4EaFtuGBgYGGixJbOY7GKX5nuhPPdqbW/qSaYArNifj6yo9ZotD4V6D6ZiHBYskhhwYwC9XkqB6ArEVOpFRSg+wkp8hBWnR+bW2es0A6maHsBUj6vGL19PlK0oKnNXZ/D8NUYwZWBwuqTnV7AqLVhlNJRF3u1T+GVHDjmlLk2Bm0JXIV45eBNHLpcxhZsQpOCF3RRlwpnhDHitwFVwOlM0qIXvt2SzZG8e399/foNm+VtF2Xn9hn6sSivgqe928NGaDJ6a0INOiRGoqso/vtvB5+uz6iwxrfDIVHhkpr6/jpeu7cNlepSqsjfBpjngPfH71f5/ZVR6If3BcMIs/vm/u9HD3K1eFt/qL0uUUKAiH357Cq54TffcDQwMzl3eXXbgmGpfIKejmaC3qs0kipS5vEQ7LPoGf5rUu3bpyM7xAU2sVT+ompdS5d5VKF4XquzDmbaeokXv1XrN6oFYXXhkhbjw+us5+H5LtmaLVOnaryn8/R2ihlxH8n1zSbp7NhH9x+Pct6bWL19XyaIK327OxtNAfiAGBucSs1eko5y0Nh1f5GPH3o2jyzBEiw1BMuHoODigFhv8PZNzV2doXltWZRQ1+O9UCpfwlftQNUoKfSU+TOGB+1myUvsDtsGpcSCvnH98t4M3buxPpE27fK2+GZoaxw8PnM+YbolMens1/5q/k//9tjekQKo6Lq/Co59vYc0BHYH2yungcwe9LKvw6hqPxgnVULyw9TNwa3uqGRgYGGixOr0Qjcd/3dml6lRPpoSCIKA5hoam3oOp+AgrIzvHa3YWRQ66mpgLb6Nk1TyyXr+JrLdupWzjfOydav+A9QRibeMcJNWjotBbS4Jl3ut6AKvty9cbZQv4+9AMDAxOHY9P4YsNWX6j72roWeTdPoUPV2UEBWQAkZZIzBpy6Y6ODgSTQOmG0oDXZZdM2dYywroHChUY0uj6KazwsPpAAb/uOMrSvXnsyymjurCS0yNzz0cbefTizgEG642BSRKZMrwDvz48ktwyF6/+vl8zkKrYuZgjHzzEoZevIWv6ZHI++weurB1V77u8Cv/35VZqE4yqwlkEu+eDRnD/p2EWXlzppthVx3UEAbZ9XvfPMjAwMDhGpVvbYuh0NBP0JFPAL5ATaWu0DqYqGuQn3jEyheX7gqWDoW4vpZoIxdQ2zCJx96jUWq6iD39DXEXQ63U9gNVW0qi3d0wUhTp9QAwMDGonr9yN1nOo3kW+0uOjzO0Las4f1HKQ5vGSQyJhYgLZc7MRbWKAmp851kz0sBMlvFbJyug2o0Oe07mMqqpsPFTEzKUHWLwnzy8WpAIC+GSVllE27h6VyoQ+rfnHd9vp2jKCGwe1bbLxtgi30jkxAot0FM9JWcpQe2jzytxsPFTMgHYxJ18+kH0LQNS+tQ9sLTG6vYkXV7r514W1iEh4K/1lggOnhDxHAwODcxuTJALBz6uno5lQPZkiiBK2Dv0QRBOug5txHdoaVEHSOzn62DgalwYJpga2i+HCbgn8vitHtzyieEw9T2vfrM5ATPAba9YXpS4vZkkM6rGo6wGsti/fmb5BV++YrKiEN0GUbWBwNlHu8qG1vupd5E2iSLlGMJUSnULH6I7sKNgRdE78+HikMImj847iyfUg2kUi+0fS5s42iOYTg1JVlWs7X6t/cucYJZVepry/lt1Hy3B6ZVSVoFLo9PwKnvp+B3/7djtxYRYWPDKqSdVQZUXlvRUHgwIpPb0ETq/MO0vTGDC5DtXwinyQa65mePoCK8Pfq+DBwXX0FFQa/XsGBgahExtm0RSAOx3NBAgtmQIQZpW4qx4TKnpokKd0QRB45bq+TH1/HRsyCnGGEFCJAoRbTTx/bW8embdFdzbGZhZ5/YZ+2Mz157xuMYkoGtvZoTyA1fblmyLjQ46yzZJIdA0StQYGBqHhsEiaddR6F3lZUXHUsMZM7TmVv634G5W+YPW02FGxxI6qWR1NQOD8pPOJszeNe/uZQlGFhwnTl5NT6qpT2v74PSS/wsPKtALGdk9sjCFqsjenDLdGpYaeMlNVhUV78th1pPSYL5WH4kovxU5vwL+H5x5gkixT012jZ4LE5Z1N/Ge5h27xtezgGv17BgYGOrhhUBteWbA36Jlfb3ZJi1Cq2syiyJhuTWOc3mApD4tJ5IOpg3jm+518su4QgoBmlkpUFSySSIfESN6ePIA2sQ7iplq59b21VHpkzQzVydjMIs9N7MWFXev3Zhleg0dIqA9gNX35tuRuIUXZFkngpsFtmyRlaWBwNtEi3FqnME6oi/zenDK6tY4MEjK4oO0FtN3ahv0Fe/GJ+jpg7SY7Dw54UN+kzjF8ssIfZ60JKZCqjsencP8nm/jirqGN3jN1nKJKD6KGVK/eMlO3T+GhTzcR7bAQ7TATbff/P8phpm2sg2iHmc6HUxHXW0EjqD/OP0fb6D+znEeH1iLWZGuaz8rAwODM5LqBbXjp172a74WSXbKaRAa2j2FjRrEukR7wxwHPTOzZZM/LDVo/JokCT13Zg4fGdmLeukzeXZ5OcaUHk+jP+KjA2CiZq9OXcNFzz1edd177WL65dzh/+mIru46U4vN4kU+62Rw3tE2IsPLcVb0Y1gAmXaIocGnPVszfmh2wq91YUbYgCEwe2u50p2FgcM5jt0iM69Ey6G8ZQi8hAFBRuf3D9Xh8Cpf1asXtI1Lo3trvISV5ZJ76OYKH+lgpcMh4ldD87mySjdcvfJ2UqJTTnueZwJHyI3y6+1N+TP+RMq9fLCLMHMbI5JHc3P1mUqK1P4eFu3M5mF+hGUjVZYLu8so89+MuPp42pEHnphe9ZaYC8MMDI2p/YEi+FNY/Vet1OsaKTOph5rW1HnolaNW/WqHLpXWOx8DAwOA40Q4LF/dI5OftRzXX6VCee1++tg9vLknjs3Whq57azCIPj+3MhD467CPqmUZpxol2WLhzVCp3jEyh3O2j1OXDZhKJspsRPW72X/BvvIcPY05KqjqnU2IE39w7nG0/LOLdH3ezJLEH5S4fquov2RnRqQXTRqbQt010g9bCTxuRwoKdOUFfqp4HsFPBahIZ3SWe5BhHvVzPwOBcp6a/ZQhdGMcrq3hlv2LRt5uz+XH7EYakxDF9YlcKHryf6MREvpj8Kw8seYidBTvxKB5NyXQAh8mBRbLw1pi36Nmi5+lN7gzgUOkhnln9DBtzN6KqakCwWemr5Jv93zD/wHw6RnfkySFPBn0mM5akUaFR/h2qgMOGjCIOFzvrVe01VGIcFk0VSL1lplazWPfOa0x7aN0fDq2s9bC/j7IyZ2stAf/A2+ocj4GBgUF1nr6iJ+vSi8gtc+mSKLebRf4xoQeJUXaemtCDpGg7L/26F1EQagyqHBYJVYVnr+rJ1f2T62kGp4ZQm9TqwIED1fXr1zf4II4+9xyi3UHMAw+QUVBBidOHWRJoEW6F5/+FrUtnYm+5pcHHURNjX1nC/pzykEoO6wOLJNIuzsG39w3HUUOpoYHBqSIIwgZVVevoYm/+nMr6NP7VZezJKUOuRyMKq0kguaKAd8LS6PDPvyFI/gzDjvwdfLDzAxYeWohZPFES6FW8pEalMrXXVC5se2HAe2cr2/K2cceCO6j0VqJQdw+tTbLx/MjnuaCtP8BNz69g3P+WBokBKe4Kst64hbjxD9UZjFgkgSnDO/CX8d1OfSKniKyonPfsbxRWBAtDlK79ipI1XxF3yb21VjmIAlzcvSUzJg+o+wfu/QW+mAKeYDXauhGg01i4qWmk0c+G9amxnp0MDJojmYWVXDNjJYUVnpBKsm1mkYfHdObOk8QjSpxevtyQxdtLD1Ds9Fe1gV+MJ8wi8eRl3ZnQpzV2S/1pJdRGbWtTs3hSd192NW8+P5cfn1mAoqj+2nLVb8DbqjSV+9oMYIJPxmpqnA/sZN66aQBXvrGcCrf+hlyLJCII/lp3MQQzMYdFomvLCN6fOsgIpAwM6plZtw5k/KvLKHF6683Yz+1TybRG8/f245gjnsga9GjRg+dHPk+pp5SMkgzKPGXYTDYSwxJJCk+q5YpnF+kl6UxbMI0Kb+gP9i7Zxf8t/T/eGvMWA1sOZG16AaJGBYIeAQePrLJoT26TBFOSKDB1eHteX7g/KCAMtcrBapaYNjLEUtCOYyCxJ2RvBjnYvLdWzA4Y+7S+cwwMDAyO0SbWwU8PjuTv327n1505iBqaCYIAdrNEXJiFv13enYt7tAy6TpTdzNTzOzBleHuyipyUOL0IAmQUVPLmov1cd16bxppSnTTp07qsqPzj2+18viELJXkQXlew4dfB8ET+sTiLfy45zPSb+jOqc3yjj7NjQjhzbxvMzbPWUu7xafrVnIzdLPG/SX3o1zYGQRBwWCR+25XDW4vTyCiowOP1IR/zTJYEv2BHu7gw7h6dyvherTAbohMGBvVOqyg7X90znEkzV1FUGdqu2XFq68vxILLhUDEbMooY2D5QtS/SEkmv+F71PZVzHoP/AAAgAElEQVQzhseWPEalN1gMoWhZEfm/5OPJ9SDZJCIHRJJ4TSJSmH/TzCW7eGjRQyyatIhSpw+fEpzR0ivgUKZxj2ksrh/UltcX7td8L5Qy04QIK/3bRtd6TBWiBDd9Ae9eBEUZoQdUZjtc/xEkNH7AaWBgcPYQG2Zh+o39KarwMG/dIT5em8nhIieSKBBuMzGwXQx3jkqh/7Fn5NoQBIE2sQ6Oh04dE8L50+dbOFLiRFZUTKJItMNcr2reemmyYMonK9z2wXrWphf4d+qEmj+E43Xyd85Zz3+v7s2V/Rp/V7df2xi+v/98nvh6G+szilBVNcgzxCQKmESBrq0ieebKnvRKDlRDurJvElf2TWJndikLP5tOXlRPiE0hPsLKhV0TqxrZDQwMGo4OLcL45aGRzFiSxkdrDqGqqmYvTnVC6ctxeWXeXnYgKJg6l9lVsItDpYdQTyqSzv8pn7yf8ki+PTnAyPjgiwfp8EQHRJN/M8mreFmcuRiT1PHYDTfwOvp9wprOa6pFuJUHx3Ti9d/3n5JS1QvX9NHXH2yLhGmLYN4fIXM1+Dyg1vBzLeEgWfwBWHIIZYQGBgYGIRATZuGu0R25a3RHrpi+nGeu7EmfNiFuCmng9Mh8vzUbWVUZ/p+FWE0SKipeWWVQ+1juHJXCyE7xmuqpDUmTBVNPfL2dtemheVAdx+VVePyrrSRG2RiS0vh+LO1bhPHRtCFkFzv5cFUG320+TJnLh4JKuNXE2O6JTBnegdT48Fqv0711JN2tP8LF4yDp7G88NzBobsSEWfjL+G48cnFnft5+lPlbj7Axo4gCjZ6WUI1VVRWW7Mkjv9zt7/c04MOdH+JRAj9T2SmT+00uSbclEdE7AgBLvIU297Rh75/2UrKyhJiRMYBfmOKdVa9yZfpETJ4WeE7qL9Mr4NDU38vdo1LJKXHx2XodSlWSysvX9WFQh1MI0q3hcPM3cHQ7rHoDdnwNUrXPUPb4s1DDH4KulwW+Z2BgYFCPlLt8hNtOPex4f2U6z/+8BzhRNlh9HV11oICtWcU4rCam39CPwY0YJzRJMLU/t5xvtxzW9J2qW+JW4clvtvPbI6Mae9hVtI628+dLu/LnS7ue2gVUFYoz/KpLBgYGTYbVJFVljAc/95vmMXr6csySyMaMIs3673MNVVX59eCvQWqGlfsqUbwKkQMCM/GSTSKidwTlO8qrgimA/ZWH6GrPRy5rxcnaFXpsKhwWiRsGta3/iepAEAT+eWVP2sQ6eOnXvQgCNRrUh1kkbJLKdNMrDG05/fR+cMuecNVbMP55KEgDV4m/pC+iFUQ3n74DAwODs4vsYicLduaQX+7mSImT7zZnc3GPRHq01udj96/5O/loTUadCZgKj0yFR+aW2Wv536S+jOvZ6nSGHzJNEkzNXpGOT6NXIVSJ28NFTrZmFdM7+dRThU2Ks8j/f3tM7ccZGBg0GjUJzOjpy1FUlRJnaP5SZwMur8zREhflbh9hVhMJEVbCrP7bitPn1JSFl8tlTOEmBCm4DMMUZcKZ4Qx4zWoLI3ziVUz43cXXmw4HKTGGKuCgqjSpD0l1bh+Rwo2D2/LNpsPMWHKAIyVOzJKIqoJXVujfLoa7RqUwqnMC0saj8PmtMO13fwB0OlgjoHXfepmDgYGBgRaqqrJifwEzl6axJr0QEXAdE955c/F+Zi5No22so0ojoC5xuXeWHuCjNYd0V7I9NG8zc8OtjVJ63+jBVKXHx1cbD+M76YYYaikN+B3t312Wzms39GuUMdc7xRkQ3c4vZ2JgYNAskGqosdbTlyPgF5M529mbU8as5el8u/kwoiAgCgKqquJTVC7unsi0kSm0bQGiIJ7c5oQULuEr96HKalBA5SvxYQoPvi15ZA+3j+jA/K3ZmrL2dQk4WCSBawYkN5qEbig4LCZuHNyOGwa1pdTpo8TpRZIEYhzmQCXXAbfCweXw0+NwxWsB1yh2FVPk9vfwRlojibPFNajvooGBgUFteGWFR+Zt5vfduZpZd79Xo8renHKe+Ho7M5Yc4OPbBxNXQwl2YYWHF3/dE6SCCqFVsj36+RYWPza6wdfFRg+mtmSWaDYB6ymlkVWVpXvzGmJ4DYaqqmSVZ1HqLkXIWkV0dCtaqapx4zMwaCZEO8yaWSU9fTmCAHFhZ2+/VEmllzvmrmdLZjFen4KWGOIP247w265c2rew4Y0O/jwdHR0IJoHSDaVEDTpR6iG7ZMq2lpF4TWLA8bIqE2mNpF1kJLcN78B7Kw7qEnCQRIHESBt/Gtcl9Ik2IoIgEOUwE+WooV9JEGDC/2DmKNj2BZ7uV/Bbxm/M2j6L9JJ0zKIZAQGv4iUxLJEpPaZwWcplOMyG4buBgUHjISsqt3+wnjXpBZptPCdT6ZE5kFvOFdNX8OMDIzTXwHnrDqH1lBxqJVtuqZtNmcX0b9uwlWCNHkyVOL2a5rd6JW4rPE0ncauHUk8p3+7/lve3v0+ppxSTaALZjVeRif9qPFN6TuHylMuNG5+BQRNz7cBkpv++v6oc4Th6+nJAODWhgDOA/HI3V05fQW6Zq1ZJeUX1NwXvzanE4UgAS07A+5JDImFiAtlzsxFtYoCanznWTPSwwPJtk2Cidbi/PO+xS7pQ7PTy1cbDIQVUZkkgIcLKZ3cNJdJ2BosrWCPg2vf57bM/8LetL6DgF+cAv+LhcTLLMnlh/Qs8v+55Hh34KNd3vb6JBmxgYHCu8eIve1ibXhhSIHUcr6KSW+Zi6gfr+PLuYQHvKYrKu8vSg+7JeirZXD6Zd5Ye4K0/NqxKaaMHU6KgXd2mV+JWy8CxufHhjg95bdNriIg45WN9ANXu/1nlWby4/kVeWPcCTwx5gokdJzbNQA0MDLjhvLa8/ru2D1AofTkWSeCmwW3PyjI/l1fmxndWk1PqCirRrglZUXHnj8Ta8jsQA32O4sfHI4VJHJ13FE+uB9EuEtk/kjZ3tkE0n/j8LKKFSV0nYT6m4icIAs9e1YsuLSN46ftteGWFSiH4NmYz+/uPLuqawL+v7l1z1ucM4tOSnbwYE47bF+zZVR2nz3+veWn9SxypOMLDAx5ujOEZGBicw1R6fLy/suaqgdpK8ryyys7sUrZllQRYCu3NLdO8np5KNlWFRXtyT31iIdLowVRcuEXT9FavxG24tUn9huvk+bXP88W+L3DXYZZ4/Mb37OpnyavMY1rvaY0xPAMDg5OIC7dyQZd4FuzKQdbYWKurL0cQBG4e1r7hBtiEfLkhi8xCp2YgVauZcWlvLInfapZpxI6KJXZU3Vm8SV0mBb02eUBrBv1lKvsf/y/vpXvYfbQMl1fGJInEhVm4aUhbbjivbY11+Gcaiw4t4qX1L+FWQ6/IcMkuPt71MYmORG7sdmMDjs7AwOBc55tNh2uUAQilJM/tk3ln2YEALYSiCq9mL7PeSja3T8EnK5ikhtvobPSIpE9ytOaHo6eUxiwJzUaVSYs5O+bw+d7PccmukM9xyS7e3vo2rcJacXnq5Q04OgMDg5p4emJP1h0sorDCo1mOXBN2s8QjYzuTFH2aamvNEFVVmbEkTXOHsM6bpGrGnT8GW/wCEPWpHNokG5elXEbLsGCZ+ZL5P+BIac/ll57H2b5ayorM31f+XfN+UrSsiPxf8vHkepBsEpEDIkm8JhEpzP+Q4ZJdvLzhZa5IvYJwS+3+hwYGBganyttLD2gKToRakqeo8MuOo5Q4vUTZa68k0FvJ1hg0ej2KSRK5eWg7rBqlMJGDribmwtsoWTWPrNdvIuutWynbOB97p8BUnigITBnevpFGrI9yTzmvbXqtxhvfvif3seOOHex+YDfZH2QjV5z45XPJLp5b81xADbyBgUHjkRBh47O7hhITZiHUTSy7WWTK8PZMG5nSsINrIjbUYWYcO/ZuHF2GIVpsCJIJR8fBAZtf3sIRqOX9MYuhZ4lsko0+8X14csiTQe+pikLBrFnE3X77qU3oDGPZ4WV45ODPP/+nfI5+fpSW17Wk+5vdSflbCp4CDwdfPIhSrcdAFES+T/u+MYdsYGBwDqGqKocKtcuP9fo0Hio4cZ1ohxlFoxqieiVbKFgksUGzUtBEPlOTh7Tj7aUHNN+ru5QGeiVF0S4urKGGd1p8n/Y9gkZRS/5P+eT9lEfy7ckBDdcHXzxIhyc6IB4LLmVVZuGhhVzS/pLGHrqBgQGQGh/OTw+O4LHPt7A2vRBVBY9G3V+YRcJmlv6fvfuOk6o6Hz/+ee7Una0ssEtbepOmAtKkiIoaS2LUYG9YUGNiSWJ+STQxiUa/iemJXbCgEY2KLXYFadIEkSa9s5SF7Ts77fz+mAF3d2Z3Z7az+7x9zcvZW88dds+d55xzn8Mvzj2BS0Z0a4aSNo1PNxygLEaLY/w3SaF4z4UMycgm1/YBQRMkEIo9XM0mNhyWgzN6nMFlPX/GPz7ewoGi8FDpjqkuzjghmz4bV2C53XjG1n5zbg2e+fqZY8kmjgqWBTkw5wBdb+hK6rBUAJwdneTclsPGn22kYFHBsYmPywJlzFw7k8sGXqbZY5VSDS5W2vKjEhmSJwJF3m87E/pnp+K0W5RUuf8kMpJNgIn9OyZ+UQlqlmAqK83NXVP68fePNyeU4hbCs9g/fPGwRipZ/RhjmLl25rfJJiISufGVBkqZ8fUMDaaUakbZaW5euGE0e/PLmPXFDmYv30VhmZ9AyJDksDG0azq3TOrDxP4dq52fqrU4WFTeABlYhTTvBTx22S38Z/1/eG3Ta99+sTfh5838IT9Tup9NN9s5vLEkxJsfLqM8EORow6QAT8/fRvuSI9x0/jS6BEO1TvZ4vAuEAqw+uDpqeemmUkL+EGkj0iott7ltpA5LpXht8bF7CsBh72FyS3LpnNK50cuslGpbnDar2mHxiQ7JqzgXoM0Srj+1F//+bHNUwBbvZO1JThs3N8GokWbL4jB9Yh/yiv3M+mJHXAGVEA6knp82ir5ZLXPs98Gyg+R586KWJ3rj23BkA76gD6fN2ehlVkpVr0tGEvecM5B7zhkIhBtM2lrrfnWZU+uSgTUnNYd7Rt3DHSPuYOWBleR78wmaIOmudHqlDOKHs9YxJ7c45j3BEE65vtuZxsNbLV55dBEv3jiaDE/rrSeLfcXYLTu+UOVhfsHiIPYUe9SkxwD2dDtlOyo36NktOwW+AjqjwZRSqmFZltDO4+RwjOHgiSSX8wVCdKny3PEVo7vz789iZ9mtbSQbQPtkJyN7NO4cU9AMz0wdJSL86rwT+NV5J5DqtpNczcz0NhGSHDb6Z6fy5u2nMqJHy53DpaC84FgK34pqu/EFiisPeXFYDgp9hY1WTqVU3bS1QAogK81FrM63RMetd0z99pkpl83FmM5jOKfXOZzX+zyGdxzDdc98zdq9hXE1roXnsSriwn8vrDQspK2wpdgIFAcwMeb7ChQEsKdEt5PGGn6ulFINobppQSoOySvduJiQ34sJBijbspwjn82otO3Qrulkp7krLeuQ4uL2yX1JciQ+CsHtsPi/S4Y1yX272fOLXzWmB1NH5vDhulzunr0Kf9BgtwmhENhswnlDO3PD+F4M6Zpe+8GamSUWsfo6K974qgZUsW58BqM3PqVUi3D24E48PX9bVJCTyLj1ZKeNC06svlfk7tmr2Hm4NGoy4NrmJtlb4OX2l1by3LRRDXvRLUSqMzXm82Wevh7ELhSuKCR91Lf3xqA3SNHqIrIvya60fSAUIN3V8u+hSqnj01VjevBENbkQ4hmSl+y0cetpfWLuf/vpfckt9MY9WTuEA6mHLxrGuD4dEr+YOmj2YArAabcY0aMdyS47y355JqWBIE6bhbsOkWhzynBlRA3HgLrd+NKcaVUPo5RSTW5I13S6tUti04HiqHXxjlt3O2xM6p8V8/h788v47JuDUWPi45mbxBcI8cXWPHbklbTYpET1YbNsjMgewbL9yyov99jIujCLvbP2YrmtSkmNHJkOMsZlVNo+y5NFtqfyfUYppRpKdpqbSf07Mu+bA/hi9JjXNCRPCM8de9qA2PcIEeGBC4fQrV0Sf/1oE0FjCFYzeXyyy4bDZvHPy09mQr/GTzxxVLMHU0efQfhiax6je7XHbrdIi9FVeDxon9SeHmk92JxfeXxnoje+kdkjcdhqzrOvlFJN5dbT+nDvnDUx5xGpbdy6224xbXzPahN1PLdoO6bKTO7xzk0CEAoZZizczm+/OziRSzpuTBs6jbV5a6My+nU8tyO2ZBu5s3PxHfBhJVmkDU8jZ3oOluPbe6jH7mHakGltcoiqUqrpPHLJiZz7j/nkFnqrDXZicTssnr9hdLX3CH8wxCvLdvHSkp3YLPD5Yx87O9XFz84ZwIUndW30VOhVNXkwZYJBiufNI++pp/GuX4/xesFup4c7mUtO+w7+s7rhyD5+W9CmDZnGA188UO8bn1JKtRQXnNiFF5fs5Os9BfhqSIMbizcQ4l+fbuarXQXcPLE3I3q0O/bFPhgyvLhkZ1RLZiJzk/hDhleX7+JX554Qc8z+8W5cl3F4HJ6oewpA5qRMMifV/BxxyIQ4t9e5jVU8pZQCIN3j4PXbxnHpE4vZV+CtMWV6RWX+EOf9Yz7tPE4uG5XDVWN6HHt2qsjr57qZy1gXx/O0R0r9/OF/GxjWLYP+2an1vp5ENOmd58jsV9h06nj2/uweylauDAdSAIEAScUFdP3gdbZMOYudN08ncOhQUxatwZzV86xqn3fKnJRJvwf7MfipwZzwjxPoel3XYzPVH+VxeBjbpW3Mn6KUOj44bBbPXn8KvTokx5xwvTZl/hAfrd/PNTOWMvFPn7FgU7h+LyjzxwzOEku7DiFjyCspT7hcxwNLLB489UHcNnftG1fhtrn51Zhf4XF4GqFkSilVWXaam3d+PIEbxvcKJ5dzxVeHB0KGg8XlPPn5Vib+8TOmPbuMffllXPHUEr7enR/Xs1K+YIgjJT4ufmwR2w+V1PdSEtIkwZQxhtw/PMT+hx8imJ9PqKSai/T7MD4fJYsWsfXC7+PbubMpitegXDYXvzv1d3W+8T084eFwIgullGpBUt0O3vzhqUwZlI3TbiUcVBkDpb4guw6XcePzy3h1+S5KygPYY2Q5rZh2PR42Syj2xp4IuDUY13Uc9425D5fNVfvGEW6bm5uG3sSFfS9sxJIppVRlKS4795wzkBX3TuH/Lh7GpP4dGdQ5lYwkR8zMsBWVB0KUB0LM33iQyX+ey8b9RTGfwSpZN5d9z93Jzr9cwu5/Xc3+V36Dd/daDFBSHuCqZ5YQSmCoYX01ybf2Q489Tv6rr2LKvPHtEAgQzMtjx1VXEzh8uHEL1wjO6nkWd4+8O6GAymVz8dtxv2V059G1b6yUUs3A7bDxryuGs+CeyUyf1JuMJEedJi32+kPcN2cNy7Yfjjm2PtG066EQeFzN/ghwo/pu3+/y98l/J8OVgcdefU+Tx+4hxZHCfWPv4+YTb27CEiql1Lecdovzh3XhuWmjmDKoE+WBEPHGN/6QwesPxRwqWLj0DQ5/8hTpY6bS7fZZdL11JqnDz6Vs0xIAQgaOlPiYt+lgQ15OjRr97uPbvp28J57AlFcegnHmls3kBYOVorn3evcmyx5JvGAMgcOH2f/ww3T94x8bu5gN7vKBl5Ptyeb+RfdTHiyPOd4dIjc+Zwp/GP8HDaSUUseFrDQ3d08ZwF1n9ueMP89l66HY9VtNqc29gXBAFSsUSyTtOoSH+bVPbr2T9x51atdTmTt1Lgv2LGDGmhl8dfAr7JYdQfCH/PTN6MsNQ2/gzO5nahIjpVSLsPlAMY/P2xIzMKrpHhFLvMmJSnxBnpi3hcnVZAhsaI0eTB1+4QVMMPZQjX937ca45BrS2QYCFH3wIcF778WWdvylCj+9++lM6jaJhXsXMmPNDFbuX4nNsmEwhEyIMZ3HcP2Q6xndabRmWlJKHXdW7y5gX0HsZ5XiSW1ugkEG+fL4SjIIVnk+Kt606zZLOH9Y5+NuKo26slk2JuVMYlLOJMqD5RSUFxAyIdJd6STZk5q7eEopVcmMBdtijkCI5x5RVSLJiVbuzGd/oTdqIuDG0KjBVKisjPzX34BAPcayW0L+G3Nof+01DVewJmSzbEzsNpGJ3SYSMiGK/cVYWCQ7kjWAUkod156ev5XyQHRjWbyth6VBOJiUgSPkIBij1bK2tOsADptw44Te9biK45fL5iLL0zQtr0oplahSX4A3Vu4hEKr79BcVJZKcyGm32HW4tEmCqUZ9Zqpk8ReIrX6thabMS8F//9tAJWpellikOdNIcaZoIKWUOq4ZY3h/bW7MMfCJtB4eCDkY2DkVR4xEFLWxW8LATqmc0Pn4G7mglFKt3ecbD8V8rjaRe0RFiSYnKi5vmsREjRpMBfIOVTvED+BHe3YzetNGRm/ayO17dld/nCNHGqN4Siml6sjrr/5h4oRaD20WP5kygMxkZ0LJLCwJz2vy5DUj495HKaVU08krKScQqv/0F0clmpwo1d00iYka9yzBYDgfbjX+WdszUxWPo5RSqsXwh0JYArFq54qth7XeLCWciW/OD0/l0ie+YH9h7ZM9Ou0WHZKdzJ4+lqzUxh/CoZRSKnGBoIkZBiR0j6ggkeREvkCInMymmWOvUYMpW1oa1HOYH4CVktIApVFKKdVQUl32mA8VQ+XWw+SB42s8TjBkSE9y0Dk9iXd/PJ4ZC7Yxc+F2/KEQJeWVQ7Vkpw2bJVwztic3TehNukcz1imlVEuV4XFgt4SqaYoSuUdUFW9yolG9Mpussa1Rg6mk4cPrl3wCwG4nefypDVMgpZRSDUJEGNIlndV7CqLWJdJ66LBZ9Gwfbj1MdTu448z+/HByXz7dcIB3Vu/jYFE5BuiQ4uS8oZ05c1A2DptObK6UUi3diB7topJPQOLTX1RVW3Iij9PG9Il96l3+eDVqMOXo1AnPyJGULFxY52OIzUbmNcdnJj+llGrNbjmtDz979StKfNGD/eJpPXTZLa4d1wN7leDIbrM4a3AnzhrcqdGvQSmlVOPo1s7DiB7tWLQlL2pdvD1MibKJkJ3mZlyf9vU6TiIa/cms9jfeQOnKlZjSypM6ftynb1z7uwYOxNWrV2MUTSmlVD1MGZRdY9KIeFKbXzW6R0MXSymlVAsxfVIfVu3KpzRGo1s89wiRGtMvVGJJOOnErBtHYyWQ0Ki+Gn2shGfMGNyDBiHOxGenF7ebTr/4f41QKqWUUvXlsFnce/4gkuowYW6Sw8YVo7uT1QRzgCillGoeE/p24ITOaTjtiYccLrtwYrcMPHHcY1x2i46pLub88FS6ZjTtBOaNHkyJCDmPP44jJwdxueLfz+2m88MPkXTSSY1YOqWUUvUxdWQO08b3TCigSnLYmNCvA/edN6gRS6aUUqq5WZYw8/pT6NYuCVcCAZXbYfH3y4bz2q3jeOjioQzolEqSw6Jqh1Oy00b7ZCc/OqMvH941iZ4d4sgS3sCaJAG7LSWZXq/MZtePfkTZylWY8nKIkXceQDweBOj6j3+QooknlFKqxfvZ2QPpmOLiofc2IIC3mtTmDptgifCDkd34zQWDm3QYhlJKqeaR5nbw1u3juem55azalU95IFjtPIUepw0RePLqkZzatwMA3zupK987qStr9xbwzlf7yC304g+G6JDiYkK/Dpw2ICuheQobWtPMZgVYycn0mDGDsrVrOfzscxR9+CHiiKS1FTCBAPasbDrcdCNp552HldS0XXRKKaXq7rpTe3HBiV2YvWwXzyzYhtcfPBYsGQPGGC4f3Z1rx/Zssrk/lFJKtQwpLjv/uXkMq3fn8/T8bXywNjc89M8AEp6TqlO6m1tP68MFw7qQ5Iwe7TC4SzqDu6Q3feFrIaaGp7pGjhxpli9f3ignDhYWUr55M8HCQiyXC3tWFs7evRHRlkqlGpOIrDDGjGzuctRXY9ZPqn6CIcP6fYUcLvERNIaMJAcndE7DXYdnq1Tb0hrqJ62blKpdQamfTQeKKPT6cdttZKW56ZvVcueVraluarKeqapsaWl4hg9vrtMrpZRqJDZLGNK15bUeKqWUahnSPQ5G9sxs7mI0CJ35UCmllFJKKaXqQIMppZRSSimllKoDDaaUUkoppZRSqg40mFJKKaWUUkqpOtBgSimllFJKKaXqQIMppZRSSimllKoDDaaUUkoppZRSqg40mFJKKaWUUkqpOtBgSimllFJKKaXqQIwx1a8UOQjsaLriKKWaQA9jTMfmLkR9af2kVKt03NdPWjcp1SpVWzfVGEwppZRSSimllIpNh/kppZRSSimlVB1oMKWUUkoppZRSdaDBlFJKKaWUUkrVgQZTSimllFJKKVUHGkwppZRSSimlVB1oMKWUUkoppZRSdaDBlFJKKaWUUkrVgQZTSimllFJKKVUHGkwppZRSSimlVB1oMKWUUkoppZRSdaDBlFJKKaWUUkrVgQZTSimllFJKKVUHGkwppZRSSimlVB1oMKWUUkoppZRSdaDBVCsiIqeJyO6m3repiEh3ESkWEVtzl0UplRitn5RSx5Pjtc4SkfdE5NrmOHdbpcFUDSI3xqOvkIiUVfj5ykY873UisqCxjt8QRMSISN9GPsd2ETnz6M/GmJ3GmBRjTLAxz6vU8UDrp+pp/aRUy6N1Vs0kbKuIrEtgn/tFZFbFZcaY7xhjnmv4Eqrq2Ju7AC2ZMSbl6HsR2Q7caIz5uOp2ImI3xgSasmxKqbZN6yel1PFE66xaTQSyALuInGKMWdbcBVLx0Z6pOjjafSsiPxeRXGBmrJaPiq2jIuISkUdEZKeI7BeRx0UkqQ7nvl5E1otIUaQFY3qMbX4pIociLadXVljeIGWIcb77ReQVEXk+Uq61IjKywvr/JyJbIuvWicj3q+x/U4VrWiciw0XkBaA78Hak1eoeEekZ+UztInKpiCyvcpy7ROStxrxWpVo6rZ+izqf1k1ItmLNM4vQAACAASURBVNZZx1wLvAn8L/K+YhkGi8hHInI4cq5fisg5wC+BSyP10FeRbeeKyI2R8uWLyJAKx+ko4R7BrMjP54vIqsh2i0RkWD3K32ZpMFV3nYBMoAdwcxzbPwz0B04C+gJdgV/X4bwHgPOBNOB64K8iMrxKuTpEjn8t8KSIDEi0DCLyqIg8mkC5vgu8DGQAbwH/qrBuCzABSAd+C8wSkc6R8/wAuB+4JnJN3wXyjDFXAzuBCyJDZ/5Y5XxvAwNEpF+FZVcALyV6rUq1Qlo/Vab1k1ItW5uus0TEA1wCvBh5XSYizsi6VOBj4H2gS+Rcnxhj3gf+AMyO1EMnVjymMaYceB24vMLiqcA8Y8wBETkZmAFMB9oDTwBviYirunKqahhj9BXHC9gOnBl5fxrgA9wV1l8HLKiyjyH8Sy9ACdCnwrqxwLZqzhV1rBrKNQe4o0K5AkByhfWvAPfVVobIvrsT+DwM0Dfy/n7g4wrrBgFlNey7Cvhe5P0HR8tf02ce+bln5Lz2yM+zgF9H3vcDigBPop+3vvR1vL+0foo6r9ZP+tJXC35pnRV13quAg4Qfv3EDBcD3I+suB1ZWs9/9wKwqy+YSHkIJcCawpcK6hcA1kfePAb+vsu83wKTm/v043l76zFTdHTTGeOPctiPhm+gKETm6TICEsz6JyHeA3xBuDbEix/26wiZHjDElFX7eQbglo8HKUI3cCu9LAbdExj2LyDXA3YS/bACkEG7pAcgh3DJcFy8BfwZ+R7jVd44xpjTSfd2Y16pUS6f1U2VaPynVsrX1Outa4BUTflYsICKvRZa9Qf3qoc8Aj4iMBvYT7kV7I7KuB3CtiPyowvZOwtenEqDBVN2ZKj+XEP7DAkBEOlVYdwgoAwYbY/bU9YSRrtfXCA85edMY4xeROYT/gI9qJyLJFf74uwNrGqoMdShzD+Ap4AxgsTEmKCKrKpR5F9Cnmt2rfsZVfQR0FJGTCLfc3BVZ3izXqlQLovVTfGXW+kmplqHN1lki0g04HRglIhdHFnsIN/p0IFwPXVbN7jXWQ5E67RXCddB+4B1jTFFk9S7gQWPMg/Upv9JnphrSV8BgETlJRNyEu14BMMaECN+w/1rhob+uInJ2DccTEXFXfBFuMXAR7goORFpUzoqx729FxCkiEwiPBX61jmVoCMmE/9gPRs55PTCkwvqngZ+KyAgJ6xv5ggPhP/ze1R3YGOMHXgX+RHis9UeR5c11rUq1VFo/xab1k1ItU1uqs64GNgIDCPccnUS4p2w34SDoHaCziNwp4aQSqZGeJgjXQz1FpKbv8y8BlwJX8u1zm0TKf4uIjI7Ub8kicp6En9FSCdBgqoEYYzYSHs7xMbAJqDqnwc+BzcAXIlIY2W4A1RtHuNWj6uvHhMfsHiE8dOStKvvlRtbtJfwQ4y3GmA2JlkHCWWker/mqa2eMWUd4qMtiwn/0QwmP2T26/lXgQcJ/4EWExytnRlY/BNwr4SwzP63mFC8RHhP8qqmcSjXRz1upVkvrp9i0flKqZWpjdda1wKPGmNyKL+Bx4NpIT9IU4IJIeTYBkyP7vhr5f56IfBnr4MaYJYR7+roA71VYvhy4iXBCniORa7mumjKqGogxtY1UUEoppZRSSilVlfZMKaWUUkoppVQdaDCllFJKKaWUUnWgwZRSSimllFJK1YEGU0oppZRSSilVBxpM1ZOIPCsiD0TeTxCRb5rovEZE+jbwMY9dS1Pu21RE5Jci8nRzl0OppqL1U/33bSpaP6m2Tuur+u9bHyLSXUSKRUQnEE9QmwimRGS7iJRFfkn2R35RUxr6PMaY+caYWtPbish1IlI1zWeDEZG5InJjYx2/vhr7+iPnOE1EdldcZoz5gzGmxX4uqm3S+qll0fpJqeppfdUyReoUIyI/T2Cf7SJy5tGfjTE7jTEpxphg45Sy9WoTwVTEBcaYFGA4MBK4t+oGImJv8lIppZTWT0qp44fWVy3PtcBh4JrmLkhb1JaCKQCMMXsIT1o2BI517/5QRDYRnggNETlfRFZFJmNcJCLDju4vIieLyJciUiQiswF3hXWVWhtFJEdEXheRgyKSJyL/EpETCE/ENjbSspMf2dYlIo+IyM5Ia8/jIpJU4Vg/E5F9IrJXRKbV9fpF5FURyRWRAhH5XEQGV9mkg4h8FLm+eSLSo8K+AyPrDovINyIyta7lqFKm7SLyUxFZHSnXbAnPTo6ItBORdyKf4ZHI+24V9s0UkZmRz+WIiMwRkWTC/8ZdIp9xsYh0EZH7RWRWZL/3ROT2KuX4SkQuasxrVaomWj9p/RTZT+sn1eJpfdUy6qtInXIJ8EOgn4iMrLL+JhFZHynHOhEZLiIvAN2BtyOf3T0i0jPyb2gXkUtFZHmV49wlIm9F3tf4Gbc1bS6YEpEc4FxgZYXFFwKjgUEicjIwA5gOtAeeAN6K/OI4gTnAC0Am4ZmnL67mPDbgHWAH0BPoCrxsjFkP3AIsjnSnZkR2eRjoD5wE9I1s/+vIsc4Bfkp4Bux+wJnU3XuRY2QBXxKe0buiK4HfAx2AVUfXR/5YPwJeiux7GfCoiAyq5vrzRWR8AuWaCpwD9AKG8e0s3BYwE+hB+A+/jPBs3Ue9AHiAwZFy/dUYUwJ8B9gb+YxTjDF7q5zvP8DlFco7KHKOdxO9VqUaitZPWj9FaP2kWjytr1pMfXURUEz4M/yAcC/V0X1/ANxPuMcqDfgukGeMuRrYSaSX0RjzxyrHfBsYICL9Kiy7IlJmqOEzbpOMMa3+BWwn/IuWT/iP8VEgKbLOAKdX2PYx4PdV9v8GmARMBPYCUmHdIuCByPvTgN2R92OBg4A9RnmuAxZU+FmAEqBPhWVjgW2R9zOAhyus6x8pd99qrncucGMcn0tG5DjpkZ+fJVxBHV2fAgSBHOBSYH6V/Z8AflNh3wfi/Peoev3bgasq/PxH4PFq9j0JOBJ53xkIAe1ibHfs36LCsvuBWZH3qZHPvEfk5weBGZH3NV6rvvTVkC+tn6r9XLR+0vpJXy3spfVVtZ9Ls9RXke0/Bv4WeX955LNyRH7+ALijhn/LMyv83DNyDfbIz7OAX0fe9wOKCDcO1fgZt8VXWxrTeqEx5uNq1u2q8L4HcK2I/KjCMifQhfAv2R4T+c2J2FHNMXOAHcaYQBxl60j4F3SFiBxdJsDRjCpdgBVxnLNGkdadB4EfRM4ZiqzqABRE3h/7LIwxxSJyOHL+HsDoo93oEXbCrUoNIbfC+9LIORERD/BXwq3C7SLrUyPXkgMcNsYcSfRkxpgiEXmXcIvQ/xGugG6KrG7sa1WqKq2ftH46Rusn1cJpfdVC6qtI7+Bk4BeRRW8CTwLnEe75ywG2JHrciJeAPwO/I9wrNccYUyoiWdT8Gbc5bSmYqknFP+ZdwIPGmAerbiQik4CuIiIVKoDuxP5F3QV0FxF7jArAVPn5EOHhIYNNeAxyVfsI/0Ec1b36S6nRFcD3CHdrbwfSgSOE/wiOOnYeCWfoySTcerQLmGeMmVLHc9fVT4ABwGhjTK6InER4SIFEypQpIhnGmPwq+1X9jGP5D/AbEfmc8FjtzyLLm+talYpF66dvaf2k9ZNq2bS++lZT1FdXEx5u/HaFwMZNeKjfnMi5+lSzb2310EdAx0i9djlwV2R5bZ9xm9PmnpmKw1PALSIyWsKSReQ8EUkFFgMB4Mci4pDww8CjqjnOUsJ/tA9HjuEWkVMj6/YD3SJjhjHGhCLn/Wsk4kdEuorI2ZHtXwGuE5FBkZbQ38RxHfbIOY++HISHjpQDeYRbFf4QY79zRWR8pGy/B74wxuwiPF65v4hcHbl2h4icIuEHQBtTKuE/2nwRyaTCtRtj9hEes/yohB8Ed4jIxMjq/UB7EUmv4dj/I9xC9DtgduTfAZrvWpWqjdZPWj9p/aSOF1pfNX59dS3wW8JDjI++Lo6cuz3wNPBTERkR+TfoK98mwtgP9K7uwMYYP+HnsP5EOBD8KLK8ts+4zdFgqgpjzHLCwyn+RbiVYTORh42NMT7CD/pdRzgF5aXA69UcJwhcQPjBvJ3A7sj2AJ8Ca4FcETkUWfbzyLm+EJFCwmNgB0SO9R7wt8h+myP/r81jhG/yR18zgecJd2nvAdYBX8TY7yXClcthYARwVaQMRcBZhIed7CU87OX/AFesk0s4O8yEOMpZm78BSYRbQr4A3q+y/mrAD2wADgB3Rsq7gXDL7lYJP7zZpeqBjTHlhP/9zuTbhyoTvlalmorWT1o/af2kjhdaXzVufSUiYwg3uPzbGJNb4fVW5NouN8a8Sng44kuEn3maQzgwAngIuDdSB/20mmt/iXAd9GqVXsFqP+O2SCoPV1VKKaWUUkopFQ/tmVJKKaWUUkqpOtBgSimllFJKKaXqQIMppZRSSimllKoDDaaUUkoppZRSqg5qnGeqQ4cOpmfPnk1UFKVUU1ixYsUhY0zH5i5HfWn9pFTr0xrqJ62blGp9aqqbagymevbsyfLlyxunVEqpZiEidZrxvaXR+kmp1qc11E9aNynV+tRUN+kwP6WUUkoppZSqAw2mlFJKKaWUUqoONJhSSimllFJKqTrQYEoppZRSSiml6kCDKaWUUkoppZSqAw2mlFJKKaWUUqoONJhSSimllFJKqTrQYEoppZRSSiml6kCDKaWUUkoppZSqA3tzF6A18wdDfLxuP09+vpXNB4vx+oM47Rad05OYdmpPLjy5Kx6n/hMopVRTKfGX8M6Wd3hx/YscKDuAL+jDZXPRJ6MP1w++nkk5k7BbWi+rtmvrwWJmLtzOe2v2UVwewBhIcdk5c1A2N4zvRf/s1OYuolItihhjql05cuRIs3z58iYsTutgjOHJz7fy7882EzSGkvJg1DYepw1jYOrIbvzyvBNw2W3NUFLVFonICmPMyOYuR31p/aQS4Q14+dOyP/HWlrcQEcoCZVHbeOweHJaD6cOmc9WgqxCRZihp29Ya6qfjtW76JreI//f6atbtLSQYMgRClb8f2gQcdou+HVP4w0VDGdYto5lKqlTTq6lu0mF+DSwYMtz+n5X87eNNFHoDMQMpgFJfkDJ/kNnLdvGDxxZT5PU3cUmVUqptKCgv4Kr/XcWbW97EG/TGDKQASgOlFPgK+Oeqf/LLBb8kZEJNXFKlmsfiLXl8/9GFrNyZT3kgFBVIAQQNeP0h1uwtZOoTX/Dphv3NUFKlWh4NphqQMYZfvv41n67fT5k/dhBVlTcQYsP+Iq6dsRR/UG/cSinVkHxBH9M/ms7Wgq2UB8vj2qcsUMbHOz7moSUPNXLplGp+a/cWcMNzyyj1xfe9BcDrD3Lbi1+yYsfhRiyZUscHHRjegBZuzuOt1Xsp80cHRSXr5lK4bA7+vN1YziQcWb1JHzcVd7fB+AIh1u0r5LlF27lxQu9mKLlSSrVOz617js35m/GHKvf+H5l/hEMfHMJ3wIfNbSNtRBrZl2RjSw4PufYGvby55U3O6XUOI7JHNEfRlWp0xhhunfVlzECqpu8tEO6luvn5FSz91ZnYLB0Sq9ouDaYa0BOfb6EsRoVUuPQNCpb8l/Zn/RB3r+GIzU7ZthWUbVpSqVJ6av5WbhjfS8fpK6VUAwiGgsxaNyuqR+rQe4c4+N5But3YjZRBKfiP+Nn7wl62P7KdXr/qhWUPD9rwBrw8u+ZZDaZUq7V8xxEOFUf32MbzvQXAGwgyb+MBTh+Y3ZTFVqpF0WF+DWRfQRlLt0V3d4fKS8hf8CKZU27FM2AcltON2Ox4+o6m3eRplbYt9gZYtCWvqYqslFKt2sK9C/EGvJWWBcuCHJhzgC5XdSF1WCpiF5wdneTcloPvkI+CRQXHtjUYFu1dxKGyQ01ddKWaxBPzohuBE/neUlIe5PF5W5uyyEq1OBpMNZD/fZ0bc3n5ng2YgA9P/7G1HqPEF+TlZTsbumhKKdUmvb7pdUoDpZWWlW4qJeQPkTYirdJym9tG6rBUitcWV1puicVHOz5q9LIq1dT8wRBzvzlI1VQTiXxvAVi58wgFpZpES7VdGkw1kNyCMsoD0c9KBcsKsTxpiBVf6vPcAm/tGymllKpVbkl0I1ewOIg9xY7YoodT29PtBIoDlZZ5g17tmVKtUkGZP+azTol+b3HaLPJK4kvuolRrpMFUA/HFCKQAbElphEoLMaH4suRUdxyllFKJCYQCUctsKTYCxQFMMDr1c6AggD0l+lFiX9DXKOVTqjn5gyGsGM9oJ/q9RUTwx/h7Uqqt0GCqgbRPcRErbYSr60DE7qB04+K4jpPhcTZswZRSqo1Kc6ZFLfP09SB2oXBFYaXlQW+QotVFJA9KrrTcJjYyXDo5qWp90pMc+GJMyZLo95ZAKER6kqOhi6fUcUODqQR4/UFyC7zsyS+juLxyi+fIHu1IckZ3iVuuZDLGX8nhjx6ndONiQn4vJhigbMtyjnw2o9K2SQ4bp/Xv2KjXoJRSbcX4ruNx29yVltk8NrIuzGLvrL0UrS7CBAy+gz52PboLR6aDjHGVAyenzcnw7OFNWWylmoTHaadrRlLU8kS+twAkO+1kpbqaoshKtUiaGr0WwZBh3sYDPD5vKyt2HMFhEwTBFwzRq72HW07ry/nDOjO2T3vS3I6YczWkjboIK7kdBYtnc+idRxBnEq7svqSNvbTSdiFjuHhkt6a6NKWUatUu6ncR/17176jlHc/tiC3ZRu7sXHwHfFhJFmnD08iZnoPlqNzG2N7dnpM6ntRURVaqSU2f1JsH310f9d0l3u8tbofFjRN6Yek8U6oNa9XB1JESH+v2FVJQ5sdlt8hKdTOka1rc8zgt2nyIH7+8kjJfkJJIRRMMfTsuePPBEn7z5hp+/eYafnnuCdw8sTd/+mBDzEl7UwZPJmXw5GrPZbOE84Z2Js2tXeVKKdUQMtwZTM6ZzEc7PiJE5Xo5c1ImmZMya9w/yZ7E9UOu17n/VKt14Uld+f0762Kuq+17C4AxcOkp3RujaEodN1pdMGWMYdWufJ6av5WP1x/AZbMwgABBY0h127lpQm9+MDKnxjG+b6/ay8/++xXeWhJCHA2yHnx3PVNHdsPU8RlMt93ix2f0q9vOSimlYrrt5Nv4fM/nlAXKEtrPEot2rnac3/v8RiqZUs0v2WXnR6f341+fbqbMH1/CiaNcDuGSkR1JS9InRlTb1qqCqZLyADc9v5yVO/MpDwQJmejseKW+IH/+cCOPfPgNf5l6IucO7RJ1nC+25vGz12oPpCoq8wd5fvEOBnROZV++l2JvgGCckVWSw8Yz151Czw7JtW+slFIqbr3Te/OP0//Bjz75Ed5gfFNP2MRGqjOVmefMxOPwNHIJa1dSHmDOqj28uXIPh0v9GGPITHZy3tDOXDyiG6k6okHVw22n9WH7oRLeWb0v/oBKfJC0hQ8K72fubAeXD7ycqQOmkuXJatzCKtUCtZpgqrg8wPf/vZCdh0tjzvdU0dHK4u5XvqKwLMBlo77tojbG8JNXvsIbY6geQMm6uRQum4M/bzeWMwlHVm/Sx03F3W0wBth2sIQ3bz+Vm59fQV5x+bGeq1iSnTbsNovnp43ixBzNFqWUUo1hTOcxPHXWU9z28W0ETTBqIt+KPHYPHZM68vTZT9MpuVMTljLawaJy/vLRRt5YuRtLpNJzLVsOlrBmTyEPv7eBC07swk/OGkCndHcNR1MqNhHhj5cMo2Oqi2cWbCNkTA2pzgMgBkfGchzZb1MeMpT7ypi5ZiYz18zkvN7ncd/Y+3BYGuCrtqNJgqlAKMCCPQvYkr+FEn8JHoeHbqndmJwzGZet/hlgjDHc8OyyuAKpirz+EPe/vZacTA+n9u0AwNJthzlSGntOkcKlb1Cw5L+0P+uHuHsNR2x2yratoGzTEtzdBgNgibBw0yE+++lpzP3mAI/P28Lq3QU47RbGgAgEgoau7ZK4dVIfzhvWGbcjvonxlFJK1c1JWSfxydRPeH/b+zyz5hn2l+zHbtkxJoT4SvA73AzrMIzrh1zPqV1PxZLmHbq0+UAxlz25mPxSP4FQ7C+2RxsGX1+5m4/W7+c/N43hhM7R6eCVqo2IcM85A7l8VHeeW7yd/yzZiYgQNMFvh8gawZGxAmfmQixnXqX9faHw96b3tr3HrqJdPDHlCZw2nepFtQ1iahiKNnLkSLN8+fI6H/xQ2SFe/eZVXlz/IoFQgPJgOQETwCY23DY3BsPF/S7mihOuoFtq3bPYLdmax/XPLouZSa+mnqSj+men8OFdkwCY9uwyPttwgKqfSqi8hN3/vpb2595J8sDxNZYnO83FF78449hDy7sOl7LlYDFF3gDJLhvd2nnon51a5+tVqj5EZIUxZmRzl6O+6ls/qbZt45GN7CveR5m/hJTXp9Pnhnl0bte7uYsFwL6CMr7z9/kUlPkTeg43zW3nnR9NoHv75h+aWFetoX5qDXVTeSDIe+s38Jv5f6Q86ENspdiSdiGWv9Z93TY347uO5y+n/UWTt6hWo6a6qdF6ppbnLuf2T28/FkRVFDRBSgIlALz8zcu8uvFVHhj/AGf3PLtO53ry862UxQik4ulJAth1uIw1ewoY0jWdzzcejAqkAMr3bMAEfHj6j621PIVlAbYdKqF3xxQAcjI95GQevzc3pZRqbfq360//dv3DP7iyIFD7l8SmcvPzyynyBqICqdoaB4vLA1z/7FI+vnuSfolV9eKy2/hg/78wyStxVPlWdGT+EQ59cAjfAR82t420EWlkX5KNLTk8ysYb9LJw70K+OvgVJ2XptAKq9WuUcQzLcpdx68e3UuIviQqkqvKH/HiDXu5dcC/vbn034XMdKPIyf/OhmD1J+QteJHPKrXgGjMNyuhGbHU/f0bSbPK3Str5AiKfnb8XrDxKqphkwWFaI5UlDrNqH5NltwpHSlnNjVkopVYP0HMjf1dylAGDd3kI2HyipNA0HhBsHD3/yFOljptLt9ll0vXUmqcPPpWzTkmPbhAzsK/Cycld+UxdbtTIHSg+wNHcppsq3q0PvHSL31Vw6Te3EoEcH0fu+3vjyfGx/ZDuhCo9ZeANenl37bBOXWqnm0eDBVG5JLrd/cnvcWZOO8ga93L/oftblxZ7voDqrdubjtEVfRiI9SUFjWLh6J/t+/wCEYj9zZUtKI1RaiAklljpUKaVUC5eRA/k7mrsUADyzYCu+YOX7UCKNg15/kKc+39qURVat0CvfvELVVupgWZADcw7Q5aoupA5LReyCs6OTnNty8B3yUbCo4Ni2BsP8PfM57D3cxCVXquk1eDA1a90s/KHoXpkj84+w6d5NrL15LRt+vIG9z+0lWFI5MCkPlvP4V48ndL5CbyBmb1IiPUkApcYi44T+WFbsj8TVdSBid1C6cXGtxwoEDRkezWSjlFLHhYzuUND8PVP+YIh3Vu+L6pVKpHEwZOCT9Qco9QUaq5iqDfhwx4fHkkocVbqplJA/RNqIyklObG4bqcNSKV5bXGm5Qxws3be00cuqVHNr0GDKF/Tx303/jQqm4u0WNhgW7llIXlle1UNXy2ETYo0MT7QnyeF20e7yyzm1X4eY6y1XMhnjr+TwR49TunExIb8XEwxQtmU5Rz6bUWnbFJedXu11ziillDoupHeH/J3NXQoKymIPD0+0cdBuE/KKY2elVSoeRb6iqGXB4iD2FDtii/7WZU+3EyiuHMAHTID8ch1yqlq/Bg2mPtn5CVWzAybSLQzh9JyvbXot7nN2SHHFfNA2kZ4k4FhP0vSJffA4Y9+w0kZdRLvTb6Bg8Wx2//NKdj92HUVfvkNSv29bC90Oixsn9MKy9OFfpZQ6LmR0bxHPTJX5glgx7meJNg6KhIf7KdWQbCk2AsUBTIw5qAIFAewplXOaSeQ/pVq7Bs3mtz5vfdRkiPF0C7eb2O7Y8vJgOasPro77nKf0zCRW3FKxJ0ksG+5eJyOWHe/2VXh3rq40ztztsI5N3Du2T3vS3I6YadYBUgZPJmXw5GrLYwxMHZkTd/mVUko1s4ycFtEzleZ2EIzx3G7FxsHapuYACIUMqW4daq7qLt2ZzqGyQ5WWefp6ELtQuKKQ9FHpx5YHvUGKVheRfUl2pe1tlo0Md0aTlFep5tSgPVNHvEeiliXaLQxQWF4Y9zmddourxvTAaY++lHh6kiAcAF0aCYBEhD/9YBhuR+IfTZLDxl1T+tMuWSeqU0qp40ZqZyg7DIGas882ejHcdpJd0W2ciQwzB7BZFu1T9D6k6u7c3ufisrkqLbN5bGRdmMXeWXspWl2ECRh8B33senQXjkwHGeMqB07+oJ8xncc0ZbGVahYN2jOV7Ix+Tqhit3DVgCpWtzBAkiMpapnXH6QwMp483ePAZf92KN41Y3vyzIJtMctUW0+S3RKmDMquFABN6NeRBy8cyq/mfI3XHzu7X3SZbVx6Sg7TJ7aMSR+VUkrFybKFA6qC3dC+T/MVwxKuG9eLR+dupjxQ+d6TNuoirOR2FCyezaF3HkGcSbiy+5I29tJK2zltwlVjuuOIkeVWqXhd3O9invjqiajlHc/tiC3ZRu7sXHwHfFhJFmnD08iZnoNVoRHaEoszepxBuis96hhKtTYNGkzlpObgsrkqzS2VcLew2OiR2gMIZzb6aN1+Hp+3hTV7CsK9TwZ8QcPJ3TO4ZVIfTh+YRad0N788dyAPvb0Wr4l/fK4IZCY7+e13B0etu3hENzJTnNz58ioCoRAl5bGH/XmcNkLG8JOz+nPjBA2klFLquJQRSULRjMEUwJVjuvPo3M0x19XWOAiACFeP7dnwBVNtSvuk9kzoNoHPdn1GyFQO7DMnZZI5KbPG/Z2Wk2sHXduYRVSqxWjQYOrcXufylxV/qbSsYrew5bZIGZSC/4ifvS/sjdkt7LAc/GDAD3h39V5+8cbXBEPmWCBTsZdoxY4j3PnySpx2iz9PPZELti1i6651zO55Kt5A7Il3K7JbQrtkJ6/eMpb2En7DpgAAIABJREFUKa6Y20wekMXye8/k43X7eWzeFtbtyccRCmK5XPiDhuw0F7dM6sOFJ3eNOTRDKaVUy3eouJz5/rEcWXYQ9m0jw+NgfN8OZKW5m7wsHVJcfO+kLrz11d64R0Yc5bZbnDEom64Z0aM7lErUPafcw9LcpTEz+9XEbXNzTs9zGNwhuqFaqdaoQSOAdu52TOw6kU93fkqIb28C8XYLA/RM78m8NTYe+fCrWm8kJb4gJb4gtz63jB9uWcKv/vJzTskTHnh3PUVePyUxkki4Is9WTezfkYcuGkqHagKpoxw2i+8M7cx3hnZm26NPcbiwhIwbbiQ9yUH7ZGfMTIJKKaVaNmMMK3Yc4cnPtzJv40HsZij+kMDXG3DYBH/IMKFvB26e2JtRvTKbtK5/4MKhbDlYwto9BXgD8QVULrtFn6wU/vyDExu5dKqt6JLShafPepobPriBUn9ppe911XHb3IzuPJrfjPtNE5RQqZahwbtTpg2ZxoI9C/AGvZWWx9MtnGRg+IFMHlnyTdw3EIByIzza7xyGFNs5/8RszhvWmcVb8nji8y2s3JlPmT+IzRLaeZxcdkp3rhjdnY6pNQdRABxYD4v/DRs/AF8RPfwButs82Fbsh9HTQXrGXUallFItgy8Q4q5XVvHZhgOU+YMYA+VEnsMNhjjaDvfphgMs3prH+L4d+OcVJ1d6VrcxOe0WL944mtte/JIvtuZVm132KI/Txsk5GTx17UjcjqYpo2obBrUfxOzzZ3PHZ3ewu2g3vpAvatgfACEHTrtw2cDLuGvEXViiz+yptkOqzgtV0ciRI83y5csTPujTXz/Nk6ufpCxQFvc+bpubc3Km8PoHoygJRad0LVk3l8Jlc/Dn7cZyJuHI6k36uKm4u33bjdzO42D5vVOw1XeOpz0r4J274eA3EPSBqXIjszlBLOgyHC74O3TsX7/zKdWERGSFMWZkc5ejvupaP6m2LRAMcd3MZSzfcTjuYXRuu8XQbum8dNOYJk3sEAoZ5m8+xF8/2siqXfm4HRa+QAhjwOWwMAZOzMnglkm9Oa1/VquY37A11E+ttW5ae2gtz617jo93fIwgWGIRCAXITMqkl/08MkLjeORizd6nWqea6qZGedDnhiE34A/5mfH1jKgeqliS7Emc3fNsTnTfyOv29VClFa5w6RsULPkv7c/6Ie5ewxGbnbJtKyjbtKRSMOULhvh0wwGmDMqueor4ffMe/Hca+Eur3yYYmVl+52J4ajJc9Rp01wpEKaVaut+9vY4VCQRSAN5AiK/3FPCrN77mj5c03TA6yxIm9e/I17vz6dsxhRO7Z5BfEr7/ZHgcTOzfkR7to7PoKtUYBncYzB8n/pFAKECRr4jyYDmpzlQ8dg95JT5Of2Quh8/xkanTw6g2plGCKRHh1hNvZXD7wTy66lE2528mEAoQrNDDY4mFy+Yiy5PF9GHTOb/3+Zz9t8+jnnMKlZeQv+BF2p97J54B444t9/Qdjafv6ErblpQHeXzelroHUzsWwavXQ9w9agZ8xTDrYrjxY8g6oW7nVUop1egOl/iYvXxXVNpxqH30g9cfYs6qvfz07AFkpTZtYor31uRy3/mDGNO7fZOeV6lY7Jaddu52lZZ1SHFxzpBOvLRkBzdP7MNXu/M5XOLDGEN6kpOTcjJIcuoQVNU6NWoKuondJjKx20S25G/hxfUvsuHwBor9xXjsHnqn9+bygZcztONQIHyT236oJOoY5Xs2YAI+PP3HRq2LZdWufLz+YOLjxoMBmH1lVCDV829FlPph2x0pJDvDQyie/tLHrNV+5l4XaRH0lcAr18DtyxI7p1JKqSYze9lOYg2Ei3f0gwAvfrGTu6Y03dDunXml7C/0ckrPmp85Vqq5XXBiF6Y/v5zH522NWhcMGS4e0ZVpp/aid8eUZiidUo2nSfJ598now6/H/rrGbY6U+nDYLHzByj1TwbJCLE8aYsUXHDlsQkGZP/FgauP7EPDFXBU08PclPn45obqkFSY82eOeFdB1RGLnVUop1eiMMTyzYFtUcqNERj+UB0I8t2g7Pz6jX/2fzY3Te2v2MWVQpyY7n1KJMsbw90828djcLeFn+qrJ+vfy0l38d/luLhrejd9fOER/p1Wr0WLSrVSXB8OWlEaotBATqjmbUTzHqtHCv4WH7MXws3FOHllUTr63hgMHvLDoX3U4sVJKqcZWUOanoMwftTzR0Q9ef5C8kvLaN2wg763J5dyhnZrsfEolwhjDfXPW8MS8rZQHQtT09SsQMngDId5YuYfpLywnFKrLlzWlWp4WE0xleBz4g9F/WK6uAxG7g9KNi+M6jj9oSE+KzgZYo5I82PdVtatHdrFxWk87jyyq4QZqQrD+bQglNsmiUkqpxlfkDcTMxJfo6AebTSiMEZQ1hr35ZezIK9FnpVSL9cyCbbz25R7K/PE3eJf5gyzcnMcD765rxJIp1XRaTDDVPtlJl4zoh3otVzIZ46/k8EePU7pxMSG/FxMMULZlOUc+mxG1/aDOaYk/5FhyAGw1zzv1u8ku/rnUx8GSGoIlESgvSOzcSimlGp3LbsVsCU909EMoRJPN5fT+mlzOOCG7SdOxKxUvrz/IXz7aGDOQKlk3l33P3cnOv1zC7n9dzf5XfoN399pj68v8QV5cspODRU3Xy6tUY2mSZ6biISLcMqkPv3tnXdQEhWmjLsJKbkfB4tkceucRxJmEK7svaWMvrbRdstPGLZP6JH7yoC8cCNVgSJaN8/vbeXiBjxM6VnNjEwuCTdNiqZRSKn4ZHiexRhVVHP2QPHB8rccJhEK0T45j0vcG8P6aXKZP6t0gxyr1BZizcg/PLNjGvgIv5f4QLodF90wPN0/szblDO+uEvyoh76zeF3N5vAldAF5aspM7zuzXaGU0xrB6dwH7Csrw+kOkuu2c0DmNLhlJjXZO1fa0mGAK4LsndeG3b8fu9k0ZPJmUwZNr3N9mCWcNrkNadHc6hAK1bvbb09wMf6KYn4yt5kYa9IWPpZRSqkVx2i3OGpzN/77eVymoqjj6QSwb7l4nI5Yd7/ZVeHeupt3kaZWO0zcrhVCdHsxNzIEiLxtyCxnfr0O9juMLhHjovfW8vHQXIlRqrCz1BdmQW8R9c9Zw35w1XH9qT+6aMkATA6i4PDZ3c1Tjd6IJXWYu2sbtp/dt8N+5Iq+f11bs5sn5W8kv9WOJYIzBsgRfIMSIHu2YPqkPE/p2aBWTXavm1aKCKY/Tzq8vGMRv316b0ISKAG6HxYPfH1K34RBp3cDuqnmiXqBvpsWlgx38Y6mPoVkxzpPZO3wcpZRSLc5NE3rzyfoDUcOS4h394LJbpLjsjHv4U757YheuGN2dEzqn1btcpb4AizbncbjER9CEn/vdkVfC5IFZuOx17y0qKQ9w1TNLWL+vsMZ76tH5HZ9ZsJ3Vewp46pqR9Tqvav18gRDbGmA6m3J/iL35ZeRkehqsbPM2HuTWWSswQJkv9vDdRVvy+GpXPt3aeZh142g6pup3N1V3LSqYArh8VHf25Zfx1PytlMUZULkdFndP6c8FJ3at20ltdhg1HRb8FYI1j9/99SQXL6yOMZTPkQzj76rb+ZVSSjW63h2TsdsEYlThtY1+EIFO6W5emT6W3EIvs5ft4vqZy+jaLokrR3ev0zC5LQeLmbFgG69/uQebJQRDBoPBblmU+AIM6ZLGkq15jOqVidQyFL2qQDDEDc8tY93ewpiTFMdS5g+ydOthfvyflTx+1YiEz6najkKvH4fNivrdSjihixWeziangcr13tf7uOuVVXE1yJf4gmw5WMx5/5jPOz8aT1Za007GrVqPFhdMAdx91gA6ZyTxu7fXRQ1LqCjZacMAD31/KN87uY6B1FEjp4WDqSq235la6eecdAvvvbFaIg0Mvqh+ZVBKKdUolm0/zN2vrGJS/47M/eYgxeW1D+2uyOO08cy1pyAidE5P4s4z+3P75L58uuEALy7ZyQPvrueik7tyxejutU5KaozhkQ++4ekF2wiGDIGoh7nCXwTX7Cnk+meXMaJHO568emRCyZVeWb6Lr3YVxAykStbNpXDZHPx5u7GcSTiyepM+biruboPxBkLM33SID9bmcs6QznGfT7UtDsuKOdy1YkKXeAIqEwwgBUcwXdLqHbyv3p3P3XEGUkcFQobDJT4ue+oL3r9jIk67JntRiWuRwRSEe6i+d1IX3lq1l8fmbWFfgReHLfyH5iv3083m5/bvjeK8YQ300GxqNoy4FlbOqnW4XxSHByb9HJwN102tlFKq/vzBEH//eBMvL9vFQxcNZcqgbNbuLeCKp5ZQ5PXHTEpRkSWQ7LIz64bR9M2qHCTZbRZnDe7EWYM7sTOvlP8s28nUJxbTPzuVK0f3YMqg7KgvZ8YYfvXGGt5YuafWHiNDuDFx6bbDXPL4Il67dVxc9ztjDE/M2xozy1o8yQFKfUEem7tFgylVrVS3PeacnokmdPH5gxRPv4FNvhKc/fri6tcPV9++uPr2w9W/H/Z27eIu00Pvbah2RFNNDQiBkCG3wMv7a3P57old4j6fUke12GAKws9QXTaqO5eeksPB4nIKSv2IgGPNKswzT9BzRAP3BJ3zMBzeBjsWxh9QOTww9BI49Y6GLYtSSql62XqwmDtnryIz2cn/7hhPVmp4GM/gLun8744JPPS/9Xy0bj8iRLVmu+0WBjh9YBa/+M4JdG9fc2NZ9/Yefn7OQO46sz8frM3lhS+2c//ba5k6shuXndL92DMhMxdu542Vic3LUx4IsflAMT/+z0qevGZkrdt/ufMIB4ujh6wnkhzgm9wiNh8ojgoglQKwLOGsQdm8vza3XgldBuVkMurzjwnm5VG+eTPlGzdR/s1GCt/9H+WbNiFOZzi4Ohpk9Q//35ZWeYTQ7iOlfLnjSMyyxtuA8PjcLRpMqTpp0cHUUSJCVqr72I0wNHo4G3+6nlB5OZarAR8atGxw+cvwzp3w9X8xAS9S3XzeNmc4FfqY2+D0e2tNra6UUqppGGN4aelO/vzhRu46sx9XjekRNYSoa0YS/7piOEdKfMxevou3Vu0lv+z/s3fe8VHUeR9/z8z2TUhIgzQISei99yqnqKhYADnLoYjY9bxHPfXOe/T0nrPcnV1s2BCwn2JDUem9ht6SkIQkpPftM88fK5Gws8luSCBl3q8XL2Cn7G92Z+f3+7bP1wkKhFn0TB8Qy7XDuxAZEtwcY9CJXDYwjssGxnG0oIolm7O4/OV1DEoMZ/awRL99eaB+77nDLbPmcCFHCypJjQlVPf4Un21XN9aCEQdwywrLd5/gj7/rGdiFa7Q75k9IZtXhwka3szHrJWYPT8ThljFFRaGLisI6alTtdkVRcBcUeA2so0ex7Umj7PPPcR49ihgS8puB1T2VRVXRqmmHwTgQ0ouqOJhfQa/OZy8q0xjc5Q5c+dUodjeCXkIKM6KPs2q1i62AVmFMnYlotWJMScGeloZl+PCmPbmkgyteRh46n4q/XU5Ykh1BZ6A2ni3g/ffQuTBiPnRMatr319DQ0NBoNEVVDh76NI2TlXY+XjCqQcOjo9XAbRNTGtejsAFSY0J47LI+PDitJ1+n5fHP7w9S7adWKxDvuVtWWLQuk39c1b/e980tt6mmYAUjDuCWFfLKtYaqGv4ZlBhOpw4mVVW/QNrZ2Fwenli+n//9aj+/69OJ+ROSGZgQVms8CIKAvlMn9J06ETL+t5RBRZZx5ebhOHoE59Gj1GzdympHX1ymSJ/3CFZdcFtm6Tk1phRZwXG0jMrV2TiOVyDoRGp9+IqCaNUTOjEBy+AYRGOrXLK3C1rtN2MZNoya7dub3pj6lcrtmVSIFxH+4L8hZyvYSr2RKEsEJIwAvab6oqGhodEs5O6CosPgqARDCER0g4ThDWYA/HzwJH/+bA9XD03gteuHtphicpNe4pqhCSxan6Ga6xCo99wtK3yx8wR/md4bi8E7fSuKQoXNTUGlnYJKBwWVdjJVFrcQvDiAwx14KqJG+0MQBF67fghXvrQWWyNvFfuvdYPf7c3j54MFJHQ08/YfhtebViuIIoaEeAwJ8TBpEgCOp3+GUpvPvsE4EFxumQq7itRnM+GpdFL41h48pXYUp/dzUM74zXmcDsq/zaD8u0wib+iNKTXwGjKNc0frNaaGDqH601dhb4y3vskYClE9IaZXk5y/dNlHRN4yD4whkFK/d0VDQ0ND4yxx1sDez2D981CR63VeyR5v+jUKmDvCmHtg4Bww1fUc25wenvp2P78cLOSlOYMZmezroW4JHMqvVH09uPQ7mXnvbsXulimocFBY5cCoE4kJNRIdaiQm1ITox+gMVhwgKsgUR432R/z+7fx9x2IeG3oDNrfirzCiQWTFG6k6VljF9JfW8tGC0UH1cNNJ6vd8UA4EAYorndQ43bXOiubCU+Hk5Es7kKvdNKSCc8rQKn5vPxFzemHu0zKfb+2Z1mdMOaog7WNC0v6NNSoH5astCIoMoggeN0SmePs99b4cdIYGT7cvt5yDeZVU2l1YDDriO5oZ5CrClZtLyMSJ5+CCNDQ0NNo5J/fBe5eByw4u9agKzmpY+b/w85Nw3SfQxVtbsSennHs/2snAhHC+u288HUz6czfuIHC4PShquXcE5z0XBYER3SIZ3z2KmFAT0aFGH8n0T7fn8Lcv99Y24609NghxAKtBYkyKtmjT8E/lqlXk/eUvXPr6QgbGJPHQZ2nsy61AlhVcDclk+kFWoMLuZs6vUuWdw+rPAlIUBceRI0TWlJOJ728/GAeCgMC3e/NYvPk4MR2MdI8JpXtMCN07ef9OjQnB2gSpdopbpvDNtIAMqTrHuWRKlh4k+vaBGOI0YZiWROsyprK3wOJrQHYjuKoRdIDzDE/fyb2w/F5Y8SjM/Rqiuvucxu7y8HVaHgtXHeVEmR1BAI+sIImCVy3QYWfOhfO4xe4hMqR1fUQaGhoarYq83fDOxV5jqSFOqax+cCWea5eyMCuBResy+NvlfVu8CpdeFFXrmCA477lBJzKhRxRDu0b43Wf6gFge+3Kv6rZAxQGMeolJPWPqvyiNdkvVmjXkPfIoia+9irl/f7oDn98xloyiat5dn8HXaXlU2F24POo3fX1iKwCVdhfP/XCI52YO9DnWmZND9caN1GzcRPXmzYgWC5ePuIQDYndqzlBGD8aBoJMEfvjjBMx6iaySGo4UVHHkZCVrjxTy9roMMoqqiLQa6d4phB6dQkmN+e3vkCCMLNueIjzlDlVDaktOGv/45TUOF2UiiiLdI7vytwvuZlBsb8BrUJV/n0n0zf0Cfj+N5qf1WAoZa2DJLHD55sT64KzyTsxvToZ5P0JM79pN6YVVXPvGJqocbr/NgEHPOzWRvPv0L7z8+8Fc0LtT01yDhoaGhsZvVBXCe1cEZkidjqsG5+LZHI5+leV3X0JcuLl5xteEiKJAqElHhd1XgCIY77nLI9cq2/rDpJeYNSyRDzcfV13MNiQOYNSJ3Dw2CUnUVMQ0fKlau47cPz9MwisvYx5Y19jpFmXl8Sv68fgV/ZjzxiY2phf7HB+I2IpHhuW7c/nbZX0wV5ZRvWkz1Zs3UbNxE7LDgXXUKKxjRhN9//0YEuJJdHt49u8rQUXgJRAHgiQKXDYgjtBfI9vJ0SEkR4dwUd/Otft4ZIXskhoOn6zkSEEVG44W8d6GTI4VVhFhMdRGsHp0CiW1UwjdY0Jqz1fn+ldl16bunU6lo5qbPv0zT114P5f1mozT42ZLzm6MUt0sK0d6GZ5yB1KYlobbUmgdxlTxMVg6JzBDqhbFmxL47nS4aytYIjhaUMWVr66nyuH26yE8hd2jgMfDnUt28K+ZA7l0QMv2empoaGi0Ora84dPTL+n5SmpckHFvCFaDdzH/1g4ni9NcrJprrd3PqLh4Pu5nhPCrz+mQz4arhySwWMXACcZ7nhRpre1ZVR+3TUzh8505uGzq6oH1EWrScf2orkEfp9H2qVq/ntyHHiLh5ZexDB7sd7/skhp2ZPn2fQpGqlxwu3hjwV+ZfuAXLMOHYx01isi5czGkpPjIhRt1EteN7MK7GzJVm2E35EDQSwLzxnfzux28BldSlJWkKCsX9v3tdY+skFNaw5GTVRwuqGRTejEfbDrO0YIqwi362ghW95gQeks6IkvsqudPL8kGYEafqQCYRYmJ3Uao7lu1MY+waUn1jlfj3NE6jKk1z/oYUoFNuIo3SrXtHSpH3MO1b2wMyJA6HbtL5k+f7KZLhJX+CWFNdEEaGhoa7RyPy2tMeXzltz0KvLDZySPj/XteRTxewYqL/+kVIGoFzB2bxJItWaBSph+I99xqkLh9UmAS7p3DTHw4bxSz39iIzekJSBhAFLwLxvhwMzrJVwkxo6iapZuzOFJQSbXTQ5hJR7/4cOaMSCSmg6Zw2xopqXby8bZs0nLKKLe5CTFIpMSEMHt4Il0jrXX2rd64kdwHHiThpRexDPFvSAGsOVKoKr4ZjNiKHYl1/S/g/kVPIUgN1xPeN7UHPx8sIKOoGncQtUhmvcSCCcmNlkSXRIGukVa6RlqZ2ue3TCZZVjhRZquNZG3NLCUnvYIrXQpGfD+c5IhEREHkj988xeW9LmBwfF/CTSrPNreC/UipZky1IFq+MWUvh31fgOKbkhfIhIvbDpte5VNhBtUOj6oh1VDursMl8+8fD/HOTeoeAg0NDQ2NIDn0HcjqUZMHxhh4Zr2DO4YbCDfVk2omCLD7IxhxSzMNsmnpGmllYGI424+X4lFZ7DXkPZdEgWn9Ovvdfib9E8L4751j+f2bm7A5PT6CFKdjNUiEWwx8eMtIXl9zjBve3sy7N40gzKznl0MFvPTTEVVxgTVHinhl1VHGpURx9wWpDO6iSTe3BvbnVvDSz0f46WABIr9JlAPoDxXw9roM+seHcdeUVCb1jKF60yZO3P8nEl58AcvQoQ2ev6zGhVMlQhSM2ApAOfqADCkAs0Fi6a2jmPX6Rk6U2lQjVD7H6CXmjEjk3qm+9fVniygKJEZYSIyw1JaLVPycRcWPx9X8KYQarXx+3cu8unkJD37/LIXVJUxOGckz0x4k2lq3RlJWSRfWOH+0jCYc9bFrqVciV4UHxhh4boODMnsDspIuO6//clC1I3zFli8o+elNwkbNIuGuxcTf/g6hQy7BdmTzb8cD648Vc7JCPTSroaGhoREk2Zu9mQMqDIuTmJSk47kNDTSNddVAxupmGFzz8cK1g+hg0qn4pevHpBd548ZhGHWBLSxP0aNTKBsfvoBnZw6kf3wYJr1IqElHiFFHqEmHUScyPKkjL84ZzJoHJ5MUZeUfV/ZnUGI4176+kce+3Msdi3ewI6sMh1v2UWlzuGWcbplfDhUw581NLN2SFeSVaZxrvtx5gqteW8+Kffk43XIdQwrA5VFwuGW2HS/l9sU7+N+3fibn/j8R/8LzZ93b83SxleYgKsTI8rvGMWNQPCadiFmv/nuxGiRiQo08fkVfHrusr0/aYHMhSCL1/fi7RyXxn0sfYeudn7Fy3rucrCrmf396yfc8Wj1ji6LlR6YyVvnk1J/i9An3ySn+Uww22rtQKXuAuj+qYHJ3AT7YeJz/uahnoy5DQ0NDQ+M0qgvr3fzEZCNjF1Vz78gGWlzYfOsyWjKxYWY+uW00s17fRLnNiadh5zlmvcRLcwYzqpH9s/SSyCX9Y7mkfyyZRdVkldRQ7XATYtLRLcpKQse6NViCIPDY9D5c9eoGPth4PKAUQQVvWvwTy/ehEwVmDkts1Fg1mpev03J56PM07K4Abjy8vZ+WHSpHuekJHh8ReHZOuEWPUSdiO+N9gu111tHacIubM7EadTx9zQD+Mr03X+w4wbsbMymsdODyyJj1En3iOrBgQgrjUqMQz7FRIobqEXSiqgDFmaRGdmVWv2ks3vWVzzapQ/Cfi0bz0fKNqZr6J8pAJtzdSgoOj+8PJpjcXadbZnOGryqNhoaGhkYj0NevwNcvRmJ6Dx3/XOekd3Q9SRT61lerkxoTyvf3jufx5ftZeeAkgoDP4lYvCoiiQP/4MB67rA8DEsKb5L1PFdA3xC+HCjiYX6lqSNWXGm9zyTz25T4GJYbTvVPrqGVrL2QWVfPAJ+qGVH3fqV0y8HGWi7H78rmwb2BpppN6xvDE8v0+rwcjtmLWS8wY1Hjxr1CTnhvHJHHjmKRGn6OpMfeOpPTzo6rbjhYf56djG7m81xRiO8SQW3GSLw/8xJC4vnX2q0HhR9FF/8wShnbteM6iahr+afnGlK5+6cdAJtxyxYpbJaMx6Nxdmyug/TQ0NDQ0GiC8K4gGkJ1+d3l8kokhr1fxp9F+5gFBhLAuzTTA5iWmg4lXrhtCabWTj7Zm8cn2HMptLirtbiwGicsHxnHT2G4BGT7NwQsrj/hNjW9I1trpkXlzbTrPXOPbI0jj/LFofQYulVBoIN+pzeXhhZ+OBGxMxYebGda1I+uP+TqhA+11pqBw5ZCERlxpy0U06zD3i8S2u9CnbspqsLAr9wBvbv2YCkcVHYwhTE0ZzaOT76izn9mow5bUgYc/34PN5eGKQXHMGBTfpM6L3dllfLwtm5xSG063TLhFz7jUKGYMjm+SxsVtjZb/iYQn4k0w9Z9o0NCEaxbdCB4F5YxE1WAaJYK3d4eGhoaGRhPQ/xpY/U+oJ9slNUJkdl89L25x0j9GxVmmM8KQG5pvjOeAjlYDt01K5bZJqYDXiHF55POaUn60oIpD+ZU+rweaGu+RFb7anctfp/dR7bOjce6xOT18si3HR+UumHKHY4VVHD5ZSY8AFu2Ky8W1NYfZ7g7FrvPNHGpYbAUuGxAXVDPc1kLo+ATs+4pRzogQxoZG89qMx+s/WCfSYXQsd1zQjdunpLI/r4Ivd+Vyw9tbiLAamDE4jssHxtM5LPiIvSwrfLYjh9dWHSOv3I7D7anTV3j14UL+/s0/fOHyAAAgAElEQVR+ZgyK587JqQG1aGgvtHwBisE3gL7+L+z0CVeNTkK5ahHi6bm7gRDfChpDamhoaLQKwrtAQsM1GI9NNFLt9ONMC0uEuPolmlsb0aFGCisbEN5oZj7cfFxVWjqY1HhREPgmLa85hqfRCFbsyz9rqXKXW+b9DZkN7mdLSyPjmpkMSFvDyJRIjLrglpoCEGYytNkadUN8CNaxcQj6IJfgIugiTYRO8UbjBUGgb1wYj1zSm/V/nsJfpvfmWEE1Fz2/hjlvbOKjrVkBZ1TZXR7mvb+Vx77cR3pRNTZXXUMKoMbpwe6S+WR7NtNeWMNmlYbM7ZWWb0wljgRrVIO71TfhXpzklVE/k9Nzd2sOb0R22VE8bmzHtlH6y6I6+1oMEteN1JoYamhoaDQZ4/7o4yzLvC+Uqcm/eaMTw0Tsf+lQp2Ev4D1u3P3nYpTnlOhQI0VV59eYOnyyStWYCiY1vsbpIbO4ujmGp9EIMouqqVGRxg/mO/UosPV4KYqfZp2eqiryn3yK7DvvJPKWeXR5601ev2UsfeM6YArQcBAF6GDWsWzBKDq14d5lYRclYRnaKXCDSiegizQTPb8/osH3u5JEgTEpUTx9zQA2P3IBfxjTlZ8PFjDunz9z2wfb+X5vHg63uoKiR1a45b1tbDharJra67s/VDs8zH1nKztVGjO3R1p+/FQQvBPuikfqqPpl3lc3zHxqwvVBbyFs4h1cZO7I12m5PpZ2oLm7oSYdY1Iap6SkoaGhoaFC6gUw4FpIW+ZXtVUVnRmSJ8OA2Q3v28qICjFQeJ6NqWqHeg+bYFPjy2q0OuOWgr8IRbDf6eH8Sq5+bQPv3jyCDqelcFauXEn+k09hHTuGlOXLkcK9gikmvcSyW0fz6Bd7+Gp3LoBq/ydRAKNOIinKwps3DvNRmGxrCIJAxxmp6GOtVPxwHMUtozh8DRnBIKIoYBkYTfjlKaqG1JmY9BLT+sUyrV8s5TUuvtubx7sbMvnz53u4qE9nrhgcx6hukbVKhi/9fITtx0tUv5f6xWY8/GHRFjY9cgEWQ8s3J5qT1nH1Q26E/V9C1kZvE95A0Vu8k23yZG41VfDD/nxVFZuGcnfNeolbJ6SccwlNDQ0NjTbPpc+RXVBMVPb3mAnAiNBboOtYmPkOiC0/uSJYokONFJ3nNL9QP3UqQctaWzT55pZCuEWvWn0e7HeqAHtPVHDFy+v58q6xmMtLOPnkkziOHCXu6aexjvRN3TXoRJ6dOZAHpvVk6eYs3tmQSY3Tg04UkBUFWYaL+3Vm/oRk+sWHBXQ9ZTVOiqqcON0yHcw6OncwoZNa3/MgZGQs1uGdsR8ppXJ1Dq68ahSXB0ESkUINWEfHYh3aCdHUuOV6mEXPtSO6cO2ILuSV2/hqVy5///oApdVOLh8UxyX9O/P2ugwfCXsITJjELSv8d+cJft/OM7dahzElSnDtElgyE05sB5et4WP0Fug1HS79FwgC/eLDuH1iCgtXpwcUxjyFUScyuEs4fxjdvm8UDQ0NjebgSGENs09cz9djh2Pe/SI4a9Sb+RpCQNLDmHtg7H1t0pACb9PRoioniqKcN8njfvFhbMooxnVGfnwwstZWg0SPziHneugafkiODsFikKg+I9UvmO/0FE6PzIlSGzc+8y1Pf/c0EXPmEPfcc4jG+tWXY0JN3Du1B3dP6U6ZzUWl3YVZLxFuMWAIoK7KIyusPlzAwtXp7MwqxaATERDwyApGnciNY7py/aiuxIS2rvRAQRQw94zA3DOiWd8nNszMgokpLJiYwuGTlfx35wluemcrVSqR6ECFSWqcHl5fnc6cEV3atUR76zCmAAwWuOFLWPV/sPl1QPE/4eqMMOFBGLmA0ysu77mgOzVOD+9vzFS1ws/ErJcYkBDGW38Y1io9HhoaGhotmUq7iwWLt/PwJb2JG3YhXHAnpP8MG16BwoPe1D+9GTomwei7oMc0kFrPtNUYTHoJo16k3OYi/DxFdq4b1YVF6zNQU9ENXNYaLu4Xe24GrNEgF/btxJ8/V9/m7zs1xHYn7737VFO8nB6ZQ1UKRc8upOe4AUGNRRQFIqwGIoJoyLsru4z5722jxumuNQhdnt8MQ5vLu6hfuDqdWUMTePyKfkhaNpFfenQK5cFpvdhwrJjS7DKf7cEIkxRWOdh7ooL+CYFFFdsirWtWknRwwV9h4oNwYDlsfAXKslDsVcgeASlpKIy9B1KneqNZZyAIAg9f0pt+8WE8u+IQRVUObC4PZ9ZSWo0SkiBw89hu3DUlVTOkNDQ0NJoYRVF48NM0RnaLYOawRO+Louh9fqdOPb+DO8+cEqE4X8ZUQkcLQ7p0ZKMfta6GUuP1osCsYYlaO5EWhFEnMWdEF97bkOkTcQTf7zSQFC+7qGfRMQdjG8gOPJhfwUdbssksrsbm9BBm0TM8KYKZQxMJszQsnb/mcCELPtjeYFbRqZqfz3bkcLykhkVzh6PX1m/1kl2iXqsajDCJKMDxkmrNmGp16IzeHiX9rwGgatUqSpcsocvjbwR0+GUD45g+IJYdWWW8uTadvTllVJwsIiQmiviOZm4e242pfTppP0INDQ2NZuKttRnklNr4z+xB53soLY6oECMFlQ5SY5quCWew3Du1OzvfKVWtM24InSQyb1y3ZhiVxtkwd0wSH27KqhPRUSPQFC8FWHukiMJKB9GhdVP8FEXh2z35vPLLUdKLqnC55TqqymsOF/LsikNM69uZu6ak+m04uz+3IiBD6nRsLpltmSU89Fka/56lPV/qw+5H4S8YYRJZ8S9a015oncbUGegiIvCUBCfPKAgCQ7t2ZGjXocg2G4dHjabX7l3NNEINDQ0NjVNsTi/m9TXp/PfOMVr0QgVvZEq9b+K5YlRyJP9zYU/+9cPhoBayJr3IK9cN1hp6tkASOlp49foh3L54e71GcjApXkadyN4T5UzuFVP7mtsj8+BnaXy3J9/vvXOq1GJ5Wi4/7D/JS3MGM7VPJ5/9/vrlXr/nqF9pTua7PfncMq6CPnEqSs8agLecpVpFRTAYYRJRENp9c+42EXqROnbEU9p4rXvF40HQtQm7UkNDQ6NFU1Bh555lO/nXrIFtXv64sUSHnP/GvQC3jE/mgYt6IgreRqr1oRMFLAaJV34/hCm9fBfFGi2DyT1jePW6IVjqkdgOJsVLlhUq7L/JriuKwp8+3s13e/ICMsJlxVvvdNfSHaw+XFhn2/HiavaeKFc9rmLLF5T89CZho2aRcNdi4m9/h9Ahl2A7srl2H6dH5u116Q2OoT2TGq0uEhNMH1aPLJMcbVU9T3uhTVgQZ2tM4XaDZkxptADyy+3klduwu2RCTTq6RVmx+pEp1tBobbg8Mncu2cHvR3RlYo/o8z2cFktLaNx7CkGA3p070C3Kyo8HTiII1IlqWAwSigLXDE3glvHd6BrZvhdVrYEpvTrx7T3jmfrv1arNmYNJ8RIEAcNpJRHLtmTzw/6TqiJf9UWS7C6Z2xdvZ9UDk2rV+N5Zn4ms0iA40DREj6zwzZ48/nZ53zo9sTR+45bxyew5Ue6j8giBi80kRljo1bl9R//axCpNNBsRBTtyYTpiaBQYQ+uo+NWHp9qFM78aMawL7iIbUkcjglYrpXEOcXtkVh4oYOHqYxzIq6iViFUUcMsylw2I45bxyfTsfP7qJzQ0moJ/fneQEKOOu6eknu+htGiiQgxkFFWf72GQXVLDiz8d4bPbx5AcHUJZjZMvd+WSXlRFpc1NR6uBXp1DmT4gDnMAzUQ1Wg5JUVbiws1kqQgQBJPipaAQ08FbL6UoCi/+fEQ1IhWIoIVHVli6OYt7p/YA4Id9+apiGcGkIepEka0ZJVzQW4uWqjG5VwxGva9k/ikCEZuZNy65uYbXamjdxlTuTtjwMsKBr0i5xIXw+hiQ3WCN8sroDr4OzB19DlNkBfuhEipX5+DMrkSQwDRgPidf3AGCQMioWKyjY9GFt65eBRqtjz055cx9Zwt2t6c2b/nMLuSf7zzB8rRchidFsPD6oVqkSqNV8nVaLj/sz2f5XeO0BugNEB16/tP8FEXh4c/3sGBiCsm/pgKFWwz8YUzSeR2XRtPx+xGJPP/TEZ/6qWB6T5l0EoMTveusTekllNtcnEmgkSSHW+adDZncOdmrolzpR9QgmDRERVEoq/Edk4YXSRS4c3IKz60IrjYSvKm/oijw8s9HiAs3Mb57+802aJ2rspJ0WHYdlGaC2w6K7FVCd//azLcyD35+En7+OwyfB797srbBozO7kqL39qG4PCgO7wNE8YCgM6E4vf+vXHeCyvUnsAyIpuNV3RECaCanoXEKm9NDSc2vndlNOiKsBtVmdpvSi7npna0NPsA8soJHVtiSUcLlL6/jizvHaikLGq2KowWVPPblPt6/ecR5k/tuTUSHmJo9zc/u8rAxvZiiSgduWaGDSc+QruHEhpkB+HhbNuU2F7doqnxtltnDu/CflUdUtwWS4mXSidwyvlutc+SdDRnYVCIcwUSSXG6Zn1ZuZzSl4HSiVtofTBoitNn+3k3GzWO7sTOrjJ8OqKdn+sNilPjvHWPJKbXx58/2MLJbBH+Z3ieo/mFthdZnTOXthnenexv2KvV86acMq23vQNERuHYp9owqit/bh9LQzfJrWLlmTxHuEjvRt/TXDCqNelEUhU3pJbyx5hjrjhahE0VEAVwehTCLnvnjuzFrWGLtQvJoQRXz3m3YkDodh1smu9TG3EVb+OS2MVpDQo1WQZXDzYIPtvPnab3oF99++5AEQ1SoodkiU9klNby3IZOlW7IQBAFZUVAUr4fa5ZEZ0S2CWcMSefq7g3w4f5TWZ7EN09FqYHKsgZXHq3GrGCUNpXgheA2yUxwrqFJp8xxcJMlls7N32X/pYy0nJHQKlRh99gkmDVEQBM2B0wCCIPD87EE8/MUevk7LUzWIT8eoE7EYJJbMH0X3TqF07xTKD3+cwL9/PMyF/1nDo5f2YsageFUncluldRlTZVnw3uXgqAj8GFcNZK7FtexRig9d3rAhVedYGWdOFcVLDxJ5fe92dWNoBM7+3Apu/WAbJdVObE4PCnU7sxdWOvjPj4f51w+HuXF0Vx6+uDdPf3+QmkbIvTrdMgfzK/nlYIGqjKyGRktCURQe+jSNYV0jmDU88XwPp9UQaTVSWuNElpUmTYl8Y80x/vXDYWRFUa1FAW/foA1Hi4npYCS+o7nJ3lujZfKg9STbRRPFgoSKFoVfTHqRp68aUCcK4c85GEwkyaPTY7rxJhInpzLju4O8vS4Dp6fxaYgeWWFkt4jAL6ydopNEnrl6AJf0j2XhqmPsyi5DlhVcp90UVoOEXhL5w5gkbhzdlciQ3wxdq1HHX6f34YpBcfz5sz18vuMET83oT5fIhhVbFUVhZ3YZmUXVVDs9hBglkiKtDEoMbzXr7tZlTK141MeQSnq+khoXZNwbgtXg/dDf2uFkcZqLVXN/VRVy2Sjb3xXF42tIbclJ4x+/vMbhokxEUaR7ZFf+dsHdDIrt7d3BLeM4XIozuxJjl/atVqLhy+b0Ym56dys1DXhyToXOF2/K4lB+JZvSi1ERKQqoSLfG6WHh6mOaMaXR4lm0PpPjJdV8etuYhnfWqMWgE7EadZTWOOssWM6GZ1ccZNG6TJ+aTDU8ikJxlZMZL6/ny7vGtvseMm0Z05F9vNu7P/Nywyiucvg1suscoxd59JLeXDE4vs7rVoP6kjKYSJJeEulg8p7nhtFdeXt9hup+gaQh6kSBq4bEY/EzLo26CILA5J4xTO4ZQ3ZJDV/uOkFumR2by02k1cjwbhFc0Cum3mj1gIRwvrxrLIvWZXDFK+tYMDGFeeO6oVc5psrh5osdOby+Op2SGm9fvdMdSJFWA7dNTGHG4PgWXyveskd3OtXFcHiFamqfR4EXNjt5ZLz6pONWonF4evi8Xumo5qZP/8xTF97PZb0m4/S42ZKzG6NUNySsuGWq1uRgvL5P01yLRpvgaEElNwdgSJ2OzeVhw7His5J7BdhzopzMomqSojQZYo1zj6IoZBbXUFLtwP1rKmtKdEidCXNrZgmvrTrKF3eM1RrzNoLoECOFVY4mMaa+2JHDonUZQdVDOD0yOWU2bnlvG8tuHdVqPMQawWE/cIDkmTP57po+PPrGT6zMcyEZDX7vFVGA2DAz4RYDTrdcqz4L0Cs2lGOFVT4RrmAiSZIo1AqexIWbGd61I+uPFauOpaE0REkUuGmsVvPXGBIjLNw1pXujjtVLIgsmpnBxv1ge/e8evtyVyz+v6s/AxPDafXZll3Hj25txy4rfNVSN08ZT3x7gmRWH+GDeCAYkhKvu1xJoPcbUjnf9yp0/MMbAM+sd3DHcQLjJd58q93TUWg6ml2QDMKPPVADMosTEbiN830AB28ESPNUuJKvmodPw8tf/7vP7EKgvVU+trwcEV6QrCrAxvVgzpjTOKdUON1/s9HoSi6qc6CTvc1VWFCRB4PpRXblhdFckUeDuJTt5duZAEiO0xryNITrUSFGlEzqf3XlkWeEf3x30uzhuKK04LaecXdllDO7iq4yr0bpRnE6c6RkYe/bEZNTxP78s5JE77+Vbc1deXXVMVQVPViCjqJo/f5bGI1/sYd64btw9pTuSKHDz2G6s3F+gmu4XaM8ii1HH6OTI2v+fTb3TlYPjSY1Rb0qr0fx0ibTw/s0j+HJXLvPe28blA+P404U92J9XwY1vbwmoZty7xvIw+/VNLL5lJEO7tsznUOsxpvZ86lXuU2FYnMSkJB3PbXDw5BRfOXObZzTg+4NMjkhEFET++M1TXN7rAgbH9yXcpN7LR5BEHOllWPq3X+lHjd/ILqlhR1aparFtIKl6agRVpOtRVCVoNTSai+W7T/Dgp3sQBH5zIpxxC761LoO31mXQ0aJn9rBEJveMOfcDbSNEhRgprFKf84Jh3dEiapzqEtOBPKscbg9vrknn1euHnvVYNM4vTrfMj/tPcvhkBaU1LixVZYT2m0KCqEdesQLBZCJqygTWf7ADRwML3VN9iV5ffYydWWW8ceNQBiWGE2MWOe7n2IYiSSa9yPzT1AH35JTz08GT6u9fjxPgFJV2bY483wiCwIzB8UzoEc1T3xxgyr9WUW5z+cjxN4TN5WHuoi38cP+EWsXRlkTrMaZspfVufmKykbGLqrl3pK/RpKDuvQ81Wvn8upd5dfMSHvz+WQqrS5icMpJnpj1ItPWMgkVZQa5Rn5A02h8fbDx+1ql6ZxJc13lvPriGxrng3fUZ/PP7gw1OgM5f63EKKx0czK9scgGF9kRtZOosWbj6WG0Pu9MJ9FklK7DyYAEl1c52KXncFsgrt/Hu+kw+3JyFgnLa/aBgjp/AS0/+yOT8vdz6h9v557JdbM4oDnixa3PJbE4v5p4Pd/D38s1ct2kn/+57JXYl+N+9XhKZNew3oZo31h6rfaacTqAOy5UHCiircWpqfi2ACKuBf80ayPz3t/Lj/gLVfRoykO1uD2+tzeCv01teyU3rMaZUYwC/0S9GYnoPHf9c56R39JmFbv6P7R6VxH8ufQSAo8XHuefrJ/nfn17ilcv/5rOvljOucYpv9+addWf2MwmmSNcgiXTUJgiNc8CP+08GZEidjqx4VeH+/s1+/naZ/2ishn+iQowUV1SDxw1S46fqPTnlqq8H86wy6kQO5lUwJjWq0ePQOD+sP1rE/Pe34fbIOH3mLAGboAOXzIqIXqzcUINAjcp+9S907W6Z1XtP8K2jjJveeIr8rcUs3ZIdVOsPs17ivdP60JXXuPhh30mf+qtgHJaCAB9tzWbBxJSAx6HRfNicHtYdUa9/C8RAdnkUlm3J4oGLera4OtzWY0yZwqEyv95dHp9kYsjrVfxpdN2CXVGoRFYaLlxLjezKrH7TWLzrK9+NooBobT0fl0bzUuEnfSCYVL0zCaZI1y0rTO6lpVBpNC+yrPDoF3v8GlL1LbBsLg9LNmdx64TkFpmW0SJRFMhYDetf4NaMtQiyG7YBehP0vBTG3AVxg4M6pb8FbTDPKkXx/8zTaLmsO1LELe9vDcgRIgsish8lv0AWunZJz8cpE5nbuTN/nd4Jo07k3Q3Hsbs89brCDZKAQSfx7k3DGXJaXd7ao4XoRIEzu60F4wSwu2Q+33lCM6ZaCMvTclWlD4LN6Plubx5XDk5ozqEGTeuxDvrMgPXP+62bAkiNEJndV8+LW5z0j/ktOmURV1HhmQVnNH87Wnycn45t5PJeU4jtEENuxUm+PPATQ+JUPKmygjGl5SqJaLQMgu3MfiaBFOmKAlzQO0ZLudFodtYfK6La0fh6G0XxpsQ+OK3XuRx26+TwD7D8XnCUg7OaOk8Plw32fQ6HvoWwRLjq9YCNKkkUVEVvgk0rNupalidYo35OlNm49YNtqoZUIPVGpwhmoZteVMWh/Ep6dg7loYt7M7FnDAtXH2PDsWIEqCPLbzVICILAdSO7MHdsko/DpbTaqXrfBuuwLFcR0dA4P2xOL1YV7QrGQK52eticXqIZU41m2M2w7j8N7vbYRCMfpNX98Vh131Phme2zr9VgYVfuAd7c+jEVjio6GEOYmjKaRyffUXdHAcwDohFNrefj0mheOpj0VNh8F5nBpOr5o6EiXaNO4pbxyY06t4ZGMCxcfay20Px0Al1gOT0yH2w6zn1Te9SRUNY4g62LYMXD9ToLUWRvE/qiQ/DOJTD7A0id2uCpO1oN5Jf7njeYZ5VHVogObZp+VxrnhnfWZ+BS6a0ZrEBSMAtdl1vm8x05PHyJt0/nqORIRiVHkl9u58tdJ8gsrqbC5iLMYmBY145cOiDWr5Hu8aN6G6zDUq22WeP8UFKtXgMarIHs7zznk9ZjHYR2gpTJ3l5TpwWNM++rq76XGCZi/0vd5rqSUI5J2oVdHlGnfCo2NJrXZjze4FsLOpHQ8fEN7qfRfrikXyzvbMjwqZsKJlUPvIL9wTzqzXqRGYPj66RDaGg0F1sz1YV/gllgKYrC4ZOV9IsPa+rhtQ32fwUrHqnfkDoTVw18dD3M/Rbih/hs9sgKWzNL+HZPHhU2l+pzJphnVbhFT984rWl9a8Hh9rB0S5bP/NQYgaRgFroexRsRq/OessLRgirWHS1i/dEidKKIIMCyLVm8tyGT2yelMLV3J59GsB2tBtU0v2AdlmFmrZ1NS8FsUL+HgjWQ/Z3nfNJ6jCmAi/4Bx9eDozK443RmwpMPczJjDIpKNKFe9CKWYZ3Qd9b6+Wj8xo1juvLexkzUTKFAUvUkAab1i2VLZonfdAY1HG6Znp20vhkazY/bI6t6tiG4BZYgCFRoMv7quGzw39vAXXcBmvR8JTUuyLg3BKvBW2Tw1g4ni9NcrJpr/e3Yz+bB3TtAEPDIClsyvAbU9/vyiQoxcmn/zrx/8wiue2tznRSrUwTyrDLrRW4dn6wJMLUivt+br+qla4xAUrAL3dMl1TceK+a+j3ZSaXfXpnc5T3um7M4p508f70YniTw5ox+XDYyr3TY6ORKXyrwYjBPAqBOZ1u8sG7VpNBmJHS1IIpw5rQRjIOtEgYSOLa93YesypiJT4LrPYPGV4KwO7Bi9GWIHofv9i0QXuCh8Iw3F4QkoHCDoRUy9Iwi/TCte1KhLQkcLQ7p0ZFN6seqt1FCqnl4ncs8F3YkMMXDVK+vJKrX53fd0ZAWe/v4QJyscPHSxVoei0XyI9Syeg11gSZo8ujp7P0etoTx4vfwvbHbyyHj/6XVKZT57Nv/Ix/mxfL/3JDGhRi4dEMvHC0bT7bSG3hN7RvPzgQJVp01DzyoQuHpoy6pP0KifQ/mVqum5jRFICjYSFBHivV+/3p3L/3y6u0Hxi+pfm7I+8OlucststWIRMR1MjI6zsuZ4BcoZz6JAGwADXD+qa4BXqtHcXDM0gfc2ZOKR694TwRjIkihw9ZCWlynWuowpgC4jYd6PsGQ22Er8GlUKOlBkhL5Xw2XPg6THEGek092DKf7wAO5CG4pHBpXfuWAQQYGQCQl0mNpF88hpYHd52HOinNJqJ4IgEGHV89hlvbnmtY2qk1Z9mPUSMwbH07NzKB9vy6ag6sxEhvqxuTy8uyGT+HAT149OCupYDY1AEUUBs15SLRgOtt5GE0vxw/rnwVmluumBMQaeWe/gjuEGwk3qc5DsslH107+JG/syn942mqQo9QyKf141gItfWENhpcNHaro+THqRl+YMJtSkpUq1JvzVlDRGICmYha7VIDEmJZKNx4oDMqROx+6S+c/Kw0SHGrm8i4mi1xYyfd0etgy5HptKz6qGnAAC3uhWpw6mgMeg0bx07xRKaqcQ9p6o8NkWqIHcJ7YDydEtLzun9RlTAJ36wn17IGMNrH8BMteAZABB9Pbk0BlQBt9E5jM/EHftXZik3yYCXaSZTvcMwZlbRdW6E1Rvy0XQ6bwSaR4FqaOJ0IkJWAbFIBpbXl6mxrklq7iGdzZk8NHW7FpPvSB4VcoEAab27sSK/fkBTxpmvcS41EienNEPu8vD41/ta5Taks3l4alvD3LV0AQshtb5M9Zo+Uzr25kvd+f6FIMHs8AKM+tJjWl5k995pzwHyrL8bh4WJzEpScdzGxw8OUV9QSihMMazlTETU1DVHP6VCKuBT28bw8yFGymudqj2yDsTk07kH1f2Z2qfTg1fi0aLooNZfU5orEBSoAtdQRC4qG8nJj27ulHtFOwumUc+2UnKqmfodNnFzFjyKt9+5VUDVEtTrQ+LQeIvLbC5a3vnzkmp3P/xbtWWDQ1Hyb31UgfyKugd27JqOFvvKkwQIHmi94+9AmqKvDnkpnAI6YQo6Qg/mUrhiy+R+NqrPocb4kIIu7AThU9fR+raDSCDaNIhaIpTGniLZv/+zX6WbM5CVhS/i4/v9ubjlmV0ooBeErC7ZNW0P5NeRFHg+lFdePji3oiiwFe781TPGajakiDA8t25zKFNi48AACAASURBVB7epSkuWUPDh3nju/Hd3jxsKuGMwOptJBZM0OptVKku8joB6xGeeGKykbGLqrl3ZAORPWc1GOs3WBMjLHx773ie/Ho/3+zJQxQEnwWNJIJeEkmNCeGvl/ZhZHJkwJej0XLoGmnFrJd8vt9gBZJOp8HUdckrc77jeBmVfnqSBTK3CbLM3sdfZuCFAwF47fpwrn1jIwfzKwN2WloMEm/PHa45cVog0/p15vu9+UE5oU9nU3oxV766nuSoEF64dhDdO4U2fNA5oPUaU6dj6uD9cwbhs2dRvGgRtrQ0zAMG+Gy3HzyIsVcvJIuWgqLxG4qicPeynfx8oKBBb9ipYlqDTmBMShSyrLD2aBF6SUQUvB27wyx6bh2fzMxhCbXd3QFeW+UrOx2M2lKN08Nrq44xa1iitljVaBb6xoWRoPdwxIlq5KOhBZaiKFyl1duoIzecHtwvRmJ6Dx3/XOekd7QfR58gghyYsFKE1cC/Zw/ib5f35bPtOSzdkkVpjRO3RyHEpGNMSiS3jE+mRwtZoGg0jukD4nhi+X7VbcHUGwWDxSAxb1w3Hvh091m1U7AJOt7cXcLvf6cgCAImvcRHC0bzwCdprNiXX69z02qQsBp1LJo7XFMPbaEIgsBzswbiXCqz6lCh36bi/pAVb0rogbwKZryynndvHsHwpIhmGm3gtA1jyg+i0UjUbQsofPElurz1ps92+/4DmHr3Pg8j02jJPPP9IX4+UBDUj9zuktlwrJj7ftedV68fSmmNE6dbpoNZT0eL3sfYqXK4SS/0rfcLVm3pRJmN0hqXVpOi0eR4ysvJf+LvPJBZwH29r6XGHVy/FpNe5Mkr+tFBq7dRxxwekBH0+CQTQ16v4k+j/QhRyC4wBpfyEmbWc/O4btw8rltQx2m0DsLMeqb168zXu/PwqPRZCiSdKhgsBokPbxlFVIiRdUeKVfcJZm7LK7eTV24nLtzbyNeok3hxzmCyS2p4f2MmS7ZkoygKkiCgAE63zNCu4SyYmMKE7tGImuBNi0Yvibx63RAWrc/g1V+OYXd5gq49V/CKl8xdtIUv7hhDUZWT19eksy2zBJvLgygIdDDrmTEojpvGdiMxonkVANu0MQUQftVVFL/xJmVbt7NK15mlW7IoqHTg9siYSsyMi+3H/DJb7Y9Wo31TWu1k0foM1YhUIHVMz/94hBtGdW3wfiqrcaKXRNxneKeDVVsySCLlNs2Y0mhaqjdtJveRhwmdcgFTn3qSt09UM++9rapiFGqY9CL3Te3BNcMSm3mkrZiOSV61WVdNvbulRojM7qvnxS1O+seoRKdiB4Oopadr1GXBhBRW7MvH42q+prVWg4RJL7Fk/ih6dg6l3OZCFMGj8pgIZm7TSyKlNU6feTQxwsKjl/bhgYt6kVtmo8LuwqSXiA4x0lGbA1sVgiAwb1wyc8d047s9edy9dKdqiURD665qp4dLXlyHQSfWmZ9kRaGk2skHm47z4eYsBiaG86+ZA5vNqGrzxlSNIvL+ZXfz6SfZCKbCutavEMrxAljy3CpGdIvgoWm9tNBwO+ejrVmqddxNXcfkT3Y6WLUlBa92ioZGUyA7nRS+8AIVy78m9qknCRk/HoDRKSY+u30M9y7bSXaJDadbVvV4n1pcPX5FX6YPiPPZrnEaogQjb4e1zzXYsPexiUY+SFOpQzGEwLg/NtMANVozfeI68MjFvfm/7w4GnUrlD7NBQicIuGWFrpEWbp+UwrR+nTHqvHOVoih+hP4bMbfVYwMadKJf5UqN1oUkCuSV2zHpRWxn1FAFuu5yywpuP44+b0qowrbMEi59cS1Lbx1F37imX+e3aWOqsNLBtW9sJKfUiEOUQeXDdsqALLP2SBFbMzfy0pzB/E5TL2qXyLLCW+syfIoiG1PH1JAxFW7R12leeIpg1ZZcHrlOHZaGRmNxHDnCiQceRJ8QT7f/foEuom4eeu/YDvzwx4nsPVHO2+sy+G5vHg63jAAIHg9DkyK4fUoPJvSI1vpKBcrQubD2WZ+XM++rW7OUGCZi/4tKKp+og56XNNPgNFo7N45JwiXLPLfikM9C9UwaigCIAsSEGnnkkt4kR1lVC//NtircfuqZgpnb3LJMmFlLD24PyLLCm2vTfe7PYNZdAb2PAhV2N3Pe2MQ394xv8ghVmzWmqhxuZi7cQE6pTbVRoRp2l4e7l+5g0R+GMyY1qplHqNHSKKh0UGn3rWEIto4pu9RGjdNdr2S5xaBjYEIYO7LK6rwerNpSz06h2qSjAXjrBn4+eJKMohqqHW5CTDq6RVmZ0isGveQ/DUyRZUoXf0jRa68R86f7Cbv66noFTfrFh/Gf2YP4z+xBuDwyHlkh95qrib3xKcy9Yprj0touIdEw6k7YvLDBdD8f9Ga4+GmQ2uw0rtEEzBuXTN+4MJ5feZjtmaW4VNZDgUQAZAVOlttJibaSGuM1pNylpdRs20bNlq3UbNmC68QJ+ky8m716XxXIYOa2cLOBhI5a6UV74NDJSqocZ7/uOkVDToEqh5sHPt3NsluDO29DtNmn8F//u5fccruqIdVQn4P5H2xjyyNTsRrb7MejoUKF3YVeEjjzdx18HZNAuc3VYP+n2yelct9HO6l21I2YBqq2ZDVI3D4pNaAxabRd8sptvL/hOIs3H0dWwOHy4JYVdKKAUS8iCQI3jO7KjaOTfBpYuk4WkPfoo3gqK0hathRD165BvbdeEtFLYExJwXnsKOZ+fRs+SKMuFzzm7Td16NvADSq9GcbdDwOvbd6xabQJRiVHsuzW0cx6fSNbMkrqbAsmAuDyyCxcto7/qdzpNZ5ycjAPGYJlxHBi//4Epj59uPdwMfct26kqKBBoO4X547tpCrXthJJqp2omQ7DrLgjcKbAzq4zskpomjU61SWuhvMbFt3vycKqICATyYSsKfLnrBL8fGdzCQqN1oxMF1TztYHO9ZYU6kYCD+RX8dKCAokoHChAdamRyzxim9IrBIIlU07jmdZIocGFfLSW1PfPLwQLuWLIDj0fxSRt1ywruXw31N9dm8M76TBZeP5QJPaIBqPjhB/Kf+Dsdr72WqNsWeJuXNxJDagqOo8cafyHtGUGAq9+Clf8Lm17D6ZExoN6nB70FFBku+j8YdtM5HaZG66bc5mJnVqnP68FEADwKfHXCxUMp0cQ+8TimPn0Q9HUzI6b0isGgE/2qszU0t8mKognXtCPObAh/iqDXXUE4BWRF4b0NmU3a1LlNGlOfbM9WLfAP9MOucXp4fXU6c0Z00bwj7YhIq7FJ6pg8soLVIPHV7lwWrjpGelEVLo9S+9CQRHj55yN0ibQya1gi723IxB5kd3eTXuRfswbVm76l0bb5cf9J7l66I6DGh063jBO49YNtvHJ1X3p/8jo1W7eR+PJLmAcNOuuxGFNSKf/qq7M+T7tFEOB3j5PZ/QZ+fO//uMW4EkF2e3tIgbcnlSkMxtwDg+Z4ZdU1NIKgsNKOQSfi8pydgqyi02G6cS5mPy0PJFHgH1f2548f7wq6KatZL3Hf1O5a6no7IsysV3ViB7vuCsYp4PIofLHzhGZMNcT7G4+rqtcE82EXVjk4kFdJn7jg+ndotF7CLHr6x59dHZMAjE2N5Ia3t7A/r0JVStojg02WOZRfSXZJDSa9hMMtq8qCqmHSizw2vY8mlNKOOVpQyT1Ldwa9WLG7ZO78cAfvG80M++JzRGvTKGIZU5JxHj3aJOdqz7yb5sAy8gGE3y2EwoNgK/MaVJZIiO6p2jhZQyMQbE4ZQUVrL9gIgCQI2FweQuvpH3dx/1jyyu08s+JgwM8os15i1rAEbp2QHND+Gm2Dnp1DUVSsqWDrx4N1ClTY/UT/G0mbNKaKqhyqrwfzYUuiQH6FTTOm2hlnW8dkNkikF1VzssKhmmZ6JjVODzV4iAk1IisKNqf/5nVWg4TJIPHsNQOY0kszpNozr/5yTDWKCg0X4LokPV8MuYIRTWRIARi6dMGVl4fscCAa/TSX1aiXaoeb/+46wTf3jAdJD537n+8habQhOph1yCqL1mAjAE6PHFAj7pvHdSM61MBDn+1BAL/zmlkvISsK9/2uO7eOT9aygdoZJr3E7OGJfLDp+K8y5r8R6LoLgncK+EsvbCxt0phy+VlkBPNhK4rXk6PRvpjSKwajTvIxpiCwOiblV8Uj5xkPhYYWuOU2F+O7R3HD6CReX32MTenFSKLXj+iWFYYnRXDbxBRNdlqDCruLb/bkqU4GgdSEeoAvdubwl0t7N5nIjmAwoE9MxJmZialnzyY5Z3vji50nGJEUQbzWQF6jGegcZlJ9PdgIQHSoEZM+MO//ZQPj+V2fzizfncvC1cfIKbXVpqa7ZW9bjwXjk7l6WEJABppG22TumG58uDkLVPJzAll3QfBOAbMhcGGLQGiTxpRJL+Hy+EotBvNhCwKEmtrkx6NRD5Io8Np1Q/jDO1uCTqEy6UXcsuzjXQlkgetwe3ud/e2yviyZPwpZVmpl2kNNOkTNgNL4lc+355xVTSh4m0Y3tciOMSUFx9GjmjHVCBRF4YONx3nssqbL4dfQOB2jzhsBWHwWEQCv0l5waXgmvcTMYYnMHJbIyQo7JdVOZEUhzKwnPtysRaI06BJp4eohCXyxM6fBfmj+CLYcY3hShP+TNYI2aS30ju3gI/8JwX3YTrdMr86+Tek02j4jkyN54drB3LdsZ8A/bItB4oJeMazYd5LTvStBK8xszOQvl/ZBFAXCLJqnTsOX3dllZ10TWuP0kJZTzu+D733oF2NqCs5jmqJfY9icUYJblhmT4tufR0OjqZg7JokPNzU+AiArCjPPQmmvUweTT3sGjfaL3eXh67Q83lybTlZxdaMNqVMEU46xYELKWb3XmbRJY2rBhGT2nShvdJ8DgNHJkcRoP/p2y0V9O/Ph/FE8/PkesoqrcboVPGfkm0uiVwK9e0wo/3dVf+a8ucmnjiVYhZmlW7J5aFovTaVPwy+lNvXC2WALcEtrnE05LAwpKVSu+KFJz9leeH9jJjeOTtK89BoNklNaw7d78sgts+Nwy0SFGBieFMG41Kh6MxjcHpmv0/JQ8PagU+vBWR9mvcjs4V00pT2Ns8btkfnXD4d5b2Om/3o6RfEruGPSiX4VkANxCoRb9IxK1iJTDTKpZwxGvdToPgcWg8StEzVFmfbOkC4dWXHfBPbllvPkN/vZmlFa+9sOMeqY1i+WeeOSSI0JpaTaiUPlxx3sAtcjy5RUOzXvnYZfrH6aQQdbgNvUTcmNKSkUpWuRqWDJL7ez/mgxT1894HwPRaOFoigKa44UsXDVMXZklSIrSp1UPatBwmLUMX98N1WDJ6Oomj99vAuzQWLFfRO4c8kO0gurVecsNUx6kaFdI/hrE0pJa7RP7C4PN7+7lZ1ZpfVHok4zpCQBLAYdblkhKsTAbZNSWHO4kNWHCxtVjvH45f2a3HHVJo0pSRS4/3c9eOqbA6rpMPWhFwWSIq2MTtbSLdozlXYXqw8XUlzlxOWRCTHqmDUskX9cpa6wVWl3oRMFzvT1By07KwpU2t100kQkNfzQLcqCXhJ86h6CqQk1SCJJkU2n5gdgSErClZWN4nL5NPLU8M+Szce5fGBcvVLTGu0Xt0fmoc/S+HZPvt/1TPWvKrD//vEwb63N4KMFo+kWZUVRFBZvOs5/Vh7hnimp3Dg6CVEU+PS2Mcx7bytpOeWq7TtOIeBNiZrcM4b/zB6kiR9pnBWyrHD74u3sOF4aVG9NURBIjQnh2ZkDSIkOQRAEZg5NZM6bm9iXWx6wQWXSizxwUa9maSvTJo0pgOtHdeVgfgWfbT8RsEGlEwUiQgx8MG+Elm7RTjmUX8lba9NZnpaLJAq4PQqy4m24K4oCh05WctvEFKb0iqkzsVgMOlV1tWAVZmTZGxnV0PDHzGGJvLk2gzPrHoJS5RLgmqEJTTouwWgkM7kfOev3Q0xnwix6esd2IKSJI2BtCadbZsmWbJbOb8LiNY02g6Io3L1sJ6sOFgRUT2J3yTjcDma8sp5Fc4fx/MojVNjdfHLbaFKiQ2r3sxp1LLllFKsPF3Lvsp3YXB50oojT7QEBDJJXrnxMSiTzJyQzOjlSWxNpnDXL03LZnFGiakjVp3jskhUOnqxkz4lyUmO8WgYGncjS+aO4/+NdrNx/EpfLjUdQL48w60UUBf5xZX+uGtK0894p2vQs9/cr+hFm0vP2+gxcHqVeXXmLQSI+3MyS+aOIDNH6pLQ3FEXhXz8c5q116bjcMh6VW0X2KGw/Xsp9y3aSHB3C4nkja0Uiwv2IRQQrOysrChFWQ5Nfn0bboWuklQEJYWzNLPXZFmhN6IikCOKaSIK73Obik23ZvLU2g/Kes5B+zEHQ5QPgkmUuHxjHvHHJ9NQEfXz4bm8ePTqF0L2T9tlo+LJw9TFWHSxUNaTqW3xW2F3MXLiR+y7ozh2TU9Gp1OCKosDkXjFc2LczyVFWwix6ympciIJAR4ueyb1itHRzjSbl1VXHVCOhgSge25weXlx5hN3Z5fx04CSVDjeiIBBu0TO7fxTZ361kQ8Ig9Lrf7nVZVggx6Zg/PpmZQxObVdSrTRtTgiDwwLReXDogjrfWpvPNnjxvKpZHRv410uCRFXrHduD2SSlc2KczBp1W+N8e+euX+/hse05A4eJqp4eD+RVc9vI6lt89jjCzHr0kMmNQLJ9uz8Gj1PXgBbrAlQSBS/vHBtzDQ6P9cufkVG5fvEM16t5QTahZL3HH5KZRMvpq1wke/DQNQRC8YxF03kZWp7Wm+Gx7Dl/tzmVSzxhe+H/2zjs+ijL/4++Z2d3sppJCSICQhCYdpPeiYgHROxsWUFBEwX6eeud5evqzt7PcKVhAbOjZQMEu0kvoSJckpEEK6dm+M/P7YwkSdjbZJSEQdt6vV16YmdmZZ+LsM8/3eb7fz+fafoQZ9Oe7lvfX5XDrqPTT3QydMxC3rPDm8kzN73hDg09V9UqSd0uO1gykjic5xoxLVrihCW0SdHROZNehSnJLrT7bg1E8zi61kbc+p454SpnVRUFJNUpSL8Z0SeCqAe0xiCJGg0h8hIkeydHNYi1zVgdTtfRoG83Lk/vxr8t7smxPMUdqnMcCqndXZ7PkrpH6EnYI88G6g3yxOT+o+jq3rFJY6eDm9zby+a0DEX77jFtyP2AxtyPjO/sRiMKMySByiz6w0gmAseckcv2QDny8ITeo59ZilJg+Io3hnRIa3YYFaw/yzHd7Gsx9l1WQ3QrL9xZz7dz1LJw5VJ8wAHYWVHKows4F3Zs+f1+n5fPT7iIfBVkIfPBpd8vMWZHJhT2T6r1OcoyFHfkVTddwHR0NFm0t0BQ8CUbxGNBUoXTinTBYtreYzTnlfDJzWLNnQoREMFVLtNnIn85td+x3VVX5aEMuB0ttpCc0bTG2TsvAIyu89NN+/4W99aRSuGSFPfmlbH3pT/RPkuj654fp8Z3EjvzKoGVnJdFbYNmzbUxT3JZOCPCPCd3xOF18uv4gDqnh9AWLUWLK0A48cFHjTXV/3VvsDaSCUFJyeBT2HK7i3k+2MWfqgEa3oaXzwbocpgxNbXDlQCc0eWtlFlZn4/zkdh2qIrfURof4cN/zeGS++62Q99Zmk1dmY/WBI0SbDYzpmsjUYalNlgasowNwqNKB1rAoWMXj+lBUKLe5uWbuOr69ZxTtmvEZDqlg6kQEQWBQWhwbs8v0YCpEWba3GLesPSAMJI/XIcPbsffx5k3jAZgz1cGEV1cddXkPrA2iADEWI+/cNLBJ7kknNBBFgdt3LKKLFMmCxEHkl9txeZQ6s9mSKGCSBDrER3DfBV24uFdyo6+rqiqPLt7pN5CqbwLC4VFYvr+Y3Yeq6NE2dCUrK2wuvtt5mGV/HXu6m6JzhpJVUqO5PZjBp8kgkllSUyeYsjo9vPLzfj7OyAX1D4+f/HI7AAeKrcxbk82gtDj+dkk3erXTJ/h0Go/bTwZDsIrHUP87BqDG4eGuj7fw5ewRTdb+hgjpYApgcHocG7LLuGbQybt667Rc5vqZ/Qs0lUJBZFmuhzKri7gIE4lRZr6aPYLJc9cdTSetP6IySgJxESY+nTlML/bVCYrq5cuxrlrFdYsXMSUqip0Flby39iD7i6qxOj1Ehhk4JymKm4anNemK5+acckqt2oa/gUxAuGWVd1dn8dI1/ZqsTS2Nzzblc373NiToYkc6fvA3WRHM4FNRVaqdf9QvllQ7ufatdeSX2/16TNUaz68+cISr56zj1Wv7NZgqqKPTEP6EtYJVPA7kHSOrKrsOVZFZUlNHxfJUEvL5BYPT49h4sOx0N0PnNLHncJXm9mBSKUwGkd+Lqo/9nhIXznf3jOaWkR2JMhuI0JA6jzBJRIYZmD48ne/vGU2avjKqEwSesjIK//koyc8+gxTlzQ3v1S6GF6/uy9d3juSX+8ey+M6RPH9V3yZPHZ27MkszLbZ2AiJu/CzCzxmOaDIjSAbCOw+po1opKypLdhymyuFu0na1FGRF5YP1OUwdphf86/jHnxjW8YPPhhAF4dj7p8bp4eo5a8kptQVs1mt3y9z9yVbWHDgSeMN1dDQY2SWBiDDfsdDxise2/etQ3A5U2YM9cxPlv86rc2yg7xjw9rPz12Sf0ns6npBfmercOpJqh5vCSgdJMfrKQKjh8FMrFVQerwpVDk+dTTHhRh66pBv3je/KD7sK+Xr7IcpqXKhAfKSJS/skc3GvJF3ZTCdoVFXl8KOPEn3ZJCIGD27262dkl6FRFx/0BMTOgsomEcJoaazYX0yMxci5Ka1Od1N0zmDatjKzv8g31S8Yuw2PrJIS503xe3TxTg5V2DXreetNzXUrzPxgExkPX0CE7hmn44fiagfL95ZQZnOhqCqtLCZGdUk49vxd2COJv3/5m+ZnA1U8DuYd41FUvtpSwJN/6t34mwuAkP9miKLAwLQ4Mg6WcVnftqe7OTrNjEEUj6U1HE9QebyC11lbC5NBZFLftkzSny2dJqLyy69w5+bR7uWXT8v1bS6P5vZgJiBUFarsobky9f66HG4clqoryOrUyy0j03n8m92avjyBDj7bx1ro2iaKSrubpTsOa6adB5I2paqwaFuBLp+uUwdVVcnILmPuyixWHzjitR7yKKiomCQRRYV+Ka24fUwnxnRtzdShqbyzOhuXxspoIIrHwYpV2NwyblnB2AwiPyEfTAEMSfeKUOjBVOgRF2GisMrhsz2YPF6PrJKk1zvpNAOu/HyKX3yRDu/NRzSdHnNnURAA30FZMBMQAjTLC+5M4+ARKzvyK5kzRVcz1Kmfy/q2419f7/a7v6HBZ4RJ4vYxXj+5zzblHf3e1iXQ2mCbS2bO8kyuH9xBnwTQAbw+aH/53zZ+3lOMwyWjAsdX0toVb8C0IbuM3woq6dM+huev7MNHq37HpQpwEs9RsGIVkiDg9DRPMBV6bzMNBqXFkZGt102FItcP6UCYRm56MHm8STFmOic2T5GjTuiiyjKHHvob8TNmYD6n8fLmJ0uMRVuGPZhaDgXVb0Hy2cyH63O4emB73WdLp0EsJonrBqdg9lM71RBGSWRiH6965wfrcjTrHINJmyq1utjtp8ZYJ7SQFZVbFmzkp91F2I8GUvVhc8lsza3glpe+4+m9XxBpkk4mlgrqHQNeIQqtmvVTgR5MAT3bRlNQYafCpq1QpXP2cv2QDn73RQ++gtjzbqFy3afkv34D+W9Oo3rLEixd/njxhJskZo3ppM/W6ZxySufNQxBF4qbddFrbcWX/9pgk3+c9mAkIs0GiT/vQqhmyuTx8sSWfKXqqlE6A/O2S7nRLjvIrRuEPi1Hi/VsGHwvaS2qcmscFkzYliQKHK3yzOHRCj6e/3cPG7PKgfAadHoUcj5FFE29n8d2jSIwKCzrQCeYdA9AlMbLZxmZ6mh9gkETO7dCKjQfLGd9Dd6MPJRIiwxh3Tmt+2VuMWyOfvKFUCkFAr4fSOeU49uyhbN580j//DEE6vasaU4elMm9NNlqpfoHUcpgNIreMTEcSQ2sCYvG2QwxIjT1WkK2j0xAmSeDDiw1MX+RmZ5mAXa4/qJIEsJgMzJ8+qM5khT8vxWDSplRVxeHRFmzSCR2qHG4+XJ+jqQjZkP+TSzSwMqucR0SB1Q+dxy97injjxz3sK7FhkD2osowgiXgkIxFhEpVOhROHZYHWC0aYJGaP7XzK/g4nogdTRxmbLKOufxMKVHDbISIBUoZA+uiTyu3UaTk8e2UfJry6iqIqZx3D04YwG0XenjoQSzMtI+uEJorTyaEHHyTxoQcxtmt3uptD21YWBqXFsS6zVPP70tAEhApcO9j/inBLxiMr/LynmO93FnKkxokgQGJUGJf2SWbB2oM8PKH76W6iTkvAUQnbFsLaV4l0VPGxKvGRNJK3lIupUCOwY0I9LrHIYpRQUbmsb1vuHNeljklv7X637CscE0xtsCAIRJm1U3x1QocvNuVr1t8FImQCXu+ztz9fy52lG+mychUvWmuoGDGe0n6DkDt1JSo2mrT4cFQVLnplJfJJilUgwCW9m88fTQ+mslfBmleYnr3KuzKRW5vqJ4LJAuYYGHY3nHsDmKNPa1N1Tg2twk18Pms4E15bRaXN3WD+L4DZKPHatf0Y3jn0pJ11mpeSf7+CKb0jMZdffrqbcoznrvJOQFQGqchnMYr867KeZ129VKXdzburs3h/bQ5uRfExAl/622FcHoWdBZUMSI3VJaZ1/JO1HD65AVQF3DYAjMA0cSk3mZayQe3GZ54xHBLa4BLNxKb2YnSPVK4Y0J5IP89V9+RoNmjUhQcjs+7yKHRLijoVd6zTgnhrla/PYKBCJuA1bf8i08rs9DjavfwSYd26IYjaK64XdG/DL3uKcAToi1aLxShx9/ldmtV6JnR7dEWBb/8K2xeC246ISl0vegVcVu/Psidg7asw/VuI63iaGqxzfPKOoAAAIABJREFUKqkdFE7qm8xPu4sRBHwkaY2SgCgI9O8Qy8MTutO7fdOaoeqEHk6PzIHiGqrsHoySQEJkGKnx4cfyvK3r11P17bekL150RtXltWtl4dPbhjJ57nqqbU4UoeGaDrNR5J4LujB50Nm1KpVXZmPy3HWUWl1+zVBrawte++V3Pt+Szye3DiVRVwDVOZF938Nn08Bj19wtCDBU2MtQ096jG0Q4Eg3nLIN6AvTbxnRkZ0El1kbIrA/tGEcb/Zk9I7G5bSzJWsLiA4spc5ShohJliuL8DudzVderSLA0zaSvwy1TpKF+HIyQCYBoNuO8eirmhIh6j3t5cl8mz13HnsPVARtNW4wSk/omM3NU847VQzOYUlVYNAv2fH1s5qde3DbwOOCtcXDbSojVC4jPJuwumbsWbuWfE3tw5YD22Fwevt52iE825nGkxolHVokyGxh7TmtuHJam1zzoNJq8MhsL1h1kYUYuAsKxTGK3rJAUbWbW2E5M7BjF4YcfJvmpJzHExp7W9mrRLSmaj9sV8chON7sjk1FUVbPuMMIkEWU28tikHlzSO/k0tPTUUVzl4E//XUO5zYWGF6oPDo9CbqmNP7+xhm/vHk1MuJ42pXOUol3w+TS/gZQmqgLOKnhvAty5CcK0V47GdE3EbJQ0gykITGb9ttGdAm+XTrNQ7ijn1S2vsjRrKYIgYD/h2cmuzObtHW8zot0I7htwH+kx6Y26XrXDg1ESfQKbYP2fREGg2tFwVkOYQeKTmcOY/dEW1meVYnfLmobxAAZRwCAK3DQ8lYcu7tbsk4+hGUxtmOsTSKW9Uo3NDdn3RBJh8v5PeGeLiw93uFk+LeKPTmvBpXDXVpBC8093NvLEkl30ahvNFf299SjhJgPXDu5w1tZ16Jw+PLLCPxbtZNHWAr/Bx8FSG49/s5t/OV08PvJyuowadRpa2jCu3FyM77zBp58spCg6kffXHeTLLQVUOdzIiorFKDEgNZbbx3RiWKf4M2plramYsWATlXa3TyBVXyG2R1EpqXZy18ItvH/LEO0T64Qey54Ed91Z/wbHJeAdmziqYNvHMOQ2zVNLosBfxnflyaV7NCXS68MoCnSID2dYp/jg70nnlJFXlcdN399EubMcj6JtpO6UvSqOy/OWs+HwBv57/n8ZmDTwpK9pMUnIGrNGwfo/KXjfD4FgNkq8e9NAtuSWM3dlFiv2lWCUBGTF61coigKyonJF/3ZMH5F+2mxqQi8iUGRY+ZzmipSswqsbXDw8Kkzjg3g7LVsZ7P8eul96ihuq0xws3XGYtZmlLLlr5Fk52NM5c/DICtPmb2RTTlmDKQs2lwyCxKNKV9SNuWdcapyqKBz+56PEz5yJKS2NFOAfE3vwj4k9TnfT/KKqKhXOCqpcVUiCRKw5lghj/Wkm9bGzoJLfi2vwnDC4CKQQ2yWrbMguI6/Mpq9060BNMWT+gpZCZoPjEvCOZ9a+BoNn+hXMumFoKnsLq/l8c37AAZVBFIiLNPHhLUP092MQlNY4Wbgxl6+2FFBhc6OqEGUxML5HG25qguyWI/YjTPluChWOChQaTn9TUbF5bMz+ZTYLLl5A9/jghHCO1DjZmF3GhuxSzWAqGCETALdHpXVUPc/zCQiCwIDUON6aGseRGicZ2WVU2t1IgkBchInhneMJN53ecCb0gqnffwSPtp/UA8NNPL/GyexBJlqZ/XQcrhpY84oeTJ3hlFtdfLoxl4Ub8yi3uvAoKuEmiSHp8dw6uiN928eQX27n0cU7mTdtkK5SpHPK+duXv7E5JzhvDodH4bGvd5EcY2F019ansHXBUfG/z1DsduJuuvF0N6VBrG4r32R+w/yd8ymxl2AUjaiouGU3fVv35ebeNzOi7QikAFNUanl7VRauE4LiYAqxFVXlvbUH+eelZ24AqtNMbH4P7zy7LwGNSwBs5ZCzBtL8D2afuLwnMRYj76zOwi2rmgPjWsweJ+0TY1g4ayTxkYEPfEOZQxV2nvhmF7/uK0EQqNPXl9lcLFh7kA/W5XBuh1Y8Nqkn3ZNPTtTswRUPUuWs8gmkyleVc+SHI7iKXUhmiegB0bS5qg1ShLdvs3vszPp5Fr9c/Yvf/k5VVfLL7WRkl7HxYBkZB8soqXYyMDWWQelxXNA9kV/3ldSZRApGyARgSMc4WoWfnAhRQmQYE87AdPHQC6bWvOYNiDQY2FZibJqBF9c6efK8egotC3+D0kyI13OIzzTKrC4eXbyTH3cXIZ7QmdlcMt/tPMyyvcUkxYQhCgK3jelI35TQMg/VaX4OFNfwzfZDJ+XN4XAr/HPRTpY/MPaMmB12Hz5Myauvkvr+gtPueVUfqqoyb+c85myfU6eewK38kau/uXgze1bswWww8+KYFxmUNCigc9tdMt/vLPSRhg+mENstq3ySkcs/JnRHDDHPLZ0TyFrurcvWIOBxiccJ+RvrDaYEQeCvF53DhN7JvLPiAEu25GI0h+FSVBRVxSh5hWQ6tY7kBvvvDCtYReuo8xtzZyHDnsNVXPf2eqo00n5r8aZ1q6zPKuPKN9cyZ8qAoCfJ8qry2H5kOx61bmrfke+OUPJdCe1ntCeyRyTucjeHPjjEwRcPkv6PdMSjxs92j51VBasYmzIWAEVROVBSQ0Z22bEAyi2rDE6PZXBaHFOHpdItKfqYL+CB4hpW/X7EZ0U+GP+ns7H+LvSCqZK99e5+YlwYI+ZZuWdIPVGzZIIj+/Vg6gwjr8zGVXPWUmZ1adaiACgq2N0y2UdsiAIY/Ehy6ug0JfPXZGvOAgfqzVFS42RLbgUDUk+vEIWqqhx+7DHibpxKWJcup7Ut9aGqKk+se4Il2UtwyNqD1FpsHps3Bebn2Twz6hkuSL2gwfMXVzs0TYeDLcR2yQrVTg8xFn1lPKSxV9S7O6BxierxliEEQI+20TzepoJbir7hwF+eoKTaiUtWiLEYGZgaxzlJUSiOQWRdOomaNWuIHDEimLsJOWoVPasc2rVLWthcMjM/2MTHtw6lf4fA+/WP936Mop4oACFTvKiYdre0I6qPV4TE1NpEyuwU9j+wn8q1lcSO9l7D5rHx+ua3+D27AxkHy9h0sIwos5FBaXGM6BzPfeO7knacouyJdE6MpEfbaLbnV/q80xoSMhGAmHAjw8/C+rvQC6bc9Svl9EqUuLSrgWdXu+je2s9Au7bgU+eMoczq4qo311JS4wxIVQu8gdULP+ylVbiRK/q3P7UN1AlZbC4PX24p8JnJCyYlzO6WeWtlFnOnDmiWNvujcvFiPCVHiJ8x47S2oyHmbJ/DkqyGA6njccgO/r7q78Rb4jk38dx6j7U6ZU3jymALsQ2iiFUPpnQM9ac8BTQuATAGXotT/cOPJI8fR89+2kbgotlMm4f/TtGTTxGxeBGC6ezyhmtK7vh4CzVO7UCqvswDh1thxoJNZDx8PgYpsIndRQcW+QhO2H63obgVogfUTRuUzBJRfaKo2VVzLJgC2F+xh47KISb17cj/Xd6LpJjgJO/fnDKAS15ZRbnd5VddT4twk8R70weflSvxoTctb2g49/fxsWbe3uKioMrPUyIIYDr5wmWdpuefi3ZSatWWJ7buXs7hBfeS+/JV5P9nKkX/ewxH/i4A7G6Fv3/5G6U1zmZusU6osOtQFQaNl0cwKWGqCmsPHDkVzQsYd3Exxc+/QNunnkQwnrmD/2JbMe/89o5mIFW+qpzfH/mdXTN3sffuvRxacAjZ+kcxvkN28OiaRxu8RiR2FNl38HR8IXYgeBSFSHPozWnqnEBMSoOHNDguMVggqk1Al1PdbmqWLydqfP2rsJHjxmHskELpggUBnTcU2V9Uzf7Cas2xR1XGV5T98jYxQ6+h/Z0f0m7WfKL6T8D++4Zjxzg9Msv2Fgd0Lbfixuq2+myXa2QMkQYEyfc9Y4gx4Kmp21dFmsxMGxPLZX3bBh1IAbSJNvP5rGEkRIRpvttORBQgymzgwxlD6Nrm7DR+Dr1ePKYdOOpfUu8cJzK5p5HXMlz0TtSINxUZWp1Z6lqhTJnVxc97inxm/iGwNCoB+HRjHrPHdW7mluuEApU2bT+NYFPCbEFKGjclqqpS+MQTtLrmasw9zmzBhP/t+5/m9kBrCgqthew8spNeCb28H/S4oHgXFGyG/M1QsJnEikIU+TWgblAZbCG2xSgRVY/Zqk6IMOAmr5qfy3egXEuD4xJVge6X+2yusLn49rdCDlfYsbpk4iKMdKoooEtaGsakpHqbJQgCSf/4BwevmUzMpEkNHh+KvLs6G7fiWwsbaOaB1SkzZ0UmF/b0/7dVVZVKu5usslJEQUI+oV5KipTw1HhQZdUnoPJUejBE+vYxTk/jJpA7to7ku3tH8cIP+1i8rQARwecdZTaKqCqM79GGBy/qRof4s1e5NPR68aGz4dsHQSO6P55Hx4TxwQ4/pmJRSZDU+xQ0Tudk+CQjV1MHKdDOzOFReHd1NreN6aRZB6Gj0xj8PVPBpoSdzkez+vvvcWUfpN3LL5++RgSAW3GzcO9CXEpdxdZgagpcsov31vwfL5rSoWCT10w1Ng3a9YeUQTB0FubE7vxp0R4+25yHfMI4KtBCbJNBZOqw1DNCVETnNNPxPG+2Sz3BFPgfl6iCiND1Ioj4oxblt/xK5q7M5KfdRYiCcEwOXRTArMqYO1/PjF8PcN3gDsRG+E/hM3XoQOz111H03HO0//e/vRsLd0Lxbq/3pjECWqVAh+EQYjXIqqqyeGuBTx8AwWUe/FZQyc+7C3F4FIqqnBRVOSisdFBY5aDo6I9REmkTbUKOk32EH8M7hyMYBKo2VxEzOObYdtkhU72jmjZX1V2xVFGJNDXejykhMoznruzDo5f2YNHWAr7Ykk+5zY2iqsRYjEzonczkgSn1Pl9nC6EXTPW8Ar59wGfzwXvrLj2mxIg4HtGQrTRFwIh7/Xo56DQ/CzNycWiopAXTmTk9CltzyxmYFncqmqgTwsRFmFA0/GOC9eaIPEUrGGVHbQRW/n6ESpsbk0GkbSszkwd1YFTnBJTKCgqffpqU119HPMPrJvaX70dWfFfwgqkpUFBYVbEXOg2D8x+Dtv0gzDc15ZaR6SzaWoCsMSvdUCE2eMdDU4emBX5zOmcvogjD7oJfnwbPH3XdgY5LnKqRgi430wnvAP8/vx7gv78ewOVRfNLPFBVsSNhkeG3Z77y1KouPZwylR1v/Mt3xt95K9qSJOD5/EnPRN1CRA4IEisf7ryCAKRyG3gn9p0J4aLxHHW7Fr9hVMJkHHlnlhR/30TEhkjbRZpJizPRIjiYxOoyko7/X+ihN/DKF3OrcOp+XwiUS/5TIoQ8PIZrFOivvxjgjrYbXVSyWFZn2kU1XJx4RZuCGoancMDS1yc7Z0gi9YMoUDgOmw6Z5dTqtgBEk6H1V07dL56Qpb6I0qiN63ZTOKaBXuxjCDBJWZ91BfjApYUZR4DI/heIny4Hial7+aT8/7yn2sRHYlgcr9pUQbjJwTeUurpswCUu/fk16/VNBpaNSc6WnoZoCe07dd4EdFXXEPfWuGnVpE0W/Dq3YnFPud0DljzCDyPndEk+qXkHnLGXoLNi7FA5tAVnbC1MTYzi5addx3VIPN1ceoMbh4b21BwPys3O4FRxuhavnrOWL2cPplqQdUInWAtLPy4LtW0Hyc15XDSx/BlY8C5M/gM4Nq2K2dFweBVEEWSMDO5jMg0izgWeu6BOQqt/0XtN5fuPzx6weamk9oTVShEThp4W4il2IFpHo/tGk3JaCaPxjxdAgGJjUaRLhQYiV6DRMaK3J1nLBv6BND5CCNKIzhsMNn+viE2cYHo2ZYajbmTWEqqqaHkA6Oo1FEgVuHpGG2ejb3UYPvoLY826hct2n5L9+A/lvTqN6yxIsXequpoqiwPThaU3WphX7S7jsP2v4fmchLo+iOfCyumRKapy8Qyp3Rw2j2uEn7bkFcHxNwYn4qykIhLlTBpIYFVgRdi1GSSA1PpyXrjnzg1OdZkQywg2fQVIfr5hEIBjDod/1dL3+Jb6+aySLtxUwd2XmsZS+QLG6ZK5/ewNWLUW6kv3w9jgEVwWiv0CqFo8d3Db4ZArsWRJUG1oikWYDHj8TKcGI0SiKSrQ5MFGfCekTUP1I6MWNiaPLU13o+XZPur/WnXbT2h0z7K1FEiWm9JgS0LV0Aic0gymDCaYu8qZvGBvutBRBxIaZisvfhw5DGjxep3mpXf4+kWA6M1EQiNbliXVOEZMHpaB4tAc4kT3HkXzTK3T4yxek3PkhiVf/C3P77sf2C4J3dSstoWkmcTZklXLbB5uwueSAbASckpGdh6uZ8s4GnH7u4UwhxhyjOdA4vqbgeGprCiJ61P3bWgyWgGqZYsKNLLpjJOkJEVg0guUTsRgleraN5rPbh2MxnbmGxzqnCXM0TP8OBs8EU6T3RwtThLd2+5LnYOJLIAi0a2VBVlS/3+n6VG0BHG6ZRVsL6n7IUQnvTQRnNYJGqrJfPHb48lYo/C3wz7RAJFHw2y8fn3lg278Oxe1AlT3YMzdR/uu8ugcLkBIXWAAdbgznxh43YpaCX9U2SSYGJQ2iY0zHoD+rUz+hGUyBt9OathTGPQJRSSjGCHzmXAwWMIQh9ryChf0+4K4N0SiBmhjpNBsDU2M1S9iC6cycskKvtjG+J9HRaSSesjIcf72X24s3YDYEX2sZYTLw/FV9mqQtVqeHWxZs0l6Jqmew5ZJV9hVV8/z3+5qkHaeKc2LPwSD6Tq4cX1NQvaMa1aPiKnGR90aeT02BKIiMbj864Gu2jgrjm7tG8o+J3WkfayH8hCBJFLxBVKfWETx+eU/+d9tw3VdKxz8GE1z4BDyQ6Q2UkvtBRIK3bi+yDXQ6HyZ/BPftgf43HvvY9rwKDlVo+6oFItFtc3lV5epMRmz9EFzVcFwglfZKNYkvVGN1/bHtnS0uxr53gniG2w6/PtW4v0ULYNaYTj7f+VoCyTwwSgLXDe5AmCHwyZU7zr2DEe1GBBVQmUQTKVEpvDTmpYA/oxM4oVczdTySEYbfCUNns/KHz4jevZD+rWzgcYIlFjqfD+dOAUssN8kKP7y9gTeWH+DO87qc7pbrHMfM0R1ZfeAINpfvrHkgylqCAGO6tqZ1VJBpnzo6DVCzZg2H//4wMZdfxn133YWw4iBvrcwKKA1HELyB1Pu3DKZT68YrLwEs2lqAorFyE4iFgMOtsDAjlwcuOgez8cxcVTGIBq7rdh3zd83HKdetgQy0psAkmpjWc1pQ1zUbJaYMTeOGIalsyinnlz1FvLUyi8v6tiW5lYWLeybRN6VVwyfS0anFaIa+13p/AmDe6mzNleNgzMFLrS625JYzIDXuqLnd696g6ARkFV7d4OLhUfW9M1U4sAxqiiEyMaB7aIlM6tuWx77e5Xd/Q2I0oiBw07C0oK4pCiIvjXmJx9c9znfZ3+GSXSi+ywHHsBgsdI/rzhsXvKHXSp0iQjaYKq5ysOtQFVUON2EGif8VdWTc8Ffo7+ehNkgir17Xj0mvr2FQWhxDOsZrHqfT/AxIjSU+woTNpS0o0lBnZjFKzBytL3vraOCywa4vYftCsB7xerlY4qDH5dDverBoD5BVl4viV1+laslS2j73LBHDvDOR943vSmp8OE8s2Y3bo2DVmAAQFRmTyUinxEheu+7cJgukVFVlzopMn0mHYAZbAN9sP8TVAxs2GT1dXHPONczbOU9zX9yYOOLG1K80lhyZTM+Enid1bUEQGJQWR/8OscxdmcXL1/RD1O0WdJqBnYcqNVP8gjUH31dY4w2msld4pc81eGC4iefXOJk9yEQrcz3PtyDApvkw9qFAb6PFYTFJzB7biTeWB1+rZjaKXNgjiZS44AMcSZR4YsQTXNn1St7b+R4r81diEA24FTeqqmKQDCiqQu+E3tzc62ZGtB2BFKAYl07whFQwpaoq67JKeWtFFuuySjEZRBRVRUDA6vSwPa8Cp0fh6oEpmmkYyTEWXri6D/d+uo0ld40kPlJfyTgTEASBf17ag7s/2RqQgtHxmAwivdrGMDC1YRUdnRDCVuaVKd72kXdAcKL/S+EO+OVxb1B13iN1TLxdOTkU3P9XDAkJpH/1JYa4uoP3K/q35/J+7Vi2t5g5KzLZnleBoqqoeAP78+0F3NSpFQOnX9Kkt7SzoIpSq69KWDCDLZtL5t3V2Wd0MNU6vDUz+8zknZ3v4PBopz35wyyZeWL4E41ug83lIdwo6YGUTrNxolpoLcGo2rplhRrnUaGZgs3g1v7+DGwrMTbNwItrnTx5Xj2pZh4HHFwJnL3BFMCd53XmQEkNP+4qCjigMhtEuidH8+LVfRt17b6t+/Lvcf+mzFHGqvxVVDgrkFWZGFMMg5MHkxJ15vbVZxMhE0xV2FzcOC+DA8U1x2ZmT1RvK7W6eOnH/bz0435eu+5cxvdo43Oececkclm/tvzlf9uZP20QHkXlx92FvLUyi6wSKw63jMkgkhxj5uaR6fypXzsidIf7U86FPZO4c1xnXvppP36EbnwwGUTax1p4d9pA3ThT5w/KsmH+BO9KlOJHotht8/772+ew/3u4cTG0PZfKxYspevY5Eu64g9gbrvf7XEmiwPgebY71MQ63jEEUMEgi1cuWUTpvHky/pklvq6DChqglGx6khUBhZXAByulgZp+ZFNmK+CbzGxxyYO01S2aeG/Uc/RIbr7Jnc8mE6/2+TjOipRYKwUl0G0QBS62gk60MVP+BwRPjwhgxz8o9QxrwnrNX1L//LEAQBP59TT+eXLqbjzNycXsU/LklCAKYDRIjOsfzn+v7YzI0jXRBnDmOyztf3iTn0gmekBCgKLe6mPjaavYcrtKsqzkeu1vG7pa5a+EWvtySr3nMXy88h2qHm2nzMxjw5E889MUOduRXUuP04FFUbC6ZzBIrTy3dw8Anf+axxTvPeBWslo6sqGzLq6R/h1aEGUTC6umgBCDcJNGvfQyL7xhBVICSpDohQHURvDseagr9B1LHo8rgqESdP5GiB2/lyFtv0+G9+cRNuSGoAN1slDBI3mc2cuRIXAcycRcUNPCp4PCq9/m+4YOxEADfSagzEUEQ+OfQfzK732zMkhlLPVLT4YZwEswJzBk/h/NSz2uS61udHiJ0tT6dZiQ1XltVLhhVW4PknWAEvLLr9dArUeLSrgaeXd1AP2kIDS81URR4dFJPvr5zJFcNTMFsFIkMM2AxSliMIpFhEmEGkfE92vDhjMG8fePAM7b2VCd4zvqpM1lRueGdDRRXO4IyVnS4FR7+6jdS4sIZlFY3TUcAosxGVuwvqfcctYHbp5vy2JpXwUczhugD91PE09/uwer0sPDWYVTYXSzckMv8tQePekCoqKq3s3N6FEZ1SeC20Z0YlBarr0jp1OXz6WAv99ZGHSXtlWpsbsi+J5IIk/d5eWeLiw93uFk+7egAxmUl3vI9rT/bhxjeuAJfwWQi6sILqfz2WxJuvbVR5zqeyDCD5srU8YOtiG4jGzxPS5H0FgSB6b2mM/mcySzNWsq8nfMotBViFI2oqopbcdO/TX9u7nUzw9sORxSabm7R5pL9Wjbo6JwKpg1PY9PBMp86zGDMwQ2iwKjOCd5fopO9AVXtKrwGj481039uDfcPq6fkoVVopZl1bRPFc1f24Z+X9mDNgSOUW10oKsRYjAzpGEeCXh5yVnLW9/a/7i0mp9SqGUhZdy+nauMi3KX5iCYLxsSOxAy/po561ZNLdrP4zj8GGKqq8sDnO8jILgu4DQ63wt7DVUybv5FPZg7FKIXEgmCTUbuq50869P11B/l1XzFfzRqBySCSGGXmngu6cse4zmzKKaek2olbVog2G+mTEkNiVGjMlOkESWmmt05A8TWubEi9ShDAYJKhcCN0HNPopkRfOpGip55u0mCqa5so3LLvqlIwgy3veZpGEKO5CDeGc/U5V3NV16uoclVR5arCKBqJCYupd8WqMVidHiLCWkbQqXN2MLpra8xGSVPUJhBV2zCDyI3DU4+tkNP9cvj+b/Ves3OcyOSeRl7LcNE7UWNcY4qEAdMbdV8tlcgwAxf1TDrdzdBpJs76YGrOykzNziUQKWCAfYXVHCiuoXOidwCxYn8JP+wq1CwyrC84c8kquw9V8uG6HKaPTD91N3wWoKoqGdllzF2ZxarfS5CPShRJosCozq25bUxHBqfHIQgCv+4t5vVlB/ji9uHEhNdd9TNIIkN11UWdQFn/JvhJdQtIvcplhTWvNkkwFT5wIHJlJY79+zF37dro8wGkJUTQPTmabXm+NQyBDLYAIkwSt43u1CTtaW4EQSAmLIaYsFPvJ6evTOk0N5IocNuYjvz7p981xyeBSHRPGZL6x4aIeOh6Mez5ps5K/Yk8OiaMD3a4tXeGRUN64J5tOjotlbO6t88ptfJbfqXP9mCkgD2Kyvw12Tz1594AzF2RpVl3FUhwZncrzF2VxbQRaXp6mR/WZZby18+2U25zYXfJdTzXFVnl133FrM8uJTbcxB3jOvHij/t5+8aBdIjXvRN0GoGqepX7FO1BQcDqVQdXewu3w+uX324IQRSJnjiBqiVLMf+laYIpT2kp1xZtZp+nPXaD7wpbQ4Mt8NZ2je7auknaczZjdekrUzrNz4yRHVmbWcr6zFIcQdQ2mo0ir113LonRJ/Rtw++B33+s4zV18N6oOoekxIg4Hon2PanRAiPu8S7b6+ic5ZzV+WZbcssxaEjTBiMF7FFU1hw4AkBemY0tueU+x9QGZ3HjZxF+znBEkxlBMhDeeYhPikyV3c26rNKTvKOzm6+3FTD9vQwKKuzYTgikalHxzvoWVNh5+KudTOqbzABd1lynsbhtINdfSP3EuDBez3BRYq1nkGIwQU1RkzQpZtIk9v24kmeW7ubGdzdwxRtruPHdDTzz7R5yS/3XMZyI4nBwZO5bZE28lJGmGuLjo5FOYoCbCwBlAAAgAElEQVRjMUrcfX4XJF3uu0FsTn1lSqf5EUWBOVMGMKJLApYAxA0EVcFsFHnhqr6a6sW0HwBD72hQjMIHKQzaDYJBM4L7nI5OC+Ws7u2r7F51vRMJVgq4xumtofhu52FNNaxgfVo+25TP8E4JAV07VFhz4AgPfrEjaJ+oTzfmcUH3Nozqos+W6zQClw1ESbNeqpbj1au6t/Y3DyXUW7AdKGsOHOHV5WVs6zMNZXU2nuO6nXVZpby39iB9U1pxz/ldGNFZuy9RFYWqJUsofuUVLD17kfbJQkxpaSwss3Hp66updrg1TT61sBhFJvVN5sZhqQ0frONdmWohQh06Zxdmo8TbUwfyycZc3lieSZnVhd0t17EMCTOIqED/qgLu7hHP8L5t/Z/wvEe8KcxbFgTUt6mSGaFdf7j+E5DO6iGmjs4xzuon3SiJmivMwfguABhE78DpcIW2ImCwwdnhFuDT0pzIispdC7UNdwMRCbl74VY2PTJenzHXOXnM0SD7yfs/jgbVq1QFGlGTo6oq/1l2gDeWZ3rrHkQDJy7Ruo8qVGZklzFjwSZmje3EXed1rpM6bN2QQfFzz4HBQLsXXiB8wIBj+1Liwvn6zhFMnrueSrsLez0TGALewdm1g1P458QeenpygOg+UzqnE1EUuH5IKtcN7sCmnHI+Wp9LXrkNh1sm2mJkUFosU4akEpVzgLzZs5CvvAgp0o+wjCDAJc9C23Nh2f+Bvcw7+XRix2SKRHF7qK7qSszD3+iBlE5IcVY/7a2jwo4GQnUHC8FKASdEeU3p/PmrBBucuXTPqTos21uMU6NgNlCREJdHYdneYu00BR2dQDCEQVQSVB2q97AG1atUBWLan3Qz3lyR+UcgFQB2t8ybyzORRIE7xnXGmZVF8Qsv4ty/n8T7/0LUJZdoBkCp8RH8fP8Yvticz9yVmVTY3Djc8rGVKrNRRFXx2giM6eRjD6FTPzZ9ZUrnDEAQBAalxfn//vbuReSIkZTOfYvE+/9S/8n6ToY+10DOWlj3OhT+5g2qDGav/PnQWZA+npLLr0Bav4HIESOa/oZ0dM5QzupgalSXBM20vGCkgMNNEtcP7gBAfKS203ewwVmr8AYcw0OMOSt8FReDEQmxumTmrMjUgymdxjHsTlj2ZIOpLH7Vq0Qj9LsejCcnvb/pYBmv/3IgaKVQu1vmP7/8TqflX9Pxx8+JnzGDdq++gmiqv5+JDDNw0/A0bhyWyobsMjbnlFNudWE2SrSJDuPiXsm0jtI9UU4Gq1MmPkL/2+mc+bS+7z6yL7uMVpOvwdS+gYkgQYC0Ed4fDUSgzUMPUvTMM0R89RWCUffV1AkNzupgymyUuHpgCh9vyPFJzwtUClhV4U/ntgNgSHo880zZjTLFsxglxupqWMewOj1s15BqDqYODWB7XgU1Tg+RemqNzsnS7wb45XGfzQGrV4kSDJl10pf/768HcJzkCq3dLbNATeaDb5diiA1OkEUQBIZ2jNdtBJoQm67mp9NCMLZJJG7aTRS/8CLtX32l0eeLPP98yj76iPJPPiVu6pQmaKGOzpnPWT/yvHlEGp9k5OKT30vDUsAmSeCK/u2OqTIN7xRPpNlw0qZ4ACoqVww4+TSgs41ymwujJOI5wd8n2Do0oyRSYXPpwZTOyWNpBf2nwdb360gBB4QhDNLHQELnk7p0UZWDtZmlPr1UwCu0gsAWoRWlkgV9ffb0Y9V9pnRaEHHTp5M5YQK2jRsJHzTIu7GmBGyl3tRlS6w3DTqAmklBEGjz97+TO206kRMnsKbEw7urssksqcHulrEYJTq2jmDGyI6M7tpar3XWOSs463v71PgI/jK+K6/8rG1k5w9JFEiKsfDQJd2ObRNFgVtHdeTFH/dpiiU0FJxJAlzWt60+4D8ORdHun4OtQxMEjpn76uicNBc9DSV7IS8DPAEGVJIJYtPh6vknfdnPN+dpbg9mhVYF/rcpj7vO63LS7dBpGmxOfWVKp+Ugms0k3n8/Rc8+SdojkxHWvgplWd6+DbwqpxGJXt+ovpMhLKre84V16cJ3509lwfMrcZrMdSagK3BzuNLBttwKzEaJO8Z1YvqIdF3cRqdFc1b7TNUyc3RHbh6ZhsUY2O2aJJHkGDP/u20Y0ea6Ob+TB6UQYzFyMpMpFpNBH+icQIzFiFv2DUyPr0MLBLesEGPR87N1GolkgBs+hy7jwRih6XVWB1MEtO0Ht/zk/e+T5PeiGk2Bm2BWaF0ehQPFNSfdBp2mQ1+Z0mlpRKc6Se2xCr65zzuhJLvAVeP98TigMhd+ehRe6ALr5/g9j0dWuGvhVuaQShlGzUwe8H5HSq0uXvhhP3d8vEVzHKCj01IIiWBKEAQeuKgbz17Rh6RoM+F+VJYsRpEwg8glvZNYevcokmJ8C8mjzEb+d9swoszGoIwvzUaR96YPIiUuSPO7s5xoi4F2rSw+24+vQ7PtX4fidqDKHuyZmyj/dZ7P8W1bWfRgSqdpMJjgmvfhuo+RY89FUUQwWvB2l4JXvcpghvaD4M9vwbTvvNLqjaDWy+5Ejl+hDYRqh3+fLJ3mw6vmpwdTOi2E9XMQFs1GlGQEpR7rFrfVu2L/y+Pww8M+u1VV5aEvdvDLnqJ6LReOx+6WWba3mAc/34GqIRimo9MSCKne/vJz23FZv7asyyxl7spMtudXYnN6cMkqCZEmLu3TlunD00hNqH+GOTU+gqV3j+T6tzdQanVidfof6ESYJBRVJTU+nD7tWzX1LbV4BEHg9jGdeGLJbmwnzGAFWocWbpKYNaaTniag03QIAnQcS1HBYMK7X0lsr3Cvv4qqgrkVdBoHCU23yuxvIiBopVB9QuG0oaoquw5VUVLtpKTKyb6iKjrEhRMTrv8/0TmD2f01/PyvwNOawat4umkexByVRD/KT7uL+G5noWYgVZ8iqcOt8MOuQr7f2YZLeic3wU3p6DQvIRVMgXfwPjg9jhqnh2pHJrsOVSGgYnfLfLEln48zchnTtTW3je7IgNRYvwP09rHh/PrXsazYX8ycFVlsz6vAZBBRVRVBEHDLCqnxEcwa04mLerbhzo+38ux3e3l0Uo9mvuMzn8v6teXxb3Zr7muoDg2849vL+tXj4K6jcxIodjs1vy6nzYPfQkLCKb1Wr7YxfPvbYZ9BSDBKoWajSM+2jVsh0wmeSrubLzbn89bKLKodbkRRoMbh4bGvd/GPr3Zycc8kbh3dkV7tTt7MWUfnlKDI8M09PoFU2ivV2NyQfU8kESbvGOidLS4+3OFm+bSjk81uuzcI63fDsZX5N5Zn+kyKQmCKpDaX1zNPD6Z0WiIhF0ztK6xm6rsbsDo9dXJ5j19d+nlPEWsOHCE9IYIFNw8mIVLbL0QSBc7r1obzurUhv9xGZomVGoeHcJNESpyFzol/FGm+dE1fLn19NQPTYpmgdxZ1qHZ4aB1lIr/cTrAaEhajxIMXn6PXJ+g0OTUrVmLp3QvDKQ6kAK7o355nv9+ruS8YG4erBqSc8rbq/MHSHYe4/7PtCAg+Ake175Rvdhzix91FDE6PY86UAVh0M1+dM4X9P3hrozSQVXh1g4uHR9XjlyaIsH0hDLmNzJIa9hyu8jkkGM/I/UXVHCiurjN20tFpCYTUCHRbXgXXv70eu0uut7BcVb2zJPsKq5nw6iq+uWskbaLrN+JsHxtO+1j/9VCtwk28cUN/ps3fSLekKDq2jqyz3y0rVNndmAwikWGGlpey5rLCb5/D5vegpsir/hMW7U2HGnI7xHfS/NjaA0e499Nt3DgslTKri4UZeQGrLlqMEtcNTmH6iPQmvBEdHS9VS5cSPWFCs1wrJtzIxT2T+GbHIc0JhYZWaEUBLuzZRk8pa0Y+3pDDE0t2ayq7Ho+ieutC1meV8uc31vDl7OH65I/OmcGaV7wCExo8MNzE82uczB5kopXZz3jEbYO1r8HgmXy5pUBTUTcYRVK3ovD55nz+dkn3oG5DR+d0EzI9en65janvbtBcgvaHR1Eps7q49q31fHv3qEbPKPZp34r7xndl9kdb+Gq210H8m+2HeHNFJgdLrZgkEeVoAebE3snMGNUCUkMclfDTv2DHQkD0FqjWUlME5Qdhy/uQ3BcufApSvB4WiqLy5opMFqw9yL8n92NEZ+/sf1KMmZd+3I8g4HeQYjaKqCrcN74LM0drB2k6Oo1Brq7Gum4dyU/+X7Ndc9a4Tvywu7DBwbkWJoPI7LEn53GlEzyrfz8SUCB1PE6PQvYRKzPf38wHtwxueRNmOmcXsgfyN/rdPbCtxNg0Ay+udfLkefVMJluPQGU+eWU2PBrBVDCKpLICeWVBevzp6JwBhEww9fqyA9hc2kpX9RVGehSVwioHi7cVcO3gDo1ux5QhHdiYXco1c9ZyoMSKIHAswDteGvmb7Yf5YVcRHeLDefOG/j4rWWcEVYdg/iXef/2kCqC4vT95G2DBJPjTG1R0vJT7Pt1GtcPD13eOrKOaOHN0J64ekMKnm/J4d1U2VpfnmKmfrKhEmAzcMiqdyQNTiI0wNcdd6pzlHKlxsubAESrtbgBiw0303p9B+ODBSDHNN5lxTusI/u7Zy9NKR5xi4F2z2Sjy3BV96J6s10s1F499vdNvIFXf+8TpUdiSW87WvAr6d4ht5lbr6ByHoxJEI8hOv4c8MS6MEfOs3DPE/7tWdrgpmHo1R9pcAwnn+OwP1jMyGD9QHZ0zhZAIpmqcHhZvK0DLxiCQwki7S+bNFZlMHpTSJLOJoiCw81BVvamGsuoVxdhfVM1l/1nDRzOG0DflDFIDtJfDuxdBVQGoAXZ+HjvKV7P4PzGTTn0n8tAl3TBKvur8sREmbh/TiZmjOpJ1pIZyW+0g10jHhEhE3TFdp5GoqsrmnHLmrsxi5f4SDJKAR1YRAIMk4nKIjOs4gdl5Fc3yvVPdbg499DfGlpURO/tyHvxmH25Z0eyzapFEMEoiL1zVh0l9253yNup42ZFfwaEKbfnoQN4nDrfM2yuzeHPKgOZsto5OXQQBGnDS65UocWlXA8+udtG9tbaTjmgxk/zkU7TbGQW/lfjsD1aRNE6fJNVpgYREMPXVlnxEjSAomMLIkmonW3IrGJDauNnEZ7/byw+7iho2Az2KqnqDwSnvbmDJXSNJjT95Y9AmZcn9UFNYJ5AKRAFIlB08x0sYxt/uHQ3WgygKeiGqTpPj9Mjcs3AbK38vwe6WUVWoY/PkUUCQ+KkUVr61nkt6J/H8lX0wNPC8niyK00nBvfeBqpIydw6pYWH0TE/knVXZfLXV23cdn54cftRu4c/ntmfGqHQ6nYmr1mcxb6/KwunxnUAK9H2iqLBsbzGlNU7i/Ygb6eiccswx3trmBnh8rJn+c2u4f5j2syqoHoxp3Rgmq3y7r8zHpDcYRdJwk8Tg9LjG3ZeOzmkgJIKpn/YUa9ZKBVMY6ThaQNyYYCqrpIb31h6sk85XS32pIQBWp4dHFu3kg1uG+Hy22bGVwd4lmql9gSgAGUQBtn8Cg289la3U0fHBLStMfSeDHQUVAQsHfPdbIeVWF+/cNOhYymlToVit5N15J1KrVrR7/nkEo1dAolPrSJ65ojePTOzO0h2HOVBSQ6XNTUy4kc6tI5nYJ5mIsJDovs841mWWaoqEBPM+MRpEduRXMq5b4ilooY5OAIgSpI2C7BX1HtY5TmRyTyOvZbjonagxoRTdHqLbMqG3wiOLdmqeI1BFUoBJfXSbE52WR0i8jStsfqQ/gyiMVFQorfGfWxwI89cc1FS7CSQ1RFEhI7uMw5V2kmMsjWpHo9nyvlcSVYOAFYDWvAqDZhxNNdDRaR7+8dVOfgsgkDqeWiW2p7/dwz8vrd8nTlVUPMU2lKOpqWK4EUNiOIJGECZXVZF32+2YOqaT/MQTCJJvPxQRZuCaQbrc+ZmEPxGjYN4nqqIeq9HT0TltjLgHCjb7VfSr5dExYXywQ+N5NUbAyHtBEDAbJSYPSuHD9Tm4Zd9xTkOKpEZR4KoB7XXrAJ0WSUgEUwY/s8nBFkbWpvlkldQwf81BftxdSI3Du0weZTEysXcy04ankRLnK5Fud8l8vjnfR+0mmFRDFfhgXQ4PXtytwbaeUjbN8+uWHrACkL0MDm+Htv1OUSN1dOpSVOVg0bYCXCexMmx3K3y4Poe7z+9CjMVXfly2urFuLKRmdQGqS/ZqlQMoKoJRJHJEOyIGJyFFeusBPGVl5M6YQfiAgbT5+98QxFOTQqjT9Eh+JoCCeZ8IApr1ojo6zUrHcWCK8AmmDt5bN70+JUbE8YiWwI0Kva469tutozry+aZ83HLD6YMnEmYQmTm6Y9Cf09E5EwiJ3tyfR9TxhZENYTKIKIrKn/+7hkteXcXCjFyKqpxYXTJWl0xhpYP31x3kgpdXMHnuOg4UV9f5fMbBMs0UoWBSQ1wehcXbDjV43CnH6ltkejxPjAvj9QwXJdZ6Zv8FCaoLm7hhOjr++XB9DlrD4KqMryj75W1ihl5D+zs/pN2s+UT1n4D99w11jhMFgc825fl8vnplPoef2UDVL7koNW5Ul4LqkL0/LgXF6qHq11wOP5tB1a+5uAoLyZl6I5GjR9Pm4b/rgVQLo5UfL69g3icgEB+pF9rrnGZEEa54Cwwnke1isMDEl8H0x+Rx21YW3rt5MBZjcKtLZmSezl5CEo3L/tHROV2ExFv8yv7tidBYOj6+MNK2fx2K24Eqe7BnbqL813l1jlUVlQ835LA1rwKnR9H0U3DLKk6PQkZ2GZf/Zw0Z2WXH9pVbXagashPBpIYAVDtOPjWkoMLO55vzeXd1NvPXZLN4W4HfFMh6aaBo9XgFIL+oqt/VLR2dpsYjK7y/LsenXrF2ZThu/CzCzxmOaDIjSAbCOw+pUxgN3nS/t1dloap/fI8rlmRR9VMOeFSoL3XQrYJHpeqXXA49/DHRl11O4r336l5DLZCrB6YQZvB9dQbzPpFEgYGNFDPS0WkSOo6Fy14LLqAyWmDs36DfdT67BqTG8snMocRYjIQ3kLIXbpKINhtYOHsUQ/umkXvjjbiLi+v9jKKo2F1ynX5YR+d0ExJpfuO6JWIyiD4qMxB4YaSsqrjdgX15VcDqkpk2P4MvZw+nW1K0t1ZK4+PBphoG238oisqqA0eYuyKTzTnlSKKAW1YQBAHjUTnoC3u0YeboTvRuH6CnjtHi31fqKA0pACEIXjUhHZ1moKTGqanAFszKMEBpjYsap4cos5HqtQVYNxxGDcZk16NiaDMAc3fdbLqlcv2QDvz31wOa+wJ5n4QZRG4clnrK1CF1dIKmzzUQ0Rq+nOmtafZTQ6VKFhSnC+GyfyP29w2kaumb0oo1fzuPRVvzmbsii1Krd7ygKOoxa5O4cBO3jenIn/u3JzLMAH/9K2JUNDlTptJh3ruY2rc/dr7iagcfrc/lw/U5lNlciAioqLSPDefWUel/nENH5zQREk+fJArcPDKd/y47gEOjXqKhwkhAU72poToLm0vmtg82s/yvY4mNMGrOQgfrwRCMgpfV6eGWBRvZkV+pUTStUuthvPS3w/y8p5grB7Tj8ct6NahYpib1RTi4st5jGlQA8jghsWfA96Kj0xiqHR4MogjU/f4HuzJslESqHB4iRJGq7w9qBlIZ+Tt4+tc32X/kIKIo0iU+lcfOv4t+yd29B6giVT/lEDEkCVEfALQ4EiLDGHNOa37ZU6TpAxbI+2TK0NRT1DodnZOk0zi4fx9kLoM1r0DOWpCOprTKbkjqgzDyXgpeW0rEdifx/es/XWSYgSlD07hhSCrb8io4WGqlxuEhIsxAWkIE56a08hkTJdw2EzEywhtQvfsOrnYdeODzHSzbW4wAxzIL5KMz07llNp75bi9PfbuHG4ak8vdLuumTFDqnhZB5k88Y2ZEl2w+TVVKDWysy8oNRElAUlRPFaQJR4AOvP9WPu4rILKnB6vRNjwvGgwEg2mJgbeYRhqTH1xv02F0yV765luwjVk0p9uOplYD+YnMB5VY3/7n+XM3Ar6DCzmeb8sgtGMOTZBCOtnFlLX4VgBC8HXdUm3o/r6PTVIQZRBSNZd1gV4YVVcVsELHv0K4brHZamf7533jqwr8wqds4XLKHjPzthEkn1McIYNtaTORQXQa4JfLkn3qx+WA5ZVZXwJ6BABajxD8mdPNbx6ujc1oRRehygffH4wJHJagKWFqBwZtl0ubB3uRMmUrMny7HENtwqqogCJzbIZZzOwSW1hp3ww1IkZFsu3U2D5z3FwodiqZoUC21E8Ufb8hhz+Eq5k8fRJhBVwTUaV6E+vJOBw4cqG7atKkZm3NqKa1xcvXcdeSX2+v9ctZiMUqYDKKPhK3itJL/35uIn3BvQKtJRlHg8nPbkV9uIyO7THOVq2bXr1RvWoy7NK9Oaoi5ffdjx5gMIlOHprIhu5TCSgcX9kxiYu9khqTH+czGzFiwkVW/H2kwkNK659vHdOSeC7oCXoPTn3cX88nGXH4rqGRSn7b/z959h0dVpQ8c/547NZOeQAgl9A4qIkiRKva+9rpi19VVd9ctP1fX3XWLq26zgohY17a2FbErTToqVXoLJCG9z2Ta+f0xk5BkJpmZkARI3s/z5DGZueXcwfveeU/lijE9Gfn6yVDVfN/mJlnj4eq3oe8pLdtfHBal1Fqt9ZgjXY7DFUt8qqzxcuIfPwuZsjdwL/+Y9HN+Ft297PfyackCbN0uRXtDvxCvy93C1W/+nE33Loh4LFOqjcxfjZVxU8eonQWVXDZrOWVOd9gWqsbsFoOfTBvI3TMGtX3hjmEdIT51tO9OjeX98WFQiqRf/4b/fZ/DB9/nUBzsypcWb+XCUT24YFQPHNaW1ddXu72c/7fP2FPhxRdlrwEAu9lg6pCuzLr2JImrotU1F5s6TcsUQHqCjQ/vmsRv3l3PZ5sOoiBstz+H1YTWcNW4LF5Zvjfk/VjHWWjgrxcfx9a8Ci6dtSzsGjfRdA0Z0T2pbp2bvUVVLNiQx98+2cKBEidnjOjGOcd1Z3z/dPYWVbO0iUQq8hTQPmYv3sW0IRl88H0OH3x/gMHdErlibBZzfjwGe+0sPac9BAt+GehfHQuTFboOhT4TI28rRCtJsJmZOKALi7Y1bFGKpWXYUHDmsAy6dL+M8oXhZ3Trn5aFoQx+9tGfuWDoDE7sOYIUe2LYbX3lbvxVnrrp0sWxZUDXBD65ZzK/+u96lu0qatANqT6H1USi3cyD5w7nvBOkJVJ0ADfeysP/N5sv//gZhmGEDCNYt7+UP3y4mYtH9+Te0wbTNbGJsdNNmLN4N/trVNhEqrnvMC6vnyXbC/lqSz4zhknPF9F+OlUyBYExR09eNZqiyhreWJ3Ny8v2UFjpxq81JkPRO93B7VMGcP4JPcgrd/Hmqmw8voaBItZxFiZDUeb0MLJnMmcMz+SzzXkxLRoKgRajhy8aWfd3n/R47pg2gDumDSC7uJqPN+by+GfbyC6uJjXOgjtMVWm0XRNdHh/XzV3Jjyf05d2fTKRPenxogU68FvJ/CKw5FW1CZbJCYiZc+44s1iva3W1T+rN6T3HIgz/aSWhsZhO3nT4cR7yNim/Whh0vlWiL591rnuKZlf/hV588RkFVMdMHjOPRs35F1/i0Btsqk4Hf6ZVk6hiWkWTnxRtPJq/Mxasr9vLWmmxKqt14fJpEm5lRvVO4feoAJg5Il5py0SHsyK/kinnrKes5Gq9XA6ET+9TG2DdXZ/PJxjzevG08AzPCVyo15vNr5i3bHbZiIprvMNVuH7MW7ZRkSrSrTpdM1UpPsHHn9IHcOX0gWmu8fh2yiGLtrHeNxTrOwlCqrlvh3y8/getfqOG7faU4PaFBKBy7xeCZa0Yzsmf42e+y0hzcOmUAt04ZwI78Cs7+95KQroSxLA7s14HP574zhzRfsDP+BHGpsPixwHTpzU2Zbo2HtAHw4w8C+wjRziYMSCfVYcXpdoaMc4nUMmwo6J3m4LheyXjLml8LZVCXvvzz3PsB2FG0l7vn/4nff/kkT1/wUKMtNSrCZC/i2JCZbOe+M4dw35lD2LC/jN+8u56P7p58pIslRKs6UOrkkmeXUe70oMOu2teQ168prnJzybPLWXDPZHqmRJ5+/csfDuIJUxkcy3eY9fvL2FtUFb4iWIg2INOeQHCa8NCPIsluwesPvaljW5wxkJQlxQW6BVlMBi/feDIXjuqBzWxgDbNeSa14m4lUh4VXbxrH9KEZUZ3L69dhB1/G2jVxT1FV5HUclIIp98HtS2H09WBxgDUxsF6FyQaWeDDbIWs8XDIXbl0IjrTmjylEG1FKMXfmGOIirH0STrzVzOzrTgLA5DCjG89I04SB6X24fORZbC3YFfKe9mqM+PDdBcWxyzDCz/4qxLHulpfWUOnyhlRGVW1eSO5L97LvH5ey/6nrOPjWQ7j2bwICwxwqXV5ueSm6MWTvfXeAqprDW8bCrzWfbsqL6nxCtIZO2zIVjYxEGw6rGZen4ZpKsc7Al5lkb7BosNlk8Mglx3PPaYN4ZfleXl2xF69fY1IKDbi9fkb0SOKOaQM4dWhGTFN9Vri8YXvQxdw1USmq3L7o1m7oMgjO+wec8XBgWtXK/MA6VPYUyDoZ0mVNHXF0GJqZxCs3ncz1L6ymyu2NuG6boSDBbuY/N4+nb5dALaeymLD2TcK9qyxk+x1Fe/ly53IuGHoq3ZMyyCk/yAc/fMnoHqHLAFh6JmDYJQR3NIZSsqCo6HA2Hihjd2EVvkb/b0fT9c6nNbsLq9h4oKzJHja1DpaHb/mP5TuMx6cpqGi+B4EQrUme5M0wDMWNk/ry1Jeh61NFO87CYTVx+7T+YbsLdk+O41dnDeVnpw8mp9RJudOL1WyQnmClS0JsAzZr2cxGqywO7NM6cKxYWONh2Pmx7SNEOzupTxof/nQSD8/fzNIdhWEnDrCbDTQwfUhXHjhvOL1SHQ3eT5zai7tlTQ0AACAASURBVOL9FWh3w/3irQ6+z/mBOavforymkiRbAqcNmMBvp/+kwXbKZiJpWi9Ex2MoFXYafiGOZXMW7wqZBTmWrndur585i3fx76tObPY8jZO1WrF+h/FG2XtAiNYgyVQEV43tzRNfhl/tPrrFfjUXndj8lyaLyWi1vr2ZSfawk0/EvDiw1Ry266MQHUG/LvG8MHMs+RUuXluxj4825FLu9KAUJMdZuHBUT64cm0V6E5Ua9kGpKKspJJnqntiVZy/6Q8TzK5PCPjS9Va5FHF0MFRhEL0RHUeP18fHGvJBEJ5audz6t+XhjHo96fc2uA5XmCN/1OZbvMIaCdJnYR7QjSaYiSE+wce243ry+KjvqCSNqxVlM3Dalf3Rd5VpJRpKd4d2T+C67tMHrsXRNtJgUl54kteai48tItPOz0wfzs9MHx7SfMhTpVw+lcN6msLP6NctikHbVUJRJJp/oiAxDRew+KsSxpKTKg2EQMnFfrMMHDCNwrMzkprc/bXg3Vu4OnXU1lu8wNrOJiQO7RH19QhwuSaai8MC5w9lbXM2yHYU4o/ziFGcxccaIbtxzWvsv0Hj7tAH8/M3vqWrhFNCGUtxwSt92LLEQxx5b/xRSLx9MyVvbok6olMUg5eJB2AfJjJYdlXTzEx2N0+PDaIWZjU1KRayUvmhUTx6evznse9F+h+mWZOPErJSI5RGitUgyFQXDUDx33RgeeH8j7323H4/P3+SK92ZDYTIUV4/rzW/PGXZE1haZMTQDu8UUkkxB5K6JZkNxQlaKTCkqRBQcx3XFlGCl+O1t+CvdgaSq8fdoFUiijHgLqZcOxj5AHvIdmaFkNj/RsSTazWG7rsY6fMDr1yRGmHQn3mbmRyf24u012XjDnDPSd5g4i4k7pg2Qdd1Eu5JkKkomQ/HXi4/jxlP6Mnfpbt7//gBmw0AHvzkpFD6/5rIxvbjhlH7063LkkhGzyWDeDWO5YvaKmLomKgVJcRaejDBAVAhxiK1fMpm/HIN7XwUVi/fj2lpCXW2LSWEflErilF5Y+ybJA74TMJSSMVOiQ0l1WLGZjZCJemKd2dhmNkh1RB7LdPeMgSzYkEuZ0xNTOQ0VqBD+37ocPvg+hy4JNmYMy+Dskd2bXYZGiMMlyVSMBnVL5JFLjueB84azZk9x3c2e6rAytm9ai9awaQvH90ph7swx3PzSGpxuX7gJ/howG4pUh5U3bxtPtyR7u5RRiI5CKYWtTxK264YDBNeh0iiZxKXTCYyZkmRKdBwmQ3HdhD48v2R3SEIVbdc7m9ngugl9MEWxUHn35Dheu3kcVzy3nOqayN9famkNFTVevtlRVPfal1sO8tv3NnLNuN7cOqV/k5MKCXE4JJlqoQSbmWlDoltI90iZOKALH9x5Cn9Z8APf7AwEl9qpTQ1bDubETVislWj8DErvwc8nXXhEW9SE6CgCk0tIK1RnUu7y8OG6HNZnl1Jc5ebh+ZsZ0i2Rc4/vTnw7TkIkRFu4bnxfnl+yO+x70cxsXHuMxvx+zdIdhSzfWURRVQ0Wk0GPFDvnHteDD+6cxHVzV1Lu9IQdtlBHawiu09lY7QLAL3yzm/+u3c+bt41nYEZixLIKEQuJ8B3coG6JzLvhZA6Wu3h5+U4+2LGACttn+E1FoLygAuFnt1fxy6Wf0SWuCzeMuIELBl6AzSQ1OEII0ZyteRU8t3gnH23IRSmFM/ilb+7S3TisJh763yYuOrEnt0zuR/+uCUe4tEK0TGaynbNGZvLpxryQdTcjsfo8TE9VZCYf6vVS7vLwxqp9zFmym+oab4NkyWwonvhyByN7JvOHC0ZgNimeW7yL7/aV1nXX035dt4+Oovu0x6cprnJz8TPL+OjuyWSlOSLuI0S0VHPdEcaMGaPXrFnTjsURbaXSXckdX9zB1pKtOL3OZreNM8fRK6EXz5/5PGn2tHYqoWgvSqm1WusxR7och0vikzjS3l6TzYMfbMTj082OkzIZYDWZ+Pvlx3POcT3asYTHno4QnzpqbHJ5fFw2aznbDlaEdPdris1sMDDFyuOLniD5hOPIfPBBsiu9XDF7BaVON64IM6E6rCamD8ngn1eMotTpJqfURbXby6yFO1mxuzhkIWGAqs0LKV/9Pp6i/RjWOCwZ/UmeeDn2XiMwFPROc/D1fdNkDKuISXOxSTr0dwIur4uZn8xkc9HmiIkUgNPrZHfZbq7+6GrK3eXtUEIhhDi2vLl6Hw9+sBGXxx9xwgmfPzC99M/fWsdH63PaqYRCtC67xcSbt41nTJ9UHFGMD3dYTYzpk8rbd09lyBv/wV9Zxarrb+OCJ5aQX+GKmEgBVLt9fLnlILe+vIb0eBujslIY3C2RlU0kUuWr3qP4yzkkj7+cXne9Ss875pE4+hyc21cCgZk28ytqWL6zKGRfIVpKuvl1Ar/75nfsKd+D2+9u8HrJkhIKPy3Ene/GZDeRdFIS3S7thinehFd7Kagu4Gdf/4y5Z849QiUXQoijz/r9pTz0v01hvww2Vyvu8vj5xdvrGJKZKOM2xDHJYTXz8k3j+HpLPrMW7WTDgTKAupYqW7Ab3nE9k7l96gCmD80ITDphNdP973/nqt9/SIXTjb/RulSR7puVu4v51xfb+MUZQ3h91b6wZfPXVFG69DXSz7kXx5CJh8o8cByOgePq/na6fcxevEsW9hWtRpKpDq6guoAv930ZkkgVflxIwccF9Lq5FwnDE/CUeMh5JYc9j++h32/7YZgN3H436wrWsat0F/1T+h+hKxBCiKPLk1/tCNvNqXzVe5St/C/pZ9yJvd9olMmMc/danNtXYu81AgCP18/sRbt47LIT2rvYQrQKk6E4bXg3ThvejT2FVSzcmk9pdWBm4xSHhWlDMugbZjKrxTsKKVY2/EbDySSiuW+cHh8vLN3NndMH8tKyPWHvv5oDW9BeN47BE5otvwaW7yqiqLJGZvcTrUKSqQ7ura1vhbzmc/rIfz+fnjf1JPH4QO2otauVrJ9kse2X2yhbVkbqlFQAvH4vr/zwCg9NeKhdyy2EEEejwsoaFm8roPFw42hrxX0aPlyfw+/OH06i3dJexRaiTfTtEs/MLv2i2nb2op0hs/JFe9/U+mh9LkVV7pDXAXzOcgxHEsqI3AXRajLILXNJMiVahYyZ6sC01ry+5fWQVqnq7dX4PX6STkpq8LrJbiLx+EQqN1XWvebTPubvnE+Nr6ZdyiyEEEezt1Znh3092lpxCCzy/r/vZeyU6Dxyy5x8t6805PVY7psqt4/Zi3diNDFxhCkuCX91OdrfzDTqQUpBZY03csGFiIIkUx1YlaeKKk9VyOu+Sh/mBHNwLZyGzMlmvJUNA4yhDPKr89usnEIIcazYcKAsbBejWGrFnR4fm3Nlch/ReewqqKqb1ry+WO4bgL05xfh94ZMlW8+hKLOF6m3LIx5H68B6oUK0BkmmOrAqTxVmIzRYmBJMeCu9aF/oDFTeMi/mhIb7GMoIm5QJIURnU+EKX5sdS604UDfGRIjOoLLGS7hVdWO9bzxmK5kp4deIMmzxpEy6huLPZ1G9bTl+jwvt8+LcuYaSr19osK3b5ycrVdaaEq1DkqkOLM4Sh0+HBijHQAfKrChf27Bm1OfyUbG+gvjhDQeO+rUfh1mCjhBCxNvC16DHUisOkBQn46VE5+GwmiBM77xY7xur2eCmSf2Is4S/D5NOvpjUU2+ibPmb7H/yGvY/O5OKb+cTN+hQN0JDwalDMkh2yD0oWoe0cXZgCZYELIYFj79hDajJYSLjogxyXs3BsBsNZvOzpFlImZjSYHuv30uXOJlCVAghBndL5Kst+XgatezXrxVXhgl7vxNRhhnXnu9x7VtP6vQb67a1mw0Gdg2d7UyIjqp3mgOPL7R7bCz3DUBmkp3LTsrisU+3NnmuhBHTSRgxvcn37RYTt0yRGYpF65FkqgMzlMElgy7hja1vhCRUXc/piineRN6bebjz3RhxBkmjk8i6LQvDYjQ4xqm9T8VhkZYpIYS4YmwWzy3eRbg+S0knX4wRn0rZ8jcpnP84yhqHrdtAkiZc0WA7DfxodK/2KbAQR4E+6fEMykisW5eqvmjvG4fVxE2T+pHssHDJ6F68++1+XGHGLzbHbCj6psczundK5I2FiJIkUx3cVcOu4q1todOjA6RNTSNtalqz+9tMNmaOmNkGJRNCiGNPr1QHJ/VJZdnOorDvR6oVNxScOjSDtHhrWxVRiKPS7VMH8Kv/rguZHh0i3zcAfq3rKiF+f8EIfsgtZ3NuedgJYcIxGYoUh4UXbxyLamJGQCFaQsZMdXBZiVmclHESFiP2vsEmZaJ3Ym9GdBnRBiUTQohj012nDmxyzEYkVrPB7VMHtHKJhDj6nTGiG3HWlt03drPBpaN71c3AZzUbvHbLOMb0SQ2Mx4q0v8WgZ0oc/7trEhmJ9haVQYimSDLVCTw69VHS49IxqeiDmIFBkjWJp2c83YYlE0KIY8/EAV24bWr/mBOqOIuJX581lBOypIuR6HwsJoNXbx4XVfJTn9WkGNQtkQfPH97gdYfVzMs3jeOvFx/HsMxE7BYDo1GDU7zVRGaSnV+dOZSP75lMj5S4w70MIUJIN79OINmWzGvnvMaNn9xIXnVexAV4rYaVFHsK886cR7f4bu1USiGEOHbcM2MQaJi9eBdOT+RpneMsBr84YzA3nNKvHUonxNFpaGYSr98ynuvmrqTa7cPrDzNfej1xFoPhPZJ58Yax2MyhSZjJUFw4qicXjurJlrxyFqzP5WB5DR6/ny4JNqYO7srEAenSrU+0KUmmOokMRwZvnv8mr2x+hdd+eA23z021t7rBNg6zA0MZXDn0Sq4ffj0pdqk9FUKIcJRS3Hv6YE7qm8qTX+1gXXYpfq0bzPJnNRmg4OS+adw9YxAn92t+jKoQncEJWSl8cu8Unlm4g3fWHkApqK43jkoBcVYTKQ4Lt07uzzXj+2AxRe5INTQziaGZSW1YciHCU1o3XSswZswYvWbNmnYsjmgPPr+PJQeW8PHujyl0FuLXftLj0jmtz2nMyJqBxSRrL3RkSqm1WusxR7och0vikzia7Cuq5rWVe9l2sIIKl5ekOAvDuidyzbg+0rUoBh0hPklsil6128uH63L4eksBxdVuLCaDHsl2LjmpF+P6pUmLkjhqNBebpGWqEzIZJqZlTWNa1rQjXRQhhOgQeqc7+L9zhh3pYghxTHFYzVwxtjdXjO19pIsiRIvJBBRCCCGEEEII0QKSTAkhhBBCCCFEC0gyJYQQQgghhBAtIMmUEEIIIYQQQrSAJFNCCCGEEEII0QKSTAkhhBBCCCFEC0gyJYQQQgghhBAtIMmUEEIIIYQQQrSAJFNCCCGEEEII0QKSTAkhhBBCCCFECyitddNvKlUA7G2/4ggh2kEfrXXXI12IwyXxSYgO6ZiPTxKbhOiQmoxNzSZTQgghhBBCCCHCk25+QgghhBBCCNECkkwJIYQQQgghRAtIMiWEEEIIIYQQLSDJlBBCCCGEEEK0gCRTQgghhBBCCNECkkwJIYQQQgghRAtIMiWEEEIIIYQQLSDJlBBCCCGEEEK0gCRTQgghhBBCCNECkkwJIYQQQgghRAtIMiWEEEIIIYQQLSDJlBBCCCGEEEK0gCRTQgghhBBCCNECkkwd45RS05RS+9t73xaeb6FS6ub2Op8Q4siR2CSEOBpJbBKtTZKpRpRSlfV+/EopZ72/r2nD885USi1tq+MfLqVUX6WUVkqZG73+olLqT0eqXEJ0FhKbwpPYJMSRJbEpPIlNnYc58iadi9Y6ofZ3pdQe4Gat9ReNt1NKmbXW3vYsmxCi85LYJIQ4GklsEp2dtExFqbZpVyn1a6VUHjAvXK1IsBZiYPB3m1LqcaXUPqXUQaXULKVUXAvOfYNS6gelVIVSapdS6rYw29yvlCpUSu2pXxPUWmWIspwzlVJLg+crUUrtVkqd3cS23ZVS65VSvwz+vVAp9bBS6pvgdX6mlOpSb/sLlFKblFKlwW2HBV+/QSn1Yb3ttiul3q73d7ZSalTwd62Uuj24TalS6mmllGqLz0KI9iKxKapySmwSop1JbIqqnBKbOgBJpmKTCaQBfYBbo9j+EWAwMAoYCPQEfteC8+YD5wFJwA3AP5VSoxuVq0vw+NcDzymlhsRaBqXUM0qpZ1pQvvrGAVuD5XkUmNv4xlNK9QMWAU9prR+r99bVBK4vA7AC9wW3Hwy8DtwLdAUWAB8qpazB40xWShlKqR7B/SYE9+sPJADr653jPGAscDxwOXDmYV6vEEcDiU2RSWwSov1JbIpMYtMxTpKp2PiBh7TWNVprZ3MbBm+EW4Gfaa2LtdYVwF+AK2M9qdb6I631Th2wCPgMmNxosweD5VoEfARcHmsZtNY/0Vr/JNbyNbJXaz1Ha+0DXgK6A93qvT8c+JrA5/hco33naa23BT/btwgEMoArgI+01p9rrT3A40AcMFFrvQuoCG47BfgUyFFKDQWmAku01v5653hEa12qtd4XLMcohDj2SWyKTGKTEO1PYlNkEpuOcTJmKjYFWmtXlNt2BRzA2noVDAowxXrSYJPvQwRqSozgcTfU26REa11V7++9QI/WLANQ28/ZUu/32r899f7Oq/1Fa10dPG9CvfevAXYA/w1zjrx6v1fX268HgWuqPa5fKZVNoLYIArUs0wjUIC0CSgkEhAnBv6M5hxDHMolNEpuEOBpJbJLY1OFJy1RsdKO/qwjcdAAopTLrvVcIOIERWuuU4E9y/YGa0VBK2YB3CNQqdNNapxBorq3fBJyqlIqv93dvIKe1yhCUS+Dm79vo9X7Uu2Gj8Ptguf6jlIo2OOUQ6CIA1NVeZQEHgi/VBoXJwd8XEQgKUwkNCkJ0RBKbJDYJcTSS2CSxqcOTZOrwrANGKKVGKaXsBP6HBwK1AMAcAv10MwCUUj2VUs31NVVKKXv9HwJ9WW1AAeAN1racEWbfPyilrEqpyQT6t77dwjKEFWx+fgf4s1IqXSllUUpdRaD5+eMYDuUBLgPigZeVUtH8P/gWcK5SaoZSygL8AqgBlgXfXwRMB+K01vuBJcBZQDrwXQxlE6KjkNgksUmIo5HEJolNHY4kU4dBa70N+CPwBbAdaLzewa8JNM2uUEqVB7cbQtMmEqgRafxzN4Ebo4TAYMP/NdovL/heDvAacLvWekusZVCBGWtmNVO+nwDFBAYm5gN3AedqrQ82s08IrbUbuJhAn+AXIgUGrfVW4FrgSQK1M+cD5wePU/vvUEkgGKC1Lgd2Ad8Eg5kQnYrEJolNQhyNJDZJbOqIlNaNW2CFEEIIIYQQQkQiLVNCCCGEEEII0QKSTAkhhBBCCCFEC0gyJYQQQgghhBAtIMmUEEIIIYQQQrSAJFOHSSn1olLqT8HfJyultrbTebVSamArH7PuWtpz3/ailLpfKfX8kS6HEO1F4tPh79teJD6JzkRi0+Hv214kNkXWKZIppdQepZRTKVWplDoY/J+31Vdw1lov0Vo3N4VnbXlmKqUaTwfaapRSC5VSN7fV8Q9XW19/8BzTlFL767+mtf6L1vqo/VxE5yTx6egi8UmIAIlNRxeJTUevTpFMBZ0fXMF6NDAGeKDxBkopc7uXSgghJD4JIY5OEpuEiKAzJVMAaK0PEFh5eiTUNfneqZTaTmABOZRS5ymlvldKlSqllimljq/dXyl1olLqW6VUhVLqTcBe770GGb1SKksp9a5SqkApVaSUekopNQyYBUwI1vaUBre1KaUeV0rtC9YAzVJKxdU71i+VUrlKqRyl1I0tvX6l1NtKqTylVJlSarFSakSjTboopT4PXt8ipVSfevsODb5XrJTaqpS6vKXlaFSmPUqp+5RS64PlelMFVjFHKZWqlJof/AxLgr/3qrdvmlJqXvBzKVFKva+Uiifwb9wj+BlXKqV6KKV+r5R6Nbjfx0qpuxqVY51S6uK2vFYhmiPxSeJTcD+JT+KoIrFJYlNwP4lNYXS6ZEoplQWcA3xX7+WLgHHAcKXUicALwG1AOjAb+F/whrUC7wOvAGnA28AlTZzHBMwH9gJ9gZ7AG1rrH4DbgeVa6wStdUpwl0eAwcAoYGBw+98Fj3UWcB9wOjAIOO0wPoKPg8fIAL4lsPJ3fdcADwNdgO9r3w/eZJ8D/wnueyXwjFJqeBPXX6qUmhRDuS4HzgL6AccDM4OvG8A8oA/Qm8DK5k/V2+8VwAGMCJbrn1rrKuBsICf4GSdorXMane914Kp65R0ePMdHsV6rEK1F4pPEpyCJT+KoIrFJYlOQxKZwtNYd/gfYA1QCpQRu0GeAuOB7Gji13rbPAg832n8rMBWYAuQAqt57y4A/BX+fBuwP/j4BKADMYcozE1ha728FVAED6r02Adgd/P0F4JF67w0OlntgE9e7ELg5is8lJXic5ODfLxIIWrXvJwA+IAu4AljSaP/ZwEP19v1TlP8eja9/D3Btvb8fBWY1se8ooCT4e3fAD6SG2a7u36Lea78HXg3+nhj8zPsE//4z8ELw92avVX7kpzV/JD41+blIfJL4JD9H8EdiU5Ofi8QmiU0NfjpTP9eLtNZfNPFedr3f+wDXK6V+Wu81K9CDwM1zQAf/Dwna28Qxs4C9WmtvFGXrSqCGYK1SqvY1BZiCv/cA1kZxzmYFa3z+DFwWPKc/+FYXoCz4e91nobWuVEoVB8/fBxhX27QeZCZQu9Ea8ur9Xh08J0opB/BPAjUvqcH3E4PXkgUUa61LYj2Z1rpCKfURgZqTvxGoabkl+HZbX6sQjUl8kvhUR+KTOIpIbJLYVEdiU3idKZlqTv0bPBv4s9b6z403UkpNBXoqpVS9oNAb2BnmmNlAb6WUOUxQ0I3+LiTQBDtCB/olN5ZL4H/+Wr2bvpRmXQ1cSKCpew+QDJQQCD616s6jArP2pBGoUcoGFmmtT2/huVvqF8AQYJzWOk8pNYpANwMVLFOaUipFa13aaL/Gn3E4rwMPKaUWE+i//XXw9SN1rUKEI/HpEIlPEp/E0UNi0yESmzpxbOp0Y6aiMAe4XSk1TgXEK6XOVUolAssBL3C3UsoSHHB3chPHWUXgRn4keAy7UuqU4HsHgV7BfsRorf3B8/5TKZUBoJTqqZQ6M7j9W8BMpdTwYG3DQ1Fchzl4ztofC4Hm2RqgiEBtzl/C7HeOUmpSsGwPAyu01tkE+jAPVkpdF7x2i1JqrAoMCm1LiQSCZalSKo161661ziXQj/kZFRhsaVFKTQm+fRBIV0olN3PsBQRqUv4IvBn8d4Ajd61CRCLxSeKTxCdxNJLYJLGp08YmSaYa0VqvIdBk+RSBmocdBAf0aa3dwMXBv4sJ9A99t4nj+IDzCQyI3AfsD24P8BWwCchTShUGX/t18FwrlFLlwBcEahXQWn8M/Cu4347gfyN5lsCNVPszD3iZQDP3AWAzsCLMfv8hcNMVAycB1wbLUAGcQaBpN4dA0/LfAFu4k6vALDCToyhnJP8C4gjUQK0APmn0/nWAB9gC5AP3Bsu7hUDtyS4VGNDZo/GBtdY1BP79TiNw3bWvx3StQrQXiU8SnyQ+iaORxCaJTZ05NqmGXViFEEIIIYQQQkRDWqaEEEIIIYQQogUkmRJCCCGEEEKIFpBkSgghhBBCCCFaQJIpIYQQQgghhGgBSaaEEEIIIYQQogWaXbS3S5cuum/fvu1UFCFEe1i7dm2h1rrrkS7H4ZL4JETH0xHik8QmITqe5mJTs8lU3759WbNmTduUSghxRCil9h7pMrQGiU9CdDwdIT5JbBKi42kuNkk3PyGEEEIIIYRoAUmmhBBCCCGEEKIFJJkSQgghhBBCiBaQZEoIIYQQQgghWkCSKSGEEEIIIYRoAUmmhBBCCCGEEKIFJJkSQgghhBBCiBaQZEoIIYQQQgghWqDZRXtFdHx+zddb8tmcW05JlZsEm5meqXGcfVx3kuMsR7p4QohmbMkrZ+HWAooqawBF10QrM4Z1Y0DXhJiOU1btobCqBrfXT1KchcwkOyZDtU2hhRBCHHXKnB4WbMglp8RJZY2X1HgrI3okMW1IhjwPOjBJpg5DUWUNr67Yx4vLduP2+ql2+9DB9xxWEw/9bxPnHNedWyb3Z3iPpCNaViHEIR6fn4835jFr4U52FVTi9Wu8/sDdazEp/v7ZNoZmJnLHtAGcPjyzyYegz69ZtC2fWYt28d2+EqwmA6UUPr/GZjaYObEvV4/vTUaivT0vTwghRDvalFPGnMW7+HhjHoahcLp9ABgK4iwmrGaDG07pxzXjepOeYDvCpRWtTWmtm3xzzJgxes2aNe1YnGPHuuxSrpu7khqvnxqvv8ntTEphMSvuO2MIN0/u344lFCI8pdRarfWYI12Ow9XS+FTm9DDzhVVsPVhBdfCB1xSH1cSorBSev34MDmvDuqfv9pVw68trqXZ7qWriODZzoCf15WOy+P0FI6RmUogIOkJ8ku9Onctzi3fyj8+34fFqfM18p7aZDWxmg1dvHsfxvVLasYSiNTQXm2TMVAus31/Klc+toNzlbTaRAvBpjcvj5++fbePpr3e0UwmFEOFUu71c8uwyNuWUR0ykAtv7WLu3hCtmr6DGe2j7xdsKuHrOSgoqa5pMpIC6ypb/rs1m5rxVeHzNxwshhBDHjie/3M4/P9+Oy+NvNpGCwPOg3OXliudWsGF/WTuVULQHSaZiVFrt5rq5q3B6In8Rq8/p8fHEl9t4a002Gw+UsTWvgsLKmjYqpRAinLtf/47s4mrcMSQ1NV4/2w9W8Jt3NgCwOaec219dG1MMcHr8rNlTwq/fWR9zmYUQQhx9vt6SzzMLd8b+fdDt49q5Kymr9rRRyUR7kzFTMXprTXaDGur6qjYvpHz1+3iK9mNY47Bk9Cd54uXYe40AoMar+fU764kPdhdye/0MyUzk9qkDOGNENywmyW2FaCt7i6pYsr0wbGtyj7uIVgAAIABJREFUpHvX5fXz0YZc/u+cofzug41Ntmo1dxynx8eCDbncPEnGUAohxLHusc+2NplIRXqmuL1+3l6bLcM/OghJpmLg92vmLNmNyxP6Zax81XuUrfwv6Wfcib3faJTJjHP3WpzbV9bdPABaQ2WNt+7vDQfK+NV/13H/ewZPXz2aSYO6tMu1CNHZvLhsD/4w3TCivXcV8OzCnWw4EL57RjTH8fg0c5fu4u+Xj2qTaxRCCNH2tuZVsKugMux70TwLnB4fzy3exY2n9MOQsbTHPGkKicGynUVU10uEavlrqihd+hppp9+BY8hEDKsdZTLjGDiO1Ok3RjxuldtHmdPDzS+vZv76nLYouhCdmtvr583V2Xh8DZOpWO7dGq+f/6zchy9MF8Foj+Pza+avz6XcJd07hBDiWDXvm90hzxOI7ZlSVeNlxa6iZs+zt6iKZTsL+WrLQVbvKaa4yt2q1yFah7RMxWDrwYqwA8hrDmxBe904Bk84rOO7PH7ue3sdGYl2Tu6XdljHEkIcUlBZQ7ixwbHeu01NOBPLcSwmgzV7ijl1aLeozimEEKLt7SuqZn9JNdVuHwl2M/27xje5rMX6/aX4/KEPlVieBR6/ZuvBCiYObNgjye3188mmQ0t3WMxGg/emDu7KrVP6c1KfVJSSVq2jgSRTMahweXCHqYnwOcsxHEkowxTxGBHHZnj83P/uer74xbTWLr4QnValy0u4IYmx3LtAoJ9umIdXLMfxa02pDDwWQogjzuPz89mmg8xatIPt+ZUNxq67vX4m9E/n1qn9mdA/vUHiUlkTfqxULM8Cj9dPpathb6d12aXMnLcKt89PVfAcrkaVeJ//cJClOwoZmJHASzecTGq8NerrFW1DkqkYOKwmzIaqW9yzlikuCX91Odrva/YGinZsxoFSF+v3l8o6BEK0EofVRJhKxKjv3TpN1ALGchytNYbUJgohxBG1KaeMH89dhcvrO5S4NBoTv3BbAav2FJOV6uCVm0+ua6mKs4SP87E8C8wmRZz10DYrdhVxw7zVEWcH1DqwbMcPueWc/cQSPrxrEl0TZSHgI0nGTMWgV6oDmyX0I7P1HIoyW6jetrzJfWMbm+FjzuJdrV5+ITqrLgm2kEoQiO7ejUYsx3F6/Pzjs2188P0B3BHWqRNCCNH61u4t4bJZyymqctclUk2pdvvYWVDJOf9eQm6ZE4C+XRxht43lWWA1G/RKjQNgV0ElN70YOZGqz+PTFFbUcPWcFU3OMi3ahyRTMTh1aAaEqd02bPGkTLqG4s9nUb1tOX6PC+3z4ty5hpKvXwBi60fr17BoW0FrF1+ITivOauL0Yd1oPGlSNPduLbOhGN49EasptFUpluMA7Cup5v53N3D+k0tlvTkhhGhH2cXVXP/CqqgWbq/l9WtKqj1cMXsF1W4v147vg80c+hU6lmeB0+1jT2EVhZU1/P2zbc1Os5770r3s+8el7H/qOg6+9RCu/ZvqynWg1MnHG/Ji+AREa5NufjGwW0xcdXJvXlq+J2QWl6STL8aIT6Vs+ZsUzn8cZY3D1m0gSROuAGIfm1EVw00uhIjslin9+WpLfsgDK9K9W8tsKB44bzg3zFtNuFqVaI9Tq8rtY1dBJec/uZQFd0+Wfu9CCNEO/vXFNqrdoTMzQ/Pj2n1+TX6Fi1++vY7vs0vDLrUB0T0LrCbFGSMy2ZZfybTHvqba7QvbFT2a4SHVbh/PLtrJRSf2PPwPR7SIJFMxun5iX15duTfslJgJI6aTMGJ62P1iHZshIyqEaF0n9Eqmd5qDHfmV+Bo9BOvfu4YtB3PytxiWNShjOdobj3YOZHjyFCYO6MLYvml8s7Mw7OyAzcWAcDx+TWFlDdfPW8UHd54iMzMJIUQbqnB5+GhDbosTF5fHz1dbCnjt5pPZlFPOXxZsCduiFOlZYBiK35w9lF6pDv79xTae+mpHSHJWOzwk/Zx7cQyZWPe6Y+A4HAPHNdh2b1EVm3LKGNEjOZaPQ7QSSaZilJXm4NFLjudX76wPu3hvU+r3o40fOini9gk2+acRojUppXjhhrGc8+8llDs9jdqWNOakdVjTF2JYi0B5UareFkmbybbM5y8rL+LuM6/k++dLGyy+HUlztZ0en2ZHfiVr95Ywpq8siSCEEG3lnbX7w1ZaxZK4KBUYjnHt+D4s3VHIom0FMX0ftFsMHrv0BHqlBsZdLdleiOcwp1n3+2HFrmJJpo4QGTPVAheM6smfLhyJPcxkFE2JpR+tyYCzRma2drGF6PR6psTxzh0TSU+wYqkdQKW82Hu+ir37u5jseSjD0zCRAjBqcPmcvL31be5Zcg2/uzSeeKspqhbk8lXvUfzlHJLHX06vu16l5x3zSBx9Ds7tK+u2cXp8PCeTzgghRJv6aks+zjDDKGJJXFweHyt3F6OU4smrRjN9aEaTs/s1ZrcY/Pmi4zj/hB51r5U5wy+VEcvwELfPT2m1LOh7pEjzRwtdOiaLARkJ/PPzbazYXYz267A1C/VFO6bCYhjcNKlfWxZfiE5rYEYCn9w7hae+2sFba/ZiZL4KcTtQRuSWJq/2Uump5PH1P+cf1z3Do/+rJLfM1eRA5mhrO7UOTMFbVFlDeoJMcSuEEG2hpIk1/mJbKxCKghMHWc0Gz1w9mrfXZPP01zvJr6zB5fE16AZuNRsoYHz/NO49bTAn9k5tcDxT45mRal+PYXiIAqzhFlMU7UKSqcPQO83Byf3TWL+/jKoaD0aw6VcRdtI/ILoxFQMzEhjULbHVyyuECOiSYOP3F4yga9YS5m7YjUc3TKRKlpRQ+Gkh7nw3JruJpJOS6HZpN0zxgQea0+vk4TX38tndn7Etz83Fz3wTtg9+LLWdNrPBzoIqSaaEEKKNWMLMxgqxj2u31pvJTynF5WN7c9mYLL7LLuW1FXvZW1RNtdtHcpyF0X1SuG58XzKT7WGP1TXRxpa8ipDXYxkeYrcY8uw4giSZaqEFG3L4+VvrgNBF3ppvn2qew2risctOOIwjCCGi4fF7eO2Hl/HohlOTF35cSMHHBfS6uRcJwxPwlHjIeSWHPY/vod9v+2EEH6Iev4dP93zKRQMvCjsZBcQ4i6cODI4WQgjRNrolhU9oYklcrCZF18TQ4yilGN07ldGNWp4iuWJsFt/uKwlZ76r+8BBlmLD3OxFlmHHt+R7XvvUN1ij1+eG04RkxnVe0HmkTbIG312Tz87fW4fL4Yxp0GInDauL5H49hWPekVjumECK8r/d9jU83fHj5nD7y38+nx7U9SDw+EWVWWLtayfpJFu5CN2XLyuq2rfZW88LGF1BKYY6itjMiRdT97oUQQsTusjG9iLeGxtmY1gpUqlXHtZ8xPBOjiZlck06+mNRTb6Js+Zvsf/Ia9j87k4pv5xM36FBvBwVMGdyFjDAJnmgf0jIVo1W7i3nwg42tlkSp4Beorgk2nr5mNCN7ykwsQrSHlze/TLW3usFr1dur8Xv8JJ3UsELDZDeReHwilZsqSZ1yqNbxYPVBthRvoUuCjdwyV8g5Yqnt9Pj8dE+JO4wrEkII0ZypgzOwWUxh1/KMdlz72D6p9GzFWG01G1w3vg/PL92N2xv63TLS8BC7xcStUwa0WnlE7CSZitFfF/zQZCLV3PTHcGjtKKvZQGuNX8OMYRncOqU/o3unyhozQrSj7IrskNd8lT7MCWZUmJYmc7IZ515ng9cMZbC3fC/Xjh/Ak19ux9XoQRhLN40+afH06xLfSlcnhBCiMZOhuHFSX576ckdIvIbIiYvDauK2qa2fuNw5fSCfbMpjb1EVvhjq6uMsJi4Y1YOxfWPrWihalyRTMdhVUMnm3PKw70Wz2JsGhnVP5OUbx2E1GSTYzU3O4iKEaFsub2hLkinBhLfSi/bpkITKW+bFnNAwZPq1n2pPNVed3Jsnvtwe9jzR1HbGW03cMU1qFoUQoq3dPKk/C9bnsT2/Ao8v+lHucRaDM0dkMnlQl1YvU7zNzBu3jueyWcvJLXOFbaFqzGJSnDo0g7/86DipjD/CJJmKwbxv9uALM2VXLIu97S6sorjKzZBMma1PiCPJbraHdPNzDHSgzIryteUkn3yoy63P5aNifQXdLu3WYHtDGcRb4kmLt3La8G58vukg7jDVipFqO00mxdnHydpyQgjR1uwWE6/dPI7LZy9nX3E1NVEkLnEWE6cMTOfRS4+POnGprPFysNxFdY2PRLuZzGQ79mbGxWYk2pn/00n85t0NfL75IArCli3eagrOHqv57blD6yrlfX7NwXIXFS4vVrNBlwQriXZLVGUVh0eSqRgs31WE9zBXqVbA99klkkwJcYT1SepDsau4wWsmh4mMizLIeTUHw240mM3PkmYhZWJKg+19fh99kvoA8NeLj2PD/jIOlFbH3E1j3syTsZll8gkhhGhTPg+4ykg1zHxw5wQe/GAz73+fE3grzPc7R3Bx9lsm9+fuGYMwIvQm0lrz7b4Snlu8i6+3FGAxKZRS+LXGrzUXjerJTZP6Nbn8TaLdwtNXj6aosobXV+/j1eX7KKqqwePT2MwG/bvEc/u0AZw1MpNZC3fx2/c28rdLjue1lft4adkeXF4fZiMwlMTt8zO+fzq3TRnAxAHpEcsuWk6SqRhUusIv6hnL9Mdev6bcGXlxUCFE27p+xPVsLd4a0jrV9ZyumOJN5L2ZhzvfjRFnkDQ6iazbsjAsDSdA7ZnQkyFpQwBIslv47+0T6mo7I6zhDQQqVx44dxgn9ZH+7kII0SY8Ltj8Piz9JxRsBbMVtMah/Tw66CwKbeMZPnYG/1ufR0FlDW6vH7vFoG96PLdN7c/ZI7s326JUa39JNTfMW82BUifO4MK9jee5eHttNu99d4AxfVJ59rqTSGqi5Sg9wcZd0wdx1/RBAPj9OiQZunVKPyY+spuJj3yFyVD1WrEO1eYt2V7It3tLSIqz8Pz1YxjRQyY5awuSTMWgNRZ7M5TCZpEZ6YU40qb2morFCP8gS5uaRtrUtGb3t7vhRxsduE/KxpqVBQRqFeOtZqKt/9PAnxf8wODMRMb2bf58QgghYrRyNnz5x8Dv7srAf72H1hZUWxcwW32BfccL/Pr6l6D78WitYx6DtCO/kkueXUaly4uvqYUHCawH5fP7Wb2nhPOeWMr7d55CWrw14vEbJ1Ien5+bXlpDtduH16/D9pqqVeX2UeX2cdms5cybOZZx/dOjvzARFflWH4OmVq+uP/1xJGaTIiNRVqkW4kgzG2ZuHHkjcaaWTXEbl5jKaanj2HPpZeT96c94Cgu5/dW17CioJIYxzVS7fcyct4pdBZUtKocQQohGtIaPfw1f/D6QRLnDx1cDP3btguJd8MKZsHtxzIlUUWUNVz63nHKnp9lEqj63z09umZNrnl9JjTeKdQgb+flb3/PtvpKoxnvVqnb7uPHF1ezIr4j5fKJ50jIVg2vH92FzTnnI+gSxTH/s9wfWORBCHHkzR85kfeF6vjnwDS5f6Ox+TXGYHTx/5vP0TB2M98prKZw1m3evuYtVo6+hRofWUUVaNqHa7ePPC35g7vVjW+3ahBCi01r6T/j2ZfBUR962lqca/nMl3Pw5dBsR9W7PLNxJmdNDuDSqudjv8Wn2Flbx4bpcLj2pV9TnW5ddyheb88Mu0xPNs+YPH27mlZvGhewrWk6SqRicNTKT+9/bGPa9aKY/NhuKy8b0Ii7M6ttCiPZnKIPHpj7GA0sf4Ovsr3F6nc1ubzEsxJnjmHPGHAanDgbAnJ5O5m/vZ37qIpx7K2jcxy+qZRN0oG97foVLVrEXQojDUZ4Lix5p0J0PoO+/Kqj2wO57Eoi3BgL189+6eXW9h4Uzg2v8earggzvh1oVRncrl8fHGqn1hp1iPJvZXe3zMWrgjpmRqzpJdYVuzol2iZ+XuYnLLnHRPlkXiW4skUzGwmU1cM643Ly3bE7ZpNdL0x2ZDccMp/dqyiEKIGFkMC49MfoSvsr9i7oa5bCvZhs/vw6sPTRTjMDswlMHlQy7n2mHX0tXRtcEx8spcrMqphkbdQ2JZNkEBr6/cxz2nDW79ixRCiM5izQtNvuXT8O+Vbu6f3Mxwi/wfAhNVdB0S8VQLNuSGfT2W2H+g1MX6/aUc3yul8WFClFS5+XzzwZAJjmI5HxpeWb6XX501NOL5RHQkmYrRz08fzKJtBewqqIxxsTcT9505mH5d4tuwdEKIllBKMaP3DGb0nsHust3M3zmf3KpcXD4XqbZUxmaOZUbvGVhM4Ses+GZHYdgFuGNZNqHG62fBhjxJpoQQoqV8Hlj1XEirVK1fTrTy6Dc1/GSslRR7E2OjfF5Y8Syc/6+Ip/t4Q17I0A+INfb7WLytIKpkauXuIiwmI6RCP5bzuX1+PtqQ26JkKr/cxa7CKipdXhxWE71SHfROd8R8nI5GkqkY2S0mXr9lPFc9t4Lt+RVRTX9stxjcNrU/N03q3/YFFEIcln7J/fjp6J/GtE+p04M3TOVKLMsmAJS7PDGdVwghRD2568Df9PIzY3qYmNbXzOPLavjTqU10qdbewFTqjZIprTU1Xj9Otw+nx0e120d2SfgxWbHEfr+GjzfkUh5m+Z3G6d4PueW4PKHJW6zPmoomlvoJWz6/ZtnOImYv3smq3cVYzUagv6ACt9fPkMxEbp86gNOHd8Ni6pzz2kky1QJp8VbuO3Mw977xPX4d6NlT3ahmQqlAa1SXBBv/d/ZQzj6u+xEqrRCirRmK0KcesS2bACG9BIUQQsSiughU81/o/zjdxikvVHHPuKanJPc4KzjnH4uodvtweQLJk8vjw2IyiLOaiLMEfnLLwk9cFGvsj7OaSW80RXq4unqH1dwqz5po5ZW5uHbuSnJLnXUtcI1bxdbvL+OX/11HnMXEKzeNY1j3pFY7/7FCkqlwvG7YMh+2LoDKAjAMSMiE4y6B/qdS7fXzhw838/Q1oxnXL50P1+Xw4rLd5FfUUOP1E28zc3yvZG6b0p/RvVNjnmZTCHFsSYu3YjEp3I0q++ovmxA/dFLE46Q6Iq83IoQQogk68lThIzNMnDfYzCNL3QzrGj7xMit4+prRgaQpmDzZLaaQ7tx3vLqWjzfmhewfS+w3G4pTh2Vw29QBEcv+6aY8lu4opNLX8GET67MmOS58l/X69pdUc8FT31Dm9OCL0A2rqsZHVY2PS55dxms3j+PE3p1rIXpJpuqrLoZlT8Lq5wM3ZON1CX74AKzxfJNyGRN7X8i0IYEpzi8fm8XlY7OOQIGFEEeDKYO6hu3mF8uyCXEWEz86sWd7FlsIITqWuFTCt+k09IdpdkbPruQXE8JPRKFs8QzulhjxOBeO6sHi7QVU1bR8yRyzSXH6sG4RzwUwYUA6Hl9owhjL+QCcbi+zFu3kzBGZYcfyV9V4uWL2Csqqo187CwK9tH78wio+vmcyvVI7z1gqSaZqFe2EeeeAswR84Qcu4q4CdxWTKucwvctXUDUf4ru0bzmFEEed1Hgrpw3rxscbc0PGUUazbAKAX2suGyOVMkII0WLdR0XVOjUwzeCKERaeWOXmuIxGrVPKgIGnR3W604Z1w2wYQOg4pmhj/6CMRAZFkbgBJNktnHtcdz5YlxPSWhTt+Wxmg/vPGcbK3cVcPns5qQ4LZ47I5MwRmYzokYRSirfXZFNcVRM2kYq8lpWXp77awSOXHB/VNXUEkkwBlO2H52eAs5RoajTicEPxDph7Oty6COydr3+oEKKhW6b056st+TjDDA6OtGwCwORBXaLqeiGEEKIJFjuc+GP06udR/uYn9PndVBuvrA+zjdkOE++K6nRmk8GPJ/ThucW7WrRkjsNq4rapsU1OdvPk/izYmBu2612k8xkKThnQhQtG9eSCUT15+MKRfJddyqeb8rjzP9/i9WlOH57Bh+tycYZZFDiatax8fnj/+wM8eN5w4m2dI83onNNu1Kc1vHoJuMqpn0j1/VcFGY9VUOU+9Nrz37qZ9mJV4A+/B8oOwPu3t3OBhRBHo1FZKZx3QnfiLLGFVUXggbp2TwkffH8g4vabcsr49xfbuP/dDfzmnfX8/bOtrN1bgo6hK4YQQnREWms+ib8Ad5jGqT33JnJa/0Nf7rOSDVwPJB1asLdWci/ocWLU57x1Sn8yk+1hl8dojs1sMCorhbNHxjZB2fAeSVw0qidxltgnmoi3mfn9BSPq/jYMxUl9Uvm/s4fy1c+nMnfmGCpdPkqq3SH71q5llXb6HTiGTMSw2lEmM46B40K6ERpK8d53+2Mu37Gqc6SMzcleCaXZoMNMNRlpcTdfDWz/AspzIKlHGxdUCHG0++uPjqO0ysPSHYVhW6gaMylFgt3M27dPwO31c/cb37FoawF/uHAEifZDrVRen5/563N5dtFO9hVVUeP113UnVMDzS3eTkWDj9mkD+NGJPbG34CErhBDHsp0FlTzw3kYqanyMGX41tu1vgyf81OVNMsfB+f+OaZdEu4U3bh3Pxc8so7CyJqo1SO1mg8GZiTx//ZiYkzCAP//oOIqr3CzZHt2zpvYM950xuG5dqIPlLl5Zvpc3Vu+jpDowyYTNbJAcZwm77E8sa1lVu318taWAa8f3jeGqjl3SMrXsiSZvtl9OtPL4shpKXRFujNVz26BgQohjjdlkMPu6k7h+Yl/sFqPJmkOTEVh/bmTPJBbcM5nB3RIZ2TOZ+T+dhM1i4twnlvLdvhIAKlwerpyzgvvf28DWvAqcHn+DB50GnG4fe4ur+eOHm7jwqW8oqmxi3KcQQhxlPD4/CzbkcuOLqzj3iSWc9a/FXPnccuYu2UWZM/Laey6Pj398vo1Ln13G6cO78f5PTqHLpf+CATPAEsMkCOY4uOhZ6DMx5mvonhzHgrsnM65fOjazgcUUPkGyWwxsZoNzj+/O27dPCEx13gImQzHr2sCzxmZu+lljBJfp6d81nn9eMYonvtzBNzsKufHF1Ux59GvmLNlFYaW7rstgjddPfkX450esa1mFa93qqDp3y5SzNNCy1MQ4qagWd/PVBGb/m/Fg25VTCHHMMAzFb84eyl2nDuS9b/cza9EucsucWEwGPr/GZCh+dGJPbprUL2TQscNq5q8XH8cnG3O55eU1XDOuN59tPsjO/ErcUdR2Oj1+dhVUcuHT3/DR3ZNlDJYQ4qhVVePlmYU7eGX5Xnxah8yIty67jEc/3crZIzP5+elD6lpU6lu6vZAH3t/A0MxAxVT35LhDb17+Mnz+IKyeAyjwhl8TCmsCGCa47CUY0PzY1uakxlt59eZx7Cuq5sVlu3lzdTY1Xj9mQ+Hxa5LjLNx4Sl+uPLk3XRKa6PEUg9pnzU+mD+CdNfuZvWQXhRU1WEwGfq3xa80ZIzK5dXJ/TshKAaCyxsN1c1diKIU3wnTnjcW6llVnWsC3cydT5TlgtjY9ex/RLe5GTXlgbSqzrBEjhAhIsJm5bkJfrpvQlxqvjwqXF7vFRLzVFHHtubNGdueE/2fvvMOrKNP/fc/MaUlOeiEhARJCkyZNOoIVRVRUFHsFFbu76+73p9t0V9fddXfVVVGRVRGxCyoW7PRepUMggRBCQnpOPzPz++MAJjmTnHMwkIS893VxeTEzZ+Z98ZyZed7neT6fTglMfH4p5U4vDduhmlJT8mk6JdVu7nprLe/eGbocQyAQCE41JTVurnt1JYUVLkPhBuB4+dqnm4r4dnsJb95+FoO7JAFQWuPhyc+3sSa/gicu78N5RtLisgzjn4RRD8G6N2DVS+Bzg2wCdPB7IK03jH4Iek0EpXkWnzonR/PHS/vwh4m9cXhVnF4/sVYzNrN8UnxH42xmbhudw62jsqn1+Kl2+4+X69UNaCqdXl78IQ9dB/8J9NhG6mWVHtdIEuI0pH0HU14HhlbSdQjH3A3ZFPCkMiU1/xgFAkGbx2pSsNoj62NSZIkajz8okApHTcmr6mzcX8nuwzVhS+4KBALBqaDa7WPyjBUcrHRiYJkUhKZDrcfPTbNW88FdI9hYWMm/v97F5CFZfPOrs0OXytlTYewjMOZXUHUgUJWkWALWNva05pmUAZIkYbeasJ8iRTtJkoi1mev129bl6S93cKTWY1iLFUruHCLzsoqxKEwenNXcU2y1tO9gymoPy48glLkbqg+s4oVFIBA0H2+v3B+01HNMTSl5wkNE9/y5rj+62zCiuw2rd6xP05i1dF+78voQCAStn//7cDPF1W7DQKqpl3qnV+XyF5fRLzOOt6cNo1d6hLY0sgKJ2ZDYLNNoUzg8fj7ZWGQojhHOAt0xwvWyiraaGN2t/fiwtu9gKj4LNH/Iw5o0d4PA6kYzpYcFAoFA13XeXJEfVP4SiZrSMa+PP1/WR6j7CQSCVkFpjYdvd5TgNSjtC+elXpEl7jw7N/JAqp0zf8NBjCoMI1mgO0YoLyubWWbamBzkE1ApbKu0n+4wI6yxcMalIIV+0fjjWGs9z6njmKJg2PSTMDiBQNBecXhVat3BCz2RqinJksTh6kaargUCgeAUM3dVcMYdwvcw8vg1Xlm899QN+DTh3TUHcHqDJdQjWaALB5Ms0SMtlltGZjfL+doK7TszBTDiPtjxeZA8ev5D9cv2jpm7BaFrMPiWkzlCgUDQznB4/JiUYLWlSNWUZEmixiAoEwgEgpZgzqoCQ8GJSF7qtx+q5mCli8yEqJDHCgKUOZpH7rwpzIqEpsPTk/thNbWvaojWHUzpOjjLwV0JJitEJ4M58h/PzuIavvzpEMXVbvyaTlqsldHdUxjRNRmp44CAmsuhTaCF9jOohykK+kwKlPkJBAJBMxFlUY77ftQlUjUlXdeJOUXNzwKBQNAUuq436oEXyUu9xSRTJIKpiDB6nkDkC3RGRJkVNF3n8gGZdE2N4cF3NvLh9JFB1hx7S2s5VOXG6VWJs5nolmYnuRkk4lsDrfMp6yyHDW/BiheOqq6YA4GV6oPcc2HUA9BlFIYFoEdRNZ0vfjrEjEXjTEipAAAgAElEQVR57C2txefXqNt398byfOKizNw1pivXTJ6L6dUxyK4jmAntJA2AYoWUHjDx2V84WYFAIKhPrNWExSTjU+vfjyJRUwKOLx4JBAJBS+PX9EZcPSN/qTcqWRM0TpzNzOHq4EA2kgU6iyKTFGOh0uXF49Mwm2Q6xFq5fVQOVw3JIu6oimBxlZvpc9bxxm1D0fSf38UPlDsDUu06IAVKNsf2SOXOs7sypEviSZGNP1W0rmBK0342WJNk8LkC29U6Lsq7v4aCpYEs1bXvQHrfoNO4vCrTZq9l/f6KRn9wTq+K06vy94U7eHmxmTjvE3wa9zRmZzH4XU2P0xwNGQPghvfB3H509AUCwalBkiSuPaszs1fkB6kvhaumJEtwQe8OIjMlEAhaBWZFRpEkQ4+jSLPusTZxX4uEc3ulkV/mCHqeRLJAJ0nwwd0j6JQUja7rjQY/f5jYm7veWse02WvYsL8SVdNxHH0Xd/vql3h+u/0wy/YcoVuanTdvG0piTNv0a20930ZNhfdugL2LAkZqjaIH/KG8Dph1Idz4IXT5WYHE69e44bWVbC2qbtQIri5un4bb58FlS6T21u+J2v0BLH8eXBVHfaiOIsmBUsPEbBj1MPS9CpTW888nEAhOL24dmc2clQVgsJYbSk0JAi8uU8d0PUmjEwgEgsjpkhxNXqkjaHskL/Vev0bXlJhTOew2z00juvDG8nyMnifhLtAN6JRAp6RogCazSIosce1ZWUx7a12QT2JDdD2Q3Nh+qJoJzy/hs/tHk9IGS/9aTzSw4OFAINVACKJJfA54+2qY9gOk9gDgic+2su2QcSAVyr/glre28vkDdyINvRMKlsHub6C2JOCiHZsRUP7LOLO5ZiwQCASN0ikpmqE5SazcW2boDdIUshQodX539X5yUmKCatcN0XUo2Q41hwLVALb4QD9pVMIJzkAgEAjqc9fYXP786VbDqqFwXuplKZBlSYhumxmMliIrMZrBXRJZnldmuD/UAl20ReHusblhXWv34RoeeHdjyECqLj5V50iNhxtmruKz+0djMbUtsfHWEUwVb4Gf3v+5rA/IfrYGpw/2PWgnxhKIgF9b72XOZh8/3lpnRcLrgC9+A7d8SpXLxwfrCg0DqVD+BX5NJ7/Mwfr9lQzukgjZowN/BAKBoIV44bpBTHh+McXVnkYbiBsiSxAfZeb9u0fw+rJ8xv9nMY9f3ofxfdKNP+CphZ8+gGXPHl08qvNYUL2BRaQR90HHAc0wI4FA0J65tH9H/vTJ1kb3h/YwUph2tsi4nwiPTjiDq19egcsXWb+ZRZHp2SGWs3ukhnX801/twNVIi01TSQ2fpnOgwslXW4u57MyOEY2xpWkdod/Kl8DvDdqs6vDcquDt9dHhwEqoPMBH6wqRDVKP4foXuHwqM4V/gUAgaCXER5v5+J5RdE6KxhrGSp1FkUi1W5l3zyi6p8Xy1BX9eO7aAfz9yx3c8/Y6SmoaeE7t+hqe6QELH4OK/EBlgKf65z9+N2z5CF6/CGZPCgReAoFAcIL4NY2sxChDr6lQmBWJ7ml2BnYS2fIToW9mPC/eMBCbOfxXfwmIsZp4846hKA1MeCudXvYdcbDviINKZ+BdvaTGzdLdRwyFRqpXz6P8u5nED7+GrPvmkDn9dWIHTcC1e9XxY5xelZd/zDuR6bUoLZ+ZclcHHtZ6cBT7yEgL/1jm4Z6zLCTYmvjp6Tqsnsms9WMNI+5w/Qt0HX7YWUKl0ytSyAKBoFXQIc7GgvtH879l+3h9WT4en3q8mfcYMRYFWZa4aXgXpo3pWq+Jd1jXZL54cAzPf7ebi59dwu8u6sXVQ7KQfvoQPr0/tOCOrgWqBvavgNfOg6nfgdXerHPMq8xj7o657C7fTa2vlhhzDN0Tu3Ndr+vonti9Wa8lEAhahs2Fldz/zgZG5iaTlRjFyr1luHyhe9shYAabbLfyxm1D27TqW0tzbq8OvHnbUKbNXltPGKIhsgRWk8IZGbEcrHSxeFcpE/t3xO1T+XLLIWb8mMe+Iw4sSiAw86oa2ckxdE2JQTeo7zuW1Eie8BDRPX/WOYjuNozobsPqHbv3SC07iqvplW7g7dpKaflgqmA5yGbAHbRrSEeFcdkmnlnu4a/nNqGap3rRt86jqGqY8e4I/AvMisyBcpcIpgQCQashxmri/nO7c8+4bvy4s4RPNhZRWutB13SS7RYu7pfBhb3TG60zt5kVfntRLyb278jvPtrMtlVf88fKR5H9wffdRvG7oXwfzJ0Cty5o0poiXBYdWMRLG18iryoPv+ZHrbOo9tORn/g071Ny4nO458x7OKdz04IbAoHg5FNa42HboWqqXT6sJpkOcTb6Z8U3GeBoms7/lu1jxo95PHF5Xy7pn4FP1fjth5v5aktxyLKzaItCZkIUc6cNb7Nqb62JYV2TWf3Y+ceDov3HJMuP4vVrXNQnnaljutIvK56tRVXcNGs1e0pqeG1JPrr+cxBW175jd0kteaW1GFWkR2LKjA6r95WLYCoiXOWGWaljPHGOlVH/c/DgsKZ/QLVlZciahioFv0xE4l8gSVDjidC8VyAQCE4Biixx3hkdOO+MDif0+d4d45h3z0gqnv1NUCAVVp+q6oFDGyB/KeSMOeF56LrOc+uf4+3tb+NWjQM6VVdRVZUd5Tv47eLfcl2v63h48MNiVVogOMXous6a/ApeWZzHkt1HsCryMasgVF0nzmZm2pgcJg/pFCR2U1br4dcfbKLS6WP+vaOOq8GZFZl/X3Mmkwdn8cqiPFbtKwc43vNuViRMskxWYhR3j83lkv4Z2MwnZiorCMZmVrhiYBZXDMxi3xEHxVVu3H6VOJuZbmn2ev8f+3SM54LeaTz77Z6Q522stTeSpIbHr1HlbFvv4S0fTIWgb5rCxB4mnl7q5YzUxus8o60ymt94f0T+BTrYhS+LQCA4TTGV7SLVtc9w37E+1UfHNCFN63UG7CN+QTD14sYXeXtH44FUQ9yqm3d2vIMiKTw4+METvq5AIIiMKpeP215fzY7iGlxeFZ1A5qIuTq/KM1/v4p9f7+T5awdy4VGxm+V7jvCr9zdxxaBMfnVBj3rZDwjIa4/qlsKobikcqnLxzbbDHKnx4FN1EmPMjOiaQr+s+FM11XZLTkoMOU1IzX+68SCfbCz6RdeIJKkhy2AWan4REpUEUtP/sI+PszHolVp+PaLxB7wSm0iKZKW0NtijKhL/Ao+q0TEh6sTnIxAIBK2ZlS+B6jfcFV6fqh6wsag+BHEZEV9+TfEa3tz6ZlAgVbGkgiMLj+At8aLYFOIGx9FhcgeUmMDzwa26mbN9DsM6DmN4xvCIrysQCCKjyunj0heWUlzlwhvCnuFYqd4D727gicv6cKDCxXtrDvCva85kTPfQKnAZ8VHcPCK7OYYtaEY8fpVH520JMtuFppX5GhJJUsNqUkhqY+WcLR9MdRkJWtPpvG5JMlP6mHl+tZd+aQbRqmKB3ldws6kLL3y/x1AaPVxTsmE5SW3SMEwgEAjCIu8H0I2DqbD7VBULHFgFfSZFfPmZm2cGBVJHvjxC6ZelZE3Nwt7bjq/CR9FbReQ/k0/OYznIR1cp3aqbmZtnimBKIDjJaJrOLa+vCiuQqovbp/G7j3+iX2Y8nz8whtRY8T7VlvlqS7GhoEQou6GGRJLU8Gs6559gKXtL0fLBlC0O+lwJm99rsnfqj2OtvLW5kaBLkmDoNK4zpfHC943XdIbyL4iJwJRMIBAI2iSe6iZ3h9Wnqqvgroz40sWOYtYdXldvm+pSKZlfQuYdmcT2jwXAkmqh0z2d2PXILqqWV5F4duLx4zeWbqSotoiO9rblQyIQnGrUKg/+cjeaR0W2KZiSolDiwlvxX7y7lN2Haw0DqVAZCV0PBGMikGr7vLJob5DiXyTKfHUJ15T5/DPSRGbqhBhxL2ydV0+iN/+h2HqHdIqXcf/eSNlDgk7DIaETKcCFfTrw9dbDhtmpppAlSLZbGZmbfAITEAgEgjaC3PRtP7w+VSmQnYqQT/Z8ErTNuduJ5tOIG1z//q7YFGL7x1K7tbZeMKXrOvN3z+eegfdEfH2B4HRH13TcuyqoWXQA74EapDq9J7pfw5odT+zYLKy5CUhy42IuRi/REH5GYk9JLbsO19CjQ2zQOQRtg1qPn12Ha4K2R6TM14BQSQ2rSWHamLZnytw6OrzS+0L/a8AcHflnLTEw4Z/H//r0lf3JSozCrISv+CQREJ146w7hXyAQCE5zokMvGD0+zsbM9V4OVjdS3iPLEBO6D6Ih+dX5eLX6RuxqrYrJbkIyuGeb4k34a+uXJPo0H/nV+RFfWyA43fEfcVH8jzWUz92Bd181+HV0t3r8D34dz55Kyt7aRvG/1uKvMBaAKaxwsn5/RdD2YxmJpAumE91zJLLFhqSYiO42rF6ZFoBP05i11FjoRtA2qHR6g0RDIDJlvkiIMitcMySLgZ0TQx/cymgdwRTAxP9AztjIAipzDNzwAaT2PL4pxmriw7tH0j0tNiyXZ7MMiTEWPpw+ki7JjauZCAQCwWnBwBvB3LTITt0+VUN0HbIjV/Nz+BxB2xS7gr/Wj25QTuSv8mOyB2fSjM4jELRnfMUODr+wAbXKg96IEesxdK+GWu7m8PMb8B0JNu3esL/ScEE6koyEqgXU/ASnH3WV+ZqLKLPCRX3T+dOlwf1WbYHWE0zJClw7F8/AW/FLFlwYl5BoOjixUW3NwHnzlwEBiwYkxlj4+J6R/ObCnqTFWomxBEfPMRaFGFnjKvc+Fj50tkhFCwSC9sHAGwPBUAj+ONaKwxt8nC6bYdDNYG5CoKIR4i3BMsfR3aKRTBLV6+r3cqlulZrNNcT0Dl7kirO2HTNHgeBko9Z6KX11cyD7FK5WhA6620/pK5vRGnj6VLt9+A0MgyLNSNR6jIVuBG2DhGgLPjW4ZaauMl+4WBupFouxKMTaTDx8QXf+fc2ZyE2UnrZmWkfP1FEKq9xM2XwBPt8gLte/Z6rpCxKoxYcJGR0TfhZr/XnVP5HN/t5kvFfBe3e6SYsLfqjbzApTx3Tl9lE5LM8rY96GQg5Xe1A1naQYMxf0Tmd890QKL76I2APnQu/eLTBjgUAgOMVEJ0GvibBtPmg/v+yE26fq0eAL8yVM0vSIH3x9UvqwsGAhrjr9sUq0QtqkNIrmFCHb5HpqfuYkMwkjE+qdw6bY6JvSN6LrCgSnMzWLD6J5grMEqws389QPM9h1JB9Zlume3IU/nXc/AzLOCBygg+byUbvqEHHndD7+ObMsIxP8247EKwjAZFAiJmg72K0meqXHsqWo/kJXJMp8AN072LlyYCazVxRQVuvFq2pYFJme6bHcPTaXC/t0MCwnbEu0mmCqtMbDFS8up9zhRdWjmclEZqqXkEgNcZITr26mAjtujqrD+HUOlDuZ9OIyvnzwbOKjzYbnlWWJ0d1TGN09xXB/8rRplL7wIp1eevFkTU0gEAhaFxc9DflLoLaE8JeyAXM0tYMf4O1dMu/sXsHTV/UnN9Ue9scndp3IM2ufCdqeOiEVJUah+L1ivCVe5CiZuEFxdLqrE3KDcm0dnUtzLw1/zALBaYzu13CsOhRw3K5DjcfBbR/+H09e+Csu7XUOXtXP6sJNWBsKx/h1apcWETu2E5rTgWvDRszLt4EnFeT6x0biFQSQ3MYU2QTB3D0ul999uDlIjCRcu6EYi8J953Tj8gGZTB/XDYAH3lnPub06MGlg5imbx8mm1QRT02avpcLpRa1XfiJRuG19oxKcfk2ntNbDPXPX8fbUE/MdSZhyDWWzZuHasJYofTtsfBscpaBrYEuAXpfAoFvAHnmztUAgELRK7Klw25fwv4vAVV4vQ9Uo5ig4ayopF/wfH+jw1soCJs9YztQxXbnz7K5NrixWu318vK6Q99cWoprPhOi1INUvH0kam0TS2KQmhyBLMhd2uZA4iyjzEwgAXFuPGJbt7i0/AMCk3ucDECUrjM0ZangOzeGi4JZHcG/9gajevRkyeAiYLdAg2RVJRiLKrHD90M4I2jYX9k7n/8k/Ge4LpcwHIEkSF/VNr/85m/m0KwFtFcHUtqJqdhZXB9XohiPB6VN11uZXUFDmOCEBCVl10GlSLJZ5F4LVBg0bm0u2weJ/QLcL4PzHIaXbCc9TIBAIWg3JuTB9GXxyL+xbFEhQqZ7g4yx2MNnggscD/VYErCRuGZnNeWek8ei8LSzYfIh/Tu5P38z6PVFltR6e+mI7CzYfQpYkXD4V2TKa6JwNSFJk9hUAFtnCHf3uOJHZCgSnJe4dFeje4N9S16ROyJLMw58/yWW9zmNgZh8SbI31hpuIGXcVnWc9hWwJZJOuW7CV2SsK8DXIeIWbkdDRuWLQ6ZN5aK9YTDL/uKo/D7+/Ebcvsnu2zSzztyv7YTUdLQlV/bD3R8ZWriC+xgtKF0jpCdmjA36xbZhWEUzNWro3yBguElMwTdd5fVk+f74sQhWQigJ4/WKsjhIkWQ0OpAD8R6VDd34Be3+E69+H7FGRXUcgEAhaI/a0gCJqdRGs+R9smA2uCtB8YIqCjAEw+iHodn5AJKgBWYnRvHnbWXy8/iC3vr6ayYM78dD53bGZFfKPOLjmlRWUO7z1Fso0bxrug9diy3wXSW7EiN0Aq2Llb2P+Rm6CMFYXtB9q3D5Kajy4vCqxNhPp8bafX04B1WH8G4q1xvDxDS/w0qq5/Parf1LqKOec3GH846LfkhoTnAGWoxKOB1IAt47MYc7K/RiVAYfKSJgViUv7dyTWZtx+IWhbXNwvg8PVbp7+akfYAZUsQfdUO5f0ywiUk699HVa9DKqPc/0eZE2F/eaA76E1DkbeDwOuh6iE0CdvhUh6E6pOQ4YM0deuXXtSB+D2qQx4/GvcDUx2XXvXUfLh43T+zbywGh2jLQo//Xk8SrgN0Y4jMGMUOEoCJX3hYo4OlMd0HBD+ZwSCVoQkSet0XR/S0uP4pZyK+1O7RdcjXiksrfHw58+2sq2omt9d1JPH5m2h3OltVDhQsW8nKnMuoAUWsxrBLJsxySb+PubvnNO56ZISQdvndLg//dJ7k67rrC2o4NXFe1m0sxSzIiFJEtrRH9PkwVncNiqHnJQYjryxFfeO8pDn3FNWwAML/kpOYhYvXvanoP3RA9NImtKz3rbXluzlX1/vwuULXwJbkSTS46188UDjveyCtsmXPx3idx9tRtV0Q0NnCPRIybLEny/twzur93NF7HauL/gDkq7+nJwwwhwNihlunAdZg0/SDH4ZTd2bWjwzVVrjMVSEilSC06/q1Lh9JESH2fD48TRwltULpLKfrcHpg30P2omxBMb02novczb7+PHWoyWEPie8PRl+tQOUFv/nEwgEgubnBEouUmOtvHj9IBZuLeb+uevxqbqhtIVj2491+mCtWLPiSLsimpjuMSB7kSSQkIgyBbywJveYzPVnXE+mXZQMCU5/9pc5ue2N1RyqcuPyqeg6NHxvfWfVft5bc4BRuSk8HRsPEiF1ZLold+GavhcxZ+OnBns1lD3vwAufQ1RioFd84E1MHdOVyionMxfvwSOHDozMikSq3cr7d48UgdRpyMX9Mji/dwe+2XaYGT/msf1Q9fFeWZ+q0TM9lunjcrmwdzoWk8xF5g2YPv4/JBrxK6yLzwk+4M2JcPOn0OmskzuZZqbFowGnV0U2eHBHKsGpyBK1Hn94wVTlAShYFihlaYCqw3OrvDw6xtr4530u2L0wcMMRCAQCwXF6Z8SBJGEUSjXWB1v+/U8o9qHIllIsFjfn98rhwh59ObfzuViVJu7FAsFpxM7iGq5+eTm1Hj8GNk/H8Wk6aDrL8o7wx1gnvzeZoEH51Z6yAr7LW8Flvc4lIy6NourDfLL9OwZ1DG6HkPAS5ZkHR/IDG4o3ww9PofecwI2rzKRaejPD3o8at88wI2Ezy+g6nN+7A09O6hv+oragzWFWZCb0y2BCvwycXj8VRz3KEqLMxFjrhBQl24n5bBqEE0jVxeeEOVfCfWsgNj308a2EFg+mYqwKqsFdI1IJTr+mEWsNcyVk9auNmlY+MtLCP5Z5uOcsCwm2RlZnvbWw9D8imBIIBIIGvLkiH6Py8VB9sP6awN+9wC4pjv+MH3NqBiwQtAJKqt1c++oKqt3hq5x5/BrfVzm5S7LTUG84xhLNxqLtzFzzPtWeWuKsds7PHcFj59wTdB5FKsIi5/+8wXfUB27rPNKiZG669RNu7jqaFXllvLwoj02FVbi8KiZFIinGwo3DOjPlrM4kCin0dkW0xUS0pZEwYtE/wV9f0Cis6i8IfG7VK3B+cDlqa6XFg6nUWONVx0hNwaymgItyWKyfDapxtDyko8K4bBPPLPfw13ODzYCPc2gzVB+CuIzwrikQCATtgA/WFgYpgAF4Du5A93uJ7jEi5DnySmoprHCSlRh9MoYoELQ6nv12NzWNBFL1S2PrW8R4NY05Wg33S9GYpJ+reDJiU5kx6fGQ15VwEWf6sJF9OpKiwrtT4I6vGdmtLyO7GXt2CgTHcZbDzs8N9QjCqv5SPbBmFpzzaKCPqg3Q4pbDVpPClYMyMRn0TcUNvZLEc++gasV7FP73Bgpn3ErN+gVEda//MLYoMjcN72LYe+X1a+w6XMOa/HI2HqjkYHktuKuaHNMT51j572ovpY4mhClMVqgpCm+SAoFA0A7QNJ1qt7G6WCR9sGaTzOFqA5l2geA0xOn1M2/DwSB7GAiUxpZ/N5P44deQdd8cMqe/TuygCbh2rzp+zHxZZ59VASXSXkcPI16+lux/f47D+/O1X1vvZdwbddSNfQ6YcxVo4QtRCNoxG99utO/2kZEWnlnuodIdoslPV2HHgpMwuJNDi2emAG4fncNH6woNbyThmIIhwU0jutTbVFjhZPbyAuau3o+OfrwvS/E7WWuSMDXRrdk3TWFiDxNPL/VyRmoT8eaxVLhA0FxU5AfKULd+Ap5qQA/Ihp4xEYbdDUldW3qEAkGj+DSt0V74SPtgPREoiAkEbZn5Gw4avnuGaxGjAr9Wa5mXnIpS7gF/OArFbizSHmSqwssWeGth9zfQ86LwJyZonxSsaPT9OOzqL28tHFwPfa44SYNsXlo8MwWQm2pnUJdELBGvqoBF8zNCKyP9aH+TX9V45MNNnPevRby+fB+1Hj8Oj0qN20+N20+l34zUhBz8MR4fZ2Pmei8Hqxs79uhLrkDQHJTuhNcnwItDYfVMqC4MBFOeGqg+GPAAemkEzBoPh7e29GgFAkOsJgWpkRXJun2wodB1iItqG+UdAsEv5fOfDuE0EHaIpDS2Bp1vByUSMzgNTDKS2fj1TsKFhBu78hWplscAPbxsgbcWlj0b7pQE7RlXRZO7w6r+goCFURuhVQRTADNuHExanM2w3K8xLIpEp9RY/uDeRP41U6jdk8fN/1vNgk1FePyaYd0+SOTroRVCuiXJTOlj5vnVjSiRqD5Iygl7rAJBo+Qvg5nnQsHyQOOlUT+f5gt4NBxYCa9dEDCQFghaId3T7Ibb6/bBOnetQPO50VU/rry1VPzwv3rH+lWN7JQYw/MIBKcbFQ7j94xISmN9qk6p00viFd3p+PthxF2UjZJorVP650PhIPGmWWRYbyTB/BqSFAjg6mYLmuTgeqgpjmRqgvaIOarJ3XWrv5rE0naeAa0mmIqPMjP/3lF0TY0hyhz6xhFlVuiVEcfH944h91//IPGmG7n/yQ9Yv+8IrhAOzS+rl+LQQ8vt/nGstV4d8XFkE/S9CqyxIc8hEDRJ8RZ4++rAql8oo5Bj+BzwznVQtOGkDk0gOBHuHptLjMX4Hh5OH6wiwWUDOmK3tooqdIHgpNNYNrduaWw4vLx4L4P+8g3/WZyHs18yGb8bStaTo8n86wiybFeSYbsLu+krZCnYPDXsXvFq0SsuCEFSDkhNv8eHrP5SrJCY3fxjO0m0qqdVit3Kp/eN5qN1hcxYlEe5w4vLqx5/xZQAm1mhQ7yVe8Z2Y9LATCymQDxYPHo8KzYuxW2QjWqohDMjLZvxY32cX7/NivyH6gdHneJl3L83KOWTzTA8WF5UIIgIXYf3bggER3UI2zz6nevh4a0gt5o1EYGAi/ul89j8LY3uD9UHazbJTB0jegMF7Ydku7GkeKQWMQDlDi+vLt7LK4v3MrZHKs9cfSbxsgdkBbTGZddFr7ig2Rh0S0CEoonvSt3qr35pjXzf+k0+SQNsflpVMAWBYOmG4V24flhn1hZU8N32w5RUe5AkSIuzcWHvDgzolBC0kjNr6T58BgFuYyaRT23/mJGdS4iWIlSMUqzQaSik9/0FsxQIgP0rwFFquCushmBPFexbBLkhBFoEglOI1aTwwLndePbb3bgiFJGwmmSGd02mRweR9Re0H64YmMm6/IogQ9xILWKO4TkqQPHjzhImPLeEj6ePoIMe+rf4+Dgbg16p5dcjGnvu6GATveKCEGT0h8QcKNnW5GF/HGvlrc1G6q8SdB0nTHubA0mSOCs7ibOyk0IeW+P28dmmoiDz36aUcPZ1O4tF2nOMlTcRLYXp0KxYIKETXPt2RHMRCAxZ9jx4nYa7wjOPdsCy50QwJWh13Hl2V3YdruGLn4ojCqhS7VZeumHQSRyZQND6uLhvBr9vJJsbN/RK5JhEqla8x5EFzyBZorB26EbciCkhz+tTdYqr3Ux5dSUL4ntir9zR5PEhswWqr02VXglakDG/hk/vD1TRHCXs6i9zFIx+6GSPsFlptcFUJGw8UIlFkY+vxhyjKSUcHZl7fQ/ylOk1LlNWYMWLIjXes6L5ZaT0Pki3zBe9UoJfjt8De76lsT6psOVD85eCpxasxk3/AkFLIEkS/5x8JgnRFt5eWYBP04MWu+oSbVHITo6muMpN/hEnvTuK1W9B+8FmVrhuaGfeXJ5vKJwVqjS2KVNfVdM5VOXmvxnT+Z3+a9zwbQ4AACAASURBVGSp8VI/aCJbICkBmWrx/iMIh75Xwa6vAl5RkZSGmqMDNjBdRoY+thVxWgRTVS4fmsFLaSglHA2Z//PfyfvqOKaZPucceSM2iyXgvqzrYLIE/pven8oD6XhLupMelXCypyNoDzjLA0ImmrHBKQQagkf9z8GDw4zr6YGAO7izTARTglaHLEv8YWJvJg/OYtbSfXy2qQiTIqGqOjpgkiV8ms5Z2YncPTaXUbkpfLHlEHe8uYaPpo+kY0LTilACwenEfed04/PNhzhc7aaJdYcgGmtlcO1ehS2rDxAo+5uzL5ZfxZiwUj+YCjtbYLLAiHsjnpegnSJJMGkGfOSH3QvrZagaxRwNg2+F8/540ofX3JwWwZQsSUgEl0KFaxK5Xu/BdF8PMi0Olk3wBl5ONT/YEiBnLKR0I66igr0TLyXh6quxnXHGyZyOoD2gehp1CD9GWA3BkmwspS4QtBLOyIjjmavP5E+X9mbxriOUOzz4NZ2EaDNDc5LJrBM0TezfkUOVbm59fTUf3D2SeOE1JWgnJERbeP+uEVzx0jIqnT78YURU4Zr6BjbG8E2Px5iY/7fwXmzrolghcwik94vsc4L2jWKGq1+HNbNgyb8Cfd7e+oJbGhKyOSrQH3XOY21KdKIup0UwlRTTPEo4flsSDDzfcJ8pMZHUBx6g+C9/pcuct3D7dRZsLmJLURWVTh92q4muKTFMGphJsj207Lrg9EbTdJblHWHhlmJKawMiJ2mxNi7ul86IrslItvgms1LHCNkQrPnAFt+cQxcITgqxNjOX9M8IedzUMTkcrHRx11trefP2oVhNoa0yBILTgU5J0Xz54Nnc+/Y6NhVWoWp6k0FVJKa+Dq/KB55hTMweA/mLwy+9UiwQlwnXzg13GgLBz0gSDJ0KZ90REMxaOQPK8sDnwKfE8H1lGhfe/DhS1pCQC8ytmdMimBrUOdFweyRKOBZFYtKAzCavkzD5KrbO+4oXn13AggoLkkQ913KrSebvC3dyTs9U7h6by8BGxiU4fXF6/by9cj8zl+zF4fEHqTN9tL6QuCgzd43J4broDGw1+5s8X8iGYGs8RKc05xQEghZFkgLlgfe8vY5HPtjMs1MGIEdg5i4QtGVSY628f/dI9pbW8vqyfD5cV4jHrxqW/kVi6gtwpNYLt7wNn9wL2z87GlA1kQGzxEByd7h5vlDxE/wypKMKfV3HHd9k0nV+/9R39I7pQ6c2HEhBKzLt/SVYTDI3Du+MRQmeTjgmkQC618v4VfPxHjjQ6HV+3FPGtNzJfFwMLp9aL5CCQF2y16/x9bbDXD9zFa8sykPXIyh+FrRpDle7ueT5pfzrm52U1HiCAikIBN/FVW7+vnAnl7uf4IjSIeR5GzWPNkUFatiFz5TgNEORJZ67diCFFU7+sXBnSw9HIDjldE2185dJfdn6+HhuH5VteEykpr6qpgdKr654Ba5/H7pfECjhM0cHSsaRwGQL/MkYAJe/CFO/hSixMCxofiRJYlDnBDYcqGzpofxiTovMFMBNI7J5fVm+4b5QSjiyBENyksnyRJF/9TXEjBxJ8rSp9XqjftxZwvQ56wKmwCFWgXQ9EGw9++1ufJrGfed0P6E5CdoOFQ4vl7+4jCM1nrBq3d0+jTzVwpX6/2OB5f8RJ/1cchF2QzA6DLr5lw5dIGiV2MwKr91yFpNnLCczMYqbhv/ssr6npJbZK/LZcrCKWo+faIuJbml2bh7Rhf5ZQiRIcPogyxLp8VGYFSlI6S/SVoaEYz2IkgQ5YwJ/aoph55c/94pHJQZ6xdN6nYzpCAT1GNg5kfUFFVx2ZseWHsov4rQJpjITorj/vG68+H1exEaRMVYTT08ZRFryGJKnTaXyvfc5cNfdWLt3J3naNCp79uOet9fj9mlBn21KktTlU3nx+zzOzEpgTPfU5pqqoBVy11vrKKsNDqSa+n74NThIIvf5HmC25e+RXdAcDUNug+jQPmwCQVslKcbC67edxdUvryA9zoZZkfj3N7vYWVwT1E+yubCSzzcfomOCjQfO685lZ3YMMncXCNoiI3KTUeTgYCqSVoYos8IFvQ0qIWLTA88SgaAFGNQ5kSe/2N7Sw/jFnDbBFMC947pRVuvl3dUHwgqoJAliLCbm3DGMLskxACh2O8l33E7iTTdS/emnFD/+OK92HosvqW/Q58ORJHX5VP7zzS4RTJ3G7D5cw+bCyqAHXTjfDxUzq/Q+HFA600ltun/qOObowMrhBX9t7qkIBK2OLskxvHrTYK6buRJNJ8hP8Bja0YqAvFIH//fRT6zIK+PJK/qhiH4rQRunT8d4OidFs+twbdC+cE19NV1n8pBOp2rIAkFY9MuMZ1dxDW6fis3cdsWGTqtgSpIk/nRpH3JSYvjnwp1omm7Yt6JIYDbJdEu18/x1A+maGuzRI1ssJEyejO2yy1nw+Nc0jM0ikSTdWlTNviMOclJimmeiAkPcPpVKpw9Zgrgo8yn7Yc5aug+fVv8FL5LvhyabeSPhXv5Q9SdQ/aA3YqoomQL17v2vgUv+LXqlBO2GVfvKUTXwqsaBVENcPpVPNhahyBJPXiHknAVtn3vGdePReT8F9WpDeK0Ml/TLEFYDglZHlEWhW5qdrUVVDO7SdittTqtg6hg3j8jm2rM6s3BrMS8vymP7oWpkSULTdSwmmUv7d2TqmK70TA/t5P3tjlIwKaA2EJuIQJJU03Rmr8jnT5f2OdEpCRohIPhRzIwf89hRXINZCaxC+1SdAZ0SmD42l3N6pZ201WmPX2X+xoM0fMeL5PvhU3XeLc3m0XsXoax+GTa/C7IZ9KMnlaRAkNX/ahh+D6QJnzNB+2FzYSXPfrvLMJAKVWb98fqDjOmewkV9Q0uyCwStmQn9MpixKI+9JbX4InH1JdDK8PAFPU7SyASCX8bAzgmsL6gUwVRrxGKSufTMjlx6Zkd0XcfpVTErMhZTZKv5ew7X4vQErwRFIknq03S2FVVHdF1BaD7deJDH5m9B03UcR/8fqXUeMusKKnjw3Q1YzQr/uuZMzumZ1uxjOFLrNTSMjlSy1qfpVNm7knTpszD+SShYAa7ywM6oROg8HKyhg3+B4HTj5R/zDEv7wi2zfuH7PSKYErR5LCaZuVOHcdkLyyipcQeVlTdGtEXhzduH0ikp+iSPUCA4MVLtVuauLgi0S2g6KTEWRndP5fwz0jAZqHS3Rk7bYKoukiQRYz2xqVa5fIYuDHUlScN5Ya5xN1K6JTghXl6Ux7Pf7jIUBamLw6vi8KpMn7OOv1zel6ubuWbc6fEbVttF+v0wyRIOjz9gQG2Jge7G5tECQXui3OHlux0lQR47kZTR7impZdfhGnp0EIsRgrZNst3KFw+M4fY317CtqLpR/ymAGItCjNXE7DuG0itdeEQJWheqpvPppoO89GMe+8ucePwa+444j++ft/EgZlnm1pHZ3DIym8QYSwuONjTtIpj6JcRFGf8TRSpJareJf+pjuLwqu0tqqHL5MCsyKXYruakxYStvfby+MKxAqi5un8YfPtlCaqyVcSEyVJqmI0mENZ5oqwnNYBiRfj/8mn7CAb9AcLry2aYijH6GEZXRahpzV+3nz5eJMmtB2yc+2sxH00eyubCS15bsY+HWYsyKfPx34vVr9MuM5+6TXOIuEJwoLq/KtNlrWb+/wrAHEDhabaQyY1Eec1YV8N5dI8g10Dc4hk/18d3+71hetJwKdwUmxUR6dDoTu06kT8rJv/eLt7cQdEmOIdqiBP0Pj0SSVJGge1rjX4L2Ql1Hd0WWjt/8/apOaqyVu8Z2ZdKAzCaDCpdX5ffzt0QsUw+BgOrX729i9WPn13vAqJrOjztLeHlRHpsLq/D6NSQJYm1mrhyUyW0jc+icbFwikWK3oBvkLiP5fgCYZUk0BwsEDcg/4jD8rUdSRqtqsPdIsApaODi9fuZvOMiS3UeocHixmGQ6JkRx9ZAsBnVOFNLrghajf1YCz183kCqXj/1lTmrcPqIsCh0TougQZ2vp4QkEhvhUjZtmreKng1WNKrPWxePX8Dq8XPHiMj5/YExQuWqZq4zZ22bz/s730XQNp//n7JaExIe7PiQjJoPb+93OpV0vRQmz9SJSRDAVgov7ZvDYvC2G+8KVJDWbZG4a0cXwHO0Br1/jkQ828dXW4iBvmGPsL3fy5Ofb+euC7Tx77QDG90k3PNdnm4sMt4fTPwEBxb/Fu0o5p1cgO/XRukKe/GI7Hr96vO8KAsbLVS4fc1YWMHfVfs7slMC/rzmTrMT6P2SrSWHSgEw+WHuAhiXsYX8/FIlrh3YWK4gCQQNqPcbl0ZGW0Rr1vTZFYYWTl37IY96Gg0gS9RbTJAk+3VREit3K3WO7MuUs8dsVtBzxUWb6ZcW39DAEgrB48vPtbCkyDqSaWhCv9fi5fuZKFj1yDvLR++2uil3csfAOnD4nXs0bdD4dHbfqZl/1Pp5a9RRf7P2C5859jihTVLPPSwRTIYiyKFw1OJN3Vx8wDAJCSZICdE2xt9uaZa9f48bXVrH5YGXIVYhjLywPvruBxy/rw5SzOgcd8/KPeUFZwkj6JxxelZcX5XFOrzSeWbiDWUv34WqiXDDQ5KuzLr+CS55fynt3Da/3//JQlYtqty8okDpGON8PWZK4dWR2k8cIBO2RhGjjOvlIy2jjIsj6riuo4Nb/rcblUw3v+boeuFftL3fylwXb+XzzIWbeMoRoi3icCgQCQWM4PH7eXbPfsNog1IK4pgd6aBfvLmVczzTyq/K55ctbcPgchtVBDXH5XawvWc9d39zFrPGzMMvNWwnUNmQyWpg7RnfFpJzYyqMiSdx5dk4zj6jt8Ov3N7L5YGXE/U1/+nQrS3aX1tteUuOmsNIVdHwk/RMAa/LLmbl4L7OW5jcZSNVF1XWqXD6ufXUlRZUuyh1e/rpgGxc9u4ROSdEM7pxwXJY9EiwmmRG5yUJpSSAwoG9mHDGW4MxT3TJa564VaD43uurHlbeWih/+V+9Yq+6nV1Uh/tLSoPM0ZMvBKm6atYoaj98wkGqIy6eytqCCm2etxhemB5ZAIBC0R+ZvOIhsoH58bEE86YLpRPcciWyxISkmorsNq9cS4fCqvLJoL37Nz7SvpxkGUhVLKtj9+91svXMrOx7YQdGbRaiOwAK8R/WwvWw7/1n7n2afm1hKC4OclBj+cVV/fvvR5oiCgiizTKekaP77/R76ZibQrZ31Te0oruab7YdPuL/pD/O38MNvxh3vS6hw+LAoMt4GGa5IZchNssw/F+7Aa5BOCjWuGncg1Vzp8nFJvwy+fvhsOsTZqHB4mfD8EkprPGG9hEGgTyo9zsbz1w0M63iBoL1xUd90Hp33k+G+cMtoddnExQVLyLvkaaw5OdjPPZfY887Fkptbr+fJ7VO5adYqw4bopu4LHr/G1qIq/rlwJ49OOH094HyqxrfbDjNr6T4OVDhx+zSiLQo902OZNqYrI3OTRQ+ZQCBolNeW7sPpC76/RrIgvn5/BZ/sWES1tzookDry5RFKvywla2oW9t52fBU+it4qIv+ZfHIey0E2ybhVNx/s+oD7Bt5HtLn5FrFFMBUmlw3IxK/pPDbvJ9x+DT3E+3KUWeGGYZ15dMIZfLiukCmvrOCvk/pycb/243cya8k+Qy+McPubDld72HigkoGdEwEaTeVG2j+h6jpGSaRwxqVqOvvLnXxw94h6BnOJMRY+uXcUU15dyaEqV8ig22aW6ZIUw9xpw4izCeEJgcAIq0nhurM68+aKfMN7SagyWkmCs3ul0e+Wp9G9Xhxr1lD73ffsn3YnksVM7LnnEXveuUQNHMiXWw4FLdRAuH5WGnNWFvCrC3pgM5+cBueWQtV0Xvh+N7OW7kOt4+kHgb7SQ1VuVu8rJ9Zm4uHzezDlrE4iqBIIBEEcrAiuLILIFsQtJpk3Ni6oJzQROIdKyfwSMu/IJLZ/wAbDkmqh0z2d2PXILqqWV5F4duBdUpIkPt/7OVf3vPoXzuhnRDAVAVcOyqJHh1ie/343i3aWIknUe2k2yxKyLNGnYxz3n9v9uMjBNWd1oldGLNPnrGdjYSWPXNizzRiRnSg1bh+fbS6qZ6ILkfU3efwqMz5dz99Ty/Ds3EHt7gN40y8BpX7wEWn/hKppqA1SzZGMy6TI/LCjNMitOy3OxucPjGbuqv28ungvDo8fR93GdQI9ePFRZu46uyvXDu182r14CQTNzW2jc5i7ej8+NTIRCQCrSeaB87oDIFks2EeNwj5qFB3+8Hs827dT8933FD/5FP7iYv479iEcxNT7fCT3BYAFmw8xeXBWxONsrbh9Kne8sYb1+yuaLIl2elWcXpXHP9vGpsIqnpzU93iTuEAgEOi6jreRUuhIFsQ1TaOg8jANbtU4dzvRfBpxg+vrEyg2hdj+sdRurT0eTLn8Lt7a9pYIplqSvpnxvHrTEMpqPby75gCbCyupcvmxWxVyU+1cO7QzOSkxQZ/rn5XAZ/eP5oF3NnDz/1bz3+sGkmy3tsAMTg0/FVZhlmXc1P/xRJLO1XRYUVCJq2gTtp496HXBhaT9WEthlafecZHKkMu6jtZg5TSScXn9Gh+vL+Q343sG7Yu2mJg6piu3j8pheV4ZX28rpqTGgwSkxVq5qG8Gw7smiZVbgSBMMhOieO2WIdz+xpqIy6yfnNSP/lkJQfskScLWuze23r1Jvf8+tv2UR9Hc7TRMfkdyX3B6VV5bsve0CaZUTeeut9axtqAiLAljCPSQzd9wkGiLwh8m9j7JIxQIBG0FSZICbRoGAVUkC+K6rxaTXoOvwXa1VsVkNyEZlB2Z4k24CupnxYocxsrQJ4oIpk6QZLuVe8/pFtFnkmIsvHn7UP719U4ue2EZL90wiDM7BT/oWws1bh8Hyl3UevxEH/WvSArThbrK5TMsyou0v8ljiabjk08e//t0qYAnP98eVHcbbv+E1SQjIeP+hX1XVa6GP+X6yLLE6O4pjO6eEtb5BAJB44zMTeH1W4cy9c01+DW9yZd7syJhkmX+Mbk/l57ZMazzFyl2zBYz7gZS7JHeF4oMBHJailqPn3nrC3ljeT4lNR68fo0oi0LvjDjuGpvLmG4pTWaP3luzn9X7yiOWMHb5VOau2s8FvTswvGvyyZyiQCBoQ2Qk2CgocwZtj2RB3CdHY49xBwVTil3BX+tHV/WggMpf5cdkrx/ueNVgKfVfggimTjGKLPHbi3rRPyuB299Yw2/G9+S6ocES4Eboun5KMho/FVbx6uI8vt52+Lizuq6DV9UYnpPEXWNzQzYb1zXlrbc9wv4muUE15KSBmfzl822Gx4YjQx7wgwkO806k70ogEJw6RuQm88NvxvHWygJmryjAr2m4PX5UJGQp0KeqA9cM6cRto7LpkhxcIdAYDq8frRnuCw0XaVoCp9fPXxZsY96Gg8iSVE9Qw+PXWJ5XxqYDlURbTPxmfA9DCwpd15nxYx4ug2bxcHrI3D6VVxbliWBKIBAcZ+roHJ76YofhfSXcBfFu6SYqLbU0jKaiu0UjmSSq11UTP/Rn3zXVrVKzuYYOkzvUO96qNG9lmAimWoiL+qbTvYOdu95ax8b9lTx+eZ+g/hld11meV8Yri/aytqAcl09FkSRibSYuG5DJbSOzyTYoKTxRyh1e7nhjDTuKa/D4VTSdoFXJxbuPsK6ggmS7ldm3D230+sl2q6FIR6T9TbENxBlirCZ+P+EMnvxie9iy5sewmWVuG5XN7OUF0KD8MNJxxQhPGYHglJMWZ+PXF/bkgfO68/2OEjZ/9h3VXo0Oo4fTKSma8X3ST6gP0W41IRus/kR6X7CZW7YXttzhZcorK9hf7mwye+fwqji8Kn/+dBvbiqr582V96i2OrSuooMwRvHIbbg+ZDizLK6Ok2k1anK15JicQCNo0VwzK4skvtje6P9SCeIxFYeqYLvxtS3BlkBKtkDYpjaI5Rcg2uZ6anznJTMLI+lVgGfbmFYM7vVUQWjm5qXbm3zuKGo+Pa15ZwcE6JSILtxxixN++Z9rstSzeXYrTq6Lr4Nd0Kpw+5q4qYPyzi5k8YzkFZY5fPJaSGjcTnl/ClqIqXL5AINUYDq9KYYWTS19Yys7iGsNj+mfYUTR/0PZI/GHMisQVAzKDznHjiGxuG5VDVAQvTTazzG/H9+LOMbmGNbuRjEuRJEZ1E+V7AkFLYVZkxvdJ52a5iIc6erjv3O5cPiDzhAVdclPthj5RkdwXAHIiyIY1N26fynUzV7LviCOiHqf31xbyr6931dv+3poDhqvHkfSQScBnmw+FNQ6BQHD6Y7eauGZwpxNadJKlgIn7pP496Z1s3I+ZOiGVDld1oPi9YrZN30beX/IwJ5nJ+W0Ocp1rRpuiubXPrSc6DUPE8noLY7eaePH6QcxcspfLX1jGs1MGsKO4mme+3tlks3VAJlhn/f4KJv53KXOnDqdfVnyjxzeF26dy7asrORKBR5KmBzyXrn11BQsfOvv46qO/tJTKjz6m8v33uSL3HN5JGYBHr7/iG246V5Ykbh7ZxfD6v72oFx3jbTz5xQ4kCUNvGAisZCDB36/qz8T+gf6J8X3SWbC5KChgDHdcFpPMnWd3DevfSSAQnDy0mhrkzuGVSTdFdkoMPTvEsqmwKmhfuPeFGIvCnWfn/uKxnCj//X43+UcchvfwUD1Ory3dy8X90unTMfAMOVDuNKwsiKSHzOPXWlUPmUAgaHl+P7E3Px2sYtuh6rAXfSQCVUlzpw1DkSVu63MbO8t3BsmjAySNTSJpbFLwSeqg6RoX51x8IsNvFBFMtQIkSeLOs3PpmxnPnbPX4vZpEQc117+2ks/vH0Pn5MhNyD5aX8ihSnfED2EINDm/9OMeHkmrpeLd93AsX07c+PFk/vd57u3cjXf+/j0Y/GBCpXNlCQZ1TiQrsfH53DgimysHZ/HJxiJm/JhHcZUb89HGQ5+qkZMSw/Rx3bi4XzpW088P/2ljuvLNtmLDMsFw+q6yEqPom3ligatAIGg+1NoalLjYZjnX9HG5/PqDTfV8lI4Rzn1B1lQu7NOhyWNOFj5V460VBYYvJ+H0OPlUndeW7OM/UwYAweXdx4i0h6yxRS6BQNA+sZhk5kwdxh1vrmHTgSrDDHi94xWJWJuZd+4cfrwP9uyss4m3xuP2u9GIsN1DsTGl5xSiTFEnPAcjRDDViuiTEY9P1U8oqHF4/Pz6g418cPfIoM82ha7rvLzoxBuNfarOu0t2M2XXW3SYMpmMvzyBEht4uYkCHptwBn/7MvL+JrvVxD8m9w953J6SWrYfqqZbmp30OBvRFoUh2YlcMSiLzATjH0sPbzkDy/JYF9cZD5GVBdnMMo9f3if0gQKB4KRR4/ZRVutlv1MiyxyDXdN/sa/R+Wd0IM5mxuVtuszZiCgFJhcs5/BvviX90UcxpaY2frDPDVs/huX/hfJ9oLpBsUBsRxh2Nwy4HmxxjX/egK+3Hg7y9IPwe5xUTeeLnw5xz7hcCsqclDuNla4i7SFLCVP9VSAQtB9irCbm3DGMjzccPL4Q7var9bLhMVYFRZK4cXgXbh+dQ0odKyFFVnjtwteYsmAKDp8D3VA7OhirYqVvSl8eGvxQc09JBFOtiffXHsDofSCcoEbTYXNhFQVljohUrNbvr6Cs9sQbjQFkq5Wtf3iOM4Z0CjrPzSOzKXN4eXXx3pArEBDISNmtJuZOCzSUG6Fp+tEf4R6KKt3HxTKOsTq/nBd/yGPy4CzuHpdbL6iq/mohxU88wb8feJC7y5LYcagmbAUum1nmL5f3ZWSu6JcSCE41mqazeHfpcUEeiyKjx49B/bEW+8pvuW10Nted1fmE/ftMisw704Zz6QtLqXX7w3w8B+4Lo7un8thjj1E+42X2Xj6J1AcfJOHqyUh15Ug1FX54ClbNCPzdW6fX1e+Bin3w3ePw7Z/gzOvgoqfBHJ54wzurC+oZhB8jkh4nj19j0ovLGJydRMd4G0WVrqPl5D8TiYRxjEVhYJfWa/0hEAhaDpMic82QTlw9OItNhVX8ZcE2ymo9dE21k2K3cHaPVMb3ScesGPdXdY7rzJwJc7h94e3Uemvxak1LnUeZohicNpj/nPMfTHLzhz4imGolaJrOzCV7gzI4kQQ1mq7z+rJ8/nxZ+JmTJbuO4PqFD2GnX2fh1sNMNgimAB6+oAfZydH85fPteHyq4UPfJEsoskSfjnE8O2Vgo+WKbp/K9DnrWLWvvNESkmPb31m9n/kbDvLG7UMZ1NFOyTPPUPPd93R69VWi+vbhPb/KQ+9u5IedJfhUjUbMuYm2BLJXz107kAt6t0wZj0DQntl0oJJps9fi8PiP3z98qgqyBVTwOLy88N0env9uDzcN78JjE844oUxVdkoM8+4ZyZRXVlLr8Yes6Y+2KFzQuwPPXH0mJkUm7de/Im7iRIr/+EeqPvmEjCcex9qtG/i98O51ULAcfMF1/sc5tm/Tu3BwPdz6GdhClxQXV7kNt0fS42SSJe4/rzt3j82l0ull2FPfYWQjEW4Pmc2sMLZHWsjrCgSC9oskSQzolEDvjDhyU2O4dVRO2J/NTchl/uXzmbt9LnN3zMWv+ev1UcnIWBQLnWI7cXu/27k4+2KUMD0DI0UEU62E7cXV1HqC1e8iCWp8qs4nGw82GkzpmobucqG5XGhuN5rTSXFBUbOY61Y0UhZyjCsGZXHZgEwW7Srh5UV72XigEp9fg6OZqEkDMrltVDZdU+2NnsOvatzxxhrWFVSElU3yazo1Hj83zlzJcwc+5wybn5wPP0BJCKyWWk0KM24czJ6SGv63NJ95Gw5ikqVAtyPgV3VSY63cPbYrkwZmEi3k0AWCU86S3aXcOXtdyMz2sXvC3FX7yS9z8OpNQ476ykXG/2fvvuOrqu/Hj7/OOXdnB5KwRwgzLAHZS0TqrkXFgQMnaLXaob8Oq7X229pqrdaFdDaO5AAAIABJREFUCxcO1NZRUeuoLNkgIhtCWCFAFhk3d5/z++MCJuTc5N6YcYPv5+PBo+29557zuUl6znmfz+f9fudkJvHFLycxf+Ve5n21B18wVCuP6viDn0GdU5gzqRdn9s+sVVbc0bcP3V9/jbIFC9h79TWkXjaDjKy1KHu+gmCUBRmCHijaCvMvhlkfgaX+5XKBCOsSY8lx0g2D4LEnSqkuG9MGZLHw20LTJY8N5ZA5LCrXje/RqJ+/EOKH50BZNRP71LM8OoI0Rxo/Pe2nzB4ym0X7F7H84HJKvaXYVBsdEjpwbva59Evv1wwjrk3uDuNESZXf9MITa1BT6fax96qrw8GSx4PuqcaoDgdPhs+H4nCgOp2oDgeKy4m3y2RI7V9nP7E31234oqmpClP6ZTGlX3h2JxjSjzX3je6C++jnO1m/zzyQqrdaVVDnrs7TWHnfOWh2a53P5mQm8efpg7jn/P5sP1RJhTeIVVPITLLTKyOxRRolC/FDEUvz8S0HK5j9asOBVE2eQIjlu4r53bvf8uDFDeddmkl12bhtSm9umZzDou1HWLm7lBK3D7tFo1OqgwuHdKp3ObWiaaRfeSVJZ55J+V9nY5QsRtG+O2/1eLSS6gDk35FIgi38s3h+vZ/5GwMsmnVsvyE/HNoEa+fB6DkRj1XlCxLppxlLjpNVU0l2fnd+vGNqb77YdiTmIhIK4LJZmDnSvBKrEEKcrOCohy5pjS8KYVEtTO0+landpzbhqGI4fqscVdQR1HWzFRUxBzW6otLuttvQElyoTieKw4nqcob/u91eew0/0PvLXVg+21Gn6EWsicaZSbHnKVgirIU14wuGeGn5HtNCFtHklAUtNhZuOswlw7tEPIbLZuG0bmkxfw8hRGQhPcRXB79i3rfz2FyyGV/Ih6ZopNhTmN57OjP6zqBDQgfTz97z3rcRb+brL/et896GAq4f35M+WY2v9qepCmf2z+LM/o1b3mvNyqJ9bjXsr3veChnw2Co/v51Qz7kz6IHlj8Go2XAsAC0s97BmTxnr9pSydm8Zu4vcpDgtqAp1ZpFiyXFSFYVhNc5/OZlJPDVzGHPmr6u3TUdNCuCya7x+8yjSpPiEECIKhmFQUOah8/cIplqbBFNxIsVZd8YEYg9qHDaNxNGjGtzuuB/lduDxL3bWCaZiTTT+yWl1m+s2pU82HcIwaXwSbU6Z2x9i7qJd9QZTQoim9d7O93hk3SP4Qr5aa9mDRpASbwkvb36Zlze/zOkdTudP4/9Ee+d3xV3yi91sPlhhut/oKo3qvLAsn782cnaqSZTkQeE3pm/dNdbG377ycevpNlIdkWfqQtVH+Xzh2yys7M26vWV4AiGGd09jRPc0/ji0EwM7p1Dq9jP5oUWmOV7R5jh1Ta/b8mFy30zmzTqdm15Zi65T7wxhgk3DadN48+bR5GQ2Tbl6IcSp72h1AFVVSHaY3we3BRJMxYn+HZNNS6LHEtQowOk96m9WdrKczET6dEhi4/doVmm3akzue1KicdF2WPkUbPsI/JXh0dmTIfeicPnf9OiTDAGeX5r/vatVFRz1srWwgv4dYys7LISIjWEYPLr+UV7f+jrekHlxBOBEBaZVhau4+IOLeeWcV+ieHF4eNm9ZPvr3KvcN728o4PfnDyDR3kqXui3vhav4mRjRSWNyDwsPL/fxpymRq/YpwWra7foX48c/zJ1Te9OzfUKdZZIdU5yM6pnOkp3FpvtoKMfJZdO4ZbJ5w+Gxvdqz/Ndn8q91B3h2yW4qvQEAQoaBpijoRrj33pxJvThvcEcc1uZJ8BZCnJrCS/xi75EaTySYihMum4WfnNaZt9burxNURRvUOG0asydlx3zsWyaFm1WaLadp6CJsR+e6Md2+y/cqWA8f/hyKtkEoCEaNohqBaljzAqx7CTqdBhc8Bhl9oxrj/jLzClix5JRpqsLekmoJpoRoZi9uerHBQKqmoBGkzFvGtR9fyzsXvkN7Z3s+3XLItLBCLA9QrKrKur1lTGpEYnOTKC8APRDx7T+eYWfcPDd3jIq8JE4FRqR5GBGhWupxd5/djzV7VsSUXwbh82Jmkp1zBnaMuE2K08r143ty3bgerNtbxv6yaty+EEkOCzmZieR2kibmQojoeYNevj7yNUd9R9mwv4yEVB/uwDASrNG39oknEkzFkevH9+Df6w+YzlA1FNRA+II3JrtdzMf9UW4H3lyzn5W7SxosBVyTRYWO/iqmvvxnAoP+hrV8A7x9LQTqqVilB8L/9q2E56bAzLehe8ONhn0R1uzHklNmGAZuk4qJQoimU1BVwFPfPIUv5KvzXtnSMor/W4z/iB/NoZE8PJmsS7LQEjQMDMp95Ty4+kEenvSwaXVTiO0Bim4YHG2g0mizCtV/7IGZGuf3sfDgMj/9M+rJIW1gPwADO6fw2OVD+dmbX0ed42RRFVJdVt68eUxUM0qKojCiRzojYlwBIYQQAPsr9vPattf4985/oynh874/qGMAkxY8ybk9z+XqAVfTO613aw81JtFXABDNLicziR/ldsBhjf3X4rCq3H9hbqMqz6mqwjNXD2dQ55Soj22zqHROc/HW7y8iY9wYCm84B+PNq+sPpGoxwF8F8y+Bw5sb3DrSuGrmlDVEURQSHfL8QIjm9PrW19GNujfzxR8Xc+jtQ3SY0YEBTw0g+/fZ+Ev87Hl4D/qxhzhBI8ii/Ys46j0acf81H6A0qLULcSY23Gfp/skOnlvvp6CinjbBrugekk3L7cDz15yOy6bhbCA4SrBp9GifwEc/m0CHlOiaAwshRGMYhsE/1v2Dn3zwExZsW4An6KEqUIU74CZgeAjiwRfy8UHeB1y58Ep+v+z3BPW28/Bbgqk489ClQxjSJTWmgMphVbnrR/2YlmteESu6fWi8cfNoLhnWBbtFjXghtllUbBaVyX0z+PD28WSmOGl/8010HVuMotd+Et3j0UoyH6rE7f/uJuH59X4mv+T+bqOAGxZcDSbFJQxdx7N5MyXPP0+n8kOm46mZU1a9YwV6wIsRCuLJW0vZl/NqbRvUdXpltM0pZCHaAn/Iz792/ovASUvbQp4QR947QqerOpE0OAnFomDLsNH11q74i/2UL/8uZ1MxDP697H6ScJ+8eyC2BygqCmmuVqwq12MC2Oo/5+Skq1yWa+WfqyPMPtkSoPe0qA85vnd7Vv32TH59Tj86pzpx2TQS7RYSbBpJDgt2i8qY7HSenDmMT++cSGayBFJCiOZjGAb3fnUvr297HV/IR9CIHCSFjBDekJdP9nzCbV/cRiiah2ZxQB7TxxmbReXVG0Zx1zvf8N9NhwjqhumyPwCnVcMwDP580SCmN0GVOqum8qefDOKus/vxzroDPL90N4fKvVg0haBukOq0ctXo7lw1ujtZNS/AOz9FwfwPPqryv5WFcGAtdD0d/4EC3Mu/wr1iBdUrV6GlpZEwZgw3DB/J77fppkUoos0py26fKFWmhGhGaw6tMX29emc1ekAneXjtfEXNoZE0OImqzVWkTQyX5fbqft4tWMJ5nYbzyl4HJ69Yi6UoT1A3Yi7K06R6TgoX3vGbB4bH3TvJzqsbI+RWGQYMujSmwyY5rFw7tgfXjOnOtwXlHDzqxRsIkey00Ccrqc0newsh2o55m+bxyd5P8Aajy6EF8Ia8rDu8jr+s/gv3jL6nGUfXNCSYikM2i8pjl5/GriNVzFuWz7+/PoBVVUEPoXt94HSS6LBw04RsLh3RNWJZ9cZKcVq5YXxPbhjfE103cPuDOK1a5L5Qyx4NL9kzEU35XyPgwfv8HApWtEevriZhzBgSJ0wk6+67sXYMJ0VfENL5wwOfRRxzQzllCXaNORGqVQkhmkapt9S0hUGoKoQl0YKi1T0HWFIsePbWXh581O5i1gXTmf/IYtAbV+7boipcPLwzTlsrVpdTVRhzG/zvT+GeUcfsubP2Q52uKSree0wK46gWGHIF2BoX/CiKwuAuqQyWjhBCiFZQHahm7jdzTYsR1ZdDC+GA6t2d7zJnyJxabTPikQRTcSwnM5E/Tx/EPef3Z2thBcW793P06acY+OjfGNAxuVH5UbFSVYWk+mr/V5fCwfUR346m/K+CgUPJp8tj87H362f6vayaypzJvXj8i10xV6tSCFdLPPt7LIMUQjQsqAcxTLqPa4kawaogRsioE1AFy4NYEmtfikJ6iK7pLoZ2TWV1fqlZP/MGH6BoqsJ142JrwdAshl0TbhNRWQgmuWT1siXAhF80z7iEEKKZfbLnE9N7uuKPiyn6uIguN3YhcUAigbIAB189yJ6H99Dzdz1RLeGH94qi8Pb2t7ll6C0tPfSYSM5UG+CyWRjePZ0p/bMYdmgbuZ1SWiSQioq7CLT6cxL+eIadx1f7KXJHvpFQNAuO7ln1fq85E3sxoXd7nDEW6HDZNd64aRQ2i/y5C9Gckm3JqErd/5+5clwoFoWKdbWb8Ia8ISo3VpIwoHZe0fHyuP/3k4G4GjGz5LRqzBzVjV4ZiTF/tsk5kuG6j8CRCkoM38WWAFe/CykyrSSEaJte+PYFPMHaKw9iyaH1hXy8tvW1uC9GIXeXbYiamIjurn/tfYsL+aGBwK5m+d+IFBWCdUsp16SqCk/OHMa03A5R3WDZLSppLiv/umWs5EoJ0QKGZA4hEKqb+6O5NDIvyuTg/INUbqzECBr4i/zsf2o/1nQrqWNTv9tW0RjTKdxDKicziRevGxlTQOUwQkztn8k95w34/l+oqaT1gDlLIb0numE1nWk7wZYICZlww+fQeXgLDVAIIZqWL+TjQNWBOq9Hk0N78n4OVx9u1rF+XxJMtSGqy4Xu8WCY5BC0GkcqRPHEoMHyvyE/OFPN36vBqqk8etlQHr1sKMO6pWK3qFjU2sFcgl0jxWll9qRsPv/FJPp1kCa9QrSE9s72jOk0BsWkJnnGuRlkXZzFoQWH2HLLFvIeyMOabqXn3T1Ra8w2W1QL1wy45sT/HtkznX/fOpZu6S5cNi3isxunVcNuUbn0yHr+6NyHqsbJ7P1xKV0ITH+fgrUdMLqOA4s9XJzi+D+LIxw8XfQ0/GILZMVRMCiEEDGq9FdiVeumiTSUQxusqn1Pqakalf7KZhtnU5CcqTZEUVVUhwO92oOWGCclvpM7gdXZYH+pmuV/B2WaxPCp3cL7iYKiKEzL7cC03A7kF7t5/+sCCso9eAM6GYk2RvZMZ2r/rMgFM4QQzWZW7ixWH1pdZ2kHQPqkdNIn1V9dLzslm5y0nFqv9euQzOK7JrNubxnPLtnNl9uPoCoKqqIQCOlkJTuYPTGbnwzrjHVvL/bNmoVryBDs2XGQM1VDyQvzsJ8xC/WGX0HFQSjZBb7K8JK+1O6QHl/jFUKIxrKqVtOeg7Hm0GKATW3FFhdRkGCqjQkv9atq+mDq0Lew+jko2go+N9iToONgGHkztK+nE7WqwahbMJb8HcWkWktNEcv/WhNg/M8bNeye7RO486w+jfqsEKLpDc8azqD2g9hQtAF/qJ6lvSbsmp1fj/y16XuKojCiRzojeqTjD+qUewL4QzrJDguJdst3+ZZ9+5Bxx88o+OUv6fHmG6j2etoytKDA4cOUL/yIXh8tDL+Q3Cn8TwghTkFJtiTTgkQ1c2hTRqaceP14Dm3WJVm1tvfrftIdrdjiIgoSTLUxakJC0+ZNbfkAFj0IZbsh6AejRqW8grWw/hXIyoXJv4XeU+t83DAMKks6k+j31cmtjrr8LzoMuqQJvowQorUpisLjUx5n5kcz2V+xH59efy7kcQ7NwR/G/oFhWcMa3NZmUclIihwkpV52Ge7lKzjy0MN0uOd3J153+4JsLaygwhvAoqq0T7TTv2NSixT0KXnhBVIvughLu3bNfiwhhGhtqqJyZrcz+WzPZ+h8N0NVM4dWdai1qvmdnEMLkNsul1RHw2kgrUmCqTZGTUxErzLv6RQTXYdPfg1fvwqB6gjbBMP/CtbBW1fBuJ/DpLtPFJzwbt/B4QceIFTtxnnRxaj7P4q8r0isLph4V3iZixDilOCyunjt3Nf4xaJfsO7wOvy633S5B4DTEl7e+/Ckh5nYZWKTHF9RFDo+8EfyfzKdyjGjOTRwJC8sy+f9DQfDOZbHYqeQbpDstDJ7QjYXj+hCcn1tIBrgC4b4ZNMh3lyzn6JKH0FdJ9lhZUrfTC7rnUDF+x+Q/Z8PmuT7CSFEWzArdxZLDiyps+w749wMtASNQwsO4T/iR3WqJA9LpuvsrrVyaBMsCVw/8PqTdxt3JJhqY9TEhKYJpj79bf2B1MkCHvjqUdCshIbcSNHjT1CxcCEZt99G6owZ4Se7b14B+Uui36fVBbnTG73ETwgRv1xWF3PPmsvWkq28suUVPtv72YlkZAWFoBGknaMd1w+8nvOyz8NlbVxj2ki0lBTa/+0hbn36f6zp6COgG4RM4rlqf4i//Xc7f/3vNh66ZDAXDOkc03HcviD//GInr63ah2EYuP21++BtP1TJU18EGTn1Z9xrOJFFyUKIH4qB7QfSMaEj+eX5dZb8RZNDa9NsTfaQrTlJMNVGeAMhluwoYmtqf5QtFWRa9jGwcwoDO6c0/OGT7fgvrHu5VtDT49FKqgOQf0ciCbbwY9vn1/uZvzHAolnHZo0C1Rj/+wsFD87HOvwCshd+iCUt7bv9Xv46LPwFfPPmd7NaZjRbuBT6qDlw5r0NllYXQrRd/dv15y8T/sJvRv2GXWW7qPRXYtNsZDgz6JXaq9mW2PmCIa5f7WFzZl98wXqLkZ9oBH7XOxspqw5wzZgeUR3jSKWXy55ZycGjHnxB85m38OsqX4VSuOjJr5h71XAm9smI5asIIUSb9cjkR7hy4ZVUB2NbueTQHPxzyj/R1Nh7DbY0Cabi3P7Sal5evoc3Vu9DURT8tn7ou4JY924BoHOqgzmTczh/cEcc1ij/4JY8ZDp7FDLgsVV+fjuhnoRt3U/nCzuhzb6/7nuqBhc8BqNvhZVPwcYFoFrBOHYjoxD+78OuCRe2kMpVQvxgJNuSo8qHaip3vrmBLQcr8BnRV/X0BnT+/NFWuqa7OKNvZr3bVnoDXDp3BQfKqk1nvE5mEJ4Fm/3qWubfOIrh3eM7oVoIIZpCr9RePHPWM8z5fA7VgWrTohQnc2gO/j757wzNHNoCI/z+JJiKYwvW7OPe9zejGwaB0PGAJBwwBY89Sd1V5Obe9zfx90+389bsMXRNb2CpTPGucOU+E3eNtfG3r3zcerqNVIf502JFAa1oNVQVQWKEp6sZfcNB1bT/gwNrwFMWnolypUOXkWB1NPzlhRCikbYWVvDl9iN4TWaL3FsWUbHmPQIlB1BtTqyZ2aSMnYGjSy4QDqjufW8TS+4+o95Zs9+++y2F5V7TQKq+Y3gCOte9uIbVv5sa/QMwIYRow4ZmDuWN897ggRUP8E3xNxiGQUCvXd3ZoljQVI0+aX343ejfkdsut5VGGzsJpuLUi1/l87dPtkdcOlJTtT+ENxDi/MeX8eHt4+sPqNa/DHrI9K0RnTQm97Dw8HIff5pSX8CjhGedxt5W/8DsidDrjAbHL4QQTemFZfnfPYCqoWL1u5Sveod2036Ko+cwFM2CJ38dnp2rTgRTACVuP+v3lUWcPSpz+/l082H8JufnaI4R0g0++raQ6cO6NNE3FkKI+NYzpSfzzp5HQVUBb2x9gw/zPqbUW47TqpFgTWBy18lcNeAqslOyW3uoMZNgKg4t3VnEXz/ZhjcQxdqRY3QjvOxkxjMr+PJXkyM/8SzZCbpJr6dj/niGnXHz3Nwxqp4GaUEvlOZFPTYhhGgpld4A//nmICG9djCl+9wcXfYa7c69E1ffsSded+WMwpUzqta2nkCIZxbv5tlrzIOpBWv2maZ6RnsMtz/E04vzJJgSQvzgdE7szK9O/xVj0mbxxP928cas0a09pO9Ngqk49ODHkQOp+paP6AZUeAIs3FjIxcMjXKR99feoGpipcX4fCw8u89M/o55cA29ltF9HCCFazKaCCmyaWmdW31ewDSPox9VnTIP7MAxYubsk4vuvrNxreo6O5RgHSj3kFVXRKyOxwW2FEOJUU+L2k55Yz4P7NkSCqTiz/VAleUXmpc+jWT7i9oeYuzivVjBlBAL49+zBu30H9gNFNJSxdP9kB8OeqeKXY+opROGS5GkhRPwp9wRM05tDngpUVzJKlJWhqr0BSl54AcViAYsFxWpFsVhRrBaKy80fNMVyDIumUHjUK8GUEOIHqaTKR7sECaZEM5j3VT4Bk3X4sSxR2V9SxbInXyZ73xa823fgz8/H2rEj9j590Lp0wwjloRj+iGPISVe5LNfKP1f7GZRpctNgTYCOQxr/JYUQoplYNQWzshGaMxm9ugJDD0UV7GgYBItLMIJBjGAAIxCAQBAjGCRgmWza0iHWYxwvyS6EED80pW4/6RJMieawdk8pJnnTMS0fIRBkS5GHgSNHknbV1dhzeqE6neH33MXwjwEQoQXUcfdOsvPqxki5VQbk/qThcQghRAtrl2hHN5mbsnfuh2KxUr1jBQn9xje4n+QkJ1n/727T9xy//8Q0EIr1GEkOuQQLIX6YStx++ndMbu1hNAk5k8cZt8/8SWUsy0dCVivKlGmkTjSpiJLQHnr/CLZ9CMZ3M2B77kyqtVnXFBXvPSZ/5KoVhs4EWwMl2IUQohUM6pyCw6LVOZeq9gRSx8+k9LO5KKqGo+dpKKoF754NePdtJO2M609sa9MUpp/WOeIx+mQl8s2B8jqvx3IMX1AnJ1OW+AkhfphKq/yyzE80D5vFfC1+LMtHNEXBYa2neMTEu2DXZxDwxD5AzQZjfhr754QQogVoqsJ143vwxBe76vSZSh45HTUhjfIVCyj+8GEUmxN7Vg7JYy6rtZ2iKFwztkfEY8ye1Iu73v4Gt7/uw69ojqEAE3u3p31iPXmpQghxitlT7GbR9iMcrQ6w8cBROqQ4GNgphW7t2vYDegmm4kynVAf7SqvrvB7L8hGLppKVXE+ZiY6D4YLH4T+3xxZQWZxw2SuQ3jP6zwghRAu74vRuPP7FLtP3EnPPIDE3cv87TVEY0T2NzqnOiNucNSALTY3c0LehYzhtGjdP7BXxfSGEOFWEdIP/bTvC3MV5bCoIz+j7g+HF2K+v2ssbq/cxpGsqcyZlM7lPJmo959Z4Vc/0hWgN14zpQYK97sxTzeUj1TtWoAe8GKEgnry1lH05r872k/pm1H+gwZfCRXPB6gwv3auPZgNbAlz+GuRMjeXrCCFEi2uXaOf+C3Nx1jdDH0Giw8JfLxlc7zZWTeX2M3NwRurnV+9nFXplJHB6j7SYPyuEEG2J2xfkqudXccebX7Nubxm+oI4v+F1Wqz9k4AvqrM4v5bbXv2bWS2vwmMz4xzuZmYozZw3IQjPrBkl0y0esmsIVI7tit0Rxkc+9KDxLtWoufD0fUMDvBgxADedFKSqMuAFG3gQpkXMIhBAinlw+shul1X7++cXOqBqga4pCosPCGzeNpktaw0tObhyfzbbCSj769lDUVfksqkL7RDuvXD8KJcJ5XgghTgXeQIgZz6xg15GqOn3/zFT7Q6zaXcIVz61gwewx0d3HxgkJpuKMVVOZNbYHzy7dbXoD0NDyEVVRuGZMj+gPmJ4N5/wNpt4PWz+Esj3gPQrOVGjXG/qeC5ZTI0FQCPHDcuvkHHq0c/GHD7bg9gVNc5ysmoKqKAztmsrDlw6ha3p0a/cVReGhS4aQ7LDy5pr9+IIhdLMGV8e4bBpd0py8ftNo0k6RpGshhIjkV29/Q16UgdRxvqDOtkOV/OZf3/LIZUObcXRNS4KpOHTblN4s3VnM5oPl+M3qpEfgtGr84cLcqG8GarE6w0v/hBDiFHLuoE6cnduR5XklzF2cx7pdh/GioqkqSQ4L04d1YdbYHo06b6qqwn0X5vLj0zrz3JLdfL71MJqqhAMrHax6ENVqJTsziTmTe3F2boeIRYaEEOJUcfCoh8+2HDYNpNxbFlGx5j0CJQdQbU6smdmkjJ2Bo0suAN6AzoffFvLrc/uRmVRP/n8ckWAqDtksKq/cMJLz/rmM/aXVJh1T6nJYVe4+uy+Xnd612ccnhBBtiaoqjO/dnvG927PvxptInXklyWdEnuGP1dCuqTw5cxhlbj+fbz1Midsfbr7+2ceM7N2RUbPObbJjCSFEvHt1xV4Mk5vXitXvUr7qHdpN+ymOnsNQNAue/HV4dq46EUxBuOLp6yv3cedZfVpu0N+DBFNxqrDcS6U3wMxR3fjgm4OEDKNO3xRNAatFJScjkV+f05/xvdu30miFECJ+bTlYwaaCciq8ASr97egbTGZqUG/yWaK0BBuXjvjugVZpwQp8+Vua9BhCCBHPgiGd+av24g/VnpXSfW6OLnuNdufeiavv2BOvu3JG4coZVWtbX1DnpeV7+NmZvdtEdT8JpuKQxx/ip6+t5zfn9mfGiK7cd2Eun285zCsr9nKw3IMvoJPosDC8exo3jO9Jn6ykhncqhBA/IL5giI++LeTpRXnsL/WgAAFdR80ahWXpEZRlnzFzVDeuHtM9qoITjWHv14+Kjz5uln0LIUQ8KnH7CYTqLu/zFWzDCPpx9RkT1X6q/SEqvAFSXfGfYyrBVBy674NNDOqcwqXDuwDhohTnDOrIOYM6tvLIhBAi/u0vrebyZ1dytNpft+iEZsPn1wGdF7/K5+Xle3jgooG1ZpSaiqNfP/SCTRi7l6DoAXCkQma/cKsJIYQ4BVV6A6Z9+EKeClRXMooaXZU+i6ZQ4QlKMCVi9+7XB1i7p4z/3D5eSucKIUSM9pdWc8ETy6jwBOqtrgccK/Bj8Pv3N+H2BZk1rokakvurYdM7aMsepcfkA/DGFaCqgAGhIAy5DEb/FDLaRj6AEEJEy2HV0E0K+GnOZPTqCgw9FFVAFdINnLa2UR5dygrFkbyiKh74cCtPXDmMBLvfCHTKAAAgAElEQVTEuUIIEQtvIMTlz66IKpCq/TmdBz/ZxrKdxd9/EDs+hYdz4ONfQ2keqqajBKrAVwG+Sgh6wn39npkAC66GgPf7H1MIIeJE+0Q7IZMTsL1zPxSLleodK6LeV6rL2pRDazYSTMUJbyCcJ/XLaX0Y0Cm5tYcjhBBtzkffFlJWbR5IubcsovDlO9n3yCUceOJqDr91H94Dm0+87w3oPPjx1u83gI1vwVvXhJufB9yRt9ODEPTCzk/hxXMkoBJCnDIcVo1puVmcvNJPtSeQOn4mpZ/NpXrHCvSAFyMUxJO3lrIv59XaVlMVzh/cEavWNsIUmf6IE39auIVemYlcObJbaw9FCCHapLmL86g2acwbbTneXUVV7Dhc2biiPnu+gg9uDwdJ0Qp64cgWeHsWXPlm7McUQog4dNOEbL7YegRPoPb5OHnkdNSENMpXLKD4w4dRbE7sWTkkj7ms1nZWTeGG8dktOeTvRYKpOPDhxoMs3VkseVJCCNFImw+Ws7/UU+f1WMrxBkI6LyzL568XD459AB/fVSeQ6vFoJdUByL8jkQRb+Nz+/Ho/8zcGWDTrWBGKoBfyF0PBOug8PPbjCiFEnBnSNZVu6S52Hqmss1IgMfcMEnMj9/nTVIWcjMQ2tUpLgqlmdLjCy/yVe3l/w0EqvAEMA5IcFs7O7cC1Y3vQNd3F3hI3972/mZeuG0myo22sDRVCiHizqaDc9PVYyvGGdFi7pzT2gx/6Fkp3m+/TgMdW+fntBHvkzwe9sPwJuPTF2I8thBBx6PlrR3D+40up8ASJNoVVUcL3yc9dO6JZx9bUJJhqBvtKqrn3g02syCsBws3Hjiv3BHh5xR5eXbmXwV1SOFod4LYpOQzqktJKoxVCiLav0hskaFJCKtZyvFW+YOwHX/EkBP2mb9011sbfvvJx6+k2Uh0RVh4YOmxbCNWl4EqP/fhCCBFnuqa7eHvOWC5/diXl1X5CDURUFnRSExy8efMYOqY4W2aQTaRtZHa1Id/sP8p5jy9lyY4ifEG9ViB1XCBk4AvqrNlTxu4iN70zE1thpEIIceqwW1Q0k2XSNcvxRrefRpTizfsfGOb7H9FJY3IPCw8v99W/D4sNDqyJ/dhCCBFH9Brr+vpkJfHx7eM4p2wbDhVcJqXOXTYNp1Xl7IMbeO/cDuRkJlJU6eO9rwt46at8Xvoqn/c3FFBS1cA5tBXJzFQTyiuqYubzq2J6shkyDG56ZR0LZo9mcJfUZhydEEKcurKSHVg0FU56gFWzHG9Cv/EN7qdjiiP2g/ur6n37j2fYGTfPzR2j6mk+aejgORr7sYUQohV5AyE++OYgzy7Zzd4SN4GQgUVVyEiyc+2Y7pxzdDu/LF7Jg4/+jPc3HOQ/Gw9S5g4AkJZg5aKhnblwaCe8/6pg2RPzuG/SlSzbWYxFVQgeC8wsmkIgZDClbyY3TcxmWLfUuKoxIMFUE5ozfx1uf91Ayr1lERVr3iNQcgDV5sSamU3K2Bknqkh5AiFufHktK39zJqpJ12ghhBD1m9gnA8Oou46kZjleRdVw9DwNRbXg3bMB776NpJ1x/YltE2wa14zpEfvBlfpnswZmapzfx8KDy/z0z4i0IESFKJciCiFEa9N1g79/up0Xl+8BqFVJNagbFJZ7efSLnTziCzJt8s38TVWYObo7M0d3r7MvXzDEbwPZLGpvx7f1CAbgq/V++D8/3XKIxTuKmDogk0dmDI2b0unxMYpTwDf7j3Kg1MPJ1/KK1e9S+sVzpIyeQZfb5tP5lhdJGnYunp2ram3n9gVZvLOoBUcshBCnDodV44qR3bBqdR9IJY+cTtqUGyhfsYADj8/kwNOzqFz/Ic7etYtSaKrCtNys2A/ubHhVwf2THTy33k9BRT2JAwntYz+2EEK0MH9Q5/qX1zDvq3yq/SHTlhQQ7t/nVy18dijIhU98Ram7bm5pIKRz9fOrWbSzGK9qrbdYhW6EJyA+23KY615cY9ocuDVIMNVEnl+6G1+w9h/T8ZK86WfdgqvvWFSbA0Wz4MoZVetpKIDbH+KZxXktOWQhhDilXDu2B2qEpR+JuWfQ8dpH6faLf9H1tvlkXvoHHF36n3jfYVG5ZkyPxj3pHHIlWOpfHpiTrnJZrpV/rjYvVIGiQLex5u8JIUScMAyDn7+1gZW7S/AE6tYFMOML6uwtcTPz+ZV4T+o9dc97m/i24CjeKPcF4SBt3d5SHvjPlpjG3lwkmGoCId3gk82H6tTSj6UkL8C6vWWUVweaYYRCCHHq65ru4v4Lc3FaY1suZ9NU+nVM4mdn9m7cgUdcD1EU/713kh2332Q7zQYjbggXoRBCiDi2aEcRX247Yhr8uLcsovDlO9n3yCUceOJqDr91H94Dm4Fw8bX8IjcvLMs/sf2RCi/vfl1gGpTVty8AT0DnjTX7KDOZ7WppkjNlwjAMvjlQzqrdJZRV+7FqKhlJdqYN6EAHk+TkCk8ABYWTL6axluS1aSrFbh8pLuk3JYQQjXH5yG64/UEe+u/2qJ50Oiwq2RkJ/Obc/nxbUE6yw0LXdBeOWAKypCzIPgN2flarqt+eO5NqbdY1RcV7j0kjSkWBkTdGfzwhhGglzyzOM13WV7H6XcpXvUO7aT/F0XMYimbBk78Oz85VJ2oEeIM685blM2dSLzRVYf6qvZitJYhmXxA+dS5Ys485k3Oa6+tGpUWCKe+OHZTNfw3vli3objeqy4ktO5v0mTNxDBkSNxU5vIEQH2w4yNOL8zhU7iUQ0k9UEnFYVP60cCtjstsxe1I2Y7LbnRi3P6SjqsBJf1s1S/JGE1ApioIvhmlOIYQQdd0wPpuczCQe/Hgr+cXh6lInr613aRDy+0lNTmJXkZsbXlqLooBuGBgGXDK8C9eN60nP9gnRHfS8R2DuOPCUxTZYqwsm/RpSusT2OSGEaGH7S6v5el/dqqPH01ranXsnrr7fLVd25YzClTOq1rbeQIglO4qY2CeDV5bvrdNCKLZ96Ty/LJ/Zk3q1aizRrMFU5aJFFD32T/z5+RiBAIS+iza8W7dR+fkXWDIyyPjprSRfeGGr/iAOV3i5/NmVHK7wmkbc3mO/7CU7ili9p5TzB3fkLz8ZhEVTSXZYCQTrLt2ItSRvSDdIdspkoRBCfF+T+mQwqU8GWwsrmLcsn/X7yqj0BrFbVJKcVvaXVmMEAhyqCNeM8p90QX9j1T4WrNnP5L4ZPHb5aQ3PVKV0hlkfwUvngrciYt+pWqwuOP1GGPezxn5NIYRoMZ9sOlSn0BrEltbi9od4e91+cjsl4wnUPU/GmiJT4QlQ4Qm26qquZsmZMgyDosefoODOn+PbuhXD660VSAGg6xgeD4F9+yj8w/0U/vZ3GCdv00KKq3xc8Pgy9pVWR6xIcpwBePwhPvzmILe/8TWGYeCwqqa9SWqW5K3esQI94MUIBfHkraXsy3l1trdZVDokN6LHiRBCCFP9Oybz0KVD+OKXk1n9u6n89ZLB5Be5qfQG8aiRL74BPdxcfdH2Ii5+ejnVJm0v6sgaALOXQrfRYHFgKBEejtkSwZkOZz8I0x4Ir1URQog4d7jSiz9UdwVVrGktRyp8VHgDWEyqr8a6L4umUuFt3XoDzRJMlTz7HCXz5oWDqCgYHg8VH39M4b33mvYJaU6GYXD1C6sodftjKrHoCYQvsk/8bxcEAlydUokjVPeXGW1JXrtF5Zox3cNNJ4UQQjS5nYcrufHltaZPQyPxBXV2Hani5lfWRXd9Su0K130Et67AmzAGHTsoKqCEC010GQnTn4Vf7YTh1zb+ywghRAsLmARSUDutJdr92C2a6SxXrPvSdQO7tXXvnZt8TZln40aKn366TiA1NW8XJaFQrejt4+xsMi3hJ4OG10vFRx+TOGEiyWf/qKmHFdGq/FL2llSfyI2qKZpmu3M/38aUB+YwKac3/8i8AEz+zhJzzyAx94x6x2EAV5s0MhNCCNE0/rRwK54Iqw/qO9/7gjrr95WxPK+EcTmRe0GVewIUlnuo9odItGeg7+xGp4t/RvJZZ4EeAk2WcQsh2q72CXZUhTrVq2NNa0lz2WiXaDMNzmLdl25AqrN1K6E2+Zm95LnnMXw+0/ee7NyFsQmRk3kNj4fiuXNbNJh6dvFu04trtJVEdF1ny10Pcul5p3Pzp9t5bml+TE89AZxWjYtO60SmLPETQohmcajcy4rdJaYFzKM531cf6wV4cjBlGAZr9pTx3NLdLN5RhE1TTxSyCCRP45xDKdxcWEFup5QW+JZCCNF8xvRqx9OLtTopMTXTWhRVw9HzNBTVgnfPBrz7Ntbqreq0aZzZPxOXzcLE3hn8b9uRWuflWPalKjAtNwub5RSamQqWllK1ZDGm83ZR8u/Zg3f7Dhx9+zR6H8VVPvKOVFHpDeK0aXRKdZpWZCqq9LEsr7jOxTWWSiIexcK8XV4uBX5+Vh/yit38b+uRqAMqh1XltG6pPPDjgbF+TSGEEFF6deUe0xK8sZzvV+WXcqjce6JFxoGyaq6dt5rCci+eQAjDOKmQhWblw60lfLpzBf07JvHCtaeTliC9pIQQbdPw7mm0S7BR7ffUeS955HTUhDTKVyyg+MOHUWxO7Fk5JI+5rNZ2hmHwk2Hh6qU3T8xmxe6SOsFZtPuyWzRumpDdxN8ydk0aTFUs/AhML1fRMwIBjr61gA6//31snzMMVueX8uyS3SzbVRyOUo3wcAIhnR7tErhlci/OHtgBuyWc1JZXVIXdotap4hRrJZH84iogXNr88ctP4w//2cyC5XkEFY1QhJ+HqoT/CKYOyOSRGUMlV0oIIZrRZ1sO1ynBC7Gd7y2qwvK8YqYP60JeURXTn1pOlTdIqJ4HiCHDwBMIsamgnHP+uZQPbhtHZpKsQhBCtD2KojBnUq/wkmmTSYOG0lo0BX48tDOJ9nD4MbJnOu0T7eHqqjHuS1WgS5qTIV1TG/VdmlKT3sH79+2LuMQP4PaCA4zauYNRO3dwW8EB841CIfx79sZ03JIqH+c/vozrXlrD/7YdwRfUqfQGqfQFqfQG8QZ0th2q5Lf//pZR//cFX+8L9wGp8ppXZ4q1kog3oJ9ITFZVhbsSD/PUrne4ZFgXHFaVRLsFl03DZdVItGvYLSrnD+7EW7PH8PgVw7BKICWEEM2qwvP9z/cB3aDcE6Ckysdlz6ygwhuoN5CqyR8yKK70cfmzKyPmbQkhRLybPqwLmcl2NDX2yROX3cLtU75rsKsoCi9cOwKXLYYm6cck2Cw8d82ImD/XHJp0Zkp3V9X7/uMN5Eyd2E91ddTHPFLp5YLHl1Fa5SfQQDU+tz8EhLjyuVW8cO0InBF+ebE22w2vkQ//Uel+P4cf/Cuj/nAfZ44byu9/PJCVeSWUVfsxDEh1WRnVs12r1sMXQogfmkjVx2M53yuAqig8u2Q35Z6A6Yr2+gpZBHWDwqMe/rV+P1eN7vG9v5MQQrQ0p03jzZtHc8HjyzhaHTAt4GbGZdN45fqRdElz1Xq9d1YSr900mqtfWIXbF6xT3OJkqgJJDguv3TiaHtE2VW9mTToloqWlNc1+UqJL1PUFw4FRSRSBVE2eQIgbX1lLSNcbrCQSjfZJ362BL335Zey9e5M4bhwAiXYLUwdkcemIrsw4vSvTcjtIICWEEC0sPUKuUizne83QSQj6eG3VPgKhutecitXvUvrFc6SMnkGX2+bT+ZYXSRp2Lp6dq05s4wnozF28u8XbgAghRFPpmOLk4zsmkpOZiMum1Zvgk2DTyEiy8+9bx3JaN/M4YWjXVBbePoGpA7KwW1TsJgUllGP/stsn8vJ1IxnYOX6K+jTpzJRjwACUhAQMt7vR+1DsdpxDhkS17cKNhRw86mlUWXNvIMTry/PprvnZHqj9Y4ilkohVU/jJaZ0BCBw5QukL8+ix4M3Gfn0hhBDNYPppndld5K6zzj+W830wGOLAQ38n1P8cUGsHZ7EUsih1+1m7t4zTe6Q3wzcVQojml5Fk5+M7JrA8r4QHP97KpoIKIFzDQFEUFAV6tk/grh/1ZWr/rAZrA3Rr5+LZq0ewv7SaX761gbV7yzAMTuRSHf/P/WXVXPbsSsb2asfdZ/ejf8fk5vuSUWrSYCrprLM4dO99pqVno2YYpM64FIBthyp48as9bNh3lCpfEIdVpWu6i2vH9GBinwyeXpxXpwIIRFfmVjfgi82H+Lm+gyedubiDtUcdbSWRQMjguSX5fL3vKJfkLWHsxRdj6y79ooQQIp5cMqIrf/vvdtP3ojnfa4rCOUO6smPgzXg2H66zj1gKWXj8IRZvL5JgSgjRZoV0gye/3MXzS3cTMozv7v0VBYNwYe9D5V7ufX8zpe4AV4zseiIlJpLiKh/XvbSG/aXVEZf7HS8ktGh7ESt3l/LUzGGc0S+zyb5XYzRpMKXabKReNoPSV+dDIFDrvc975UT4VA2KQsKECSwtDvHQ/KXsLq4iEDII1fiJ5hW5WZNfikVTcfvqJhTH8nRQs1sJTv4xtuV7cAcDJ+8qqma7AP6QzvK8EtaHetPX3o5XPAFSnLKUTwgh4kWK08o5Azvw4cZC09UMDZ3vbRaVGydk838Lt5i+H0shCwMoqopcrEkIIeKZNxDixpfXsnZvKd5A3XSZ49z+EG5/iAc+3MKG/WU8OH0waoTCFVW+IJfOXcH+0uqo8rAMwmk7t7y2jpeuG8no7HaN/TrfW5OXkUu/5hpUW+P6aOhWG4tGnc8t89expbACb0CvFUgd5/aHKPeYJ73F8nTQGzRYsrOYV64fhdMaeyWROvvTbGw9XMWFTyyj3FM3OBNCCNF6fntef1Kc1pgbeDitGhcP68ygLikRbwRqFrKIhtbAE1ohhIhHum5wy/x1rNlTfyBVkycQ4j/fFPLHDzdH3Obe9zdRUE/qTuHLd7LvkUs48MTVHH7rPrwHwvvyBnRufHkt1X7ziq0tocmDKWuHDnR5+ikUR4x9NOx2nh4xg/u36VH/cszEWtb8qCfAoC4pvHz9SBLs2ve+wPlDBoVHvVz/0mpJMBZCiDiSmeRgwezRpLqsUZf1dVo1JvfL4P5jjdUzkuym28VSyEJVICvZfD9CCBHP3ll3gJW7S0379tUX9HgCIRasOcDyXcV1PldeHWDhxsI6fV8husI+umHwn28ONuG3jE2TLvM7LmHkSLo9/xz759yCEQjU23sKqxXFoqH85j4+22gnGOGXU18xiZpiLWtu1cIX1JE901l4+wT+8dkOPtl8CAUDbzByMFTfmPwhna2FlZJgLIQQcSYnM4mP7pjAba9/zaaCckK6Yfok1GXTMAy4aWJPfj61z4m1/hcN7cznWw/j9jW+kIXNojItt0PzflEhhGhihmHw5KJdpg17o6lX4AmEmLs4j7E57Wt99u11+1FNJjOiTd2p9od4elEeM0Y0nJfVHJolmAJwjRhBr/9+QtmbCyh79VX040FVMAgWC6o9/FQudcalpM28intXlRDU99fZTzS/nJpqPh1M6De+wXHW7ETfo30Cj11xGuXVAf76yTbeXLPPNAEuqj8Yf4hnFudJMCWEEHGmY4qTf90ylryiKl78Kp9/rSvAFwyhqQpB3aB7uos5k3px4dBOuGy1L5MT+2TgsGh1gimIvnBRz/YJcVGBSgghYvH1/qMUVdadIImlXsGq/FIKyz10THGeeO2VFXtNA7RYUneOVPrYfriSfh1a/tzabMEUgKVdOzJ+eivt58ymaulS/Hl5hCorURMSsHXtSuKUKag2G1W+IO9t2MjJLZ9i+eUcF8vTwQS7xowRXevsI8Vl5UCZeSWRaMdkAEt3FlNU6Yu4LEQIIUTr6ZWRyJ8uGsSfLhqELxjCG9BJtFvqXQKoqQo3jO/JP7/YiddkJUVDhSxcNo05k3o1yfiFEKIlLVizH+/3DHoMw+DDbwq5aWL2idfMAjSILXVHUxUKj3pPvWDqOEXTSJo8GSZPNn3/sy2HTHOVYvnl1BTt00FNUZiWm2W6j28OlJu+HsuYbBaVLYUVTErKiGn8QgghWpbdomG3RJdre924nrz7dQH5xe6oqk59dwyVIV1SOW9Qx8YOUwghWs2BCCXLYwl6/CGDgqOeWq8FTp5NOSaW1B3DwHR2qyW0SDDVkINHvaaRbqzFJGpqsMytpnDV6O5YIzQR85j0r4p1TIZhUCFV/YQQ4pTitGm8ftNoLn56OYcqvKZJ0yezWRS6pbv4/fn98Yf0BhtYCiFEvDErOgGx1yvwBEIcqfRytDp8j2y3qARN7rtjSd1RFEhytE5YExfBlDcQImQS6cb6y4lFyAj/0Msj9ISyaApm8VQsY1IUBbtFLphCCHGqyUiy8+HPxvOLBRtYurMYA+oNqkI6FJZ7uWTuCkK6wbQBWdw0MZvBXVJbbtBCCPE9pLrMe6jGWq/gv5sKeXd9AVaLgoJCdYQJjFhSd/xBnT5ZSY37Yt9TXNzpJzusJ6rq1RRLqdlYhXSDf36xi5H/9zmPf7GzThnz9ATzXlmxjEnXDbKSYywRL4QQok1Idlh5/trTWXTXZG4c35MUpwVVAatJzlVIN6jyBan2h/AFdRZ+W8hlz6zk3MeWcqjc2wqjF0KI2Izt1d60L2vNoKd6xwr0gBcjFMSTt5ayL+fV2f6oJ4g/pOP2hajyBalvsXTyyOmkTbmB8hULOPD4TA48PYvK9R/i7F073WZ0drtWu+eOi5mp3M7J2DSVQKjxpWZVJZx8ZtNU3BEi3JMdX1v51KI8DpR5ePDiQSdKKs4c1Y3HPt9Rpzx6LGNKclgZ3CWlUT8TIYQQbUPHFCd3n92Pu8/ux6EKD5c8vYIjFT78EfIAAPRj6/u3H67knMeW8O9bx9GzfUILjloIIWJz8fAu/PWTbabvRVuvoDEaSt1JsGnMnpQd8f3mFhfB1JjsdiQ5rKZBULS/HAvw0W3jeG/jIZ78cle9Ue7JPIEQH3xzkI4pDu48qw/+PXuY9OmrPBIYBlrdH1E0Y3JaNW6emN0q9e6FEEK0PF8wxHUvruVwhZeA2dp1EyHdoNwT4PJnVvDxnRMjrooQQojWluK0cu6gjnyw4SAho+45rqGgpzlYVIVu7VyMyW7XosetNYZWO3INiqJw88RsHvrvdtNKHNH8cvp7DsP1V/CvkbdiYB7A1Ndo1xMI8dSiXUz55CUcyxeRMXMmZyVm8dnOUtOLYjRjunh4l3rfF0IIcep4f8NB9pS4Ta8Z9V1/dANKq/08sziP35zbvxVGLoQQ0bl9Sg7/3XwoYp5TtOo7J0bLoiq0S7Tx6g2jWnXyIi6CKYBLRnThiS934Q2EYppVAnBYVX43+8cUbMrh6NJSUOsmyEXV/NfvZ2HmIO7+7LdoiYn8udrP148t5UiFzzQCr288j10+1LSwhRBCiFPT3EV5ppVgo7n+BEIGr63axy+n9cUmhYuEEHEqOyORZ64ezk2vrMUbaLiSqZmo7skb4LJpdExx8MbNo2mf2Lr9XOPmjJ3ssPLmzaNx2bUI80rmHFaVP1yQy8ie7XilxInXJJA63mg3/axbcPUdi2pzoGgWXDmjalcC0aws8LXDcLoASHXZeOeWsXRIsZsWyIg0nj/9eCDTcjvE8C2EEEK0Zd/sP0qhSSGJaK8/EG6n8cnmQy01ZCGEaJQJvTN4+bqRJNotpgUp6hPLORFAITwDpSrhEuoOq0r/Dkn8ZfogPr5jIplJrV/oLW5mpgD6ZCXx7q3juOLZlXgCoXqnEK2agqYq/HX6YH58WmcAlu0qNp3ViqXRbiCks+NwFQM6hTsod0518tHPJvLAwi18+M1BFEWpsxRRU8GqquRkJnLP+QMY3YrrNoUQQrS8RduP4AvWvWbFcv1x+0N89G0hFw7p1BxDFEKIJjMqux3LfzOFf687wDNLdof7qirhPFBNVQjphunMVSznRABVVfjpGb2wqCopLisjuqefuEePF3EVTEE4oFp89xm89/UB5i7eTanbj64bBI41ObRoKgrhantXj+lOl7TwLFIwpEfs8RFLo11VUTjq8dd6LcVl5eFLh3DfBQP49/oDvLF6PyVuP8GQTqLdwtic9twwvmer1bcXQgjRug5X+tBNnubF2ny+pMrXxCMTQojmkeywMmtcT64d24P1+46yv7Qatz9Iot1CstPKrfPX15mAiPWc6LCoXDi0M70yEpvjKzSJuAumABLtFq4a3YOZo7qzbm8ZWw9VUukN4LRqdEp1MrlvBnZL7V+CoijhucAmaP6rRkhiS3JYuXZsT64d27MxX0sIIcQpyqS1FBD79UcqwAoh2hpFURjePY3h3dNOvFZY7jHdNtZzIkS+L48XcRlMHacoCiN6pDOiR3qD22qqgsOimVYDjKUzc0g3SHNJaVohhBDRy0xyoClwciG/WK4/ABmJcv0RQrR9qU6baa+9WM+J/pBOmiu+C7rFTQGKpjClX6bp08FYOjMn2Cz0zozfqUQhhBDx56wBWVhNqvDFdv3RuHBo55YashBCNBunTWNw55Q6r8dyTgTIyUgkNc4nOeJ6ZipWN03M5n/bjpjOTkXTaNdhVblhQk/USOs1hBBCCBP9OybTs10CWw9V1nkv2ubzVk3lzH6ZLTVkIYRoVnMm9+IXb23A7at9Xx7tOTHBpjFncq+WHHKjnFLB1JAuKXRMcbC72G36fkONdg0DLhvRtbmGJ4QQ4hQ2Z3IvfvPvb00r0TZ0/bFbVK4Z2x2LdkotGBFC/ICd2S8Tq6oCsZ8TIZzuc/bA+G81dEqdtRVF4eEZQ3BYY/9aTqvG/zu7H2kJ8T2VKIQQIj6dN6gjuZ2SsUXZl/A4TVXISnZw04TsZhqZEEK0PIum8vdG3pc7rCp/vXhQnYJz8eiUCqYAhnVL44krhsX0i3NaNa4d253rx0uVPiGEEI1j0VRevG4kvTITsZvkT5mxagoZiXYWzB5NkiO+k6yFECJWZ/bP4r7zB8R0X+6wqvy/s2XX/dQAACAASURBVPtx3uC20XPvlAumAKYOyGL+DaPokOLAZYsc0bpsGgk2jd+d159fn9O/BUcohBDiVJRot/DureOYlpuFzaJGDKqsqoLdonJ6j3Q+vmMCHVOcLTxSIYRoGVeM6s6TVw4jzWUloZ778gSbRorTyj9mDOW6cW1nguOUypmqaUSPdFb8egrL80p4ZnEeX+WVYFEVFAX8QZ1emYncOrkX5wzsiMMa/1OIQggh2gaHVePxK4ZxpMLL/FV7eWX5Xtz+IBZVJajrWFSVGad35bqxPejRPqG1hyuEEM3uzP5ZrL3nLL7cdoS5i/NYv+8oVk1BIVz+fHCXVOZM6sXU/pltLnf0lA2mIJxDNS6nPeNy2hMI6VR4AgR1gxSnVQIoIYQQzSoz2cEvzurLz6f2ocoXpMoXxGWzkGS3SNVYIcQPjqYqTB2QxdQBWfiDOkc9fgBSnNY2kRsVySkdTNVk1VTaJdpbexhCCCF+YBRFIclhlZwoIYQ4xmZRyUxytPYwmkTbmkcTQgghhBBCiDghwZQQQgghhBBCNIIEU0IIIYQQQgjRCBJMCSGEEEIIIUQjSDAlhBBCCCGEEI0gwZQQQgghhBBCNIIEU0IIIYQQQgjRCBJMCSGEEEIIIUQjSDAlhBBCCCGEEI2gGIYR+U1FKQL2ttxwhBAtoLthGBmtPYjvS85PQpyS2vz5Sc5NQpySIp6b6g2mhBBCCCGEEEKYk2V+QgghhBBCCNEIEkwJIYQQQgghRCNIMCWEEEIIIYQQjSDBlBBCCCGEEEI0ggRTQgghhBBCCNEIEkwJIYQQQgghRCNIMCWEEEIIIYQQjSDBlBBCCCGEEEI0ggRTQgghhBBCCNEIEkwJIYQQQgghRCNIMCWEEEIIIYQQjSDBlBBCCCGEEEI0ggRTQgghhBBCCNEIEkydQhRFmawoyoGW/mxLURSlm6IoVYqiaK09FiFE/drS+UhRlEWKotzYUscTQsSPtnSuagy5d2p+EkzV49gf3/F/uqIonhr/e2YzHneWoijLmmv/TUFRFENRlJxmPsYeRVGmHv/fhmHsMwwj0TCMUHMeV4h49P/Zu+/4qOv7geOvz61c9oAECIQRhkzZIihDq1VRsQXEuqpYd63a/qq2ah11tY5qW6viQgUnuBVwi2xkKLJHCARCICE7d5e7+97n98cdeEkuyV0g+/30cQ/vvvPzPXKf7/f9mZIfhaaU6hnIjyzVlr+ilHqwudIlRHsleVXt5NmpbbLUv0n7pbWOO/JeKZUNXK21/rL6dkopi9ba25RpE0K0L5IfCSFaA8mrRHsjNVMNcKRaVyl1h1IqD5gdqkQkuARCKRWllHpcKbVXKXVQKfWcUiq6AeeeqZTaopQqU0plKaWuC7HNnUqpgkDpxKVBy49LGkKc7z6l1DtKqdcC6dqklBoVtP4vSqldgXWblVK/rrb/NUHXtFkpNUIpNQfoDnwcKM26PbgEWil1kVJqTbXj/FEp9VFjXqsQLY3kR2Gl80ql1NLA+YqUUruVUufUsm0XpdQGpdRtgc/fKqUeUEotC1zn50qpjkHbTwnkecWBbQcEls9USn0ctN0OpdS8oM85SqlhgfdaKXV9YJtipdT/lFKqMb4LIZqL5FU1zifPTm2EBFMN1xlIAXoA14ax/T+AfsAwoA/QFbinAec9BJwHJAAzgSeVUiOqpatj4PhXAM8rpU6INA1KqWeUUs9EkK4pwFtAEvAR8HTQul3AeCARuB+Yq5TqEjjPhcB9wG8D1zQFOKy1vhzYC5wfqJ5+tNr5PgZOUEr1DVp2CfBGpNcqRBsg+VH9xgDbAul5FHipesCilOoFLAae1lo/FrTqEvzXlwbYgD8Htu8HvAncCqQCC/A/xNgCxxmvlDIppdID+40N7JcJxAEbgs5xHjAaOBGYAZx1jNcrREskeVVV8uzUFmit5RXGC8gGzgi8nwS4AXvQ+iuBpdX20fj/GBVQAfQOWjcW2F3LuWocq450fQDcEpQuLxAbtP4d4G/1pSGw774Ivg8N9Am8vw/4MmjdQMBZx74/ABcE3n92JP11feeBzz0D57UEPs8F7gm87wuUATGRft/ykldre0l+VOWcVfKFoOWvAA8GXcPOoHUxgX06Bz5/C/wr8L1eXO043wJ3B32+EVgUeP834J2gdSZgPzAp8DkHGAH8BngeWA30x/8w91G1f5tTq31Pf2nuvzN5yetYX5JX1TivPDu1wZf0mWq4fK21K8xtU/H/oa4NKghVQMQjqwSaptyLv+TAFDjuT0GbFGmtK4I+7wHSj2caapEX9N4B2FWgPbRS6rfAn/D/oMFfInukmUwG/tKXhngDeAL4O/6SlQ+01g6lVBqNe61CtDTtOT860ufCGvT+yGdP0OejeVQgnwB/XnTEpcBOYH6Ic1TP347sl47/mo4c16eUysFfmgv+2qlJ+B8MFwPFwET8DyiLwzyHEG1Je86rQpFnpzZAmvk1nK72uQL/HyEASqnOQesKACcwSGudFHgl6qBOmuFQSkUB7wKPA5201kn4m5UEN1VJVkrFBn3uDuQerzRESinVA3gBuAnoEEjzxqA05wC9a9m9+ndc3RdAaqDfwcX8XE3dLNcqRDNqz/nRAfxBU89qy3sRFOiE4b5Aut5Q4Q8hnIu/uRIAgWaDGfhrp+DnYGp84P1i/MHURGoGU0K0B+05r4okzfLs1IpIMHX8/AgMUkoNU0rZ8d+YAX9pJf4fxZOByB+lVFelVF1t4pVSyh78wt/mPgrIB7yBkpZfhtj3fqWUTSk1Hn8b4XkNTMPxEIv/h50fOOdMYHDQ+heBPyulRiq/PoFMBOAgkFnbgbXWHmAe8Bj+NthfBJY317UK0VK0m/xI+4f7fRd4SCnVQSllVUpdjL/JzMIIDuUBLsSfZ72mlArn/vgOcK5S6hdKKSvwf0AlsDywfjFwGhCttd4HLAHOBjoA6yNImxBtVbvJqyIkz06tiARTx4nWejv+KtMvgR1A9bkO7sDfhGSlUqo0sN0J1G4c/hKC6q+b8d/Ai/BXz35Ubb+8wLpc4HXgeq311kjToPwjuDxX91XXT2u9GX918gr8P/AhwLKg9fOAh/CXjJThb8ecElj9CHC38o9u9edaTvEGcAb+TC+4iU+k37cQbUY7zI9uBArxD+hwCH9p7rla64N17FOD1toNTAU6AS/XF1BprbcBlwH/xV+qez7+jt/uwPrtQDn+IAqtdSmQBSzTMueLEO0xrwqLPDu1Lkrr+moDhRBCCCGEEEJUJzVTQgghhBBCCNEAEkwJIYQQQgghRANIMCWEEEIIIYQQDSDBlBBCCCGEEEI0gARTx0gp9YpS6sHA+/FKqW1NdF6tlOpznI959Fqact+mopS6Uyn1YnOnQ4imIvnTse/bVCR/Eu2J5E3Hvm9Tkbypfu0imFJKZSulnEqpcqXUwcAf73GfhExrvURrXe8QkkqpK5VS1Yf/PG6UUt8qpa5urOMfq8a+/sA5Jiml9gUv01o/rLVusd+LaJ8kf2pZJH8Swk/yppZF8qaWq10EUwHnB2ZxHgGMAu6uvoFSytLkqRJCCMmfhBAtk+RNQtSjPQVTAGit9wMLCcwkHajy/b1Sagf+CeNQSp2nlPohMOHZcqXUiUf2V0oNV0qtU0qVKaXeBuxB66pE9EqpDKXUe0qpfKXUYaXU00qpAcBzwNhAaU9xYNsopdTjSqm9gRKg55RS0UHHuk0pdUAplauUuqqh16+UmqeUylNKlSilvlNKDaq2SUel1BeB61usfp5RG6VU/8C6QqXUNqXUjIamo1qaspVSf1ZKbQik623ln7UcpVSyUuqTwHdYFHjfLWjfFKXU7MD3UqSU+kApFYv/3zg98B2XK6XSlVL3KaXmBvZbqJS6qVo6flRKTW3MaxWiLpI/Sf4U2E/yJ9GiSN4keVNgP8mbQmh3wZRSKgOYDKwPWvwrYAwwUCk1HHgZuA7oAMwCPgr8YG34Z5meg3+m6XnAtFrOYwY+AfYAPYGuwFta6y3A9cAKrXWc1jopsMs/gH7AMKBPYPt7Asc6G/gzcCbQF/+s1Q21MHCMNGAd/pm+g10KPAB0BH44sj7wI/sC/6zZacBvgGeUUgNruf5ipdSpEaRrBnA20As4EbgysNwEzAZ6AN3xz2T+dNB+c4AYYFAgXU9qrSuAc4DcwHccp7XOrXa+N4GLg9I7MHCOTyO9ViGOF8mfJH8KkPxJtCiSN0neFCB5Uyha6zb/ArKBcqAY/w/0GSA6sE4Dpwdt+yzwQLX9twETgQlALqCC1i0HHgy8nwTsC7wfC+QDlhDpuRJYGvRZARVA76BlY4HdgfcvA/8IWtcvkO4+tVzvt8DVYXwvSYHjJAY+v4I/0zqyPg4wgAzgImBJtf1nAfcG7ftgmP8e1a8/G7gs6POjwHO17DsMKAq87wL4gOQQ2x39twhadh8wN/A+PvCd9wh8fgh4OfC+zmuVl7yO50vyp1q/F8mfJH+SVzO+JG+q9XuRvEnypiqv9tTO9Vda6y9rWZcT9L4HcIVS6g9By2xAOv4fz34d+AsJ2FPLMTOAPVprbxhpS8VfQrBWKXVkmQLMgffpwNowzlmnQInPQ8CFgXP6Aqs6AiWB90e/C611uVKqMHD+HsCYI1XrARb8pRvHQ17Qe0fgnCilYoAn8Ze8JAfWxweuJQMo1FoXRXoyrXWZUupT/CUn/8Rf0nJNYHVjX6sQ1Un+JPnTUZI/iRZE8ibJm46SvCm09hRM1SX4B54DPKS1fqj6RkqpiUBXpZQKyhS6A7tCHDMH6K6UsoTIFHS1zwX4q2AHaX+75OoO4P/jP6J77ZdSp0uAC/BXdWcDiUAR/szniKPnUf5Re1LwlyjlAIu11mc28NwN9X/ACcAYrXWeUmoY/mYGKpCmFKVUkta6uNp+1b/jUN4E7lVKfYe//fY3geXNda1ChCL5088kf5L8SbQckjf9TPKmdpw3tbs+U2F4AbheKTVG+cUqpc5VSsUDKwAvcLNSyhrocHdSLcdZjf+H/I/AMexKqVMC6w4C3QLtiNFa+wLnfVIplQaglOqqlDorsP07wJVKqYGB0oZ7w7gOS+CcR15W/NWzlcBh/KU5D4fYb7JS6tRA2h4AVmqtc/C3Ye6nlLo8cO1WpdRo5e8U2pji8WeWxUqpFIKuXWt9AH875meUv7OlVSk1IbD6INBBKZVYx7EX4C9J+TvwduDfAZrvWoWoj+RPkj9J/iRaIsmbJG9qt3mTBFPVaK3X4K+yfBp/ycNOAh36tNZuYGrgcyH+9qHv1XIcAzgff4fIvcC+wPYAXwObgDylVEFg2R2Bc61USpUCX+IvVUBrvRB4KrDfzsD/6/Ms/h/Skdds4DX81dz7gc3AyhD7vYH/R1cIjAQuC6ShDPgl/qrdXPxVy/8EokKdXPlHgRkfRjrr8xQQjb8EaiWwqNr6ywEPsBU4BNwaSO9W/KUnWcrfoTO9+oG11pX4//3OwH/dR5ZHdK1CNBXJnyR/kvxJtESSN0ne1J7zJlW1CasQQgghhBBCiHBIzZQQQgghhBBCNIAEU0IIIYQQQgjRABJMCSGEEEIIIUQDSDAlhBBCCCGEEA0gwZQQQgghhBBCNECdk/Z27NhR9+zZs4mSIoRoCmvXri3QWqc2dzqOleRPQrQ9bSF/krxJiLanrrypzmCqZ8+erFmzpnFSJYRoFkqpPc2dhuNB8ich2p62kD9J3iRE21NX3iTN/IQQQgghhBCiASSYEkIIIYQQQogGkGBKCCGEEEIIIRpAgikhhBBCCCGEaAAJpoQQQgghhBCiASSYEkIIIYQQQogGkGBKCCGEEEIIIRpAgikhhBBCCCGEaIA6J+1tK1wegyKHG49XkxBtITHailKquZMlhBCNSmtNmaeMksoSrCYriVGJRFuimztZQogIaa0pcXoodXqxWhTJMTbsVnNzJ0sIQRsOprTWrNh1mFnfZbFsZwFWswmlwGP46BgXxbXjM5k2qhsJdmtzJ1UIIY4rh8fBp1mf8vLGl8lz5GE1WdFovIaXEZ1GMHPwTMalj8OkpHGCEC1ZidPDvDU5vLAki8IKN1azCa39zzIT+qZy7cRMxvRKkQJiIZpRmwymftpXwnVz11Ds8OBwGwB4fcbR9QdKXDz62Tb+sWgr103I5I9n9pOMSAjR6mmteWXTKzzzwzMopXB6nQB4fd6j26zOW83Ggo3EWGN4YuITjOg0ormSK4Sohc+nefzzbby0dDcmpXB6/M8wHuPnZ5lvth1i5e7DdIi1MevyUQxMT6jzmFpr1u4p4rvt+eSXuzEp6JRg56xBnTmhc3yjXo8QbVmbC6aW7Szg6lfXHM14anNk/QtLdpN92MFTFw3DZJKASgjROmmteWDlA3y862NchqvObR1eBw6vg+u+uI5/Tvgnp3c/vYlSKYSoj+HT/P6NdSzelk+l11frdhpwuA0cbifTn1vO7CtHMyazQ43tXB6Dd9fu47nFuyiocONyG+jAOrOCZ77dSWbHOK6f1Jtzh3TBLM9CQkSkTbXx2HKglGteqz+QCub0GHyx+SAPLdjSiCkTQojG9fyG5/k4q/5AKpjLcHHHd3fwY/6PjZgyIUQk7v94E4u3HYroWcbhNrjqle/ZcbCsyvJDZS7O/c8SHvx0CzlFTpxBgRSAocHl8bH5QCl/eXcDl7+0CofbixAifG2qZurO93462qyvuorN31L6/Qd4Du/DZIvGmpZJ4rgZ2LsNwukxmLtyD78d24MeHWKbONVCCHFs8h35PL/hedw+d411RUuKKPisAPchN2a7mYSRCXSa3glzrL/zustwcc+ye/jwVx82dbKFENXsPFTGO9/n4ApRI1XXcwz4A6q/fbiRt64dC0BRhZsLnl5GflklXp+ucbzqHG6DtXuKuPj5lbxz/ViiLDLAhRDhaDPBVFZ+OZsPlIZcV7r6fUpWzafDL3+PvdcIlNmCc/danDtWHc2EfD7N7GXZ3DdlUFMmWwghjtk7294J2e+zYGEB+Qvz6XZ1N+IGxuEp8pA7J5fsx7PpdVcvTBZ/44Tc8lw2FWxiUEfJ/4RoTi8tzcbjqxlIhfMco4H1e4vJKXSQkRLD1a+toaC8ZiBVV1BW6fWx7WAZf/tgI49OH9oUlyxEq9dmmvnNXpaNEaLkxVdZQfHS10k58wZiThiHyWZHmS3E9BlD8mlXHd3O49O8syYHVwTV6kII0dy8Pi9vbH2DSqOyynLDaXDog0OkX5ZO/InxKIvClmoj48YM3AVuSpaXHN3W7XPz6uZXmzrpQoggFZVe3l+/D6NaLBXucwyAT2teXZHNptwSNuWW4DGqPheVrn6fwq9eIPHkGXS7aS5db5hN/IjJOHesOrqNy+Pjwx9yKaqoWdMthKipzQRTizblhazGrty/Fe11E9NvbL3HMCnFuj1FjZE8IYRoFNuKtmH4ahYCOXY48Hl8JIysOsKX2W4m/sR4yjeVH13m0z4W5yxu9LQKIWr3fXYhFlPNx7JInmM8hmbBhgO8tGQ3Hm/VZ6JIgjKl4O3v9x7bBQnRTrSZYKrcFbrDpOEsxRSTgDKF1/a3yOE5nskSQohGVeIqCdnEzyg3sMRZUOaa6yyJFrzlVfNMl9eF1vX3qxBCNI5ihwdNzd9gpM8xpS4vn/50AKPa7zmSoMzl8fHS0uywzidEe9dmgqlQGRCAOToBn6MUHaLkNtRRGnx+rSlxeMgpdFBQXomnej29EEI0IXOcGW+5F23UzNe8JV4scW2my6wQbYJGh3wMiew5BgynE7OnZhO9SIOywxWVuOsYml0I4ddm7qbxUVZcnsoay6O69kdZrDi2ryC2/6n1HEWRFGON6LyHyyt58/u9zF6aTanLg8Vkwqf9od15Q7pw9fjMeifSE0KIhkq0J4YsTIrpE4OyKErXlpJ4UuLR5YbLoGxDGZ2md6qyvd1il8nLhWhGSTG2kL/ByJ5jIDrahturoVrsFRyUhRNQWcwmKiq92Cy2sK9BiPaozdRM/XJQJywhJpozRcWSdOqlFH7xHI7tK/B5XGjDi3PXGoq+ebnKtobXy/BuiTWOEYrH8HHnez8x7h9f8/TXOzlc4cZjaJweg0qvD7fX34Fz6rPLOPc/S9hf7Dwu1ymEEMH6JffDpGpm5eYYM2m/SiN3bi5lG8rQXo07303OMzlYU6wkjUs6uq1JmRjfbXxTJlsIUc3onil4Q4zkF8lzjNWsmDgwHSw1y8qDg7JweA0f0TYZHl2I+rSZmqmrTu3F/LX7Qg5CkXDSVEyxyZSseJuCTx5H2aKJ6tSHhLEXHd3Gguasgs3knvcUSdOnkfjrqVg7pYU8V6XX4PIXV/PT/uI6Zyc3tMbwaLYeKGPyv5cw//qx9O0Uf+wXK4QQAVaTlUv6X8LsjbNrzDOVOjkVc6yZvLfzcB9yY4o2kTAigYzrMjBZfw7AbCYbVw66solTLoQIFhdl4YJhXZm/JofqrXPDeY4B/0BaV4/vxacbDtQ4fnBQpkxm7L2Go0wWXNk/4Nq7ocYgFHFRFuxWCaaEqE+bCaZ6p8bRv3M8P+4rCbk+btBpxA06rdb9zRYzNz94I13yJlM8bx5Z559PzMiRJF14IXETxqMCpTxaa256fT0b9hWHnFQvFENrSp0eLnp+JYtuGU9agj3yCxRCiFrMOGEGszfODrkuZWIKKRNT6ty/S2wXBncc3BhJE0JE4Hen9uLDH/ZjeH5+vrDiJZpKTIMm1vkcAzA0I4lB6YmM79eRr7ccqtEAONygzGY2cdnJPY7XZQnRprWZYArg4alDmP7sCpwRzhUVbTVz0ehuZKbGQepgoocMptMdt1O6aBGHZ80i7/77SZz6a5KmTWO1086yXQURz06ugVKnh8c/3yYT4Qkhjqu0mDSuGnwVr256FacRWZNiu9nO30/5eyOlTAgRiX6d4pk6vBur1v/AhXoRF5u/Jg4nBiYsGBzQHXjBOJd3jQmUEVNl3xibmQcu8BeKXDehNyt2Hcbhrvk8VF/hMgAKLh/rD6Y8ho+1e4r8EwAbmsQYK8O6JZEcK32phIA2FkwNSk9k1uUjufaVVbh84XWkjraaOa1/KvecN6jKclNsLEnTppE0bRqu7dspnj+f7OkX8u+x1+CI6lTjOOHMTu71aT76MZd7zh9EXFSb+uqFEM3sxmE3ku/M59OsT3EZrrD2sZvtPDL+EYalDWvk1AkhwuIo5KGK+/BYvkP7fEQp/xQG5sBoEl3VYW5Xb/MXy5u8YZzOQ97LMDATbTXz4hWjOKGzvyvB6J7JdEm0s7ugghC9H+pkM5sY36cjCsVjn21lzoo9gWNotPY3Jaw0fJwxII1rxmcyLCNJBq8R7VqbGYDiiDHmUh5bM5tOMWZi6+g4GW01EWUxMfOUnvzvkhGYQgxecYS9Xz8633knsR8t4kd7zX5UkUyEZ1KK99fta/gFCiFECEop7h17L9cNvY4ocxTR5uhat42xxJBiT+GZM57hjB5nNGEqhRC1KtkPz52Kyv4Om3YfDaSqi1GV2JWH35i/4fWof9Ajwcy868cyrnfHo9sopXj1qpOIt1uJJMwxa4OOPiejeiYz8bFveHHJbkpdXsorvZRXGlS4Dcoqvbi9PhZtzOOSF1Yxc/b3OEPUgAnRXrSpYMrncLD/1ls59ZrfsOLus3ju8pGc2qcjVrMi2uoPrmwWE50Sovi/X57A6jvP4Paz+4ddovJddgnmECPkRDIRnsNt8P763IivTQgh6qOU4uohV/PtjG/548g/0iW2C1aTlRhLDNGWaCzKzEi3wWPj/8nXF37N6M6jmzvJQggAVwm8MhnK8sCoOUdUKDHKzWjLLr7NnMPgEFOwdEuO4d0bxtEhzoY1xOTd1dktJnqlxnNq4Q7+vWgzlV5fnYNs+TQ4PQYrsg4z/bnluCLsYiFEW9Fm2ppprcm7/+/YBw4kafp0lFKM75vK+L6pVFR6Kaxw4zF8JERb6RAbei6H+hQ53Li9NTOLSCfCK3KEl1EKIURDxNniuHjAxfym/28oqiyitLIUi8lCsj2Z2OdPB1M8hJlfCSGawJInoTQX9M/PGD2fKsPhgd23xBFr8z+zvLjOzdwNHr69MhYAs+GCXV/Dzi+h75k1DtsnLY7Pbp3Ac4t38ebqHLTWVFSrRYq1mbFbzcw8tSdpcVHcU+wKe4AtgEqvj12Hyvn9G+t46QopoBHtT5sJpkreew/npo30euedGoFSbJSF2OPQR0mHnpw84onwhBCiKSilSLGnkGIPGs2v31mwfSF0G9l8CRNC/MzrhjUvhayRMjT8e5WbO8dH1b6/xwHL/h0ymALoEBfFXecO5M9nncCijXks3JhHYYUbs1Kkxkfxq+HpTOyXhgLGPPxVrYFUXYNsubw+lu0sYHNuKQND1JIJ0Za1iWDKtW0bhx5/gh5z52CKial/hwZKjrFhs5hweapmNJHOTp4UY22sJAohRN36nQMLb4PT727ulLQ5hs/ApEzSGV9EZuvHoEMHMLeNs/HoskpuHG0jyV7H31XOaijaA8m1D2ceZTFzwbCuXDCsa8j1i7fn43CH7qcVziBbHkPz0tIsnpghA9qI9qXV95kyysvZf8utdPrrX4jq3btRzzW+b0d0iKqpSGYnj7aaOf/E9EZNpxBC1KrbaH9H95L9zZ2SVs+nfSzfv5xrPr+GkXNHMnzOcIa+NpSxb4zlvuX3sat4V3MnUbQGP7wJ7vKQq0alm5nU08LjyyvrP862hceUjOcW76rRBBDCH2TL8Gk+2XCAUpfnmNIhRGvTomumPIaPvYUOSpwebGYTHeOi6Jz484S3WmsO/O1vxJx0EolTpjR6ejJSYhjePYmVWYU11oU7EZ5Pa6aN7NboaRVCiJDMFpyZZ/PJ54vZYh9KkcNNvN1KZsdYpgzrSorMHROWL/d8yUOrHsLhceDwOqqsK/eU88HOD/gk6xN6J/Xm3kunYAAAIABJREFUkfGPkJmY2UwpFS1e+aE6V//9tChOebmCW8bU8ds0KqGi4JiS8WNOccjlkQyyZTOb2JZXxuiedU8ULkRb0iKDqbwSF3NWZjNnxR68Po1ZKTT+4KpXx1humNSbswd3puLtt3Bn76HnW282Wdqum9ibDftKGjQRnknB5CFdSIyWZn5CiKa397CDF5ZkMf/Hc1E+Lw5f9tF1dquJRxZu5RcD0rh+Ym9O7JbUfAlt4V7Z9ApPr3+aSqP22gJDGxiGwZbDW7j4k4uZdeYsmc9LhFZLE78jBqeZOa+fhX8sdTMgtY4GRb7QTfTC5axlNL5IB9kqdUrNlGhfWlQzP8On+dsHG5n42De8EJjbwBGY06C80kul18fWvDLufO8nRt7/GYteX0i3p57EFFVHx8zjbGLfVEb2SCbKEvlXF2+3cttZJzRCqoQQom7fbD3EWU99x5ur9+L0Khy+qoU6Lo9/GORFG/OYMWsFLy3NaqaUtmwf7fqI/63/X52BVDCNxuF1cP0X15NVIt+pCCEmud5N7p9k54V1bvaX1jIDr9kGsR2OKRmWWubbDB5kq14KbA14PhKiNWsxf/GGT3P1q98zf+0+Kr0+3HUMy1nhNij3aO4fdjELC5u2cs1kUjx/uX+WcXuYGYby+YizmnjjmjGkJ9U+kaYQQjSGb7Yd4obX1+L0GHh9tTyMBfi0P7B6/LPtzFosfX6ClbnLeGDFA7gMV411RUuK2HH3DjZdu4mtN28l99VcjIqfHz4dXgd3Lb2rKZMrWosBF4C17sGz+qSYuGiQlf+srmVqFWWGzEnHlIzaWs0ED7JVH8On6RjXdAXcQrQELaaZ390f/MTKrMM4PRHMbeCDO97bQKdEOydnHluJTCSibf7Zxm+bt4FFm/JAg9uomW6lwG41k2bSPLD+NfonTGiyNAohBMC+Igc3zl1XYxRSqHuoY6fH4MkvtzOkayLj+nRshpS3PB/s/CDk8oKFBeQvzKfb1d2IGxiHp8hD7pxcsh/PptddvTBZTGg0O4p2kFWcRWaS9J8SQYZeBJ/XP7rmPROjmLOhliZ0KZnQeUiVRVprVuw6zKzvsli7pwinx8CsFAnRFqaN6MZvx/Wka1AB74UjM3hpaRZuo2qBS/AgW8pkxt5rOMpkwZX9A669G6oMQpEcY6N/5/gILl6I1q9FBFNZ+eW8t25/yJm267rZg78E9e4PNvLlnyY2aZqjLGb+c/FwcoudzF25hzkr9+A1fJhNJnxa4zU0k/qnct2ETEZ0T+bA31aSd999pD/2qAybK4RoMrOXZeP11cxbwxnq2OXx8dRXOySYwv9g+sqmV2rUShlOg0MfHKLr77oSf6L/IdKWaiPjxgy237adkuUlJE/wN+Py+rzM2TKHe8fe2+TpFy1YVDwMmeYf1S9o0t7sW6sGJRmJJlx315zDyaetMPqGKk2NFm08wH0fb6bU6anSx9tAU1DuZvaybF5Zns2onsk8Nn0o6UnRXD62By8v202oGTXDGWQr2mrm2gmZ8owj2p0WEUzNXpaNEaLpSTg3e4D9RU427Ctulg7T6UnR3H52f/50Zj8Kyt2UujxEW810iLMRY/v56+18111kz7iI4vnzSb7wwiZPpxCi/XF5DN5avRdPtZLmI0Mdd5h8KzEnjDu6PKbPGGL6jKmy7Y85xew97KB7h8abw6812Fm8kzJ3WY3ljh0OfB4fCSOrPuSa7WbiT4ynfFP50WDK0AaLdi+SYErUNOmvsOUTcIUeUa822mTF601k792vkf6PwcQMH86sxbt48svtIWujjzjSmmZlViGT/7OEt68dywmd4zmpVwrLdx0O+UxW3yBbAFNHhJ7DSoi2rNn7TDndBvPX7qvRjj/ceQ0AKr0GLy7Z3ZTJrsFiNtE50U6/TvFkpMRUCaQATNHRdP33U+T/60lc27Y1UyqFEK2G4YEQNUqR+GxTXsjlkQx17NOaOSuzjykdbUGhqxCzqjmamVFuYImzoMw1S+MtiRa85VVHWHN4HOhQExaK9i2xG/z2Q38tFWHW7JhtqMSu2P7yPWl/uo19f7iZFx9+iafqCaSCGT5NscPDRc+v4ECJk8emDyXRbg03BUfZrSaevGgY8XYZrVi0P81eM7VhX3HIEWQiu9n7Z+5u6aIyM+n017+w/9Y/0nPePMxxsc2dJCFES+EzYPtnsOwp2L/u52GObXH+JkAn3wipkY0GuvNgechJOCMZ6thjaDbllkZ03rbIqGUkM3OcGW+5F23oGgGVt8SLJa7qbVYH/lMRP66KNi99GFzzDcz5NTiLap3IF2X2j97XbRT85nWwJ5Jw1i9x9R/Mo8+swa0i7zJR5vJyx/wNvPa7Mbxz/cnMeG4lJU43Rhhxv91q4v4pgzl7cOdjuXohWq1mr5kqdnpCtM6NfF4Dh/vY5ldoKolTphA9cgR5990npZNCCL8f3oTH+8B710DOKvB58Pdb0OAug/VzYdYEeP40OBz+CHvFtcz3EtFQx/gftNq7hKgEdIi7VUyfGJRFUbq2asBpuAzKNpQRO7BqoZnNbMOkmv3WK1qqjn3hlg1w0RzIPA3MUf7aqqgEsMWDxQ5DfwPXfAVXfgL2xKO7zstyYLLWrBkqXf0+hV+9QOLJM+h201y63jCb+BGTce5YdXQbw6dZtbuQAyVO+qTFs/DW8Zw+oBNRFlPIqWBMCqIMD73jzbx0xWguGp3RON+HEK1As9dMmVTo8rngm304AVVr6vAo/aeEEEd9/SCseBo8ztq38Xn9rwM/wPMT/c2Buo6s99AJYQx1HNv/1HqPExfV7LeKZtc3uW/IAjBzjJm0X6WROzcXk91UZTQ/a4qVpHFV+/KOSBvRVEkWrZXJBL1P978qCqA0FzwOf0CV1B2i4mrsYvg0Ly/djavaQF6R9I/UwJwVe7j97P50SrDzwm9HkV9WyRur9jJvbQ4lDg9enyY2ysy43h25yLGd7ms+J6PP2Y3yNQjRWjT7HTIl1hayZirSm318K7rZH+k/tefSy4g+8UTsJ8hEvkK0S6tmwYr/1R1IBdM+qCyD134F1y32D4dcC4fbS0WlF4tJ1eiTGslQx2YT9O1U8+GtvYkyRzG171Te3vY2Hl/VGr/UyamYY83kvZ2H+5AbU7SJhBEJZFyXgcn6c6l+jCWGmYNnNnXSRWsW29H/qseGfcW4vDVrmiPpMuH2+nh33T5uP7v/0WWp8VHcckZfbjmjb43tfa4B7Hz5v2xes5nXcwzW7S2i3OXFbjXTPSWGK8b15NQ+HTHVMhmwEG1Fs0cgQ7slYg7xQ4vkZm81Kc49sUtTJvuYSf8pIdq58kPwxd/AW1llcc+nynB4YPctccTa/Hnji+vczN3g4dsrA/mEuxw+vAlmLqiyb06hg6+3HuKrrYdYt6eIQekJ1FZpH85QxwBWk4nLT+5xfK65lbtkwCXM2z4v5LqUiSmkTEypc/84Wxwndzm5MZIm2rmCcjemED/2SLtMlNTSNDiUZTllPHz6n8iavxOvyUzwdJs7DpWzMuswMTYL103MZOYpvUI+6wnRFjR7MGUxm/jt2B48/11WjXmmwr3Zm0yKmaf0aspkHxeJU6ZQsXq1zD8lRHu05mVqG7XL0PDvVW7uHB8Vel/tg/1r8BTsZm1pAl9vPcTXWw9R7HAz6YQ0Lh6dwf8uGU683cpf3/uJd9bsJcS84mENdRxlNWO3hvcg1tZlxGdwXuZ5fJr1aY35pupjN9u586Q7JZ8XjcJr+AjVDTvSLhOhhkQP5cUlWTz++TZcnkAeFSJ/qXAbVLgNnvh8G0t2FDDr8pGSl4g2qUX0gq2r1DNu0Gl0ueIpuv/pXTJumkvahfdh7zbg6HqlYFB6Ar06ts6anc533UXltm0Uz5/f3EkRQjQVw+tv4ucN/UB+2zgbjy+vpNhV+4ON12vw5tN38/CCLURbzTxx4VBW33kGj184lHOGdDk6RPHV43thNTcsq4+2mpnYN5Xzn17Kg59sptjhbtBx2pK7T76bkZ1GYjfbw97HbrZz8/Cb+UWPXzRiykR7lhBtDVkLHdxlIhzVp3UJZe6KbJ74PPzh150eHyuzDnPdnLVhB2tCtCYtIphKS7DzxzP7Ed2AEosYm5lHp5/YCKlqGjL/lBDtUO46MGoPTEalm5nU08Ljyytr3caCl0tjVvHRTafyxzP7MTQjKWTfhN6pcTz86yHYrZFl99FWM5eO6c5/LhnO57dOoMJtcPoTi3lxSRaVIfpmtBcWk4Wnf/E052aei81sw2qqfV6daEs0UeYo7h17L5cPurwJUynam8FdE3F7awY3wV0mHNtX4PO40IYX5641FH3zcpVtlfYxKG8bBx95hPJly/C5a+ZRu/LLeXDBFpyemnlAxeZvOfDqrez913T2PX05B9+5F9e+TQBUen2s3l3Iayuyj8v1CtGSNHszvyOum5DJ4XI3c1fuCfkjrU7hD6RenXkSfdLiGz+BjUj6TwnRzlQUQD3DY//9tChOebmCW8bYat3G7C4L63RTR3TD8Gn+9uFGKj2+kIP+BIu2mrns5O789Rx/K4C0BDuPTB3CzFN68o+FW3l1RTZ3nN2fc4d0CbvZ2p7DFby1Ooes/HLK3QaJ0RaGdkviwlEZpMTWfo0tkcVk4b5x9/G7Ib/jzS1v8u6Od1GGB4VGWaLxai+JtkSuHHwlU3pPId7Wuu9RouVLjLYyeUgXPvohF6Nae79wu0xER1n5/YxxmLeupuC/T1O5cycxJ51E3IQJxE0YjzU9nZeX7sYbYvKp0tXvU7JqPh1++XvsvUagzBacu9fi3LHq6FxWTo/Bc4t3ceW4ntLcVbQpqq65jkaNGqXXrFnThMmBuSv38M9FW/FpTUVlzaDKrBQ2i4mMlGievmQE/Tq1nZtU7t13o12V7ab/lNfnpbiyGIfHQYw1hsSoxDpLecXxoZRaq7Ue1dzpOFbNkT8dN5s/gg9/D5U1J8Pt+VQZL06J5oxMC5e+56BzrIkBqaaqA1AcYbLCPQVhn3bDvmL+89UOluzw7xPcT9VqUphMikHpCfzh9L6c1j+t1uMs31nAQwu2YDWbuOvcAYzuWfvAC99uO8R/v97Jxv0l+LTGE/QgZrea8Gk4Y0AaN53Wl4HpCWFfS0tSaVSyc+GfKImKwTL4QlLsKfRO6t0u8vHjrS3kT82VN23KLWH6syvCKpAOJSM5mu9uP+3o3623qIiKpcso/+47KpYswZPWhWkDZ+Kq1qjJV1nBvv9dQYfJt9Y7+nKszczzvx3FKX3qH6FQiJakrrypxdRMHXHZyT2YMSqDzzbl8ey3u9iSV4rFpPD5NCafweRhGVwzPpPBXRPrP1gr017mn8opyzlammtoA7MyY2gDheL83udz2cDLyEysfchnIVq96OSwNrt/kp0Rs8r5v7G1DERhi6wW+8RuSbx4xWgOlbl4+/scftpXQonTQ2yUhT6pcVw8pntY/U/H9enIxzedygc/7OeWN9czpFsid5zdn8zUn4dQ9/k0jyzcwtyVe2t9uDvS52LRxjy+3nqIR6edyJRhXSO6ppYgyhzFoMpKyJgAXU5q7uSIdmpQeiKn9unAkh0FNeabqo/dauL+CwZVKQCwJCeTeP55JJ5/HtowePeTVZhWFlC9ajuS4dcr3AavrciWYEq0KS0umAKwWUycPzSd84em4/Npyiq96OwsCv70R/r8c0H9B2il6pt/SmtNpdeH1WxqlUOMVngquP2721l5YCVa6xpztQC8v+N9Ptr1EYM7DubJSU+SbA/voVOIViV9OBj1D0HcJ8XERYOs/Ge1myFp1ZoFKhP0mtig06fF2/nD6TXnjYmEyaSYOqIbk4d04eVlu5n27HKmDE3n5l/0pUNcFA8t2MIbq2oPpIL5tD+wuv3dDZhNinNPTD+mtDWLsjyI79TcqRDt3H8vGcFFs1awLa8s7IDKbjXxl3P6c3r/2v9+ldlMQVInKlVRjWAq0uHX9xY6wtpOiNaiRQZTwUwmRWK0FV/vXuTt3492u1G21tW+PhLV+0957XYW/pTHs4t3seNgGQqFD01clIXpI7ox85RedO8Q09zJrlexq5jLFlzGgYoDuH21d7z3ai9ew8uG/A1M/3g6r09+nc6xnZswpUI0gag4GDIdfngTtLfOTe+ZGMWcDSECL4sdTrm5kRIYPrvVzI2T+nDRqAz+89UOzvjXYiadkMqijQdr7aRe+v0HeA7vw2SLxpqWSeK4Gdi7DcLl8fHneRsY3DWRHh1aWd/R8oMQ17LzKq/h48sth5i9bDc5RQ5cHh8xNjOD0hO4ZnwmI3skS9PEVs5uNfP2dWO5+c31fLcjH4+hax1BL9pqQkPYNcJOtzfkFAuRDr8e7iiAQrQWLT6YOsJks2Ht0gX33r1E9enT3MlpVIlTplC+ajWP3/8Sr8X0A+2vGvfzZ4plLi9zV+3hjdV7GdY9if/8ZjidEsIfqrcpuQ0313xxDbkVuSFro0Lx+Dwcdh5m5qKZzDt/HnG2uPp3EqI1Gft7+Gk+eKsGU9m3Vu0HmpFownV3iL5E8V2g68jGTGFEOsRFcf8Fg7liXE8ueHpZyEAqnE7qXsPH7GXZ3DdlUFNfQsQKK9zsPFROmcuDvSiJ9Mp4WuKMh17Dx7Pf7uKlpbvx+HxV+iMXVsD+YidLdhTQIdbGbWed0CqbWoqf2a3+fklbDpTy0tLdfPxjLjbLzzXbhk8Tb7dwzfhMLhyZQWJMeH2V4+1WrGZVpd8jVB1+vb4+U/7jtJpHTyHC0qr+om2ZmVRmZbX5YEprzb8G/5pP1udQGWIQjiP8GZpmTXYR5zy1hHk3jKV3assLOj7e9TF7SvfUCKSKlhRR8FkB7kNuzHYzCSMT6DS9E+ZYf8mWoQ3ynfm8vuV1rht6XXMkXYjGkzYABpyPd9NHWHyRTQCLJRrOfYKQE8s0M7fhw+OrWfLsq6ygeOnrdJh8KzEnjDu6PKbPGGL6jDn62ePTvLMmh7+c079FTvCptWbd3iJmfZfFt9vyibKYQGuovAbP8xvonrKTGyb15pzBXVpE+p1ug9+9+j3r9xbhrKVGQGtwuA0cbid3vPsTG/aXcNfkAVJL1coN6JLA4xcO5d7zB7I5t5QSpwerxURqXBQDuySEnEqhLgPTE7CaTXiMqs8lwcOvK5MZe6/hKJMFV/YPuPZuIPm0q45uazUpRvaQ5vuibWlVwVRU70zcWVnNnYxG98iCrSzYnE9lmCPbGT5NkdPNjFkrWHjLeNLiW04Nldaalza+hNPrrLK8YGEB+Qvz6XZ1N+IGxuEp8pA7J5fsx7PpdVcvTIFStEqjkrlb5nL1kKsxh9keW4jW4sUOf2YEmxlm2YWp2m+kVpZoOPdx6H1a4yaugeas2BNyvptIOqkr4LNNeVzQwmpIih1urpi9mu0Hy3F5DLQm6FpjweNj+8Fy7np/I/d+tInZV45mZI/aRzpsbIZPc+2cNazdU1Rl5Ma6OD0Gr6/cS6zNwh/P7NfIKRRNId5uZUxmh2M+ztjMDiTYrTjcNQt5wx1+3WRSXDmu5zGnRYiWpFUFU7ZemVSsWtncyWhU2/LKeG1ldsg2xXX1NdAaShweHvpkC/++eHgzpDy0H/N/pMBZdehmw2lw6INDdP1dV+JP9DdpsqXayLgxg+23badkeQnJE34uuXIbbpbsX8KkjElNmXQhGtX/vtnJ/LUHOPemhZgW3wGb3gOfAbU1hbUG+hBNewH6n9t0CY3QtrwyQnXRiKSTusNtkF1Q0Qipa7jD5ZWc/9+l5JdX1mjmVN2Rh83LXlzF878dxfi+qU2RxBrmrtzDmuzQgVRd9xOnx+D577I4rX8awzKSmiHloiVSSnHthEwe+2xbyGa8cYNOI25Q3YU8mamxra8/pBD1qHvWyBYmqncm7l1tu2bqpaVZeELc+EpXv0/hVy+QePIMut00l643zCZ+xGScO1Yd3cbr0yzalEeJI7x+SU1hcc5iXN6qTZgcOxz4PD4SRlbtB2K2m4k/MZ7yTeVVt/c6+Dz780ZPqxBNQWvNU19u5711+3j72pPpkpIAv34WblwJJ10DtjiwxUNUgv9ljYaUTDjnH3DbjhYdSAFUuEMPqBHcSb0+Gihxtpx8zGP4uOTFVeSX1R9IBXN6fFw3Zy3bD4Y3ufLxpLXmucW7au27Vt/9pNJr8Px3u5oyyaIVmD6qG9E2Mw1pAGo1K3KLnTyycAuuBs6FJURL1LpqpjIzce/ejda6TbblLq/08tGPuVS/V4fb1wDApBTz1uZw9fiWMU/TIechdLVxVI1yA0ucBWWu+W9oSbTg3FOzudNh1+FGS6MQx0prjdNj4HQbxNktRFlC175orXnss218teUQb107ltT4oPmjUnrB2Y/AL+6Fgm3gLAazDeLS/MFUK8nz4qJC31Yi6aSugKSYljNq62eb8sgpdOAJUeVWVw0P+Pss/XPRVl66YnSTpnnV7sKQAWm49xOfhq+2HKKowk1ybMv5txDNK8Fu5Y1rxjDt2eU4Ko3qo6TXym41cffkAZw1uAv3f7yJc/69hEemDuHk49D8UIjm1qqCKXNCAio2Bm9eHtYuXZrsvCUOD/PW5vDJhgMUO9yYlCI51savh6fz6+HdiK3l4SFSS3fkYwnRITSSvgZOj8G8tftaTDAVijnOjLfcizZ0jYDKW+LFElfz+9Q6/NJgIZrKnsMVzF6Wzbw1Obi8Pswmhdfw0TEuiqvH92LGqIyjQYHWmocXbGHZzsO8ee3JpNT2gGq1Q5ehTXgVx9fALgms31uMt1rgEUkn9ZgoM5mpLacp0LPf7grZTySc0Qk1sGRHAYfKXE3an/WNVXtxhkhzJPcTk4IFGw9w6ZgeVZZ7DR9Oj0GszRLxIAai9evfOYF3bxjHJS+swun21jqwCfhro8wmxYMXDGb6qAwAnr5kBF9sPsitb/3A6QPS+Ms5/Umwh9dHXIiWqFUFUwBRmb2pzMpqkmBqX5GDxxZtY9GmPExKVW0uUVDBlgOlPPjpFn41rCv/98sTqpYyN8DhCneNBxCIfEK8Ykft8zg1tdToVBSqSu1UTJ8YlEVRuraUxJMSjy43XAZlG8roNL3mxIEdo2W2dNFyHC6v5A9vrmftniJ8Wh9t+nVkPpdDZZX864vtPPH5di4c2Y17zh/Iwwu2sm5vEW9cM6ZF1bocb78d15O3vs8JmZeF20ldoThzYMuYAHf7wTJ25ZfXWB5JiwEFvLFyL7c24YAO+4ocIWsNIrmfOD0+8kr8zbQPlbqYs3IPc1fuodjpwawUhtZ0SbBz9fhMpo3sRmK0PBC3F/07J7D4tkm8u3Yfzy/JotjhwfBp3F4fFrPCZjGhNVw0OoOZ42rOh3nmwE6MyUzhkQVbOevJ73jggsGcUcdvXmvN2j1FPP9dFuv2FuFwG1hMipRYG5eO6cGMUeEP8S7E8dbqgilbZi9/v6lTTmnU8/y0r4RLX1xJeaU3ZGdq+LmT8fy1+/hyy0HeunYsfdIaPjS5z6drzCwOkU+IV9sEfc1hUsYk3tj6RpXR/MwxZtJ+lUbu3FxMdlOV0fysKVaSxlXt8BxjieHMHmc2ddKFCGlfkYOpzyynqMIdstnXEUcGkZm/bh+fbT5I1yQ7c68e0+ZLYHunxtG/czw/7isJub6+TupWs+LSMd1rbSrZ1FbtLiRUxXgkNTyVXh9fbT0UUTCltWbV7kJeWrqbbXllONxeoq1menaM5apTezGxb2qdtUK1TYwa6f2k2OHh2jlrWLwt/+i1AHgDX0puiYvHPtvGPxdtZcaoDO45fyBWc6vqji0aKN5u5cpTenHFuJ6s3l3I1rwy/5xrVjPpSdGc3j+tzukBEuxWHpk6hBW7DvPX9zbwwQ/7uW/KIDrGVS2Y/mxTHg9+spnDFW6cgVE0jyh1efnXF9t4/PNtnD2oM3+/YLAEVaLJtbpgKiqzN5W7djbqOXbll3PxC/5AKhxen+ZwhZsLn1vOglvG0yUxOuJzegsLse3ajsnrAapmPpFPiNdyMpKhqUPpGN2RnLKcKstTJ6dijjWT93Ye7kNuTNEmEkYkkHFdBiZr1RuxzWxjQrcJTZlsIUIqcXq4aNYKDpe7McJseury+HB7KxmUnkCcrdVluQ3yxzP7ccPcdSEHP6iP4dP8YkBaI6SqYUocbjxGzcAk0hYDpa7wB9SYvzaHJ7/YQZHDjdNdtV9KTpGTdXuKiLaZuXFSb64c1ytkUFVbLVGkfdc+/jEXh8cIOdz9EUf+neevzWFrXilzfjemRcyxJZqGUooxmR0aPPz62N4dWHjLBJ76cjtnP/Udd04ewK+Hd0UpxbPf7uTfX+2otXAAONrMcMHGA6zZU8i868eRnhT5c5gQDdXq7uy2zF6UffFFox1fa83M2d+HHJGqvqHJS11ebpi7lg9+X/cNSvt8uHftwrF+Pc5163GuX4+3sJC+w8fgTfllje0jmhDPrFrUg4hSipmDZvLomkdrjOqXMjGFlIl1z8FiM9u4dMClMseUaBH+9/VO8ssqQwZSdeUPPg2rdxeyeHs+p/VvOb/PxjLphDSum5jJrMVZEQVUdquJaSO6ccPcdTzwq8FMHtJ0fWNrYzGbMCmFr9q/eaQ1PKH6w1bn82n+9uFG3lu3v87vrcJtUOE2eOyz7azaXch/Lx6BzVK1EOrkzA6s21tzWPRI7ifgDwLDbezg9PjYsK+E6+eu5eUrRkt/KhG2aJuZv04ewLknduH2+Rv48IdcRvdM5n/f7KozkArmMTR5JZVc+NwKFtwyXpqdiibT6oKpqN7+PlONZdXuQgrKK2s06wino7Hh02w9UMb2g2X06xR/dF9fRQXOn37CuX69P4D64UfMyUnEDBtO9PDhpFw1k6g+fVAmE+NeXs232/NrpCvsCfFUy5sQb0qfKby97W2ySrLw1DaHTghmZSYtOo3LBlzWiKkTIjyVXoN3LPzmAAAgAElEQVQ3Vu/FHWJo7HDyB4fbYNZ3u9pFMAVwyy/6YjWbePrrnf4JbuvY1mLy97H43yUjOK1/GjNGZXDTm+tYlXWYO88d0KxN/jrE2rBZTHirDeYQaYuB5DD6yT20YEu9gVQwp8dg8fZ8/vj2Dzx9yfCjo9yWV3oxKWqdqDfc+4mGkE0c6yo4qPT6WL27kM8353H24OYPhkXrcmK3JD7+w6k8/tlWHv98e8ht6vr7M7TmUJmLv3+8iSdmDGvi1Iv2qtUFU5ZOndAOB0ZJCebExPp3iNDzi7NqjIAUSUdjj8/HC59t5O6UQpzr/bVOldnZ2Pv3J3r4MJJnzCD94YexdAw9oMK1EzJZnV0YcuSocCbE69spjm7JMXVu09SizFG8+MsXuXTBpeRV5OH21T9AhkVZSLYnM/vs2cTZGt4PTYjjZeFPeSFHlYwkf1i/t5icQgcZKS3rN9oYlFL8/rQ+nJzZgWe+3cnSHf7Ju4Mf8GNsZrSGqSO6cs34THp29I/gNzQjiU/+MJ7b5//I9GdX8PQlw5ttos8zBnTi7g821lgeSQ2PWcFP+4q59MWVnN6/E2cMSKtxPct3FfhH4AsRSNX18Ojy+Ph62yE+/CGX4d2TeHX5Ht5bv4+xmR0Y2SOZdXuKQgay4dxPQgm34OC5xVkSTIkGsZpNmE0mrGZVY163cP7+PIbmkw0HuHfKoDbfR1W0DK0umFJKYcvMpDIri5jhw4/rsctcHpbuLKhx44mko7Hhgw9/OsgfXF8RO3wECeedi33QIEy28EbvGtu7A8MzklizJ/Ss9XWxmRU5hx38+8sd3Hha7xbVCTjJnsTb573Nbd/dxve5K/BpH54Qt3izMmMxWRjUYRBPnvYkKfa6mwEK0VTeX7+fimMcahrgyy0HmXlKr+OdvBZrZI9kXrpiNIfKXHywbj+7Dzsoc3lIjrExuGsC5w9NJyZEX7LEaCvPXTaSV5ZnM/WZ5c3W7C851sYZAzqxcOOBGs3dwq3hsdvMLL39dFZnF/LVloM8++0uEqMtnDGgE78Y0IkR3ZN47tvaJ9it7+HR6Ta464OfsJlNzBidwac3j6drUjSbckuY9uzysJtJ1SeSgoOtB0rZlV9O71QpDBOR8Rg+5q7cUyOQinTOzffW7uPKdpTXiubT6oIpCIzol7X7uAdT+WWVWM2K6s9LkXY01jYbiX//V4Pa6yqleOGKUUx7dgVZ+eVhB1R2q4lnLxtJ/87x3D5/A9OePci/ZgylT1p8/Ts3kThbHM9OfJKc/57I6yN+xfv7F2NoA7My49P+6zwv8zwuG3gZvZN6N3NqhaiqoLwy5PJI8odKr4/C8pYzdUFTSou3c+3EyH7XSilmntKLEd2Tm7XZ3+n901jw04GQ6+qr4bFZTFx8UneSY22cNagzZw3qjM+n2bC/hK+3HOS+jzaxr8gRcsCjSB4e3V4fr101hpE9ko8uG5SeyFMXDefWt9eHHVAp5X8QDTUqbCQFB16f5rNNedw4qU9Y5xXiiFVZhTX6KELkc26+vmqvBFOiSbScqosI+Oea2nXcj+v0GEfbnAcL7mgcDrNJ4WrAKFZHxNgsvH/jOMb37YjdYqKuCqbYKDOJ0Vbm/m4Mp52QRpfEaF676iRmjMpgxqyVvLR0t3/I9ZbixzfJ6DSUv0x6lKUXL2Xh1IW8dd5bfDr1U5ZfvJx7x90rgZRokULd3CHy/CHcUQDFz440+8srdTH92RXsOVzRJOf1Gj7++9UOHl6whVP7diQ6whHqzCZFeqKdW37Rt8pyk0kxLCOJP/3yBBbcMp7Lx/YIee+J5OHR54MP1u+vsfzswZ159rKRRFvN2C113/JjbWY6xNoY0Dl0IVwkBQden+Zgqave7YSo7lCZK2RfvUgLtg9XtM+CK9H0WmUwZeud6Z9r6jhLsFtDPjAFdzQOh8fQxNuPrdLPbjXz4hWj+eTmU5kxqjt2q4m4KAvxdgvxURaiLCYGpSfwz2kn8v1dZzCq58/N4ZRSXHZyD96/cRwLfzrAJS+uJKfQcUzpOS58Biz/L5xyCwBWk5XUmFR6JfYiLSYNq1naNouWKyU2dFPdSPIHq1mFNRCBqOlIs7+pI7oy9ZnltdYU1eAqhaJsKMwCR2HoERVC2JVfzrTnVrA6u5BPbj6VV2ee9P/snXd4VGX2xz/33mmZ9E5CEpIQeu+9KirYEBV1FQtIU1fZXcuuu2tb/bnFtmsXERXBBqiIXZQaOtI7IUAIgfQ6fe7vj6GFTCYzIZNMwvt5njw+zr0zcyZM7n3Pe77ne7iqayuvEyqdIpMQbuDT6YPqHFdxZuDphfiyeHSoKseK3V/nR3WIY/Vjo3jwsnZEB+sI1isYtQp6jYxRpxCkVWgbG8xT13Vh9WOja638+bxx4MasRSCoC5vD6dZB0tfvn93NSAOBwB80S5mfPj0dy+GGT6biwwxudwd9tZKNCtb5vINZGxlxoTw/oRt/v6YThwsqKTPZMWhl4sIMtK5jjkKb6GA+nT6I2auyuP71NTx2VQcm9k12+xkbhb3fQFAEtPHvwGWBwB+M7dqK346W1DCH8cmIQJYY3j62sUNvMXgt+3PYYf93sPoVOLEVFJ1Lv+awQngyDJkF3W4GXU0jEKdT5YO12fxv2QH+OKY9dww8VzV6aWIPOiWE8uovB3GqKpWWmos6g1ZGVeGyTvE8P6GbV3LvC02PzuCr/bo746IzRIfouW9UBtNHtGV9ViE5JSZMVgcheg3t40PplnTO0CnSi40Db2ZUXTh8VSDwhjCDFsWNrb6vDprB+ma5xBU0Q5rlN02XkoL9RB5OiwVZ33AXa63TzoQoKx8fV7FfcOPyttE4SCszdVhagycsRp2GLom+uxcqssSMEW0Z2SGWP366jR93neT5G7sRF2po0PjqRFVhzSuuRUxTJXMCwUUwvlcSz36zx+0xb68P6bEhdKhFQiXwHo9ufwd+hsX3gsMG1grXY87z+pGKsuCHv8D3f4bLn4IB088eOl5i4uHPtmGxO1h83xDSYqo77kmSxLThbblnSBo/7jrJWysOceBUOWabE60iER2s545BKdzWL4VoHxKJqJCLT14AoryoeiqyxOAM926yZ7iySzyZhwouauMgSKcwTGwcCOpBj+QIrG6qSr45aEoMSBMGVoLGodklU6qqsvZIKR8NuJ3y2ZkQZCQ6WM+YzvFc0SW+Xg52jvJySj77jKIP53F1Rhc+ix+HO98Hb6xknSpM7Jvscwz+pmOrML68fwiv/nKAcf9dzdPXdeHq7o3ojHUkE0wl0PHqxntPgaABCdFruL5naxZtzsHuRoNS1/XBqFOY6aMBg6B2LnT7e3Z8V8Y6lsPSP4Dd5PnJ1tM9Vz8/BaU5qGOeYeHmHJ7/bi/3Dktj+vC2bnfGz6BVZK7unnD2Gqqq6kVtoPVOiSRYd7SGW6RPyYtWoX8DLR6v7ZHIk0t2uT3m7cZBbIie3ikRDRKP4NIiMSKIfqmRrDlYWOOYt98/naIydXh6Y4UsuMRpNsmU2eZgwfqjvLMqizKTjaqozpBTBbg04j/tzuMvi129QlOHpdcqUzgf24kTFH04j9LFiwkeNoykN16nXZcuXPbRZpbtPeWzNXmQVuHmPklEBGhPhE4j86crOjC6Yxx/+mwbP+zK45nruzROvGtegcG/By8bRwWCQGTW5e34YWceJSbvh0+Dq1eqXVwIV3Vt5afILk3Ol/3N/Wgul1mfR6e6d110i60KdcNsPt1r533nWD6aMoDOiWH1iuNiGNM5HrmW5M3bxaNTVbmxT9JFxXEGw+l72fz1R+u1cRCkVZg+Ir3pJOWCZs/04W3dyqrBu43tJDWPLqsegNF/h9j2/gpTIACaiQFFSZWVG9/M5N8/7CWv1Oz647rgIl1pdVBmtvPuqiyufGUlWfkVtb6eee9ejj/6KFnjbwCnk7QvFtP6hf8Q1MU1s+OlW3qSEReCvg7no/MxaGV6pUTwxLWd6/chG5FeKZF88+AwooJ1XPXKKpbvO+XfNzy5C3K3Qo/b/Ps+AoGfSQgPYv7UAYTqNXgoXFRDr5FJijTy4eQBATX7rSXRIymcl7Rv10ikUl8pJ+4/5VRazyUE726xMvL9c26Akt3EjcXv8tW9XeuVSDUEWkVm0sA26Gq554R0GUXCXa+Q8sdFJD/wEXE3P4UhqdPZ44okMa5bQr3GcdTG/aMy6mWkpJElEiMMTOjdMImd4NJkaEYM7eJC0Cq+J+QGrcwTd1wFrfvA3Ktgye+hLNcPUQoELgL+zl5ltXPTW2vZf7LcqzkZVodKfoWFG97I5HjJOamHqqpUrF7D0clTODZtOob27cn46Ufi//JntImJ1V7DoFX4fMYg+rSJxKjzXEmRcMl3RnWI4/17+qNpJoulIJ3CU9d14aWJPfjrFzt5/IsdVLqZc9IgZL7q6kvQNnKflkDgB7okhrPk90NJjjR6vD4oMhg0Mv3Tovj690MJNwq3Sr9xeAWytdztIYcK/13v2SJZq9Gg377AH5F5zd1DUuu0Lq8NnUbm/lENO88pLszAx9MGEmrwfuNAq0jEhOj4eNpADA1kwiS4NJFliQ8nDyAhPAidDwmVQSvzt3GdGNapNQydBb/fDEFR8OZg+OkJMBV790IOO+xZCu+OgecS4OkI+EcsvNgRlv8TKvy8CS1oVkiqB5vYvn37qps2bWrEcGrywIIt/LT7pFvJXeXu5ZRt/BJbYQ6yLghtXDrhgydiSOqCIkmkxhj56YFBlH//PYXvzQWnk6jJkwm/ehySrm5pm9OpsmJ/Pm+tOMTWYyXIkoTF7kACdBoFp6oyKD2aacPTGdQ2utlKGsrMNp75ejcbDhfx4sQe9EttwKbN0hx4cwg8tBWCIus+X+B3JEnarKpq36aO42Jp6uuTqqqsyypk8vubsDud6DQysiThdKo4VbihV2smD00NqMHZLZZ5N8ChX2o8nPpKOTP66vj3GgtZD4USYZB4d4uVj7bbWH53dXMJQlvBH/aA3HQbYtuOlXDr7HW1uvu5w6CVeeuOPozsEOeXmLILKrnxrUyKK61u7aoBZAn0GoWOrUKZc3e/WscIeENLuD419bWpJVFqsnHXexvYf7Ick9VBbStWg0YGCf59Y3eu69m65gllua4kaO9SGPQADJjh1s0TgI3vwi/Pudw/rW5UThqDy1Sr3eVw3WtgFEYXlwKerk0B3TOVX26pNZEq2/AFpesXEn3F/RjSeiMpGkyHN2M6sB5DUhccqkpuQTlf3jSFPjFa4h7+E8FDh/qU8MiyxKiOcYzqGMexoipWHsinpMqGLElEGrWM7BBHq/DmX20JM2h54eYe/LT7JPfP38INvVrzhzHtG2Znce0b0OsOkUgJWhySJBEVrCcqWMc3Dw7lVLmFKquDUIOG1hFBYme+sVBVOLyy1sN9ExVGpmp4IdPCs6M9XK/NZVB8GKKbziSkR3IEn08fxB3vrsfqcHq0OjdoZRRJYvadfet056uL7IJK3s/MZsvRYsrNrvEbKVFGJg1MpXdKBHpF5vFxnVi5P591h4vQKTKS5PrV2xxOrugSz9Rh6XRPEoYTgoYlPEjLF/cNZm1WIe+syGJtViE6jYzT6TJ9UVExaBTuGZrq2UUzLBGu+5+rd/uXf8CrfWDEo9BrEiinl8KqCt/8CbZ9DDYPszntp4dRH/gJ3hoCk3+EiMAzHmvOWO1ODuVXUGqyoZElYkL0tIk2BmzRIqCTqQXrj7p93GmppGT1fKLHzcLYYfDZx40ZAzBmDDj7/2anxNdXTmbCA6MvOpbkKCO3D2hz0a8TyIzpHE/vlAj++sVOrnttNS9N7EnX1t7ZsauqSqXVgeP0wGJZllzl9K3zYeYaP0cuEDQN3+/M46qurYgw6gLWeKbF42nRc5pnRukZ8l4lDw3w8G8ka1xDfZswmQLo2jqc1X8ezZe/5fDW8iyKqqyggt3pRJIkbA4n0cF6pgxL45a+yV6ZLdXGuqxC/vPDPnYeL8XpVLGdV3rac6Kc1QcKkCSJpMggJg9J495h6RRVWsktMZ3bOIgMIqyOocQCwcUgSRKD28YwuG0MJ8vM7DxeSpnZhk5RiAvT0zsl0qP7ZjVi2sHED+H4Zpeb59rXXCYVna93VaPqSqTOx2GF8pMwdyzMWCU2jRuA4yUmPszMZv76o6ioyKeTJ5vDSXyYgRkj2nJ9z0SMusBKXwIrmgv4YG2226qU5fheVLsVY/tBHp+vShJr8iyUVFnFQsdLokP0vHlHb77cepy73tvAXYNTuW9kW7e9YKqqsvlIMe+szOKXvS79sCSB3anSLTGcGXG7GNNuHNpw0YgsaJl8f9oRUxDYdI1TuKa9hn+uttIpNvD7WkP0Gu4YmMrtA9qw5Wgxh05VUm6xo9fIPPP1bpY/PJLgephDnM+8tdk89+0ej73IZ6zaswsqmfz+Rt6a1IeoYN1FyfgEgoshPsxAfFgDKIJa94E7l7jkwT8/Bcufd82gc5zrr0x9pZwqGxx+KIRgnWtRX0MmrDpcCdXPT8G1/734uC5RHE6VJ7/ayeebc1BVFaujpqDzSGEV/1i6m38s3c2LN/dgbLdGHO9TBwF7V7E7nBRXuW8adpjKkI1hXk2E1ykyuSXmhg6vRSNJEjf0SmLpg0PZmF3EjW9mcvBUdd3wtmMljPzPcu58bwM/7TmJ3alid6rYHCqqCtuPl/LIbzH02X4dizfnNNEnEQj8x5HCSvLLLfROEbuRTYq2lr6HC3h6pIHZW6wcL6ul68JpD7jeB0mS6NMmion9kpkyNI07BrahY0Ioe/LKLup1P990rM5E6nzMdifrsgqZMW8zjtoapwSC5oYkQcZlMG0FBMdWS6TO4I2BDU4rbP/03Pw6gU84nCr3frCRRVtysNidbhOpM1RZHVRZHfzhs63MX3ekEaP0TMAmU1U2B5payrZKUBjOqjJUZ91NupIElVY/udS1cBLCg/hwcn9u6pvMzW9l8t7qw2dNOW59Zx1HiqqosjqozcOkEgNlFiePf7mDV37a37jBCwR+5oddeVzRJd57eYnAP0gSpI2o87SMKJlbumj534ZaFkaGCIgK/CGfPZMj2HqspN7PP1pYxd+/2uk2karcvZwTH8zi6Es3kfPaJE5+9iTmHNfwXrPdyfrDRbyfmV3v9xYIAhJbJeS4Nwx5ZLCOFzItlJjr2kSQYfvnDR/bJcDfvtzBuqxCTF5u7gCYbU7+8c1ufvX3aB8vCdhkKlinwV5Ldqpv3RFJo6Vq/9o6X0dVXZIJQf2QJIlJA9vwxX1D+GbHCa5/fQ3TP9yEyea925TZ5uTtlVnMXx84uwgCwcXy/c48ruoihvAGBEMeAl1wnac9MUJfbebUWbRGGPxAjfmFgUiPpItLpuZmHnZbXSrb8AVFy2YTPnAiSQ98ROuZcwntPQ7TgfVnzzHZHLy94hBOUZ0StCQOLoNalE7nG9h4xFYJWz7wQ3AtmyOFlSzectxtIuVpcwdca8u/f7ETT67kjUXAZhmKLBEbqudUec0vsKwPJmLo7RT99BaSrGBI64UkazBnb8V8dDuRoyafPdfmcJIYEdSYobdIUmOC+Wz6IPo/9zNmN31s4Nmq3mRz8MzXu7mmW6KYtyNo9uSVmskqqGRgenRThyIASBsOhvAaMpvsWdVt6ZPDZcx/czOYV3VAz9v9GWGD0TMlgpd/rl+l32xz8OnGY9gu2Kj01tQJoNJiZ82hAoa1i61XDAJBwFGZ75L51oJXBjZnXkfgE+9nZuN0kwzV5dh9hqIqKxuzi+mf1rQS7YCtTAFMHpKKQes+xLD+E4gcPYXStZ+S8+rt5Lx5N+VblhLU7pwphSzBZZ3iGnQq/KXMrtzSWq16vdnVlCWJzzcfa6xwBQK/8ePuPEZ3jENXzyGrggZGkmD8m6Cpx8aZNggufwaCmoetd1p0MGUmGwUVdeyUu2HZnlNui2/emjqBy5TiAyH1E7QkHDZQa5eYnW9g4xGnrYEDa9mYbQ4+87C5EzVmJsYOg5F1BiRFgzFjQLViCYDJ6qqWNzUBW5kCuKVfCi//fKDW4yFdRhHSZVStx/UahanDAl8D31x4d9VhLPaayZS3u5omm4N3VmYxeUiayzpdIGimfL8zj7sHpzZ1GILzSR/pctP6+iGwm7x7jtYI/afCwBn+jKxBkWWJHskRbM8pYXTHeJ+em1tiwupGWeCLqRPAsSIvf78CQXMgKBIUrVsDijM8PdJA77cr+NOgWuZYAejdVL0FtbL7RJnbuVG+bO6owNqsQrLyK5i7JpuVB/KpMNtRZInoEB239E3mxj5JhPp5fENAb6tGBuu4rkeia7K1j2hkifTYYHomN4/dxubAT7tP4k4q78sXv8Ji51C+m4niAkGAYrU7OVlm5nCBy70vv9zCjpxShrcXMqeAo8ctcOt8MEaDLqTW02yK0ZVIjfkHjHmmEQNsGHokRbD1qO99Uyabw20vsi+mTmdeRyBoMaQM9Cjzg7oNbFRFB+3G+CO6Fkupyea2Uu7r5k6V1cHY/67i4w1HOVJYRWGllVPlFvacKOffP+yj33M/8/Dn2yipxSG8IQjoyhTAszd0ZW9eOftOlrvdUXOHIktEGLW8f0//gJ2W3NxwOFXMtdxAffniK7JEUaX/vtACQUOxN6+MOasOs2RbLpIEiiRhP72b0DoiiMJKK61FP2bgkXEZPHwA9v8Aa15xDedUdIAEDisVwUl8IN/A/Q886pL4NTfsFq5UV6P+tgAOWVz9XkGR0Ok66Hmbq3esFkINGrQauca99HxTp+COQ+sMIfQiZ1wJBAFFVBok9oajnk3NnhihZ95291I+q0PlG+3VjLU6CNJ5lwRc6ii1rM/P39zxNqFyN5MWONua8tVvx1lzsIDPpg8iOcq7cRq+EPBXRL1G4ZNpA5n8/ka255TWuSNm0MrEhRr4ZNpAYkM9lGMFPuFUVZBw1VQvwNcvvjCCEgQyp8rMTJu3mb15ZdgcqlvnsyNFVYx+YTmjOsbx8sSe4uYZaMgKdBzn+rFUgKkInK6kQ6cNY87zy7iuXCU5sMZKecZcBiv/A5vn0sXpRLZVwvkjFHN/g5+fhM7jYdTjENmmxkt0TghDI0tcuJ3li6mTRpaE4kPQ8hgyC/K2VzOx8drABrDE9+H7HD3P/esXbh+QwqRBqRe3Bs39DfJ2uP7utUEQngxtR7nkiAGIqqpUWOzYHCphBg0apW5FWVSwzq0zqK+bO95gc6qcLDNz01uZfPfQ8AYfPB7wyRRAsF7D/HsH8NXWXN5acYgjhVVYHdWzUKNOIdSgYeqwdG7tnyLs0BsYrSKjVWruaIJvX3ynqgpDEEHAcrSwihveWEOpyXa2CuUOh9OVZP269xTjX1/DwpmD/K7JFtQTfYjr5zQ64NruCSzecpyHLm/XdHH5QulxeH8clJ0Ah8W9Pt9W5frvjs9h37cw6UtI6lPtlP5pUQRpFbdGQmH9JyAHR1K69lMKlr6ApAtCH59B2KBbqp2nkSXuGZLWQB9MIAgQ2o2BmA5wcqfH3im3aIMIu/6fvJPYi6z8CuasPsxlLy5nXLcE7h2WRkZcaN2vAWAzwc7FsPplKDvuesxpB0kBRQOSDP2mQr8pEJboW4x+YkdOKe+uyuLbnSdQVVdPp83hJCM2hBkj2nJ19wQMWvebjZ0TwgjSKVRecD3yZXPnQjy5SjtVKKyw8tii7cy+s2+D/h4kT/7sffv2VTdtcj/IrCl5+PNtZBdU0ibaiMOpEhOiZ2SHOIZkRAtZnx+5670NrNjv3vqzbMNiStcvJvrK+z1+8cOCNGz+2xi0XuxaCPyDJEmbVVVt2CtJE9DQ16eSKitX/XcVp8rMPlVPdRqZbq3D+WTaQPG9biZsO1bCg5/8xvKHRwb+PaOqCN4aCuV5Lkmft+hC4N6fIa4TqqqyNquQV5cdZPcJlyvrhQ5a3tI9KZwlDzTMbrE7WsL1KVDXToI6MBXDO6NdiYzDS7dMTRDcNAc6Xl3t4cIKCx+tO8q8ddn0SIpg6vB0BqRF1X69KTzk2jCxlNcY8VANRe9Kqsa/AV0nePnBGp6Dp8qZOX8LOUUmLHaH23tm8GnFxh/GtGfK0DS3n/3Vhet4beMpLFLNhKti16+Ub/oKW+Gxaps7hqRObmOqzU7dcmxXtXWoTiOz+rFRxIUafPrMnq5NzbJ8c6SwklmXt2dou5imDuWSYvrwdDZlF9XYRQDvdjX1Gpm7BqWKBacgIHlz+SGKKixubwqedrusdid7TpTx3c48rusRGLuFAs90TwpHI0tsOVpMnzYBrvX7Yrprfs15iVTqK+VU2eDwQyEE61wLlHe3WPlou43ld58eXmytRJ03gRVX/8prv2ZRWGnlvpFteb1TLy5/aWW9elcNWplHruzQIB9LIAg4giJh+nJYcCuc2Ap2c+2W6dpg10iGWxdA+ogah6ND9Dx0eTumj0hn8ZbjPL54B8F6DVOHpzOua6vqMriCgzB7lCuRctdLcT5nkrwv73MlXb0n1eujXgxbjhYz6d31VFkdHqM9s1Z88cf9HC6o5NnxXc8mVLYTJ8h/9TUGrdnAa4Mfcvux63LsPh9fZuVJwIJ1R5k1pr1Xr+0NzS6Zstgd7Moto2eK0Gw3NoPaRhNq0LpNpqDuL74K3DGwpo5fIGhqLHYH89cfxepmt96b4YFVVgdvLT8kkqlmgiRJ3NgniYWbjwd2MlV6HLJWuJUdOVT473orjw+rrS9DxVRRzI9fLWDSlbdwTfdElNMjKRZMHcCNb2ZSafG+0hWkVfjTFe3FsF5By8YQDvd86zKuyfwf7PseNOf9jTntLondkFnQ9UbQeTYzMGgVfjcghVv7JfPL3lO8syqLf323l3uGpLpaUjCfq0jVlUidj90E3z4C0RnQpm4n5YYiK7+CO+dsqHUd6A6TzcHiLceJCdHzYP94Ct55h9KFi1Y3m6QAACAASURBVIi49VZ6LV3ME7uLeHbpnotyCfXFVdpid/Lxhks8mdp5vJT02GDRE9UESJLE8zd2Y+ZHmzHbvHNWPEOQVuGuwW2ID/OtrCoQNAbf78xDdXMj82W3K6uggr15ZXRsJWaNNAfG92zNuP+t4slrO9eq6W9yNs6mtgXWI4N1/HuNhfv66YgwuJcOGVUTz8UtQ+r5x2qPd2wVxqKZg7ntnXWYbU6PixiNLKFRJP52TSduHyA2wwSXAJIESX1h4ocumW3+XjCXgsbgSqRifa/OyrLE5Z3jubxzPFuPlTB7VRav/3qQ51M2cIWlHPm8v3OvKs/gSqiWPQOTv7voj+wtj3+xg0qrext5TwoOk83Bm7/sp89zD9FuWD/SlixBGx8HwO0DQimssPLG8oM+ry3P4KuderGpYQcsNzu91cbsYvoG8k5iC2dUhzieuKYzBq33Xx2DVmZM53geu6qjHyMTCOrP0u0n3O7S+7LbZXOoLNtzyh/hCfxAYkQQXRLDAvvfbPP7tTbD901UGJmq4YVMz70dUs5GKD9Z4/GOrcJY8egoHrmyA63CDATrFHSKa+GmkSWCdQpBWoVb+iXz7YPDRCIluDQxRkGbwdBhrMtNrx6J1IX0TI7g9d/1Zsn9QxiQ+xHyGfOY8zhTea6T3C1QdPiiY/KGY0VV/Ha0BHdWC2UbvqBo2WzCB04k6YGPaD1zLqG9x2E6sP7sOarTyfJpT5Dwj2fOJlJnePCydjx/QzcijdqzvVYXosgSWkVCdrN35OusPLujfklbbTS78s6m7GLG9xJSmqbkdwPaEB2i54+fbQWVWsu9Bo2MzanSKszAK7f0CPxGb8ElS2GF+wWpL7tdDqda6+sIApMJvZJYvCWHq7snNHUoNXHYweR5MO8zo/QMea+ShwZ4sPnV6F0N9aHxNQ6FGbRMHprGPUNSWZtVyO7cMsrNdoJ0CgnhBsZ0jseoa3bLBIGgWZBcuQMoc3vMm8oz4Orp2vAOXPW8f4I8jw/WZrvG5FyAtwoOm6Tw2b4yHrM70GvO3VNLqqzkl1toFx/K4vuGsC+vjNmrDrPtWAlO1aUZMWgUru6eQL/USJ79Zg/l5urVMV/t1BtajdCsrpKqqrL5SBHPju/a1KFc8lzZpRWb/zaGb7af4K0VhzhaVIVOkUECu0MlSKcweWgqN/VO4u65G/lyay4Teic1ddgCgU/UZ3igoPlwVddWPPX1LvLLLYE3l9Bucs3LcrqX1AB0jVO4pr2Gf6620inWg1rAZvL4VpIkMbhtDIPbClMngaDRyP0NHO7lZudXnp8d7aE9wmGFI5l+CrA6X2/LdesC6ouCA2DzkWIGpkWz4kA+b684xJYjxWg1MhISTlVFkSUmDWzDa7/rRaRRd1pq7Lq+FVda+ftXu2q8pq926l0SG1aO36ySqUP5lQTrNbQKF303gYBBq3BjnyRu7JPEiVIThRVWHE7XHKnkKOPZZucXbu7BXe9tYEhGjOiZEgQk0SHuF9K+7HYpslTr6wgCk2C9hjGd41myLZcpQwNsdpI22Csr9KdHGuj9dgV/GlS7EQUG0ccnEAQc5jKPM628qjwDWNxXtxqaC6tBZ/C1X2n7sVJmfbKVSov9rLLJ6qh+rXt39WHmrD7MdT0S+b8J3c4+Hhms4/JOcXy/M6+G8663s/KC9QozRrT1KlZvaVbJ1KbsIvqlin6pQCQhPIiE8CC3x7q2Duf2gW34y+IdzLmrr5D7CQKOa7onkHmooEbflC+7XVpFYnTHuAtfWhDg3Ng7if/7dk/gJVOyDBFtoDjb42kZUTK3dNHyvw1WusW5qU45bBCZ6pcQBQLBRaDRg6yptfrsdeVZ0zib1LWNpfVFwWF3Onnpp/1Y6+hZstpdx7/ensuRwirm3dsfvUbBWVXFzbkb+cUej1mpmWR6Y6eu1yiM7NCw9+pmZUCx6UgxfdpENnUYgnrwwKgMTpSaWbg5p6lDEQhqcFXXVki4T/LD+k8gcvQUStd+Ss6rt5Pz5t2Ub1lKULvqkoa0mGA6JYgKQHNjUHo0xZVW9uY1zu6uTwx+ELSerZcBnhihp9LqZqUjKdB5POhD/RCcQCC4KMJa15kIPT3SwOwtVo6XebBND09u4MDcE6x3nyidr+CoC7PNWWcideH524+XMOvj3yj5agmHxo6j3clDdE+JQqfxPYUJ0io8emWHs8qphqLZVaamDktv6jAE9UCnkXnh5u5MmrOBoe1iaq1iCQRNgV6jcPuAFOauOex21lRdu11GncLMkRn+DFHgJ2RZYnyv1q7BmuMCLBnuPhF++GuNh7NnVU+OksNlzH9zE7tGB4Pu91d0AoHgYugwFpb83uMpdVWeTVIQO2In0NPurFdy4QujO8axaHMOF94ife1XcocnW3WzzcmvO3JYc/wHRr38MsbevZhjtnHda2s4XmzyOjkLOjPzq39KfT6+R5pNZSq/3EJxlY12cSFNHYqgnnRJDOeuQan8edEO1NrqxQJBEzFjRFsig3VubVc9odPIdGwVytiurfwTmMDvTOidxJe/HW9wu9yLRh8K/ad6VZ2qgaKDxN6Q0L3h4xIIBBePPsS1YSJ5rmvUWnkGJF0QLx9OYvA/l/F/3+4hK7/iosM6XmLin9/t5aY3M7ni5RVc99pqfr9gC/1To9AoF6fgcIc3tupWWcPSsfdi7N0LgFCDlq8eGELX1mEYa7FSP4NGltBrZGaObMvfru7kw2/Ce5pNZWrzkSJ6p0QgN3BpTtC43DeqLeNfX8Pnm3KY2K9xStMCgTdEBuv4bPogbngjkzKTDfuF3a1u0EkqqdFG3p/cH63SbPamBBeQERdCQkQQqw8WNLiW/qK5/Ck4ucvl2GX37Mp3FlkLoYlw28f+jEwgEFwsg+6H7Z+B/VzflNeVZ20QhmEP8fHQoRwuqOSTjUeZ+PZa2saG8LsBKVzZpZVPFuBbjhbz4g/72HSkGFVVq6k0duSUsmzvqRqmD+fjTb/ShXhrq+5E4pe9pyiqtBIV7OqVCjNoWThjMKsOFvD2ikNsPlKMRpZwOFUkSUKRXe6AN/ZO4p4hqaTH+q8Y02ySqY3ZxfQV5hPNHq0i88LNPbj93fUMbRdDYoSQ+wkChzbRwXz/0DCmfriJfSfLsdkcONz0UmkVCUlV6XtqP289cBthBm0TRCtoSG7s7ZL6BVwyJStw2yc4F0/Dsvs7DJhr6e47jTYYotPhziVgCG+sKAUCQX2I7QCXPwnLngE3w3trRdFD6z4w6AHA1bP7l7Gd+NOYDvy85yQfbzjK01/v5vqeidzWP4X28Z77JhduOsbfvtqJ2ea+Oq8CVbXMFPUG6fRrXIgvtupaRWJdViHjup2bCyjLEiPaxzKifSw5xVVsPlJMmcmGVpGJDtEzNCOGoDoqVw1Bs0mmNh0p5vGxHZs6DEED0CkhjHsGp/LnxTv44J5+wt1PEFDEhRn46oGh7M4t43//WcAvhmRkjYwiSdidKlpF5vYBKUwa1AbdBwcofe5ZQl9/TXyPmznXdk/kPz/so9xsIzTQkmONjo+SnuRoXj/+GvGTq0olSWA3u44rOpfZRHQGDJ0Fna5z9UsJBILAZ+BMsFTAqhe9qz5rgiCxJ9z2KSjVl/E6jcy4bgmM65bAsaIqPt14jDveXU9ylJHb+qdwdbeEGsnF19tyPSZSblFV1zXIC2TJdao7FbUvtuoOJ5RUuZ/LBZAUaSQpsh6S6AagWSRTJquD/Xnl9EiOaOpQBA3EjJFt+fGNTD7deMwvzYACwcXSMULDH1a+y0s//UyZYqDKaifEoCHKqDs7QNB5331k33gTZUuXEn7ttRwrquKDtdlsyi6m3GzDoFVIjjRyx8A2DMmIFglXABMZrGNw22i+3XGCW/oF1jUpv9zCK8sO8sm0O5Hi74fSHNj3HVQWuGyVjVGQNgJaiYH2AkGzZMQjrr/fn5+CkqNgt9ScM6cLcVmpD7wPhv0RFM+bPslRRh6+sgOzLm/HL3tP8cnGYzz7zW6u7Z7Irf2T6ZIYzvESE48s3OY2kfJkCnEmkdLIYK8jB3OquC9L4ZutOuBzT3Nj0SySqa3HSuiYEOqT9lMQ2JyR+902ex3D2sfSWsj9BAFGZWYmhu7dMEZFUNtel6zTkfD883z/yDMsOBLO9rxKnKpabUr8rtwyVh7IJ8ygZcaIdO4clCp6PwOUCb2TmLP6cMAlU89/u4eb+iSdk+qEJ7mMKQQCQcuhw1jXT+5WWPcGHN8C1nKXfXp4MgyYDu3H1qhG1YVGkbmiSyuu6NKK3BITn206xtQPNhEbqic8SIvTTSNU2YYvKF2/kOgr7seQ1htJ0WA6vBnTgfWuZOo0qiohodaWK9XJ+bbqwR2HejxXkSHCGJgV92aRTIlhvS2TDq1CmTI0jccWbmfelP5i114QUJT/8iuho0bXed4SUyhP9LgDS055redUWR1UWR386/t9/Lovn7cn9RGbQwHIqA5x/HnRdj5cm822YyUUVFhRZIlW4Qau75FI/7SoRr9Orc8qZG1WIT//cUSjvq9AIGgiEnvChHf889IRQcy6vD2/H92OZXtOMnP+FhwXJFPemkIAOC7SmdkXW3WbQ2VIRvRFvZ+/aB7J1JFifjcgsHYKBQ3D9OHp/LgrjwUbjnL7gDZNHY5AAIDqcFCxfDkx9830eN5Xvx3nySW7sHg5ZcJkc7A+q5Bp8zYz9+5+DT44UFB/iiqtzFmdRaXFzj+W7q5WXZSAL387TqRRx7Th6dzWP8XvM10AbA4nf/9qJ3+/pjPB+mZxuxYIBM0ARZZQAYNWptJSXU7oiymEOzzKA90Q1n8CcnAkpWs/pWDpC0i6IPTxGYQNuqVavNf2SAy8ftbTBPzV2eFU2XK0mJcm9mjqUAR+QHNa7nfLO+sY3i6W5KimaR4UCM7HtG07mpgYdElJtZ6TU1zFY4u3+6w1N9udbDxcyJzVWUwb3tafH0PgJYfyK7j17XWUmqxuhzafcbKqspr453d7WLwlhw8nDyDc6N8b+/trsokPM4gZZgKBoME5XmzCZq95vfPFFOJCvJUHXkhdtupaRWLK0DSf42ksAn4wyv6T5cSG6okO0Td1KAI/0S4+lKnD0nls0Xa32l2BoLGp+PUXQi7zLPH7MPNIDXkEeDeA0GRz8s7KLPF9DwByiquY8EYmBZUWt4nUhZhsTnafKGPi22sxXYRVcF2cKDXxxvKDPHN9VyGBFggEDUZJlZV3Vh7ixR/3YXVjsXe+KYQvnJEHRo2ZibHDYGSdAUnRYMwYUE2u5ytBWoWb+iTRKcHNrK0AIeArU5uyi+jbJrKpwxD4manD0vhhVx7zNxxl0kAh9xM0DiargyXbjrN0+wmKKq0ARIfoGLw5h1sfvafW51nsDhZsOFpNCga+ac1NVgcrD+QH3lyjSwhVVblzzgYqzHYulP57qi7aHCrZhZX8ZfEOXrm1p19ie3bpHiYNbENaTLBfXl8gEFxa2BxOnl6yi8835yBLro0hd/hiCnE+FysPdEeQVmFkh1ievi6wnUoDPpnamF3M0HYxTR2GwM+ckftNfHstI9sLuZ/Av+SXW/jfsgMs3JyDJNUcRrgxcRivLs7llmMKvx/djsjg6g5Cv+7Nd/u6vtxMKq0OPsjMFslUE7I2q5C8MnONJmpvpCoWu5Nvd57gicrORAU3rMPUyv35bD9ewotC3i4QCBoAs83BnXM2sON4CZY6vMx9MYU4H1/lgRIgS5JbEwujTsGpqkwZmsafrmgf8NX5gE+mNh8pZtbl7Zo6DEEjkBEXwvTh6TyycBsL7h0o7KMFfuHgqXJueXsdpSYb9lpkdiaNHqwO5q07wrc78vhs+iBSos8l+LklJrfyCF9vJkeLfJh4L2hw3lmRVUOq50t1UQY+2XCU+0Zl+PS+5WYbK/cXUFhpwe5QCQ/SMiA9iqRIIxa7gyeX7OKpa7sIx0eBQHDROJ0q9y/YwvacEsx1DYU6jTemEBfi68yoxIgg+qdFsXR7ruv5koTV4SQxIogZI9IZ3yuJkGZivBPQUeaWmDDbHELmcAlx77B0vt+Vx0frj3DnoNSmDkfQwjhWVMWNb66lzGTzai6GzaFyqtzMDW+s4btZw4gLNQAuVz5HHVpzb24mde0QCvxHUaWVzKzCGt8DX6qLZruTuZnZXidTe/PKmLPqMF9vz0Ujy9gdTpyAVpawO1V6p0SSEG6gbWwwl3WK9/1DCQQCwQUs33+KtYcK3SZSnuTMdZlCXIiv8sC4MD0v39KTF2/uQbnFjsXuIMygbZabSAGdTG06UkyfNpEBX94TNByKLPHCzT246c1MRrSPpU20SKQFDYOqqkz5YKOrP+aCY55uKE4VSk02ZszbzOL7hgAQZtCgVWQcF9ycfL2ZNJddt5ZITnEVekXGesG/oa/VxYIKC06n6rGSrqoqz3+7lw/XZWNzqKeNS869r/X0f9dmFQLQrXU45WZbwNoACwSC5sNbK7JqSNmh/s57teGLPDBIq3B9j0TX82SJ8CAt0HyvdwF9JxfDei9N2saGcP+oDB5ZuJ1Ppgq5n6Bh2HqshGNFpnr1x9idKrtzyzh4qpyMuFA6JoS5/V76cjNRJOieFO7fDy2olUqLwyXavwBfq4uyJGG2OzDq3N9OVVXlsUXb+XrbCbc2+u7Yd7KcG97I5Kv7h4j5UgKBoN4cLaxi27GSGo/7Imd2h4wTJxIXXkS9lQeqqNzYp/bRI82NgL5Kb8ouZnyv1k0dhqAJuGdIGt/vzOODtdncMyRwZwsImg+zV2Vhsde/P8bmdDJndTbPT+hG3zaRRBp1VFlNNd7H25uJViMzOYDnZrR0QvQa3Gk9fa0uOlWVIA+ylHdXZ/H1thOYbN7bDFvtTo4VVTFj3mbm3Vv3wkYgEAjc8fOek24fr6/zniJLaBWJCd1b8eWWo1SpNc136pIHKjJc2z1wB/DWh4BNpsrNNrILK+maKHZuL0UUWeLfN3XnxjczGdUhjlTRNye4CKqsdn7efYoL/SZ8uaE4nPDFbzn84/ouaBSZGSPS+b9v97pdJHujNU+PCaFjq8Cdm9HSSY4Kcmsi4quTVbS1gpPPPodx4ACC+/VDiYg4e8xqd/K/ZQdrTaQ8yUstdiebjhSxO7eMzonieyIQCHynqNLqtje3Ps57Oo3MtT0SuXdYGh3iQzlZbmXNgZOYVV+SIifIFu4Z3rJcugMnmaosgP0/QFUBOB3klWkYE5+IThPwc4X9hqqqbDpSzLI9pzhVZkaSIC7MwBWd4+mZHNHie8nSY0N4YHQ7Hlm4jU+nDRJyP0G9KaywolEkLpSN+3pDcapQbrYTGazjht5JvPLzAcw2h1dmFudj0Mo8clUHH58laEgijDqGt4/l5z0na8yY8ra6GKSVmdwrCW15OSWffsaJP/8FXZs2GAcOJHhAf5YZUnC6sf0F7+SlVofKnNVZvDjRP7OsBAJBy8bdYHnwXc7cKtzAj38YXq2a9Pqk/tzy9gq2Hy9FdVOhchMNyBaCU95l1qp3+PzazwnXt4yCSdMnU8c2wJr/wcEfQdKAwwKqSqqk5QVVhXffhiEPQfurQGn6cBsDs83Bws3HeGtFFkWVVkzWc4s1CXh/TTbxYXpmjmzLDb2SWnTCec/gVH7YmcfczGymCEmUoJ5UWR24y8V9vaEokkSVzUEkLpnYx9MGcsMba1z9N15ikFUevKwdo8R8qSZn2vB01hwscNuc7U110anC7df2J9w4hOh770W1WjHt3EnlunUUznmPV0OGUhmWWPN5XspLHU6VpdtP8NR1XVqUJEYgEDQOUcE6tIpUY8C8r3Lm9JjgGtcgg1ZBm/w22opOWEt7AiqcTqoqdvxM2YYvsJfkIemDMLYfQNSYsYSkLULVFVNg0vLoykd5e8zbDfZZm5KmW4U7HbDkQfjweti7FOwWsFWC0w6qA63TjFa1QM4G+GIavHclmEubLNzGoqDCwvWvreG5b/aSU2yiylp911vFZcucXVjFU0t2c9ObmZRW2ZoqXL8jn5b7vfbLAbLyK5o6HEEzJcSgwY2iq9oNxRvsTmc1B7728aEsmjmYqGCdx74ZcOnEDYrE5H0/ck+UmC8VCPRtE0lypBGlHlVvg1ZmfK/WhBvPLTAknQ5j797E3ncfbT78gCMR7nt+fZGXahWZrPxKn+MTCASCYe1j3F7fzpczV+1fi9NmRnXYMR3aRPGv71U716hTGNctocZr7CrYxZHyLPQJiwjO+Be6mF9BqaRsw0KKl88lcvQkkv/wEUkzJ6NaN5P/+V9AdjmW2pw2Np3cRE55jn8+eCPTNMmU0wmf3wU7PgdbFW67gM/HWgl522H2aLCUN0qITUGpycb419eQlV/hVbOyyeZgT14ZE95cQ6XF3ggRNg2pMcE8dFk7Hlm4vdaStUDgidgQvdvKlC83FIBgvYYwQ/UKecdWYfz68EgevqI98WF6gnUKWsX1ZrIEwToFg1bmpj7JfP3gMGbeP56c3z+I7eQpv3xWgfdIksSHU/oTHqR1+/2oDb1Gpn18KM9cX7t9sM3hxF6LxM8Xeakkue4NAoFA4CsdW4WRHhPi9lhY/wlEjp5C6dpPyXn1dnLevJvyLUsJald9k8epqm7N4D7Y/QFWp2uwg6ypQB/zK0FJT1Oa+SEJdyQSc8XPhHZ+lrCuP5Ly+1isBVZKM88VRVRV5eO9Hzfgp206mkY3t+KfcHDZ6UTKSxxWKDkGn94Bd37lv9iakPvmb+ZUmQXbBQmDpyZlm0Mlp9jEHz7byjuT+jZR5P7nzkGpfLczj7lrDnPvsPSmDkfQzNBpZG7pl8y8dUdqyB287Y/Ra2TuGZzmtlcxPEjLlGHpTB6axtpDhezKLaPMZMOgU0gIN3BFl1bnKlrxo7EcOEjOAw/QZt6HyAaD3z63oG7iwwx8df8Qbn1nLYUVVreDLc/HqFPonhTOnLv6odfUngxpZAkJ91uFvspLg3TNb4ilQCAIDGaMbMufF22vl5xZkWF8z9Y1RjRU2apYdnQZTrX69dJ0sBKnzUnEABlJOacmUgwKod1DqdhVQeTwSMBVnVq4fyEP93242XsANH4yZa2CzFdrJFKpr5RTZYPDD4UQrHP9Ut/dYuWj7TaW333ayc1hgWPr4eQuiPd9oFggc/BUBZuyi2u4S3nTpGyxO1mxL5/cEhOJEUFNEb7fkWWJ/9zUg/FvrGFUxzjaxrrfaREIauPuwWnMX38Ud8tbbye9/25AisfjkiQxOCOGwRmenYqip03Fsm8fJ/7+BIn//tfZG4nJ6uC7nSc4XFBJmdlGRJCWtnGhXNkl3uPCXXBxJEcZ+X7WcOavP8qcVYepstqpPG/hIUug1ygkRQYxY0Rbru+ZiEbxLOyQJNcgymI3Mmxf+hVsdifxoSLhFggE9WNs11bMWZXF7hNlNTYT6yLUoGXW5e1rPF5gKkAjabCeHTnuwlHhQBOiQVJqJkeacA2mI9XHiVidViptlYTomvearvGTqV2LcTspEXCo8N/1Vh4fpq/9+XYrrH0Dxr/un/iaiLlrDteQsPkyA0dVVT5ce4Q/j+3YKPE2BSnRRmZd3o6HP9/GwhmDXTrgU3tg26dQehTsZjDGQNpw6HQdaLxxlxFcKqREGxnePpaV+/PdWsV6wqCVubpbArGhHq5NPiBJEgnPPcuROyZRNGcOZeNvY87qLBZtPo4sUW0hH6xT+MtiuK1/CncPTiUp0tggMQiqE2rQMmNEW6YNS2flgXxW7Mvn2Ibf0BmDSOnVhWu6J9C1tW/OU7f2S2HO6iysFyxgfLFfT40JJiVa/JsLBIL6oVVkPpw8gMtfWkF+hcWr58iSS9a+4N6BtAqvuZlTZa9ClmpuKCkhCvYKO6pDrZFQ2UvtaEKqpx0aSSOSqXqx+hWX0YQbHhms499rLNzXT0eEoZaSn+qAnYtg7D9BH+rHQBsPi93B4i3HsV+QTPnSpGx1qMxff4RHr+zQoi3E7xjQhu+2n+CXRW8zpnABFOx3SUDV88rXOxfB0lnQZzIMuh9C45suYEFA8d9bezL+9TVkF1Zh9TKh0mtkOsSH8n8TujVoLHJQEEmvv8bcmU/yUlYCdqQa1wA4l1h9kJnN/HVHeeOO3sIJ0I/IssTIDnGM7BBH3s6v0CYmEl3PTapJg9owZ81h3FVDvZGXBusUZo5sW9+PIhAIBAAcLqzE7nTSOyWC3bll2J2q2/uNhEtWHBuq54N7+tc64zNYE1xD4gdgzDAiaSTKNpcR3v/c5pPD7KB8eznxN1Vfj9md9mafSEFjJ1OqCkWHaj3cN1FhZKqGFzItPDvag6xB0UBRFiT08EOQjU9BhdXt477OwLHYHC5pkLHlVmRkp405xldh189ALTss1tM63fVvwpYP4K4lLea7Irg4jDoNC2cO5u73NrA3r9ythrz6+Qq9UyJ5584+fpHZLc2181Ln6zA7oC4jHptDxeZwMPOjzbxxe29GdxSbBP5GkqQ6/ZE8kRgRxMC0KDIPFbpduNTdryBxVddW9Q9AIBBc8uSVmpkxbzP/urE7V3RpxcFT5by3OpsvfnMpIWRJwomKza4yNCOGaSPSGZAW5bGPKdYY6zaZUowKcePjyP0oF9kgE9I5BFuxjdx5uWijtEQMjqh2vlFrxKhp/pX3xk2mbCZqk/id4ZlReoa8V8lDAzwlBBKYyxo0tKak0mJHdiO/93kGjixTbra33GTK6YRPb8d4ZBW1JlLn47C6fuaOhSk/Q3xnv4coCHzCDFo+mz6IH3ad5K0VBzlwsgKHqp7Vkus0MhLQtugo90/oz1Wje9bLOrsudh4v5S9f7DidSFXHk+mM2ebk/vm/8d1Dw2rdNRQ0FBI1Jvr6yIsTezLuv6sorLTgixmpQSvzbh0mFwKBQOAJs83BtHmbmDSoDVd0cW3MZMS5lBZPXNuZ3BIT5WY7Rp1CXJiB8CDv5tkZqnuIagAAIABJREFUNAauTr+aLw9+iUOtfhOLHReLEqyQ92ke1lNW5CCZsN5hJE9PRtaeW+zqZB23dbyt2ZtPQGMnUxoDuMlkz6drnMI17TX8c7WVTrEeGny1zT+TPYNRp+CsYwaON0PVHKpabQZOi2PVi5C9GuznGhi9Mi6xVsGH18EfdoGmYXpeBM0bjSJzdfcEru6ewP6T5azYl09RpRVJcg05vKxTPCGffYBt+SKUy3v5JYbXfjnotnfLG9MZm8PJ7FVZPHdDw0oPBRcgXXwyFRuqZ9HMwdz8diZFlVavGsCDtApv3N6b/mlRF/XeAoHg0kVVVR5duJ3U6GDucyMXNmgV0i/CzGtS50kszVqKw1FzRzBqRBRRI+q+fk3sMLHe7x9INO6cKVkGQ90NvE+PNDB7i5XjZbXcdBwWCKs5QKy5EhOix+nmhu3rDByNLBHm5a5Cs8Nhc+sCCeeMS2pHdVVFdy/xX3yCZkv7+FCmDk/nsbEdefSqjtw7LJ20mGAib7uN8h9/wl5Q0ODvWVhh4Zd9p2qs08+YzkSNmYmxw2BknQFJ0WDMGFDNlMDuVFm85ThV1pY7Xy4gkGozN/eNlGgj3z00nPE9W2PQym4HPOsUCb1GZnDbaBbOHMSojqIvTiAQ1J83lh8iu7CSf9/U3S/Vn7YRbekW0w2t7Pu6U6/oGZUyijhjy7jONf7Q3t53guJZhpYRJXNLFy3/21DLArlVdwhL9ENwTYNBq3BN9wS3UiJvh6ppZYmJfZP9IkcKCPYurW4ycR6PDNbxQqaFErOHRY+1Ala/7KfgBC0RTWQkYWPHUrxgQYO/9uebjrm9+PpiOiNJsHTbiQaPTXAeDVCZOkNUsI7/3NyDTX8bw+PjOtIlMYzEcAMxWEnT2bl7cBrL/jSCBVMH0iXRN9dAgUAgOJ8fd+Uxb+0R3pnUF4ObzZuG4pVRrxAbFItG8l4VpZN1JIcm848h//BbXI1N42vC+k+FDW/XedoTI/TM2+5m6rsuBIbM8kNgTcuUoel8s+NEDXt08G4GjixL3DMk1U/RBQCZr54zlrgAr41LirPg1F6Ia7n28YKGJequuzhyxx1ET52KHNRwM9y25ZS6HQ7ri+lMldXBztxSJpLcYHEJLkBySWUakhC9hkmDUpk0KBWA/DfeQLVYibv6+gZ9H4FAcGmyN6+MPy/ewXt393Nra96QhOvDWXD1Aqb8OIXj5ccxO8wezw/SBNE+sj1vXv4mQZqWMxe18ZOpiBRoMwQOrwLnuWQpe1Z1m/PkcBnz38JqPl8bBO2v8neUjU7nxDA6tgpj5/FSt65PntAqEn1TI2kT3YKb0YsOezzslXGJrIHCgyKZEniNPj2NoB49KP1qCYYJN7H5SDHFVVZUINKopXdKZI3J8N5QanKzUYTvpjPuBsIC5Ffl89m+z1h4YCGlllLsTjsGjYGOUR2Z3HUyw1oPQ/HSJfRS5mLd/LxBEx2Dacd2/76JQCC4JCissDD1w008eW1neiZH1P2EBiA6KJpPrv6EJYeWMHfnXArNhZjtZtTTF08JCYNiICEkgcldJzMufVy9pIGBTNO4FUyYDW8OgcpTdRpSVENrhNsXuqzRWyCz7+zL2P+upKjCgrMO18MzaGSJ2BA9r/+ut5+ja2LsJo+HvTIuUZ21VrcEgtoov/kOXpi/nB93/XRWRquiIiFhdzq5oVcSU4amkRHnfSNvbQmYr6YzoRe8TpG5iCfXPEnmiUxQXdPlz2Cym/jt1G88tvIx9IqeB3s/yE3tb/I65kuSBpT51YYmNgZHQaFf30MgEDRvKi12ftp9ktxSEyargzCDls6JYQxuG322H8pqdzJz/hau6Z7I9T1bN2p8Bo2BiR0mcnP7m9mav5WVx1ZSYCpAlmRijDGMThlNl+gujRpTY9I0WUlwDEz50WVZXZnvsq/2iAS6YPjdp5DYs1FCbApiQ/V8elUrfjdvOyUhkVjqcH3Sa2QSI4L4ZNrAlmuHfgaN4bS1fu08PdJA77cr+NOgWhz7JNklExUIvMDpVHnu2z18tK4UR0xX7Db3PXufbzrGF1tyuLFPEs9c39WrvsWM2BB+lU/VqEKfbzojyQqGtF5IsgZz9lbMR7dXM6HQa2TSY89Vo3PKc7jzuzspNhdjV2s3pqiyV1Flr+JfG/5FVkkWj/R7pEVY0/qHRkimoqP9YnIiEAiaPwdPVTBn9WG+PD0Tymxz4FBdiiSdIhNi0DB1WDo390nin9/vI8yg4ZErOjRZvJIk0SuuF73i/OOCG6g0XYknsg3MWA2/PgdbFwAS2Cqrn6MxuG5k7S6Hy5+GmHZNEmpj4SgtRX38YRbd9yCLQjKYm5mNzeGk0lJ9EResVzBoFKYMTeOuwan1khk1OyLTwFTs8ZTzjUu6xbmpTjntEJXupwAFLQlVVfnDZ1v5cddJl325VLsk7swk+cVbjpNfbuGtO/og15FQ3do/mffWHHYr6Q3rPwE5OJLStZ9SsPQFJF0Q+vgMwgbdUu08h1Pluh4uI55iczF3fX8XhaZCnHhX7Tc7zHy+/3PC9eFM7zHdq+dccjSQm58nlOgY7IWiMiUQCKozb202z327B7tDrXGvODPEvdLq4MUf9/PST/uJC9Gz9KFhdd5/BA1P067CjVFw9Ysw5h+wcxFs/QgqCyiuMGPRhtKq343Q524IiW3SMBsD1enk+KOPEjJyBK3GX81DwP2jMvhl7ym+25lHQYUFW04OUQ4zN982huHtY1uuc587Bj8ASx6sU6ZXq3EJQESqGNwr8Ir/LTvAj7tOYqqlGuUOk83BqgMFPP/dHv56tefvWZvoYLq1DmfTEfcbBHWZzkhAkE7hd++uZ+qwNH4zvUWRuahGIlW8qpiCHwqwnrKiGBTC+oQRf1M8SrArOTQ7zMzeMZur0q6iTVgbrz/rJUNjyPxionEUFKCqqqgQCgQCAGavzOKln/ZjttW9OXbmPnWqwkJOcRUdW7nxGxD4lcAoaeiM0HuS6wdYuu4Iu3PLeH7EpTOQsuD1N3BWVhL/yCNnH9MoMld0aXV2anXp1/lU/LqG1h1vb6owm46O18LXNV0cvTUuMUtBWPo8gDAcFtRFhcXOmysOub2JVe5eTtnGL7EV5iDrgtDGpRM+eOLZYbomm4MP1x7hvpEZRAZ7lt4+MDqDmR9t8SlhO4NBK/PBPf0x2xy8uWonW6XvQKou7Sv4roD87/JJujeJkM4h2Ipt5M7LJfuFbNL+moascVVvHU4H83fP5/GBj/scR4vHD25+FyIbDEh6Pc6yMpRwcYUSCC51Vu7P58Wf9nmVSJ1PldXBbe+sY+Wjowg1tCyDh0Cn8edMeUF8qJ5TZZ7tFVsS5b/+SsnChSS9/DKStvY/AG1iArbcS3SujEYHg+53mZD4iAo4FT1X/hjJ3DWHsTt8u0AJLi2+2JKD5MYApmzDFxQtm034wIkkPfARrWfOJbT3OEwH1lc7T5Lg041H63yfkR3imDSwjdsBrp4I0io8dHk7ereJZHBGDFcNPI5eU31fzGFycOrLUyTekUho91AkjYQuVkfyfclYC6yUZpaePdeu2vny0JeY6jB5uRSRZNnvlSk43TclpH4CgQD4zw+1J1KVu5dz4oNZHH3pJnJem8TJz57EnLPr7HGzzcmizTmNFargNAGZTMWFGThVbmnqMBoF65EjnPjr32j9ystoYj3LGbWJidhycxspsgBk+COQMsjVS+cDki4Y473f8NGM4fy46yTXvbaGLUc9918JLk1UVeWtFVk1qkVOSyUlq+cTNWYmxg6DkXUGJEWDMWNANVMIcN3M3l19GKcXIw7+n73zjo+izv//c2a2ZdM7CQkJEDqE3gJIEVEBsWFBVFBsWE7OOz3v9PR+enffK56Hp6dgQWwIcvYuFpoEUJDeAwECpJGQsiW7OzO/P1ZK2E12FxOygc/z8cjjAbMzn/lMMjvzebfX+/fju3LD4HZBG1QWo8Ldozpy5wUdT2x7d9e71J3W28O+247m1ojpXz9Kq1gUonOjqd1aP11WlmRWH14d1BzOLyQIsVXFmaAkJQkRCoFAwJ7SGnaX1vj9LBiHnsOt8uKKvc0eURfUJyyNqdQYMyXnQWRKs9spuvc+ku69B2vfwMonhuRkPJWV6K5A6ofnKLICU96GjhcGF6FSTGCJhWmfQJue5KREs+D2wdw5sgMz31zHw+9uotIW2u/SVufhYIWdPaU1lFQ7/TZZFrReauo8fp89dYd2oHtcWDsPDW4cp4fy2sAOIUmS+OPE7jxzfR96pMdgMcjIWn1DzihLmA0y/bPieenm/tx3Yad6tTWVdb6OAbVWxRBlQFJ8I2yGWAOe2vopgaqmUuGsCOrazivOQs0UgCEpCVVEpgSC8555K73CY6cTikPvmN3N2n3ieX42CY+aqdNIijJTYXOhavo5K7Kg6zpH/vgYlh49iJ8yJahjJIMBQ3IS7tJSTBkZzTzDMMVghuvehK3vwcp/Q0UBeFygn1yA1skRgIR50K0w9F6ISTvxmSRJXN6nLaO7pvD0V7u46N/LePDiLlzTP7NBBRxd11mzr4IXl+1lxZ4yjIqMJHnV1CwGhZvzsrhxcBYpMc3baVzQ/FQ73BgVGc9pBo3qqEa2xgTVSBfAoEhUOdxB3xPHayPXv/sFr684wJFu/amt8xBtMdI9LZqbhmbTPsl/U25V8625UqIUPLUedFX3Mag8VR4MUfUf/To6bq0B4ZbzGB0ptF6IZ4ghMRFPmYhMCQTnOz8drMRfJUIoDj1V09lRXMPgDonNMEOBP8LSmDIqMnFWE0dr687ZBWrl669Tt28v2QsWhKTgdDzV77w1pgBkGXpN9v4Ub4FNC+HYQfA4wZrE0YSBTFnZhu8uurhBAynGYuRPk3owuX8Gj36whUU/HOTJK3rSI71+AfiukhpmzP+BozYXDpeKDrjVk4tXp1tj7rK9zFm2lyv7tuXPV/TEqIRlwFcQBCZF9pseoUTEoNmr0TU1KINK18FkaPw+KK+tY+EPB1i//xjVDjcRJoWUvUe4rkssQ+8MLgIGEGWK8olOWXOsSAaJ6nXVxA46eU+rTpWaTTWkTk6tf32SQozJV7jFo2oosnReqcxtPVzFyyv2sWRbCfa6LHQg4rEvGNEpiTsu6EC/dvFN9/vQNNi3lGjLj8gHv4FP1kNMBvS4AhI7Bj5eIBCcU9Q6/fcIDMWh5/Jo1DiFc+xsEpbGFEBKtJmS6nPTmLKtXUv5iy+RvWghsiW06zOmned1U6fTpie0+XO9TemAed0y1h2oZGB2QqOH92wby3sz81j040GmzVvLZb3TeeCizkRbjKw/UMlNL6/B/rMR1RB1Hq8b6aMNhygst/H6jEGYDaGJCgjCg1ir0W/vJ3PbrkgGI/Zd+UR2HR5wHLeqkdCAmt/Ww1U8881ulu0sA07ePwCynsSnRQrZs5dz75gcJvRKC7hwz0vP491d79Zr1KtYFVKuSOHwm4eRLXI9NT9jgpG4vLh6Y3g0D/1S++FWNZZsK2HOsgK2H6nGo+pIEsREGLm2fyY352WRER+6CExrYHNRFb9dvJEDFXZcHg1V1+FnIRK7S+WrbSWs2F1OcrSZv12Vy9COv8Dr66yG9a/BqufAVYPVZfOe6cefQDbC8n9Am1wYPgu6jP+535VAIDjXsTRQPxuKQ8+gSA2OI2gewtaF3irrpir2wue/g5fGwHMDYM4IWHQz7Ft+Iu/eXVLC4d/8lvS///2MokvGtDQ8R85TRb8QuLxPWz7ccCiofWVZYsqgdnz165HY61TGPr2MV1buY9ora7EFMKROxeHW2Fh0jFkLN4jiz1aK2aAwqkuyj5afbI4kbvhUKpbMwb4rH83tRFc9OAp+pPK7eT7jDMiK9ytN+/HGw1z9wiqWbPM2Az7VkALQJBmn6k3ReHDxJn67eGNA9ckbu92I4uflmjw+mdSrUyleVMy2mdsoeLIAY4KR9g+1RzbWf/T3S+3HV5uc9H9yCQ/+byObiqpwq7pXCVP35uDPX7WPC/+1jKkvrz7n1FaX7Srj2rn57CypweFWfzak6qPrXqNq/1E7t8xfywc/Bfd88aFiHzw/BL79C9QWw3FD6jia2xtlL1oL794Oi6d5U5kFAsE5T1aif2fVqQ69QJgMMulxEU09NUEjhG1kKjXGQklNK3lhF66Eb56EIxtAU70vw+MUb4KCr8ESizb0fg7N/o74qVOJGj7sjE5lTE/DuXVbE0383OWy3HSufP57Hr+sR9BpdwmRJv4+OZd1+yuYNm8ttXX++/801mvI6dZYuquMdfsrGRAgKiYIT+64oCOrCo5id9X/+8cMugo5Mp6q/EWUf/IUkikCc2oOMUOvq7dfhOriiuVvU92uiuiLLkJSvIbOl1uLefB/G4PuHeJwq3y2+QiaDk9f27vBCFV2bDZd4ruwqXyTz2cJIxNIGNn4fWhRIqD8Wv5vxXYcjczNpeqAzuq9R7n0mRUsvmsoHZKjgrqWcGbjwWPc9ca6kPp9Od0aD7+3ifhIEyM7h9BUvqrI62xzHguuFsttg11fwdvXw9TFXhEegUBwznLz0GzyC45iO+39c6pDT5IVLO37IskGnIUbcB7YVF+EQocxXVPO8szPb8LWmEqJNlNa3Qrk0de+DF89Co31aHHZvD9fPkJqh1Qst756xqczpqdTs+TrMz7+fKFdopV2iVZW7ilndJfQHirtEiJ/Xjj6Ur32farW/I/Ecfdgad8PSTHg2LcOx+41Jxq3Ot0qLy7fK4ypVsrA7HgSo8zYK+w+n0X1GE1Uj9GNHh8dF80ll1xOxUsvU/rvf5N46wxsoy9m1sINITcCdrg1vthSTF7HRK4ZkNngOR8a9BAzvpxBnRraM1PGSHXReJbXenyiZA2halBhd3Ht3Hw+v/8CkqPNIZ0znNB1nXve8t84OVCDZqdb474F61n3x4uCc9hoGrx+OTir6hlS2bNrsLth3/1RRJq8BvPL6128ucnN0umR3nfLgXz47i9w4WNNc+ECgSAsGZ6ThNVs8DGmIDiHnlGRuG5gpkjzO8uErzEVY2Hr4arAO7YkP70FSwIYUqcg48ESeRRp8XSYstArpBAixrQ03CLNLygm9U7n4w2HQzamFqw54Kdl60lp0sTxs7B2yTux3ZozGGvO4BP/13VYuquM8to6kqJa70LzfEWSJObe2J/Jc1b5RKcCYTHKzL2pP7Ht4okZMwbHunWUv/QSz328GU/GYDjtzgrGOHe4Vf7z7W4m989oMDrVO7k3fxn+Fx5d+ShONbiIvkk2kah0Y1/NIL9SvI0ZE/rPqX9PfLyVZ2/oF9LvKJxYu6+CCrtvCl0wfxfwqmZ9tbWECblpPmP4sPc7qCmupzx6Yhwdnlnj4g8jGnheuO2wZo63155RpO8IBOcqsixx18iOPPXlTr9OnkAOPUWWmJ7XvjmnKPBDGNdMWSgJ58hUxV749DfgPmlIZc+uIeWfNdhcJ6MaL693MWq+7cT/JY8TClfADy+Ffs66Wozly4hL2Iz+/kz44vfw46veYmaBDxNy0/h6ewnOENJ3AD7dfNivlz4UaVKjIvH9HiF13Frpnh7DvOkDsZoUv4a1P6wmhbk3DaBvu3jAa5RZBwygzX9f4LP2Q3GfNlIofUOO1rpYt7/xRtMXZ1/M7NGzsRqsRBgaXnDLkoxFsdAnpQ/9LL/x20YpmOaQHk3nq20lVNlbr2rU3OV7cbjOvEGzzaUyZ9me4E72/TPgqvX70YN5Jp5aVccxZ4Bay63vB3cugUDQarklL5vBHRKwBFCEPZ0Io8xfr+hFuwbqrgTNRxgbU2ZKw7lmavUc0HwlLI97GBvFbff2SApWpKB8D3x8PzyVg/ztYyTmVCFtXACrn4cv/wBPdYL3Z0Lp9jO4kHOXlGgLuRlxfLO9NKTjqhz+F4ehSJN6VL3BcQStgyEdEvngnmEMyI7HbJAx+muA+3ND3b6ZcSy+a6jf+plvd5Si+7lnQjHOHW6V+asKA+43rO0wvrv2Ox4a+BCZ0ZlYFAtRxqgTPybZxLiscbx6yas8N/pFPvypxEe9MBRjQpbgnR8PBpxXOOJRNZbtKvMRmAm1QfPO4lqOBmrQXH0EDq5u8OMB6Qqjsg08taqRcVw2+P4/Qc1JIBC0XmRZYs6N/clMsAbtzLMYZR6Z0J2r+p/HbXNakPBN84sO48iU2wEb3qwvNPEzD+aZ+Mf3ddw90EScpZGvQV017FsGHUY1fq6tH8AHd4Hq9mu84f65rmPTIq/XcsLT0PeGoC/lXGdS73Q+2ngouDScAITaa0jQ+umcGs3iu/LYf9TG/FWFfLb5CDU/9wGJMhsY1yOVW4e1b1SI4UCFDZfHT2pXCMa5rkNBmf+oxulYjVYmd57M1Z2upuBYAaWOUlyqi2hTNDlxOcSavX2nvt5W4rcpemhGnsY7Px7k9gs6BDW3cKLa6UGRJdTTjMlQGzQbDRIVNheJjaX0Ht0Nihk8Db/TnhhtZtg8G/cP9i+pD0BlYVBzEggErZuCslrKa+u4d0wOi38sosbp9qmjMioSsiSRmxHLb8d1EU16W5CwNaaSokxU2lx4VA1DuDVB3fk5p9c+HOdUD+OfxzTSQ8plg7UvNW5Mbf0A3r8ruJosXfXu99kDgAZ9bwx8zHnAxT3b8OQn26hyuImN8JWq9kdchMmvIR9KryGDLBFnbWRRJGhVZCVG8vhlPXj8sh6Bdz4Nu0vF7UfQJFTjPNT6LUmSyInPISc+x+/nR211aH6i46EaE5V+ao5aAx5N8/sUD/XvIiH5/fvWoy6wIdwzRWFiZwN/W+miW3ID7zxPGGdrCASCJqG2zsO9C37iT5N6cHmftjxwUWdWFRzljfz9HKy043CrxFiM9GsXx7S8bLISI1t6yuc9YWtMGRSZ+EgT5bUu2sSGWePeqoO/3MMIjXsZy3bCBzPrGVIBVZ/AGzX79LeQ2gPS+wZ7RecssRFGhnZM5MutxVzbiBraqUzqk86z3+zGeVrdVCjSpB5NZ3hOUpNei6B1EmU2YFR8F9yhNgKOMjf8uHarbr458A1vbX+LI7YjuFQXVqOVnkk9mdZ9Gj2TevqIV3g03W8/tFCNCS04EcCwI8Zi9Cu8EerfxaNpxFoDOGpMwdUw/L9RFvrNreU3QxuIchmEg0YgOJfRdZ1H3t/M4PYJXN6nLeB1jA3LSWKYWFOELWEW8qlP2NZNuR1+U/yOc6qHsfFxfKWXT/D9M34NtqBqsjxOWP5U4/ucR0zqk87HGw8Hvf+UQe0abNQbM+gq4sfMoCp/EUXPTqXohenUrP+EiE4nU6JkCS7slkJCpFj4CCAnJQqTn0LiUBoBK7JEj/QYnzHcqpv/rP8PFyy6gMdXPc6Gsg2U2EuorKvkUO0hlhQuYcZXM5j4/kS+2f9NvWNjI4zIftL8QmkOCRBlCVufXKNYjAo5Kb7pmaE2aI42G0mLCeDwi88GNXAELydB5roeRv6ztoF9o9MDjiEQCFovi344yI4jNWeUBSFoOcL6LRi2dVPmGFBMjb4cA3oYj4/jD2c1bHnXr4RucDVZOuxeArVlEBVCQ8lzlAu7pvL79zZTVlMXVE+chEgTY7qm8NXWElQ/nvtA0qRmg8LtI1pfDYmgeRjRKRmLQcHmpwl0sI2A/cnd1rpquXPJneyq3NWgHLqGhsPj4EDNAR5e8TC3VN7C3X3uBmBQdoLf9LRQIrBGWQqtaW2YMXNURx59f4tPLUKwfxeLUea2Ee39GqX1iM/2ZgscWhdwTo+NNPPGJj/OOqMVhswMeLxAIGid7Ciu5h9f7uSdO4cQYRJ12a2JsDamUmPMlFSHYWSqTc+AxtSpHsZeKX4CgLIB2vb3f/DmxSD5DxoGXZMlSfDTmzDi141dyXlBhElhbLdUPt10mOnDguu/8MiEbnxfUE61w4/oR2PnMspc3COVPplxZzJVwTmIIkvcMjyb577Z45M6CsE1AtY0nae+2sn9F3aid2YcbtXNXV/fxY6KHbi04GqWnKqTV7e8itVgZXrP6SRFmclJiWLbYd/WCsEaE7Iscevw1tvT5NKeaTz6/ha/nwXzd9F1gk4fZtj98MHdPvLohbOi6/0/M1bG+agfR5uuQe8pwZ1LIBC0Kmx1Hu55az1/GN+NnJTowAcIwoqwNqZSoi2UhqMxlT3CG1VqoGfIcRr0MILXmGrIy1i8qdEUwKBqsjxOKN7Y6PzOJyb1TufZb3cHbUxlxFtZcNsQrn9xNXaXBy1IFfsubaL55zW9G2yuKjg/uWFQFnOW7gU/xlQgLEaZp6/pQ7mtjjvfWEf39BiyOuSzs2KnjyFVuaKS8i/LcZW6UCwKMf1jSJ2cihLp9XI6VSfPbfgvHlsXFn3vQtN1zAbZb1+1YIyJHukxtE9qvcXPFqPCw5d25a+f7fDbILMxIowKM4a3Jz7YdN4u48EU6RUfajCRuAEMFuh9PVgayGYQCAStmsc+3EqfzHgmC2nzVklY10ylxJgprQnDND9Jgrx7vWkXp1A4K5qxHU7ap8c9jCfEIU4ltQckdfI/vuNYo6cPuiYrwDjnE8M7JVF41M7Bikbq1E6jZ9tYPr5vOB2To4gwKjSUyWMxypgNMhd2TaG4ykmlrXWqmwmaj4RIE6/dOpAIY2ipGxFGhduGd2B8bho3D81m6YOjGNk5gf/tecsnta/883KKFxfT5to2dH++Ox3+2AHXUReFTxWinWIs1XlczNv0Go9M6MaSX19A/yxvH61QsRhlHjsH8vpvGprNjUOyQvrbRBgVxvVI5TfjOgd/IsUIN38YtBjFyeNMkNQZLvl7aMcJBIJWwf/WFbGx6BhPXtH6n6fnK2FtTKVGW8IzzQ+gz1Svt/BMMETAmD82/Lk5cIj3/42y8NJ6F4eqG/FwBjHO+YJRkbm0Zxs+CkGIAqB9UiRLHhjJwjuGML5XGkbF26TValJTl0rJAAAgAElEQVQwKhLJUWZ+PbYza/5wIa9MH8iUQVnc9eY66vz0FRKc3/TPSuD1GYOIMhsw+WkAfCqS5F2w3zO6Y70Fu8Wo0C7jANbTSv9Uh0rpB6Wk35hOdG40kkHClGwi8+5MXOUuqlZVnTK4hjtiHQM7RCLLMi9PG0D7pMiQDCqLUWb2dX3PmXTWRyZ049cXdcJskLEYG/49mAxex8m0vCxmX9cn9Ah0SjeY9glYYr3ZCYEwRkBab5j+CRjDTNVWIBA0jssO61+HhVNh3qXw2iRvu5t9y705wsDukhr++tl2/ntDP6ymsE4WEzRCWP/lUmPCVIACICIOpn0E8y5Gd9mC7lKN8WdDqmMj6TPJXb2GWiM9RQLWZCkm7ziCE0zqnc7jH23lntH+++40Ru/MOJ67oR8uj0aVw43TrRITYSTGYqi3oLpvTA5bD1fxp4+28X9X9WrK6QvOAQZmJ/D1AyOZv6qQt9bsR9P1esIUxw2a0V2SuXNkR/q2i/cZ451d72D31I+w2nfb0dwaMf3rp4EpFoXo3Ghqt9YSf8HJsWRJ5ruD3zGxw0SsJgPv3z2Mu99aR/7eo7hV3aeR7XEiTQqSJDH3pv7nnEzvHRd05Kp+GSxce4B53xdS51aRZQnN6QRZQTYZuWlIFjcNzSItNuLMT9S2H8zMh5VPw4YFaHUuZPnU2kzJG72yxEHer2DArUISXSBoTVQdghX/go1vez1jLlv9z7d/BJZYXIPv4/41nfndJV3o0kY4v1szYW5MhWma33Ha9EKf9jnaCxciGSVkvZH0LtkAshEu/qv35dgYva+Hb58IePpGa7IkCfrdHHCM84mB2QlUOdzsLK454weXySA3qggoyxJPX9eHK//7PW+t2c/UwVlnOl3BOUqbWAsPX9qVBy7qzFfbitlyqIqjNhfRZgOZCVYm9U4nMarhe+yI7YjPNrVWxRBlQPIT8TLEGnDsr9/426W6KLeXn/h/hEnh1VsGsaO4mldW7OPjTYcxyDLH/QRuVSMz3spdIzsyITcNS4jpiq2FpCgz947pxMxROWw/Us0xu5uK998j1igx9L4ZfiXuz4jYtjDhX+hjHqd06gBSruiD7KnxOsFi0qHPFG9trqi9FAhaF4fWwxtXeKNSDbXQcdm8P1//if+aO5Ld8/OzO0dBkxPWxlRilJljdhduVcOohGdGYsXnP2A7NI7M24ZC/nNeUQrV5f2RFG8kSlchdwoMvbvhOqlTiUyCnItgx6ecWqgctOoTQOYQiAtSZeo8QZYlLuudzkcbD/Fgm+aL2kWZDbx08wAmz1lF59RoBmYnNNu5BK0Xk0FmYm46E3ND6x3k8qMiqkQpeGo96KruY1B5qjwYouo/6lVd9Sun3rVNDP+8pjePT+rBntJaqh1uzAaZlBhLqxaaCBVFlujZNhaAysxonDt2Np0hdQrOPfuxu7sg3/xOk48tEAjOMiXb4LWJvpGoBjDpdWS79yC9dhnMWCJSeVsxYW1MKbJEQqSJ8tq6X5ZW0UzU7d7N0Vfmkb14MVJGWxh6DxQuh+LN4KjypmrEZkDXCV4Vp1AY/gAUfONtEBwqxgi44LehH3ceMKl3OjPfWsdvx3VpVsW97KRInrqmN/e8tZ4P7x0WlvevoHUSbfKNqlpzrEgGiep11cQOij2xXXWq1GyqIXVyar39TbLJ7zjHiTIbzpl6qF+KITkZz/IVzTK2fc0arIMGN8vYAoHgLKJ6TkakTiF7dg12N+y7P4pIk3fN8fJ6F29ucrN0eiSS6oLy3fDFw3DZ7JaYuaAJCM9wzymEa92U7nZz+HcPk/zrWZgy2no3yjJ0GAV598GFj8KIByD32tANKYCM/nDh4z6KgQExWmHYr6H9BaGf8zygR3oMRkXmp4PNr3Q4qksKtwxrz51vrMMZouyyQNAQA1IHYJSN9bYpVoWUK1I4/OZhajbVoHt0XGUuDj5/EGOCkbi8+oaRLMv0TOp5NqfdajGkpOApLW2WsW1r1mIdPKhZxhYIBGeRnZ822PZA1eGZNY2UgXgcsHEhOH17/glaB2FvTKVEm8Oy11T5iy+iJCYSd801zXeSITPRhj6ApkrBdSUxRngNuZEPNd+cWjmSJDGpdzofbQhN1e9MuWtkBzITrDzy/hZ0PcTeMgKBH6Z0neI3qpo8PpnUq1MpXlTMtpnbKHiyAGOCkfYPtUc+TaEuOSKZ3KTcszXlVk1zGVO6241j/XqsAwc2+dgCgeAss3J2g71HH8wz8dSqOo45G1kDSJJXsELQKgl7Yyo+0sTOkhpKqp1h4913bttG5VsLSPvzk83enPXIF2VUKpOR2o8Eg9lboHwqstGr/NcuD65fAKP/IIqWAzCpdzqfbj7SoGJZUyJJEv+cnMu2I9XMX1XY7OcTnPtkRGc0aAgljEyg01860eOlHnT7TzfaTm97omHvcSIMEdza81bRWDpIDImJeI4dQ1eb9v3j2LwFY7t2GOJ9FRsFAkErorIQSrc1+PGAdIVR2QaeWtVIlpXbDqtfaPq5Cc4KYVkz5XCpfLzxMHOWFbDvqA1Fkvjvd3twqzpDOyRy58gODOuYhNxQF9VmRHO5OPzw70n93UMYU1MDH/ALqP78c5wbN5H2/ntgtcKxg/Djq1CyBZzHwBIDyd1gwC2Q0KFZ53Iu0SE5ijYxFvILjjK8U/PLO1tNBl68qT9XPr+KLqnR5P0sKe3yaHy5tZjlu8o4anNhUmTSYi1c0bctvUW9iqAR7u93P7d9dRt1amgp0DIy0aZoxncY30wzO/eQjEaU2Fg8R49iTEk543GKKu1sP1JDjdONxahgXfETXQeKFD+BoNVTWeh1dDfSzuaJ0WaGzbNx/+BG2hzUnJ2MGUHTE3bG1LyV+3jqq50A2F1eT6BH1/H8HEVYuaecnw5UEm0x8t+pfemfdXaV0sqf+y/GzExiJk1q1vO4S0oofvLPZM55Adn6c91UXCaMfaxZz3u+MOlnVb+zYUwBZCZYeeb6Pvxq4QZemTaAL7YU8+aa/Wiajs110uMtAQt/OEharIW7RnXk6n4ZKC3gNBCEN31S+vDY0Md4Mv9Jv6p8/pAlmShjFPMvnk+EQQiihIIhJRlPaVnIxpSm6SzbXcbcZQX8dOAYJkVG1XVkSUJ1JBAdkcpdK/dxdf8MYiOMgQcUCAThh8uOv1qpU+mZojCxs4G/rXTRLbmBpDBPI3VVgrAmbIwpXdd54pNtLFx7EEeAdD6bS8XmUpn68hqem9KPsd2bN0J0HMfGjRx77z06fPB+s6bI6JrGkd//nvipNxCRK+oamoOJvdO4ZPYeZgyvZk+p7YS3OC3WwsDshGaJeg7LSeKqfm256vlVKLKES9V89tEBh1tlb7mNxz/cyoc/HeKlaQNEZ3SBD5M6TsJqsPL7Fb9HR280SmU1WIkxxfDqJa+SEZ1xFmd5bmBMPl431SPoY0prnNz48hoOVTpOOEzqPKd852Ujjjr455c7eeqrnfx3aj9GdznzyJdAIGghzFF4XaGN8/9GWeg3t5bfDG2gj6BBSKO3VsJmhTZnWUFQhtSpON0a9769ngW3D6Ffu1+ed15WU0dJtROHWyXaYiAz3kqk2fsr0pxODj/8e9o8+giGpOaNZlS++RaazU7SnXc263nOVzyqxsaDx/CoGhOfXYnZIKNqXjFGdG8D0xnD23P9wHbERzYSkg+RvWW1LFhzAFXXUdXA9VoOt8qP+yu56ZW1LLxjSNj2WhO0HGOzxtInpQ+Ldy7mrR1v4dE8aLqGpmkosoKma7SLacetPW/loqyLMJ1ecykIilBFKEqrnUx4diWVNteJrIqGOP7Om/nmOv51TW8mhNh3TCAQtDCJOeAJnHKdkyBzXQ8j/1nroleKn/d5fHbTz01wVggLY6q8to7ZX++u77X7Gdu2pVT/8AHuo0XIpgiMKR2IzbsWS4bXQ+h0azz0v018/cDIMzq3quks31XGnONpGAYZSQJdB4+mcVluOjNGtCf+1eexdOtKzCWX/KJrDUTdnj2Uv/AC2QvfRjKExZ/nnOJIlYPrX1xNeU3dCW+x+7TCcptL5ZlvdvOfb/bw7JS+TRL5VDWdqS+vweby+HzW2D1e59HYeriKv362nccvC94rLjh/SIpIYmafmdyeezs/FP9Asa0Yh8dBtCmaLgld6BzfuaWn2OoJxZhyeTSuf2l1UIbUqTjdGr9ZvJGMeKuomRQIWhMx6dBuCOxbFnDXx0aaeWOT2/cDUyQM+1UzTE5wNgiL1frbaw/43V699n2q1vyPxHH3YGnfD0kx4Ni3DsfuNSeMKYBDlQ42FR0jNyO0F9CO4mqmzVtLrdNzYmF9eurVe+uL+GhDEd2PRjPvTzNCvLLQ0F0uDj34kLd3VVZWs57rfOTQMQcT/7OCaqcnoJKf0+29D+59ez3/uDqXSX3a/qJzf7ujlGqnm9PV0YO5x51ujYVrD/LgxV1Eup+gQQyygaHpQ1t6GuckhpQUnFs2B7Xv51uOUFzl9GtIBeMc/NvnO3j7jiFNOn+BQNDMDLsfDq3zkUcvnFW/OXpmrIzz0Rj/Y/S4qrlmJ2hmWnxlpmo681bu84lKaXU2jq18i8Txs7B2yTux3ZozGGtO/Y7xdR6Vl5bv5dkb+gV93vUHKrnx5TU4XGqjZYOqDqoKWxOyuerNrXxw9zBirWdQKKxpsPc7OLQe7EfBZPV6M7pfCZGJAJQ9+xzGtLTm7V11nuJ0q1w3N59qhwc1hH5PTrfGQ+9uIjPBSt9fkEo6Z1kBtrr6EbBQ7nFJgg9+OsQNg4WRLRCcbQwpybiDjEzNWVZwQjzpVIJ1Dq4/UMnBCjuZCSE2bBcIBC1Hh9Fo5gRw1npLBkLBaIXBM8EoaqZaKy1uTG0/Uo3LT3pf3aEd6B4X1s6BPa2aDt/sCD6f/WCFnWmvrPX7wmsIlwaHKu3cPG8N787MwxBs/YrjGKx/HfKf83bHdtmBn6/XGAFf/gE6XYwzYRzHPnifDh98IPq/NAOfbDpChc3l15AKxlv8jy92nrG3uKjSzpZDVT7bQ7nH7S6Vl1fsE8aUQNAC1MQksrvKTfWhKqItBtLjIvzWMO4ormZfuc1neyiOE03XeT2/kEcmdG/y6xAIBM1DzbffUvaJhewLo0C1g+67rvWLIQLaDfX2CBW0WlrcmDpmd/tVTlMd1cjWGCRZ8XOULw63iqbpQamw/fPLnX5rV6DxhbVL1dldWstX20oY3yst8KRKt8P8iV4jyuPw/dzt3abv+AST5xOypl+NIeHsSr2fL8xZuqfFvMX7j9oxGWSf6Guo9/jhKj/3kEAgaBZUTWfZrlJeWFrAhgPHMGRMRHlxNZquo8gSNw7O4qahWaTHnZSZX7e/0u9YoThO3KrOyj3lTXYdAoGg+dB1naNzX6Ry4UIynn0VOSMKXh0P9gr/675TMVqh0zi46iUIch0gCE9a3JhqCCUiBs1eja6pwS02dSg6ZqdtnLXRvjxVdjdfbi3GX8lMMAtru0tlztKCwMZU2S545SKoqyVQ/wFJ15AUMBd/6o1UXfJ/ga5WEAKbi6o4dMy3F8/Z8hbb6jx+b4FQ73F/Ai0CgaDp2XjwGLe9/iP2upP1tG6DGepOOuFeWbmPed/vY2JuGn+7OhejIlPt8OD20/IgVMdJjdO/s08gEDQfuq7jcKvUOj1EmBSizIZGM4U0p5MjjzyK68ABshctwpj6c2uDu1d7M5JWPQuuGq9D/cQ5ZCSjCdJ6w7BZ0OVSbx6/oFXT4sZUnNWI5if1yty2K5LBiH1XPpFdhwccR5LgurmrqbC5yIiPoH1SJFmJkWQnWslKjKR9UiRpsRbe+fEg/mytUBbWu0pr2FNaQ05K9OnDeHE7YP6EoAyp+sfZYd2rkN4Xcq8N/jhBo2w4WOn3HgvVW7yq4OgZnT/KbPDbgiLUe9xsENLoAkFz8/2ecm577ceAbTqOixV9uvkIByscvHHbIIyK5G3Ie9pzP1THiWiDIBCcPcpq6liwZj/zVxVS4/RgkCU8mo7ZIHPdwEym57WnXWL9rBR3SSlF996LKSuLrDdeR7acUu9kiYG8e2HI3bBvKRSuhNpSPFU2qlf+RML/vQuJHc/uRQqalRY3prq0icYgy0D9F5dsjiRu+FQqlsxBkhUs7fsiyQachRtwHthE/OhbT+4rwbjubZhzU3+cbpX9R+0UHrWx/6iN7cU1fL6lmP1H7ZTV1iHh38MfysJa12Hl7vKGjakt7/3siTj5Qs2eXYPdDfvujyLS5F1Zv7zexZub3CydHnnyWLcDvnkSel0jvBVNRLWzZb3FWUmRfusCQ7nHAdqekk4kEAianp3FNdz+emBD6lScbo1Nh44xa+EGLunZBpMi+7RbCNVxkhLTQFNPgUDQZDjdKr9/bzOfbT4CnFwbHlfi9LhU3li9n7fWHGBgdgLP3dCXOKsJx+bNFN33K+KnTCHxjtsbjl7JMnQc4/0B5IJVqJ9MQc9/HklRIKoNdBkPKV2b/2IFzUqLG1NGRebmoVm8uHyvj5ETM+gq5Mh4qvIXUf7JU0imCMypOcQMva7efhajwh0jO5z4d5c20XRp42voON0q4/69jAMVvnmsoSysXR6NY3Y/fQKO8/1scPsWIas6PLPGxR9GBHhR2o/CgXzIymt8P0FQmBQZRZbQ1F/mLTYpZ2bcto2LIDcjlh8Kfespgr3HrSaFOy7ocEbnFwgEwfGnj7Y2KEzUWD2t062xdGcZU/uk4nF7OD0UHYrjJNKsMFUIzQgEzYqtzsO1c/MpKKttNIXereqAzpp9R7n0mRW81sUJ//wLaU8+QfTYsYFP5KmDzf+D72cjVxWR1M2O9OPL3s9kAyz7OyR3heGzoNskUTvVSmlxYwrgpiFeY8ofUT1GE9VjdKPHp0Sb6RtEk0OLUcFi9H+jhrSw1nXsy5ZSXvAdhjapGNukYWyTiiE1FfnYLqgq8nvYg3km/vF9HXcPNBFnaWRh7rbDqueEMdVEpMSYm8Rb3Cb2zCNDd43syLbDP52ovziVYO5xXYdJvX9ZryuBQNAwByvsrD/gX0AiuH5wHuY+/TZj46L5wpTJ6d/0YB0nEhKX9GjTHJcoEAgAj6px6/wf2F1a6zdrxB9uVae0ys4ty6v48MWXie4VRP207Si8Pgkq9nrXdZyWcKR5vD9HNsAH98CPr8KUt70NfAWtirAwplJiLNw1siMvLt8bUnoFgMUo8/erc4OWE0+KMrOrpNZneygLa7MikZQSj2Y7gG3lHjzFxbhLSvCUlBDb0UVqzzq/zoUB6Qqjsg08taqOP49prJ+ADiXBNYgUBObCbqmo+iaf7aF6i28c0u6M5zCqSwpxVqNXdTKEMjqACKPC1MHtiDAJj5VA0Fy8nl/ot7Yy2HpaHYk1aT1487YhfPPKGlQ/i7RAjhOzQWbq4HaYRH2kQNBsfLalmM2HqvwaUo1FoFVkyiMTeLPYwKxeAU7irIKXx0DVYdBcgSfltsHB1d56+1u/BINI9W1NhIUxBTBrbCeKq5x8tPFw0AaVxSjzz8m9GdwhMejzXN0vg40Hj/lECEJZWOuyzOXTLyMlpr5BpOs62tLZSCv+Cpr/a3hitJlh82zcP9jU+ERdvmmCgjMjymxgUu903l1XxGmZfkF7iw2yzNhuqWc8B0WWWHD7ECY+u5JapydoWRKzQSY3I5aHLxU51QJBc/Ll1pKfU3rqE0o9rcmgUGGr4+7RObywtCAk56BBlshOjGTW2M4hzVsgEITGnKVn3li7TtWZn1/IvWNyGu83uugmqK5vSAWsnffUQekO+HgWXPlCk16zoHkJG2NKkiT+dnUvMhMiePbbPciaC4fmf3qRZgWjIvPfG/oxLCcppPNMyE3jjx9u8ftZsAvroR0SSY3xjSxJkoQSlwKKsUFPRM8UhYmdDfxtpYtuyY18EQ2iE3ZTctuIDny08TCqO3RvscUgMy0vK/hGzQ2QlRjJ+3fncd3c1dTUeQKmF1hNCoPaJzDnxv6/+NwCgaBxapz+62BDqadVNZ1Ku5v7xuRQ43Tz5uoDQRlUZoNMZoKVBbcPFhFogaAZ2VFczd5y3+ykUBSd3arGtztKGddQOm7ZTji4BlTfdWDA2nmPA7a+C+OehMjQ1reCliNsjCnwGiP3junEzb2jeffZ3/Gi6UaO2j0YZQkd7w3cMz2Wu0Z15MKuKWe0wLQYFa4bmMmbq/f79UIGWlgHFAKIaxdQhe//jbLQb24tvxnaSBg3XhQgNyWdU6O5b0wnnvt2T0jeYqMi0Sk1intG5zTJPHJSolnywEheW1XIa6sKcWsatrr684kwymQlRjJzVEcuy00PqhG1QCD4ZTSUKh6qUI0secd6ZEJ3urWJ4e9f7qDW6fFbL2k1KWi6zpV92/LHid2xmsLqlSwQhC3H7C7e/+kQe0prqXF6iI0w0i0thkl90r3tSBrg620luH+horOtTuXDDYcbNqZWPw+qf/Xf4GrnJVj3Glzwm4BzEYQHYfnkjtnxDrf0iWT65WMpr3VR5XBjkCXiI03ERhh/8fizLuzMV1uLKa6qQ/WTI98QFqM31Wtox0bSCrOGgTECXL6ej+PkJMhc18PIf9a66JXixyA0RcHgmUHPSxAcd4/qSG2dh1e/34fTT4TqdIyKRKeUaN6cMQSzoem8xQmRJn59UWfuG5PD19tLWbm7jHKbi2qHm71ltcybPoju6TFNdj6BQBCYuAgjFTZfT3Io9bSKLBFvPZnCfVX/DK7s15ZVBUeZu6yALYercbhUjAaJlGgz04Zmc2W/jEYXfwKB4CRbDlUxZ1kBS7aVIEnUe5dbTQpPfLKVSb3TueOCjuSkRPkcX1zt9En3h9BbpZTV1vn/wO2ATYtA929MBVU773F6DTJhTLUawu8Jruuwbj5cOQdJkkiONpMc3bSFeLFWI4vuHMrVz6+iwubCHYQiQIRRYXCHBP51be/GxS5kGYbc45W79PhKsB/nsZFm3tjUgLy6bPD2HhA0KZIk8btLutI9LZq/f76TCrsLh+v09poQaVLwaDoGWeKt2wYRa/3lBrw/DIrMJT3bcElPr3ertMbJxf9eLgwpgaAFuLxPW55fusdHJjmUelqPppF3Wuq5JEkMy0kKOSVdIBDU5838Qv782XZcHs2vkNPxOqh31x/i441HeOqaXCbkptfbpyH/eagRaL2hgaoPg9T48UHVzjsqvYaZUfSXbA2EXyFG4QpvvVDGwGY9TUa8lc9nXcCQjomYDTLGBnoIWU0KFqO3ZuaVaQOD60zffxpI9fcrnBXN2A4nbdfMWBnnozH1G/aC94sz9B5Qws/OPVe4rHdbVvxuNPNvGcSF3VJIjTETbTGQFGWiR3oMf5rUg42Pj2Ns91ReWrHvrM0rOcqMR9P9escFAkHzcsPghtU6YwZdRfyYGVTlL6Lo2akUvTCdmvWfENHpZEqQIsFluY2nGAkEgjPj9VWF/OWz7Tjd/g2pU1E1HYdb5TeLN/LppsP1PkuJNuMvc/7UCHQwJEQ2YAjVVfus/07n1Nr5BlGMUFcT1FwELU/4PfXXzYf+0wPWHTUFCZEm3pgxmKJKO2/k7+fttQeoqfNgkCU8mk5mvJU7L+jAFX3bEhnKC9KaANe9DgtvbDQ65YPBAm0HwvAHQr8YQUhIksSg9gkMap/Q4D6PjO/GJc+s4Kp+GX7TBZpjTjkpURSU1ZIQ2fC8BAJB05McbWZEpyS+3VHqd7EWqJ7WqMjMGNG+GWcoEJyfrNtfyV8/3+43PT9QM+3fLt5ElzYxJ97h3dJikCXJpw1CSK1STArje6X5n6wxEvTAZQQBa+c1DxitAccRhAfhZUzZymH31zDh6bN62ox4K78f343fj++GW9VwulUiTYZfVvifMxaufhneuw3cTggkhm20QuZguH6BiEqFCSkxFu4ZncOfPtrKGzMGBd3L7JfQMTmKPaW1DMwWxpRAcLb506QerC2soNrhv96hISKMCtcOyKBrG5GiKxA0Nc9+s9uvIRWMlLnLo/LC0j3kdUzi3fVFbD1cRYRJocbp+x0PurG2JJ1Iz/chug2oDZRwnELA2nmDWTTvbUWE16p949vQdQJExLXYFIyKHFwqXzB0mwgzlsDS/4M9XwOSt7DwBBKYrBARD3n3w8AZ+O32K2gxpg3NYvGPB/l08xEmnpZ73RzkpERRUNqweIlAIGg+MuKtLLhtCFNeWo2tzhNUg+0Io8KF3VJ4/LIezT9BgeA8o6TaSf7eoz7bg5UyV3VvDVVZTR1TB2dxYbcUPtpwmMc/2upX2TdwBFpi6uB2DYtSWWKg00Ww41MCOdEbrJ1XTNBv2lnJ0BI0DeFjTB0Xnrj8+ZaeSdPSppc32lRb6pW6LFwBjmNgtEBsJgy4xasAKL40YYlBkXnyip786u2fGNUlpdnrITomR7Haz4tDIBCcHXq2jeWT+4Yz88317Cuvpa6BYnerSUHX4c4LOnD/2E5nJXItEJxvLFx7wO/2UKTMLUaZS3ulMSHXm5o3qU86877fR0FpbVACZMeRJIizmrhzZMfGd8z7FRR8B25bvc2Fs6Lr/f947bzfEw26I+h5CVqe8DGm9n8PshEyB7X0TJqHqBQY+aD3R9CqGJidQF7HJP7zzW7+ML5bs56rY3IkBWUiMiUQtCRZiZF8dv8Ith2u5tb5P1BeW4ckgSxJuFWNjHgrd47swBV9QqynFQgEIbHlULWPwiaEJmXudGvsLD4p5mAxKiy4fQiXP7eSkuo6XGrgGidFloi2GFh0x5CGxSeOkzkI4jKhfDfowfe19J7IBO3yIEHUX7YmwuctcBaFJwSCUHn40q5cPHs5k/tn0Dk1OvABZ0i7BCul1XU43SoWo0j5FAhakg7JkdTWeVj7yIUYFBmXRyPaYmjSvnMCgaBhapz+649ClTKvtNdXzkuINPHp/SO46411rD9QiVvV8GdTyRKYDQrtEsX5ObMAACAASURBVKy8estA0uOCkCqXJJi6GOaMAGcVAWvmTxynQGQyTJ4X3P6CsCE8pNFtR2HXV5B7bUvPRCDwS3K0mVljO/HHD7Y03F+iCTAoMu0SrOwtswXeWSAQNCs/FFbQOTWKhEgzMRYjSVFmYUgJBGeRKIt/n3+oUuZxEb79ImMsRhbcPoRP7hvOtQPaEWFUsJoUoswGrCYFRVNJjzaRmxFL1zZRLFx7gMLyIN/Nce28NfORSd6sq0Ao5pPHWIUAVWsjPIypjW9Dl0vFDSQIa6YOzqK2zsNHGw8H3vkX0DE5SqT6CQRhwLKdZYzqktLS0xAIzls6p0b77QN6qpS5fVc+mtuJrnpwFPxI5Xf1IzsRRpmOjbQ3yUmJ5v+u6sVPj13Eq7cMZFSXZDRdxyhpFFW7WLOvgg83HuH5pQVcPHs5Vz3/PV9vKwnsWE3uDDPzqe51C27FilvxI3VuigJLHAz7Fdy5HGLbBvV7EYQXLZ/md1x4YtKzLT0TgaBRFFniict7cvdb6xjdNYUYSxDepjPgeK8pgUDQsizdVcbT1/Zu6WkIBOctUwa145WV+/CXKheslLmmw+W9Axspm4qquHX+D7g8Gm5VB6n+O96j6Xg0nfUHjnHfwp+4vHc6f7myF4qfNjqqprN0ZylzlhWwqWgMEfIILtLzuZjVJEo1aJKMOa4NSUNvpM3Aq0RLnFZOy//19q/yyoG3G9LSMxEIAtI/K55RnVOYvWQ3j13WvVnO0TElkm+2lzbL2AKBIDiKKu1U2lz0TI9t6akIBOctmQlW+raLY/XeCr+fB5IylyW4pEcbYq2NOz/XH6hk2ry1fuXS/eFwqXy44TB1Ho2nr+1dT83zmN3FTa+spaCsFrvLO14dCosZzmKGn9jPUCZh+FTibts+7huTIxRBWzEtn+YnhCcErYyHLunChxsOsf1IdbOM703zEzVTAkFLsnRnGRd0Tv5lzdsFAsEv5ldjOhFxhoJMJoMcUMq82ulu0JCybVvKkddmceDpyRQ9dxMl7zyOs2grAA63yhdbillwinx7lcPNZc+tZEdx9QlDqiE8mo7TrfHC0gKe/GTbGVydIFxoWWPKXgG7voTc6wLvKxCECYlRZh4Y15nHPmweMYqOyVHsK69FC6H/hUAgaFqW7ixjVJfklp6GQHDek5eTxJ0jO4RsUEUYFR6/rAfd0/30cjqFd38swqP6vm+r175PxTcvETvkWjLufZO2M18lut94HLvXnNjH4VZ59ps96LqOruvc8upaSqqc3jTBIHG4Vd5ee7DBnlqC8OespPkVHCvgre1vsb1iOza3DavBSofYDkxxG+nV5RIhPCFodVw/sB2LfjjIe+sPMb5XGh9vPMw7Px6kvLYOTYeYCAMXdUvlhsFZJEebgx632lXNhwUfYm63kIvffQaDLBFnjuPi7Iu5stOVxJpFypFA0JS4XC62bNnCpk2bsNvtAFgiIigtlBk8qWsLz04gEADcf2EnAOYu2xswFU8CzEaZR8Z3Zcqgdo3uq+s6L67wHVOrs3Fs5Vskjp+FtUveie3WnMFYcwbX27fa6Sa/4Cgmg8yO4hpcfgwp27alVP/wAe6jRcimCIwpHYjNuxZLRg/Aa1D988udXDMg80QNlq5767Nezy9kb5kNh8tDlMVIbkYs0/Oy6ZDcsKiG4OzSrMbUiqIVPLfhOQqOFeDRPKinNC/bfnQ7S3SNVGsqd+39hAntJ4h8UUGrQZElfndxV2a8/gOPfrAFScInpL+7pJbnlxZwQedkfndJV3IaURMqsZUwe/1sluxfgoyMZnJQ7F3XUVRbxJ5je3huw3OMaTeGX/f7NWlRac15eQLBOY/NZmPp0qVs2LABALe7fj+b3pLCy/+dTe/evRk1ahTR0c3XX04gEDSOJEnMGtuZQdkJPPvtHtYfqETT9XoRILNBRgeG5yRx35gc+raLDzjuj/srqXb49rKqO7QD3ePC2nlowDHsLpWXV+5DkSW/hl712vepWvM/Esfdg6V9PyTFgGPfOhy715wwpgCcbpXlu7wR8f+tK+K5b/dQVluHw61yahLMlkNVLPrhIN3SYvjtuC4M75QUcI6C5qVZjCld15m7aS6vbH4Fp+r0u4+GhlOC/Y4Snsh/grVH1vL40MdRgmi+JhC0NEWVdh58dyN1Ho2GMv2Od23/ensJ3+8p5+VpA8jr6PvQ21W5ixlfzqDGVVPP4XAqx79HX+77ku8Pfc/L416mW2K3prkYgeA8o7y8nPnz52O329E0P506AQUVjwd++ukntm/fzvTp00lJETLpAkFLkpeTRF5OEgcr7Cz84QC7S2qpcXqIjTDSPT2G6wdlkhJtCXq8fWU2/GXUq45qZGtMUA2BAXYW11BeW+ezHgglwmVzqbywdA8fbDjEV1tLGozAHVcV3HDwGLe9/gO/Hts5YF2YoHlpFmNq3pZ5vLKlYUPqdBweB5/v+xxZknl86OMiQiUIaypsLq5+fhVlfh6c/tB1r+dqxvwfWXjHEHpnxp34rKimiOmfT6fGXRPUuTU0ql3V3PLlLSyauIismKwzvQyB4LykqqqKV155BYfDEXDfpUuXUlFRwVVXXcW8efO48847iY8P7O0WCATNS2aClQcv/uVpuLV1HlQ/DhUlIgbNXo2uqUEZVNXVNoySTN1p20OJcAGsP3AMw6EqnG7/Tp7Tcbo1Zn+9G5NB5pZh7YM6RtD0NLkxtaV8C3M2zvExpCpXVFL+ZTmuUheKRSGmfwypk1NRIr03qVN18tm+zxjWdhgXZV3U1NMSCJqMB97ZQIXN5deb1VhetMOtMv3Vtax9ZCxGxav9cv9392Pz+Cr3Bfq+ONwO7v3mXj664iPhfBAIQmDRokXU1dVf8mzevJn8/HzKy8sxm820adOGESNG1Nunrq6OBQsWcPfdd4vvnEBwjhBpVlBkGdT6USBz265IBiP2XflEdh3ewNEnMWgqLl316U0VaoTreNTpdAKtLf7+xQ6GdEikW1rjYhuC5qHJjalXNr+CS3XV21b+eTlln5eRcVsGUd2jcFe6OfzGYQqfKqT9I+2RDd6FpcPj4MVNLwpjShC2HKlysKrgKG4/D7tg8qJdHo2vt5Vwaa80th7dyoHqA2h6fQ9UMN8XDY0SewkbyzbSJ6XPWbl2gaC1U1xcTFlZWb3Uvvz8fFauXMnEiRPp2LEjiqKwZ88eduzYgclkOrGfruscO3aMQ4cOkZGR0RLTFwgETUx2YiT+uh/I5kjihk+lYskcJFnB0r4vkmzAWbgB54FNxI++td7+GW2TKDxqA6en3vZQI1z+CGZt4VZ1Xlq+l6evE+uBlqBJpdErnZUsP7QcjZMvKtWhUvpBKek3phOdG41kkDAlm8i8OxNXuYuqVVX1xthXtY/dlbubcloCQZPxRv5+f43YT+RFJ1w0E2uXPGSTBUkxYM0ZXO+ha3OpvLCsAIDXt76OS6vveAjl++L0OHlt62vNc6ECwTnI6tWr8XhOLnacTiffffcd48ePp1u3bphMJhRFoUuXLowbN87neI/HQ35+/tmcskAgaEYGtU8g2uK/oW/MoKuIHzODqvxFFD07laIXplOz/hMiOtVP2Ys0KVw/MBPVj5P11AjXmRDs2kLVdD7dfIRqp6+YhqD5aVJj6rN9nyGfNqR9tx3NrRHTv37oUbEoROdGU7u1tt52j+Zh8a7FTTktgaDJeHvtAVyqby5zKHnRO4tr2F9Rzdf7v/aJSoXyfdHRWVa0DIcncO2HQHC+o6oqW7bU7w1XVFSEx+OhW7fgxFx0XWfHjh24XK7AOwsEgrBHkiRuH9G+wR5WUT1GkzZtNu0eeJfMe98k5Zo/Ycmo/7yIshi4fmAmiZEmn+NPjXDZd+WjuZ3oqgdHwY9Ufjcv4PxCWVvIksRnm44E3E/Q9DSpMXWw5qBPrZRaq2KIMiApvnFUQ6wBT239kKiqq+yv3t+U0xIImgRN0znmR0IVQsuLNsgSy3YX4u2Gcdo4IX5fDLKBSmdlUPMXCM5n/AlO2O12rFYrshz8q1CW5RP9qAQCQetn8im9nUIlwqhw7+gcFEXmrpEdsZp81wDBRrj8EcrawuFWOVQpnKstQZPWTNlcvoX0SpSCp9aDruo+C0RPlQdDlO8UHG5xMwjCD5eqIeE3yy+kvGiHW2XO8u24k3UfeyrU74uEFLRqpkBwPuN2u32EI6xW6wl59GANKkmSRGRKIDiHiI0wMv+Wgdz0ytqADYFPJcIoM657KjcO8arqXtG3LX/5bLvffaN6jCaqx+iQ5xZqzVVNnSfgPoKmp0kjU/EWX8lYa44VySBRva663nbVqVKzqYbI7pE+x8SYhRqJIPwwG2QkP9EkCC0vOtJk4KnJQ/GXVRDq90XVVaKNopmoQBAIs9ns01MqIyMDg8HAjh07gh5H0zQsluD72AgEgvBnQHYC86YPJNKkYAgiSmU1KUzMTedf1/Y+4aSJNBt44cb+WIyhLa0VScJk8H9MKGsLCYi3+q//EjQvTWpMdU3oitVgrbdNsSqkXJHC4TcPU7OpBt2j4ypzcfD5gxgTjMTlxdXb36yYyU3KbcppCQRNgiRJdEj2Nf4htLxol6rRo00qkUbfsc7k+5JgSWi6ixQIzlEsFks9db7j20aNGsVnn33Gjh07cLvdqKrK7t27WbJkid9xFEUhMtL/c0AgELRehnZM5ItZF3D9oEwijIpPyp5BlrAYZHpnxPLv6/rwj8m5GJT6y+iRnZP51zW9gzaojIpEWpyF6wZkYlJ8jwllbWE1K3QV0ugtQpOm+Y3NGssTq5/w2Z48PhklUqF4UTGuUhdyhExMvxgy78xEPu2G03Wdqztf3ZTTEgiajLtGduSxD7dgc/mmAsQMugo5Mp6q/EWUf/IUkikCc2oOMUOvO7GPIktM6JVGbISZG7vfyIubXqROrd/zJtjvi0k2MaXrFJQzlFsVCM4nZFlm8ODBrFy5sp6iX15eHlFRUSxfvpz33nsPk8lEeno6I0aMoKCgoN4YiqLQv39/FEV85wSCc5HMBCt/vqIXfxjfjY82HOaHwgoq7W4ijAqZCRFM7p9JTkpUo2NMyE0nPS6Cv3y6nc2HqtA03aeditWkoOtwZd+2/O7SrthdHt758aDf8YJZWwAYZZkLu6b8sl+A4IxoUmPKpJiY3GkyC3YswK3VL9RPGJlAwsjGPegSEsPaDiMxIrEppyUQNBkTctP444dbGvw8UF60SZGYMcLbpfzqTlczd9Ncv/sF830BuLbLtQH3EQgEXvr378+KFSt8tufm5pKb65sRkZmZWe//kiQxaNCgZpufQCAID6wmA9cPasf1g9qd0fF928Xzv5l5FJbbmL+qkBW7y6hxejDIEolRZv4/e/cdHlWVP378fe6dmUw6BJJA6CEUAbGAdARE0bWXtYNir7u6xd39ubvqurrfXdddd+2KCgo20LX3RkeaItJLCC1ASG+TTLnn98dMMMlMkpmQRvJ5Pc88mbn13MnMmfs59cpRvbjgxB7ERvlvwxOj7Yzql8SS7bkhj9fQvUWUzeCacX2CaspEy2jySXunD5nOgm0LgoKpcESZUdx6wq1NnSQhmozTbnLbpP48vXBnRB1VARw2gxN7dWZoWiIAXaK7cH76+XyY+WHEg0g4TSdn9DmDlBgphRIiXPHx8QwbNoyNGzfWqJ0Kh81mY+DAgXTq1KnhjYUQAujbNZYHzh8a1rZ3nz6Q1Vn5VHiCp19piN00jgyEIVpek4ew3WK78eTUJ3GakXXQdZpO7h97P0O6DGnqJAnRpO48LYMpg5PrnJciFIepSEt08vw1I2osv3fMvQztOpQoMyrsY0WZUQxMGshfxv0l7H2EEH7nnXceqamp2GzhlyWapknXrl256KKLmjFlQoiObESfzjxw3tCIB7CItpu8fP0ppMTLwDitpVnqA0/pdgrPnP4MsfbYBm8S7YYdp+nkofEPcW7/c5sjOUI0KaUUT1x5Mhed1INou0lDA//EOEyGpCXw3p0TgmZatxt2nj/jeSb0mEC0LbrBc0fbohnVbRQvTnsRuymj9ggRKZvNxrXXXkufPn2w2xv+DjkcDnr16sV1110X1vZCCNFYV4zqzd8uOh6n3cDewM1FlM0g3mlj3o2jGdFHBqJqTar6bPC1jRw5Uq9Zs6bRB8915bJg6wJe3fwqXu2l0luJV3sxlYnTdKLRXDLgEq467ip6xvds9HmEaC3f7SnghcWZfLUlB9NQuL0WWmvsNgOtYUj3BG6b3J/TBqfU25ZZa83qg6uZvXE2qw6swjRMPD5/U1m7YcenfYxIHcF1w65jTPcxQfPlREIptVZrPbLRB2gjjjZ/Eh2bZVns2LGDpUuXkp2dDXCk6V9VrVW3bt2YMGECAwcOjGhiX9F47SF/krxJHK2s3DJmL9vFgrX7UPjnp7S0f0TBKLtBlM3k+vF9uXJUb7rEhd+yRTRefXlTswZTVbyWl6X7l7KzcCcl7hLiHHH0jOvJlN5TImreJERblV/m5ustOeSXVeK1NInRdsamdyE9uf5Rf0LJKc9h2f5lFFYWApAYlci4tHF0i+3WJGltDzcrIDcsounk5+ezfft2ysvL0VoTExPDgAED6NJFBkNqae0hf5K8STSVCo+PzzYeZF+Bi7JKLwnRdgalxnPqwGTMMObDEk2nvrypyQegCHkSw8bkXpOZ3GtyS5xOiBaXFOvg5yOapnY1JSaFiwZI3wwhWkpSUhKjR49u7WQIIUQNTrvJBSf2aO1kiAZIuwUhhBBCCCGEaAQJpoQQQgghhBCiESSYEkIIIYQQQohGkGBKCCGEEEIIIRpBgikhhBBCCCGEaAQJpoQQQgghhBCiESSYEkIIIYQQQohGkGBKCCGEEEIIIRpBgikhhBBCCCGEaASlta57pVKHgd0tlxwhRAvoo7VObu1EHC3Jn4Rol475/EnyJiHapTrzpnqDKSGEEEIIIYQQoUkzPyGEEEIIIYRoBAmmhBBCCCGEEKIRJJgSQgghhBBCiEaQYEoIIYQQQgghGkGCKSGEEEIIIYRoBAmmhBBCCCGEEKIRJJgSQgghhBBCiEaQYEoIIYQQQgghGkGCKSGEEEIIIYRoBAmmhBBCCCGEEKIRJJgSQgghhBBCiEaQYEoIIYQQQgghGkGCKSGEEEIIIYRoBAmmhBBCCCGEEKIRJJhqR5RSk5VS+1p635ailOqtlCpVSpmtnRYhRP2OpfxIKbVQKXVjS51PCNF6JG8STU2CqXoEbtyrHpZSylXt9dXNeN6ZSqmlzXX8pqCU0kqpjGY+R5ZS6vSq11rrPVrrOK21rznPK0RbJPlRaEqpvoH8yFZr+Ryl1EOtlS4hOgrJm0KTvKnjsDW8SceltY6req6UygJu1Fp/WXs7pZRNa+1tybQJIToWyY+EEG2R5E2io5OaqUaoquZVSv1eKXUQmB2qhKR67Y1SKkop9ahSao9S6pBS6lmlVHQjzn2dUmqzUqpEKZWplLolxDb3KqVyAzU7V1db3iRpCHG+B5RS85VSrwTStVEpNbLa+j8opXYG1m1SSl1Ua/+bql3TJqXUyUqpuUBv4INA6dbvqpfyKKUuV0qtqXWcXyml3m/OaxWirZH8KKx0zlRKLQ2cr0AptUsp9bM6tu2ulFqvlLon8HqhUuqvSqllgev8XCnVtdr25wfyvMLAtscFll+nlPqg2nbblVILqr3eq5Q6MfBcK6VuDWxTqJR6SimlmuO9EKKlSN4UVjolb2oHJJhqvG5AEtAHuDmM7f8ODAROBDKAHsB9jThvDnAukABcBzymlDq5Vrq6Bo5/LfC8UmpQpGlQSj2tlHo6gnSdD7wBdALeB56stm4nMBFIBP4CzFNKdQ+c51LgAeCawDWdD+RprWcAe4DzAk37Hql1vg+AQUqpAdWWXQW8Fum1CtEOSH7UsNHA1kB6HgFerH1ToJTqBywCntRa/7PaqqvwX18K4AB+G9h+IPA6cDeQDHyMvwDIETjORKWUoZRKC+w3NrBfOhAHrK92jnOBU4DhwGXAmUd5vUK0BZI3NUzypmOd1loeYTyALOD0wPPJgBtwVls/E1haax+N/4uogDKgf7V1Y4FddZwr6Fj1pOtd4K5q6fICsdXWzwf+3FAaAvvui+D90EBG4PkDwJfV1g0BXPXsuw64IPD8s6r01/eeB173DZzXFng9D7gv8HwAUALERPp+y0Mex9pD8qMa56yRL1RbPgd4qNo17Ki2LiawT7fA64XAvwPv65W1jrMQ+FO117cDnwae/xmYX22dAewHJgde7wVOBq4AngdWAYPx3/y8X+t/M6HW+/SH1v6cyUMekT4kb6pxTsmbOshD+kw13mGtdUWY2ybj/4KsrVbYoICIR6ULVP/ej7/UxAgc98dqmxRorcuqvd4NpDVlGupwsNrzcsCpAu2jlVLXAL/Gn7GAv9Sjqiq6F/6aq8Z4DfgX8CD+0pl3tdblSqkUmvdahWhrOnJ+VNUHw17tedVrT7XXR/KoQD4B/ryoytXADuCtEOeonb9V7ZeG/5qqjmsppfbiL8kGfwnwZPw3iouAQmAS/puzRWGeQ4hjmeRNkje1e9LMr/F0rddl+L+AACilulVblwu4gKFa606BR6Ku1mkzHEqpKOBt4FEgVWvdCX/VbfXq4M5Kqdhqr3sD2U2VhkgppfoAs4A7gS6BNG+olua9QP86dq/9Htf2BZAcaNt7JT818WuVaxWiFXXk/OgA/huTvrWW96PazUQYHgik6zUV/vQL2fibLwEQaJrTC38JMPx0wzIx8HwR/huWSQTfsAjRHkneJHlTuyfBVNP5ARiqlDpRKeXE/+EH/CUC+AOKxwK1Jiileiil6mt3qpRSzuoP/O1ao4DDgDdQ8jItxL5/UUo5lFIT8bd1XdDINDSFWPyZ6eHAOa8DhlVb/wLwW6XUCOWXEQjAAA4B6XUdWGvtARYA/8TfJvuLwPLWulYh2ooOkx9p/1QJbwMPK6W6KKXsSqkr8Tc3/iSCQ3mAS/HnWa8opcL5fZwPnKOUmqqUsgO/ASqB5YH1i4ApQLTWeh+wBDgL6AJ8H0HahGgvJG+SvKndkWCqiWitt+FvbvYlsB2oPffB7/FX036rlCoObDeIuo3DXzpS+/FL/F+SAvxN296vtd/BwLps4FXgVq31lkjToPyj1zxb/1U3TGu9CX9TvBX4g6PjgWXV1i8AHsZfq1SCv11zUmD1/wF/Uv4RZH5bxyleA07HnwlWr0aP9P0Wot3ogPnR7UA+/k7TOfhrws/RWh+qZ58gWms3cDGQCrzU0E2L1norMB14An/J8Xn4B81xB9ZvA0rx36igtS4GMoFlWubLEx2Q5E2SN7VHSuuGWlIJIYQQQgghhKhNaqaEEEIIIYQQohEkmBJCCCGEEEKIRpBgSgghhBBCCCEaQYIpIYQQQgghhGgECaaOklJqjlLqocDziUqprS10Xq2UymjiYx65lpbct6Uope5VSr3Q2ukQoqVI/nT0+7YUyZ9ERyJ509Hv21Ikb2pYhwimlFJZSimXUqpUKXUo8OFt8glctdZLtNYNDr+tlJqplKo9HGiTUUotVErd2FzHP1rNff2Bc0xWSu2rvkxr/TetdZt9X0THJPlT2yL5kxB+kje1LZI3tV0dIpgKOC8wg/XJwEjgT7U3UErZWjxVQggh+ZMQom2SvEmIBnSkYAoArfV+/DNPD4MjVb53KKW2459ADqXUuUqpdYHJYpcrpYZX7a+UOkkp9Z1SqkQp9SbgrLauRkSvlOqllPqfUuqwUipPKfWkUuo44FlgbKC0pzCwbZRS6lGl1J5ACdCzSqnoase6Ryl1QCmVrZS6vrHXr5RaoJQ6qJQqUkotVkoNrbVJV6XUF4HrW6SU6lNt38GBdflKqa1Kqcsam45aacpSSv1WKbU+kK43lX8Wc5RSnZVSHwbew4LA857V9k1SSs0OvC8FSql3lVKx+P/HaYH3uFQplaaUekApNS+w3ydKqTtrpeMHpdTFzXmtQtRH8ifJnwL7Sf4k2hTJmyRvCuwneVMIHS6YUkr1As4Gvq+2+EJgNDBEKXUS8BJwC9AFeA54P/CFdQDvAnOBJGABcEkd5zGBD4HdQF+gB/CG1nozcCuwQmsdp7XuFNjl78BA4EQgI7D9fYFjnQX8FjgDGACcfhRvwSeBY6QA3+Gf+bu6q4G/Al2BdVXrA1+yL4DXAvteATytlBpSx/UXKqUmRJCuy4CzgH7AcGBmYLkBzAb6AL3xz2z+ZLX95gIxwNBAuh7TWpcBPwOyA+9xnNY6u9b5XgeurJbeIYFzfBTptQrRVCR/kvwpQPIn0aZI3iR5U4DkTaFordv9A8gCSoFC/F/Qp4HowDoNnFZt22eAv9bafyswCTgVyAZUtXXLgYcCzycD+wLPxwKHAVuI9MwEllZ7rYAyoH+1ZWOBXYHnLwF/r7ZuYCDdGXVc70LgxjDel06B4yQGXs/Bn2lVrY8DfEAv4HJgSa39nwPur7bvQ2H+P2pffxYwvdrrR4Bn69j3RKAg8Lw7YAGdQ2x35H9RbdkDwLzA8/jAe94n8Pph4KXA83qvVR7yaMqH5E91vi+SP0n+JI9WfEjeVOf7InmT5E01Hh2pneuFWusv61i3t9rzPsC1SqlfVFvmANLwf3n268AnJGB3HcfsBezWWnvDSFsy/hKCtUqpqmUKMAPP04C1YZyzXoESn4eBSwPntAKrugJFgedH3gutdalSKj9w/j7A6Kqq9QAb/tKNpnCw2vPywDlRSsUAj+EveekcWB8fuJZeQL7WuiDSk2mtS5RSH+EvOfkH/pKWmwKrm/tahahN8ifJn46Q/Em0IZI3Sd50hORNoXWkYKo+1b/ge4GHtdYP195IKTUJ6KGUUtUyhd7AzhDH3Av0VkrZQmQKutbrXPxVsEO1v11ybQfwf/ir9K77Uup1FXAB/qruLCARgYPVtAAAIABJREFUKMCf+VQ5ch7lH7UnCX+J0l5gkdb6jEaeu7F+AwwCRmutDyqlTsTfzEAF0pSklOqktS6stV/t9ziU14H7lVKL8bff/iawvLWuVYhQJH/6ieRPkj+JtkPypp9I3tSB86YO12cqDLOAW5VSo5VfrFLqHKVUPLAC8AK/VErZAx3uRtVxnFX4v8h/DxzDqZQaH1h3COgZaEeM1toKnPcxpVQKgFKqh1LqzMD284GZSqkhgdKG+8O4DlvgnFUPO/7q2UogD39pzt9C7He2UmpCIG1/Bb7VWu/F34Z5oFJqRuDa7UqpU5S/U2hzisefWRYqpZKodu1a6wP42zE/rfydLe1KqVMDqw8BXZRSifUc+2P8JSkPAm8G/g/QetcqREMkf5L8SfIn0RZJ3iR5U4fNmySYqkVrvQZ/leWT+EsedhDo0Ke1dgMXB17n428f+r86juMDzsPfIXIPsC+wPcDXwEbgoFIqN7Ds94FzfauUKga+xF+qgNb6E+A/gf12BP425Bn8X6Sqx2zgFfzV3PuBTcC3IfZ7Df+XLh8YAUwPpKEEmIa/ajcbf9XyP4CoUCdX/lFgJoaRzob8B4jGXwL1LfBprfUzAA+wBcgB7g6kdwv+0pNM5e/QmVb7wFrrSvz/v9PxX3fV8oiuVYiWIvmT5E+SP4m2SPImyZs6ct6kajZhFUIIIYQQQggRDqmZEkIIIYQQQohGkGBKCCGEEEIIIRpBgikhhBBCCCGEaAQJpoQQQgghhBCiEeqdZ6pr1666b9++LZQUIURLWLt2ba7WOrm103G0JH8Sov1pD/mT5E1CtD/15U31BlN9+/ZlzZo1zZMqIUSrUEo1ahb4tkbyJyHan/aQP0neJET7U1/eJM38hBBCCCGEEKIRJJgSQgghhBBCiEaQYEoIIYQQQgghGkGCKSGEEEIIIYRoBAmmhBBCCCGEEKIRJJgSQgghhBBCiEaQYEoIIYQQQgghGkGCKSGEEEIIIYRoBAmmhBBCCCGEEKIRbK2dACE6Aq/lZdG+RXy480NyXbn4tI/Ozs6c3vt0zup3FtG26NZOohAdktdn8fWWHN75fj85JZX4LE1SrIMzh6Zy/gk9iHaYrZ1E0cFtOVjMvG93syOnjHK3l3injeN7dGL6mN707BzT2skTosOTYEqIZlTuKefljS/z6uZX8Vgeyr3lNdavObiG/1v1f1zQ/wJuGn4TKTEprZRSITqWcreXWYszmb08C4/PoqzSV2P9t5l5PPD+Ji4Z0YNfnDaA1ARnK6VUdFSfbjjA41/tIDO3FI/Xwqd/WrdqVz6zl+3i5N6dufv0AYxO79J6CRWig5NgSohmcrj8MNd/dj0Hyg5Q6asMuU1VcPXWtrf4NOtTXpj2AoOSBrVkMoXocHJKKrjy+W/ZV+Ci0muF3Kbc7Q+u3li1lw/XH+C1G8cwJC2hJZMpOijL0jzwwUYWrNmHy+MLuY3HpwHNisw8vp9dwO/PGsx14/u1bEKFEID0mWrzLG1R7C6m3FOO1rrhHUSbUOwuZsYnM9hbsrfOQKo6r/ZSWFnIzE9nsqd4TwukUIiOqbjCwyXPLGd3XnmdgVR1XktTWO7hsudWsCu3rAVSKDq6hgKp2io8Fo98upV5K7KaNV1CiNCkZqoN8vg8fLXnK17a8BJb8rdgM2xY2sJQBlN7T2Xm0JkM7Tq0tZMp6vHHJX8kpzwHn675Y1iwpIDcz3Jx57gxnSYJIxJI/XkqZqy/X0a5p5xbvriFjy/+GKVUayRdiHbtN/N/4FBRJV4ruHCqbNNCile/iydvH4YjGntKOonjLsPZcyhlbi8zXlzJ4numYBjy3RTN47ONB+sMpOr7fLo8Ph76eDMj+iZxXHepQRWiJUkw1ca8te0t/rXmX2itKfP6S0E9lgcAn/bx+e7PWbh3IWlxafxz0j8Z2Hlgaya3SbjcPj5Yn82P+4ooLHcTG2UjPTmWi07qSXJ8VGsnL2IHyw6yPHv5kf9bldxPcjn8yWF63tiTuCFxeAo8ZM/NJuvRLPr9sR+GzcDCIr8in5UHVzKm+5hWugIh2qdDxRUs3nYYty+4Rqp41TsUrXyLLtPuwNnvZJRpw7VrLa7tK3H2HIrWUFDmZvnOPCYM6NoKqRcdweNfbQ8ZSDX0+QR/079ZizP59+UntnSyhejQJJhqQx5b+xivbX6NCl9FndtY2qLCV0FmUSbTP57O01OfZmS3kS2YyqazO6+MWUsyeXvtfpT6qY8CQJTN4NHPtzFxQFdun9yfEX2SWjGlkZm/dX7QMp/LR867OfS4oQfxw+MBcCQ76HV7L7bds42i5UV0PrUz4O9HNXvDbAmmhKhFa83qrAJ25ZZSVukjLspGv+RYRvbpHFZN7twVu0MutyrLKFz6Kl3OvpuYQeOOLI/JGE1Mxugjr8vcPp5bvFOCKdEsth0qYefh0qDl4X4+fZbmox8P8MAFQ0lw2lskzUIICabajJc3vtxgIFWby+vijq/u4NWzXyWjc0Yzpq7pfbM1hzte/Q631wrZ3KaqL8PXm3NYviOPO07rzx2TM46Jpm/zt87HbblrLCvfXo7lsUgYUbP5hek0iR8eT+nG0iPBFPhH+SuoKKCzszNCdHRFLg9vr93H80syKXF50IDXp7GZCgUkRNu5eWI6l4zsWe9N5Ksrd4fsJ1W5fwva6yZm4NgG07JyVz55pZV0iTv2as1F2zbv2914jvLzaRiKD384wFWjezdHEoUQIUgw1QbkunJ5/PvHcfvcQesa6mPj8rr48/I/8/o5r7d0shtt8bbD3DZvLRWehjt/a8Dl8fHU1zvxejV3n9EEzRo9FVBZDLYoiEqAJgzQPJaHYndx0HJfqQ9bnA1lBp/LlmjDtdtVY5nDdHCo/JAEU6LDW7s7n5mzV+P16aDmT1WV2WVuH498tpV/f7mNl68fxcm9g783Xp9FocsTtBzA5yrGiElAGQ3PKeUwDQ4UVUgwJZrcjpzSGsOfV4nk8+ly+9iTLwOlCNGSJJhqAxZsXYC/fLWmcPrYaDTbC7aTWZRJemJ6K6Q+MgeKXNxaRyDVUOfa5xbv5ITenZgyqBFzMVUUwbrXYcUTUHwATDtYPjBMOP5SGHsHpBx31NdX4a3ANEy8lrfGcjPOxFvqRft0UEDlLfJiiwv+Krq8rqBlQnQk32bmMXP2airCGNXM5fGBB66etZKXrx/FqH5JQetNpfCGGBXVjE7AKi9GW74Gb1hrN0kWoqnU9bmK5PMJUOzyNriNEKLpyNDorcxreXl1y6tBw2dX9bFJm55G/PB4lE0d6WPjznVTtLzop20tH/M2zWvytGmt2ZRdzBebDvHeuv18veUQ+wrKG96xHnOWZeGto/N3/lezSBxzGT3vnEeP22YTf/LZuLavPLKNy2Px2BfbIjuhZcHnf4ZHB8JXf4GifaB94K0Ay+P/u+41eH4KzDrNv/4oxNhi8FnBP4gxGTEom6J4bc1aK1+Fj5L1JcQOia2xXKOJtddcJkRHsje/nBteDi+Qqs7l8XHdnFXsL6xZGBHrsOGrY3qJqB6DUTY75dtWNHh8rSEuSsohRdOLd4b+XEXy+QToHOtoymQJIRogvwitbHPe5qBaDIisj41Xe/l89+fcN/a+JkmTy+3j/R/288zCnRwqqcRmKCytMZTC7bU4oVcnbp2UzqSBKZgRDBHs9lq8unIP7lrtGMLtXAuw7WAJO3JKyUiJa/iEPi+8cRVkLfEHTXXRPvC6IHsdPDsBrvsUUgaHfV3VmYZJakwqB8sP1lweY5JyYQrZ87IxnEaNmkZ7kp1O4zrV2N7lcfHcD89hM2wkxyQzuedkRqSOOCb6jAnRFJ5fnEllHU2B66vFBn9eM2txJg+c739tlZVRtmIFKbqSQwQ3zzOiYuk04Wryv3gWZZg4+52EMmxUZK2jYs96Ok+5/si2Hp9Fr6ToZrhi0dEN75nIysz8oNEmI/l8xjpMBqaG8fsohGgyEky1soLKgpBN/CLtY1Pmbpo20quz8rl+zmp8lq6zycGqXfls3F9ESoKT124aTffE8G4svtp8CE2IwSYi6FzrtTSvrMjiwQuG1b+h1vDenZC1GDxhNpfTPnAVwpyz4bblEN8tvP1quWboNTzx3RO4fDXPm3x2MmasycE3D+LOcWNEGyScnECvW3ph2GtWEltYfL77cwAUigVbF9ApqhMzh87kogEX4bQ5G5U2IY4FLrePt9buCzk4TbhDRM9ftYeb89bgXbII17ofcJ4wnBnHn8lT+U5c3uDjJoy6GCO2M0Ur3iT3w0dRjmiiUjNIGHv5kW0MBWcN7Ua8jJQmmsH0MX14YcmukOvC+XwCoBRnDWvcb5cQonEkmGpluo5mJ5H2sQkVpERqyfbD3PzKGlxhDAxR5vaxJ7+cs/+7hA9+MYGenWMa3Gfn4VJcIQK0SDrXei1/08MG7VsNm9+rEUj1/U8J5R7YdVccsQ7/e/rCd27mrfewcGZVkzoNriL44s9w8ayGzxPCBRkX8N/v/htyXdKkJJImRTbMu0ZT7i2n3FvOv9f+m/nb5vPCtBfoEt2lUekToq374IfskOPCRFKLrStcfLy9kMuuvJIejz+BGRfLNRUennjoS6gjv4wbOoW4oVPqTFeUzeSmU9t+31RxbOqeGM2ofkks2Z4bcn1Dn0+7qbhyVC+ibA3/lgohmo70mWpliVGJIZdH2scm2nZ0zU525JRwy9y1YQVSR9JiaYpdXi5/bgXl7oY7vBa6PIQoaK7RuTYcJRVhdK5d/njIGimfhv+uDB41sQbthU3v+wetaIQERwIXZlyI02z62qMKXwVZRVlc/fHVIUcNFKI9WLkrL2TNeCS12C5bFDtGTCHhjDMw4/z5ZYLTzhWn9CLaHvlPn91UDO4ez7AeofNsIZrC3acPwNmIzyeA3TS4bny/Jk6REKIhEky1ssFJg7F0cABTvY9NyfoStFfjPuxm79N7g/rYGBiMTWv45qI+//58W50dvcs2LeTAy3ez598/Z9+TMzg0/34q9m0EwKc1BeUe3vl+f73HtyyNL1QkReSda/cVlPOrN9fx7y+28fbafazOyienuOKnWr6yPNj+OaFKn+8Z5+DR5ZUUVjRQk6cM/+h/jfT7Ub9nSJchRJmNGz65YEkB2/+0nY03b2TLL7eQ/XI2vjL//8ervRwuP8xvFv6m0ekToi0rKD/6IcwB8suCC07+dO4QhvfshNMW/s+fqS26xkUxe+YpYe8jRGOM6JPEn84+LuKAymk3eG7GCHp0kv58QrQ0aebXypw2JxcNuIj5W+fjsWreQITbxybKFsXMoTMbnYb8MjdfbckJWWsUTv+EcrePZxfu5KpRvVFKYVmarLwyftxfxIb9RazfV8Sm7GLsNoVpqKCgKpLOtaZSjEnvwviMruzJK2PJ9sPsXlnOnrxyyt0+eiVFc5ljOTMsI0Q3cxiZZjK5r41Hl1fy0Gn11Bx5yuGH12HMrY15S7Ebdp474zl+vfDXrDm0JqJhzsMZEt9tufk+53uyirLom9i3UWkUoq2KtocOliIdIjraEbyN3TR4+fpR3PHad6zYGboGrHZaUsqLeNaWSaeYqYC/efaa3QW8tXYf2YUu3F6LpFgHpw5M5oIT04hxyE+raLzpY/titxnc//5GPF4r5NxTVUzlLzacdc1IJg5IbrE0CiF+Ijl+G3D14Kt5a9tbIdeF08cmJTqF47se3+jzv7l6T4ghMCLrn3CopJK73lhHTkkFG7OLSXDaOb5HIsf3TOSOKRkc3yORaIfJyX/9IuTNS7ida+02xW+mDWJIWkLQMUorvezJK0etXIGZW3dTwAenRDH+pTLuGt3A8LGu/PrXN8Bpc/Lk1Cf5Zu83zF4/iy25G7BMO54QozdWqRoSv8cNPYgfHg9wZEj8bfdso2h50ZFRHH2Wj1c3v8ofx/zxqNIpRFvTOykGm6GCBqCoXosdO3hCvcewGYpedfTldNpNZs0YyZebD/Hc4kw27C/CsjSewPkM5e8flZoQxe2TMzinbwwHrryCvMEZfJo8jGcW7SS/zI3L46N6t9dF2w7z4AebuOikHtxxWobUEohGu/yU3pzUuzMvLMnk/R+yMZSq8dsZ6zBRSnHFKb1YvP0wB4vqGbFWCNGsJJhqA3ol9OKCjAt4f8f7VPgiyxCdppP7xt53VENmL9p6mApvcFPDSPoneLwWheVubpvsD5yS6pjn4rKRvXh15W48IYraGupcC9C3S2zIQAr8c78MSUuA1Dj/3VAdBc7DUkzOHWjj70vdHJdcT1OKOgYHiYShDKb2nsrUglx2F7v5avi55Lpy8fg8vL397aDayEiHxH9v53v87pTfYTdldDHRfvx8ZE9mL9sVFExFVIttKC4Z0bPOcxiGYtrQbkwb2o1duWV8vvEgh0sq8VgWXWOjmDgwmRN6Jh7JWzv/53Gue+JrtndVuOqoKqi62Z2/Zi8f/JDNyzeM4uTenY/27RAd1MDUeB75+Qncd95QPlqfze68coorPHSOcTAwNZ5pQ1OJsplsyi5mxosrmTggmW6JMtKrEC1Ngqk24t5R95JTlsO3B74NO6CKMqO4b+x9jOo+6qjOXeg6+v4JGhjcPYFJA+tvZnD9+H68uXovHl9kE3GCv7nN3acPCGPDzmDYwVf3QBN/mezk5OdK+c3Yevo0OTvVvS5Sm9+nz5Cfc/2wqwDIr8jnfzv+F7RZpEPia63Jr8gnNTa16dIqRCvrnxzHoG7x/LAveBCYcGuxh6Yl0K9reBNf9+sayy2T+te53u21uGlxPluS+gbNkxeK19KUVHqZ/sJKFtw6lqFpMmiFaLy4KBuXn9K7zvVD0hKYPqYP977zIy9eO1LmIxSihckAFG2EaZj897T/csmAS3AYDhxG3U3QYkwnsfZYHpv8GOf1P++oz203Q38MIh1lL5wO3b27xPCvy4ZH3Lk22m5yxSm9OGtY94Y3Tp8M9TSlA8hIMrh8qJ3HV9URcNmcMOT8iNJYRVsaX5kHb34FvjIPuqIUdi2GgWcd2abMU4ZNBZdlVB8Sv7ZQQ+IbyqDM0zRzjAnRltw+JYOYEH2ewF+L3f3a/9D712/T6855pFz6AM6exx1ZH+MwuX1yRpOl5R+fbmHTgWLcIeKo+gboKXf7mPHiKiq9kRceCRGJO6ZkkF3oanAwKCFE05OaqTbEUAZ/GP0Hbjj+BuZvnc9rW17DY3kwlf+GwmN56I2d6zsNZ9rPnsBhNtDnJ0zdE538uD+4BDiS/glOu0FyQnjNC84+Pg2fpbnnrfW4vVbIgS+qi7abXDW6N388+7j6N6yS2AP6jIPMb+rd7L5JUcxdH7pWzuP18UPXCxihddilfJ5cF6XL9lO+9hDa8u+nLY0yLGIT7iauMgZboAtHjC0Gnw6+wao+JH7iqJ9Ks6uGxE/9ec0aKEtbxNgbnuNLiGPNtCGpnH18dz5anx3RlA3RdpPzTkhj6nEpTZIOl9vHayv3UBEiDeEM0FPp8fHphoNccGKPJkmPEKE4bAaPXnoC1760igkZXUkJ8/dYCHH0JJhqg5JjkrnjpDu45YRb2Feyj2J3MTbDRpIziW6F2fDWdVRaBptziil2ebCZBslxUfTu0rib6itG9WLpjtyggSEi6Z+gNZw1NPxZ1887oQf9k+N56psdfLH5EIaixs2K3VAYhmJYj0TunJLBlMER3hiNvwv2roJqtTZZd8fX2KRXokHFn4L7X2llcChlIr/9OJvkJXncMSWDSQOT6wyqrHIPea9uoXJ3EVhQFR1WTaSsfYrSgpGU/msNzoGdSbp8MIlRiRgquHau+pD4htOoMZpf7SHxq86R5IxsEmAhjgVKKf5+8fG48gv4ckcBlWEUHkXbTc4cmsrfLjq+yZo6vf/D/qOaQLjM7eOZhTslmBLNbliPRK4a3Zt739nArGtGSHM/IVqIBFNtmM2wBQ17va8ygZdLzuW1Bz8Dw8QIZJYen0WPTtHcOqk/552QhrOOoYVrO1hUwdLtubjqGB44nP4JCjh1YDLJ8ZHNqTQkLYGnrj6ZgjI3C9buZd3eQgrLPcRF2UhPjuWKU3rTN8w+D0HSJ0PvMbB7GXgjG9RDOWLpedmjfNmpHx/9eIC/fbyZf32+jTtPy+CM41IxjJ9+oHzFbnKeWoev1E2949daBliaim0F5DzxPSm3n8CF/S9kwbYFeHXNJonhDolvKpNz0s9pshpKIdoao8LFr975BydcfDsvHI6irNJLWYi8KsZhEu+08YvTBnD16N5NehP54tJdRz2B8O68cnbklJKREtdk6RIilDtPy+D8J5bx/g/ZEsAL0UKUrmfEspEjR+o1a9a0YHJEXbw+iz++u4F3v9+P5fPi0aH7HMU4TBTw+JUnMfW4ugcl2JVbxnOLdvLJhoNccnJPHDbFnGVZIUf1a0i03eSVG0ZxSt82VkPiLoc5Z0POFgh3nidHLEx/B3r/VLJsWZovNh/iya93UOn1cceUDM4dnobyWOQ8+T3evAoabKtYnamwd4vFdeZhfr78t1Sqxo0aGGVG8cY5b5DRObK+IUqptVrrkY06aRsi+VP7l/37P4Bpkva3h7EszbKduby4dBc7c0pxeXxE2036p8Rx44R0xmd0aZaS+GH3f0ZpZXAfzNKN31DwzYv0unNeg8eId9p4/IqTIq9h74DaQ/7U2nnT+n2FXD9nNR/fNZGUeGnuJ0RTqC9vkpqpY4DHZ3Hd7NWs3V1ApdeivnFDqkpQ73jtO/524fFcXGto4I3ZRTy9cCcrduYxfUwfvvntZJJiHVR6fSzfmcfmA8Uhhy2vS7Td5PJTera9QArAEQPXfQrv3gpbPgZ03SP8OWL9o/ddvQBSh9ZYZRiKM4d2Y9qQVBZtO8yTX+/gsS+28VC3rvQtCA6kVu1bz9++eYZtuVkYhsGALn24f+ovOLF7oM+XT+M9UEjym29w4uAefFdxKGiI9IbYlI0hXYZEHEgJcawofOddXBs20G/BfMD/PZw4ILnFJyata/CISCYQ1lqHDMiEaA7De3bispG9+PO7G3h2ur+534b9RXy64QAHiyuxLE1yfBSTB6UwJj1JmgMKcZQkmDoG/OHt9azdnR9RJ+wKj8W97/5It0Qn4zK6smpXPk99s4MtB4u5cUI6/7hkOHFRP/37o2wmc68fzZWzviXzcGlYNVTRdpOzj+/GfecObXDbVmN3wqVzoCALVs2CtXP8HbwME7TlD676jPP3seo3GYy6A1WlFJMHpTBpYDIrM/NIfGlz0FxWJZVlXPfWH3h42q85b/AU3D4vq/b9QFStpnjaslHsuIVHzxvEpR9eSm55blBzv7qYyqSzszP/mfKfiN4KIY4VlZmZ5DzyCL3nzMGIiawvaE5JBa+v3MPnmw5R5PJgKkVSnINLR/TkwpN6EOOI7GfPaTPx+IK/m5EM0KOUIs4pP7ei5fxy6gDOeXwJf35vAysz89lX4KLS6ztS9qeAud/uJjHazi2npnPZKb0i/m4IIfzkm9PG7cgp4aP1B0IGN2WbFlK8+l08efswHNHYU9JJHHfZkVGkKjwWv56/jp6dYzhcWsktp/bn+WtGEGULXYqaGGPnf7eP46GPNvPW2r1BM65XiXWY2G0Gd00dwMxxfY+NUq3OfeHMh2Hq/VC8HyoKwR4DcSn+eakioJTiBK9Bvmmia82XlZm/F4ALh5wOQLRhMqlf6HnArFIv0TkGr539Gtd9dh0Hyw5S6aus99xRZhQpMSm8dOZLMvCEOCZtzC7i0w0HOVhcgWVpusZFMWlQMmPT/c30rIoK9v/q1yTffTfOQQPDPu6OnBL+75MtLNmei4JALb7f7vxyth4s4a8fbuaik3pwz5mD6FzHxOK19e0aG3K000gG6HF7Lfp3lf5SouV4fBYO02Det3tCrtf4W7KUu33849OtvLxiN2/cPIZUGQVQiIhJMNXGvbg0C4/VuCF5AQ4VV3L1mD7cNqk/tjrmk6rOaTd56MJh/L+fDebddfuZvWwXh4oqqfRZRNtNBqTEccuk/pw2OAXTOAaCqNpsDkjqd9SHca0/jA4RaKYn9cJQBr/66GHOHzyVk3oMpZMzPsQRQHssXBvzSP5ZP+afO583trzBK5teweV1Ue4tr7FtjC2GKFsUM46bwVXHXUWsvZEDcwjRCrw+iw/XH+CZRTvZnVcWNCVCVQn5zRPTOfWLuTj7p9PpskvDPv7ynbnc+PIaXB4fdXUDrioYWrBmL19vzWHBLWPplVR/rZdlaZLj6g66IplAuLGjrQoRqQqPj8uf+5adueHNQejy+NibX875Ty7lk7tOJSnMggYhhJ8MQNGGlVV6GfHQF0Hzm1iVZex76lq6nH13g81LDAVnDevG01ePaM6kdjiHX9pA5baCkOu252bx9MrXWJq1lsNl+UzpP5pHzvodybHBNUkxJ6WQdPmgI68tbbE8ezlf7P6Cw+WHAeji7MIZfc9gfNp4zAb6ZoSjPXTwBsmfjhVllV5ueHk16/cVhazprs5pQFJ5IQt+dxY90rqEdfwf9hZyxfPf4vKEPzGuoaBrXBQf3zWRrnGhRyH1+ixumbeWZdtzGzUwT5XYKJPHLjuRaRFMHdGRtYf8qbXzpl+9uY5Pfoy8RYvd9E9H8s7t41sh1UK0bTIAxTHqx/1F2A2DCmpmiJEMyWtpWLYjr7mSKEIY0LUvj51zLwA78nbzyw8f4oGvnuCp8+9vcF9DGUzoMYEJPeoPkoU4FlR6fVw561u2Hiyp0eyuLhUWHIruxEWzv6830Kni9VlcN2d1yECqvptGS0N+mZu731jHvBtHhzgy3PvOjyzfcXSBlKGgU7Sd02QUP9FCcksr+ejHA7hDfG4batHi8Wm2HChhw/4ihvVIDHF0IUQoEky1YUUuD6HqDX2uYoyYhAZHkKpS7pZRpJqaGWcPa7uMLn24bNhZzFv3fsj1Rrw0pxByEGMIAAAgAElEQVTt15/f3cC2OgKp+oKd/DI3189Zzft31l+o8OXmQ1SGCKTCaQbttTSrs/LZV1BOz841m+Bt2F/E+z9kB7UKaCjd1SkFcVE2Xr9pbFhNrIVoCq+t3EOoBvjhTjLt9lq8uHQXj11+YgukVoj2QXL4NsxUKmSmWH1I3nAYx8IAEceY6GFdUVHBweyOvN08t+oNDhTnAJBdfIj3Nn/FyWnBIx4qh0H0EBlEQrRPBWVu3luXHbJmp3jVO+R/NYvEMZfR88559LhtNvEnn41r+0rAH+hsP1TKD3sL6z3HMwt3Bk3iW3XTmHTGbcQMGofhcKJMGzEZo2sMCgFgac3cFbuDjvvCkkw83uCirIbSXV1SjIN37hgvfaVEi5qzPCtk4UW4LVp8WvPxjwcok6H8hQibBFNtWFKcI2TNVPUhecMhQ/I2PefgJJQZHKTGOmJYl72Z8+beysB/T+P8ubcxqGs//nzaHUHbGrF2HH0SWiK5QrS4N1fvIVQ5TrjBTqXXx6wlmXUe/0CRiy0HS4KWR9IM2uPTvLF6b41lRS4Pn2w4iK9Wf+JIgjRDwcUn96B/sozgJ1pOhcdHYXnouRQjadFiMxUHisKc6F4IIc382rLhPRKxhRgxL5Ihee2m4vwT0loy2R2CMhRx43tQ/M0eqFaC3T0+mWcu/EvD+9sN4if1PDaGlReiEV5alhWymVy4wY6l4YtNhyiu8JDgDG5We6CoAofNCCqFj7QZdHGFB8vSGIG8dsXOPOxm8HEj7av65eZD/PGcIWGlQYimUFbpDfnZhcgmmTaUoqRCaqaECJcEU22YzTSYOa4vzyzaGZQ5hjskr6EUM8f1bcFUdxxx49MoW3MQX2ElIasQ62KAmeQkdkRqs6VNiNbk9VkcLg09Z1pEJeSWj01z3yLDKsEqK8MqLz/yd487Git+FBg1+x1GctMIoCyLTdPOItpuoqKi2NFlCN4uo8Co+fMYaZBW5JKbUdGyYqNseHyhB0yJZJJpS2viouT2UIhwybeljbtqTG+eWbQz5Lq4oVOIGzqlzn2VguE9E+nTReYkag6G00byLSeQ89T3WGUeCGfQL1NhJjhIvul4lP3ohzkXoi0qc/uwGQqPL7iUIaJgx+ejcF82VgIYsbHYunbBiI3FiIkhzetErXTVqBmGyG4aATAMBsx+ETxudGUl8Zvy4btCqNUlNdIgTYiW5rSbxEXZKA5RqxRJixaPT5OaKJP3ChEuCabauJR4J787axCPfrYtonlUAGIdNv5xyfBmSpkAsHWKIvWuk8l7ZROeA2Von6bGbKRVDMAwiOqbQJerj8OIlq+eaL9iHCbeUN8DIgt2lNNJ71tuIiUtuG/hkEov3lVfULtaOJKbRoB+XWOJ6tXzyOtU3wFs69eDr+YNaaRBWoL0VRWtYPqYPrywdFfIodHDadFiKJg6OCVk01ohRGiS2x8DbpiQzsrMfD7fdCis7ZXyB1JzbxhFunSAbnZmnIOU20/Ec7CMkiX7Kf/hsL/jvaGocPuwm4r4kd2IG5+GPVlG9hLtn9006BRtp6DcE7QushJyi7ROoUvI46JsnDc8jf99vx9frcAt3GbQMQ6T2yZn1Fg2Nr1ryKZSkaTbYTM4Z3j3ht8oIZrYjLF9eHHprjrXN9SixWk3uenU9OZImhDtlgRTx4Af9xWxZncBd00dwJzlWXgti7LK4FoqU4HdZtA/OY7HrzxJRpJqYfZusSRdOpDOF2bgK3OjK3088vV2+vTuxPTxfVs7eUK0qBlj+/DcosyQneHDCXYUMCGjK51i6p6L7YaJ/fhgfXZQMAUN3zRWObdW0JMYY2fakFQ++vFAUCVzuEGaAmaM6dvguYVoat0Tozl1YDKLth0OWTtVHxMfvRwuTuoR30ypE6J9kmCqjTtUXMHNc9fwt4uGcdaw7vzitAy+3JzDMwt3sGF/MQAWGqfN5LwTunPDhHQGdZOMsDUpu4EtUJqe3COenfllrZwiIVre9NF9eHZR3UObNxTsRDtMbm6ghHxwtwRO6ZvEql35IYO2+kTbTW49tT/OEH0Xbzo1nS8354RsWt1gX1VgdL8kukmfE9FKHrnkeCb84xtCD5IemqEg3ulgTrfXUHPnwCUvQHy35kqiEO2KBFNtWIXHx82vrOGqUb05a5i/9NRmGpw1rBtnDeuG1poKj4XNVNhNmTKsLUrvGseyHXmtnQwhWpzTYdIl1sHBooqIBrsE/4Tl3RKcjOrX8KTWz04fwflPLmVvgSvskvhou8mUwcn8YmpGyPXDe3bizKGpfLrxYMjh3esTE2Vy//nBk3QL0RK8Pos/v7eR43t2Iq+0kr355SEnzq7OYSo6xzh485axdE86Axb/E56bBBc9C/0brt0VoqOTO/A2SmvNPW+tp0+XWO48LfQPvlKKaIcpgVQblp4cS2ZuaWsnQ4gWtSYrn7P/u4RJA5Pp0yUGe4gJruuilH+i8VduGBXWPGyxUTbeuWM8w9ISiHE0NH+OP5C66KQePHHlyfUe/5+XnsCofklE28PPX2McJnOuGyVNrEWr8Pgs7npjHSUVXuZcdwof/GICv542kOT4KGJDfDdio0wSnDZuPrU/n/3qVPp2jQXDhMl/gIufh3duhW/+BlZkg18J0dFIzVQb9cTXO9iTV8abt4yViV2PYb2SYjhUXEmFxxeyOZEQ7YnXZ/H41zt4beUe/u/i4zljSCr5ZW6ufuFbduWWNVjLYzcVnaIdvHHLGHp2Dn+wlgSnnQW3jmPh1hyeWbiTH/cXYSj/EM9K+QfE8FqaM45L5caJ/Tipd+cGj2k3DWbPHMWDH2zkjdV7URCyhF/hb5KYFOvghWtHMrhb8MiDQjSW1ppSTykV3gpi7bFE26JD3hN4fBa/fP17Kjw+npsx4sjvzc2n9ufGCeks25nL++uyOVRSiWVpusQ5OGtoN04fkhq6QDZ9EtyyGP53I7xygTT7E6IeEky1QZ/8eIA3Vu3h3TvGyw34Mc5uGvTsHM3uvHLpyybatb355dz1xvfERtn4+JcTSEnw9xlKinXwzu3jmbtiN7OWZFJW6aXMXbOkO9ZhYhiKGWP6cOPEdJJi6x50oi6moZh6XCpTj0tld14ZK3bmUejyYDMUSbEOpgxKoXOExzUNxV8uGMYvpw7g9VV7eGlZFmWVXmyGwtLgtSwmZHTl5lP7MyY9SQq+RJPJKsri1c2v8t7O9/BYHkxl4rW8dIrqxIwhM7h4wMV0dvoLBdxei1+8/h1en+bZGSOIstW8bzAMxcQByUwckBxZIuJTYca70uxPiAYoretuzT5y5Ei9Zs2aFkyO2LC/iGteWsUr149iWI/E1k6OaAI3vryGS07uwc+ObxtDJSul1mqtR7Z2Oo6W5E9tx7vf7+evH27itsn9uX58PwwjdFBhWZrlO/N4d91+coor8FmaLnFRTBuayrQh3XDY2naTZcvS5JZVUuzyEmUzSIp1EBslZZJNqT3kT0eTN+W58vjNot+wIXcDPsuHVwdPwOs0nVhYnJ9+Pr8d+QfufuNHtIanrj4pKJBqMpkL4X+3wIhrYdLv/c0BhehA6sub5FegDckpruDmV9bw0IXDJJBqR/onx5KZKyP6ifanuMLDfe9u4Mf9RbwcRgGQYSgmDOjKhAFdWyiFTcswFCnxTlKkklk0g/2l+5n+8XQKKwpDBlFVKnwVAHyY+SGfbl/HEO7hmavHNG9hRPpkf7O/t2+QZn9C1NK2iwHbC62h5BDkbIHDW6EsN2iTCo+Pm+au5YpRvTm7jdRgiKaRnhxL5mEJpkTblFNSwecbD/LW2n28t24/y3bkhpy0tra1u/M55/ElxETZ+PAXE6UASIijUFRZxMxPZ5JfkV9vIFVdha+CcnZjT5tLc1VI1RCfCte8B33G+5v9ZS4Me9cKj4+ySi/1tYYS4lglNVPNqbIUflwAy/4DxQfAtPuX+9yQlA4TfgVDLkTbovjdW+vpnRTDL+oYuU8cu/p1jeON1XtbOxlCHKG1ZuWufJ5flMnSnbk4TANLaxT+UUKVghlj+jBjbB+6J0bX2Nfrs3jymx3M+3YPD180jDOHSum0EEfr+fXPk+fKw9I1CzIKlhSQ+1ku7hw3ptMkYUQCqT9PxYz1R08WHr7LWcvifYuZ3Gty8yfUMGHK/4M+YwPN/mbCpN8FNfvTWvPDviJmLc7k800HsSz/SJ2W1pzUuzO3TurPaYNTMOtoEizEsUSCqeay+kX4/E+AAk+gVsJX+dP6w1vgo1/DR7/hs/R72Z13gozc105V1UxpreX/K1pdaaWXG+as5sf9RbjcPjSEnJ/phaW7eHHpLu49ezDXjusH+AeZ+NWb64iyG3z0ywmkJsjEtEIcLbfPzdvb38ZjeWosz/0kl8OfHKbnjT2JGxKHp8BD9txssh7Not8f+2EEmvW5vC5mb5jdMsFUlfTJtZr9veivucLf9/uXr3/PgaIKKr0+rKrKqMDftbsLuPuN73HYDB68YBjnnZDWcukWohlIM7/m8OWD8PkfwVP+UyAVirsM3KVM3vIAc4d9JyP3tVNdYh1orckvi2Q+eiGaXlmllwufWsa6vYWUBwKpuri9FpVei79/spX/frmN99bt58KnljFtaCpzrx8tgZQQTeTz3Z8HNX/zuXzkvJtD2vQ04ofHo2wKR7KDXrf3wp3rpmh5UY3tN+RtYG9xC7eAqNHs71TIXMiS7Ye59NkVZOaW4fJUC6RqKXP7KCj3cM9bP/D0wh0tm24hmpjUTDW11S/CyqfB4wp7FydunEv/Bsm9YeiFzZg40RqUUqQnx5GZW0aXuKjWTo7ooLTW3PTKGvbml1MZoiaqLi6Pj8e/2k6XOEdYg0wIISLz4c4PKfeW11hWvr0cy2ORMKLmvGWm0yR+eDylG0vpfGq1+dI0LNq3iOlDprdEkn9SrdnfhvkPcnPxr3BZ4ZfTV3gsHv9qO11jo7jslF7NmFAhmo8EU02psjRQI1UzkOr7nxLKPbDrrjhiHf5mXi9852beeg8LZ8b6N/K64IO7YPC5YMq/pb3xN/Ur5ZS+Sa2dFNFBrdtbyPd7CkMGUmWbFlK8+l08efswHNHYU9JJHHcZzp5DAfBpf6HA0DSZkFaIppZfkR+0zFfqwxZnQ5nBTcNtiTZcu2veZ7gtd8jjtJj0ydxt8+GyKkKuri+PqfBY3PfeBs46vhsJTnvLpluIJiB37U3pxwXU1XLSp+G/K93cO7GemgnLC9s+gePOa570iVbTP1AzJURrmbUkk0qvL2h58ap3KFr5Fl2m3YGz38ko04Zr11pc21ceCaYASiq8rNiZx7iMY3NYcyHaqtqDTgCYcSbeUi/ap4MCKm+RF1tc8O1bqOO0lPX7Ctlf5Am5Lpw8RhmK/63dx8zx/Voy2UI0Cekz1VS0hqWP1dlH6p5xDh5dXklhRT29FNylsPQ/zZRA0ZrSu8rw6KL1FJS5+XJzTlD/BauyjMKlr5J0xm3EDBqH4XCiTBsxGaPpPOX6GtuWu308t3hnC6ZaiI6hs7Nz0LKYjBiUTVG8trjGcl+Fj5L1JcQOia2x3G7YQx6npcxaHLqwJtw8xuX28fziTBk6XRyTJJhqKiUHofRQnatHpplM7mvj0eWVdW4DQPZ34JWBCtoTy9LklVWyNiufBz/YyD8/28rcb3fLgBSixWw6UEyUGZzdV+7fgva6iRk4NqzjfL+nsKmTJkSHN63PNGJsMTWWmTEmKRemkD0vm5L1JWivxn3Yzd6n92JPstNpXKca2xvKYHza+JZMdg1fbD4UcrCJSPKYgnIPO6XQURyDpJlfU3EV+OeR8oZuLwzw4JQoxr9Uxl2jHXUfx3T4jxUYYlQcu4rKPby+ajcvLN1FudtHudvHS8uyAIi2Gzz04SamHpfCzaf258Reneo/mBBHodjlCTlyn89VjBGTgDLCG0nU5QkueRZCHJ1z0s/hkdWPBC1PPjsZM9bk4JsHcee4MaINEk5OoNctvTDsNQtH0hPTyejcOvNUenxWnYPaRJLH2ExFQbkUMopjjwRTTamB2ulhKSbnDrTx96Vujkuup1JQ5iI65m05WMxVs1ZS7vZS4Qn+kXEFln264SDfbDnMTRP78aszBso8VKJZOGwGoT5ZZnQCVnkx2vKFdbMjE2wK0fRi7DGc1/883tn+Dl7trbEuaVISSZPqH7go2vb/2bvv+CjK/IHjn5nZ3exuKgkkhBBKCITeO1IFxYJ6FtRDFBVF1BPU8353ljvP655dT0CKDVTsWLArKD3SCT2UNEJCSN8+M78/VpCwk2QXUuF5v168XmbL7BMYZ+f7PN/n+7VxW6/banxNfdJ1kDC+BQr1GqNWV0tdEJowkeZXV+xxoNY+o/LXMVbmb/KQW1bNBUP1gFWsUjRn+46Wc82cNRyv9BgGUqfSdP9s//yfDvKP5bsaaITC+SY+0opmsBchLKkrksmMY+/aoI7Twl7DqrogCGfsrj53EWGJCPl9FtlCakwqF7a7sB5GFeQYTHK1Ey2hXGM0XSfaJqr5Cc2PCKbqSmQCxLSr9WWpsTLX9zDzwoZqAq/koWASNyzNldOjcsP8dTjcgelQlTtXcOT12WQ9cy05L03l6Lt/wZWT4X+fV2XJuiw+2ZLb0EMWzgM9k6KIsAYmIshh4cRcMIXj38zFsXctmteFrvpwZv5M8Q+Lqrw2zCRzw6Dar3GCIIQu3h7PwosXEmmORA7y1ixMCaNNRBvmTpiLWW7cIKS6th+hXGMUWSI1PvSAUhAam0jzq0sX3A9fPASemjdQ/nl0GG9uMyghaomAC2bX0+CEhvDJ1lycHjUg3SGY0rBOr8p/v97DpD5tRLqfUKckSeKOkSk8/fXegH1PUYOvRg5vQenapRz77Ckki42whFSihl0fcJzfDhHBlCDUly4turD08qXM/G4mBY4CXD4XukHynIwCyAxqPYinRz+N3WwPPFgDu2t0J7Zml1DpCZxIDOYaE2aSuWVYB8wGhXIEoakTwVRd6nk16ue/5/Ss4EOzI6v8nBwt43rUoPmlyQqp4+tvfEK90nWdOSsycZz2ZXKiNGzcpbOxpw0/+bg9dQj21CFVXltU4WFTVjED2ovmvkLdum5gMk99vcfwuYgeY4noMbba95pkiVFdWtEqsoY+eYIgnLXkqGQ+vepTNhds5rWM11iduxqTbEKSJDRdQ9d1Lku5gk9+7MTsiyc1iUAK4ILUloSHmQyDKaj9GqMDNw1tX0+jE4T6JYKpOqJpOvNW55GrTuevyisoavVV/QyZrPCbeRBkVS2h6dmeW8rRssDS96GUhj2xf0oEU0Jdi7aZeWZyH+57ewu+EDZ5SxK0CLfwr6t71ePoBEE4QZIk+if0p39Cfyo8FRQ6C8nIK2HzQRWH24z3qEz3lpU88VkGS6YPbezhAiDLEv/8TS/ufXtTrXuFT2czK0wd1p6EKGs9jU4Q6pcIpupAQbmLB9/disur8tysh1EyWsEP/wSfM7gDmKxw2TPQWaxKNWd78ssNCzGGUhpW12FnXlmtrxOEM3GoyEGM3Uy5y1dtKeNTmWSJ2HAL7901jJYRYlVKEBqST9X4bmcpc1Zkc6ioEq9PQ/1lHkSW/AWMLvjP99x3YWeu7NuGMFPjTsaO757AHy/pyr+/2B10QCVLMK5rK/44sWs9j04Q6o8Ipn7h9qkcLnJQ5vRiUmRaRYaRFGOr9X0r9xby0HtbuWFwO+4bl4pJkWHEfRCVBJ/fD5oKngrjN1siwBQGv3lFBFLngHKXz3DGP9TSsA6Pr9bXCEKoFvx0gHfTs/n8vpFkFlTw+KcZZB934vFpqKdV+rOaZXQdJnRP4K9X9CBOBFKC0KDKXV6mvZrOriNlAanjwMkGuTnFTv7ySQavrT7EkulDaBHeuAWspg3vSFx4GH94fxuShOHYwb9HSgcSosJIiLIii7YLQjN23gdTuSVOXl99iLc2ZKGjI0sS6OBRNdrH2Zk5phOX9EzEaq56E+zxaTz99R6WbcnjuRv6MrxTy6oH7nUNdL8Cdn+Ovuo51LytKGaLv9eL6oXkwTBiNnSeIFL7zhF2i4JisDR1amnY8K4X1Hqc0881QThbr685xOtrD7H0zmEkRFlJiLLy9f2j2ZFbysJVB1mTeYxKt4pJlmgRbuH6QclcPzC50W/MBOF85PKqXDd3LQeOVeIJYgXZ6VHZV1DOVS+v5rPfXUCktXEr+03q04YLu8XzyZY85q7MJL/MdbKwhKrpWEwytw7vwI1D2mFRZCa9tIp+7VowqU+bRh23IJyp8zaY8qkaj3y0g4+35KLrOh41cEVh79EKHv1oB499nMH/pvRndJdWAGQVOfjd25toGRHG8lkjia3uhkMxQ4+rOJQwgdsX/Mj3d/cFSQZbjH9FSjinJMfakQ0KEZ1aGlaSFawd+yHJJlyHtuDK2kaLsVWbLSa3aBobioVzw5L1h3nlxwO8c+dQ2py22t4zKZpnr+/bSCMTBMHIQ+9t5aBBIFW5cwVl6R/jLcpBttgwx6cQPXwy1rY98Ko6R0pd3LNkE2/cPqSaIzccu8XEDYPbcf2gZHKKnRRVelA1jWibhQ5xdn8Wzy/mTBnAzYs2kNY6ki4JkTUcVRCapvMymPKqGrcs2sDmrJJa9w2cqEwz482fefKa3iBJPP5JBveOTeXWER2CKmG972g57RPiICqxTsYvNE1DU+IIMylUGvSYCrb8dHiYwq0jOjTQiIVz3bvp2bz0/X7euXMoybEiSBeEpu5omYuvdh4NCKSCaa/h8WmsP3icA4UVpLRqGv2aJEkiOdZe4/WnZ1I0j1zajRlvbmTZvSOIauSVNUEI1XkZTP3h/W1szirGGULFGZdX4/6lW4mPCuON2wbTMyk66PfuK6gQsy3nAUWWuG1EB176fj8ugyC9ttKwAGZFZlzX+PoaonAe+XBTDk9/s4e37xhK+7jwxh6OIAhBWLzuMKdP0YbSXkPVdF5dfYi/XdWzAUZbd64Z0JYt2SU8+O5W5t00QOyhEpqVZhdMubwqn207woKfDpBT7MTtU7GaFNrF2blzVAoTe7ausaLNnvxyvthxxLDSTE1L6ACqrhMeZqoxkHJ4fCzbkseiVQfJL3XhVv19IdrE2OjXrgXju8VXWd4Wzi03Dm7H/37IPKP3Ws0yd4xMEeeHcNaWbcnl31/s5q07hjSZGWpBEGqm6zpvrD0UkDETSnsNn6bzwaYcHr28W6NX9wvVY5d35/pX1jJnZSb3jE1t7OEIQtCaTTDlUzWe+noPb6z1z9qc2hjOq/rIyCvj4Q+388hHO7j9go7MurCz4czGwlUH8BqsGgSzhA6QW+xkR25pQEDl9qn8c/ku3k3PMaxgc7jIwYPvbcEky9w7NpXpIzsGlSIoNC9xEWG8fFN/Zi7eGFKvjTCTzMD2sdw1ulM9jk44HyzffoS/f76LxbcPITVerIgLQnNR7vYZVr8Lpb0G+FtsFJa7advM9t9aTDJzpgzgipdW0bttNCM7t2rsIQlCUJrFFLjTozJlwXpeW3MIh0ettsN2pUelwu3jlR8zue319ICc4wq3j0+25HF6rYkTS+ixE2ZiTxuObLEiKSbsqUMCigN4fBoLVx2s8liZy8t1c9aydEM2Tq9abSnQSrdKqdPLM9/s5b53NqOG0DhTaD7GpsXz3PX9sJplw75Tp7OZFYamxLHgloEoIrVBOAtfZ+Tz52U7eO3WQaS1FoGUIDQnFS4fJoMqRqe21wiGjs5fP93J1IXrmTJ/Hfe9vZnPtx3Bq4bWTLcxtI628vwN/bh/6VZyih2NPRxBCEqTD6ZUTWfG4p/Zkl0S9Ey/06uxLrOIWe9sRj+lf8rW7BLDFKpQltBVXWfl3sKTP3t8GtMWbWD30XLDfTLG41P5dmcBj368o8r4hHPHxJ6t+XDmCMZ3SyDMJGPRqvaOktCxWxSSYmw8cllXFk0bJEqiC2flh90FPPzRdl6dNpgebYLf0ykIQtNgtyj4tMD7iFPbawTD5dX4ZudRftp3jNWZRXyyNY8/vL+V/n/7hie/3E1xpaeuh16nhnWKY8aoFGYu3oTLG1wAKQiNqcmn+X24KYf0g8WGVfdq2uPk8mms3FvIlzvyuaSXv4peqdMLBAYvoS6hV7p/vTF+Y+0hdh4pM+wFUdP4nF6VjzfnMql3IsNTWwa8V2j+ureJYv7NAzmaV8j/7vs3+y+6hlKXisVZSVxuJnfceQODO8aKdE/BkKZrbDy6kZzyHBw+B+HmcFKiU+jVslfAOfPj3kJ+/95WFtwykF5tRSAlCM1RlNWMSZbxqlUDiFDbaxg5kdGz4KcDvL8xh6UzhtGxZdMtTDN9ZEe25JTwl2UZ/Ofa3o09HEGoUZMPpuauzMRpMDMRzB4nh0dlzsrMk8GUP4Mq8Mb11CX0YAIq+ZcbGU3TeeXHA4YrZsGMz+lVmfdjpgimznH2XduYHlVCuzv9VZh8hYUcmPRXOne8WwRSQoBSdykf7fuI13e+jsPrQEdH0zTkX9J/4qxx3NbzNi5LuQy72c6a/ceYvXQLr0wdQL92LRp59IIgnClZlrhuYFveXp+F97RtAMG216iNR9U5VuHmNy+v5otZI0mMttX+pkYgSRJPXtObK/+3mnc2ZHHD4HaNPSRBqFaTDqa2ZpeQV+IKeDyUMqF788vZsTKdNod3o27PQqMbKFUb5p66hB7e9YJaxxVp9f+1rc48RoXbF/B8KONbd+A4+aUuWkdba/1coXmqXL8e+5Bf/92Vlv7gWT12DFMrscFW+FV6fjr3fncvmq7hUk+79v0yZ5NTkcNTPz/FC5tfYFb3p/jnsjJentKfgR1iG37AgiDUqVtHdGRpejYY7KkOpr3GqWrKjil3+Zi2KJ0vZ49sspN64WEm5k0dwOS5a+mWGEWf5JjGHtJ56UjFEd7Z8w6bCzZT4akgTAkjOTKZyWmTGZAwoMmeP5fJiukAACAASURBVA2pSQdTH27Oxe0LXJUKZY+T1+PlncVfcU+Sj8GDe6Nst4Gn6kpSKEvoZkXiyr5tAHj352zDYhOhjA/81bduu6BjUK8VmgevquH2aYRbFBwb0kl8/C8nn5MkibDOnXHv2yeCKeGk1bmrmf3D7MAgyoDD58Dhc/DXjXfz8KQXGZoS1wAjFAShvnVsGU6/djFsPFyM9/RqWSGoLTtG1XSyix1syiphQPumu6LdqVUE//hNL+5esolP7h1BXERY7W8S6sTWwq28tPklNh3dhI6OV/OefC6jKIMVOSuICYtheq/pXNvlWmSpyZdhqDdNOpjKK3YYTc6EtMdJlRUcEy4ncXJfAKZG7mbBqoMBe5yCXUKXJYlbhncA4IjBqlmo43P7NI6W1X7zJDR9BworeHX1IT7cnIPToyJLEqqmk9z2Cu7ztOByr3qyyMSJYCp8+PBajiqcDw6WHuT+FfcHFUhVIXt4addDTEhbRkubSBcWhHPBy1MGcMnzP3KswnNGVX+DzY5xelXm/3iAAVMH1Mm468vEnq3ZmlPCfe9s5o3bhoiqtw1g2f5l/H3d36v9TtLRcfqcOH1Onkp/ih9zfuTpMU8TppyfwW6TDiPd1ZTxDLVM6KnFK6YOa2+wa8ovosdYEm95jnYPfEDyvYuJv+5xrG27nXxelqBfu5iTvRs8dTQ+l8Hqm9B8HCl1cu2cNVzy/E+8vSGLSreKpvubJ+pAVkQ8f/lsF/3/9g0vr9iPruuEdemMa9++xh660ETM2zoPt+oOeLz4p2L2PbqPjDsz2H3fbvJez0OtrHq9cPlcvLP7nYYaqiAI9Sw23MLH94wgKcZGmCn027Rgs2N0Hb7fU0CZy1vj65qCByd0AeCpr/c08kjOfV8e/LLGQOp0TtXJuiPreGDFA2h60y+/Xx+adDAVF24xfDzUMqEtTzlOYrSN2eM7YzuDMtThFhP/vvrXqjIxdvNZj0+WIC78/IzkzwX7C8q55Pmf2Jxdgtun4atmFrHS4+8/9uJ3+3nw3a2YU1Nxi2BKAMo8ZXyT9U3Al9CxL46R/14+rSe3pvvL3Ul5LAVPkYdDTx1CO2WCyKN5eHv321VSMARBaN4So218ft8FzBidQozNTLgl+HuWULJjzIpEQTPIjjEpMi/c0I9PtuTx5Y78gOc9Po3Ccjf5pS6c1fT6FGp3pOIIj61+zDCQqmlyz626Sc9P561dbzX0kJuEJp3md0FqK77KOBqwLymUPU7hYQrDOlVNf7lrdCeKKj0sWZdlWCnwdJLkD6TenD6EDqeUEh3VuRXpB4sDjhHK+KxmhYEdmm6+slC9gjIXk+eto9ThNSi4b8zpVfliRz4tzAlM3rcfXdOQDJo0CuePZfuXIZ82r6U6VQo+LiDp9iQie/ub71paWUi+O5m9D+2ldE0pLUb9et3waT5WZq9kfPvxDTp2QRDqT6TVzAMT0rhvXGe+3VXAp1vzKKxwo+s6kVYzP+4tNJzAC6VCsYxEhbt5BB9xEWG8PKU/t72WTueECFJahrP2QBGvrDzAqv3HUGQJSQKvT6dDSzt3je7EpD5tRA/HELy9+21UPfB8OPbFMQq/KKTt9LZEdI/AW+wl7808Dj11iI6PdEQ2yTh9ThbtWMRvu/32vNs/1aSDqct6J/LYsh2GzwW7x8ksy4zvFl/lMUmSePSy7nSIC+fJL3ejqhqVBuXNFVnCrEh0ahXBizf2I6VVRJXnrxuYzH+/Ml5yDnZ80TYzw8Tm8Wbpb5/tpNTpMQykausxtmRzPkPiO9Ix7wiWtkkNPnah6UjPTw+YBXTsc6B5NaIGRFV5XLEqRPaOpCKjokow5fA52Fa4TQRTgnAOMikyE3u2ZmLP1icfc3lVuv/5S8PXh1KhWNN1IsLq7lZQ03TketzT1Cc5ht9fnMa0RRvwaTqlTi9Oj4oOVQLLzMJKHv8kg798ksFDF6dx6whR5Ks2HtXDu3vfDchyCGVyr9Jbybq8dQxPOr/2gzfpYMpqVrh+UDKL1x02rGpTW5nQMJPMzcPbY1KMI+Sbhrbn+kHJfPTGchZklLLfHIMkSei6jsUkM6l3G24f2ZGuraMM3x9tMzOxZ2s+23oEVQ99fDazzIxRKaKsZDNU4vDw9c6jGG2bC6bHmNensazrWIbt2yuCqfNcqbs04DG1QsUUYUJSAq8NpmgTzsPOgMeLXEX1Mj5BEJoeq1mhhd1CUaUn4LlQsmN8mk7iWbRmOV7pYWl6Fq+vOcyxCjc+TcesSLSLtTNjdCcm9W6DLYQUxWAkRlvJK3XVWpzjRKPiJ7/cQ9ZxB3++vLu436rBT7k/YTQ7HMrknsPnYMnuJSKYampmjOrEBxtz8KqB/ZxqIuHvUXDLsA41vs6syIzY/j0Thg0j5tpLcXhUzIqMJchNn7PHd+GbnYGpiLWRJYi2Wbh2YHJI7xOahqXp2Rhdk4OtoqTq8K21Hcf3ZBI5Nvi+IcK5x2oKvJFRIhR8FT50VQ8IqHylPkwRgZduu8leb2MUBKHpmTaiAy99v79Kka0TgsmOUVC5VPmZ8J82QN8p0DI16M+udPv404fb+TIjH1kC1ynZPV5VP7ky9PgnGdw2ogMPTEirkxWrbTklzFy8KaQqh06vyjsbsomPDGPmmOB/x/PNkYojhntvQ53cyy7PrrcxNlVNMpjSf1nlkSSJ1tFWFk8fwg2vrAs6YDmxx+ntO4bW2pNA93ioWLWahD/9CUmSCA9xubtjy3AW3jKI215LD2r/FYAiSURaTbw7Y1idLq8LDefDTblVvjxOCKXHmEmWWJ1ZRPv6GKDQbCRFJCFLcpUCFPZUO5JJomxjGdGDo08+rrpUyreVk3BtQpVjWGQLSRFihVMQzic3Dm7HS9/vr/b52rJjLGYz06+9AvLeh1cvgdiO0Pe30ONqsBpn5IB/NerauWvILXYGtJk51Yl7toWrDpFxpIxXpg7EXE2mULAeem9btfdataXXP/ftPq4Z0Jb4yDNfiTuXOX1OfFrgwkWok3suX9MvaFLXmsSdvK7rbDh4nHk/HmDdgSKcXtW/smQxMbFna24f2ZH3p/bmt3NW4bFH4PBVPyMRblGItJp4646hAXucjFSmpxPWseNZNU8d1imOt+4Ywo2vrMOratTUZy/cohAXEcZbdww5WWJdaH6KHYGpFRBaFSUvEu+7Ylj/zmZMskRitI1JfdqQ1jqyrocrNGHXdL6GTzM/rbJvSrErxF8VT97iPGSrXGXDrznWTMzwmIDjTOw4sSGHLQhCI2sZEcalvRL5YvsRXDUENUbMikS3xCh69OoHvfrB+Mdh/7ewZQl8/WfocjH0mwIdRsEpRZJcXpUb568ju8iBN8jVIadXZW1mEQ+8u4UXbuh3xql2O3JLOXy80vC5YNLrAd5en8Ws8V3O6PPPdRGWCMyKOaBNR6iTexHm2u+9zzWNHkz9uLeQP324nWKH5+QmQvCnbZa7fXy4OZdPt+WR7C3j2chsyq/+LXNXZHLgWAVmRUbX/StRXlUjLSGSmWM6cWG3hKBnPyq++56IceOCem1BmYsvduRTUO7G7VWJjbAwqEMsA9u3oE2MDatZZtb4zry/MYe8EhcmWUJDR5YkPD6N3m2juWt0J8akxYumc82cZrBHDkKrouTRYFN4WzZtyfO/V4IFqw6Q0iqCmaM7cVmvxHrdyCs0DT1a9iAxPJGDZQerPN7q0lYo4Qr5S/PxFHiQbTJR/aNInpGMbK56fRvYeiCtw1sjCML55V9X92JPfjn7C8rx1DSTewqTLBEXbmHBLYN+fVAxQ9ol/j+VRbD9Pfj6UXCWQJ8boe+NEJvC/J8OcPhYpWEgVdPKkMur8d2uAlbuLWRMWnzAe4OxcNVBvAaT6cGm17t9Gq+uOcS94zqLezADnWM6G1bhC2VyT5EUesT1CDjGua5Rg6ml6Vn85ZMMw3SpE1RNR9V09unh3KP1YEGEheWzRpJZWEFWkYNyt49Iq4mOceFVypYHQ9d1yn/4gXbzX6nxdWszi3jlx0xWZxYh8WsTYEWCMLNCbLiFluH+/U8zx6Qyc0wqu/PLyDnuxOFViQwzkRofQXKsWIk6V0RaTRyrCFydCqWK0ulUHVSvxs68Mv7vg218sCmHuTcNEGVdzwPTe0/n7+v+jtNXNfc8dnQssaNja3yvzWTj9p631+fwBEFooqxmhXfvGsatr25gR25ZrdsNrGaZNjE23rljKLHV9PIkPA6G3uX/c2QbbHkLFkxAbdmVRYfuwuULvOEOZmXI4VGZt/LAGQdT3+48aljsK5T0eq+qkZFXSu+2gav757sBCQOItkQHfA9B8JN7ZtnMTd1vashhNwmNFkx9s/NorYFUFZKEw6sx/fWfef+u4XRvE0WnINL4Tsg+7mDxusNsyS6h3OXDblFoZ/ZyYXQSqZ06Gb5H1XQe+3gHH23OxeVVA4qcqLr/4uDwOMkpdlJQ4Wba8A60bWGna+uoaqsACs3f+G4JvLbmUECVyVCqKNXE4fGnRdyyaANLpg+ptiKlcG6YlDKJ77O+Z3Xu6qC7zgNYFSvXdL6GwYmD63F0giA0ZRFh/j3in2zNY+7KTLKPO3H7VE5dPLJb/NX/7hyVwuSBycFX2Evs7f8z4Ql++P4rPPt9QNUgLNiVIYBNWcVkH3eEPLms6zqVHuNCZKGk18uSRLFDNDg3IkkS03pO47mNzxl+DwUzudcuqh1psWn1NcQmq1GCKZdXZfY7mw0DqZqWicF/k3nf25v49sExQX3WhoPHeeabPWzOKkHT9So3v5vQWd51Ms88s5LfjevMlX3bnMzl1XWdB5Zu4eud+TiDDPiOlrqZ9OIqls8aSWK0Laj3CM3TLcM78MbawxjVEQ22xxjUfL67fRrbckr41xe7eezy7g3wWwmNRZIknhz1JA+seIBVOetQcdf6Hqti5fJOl/PQoIcaYISCIDRlJkXm6v5tubp/W3bklvJVRj5Hy/zlw1tFhjE2LZ7BHWPPvDS4ycL7R9tQqeUHPBXKypAOfJWRz01D21Pq9P76x+Gt+rPTS9kp/13i8FDdFq1Q0uvh1yJnQqArO13JnK1zQprUO8GqWJnVf1Y9jKrpa5Rg6rNtRwwfD3YDYW6Ji63ZJfRJrnmZ9s21h/jH8l3Vrn5pSLiQyCys5E8fbuenfYX855remBSZ+T8d4OudR4MOpABUXafM5eO389fz3QOjxX6Xc1jbFnb6t2vB2gPGvX1qq6IEwZ3vTq/GkvWHeWBCl5ArTQrNi0WxcH27P7M243ki43/Eo3pw+BxVXiMhYTPZiLJEcXffu7kq9SrRN0UQhCp6JkXTMym69heG6Gi58Q12KCtDHp/GP5fv4smv9hBtMxv+ibKZSY61Bzz+m5dXG1Z1DiW9Xtchxl5NeqNAhCWChRct5ObPbsShef19hoJgVaxM7z2dUW1H1e8Am6hGuTubu2L/yWZqJ4SyTOz2qSz46QAv/rZ/tZ+xND2rxkDqdE6vyvLt/iDvX1f34n8/ZJ5R+U1V0ykoc7FyXyFjzzAvWGgeHru8O9fMWRN0SfxThXK+y5LEx1tymTJEFFE/l1W6fTz80Q6eumIWo7o8warcVby5802yy7Nx+pzYzXZSY1K5pcctDEwYKIIoQRAalK+aAhehrgzdNLQ9T1zZM+TPH5PWii935AesUIWSXi9J0D1RbMGoSVrBfl4vcjA9IRaX5gmo7ncqRVIwy2Z+1+933Nzj5gYcZdPS4MFUbomT7OLAzW2hLBNrOnyZkY+u64Y3FPsLKqrdj1VzHwKN5dvzibKa8KnGQVgwqwmVHpV5KzNFMHWO694mirlTB3DXmz+HtIIJoZ3vDo/K/B8PiGDqHPffr/YwqH0sY7v6rxujk0czOnl0I49KEATBr4XdbPh4KCtDsgQJUWfW5+mOkSn8sLvQcAIzmPR6iyIxdWh7LCaxB7la2enwye/oOuU9Pm+Zygf7PuCNjDdw+py4VTeqrp7MkFB1lYkdJnJzj5vp0uL8Ljff4MFUUYUbiyIHdOwOZZkYQNP8HbeNNlEuXHUQr0EwFFxalcpbG7INO4qHspqwOauE3BInSTFi79S5bHSXViyePpTpr6fjUTUq3cGtUoV6vueWBE5ACOeOnw8dZ/n2I3x9//mZIiEIQtM3vlsC6YeKA4KZUFaGwswKQ1NqLmJQnb7JMbSOtnLwmHGvqdrS6yXVw1TrKtA6Q5DfveeVokx457dw1RxIGkA0cFvP25jWYxrr8tax8/hOytxl2Mw2WttbM6H9BCIs519PKSMNHkx5Vd0wBzPUZWJZBq+mYaPqayvdPj7anMPpsVRoaYTGqwyhrCZYTDL7CypEMHUeGNC+BRseGc+3O48yZ2Ume/JKMfk8yHY7XlUzXCEN9Xz3qjqapot9eOcgl1flD+9v44kre4hcfkEQmqzf9E/in1/sMnwu2MJL8RFh9G/X4ow+X5Ik/nNNb25etD74StC/sJkVpvWNIfHgfNjzJlz+LCRVv1XkvFNRCIuvhnGP+Bs2n0KWZIYnDWd40vBq3iw0eDAVYzejGpRkCbU/j0/TibAEDv/LHfnIBql/oQRC1QllNUHXdSpcxmU8hXOPWZG5pFcil/RKZO97n5C1bitRt9xHVlElj3+aQcVpK1ahnu9mRRKB1DnquW/30S0xiok9Ext7KIIgCNWKtJqZ1LsNH27KwWj7VG0rQzazwozRKWe133Nwx1j+e20fHnp/a9ABlc2scHGPBP5wdV/gM9j6Drx9A3SbBOMeA9t53nPKUwlvTYZek2HAtMYeTbPU4Imj7WPtWAx65py6TOzYuxbN60JXfTgzf6b4h0UBr++VFG14c5lT7MBpUO0l1LQqI6euJtRGQsIebB8H4ZwSW5RH7zaRDO4Yy4jOLQN6UUHo5/uZ5pgLTdu2nBLe35jN41ecfx3jBUFofmaN74z9DCrLKrJEQlQYv+nX9qzHMKlPGxbeMogYm5nwsOrvs2xmmTCTzB0jU3j2+r7+IE6SoO+NcPc60FT43xDY9p6/zN85xuHxsSe/nJ8PHWdHbilFFQaFJFQfvHcrxHeDsQ83/CDPEQ2+MmVSZG4e3p55Kw8EpNMFu0wcHqZw9xjjRrvlbp9B55/Q06qMhLKa4NN0klqIFL/zkTc3D2v3bgAkRtvo3Taa9EPFAa8L9ny3mRVuG9GxQcYuNByPT+MP72/j0cu60yoyrLGHIwiCUKu2Ley8cdtgpixYb1im3IhJlmgRbmHpjGHBNwuuxYjUlvz86Hi+3eVPr9+ZV4ZFkZEkf1p8tM3MHSNTuG5gW+P0aXssTHoO+k6Bz+6HzW/CZU9Dy851Mr7GtCe/nAU/HeDTbXmYZP/fia6DR9UY2L4FM0Z3YmRqS2QJ+PwB0Lww6Xl/oCmckUYpjX7TkPbMXXnA8Llg+vOYZJnx3RIMn4uxmZElAkpnhppWZXSMUDZZtomx0iUhstbPEc493iNHiBx/4cmfZ47pxM63Nge0A4DgzndN17l24NnP5glNy8sr9pMUY+PKvm0aeyiCIAhB69euBR/ePZybFqzH6VWrLbwkATaLQoe4cN64fTAtI+p20sikyEzsmcjEnomUODwUVXrw/RJIxUeGBZcanzwI7lwBG16BhRfBoOkw8gEwN7/J8HKXl7ve3MjGrGK8qv7LlpqqixZrMovYml1CtM3MG713kJq3CW79AhTjSo1CcBo8mFI1nfgoK1OHtuet9Vkh9+ixmWUeu7w7JoNUQYAuCZHYLErA/9yhBEJmRUKrptV2MKsJ4RaFu0Ybr5wJ5z5vXh6mxF/3v4zuEk+0zYzDq4acSWA1y1zVN4koq7jQNRdun0qZ04fFJBMZZjL8Qt91pIw31h5m+X0jRb8oQRCana6to1j3pwv5dlcB81ZmsvNIWZWS4x6fxsjOLZkxuhMD27eo9+tcjN1y5gV8FBMMuxt6XAVf/hFeHgqXPg2dx9ftIOtRqcPLlf9bRV6pC081RdROqPSoODw+rlzVlrenvU7vMDHxf7bqPZhyeHx8vDmX+T8dJPu4A5+mo8gS8REW2sXayDruCLpHj82scOeoFK4dUP0s/biu8ZhlGTizPgTgb5J6Se9EvsrIN6zsV9tqgqJITOojZpvPR7qu4z1yBHObpJOPKbLE4ulDuOKl1VS4gy9KYjHJpMZHiP00zUCl28fHW3KZt/IAOcUOzIqM9kvkfGmvRO4YmULPpGgAfKo/ve//JqbROlrshRMEoXnyrwy1ZmLP1uSWOMktduLw+Ii0mugQF05cHa9E1buoNjD5Ddj3LSx/EBL7wMR/+x9vwryqxk0L15Fb4jTco21ER6JSt3LT2/tZPiuRti3s9TzKc1u9BVM+VeM/X+5m8bosJIkqubWqpnOkzE2xw4PHpyPrGiaTgqeak8BqltF1ePjSrkwd1qHGz61pTxYEl1bVu200T17bm30FFewvKA/65AR/wPfarYOxmkXxifORWlyMZLGgRIRXeTylVQTvzxzGja+so8Ltq/WcspkVureJ4rVbB4lzqQnTdZ0Xv9/PnBWZVa5zp157Ptt6hK8yjtI+zs7cmwbw5Y58om1mJg9MbqxhC4Ig1KmkGNu50wqm83h/gYqfnoE5I2DUQzD4Tv8KVhP05Y58MgsrDe8rKneuoCz9Y7xFOcgWG+b4FKKHTz7ZW7XC7ePZb/by9OS+DT3sc4qk15B3NHDgQP3nn38O+aBun8qtr6azOas4qFUniwThNjMeVUNCOrkHTtN0rGaFXklRHK/0UOFWkWWJlhEWrh2QzOW9Ew1vNI8WljLuqRVUSqGf+FazzKvTBjOsUxylDi9TF65nb0F5rSU4Zcl/Azz/loEM79Qy5M8VmqcKt4+PNuWweF0WhRVuvB4vVkc5QwZ05o6RKfRJrlpytaDcxYIfD/LWhix0Xa+yj+pEfnlsuIUZo1K4YXA7zNWks54NSZI26ro+sM4P3MDO9PpUVzRN5/53t/B1xtGg0pWlX64RigTLZ40iOVbMBArC6c6F61NjX5uEOnRsn79Ig7MYLnvWv8eqibn0hZ/YmVcW8HjZho8oXf8+cRfdg7VjfyTFhPPgRtzZGVUbKZtk0h8dL7YT1KKma1Odh9mapnPvks1sOlyMq5a8zRM8OuhuH2kJkfzlih44PSqF5S4+2pzH+oNFrDt4vEows78AtueU8udlO5g8MJn7x3ch2u4/CTzZ2VT+7j6e6jyA+60DQmrsZjMrPHRxGsM6xQEQbTfz3sxhLFp1kAU/HcTlC9xoaTPLaDpc1COB+8d3IaWV6AZ9Pih3efnH57v4eEsusiRVWXktN0ewfPsRvttVQGK0lUcv78a4rv6CKfGRVh6+rBsPXtyFL7bn8+ncd3H16IMlMoLW0Tau7p/UIPnlwtn7++c7gw6kwF9NyeFRsZrkc7EKryAIwrmnZWe4+RPY/j4svQnSJsKFf/FXA2wC9uSXc6CwIuBxzV1JyaolxF06G3var8127alDsKcOqfJaWZL4cGMO00TV4DNW58HU1zuPsjrzmGEgVdNyo1fV2V9YwdbsEvq1i+HuJTtxeHxoOsZLl7/cvC5Zf5hvduazdMYwYrb/TN6fHqblXXcx8aYptDpczLRX0/GqmmHK3wmKBGaTzKOXd2PKkPZVngszKcwck8qMUZ1Yua+QpRuyyS9z4VU1om1mxqS1YvLA5DPf+Cg0OwVlLibPW0teibPa1FRNB6dX5cCxSu5esonfX5TG9JEpJ58PMylc2SeRrutep8uza5DtYpWiOdlfUM5bG7IMJ2tqS6vwqBp//mQHr906OPgPVH0gySA3eGtAQRCE85skQe/roPME+P7v/gIV4/8KfW4IqZy4rus4PCqyJGE1y3UyafrTvkI0g9tbd+5udJ8He5dhtR7D6VX5cke+CKbOQp0HU3NXZhr2HqhuudG5b/3JmwyXV+PlFZk4Pb6gi1J4VZ38Uhe/eeob5qydQ/cXnsc+YAAAAzvE8v3vR7N43WFeX3MYn6bhcKsn+1DZLQqarnNFnzZMH5lSYylzWZYYmxbP2LT40P5ChHNKucvL5HlryS52oAa56Onyajz99V4irSauH9Tu5OO+wmPIkZEikGqGFq46hNfgBAjmOqfp/vK0R8tc1Tdj1jTI/A5WPwfZG0D1+h+3hPsrTg27199kURAEQWgYthi47Cl/09/P7ofNi+HyZ6BVWrVv0TSdlfsKmbcyk/SDxSd7PkkSjO0az4xRKQw4i2yUUocXj8F3keosQ7ZHBd1XtdjhPaPPF/zqNJjKLKxg15HAvM1QlhuLKz2GTXdrmu1VdSj2STx/3SO8+UsgdUJ8pJUHJqRx37jOfLe7gD355ZQ4PERaTSS1sHNpr0QizqCbt3B++ufy3eSVOA0DqZrOUadX5c/LMhjdJf5kBTdvTjaWtqJ/VHNT6fbx0eacgHMglOucBCxed5gHLzL4Et7xEXzxEHid4DktfcNTAVvehu0f+L/Ar36lxi9yQRAEoY4lDYA7foD0BfDqJdD/Fn+RCkvVidFvdh7lTx9uw+lRf90ffeIGV4dvdx1l9f5jtIwI49nr+zCgfeipg9UlKyi2KDRHGbqmBhVQiaSHs1OnUcQX24/80iSsqlCWG40CqWBme32SwvrcCnJLnIYVZUyKzMU9WnNxj9Yh/16CAL8Um9icY5jaF8w5CvDmukM8dHFXwL+/z5wsKro1N+sPFmGSZU5vhhjKdc7t01i2JS8wmPrpGVj5JPic1b9ZV/3PH9kK88fBTR9CuyHVv14QBEGoW7ICQ2ZA9yvhq4fh5SFwyX/9e6qA19cc4l9f7Kpx3/6JfbRZxx1MWbCeF27ox0Uh3qNGWs2YFSlgO0xYUlckkxnH3rWEd72g1uPUdUPl802dxqL5pS58BsFUqMuNpzox2xs7YSb2tOHIFiuSYsKeOqRKNRIAdJ031hw6w9ELQs0+3pyDbLAUH+w5glewmQAAIABJREFU6vZpvLn28Mn0MG92DpZksTLV3Byv9OIzSFIP9TpX5jotrWLTG/BjLYFUFbp/pWrx1VC4N8j3CIIgCHUmsjVcuwgmPe8Pqt6Zwudrt9caSJ3O5dW4753N/HzoeI2v03WdzMIKFq46yNSF63n66z2G991yWDgxF0zh+Ddzcexdi+Z1oas+nJk/U/zDoiqvDbcoNfZvFWpXp8GU1+AfFKouN4YqlNlej6rz9oaskD9DEIKxeF2W4X7AUM5RTddZd6AI8Kf5mduKlanmptTpMfySDPU6V6Win7MYlv+S2neKDs+VE//fcio9v754wSYPY16r/PVFnkr4aEZIv4MgCIJQhzqNg5lrcLXqw0PL9lZbnOjI67PJeuZacl6aytF3/4IrJ+Pk8/6AaguntyxyeHx8t+soj328g1H//YEp89ez72g5U4a0Y/0j4xn+SwXq00UNvpoW426ndO1Scl6cQs6caZRv+gxb59PuVSSJiT1F1tbZqNM0v+qWCUNdbjxVqLO95S4fqqajyKK0tFC3jlW4DR8P5RzVdCgo8x/Hk51DzLViNqi5eWu98YSNr/wYoJP1zLUo1oiAKn6nMysSLq/q75W3eXG1VaFUHZ5f7+HhkdWlYehQsNO/OtWqyxn8RoIgCMJZM1v5NPpGMG8Hb9WAKNitACUOD+sOFNEq0sqKPQWs3FvIpsPF9GobzZi0eObfPJC0hMgqBSvuGt2JzVklhpO9ET3GEtFjbLVDtigSNw5OJswUeuaY8Ks6DaZGdIrjtdUHqzQiharLjZKsYO3YD0k24Tq0BVfWtsB0vVOEvonOf4MSLopKCHXMqEQ/hHaOarp+svKOV+yZanb2F1SQUxyYhnfiizK8x1gcmRtpMd6/UuTYsxbd7Qy4zsmSvwhF/799Q6/ESBYef44Ir3F630PDLTy52s3dgyzEWKuZJNJ8sG4OTHr2rH9HQRAE4czMXZmJ47RAKpTiRA6Pyq2vptMi3MKYtFZMGdKel6f0J7KGhroXpLZkTForvt9dEFJqoSJJxEdZ+d24zkG/RzBWpxHHsE5xRFrNAcEU+Jcb5fAWlK5dyrHPnkKy2AhLSCVq2PU1HjPUVS1N17FbRIQt1L3wMIVSZ2D50FDOUZMsEWU1ozmdqKWlmOJFqf3m5NXVBwOK7Jz6RRne9QIqMn6gbP37eIuykSw2fMdzAq5zZkXmg5kjiIuwsGvLGixfVVKdgW0UxnQw8dQaN38fV10pdR9kfCCCKUEQhEaSX+oynGwLZSsA+HsRrvrDWBQluJ04kiTx3PX9uO21dDYePh5UayGzIhEbbmHpjGFE26oP1ITg1GkwJUkSd45K4b9f7Tb8x6xtudFIqKta7WPtddIITRBON6h9LJ+VHkE9LZ85lHPUq+r0bhuNNzcXc1ISkqhH2mxoms4Hm3ICNvue/kUZzHWuZ1I07eL8ZXQHtlTBbAZ39YUnnhgbxohFlcwaUkNzcHf5rw1MBEEQhAZVVOnGrMi4fVXvf0PdrqLIEhUelWhb8PcHFpPM67cN5p/Ld7Fk3WEkCcP7cLMiIUsSgzvG8vwN/YgNr+E7RQhanefCTR6UzKLVBzlS4gq46ayNWZFQNZ3T61gEu6pltyjcOSrlbH8FQTB0x6gUvt55FKf3zFde+ybHkBxrp3xrNmZRya9ZObEf83RnUq30rtGdfv1B89X6+p7xCpd3MfHvVR66taruCza0660gCIJQd6q75Q11u4qEhFZNQbeaKLLEY5d3Z/b4znywMYf5Px0kr9SJSfaXTreZFaYObc/UYe1JjrXXfkAhaHUeTEWEmVg6YxhXvLiKUqfXsGTj6STJHwj9+ze9eeiDrYY5n8HM9uo6XNUv6YzHLgg16ZkUTdsWNvYVVBg+X9s5Gm5RTt5Ee7NzsIhKfs2K06uiyIH9PEL9orSZlaq98GwtCCYQ+usYK/3nVfDgsGoKUZhsYlVKEAShkUTbzIZtM0LdruLVNKLOIvUu0mpm2oiOTBvREVXTqfT4+NfyXaQlRDJtRMczPq5QvXrJMUqKsbF81khSWoVjtyjU9PUeHqaQEGll2T0jmNS3Df3btSDMFPqwbGaFW0d0wG4RhSeE+vPIZd2wmkM/P02yhNWs8Mw3exjz1A9cu9vG/fTgyx35+NTgN4wKjSfSasJnUITk1C/KYEiS/1gnJfQAvfZzIDVW5voeZl7Y4DE6KrQfbvC4IAiC0BCSYmxEhgUGQaH0fALonRRdZxWplV/2aXdsGU7W8WB7GAqhqrcNGwlRVr6aPYqFtwxiTForLCaZSKvp5J8wk8zA9i14/vp+rP7jOFLjIwGYN3UAbWKsmJXgTySrWWZ4ahy/vyitvn4dQQBgTFo8f5zYNaSASgJ8mk6Zy8v23DIOHXNwQLOx1hHGg+9tYeDfv+XZb/biMkgfFJoOu0Ux3Kgb6hclQHzUKatLZhv0uwnk2mci/zw6rErPqZMsdhgxO+jfRRAEQahbsiwxfWRHw/uDYHs+hYcpzBzTKeD9Z6tdbDhZx6svdCScnXpdxpEkiWGd4hjWKY5jFW4OF1VS7vIRHmaibQsbidG2gPdEWs18fPcF3LxoPfsKKgzr5p88PmCzKFzcozX/vbY3sugtJTSAaSM6EmU186ePtgMEbDY93YlbX6PS6pVuFVCZtzKT73cfZfHtQ4m2i8o6TZEkSdx2QUde/G4frtP+zYPdM2dGZXLC0cCeHkPugo2vg1a1WuSh2ZFVfk6OlnE9GhU4OGsMdAith58gCIJQtyYPTOaZb/YaPhfMdhWTLDO+W0Kdj6tdrJ2s4446P67g12A5cS0jwqpt6nu6aLuZD2YO55udR5mzMpO9+eXogMenIUkQZlJQdZ0LUlsyY1QKgzvGigp+QoO6ekBbLujSkiXrsnh9zSF8moYOqC43ismEBwmfple7IfV0Lp/G7vxybpy/jg/vHu5v5Co0OTcObscL3+0zfC6YL0rZZOLWsnlwZAAk9vn1idgU6D0Ztr8H1fSbqo5usiJd9ozYLyUIgtDIWoRbuGdsKnNWZBoWq6qJ1SzzxBU9MAVZEj0U7eL8wZSu6+J+uR402Q1GJkXmkl6JXNIrkf0F5azeX0Sp04si+2vjX9gtnvjIanquCEIDiI+0cv+ELvxuXCobDh7naLmLgvc/wtYmkX/mhxuvRO1cQVn6x3iLcpAtNszxKUQPn4y1bQ+8qk5mYQX/+mIXf72iZyP8RkJtYsMtXN0/iY8254bUHBEgzCQzsnNL2veZBcvugTt+ANkEuZv8PaJcpWCJANUbVIU/AI8UxtsRt3N9ygTE1VAQBKHx/W5cKnklTpZtyQspoLqqbxJX1lMRtYgwE+EWE4XlbuKjxLdFXWuywdSpUuMjT+6pEoSmxqTIDE9tCUDB9x7eqXQjERHwurINH1G6/n3iLroHa8f+SIoJ58GNOPetx9q2B+BPGXw3PYc/TuyGTTSfbpKeuLIne3NL2J59HE8Q+5zA3wOkfZydF27sB+aB/hWo926Fwl1QlvfLalQIpXDNdkBHmfQS6TtSWLF4I3OnDghMHxQEQRAalCRJ/OvqXiRGW/nfikxkCEgNP8FuUdB0nalD2/PR5lz2F5TX2/1uuzg7h487RDBVD0THUEGoQ3JUNG8WhwfMRmnuSkpWLSF2wkzsacORLVYkxYQ9dUhA42lJgk+35jXksIUQKKqPf214lYEWF/YgAl67RaFXUjQfzBzurzbqqQBnCez+FIr2g9dBzYGU5F+xCovyB1Ex7eCif8Dv96H0vpZnr++LxSRz71ub8YrKkIIgCI1OkiRmje/C2j+O474LOxMXbsFuUU4WYbOaZdq2sPGnS7ry86MTeOSy7vzfxK5MezWdgnJXvYypXaydrCKxb6o+NIuVKUFoLnaYWlCpB+Yju3N3o/s82LsMM3hXVQ6PyqLVB5k8SPShamp0XSf/scewx0Ty5uPXsXL/MeatPMCW7BJ0txuP4l+psphkZCCtdSR3je7EhO4J/jx4nxteuxwKdoX4yRJc8RLEd4OWnavsjzIrMi/e2J+Zizcy653NvHBDv3rJuRcEQRBCExcRxt1jU5kxuhOHiiopdXqRJYkWdjPtYu1V9i9dNzCZ3BInt7/2M0tnDK3zVj/tY/0rU0LdE8GUINShQlM46C5Ob66mOsuQ7VFBNXUFOFpWPzNTwtk59tL/cB88RPvXX0M2KYzrmsC4rgkc3J/De3/4D6a77kXX/ZuQR3VpRWr8aemenz8IhbtBdZ98qMNz5Ti8cHBWBOEW/4mzYJOHxdu8rJgWDuj+1asf/gH3rDcsNGExyfxvSn/ueONnHnxvK89M7mvcp0RT/XuzVA9Yo/1l2QVBEIR6pcgSnVoFpv+fbtaFnckpdvK7tzYzb+qAOp0YS461syazqM6OJ/xKTF8KQh1yh9nRDDK2FFsUmqMMXQtuM6qnlnLrQt3TdR29hvKLJR9/TOmyZSS//D9kW9UgJL4gi+vDS3ngojQevDiN2y7oGBhIOY7790r5AgNlVYfn1xs14z0xOBVKc+Dw6mpfYjUrzL95IAVlbv7vg21op56IuRv9e7T+ngBPp8HzfeCfSfBcb0hfBO6K6j9bEARBaBAn9lt5VI3HP82o8TspVJ1khejsCio3HsW54xjefNF3qq6IlSlBqEPRMeHI+rGAx8OSuiKZzDj2riW8a+39gETxifqn6zpbskuY/+MBVuwtPLnPLdxiYny3eKaPTKFnUjQAles3UPDfp2j/xuuYWrYMOJZrz17CunSp+QO3LKm2fPlDwy08udrN3YMsxFirKVvrdcCaF2vsJ2U1KyycNpBbFm3g0WU7+McFFqSlU6E02x/E6acF6SWH4etH4auHYcQsGPNHUWJdEAShEZkVmZen9Oe6uWuZ9+MB7hp95k18dZ+Gc8cxylZkE3/MxY0+jZJl+/3XeU1HaWElcnRb7L1bIRk0GxaCI4IpQahD3dq1xEd2wONyWDgxF0zh+DdzkWQFa8d+SLIJ16EtuLK2BRSh6NEmuqGGfF7acPA4f/hgK0dL3bh9apXVxAq3j0+3HuGrjKMkx9r4x7A4oh98gKSnnyask/GXmnvvXmx9+9b8oeterraH1MA2CmM6mHhqjZu/j6uu0pIOmT9AZRGEx1X7MXaLiUXTBvHE3DfwbH8Ui+ZEqqnAhfeX2ck1L/gLYlw9H2TxpSoIgtBYIq1mXr11ENe8vIY2MTau6NMm5GN4j1ZSOH87ukdF9/gn0sKRTv43gK/AQcmy/ZR+foCW03thaVN7KqIQSHxjCkId8Pg0lm3J5YlVeUinz/7/Imrw1bQYdzula5eS8+IUcuZMo3zTZ9g6Vy1KEW5RmDEqpSGGfV76bFseNy9az6FjDpxe1TAtU9V1nF6VvUcruOmD/eyZ/hDhQ4dUe0z33r1Y02pYmdJ1KD9S47ieGBvGixs8FFbWkOJpCvOvJtUisjKLJ51/IUxz1BxIncrrgD3L4etHgnu9IAiCUG8So20snDaIv36SwfoDoe118uRVUPDyVrQKb5XgyYju0dAcPgrnbsWTXX42Qz5viZUpQTgLZS4v//t+P0vWZ6GjU+lW/Te81YjoMZaIHmNrPGak1cywTtWvPAhnbvX+Y/z+va0hNdx1K2YeOiCTlFVM/3YtAp7XVRV3ZiaW1M7VH8TrBEn2732qRs94hcu7mPj3Kg/dWtUwz+UJYn/T8t8jearmw9de6AJ/QPXzIhhwK7SqJW1REARBqFfdEqN47oa+3PPWJt65c1jgXlwDarnHvyLlDr5hMPiDqsKF20m4fwCm6OrvY4RAYmVKEM5QXomTy57/iVfXHKLC7fMHUmfJZlZ46OK0KuVShbrhUzXueWuTYSBVuXMFR16fTdYz15Lz0lSOvvsXXDkZJ593eTXuXrLJcDOw53AWplatUCLCq/9wkzVwv5KBv46xMn+Th9yyGlaTLLV8mZbmwqHVhp9Xa6ELAM3nT0kUBEEQGt3Izq34v4ldufW1DRSWu2t9fcXqXHSv8f3IhpxtXPXmTLo/ewk9n7+M3yy+my1Hfm3VoXs1ylcGblUQaiaCKUE4A8WVHq6es4a8UledVd6zmRV+O6Qd1wxoWyfHE6r6dlcBXoN/q7INH3H8u/lED51M23sXkzTzVSL7X4pz3/qqr3N6WWtQVta9d0/txSdkGSJa1zrG1FiZ63uYeWFDNQGPzw0x7Ws+SPqCap96aLiFp9a4KXHVEKxpPtj2DnhEpSdBEISm4LqByVzdry23v56Ow+Or9nW6qlGx7gj4Aq/x5e5Kbn3/j0wbcA3bZ31G+t0fMnvENMIUy68vUnUcPx+tNhgTjIk0P0E4A7Pe2UJRhRvVYMNN5c4VlKV/jLcoB9liwxyfQvTwyVjb9jA8liJJmE0Sd45KYfb4GlLFhLMyd2UmlZ6qXxCau5KSVUuIu3Q29rThJx+3pw7Bnlp1j5TDozLvx0z6t2/BD7sLyC1x4vSosOkonTv0pq2u17yiOGQmrPgX+IyLUJzw59FhvLnNa/CMBB1H1Vh8AoAdH1TpY3Wq4ApdALLJv7rV5aKaP0sQBEFoELPH196DyplRRHXbZA8c9684XdV9PAA2WWF0x8GGr3VsO0b4gIS6Gfh5QARTghCi7OMO1h8swqsGXrHKNnxE6fr3ibvoHqwd+yMpJpwHN+Lct75KMCVL/jLWqqZzaa9Epo/sKCr41aNSh5eMvNKAx925u9F9Huxdhhm8K9CP+47R/29fI0sSbp+GqumYtBaYTArR//6eO0emcM3AtkRZzYFv7j8VVvwz4OFDsyOr/JwcLeN6NCrw/Wa7v3x5bVyBv+epnhgbxohFlcwaYqn+RboGzuO1f5YgCILQIE70oLrttXQe/zSDv13ZM2ACz51ZUu1eqZTYZGRJ5v7P/8EVXS+kX1IPYqyRAa/TPRru/SUimAqBSPMThBC9ufaw4d6ZE6scsRNmYk8bjmyxIikm7KlDAkqft4u187cre7LhkfE8e31fEUjVs6JKNxaDWTzVWYZsj0KSg+vrpevg8GhUuFW8qo6mg0dScKhwpNTFk1/tYdSTP7AzryzwzfZY6HmNf/9UiHzIuMNb19hj6pRR1vjsqYUuaj6MaBwtCILQlFhMMi/f1J+fDxXzyo8HAp5XK42yGvwiw8L5cMpLgMQfvvwvfV+4gls/+COFlYETZzUdRwgkgilBCNHb6Vl4DFalQlnlyCt1MbZrPNE2gxUMoc6pmg4GGXiKLQrNUYau1U1+uNOrUuLwcu3cNcYB1WXPQMsuoNSwKnQ6SUYzR/Kbst/zZUZ+7a8PC5xpPF2thS4kGWyBlQsFQRCExhX1Sw+q19Yc4tOteVWek801Twx2btmBZy97mPR7PuDb21/jaEURj3/3YsDrRAPf0Ii/LUEIgcenUek23vwZyiqHRZE5Ulrz3hmh7kTbzPgMAuCwpK5IJjOOvWvr9PMcHpXfLlhHqeO02T2zFaZ9Dm36+tP2amMKg4h4LHf9wL9vu5QnPt3Js9/sRTNqjnVC18tBrjlIr7XQheqFdsGlPgqCIAgNKzHaxsJbBvH4JxlsOPjrypISExb0nX1qXHsm95zInsLTVrikX44jBE0EU4IQAqdXRZGNiwyEssohSf/P3n3HR1VmDRz/PfdOTS+kQQIk9N6rVBV7b6hYEBDUVVd31y3uvu67zfVddd1dXVGxIboq6toVO6CANKV3QiAQIAkJaTOTaff9YwIkzCSZSaGE8/18+EBm7tx5JiT3Pucp5xBIXiBOiJRYK4lRwbNBmjWahDFTKPniGRzbluH3uDB8Xpw7V1H6zYsNnrPxdOo+3li5J/iFtjiY+gmc/1dI7FwTVB33M2WJCcwMnXU/3LkMkrvQPzOB9+4+i+92FHPHq6uprCeoZ8QsCCOgf2i8lSp3iKBM6dD7crAnNHoOIYQQJ0fv9nE8MXkgd722mh2FgfqDUYNSA9ljQ9hxaDfPrniD/eWFABSUH+T9zV8xuH3d5FjKpMl+qQhJAgohIhBjNYWc4YC6sxzRPRve22IYgeK84sRQSnH72Gwe+3wbzuNSvsYNvwotOpGyZW9S/NFjKIsda1pX4kZNrvd84SQacXn8PP/dLm4fm4N2fACum2HoVBhyK+xdCRvehvL9gbTk0SnQ7TzofgHodS/RqbE2/nP7CH7//kauenoJc24ZSqfk4+pbJeUEZr72fF/n4bATXZgsMOruej+7EEKIU8O47in8sqYG1X/vPIuU1CjM6VF49gYXd4+2RLGmYDNzVs6nvLqSOGsM53YZxW8n3lXnOD3JhqV948WBxTESTAkRAV1TZCTYKDjsCnqu9iyH0nRs2YNQmglX3hpce9bVSULh9fvJTLSfyKaf8a4ZmsXfPtsa8rmYPhOJ6TMxrPNElE692suSncWM7ZYS+mRKQdbwwJ8wWU06f72qH/O+383Vs5fyj8mDGNOt3bH2+Q3W9v0dvfKvwWYE/5w2yGSHXpdCRv/IXieEEOKkuG5oFvtKnUyfu5I3Zo4kbkIWJW9uxjhulXlGbAqzr/hDg+dSZo24CVmt2Nq2SZb5CRGh28fkYLeEXkYVN/wqEs+eTtmyN9n75BT2zp5KxQ8fYe92bP+Jriku7d+eaKuMZZxI8XYzd03ogr2RDbqNiSTRSLXXz8ZQiSiaSSnFLaM68+QNg7nvzTW8+N0uDMOgsMLFBf9czJSPqri9+j4cRvjr3l1Y8WeNgMufbvH2CiGEaD33nduNbqmxPPjqQqzLrsdmLAEiHUxTWLsmYB9Qz+CfqJf05oSI0NVDM3lkwZZ6n29slsOsK6aPzW6NpolG3HtON/aUOPhk/YGg5X7hiiTRiNdvBCehaEGjuiTz7l2juf2VVazeXcr3uw5R5vDg9Rt8S39ucP+WOZbHiaKaGBX6xurTrGgKvo+axOLEX/KQLstPhRDidKKU4q+TUqh46ir8vsMk6T9yyP9rqv0DMWi8HIcya1g6x5F8Yy9UPfvCRf1kZkqICMXZzNw4oiP2JqQOtZg0BndMpGd6iL0qotUppXjs2gFMH5ON1aRhNdX/f6jqqbMUSaIRTUG0tXkzYY3JSoriP7eP4JuthRyqdOOtlelvrdGVkdX/5m7PPXzv60m1YaLSsFFh2HEYVkqNaGZ7LyXvpmUMuvNlvthyiE/W72/V9gohhGhhbgeWVy8jyTiMCS9KeUk2/4VY/Q0UlSgcIV+mcKCUg9hRybS7ra+kRG8imZkSogl+e1Evth6o4Ic9pbg84RU3NeuKjHgbz948pJVbJxqilOIX5/fgppGdmPd9HvOW7cZnBAIfCOw5Mps0EkoPkmdNCnp9JIlGbGad9PjW3xv35abCep+r2LSY11e+xyuH9qJbbMSlpdBt1Hnomf05QBKG0tm05DBP39SFp28cwq0vraBneiw5KbIBWQghTgs/vgrle1H+Y1les/9ZjsPzErn3/hfNNJJK75W8smY972xczDs3/gVdHSDW9B5282qU5U7QHjqJH+D0JsGUEE1g0jVeum0Y977+I99uL8bRSJpzTfOQkWjhnTtGSBa/U0R6vI0Hzu/Jfed2Z/P+cg47PCgFiVEWeqbHsmj+Au7+0Y3zuJpNkSQa8RsGF/RNb9XPYRgGTy/aEfJnsL6sg5u3byQx85yaE8CXWwopqXLTLzOen03qzl2v/cC7d51V795AIYQQpwjDgKX/BE9w7UqfAf9a4eTBsUuJ0peSYHZj1Ty0t005dpAfWPk8TPxNINOsiJgEU0I0kdWk88xNQ3h7zSYe+uRbnJXpgAGGmUDdIC8oP5qlGEvyIlyJO5n86eM8f97zZMfLnqlThVnX6J8ZXFNp4jWTiF3xNk5r8M0lnHTquoLLB3YgppUTjazdW8bBsuqgxyPJOqgBb6zYw10TuzJlREdW5pXw0PsbePTaAUePcbp9fL/rEKVVbvwGJNjNDOucRHyU3HzPFHtLHby5Mp/thZVUubzE2c30y4znuqFZJEUH13ETQpwAu5eAozTkUw+MtvC3JdXcNcxCgq2BvVCGD7Z8BH2ubKVGtm0STAnRDLllufx98wwsmQ40Txzeit74vdGAhqZXoUfvQLcdAMDpB5ejihs+voFXLnyF7ondT27jRYM0s5lpnXT+sd+LSwVfKhtNNGLSmD6m9YPmlbtK8PqDl5pGknXQ5fXzzdZC7prYFaUUD1/Zj8v/vYT5K/MZ2jmRF5fs4p3V+9A1hWEYGICmFB6fnwv6pHP7uBz6dohvhU8nmsIwDEqrSymvLsekmUi0JRJtjm78hfVYsqOYp77ewQ97SvEbBp5atfa+2nyQJ77Yxtk9U/nJxK7ycyDEibb9S/BUhXxqaHudCZ1NPLa0mj+f3UAiCnclbJZgqqkkmBKiiUpdpUz7bBpVnioMDDTzYSxJSxt8jYFBlaeK6Z9N593L36WdvV2Dx4uTa8aMi1h43xzWpPXAVU+x5lDsZp1fXdCD7mmxjR/cTIednjqd2yMiyToIUOY8lnUw2mri6RsHcdlTS/AT2EdWO7FFbR+uK+DzTQc5p1cqf79uIJYGknqI1lXhruCDnR/w8oaXKXGVYNJMGBh4/B4Gpgzktr63cVb7s9DD/JkwDIMnvtjGnG931Zv90uUNBPKfbTzAN1sLefiKflw1JLPFPpMQohGVBxt8+o8TrZz1YhU/HdHI7HFVcQs26swiwZQQTfTa5teodFcSGKc/pvTbUoo/K8Zd6Ea36cQNiSPtmjT06GMdmCpPFfM2zeP+Ifef6GaLSNijuCQnhg3lHsK9XNrMGvee05WpZ52YpZwWXaGA40Od2lkHwwmozPqxIMgwDJ5dnIvfMHA3EkT6DXB6fHy5+SC3vricedNHYNIloDqRDMPguXXPMWf9HDSl4fQG9k64/e6jx6w6uIpNhzZhN9l5fMLjDElrPBHO459v44Xv6g+kavMb4PL4efC99WgaXDFIAipBiLq0AAAgAElEQVQhTgit4ett31SdS7qbeOQ7N71SGji2kfOI+sl3Togm8Pg9vL7l9TqdFYDiT4s58NYB0q9Lp/fTvcn5nxzch9zkPZaH3+uv8/r5W+fj8bVeDSLRPIECuN/ysDOTysYCKcPApKBP+zieuWkId07oemIaCbSLtWILkc62dtbBcKTFHVsC8vTCnXyy/kCjgVRtLo+fNfmHefDdDWG/RjSfYRg8tOQhnt/wPNW+6qOBVCgOr4NDrkPc8cUdfL376wbPu3BrYb2BVNWmheyfex97/n4Ne5+6mYPzf49r70Yg8HPwm/+uZ0dhZfM+mBAiPLEZoBruzv9hgo05P7jZV97ANT0mrYUbduaQYEqIJliUvwhvrRSkAD6nj8L3Cml/U3ti+8eiTApLioWsu7JwF7spW1pW53i/4efLPV+eyGaLMBVVVHPJv74jr7gKRzip75VC1+CifhlM6JHa+g2s5bze6YRagVc766Bj2zL8HheGz4tz5ypKv3mxzrHRVp3Jw7KAQKKJp77eUe9sREMdaafHz/tr9rG/rP4OvWhZT/74JJ/lfYbLG7oocygun4tfffsr1hatrfeYf365PeTPQPmKdyn5ag7xI68j8+5X6XDnS8QOvgjn9uVHj3H7/Dz/bW5kH0QI0TQ9LwGTtcFDuiZpTO5j5l8r3KEPsMRAv2taoXFnBgmmhGiCZQXLcHjrFsFzbHfg9/iJG1K3IK9u04ntH0vlxrojtQ6vg6X7Gt5jJU48n9/gxjnfU1LlrnefUCjVPnjy6+18vvFAK7YuWEqslXHdU1AhEjXFDb+KxLOnU7bsTfY+OYW9s6dS8cNH2LvVTUph1jTO6RkIAj9cWxDyXBBeR9oA5i3b3VIf77RnGAZOt4+qai+GEf7PUzgOVB1g7sa5OH3BwWvpt6Vs/912Ns7cyJZ7t1AwtwBf1bHgyOVz8fslvw953tyiSjbtLw96/EiGyKRJdxLVYzSaxYbSTUR1HVGnJIDPD++t2UdVtTfoHEKIFtZ+ICR0avSwh8ZbqXKHvgaV+cxsiR7a0i07Y8ieKSGaoKS6JOgxX6UPU4wJpQf3RE3xJpy7gzs8oc4jTq5vthRScNgZMpCq2rSQ8pXv4Tm0F81ix5yaQ/zo67Bl9gECS5z+8slmJvVOQ9UXkbSCmeNy+G57cciZhMayDlpNGlNHdz66z2n2op0ha1aFm2rd7fUz7/vd3D+pe519WK0pv8TBjsJKKqq9RJl1spKi6JHe+sk/6mMYBj/sKeW5xbl8vaUQvx+UCtQdG9IpkTvGd2FCj1R0rXk/I29seSNozyYElhsXfVpE5oxMYnrH4Cn1UDCvgLzH8sj+bTZaTZKQfZX72HRoE72Te9d5/WvL9+AL8fMfSYZITSk+Xref62pmPIUQrWjM/fDR/eA5Nsibd1/da2BWvIbrd3HHvxLDZGNzp5u5+4VVDOucyL3ndKNXRvBxon4STAnRBGYtuLaOHqPjrfRi+IyggMpb5sUUE/zrZtGkNsup5plFO6mKoACuc/vyo8EUQGF5NT/mH2Zwx8QWa5Pfb6A10PEe1jmJa4Zk8vbqvWElCzjCjJecGI07JnQBwOPzk3codIrdSDrSPr/BvlInnds1PR13OO/xzZZCnlm0k/X7yrCYNAyjpsKb36B9go07J3Tlkv4Z2MwnrvjwD3tKuf+NNRRVVuP0+Dg6GVXz98q8UjYV/IjVrPOXK/pyYb+MJr2Px+fhza1v4vHX3Xd5ZLlxh+kdiO0f6EwdWW687YFtlC0tI3Fc4GfT7XfzysZXeGTcI3XOsfVARcjBhEgyRDrcPnKLZd+UECdE32tg9VzYtxp8wXUH66WZUYnZjJz8GxZj4bXv93DLiysY0jEQVPVuL0FVOCSYEqIJMqIz0NDwc2w/TVTXKJRJUb66nPjhx2qt+Fw+KtZVkHZN3c2dmtJIi5YNn6eSPYccrN9XFvR4JAVwq70+nl+cy9M3NZ4trT5lDg9vrc7npSV5FFa48PgMLLpGVpKdmeNyuGxAB+yWuh3aP1zWhwqXh882HgwroLKaNDrG2fmP+jW2dSUw5FYqXF4suka1N3ifWCQdaV0pyl2tl1xlb6mDG+cs51Bl9dHA9/g27yyq4qH3N/CnjzYxd9pwBmYFF2ZuaV9vOchdr/2Aq5F9dlVuH1VuH/fPX8O+w05mjM2J+L02l2wOOSsVznLjI8GU3/CzMH8hOworOOzwUOYM/NlVTxAUaYbIww5JsCPECaGbYMp8eOkiKN4G4eyh1C0Q1wFu/RAsUUQBt4/L4aaRnXht+W5ufWkFgzsmcO853ejTXurHNUT2TAnRBBdlX4RFrzurpEfppF6RSsGrBVSsq8DwGriL3OQ/nY85yUzC6LqdOatu5dIul57IZotGrN9XFnJpWiSzMn4DFm4r4vHPtvLJ+v14fGEksKjh8vh44K21DH/4Sx7/fBv7DjuP1pBy+/zsLKriDx9uYvCfvuCRTzfXWYqlaYonJg/k5+d1J95uJtoSurNrN+tYTRqXDWzP+/edTeK0t2Dxo7DsaSwmLeTyLqjbkW6M12/QwtuDjtp9qIpL/vUd+0qdIWcQa3O4fZQ5Pdzw3Pd8n3uodRpU44c9pWEFUrW5PH4e+3wr7/+4L+L3O1x9mEBS/LoaW27sray7j6nK4+T2eav4yyebmff9bhZtKwqZ0AQizxCZGCUz70KcMNZYmP4F9LkikJDCVE+RXt0a+NP1XJi1GGJS6jxtt+jMGJvD4gcmMqxzEre9tJKZr6xiQ4iBxnBVVXvZVFDO8txDrNt7mMKK8BPmnA5kZkqIJuiR1IOs2Cy2H95e5/GUi1LQo3UOvHkAd6Ebza4RNziOrFlZaMelr06PSqdPch/EqaPC5QkZTERaANfh9vHkNzuItupoSnHzyE7cOrpznfTjxytzeLjuuWXkFVeFnBmqfW6AuUvz2LCvnBemDsVqCrRLKcWMsTncOrozX20+yDOLctlRWInL48Osa6TEWrltdGeuHppJnK1mqWpyF6qmfMgPLz/Aig1R+PwpId+3dkc6uueYBj+/y+Pj2meWkpUURa+MOHplxNE7I46eGbGkx9mavJ+swuVh8nPfU+7y1NvhD8Xp8TH95ZV8fO/YVll6aBgG972xpt5AqqG9di6Pn1//dz2T+qQRZYnslhwqYI10ubGuKb7+2YQ6/yd//WQzLy7ZFVQMunaGSKXp2LIHoTQTrrw1uPasq5OEItqi0zU1JqLPI4RoJrMNrnwWzvtLYNnf8tngOES1oWPCh2GJwTR8BgybAfEdGjzVkaBqyohO/GfFHqa9vJL+mQncd243+nYIb6Zq8/5ynv82l4/X78ekaYHkRgZU+/wMykpg1vgcxndv/v7Rk02CKSGaaHq/6fxh2R+C6rokjU8iaXxSg6+1m+xM6zetwWPEiWc1a4S6pke6vOmIqupA4PP8d7uYuyyPl6YOZ3h28M9GtdfHzS8uJ7eoMqgDWx+nx8+q3SXc+/qPzJ4ypM6eKrOucUHfDC7oG3o/TmGFi++2F7Myr4RVeaXsKKykT9rdDCv8nJGJFpYfjg8KViLpSA/qmMAbM0exs6iSzfvL2by/nBe+28Xm/eX4DINe6YEAq2dGLL0z4uiaGhPWvqa3Vu2lzOEOGUg1mhzE6+dfX23n75MHhvHdjczq3aUUV4bepxDOXjul4P01BdwwvGPIc5Q5PGwrrGDbwQq2Hahg68EKtpRuw5fqCVpfEuly4yhTVFBwO2VEJ15emkdwKehAhkgtOpGyZW9S/NFjKIsda1pX4kZNrnOc3wiUChBCnATR7WDczwN/vNVc+shH+EzRPDttTMSDHHaLzvQx2UwZ0ZH/LN/D9Lkr6dchnp+e051+maGDqgqXh1nzVvPDnlI8PqNmkLLuYNPyXSVs2FdGjM3E3GnD6Zl++u7PkmBKiCa6MPtCPs79mBUHVlAdwYZPi99goNPJpYl9W7F1oinS6pk1iWRWJhS314/bC7e8uJxXpo0ICqjmLdvNtoMVIQOpxmY1vt1ezJebD3Jen/SQ720YBrnFVazKK2FlXimr8kooqXIztHMSQzsn8tClvenXIT4QzDj6s/H5WVxTdjNOIzi4CacjHW3VuWN8Fywm7eisVO22FFVWs3l/BZv3l7NkRzEvfLuLvENVdEqOomdNkNWrJshKibUe/f8wDIPnFufiDDH7E07A4vMbfLx+P/97eZ9js3It5NnFuSH3qYW7187h9vHMwp1cNiCDHYVVbD1YwfaDFWw9WMm2AxVUuDx0S4ulR1os3dJiOLd3Gl1S+nHNpy9S4a577am93FizaXWy+R2/3FhTGhOyJgS1u2NyFAOyElixK3S20cYyRJo0xTVDMoP29QkhTgKTlXItgSqnl+Topi+9tZl1po3J5sYRHXl9xR5mvLKSvu3j+em53eifeey6ctjh5op/L6GgzIW7gVUWcGz/6FVPL+XVGSNaNHHTiSTBlBBNpCmNv0/4O3d9eRfri9bi8tdTDK8Wm99Pr2o3/yw8hD7nbJj+JaR0PwGtFeEYkZ2MOcRek0hmZRoLfm57YRlf3T2K9PRAQOX3B4KEUEvEwgkSHG4fzyzaeTSY8vj8bCworwmeAjNPVpPGsOwkhnZOYsbYbLqnxobODhiVRJ+Zz9Pp4ffY6k7CCLGttvFU6zpn9wxduFgpRWqsjdRYG+O7H1tOWO31sf3gkVmsCr7dXsTm/eVoStEzI5Ze6XGYdMVhZ/DvWCTJQTSleGfVXm4bk11v+yNlGIGsgqGW3EWy125PiYNBf/qCrimx9EgPBE23jupE97RYOiTYQ/5/Xd/jeuZunIv7uGtPuMuNLZqFW/rcErI9953TjelzV0WUHfIIs64xvQW/x0KI5lEYODw+4u3NH0iymXVuOyubG4Z35M2V+cx8ZTW928fx05rsfze/sKLOft9wONw+bn1xBR/fM5aOyVHNbuOJJsGUEM1gM9l4bvRf+Ne8CbwRZUIBDi24A2r3+zGAa8srub/0MGYAVzm8fDH8ZDlENbwsUJwYuqaYOjqbpxfuCNq3FM6sTDjBj8ft5am7/8K06u3YBw1ibfYgKl3Bo4WRBAnr95Xx0Psb2FFYydr8w2QlRTG0cyIX9cvgoUv70CHBHv43wRbP03dcxGVPfUelP7JRTJvyMueidkdrVoXLatLp2yG+zjp8wzA4WF4dCLAOlPPej/tCBpyRBCxOj4+F24rqDaYMI7AcxXvkj8+Px2fg9fvx+kI85jeodHnx15NtI5K9dlEWnVenj2BQp/BHZif3mMzcTXNDPhfOcuOs2Cx6JvUM+dzoru2455yuPPnVjogCKptZ4x/XD2zVtPhCiMbtKKzkpSW7WL27lIMVgRnsqS+t4JZRnZnYs/n7lGxmnVtHd2bysCzmr8pn1rzVJMdY2FnPcvXGlmJXVXt5/Iut/PP6Qc1q18kgwZQQzWRa/gw/KynhzmIPC6KjmBsfy6e/2YXP7af/37qRqRtMLSvnwOLDzF/n5pdTj3QyDHCVwYrnYMKvT+pnEMfcOKIjsxftCPlcQ7MyYRe11Ux81O98fnXFVDxr1/De+mIcWnpg40wtkQQJHp/B1gMVzBqfw5COScRHNW/0Mad9Cv+ZNYYpz35Hld+EP4zEr3aqedr0D4Ys2AZRc6BX8zJVKqVIj7eRHm9jYs9UDpa52HYwOGV3pMlBlu4sZvRfv8JTExgdDZL8gSBJU2DSNUyawqQpzLqGSVeYtCN/H3tM1zQ0QieCgMj22mmaIkRyvgalRacxve90Xt74ctDezcbYdBt/POuPDR5z14SumDTF37/YRrXX32CGRl2BxaTzxOSBnF/PklMhROtbuqOYRxZsYVtNvbjaNeMWby9m9e5SLCaNGWNzmDUuJ+LBr+PZzDq3jOrMdUOzmPjYwiavsvAbsGDDAcocnmbfw040CaaEaA6vG1a9AD43duDKyiqurKyis9dLhReueWsXD461AvB8iM3c+Kph+bMw9heBOhHipEuJtfLoNf154O11EaW5jij48fv53pzKpFtupmzO97AzOG13pEFC97RYzu7ZcnXL+ndK4eO7R/M/T8/le283DMBN3ZkqMx40DPqoPP5ofpm+Wh54gP/eDpNfg67ntFh7LKbQN/xIk4P0TI9j9k2DAwGRpjDpGmZdoWsKs6Y1WBw5FMMw6PLgJyEDjUj22vn8RpOW4Nw54E6KHEV8vOvjsAMqm27j8QmP07dd4/s2Z47rwtDOScxeuJPF24qAujW97GYdA4PLBrRn5rguksFPiJPo5aW7eOTTLQ3eu47sU3ry6+0s3lbEi1OHEW1tfv9jT4mD0qrmL8V+a3V+k2rvnUzSexOiOTZ/UO+w9AOjLfxtSTV3DbOQYGugg+Zzw7YF0OuSVmqkiNSlAzpQWe3jDx9uDDugiiT4cVV7+fH1D+jt2Eql0Q8s7YKOiTRIcEdQzypcHdc8xlzbKxzw2JnnPZcPfGexfPYv8XvdDJn1KOfbN3Cb/hlfrdnD3es8LDwy6+pxwvyb4WebwdYyxR7T4+1YTcEFhSNNDtK5XRSZiS23Jl8pxeCOiazaXRr0XCR77exmnU7JkS+NU0rx0KiHyIrNYvba2Sil6g2qokxRxJhjeHzC4wxMDT+r4eCOicy5ZShFFdW8++NecouqKHd5SIyy0DsjjssHdSCmBTpjQoimm78qv9FAqjaXx8+a/MNMe3klr80Y0ewZqiU7ign1zpEuxV6w4YAEU0KcUXZ+De7gpUcAQ9vrTOhs4rGl1fz57PrrC+GuhNxvJJg6xdwwvCOdkqP46yeb2X6wEo/fT0PxSiTBjw9FVVwiMcPGkbrTCge9QcdEGiS0i7E2ekxE3FXw46vgdZGuXDxgfosHzG/RWVVQYcDla+89OusakmHAmtdh5B0t0pyL+qXzfwu2BD0eae2j64eFTj/eHHeM78JP3/zxaCr82sLZa2czaUwbkx20h2HjoY28tuk1tpRsweF1YNNtdI7vzJReUxiaNvRopkOlFNP6TeP6ntfzUe5HvLThJQ44DmDWzBiGgdfvZUjaEKb1m8bIjJFoqmmdppRYKzPHdWnSa4UQrSe/xMFD728IGUg1tFep2utn3d4ynl2cy08mdm1WG8ocnpDZ+yJdZXHY6WlWO04GCaaEaI6q4gaf/uNEK2e9WMVPRzSykb+yqAUbJVrK6C7t+PCesWw7WMGL3+1i+a4S8kscddagHxFJ8KMpSBkygPizu3H+ynyWfbjxaDHeo8dEGCSc1SW5ZT70Eevfpr5NPGHNunocsPSfMGJW0H6wpsiItzMiO4lvtwf/zoVb+yjGZmJ0S3+fgIk9U7HoGlWETtTQWAZEwzDq1Jj6cveXPPnjkxRUFuD2u/EbxzoouWW5LCtYRpw1jpn9ZnJN92uOBlVR5iiu63Ed13a/lgpPBeXV5Zg1M/HWeGymBgZ0hBCntZeX5oUsOB/OXiWnx8fz3+Zyx/guzUpKoWuKmpq8dR+PcJWF6TQs4CvBlBDNoTe8x6Fvqs4l3U088p2bXikNjAabWnhWQbSo7mmxPHJ1fwB+Pn8N7/64r1lFbe1m/ehSs0sHtOf3H2wM+b7hBgmxdjOjWjpIWP4MeKpCPhX2rKurDPathsyhLdKkWeO6sGp3KU53cNDSWMBiM2vcPjYnZB2x5tI1xZ+v6MvP31ob0T47ALvu53bL5yQVmDG6nssTq5/g9S2v4/K5Qh5vYODwOnB4HTy68lFWHFjBw2MfxqwduxYppYizxBFnOX2LYAohwuPy+HhjxZ6gDHqR7FVy+/x8s6WQc3s3fd9tYrQFm1kLqgUY6SqL5Jim18I6WZq3QFKIM11cB2hkycwfJtiY84ObfeX1pMJSOsS1b4XGidZw86jOWE2hR9fihl9F4tnTKVv2JnufnMLe2VOp+OEj7N3qrhX3GxzNeGa36FwzJLPe0biYPhPJuPUfdPzZO2Td/Sqp1/4vtsxeR5+3mTVmtUaQULG/waf/ONHKkyvcFFU1EDwoDcr2tliTzuqazGUD2mM3R3brspg0+rSP59bRnVusLce7uH97fjapO7YI2mY361w8sCP333o9fPhTnv7vdQ0GUsdz+pwszF/IQ0sewmgo1Z4Qos1auLUo5OR/JHuVqqp9zF2W16x2nNsrLWiQEeoONDq2LcPvcWH4vDh3rqL0mxfrHBtdcz883cjMlBBNYBgGO4sqKWl3JX6Wk+A/RFe1D5MK7lh2TdKY3MfMv1a46ZcaoqNlskC/a05Aq0VLGJAZT3q8jV3FoWdtGpshMWmKa4ZkYrccC8jumtiFD9buo8wZvHeqIbpSJEdbuHZYVkSvC4u3usGnw5p1NfyBZBQtRCnFw1f2o7Kyiq8378dJ4zO6NpNG97RYXr5tGOZmbrBuzMxxXWgXY+W3725AKYKWbtZukwHcPjaH+yd1QynFmqtn8/LCe3Adt0im9NtSij8rxl3oRrfpxA2JI+2aNPTowM+Py+fiqz1fsSBvARdmX9iqn08IcerZX+bEHaKuU6R7lfaVNu9anR5vY2ROMou2BW9bCHeVhdvnZ0zX4IRMpzoJpoSIQJnDw1ur85nzbS4VLm9gfXH1PfgNMOPlNn0BN5q+AirqvO6h8VbmratnU2VSDqT3a/3GixahlOKn53TjN/9dH1Ex0yNMumLacUVjM+LtvDp9JNfPWYaj2hcqiX4QXVPE2828OWtU62RSs0QF9j014A8TbAx+tpKfj6onqFEa2Fp2qZmOwVPqbzzX7Xye3puD1+enKkTQEmXRMQy4bmgmD17cq97ZxJZ21eBMzu+Tzvtr9vHMolwKK1xHgzif38Bu1pk2Jpvrh2WRXCtpyAs73+b48LX402KKPi0ic0YmMb1j8JR6KJhXQN5jeWT/NhutJl280+tkzro5EkwJcQZyefz4/MEDuZHuVXJ5I7+fHW/W+BxW5pWEHEhqbKDRrCs6JUVz4T+/Zda4Ltw8qhM2c3jX7Q37ypizOJeluYeoqg70zRKjLFw7NJMbhnds+QRNx5FgSogwvbosjz99vBlNqeM60faj/3radzlP+y7n3p98yDmmd44+nhWv4fpdiE6lOQrOur/1Gi1axeUD2/PdjmI+Xrc/ooDKbtb5y5V9yW4XnAK7X2Y87911Fje9sJzKam/IzHBHRFt0MhLsvDp9BOnxrZRYIL1/IFtlAxqddfV5ILV32G+5eX85Ly3ZxZr8w1S6vNjMOllJUdw6uhPju6cGBi+W/hPlcTBr6gymofHV5oM8uyiXnUWVuDx+LCaNjHgb08Zkc/nA9kRZTvxtLtpq4sYRnbhheEfyS5yUONz4jUAdqc7J0UGbvIudxSzdtxSjVhjtc/oofK+QDtM7ENs/FgBLioWsu7LY9sA2ypaWkTgu8ejx+RX5bCnZQs+knifmQwohTgmxNhNmXQsKqCLdq9QSg3KjcpKZ1CuNzzYdiGj/qK4UaXE23rlrNAfLXTz22VZeXLKLn57TLbAMvp5VBUt3FPP7DzaSX+rE4/Xjq7XcucLl5amvd/Dk1zuY0D2FP1/Zl9TY1rlfSjAlRBie+GIbzy3ODapxc7zqmqKmc3wXs99I5m/m5+pPZGayQaezoO/VLdxa0dqUUvzf1f1RCj5aG15AZTNr/P7S3lw1uP714N3SYlnyq7P5ekshzyzaycaC8jrFat1ePyNzkpk1PodROcmtkkzhqNH3Qv7yQIr0BjQ061oS35v4hM40Nrb4zdZC/vbpFnYdqsLjM+pkpcotrmJVXglWs87M/iZu3/o0+syvQTdhBi7om8EFfTMi/HAnhlKKjslRdExuuK7Vgl0Lgv4vHdsd+D1+4obUHYTRbTqx/WOp3FhZJ5hy+9y8ve1tfjfydy33AYQQp7w+7ePQQtwLIkmKpGuBenLNpZTisesGUPaKh+W5h4KSUYRi1hXJ0VbenDWKeLuZeLuZ524Zyo97Svnbgq08tziXn53XnYv6ZtQpqv7mynx+/0HodPBHHOmzfbXlIKv/Ucr8O0bRJaXlC4tLMCVEI+avyue5xTvDuigc4cTGR/6RdPAWcZ/53aDnXVixtB+CNnkeaJIH5nSka4q/Xd2fsd1SeOrr7eSXOHAfFwhYawKhkTlJ/PTc7mHdrEy6xnl90jmvTzr7DjvJL3HgcHuJsZrp3C6q1UbWgmSPB2tcUDCVd19sna/rm3X1mqKZ7bmUL/++iFnjcrhycIeQS+1mL9zBP7/a3uANscrto8rt4x/LKvku4x/MiWpfaz749JdfkU+1r+4iP1+lD1OMCaUHd5JM8Sacu+vub/DjJ78iv1XbKYQ49QzMSiAl1sruQ8HLssPdq2TWNaYft/y8qcy6xgu3DuPRz7bw8tI8NKVCLvuz6AqlFCNzkvnH5IEkRtfN4jeoYyL/uX0E3+0o5tHPtjJ74U5+eUFPxnVrx2cbDzYaSNXm80NJlZtrn1nGgp+OJTWuZe+jEkwJ0QC318+fPtoUMpBqqBAeBAKq2b7LmWr6nARV0yG1RGMYBl9bL2Rb5i+4z9yWuoRnHqUUlw1oz2UD2rNhXxn/Wb6H3OIqHG4vcTYzA7MSmDKyIxnxTft/7pBgp0PCSfoZ0TQ4+3fwyS8iTyKhmTDFpfPgT+7lnN1lzF64kye+3Mb0MdncOKLT0eUkLy/Zxb++2hH2DdGFhZVFGre/soq504Y3qybKqcQRYm+aHqPjrfRi+IyggMpb5sUUE3z7dnpbLtmHEOL0oJTijvFd+NNHm5q0VwmgW2oM3dJiGzwmErqm+PWFvbjn7G689+M+nl2cS36pA10pfIZBrNXEDcM7csvozg3e45RSjO2Wwpiu7Viw4QB//HAjCVEW1u8rC1kguKF+mQGUOd088PY65k4b3mKfFSSYEqJBn208gL+JhfAANAzmeycw0/41JGTB6HtQfa9miEvnoX99x9geGQzplMj+Mifzlu1m4bZKhoAAACAASURBVNYiyl0ezLpGSoyVG0ZkcWHfjLA3YYqTp2+HeB6+qo0lEhl0ExxYDz+80mgyiqOUDrZ4mPoRSjcxMieZkTnJbNhXxrOLcxn7f18zZUQnxvdI4ZEFW0IGUg3dEKu9flbvLuGlJbuYMTanhT/wyZFoC56xjOoahTIpyleXEz88/ujjPpePinUVpF0TXA8m3hIf9JgQou27fGB7nvhiG06Pj0irJNjMGr+8oHX2WkZbTUwZ2YkpIzvh9xs4PD5sJq3ePVD1UUpxYb8MJvVO4+fz1/LD7tKgY8Lpl/n88H3uIfaXOZs8yBmKBFNCNOCZRTuDMoVFUgjPiZXno6Yx4zev1Vnrm2aBP1/Rl7v/8wPZ7aJZvbsUA+qMtOwqrmJjQRm/e3cDNwzvyH2TurdO1jYhGnLBI2CNhaVPgc8NRgP7w8zREN0ObvskqHZa3w7xPHnDIHYfquK5xblMmfN9UJFJCO+G6PT4eW5xLtPOyq7ze3W66tOuD9GmaKq8x5ZU6lE6qVekUvBqAZpNq5PNz5xkJmF0Qp1zWHUrA1MHnuimCyFOAVEWE2/MHMnl/15CZbU37IDKbtb4xfk9GNstpXUbCGiaanYfRtcUK/NKgjLeRtIvM4B5y3a3aAApmzWEqIfH52fz/vKgxyMphAdQ7vJQUBa8/MasKworqlm68xDVXn/oKeuavSLzvt/NpU9+R2F5eMU8hWgxSgWW+03/LFAPzWQDSzQoE6BAtwayUiZ3g4sehZ8sh/j6k2x0So7mwYt6oWmq3hti0qQ7ieoxGs1iQ+kmorqOqLNZGqCq2suSncUt/3lPgrOzzkYLUfw75aIU0q5O48CbB9h05yZ2/mkn5iQz2b/MRjuuOLCBwVXdrjpRTRZCnGJyUmJ4/ydnkRJjJcrS8GoWs65qkiL1YfqY02eGf/P+Cg47gxMeRdIvc3v9vLWq5YrJg8xMCVGvcmdgud3xGfwiLYRn0jTKnB4ya63kWbqzmJ/854c6yQoaUu31s6ekimufXcZH94wh1mYO+3MI0SIyBsBVz8GFf4Otn0DlQfC6A0v6soZBhyFhn+qLTQfRQ2SfiuSGWOX2MW/Z7hMyotrazLqZa3tcy7xN8/D463YUksYnkTQ+qcHXa2iMzxwfcrmgEOLMkZMSw8IHJvD+mgKeWbiTwspqDL8fj9eLrpsxmxSGAZOHZTF1dGc6JQeX6TiVFVVWh9wrG2m/rCxEQNYcEkwJUQ9dUyGnyiMthAeBgOoIp9vHzFdWR7xXxOeH/WUufvfeBv55/aAmfy4hmsWeAANvbNYpCsqcIQtERnpDzC8Ncx/XaWBKrym8ufXNoGAqHBbdwqz+s1qhVUKI002UJZDc4fphWfyYf5iNW7ZSvno+trH3kBFv4+yeqaftPmy310+oqvaR9su8IYocN4cEU0LUI9ZmrlMA7ohIC+G5vX4So4/NJH2wdh/+EOcNZ6+I2+tnwYYDHHa4SYiyBJ1DiNOBy+3DF+JeFukNMZKikKe61KhUnjn3GW7//HZcvvCX89p0Gw+PfZgeST1asXVCiNONUorBHRMZrJsh7wdoodTnJ1OczQQhtslG2i+zt3AwKXumhKiHrinGdWsX9HtbuxCeY9sy/B4Xhs+Lc+cqSr95Meg82SnRR2sDGYbB7IU7g9KXRrJXRFOB2ldCnK7i7GbMIeon1b4hhnUeW9saDxyYOpAXzn+BGHMMNr3hOigWzYLdZOex8Y8xqdOkE9RCIcRpx+MEU9sow9IzIy7k/vJI+2UDshKCHmuOtnUnEqKFzRrfheW7SoKCn3AL4UVbdO6a0OXo11sPVlBYUbc4J0S2V8Tp8fPKst3MHNel0WOFOBX17RCPWdfw+Or+XtW+ISpNx5Y9CKWZcOWtwbVnXZ2BBbOuGNKp7e0R6p/SnwVXL+Dd7e8yd+NcHF4HPr8Pj9+DSTNh1sxomsYNPW9gco/JpEalnuwmCyFOZR4ntJGalvF2Mxf2TefDtfuDVg5F0i+7Y3zL9p8kmBKiASOyk0iMsuBwB2fjC6cQnlKKC/qmH/16f5mrRTZPHqp0h3WcEKeiEdlJxNvNIQtMhntD1JRi6ujTf9lKKPHWeKb2ncotfW5h+f7l5JblUuWpIsoURWZsJmd1OAuzJklohBBhaEPBFMCMsTl8tvEgTk/TChRHWU2M6dquRdskwZQQDVBK8a8bBnHT89/jjHB/hs2s8cTkgVhNxwIklzt0Qb1I94p4Qm04EeI0oZRi1rgc/m/B1ibfEAdkJdAxOaq1mnhK0JTGqPajGNU+vDIMQggRpI0FU307xDMiO4lluYeCsi03xmbW+F1NaY6WJHumhGjEkE6JPHXjYGzm8H9dbGaN31/Sm0m90+o8HmszEyIjdMR7RU7XTDxCHHH1kEyiLHqovcSNspk1fnGeJFwQQohGeRxtKpgCmH3TEHJSorGawu+X2c06s8Z14fJBHVq8PRJMCRGGc3ql8frtI+mcHIXdrFPfoEaURSc9zsYzNw3hhhGdgp7vkR7bIpsne6bHNvszCXEyxdrMvDFzJNFWU0QBlc2s8b+X9mF4dsO1l4QQQlAzM9W2ZvHtFp137hzN8OykwKBcAzcRi0nDatL41QU9uH9S91ZpjyzzEyJMgzom8s0vJrAm/zDPf5vLl+v34MaCUoE6UqO6JDFrXBdGdUlG1fObnRJr5awu7fhma2FQqYRINk/OauHNk0KcDN3SYnnvJ6OZ/Nz3ON2+kHuojrDoCk1T/N9V/VtlZFEIIdokb9ta5ndElMXEK9OG88Oew8xZvJNvthZh0TX8hnG0D6YpuHlUJ24e2Zn0+IYzpDaHBFNCREApxaCOifz7onZwYAqe+zZhGIGRj3DNHJ/D97sOhew4hrNXxGLSOLunZPASbUPX1FgWPzCR99bs45mFOzlU5cZvGHh8fkyahklXOKp9XDk4k3vP6UaHhLbXKRBCiFbThlKjH0+pQFbXITcP5VBlNWvyD1Pm9GDWNZJjLAzrnIRZb/1FeBJMCdEUhZshtXeTfklHZCeRkxLN1gMVeHwhslE0wG7WuX9S95AZAYU4XUVbTUwZ0Ykbh3fkhz2H2by/nAqXF7tZo32CnSU7DhFnN0kgJYQQkfI4ISbuZLei1SXHWDmnV1rjB7YCCaaEaIrCjZDaq0kvVUrxyrQRXPyvbymqqMbrDy+gspt1rhzUgZtHBu/FEqItODrKeFz9qPYJdu54dTX3n9u9xbMwCSFEm9bGsvmdiiSYEiIMPr+PtUVrKXQW4va5iStYSq/sSTR1DCQp2sKH94xhypzl5Jc6GtwroiuF2aSYelZnfnl+j3r3YwnRVvVpH4fdrLN05yEMDEpqlgIm2C0M7pRIvF1qLgkhRG2eg1V4i5z4C9LRjCT0/VVYMqJPdrPaJAmmhGhAiauEd7a9w7xN86j2VaNQGBgojwPP1l0MLfuR2/rexoj0EREHOe1irHx07xg+33iQ2Yt2sKOwEgjUkNKUwqJreP0GF/fLYPrYbPq0j2+NjyjEKW9PiYM4u5mpL63AZtYxMMAIzGR5fH4u6pfBDPkdEUKc4QyvH+eGYsoX5uM75ApkYPAMgXwTrFiDnmAldnwWUQPaoaTESouRYEqIenya+yn/s/R/AKj2Vdd9UgF+D0sLlrKmcA3Z8dk8O+lZ4q2RdebMusbF/TO4uH8G2w9WsCKvhDKnB4uu0S7Gytm9Uomzyai7ODMZhsHjn29jzre5+PwGXr9BZbU36LgP1hTw6Yb9TOqVxt8nDzwhG46FEOJU4jlYRdGc9RhuH4a7dgkWC/gB/HiLnBz+YAdln+TSbno/LB1iTlJr2xYJpoQIYf7W+Ty68tHgICoEh9fBttJtXPfhdcy/dH7EAdUR3dJi6ZYm9aOEgEAg9eC7G3jvx32NVrn3GQY+j8EXmw9yywvLeWX6CAmohBBnDHdBJUXPrsOorn/LwBGG24/h9lP0zFra3d4Pa8e2n5yitcndRojjLCtYxqMrH8Xlc4X9Go/fQ5GziJlfzMRvNNzxE0I07vnvcnnvx304PY13Do5wefysyT/Mb9/d0IotE0KIU4ev0k3x8+vDCqRqMzx+il/YgPdw+H0dEZrMTAlxnPoCqdJvSyn+rBh3oRvdphM3JI60a9LQowPrjj1+D3lleSwtWMqYDmNOdLOFaDNcHh9PfLG93kCqatNCyle+h+fQXjSLHXNqDvGjr8OW2Qenx8/7a/Zx37ndaC+p1IUQbVzlkgL89SSxWrF3HQ9/M5ttxXlomka35E78/px7GJgRyEZsePxULNpL4uVdT2ST2xwJpoSoZUvJFvIr8oMeL/60mKJPi8ickUlM7xg8pR4K5hWQ91ge2b/NRqsp2uvwOnhxw4sSTAnRDB+v20996VzKV7xL2fK3ST7vJ9iyB6N0E85dq3FuX44tsw8AhgHzlu3mVxf2PHGNFkKIE8zw+alcVgDe4BIrFdVV3Pb2r/nLeT/j0p4Tcfu8rNi7FqtuOXaQ38Cx+iDxF2ajWSQhRVPJMj8hanll4yt4/J46j/mcPgrfK6T9Te2J7R+LMiksKRay7srCXeymbGlZnePXFq2loLLgRDZbiDblmUU7qQox0uqvruLwd6+RNOlOonqMRrPYULqJqK4jSJw47ehxbp+fV5fvxt3IXishhDidOTcdgnpKVeaWBAaGr+h9LrqmYzdbGZ89nF6pXYLPs66oNZvZ5kkwJUQtqw6uwmfU7cQ5tjvwe/zEDam7SVO36cT2j6VyY2Wdx82amfXF61u9rUK0RT6/wY6iypDPVe/bguF1E9V9VKPn8fsN8ksdLd08IYQ4ZVTnltW7VyonKQtNadz/8V/4Zuf3HHZVhDzOcPtx7Tjcms1s8ySYEqKWKk9V0GO+Sh+mGBNKD154ZIo34a2sm6rZ5/dR7i5vtTYK0ZZVuryYtdCL/HzOcrSoOJTW+HIUTVOUOz2NHieEEKcrf1X917hYazT/nfIUoPjlgkcZ+K/LuO2dX1NUVRLReUTjJJgSohaTFryNUI/R8VZ6MXzBc+neMi+mmLqvUUph0SxBxwohGmc2Kfz1LFvR7XH4HeUY/jDS/xpgMcktTgjRdilzw9e4bu0688TFD7LyJ+/w5fSXOVh5iP/96sng88i1slnkuydELcm25KDHorpGoUyK8tV1Z5t8Lh8V6yqI7h1d53FNabSzt2vVdgrRVtnNOlo9M1PWDj1RJjOObcsaPY/H5yclxtrSzRNCiFOGnmALuyffNbkT1/W9gK1FuXWfUKAn2lq+cWcQCaaEqOXa7tdi1+umU9ajdFKvSKXg1QIq1lVgeA3cRW7yn87HnGQmYXRCneMViuHpw09ks4VoM5RSXNIvAz1EQKVZo0kYM4WSL57BsW0Zfo8Lw+fFuXMVpd+8WOfYHumxpMZJB0EI0XZFDUwBLXRXfseh3Ty74g32lxcCUFB+kPc3f8Xg9n3qHKdMGtFD0lq9rW2ZpEYXopbLul7G46sfD3o85aIU9GidA28ewF3oRrNrxA2OI2tWFlqtaXazZmZyj8mYdfOJbLYQbcqMsTl8smE/vhDr/eKGX4UWnUjZsjcp/ugxlMWONa0rcaMmHz0m2qpz5/jgjFVCCNGWmFOiMGdE48kPTi4RbYliTcFm5qycT3l1JXHWGM7tMorfTryrznF6og1Lh5gT1eQ2SYIpIWqJNkdzcfbFfJj7YVCK9KTxSSSNT2rw9ZrSuL7n9a3ZRCHavN7t48hpF8PmA+UYIfZPxfSZSEyfifW+3qxpnNtbRlqFEG1f3IQsSt7cguGuWwoiIzaF2Vf8ocHXKrNG7ITM1mzeGUGW+QlxnF8M+wWpUanoKrICdjbdxgNDHyA9Or2VWibEmeOpGwcRbYl8vM9m1nj25iGYdbm9CSHaPlvvJGy9kqGRZBRBTApr1wSiBqa2TsPOIHK3EeI4sZZYXr7gZdKj0zFr4S3Xs+k2Zg2YxeSekxs/WAjRqJyUGF6bMYJYm4l68lEEsZt1nrxhMCNyghPJCCFEW6SUIum67ti6JTSa3e/oa8wa1s7xJN/YCxXuBVbUS4IpIUJIj05n/qXzuaDzBVh1KzY9eCO7QmE32cmMyeSRsY8wo9+Mk9BSIdquAVkJfHzPWM7q2g6rScMSYrZJ18Bm0uifGc8bM0cySZb3CSHOMErXSL6pN7HndETZTShr6JU1yqKjbDqx4zJpN61v2MGXaJjsmRKiHnGWOB4e+zC/HvFr3t/xPvO3zqfEVYLX78VusjMgZQC39b2NASkDUEpGdoRoDR2To5g3fQT7y5zMW7ab99bso8LpxY9BjNXEOT3TmDamM11TY092U4UQ4qRRmiJuQhaxYzNxbT5ExeK9eIudGB4/yqyhJ9mIHZuJvU+y1JVqYRJMCdGIOEscN/e+mZt733yymyLEGSsj3s4vL+jJLy/oebKbIoQQpyylK+x922HvK/UuTxQJTYUQQgghhBCiCSSYEkIIIYQQQogmkGBKCCGEEEIIIZpAgikhhBBCCCGEaAIJpoQQQgghhBCiCSSYEkIIIYQQQogmkGBKCCGEEEIIIZpAgikhhBBCCCGEaAIJpoQQQgghhBCiCSSYEkIIIYQQQogmUIZh1P+kUkXA7hPXHCHECdDJMIyUk92I5pLrkxBt0ml/fZJrkxBtUr3XpgaDKSGEEEIIIYQQockyPyGEEEIIIYRoAgmmhBBCCCGEEKIJJJgSQgghhBBCiCaQYEoIIYQQQgghmkCCKSGEEEIIIYRoAgmmhBBCCCGEEKIJJJgSQgghhBBCiCaQYEoIIYQQQgghmkCCKSGEEEIIIYRoAgmmhBBCCCGEEKIJJJgSQgghhBBCiCaQYEoIIYQQQgghmkCCKSGEEEIIIYRoAgmmTnNKqQlKqb0n+rVNfL+FSqkZJ+r9hBCnttPp+nXce3+qlLr1ZLy3EKJ5TqfrjvSbTg8STB1HKVVZ649fKeWs9fWUVnzfqUqp71rr/M2llOqslDKUUqbjHn9ZKfXnk9UuIcQxcv1qmArIVUptiuA1/6uUerX2Y4ZhXGgYxtyWb6EQpx+57oQm/aYzh6nxQ84shmHEHPm3UioPmGEYxpfHH6eUMhmG4T2RbRNCiIbI9atR44BUwKSUGmYYxsqT3SAhTndy3RFnOpmZCtORqV2l1K+UUgeAl0KNitSMQnSt+bdVKfWYUmqPUuqgUuoZpZS9Ce99m1Jqs1KqomZUdVaIYx5UShUrpfJqjwS1VBvCbOdUpdR3Ne9XqpTapZS6sJ5jM5RS65RSD9R8vVAp9Sel1JKaz/m5UqpdreMvU0ptVEodrjm2V83jtymlPqx13Hal1Fu1vs5XSg2s+behlLqj5pjDSql/K6VUa3wvhDiVyPXrqFuB94FPav5duw19lFJfKKVKat7rQaXUBcCDwOSaUfa1NccuVErNqGnfYaVU31rnSakZmU+t+foSpdSamuOWKqX6N6P9Qpw25LoTVjul39QGSDAVmXQgCegEzAzj+EeA7sBAoCvQAXioCe9bCFwCxAG3AU8opQYf1652Nee/FXhOKdUj0jYopZ5WSj3dhPbVNgLYWtOevwEvHP+Lp5TKBhYBTxmG8Witp24k8PlSAQvwi5rjuwOvA/cBKQQ6Qh8qpSw15xmrlNKUUu1rXjeq5nU5QAywrtZ7XAIMA/oD1wHnN/PzCnG6OKOvX0qpKOAa4LWaP9fXXENQSsUCXwILgPY17/WVYRgLgIeBNw3DiDEMY0DtcxqGUQ38F7ih1sPXAYsMwyhUSg0CXgRmAcnAs8AHSilrfe0Uoo05o687YZJ+02lOgqnI+IHfG4ZRbRiGs6EDa34RZgL3G4ZRYhhGBYGb8vWRvqlhGB8bhrHTCFgEfA6MPe6w/6lp1yLgY+C6SNtgGMZdhmHcFWn7jrPbMIw5hmH4gLlABpBW6/newDcEvo/PHffalwzD2FbzvZ1P4EIGMBn42DCMLwzD8ACPAXZgtGEYuUBFzbHjgM+AAqVUT2A88K1hGP5a7/GIYRiHDcPYU9OOgQhxZjjTr19XAdU17/8xYAYurnnuEuCAYRiPG4bhMgyjwjCM5WF+xP8c16Ybax6jpv3PGoax3DAMX80+q2pgZJjnFuJ0d6Zfd8Ih/abTnOyZikyRYRiuMI9NAaKA1bUGGBSgR/qmNVO+vycwUqLVnHd9rUNKDcOoqvX1bgKjqy3WBuDIOmdzrX8f+dpT6+sDR/5hGIaj5n1jaj0/BdgBvB3iPQ7U+rej1uvaE/hMR87rV0rlExgtgsAoywQCI0iLgMMELgijar4O5z2EaOvO5OsXBEaf59fs2fAqpd6peexdIAvY2cTzfgNEKaVGAAcJdDTerXmuE3CrUuqeWsdbCHw+Ic4EZ/J1R/pNZwiZmYqMcdzXVQR+6QBQSqXXeq4YcAJ9DMNIqPkTX3ujZjhqloO8Q2BUIc0wjAQC07W1p4ATlVLRtb7uCBS0VBtq7Of/2bvv8KjKtPHj3+dMyUx6gQRJ6EgHaYqAiNgb9i5iw7bWd13ffVf3dXV33XV3fX+2ta2gqIhrL6vi2lZ6ERDpNYSSENL79HN+f8wASWYmmUkmIST357pykZw558xzhplnzv2U+/F/+Ps22t6Peh/YCDwWKNd8pVSklVMB/psS4HDrVS8gP7DpUKUwJfD7QvyVwlSCKwUhuqouW38ppXKA04EZSqlC5Z+/cQVwfmCOwT6gf5jDG79uDR/0tya/h3+o37XA54HWbALnfaJe+VMNw4g3DOOdaK9BiGNUl613kPumLkOCqdb5GRiulBqtlLLhf8MD/lYA4FX843QPTUTOVko1NdZUKaVs9X/wt2LGAcX4W1PPA84OcezjSimrUmoK/iEr77ewDCEFbhg+BJ5QSmUopSxKqWvxdz8viOJUHuBKIAF4UykVyXvwPeACpdQZSikL8CD+oTLLAo8vBKYBdsMw9gOLgXPxz1H4KYqyCdGVdJn6C7gB2A4Mxt9zNBp/i/V+AgEQcJxS6gHln3yeFOhpAn9vU99m6qr5+IfVXM+RIX4Eyn+nUmqC8ktQSl0QmKMlRFfUZeoduW/qOiSYagXDMLYDv8c/cXkH0Hi9g1/j75pdoZSqCuw3mPAm4W8RafxzH/4PRjn+8fifNTquMPBYAf6J1XcahrE12jIof8aal5so3y+AMvwTE4uAe4ALDMM42MQxQQzDcOOfv5AFvNZcxWAYxjZgBvA8/taZ6cD0wHkO/T/U4K8MMAyjCsgFlgYqMyFEI12s/roReNEwjML6P8DLwI2BnqSz8NcthYHXY1rg2ENZrkqVUmtDnTwwv6oW/9CaBfW2rwZuA/4euMadwE1hyihEp9fF6h2Q+6YuQRlGkyMYhBBCCCGEEEKEID1TQgghhBBCCNECEkwJIYQQQgghRAtIMCWEEEIIIYQQLSDBlBBCCCGEEEK0gARTraSUmquU+mPg9ylKqW3t9LyGUmpgjM95+Fra89j2opR6WCk1+2iXQ4j2IvVT649tL1I/ia5E6qbWH9tepG5qXpcIppRSeUoph1KqRil1MPDmjfkKzoZhLDYMo6kUnofKc5NSqnE60JhRSv2glJrVVudvrba+/sBznKaU2l9/m2EYfzIMo8O+LqJrkvqpY5H6SQg/qZs6FqmbOq4uEUwFTA+sYD0WGA/8tvEOSilzu5dKCCGkfhJCdExSNwnRjK4UTAFgGEY+/kUVR8DhLt+7lVI78C8gh1LqQqXUOqVUhVJqmVJq1KHjlVJjlFJrlVLVSql3AVu9xxpE9EqpXkqpj5RSxUqpUqXU35VSQ/EvFDkx0NpTEdg3Tin1lFJqb6AF6GWllL3euR5SSh1QShUopW5p6fUrpd5XShUqpSqVUouUUsMb7dJNKfVN4PoWKqX61Dt2SOCxMqXUNqXUVS0tR6My5SmlfqWUWh8o17vKv4o5Sqk0pdTngdewPPB7Tr1j05VSrwdel3Kl1CdKqQT8/8c9A69xjVKqp1LqMaXUvMBxC5RS9zQqx89Kqcva8lqFaIrUT1I/BY6T+kl0KFI3Sd0UOE7qphC6XDCllOoFnA/8VG/zJcAEYJhSagzwGnAHkAG8AnwW+MBagU+At4B04H3g8jDPYwI+B/YAfYFs4J+GYWwB7gSWG4aRaBhGauCQJ4FBwGhgYGD/RwPnOhf4FXAWcDxwZiteggWBc2QCa/Gv/F3f9cAfgG7AukOPBz5k3wDzA8deA7yolBoW5vorlFKnRFGuq4BzgX7AKOCmwHYNeB3oA/TGv7L53+sd9xYQDwwPlOtpwzBqgfOAgsBrnGgYRkGj53sHuLZeeYcFnuOLaK9ViFiR+knqpwCpn0SHInWT1E0BUjeFYhhGp/8B8oAaoAL/B/RFwB54zABOr7fvS8AfGh2/DZgKnAoUAKreY8uAPwZ+Pw3YH/h9IlAMmEOU5yZgSb2/FVALDKi3bSKwO/D7a8CT9R4bFCj3wDDX+wMwK4LXJTVwnpTA33PxV1qHHk8EfEAv4GpgcaPjXwF+V+/YP0b4/9H4+vOAGfX+/ivwcphjRwPlgd+PA3QgLcR+h/8v6m17DJgX+D0p8Jr3Cfz9BPBa4Pcmr1V+5CeWP1I/hX1dpH6S+kl+juKP1E1hXxepm6RuavDTlca5XmIYxrdhHttX7/c+wI1KqXvrbbMCPfF/ePKNwDskYE+Yc/YC9hiG4Y2gbN3xtxCsUUod2qYAU+D3nsCaCJ6zSYEWnyeAKwPPqQce6gZUBn4//FoYhlGjlCoLPH8fYMKhrvUAM/7WjVgorPd7XeA5UUrFA0/jb3lJCzyeFLiWXkCZYRjl0T6ZYRjVSqkv8Lec/AV/S8ttgYfb+lqFaEzqJ6mfDpP6SXQgUjdJ3XSY1E2hdaVgqin1P+D7gCcMw3ii8U5KqalAtlJK1asUegO7QpxzH9BbKWUOUSkYjf4uwd8FO9zwj0tu7AD+N/8hvcNfSpOuAy7GUwtauAAAIABJREFU39WdB6QA5fgrn0MOP4/yZ+1Jx9+itA9YaBjGWS187pZ6EBgMTDAMo1ApNRr/MAMVKFO6UirVMIyKRsc1fo1DeQf4nVJqEf7x2/8JbD9a1ypEKFI/HSH1k9RPouOQuukIqZu6cN3U5eZMReBV4E6l1ATll6CUukAplQQsB7zAfUopS2DC3UlhzrMK/wf5ycA5bEqpyYHHDgI5gXHEGIahB573aaVUJoBSKlspdU5g//eAm5RSwwKtDb+L4DrMgec89GPB3z3rAkrxt+b8KcRx5yulTgmU7Q/ACsMw9uEfwzxIKXVD4NotSqkTlX9SaFtKwl9ZViil0ql37YZhHMA/jvlF5Z9saVFKnRp4+CCQoZRKaeLcX+JvSfk98G7g/wGO3rUK0Rypn6R+kvpJdERSN0nd1GXrJgmmGjEMYzX+Lsu/42952ElgQp9hGG7gssDfZfjHh34U5jw+YDr+CZF7gf2B/QG+BzYBhUqpksC2Xweea4VSqgr4Fn+rAoZhLACeCRy3M/Bvc17C/0E69PM68Cb+bu58YDOwIsRx8/F/6MqAccCMQBmqgbPxd+0W4O9a/gsQF+rJlT8LzJQIytmcZwA7/haoFcBXjR6/AfAAW4Ei4IFAebfibz3JVf4JnT0bn9gwDBf+/78z8V/3oe1RXasQ7UXqJ6mfpH4SHZHUTVI3deW6STUcwiqEEEIIIYQQIhLSMyWEEEIIIYQQLSDBlBBCCCGEEEK0gARTQgghhBBCCNECEkwJIYQQQgghRAtIMCWEEEIIIYQQLdDkor3dunUz+vbt205FEUK0hzVr1pQYhtH9aJejtaR+EqLz6Qz1k9RNQnQ+TdVNTQZTffv2ZfXq1W1TKiHEUaGU2nO0yxALUj8J0fl0hvpJ6iYhOp+m6iYZ5ieEEEIIIYQQLSDBlBBCCCGEEEK0gARTQgghhBBCCNECEkwJIYQQQgghRAtIMCWEEEIIIYQQLSDBlBBCCCGEEEK0gARTQgghhBBCCNECEkwJIYQQQgghRAs0uWivELFQVO1ke2EN1U4PNquJ41JsDM5KQil1tIsmhAjDpxv8vL+Csho3Xt0gxW5hZE4KiXHytSGE6DjKneVsK99GjbsGq8lKZnwmg9MGyz2GaDfyrSjahGEYLM8t5R+Lclm2q5Q4s4ZhgFLg9RlkJcdx59QBXDw6G7vVdLSLK4QIKK528c6qvby+dDcen45CYeD/7Hp8OhePzubWU/oxKCupReevcXnZW1pHjcuL3WKiR4qN7klxsb0IIUSnZhgGPxf/zNxNc1m8fzFWkxUDA4XCZ/hIjUvl5uE3M33AdBKtiUe7uKKTU4ZhhH1w/PjxxurVq9uxOKIzKK1xMfO1VewuqcXh9hHuHRZvNaEpxaszxzNxQEa7lrErU0qtMQxj/NEuR2tJ/RR7by7L44kvtwDg8uoh9zFpYDFpnD2sB09deQJWc2SjxTcVVDJnyW6+WH8Ai0lDKTAMcPt0xvZO5Y6pA5h6fHc0TVqTu7LOUD9J3dS2qt3V3PPdPWwp24LT68QIc5dhN9sB+Nupf2Nqr6ntWUTRCTVVN8mcKRFTxdUuLnhuCdsLq6lrIpACqHP7qHF5uXnuKr7ferDdyiiECPbMN9v584KtuLx62EAKwKeD06Pz9eZCZs5ZiccXfl+ASoeHq19ZzhUvLePTdQW4vDo1Li/VTi81Li9ur86K3DLueXstk//yPTsOVsf60oQQnUSVu4prPr+GjSUbcXgdYQMpAIfXgcPr4FcLf8UXuV+0YylFVyPBlIgZt1fnuldXUFLjwqM3FUY15PTo3P32T2wuqGrD0gkhwvl47X5eXrQLh8cX8TFOj866/RX89wfrw+5TXuvmwucX89PeChweHV8T9UKt20dhpZNLXljKz/sqoiq/EKLz0w2du765iwO1B3Dr7oiPc/qc/G7Z7/ip6Kc2LJ3oymTOlIiZrzYVkl/hwBvihql28w9U/fgJntL9aFY7lsz+pEy6ClvOcAAcHh9//Worc285qb2LLURYhmHg3LwZz549+Gpr0eLjsfbpi234sE4zudmnG/zhiy04PaF7mJr67Do9Ol9uOMADZx5Pn4yEBse5vTrXz15JYaUTjy+yxhUDf1B1w5yVfHn/FHLS4lt7eUKITmJ5wXJ2VuzEo3uCHitfXE7Jv0twF7kx2Uwkj0sm64osTAn+Odkun4u//vhX3rngnfYutugCJJgSMfPyD7uocwe3bFet+pjKlR+Qcfbd2PqNRZnMOHavwbFj5eFgCmB5bikHq5xkJdvas9hCBNEdDqq+/JLSV2fjOXgQpWkYPh/KZMLQdcyZ3cmYNYuUCy9Es9uPdnFb5YdtRbjC9EhF8tnVdYPXl+bx2EXDGxy7YOMB8kprQwZSzTWu1Li8PP3Ndv7vqtExvlohxLHqtY2vUeetC9pesqCE4gXF5MzKIXFYIp5yDwVvFZD3VB79HumHFpjXuaN8B7mVufRP6d/eRRednAzzEzGxtbCK3JKaoO26q5aKJW+TftZdxA+ehGa1oUxm4gdOIG3aLUH7v7k8r+0LK0QTnJs3s3Pa6RQ+8SfceXkYDgd6bS2G0+n/1+HAs2cvB//8JDtOm4Zj46ajXeRWeXnhLmpDNIJE+tn16Abvrd6Hs1FA1lTjStl3r5Jy8lXk3DOP7LteJ2ns+Th2rDzy3AZ8vv4AVc7gFmghRNdTUFPAuuJ1Qdt9Dh9FnxTRc0ZPkkYlocwKa3crvX7RC3eJm8pllUf21X3M2zyvPYstuggJpkRM/JhXTqjEkK78rRheN/GDJjZ7DpdXZ+H24jYonRCRcaxfT96MG/BVVGDUBbeA1mfU1aFXVrLnhhtwrAv+kj9W/LyvMuT2aD67mlLsOHikMWXLgSp2l9YG7RdN44qmFB+t2R/l1QghOqN1Reswq+DBVHU76tA9OsnjkhtsN9lMJI1KombTkXrJa3hZXrC8zcsquh4JpkRMVDk8IbN6+RxVaPHJKC2ytaSqHN5YF02IiHgOHmTvrbOaDaIaMxwO9s66Dc+BA21Usrbj9elhs/FF89lVPh9Fm7fjys3Fc/Agizflh0w2EU2A5vD4+GpjYfMXIYTo9Krd1fiM4J5uX40Pc6IZZQqew2pOMeOtaXhPUesJbuQRorVkzpSICYtJoSmF3qh7ymRPRq+rwtB9Ed2UWUJUiEK0h7K5c9GdzqDtZ+7aSanP16DlaUH//mSaLYf/1p1OSl97nR6PPNwOJY0dTSlQECq7cDSfXd3tovqNuewv241eW0tujwl4+k0B1bC9LtrGlQqHDPMTQoBZM6MIvj8wJZrw1ngxfEZQQOWt9GJObHiba9bktlfEnvRMiZjolhgXcvHOuOwhKLOFuu2Rda13T5LkE6L96W43Fe+9D57QN+8vZOewZtDgwz/1AykAvF4qPvwwZDDWkWmaIjEu9M1FNJ9dPc7OyGf+woAFX3L8ooVkzro1KJCChgFaROXrJBkThRCt4/Mm4PUF1wfxA+NRZkXVmoZLq/icPqrXV5MwrGGW0TRbWpuWU3RNEkyJmDhjSFbIYT1aXAKpp1xP2TcvU7d9ObrHieHz4ti1mvL/vNZg3wSriWtP6tVeRRbisOp//xsj1KS/KFV99VUMStO+LjqhJ2Yt+CYlms9uZlIc/bsduWnJSIzDZml940q3JGuUVyOE6GwWbS/msXc9eEI0wpjiTWRekknBvAKq11djeA3cxW72vbgPS7qF1EmpR/bFRE5iDiWOkvYsvugCpL9TxERKvIVzh/fg8/UH8DW6KU0+6TK0hDQql79LyedPoax24rIGkjzx6gb7KaU4d0SP9iy2EADUrVwV9Vypxoy6OupWriL1kktiVKrWcXl9eH0G8VZTk2ti3XJKPz5csz/k+nCRfHbjrSbumNq/wXOcNSyLP325Jeh89QM0pZmw9RuD0sw489bh3Lu+QRKKhDgTV4yTxhUhurJlO0u4/a3VOD0acRXjsaSuQGkN53l2P787pgQThe8W4i5yo9k1kscm0+uOXmj1GnV8+FhSsIRzPjiHydmTuXXkrZzQ/YT2viTRCUkwJWLmtlP78+/Nhfg8wTdlicOnkTh8Wthj48wa10/oTZw5srkUQsSSt6K8ycfvzd+PORAsnBgfz9+zc0Lu5ytv+jxtbcuBKmYvzuWL9Qdw+/TD8xiH90zhzqkDOHt4FhZTwx6jAd0TGZRqYUOxAyPE0LzmPruGAZeMyW6wrWeqnfF901i6szRo/4gbV1CcO1waV4ToqkpqXNz25urDC4q7yyZjSfUvoVCxvMK/SO8BN5pNw9bbRs8be5IwKKGpU+L2uQH4Yd8PLC9Yzt2j7+amETe16XWIzk+CKREzI7JTuHFiX95cvgdHmEVAQzFpkJNm574zjm/D0gkRnhbX9Fy957NzmJTQ9Jc0gLIdnTl/u4pruHf+T+SW1ODxGYeH3B5KCLMhv5L//uBnfvOR4jfnDeHaCX0AMAyD8rfn819fvsO9E++iJvKPLQA2i8Yz14wm3hr8VXLn1AH8tLci5FpTzQVo1kDjSqh5mEKIrmH+yr0Ne8wNM2Ci+KuDlHxRTM8be5I0MgllUlRvqKZ6bXWzwdThU2Hg9Dl5Yd0LGIbBzSNvbpuLEF2CBFMipv7nvCEcrHLx+fqCkMOGGrOaFD1T7bxz+8kkhJkIL0Rbs+Rkg9kM3lak5jeZsGRnN79fjK3dW87MOauodXtDrvV2yKGFeX//+RbySuv478k9OfDb3+IpOMDk11/gHVMq181eQa3LSwQfXWwWjccvGs45YXqPThnYjWmDM/lu68HDLcuRMGmKHsk27j59YMTHCCE6F59u8PrS3bi8gbpDebH3+Qe+OhfFHxeRfWs2KeNTDu+fPCaZ5DHJIc9Vvrjc34tV5MZkM5E8LpmsK7IwJZhw+py8+POLjOg+ghN7nNgelyY6IWn2EzH1/dYi/r0p8rVhJg7oxuf3TSFTsviJoyjl4otRrRxiqiwWUi+5OEYlikxucQ0z56yixtV0IFWfw+PjzaW7+ctdT2Dp2ZM+78zH2qcPI3NS+OLeKZzUL504sxZymQJN+edI9cmI59WZ47n6xN5hn0cpxdNXj+bEvunYLZG9thaTIjMpjnfvOJlkm6X5A4QQndKiHcW4662BZ07aiGaqxpFbE7RIb/nicnb8dgebbt/E1vu2UvBGAb5af+NRyYISCt8vpMdVPRj24jD6/29/3KVu8p7KQw8Eak6fk5d+fql9L1B0KtIVIGLmq40HeODddVG1Qq/aXcbOohpG90ptfmch2khc//7EHX88zg0bW3wOa+Ac7enXH66n1h26N6128w9U/fgJntL9aFY7lsz+pEy6ClvOcBw+gzf6TmXW3WegWY9kzOudEc8/b5/IvrI63liWxyfr8ql2etENg3irmYn9M7h9an/G9EptMqnFIVazxtybT+LJBVt4a8UeNKVCDvuzmhS43ZzcK5VZZw3j+e92sr+iDpdHJy3ByuQBGVw6NidsGnchROdQXO2itNbFsp0luL1H7iWsGT+gTO6gRXpLFpRQvKCYnFk5JA5LxFPuoeCtAvKeyqPPL/tQ9Im/FytpVJL/PN2t9PpFL7Y/tJ3KZZWknepPlf7TwXXk1+STndj+owvEsU++mURMbDlQxX+9+3NUgRT4W8lnzlnJdw+eRvekuDYqnRChGV4vNQsXUvX1NxgeL2ga6A3fw98OaH64mbLb6Xb7bQ227Syq5tN1Bewvd+Dy+shIiOPk/hkhk0C0xJ7SWtbvrwzZI1W16mMqV35Axtl3Y+s3FmUy49i9BseOldhyhvt30jTmrdzDL88aHHR8r/R4fnvhMH574bBWl9OkKR65YBj3nzmIT37azysLc9lf4cCkFD7DINlm4bqTemHdsJYP97i54y0nLq+vwVDDRduLeeLLLVw8Opu7TxtI74z4VpdLCNExeHw6X286yMsLd7GtsBqrWcPt0/H4/JWAFncAzepPZlN/kV7drTcZLBX/qzioFwvAZDORNCqJmk01h4Mpj0/nvs9f4v2r/oAWYqkIIZoiwZSIiWe/3Y7TG3r2elMt5ABOr86by/N48Ozgmzoh2oKvupqyN96kfN48dLe7VWnRXSYLjlEnMvicczAMg682FvLSwl1sL6zGqxsN5g5+9NN+fvORYsbJfbh5cr9WNSDMXZZ3OMFEfbqrloolb5Nx/gPED550eHv8wAnED5xwpNxenTeW7eG+04/HHIPgrjmJcWZmnNyXGSf3RdcNat1e7BYTXt3gznlrWFmdjkPTIUTymkO9WR+s3sfnPxcw56YTObl/RpuXWQjRtpbsKOHu+Wvx+vTD8zrrD+8DMNn3Hf798CK9a6vQ4rQmg6W6XXUNerHqM6eYcexxHP5baT62lm3l4Y838OfLRkbU8y7EIRJMiVYrrXHx/bbiFreQu706by7fw31nHB+TFnshmuLJz2fPDTPxlpZiuFytOpey27GPO5H7+l3OdQtz2ZBfyX+2FYUcygZQ6/Jvn71kN/NX7WX+rJMZ1jP0pOnm/OvngsMtt/W58rdieN3ED5rY7Dl8usHP+ysY1ye9RWVoKU1TJNkseH06t879kdV7yo9MNG+Cz/An0rj59R+Zf9sExvROa4fSCiHawhfrC3jw/QhGtGhOwF93muJNZF2aRcFbBaROSMWUaAIDqtdXU7ullh5X+xPimFPMOPc5D/diHQqoDqVUd+51okyKvP/Lo/v07iQMSkBXDj5dV8D4PmlcMV7WuBORk2BKtNo/f9xHqDacSFvIAbw+ne+2HOTcEce1cWlFV+YtKWH3VVf714PSoxuSWp+WkICyWMiYdSvpt9zC/EonZz+9CLdXjyiLpdur4/bqXPnKMj69ezIDM5OiLkO1M/RcKZ+jCi0+GaU1n/RBKSitcUf93LHy/Pc7Wbs3dCDV5Jwvj48bX1vFiofPCJmWXQjRsf2YVxZZIAVgmPDnS/Pv2+28bphTzBz85CC+Kh9b/2sr9n52MqdnHj7EW+nFmmnFVeiiak0VKSelUPJVCcVfFNPj2h4UvlNI1qVZmNPMR1Kq6zYcHh9Pf7uDy8flSO+UiJh8C4lW+zGvLOTNUDQt5LVuH+v3V0owJdrUvl/cja+yMiiQOnPXTkp9vgbpTRf070+muV5GOZMJU1oa1r59yLjpJhJPOw1l9lehn/yUj24YIQOppoKCOpePa19dydJfnx71mkqhhvgBmOzJ6HVVGLqv2YDKMMKfp625vTqvLd2NI8TNVCQ92l7d4NN1BVx7UviMgkKIjunhjzaEDaQa15nWnt3JulyROORIfZw6KZWk0UlsfWArPW/oScpJR9Kk+5w+qtdXk3VFFvb+dgrmFYCCgx8dpMdVPahcUYkl3ULqlFQ0i0bymGQM3YTPmQVAeZ2bVbvLmCBDiUWEJJgSrVbl8ITcHk0LOUBZ7dFrIRedn3PLFlzbt4ddS+qFZhbmVWYzfee/jbV3w5t3j0/nlUW5IW8MmgsKDKDO5eWrTYVcdELPsM9tGAb5FQ62Hqhma2EVWw5UH16Yt7G47CEos4W67ctJGHJK2HMCKCDZfnRSkP97UyF6iGuItEe7zu3j5R92cc2JvaQFWYhjyIb9lewvd4R8LHSduZLqdXPwllv960UdcKPZNGy9baROTqVgXgGaTWuQzc+SbiF1kj9YMiWYKPxnIYbboOizIpLHJtPrjl5olvoNWApPxckAONw+Zi/ZLcGUiJgEU6LVbGHWkImmhRyQRXtFmyqd+waGJ3TgHwlD1yl7ax49Hnm4wfavNx1sVVBQGwgKDgVTdW4v2wqr2VpYzdYD/sBpS2EVdouJocclM+S4JM4alkWNy8viHcVBC+xqcQmknnI9Zd+8jNJM2PqNQWlmnHnrcO5dT9q0Ww7v69WNiJYlMAyD5btK2XygimqnF7vVRHaqnTOHZmG3tmx9rjlLdh+ecF5fND3axTUuNuZXMTInpdl9hRAdw+zFubhCJKwKX2dOpm7L9xyYv5qeN/YkaWQSyqSo3lBN3bY6si7PovDdQtxFbjS7FhQspU9NR7NoFP6zkCHPDQlZJp8zB8PjD54MYMfB6thfuOi05O5VtFqfjARW5JYG3dRF00IeZ9bISbM3uY9P97E4fzFvbX6LvKo8nF4nNrONfin9mDlsJpN7TsYUYS+Y6Fp0p5Pqr74CX+jEEBHxeKj48EOy/ufXKNOR99mri3NbHRRsP1jNDbNXsr/CwYFKBwMzExnSI5khPZI4Z3gPBvdIIiOxYea//t0TWLW7DEeI7HfJJ12GlpBG5fJ3Kfn8KZTVTlzWQJInXn14HxMGlwzv3uSco0qHh/dX7+PVxbnUOL2H0xWbNUWcWUM34MrxOdwyuR99u4Xv1Qtlb1noDIrR9Ghryn8eCaaEOHb8sD24EQjC15m6q5aSBevIntWLlPGJh7cnj0kmeYw/gU/61KaT6NRPqd44u5+hW3CXnNFgW7gkQkKEIsGUaLXrTurNJz/lB93URdNCDjA9zDAnwzB4c/ObzNkwB5fPRZ33yE1YlbuKoroiNhRvwGa2MWvkLGYMnSHDfkQD3pJS/xpSTbg3fz/mwPvmxPh4/p6dE+JEXtwVlZjT/FnkFJBXUhvyfNEmghhyXDKPTh9Gv24JEaUqH5WTSs9UG7uKQz9/4vBpJA6fFvZ4Mzpn/OMxSmsuJm3GDDSbrcHjmwuquG72ClwePeiz7dUNvIGbjXdW7uW91ft4/KIRXH1i5BmwXCGCQIiuR1vXodYVetimEKJjCveZDVdn+oMsD3E5V2Lo36K06EcY1E+pnnLikcYXQ7fgLj0VX23DBdcrHR7e+3Ef00/o2eLed9F1SDAlWm1kTkrYm7pIWsiVgqmDu9MtMXjNHY/u4aGFD7EsfxkOX+gx1gB13jrqvHU8t/Y5fi7+mSenPIlZk7e38NPralGaRlOpFp5vZs4UgMMHp/7+C4rj0zEg5HIAh0QTFFhNGsN6JnF8VnRZ/f54yUhunrsq6sWy7RYTF4zKYcrtz1L89DOUzTuP7vfeQ8rFF6PMZjYVVHLly8sjap316AYe3eB3n23E5fExc1LfyMpgNYXs0YumR1vTFIk2+ZwLcSwJ19YZrs48FGR5q6bitnqwZiyMOqCqn1JdaYqE4YkozUrFit7UbNxHWqN2J5dX57F/beKxf23i8rE5/PKsQaQlWKO9VNFFyLeQiIl7Th/Iwx9tDDnkqLkWcpvZxB2nDgjabhgGDy9+mKX5S3H6nBGVw+lzsnDfQh5b9hh/mPwH6aESAJgSEjBakQr9ELtmsOz3F2FKOdKyecLjX1MZIglLVIkglCKhBSm+Jw7I4M+XjuQ3H4fPjNWY3aIxrk8aT142ErNJI+f553CsW8fBp56i9PXXibvnAa5fpUc9zMXp0fnTgi0cn5XExAHNT9we0D2RkpqyoO3R9Gj7dIOBmYlB5xBCdEyGYWC3mPD4gnunwtWZ9YMsd8mZ6K4s4jIXoMw1oDwo1bBVy9AtoALnNzSU5q/Lup3XDVOylaLPinG9ko9mjceaaW7QuFvfoTrwnz/u5ZstB3nv9on0zoiPxcsgOhkJpkRMXDI6m2+3FPHdloNRtZLbLSbuOLU/4/oEL7759Z6vWbh/YVAgVb643J/Rp8iNyWYieVwyWVdkYUrwt2Q5fU6+3vM1Z/Q+g2m9wwdxouswde8ek/Momw0tqWHvUe/0eDbkVwbtG21QEO2co0MuHZtDaoKVe+avhcCitqFYTQqU4uLR2fzxkhENhhLaR4+mz1tvUbNwIS/M/RZH5okQome3qTTv4A+onvp6Gx/eNSno2MZum9KfjfmVIcsbSY82QJ+MeAZF2ZsnhGh/1U4PH67Zzz8W54ZtqAlXZ+puByh1OMjyVo/EWz0Czb4Xa/oiTPb9KM2FYZgxvMlkGmdQWDAYl8+FOWU1JlsByuTA0OOIH5SJJfNEDG/ziXcO8fgMiqqcXP7SMr68fwrdk4JH0YiuTYIpERNKKZ6+ajT3//MnfthWHLKHqjG7xcTMiX24/8zjQz4+Z8McHN6GQ/tKFpRQvKCYnFk5DdKg5j2VR79H+qEF1upxeB3M2ThHgikBgGa1knLRdCo+/ChsavRmWa2kXXsNqtHcq1lT+vHwRxtaFRT0Sre3KiiYNjiTNb89iy83HOClhbvYW1p3eN0qXTfQNMX1E/pww8Q+ZKeGTvSilCLh1Kl8tNSLq9oV9Hgkaz8BbMyvJK+kttngcNqQTKxmLWzw11yPdoLVxF2nBfdoCyE6li/WF/Dg+z+jKdVsj3e4OjNp9LlBQVbtxmJK93pIm/abw8fHmTUuOnUAcwpyMXxmPGVTiXRAYFONRbrhX3/qvnd+4p3bT27FqyE6IwmmRMxYzRovXj+Wt5bv4cUfdlHt9ATdKGkK4swmeqba+NXZgzlvZOhFeneW7yS3MrfBNp/DR9EnRWTfmk3SKP+Np7W7lV6/6MX2h7ZTuayStFOP9HBtLdvKnqo99EnuE+MrFcei9BtvpPKTTzFCBFPfDhjY7PFKKdKuuy5o+7kjevDIxxvDHtdeQYHNYuKysTlcNjaH/eV1lNa4cft0UuwW+mYkRLQo8PLc0pCTwyNN8w7+4O2N5Xn8bvpwmmLSFPecPpCn/r09osaX+hQQH2fm3BE9ojpOCNG+3l6xhz98sTmqESvh6kxrzyHNNkwZwIyJvZk4ICOq+aSRLhS+dm85e0vrZLifaECCKRFTSilmTurLDRP7sGxXKa8v3c3uklrKav03dmcMyWLWlH6Mymm6i/3TXZ/i1Rve1NXtqEP36CSPS26w3WQzkTQqiZpNNQ2CKZ/u47Odn3Hv2Htjd4HimBXXvz/20aOp++kncEe3QLSKiyNhyhQsPYJv3uPMJm6e3JfZi3NxRJkIAvyNEOeNCN2o0FI5afHkpEX/Zb+5oAq3L/gaoknz7tEN1u6tiOguQ6YoAAAgAElEQVT5bpncj3X7Kvh288GoXrv4OBPv3DaBOLNk2RKio1q4vTjqQKopzTVMKWDKwG5kJtnITLLx6szx3PHWGlweH74mkgVF1VhkGMxdtptHpw+H6kKozAdPLcQlQVo/sEc+fFB0HhJMiTahlGLywG5MHtgNgBW5pfy/r7fz3LVjIjo+vyYfn9GwtdpX48OcaA5aIwLAnGLGsafhkECv4SW/Nr+FVyA6o5znn2P3pZfiOVgU+XA/iwVLdjbZf/1L2F0eOHMQP+2t4Me8MlzeKIICq4m3Z50cduHr9lbl9OAJcdcRTZp3gOoQCTlCUUrxzNVj+M1H6/l8/YFmhwDFmTXirSbeuf1kBmbKXCkhOirDMHj0041hA6nm5l+2hM1i4hfTjvTyTzm+O1/cN4WXftjJZz8XAIQsT1SNRT6D91bl8WjJQ7B/NZgPzZ8ywOeGwRfApHsge1yLr0MceySYEm2quNrF/JV7eG/1Pg5UOjnh8a9Jspk5c2gWN0/uS5+M0PMqnN7g7H1NLbrnrfRiTgx+O4c6j+i6TMnJ9H3vPfbceBOe/fsxnE2/P5TdTtyAAfSeMxstPnxPj0lTzL5xPL94ey3Ld5U2O2zNYlLYLSbevHUCw3omN7lve7JbTZgUQa240aR5P3SeSJk0xV8uH8X5I4/j5YW7WJtbggF4ODIsMcFqwmLWuHlSP26Y2Id0SVEsRIe2dm8FRSHmXkLTQ+p8VcUtCrLsFhP3njGQcX0aLt7br1sCf73iBB6dPpwX/7ODVxbtxtdoxeBoG4vqPDquvJXEKS/4Gl3j5k9g+1eQNQyuex/im15MWHQOEkyJNpFf4eDRTzeyeEcJCg631lc6PFQ6PLy9Yg/vrNrLiOwUHr9oOCOyUxocnxoX3FV+eNG9NVWknHRkf5/TR/X6arKuyAo6Js0WnCVQdG3mjAz6ffA+FR9/TNns2XjLKzDqjiwEjVIomw1zZibdbptF8kUXoVmbv3m3WUzMnjmeT9bl89LCXewvc+Dy+qj/vZ1gNWEAV5/Yi9tP7c9xKaGTQRwt2al2bJbg9Z+iSfMO/gyH0VBKcdrgTE4bnMmyq2ay6qJbKU5Ix+nxkZEYx/g+aZw+JDOixYyFEEffq4t24QzRqNTUkDpvWQFl373abJKbxuwW/7zTu6aGn3uaGGfmjKFZvLViL9XOhqMSom0sMuHDSRxxhBjdYOjgqYMD6+HlU+D2HyAxs9lzimObBFMi5jYVVHLdqyupcXrxhVnV1KMboBus2VPOlS8v58UZY5k2+EiFc2KPE/lu73fUeY/c5JriTWRekknBvAI0m9Ygm58l3ULqpIYBWLw5nnFZ0tUugmk2G+nXXkvaNdfgWLuW6u+/x1tcAkphyexO0llnYRs5Mup1yjRNHU4CsTG/kg/X7qegwsHSnSWc2DedC0b15MJRx3WYYX2NnTUsi//5cEPQ9mjSvMdbTdwwsWVJX7wlJWTkbeP+ayajzPL1JMSxakVuWchFzcMNqYtm3hKAWVOYNMWI7BTuP+N4Th3U/PIXSTYLeohCRdtY5MVEInVN7+RzQ00RzL0Q7lgEFluz5xXHLvm2EjG1t7SOa/+xgipn5OmnHR4fd81bw9uzJhzuoj83ZTB/DjFEr/v53TElmCh8txB3kRvNrpE8Npled/RCszRstVZKcXafs1t3QaJTU0oRP24c8eNiH3SPyE5hkLeCyn/9h282rGPYgTh67D8Od9E44s44AxVBb1d7i7eauWxsNu/+uA9vo6EwkaZ5T7aamNi/+UV7AXYW1bCzqJpqp5d4q5n0DavImTBBAikhjnHh5j+GG1IXzbwlgIGZCbw8Y3xU6/P1Sosn1Nrt0TQWAWRTwoBnq6jzwO77E0mw+hvdZq91M2+9hx9uCpRJ90DlPtj4AYyZEXE5xbFHvrFETN09fy01IVIrQ9MTTp0enVlvrGbVTRlYVjyLPW8JF/c/gQ/q8vAaDc+XPjWd9KlNj0O2aBYuP/5yrKaOd8MqOjfDMKj5zw+Uvvoqzs2bMXw+xgSSXZQDFR9/jPrfR0m77jrSb5iBOUYLCsfKLaf048O1+4OCKWg+m5YNnct/+oyKDypIvfzyoDW5ANxena83F/LyD7vYWVyDWdPQDQNNKbxORbe0s7hr5R4uGZ1NQpx8RQlxLNI0IEQ8FW5IXbTzllLs1qgXOrdbTVwyJpv3Vu+lcdLSSBuL7Di5w/w5/4t/bumzK908PKWJRXw9dbDkGQmmOjkZgC5iZlthNTuKqglxD0bVqo8p++5VUk6+ipx75pF91+skjT0fx46Vh/dxO2r5bt5fodfJcP96bj73BeLMLVtpPM4Ux8xhM1t6KUK0iOH1UvA/vyH/wQdx/PQThssVlDXQqKtDr6mhbO5cdl04HcemTUeptKEN6J7IoxcOx26J7uvBZtaYPLQHd//xF1S8/wF7bpiJa+fOBvvkldRy6t/+w68/WM/GgiqcHp0al5c6t48alxenMrHfrfHEF1s4+c/fsTqvLJaXJoRoJ8k2S8jt9YfU1Vc/yIpErbtli6/fekpfLCEaecDfWHTcjc/Q+5cf0uueeWRe+Ri2nKEN9jFQXGpaAsBDk6w8tcxFhbOJvOsAVfmQv6ZF5RXHBgmmRMzMWZKLJ0Ra6ENjodPPuov4wZPQrDaUyUz8wAkNus9rDSsvJd8DE38BcYn0TOzJi2e8iM0U3Vhju9nOK2e9QlZCcEIKIdqKYRjkP/QQ1f/+N4bD0fz+bjd6ZSV7bpiJc9v2dihh5K6b0JuHzhmCLcKAym4xMXlgN164bizxQ4fS9535JJ9/HntumEnRs8+iu1zsLKph+t+XUFTlDEpw0Vid20e108uMOStZurMkFpckhGhHl43NxhpiGZP6Q+rqti9H9zgxfF50twOUCgqywkkKE6w1Z2BmEhP6ZxAXwSLmjdlxcYPpGxKVfwrC+J4mTutr5qllobMWHuZ1Q+7ClhRXHCMkmBIx89m6gpAL40UzFnprYQ1FVUfmSo3NGstr57xGkjUJu7npzGd2s50Uawpzz53LqO6joi6/EK1R9sab1PywMCjd+pm7djJm+zbG1fsp8h5Zh8moq2PvTTehN5Omvb3dcko/5tx4ImN6pRJn1jBrwTdG8VYTPZJt/Prcwbw6c/zhxBrKZCL9+uvp98nHuHflsv7SK7j6hUXUuLwhe67DcXp0bntzNbnFNbG6LCFEO7hhYl8Ik8An+aTLSDv9ViqXv8v+569n/0s3UbvhW5JGnxsUZDl2rab8P681ON5iUoxoxZISL14/lpw0O9YosoPacDFB28z/mN9psP330+J4fpWb4tom1hc0vFArjUKdmQxIFzHh9PhCLvYJ0Y2Ftpg0iqpdZCYf6Y0a2X0k31zxDV/kfsFrG1+j3FkOgG7oaErDwKCbvRs3D7+ZC/pfQLwlurTMQrSW4fNR+sorYXukXsjOYVJC+PH9hstF1YKvSL30krYqYoscWng7t7iGucvyWLe3gmqXF5tFo3d6PDdO7MvEARlhsx5asrLIee5ZPpzzb2q2OjBMwa3JzS3e6fLoPP/9Tp6+enSbXqsQInaSbWa6J8aRXxG6Tgw3/9Lac0iz85Y0pbjh5L4tLltCnJmP757MTa+tYmthdZOLhSt07Lg5S1vNU5ZXMKmG9zkjMk1cOMjMk0vcDO3eRHAWou7r0rwu8HnAmhA26D6WSDAlYsLl0dE08EUx4TQUpcDdeGYokGBJ4KrBV3HloCvZULKBPVV7qPPUkWBNoG9yX4ZnDI86jbUQsVKzaBG6293i4/W6OkpffbXDBVOH9O+eyO8vHtGiY3XdYN4BE64QNxNNLd55KJjyGQZfbjjAYxcNJ8UuNyRCdHSFlU6u/sfyBqNMItVckhuAE3ql0jujdY2myTYL7985iT99uYXXl+4O2WNeu/kHPD++y8HSfN6LM9jeQ+ORKXGc0rvhrfPjp9kY+0oND04MM8fbZJW1pgAK1sHyv8OWf/lTxyvNvy5Xz7FwygMw6DwwHZthybFZatHhJNrMIbN/QXRrOOi6EXbiKvhTWY/qPkqG8YkOpez1uRi1ta06h+fAAZxbtmAbOrT5nY8hi3eWUBdisng068poSvHB6n3cOqV/m5dXCNFylXUeLn9pKYVVLnzRjOmNkM2i8cAZx8fkXCZNsWRHSdikWZUrP+CMc8/hrUH/Id3s5KudXj7d6g0Kpgama1w93MJzq9yMzAzRO6UUDLkgJmU+Jh3cBB/cAhV7/PPHjECruxFoOM9fDR/f6e+9O/dJOOGao1fWFpJgSsSESVP075bAruLgG8qo1nBQil7pTc+NEqKjceXmNvn4vfn7MQd6Tk+Mj+fv2TnBO2ka7ry8ThdMrckrC5lwIpq5lA6PjyU7SySYEqKD+9UHP1NcHTqQam5Ib3PsFhMPnHk8kwZ2i0lZtxyoYm9Z8OK79Rt6dh0/Eat1MRalmD7YwvTBoRt7H50ax1vrPSEfI3s8pPWNSZmPOXlL4e0rwdNMY6M7MC/28wegbDdM+03bly2GJJgSMXPn1AH87rNNIccfR7KGg8WkuPakXsSZI1tnQoiOonHSicaeb2bOFAC6jq+6Ooal6hhKakMPf4x2XZmKujA3KkKIDuFglZOF24txh5g/HcmQ3nA0BVazxq/OGcStp8SuQWVTQVXI6Tr1G3p8mHjDew53mv+FTR2pg/IeSGpwTK8UDedvQyTFsMTD5AdiVuZjStEWmB9BIFWfxwHLnoOE7nDSrLYrW4xJMCViZvoJPfndZ+HXzGluLLSmFDdO7NsGJROibSmbDWpal3FOaRqmxMQYlajjsFtCB0vRzKUE/82UEKLjmrdiD6FmLkczpLc+m0XDMOCsYVncceoARuakxLS81U5PyB60xg09r/gu5ELTCvpwEIuKbB0sAMx2GHgmHH9WrIp8bPn4TnA37Pnr+0w1dR7YfX8iCVb/u2X2Wjfz1nv44aZAg6OnDr5+BEZcBvHp7V3qFpFvJxEzNouJe08fGPbmqeljNc4bcRy90iUTnzj2WHv3avU5DF3H0qt3DErTsRyXYgu53ky4xTtDUUDPVBn+K0RH9vbKvbhCrDUZzZBeTcH5I3tw06S+/Oa8oaz4zRn8/bqxMQ+kwN/Qo4Xommq8gLADG9e6f8sBIwOXEWEfhCUeek+Ay+d0imx1USvaCsXbgBDBqgHPrmwmYZPPA4v+1jZlawMSTImYunPqAC4YdVxUAZXNrDG8Zwp/vUKSSohjU/pNN6GaG8bXDHO3bthGRDZ34FhywajjQmbaDLd4Z6h1ZexWE9ec2PqAVQjRNnTdoLyu9UN6E+LMzJrSn8cuGs6Nk/qSlmCNdVEPOy7VTqilpkI19BSTygXuP/G9PganYcFhhEmUZUnw90iNvwVmfATmtit/R+Xy+nAseg7DF3po9kOTrDy1zEWFs4kEJYYPVrwIb1wMro4//F2G+YmYUkrxtytGkZ5g5c1lebh9ethFOjUFcWYTpw7qxnPXjpFhPOKYlXT66SiTKUQbHHw7YGCzxyu7nYxZszplev/jUuxM6JfOoh3Bi1ZGMpcSIC3eykn9jo3hHkJ0RW6fjiJUP0T0Q3qdTaz7FEuTB2Rg1jSg4fOFS5pVlLeZa/amMGjas1xn+o4bzN+QTjWG8t/PqNQ+MOl+GHUVxHW+IdtN0XWDpbtKeGVhLstzS1lr+Ri7Cs7iCjC+p4nT+pp5apmLP55uC7nPYXuWwj9Og1nfgT019gWPEQmmRMwppXj4/KFcNb4Xt8z9kQOVDuLMJnTDX81qSuH26Zw1NItZU/oxuldqp7yJFF2HMptJv+WWJhfubfJ4i4WU6Re2Qck6hjunDuDHvHIcnuCbpObmUtotGref2l/qCCE6sDizhgoTTkWzPAoGJDWxPEosmU0aMyf14ZWFuUHDE5tq6CkmlWd9l/Os73KUVkv6wD+iWaz8fvIfOK//ee1S9o5kdV4Z98z/iWqnJ5C51SCRppMy/X5aHJNfq+X+Cc303OkeqNgL8y6DW77usOtQdcxSiU4hI8FKeZ2bRQ9NY0thFWW1HnTdICXewkl909u0+16I9tbt9ttw/LSWupUrMZyuiI9Tdhu958xGi++88wUnDezG5eOy+XDNfhye4DkV4cSZNcb0TuP6CZ1vLpkQnYlSij4Z8eSWtG55FLdPp0+39qsLZ0zowz8WhV7aotkFhJUbS9oK3CYDdBePLnuUlLgUJmVPCn9MJ/Pt5oPc885anI3qdRWyj/KIEZkmLhxkZsBzNRjA6B5HeiyDElL43FC8FbZ9CcMuivUlxIQEU6LNfLIun9OHZHJcqp3jZPK46OSUppHz/PPk//JBapcubb6HymxGi4uj1z9ewT5yZPsU8ih6/KIR1Lp8fLWxMGQPVWM2i8bI7BRm3zgec6iJDUKIDuWOqf15/F+bW7w8iqbg3OE9SG6nnimAzGQb/3flCTz4/s9BAUFjNRu+pWrVx3grClFxdhKGjiD75orDjzt9Th5c+CALr16I1dT5G4vX7i0PGUiBwoGVBJpuVHz8NBvvbKjBZob8qmYa2dy1sPQZCaZE12IYBu/+uI9HLxx2tIsiRLvRrFZynnuWqi++pPTVf+Deuw/D7Qb9yBeFio8HwyDlkovpNmsWluzso1ji9mPSFP/vqhM4ISeF57/fidPjC7mYb4LV30I5c2Iffnn2YCwSSAkREcMwWLW7jNeW7mbHwRrq3D7sVhPHZyVy6+R+nNQvvU2Hy150QjaPfbY57OPN9fTEmU3cdmr7Lsx9oNLByt1lIVOk11e16iMqV35Etwv+C1vfoejuLVR8/3/secpDv0f6oQXmfOuGzrd7vuX8/ue3R/GPqoeaCED/uL4b36zcxtYSnaQ4xegeGo9MiWuwz8B0jQSLP7vfvkqDCqdBqq2J9+fBTVC8HboPiuVlxIQEU6JNrN9fSZ3bx8n9M452UYRoV0rTSJl+ISnTL8S5ZQsVH32MpyAfw+nClJpKwsSTST7//E49rC8cpRQ3Te7HDRP7snB7Ef9YlMvOohocbh9Ws0aFw8NvLxzGZWOzZfFuISJkGAYfrNnPM9/uoLzOjcPtazDIKq+kliU7SkiLt3LfGQO5anyvNgmq7FYTt0zuy2tL8yLqfa7PalIM75nMiOzYp0APZ92+CmbOWUmd24c3EEwpSxGGNw0MxaFbZN1VR8WS+WSc/wvsxw/Bmr4Ia8YikoZms/2h7VQuqyTt1DQA6rx1zNk4p9MHUz/vq6CgIvS8qKpVH/P0yr28dkEiFw80sJrgq51ePt0anJAixQZFtZBqI7KEFLsXSjAluo53V+/jqvE5aJpMGhddl23oUHo8MvRoF6PDMWmK04dkcfqQrAbbr3hpGb3S4iWQEiJCPt3gNx+t518/HwgbwBhAndtHndvBY59t5se8cv5y+ShMbfD9/ODZg9l6sJplO0sinh9p0RSZyTbm3HhizMsTzpYDVVz36opGQxINEvo/h+FNxl0+EV9dPwyfDfe+dRg+NxnT9mBO+Qal/IGXyWYiaVQSNZtqDgdTAHlVeRTWFtIjoUe7XU9MGIY/WNnyOdQUAgoSs2DodOh3aoP1smYvzsXlDX6/HVmg+X6mDHmfBM2fxXX6YAvTB1v429kNgyWzpvjyejs9ElXzCSm8LnCUx+RSY02CKdE65Xtgx9dH3uD2NJx9pvHF+gP8+4FTj27ZhBDHlHF90li7t5xTju92tIsixDHhd59ubDKQaszh8fHF+gLizBpPXBr7uZqapnh5xjgeev9nvtxQiNvXdEAVbzXROz2e+bedTEp8+8yVcnl9zJi9Mnhul/JQsaKU0q+34j6wBM2mYettw97PjjlRw5K6Lehc5hQzjj0N58daNAtlzrJjJ5hy18GaubDseXBV+ucn1ffzfLClwMT7YNyNYI3nu61FIZe9ObJA8yQe8SbziuVp7KqZBXo5kpDiySVuhnYPM7RbKVAdc9i3BFMieroOu77zTwbcvxpQ4A1095rjsOgG71uPp0fhbyHpHNA65ptfCNGxjOmdxj9/3Hu0iyHEMeE/W4v4cG1+yECqdvMPVP34CZ7S/WhWO5bM/qRMugpbznAcHp2PfvIniDpjaFaIM0emqNrJ/BV7WbCxkCqnB6X8a8JdOiab/71gKCt3l9EtMY5tB6vRFLg8OgZgNWloGvTLSODO0wZw3ojj2nWdya82FuIM8ZpV/fgpVasO0PPG40gamYQyKao3VFOxpAJvjRfDZ6BMDXvzvJVezIkNb6UVCo8eesHaDqf6ILwx3Z9+3BsmaZK71v/z3eOw9nX0mf/CEWYtsPoLNC/ST+Ax70weM78ZUUD1+Gk2xr5Sw4MT40LvYLZBfMecOiLBlIiOxwHvzoC9y4NbLwC8TkzAINdG+PBWyB4H1/4TrAntXlQhxLFlbJ9Ufv3henTdkCHCQjTjhf/sDBlIVa36mMqVH5Bx9t3Y+o1Fmcw4dq/BsWMltpzhADjcPl78YVeLgqldxTX86YstLN5ZgoIGazQVVDjJLa7lyQVbSY238OrMcYDimy0HKatx4fUZpCZYmTQgg6HHJbf00lvlpR92BSW/0V21VCx6l+zbepIy/ki5ksckkzA4geoHqqlaU0XKSUfmdPmcPqrXV5N1RcPXUDd0kq1H59qi4iiH2WdA9QHQQy+w24DXAaW7YPZZwJ9D7tJ4geZ3fadTbiTxvOV5rHhpaqrewHSNq4dbeG6Vm5GZIYJr3QeDzo3s2tqZdBmIyHnd/haMvCWhA6nG3LWwbxW8fh54ml7ATQghMpNsJNnM5JbUHO2iCNGh7SmtZUN+ZdD2Q3NW0s+6i/jBk9CsNpTJTPzACQ3WcwLYmF9JXoh1oZryY14ZF/19Cd9vK8Lt1YMWuwX/UEKvblBW6+b855ZQ7fRww8l9uP/MQTx4zmBuPaXfUQukdhbVkFcafM2HhqclDj8+6DFTvInMSzIpmFdA9fpqDK+Bu9jNvhf3YUm3kDopteH+ykROUk6bXUPMvHcj1BxsEEj1faaazL9VU+s+MoZv9lo3p80NvGa6F63mADYVuuet/gLNh3ytn8hJrhdwR9B/8+jUuAbPfYSC/lMh+bjIrq2dSTAlIvev+6Fw45EhfZHwOqF4G3z6i7YrlxCi0xjXJ421eyqa31GILuyDNftDpvM+MmdlYrPn8OkG76/ZF/FzbiqoZOacVdS6fBhNZxIHQDegvNbNFS8vJ7+imXX32snestqQyy0cGp7mqTgdwxc8zKz7+d3JujyLwncL2XzXZnb9YReWdAv9/rsfmuXI+ayalWuGXINFa7+1slqkZAfsW+lfELcRnwHPrmxiWJ7PzRRtfciFeesv0Fy3fTm6x4nh81K4axuXfN0dt9EwuVDe/2fvvMOjuK4+/M7MNu2qV5pE75jeu4HYuHfjAu69O4mTL44TJ07iOHbsuOACtrHjTtwbYBtM772ZjgQC9V62z8z3xyIhsSPtLkiIct/n0WM8O7N7d3d25p57fud3HolhUqejQVZ6nIz7idijDXtrMEfByAcje48nESHzE4RHZT5s+xzU+k3YOrxYidMHmQ9H47AE8rdB3av97oA7TPlhiDs7euoIBILjY2BGwITi2iHpLT0UgeCU5WCJs9bOuy51a1ZC4dd0DhY7w3o9TdO57d21Eddn6UCV2899H6zn6wdGh/VazUm1R0UziARr5Gm+sh7YGvCNSByXSOK4xJCvcW33a090mM3PqtcDsjkDHhtp4dnlHu4bYmmw79Pdyjcs0/rg1IMDz4YaNO8ecQOFvEuaXoJJCiMar8EUBT0ugg5jwj/mJCOCKUF4rH+nwYdqVjEeH9NA0SAAOqx9Eyb9palHJhAIziAGtoth57IvYd36QI2mNQbSekObATQquBcIziKcHuMal2NrVkI+TwNGAseyZE8hVe7g1wynPkvVdXblVbInv5KuaTFhvV5zEW01IRtcR47K09ZiSZyMNW0OkhyZiYRNsXFp50tPuoufruus2l/C+6sOkFVUjdunEmMz0T89nptHdqBTSnT9A1QfbP4YGjDJGNxGYXwHU6N9nwZKe0imnIOkGj5u1KBZBaZ4M/jM8iSJeiUWKYxzz2yH9qPg8tdP6eu/CKYEodE0WD0jKCtVQzirGKheWPs2nPsEKOK0EwgEx1CZD+tm0WfNDP7odKPPA0lXQT5yvYhtA6MegT5XgeXsa3gsENTFZ+RLTf2aFUeP0JmgBEcjfX3qMGPxfmPThmUfknThI9i7j6zdbu8yDHuXYceMV+PtZZk8c1XfsF6vueiU4sBrUOdVV54myfcTO3Qg1qT1VO8spXpHNa2mNB4g2RQbg9IG8fiwx5tr6EFoms5Hqw/y2uK9lDl9Qc2at+dU8MnabHq1juWx87szssuRlhPOYtAbt6x/6lxro32fJAmetc3iFt/juA0+z4Y4pKdwpfYcHye+SfvqLYFxGEgNMdsDfa+G3AmTnoQwFgZaElEzJQiNu6xRw4m6qxiNonrBWdTEgxMIBKc9O76Dl/vBsv8guUqJllxIflfgmuFzBv6K98K838NL/aBwd0uPWCBoMX76JZ/le43vpQ3VrLj2raN04ax6+9otCoPaJxg+T13KnF7WHygJ2h5ZfRZ8ufEwejjFVs1I+yQH3VsZZ8dih15JwoTbKV85m/1//oBdv95N8fxSYgY0/BmZJBNWxcoFHS9g+sTpKCdp0u/1a9z74Xr+MWcHOWVunMcEUhCQcXr8Ghuzy7jtv2uZtSwz8ICn6ugiVQPU7fvUEMPN+3h2cho2c/ihRJRZYXS/7mQ88iPcvwaG3gXW2ED/KMUMSBDfHs77Ozy2B8576pQPpEBkpgTh4C4P/PAayExB6FUMIPAc7nKIOU0a2QkEguZny6fw7YMBSV8ovNWBBpNvTYA7FkBK9+Yfn0BwCrEpu4wHP96AT204KGmoZiV2xJR6++k6XLUfhmkAACAASURBVNa/TcjXLKj0YDEpeNX6Mr9I6rMgYHhR7VWJtrbs1PPecZ357WebqfYEy8yOladJpnJMCavQ1ZVQJ1yJtplRNT+XdL6Eqb2m0imu08kYOhDISD3w0QaW7CnE7QsvK+T2aTz3w05sZoUbejrCskIP2fdJU7m0X1sS0tJ58KON+DTN8DOFQBCl6Tr3je/MAxO6IEkSJLSH8/8R+PN7A9brlujTIng6FhFMCUJjtoPeuLY1rO7VuhZ4LoFAIAA4vB6+ebDhZpGG6IGV1XcvggfXgy0u9CECwRnCn77a2uAE2sgIIvXaO2trl+pikiWuHNgWuyX0NNDtUzES8Edan6XIEm5fywdTk3qlEWMz4/KqNKCWrEX3x+EtPB9v4SSUqGysVhfn9U7lxiE96Z3UG3sLzGk+XnuQpXuKDM+DUM2an/puO0MzhtNFCp1NCtn3SZLBnsiYrmbWPTGJ+TsKeGPxPrbnlGNWZCQJ/KpOosPCXWM6cdXgdsTaGnA5NFkCf6cpIpgShMaeCAYWmMcSehXDD47kph2bQCA4fVnwt6BAKiyHUPRAlmrTxzD8npM8aIGgZdiTX8mefOMebOEYQdTFYpK5e2znsF431mZGNZDnRVqf5VO1hifTJxFzVS4fDzvApT8nU+WX0Q1DxWNRMPs7Mbx9Mi9cNLjFmorrus5rC/cdd7Nmv6rz5He7uNUygXG+uZhpfKH8z+OsvL/FwKhCNkP/G49I88CkyEzu04rJfVpR7vRR4vSiahqxUWZSoq2BTNQZjKiZEoRGMUOvK0BqfOWp7ipGEJIM3S8M9AoQCASC8sNwYIXhQyH7nECgjmrFS4TV8EYgOAN4e1kmPi04GxFJo14Am0nizZsGk5EUXlalTXyUYbgRSX0WQOu4KCymFpp26jrsWwjvXQ4vD6Djisf5Qv4/EqnASohrDYH6sok903hj6qAWC6QA1mSWUOoMHm+454Bf01m5vxjfkLsxmYID27D7PskKDL/XcIxxdjMdkx10SY0hNcZ2xgdSIDJTgnAZcT/s+DpkXUODqxgmG4x8qJkGJxAITjvWvt3gQ2E5hEKgBjNrKXQc2wwDFAhaHo9fpaTai8ur8tMv+agGCr9IjCAAHv1Vd0Z1CV8lYjHJXDc0g/dWZgXVaoVbnxVlVrhn3MmrK6qH6oMv74FdcwKLMEfoKh9mvvUx3vNP4h31AnyYqObogq9JllBkiZ6tY7lnXGfO753WpIFBtcfPt5tz+CW3gjKnj2ibiS4pDi4f0I7EBlwWP1h1AJeBnX0k50CUWcEV2wkpfSgcXGXsptcYigXSh0NSeJnNswERTAnCo3VfSO4Gedvq1U9lPVLfFadmFaMekgzxGdB24MkYqUAgOB04sKxBU5tw+pwA4PdAzkYRTAnOOLYdLuftZZnM2ZqLLIEsSUHW5DVEYgRhtyi0imvkN9UAt4zswAerDmAk+TfqKXQsOjpXDGwX8eueMJoKH18PWcsMazMTpCoeNn/F/aZvWKj1Z5XWkyISsFrMtBl6BZcO7hzcp+kE2V9Yxcwl+/lq02FkSarX68tmlvnXvF1M7JHKPeM707ddfL1js4qdhkUXkZwDLp9KbrkbrvkvzBgDlXlhGVIAASOx6DS49t3w9j9LEMGUIHyu+wjeGA2uMsKpoarFGgM3zD6lG64JBIKTjLui0YfDcgjV/OAsbeKBCQQtR265izv+u479hdV4/ZphrdKxRGIEUROYRUp6op1JvdJY8Et+RH2FIJAJuXlkh5YxnljwFBxYHtLkxiRp/ErZwK+UDYENsgVyV0LKD006nB+25/HIJ5vwqqphlrHGVGLe9jx+3lXA7yf34NZRHWsfN6qVgsjOAU0PZMWwJ8IdP8N/L4ay7NBGQKaowML4Ld9BVGhL/bMJUTMlCJ+4dnDbjwETCTmMIlLZBFGJcOs8SOjQ7MMTCASnEaYGjGqOEE6fE5DA4mjkcYHg9GF/YRUXvrSUnXmVuHxqWIEU1DeCCI3UoIQsFM9f049uaTFYI6h7ijLLjO6SzO/Ob4E2Bt5qWDOjnrQPAiY3qc9VUu09+vm+tcHL+Hfr9NNUvZC7JZD5biJ+3J7Hw59sDHy3IeJRTQ8EVs/O28WbS/bXbo+1GQekkZwDJlkiLurIHC4mDe5aBBOeCLStsRhk4SzRENMaJv4J7loI0akhX+NsQwRTgshI6Qb3roABU/HKNryygaGE2R5Yweh3fWDftF4nf5wCgeDUJr59yF3+Ot7Gmxu8HK5oYFJptkNs6yYemEBw8imu8jBlxirKXD7UUH7dxxCJEYTXr9KztXHT2lDYzAqz7x7BsE5JmJXGs1uyFMhIXdyvDa9PHdgypg1bPqWhaW5YJjd+N6x8tUmGcqC4moc/2dSgnXnufx/h4AtXc2j6NPL/9yTuQ9uBQCbq+Z92sWp/MQADMhIwGXyWkZwDVpNMj9Z1yjEsDhj5ADy6A6Z8AINvhx4XB/4G3w7XfQi/3hGonReLV4YImZ8gcqJT8V7wAudunMjXY3NIzvwanIEfOlGJ0Ocq6DclIO8TCAQCIwbfCvsWgNfY6hnC6HOiq9DzkmYcpEBwcnh5wR7KnF5Dc8rGegfVEK4RhKbDqGcW8rvJ3evJx4woqHCzI6+SSrcPm0mhdbyNXq1jefWGAYx6ZiF920Wz7XA5iizhV3UkKZD18Gs6v+qVxh1jOtE/Pb7R12hWVrwEvmrDh8IyudE12PFtwOjmBPvZvbU0E59BOiocO3O3T+OFH3dzSf82LN9biL+BYDtsMxCLwhgjAxJZhs7nBv4EESGCKcFxsWBHPu3SkkkefwmMv7ulhyMQCE43Oo4LLLg0EkxBIw6hkgK9LhdNewWnPW6fyqfrD+EzmCSHnGzLbiRTJZLkI6Z/P6J7jwUarpnxazp+TeXZebvILnHyp4t71XOo0/WAdfbMxftZub8Yi0lG13UkSULVdJIcFrqmxjC2axLTbxxEXrmb5XuLKHf5kCVIcFgY2zWFhOOUEjYZug6lWQ0+HLbJjWKBkkxo0/+4h+Lyqny2/lBQEFRjZ5504SPYu4+s3W7vMgx7l2H19l2TVYLDpvB/F/Rk+s97WXfAuFY0lBmIzSxzx+iOLWrvfiYiginBcTF7XTZThqS39DAEAsHpiiwH2iX8/FS9lgthOYRCYJIz8oHmHqVA0Ox8uznHcHvDk+2hxPRNxJL0Nop9P+gKIIGkgS7jLR2Gr3Qkur/hrJDLp/LxmmzSYm3cPS5gcV1S7WXa26vJLKqudZjzHGM24fS6yC51EWVWWLaniNFdk7lqUAu49IXC54IQzXjDMrkB8ATMcnRdZ3VmCd9vySW/wo2m6yRHWzmvdxrjuqWiNBCgzNmaa+i/FYmduVmR6NkqlvHdU7GZFW55Z42hZDAUVpPCdUMzIj5O0DgimBJETG65i40Hy3j9xkEtPRSBQHA6M/Qu2D0PslcH6hPCxWyHsY9Bq3Oab2wCwUnim8059eyxazCabMvWPKLavYukOEH2Hpmk1z/WkrgcS+IK/JW9KZ6fSsXabw0lgi6fygs/7ebawemous7FryyjuNJjmCE7FpdP5Y731vLSdQM4v3erE/0Imh6TLSDTa4S6Jjc9Uxq2EPBIVj5ensnMpfspc/pwedV6fsbfbsnBalK4bVQHbhrZgVhbfYOu3fmVht9vJHbmPlVne04gqBveKYnHL+jJ03N3RBRQ2S0KH94xjHh7C2cNz0CEAYUgYj5ff4iL+rYmyhL6AiAQCAQNopjg+k+g3RAwG5jZGGGOguH3w+hHm3dsAsFJorjK2Ajh2Mm2bDuIvf1rSOYyJMXbYLcRSVaRZD9lKxZRuvBV4kZcSbsHPqDtve8QM/BCXHtW1+7rUzWm/7yHG95cRVGYgVQNbp/Gw59sZOuh8rCPOWnIcsD6OwShTG50v4c7vsrlmXk7ySlz4zwmkAKo9gQaK7/y814ueHEph0rruweWuQxkytS3Mw+HCvfR57lpZAeeurQPNrNsaEhRF5tZJi7KzP/uHkGftkIW3RyIYEoQEZqm8791h5gyWEj8BAJBE2Cxw7SvYPRvAr1LjKx5JTmQjUrqAlfMgIlPiL51gjOeer2DzMXYM95uNIiqi+pUKfwql9bT2pA0cS+yxYqkmLB3GUbCubfV7qfp8PbyLA4UOw2NDRpzmoNAQPXPuTua5P02OYNuA6XxFgx1TW6M2Kp3ZlWhJawMkMevkVfu5rJXl1NQeTTT3hR25gAxx/TpunZIOnMfHssNwzKIMis4LArKkVm9WZFwWBRSYqz8+lfdWPzYeBFINSNC5ieIiFWZxUSZFfq2Ez9KgUDQRCgmGPdYINu0ex6sfgPKDqL7XByokmnXczimUQ9BOyEtFpx5JEUby67qTraTJmaCHDzhL11aStEPRXgLvCg2hdhBsaRdnYZznxPNpxE32AFSFop9H6qzS4NjOLY2CsJzmgNYf6CUw2Uu2saHmV0+WQy5HVa+HHK3hkxuXFIUr/ouwqfWDzIbc1dUdZ1yp4/b3lnLdw+NAaBzSjQ2sxwUkNW1M5dkBVvHAUiyCXfWJtwHt9QLek2yRNe0YIfkjskOnrqsD3+4oCc//pJ3JHvmJy7KTPdWMYzqnCzMJk4CZ1wwtatkF+/98h5LDi2h2leNhITD7GBS+0lM6zWNjnGNW4EKAlS4fczZksvBEicVbh8Jdgs9WsXy4/Y8rh2SXs/9RyAQCJoExQQ9Lw78ESgfv+3fi5gxfpDhREIgOBO4tF8b1h8oDaqrOTrZfp2o9ARizrGDIlH1SxXVO6oxxZoonFtIuzvaEd0rGl+pj5z3c8j6dxaJExIxRZuQFAld92JJWkzRukMhLdZriMRpTtd13luRxR8u7Nk8H9DxEtsaOk0ItGBQjwai4Zjc6EClZuUnX79628MJMP2azr7CapbtKSSz2MkX6w81mNkK185ckSWmDm+4N1+UReGy/m1DfiSC5uGMCaY2F27mryv+SnZlNj7Nh6ofvSh5PV6+3PMl3+z7hm4J3fjryL/SNaFrC4721GVnXgUzl+zn+y25yJKEy3f0c3RYFKq9KvEOM7nlLlrHnWKrUAKB4IwjPdFOdqlTBFOCM5ZL+rXhyW+2Gz4WO/RKLGm5FH67iEMz3Sg2BVsHG8nnJXPwlYO0vb0tMX0Dvw1LioX0+9LZ/dhu3Nlu/FV+dFVHUiRKl66mZMECks57oNEsUw2ROM15VZ3FuwtPvWAKcF/8Kvobo7G68pH18GqTAHxyFLd7H0erUw0TSYDp8qnc8s5aJvdpxf3ndmHutly+2pRj2JA5lJ05QO82sXRMFg1zT1XOiGDq5wM/8/ulv8etNuwG5df9+FU/W4u2cuOcG5k+YTpDWw9t9rG5fSplTh9ev0ZslIm4KPMpm9V5f2UW/5izA5+qG/7gq4+smn28OpvP1h/irZuGMKJz0kkepUAgOJtIT4ziYLEz9I4CwWmKzaxwzaB2fLT6oKEBROJ4PymTO9fbVrmlEs2nETuofkZFsSnE9I3BV+JDMklUbKggunc0hV/lkXLl1dgyGg8CaojEaQ6g0u0Pa7+TRVZRNbOWZ/LZ+kO0kv7MO/yFVhRjlYzNIGqRZHRLNLd5fsdWf33L90gCTABZlnjmqr5EW020T7Lz/dZcw7lVKKLMCg9NFAmAU5nT3oBibd7akIHUsbj8Lh74+QF2luxsljHpus7arBLufG8dfZ78gfH/Xsjkl5Yw5B/zGfr0AmYu2UeZ07jYsaWYtSyTp+fsxO3TQv7YvapGtUfltnfXsGp/8UkaoUAgOBvJSLSTXeoKvaNAcBrz0MSuxNsthuYSkhJ8/qtVaq2M71hMcSZUl0raFWnkvJ9D0bwiNJ+Go1c7XPvWUbpwVsjxROo0ZzYYR0ug6zov/rSb819cwserD+L0quz3xHKR5++8rU6mQo+iSjdo0mu2B+zUz7kGz+2LWenpHLRLpAGmWZbIKw/MTbumxfCPy8/BZo5s2h1lDliuj++eGtFxgpPLSctMaZpOhdtHlceP3RLI0DTU4Cxc/JqfXy/6tWEg1VBRpuII/AhcfhePLnyUOVfOadJM0e78Su58bx2FlR5cPhVdp55DTmGlh//8tJvnf9zNtOHtefzCni1eHLhyXzHP/rDTUNPbWKGly6dx+3/X8vNvxpMW20gHcYFAIDhO0hPsrM0qbelhCATNSlK0ldl3D+eq11dQ4fbXX9TUg+cISrRST8ZXF3+5H1O0ieQLkjHFmcj/Ih80OPCvGVhSewbV4xhR1/zC0WN0yP1TYhp3zTtZ/Pnr7Xy2/lCQoUYVdp71X89//NdwvryWq5XFJEvlmCUNrzmOHuOuxzxoKkTFU1XlwaTsQvXXX1iu564YRkAlSxLVnqMZu6sGtcOvaTz5zXY8fg09RJIqyqxw++iO/Oa8buF/AIIWodmDqbxyN++tzOKDVQdw+1QUWUbVdRRJ4upB7bhtdMfj1oEuPrQYnxqcsi2aW9RgUWbHP3ZENgVWBordxWwq3MSA1AEn8hZrWX+glJveXm3Yh6AuriNBy4erD5JZXM2MqYMwKS2XJPzP/N2GgVQ4hZY+Vee9lVk8dn6PkzxqgUBwNpCeaCe7RMj8BGc+nVKimfPwGO787zr2Flbh8+uouo6uRgH1FxTsXewBGd/6CuKGHnXXVd0qlVsqSbs6DYD4kfEo0QoHXjxA+9/9FrV6SFhjicRpzmFRuGFYxol/ACfIf1dk8dn6Q/VqvY/Fh4nvtBF8px2V6llVmXH7U5g5Oh6AaKsJvxo8i4s0wNR0nehjbNGnDMmgR6tYXvl5D0v3FAH1nRTNioQsSfRPj+eBCV0Y0zUl5OsIWp5mC6ZcXpXffLqZBTvy0QFvzcmiHj3JP15zkP+ty2ZARjyv3TiIREdkXZlnbZ1Ftb+63jbVpVLwVUGDRZnlK8pJGJsAgNvv5t1t7zJgwokHU5lF1dw8a01tXVE4uHwqK/YW8fiXW3n26n6hD2gGDhY72ZxdFrQ93EJLr1/j/ZUHeGRSN8wtGBAKBIIzk4ykQDCl6/opW28qEDQVreOi+O6hMWzPKeetpZnM2ZoLVQPRrYVI8tHFY8WukHp5Kjkf5CDb5HoLx+ZEM/Ej42v3tXexgwaZf3+Wdve/j2wJKEkqN/9A9faFtLrhGcOxhOs0B3BBn9ZN/ElEhk/V+PePuxoMpBpT2Xj8Got3F7K3oJIuqTF4VQ2rWW7EXTF0gAkBVVIrA9VOv/R43rp5CIWVHmavO8i2QxWUu33EWE10SY3m+qEZpCfam+7DETQ7zRJMVbh9XPP6SrKKqw17F9Tg13T8ms76A6Vc8NISvrp/VNgOcWXuMnaUBDeKc+5xNlqUWbW9qjaY0tFZfGgxfs2PST6xj+Iv32yj2mtcgBlKKvfN5hxuHtmB3m1Ofu+m91dloRnkmiMptFR1nQU78pncwhdTgUBw5hFrM2NSZEqqvSRFnxpSIoGguendJo7/TOnPM1edw6GywVw790e8x0ynUi5MQXEo5M3Ow1vgRY6SiR0YS/rd6ch1anPkKBNylBnN5ab4h+kkTX4ASTbhLczCX5bX6DjCcZq7qG9rbObw6oiaix+356M1UO8dlp25qvF/n2/FbjWx4UApqTFWDpW6ghoZh21lLklM7t0Kh7XhuWVKjJUHzhXGEmcCTR5M+VSNW2atIbOoGq8aumN04BidoiovU2as4vuHRhNjM4c8psRdgkW24NPqy/xCFWW6DtQv5JQlmSpvFfG2+KD9wyWv3M2q/SWG+tdwpXJvL83khSn9j3sMx8uWQ+VBDekgskJLp0dld34Vk/s0xwgFAsHZys68Ct5emkm1x8+YZxdiVmTi7WauGdyO64dkiOBKcMZjNSl0Tk7hV+1/xbysefXavgAkjkskcVxi40+im5AtMdja98a5Ywmu3SuRrFEo9jjkqNjGj6XxBWGAlfuLWzxz/MbivYbKoHBVNqoOGw+W8u9r+/HajQNxHrnmHBtMQXgBpsUkc+fYTifwjgSnE00eTH23JYcduZWGgVSjXaM1nfwKN28vzeSRX4UutvNpvkBHx2MIpyizLpJfJfe1l1Hj2qEkJmJKTEBJTERJCPxbtodOtb6/KsuwRirsH7Gm8/3WXP5yWW9iwwgkw0XXdZxelWqPnyqPn2qPSqXHR7Xn6LbMomrDYyMptNThlHMnFAgEpy9rMkv4yzfb2V9UVduqwe9VAZVyl4/pC/byyoK9nNs9lacu701qjDDAEZzZ/Hrwr1mRs4JST2RmLLpmxl/ZG10/THS/89E1P+akDBLGTquV+dVgwoeMjpejJRfhLAgXV3lZua+YkV2Sm+bNHgc7cisNt0eisomymOiSEkO01US01cSknmnM35HfqMLKCJMs0S0tmj5tT77aSNAyNHkw9fqifYaa1XB+kB6/xrsrs3hgQpcgQwZd1/EdzsGzezee3btwZm7G180Z9A7CLcqswS/rxDmS8RcU4N61E7WkFLWkBH9p4L9IEqbExECAlZiAKaHOv48EXd+u9RytCatDJD9isyKzcl8xE3qk1gt+qjw+qjx1A6LAX2Xtv9V626vqHOv0+rGZFRxHLgwOq4LDUvNvE9E2k6ENK0ReaBljOyNalgkEghbmq42H+b8vthia4tTgPnK9/WlHHmuzSph99wi6pEafrCGe1eSWuzhc6sLpVYm2meiQ5Ii43lkQOan2VGadP4ub591Mla8KTQ89wdc1M6qzI+6ca4AfAYgffSN5H/yO2MGX1ttXQSVDKuAO+TueV6+jiihcHl9YC8JOr8qMJftaLJjy+jXDcgWITGUjSYEylRqeu6Yvl01fzoHiarwGCh4jFFkiwWHh7VvCM/oQnBk06Qx42+FyskuC+yFE0jXap2rM35DFWIpx796NZ/eeIwHUbmSHA2v37li7daXDiPNwVG/B66+od3wkRZkA6bEZtLn5PsP3o+s6WrUTtbQkEGCVlASCrdIS/CWlePftx19aQlnseaAEr4xG8iOu8vi578MNALWrIg5rnUDIEgh+arbHRZlpGx+FwxIIjGJsRwKkI8c4jhwTyn7+95rOZ+sOoR5zIYqk0NJuUUhPFJ25BQLBibFwZ0HIQKouqgYl1V6unbGSuQ+PES0amglV01m4s4A3Fu9j6+FyLKaji51ev8aYrsncNbYzQzokCJOQZqRLQhc+veRTHl34KPvL9+PTfEGyPwBZk1GR8ZYOw1twIXVbilpSOhDVZQjlqz7FnNSOaFy8an6RifJGrPiQJLjetJgDeiq3Z44mO8wF4VX7Syiq8pDcAtJbkyw16KAcqZ25tc65bbeY+OyekUybtZo9+VWNugQC2EwyqbE2Zt89vEU+B0HL0aTB1IId+Xj8wSdbJBmaao/Kx9P/Rw/3JqzdumHt3p3Yiy7E1q0bSnz9QGjaliJmbpkZ1Gcq3KJMNCt9o6/A7VMNiyclSUKJdqBEOyA9vcExK0/9CM5gi/ZIfsQ2s8yfLurFDcMyTurNaNrw9ny9KQfV4CIRbqGlrsMFfVqdrCELBIIzEKfXz/0fbYi4350OlDu9/OZ/m/ngjmHBTyw4IfbkVzLt7dVUuv21NSnHyp4W7Cxgxb5i2ifZ+e9tQ4XsshlpE92G2ZfMZnfpbt7f/j5zMuegoyNLMn7NTxIyN5cU8EvZJXztm4QEQYFG/OgbKXj3QW4dHs1GuZzJchGKdHQvSYL2FHCD7wu22mGSaRMLtEGNjstikjlc6mqRIEKWJWKsJircwSZgkahsfKoWdO7G2c18es8IPlt3iDeW7KO4ylvbQ7QGh0XBbjVxx+iO3Di8PdGNmE4Izkya9BvPq/BgZKYSaddo98hxdLj7/0Lud1W3q3hjyxuGj4VTlKmj89WyFL5e+hOv3DCACT3SGt2/IWJtZkoNgqlIfsQmWSI5xnrSV/X6tI0jPSGKPQVVho+HKrSUJbh8QNtGHWsEAoEgFF9tPGy4PRyJuKrDmqwSDpe5aBsfniOsIDTbDpczZeZKnJ7GeyfqekDqtSe/igtfWsq3D44O25lXcHx0S+jG30b/jb+O+iuV3ko8qocYSwy25dORDj4HykdcIW1ghv9ilml9qJu7OS8pD3rDB2tLOCdVrhdI1SBJ0Nbup8Sp8aLyMv+SbuQD9bxGx1S3Qe3J5qpB7fhg1YEgQ61IVDbpiXYykoLr5K0mhRuHt+eGYRlsOFjK91vyKKh0o2k6yTFWJvVMY3SXZOQQSiDBmctJmQFHmmYNT5kKibZEbuhxA7N3zsalBssLG30NzYyncBI+jwKo3PfhBp6+/ByuHNQurOM1rxfnmrVULVzIiN1ecloPwneMvXokP2KfqjO8U1JE76GpeGhiV3732ZaQKWwjNB1UTWswuycQCASh0HWdGYv3B/V1iUQijq7z3oos/nBhz5Mx5DOe3HIXN7y1impP+PcFv6ZT6vQxZcYq5j0yBrtFLLI1N7IkE2etY3QwYCoseQ6AYfJOhll2UqDH00uq5lJ5OcNNm7lJmU/eODMfb/E0+twj2ilYTfDDLid/7PkRRXoc87QGsr86QQ1qTya3jOzAR6sPYjSDDEdl47Ao3Duuc6OvIUkSg9onMqh9CPdEwVlHk575KdEWZImg7FSkZgbJ0eEXsz466FGyK7NZfnh5kNyvIXTNjK98AL6SMbXb3D6Nx7/aSqt4GyM7GxdR+ktLqVq8mKqfF1K9ciXWzp2JPvdc7rl4FF9+fhAMTCjC+RErksQFfVoRF9V0Tn6RcHHf1izfW8TXmw7jCrNWAQLSxN+d34O1WSVc8soynr+2H33bHb/FvEAgODvZW1BFQWXwxC4SibhX1fl0/SERTDUR03/ei7OBQCqUM29hpYfP1x9i2ogOJ3fQAohtDR3Hwt751AQWqVIZRY+YgOW1u6XHybifCNiid3ixEqcPMh+OxmEJZFfe2uDlgy0+nhpv5f45bkwy/K3TcxkveQAAIABJREFU6/ysnkN51i9BC8JeVSM9oeUazbZPctA/PZ61WSWGCqlQKhtJkrjwHNErU3B8NGkwNaFnGm8uzQzKcETaNVqRJEqrvSSE4RAkSzIvjH+BZ9Y8wxd7vsCvqai6caq5dEUFxT8U4sn1I5uzMKdurtcrwe3T+PNX25j/m/FAYLXUu28flQsXUrVwEZ7du3GMGE70+HNp9eSfMSUFMknJwNANVSzbU2SYVQv1I7aYJO4Y03L9CCRJ4h9XnIMkSXy16TAug14Nx2IzyzxxUS+mDm/PraM68M3mHG57dy3XD83gwQld6xUoCwQCQWMUVHowKRIco5aOVCJe4QqWWwsix+n188WGw4Y9dsKRXbp8KjMW72fq8PbCkKIlGPtbOLAMfOErdlQdXlrt5fEx9WuefjPSSqtomb8v8bDjCxdYbkFP61VvQVgCxnZLCWvOBoHFk282HeZwmQuvXyMp2srwTklM6pka5OQcCdcOTmdNZknEx9nMMjOnDRLqGsFx06TBVL92cbSKsxn2LgrXzADgh+15LNtbxEd3Dqdn69AN5WRJ5vFhjzO151Tu/fYlDngXgV7/Al70Qz5F3xeTOPkuUjLGGN4AAA6XuVg9dxkdNi2lcuEidK+XmAnnknzvPdiHDkW2GhdX/uXS3lw6fVlEkgiAKLPChee0avF+BIos8fQVfRjdJZmXf97DwSNWoGqdm6lFkZEkGNIhkUcmdWVwh0CqW5IkLuvflhGdkvjDF1u57NXlPH9NP3q1Cf3dCQQCgVFrCYhcIq5qeos3Dz0T+HpTjmHbjEhklyVOL2sySxjWQvL1s5qM4TD+cVj0dNgB1WMjLTy73MN9QyzE2+p/+Tf2NXNj34ByZq/Whknev9R7XAeirQoVbl+DvTJ1XeeH7fm8tmgvu/Mq8Wt6vWD90/XZmGSZm0e2Z9rwDmw7XM7stdnkVbjxqRrxUWbO7ZHKNYPSibMHv8ZHqw/yn/m7efG6/vz1218od3kxaHcaRJRZ5j9T+rdojyzB6U+TBlOSJHHvuM48+c12w/qbcLpGQ0Cu4XX6uPr1FXxx3yi6t4oJ6/UTLK3Zt2MSbnUMSlQ2kuIEXUKtlin48k+BG0CXozJDoxuAx+Nnxncb+VfvONq99CLWHj3CujF3Tonm3VuHcvOsNUG6/4aIMisM65TIv67qG9b+zY0kSVzUtzUX9W3N9pxyPlx1kP1F1Ti9fmJtZvq2i2Pq8Pa0aaDAOzXWxls3D+bT9YeY+vZqbhvVgXvGdT6hlSaBQHDmExtlfCuKVCJuNcsikGoC5v+Sb3gfi0R26fKqrNxXLIKplmLUQ4H/Lvon+Jwhdx/cRmF8BxP/XuHh7xMadmNsJxXSmmJyCXyviizRPsmOIklM+Pci7hzXhtE9FNxaNXaTnTRHGtGmeB6dvYmFuwoanB8FFqJVXl24j1cW7MVmloPKDjYcLOW5H3YxuXcrHvlVNzomO9A0ned+3MXcrbl8evcIOiQ7GN4pib98s50FOwuQOdqXrgaTLKHIEj1bx/LkJb0YkJEQ8vMRCBqjyasFLxvQhndXZLGnoDLIVcWIxrTXTq/KDW+uYtnvJxBlCb0queVQeWDi7jejOo/K5lz714d9A9BkmY1p3Um5r3HXGiOGdEjk83tHcud76yip9uLyGjsg2UwyOjBlSDp/urhXyF5QLUHvNnE8feU5ER8nSRLXDk5ndJdkfv/5Fn76JZ/nr+1Hl9TwAmKBQHD20S0tBp/BMnKkEvFzjmT4fZqPA+UHKPeWY5JNJFoTaRfT7qwKtPYWVLJoVyEl1V4AkqKtTOyRSofk0D0BS51ew+2RyC51oLCqcYMDQTMz6iFoNwSWPAuZy0Az/l5reOpcK6NmVfPwsIblel5MJEoV5OpJKJJEXJSZj+8cTrFvP9Vx83h130Km71WwmRQUBbyqlyi1KyU5I3F5O1G355URNWoYo/rtmm3fbslh/o58Xp86kE/XH+ZwqZMv7htV2zw6LdbG61MHUVLtZfbag3y6/hBlTh+aphNtMzGuWwq3je5I5xTR6FvQNDR5MGU1KXx4xzAuf205ueXuBuUbEFp7rRPQXn+7JYdrBzfc56mG8gb08pHq7qvDzCwZ0bN1LEt/dy6r9pcwc8k+lu0twiTLyFLA6SjaauL20R25bmjGGd01vk18FO/dNpQPVx/k2hmruGdcJ24f3emUDBwFAkHLEmMzc3HfNny58XA9aTGELxF3WBWuGxHHyxte5pNdn6BqKrIUmLj5NB+p9lRu7X0rF3W6CLu55QrlmxO/qvHjL/m8sWgfu/MrUXW9dlHTYpJ5dt5O+rSN455xnZnYI7VBK+eGrtORyi4tQpXQ8rQfAdO+hPwd8MYoMGjyW0OfVIWLu5l4ZpmXnikNf3cmVGxmmSSHhTdv6cVvl93JzpKdeDUvuqSB5MOtAUemf15+QWm1F0eKA2f2bejeFKDxxfTG0PTAPO2Wd9YyrGMiH9053LDeKdFh4d7xXbh3fJfQn5NAcAI0i49lgsPCdw+O5v4PN7Jyf5Fhhipc7bXTq/L6on1cM8h4VVHXNPyFRfiyD+LefBjdKwH1f1SR3gBOdL4vSRIjOicxonMSTq+fkmovXr9GbJSZRLvlrOlFIEkSU4e3Z2zXFH772WZ+3J7Pv6/pF9bKqEAgOLu4fXRHvtuSExRMQTgScQ0leQ5Pb1kFgNdgBT67Mpvn1j3Hc+ue428j/8b5Hc9vqqGfElS6fdzyzlp25FYYSqlqFjbXHyjl4U82Mqh9AjOnDTZUfbSKM5Z5RSK7NMsSqbGiee8pQ0p3aqObRvjreBsDZ1TxmxHG9eEmNGRHIo+O6cbkftHcOX8aRa4ifFrD5i+SBChekH04OkzHeeBuypauDmlkAo0HXJoOv+RWRvpJCARNTrM1BYixmXnv9qFcP3MlK/cHu6tEor3OK3exftV2ursL8WVn480+dOS/2fgOHUKOjsaSno4tvTeY+gUdH6nuvqECyuPBbjGd9b02MpLsfHLncN5ZkcUVry3nkUndmDa8fcig0q9qrDtQSlGVJxCM2sz0TY8L6lAuEAhOf3q2jmVk52SW7y3C04iiIRgNR7sPMcXuNQyi6uLyB4rxn1j+BCWeEq7vcf0JjPjUwen1c9XrK8gqdjaqBjm6v8qazBKuf3Mls+8egdVUP6A6v1crftiWh/cEGqDKcqDlh+AUQZahzSA4vK7R3bokykzpbeblNV7OSQ3OTplsDj5/7Hp8ksqU76ZQ6CzE34CD8rFIko4ue7Akv0HZsu0kXfhoo4vp4ThH+lWNOVtzuXJgeD1CBYLmoNln+dmlxk4yEWmvXS7WvfwWqTFOLO3SMaen4xg+DHO7dCzt2iI7ApmOtqqG6e/z4Ri5XyQ3ALMicXn/tifwjgVGyLLE7aM7Mr57Cr/9dDPztuXx7NV9SU8MltsUVnr4aPUB3lmRhf9IHYWugyxJeFSNMV2SuWtsJ4Z2TDyraiAEgjOdV28YyJWvr2B/YVXYAZWj9feYHTvxNSJfOha36ub5dc/TxtGGcenjjne4pwyPfLKJA2EGUjV4/Bo78yp5/IutPH9tf3RdZ2N2GR+uOsgP23OPXFuPrwEqQJ+2cUKFcKox+hH48h7wVjW625/HWXl/i0GmyWTDMvoBMJmYs+dbcqpyDAOp0qWlFP1QhLfAi2JTiB0US9rVaSgOBUkCV1Yputr4Ynq46qXqI+olEUwJWpJmD6YaavoXifROt0Vhf+gRMkI0ADQrMjeNaM+MJfuDbirh3gBkSeKmke1DvzHBcdE5JZrP7hnJzCX7uezV5Tx2fneuG5JeGxT9b102f/pqG0CDk6mfdxawcn8xfdrE8vYtQ4hpwkyiQCBoOaIsCp/fO4K73lvPhoOluHwqegM+RhZFQjKXY05Yg2oQSDU2oQPwqB6eWvUU89vNP60XZQ4WO1m8u9DwehmqJsXt0/h2Sy7d0mL4elMO1V4/Nw7L4I8XTWD22oO8NH9PkBMahJZd2i0K94zr3HRvUtA0dLsAlOD7ZdYj9Q2i6jb0rYeuw8BbAJi1bVZtprcuRXOLKJxbSLs72hHdKxpfqY+c93PI+ncWHf/YEdkko1X5MEUrSIoKuvH8LxL10qFSF3sLKoXRlaDFaPZgymo2LmKMRHqnyBJRYTZTmzq8PTOX7Dd8LNQNQJZgYEYC7Vqwi/fZgCJL3Du+MxN6pPKbTzcxb1se/7qqL99uzuGFn3aHXJHWCchUNh0q57JXl/P1/aNEQCUQnCHYLSbev30o6w+UMnPJfhbtLsSiyFR5/DgsSm3gc8PQDHxx3/F1ZqDhaF3CmdABVHmrWJO3hmGthx07jNOGd1dkohlEnOFIpCBQS/XpukM8eWkvRnVOrpVf3zqqI99uzg3bmbcGm1lmTNdkJvVMPfE3J2haFBNc8Bx8+2BEDX0BMNthxP3gSGJb0TbynflBu6gulYKvCmh7e1ti+gYCG0uKhfT70tn92G7KV5STMDYBJVrBX+XHFL0Bf+Vww5erVS8pKpKlEN2TSkNOgCZFIrvUJYKpZqbIVcSnuz5lyeElVHorMUkmkqOSubLrlfyq/a8wGwTqZwvNHkxlJNrJLXcHbY9EeieBoRzMiLRYG789rzsv/LTbsNdVYzisJp65KnI7cMHx0b1VDF/eN4rXFu5j0vOL8KhaRDdtr1/jUKmL2/+7jtl3DT+tV5cFAsFRJElicIdEBndIpKjKw6p9xTw8exP/uOIckqItDO2YiCxpjJ39RVDhe7gTOgCn38msbbNOSjDlUzVcPhWHxdRkrqZev8bstdlB181ImusCFFd7GN0lud411GYOOPNeM2Ml2SXOsGSXUWaZQe0Tefn6AeJ6fKrS9xooO0Dpktcp9lrQkIiTqkmlzLBRMxAIpHpeAuf+EYCVOSvxqsH1ic49TjSfRuyg+lktxaYQ0zeGqu1VJIxNwN7FjmSScGX+hDnZOJiqVS/pXmxJ81HsB3DnXIfqDM546nqgr5mgedhTuoeXNrzEypyVSJKERz3a8mBf+T62Fm3lb6v+xrXdruXufndH5JZa7vLx+fpDfLslh9JqL5IkkWA3c1n/tlw5sO1ps1De7MHUbaM7si2n/EhDtvqEK72zW00M7ZAY9mveObYTJdVe3l2RFVZAJUvgsJj48I5htE8SGu+TiVmReWhiFz5ac4CqCuOeJI1JVbx+jW2Hy1l3oJQhEZwjAoHg9CA52srILsnE2ExcPuBoPesvxbvQDWp6wp3Q1bAur/GC/BMht9zF+ysP8OHqg1S4fZhkCb+q0ybexp1jOnHloHYnZHhUWOXBwPwwIokUQJXHT4XbT1xU/bEkOCx888AoHv9iK3O35SFJAWngsdgtCroON41oz+8m9xAtME5RfKrGj9vzeWPzMHY6e2DR3EiAD4VEKrjL9D1XKUuIlY5krUw2QIfh98GEJ6iJtkrcJYbSWrVKxRRtQlKCv39TnAnXgcDzKnaFtCvSyP/fRhInrTRcTK9VL+1ah6NTPrK5kqj0d3HnXIO/sm+955YkiLEdM53VVNg7H7bMhorcgCW8PQl6XAS9rwSLUCCFw4rDK3hk0SO4/W7D6y0EFqUAPtjxAT9n/8w7k98hOSq50efNKXPx3A+7mLM1F1mS6s3VM4EduZX8c+4OLunbhsfO737KO4M2ezA1sUcqZkUGjIOaUNK7KLPMnWM6Rmwn/vsLetAh2cE/5+zAp2qGvaMUGcyyTLdWMbx83QBRLNtCbDhYSqXb2A0oHKmKy6cyc8l+EUwJBGcoTq8f+zFS7wpvBRLB94VwJ3Q1+DQffs2PSW6622G508ejszexbF8R6OA9YqRTk0E6XObmX/N28c+5O7l+aAZPXNQz0HA+QirdPsPAJdLeiiZZptLtCwqmICC7fPG6ATx5pAHqOyuyKKr0ouk6iiyRnmjnrrGduKx/m7PeufZUZuGuAh7+eCOqrh9Z3JbwEVX7eC7JPOu/jmf81/OQ9Vvui1qINPxuGHQrxKTVey5FMj6vauR7uqoH/f785X5M0UfPj+QLkpGi0iieZ7yYflS9NAN7egLRfaKRFC++orcoX9uB+NG/rn0ur1+jS+qRBrxeJ6x6HVa9Bn53sNlG5mKY8xj0vxHG/AZiW0f6UZ41bCzYyMMLH8atBqvLjPBqXrIrs5nyzVQeO+cNUqPj6JwcTZy9/nVle045N7y5mkq3z3AxCKgNrr7ceJgFOwv45K7hdEs7dWWczX7lMykyd43pxCs/741YdgcBuUc4DXuNmDIknasGtmX+jgLeWLyPbYfL0fRAbG23KFzary23j+4gdLYtzIwl+w3PjXClKroOS3YXUlTlITnauDeGQCA4fXF51aB+SIqkGAZTkUzoaqhp7tsU5Fe4ufK15RRUehqVLddc82avzWZXXgXv3jY0yKI8FHazybBeKtLeiqquhwyEEhwW7hnfhXvGd0HXdfyafmShVHCq89n6Qzzx1VbDrGJdXATun9PVqznc9WH+Ma6voVwzOSoZs2wOktjWyPcq1lcQNzSudrvqVqncUkna1fWDstiBPbCkTau3rWrrfHJ+mI6/LA/JGoWtrZ38r/LJnpGNYlOwdbCRfIHjSJYkMLZB7RNoHRcFVYXw30ugNAsMzDEA8FYH/rv+Xdj2Gdz8LbQS5R3H4lE93L/gfsNAqjFzH1VXKXDm8djCP0PBjXhVjUk907hzbCf6p8eTVVTNdTNXNbiAfix+Tae02ss1b6xkzsNjaBsfFfqgFuCkLCPdPa4zK/cXsyazJKL+ITazzMxpg4m3W477tU2KzOQ+rZh8pN+F26diVmQhQziFWJdVaujYFYlUxWKS2XqonHN7iKJngeBMw2kQTCVYEwxtmSOd0EWZoposmKry+JkyYyV5FR7D5sNGuHwqGw+W8cCHG5kxbVBYKgzN7ca1aTOsWoPP0x6OyRRE2ltRAmKPlUk1tr8kYTbI/AlOPZbuKQwrkKqLy6/z5aZcWsXbeWhi16DHJ2RMYPqm6UHbFbtC6uWp5HyQg2yT65m/mBPNxI+Mr91XV634KwbUO75izReUr/6C5Isexdb+HFTnIcqX/hHdJ9PztZ61xjG6asZ1aD+qszMOi8Ld4zqDpxJmnQdl2dBIA+FaNB+4SuGdC+CuxZAk3Cfr8s3eb6j2VQdtD8vcR1bR7Vup8peD6mDutlx+3llAz9bRlFT7qPYEX7cbK+fQgSq3n7veW8f3D405Ce8+ck5KMKXIEm/eNJj7PtzAyn3FITNUsgRWk8KrNw5gdNfGdZeRYgvTFVBw8nB6jVcoIpGqaJpOhTuMC6hAIDjtcPlU7Ob6t6vO8Z2JscQE2TNHMqFTJIVJ7Sc12ThfX7SX3HK3YSDV2GTB7ddYvq+IhbsKmNgzLehYzevFtWkTzjVrca5ejWv7dmxdu2IfNozxrW0syK8vl4nE4EmR4YoBbY9LZig4tdF1nd9/tqXBQKqxc9LlU3l14V5uGJYRpPjIiM2gZ2JPNhVuCnrOlAtTUBwKebPz8BZ4kaNkYgfGkn53OnI9d2cZf2XP2v/TPE7Kln1E0gUPE9VpEEgeYvt+S2zfdkHGMcheLEmLcTs7k+CwMKZLMnx2C5QfrhdIdXixEqcPMh+OxmEJBP9vbfDywRYfi245UtbhqYL3LoeHNwcaGwso95Tzz9X/RNPrnzeRmPuAhDluHb6ScWh64Bq++VA5mhbcvS6ccg5V19lXWMX2nHJ6t4njVOOkCZxtZoW3bhrM7HXZvL5oH0VVnqAeIjazjK7DpF5pPDyx6ymtjxQ0HSZZBoIv9pFIVSRJwiImAwLBGYmRzE+SJG7pfQuvbHgFl1o/oAp3QmeWzdzU66YmGaNP1Xh/5QFD9UU4kwWnV2XGkv1M7JmG7vXi2rqV6tWrca5Zi3vLFiydO+MYNpSku+4kasBAlOjAZPC+g6Use3N10CJluAZPZlnmjjEdm+QzEJxarNpfQpnLeJExnHNSkuCTNQd5YEJwduq2nlP5fcEmXAYJysRxiSSOa7iGWddMeEtGAEd/057DOwJKlO5DQXZhz3gbxZYLBBvHSBIojj3YzBLVHpX7Z85lesEcFC3YYVDV4aXVXh4f01AJgA6uEti/ELpMbHDMZws+1cdNc2/CpwefN5GY+0iyD0vCSnwlR5uiqwYxfSTOo26fxnPzdvHubUNP5C02Cye1WlSWJa4fmsF1Q9LZcLCMT9YcJLvUhdunEh9lZnjnJKYMTifBcfyyPsHpR6LDQpVB2jcSqYqOTkqMqJcSCM4Yqgpg3TuwbwGDy4ro6gE+7AwDb4Zuk0ExcVmXy3hpw0uGh4ea0EFghb17YvcmGe5Pv+SjGuiVI5ksbMoqZtUdD5KwYQWWDh2wDxtG4i03Yx88GCU62vB1+6fHk54Yxd6CqqBi7lAGTyZZolebWFE3fIYyc8k+Q8vwcM9Jt09j1vIs7h3fJag0YlzmOoa6PayymvEcyejs+s0uNK9G9393R7YGtpUsLqFsRRmd/tDpyIvL6L54vMX17+mquxjZHoslaR2WpMXI5orax4yMYyQkvn5oCB0Sktj32ROoeTpGS66PjbTw7HIP9w2xEG9rQJrqrYLlL53dwZTPBdu/Yk7m9xyqyDTcpWJ94Dv55d5fgmql6n5HZSvLAjVVuV4k07SgZuF1idR5dNHuQj5cfYAbh7U/zjfaPLSI9Y4kSQxqn8Cg9gmhdxac8UwZks4rC/bgPmZFNxKpitWkMCBDnE8CwWlP/i/w899g74LAErTfTQwQA7BnHxxYAYoZht5N7OhH+cPQP/DMmmfCdpyqIcoUxdOjn26yYX+x4ZBhC5BIJgu6rrN+1CXc88I/UGJjQ+4Pgfvp2zcP4aJXllLp8jdgXmyMSZZ46rLgCY7gzGD53mLD8yGSc9LtU9lbUEX3VnUCbtWHvGYmz3vKuCctlW1WC+4aiZwGRT8WkXpJcP2yWdOI1/yUqQreuE1o/njQTUhKNVFtfqHYVYY15ZuwjGNMskzrOBsWk0zP7P8Bxhm4wW0Uxncw8e8VHv4+oRF77YOrAgs40WdZ3XXpAVj5Kmz6AIC3U2LxWoJdPYvmFlG+thx06PlKT/wV/nq1UjXfUdG8Igq/L6TNzW2I7hNL9d5/GDYLryFS51GAp779Basic/VxmtM1B8LHVNDiXD80g5cW7DF8LBypikXzcUOSjlzH3ed0R9f12uJPh9khGmAKzg72zIf/TQuskjYUFtRYHS9/EXZ9z1U3fUOpp5QZm2eEHVBFmaKYPmF6k2WlAAoqjfvkRTJZ8CFT1Soj7ECqhvREO5/fM5IpM1dR4fLhD9P8QtV1rp2xihnTBjGma0pEryk4tSl3V0PMauwJy5DNZSCpoJtRPSm49sWEfU4qskSZ8xj53K65oPmx6vBmXgHPJcbzeUw0EgHL86K5RSRNSEJxBJ5fBiyazkiXm/tKy7mpjYzN/k29p7QkqBEZx2i6hsPsAE2D6qJG38NT51oZNauah4c1onoyWaE8+5QLpjRNZ8meQuZszaWgwoMOpMRYmdy7Fef2SD0xM7W9C2D2NFC9oPnYbrGQZ+AoWlMr1Xpqa3Lez6FycyVxQ+Nqa6VKF5dSuaWSlEtTKPgsUFMVNzgOXbMgKaYGm4VDw+UcjdXzefwaT3y9jf4Z8adMVl0EU4IWJ9FhYVKPVH78Jd9wEhBKqiKZzExcMpvMn2aR+utHcYwefVoGH5qusTp3Ne9se4fVeauRkUEKbB+cNpjb+tzGiDYjmtTGWSA4ZchaBv+beiSQCgO/Gwp2wrsXcccdC2jjaMMza57Bo3pqm0jWRZZkLLKFttFt+dfYfzVpIAUYWpRD5Dblfi1817W6dE2LYd4jY/jtp5tZsrvxyWUNPlXHp6rc+d463pg6iPHdT62JpCByvKqXF9a/wBd7vsCSqiIpdQMhPyb7QWzt3GiuMkyxi/BXjCPUImTQ/XT9O7WLGibgDyVlPFBaTidNo226BVd3OyVzC0m9Ko0oTcPuV5lzKIc0VUUDYjTtaCbrCJEYxwAMSBsQuBf6qgPGEVrDxmZ9UhUu7mbimWVeeqY0cv/0Bl83Wgqn1897K7J4e1kWTq8/qFfq3K25WM0Kt47qwC0jOxATafPvfT/DJzfWs5BfYrfhMZg71dRKxY+Mx1/pr/cd2bvZKfymEHOiGUuSpV5NleYOfT0xKucIp57P59d4c2km/7qqb2NPf9IQwZTglODvV5zD+oOlFFZ6GmziZoTNLPP0FefQf8CFVP70E/lP/xNTSgqpv/k1Uf36Nd+Am5iVOSt5YvkTVHmraieCGkdtb9bkrWFb0TbsZjtPjXyKMe1OTXtQgeC48FbDx9fVC6TCcuLSfFCyH354nAsveZHzO5zP0sNLmbVtFlsKtwR20TUsioWJGRO5qfdN9E5qHllbQgMtPCKp/VRkiUTH8dd+llb7WJNZEvFxbp/GfR9uYO7DY2ifJJrXn65Ueau488c72VO2B4/qOdYxvxZHNzOSScJz+DMc3XNx50whkD8KRtV0Eo5pukpFTtB+/8/eecdHUef//zkz27KbSgohhRJqKAHpLYCoiAVRPEHBjv1rL1f8nXfe6XlNPc+znA1RUc+KBXsB6S1I7wkBAoSQXrbPzO+PNSGbnc3uhKCUeT4ePHxkd8pn3Z3P5/Nur3ecquJQFP5eXkHqCIFxc8r5qo+bzyr9zPP76SgHjAERuKamjmeSEkIMqmiFY+wmO9f3/ynN32wHNbID4k8TbAx+oZ77RrXyfFnbFuVQFJVlheUs2hHodykKkB4fw/kDOjEgS7/yXFmdmyteXMmBKldI+UMjDV6ZBq/MM9/v5t01+3nn5lFkRNuDqeYAvHNlkCHV9ak6yuV6ujwRj2gNzLeN9W4dxndoaoTe8jtCAMkh0e3X3ajbUNd0nCpb8DYTnwhHy3IOS0Yvqpe+SdyQKXgObm+q6WuucC59AAAgAElEQVQZ3ZJV+Hj9AR66sC+x1l/elPnlR2BgQCA69d7No7nsheVUNnhbbXbZiM0s8pvJfZg2OAuA+EmTiJs4ker58ym58y5i8vJIvedurDk5x3v4x8SCwgU8vOJhPLJ2mlAjTr8Tp9/JPYvu4cERDzKt57SfaYQGBseZje8F0nVaEFmJi0CEasP/YNIjSNY4JmRPYEL2BFRVxS27MQkmzJJOr20bmNwvnYLiSpy+ttd+miWR/GNoB/Lswt34/NpzZ2tpMwBev8KLi4v4yyVGA9OTEZ/i45Zvb2FH1Y6QZrotkewSHS/pyKF5++l01QosHTvir83HXbwh5DcZazXRPbWF+IkcqprXnIE/RYIeDxMJuri+nv8kaRsZ0QjHWE1WxmSMCfwhCJCQDdX7Wj2nRweRGf3MPL3ay4C00DGpsgchSZ+oQb3Hz9ur9vHSkiIaPMHRI1GAucuLyUy0ccuEHkwdlBFVk+sap49Lnl1Gaa1bU/2uJR6/wsFqN1OfXcYXd+WHyNhrsvpFkEN/I0qYereWjdCbf0clL5Wg+lUkhxR0XED6vm/ksRBczuE9shdkH97DRSSMvrzV80RB4LONB5kxrHNU9zmeGMaUwQlD52Q7X9w1joc+2sy32w4jCIT0x2jsQdYxwcofL+wX0qRXMJlIuuwyEqZMoWrePPbOupK4s88i5fbbMXcM7d/yS7P8wHL+tOJPEQ2p5nhkD4+teoxkWzLjsyN7fgwMTmhUNVD/pNEgMiolLghsqDa+C8NmN3tJIMYUpaf2GPGXlzN25afI7qyAOEYLopUpz06KoX9m23qo1Dh9fLWlVFNRMJq0Gb+i8uG6A/y/C3KxW4ytwcnG29veZkdlqCFVtaQqoKxW5g1SYEs+tyOq1JeyD/bgq3gWwTJH8zfZuYOd8noPafHNxBuskX+jrUWCEhSV/6tp4LkOSbg1Gm+3hk2y8fCoh5Gap8yOuh2+fRh8rafp/WG8lTc2hhoRMgLf+Qez/JuD/GqIGNUzeLDaxYwXVnCk3qPZx6uxt9LuIw089NFm3lmzj1evGx4xinL72+soq29AkWoQzR5UxYLqj6Nhy/KwzhBZValq8HLj62uZf9uY1gfu98LaOZoGcY8JPdn65e6gejf4qRG6KLDztzuRa2VEm4its43kc5OD6tkaG6bXFNRh63whaOoratNYzlG/ZSFVC1+h4/Q/RTzH6ZUpLj8xUjONGdPghKKDw8KzswZT1eDlnTX7eGv1PsrqPHh9MgIqsTYLPdPimDWiM6N7JIe9jmizkXzDDSRedhkVL7/Mnoumkjj9MpJvuAEpIfrNyqaSGt4v2M/+Khcev0IHe0DC/+JBmTiOMbQsKzK/WfIbzaL5cAtg4wTnkT38bunv+GHGD5jF4+91NzA4bpRugvrDmm9FrcTlc8LK54OMqZbsr3Syp7wBp9eP3WKiW4qD7A52zWPr3D4+XHeAr7eWUtXgQxIFUuOs/GpIFuf07djkYXbv3Enla69R9823xJ9/Hhf1O4P5O2vaVPtpt0jcMr47qqri8St4/ApxVhNilAXm7xfsR9Sod9AjzS4I8OmGE8PTaxA9iqowd8vckLWk/ItyjnxxhKwbsoJqkPb8cy/pV/4DW+cMOl3Teiua9furmfTUYt66YSR9M34SRuk+Acq2tBqhihQJurbBTWnemcw/tDSk8XY4bJKNe4fey8TOE4PfGHgFfPOHkOOL7w5O28tOEHH/PlTcRTLHkHfxQ2w5aObmNwqIjzHzqyFZXDwog2SNSM+ROg8XPbOUKqdPs0F3S1w+mY0lNVz+4grev2U0NrO2kbF07xYK6l/C2n0dVgRQBUCl/OtSqhZWkHz+tVgzzwnrDNl2qJbtpbX0SW9FwGbnF5ppkSoCNR3Oxt67lPIvy+l46VHnc9XiKhBBrpfJuimL2H6xVK+q5uDc4Ho2yS6RdklHDr1RSoezJWzd3AiiierFr1O/+XtUn1szKt4cvTWm4fqo/dwYxpTBCUt5vZeKei+SIKAgAAI1Lj9r91ax7VAtv/94C1cMz+b/JvQI25tMSkgg7b77SLrySsqfeZbC884n+frrSLrySkSb9uZMUVQ+Wn+A5xYVcqDKhdevBHl7v9texiMLtnLJGVncNqF72A1ZJJYcWIJXYzEKtwA2SpCKpsDCJCsyC/ctZFLXSW26v4HBCUFNCYjhl6KolLgA6ktDXvLLCt9tL+O/iwrZeqgWiynQGF4QAmltfTvFc8uE7pzVJw2TJFJS5eSpb3fx6YaDiIIQ0gh31Z4KJEHg8kyRaSvfQ9q5naRZM+n+1ZeYkpK4t8bF108toUbnAm8WBdLirCzedYQH52/CL6uIYmCD1CstllsmdOe8/p3CbsIA1pfUhIwX9MlgO72BTd+MYbqGb/ALs+rQqib110YaFdgyZ2cSlxcwKiypFrJuyWHn/UXUFWwnNq9rxGv7FZVqp4/pL6zgk9vHkJMai3/IbITlz0aMO4SLBAEICdn89swn6LTlNZ5Z/wyiIIY1quymwBr75zF/5tyu54YeYIsP9J/78fXoBWwaEc2Q0pv0vqO5p5/AXWf1ZOWeCt5fW8JT3+5kVE4ylw3NZkLvVMySiKqqXPvqaqo1DKlICnS7y+r5/Uebefyy4HruWm8t9y26j9WH1iEm+BGEo8aO7JQ58lEpGddnkTBsNYqnGFfJNZrOEJ+s8vKSPSHXD+LIzkCNagtcWInxdSDlgj7sfXwNyecEnNWqrFI2v4ysG7KQXTKHPzzM/v/ub6pnS78sPaieLeXcjvgbLqdmxbuUL3gCBAHV7ydhzEzih16kaQg2R0+NKRBaz/cLYRhTBiccxeUNzHhhBZXO8LVTjbnJry0v5tMNB3n35lGtFk6bO3ak0yN/psN113LkqX9TOfk8Uv7vNhIvuQTBdPQxcPtkbn9rHct2V2huTCCw4QB4b+1+PtlwgFevHc7wbq3neGsxZ/OcENWxcAtgowRpzfKapg7jTr+TOZvnGMaUwcmNz9lqAXnUSlz+YMdE0ZF6Zr68ijq3r6n/k6dFMfeP+6u59531xMWY+cOFffntBxtp8PgJV7LZeJ1Xd/n5Kvsi3vnHE6QkH/V+d0qI4c0bRnD5iytp8PoJI/AXhPmnnjqHalzsr3I1bdB+qtdnx+F6fj9/M7+fv5l7J/Xm+jFdNdVKQ+Srf0JvH5eqMNcxOHGZv2t+yFrSqMDWqKzWiK/qUmJyFuIq3khs3uSg91ozBhq8fq58eRX3n9ub/3xfyFOm/uR51wXpAEYbCcLsgLH3IAgC1/a/lkt7XconhZ8wd/NcqjxVmEQTqIE6sM7xnbm+//VM6joJq9RKPdC5f4HDm+FAQaCOMhpEEziS4cr3Ax4WQBQFRndPYXT3FOrcPr7YVMqLiwv53YebuHhQBv0z4yk60hASfY4mldbtU/h0w0H+3/m5TQ7gSnclMz+bSZmzDAUfLR9t5+7A95gwNA5B9CFaD2Hv9m+cxf+H6gved8iKyoINB3n04v5NjhdVVXF6ZaqcXqqdPhIPHSBLo+1Eg2pDwEqc7ULiBu7kyGdHsGZYUdxK0++osV4qHKoi4a0YT2y/s4ntdzaKp4GSZ68hZcr9QYZRa1LpempMHRaJXh0NaXQDgxBKa9xMe2451S5vVKp+PlnlSJ2HS55bzpd35QfndWtgzckh6+l/49qwgbInnqTy1bmk3n0Xceecg6LCTW8UsLqoIqyCTnP8iorfI3PNnNW8fdNIBmUnRjynEZff1aQ21pxwC6Bkk4jLi6N+S32TMQWwo2oHtd5a4i36+tIYGJww2BIggtx/NEpcsmQDRUUSBbaX1vKr51dEZdA0eGWcXpnb3lwX9ZB9ookDXoFfvbyGz+7MJyHmqHe0f2YCH98+hqtfWU210xsiadyIJAqYRAGfrESc6xqv8fhXO9hT3sAjU/uFGFSOMHVOetNmTgRlLAN9lDpDo7JyvdykrNaI6rfjr8tDcmzCW7o76PhIxoCqwqFaN88vKuTRi/uTZ38S4dXz9EeCBAliU6HfJU0vxVnimJU7i5l9ZnLYeZhaby0iIom2RFJiohRkkcxw5Yfw7tWBNgsaNZhBmO0QnwHXfgYO7XvE2cxMH5bN9GHZFJc38H5BCb/9YFPI/kBvKu07a/Zxy4QeeGQPN3x1A4cbDuMPUzvW8nsURAUEJ/YuL9BQdBcowZkxPlnlihdX4PIpVDm9TWnKSXYziXYL1/v9TCNUu1H56RV/bR4pk/PY87clpExOQZXVkN+RFqoioPrj8VaOa3pNT1S8OdHWmKrA5P7pTX83+BrYWrGVOm8dJtFEckwyuR1yf5Z2MsasaXBCcePra6lx+0I2F615zBQVal0+bnx9LR/fHjksDBAzcCCdX5tLw9KllD3xJBWvvMKH59/Cmj1uTUOqtfu7fDJXv7KKlQ+eFXXhdo2nBrNoRpaDN1paC2AjpgQTrr3BC5dZNFPjrjGMKYOTl9Q+4G9dgCVS/QXAZn8mV/7pa3I7xbPxQLVmUXg4wtkyrT33fkWlrNbDA+9t4MWrhwad1z01liW/PpNlheX894dC1hRXYZFEBCGgt+GTFcb3SmXp7nI8OurvXT6ZDwpKSI21cNfZvYLey0l1YBKh5fSlJ23GYhLpmmJIo59saKWLt1RgA/BWDwVU5IZKRPvRNSNaY0BVA3XNY3qkACkw7SX44MYgie1WEU1gSwwYMOZQcRhBEEh3pJPuSNc4OQrMNrjif7Dj84CoTemmQP+pJlEOIWBEOVJg7D2QNwMs0aXpd01xcNP4HF5cUhTynh6jwe1TeGVZMbdM6MEnuz9hf93+EEOqec20YBJQnAr+Oj+muMD+QhBUkBqwdFiKtzw4M8UsCVw+rDMDshJJcphJsluC04M3HoIF8482P/+Jo3OghN9zB/FDdlHxzWFMSaaQ31FLVEVElR04994MylGHtt6oeHMi1ZiaJYEZw7KxmSV2Ve3i9S2v80XxF0E15IqqYDfZubrf1UzrOY2EKIRT2ophTBmcMGw5WMPusrqQPORolah2HK6LXHzZDEEQiM3PxzFmDBWffs6cJZW4TKGRrWjv/9GPB5g5IjppVZ8cxlutsQA24q/xY4oNfmQFhEA/KgODk5XEbMgeFvAmt0Jr9RdYYhl46R/5IetMfj9/E95wvVkiyIM3J5rn3isrLNp5hMO1bjq2iIqLokB+z1Tye6ZSUe+htNaNyysTazORlWTnptfX4g6TShzJefPcokKmDc4iu4Mdf1UVdV99xegvFvFy2uQQNUE9aTMCcOlPrSYMTh4SraFZEY3KarUFtSQMD2wifVWjUTwyrqICEsdd3XSsHmNg/f7qo7/33Clw+Tx456rAm62p6VliIS4drvk0EBE6Xogi5F4Y+Fe+G7Z9EqinlP3gSIXuZ0L2CELy6aLgQJULiySGzC96jYYjdR7OemIhZQnPgrl10RD3ITeFfyik8JFCej7Ws6lmWhD9WDqswFt+Fs1V8wRBYGyvVDLD9ZzqcyF8elfIy2Jzl5JiwzHwYWpW3YJokxBMIrXrakkYFmyMqKoIqojsysZ9YBaqHCyhrzcqrgdBELhqZCb3LbqPH0p+wK/4kVU5RBnZ6Xfy7PpneXb9szw08iGm9pjaruNoxDCmDE4YXlmyB2+LYgU94XOfrPLKkj38s7XiSw0EUWR1zhAo2AAtUnKivb/TK/PCD0VcMbxzUPpNjdNHYXk9RUcaKDry03/L69lbWYmluzek8bzWAgggu+UgCdKmz6z4jKiUwcnPmLtRDvyI2Cw1J+r6CwCTFXqeQ5wqsGR3uWbaXDTGUSO60naAeSv3ct+k3mE/XnKsNUgVrKTKScHeqjaPU1FVXpr3PTds+RTXuh+Jzc9n4DUzGLDdTMG+mpBrRpM2IwCjcpJDjEKDE58xmWNYd3gdLvlohEiyS6RdnMbBeQcRbSKO3Fh8FS4qvv4XprgUYvsdVcTTYwxYTAGhlqbfSY+z4d5tsP5tWP40eGoBIVAHKYqBfkadBsHYu6HnJGjnTXWrpPSA/Hvb7XINHr+mDabXaJBEgRvPlnhikxN3sy2HVs10THYMaZekUTa/jNK3S0mfno4gCdRvrad+6xESx27DX9e/6RoqKimxrYj1WOwwaBYUzG0WsYO/3H4R/5X70miKmOIy6HzfJwimGuo3PMfBN1aBIBHbLxFBUqnfWkfdBgcJY+9C9aZq3qoxKu7buYTMPoPwYKIOO2qYBtHNiej4Evz8dsVt7KndFbG1TOP7j658lBpPDVf3u7rV49uCYUwZnBB4/DKfbToUEpXS4zGTFZVPNhzkL5cMwGLSlyP7ypI9mrUNeu5fWuvmT59uweVVKPrJgHL7ZHJSY8lJdZCTEssFeZ3ISXXQLcXBzM9fo7CmMOgaLRfA5mp+zSVIG+nk6KTplTQwOJlY4Mylny+RzoIbSdWO1oTFbIcJvwVR4uuNh1A0LBQ9xhHoe+49foU3V+1r1ZhqyRsr9qJqFHNFO06frPL+AT93nD+FrCefRHQEUvNuzyzjtnnrNMVzIqXN2MwSt07oHvVnMDhxmNpjKv9e9++Q11PPT0VySJS+UxpIGTM/gL3nSBKnPIBgOhrB1GsMNAqxNBGTCKNuhZG3QMlaqD0QqKWyxUNaX+jQ7Zg/44mA3WLSrMHUq0CnqCql/h/xtJCyD1cznXZRGvWb6qleWU3Vkiokm4Stq420KWmY4jY1GVMmUWDaGVlYTRG+w5G3wvp5QcbULNN3PC9fFHKo6k/A0e93qHzP4fc/ouSFrT85Y3oSP2pGWEPKjpup9pVI42N56dsnuN8ay6TuJiyiyj925zCvOAl5wj2a50bjUDJ1fJudVTuQiV4wxy27efrHp8mKzWJil4mRT9CBYUwZnBBUNYSq2ID+8LkgBFStIglRtGRvpXaxqp77+xWV4nInZ/ftyNQzMuieGktanFVTeQtg9oDZPLry0RAVppYLYKMEafbN2UESpHaTnev7Xx/2+gYGJzr1Hj8Pf7KFgr1VPHvFfKSPzwdXFURrUJnt0P9SGHYjAEt2HTlmpwi0TQFPUdSoe0J9vOFgSBRe7zhFWwy7+g0h1XG0xunM3mlcMTybt1fvD6tGqkWMWeKmcTmMyAnfu8/gxCXeEs/Znc/mi+IvUFooY3YY36FJga1u21/QaqSqyxhQIdYWZusoCIGUXU5Nbf3MxBi8cmgKsZ5UWoAku4UKdzlqi2rN1mqm7T3tCBaBbg8EG6b+hqO1TyZR4PqxURiuyd3homfhk/9rEhBJE2rIFzfxvXKGZuQott/EoGhmS6w/ObA9fpk7pPncZvoEGYHYMT5y42z8c0kD189XiLMIDO60kZfyHWRZHuBW390UqplN14nGoSRaSxEd25EJTfuOpkfnX1f/lTM7n9mueyfDmDI4IXB6/UgaP2zd4XNBaJIu14MnTLG6nvtbJJHJ/dO5fHh0DS8ndZ3Eoysf1Xyv+QIYDkVVOK/beVHdy8DgROPHfVXc/c56RuUks+COsYEm2Df/AHMvgPqy1usvBDGQ2jf0Ojjn0ab6h4p67XQPvcaR3nlHlBW2nD2JGElAMJsj/qsxjwMhdPnVN06VSg0Z899f0BdZUXl3bUlUBlWMWeLa0V25++yeUdzT4ETl1kG3snD/whDnXHMEcxWqL1S5To8x4JUVurSxt+LJToLdTH7PFL7fVhYiWhOtAp3VJHLNqK5UElojqrdmujkWSWBwlyR6pMWGPSaIAZcGIlOf3v2TlLzKb03/Y4W3H070OaPNkkD31FgOVtbzd55lklRAjHB0bpqVZ2ZWXst+UCqyepCPLQ9xpfdB1qs9gOgcSuakpSCEKvdE26Oz1lvLmtI1DO80XNfnbA3DmDI4IYi1mUL6NoD+8LlfUcN7zVrBZpGo05DV0nN/SWzFY6eBVbLy62G/5m+r/xbSuT7ieCUb9wy5B7v59FzUDH55FEVldXEl+yqduLwyDquJHmmxDMxKaNXjJysqzy3czWsrinn04v5M7t/p6JsJWXDbStj8ASz9F9QeDBSOK15ADKh1KTL0ngyj74SsYBU9S5j0Fr3Gkd55R5VEes97Dfx+VJ8v+J/XF/Iay1W0dGP0jFMFzZorURT409T+DO+WzFPf7aSkMrTxuCSA2SSSkxLLPef04py+HUMvZHBS0SW+C8+c9Qy3fXtb2PXE0mEJnrLzQQ1tMRBtXd3YHilB9X+nGzeP686KwgpNp22kVNpGZo7ozDu7UxEQgqJTemumAVR/HCZRIDXOxvOzhuj7MAMvD6ipLv4H7PqWnqYKnlKf4VbfPcgRWzIHsEgCafE23rxhBO4FvyVhawF2ofUapkYkQSUWN69b/soU71/Yq6ZHdiiJHswJGwKKhs3Q06PT5XcxZ/Mcw5gyOPXoYLdglsSQppp6w+cWk0iSvZXiyzD0SHVwpC50AtBzf1kJSCLr4dJel3Kw/iCvb309aoPKJtm4vM/lzMydqeteBgbtQbXTy7tr9/Pykj00ePwoaqAGQPopxS01zsot47szdVBGSKuAkion97yzHrMksuCOfNITNDyg5hg448rAv5IC2L8SnFVgtoIjLaBG5dBOR8tItCEKoUaGXuNI77yTnWTHkhG9Qlnshm9xacw3esYpIgT1t2rJBXmduCCvE5sP1PDa8mJ2ldVT7/ETazWR2ymOa0Z3jVr51ODkYFDqEK7r9jRPfbcevzcRRTEhiF5EaymW5CWY4n/EU3ZB2PMjGQMxlkA66OnMsK5JZCTGsKe8IaTGOxJWk8g5fTuSGmflzOwzeXXzXLzK0XlAb820KlsQGwbRLcXBWzeOJMEefj4IS8YguPwtqD9C2fI3eXRpOkLYZhGhJMRY+OyOsSQ498KuN6GZIdX1qTqcPthzVywOS2B9eHmdl3kbfSy69mh6sgM3j5he5Wrf7yI6lCTrQVBD0xD19OhUUVl/ZH3UnzEaDGPK4ITAJInMGtGZV5eFKvpFGz63SCJXjezStKnTw+yxOWw8UBNaWKvj/llJMeR20r85uWPwHaTZ03h87eMICEGKTM2JMcWgqAp3Dr6Tq/pepfs+BgbHyvLCcm58bS2KquIKkxq7t8LJIwu28o8vt/PWjSObnomP1x/gz59u5aZxOdyYnxNdfVHWkMC/KJk6KJM3Vu4LkRzXaxxB9M+9vQ0bzAm9U/mwoISWZVN6060Gd44sPtM/M0G3wqnByYWsqDzz/S5eWboHWVXxeo6mmquKHdmfgNvVDUH0YLLvQ3Z2RVX1bf9MokB2kp3h3VpPPz/VEQSB168fzvlPL6HWFdoTM+x5BFT8/jy1Hwu3l/HQR4fxJCcgWMqCjou2ZhpAxMz94y5mxrCuxFiOTSWx3pzEjI0DOeBzoqdQos7t45WlxdzreR6U0OweWYV/r/LyYH74aKYkqAwXt5NOBQcjOZQk7f2R3h6drmh7o0WJYUwZnDBcNaoLc5cXo9VCM5rwuSAErtEWzuyThkUSaQgzjUS6v8Miccv4tithzegzgwtyLuCTwk+Yu2UuVe4qTGLg8fQrfhKsCVzb71qm9phKnCUuwtUMDNqfRTvKuGVeQVTNcJ1eGadX5tLnl/PqtcP435r9bCip5rXrh9M/8/g1TuyfmUB2Ugy7yupD3ovWOGpONPOOqsLFZ2S2ekxLZo/txqcbDiJr/L+MZpySKDBlYAZxtjZ4og1OKdw+mdlz17BuX1VYBwcIqIoVVbFiIQGzSWjl2FAkIRAFfX32cEPwCMhIjGH+bWOY8cIKql2+sH3tGrGZRbqnxpKRGMNl/13BgSoXbr+CSR2PLf0jBDFYSCGammmzaOGmQddy7cD2UeB84usdHKh2hzh4oHWZcrdf4cXFhVxgWUxvDWPqgdEW/rHMw23DLCTaWv/tXGn6lseZ0apDKeXCMYT0lEF/vZkktK9Ev2FMGZwwZCXZObN3Ggt3lIWk+0XCahI5K7cjnRLCNKqLgCQK3D6xB49/tVOXChYEHmubWeKCvE4Rj22NWEssM3NnckWfKzhQf4AaT6BfTLw1nqzYLGMRM/jF2F1Wz21vrovKkGqO0ytz+UsrueSMTBbcMTYk7a+9qXH56OCwIKDlkom+piFaYswSV43qEhDP0EGf9HhyUmLZeqhW8/1I4zRLArOjUe0yOKWRFZWb3ihg7d6qqNdMt19FEFQyE21UOX0RBZtsZpGUWCvv3DzK6EHWjG4pDr6+ZxyvLN3D6yv24leUoMwWgUDUOj7GzE35OVwxojNfbynlnnc2NNUv+msGoSSuRLQdRBCj33dIgkSGo1O7Zai4fTLvrNmvaRRGI1PukxVe8U3iH9LukPOHZkhM6Gri8eUeHp0Y/vdjE3xcJC7ncWa06lAKNAYOHafeerP2dkobxpTBCcW/Zgzi4meXsaeiIaK3pxGrSaRbioMnpx9bKsv1Y7rx475qvtt2WJfXLsYi8daNI7GZ28fTIQgCWXFZZMVltcv1DAyOlf98vyus4mWk5oomIZAedLwNqS83H+KPn2zhrNw0zJLAmuLoN5htwWYWGdo1id9M7tOm8x+9pD8zX1qp20C1mUUm90tvU0qxwanF/9bsY82eSs3feWvPpapCZYOPW8fn8O32MnaW1uFT1KAaoOaGwPRh2cTqdBicDiTaLdw3qTd3ntWT77Yd5uuthymv92ISBdLirEwdlMnInA4IgkCd28dvPtgUJAQDEs79s/Ed/iOV3xXhLfUg2kRsnW2kTknF0csRck+zaCY5Jpk5k+fgMIe+3xY+23hII9YTfd87WYVPfMP5gziHWCG09vvPZ1oZM6eBu0a0Xs8eLxxVogznUFLcCqpiRZCClUz11JuZRTMX5lzY6lj0YjwdBicUMRaJ928dxbWvrmHbodqIXjO7RaJfRjyvXjf8mI0ZQRB4asYgfvfhJj7bdCjivS0mkRizxFs3jqB3upF6Z3BqUuP08eXm0mcFDSAAACAASURBVBabgABReS0VldeWF3PHxB6YJH3NtKOhrM7NHz/ewo7SOv5zxWCGd+uA2ydzzZzVbCypiVoePLdTHHVuPweqXVHNO2f2SeNf0we1qUYTYHDnJJ6aMYi731kftUEVY5Y4o3OiUQN1GlHV4OX9ghLW7aui1uXDbjHRPc3B9KHZPL+oUPP3Hc1z6fbJrN9fwye3j2V3WR1fbC6ltMaNV1ZIjbUytkcKo7onGxkRUWCWRCb37xSsTNqCDwpKNHtp1q78gppV+0m/YggJQ2sQJJH6LVXUrasLMqZURcJqMjG44xk8Pv5xEqztly791up9x9yfT0JmoTKIKdLKkPf6p0lc2MvE35Z6yU0NvwZEJ3wh4q3Mx5r6TUh6ZLT1ZqIgtruAl2FMGbQ79fX1bN68mcrKSjweDw6Hg/T0dHJzczGbI+f4x9nMvHPTSL7cUsrziwopPFKPX1abpNNNqEiqQs/MJG6Z0J3J/dLbbZNmkkT+8as8zhuQzvOLCtlYUoOiqviaJRI7rBKSIHD1qK5cM7orqXGnr0yswanPewX7ETV2AdF6LSGQBvLttjIm909vt3Gpqsp7a0v4+5fbuXx4Nv+aMajJoWIzS7x5wwj+8dUO5q3cC6BpINl/Ktq+cmQXfn1ub2RVZcGGQ/z3h0JKqlz4ZKVp3mlsSjmkSxI3j+/OuJ4px7zRnNy/E3NizNzyRgGyqmoK4EBAXEcQ4KJBGfzl4v7HxSg1OLHYdqiWZ77fzbfbDiMIBBnci3YIAbEJDfWDaJ9LFVhaWE5ZrZseaXHcMdFwCB4vVFXlxSVFIXNQ8+/KnDIWV0kN5qSVxOWtIG6gD1URCaS0ScjVwxjf5TKenHRWu49PS8kY9PW98yFRroY38P40wcbgF+q5b1T4/VK9Gl2ZhtQwHEunb9HyQUWqNxMFkYGpA8mM1VfnGgnDmDJoN0pKSli6dCm7du1CEAT8/qPFiBaLhQULFjBkyBBGjBhBYmLrKlQmSeTCvAwuzMtgR2kdC3eUNTXkTPQ66fXfxzj7bx8cF6+ZIAhM7NORiX06sreigY/XH+RAtQuX14/4yYecc9e1TBqWg9nY0BicBqwsqtD0fuvxWjZ4ZTbsrw4ypvZXOtl8oIY6tx+rWSQtzsbwbh2iivTsq3Dyu/kbqXH5eH32cPplhC7iJknkwfNzuefsXny68SAvLS7iQLULj0/BahbJTIzhxnE5TMnLaFLCMgGXDsni0iFZbCqpYVlhOZUNXsySSLLDwqR+HclKat/ebqO7p7D29+fw9daA82h3WT3mn4wnRVERRYGrRnbhypFdyEhsW02owcnFgo0HeeC9DXj8iqZanK8VCTk9z6UALNh4iOuN+rvjys7D9VQ7fSGvt/yuVH8C3iPn4j1yDogeBNGDqlhAsQEi39XKMK39x+dXtCPjevreKYh4Ce8s79FBZEY/M0+v9jIgLXTv5FVNfKUM1TjzKCZRwGISeW5mPrWmP/LIikd09+iMNcfyyJhHdJ0TDYYxZXDMqKrKkiVLWLJkCT5f6IQB4PUG8ltXr15NQUEBV1xxBd26RTeB906PC0qj236olqc+Hc7v/vwVTjmwIMTHmLl4UCZXjWrfDUeXZAd3ntWz6e897z9GekyDYUgZnDZobQJAn9cSoLzBg6yo/LCzrCnqa5ZEFFUN5OsLYDVJXDemK1cM70yKRmNQWVF5ddkenl24m1vGd2f22G4RozQxFonpQ7OZPjQ7qnE2MiArgQFZx095sDkW01HnUWmNm/J6D15ZId5mpnMHOxaTMd+cLny5uZT739ugu5auET3PpcevcLCmfSWiDUKpqPdoOonCf1ciKDGoSvBeps7tQ1XVdncix1nNQKhRoqfvnVmEhGY1T1r8YbyVNzZqrycmk8S2jJlY9gfmuuY183aLhKrCtMGZ3DQuhy7JDuAiajw1PL3u6agMKlEQiTXHMufcOWTERt8TMFoMY8rgmFm0aBHLly8Pa0g1P66yspJp06bx1ltvMWvWLLp27Rr1fdbtq+KhjzZTeKQeb1o/FNdRb3mDV+aVpXt4ZdkehnftwGOXDKBzcvt6kAEsXbviLS4mZqBRs2BwehCuFlGP1xICTo/z/r2YA1Wupvz8loXzDR6ZZxfu5tmFu/nnrwYyZeDRRW97aS2/eX8jMRaJ+beNoWtK+xRfn2ikJ9i0mxkbnPLsrww0tdYypCIJvTSi97mMVB9ocOyEiyTq/a5UNfCvvRNyxvVKoai8PqicAfT1vVMFM0PNxUFCe8V3B6eOZieIuH+vLZwjZpzBP2dfzP21bj5YV8Lecid1Hh9Jdgt5WQlMGRjaBP6qvleRGZvJX1f/lVpPLS6/C7VF3ZVZNDel9j0y5pHjYkiBYUwZHCM7d+4MMaQ2bdrEihUrKC8vx2q1kp6eTn5+ftB5Pp+Pt956izvvvJPY2NiI9/li0yHuebf5IqMRJpYD7y0vLOeC/yzhzRtGkJcVuamlHixduuDdu7ddr2lgcCKT3cGOKBCSbqTLaykJfL7pEG7f0RqkcDQ+4w+8v4EGj59LBmfyzPe7eXPVPh44tzeXD8s2iuINTkleXbZHM+UqGkGJRvQ8lwApjtYV1gyOnXibSbNXg97vymoWo2t2rpNrRnfl9RV70RpktP35enWKp0fGUNh6CPQ2xDXb4fx/AtAx3sZtE3pEferEzhM5M/tMCg4X8OqWVyk4XIDb70YUROIscUzJmcIVuVe0e41USwxjyuCY+P7774MMqRUrVrB06VIuvPBCunfvjiRJ7N69m+3bt2OxBE/aiqKwdu1aJkyY0Oo9lu8ub2FItY6iQp3bz8yXVvHpHWPp1o4ebEvXLtQvWtRu1zMwONG5fFg2H/14IKRuSo/X0i+rKIqs2RAyHG6fwh8+3sLT3++if0YCX9yVb/S5MThlaez10zI6oEfoBfQ9lw6LxBmdk47fhzIAAn3ltJxIer4rgDOyW/+uPH6ZBo+M3SLpUjfOSrIzuHMSK4oqNN+P1PfObpG4dXx36Psc1JfB/lXgi9KgMsfAjDegU17U422JIAgMTR/K0PTWa66OJ4YxZdBmysrKKC8vb/rb7XazcOFCpk6dSm5ubtPrvXv3pnfv3ixqYYT4/X5WrVpFfn4+kqT94Hv9CrfMK2hT2oPT6+f2t9bx2Z35Iee2lUCanxGZMjixKat189mmQxyqceP2ySQ7rAzpksTo7sm6PJuqqlLt8qFoyKJD9F5LFTQNqUjPsFdWMIkCL1w1xIhGGZzSfLP1sObregQlGon2ubSZJcb1Sj2mcRtEJsYicemQTP63en+IURXtd+WwSNwyoXvItY/UeXhr1V7mLi+mxuXDLIn4ZIU4m5mrRwWEa6JxQj14fi6XvbBcd62eWRTI7mDn7L4dQRJh1gfw+f2w4a3AAX5tpUAssYGI1BVvQ9YvZwS1F4YxZdBmVq1ahSwf9VaXlJTg9/uDDKlIyLLM7t276d27t+b7X21pe38bRYXCI/VsL62lT/qxNbhUVZWVRZW8uKqedV2n433oS8ySQHKslatGduHSIVkkxESWfTcwOJ6sLKrghR8KWVZYgcDRmiSBgPfQYTVxY34OM4ZnE29r/fe6fHc5T36zk2qXj8uGZvFBQWh0CiJ7LQUhkOffkmhTl8pqPewuq6dnR0O62eDUZW9Fg+bzFUlQIpxDItJzaTOJzB7brc190gz0cf2Ybry3tkQzQhXpuwKQJIH8HilNf7u8Mr/5YCNfbi5FEI7O9Y3/rXH5eHFxES8sLmJi71SemD4Ih9WE16/w9dZSPt1wsEkSvYPDyoV5nfjnrwbywPvRi5+YJYGUWCtv3TDiqCiXZIIpT8GE38LaObDqBZC9IEqBhUD2QuZQGHs39Dg78PopgGFMGbSZQ4cOoTbbJTmdTux2O6IYvfKU3+/nyJEjYY2p538oDOm9oq+/jcorS/YcU5PLBRsP8tjn26h2+nB5ZVRLLPhkXD6odfv551c7+PuX27kwrxMPX9SPuAibVAOD9kZRVP68YAvvrCnB7ZNDMt9VAiItDV6ZJ7/ZwYtLinjnppHkpIbWK67eU8mT3+ygtMbN3Wf3YsrADAQCRs3inUdw+6P3XNpMIl5ZCRmPrmdYUXhl6R7+dmnb00AMDE506tx+TRn01kQK9NRSNUcA7FYTM0d0budPYRCOnNRYzuufzpdbSnVHfyySiEkQ+M0HG3loSl9UFaa/sILi8oamWnEtGg2rhTuOcP7TS5jYO43315WgaPS0W1FYDoLAwKxEVu2pJMYshW14LhCItuWkOnjj+hEkadXdxaXDmQ/CuF9D7QFw14DJCo5UsIfvA3WyYhhTBm2mUe68EbvdjtPpRFGUqA0qRVHYsr8CYXc5CTFmEmLMxMeYibOaKKvzUFhWH3KOnrQHWVH5dMPBNhtT//pmJy8sLmx18muccD7dcJC1e6t47+ZRpBm1HQY/E6qq8psPNrJg46Gwi19zXD4Ft9/D1GeX8entY5tU8dbtq+Jf3+xkT3kDd57Vk2lnZAbJjv9n5hncOm8dKwq1+061JMYsMSg7kfX7q3C1eH70PcOBiJuBwalMfIxZl9CL3lqqRgKGlMTbN44k0W6IT/yc/POygRyqcbOhpDpqgyrGLHL7xJ5cM7orf/18G+c++QN2q4n9la5WDanmePwKeyuczF1erKWDAdCksLpqTyUJMSauH9ON99eVUFHvRRIEVGiKgI3tnsJN43MY0a1D5PRryQRJXaIa58mMYUwZtBmzOTgCk5WVhclkYvv27fTt2zeqa6gI7DziYtH3u6lx+Zr+uXwyMRYpqNdAI3r723hlBY9fxmrSF05+eUkRLy4uinrS88oqJVUupr+wggV35hNrNR4vg+PPvJV7wxpSrdUkNXj8XPHSSp6bOZinv9/FjtI6bp/Yk18NydLsa2Q1Sbx89VBeWlLEf38oxOtXmhbgRho9lgkxZu45uxcNXj/r9lWFXEvvM1zv8Uc+yMDgJKZ7qoMYsxTyTIUTKahd+wmqz6OrlspuFrFbTfzvppH0SDPSZn9uzJLIvBtGcP+7G/hqayk+WSGcPWQzi6gq/HFKPy4fHogg/uWSATxilpizdI+mURSpBjVa/Z9at5+5y4v57M6xVDv9lNa6cPsCfe96pceSFmc4i1ti7PYM2kxqaiqlpaVNqX42m40JEybw+eefI4oi3bt3RxRFioqKKC4uDjG+AKwWM9PPGkC/fsEpCX5ZYcmucm5/a13I4qK3N4MkCnj9ii5jam9FA//8akdIHxxofcKSFZWD1S7+9sU2Hr14QNT3MzBoC4qi8tR3uzQNqUgpQIoaEKq4du5q7pvUm/9eNSTiMyKKAjeP784N+Tks2lHGK0v3UFzRgNunYLdI9O4Yxw35OYzMCXgs31y1F62SDL3PsMVokm1wijOxT8ew4jBaIgWSPQnR5ojaISEIcNfZvbhmdFddSm8G7YtZEvn3FWew63Adc5YV89GPBzCJQsATRWBOj7FI3DC2GzOGdQ5KoVNVle+2HdY0ivSmfLa2j1HVgEF1w2sFfH5XPn0zjq3m/HTAMKYM2szw4cPZtm1bkDT66NGjiY2NZfHixXz44YdYLBYyMjLIz8+nsLBQ8zpa9VImSSQrKUbjaP29GfyKisOi76c+d3mxpoJZNBOWV1b5oOAAD56fG9JkzsCgPflh1xHcGk03o00BktWALO7Vo7rquq8kCpyV25Gzcju2elxanA1JFAnq5Ij+Zzg1zqprfAYGJxsWk8hVI7vw8tI9mhkZLUUKXEUFlL3/p6gdEjaTxAV5nQxD6gShZ8c4/jptAL+/IJctB2upcfkwiQLJsRb6ZyRoGtbr91dzuDZUHU9vymc0+xhZUdlT3sCmkhoGZCW04yc/NTF2egZtJjMzk/j4eCoqgusZ8vLyyMsLLRbPzs4O+luSJIYMGYLJpP0z7JLs0MzH1duboXfHOF1y0O3R70MQAjVUM4YZBb4Gx48XFxeFRG5BX01S4ZF6dpfV0yMtcvNsveT3TAkSqWlEby+cK0ee+jn3BgZXj+rK3OXFeCMfqtsh4ZMVo0bqBMRhNTG8W3SCDF9sPoRbIwtBz3yvZx/j8cu8tKSIp684I6rxnc4YuRMGbUYQBMaPH6+ZvhcNoigyfPjwsO9bTCKzRnTGLIUaQvHDp5E0cTY1K96h5D+zKHn+WurWLSCmZ/Bk4rBI3KrRm6E1luwqR9Qw4vRMWE6vzLyV+3Td18BAL7sO12m+rqcmySSK7C7Tvs6xYjNLzBiWfUzPsApMGZhxXMZnYHAikZ5g44WrhmAzR96aNXdIOHeuQPG5UWU/rsK1VC2cE3L8wOxEo473JOdwrUczxU/PfK/L8FLhyy2luDQcdgbBGE+WwTExYMAA9uzZw+bNm4PS/SJhMpm45JJLSEpqvaP3VaO6MHd5MVqlk9H0ZkCRmdw/PepxQaAJnqwcu/BFeX2YZnUGBu1EuEVOT02SrKjUuY+fwMN1Y7rx1qp9+NrwDFtNItOHZhupSQanDfk9U3nhqqHc8kYBPlnR7EvUiJ6GrzePyzneQzc4zshhfgu65nud+xifrHCgykkPo89fqxiRKYNjQhAEpkyZQl5eXtQRKrPZzNSpU6NS/MtKsnPRoAxiovDUtSRGgqt3f0/lo4+gNDREfZ7XL0fs9xENvihlSw0M2ko4wYjmKUCRkESOa21fdgc7UwfpjyyZRIGspBh+PVm7B52BwanK+F6pfHl3PpcPzybGLLW6/sX2O5NO1zxF53s/IPv2eaRd9jC2rNygYywmkYl90o73sA2OM2lhakd1zfc69zGqCrNfX6OZrm1wFMOYMjhmRFHkwgsvZNq0aXTq1AmTyRRS62QymTCZTPTp04frr7+eAQOiV7p77JIBDMpOjCr1oZEYs8Slw7pw/4t/RPV4KZp6MQ2rV0d1bnyMOaCu0wI9ExZgpFQYHHcykrQlavWkACkqYcVe2oO3V+9j4Y4jXDWyS9TPsNUk0jnZzjs3jzJEXAxOS7okO3j04gGse+gc/jS1P2fnpmkqY0bCZhZ5+oozgnrGGZycnNknDYcl1IGmZ77X2sc0bF3EodfuZt+Tv6Lkmas4/O4fcZdsaXp/b4WL33yw8fh9sFMAY5UyaBcEQSA3N5fc3FzKysooKCigoqICr9dLTEwMWVlZDB48GIfDofvaZknk9dkjuO/dDXy9tRSfrIYNd5slAVEQuHl8Dned1RNBEMj462PULVzIwfsfIG7yuaTdcw9ijPbmsbi8gSU7j2gW9espmjeJAqO6p+j+rAYGepg9thu/n79Z8/cabQpQB4eFvOOg1qSqKs8tKuR/a/bxzs2j6JbiYGyPZB77fDtH6j24fDItnZ12i4Siqlw6OIsHz8/FYTgkDE5zYiwS04dmM31oNp+sP8CvP9gYde9Dm1nk79PyyO+ZepxHafBzMKZ7Cg6r6Zjm+5b7GF9FCTVr5hM36Dysmf1IOvM6TUn19wtKGNKlAzOGZePxy3y5uZRXlxVzqMaFx6/gsJgYmJXADeNyOCM7MXIz31MMobXQ3dChQ9W1a9f+jMMxMGidrQdreWVpEQs2HsJiElFU9af2DIEO3TOHZ3P1qK5kd7CHnOuvquLwo3/BvWULnf76GPYzjirUHKnzcPvb61i/rxpFUfG1kqdev2UhdWs/xlexP2jCap5aYTOJLLgz/7gopB0rgiAUqKo69Jcex7FizE8B5ckhj3yjubhGg90i8eD5fbhyZNd2HZeiqDz62TaW7S7n9dnD6Rh/NIKmqirr9lXz0uJCftxfjdMjYzaJpMZZuXpkFy4+I9Mwok5jToX56XjOTcsLy3ngvY1UOb24vHJIJWJj4+xEu5nHLxvIaMOpd0rx0uIinvhmR9QGdTjqtyykdvV8fGVFCFYHpoQ0VL8Xua5Cs+EvgFUSuHJUF/63pgRUNbTBtABWs0RanJXfndeHyf07HdMYTzRam5sMY8rgpKTW7WPd3ipqXD5EQaCDw8LQrklRNeat/eprSh95hMSLp5Jyxx0caJC55LllVDt9rRb76iEvK4FPbo8sV/tLcCpsVsCYnxr5+xfbeHV5cZsW11iriZUPnhVVSqrHL1Pn9hNjlrBbpLCeR5+s8Ov3N7K/0skr1wwjwd42tU+D05NTYX463nOTqqqs3lPJC4uLWLLrSFOmhiQK5PdI5abxOYzo1uG0iw6cDtR7/Jz/7yUcqHaFzdCJlsZeZYnjrqF2zYchfac8+7cEZdwIgCQJ+OXI97WZRW4e1517zul1TGM8kWhtbjLcfwYnJfE2MxN6t62gNv7cSdiHDqH04T+xYfpMbh1yI5UubdGJtmAzi/y/83MjH2hg0A7cO6k3q4ur2FRSjTeKRa4Rm1lkzrXDWjWkat0+PlhbwotLijhc68YsiciKilkSuXRIJteP6UZO6tHoq8src9ubBQiCwBuzRxCjkd9vYGBwbAiCwIicZEbkJAM09R4yVC9PfWKtJt65eSQXPbOMqgbvMTmAZVctYkwcNcvfjqrvlApRGVIAbp/Ci4uLcFglbhqnrz3NyYhhTBmclpiSk8l8+t88//RnlB/woIihj0LD1kXUrvkIX0VJ2LB3S2xmkUen9m9a5AwMjhWfrPDN1sN8vaWU8novoggd42xMGZjB2B4pmCWROdcMJf8fC1FUOeLiKgqBTdcLVw0J2yxSVlQe+3wb81buRRQEXD9t1jz+QPTLr8j8b/V+3ltbQl5WAs/OHIzFJDL7tbV0Sbbz90vzMBsF7wYGPwuGEXV60Skhhi/uymf23DXsPFyPJ4wCsSiARRJx+7WzFqSYeBRXHahqVH2ntGhtn+TyyTz5zU4m9kmjR9qpLa1uGFMGpy0ev8L8ChN+DYmk2tXzqVn1fkjYu2VRZiMxZhEVeGrGGbr7WhkYaFHt9PLSkiLeWLEXWQnNT/980yHsVhM35nfjcK2HvKwEJvbpyEtLiqhx+ULqKWxmEVWFSX07ctfZvcLW8/lkhevnrmFtcWWT8aSFX1HxKyo/7qtm0lOLSYqxcFZuGg+en4vYFtkxAwMDA4OoSIm18vHtY9lUUsPLS4r4ckspkggqAn5ZQVZUunRwML53Kt9vO8y+KlfINayZfUAUEUyWqPtONSeafZLPr/Dykj08eEEun288xL5KJ7VuH0l2C33S4zmnb0csppPf8WYYUwanLQs2HkJry6d4Gqhe+mZUYW+AZIeFm8fnMH1oNol2y3EcscHpQnF5AzNeWEGV0xs2da/BK9PglfnnVztQVfjy7nH0SIvlujFdWVlUyf/W7ONgdUBpKTHGzNieKRF/o6qqcu8761lTXBl1DZZfUal2+vDJCreMzzEMKQMDA4OfiQFZCfz10gEMyEoIONKcPvyyigrsqWigZJUTRQ3UO7VcSUSrg9h+E6nf+DUNO5YRkzMkrEJxS6LdJ8kqvLt2Px+uK0ESxaYsBwCHVUL8UODKEV24ZnRX0hO0W32cDBjGlMFpy/wfSzRV0DwHtqP6vVGFvS2SwB1n9eDa0d2OxxANTkMO1bi45Lll1Lh8UdXx+WQVSYDZr61hwR1jibOZGdU9mVHd9aeartpTyXfbyzQNqUhpr16fwpPf7OKxadH3kDMwMDAwaDuHalxc/uJKymrduDTmbV+EGqekibOp3/I91QtfpeKzJ8NKqrek+T4p0tqgqAScgnLwfqvBE/j7laV7eH1lMa9cM4yRJ2mJhGFMGZy2VNR7NV+XXbWI9viowt4+WaXG6W/voRmcpqiqynWvrqHW7Q8xpFpbsGQVDtW4uffdDbx0dduF0F78oQiXhoMhqnQORWX+jwf4/YW5RqNdAwMDg+NMWZ2bKf9ZSpXT12ZlP9HqIGnc1dSs+pCUC+8P6p9ZtXBO2OhU4z6pbu0nukoitPDKCl4Zrnt1Na9eN/ykNKiMFc/gtCXc5CPFxKE4a8lUD+EVbNQQixdteWcVUNRj6/dgYNDIj/ur2VfpDPltRmPMeP0Ki3ce4VCNi04J2k2pW6Os1s2ywvKQVBA9aa+CAB/9eICZI7rovr+BgYGBQXSoqsrVr6ymWsOQ0iueFW3D3+ZIMfEozlqqlswj5YJ7oiqJiDQ2l09h9mtr+P6+CUG9CU8GDGPK4LQlyRFcO5JAPdOlRUzvtoBBJoVbi+7j0r4WJBS+VQbzkv8CflR7QLNKK4tJJCHGqJMyaB9eWlwUlFMO+owZFXhjxV5+PbmP7nt/v70MUaMvjZ60V6dX5sN1hjFlYGBgcDxZvaeSfZXOEPVWveJZjcT2O5PYfmdGff+AeIUEUa4N0Y7NL6u8vqKYByZkQvEycFYETo5Jgi6jISYx6jH+nBjGlMFpy/n909l0oAaP18eDpje5UvoWBQG72cufJ1i574s6YiUbk7qbmCisxrN3NSnFEkXj/0ahmgmACIzrlfrLfhCDU4I6t4/vtpfRso+6HmPG61eYt7JtxlSl04tXDk3x05P2ClDl1E6fNTAwMDBoH15cHJqSrVc861gQrQ7svUbj3LYY1+7VQemBWuIV0Y6ts7yPnBUvoq5ZgSCZQPkp80cUQfZB36kw6nbolNeun+dYMYwpg9OWaUOy+McXW3jZ/Dgjxa3YBF/Te/eNtpIeK/LoYg+zPnQRZxEYkiHx27FWzrA8xEzv/2Oj2p1e6XFhJaYNDPRQWuPGLAl4W5Tg6TVmGjwybp+su/eMqhJiyMHRdA5VkaMaQ1tz9w0MDAwMIlNR72Hp7tCUbD2Ot/Ygtv9EnNsWU7P8fxHTAyONTUDhYdNrTJd+wKz6EfwKaJWjb3oftn0KuVNh6jMgnRhmzIkxCgODX4B4m5nXkl+nX/VW7EKoN31WnplZeVq1Um7mWR7jV+rfuXX84OM/UIPTggavrJlmp9eYMUkC9R6/bmMqIcaMxSSGKPlZM/sgmMw4d67A0WdsxOsY7QEMDAwMjh/FFQ1YTGJIH0C96YuTvAAAIABJREFUjjcAh0XSVDVufE8lkL6thTWzD4LZQvzIyyKuDa2PTeU/5v8wUfyRGI29WPChMvicsO0jaDgMs94PpBv+wpz8nbIMDNrKgQKGNCwOMqS6PlVH2j/raPAe9fm8vM7LhLkNQac6cPOo7U3O6dvxZxuuwalNrFVC0QgNNTdmosEnK8Ra9fvJ8numaEamRKuDxLGzqPzmvzh3rkDxuVFlP67CtVQtnBN0rM0scv4Ao2m1gYGBwfGizq2tINzc8RYtf5uWx4DMeEyigN0iYbdImCWBnmmx/Hlqf9Y9dA7p8VbNc/WsDa2N7V7pPSaK6zWd2mHxuWDfSvji19GfcxwxIlMGpy/Ln0WUPSEvyyr8e5WXB/O1JxAASVAZJq9HcB6BOMOgMjh2OiXE4NfoCdJ8wRJEKWJueqLdojsqBdAl2cGAzATW7q0KeS9atSdVhelDs3Xf28DAwMAgOsK1ntCbRWASBaYMymDKoAxq3T5qnD4UVSUxxkKC/WhWzh0Te/Lwp1s0e1ZFuzaEG1ssTm4yfRZUZgEBx7bTB3vuisVhCWRsvLzOy7yNPhZd6wgc5HPCujdg/G8gNi3i5z2eGMaUwemJsxJ2LAANWfMHRlv4xzIPtw2zkGgLTbtqRFFBKngVJvz2eI7U4DTBYTVx/oBOfLL+IHKLEFHUC5ZJ5NrRXds8huHdOmgaUxBZ7UkUYFLfjkaan4GBgcFxYH+lk3kr97J2bxX1GtEpvY63lNijDuN4m5l4m3YLmIvPyOTRz7bh0xAoguiUAMONLW//G/xun5N/TQq9dzSObQQB1r4KE37T6v2PN4YxZXB6smcxiBYgNDI1NENiQlcTjy/38OjE8L0OJMWD68d3iTGMKYN24ob8bny5uTREHh2ilK6VvcxcPwsKKsBsh465AeWjrvmBRScMh2vd/PHjLeworaVPehx7yhtC8vEj4bCa2qQiaGBgYGAQnpVFFTz17U5+3FeNoqqaEaJGonW82cwi146OroWFw2ri8uHZvL16X0hNrR60xna4k4fbx2mvTVE5tv1uWPVfGHf/L1o7ZRhTBqcnrkpQtPOOAf58ppUxcxq4a0TrXvaGmgrKK51kd7ADUO/x89GPJXy1+TCVTi+iIJAaZ2Xa4Ewm9U3HYjLKFA3C0y8jgdxOcWw6UNPqgqmFDS+ThVWk1O04+mLNfiheCrb/z959h7dV3Q0c/56rYdnydpy99yJkkZBFEjZhlhEg7A0FCrTQ9oW+LW0ppS19KaMByghl7713CJkkIQlZZDrLcWI73ta+5/1DSrCta1tSvBL/Ps+jJ/Kd5yrSufd3ZiZMuxNGXVQrqDJNzStLd/CPT37kovE9+dcFIwmamnNmLyCvOLaASgEpThvPXTX+wO9ACCHEwXvq2y3845Mf4wpiYil40xrOP6pnzMe8c8YQfthZFp5OJs6CtvrS5sbDx0nX4lDWNV6xFmwT8EDlHkjvmnC6DpY82Yn2yaqnfQ3DO9o4baCd+75tuENkitPG1f9dysY9FfzmjVWMvecz7v1wPfM2FbEmv5wfdpXx5fq9/OaNVYy55zP+9tF6KryBBo8p2rcnLzuKbLcTu1F/TVJdTvz0Vfnc53iyzhoN/ioo3wUf3QHv33Zg3o4thZVc+MQiXv5uBy9eM55fnTgIl8NGapKdt2+cxDEDc0myGzgaSIfbaaNLpou3b5zEyB5tczJFIYQ4FD27II/7P9lwULVBVpw2xSnDu5Dljr1JtsNm8NxV4xnXJ5uUBPrkWkmnmgANH+tP05N4eImfwqoGPgPDBt6yJklToqRmSrRPKdlgNPz1/+M0F6Mfr+RXE+pvr5uclk2/bDenPDgPjSZUz++9yhcueXnq2618uHo3L197NF0ykhNOvjh8ZbudvHvTZC74zyIKyryWTf5qSsbLMJXHHOc/ojrx1hKohlWvELK7eMx1NU/O28LNxw7gsom9sdUJmJKdNp64dCxbi6p4Zv5WXlu2EwUYhkLr8OTA4/tmc90x/ZjYLwcjjsBPCCFEw9bml3PvR+ssA6mqtV9T/t3bBIp3YjiTcXTsS8bEmbi6D4vp2O4kO/eePTzuNCU7bTxzxTjeW5nPo3M3sz3SeqHm1ILuyGiAPbJT2Li3ssFA0I8dI2q2rNpqFmwPya2v/keDrXX76kowJdqn3sdAqOFap/7ZBucPc/DQEj9HdLT4EduT2NL7Ar5ZWkQwxolK/SGTnSUezpm9gA9vmSKd9YWlTukuPvjFZF79bgf/mbeF0uoAHn/owG3HQJOEn+5qL9fb3+cMYwEDHiyNafSjwOKnKe/Yj3dvuqjRZnl9Orj545nD+Z8ZQ8gv9VDhDZLstNExLUm+u0II0Uwe/2YzgWD0c0X5krcoW/w6OSfeiKvPaJTNjmfrMjwbFx8IphoLtjyBEN9uLOLEYfFPY2EzFGeN6sZZo7qxJr+M91bmU1DmxR8y6ZCaxJQBuRw7ODyy3l1v/cA7K/LrLRAsw41B47VujRZsh/zg7hD3tTQlCaZE++TOgUEnh2fSthjRb7/fT03iuVXWpf2mqbhs9ZFU+aODsoYys5CpKaz0cesrK3jminFNdkni8JLitHP5pD5cNrE3i7fu48v1eymq8GEzFJ1+fIGT/Z8w3MirtU8sox8l4eO3aR+jsq+JOS0uh42+uamJXooQQogYlXkCfLy6IGpUV9NXRem3L5Az41ZSBk08sDyl/3hS+o8HYgu2vAGT+z/9MaFgqqZhXTMY1jWj3vV/PfsIpgzI5eEvN7KtuAp/SBOqUfBs2Jx8ocdyolqCrYEaqkYLtntOBFf96WgJEkyJ9mviL2DjZ+HmTxF5t6bV2qRHhoH3d+kWOyvmdzyfkvxQVPerWDKzQEizcHMxu0o9dMuU5n6ifkopju6bw9F9c8ILdi2DDc+BUR21bSyjHymAvHlQnt+qHXaFEEJEe29lPobF6Ku+XevRQT8pAydY7hdLsLXf9n3VrN5VxvBuzReEKKU4dUQXTh3RhbX55byweBubC6uo9gdJdzkY0T2Dsb3+gO3Nc2s9h1mpt2DbmQqTbmmmK4idBFOi/eo2BgbNgPUfQNAT375JqTyuf0aVv3YGEE9mprXm2QV5/M+MIQlfgmiHlj8XHg7WQsyjHwGsfgMm3twMCRRCCJGovKIqy6ZxIU85Rko6qp4hwBsLtmoKBDVz5m/lnzNHHnR6YzG0azp/+dkR0Su0hvRuULwJatROxVywnZQGfRuZMqQFyGh+ov1SCn72GPQYB45Ya4cUON3sO/s1luyMfqCNJzPzhzQvLdkeZ6JFu1eS12DT1JhGPwr5oFS+e0II0dZU+KynbbElp2NWl6NN6z5IjQVbtbbVmjX55QeVziahFFzwAjjd8e/rSIELXwKj9UOZ1k+BEK3J5oCL34QjZ4E9Cez1leaHgygye8JVn1OQNsxyzqh4MjMIz0sVqG8IQCGsNFKLGuuw/vgbblYhhBCi5WWlOCyXJ3UbjLI7qN6w0HJ9Y8FWXZX1BG0tLncQXPJ2uJaJGEeGdbjhwpeh66hmTVqspJmfEDY7lSf8nbdTr+ClBT9S7PMS1AZu5WGqfQ1X2D+lT/8h4Xa5vSaBUni27cOiSXOtzCyWgMpmKKr9ITKSpVxDxCg5q9FNYhnWH3duEyZKCCFEUxjQMQ2nTeGvM3G7keQmc/JF7PvsMZRhw9VnFMqw481bgXf7KjImnn8g2HIPntzoeZKbaL6oJtHjKLh2Lnx4O+TNDy8L+WpvY3OCMqD7WDjlH9BpaMunsx4STIl2raw6wF8/WsfbK3ZhKEW1PxkIN/kr0hm8FOrMK/oEhpWl8ztjKKMjEVRqkgNtMfFvzZKjWDKzYEiTmiQ/QxGH3lNgy9wGO+w2OvqRMxW6j2nGRAohhIhXhTfA84u2RQVS+6WPOxvDnUXZwlcoev9+lDOZpE79SZ9wfqPBVtb0K2sdq3dOAk3rmlNOP7jkLSjbBd89Fe7X6y0Nr3Olw5AzYNw1kNW7VZNpRZ7iRLu1q9TDzMcWsLfCR6CejCtgAqbJ8u2lzHpiEfefdySnjehK96xkghb7xJuZdcl0RU2YKkSDRs6CL/7U6GYNDeuPYYeBpzRxwoQQQiTKGwhx3mML2VJU1eB2qcOmkzrMetCFhoKtmlKcNi6f1Lupkt60MrrB8b8Pvw4REkyJdqm02s+5jy5gb7kvai6H+ngDJre/tpI0l4OpA3M5dUQX3l6RX2veBIg9M0t22Lhmct8muybRTiRnwZDTYfWboH9qGx/z6Ee2JBh/Hdgk+xdCiLbiN2+sYmtRFf5gdD/qxibiramhYGu/dJeDif1ymizt7Z3cTUW79Ns3f6Co0jqQaijT8gZMbnh+Gd/ddTxXTe7LBz/sjgqmILbMzNSac8Z2b7JrEu3IMb8OD+nfyNwclhwpcFTsE/YKIYRoXoUVPj5aXWAZSMUydyXEHnAlO2zcMK0fyqrjt0iIBFOi3Smq9PHV+r2WTftizbTe/n4XFx3diyO7Z/L99lL8cY7Il+ywcf5RPUh3WY/aI0SDcgfCzOfglYvjmyPN6YbL3oFUGXxCCCHaihcXb7ccxy7WuStjfXYxFBwzsAOXTujVnJfT7kgwJdqdl+uZ2ynWTKvaH+KxuZuZNb4n/7l0LKc9PI+CMm+9/a7qctkNRvbI4HenymS9on6mqfl2UxGPf7OZVTvL8AZC2AxFttvJrHE9uWDcFGxnPY/z9Ytw2RS2UANBlcMNDhdc+i50Ht5yFyGEEKJR/12Yh8+iViqWuStjfXbZ784ZQ6RWqolJMCXanRcWb08409qvuMrPut0VDO2azjs/n8QpD81jT7mv0f3shmJS/w7Mvng0dpsMhy6svbV8J3/9aD1VviBV/p/6RQVCmvxSL498tYmHv9xEerKdmSPf4Y6O38GCR8BfGZ5RXgfDg0xoILUjTL4Njjg3sYkRhRBCNBtvIERptfW8gLHMXRnPs4tNwWtLd3L7SYMSTq+IJsGUaHf2VSWeae1nMxR7KrwMNtP48wfrKKuuZ9Q0C9v3VVHtC5Fkb0NzPIg24+8fr2fO/K14AvU3HfVG1hVW+HljbRUXHHMVPY6+EbYvgJI88FeFJ0DMHQRdR2M5KZoQQohWV+UL4rAZloW8scxdGc+zS8CE9QXlB51mUZsUjYt2J2gxYATEN3u41uALhLj7vTV8vLoAr0UmWN+584qruejJxXgDsc1SLtqP/3yzmTnz8xoMpOraW+Fl5uMLKfUGofdkGHVxeLS+kbOg2xgJpIQQog1zJ9kJ1NPvuubclfWJ59kFoMwTe+GviI3UTIl2x+UwqPJFZzrxTLirFOwu9fLa0p14LIKihkbVCYQ0WworeXzuZm45fmCTXZc4tO3YV80/P91gWTrZ0PfJ1FBU4ePP76/lnzNHtkLKhRBCJMrlsJHqslPuCUats5q7svrH+ZQteJlg2V5syenYO/QCwxbTswsgA181AwmmRLszvGsGi7fui1oez4S7Hn+IT9YW4A1GB1KxjKrjDZrMWZDHjdP7S98pAcCc+VsxLYbqj+X7FDA176/azR/OGCY3SiGEOMRcPL4XT3671XJo9JpzVxa+cx+YJs5Ofel43p9w9RiGZ+syys1QTM8uDptiUOe0qHOIgyPBlGh3rp/aj9W7ymp17N8vlgl3FZDjdrJoS3RAFs+oOoGgyZfr93LisM5Nd3HikOQNhHjlux1RI0LG830yDMWby3Zy+aQ+LZJmIYQQTeOSCeFgqj6pw6aT0n8cO/99GTmn31qrBmr//aByzVcNPrsAGEoxa3zPZruO9kqCKdHuHDMwF5fDZhlMQeMT7iY7bZw5qhvPLYzu2xLPqDpV/vADtARTYuHmYsuhauP5Pnn8IV5askOCKSGEOMR0yUhm2sBcvv5xL/56pllp7H7Q2LMLwOieWXTPSjno9IrapH2RaHdshuLXJw8m2RH/aHpOm8GQLukkO4wDI6rVFM+oOgAF5d640yAOP4WVPkyLgVHi/T4VVzU+PL8QQoi25//OH0m3rOR618d7P6gr2WFw83H9E02eaIAEU6JdOv+oHswa3zOugMppU3TOcDHniqMIhDRWZUfxjqpT3wg+on0JhExMi29UvN+n+kaqFEII0balJtl544ZJGPUMwBrv/aCmZIeNG4/tz8R+HQ4ylcKKBFOi3frdqUP4+bR+uOwG9vpyr4gUp40hXdJ576bJpLscZKY4cNii94llGNOaMpJlsAAR/h7YLL6D8X6fUpOk5bYQQhyqst1OOqQmWa6L936wX7LDxi+O68+N06RWqrnInVe0W0opbj5uADNGdGHO/K28sWwXNkMRNE20BruhCJqakT0yuX5qP44ZmHvggfeo3tnYDYNAqHYJUTwjArrsBtMHdWzRaxZt0+ieWVGDT0B83yebAZP7S6mjEEIcyib2y+G9lbsJ1RndNZ77AYS7JUzol82N0wcwrk92S15CuyPBlGj3+uWmcs9ZR3DnjCHM21hEcaUffzBERoqDsb2y6ZEd3VlzZI9MOqUnkVdcHbUulhEBATRwwTgZVUdA18xkjuqVxfzNxVHrYv0+OQyDqybL4BNCCHEou3pKXz5Zs8dyDstY7wd9OqTw0jUT6Jzhaqlkt2sSTAkRkeK0c1KMI+sppbhhWj/++N5aqi1GBWxsVB1DwbGDO5LtdiacXnF4uW5qP77fUZrQ9wlgQKdUBnSS+UOEEOJQNrxbBt2zk9m4p9JyfWP3gxSnjfvOHiGBVAuSPlNCJOjMkd3ITU2y7OvSGJfDxq9OHNgMqRKHqsn9OzCoc5plX7zGuBwG/3vasGZIlRBCiJZ29+nDcDnif0RPshsc0S1DmvW1MAmmhEiQy2Hj5euOJjPF0egAFjUlO2z855Kx9O8otQjiJ4ah+O+V4+ielYLTHnvW7HIY3PuzI+TmKYQQh4lJ/TvEHVAl2Q16ZKfw1OVHWc5bKJqPBFNCHIQuGcl89Isp9Ongxu1seJh1t9NGusvO81ePY/IAGShAREt3OXjnpkkc2T2DFKeNhu6HyU4byQ4bD184mrNHd2+5RAohhGh2F4zryf/NPBKXw2hwGhdDhQtpx/TK4p0bJ8morq1APnEhDlLHdBef3HoM8zcX8djczSzNK8FpN9AalIJgSNM108UN0/pz2oguuBKYLFi0H+kuB69eN4Hl20t54pvNfPVjYa3vU8jUZKY4uHZKX84Z0500lwyvL4QQh6MZR3RlUv9cXlu6gyfmbaHSG8Qw1IH7gT9ocvyQTlxzTF+O7J4hNVKtRIIpIZqAYSimDMhlyoBcdpd52LS3kgpvkGSHjW5ZyQyUgQFEHJRSjOmVxZhLxlJc6WPd7grKvQGS7Aad0l0M65ouN00hhGgHMpIdXD2lL1dN7sMPu8oorPDhD5qkJzsY2iWdLBnIqtVJMCVEE+uSkUyXjOTWToY4TOSkJjF5gPUkjkIIIdoHpRQjume2djKEBekzJYQQQgghhBAJkGBKCCGEEEIIIRIgwZQQQgghhBBCJECCKSGEEEIIIYRIgARTQgghhBBCCJEACaaEEEIIIYQQIgESTAkhhBBCCCFEAiSYEkIIIYQQQogESDAlhBBCCCGEEAlQWuv6VypVCGxrueQIIVpAL611bmsn4mBJ/iTEYemQz58kbxLisFRv3tRgMCWEEEIIIYQQwpo08xNCCCGEEEKIBEgwJYQQQgghhBAJkGBKCCGEEEIIIRIgwZQQQgghhBBCJECCKSGEEEIIIYRIgARTQgghhBBCCJEACaaEEEIIIYQQIgESTAkhhBBCCCFEAiSYEkIIIYQQQogESDAlhBBCCCGEEAmQYEoIIYQQQgghEiDBlBBCCCGEEEIkQIIpIYQQQgghhEiABFNCCCGEEEIIkQAJpg4jSqlpSqmdLb1vS1FK9VRKVSqlbK2dFiFEfCR/EkK0RZI3iYMlwVQDIl++/S9TKeWp8fdFzXjey5VS3zbX8ZuCUkorpfo38znylFLH7/9ba71da52qtQ4153mFOBRI/lQ/yZ+EaD2SN9VP8qbDk721E9CWaa1T979XSuUBV2utP6+7nVLKrrUOtmTahBDtm+RPQoi2SPIm0d5IzVQC9lfrKqV+o5QqAOZYlYjULIFQSiUppe5XSm1XSu1RSj2mlEpO4NxXKKXWKaUqlFJblFLXWWxzp1KqKFI6cVGN5U2SBovz3a2UelUp9WwkXWuUUmNrrP+tUmpzZN1apdTP6ux/TY1rWquUGq2Ueg7oCbwXKc36tVKqd+QztSulzldKLa1znNuUUu8257UK0dZJ/hR1PsmfhGgDJG+KOp/kTYcJCaYS1xnIBnoB18aw/X3AQGAk0B/oBvw+gfPuBU4D0oErgAeUUqPrpKtD5PiXAf9RSg2KNw1KqdlKqdlxpOsM4GUgE3gXeKTGus3AFCAD+CPwvFKqS+Q85wF3A5dGrukMoFhrfQmwHTg9Uj399zrnew8YpJQaUGPZLODFeK9ViMOQ5E+1Sf4kRNsgeVNtkjcdDrTW8orhBeQBx0feTwP8gKvG+suBb+vsowl/GRVQBfSrsW4CsLWec0Udq4F0vQ3cUiNdQcBdY/2rwP82lobIvjvj+Dw00D/y/m7g8xrrhgKeBvZdAZwZef/J/vQ39JlH/u4dOa898vfzwO8j7wcAFUBKvJ+3vOR1qL8kf4o6r+RP8pJXG3hJ3hR1XsmbDsOX9JlKXKHW2hvjtrmEv6jLlFL7lykg7pFVlFKnAH8gXHJgRI77Q41NSrTWVTX+3gZ0bco01KOgxvtqwKUi7aGVUpcCvyT8gwZIJVwCBNCDcOlLIl4E/gn8iXDJytta62qlVEea91qFaOskf6pN8ich2gbJm2qTvOkwIMFU4nSdv6sIfwkBUEp1rrGuCPAAw7TWuxI9oVIqCXiDcLXuO1rrgFLqbcJf9v2ylFLuGplCT2B1U6UhgTT3Ap4AjgMWaq1DSqkVNdK8A+hXz+51P+O6PgNylVIjgQuB2yLLW+VahWhDJH+KLc2SPwnRsiRvii3NkjcdQqTPVNNZCQxTSo1USrkIV98CoLU2Cf8oHohE/iiluimlTmrgeEop5ar5ApxAElAIBCMlLSda7PtHpZRTKTWFcBvh1xJMQ1NwE/5hF0bOeQUwvMb6J4HblVJjVFj/SCYCsAfoW9+BtdYB4DXgH4TbYH8WWd5a1ypEWyX5kzXJn4RoXZI3WZO86RAiwVQT0VpvIFxl+jmwEag718FvgE3AIqVUeWS7QdRvIuESgrqvXxBuy1tCuHr23Tr7FUTW5QMvANdrrdfHmwYVHsHlsYavunFa67WEq5MXEv6BHwHMr7H+NeAvhKudKwi3Y86OrP4r8DulVKlS6vZ6TvEicDzhTK/mEKvxft5CHLYkf7Im+ZMQrUvyJmuSNx1alNaN1QYKIYQQQgghhKhLaqaEEEIIIYQQIgESTAkhhBBCCCFEAiSYEkIIIYQQQogESDAlhBBCCCGEEAmQYOogKaWeUUrdE3k/RSn1YwudVyul+jfxMQ9cS0vu21KUUncqpZ5s7XQI0VIkfzr4fVuK5E+iPZG86eD3bSmSNzWuXQRTSqk8pZRHKVWplNoT+fKmNvV5tNbztNaNDiGplLpcKVV3+M8mo5T6Wil1dXMd/2A19/VHzjFNKbWz5jKt9b1a6zb7uYj2SfKntkXyJyHCJG9qWyRvarvaRTAVcbrWOhUYDYwFfld3A6WUvcVTJYQQkj8JIdomyZuEaER7CqYA0FrvAj4iMpN0pMr3RqXURsITxqGUOk0ptSIy4dkCpdSI/fsrpUYppZYrpSqUUq8ArhrrakX0SqkeSqk3lVKFSqlipdQjSqkhwGPAhEhpT2lk2ySl1P1Kqe2REqDHlFLJNY51h1Jqt1IqXyl1ZaLXr5R6TSlVoJQqU0p9o5QaVmeTDkqpzyLXN1f9NKM2SqnBkXX7lFI/KqVmJpqOOmnKU0rdrpRaFUnXKyo8azlKqSyl1PuRz7Ak8r57jX2zlVJzIp9LiVLqbaWUm/D/cdfIZ1yplOqqlLpbKfV8ZL+PlFI31UnHSqXU2c15rUI0RPInyZ8i+0n+JNoUyZskb4rsJ3mThXYXTCmlegAzgO9rLD4LGA8MVUqNAp4GrgNygMeBdyM/WCfhWaafIzzT9GvAOfWcxwa8D2wDegPdgJe11uuA64GFWutUrXVmZJf7gIHASKB/ZPvfR451MnA7cAIwgPCs1Yn6KHKMjsBywjN913QR8GegA7Bi//rIj+wzwrNmdwQuAGYrpYbWc/2lSqnJcaRrJnAy0AcYAVweWW4Ac4BeQE/CM5k/UmO/54AUYFgkXQ9orauAU4D8yGecqrXOr3O+l4ALa6R3aOQcH8R7rUI0FcmfJH+KkPxJtCmSN0neFCF5kxWt9WH/AvKASqCU8A90NpAcWaeBY2ts+yjw5zr7/whMBY4B8gFVY90C4J7I+2nAzsj7CUAhYLdIz+XAtzX+VkAV0K/GsgnA1sj7p4H7aqwbGEl3/3qu92vg6hg+l8zIcTIifz9DONPavz4VCAE9gPOBeXX2fxz4Q41974nx/6Pu9ecBF9f4++/AY/XsOxIoibzvAphAlsV2B/4vaiy7G3g+8j4t8pn3ivz9F+DpyPsGr1Ve8mrKl+RP9X4ukj9J/iSvVnxJ3lTv5yJ5k+RNtV7tqZ3rWVrrz+tZt6PG+17AZUqpm2sscwJdCf94dunINyRiWz3H7AFs01oHY0hbLuESgmVKqf3LFGCLvO8KLIvhnA2KlPj8BTgvck4zsqoDUBZ5f+Cz0FpXKqX2Rc7fCxi/v2o9wk64dKMpFNR4Xx05J0qpFOABwiUvWZH1aZFr6QHs01qXxHsyrXWFUuoDwiUnfyNc0nJNZHVzX6sQdUn+JPnTAZI/iTZE8ibJmw6QvMlaewqmGlLzB74D+IvW+i91N1JKTQW6KaVUjUyhJ7DZ4pg7gJ4fkcTfAAAgAElEQVRKKbtFpqDr/F1EuAp2mA63S65rN+Ev/34967+UBs0CziRc1Z0HZAAlhDOf/Q6cR4VH7ckmXKK0A5irtT4hwXMn6lfAIGC81rpAKTWScDMDFUlTtlIqU2tdWme/up+xlZeAPyilviHcfvuryPLWulYhrEj+9BPJnyR/Em2H5E0/kbypHedN7a7PVAyeAK5XSo1XYW6l1KlKqTRgIRAEfqGUckQ63I2r5zhLCP+Q74scw6WUmhRZtwfoHmlHjNbajJz3AaVURwClVDel1EmR7V8FLldKDY2UNvwhhuuwR865/+UgXD3rA4oJl+bca7HfDKXU5Eja/gws0lrvINyGeaBS6pLItTuUUkepcKfQ5pRGOLMsVUplU+Patda7Cbdjnq3CnS0dSqljIqv3ADlKqYwGjv0h4ZKUPwGvRP4foPWuVYjGSP4k+ZPkT6ItkrxJ8qZ2mzdJMFWH1nop4SrLRwiXPGwi0qFPa+0Hzo78vY9w+9A36zlOCDidcIfI7cDOyPYAXwJrgAKlVFFk2W8i51qklCoHPidcqoDW+iPgX5H9NkX+bcyjhH9I+19zgGcJV3PvAtYCiyz2e5Hwj24fMAa4OJKGCuBEwlW7+YSrlv8GJFmdXIVHgZkSQzob8y8gmXAJ1CLg4zrrLwECwHpgL3BrJL3rCZeebFHhDp1d6x5Ya+0j/P93POHr3r88rmsVoqVI/iT5k+RPoi2SvEnypvacN6naTViFEEIIIYQQQsRCaqaEEEIIIYQQIgESTAkhhBBCCCFEAiSYEkIIIYQQQogESDAlhBBCCCGEEAlocJ6pDh066N69e7dQUoQQLWHZsmVFWuvc1k7HwZL8SYjDz+GQP0neJMThp6G8qcFgqnfv3ixdurR5UiWEaBVKqYRmgW9rJH8S4vBzOORPkjcJcfhpKG+SZn5CCCGEEEIIkQAJpoQQQgghhBAiARJMCSGEEEIIIUQCJJgSQgghhBBCiARIMCWEEEIIIYQQCZBgSgghhBBCCCESIMGUEEIIIYQQQiRAgikhhBBCCCGESIAEU0IIIYQQQgiRAHtrJ0CIQ8G24irmzM/jkzUFVHqDAKS57MwY0YXLJvSmR3ZKK6dQiOa1t3ovr6x/hQ+2fkCZrwyNxu1wM6XbFC4deil9M/u2dhKFOGx4AyHeW5nPnPl55Jd58AdNkh02BndO45pj+nLMgFwMQ7V2MoUFrTVLCpbw1qa3KKgqIBAKkJGUwcSuEzmz/5mkOdNaO4miiSmtdb0rx44dq5cuXdqCyRGibfmxoILfvf0Dq3aWYWpNIFT79+KwKQylGNUzk3vOGk7/jm0/k1RKLdNaj23tdBwsyZ9axq7KXdyz6B6WFCwBDX7TX2u9XdmxGTb6Z/bnrvF3cUTuEa2UUnE4OBzyp4PJm3zBEH//+EdeWrIdBVT5Q1HbuJ02Upx2bjthALPG9zrI1IqmEjADvLL+FZ5Z8wwV/gqqg9W11ifbkzG1yYm9TuT6I6+nZ3rPVkqpSERDeZM08xOiHgs2FfGz2fP5Lq8EX9CMCqQAAiGNL2iyeMs+znxkPku27muFlArRPNYVr+O8985jQf4C/CF/VCAFENRBfCEfa4rXcOUnV/Ll9i9bIaVCHPoqvAFmPraQFxZto9ofsgykIBxgFVb6+PP767jzzR9oqFBctIyqQBVXf3I1Dy5/kD3Ve6ICKQBP0IMv5OODLR9w3nvn8V3Bd62QUtEcJJgSwsIPO8u46r9Lqa7nZlaXJnyDu3zOEtYXlDdv4oRoATvKd3DlJ1dS4a/A1GZM+3hDXn79za/lIUGIOPmDJpc9vYR1u8vxBmP7vXkCId76fhd//XB9M6dONMQf8nP1p1ezumg13pC30e1NTKqD1fz885/zQ+EPLZBC0dykz5QQdZim5ppnl+IJRAdSVWu/pvy7twkU78RwJuPo2JeMiTNxdR8GQLU/xLXPLmPuHdNQKvH27FW+IKWeAACZyQ7cSfJTFS3rjm/uoDoQXbpaMq+Eok+K8O/1Y3PZSB+TTqdzO2Fz2wDwhXzc8tUtzJ05F4fN0dLJFuKQ9PT8razdXY7fogVEQ/cdTyDEc4u2cdLwzozpldUKKRf/WvYvNpVsiqq5byyv9Ia8XP/59Xxx3he47K7WSLpoIvKEJkQd8zcXUeENRC0vX/IWZYtfJ+fEG3H1GY2y2fFsXYZn4+IDwRRAUaWPpdtKOKp3dlznDYRMPlu7h8fmbmZtfjkOm3Fg+aDOaVw/tR8nDeuM0y4VyuLgbSrZxHPrnuOr7V9RFag6MKDE9B7TmdZ9GptKN2FSu4S86KMiCj8qpPvV3UkdmkqgJED+c/nk3Z9Hn7v6YES+myEzxBfbv+DkPie3xqUJcUgxTc2T87bgDUTXSMVy3/EGQzzxzWbGXHJIdzU7JHmCHl7f+HpUjVSseWXQDPJJ3iec2f/M1ki+aCLyVCYOSQVlXj5bu4c3l+/k/VX5fJe3j5DZNO3GH5+7JaqtuumrovTbF8g+4QZSBk3EcLpQNjsp/ceTNf3KWtt6/CH+882WuM758eoCxtzzGXe8vpJVO8sImhpPIIQnECJoatbkl/PbN1cx5p7PeH9V/kFfo2i/1hWvY+Z7M7nggwt4Z9M7lPhK8Jt+AmaAUl8p725+l9u+vg1fyFdrv5AnxN6399L14q6kjUhD2RXOXCc9ft4Df5GfsgVlB7atDlbz1OqnWvrShDgkfbOxEI9Fk/JY7ztaw5c/FlJc6Ys6hmheH2/9GEXtViiSV7Y/UjMlDhmmqVmwuZjHv9nMkq37cNoMTK1RKtxnyWW3cdXkPlwwrifZbmdC56jyBVm8tThquW/XenTQT8rACY0eQwNfrd+LLxgiyW5rdPv/Lsjjrx+tsyyVrJ228M329tdWsrvMwzVT+jV6bCFqmr9rPrd+fSveYP3t+kPaup9g9cZqzIBJ+pj0WsttLhtpI9KoXFNJ1jE/NTPaWraVgqoCOrs7N03ihThMvbxkh+VgE/Hcd2wKPl5TwEUyul+Lenbts1GDTcSbVxZUFbChZAMDswa2SJpF05OaKXFIKK32c9a/53Ptc0uZt7EIX9Ckwhekyh+i0heiyheiuMrPQ19uZOJ9X/DeysRqb/ZV+Q80r6sp5CnHSElHGY0HRwB2m6KsOrqpYF0fry6IKZCqyRsw+eenG3h3xa6Y9xHih8IfuPWrhgOphoQqQ9hT7ShbdF9Ae4adYGWw1jKH4aDYE10wIYSoLb/MY7k8nvuOJ2Cytzyx37ZIXEFVQdSyePNKm7Kxq0Lu54cyqZkSbV5JlZ/TH/mWPeVey+HJa9oflNzx+koqvEFmjY9vHgdf0MRq2AhbcjpmdTnaDMV0YzOUwtfIiEyBkMmvX19pGUg1NtCFN2By51urOWl455hqv0T7ZmozXCNlMdJUY52k97Ol2ghWBtEhHfWQECwLYk+Nvp1YDaUuhKjNX8+9It77jieOQjnRNKzyuHjzSlObeILWAbU4NEjNlGjTTFNzydOLYwqkavIGTP70/hrmbyqKeR9tmrj27CIQCEatS+o2GGV3UL1hYUzH8vsC+J96nLL3P8C3aRM6GH3Mz9buIWQxP0j5krfY98UTZBw9k+43PU+3G+aQNnoGno2La6dXaz76IbpUTIi6Fu1eRKW/Mmp50UdFFLxWQOeZnRk6eyh9/7cv/mI/effnYdZ5wEvpn4KyK8qX1R76P+QNUbGqAvdQd63lWmvSHG1/EmshWlt6svWol/Hcd+yGIislsebtInEuW/QofPHmlYYySHWmNms6RfOSminRps3dWMjWwirLQCqW2pt7PljLR7ccY3nsUGUl3lWrqF6xAs/3K/CsWoVKT8c9+gb8qvZNyUhykzn5IvZ99hjKsOHqMwpl2PHmrcC7fVXUIBQ5LhtpKUlUfPophQ8/RHBvIUl9+5I0eBCuQYNJGjyIRxd4DvSD2m9/h+OcGbeSMmjigeUp/ceT0n987ev3h3h07mbOGtUt9g9UtEvPrH4mql3//k7S3a7qRtqIcNCzv5P0hjs2ULagrFa7fluKjY5ndST/+XwMl1FrhCpHtoPMiZm1jq/R9EyPr2ZYiPZoyoAOrNxRGtWaIZ77jtNuMLpnZt1Di2bWJ6MPKwtX1loWb14ZMAP0yejTkskWTUyCKdGmPT53s2XH3FiHKd9aVMW63eUM7pxGYNu2nwKnFSvw79iBa+gQUkaOJOv8mXT9673YO3Tgmq828dAXG6MmTkwfdzaGO4uyha9Q9P79KGcySZ36kz7h/FrbJTtsXH/CIDpOnnFgmVlVhXfDBnw//oh3/XoKPvqUtT3OgzpNN+LpcAywpbCS0mo/mVIiKerhCXosJ9GNt5M0QO6MXGxuGwWvFODf68dINkgfnU6P63pgOH5q6OAwHJw78FycNvleCtGYC8f15JEvN1mui/W+k5XiZFyf+KbjEAfv8mGXc9e3d0UVVsWaVwIMzh5Mj7QeLZls0cQkmBJt1o591Xy/vTRqeTy1N4FAiAf/8SK/XPgsKtlFysiRJI8cRea55+AaNAjljH7Yu2BcT/71xUbLNKUOm07qsOkNplujOWds91rLDLeblFGjSBk1KnwNpR6c/5wbNTFwvANdOGwGpdUBCaZEvcp8ZdgNO8FQ7aamjXWS9myzbsOfPTWb7KkNP7QppZg1ZFbiiRaiHemQmsTUQbl8tnYPFi2/G73vJDtsXD+170FNFC8SM63HNOyG9aN0LHml2+7myuFXNriNaPskmBJt1oodpdhtCl+d7kbx1N6EUKxK60Gft9/C0alTTOfNdju55OhevLh4e1Sw05hkh40rJ/cm3RXdBj4YMlm3u4KSaj9FlT5Mi3mx4u1wLERjAmbA8iErkQElYuGyuTih1wl0S5Xmp0LE6lcnDGLehqK47zmGgoxkB2eP7t74xqLJ2Q071xxxDbNXzMYTim8QCUMZZCRlcEx3664I4tAhwZRos8q9AcuJeOOtvam2OWIOpPa7a8YQ8oqrWLCpKOYRkpIdNqYPzuX2EwfVWr63wssLi7bz3wV5BEwTQym0Bl8o+rg1Oxy7B09u9JyBkElminXnZSEA0p3pBM3oAVBqdpLOGJdxYPn+TtKdzo3vNwPhQGpozlD+OOmPB5VmIdqbQZ3TmH3RaG54YVnMU2XYlCLVZefV6ybgTpLHudZy2bDLWF28mrk75lqOmGrFwCDVkcpTJz1Vb82WOHTIaH6izXLajKiZxaF27U0srOaNaoxhKP5zyVjOHNkNl8PAoiXUAXZD4XIYnDe2O49cOPpALYDWmke+3MiUv33FY3M3U+oJUOULUeENUlm3um3/eWt0OK7esBAz4EWHgng2L6Xkq6ejtu+d45YmfqJB6c50cpJzopbX7CRdsaoCHdT4C/3smL3DspN0Qwxl4LK5mNJ9Ck+c+AQO49AL8L2BEDtLqtmwp4JdpR58wfhqCIQ4WNMHd+Tpy4/CnWQj2dFwYaE7yUaXTBfv3zyZnjkpLZRCYUUpxd+m/I0ZfWeQbE9udHsXNnKSc3jp1JfoniY1iocDCYdFm5WbloTNiI5i4q296ZCaWLBhMxT3nTOCKyb14alvt/DuinzsNgMdadSulCJkas4a1Y2rJvemf8efhoHWWvP7d9bw+rKdjc43VVesHY7dThvXT+2X0LWJ9kMpxeXDLufB5Q9GzWUSayfpJFsSw3OGs7ZoLYZWmKEgKDAMOwEV5Lhex3HZsMsYljOs7unbvPUF5Tw5byvvr8xHKYVhgBn5yZ4zphtXTOpDv1wZtli0jIn9OrDof47jzeU7efybLZRWBzAUB+4jhlIM7pzGDdP6cdyQTgkVFoqmZzNs3D3hbk7qfRJzVs9h+d7loH+ah0qhSLYnk2pP5tK9uzj71FdJk9FODxtKW/V2jBg7dqxeunRpCyZHiJ/4gyaj//yZZS1O+ZI3KVv8Jjkn3djgcLEpTht3zhjMxUf3Puj0VPqCfJe3j7LqAACZKQ6O6p1t2bziyXlb+OenG+Ju/x6PFKeN5f97Aq5GSjDrUkot01qPbaZktRjJn2JX4a9g+qvT8YV8ce/rMpOYUTaFa6vPx+fz8IN9A+W2KkxMUnUKw3z9yeqUS/rUHriG5qAsCkDaoqJKH9c8u5R1u8sJhLRlk2K7obAbijG9s5h90Rgy6pkPSDSdwyF/aqq8SWvNyp1l7CrxMH9TIT/uqeQf546grwT3bV5BVQGf5H3Cnuo9+II+sl3ZjO40mqO7HI36/G7w7IMzHm7tZIo4NJQ3Sc2UaLOcdoOLj+7J099uxV9nnqlYa29MrTlr1E/V6N5AiPdX7ealJdspqvQRMjXpLjvTBnXk0gm96ZwRPQHffqlJdqYP6thouqv9wXoDqcbmxopVssPgz2cOjzuQEu1TmjONq4ZfxZw1c6JqpxqSG8jiH9t+RY6ZiQ4FcGJnjG9o1HaBHZXse/VHkvplkjNrCMrRtkvLd5V6OOvf8ymp8hO0CKL2C5qaoKn5bus+Tn1oHu/cOImc1KQWTKloz5RSjOyRycgemWQkO9hatEkCqUNEZ3dnLht2mfXKybfCw2Ngws2QO7BlEyaahQRTokUVV/rYXebFEwiRmmSnR3YKqQ10nL1kQm/mzNtiua6x4WIdNsWZI7uRmmSn3Bvg/z7dwKtLd6Agau6qTYVVPPXtVsb3yeZ/ZgxhSJd064PG4L2V+ViNUBvr3FiNcTkMbjp2AOeMkbbWInbXH3k92yu288W2L2IadSorkM5DW39DesiNEUP3Wu038W4soeiZ1XS48gjLIdcPVmGFjxcXb+ON5bso9fgxzXDfkWMG5HL1lL4M6pzW6DHKvQHOf3wh+yr9hBpomVGTP6QpKPMy64nFvHPTJCnEEC0u2+1kX5W/tZMhmkJyFky8Gb66B2Y+29qpEU1AginR7ExTM39zEY/P3cKSvH3hgSUUaB0eje7UI7pw9ZS+DO1aO4DRoRDOZ5/g2o2beWLgiXjjaDFnMxSdM1zcdeoQCsq8zHx8IQVlnqgarv38kfbo32ws4rvZC5h98eiYaqGsPPr1ZqrrBGvxzI1VH7fThgb+fOZwCaRE3JRS3Dv5Xh5IfoAX17+IqU0CZsByW4dy8Ncdt5BmpkYFUkt2ruLerx5lQ1EehmEwIKcXfzjuZkZ2GQJBjX97BWUfbSXztL5NlvaCMi+/e/sH5m0sAqjVD7HSF+TN5Tt5b1U+fTuk8qczhzG2d/1zu8z5No/CCp9lINVQzXHQ1GzfV80by3Zy0dG9muzahIhFh1QnxRJMHT7GXReundq1DLqNae3UiIMkwZRoVpsLK7n0qSWUVvsP1Ab56wzI8M6KfD5cvZsR3TN54tKxZCQ7CO7bR/7td6CDQW6ZfT+u1WU8/OXGmIaMddoMOqUn8dp1E9EmnPPoAgrKvZZ9Iqx4AiFueH4Z/71iHOP7Ro+C1pAKb4CdJdGl/vHMjbVfssNAKUUgZNI7x831U/tx6oguUiouEqaU4pdjf8n5g8/nxXUv8sbGN8LLa4yaaWqTG7KuoNfGrtSdQbTCV8UVr/+Wv5z4S04fPB1/KMiSnStJsv00yIsOmFQt3k36Cb0wkg7+u7phTwXnP76Qck+AespCCGkIBUzW7i7n4qcW849zR3D6kdHzXIVMzZwFWy0HhYml5tgTCPHY3M3MGt9TJkgVLSrL7aS02o9paoxDpF+iaIAzBab+Gj7/I1z2bmunRhwkCaZEs1mTX8b5jy+iyh+0nNV9v5DWhAKaFdtLOPWhebx8TDqeO39NxmmnknvLLSi7nRun59Iv181fPlhHcZUfTyAUdcxkhw1Ta04b0YU/nDGMdJeDnz+/jL0V1oFUQ6XQ3oDJ1c8u5bu7jo8reCnzBHDYDIJ1hm2Pd26sFKeNRy8aTa8cN5kpDhn+XDSpbqnduOOoO7hl9C2sKlxFmb8MgAxnBkfkHkHFsxvxBkqi9tuybwcAZw09HoBkw8bUPuOiT6CgesVeUsd3Oah05pd6mPn4QkqrrWvQrHgDJne8vor0ZCdTB+bWWvfl+r0ELOZ3i6fmuLjKz7JtJQ3WfgnR1Bw2gxSnjTJPgCy33A8OC6MuhgUPw+Yvod+xrZ0acRAkmBLNYneZh1lPLK53PiUr/pCmoKSaS17I4/W77iLn+ONqrT95eBdOGtaZZdtK+M83W1i+vYRqfwiHzSDH7eSSo3txztjupLvCI24VVfr4Yv1eAhbF2bGUQpum5oNVu2NqUqe1JlRaSmDNZnQoCHXmx6o5N1asAVXvDm565bhj2laIRDhtTsZ2rj04UajSj3dzqeX2fbN7YCiD2z74C2cMPo5R3YaR6Yrup6T9JhXf7DzoYOrWV1ZQ4bXOQxorDPn588tY+rsTSHb+9Hv7ePVuqnzR7YXjqTn2+EN8tnaPBFOixXVITaK4yi/BVFtnhmDT51C0AXyV4HRDdh8YcBLYa/zf2Rxw7O/CtVN9poHRtgfuEfWTYEo0i4e+2ESlz7o0ucF+CSj2ZHZhbvYgzrbYVynF2N7ZMT3IvLRku+XyWEuhq/whHv16U61gyvR48G/fjn/rVvx5efi35uHPy8OXlxeenKZvPwK9Z1F3BIp458byB025YYpWESz2ouwG2mLS2rQkN29e9AizF7/Irz/+B4VV+5jebzx/P/nX5Lpr/yZDJfEPw17TtuIqVu4otaxVjqUwRAPvrcpn5tgeB/YrrLROUzw1x5rwQBiNCYZMFm4ppqDMiy9okuayM7xbhsxZJRImg1C0cZWFsPRpWPIYBP0QirwMB9iTQBkw9ioYdzVkRJ4rhp4F8x+Ede/AsJ+1bvpFwiSYEk2uyhfk7e93YdGaJqaHoOqgyaNzN3P26IMbZOGFRdst+0bEUwq9s6iSRXf/ja7bf8S3NY/Qvn04enQnqU8fnL17kzLuKDJnnoezTx9sWVkopTj6yUXM31Rc6zhGkpvMyRex77PHUIatwbmxAI7skXmghk2IlqQtam5qGtChNw+ceicAm4q38Yv37+HuLx7m32f8ofZxTM2Pu8rJTk8iK8WBPc7JRZ9ZkIdp0T441sKQan+Ix77efCCYMv1+zKpqy3PFXXPcQJeVveVenl+8jf8u2EbI1JhaY5oam00RCmkGdk7jhqn9OH6oTLgq4pPtdlJcT4GAaGV58+HFmWAGIeitvc4MgD9SuLzo3+Fg69w5MOiUcG3U8X+AD26HwaeFa6vEIUeCKdHk3l6xy3Jo8Hj6Jezc52H1rjKGd8tIOB3FVQdfCm3XIUo69mT4cRNx9umDo0sXlK3h/a6f2o/vt5dGjegX69xYbqeN66f2azRtQjQH5Yz9Ab9/Ti9mDj+Z51dEd6A2gRtf/p6Saj+lngCpSXZy3E6y3E6y3U6yU5xkpzrDyyLvs1Mi69xO3ly+y7KJbjyFIbuKKph/wy/ptHkNwYIC3OMvhpzoebLiqTlWQKc06/no3l2xi1+/sQqtiS7IibRWXLWzjNtfW0luWhIvXzuhwbnthKgpJ9LMT7QxW78JB1KBGObwC/khBLx2BZz9OAw9M9xfKqM7fP88jL2i2ZMrmp4EU6LJfby6ICqQgPgeggIhk3nr9zAkFczq6p9eVTXeV1cdeK8t1gc7nBXV3A7iK4VWSS4c044mdWinmK9/Ur8OpLnslp9BY3NjASTZbRw7OLFh2YU4WPYsF9qiRhfCNVFfbF7IGYOPpUt6R/LL9/DOui8Y3TV6njRnhpPPfxUuJAmZmnJPgOIqP/tqvEqq/eSXelmTX05xlZ+SyPKiSi++oPWoNfEUhjgMCBx3Cj3uuAln9+5csKWUb19cHjXPXDw1xy674sRh0fnBi4u386f318Q04miVP4S3xMOpD83jg19MkYBKxCRHmvm1PaXb4aULYgukagp64K3rIKc/dBoWrp16+WIYcX54pD9xSJFgSjS5kmrrzD6eh6Cgqdn46JNs2vY1RkpK9Mudgqr5d0YG9i5dMFLcB5YlfVxmOTdVXP2XFKS54vuZGIZi9kWjufjJxXhieLCqyeUweOSiUdhk6FvRSmwZSTi7p+HPK49a53amsCJ/HU989yrlvkrSk1I5vt8E7pr+81rbKYeBe9JPQ5PbDEVWpFYqFsGQyYC7PsIqnIqrMMTpxHbkSJL6hkf1O2ZgLi6HLSqYgthrjrMrisn8zY2UXXQR6SeegHI6WbC5KOZAar+QqSn1BLjgPwv59LapOO3S5E80LNvtZFtxVWsnQ9S0cDaEavcP7/2vCqoDsPWWVNzO8L38yeV+nl8V4OvLawwqFfTBN/fDeXPCc011HwtLHofJt7XkFYgmIMGUaHJGPfOvxNsvodNVVzDolL8nnI5BK79l5c6y6PTFUQrtC5oM6BQ9WlljxvTK5pFZo7nxxeUxP2C5HAb/N/NIJvbrEPf5hGgKWmu+3VTEp1UVXKo0ybr2b7lLWi6PnvXHmI6TOjb22ty67DYDu01ZNvOLqzBEU6vvoc1QXDW5Dw99sRGvRe1bYzXHLrvBzRdOIaeiMyUvvMiev91H1nnncW/oiHp/5w0NuBMyNXsrfHyypoDTj+za8LWIdi8n1cny7dFTFohWEvDC98+Gm+7VEdLw4GI/d05Jqn9/bcKPH0L1PkjJhuN+D0+fBGMuh+Ss5ku3aHJSFCaaXMc068yj5kNQY5x2g5zUBjKhGFw/tR/ueiYNTR93NlnHXkXZwlfY+fBF7Hz0ciqWv0/ygJ+aIBoKjh3UkewER9U7bkgnXr52AgM6ppLssGGzCDINFZ4fq28HN89dNZ4ZR8gDlYhDeT6sfReWPwerXoUtX0eVksZqweYiZj6+kD+8s4ax03vjTk9qcKCFetkVKSNyMVIOriP1sK7W/SVrFoZUb1iIGfCiQ0E8m5dS8tXTtbYNmCYDOtUePe+yib3pkpmMPYHaX1/IZGV+OUUjJ9DrmTn0mjOHTft8bMiPLrSB8IA7+754goyjZ9L9piPdFTsAACAASURBVOfpdsMc0kbPwLNx8YFtqv0hHp27Oe60iPYnx50kzfzakrXvUF8mecdEJ/cv8FHqbWCSTQjvv+KF8NsOA2DwqfDtv5o0maL5Sc2UaHKnH9mVeRuLojpgx1MjpIATh3Y+qHQcP7QTdsMg3NszWqOl0A4b1xzT96DSMLJHJp/9ciqrd5XxxLwtfLFuL9X+cE/0ZIeN6YM7cs2UvhzZI/OgziPaEa3DQdP8B2HbArA5wyWcSkVeBhx1LRx1FaQ3Ps/Toi3FPPDZBvaUe/nFcQM448iu2G0GgR7Z7H1kRaOj+9ViU9hzksk8q3/i1xdx/dR+/Oq1FZbzQsXSJM+mTY7buZLqp/JwnHMOjo7hfojuJDsvX3s0P/v3fAorfZa1X/XRGl5bupO3lu/irlOHcMmE/rw74mRC322nbpvEeAbc2VJYyY8FFQzqHH8tuGg/wqP5STDVZuR/D/5Ky1Vju9qY1tvO/Qt83HNsA30igx7YsRi4Ofz31N/CY5Ng/HWQLoWrhwoJpkST2bGvmheXbOfV77YTtJgbBmLvl3Bkj0x65hxcJ0yHzeCW4wbwj09+xBOI44EQcNoUgzqlMbpn0wQ5w7tl8OAFo4BwEygIz5klRFw8pfD8OVC4DvyRvhMhi1ErFzwECx+GU/4WbjJi4bu8fTzw2QZ2lni4+dj+/GxUt1rDlztyU8i9bgSFT/wQDqjq+U3/tIOBo2MKuVcNx3DGNjF1Q44f0hHHQRSGOBx2fn7NDAKfvM2W007HPWECWRecT8rRR9Mp3cWHt0zhxheXszSvhFDIpJ7xLqIETU3Q1Nz74Xqq/SGWbC3GKh6LZ8AdhWLljlIJpkSDclKdMppfW1Jd3ODqP01PYtLTVdwyvpHWLZ4ak6RndINRl8Dcv8PpUkN1qJBgqp0o9wZ4Y+lOvtlYSGl1AIfdoGtGMjPHdmdCv5yEH+xNUzN3YyHPL9zG8u0lnD26O69dP5E3lu/kyXlbLed5auwhKMVp47qDrBHa74pJvflxTwXvrsiPOaBy2BQd0108c8W4Zgl4JIhqX0KmZsOeCkqq/SgUWW4HAzqmxT/IiKcU/jMNyndZttGvfdJIgPXxb8FbBpNuObBq2bYS/vX5BrYWVfGLYwfws9Hd6p3vyNk1lc63jab86x1UL90DgPbX/k0rp4FKspN2TDdSj+6KcjRN63G7zeA3Jw/iT++vi7swxGU3mDqoI8MmjoSJI+l4x+2Uvfsue+79KzoQIPP888k460xeuPpo1u8u52ezFxCM8xyeQIgHPt+Ay2EdOMY34I5JuTex5pmi/chKcVJa7cc0NYYMUtT6khqegHt4RxunDbRz37d+huQ2kC863bX/nnwbPDwGJtwEHQ6+ll80PwmmDnN5RVU89MVGPvhhN4ZSUQ8ln64tID3ZwXVT+nLxhF4xTyK5r8rPa0t38MLi7aQn27n06N48Mms0yZES6Rum9eeDVbvZUeIh1FiJdg0uu8Gk/h2YPqhphgZXSnHf2UeQ4XLw7KI8AkFNyGIi0P1SnDZ657h54erxZBxknw/RvhVV+nhx8XbmzN+KP2geePgxTU2Sw8aVk3tz4VE9Y+sbqDW8cG5sgVRNAQ989VfI6c/3KRN54PONbN5byU3H9uec0d1jGkHOlp5E1hn9yTylD9UrC/Gs34dZFUAZClt6EimjO5LULxPVDA93F47vxeaiKl5YtD3mgCrJbjCwUxoPXjDyp2tITSV71iyyLrwQz/crKHn5JYpmzybt2GP5btxplvPiQcODRwB4A2a9+Vs8A+7YlCKpkf8L09QUlHsp8wRw2Aw6pDrJTEmsP6c4NDntBilOG2WeQMwjY4pmlN0XbEnWrQMi/jjNxejHK/nVhHryeWWHnDrzSqZkw4Qb4at74Lxn0FqzubCS4ko/poaMZAcDOqXKpN9tiARTh7FFW4q56pnv8ARC9bbQqfaHqPaH+Psn6/lw9W7mXDGO1CTrr4XWmhU7Snlu0TY+X7uHE4Z25qELR3Fk94yo2pbUJDsvXzuBcx6dz96K2PolJDtsHNkjg0dmjWrSUjelFHeeOoQzRnblqW+38uEPu7EbCn/IROtwc0BTawZ3TuOGaf04fkinWs2dhIiH1pp/f7WJh7/cBFhM3kp4nqGHv9jEQ19s4pcn/D975x0mRZX14beqOvfkCAzDBILkHIdsFrMuJkRJCkZQ191v1V3TuqsrrnlFRFERXdacMaBIBhUByTCRgYHJuWN1fX80aaa7Z7phEnDf55kHpruq+nZPV9X93XPO73Rl5pjODUcsc1ZC4fY6Qip4+10bhR/cy+3yy9xxdlcm3jQIoy70NDxJr2Ad3A7r4JOrZQyVByf0IMZq4Pnv9+DRtIDXElny9mgb0TmW/0wa6DdiJEkSloEDsAwcgLusjPIPP+aVlbnUmnzTeSs3fEzF+g+IPf8OTGkDkRQdtuxfse1Zf1RMAQFTmkNxHdQdjob7o7TGyXsb8nhjVTY1Tjc6WUbTNJyqh94dIpk1rjPndE8Q16wzhCONe4WYagP0mQjLHm9wky4xMtf20vPCBid9Evyco4oOBt7s+/jw26h4LoMPv1rJa5udlNe60CnHFuQUWeLG4SlMHpFC+0hzU7wbwUkgxNRpyqZ95Uxd+HPQq7k2l4fN+8q5ccF6/jdzRJ0Va5tT5fPNB1i0LpcKm4sbh3firxf3bPRi3i7SxFd3j+GeJZtYlVkMGjhV34mlWa/g0TSuHZLMQxf3aLZJQe+kSJ69tj+PXNaLH3cWUlztwO3RiDTrGZIaQ5eEhkP2AkFjaJrGXz/Zyocb9/sVUcdz5Pnnv9/LoUoHD1/q2/j2KKufB2etz8NB2e8C0VoFP11vwZCW0vibaGNIksTt47pwad8OvL02h/c27ANAQ0PTvFEdp+rhnB5eM5f+yVFBpdLqoqPJv+AqynLXQb3rZCjmERKgyFD/zx2K4Y6mwdhu8XXH4NF4aulO3lyTgwTHWbkfe6Hf9pVz75JNGHQyr9w4iOHpsY2+b8GpTYxFz/4yG2lxVtGPsLUJbwfp42DPt/g40BzH38YaWbQlQBpvQg+IP8vn4SWbS3m4/AmklSXYtMNZMvUO8fqqbF5flc3kESk8cFEPkfrZiggxdRpid6nc/MYGv0KqobQVp6qxs6CSp5bu5K+X9CSrqJrF6/P4aGM+g1Kiue/8bozpGh/SCRtp0fPG1CEcKLfxzrpcFq/Po9LuQidLuD0aHSJN3DI6nasGdazTD6Y5iTTruWJAUuMbCgQhMn9FFh9u3B9SjY/NpfLfDftIjrYwbVSa7wZVByF7Jf5u1vdnGPjXage3DzEQZQp8XupVO6x/GdIyAm7T1kmOsfDgxT354wVnsSG7lNIaJ27VuxgyMCX6hFoYbD9Qiebncw3FPMKjBZ5GBWO4o/O4uUwuQVddCdHe3jIej8Yd725k+a6iRkV5jVOlxqkyZeEGnr+2Pxf0btzBUXBqUVHr4n+/7OP1VdkcrLQz9c0NaECs1cDNGalcP7QTcSfZSkRwgoy6B3JWeFOqD5Mzp66RTHKkjP2hCN999RYYfZ/Pwy//uJeXftiL3SMBgedFR64Ni9flcaDcxkvXDxSCqpUQYuo05PPNB3D5iQAFk7Zid3tYtC6XHQWV7D5UxTWDk/nszlEkx5ycs16HKDN/urA7f7qwOy7Vg82lYjXoxMqa4LSh2uHm2e93+23e2ljtjc2l8vQ3u7h+aKejdYdHObAJdAa/eflB2++iHbbfPfUx6hRGd41vfMMgqLK7cPtJGwzFPALAoldwe7QTMtxR9Hqu9uwn66IJRN80mZibbuaxH7JZvqsoJFFud3mYs2QTiyNMDOwkGn6eDjjcKg9/uo2Pf9uPJHH02nIks7S42slLP3hTii/olchTV/fFYhDTuhYlZQRk3A1rXgSXb/ZAQPRm6HstdL+kzsOfbzrAiz/sCdgE3B82l8qPO4t48uudPHBxj+DHIGgyxFl3GjLvp0xqnSeetuJWPaTEWFg4dcgJ1VY0hl6RReGk4LTj4435yH7Sy4KtvZEk+Gzzfq4d0qnuAewV3j5SAQjafveIlbrgKCa9gnI4Sn48oZhHAMSFG7HoFfYWVYfUt8qsV3hgQncGj5iAc/oNFL34Eisvv5b3hs7Cqfl+lxoX5R4e+Oh3ls4ZE/QYBG2TGoebSQvWsfNgVYPRySPPfbvtEDsLVvP+rBHCmKSlGfcXcDtgw/zgBJXeAr2ugouf4Xj3G9Wj8bfPtgYUUg2d/zaXyptrc5gxJo2E8IYW1gTNgZjRnmbsPFjJgXK7z+Ohpq1s3FfeLEJKIDgd0TSNV3/KCriIEXPebVjOykA2mJAUHZYuw+rUzIDXDOaV5ZlomkZFrYst+eV8vvkAS3eVYm+gCdLx9rsNogh3yvp0iDIfLeo+nuPNI4IhOdrCu7cMp3N8WKOufEcw6WXuOqcLk0ekAmBISSFp7tMsu+n/8Pgxtajc8DGly14jcvg1dLzzHZJuW0j4wAnY9tSNOOaW1LD9QGVQYxC0Tdyqhxlv/cKOgqqgIxQOt4eckhpufH099hBt/gUniSTBeY/CVa9BQk/QmUGqN3+SZK+Iik6DS56Fy1+Cegs1P+ws9FtXDsGd/xLwzrrcpn53giAQkanTjLySWu/koF6hYqhpKwXltsY3EggEAJTVuiis8k3DC2URAyCnpJZ+j36LR4NOMRZSYi1kGKJQGjFUaNR+F8ASF9QYziTGdItDwvezDcU8wmpUuDkjlWirgU/uGMnT3+zivQ15SHjrmeoc97DjYIcoE3+5qAfn9kys87zdpfJhZjVuqa4gCyWzwKlqvL4qi2eu6Y/g1OSLLQVszi/37wTaQHTCpWpkFlbz3vo8pvqrvxQ0Lz0u8f4UbIH1r1KU+Suys5rY6BiI6wrDZkHHwQTqxTDvp0xqHL5CONjz3+H28NaaXO4+u6tw92xhhJg6zfDaoPuuaoaattJY0bNAcLqzJb+ctZkllNY4MSgyceFGzu+V6NeGtrzWiV6RqDd3DnkRw6iTWTxjGL2Tjms34OkPcx+E2sBpeo3a7+rNMHhqUGM4kzDqFCYN68Qbq7N90vOCMY8A0LmcjEuLBLxpg3+9pCf3X3AWX24p4M01ORRW2XG4PYQZdQxIjmLG6HT6JftasQNs3lfuN1U0FFGuejSW7SgM9iMQtEFe8ZOqD8GlDNtcHuavzGLKyFTRIL61aN8XrniZRd/tRpZgzrndGt1F0zR+yyvz+1wo579L9UYouySEN7qtoOkQYuo0I8yo83szDqXnCXhz+QWCMw27y9sGYN5PmRwot+NSPUfraYw6mSe+2sGwtBhmjunMyC6xRycrOln26+gW6iKGLEnEhRvrToJkBYbfDiueBrdvCu8RGrTf1TQYcGOjr38mclNGKm+tycHl5y/YmHmESSdxTc1eci99kYQ//pHwC85HkiRMeoWrB3Xk6kEdA+5b63Tzw85CDlZ4xVaESUdtgPSsUEV5jdMd1HaCtse2AxXklfgunIQSnaywuVibWUJGFxGNbk2q7W46RAVXv2RzqciS5HcxPJTzX5ElKmwB7gOCZkOIqdOMLglhuPxElUJJWwFIj7f6HEMgOJ0prLJz/fx1FFTY/a4KH4nWrtxTzK+5ZZzfM5GnJ/ZDr8hEW/U4/Zx3oS5iuFQPUWY/xeODpnjF1HEEbb8rG6DHpWAWDm/+SIoy84+r+vDAx7+H5KBl1Mn07xTFn6fPwfFzBof++SSl7ywi8S9/wdwrcM+wzKJq3liVzUcb9yPLEk63iurRMCgyHg1cHt8xhCrKG4tIuFQXB2sOUumqxKSYiDPHEWmMbPxNC5qdr38v8JsZEkp0wuZU+XTzASGmWplqh4twU3D9KwMJKQj9/FdkkeLX0ggxdZqREmulR/sIfttX7vNcsGkrVoPCrWM6t9SQBYJWp6TawaUvrqKk2unj7OaPWqfK0m0HqXa4mT95MOEmPX07RrIxr+55F+oixuDUaF9rdABrHFz2Iu5P70KnBo5O+SApEJ4AE55ufNszmKsGdsTmVHn8y+1BCSqzXqZfchSv3zwEnSKjGz6ctI8+pPzDD9k3axZho8cQP2c2+oSEOvu9sSqbf32zE7eq+XzP7A2kVocqysOM/m/t+VX5vLfzPT7Y/QEaGoqkoKHhVJ0MTBjItN7TGN5hOLIkJmOtxcFKB/4uQaFEJzTgUGUI1wlBs1DtcBNmDM74x6RX0MmyXwOKUM5/l9tDjHBzbHGEmDoNuWJABzbnl/u9IDeWtgKgU2TO7ZHQ4DYCwemCpmlMWbiB0iCF1BHsLg+r95bw3Pe7mTg4mXCT/8tpsIsYJr3Mhb3aUWFzEWn2vQF/wSi2azfyR907yA2k+x1FNniF1NSlIioVBJOGp9A5IYynv9nF1v0VeDwarnrfB6tBwaRXmDE6jVtGp9cp8pYUhehrriHioosoefVVsi+7nJgpU4iZcjOyycQLy/bwyvLMkKJfRwhFlOtkiYt6t6uzv0N18OCqB/lx349omobL45sGtP7gen4v/p1wQzgvn/MyZ8WcFfI4BSePGsDNLdTohL/+aYKWpcruDnhf8Mf5vRL56vcCn7lbKOd/u0gTyTG+db2C5kWIqdOIWqebF5btZcnPeYSb9FTZXX4FVUOY9Qp3jO8snGAEZwwb88rILKrxmThDcM12X16eyVtrc7h+aCc251dQXus7UQ1mEcPh8vD0N7v459c7ObdHIreMSaf/YaOCZTsO8chn21h0y8PINRfBNw9A+T5vDVX9HlR6i/exnpfDRU8JIRUCw9Nj+fC2DLKLa3hzdTa/5JZRZXdj1Ml0jDZz04hUxnSLb7DZuBIeTsIf/0jUtddS+PRcsiZczLYp9/KfPfIJCakjBCvKFY/KzX1ij/5uc9uYsnQKmeWZONWG7fNr3bXUumuZ/PVk5p83n/4JwhGwpYkL9+/IGWp0MjZMRCdamyq7m7AQxNQto9NZtqPQb7PuYM5/i0Fh1tjOwnikFRBi6jRA0zS+2XaQx7/YwZDUaL6ZMwaH28MlL64KSVCZ9TLn9EjgltHpzTtggaAN8eqKLL83r2Cb7coS/PH8s5g8IpVRXeKZ8fbPJzRp1jhmpf311gJ+2FlI18Qwbh/bmQc/2crrU4bQo30EcC50ORcO/AZrXoLcNeCs9vaRssTCoGnQ/3ow+3eMEzROWpyVRy/vfVLHMCQn0/GF56nZsIGp7+3BbvT/92hMsB9PMKI8XaqFG6+m+OabiJ48mXvW/pHM8kwcqq91fyBsbhuzvp/FkkuWkBKREvR+gpNnbLcE3l2f52urH6Jd/7k9EusfWtDCVDvchAdIufVHv+QoOkSZyCzy79za2PmvaXBZ/w4hj1Nw8ggx1cqo1U6c+dVoNjcoEkqYAUNKBJKfRpL+yCmu4ZHPt7G/zMbcif0Y0fnYiuQnd4zk2lfXUml3NTq5sxgULu3XgX9c2UesagjOGMpqnCzfVUT9ut9QnLNcqsZba3O9YqprHE9f3Y/7P9x8UlEIj+aNem07UMFt727kP5MGHo1SHaXDAPjD6yf8GoKWYW/7bhSHlYCf70Owgj1YLAaFZ2+/iNRZwyl64UU+mn4uGyc4cdRvPAiUrSyj+JtinIVOFJNCxKAIEv+QiGL1ppDVump55pdneOHsF0J/04ITJqNzLGEmnY+YguCjk7IkcUGvdj77C1qW6hAjUwDPXzeAifPW+l3gawiTDE9P7IvFIKb1rYH41FsBTdNw5lZStSIf++4yJJ3M8a68kiIRltEB67D2KOH+Q/V2l8oryzN5e20OM8d2ZtrINAy6uql5aXFWvr9vLP/7eR/zV2RR43DXuUAbdDISMNSP1bNAcCaQVVyDUSf7OPGF2mw39zgr40v7dyA+wsifPthCcbUDm0v1EWvBonpAkST+tXQX47ol+DenELRpFqzM8uvOFopgDwaLQWHBzYPp3i4CiCDpmbl88/EkbBVbqN+XuPjrYoq+LqLjjI6E9QzDVebiwKID5MzNIe3BNGSdjIbG6v2rKbGVEGuO9fuagqZHliVuGZ3O3G93+V2QaSw6YVAkJg9P8ZkPCFoerwFFaNPsXu3DedL+K3+We2OXgtvXpMD0nUsZXx0NiMhUayDEVAvjcaiUvLUNZ34VmssDGmjuuisQGlC5PJ/K5flEXZZO2ND2dZ7/cWchD3+2jV4dIvjy7tF0iApcbBhh0jNjdDrTRqaxJrOEn3NKKK52YtIrtI80MaFP+wb3FwhOZ6od/vvxhNrXx6VquFXP0VrD4emx/HT/ODbmlfHqiix+2lXUYCPshlK9VE3jYIWdjzbmM2m4SLk61di0z78ZUKiC3R+S5K1zjQsz8p9JA+mddMze/GDNQX6r3ukjpFSbSuEnhSRNTyK8r9de3xBvIPn2ZHbfv5uKNRVEj4k+fHyJD3Z/wMx+M094jILQuW5oJ95ck0NBuR01hJUYSYIIs57po9KacXSCYPB4NGqdbqwhRIo0TePQP5+k34EdfPD3W7n3k+3sK7XhdHv8fg+sBgWrUcc/ruxDRm0i+XfdTdLzz2EdOrQp34ogCISYakE8DpXClzfhLrWBu5EL5OGJV8XnWWg2N+Fjk8kvq+Wxz7ez+1AVj1/Rm7Hd4oN+bVmWGNU1jlFdRd8JgeAIgZpTh9zXQ5J8TAkkSWJQSgzzJ8dgd6kMeOxbbCeY6mVzqcz7KZMbhnUS0eNTDH/pWhCaYNfJEpFmPeW1LvQ6CU0D1aMxvnsCM8ekMygl2ud7sa5gHYqsQL2vXO2eWjwuDxGD6vYkU0wK4X3Dqd5WfVRMOVQH3+V+J8RUCxNm1LFk5gguf2kV5bWuoFxGZQnCTXqWzBxBbJh/EwtBy1F9WEjJDZjV1Kfk1fnUrl9PyjuLUCIi+PaesWzdX8GClVl8s+0QdpcKkvd6MDw9lpljOpPROfbwaySS9Oy/2T/nHjq+9CKWgQOb780JfBBiqoXQNI2St7cFJ6SO38/loeL7PL4tKOfx3QVMG5nGC9cPwBRgEigQCIInKdrcJM1248IMDYqczzcfQPbzfCipXiU1Tn7JLWNIakyj4xG0HYwB0q1CEewGncyfL+zOZf07UGFzoVdkIky6Bl1XKxwVuD2+kVe1WkUXpvNbl6uL1GHLtdV5rNJZ2eDYBM1DUpSZr2aPZsobP5NTUhMwXVgCzAaFxAgTb08bSnKMpcXHKvAl1HqpsiX/o/yDD0h5dzFKxLGFjt5JkTx33QDAu4Di0TT0Ac576/DhdHjqKfLvvIvkea9g7tv35N6EIGiEmGohnPuqcO6r8iukNuRv4R8/vsLu4hxkWaZrbAoPn3MX/dv38G7g8tBtazmfzRlJpzhrC49cIDh9SYoy071dOJvzK+o8HopzlkkvMzmj4fS777Yf8huhCCXVy+ZUWbWnSIipU4z2kSYKKnz7goUi2GXJ2z/GpFeCXkgL1HhXCVNwV7vRVM1HULkr3OjC6k4LpPp5goIWIyHcxJd3j+K3feXMX5HFjzsL0SsykuStm/ZoMKZrHLeO6czw9BgRtW5DhFIvVfnNtxS/9BIp7yzyafR9PIosoTRyPoaNHkX7J/7OvttuJ3n+q5h7hW5kIwgdIaZaiOoV+d4aqXpUOWqY+sH/8cT593Jp9/E4VTcb8jdjVOoaT8TIMnFVbhBZegJBk3LbuM7c9/5mahx1xU6wzlmaBtcP6dTga5TW+O/vE0qqlwYUVTfcJ0jQ9rg5I5VdB38/KatrnSyT0Tk0E4goYxR6We/ToNfSxYKkk6j8tZLIocdqrFS7StWWKhL/UNdSO8okLPZbE0mSGNgpmnk3DqKi1sW+slqqHW5W7ikmr7SGF68X6VxtkWB7TNWsW8fBRx+l04LXMKQ0TU1s+PjxaI88zL6Zs+j0+gJMZ4kG3M2NEFMtgFrjwraztI5j3xGySvcBcEXPcwEwywpj03yLBzWnh6oV+RjTIn2eEwgEJ06fjpEBV9+Dcc46t2diozUKgZq8hlqbpQ8h/17QNriwdzse/Hir3+eCEexGnczNGakhN1IflTQKt+ab5qdYFBKuSODAOweQTXIdNz99jJ6ojGPiyaSYuDT90pBeV9B8RFr0RFq8cwCdLPH4F0WtPCJBIIKJTNm2bmP/vfeR9OyzmHr2bNLXjzjvPHC5yJsxg5SFCzF26dKkxxfURYipFsC1vxpJJ/u49gGkxyQjSzL3fPkEl3U/hwFJvYgyhfs9jjNX5K4LBE2Bpmmsyypl/opMVmeWcCImwjpZIjHCxCV92vPppv3oZJnYMAODUqJ9ctoTIvyLrVBSvXSyREKE6QRGKmhNjDqFycNTWLg6G7uf+rzGBLskwaThDUc+/RFtimZM0hh+yPsBTz0XivgJ8ShWhYNLDuIsdCKbZSIGRpA8MxlZf+y7q6FxeZfLQ35tQfPTNTGcvYXVeDxaSCYHgpahyu4iwqQP+LwjO5t9t82i3aOPYB3WPO57ERMmoLnd5E2fQac3F2JMEy6PzYUQUy2Ax+b2G5UCCDda+WjSS/xn/bv8aenTFNWUMr7zMP514Z+It9atjfAEcIUSCATBY3ep3PnuRtZklmBzqoFOzQbRKxJGncyhSjv3f7AFDe1wbEtClmDyiBQmD0+lXaRX/Fw5IIkfdhb6pBKGluolGnGeqsw5rytrMovZUVCFUw2+mbNJL/PMxH4khJ+YiJ7Sewqr9q/CrvrWbMWMjSFmbOD6O0VSOC/lPMIN/hf3BK1LpFlPmEnH/nKbMJ1oKxzaButegf2/Mq6qnCEePbzbA4beCunjQfYuVLgOFbJvxi3E3323N4LUjERedhmay0XetOmkvPUmhk6hL8wIGkeIqRZA0jW8atQ1LpVnL34AgL0ludz9xd95ZNmLvHzZw3WPI1afBIKTwun2MGnBerYdqPDbELMxLAYFp9uDpkGNwyvEnKpvKtWCldksggXuEQAAIABJREFUWJnNny/szrRRaYztloBJp/iIKQi+NuusduF0SQgLecyC1seoU1g0Yxg3vb6enQergvruGVUXDw1L4OK+J96Es29cX85LOY/vcr/zK6gCISERYYjg3kH3nvBrC5qfbonh7CmsEmKqtdn9LSx7DEr3guoCj5swIAxgdy7krAKDFUbdg9r9OvbNmEHUtdcSPXFiiwwv6uqrvYJqylRSFr2NPimpRV73TEKIqRZADjc0vtFhusSmcE3vC3ln02e+x7EEDhkLBILG+ctHW05ISElAhygzOlmioMKGU204nnWkQe/T3+yirNbJ3ed0pXu7cFZnlvjdvrFUL4tB4bZxnUMas6BtEWHS87+ZGbyyfC9vrM7B7fH4iGu9IiFLEn2SIrkrvoYOLz6Ee9yH6KKjT+g1JUnisZGPUeGsYEPBhqAElSzJhOnDWHjhQuItwfcyFLQ8XRPC2X2omrO7Jza+saB5WDEXVs4Fly3wNs5qcFajff8ojg/+jXXEzcTeMqPlxghEX3cdmtNF7hFB1U5kOTQlQky1AIaO4Uh6Gc3PqvTeklyWZa7lsu5n0z4igQOVh/h0xzIGdqhnZ6mTsA4RX36B4EQ5WGHniy0FR4XO8dRsX07lz5/gKslHNpjRJ6QTmXHN0aa5GlBQYUMvS40KqeOxuVTmr8jig1/z6ZoQRnqclbzS2qCacB7BqJMZmhbD+T3F+X+qY9DJzD63G3eM78L3Owp5Z10uBRU2HG4P4SYdQ9NimZqRSurhFhiFuVs48Mf7SZ7/KpJyYr0FdbKOF89+kWd/fZb3dr6HLMnY3L4TP0WT0OkMdI3uyrPjnqWdVXzf2jrdEsPYkFPa2sM4c1n7cuNC6jgktw2T1Yk5ZQuSpnkLIluQmJsmo7mc5E2ZSqe332rQhl0QGkJMtQCSLBE2KonKZXlQb0XcarCw6cAOXvv5f1Q6qokwhnFu5xE8OP52n+NYh4mbm0Bwoixal+P38coNH1Ox/gNiz78DU9pAJEWHLftXbHvWHxVTAB4NHH6EVGNCzOH2UGV3s3DKEMpsLibOW8uBcptfUVcfs16md1Ik824cJIrMTyN0isyFvdtxYe+Gr+nxs2eTN206xS//h/i77zrh15MlmfsG38fMvjP5LPMzFm5bSGFtIXpZj+pRkZEYu9fAHfctpnOUiICeKnRNDGfx+rzWHsaZyaFtsOxxOG5hIvW5KmpdkD07DKvBe71esNHJO1tcLJ/iXSCRJRVyVsKvC2HI9BYfduz06WhOJ3nTppHy1lvoYuu2XKhxuPlpdxHF1Q5cqkaEScfQtBhSYkWP04YQYqqFsA5p5xVT9WgfHs8rVzza8M4ymLrFoIQFny4oEAiO4VY9LFqb6yNgPI4aylctJnbCHCxnZRx93NJlGJYuwxo9brBCTENjxd5ixp+VwOd3jeL+9zfzw85CAL+iyqxX8GgaEwcl89dLewbseC84vZF0OpKemUv2HyZi7teXsLFjT+p4YYYwbuhxAzf0uAGH6qDKWYVJMWHRWcgcfzYdp8og2kqdMnRNDBOOfq3F2pdB9e37p2rw/HonD4xuoF2GqxZW/RsGT2vx6BRA3G23eWuopk6j01tvoouOZs+hKl5flc0nm/ajyBJuVUPTNHSKjOrR6J0UyayxnTm7e0LAVh9nMkJMtRCKVY/znI5oS3MxhtJRXgLZqif6KtEjQCA4UYqqHX5d1Bz7d6K5nVi6jQj5mKEIsRqHyrzlmYw/K4Ewo45XbhxEYZWdd9fl8fa6XMpqnMiyhEfT6BBpZsaoNK4e3LFBa13BmYEuPp6kfz9D/t2zSV2yBEPHpikeNypGjGbvhE/TNLLHXsKTi39llyGbGocbk16hQ5SZKRmpXNCrHQadEPRtjQiTniiLnvwyG51ihQlFi2GvhK0fguZbunF/hoF/rXZw+xADUaYG5nq2MshdDakNt8RoLuLuussboZo+g0+mPcL89fm4VA/1b5NO1fsef80tY/Z/fyM9zso7M4YRZRGL+8cjxFQLsXlfOdNX7eH5ge1J/b0cLZgCeNkrpBJm9hNRKYHgJKiyu9H5WU1TbZXIloigGubWJ1Qhtjm/vM7vCeEm5pzXjTnndcPj0bC7Vcx6BakVVioFbRvLoEHE3jKD/bNnk/LuYmTjsVXvKruLgxV2qh1uwk062kWaG20WejxfbjnAv5buolDtgd3hQZOOpC25KKiws/NgJQ98/Ds3j0hl9rldRZS0jdE1MZzdh6qEmGpJdnwOAe4ZgzsojEvVMXeNg7+f3UBLA2ct/Px6q4kpSZKIv+8+5la8xWcrMrHLjS/c1TpVdh2q4pIXV/HlXaOJPElTNJfq4dtth1ifXUJxtQOTTqFjtJkrBiSRHn9qOdcKMdUCrM0s4c53N/Lk1X0Z2TMRW79Syj/PxFPhRHN7fHtQ6WRAw9w9hqgrugghJRCcJEadjD/PB8Ucgae2Es2jhiyoQhViDpcnYDqOLEtYDOJyLAhMzM03Y/ttE4f+8U/aPfIwm/aV89qKLJbtLER32AXQo2m4VY3zeyVy6+jO9OkYGfB4mqbx9De7eGN19jF3S8lXKB1xHFywKou1WSW8NW1oSGJN0Lx0SwhjT2E15/YUjn4tRuUBrxgKwGPjjYx8o4bZwxqau2lQltPkQwuF9zbk8TmJ2OXg3W1dqsahSjtTFm7go9szTmjxr6jKwcLV2Sxal4vHo1FzXA9VnSzx6ooserSP4LZxnTm/Z+IpscAorogniMej4TmcT9oQ328/xJ8/3MKLNwwgo3McAOazYjB1i8aVX03Vinwc2RVoThVkCdmixzokEeuQdkJECQRNRGyYEZefND9jUncknZ7a3Wuxdg9thTBUISZLUmukxwtOEyRJov0Tf+eX66cw5fEvyHPpcbhVrzFKvVZnX24p4PvthXRNDOONKUOIC/Ot3/jP8kwWrs4Juk2A3eXh9/0VTF24gXdvGS4iVG2EbonhrMvy33JB0Ew4q/FdBT9G7wSFS7rpeHKVkx7xDZwnQboANgeqR2Put7uxBTj/GzJWcqkauw5V8WtuGYNTAzf+9sfW/RVMWrAem0vF6ade2O3RcHs0Nu0r554lmzivZyJzJ/Zr89cbIaZC4EiB3hdbCqhxeu9eekVmQHIUs8Z2Zmy3+Dqrzh//ls8TX+5k4dQh9O1Yt6pXkiQMyeHETurRou9BIDgTCTPqGNUljuW7iurcAmWjlahRkyj9bh6SrGBKG4Ak67DnbMKet4Xo8dMCHjNUIRZh1p0SK2yCtssBp8ysQTOorHaiyr71GkfwaF5b/u0HKrno+ZV8esdIOkSZjz6/o6CSF3/Y41dINTSJcro9bN1fyesrs5g1TtTxtgW6JobxdgCnUkEzYY4CSQeab8P2Izw6zsTAV6u5b0QDRhSmwJHj5mb5rkIcLv/XkGCMlY60/QhFTO06WMW1r66tE4lqiFqnyrfbDnKnS+WVSW3b0fbUEFMuG2z7GHZ9DbUlIOsgogP0uw7Sxja7G0pmUTX3LNnE7oNVuDx1C/Scbg/rs0vZur8Ck0HhkUt6cWn/Dry1Jod5P2Xy3i3D6JoY3qzjEwgEjXPrmM6szy6ltt6FPGLoVcjWaCrWLqH4i7lIBjPGxC5EjLi2weOFIsT0isTVgzo2y/sSnBlU2Fxc8+paKhwqniBTS90ejdIaJ9fNX8dXs0cfTc9bsDILlx+b/2AnUa+tzObWMZ3b9OTmTKFrYjiZhTXC0a8ladcX9KbDESr/dImRubaXnhc2OOmT4CeqohggeWgzDrJhXl2R5VfUBGuspGkctVD3F/muj92lMmnBOp/7LzS8gGNzeVixu5j5K7OYNbbttm1o22KqsgBWPwu/veP93VlT9/kdn4MxAjLu8vr16xr/g4bKb3llTH59AzVON1oDfTZrnCo1TpX7P9zMkp/z2Fdu438zR5AcI4pCBYK2wPD0GGKtBmxOm0+CRliv8YT1Gh9w3yNTlPr7BSvEZEni5hGpJ/sWBGcwb67OobTG6bf2r6HJiOrRKKy08976PG4Zk06l3cUXWwpQ6x0oFHdKu0vlpz1FjD9LNP1sbcKMOmKsBvaV1YpeQC1F+ngwWBoUUwB/G2tk0RaX/yclCYbe0gyDC47tByr9Ph6KsZJBJ7P7YBVxXRqfe3/1ewE2p+pzDw12AWfe8kxmjEprtLSmtWi7YqpgC7x9GTiqwRPgy+is9v788Bj8/j5M/gjM0U02hKyiaia/voHq+gnpDWB3eVidWcJfLuouhJRA0IaQJIk3pgzh8pdX+10da4gwkw6n2+O3J1RjQkyRJQZ0ihLXA8EJo3o0Fq7J9vv9C2YyYnd7eG1lFtNHpfHdtkN++8SEMomqcaosXpcrxFQboWtiGLsPVQsx1VLIMoy4E378Z52mvTlz6mYhJUfK2B+K8H+M5OEQ1ak5R9kg9gApfqEYK2kaVNoDzM/rMe+nTJ9IWCgLOC6Ph2U7C7mgV8ONzluLtimmivfAmxPAURXc9i4bHNoKCyfALT+A3tz4PkHwx/c3H62Nqk9DK4Ea8Mx3u5k4OJloqzCREAjaCl0Tw1k8Yxg3vbGBGofb7yr/8ciSt5fLu7cMJ7+slrv/+1vQBftHCDfp+Pc1/U9i1IIznR92FuLyI6RC63XmZnVmMQcr7X4nUqG6U+4vb73ieUFduh22Rz9POPq1HAMmw6rn6oipoNGZ4eyHmn5MoQxBkXD7uQGGYqykVVdR+ugj5Jrt6BISDv/Eo09IQJeY6P09Pp7d5U72lfp+TiEt4DhUXluR1WbFVNuLl3nUYxGp40h9roqEp6uocR774y/Y6GTcm4dT/1QnlGbBF/c2yTCyiqrZdqDSb2pf5YaPKV32GpHDr6Hjne+QdNtCwgdOwLZn/dFtJAmW/LKvScYiEAiajgGdovnirlGc3zMRo07G5KcZqVkvY9TJXNynPV/NHk3PDhGc36sdj1/eG5M+uMumLEG0Rc+SW0fUKf4XCELlu+0H/dY3hBpN+uqr9ZT/8hsePze24ydRweAIcVFB0Hx0TQhjz6EgF58FTYMlBm7+DAwhRgN1Zrj4mVatlwKICbDQf7yxUmNoFiu9/jSbuNtmETZ6FEp0FK79B6j85lsK5z5D3s1T2D1kKKtm3I1k97WSD3UBJ7u4pvGNWom2F5na8523u7Qf20lVg+fXO3lgdID8TLcdtn0IF/7jpNP9Fq7O8ckph+BXAu0uD6+vzObW0emiKFQgaGOkxFqZN3kwxdUO/rshjy+2FFBpc4EEUWYDVwzowDWDk326vE8cnExStJmHP91GfpntqDX18Rh1MhowqnMcT1zVm/aRQkgJTo7iaqffx0OdjBTmF5JgcaPHTP3knNDdKU+uYaeg6eiWGM7C1TmtPYwzj3Z9YNq38Pal4LKDK3DvKRQDSApc/jL0ubrlxhiAG4Z14qVle7HXi3iHYqwUF26i7+AeDbrUaprG5pW70b7NBHfdm2Wo7UVsAVIT2wJtT0ytfi5gUd/9GQb+tdrB7UMMRJkC/fFkr2FFxl0nNYzPNx/wGwINbSXQzfaCSnontZ79pUAgCExcmJE7z+7KnWd3DXqfjM5xfHfvWH7Pr2DByixW7S2m1qkiy14hdvWgJG4clkJChKkZRy44kwi0HhfqZCRs0ABGDU/htTc24KoX6QplEmVQJEakx57UexI0HV0Swsgqrkb1aH7r4QTNSLvecPdmtI3v4P70b+jCdUho3mwpWecVUUgweKrXcKIV66SO5/ohnXhh2V6/zwVjrGQxKMwcm95ouw9JkgiPjkBRFHDXLZsJdQHHpA9u0ag1aFtiqmI/HPgt4NODOyiMS9Uxd42Dv58dYKLitsG6eScnpjweqgKYToSyEqjIEqU1/lcUBQLBqU2fjpE8f/2A1h6G4AwgIcKEhG++RiiTEQlICDcyOCWaGKuBWqdvDUOw7pSSJDF5RMrJvSlBk2E16oi1GskrrSUtTphQtDimCKqdPSnJOYfUp+6Gwu1gr/CmAEYmQ9fzmsVt+mSIDTNyXo9Evt1+0G+bhGAcbq/onxTUa6XEWlD9pBaH2ucxpQ2bOLUtMVWe61XxbnvATR4bb2TkGzXMHtaAsUPVAa+gctV6j+Wq9ZpUuGzH/b/2cFi2/mM2NLcTzbMIfyVlIa0EauD2iLxygUAgEJw4l/btwCe/7fdxoQxlMmLSK0zo0x5Jkpg1Np0nvtzpN22msUkUwODUaFEH2MbolhjG7kNVQky1EhWffkrkFVdA6kjvzynAP67sw6Z9ZRysdPgtawmESS8zb/IgrMbgJETP9hG0izD5rXkKdgHHalCYMTo96DG2NG1LTDkbyDc9TO8EhUu66XhylZMe8QEKwTUPlOzxrgroLWCN9/6rM3md/vSWAP96fySdCdPD3/i1Tw5lJVADIkVeuUAgEAhOguHpMUSa9X7vScFORhIjjPRPjgLgygEdeXHZXr81f41h0svcf0H3E34vguahW2I4ew5VtVm3s9MZtbKSmpWraP/ww609lJCItOj54LYMJs5bS2GlHaefCFV9THqZZ6/pz+iu8UG/jiRJ3Da2M498vs3vNSyYBRxFlji/V9t1q2x2MeV0e/h6awGvrcwiu7gGh8uDQSfTPtLEtFFpXNE/6Zi6NYYFdcxHx5kY+Go1940IEDaV9V63lJNgeHosP+4s9EmrCGUlUPVo9Gwv6qUEAoFAcOJIksTMMek8tXQnNj8ueo1NRsx6hVljOx+tb7AadSyZOYLLXlpFdRAtAo5g0sv86+q+R0WZoO3QNTGcFbuLWnsYZySVXy/FmpGBEnXqnRftI818NXs0T329k4827keS8BE8ekVCliT6JUfxwIQeJ3T+X9qvA//4age1hG4iYdYrTB+Vhr6NNuyFZhRTHo/GSz/uZf6KLDQ0ahzHPkC3UyWzqIYnvtzB419s59rByTxwcQ+MMZ3B7Wj02F1iZK7tpeeFDU76JPj5cJugwG/mmHTWZZWc8EqgTpaYOLgjZkPbLZgTCAQCwanBdUM78b9f8tlTWOW3xiEQBkWme7twrhrYsc7jqXFWPr9rFNfNX0elzeXXev0IRklD1ul44foBopdRG6VbYhgLVma19jDOSCo+/ZTYGdNbexgnTIRJzxNX9uHBi3vw2aYDLF65h8L8QoiPJ9ykZ0zXOKZkpNEp9sRrlswGhUXTh3Ht/LV+59WBMOJhSFpcSCZRrUGziCmX6uG2d35l9d6SBq0Mj3ygS37Zx6Z95Sy+ZThhaWPQ9n7vdUNpgL+NNbJoi5/Oy3orZNx9UuMHGJoWQ7TFf5EuNL4SqMgSU0emnfQ4BAKBQCAw6RUWzxjG1fPWsL/MhsNPE9/6GHUyqbFW3pw2FIOffmopsVaW3z+OpVsP8sryTHJLapAlbzNPRZaQJdCjcUXmSm5/+SHiI9tuAfiZTpeEMHJKanCrHnRteAX/dMOZl4czJ4ew0aNbeygnjcWg47qhnbhMLqTw32+Q+uziJj1+n46RLJo+jClvbMDuUnE1EhI362UGFO7lmVjavEtlk4spTdP44/82s2pvMfYgm/rZXR52FFQydeEG5g6eQULmKsxaXRGTMye8zu/JkTL2hyL8HM0DfSee6PCPIkkSf7+iN7ct/jXo93EEs17h4r7tRSGoQCAQCJqMaKuBz+8cxX3vb+bHnYUAfkWV8bBwuqBXO566um+DGRJGncLl/ZO4vH8SOwoq2Xmwkiq7G5NeISnKzPD0WPbd8DbGdavggvOb540JTg5Nw1LwM08a3sDx9mvodBpY4rwucj0vb3NOcqcTFZ98SsTFFyPpT5/6eHdxMbq4uGY59qCUaJbeM4b5P2Xy/q/5QN20QlnyXpOSY8zcPq4L51vSyZ86DWtyEpaBA5tlTE1Bk4upH3cV8t2OQ34FSM325VT+/Amuknxkgxl9QjqRGddg6tgLp6qxMa+cC/fDmogOmGy5SB7/9uQB0Zth8PTQO1IHYHz3BB6Y0IN/fLUjaEFl1isMSonmyav6NMkYBAKBQCA4gtWoY96NgzhUaeeddbksWptLpd2FTgK3qhFpNTAlI40bhnUiPjy0SXSP9hH0aO+7SBkzdSqlCxcSccH5HCi3kV9mo9bpJtykIzXWSmyYmKy3CqoLNi7y9uesKeIy1Yace9xq/66v4It7YNAUb8ZOuEjRbEo0j4eKTz8l6fnnW3soTYq7qBhdXPP1kUuKMvPo5b35y4QefLb5AGv2FlNa68Sk8y7gXD2oY53+rB2eepL82bNJffddDMnJzTauk6HJxdQryzP95kNWbviYivUfEHv+HZjSBiIpOmzZv2Lbsx5Tx16A17AhMsxE1MwvkOaNAls5aEHmVupMkDQEzn20Kd8ON41IJcZq4P73t/gtzDseo07msn4deOLK3iLMLhAIBIJmIzHCxH3nn8V955+F3aVS+vt2Kh9/hO4ff9jkr2U5+2w+mf8h9zz9Ldsr1Dopg063h5Fd4pg5Jp2haTGNNvEUNBH2Snh3IhRs8bZ1wU8zF2e199/182HTu3Dz594ms4ImwbZxI7LFjKlXz9YeSpPiLi5GaabI1PGY9ArXDE7mmsENC6Sw0aOJmzWLfTNnkfrf91Ai/GWltS5NKqZyS2rYkl/h87jHUUP5qsXETpiD5ayMo49bugzD0mVYnW0r7S7WFxsZfutyWDgBaoq9jXgbQm+F9LHwh4WgNH0Z2CV9OzDurAQ++S2feT9lUVLtRKd4bxia5k1tVGSJf17Vh4v7dmjy1xcIBAKBIBAmvUJ8UgI1xU3v5rbnUBU3vr6eqrMuobbEW6dcP7Xwx52FrMsqoVOMhbenDyUh3NTk4xAch9sBb10KhTtAbdy0C48TbKWw8EK4ZTnEdWn2IZ4JlH/yCZGXX37aLSC4S4ox9+3b2sOoQ8ykSTizc9g/Zw7Jr77a5tIqm1R5fPl7AR4/BWWO/TvR3E4s3UY0egybU+WDX/MZPrEf3L4WNr4Da14ARwU4j2v4Jeu8P4m9qB16F9+ogyhaux+n20OEWc/ATtF1woQnS5hRx43DU5k0LIXMohpKqh24VI1Is56uiWG8uz6Pr7ceFGJKIBAIBC2OLjoad1kZmseDJDdNZsTv+RVcd9h9S/PTxP4IGt6sjb2F1Ux4fiWf3TlKNPVtTr7+PyjaVUdIpT5XRa0LsmeHYTV4J/cLNjp5Z4uL5VMOlz44quHty2DO7yALp+GTwWO3U/Xtd8R//llrD6XJUYuK0cUF30eqpUj8vz+z7447OPj3J2j3yMNtSsQ2qZgqKLf7dedQbZXIlgikIE5eDSgoPxyJMobDiNtg+CzI/gmyVkBNISh6CO9AVsI5vPy7zBfvF6CTt+Nwe/BoGgZFBkkiOdrMbeM6M6FPe0z6prlwSJJEl4QwuiTU7Yl19aCOPPf9bg5V2kmMEKtyAoFAIGg5JIMB2WLBU1nZJP1uCipsTFqwrkHL9Pq4PRpltS6um7+OpXNGYzE0eyvLMw97JWx+F9x2n6dUDZ5f7+SB0YFq2DSwV8Ce7+CsC5t3nKc5VcuWYe7TB33i6VeH5i4uRhff/Gl+oSLpdCQ9829yb7iB0jffInbqlDrPa5pGUZXDW0Mqy8SEGYgwtUwEq0mvdM4AVq2KOQJPbSWaRw1KUDnVeseRJEgf5/3B+4E9v2wP897NxKVqqB6N4wPd9sPj2FNYzUOfbOX5ZXv4763DaR/ZfCtlkWY9l/brwPfLvmWS9gXk/+LNV1aMENURht4K3S/xCkGBQCAQCJoYXUwM7tLSJhFTL/2wN2CNcENmUqrHO6H58Nd8Jo9IPelxCOqx+T2Q/EcJ788w8K/VDm4fYiDKFGDV3lntNawQYuqkqPjkUyKvuLy1h9EsNKeb38mihFlJnvcKOdddjyGlE+Fnn02FzcX7v+zjtZVZlNe6vGU4mldL9OsYxayxnRnfPaFZ7dWbVEzFhhn8Pm5M6o6k01O7ey3W7qMaPU60xf9xjvDEVztYvC4vKIe9WqdKfpmNS15YxVezRzdf1GjXUv6W91c8ZTlokop0vHFGRZ63SPSz2TBsJoy5H3QNv0eBQCAQCEJBiY1FLSmB9PSTOk6t081HG/fj9pNpEoyZlM2lMu+nLG4cntKmUnFOC9b956jhRH0Gd1AYl6pj7hoHfz+7gbnOgY1QkQ+RHQNvIwiIq7AQ2+bNdHz+udYeSpOjeTzeBZnY5nPzO1n0HTrQ8aUXyZs5i8/uepoF2yuRJeloX1vHcUbgv+SWMfu/v2HUK7wyaSDD0pvnfTWp5dyw9FisfvpZyEYrUaMmUfrdPGp3r8XjsqOpbmyZv1D24xt1tpUAh1tlS345muZ7If/o13wWr8trsBlwfVSPRrnNxfWvrUNtpEnYCbHyGXh/Csay3Zhx1hVSR3BWe+u+1rwIb17sDdULBAKBQNBE6GJjcZeUnPRxPt10AH8a6IiZVMx5t2E5KwPZYEJSdFi6DCN6/LQ625bVOtmQXXrSYxHUo/JAg08/Nt7IixucFNU0sNisGKEst4kHduZQ+cWXhJ97LrLl9GtirVZUIFssyIa2veBv6tOHFyY+xOu/FeJwexrUBDVOldIaJzcv3MA32w42y3iaNDI1ukscFoPOb451xNCrkK3RVKxdQvEXc5EMZoyJXYgYcW2d7Qw6mc4JYcz+7yYcLpXze7Xj/F6JDE2NQZElnvpmZ8APrbHUg0MVdpbvKuScHk2Y47puHqyY27jj4BHcNijYDIv/AFO+FGl/AoFAIDhpNE2jKiaB4gNlJFbYiLYYTrhW+Pvth/ym+IVqJrU2s6TZVoLPSDQNVGeDm/ROULikm44nVznpEd/AermjqokHd2agaRoVn3xC4oMPtvZQmgW1Daf4Hc9TS3fyQ5mMXQle9NldHub89zfemTGcQSnRTTqeJhVTsiwxY3Qaz363+2jd0vGE9RpPWK/xAfdXZLhqYEf+dkkv/nodEsJQAAAgAElEQVSxxp7Car7ZepB/frWT/LJaeneIpLzW5XffYFIPapwq837KbDoxVbwXvn+4TiFoUI46qsOb9rfqORh7f9OMRSAQCARnHGU1Tpb8nMeCVdlU2Puh26PB3J9wezyMOyuBW8ekMzglOqR0u7Ja/xP2UM2kiqqDsO0WBI8kgWJoVFA9Os7EwFeruW9EA82UTW2vV8+pgGPnTjzV1ViGDG7toTQLbble6gj5ZbUsXJ3j06IBGg6qANhcHv7y0Ra+vWdsk46pya12rhvaiddXZeOsdhBqRp1Fr+OO8Z0Br2tet8RwuiWGc9c5XdlfbuOm19f7/fBC6WO1Jb+CfaW1JMc0QXh23X/A4/Z5uHFHHbwRqnX/gdH3CotSgUAgEISE6tF44svtvLM+D1nicA2xhEuT4HD2xvc7DrF6bzHx4UZeu2kw3RLDgzp2oELtUM2kDKJ5fdMT1QlK9ja4SZcYmWt76Xlhg5M+CX7+Bm4HRKc10wBPbyo++ZSIyy9rsvYDbY1TQUwtWpvrtwwomKAKwL5SG1v3VzRp+6Qm/zZEmvX8b+YIwk16QjHOsBgU3po+lI7R/kVOUpQ5YFQqlNQDgyKzo6AJ6pWcNV5XHT9i6v4MA3PXOCi3N6ImVSfs/ubkxyIQCASCMwa36mHGWz/z3oZ9ON2egGZMmuY1YcorqeXKl1fza25ZUMdvF+nfvOB4M6nG0MsSCaJNSNMz4i7QWxvd7G9jjdQ4A8xBOg2HiPZNPLDTH83tpuLLL4m87LLWHkqz4S5q22LK4VZZvD4Pp1r3ux1KPafT7WHByqwmHVezSOvUOCtf3DWKpCgzVmPDq1dWg0KM1cD7s0YwsFPDOYyBaqVCST2wu1WWbj3Ih7/ms3RrASt2F/FLTinbD1SSW1JDUZWDGofbb/PhOuxeGtCe9HhHnQZxVsPPCxods0AgEAgER/jzh1tYl1UatBGThjfN/eY3NpBdXNPo9lcP7HjSZlKyLDGhT7ugxicIgb4TAV/xnDMnnHPTjyUbJUfK2B+KOFZecASDFUbObuZBnp5Ur1qFISkJY9rpG9VzFxejxLXdOset+/0HQ0IJqqiaxg87C5t0XM3WUS85xsLy+8ezfFch837KZEt+BQadjKZ5035dqoe0OCuzxnbmwt7tMOqCSxmoxffmEWrqQVG1g9V7i6lxuql1qtQ6VWocR/7vpsahYnermPUKFoMOi0HBYlCwGo/9f0LVL1zstAf8AB8bb2TkGzXMHtZIcVxFfqPjFQgEAoEA4Pf8Cr78vcBvNKqxeoEap5tHP9vGm9OGNvgaY7rGYzYoJ2Um1ScpkpTYxiMoghAxWGHwNPjlDXAFaXx1GM0DqtOAkjwSYVgfOhWffkrklVe09jCaFXdxEcauXVt7GAGpsDn9Oo2GElQBAvbQO1GatT25Ikuc0yORc3okkl9WS2ZRDdV2NxajQqcYC53jw0I6XkKEkXKbb6pfKH2sDIrMPed1azQK5vFo2FyqV3A5vP/anCo1TpVah5vkbXrkwsDWo0E76vjpYi4QCAQCgT8WrMzC5T6xegFNg7VZJRyqtDfYc1GWJWaMSuO5ZXv8irbGzKQsBoVZYzufwLsTBMW5j8GBzbD/l+DnEJIM5nAK943BedMUkp57Fn17keoXLGplJTUrV9H+4YdbeyjNilpc0qbT/ABvqL0eoQZVmpoWq6DrGG1hbLd4Lu7bnvFnJYQspACmZKRiOcnUg3CTngHJjXeHl2UJq1FHQriJ1DgrvTpEMjg1hrHd4rmoT3t6pycjN9J499FxJl7b6GR/ZQMpg8JRRyAQCARBUFHrYum2g6jaidcLgLeAuzGmjEwjPS4MvRJaDMOklxnTNZ5zeiSEtJ8gBBQd3PgBpI0Jqn4KnQnCEpBu/ZH2L71J2LnnkD3xGqpXrgzu9TQNnLUhR8JOJyq/Xoo1IwMlqvH546mMu7gYXXzbFVNRFoM/LRVSPSeA1di0saRTyo7kigFJ+DHwALypB9FnT6di7RLyX5xE/itTqNr4Beaux/InzXqFW8ekN01H9g4DA9ZMHeF4Rx2/yHro1Hh+p0AgEAgEa7NK0PtxyAulXsDh9vDFloYbvwKY9AqLZwwjJdaKURfcVMGsVxicEsPz1/dvmvusIDB6M1y/BC5/Cdr1BZ0ZpHqLzYYwsMbD2D/D7esgtjOSLBN3yy10fPbfFDz0Vwqffx5N9ZPy5PFA5g/w9hXweBw82RH+kQSPxcG710HuGgJOyE5DKj79lMgrTu8UP2j7bn59kiL9mtuFElRRZInzejZhv1maOc2vqbEYdFw9qCPv/7LPr0V6Y6kHkgRXD+rYNINJGgjh7aE0s8HN/jbWyKIt/l0INVlBGjaracYjEAgEgtOa8lonqh9zpFDrBSrtvi60/oi2GvjszpE8+PFWvvy9AMnpxCH7ThssBgVNg5tGpPCnC7sHtFYXNDGyDL2v8v4c2g5bP0I9sJPadWsIv/QaSB8PXc71237FMmQIaR9+wP4/3k/e9BkkzX362CR6z/fw2R3exr7OeoYlmuo14MpeAdZYuPJVSMnwOf7phDMvD2dODmGjGy4jOdXRXC7UykqU6KZtaNuU6BWZG4ensGBVNs56OiDYek694u2J25ScUmIK4MEJPfg5p5SsompcavCrIia9zH8mDSTSrG+agUgSjJoDX/8fuI5dbHLm1O3jccRRxx/btVTKSiMZ1XaNUwQCgUDQRgh0x2vOegGLQcez1/bnL0PjmH//XL4ceAmF1U48Hg2dIpEcbWHm2HQu7dcBi+GUm1KcPiT2hMSeqPn5HHpzCuEvP9noLrq4ODq9voCil14i++o/kDT3aSzydvj6T43UYmneeU95DSy6Cq6cB71O36hNxSefEnHxxUj6Jpo/tlHcpWUoUVFIStvufTp5hFdM+aOxoApAelwY3ds1bYnNKXflMxsU/nvLcG5YsJ6somq/ESqfffQyz1zTj3FnNXEOd+8/wI//8Dbg1Rofx/FoOjPVIx/igY9/p2tCGA9c3OOE6shaizJ7Gd/lfkdBdQG17lqiTdH0jO3JyA4jUUQTYoFAIGhyoi16v1GfUEyYACJMod/6leXfM62blYceOBfw9rrSiaa8bQ5Jr0dz+c+G8bu9opAwezaWgQMpe2IG5v77kQh+f9w2+HgWWOMg9dSO3PyeX8E763PJKqrG5lIJN+np3zGSUV//wMC5f2/t4TU77uKiNp3id4T2kWZuH9uZV1dkBd0e4ghmvcJTV/dt8jGdcmIKvKkHH9+ewYs/7GHR2lxUTaPGUfcDNSgySDAkNZo/X9idvh2boWjQYIGpX8H88eCoDF5Q6c1IFz7FsEGX8t1olbfW5PCHV9Zw5YCOzD6nK5GWtrv6sbV4K29sfYOf9v2ELMnYVe/qlYSEWWfGqBi5seeNTOw2kWhT2w0VCwQCwanGiPQ4XKrvfeb4egFJVjClDUCSddhzNmHP21LHhMKok7m4b8MubkVVDvJKa6l1urEadXSKsVD15ZfE3zPn6DZCSLVNQhVTRwgbMQzr0BIkZ919U5+rotYF2bPDsBq8Qn7BRifvbHEd62HltsGH0+GeHd7Uw1MITdP4bPMBXvxhL/vLbDjcKsdn0v6aXcKCPlMYuqaKe/+fvfsOj6pKHzj+vfdOT09IgYSSQg29Ezp2VFRA1EVWwYKiq7jW1f2t64q67qKsHSyoiF3BtmJdkA5Kb1ICISSQ3pPp9/7+GEHCTJKZkJB2Ps/D7uPUMxom9z3nLcFFDOoc2XSLbWTuwhbQye83d5/fldxyG59tO+53QGXSy7xy/UD6JIQ1+HpaZDAFnuLY+y/qwdzzu/H93lze25RJTpkNh0sl1KxjVEo7ZozoQny4uXEXEpkEt66CtyaCrcwziLcmitHTtGLSS9BnCgBGncKtY5KZPDCBBd8fYMIzq/jThBSmD+/ss9C4qWiaxsKdC1m8azEOtwP1jKGBGhpVriqqXFW8uvNVluxZwusXvU6PyB5NtGJBEITWJcyi56LUOL7aeZwzS6f8rRcAmDG8i9dtqqqxPr2QhT+l83NGEYbTmk44nG76RI7mnnYppGmaaC7RjEk6Xb2CKfZ9gVRDIqlbg+c2OXh4tLHm59sr4PBKSDkv8PduIi63yoOf7uTrXTk1XpA7VEDWsS69kC2vb+Jvl6Xyh2Gdzu1CzxFXfvNuPnE6SZJ48qo+dIkK4j8/HESSap4dFWRQCDHreWX6QAbUMRapvlpsMHWSXpGZ2Kc9E/s04byEyET40zbY+zmsXQAlRwEJVKen8FPWeYKoIbfAkJsgtIPXS7QLNvLEVX2YMaIzT/x3H+9sPMpfL+vF+IZOTaynBVsW8MH+D06dRNXG7rZjd9u5YcUNLLlkCd0ju5+DFQqCILR+t45J4vu9uT4v/vxpwjQ8KYq4sOozpo4VVTHjjU3kl9tPDeo9M4X+l8gkbnl3G3GhJpbePIwOjb1RKdRLfU+mWPefGjeD708z8K91duYMMRBuqiGQdlTAuudaTDClaRoPfLqTFbtOYPUxS80Xm1Pl8a/2YNTJDdfMrBlp7m3RzyRJErPHJjNjRGc+25bNop8Ok1VUiV7SQKfD5dYYmRLF7DHJjEiOatRNoBYfTDUbehP0u8bz58ROyNnpOanSmyE0HpLHg1J3+l6PuFCWzBrKyv15PP7lXt5cl8FfL+1Jt9iQOp8LcKLUyrsbM9mZXUK5zYXFoCMlOojrh3emq5+vcabvMr7j/V/f9wqkitcUU/BtAY48B4pJIXRQKLFTY1GCPDVTVa4qbv7uZr6d8i0WvaVe7y0IgiD8rnd8GJf0iQvoIvCkIIOORy/vVe229PwKJr+8jnKby+u0qzqJKoebo4VVTHx+DZ/NGUmXdn7MOBLOKUmvR3P5163xlMpCyN9f492DOyiM66Jj/no78ybUPOyZjLXgckAdMzibg8+2Z7NiV47Pv0OVe1dR9vNnOAuzkA1m9DFJhKVNw5SQitWp8shnuxjQKZykFlTn7g9XQQH6eO/N/ubOYtDxh2GduW5oJ9L/7zHs3XoRfeUkQs16TPpzU8MvgqnG0L6v5089SZLEhB6xjO4azTsbjnLdqxu5pE8c95zfjahg38fsW44W89wPB9h0pAgNqrWM3Hi4gA9/PkbX2GD+NKErF6bG+b0WTdN4fuvzXoFUwYoC8lfkk3BzAsG9gnEWOzn+znEy5meQ+Egi8m8pIna3na8Of8W07tMC/xchCIIgePnXlL4UVjjYfKTIr3oBCU/zprdnDal2AVhU6eCaRRsos7n8Hhnk1jTKrE6ueXUD380d26xrfNskRQFVRXO7/e/KVlUIigHcNczEBP4x3sjIxZXcPayWQEnRg60UgqMDXPS5pWkaz/1w0OffnbLNyynd9AlRF96BKXEgkqLDemQL1oObMCWkAuByayxed4R5V/Y510tvFFaHmxKrg9yCUjr06dfUy6k3SZIw5ucQO2EMIaG1BP2NoPkU5Qhe9IrMrFGJ/HjvWHSyzAULVvP6msNevfWXbjzK9Nc3svpgAXaX6nW/WwWbS2VXdhl3f7Cdv3+xG7X2LchTdhfsJs+aV/31rG7yPsujw/UdCOkbgqSTMEQb6DinI44CB6XrS0891uqysnj3YrQ2NNxPEAShMekUmcU3DuGaIR0x6mRMNQzVlfDMgEqItLB8zkiv4vk31x2pMZCq3LuKE2/PJfPZqWS9OIPcjx7FlrUHAFWDkion72zMaOBPJpwtSZI8dVOBnE5pdQfkvWMULuum459raw64AFADPBVrAjuySskts3vdrtorKVn7LpEX3I6lexqywYSk6LCkDKvWxMWlany6JZsqR/P/rDVxuFS+3HGcS59fQ+9Hv2XC/FVco09jxAaY8sp6vt+b63OmXXPnys9HF33ug3kRTLUA4RYDf5+UykezR7A+vZALF/zEt3ty0DSND3/OZN5/92LzM93D6nTz4c9Z/OOrPX49fsneJdhd1b90qg5WoTpVQgdV79OvmBRC+oZQsad63nWxrZgd+Tv8ej9BEAShboos8fdJqWz4y3ncfX5XooIMGHQyZlyYZA29IjGhZwxvzRzK6vvH0T2uepq3063y9voMr8038OzOF/34GmHDp5Fw51Lib3+TkIETsR7cdOoxdpfK4nUZLfKCq7ULuG7KFO6p8a7DY+NMvLbVQXZZDf/N3U4wN0Ln5Aa2ZEMGdpd3AGnP/hXN5cDSbUSdryFJ8O2enEZYXeP7fFs2g+Z9z0PLdrLneBluTcPqVLHLOtyaJ9Np7ofbGDzve1buz6v7BZuRpgqmRJpfC5ISE8ziG4ew+kA+8/67l5dXHmJfTrnPX4a15/x6AqrhSe24uHftKX+/Fv3q1bnPXeFGF6xDUryL+XRhOqxHrdVuUzWVQyWH6B/Tvx6fWhAEQahJZJCB28elMHtMMkVVDg499W9Ce3YncdpkzIaa07x+2JuL28eR1Mnd+aiJc7F0Tzt1uyVlGJaUYdUea3e6WflrHuf3im24DySctYCDqZA4sLSDsuxaH5YSKXNNqp7nNzvoE+NjLz6mh6dOvJk7nF/psz7QbS1DtoT6Nfja6nCTVWSt83HNzSurDvHcjwfr3ICvtLupxM3tS7fwjyt6M21wx3O0wvrT3G5cxcXooqLO+XuLk6kWaEy3aL6+azSyLNV7V9HqdPPC/w7W+V5Vriqv25RgBVeFC83t/W3kKnWhC64eo7s0F5XOSn8+miAIglAPsizRLthIZ72L9oqr1kAKYOX+PK/5jBDY7nylw81PB/LrvWahkRgCDKYkCdLuAj8aRf1trJFKh49IxBAMI+8JYJFNp6b0PMUcilpVhqbWnfaoAeX2lpXmt2xLll+B1OlsTpW/fb67RZxQuYuKUEJDkfTnvo5TnEy1UFanm30nyrxuD2RXMT2/gv055dXTP1wOT2v3wkNQeAiTrdzrPSwpFiSdRNmWMsKG/j78zG1zU76znNip1XcpdZIOs67571YJgiC0dJKiQ/ORwnSmggrftS+B7M57Xse79kRoWpJOD4G2R+9/Hfzwd6+bM+ZWTw/tGCZj+2uo1+OQJOg1KbD3bCKhJt8X28b4Hkg6PVUHNhDUY1Str6FIEGZuOc1XrA43f/18t89AqrZMJvAEVPd9tIPNj5yPIjffGXOu/Hx0MU0zTkgEUy3U17tOIPvomR/IrqLTpbL0ixU83mnbqeCJ0mzPHKyoFGjXlQRzNJmVmdWep1gUYq6M4fjS48gmuVo3P32knvC06jnTiqzQIbjltdsUBEFoaSSdguaue8dc7yNNG6rvzvsTUDWn4fKCR71mTZnCYNxD8NPT4PTOSKmN6pZRhz2ATlfLUN9mpE9CGDuOleA8I9dPNgYRPmo6Rd8vRJIVTIkDkGQdtozt2DJ3VmtCYdIrdK/nuJkauZ2ekTqyDMYwz/83kC93Hvd5uz/dCwFsTjc/HchjQo/mm9LrzMtrknopEMFUi5VRUOVz2nMgu4puDQ4V2CC1PXQZ7QmgIrpUmxExPWs123+63yvdL3piNEqQQs6HOTjyHMhmmdCBoXSc3RFZX/0LQCfrGN5+eP0+qCAIguA/RQd+dHKLCzUhSXh18gtkd16WIDa0ZVxAtyX1mjUFMPJuKD0G29/zP6DSm7GFjSH7sWUkxJ6HOTW17uc0sT8M7cQ7G476vC906GTkoAhKN3xIwVfzkQxmjLEphI64ptrj5Ipyur4xn8opV2EZPhypvoGP2wn7vvQMTT6x09OiHg1Ut2cActpd0GWU5+TvLCxcle51zRhIJlOlw83Cnw4362CqqZpPgAimWqxym+9dp0B3FStDk2Bkzb8wR3YYiUln8lk7FTk2ksixkT6e9TujYmR6j+noZPGjJgiC0Ng8bbHrTvOb1D+ej7dkeV1gBbI7b9QpXNE/vsE/g3B26nUyBZ4L9onzITTec0KFBK4amizoLaCpcOETWIbcRGzv7zl2y620f/IJQsaNO5vlNxqb080HmzN5dfVhLAaFMpvvgDM4dTzBqeNrfB2jTmbmqO6ElDjInT8fd0kJ4VdeRdjkqzAkJPi/oO3vwYqHPK3pHb91QXafljZ78Hs4ug7MkTB1MXQc6v9rnyavzEZWifd/x0AymQB+ySjC4VIx1DCKoamJYEoIWITF9+C8QHYVAUJqyB0+SZEVZqbO5KUdL2Fz2Wp9rC8SEld3vzrg5wmCIAiB8zfNb2CncKKDjRwt8t4o83d3Pj7CTO/4MK/nC02r3sEUeAKq0X+GQTfC1iWw4SVwVMLJzVnVBeYIz4lJ/+s86YFA6AUXoI+J4didd+KaM4eI665rmA/TAMptTt7ZeJTFazMY2Cmcl68fRJXdxay3fw6oGcNJekVmxtjuRIb0JfL66dj27aNk2XIyrp6GsXt3widfRcgFFyCba6kVX/VPz2mUs7aOgJrn372jEpZMgqlvQfeLA15vcZUTgyJxZt+NQOsj9YpMmc1Ju+DmeRrtys/HmJzSJO8tgqkWqmtcCEEGhcqz2FXUKxJ9/PhF+MfUP7I5ZzObczZjd/tfbGxSTMwfO5925nb+fzBBEAQhYKqqsSu7lMOOYFxulYT0QnrHh9a4YSZJErPHJvGPz3Zh07xTiOranbcYFG4bm9xg6xcajqTT1T+YOskSCaPmeoKmwkNgLfYEWuZIiEr2mXZm7tePLu++y7FbZ+M4lkXMfffWP/2tARRVOnhz3RGWbjzK2G7RvHvzsGoNt+69oDvPfn8Aq7Puk9yTTHqFt2YOITrk94DC1LMncY/0JOb++6j43/8o+XQZOU88SejFFxM+ZTKmPn2QTv/39cubfgRSZ3Ba4ZMb4Y9fQsch/j/vFO//XoFmMoF3WnBz4srPJ2iEf6dsDU0EUy3URamx/GWZ7xxav3N+JYnrh3eq871kSWbB+AXct+o+Np3YhNVd+xeAhIRRMTJv1DzGdhzr/4cSBEEQAlJc6eCDnzN5Y+0RrA43kjMaDQ35nV9wuFQu79eBm0cn0iOuegc2d3k5ae8uIEXqy4HgOBw+Rl3UxKiT6ZcQzpX9RWOh5shzMtVAbbtlGaK7+f1wQ6dOdPngfY7deSfZf76XDv98Ctlkapi1+Cmn1MZraw7zyZYsJvZpz2d3jKRzVJDX424Zk4RRJ/Pkin043VqtA6iNOhm9IvPmzCEM7uK7vEE2GAi9+GJCL74YZ04OpZ99TvZ99yMZ9IRPnkLYpMvRBevhm4fgjEyfLv8pp8oJR+4OJsjgubZ7fauDpTudrLrxt7U7rbD8VvjT1jprqDSnE/uhQ1h378a++wAOa284o9wi0Ewmp1tt1h0MRZqfEDCjTuG6oR15e30GTh+/BOvaVQQY2CmChIi650qAp/bpuQnPsezgMt7Y9QaFtkJsLhsaWrXHaJpGWoc05vSfQ8+onoF9KEEQBMFvH/9yjL9+thtJ4rR0pd9OAn6rB1m+LZuvdh5nXPcYnru2P0adgm3fPrLmziUoLY13517D9CVbOZRbgc3H3MIzmfQyPeJCeP2GwehEJ79m6azS/BqAEh5Opzfe4MRfHiZz5iwSXn4JXUSE1+Mq7S6+2nmcg7kVlFidhJn1dI0J5rJ+HQg2Bn55erSwkoU/HebrXSeYOiiBb+eOIS6s9kDuj2ldGJoUyWurD/PVzhMosnSqjlACLEYFnSwzY3hn/jiiMzGh/gWG+rg42t02m6jZt2LdsoWST5eRPvFSYsYEEx6q+Tgn8jQFe26Tg4dH15JGV54DWb9UO506GTjZ9uzBumcPtt17sB86hD6+A+bUVOJ7pRKTaSSrsv6ZTAB9E8Kbbb0UgCsvH120aI0uBGjmyETe25SJ0+3/EfVJJr3M3PO7BvQcWZKZ2m0qU7pOYUf+Dj7a/xHZFdnY3DZCDaEMih3E1G5TRVqfIAhCI1u0Op0F3x/AXkcA5FY9O+6rfs3j2kUbWdjuOKX/WUDsww8TdvllAHxyWxqPfrGHz7ZlI0tg9VFHYtYrqJrG5AEJ/H1SarO+qGrrmjqYApCNRjrM/zf5/3mOjGuvpdOiRRi6dAE8My5fW32Yz7ZnI0tStSYoFoPC37/cw6R+Hbh1TBIpMXW3H9+fU84rqw7x04F8rh/emZX3jSMyyHdduS894kJ5Zlp/Hp2Uytc7T3Cs2Eq5zUmkxUCP9qGc1zOm3iMAJEnCMngwlsGDcZdXID3XG8nlu1zi/jQD/1pnZ84QA+Em3ydPmsuG+u0TlEfM8ARPu/dgP3gQfYcOmFJ7Ye7dm7DLLsPUowdy0O+ncbdvPMoTX+/zajjjbyZTUDNP69U0DXdBAbroprn+FMFUCxYfbmbRjEHcsuSXgIoodbLEAxf1YFhSVL3eV5Ik+sf0p39M/3o9XxAEQai/b3efYMH3BwL63re5VPZmFnL/7hO8tvQdjMm/XxiZ9ApPT+nLw5f05OMtx1i89gi5ZXY0NCQk4sJM3DwqkSmDE2oceCo0H5K+AWqmGmIdskzMn+9BnxBPxvUzSHj+eX5UYnjgk1043G7cPn58T17sf7oliy93HOefk/tyxQDfHSN3HCvhxZWH2JZZwk2jEnn8yt51NtWqTahJz7VD6y59qC/FfgKkmuvOB3dQGNdFx/z1duZN8H0CJmkqUuYqKvd3w9w7ldCJEzH26IkS7J3GeLorB8Qz77/7fN7nTyaTXobzezbNqY8/3CUlSBYLsrFpmmOIYKqFG901mldnDOa2pVtwulWfKX8nKZKETpEINioY9WJXURAEoaXRNI1Hv9hbYyBVuXcVZT9/hrMwC9lgRh+TRFjaNEwJqdglhc1RXckIiqG7j+eGWfTcPDqJm0cnoWkaTreGXpGqF88LzZ5nzlTTB1MnRUybhr59e9557GUWpF6JzY89ALfmOSF9cNlO3JrG5IGeluOaprHhcCEvr0znSEElt45J4vlrB2A2+NdAoUlVFXrVLZ3pH+ONjFxcyeZVFd0AACAASURBVN3Daj5ZkxWJ+KfmVZsJWpcgo45HJvbgia9/DajhBoARlbu2Lce+KRRdWlrdTzhHDuSW886GoxzILaeirBL9kJkM+e9eZgzvQqco/0pYGooIplqBMd2i+e6eMby1LoP3N2cCVOvydzI9Y2Kf9twyOgmLQWHqwg10CDMzvkfz3WkQBEEQqtuQXkhZDXMGyzYvp3TTJ0RdeAemxIFIig7rkS1YD27ClOAZpup0q7yx9jD/mtqv1veRJAmDTgRRLVFzSPM7U1bXfizofRU2Hxu+tW0A2JwqjyzfTfe4EHJKbby08hAlVU5uG5fMlf3jW1a6qVZ3FNk7RuGybjr+udZBz+iaPpvk12ud6foRXThRZmPx2gy/AyqTXubBi3sx7Yp2HH/wIcKmTCb6jjuQdE0XPny3J4fnfjxIen4FTpfKqR+p4AT2rM9gyYaj9E0I457zu5GWcm7S/kQw1UokRFj462W9uO+i7ny7J4f9OeUUVzkIMenpHGXhsr4dqnVhWTRjILcs2cKSWUPFnBBBEIQW4tXVh73qHgBUeyUla98lauJcLN1/3z22pAzDkjLs1D+7Nfhix3H+dnlqvYr8heYrr9zG//blkaF0Rs7UiNuQwciUdiRHBzf10nhlZTq+DlP92QCwudxcu2gjCZEW7hifzCW926PILTDQN0d4BvTW4bFxJgYuquDeETWkrEky6OvXIfH+i3rQPszME//dhyTh87sEPDVSSPD0lL5c1rcDkEjisk85/sADZN44kw7PzEcfG1uvNdSXqmo88fU+3tuUWWMw6MnO0vg5o5hZb//M3PO7ctvYxp89Jb5JWxmT3r+J9IM6RzLvyt7c/PYvfDonjfjwWobLCYIgCM3CzxlFPm+3Z/+K5nJg6Vb3nBW9IrP3eBlDE323eBZaDk3T2HykiEWrD7P2UAGKJGFXO6EdB8N/9yEB3eNCuH1cMuf3jG2SDoylVU6+2ZOD+4whRf5uAGiaJ6B696ahRDbTgbF+adcNdEbPEN5apETKXJOq5/nNDvrE+Pjv1aXuNua1uX54ZyYPjOeL7cdZ+FM6x0ts6BVPcOpwqyRHB3Pb2GQu6ROHUfd7+qSuXTs6vv46ha++ypEpU+nwxDyCx5678TdPrag9kDqTzany3A+H0MkyN49OatS1iWCqDZvYpz3ZxVZmvfkzH98+QhQWC4IgNHM1XUi4rWXIllD/hm9qUGptXmlgQuBcbpX7P9nJN7tzsDndpw0q8VwYn+z0uCOrlD9/tINuscG8PWvYOZ8VtHxbFrKPurtANgAUWWLZtuxGvyhuVLICw26HNc94zZk609/GGnlnp4+/o4YgGHn3WS/FYtBx7dBOXDOkI4WVDkqqnMgShFsMtXZClGSZdrfdhmXwYLLvu5/QSycSM3cukr72n6nsEisnSqxYnW5CTHqSooMCuuZc+WseSzf6DqRqSxO1Ot3M/24/QxMj6ZsQ7vf7BUoEU23czaMTySyq4valW3jzxqGn8o+PFlby1voM1h8qpNzuRK/IxIWamD68MxenxrWsPGVBEIRWQpElVB91J4o5FLWqDE111x1QSZzaiRZaJlXVmL10C+sOFfjV1bHK4Wbv8TKufGkdX9w58qy63gVq9/EynxfBgWwA2Jwqu7NLG2N559agmbBmvtfNGXOrt4DvGCZj+2uo1+MwhkJiw50GSZJEu2Aj7QI88bMMHkzi8mWceOgvHL1+Bh2eeQZDQvWsKKdb5fu9uSxclc7+3PJq140Ol8rFqXHcMibJr1KT53886PNnyJ80UYdLZeFP6bw8fVBAnzEQ4oq4jZMkiUcv74VJp/Dw8l1sySji6oXruXDBapZuPMr+3HKOl9g4WljFpiNF/GXZTgbN+56nV/yKLcCOMIIgCMLZqWk31xjfA0mnp+rAhjpfw61qAV88Cc3LM9/tZ/2hwoDa4zvcGtklVm5buqURV+attMr3KejpGwB+vU5rOE0NjoaR94C+Ht3mdGa4/HmQm8eluy4igoRXXibkoovImDaN8h9+OHXfjmMlDH3iB+7/eAc7s0uxu1TKba5Tf+wulS93HufqhRuYtmhDjU11AA7nV7D3RJnX7SfTRCMvuB1L9zRkgwlJ0WFJGVZt4LCqwQ/78iiqdDTsv4DTNI//IkKT0ikyL/xhABvTC7nm1Y38nFGM3eW7zXql3U25zcXidUe48qV1FDfiD6cgCIJQ3VUD432eKsnGIMJHTafo+4VUHdiA6rShuV1Y03+heOXiao8NMeno1d7HrrfQIlQ5XCxeV3NHtsq9qzjx9lwyn51K1oszyP3oUWxZewDPLv2Wo8X8muN9cdpYQky+k6AC2QDwvE4rKUUY9xD0nhxYQKUzw8VPQbcLG29d9SDJMlGzZtLxlZfJfeqf5DzxJOv353DtqxsprnJW6yx9JlXzpC1vzyzh8hfW1hh0f7D5GG7V+3o0kDRRWYIvtmf7/8ECJIIpAYC1BwvIr7Dj8vED64vdpZKeX8G1r27EWstfFkEQBKHh3DCii8/6E4DQoZOJmHATpRs+JOuF6WS9ciPlW7/C3PX3iw2zXuaW0UnILbEbmgDAZ9uyqWn0V9nm5RT9+Bphw6eRcOdS4m9/k5CBE7Ee3HTqMU6Xyhtrjpyj1UJKbDBGH6UBgWwAGHUyKTFN35WwQUgSTHoRRs71NKTQ19IAzBDs+TPldRg889ytMUDmfv1IXPYp6bll3PTGhoBmWTncKsdLrFz/xiZcPiY5H8wr93ltGmia6JGC2ht/nA1RMyVQXOng7g+2nypWPV1thX1Ot0ZGYSVPrdjHP67o3QQrFwRBaFs6RloY2CmCzUeKvLqjAQSnjic4dXyNz9c0uHpQx8ZcotDIFjVAe/wvdx7n75NSCToH7fGnDkzgPz8c9Hlf6NDJyEERlG74kIKv5iMZzBhjUwgdcY3XY6cNbkU/t5IE4x6EobfAtqWw4UWwl/0+1NflgKgkT0pgryvq3Qr9XFLCwnhr8FRse3N83l/X9WR6fgU/7Mvl4t7tqz+vhg37gOpEodZUwrMlgimBD3/ORMP7l7I/hX12l8rHv2Tx0CU9sBjEj5MgCEJj+/fVfbn0+bUB15CY9DL/nNyHMEsrSZdqg1RVI7Ooyud9gaQ96WSZzKIqep6DdM+YUBOjktuxcn+ejyuNujcAAIYnRREX1vwDioBZImHkXTDiTig9BrYST0BliYKQuKZeXUAKKuysPpCPhvexqT/Xk1UON6+sSq8WTLkKCgiq9N145PQ00aAedbeLj7A0Xp2oSPNr41RV4/W1R7yKWP0t7APPBsuXO46fy2ULgiC0WQkRFt6/ZThhZj3+ZuuZ9DIPXtyDKwckNO7ihEZV5XSj1JDjF0jakyRBuc3V0Mur0R0TkjHq63fJadYr3DG+8QevNilZhojO0L4fxKa2uEAK4L1NmT5vD+R68tfjpfzy/Gscu+NODo4dR/qll5F8eCdGvDOnAkkTtRgUerYP8XqNhiKCqTbul6PFPnNbA9nhqnK4eXNdRiOsThAEQfClV4dQ/nvXKMLMehRZwuCjKYWE5yIiPtzMi9cNZObIxHO/UKFBmXSyz/ROCLA7nuYJUs6VQZ0jmXt+N8wBjlUx6xXunJAiBky3AD/sy/VZLhLI9SQuF5tLJcIuv4zOS9+h28YN3PrUXaDznfnkT50oeNKbL+vboV6fyx8iL6uNO1Fq9Xl7QAMggbxye0MuSxAEQajD0cIqQkx6Pr09jQ9/PsayrdmU2Zy4VQ2zQWFol0hmj01mSJcIpJo6Fggtik6RCTHqKPNxqhRI2pPdrRIbdm7b498yPIHC997nnch+2LS6gyqzXmHOuGTmjEs+B6sTzlZNaceBXE86FR3aqHGEnnYSGRNiYnTXdvy4r35pojpZYuqgBMyGxts8EMFUG2dzuvG1yRVoYZ/Dx26EIAiC0Djcqsa8/+7joUt6kBQdzF8m9uQvE3s29bKEc+C6oZ1YvO6I1/iS09OeJFnBlDgASdZhy9iOLXNntZSqfglhxIScuxokTdPImzePP+qLGT9rNi+sTGdrZjGqplX7HHpFQpYk+ncM567zujIypd05W6NwdnQ15BwHcj0pSxKKj9e567yurPVzQPWZ9Iqng2ljEsFUGxdi0vtssxtoYZ+lESN+QRAEobpPt2YRZFC4pHfLq60Qzs6MEZ15a30G+Nin96c7XpBR4bax5/a0p/i997Bu307n9z+gY3AQaV2jOVZUxfubM9l3ooxym4tgk44ecaFMH9aJjpH1GGorNKmoICPp+d7txwO5njToZCKDDF63900I5/ErevN/n+8OKKAy6WVenj6QTlGN+/Mkgqk2LrVDKE4fff0D2eGSgD4JYedw1YIgCG1Xpd3FM9/tZ+H1g0T6XhuUEGFhWGIkG9ILcfqYv1NX2lOQQce47jGNucRqKjdtpuDlV+jywfsowUGnbu8YaeGBi3ucs3UIjWvakI7sPl7q1bY/kOtJt6pxfs9Yn69/9eCOKLLEI8t34XCr+Lh0PUWvSOgVTyB1Ln7WRTDVBpXZnBRVOHCpKqEmPakdQtmaWeL1OH/nP5gNCrPHiJxmQRCEc+HV1YcZmhjFgE4RTb0UoYksuKY/E59bQ36FHR/xVI0sBoW3Zg71mUrVGBxZWWTfey/x8/+NoWMrmhMleLmsb3v+9vlun/f5cz0pSzChR4zPk6mTJg9MoF/HcN5Yc4Tlvw2vPj14CzIoIHlSYW9M60JCxLk54RTBVBvhVjVWH8hn4U+ePGW9IiNJ4HJrmPQKelmq1w4XQITFwJAu4pe6IAhCY8sts/H2hgy+vLPu9Guh9YoKNvLpnDSuWbSBvHK7V/3UmWTA5LLz1vRB9OrQ+LOlANTKSrLuuJN2t95K0Ag/OrkJLZpJrzBtcEfe3XTU589jXdeTRp3CzX7UNiVHB/Pk5D48cmlPvt51gqOFVZRanYRb9HSNDeGi1FiMunNbeiKCqTZgy9FibntnC1UO16lJ0k7375G83aX6GLHmH5NO4qFLeohUE0EQhHNg/rf7uXaIqCkRPOl+X981hmd/OMDHvxwD8EqxMullNA0u6BXLzMzVRD//DdqihUhy407G0VSV4395GFNqKhEzrm/U9xKaj3su6MYP+3I5UWKrsYW/L2a9wlUD4hnU2f+N+SCjjqsHN4/TThFMtXIrf81jzrtbsNZRsHfqR17TPNP8/GCSNCZnbOACcysfpicIgtAM7Dleysr9+fzvvrFNvRShmQiz6HlsUip/uaQHX+08wSdbjlFY4cCtaoRZ9FzUK5ZrhnQiIsiA5uzD0ZkzKVy0iHa3396o6ypYuBBXbi6d5v9bbLa2IWFmPR/OHsGUV9ZTWFH3iSl4AqnxPaJ5/Mre52CFjUMEU63YzqwS5ry7tc5Aqho/vvR0soROlph7QVeuLYXMWTfR8ZWXMffpcxarFQRBEGqiaRpPfr2Pu89LIdSkb+rlCM2MSa8wdVACUwcl1PgYSa8n/plnyZg6FXP//o2Welf+44+UfPQxXT76ENlQc/2L0DrFh5tZcddo7v14B2sPFQC+x+cE/dYF+raxydw5IaVFB90imGrFHlm+C6vT9yT0yr2rKPv5M5yFWcgGM/qYJMLSpmFKSAU8Hfr0Os+8B5dbQ5Yk9DoJVYWrBsQza1QiKTHBQApKcBDHZt9G/IIFBA0beu4+oCAIQhuxcn8eOaU2rhvaqamXIrRg+tgYOvzraY4/8CBdPvkYfazvzmn1ZTtwgBP/9zc6LlqIPubcdQwUmpeIIAOLbxxCTqmNpRuP8t7mTEqqHIBnllRSdBC3jU1mYp/2mPQtf7SOCKZaqQO55RzMq/B5X9nm5ZRu+oSoC+/AlDgQSdFhPbIF68FNp4IpnSxxcWocw5PaUWZzYlBkYkKNTOgRg8VQ/ccmZMJ45GefIfuee2j/5BOEjBtX7f5ym5NPt2Tx9a4TFFc5kSWJyCADVw2I5/J+HRp1KrUgCEJLoTnduMudaA43kkmHEqxH0sm43CpPfv0rD0/siU5p3FoXofULGjGCiD9cR/af76XzW28i6RvmpNNdUkLWnX8i9sEHRKaKAEBcmIn7LurOfRd1R/ttQLNB1/q+w0Qw1Uq9sfaIz/lRqr2SkrXvEjVxLpbuaadut6QMw5Iy7NQ/O1WNH/bl8a+p/fzaNQgaPpyOr7zMsTl3oD78F8IuvZQTpVae+e4AX+44jixJXqdkO7JKePSLPUwdlMA9F3SrtR2mIAhCa+XIKqd8TTbWPQVIsuRJt1Y1kMAyOJYfLRAT4tnMEoSGEDV7NlXbtpH3n/8Qe//9Z/16mstF9p//TMh55xF2xRUNsEKhtZEkCYOu5aby1UYEU63UD3tzfQ40s2f/iuZyYOlWd660IklsyyxhRHKUX+9p7tePTosXc+yWW9iXZ+W2zFDK7c4aB6ud7Dr0weZMvtuTw4ezR9ClXZDvBwuCILQy7lI7BW/vwZVvRXOpoIFG9YLtyo0nGKhqDOoc6jmxMopf28LZk2SZDk8/zZEpU7AMGEDI+eef1evl/fvfIMnE3HdvA61QEFqO1nfWJgBQYXf5vN1tLUO2hCLJdZ82aUCp1RHQ+5q6d0N74TVm7ZEpsTpqnVB9klPVyK+wM+WV9eSV2QJ6P0EQhJbIVWAl97mtOHOq0Jwq1NT0SgUDEvrsCvJe2I5a5Tyn6xRaL11EBAkLFnDib4/iOHas3q9Tsmw5Fat+Iv7ZZ5AUkbYvtD0imGpjFHMoalUZmuq7McXZ0jSNW7/LxqY3whnTqyr3ruLE23PJfHYqWS/OIPejR7Fl7QE8GS2lVie3v7u1UdYlCILQXLgrneQt2olqdXm+/Pzh0nAV28hfvNtziiUIDcDcrx/tbr+drLvvRrXbA36+dft28ubPJ+Hll1DCwhphhYLQ/Il8gVYq2KjD7vI+VTLG90DS6ak6sIGgHqPqfJ1wS2B1TJuPFJFXbve6PvCn6YVL1diTXcqhvHJSYkICel9BEISWonzVMc8J0xnfk5uzdvLkylc4UJCBLMt0jerMo+f9if7te3oe4NZw5VZRtTOfoIEN24VNaLsirp9O1ZYt5D7xJO3/8Rh2l5sVu3JYtPowGQWV2F1u9L81oZqZlsjUwQmEmvQ4c3PJunsu7Z+YhzE5uak/hiA0GRFMtVIX947jw5+P4TojqpGNQYSPmk7R9wuRZAVT4gAkWYctYzu2zJ1EjJ916rGaptG/Y3hA7/vq6sNYz5jA7m/TC/AEVIvXZvDkZNEJSBCE1kdzqVRuzoEzhlmW2yuZ+clDPHHhn7m8x3gcbhebs3ZgVKpvaGlOlfKfskQwJTQYSZJoP+9xDk+9mqdf/IJ38o1omkblab/L7S6VY0VW/v3tfp7+5leu6BvHrE/+Rbs//IGQ8eObcPWC0PREMNVKzRqVyCdbsryCKYDQoZORgyIo3fAhBV/NRzKYMcamEDrimlOP0SsS1w3tFFD//3KbkzUHC7xS/wNpeuFSNZZty2Lelb2R5dbZ9UUQhLbLursANO/v5cNFnpqVK3t5GgGYZYWxib7n9rmLbDiyKzDEBzfeQoU2RTNb+OekB1l7pBi74rvmGjjVlffzLcfY0uUKPp0x+VwtURCaLRFMtVLJ0cH0bB/K9mMlPu8PTh1PcGrNu0myJHFDWpeA3jO/3I5ekTjjYCqgphcAblWj3O4izNwwsy8EQRCai6qdBWgO75qnpMiOyJLMPf99gkk9zmNAfCrhJt/pzppLxba/SARTQoPQNI37P9nJuhNW7Ip/qf12SSFTF8qMNzbzye0jMOpE4wmh7RLBVCv21OQ+TH55HVZnYMXKZr3CjWmd6RhpCeh5VqcbSfI+TTq96YU/AZUiS9icbhFMCYLQ6qiVvrvxhRiDWDb9RV7e9B4PfPNv8iuLGJ88jH9d/ADRQZHVH6yBu/z3mthSeynLDi7jm4xvKLWXIiERbgxnYtJErky5khCDqEEVarbmYAHf7snB5uNaoXLvKsp+/gxnYRaywYw+JomwtGmYElJxujUO5pXz5toj3DYupQlWLgjNgwimWqkqh4ttmcWEmvRYnf536DHrFSb2ieOBi3sE/J6hJj2qj/SVQJteuNwaISbxoykIQtvStV0XFlz6MACHCo9y11fz+PuPL/DSpEe9HyxBdkU2C7YsYNWxVUhI2Ny/j5bIqsgivSSd57Y+x4WdL+TugXcTGyTqrARvC39KPzX38XT+NI6yOVVeX3uEW8cki9R8oc0SV6yt0Hd7cpj74XYAn1+QvlgMCqqmMWdcMndOSPF5wlSXmFCjz9sDaXoBEG7RYw6gVksQBKGlkIP9O3FPierMtN4Xs3T7Fz5eBA4ajzL3y/+j0lGJiu/sA6vbCsDXR75mTfYaFl+0mK4RXeu9dqH1ySquYsvRYq/bA2kcZXW4WX0wn3HdYxp9vYLQHIk5U63MR78c464PtlHlcPsdSBkUiQcv7s4vf72AP53XtV6BFIBRp3DN4I7oFe/nhw6dTMSEmyjd8CFZL0wn65UbKd/6Feau1ZtSmPQyN49KrPcaBEEQmjNLv2gkg/ev3kOFR1m0+QNOlOUBcLwsl8/3/cjADqlej80y5XFX/sOUO8prDKRO59bclNhLuPGbG8muyD77DyG0Gt/szvE5LzqQxlGVDjcf/5LV8IsThBZCnEy1IusOFfC3z3f7zHuujSxL/JxRzA1piWe9hhtHduG9zZl4DVCh7qYX4Glydc2QTme9DkEQhObInBpF8TLvzaIgg4Xtx/fx2s8fUWavINQYzPnJI3hk/Jxqj9PQeDT+Jap+O3U6XfGaYgq+LcCR50AxKYQOCiV2aixKkOekv8JZwd3/u5tPJn3SOB9OaHFyy2w4fAyBDrRxVG6Zre4HCUIrJYKpVuT/Pqs5kKqtiNTmVPlhXy6/5pTRIy70rNbQOSqIMd2iWX0gH7uPL+jamPUyVw1IICIosEHBgiAILYGmaXy5O4cjqoMrJAXdaXtO7UOieeXKx+p8jb3BBymWC9HO2LAqWFFA/op8Em5OILhXMM5iJ8ffOU7G/AwSH0lE1smomsrRsqPsK9xHz6ieDf3xhBbI4fb9ezrQxlE1vY4gtAUiza+V2HGshBOlvneGyjYvp+jH1wgbPo2EO5cSf/ubhAyciPXgplOPcbhU3lhzpEHW8vy1A+gSFYRB5/+Pl0kn0yc+nMeu8E5pEQRBaOn2nSjj2lc3snBVOqOmp2IMM0Kg2cyKxKcdfsImVw+k3FY3eZ/l0eH6DoT0DUHSSRiiDXSc0xFHgYPS9aWnHutQHSzZu6QBPpHQGkQFGfHVN+L0xlH+CLeITVCh7RLBVCvx2prD2F3eNVIni0gjL7gdS/c0ZIMJSdFhSRlWrfGDqsEXO45TbvPdtjcQZoPCp3PS6JcQhsVQ946WxaAwqms7ltw0FL0ifiQFQWg9SqucPPr5bma8sYnL+3Xgyz+NYlD3aKJn90UONvj/W1gn44qR+UW/2+tUqupgFapTJXRQ9cwCxaQQ0jeEij0Vp25TNZXvMr7DqZ79d73Q8g1LjMTko+HT6Y2jqg5sQHXa0NwurOm/ULxycbXHmvUKE7pHn6slC0KzI9L8Wom1hwpQfVSRBlJEalBkdmWVkpbS7qzXE2zU8cGtI/hhXy4LV6Wz90QZwKncbIPqQtPpGJwYxa1jkhjbLVo0nRAEofk4sQPWvwgHvwNHJUgSGIKh1yQYfgdEd6v16W5V46NfjvHMdwe4uHcs398ztloKsy7CROzdAyhcug9nVgWaqnJmL4lsfR6fR/2P7UEHsBrsyGY9bqv3ppm7wo0uWIfko/mPLkyH9egZ9VUSlNnLiDJH+f/vQ2iVhiZGEmExUOXwrsELHToZOSiC0g0fUvDVfCSDGWNsCqEjrqn2OFXTmDIo4VwtWRCaHRFMtRJVdt+d+wIpItWAEmvD7VYqssRFqXFclBrH4fwKVh/Ip6TKiSxL6DasYZTFTupNE3GrmgikBEFoHo5thi/ugpIMcDlAO+271VoE25bCjg8gpidMehHienu9xNbMYh79fA9GncxbM4fQOz7M51spwQZibuuHM6+KinXZVG3NQ9M09pkP81rUp6Qbj6FKLlzSbztl3te7v72OgqvChebWvAIqV6kLXXD1X/WKpGB3+z9/UGi9JEli9pgknlrxK1an93VEXY2jZAku79eBEJN/Lf8FoTUSwVQrIcuAj3gqkCJSCdA10tC9pOhgkqKDAU/twMI9HXg+y4794a+RJE/gNahzBLPHJjO2a7QY/icIwrm393NYPhucNUQtAKrL8+f4NnjjQrjufUgaC0BeuY2nV+xn3aECHrqkB1f07+DXRpE+xkLEVV0Jn5TCl/u/4PGtL1QbwFsXS4oFSSdRtqWMsKG/B25um5vyneXETq0+rNelugg2BPv9+kLrNmVQAq+uOcyJEhtuzVej9JpZDDruPk/MLhPaNhFMtRKhJj02p/dO4+lFpEE9RtX6GhoQFdx4RaQHcsuZ+8F2DhdU4HSpuGXPj5+mgerW2Hi4iJ1ZpVgMCo9NSuXSvh0abS2CIAjVHFkNy2aDq5ZA6kzOSnj/Wpw3rOCtw6G8vOoQ04Z05Id7xxJsDPzX66rsVTy+dV5AgRSAYlGIuTKG40uPI5vkat389JF6wtPCqz0+1BBKiD4k4PUJrVOQUceHs0cw6YW1lFiduH3VDPhgMSi8PWsIHSMtjbxCQWjeRDDVSlw1IJ7F647gdFf/Ejy9iFSSFUyJA5BkHbaM7dgyd1ZrQqGTJfolhJ/50g3i54wibly8mSqH2+eAwJNODhu+9+MdZBVbmT02uVHWIwiCcIrbBR/d4BVIdflPOVVOOHJ3MEEGzwnT61sdLN3pZNWNQQBoziry3riGtQlL+OT2NJKj63fiU2ov5YHVD/gMpOqaHwUQPTEaJUgh58McHHkOZLNM6MBQOs7uiKz/vcuFLeJCAAAAIABJREFUSTHxx15/FKnVQjXx4Wa+vns001/fxIkSK5UO36UDAEEGBaNeYcmsoTWmsApCWyKCqVZixojOvLU+A1/Dcv0pIjXqZG5I64KuEbrpHcgt54bfAil/2ZwqC344QGSQgasHd2zwNQmCIJxyYAW4HT7vcmvw3CYHD482+rxfAmLlMt66QEOqZyAFsPzgcp+3+zM/6qTIsZFEjo2s9X1UVCZ3nVzvdQqtV2yoie/mjmF9eiGLVqez9lABJp3iSf+XwOlWSY4O5raxyVyUGhfQ+BNBaM1EMNVKJERYGNQ5go2HC3129auriBRgeHIkz/1wgJwyOy5VJTrYyOiu0QxPijyrXcx7PtyOtYZAqq5hwv/32W4uTI0jzCyKWwVBaCRrF4Cjwudd96cZ+Nc6O3OGGAg3+f4e1LmtsP556DS8Xm+vaipL9i7xOpU6OT8q/qZ4Qvp60vJOzo86cP8BSteXEjEmwu/3MSkmrki5gnBT42QgCC2fLEuM6tqOEclR9P37tzx7TX8kPLVRCRFmurQLauolCkKzI4KpVuTpKX259Pk1lNlcAT1Pr0hEWAzMfPNnTy3TacHYW+szCDPrmT0miasHdyQowDqA/TnlpOdX+EztK9u8nNJNnxB14R2YEgciKTqsR7ZgPbgJU4JneK8kSXy6JYtZoxIDel9BEAS/VBVBzq4a7x7cQWFcFx3z19uZN8FUw6M0OPCtJ11QCfzX6q6CXVQ6K72X5sf8KH+DKZNiom90Xx4a+lDA6xPanvT8CqJDjFyUGtfUSxGEZk+c0bYiHSMtvHfLcEJMOp8TzX2RNBVN08gps2FzVg+kwFPDdKLUxtPf/MrE59ZwojSA4mzgjbWHcbpUr9v9HSZsdbp5dfVhtAA7DAmCIPilMh+U2hvv/GO8kRc2O8iv9P4uO0VSwF5WryXkVub6PP2va36Uq6LujTMJCbPOzOiE0Sw8fyE6WeyhCnXblllM/47iBFMQ/CGCqVamd3wY//3TaPp3DMeok2tsdW7Re+5TZAkfsY4Xq1Mlq8TK5S+sJb/c//kkX+/K8QrQILBhwmU2JwdyfafgCIIgnBW30zOQtxa9YxQu66bjn2t911UBnteooe6qLna33eeG0enzo87ka34UgE7SoZN1mHVmDLKB0Qmjeem8l3hm7DPoFZEuLfhn+7ESBnTyP4VUENoysUXVCnWKsrBszkjS8yt4c90RPt9+nAq7ZwdTL8v0SQgjxKhj45FCbE7vi4ja6phKqpzc+OZmvvrTqDrrqFRVo9Lhe+c0kGHCiixRWGkHRCtfQRAamCnMk55Xh8fGmRi4qIJ7R/huRIHbCfWsRQoxhPj8Pg10flSQLojb+t2GJEmEGcMY2WEk0Zboeq1JaNu2ZZZw7ZBOTb0MQWgRRDDViiVHBzPvyj7Mu7IPqqrh1jT0ikyZzcmQeT9g93EkVVcdk0vVOFJQydbMEgZ1rn3XSjv1P94CGSYM+D33QhCEtiWvzMa7mzL5cudxyqxOwDN3b2KfOK4f3oW4sJrqnH4TGg+GoDrnS6VEylyTquf5zQ76xPhI6ohMAn0d71WDHpE9cKpOr9sDmR8lITEgdgA39r6xXmsQhJMq7C6OFlbRs31o3Q8WBEEEU22FLEvIeHY+P/0lC9nHLujJOqaoiXOxdE87dbslZRiWlGGn/tnqdPP6msMM6jzo1G2a04nj2DHshw7hSE/Hnn4Ye3o6+pQ/4vCRWhLQMGEN0c1PEIRqjhZW8tiXe1l7qAAJqm0OFVQ4eG3NEV5dc4QRSVE8enkvkmpqWy7LFPe7haCNz2LQak9h/ttYI+/s9A56MATBqHvq/VniguLoH92fzTmbve7zd36UWWdmVu9ZXs8XhEDtzCqhV4dQ0fpcEPwkgqk26PW1h7E6vVuV+1vHpGnw494cDj33MqYjB3GkH8KReQxdbCzG5GSMKckEjRpJ5A03MHRDGWsPF3u9RiDDhDWge5xI8RMEwWP7sRJmvLGJSrvL5ygI+D24+ulAPuc9+xORQQaCDTpiQk1cP7wTF/eOI6/MzksrD7FpdzLf+zhGz5hb/XunY5iM7a817Nb3PrvZTbN6z2J3wW6qXFVe9/kzPyrUEMrg2MFntQZBAM/fL9F8QhD8J4KpNkbTNI6X2nzeF0gdk97tJNsh0f+88zDOvhVDYiKyyTvF5TZTAVuzfvE5sNefYcJ6ReK6oR0x6upekyAIrd+hvAqmv76RSrv/Q8A1DQorHBTi4GhRFbuzS7n3ox0oMtyYlsjy+6+g6oct6Lcuxoz/DXYA0Ftg/COgNwf4Saob0WEE8cHxHCk9gksLbLyFSTExd9Dcs5oHKAgnbcss4Yr+HZp6GYLQYohgqo2xOt3IkoTbV+eoAOqYZLMZ/eSrCUtuV+vjRqZEEWrS+wymoO5hwrIkcWNal1rfQxCEtkHTNG5d4ntzprbGOWc6eTIvyzIr9+eRltKOB3ZO4MOYA3Qp3gBO79Mhn/QW6HcdDJ9zVp8LQJZkXrvwNaZ9NY0ia5HfAZVJMXF9z+u5NOnSs16DIGiaxrbMEh69vFdTL0UQWgyRENvGmHQKag0zm06vY6qTBsF+DPCVJInHr+yNSR/4j5pZr3D1oAQSIiwBP1cQhNZna2YJJ8psnPkVVrZ5OUU/vkbY8Gkk3LmU+NvfJGTgRKwHN9X6eg6XSnpeBTPf3Mz/Xd6bLrd9AgNngM4Eci11mooRdEYYORcufabO1ur+ijJH8eFlH9I5tDMWXe3fezpJh1ExMqf/HO4edHeDvL8gZJdYkSSIDz+7k1ZBaEvEyVQbI8sSUUFGCiq8U1kCqWNyuFU6+Plle0GvWB68uAdPf/MrNqcfQ60As16jR3ImKV1zWbhjHUH6IJLCkhjefjiKH2mIgiC0Pq+tPoztjHpPfxrn1HZq5dY86cTr0wu5rF8HuORfMOw22LQItr0DkgycDJY0QIYhN8OQmyAsvsE/YztzOz6e9DGrjq1i8a7FHCw5iCzJuFQXsiSjk3SoqExKnsT1Pa+nS1iXBl+D0HZty/TUS4mUUUHwnwim2qA/jujMSysP+WyN7k8dE8DQxEjaBdcwb8WHmSMTiQoy8MCnO5Elqca0P52+FCV8A6Z2mzmhk3h+mwOX6kIn6zDIBow6IzN6zWBK1ylEmMRAQUFoTSrtLr7ZnUNWsZVyu5MIi4HusSGM6x6NqsGPv+Z6nUrV1TinrnEPAE63xrJtWTxyaU+CjDpPm/NLnobz/w7Ht4G12BNUmSMhfiA08vBbvazngs4XcEHnCzhccpgd+Tsod5RjUAxEmaMYFT8Ks06cHAgNzzOsVzSfEIRAiGCqDfrDsE68uPJQjffXVccUZFCYPSY54Ped1D+eCT1j+WxbFgtXHaaw0oFO8ex+aZoGlr3o4t7Drbpx4cJ1WsmAU3XiVJ1UuipZtGMRb+x6g1fOf4X+Mf0DXocgCM3LobxyXl9zhM+2ZyNLElaH23MGJIHZoKBXZK4elPDbSIfq0VRtjXP8HfcAnvrMz7ZlM314599v1JuhcxpNKSk8iaTwpCZdg9B2bMss5v6LejT1MgShRRHBVBvULtjIhb1i+X5vrud0SnKhC96LbCgA2QaqCdXRDldFL9Cq/4jIEkQFG0lLjqrXewcbdVw/vAvTh3UmPb+SokoHLlVlX+laXtn7Hna3vc5KPpvbBm645btbWHjBQgbFDqr9CYIgNFvvb87ksS/34HSruM84LFc1fuva5+bt9Rk43IE1zvF33ANAlcPNWxsyqgdTgtCGOFwq+06U0zchrKmXIggtigim2qinp/Rld+4X5PIjSthvRdqyA0nS0DQJVAPwKY7iYTiL09Bc4Uh4gqF3bhqKLJ9dPrUkSaTEeIZo7i/az8INT3gCqQDY3Dbu+PEOPr/ic2KDYs9qPYIgnHtL1mfw1Ar/ail9BVJQ+wDwQMY9AOSXBdgWXRBakX0nyugcZfGkugqC4DfxN6aN2pK/nsrop9C5nCBVr1+SJA0Uz0WFIXIdhsgNOE9MJ8jdlw9uHU7nqKAGXcuinYtwuB1etxevKabg2wIceQ4Uk0LooFBip8aiBP1+YeRwO3h337v8efCfG3RNgiA0rk2HC3lyxT6/m9LUpLbGOdaDG/0e9wCexjqC0FZtyywW9VKCUA8imGqDVmau5IHVD3jS5eo4YJJkT4qNOf49/ja8D91iQxp0LcW2Yn7K+gmV6hcxBSsKyF+RT8LNCQT3CsZZ7OT4O8fJmJ9B4iOJyDpPLqBTdfLRgY+4c8CdGBRDg65NEITG8+z3B2oMpAKZGQU1N84JGTwJ6+FffJ5a+WIxiE6hQtu1/VgJaXXMjhQEwZsIptqYw6WHfw+kAuDGwRM/P0LfuBSSwhquGHr5weVIZ0R0bqubvM/yiL8pnpC+nuDNEG2g45yOHLj/AKXrS4kY83snP03T+DHzRy5JvKTB1iUIQuM5VlTF9mMlPu/zp/ueLzU1zvF33IME9IkXtSJC27XtWAlzxqc09TIEocURwVQbs3jXYpyq0+t2f1LqnKqTN3e9yeOjHm+w9WzL2+ZVK1V1sArVqRI6KLTa7YpJIaRvCBV7KqoFU1WuKvYV7hPBlCC0EEs2ZPgcHh5I9z1/+TvuwWxQuLUeXUoFoTUoqnRQVOEgJTq4qZciCC2OCKbakApHBd9kfINbq14j5W9KnVtzsyJjBQ8OfZBgQ8N84ZY5yrxuc1e40QXrkBTvHERdmA7rUavX7UW2ogZZjyAIjW/7sRKcPhpKBNJ9LxD/z959x7dVnY8f/5yrYXk7ju1sZzmbhCwSEhJCgLbMUvZIWCGMFgq00PbbFmhpgVLKj10IlB0IDbvsGQJkb0L2Tpxlx4m3LWud3x9XTmRLsiXZTmzreb9eekXWXecq0tF9zjn3OY1N9wCQnmjjxD6ZzXpcIdqKVfnFHN8jo8nJpYSIRxJMxZGPtn2EoermHY92SJ2hDD7a9hGXDbyswWM53V4+X7ufLYUVlFa7SXfY6JmVzFlDO5NkP/KxCzXxpCXFgqfCg/bqoIDKU+rBmhL8sU22NW9SDCFE7FweH1+tL2Dj/jKKq9ykOmzkZiZy5tAupDlslDs9IbeLNvue1VD4tMYXOtFfxBw2gz+cMRCl5EJSxIeKys0cLJpDTU0RoNmbrxnXQ6YZESIWEkzFkR+LfqTaU7dXJ9ohddWean4s+pHLCB1M5R+q4sX525m9NB8FVLqO9IIl2S3c/f4azh/ZjekTetMnO4XctFwW7VtUp7csKS8JZVWULS8jfcyRexi8Ti/lq8vpdFHdNOh2i53uqd2jfj+EEM1rf6mTVxfuYOainfi09s8RZUqyW7jnf2s5Z1iXsNs3NGdUfQq4ZHR3Kmu8fLZ2PzaLgcfnizo7YKLNwjXje/KLEd2i2k6Itsbn83Cg6At27nyWysotaO1Fa3PYf45hoYvlDRYtfopePW8kJ+csDMN2jEssRNsgwVQcKatpniF1pTWlIfc/Z0MBN7++Eo/PF3IIT5U/sHpzaT7vrtjNQxcO48J+F/Le5vfweo9cdFmSLOT8Ioe9r+3FcBh1hh7aMm1kjK+XulUj90sJcYwt3HqQ6a8sxe3TuDzBAU3t9//9VXvRYbqSGpozqr5Eu4WRPTO5aFR3SqvcLNhaRHGVm3Knmxfmb6e82k11A4GVxQCbxeDW0/K4aZLcKyXaN4+nkh9WX09Z2Y/4fFVByy3KzNxbWbmRDRvuYlf+y4wY/jI2myRlEaIxEkzFkSRbUtBrzTWk7uv1Bdw8a0VErcIen8bj0/z+ndU8cP5Quqf0YEvp5jrrZJ+VjSXZwv7Z+3EVujASDdJGptHjxh4YtiNDFRWKk7qdRFaipHMV4lhZtO0g015e0mDwUsvbwJi86q1LMRwpFP3vQQ59loy9Sz/STrwEPK6g7Hs+rTnzuM4ApCfZOHPokR6va0/qzedr9/PMt1vZdqAChcLt9WEYCpvFwOfTnD+iG9Mm9CIvp3mnexCitfF6a1ix8nIqKjajdfCcjkHr+6qoqNjAsuUXccLo97FaZRi9EA2RYCqO9Errhc2w1cnmF+2QOquy0SM1t85r24squWXWypCBVEPzxTjdPv7wzmqS0iZj6bwLL3Wz+mVOyiRzUsM3hCdYEph23LQG1xFCtJyCMifXvbI0ZCAVzXxRgSnRvc4Kyld+jDN/Dc4960nsMbRO9j1DwS+GdyM5IfRPmN1qcO7xXTn3+K5s3F/Oil3FlFW7sVsNslMTOHVgTp17N4Voz9Zv+COVlVvR2sWUK3ZRXOzFMMBqVQweksDtt2eTk1P3+6C1i+rqPaxZexvDj3/+GJVciLZBfk3iyHl55/H8mrqVYrRD6jw+zYufZrNn+4+cPbQLJ/bJ5LnvtuH2Bl9IRTJfjNurGZQxmry+BXy+4/Oo5r9yWBxMO24aw3OGx/BuCCGawysLdoQc1hvNfFGhUqKnHv/TsMfUwPSJvSMq34DOqQzoLL1PIj7V1BRyoPBTfAE9Un+/rxOjRiXhcvl4/PEinnqyiL/9vXPQth6Pk+LiBVRV7SApqddRLLUQbYsEU3Gka0pXRmSPYPH+xXVej3RIHcDYLqO459xz+GTNPh76fAO7D1VR6vQEDd2JZr6YlfklPH75n/BpH1/t+iooSUYoDouDKYOmcNPxN0X7Ngghmonb6+O1RTuD7pGKdr6oaFOiG0rRMTmhaYUXIg7s2TMLlDJbIOqx2w1OPjmFp/99EICH/lmIPUFRWOBh9Won9/6tE6NH28nf/QoD+v/lKJdciLZDgqk4M33YdH448ENQD1AkQ+ocFgfXD7ue3I5J3DSpLzdN6sszc7fwyJeb8NZbN5qLIwW8u2Iv90+6nxO2nMCzq5/lkPMQTo8THfALYCgDu2GnR2oPbh5xM6flnhbpaQshWsDX6wvwhph8N9rgKNqU6A6bwcHKGjok26MqrxDxRGtN/u6Z+Hw1IZc7nT7mflPBoMFHGibmzKnggQe6cN/9CXg8Gq3d7Nv7Nv3y/ohhyPdNiFAkmIozJ3Y5kamDp/L6+tcj6gGqlWhNZMqgKYztUrdVeV+pM+QQn2gujpweH6t3l6KU4vx+5/OLvF+w6sAqXl/3OtvLtlPlriLZlszAzIFMHTyVgZkDIy63EKLlbNxfTlVN/aaU6IOjaFKig5l4xuVp4uRSQrRzHk85Xm9l0Ot/uacAi0XhdPrIyLDwjwePJG8ZPz6Z445zAGC3m0mpNBqXqwiHo+vRKbgQbYwEU3Ho1hG34vF5+O+G/0Z0j5LD4uDSAZdy64hbg5YVV7lDbBH9xVFp9ZH9KKUYkTOCETkjGt1OCHHsFFe5Q40eivr7H01KdDAzAqYlys+XEA3xeitQynp4Lqla9/7NvGfK69UsWFDFHb/dywsvmnM15mQHf6+UsuDxlB+VMgvRFhmNryLaG6UUd4y+g3+e/E8GdBiAw+LAqPdRMDBwWBwM6DCAB09+kDtG34FSwXNRpYbJphV4cRSJlDD7EUK0XqmO5vn+GwnJZEyYwqEvZ1C1aSE+txPt9VC9dRnF37wYtL7DZtAlPbFJZReivbNYEtE6uOf4yHLFxInJGAas+dHfsBr8M4/WPiyW4KlVhBAmuYKNY6fmnsqpuaeyqXgTb6x/g80lm6l0V5JsS6ZfRj8uH3Q5/Tv0b3AffbKTSbAa1NS7AT3w4kgZFhy9R6AMK84dq4Lmi7FZFH2zU1rkHIUQLSc3M4kku+XwhLy1ovn+10obcwFGcgdKF86m6KOHUfZEEjrl1UmJDpBgNbj2pN5YjBBXfUKIw6zWNJQyCHFbI2DeU7VgQRXl5T5ye9pZtCh4Ml9zRQ92u8zlKEQ4EkwJ+nfoz1/Gx5ap5xcjuvHQ5xtDLov04shQisvG9Ijp+EKIY+fMoV24+39rQi6L9PsfKGXIZFKGTG70uFeMzW10HSHinVIWunS5gD173gQ8h1+/+64CDMNM8tepk5U//CGbXr3CJZcwyMr+CRaL9AQLEY4EU6JJslISOKV/Nl+uLwjZ+hXJxdGw7un07CgzrAvR1qQkWDlveDfeXpZPiDw0EQdHkUq0Wbj0hB5kpUhadCEi0aP7tezb9y4+nxlMvT4rfEPE7/+QE/SaYSTQM3d6i5VPiPZA7pkSTfbLU/qSYI3to5Ros3DLqf2auURCiKPl+om9scX4/Y9Gos1gXN+O3H3O4BY/lhDtRXJyH9JSh6JU9G3nSllITMwlLW1YC5RMiPZDginRZCNyO/Cb0/uTaIssDXKtRJuFa8b3ZFL/7BYqmRCipeXlpHLvz4fgsEX3c+KwKgZ1TiXBamCzhL//yWKYCSfOG96N/1w1Wu6VEiJKQ4c+hc3Wgegu+RQWSyrDj3++pYolRLshw/xEs7jh5D54fJon52zG6fY1un6izcKV43ry+zNkzigh2rpLT8jF7fFx3yfrI/7+nze8K/efP5Q9xdW8vGA7s5fm+zOGmuMFFQqPT3P+iG5Mm9CLvJzUFj4LIdonuz2L0aPeYcWKy6hxFaG1q8H1lbJhs2UwcsQsmVtKiAhIMCWahVKKmyfnMbxHBo9+uYkf95Ti07rOhL4WA2wWg345qdx2Wj9OH9zpGJZYCNGcpo7rxcAuaTz65SaW7SxGa42rzvdfYbcocjsm8+tT8zh7aBeUUuR2TOKec4fw+zMGsmJnMcVVbnxak5FkY2RuB5Jl2gQhmiwxsRtjxnzMzp0z2L1nFuALmtDXYjHvXe7a5RJ69foldnvHY1BSIdoe+ZUSzeqkvCxOystiR1ElMxftZP2+MsqcblISbPTvlMKVJ/akXydpYRaiPRrdK5PXrz+RPSXVvL5oJz/kl1Dm9JBkt9AnO4WpJ+YypGt6yG0dNgvj8yT9shAtxWZLIy/v9/TpczsHDnzB/oIPcbmKAI3NlkmnnLPJyTkLi0USvAgRDQmmRIvolZUsN4oLEae6ZSTKEF4hWinDsNOp0zl06nTOsS6KEO2CJKAQQgghhBBCiBhIMCWEEEIIIYQQMZBgSgghhBBCCCFiIMGUEEIIIYQQQsRAgikhhBBCCCGEiIEEU0IIIYQQQggRAwmmhBBCCCGEECIGEkwJIYQQQgghRAwkmBJCCCGEEEKIGEgwJYQQQgghhBAxUFrr8AuVOgDsPHrFEUIcBT211tnHuhBNJfWTEO1Sm6+fpG4Sol0KWzc1GEwJIYQQQgghhAhNhvkJIYQQQgghRAwkmBJCCCGEEEKIGEgwJYQQQgghhBAxkGBKCCGEEEIIIWIgwZQQQgghhBBCxECCKSGEEEIIIYSIgQRTQgghhBBCCBEDCaaEEEIIIYQQIgYSTAkhhBBCCCFEDCSYEkIIIYQQQogYSDAlhBBCCCGEEDGQYEoIIYQQQgghYiDBlBBCCCGEEELEQIKpNk4pdYpSavfR3jbG481VSk0/WscTQrQebamuioVSKlcpVaGUshzrsgghwmtLdZFcN7UNEkzV4/8xrH34lFLVAX9PacHjXqOUmtdS+28qpVQvpZRWSlnrvf6yUuq+Y1UuIeKV1FXh+euqvBY+xg6l1Om1f2utd2mtU7TW3pY8rhCtjdRFocl1U/ywNr5KfNFap9Q+V0rtAKZrrb+qv55Syqq19hzNsgkhRC2pq4QQrYHURSLeSc9UhGq7dpVSf1BK7QdeCtUqEtgiqpRKUEo9rJTapZQqUErNUEolxnDsa5VS65VS5UqpbUqpG0Os8yelVJG/tXRKwOvNUoYIy3mNUmqe/3jFSqntSqkzw6zbRSm1Win1O//fc5VSf1dKzfef5xdKqayA9X+ulFqrlCrxrzvI//q1SqkPA9bbrJR6K+DvfKXUcP9zrZS6yb9OiVLq30op1RLvhRDHitRVQcf7q1LqTaXUq/5yrVVKjQ5Y/n9Kqa3+ZeuUUufX2/76gHNap5QaqZSaCeQCH/pb338f2AqtlLpUKbWs3n5+o5T6oCXPVYjWROqiiMop103tgART0ekMZAI9gRsiWP9BoD8wHMgDugH3xHDcQuAcIA24FnhUKTWyXrmy/Pu/GnhOKTUg2jIopZ5WSj0dQ/kCjQU2+svzEPBC/S+eUqo38C3wlNb6XwGLrsA8vxzADtzpX78/8AZwO5ANfIJ5EWP372eiUspQSnX1bzfOv10fIAVYHXCMc4ATgGHAJcDPmni+QrRGUlfV9XPgv0AG8AHwVMCyrcBEIB24F3hNKdXFf5yLgb8CV/nP6efAQa31lcAu4Fz/0L6H6h3vQ2CAUqpfwGtXALOiPVch2jipixon101tndZaHmEewA7gdP/zUwAX4AhYfg0wr942GvPLp4BKoG/AsnHA9jDHCtpXA+V6H7gtoFweIDlg+ZvA3Y2Vwb/t7giP2ct/btZ6r78M3BdwDlsCliX5t+ns/3su8Ij/fb283n7mAncF/P0r4DP/87uBNwOWGcAe4BT/3/nASOAy4DlgCTAQs4L5oN7/zYR679P/HevPmTzk0dSH1FVBx9VAnv/5X4GvApYNBqob2HYVcJ7/+ee15W/oPff/XaeOBF4D7vE/7weU++vEqN5vecijLT2kLqpzzDp1QsDrLyPXTe3qIfdMReeA1toZ4brZmF+K5QENDAqIOtOTv8v3L5gtJYZ/vz8GrFKsta4M+Hsn0LU5y4BZ8QDYAp7X/u0O+Ht/7ROtdZX/uCkBy6cAW4C3Qxxjf8DzqoDtumKeU+1+fUqpfMzWIjBbWU7BrIy/BUqASZgV4LcRHkOI9iSe66pQ6n/vHcp//4ZS6irgt5gXPmDWCbVDZXpg9lzFYhbw/4C/YbYev++vE3No2XMVojWJ57pIrpvihAzzi46u93cl5pcOAKVU54BlRUA1MERrneF/pOuAGzUjoZRKAN4BHgY6aa0zMLtrA7uAOyilkgP+zgX2NlcZ/PZhfvl71Xu9NwFf2Aj81V+uWSryFMJ7MYcIAODv/u6B2coCRyqFif7n32LozJh3AAAgAElEQVRWCpMIrhSEiAfxXFdFU+aewH+AW4CO/jKvCShzPtA3zOb13+P6vgSy/fceXM6RIX7H5FyFOEbiuS6S66Y4IcFU0/wADFFKDVdKOTA/8IDZCoD5I/2ovyUSpVQ3pVRDY02VUsoR+MAcy5oAHAA8/taWn4bY9l6llF0pNRFzfOtbMZYhJG2m+30HuF8p1VEpZVNKXY45ZObTKHblBi4GkoFXlVKRfAbfBM5WSp2mlLIBdwA1wAL/8m+ByUCi1no38D1wBtARWBlF2YRor+KmropSMubF3gH/Ma8FjgtY/jxwp1JqlDLl+QMwgAKgT7gda63dwFvAvzDvGfnS//qxOlchWoO4qYvkuil+SDDVBFrrTZhDOL4CNgP15zv4A2bX7CKlVJl/vQGENx6zRaT+41bML0Yx5nCRD+ptt9+/bC/wOnCT1npDtGVQZsaaGQ2U71fAIcwbEwsxW3PP1loXNLBNEK21C7gA6AS82FjFoLXeCEwFnsRsnTkX88Zvl3/5JqACszJAa10GbAPma5nzRYh4rKsiorVehzkUbyFmcDQUmB+w/C3gfsxepXLM+y4y/Yv/AdylzAxXd4Y5xCzgdMyLtMBhPtG+30K0C3FYF8l1UxxQWjc2UkEIIYQQQgghRH3SMyWEEEIIIYQQMZBgSgghhBBCCCFiIMGUEEIIIYQQQsRAgikhhBBCCCGEiIEEU02klHpZKXWf//lEpdTGo3RcrZTKa+Z9Hj6Xo7nt0aKU+pNS6vljXQ4hjhapn5q+7dEi9ZOIJ1I3NX3bo0XqpsbFRTCllNqhlKpWSlUopQr8H95mnyBRa/291rrR9LZKqWuUUvXTgTYbpdRcpdT0ltp/U7X0+fuPcYpSanfga1rrB7TWrfZ9EfFJ6qfWReonIUxSN7UuUje1XnERTPmd65/BeiQwGrir/gpKKetRL5UQQkj9JIRonaRuEqIR8RRMAaC13oM58/RxcLjL92al1GbMCeRQSp2jlFrln4xxgVJqWO32SqkRSqkVSqlypdRswBGwrE5Er5TqoZR6Vyl1QCl1UCn1lFJqEDADGOdv7Snxr5uglHpYKbXL3wI0QymVGLCv3yml9iml9iqlpsV6/kqpt5RS+5VSpUqp75RSQ+qtkqWU+tJ/ft8qpXoGbDvQv+yQUmqjUuqSWMtRr0w7lFJ3KqVW+8s1W5mzmKOU6qCU+sj/Hhb7n3cP2DZTKfWS/30pVkq9r5RKxvw/7up/jyuUUl2VUn9VSr3m3+5TpdQt9crxg1LqgpY8VyEaIvWT1E/+7aR+Eq2K1E1SN/m3k7ophLgLppRSPYCzgJUBL/8CGAsMVkqNAF4EbgQ6As8CH/i/sHbgfWAmkAm8BVwY5jgW4CNgJ9AL6Ab8V2u9HrgJWKi1TtFaZ/g3eRDoDwwH8vzr3+Pf1xnAncBPgH7A6U14Cz717yMHWIE583egKcDfgSxgVe1y/5fsS2CWf9vLgKeVUoPDnH+JUmpCFOW6BDgD6A0MA67xv24ALwE9gVzMmc2fCthuJpAEDPGX61GtdSVwJrDX/x6naK331jveG8DlAeUd7D/Gx9GeqxDNReonqZ/8pH4SrYrUTVI3+UndFIrWut0/gB1ABVCC+QV9Gkj0L9PAqQHrPgP8vd72G4FJwMnAXkAFLFsA3Od/fgqw2/98HHAAsIYozzXAvIC/FVAJ9A14bRyw3f/8ReDBgGX9/eXOC3O+c4HpEbwvGf79pPv/fhmz0qpdngJ4gR7ApcD39bZ/FvhLwLb3Rfj/Uf/8dwBTA/5+CJgRZtvhQLH/eRfAB3QIsd7h/4uA1/4KvOZ/nup/z3v6/74feNH/vMFzlYc8mvMh9VPY90XqJ6mf5HEMH1I3hX1fpG6SuqnOI57Guf5Ca/1VmGX5Ac97AlcrpX4d8Jod6Ir55dmj/Z8Qv51h9tkD2Km19kRQtmzMFoLlSqna1xRg8T/vCiyP4JgN8rf43A9c7D+mz78oCyj1Pz/8XmitK5RSh/zH7wmMre1a97Nitm40h/0Bz6v8x0QplQQ8itny0sG/PNV/Lj2AQ1rr4mgPprUuV0p9jNly8k/Mlpbr/Ytb+lyFqE/qJ6mfDpP6SbQiUjdJ3XSY1E2hxVMw1ZDAL3g+cL/W+v76KymlJgHdlFIqoFLIBbaG2Gc+kKuUsoaoFHS9v4swu2CHaHNccn37MD/8tXLDn0qDrgDOw+zq3gGkA8WYlU+tw8dRZtaeTMwWpXzgW631T2I8dqzuAAYAY7XW+5VSwzGHGSh/mTKVUhla65J629V/j0N5A/iLUuo7zPHb3/hfP1bnKkQoUj8dIfWT1E+i9ZC66Qipm+K4boq7e6Yi8B/gJqXUWGVKVkqdrZRKBRYCHuBWpZTNf8PdmDD7WYL5RX7Qvw+HUuok/7ICoLt/HDFaa5//uI8qpXIAlFLdlFI/86//JnCNUmqwv7XhLxGch9V/zNqHDbN7tgY4iNma80CI7c5SSk3wl+3vwCKtdT7mGOb+Sqkr/eduU0qdoMybQltSKmZlWaKUyiTg3LXW+zDHMT+tzJstbUqpk/2LC4COSqn0Bvb9CWZLyt+A2f7/Bzh25ypEY6R+kvpJ6ifRGkndJHVT3NZNEkzVo7Vehtll+RRmy8MW/Df0aa1dwAX+vw9hjg99N8x+vMC5mDdE7gJ2+9cHmAOsBfYrpYr8r/3Bf6xFSqky4CvMVgW01p8Cj/m32+L/tzHPYH6Rah8vAa9idnPvAdYBi0JsNwvzS3cIGAVM9ZehHPgpZtfuXsyu5X8CCaEOrswsMBMjKGdjHgMSMVugFgGf1Vt+JeAGNgCFwO3+8m7AbD3ZpswbOrvW37HWugbz/+90zPOufT2qcxXiaJH6SeonqZ9EayR1k9RN8Vw3qbpDWIUQQgghhBBCREJ6poQQQgghhBAiBhJMCSGEEEIIIUQMJJgSQgghhBBCiBhIMCWEEEIIIYQQMZBgSgghhBBCCCFi0OCkvVlZWbpXr15HqShCiKNh+fLlRVrr7GNdjqaS+kmI9qc91E9SNwnR/jRUNzUYTPXq1Ytly5a1TKmEEMeEUmrnsS5Dc5D6SYj2pz3UT1I3CdH+NFQ3yTA/IYQQQgghhIiBBFNCCCGEEEIIEQMJpoQQQgghhBAiBhJMCSGEEEIIIUQMJJgSQgghhBBCiBhIMCWEEEIIIYQQMZBgSgghhBBCCCFiIMGUEEIIIYQQQsSgwUl7hRDRyS/L5/s931NSUwJARkIGE7tPpEdqj2NcMiFEXKo6BBs+gooCcNdAYgbkngjdRoFSx7p0QrSIbQcq+G7TAYqrXBjKoEOyjVMH5tC9Q9KxLppohySYEqKJfNrH97u/58U1L7L24FoAXF4XAHaLnUeWP8KQjkOYdtw0JnafiKGkQ1gI0cL2rIAFT8LGT0AZ4HGC9oHFDoYVUrvASbfB0IvBLheYou3z+jRfrS9gxtytrNtXBoDL4wMgwWZw/8frGZnbgRsn9eHkftkYhjQmiOYhwZQQTVDlruK2b27jhwM/UO2pDlpe460BYEXhCjZ8t4Hjs4/n8VMfJ9GaeLSLKoSIB1rDV3+FJc+Cp8YMoAJ5Xebj0Fb47I/w3cNw7SeQIb3nou2qqPFw7UtLWLu3jCqXN2i5021+DxZuO8jCbQcxFCTZLGSlJjD1xJ5cPKoH6Um2o11s0U5IE7kQMXJ5XUz7fBorClaEDKTqq/JUsbxgOdd+du3hnishhGhWH98BS54Dd3VwIFWfuxLK9sBzk6B0z9EpnxDNrNrl5cJnFvBDfknIQCoUn4YKl5cdB6v4f19sZMwDX3H7f1dRWu1u4dKK9kiCKSFidM/8e9hSsgWXL/LAyOVzsaVkC/fMv6cFSyaEiEvLX4Yf3gB3VeTbaC9Ul8Ar54AvsgtRIVqT22avZEdRJS6vjmn7arePGo+Pj3/cy9lPfM++0sYbR4UIJMP8hIhBQWUBX+78MmQgVfx9MUWfF+EqdGFxWEgblUanizphSbYA5tC/L3d+yW+rfktOUs7RLroQoj3y+WDO/UGBVK/Hyqlyw/bbUki2m/eIPL/CxWur3cy9JtlcSXuhohA2fwEDzjzaJRftjNPj5MeiHymtKUWhSE9IZ1j2MOwWe7Mfa9fBKr7deIAaT3AvbOW6uZQtfR/3wd0Y9kRsOX1IH38Jju5DQu7L7dXsK3VyyYyFfHzbRNIcMuxPREaCKSFiMHvj7JCvF31axIFPD9B9endSBqfgLnazd+Zedjy8g95/7o1hPdIZ/ObGN7llxC1Hq8hCiPZs65ywPVJeDY8vdvGniQnht3dVwLzHJJgSMdtVtovX17/Oe1veO5xoSaHQmD1GF/W7iCsGXUHXlK7NdsyXF2zHp4N7pMqWvEfp4rfp+NObcfQeibJYqd6+nOrNi/GWHQgbZHl9mv1lTh74eD0PXjis2cop2jcZ5idElLw+L//d+N+gXilvtZfC9wvpOrUrqcNSUVaFPdtOj1/1wFXkonRB6eF1XT4Xb2x4A68MqxFCNIcFT5gBUQi/G2/n4QU1lDgbGQa1bxUc2t4ChRPtmU/7+Mfif3DBBxfw5sY3qfZUU+mupNJdSYW74vDzWRtmce775/L48sfRIQKgaLk8PmYvzcddb3ifr6aSknmvk/mTX5I0YDyG3YGyWEnKG4sluQOHvv4P6SdeQvdbXqPbL18ideRZVG9efHh7t1fz/qo9VNZ4mlxGER8kmBIiSsU1xSETSFRtrsLn9pE2Kq3O6xaHhdRhqVSsrXuh4/Q6KXWVIoQQTVawNuyi0V0tnNLLysMLahreh8UOheubuWCiPfNpH3d+eyfvbnmXGm8NHh0+AHH73Li8Ll5f/zr3zL+nyQFVQZmTUHuo2bMB7XGR1H9c3bI2EGR1mDytzrqGUry/UpKyiMhIMCVElCpcFViUJeh1b4UXa4oVZQmeu8KabsVTUfdHxqqsVIRpSRZCiKi4Kxtc/LfJCTy5xMWBygYy/Gkv1JQ1c8FEe/bEiieYt2ceTo8z4m2qvdV8tvMznv/x+SYdu9zpwQgx8bS3ugwjKQ1l1P2dDhdkhVLl8vLa4p1NKp+IH3LPlBBRclgd+EKkHLakWPBUeNBeHRRQeUo9WFPqft282ovD6mjRsgoh4oQlwUyHHsZxORbO6W/lwXkuBmWHaUdVBthkDjwRmRJnCTPXzYwpEZPT4+S51c8xZdAUkmyxTRqdaLeE7N2yJKbhqypD+7x1AqpwQVY4BytkChMRGemZEiJKHRwdQgZTSXlJKKuibHndll2v00v56nKSByfXeV1rTUZCRouWVQgRJ9K6NLrKvac4+M8KF3vKwgyv0j5I797MBRPt1btb3kWF6Bkq+rSI/W/tp/MlnRn89GD63N0H10EXOx7egS8g655Sio+3fRzz8bNTE4LulwJI6DYQZbVRtWlhndcDg6xIuL2NzNMmhJ8EU0JEKcGSwGk9TzucraiWJclCzi9y2PvaXspXl6M9GtcBF/lP52PLtJEx/kjgZCiD03ue3iKpYoUQcWjMTWBLbnCVvEyDS4fYeGJJmBb3xEzoOrIFCifaG5/28eraV6nx1r0PL5pETNWeal5c82LM906lJFiZNCCb+vGckZBMxoQpHPpyBlWbFuJzO9FeDz5XNSgVFGQ1tH8hIiGfFCFicPWQq/lm1zc4vXXHiWeflY0l2cL+2ftxFbowEg3SRqbR48YeGLYjwZfdsHP1kKuPdrGFEO3VsIvh8z82uto9kxKYudodvMCWBCfdRtCVqRD1ON1e/rt8LQery6DexyWSREwdTu5w+PX9lfup8lSR3EhDQDg3nNyH+VuKqHLV7W1KG3MBRnIHShfOpuijh1H2RBI65ZE6/AwOfTkDZVhw9B6BMqw4d6zCuWt1nSQUFgXj+naMqUwi/kgwJUQMhnQcQs+0nmwu2Rw05C9zUiaZkzLDbmsog55pPRnccXBLF1MIES/syTBiKqx4FQKSAey4PbXOaj3SDZx3pdXf2rxfatilLV1K0Yb5fJpHv9rEC/O2o2wFqG4G9XMxNZaIqXpn3fv6rBYrZTVlMQdTo3t2oHOag+0HK6nfwZUyZDIpQyYHbWPvOjAoyEobV/ezb7MaXDehT0xlEvFHgikhYvTEqU9w8YcXU+4qPzwpYWMUihRbCk+e+mQLl04IEXd+8nfYuRAObIQQSQHCsiXCZbPAESLIEgJzTqcbZi5j8baDVLt9KA3JIX73ok3EpLXGZrHFXC6lFE9NGcnZj38f8TbhgqxAfbNTGNA5tcF1hKgl90wJEaOuKV159cxXybCnYYlgzLdVWeng6MDMM2fSJaXxm8WFECIqNgdc8yF0GgwRZArVYAZSF70EfSa1ePFE26S15jezV7HIH0gBaG8KqOBEDtEmYvL6vKTZmxbEL9t+CLs19OVs5bq57HvldnY9chG7n7qSgjf/gnN3+DnZABw2gz+fPahJZRLxRXqmhGiCvul9eNeVxtNpnfmwaidKKao9dYcxJFoT0Vpzbt9z+dXwX5GVmHWMSiuEaPcSO8B1X8D8x2HRM/gqyzCMehOpWh1on4/KwgSS7nwPo9cJx6asok34en0h32wsxOkOGNLuc+CtzqVi9UozBfo+F4bDwJHrIOOkDPa+thfDYZAyOAV3sZu9M/cGJWJSKE7qdlKTEjFprXn2u23UeIIz75UteY/SxW/T8ac34+g9EmWxUr19OdWbF+PoPiTk/hw2g3t/PoTxfeV3WkROgikhmmLFq2Q5K7ln6gfc6XPx8baP+WT7J5TUlACQkZDBWb3P4uw+Z8c8l4YQQkTFmgCTfo/3+Onsu2I83S4bjKoqBI8LEtOhz2TU6OsouedBnJ+vJOtGCaZEeM98uzUowQNA4YdQOn8/Xa/uQurQVJRFUf5jOVUbq+h0YadGEzElWhO5Zsg1TSrb0h3FFFcFD2n11VRSMu91Op51O0kDxh9+PSlvLEl5Y4PWtyjzPqlHLjmes4Z2bVKZRPyRYEqIWJXtha/vhas+AIuVJIuViwdczMUDLj7WJRNCCCrmLUT3OAV11YyQy7NvvZWdV0yhw6WXYMmQOe9EsB1FlazZUxr0uq+mkuKvv6Tbdb1JH51w+PW0EWmkjTCH7TWUiAnMxsZRnUY1qXw/7inFE2I+qJo9G9AeF0n9x0W0n5w0B5/cOpEOyTJdiYie3DMlRCy0ho9+CydMh87HHevSCCFEkIo5X5Ny2qlhlyf07k3q6adz8IUXjmKpRFvyyZp9+ELcE1wbrNiyb0H7ok8g4bA6eHTyoyEn/Y1GudONK8TEvd7qMoykNJRhCbFViPLYLBJIiZhJMCVELNa8A8U7YOKdx7okQggRRLtcVMybT+rkhrOWZd38K0refAt3QeFRKploS/aWVONuIFjxufpSveeKqAKqRGsiT576ZLNMD5JgtWAJcSVrSUzDV1WG9gUPTwy9H7kcFrGTT48Q0ao4AJ/9H5z3b7BKS5YQovWpXLKUhN69sWZnN7ierXNn0i+8kKJnnj5KJRNtSZ2kEwECgxVvxSCqdt6At7or2mdD6xC9TdogwZLA0KyhzDxzJid2ObFZytcl3YHDGtz7lNBtIMpqo2rTwoj2071DYrOUR8QnCaaEiNanvzcnt+zetLHeQgjRUswhfqdFtG7H66dT/tnnuHbubOFSibbC69M88sUm/rdqT8jl9YMVn7MHVTtupWrHzbhLR+HzJKN9VrTPis+TTCc1ibfOfYtZZ89iQOaAZivn6YM74Q0xDNFISCZjwhQOfTmDqk0L8bmdaK+H6q3LKP7mxTrrJtstTD2xZ7OVScQfSUAhRDQ2fAz7Vpm9UkII0QpprSmf8w25Lzwf0frWDh3IvPoqDjzxJN3+38MtXDrR2jndXqa/sozlOw+FHOIHdYMVZVhw9B6BMqxUrtuNc1cZHSbffXhdu9Vg+s+H0Ds9t9nLmpJg5bzh3Xh7eT7181CkjbkAI7kDpQtnU/TRwyh7Igmd8kgbd2md9ZLsVk7u13APrhANkWBKiEhVl8DHd8KFz4Nd0pwLIVon59p1GAkJ2Pv0iXibzKuuYssZZ+Bcvx7HIJmwNF75fJqbZ61g2Y5DOEPM3RSoNlgpnvM8nrJCMzGTxYo9py/O3WsPz+Xk8vg4bVCnFivz9Am9+d/KPXh9weVNGTKZlCHh7xt02AyuP7k3htG0RBgivkkwJUSkvvgzDDgTep10rEsihBBhVcyZQ8ppp0WVKc1ITibrhhspfOwxcp99tgVLJ1qz//2whwVbD4YMpCrXzaVs6fu4D+7GsCdiy+mDLaMzPreT7PP+GHZiXIsBn63Zx5XjerVImft1SuV3Zwzg4c83Ue2OLOEEmEknRvTowLSTerdIuUT8kHumhPCr8XgpLHeyv9RJlctTd+HWObDtWzj9r8eiaEIIEbHyOXNIbSAlejgZl16Ca8tWqpYta4FSibbg6blbqQ4xQW/Zkvc49PV/SD/xErrf8hrdfvkSKUNPo/yHz8j8yS9JGjAew+5AWawk5Y2lw+Rph7f1+uDZb7ehQ9zb1Fyum9CHmyf3xWGL7LLWYTMYmZvBC9eMxhoqHaAQUZCeKRHXfD7NvC1FPPvdVhZtO4TNUKDA7dX0yUrmpkl9OXtAGo4Pb4NzHgNH2rEushBChOXavQdPYSGJw4dHva1ht5P161sofORRer7+GgqgaBNUHgCvGxIzIHsg2CTzWXu0Zk8puw9VB73uq6mkZN7rdDzrdpIGjD/8upGQDFpHNDHuwUoXWwor6NcptVnLDFBa5eat5fl8sa6AlAQrHq8br08TKnRLTrBgsxhMn9Cbmyb1lUBKNAsJpkTcWrL9ELfMWkFljYdKf0uc13ek+t1cWME9/1vD3Z4a/pR7BVP7nR7zscqdbuZvOcihShderclItDG2TyY5qY4mn4cQQtSqmDOHlFNOQVkim6y0vvRzz6X4pWepmfUHHIWfQlURGLWXChq0D0ZMhbG/hEwZHtWevLtiNzWe4F6p2gl66wdN0UyMa7Uoiipc9GvGW6f2lVbz4Kcb+GzNfgylQg7xU5gJMLKS7fTKSubq8b04dWCOBFGiWUkwJeLSZ2v2cfvsVWHn0KhlBllW7t8znPxP1vPHs6K7MXvD/jJe+H47H67ei9VQZmuZBouhcPs0E/KyuPHkPozpndnkmeCFEKJ8zhwyr5wa8/Zq/Qf0GrUcvWERWMLUj0tfhOWvwHEXwrlPgEUuJdqDPSXV+EJ054QLmgLnmookoHLXT7fXBOv2lnHFfxZR7nQTJuEgANp/3AqXlz+eNYjjuqU3WxmEqCU1oIg7y3YciiiQClTt9vHqwp10SnMwbULjrbE+n+a+j9cxa/Eu3D5dp8cr0DcbClm07SBjemcyY+ooHLbYWpOFEMJbWopzzRqSx49vfOVQlj4Pn9+F8tWgGqqKfG7zsfZdKN0NU98Biy22Y4pWoyaCCXoDg6bAuaaSB05ocN9aQ3pi83xGdh2s4rLnFlLm9DS+MuDTUFrt5vLnFvHBryfQOyu5WcrRVmmtWbGrhE0F5VQ4PSTaLXTvkMiEvCzpsYuRBFMirmitueOtH8IGUqGyFaWPvwRH9yFUu73887MNXDCyGxlJ9gaP8ds3V/H52oJGU8tqoMrlZdHWg1zy7ELevHGcBFRCiJhUfPcdSWPGYCTGcE/Tpi/g87vAE3zPTFjuati9BD74NZw/I/pjilalQ3Lo37VwQVO4uaacO1bh3LW6ThIKn9b0b6b7pa5/dRkVNcGBVEO/3wAVLg/XvbyUr++YFJcjQSpqPLy/cjczvt3GoUoXWoPH58NiKCyGwmYYXDO+F1ecmCu3IERJgikRV1bsKuFAeU3IZWVL3qN08dt0/OnNYVO8KgWzl+Zz46S+YY/x9NytfL62IKoUrU6Pj037y7njzR/495SR0Z2UEEIA5V/HlsUPreHj3wYFUr0eK6fKDdtvSyHZbl58Pr/CxWur3cy9xt+6766Gte/BxDsgq19TT0EcQxPysvhi7f7D9xDXaiho8laW0OHU6xqcGNdqKC4a1Z1Ee9MbCn/IL2HXoaqg4YiR/H5rDftKnazYVcKonh2aXJa2ZMP+Mi5/bhE1bh9V9a5NjkzM7GXGt1t59rttPHn5CE4f3HJzg7U3EkyJuPKf77aFDHLCZStKyhtLUt7Yw3873T6en7ed6yf2CTnJn9Pt5d/fbAkbSDXUcub0+PhqfQE7D1bSs2N8D0MQQkTH53JROX8+ne/6c/Qb71wA1YdCLvJqeHyxiz9NTGjg4F5Y9Ayc80j0xxatxtnDunD3/9aEXFY7QW+ooMnRfVCDE+N6fBqXx4fWusk9Qs9/vy0oSUakv98ATo+X/3y3lVFXjm5SOdqStXtLuXjGQqpCpLyvr3Y0zS1vrODhi4/nnGFdW7p47YIEUyKuzN1YSKipLsJlKwqlssbDlgMVIYcsfPDD3rDbRdJy5tOal+bv4K8/HxL5SQkh4l7V4sUk5OVhzcqKfuMFT4KrKuSi342389D8Gn51gp0MR5gLYZ8bfngDfnof2JOiP75oFRw2C5ee0IPXFu0M6K04ImXI5AaDpob8b9VelFI8cP5xMQdUNR4vn68tCOqViub3W2v4ekMhTrc3LobUF1e6mPL84ogCqUBOt4/fvfUDvTomS9KOCEgwJeKGx+ujJkw2oWhSvFoMRXGlK+SyGXO3hqy0Im05c3s1by7L5//OHBgXFb0QIko+L2z+Ala8CmV7weuCxAx8OyHtlJ/Fts8d30PIWXlgdFcLp/Sy8vCCGu47tYH7KAwL7F8NuSfGVgbRKtx4cl/eWb4btzey5A6RqnZ7eX/lHgZ2TuXq8b1i2kdJlRvDAOr9xEbz+w1gNQwOVrroltH+50ubtWQXzjCBVGP3mDndPh79ahMvXH3C0SxymyTBlIgbDc29Hk2K1wqnh9+/vZpuHRLJSLKRkWQnI9FGeqKVbUWVITzkSkUAACAASURBVLeJpuVMKcg/VNUikxuKJirbCwVrwVlqTlya1g26HG/+pwnRktxOWPgULPo3eFzgqqizOMWrUAXLYPYymHwX5AyMYt+he6Vq/W1yAie9WMltY8Mn3gEF1SWRH1O0Sp3THcy8biyX/2dR1L0ZjV2cV7u9PPrlJqaMzY0pa5zT7cUIUddGm6LdUFAd5bm1RV6f5sV520MmworoHjNg3uYiCsudkpCiERJMibhhsxjYLAauEBVLNCleHTYLv/vZADKS7JRUuyipclNa7WZviROlCDmMMJqWM0MpypzuiM9LtDCfD7Z/C/MfN+8tsSaYE5cqZf6b2BFOug2OvxQSJAAWLaDqEMw8Hw5sAI8z5CqGRYPPBRs+hi1fw6WvQd5pke1fGeZnOYzjciyc09/Kg/NcDMpu4CJY5ptqF47vkcHbN41n6guLKa12EWpAR/3ASSUk460sJuvM28JenIM559PXGwr52ZDOUZcr1WHDE2KakWh+v8G8hystsf1/Vj9fuz/ktUQ095gBzFq8i9tP79+iZW3r2v+nSYgAE/pm8c3GwqBeqmhSvCbaLZw5tAuWegkonG4vMxftDNkDFk3LmdaQYJUhfq1CxQGYeR4U7wCXv9fRWy8bpKsSvrwbvrrHvIDtG0M2NSHCcVfDK+dC0SZzSF9jtM/saZo9Baa+Bz0b7w0nIS1sAopa957iYOSzFdwxLkwiCu2FpBju1xKt0uCuaXz06wmc/NA31B/XUb9XQ3tq2P30NST2HtHoxXmly8uMuVtjCqYyEm2kJFg55Kn7PYjm9xvMBtGOyQ0kVGnjdhdX8fhXm3lv5Z6QwWc0I2VqPD6W7mi4bhAgs3OJuHLDpD5h07OmjbngcIrX3U9OYfcz11C+4iMS+x2pcBKsBtNO6hUUSIFZQYcLggJbzhrj9vrISW2/FX2bUV4Az06AA5uOBFLhuKvMdd64AtZ/dHTKJ+LD53+Cg1vqBFK9Hisn51/lVLqOXCg9v8LFKS8HfE7d1TDrksY/uwDHXQhGwxOq5mUaXDrExhNLwgR0tmToPKzxY4k2492Vu7HW+62r7dXI/MkvSRowHsPuwLV/C/i8ZJ8fWSbJjQXlMZXHMBTTTupFgjX40jWS328wf8OvHt8z5G94e/BDfglnPv49767YHTKQgujvMSutlpEyjZGeKRFXxvbOJCPJFnYseCTZii4bkxt22fkjuvHmsvygSiyalrNBXdLISZPxyceUx2X2BlQeNDOVRbxdNbw7Ha79DLoOb7nyifhQUw6r3gg5tC/ilOVr3oGRVzV8nBN/CStnNvpZv2dSAjNXh1jHmgjjbsbMDiDaA601L83bEXS/TahejWgvzp1RzMFY32VjcnlyzpaQyyL5/dbA1LE9Yz5+a7Zxf3lE97pFe4+ZQ0bKNEpqPhFXlFL866Ljcdii/+gn2gxunpxHVkr4i5dpE3oHteTViqTlLDnBwk0NTAgsjpL1H0DZ7joXlxH1BoDZI/DVX45WSUV7tvrNsMlNfjfezsMLaihxNpBax10J8x4NfSNnoI59zUQq1D3WjttTOb3PkTbXHukGzrvSjkzYe5iv8YBNtClOt4+SED0SoQKnwIvzSMSSfKJWVkoCF47sTmIMv+EOm8HPj+/aLhsrvT7NVS8uDplYo3LdXPa9cju7HrmI3U9dSdmS98CwRDRSRinokSnTHTRGeqZE3DkpL4sLRnRj1pL8iLdJtFk4f2Q3fn1qXoPr5eWkMKhrGqvzS/GGuIBprOXMbjE4fVBOxOUSLWT+YyGHR0XUGwCwcyGU5ENGjxYqoIgLi2eEzbQXccry8v1QsAY6D234WL94Bp6bZPaGRcOWCGf8E5Iyo9tOtGrlTjc2Q+GtN8oiVK9GtAkgOiY3lBWycfeeN4RtBypYtbsEpzt84pRADqvBkK7pPHB+I9+DNmruxkIqnJ6ge7bDZe3TPk9EI2UcNgtTxoYfjSNM0jMl4s53mw7wxboC7j//ONITbSQnhO/CTrRbSLAa/Pq0PO7/RWSTDT51xUhSHNG3UzhsBi9cc0KTWu1EMyhYB0Whh5FE1BsAgIYl/2n+son4UhZ+EnAwU5Y/ucTFgcoGLigNqxnYN6ZjX7jyfX9GygjvJ7ElwsTfwairI1tftBlJCdZGM+fVChzGXrVpIT63E+31UL11GcXfvFhne4fVYOqJTRtmZ7MYvHLdGCb1zybJbmn005pkt3BSvyxenz4We4j7rdqDGd9upbJer1So+9uUxUpS3lg6X/EgSQPGU/Thw+Q/chG7/t/5FH38CJaMTnX2YVEwqmeHo3kqbZL0TIm4smZPKb+ZvYoZV47ihF6ZXDK6B1+tK+CZb7eyfl8ZdouBUubkuR1T7NwwsQ8XjupOqqPhm7MDdctI5K2bxnHZs4vMtLKNXXdjVvYzpo5iZK5UWsdc/qKw15IR9wZ4XbDtm5Ypn4gfjWTviyhleW12v0h0Hw3Xz4V3r4fCdeDzmI/67ClmIHXmQ3DcBZHtW7QpyXYLdquBp94Ferj7f20dc7F3zqN04WyKPnoYZU8koVMeaeMurbO9Bi5v4L7jSCVYzd/MxdsP8czcrXy76QCJNgsen9mwYLMYeH2aE/t05MaT+zCub8eIGkPbooIyJ6t3lwa93lDWvrIl71G1cT5Z59xRp8eqJn9tnfWqXV7KnB7SEyO/BopHEkyJuLHzYCXTXl7K/ecP5YRe5pAUm8XgzKFdOHNoF0qqXByqdOH1adITbWSnJsRc+fbvlMont03k7x+t46v1BShF0HAEq6Hw+DTj+nTknnMHM6hLWpPPUTQDZ6mZgCKMyCYw9e/nGPH6NIu3H2RviZNqt5c0h5VBXdLoLxNBty22xEYDqkZTlivDTH0eqaw8uOEbOLARFj0Da95G11SA1ihrAnQ/ASbcDn1Pk4QT7ZhSiktG9+D1RTtx1+uhShtzAUZyh6DAKX3cpTi6Dwq7T5tFcfqgTmQ2cZif1poql5cql5fhPTK4aVJfCsqcTJvQm9Iq8z6v9CQbJ/fLpnN6+7s/qr49JdXYrQY19ZKFhEsMEs08UzaLwTvLdzNtQu+WO4F2QIIpEReKKmq4+sUl/Pq0fpxxXOj5LTKS7GQkNa2SD9Q53cG/p4ykuNLF7GX5vLN8NyXVbnz+CQN/Orgz3206wHUTeksg1ZoYNjAs4A19M3XkE5ge/Za8gxU1vLF0Fy/N20GNx4tPg09rLErh1Zo+WSncdEpfzhjSud0Od2lXcgbDroZvEg9MWT40J8T/qdcFOQOjP3b2ADj3MTj3MWrWrmXvXX+mz3vvR78f0WZde1Iv3liyC0IM94skc159bq9mS2EFH/ywN6Y6KP9QFa8s2MEbS3fhdPvMBkmvxmZRDM/N4NSBOQ0miGqvQiWdgPBZ+6KZZ8rp8fH8vG0STDVCginR7lXWeLju5aWcM6wrVzZxrHYsOiTbuWlS35BZ+vp3SuWVhTs4fXCn4A3FsZHSCSz2BnsEGu0NqN3PUfTlugJufWMlWuugdMa11u0r44/vrObBT9cz+4ZxkqWptRt/K+z/EVwVDa4WNmU5mEP3Mpo4rMpqgzCfKdF+9eyYzKieHVi64xDuSMarR2BjQTl/fHc1f35/OVedXsL8A++wt3IvLq8Lu8VO1+SuXD3kas7ofQaJ1kQASqpc3PrGKhZvP4hP68NlqU2O4fVoVuws5qQH53DOsC48cMHQuJr4PiXBWn9eZSB8YpBoU9nvK3WitW63wySbgwRT7dz+yv28t/k9tpVuo9JdSZo9jcEdB3Ne3nmkJ6Qf6+K1OLfXx82zVjCgcyp3/LT/sS5OkLOHdeGBT9az7UAFfbJTjnVxBED/n4a+TyRAo70B9hQYdW2jh6rttVyw9SBl1W7sVoPczCQuO6EHo3p2iPjH68Mf9vC7t1dHlNmq0uWl2u3lnCfn8eEtE8jtKAFVq9X/Z2ZgX8+O2+sO16xNWR7EngLjb2tyMZTVgg7TUyvat6euGMmZj39HUbkrZIba6Hlxp32CPXMhr2xSKONIo1W1p5qtpVt5cMmD/GPJP7h0wKVcnncjFz2zmAMVNQ0GdC6vBjQfr97HpoIK/nvDiSQnxMclbq+Oybi8wXV/uPvbPCUF+CpLIp5nylCKareXJHt8vJ+xkHemnVq6fynPrX6OFQUr0GjcAfPlfL3ra55Y+QSTe0xm+tDpDMgccAxL2nK01vzx3R9RwP3nD22VrSoOm4WLR/fgtUW7uOfcwce6OALAkQ5Dzjfn+NHhLyAb6g3wuT3Q92dh06VuLijn8a838+W64Pvplu44xCc/7qNjsp1fntKXS0/IxRJm7jKAVfklEQdSh8unzdTHlzy3kDl3TJIfydbKsMCkP8DX90aeROLwtlZI7Qx5pzW5GMpiQXsbbmAQ7VNmsp33fnUSlzy7kMLyGlxN6aFUbhJ7vIQlMR9lhJ8guspjftb/u+G/vLZyHmUV1+D1RlZHOT0+NhWUM/2VZbw2fWyDdWd7kZ5k4/RBnfh0zb6gEZmh7m+zZ/cGS+Sp7H1ay8S9jZBf0HZGa82zq5/lhR9fwOl1hlyn9vUvdn7B3Py53Dv+Xs7qc9bRLOZR8fAXG9lcWMEb14/F1orTjU8Zm8u5T83jzp/1l4va1mLcLbD2ffBUH34p0t4AbdipKOtN4c8vIPv220k7+yxUwI3632wo5Fevrzh8T1PQ9hr/zdXV/P2j9XyxroAZU0fhsIX+MXvosw1hA6nKdXMpW/o+7oO7MeyJ2HL6kD7+Ehzdh+DTUFbt5v2Ve7hi7NEf/ioiNPZG2L8a1r4XeUClLGajwNUfmgFZU1mt4JGeqXjVNSORT26byIy5W3lt0U68Ph2UhjucOnVQoiYx1072z7NI7l9/4udgTq8Tbd2BveurVOdfQ+BsPg3VbTUeHz/sLuGTH/dx7vFdYzvpNub6k/vw1fqCoCQUEPr+trIl70Y0zxRAVnICRhwEpU0hV27tzNOrnubltS+HDaQC+bQPp9fJXxb8BUMZnNH7jKNQwqNj5sIdfPLjft6+aVyrD1B6ZCYxumcm76/cyxUyOV7r0Pk4GH8LLPx3dD0ChhXVsTdp//c1llVrKXzoXxx65RVyfv87kseMYf6WIn75+vKIe5Gq3V4Wbj3I9FeW8cq0MUGtrHtKqlm2szjktuEma6zevBhH9yGAGbQ9++02Lh+T2yp7bgWgFPz8KXP+pxWvgruakDdI1LIlQXI2XPMxpDXPhaTZMyXBVDxLc9j4/RkD+c1P+vPF2gLeX7mHoooafFpTWF7D/jIn9UcBBtZBKcdbSOz2IRVrD1G+ovxwMFX8fTFFnxfhKnRhcVhIG5VGp4s6YUk2GwGU4cGStB1r2g94ykYE7behum3Gt1vjIpjaUljBvz7fEFWvYbiMjPVT2TusBlePl8a2xrTuq0wRlfl75ocMpBqrrJxeJ3fPv5uBmQPpld7rGJS8eX22Zh9PfbOFt24cT8c2ktnnqnE9eeCT9Vw+podc1LYWk/9spjdf+VpkAZUlATJ6wNUfgz2Z5DFj6PXmbMo++ZR9//dHXAOGcEPWGTjdwRfCjbWyLt95iKe/2cKvT+tXZ7uZC3cSdAVDdKlvD1TUsGJXiUzM2JoZBpz5TxhyASx4ArZ8afY++QMrnxeUIwmV0gkm/AaGXgz2ZrwXTob5CT+bxeDsYV04e1gXwEzwNPq+r4Kqofp1kKPLoxh2D2kj0kgbYfboF31axIFPD9B9endSBqfgLnazd+Zedjy8g//P3nnHR1Gnf/w9M9uy6YEUAgmk0HuXJsXesTcQKTbuLHfqeXfe6RXP83fnnWdv2PFU7HoqRQUE6UWQFiAECIT0tsn2mfn9sQSS7GyyG7IpMO/Xy5cvZ2dnvrtuvvP9fJ/n+TwZD2cgHnf7E0QPpi4r8FYPD2luyy2pYW+R7bRuCbEhr5zZb2zA7pab2mLRJBhHRgW4oRX6gp3udNzcJ52QeXHbi35CqvSbUgo/LCTluhQGvDCAzD9m4i5zc/DJgyj1djG8ipe3d73d1kNudTbklfPwpzt4bdboTlVYPzG7Ky6vEjDKoNMOCAJc/E+48AmITEIVA/QrMVrBYIHB18DtKyAq8eQlRJHYSy8h85uvWZo1Hq/T5ff26g2fUv7dq8SedR09frmQ7ne9QfSIi3HsW3/iHIdHYcHqPDyNiozX5pYeL7xuSCjWt15Z5af8ymbP0+kApI+FG96F+3bCBX/z1VNN+BXl+el4zn0B7tkKI2e1rpACBD3NTycAP+VXYtBIAas/B4mWo4im8gavyw6Z4s+KSZ2RSvSQaASDgCnRRNr8NNylbqrWNOzTJ5rKEc1HQ5rbPLLC0p2Fp/YBOzC7Cqq59Y0N1LZASAWD2SBywcDkM9JuPlT0yNRpwqHqQ+wp39PgWN1k1X1ud6KH+HZm6iarvQ/upWpNFfFn+3ajvaqXL3O/5IFRD2A1dh4RUp+9RTbmv7uZ/9wwjEHdO5dToSgKzDyrJ2+vPcSQHrFsy6+iwu5zOYq3mhjSIzZgzYxOmBk5C4bPpPS+K4nvWYzBU+CLCEhGiEqCMXfAsBt9NSqBMJr4r6srLqmhmApll9WrKHy3u4gLB3U7cazaqR0tCMX61i0rVDsCF4PrdECiEmHUyboGx8dFmG3RmMIU1dbT/HQCUWn3oGos5evPQYaYn0BoOFfZ99lRPAoxIxvWnUoWiegh0dTsrDmxPgFA8GKI/QnZYQ16bpMVKLb5b2CdDqiqyry3N2LXqF1rKtMhWIySQI/4CP7v6iGtOezTFl1MnSYsylmErDT8owp1shIFkaWHljI9e3qbjLk1Kah0cOvrG/jDJQOY1Dux+Td0QMZldeGJb3bz3e4iREFAwFcZIeBz07lhTDq3ju+l9wZqB+TaWsp/PErC4ysgKnQL+02HKqh1+QufUHZZa10yry7bxYSSPXiOFuA5ehSxLB1E/xSWQM0atZBEMBv1JIXOjCkrE1fuAaLPDdMNJAN49TQ/HX98QSl/EV9/DhINlQhCQ8El18gYogwIkv97DbEGHIccDY4JgopoqESKSAl6bgNQtFx+TgPW5pZRafffBAumnqw5LEaRrMQoFs4d2+FrzjsK+rd0mrC3Yi9eteHDLtTJyu61c7DqYDiHGRaq7B5ufWMDt07oxfTh3dt7OCGjqir/XJLDa6vz8Coq7gA7wG+vPcjCdYeYOa4nv7+ov+6u04bUrFiBdcwYpBYIKYD8crtmGkaozRMP5ZdSsfs7jN27Y0xNpbsQTW65/3mBmjVqYTZIJEcHSGHU6RSYM7Owb1jf/IktRO8zpROIhEj/PmjQcA6KSPP/7UhREt4aL6qs+q1RvFVeDFEay1PRG9LcJgqQGH16pqi9/EOuX1SquUyHYCJWSdFmbj87kxln9dSzYUJAF1OnCbWeWr9jLZmsqtxVfsc6Mk6PzG1vb2JidiK3Tcps7+GEjKqq/Oaj7fxv+zFNS9P6eI43JXx33WGKq108fcMw3ayijahesoToC85v8fvtbq/mDmkoESQAb1w86f9+5cR/z9pdxOb3tvrZFAdq1qhlfauoKucPTG7xZ9Npf0yZGVS8917Yrq+n+ekEYnh6PFqPofpzkKnLCOLGKgiSQM2uGmp315J4WSKCQaB6czWxY06mSMtOGdt2G8nX+M9JqjcqpLnNbJCY3DcpLJ+7Pal1eVmTW+Z3vKlMh2AiViZJ5M6zM5nTCddS7Y0upk4TIo3+PRus2daQJ6s4U1xYx9mayIrKfe//RFKMmT9c0r9TCovnl+/nf9uP4fAEv1BxeGSW7Sri38v2cv/5p2fD5Y6EUluLfd16Uv/2txZfI9pi1GweGcouK4DV1FBwTembhNkoafZ8Ccb6VhIFLh2SSrTF2IJPpdNRMGdl4T5wAFVVwzMPShJ4veG7vk6nxWQQmTG2J6+tPuBnhlM3B5Utfp9jC48hWUQsvSwkXZaEZJVImp5EwcICRIvYwM3PmGAkbnzDtYgqm5Breze4bnO23imxFob26Fz108FQXuvGKIl4Gm1wBMp0CLY21y0rvLI6j9kTM/S/8xDRxdRpQt/4vmwu2oxHOZlDG+pkZTVYO4Q1+u5j1fywt4TyWjeiINAlysR5A5Lp2eWkYFRVlT9/uZMqh4c354zulClvNS4vzy3fr9lzqLlwvMMj88oPB5g3KZPYCH0hHE5qfviBiOHDkWJb/lDOTorSbNAbyi4rQN/kRvWPosDciRk8+/0+zd9Rc9a3Rsn3fp3OjRQTgxBpxVtUhDElpdWvL4iiz55dUXzCSkenHjPH9eT1H/PQ6n/mm4MmEdXnMQSpodtw4sWJSJEShR8U4i52I0aIxIyIIe2ONES/Ok4Rr21go+sGntsijBJ3Tc46LUWByytrVKkFznQIpTa3yu7hcLm9wXpLp3l0MXWacH3f63k/532/46FMVoqqcH6vlqcynQoeWeHrn4/x4opcDpbVIivq8bQ2387XP5fkMLh7LHdOzmJavyReXJnLhrxyFt05DrOhcz7cP9t6BFFjog+2gFQUBD7alM9cPSQfVqqXLCXmFFL8AAZ1jyU1zkJuiX86brC7rFaTxG1n+wuf2yZl8t3uInYcrdK0SQ9EhFHil9Oy6d8tpvmTdTo85swsXLm5YRFTcDLVT9DFlE4jUuMi+OXUbF5YkRsgy8KAu+IsTAmrEcSGtd0JkxNImJzQ5PVVxfd+CO63ZxQF0rtYuXzY6dmwNybCiFdjdy5QpkMotbkGSaDC7qFnl1Yd8mmPLqY0UFWVbSXb2F6ynWp3NWbJTNeIrkxLn0asuWOGjNNi0hjQZQBbi7f6vRbMZGUQDFyRfQURhohwDTEglXY3M1/bQG5JjabNZ11X702HKrjn/a30TLBS7fTwyfwJxHTS9CRVVXlp5YGQC0jr4/DIvLLqAHP0kHzYUBwOalevJuXRR075WndOzuLRL3Zq/saDaZ4YG2FkXKb/E85kEHlrzhhmvraePYU2zQhVYyySwK0TejF/SlbwH0CnQ2PKzMCdewAmTAjPDQzHHf1M2oYDOmc2v5yWTWmti0Ubj2gKKk/5RIxxG0DwatZY1adybSWlS0pxH3MjWkTMaZFED5+OOQhtZJIEkmIs/Hfe2NPWQKFrpJlIswGX193geKBMB29lEUptZdC1uTqho4upetg9dr468BWv73idMmcZXsWLR/EgImI2mPnb+r9xTvo5zBo4iwFdBrT3cP24a+hd3PP9PX6Ne4PBIBq4ZcAtYRhV09icHq58YQ1HKuwnIlFNYXfL7C60MSg1JqCLUGeg2umlqNr//1Mo4XiACruHslq33lQvTNSsWkXEkMEY4uObP7kZLhuayv8t3oOjBQ0WI4wSvz6vT0DRHG0xsuiO8Ty1bC/vrDuEqqp+dVSC4LtOPB7mHv6B2eddoIvw0whzZhauA7lhu75uQqHTFIIg8OfLB9GrSyT/WrrXbw5S5Sgch2/H2vNFVNHtZ5VeR+niUkq+KiF1VipRg2IQBDNl34/HnrMdc+qogPeXRDBKIsPS4nh55qjTOv1dFAXmTOjFs9/v9zOu0sp0MCVmgBRcba5XVom3nr7fXbjQxdRxjtiOMHvJbKpcVTi8DS3DFZQTxxYfXMz3h79n1sBZ/GLYLzrUYmRc6jjmDZ7Hgh0LcHqDF1QWycLjkx4nPSY9jKPTZv67Wzha6QhKSNVnf0kNj36+k8evGhymkYWXaocnpALSQBhFgSqHRxdTYcK2ZCnR51/QKteyGCXeu+0spr/wI7Wu4BelEUaJq0Z059pRaU2eZzKIPHRRP+47rzeLdxTy5o8HKahy4PIqRJoMDOwew+2TMhmRHkf+LYuoeO99EmbcfKofS6eDYM7KxLZsWdiuL0gSqt5rSqcZZk/I4Kax6SzeUciLK3LJKbJhEAVfuZ03lfERf2Wr5wkc3lqERk3MZbtM8afFdJ/bnZjhSaiKGfvh27D0SsTSq+F9InDxqvFJCsQU3lQuZsiwMcyZmEGfZP++e6cjN4xJ59nv92u+ppXpUL3hk6Bqc2OtRtL1XpYho4spoKCmgBv+dwM2jw1FbTpFRlEVnLKTt3a9hd1r5zejf9NGowyO24fcjiRKvLztZVyyS7MzeR2SIGEUjTw28THO63leG47Sx/5iGxsPlp9I46tPcwYMTo/Cx1uO8JsL+xJn7XwRKqMkoqqnbpWt4rMz1Wl9FJeLmlWrSP7db1vtmr2To/n4rvHc+Mo6HG4ZZxN2+AI+ATbzrJ789qJ+Qd/DbJC4Ylh3rhgWuOdayp//xKGbZxB93rkYk3Vb9NMBU2YmrgMHwncDgwH0yJROENSfgxRFxebyYpJELEYRQRC45qUofipdi6nLCkTLMVTVgIBK7d5qFI9CZP8BOAum4a3pR6A6KS8ie9SezBWXcK3xR4TSfqA+B3TODdZQ6Rpl5vpRaXy4OR9HEKndwdTmRhgl7piU2aGCBJ2FM15MeRQPc5bMCUpI1cfpdfJhzof0S+jH5VmXh3GEoSEIAvMGz2Nk8kgW/LyAdQXrEAQBl3xyB8gomBFFOL/n+cwZNIfs+Ox2Getrqw/ikf2/82ANGAQB3t+Yz52TO1/dR5zVqBmNC9Uq2y0rxHfidMeOTO2Pa7D07Yuha9dWvW6/lBi+v38K7204zGur83B65QaRKrNBRAUmZXfljslZjMlout6xJZizsoi/6UaK/vY4PZ55utWvr9P2GJKTUR0O5KqqU3KeDISe5qfTEkRR8Eu5q7IreGsG4q0ZiGAsRzSVIYguHPk7ECOKcRy5u9nrepCoVCMRVBm8Mhz7CV47H258DzKnhOfDdDAevXwgB8tr2ZhXHpSgaq42V1FVrh7VozWHeMZwxoup5YeXU+Gs0BRSFasqfEWQxW4ki0TMyBiSr0lGivTtlDhlJ89seYbLMi/rcEp+eNJwnj/neUrsJXye+zl5VXnUeGo4WioQKQ7H7gAAIABJREFULfbkmcvmEW1qv3C40yPz2dajNNZSoRgwOD0Kr63O65RiymKUmJDdhZV7SxscD9Uqe1TPeKLMZ/yfcViwLVlC9AWtk+LXmPhIE/OnZnPH5Cx+2FvCT/kVlNa4cX20iP63XMdl43uTFG0Jy73r6HLHHRy4/HJs3y8nelrT5hc6HR9BEE5Ep6zDh7f+DQy+XlM6OqeKQTq5XlI9Ccge34aRILpQ7LagMjMEwCQ0+j167PDejTD7a0gNw99AB0MSBRbcMorb397Mqn2lyBrZLsESYZR46MK+ndbUq70541dhr+94HbvX7ne89JtSSr4poce8Hg36Mx188iAZD2cgGnypVTa3jY2FGxnTbUxbDz0oEq2JzBs878R/L9tVxLvrD7WrkAI4WulAqzVUqAYMZTUuXF65U9qj33F2FhsPVvi5uwVrlR1pkrijEwrJzoDqdmNbsYLEX/86rPeRRIGp/ZKY2i8JgANv/pHU9JuwhFlIAYhmM93+9CcKHn6YyLFjECP1viKdHXNmJu4wiSlBMuiRKZ1WISXGwu5jNr/jzWVmHHlxDqrXRfc7XsNqEuhCFQu2uFm43cOKW4/PXx47LJoF926jWdvATs6X2wr497K9FFY5UU5RSN06vhe3TtB7DraUM1pM5VXlsb/Sv4BPdsgUf+Yrgowe4hMdpkQTafPT2PvgXqrWVBF/ts/dy+F18MbONzqsmGpMv5Ro9mhMYm1NjdOr2WMpZAMGScTm9GKO6nxialxWF2IjjC22yraaDEzunRiu4Z3R1K5bhzkrC2NyUpve19C1K96SEujfv03uFzluHJGjR1Py3PMkP9Sx6j91QseUmYkrNzx1Uz4DCl1M6Zw6149OY0NeuZ/jaHOZGQAoCrbNX2AYdw3ni5v4XOsGtaVweC30HK/1aqdHVVX+/OVOPghgQx8skSYJFfj9Rf2YMa5Xq43vTOSMFlP7KvZhEA0N6okA7PvsKB6FmJENm1lKFonoIdHU7Kw5IaZUVHLKc9pszKdK97gIalxeKu3udjVusJokFA1zjFANGGRFxWrqfEIKfGk5L80YyQ2vrAt5QrQYRV6aOQJRK7ynEzyKArnfwe4voabYlzsSlYJzQwXR553b5sPxianS5k9sRZIeeogDl11O7GWXYhnQ8Vo+6ASPOSuTyg8/Cs/FDRLIepqfzqlzbv9kDJII+D/36jIzKr5fgLe6GFQVJAOmpCxUr5uYsVdRvf5jYoZfwFFTIqAxX3rs8OMzp62Y+vvXe1okpCRRwGwQ8cgKaQlW7pqcxWVDU0/bflxtyRktpmxubdMJuUbGEGVAkPwXqoZYA45DDa3TG1upd2REUaBPchR7Cm2cpdEAtK1IirHg8Z66AYPJIBLRiSeCoWlxvDxzJHPf2hi0PXyEUeL5m4czsmfrGxOcMbhssGEBrHsePA5w15x4SQUSVAGhZC2stcPI2WBqG6tYQ1Ii3tK2FVOGhASSfv0rjj36J3q9/x6C1Hn/ns50TJlZuPLywnJtPc1Pp7UwSCJXDEvl7bWHNF9XaitRPE4Sr/hdAxOq0i//hSmlN+b0wVRt+IxXp1zMcHZrXEGF/ctA9oB0etUArdlfyjvrDmkKqeZckA2SwB8u7s8lQ1NP6z5c7cEZ7akcYYhAwF8wSVES3hovqsbi1lvlxRDVUIOapM7lptavWww5he2b6hcbYWRidle/b79+mN++dy2Kx4kqe3HkbqJi+esNzjVKAtePTutw5h+hEhthxGqU6JsSjcUgouV0Lgm+aNSAbjG8f/tZTOun21m3mKqj8NIkWPkE1JY0EFLgC06JBhXBXgTf/RVenQY1JW0yNEPXrm0upgBir7oK0Wym4r332/zeOq2HKa0H3qIiFJer+ZNDRE/z02lNFEV787DOhCrhvLuw9h2PaLIgSAas2WMRLVEAxE28merNX/FNTTa1aoAei6IEzqpwDb/deH7Ffk0hVb3hU8q/e5XYs66jxy8X0v2uN4gecTGOfetPnOPyKHywKV8XUmHgjI5MdYvqpnncmm1FMAhUb64mdsxJi1nZKWPbbiP5moYL2WRr51rY9k+JZtex6nYdQ3mtG6tJ1OyCFawBgygI3Dq+V5uMN1zUuLzc+/5WHr9qCJcM6ca+Ihuvrc5j6a4iapy+lJpIs8S5/ZOZOymDfikxzVxRp0lqS2HBcXGkBrEw9DqgbD8sOAfuXAWW1recro/UtSven34K6z20EARB7z11GiAYjRh79MB98CCWvn1b9+IGg57mp9NqrNirvUEVjAmVKbEXEdmjqV73MXlJ3YB9GmeJILtbZ7AdhIJKBxsPVvgdD8UFOafQxv7iGrKTosI+3jOJM1pMDU0cSqQx0s/NT7JKJE1PomBhAaJFbODmZ0wwEjc+7sS5VoOVGQNmtPXQT4m+KTF8vOVou9zb5vSwYFUeb689yCWDu5EWH8HRSgeNN6maM2AwigIje8bTs0vndiB75LMdjMvqwiVDfMK+d3I0T1w9hCeubueBna68dyPUljUQUr3+Y8Pugbx7o4g0+aKcDRyiFA/YjsFHc2FGmOpRjmPomugzoGgH9N5TpwfmzEzcubmtLqb0PlM6rYnNqS3MgzWhipt4M4Vv3svRsWnaJ8husMRpv9ZJ+XTrUbR2oENxQfYqKos25vP7S9rG5OhM4YxO8xMFkVsG3IJF8rchTrw4keSrkyn8oJBdd+0i96+5GBOMZPwmA9HY8Gu7oFd4etGEi74p0ewrsgUMs4cDh1vm5ZW5TH1yBUcqHHzxy4k8duVg3pk7lsgQ+yRJgkBCpIlnb+zcfSQ+2XKE7UereOTSgc2frHPqFO6Awp994qgRsgpPr29iF1N2w8EfoOJg+MYHGBK7IrexAUV9utxxB86cPdi+X95uY9A5NUxZmbgOtH7dlC/NT49M6bQOgbLz65tQNYUxPpWYfuNZvCnAbz2+V5vVurYVB8tqcTduzkloLsheReVQeW04hhd2VFXF7vZSVO2kyuEJnCqqqKinYBXfEs7oyBTAlb2v5IVtL2i+ljA5gYTJgYv8zZKZa/tci1kKkLPbQYmNMBIbYSS/wh72yI7b68vRfe77fYxIj+e9286id/LJHle9ukby4Z3juPGVdVQ7vcjNCDyTIJMUF8UHd4yjS1Tn+t7rc7C0lse+2s3CuWOJ6KRuhJ2OdS9oCimAB8eb+MePLuaPNhFnCfCUV1XY8Apc8HjYhtheNVN16L2nOj/mzExqVqxo/QsbDKBHpnRaidgII5V2DymUMcPwLRPFn4nBTnkvgUkG6LN/IQf6zEAm8PMxacLV5O38Hr+4gDESJv4qvB+gHbC7tP/+QnVBDnSdjkql3c2iTfksWJVHWa0boyQgKyqiIDB9eHdmT+hFUbWLl1fmsvlQBW6vgiBAtMXI5UNTmT2hF5mJ4U1rPOPFVKw5ln9P+Te/Wv4rnLIz6PcZRSPZcdncM+KeMI4ufPTrFsOeQlvYxJSsqHz+01Ge+nYvGV2jePWWUQzpoR1y75cSw+L7zubpb/fx6dajCAJ+vZciTRKCADdLK5g//XJi4yLCMu62wO1VuPf9rdwzLZsBqXoNVJvgccKOj0HR3lkflSoxpZeBJ9e4eGxagIa5shs2vwXnPQZieIL6YnQ0qteLYrcjWttnV1Wz95SjEgq3+wq6JTNEJUG3oad9U8zOiCkzC9drrzd/YojoBhQ6rcnt2TbSbf9ktOBz47MIvo2uTCs8NsXI/y39iKeM33Kk5yW8pl5O1cFdRPabSESvYSeuIcYkU/lwV6xCY8MVFQadfrny8ZHaxhGhuiDHWQMbULhkF0sPLmVRziKK7cW4FTdRxihGJY9ixoAZZMVltXj8oeKVFf76v128vzEfQQCnxxeVO7np7ktZ/GBjPqJAg3IRVYUqh4f3Nhxm0aZ8BqbG8NT1w8K25j3jxRTAxO4TeWLSE/x21W+DElRmBPqKkbx0/qudzsmvjr7Hm/deMDClVa+rqipLdhbyr6V7ibMaefKaoYwNwoI9OcbC41cN5g+X9ufznwr4evsxyu1uJEGgS5SJK0f04MKBKZhyHPDdQ5D9A0id8+f7r2U5dI0yM6uTm2d0KmqKQGhaAP1lqpkJr9dy79gm/qZlNzgrwRoeW3pBEHzRqbIyTO0kpuBk76m4ib0xF37l68MlmfAl7Au+mjNzLIy/G4bdBBGnV21CZ8acmYH70CFUWW5dm3uDhKobUOi0Bru/5Kadt6MKDrRaJd4/3kxKlMhTq6rY/elCIkzv404ejHXcTSfOMeDlSmmVv5AyWmHK7067FD+A4WnxfLLlqN9mc3PNjuOnzjlxboRRYmTPeL9rO71Onv/peT7c+6Evna6el0Cpo5R8Wz5fHviSrLgs7h95P2O6jQnfB8W36Tz7jQ1sOVyBy+uf2lhHnX4KlNTkVVS8ispP+ZVc+uxq/jvvLAb3aH0jqc65Gg0D5/Q8h7ej3uaZrc+woXADqOBWGtZQWA1WzJKZmb2v4dZVr2LMXQH9L2ufAZ8i/VKiWbKzsNWup6oqq/aV8uTSHGRF5feX9GdKn8SQbcutJgM3jknnxjHp2icMmA6b3oCNr8JZd7XCyFuXQ2W15BTasDm9WE0S3eMjGNw99sT38MPeEj7fWsDX907q9JbunQp3LQhNLywHJUlc2sfAE6vd9E8MILxEg89KPUxiCk427jWlBSisbgMMURYyrgRpyRxUSUBQZfA22mhy18L3f4Xv/gJXvgQDp7fPYHUaIFqtSAnxeAoKWvU3JEh6mp9OK7D/O/j4NgSvo8nA9s1DjNw8xBdBcakSuaqHK91Z1EknAzJzpMUN32S0wpDrfZs8pyGXDOnGHz/foflasC7IKipXjezR4Fils5I5S+dwuOowLkW7rYKsysiyzK6yXcz/bj4PjX6Ia/te2zofrBGqqvKrD35i8+GKE9GoU0VRfaYnNy1Yx//untjqESpdTNWjf5f+vHjuixTVFrFo7yLWH1tPtbsak2giJTKF6/pex4TUCUiiBMkT4L/XQfIgSMho76GHTL+UaP61JIddBT6L9DirkW6xlhYt8DcdLOefS3IoqXFx/3l9uWhQCqLWdlNrIAhw8T/hjYtg4FUQ3f4Wzl5Z4bs9xby0Ipddx6oxSiKqqiIIAoqqEm81cefkTM7uk8iDH23jqeuGkRDZOSOanRZzVFBW6H+eYmHEyzXcPy5APZ7iBXO09muthJTYtd0c/QCfSHrtfAzO/QiioukedQLP8d3LT+8ERwWMmt0mQ9RpGnNGJq7c3FYWU6Ke5qdzajirYNFMX8uJ4zTrpgqYBZkMjvF7w7s86p1NBE5mSsvIFgt8F5FMvrXBhHth8kOnbfqxxShx3ag03l13CI9GKKY5F2RJFLhkcDdiLCfT/BxeB3OWzCGvOg9vgDT4xrhkF//Y+A8ijZFcnHlx6B+kGdYdKGd5TrGmkGquMXGzr7u8PPjRdhbd0bzzYSjoYkqD5Mhk7h5+N3cPb2J3o8dIOPsB+PBWmLsUDJ3DDMHu9vLFTwW8uDKXwxUOrnt5LQAeWSEx2sydkzOZPrwHUUE47O04WsW/luawt6iG+87tzZXDu2PQ6jjb2iT2hWE3w7eP+nbE25GjlQ5ufGUdZTUuao+H3huHpO1uB3//eg9/+nIXlw3pxvjsru0x1DObyCRfEnUzZCeIXD/QyDMb3AxO0uqebPalt4URnwlFO4kpVYUPZkLZfoTGkaim8Dpg8e8gLh2yzwnf+HSCwpSViTv3AEyZ0noXlfQ+UzqnyE/vac7DdW6qv58UeB0VIXi4VlrJ096rONuwi99ZvwIhGhBh9FwYPQ9iu4dx8B2DeZMy+HBTPh536BsbZoPIXVOyGxx7esvTHLYd9hNSFasqKF1SirvYjWSRiBkZQ/I1yUiRvgwPp+zk0TWPMqbbGLpGtHxNs6ugmrzSWmpdXqxmiV5dInl5ZS4Ojc9XveFTqtZ/RJfzf4ElYwSCZMCRtxnHvvVYegxs9nXwRai25VdyuMxOepfWSwXVxdSpMPZOOPQjLP2DL1rSwfloUz6PfLETOGnwUOM6+Qd0pMLB41/v4bGvdvPwxf2ZOa6X5nVyS2r497K9bMwr5xdTs3lp5kjMhjZ2pJv8G3huDBxeB+lnte29j5Nfbuey51Zjc3iRm1mo2493LF+8s5Dle4qZ2i+pLYaoU4fRAoOvOf4wb3pB+MhkM+9s13D9k0y+yEuYzCfqMCQmtp+j39HNcHhtg5S+YHaOAZ+g+uYhuHtTW49apxHmzCwcP29v1WvqfaZ0TglVhTVPn4xm1yMoN1VAReA6aQWFhp7kT3mK9B49oPtIkAIbKpxu9Ii3smDWaOa8uRGHJ/i/R4tR5Lmbhjdo1uvwOvhk3ye45IapfaXflFLyTQk95vVo0Gf14JMHyXg4A9HgewaqqHyY8yF3DQut5MLpkfnf9mO8tGI/RyudSCLICkgieGUVp0aNVHONiUNpXKyoKm+syePRy1qvLc0Z3WfqlBEEuPw52LcUdn7a3qNpkueX7+cPn+/A7pb9ihfrY3fLOD0Kj3+9m38s3tPgtSMVdh78cBvXvrSWQamxrHhwCrPG92p7IQW+VKvz/wpfPdAuu6V2t5frX15LtcPTrJCqj9OjMP/dLeQU2sI4Oh1Nxv1C07Tk4H3RnJt58nharIjzDzENhQL4/t7H3BbuUWLomojcXmJqzTP+tVEE0YerjuqjPkGm066YMjNwt3KvKcGgu/npnAJHN6M6qzRfqu+m2hSRgovp0hq2R01g+nexFMUNO6OEVB3jsrrw1pwxRJkNWAxNL+PNBhGrSeKVmaOY1q9hWcTivMUINBSvskOm+LNiUmekEj0kGsEgYEo0kTY/DXepm6o1J/8fumQX7+5+N+j0QIDdx6oZ/8T3PPr5DvaX1OLwyNS45BP/1hJS0Hxj4lAaF3tklc+3Hg16zMGgi6lTJSIOrn3Tt6gvy23v0WjyyZYjPPv9vpAK+RwehTd+PMjCtQcptjn50xc7ufTZ1aTEWlj+wBTumpKF1dTOgc1BV/u+/02tbwPcHJ9sOUKF3aPpIFO7awXH3rqPw/++hiPPzaRo0aM4j+w88brTK/PvZTltOFodAJL6Q/cRx13pQsRghsxpvjS2MGNI9BlQtDm1ZbB3Maj+88SD4008ucZFpbOZjQOvE9Y8F6YB6gSLOSsL14EDrdu4Uk/z0zkVqvJxyoHF0l+mmnl2g5uS2qbXKclCBQdKa7E5PUx44nu+213U4PWiaifLdhXx8eYjfLGtgHUHyvBqNLrt7IzJSOCH30zl3nN7E2U2IIkCZoOIKIDJIBJploizGpk/JYsVD07h7D6JDd5fWOXklZ/eaeDaB2DfZ0fxKMSMbNi2RbJIRA+JpmZnTYPjHsXjM20Lgp+PVHH1i2sor3WfKIsIluYaE4fSuBjA5mrduUxP82sNUofDlN/Ch7Ng7re+lKIOgssr88jnO1tUyOfwyPzpy138Y0kO14xM49tfT6ZrR2qUW2dG8eYlMPBKiEps/j2tgKqqvLTygGaIPZicXVWFFTkllNa4Otb3eSZw/bvw0kSfVXqwu2miEWJ6wNWvhndsx/G5+bVDzdTRzb6aMK//gieoPlzgE2IHV4VxkDrBICUkIAByWRmGrq1To6n3mdI5FY5VHiBG0f79LNnvYXGuTK0bev2nhtHdJQK5YZvwpWB7ZN9Gwe1vb+aJqwfRPd7KKysPsOZAGSZJRFFVBAEEwCCJzB7fi5vG9iQx+vR55iZEmrhrSjY/H62iW6yF9IRIbE4PUWYDPbtEMql31wZ17Kqqsia3jJdX5rI+rxxTRqGfCpBrZAxRBgTJP93SEGvAccjR4JjT6+RYzbFmx1psczLjtXVNZkY1RXONiUNtXCwH8lJvIbqYai1Gz4ODq2Hxb+Gy/7T3aE6weEeh5u5kMIv+OuZPzebOyW3XqC0kkvrD0Bvh2z/B9Ofb5JabDlVQXuuf8hRKzq4gwHsbDnP3tN5hH69OPawJcNv38OalUHWkgauUFqoxAiEhG2Z9EXYXvzp8BhTtEJlyVmlGpeoIqg8X+KzjddoVQRAwZWXhyj3QamJK7zOlcyp8V7SJKzSOV7tUFmz18N7VVnrGCoxZUMtRm0phjYJbVjE1WtTbaSiGZFXlwY9+xmwQT5g/uf2UmMwLK3J5YUUuf5s+iGtGtV/bidbGKyus3lfKsl9PJjkm8EZXld3DrDfWs7eo5oSgMSD7padJURLeGi+qrPoJKm+VF0NUQ9mgoLAoZxFX9b6qSTfo11cfDCikmtvYh+YbE4fauDiilTOrdDHVWggCXP4svDIZfv4IBl+DqqqszS3jrbUHOVRmx+GWiTQbGJoWx+wJveiTHP7F2Usrc/3CqaEs+r2KynsbDnPH2Zkdty/S5Ifg+TGQvwHSwttIDmDzoQqNyTq0nF2nR2HVvlJdTLUH0Slwx0rY/LavRshV5bMDP4GAarTirfEi952N5ZpH2zTaLHXtire8HFVREMJsdtHwxgYg8N94UH24wNeLS6fdMWVm4M47QOTYls+JFbVuFm3K55MtRylzDkXdJRH3rxVcMDCFmeN60i02ohVHrHO6UuOu4cOqXVzT6Hi1S6XSCb+fZOLCbN+8ccNAIx/v9uD0wsLtHuYMN/Hfnz38e62LPaUKZrMDe+Kjfovtphq71n/9D5/voNLhYd6kzBZ/nmp3NV/s/4JNRZuoclVhNphJi0rj6j5X0y+hX4uv2xK2HK4kLcHarJC69LlVFFU5ccv1NtcVE9Awzc+abUUwCFRvriZ2TOxJV78iN6pXxdrHilwrn3D1A9hfuZ8V+SuYmq5ty+72Kry7/tCJaGJ9gt3Y12pMbM/5kao17+OtKkaKiEGK7krZ4ueabVwsAGMzWrdXpP7Ua00sMXDtWyhvX8nCwnRe2Gij2unxU+M5RdV8uvUI2YlRPHBBX6b0DY+zm8Mts7fIf5c4lEU/wLFKJ2W17o6bkmaJgfP+Al8/ALctB1FCVVWKbS6qHB5EQSAh0tRqvZ0q7W68GiHiUHN2qx0ajnE6bYMpEsbdBWfdCXk/QM7XvtQ/BIhOQeh/OfatRVR++hk9b2zbtF3RZEKyWpGrqjDE+3eqDxuRSU1pKSCIPlwAEW04Zh0/VFVl06EKXrYM58B2Ffeh74k0GeiXEs3siRkMS4tr9hqFVU4e+2oXy3YVIQgcTxM3ggxlJbUsWJ3Ha6vzGJuRwB8vHUDvNtgY1Om87K/cT7HZym6TkeGuk1kda/JlVGBS+slnZp2bapcIgWUHvFQ6VZ5Y7ealSy1MzDZyn/c3fJvr0syiCQanR+HJpTlkJkYyrZfFZx5Wug+clRCR4Mt2GTAdTP622XlVebyy/RWWHVqGiIhDPpnZIAoin+3/jB7RPZg3eB4XZ1zcJhvQ3+8pZmoTa0hVVZn1xnp/IQXIznQEYxWCcPK4ZJVImp5EwcICanbUUL21mpTrUqjcUImn3INoFP1c/dyKm9d3vB5QTC3ZWYiisWYKZWMfGjYmLvn8CVAUTMmZJF37FyxpA3Hkbca29ZtmGxdHmCTuOLvlYloLXUy1Mq7Egdwd8Q9WfX8Uh6q9eJcVkBWFHQXV3LlwM/ec05v5jbz/W4NKhxuTJOJolKcc6qLfKAlU2j0dV0wBDL4WNr9J1do3+VidxqurDlBe68YoCagquGWFvsnR3Dkli/MHpGBqxgGnKSwB3AtDzdk1tUVPLp2mEQTInOz7pxExqR6Kn3kGx7ZtRAwd2qbDqmvc26ZiKm1ss1GlZvtwGSwwbEaYBqjTFKqq8tHmIzz93T7Ka9043BZfv+UK34JvX7GNpbuK6BZr4dfn9+HSIama18kptHHDK3VOpdr3qovMr9pXyhXP/8iCWaMYn6X3z9PRptpdDcDrsTH8vaSMqOOlB6V2heRIgQuyTzry1bmp/vZbJ+uPyDyy3MUbV0RwVX8jxZLI2trBWLMlzcV2Hc2ljaV7D+H+6GVU1iIIYkO7dmMkfHW/r3xg3C+gi6/EYfXR1fx6xa9xyS4UjXRoRVVwyk72V+7nz2v/zPL85fx94t8xhtltcEVOMX+7cnDA19fklrG3qMZPSAGUrzBj27wfd6EL0SJiSbeQeFkiiRcnIhgFCv9bCCIUflhIzIgY0u9MBwn2PriXqjVVxJ998vm0u3w3h6sPkx7jb9K0Pq9M03Ai1I198DUmtmaP4cjzs+hy2X0NUvoCibDGJFhNjNEjUx0XRVG5572t/FBixdlE7UF9nB6FZ7/bT4RRYvaEjNYdkOrz029MqIt+ALGDZvjVoQIvdH2YZ74sRjTuwXHccKO+YcuOgmoe+ng7v//kZ56+YXiLez11iTI3yM+uI9Sc3ZS4jmNUouOPYDTSZfYcyhYsoMezz7bpvU/Yo/fp03Y3lQww5nb48T+aJhR1BOzDBT53lVGzwzRAnUDIispvPtrG1z8XBuw9o6g+U6EDpbU8+OF2Nh+q4JFLBzTYPT9SYee6l9dSFWTUXMXXTmPum5v44I6zGNKj+aiXzpmH6biD6g/WCEolCbPXixHoahUptat4FRVDo0XGsRoVFXB64cr+BhyCwDNx8ai1Ta9Zmksbu0ZawV8Nb2L0eBEEjXWa53jK95a3Ydt7cPUC1sd25VfLf4VTDq6RucPrYEX+Ch5Y+QBPTX0KUTj1jVOPrPDtriLWHiijrMaN2SASE2GgoNLRZLT5pZW5mrVKvu/pE7rNzCR2hK8+yvazDdsWG5F9IjEnm0GAga8M9KudqnP1qy+mPB4jf/p6NZXlPam0exAFSIgyc/WI7hRXaz9PQt3Yr6MlIqwOi1HkL9MHtnrUUBdTrcgHm/L5YW9pyM55Do/M/y3ew7isLvRLidG4ctOoqkpJjYu9hTXsLbKd/KfQpplHHOqi3y0rxFtbJ0UuHKiqyu8//ZnPttpwYYImLOBrXb5J5a5aHdQpAAAgAElEQVR3N/PXKwZxbZCFqA63zKp9JSzbVcS3u4s0a6a0cnoD5exGmiRuHB1+m22dUyPu6qsoffFFXLm5mLPazoSl3Rz9Rs2BH59ucOjgfQ1TuOp2jv0QDJB9DkTpDanbkrr57+ufj53YRGoOh0fm/Q35WAwiD13U/8TxX7y7hRqntslEc8+wOW9uZP3vz0Xq6DtvOm1OUkQSsiqjCAJzuyXx4dFCYhSFcT0kzAb4ZLeX6waejODUuFW+2efl8r4G9pQquCWRT6Ij+dTa9NzSXNrY9dJyHjW8RYQQRM88xQOKh6pP5nFPWipOxf89J+qJit1IFomYkTEkX5OMFCnhkl2sLVjLu7vfZeaAmcF/WY0orXHx1pqDvLXmILKiNojw1P2pXfnCj9w1OYsLB6U0EAnHqhxsyCv3/2j1vidzjwQE0zsIgkrM8Bhihvvm9mBd/RR3PK6SC/DaBrJSFJHlypMnltSy42gVrgAbPI039oMxooCWizCLUeR3F/X367fVGuhiqpVQVZUXlu9vsV22R1ZZsCqPJ69tOp2ootZNTpGNfUU2cops7C3yCSiAPsnR9EmOYmBqDFcO706f5Gjmvb2JzYcqGlwjlEU/QHZSFPGtVG8UDp79fj+fbS0IqRu406Pwx893kBJrYVJvbUv10hoX3+8uZumuItYfKGNwj1jOG5DMPef05m9f7WLJriIaB/7q5/Q2lbNrMUp+fR90Oh6i1Ur8zTdR9trrpD7+tza7b7s5+kWnwAWPw9KHwdO002FDBF/Pt0ufCtvQdLRZsrOIL7YVaAqp5gTQm2sOMalPIuOzurK3yEZOoU2zCXkwzzCHW2ZFTjHn9G/9hYpO5yYjNoOuEV3Jt+VTbDBwXfcU3jhWREKEwqOTzdz9jZMYs8A5GRJHbSrzv3LSI0bg0r4G3vzJwxuRUbwY1wVPxagm79NUxGKIkMujhrew1hNSvf5jw+6BvHujiDT5RMOCLW4WbvecaNr+sUVC8TqhkRlQ6TellHxTQo95PYgaEIWnwkPBOwUN6okcsoPXfn6Nm/vf3KLo1O5j1dz0qs9OXGtjvK4MafuRKu59/yeykiK5eWw6vZOiGZORwLb8KoySfxZN/e9JrpFQPHFIpobrxGBc/WRHD+yH54JiBkRkjSVYU1bo9Tf25eqSoB2mQy6pON5/65/XDOGyod2bPb8l6GKqldh0qIKyU7DLlhWVL7cV8OhlA4i2GKl2eth3XCzlFNrYV2wjp7AGp0emT3LUceEUzYUDu9EnJYrEKLNm2PLOyVnc9/5Wv3zVYBf9kWap49qi4xM8zy/frznRNLfL4fQoPPTRdn787bQT392BkhqW7Spi2a4icopsnN07kUuHdONf1w4l1npy5+z2yVms3FuqKeCiBk4laqB2ISb4dkfmTszQd3A7CfE33UTuhRfhKbwbY0pKm9zTkJjYPpEpgNFzwVEBPzzZrHU84ItIRcTB7G98YkynTXlh+X4cAdN4mhFAHpmXVhxgfFZXXludh0fRWLAF+Qyrdcu8tDJXF1M6fgiCwJxBc/jHxn/g8DooMhi4sns3Lqy1M2dKNdFWkfuXOjlQoRBjFrikn5HXro5gpTUCjE7+76BITBcRT0XTaV1NRSzuNnyKBf81mqzC0+vd/H6Sf024ArwdG4OzkZCSHTLFnxXTfW53oof4IvemRBNp89P86okcXgdrCtYwsXvzWUD12V9cwzUvraU2yOaybllh9zEbj3y+E7NBJDbCxOhe8cgaf9ONvyfF2d1PTDV29TvxXqeMbbuNxOm9sR++7biQahl1G/tlS19EddvpctHdRGSPabCxr1UDFUp2lQDcclZP7j23N9GW8NWv6WKqlVi47pDmwjqU3E5ZUbnmxbXYnB4q7B56J0fROymavilRTO6bSN/kaLrFWkLK9ZzWLwmTQdIs/mtu0Q8gIHDhoI67QPrv+sOax4O126x0eHh77SGOVTlZtqsQm9PLuQOS+cW0bMZndcEcwGxiRHo8M85KZ+G6wyFFxEySSL+U6FOyZdVpWwzx8cRNn075m2+R/NuH2uaeiV1x7t7dJvfS5OwHICETljysYR3vQ5XMCALQaxJc8bwupNqB/cU29hbb/I6H4pK1Pq+Mgko7n289iqyRJRjKM2zbkSrKalx06chmRTrtwsUZF/PPjf888d8uUeTz6Cg+j45iYKKLGee4iVEU3IJAqSRytdWKTRJJvFKk4J1jKO5UxGgLgugNmEUTKGKRSCWTxJ81a78fHG/iHz+6mD/aRJyl4QmbLGacGm+y77OjeBRiRjZMd5Yskl89kd1r551d74Qkptxe5XhEyl9INbdJ7KuNVHB4nCzeWahpR974e5Jr+mKI3IsgnRSb9V39RIvYIPpmjDdiTn0IFH9xEmyqXh0xY67CW1uJbcMnlC1+jorlr2tu7NcnlOyqbrEWHr6kf9idFXUx1UocLrP7pXxBaLmdXkWlb0oUD5zfjx7xEYitELmQRIGnrh/KnQs3a9ZyNYXFKPLktUMCCor2RlZU3vgxzy8qFcpCwu6W+cfiPdw6oRf/um4YQ7rHBv29/+6i/tS4ZD7bejQoQWUxiPROjuatOWNPyU1Qp+1JmH0rB66YTtc770CKC3+Rfbul+dVn0FUw8ErIW+mrozqyyVecLRrwOgXktEswX/sYxGi7wumEnw825uPRUEChCCAV+GjTEc12DxDaM8wkiZToYuq0YWfpTtYUrKHUUYpBNNA1oitT06bSK7ZXyNeyGq08OflJfrXiV7jkhoYEO81mdpq1fzNdLuyKFB1JyZfFeMpuDphFA4EjFtdJywOOa1SqxJReBp5c4+KxaQ1NofINBrSe7MHWE9VxuFp70zcQS3YWUuv2+q0pg90krkNLSIH/9+SpHoo55Qu/8xIvTkSKlCj8oBB3sRsxQiRmRAzdZkzGXR4NasN1TKjjq8OUlIEYGUfaLxcG8e34CCa7KsIkcXsb9UjVxVQrEWgxHWpuZ3KMhfQu/v0NToUpfZP46xWD+OPnO4IWVHWFehcO6taqY2lN9hfXnHLzXPBNOA9eEHqjPVEUePzKQQxMjeGpb/fidMuaEUCrSUJV4YbRafz24n4dVpzqBMaYkkL0OedQ/t//kjh/ftjvJ7WXAUVjBAEyp/j+qUfF88+jlNtJ1oVUu3KwtFYzmhSKAHJ7FXZ9sxzB0A0E//NDeYYJEPKmnU7Hwi27+Trva17/+XWO1R7Do3iQVd9zzSgaef6n5+kb35c5g+cwNW1qSLVAk3pM4q/j/8ojax4JyhlPVUVU2Yqx2/10m9V8jXHAiMWhTTxy2MY/ztN20P3LVDMTXq/l3rENa8PtoojWrzmYeqL6OIJJl67HSytzT5hl1RFqT6am0Pqe3GXDcRV8j31PDSnXn8wySJicQMLkhjbi9kPnQqPWP6cyvpY4TEPz2VUCcPXIHkFf71TQxVQrEW3R/ipDye2UBIgLk2vetaPSSIqxcP+in3AEWPSDz2XOZBB54uohXDCwY6ftVNrdmlGkUJ1ePIqC26u0KFokCAIzzurJTWPSWbmvhFd+OEBOoQ2HW8ZkEEmJtTB3QgaXDU0lwqSLqM5Ml7lzODTzFrrMno0YERHWexkSE9s/MtUE1hEjKHmmbe3idfwJVNwd6uJETU3DWxLcLnZTKKpKTIBnoU7Hp9xZzrwl8zhSc0RTAHgUn2X+9tLt/G7V7xiZPJKnpjyFxRB8m4+LMi8iJSqFJ9Y/QW5VLl7Fi8cZj7t8ArKjJ6piBkFGNFYgWgp8phNK8E2htSIW33SDJ88O/LsclCRxaR8DT6x20z/x5DogUlHQ+utprp4o+ZqGdYNWY/Ab5LklNeQW1/gdb26TuCXpdQ2/JwvWTEi6vEuT41O8VmRHBtBwvXQqduWhOkwHg8Uo8vLMkWGtk6qPPuu1EqN6+pxT3I22CUPJ7bSYJAakhm6NHiyT+ySy4ffnsnJfCS+vzGVDXjmq6ksFVFSV4enx3Dk5i2n9kjqFOUKg0G1LdjlONQosigJT+yY12Ylcp3NjzsoiYsRwKj/+hIQZN4f1XlJsLIrdjuJ2I5o6npNmxJAhOPfsQXG5EAOk5+iEn/qmOPUJuedddjppcgmHy+1+r4Xq/to9PrwbDTrhocpVxY3/u5FiezFetXnTA4fXwcbCjdy29DZev/B1jGLwi9bhScP54LIP+Oin7TyxeDdVVRZUVYB60kV2JyHX9m3JR/GLWPzZ+CTjpS1NvufPUyyMeLmG+8ednM88rmQcqgA07LvWZD1RgpG48SdTwQUEsuOygx57XkktRknE2SjrpqlN4pam1zX+ngRzARE9X0JV3QHXRKo3BgQvqA3lQ0vtyiH0OaY5LEaRp28YHtCpORzoYqqVuPmsdF7/MU/ztaDtsg0SZ4f5f379Rb9XVrj2pbXcNTWLaX2TMEidq44n3mpE1sjzD3UhYZJEjJ3ss+u0D13nzePor+8n/vrrEIzh2/ESRBGpSxfshcVEpnVvk5zvUBAjIzFnZuLcsQPryJHtPZwzltE9E1i+p8QvzTzUnndD0+LokxLN377arRntCuYZZpQEbhyTrqcxd1LuW34fJY4SPyHVXC+lPeV7eGL9E/xx3B9Dut/baw7y+DdHcXqiWjTeSLOE57jgcAeoDapjr9qDyep2TEJgkZidIHL9QCPPbHAzOElkrdyfPzkfREz6D6JU6Xd+oHqitDvSEI0n1xMRhghuGXBL0J+r1u1F0SjAD7RJ3Jrpf6orlZp9D2Pu9jGGyBxEg0azXdWIr9IyuPEFS7DrZACzwd/u3WQQEYCxGQn85sJ+DOoe6/e+cKKLqVaiR7yVkT3jWZNbpvl6s3bZhra3yzZIIoqqkhhl7nRCCiArMYpIs8Hv4R/KQkIU0KNJOkETMWwYxh49qP76K2L7R8Cuz6GmyPdiVDL0vxx6nwcteJgAVDk8fLz5CK+tPsCxsb+GF7cB20iKtjB7Qi+uH50WtlTgUIkYMQL7li26mGpHrhrZgycW79F8LZTFycWDuyErKo/9L7CDZHPPMFEQmDW+V8ifQaf9ySnPYUfpjhNpfHUE00vJKTv5bP9nWGsvJa9Yxeb0EhNhoF9KNNeNSicx2j9y/eGmfP7+ze4W1dcZROiTHMOdU7K4cGAKb689yJNLc5q81vvyVOZK3zR77Ucmm3lnuwcnJuZ5HsCBBWPZZMzJXyOIHr/zteqJGhNrjmVkcvBzpNVk0CxfCLRJfCrpdZqoZlwFN4Klmv5DP+NobR4G8aRUcHgUXIK/MUdrpOoF4zAdZzXyt+mDeH9jPkXVTtxehZgIIxOzu3LLuF6kxAafctqa6GKqFbn//D5sWbC+RROE2Shx45j0MIyqaVwtrBXqCIiiwG2TMnhq2V6/hpWhNM+97WzdplwnSFw2Us6NwbBhHupOE0Jjy/Bdn4PBDGfNh7F3gDm4XH+3V+HPX+7ko81HEAXBF2kQxBMbgIXVTv7z7V7+vWwvlw1N5bHpg7AY2zcCYB0xnKovvmzXMZzpxEYYuWhQCl9uO6bZbLe5xYlRErhudNqJ39JtZ2fw6g95IbV7AF9azXn9k0lLaF3zJJ224Z1d7/gJqVB6Kbm8Kgt++gBH6aQT71+6s4hnvtvPpN5d+cXUbIan+87NL7cHNMNqru5HAOZPyebX559M/5s3KZPMxEj+8uUuiqpduLwyjRNWSg3d2K5mMUZouPFw8L6G83NarIjzDzH8yn0XXyi+TStP9QjMiUvhuJiqXFvpi9QdcyNaRCzpFhIvSySyT6Tf57EYLMwfNr/JzAJVVVmfV86KnBJKbE5qXF7NvnGBNokdh7aBwdiiaFBgBAQ5niTbb/jPFYkU1B7F7rETaYykm7UnF/1rOw655ZvYLSXCKHHPtN5cMiSVS4Z0LPMjXUy1IiN7JvDnywfy6Bc7QxJUVpPEu/PGEh/Z9jvOblnp1GkZ141K419L92q+FswuR2K0mRHp4be61jkNqD4Gb16MqboAwegBt/9OJe4a3z8//BO2LoRbv4LYpjuu291ebn51PbsLqzWbT9dRt2Hw5bYCdhZU8cEd44hpo+JaLSJGjKDwz39BVRQEsXNuyJwO/HJaNkt2FoUsgMCXGjN3YsaJ//7VuX3YX1zD8j3FfhtUgbAYRPomR/PkdUNDvr9O+2P32Fl8cPEJx74Tx0PopSSIHsS4VVBPTNXNZd/tKWb1/lJ+d1E/Zo3P4I0f8zTT84Op+1GBhesPc++5fRpk8Uzrl8zUvklsO1LFKz/ksvFgBXaXF4MkkhBp4uax6Qzs8hjuD2dgUpt2EaxSrXytjEWuq99SzNjz52Ht+TJlSwso+aqE1FmpRA+ORpAEbD/bsG2x+YkpixTBZZmXckXWFdrfu9vLoo35vLLqAJV2T0AzmfpobRIbohPB425Rel1T4tXpUVi1r5TPN8Zw77m+9MH8cjsfbjwSsJ4q+JIWAUn01YVp/Ra0iDBKXDw4hdkTeoX0GdsKXUy1MtePTsdikHjwo+1+ZhSNiTCKmAw+IdXW+Z11uDwK5k4amQKf++HvLurH41/vafb7bozFKPLv64Z2uHoUnQ6IvRwWTIOaYgQliI70XidU5fvec+caiNR2SPLKCvPe2sSuY00Lqfq4vAq5xbXMen0DH9w+rt0iy8bkZESrFXdeHuasrHYZgw5kJ0Xz3E3D+cV/t4S0iWcxirw2azQ94k9GkwRB4NkbR/DoFzv4aPMRPF5VM+IFvihBhEliTEYCL80Y2ak35c5kjtUewyAa/Ho/hdpLSTDYABka+d+pqs8u/4lvcpAV9XhvtIa/qVDqfspr3Qx8ZDFXjejBnIkZZCf5aq4EQWBYWhwv3BwopS4Tx9i7cKx7gQg0aoGO87H3bBoboivO7tTkzKD40wdInZNO7KiTdV4xw2OIGX5ScKoqoBoZ3+0S/nDWHzTXF0XVTm54ZR2FVY6gNy3qaLxJrLhqOfL8LSGn1wUjXh0emVd+OEDPLpF8vOUIPx+t4oqhqfzn+mHc895WP5MMrfFpoSLw2S8mMO/tTZTYXE0KSUHw+QncMDqNP146oMOu13QxFQYuHtKNfy3LoV9KDD/mliIg4PLKeGUVo0HEKApEWQzcNimTa0elERvRfrvLvshU5xVTABmJUYiCdlFiIOrcXkb2bDrfWUcHgEW3QE0J1BNSvf5jw+6BvHujiDT5JvgFW9ws3O5hxa2RvnNry+GDm2HOYs3LfrGtgK2HKzV/t03tGrplhT3Hqvlg42FmjusVlo8cDBEjR2LfskUXU+3MOf2TeXnmKO5auBlZUZucBy1GEYMo8sbs0Yzu5T//SaLAY9MHc8PodF5bncfXPx/DIAnIioridGGMsOBRVCZmd+WOszMZk5HQYRc4Os1T46lBwP//X6i9lFBFand/R/WGrzXnLIdH5olvcjTrwkOt+3F6FRZtyueTLUfo1y2G/2fvvOOjqNM//p6Z7dn0SgkkhN57l3YqJxY8+4lwVuzlbOdZzu5P7/TEXlAs2Lg7RbFgQQHpLUivCSGBAOlt+87M748lIclOkt2QEJB5v168XmR3duY7m+x3v8/3eZ7PZ9blA0lLCC6zq4918qNsK6wkI2sulgYCql/V/rgI7rtxZVeieMGSNgVVXgsoCJL32O0rEiAgO9PxFo8nveM5mh5cJQ4vU19dTlGVt0Gj7HBoTnldOMGrwyvz6i97uG1SN2bPGFpTEjwqI56VWcUhr7mqsRolrhzRiW7Jkfxw1zgWbMrnjSVZHC5341cUfLKKQKAFQ1FVJvZI5IZxGQzpHBv+m3MC0YOpVuCj1fvpHB/BW9OH4JUVftlRwMEyF06vjN0caMwclRF/UnwBeXzyKdszBbB8TxF3z/uNj28Ywa7DlTz93U5Q1UZ9tCItRl65cpDmQkJHJ4iiPXBgHSjBZX2yCi+t8fLgGQ3IgyteyN8IhbsgMVjm942lWZrlWaHtGiq89Ws2V43s3GZziW3wIFyZG4m99NI2ub7OMcZ3T2TJfRP4ZHUu76/MweeX8Xq8yJIBoyRikAQsRonrx6Zz+bBOxDVRVt63QzQvXj6Qxy7ow4q9RZQ4vBx+7p9k3HsnY/p3IimybRq9dVoWq8GKqqHOFq6XUtGPRyhZ9EGjc5ZXVhCP02S6Gr+i4ldUNh8o4/xXlvPxDSPo37GJkn1BQJ7wCLfsiuc28Qt6CzmIKJiEY3NwqaqtLlg9Rl/JFHwlkzFEbkO05iJKDlTFhOKLw18xENUfeK+KHNrB2vUfrKPYERxIhesTVZtwxGYg/ODVKIlcOKhuufpr0wZzwasryC1x4g1jE3tYeiwPTul19GeJy4amctnQVH7LK2N9TgmlTh9mg0iC3czkPsnE208N6w09mGphSh1eXv1lL5/OHIkgCJgNEuf0a9fWw2qQU7lnamVWEXd+tpE3rhrCkM5xDOkcx8VDOrJwy2HeWLKXPQVVmAwiqhqYeEdnxHPT+AxGnySBrM4pwuo36mSkanPfaBP/XOHhlmEmYiwN/E3Jflj9Opz/Up2Htx4s50BJsDFmuCUv6/eXttnGgHXwYEre/6BNrq0TTFKkhbvO6h7oo/rfz+xesRnjhRdhNxvokmhnbNcETaWwxoi2Gply9Dtsb8UuOnW2YdIDqd8NybZkfHLwRlE4XkqyU6ZwfgFxf7y/yTlLKxlzPLLaigqVHj/T3lnDN7ePpXN8wxmqlVlFXPf+elz+AfzCANKEQ1wp/Ux34QCRgosK1Uaxqu31WX+M/sr+UNm/wWtZNQSCth4sZ8ehyqAyx+b6RNUmlPK6asINXvcVO4Ies5kMfHHLaK6es5adhysbL9UjUBL8h55J/PvygZrZyYGpMQxMPXX71/VgqoWZtWg3U/q1o3ty6I7dbYWqqqesmt+a7GJu/2Qjr145mOHpxxaSZoPEhYM6cOGgDrh9MuUuH5IoEGUxnpL3qdPGyD7Y9GmDwdTQ9hIT0gw8v9LDU5MaWGCqftg8D875FxiOZQO+2ZyPxx/8BRTOrqHLJ/PlxoNtFkyZu3bFX1KCv7gYQ7x2X5jOiccgiYx0HmRIskzKpG4tdl7Rbkepqmqx8+m0PdHmaIa3G86KgyuCMlSheik59rhRfGqz5blDkdVuKnPj8Pi5/3+bmXej9hh2H6nk+g/W16kEyFHb8Yz/qhYbYzUmSSA5Kvj74J1l2UFZnJb0iQqVcINXj09BVdWgTegoi5H/3jSaRTuO8ObSLLbnVyAI1Nyj2SAhq4GN7JnjujCqy+93I1sPplqQPUcq+XrzIRbdPb6thxISfkVFFIQT6m3VEqzPKeGWjzN5+c+DGJXR8ALOYpTaXD5a5xTHWQJq4yUMT0w0M2aOgztHNFw2pSoKni1rUU2xIMuofpmD2cWau7Th7BqqaqCZua0QJAnrgAE4MzOJOuusNhuHTjC+vFxMaelNHxgGejD1+0NWVAZGXsgKdR0IwaVpoXgpyVUyojW8Mr3aNNX3I0XENpm5UVT4La+MvBKnpkT/o19t05Qch9BK7MLpTRIEgXP7161Icnj8LNx6OEjQpcV9okIgXE8oo0FsMAiSRIHJfVKY3CeF7MIqVuwtotzlQxAE4iJMTOqZpBlY/t7Qg6kW5Klvd3DrxK5N1qKfLHj8p574xIb9pdw4dwMvXj6QMV0T2no4Or93vFVNGvD2TZI4r7uBZ5d76ZWo/XlSvX4Kn3kMnxyFIBkQJAlH0liwpQUdG+6uYf2SkRONbchgXJkb9WDqJMObm4d9fMtu7In2CGQ9mDop8csBKevcEicOr59Is4GMJHuj2YA12cXc+kkmLq8CqXYEoxdBCG8+URUJwZiK4jzQrDK9ahrq+4kcegGF858JKXOjqCrvr8zhkfN613k8r8RJZm6pRmdYeCV2ofYmDU+Po120tc5jh8pdGEQhSPaiOf1ioRBhkrCaJcocvqD+rHBFK5I0jJe16JJop0uids/Z7x09mGohFu8sIK/UyYxRndt6KCHjPcVK/H7LK2Pmh+t5/rIBjOue2NbD0TkdMEeC0rT/x+MTLAx+q4p7Rml/6YhmI6lzPoLIlJrHOn+9DVbkBF8yzF3DeHvbbt5YBw2m8MUX23QMOsF4c3MxprasEbwUYUepCu6f0Gk7CirdfLR6Px+u3I9PUfDLKn5FwSiKiKJAZLVy8JBUom3HlIN/2HqIO+f9ViOnL+ReS0T6K6iiu0Efofqoiojqj0QwX49gWF9nzmqOoIJW348re0PImRufrPLlbweDgqm5q/ajakj8N6fErqneJKtJ4sZxweqmVR4ZUeONPZ5+MeCokrFEgt2EyycTYTKQkRTBpJ7J7C2oZO7qXM3XhRoYWo0SV49OC3tcpxt6MNUC+GSFJ7/dzsPn9sIonTrBiccvnzKZqS0Hyrn+g3X885L+TOyR1NbD0TldsMaB2PQ02TVO5PI+Rl5e66VfksZnShDBVrckdVy3RP6zLi9IeTKcXcMIk8Sknm37ebD274d7924UtxvR8vsv5zgVUH0+/IcOYezYuGF0uOhlficQnxu2fwVr3gx41vk9YLJBcl8YfTukj2fxrkJu/SRTUw5fVgI/O70yL/y4m5d/3sPc60YwIDWGDftL6wRSAKovHuf+m7B2ehtEN4KW7F4tVMWA4ovBlTsT0RRVZ87yFR+gfN18Igeeg7lDH2InXhO2oELNfYSZual0B/e3fr/tMF6NDH5zSuzMeBkp7iCOCiRBoVyNYL3SnVKisBolLhzYgTFdg9sPIkwSikZu7Hj7xRQ1EFDd/8eeDOoUw1e/5TN/40He/jWbCwd14KEpPXnhx92aCsehiFYoqsqlQ1NDfHdOX/RgqgWYu2o/HWNtp9wi/1TJTG09WM4176/lmT/14w+9kpt+gY5OSyEZYPAMWDtbUxq9Nv8Yb2buZo1jRCMMmg5SXT+5cd0TsZgkzS+5UHcNq+vV2xLRZsOckYF761ZsQ2RSQ5IAACAASURBVIe26Vh0AvgOHcKQmIhoatmspWi3ozj0YKpV8Xth8dOw7p3Az95a77enAioPQ+4q3JKdHxx/wukd1+QpXT4Zlw+ueHs1H98wgr99vlnT4FnxpFDw9TAqN8zDV1iCaBGxdLKQeH4iEd0DKnmqbAZEvCWj8JaMByWQja+es8pWfIqvIBvBHIHvSBZRoy5HkAzNFlQIN3OjaDSiVrg9SLYsRGMJiB5UxYziTQorUOskHOFq6QculxYjIyKiIqCiIGLEzy/KIN6Vz+XzTIXvtx5idEYCM8d1YcBRhbrkaAs+f/DYWqJfzOGV+dvnmzEbRM7qnczkPsmgBtSl/bIScqaxPmaDyPkD2repF+qpgh5MHSelDi+vLd7LZ0el0E8lvP6TXxZ9e34FV7+3jien9uXsNl406pymDJ8J6+cEBVM5d9VV7EyNFnE/rCGrK0ow4saghyVR4Pqx6cxatEfT+LCpXUOTJDB9VOeTIhtuHTwI54ZMPZg6SfDm5mHs1LIlfqD3TLU6nkqYexEc3gz+RoRlvA4sOPiHMIeehiwe818NGsa79XH5ZK56Z41mwAH1+oe69MUYuwNP3iLK1x7G3LETqj8Sf8Ug5MqeqBrLR3ufiUjWKAr+9zipd3zSIn1A4ZY920zHrlnuKeeLPV8gd5iNVfCCoAAqqELAc6rEQ/F35ag4EQgWrQigcof0BTcbFiChYhK0lV0ni+uYIG5ipdKbW513snCrj192FtAu2sKTF/ZlTNcEJvZM5MftR6hfcdjQ5pl9yAUUfRlav5jXrzC2awILNuUjCkIdqXJLMzbNjZJA53gbT07tG/ZrT0f0YOo4eXHRbs7r345up4AUen08fgXTCV6IVXmr+Drra349+Ctl7jIMooEkWxJTu05lTPsxSLUm312HK/nLe2t57ILeJ7VXl87vnLh0SB8H2UtB1jZibBDJDJ1HQ3xwDT3AtJGd+XDVfo5UuDWV/RpCUBXsqsq1Y1pWra252AYPoXz+/LYehs5RfHm5mFJbvjRHstvx5h1o8fPqELBh+PhSOLQp5HnGJni5VPoVB1b+5b8CaLpXye2TNecarf4huWoEhtgRRA4B936FGdJPjBO/5Va64W5g+RhuWZ5RFIi3myh3+TUNzMMpexYFGJ0REKZae2gtt/9yO4qqgMGtGWpG9FQQDAJK+d8xJt2P4gkui33c8AGXSkuxCo1XJkiCig0PY8StzDM9wWXeR3H5ILvIwXUfrOOJqX2ZOa4LS3cXamYFj7dfzK+o/LQjOFADcIdoqluNxSDSJdHOx9ePwGo6uTfcTxb0YOo42H2kkm9PISn0+nj8CmbjiQmmDlYd5M3f3mRhzkJEQcTlr2tWuvzgciwGC9N7T2d67+nkFXuZ/u4aHj63F+f1b39Cxqij0yAXvwtvj4eyvCbL/WoQjRDVDi59v8FDoixG5s0cxdTXllPh9iOHEFGJAtjNRl7c+R88/7cG9bHHEIxtW4ZhHTyIQ48+iqooCGLbZ8pOd7z7czF1boXMVITeM9VqrHkTDv1WJ5BKm1WJ0wf77rQTYQqEA+9kevlos48lVwfK7iIED9dI3/OTPJRf12wJSUJci8b6h0z4eN04i9HidmyCh4fVuTztn46LYMGdcMvyzEaJVX//Axv2l/LMdzvYmFcWcuamftmzqsLewiqu++gnMj3/xm9wNnptySaR/KckDn2SQ/sZzyFG3YXq7VQTqN1xZkculZZi05CMbwir4KMHebxofJ1bfXcC4PYp/OOrrVw2tCM+f6DsTivoqU+4gWko56zGbBCDKiIizBIWg8S1Y9O4bmwX3VomDPRgqpmoqsqT32zn1oldiT1FpNDr4/HLJyQztalwEzf9dBMuvwtZ1VZGc/qdOP1O3tr0Fgv2LuTwrhk8cM4Qpg5s2QZqHZ1mYYmC63+GDy6AkizwNf4ljdEWyGjN+Bos0Y0e2inexsI7xzH93TUcLHPh8sqaEr4QaGJOiDTz0XUj6GA+g4N330PejTfR4aVZSJFtlx03JiUhRUbizc7G3LVrm41DJ4A3Lw/roEEtfl5dgKKVUBRY+Qr4XEFPySq8tMbLg2c0LE9txseV/i9ZsHx1s81fG1q4Cyi8bpzFGHFrTXbmKsMviKg84Z+BFwMKx14TbllehFlCEASGpsXx3MX9Of/V5SFnbuqjAnsLqthbIINwI6L5MObEHzHY9zT4moRzEjBEGyj8Jh9P/l0IxkjMyd2IG3UJ9xpmBQVSjQW4M4eY+PcqDzuLFOzmRfiSDmEcNR1Lxz64fQpzV+fy9lVDuO/zzVS4fE1WIxyP0l9TGUoFFx0T/RQ4KpFVD4KhHCluM91SBbp3uQZJSgP0YCpU9GCqmSzeVcDBMhfTTyEp9Pp4/QrmVt552FWyixt+vCEoE9UQbtnNvvIsEtPe5o/9/tuqY9PRCQtbHFy/CDZ9CitmQVXh0aCq+htRCARREQkw5i4YeCUYQ1O3S4m28ONfx7FmXwlvL81mRVZRHXEYj19hRHocN47LYHRGPOJRo+2Or73KkWeeYf+0q0h9602M7dquHNY6eBDOzEw9mDoJ8OW2UmbKHoFcVdni5z3tyV4MXm3J+ftGm/jnCg+3DDMRY2nAOFVQseaHXhKmeY4GFu7TpEWMFrfXKXMLBBQL+Ob2LD4Qp7JIGUL5pp8o27aMlCufJWbsNEp/egOTqBKV3henGKlZlieJ1BHu6pYcyeQ+Kfyw7bBmQBXG3YAqobhTcR2YjilhEc4dX1H0QxHeAi+SRSJqSBTJlyQjRUjEjI4hZnQMqmLAW3wG3qLJTBFXI6E9Bq0A90CFwl3fu3nzPAuTMwwgSjy8O5I5tbKCFoPEhtxSrh6dxltLs3H7Gt44g/AD02oa9c7qnIal3f8wROyhUhQwJxz7varAtmJ4ZPkjiKLIrQNu5cpeV55yegBtgR5MNQOfrPDUNzt45LzeJ0Xzd3Np7Z4pn+Ljxp9u1AykSpeVNjixIciU+w7z9JqneXrs0602Ph2dsDFaYOg1MORqOLAOdnwDVYcDz9mToee5kDqC5sgnCYLAyC7xjOwST3GVhwOlLqo8fiLMBjrEWEnUME4UDAaSH3mEkvfeJ+fPV5L6xutYevU6zptsHrbBg3FtyCT2ssva5Po6AVRVxZuXh7Fj6/RM6T5TrcCG9+uq9tViaHuJCWkGnl/p4alJDW/OFLogwmZptuiD9sJd5WbD15plbrIKi9bt4bUzXqFUtXOX2J7FQhmTxbXEjIylNLIna1f/m53fKPhMkZpleUZR5Poz6vZ9Pn/pAI5UuNmUV4bruAKq6lswUfStg4p1JXScmYK9tx1fqY/8ufnkPJ9D+kPpiEc3rgTRjyluFd6iM7nJ8DV2QVsEpH6A6/Kp7CtT+exiKxf1qi65Vni61x4WdrmX6rO4fDJzlu9j6sD23DYpg9VZxSzbW9zg0MM114XGvbMienbHlvYSguRAEOUGQkVw+AOf8VmZs9hbtpd/jPqHHlA1gR5MNYMPV+0nNc7GxDb2dzleAmp+rRdMLc5djFtDkahoYRGFCwvpeH3HBic2r+Ll+5zvuX/Y/USbGy+T0tE54QgCpA4P/GsF4u1m4u2huc4LgkD8tddgbN+e3Guvo/1zz2If17RccktjHTyY4jnvnfDr6tTFX1CIGBGBZI9o8XPrZX6tRHnjoh5PTDQzZo6DO0c03FKQYpVxOV3NN3/VWLiPNOxhzZ4SVu9388+z6gZytQOKWEsVZ0hb2S/4eMs0K3DAIGCQFZdqYojnTZwEB4IWo0Sp04eqqjWLdaMkMve6ETzw+Wa+2XwIRVHxhdBL2lBZmykxnbLl84ifciu27psRDIcwJZpIvSWV3fftpnxlObHjYmudScUQuZ3e3v0NXqt+gJtVqqKo8KdedZfUMhIZwiG2qWk1j0miyI3ju5KRGMHcVdqGuvXvRYqIoXTp+8iN9ItV02Dvm+jElvYmgqESQQitucotu/lm3zdEmiK5e+jdIb3mdEUPpsKk5KgU+ryZI9t6KMeNp5WDqTlb59TscFQju2QKviygw3UdiOwf6PFoaGITEZm/Zz5X97261caoo/N7IeqPkzEkJXHgzjtIvO12Yi8/sRkic9euyGVl+IuKMCQknNBr6xzDl5eLqRVk0QFEe6QeTLUG/sYFDvomSZzX3cCzy730StT+zh6TKmEwiGGXhNWmvtDDTyaVqnYKD48LDuJCzZjJiJwtrudLJXhM5S4ff5mzlvgIEzdNyODPwzohigJGSeSFywZy+6RuvL8yh/+sz0MUBBxev6bIQmNlbarXfTS4OANPUTy2jp8AIFkkIvtHUrWtqk4wJUgeIi17EbyNBxy1A9wqr4pRBINYN3ujIhAlOKhdy2eQBIqqPBypcFPhDhYzauhePHnbNDNR9Wmo982c9D2CVBUUSDVaJQS4/W4+3fkp52WcR/fY7k1e/3Tl1K1RayNe/Gk3Fwxof0pKodenNU17D1QeYG/Z3qDHnXucKD6FqCF1/XhqT2zVuGU3n+z8pFXGp6Pze8Q2eBBpH31EyZw5FLzwAqrSSJlM+YGA3PvObyF7CZRp75KGiiCKWAcOwJmZeVzn0Tk+vPtzMXVq+RI/AMkegezQy/xaHGvT1RePT7AwO9PLwQrtRb7NLHHmuGGU/PQmzt2rUHxuVNmPK2s9pYvn1Dm2sYIte5+JtPvLLDrd/Tm/3t2T76bZGJ2qve/+xEQzr6z1UuhoeJ6x4KW9oF3KpgJOr0xeqYunvtnBtR+sw11LHj0tIYLHLuhD5iNn8fgFfTTbEqrL2uLOuhlbj9GIJkuNSXDsxGtrBRcG5Mp+7LpnDztu34HiUTBEG/BX+SlZWkL2/2XXnLOHcTtiE9mb2gGu3STgUwLy5PXv0KMGK636ZIW3lmbV8YIK5V7qc+CNa8l7ZRqK91gFkOfQHhRHGapS69yCF2P0RgSx7vWKFhZx+L+HSbkshd6v96bLI13wFnvJeT4HpZbSn0/xMXfb3Ebfj9MdPZjSwOtXKKryUFzlwScf+4PadbiS77Yc4s4/dGvD0bUcHr/capmp/Kp8jGLwJCJXyRjsBgQpeDqvnthqU+QqapXx6ej8XjF17kznzz7FuSGT/HvvRfHU2vVWZNj1Pcz5I7wyBOZdBfNvhHnT4dVh8O7ZsPO7wHHNwDZ4MK7MjS10JzrNwZuXizG1dTJTgs2G6najys37+9BpgPQJYGhcrKZrnMjlfYy8vNar+bwXI/ZhFxM76TrKV83jwCvTOPDG1VRmfoO127GSLwEY0jkGMYQWGAva16qmdkDREAZBwdpA71FtXD6Z1VnF3PDh+iCLCItRYs2+4jrrsWpql7U5ti/h0Ad3kfvvSzjw6nSO/OdRZEdpjbAGgCpbQYGiH4vwl/sx2IMDxb7kN2mBnDarku/3+Jmd6cVqCFhW3P6dmwnvH9tsMOKngJg6r1PVgCrriqzgALMxefoGURQqNyyo+dEQkwKCgHP3qmPjiN4YMCquRXWVUPur2hPZPxLBINRUCXmLvJSvLD92rCqzMGchVQ309enoZX41+GWFRTsKeHNpFpsPlNUIS/hkheHpcdw4rgvvLt/H7ZNOXSn0+rRmZqoh9T7JLuGv8qPKalBApTWx+ZS69dQ6OjpNY4iNpdN7c8h/4AFyr7mWjq+9ikEthQ+mgqvkWLN7/Z7GvDXwxQ0BOfcZCyAhPGU+66DBFLzwQgvdhU5z8OXmYp/YuIR0cxEEATEiAsXhQIqKavoFOqEx9BpY1vTn5h/jzczdrO1zV6baWaP2xN6nV6MS4hajRJcEO5sPlOOVg7Mvtft1Rpn9DEtReegMM2M7aS8XH59gYfBbVdwzSrvH06tKVKjH+vcak+x2+xXW55TyxpK93Dap7qb1viKHppR4deapcv0CzfI4d85vAWGNXSuJ6HUGqmIk/pwEir4rRBAFki9NrnM+k6LQXg5s6qo0nsVDEOgRL/B2po+0GIGPt/joFC3i9AXK/j7KsrEl65s6WSVZUUmKtGAQhaCgMRRfKREZMz6MyAhA5IiLqFjzOZGDpiBa7IhGM1JUcp3eN0PkGqq2FePY4SDl8hQgtCqh2uWPkiixPH85f0z7Y2PvyGmLHkwBC7cc4u9fbMGnKDg8gd2L2mZmq7NL2Jhbhl9WuHZsekOnOeUICFC0jjR6hDECQWMasnW1IRgEKjZUED38WGmD7Jap3FxJ8iV1JzazZNYDKR2dZiCazXR44QUK//1v8q+ZSurwbASfA9QmFLK8VQGZ5tkT4NofILlPyNe09u+HZ88eFJcL0Wo9vhvQaRbe3DxMqa1T5gfHRCj0YKoFsSdB1z/AroXUbrDJuatuO0FqtIj74eD33amaeEs+jyaW/liNElePTuPD1fs1A6n6/Tp/N88jed9CvtrpbTCYqp0x65cUvDnrxcgOtbPm+bVMhV0+mXeW7eOm8RkYapX11S+Jq6Za0r102UcknPtXTY8tKTKekkVvIZptqKqKZEtFNAQyQzGjYyhbWVbrjAJTqpxH/9c494028exyD24/9EuSmJgmMm+bj8R/VWI3idhSkrCOPJZlMogClwzpiMkgagq+NiRPn0AZPgxE4GaguJcbDN8SqTrphYoppRvmTv0oXzuf2HHTA9eJSsA+YHJN75toVbCmmUk6/5hoWlNVQq79dTfEZVmmxFXSxDty+nLaB1PvLt/Hv37Y2aSnQXVwdcvHmTx9YT8uHtLxRAyvVXH75FbzmUqLTsMjBzfVSjaJpAuTyP8oH9Ei1lHzM8YZiRldNyWeGtl6iwIdnd87giiSdMu1KK5Z4KlsenVQgxo4/v3z4NY1gcVeCIhWK+Zu3XBt2ULE8NZROtRpHG9eHsZWEqCAo31TVVUEF3HrHBcTHwr0LjZlCF4PvypQgY0v5DMaPc5qlLhqZCd6pNjRUnHQktT+lD/yY49fuKhn4xUsjWXMHFhYofRpVLK7vqmwT1H4eWcBk/sEsigHqw5S5S8ADUVAc4eeIErQSHlc9IhLEC1RlC5+F6WyiCP/qcLe10bl1koUV621n6oy1O0mpbqMVRCPvlfaAe7Q9hKT0g30ThR5apKFdzK97C1RWHJ1BJWqlaGef6FyrJJJEgWuGZNGlNWI1x+85tSSpxdQWGu+VbOHy3jUVDdm7DQOf3Q/UUMvqHmutslxRNdnEI0VdV4bbpWQgoJfqduGoXOM07pn6utNB0MKpGrj9ik89OUWlu4ubMWRtQ4+WeG7LYeY+upyejy8kNeWZDFr0W6GPrWIF37cxZGKpuuam8Lllfnqt4Pc+1k23qoumsckTkkk+eJkDs87zPabt5P1ZBbGOCPp96cjGo/9SdoMNq7pe81xj0lH57RmzVuIeOvshKbNqiTpX5U4ailWvZPprVPvDwSyVKteD+tytkGD9L6pNkIuKwNZRoqNbfrgZiJG6F5TrUJKX7hkDhhCz+iqgohitHOb6WlUU7AoligEgqguCRE8f+kAHjq3N28v24dDI8uj1a+TqyazWQn+Hs+5K5IzuxxbbFdnzJZcXVeO36mamO0/FxUxrH4ghyfgx7Ti4Aqu/f5apn45lSPKahCCAzbRHIGteyA4c+1d26DwRuSAs2l/3etIUUkkXXY1nW7rRNTAKAq/PbaWs6oq15XXCjqMNpAab+vQEuFwqSZu9P0VT61AyoKXyX1S6JJox2KU6Jak8fuqJU9fW0Rk4R4v9/8UvD6LEpyY8WJKTMPadRjlq/+rOUZVDg5Ca1cJ1aa6Siiid93fpUE0EGXWs9ENcdpmptw+mQc+36IZSDVW0xt4rcLd835j7UNnIoXSxdnGqKrK+ytzePGn3ciqWlPKCKCoUFTl4e1fs3nr12zGdk3gX5f0D9njBkBRVNbsK+GLzAP8sO0wA1JjuGhwB65NuJv7fr0Lpz94py1ufBxx4+OaPPfktMkhj0NHR6cesh/WvgUaWWJZhZfWeHnwjEY+67IX1r8LEx8EQ2i9otbBgyn74vPmjljnOAhkpVJbtTRatNtRHHojeqvQ4xz486cwb1ogI9JYlspkR7DGYPrLN/w3No1VWcW8tyKH7KIqXF6ZCLOBXu2iuG5sOgNSj1V87C/SDoQb6td53n8ZH4jPYRUaF6Ooj6KCBxP/kcc3ev6GyDyQzy2LnkE5KoIhxayE4jGax9r7TsK541fKV35GUQheTNV5hKQ/JZH1aBYJf0xAAi6prGK4u9ZcqSow8ErY0LB/Xn3ZegWBu3y3sFLpW3OMBQ99pTyev3RqzWM3T8jgoflbggLb+vL0gsnKK+38/GNccC7YgpcJ4kbW0J+YsdM49P6dRA37U9Bx/qqeiKbiOmp+4VYJyarMkKQhDb4PpzunbTD13ZZDmo+HUtMLgWDs192FJ71xr6qqPDR/K/M3HsTla1iBqbqMcdmeQqa8tIz/3Tya1Dhbo+feW1DF/I0H+HJjPpEWAxcN7sC9k8eTHGU5eu0OtLe3J6c8B78aXnrYarAyvfd0zFLoQZ2Ojk49dn8Psnb5TW3TzRhLI4tvVYEdC6DfJSFd0jZ4EIceeQRVURDE07r44YTjzc3F1EpKftXoxr2tTMZE+Ot2+O1jWPlKoNwWAp9DQQLFDwndYOxfoed5YDAhAKO7JjC6a9P+bh6N8jJouF9nrdqLJ/1X8bDhI2whBlSKCi4sTPM+SAX2Rs/fEH5FrgmkAERjBZJtH7KjK/WLqswdeiIYTUSNvDQkjy1BDPQDmZPNRA+PpvinYlLbGbi3pKzugVEdIG0sbJ7XaGD78HgrI96uYtrIZLYrEeQqgRJnCT9GZCaJG5kVvwCT4c6a15zTL4WHvtyqeb7aJXoA9xifZ6SobTlxpfQLXaRUPo2dREXPM6jc8DXGxM51jvGVjsIUtzLotYlTEpEiJA7PO4y3wItoFYkaHEXqjal1qoQA+sb3JTVKb7toiNM2mHpjSVbQjkA4Nb0Or8ybS7NO+mDq+R92NRlI1cYnqxRWebjsrVUsvPMMYmx1d6OLqzx8vSmf+RsPkl/u5sKB7Zk9Yyi92wenfwVBYPbZs7l4wcWUe8qR1dDGYJEsDE8Zzi0DbwnpeB0dnQbIXnxMua8eoZpu4q2CrF9CDqYMiYlIMTF49u7F0l03eTyR+HJbz7C3GvFoz5ROK2KNgVG3wshbAgqbZXngc4A5CpJ6Q1LPZp/aZBDxa5T5afXrVPOJfCZu1cTTxjlIKPxvq4t/r/Kws0gh0iwwMEWsUfxzqiY8mLjS+1CN8ERT59dErJtNL11WStHCf+At8CKYbdi6jSJ2/F8CCna1yuOqFewE0YA75zfcuZvrejQJKlLEfgRVxaKq9Dk3jjUrS+nsA5Fjm7dOLOR1u54e7QY2KNqjqFCpWtkWPZTYXsXMXrsVc2Ic7XDgxcj54iquM3xHL+kgpF5Y9/02SDw0pSdPfrOjyfXZ2/7zGG3aRgTBFQaioPIP40eMEHfy3NgJLNm2OKg11uC3o7o6I9iygvpmQ6kSshlsXNu3acPg05nTMpgqqvKwvzh4lyFcjf/1OaVHvZpaR8TheNl1uJJ3V+wLu5SxuvTv/xbu5LmL++P2yfy8o4D5Gw+wJruESb2SuPvsHozJiK+jtqNFgjWBeefN45rvr6HIVYRbbrwvy2qwMil1Ek+OfRJR0He1dXSOi6rGezufmGhmzBwHd45oooTPEbrf2/b8Cl4ZdBlrPtyFS8hGFCDSauTCge2ZPiqNDjG6yl9r4c3Nwzp4UKteQ4qwo1TqwdQJQRCg08jAvxaioc6EpgKSLyZeyzpvD9qvf5lfVm5m1rmRTMkQsUky3+yF/+6UyUiN5W3/eXwuj6OCiLDOX9eUVkE0H6seKlpYROHCQjpe3xFT+3E4946h+MfZHJn3CClX/RNBMmqWx2mV+qXe8jo90h+lm9PNjIpKhro9CA8FbwYbjUZu25xGx0MlvB7fC+uR4D5QUYBIwcWF0gos4ztx0VYnGUI+bxhn0U/MJkoIZMAUyYY4+rag1185ojN5JU7eX7m/0YBqndqDEjWKCKHufF5bDGOytJ7JievZ/WBX5sqT2KNspwoLUTjpY8hjzPjbuP+3Bxu0rWkIo2ikS0wXxnYIIQA+jTktg6lShxejQaD+5ky4Nb0GSaDC5Scx8uQMpt5dno1PI6UfSimjT1aZn3kQWQn4b/VuF8VFgzsy64pB2M3h/dmkRKTw+QWf83XW17y37T1K3CV4/B4UAmMziSYQYEDiAK7rex2j24/W5dB1dFqCJpqn69f7N0gI/VLrckp45Mut5BQ78AnJyH4BCEyyDq/MnOU5zFmRw/C0OJ75Uz86xTdeRqwTPt68XKKnTm36wONAL/M7ddlbUIm7AYlxx/YlOHb8iuJxUPjVs6CqiGYb5vY9awKS/W47K5buImXKPazoamGfUIxV9VCZYaOgSyqjPL1pTDI01IAHwYcpfjlwzFy2w3UdiOwfCWxEMKiItrs5+MbNOLYtxt7/bCC4PK4+RnxcL/zI3w4eaPyNMtowzviCb9sN5cNVOfz95z/wlLoDey3z4epA5pMtvqNZus3EWwU6mCpRD24iqpaUfIEQT0qHQL9RlcfPlxsPMHvZPvLLXPhkteYdkzS8p46+ITyszuRN9V9N9q51Fw/ypPh+nXth+Ezocyb/ijJy79J7m9zUrsYkmkiJSOGts95CCnFdfLpyWgZTASO24A98uDW9AoGepJORKo+fBZvyqW8lEZY8qaxQ7vLx3R1n0P44d5NtRhuX97ycy3pcRmZBJusPr6fEXYJJMpFgTeDMzmfSwd7huK6ho6NTj+gOgT6LRkpsmzLdVAURIarxz+bXm/K573+bamXBg+dXrxx4bmVWEee9soyPrx9Jv47RQcfpNB/f/lxMnVu/Z8p/WLvnWOfkZs7yHLRWLA1tsHryttXJGFVX7xi7n8F3SvMW100FPACC5ESy7gO0zWVNMb9h5gGAGgAAIABJREFUsOUS0asTrpwN2AeMA9UISICfQE9V3c0hAYUYqrjB8F0jFxYDaopXfAypwzEB149N54GcsRTumYsZL0bh2Ab1v1d5eHa5lzfPszA5w4BJgu/3+vlqp7/Gl8ulmnjcfTmPVbh5Y0kWn63LRRSEOr5Z1b+T6lnTIApYjRIIAYEvk0FkyJg/4benwE/3Q6jZJaMt0Fd35mMAjE8dz+tnvs4dv9yBoiqa4mAAoiBilsz0ie/DK5NewW6yh3a905jTMpiKsRrxyaFp/DeGV1aIsp6cbhtLdxUiiSJQ9z7DKWVUgcPlnuMOpGojCAJDkocwJFlXhdHRaXX6XRJQ8/M1/OXblOmm6ofyvQZsBw5g6hjsr/fr7sJ6gVTjKCpUuP1c+c5qvr5tLGkJEU2/SKdJFKcTuaICQ3Jy0wcfB3rP1KmJ0+tn/saDx7XBGm71TrMQvJgSf6yxcmjIXFY0lWBOOYwrJwdL+y9QvLGoiglB9CE705EdPY4dK4BJ9fNJ5CvECSrUT+4YrIG+qB7nBPy+ErpBznJY8RLq3p95UgFVEDCgoKiB85W7Vf6x2MN7U61c1OvYOvD8HkbO7xH42amaeNn/J76XB7PptRWUOL2NzpP+o1kpUYSkKDMPn9ebdtEWuiVFHlWO7gaxKfD5dYHxNtAPW3M/I2+BSQ9T2xdjWMowFl+2mB/3/8i7W94lvyofg3gsFPApPiamTuQvff5C34S+WmfX0eC0DKYSI80kRZrJK627wAivphf6to/G0kqmt8dLscODrBEwhjsZFjuCGx51dHROEVL6QWwXKNjW6GGNmW4S1QFXiYmCSy/D1KUL0eefR+TkyRhiY3H7ZG75OLNZFhMOj5/bP93I17frtfgtgTfvAMYOHVpdQVGy6z5TpyK7Dldi0GiYCmeDNdzqnXARRC+G6HWYYo71JzVpLhspYYzaXPOYqkh4ZBuyowfCUZ+tDjFWpo3szA3Lnubbc9xEbPsEKvNB9lFlSiQv5Swqu5yLLSqWlKKtJHx0EbhKwOtEQMVY67KBiiRYdUDG7Yc/9QpeRntUAwoCT/un8bH8BxAd5JfLQGjvmdevklfq4o3FWXx8w4i6Fjzdz4b7sgIKqytmQeGuQDm3IASUHk2RAfGSwTPApi0sYTFYuCDjAi7IuIB95fsocBbgkT1EmiLpEt2FaLNeMRAup2UwJQgCN43P4Olvd+D0Na3xr1XTG2GSuGlCxokcdlj4ZRWt/Y+w5Unrb2Pp6OicWoz9K3x9Rx1p39qNy3DMdDMIow3xzAdoN+gqUv7+d6qWr6Dim68peP4FbCNGsGzYFFQ1ePEeSl+mosKeI5XsOVJJt+RgA0ud8PDltb6SH+g9U6cqFW6/ZjtTOBus4VbvGCUBm0nC41NwNyDJDoE4wGKQiEjYgjv66zrP1TaXjR5+bJFfbS6bfEndTKyAiEm0IhhEJvZIZOa4LgzuFIsgCOQUObgj08nsGR+z83Al7yzL5tsthzDmigjrdqHKPrw+P0OF6dxo+IYzxC2agh2CAMVOlQSbgCwYkFUBGREZCQGVT+WJfKiewZGYbMyGBXgKpqAVSDW24eT1K2w5WM4XmQe4fFi9z7XBFKg66HcJVBUEBIIUf0AFMqpjILUVIunR6aRHp4d8vI42p2UwBXDhoA48uSA0jX8tJFHgrN6tW05xPERbjRhEISibHe5kGGU5bf9EdHR+H/S9OOCTkrMM/KE1HgNgsEDHYdD/CgAEk4nISROJnDQRuaqKyh9/YvbKMhyGuruYYfVlKgrvLt/Hsxf3b/796QABJT9jp9b3gREj9GDqVMTYgIxfOBus4VbvZCTa+WzmSD5Zk8u7y/fh9ss4PXJNj5DFIKIC47onctP4Lry39xuW1tOGCNdc1mI08Lcp4zm7y1iiLHXbMB6c0otL31rJuH8tpqjKg09WkRW1lveWABhZqfZlky+DaBx8aHqWrmJ+0HsRbxMocqrICnyvjmCD2oMjaiy/Kv2RbfuxdpyDSVBx5twKarCATygbTi6fzBtLsrhsaCNG3PakwD+dNuW0XCkrbjeV/3qeG7KKmN31LNyh2R/VYDGKPP2nvhibkAVvS4anx2mqwoQzGRol4aT30dLR0WkCUYTL58JHF0N+ZqP9UzUYrZDSH/78KUjBXxOS3U7lhMnkZy4Ff/P7MmUFFmzK14OpFsCbux9zRtdWv44UaUd26MHUqUZipFmz0iTsDdYwqndunpBBjM3ELRO7cuP4DH7dXcimvDKKHV4izBLtoq2c278dCfaA+I3PdBXrDq8LEkYIx1w2wmjjol7jNa1VHB4/hZUe8sua3lRyYGXXG7fS0+9mxe2dGGHJBeCdTC8fbfbx1RU2zAb4bpebKb3W85DnOlxYkCJ2Y+04F9UXjbvwLBRP8KZ7OBtOBZUefssrY1Cn2CbHrNN2nHbBlHvHDg7edx+W7t256+3H8a3M54MmNP5rY1b83Giv4vwBJ7fyXGqcjYGpMazZVxL0XKiToSgIzBiVdoJGrKOj02oYrTBjASx6FNbPAYSACWjQcbZAQ8Dgv8DZT4HUsMBOYZUHk0GstasbINy+TJdPxicrJ/Xm1KmALzePyEmTWv06ot4zdUrSNclOgt103L3iEKIinyDwx74pNT9LYmBztrEN2hEpI7Cb7Joqc6GYy1okCzN6z9AMpLx+hWnvrKGwMrw+cFVRmbqqDxsmOOggFNc8Hm0ReGKCmVu/c/OSYGRK5yV8bhiI9/BnFH8fT8y4W0GV0KqtDGfDye2T+XLjQT2YOsk5bYIpVVEoef8DimfPJvmBvxF1wQUIgsAD50TTPtrKMwt3BMlV1sZmCiwMHp+UzqCn7qQsw0zMJZecyFsIi0PlgQlTAE0p1FAmw0GdYkiN071gdHR+F0gGmPw0THwQtn4OK1+BsrxA6Z/BEpBRH3Ub9LsUzE1L4Wp52EH4fZmSIOjBVAvgzcvDmHoCyvz0nqlTkppe8e92BK1zQvZ/ChGLUWTW5QMxG8ITqRAEgZn9ZvL8+udD9kKqjSiIXNTtIs3nvttyKOCBp5Gda6h3CSBqxEUUrvmcZ4dczStRH9Z53T2jzaTYRZ5fVsX2+a/hNdswJfcgauTko1LtAaq2LKJi7Xz8ZYcRzFaM8Z0QrZEhzY+KCofKw38vdE4sp0Uw5TtyhPwHHkB1e0j773+C5H1njE7joiEd+XLjAd5cmk1hpefoF7uKV1ZJjbVy0/gMzh/QHotRwtP5bfbPmIEUF3dCdgLDwe2Tmf1rNu+u2Mf0EZ1w+2R2HKqs8XgJFatR4uFze7fSKHV0dNoMU0RA6WnwjOM6TZTViKLhsxdu2ZCiqgFPFZ1mo/p8+A8fxtSh9SsmRJsNxelEVZRWVw7UaVkuHNSBp77doflcKBusoWAxijx+QV/ObGZPebUX5S+5v4QVUFkkC6/+4VViLDGaz7+xNEtzs7yx3iUAU0o3zJ368cnqQzxzlpX6uurT+huZ1t/IL/IArvP/FbVef1TF2i8oX/MFCef+FUvnAchVxRR8+X8oznIUnxvRaGny3rSsfHROLn73wVTFDz9y+IkniJ12JQkzZyIYtG/ZbjZw1cg0po3ozJEKD2UuLwICsTYjSVF1/9jN6emkvv46eTNvRHrtVWyDB5+IW2kUVVX5YdsRnvp2O33bR/P1bWNJjbNxw3gfU19dwcEyF95G1HRqYzGKvPLnQfTtoMtj6ujoaNMlUdsfKtyyod7tohpurtYJCV9+PoakJARTcKN7SyNIEqLFguJ0Idl1j7BTiQizgVlXDOTOzzaG7AtXjSiAURLxyQGvpdoIgNUkEWMz8uxF/RnXPbHZYxQEgafHPs1jKx/jx/0/4mrCoFYSJEySiVkTZzEsZZjmMdvyy9lfHFya2lTvkmPncgBixk7jyEf388GIiVj4Oeg8btXI7b47UKn7+VM8Tkp/nYvBHk/hl/93LOs16jKK5j9D6S/vEj/51ibfk+qeMp2Tl1MimJIVlSqPH7NBDNnXSXE4OPzMMzjXrSf1tVexDhwY0usEQSAl2kJKdOO7BdZ+/Wj/z39y4PY76Pz+e5i7dQvp/K3B7iOVPP71NgorPTx3cX/GdE2oeS7KYmTBbWOY+eEGfssrw+OXgybCaiJMEqIoMHvGUEZ2iT9Bo9fR0TkVMRskrhjWiQ9X5QSVzoTXpN76ogm/SxQZSrLBVYp/y2asXWID/W4nIDAV7XYUR5UeTJ2CTO6TwuMX9OXRBVtDCqgkQcBuMfDpDSPxKwqzl2WzaHsBbp8MApgkkTEZCcwc34UR6XEtsjFiEA08OeZJJqRO4J0t77C3bC9+xY+sHsssWQ1WVFVlSpcpXNf3OjpFNWwLsHJvsaYgV6i9S6bENCxdh/HqinLuTQl+/mt5lGY7Renid0H2ETPhaqxdhtZkvTx52zCldKVqyyKs6YMb3XCKMEmM79F0cOqTFQyioG9MtREnbTDl9Pr5cuNB3vo1m9wSJwZRQFZULEaJS4Z05Jox6aQnaE/krk2bOHjf/diGDSX9iy9abcK3nzGW5Af+Ru7MG0n75GOM7dq1ynUaotzp48VFu1mwKZ87JnXlqpGdMWj0HURajHw6cyRbDpQze1k2P2w7jMlw7DiHx09SpIX7Jvfg3P7tTlojYh0dnZOLq0en8dHq/Wh1ZoZSNiSKAmf3OXktJk5Kqgoh8wNY/Tr43CBKWHxe2rX3w0v9Ycxd0P8yMLeed5dot6NUVkKy/rs7Fbl8WCodYqw88tVWDpe7NTdZTZKIIASUgf/von50jA30T7/y50AljqyoqKqqueZoCQRB4MzOZ3Jm5zPJKsti/p75HKg6gMvvIsYcw+CkwZyXcR4RxqbXdyUOr2avVDhiOTFjp7H3/Ts4aAs+z5vy+Tix1nlM8Tio2voLgtlWp9y5OutVuuR9VJ+3aSEwUWByn+AIzu2T+W7LId5cmkV2oQNZVRGAGKuRK4Z3YvqozrSLtga9Tqd1OOmCKUVRmbVoN7OX7UMQqKlxrf4gOL0yn6zJZd66PPp3jObVKweTfLQMT/X7KXr7bUo//oSUf/yDqMlnt/p4o88/H39RMbnX30Daxx8hxWjX67YksqIyb10e//5pN2f3Seanv44jPoQ0cL+O0bz850GUO33sLaykwuXHbBRZtruQCrefi4d0bPIcOjo6OtWkxtk4t387vttyKOyyIatR4v7JPXThiVBRVVjyf7B8ViD7VMszTIRArVVZLvz4MPz4EJz7bxh4ZasMRRehOPUZ2y2BxfdOYFNeGW//ms3KrCKcXhlJFIiyGLhocMdGF+SSKKDpAtwKZMRkcO+we5v9ekMLeGwZY9vTvtcgXl67mn5Jx+asbCWFg2pC0PGegztBkVE9fs3zy44SjElpJF5wf6PXndw7uc4cqaoqry3eyxtLsgBw1OoDU4ESp493lu/j3eX7GJ0RzwuXDSQuovXLf093TqpgSlZUbv04k6W7CxuVKvcrKn5FJTO3jHNmLeN/N4+io6eM/Pv/hmA2kf75/zCmaORiW4n4a67GX1RI3k030+m9OYjW1tsNWJdTwmMLtmEzSbx/zbBm9TVF24wM6XxMYjTKYuSOzza25DB1dHROE567uD8HS11sOlAWckBlNUr8eXgq03XrhdBQVfjqNtj2BchNSDv7jspKf3s3VBXA2LtafDiSPQJZl0f/XTAgNYbXprV933drEm83YdawcQhXLGf42DF8s3VVnceOqLEY8eOm7oa27KpAtEai+tw4d60kotcZNc8pXheu7A3EjGtYBEgAIswSv+wq4NmFO7nrzG4YJZE7P9vIzzsKGl0jV/fHL99bxDkv/crnN4+uySzqtA4n1ZbgQ/O3sHR3438ktZEVlVKXl0tfXsLGK68h8swz6fTuuyc0kKom6Z57MHXuxMG/3o3q99d5TlFU1u4r4YvMA3y0ej9f/XaQnYcrwjr/oXIXd3y6kTs+3ciN4zP4z42jWkwgone7KIqrvByp0OU3dXR0wsMoicy9bgRn9UrGYhSP7lhrU106NCI9jkfO09VCQ2bpc4FAyhfsv9MgPhcseRa2fN7iwxEj9MyUzqlDQ8qCtcVynLtXofjcqLIfV9Z6ShfPoePNc7CmBfrtI3Dxl9jNuB+OYsnVEciqgKqCByNaGTrJGoXiqiRq1OWULHoLV/YGVNmPv/wIhV89iyEyAXsfbTVoUQC7xcD8W8bww13j2V/s4NyXl3HLxxv4eceRkNfIPlmlsNLD5W+totzpC+3N0mkWJ01malNeGV/9lo9LY2ezIQ8AS8c+qCqUeRQ+m/4wL147vg1GHkAQRdo99RR5t97KoUcfpd1TT1Hh8vOf9XnMXpaNwxMIsGRVRRIFFAU6xdm4aUIXzunbcJ9SHanzkZ159uJ+2Ewt+2sTRYER6XGszi5m6sCT24xYR0fn5MNkEHnlysFsz6/g3eXZfLP5ECZJrNNJJQhw1YjOTOqZxI0fbSCrsIquSa3X1/O7wVkCy1+sU9YHkDarEqcP9t1pJ8IUWMy9k+nlo80+llx9tI/E74Lv7oXeUwM+Yy1EtQCFjs6pQMdYG4M7xbIquzjouVDFcgQUInAxzfN31qs98GIEVKx48QmmoLbR6qyXMbY9MeNmULr43aM+U2YiByTT7opIpMhHAwcrZnyVffCVjMWqdiDGZuKTG0bQOT7wOX7jqiG8vngv//xhl+b9NbZGVlQoqPTw/I+7ePLCvhRWesgrdeL0yESYJTrF2UJqE9FpnJMmmJq9LBuPPzwPAEvHPgDIosTCXCdPuH1EWoxB5zhRCEYjHWfNYv/V1/Dtc29zv7MzqqpqBogAu45U8vD8rTy7cCef3jCSLonHjDJrS533aR9VI3XeWozKiNeDKR0dneOid/soXrhsII9e0IeNuWWUu3wYRIFYm4khnWNrhG/uPqs79/xnE5/fPLrVGth/N2R+SEO9KbIKL63x8uAZjSyGZC/s+QF6nnvcQ9l8oIwPV+1nhzAQz2YDMQUr6dshmhmjOtf5/tLROdm4cXwXfjtQhkvDa6opsRwJHwoSt/vuwIGZY0VdAk4sWvo7dbJe8ZNvpcPMu7B0XIBzdy6OXWUYalvuSC6MUZmYozfRIaITz098piaQqmZbfgWiQJBQSChrZJ8c6LPfW1hF5v7SOgJkHr/CmIx4Zo7LYGSXllFjPB05KYKpUoeXn7YfCfojacoDoDaiIPBF5kH+MjrtBIy4YUSbjd33PcM9/92CR2w6Fevwyjh9MlNfW8H8W0bTNSmyUanz1mJkl3g+WJnT6tfR0dH5/RNlMTK+Ea+ZaSM68cO2w7y5NIvbJrWdrcRJj6LAqlcDGSYN7htt4p8rPNwyzESMpYFFkLcqIFrRzGBKVVUWbMrn5Z/3kF9WrfwWAR5gfymb8sr4dG0uvdtHcfdZ3TmjW/M9hnR0Wovx3ROZ0D2RxTsLcIfouQkgIKMg4ST8jfrqrFfFmvcp+vYQkkXEkmYh6fyk4OuICioKB5xZXPPD1cyaOIvR7QPr3lKHl0U7jm+N7JUVVmUFMnP1e8eW7Cpkzb4S2kdbmHv9CF0FsBmcFFuCS3cXaqqthOoBAAGVv/9tONAawwuLnYcruGvBHjxi6B88VYUqj5/L3lrNg19s4Yq3V3Nmr2S+veOMExJIAfRIjqTC7edQeeMGeTo6OjrHiyAIPHdxf95bkcO2/PK2Hs7JS+k+8DYs9DC0vcSENAPPr2xClOLgukBgFiayovK3zzfzwOdbyCp04PIFS2j7FBWPX2FjbhkzP1zPG0v2hn0dHZ3WRhAEXrpiEEPT4rAaQ1v6CsioiKjHsVSOHtaJrk8l0eft3vR8uSdpd6dh69Z4lZHL7+LOX+5kW9E2AH7ZWaDZixrOGrkxVAJr6H3FTqa8tIzc4jB6M3WAkyQzVdwCHgAApU5vSw8tbF78aTdujXJFaLr3q8ThZfuhipClzluS6r6pVVnFXDRYl0jX0dFpXdrHWPn7lF7c859NLLhtbJ3SE52juMtAbPxr+omJZsbMcXDniEbkj0UjeCrAGrp1h6qqPDR/C19v0u5l1sLlU3j5572YDRLXjk0P+Vo6OicCk0Hkg2uH88x3O/h4zX4EhAbFHEQUFES0SmwbW8vVRcbS4WMEMVj8oXRZKUU/FOEt8CJZJKKGRJF8STJSRGC965bd3L3kbr6/+HuKHR58cvBnMNw1clPIikq5y8flb6/ih7+OI6oN22ZONU6KYEpVVVSNotNwPAAAFFXLg/rEUVTlYcmuQrSGEUpdK0B+mYtYW9t4AozKiNeDKR0dnRPGxYM78P3Ww7z88x7undyjrYdz8iE0HWD2TZI4r7uBZ5d76ZXYwPGqCmEuuBZuPcxXDQRSjS0mXT6Zf/6wk+HpcS2mOKuj01KIAkzpl8LBUicb88rwygqyoiIIgBp4PjbCRJnTq5nMDXUtByBFZCEI/qBzFC0sonBhIR2v74i9tx1fqY/8ufnkPJ9D+kPpiEc3lko9pfw/e/cdJmV1Nn78ez/P1O27sEtdegfp0ixgN2rUmGA3ghprbNHExNcaY15fY37mjQZ7SfSNBWMDu0awoBJRQBYUpMPSWbbvTju/P2YWZ3dmd2cWdncW7s91zcXMPOWcZ9g585x6L9q2iEAoJ+79bbL3yHWaW7CipNLH/32+niumDUj4nAe7lGgKzPY64wZujI4BkIj2rkU/v3AD8ebu1Y1rzTvuCtIGT8FyeRDbQdqAieQedVG9fStrAyxYHbviTFuY3K8Tn69tn7SVUgcfEeGPZ4zg+f9sYPHGPe2dndST1jm8gEQz7pzm4bGvfGwua6xB0YAruQUiHvhgVdzJ+mULX2H3B4+RPelMev7yWXpc8RSZY0+ietUXe/fxBUI89tGapNJTqjUZY3jpy41M/dM8LnhiIe8s38a2slqCkXGrArgs6GJ2UlFZSbxpVcncy4VPWovY9b+/weog21/dTvfzu5M5MhNxCK58F4VXFuLb6aN0wQ/DnmsCNTxV9BQ5XhcuO7aylOw9MiT2/a0JhHj8k7WEGo7pVY1KicrU5P6d9v5BR2suBkA0t8PihOFtH18q2oLVu+IGrUxmXGu1P8iSTe1zUzGgIINqX5BNJTpeVinVNgoyPdxx6nBueHExNQnGTzloZPcMP5oxIM/irOFO/rowXsVLYMBxxG3pa8SKLWWsizNvItGbyZCBt4u2amwblRKCIcONs5dw62tFbNhdRZUvGDOCyGNquFaeo4QMakxsw3zl8nls+fv1GH8Nu99/hG0v3k7NpqKm063qH/Ne1aoqQv4QWeOy6r1ve2wyR2ZSUfRDyAGD4dPNn3Jo3+y4PVPJ3CNDcpXBGl+Q+St3NHl96gcpUZnqmZvGuN65cbdlTTiD3KMvpvSzF9j0wHlsemgG5V/NxTuwfsXEAOdN6tUGuW1caXX8H45kxrWGDOyqaJ+5XyLCxH6d9q74opRSbeGUkd0Z0i2L+xqJo3LQEoHDrgNnerO73jbVTaUvTkuyMw0OuyapZP/vi/VxQ5Uk0zBoiTBnaXFS6Sq1vxlj+N3LS3nzm62Nzo9y4+Ml1x18EBpLNZ6Y7XW9OZ5+45G0bHrG6c2JK+SJqbQFK4I4MhyIHdu44ch2EKioPyzQtmw6ZQUZ1i0rZn9I/B4Zkvv+VvqCfKSVqYSlxJwpgMum9mfxxj1UtSAGgAgcMaAzBZmxX4K25G5kAnWy41q9rv0zmbAlJvfrxGdrdjF9fGG75UEpdfC567QRnPiXjzh+eFcm9M1r7+ykjhE/hbd+E/P2uuvqBzwuzLaouSXODVd6J+iV3Gpfa3ZUxqzaB8k1DFb7g6zXVcFUO3t72VbmLN0StyJVN3fI2rWWaW6ozH+WzCln15v7FL38uOXyUvHVXBCJu/x4LIk8fvgy2Rk2gYoAJmhiKlSB0gCOjPq35RYW/qCfy6f15/oXFrfoHrnuWks+egZMiM2zZjSxaMYPdlQ0s0qo2isleqYgXBka1TOn0QpJU9KcNr87aWgr5Co5PXLjr82fzLhWr9Oia1b7RaOe3L8Tn6/ehWnnxTyUUgeXvHQXd//kEG6cvYTK2thJ2wctVxqc9jdwtCD2i8MLP30iqSF+QNwbNqjfMJiI8hod5qfa1wP//r7JuX/9Jh/Pphuy+d1V55A99kd7e5s2PXQRGx84j+r1S/f25vh2bwZMEnOUQlRvOp+q9RdTvelcancdjrdvHuIQyhaV1dszWBOkfGk56cPq90L7Q36y3FkcM6SAgQUZOOP0aDWn7lrTh00DhO6XPZ5Qz1pL0jpYpUxlyrKExy8cT9/O6UlVqLxOmydmHMqAgvaPvn7Oob1Ii9OrlMy41qCBkw7p1lZZjtG3Uxq1zlXc9el93PLJLdy+4HZmLZ7FypKV7ZYnpdTB4bhhXTi0Tx73vPVte2cltYw4A469I7kKldML05+GwglJJ5fliT9oJdkJ73np7bMyrVIQjvu5ZmdFzPvRc4duGLaLTBdstbrgGjCl/tyhUIiq5fP39sZaTjd2VpeE5yiBRbBiOMGqgQTKR+LbcTzVm28j96jjKX5mK+VLyzEBg2+Hj42zNuLMc5IzpX74gp6ZPfE6vDhsi39cPJEeOd6kwkhEX2v2xDMQp4vq1f9pfNGMupwLdM3S4L2JSplhfgDpbgevXnUYVz/3NfOXbyGEEIizxj9AusvG67J5euaElFl+dXL/TmR7nXFb9eoiYZd+9gI7596HuLy4uwwga/JZe/exBI4eXNDmMaYAqvxVvPL9Kzy97GmC+SXMXv1D966FxVPLnqJ3Vm8uPuRiju99PPZ+imuglFLRbvvxMH70l484fngXjhiY397ZSR2TLg8vRjH3OvBXgy/2JhEkPEcqowDOeAwKD21RUqN75fDF2t3UNljSLLphUCwbT98xiOWgZt1iajYsrXdzlu62GdrIPA+l2sJzCzfij7MsX931niulAAAgAElEQVTcoexBh3K2/SQuCVBFbMU/a+IZlC54ARPw7e2NdWR1JmPUCU3ey/2gwf2rCaeROfpKsPqz5bln8e/cgOW1yBqbReFlhVhRAYXTHGlcPOLiva+zvU7mXH04l/5jEYs37sEXCBFsZhRR9DwpseyEv79uh80po9qvYb+jSanKFIDHaTPrmG58/ORtfHDF73l52U4sKzyZ1ZjwkqujCrO5fGp/pg0uiBsVur2ICJcd2Y//efvbuLE5mhvX6nbYXDq1X2tmMa5tlduY+c5MdlTtoCZYE7M9RIiaYA3flXzH7Qtu5+VVL/O/R/0vac6mo3grpVSysr1O7vnpSG56aSlva+DI+oaeAoNPgtUfEJh7B/aeIsS2w3GkxIJBx8OUa8O9UUkO7Yt23sTePDI//tLmiTQMAgjS7ivsqoPbup2VBJuY+9fN3oNE5jN1lrKY/VxdB+LudQg1a7+q1xubyBylptlkjjyJzFHH4u3xTxyZ8XviDYYT+55Y771Mj5PnLp3Ess2lPPHJWt78Zgsu26I2GMIXp+LYcJ5jot/fwjwvw7unRkdFR5BylSmAnY8+xohTj+Pos8Zz6xnhpbpLq/14nDYFmR7yM9tvTlFzzp/Um3eXb2PR+pKYVr2meJ02M6b0Zmyv+KsatpZd1bs4541z2F2zm6Bpfhx8daCar7Z9xSXvXsLfT/w7TltvdJRS+9eRg/I5akgBd81Zzp+mj2rv7KQWy4KBx7HH8z2hjAoKrr0yXJFy7r8hOV2yPEzu14l5jazm1dzNpMsWzpvYK6nhSErtb1W++HMv6+b+ZYQqCEVmuxxqfceLwWlUUv97lDv1QrasW8zud2eRNuRwTCiECQbi9ubE01SAXIyL6s3nkt7/z1jO0nrHeR1erh59Nd5GhvaO6JHN/WeN5s7ThrN2RyWPzF/Nm8u2Nnqt0QugNff99TptrpgWu6y7alzKlXT+zZspf/tt8mbOAMI9VQMKMhnXO4/h3bNTuiIF4LAtHr9wPCMyDO5gYpNvvU6b6eN78psTh7Ry7uozxnDlB1dSUlMSU5Eq+biEVbesoujSIr695luK/15MsDK8jy/kY1XJKu7+4u42za9S6uBx80lD+XztLj5Ysa29s5KSqouK8AwfDq70/VqRqnPdcYPwOFt2i+B0WMw4rM/+zZBSSSit9uNsZORS3dy/jSuX7+2ZOs5ahE1sA7grvw9pg6fg6jaQ6tVf4ite0eTy49ESCZCLEXwlk+od53V4Oa3/aVww/IJmrzPL42RUYQ5Du2XFnRST7DxHAcb2yuHUUT0S2l+FpVxlaucjj5Jz9tk4ctu2h2Z/cmzbwh9e/yMXjcwj0+Mg3R07v0iANJdNt2wPd502nN+fNgLZh2EZLbFs5zLWlq4lYOq33ux8aydbZ2+l65ldGTZrGP1u7Ydvl491960jFOltqwnWMGf1HEprS+OdWiml9km628GffjaKm1/5hpLK9om9l8pqipaHK1OtZHRhDneeOiLpCpXXafHEhYfSLVsnr6u2FQoZPlm1k/Mf/4Lxf3iPL9eXxN2vbu7ft+8+z1vfVlHlNxAKMH7dk5R9+HjM/jmHn0ftpuVkHHIs7p7DKfzlsxRMvwNPz8ZXkU48QK4Tf8kkjLFxWS5clouLD7mYmyfenNS19+6cvs8LoEF40ZjHLhyfUlNoOoKUGubn37yZ8nffpd9bb7Z3Vlos5POx+fpf0eWyX3DTBUdyfSDEu8u38uQna9lUUk1tIESay2Zo1ywuObIvk/t1avNKVJ2ni56mNlA/jkCwOsj2V7fT4+IeZI4MxzFx5bsovLKQlb9eSemCUnKPDFd0LbF4ZdUrzBgxo62zrpQ6CEzq14lTRnbntteLeOCcMe2dnZQR2L2bUEUFzl6tG6j+rEMLcdnC7175hkDQEIgXfCrC7bBw2MKTFx7KxH6dWjVfSjX07dYyLnrqP+yp9je6tP8PQmRN/BFWeg63fPwoM14pJ9MlHNJtPnkTf0egQQ+VM7c76UOOoHzRHJz5vRPKTzIBcgULV80YZk4cx5mDz6QgrSChNKIdP6wLv/tX/G2JzpOyLeHpmRNIc6VU1aBDSKlPbOfDj5Bz9lkduldq+71/wtG1C7k//zkALofFKSO7c8rI7u2cs/rKfeXM2ziPUINCo2pVFSF/iKxx9Vdhsj02mSMzqSiq2FuZqgnW8MyKZ7QypZRqNb8+YTAn/fVj3li6hZNH6upSEOmVGjasTRrifjK2J2N65fL4R9/z0mdrsD0eqgMhQia8Aq3XZeOwLH4+uTcXTO5NQaan1fOkVLRF60u44IkvEqhERUgAy72TLj/5kPEnDGDWzo2kE56WsSz0L870DYs5JPuwc6go+rDZU9fNkfLtWA8Ytr/0+2aD41q4+cXQm7lizMDE8h+Hx2lz1qGFPPP5evxxVt1IZNGMwV0yOaSnLjrREilTmfJt2kz5e+/R/+232jsrLVb2zrtUzJtH33+91G69TYnaXLEZl+3CF6o/fCZYEcSR4YiJzA3gyHZQvb663ns7qnYQMiEsSbkRo0qpA4DHafPn6aP4xT8WMaFvXsrPm20LNUVFeIbF3vC1lj6d07kpawcXbHudFdf/gY27qyivDZDtdTKgIIOjhxTgtPU3QLW9dTsrufDJhYlXpACMi1BtV2o2n8eink+xa4/BEwAbGGGt42XXHZx7xV+oxUll5BBHVj69b3ylydOWLXyF0i9eotPxV4EIO175IxmjT6R61RdNVqaCIdhTue+Bymce1je8HHwwic8iwuu0ufbYllfmDnYpU5na9cjD5JxzNnZOTvM7pyDfhg1svfNOCh95GDs79Wv2lf7KuO/bGTaBigAmaGIqVIHSAI6M+n8ytmVT5a8iw9X+QZOVUgemMb1yOXN8T25+5RsevWAcGKhdvYeqr7YTKKuFoMFKd+Idkod3VD5WnLkDB5KaoiIyTzihTdMsnTOXglN+xOAxOjFdpY6731zR6Kp9Ta+k5yRY1Zdg9QCu6hLgueKtZERiNg2xNrLAfTVzgpN5OPhjik0nbEIEsQhi4cNJwxhSdXOkOp10HWmDpxCqrUScLggFm13xD+DJT9cyqX8njhqc/BC/OoV5adx/1miue+FrauKE52mM12lz7sReGspgH6REU5Jv0ybK33ufThde2N5ZaZGQz8fm666n8+WX4z3kkPbOTkLSHPFjRKUNSEMcQtmi+jEXgjVBypeWkz4svf77oWCjS3cqpdT+cu2xA9myq4oFzxex5Z6F7HpmOVVfb8e3uhTfujJqinaxZ85qttz1OSWvfk+gtLb5k3ZQNUVFeIa3Xc9UsKKCyk8/Jev449ssTaWas7Oilo9W7iDeVL7EVtJz4tt1JOtcTi7q1oVSS6irlnnEz3THR3zg/jWzXXfyR+cT3Op4liwqiQnGS+wcqWQXfvAHDVc8u4i5S4v36TM5cURX/jx9FB6nRSJrSHidNj+f3JtbTm58MQ3VvJTomdr1yCPknntOh+2V2n7P/+Ds0YPcC85v76wkrFt6N3zB2BWy7DSbgtMLKH62GMtjkTEsA3+Jn+JninHmOcmZUv//KMedg20d2K3ASqn25/SFeNCkEViym1DcRYDB+MKtsZULt1K1ZAf5F4/A1TOzLbPZ6gIlJQRLS3H1Tmwi/P5Q/t77pE2Y0GF/o9WB6Z9fbIj7fsNeojppAyYSqi6n+ImrCOzZiri9pA2aSI8ZaazIhp/16MY1u/dwXFU1TtuFHVmga4S1nhGsB2c6d5THb/RvGBwXEl/4oU6NP8SNs5fQI8fLmH2IOXryyO4MKMhk1off83bRViyB6qieKqctWCKMLszhqqMGcOSg/BanpcLavTJV1yvVUedKlb39NhUff9wh5klFy/HkMKHbBD7d/CmG+s06+SflY6fbbH1hK77tPiyvRdbYLAovK8SKWibXbbs5Z8g5bZ11pdRBxviDbH90KY5dNTgaqUjVEzKY6gA7HvuGgitH4eyS3vwxHUTN8uV4hg5FrLYbWFI2dy45Pz2jzdJTKhFvfrOF2kDscLbGVtIrW/gypV+8TOeTr8fTexTBil3sevdB1t+3gX63d2Grw8HNBZ3572CIa6zOnB10Q9VuEIG0TphDzsb3L2fcvMQLjguxCz9ULp9HyQePxB96SLhC9cc3VzD78ikxaSRjcNdM/vecMZRW+Zm9aCNLN5Wyp8pHuttB//x0zjq0F4V58UcoqeS1emUqFDLMX7WDxz5aw6rtFdT4gridFoV5aVx8eF9GvtBxe6V869ez9fd3Ufjoo9hZWc0fkGIuGnERX237iqpAVcy2vKl55E3Na/J4YwzTB09vrewppRRAeNjezhposErVwk1L+eOHD7Fy5zosy2Jgp97cfszVjO4WHrJiaoPseGIZ3W6aEHdRnY6oteNLNRTYuZPqpUvp+eADbZamUokoq/bHfT9eL1Gotoo9n/yTTj+6Fm+/cQA4sruQf9pvKH50Zr2wL+W2xZ+sCqZNf4Gu6T/MIxLA9dpb+OJU4KKD46YPOTx+fqMWqPD0HYvYDqrXLopZoGLpplLW76qkd6d9bwTKTnNyyRH99vk8qmmtVpkyxvDMZ+v5679XUe0LUhm10kp5Leys8HHT7CVQM5qL+w3k2pDpUEHCQrW1bLruejpfdSXeEW33w7Y/je8ynvy0fDaWbyRkEp+sCOFeqak9p9LZ27mVcqeUUhCq8lO1ZAcE6lekymsrmfnSb7n7+F/x4yFH4QsGWLhpCW7bVW8/Uxuk5tvdeIcfGLGPaoqKyDz2mDZLr+ytt8k4ahqWV+fGqtTS2GigeL1EtZtXhHurBtfv8bFcHtKH9qKiaNveyhSE42jO2ziPs4ecXW//Hjle1u6MXcAreo6UWDaevmMQy0HNusXUbFhK9pSzGh16mDZgYr1zhYzh6QXruP3HHfPe8mDUKuMEgiHDDbOX8N9vfcvOCl+9ilS0Sn+ISqeHx77YzM+f+IIaf/LLObaXbffcg6t3b3LPPbe9s9JiIsLDxz5MhjMDSWToTITTctItvRt3HXZXK+ZOKaWgctG2uDdNa3ZvBOD0YcdiWzZep5upfScwtKB/vf1MbZDy+RvbJK9tIbz4RNvdZJXNnUv2Kae0WXpKlVb5+X57BcuLy9i8p5pgI8Gic9PiD7mL7iWqE6+3CgAJ4sh2EqiovyJgbbCWkpqSmHNfcnhf0hpZLTRrwhnkHn0xpZ+9wKYHzmPTQzMo/2ou3oGTkwri6w8aPl61s9n9VOpolZ6p219bxlvfbKk34a0p1f4QX64v4bJnFvHkjENTvoeq7M03qfx0QYebJxVPz8yePPOjZ5j5zkzKfeX4Q/G7zet4bA+9snrx+PGPk+bU8bZKqdZV8WkxJs5vSb+8QiyxuP6Nuzl1yDGM6TGcHE/8xSZ8xZUE9tTgyOnYAWWDpaUEd+9us8UnfBs34tu4kfTJzd8AKrUvgiHD/JXbeXj+Gr7eUILLYSEIwVAIt9Nm5pQ+nDuxd704c6eP6cHqHZVUN2iIj9dLZLnTCVWWsvvfT5B39MVRe9sEq7bGhH0xGIImtoH/9DE9uOuN5Y1eR2PBcSuKPoxfmWtEeU3T92Iqtez3ytQnq3byr682x61INbXmf20gxMK1u5n95UbOntBrf2drv/GtW8fWu/5A4eOPYWceGKtE9cvpx8unvszTRU8ze+VsjDEx86jSHGmkOdO4cNiFnD3kbDyOjn1TopTqGILlsauOAmS603n5vAeZ9cU/+c3bf2JH5W6O6j+Re0/8Dfnp9ed7ikMI7u74lama5ctxDx2K2G2zgmrZG2+QdeIJiDN+D4BS+8NXG0q49B9f1psSEh14ttofYta81fxt3mrOHl/I7acOx7aE6eML+dM738U9Z8xKek4PiGC5oxuBQ1jOJVQs20OXn3Wpd7zLcpHjjp3Ln+52cM6EXjy/cEPCHQbQ+AIVjdEg2B3Lfq9MPTT/+5hWAkhs4l21P8hD81Zz1qGFKdnjE6qpYdN115N/zdV423CYRVvo5O3EDeNv4OoxV/P++vf598Z/U1JTgi02+Wn5nNz3ZCZ1n4Ql+gVXSrUNEzIxi05EG9i5D/effDMA3+9azzVz/8AdHzzA3069PWbfkC+5eaGpqC3jSxljKJ0zl253/b5N0lMHp/krd3DZM182G2S2btW+2Ys2sX53FU9cOB4MdM3ysH537CJaENtLVPrFS5QtfBV3t0GR1fy2sufjv8UN+yIiTOo2Ke55bz5pKN9sKuWbzaVxVxOMJ5EFKqJ1znA3u49KHfu1MlW8p5ov18WOMW1qzf+GE++2V9Ty9cY9jN2HNfZbotoX5PM1u9hd6SNkDDlpLib0ySM7akzutv++B3e/vuScfXYTZ+rYXLaLk/qdxEn9TmrvrCilDnJiCdjSZIWqzoBOvTlzxIk8u/j1uNstd8ePh1ddVETG1Kltklbtd99hqqvxjh7dJumpg09RcSmXP7Oo2YpUtGp/kIVrd3P+E1+wdmclRwzsTFmNn5Kq5ofFZU/8GZYni5IPnyCwZyuW20nWeBddr+xbL+wLQL/sfgzIHRD3PE7b4pmLJ/KLf3zJovUlcTsQGmpugYrcoy7au2+ay+a8iak7QkvF2q+VqblLizFxfvOSmXhX4w/ywn82tlllavWOCp76dC3/WrQZ2xKMCUddskXwBUMcP7wLlx7Rn16LP6Hq88/pcwDMk1JKqY7CkechsKM65v3vd63ng9WfceqQo+mWVUBx2TZeW/EBY7vHjhowgRCOzh1/Nbqa5cvJv/LKNkmrbO5csk4+uU3jWamDy62vLmu0ItLUtJBqf5Av1uzmvukj+em4QlZuK2f6w59RURMgGO8mNErmqOPJHD0NO2013sJnEImtyKU50rjokIviHP0Dr8vm7xdN4LXFm3lo3mo2lVRTGwjSyFoZQHJBfH88qnuT6avUsl8rU5tKqvEFY/8wG11FJQ5jYHNJ7A/n/maM4b/f+pZ/LFhHIGQINPINeGPpFt4v2sqE4uU88Of/h52R0ep5U0opFZZ5ZE/2zFmNaTBML92VxuLiFTz2nxcpq60gy53Bsf0n819HxVY23P1ysDNdMe93JMHycgI7duLq1/oxY0woROkbb1L48MOtnpY6OK3dWUlRcVncbYlMC7Es4fM1u/npuEIGdcnkjWsO54pnv2LV9nL8QRN/BUCpBSycuZ/hLngbkdh9nJaTPtl9OKZX8+EHbEs4Y2xPzhjbk2WbS3lt8Wa2lNZQUunj8zW74naoN7ZARR2XLUwf1xOPs+P3pB9M9mtlqrGlzZOdeNfaS6QbY7hx9hLe/GYrNc2Mdw0ZqA4YvugylEs/LeWZoSGdGKiUUm3EOyqfPa+vjnm/W2Y+D51+Z7PHi8si88ierZG1NlVTtBzP4MFtsvhE9VdfYWdk4Bk8qNXTUgenpz5dG7fCk+i0kGDIMGdpMbf9eBiZHic9c9OYc/XhfLe1nCc+WcOri4sJhUIEjB+MjTjKcOZ9hCtnEWLXxs2Ty3LRNb0rjx73KE4ruUVXRvTIZkSP7L2vH/hgFbPmrU5oCGAd2xJ65Hr59YlDkkpbtb/9WivolB5/wly8Nf+b0ljsgP1l1rzVvPnN1qT+yGuMsGTTHn738jetmDOllFLRLJdN+pRuiDP5n6sgYGW7cffPbnbfVFdTVIRnWNssPlE6dy5ZGltKtZA/GGJLaTWrtpWzqaQqbgP5u0Xb4o4ISmZaiMOyYubpD+6ayb0/G8W3vz+RRbcez4tXDaXfmAcpGPpX3J0WxK1I2WLjsT2MKRjDC6e8QLZ738uLXx49gPMm9cKbYA+Ty2HRI8fL85dOJsPdKlGLVCvar/9j43rnku62qaxtfs3/xibeeZ02hw3ovD+zVU+1L8iD/46/4iA0N043xJwlxfzquEF0z+n44++VUqojyD6hL/7NldSuL4NEJ6sL1Fhwt9fHn2oCZHs79vLeNcuXk37YYa2ejvH5KH/7Hfq89FKrp6U6jppADevL1lPuK8dlu8j35tMto1u9fTburuLvn63j+YUbCYZCWJYQCkHQGE4Y1oVfHNmPkT3Dq+ZV1AbipJLktBAMe6rjh06wLCHb62RCz8G8f+ZbfFb8GU8ue5LF2xfjsl17jw+EApzU9yQuGHYBA3MHJvGJNE1EuOXkYQzuksm973xHVW1g77Lv0bxOm5AxnHRIN+48bThZno5dTh2s9mtl6qghBbhsi0pi/2ASnXgXMoYzxrXekIw5S4ppbP2IRMbpGgPPfLaem36k3bBKKdUWxBI6zxjOrue+pXZVScz8qRhOCzvdSZ9fHEL3T9cw/eEFPDVzAj06cCNYTVERnS69tNXTqfj0U1z9+uHq2aPV01Kpb33Zev5v+f/x6upXsSQcSBfAH/LTO6s3F4+4mMO6HcVNLxXx4Xc7CBmDP85koTe+2cL7K7bTp1Maj884lMaW8UpqWogBK4EFwSyxOKzHYRzW4zB21+xme9V2qgPVZDoz6Z7RnTRnWrPnaKnp4wv56diefPL9Th6Zv5qi4jKq/EGcttA5w80Fk3ozfVxhvZWjVcezXytTtiXMPKwvf/vw+7hr7zc38c62hFNGdmvVmvnD81dTFad1INFxur5giGe/WM/1xw3C5dC5U0op1RbEYdHpvKHUrNhN+fyN+Iorw5Nao4YKictCPA4yj+xJ+qFdsNwObv/xMJ74ZC0/nbWAJ2aMZ3j3jjfkL1hRgX/bNtz9W3/xibK5b5B1ysmtno5Kbb6gj5s/vpl5m+YRDAUJmNiepJUlK7nzszup9t2Br/jn1Ab6Nnq+kAkva75yWwUn/e/HeFwW5XGmLiUTj8kSISctuYVl8jx55Hnymt9xP7Is4chB+Rw5KL9N01VtZ78PzLxgUm+e/nQdtYH4Xa9NcTssfnn0/utmbSgQDLF2V2XcbcmM0w2GDJv3VNO3c/r+zqJSSqlGiCV4h3fCO7wT/h1VVC/bRbC0FhMMYWe4cPfPxt0/p174ChHhkiP60S3by8+fWMj9Z43ucDc1NcuX4x40EHG07lyKUGUlFfPn0+W/bm7VdFRqqw3WMvPtmawqWUVtMP5iDXWqAlVggaP7k9ibzyFY0fS8vqAxlNX48TgsHBY0bHdPZlpIIGSY0KdtK0ZKxbPfS+bcdBf/vHQiP31oAVW1QZoPtRjmsYXHfj6+VSsoZTUBnLaFL06vWTLjdG1LKKtuPkCcUkqp1uHMT8N5VOLDc04e2Y2CLDdXPPsVvzlxMGeOL2zF3O1fNcuX4x0eGz9rfyv/94d4x47Bkac3qAcrYww3zLuBlSUrm61IRRPLj7fHc1Stv5RQTWGT88+NAV/QECeSDpDYtBCHJfxsXE+8Ll1CXLW/VmnmGtI1i9euOoyzH/2cKl8w7rC6OmkuGysY4O4VLzGpoOku3X3ldliEGoknlcw4XWPQIX5KKdXBHNonjxcum8SMpxayuaSa644d2CGCsNcULSd90sTmd9xHZXPnkq2r+B3UluxYwsKtC+NWpEo+LmHnOzvxbfdhe2yyxmXR5WddsNPD90xi+fF0eZ2ts7s3O/88GDIINNrg3ty0EIclzDyszz5erVL7R6vVCAYUZPLpb4/mjz85hEFdMvA6LTLcNl6nTbrbJs1l0zPHy80nDeHz205k4uh+FN/0W0wowZWaWiDNZTc6WTGZ5dt9gRCdM+IvA6+UUip19c/P4F9XTOHf327nNy8txd9Y83gKqSkqwtPKPVOBkhKqFi0i4+jmg5WqA9ffi/5OTaAm5v2db+1k6+ytdD2zK8NmDaPfrf3w7fKx7r51hKJG+5jQJvZ88ix5x11B2uApWC4PYjtIGzCx3hA9AI/Twt2Chmmv0+b8yb3pl5+R/AUq1QpadQC222Fz+pgenD6mB99uLeP77RVU1ARIczvolZfGqJ7Ze1sF0397E+svnMGuRx6h8xVXtEp+RMILXLy2uJigqd8eksw43eHds8jP1MqUUkp1RAWZHp6/dBJXP/c1Fz39H2adN5bMFF2SOFhRib+4GHf//q2aTvk775Bx5BHYGToX+GC1u2Y3H236CNOgvyhYHWT7q9vpcXEPMkdmAuDKd1F4ZSErf72S0gWl5B6ZC0DVmgpMMLH559X+EKMLs1m1rYIqX2LTQrxOmxOGd+HmHw1N+vqUai1tFhlsSNcshnTNanS7uFz0+MtfWDd9Op7hw8k48sh62ytqA5RU+ggZQ7bXSbbX2aLhGedml/NG0E/Qir30RMbpprttLp/Wuj9qSimlWle628GjF4zjtteLOPORz3l65qF0yfLE3dcEDYFd1YSqA4glWOlO7Fx3qwwRrKgN8MbSYlbvqKS0yk96+S7yx55ILyPstwWcAz6o2gm+KnBnQlonSufOpdNFFzV/rDpgfb39a5y2E1+o/gJiVauqCPlDZI2rfw9ne2wyR2ZSUVSxtzIVqvBjZzgSmn8OEAgaXrnqMC5/dhFb9tRQGwgSbzZGmsvGGLhian+uPmZAhxieqw4eKRVm2dmlgB7/789suvY6+jz/HHaPnvz72+08PH81Szbu2TtPyR8M0S3by2VT+3H66B6kJxAtuvqbb9hx//3kbN5Mn8OvYVUVcb+wzY3TddkWxwwpaPE1KqWUSg0O2+Lu00cwa95qzpi1gKdmHsqgLpl7twfLaqn4fAsVC4rDPxiWgAlXrhw5bjKn9sQ7Kh9rP0yC/357OY99vJbXFm/GQqiKCizvKZjMX+96nzPG9uCSI/q1fKGmbcvhs7/BspcAAcuGUAAjNlm2m4wRd+zzdaiOq6y2jJCJs0BXRRBHhgOxYyswjmwH1eur9762M2yCFf7E4kQBCAzqkskHv5rKkk2lPPbRGt5dvhURwRLwBwzdczxcPq1/wvd7SrW1lPurTBs/ns6XXcbrN/03fxxyOrWB0N6o0YGohSw27K7i7jdW8Ie5K7jxhMFcfHj8+AIlvuYAACAASURBVAa1a9aw4y//S/WSJXS+8kpyzvgJj5T6OOWBTxqNwN0Yr9Pi8QvH47B18QmllDoQiAhXHTWAHjleznn0cx44dwyT+3ai9I01VHyxJbxTILblLbCzmj1zVrNnzmryzhyMd0TnFufhla828btXvsEfCBEn3ik12OAP8sJ/NvLyV5u5/6xRnDiiW+IJVGyH58+Frcsg5IdQ/d8+AXL7VCOPHgb9j4YzHgO3zkc52NiNVH7sDJtARQATNDEVqkBpAEfGD7eSaQPSEIdQs+41vP3OaDbNTunhKRMiwujCHP523lgCwRDlNQF8wRBZHqeu2KdSXkrWCj4ZeTT/VXgiu6v8eytS8VT5glT7g9z3zrfc8XoRJmoelH/LFopvuYX1552Pd+Qh9H/7LXLPOhNxOunTOZ3nfjGJLI8DK8GeYq/TZtb54xjXW5eMVUqpA83pY3rwwLljuPafX7Ns1ldULtwarkTFqUjVMb4Qxhdi9wvfUfH5lhal+/KicEWqxh+/IhUtEDJU+4Nc98Ji3vymOLEEStbDQ4dB8dcQqI6pSNURCUGgBlb/Gx47Cqr3JHklqqPLdediSextYV0FqWxRWb33gzVBypeWkz7sh55SO82my0+6sOudf1C74UVC/hpMMED16i8p+fDJesenu2xOG909Jj2HbZGb7qJLlkcrUqpDSLmeqc/X7OI3/1pKrST+Bar2h3jhPxvpkuXmF6M6sevRxyh9+WVyzjqL/m+/hZ0dG/H+kJ7ZvHHNEdz62jIWrN4FEBN/ymkJliWM6JHNnacOZ0SP2PMopZQ6MEzp35kXBvYktHgH4YWbE2P8IUrfWIOd68Y7OPEGt2WbS/mvV8MVqYaaitNT4w9xw4tLGdQlkwEFmXHOHFFdAk+fFJ4fFWf4VlyBGihZB8/8BC56BxyuhK9HdWzju46PO8zPTrMpOL2A4meLsTwWGcMy8Jf4KX6mGGeek5wpOfX27/yjzjiyHex8Zza1xc8jrrSY+ed1TjokiR5WpVJUSlWmjDH8evaSuD8s0PSPS7U/yP1vr2D8rfdTeNw0+s55HWdB03ObCvPSeHrmBLaV1fDs5+t5bfFmyqoDhIwhw+PguKFdmHlYX/q0YiBhpZRSqSFQUoNr2S6IU5FauGkpf/zwIVbuXIdlWQzs1Jvbj7ma0d3Cq4oZf4g9L3+P57eHJjw5/m8ffk9tnCDyZQtfaTZOjy8Y4uH5a7hv+qjGE1jwIFTsqFeR6vOXcqr8sPbaDNJd4Xw+/pWPZ5f6mTcj8lsX9MGOb6HoZRh1dkLXojo+r8PLaf1P46WVLxEw9Xsw80/Kx0632frCVnzbfVheC3c3NyF/iBVXrYiJO5UzJYecKTkEawuoWvOrmLSctnDWoYV4nNrzpDq+lKpMfbm+hF2VvrjbEvlxEWP4/Ff3cuhPxieVbpcsDzccP5gbjh+8z9eglFKqY6r4rDhuFNHy2kpmvvRb7j7+V/x4yFH4ggEWblqC267faxOq9lO7phRP/5zYkzSwu9LHB99uj1kIKVRbyZ5P/o9OJ11H2uApe99PGzCRtAE/BO4NhgxzlhRz+4+HxV/WPeiH/zwGcYKvBg387xc+bj6iiRAf/ir45H6tTB1kzht6Hq98/wqBYOxw0LypeeRNDfe87nxrJzve2kHPS3rW66lad986+v5XX6zIgmGWswTLs4lQTc+957EE8tJd/PLogW1zUUq1spSaM/XoR2uo9sfOkar7cWkuCFyt2Dy9rIRgvGX6lFJKqUaYQIjKL7YSb+LSmt0bATh92LHYlo3X6WZq3wkMLagfJsP4QlR8tCmh9GZ/uTHuD3Dt5m8xgcTi9FiW8NrizfE3fvsGhOLPOf71FBf3LahlT00zv5V7NkDx4mbzoQ4cfbL7cOHwC/E6vI3uUxd3qvv53ckcmYk4ZG/cKd9OH6ULSn/YWQK48j7e+9JhCTlpLl64dDJ56TqEVB0YUqoy9dHKHZg4ZXsyPy7V/iDfb69ohdwppZQ6UPm3Vja6rV9eIZZYXP/G3Xy4+nP21JQ3um/tmtJGt0X7ZnMpNXGG+AWry7DSshJaVrraF6SouJG8FL0Mvvi/heO720zr4+C+BbG9VvUEamDl283mQx1Yfjn6l5zS7xTcdvyey0TiTtURMTgyVmALeJwWI3pk8da1R+j0CXVASZlhfv5gCF8w/lypZH5cbEvYUxV/qKBSSikVT6g6EG+qFACZ7nRePu9BZn3xT37z9p/YUbmbo/pP5N4Tf0N+ev0FJ0L+EDOf/IKACQ/FC4YMIRP+N2ggFDIEQob1u+JX3mxvFqGqsoTj9Gz+8BPWv3wvOB2I7UAcDsTpID/7c5oKyfP7o9wc9mQl105sonfAhKBiW7N5UAcWEeHWSbcCMHvl7JjtycSdAsDyR2Kk9Wdw1yYWTFGqg0qZylRT03WT/XFRSimlktLMmhEDO/fh/pNvBuD7Xeu5Zu4fuOODB/jbqbfH7HvB5N7Yto0tgmWBLYIdWR227vk9b33LJ9/vjDnW3WMI4nBStfIz0occ3my2Ow8fTOcx4yAYwAR+eNhLFzZZERpRYHPKIAf3fOJjaH5KDVJRKUBEOKbXMby55k0qA/Ur/snEnYLw3/+9PxuV8MIsSnU0KVOZctgWLtuKu7JRMj8uwZAhJ03H4SqllEqcle4iZjWIRgzo1JszR5zIs4tfjz2P2+booV2bPce43rl8sXYX/gZztCx3OjmHn8fu9x5GLBtP3zGI5aBm3WJqNiytN0/Y7bAYMrQ36RP7Nzw97HkWln/bZB7unOZh7CMV3DC5kYUoxIKMLs1eizow5bhz4jYyRMedyp7wQ8iYurhTXX5W/2/G4/BoRUod0FKqOerIQfnE+75F/7hUrfysySBwXqfNgAKN3K6UUipxzi5piDv+yIfvd63nkYXPs6VsOwDFZdt4bcUHjO0+vP6OluAd0Smh9KaP74nVyA1m1oQzyD36Yko/e4FND5zHpodmUP7VXLwDY+cN/2RszzhnAIafAa6mfwsH5FmcNdzJXxc2MjTe4YZBJzZ5DnXgGpw3GDtOzM/ouFPlS8sxAYNvh4+NszbGxJ2yxGJq4dS2zLZSbS5leqYALj2yH59+v5MqX+wKRFkTzsBKz6X0sxfYOfc+xOWNCQLncVhcfHhfbEtbQJRSSiVOLCHjiJ6Uv7ce0yDWYborjcXFK3jsPy9SVltBljuDY/tP5r+OurL+OWwh4/BGKjcN9MxNY1zv3L1B4xvKGH4UGcOPajy/Em6AzM9spFdpyMkw55pm83HbVDfPLPXH35jTG7qPbvYc6sDksBycM+Qcnlr2FL5Q/Qp3vLhTWWOzKLysEMv5Qzu9y3Jx4fAL2zrrSrWplKpMje+dS+cMNxt2V8Xd3tyPiwHOntCrlXKnlFLqQJYxvgtl766Peb9bZj4PnX5ns8c7OntxdUt8lbJfHj2ArzfsiRsSpDkeh80V0+IM76tjO+HQX8CCB+rFmlp3Xf0FAAqzLWpuyWp4NDjT4PDrk86XOrCcOfhMnlr2VNxt0XGnGtM9ozvDOw1vch+lOrqUGuYnItz7s5F4nMlny+u0uf64gRq3QCmlVItYaU5yTumHtOA3SFwWuWcmF/h9Sv/OXHZkP7zO5BZWqvu9G9srt5kEroaMgvDcp2TYLigYGh4qqA5qBWkFXHLIJXgcnqSP9dge7pzSfCOEUh1dSlWmACb168R900clVaHyOm3OnlDIZUc20UqnlFJKNSNjUjcyp/VMqkIlLotOFw5PqleqzrXHDuTSJCpUXqfNNccM4NJEfu+8OTDjDUjrDJYzsQw5PJDbB85/GRzaOKng8lGXc0q/U5oM5NuQx/bwP0f8D6MLdJioOvCl1DC/OqeM7E5emotfPvc1tYEglbXxh0CkuWyMgRtPGMzFh/dt41wqpZQ6EGUd0xs710Pp3DWYoMHE+w0SEIeFnesm7+whuLq3bOEjEeH64wYxrncuf/1gFd9sLiUUMvijVhZ02oIlwpheOVxzzECm9O+ceAK5veGKT+H582DrNxD0gwnE7me7whOx+h8NP30cXBpUVYWJCLdNuo2eGT15eMnDiAjVgeq4+6Y50khzpHHftPsY12VcG+dUqfYhxjS+FOz48ePNl19+2YbZqS8YMsz7bjsPz1/N1xv24LQtRMIBfrvneLlian9OHd2dNFdK1gmVSkkissgYM76987Gv2rt8Ugc+EzTUfLubsnkb8W8qD1c2jIms2teZzCN64Oq5f4OQrttZybNfrGfl1nLKawNkehwM7ZrF+ZN6U5iXtm8n374CPvsbfBMJxGrZEAqGe63GzwjPscrtvc/XsC8OhPLpQC6bKv2VvLHmDZ5c9iTbq7bjsBwYYwiEAozpMoaLRlzElO5TsJIdWqpUimuqbErpylS0ytoAe6r9BIOGbK+TLK9D4xYo1QIHws0KpFb5pA58xhiML4RY0qI5VSkl4IPq3eCrBHcWeHPBTo1GyQOhfDoYyiZjDGW+Mspqy3BYDrLd2aQ597Gyr1QKa6psSo3SMwHpbgfp7g6TXaWUUgcQEWk0DlWH43BBZvOBhZVqjIiQ7c4m253d/M5KHeA6ePOaUkoppZRSSrUPrUwppZRSSimlVAtoZUoppZRSSimlWkArU0oppZRSSinVAlqZUkoppZRSSqkW0MqUUkoppZRSSrWAVqaUUkoppZRSqgW0MqWUUkoppZRSLaCVKaWUUkoppZRqATHGNL5RZAewvu2yo5RqA72NMfntnYl9peWTUgekDl8+admk1AGp0bKpycqUUkoppZRSSqn4dJifUkoppZRSSrWAVqaUUkoppZRSqgW0MqWUUkoppZRSLaCVKaWUUkoppZRqAa1MKaWUUkoppVQLaGVKKaWUUkoppVpAK1NKKaWUUkop1QJamVJKKaWUUkqpFtDKlFJKKaWUUkq1gFamlFJKKaWUUqoFtDKllFJKKaWUUi2glSmllFJKKaWUagGtTCmllFJKKaVUC2hlqoMTkWkisqmtj21hevNE5JK2Sk8p1X60bFJKdQQdqaxqCRHpJSIVImK3d14OVFqZaiDyB1f3CIlIddTr81ox3Rki8klrnX9fiUgfETEi4mjw/tMi8of2ypdSBwstm+LTskmp1KJlVeMiZdWAVk5jnYgcW/faGLPBGJNhjAm2ZroHM0fzuxxcjDEZdc9FZB1wiTHm/Yb7iYjDGBNoy7wppQ5eWjYppToCLavUwUZ7phJU15UrIjeJyFbgqXitINGtDiLiFpH7RGSDiGwTkYdFxNuCtGeKyAoRKReRNSJyWZx9bhaRnZEWifOi3t8veUgwnzNE5JNIeiUislZEftTIvt1EZKmI/Dryep6I3CUin0au810R6Ry1/6kiUiQieyL7Do28P1NE5kTtt0pEZke93igioyPPjYhcHtlnj4j8TUSkNT4LpdqKlk0J5VPLJqXamZZVMendISIvisg/IvkqEpHxUdt/KyKrI9uWi8hPGhz/i6hrWi4iY0XkGaAXMEfCPYG/kajeexE5S0S+bHCe60Xk9da81gOdVqaS0xXIA3oDlyaw/z3AIGA0MADoAdzWgnS3A6cAWcBM4H4RGdsgX50j578QeFREBiebBxGZJSKzWpC/aBOB7yL5uRd4ouFNgYj0BeYDDxpj/hS16VzC11cAuIAbI/sPAp4DrgPygTcJFxSuyHmOEBFLRLpHjpscOa4fkAEsjUrjFOBQYCRwJnDCPl6vUqlAy6bmadmkVPvTsqq+U4HngRzgdeDBqG2rgSOAbOBO4FkR6RZJZzpwB/DzyDWdCuwyxlwAbAB+HBnad2+D9OYAg0VkYNR75wL/TPZaVRRjjD4aeQDrgGMjz6cBPsATtX0G8EmDYwzhP0ABKoH+UdsmA2sbSSvmXE3k61Xg2qh8BYD0qO0vArc2l4fIsZsSTLNP5NocDd5/GvhD1DV8H7UtLXJM18jrecD/i3yu5zQ4zzzglqjXVwJvR57fCrwYtc0CNgPTIq83AmOBs4FHgYXAEMIF5usN/m8Ob/A5/ba9/870oY9kH1o21UtTyyZ96CNFH1pWxaRrgAGR53cA70dtGwZUN3HsYuC0yPN36vLf1GceeV2vjASeBW6LPB8IlEfKxKQ+b3388NA5U8nZYYypSXDffMJ/nIuiGj8FSHo1lchwlNsJtxZYkfN+E7VLiTGmMur1eqD7/swD4YIGwBn1vO61P+r11ronxpiqSLoZUdvPA74HXoqTxtao51VRx3UnfE115w2JyEbCLSYQbgGeRrjwnQ/sAaYSLgTmJ5iGUh2Zlk1aNinVERzMZVU8Db/3HonMJRORnwO/IlwZgnCZUDfEuJBwz1VL/BP4M/B7wr1Sr0bKxAJa91oPWDrMLzmmwetKwn94AIhI16htO4FqYLgxJifyyDZREzMTISJu4F/AfUAXY0wO4aEk0cNTckUkPep1L6B4f+UhYgvhG5M+Dd7vS9TNRALuiOTrn5L4Mp3FhIcEABAZmlNIuAUYfrhhOSLyfD7hG5apxN6wKHUg0rJJyyalOoKDuaxKJs+9gceAXwKdInleFpXnjUD/Rg5v+Bk39B6QL+E5m+fwwxC/drnWA4FWpvbNEmC4iIwWEQ/hH2Mg3EJJ+Itwf6S2j4j0EJGmxsGLiHiiH4TH2buBHUAg0rpyfJxj7xQRl4gcQXhc8OwW5iEuE15S81/A3SLSSUScInIO4W7pt5I4lR+YDqQD/xCRRP4GXwROFpFjRMQJ3ADUAgsi2+cDRwFeY8wm4GPgRKAT8HUSeVPqQKFlk5ZNSnUEB01ZlaR0wpWiHZE0ZwIjorY/DtwoIuMkbECkAgawDejX2ImNMX5gNvAnwvPX3ou8317X2uFpZWofGGNWEu4mfR9YBTSMb3AT4WEjn4tIWWS/wTRuCuFWgYaPawj/aJcQ7pJ9vcFxWyPbioH/Ay43xnybbB4kvGrLw03k70pgN+FJ09sJt5icbIzZ1sQxMYwxPuAMoAvwZHM3LcaY74DzgQcIt5z8mPDkSl9k+0qggvCNCsaYMmAN8KnRuArqIKRlk5ZNSnUEB2FZlRBjzHLCQ/E+I1w5OgT4NGr7bOBuwr1K5YTngOVFNv83cIuEVwa9sZEk/gkcS7jCGD08OtnPWwFiTHO9gUoppZRSSimlGtKeKaWUUkoppZRqAa1MKaWUUkoppVQLaGVKKaWUUkoppVpAK1NKKaWUUkop1QJamdpHIvK0iPwh8vwIEfmujdI1IjJgP59z77W05bFtRURuFpHH2zsfSrUnLbP2/dh9ISK9RKQiiVhWSh1wtBza92Pbit47Ne+gqEyJyDoRqY78gG2L/PHu9yBkxpiPjTHNLiEpIjNEpOHyn/uNiMwTkUta6/z7qrWvP5LGNBHZFP2eMeaPxpiU/VyUqqNlVmqKlCtGRG5K4ph1InJs3WtjzAZjTIYuja5SnZZDqUXvnVLXQVGZivhxJIrzWGA8cEvDHUTE0ea5Ukqp+LTMSj0XEo5n9fP2zohSbUTLIaWacTBVpgAwxmwG3iISSTrSyniViKwiHDAOETlFRBZHAp4tEJGRdceLyBgR+UpEykXkBcATta1ejV5ECkXkZRHZISK7RORBERkKPAxMjrT27Ins6xaR+0RkQ6QF6GER8Uad69ciskVEikXkopZev4jMFpGtIlIqIh+JyPAGu3QWkfci1zdffoiojYgMiWzbLSLficiZLc1HgzytE5EbRWRpJF8vSDhqOSKSKyJzI59hSeR5z6hj80TkqcjnUiIir4pIOuH/4+6Rz7hCRLqLyB0i8mzkuLdE5JcN8rFERM5ozWtVKllaZqVGmRUpV34GXAUMFJHxDbb/QkRWRPKxXETGisgzQC9gTuSz+42I9In8HzpE5CwR+bLBea4Xkdcjz5v8jJVqK1oOpUY51CBPeu+UIg66ypSIFAInAV9HvX06MBEYJiJjgCeBy4BOwCPA65EvrItwlOlnCEeang38tJF0bGAusB7oA/QAnjfGrAAuBz6LDPXIiRxyDzAIGA0MiOx/W+RcJwI3AscBAwlHrW6ptyLnKAC+IhzpO9p5wF1AZ2Bx3fbIl+w9wlGzC4CzgVkiMqyR698jIocnka8zgROBvsBIYEbkfQt4CuhN+KakGngw6rhngDRgeCRf9xtjKoEfAcWRzzjDGFPcIL3ngHOi8jssksYbyV6rUq1Jy6yUKbPOACoIf4bvEO6lqjt2OnAH4R6rLOBUYJcx5gJgA5HWfWPMvQ3OOQcYLCIDo947N5JnaOIzVqotaTmUMuVQQ3rvlAqMMQf8A1hH+EdwD+Ev6CzAG9lmgKOj9n0IuKvB8d8BU4EjgWJAorYtAP4QeT4N2BR5PhnYATji5GcG8EnUawEqgf5R700G1kaePwncE7VtUCTfAxq53nnAJQl8LjmR82RHXj9NuNCq254BBIFC4Czg4wbHPwLcHnXsHxL8/2h4/euA86Ne3ws83Mixo4GSyPNuQAjIjbPf3v+LqPfuAJ6NPM+MfOa9I6/vBp6MPG/yWvWhj9Z+aJnV6OfSLmVWZP/3gb9Enp8T+ayckdfvANc28X95bNTrPpFrcERePwvcFnk+ECgnfJPT5GesD3209kPLoUY/F7130nuneo+DaZzr6caY9xvZtjHqeW/gQhG5Ouo9F9Cd8Jdns4n8hUSsb+SchcB6Y0wggbzlE/7xXCQide8JULfaU3dgUQJpNinS4nM3MD2SZiiyqTNQGnm+97MwxlSIyO5I+r2BiXVd6xEOwq0b+8PWqOdVkTQRkTTgfsItL7mR7ZmRaykEdhtjSpJNzBhTLiJvEG45+R/CN0e/iGxu7WtVKhFaZqVImRVplT8K+F3krdeAR4GTCbe4FwKrkz1vxD+BPwO/J9wr9aoxpkpECmj6M1aqLWg5lCLlUCP03ikFHEyVqaZEf8E3AncbY+5uuJOITAV6iIhEFQq9iP8juhHoJSKOOIWCafB6J+Eu2OEmPC65oS2E//jr9Gr8Upp0LnAa4a7udUA2UEK48KmzNx0Jr9qTR7hFaSMw3xhzXAvTbqkbgMHARGPMVhEZTXiYgUTylCciOcaYPQ2Oa/gZx/MccLuIfER4/PaHkffb61qVSpSWWT9oizLrAsLDZuZE3bR5CA/1ezWSVv9Gjm2uLHoPyI+UbecA10feb+4zVqq9aTn0A713OojvnQ66OVMJeAy4XEQmSli6iJwsIpnAZ0AAuEZEnJEJdxMaOc9Cwl/keyLn8IjIYZFt24CekXHEGGNCkXTvj7RGIiI9ROSEyP4vAjNEZFikteH2BK7DEUmz7uEk3D1bC+wi3JrzxzjHnSQih0fydhfwuTFmI+ExzINE5ILItTtF5FAJTwptTZmEC8s9IpJH1LUbY7YQHsc8S8KTLZ0icmRk8zagk4hkN3HuNwm3pPweeCHy/wDtd61KtYSWWa1fZl0I3El4qEzd46eRtDsBjwM3isi4yP/BAPlhAvo2oF9jJzbG+AnPIfkT4Ruw9yLvN/cZK5VKtBzSe6eD9t5JK1MNGGO+JNxl+SDhlofviUzoM8b4CE9CnkF4edyz+P/s3Xl8VNX5+PHPvXf2bCSBBEgCSUBAQGSTTRaxSl0rat0X3Kjrt1qt3dvvr63VWv1W27rgirt1ReuCS1X2HZFVCJAESEII2ZfZ772/PwYCyUySmZAACc/79eKlzNy59wxwb85zznOeA++3cB4duJDQgsjdQNGB4wG+BjYDpYqilB947ZcHrrVCUZRaQvn5gw+caz7w+IHP7Tjw37Y8TehGOvhrLvAKoWnuYmALsCLC594gdNNVAmOAaw+0oQ6YQWhqt4TQ1PLDgD3SxZVQFZgpUbSzLY8DTkIjUCuAz5q9fx0QALYCZcA9B9q7ldDoSb4SWtDZt/mJTdP0Efr7O4tDC75j/q5CHEvyzOrcZ5aiKBMIdRyeNE2z9LBf/znw3a4yTfMdQmlAbxBa8/QBocAI4CHgdweeQz9v4bu/Qeg59E6z0fgW/4yFOJ7Ic0j6Tidy30lpmsIqhBBCCCGEECIaMjMlhBBCCCGEEO0gwZQQQgghhBBCtIMEU0IIIYQQQgjRDhJMCSGEEEIIIUQ7SDAlhBBCCCGEEO3Q6qa9PXv2NLOzs49SU4QQR8PatWvLTdPsdazbcaTk+SRE99Mdnk/ybBKi+2nt2dRqMJWdnc2aNWs6p1VCiGNCUZRdx7oNHUGeT0J0P93h+STPJiG6n9aeTZLmJ4QQQgghhBDtIMGUEEIIIYQQQrSDBFNCCCGEEEII0Q4STAkhhBBCCCFEO0gwJYQQQgghhBDtIMGUEEIIIYQQQrSDBFNCCCGEEEII0Q4STAkhhBBCCCFEO7S6aa9omT9oUOsNYJqQ6LRgt2jHuklCCNFlGYZJnTeIJ6AT77AQZ9NQFOVYN0sI0Y0FdIMaTwDDNElyWqUvJ9ql2wZTm0tqeP/bIoqqPPiCBikuG5MG9uSCEX1wWNt3s+iGycK8MuYszGdtYRUWLfSDPqAbjMjswW3TBnDWyWlYNJnwE0KIaOwoq+fFJQW8v64I3TDRVIWgbpLksnLz6TlcOa4fKXG2Y91MIUQ3oRsmi/L2M2fhTtY068udkpHE7WcM4KyT06UvJ6KmmKbZ4ptjx44116xZcxSbc2QMw+Q/60t4asEOdle68QcNjMO+XpxNwwQuH5vFrdNy6ZPkjPrcC7aVcd/b6/EGdBr8esRj4uwaFlXloUuGc94pfY/w2wjRORRFWWua5thj3Y4j1dWeT6Kpsjovd7z2LZuKawgaJkEj/GeRw6JiAhePyuDPM4djlc5Nt9cdnk/ybDp+Lczbz71vf4fX33Zf7sGLh3P+COnLiZDWnk3dZmbKF9S56411LNlRjqeFG+TgjfP6yl28920Rr9w0jlH9kts89ztr9vD7DzfhDRitHtfg0wGde99eT2mNj5sm58T8PYQQorvbXeHmF0uU0wAAIABJREFU4qeWUuMJRAyiDvIGQ8/cD74rZkdZPa/dMr7dmQVCiBPbu2uL+N0HG6Puy933znr21ni5ZUru0Wmg6LK6xTCfbpjMfmUNi/P2txhIHS6gh3Lzr3l+JZtLalo9dsG2sqgCqcN5AwZ/+3wrH60vifozQghxIqhq8HPZM8uocvtbDaQO5w0YbCqu4fbX12JE+RkhRPdmGCYev05rGVYHLdhWFlUgdThvwODRL7bx0friI2mmOAF0i5mpf329ndUFVY2jmIdr2LKA2tUfEKgoQrU5sablkjTpchyZw3D7da59fiXLf/2DiKOdumFy79vrI958rZ0XQjfhr97fwIxh6bKgUQghDvjn19upbPATKSZq7bnqDRqszK9kYd5+pg9JO/oNF0Icc/vrfLy5ajevLt9FeYMPFQUDk75JTm6ZnMOlYzNJdFibfEY3zAPLNNrbl9vI2UN7y6y4aFGXD6b8QYMXlhTgCYTPSNWumkfNyndJnXEnjpzRKJoFT8FaPNtXNt4o/qDBJxv2cumYzLDPf7O1DF+wfecFwIRPN+7l4lHh5xZCiBONN6Dz1uo9BPTwSCqa56rbrzNn4U6mD0ljT90e3s17l53VO6kP1JNgTWBo6lAuHXQpaS4JtoToThp8QX753ga+2LIPBfAdGDzXCT1Liqs9/O2Lbfz1s61cMTaL3184tHGN5cK8Mrzt7CMCjX25S0ZLX05E1uWDqc83l0ZM+zB8DVQveZ3U8+7BNXhS4+uugeNxDRzf+PuGAz+cIwVTcxbtPJA7G/t5D5776QU7JZgSQgjg4w17iVTsPJbn6vqKVVz10VPk1WzBMA2CRrDxvWUly3h+4/OM6zOOW0fcysi0kZ31VYQ4bnn8OlVuP7phkui0kuiwdOltBiob/Px4zjKKqzz4I2QgHXRwmcc7a/ewpbSW124OrbF8esHOsGIT7enLSTAlWtLlg6kXlhRErMjiK96KGfTjGjSxzXMUVXnYWlrLkN6Jhz4f1Pl2V/URnRegsNxNRb2P1Hh7VMcLIUR39cG64iN4XpvYen6JJXUxmyoDEY/wG34AlhQvYU3pGu4dey9XDbmqI5ouxHHNMEyW7CjnmYU7WVlQiVVTUZRQue8+SU5unZbLzJEZxNm7VrfPG9C55vkV7Kl0R5zRjsQTMNhYVMNtr65lzrVjOqQvt6vCTXm9j57SlxMRdK27KoJdFQ0RX9c9taiuRBS17RxXDYNt326lf18bZjCIGQyyv86PVTHxNbt3YzkvgNWiUOUOSDAlhDjhldf7Ir4ezXPV1vNLbKmLUdTIgVRzXt3L39f8HU3RuHzw5e1qrxBdwdpdldzx+rfUe4ONgxVB49Cgxe5KN3/55Hse+Ph77jnrJH4yNbfLzFS9tKyQgvKGiIFUa+udfEGDVYWVfPBdMRZNQQ82/XysfTm/bnDmows475Q+3Dw5h5PSEzrk+4nuocsHUy1VZtGciRjuWkxDb/Nm0b0+St5+l1L3LhSLBcViodIWD6k/BKXpZ2M570HRVJoRQogTVVvPVS1ue8RAqmpxFeWfl+Mv86M5NBLHJJL+43S0uNA5vLqXR1Y/woheIxiSMuSofBchjqavvt/HnW9822aVOveBIOvx/25nV6Wbv8wcHnVAFTACLNizgJc2vURhbSE+3Ydds9Mnrg+zhs1iRvYM7FrHDxgbhskLiwsifrdo11i+tKwwYmpxe/pytd4g76zdwwfrijkpPYGHLx3B0L6JbX9QdHtdPphyWNWIxSfsGUNQLFbcecuJGzK51XNoLhc5P7+XnOG9G19LD+gE/vAZNIuDYjkvhMqwJ7msbR4nhBDdXbLLFvH1tp6rttSvwwKp8vnl7J+/n8xbMokfGk+gKkDJqyUUPlpIzm9zUC2hxed+w8/cTXN5eOrDHf+FhDiGvttTzV1vrIup3LcnoDPv22LSE+zcfdagVo81TIO5m+bywqYXMAyDhuChTCCf7qPWX8sDKx7ggRUPcOWQK7lr1F1Y1dj7Ow2+IB98V8zivHIq3X5smkpGspOBveJp8IXPRMey3il/Xx1B04RmIVWsfbmDdAN0w2BjcQ0/nrOM564fy+kDe8b2hUW30+WDqeyecVTtDs+HVe1x9Jh8DZVfzkFRNRw5o1BUC97C7/Du3kDy9Jsaj9UNk9xecU0+77BqDO2TyKaS2nafF6BPooNekuInhBCcP6I364uqG0fJD2r1uVq8ivjZTffs0z06ZR+UkXFzBgkjQuk2tl42su7IIu/+PGqW1ZA8NbQhu2Ea/HfXf6nx1ZBkTzo6X1SIo+AX766POJgMrafAeQI6Ty7YyZXj+pGe6Ij4+YAR4P6F97OsZBmeoKfFNriDbgDe+P4N1u9fz9NnPY3T4mx83zRNlu+s4IUlBWwvq8cT0HFaNQb2iudHp/ZlRUEFH3xXjKooTZ4LigKqoqBHKDAWy3on3TBI1XT2G02DvFj7chG/u19n9itrePvWiQzPkGfLiazLB1Ozp+Ry/7vrw6ruASSOuwQ1Lpma5W9R/vGjKDYn9vSBJE68oslx/VJdDIqQ/3rbGQP45bsbwhZMR3tel03j1mldJzdZCCE608xRmTzwyfcR32vpuZp6Xj+apwi4t7sxAgaJY5qm2GgOjYQRCdRvrm8MpgBURWV+wXyuHHJlh38nIY6FTcU17KmMHOREW/L7tRW7uG/G4LDPB/QAv1z8S5YUL8GnR17n2JxX97KpfBN3f303T5/1NKqi8u9Ve/jH19up9QRCm+sedvzuSjdfbytr8XymCXoLSyRiWe+kKypDctNw76kK6ydG05drax8qt1/n9tfWsugX06WvdwLr8sHU2UPT0Vr5Bxw/bDrxw6a3+H6cTeP2aQMivjdjaG9+rW5s13kh9DCYOSqj1WOEEOJEEW+3cNHIDN5bW0QwwohzpOeqI/MlFLVpJ0iv17HEW1C08Ge/JcmCZ1fTTqZX91JQU9AB30CI48Nzi/MjlgmPNgXOHzR4ZfkufvqDkxr3Y8qryuPVza/yccHHTbYcOKitNYo+3cd3Zd/xXt48Vm4YyKcbS1ucOTsSsa53SnZZUVpYut5aXy7aoLSiwc+aXVWclp3Sru8jur4uH0xZNZWfTM3lyW92xnzTKoDTpnHuKb0jvm+zqDx48XDuf3dDTDnJAE6rxm/PPxmX7dAfsWGYuA9McWuqjGAIIU489549iC82l1Lljq4qn6KGj4xr8RrB+iCmboYFVMGaIJb48B9tNf6a9jVYiOPQV9+XRZy5iSUFLmgYbCmppXeKj3u+uYcd1TsIGAF0M7wvFe0aRY/u4f9WzaF2+714ArEV32prFuigWNc7fbKptNVB90hiWZfl8es8s3CnBFMnsC4dTNV5A5TX+znz5DRW5FeyZldlTEGP06bxxuwJ2C0tj2xceGoGJdVeHvtvXtTntut+rh/Vh2sn9KfWG+C9NUU8tySfvTVetAM5wD1cVq6d0J9rJ/RvMWdZCCG6m/REB2/MnsDlzyynwRckwgRVE6YR/nx0DXShWBRq19aSNO7QWgXdq1O3oY70H6eHfSbZnhz2mhBdkWmauP3hM0cQWwqcqihsq9jJXUvvoc5fFzGICp0z+jWKAA16FT5tFwT6NTlPa8FStLNAEPt6J90w0ZtXE2tDLEGpCSzYtp+gbmA5MMsnTixdLpgyTZNlOyt4ZuFOludXYNVUVEUhqBsEDROLqkRMHzmcTVNx2jRev2V8xLVSzd06bQBpCXZ++8EmgLDF0we5bBqGafKzTIMzXnqA36t/4u31+1AVpXHWLHhgJKnKHeDZRfk8syifMwb14tHLTyXRIVX/hBDd38l9Evn4fyZz49zVlNZ68QR0Ii2PsGoKii8DJWE7Joc6j5pLI21mGiWvlaA61CYj5dYUKz0m9WhyHqfFyUnJJ3X21xLimIspBU6t4/Etf6E+UIPZSrAR6xpFlADWHmvwlR4KploLlmy9sqOeBToo2rXr7RXrPlSaqlDrDZISF7liqejeulQwtbW0lpteWk21O9AY0AT0poGNBRPFBFVVQAmVsTwozq6hKgrXTejPDadnk5YQ/YzQxaMzOfeUPny0voQ5C3eyp9KD7cC0dkA3SEuwc9u0AcwclYGmKlxZEuT71XvwqS3/EfsO5DsvyNvPBf9cwnu3T6JXglT+E0J0f/1T4/jqvmms2VXF0wt2sjhvPwYmCkooiFIULh+bxY/G3M3sbxbRfAyr13m90OI0St8qxV/mR3WqJI5OJOvWLFRr09Fh0zQ5J/uco/jthOg8iqLgslmo94XPTsWSAmf2+AJ3sC4skGq+NsqeYUeL06Jeo6goJqr1UJXltlLmPPlro54FOlw0a9fbK9Z1WYoS6guKE1OXCabWFFZy/YurWpwVOihIKIiyaAopcXZOy04hoBukxtmYMCCVGUN7NwZBsXJYNS4bm8VlY7Moq/NS7Q5gmCbJLhtpCXYURcEwTG55ZQ3fq4n41OhuLH/QoKTaw1XPreA/d53eZJ2VEEJ0V3n76nlnzR6W7SzHZlHxH+iMBA2TibkpzBiWzqi+qYzrPY4lxUvCPp8yLYWUaa2vU9AUjQtyL8BldXXKdxDiWJg+pBefbNgbliYbdQqc4oeEtehm04As0tqo3U/tRq/T0X06mr1pYNHSGkWUQ+dtK2Uu1lmgWEW7FutwMe8pGjRJckp20YmqS/Ta8/fXc8Pc1W0GUofzBU2q3H78QYOnrx3d4SUr0xIcEWe2vvx+HyvyKxpnnQ7X2g0dNEz2VLp5blEBd58l6ShCiO6r1hvgtlfX8u3uKgJBAz1ChtGi7eWs2VVFz3g7v754FqtLV0ddpvlwVtXKrGGzOqDVQhw/fjJlAP/dUhax8FY0KXD2HuvD+kUtrY3q99N+5P0sj9I3S8m44VCF4tbWKJq667Dzth4sxToLFItY1mIdLtZ1Wf1TXTisnRMMiuNflwimHvj4expaWGzZWoDiDRgs2r6fb3dXMab/0amyMmfBzohBXzQ3tC9o8NKyAu46c6BU+xNCdEtVDX4uenIppbXeiKWdD+f26xRVubnvVR+OlPOxpHxC0Iw+oHJoDh6a8hDZSdlH2Gohji+nZCaRkexkR1l9xPfbSoGzJK4Dpem91NLaKFuyDUd/B9VLqkkcndjmGkVTt6E3HBoUbitYinUWCNMM5dW1IZaKfKEyEk3PGcueoredEXmLHXFiOO6DqbJaL0t2lkdcnBxNgOIJ6DyzKJ9nr+v8YCp/fz1b9taGvR7LDe3XDb7ZWsZZQ8NHeoQQoivzBw2ueX4le2s8BCJNR0VgmAeK/lRO5N4zBvHMpsfxG34Ms+VATFM0bJqNhyY/xA/6/6Cjmi/EceXhS0dw5bPLo76XDme3e2i+OUFr+7fFD4vHCBhRrVFEMQnUjDp0rTaCpVhngeI0CBgGVkMHQ8erWtAjrE+PpSJf80Dq0PeObl3WhSP6RnEN0V0d98HUayt3RfwnHm2AYpqwcNt+yut99Izv3OIOX2zZhxGhkmAsN3SDT+f9dUUSTAkholLvr+ejnR/xWeFnVPtCi75THCmcn3s+5+Wcd1ytFfpkYwmFFQ0RO39trWvwBw3y80fw6nmvMnfTXL7a/RUKCl7d23gOp8WJaZpckHsBs4bNkhkp0a31cFmJONJMdPeT0mySqK3925xZTrJuz2q1TaapEqgeDeahqnbRBEvRzgLF2TT+9uNTGZ+bQnGVhwZ/kIc+3szGveEzdDFX5FOImHLcGqdV42dnn4TTJil+J7LjPpj6dENpxPVHsQQoFk1h6Y5yLhqZ0eaxR6Ks1ksgQjAV6w1dVhv7ugAhxIllX8M+nlj3BPML56Oi4tEPVdTKr8lnc8VmHl71MBcOuJA7R95JqjP1GLY25OkjSIMOGibz1hXzuwvO4uGpD1Pjq+HTgk8prCmkxl9DD3sPBiUP4pzsc46rAFKIzvLC4gIizc9Gcz/pQSeWZuPL7dm/LYxhw19xRtjL0QRL0cwCqarCjGHpWDW1cYA8zmkHwoOpWNdi9enhpKLeH3EdWiROq8alozO5ZXJuVMeL7uu4D6Zqvc0nokNiCVB03aTGE/k8Haml/a1ivaEj7WouhBAHbavcxs2f30x9oL7FjTY9wVBwNW/7PL7Z/Q0vnvMiOUk5R7OZTWwqrmF3pTvs9VjSoBUFPlhXzNXj+5NkT+KqIVd1eruFOB65/UHmrSumeTXuaO+nYO2paM5iFPVQ3yjW/duac1qcjE/4FZ+ZDgIR9q060lLmTqvGrVNzsTbbGLdXQuS9nWJZi6UAM0dm0OAP8sbK3Rim2WL6pM2iogC3nzGA/zlzYIcXOBNdz3EfTLUk5j0AjkKbesXbURXCSpXGurgy2SWbvgkhIttTu4cbPruB+kDkhefNBc0gFd4Krp9/Pe9c+A6943p3cgsjW7KjnGCEzkksWQZuv878TaVcPb5/ZzRRiC7j042lEWswRHs/BWpHY0//NOz1WPZvO8hlcWFVrTwz4xl6aLl8vX5h2B6gR8phVZlyUk/uOGNg2HvnDe/D11vLaPA1vWYsa7EcVo2zhqYzMqsHN52ew9ylBby1eg+KooT6j0ooo1JV4LqJ/bluQja9k6Lfq1R0b8d9MNXDZaWsLjztLZYAxaKp9DgKAcrpJ/XkqYU78fjbf0O7bBrnDD82nR0hxPHNNE3u/OpO3IHwGZ7mG20mjkkk/cfpaHEaJiZ1/jp+9s3PePOCN49By6Gi3hdx9j7WNOgqt7+jmyZEl7OroiFiymzU95NhJ1B7Krakb0FpOr0Vzf5tADbVRmZCJjefcjMz+s/AYQkFF3NvGMesF1dFnS7XFpdNY8bQdB657FTUCJWOzxqajtbC7FC0a7F6Jzk4NTOU2piV4uIPFw7jF+cMYWNxDdXuAKoCPVw2RmQmhc2MCXHcB1M/GtmXf321I2zdVCwBSkA3mHJSz05v66isHqQl2NlVEd7RifaGNk2pCiOEiGz9/vWUuksxmq2UiLTRZsmrJRQ+WkjOb3NQLSq6qbOjegfbq7ZzUvLR38vOokbugMSaZaBKSo0QLS6BiOV+8u+fgSPxewylIaZra4pGTmIOf53yVwanDg57f1xOCq/PHs+sF1cR1M1WgyqLCroBdquKN3DouWZVFVRV4dSsHtw2LZfpg9NaTKezairXT8zmucX5EdfYt5Ve6LRq3D5tQNj5HVaN07KPzrY6oms77oOpq07rxz+/2hHxvWgCFFWBGUPTj8rMlKIo3DZtAH/6aEvEh0eb+z6oCpeNzZSqMEKIiF7a/BLeoLfJay1ttJl1RxZ59+dRs6yG5KnJAASMAK9ueZU/nf6no9721HgbNosatrdUrGnQnV2VVYiuoIfTGvH1WO4nM5hITvDnlMX9nYZAQ9ggTcTza3ayE7N55dxXWi30MrpfMkt/dSbvrSni2cX51HoCGCbohoGmqqgKxDsszJ6Sy6WjM/muqJpV+ZWU1/uwW1T69HBw4YgM+qVGV0xm9tRc3v+2iNJab9hSi9ZYNYWcnnFcNEoGsUX7HffBVGq8nTMH9wqVHY9wg7QVoNgtGrdMOXqVVi4elcELSwooLG9osSBFSxIcFu6cHp4PLIQQ7oCbRUWLMJst7G5po03NoZEwIoH6zfWNwZRu6nxa8Cl/mPgHLBH2ZelMZw9N55HPt4W9HkuWQZxd45LRnVuVVYiu4KT0BOLs2hGtE7JqCmP7DmXW1Le486s7KW0oxRv0hj1jAKyqFUVRmJo5lQcnP9iY0teaRIeVGyfncMPp2awurKKgvJ46b5B4u4XsnnGMz0lpnA2aPjiN6YPT2v3nkeS08tatE5n51FJq3IGo+l82i0qfJAev3zIeu0UGsUX7HffBFMDvLhjKsvwKaj3BmD7ntGqcP6IPp2a1XoGmIzmsGm/MHs+PnlhKRb0vqs30VAXibBbemD2B9ERZ0CiECFfhrcCqWgkYTdN7Wtto05JkwbPL0+Q10wytn0p2JHdqe5vrnxrHKRlJrNlVFfZetGnQqqIwY6isKRXi7KHp/Oq9I1snpCoK103oT1ZCHB9e9CHr96/npc0vsahoERbVgqqo6IaOpmpcNugyrhpyFX3jY5/BURSFcTkpjMvp3JS5rBQX8++ewuyX17BtXx0B3USPEFQdTCGckJvKE1ePIsEReZZPiGh1iWAqM9nFG7dM4KpnV9DgD0Y1hevUYOKAVP56ySmd38Bm0hIczP/pFG56eTVb99bhC+ottjnOppESZ+OVm8eT0zPu6DZUCNFleIPeiGsG2tpo0xLf9DGvqVpYquDRctu0Afz03+siLpxvK8vAZlG5bkJ/bBZZ/C2E3aJxzfh+vLi0IOKgbTRlyEdkJtE/NdTvUBSFkWkjeTztcdwBN5XeSrxBL/G2eFKdqVjVrhFwpCU4+PCuyWwpqeWFJfl8vGEvqqKgqqAbJqqicNnYLG6clE229LlEB+kSwRTA8IwkPvqfydz62lp2V7hbDFCcVg1D17mg9Dse+t19WI5R1ZXkOBvz7jidjUU1PLc4n883l2LVVMyGehRXHH7DZHxOCrdNG8CkAamyT4EQolUJtoSIe0rFutFm0AgSb4vv9PZGcuaQNCYOSGXp9nK8ERaKt0RTFXonOrj9jAGd2DohupZZk7J5dcWudpUhd1hV7jlrUMT3XFZXl9/4emjfRP7v8pH85eJTqGzw0+ALEu+wkBpnlwEZ0eG6TDAFkN0zjs/vmcrGohqeX5zP/E2hfRZURSGgG/SMt/OTqblcOjqDul9+RuXTT5N278+OaZtPyUzin1eNotYboKjSw9bZt5P70J/Jys0gVRZSCyGilOpMRVPC8/pj3Wgz3hpPvPXYBFOqqvDk1aO5/oVVbCyuxhNoO6Cyagqp8XbeunWCpOMIcZi+PZzMuXYMP3l1TZNKeG1xWjXuPuskTh/Y+VWOjzWHVaNvD+exbobo5rpUMHXQKZlJ/OOqUfyfblDrDeIN6CQ6rcTZtMYZnrg//IH8mReTcPbZOE8ZDoRKiW4rraPGHcBuVUlLcDAoPf6ozAolOqwM7WvFWl9Mdi8nVgmkhBAxsKpWLht0Ga9//3rYuqloN9q0a3auG3rdMZ0Jd1g1Xp89nnv+/R2fbNyLw6JGnKWyagrqgbUW/7xyFMlxspm5EM1NHdSLZ64by80vrUY3IpWOOERVQumyP//hIG6efPQKcwnR3XXJYOogi6aS0sIPWEuvXqT/6pfs/c1vaHjiRV5Yvof5m0qxaSoc6EcEDZPUOBu3Tctl5qhM4u2d88dR1eDnnbV7+HZXNXuHXkHKhzsYkFnFleOyGJiW0CnXFEJ0P1cNuYo3tr4R8b1oNto0Mbl00KWd0bSYWFSFao+f35w7BFVVeG5xPmV1PqyqStAwcNksXDUui+snZpOV0rXTjYToKIZhEjTMsDS1Hk4rCQ4rPxyWzn/WlwA0WZfosKqYZmhz29umDuCUzCSEEB2nSwdTbbHMOIdfLyhn/ZzlBNDQTTNsQ7civ4cHP93KXz7ZyhNXj+IHJ6e3cLbYbS2t5Ymvd/Dlln0oCqFp+KR+UFjL4t11vLZyF4PSE7hz+kBmDE2XdVNCiFb1je/LtMxpLCpahE/3xfRZh+bg/NzzSXEc+00oF+TtZ2+Nlxsn52DVVG6ZkktAN2jwBXFYNRxWKVMsBED+/npeXFrAB+tKaPAFQQkNRpySkcRt0wYweWBPfvneBv7fj4Zy0cgM/t+PhvHpxr1s2VtLVYOfBIeVnJ5x/OjUvjK7K0QnUUyz5UnhsWPHmmvWrDmKzek4bn+Qi59aRuH+enxRlCeH0OjNgxefwiWjM4/4+vM37uXet9e3WsnvIKdN4+KRGfx55nA0VQIq0bkURVlrmubYY92OI9WVn09Hwhv0ct2n15Ffm49f90f1GbtmZ2jqUF744QvHvCpXUDc49x+L+cU5Qzh7aMcNXonuoTs8nzri2bS7ws09b61jS0ktwQMzUs3F2TV0wyQ7xcX8e6bKgKwQnai1Z1O3LGlimiazX15DYXlD1IEUhGaOfjNvIyvzK47o+l9u2cfP3v4OT6DtQArA49eZt66Y38zbSGvBrRBCOCwOXj73ZU7tdSouS9spcE7Fwvg+4/nX9DmsKajl0417+WTDXpbuKMftj23vvo7wztoiUuJsnHVy+zfoFKI721RcwwX/Wsx3e6rxBo0WN6Bt8Ol4AwYFFe6IG2ILIY6Obpnmt7qwinV7qsNS+gAatiygdvUHBCqKUG1OrGm5JE26HEfmMCAUUP3xoy18eveUdl27pNrDT99cF7GyTmvX9gR0/vNdCRNzU5k5KqNd1xZCnBhcVhfPnf0ci4sX8+KmF9lSsQWgcabKptkwMRmZPJgLtm5lY/wsJj246LCRaxMFhaBhcsnoDG6anMOAXp1f4a/eF+SxL/N4ftZYGUUXIoI9lW6ufm4Ftd7oBzp8QYO5SwtJibNxyxQpLCHE0dYtg6lnF+3EEwjfd6F21TxqVr5L6ow7ceSMRtEseArW4tm+sjGYAsgvr2draS1DeifGfO1Xlu9CN8IDqWiu7QnoPP7fPC4a2Vc6GkKIVmmqxhlZZ3BG1hnsrt3N0pKl1PhqAEi2JzM5YzJvLqvnVxXbMMr3ETAjJyK8tXoP764t4pLRmTzQyanGzy7cyaQBqYzI7NH2wUKcgH79/kbqfZEDqbYGZB/5fBsXntqX9ETHUW61ECe2bhdM7a/zsWh7Oc2z5QxfA9VLXif1vHtwDZ7U+Lpr4HhcA8c3OTagmzy/uIBHLzs1pmv7gwavr9yFv1lqYSzX3lfrY31RDSOzpLMhhIhOv8R+9Evs1/h70zT51Xsb+c/6EnxG68UcDq7H+GBdMfvrvDx73VjUTgioSmu8vLJiF5/8tH2z/kJ0dyXVHlYXVkZcHhDtYPBrK3Zx34zBR7HVQohuF0yt212FTVPxN0vx8xVvxQz6cQ2a2OZTwpJGAAAgAElEQVQ5dMNkyfbymK/99dYyjAhPwViu7QvqzF1awD+uHBXz9YUQAmDOwp38Z31JxBn6lngCOkt3VPCXT7/n9xcMjfmadd4A760tYt66Yiob/JhAktPKucN7c9W4fvzfF9u4alw/MmQDTSEienX5rrCBYIh+QNYXNHhl+S5++oOTsGrdckm8EMelbhdM1XgCGBGeRrqnFtWViKJGV3K3rs5N+XPPoVisKBYLitWCYrGAxYJitYZeP/Dawdd37PDhi9B5ieXahgk7yuqjaqMQQjTn9gf551c7IgZSba0Z9QR0Xluxi9vPGEDPKDcWL6vz8rfPtvHxhhIURcFz2P42RVUedu6v5x9fbUdB4f07JrVyJiFObPPWFeHXw5cJxDoYvG53NeNyjv0WCEKcKLpdMGXVVCIlqGjORAx3LaahRxXUWDDRq6shGMQMBDGDB38FMAOBiK+XxQ0l2GMENFvvFOu1G45BhS0hRPfw0fqS5o8gIPo0IYB/r9rNXWee1Oa1dpTVc8Uzy6nxBFqsOHawGI+CyZXPruDlm8Yxpn9y7F9MiG6upaITsQ4GVzbEtgedEOLIdLtgKjXeFrF4gz1jCIrFijtvOXFDJrd5nuSUBNLvvz+ma2ctysf6+VYCzdZMxXrteHu3+2sRQhwlTy/YidvfdFYqlnWbvqDBi0sLuf2Mga0Wo9hb4+GyOcuodgeIZkMHk1A1v+teWMm8O05ncO+EWL6WEN2e3sKARGwDspH3pBJCdJ5ul1Tb0tS2ao+jx+RrqPxyDu685RgBL6YexLNzDVXfvNjkWIdV5YrTsmK+9oC0OGwR8pRjubamKgztE3sVQSGE8AZ0dle6w16PJU0IwOMPUlrrbfWYO17/llpvMCyQatiygL0v38Puv/+YoieuY9/b/4u3aHPj+26/zg1zV0VcXyrEiSyuhYHUwwdk26KgkOQ8thtzC3Gi6XZTIHaLxtXj+vHSsoKwqnqJ4y5BjUumZvlblH/8KIrNiT19IIkTr2hynGnCFaf1I1ZTT+qF1aKCP3ytQrTXtmoKN0zKifnaQghR6wlgs6hh+9zFmiakqSo17kCLxSJ2lNXx/d7asJH0aFMJaz0Blu4sZ8pJvWL8hkJ0X5MH9uTjDSVh1fwOH5BVVA1HzigU1YK38Du8uzeQPP2mxmP9usGpUg1YiKOq2wVTANdP6s/LywshQvJJ/LDpxA+b3uJnLarCD05OIyXOFvN1LZrKjZOyeWrBzogbBrd1bYDs1DiG9pWZKSFE7KyaGrGscqzrNn1Bna+37qPBH6R/qote8fYm6dMvLCkk0GyhfCyphA1+nTkLd0owJcRhZk/J5cst+yIWj4lmQFZTFS4Y0YdEh8xMCXE0dctgKjPZxX0zBvHYl9tjKg2sAD1cVv74o+HtvvbV4/vz7OJ8Wthzr1VOq8q9Zw9q97WFECe2RKc1YvpcrOs2dcPk+721fL21jF0VbjwBnf6pcWSnuuiX4uKdNXtoXnQs1lTC1QVVVLv99HDFPnAlRHd0SmYSGcnOFiv6tjUga9NUbp6c21nNE0K0oFsGUwA/mTqAaneAuUsLowqoNBV6OG28fetEeiVEVxI4kl4Jdl6YdRo3zF0VlmrTGocR4FL/Xs4efHa7ry2EOLFpqsJZJ6fx+ZZ9TfariSVNCGBUv2SevGZM4+9rvQF2V7gprGggb18degdsP2G1KJTV+SSYEuIwf75oODe+FFv/AUJrvc88OU0yW4Q4BrpdAYrD/eKcITwwczipcVacwcilQm0WFbtFZfLAnsy/ewq5veKP+LoTclN5cdZpuGwaVq3lalgQmg1zWjV+Mv0kZu9fQ/F9P8f0+4+4DUKIE9PsqQNwWMIDmsRxl5B85s3ULH+Lon9dQ9HTN1D37cc4T2o6kxRn07ht2oCmn3VYGZ6RxAUj+nLdhOyIhXYOTyWMhoqCN4bMASFOBBMHpPLgzFNwWKPvnjmsKqdkJPHY5SM7sWVCiJZ025mpgy4dk8m0Hcv4eslm3jvpHDYW1+AJ6FjUUMWby8dmcd3E/vRJirzQur0mDezJFz+bygtLCnhr9R4UQusEDrJbVExgysCe3HbGAE7LTsE48ymK7/kZRXffQ8bjj6Ha2z9DJoQ4MY3u14P0RDuFFeFV/aJZt2m3apw5JK3F9xMclrD1UhB7KqFhmiTI2g4hwlwyJpMkl5X/eXMdQNhWBwdZVQVVVfjh0N48ctmp2CzdenxciONW9wqm9CBU7wJvNWg2iE/HjOtF7Ztvct5993H5lNAIrGmaEfei6miZyS7+98Jh/PKcIXy6cS+bimuobPATZ7fQP9XFzFEZpCU4Go9X7XYy//kPiu//BUW330Hmk0+gOjs2yBNCdG+KovD0tWO49OllLXbCWuKwqjxz3ZhW95dyWDV6JzkoqW5aOj3WVEKAvj0cYa8JIeAHJ6ez9ndn89H6Ep5auIPSGi/egEGCw4JphgYjLhubxQ2TssnpGXesmyvECa17BFO1e2HNC7Dq2VBApWqACUE/RkIOTqdO3ITTGg8/GoHU4RxWjUtGZ3LJ6Mw2j1WsVjIefYS9v/0de2b/hMw5c9Di5UEphIjeyX0SefSyU7nz9W+j2lAXQunGT1w9itOyI+/Vd7ifTMnl4c+2ha1HjWULiCtOy8IeIR1RCBHitGlcfloWl43NZPH2/dz/7gYeu3wkiU4rA9PicVjl/hHieNC1gylDh0/vh+9eC1VB18PXRWnVW0kfYkf5+2C4/FXImXL02xkjxWKhz0MPUvrHP7H75pvo99xzaInhi0o9fp2tpbXUeAJYVJWeCTYGpycc9WBRCHF88QZ0nlqwg1sm57BuTzUbi2swTJNAs733NDVUTn1gr3j+cvEpUe9Pc8mYTB6avzXie9GkEqqK7KcnRLQURcFhtZCZ7GLSwJ7HujlCiGa6bjClB+GNy2H3cmihuMRBqukDjw9evwwufQ5OvvAoNbL9FFWl9//7X/Y99BC7briBfi+8gCU5GYCC8gbmLi3gnTVFaKrCwdhJN0ySnFZ+MiWXS8dmyl4TQnRHVYVQtQv8DWCPh9SBkNi38W3TNPnVexvI7RnPb84/GUVR2FXRwNylhXyycS/13iAmJnE2C2cNTefmyTkMSk+IqQmJDivXT+zPayt24Ym16phF5cwhafRLdcX0OSFOKHoQ8j6DlXOgqpBTvW6e1m3w9niYeBdkjgUZOBXiuKCYEUrcHjR27FhzzZo1R7E5MfjgTtj8HgQ8sX3O4oRZH0HWaW0fexwwTZP9jz1O/Tdf0+e55/ndor18smEvQcMkGGl3TkLpOiYmD18ygotGZRzlFovjnaIoa03THHus23GkjuvnU0fTg7DtU1j6OOzbHFoT2vieD7ImwOk/hdwzeW5JIR98V8y7t03Caeu8NCDDMLn55dWsyK+IOqCyW1QGpsXz3u2TJEVJRNQdnk9H9GzSg7D0MVj+JOgB8Dfbc0pRweKAhN7wg/+FYTOPvMFCiDa19mzqmjNTZd/DpvcgeCiQyn68DncACu6OJ84WGq15/ls/r20IsOCGw9YcBT3wyb1w2+Ijb0dtCax+EQoWgrcKNDskZcKYG+CkGQfWbh0ZRVFIu/dnGA4HV/15HluT++MNtt5xObiO4Zfvb6Ciwc9NkyWdRoguq+x7eGVmqFN1sGMVbFr8gYKFULwWt70nH3p+zbN3XdSpgRSAqio8e/1YfvHuBj7bVIo3qNPK2Bwum8aorB48N2usBFJCROJ3wxtXQPHqlgeKTQMCbqjMhw9ug5J1cNb/k1kqIY6hrhlMHRyxaUY34R8r/fxmShslxcu3w74tkD60fdffux7++ycoPBCQHb5Wa98mKFwCFjtMuAMm/RQsR74p5d9SJ7I1sajNQOpw3oDB3z7fSlaKi7OHph9xG4QQR1nJd/DS+aGUvrZKSfjrsfndzLP9EqsxDsjt9OZZNZW/X34q14zvx3OL8/lm234sqkJANwjoJg6rimnC2P7J3DptAJMH9kRtpVKgECcsQ4d/Xw1FK9tcutAo4AkV3rLFw7T7O7d9QogWHbVgqtrtZ1eFm3pfEJdNIyPZ2aQseNR8dbDxXTCDYW/dP8nG35b6uOM0Gz0crfzA1v2w4km46MnYr//9R/De7CazYmEOjiAveiSUmnPt++CMbmF3JDv31/PJxr34zPDv1LBlAbWrPyBQUYRqc2JNyyVp0uU4MocBoYDqDx9u4qyT06QwhRBdSW0JvPKj8DSfVlgwIFALc8+HO1eAI6kTGxiiKApjs1MYm51Ceb2Phdv2U+X289SCndw4KZtLxmSS0UO2eBCiVaufhz1NA6moMm4Cbljyf3DS2dBXNu0V4ljo1GDKNE3W7qri2UX5LMjbj/2wDeX8QYMxB0Yrp8QyWlm0GjRrxGBmbF+NM7ItPLrMxwNnthKomTrkfR7r14Ed/207kDpc0AulG+HVmXDT56HZqnaYu7QAPcL6qNpV86hZ+S6pM+7EkTMaRbPgKViLZ/vKxmAKoMYTYEV+JRMHpLbr+sdC0AiyqGgRG/dvpMpXhdPipG98X87NOZeeTqlmJE4Ai/8eSvs5TFSdK9MATxWsfTm0juoo6hlv59IxoS0gPtm4l4kDUiWQEqItphlaDxkI32g7qoyboA+WPwGXPt+JjRRCtKTTgqn9dT5mvbiKgooGvIFQLr2/WYrasp0VrN9TTQ+XjdduGR/dxnOeKlpLzP/TdDunv9jA3ePbSK2LYbS38bpvXx/7Oi3dH1rz8NUf4YcPxnZNQuXP31tbHFZswvA1UL3kdVLPuwfX4EmNr7sGjsc1cHzYOZ5ZuLNLBFOV3kr+vfXfvPH9GwSNIA3Bhsb3bJqNx9c+zoS+E7hp+E2MSR9zDFsqRCfyu+G718FoZzpz0BPqXE28C1S15eM6UZLTSo0nvP1CiGYKF4O3JuJbUWXcmEYoa8ZTBc7kTmyoECKSTvkpu7fGw7n/WETevjo8/tYXJTf4dUpqPFz4ryV8v7e27ZMrKrQyiTU8TeOCQRb+usTf6mlMwyRYXk5r1QybWPd6xCDuYMemVUEvrH0p9sqDwJa9tVgizNr5irdiBv24Bk1s8xwmsKqwMuZrH21bKrZw4bwLeXHTi9T4a5oEUgB+3Y/f8LO4aDG3fXkbf1/z9+j//oToSja/3+KC8vsn2Xh0mY9qb9trqMj/uhMaFx0JpoSI0poXD6yLDHd4xk2rFBW2/KcTGieEaEuHB1Nuf5ArnllBlTvQYunu5kwT6n1Brn5uBWW13laP9Vh7ENRbP+8fz3Dw3Ld+imtbPk73muSffwF548ZTcMUVlPzyV5TPeYbaz7/Am5eH4TvswWUYsPxfEafgo+7YAGye1/YxzdS20BnRPbWorkSUKCsGegL6cR14bKvcxo2f3UitvxZfhM2XD2di4tW9vLn1TR5a9dBRaqEQR1HeF0feufI3QMGSTmhcdCSYEiJKVbtafftP0+38a5Wf/Q2tFKAKuEPrLIUQR12Hp/m9vXoP++u8Edf4tFUsod4X5OmFO/nfCw+t9zEMk80ltSzavp9FefvZVlzPck1vteEDU1SuGGbln6v8nJIWIV7UbFgm3sKgvzxIsKoKf0FB46+aDz/EX1BAoLgYS1oattwc4vtb6RGsihh5Rr1Oy98Q2nxv5NWYuo5eXY1eWUmwqgq9sgq96rD/r6xEr64iWFlFqZ6IMehCsDZdd6A5EzHctZiGHlVApSrKcVuAwh1wc8sXt+AOhgerVYurKP+8HH+ZH82hkTgmkfQfp6PFaXh1L/O2z2NEzxFcMOCCY9ByITqJp6LVt6NOZ24o68BGxUaCKSGi1Eb1vsMzbk7u1coYeCDyAIwQonN1aDBlmibPLs6PuIFjNMUSArrJW6v3cMOkbFYXVrEobz9LdpST7LIydVAvbps2gPG5p+FccAusfCa0HqkFf5hm59UNLfwgVxQYdysAluRkLMnJuEaPbvpdAgH8RUX4CwpDe1pVGC2mF0bbsQnu3kL++Ano9fVoCQloKSloKaHra8mh/7dlZaKdOgKtRzJaSjJDdQfmu9vBrzc5lz1jCIrFijtvOXFDJrd6XYAEx/FbBf/j/I8jzkaVzy9n//z9ZN6SSfzQeAJVAUpeLaHw0UJyfpuDalHx6l6e+O4Jzs89/7gNFoWImWpt9e2oO1da+4redIREh5W9Na1nGgghiKrq5h/PcDD6mXrum9jCPa1o4Dr+10UL0R11aA97ZUEl1e7wACamYgkBnXP/sZjpg9OYclJPfnnukPBqUON+Aquea/JS4T0JTX6flaTi/V1ihFYqkDEWkvu3+l0UqxV7Tg72nBzosQfmvweByMFZtB0bzWUjd/6naElJKFp06XmnGCbxHxXgbhZMqfY4eky+hsov56CoGo6cUSiqBW/hd3h3byB5+k2HrqvABSP6RHW9o800TeZumounWYVE3aNT9kEZGTdnkDAi9Hdr62Uj644s8u7Po2ZZDclTQwttK72VrCtbx+j00WHnF6JLSswgNHrTcmpum50r1QKJfY+oGaZpUu8L4vHrxNktuGxa1IMWSU4rW0vrjuj6QpwQsidD8dqme1Y202bGjdUBfaQ0uhDHQocGU0u3l4d1+iHGYgkmTBnYkyevaaVj3KMfnH43LIu8jqlV9ni48J8xfiYR2kila7NjAyj2BCwpKTFdWlUVZk/J5e9fbgub8UscdwlqXDI1y9+i/ONHUWxO7OkDSZx4RZPjTOCTDaWkuGxcO6E/aYnt2N+rk2wo30CFNzylyb3djREwSBzTNCDWHBoJIxKo31zfGEx5g15e2fKKBFPiuKQbJgvzynhmYT5bSmrxBHSsmkrPBBvXT8jm8rFZJLmazUSdeiVsmdfiuimIonOlWmDYzHa1eU+lm5eXFfLm6t34AgYWVSFomMQ7LNw4KZurx/enV0Lrs16JkuYnRHTG3gTL/tHmYa1m3NgTIWdaBzdMCBGNDg2m9tdHHlWJtVhCpbuN6ngAZ/wa6stgw1tRBlQK2OJCG+j2HBhVOxqlDwe99U5Bmx0bgN7DY7vuAZePzeLRL7ZFfC9+2HTih01v8bOqAqP6JfPwpSN4eVkhZz+2iOmDe3Hj6TmcmtX+jYQ7Sn51fsTX9XodS7wFRQsfBbckWfDsOjSTZWKSV5XXaW0Uor3eWLmLR7/IwxfUafAdGmgKGjp7Kj38/cttPPrFNs4d3oc/zxxGguNAUJU9OVTiuJVgCtroXKUPh54nxdTeGneA/3lzHSsLKjBMk8CBYj8HiwlVuwM8tWAnTy3YybnDe/PXS0fgsEZ+ric5rS0W0BFCHCaxD2RPgR1fcfhsdNQZN1bXMd0GQYgTXYfeeVoLG+8eXiwhGhYtimYpClzwGEz/bShIsrWwR5Wihoo39BoMt3wFWeOiakMTPQdC2tA2D/vDNDsN/hbScmxxMKl9G2gmuaw8MHM4Dmvsf11xNguPXnYqA9Pi+fPM4Sz6xXSG9U3izje+5ZKnlvKf9SUE9FYqBHWy+kA9eoR/F1q8RrA+iBmhcmOwJoglvuk4gDvWGUohOpFpmvzhw038+ePvqWzwNwmkDucJGPiCBp9uLOGCfy2hrC60xmhfnY83LTPx0nT2p/CeBM7KPfRv/2DnqnFfu4OscTD5ZzG1uazWy3n/XMTy/HJ8QaMxkGrOFwy1+bNNpVzy1DLqfcGIx0kBCiFicObvQql6MVNCwdSoazu8SUKI6HTozFTvJAeaqoRV8ou1WEJ6QpQPFEWBSXfBaTfD5g9g6WOwPw80W2gTOwU4+aLQMX1HteMbHWbyPfDBHU02+41+nRZgT4Kcqe2+/GVjs6h0+3nsyzy8EQp8NKcqoUCq+WbISU4rs6fmctPkHL7cso+5Swt48JPvuW5if64a14+UuDaqg7WiuNrD/I172VfrxR80SI23Mz4nhXE5KS2us3BZXKhKeJDoGuhCsSjUrq0ladyhxbm6V6duQx3pP05vcrzDcvykLgrxf1/k8c6aIjyB6AaQ/LpJUZWbK55ZwXUT+/HE1zu57rSrsTlWQel3rRbbCWNxwsAzYcj5UX+kwRfkimdXsK/WF/WWFt6gwc799dw4dxVvzp4QNgiW5JJgSoio9R0FFz0V6mcEo92TUgktXbjhE3A2zTSp9lbz/vb3eW/7e1T7qtFNHZfFxdjeY5k1bBbDUoe1cE4hRKw6NJg6d3gfnvxmR1gwFUuxhDibxiWjM2K7sNUJI68K/dKDoZ3ELTawxbe48WXMBp8XqpQT8IAZXQepSft+8PsjbsutUwfQL9nF//5nMw2+IA0R1qdZVAVNVRjaN5G/Xz6ySSB1OE1VOGd4b84Z3pvNJTW8tLSQMx75hnOH9+HGydkM6d1CUNiMaZos3l7OnIU7WburChPwB0PBnqqAw6qR7LJx69RcLh2TSZy96T+5zITMiMGU5tJIm5lGyWslqA61STU/a4qVHpOa/uDISsiKqr1CdLatpbU8vyQ/4qBHa9tD6AYUljfw9Df5vDF7fOge9L4LL18I+7eFNv9ug09xYO8/ES59IabnzYtLCyip9kQMpFprsy9osLmklo837GXmqNBz2zRNVhdW8eF3xZTVebnt1bX0SrBz1tB0pgzsidpCBoMQJ7zhl4T6C+/eFBoQbu2etx7IyLnhE+g1qPHlSm8lD658kAV7FqCg4NUPnaMh0MDnhZ/zze5v6BPfh1+N+xWT+k6KdHYhRAyU1jZyHTt2rLlmzZqYTvijJ5awoagm4nv1m7+hbs2HBCr2NCmW4Mg8ufGYXgl2Vv76B8fnD9zq3fDM1FCwZkaZGmd1wejr4dyHO6wZhmGybGcFcxbuZM2uSnwBA0WBBIeVH53alxtPzya3V3zM562o9/Hmqt28umIXuT3jufH0bH5wcnqL6ZtB3eD+dzfw+ebSiIVHDue0aqTG23jr1olNqjPqhs6Z75xJpbcy4ucqF1ZS8UUF/jI/qlMlcXQivS/rjRZ3aJ2Gy+LikWmPMDWz/TN/JxJFUdaapjn2WLfjSLXn+XQ0/Pyd9cxbVxw2qNTS9hC+PZvDBpTW/v7sQ2uRAl747Few/s1Q2nKklFZbHKZp8qpxDv0ve4hpQ3pH3V7dMBn7wJdURajEGm2bh/RO4IM7T+edNXuYszCfKrcfj19vUoswzqbhslm4ZUoO10zoT7z9+N2uQRw73eH5dMTPpvoyWPsSrHwagv5Qf8MIYqpWvIEgWlIGtjPug2GXgM3V+LE9dXuYNX8WVd4qgmbk9NvD2TU7vzjtF1w++PL2t1WIE0Rrz6YOD6bmb9zLfe+sb7NzHYnDqvKzswZx67QBMX/2qKkqhJfOB3dl64UvFBUsdpj4PzD9Nx03QxaBYZgoCh22z5I/aDB/017mLi2kosHHrInZXH5aFomOQxXHDMPkttfWsmj7/qjSDgE0RSHJZeXTn06hd5IjVLpx59c8v/A3PKO58baz+cmOZL657Bu0KAucnOi6Q2cFjs9gqs4bYOwD/8UXbHpPGL4Gip6cRep597SZ6hxn03jg4uFcPCqz6Rve2lBAteJpqNsb2ujT4ght8zDppzD8EpbuauDn76zns7unkuCwsGxnBS8vL2RXRQMev068w8LIrB7ceHoOg9JDacr/3bKPu/+9LmymO5Y22y0KGcku9lZ720xtdFhU0pMc/PsnE+iT5Gz1WHHi6Q7Ppw57Nhk6FCyE6j2h/oY9kX9sstJjwDhmTcpucmiFp4LLPrqMCk8FBtGvg3ZoDv58+p85J+ecI2+vEN1Ya8+mDh8a/OGw3ry9Zg/LdlaEdShaY9UUcnvGhz0gjjvJ2XDnKtj4Lix5DOr3gREMrWlQtFDnxtTh5AtD1XX6dv6+Dx09i2ezqFw0MoOLRmawbncVc5cW8q+vd3DRyL7MmpTNgF7xPPnNDhZvL485lanGE+Ca51fw5cUq6oIHoaGcSyf/lOe2PBFDnvghDouDG4fdKIGUOC58vbUMi6rQvK5pLNtDNPh13li5OzyYciTC+FtDv1pw+kAnZ5+cxnUvrGRfnZc6bzBsYCuvtI5564oZ0Cue+384mP+sL4mYMhxLm31Bk8LyBqJZbuUNGhRVebjoiaXMv3sKqfHHbmNhIY5rqgYDzmzyUq5awoffFYf1lf684s9UeavCAqmqxVWUf16Ov8yP5tBIHJNI+o/TG7M7vLqX3y/9PRP7TiTJ3vbmwUKIcB0eTKmqwtPXjuH6F1axobg6qlkLu0UlM9nJ67eMb7HM7nHFFgdjZoXS94pWh355qkMzUQm9Qwu/ncnHupUdYlS/ZEb1S6a0xstrK3ZxxTPLGdonkVWFlRH/bltKC/JsX3lgXYhJaUUVS999jSkzZsEpl5GsWXgsbSB3f3N3k/zuttg1O+N6j2PWsFkd+ZWFaLfyen/E6pixbg+xv67lzTtb4wvq7K70sLG4psXtfnUT9EBordNtr60lrYX9omJtc/NAqvX1YSZVDX5ufnkNH9x5egzfUIgT2+kDe/Kb9zfiDxrYLKH1xhWeChYXLQ5L7SufX87++fvJvCWzybrjwkcLyfltDuqBz5uYvL/9fW4cfuNR/z5CdAedkrTusGq8Pns8f/nke/69ejcKSsTUD4dFxSQ0m/XXS0/BZetiOfSKEiq13p5y611M7yQHP//hYO46cyB//Ggzvh3hHUbD10D1ktdJPe8eXIMPLWp1DRyPa+D4xt83GFaeSf45U0YeGvGelDGJR6c9yv2L7scX9LWZpuDUnEzsO5FHpj0SsYCFEMdCUDcizs4cvj1ENMFJS2XJW2MYJne+/i0r8itaDKSa8wZCs0SRxNrmw7U1qAIQMEy2ldayuaSGYX1lRFyIaCQ5raTG27j2+RX4dZOgYeCL+wK9WfdJ9+iUfVBGxs0ZJIwIpfTaetnIuiOLvPvzqFlW07jxvU/38cqWV5g1bJb8PBWiHToterFqKv/vR8O4b8Yg5iz5gs8AACAASURBVH1bzHOL8ymuDv3QNoGe8XZm/X/2zjs8jurqw+/MbFevbpKr3HvvlWYDoRhTDYRiwEAoCSEhQAgEkvARklAC2KaZjgPBOKZXd1lu2MY2uEhykWVbvW2f8v2xbvKOpF0h2bK47xM/edidnb272rlzf/ec8ztjOnHl8I4izeM0wmFV2FRQiVmpXeRpQRJr91RwqMpHm/hjluYTMyfy9rlvM2fzHFMnIhkZu8VOuiudWf1ncUG3C8TEL2hRxDutWBUZ9YTeadG2h4hzRD81v7N2Lyt3leIzSa+uL0pUV2petGM+QqSbKhCyhH95eT7/vLz506EFgtMZX1DjpeV5vLJyN9W+ILtLj9Vsx3RbgqzUbp/g2elBD+rED63tzKs4FOIGxFGzteaomIJQr8YtJVsYkDageT+IQNAKafZQUJzDyrVjOnPtmM4YhoFf1bFb5CYzSxCcfPKL3aaPR5MWZLPI5BbX1BJTAFlJWTw58clQj4xdH7D+4HoqA5U4FAcZcRnM6DGDfqn9muRzCARNzbBOSZiZ+kTTHsKmSEzskRbV+xqGwfPf5ppmAEQSJTIjmjEfTzS1Vppu8PH3B3js4n6nX2aCQHCSKHcHmPlSDnnFNaabJZIl3AxLq9GwxFqQlPC1liXBgndP7Yi0LMmUeEuabtACwc+Ik3r3kiTp9KiJEtSLXzV364omLcgwoMZXt3VroiORG/rdwA39zBdsAkFLpHubOLq3ieP7/eHtIeJHTEeOSaIyewElHz1Zqz3E8UiSxLVRGvGsyS+j3BPe2DeaKJEZkY75eKKttbLIEkVVfjqnCjElEJyIJ6By2dxsdpe660n/DRdYSqyCWqNiaEaYoFIrVSyxta83wzAI6FE0BxcIBEcRdy9B1FgV2dSpMZq0IElC7EQLWiWzJ3bj3vfN20PE9p1MbN/J9b5+WKekWr3YIuHNnD14f6IjX11EMubjibbWSpIkavwN98QRCH6OPLRoK3vLPKZC6mj6bnkeikPG0dFB2i/SiOkRgyvLhWSRqFpfRcKIYzWJmk+jenM1bWa0qXUuSZKIt8af+BYCgSACRMGJIGpOTM07wvFpQZ4d2ehBH4am4s1dR/m3r9Q6NqjpdEgSPWYErY+z+7YhM8mFpREtCxxWmfum9W74wBPYW+oxNZ2INkrUFBy/qRIJhmEQIxr4CgRhVHqCLN5UaLp5WbVmIWVfv0jCqMvI+ssv6PFkT5KnJFO9oRoAxaWQflE6hW8WUr25GkM1CBQH2Pf8PqzJVhLHJNY6X0AP0Cul10n5XAJBa0PcwQRRc/2Yzjzx+XbT+oxI04K6psbSJTXmZA1ZIDhpWBWZt24ayfnPrqC0xh+xM5/DKvP0FYPpnxG9s11djXKjjRLZLTKqbqBF0DDKKoOqEybioq21UnWDNvHChEggOJH31u9DNqkvPzF9V3O3wZa0j/jBEvGDj0WX0s5NQ4lROLjgIIGiALJTJn5IPJm3ZCJbj+2ly8hM6DCBZEfySflcAkFrQ4gpQdRcMiyDxz/7sc7nG0oLirErzJ7UrTmGJhC0CFJj7Xxy53iueTmH/BI33oBWp125y6YgAfOuHcbYrNRGvV+8w2r6eDSptxZZ4oZxXfj2xyL2lHrwq5q5zbsUEoyDOibSu108b63eS+CE3lqRbqooksS5/duJlF+BwITXVu023Sg5MX1X83TF0FxISnjNU/LEZJIn1i+S7BY71/W7rknGLBD8HBF3MEHUxDusXDS4Ax9+t980/aAhbIrM1L5tm2FkAkHLITnGxkd3jGN1XhnzluWyMrcUu0XGF9SQJAlFlkiNsTF7UjcuGtThJ6W6De2UxKaCirAoWDRRIrtFZnTXFH53Tk827qvgxeV5fPVDEVZZQpIkDMNAMwwuHNSBG8d1oUebOA5Uenk7Z6/pmCKptbJZZGaN79Lozy0QtGYO1dG8Ozx9V8JfdB6O9v9BkoNRvYdNttEvpR8DUoUlukDQWISYEjSKRy7oy8Z9FeQV10TVYFSRJF69fvjRzu0CQWtGkiRGd0thdLcUiqp8fL+/kjez92BIMGtcV0Z1Tcai/PRr4ZrRnZi/ajfhSXeRR4lcNgvjslKRJInBHZN4fuZQavwqh6p8ePwasQ4LbeMdOG3H0gXbJTiZ0CONZTuKo95YscgSPdvGiYa9AkEdBOu4pszSd9Xq/vhLSrCnfhOxoLLJNjrEdeDZM54V7WoEgp+AEFOCRuGwKvzn5tFc83IOO4qq8QXrX0jJUiidqXt6HC8uy+fpKxKaZBEpELR0DMPgu30VvLgsj69/KMI4LHhW55ZitypcO7oT14zqRHodxi6RkJHkYkjHJLLzSk2fbyhK5LDK3Di+C/IJphmxdguxabH1vve/Lh/EL55dQUG5ueOYGYoskRRj46VfDovoeIHg54jdKpveW+tK3w2WTgbNhb3NYgAk2dwlU0LCYXHQJ6UPz53xHDFWUb8sEPwUxGpW0GgSXFbeu3U0d5/RndRYGzG28AJ3p1XGbpGZPiSDT+6cwDs3j6Lar/Lr/2xC1aJPERQITid2l7g5859LufqlHD7fepCAphPUDIKagU/VqfQGmbcsj/FPfMs9/9lIoBFps0f47Tk9cFgbN6XbLQpXDM9s1Gtj7RY+uHUMWelxOCN4f4dFpkOik0W3jyU1VhhPCAR10SXVfCOjPufcooVbcef+jkDpBAzNgaHZcVlcuCwuYqwx2GQbEzIm8PwZz/PqOa8KISUQNAGSYdS9kzhs2DBj3bp1J3E4gtMVXTdYtrOY99cXcLDSR1DTSXTZmNwzjUuGZhB3XIG8L6gx67V1pMXZefLSgSiNsJAWNB5JktYbhnHahwRa+vy0tbCSK+atxu1XTY0cTsRhlenTLp63bxrV6ObmC9bu5U//29pgpPh4XDaF92aP/snpdr6gxrtr9jJveR4VnmCY6UaMTcFpU5g1rgszR3WqNScIBEdoDfNTU81NH363nwcWfo/bpIccQM3Wb6let4hg6b5a6buOjCPtFTSG9KjgpslpBLQA8bZ4+qf1J9XZOKMbgeDnTH1zkxBTglOCN6Bxw/y1tE908vcZA8LSiwTNR2tYrEDLnp8KK7xMe3o5ld7oisEdFpnRWSm88svhja5h+PC7Au774Hs03ag37c5pVbBbZd6aNbJJ65YMw2BNfhkff3+AvOIaNuyt4KLBHTi7TxsmdE8T17qgXlrD/NRUc5Nf1Rjy6Fe4G9nUOsam8OK1wxjTSJdQgUBwjPrmJlEzJTglOG0KL183jOteXcsfPviev03vLxZZglbDE5/9SLXPXEi5ty2hau2HBEsLkG1OrOldSRhzGY6MvvhUnZy8MlbuKmVc98YtgC4anMHQTsm8lr2bd9aEnPYCqo6qGbWMX+49pyeXDM0gwdm0ESJJkhjZNYWRXVPYV+bhyhdX89eL+zfpewgEPwfsFoVbJnTlhSW5dfaSqwuLLNEhycnobinNNDqBQHAEIaYEpwyXzcKr1w3nl6+s4YEPt/CXi/oJQSU47an0Bvl0y0HT1L6qNQupzHmflLNvx9FlCJJiwZu/Hu/OHBwZfQHwBDTmLstttJgCyEx28eB5fbj3nJ5880MR+yu8eAIacQ4Lmcku7n7nO2YMy6izP1VTEbKCF7WRAkFj+dXkLLYWVrJsR0nEgkqRJRJdVt6cNVK49AkEJwFhQCE4pcTYLcy/YQTbD1bx0P+2UF/aqUBwOvDeun3IJgsY3e+mYsVbJJ91K66eY5BtDiTFgitrZK1+TwBr8ss4UOn9yWOxWxSm9W/HrPFdufOM7lw/tgtn9m7DqG6pfLn10E8+fyTv71ej21EXCATHkGWJ564awrn92+IyMXk6EadVoUOik8V3jCM9rvEOoQKBIHKEmBKccmIPC6rv91fxyOJt5oJK9cPuFbBtEWz9EPKXQcBz8gcrEDTA4k2FpjvI/v0/YqgBXD1GN3gOWZJYtqO4OYYHwC8GtuOjzYXNdv4j2K1yoxp7CwSCY1gUmScvHciL1w5jfPdU7JaQS+4RZAkkoGOyiz+e35vP755AuwTnqRuwQPAzQ6T5CVoE8Q4rr98wgmtezuHRj37gj+f3DqUnlO+GNfNg/WtQa7dfAl2FgVfCqNsgNetUDV0gqEWFx7xWSvNWIbvijzbZrI+gptd5nqbgjN5teHDhFsrdAZJibM32PnaLTEDV0XVDpPAKBD8BSZIYm5XK2KxUCiu8fL71IMXVfvyqTnKMlReX5fPGDSPolCqszgWCk40QU4IWQ4LTyhs3jOSql1bzt09+4A/295BWPweGAVrA/EUbXoeNb8Hgq2HaExDBQlUgaFbq0AyKMx7dU4WhaxEJquYk1m5hXPdUPt96kAsGtafcE0TXDeKdVuIdliars5AkCZtFJqDpOMS1KRA0Ce0TnVw/tkutx3YVuVmRWyLElEBwChBiStCiSHBZefOGEax5ZiZBdTk23V//C/Rg6N/Gt6H6IFz2Bsgie1Vw6kh02dhTGp6Cau/QC8lixbMjm5he4+o9h9Uik+hqPnMITTfomOziL5/8wIMfbsGqyEhSKCKWkeTi1ond+MXA9jgjqNFoCIdFxh/UG907SyAQNMy4rFS++uEQM0d2OtVDEQh+dohVp6DFkfTd85ytr8Sm+yJ/UdADud/AVw8138AEggi4aGB7nCbCQbbHkDhuJmVfzsGzIxs96MPQVLy56yj/9pVax+q6wcQe6c0yvuU7ixn+l694c/Ueqn0qqm7gDWp4AhpBzSC/xM3Di7cy9LEveX3V7p/8fnargk+YUAgEzcq47qmsyi1Fi6RDuEAgaFKEmBK0LAJuWPZ/SGrtnf3OT1WT/vdq3IFjN4qXNgSYNN997KCgB3Lmgbv0ZI1WIAjjkmEZ6HW4UsaPmE7SlBupzF5AwbMzKXjhOqo3fISze21TilFdU2ib0PROXIu+289Nr6+jzB3AHahb4HgCIXH1t09/5K+f/PCT3tN+ODIlEAiajzbxDtrE29myv/JUD0Ug+Nkh0vwELYvN71GXxtcMeDonwP3j7XW/XpJhw2sw/jfNMz6BoAHiHVbO69+ORRv3o5loqti+k4ntO7nO17tsoUadTc2q3BJ+/8HmqPo+eYMab2TvoU28nRvHNW5MDmvj7dH3lXmYv2o3n3x/gGqfioFBnN3KOX3bcP3YLnQW9SECAQAVngAJTivXvrKGoKYT1HScVoW+HRK4ZUJXJnRPEyYwAkEzIcSUoGWx8ikIuk2funeMjSdW+rltuI1ERx03BdUL2c/B2LtF7ZTglPG7qb34dnsR5Y1w5BuYkcDobilNOh7DMPjtfzbVKaTc25ZQtfZDgqUFyDYn1vSuJIy5DEdGX7xBjSc+286MIZkkNKKOqzGNe3cVVfPAwi1s3FeBbhgEj1Olbr/G2zl7eXftPvp1SOCxi/rRu1181OMSCFoDVb4gDy7cwudbD2IYBoHjrpWgppKdW8rmfRU4bQq/m9qLy4ZlnsLRCgStEyGmBC2HgBsq99b59LD2CpM6W3hylZ/HptSTAhWogZqDEN++GQYpEDRM2wQH79w8isvmZFPjV4mkjMFplWmX6GTHoWo27C1naKdk0+P2lXnILa7B7ddw2RQyk11kpcfWe+6c/DIqvObCrmrNQipz3ifl7NtxdBmCpFjw5q/HuzMHR0ZfINTH5r31+5g1PvroVLSRqZy8Um6YvxZPQKOury2oG6AbrN9TziUvrGLeNcMY1z016rEJBKczh6p8zHhhFYeq/AS0ujcs3AENd0DjT4u2sv1gNQ+e17vJHDsFgpZCYU0h6w+tpzpQjUW2kOJMYXS70bisrmZ/byGmBC0HXyUoNtC9dR7y58l2xr7i5q6R9fTGkS3grRBiSnBK6dU2no/vHM+s19ext9SDX9VMRdWR5psXDurAoxf1Y+WuEm56fT3/uHQgk3uFTCg03eCbH4uYszSXLfsrsVlkDMNAkiSCmk7nlBhundSNqf3aYreEm1/MW5qH16RGSve7qVjxFinn3o2r55ijj7uyRuLKGnn0v71BnXnL8rhhbJeoU4Xslsgb9/5woIrrDwupSPEENG56fR0LbhnFgIzEqMYmEJyuVPuCXD43m8JKX8SmE95gKKqb4LRy5xndm3mEAkHzoxs6K/ev5NUtr7K5ZDOKpKAZGhISiqyg6RoXdLuAq/tcTZeELg2fsJEIMSVoOchWMOpfdPVLVzi/h4XHVwTonVZXGp8REmUCwSkmM9nF53dPYHNBBS8uz+PzLYeQJZBlCVUziLErXD+2C1eN7EhqbKgWcFLPdF68dhi3vLGOB8/rw9BOSVz54mrKjzONOFGc/Hiwmvs/+J6H/7eVN2eNpG/7hFrPL9tZbBrl8e//EUMN4Oox2uTZ2lT7VPJL3XRLqz8KdiKhNL+GxZFhGNzyxnpTIVVfGiKEFok3vb6O7PvOEHUhgp8FT36+ncIKr6mQaiht9/kluzi3f1uy0uNOwcgFgroxDIM9pR5K3X5UzSDBZaVbWixWJXy9VxWo4tYvb2VXxS48ang7Eg7fSj7Y+QGLchdxY78bmT1wdrNEZYWYErQcnImgN7zoemSSgyFza7hndB1GFFoQYpq25kQg+CkMyEjk2SuHEFB1KrwBfAGdOIeFBKfVdPE/tFMSb980ipkvrqbarxJQ9QZTBY+k8lw6J5s3Z41kSMckAHxBrU53Qc1bheyKj6iJsEWRqIi0BkzXIe9b2LyA+4q2k/aVApvaQc9zod90sDrDXrJuTzklNeF95SJJQwSo8amszC1hfPe0yMYoEJymeAMa760vqFUfdYRIrpegpvPyit38bXr/kz10gcAUt1/lw437mbs0j+JqPxYldF/UDQNFkrhmdCeuGdX5qMttdaCaqz66ikJ3IUG9/vuSaqiomsorW1+hKlDF70f8vsnHL8SUoOWgWKH7ObD9E6izWgKykmUu72vlmTUB+qebRKc6DAVnUvONUyBoJDaLTHpcZJbnbeId6AZRmzd4Ahq/fHkNn949noykUK64hITZNaU449E9VRi6FpGgapCgD9bMg+x/h2ogAzX0BPACpYQE1qf3wqCrYdyvIb7d0ZfOWZoblooYaRoihMTknCW5QkwJWj2LNxeaPh7p9aLpsPC7Ah48rzcxdrEMFJxaFm8q5Hfvb0aSOJaZcII+enF5Pi8tz+eqkR158Nze3PHNHRxwH2hQSB2PT/Xx/o736ZbYjRk9ZjThJxBiStDSGHsn5C2p09HvCA9NtPPGZpOLyBYbcvITCE5z3snZg9uvmj7XUNqbJ6jy/Le5/HV6fxxWBclcS2Hv0AvJYsWzI5uYXuPqHY+qGSTW5+bnKYPXL4CSXSFXTTMCh6/rda/A9/+BXy6Gtv3xqxpLt4enIkaThgiwZncZ1b4gcY7oXQcFgtOFd3L2mqbDRnO9KLLE0h3FnNu/XYPHCgTNxfxV+Tz+6Y8NbhoGDqe2v7tmHz8cOkiuYxsBPRB2XPnycko+LyFQFEBxKMQPjafNjDYoMaHNQp/m45kNz3Bx1sUoTbGBeBghpgQti8yREJsO5fm1Ht59d+3c7swEGd+DJnbIFgd0P6s5RygQNDu6bvDSinx8JsYNkaTxaDos3FDAHUYu2vKlDKrpxPqUbhhS7UiubI8hcdxMyr6cgyQrOLoMRpIt+HZvxLd3M0mTbzh6bKzDQpeUOvo6BdzwylQoy4NIdgr1IHjL4dVpcNMSKm0ZWBQJ9YRcxmjSEAGsiky5W4gpQeum2CQdFqK7XlTNME2rFQhOFl9tOxSRkDoeb1BjbV4AJeFM7G0X13qu5NMSij8tJmNWBrF9YgmWByl8o5DdT+6mywNdkA+bPfk1P8v3L2dS5qQm+yyiEY+gZSFJcPkb0BgrS6sTrngLmnC3QSA4FSzbWVyv+17yWbfi6jkG2eZAUiy4skbWEj4Ahs/LByt2EDthAnfd9gucdnOBET9iOklTbqQyewEFz86k4IXrqN7wEc7ux3a3HVaZm8bX4+S3+C6o2F1LSHV+qpr0v1fjDhwTSC9tCDBp/nFRZ38NvHEhgaB2OBWxNsenIUaCJEFAa1yDYIHgdEE16wZOdNeLbhhHd/sFgpONrhvcv/D7ensfHnjtbvb+cwYF/76GQ//5E76CrQBoukKgYgR68JjRkubVKPqwiPZXtyduQBySRcKWZiPztkwCJQEqV1UePdajenhlyytN+nlEZErQ8mjbH2a+B29d1mC63xF0ixN5xnzoOKp5xyYQnARy8sqOOvcdTzRpPD6Lnc0Dp3D79GGMMwwSPt9dp+V4bN/JxPadXOe59KDKjD51mLq4S2Hb/0AL3+XWDHg6J8D94+swi8EAbwXxB1ei6uE31WjSECG0yIwXUSlBKyfOYeFgVfjj0VwvFlkmwSmuFcGpYVVuKTV1pLFHZDpkSATKRuNo8xkAnp0e9KBO/NDaGUuKQyFuQBw1W2tImnCsln572fYm/TwiMiVomXQeB7O+gozhodQ92WTSlxSwOCmL782vnX8hmHX2yR+nQNAMlLh/ehoPQJk7lFMuSRJPXjoQhzX6Kd9plbmF3ZRdeRmedevCD9gwPxQSMuHeMTaeXOWnwlePFWGghvLP/oaZ4eDxaYieHdnoQR+GpuLNXUf5t+E7i7EOy1GLeYGgtTK+eypWJfyai+Z60Q2DoZ2EUZPg1DB3Wa7p5l7k2RcWghWjMIzQvVCr0bDEWpBMrgtLggW1prZw86m+Jv08QkwJWi5t+oQE1a2rYOh1EJMOFmdIXLlSYdBVcPO3JP06m6rk/vz7m12nesQCQZNgN+mpAdGnvR3fwHdsVip/u7g/Dkvk077TqjBzZCfuefwO2vzhPvb/5h4O/e1xdN9xN6I186COG9Ow9gqTOlt4clX9tRkdarZw/bBknNZwkRhJGiKAwyIza1z0TYUFgtON68Z0Qa5jAyPS66VPu3i6RtkzTiBoKtbkl5k+Hq3pkO5rA4ASq6DWqBgmKbBqpYoltnYinkVu2sQ8keYnaPmkdIPzngz9M0ECHr9kAOc+vZyz+rShX4cE0+MEgtOF9olOrLJE8ARDhmjSeCSgXUJtG/aLh2SQHGvnrne/I6jpuP3mosxlUzAM+P3Unlw3NtQ1Pm7KFJyDB3Po0UfJv3g67R//G84BA6CmqN5x/HmynbGvuLlrZN2NtC02B3cMgNc3mEewGkpDhJBZ4eXDO9Z7jEDQGuiY4mJgZmKdC9KGrpcYm8LsSd2aa3gCQb2oml5nvV502RcGhh7qWejKciFZJKrWV5Ew4rhaKp9G9eZq2sxoU+uVifbERo/fDBGZErQK2sQ7ePD83tzzn034VVGALji9OW9AO9MISzRpPE6bwqXDMsPOMbFHGuseOJN/XjaIwZmJWGQJl03BZVOwKhKdUlw8dH4f1v/xzKNC6giWpCQ6/POfpN15B/tu/xVF/3yyno5wIfqlK5zfw8LjK8JtbI8nXvJy84SuptGphnBaFa4e1YnkmLoFm0DQmrj/3N6NStu1KhKZyS7O6JXeDKMSCBpGliRM/IaA6LMvICTKFJdC+kXpFL5ZSPXmagzVIFAcYN/z+7AmW0kcc0w82RW76DMlENTFRYM68On3B3n6q538bmqvUz0cgaDRdEqJYUBGAmt3l4c9Fz9iOnJMEpXZCyj56EkkmxN7myziR19e67gkl43hnc1rIiyKzDl923JO37ZU+4JUeIJoukGC00qiy4pURwrR0TFMm4Zr+HAO/OlPkKrXeWM8wiOTHAyZW8M9o+upZ7LH85uzepBf4ubrH4rwBiO7mTqtMmOzUnjg3N4RHS8QtAYGZSbyj0sHcs97myK2lrYqEmlxdt6+aRSWOlKJBYLmRpYlXFbF1GQpKtMhQ0ayHDMpSzs3DSVG4eCCgwSKAshOmfgh8WTekol83MaDYRhc2vPSJvs8IMSUoBUhSRJ/ubg/0w6n+w3uKIprBacvsyd2Y2vhd6ZFug2l8TitCrMndm1QFAHEOayN6stkSU0l49//Rv/bFyiB+lP9spJlLu9r5Zk1Afqnhy/itKAPJakzkiTxzBWDefSjbby9Zi+aboT1njqCIof6Sl00OIPHLuonaqUEPzvOG9CeWIeVW99cj2FQ7waE06qQlR7L6zeMIElEcAXNjKYb5BXXUOENIkuhzb0uqTFH70lT+7Xlw42FaCfM79H0PpQUH7Kt9r0neWIyyROT6xyXRbYwrsM4Up2pTfhphZgStDLS4uw8fEEffvveJj6+czyORqQMCQQtgSm90jmzdxu+2HYwqqaGVl2lb7KDK0c0f/2QJEkoZ92H8fkDSKq3zuPe/j5Izn6VMq/Bqn0a095y88B4O+M6WtAMiW/VAfRTY2hLaNfyTxf05erRnXh1ZT7/Xb8fiyyhY4ARShEJ6joXDerAjeO60L1NXJ3vKxC0dib2SGPNA2eycEMBc5bmUe4JoMgShmEgSRIBVadDopOMJBev3TA8og0WgaCxlNT4eWfNXl5ZkY9f1VEOb3KpmkFSjJWbx3flkqEZ3DiuKx9/fyBMTEFk2ReKbKAkrajLSNYUWZJJcaTwyJhHfvLnPBHJMPOjPcywYcOMdWZWuAJBC+f2tzfQPsHBA+f1OdVDaXFIkrTeMIxhp3ocP5Wfw/wU1HRufXM9K3eVRpT25rDKdHfoPPrFP+n53DM4+/U1Pc4TUClzB46m9iU4G07tqxN/Ncbfu9cppv6Z7efxFQHmnO/gnG4WbAp8tktl2R6Nv5/twG3YmaXdz5DxU7n3nPD0XE9AZe3ucio8oZqrBKeV4Z2TibGLvcDWSGuYn07V3GQYBlsLqzhY6cOnasQ7rPRqG4fNIjPx70v4/O4JtD3BlEYgaAoMw+Cpr3YyZ2kuAP46DCacNgXDMHjkgr7MX7WbHw5UN/o9lditODu8iyQHGz5WspDmSmX+1Pl0iO3QqPerb24SdyNBq+TRC/sx9all5BXPKAAAIABJREFUnN23LcM71x3yFQhaMlZFZt41w3ju213MW56HrhumeeZH3PeuGJHJH6b1xj/Szr5bbiHz+edwDhwIhDrOL91ZzNyluazbXY5VkZGkkGBrG+/glolduWhwBrHRihR7HIVZV5D8w9s4pdoW6JU+g4e+9fPqhU6m9z6WSviLnlZ+0dNKwFDYY7QhW81ia/Ye7j6zB9YTajlcNgsTe6RFNyaB4GeIJEn065Bg6mh78eAOzF+1m/umiXpiQdNiGAa/e38zH20+UKeIOoL38P3r4f9tZWSX5J8kprSavnj23Iy9zUcojv2AjiTXfn9DtwEGwepBPDbxkUYLqYYQYkrQKkmOsfHoRf24971NfHLXeGRJYvGmQt7K2UtRtQ9VM4i1WxjdLYUbxnWhm+i3IWihyLLEHWd055aJ3fh860HmLcsjv8SNL6hht8i0T3Qya3wXfjGwPS5baEq3nXkmWCzsu/U2Mv79LFuSOnP7WxvwBNSjYkw9zi1pX7mXv37yI499/AN3n9mdWyZ0iypS9YjvCq7XVzNIzsUpHXPtyy7Q8Klwce/wW41qyFQSyy8D9wESmmGwdHsxZ/ZpE3asoHkorPBSUO7FE1CJc1jonBJDimh63Cq5cVwXLvj3Cn41JSv6DROBoB6e/monH20+ELFpEIA3qLNkRwkzhnTgky0HTWuDI0H3ZeLdcyvu7YupXvdfAsVlKA4Fe2YMqVOzsKacQ7BqCOh2/vrRHhbdHu5w2xSIK0rQajmnb1s+2lTIJc+vYk+ZB6DWBVtU7WdvmYf31xfQq20cD5zXhxFdRBRL0DKxWWR+MbA9vxjYPqLj4yZNQnriCRb86Rn+b8Cl+Bq4Vx25Np7+ahd7y7z85aJ+EQuqvFIf1wV/z7+tzzBG3krM4QhVqccg1SVhOcEcwmPYKDYSuTzwR4oJWdaqmkFBuSei9xM0HlXT+ebHIuYszWVrYRW245o4B1Sdcd1TuWVCN4Z3ThL1Na2IzGQXY7JSeXfNXmaN73qqhyNoJZS7A7ywNNc0IuXetoSqtR8SLC1AtjmxpnclYcxlODKOpZ9/se0QC24exT3vbWJfmZeAqqOZlB/F2BVUzSCg6mHtOKrWLKQy531Szr4dR5chSIoFb/56ypdtJWnysea/Px6oZldRDVnpTb95LsSUoNVSXO3n+/2V7C6te4GmHnYL21RQybWv5PDYhf2YYdKbRyA4Hdme2SciIXU83qDGwg37aRNn564ze0T4Gh0/Nm4O/oaz5XXMtiyml7SXBKdOiSd0jcmShAcHlUYMc7Tz+a82EQ/H6jeCmo4nip1NQfTsOFTNNS/lUBNQjzZsPnER9M0PRWTnltIpxcXrN4wkLU5EqloLt0zoyq1vbuCXYzqH0mnVAPirwOIAWwxRVfMLBMCCtXtNfzZ1CRzvzpxaYkozDArKvXzx64ls2V/Ji8vz+HzLQfyajkyoH9WQjoncMLYLv16wMUxI6X43FSveIuXcu3H1HHP0cVfWSFxZI2sdq+kGr67M5y8X92/KrwAQYkrQSqnxq1w2N5uC8rodxk7EF9R5cNEW4pxWzunbthlHJxCcHH73/qY6hVR9u4beoMbzS3K5ckRH0uMbLliPsYdcMw1kPtdH8HlgBF2lQia2XYlieZN7tvaia+8+ZOt9Wa33xqwxlVWRiRPpR83G5oIKrpy3Gk9Aq7fRskEoSrnzUA3Tnl7G4jvG0S7BebKGKWhGBmQk0j1BZ9uifzBwz+tQVQCKFXQNZAsMvAJG3QZpPU/1UAWnAbpu8PKK3WFus9EIHLdf44WluUzr345+HRJ4+orBQChKrhvGUUfmVbtKQr3RTtj88e//EUMN4OoxmoZQdYPPtx4UYkogiJRHF29jf4XXtEdNfYtIX1Dnrne/I/u+M0QvDsFpzaZ9FRRW+Eyfi3TX8M2cPfzmrIYXVn3bJ5Bb5K6VnpFntCfPeikx4xT+/dkHpBhn4+jSFUnWTHuGWGSJrHRhc14fBeUeVuWWUuUNIkkSKTE2JvZIa3CuKqzwcvVLOabmJXWh6gblniBXzFvNp3eNP1qPJzhN0XX4+mFeLplL4JABHDaLUQ//v67Chjdh07vQtj9c+hokNE+xvqB1sL/CS41fDXs8GoED8P3+SlRNr9VI+vj0Y4ByTxBMtoE0bxWyKx5JjqwNjtl4mwIxOwpaHdW+IIs27SdgksMb6SJywbp9zJ7Y7WQOWyBoUl5anodfDV88R7pr6Fd1Xlu1hzundK91kzPjxnFd+GzLQdMC5Eh6hgDEWmBU1+hrFlVNZ8n2YvJKaqjxqcTYLXRKiWFKr/SwG/LpiK4bLNtZzNyleWzYW44iSwQ1HYlQNC+oG5zVuw03TejKoMxE03M8+82uOgu869tc0nSDoiofH2wo4OpRnZvvQwqaF02FBVdD/lIUzUedcUZDBVWF/Rtgzli44XMRpRLUSYUnGFYPC9ELHKssU+1TI9jADn8vxRmP7qnC0LWI3k8yOUdTIMSUoNXxwYb9phdMpItIX1Dn5eX53Dy+K7LJRCEQnA58s70Yk8BsVLuGmm6w7UAVAzLMF+lH6NchgYwkJzuLakyfj+07mdi+k+t8vUPSuXjLVxTcvpjUW2fj7N9wGkZRlY83Vu/h9ew9qLqOP6ij6gYWWcJulZEliatHduLaMZ1O2zQ1t1/lxtfWsrmg0lQMBbTQY59uOcA3PxZx4aD2/OXi/kcbZR45x8LvCkyj9JFsLnmDOnOW5jFzZCcOVfl5b/0+cotqqPGrJDit9GufwPQhGSS4rGHnF7QQFt8J+UsgGGHau6GBtwJePRduXQVxwmFTEI4sm8WKohc4AU1n8j+WIEsSCU4L5w9ozzWjOtVKMU9yWU1lkL1DLySLFc+ObGJ6jWvwvZrLyVKIKUGr443Vu013yKNZRHoCKhv2ljNM9KgSnIYYhoE3YJ7OEM2uoSSFdh8j4Z6ze3L72xtMO9o3hN1hZ/bcP6N//D8K7rgTe1YWqbfOxjV0qOnxK3eVcPPr61B1I8xAQdUN1MPmCi+vyGf+qt08N3MwU3qdXgtCb0Bj+gur2F3ibrB3i26EjEMWbSykwhPk+ZlDjm4EfbhxP7JJhXg0dQ0l1X4uem4lPxwM9YQ5Pur/yfcHefyzH5nary23TupGr7bxjf7MgmZg31rYurCWkOr8VDWeIOTfFUuMLfTbeGlDgDc3B1lyXczhowzwVcCXf4Lpc07BwAUtnZQYO0EtfG6KVuDAsftMmTvAvGV5zF2Wx9huKTx8QV86pcQwpFOSqcufbI8hcdxMyr6cgyQrOLoMRpItdaaST+3XPPXwp38OhEBwAsXVAdPHo11EHqryN3icQHC6cfyuYVMyNisFvRFCymlVePPGkSQmxZF89Uy6ffE5cWefReHv72PPNdfizs7GOO4mumxHMTe+thZ3QGtQZAQ0HW9Q47a3NvDF1oNRj+1Ucsc7GyISUsfjDWos3VHMU1/tOPrY1z8UmUa1otlc8qk6mwoqCah6WPq0Nxj6O3y06QAXP7eSxZv2RzxewUlg1TOghtdOagY8nWN+rzyKrsK2heCrbKbBCU5n2iY4yEwKj/ofL3A8O7LRgz4MTcWbu47yb19p8Lz+w/PM0h3FnP/sCr7bW47DqnDZsEzTtML4EdNJmnIjldkLKHh2JgUvXEf1ho9wdq89tymyxPVjOzf689aHiEwJWh1mOyUQXej5yE6vQHA6IkkSTpsFt0mxbTS7hoZhkBhB+lZQ07n4+ZVRj9NhlVlwyyj6ZyQcfUy22Ui67DISp0+n6uOPOfjnR1Hi40m5dTZl/Ycz+831Ye5RDREyltnI/341lu5tWr7JRW5xDct3ljSqd4s3qPHi8nxmT+qGy2ah3P3TN5ciQTMMvEGDe9/fjCRJnD8gsn5ogmbEXQo7Pwcj/Hd07xgbT6z0c9twG4mOetLZJRk2vgOjZjfjQAWnK7MnZfGnRVvCzG0irZWtD92Aap/K1S/nsODmUQRU3TRdGRpOJQfo3S6ermlN32MKhJgStEJcNsV0JzaaRWQod1fUAAhOXyb3TOOT7w+E1U1FkxZBTQ3x//gzlVOmEDthPEpCAmb8/v3N5BW767XcNsMwDEpqzCPAksVCwoUXEn/++VR/8QXF//wXT7Ubhz+pj+nxDYmMgKrz/JJc/nX5oChHGTk7D1Xz8op8lu4opsavosgSSS4blwztwJXDO5ISG1nPpldX5pumS0ZqoCNJsGhjIVeO6Firfup4oq1rOJ6GHFHvfW8zfZpx4SKIkF1fhSzPCb/GhrVXmNTZwpOr/Dw2pZ72B0FPyOFPiCmBCecPaMefFm0xfS4SgQMNz90ev8ZFz63i3P5tuWpEJgu/K4x6s9tlU/jb9Ka3RD+CEFOCVsfAjES++bEobGEXzSIyqOn0aS9y/wWnLzeN78rXPxQ12mHPbpH55ahuJKpQ9emnHHz4YRz9+hF3xhRiJ0/Glhlqbp1bXMNHJqINGr5J+lWDP364lcm/T0eqo2GopCjET5uGMvkMvvjzl5gFniMRGZph8Mn3B3j4gr5NvlGyfk8ZD/9vKzsP1RDU9VpjrPAE+ffXu3jm611M6ZnGIxf2o009vbt8QY3/rt8ftgMbTY2TJ6Dx3KINjHv5r8SoWZDWO+x9GlPXAJF910FN5+UVzdMcUxAFnlLQ6q55/PNkO2NfcXPXyAZc1LxlTTwwQWvBYVV48tKB/Po/G6POGIDI5hODIyl6XRiYkUiNX+PLbYciFlROq8KL1w6jd7vmW9MJMSVoddw8oSvZeaWm0alIQ89DOyXRIfH0dAATCAAGZibSPtFBbrHb9PlIdg1/OaU3ifGDSZwxA93rxZ2dTfU331Aydx6W5GRip0zhxbhBBH9CG4JyT4A1+WWM7JpS71g+3nII2aLACdd1NCJDliQ+WF/A9eO61Pte0fDR5kJ++96mehcSvsPfz5fbDrFmdzkLbh5VZ7rh3jIPskk1c7S9W/ZrFuIumc7VMZms+3xPWBpOVBHKw0T6Xau6wQcb9vPAeb1Ff6pTiUl63/H0S1c4v4eFx1cE6J1WTwl9A+cRtHLK8iFvCXjLQ2mfrhTocQ7EpgMwrX87NhdU8MLSvKhOG83cHdR0XlyWx/NXD+XpKwbxr692MHdpHrIk1SmqXDaFRJeVuVcPq5VK3hyIWU7Q6hjRJZkklw1PwNwGtqFFZIxN4RbRY0rQCnhixgBmvpQT9Y6h06pwy4SutaxpZaeTuClTiJsyBUPT8G7eTMnXS1i4owJDrh3pieYm6Q1ozF2a16CY2rSv4icbKXiDGhv2VXB9g0dGxtIdxQ0KqePRDCh3B7hsbjaf3DXe1LK92hc0dd+LtsbJoshI4yczxW7B8U2BacPeaOsaovmuJQk+3nyAS4dlRjReQTPgTALZClrdRhOPTHIwZG4N94yuJwXV0bwLUUELRNdCaaIr/gWF34VElBoItXpS7PDxPdBtCoy9k6LEwSzaWMhtk7rx/voC3H41ogbh0cwnugFf/nCIgKpjs8j85qyezBrflQ/WFzBveR7F1X6sioxhGAQ1g3FZqdw8sSuju6bUmfXQlAgxJWh1SJLEH87tFdUi5wiyBBnJLsZnpTbT6ASCk8fQTsk8ffkg7loQeQqG06pw0eD23HVm9zqPkRQF1+DBFCd3QnoxB06ITEVzkzSAnPxS1PJydLcH3ePG8HrRPZ7j/nk5uFMBwtPjohUZFZ4GHMwixBfUuO0tczOM+tIbDaDKp/LrBRt59+bw78duUTBxAI66xknTDRxWGUWWuHFcF575eufRCNnxRFrXANF9156ARm6xed8xwUmiy4SQI189ZCXLXN7XyjNrAvRPN4lOWRzQ58JmGqCgReKvhrcvhwMbIXBCZoPBsd/Ujs8w8peyTR7JlUP/xh1n9+K3Z/dkxa4S5i7NJSe/DN0wMMB0Tot27pYliUpvkLS4kPCPd1i5bmwXfjmmM5XeIFVeFaslVKfqsDaNsU6kCDElaJWcP6A9Ow7W8OLyvIjzai2yhAFM6J7KSdjIEAhOCuf0a8drLhu3v70Bb0Crc8fQZVPQDYM7zsji1ondItrNq/AECerhC/Rob5Ief5DcaeeiuFxILieyKwbZ5Qr9czqRXS5cdMdMTEUrMn48WM3D/9tKx2QXmckuMpOdZCa5iImymePiTYWmC4SI6rd0g+/2VrCvzENmsgsAtawMz/r1yGs24vf2OGwccIxoa5ycVgW7JfR93DCuC4s3H2BXUTVBLXr7+iNE+11H2qNM0EwkZkKnMZD3bb2HPTTRzhub6/hbGQYMva7pxyZomQQ88PLZUJoLWkPtYQykoIfR0iomFv0R9AXIssKEHmlM6JEGhDad8opruHRuNm5/7XtPtPOJLEn41fD7lyRJJLpsJLoaqP1rRoSYErRafnN2D+IcFv7xxXY03SBYTw+cGJtChyQnz145hLve/Y6/fPwDD5zX+6SEhwWC5mZk1xRy7j+TZTuKmbM0l/V7yrEqMrIU6sXUJs7B7ElduWhwRlQd4svcAUy0VNQ3SUVR6Lk6u95jen+7i8+/3hnW5ygakWGRJYZ2TCQjycmeUjcrdpWwt8xDQbmHGJuFjGQXmUnOY0IryUXHZBftEh1Yldq79nOW5oalHUaT3qjrOvPe/IbbStfhWbcO9dAhnIMHkz5sGP1tTjaU1l7cRlPjZJUlZgzNOPrfDqvCW7NGctncbPaVeaLqXXU80Qq6SGz1Bc3M2Dth3xoIHosw7L67dr1eZoKM78Hw4nwDCSnrjKO1MYKfAe/fAGV5EQipY9gNP+xZCV/9Cc5+rNZzDqtCcozd1J002vlE1XXiW6jLshBTglbNTRO6clafNsxftZv/rNuHLEm4/SoWRcKqyGi6Qe928cye2I0ze6djUWTevXkU1726lvsXfs9jF/Wv01oYONpMVIguQUtHkSUm90pncq903H6Vck8AVTNIcFpJdFkb9Rtek1+GIkthznPR3iRjHQ3fii4e0oGnv94Z9ng0IkORJe6b1pvOqTG1zmEYBsU1fvaVedlX5mFfmYfv9pazaON+9pV5Ka72kxZnJzM5JLRibRb2lnnCxhJNemNQhw/3B/l1ny4kXnYpjp49kSyh7+G2bYe4693vGt27RT7sfHU8yTE2/versdz/wRY+3XIAScI0RVGRJXTDgND/Gv1du2wK3YQ1+qmnyyTIHA57V5s2760PAyuc+Qji7vYzoWRXKIp5wu+k81PVeIKQf1csMbbQr+GlDQHe3BxkyXWH59KgB9a8BBN+B46QMC+t8bNiVwlLthfhN5lrojXBSY6xERdlBsHJomWOSiBoQjqnxvDwBX25b1ovlu4o5sEPtzBjSAbd0mMZ3DEx7Iaf6LLx5qyRzHptLb/5z0aevHRgrV3prYWVvLw8ny+2HcITUDEIpdSM757KzRO6MaRjohBXghZNjN0SdVrbiQQ1nW+3FxHnsFDu+WlRlAsHNtzgtV2Ck5Fdklm2syTsuUhFRr8OCWFCCkKbIelxDtLjHAztlGT6WQsrvCGxVe5hdV5pk9QAuGUbSddeGzZfTO6VjtOmmKZkNlTjZJElBmaYf06XzcJTVwziT+4+zM/ezcvL86k5obGz2Q7y8UT6XRsGnDegXb3nEpwEZBmueAdenQrF2yMWVIbVRVH+YNTHnqf9E/+HbI+sR5rgNCbnhZDxhAmaAU/nBLh/fN2/A0OS2PPNK7yvTGPZzmLyi92M7JrMxB5pJLlsvJWzNywqHul84rQq3Dy+a4tdWwkxJfjZ4LAqnNO3LX/8cAvXjulk6qR1hFi7hfnXj+DWN9dz21sbePbKwewqquGe9zaxp9RNUDXQjltNeQIaX2w7xPKdJaTF2Xl8+gBGd6vfnUwgOF04VOUjv8SN26/islnomOJiy/5KOqXEcHafNvzrqx1hUY6fEkWpi1snZbF2d7lpHWRDIsNpVbh9cuNcOq2KTKeUGDqlhARKh0Qn3/xYRLWvthCJNr3RMELixaLUXiAossT860dw6ZzsqJtTxjksPHvVkHqPqfIFeXfNXoJmTbsiIBJBN31IB2GL3lKwueCGLwi+fxPSzs8xdA0rdRhT2GLBHoc0833Sk7pTeN997L3xRjKfe860affBSh+LNxVSUO7BE9RIibExuGMSZ/QKZXoIThOCXtj4NujmtXP3jrHxxEo/tw23keiooydg0INz3b/RRkzlD9N6M7RTEjZL6DdQXO3nrZy9pq+LxARHNwxmtGBnUDHTCX52VPtU4h0N5906rApzrxnGrxds5JIXVpFbXFOvI5phhETVnlIP189fw+PTB3DR4A5NOXSB4KSh68ZRV6Z1e8qP3hQBAqqO3SIzfUgHZgzN4J9f7jA9R0M3SVmqO1pkxuhuKdwwtjOvrNwdlchwWhUuH57BlF5tIn5NfcQ7reE5cESf3mhV5DoXnP06JPDq9cP55StroqpxsspSvSJpf4WXC59bSZU3aNpouSmwKjI3NmEvL8FPY3+Fl2e+3smibVfRUTqTy/RPuUIJmVLoSMiATQpyKHEQbab+HluPM0GWkYEO//gHRf/3BLtnzqTjvHlY24eiyNm5pcxZuovVeWUYhkHgOGOTGLuCVZb55ZjOXDO6E6mxIqrV4qnYF7I/r4Nh7RUmdbbw5Co/j02pu+l4G72E35/VDZTaa6y0ODuXDs3gvxsK8DaiVcc1ozs1ebP1pkSIKcHPgk37Kli2o5iDVT68QY3Xs3dzVp+2ZKXXn9Nvs8jMGt+FGXOyG0x/OR5fUOe+DzaTFGNj4mFXG4HgdGF/hZerX1pNUZX/aKrZiQt6v6rz3roCPttykF9NzuL5JblRR1FibBb+cenAqF7z23N6EtQM3li9J6L3c1oVLhnSgYfO79vgsZHSPT3W1MUw2hqAnvYgank5lqTw1EIIbfwYZqqtHkrcAWbMyebTO8eTFFPb3cowDK5+KYdqrxompOqzc48Gh1Xm75cOoKuol2oRbNxXwbUv5+AOqGg67CCVx7iGJ9QraC+VEI8HHzaKjER8JQl0/szB25kqyYd/O5Is0+YP92F5dT67r5pJhxee58ldOu+s2Vfn9RdybdOYszSX+at28/ZNI+nbXvSqatH4q+oVUwB/nmxn7Ctu7hpZj2ueYg1Zq7uSw556+IK+7C71sG5PWRStOmTGZaVy39ReER1/qhBiStBq8QU1/repkDlLcjlQ6cOvakcXEP/4YgdPf7WTHm3juHViN87u29bUaMIwDO545ztTIdXQ4sMX1Lnj7Q2se/CsWrv6AkFLZk+pmwufW0m1V62VymqGO6DhCWrMXZbH5cMzeHftvohukpIUElJvzBoZcVTq2Gsl7j+vN0M6JfGvr3awt9RNQDNqXaOyFOrX1D7RwV1n9uCCCGqyoiHGbuHCQR14f/0+TgwCRZre6FLgisot5J71KM6BA4mfNpXYM844KqyqfEHufOc7Amp0c49uhAq/71/4PS9cPbTW61blllJU5Qv7u0Zi594QiiRhs0g8MWMA5w9o2u9b0Di2H6zmqhdXmza7DmBlt3FCTZuqk1tUw4w5q1j8q3G16ipTrr8OS3oav/3b+3ybOQRvBPsmflXHr+pcOiebD24bQ6+24Y6BghaC1YlpuP04+qUrnN/DwuMrAvROq2NNo2uHzxWORZF59frh/Pa9TXyx9VCtNdmJKLKEVZG4cFAH/nJxf+R6jMBaAkJMCVolpTV+rnoxh71lHtPdM1U3UHWDzQWV3PPeJoau2cu8a4bhtNWuc1i7u5wyd3iTz0gXH5pu8MW2g2JxITgtqPQGuWxudlQpYIYBNX6VDzcW8tD5ffjXlzvxBFRT8wSLLKHIEr3bxfOvywfRJUohdTxT+7Vlar+2bC2sZP7K3fx4sJoav0qMXaF7ehzXjenMwMzERp+/IW4c14VFG/ejmUSoIqkBsNqsXP7n+5D9d1OzbDlVn33Gocf/D+egQcRPm8oHCb1NXxfJ3BPUDL7+sYjSGj8px6VYzV2aG/Z3icbOHUJ/w+PdG53WUH+yqf3actukLHq2jQt7jeDko2o617ycg9fkOqxPjAd1g/3lXu5f+D1PXzG41us+S+/Pt5mYCqn6zukJaFz1Yg4rfz8l7B4raCHEtgW1YTv0RyY5GDK3hntG15G6abHXKaYglAL89BWD2bSvgpeW5/HFtkNHnZUBFDk0f53Xvx03ju9y2kQ0hZgStDoqvUEufG4lhyp99faWOoInoLEmv4yrXlzNgltG14oizV2WG3Yzimbx4Q5ozFmSK8SU4LTg7Zw9VHrMhVRDkVi3XyW32E3O/WewfFcJc5bksmFvOQFVD0Wi7BYuGNie68d2aTC9Nhr6tk/g71GmCjYFPdrEMb57Gst3FOOLsm+T06rwu3N6huqlXC7ip55D/NRz0D0eapYupfLTz3hB9uN11F5IRDP3yMA7a/fyq8ndgVABeE5+WdhYorFzh1BPvsm90qnxqyQ4rfRrn8D0oRktup7h58jXPxbhPuw2ezyRiHG/qvPZloOUuQNH0/0Mw+Cpr3aaCqlIzhnKFNnP5cM7NufHFjSWmBTIHAm7l9d7WFayzOV9rTyzJkD/9BOiU7IVBl0V0dsNzEzk2auGUOEJsDqv9GiD70SXjTFZKRHVtbckhJgStDpufmMdRVXhQqq+xaBf1fnhYBUPLdrC45cMAEJRpSXbi8NuRtEuPnYcqqGkxi+KcAUtGl03eGlFvqkwiDQa8u6avdx7Tk8m9kg7Wiuo6Qay1Dp7sT175WBmvLCKnUU1EZtEOK0KV47IZOaoTmHPyS4X8dOmsaf/GNzzsuGEjZxo5h6fqvNOzr6jYmpvmRubRQ4bZ9R27gGNp06IWAhaHnOW5B6uXTpGNGJckmDB2r3cOikLgOy8Uiq94U5vkZ7TE9B4YUkulw3LbJVzQatg3N2oBeuxqOE99I7noYl23ths4vqYLHdDAAAgAElEQVQnKzDqtqjeMtFlY2q/07+FghBTglbFjwer2LSvopazEES6c6az8Lv9/H5qL5JibFR5gyiShHaCnIp28WG1SJS5A0JMCVo0S3cW4zNJCYo2DezjzQe4ZGjG0f+ur+n16Y7DqvDe7DHMfnM9a3eX4Q1qpv2nAKyKhCxJ3DapG7+aklXveQ9V+Uy/t2jnntIKNwf//Gd0n589fhe6pS9ItXd8o7VzV3UDVdOF7XULprDCy7YDVWGPRyXGgzrzV+0+KqZeWp5vmjIYzTmLqv1sLqhs1vRbQeM4VOXj0dWJPKQ5ScOLdNy6Z/fdtVN3MxNkfA+eUP8mW6D9EEhpXPuJ0x0hpgStipeX5xM8QUhFvRu3bh+zJ3ZD1Q3MNtCiXXxI1G9VLBC0BFbtKjGtc4pmseQOaCzZUVRLTLV2nDaF+dcPZ8PecuYuy2Pp9mKsiox+eP6QJAnDMLhiREeuG9OZzGRXg+cMaLqpKIta+CBh69IVyWEnOeBA2ibDCX/iaO3cLbIkhFQLZ3+FN4oopIYl7gcssduQLDUYhoShxqNWDaSk5pjo336w2tSeIBqBLwG7imqEmGpBaLrB2zl7+NdXO5k5siMJkxchvXYOBNxRnEUCRyJc+mqzjbOlI8SUoNXgCags3lQY5rwX7W7cyyvymT2xG/FOi6kIinbxoeq6qCcQtHhKasyLj6ONhpTVhBu2tHYkSWJop2TmXZNMSY2f9XvKqfQGscgSyTE2RnVNwWGNvPA+zmEx3ciJdu5xOawkX3M1AL2rfAR/+BaoPadFa+feJl5E2Fs6br95Q95aYtzix5a8AltyNqAhKceuW8MAa/wmDN3Ja1tKuaL35aZRqbBzNjBHqLpBTR1jE5x8thZWcv/CLdgUiQU3j6J7m8MRqGsWwZvTIVADRgMbwbIVXClw/ScQ17b5B91CEWJK0GoorGii1Jga/+GmpCFXsO2Hqms9H+3iI9ZuoX1C3e42AkFLoK5oQ7TREOvPPGqRGmvnnL4/bVHRu108AZMarGjnngEZxyIAbeIdDOmYRHZeadh5I7Vzd1oVbhgrmvG2dOLqKN4/Isa9eV+QNvU7JMWDJIeLG0kClACSEuDfG5/lo7zF2GzXgCf8+o9G4CuyhEu4+f1kDM3A90Mp7g2H0KoDoIPssuDonULMkHRkR/1Le09A5amvdvLf9QX8bmpPLh2aWdt6PHM43LIUvnoYtn8W6j+lemufxOoKqe4Bl8EZD0FMatN/0NMIIaYErYZqX9C0F0FjFoM1fpVki41bJ3XjgYXfh6U/Rbr4cFhlZo37f/buPD6q6nz8+OfcO1sme0IgLCEhIYDs+yaL4oai1hVcUXAB22ppq21/3dS21tZq+61aEFFEoVhwF1dUBEEUEARkJywhLElIyL7Men9/zABJZpLMJGELz/v1yotk5t65Z5Lh3Pvcc87zdDnrayQI0SE2Al0joG5SOBdLCkiOtZ26Rp4n2kRZGdstic+25wVM9wu174m06Ewbm17rsWlj09l8sDjodM5Q0rl7DYObBqc07U2J06ZLm8igCVE0ayRxY27k2NLZ2Dp0ILp3FOiK8m3lVGyvIHlS4E2Aak81WSVZaO2ehZIfg2ENfM0QA3wFdA5hmqsIzuv0ULbiIBWrD2N4DYw6CUYc2aWUfLQPe/8kYi7pjCk+sC/+Ynsef3xvK0O7JPDpz8fUv5Y7IR0mvgYVhfD9fNj2HlQX+wIreyL0uxX63AxWKc4NEkyJViTSagq6ziDsaXke48Tdsyv7JPO7d7cE3S6Uiw/DQFLBinPC1X3bM3vFnoC6SeFcLEVYdG4YeP6slzqV7h+TzqqsgqAFV0PpeyKtJi7MqH23eExmEvGRFqpcVSHXETvOZtK4bkBHmbJ8DkiItDAmsw1fbM+vs87JIPnmPOzp7Tm6JJ+Ds3PQbTq2NBsRnSPY/fvdOPOd6DadmEExtLupHXqkjtvrxqQXEdXpTcpzbg84XsiFqjWDwanxTXpPhmGwdt8xNuYUU1LlwmbWSY6xcUWvZGLtrf8z6SlzcnTOZtzHqiFIIW8AnL6+u3J9HlVbCki6pw+WFN/UvdySah5fspXtR0r5+419GZUZ4khSZCKMmuH7EvWSYEq0Gu1ibC0yNSbCop9Y32A16fy/8d154qMdQYv/NiTCrDN1VBrx/jodQpzNMttFk9kumh8OlQQ8F+rFUoLdwpC0pl0sidoGpcbTs30Mmw+W4AwzgU2EWeM3V/YIGBHXNMV/7x3G1c+torw6sAZRfSy6omu7KB67tlfjG4uzwv1jMli9p7BWMK5FHECzFBF/YSzxF56sYVbwcQFHPz5Kp3s7EdUzCleRi8PzD7P/6f10+V0XNJOG23Ci7NtQplIMd0zA8RoL8G3K4Ias5Ry4bR6JU6YSfdmlKL3xmSLlDjdvrT/Ii1/tpajSidPtxe0vt2Az6/zhvS1c0SuZ+8ek07vjuVHgNVzeajf5L2zCU1Rdd8ljcAYY1R6OztlM4gP9WLSvgH9/sZs7hnXmX5P6h7V+U4RGginRasRGmBmV2YYvd9S9Gxf6xaBFV9w+rPZI0h0j0sg+VsWCb7NDDqgizDqX9WzHw5d3b85bEuK0mj42g0fe3NSk0ZAIs29amdSQaRlKKeZOGcK1z63icHFVQLmH+kSYdaZcmFbvCGFqYiRvPzCSSS9+S3m1u9FALcKs06tDDK9MGSIXYeeQIWnxpCVGsiuvDLd/GNKS8BWo2vWBPFUe8t/Np+M9HYnu6xvFsCRZSPlxCrse2UXJ6hLix/hukOiaIiJhLZX5l4bdHt1i4oEXHkf7ZiXHXp5L/j//ScLddxF3/fVoEcHXFO8rqOCW2d9QWu0OOPd6DU70Ux9sPszSbblMH5PBzy7NbHV9UNFbu/GUOAICqbUHN/PXL2exq2A/mqaRmZjKo5c8SP/2FwDgdXrJeu57Puqss3jacLq2jQ7y6qIlSDAlWpVpY9L5dm9hk6fGKKW4c0RgMc3fTbiApGgLzyzdBVBvgU6LrqEU3Dkilf93ZY9W16mL1u3K3sn858usoDVqGmLRFT07xHDLUJnS2pJibGbe++ko7pq7ll15ZUH7teOO17H6+WWZ3D+m4Vovme2iWfrzMcxdtY/532bj9Rq11lEp5QuiEiMtTB+bwcQhKed9YpFzjVKKV6cO5apnV3KswomHSkxRO1GqdlBeubsSr8tLzKDao026TSe6bzTlW8tPBFMe3FgSvsF77PKgxb3rYzNrvDR5CHFRNrjsMmIuu4zKDRsonDuXguf/Q/wttxB/x+2YEhJO7JNdWMGPnl9FucPd6JRUr+HLxDv7q71UOj38dsIFIbftbOcpc1K1vTBgal+Zo4Ipb/6GJy7/Bdf0uBinx83ag5uw6idnwiggSlO8MqYbdgmkTikJpkSrMrRLAu1ibGQXVoS9JsCsK4anJ9IpPvgC2fvHZHDDwE78b+0B5n69H4fbg6YU1U4PmqawmDTuHJ7KHcNT6RAn2fvEuWfe6v0UVTjonxLLztwyqlyNXzDZTBpd20Uxb8oQzLqG0+PErJnlRkILiY0w8/YDI1mVVcDsFXtYl12EWVd4vAYaCk0Lv44V+JJc/Gp8D2Zc2o1Pt+by5c58CsudWHRFu1gb1w/oyMDO8fJ3PIclRVtZ8tNR3PLiN+RV54GhA7Wz93nKPZiiTCg98O9sijVRlV07i5tXVXBhu0pW59sb7R90pbCaNWbdMYgRGYm1nrMPHIh94EAce/dxbN489oy/kpgrryRxyt0YnVJ8I6chBFI1Vbk8zP82m14dYvjRgI6h73gWq1hzJOjje4/lAHBdT98oYYSmM7bL0IDtTG6D8hWHsPc8uUaqyl3FJ/s+YcexHRQ5iogyR5EWk8aE9AkkRiQGvIZonARTolVRSvHa1KFMeHYlZWGsCdA1RdtoK8/eMqDB7dpEWfnpuEweuKgr24+UUlTpZNG6HGIjzDx6TS8sJrl7K849hmHwz8928eHmI7zxwEiSY2w89elO5n+TjVIEHRGJMOt43Q6u7VzJoDEV3PbxDeSU5eD11yWJt8UzsdtEJnafSJI96XS/pVZF0xRjuiUxplsSh4qr2OCvY2UxaSRFWRmREV4dq5osJo1r+nXgmn4dWrjV4myQHGvj45+N4bnVDubvC5LtNkrHXe7G8BgBAZW7xI0pqvZlotkw+Ef5DL7QhzHTdAu5bl/mwJpBT4RZx2sYjO+dzIPjMunatv6Mb9b0LrT/0+Mk/ewhiv77X/bfdjurhk6gNKp/0ECqYttySte9i6vwIJolAnPbdGJHTsTWybeer8rl4R+f7uTa/h1a/EZAhcPNhgNFFFe60JQi3m5mYGr8KZ3+Wv7tkaAJJ9ITUtCUxs8/fIJre1zCgI69iLMFH31yHirDU+LgsMrn1W2v8l7We2hKo9JdeWIbq27l3xv+zaiOo5jaZyr9kvqdsvfUGikjWPozv8GDBxvffffdaWyOEC1jd14Zk178lrIqF65Gbm1ZTRod4mz87/4RtIsJP63zvK/3sedoBX++rndTm3taKaXWG4Yx+Ey3o7mkf2oZHq/BH97bwg8HS5g3ZQiJNVLlHi+E/eJXezlYVIXD7cVi0kiOsTFlZGcKnS+x+MASlNlOZd06JPhO0IZhMKrjKP486s/EWAIXrgtRU2von87Gvml30W7u/PhOKlwVtR73VHrYMWMHne7tROzQkwkcPNUedj2yi3Y3tSNh7Mnpd8owWL8/BzOAbmET3fhf6p/IcUZR7fIQZzczPD2RmwenNCnzo7eykiv+/hm7HYH3+kvXvkPJmjdJvPwn2LoMROkmqvatx5GztVYiKbtFZ96UoQztkhDwGk2RlV/Gy6v28873BzFrGsaJ27QKr2EwaUgKU0Z2oXNiy6Z9NwyDQ79dRX13hXcX7GfmmoWs2r+eoxXHuDhjGE+N/xVJkbXft7Lp7PhRKf9v2x9xeVy4jfoLJysUVt3K9H7Tmdp7qoxM19BQ3yQjU6JVymwXzaczxjBzeRaL1uWgIKC2SqRVx6xp3DUyjfvGpBNlbdp/h3YxtqCFMIU42zncHn6+aCNFFS4W3jcsoNin3WJi0pDOJ9L7G4aBUgqX18WML2ewNnct1UoFFnQ8/voeBwArD61k4pKJzL9yvoxSCXEGtI9sj9sbeBGt23XaXteWwwsOo9m0Wtn8zAlm4kbG1do+3uPlRC/hcdKPLfTLuQPu/gA6Dmt2O3eXesjxWqibbcHrqKB41X9JvGoG9u4jTzxu7zoMe9fax61yenjxqz3NDqY8XoM/vLuFtzccxO01cHsNqoOk01vwbTYL1xxg6qgu/OqK7i0XgDSSdCazTRr/mvBbALIKs3nog7/w2BfP8Z9rH6213XcRW/jzlhdxeB2NHtLAoNpTzezNs3Ebbqb1ndb09p9HJJgSrVZStJVHr+nFr8f34MPNR/hoyxGKKpxoSrH9SCkPX9adO0ekYmrmwuq2MVbyShvvpIQ4m1Q43Eybv55Iqx5ypjalfOtz/vD1H1h7ZC3VnuqQjuXyusityGXKJ1NYdM0iIs2RzW2+ECIMUZYoLk65mKXZS09MxT0u6aok9Eid3EW5OPOdaBEaMQNjSJmWwu7f7Mbr9NL96e5EmOGO0lJe2uBkwWYXy+/2/z92VcL862HGD2BrXnryHbml6EGK3DsO7cBwO7F3G9HoaxgQdhKdujxeg/vnf8fqrMJGk224PAZgMO/r/RSUOXjqpr4tElDlVzow8CWSaEzXxFQm9h7Pgo3v134N0zH+0nY2Dq8zYJ+ilUUUfFoQtLZYlbuKlza/RJ82fRjZYWTAvqI2CaZEq2cz69w4qBM3DjqZKvinCzcQE2FudiBV7fKwI7eMfQUVzPlqL9E2E307xdGzg0xnEmevYxVOpsxbR4920Txxfe+w/h+szV3LsgPLggZSDZ2cPYaHIxVHmLN5DjMGSQFIIU63u3vdzYqDK6gKMpKcMDah1nS+WrxQsLSAThOSuLGsgreDbeNxwcbXYfj0ZrWxrNqNJ8jUfE9VKZo9BqWFtj6p0hFeXci6nvhwG6uzCsOqL1nl8vDB5iOktYnkJxd3Det45Q43PxwsYdPBYjblFLP5YAll1S5eNdlJcgVun1WYzRd7vuHaHuNoH9OWw6V5vLf9CwZ2qF0L7v345bgJfA+h1Bar9lQzc+NMCaZCIMGUOC/17hjLlsMltQKscOQcq2Te6v38b+0BlIJyh4enPt2BSfNdlKYkRPDARRlc2bu91GYRZ5XDxVXc+fIaLuuZzK/Hhz8l5ZUtrwS9GAvl5Oz0Olm8czE/GfATzFr46ymEEE3Xq00vusR2YdexXQ2um6mrzZVtKPj4KLcNtZLgrWeUxlUJq5+FYdN8ufWbKMKsowXZX4+IwVtZiuH1hBRQmZ1VlH72GdaMDCwpKShz6P1NQbmDBWsO4AwyIhVKAoznl2Ux5cI07Jbgl9hOt5eduWUnAqdNB4vJOVbFBe2j6dspjit6JfOr8T1IS7RTsTaXkg/3YjhrtyXSYmfj4e3MWbeYUkc5MdYoLs0Ywe8u/vHJ4ygXH8WvxK3VDqbCqS2249gOskuzSY0JLBkjTpJgSpyXeneIZdmO/Cbt+/raAzz+/lY8huEf3vdxeQxcHl+ntSuvnN+9s4WnP93FomnD6023LkRzVbs8FFU6qXR6iLaaSIi01DvSlJVfzl1z13LXyNRGaxEFk1eRx7rcdQGPh3Ny9hgelh1YxhVpV4R9fCFE88y8ZCY3LbmJouoiPEZooy7RqVb0TDvWD/JgXANJmqpL4MC3kNr4VDyv16DM4UbXFJEW/cRNnY7xEUFjMWvHHiiTmcpd3xDZY1Sjr99OOSl56xMce/fizs3FnJKCNT0dS0a6L8BKT8fapQuaPfDc/PraA0Gn1tWXAKNq95oTwRT4Ysn3Nx7mlqGdMQyD/YWVbMopZqM/cNpxpIyUhAj6dYqjX0ocd41Mo1u76KDZgO0D2lLywd6Ax9tHJzHruseDvvcCUzHFplLWRW7FqwJH+cKqLeb18PqO1/nN0N8EPZbwkWBKnJeirDrfHyjiV29uwun20ibKyrD0RMb1aBt0vvZxL63cyzNLd4ZUsLDS6cHhqubq51bxwYOjJKASLWrLoRJeWrmXj7fkoimFpnzz/E26xu3DOjN5ZBoda9Q723ywmHte/Y5HrujOxMEpTTrmlzlfBh3JCufkXOmu5N2sdyWYEuIMSIxI5PUJrzPlkykUVBU0uu5RAWkuN0+MgMvmOvnlMEv9GxteKNhVbzDl9nj5fHses5bvYfOhEkya4nhC6XE92jJtbDpD/KnGK+pM09OskcSNup1jn72A0nRsXQagNBPV+zdSfWBzrWx+kVad+2+4kJR+NwPgdThw7t+Pc88eHHv2UrZsGc45L+HMzsaUmIglI+NEoGVKT2fuykIc7qYnwKh0enjy4x18+MMRNuUUE20z0y8lln6d4hjfuwe9O8aGnPBKs+hEjmxPxeojGA3U9XLh5uuYjbyRuJQcSy5mw4RLuXCpwIA5nNpibsPNjmM7Qmrr+UyCKXHeMAyDpdt8Hfn2I6W4PAaLvzt44vnX1x7AYtK4e2QXJo9IJT6y9klj+c58nl66k+oQCpke5zEMyqrcTJr9LcseHovVJFP+RPMcKani3le/Y+/RCpxuL5665S3cXl75eh/zVu/nkh5t+eek/qzPLuKh17/nyRv6cHmv5CYf+1j1sRMZ+moKt/BnQVVBk9sghGie5Mhk3rr2Ld7c9Savbn2Vclc5la6KE9PzlGFgMwyS3R7yvF5+WVTM0HQTV3cz8bdVTi5IqmeNpdcNjrKgT72z4SCPLdmG2+M9kVm35syOz7bnsSqrgKQoKxN6t2fx+pyAc23M0BvQIuMp+WYRBR88jbJEYG3XlZgRk2ptp1CMr9HPaVYrtu7dsXXvXms7w+3GdfAgjr17cezZQ9WG7znwzidUdroG9Nrn/3ASYACUVrm4ZUgKz0zsR9vo8Euu1BR7RRdcRypx7CuBINcf39t38ESnOXjxUqX7+mcX9U/jDLe2WLmzvFntPx9IMCXOC063l58v2siXO/ODFiAFX+r0CqeHmcuzeO2b/SyaNpyubU8WwXvy4x31BlINzaP2GAbFlU4+2ZLLj/q3jqrs4szYe7ScG2etprSeRdrHOf3Zpb7Ykc8lz6ygyulm5h2DGJ7evOr2wVIrQ/gnZ4+3eYvDhRDNYzfbmdxrMnf2vJM1uWv46v37KHRXoBkGbT0eLq+oopfTSZrXe+JC8fGLbAycXc4vR1iDv6hmAmtggd7nl+3m+S+zGrwRaRi+EZ3sY5Us/i4Hbz01UKN6XUxUr4vrfR2bWeOuEalBp8zVpUwmLGlpWNLSiB43DoCSI6WYX/gGh6N2XxduAgyrWWNwWkKzAykApSnaTO7JscU7qd5xrNb6qZXRG3i6w6s4tSBZKuph72pHmRSl60sDaouVbS6j3U3tam0v2VcbJ8GUaPW8XoPpC9azek9BSKNKDrcXp9vJ9TNXs+Sno0hrE8nWwyUcKKwIun0o86grnB5mLd8jwZRossJyB5Nmf0txlYsGaq3X4nB7OVRcRc8OMQxOjW92G+KscZg1My5v7RN3uCfnGKtkuxTibKCUYnj74QyP7A5ZnzW4bdcEjUm9zDy71kmftoHBiqE0VHyXWo/9b+2BRgOpuqrdXkweFxaTCacRejILi0mjV4dYZlzWLeR96jLrNYvynhRuAgyv4XutlqJMGgm39sCxq4jSFQdxHihjmy2LZ8IMpCC82mK60smIC3997flGginR6j23bDff7CkM2pk3NKJU4XBz65xvWfmri5m7ap//bn9t4cyj3l9YwfYjpVzQXi4kRfieW5ZFcaUzaCDVWIap/QUVLN2Wx1V92jerDUPbD0VXOi5qn7zDOTlH6BGMSxnXrHYIIVrY8AfgwGpwBr9peNwfx1qZvzn4xbuntIr82R8S+yMz9mHDKHN5eez9rUHXGDfWZ7l1M22jLJRWu3G4vEHCm9oizBq9O8byypShzQpi2kRZgmbxCzcBhtdrEGNr2UtspRS27gnYuifgLqrmwU/+jsMZ+LdoqETFcQ3VFtPMJ39/Zs3MrT1ubdH30RpJMCVaNafby5yV+4LWimhsRMlrQGm1i8+357M+uyjotKpw5lFrSrHlUIkEUyJs1S4Pi7/LwRXkMxjKyGil08MLy/c0O5jqkdCDTtGdyCrOCngu1JOzFy/XZV7XrHYIIVpY+sVgiQoIpvbPiK71c0qsRvXvg5zDTBGocQ9hy+1E/tPP4C4s5ONL7kQR2OeEmhWvrNrDX6/vw/sbD7N6byFAQKATadGxW0zcO7oLU0d1afZoUJzdQt9OcazPLqr1eDgJMDQFl17Qrtl1LBuyx8jmoOdIwOOhlKg4rsHaYn5psWlkxme2ePtbGwmmRKv2ydZcjCC38kMdUapweHhhxZ4TC2brCmcetdtjUFodem0PIY77YPORoKl6wxkZ3ZVfRlZ+Wa11gE0xtfdU/vztn8Mv/AloaFza+VJiLHJDQYiziqbB6Ifh80d9NaPC3l9HHzWdBHsCCZMnU7VzJwte3U5VnRtA4fRZDreHL3fkM2/qUHJLqlnwbTbf7i2kpMqF1azRMS6C24alMrprG7QGsvCGa/rYDGb87/uA836oCTBsZp37xqS3WHuCmb9tfsB063BKVITCptuY3q95RZjPFxJMidbD44bqYt+JwBoDtljmfLU3aCAUzojS9iOlxNmDF/wLZx61rvkWx4rzT1m1i8+355FX6sDp9hJtM9E/JY7+KXEhFc1dsulwsz/HHq/Bsh35zQ6mrki7gle2vMK+kn1hFf4E36L3nwz4SbOOL4Q4RYbe55vqt/MTCHKzpF6mCLhtEdhP3kg5FN+RMn0f1Ek2E06f5TXgs215ACTH2nj4iu6N7NEyxvVo60vPHqTPbSwBBkDbaCsDO8c1uE1zrTy0MqBOWDglKhpjM9m4KfMmLul8SYu1uTWTYEqc+45shm9mwta3fUUxlA5eF1hj2F/6f0BgkBPOiJLFpBFjM5NXGpgSOpx51LqmkRzT/Mw+4tyxM7eMl1buZcnmw+hK4XB78XgNLCYNXVO0jbbywEUZXNuvIxGW+j+LhRWBnz0I73Ps8hgcq3A2+b0cZ9EtvHTFS9zywS0UVBUE3B2tT4QpglmXziIlumk1roQQp5hScMMceO/HsP2DxkeolA4mG0yaD2m1z39FlU5MuqLO8sqws+I5PF5cHm+LJnNojK4pZt85iDteXhNW4gwAu0XnhTsHhXSTrDkqXIFr28ItUVEfm8nGpG6T+MXgXzS7necLuU0uzl3FOTB7DLx8OfywGDwOcDt8JwCPCyoL6y2uW3NEqTEOl5fCCkfQaVY151FX7voGr6saw+Omas93FH05t9a2SsHozKSmvFNxjjEMg/98mcWP/rOKt78/RLXLV1vF7fXliXK4vVQ6PewvrOTxJdsY98xyco6FP7UmnM9xS0qwJfDGNW/QM7EnEaYINFX/qcRuspNoS2T+lfPp37b/aWylECJsuhmufxGumwnJfX2jTqpO4GO2+x7vdytMXwldA0cv6ivdEG6fpaDeNOmn0uC0BGbePpAIc2hBn8K3fmvelKH0SD7105iD9bk1S1TUFaxERU0mZcKqW+mf1J9nxj7Dw0MebrBfF7XJyJQ4Nx3dCXOvgOpSMOrvlK24cBE4RS+cESUDg7tHpvHiV/sodwROawplHrVF17hjWGi1L8S57+mlO5m7an9IdzUrnR6qXR6ueW4VHzw0ik7x9oBtEiOD13UJ53Ns1lW9r9MUsdZYFly1gC0FW5i3dR5fHvgSS41Cl06vkwsSLmBq76mM6TQGkyanGyHOCUpBr+t9X3lb4fv5UJQNzkrfVL7OI6DfLWCrP2iIs5vxBgmows2Kp2vqjBW7HxInYgwAACAASURBVNejHS91d/Cn74rIjmqHy+vFU6dLN+sKTSn6dIrlyev7kNmuedOoQxVljgpYtxpuiQpd6WTGZ9I+sj1dYrpwY7cb6RzT+bS0v7WRs5s495TlwStXQVUxNJIwtaMqYKcR2DmEk5nHpGlc1acD1S4vc1ftCzra1dg8arfXyw0DpcbU+eCDzYd5edW+sKaHHM8cOWn2t3z58EW1gm7X4cNcUrKbtV47VZql1n7hfI51pRh3Qdvmv8E6erfpzdNjn6bEUcLBsoOUucqw6TaSI5NJjkxu8eMJIU6jdr1g/N/C3q1rUpQ/m13tm53h9FlArULj3koXFevzqFyfj6fSBYaBZjNhuyCBqJEdMMWFP42+pNLFhpwiSqtcaEqRGGlhUFo8VpOOp7iYpNlP8+7MWRxsm8rcr/fz+bY8yh1ulIJom4mr+rTn7pFppCae3sK2V3a5ktd3vF5rmnU4JSrAN/V64VULMevB14SL0EkwJc49y/7kSzRRI5BK+78yKl2w72dRRFp8E/Je2uDk4KbfY791NpVEBLxMqJl50hLtdG0bxc8uzWTV7gK255biCjKMXh+bWaN/Shz3vvYdz94ygH4pDS9M9XgNduWVUVTpBMOXqrVbu6hTmmZVtAzDMPj7xzvqDaQaqq3iNaC40snSbbmM72Ch9JNPKf3oI5x79zLm0sv5m2UoBMn3EOrnuHtyNBlJUafibQO+kapYa2zjGwohWj2TrnHXyFRmr9iLo84NyFD7rEiLzrQxGXjKnBR/uJeqLQUopTBq9K/eMhflxw5Tvvow1tQYYq/OwNK+8cBmy6ES5qzcyydbcn2Feg0DlG+6noHitqEpXPntO3S47HIi+vQmE3jyhj48eUOflvj1NNttF9zGop2LAh4PtUSFRbMwsdtECaRaiART4tziKIMf3gJv4FWlx4B/r3Hy29EnpzIlqlIKgq528mlsRCnSojP9Il/1b6tJZ/69w7jz5TXsyi2rdz1WTTazxm+vvIDJI9P4cPMRps5bx/SxGdwzqktAKtfCcgevrz3A3K/343B5Tjzv9RqYTRpTRqZx27BUkqJbbqqWaFnrs4sorCfJQyi1VSqcHp59dRmZXz1L1MUXkXj/fUSNHImyWJi0ZCsLvs0OGsg39jm2W3Smj5Uq9kKI0+eOYam8sGJv0OdCyYoXZdEZEhNB3r834K10gdc37T6Av0907Cnh6MyNJN7ZE1u34FnrHG4PM/63keU7j+J0e/EYRkCwBzDv63286u7NT0dm8KBhnPKEEuHqGNWRvm36sj5vPV5qtz+U+lEouKXHLaewhecXudUtzi2bFvnmcwfxyEgLT692UFx9srPVMLhD/wwbwbOhNcZq1hnf++RUpdgIM29MH8G9o9OJsZmIDJKBTdfAZtLo0zGGl+8awuSRaQBM6Nued39yIR9tOcLUV9dRUO5rk2EYvLAiixF/W8bzy7I4VuGkwumhrNpNWbWbCqeH4koXM5fv4cK/L+PZL3YHrZ0lWphhwMH1sP5VWP0crHsJdi31JTepx5yVe4MWiD5eWyXhsgewdx+JZrGhdBP2rsMCprVkW+PhrY/o+NRTRF90Ecrim9r34LhM4u2W+j7+9bKaNPp2iuXyXjLlTghx+rSNsXHvqC4hJ3GoyYqXBze9T8F/NuAt9wVSoTBcXgrnb8NxoDTgOafby+1z1vDlznyqXB48DZxHXV5waiZmrT7I40u2hd3+0+HRkY8SYQ6cddOYCFME9/e5n/ZRzSviLk6SkSlxbvlhcb3pWgd30LkozcTTqx38ZdzJudOPmBazwZvJD0Y6DixB9w3GbtFZcM+wgMWvVpPOw1d052eXZrJ0ax6vfbOfIyXVON1eomwmBqfGc8+oLkEXoqYk2Fk8bQT/+mwXE55dyT8n9ufz7Xn8b21OQGX3uo7fPZu1fA95pdX85breZ93dslbBWQlb3oRV/4KyXN9jXrcvo5Vm8hW3HHo/DL4HYmqfjH44VEKw83M4tVVMZp0dxxx0r7PULyHSwqJpI7hh5teUllfjCSG1sM2kkZ4Uxct3DUFvwaKWQggRikeu6M7h4io+3ZoX9EZTMDazxu+u6sUlK2JwHwus4bj24Gb++uUsdhXsR9M0MhNTefSSB+nf/gLAH1DN20r73w1D1Zge/8ibm9hyuCSs9axVLg+L1uWQkRTJnSPSQt7vdEiNSeXFy17k/qX3U+muDD5qV4dNt3Fd1+u4v+/9p6GF5w8JpsS5pbKwwaf/dLGVC+dW8LNhJ4Mms/LwmuXv3Of8BRuMTKpoeJGqSVN4DYOpF6bRs0P92YrMusaEvu2Z0De8uztmXeNX43swMqMN0+Z/R7XLi7ueNLLBVLk8vL3hECkJEUwf2zWsY4tGHN0J864GVwU4A+t4nPD1s/DN83D9bOj5oxMPVziCXyyEU1vF4zUoqw5eDLdLm0jmOb/lN1o39pljcXmMoCmILSYNBVzWsx3/uLkftibcGRZCiOZSSvGvSf15eulO5qzchwb1TpE/PtPj6Zv7cVn7OHI/PITSam9b5qhgypu/4YnLf8E1PS7G6XGz9uAmrHrtG6WGx6Bq2zHsfdoAkF1YwSdbcoNO6WtoLSv4zrlPfbqTSUM6n3UZefsm9WXh1Qt5ZMUjHCg9gNPrxGsEvke7yY5SiocGPMRtF9x2BlraukkwJU45j9dg2Y58Xlq5l70FFVQ7PdgsOmmJdu4dnc4lPdq2WHKF3m11ru5m4m+rnFyQdPI17crBa5a/8aZnLLM815KvtaXKq9UaRbD7O/KbBnVifO9kHlz4PZf2TKZ/IwkjmmpQajwerxE0kAqlc/+/z3dzx/A0oqzy37hF5O+Aly8FRzmNZYnE4/AlqXpnGriqfGmC8U2pC6ZmbZXGAipNqXqnxZR+8imRa7/mg7d/yc5SDy+v2ssHm4+gAE1TuL0GNpPGHcNTuXNEKu1jw58CIoQQLUkpxSNX9GDqhV14fd0BXlm1n0qXB5NSGIDL4yUlPoLpF3Xl6r7tsZl1it7P8qU5rWPvsRwArut5KQARms7YLkMDtjMcHspW5JwIpuat3h+0XlUoa1nBV+tq6bZcru7boSV+JS0qPTadt659ix3HdvDa1tf4dP+neAwPmtJwe91kxmcytfdULku9rFb5CtFy5CpMnDKGYTBn5V5mLd+D0+0rWnpcmcPN0TIH2w5vxKxrTBubzrQxGQFJGQJEBF9UWtPjF9kYOLucX46onahBVwaTTMuZaFvLxiF/543KARwprsLp9hIfaWFMtySu6duBCH9Q9eQNfXhgwXre++mFtI0OP+VqY5ZsOhx0ml6onbumFO9sOHjWTT04J1WXwLwJoQVSNbmqYMkMSOwKnQbTPtZGflng+rxwaqt4vAY7c0tZujWXkV3bnAiWXUeOkPvnP5MyayZ6VCQ9o+CZif158oa+FFc5qXJ6iLKaiLNbZEqfEOKskxhl5acXZ/LA2K7kHKukpMqFrikSIi10iKt946dyfX7QYCo9IQVNafz8wye4tsclDOjYizhb8NpOrtwK3CUO3HYTi9blBCTvOb6WNfGqGdi7jzzxuL3rMOxdh9XatsLhYdbyPWdlMHVcj4Qe/HX0X3li1BNUe6pxeV1EmaOk+O5pIMGUOCVcHi8PLvyeFbuONjhP2hdgeXj2iyzWZxcx8/ZBDQ6jl2ZcS8ThHzB7q+vdpmuCxqReZp5d66RP28DXUngYMHQsA+JTG3wPl/dKZsvhUn7y3w38997hJ9q15VAJi9blcOBYJdUuD/F2C8PTE7hhUCdibKGnGZ21Yg+Vztq/m3A690qnhxdW7OWO4amydqq5vl/oX4vXeLr9BZtdLL+7RupddxXGsj9TPeivXH90M7vcyVSZagfy4dRWcbi9vPZNNgvX5uD2erm2Xwemjkwl4te/JmHyZCL69q312haTdkqCfSGEOBV0TZHWpv705YbHwHAGv26Itkby9u3PM3PNQn71yT84WnGMizOG8dT4X5EUWTuDndI1vKVO9lU6gibuCWctK8D2I6UYZ2Fmv7qUUkSYIogIUhJGnBoSTIkWZxgGD7+xieW78kNe6Fnl8rAqq4CfL9rI87cNqNVZOd1elu3IY9G6HHZnp/Clavw1/zjWyvzN9WRd6zQEGgmkjptxSSZbD5Xwlw+2MiA1nv98uYdDRVU43J5aN81W7DrKkx/v4Oq+7fnJxV1Jb6SeT1m1i5xjgYk0wu3cj5Y5OFbhJDFK0qU3mWHA6meDJjYJlm4/6EtkfUXe69O57MpJ/MMb4UsFVUeotVXAv6bAP7f/rfUHeW/9Aa6OHcBTU6cGbCuEEK2J4fH6svbWk20vs00a/5rwWwCyCrN56IO/8NgXz/Gfax+ttZ23rIz9t09miwkYcAeYa990CmctK/iClGqX98TsFSGOk2BKtLhPt+bx2ba8oIFUQ2uBql1evtyZz4c/HOHqvh3Iyi9j0boc3vn+EOltopg4JIWrbh+I+YPr4Yc3wTh552r/jNrD/CmxGtW/D5I8whwJF84I+b1omuLvN/VlzFNf8nqQaQLHHR99e3fjYT7eksvsOwcxOjOp3tctrnRh1jXc3tp338Lt3M26oqTKJcFUc+xfCY6SoE89MtLCU187+PEQC3G2+u9GKl0n9ecXoy7/KXd8tJ1XV+8Pusg6lNoqdXkM8KDxUVwPnG9s5tlbB5z1d0aFEKKplFmrN5Cqq2tiKhN7j2fBxvcDntOio0n77yuUV1Wh/XcT1EkQFM5aVvCtmzLr0veKQDKRUrS4WcuzAqavgW8t0LEv5hA7fCKdfrqAjg+8QvTAq6javebENpVOD3/9cDs3zPyaW+esQdc0Fk8bweLpI7hpUCfsFhNc8kew1p9lr14mG6QMhYxLQt7F6zX49VubcXuNegOpmjxeg0qnh/te+45v99afedCk++qs11Wzcw+FAZg0+W/cLIc3gjt4HbKa6fYbogwXKucbAH5xeTe6J0djaaGkKsdVub18vj2fZ7/Y3aKvK4QQZxOlFKYke9Dnsgqzmb32fxwpzQfgcGke723/goEdegVu7DUwtYumbVIsziDn75prWUMRYdZbLFmWaF1kZEq0qKz8cnbmlgU8Hs5aoNzSau4fk8HtwztjDtZxxXaCu973JQxwlkOQNKABTDZI6gG3LPTVCQrRS6v2sjqrMGgNqMZG2e599Tu+/vU4Yu2B66ji7ZagwVk4iQrANwUyLjL0dVoiiKpiXx2pegRLtx9UtW90y2rSmX/vMCa/vIYduWVh1TQJJYvjCyv2ct+YdN+NBSGEaIWix3ai+L0sDGft/jPSYmfj4e3MWbeYUkc5MdYoLs0Ywe8u/nHtF9DAPrAtmkWnk8VOaqKdXXnltTcJYy2rSVP8qH/HU/Z+xblNzsaiRb3z/aGgqb7DWQukgMMlVcEDqePa94X7l8PCiVB6JCB5wAm6xTf3usfVcN1MMIU+Hc7jNXhhxd6gCTRCybjn8Rq8uT6He0anB+xvM+sMSUvgmzqjV+F07gB9O8WGlfRCBGGJwDdIHzzoqS/dfgDTycW+MTYzi6eN5IUVe5j79T5cHm+9NaiOCzWLo1Lw7veHuG1YaOv+hBDiXGPv24bi97ICHm8fncSs6x5vdH+laUSNOhn8PHBRBr9/Z0utrMIQ+lpWk6a4Z1Ra096MaPUkmBIt6mBRZdBgKqyipYbvdRqVmAE//Q5y1voSCOxeCpoZlAaG2/f9kHtgyL2+0awwLd+ZjyNIIBXqKFuVy8OLK/cydVSXoGtcpo1NZ/PB4iZ37pFWneljM8J+X6KO6A5gjvAV6q1Hfen2a4nrXOtHi0njoUsy+fFFGSzbkc/CNQc4VFxF1tHygOUA4WZxnL1irwRTQohWS5l1oi/uTNmyAxhhjO4DYNKwdovDXGOq4JW92/OHd7cG3byxtayagu7J0XRtGzwFuxASTIkWVV1PGvRwF3qGPDVKKeg8DDr/F5wVUHHUV/vHGgNR7UBv+kf85VX7AgIdCG+Urazazbr9RQztkhDw3JjMJOxWU9BjhJKowKJrjOvRttE2iPqVO9x8WN6X61wuGhqzbCzdPpYoX+AehEnXuLxXMpf3Sia7sIIr/70yYE1huFkcc4oqcbq9DZYREEKIc1n0RZ1w5VVQvbUw9IDKpDAnRZB4a49aD9vMOrPuGMh9r30X1tRrgEirieduHRjWPuL8Imdi0aISIoNfkoa70DMhsglVui2REJ8GbS+A2I7NCqQA9hcGH6kIN+PegSAp0MGXKfD5Wwdga8IFsc2s8fxtA2UxbBPllVbz67c2M/gvn/H4Z4dY4h6G22j4d/nHsVYqnMGTkBR5bNzwsYlrnlvF5JfXMO/rfZRWB6bmL61yowcZpQw/i6MW9PWFEKK1UEqRMLE7kUOTfRn+GjndKYuGNS2WpAf6ocyBfenozCSeubkfNnPo5027RWfBPUPpnBg8IYYQICNTooUNT0/g/Y2HAkZbwlkLFGnRGZ6eeLqbHqC+u1fhjLJ5vAYVjuDJDXKOVfLG+oO4veHdJbOZNf5xYz8u7NomrP2Ez47cUm598VvKqt0npqS+pCYwQV+DCeeJ7UJNt19lWPh39VVsyDmZXn1ddtGJumMzLu1GSoLvRGwxacFW9jUpRa9VRqWEEK2c0hRx12RgH5xM+aqDVG4qQOkKw2uA4Z/V7zWwpscRPaYT1ozYBktHTOjbgeRYG398byt7jpbjcntpKFGvgcFdc9cxeWQq08ZkEGmVy2YRSD4VokWN753Mb9/ZEvS5UNcCGcDVfdufhtY2LCLInS0IL+OerimignS+67OLuGvuWiqdboIsMau3PUnRVp66qe9ZEWyeiw4UVnLzC99QVl07wN1hdOZf7huZYXobu2o4DXpN1YaZ77zdeM1zea3Hq/w3E975/hCfbs3j1alDGJSaQFK0NWhmyHCzOCqCf66EEKI1srSPJOHm7sRdk0H17iK8FW7wGmgRJqzpseixoSeXGpSawIcPjeaDzYf5+aKNeBqIpqqcXqrwMnvFXj7YfIT/3TectjG2ercX5yc5G4sWZTXp3DY0hXmr9wdN/d3YWiCzrpg4OAVbPYHM6dStXTSHiqsCHg83415G26haP289XMKdL68JWourPgpIjLLw8c9GEWmV7H1NYRgGU+etDTpSWLFtOY+v+4LHC4uIt3oYkKzxu9FWRnWuv4usNCxs8GZyn+uXeOuZf+I1fOuy7nx5LW9MH0GvDrH0S4ll3f6iWtuF85nSleLqvu2lcK8Q4ryj2UzY+yQ1+3X2HC3nN2/9EFL9SACH20t2YQU3zFrNhw+NJjZCzsPiJAmmRIu7d3Q6i9bl4PLUX7unPjaTzv1jAlOJnwn3je7C2n2FQRNEhDrK5nB5WLr1CHERZtLaROLyeJn88tqggVRDNYYM4GiZgz+8t5V/Tux/qt5yq7bhQDGHS6oDRgLrpiS/zLyJIdmv8OaOQoanKEyq9g7lhg0nJl5yT2C252o8+AL/hv5+lU4Pd768luUT2nDjlk/YYutPVZ00/aF+pswmxT2ju7T8L0gIIc4DXq/B5JfXUuEMfmOtvn7c44X8UgcPv7GJOZMHn4GWi7OVBFOixbWLsTH/nmHcOufbsEZfIsw686YOpUNcROMbnwYjMhKJtpmDBlPQ+Cib3aJz3+gulDs83DhrNV3bRtEjOTpoxsNQagw53F4+3HyER6/uFbQQsGjYnK8Ca4YFS0n+BcP4Im0Yvbvs5QPPx/TW9hKlqnBgJcdowzz3eJZ5B9QajQrl71dVVsHiJ97iuhvH8vSeaKoqnNQVSoreLm0i6dUhtiV+JUIIcd5ZlVVAcaUzoERFKP240+Nlxa6j5JVW006m+wk/CabEKdEvJY43po/gjpfW4HB7Gwyq7BYdi64x/55h9Ol09lwkKqWYcWkmjy/ZFrRwb2OsJo1pYzOwW0z8enwPPt+ex6/e3BwQnIVTY0gpWPxdDvedJaN354pKp5tlO/IDTp4NpSTfYqQzw/2TRl875Lpjmpl3x97GvbeN4eWcYm558duwP1eRVhMv3DEorH2EEEKc9MKKPc07DwPzv8nm4Su6n47minOApIMSp0yvDrGs/s0lPHZtL7q0iSTCohNlNRFh9v1rt+ikJdp59JqerP5/486qQOq4SUNSuLZ/h3qTUdTHbtFZcO8w7Bbf/QqLSaNPx9igmfvCqTFU7fLyyup9YbVFQEGZE5Pe/JTkwYTz98vKryC7sIJ+KXHMmTwYuyW042oKYiJMvH7fcFITI5vcViGEOJ8VlDv4Lrso4PFw+nGH28uCNdmnonniHCUjU+KUirDoTBycwsTBKWw5VMK+ggrKHW4irSa6JEbSu2PMWb2QXinFk9f3wWbWWLzuYKMjCRaThs2sseCeYQFTsXKKKjHrWkDK9XAv6AvKAqeHiYZVuTxoQT5m4aYkDyacv5/ZpHGoqIrUxEhGZbbh7R+P5PH3t7LhQDFewwhYDH08/fmozDY8fm0vOsVLrRMhhGiqI8XVWHUtIKtquOfhkioXHq+BHuzEIs47EkyJ06Z3x1h6dzz7Rp8ao2mKx6/tzWUXJDNrRRbf7S8KuPCNtOroSjF5RBp3jUwjKTowTWtVPVMdw72gd3q8GIZxVgehZ5tomwlPkHJe4aYkDyasv59BrSmvPZJjeP3+EeQcq+S1b/bz6dY8yqpdKKWIjTBzXf+O3Dasc9DPkxBCiPBUOt2+eXp1hHse1pWi2uWRulMCkGBKiJCNymzDqMw2HC6u4r2NhzhYVEWl00NilIVBneO5tGc7zHr9M2frqwsU7gW91aRJIBWmpGhr0DuI4aa5Dyacv59SEGUL/BykJNj53YSe/G5Cz9DflBBCiLBE2UwBa2ch/POwxzDCnv4vWi8JpoQIU4e4CB64qGvY+6UnRQUt2BruBX2arJkJm1nXuGVoCq8GqX8Wakpyha+gdF3h/P0cbi/pSfL3E0KIM6Fzgh1XkGkK4Z6HO8RGoMkUP+EnwZQQp0lStJWRGYks33k04KI81Av6SIvOtLGSya8p7hqRxvxvsgkWEjWWkhxA1xQWXVHpCjwRh/r3G5GeSNtoSacrhBBnQrTNzPheySzZfDig5mCo/fjxsidCHCfBlBCn0f1jMliz71jQVPGhXNADXNWn/aloWquXkmBnXI+2LNuRjyPICGFDbGaNCX3a88mW3Hq3aezvF2nRuV8CYSGEOKPuG5PO0m15QRNKhXIe9ngNbhzU6VQ1T5yDJDW6EKfR8PQEUuLtmJowPSDCrHPv6HRsMk+7yf41qT9d2kSeyJIXCqtJo2f7GJ68oS/3jU5v0jx5k6boGB/BiPTEsPcVQgjRcnp3jOWC9tGYg5TLaEyEWefWoZ2JtplPQcvEuUqCKSFOI6UU8+8dSpzdHFZKVZtZY2TXRH52SeYpbF3rZzPrvPnASPqnxIVU48lu0RmRkcjC+4ZjMWk8dEkmozLbhBVQ6Zoizm5mwT3DJHGIEEKcBebePYSkKGtYNzZtJo3+KXH8fsIFp7Bl4lwkwZQQp1nbaBtLHhxFx7iIkC7K7RadK3ol88Idg2TBawuIsppYeN9w/m9SfwZ2jsNq0k6MVCk48fPw9AT+c9tA5t415MRooKYpZt4+kPG9k0MKxiLMOh3jIljy4CjaxshaKSGEOBvE2S28/+AourSJDPnG2uhuScybOgRTA1l7xflJ1kwJcQa0j43g0xljePv7g7ywYg+F5U6qXZ4TC2ItukIpxeDUeKaNzWB0ZhsZ1WhBuqa4vFcyl/dKZl9BBct35lNc6UShiI+0MK5HW1ISghfINesa/5zYj1VZHXlh+R7WZReh4MQ6LE35RsASIy1MH5vB9QM7YrdIVyuEEGeTNlFWljw4iiWbDjNrxR6OFFfjdHs4nvDVrCs0pejbKZZpYzIY16Ot3NAUQckZXogzJMKic/uwVG4b2pn12UWs2XeMwnIHFpNGUrSN8b2T6RgXcaab2ep1aRNJlzbhZWZSSjE6M4nRmUkcKq7iky25FJQ5cLg9JEZZGdolgcGp8RIACyHEWcxm1rl5cAo3D05hU04xq7IKOFbhxKQpkqKtXNazHalSjkQ0QoIpIc4wpRSD0xIYnJZwppsimqBjXAT3jJI0uUIIcS7rlxJHv5S4M90McQ6SiZ9CCCGEEEII0QQSTAkhhBBCCCFEE0gwJYQQQgghhBBNIMGUEEIIIYQQQjSBBFNCCCGEEEII0QQSTAkhhBBCCCFEE0gwJYQQQgghhBBNIMGUEEIIIYQQQjSBBFNCCCGEEEII0QTKMIz6n1TqKJB9+pojhDgNUg3DSDrTjWgu6Z+EaJXO+f5J+iYhWqV6+6YGgykhhBBCCCGEEMHJND8hhBBCCCGEaAIJpoQQQgghhBCiCSSYEkIIIYQQQogmkGBKCCGEEEIIIZpAgikhhBBCCCGEaAIJpoQQQgghhBCiCSSYEkIIIYQQQogmkGBKCCGEEEIIIZpAgikhhBBCCCGEaAIJpoQQQgghhBCiCSSYEkIIIYQQQogmkGBKCCGEEEIIIZpAgikhhBBCCCGEaAIJpoQQQgghhBCiCSSYakWUUhcppQ6e7n1PF6VUZ6VUuVJKP9NtEUI037naZymlPlZK3XUmji2EaJ5zqd9RSi1XSt17uo4nmkaCqQb4L9yPf3mVUlU1fr79FB73bqXUqlP1+i1BKWUopbqe4mPsV0pdevxnwzAOGIYRZRiG51QeV4hzlfRZDVM+e5VS28LY5zGl1IKajxmGcaVhGK+2fAuFOPdIvxOcUirNf61kqvP4PKXUX85Uu0TLMzW+yfnLMIyo498rpfYD9xqG8Xnd7ZRSJsMw3KezbUIIUZf0WY0aA7QFTEqpIYZhrDvTDRLiXCf9jjjfychUExwf5lVK/VoplQu8EuwOSc3RG6WUVSn1tFLqgFIqTyn1glIqognHnqKU2q6UKvPfYZ0WfU1T6gAAIABJREFUZJvfKqUK/CM7t9d4vEXaEOR4jymlFiulXvO3a6tSanCN53+jlNrjf26bUur6OvvfV+M9bVNKDVRKzQc6A0v8d7d+VfMuj1JqklLquzqv83Ol1Pun8r0KcS6SPuuEu4D3gI/839dsQy+l1GdKqWP+Y/1WKTUe+C0wyd8PbfJvu1wpda+/fcVKqd41XifJf2e+rf/nq5VSG/3brVZK9W1G+4U4Z0i/E1I771ZKrfIfr0gptU8pdWU927ZXSm1WSj3i/3m5UurPSqmv/e9zqVKqTY3tr/VfjxX7t73A//gUpdSSGtvtVkq9UePnHKVUf//3hlJqun+bYqXUf5RS6lT8Ls5lEkw1XTKQAKQC94ew/d+AbkB/oCvQEfhjE46bD1wNxABTgH8ppQbWaVcb/+vfBbyolOoebhuUUjOVUjPDaNe1wP+AOOB94Pkaz+0BRgOxwOPAAqVUe/9xbgYeAyb739O1QKFhGHcCB4Br/FP7nqpzvCVAd6VUZo3HbgMWhvtehThPnNd9llLKDtwE/Nf/dYtSyuJ/Lhr4HPgE6OA/1heGYXwC/BVY5O+H+tV8TcMwHMDbwK01Hp4IrDAMI18pNQCYC0wDEoHZwPtKKWt97RSilTmv+50QDQN2+tvzFPBy3YBFKdUFWAE8bxjGP2o8dRu+99cWsAAP+7fvBrwOzACS8N1AWuLv81YAo5VSmlKqg3+/Ef790oEoYHONY1wNDAH64uvfrmjm+219DMOQrxC+gP3Apf7vLwKcgK3G83cDq+rsY+D7j6iACiCjxnMjgH31HCvgtRpo17vAz2q0yw1E1nh+MfCHxtrg3/dgGL8PA+jq//4x4PMaz/UEqhrYdyPwI//3nx5vf0O/c//Paf7jmvw/LwD+6P8+EygD7OH+vuVLvlrjl/RZAce9AziKb3q7DSgBrvc/dyvwfT37PQYsqPPYcnxTmQAuBfbUeO5rYLL/+1nAn+vsuxMYe6Y/H/IlX6fiS/qdWsdMo8Y1S43H5wF/qfEesmo8Z/fvk+z/eTnwT//v9dY6r7Mc+H2Nn38MfOL//g/A4hrPacAh4CL/zznAQOAW4EVgLdADX2D2fp2/zag6v6ffnOnP2dn2JWummu6oYRjVIW6bhO8/yPoaNxsUEHZWOv/w76P47ppo/tf9ocYmRYZhVNT4ORvfndYWa0M9cmt8XwnYlH9+tFJqMvALfB0L+O56HB+KTsE3ctUUC4FngD/huzvzrmEYlf7pNafyvQpxLjrf+6y78F1cuAG3Uuot/2Pv0Lx+6EvArpQaBuThu5v9jv+5VOAupdSDNba34Ht/QpwPzud+5/j6MHON74//7Krx84nrJ/81DPiuk467HcgC3gxyjLrXXsf364DvPR1/Xa9SKgffKBv4RqcuwhfErgCKgbH4AscVIR5D+Mk0v6Yz6vxcge8/IABKqeQazxUAVUAvwzDi/F+xRo1Fm6HwTw15C3gaaGcYRhy+oduaw8HxSqnIGj93Bg63VBvCpZRKBeYAPwUS/W3eUqPNOUBGPbvX/R3X9RmQ5J/beysnp/idkfcqxFnuvO2zlFKdgHHAHUqpXP/6jZuAq/xrDHKA9Hp2b7AfMnzZRRfj64NuBT4wDKPM/3QO8ESN9scZhmE3DOP1cN+DEOeo87bfAY7gC5rS6jzehRqBTgge87droQq9NMxhfDdzAF8mU3w3jQ75HzoeTI32f78CXzA1lsBgSjRCgqmWswnopZTqr5Sy4fvwA747AvgCin+pk4uSOyqlGpp3qpRStppf+O5oWvFNVXH777xcHmTfx5VSFqXUaHxzXd9oYhtaQiS+zvSo/5hTgN41nn8JeFgpNUj5dPUHYOC7y1vfBQ6GYbiAN4B/4JuT/Zn/8TP1XoU4l5xPfdadwC6gO76Ro/747lgfxB8AAe2VUjOUb/F5tH+kCXz9UJpSqqHz5UJgEr47yAtrPD4HmK6UGubv3yKVUhP8a7SEOB+dN/2O/0bLW8ATSqlEpZRZKXUrvqUQH4fxUi7gZnzXU6810hcdtxiYoJS6RCllBn4JOIDV/udXABcDEYZhHARWAuPxre38Poy2CSSYajGGYezCN93sc2A3ULf2wa/xDdN+q5Qq9W/XnfqNxHd3pO7XQ/j+kxThm9r2fp39cv3PHca3yHq6YRg7wm2D8mWveaHhd904wzC24ZuK9w2+i5I++NYUHH/+DeAJfBcgZfjmNSf4n34S+L3yZZB5uJ5DLMS3ZuENo3bK1XB/30KcV86zPusuYKZhGLk1v4AXgLv8I0mXAdf427Mb34UG+G7YABQqpTYEe3HDMNbgu+PegRoXSYZhfAfchy8hT5H/vdxdTxuFaPXOs34HfOuYjuFL6JCPb5bOBMMw8hrYJ4BhGE7gBqAdMLexgMowjJ341ok+h29U6xp8Cb2c/ud3AeX4gigMwygF9gJfG1LLM2z/n737jo+qShs4/jt3SiY9oYTeIy2ICihNBFx7WXtBZcG6ll3L7rp91y2u7666q+u6CoqiggUromtHQZAmICC9hgCBFNIzfe59/7hDSDIzmZmQhIQ8389nIJnbzp2Z3DnPPec8RxlGtJ5UQgghhBBCCCHqk5YpIYQQQgghhGgECaaEEEIIIYQQohEkmBJCCCGEEEKIRpBgSgghhBBCCCEaQYKpY6SUekkp9XDw5wlKqW0tdFxDKZXdxPusOZeW3LalKKV+q5SadbzLIcTxJNesY9/2WCileiulquKYL0aIE45ch45925Yidafo2kUwpZTKVUq5gl9gBcEPb5NP4GoYxhLDMKKm31ZKTVdK1U8H2mSUUouUUrc11/6PVXOff/AYk5RS+2s/ZxjGI4ZhtNrXRYgj5JrVOgWvK4ZS6ldxbJOrlDrnyO+GYeQZhpEi6YdFayfXodZF6k6tV7sIpoIuDc5gPQIYBfy+/gpKKWuLl0oIIcKTa1brMw1zzpgfHe+CCNFC5DokRBTtKZgCwDCMA5iTKg6Dmibfe5RSOzAnkEMpdYlSal1wsthlSqnhR7ZXSp2mlFqrlKpUSs0DHLWW1YnolVK9lFLvKqWKlFKHlVJPK6WGYE4UOTZ4t6csuG6CUupxpVRe8A7QDKVUYq19PaiUOqiUyldK3dLY81dKvaWUOqSUKldKfa2Uyqm3Siel1OfB81uslOpTa9vBwWUlSqltSqlrG1uOemXKVUr9Qim1IViuecqcxRylVKZS6sPga1ga/LlnrW07KKVmB1+XUqXUfKVUMuZ73D34Glcppborpf6klJob3O5jpdRP6pVjvVLqyuY8VyHiJdes1nHNCl5XrgbuAU5SSo2qt/x2pdSWYDk2K6VGKKXmAL2BD4Kv3S+VUn2D76FVKXWdUmp1vf08oJRaEPy5wddYiJYi16HWcR2qVyapO7US7S6YUkr1Ai4Cvqv19OXAaGCoUuo04EXgx0BHYCawIPgHawfmA3OADsBbwFURjmMBPgT2An2BHsAbhmFsAe4Elge7emQEN/k7MBA4FcgOrv/H4L4uAH4BnAucBJxD430c3EcWsBZz5u/abgT+CnQC1h1ZHvwj+xx4Lbjt9cAzSqmhEc6/TCl1Zhzluha4AOgHDAemB5/XgNlAH8xKiQt4utZ2c4AkICdYricMw6gGLgTyg69ximEY+fWO9zowpVZ5hwaP8b94z1WI5iTXrFZzzboSqMJ8DT/FbKU6su01wJ8wW6zSgB8Chw3DmArkEby7bxjGo/X2+QEwSCl1Uq3nbgiWGRp4jYVoSXIdajXXofqk7tQaGIZxwj+AXMwvwTLMP9BngMTgMgM4u9a6zwJ/rbf9NmAicBaQD6hay5YBDwd/ngTsD/48FigCrGHKMx1YWut3BVQDA2o9NxbYE/z5ReDvtZYNDJY7O8L5LgJui+F1yQjuJz34+0uYF60jy1OAANALuA5YUm/7mcBDtbZ9OMb3o/755wI31fr9UWBGhG1PBUqDP3cDdCAzzHo170Wt5/4EzA3+nBp8zfsEf/8b8GLw5wbPVR7yaO6HXLMivi7H5ZoVXP8L4Mngz1OCr5Ut+PunwH0NvJfn1Pq9b/AcrMHf5wJ/DP58ElCJWclp8DWWhzya+yHXoYivi9SdpO5U59Ge+rlebhjGFxGW7av1cx9gmlLqp7WeswPdMf94DhjBT0jQ3gj77AXsNQzDH0PZOmN+ea5RSh15TgFHsj11B9bEcMwGBe/4/A24JnhMPbioE1Ae/LnmtTAMo0opVRI8fh9g9JGm9SAr5t2NpnCo1s/O4DFRSiUBT2DeeckMLk8NnksvoMQwjNJ4D2YYRqVS6n+Yd07+gVk5uj24uLnPVYhYyDWrlVyzgnflJwO/CT71PvAccDHmHfdewK549xv0GvBP4C+YrVLzDcNwKqWyaPg1FqIlyHWolVyHIpC6UyvQnoKphtT+A98H/M0wjL/VX0kpNRHooZRStS4KvQn/JboP6K2Usoa5KBj1fi/GbILNMcx+yfUdxPzwH9E78qk06AbgMsym7lwgHSjFvPgcUXMcZWbt6YB5R2kfsNgwjHMbeezG+jkwCBhtGMYhpdSpmN0MVLBMHZRSGYZhlNXbrv5rHM7rwENKqa8x+29/FXz+eJ2rELGSa9ZRLXHNmorZbeaDWpU2B2ZXv/nBYw2IsG20a9HnQOfgtW0K8EDw+WivsRDHm1yHjpK6UzuuO7W7MVMxeB64Uyk1WpmSlVIXK6VSgeWAH7hXKWULDrg7I8J+VmH+If89uA+HUmp8cFkB0DPYjxjDMPTgcZ8I3o1EKdVDKXV+cP03gelKqaHBuw0PxXAe1uAxjzxsmM2zHuAw5t2cR8Jsd5FS6sxg2f4KrDAMYx9mH+aBSqmpwXO3KaVOV+ag0OaUinmxLFNKdaDWuRuGcRCzH/MzyhxsaVNKnRVcXAB0VEqlN7DvjzDvpPwFmBd8H+D4nasQjSHXrOa/Zk0D/ozZVebI46rgsTsCs4BfKKVGBt+DbHV0AHoB0D/Sjg3D8GGOIXkMswL2efD5aK+xEK2JXIek7tRu604STNVjGMZqzCbLpzHvPOwkOKDPMAwv5iDk6Zjpca8D3o2wnwBwKeaAyDxgf3B9gC+BTcAhpVRx8LlfBY+1QilVgdk/f1BwXx8DTwa32xn8P5pnMf+QjjxmA69gNnMfADYDK8Js9xrmH10JMBK4KViGSuA8zKbdfMym5X8ACeEOrswsMBNiKGc0TwKJmHegVgCf1Fs+FfABW4FC4P5gebdi3j3ZrcwBnd3r79gwDA/m+3cORwd8x32uQhxPcs1q3muWUmoMZsXhv4ZhHKr1WBA8tymGYbyF2Q3oNcwxT/MxAyOA/wN+H7wO/SLCub+GeR16q97d+IivsRCtiVyHpO7UnutOqm4XViGEEEIIIYQQsZCWKSGEEEIIIYRoBAmmhBBCCCGEEKIRJJgSQgghhBBCiEaQYEoIIYQQQgghGqHBeaY6depk9O3bt4WKIoRoCWvWrCk2DKPz8S7HsZLrkxAnnhPh+iTXJiFOPA1dmxoMpvr27cvq1aubp1RCiONCKdWoWeBbG7k+CXHiORGuT3JtEuLE09C1Sbr5CSGEEEIIIUQjSDAlhBBCCCGEEI0gwZQQQgghhBBCNIIEU0IIIYQQQgjRCBJMCSGEEEIIIUQjSDAlhBBCCCGEEI0gwZQQQgghhBBCNIIEU0IIIYQQQgjRCBJMCSGEEEIIIUQjWI93AYRojfZW7OXt7W+zq2wX1b5q0uxpDO04lKsHXk3npM7Hu3hCCAFAlcfP/O/2s2RHMaVOHwlWjR4ZiVwzqhcjemeglDreRRRtgdcJG9+GHZ+DswQsNkjvAafeCL3HgnyOhIhIgikhalmyfwkzN8xka8lWAnoAv+GvWbYsfxmzvp/F6G6jufOUOxneefhxLKkQoj3bV+Lk6a928v66A2hK4fQGapZpChasz6dzagJ3TRzANaN6YdGkMizCKN8PS/4F6183AyZvda2FCja+B0kdYfx9MHKaGWQJIeqQYEoIwDAMnljzBK9vfR13wB12Ha/uBWDpgaV8e+hbfnX6r7h60NUtWUwhhGDVnhJufmkVbm+AgBG6XDfA6Q2w97CTP3+wmY82HmTmTaNItFtavrCi9dq/BuZcDj4n6P4wKxjgq4byavj8D7DxHbjxTUhIbfGiCtGayZgpISBqIFWbgYE74OYf3/6D93a81wKlE0II07p9ZUx7cRXVnvCBVH0uX4CVu83gyx/Qm7+Aom049D28fCl4KiIEUvX4nHBgDbxyGfi9zV8+IdoQaZkS7d7X+78OG0iVLiml+NNivIVeLA4LaSPT6HJ1FyzJ5t1dd8DNIysf4eROJ5OdmX08ii6EaEdc3gDTXlyFyxcIWVa9eREV387Hd3g/mj0RW1Z/0sddi6NnDh6/zvp9ZTz5xQ5+cf6g41By0ar4vfDK5WarUy19n6zE6YM996WQbDe7hc5a62XuBh+LpidDwAMFm+Hz38OFjx6PkgvRKknLlGj3ZqyfERJIFX9czKG3DtH12q4MfWYo/f/QH+9hL7mP56L7j97d9ek+Xt70cksXWQjRDn2wPh9fmNalilXvUbLwedLHXEvPn8ylx12zSR1xEa4dK2vWcfl0XlqWi8cfGoiJdmbrB+AP3wsjYMC/VzbQ8uR3wdpX6o2tEqJ9k2BKtGt7yvewvXR7necCrgCF8wvpflN3UoenoqwKe2c7ve7uhbfYS/my8qPrGgE+zv2YKm9VSxddCNGOGIbBM4t21kk0AaB7qilb+iodzr2LpEHj0OwOlMVKUvZoMiffErKPTzYeaslii9Zo6RMQ4TvrwXF2Hl/moczdUB9SDTa81TxlE6INkmBKtGtvbnuTgF63cuLc4UT36aSNTKvzvMVhIXV4KlWb6n4JaWh8kvtJs5dVCNF+bcqvoLDSE/K858BWDL+XpIFjo+6j2htg1pI9zVE80VYU7zQfEYzqbmFSXyuPLwv9rNXwVcOK/zZD4YRom2TMlGjXdpXtqpP+HCBQFcCaYkVZQlMJW9OtuPa66jznCrjIL9kJa16GDfPAeRgwILED5FwBp0wBR1rIvoQQIlb7SpxoYeb6Cbgq0JLSUFpsmfr2lzqbumiiLSnNNdOb+10RV/nL5ATGv1jNfaPtkfdTcaDpyyZEGyXBlGjXqn2h/b4tKRb8VX6MgBESUPnL/VhTjv7ZZAYC3FNazhV7HwNLQsiAXg6uh88fgmFXwuTfmZMgCiFEnKq9AXQjtOuVJTEN3VmBoQdiCqg8fsno1655q4CG00AOy7JwyUArf1/qZUjnCB2Y/A20XAnRzkg3P3Hi0HVz9va5V8FTp8E/B8N/RsHbt5jzaYSRYk8JeS4pOwllVVSsqajzfMAdoHJDJclDkwHo4/PxzoGDXFFZhV33hwZSYKaT9btg/RswY7yZjlYIIeKUkmAJ2zKV0GMwymrDuX15TPtx2GSuqXYtIQWIPoHznyc5eH6tlwMVEQIvq6NpyyVEGyYtU6Lt03VY8Sx88wT4XKEDa0t2wbaPIa07nP0HyLm8ZtGQDkP49tC3+HRfzXOWJAtZl2eRPzcfzaGRMjQFX6mP/Dn52DrYyBiXQZbfz5z8AtJ0nZiqJkYAXKUw+yK4YxF0HNAUZy6EaCeys1Lw66GtSlpCMhln3kjJ5zNQmgVHv9NQmhV37jrceRtCklAM6JzcUkUWrVGngRCIPk9UdgeN63JsPLXKy8lZYe67d+jfDIUTom2SYEq0bX4PzJsKuUvMVqBwDN1cdngnzL8T8r+Dc/4ESnHtoGuZs2VOyCadL+qMJdnCoXmH8BZ60RI10kak0evHvdBsGv/KLySlXiAVdY4OMAO9166Fn6yGMHeZhRAinOysVPp3SmHzwYqQZWlnXImWnEn58nkUf/g4yp5IQpds0sZeV2e9ZLuF2ydIJbhdy+gN3U6FfSuirvrHiQnM2eALXWBPgXH3NkPhhGibJJgSbZeuw1vTYc/iiHNmhPC5YNVzYEuCSb8iY2s+gwusbOjoDen50GFiBzpM7BCyi2yvl0FeH7Ywuz8yR8dvJySEP76hQ8VB2LcSeo+JrcxCCAHcOWkAv3l3A9We0LmiUnImk5IzucHtbVaNswdnNVfxRFtx5v3wzm0hvThy70+t83uvdA337yMkTxr6w+YqnRBtjoyZEm3Xd3Ng96I6gVTfJyvJeqySau/Rft6z1nqZ9FKt8Uw+J8aSf1Fw73Xk//JX3NHvJhxx9P++qbwSa5iB4BDjHB0+Jyx7KubjCSEEwAU5XUm2WxvVqJ1os3DPpGysFvnab/dOOg+SOoKK/7PgUortQy4Aa4Qbhu1ZyW74+FfmeO1HusPfusHjA+GDB6B4x/EunWhGclUVbZNhwNJ/he3aF3UGdwCfm4weh+j/8UdMvPKn3DviXhyW6AGVZhhcXF0dsUk3pjk6MMxEGR6Z6FcIETu7VeP1O8aYAVUc2yXazBap2yb0a7ayiTZEs8C0BZCQRizJKGpYE/H3GcsDej4Pr3gYbwxjr9qFwq3w4gXwzFj49gWoPAjearN+UlUA370CM86E5882M/yKE44EU6Jt2rcKqorCLoqldUhpkODdghYwW6ymDp3KAyMfwGFxoDVwty5V11ENZ5XlL5MT+M8qL0XVDaQgttigurDhHQkhRD0DOqfw7t3jyEy2k2CN/hWeaLdwySnd+ff1p6JknKY4IrMv3LYQUrJiy8xnS4ZBF5B60/u8cembHHYdZurHU9lfub/Zi9qq5X4Ds86GvBVmLxk9zBgz3W8uO7DGDLp2fNHy5RTNSoIp0TatmR0x4URsrUOYEdWmd2t+vWHIDbx84cuc1+c87Jo9pKUq0ZpIhpaA0hoealh7jo4Gj+2LPGmiEEJEMrBLKgt/NpF7JmeTmWTDUS+osmqKBKvGmP4deObGETx61XDp3idCdcqGe1bBpN9AcmczsURtms3sztdrNFw1C66eDVY7qfZU/jXpX1za/1Ju/OhGvsr7KuIhdhRU8ut3NjDmkYUMe+hThv/pUyb840v++dk2CitiHOvcWh36Hl69xmyFijJ3Vw2fE968CfavbtaiiZYlCShE21SaS0MXr5hmcPe7oKzuXbWhHYfy2MTHKPeU89Gej8gtz6XSW0lGQgYDOwzk/G7jsP9zaNTi/XmSgxEzq/j52Aj9yvVAsIuFEELEr9rrp9rjxxuchNeqgV83AyndMDijTwfuP2cgp/fNlBYpEVlihpmQYty9sGsh5C6FqkIziErvCTlXhJ3KQynFTUNv4uTOJ/Pg4gdZW7iWe0fci00zUzOtzSvlofc3saOwEl9AJ1Cro0aF289zX+9m5te7GT+gI3+9fBg9M5Na6oybhmHAvJtC5peMKauvzwVv3AA/2wqa3OQ4EUgwJdqmKK06Mc3gDqFzUgWlJ6QzZfCU0AWGAY50cBY3ePyoc3QoDVK6NLgPIYQI57mvd/HPz7ajGwa+QN2bSn7d/H35rsOs27eKU3tlMGvaKJLs8nUvGqBpcNK55iMOp3Q+hTcveZPfLP0Nt356K4+e9Sjr9hjcP28dbl/kru6e4E2Ar7cXc/FTS3nt9tHkdE8/plNoUXkroDr8UIOoWX3BbM3atTDu11u0ThISi7bJkRF1lagzuKMguVN8x1UKxtwN1sSoq/5xYkKdrII1LHYYOR2sDbSaCSFEGI9+spUnPt+Bx6+HBFK1GYDTG2DN3lKueGYZTq+/5Qop2pUMRwb//cF/mdBjApe9/kvufWNtg4FUbQHDoNzlY8pzK9hXEmGuyNZo2VPgDV/emLL6eqvgmyebqXCipUkwJdqmPuOiDpqt3ToUlj3ZnLwwCsMwMGqnQh85DQj9osi9P5Vz+h+9+3tkjo6apv0jlAZn3BH1uEIIUds7a/Yz+5tcXL7QeaYi8fh1couruWvu2mYsmWjvNKXxo6G3UJl3PeHi9urNizj48v3k/etq9j89lYI3H8K9f1PN8iqPn5++/l0LlvgY+L2w4zMiDTWIedz2vpXgLm/68okWJ+3+om0aOR2W/DPqahFncAezT3iYJnbDMFi1p4Tnvt7N8t2HcfkCaChSHFYuPrkbt07ox4BTpsCGefEnkbA6zDk+MvvEt50Qol3TdYNHPtoSMZCq3ryIim/n4zu8H82eiC2rP+njrsXRMwePX2fVnhI2HihnWI821JVKtCmfbDyEQgPqfkYrVr1H+cq36XjePTj6jUBZrLj2rMG1YyWOnjkA6AZsPVjBzsJKsrNSw+y9FXGXgWY1s/RFENO4bYsdqovNoQOiTZNgSrRNqV1gwNmw/RNq3x2KeQZ3q8PsrqdZ6jz91bZCfv/eRkqdXlzeQM2eA5hdEd5cvY931u5ncNereSx1P9lFC9EsMWbxsSRAx5PgyueOPmcYwXFbymwpk4HiQogwFu8owu0PH0jFUln1+nVeWLqHJ66L3hovRGM8u2gX1d66n1HdU03Z0lfpeNH9JA0aV/N8UvZokrJH11nXrxu8uDSXR648uUXK22h+T9QJj2Mbt62BzNV1QpBgSrQqm/LL+WJzAQUVHnTDoHNqApMHZ3Far4zQjFSTf4OxcyFKb8TFyJYII2+u89SrK/by1/9tbrCvt1838OsGG/ZXcFngJmZn2hhjfB2S0SeEPRm6nwZT5pl3o7Z/ZvaXzlsOykJNQDjgBzD+XugzXgIrIUSNmYt3Ue0JDaZirawGDIOPvj/In36YQ3qirUXKLNqPA2Uu9hSHfg96DmzF8HtJGjg26j78usH8dQdafzDlSI8pCIqe1dcX0/hv0fpJMNUWGIY5N4HfY6bTtpxYb5svoPPhhnyeXbSLvBInXr9OMCEVCnhh6R6yUhP48cQBXHFaDxw2C4ZhcHjBSnwbe9D15HxUIErf5NrsyfCj9yG5Y81TH32fHzWQqs0AXBY7t7qnM/+SWzhpy3/M/s9o5uR8YHYjBOh2Coy/Hwaeb/azXvAxEW3ZAAAgAElEQVRTs3vgkUyCRq1j7vjMTE2b1MFsweozDiGE2LA//NiKeCqrdqvG9oJKTu/boamLJ9q5wgo3dqtWk6XviICrAi0pDVWvF0gkLl8AX0DH1prnRUtIhdTuUJ7X4GpRs/ompEpW3xPEiVUrP9EUboUVz8CGN807GMpi3g3J7GtWzodfYwYGbVil28f02d+y5WAFTm/oXdcjGalyDzv5ywebeWlZLq/ePArf43/HtXEjvWZ8hipbA2/fYs7d1FBQZUsyu/dNWwBdj975cnr9/PzNDWEDqYbGIRAs230rkvno3g+gbB/s+BScpWbJEzPNrohH5uhY9Tx89gdzfquIDLOVq7wa5lwJV86EoZfF8EoKIU5kkcZKxVVZNaDCFWEMqRDHwOsPfyPSkpiG7qzA0AMxfUYtSrX+YEopGH8ffP4H80Z3AyKO264ZatCKz1PETIKp1qhsH7w1DQo2QcAPxpFBjsH/S/fAp7+FT35tdgmb9Js22SXM7QtwzYzl7C6ujnghrs3lC7C7sIpLH/6QFyqKGTRnDpaUZOhyIdy7Dla/CKtmQsBnDgzVA2YrntIgsaM5MeHwa827QbXM/+5A2JcvlnEIBrC7qIqthyoY3LUXnH5b+MJvfj+GQKoevwve/bE5M720UAnRrlk1FTYVelyVVWW2TgnR1NISbRhhhg8n9BiMstpwbl9O8uAzo+5HNwwSbbG1Yh1Xp1wHn/0+5OmYx20bBoyY1lylEy1MgqnWpmAzzL4IPBVgNJD+9sgYnWX/gaJtcPWLIckUWrsH315PbphAqqHWIJ9uUGxJ5OGRU3ktpVarXGoXmPwbOOtB2L0Iyvaak+I50qDzEOh1RtiA0zAMZi7eHdIqFs+gWa9f54Ule3jsmlPCn6jPDfPvDgmkYpop3e+Cd26FBza3yYBZCNE00hNtFFeFjtOIp7Lq181xqEI0tX6dktHDRFNaQjIZZ95IyeczUJoFR7/TUJoVd+463HkbyJx8S531B3dNCx0f3RolpJo3aL/5d9TWqRC2JBh1c52hBqJtk2CqNanIh5cuBndp7Nv4nOY4m49/CRdHTxXeWhRUuPl0U0FIIBVLa5APjbV5ZeFTqFqscNI5MZcj97CTwsrQroHxjEPQDfjf9wcjB1Ob50fcNqaZ0t0VsOdr6D8xalmEECema0f1YtbSPSHXzHgqqx2T7Qzq0srTTos2yWGzcO2oXry6cm9IC2raGVeiJWdSvnwexR8+jrInktAlm7Sx19VZL9lu4c5JA1qy2Mdm4q+gZDds+SD2gMqWBP0nwbkPN2fJRAuTYKo1+fR3IRO4xdR64XPCuldhxI/MZAdtwNwVe6l/7yme1iCfrjdJCtXDVR6sFgX1ujQ3ZtBsQDewaGHuqC194miyiXoeHGfn0W883H26nQxHhLtx3irz7pcEU0K0W1PH9uGFpXvCLoulsppkt3DnxP5t466/aJNuHt+X11flEW4y25ScyaTkTG5we6UUF+R0babSNQOl4PIZkJJljonWA+b49rDrWs2bvcOvh4sfl7FSJxgJploLZwls/V/Yrn0xtV74vbD8v3XnMArHMMwkFhb7ces2ZhgGryzfG5L1J57WoIAO7313gD9eOhRHmP7Vht9PoKKCQHk5enk5gZpHRc3PekU5+VVW9KRR5utRS7yDZg0D/jB/I51TE+iQbK95ZFFKdsnukMDxiNozpT98tiPyAXYvMsfPnWCZHIUQsemWnsiY/h1ZurOIQJghptEqq4YBl5/Ws4HlBh6/ToJVk4BLNEqfjsmcn9OVzzYdwh3DOOjaEm0Wfn7ewLY3pk/T4LyHzalWVs6EdXPxGwqvP0CSzWrWs3SfGUSNuQs6DzreJRbNQGpmrcXaOREngYup9cIImEkOLnwUEuvNW+Aqhe9ehRXPQmX+0efTesCYe+DUKaHbNKMqj59qT+jM4fG2BuHzsfEv/6BrZRGB8jIzQCorJ1BRge5yYUlNRUtPw5KegSU9HUtamvl/Rjq2Ht2xDB1Kdy0VvvWAv+6dtHgHzVqUYmCXFEqcPnYUVlJS7aWk2ktaxXb+FbCQ0sC2Mc2UrlnMcXRJktJYiPbqsWuGc+GTSyhxesMO9o/EYdP4z5TTSEmo+5VfXOXhjVV5vLx8L8VVnpqbPj0zk7jjrH5cflrPkG3KnF7yy9y4fH5SEmz0yEwMWUe0X49fcwpTylxsOlAeV0B1fk4Xbh7frxlL1sw6DoCLHoVz/8yKLz9k5948po/tY2b17T2mzWdeFg2TK2BrsWFexExvMbdeaDbY/RXkXGH+7nPDRw/C92+agVr9Pr3l++DLv8DCP8EpU8xAzNpAhb6JVHsCWC0Kv163NhBva5CmwN8vm7Rup6GlBYOmDDNo0lJSUDE0ow/366i1n4O/bnAX76DZ03pnMD3cF8GhFJhthQYytsc0U7pSZhcCIUS7lZXq4K07x3LtzOWUOX0h19BwHDaNR644mXOGHp3Pptrj59fvbOCzzQUANb0Ejuwtr8TJIx9t5eH/beHG0X349QWD+HZvKc8t3s2yXYexWzWUCk6BGNC5cFhXbpvQn2E90pv8nEXbYrdqvHb7aO59/Tu+3l6Mxx8g0sfUZlFoSjGmfwe+P1BOudNHelIbn1DalsjmxBEUds+Bk4ce79KIFiLBVGvhKmlwcUytF7ofnIfNn90VZjKL4h1HJ5EN50iAtf4NOLQBfrQAEhpqRzk2Tq+fnYWV+PyhV9d4W4MMm42uPzybtE6Nv+Njt2rcOLo3L36z55gGzd4VadCsI8PsnhdF1JnSA74WbT0UQrRO/Tun8PF9Z/HQgo18saUQTREyR54WTIHev1MKD106lNH9j2YNK6n2cvWMZRwodYV0ta7tSIbTV1fsZe6KvTXHMQBvvX6GH6w/yKebChjaPZUXpp1ORlLz35QTrVeC1cLMqaNYv6+MWUt289nmAmwWDSPYnKqUwjAMrju9F9PH9aN3xyT++uFm7pizmlduPYMEa9vKTFxfYYVHsma2MxJMtRZGw83hMbVeYATHRPng1auhaKs5PioWfhcc2givXwdT3z/msTmGYXCgzMWWg5VsOVjB1kMVbDlYycFyFwM6J4MiZIxqvK1BAd0gK+3YL1hTx/bhpWW5oQUitkGzCTYLkwZlhV+Y3tNs5q9seH6pqDOldzsFLG38jp0Qokl0Tk3gmRtHUlLtZd63eby5ej+lTi8ubwBNKS4e3o3bJvRjcNe689u4fQGmPL+CfYed+GJo1QJi6qoVMAxcvgAb9pdz4b+XsOAnZ0plUnBKrwz+c8MIypxeVueWUu7yYbUoMpPsnNGvQ53xzr+7aAh3v7qWX769gSevO7VNj9srqvIwtHuYuaXECUuCqdbCkQ5VBQ2uErX1QrOaLSHfvQqHvq8TSMWUFTDggQNrzW6Bp94Qc9GdXj/bDlWy9VAwcDpYyZZDFSTZLQzplsbgrmmcn9OV+88ZSP9OyVgtGn/+YBNzVzQ+haqm4LycLiTZj/0j3DMziatH9uTdtftx+eIbNOuwafzxkqHhs/iB2T1v3E/hy782fqZ0ewqMvz+ucgkhTnwdku3cNSmbuyZlA/D55gJeX5XH4xGmaXhx6R72FleHDaQamt8vVr6AQVGlhxueX8EHPz0zbHIg0f5kJNnrdDMNR9MUT15/KlOeX8Hjn23jwfMHt1Dpml5RpYes1AaGZIgTjgRTrcXgi2H5XjOgiSBa64XhcVK+voT04qdRYSruMWUF9Dlh6ZNhg6narU1bD1aw5ZAZOOWXuxjQOSUYOKVyQU5XBndLo0Ny5K4e08f15bWVjU+hmmC1cPuE/g2uE4+/XDaM/aUuVu05HHNA5bBp3DMpm8tP69HwiqfeYI5LqyfmmdI1i/n5EEKIBnRNc3CwPHy37oBu8MLSPWFbmmKZ36+2hgIvv26wv9TF/O8OcP0ZvZv8HMWJy2GzMOtHo7jy2WX0zExiShv9/BRVSje/9kaCqdZA12HY1WZq8ygitV4YKPyZp+Jd/zWGYx8qzDsbU1ZAgPI8XHtXs03LDrY0mV30jrQ2De6axpBuR1ub+nVKxmaJL51pn47JjOqbyao9JSGtU9FYNEW/TskM79l0Y4gsmuKFaaP4/fyNvPfdAXTdiNgNxmHTzFTolwzlxtF9ou88MQMm/QYWPxr/TOnWRLjwMeniJ4SIqmu6g4KK8MHU19uLcPtCk9jEM78fxBZ4uXwBnl28i+tO79Wmu2uJltcxJYHZ00/n2pkr6JbuiNyFvhUrlGCq3ZFg6niqKoK1L8OKZ8BTZSaQqCfW1gtlS8J2+cNkbZ6PsSr84WLNChjwuXn1hSd5r9OPg4FTKufldGVw11Q6pjTdBeK/N4zgwn8vobDSQyDG/vuagnSHjdk3n95k5TjCatH4+1XDuWvSAGZ/k8ubq/ehKVUzHZeuGyTZrdw6oR/XjepFZgMtbyHG3w9leWaij5hnSk+ECb+AU66Lvq4Qov0yDNi3kk4rZvC6fzX6kxqaPQW65MCYO6HHSN79bj/V3tBgKp75/eIJvIoqPXy3r4wRvTOP/fxEu9K/cwozp47gjlfW8MqtZ5DTve1kifT4Azi9fjIS5QZoeyLB1PEQ8MNHv4D1rwGq4Wx7sbA6oM8487Hiv6gwXeeOiCUroAWDW4YncNvVE46tXFFkJNmZf894rn9uOfll7gYzSwHYLYrMZDtv/ngsXdKarz9yn47J/OmHOfz6wsFsyq+gwuXDoik6JNsZ2i0NLdL4qIYoBRf/C9J7weJ/YL7vEZJS2JLMytGFj8GIqcd0LkKIE5hhwLrXzGtKdTHK52SQMqAsuLxoC2z9ENJ70s9/NTAsZBfxzO8XT+Dl9gVYsr1IginRKCP7dOAvlw3j1pdW8+7d4+iekXi8ixST4iovnVISGldPEG2WBFMtze+BOVdC/hrz52NlTYSsIXDdHLPCHiUNd2xZAUGLNQvgMeqS5uCDn07ghSV7eGnZHrx+PeTuaXKCBU0ppo7pw+0T+sfXInQMHDYLI/s0YUVAKZjwMxg53ZykecXTZoukFvwz1H2Q3NlsxRp+LSSkNrg7IUQ7pgdgwb2w6d3Ird2Gbi4r3s7d/JN062T+6p8KHK3oxTO/XzyBl26YWc2EaKyLh3fjQJmTm2d/y1t3jSXN0fpbe2S8VPskwVRLMgx45zY4sCZyq0SsLHazcj7oArhiJliDf7wpnaNuGjUrIEBKw5l3mlJKgpX7zjmJn5ydzZdbC/lgfT5FVR4M3aBjip0LT+7GeUO7YrfGNy6r1UrqAGfeZ2b5K90DrlLzvUzsAJl9QcYYCCEaYhjw4QMNB1L1OPAwxfIVbuw85r++5vl45veLd2L1iFlOhYjR7RP6s6/Exd1z1zL75tPjHp/d0gor3HRuwuEQom2QYKol7V0GOxfWCaRiSll+hGY1u/QBjJgGo+8wK9+1+DufgabeQjManxUQewpk/6BRp3gsLJri3KFdODdKCtUThqZBxwiT/QohRCRb/wffv1UnkIrluyRJebjZ8glL9OGs0IcC8c3vF0/gZdWUpIcWx0wpxUOXDuXOuWv4zbvf89jVw1t1UpOiKmmZao9ad4h/oln2VNi7iEdSlkelNPjR+/DL3XDBIzWBlO50Ur5gAXm33MKunz0DMWT2/uPEBKq9EcZWWRPgpPOi70QIIUTLW/LPRn+XOPByp2VBnefSzriSzLNvpXz5PPb/50b2PzudyrUfknhS3bFRtQMv5/bl6D43RsCPa9dqSr96sc66Vk1xfk7XRp6gEEdZLRpPTTmN7QWV/HvhjuNdnAaZc0xJMNXeSMtUM9lRUEl+uRuXN0Caw8rAFCeddn9FuHmVYk5Zrlkgbzn0HIWh6zhXfUv5++9TuXAhiaedSsbVV9Pz7GfQvnncDNxqzVkV85xGlgQYc7d5LCGEEK1L0TYo3Bx2USzfJZqCMdoWulBCAR1qno9lfj+IfWL1gV1Tyc5KiePEhIgsyW5l1rRRXPmMOQfV1SN7Hu8ihVVU6WFwVxnv3BK8fp1PNh3i040HOVzlRdMUnVMTuPzUHpw1sHOLdjOWYKoJuX0BFqzPZ8aiXRwsd2O1KAzDHAJzUeAr/mzVCNfpIdaU5fhc6KvnUrwqQPmCBVhS00i//HKyfvYA1s61xkqNuctMuV5dRLjgLSKlmeN5Tr819m2EEEK0nO/mhJ1GA2L/LlEYXGFZyozADxtVhGiBV5Ldwp0TpQuzaFpZqQ5euvl0rn/OnINqfHYnAAzDYP3+cl5buZfcw07zJnaijZF9MrhxdJ9mzf5bX2GlhwkndWqx47VHJdVeZi7exasr8zAwqPbUTVr2xZYCHFYLt57Zj+nj+5Jkb/5QR4KpJrJy92Fue2U1Ad3AeSQbXa25dVMsFSjdVzuJUh2xpCwHCBzYhZHpodczz+AYPDj8Sskd4eaPYNY54KkwMzpFoyzgSDO3S5RUtkII0Sod3h0xmILYvksSlJ8+WgGETjt1zBKsGqP6dpAufqJZZGel8vQNI7jn1bW8ettoNuVX8PRXOzlU7sbjD1B7ysrVuSXMWLybcQM68sA5AzmlV0azl0+y+TWvXUVVXD9zBeUuH95A+LpttSdAtSfAvxfu4N3vDvD67WOa/T2RMVNN4KtthUybvYpKt/9oIBVW5Fai2inLG2Lt2pUuv/5V5EDqiE4nwY+/NsdV2ZMbWFGZyzv0hzuXmv8LIYRonbzVDS6O9bukk81LckLk7txKQaLNgs2isMXYXcZh08jpnsbMm0ZKJj/RbMb078jvLh7C5c98w+/mf8+e4mpcvrqBFIDHr+P16yzeVsR1zy3n3bX7m71s5pgpSbzSHA6UubjymWUUV3siBlK1efw6ucXVXPXsMircvqjrHwtpmTpGWw9VcPfctbh9Db+xZUYqfqwkNHArMJaU5SqxQ8RlITL7wE/Xwp7FsPRJyFtmjok6wu+BfhNg/H3Qd4Kk5BZCiNYuKXrPgVi+S84ZOZSns0cwc/EuvssrqzP1hNevM3lQZ24/awDZWcn84q0NfL29CCO4rL5Em4ZuwFUjevKnH+a0+vTVom0zDIMvthTgDxj460dQ4dYH3D6d3773PQlWjYuHd2+2chVVeegkqdGbnGEYTHtxFVVuP0a9t7x68yIqvp2P7/B+NHsitqz+pI+7FkfPHPy6waFyNw+8sY4Xpp/ebOWTYOoY/ePjrbh94QOk2m/wf+wJbOzq5A9n2Tizd/iXPWrKcmsiDLsqvgIqBf0nmY+qQijLM7v+JaRBRp+Y5qUSQgjRSvQcDds/bXB+qejTXySjeo5i8qAsJg/KorDCzf4yFy5vgJQEK306JpGRdLSb4PM/GsWhcjdzVuQyd0UeFW4fVk3h1w2yUhO4fUJ/rhnZi/Sk1j+pqmj7XluZx1dbi8IGUg1VrN0+nZ+/tZ5Te2fSIyPxmMvh9et8uukQM7/exa7C6ppuhuc9sZhp4/rK30QTWrO3lPwyF4F6kVTFqvcoX/k2Hc+7B0e/ESiLFdeeNbh2rMTRMwcAb0Bnyc5iDpa76JZ+7O97OBJMHYOCCjfLdh0O23kv3Bs8Ie+3zN+6M2IwBWbK8jkbIjRHGjqMmNr4AqdkmQ8hhBBt06lTYOFDUVdr8LsEYOjlNT9mpTnIijJIv2u6gwfPH8yD5w/G69dx+QIk2y1YpRVKtCDDMPjvVztxhbmJHUvFWtcNXlmWy28uGnJMZXhm0S5mLN6FrhtU1xvesa/UxT8/285jn27jkuHd+Ovlw1okCcKJbObXu0Pec91TTdnSV+l40f0kDRpX83xS9miSskfX3YEBc5bv5ZcXRBki00jy7h6DuSv2hn0+0hu8ccB0Xhj8GNCIlOVKg8EXm9n2hBBCtE+JGTDkMtj4DhhHKxcxf5doNjhtKtgaP67DbtXqdAsUoqWs2F1CmSv0JkGsFWtvwODVlXn87LyBJFjjnwImoBvc+8Z3fLmlMGxAd8SRZR9uOMj6/eW8+eOxdEhuOMGYCK/C7WPxtqKQ7n2eA1sx/F6SBo4Nv2Et3oDO3BXNF0y1yNWw3OXj5WV7+OXb6/nxnNX8/M11zFi8i6JKT/SNW7FF2wrxhOk/HukNXqEPYak+DJfRiD+ohDQ498+NLaoQQogTxVm/MCdXbwxrAoy9p2nLI0QLeXlZLq4wib7iqViDwVdbC+M+tmEY/O6971m4paDBQKpOuYJJEG54fkXEISGiYYfK3WFv3gRcFWhJaagY50Wt8vib7T1o1pap7QWVPLtoFx99fxBNqTofvgSrxhOfb2fCSZ24e3I2I3q3vXTcFa7w6Wkjv8GKn/ruZa79EYaxhyTVcLalI9uQkAI/mg8ZvY+5zEIIIdq4zoPg6hfhrZvB74p9O2si3DBPvktEm5V7uDrs0Ip4KtYev87+0jj+boJW7inh/XX5YROORUuCsKe4mpmLd3HfOQPjPm57V+3xh82PZklMQ3dWYOiBmN53m0Wj2uPHYYu/RTKaZmuZ+mD9AS57eikL1uXjCfavrs3j1/H4dRZuKeTG51fy3Ne7jvmYHn+AokoP5U4fegwZXo6V1RI++13tN7g+LzZu8P6OBYGxuA0bbiPS4EQFtiToOADuWAzdT2vCkgshhGjTBl0I180BW3L0ViprIthTYOq70PfMlimfEM0g0vQzDdW76vMHDJyeyHO1RTJz8a6wLRsVq96jZOHzpI+5lp4/mUuPu2aTOuIiXDtW1qzj8evMXpZLoAXqpieaVIcVvX4fPyChx2CU1YZz+/KY9uML6KQ4mqcNqVn2+uGGfB58e0PUdOFgpqx0+QI88fkOdMPgzonZcR3L6fWzYF0+MxbvIq/Eic2iYRgGBnB+Tldun9C/WSZqMwyDlITwgVDtNzh5cOgXlx8rv/b/mH/6r+NH9q+4w/4ZCYYHNCu624myWPBmn4flzHux9h4tKcuFEEKEOulcuHctrH4RVs7E7fVixY9V94HFDhYbWB0w5m4YOQ2SOx3vEgtxTFISwldbo9W7arPhI2XJX2DrdrOVtubRy8xynNEbkjvXqXsVVrj5JkzCsXiSIPgCOl9uLeTcoV3iO+l2rlt6Iv5AaDClJSSTceaNlHw+A6VZcPQ7DaVZceeuw523gczJt9RZPzPJ3qhxcrFo8mBqV1EVv3hrfdzNoC5fgCe/2MEpPTMZO6Bj1OMcyaby9Jc7Uero3YraY5g++v4gC7cU0iMzkRk3jSA7KzXS7mLi8gZYvruYRduK+GpbIRUuf0162NpifYOLyOBp/SquvfvfrNy8nXeWbWJXhZ9KLQ33lkQCm0qYPHgNd5zVn1F9MlESVAkhhKgttStM/i2c9UsefvxJ7jlVo5sjAAmp0Ckb+k+GGMcUCNHandwjnW2HKqhft46nYm1LcDDwukegY5U5XUxZHpTthfy1R3/3OiG9Z02g9WnZcDQjC6hbD4tnrFa1J8C8b/MkmIpTcoKVi0/uxvvr8kNSo6edcSVacibly+dR/OHjKHsiCV2ySRt7XZ31HFaNW87s12xlbPJgataSPfjCzEwcS8pKt0/nqS93RA2mDMPgwbc38L8NBxscBKgbZqvXrsIqLnv6G165dTQj+8Q3Niu3uJqvthXy1bYi1uSWkNMjncmDsnj+R6Po2zGJkQ8vxB+muTiWN1hTkN05hclPLAHA6Q2WTQd08zX8YksB3+wspmOynWdvGsmwHulxlV8IIcSJzxmAt6ty+OMPzgfJtCdOUNPH9+X99QcIhLlhH2vFOsluZdzgXmYlrMvQ8AfyVEH5PijbB2V7KTzgwx2muhlvEoSCiradeO14uXVCPz7aeJCAL7SFKiVnMik5kxvc3gCmnNF8Y0WbNJiq9viZ/90B6sdS8TSDrt1byv5SJz0zkyIe59FPtkYNpGozgGpvgGkvrmLBT8bTv3NKxHXdvgArdh9m0bYiFm0rpNobYPKgzlx/ei/+M+U00hPrdu274YxevLJ8b9isftHeYAXsLKoKu21N2Q2z1c3pdXHNjOXMmjaK8dnSVUMIIcRRWw5WclJWqqQsFye0Id3S6NcpmS0HK8Muj1bvSrRp3D6hH5oWpadPQgpkDTEfgO/wFsjdHbJavEkQwjU2iOhyuqczuGsa6/eVhU1A0hCHVeOSU7o3a2r6Jg2m/rfhYNjhPfE0g+qGwWsr8yLmgs8trubFb3LDBiANdSMEqPb6+d17G3n9jjF1tttX4uSrbYUs2lbEqj0lDO6ayuTBWTx9wwhyuqc12L3u/nMGsnBrIXsPV4cEkQ2xoINu4InjU+HyBbj9ldW8dedYcrpLC5UQQgjTpvxyhvUIM6+UECeYX5w3iHteWxvTuPz6bBaNa0f1inu7zCR72GEd8YzVAshIjJR0TDRk7+FqypxekuwWPH495H2IJMGqcVKXVP52xbBmLV+TBlPbCyrDZlqJpxnUFzDYfLAi4vLZy3LDZuqLpRuhYcDavFJ2F1VxoMxVa+yTj4kDs7jitB7869pTyEiKPXpNTrDyxh1juGbGcg6WufCGGSRXn8Oq4Q2YvfnqixYQOr0BfvX2Bj68d0LMZRRCCHFi23igvFmSLQnR2vxgSBfunDiAmYt3x9xDCSDRZmHOraPjquMdcUa/DtgsGv562QLjGavlsGmcPSQr7mO3d6tzS7jr1bX89OxsLhjWlSnPrSC/zB31vU+yWxjeM50Xpp3ebIknjmjSYCrcrNQQfzNomTP8/EtuX4C3Vu/DVy+Yiqcbodevc/4TXzOspzn26d/XnUZO97ToTb4NyEp18OFPz+R3723k002HUIqwd0yS7RZsFo2Te6SxKrc0pHUtloAQYGdhFTsKKjmpy7El1BBCCHFi2HigghtH9znexRCiRdz3g5NItFl48osdePwBGmqoSLBqJFg1Xrl1dKNvOJzaK4MuaQnkHnaGLIt1rJZh0KhWsRNFabWXN77NY963+yh1+gjoBskJFsYN6H7xen4AACAASURBVMRtE/qF7XE1/7sD/PXDzfzz2lOYNMgMRD/86QTeXbufGYt3cbjai8sbqNP1z2HT6NcpmTsnDuDik7thtTR/1+cmDabqjyc6It5m0A37y8n54yd0SXfQNc18ZKU5qPb4wuaaj6cboQF0SEngvbvHR103HqkOG09NOY3Sai/zVu9jzvK9HK724PMbOGwaA7uk8uOJA5g4sBOjH1kYEkjFlV5T15m1dA//uGp4k56DEEKItsUf0PHrOruLqxjUVW6wifZBKcWPJw5gfHYnZn69i882FaApVae1IjnBgk3TmD6uLzeO6UPn1CjzsUU53l2TBvCnBZvDtohEG6ulKThvaJdGtYq1dcVVHh56fxNfbCkIaWyo8pjTG32y8RC9OiTyp0tzGJfdCcMwePKLHbyzdj+v3T6mzrUt0W7hxjF9uGF0b1bvLWXxtiKKqzxUuf0s2VnErGmnk9M9jUSbpcWyYDdpMJWdlUKS3RLS1S+ulJUWxS3j+3HP2dkUVrg5VO7hUIWbggo3Gw+U4wvTjS7ebCrVjZisLVaZyXbunDiAOycOCLt8bV5p2Dso8QSEAR0+3XhIgikhhGhn3L4AH244yMzFu9hTXE1AN1DKrOw9v2Q3U87oTaeUxlcahWhLhvVI5z9TRlDm9PLhhoPkl7lYuqMYh93CrWf24weDs5qsZeKyU3vw7KJd7Ct1xT35bqLNwgPnDmyScrQlucXVXDNjOaVOb8RxTgHDwOULsL2gilte/pbfXzyUVXtKyCtx8t7d4yMGwUopTu/bgVN6ZvDppkM8s2gn5S4/Nzy/At0wA9jzc7pyx1n9Gd6zebtAN2kwdcnwbvz5g01hl8XaDKopxY2j+5DmsJHmsNWZG6pHRiLf5ZVSXS9Yi7cboRGmdaullDm9YZN0xB0QehsOCEvdpXy05yPyKvKo9FaS6chkcIfBnNvnXBxWR2OKLoQQ4jjRdYN/L9zB80t2o6DO96BhmP88/eVOnv5yJz8Y0oV/XHUyqQ4Z7C7ah4wkOzeNMbu5ZqXuYXdxNefndG3SYzhsFt64YyyX/GcJZU5fzEkQEm0WZk07vcFM0ieiokoPV89YxuFqL7FWu90+nT+8v5ERvTN5444xOGwN14lfW7mX//t4K7phUO0xr4lHGl0C1J5v1sHTN4xgcNfmSdLTpMFUqsPGpcO78+53B8JG7bHkgu/VIZFuGeEr++lJtrBjm+LtRpjiaPLptWKm6xAur2P8AWH45zcVb+KFjS+weN9ilFJ4AkfnNEiyJvHwioe54qQrmDp0Kj1SejTyLIQQQrQUf0Dn7lfXsmRHcYODro90H/9iSwEXP1XO23eNJStVbp6J9iUz2U5pXlmz7LtruoOP7pvA1Fmr2FfqDJt07QibxayvvnHHmHaZHOZX72ygzOkLW19tKNmaYcDm/Aq8Ab3BYOr/PtrCy8v34o5hvtmdhdVc+cwyZk8/ndH9G57LtjGafFTWHWf1r/kANcb+UhcjH/6cj78/GLJsVJ/MsDn6a3cjdG5fju5zYwT8uHatpvSrF+usa9UU5ww5frNPZyTZwubIrx0QxiLJHvoBe2njS0z/ZDoL9y7Eq3vrBFIATr8Tp9/JvG3zuPL9K1l2YFljTkEIIUQLMQyDX76zIWogVZvXr3OgzMn1z61o1m7tQrRG6Ym2iInMmkJWqoNP7p/ArB+NYuLAztitGqkOKykJVlIdVhw2jcFdU3nkipPplu6gtBnL0loVVrhZurM4bOtdxar3KFn4POljrqXnT+bS467ZpI64CNeOlXXWe3fN/oj7f/7r3bwSJZCqz+kNcMtL37K9IPwcZceiyZtoTuqSyv9dcTK/ee/7Rs0B4PbpuH06D7y5joIKN9PH96tZlmy3ckrPDFbuKQnZLtZuhBZNccuZ/UK2bynDeqSH7WYYz7gyi1I1WU2OmPX9LGaun4k74I5aBr/ux6/7ue+r+3jq7KcY2z36OC0hhBAtb/nuw3yy8VDYQKqhu7sB3bw5+d+vdkact1GIE1Fmkp0yZ/js0k1FKcW47E6My+5EYYWbnUVVVLr9JNkt9MhIrOnS1yk1gYcWbOLT+ztG7bJ2Ipm7ci/hmlViTbbm8gWY+fVupo3rG5JEorDSzeOfbWvUfLNOb4AH31rP+z+J3ostHs3S3+2KET3RDfjd/O/x+vUGU1ZG4vbp/P2TrXRJc3B+Tlc+2niQf3+xA00p7BYVdj6nWLoRDumWxoDj2G/VYbNw/Rm9eGX53pBkGrEGhHarxh1n9a/5fXn+8rCBVOmSUoo/LcZb6MXisJA2Mo0uV3fBkmz+QbsDbu7/6n7ev/x9uiY3bd9iIYQQx+65xbvDdiWKZSoNr19n7oq9PHDuQGwtkB5YiNYgI8lGmavlWoOyghmnw5k8KIs3u+3j2UW72lUCije/3Rc22Ikn2Vq5y8em/AqG9aibMv31lXlh149pvllg26FKdhZWkZ3VdLFAsw0eumpkTwZ1TeXpL3fyyaZDYdeJFkG6fTo/e3M9vTK3k2i38NuLhzBpYGd+8dZ6Pvr+IK44W74SbRb+elnzzoIci2lj+zF3RR7hBk/FEhD2zEys8+H677r/hgRSxR8XU/RxET1v60nK0BR8pT7y5+ST+3gu/X7XD81qfrH6dB9vbH2D+0fef+wnJoQQoskUVLhZvvtwyPPxTKUR0A0+31zARSd3a/byCtEaZCTZKatu3papePzhkqFc/NQSLj+tB/06JR/v4rSI0ggtg/EkW7NoiqLKusNV/AGd2ctyj2l6Ib9u8OLSPTxy5cmxnk5UzXqraliPdC4a3o0ke+hhYu0z6fEH+MHQLObfM57Jg7JQSvGPq4Zzet8OJNpiL36iTeOZG0dwcs/QScFaWu+OSdw4ujeJjWjyddg0/q/WB2BvxV62lmyts07AFaBwfiHdb+pO6vBUlFVh72yn19298BZ7KV9WXrOuT/cxb9s8fIHWc+ERQggBnwUnga8vnru71d4Ab67e1wylE6J1SnNYcfkCYcfYHw/dMxK5a9IAHlqw6bhmk25JkTId1k62Fo3LZ167Xly6h/fXHWDJjiIWrM/He4wtXn7d4KMweRmORbO3+89YtAunN3wE2eHcu0gaNA7N7kBZrCRlj64zNgjMTBxLdhTX6TNptWjMvvkMrhrRE3twZutIkhMsdEi2M/e20UwenBVxvZb2+4uHMnlw57gCKotS/POaUxjVt0PNc69vfZ1AvQ+lc4cT3aeTNrJuCkiLw0Lq8FSqNlXVeV43dBbtXxT/SQghhGg2xVWesGOP451Ko/7dXSFOZEop0hJtlLtaz03im8f341C5i483hu+pdaJJilC3jSfZmkUpku1W9h6u5vPNBcxYvIt/fb4dV5huz/FeE6uaODFPs+YIr3D7wmbNiCeCBNh6sJJKt6/OnBkWTfHwFSdz7w9OYu7Kvby8bC9ev45FUxgY+PwGJ/dM586JAzh7cBaWMCnVjydNUzw9ZQSPfbqNF77Zg1ZvVujakhMsWDUzW0xJvabTLYe34DfqfigCVQGsKVZUmKyK1nQrrr2uOs+5/C52l+2GPsd4UkIIIZpMuEnqIf6pNPwR9iPEiSojyczo11omsLZZNB6+/GTue+M7zhrYmZSE4zdFT0sY2CWVNXmlIc/Hk2wN4PeXDCEjyV7z+8ItBdw/bx2V7rr13rinF2rkeUXSrO9mWbUPu1XDXy+KjDeCtFk0ypy+sBMQZqU5+Nm5g7j37JM4WO6m3OUjwarRMSWBDsn2MHtrPTRN8asLB3P7Wf2Z920es5buodrjx6ppGIaBN6CT090MCM8ZksX+UhdXz1jG4K6pnB5snaryVYXs15JiwV/lxwgYIQGVv9yPNaXu225gUO4pRwghROuRmWTHZlEhQVW8cyumJ8nkvaJ9yUi0NXtGv3id0a8D4wZ04t9fbOd3Fw9t0WMXVLj5+PuDFFZ68Pp1MpPtjO7XgZF9MkOy5TVWabWXDzfk887aA+wu+n/27ju8qfNs/Pj3nCPJkrxtbPACG7PN3hBmdtIMSAhp9iABMprRvmnfpvmlSdOkfds0zWpCErJJAmQvaDZ7771sDAaD99Y+5/z+EDYYSbZkvMDP57p6tZWOjh4ZnaPnfsZ912CQJb/L/YJJtiZLcHG/zvUCKThRXshPJBTqPTHQzFlTtWgwpQeI/UKNIIP5dzYoMmlxVtJCbWQ7EBdu4p5JPZg1IZPjld6A0KjIxIebiD0lIEzvFM4/rxvE/R9u4sv7xtEl2ozVYPU5n7WHFckgUbmxkuiRJ/eIqQ6Vqm1VdJ5Wv86WhESkKbLlPqAgCIIQsjGZ8RhkGbdaf0AylNFdi1HmojasrSgIbSHWagqYBKEt/fHyPlzy72VcOyyVPl2iGn/BGdB1ndU5Jby+NKcukU1t4gZFgjCjQny4idkTM5k6NAWrKfSQwOlR+Xl3IZ9tPsqanBIm907kwQt7cl73eMb94xcKAywxbizZWphB4e5TslbXykqORjvD8kKyBON7dQr5szakRYOpGKvJ7wbAUCNIl0cjpgOMrMmyRHKMheQYS8BjJvdO5JbR3bjng43MnzmazJhMthVvQ9NP/p0Vq0LilETy5+Ujm+V62fyMcUZixtavxG0xWEiNTG2xzyUIgiCErn9KNKmxFvYX+q5ACLaUhqbD9OFn4zCjIDRdtLVlC/c2VaeIMB6+qBePfb6DhbPGILfQFhRV0/nT59v5cku+3xp1qu6tuWRz2fnrt7t5dWk2C2aNIaWB/mctXdfZcKiMzzYdZfGOY/TtEsXUoSk8N31QvRVkf7q8L3/4bFvINWfDDDJDu8UyMDXG5zmzUWH68DQ+WNv08kJhBoW7x/sGameiRYOpaIuRjE7h7Cuo/0MQ6prJnokRfpf4dVT3TurB9qMVPPn1Lm6d+Gu+zfnWJzV6wuUJKOEKxxccx1XoQrbIRA2NIm1WGvJpWRB1dC7sdmFrfgRBEAQhCPdMyuSxL3b4rTXV2OiuLMGlWV3EMj+hw2mNwr1NdcPIrny8IY9PNx3huhYY6NB1nQfmb+bn3QVBlRCyu1WOlTu48qUVLHpgPF2i/dfMOlhcw+ebj/LF5qOEGWSmDk1h0QPjA04AXD0khZziGl5fluM3oPMnzCDTLd7K67cMC3jMHeel89G6ppcX6hwdxuA030DtTLT4Drh7JmXy2Oc7qDnthyDYCDLcpDB7UmZLN/OsIssS/5o+mCn/WUmUORHVFQdKvs9xcRPjiJsY5+cMJxkkA1dlXoXF0PhohCAIgtC6fjUwiTlLs8kpqgmYbjgQq8nAby/uOIVCBaFWrNVIWTucmYITCdSmDOCOd9ZzkZ99QWdqztJsft5dGFItVlXXqbC7uWnuGn54eGLdjFntPqjPNh8lr9TGlYOS+c+NQ+mfEhXUXquHL+pFrNXI3xbvQdfBFSBdvQRYTAqDUqOZe9sIwhtI0NEtPpzrhqXy6aajQQdptcxGmb9NHdhs+8RqtXgwdVn/JB77Yoff54KJICVJ4tL+XVqiaWe1iDADN45M4y/f7MYQORlz8kIkOfRRGINs4JZ+t7RACwVBEIQzFWZQ+PDu0Vz50gqKq50BM/ydzmJUeOeOEXSL7xhFQgXhVNFWE/kVlW3djIAGpEZz+YAu/OO7vTwztfmKx7pVjVeWZAcMMmp2LaFy/Re4S44gmywYE7sTPXY65tQsVE3neIWDX/YW4lY1Ptt0lNU5JUzqncgD5/dkfM9OGJTQKyrdfl4Gl/TvwvurD/H+mkPounf2TAdkScKtaozv2YlZEzMZHmRCjCev7s/xSgcrD5QEHVB5A6kBjMmMD/kzNKbFgymzUeGpq/vz6OfbQ143aTbKPD2lP2GG5s26cS5Yvr+If3y3FwBP1QDc5TkYYzaEFFCZFTNPjn2SblEiJ7ogCEJ71SkijEUPjOf2t9exv7Aah1sl0CRVeJiC2aDw3oyRZCW3fZF6QWgLse10z9Spfndxby58binXDUtlSNfYZjnn9zsL0ALcHCrXfU7F2k+Iv/g+zBlDkRQD9oMbse9fizk1C/AW+Z71/kaGp8dyzdBU/nXaPqimSoq28PtL+/DQhb1Yn1tKcbU3q2C0xcjgrjEkRvpfWhiIIku8dstwnvpmFx+uO4zEyeQapws3KUiSxEs3DmFy75apN9sqie6vGZrKsQoHL/28P+iAymyUeejCXlw9JKWFW3f2cbhV7pm3qd7f0llwJSBRs+c7Sr4vwHXMhWyWMXc1k3BlAuG9Thmd1CXMhjAeG/0Yl3e/vPU/gCAIghCS2HATX9x3HpvzynljWQ4/7CpAAswm72Cjy6PRp0sk90zK5IK+nTE2YQRZEM4VMZb2u2eqVrTFyKOX9+GxL3bw1f3jmqUe6mvLsn221QBozhrKV3xA/OUPYe09tu5xa49RWHuMqnesLEn8+/rBJEU3//YPk0HmvB7Nk0lPkSWeuCqLeydl1tWbdasn6s3q3lm69E7h3Dspk0v7d2nRiZlWqxp23+QeJEWZeezLHUjg9x8bvBEkwF+n9GfqUJFhzp/FO46h+6SGlCn6WqViXTFJN/UhaoiKZIDqHRVUbaoivFc4uub959ZtvXjq4t9xaY+Rrd94QRAEoUkkSWJo11hevXkY/7NwC9FWIxN6JWI1KSRFm0mN9S2VIQgdUYzV2C5To59uyuAU5q/LY96aQ9w2Nv2Mz5ftJ/MngPPoHnSPC2uvMY2eI8woc6CwukWCqZZwar3Zw6W2k+WFIkyt9hlatQTzNcNSuXxgEt9uO8acpdnkltRgOjF65lI1MjqFM3tiJpcPSMLczAW1ziWvLvEdeTh11MEQOw5HfiGGmLVYehRj7e3EXWVBs6fhLh+BrEeybLuFS3u20QcQBEEQzkhOcQ2/v7QPo7s3//p/QTjbxYab2nyZX22dp425ZZTWuAgzynSJMnP5gCQSo7zL2iRJ4q9T+nP962u4bECXkJe7nc4RYKmbaq9EtkYFVdtV16Ha4TmjdrQFgyLTPSGibd67td/QbFS4dlgq1w5LpazGRbndO3IQYzHWK1Ar+JdXauNwqc3n8dNHHTRXIq7CK/2eQwW+2HKUZ65pvk2PgiAIQuvQNJ29x6vo28JFPwXhbBVjMbbZMr9qp4eP1+fx+vIcKuzuenscwwwyzyzew3mZ8cyamMno7vH07BzJ9OFpPPPtbp7/9ZAzem+TImPXfFd+KZYoNFsluqY2GlBJkjeznhC8Vg+mThUbbhIBVIiKqp0YFdln71koow7grSvgUbUmZWYRBEEQ2k5emY1oi1HUjxKEAKwmBVXTcbjVVl3pdKikhumvrabS7vabmrw2ScKSvUWsySll2rBUnrwqiwcu6MFFzy1jVXYxYzND31OUV2pjVXYxgbZdhaX0QTIYse1bTXifcQ2ey6PqpMaeHUv82os2DaaE0LkCTOGGMuoAtekodUSiREEQhLPL7mNV9EkSs1KCEIgkSURbjVTY3a0WTOWV2rjq5ZVUOdwBs23W0vEOan+y8Qg1Lg//um4Qj1/Zj//3xQ4WPzgBk6Hhge6SaiersktYlV3MygMl2FwexmR24vy+ifywq8BnwF0OCydm3E2U/jAHSVYwZwxBkg04crfgOLyN2Ml31h2bFmehR2JkU/8MHZIIps4yUWYjPrknCG3UAQDdmzFREARBOLvsPlZJ3yTR2RGEhtQW7u0cdWb7kILhUTVueGON30CqodpOdrfK4u3HGZASze1j01mwPo+5K3K4d1KPeueodnpYd7CElQdKWHmgmKNldkZmxDG2RyduG5tO786RSJJElcPN9zsL/LYxauQ1yOGxVKxeQPE3zyKZLIR17kHUmOvrjgk3KdwzKbPZ/z7nOhFMnWW6J4T7yeQX2qgDQN+kyGavAC0IgiC0vN3HKrlqcHJbN0MQ2rUYa+ulR/9pTyFlNpdPIBVMbSe7W+XFn/Zz65h0nrgyi6v/s4JLs7pQWOVk1YFiVmaXsPtYJQNTozkvsxNPTx3AoNRov9s0Is1Gpg5J4fONh3Fqvn28iKzJRGRNDvg5jIrMZf2TzuyP0QGJYOosYzYqTB+Rxrw1h3Cr9a/aYEYdoHbkof6ohyAIgtA+qZrOweJqymxuZEli+9EKHrmkd1s3SxDatRiLkbIaZ6u815wl2dQ4A2dZbqy2k8uj8c7Kg7g1nSiLkYv+vYz+yVEMT49jUGoM3TuFU+30cLjUhnKwhLRYS11GwPpvqvGE9RM2y93IoQtBlnYFwGJUePfOkSKbdhOIYOosdPvYdD5cexjvqtv6Ght1AJBliYuzOrdQ6wRBEITmUFzt5KN1h3lrxUGcnpPFKKudHm57ax0zJ3Tn2mGpRJpFIgpBALC5PHyx+Shzlx8kt6SG73cVYFK2khRj5q5xGUwdmkpEWPN2fQ+V1LD7WKXP46HUdqpxqfzz+71cPzyN31/Sm6e/3U2k2cgHaw8hSxK2U8rhhBlknv9xP+dlxnPv5B6MSI/zPuGqgc9mYraXseDh+7j5w30cKKz22T91OlnyBlJzbxvBoLSY0D68AIhg6qzULT6cS/t34budxxu9SE5nMSr8/pLeGEUWP0EQhHZJ13X+/cN+5izLRuJkBrBT5Vc4+L/v9vL3xXt48uosrh/RtfUbKgjthEfV+PviPXyw9jCSRL3gw6VqHCqx8bfFe3h60W6uG5bG/7uiX6NJHoK1r6Dam2XZc2ZZls1GhSev7s8PuwoornGRX1Hs97jTMwLeMymT3wwPR5r/a0jMgmlvE2Mw8cnsTry54iBvrTiIw6361Ce1GGU0HS7N6sJDF/Uio1N4Ez69ACKYOmv9c9ogjpbZ2ZFfEXRAJUswfXgqt4xJb9nGCYIgCE2i6zq//2Qb32w7FjB7ay37ic7RE1/tpKjKyf3ni0rsQsfjcKvc+uY6th8t95uOvFZtgPXxxjy2Hy3ng7tGE94Ms1TVTjean5VCoWZZdrhVftlTyG8+2tTotQ8nMwK++st+1JXf8/Dkq2Hcb72FovAGZ/dN7sHsiZks21fEgvV5HK904FY1oi1GJvdOZPrwNFFioRmIYOosZTLIfHD3KB6av4Ule4tweTRUf2n+AJMiIUkSSdEtn9FGEARBaLoXftrPN9uOYXf7Ft4MxO7WePmXA6TEWJg6NLUFWycI7Yuq6cx8fyNbj5T7ncH1x+HW2H2sihnvrmfejFFnVG9T03SqHR5U1bf/FWqWZaMsc+8Hm/wOkDeYEdCj87p+IYMTRzLZT2IxRZaY3CeRyX0Sm/YhhUaJYOosFmZQePXmYew4WsHc5Tks3nEckyKj6TpIIOG9qG4a1ZVbxnTzZnn5z0o+WneYG0aKJSGCIAjtSVmNi1eXZPvtFDbUmQJvB/GJr3dxxaBksYxb6DC+2ZbPhtzSkK8Zp0dja14Fn246EvQSWbeqsb+gmp35FezMr2RXfiW7j1ViNsq4Vd/3DzXLskGRcHh8B1GCygioSjz/wz4m9xYBU1sQwdQ5oH9KNM//eghP2txsyiuj0u5GkSXirCaGpccSdkpl3rm3DWf6a6vJ6BTO6O7xbdhqQRAE4VQLNuThr2JFMJ0p8O4b+Wl3AZeK1MZCB/Hqkux6+6NqBZuS/NUl2UwfnuZTKqbG6WHP8Up25ley82glO49VcKCwmpQYC1nJ0WQlR3Fh3870S44i1mpk0rNLOFRi82lHsFmWrSaFGpfqk6U5lIyAe45XcaCwmh6JEaH9EYUzJoKpc0i01djoqET3hAj+ff1g7v9wM5/fO5a0OGsrtU4QBEEIRNN03lx+0GeJTyidqRqXyitLskUwJXQIu/Ir/QYwoVwzhVVOlu0vRpbwBk75lezMryC/3E7PxEiykqPonxLF9BFp9E2KxGry322ePTGTp77Z5TewCybLskfVMcjgOu3xUDICqprO+6tzefLq/o0eKzQvEUx1QON7JnD/5EzuencDn947tn6aUHs5HFnv/W9ZBmsn6DoaDGFt12BBEIRz3JEyO9VOj8/joXSmALYfrUDVdBRZFGUXzm1fb83H6WdZXCjXjM2lctc76xnSLZas5Cgm9krg3kmZ9EiMCGm57NWDk3lm0e6Q2l/LbJTp2yWSLXkVPs+FkhHQo+nsPl7VpDYIZ0YEUx3UbWPT2VtQxUPzN/PaLcNRCrbBqpdh91egmEDXOLHxymvYHTByJsSktWWzBUEQzkkVdjcGPwFQqOmVjbJMlcNNjNXU3E0UhHblWIUdzU/erVCvmYm9E5h724gzaovVZOCdO0Zw89y1DWYUPJ1JkejdOZKBqTFs9hNMhZoRsMrhDqndQvMQu1Q7KEmSePKq/tTY7ex/+Vp46xLY8Sl4HOCsBFc1uKrAeeI/a+fAy8Ng2bMQIGugIAiC0DSy7K8Me/3OVDB0dGQxKyV0AKfvL6oV6jXjCXCeUA3rFscbt47AalII5hI0G2X6p0Qz765RxFpN+HvJqRkBg9HcBYmF4IhgqgMz4eE95WnSS5eD2w56Azce1QUeJyz/F/z3j63XSEEQhA4gPjzMb0awUDtTmgYRAfZ1CMK5pFOE/9nXUK+ZuADnaYpxPTvx9W/G8asBSYQZZMxG3252uEkhMTKM313UmwWzxhBpNpKREI7V5DvzdGpGQNu+1WhuB7rqwZ69gbJf3qp3rCJD7y6RzfZZhOCJO25H9tX9GAu2YvTZ8tgAtw02vQOdesCIu1qsaYIgCB1Jl2gzabEWDhTV1Hs8lPTKEjCpd4KYmRI6hHE9E/hk4xFqTkv6EMo1Ex6mMKmZ04lnJkTw0o1DKbe5WLghj2X7iim3uwhTFJJizFw/Io3zMjvVu04vyerCHz/b7vd8wWYENCoyt45Jb9bPIgRHBFMdVWkO7PrSu6zvFOnPV2Fzw8EHIwg3eS/0uZtczNvmZsnt4d6D3Hb48S8w9DZQROVsQRCE5jB7Ug/+/OUOn85hsJ0pi0lh5oTurdlkuhL6AAAAIABJREFUQWgz5/dJxGSQfa4XCP6akSWJS7O6tEj7YqwmZk7IZOaEzEaPNRsVpg9P44M1h3D72QgWTEbAHgkR9OosZqbaggimOqq1r0GA9cSqDi+sdfHo+AYy+Okq7PkGsqa2UAMFQRA6lisGJvHnL3f4fS6YzlSs1cTIjLiWaJogtDuKLHH72AxeWXLAb9Hexq4ZoyJx86humAztY8fLnedlsGB9Hu4g93qdymJUeOCCni3QKiEY7eMbJLQutwM2vw+a/6wvj4w18ewqJ+WOBjZluqphxfMt1EBBEISOx2xUePa6QX73WTT+WpkXbxjiU3xUEM5lt43tRpTF6LfYdUMkIDLMyJ3jMlqkXU3RNd7Kc9NDv/4tRoVbRnfj4haaYRMaJ4KpjqhkP/jNG+M1PFlhUrqBZ1c5Gz7P8e0is58gCEIzumxAEn+6vG9IHSqLUeE/Nw5lWLfYFmyZILQ/MVYT82eOJjLM0ECvpj5JgvAwAx/NHE1CZPuqoXnZgCT+OW0QFqMcVEZAi1Hh9rHd+OPlfVq+cUJAIpjqiBwVIDX8T/+XyWG8tM5FUU0D9RIkvPunBEEQhGZzy5h0Xr5hKAmRYYSH+a8tI0lgNSl0i7fywd2juKBv51ZupSC0D5kJETx5dRbyiWuiIeEmhS5RZr7+zbh2m/nuykHJfH7feXUZAS2nDazIEoQZZEZlxPH6rcP4w2V9xYx0GxN7pjoipfGRmP6JClf0MvD3FS76JgQIvDTNW+BXEARBaFYX9uvM+X0SWXGgmDlLs1mVXVI3Um2QZc7vk8DMiZkMSYsRHSmhQ6uwufnnf/cy9/YRSMBrS7PZdLi83l4ol0djYGo0sydmMql3Iko7z3jZp0tUXUbATzYeYduRCsrtLjyqzr6CKj6/9zzS4qxt3UzhBBFMdUQRid66UY14cpKZoa9V87sxAYIvkxUU8RUSBEFoCbIsMaFXAoNSYxj795/Y8NiFyLJEmKHh0XdB6Ch0XedPX2zn4qwuTD6R4nxS70SOVzg4VFJDlcNDhNlA1zgryTGWNm5t6GKsJu4afzJDp92lMuSp7+kcZW7DVgmnEz3hjii2G8RmQNHuBg/rESdzfZaRF9e5GJB42uyUbIAB01uwkYIgCAJAdnE1mYkRWEQxXkGo58st+ew5XsU3vxlU7/Eu0Wa6RJ97AYfFpJAaayW7qJq+SVFt3RzhBLFnqqMa9zCYIho97PGJYdS4/CSZkA0w+t4WaJggCIJwquzCajITGr9fC0JHcqTMxlPf7OL56wdjNnac2dqs5Ch25le2dTOEU4hhro6q39Ww6H98Hs59qP6GzLRoGcdjp41+SDJ0zoKEXi3ZQkEQBAHILqohMyG8rZshCO2Gqun8buFWZozPoH9KdFs3p1X1S4piV34lDGvrlgi1xMxUR2U0w3XvgCH0aXC3Eg7Xvtn8bRIEQRB8ZBdV013MTAlCnTeW56DrMGtCZls3pdX1S45i17GKtm6GcAoRTHVkPS6AKXPAEOSmTElGNUVzO4/z5WGRxU8QBKE1ZBeJZX6CUGtnfgWvL8vhX9MHtfusfC2hdmZKF3U+2w2xzK+j6z8VopK8S/6KD3iz/Olq/WNqU6mnn4dyxfM87ozj5jfXAnD14JRWbrAgCELH4VY1jpTZ6RYv0iALgsOt8tD8LTz2q74dNjV4fEQYVpOBI2X2Dvs3aG9EMCVA19EwewUU7IRVL0P2T+CsAlkBcwwMvB5GzICoZAB6A/NmjOIWEVAJgiA0uwOFVSzfX0y5zUW53U1EmEJJjYuUszC1syA0p//77x56dYlk6pCO3e/wLvWrFMFUOyGCKeGkzlkw9dWgDu3dJZJ5d43i5rkioBIEQThTHlXj+10FzFmazb6CKnQdnB4NAFmC859dwvD0WGZNyGR8z06iUK/Q4SzfX8R/dxxn8YPjO/z3v1+SN6PfJVld2ropAiKYEs5Ar84nAypdhymnjhQ5q2DrAlj3GlQdB9UJRit07g/nPQiZF4AstuwJgiBUOtzc9tY69h6vwuZSfZ7XTgRWKw+UsPlwORN6JvDiDUMwGcQ9VOgYympcPPLxNp69bhAxVrFnOys5ik83HW3rZggniGBKOCOnBlQAU7Ji4bs/egMpSQK37eTBHifkLof8zWC0wOQ/wfA72qjlgiAIba/G6WHqf1aSV2rHpWqNHm9zqSzZV8jtb6/jvTtHYlBEQCWc23Rd59HPt/OrgUmM69mprZvTLvRLjuKv3+5u62YIJ4i7sHDGagOq/3y7jsqXxsOWj8Bjrx9IncpVDTVF8N2j8M1vQWSkEQShg7r3g00cKQsukKrlcGtsPlwmOlNCh/DppqPkFNXwyCW927op7UZarJVKu5tym6utmyIgZqaEZtIrzsDX0f+HXJILeIJ7kdsGWz8CcxRc+ETLNU4QBKEdOlBYxdqckrq9Uaeq2bWEyvVf4C45gmyyYEzsTvTY6ZhTswCwuzU+WneYhy/qRbTF2NpNF4RWkVdq45lFu5k3YxRmo9LWzWk3ZFmi74kU6WN7iNm6tiZmpoTmsfxfmCsOYjolkEp/vorEf1ZR4zo58zR3k4tJ79ScfJ3bBmvneJf+CYIgdCBvrcjFo/nOzFeu+5zSn94gevR0Uu+fR8o9bxM59HLs+9fWO06WJD7deKS1misIrUrVdB5esIXZE7vTLzmqrZvT7tRm9BPangimhDOnumHd6+Bx+D6lwwtrG5mG9ji9KdkFQRA6CIdb5fPNR32CKc1ZQ/mKD4i76B6svccim8xIigFrj1HETr6z3rF2t8rry3Nas9mC0GrmLM3GqMjcNa57WzelXaot3iu0PRFMCWdu7yLQfDNQATwy1sSzq5yUOxrYF6VrsOcbsJe1UAMFQRDal6PldmQ/2Z2dR/ege1xYe40J6jyFlQ7cIey3EoSzwfYjFby14iD/mj4I2d+FItAv2ZseXWh7IpgSztzGd71JJfwYnqwwKd3As6ucDZ9DUmDPohZonCAIQvtTaXf77SSq9kpkaxSSHNz+EKMiU+UIcp+qIJwF7C6VBxds5vEr+5EsClUH1LNzBIdKa3C4/Q9mC61HBFPCmavMb/Dpv0wO46V1LopqGhg99TihuqCZGyYIgtA+WU0Gv4lMFUsUmq0SPcBs/+k8qo5FbMwXziF/W7ybASnRXD04pfGDO7Awg0J6fDj7C/wPZgutRwRTwpnT3A0+3T9R4YpeBv6+ooG9U7rq3XslCILQAXSOCsPlJ4tfWEofJIMR277VQZ0nzChjNoqfcuHc8MveQn7aXchfru7f1k05K3iX+lW0dTM6PHEHFs5cWHSjhzw5ycwbm1wcrQywd0oxgbnx8wiCIJwLYqwmRnWP83lcDgsnZtxNlP4wB9u+1WhuB7rqwZ69gbJf3qp3rFGRmD48DUkSe0qE9k/VdHYcrWDpviJ+2VvI5sNlOD0nZ2BLqp3876fb+Od1A0W6/yBlJUeLjH7tgKgzJZy5HudDwQ5QA++L6hEnc32WkRfXuRiQ6CeGl2VIG9mCjRQEQWhfZk3IZNOhMmpc9Zf0RY28Bjk8lorVCyj+5lkkk4Wwzj2IGnN9veNkSeKO89JbscWCELqSaicfrT/MWytycXpU5BPBv66DDtwwMo3bxnTjqW92c/XgFMZmirpJweqXFMXi7cfauhkdngimhDM3fAasfLHRwx6fGMb72wIs5YtOg5ShzdwwQRCE9mtsZjyx4SZsbrvP/qmIrMlEZE0O+FqDLDGkawzd4sNbuJVCR6Wf+FKeyczn2ysP8vfFe5AkcLj975t+d1Uu76zMJcJs4IVfD27ye3VE/ZKi2HO8Ck3TRdbDNiSCKeHMRSVBxng48BPecSav3Ici6x2WFi3jeMxP4T1jOIz7bQs3UhCEjkzVdHKKqim3u5EliLWayOgU3qZL5GRZ4p07RnL1f1ZQ4ww+I5csQWy4iZdvFANQQvM6UFjFmyty+XZbPtVODzpgMSqMzohn5sTujMqIC/qa+cd/9/D2ylycfvYGnsqtevsNdpfKwwu28spNQ0VgEKRoq5EYq5FDpTYyOomBlbYigimheVz4BBxaCW576K8Nj4esqc3dIkEQBIqqnHy49jDvrDqIy6PVddI8qk5suJGZ47tz7bBUIs1ts0ejR2IE8+8ew01vrqHG4UFtoCQfePdJdYoIY+GsMXSKCGudRgrnvAOFVfx24Vb2Hq/Co2mcWrrM5lL5ZW8haw6WEGM18vTUAUzundjg+eavO8zbK3Oxh5C22+nRWLqviGcW7eaxK/o19aN0OLXFe0Uw1XZEMCU0jy4DYNrb8MkdwQdUsgEMZvC4oDQHOoubpyB0VMdrjrOlcAuVrkoMsoF4czyjkkZhNpibdD5d13nuh328viwHwO/ouL1c5R/f7eXv/93Dk1f15/oRaWf0GZpqQGo0ix+cwHPf7+WLLflIgEerH1WFmxR04PrhaTx4YU9irKY2aatw7tmQW8ptb63D5lIJFMvreIMqm0vlnnkbeexX/bh5dDe/x7o8Gk8v2h0wkKrZtYTK9V/gLjmCbLJgTOxO9NjpmFOzsLtV3ltziJkTu5MY2bRrv6PplxzFrmMV/GpgUls3pcMSwZTQfHpfBjcsgPk3gq6B2xb4WFMERCXDbV9D7gp47yqY/h50G9t67RUEoU1pusaaY2t4e8fbbCrYhFE24tE9yJKMLMlousY1Pa/hpj43kRYVfKCj6zr/8/FWFm0/3ugSI9uJ5A9PfLWD4mon903ucUafqalSYiz837UDWbqviOtHpLH+YOmJJYkS8eEmrhmayq8GJmEWNaWEZrS/oIrb3lrnkwSlIQ63xl+/3UVcuJHLByT7PP/9ruNomv+wrHLd51Ss/YT4i+/DnDEUSTFgP7gR+/61mFOzAJCAj9Ye5sELezXpM3U0WcnRfLD2UFs3o0MTwZTQvLpPhN/tgW0LYeXzYCsBSfam7ZFkb8a/1JFw3kOQeb43i9+AaWCNgwW3wFUvQp9ftfWnEAShhVW5qrjnx3vYX7Yfm8c78OLSfGvRLdizgE/2fcLsQbOZ0X9GUPs1/v3DfhZtPx7SEiO7W+Oln/eTEmNhypC2KRa6/EAxqbFWHrmkT5u8v9DxPDh/S92AwukamkFyuDV+t3Abk3onYjXV70q+uiTbb3CmOWsoX/EB8Zc/hLX3yYFTa49RWHuMqvv/To/G26tyuW9yDwyKqODTmH7J3mV+QtsRwZTQ/MIiYcQMGH4nHNsKlUe9S//CoiCxL8T4GWHOPB9u/gQ+/DXUFMGw21u92YIgtI5qVzU3fHsDx6qP+Q2gTuXRPXhUD69te41KZyW/Hd5wsprSGhdzlmX7LYjbUOcQvCPuf/5qJ1cMTGqTTtwnG48wbVhqq7+v0DHtPlbJweIav0v7gppBkuDLLfncMLJr3es0TQ9Y98h5dA+6x4W115hG2+byaBwutdE9IaJJn60jSY424/RoFFU5iQ83ieQdbUAEU0LLkSRIHuz9TzCSh8Adi2DeNVBVABN/7z1HA3RdFwUrBaGN1Tg9fLHlKF9vzae0xhscxYWbuHpwClcPTq43cq3rOvf/fH9QgdSpHB4HH+35iO7R3ZnSc0rA4+avO4x0ep5xguscAng0jR93F3Jp/y5Bt605VNjcLNtXxDNTBrTq+wod15srDuJSfQcdgp1BsrlU5izJ5tcjThaOrnJ4MMhSXYa+U6n2SmRrFJLc+FJVRZaosAcopSLUKatxMX/9YewuldF/+wlV0zEpMt3ircyamMkVYmlwqxDBlNC+xGfCnd/DB9dCdQFc/k845cZbYXOzcEMeb686SFGVE7eqYzLIdIsTNw5BaG2FVQ7+/cN+vth8FEnCZ7nQtiMV/OXrXUwdksLDF/UiITKM7cXb2VWyy28gVba8jOLvinEVulDMClHDoug8rTNKuPeadqgOnt/0PFf1uApZ8p05yi2u4bkf9vkkbwi2cwhQ41SZszS71YOpr7flM6FXAtHWtskqKHQ83247hupnb1MoM0gFFTZ2fraIpMoiPAUFVBYWoxkneJf1n0axRKHZKtE1tdGAStfBKJb4BVTj9PCnz7ezeMdxJIl6QbFL1dhfWM2fv9zB41/u4K5xGTx0YS8xY9WCxDdVaH8iO8Pti6BkP3x8O7gd2F0qv1u4hZHP/Mi/fthLfrmjbuTL5Tl54xj61A/86/u9ATe/CoLQPPYXVHHp88v5eEMedrfqd9+FzaVid6ss3JDHZS8s40BhNe/sfAenx+lzbPHiYo5/fJwu07vQ75V+dP9/3XGVuMh9NhftlCV7do+dVfmrfF6//UgFV7y03CeQgtA6hwBbj5S3+j3kk41HmDZULPETWodb1XB4/O+VCmkGye0ib8U63MfyMSR0ImHSRGTZf9cyLKUPksGIbd/qoNoXHyEyVvpTUu3kypdXsHiHN8FOoGLINSeyL76x/CAz39+A288spNA8RDAltE/mKLjpE5AVyt+9katfWso3244FdeOYu/wgM95dL24cgtBC8kptTJuzmrIal9/g5XQeTaek2sW1r67k55zNaNS/NlW7SuEXhSTfnEzkwEgkg4QpwUTavWm4il1UrKqoO9bmsfHWjrfqvf5QSQ03zl1DdYDCt6F0DgEMskSVwxPUsc3hQGE1R8vtjO/ZqdXeU+jYVE0n0DzFqTNIjZGtVmJ/8wBdHn2U+BkziLnqCi7un4S/SRA5LJyYcTdR+sMcbPtWo7kd6KoHe/YGyn6pf02nx4eTFG1pwic7t9ldKjfOXUteia3RTKV1r3GrrDhQzP98vBXdzxJo4cyJZX5C+2UIw3HVG9z0j085WFOFO8ivq92tsjqnhN8u2MKLNwwJak+VR9X4aU8hc5fnkF1Ug8OtEmZQSIuzMGNcBpf1T8JkEGMPggBw93sbqHZ4fDauN5TgQce7n0LOuwlL+gv1Xmfbb0Nza0QNi6r3uGJWiBwYSfXOamInxNY9vqtkV73jHvhoMzVO/8FPza4lVKz8CK2mnLyXb8F0WtIJf3QdFKX1lsR8uukI1wxJEZnLhFZjNirIkoTmp3N96gxSeJ9xDZ5H03WiLfWXpt49oTs/7yn0m00zauQ1yOGxVKxeQPE3zyKZLIR17kHUmOvrjgk3KdwzKbOJn+zc9vqyHA4V1+D2M4jVWPbFH3YVsOJAMeN7JrRBy89tIpgS2rW3Vx8i2xWNm9Aycznc3uBoyd4iJvcJXKld13XeWJ7DK79k41a1eulcbS6VMpuLRz/fzp8+38Ed56Xz0IW9UMS6Y6ED25pXzqESG+ppnbBgEjxoOmiOeFRHEor5WN1r1WoVQ4QByU8AY4g2YD9UvxC4w+Oo+98HCqvYW1CFvwmy2jbFnj+Dkv++TOz5dyGbzD5JJ06n4+3QtRi33VuLz2hF1eHzTUd5b8bIlns/QfBjcFoMGw6V+Tx+6gySJCuYM4YgyQYcuVtwHN5G7OQ7646VkMg8LePeoNRokqLN5BTX+H3fiKzJRGRNDtguSZK4bEDr7lk8G6iazturDuLwMyMVzP23NmGICKaanwimhHZL03Tmrjjod1lf0DeOpdkBgymPqvGb+ZtZsqeowXo0NSeWDs1dnsOmQ2W8efsIkeRC6LDmLs/Bedpei1ASPKAruErGYUn5uO4hJULBU+1BV3WfgMpT4cEQUf+nSkIhr9RGSoyFt1bk+l3Se2qbwvuMQ6spp+ynN4i/5D6ix92Irnr8dw4lOL9PYvNmCdV1OLwGVr0IB34ATfNWJkWiMmEEl5gvo1dC4M6lILSE2RMzeXD+Zr81oYKZQTIpEjeP7uqzakOSJF6+cSjXvroqpFpvAGajzIs3DCbMIH5jT/fznkLcfgKpUO6/6w+VcbTcTkqMWELZnEQwJbRbS/cX4TiDwn8AW/LKySu1kRZnrfe4ruv84dNt/LKnMOAerNPZ3RobDpUxe95G3rpthMiMI3Q4DrfKdzsLfGaBQkvwoOCpGoiuf4Ykea9vaw8rkkGicmMl0SOj645UHSpV26roPK1zvTNIWjjXzVlNlcON3a36nZU6vU3BdA4BLEaFmRO6B/E5gnR4LXw+E6qLwG2D2sWRJ/4rumA1jylb4bmX4aqXodclzffegtCAyX0SMRlkv8EUND6DhCRxy5h0v0/1S47izduHc9e7GwIWBT6d2Sjz1ykDOL9P58YP7oA+2Zjn998qlPuvBHy34zh3jstogRZ2XCKYEtqtTzceOeMbh6br/HfHce4+rXP0y95CFu847jeQamj5oNOjsTanlE82HWH6cD/FhwXhHFZa4/Iucz3tsgw1wQPo6B4rkrEKAMWqkDglkfx5+chmmYh+EbjL3OS/n48xzkjM2Ji6V4YpYdw14EZm334BpTUuhj71g9938NemRjuHQHy4ieHdYhs8Jmi7v4HP7vIu6wtABmTVBtU2WHgbXPK0t+i5ILQwRZZ4/NKe/O+n23FKoc0EWYwK1w5NaXCGY2xmJz69Zyz/8/FWDuSX40FCPS3thSR59291ijDxt6kDGSeSsARUUOmbBRVCu//WFvcVmpcIpoR263iFw+/jodw43KpOQaXveV5dku13tCyY5YN2t3fdsQimhI7G5lL9ZukKpX4MgCTpoNfftJ5weQJKuMLxBcdxFbqQLTJRQ6NIm5WGbDy5jEjXdab1mgZAeJiCLOF3ZirUNkHtEqPGk9Zomk6Vw0OV043FqBBtMfomjzi0utFAyofHDt/9CcIToN9Vwb9OEJrAU1bGkOce5dau5zHPlIk9yFUaFqPM6O5xPHl1/0aP7ZsUxdezRvLD5dP58dY/8N3BKmqcKjo6FqPCeT07MWtCd4Z2jW3epbXnII/m/98n1Hudv0LNwpkRwZTQbgVKuRzqjaPc7kLX9bob9aGSGrYdqfA5LpTlg8cqHGzNK2dQWszppxGEc1aUxeC3yGco2b8AdF0B2XeQI25iHHET4wK+ziAZGJcyjk4W7+h1mKF5MpIBmAwyr9w0lCFdA89KFVU5+XDtId5elUuN04NBltF0HUmCKYNTmDEug56dI717pD6d4RNIpT9fhc0NBx+MINzkvR/N3eRi3jY3S24P9x7kscMX93qX+xnCGm23IDSFKy+PvLvuJvLii3n04XvJWH+EJ7/eiSQRcOm7UZGQJYmpQ1J5akr/oJMxVX33HX3TE7jktrH8E+rSc4vgKTSxFv91t0K51ykyxIWL+l3NTQRTQrsVazX6fTzUTtK3246zaPtxusZZ6RZvpbTahcfPyEwoywedHm8hUhFMCR1JfHgYRoPsk00qlOxfAOEmAxaTjjOEAVIZmThLHE+MfaLe4/1TotmSV+57fEgZyeD9O0cyqnu83/d2elQe/Ww7X287hgR19V3c6snZ7Y835vHF5qP0S45izgQXiQ7fNgGoOryw1sWj4xsKlHTY+QUMur6BYwShaezbt3Pk3vuIv2c2cTfeCMCNo7pyWf8uLNiQx5vLD2Jzeer2Beu6NwC6fkQat4/NoGu8taHT+yibv4C422+r+/8iiGqaC/omsuFwGfbTVtWEcq8LMyiMDnCfE5pOBFNCu3VB30TWHiz1WY4Xyo3DalJ4+/YR9EuO4lCJjUMlNuYszUb1M+kVyvJBTYej5YGX76iaTqXdjVvViLIYRfY/4ZygyBK3jO7G3OUHfZaKBJvgIcwgM2Ncd/pkPsGfV/0Zh+p/Oe+pDLKBeHM871z6DrHm+jNHsydm8ruPt9Rl3Qy1TYqmMtVzmP5FCegZcT4dPZvLww2vr2FvQRWuBopkqhqomsa2IxVc9pGNz5Rwusm+qaEfGWviHyud3DvCRIw5QKfSVQ0r/y2CKaFBmqazKruE15dls+d4FXaXSphRJiXGwp3jMri0fxefrHhVS5Zw7I+PkvTXp4i84IJ6z8WGm5g9MZOZ47uz61glZTZvUe5oi5F+SVFN+h1z7t+POy+PyMkiW+WZumZYKn9bvMfvc8HefxMjwxjaVQwCNzcRTAnt1pQhqfz1291+nwv2xhFnNTHyRAepf0o0/VOi+WHXcbYf9V3mF+rywdOrj+u6zua8ct5YlsMPuwqQJQlZ8q5P7hpnZfbETK4anIzVJC474ex1y5huzF1x0O9zwSR40IGbR3UjMao3SRFJ/GP9P9hfth9VU/Ho9QvvWgwWNF3jkvRLeGT4I8SYfTsBF/ZNxCjL+GTFCLJNBpOR6/sm8uELH3DMshh1wBA6ZfWiV3IME3slMPO9jew5XuVzvQfi0XTKMDFdfZz/hv0vsVJ1veeHJytMSjfw7Confz3fHPhEJTlQdRwiRb0dwdeC9Yd57od9VDs89RI1VTmhuNrFo5956yPeOqYbD1/UC6MiU7ZgIUUvv0Taq69gGTw44Lll2ft72RzKFiwketq1SEb/K02E4EWZjVwxMInPNx/1u0+0sXudxagwe2KmmBlsAaJXJ7RbEWEGrh6cwicb8/C3XzKYG8esid19bhzxEf6X14S6fPDUZYh7j1cxe95GCiodOOpSNZ+82+WW2PjLN7t48utd3H9+D+6dJG5owtkpKdrCJVmd+WFXQdBlBWqZjTKX9U8iMcobRAxJHMJHv/qInIoc5u2ax+r81VS7qzHIBmLCYpjWaxpXZV5FpCky4DkNiswfL+/LE1/tDLmmjckgkRJj4baDElLPy7A5VfTjIB/bj1kG3WDArWr4i6MayvqpoVBKJC95pvC4cZ7Pa/8yOYzz3qrhwVEN7F1QTGArEcGUUI+u6zz2xQ4+23S04fqIJwKst1YeZN3BUp51bMCz+FvS338fU3p6q7RVs9up/PprMj77tFXeryN46MJefLezgGqnp/GDT6HIEp2jwpgyJKWFWtaxiWBKaNd+c34PFm0/RpUj9BtHp0gT1w5L9XluZEYc89cf9lkWFOrywXE9vJvg1+eWcttb67C7VPynzPCqXa748s8HyC2u4R/TBoqASjgr/XPaIKbNWcWIllMVAAAgAElEQVT+guqgZ2zCDDK9O0fy92sH+DzXPbo7j495vMntuX5EGrnFNbyzKjekgMrl0TlcavNJdqNJMjYdCBAsBpP1042RBepkfm9YgFly13t9/0SFK3oZ+PsKF30TZH9v4d3IpYUWHArnvr8t2tNoIHUqh1tja24xDzrCWfDBB5gSWi/1eOWixZgHD8KYIjrwzSUtzsq7d47kljfXBl2/yyBLxIabWDBrjNhy0EIC3MUFoX1IjbXy3p0jsZqCvwFIQIzFyIKZY/wuqbugT+2yIF9RI68h9vwZVKxewJGXbuLIq7dTtekbLD3rJ6VwqxpXDEziQGE1t7+1DlsjgdSp7G6Vb7Yd49nv9gb9mQShPTEbFRbOGsPwbrFBXZtWk+IdxJg5xmcPR3P5w2V9+O1FvQgzyJgMwf+0BcoaGkht1s+4i+7B2nssssmMpBiw9hjlk2gDYLE20u95npxk5o1NLo5WBnh/1QMWsbdBOGnjoVLeX3PIbyBVs2sJx959iMPPTePIy7dQsPDPOI7sBMAtKeyNSuajvVWt2t6yhQuIvV7s+2tuw7rF8snssXSKMBHewP1Xwnvv7dk5gkUPjKdzVAPLioUzImamhHZvSNdYPr1nLLe+uQ6b2+N3ozl4bxwWk4Ku60wfkUZygGKCBkXmtrHpzFma7XdUvbHlg4rkDdam/GcVSAQcHWpoGZDdrTJ3xUF+PbIraXGhZUYS2taRMhu/7CmktMY72xBjNTKxVwLpncLbuGWty2oy8N6MUfyyp5A5S7Pr9iE6PRoS1AU0g9JimDWhO5N7J9ZlB2spd0/ozlWDk3ngo82sPVjaIu8RStbPGix87hnHVGWlz3M94mSuzzLy4joXAxL9BH+WGIjynVkXOq45S3NweJpWH9Hh1nhtWQ63n5feKisiHHv24CkoJGLChBZ/r46oX3IUq/94AT/uKuDVpdnsPV5Vd8+1uW3ImJjQK5FZEzIZkS5qeLU0EUwJZ4W+SVGs+uP5/LS7gFeXeDMXmQyyd1uS5O3AnZcZz8wJmXSLt3DVy6u4oE8iw9P916y5dUw33ludG/QSpVOZTQqf3TuWDbll/PbjrX5npIL5cdN0nbdXHuTxK7NCboPQujRNZ9n+Il5bmsOmw2VIEjhPLAEzGWSeWbSb/inRzJ6Yyfl9EoOuv3K2U2SJC/t15sJ+ncktrmHpviLKalyANzPY5N6JIadRPlNrc0rYesR/WvLGNDQAUiuUrJ8AJUQFfO7xiWG8v83t+4TRCmPuhwAz6ELHU1ztZNm+Ik4vqRZKfcRKh5vV2SWM7dHyS/3KFiwgZto0JIPoZrYUoyJz2YAkLhuQRF6pjfxyOza3ynu7X+PinkO4IWtEWzexwxDfcuGsYVRkLu2fxKX9vTeOo+V27C6VCLOB9PhwEiJPJpb4+zUDeHD+Fr59YBwxVt9N3vERYXw0czTXvrrKu+k8yDZYjArv3DGStLhw3lqZi0GWcJ+WZz3YHze3qrNgfR6/v7SPWMfcjjncKvd+sIk1OSV+ZyFrA/KNh8p4cP5mBqZG8+ZtIwgP61i31/RO4W0+O6dpOn/5ZlfAxBgNBUvBDIBA6Fk/T5X7UP1EGmnRMo7H/ARbugZDbg7p3MK5bfH2Y/ibXAhlptTmUvlw3eEWD6a0mhoqFy2m+5dftOj7CCelxVnrVrnsssdR7PKfcVVoGWLYSzgrpcVZGd09nsl9EhmRHlcvkAK4sF9nLsnqwu8/2VZXbf10fbpE8eV959EpMqzBdcfgXXccbTGyYNZoRpyY7Vq0/ZhPIAWh/bhJksSmQ2WNHie0DbeqcfPctaw6UBzUZl+bS2Xz4XKmzVmNI8TMcsKZW7a/yKegZa3KdZ9T+tMbRI+eTur980i5520ih16Off/akPZBnZr1MxixhLhPxWiFCY+I/VJCPccqHH4HCUKdKc1voD5ic6n49lusw4Zh7CIyUbaFjKgMDlaIYKo1iWBKOGf94bLeHKtw8O6q3IDH9EiMZMUfJvPMNQPomxSJ2SgTEWbAalLq/jujUzhPXJXFmj9ewMDUkx2cQBkGQ/tx0ymz+VnmI7QaTdM5WFzD5sNlbD5cRm5xTV0A/v++2MGO/AocISwHdXo0coqq+e2CLS3VZCGA15bm1Ku5U6uxYCmUAZBTs37a9q1GczvQVQ/27A2U/fJWvWOtisqVYRsbPacbWG4x80lMLB/0HMNXyT05UHYg6M8tnPsCDc6cOlMajKYsbQ9V+YKFxF4/vcXfR/AvIzqD3Mrctm5Gh9Kx1qEIHUqYQeHlG4cw9ZVVDE+Po39KNLvyK3lzRQ4/7ymk2ulBkiQiwgxc1r8LL984FE3TyS6qpsrhISLMQNd4K1nJ/osXBpjwCnkZkBboREKLKre5WLghjzeWH6Ta4cEgSyCBR9WJsRq5aVRXPtt0BJef2cfG9tY4PRo/7SnkSJmN1FiRYKS1bMnzv1eqsWAp1NH9YIuGa7KRuF5jcO5bisloQHLb6j1foCgsiIpgflQkGqAqRjTnIQxrn0HTNbpFdWPGgBlc2PVCjIooetqRxYWbkCV8irWGWh8x2tKy3yP7jp2oZWWEj2u8LULL6BbVjcOVh1E1FSXEpchC04hgSjindYv3zird9e56oixG8kptuFQdte4XSafU42LB+jw+3XiE3kmR/N+1A+nTJfCm8VoRYQa/KWpD+XGTkIixik5Sa9J1ndeX5fDcD/uQJbD7WTpjr1B54af9fpdxBru3RtN13lt1iEd/1bdFP4/gpet6wNH7xoKlpuyDaizrpyyBqun8Zlc/Pp+xmb6F38KqF8FWCrKBry1Gnoy2okvgqt0Mo6ugq7g0bxKPvWV7eWLVE7yw6QXevuRtkiKSgmqbcO4Z0jUWi1HxmXkNpT6i2SBzXgvvlypfsICY6dchKaIT31asRiux5ljya/JJi0xr6+Z0CGKZn3DOCzPIFFW72FdQjd2tnRJIneTRdBweja15FVzzyipWZRc3et4L+iai+LmCQlkG5NF0hnaNbfJnE0L35Ne7eP7H/Tg9mt9Aqpa/QCqUvTVuVefDdYdxqy2/rEbwkgOk/21sKVSo+6AaIwExVhO9OkfSu0skt364hw+5DPWB7TB7BQsn/Ya/xMfilKWTgVQANo+NYzXHmP7NdI5VH2uW9glnn7GZ8USY/Y9/B1sfUQduGNm1xdqoVldT+d13RF9zTYu9hxCcjKgMcity27oZHYaYmRLOaWtySnhw/ma/AVQgNpfKXe9u4OPZYwIu8QOYMS6DLzYfRdV8O8vBLANSZJgyJKXDZX1rS3OX57BgfZ7fGcVghLK3BrwzE0VVTp+aZ6qm4/JomI2yqP/RTCRJIspi8LsHsbHZ4lBG9xsjSxBpNnDDyDTW5JSycNYYdh+r5KlvdvHe6lymjXPwWvbHOE7MPgVD0zUqXZXc8d0dfDXlK0yKb4ZS4dwmSRJ3j+/Ov77f5/f+FcxM6fl9EokLb7nvTuXXXxM+ahTGxMQWew8hOBnR3iQU41PHt3VTOgTRixPOWR5V4555G/1mQGpsz4vNpXLPvE0sfWRSwM5uz86R9OwcwfajlX6fb+zHzajIzBiXHvoHE5qk2unh2e/3Nun7UCvUvTWKLNUlKskrtfHuqlwWbsijyuFBliV0Xad7QgSzJ2ZyxcAkkSL/DE0dksL7aw75zCoGEywFuw9KAr+lFGTJu08zLc7C01MGMGveRj6ePQZFluifEs38maP5bmcB/7vmTlSjw+f1ZcvLKP6uGFehC8WsEDUsis7TOqOEe78Tmq5R5ijj+0Pfc0X3K5rrTyacRaaPSOO1Zd7CvaFutQ0zKDx8Ua+WaRjeZbZlCxaS+Mj/tNh7CMFLj05nX9m+tm5GhyGCKeGc9dOeQlx+llgFu+eluNrJxkNlAQv/Ajw9dQDTX1sTchpsi1Hm6sEp9EiMbPxgoVl8vumI32VgwX4foGnJRWwuDze+sYaNh8rQdL2uo187W3qgsJo/f7mDx7/cwb2TMrlvcg8xW9VEt4/NYN6aQ36fCyZYamwAJMwgc3FWZ37eXYime4NlTdfxaDoX9e3M3RO6MzAlipvfXMesCd3JTIioe60kSfRKtWEwF6KedrsoXlxM0eIiUu9KJaJfBO4yN/nv55P7bC4Zf8pANnjXE9s8Nt7a/pYIpjqoKLOR+TNHM+U/K6lxenySUQRiNsq8ctNQenVuud8bx7ZtaDU1hI8JbtZeaFkZ0Rl8n/t9WzejwxDBlHDOmrMkmxpn/V5LKNXi7W6V15flNBhMDUyN4ZWbhnDvB5sCFgr1p2uclaenDgj6eOHM6LrOa0tzfGpFhfJ9gNAzZ7lVjRnvrKfS4cHTQM+ndlP5f37JZu/xKl749RBkWQRUwfKoGj/uLuDVJdm4G/g7NxYsNcSiqDwwIY17Lh6Iy6NRVO2kyuHGajTQKdKE1eT9Of1g7SFqXCp3je/uc473d7+PR6tfUkG1qxR+UUjKjBQiB3o7u6YEE2n3prHvkX1UrKogdsLJfZV5VXnsLd1L77jeTfocwtktMyGCr+8fx6/fWEOV3e23FEAti1FBluGNW4a3eKHesvknEk/IYit+eyBqTbUu8a0XzkllNS525Ff4PB7Knhddh5/3FOJpJIHA+X0686sBoWXZOlxq46Wf9of0GqHpiqqdFFY7fR4PdQ9UKMlFJLzZGsvs7gYDqVPZ3So/7i7kia93BnW8ANuOlDPymZ/43cdb2XqkIuTlT8Gw4GS2/CWzN14JRzdhMsikxFjo0yWKrvHWukDqaLmdf32/j2enDUTxEwxvKtiEqtfv/Nr229DcGlHD6mcQVcwKkQMjqd5ZXe9xSZLYUbyjmT+hcDZJ7xTOskcm8/drB9Lb6CIMjYgwBYtRITxMIdykkBRt5g+X9mb1Hy9o8UBKrayk6scfiRGJJ9qNRGsido+dSpf/bQhC8xIzU8I5qaTGickg4z5tPU2oe15kWaLS4Wlw0+4Haw6xaPvxkNpnd2u8tiyb5FgL04eL1KUtrdLuxqhIuE6rsxzq9wGCrzEk6xqq6n9/TUN7tOxulYUb8rhuWBoDUgMnQBFgdXYJd76zvskJRU7SkdDRTxlfNODBgEovKY/fGj5hkrINnMA7V8AdiyB5cP0z6Dr/++k2ZozLoGeA5VTV7mqfx9RqFUOEAUnxDb4M0Qbsh+z1HnNrbqpcVfUeK7e5yC93YHN5CA8zkBprIdIsSi6cy0wGmSsHJdP/Lx9QOethjnTpTrXTjcVkICXGwtCuMS2yXLjc5qKwyonTrRFpNpAcY6H6y6+IGD8OQ3x8s7+f0DSSJJEenU5uRS4DEwa2dXPOeSKYEs5JblVHwveHJNQ9L7JEg6mt7S6Vpxft9tuZayypgd2t8eRXO7lqULJIPNDCpP/f3n3HR1Vmjx//3DI1PZCEkkAIoTfpRaooKhYUsWBXFLC767q74q7fxZ/rqrvuomvDgg2xi110QUWkIwgoIBBqqIGE9On398dQEubOZCYktJz368XrtZuZuXPj3Mx9zvM85xxFMY1qatNfCKKrnBVAMQ2kosnR8voMXp63iafGdo/6nBqajXvLuOX12gdSCuDARTpF3KTP4pdAK5599hV8Pi/nT5xEV/tubtC+4fuft/C3VV6+vzEOAMNbju/Vi/hlzDw65GQd/tt9b9l2iio8jB8cur3vEKsaOimjxWv4ynwYfiMkoPIV+9Djq9+mNUXDptswDIPFmwt5ce4mfszbh1VTUZTgirrXH2BExwzGD24tAflpzLt7N978fLqe1Zduev0N5wIBg3kb9zF1bh7LthRh0RVUFPyGgQKcv3Mbt44dU2/vL2rnUEU/CabqnwRT4rSU6LDgMylZHmvOi8fjw/3KVEo6dsDesQOWzMxqs32frdpp+rpYihp89csuLu2eGeNvKGKR6rSaFiOJ9XqIluVgAzK3r/p7Rpuj5TcMZv26m+IKL0nS1NnU47PWUREmkIqmOmN363b+youcoeRx6E/6faWUUgP6/jyZSYNsAHx/1LEVwOt18/WMJ7nWuIAre2cxsktTHp/1G2/d0vfwZ28mIy6D/LL8aj9z5jpRdIWSn0pI6nMk8PG7/JSuKiVjTEa15+uqjuptzPAn57K7xEWlx48BeI661r5YvYvZa/fSrkkCr9zQi0bxtrDnJU5Npd9+S/zQISj1GEit3VXCza8tpaRKftbRaVofp3bis++KOXvPcp68vJtMDp4kshOzJW/qOJGcKXFaappoJ84aeoOJJecFIMup4tAUij/+mK3XXc/6Pn3Zeu117H70UQ58NJPnv1kbtqhBNI1dyz1+nv8+r25/eREiJc5qWskq1uuhJpoKKU4LLVKdIYEUxJajpasK36/fG/M5NAQFpW5+WF9gmh9VsmQmhXNeIqnfFWTeOZ3mt71KQo+RVG5YXO15is9Fd/VIIHXI/QOs/GuBmwOu8MlXTtzcyGeUe3y8uWgrV0xdSFaqg3Y1VEu7qv1VOHVntZ9pTo30S9LZOX0npatKMXwGngIP25/bjiXVQvKA5GrP91am8vcPA2zZX07FwUDKTMAI5uD9uqOYkU/NY09JaDl2cWormz2H+OHD6+34y7YUctnzC9hV7IpY6MKn6rh9Aeas3cMVLyykMsJzxfFzaGVK1D9ZmRKnJVVVuGlgNs/M2YjrqEFttDkvTqvGHSM7kdb7gsM/8xUV4V67Ftfateyav5jt2gA4antYrEUNNhWUU+Lykig5DvVq4pDW/PmjVSEVHiNdD3FUcqb6C6lKKbpiUEo8K9SObPNUHzQ7LBoGBhd3a8bvz2nHmBcWmJ5DLDla3oBBYXn0jV0bkhmLzcufx1Kd8ZdASzYHmtBKrZ7v2KuZxtBsnX8tcPPIWfaw5xBPJf3VNSz0B1e71u8u5fYZy3nu6h5hKzEOzxrOZGVyyM/TRqahxWnsfnc3nr0eVIdKYo9EsiZkoVqqzHn6EqnYNh6X2xc2iDqaN2Cwr9zDVS8u4ou7Bx4uliFObf6SEipXriTzv0/Xy/E37yvnxleXhkwWRuLyBvhtTykTpi/j9Zv6SIuHE6xVUiu2lGw50afRIMi3qjhtje3dgv/O2Wj6WDQlkg0DLurWrNrP9JQU9AEDiBswgJJ95dienofPc2xFLqy6SnGFBFP17dxOTZg0c7XpY0dfD22UfMZpXzJKexwfGhr+YFNWq5WAz8NyRw9maKP4RetIcpyVi7s149IemcTbgl+p/jDV+2LJ0TIMI+xxGrr/rd1zzCt/KgEWBjqGBFMADw+zcea0cu7pG77wjI6fNko+CzmSAzn3twIem7WOSSM7mL7GolkY234sb6x5A7e/enXJ1CGppA4J34YBwL1/EB63+W070tZGf8BgV3El7y/L54YB2RHfQ5wayn6Yh7NXL9S4uHo5/qNfrqX86Io9B0W61ty+AMu2FLEgbz9n1nMVQRFZy8SW5Jfm4w14sagyvqhPEkyJ01ajeBs3DMjmzYVbY05Sd1g07j2nDQ5r+AGval7TIOaiBoaB9BQ6Dqy6yn+uOIM7ZywPWa08RCHAA/rbXKf9Dx0fFuWo53k8aEBv92J6W1ZDswFw5ZtgcVR7WqLdwq7i0G1VseRoWTSVJIfcAM2UVJoP8mKZyPCgU4z5QLRzusaFbXUe+9FDh7QjK0MzVnv590I36/YFSLApJKbPxNU/s0pRGT+vL9jC+ME5NA6TozSx20Tm75zPhqINeAPeGs/zEJsSh6d0IJgU1okmR9PlDfDiD5u4vn9LWTE4DZR9O4f4s+tni19N22hrutYqPH6mzs2TYOoEsxVtY2K5l4rPf0+SNR7i0iB3ODTtdqJP7bQjOVPitPbn89ozpG0ajhgSYh0WjUu6N2O8SdPNqlLirCFJ31B9wBwNrz9Asgyaj4uzO2bwt4s7YbeYffUZ/FOfyrXabByKJzSQOuq5eMthyzx4dST4qq8ynNMxA5se+h6x5Gj5Awb9cqTUsJlwcw9VJzJqohBcXQpn8lA7Ly33sKMkOKL890I3985yMWmQjT1/SGDDPcn07Nk1NBdLgXeWbAt7XKtm5cVzXqRNchtsWnRFIeyancEpd6MrofOfseRoFlV4WLK5MKr3FCevgMdD2Y/zSRhWuwbUNZmx2Pz6jeVaW7y5kF3FlabHEfUo4Ic1n8JLZ8ELA7lxz3aSlr8Bi56D7/4O086FZ/vBynfAJ9vI64oEU+K0pqoKz13Tg8t7ZWLTVSwRVoAsmoJNV5kwOIdHL+1S4+xtot1C+6bHXtSgS/Mk4myySHy8XNWnBc9f25P0BBtxVVYeJ2qfMlJbglMJbe4bls8Fe9fAR+Or/fi6/i3DviSxz2hSzhpH8cJ3yf/vNeQ/fyOlyz/H0ab61rQeLVLISnWGOUrDFq4yXSwTGVa8pCilYR/PTVW5spOFp5d48AXgoe/cPDvSzugOFuKsCmgW4nN7hwwiXd4A0+ZvibhFM8mWxBsj3+CaDtcQZ4kLKUoBoCoqds1O6+TWPDn0Sdwl7UyLAMSytbHC4+eH9QU1Pk+c3CoWL8HWujV64/pZ+fnql13HvI1WUxXmb9xfH6cnwvFUwPTLYOZE2PET+FxYjCqfY8AH3kooWAuf/x5ePhsqZHKlLsgITpz2VFXh4VGdufnMVry6YAvvL9uOejBQUgB/ZSWq1cJ1Z+ZwXf+WNEt2RD5gFbcNyeWPH6wMGeREW+QizqoxcWjrY/4dRWyGtUtn0QPDmZ+3jxfm5rEibzd36Z+EBFLZU0qp8MLme+KDA2jg5eUeplfpO4TPBetnwf48aBT8LDMS7fTPacTc9QWmW0FrytmLs2pMGBJ5ZbQhu7xnJmt3lYQkx1edyFBUDXur7iiqjmvLz7i2raoW+PjQOUv9OeL7PDTExpurvJS4DVw+uLTDkVtmAIXvAuZ9wCo9fnYeqIwYDNs0G7/r+TvuOOMOZm+dzYx1M9hbsRe33028JZ4ujbtwfafr6dioIwDPlS8yPU6sOZp7y2KYLBAnpdJv55Aw/Kx6O35Jpfn201iuNZ8/wIEKWfk4bnweeP0i2PNL8J5UE295MKh6aThMmAv2xPo/x9OYBFOiwchuHMfkizvxwPnt+Xn7AQ5UeIPbhb78hE4WF83OGxnzMUd0yuCBj8xXsKIpcmHRVIa3T4/5fcWxU1WFQW3SGNQmDePnGfCFBiZjCL8BTy32HO47ZCrgh8VTYeQTh3806YIOLNlSGFM1LACr4adTkyQGt0k7/DOX189nK3fyxsKt7Clx4fUHiLPpdM9K5tbBOXTNTI5wxNPPqDOaM/mzNaaPRTORoSpwtmUNqUetTG25t/pKc1aSiusviby1yst937jQD65suw2d6f6z8Ya5hWqqQnGll6wofherZmVkzkhG5kT+/tGi2NoYzSBXV2VDyqnMCAQom/MtLV5/rd7eI9yujNiuNeXwpKU4Dr74Pez5tVogVeNkoN8DJTvgvevh+o9P1JmfFiSYEg2O3aIdzkXJKyjjBWc2922toGLSVwQwcBx8fMLgHPq0So243c+iqTx5xRnc9fZyXN5IOTahbAT49xVd0as0+dxUUMa0+ZuZ9ctuytzBJPt4m845HZswbmA2uemR+9iI2lHmTwnO1Jm4f4CVJ+a7ub23lWR7mGsh4IUVb8I5k8HiwO3zs3ZXCY3irVQURp83YNNVmnnLmbzwFYzrz8Blc/DErHW8s3Q7QLXArKjCy84Dlcxeu5dmyXYeOL8DZ3fMCHfo04rDqnFZz+a8u3Q7Xn/o2l9NExk2XePWAc3gJyd4K2p8v0ZOhX0VBr6Aga4q+FF50zci4msiNe+tjfRE8zLtsRQ1URXISJDmvaeSEpeXvSUuKj0B4u06Kfl5qPHx2Fq1qrf3THFa2HEg9HsrlmtN1xRS4iQX+Lgo3w+r3wvJ3YUoJgP9bti2EAp+g7R29Xyipy8JpkSDtHFvGfe9/zO/7SrFFwjgU60cKl1U4fHz3bq9LNq0nxSnlUcu7cywduFXj87pmMFDF3bk4c/XRB1Q2XWVu3cvpP0bCzEe+X+s21vOAx+tZu2uEvyB4KDtEJfXw/vLtjNzRT5tMxJ49NIudG6edGz/AcQRgQDsWx/24Wj7DqFoULiZL/cm88cPVmNghPS0CkdTg4PvfjmNeHbscEoeW8Pqmyfyp77j2FLkMs1fgCONWfMKyrnz7eXcM7wNtw3Njeo9T3V/GNGOOWv3sqfERSwV5IMFZprTfcQIOPA/2Di7xm0x/TM1bDp8vM7HyA5x/N57G7sIXxzE4wuQGhe+rHptXNK9Od/8ujtkS3EsWxutusp5XZrU6XmJumcYBku3FDH1hzzmrd+HRVdQUILtErxeLugzljsKyshJi6+X9x/dI5O8gt9CquDGcq35AwZD2squi+Ni+WuEK4EQ3WSgDxY9DxdNqbdTPN1JMCUanJ+2FnL9K0uo8PjDNr40CAZVFZ5Kbpv+E3+5oCPX9gtfVODqvi1pkmTn/vdX4fL6w3aLj7NqOKwa/7y8G0OyhrD9jjv49P6/80BcHyoilG/3HQywVuUXc/kLC3n+2h4MjRDgiRh4y0HRwQhfpjqavkMoCq8t3cNji7fGtEqpqTCmZxbjBraibUZw5VGd9Beun/wJm/eU4osyF8blDfD0nI0k2HSu7Z8d9fufrLbuL+frX3ezt8SNL2CQlmCjf+tGdM9KRlEUkp1W3rm2K6OnfEexxYnPpGT40RwWjWHt03jkks7BZZoxr8KH44IBVYQVqiS7wuShNm7/0kXfwHmsbNkNRfWZDiIB2jVJIK2OV4AG5TbGadVNv1uizdHMaRxP+yaSG3Ey27o/2Cx3T4mLyoP3qOofucon7hQ+f2oe/Vs34tmre9R5AaPLemby+Kx1po9Fu412ePv0Op9QECYMIxgI+cx3QEQ1GRjwwap34Lx/hLT5ENGRYEo0KBv3lnL9K0vCBjtmXN4Aj3yxhi8skGsAABxBSURBVNQ4CyO7NAv7vLPaZ7DkwbP5YX0BL8zNY+mWwoP5CcFAqG+rVCYMac3gNmmH+0odeOgJ/vTCAlwx9MGq9PqZOP0n3r61H91bpET9OhGGZgMj8n//cH2HqprlPYPHFlXgMm+BFJZFU+mWmXw4kAL4z+wNbFOc+NTQoCxSw8xKr59HvljLkHbpp2QlQMMwmLN2Ly/MzWP1jmIChnF4G5+qBLfnpSfauG1Ia0Z1yUCZPInpLVrzWOYwFm0qxDAMPCbb/pwHqzZOHNyau4bnHtm6q1vhijdgxXSY9ySU7cHwVqJUmWZxGTqg0LNvNzLt2XwzfyHeT6+JXFRmSN0XlVFVhXGDWjHlf+tN+6TVtLXRKUVNTnrrdpdw+QsLKXf7Iq62+gzw+QIszNvPRc/8yMzbz6zTnnRJDgsjuzTl05U78JvMC0WzjfaWwXKtHRfeyhor8kU3GahCyc7DRZREbCSYEg3K3e/8HLYgQKRBqssb4L73VjG0XTpOa/g/G01VGNY+nWHt0/EHDEpdwdWOBLsF7aiy7IZhcPu7q3GZ9I6JdC4QDPAmvPkTix4YLg1/j5VuBd0RNmfqkMlD7fSYWsZ9/UNXHHyGyp8rrsVss1g0n+X/+3wNl3RvhtOq4/b5eWvxVtOtfdE0zPQbBq8t2MJfL+wY+3+LE8jt83PXjBX8uHGf6d/ooS2NW/dXMPmzNTw/cylTFBtnTPoDr+s6Ow9U8sbCrbyzZBvFlV4OxUs5jeOYODSXC7s2xW7Wb05RoMd10P1ayF+Gf86/+XXjb2iqn2LDyU+Bdkz3n81eUqAjNO14ZegxqtBUhRGd6id37cYB2cxcsYNNe8vwxrC30aardM9K5sKu4SeDxIm1p8TF2BcXURrDbIzbFyC/sIIbpi3m/YkD6jRP74Hz2zN3fQFF5Z6wOzjMOCwaI7s0oYdM9B0f7lLQLOALPyEYzWQgigbukno6ydOfBFOiwVi3u4RNBWWmN4ZoBqmKAh+v2MHVfcNv96tKU4NbkcJZtKmQIpPSsdGcC0C528e8jfsY0jYt5BgiRt2uhOVvBgtJhFG171CX9Oo3pDmBHqaV3aL9LKteW1+t3o3ZRXqoYWajkffibDfg8M+duX1x5vY9/P+9foN3lmzj/nPbmQcPJyGfP8BNry5l+dYi01WXo1V6/eQHLEzMHsVX7gCpOjRLdvDn89vz5/PbEwgYePwBbLpaY7+4Qwyfj32fLaXoza1UTJjEzXnxITkjNbFbVJ67pmedF584cnyNGbf0ZcwLC9l5oDJsLt3R59QuI4GXbugVMqEjTh5TZm8IG0hFmpDx+A3W7ylj1i+7uahb3QXL6Yl23rmlD5dN+ZZyRScQRVtSh0VjQG4jHr+sa52dh6iB1RmsJluDSJOBABgBsNZPDl5DIDVSRYPxyrzNppW/ou3qXuHxM3XuJgwjlnm68F78IY/Ko2bgY+kwX+7xM3VuXp2cS4PX73aIIjfpoSE2yj2hn/8L/lGUU30/eiyfZdVr682FW4+5OauiKKdUc9bHvlrHim3mgVT5mu/Z9fq9bPv3GPKfuY497/0frvxf8asahRVebnp1SfCJe9fCt4/AzNtQP7oF++xJKOu/jmqg4Vqzhs2XX0Hlzz/T6qMP6X/jGF65odfh7YHRsFtU/jmmGwPb1E8j1UMaxdv47K6BDGuXhlVXsenmt3G7JfjYqDOa897E/hFX1MWJVe72MXNFfrXCQ4eULJlJ4ZyXSOp3BZl3Tqf5ba+S0GMklRsWH35OhcfPC/VwL0h5/3Ve3vUVXZonY7eoYcvzOw5ea9f3b8lL1/WqVqFW1DNrPGg156ZVnQw05fdAfMOoBlsf5NtVNBhfrN6F3+RmFcsgdU+piy37K2jVOO6YzsXl9TNvw76QBYhYzgVg6ZZCytw+4us4AbnBadwGmnQNdo2vkj8Vru9QVR5DY2UgtExxrJ/l7hIXe0vd7Ck1rywXU8PMQIA9padGc9Yyt4/pYYp21LSy5/P7ydn7Pyr++yecxXnBAUHV/LcV00G3Qd/boPc4cKZWO37A42Hfc89x4L33Sf/j/SSNGnV4JWtAbmM+vG0Af/xgFRv2luL1B0LyRw7lcTVLtvPopV3omxO+wl9dirfpvHBdL/aUuJi+aCtvLtpKcaUXXVXw+YPFOm4e2Iore2WRIkUATnof/7zDtCdTtKvREGzz8dvuUto1qZv2GSWzZnHg44/p+f77fNKoEev3lDLtx83MXLEDf8BAUxW8/gBNkuyMH5TDZT0zSbBLKfTj7tA25aWvRNxZAUeakJscBNqcI417j4GMwESD4PMHwm7ZiWWQatFU9pe5jzmYOlDhRdeUkJnIWM7l0PkUlXskmKoLY6bB1EFQWRTTy4r1xlh8KkdXQY/1s9TcLlZdfQMVbUaDLXS7RSwNMwMBA3eMW9QA9pe5+Wj5DjbuLaPE5SXZaaVj0wRGdW9OYjQDpZ0rYP3XULo7eJOPbwJtz4VmZ4R9yczl+bUaSNrw8IzlaQYov+LcHyZw9JQF/837Fyx9CW78EhoHS8dXrlzJzkkPYm2VTauPZ2JJD62O2aFpIp/dNZANe0qZNn8LX63eRbnHh2EEe10NbZvGLYNy6JZ1YpomZyTauW9EO+4b0Q6fP0C5x0+8TZftfKeYL1fvMs0TjGVCxh8wmLehoE6CKdfateye/DAtXnkZvVFwgqBtRgKPXdaVf4zuQrnHT6XHT6JDx6afGluJT2t9J8BPr4UEU9FMBgJgccKAu+vxBE9/MgITDYLfMFAwTUWJsas7plsFY+X1B0wHkLGeiwJ4zMotidglZ8FNX8GrI8FVXGOFPyB4E7roWXjPB1T/HGL+LO02mv79EZJmbaOwKHR1KraGmSqJMVT3Wrn9AM99v5HvfitAgWq5OA6LxiNfrOWCLk2ZOLR1taqDAPg88MsH8ON/oHh7sHGkcfD1igrz/wPJLeHMe6HzZcGCH1W8NG9zzANJDT+vWP5JT3U9DiXybGzwHF1Q5oaXzyJwwzcUvPEpxZ99RpMHJ5Fw3nk15lW1yUjgH6O78I/RXWp+rxNE11SSHLK96lRUVG6+9SqWCRmv32B/WZgtXDHwFRaSf+ddNPnrX7B3DC1ioygK8TZdJvBOJqk5kNUXti6EQIzXgKJCUvPg60WtyTevaBBsumYavED1QWpNAoZRJyVoEx0WvCZBUCznAuAN1M35iIPSO8Bt86HDRcGS6bpJzw3NCrodsgfCzV+T1HE4PpMAO+bP0lBIb5VF39w0zFIOqjbMrFi/kIDXheH3UZm3jKLvplV7rs9v0DXKxs7TftzMlS8u5Js1e/D4AiFFDSq9fty+AJ/8vJNRz/zIZyt3VnmwCKaNgC/uCzY+9lYeCaQg+L+9lVCwDr74Pbx6HlQeOPKwYZBfZN7fKdJA8g/6e/RQN1QLpLKnlJL+z9JqOW0vL/cw9LVDVRoNDFcp/qcG49u7m5xPPyHx/POjLlAhRH0Jdw1WnZCJxrFWdjU8HnbcfQ+JF15I4siRx3QscZyNeRXiGgWr8kVNAVsiXPM+yPfgMZFgSjQY3VuYb8WJZZCqoNA6/di2+AEk2nXSE0Ib6MVyLgDJDgupESoGilpIbAZXvA73rYMhf4LG7SAuHZyNIaUV9LkV7lgCN34BTbti1VV6tAy9tmL9LDOTHaQn2hk3sBUW1fyrObHPaFLOGkfxwnfJ/+815D9/I6XLP8fRpvrqjaLAuNeX8eQ3v7GpoCzsr/ryvE388+vfcHkD1FRXxW8YVHoD3P/BymBA5S6Dl8+BPb9GbHh7mLcCdq+GaeeCJxjguH0BlDDNdsMNJO24uV77BqcSOgPrN+CpxeFnZhUC6HEKzSech56aGvZ5QhxPjePNK6zFMiFj1VUaHWN+3O5HH0VNSCDtHtnydcqJawS3zIakzOBEYE1UCzgbwc1fQ0p2vZ/e6U7WaUWDMXFIa9bsXGFaKS2aru4WTeGavi3qZI+4oiiMH5zD47PWhWxxiuZcIFhB6dZBOdJnqr44U2HQ74L/ahDu2or2s3RaNSYODTZLbJORQG56PL/sNO/5UVPDzDirxmOXdSEnLZ6Plu/giqmLyEp1MLpHJhd1bXq4XP+SzYU8+c1601zCmnqu/fGDVXRo9SG5B7YFiz4clD2llAovbL4nnjhr8Lp8ebmH6au8fH9jXPC5RVvgo/Fw1VtYNZVAmC424bY1XqSFH1jeP8DKE/Pd3N7bSrLd/O9C8btg/hRoOyLscYQ4ni7t3pxlWwpDvj+qTsgoqoa9VXcUVce15Wdc21ZVqwqqAMM7hOb9Ravo7bepWLaM7HfeQQkzmSNOckmZMHEe/DglWJDC8AdzRquyHpwMPuNaGHQfJEgFv7ogwZRoMIa2S8dm0UyDKah5kKoqCtf1D/aYKih1M2PxNhZt3k9JpRebrpGZ4mBsnxb0y0mNauvQ6B7N+cdXa2t1LhBsYnpFr6wa30fUv0jXVjSfJcBFVRqqThrZgZtfWxpVz6WqNFUhLcHGuZ2aYtVVOjVL4oHz2zNvwz4+XJ7PE7PWMTC3MaN7ZPLq/M2mgVQ0vbE8Pj8vbUrhcT208MOh1aFJg8LMjvpcsH4WrHgLpXkvUqwKhe7QgCrcQLLttneZvLWYJ84JXdnt1UxjaLbOvxa4eeSs0McP2/ETFOcHBx9CnGDnd2nCgx//YvpYtBMyXTOTaNkodNdEcYWX1TuKKXEFqz02TrDRLTO5WpGS8iVLKHjmWbJnvIUWL72GTmn2JDj7/2DYJFj3Oax6D8oLIBAIrl51vAQ6jwaLyRZ2UWsSTIkGQ1MV/u/Cjvzpo1WmZZgjsaswukcmJZU+bv1s2eEePlXzS1ZsK2L22j0kOSxMHNKaa/u1jFhVK8FuYeLgHJ77flPMRSQcFo2bzswmySn5UicDTVV48vJu3PbWT7FfWxaV/zeqM44qPY0G5DbmLxd25JEv1kR9PE1VSHZYeGd8f6xVeg/pmsqw9ukMa59OcaWXL1fv4r/fbmBVfnHIMaItxew34BNfP/6qvU68Ur1YRjSrQ0bABx/fiREw+NyfxhPqFXwR6BPS+NhsIPlh03IeGhz+un94mI0zp5VzT98IW540K+zPk2BKnBRsusbYPlm8sWALHpP8y5omZJy6woTBrav97Jcdxbz0wyZm/bobq6ZiEFy9CmBg1zXGDWzFVX1aEF9UwI7f30ezJx7H2jK6hvTiFKBZoNOlwX+i3kkwJRqUUd2bs62wgue+30hllINUu67QpSCPHoU+Rj+/A7fPb5pfYhBsnljh8fPYV+uYvXYPL17Xq9oguapSl5f5efvxx9gE2GHRGNY+jfvPbRfT60T9GtY+nb9d1Im/ffZr1AGQ3aJy11ltuKxn6KD+2n4tSbDp/OmjVRgGIYUhqoqzaaQn2HlnfD8yEsOvyCQ5LIzt04KdBypZu6skpDJlLKWYVQw+9/fjKv37aj+PZnVIAVACKBo00/byiPEyk3iTazwPssGo/t/i6IHkl7ZrUJXwfzOd0zUubKvz2I8eOqSF265kgLu0xt9RiOPlrmFt+Gr1LnYVuzBphxiWTTHovGcDvTYb0PECXF4/t7+1nIV5+3H7/ARMvjvK3X6enrOBp+Zs4J5t33LVrbcQf+aZdfwbCdFwyMZY0eDcNbwNf72wEzZdxW4J/ydgURVsusol3TMZd9VgHvzNwOU1D6SOVun1s2RzIeNeX4rPZNWpxOXl4mfmszK/2LSRcDiaAlf2zuSZsT2kCtlJ6Ko+LXjumh6kOC3E2cLn1sXZNOJtOo9e2oU7huWGfd6o7s354f5hTBiSQ5LDQrxNx2nVsFtU4mwadl2lW2YST17ejf/9bnDEQKqqNSaBFMRWirkCO+sN822mDw+z8d8lHgrKowsq4xUXjSlmpvUhOimbIz7XTc2rsZOH2nlpuYcdJeH+thTTXl5CnChJzuCqcuN4G3qUebB2i0rnFqm8eNfZ7PvPFLY++jhjnpvP/I37qPT6IwZlroOVO59qNogPWkVutSCEiExWpkSDdHXfFpzfuQnvLtvOK/M2U+7xHd6SZxjBEuhX9MripjOzSbBbGPj4t7i10EFcpER9ty/Aim1FPPf9Ru4e3vbwawIBgxunLWFHUWXM2/t0TeW8zk2l6MRJ7Kz2GSx98Gy+XbeX5+fmsSq/GKumAAoef4C2GfHcNjSX8zo1qbYdL5z0RDu/P6cdd5/VhgV5+9lT4sLlC5Bo1+ncPInWabEHBaUun+nPY+2NVWSYV7aMbnWoOlUBp+HiLeujjHA/wV5STJ+300iltbI74rFyU1Wu7GTh6SUeuqSbvL/fIxWsxEknK9XJl/cM4o4Zy/l52wECAQOvSURkt6gEjGDhiodHdcaiqSR88D7XPvwB6w8U4VGjH9q50fj3/9bTslEc53ZqUpe/jhANhgRTosFKibMycUhrxg/KYc2uEgrLPfgDBklOCx2bJmK3BAeTz3+/kYDJclQ0ifqV3gCv/LiF24bmYjnYPGhB3n7W7S41DaQiBWcQ3K7x8Odr+PLuQfX1n0XUAV1TGdGpCSM6NaHS46e40nu4R1lcLZtd6prK4LZpdXJ+4RpuxtIYGCBZKQ/72OShdnpMLeO+/sFCFDVW+uNIQHWL/gWP+q41Pe4r/pE8qLxFnBJa/KKqh4bYeHNVmIa+TbpIMCVOSo3jbbw7vj+b95Xz6vzNvL8sH68/gK4qeP0GKXEWbhnYiit7tyClSin0teUKK+Ka4zHZYlzTfcXlDfDwZ2sY0TFDdjwIUQsSTIkGT1UVOodpcBoIGLz84+aQHJhoE/UBfP4Ac9bu4bzOTQGY+kNeSDl0iC44A9hUUMZvu0tp1ySh1r+zOH4cVi1s3tyJ0r5JAvM2FIRs9YulFLMDF62VnUcf+jCz1aEaK/0BVsXP1dq3POcfwwGj+rZFp0Xla4bwN/Utjq6ovuXe6n8PWUkqrr8kmrxBPJx5b9j3F+Jk0KpxHA+P6szkiztR7vFT4fGRaLccnuQ72kvzNuH11W7SD6CowsOSzYX0zWlUb7+TEKcryZkSIoLFmwtxmZSPjiVRv9zjZ9r8LQDsKq5kyebCkOccCs5Sz7kNZ7sBqFY7iqbjzO1bbQAL4PUFeHneptr9QkIAY/u0QA0zAx1tY+AAKhdrCyK+z0NDbJR7jgzw7h9g5V8L3BxwRc4TdPg9TLV9ydlNLfTISqJbZhJD26bx4IUdmffXC7H2vgH02pT2PZgr1fa8WrxWiONPURTibcEm7+ECqeJKL7N+2R1SzCiW+0qlx8/UuXJfEaI2ZGVKiAjyiyoImKQ1xZKoD7C9sAKA5VsPYNHUkOpKsQRnfiO4VVCI2spKdXJGVjKLTQJ7qLkUs6bABdpSEpXKaj+PtDqUPaU06j5Qmh6gbzedvpeGaaw74u+wYznsXg3+yNv9qrHGwfWfgia3PnH6+HVHMVb92O4rBrBsq/n3gRAiMlmZEiKCYEWk0Fn0qon60ago2M/6QYPY+MhjeCsqQh6PNTgrc5sXEBAiWvcMb4MjzEx3Tay6xoQmG0CNvc9Z1JX+yveFf0y3wvWfQFZvsDhrflPVAo4UuOkrSJOWAuL0UuLymlaZjfW+Yrb9XAhRMwmmhIgg3qabNt6tmqgf1XHSUmn1wYdkjB+HZgvNF4k1OLNokiQsjs2A3MbcNrR1zAGV3aLy8KhOtLvhGXAkgxLbbaRqpb+ItPB5VUBwu951n8C5/wgWk7A4OdjB6ghrXDBHqs94uH0RNO0a07kKcSrQVRWzXbux3ld0ua8IUSuy10GICNpmJJjO+MWSqK8AHZsnY8lIp8kBULVd4Ku+shRrFbXUKlWchKitu87KRVXg2e/yqDTJDaxKUcCmq0y+uBOX9zrYX+qW2fDaBVBWENN2u6Mr/YW+mQpJoY2MQ2g69LoRet4A+Uth5dtQsjNY+tzZGHLPhk6XgF5DYCbEKSwtwWa6gyLW+0qSI/aVZiGEBFNCRNS5eRLNku3kFYSWgE7sMxo1LoXihe+y7/N/oVgd2DJySex/ZbXnOawatw7KAaB/60ZHFyEDYqyiZtEY27tFnf6eomFSFIU7z2pDr+xUnvl2I0u3FGIYVCvbb7eoGAYMaZfGncNy6ZqZfOQAKdkw8UdY8AwseQkMP3jKanzfGvtA6XbodlUsvwhk9Qn+E6KB6dI8CadFp9xdfUIklvuKVVcZ0zOKCQwhRAgJpoSowcQhrfnbp79SbrKfvKZEfYBUp5Xe2cEGpDZd4+o+Wby2YEtIWepog7OAYXBZL7npibrTL6cR/XIasfNAJe8u3c6GvaWUVvpIdFro3CyRy3tl0Tg+zOqOIwWG/xWG/hnWfQG/fgQbZoM3fA8qqKEPVFImNO9xjL+VEA2DqiqMG9SKKbPXh7TxiPa+ogDX9cs+fictxGlEgikhanBRt2Y8Mes3Krx+0y1/kTgsGn84t221Rog3DMjmzYVb8ZqsUdUUnFlUhQu6NCXRLtsxRN1rluzgd+e0rd2LNUtwS12nS+DXmfDxHdUCqqj7QFmcMPB3tTsHIRqoK3tlMWX2etPHoqnO2S+nEU2SwlfYFEKEJwUohKiB3aLx9vh+xFn1o9PbI3JYNK7u24JLuldfRcpMcXL/ee1iTvxXFWicYOOhizrG9DohjrsOo6D1sNh7Qel2aNEful5Z83OFEIelxFl59NIu2C2xD+uSHFaeGCPFWYSoLQmmhIhCbno8H90+gNQ4a403K1UJBlK3DGrFXy7oYPqccQNzuGVQq6hvfBZNIS3BxnsT+pPslOIT4iSnqjBmGrTsD5YoAyqLA5r3hCunQ5SlnIUQR4zukckfz20f9X1FUyDFaeHdCf3ISJRVKSFqS4IpIaLUNiOBb/8wlPvOaUdago04q0bVqukOi4pNVzmvUxNm3NqX+0a0q7a972j3jWjH46O7kpFow2k1HzzadBWrrjK8QwZf3TOYrNQoeuoIcTLQbXDNB9D/zmB5cmu8+fOsccF/vccHG+pa5RoXorZuHtiKZ6/uQWaKA6dVMy2Zbj14XxnUJo2v7hlMm4yE0CcJIaKmGBGSQHr16mUsW7bsOJ6OEKeGQMBgft4+Vm4/QGG5B6dVp0mSnZFdmsZcttwwDBbk7Wfq3Dx+zj9ApcePRVNJcVoZ2zeLsb1b0Chc8n8tKIryk2EYversgCeIfD+dQnxuWPMJLHgairaA1wUWOyRlwYC7oNOl0a9gidPa6fD9dDJ8NxmGwU9bi5j6wyaWbC6kwuNDV1WSHDpjemZxXf+WsholRAwifTdJAQohakFVFQa1SWNQm7RjPpaiKJyZ25gzcxvXwZkJcRLSbdD1iuA/IUS9UxSFXtmp9MpOPdGnIsRpT7b5CSGEEEIIIUQtSDAlhBBCCCGEELUgwZQQQgghhBBC1IIEU0IIIYQQQghRCxJMCSGEEEIIIUQtSDAlhBBCCCGEELUgwZQQQgghhBBC1IIEU0IIIYQQQghRCxJMCSGEEEIIIUQtSDAlhBBCCCGEELWgGIYR/kFFKQC2Hr/TEUIcBy0Nw0g70SdxrOT7SYjT0in//STfTUKclsJ+N0UMpoQQQgghhBBCmJNtfkIIIYQQQghRCxJMCSGEEEIIIUQtSDAlhBBCCCGEELUgwZQQQgghhBBC1IIEU0IIIYQQQghRC/8f3N5qxqo3NRwAAAAASUVORK5CYII=\",\n \"text/plain\": [\n \"
\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"# Take a batch of graphs from the test set, and compute predictions.\\n\",\n \"graphs = next(datasets['test'].take(1).as_numpy_iterator())\\n\",\n \"labels = graphs.globals\\n\",\n \"logits = train.get_predicted_logits(eval_state, graphs, rngs=None)\\n\",\n \"graphs = jraph.unbatch(graphs)\\n\",\n \"\\n\",\n \"# Create plot to visualize individual molecules.\\n\",\n \"ds, ds_info = tfds.load('ogbg_molpcba', split='test', with_info=True)\\n\",\n \"fig = tfds.visualization.show_examples(ds, ds_info,\\n\",\n \" node_color_fn=node_color_fn,\\n\",\n \" node_label_fn=node_label_fn,\\n\",\n \" edge_color_fn=edge_color_fn,\\n\",\n \" rows=10)\\n\",\n \"\\n\",\n \"# Update plot titles with true and predicted labels.\\n\",\n \"for graph, graph_labels, graph_logits, ax in zip(graphs, labels, logits, fig.axes):\\n\",\n \" label_for_task = get_formatted_label_for_task(graph_labels, task)\\n\",\n \" predicted_label_for_task = get_formatted_prediction_for_task(graph_logits, task)\\n\",\n \" ax.set_title(f'True Label: {label_for_task}\\\\n'\\n\",\n \" f'Predicted Label: {predicted_label_for_task}')\"\n ]\n }\n ],\n \"metadata\": {\n \"accelerator\": \"GPU\",\n \"language_info\": {\n \"name\": \"python\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 2\n}\n" }, { "path": "examples/ogbg_molpcba/ogbg_molpcba_benchmark.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Benchmark for the ogbg_molpcba example.\"\"\"\n\nimport time\n\nfrom absl import flags\nfrom absl.testing import absltest\nfrom flax.testing import Benchmark\nimport jax\nimport numpy as np\n\nimport main\nfrom configs import default\nfrom configs import test\n\n\n# Parse absl flags test_srcdir and test_tmpdir.\njax.config.parse_flags_with_absl()\n\nFLAGS = flags.FLAGS\n\n\nclass OgbgMolpcbaBenchmark(Benchmark):\n \"\"\"Benchmarks for the ogbg_molpcba Flax example.\"\"\"\n\n def test_1x_v100(self):\n \"\"\"Run training with default config for ogbg_molpcba on a v100 GPU.\"\"\"\n workdir = self.get_tmp_model_dir()\n config = default.get_config()\n\n FLAGS.workdir = workdir\n FLAGS.config = config\n\n start_time = time.time()\n main.main([])\n benchmark_time = time.time() - start_time\n\n summaries = self.read_summaries(workdir)\n\n # Summaries contain all the information necessary for\n # the regression metrics.\n wall_time, _, test_accuracy = zip(*summaries['test_accuracy'])\n wall_time = np.array(wall_time)\n sec_per_epoch = np.mean(wall_time[1:] - wall_time[:-1])\n end_test_accuracy = test_accuracy[-1]\n\n _, _, test_aps = zip(*summaries['test_mean_average_precision'])\n end_test_mean_average_precision = test_aps[-1]\n\n _, _, validation_accuracy = zip(*summaries['validation_accuracy'])\n end_validation_accuracy = validation_accuracy[-1]\n\n _, _, validation_aps = zip(*summaries['validation_mean_average_precision'])\n end_validation_mean_average_precision = validation_aps[-1]\n\n # Assertions are deferred until the test finishes, so the metrics are\n # always reported and benchmark success is determined based on *all*\n # assertions.\n self.assertGreaterEqual(end_test_mean_average_precision, 0.24)\n self.assertGreaterEqual(end_validation_mean_average_precision, 0.25)\n\n # Use the reporting API to report single or multiple metrics/extras.\n self.report_wall_time(benchmark_time)\n self.report_metrics({\n 'sec_per_epoch': sec_per_epoch,\n 'test_accuracy': end_test_accuracy,\n 'test_mean_average_precision': end_test_mean_average_precision,\n 'validation_accuracy': end_validation_accuracy,\n 'validation_mean_average_precision': (\n end_validation_mean_average_precision\n ),\n })\n self.report_extras({\n 'model_name': 'Graph Convolutional Network',\n 'description': 'GPU (1x V100) test for ogbg_molpcba.',\n 'implementation': 'linen',\n })\n\n def test_cpu(self):\n \"\"\"Run training with test config for ogbg_molpcba on CPU.\"\"\"\n workdir = self.get_tmp_model_dir()\n config = test.get_config()\n\n FLAGS.workdir = workdir\n FLAGS.config = config\n\n start_time = time.time()\n main.main([])\n benchmark_time = time.time() - start_time\n\n summaries = self.read_summaries(workdir)\n\n # Summaries contain all the information necessary for\n # the regression metrics.\n wall_time, _, test_accuracy = zip(*summaries['test_accuracy'])\n wall_time = np.array(wall_time)\n sec_per_epoch = np.mean(wall_time[1:] - wall_time[:-1])\n end_test_accuracy = test_accuracy[-1]\n\n _, _, test_aps = zip(*summaries['test_mean_average_precision'])\n end_test_mean_average_precision = test_aps[-1]\n\n _, _, validation_accuracy = zip(*summaries['validation_accuracy'])\n end_validation_accuracy = validation_accuracy[-1]\n\n _, _, validation_aps = zip(*summaries['validation_mean_average_precision'])\n end_validation_mean_average_precision = validation_aps[-1]\n\n # Use the reporting API to report single or multiple metrics/extras.\n self.report_wall_time(benchmark_time)\n self.report_metrics({\n 'sec_per_epoch': sec_per_epoch,\n 'test_accuracy': end_test_accuracy,\n 'test_mean_average_precision': end_test_mean_average_precision,\n 'validation_accuracy': end_validation_accuracy,\n 'validation_mean_average_precision': (\n end_validation_mean_average_precision\n ),\n })\n self.report_extras({\n 'model_name': 'Graph Convolutional Network',\n 'description': 'CPU test for ogbg_molpcba.',\n 'implementation': 'linen',\n })\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "examples/ogbg_molpcba/requirements.txt", "content": "absl-py==1.0.0\nclu==0.0.6\nflax==0.4.1\njax==0.3.4\n--find-links https://storage.googleapis.com/jax-releases/jax_releases.html\njaxlib==0.3.2+cuda11.cudnn82 # Make sure CUDA version matches the base image.\njraph==0.0.2.dev0\nml-collections==0.1.0\nnumpy==1.22.0\noptax==0.1.0\nsklearn==0.0\ntensorflow==2.11.1\ntensorflow-datasets==4.4.0\n" }, { "path": "examples/ogbg_molpcba/train.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Library file for executing training and evaluation on ogbg-molpcba.\"\"\"\n\nimport os\nfrom typing import Any, Dict, Tuple, Optional\nfrom collections.abc import Iterable\n\nfrom absl import logging\nfrom clu import checkpoint\nfrom clu import metric_writers\nfrom clu import metrics\nfrom clu import parameter_overview\nfrom clu import periodic_actions\nimport flax\nimport flax.core\nimport flax.linen as nn\nfrom flax.training import train_state\nimport jax\nimport jax.numpy as jnp\nimport jraph\nimport ml_collections\nimport numpy as np\nimport optax\nimport sklearn.metrics\nimport tensorflow as tf\n\nimport input_pipeline\nimport models\n\n\ndef create_model(\n config: ml_collections.ConfigDict, deterministic: bool\n) -> nn.Module:\n \"\"\"Creates a Flax model, as specified by the config.\"\"\"\n if config.model == 'GraphNet':\n return models.GraphNet(\n latent_size=config.latent_size,\n num_mlp_layers=config.num_mlp_layers,\n message_passing_steps=config.message_passing_steps,\n output_globals_size=config.num_classes,\n dropout_rate=config.dropout_rate,\n skip_connections=config.skip_connections,\n layer_norm=config.layer_norm,\n use_edge_model=config.use_edge_model,\n deterministic=deterministic,\n )\n if config.model == 'GraphConvNet':\n return models.GraphConvNet(\n latent_size=config.latent_size,\n num_mlp_layers=config.num_mlp_layers,\n message_passing_steps=config.message_passing_steps,\n output_globals_size=config.num_classes,\n dropout_rate=config.dropout_rate,\n skip_connections=config.skip_connections,\n layer_norm=config.layer_norm,\n deterministic=deterministic,\n )\n raise ValueError(f'Unsupported model: {config.model}.')\n\n\ndef create_optimizer(\n config: ml_collections.ConfigDict,\n) -> optax.GradientTransformation:\n \"\"\"Creates an optimizer, as specified by the config.\"\"\"\n if config.optimizer == 'adam':\n return optax.adam(learning_rate=config.learning_rate)\n if config.optimizer == 'sgd':\n return optax.sgd(\n learning_rate=config.learning_rate, momentum=config.momentum\n )\n raise ValueError(f'Unsupported optimizer: {config.optimizer}.')\n\n\ndef binary_cross_entropy_with_mask(\n *, logits: jnp.ndarray, labels: jnp.ndarray, mask: jnp.ndarray\n):\n \"\"\"Binary cross entropy loss for unnormalized logits, with masked elements.\"\"\"\n assert logits.shape == labels.shape == mask.shape\n assert len(logits.shape) == 2\n\n # To prevent propagation of NaNs during grad().\n # We mask over the loss for invalid targets later.\n labels = jnp.where(mask, labels, -1)\n\n # Numerically stable implementation of BCE loss.\n # This mimics TensorFlow's tf.nn.sigmoid_cross_entropy_with_logits().\n positive_logits = logits >= 0\n relu_logits = jnp.where(positive_logits, logits, 0)\n abs_logits = jnp.where(positive_logits, logits, -logits)\n return relu_logits - (logits * labels) + (jnp.log(1 + jnp.exp(-abs_logits)))\n\n\ndef predictions_match_labels(\n *, logits: jnp.ndarray, labels: jnp.ndarray, **kwargs\n) -> jnp.ndarray:\n \"\"\"Returns a binary array indicating where predictions match the labels.\"\"\"\n del kwargs # Unused.\n preds = logits > 0\n return (preds == labels).astype(jnp.float32)\n\n\ndef add_prefix_to_keys(result: dict[str, Any], prefix: str) -> dict[str, Any]:\n \"\"\"Adds a prefix to the keys of a dict, returning a new dict.\"\"\"\n return {f'{prefix}_{key}': val for key, val in result.items()}\n\n\n@flax.struct.dataclass\nclass MeanAveragePrecision(\n metrics.CollectingMetric.from_outputs(('labels', 'logits', 'mask'))\n):\n \"\"\"Computes the mean average precision (mAP) over different tasks.\"\"\"\n\n def compute(self):\n # Matches the official OGB evaluation scheme for mean average precision.\n values = super().compute()\n labels = values['labels']\n logits = values['logits']\n mask = values['mask']\n\n assert logits.shape == labels.shape == mask.shape\n assert len(logits.shape) == 2\n\n probs = jax.nn.sigmoid(logits)\n num_tasks = labels.shape[1]\n average_precisions = np.full(num_tasks, np.nan)\n\n for task in range(num_tasks):\n # AP is only defined when there is at least one negative data\n # and at least one positive data.\n is_labeled = mask[:, task]\n if len(np.unique(labels[is_labeled, task])) >= 2:\n average_precisions[task] = sklearn.metrics.average_precision_score(\n labels[is_labeled, task], probs[is_labeled, task]\n )\n\n # When all APs are NaNs, return NaN. This avoids raising a RuntimeWarning.\n if np.isnan(average_precisions).all():\n return np.nan\n return np.nanmean(average_precisions)\n\n\n@flax.struct.dataclass\nclass EvalMetrics(metrics.Collection):\n accuracy: metrics.Average.from_fun(predictions_match_labels)\n loss: metrics.Average.from_output('loss')\n mean_average_precision: MeanAveragePrecision\n\n\n@flax.struct.dataclass\nclass TrainMetrics(metrics.Collection):\n accuracy: metrics.Average.from_fun(predictions_match_labels)\n loss: metrics.Average.from_output('loss')\n\n\ndef replace_globals(graphs: jraph.GraphsTuple) -> jraph.GraphsTuple:\n \"\"\"Replaces the globals attribute with a constant feature for each graph.\"\"\"\n return graphs._replace(globals=jnp.ones([graphs.n_node.shape[0], 1]))\n\n\ndef get_predicted_logits(\n state: train_state.TrainState,\n graphs: jraph.GraphsTuple,\n rngs: dict[str, jnp.ndarray] | None,\n) -> jnp.ndarray:\n \"\"\"Get predicted logits from the network for input graphs.\"\"\"\n pred_graphs = state.apply_fn(state.params, graphs, rngs=rngs)\n logits = pred_graphs.globals\n return logits\n\n\ndef get_valid_mask(\n labels: jnp.ndarray, graphs: jraph.GraphsTuple\n) -> jnp.ndarray:\n \"\"\"Gets the binary mask indicating only valid labels and graphs.\"\"\"\n # We have to ignore all NaN values - which indicate labels for which\n # the current graphs have no label.\n labels_mask = ~jnp.isnan(labels)\n\n # Since we have extra 'dummy' graphs in our batch due to padding, we want\n # to mask out any loss associated with the dummy graphs.\n # Since we padded with `pad_with_graphs` we can recover the mask by using\n # get_graph_padding_mask.\n graph_mask = jraph.get_graph_padding_mask(graphs)\n\n # Combine the mask over labels with the mask over graphs.\n return labels_mask & graph_mask[:, None]\n\n\n@jax.jit\ndef train_step(\n state: train_state.TrainState,\n graphs: jraph.GraphsTuple,\n rngs: dict[str, jnp.ndarray],\n) -> tuple[train_state.TrainState, metrics.Collection]:\n \"\"\"Performs one update step over the current batch of graphs.\"\"\"\n\n def loss_fn(params, graphs):\n curr_state = state.replace(params=params)\n\n # Extract labels.\n labels = graphs.globals\n\n # Replace the global feature for graph classification.\n graphs = replace_globals(graphs)\n\n # Compute logits and resulting loss.\n logits = get_predicted_logits(curr_state, graphs, rngs)\n mask = get_valid_mask(labels, graphs)\n loss = binary_cross_entropy_with_mask(\n logits=logits, labels=labels, mask=mask\n )\n mean_loss = jnp.sum(jnp.where(mask, loss, 0)) / jnp.sum(mask)\n\n return mean_loss, (loss, logits, labels, mask)\n\n grad_fn = jax.value_and_grad(loss_fn, has_aux=True)\n (_, (loss, logits, labels, mask)), grads = grad_fn(state.params, graphs)\n state = state.apply_gradients(grads=grads)\n\n metrics_update = TrainMetrics.single_from_model_output(\n loss=loss, logits=logits, labels=labels, mask=mask\n )\n return state, metrics_update\n\n\n@jax.jit\ndef evaluate_step(\n state: train_state.TrainState,\n graphs: jraph.GraphsTuple,\n) -> metrics.Collection:\n \"\"\"Computes metrics over a set of graphs.\"\"\"\n\n # The target labels our model has to predict.\n labels = graphs.globals\n\n # Replace the global feature for graph classification.\n graphs = replace_globals(graphs)\n\n # Get predicted logits, and corresponding probabilities.\n logits = get_predicted_logits(state, graphs, rngs=None)\n\n # Get the mask for valid labels and graphs.\n mask = get_valid_mask(labels, graphs)\n\n # Compute the various metrics.\n loss = binary_cross_entropy_with_mask(logits=logits, labels=labels, mask=mask)\n\n return EvalMetrics.single_from_model_output(\n loss=loss, logits=logits, labels=labels, mask=mask\n )\n\n\ndef evaluate_model(\n state: train_state.TrainState,\n datasets: dict[str, tf.data.Dataset],\n splits: Iterable[str],\n) -> dict[str, metrics.Collection]:\n \"\"\"Evaluates the model on metrics over the specified splits.\"\"\"\n\n # Loop over each split independently.\n eval_metrics = {}\n for split in splits:\n split_metrics = None\n\n # Loop over graphs.\n for graphs in datasets[split].as_numpy_iterator():\n split_metrics_update = evaluate_step(state, graphs)\n\n # Update metrics.\n if split_metrics is None:\n split_metrics = split_metrics_update\n else:\n split_metrics = split_metrics.merge(split_metrics_update)\n eval_metrics[split] = split_metrics\n\n return eval_metrics # pytype: disable=bad-return-type\n\n\ndef train_and_evaluate(\n config: ml_collections.ConfigDict, workdir: str\n) -> train_state.TrainState:\n \"\"\"Execute model training and evaluation loop.\n\n Args:\n config: Hyperparameter configuration for training and evaluation.\n workdir: Directory where the TensorBoard summaries are written to.\n\n Returns:\n The train state (which includes the `.params`).\n \"\"\"\n # We only support single-host training.\n assert jax.process_count() == 1\n\n # Create writer for logs.\n writer = metric_writers.create_default_writer(workdir)\n writer.write_hparams(config.to_dict())\n\n # Get datasets, organized by split.\n logging.info('Obtaining datasets.')\n datasets = input_pipeline.get_datasets(\n config.batch_size,\n add_virtual_node=config.add_virtual_node,\n add_undirected_edges=config.add_undirected_edges,\n add_self_loops=config.add_self_loops,\n )\n train_iter = iter(datasets['train'])\n\n # Create and initialize the network.\n logging.info('Initializing network.')\n rng = jax.random.key(0)\n rng, init_rng = jax.random.split(rng)\n init_graphs = next(datasets['train'].as_numpy_iterator())\n init_graphs = replace_globals(init_graphs)\n init_net = create_model(config, deterministic=True)\n params = jax.jit(init_net.init)(init_rng, init_graphs)\n parameter_overview.log_parameter_overview(params)\n\n # Create the optimizer.\n tx = create_optimizer(config)\n\n # Create the training state.\n net = create_model(config, deterministic=False)\n state = train_state.TrainState.create(\n apply_fn=net.apply, params=params, tx=tx\n )\n\n # Set up checkpointing of the model.\n # The input pipeline cannot be checkpointed in its current form,\n # due to the use of stateful operations.\n checkpoint_dir = os.path.join(workdir, 'checkpoints')\n ckpt = checkpoint.Checkpoint(checkpoint_dir, max_to_keep=2)\n state = ckpt.restore_or_initialize(state)\n initial_step = int(state.step) + 1\n\n # Create the evaluation state, corresponding to a deterministic model.\n eval_net = create_model(config, deterministic=True)\n eval_state = state.replace(apply_fn=eval_net.apply)\n\n # Hooks called periodically during training.\n report_progress = periodic_actions.ReportProgress(\n num_train_steps=config.num_train_steps, writer=writer\n )\n profiler = periodic_actions.Profile(num_profile_steps=5, logdir=workdir)\n hooks = [report_progress, profiler]\n\n # Begin training loop.\n logging.info('Starting training.')\n train_metrics = None\n for step in range(initial_step, config.num_train_steps + 1):\n # Split PRNG key, to ensure different 'randomness' for every step.\n rng, dropout_rng = jax.random.split(rng)\n\n # Perform one step of training.\n with jax.profiler.StepTraceAnnotation('train', step_num=step):\n graphs = jax.tree_util.tree_map(np.asarray, next(train_iter))\n state, metrics_update = train_step(\n state, graphs, rngs={'dropout': dropout_rng}\n )\n\n # Update metrics.\n if train_metrics is None:\n train_metrics = metrics_update\n else:\n train_metrics = train_metrics.merge(metrics_update)\n\n # Quick indication that training is happening.\n logging.log_first_n(logging.INFO, 'Finished training step %d.', 10, step)\n for hook in hooks:\n hook(step)\n\n # Log, if required.\n is_last_step = step == config.num_train_steps - 1\n if step % config.log_every_steps == 0 or is_last_step:\n writer.write_scalars(\n step, add_prefix_to_keys(train_metrics.compute(), 'train')\n )\n train_metrics = None\n\n # Evaluate on validation and test splits, if required.\n if step % config.eval_every_steps == 0 or is_last_step:\n eval_state = eval_state.replace(params=state.params)\n\n splits = ['validation', 'test']\n with report_progress.timed('eval'):\n eval_metrics = evaluate_model(eval_state, datasets, splits=splits)\n for split in splits:\n writer.write_scalars(\n step, add_prefix_to_keys(eval_metrics[split].compute(), split)\n )\n\n # Checkpoint model, if required.\n if step % config.checkpoint_every_steps == 0 or is_last_step:\n with report_progress.timed('checkpoint'):\n ckpt.save(state)\n\n return state\n" }, { "path": "examples/ogbg_molpcba/train_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for flax.examples.ogbg_molpcba.train.\"\"\"\n\nimport functools\nimport pathlib\nimport tempfile\nfrom typing import Dict, Optional\nimport warnings\n\nfrom absl.testing import absltest\nfrom absl.testing import parameterized\nimport flax\nfrom flax.training import train_state\nimport jax\nfrom jax import numpy as jnp\nimport jraph\nimport tensorflow as tf\nimport tensorflow_datasets as tfds\nimport numpy as np\n\nfrom configs import default\nfrom configs import test\nimport input_pipeline\nimport train\n\n\ndef average_with_mask(arr: jnp.ndarray, mask: jnp.ndarray):\n \"\"\"Returns the average over elements where mask is True.\"\"\"\n arr = jnp.where(mask, arr, 0)\n return jnp.sum(arr) / jnp.sum(mask)\n\n\ndef get_dummy_raw_datasets(dataset_length) -> dict[str, tf.data.Dataset]:\n \"\"\"Returns dummy datasets, mocking tfds.DatasetBuilder.as_dataset().\"\"\"\n\n # The dummy graph.\n num_nodes = 3\n num_edges = 4\n dummy_graph = {\n 'edge_feat': tf.zeros((num_edges, 3), dtype=tf.float32),\n 'edge_index': tf.zeros((num_edges, 2), dtype=tf.int64),\n 'labels': tf.ones((128,), dtype=tf.float32),\n 'node_feat': tf.zeros((num_nodes, 9), dtype=tf.float32),\n 'num_edges': tf.expand_dims(num_edges, axis=0),\n 'num_nodes': tf.expand_dims(num_nodes, axis=0),\n }\n dummy_graph_spec = {\n 'edge_feat': tf.TensorSpec(shape=(None, 3), dtype=tf.float32),\n 'edge_index': tf.TensorSpec(shape=(None, 2), dtype=tf.int64),\n 'labels': tf.TensorSpec(shape=(128,), dtype=tf.float32),\n 'node_feat': tf.TensorSpec(shape=(None, 9), dtype=tf.float32),\n 'num_edges': tf.TensorSpec(shape=(None,), dtype=tf.int64),\n 'num_nodes': tf.TensorSpec(shape=(None,), dtype=tf.int64),\n }\n\n def get_dummy_graphs():\n for _ in range(dataset_length):\n yield dummy_graph\n\n datasets = {}\n for split in ['train', 'validation', 'test']:\n datasets[split] = tf.data.Dataset.from_generator(\n get_dummy_graphs, output_signature=dummy_graph_spec\n )\n return datasets\n\n\ndef get_dummy_datasets(\n dataset_length: int, batch_size: int | None = None\n) -> dict[str, tf.data.Dataset]:\n \"\"\"Returns dummy datasets, mocking input_pipeline.get_datasets().\"\"\"\n\n datasets = get_dummy_raw_datasets(dataset_length)\n\n # Construct the GraphsTuple converter function.\n convert_to_graphs_tuple_fn = functools.partial(\n input_pipeline.convert_to_graphs_tuple,\n add_virtual_node=True,\n add_undirected_edges=True,\n add_self_loops=True,\n )\n\n # Process each split separately.\n for split_name in datasets:\n # Convert to GraphsTuple.\n datasets[split_name] = datasets[split_name].map(\n convert_to_graphs_tuple_fn,\n num_parallel_calls=tf.data.AUTOTUNE,\n deterministic=True,\n )\n\n # If batch size is None, do not batch.\n if batch_size is not None:\n budget = input_pipeline.estimate_padding_budget_for_batch_size(\n datasets['train'], batch_size, num_estimation_graphs=1\n )\n\n # Pad an example graph to see what the output shapes will be.\n # We will use this shape information when creating the tf.data.Dataset.\n example_graph = next(datasets['train'].as_numpy_iterator())\n example_padded_graph = jraph.pad_with_graphs(example_graph, *budget)\n padded_graphs_spec = input_pipeline.specs_from_graphs_tuple(\n example_padded_graph\n )\n\n # Batch and pad each split separately.\n for split, dataset_split in datasets.items():\n batching_fn = functools.partial(\n jraph.dynamically_batch,\n graphs_tuple_iterator=iter(dataset_split),\n n_node=budget.n_node,\n n_edge=budget.n_edge,\n n_graph=budget.n_graph,\n )\n datasets[split] = tf.data.Dataset.from_generator(\n batching_fn, output_signature=padded_graphs_spec\n )\n return datasets\n\n\nclass OgbgMolpcbaTrainTest(parameterized.TestCase):\n\n def setUp(self):\n super().setUp()\n # Make sure tf does not allocate gpu memory.\n tf.config.experimental.set_visible_devices([], 'GPU')\n\n # Print the current platform (the default device).\n platform = jax.local_devices()[0].platform\n print('Running on platform:', platform.upper())\n\n # Create PRNG keys.\n self.rng = jax.random.key(0)\n\n # Create dummy datasets.\n self.datasets = get_dummy_datasets(dataset_length=20, batch_size=10)\n self.raw_datasets = get_dummy_raw_datasets(dataset_length=20)\n\n @parameterized.product(\n probs=[[[0.8, 0.9, 0.3, 0.5]]],\n labels=[\n [[1, 0, 1, 1]],\n [[1, 0, 1, jnp.nan]],\n [[1, 0, jnp.nan, jnp.nan]],\n [[1, jnp.nan, jnp.nan, jnp.nan]],\n ],\n )\n def test_binary_cross_entropy_loss(self, probs, labels):\n probs = jnp.asarray(probs)\n labels = jnp.asarray(labels)\n\n logits = jnp.log(probs / (1 - probs))\n mask = ~jnp.isnan(labels)\n\n loss_array = train.binary_cross_entropy_with_mask(\n logits=logits, labels=labels, mask=mask\n )\n loss = average_with_mask(loss_array, mask)\n expected_loss_array = -(jnp.log(probs) * labels) - (\n jnp.log(1 - probs) * (1 - labels)\n )\n expected_loss = average_with_mask(expected_loss_array, mask)\n\n self.assertAlmostEqual(loss, expected_loss, places=5)\n\n @parameterized.named_parameters(\n dict(\n testcase_name='no_valid_tasks',\n logits=[[-1.0, 1.0], [1.0, 1.0], [2.0, -1.0]],\n labels=[[jnp.nan, jnp.nan], [jnp.nan, jnp.nan], [jnp.nan, jnp.nan]],\n expected_result=jnp.nan,\n ),\n dict(\n testcase_name='1_valid_task',\n logits=[[-1.0, 1.0], [1.0, 1.0], [2.0, -1.0]],\n labels=[[0, jnp.nan], [1, jnp.nan], [1, jnp.nan]],\n expected_result=1.0,\n ),\n dict(\n testcase_name='2_valid_tasks',\n logits=[[-1.0, 1.0], [1.0, 1.0], [2.0, -1.0]],\n labels=[[0, jnp.nan], [1, 0], [1, 1]],\n expected_result=0.75,\n ),\n )\n def test_mean_average_precision(self, logits, labels, expected_result):\n logits = jnp.asarray(logits)\n labels = jnp.asarray(labels)\n mask = ~jnp.isnan(labels)\n\n mean_average_precision = train.MeanAveragePrecision.from_model_output(\n logits=logits, labels=labels, mask=mask\n ).compute()\n\n if jnp.isnan(expected_result):\n self.assertTrue(jnp.isnan(mean_average_precision))\n else:\n self.assertAlmostEqual(expected_result, mean_average_precision)\n\n @parameterized.parameters(\n dict(\n loss=[[0.5, 1.0], [1.5, 1.3], [2.0, 1.2]],\n logits=[[-1.0, 1.0], [1.0, 1.0], [2.0, 0.0]],\n labels=[[0, jnp.nan], [1, 0], [0, 1]],\n mask=[[True, False], [True, True], [False, False]],\n expected_results={\n 'loss': 1.1,\n 'accuracy': 2 / 3,\n 'mean_average_precision': 1.0,\n },\n ),\n )\n def test_eval_metrics(self, loss, logits, labels, mask, expected_results):\n loss = jnp.asarray(loss)\n logits = jnp.asarray(logits)\n labels = jnp.asarray(labels)\n mask = jnp.asarray(mask)\n\n # Ignore RuntimeWarning caused by MeanAveragePrecision calculation.\n with warnings.catch_warnings():\n warnings.simplefilter('ignore', category=RuntimeWarning)\n eval_metrics = train.EvalMetrics.single_from_model_output(\n loss=loss, logits=logits, labels=labels, mask=mask\n ).compute()\n\n for metric in expected_results:\n self.assertAlmostEqual(expected_results[metric], eval_metrics[metric])\n\n @parameterized.parameters(\n dict(\n loss=[[0.5, 1.0], [1.5, 1.3], [2.0, 1.2]],\n logits=[[-1.0, 1.0], [1.0, 1.0], [2.0, 0.0]],\n labels=[[0, jnp.nan], [1, 0], [0, 1]],\n mask=[[True, False], [True, True], [False, False]],\n expected_results={'loss': 1.1, 'accuracy': 2 / 3},\n ),\n )\n def test_train_metrics(self, loss, logits, labels, mask, expected_results):\n loss = jnp.asarray(loss)\n logits = jnp.asarray(logits)\n labels = jnp.asarray(labels)\n mask = jnp.asarray(mask)\n\n train_metrics = train.TrainMetrics.single_from_model_output(\n loss=loss, logits=logits, labels=labels, mask=mask\n ).compute()\n for metric in expected_results:\n self.assertAlmostEqual(expected_results[metric], train_metrics[metric])\n\n def test_train_step(self):\n # Get the default configuration.\n config = default.get_config()\n\n # Initialize the network with a dummy graph.\n rng, init_rng = jax.random.split(self.rng)\n init_graphs = next(self.datasets['train'].as_numpy_iterator())\n init_graphs_preprocessed = train.replace_globals(init_graphs)\n init_net = train.create_model(config, deterministic=True)\n params = jax.jit(init_net.init)(init_rng, init_graphs_preprocessed)\n\n # Create the optimizer.\n optimizer = train.create_optimizer(config)\n\n # Create the training state.\n net = train.create_model(config, deterministic=False)\n state = train_state.TrainState.create(\n apply_fn=net.apply, params=params, tx=optimizer\n )\n\n # Perform one step of updates.\n # We use the same batch of graphs that we used for initialization.\n state, train_metrics = train.train_step(\n state, init_graphs, rngs={'dropout': rng}\n )\n\n # Check that none of the parameters are NaNs!\n params = flax.core.unfreeze(state.params)\n flat_params = {\n '/'.join(k): v\n for k, v in flax.traverse_util.flatten_dict(params).items()\n }\n for array in flat_params.values():\n self.assertTrue(jnp.all(~jnp.isnan(array)))\n\n # Check that the metrics are well defined.\n train_metrics_vals = train_metrics.compute()\n self.assertGreaterEqual(train_metrics_vals['loss'], 0)\n self.assertGreaterEqual(train_metrics_vals['accuracy'], 0)\n self.assertLessEqual(train_metrics_vals['accuracy'], 1)\n\n def test_evaluate_step(self):\n # Get the default configuration.\n config = default.get_config()\n\n # Initialize the network with a dummy graph.\n _, init_rng = jax.random.split(self.rng)\n init_graphs = next(self.datasets['train'].as_numpy_iterator())\n init_graphs_preprocessed = init_graphs._replace(\n globals=jnp.zeros([init_graphs.n_node.shape[0], 1])\n )\n init_net = train.create_model(config, deterministic=True)\n params = jax.jit(init_net.init)(init_rng, init_graphs_preprocessed)\n\n # Create the optimizer.\n optimizer = train.create_optimizer(config)\n\n # Create the evaluation state.\n eval_net = train.create_model(config, deterministic=True)\n eval_state = train_state.TrainState.create(\n apply_fn=eval_net.apply, params=params, tx=optimizer\n )\n\n # Perform one step of evaluation.\n # We use the same batch of graphs that we used for initialization.\n evaluate_metrics = train.evaluate_step(eval_state, init_graphs)\n\n # Check that the metrics are well defined.\n evaluate_metrics_vals = evaluate_metrics.compute()\n self.assertGreaterEqual(evaluate_metrics_vals['loss'], 0)\n self.assertGreaterEqual(evaluate_metrics_vals['accuracy'], 0)\n self.assertLessEqual(evaluate_metrics_vals['accuracy'], 1)\n self.assertTrue(np.isnan(evaluate_metrics_vals['mean_average_precision']))\n\n def test_train_and_evaluate(self):\n # Create a temporary directory where TensorBoard metrics are written.\n workdir = tempfile.mkdtemp()\n\n # Go two directories up to the root of the flax directory.\n flax_root_dir = pathlib.Path(__file__).parents[2]\n data_dir = str(flax_root_dir) + '/.tfds/metadata' # pylint: disable=unused-variable\n\n # Get the test configuration.\n config = test.get_config()\n\n # Ensure train_and_evaluate() runs without any errors!\n def as_dataset_fn(*args, **kwargs):\n del args\n split = kwargs['split']\n return self.raw_datasets[split]\n\n with tfds.testing.mock_data(as_dataset_fn=as_dataset_fn):\n train.train_and_evaluate(config=config, workdir=workdir)\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "examples/ppo/README.md", "content": "# Proximal Policy Optimization\n\nUses the Proximal Policy Optimization algorithm ([Schulman et al., 2017](https://arxiv.org/abs/1707.06347))\nto learn playing Atari games.\n\n## Requirements\n\nThis example depends on the `gymnasium[atari,accept-rom-license]`, `opencv-python` packages\nin addition to `jax` and `flax`.\n\n## Supported setups\n\nThe example should run with other configurations and hardware, but was explicitly\ntested on the following:\n\n| Hardware | Game | Training time | Total frames seen | TensorBoard.dev |\n| --- | --- | --- | --- | --- |\n| 1x V100 GPU | Breakout | 9h 15m 15s | 40M | [2020-10-02](https://tensorboard.dev/experiment/pY7D2qYQQLO9ZT5lA9PFPA) |\n\n> **Note**\n> It is possible to improve training efficiency through further optimizations. For example, CleanRL's PPO ([ppo_atari_envpool_xla_jax_scan.py](https://docs.cleanrl.dev/rl-algorithms/ppo/#ppo_atari_envpool_xla_jax_scanpy)) can achieve the same level of results in just 30 minutes with an RTX 2080 TI, 8 CPU, and the same hyperparameters — **a 1850% speedup end-to-end**. It achieves this by using [EnvPool](https://envpool.readthedocs.io/en/latest/), a library for fast parallelizing environments, jitting the entire rollout through [EnvPool's XLA interface](https://envpool.readthedocs.io/en/latest/content/xla_interface.html), storing data more efficiently, and `jax.scan`.\n\n## How to run\n\nRunning `python ppo_main.py` will run the example with default\n(hyper)parameters, i.e. for 40M frames on the Pong game.\n\nBy default logging info and checkpoints will be stored in `/tmp/ppo_training`\ndirectory. This can be overridden as follows:\n\n```python ppo_main.py --config=configs/default.py --workdir=/my_fav_directory```\n\nYou can also override the default (hyper)parameters, for example\n\n```python ppo_main.py --config=configs/default.py --config.game=Seaquest --config.total_frames=20000000 --config.decaying_lr_and_clip_param=False --workdir=/tmp/seaquest```\n\nwill train the model on 20M Seaquest frames with constant (i.e. not linearly\ndecaying) learning rate and PPO clipping parameter. Checkpoints and tensorboard\nfiles will be saved in `/tmp/seaquest`.\n\nUnit tests can be run using `python ppo_lib_test.py`.\n\n## How to run on Google Cloud TPU\n\nIt is also possible to run this code on Google Cloud TPU. For detailed\ninstructions on the required setup, please refer to the [WMT example readme](https://github.com/google/flax/tree/main/examples/wmt).\n\n## Owners\n\nJonathan Heek @jheek, Wojciech Rzadkowski @wrzadkow\n" }, { "path": "examples/ppo/agent.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Agent utilities, incl. choosing the move and running in separate process.\"\"\"\n\nimport collections\nimport functools\nimport multiprocessing\nfrom typing import Any\nfrom collections.abc import Callable\n\nimport flax\nimport jax\nimport numpy as np\n\nimport env_utils\n\n\n@functools.partial(jax.jit, static_argnums=0)\ndef policy_action(\n apply_fn: Callable[..., Any],\n params: flax.core.frozen_dict.FrozenDict,\n state: np.ndarray,\n):\n \"\"\"Forward pass of the network.\n\n Args:\n params: the parameters of the actor-critic model\n module: the actor-critic model\n state: the input for the forward pass\n\n Returns:\n out: a tuple (log_probabilities, values)\n \"\"\"\n out = apply_fn({'params': params}, state)\n return out\n\n\nExpTuple = collections.namedtuple(\n 'ExpTuple', ['state', 'action', 'reward', 'value', 'log_prob', 'done']\n)\n\n\nclass RemoteSimulator:\n \"\"\"Wrap functionality for an agent emulating Atari in a separate process.\n\n An object of this class is created for every agent.\n \"\"\"\n\n def __init__(self, game: str):\n \"\"\"Start the remote process and create Pipe() to communicate with it.\"\"\"\n parent_conn, child_conn = multiprocessing.Pipe()\n self.proc = multiprocessing.Process(\n target=rcv_action_send_exp, args=(child_conn, game)\n )\n self.proc.daemon = True\n self.conn = parent_conn\n self.proc.start()\n\n\ndef rcv_action_send_exp(conn, game: str):\n \"\"\"Run the remote agents.\n\n Receive action from the main learner, perform one step of simulation and\n send back collected experience.\n \"\"\"\n env = env_utils.create_env(game, clip_rewards=True)\n while True:\n obs = env.reset()\n done = False\n # Observations fetched from Atari env need additional batch dimension.\n state = obs[None, ...]\n while not done:\n conn.send(state)\n action = conn.recv()\n obs, reward, done, _ = env.step(action)\n next_state = obs[None, ...] if not done else None\n experience = (state, action, reward, done)\n conn.send(experience)\n if done:\n break\n state = next_state\n" }, { "path": "examples/ppo/configs/default.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Definitions of default hyperparameters.\"\"\"\n\nimport ml_collections\n\n\ndef get_config():\n \"\"\"Get the default configuration.\n\n The default hyperparameters originate from PPO paper arXiv:1707.06347\n and openAI baselines 2::\n https://github.com/openai/baselines/blob/master/baselines/ppo2/defaults.py\n \"\"\"\n config = ml_collections.ConfigDict()\n # The Atari game used.\n config.game = 'Pong'\n # Total number of frames seen during training.\n config.total_frames = 40000000\n # The learning rate for the Adam optimizer.\n config.learning_rate = 2.5e-4\n # Batch size used in training.\n config.batch_size = 256\n # Number of agents playing in parallel.\n config.num_agents = 8\n # Number of steps each agent performs in one policy unroll.\n config.actor_steps = 128\n # Number of training epochs per each unroll of the policy.\n config.num_epochs = 3\n # RL discount parameter.\n config.gamma = 0.99\n # Generalized Advantage Estimation parameter.\n config.lambda_ = 0.95\n # The PPO clipping parameter used to clamp ratios in loss function.\n config.clip_param = 0.1\n # Weight of value function loss in the total loss.\n config.vf_coeff = 0.5\n # Weight of entropy bonus in the total loss.\n config.entropy_coeff = 0.01\n # Linearly decay learning rate and clipping parameter to zero during\n # the training.\n config.decaying_lr_and_clip_param = True\n return config\n" }, { "path": "examples/ppo/env_utils.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Utilities for handling the Atari environment.\"\"\"\n\nimport collections\n\nimport gymnasium as gym\nimport numpy as np\n\nimport seed_rl_atari_preprocessing\n\n\nclass ClipRewardEnv(gym.RewardWrapper):\n \"\"\"Adapted from OpenAI baselines.\n\n github.com/openai/baselines/blob/master/baselines/common/atari_wrappers.py\n \"\"\"\n\n def __init__(self, env):\n gym.RewardWrapper.__init__(self, env)\n\n def reward(self, reward):\n \"\"\"Bin reward to {+1, 0, -1} by its sign.\"\"\"\n return np.sign(reward)\n\n\nclass FrameStack:\n \"\"\"Implements stacking of `num_frames` last frames of the game.\n\n Wraps an AtariPreprocessing object.\n \"\"\"\n\n def __init__(\n self,\n preproc: seed_rl_atari_preprocessing.AtariPreprocessing,\n num_frames: int,\n ):\n self.preproc = preproc\n self.num_frames = num_frames\n self.frames = collections.deque(maxlen=num_frames)\n\n def reset(self):\n ob = self.preproc.reset()\n for _ in range(self.num_frames):\n self.frames.append(ob)\n return self._get_array()\n\n def step(self, action: int):\n ob, reward, done, info = self.preproc.step(action)\n self.frames.append(ob)\n return self._get_array(), reward, done, info\n\n def _get_array(self):\n assert len(self.frames) == self.num_frames\n return np.concatenate(self.frames, axis=-1)\n\n\ndef create_env(game: str, clip_rewards: bool):\n \"\"\"Create a FrameStack object that serves as environment for the `game`.\"\"\"\n env = gym.make(game)\n if clip_rewards:\n env = ClipRewardEnv(env) # bin rewards to {-1., 0., 1.}\n preproc = seed_rl_atari_preprocessing.AtariPreprocessing(env)\n stack = FrameStack(preproc, num_frames=4)\n return stack\n\n\ndef get_num_actions(game: str):\n \"\"\"Get the number of possible actions of a given Atari game.\n\n This determines the number of outputs in the actor part of the\n actor-critic model.\n \"\"\"\n env = gym.make(game)\n return env.action_space.n\n" }, { "path": "examples/ppo/models.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Class and functions to define and initialize the actor-critic model.\"\"\"\n\nfrom flax import linen as nn\nimport jax.numpy as jnp\n\n\nclass ActorCritic(nn.Module):\n \"\"\"Class defining the actor-critic model.\"\"\"\n\n num_outputs: int\n\n @nn.compact\n def __call__(self, x):\n \"\"\"Define the convolutional network architecture.\n\n Architecture originates from \"Human-level control through deep reinforcement\n learning.\", Nature 518, no. 7540 (2015): 529-533.\n Note that this is different than the one from \"Playing atari with deep\n reinforcement learning.\" arxiv.org/abs/1312.5602 (2013)\n\n Network is used to both estimate policy (logits) and expected state value;\n in other words, hidden layers' params are shared between policy and value\n networks, see e.g.:\n github.com/openai/baselines/blob/master/baselines/ppo1/cnn_policy.py\n \"\"\"\n dtype = jnp.float32\n x = x.astype(dtype) / 255.0\n x = nn.Conv(\n features=32,\n kernel_size=(8, 8),\n strides=(4, 4),\n name='conv1',\n dtype=dtype,\n )(x)\n x = nn.relu(x)\n x = nn.Conv(\n features=64,\n kernel_size=(4, 4),\n strides=(2, 2),\n name='conv2',\n dtype=dtype,\n )(x)\n x = nn.relu(x)\n x = nn.Conv(\n features=64,\n kernel_size=(3, 3),\n strides=(1, 1),\n name='conv3',\n dtype=dtype,\n )(x)\n x = nn.relu(x)\n x = x.reshape((x.shape[0], -1)) # flatten\n x = nn.Dense(features=512, name='hidden', dtype=dtype)(x)\n x = nn.relu(x)\n logits = nn.Dense(features=self.num_outputs, name='logits', dtype=dtype)(x)\n policy_log_probabilities = nn.log_softmax(logits)\n value = nn.Dense(features=1, name='value', dtype=dtype)(x)\n return policy_log_probabilities, value\n" }, { "path": "examples/ppo/ppo_lib.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Library file which executes the PPO training.\"\"\"\n\nimport functools\nfrom typing import Any\nfrom collections.abc import Callable\n\nfrom absl import logging\nimport flax\nfrom flax import linen as nn\nimport agent\nimport models\nimport test_episodes\nfrom flax.metrics import tensorboard\nfrom flax.training import checkpoints\nfrom flax.training import train_state\nimport jax\nimport jax.numpy as jnp\nimport ml_collections\nimport numpy as np\nimport optax\n\n\n@jax.jit\n@functools.partial(jax.vmap, in_axes=(1, 1, 1, None, None), out_axes=1)\ndef gae_advantages(\n rewards: np.ndarray,\n terminal_masks: np.ndarray,\n values: np.ndarray,\n discount: float,\n gae_param: float,\n):\n \"\"\"Use Generalized Advantage Estimation (GAE) to compute advantages.\n\n As defined by eqs. (11-12) in PPO paper arXiv: 1707.06347. Implementation uses\n key observation that A_{t} = delta_t + gamma*lambda*A_{t+1}.\n\n Args:\n rewards: array shaped (actor_steps, num_agents), rewards from the game\n terminal_masks: array shaped (actor_steps, num_agents), zeros for terminal\n and ones for non-terminal states\n values: array shaped (actor_steps, num_agents), values estimated by critic\n discount: RL discount usually denoted with gamma\n gae_param: GAE parameter usually denoted with lambda\n\n Returns:\n advantages: calculated advantages shaped (actor_steps, num_agents)\n \"\"\"\n assert rewards.shape[0] + 1 == values.shape[0], (\n 'One more value needed; Eq. '\n '(12) in PPO paper requires '\n 'V(s_{t+1}) for delta_t'\n )\n advantages = []\n gae = 0.0\n for t in reversed(range(len(rewards))):\n # Masks used to set next state value to 0 for terminal states.\n value_diff = discount * values[t + 1] * terminal_masks[t] - values[t]\n delta = rewards[t] + value_diff\n # Masks[t] used to ensure that values before and after a terminal state\n # are independent of each other.\n gae = delta + discount * gae_param * terminal_masks[t] * gae\n advantages.append(gae)\n advantages = advantages[::-1]\n return jnp.array(advantages)\n\n\ndef loss_fn(\n params: flax.core.FrozenDict,\n apply_fn: Callable[..., Any],\n minibatch: tuple,\n clip_param: float,\n vf_coeff: float,\n entropy_coeff: float,\n):\n \"\"\"Evaluate the loss function.\n\n Compute loss as a sum of three components: the negative of the PPO clipped\n surrogate objective, the value function loss and the negative of the entropy\n bonus.\n\n Args:\n params: the parameters of the actor-critic model\n apply_fn: the actor-critic model's apply function\n minibatch: tuple of five elements forming one experience batch:\n states: shape (batch_size, 84, 84, 4)\n actions: shape (batch_size, 84, 84, 4)\n old_log_probs: shape (batch_size,)\n returns: shape (batch_size,)\n advantages: shape (batch_size,)\n clip_param: the PPO clipping parameter used to clamp ratios in loss function\n vf_coeff: weighs value function loss in total loss\n entropy_coeff: weighs entropy bonus in the total loss\n\n Returns:\n loss: the PPO loss, scalar quantity\n \"\"\"\n states, actions, old_log_probs, returns, advantages = minibatch\n log_probs, values = agent.policy_action(apply_fn, params, states)\n values = values[:, 0] # Convert shapes: (batch, 1) to (batch, ).\n probs = jnp.exp(log_probs)\n\n value_loss = jnp.mean(jnp.square(returns - values), axis=0)\n\n entropy = jnp.sum(-probs * log_probs, axis=1).mean()\n\n log_probs_act_taken = jax.vmap(lambda lp, a: lp[a])(log_probs, actions)\n ratios = jnp.exp(log_probs_act_taken - old_log_probs)\n # Advantage normalization (following the OpenAI baselines).\n advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)\n pg_loss = ratios * advantages\n clipped_loss = advantages * jax.lax.clamp(\n 1.0 - clip_param, ratios, 1.0 + clip_param\n )\n ppo_loss = -jnp.mean(jnp.minimum(pg_loss, clipped_loss), axis=0)\n\n return ppo_loss + vf_coeff * value_loss - entropy_coeff * entropy\n\n\n@functools.partial(jax.jit, static_argnums=(2,))\ndef train_step(\n state: train_state.TrainState,\n trajectories: tuple,\n batch_size: int,\n *,\n clip_param: float,\n vf_coeff: float,\n entropy_coeff: float,\n):\n \"\"\"Compilable train step.\n\n Runs an entire epoch of training (i.e. the loop over minibatches within\n an epoch is included here for performance reasons).\n\n Args:\n state: the train state\n trajectories: tuple of the following five elements forming the experience:\n states: shape (steps_per_agent*num_agents, 84, 84, 4)\n actions: shape (steps_per_agent*num_agents, 84, 84, 4)\n old_log_probs: shape (steps_per_agent*num_agents, )\n returns: shape (steps_per_agent*num_agents, )\n advantages: (steps_per_agent*num_agents, )\n batch_size: the minibatch size, static argument\n clip_param: the PPO clipping parameter used to clamp ratios in loss function\n vf_coeff: weighs value function loss in total loss\n entropy_coeff: weighs entropy bonus in the total loss\n\n Returns:\n optimizer: new optimizer after the parameters update\n loss: loss summed over training steps\n \"\"\"\n iterations = trajectories[0].shape[0] // batch_size\n trajectories = jax.tree_util.tree_map(\n lambda x: x.reshape((iterations, batch_size) + x.shape[1:]), trajectories\n )\n loss = 0.0\n for batch in zip(*trajectories):\n grad_fn = jax.value_and_grad(loss_fn)\n l, grads = grad_fn(\n state.params, state.apply_fn, batch, clip_param, vf_coeff, entropy_coeff\n )\n loss += l\n state = state.apply_gradients(grads=grads)\n return state, loss\n\n\ndef get_experience(\n state: train_state.TrainState,\n simulators: list[agent.RemoteSimulator],\n steps_per_actor: int,\n):\n \"\"\"Collect experience from agents.\n\n Runs `steps_per_actor` time steps of the game for each of the `simulators`.\n \"\"\"\n all_experience = []\n # Range up to steps_per_actor + 1 to get one more value needed for GAE.\n for _ in range(steps_per_actor + 1):\n sim_states = []\n for sim in simulators:\n sim_state = sim.conn.recv()\n sim_states.append(sim_state)\n sim_states = np.concatenate(sim_states, axis=0)\n log_probs, values = agent.policy_action(\n state.apply_fn, state.params, sim_states\n )\n log_probs, values = jax.device_get((log_probs, values))\n probs = np.exp(np.array(log_probs))\n for i, sim in enumerate(simulators):\n probabilities = probs[i]\n action = np.random.choice(probs.shape[1], p=probabilities)\n sim.conn.send(action)\n experiences = []\n for i, sim in enumerate(simulators):\n sim_state, action, reward, done = sim.conn.recv()\n value = values[i, 0]\n log_prob = log_probs[i][action]\n sample = agent.ExpTuple(sim_state, action, reward, value, log_prob, done)\n experiences.append(sample)\n all_experience.append(experiences)\n return all_experience\n\n\ndef process_experience(\n experience: list[list[agent.ExpTuple]],\n actor_steps: int,\n num_agents: int,\n gamma: float,\n lambda_: float,\n):\n \"\"\"Process experience for training, including advantage estimation.\n\n Args:\n experience: collected from agents in the form of nested lists/namedtuple\n actor_steps: number of steps each agent has completed\n num_agents: number of agents that collected experience\n gamma: dicount parameter\n lambda_: GAE parameter\n\n Returns:\n trajectories: trajectories readily accessible for `train_step()` function\n \"\"\"\n obs_shape = (84, 84, 4)\n exp_dims = (actor_steps, num_agents)\n values_dims = (actor_steps + 1, num_agents)\n states = np.zeros(exp_dims + obs_shape, dtype=np.float32)\n actions = np.zeros(exp_dims, dtype=np.int32)\n rewards = np.zeros(exp_dims, dtype=np.float32)\n values = np.zeros(values_dims, dtype=np.float32)\n log_probs = np.zeros(exp_dims, dtype=np.float32)\n dones = np.zeros(exp_dims, dtype=np.float32)\n\n for t in range(len(experience) - 1): # experience[-1] only for next_values\n for agent_id, exp_agent in enumerate(experience[t]):\n states[t, agent_id, ...] = exp_agent.state\n actions[t, agent_id] = exp_agent.action\n rewards[t, agent_id] = exp_agent.reward\n values[t, agent_id] = exp_agent.value\n log_probs[t, agent_id] = exp_agent.log_prob\n # Dones need to be 0 for terminal states.\n dones[t, agent_id] = float(not exp_agent.done)\n for a in range(num_agents):\n values[-1, a] = experience[-1][a].value\n advantages = gae_advantages(rewards, dones, values, gamma, lambda_)\n returns = advantages + values[:-1, :]\n # After preprocessing, concatenate data from all agents.\n trajectories = (states, actions, log_probs, returns, advantages)\n trajectory_len = num_agents * actor_steps\n trajectories = tuple(\n map(\n lambda x: np.reshape(x, (trajectory_len,) + x.shape[2:]), trajectories\n )\n )\n return trajectories\n\n\n@functools.partial(jax.jit, static_argnums=1)\ndef get_initial_params(key: jax.Array, model: nn.Module):\n input_dims = (1, 84, 84, 4) # (minibatch, height, width, stacked frames)\n init_shape = jnp.ones(input_dims, jnp.float32)\n initial_params = model.init(key, init_shape)['params']\n return initial_params\n\n\ndef create_train_state(\n params,\n model: nn.Module,\n config: ml_collections.ConfigDict,\n train_steps: int,\n) -> train_state.TrainState:\n if config.decaying_lr_and_clip_param:\n lr = optax.linear_schedule(\n init_value=config.learning_rate,\n end_value=0.0,\n transition_steps=train_steps,\n )\n else:\n lr = config.learning_rate\n tx = optax.adam(lr)\n state = train_state.TrainState.create(\n apply_fn=model.apply, params=params, tx=tx\n )\n return state\n\n\ndef train(\n model: models.ActorCritic, config: ml_collections.ConfigDict, model_dir: str\n):\n \"\"\"Main training loop.\n\n Args:\n model: the actor-critic model\n config: object holding hyperparameters and the training information\n model_dir: path to dictionary where checkpoints and logging info are stored\n\n Returns:\n optimizer: the trained optimizer\n \"\"\"\n\n game = config.game + 'NoFrameskip-v4'\n simulators = [agent.RemoteSimulator(game) for _ in range(config.num_agents)]\n summary_writer = tensorboard.SummaryWriter(model_dir)\n summary_writer.hparams(dict(config))\n loop_steps = config.total_frames // (config.num_agents * config.actor_steps)\n log_frequency = 40\n checkpoint_frequency = 500\n # train_step does multiple steps per call for better performance\n # compute number of steps per call here to convert between the number of\n # train steps and the inner number of optimizer steps\n iterations_per_step = (\n config.num_agents * config.actor_steps // config.batch_size\n )\n\n initial_params = get_initial_params(jax.random.key(0), model)\n state = create_train_state(\n initial_params,\n model,\n config,\n loop_steps * config.num_epochs * iterations_per_step,\n )\n del initial_params\n state = checkpoints.restore_checkpoint(model_dir, state)\n # number of train iterations done by each train_step\n\n start_step = int(state.step) // config.num_epochs // iterations_per_step\n logging.info('Start training from step: %s', start_step)\n\n for step in range(start_step, loop_steps):\n # Bookkeeping and testing.\n if step % log_frequency == 0:\n score = test_episodes.policy_test(1, state.apply_fn, state.params, game)\n frames = step * config.num_agents * config.actor_steps\n summary_writer.scalar('game_score', score, frames)\n logging.info(\n 'Step %s:\\nframes seen %s\\nscore %s\\n\\n', step, frames, score\n )\n\n # Core training code.\n alpha = (\n 1.0 - step / loop_steps if config.decaying_lr_and_clip_param else 1.0\n )\n all_experiences = get_experience(state, simulators, config.actor_steps)\n trajectories = process_experience(\n all_experiences,\n config.actor_steps,\n config.num_agents,\n config.gamma,\n config.lambda_,\n )\n clip_param = config.clip_param * alpha\n for _ in range(config.num_epochs):\n permutation = np.random.permutation(\n config.num_agents * config.actor_steps\n )\n trajectories = tuple(x[permutation] for x in trajectories)\n state, _ = train_step(\n state,\n trajectories,\n config.batch_size,\n clip_param=clip_param,\n vf_coeff=config.vf_coeff,\n entropy_coeff=config.entropy_coeff,\n )\n if (step + 1) % checkpoint_frequency == 0:\n checkpoints.save_checkpoint(model_dir, state, step + 1)\n return state\n" }, { "path": "examples/ppo/ppo_lib_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Unit tests for the PPO example.\"\"\"\n\nfrom absl.testing import absltest\nfrom flax.training import train_state\nimport jax\nimport ml_collections\nimport numpy as np\nimport numpy.testing as np_testing\n\nimport agent\nimport env_utils\nimport models\nimport ppo_lib\nimport gymnasium as gym\nimport ale_py\n\ngym.register_envs(ale_py)\n\n\n# test GAE\nclass TestGAE(absltest.TestCase):\n\n def test_gae_shape_on_random(self):\n # create random data, simulating 4 parallel envs and 20 time_steps\n envs, steps = 10, 100\n rewards = np.random.choice(\n [-1.0, 0.0, 1.0], size=(steps, envs), p=[0.01, 0.98, 0.01]\n )\n terminal_masks = np.ones(shape=(steps, envs), dtype=np.float64)\n values = np.random.random(size=(steps + 1, envs))\n discount = 0.99\n gae_param = 0.95\n adv = ppo_lib.gae_advantages(\n rewards, terminal_masks, values, discount, gae_param\n )\n self.assertEqual(adv.shape, (steps, envs))\n\n def test_gae_hardcoded(self):\n # test on small example that can be verified by hand\n rewards = np.array([[1.0, 0.0], [0.0, 0.0], [-1.0, 1.0]])\n # one of the two episodes terminated in the middle\n terminal_masks = np.array([[1.0, 1.0], [0.0, 1.0], [1.0, 1.0]])\n values = np.array([[1.0, 1.0], [1.0, 1.0], [1.0, 1.0], [1.0, 1.0]])\n discount = 0.5\n gae_param = 0.25\n correct_gae = np.array([[0.375, -0.5546875], [-1.0, -0.4375], [-1.5, 0.5]])\n actual_gae = ppo_lib.gae_advantages(\n rewards, terminal_masks, values, discount, gae_param\n )\n np_testing.assert_allclose(actual_gae, correct_gae)\n\n\n# test environment and preprocessing\nclass TestEnvironmentPreprocessing(absltest.TestCase):\n\n def choose_random_game(self):\n games = [\n 'BeamRider',\n 'Breakout',\n 'Pong',\n 'Qbert',\n 'Seaquest',\n 'SpaceInvaders',\n ]\n ind = np.random.choice(len(games))\n return games[ind] + 'NoFrameskip-v4'\n\n def test_creation(self):\n frame_shape = (84, 84, 4)\n game = self.choose_random_game()\n env = env_utils.create_env(game, clip_rewards=True)\n obs = env.reset()\n self.assertEqual(obs.shape, frame_shape)\n\n def test_step(self):\n frame_shape = (84, 84, 4)\n game = self.choose_random_game()\n env = env_utils.create_env(game, clip_rewards=True)\n obs = env.reset()\n actions = [1, 2, 3, 0]\n for a in actions:\n obs, reward, done, info = env.step(a)\n self.assertEqual(obs.shape, frame_shape)\n self.assertTrue(reward <= 1.0 and reward >= -1.0)\n self.assertTrue(isinstance(done, bool))\n self.assertTrue(isinstance(info, dict))\n\n\n# test the model (creation and forward pass)\nclass TestModel(absltest.TestCase):\n\n def choose_random_outputs(self):\n return np.random.choice([4, 5, 6, 7, 8, 9])\n\n def test_model(self):\n outputs = self.choose_random_outputs()\n module = models.ActorCritic(num_outputs=outputs)\n params = ppo_lib.get_initial_params(jax.random.key(0), module)\n test_batch_size, obs_shape = 10, (84, 84, 4)\n random_input = np.random.random(size=(test_batch_size,) + obs_shape)\n log_probs, values = agent.policy_action(module.apply, params, random_input)\n self.assertEqual(values.shape, (test_batch_size, 1))\n sum_probs = np.sum(np.exp(log_probs), axis=1)\n self.assertEqual(sum_probs.shape, (test_batch_size,))\n np_testing.assert_allclose(\n sum_probs, np.ones((test_batch_size,)), atol=1e-6\n )\n\n\n# test one optimization step\nclass TestOptimizationStep(absltest.TestCase):\n\n def generate_random_data(self, num_actions):\n data_len = 256 # equal to one default-sized batch\n state_shape = (84, 84, 4)\n states = np.random.randint(0, 255, size=((data_len,) + state_shape))\n actions = np.random.choice(num_actions, size=data_len)\n old_log_probs = np.random.random(size=data_len)\n returns = np.random.random(size=data_len)\n advantages = np.random.random(size=data_len)\n return states, actions, old_log_probs, returns, advantages\n\n def test_optimization_step(self):\n num_outputs = 4\n trn_data = self.generate_random_data(num_actions=num_outputs)\n clip_param = 0.1\n vf_coeff = 0.5\n entropy_coeff = 0.01\n batch_size = 256\n module = models.ActorCritic(num_outputs)\n initial_params = ppo_lib.get_initial_params(jax.random.key(0), module)\n config = ml_collections.ConfigDict({\n 'learning_rate': 2.5e-4,\n 'decaying_lr_and_clip_param': True,\n })\n state = ppo_lib.create_train_state(initial_params, module, config, 1000)\n state, _ = ppo_lib.train_step(\n state,\n trn_data,\n batch_size,\n clip_param=clip_param,\n vf_coeff=vf_coeff,\n entropy_coeff=entropy_coeff,\n )\n self.assertIsInstance(state, train_state.TrainState)\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "examples/ppo/ppo_main.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# See issue #620.\n# pytype: disable=wrong-keyword-args\n\nfrom absl import app\nfrom absl import flags\nfrom ml_collections import config_flags\nimport tensorflow as tf\n\nimport env_utils\nimport models\nimport ppo_lib\nimport gymnasium as gym\nimport ale_py\n\ngym.register_envs(ale_py)\n\n\nFLAGS = flags.FLAGS\n\nflags.DEFINE_string(\n 'workdir',\n default='/tmp/ppo_training',\n help='Directory to save checkpoints and logging info.',\n)\n\nconfig_flags.DEFINE_config_file(\n 'config',\n 'configs/default.py',\n 'File path to the default configuration file.',\n lock_config=True,\n)\n\n\ndef main(argv):\n # Make sure tf does not allocate gpu memory.\n tf.config.experimental.set_visible_devices([], 'GPU')\n config = FLAGS.config\n game = config.game + 'NoFrameskip-v4'\n num_actions = env_utils.get_num_actions(game)\n print(f'Playing {game} with {num_actions} actions')\n model = models.ActorCritic(num_outputs=num_actions)\n ppo_lib.train(model, config, FLAGS.workdir)\n\n\nif __name__ == '__main__':\n app.run(main)\n" }, { "path": "examples/ppo/requirements.txt", "content": "absl-py==1.0.0\natari-py==0.2.5\nopencv-python==4.5.4.60\nflax==0.3.6\ngym==0.18.3\ngymnasium[atari, accept-rom-license]==0.29.0\nml-collections==0.1.0\nnumpy==1.22.0\noptax==0.1.5\ntensorflow==2.11.1\n" }, { "path": "examples/ppo/seed_rl_atari_preprocessing.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# coding=utf-8\n# Copyright 2019 The SEED Authors\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"A class implementing minimal Atari 2600 preprocessing.\nAdapted from SEED RL, originally adapted from Dopamine.\n\"\"\"\n\nimport cv2\nimport gymnasium as gym\nfrom gymnasium.spaces.box import Box\nimport numpy as np\n\n\nclass AtariPreprocessing:\n \"\"\"A class implementing image preprocessing for Atari 2600 agents.\n Specifically, this provides the following subset from the JAIR paper\n (Bellemare et al., 2013) and Nature DQN paper (Mnih et al., 2015):\n * Frame skipping (defaults to 4).\n * Terminal signal when a life is lost (off by default).\n * Grayscale and max-pooling of the last two frames.\n * Downsample the screen to a square image (defaults to 84x84).\n More generally, this class follows the preprocessing guidelines set down in\n Machado et al. (2018), \"Revisiting the Arcade Learning Environment:\n Evaluation Protocols and Open Problems for General Agents\".\n It also provides random starting no-ops, which are used in the Rainbow, Apex\n and R2D2 papers.\n \"\"\"\n\n def __init__(\n self,\n environment: gym.Env,\n frame_skip=4,\n terminal_on_life_loss=False,\n screen_size=84,\n max_random_noops=0,\n ):\n \"\"\"Constructor for an Atari 2600 preprocessor.\n Args:\n environment: Gym environment whose observations are preprocessed.\n frame_skip: int, the frequency at which the agent experiences the game.\n terminal_on_life_loss: bool, If True, the step() method returns\n is_terminal=True whenever a life is lost. See Mnih et al. 2015.\n screen_size: int, size of a resized Atari 2600 frame.\n max_random_noops: int, maximum number of no-ops to apply at the beginning\n of each episode to reduce determinism. These no-ops are applied at a\n low-level, before frame skipping.\n Raises:\n ValueError: if frame_skip or screen_size are not strictly positive.\n \"\"\"\n if frame_skip <= 0:\n raise ValueError(\n f'Frame skip should be strictly positive, got {frame_skip}'\n )\n if screen_size <= 0:\n raise ValueError(\n 'Target screen size should be strictly positive, got {}'.format(\n screen_size\n )\n )\n\n self.environment = environment\n self.terminal_on_life_loss = terminal_on_life_loss\n self.frame_skip = frame_skip\n self.screen_size = screen_size\n self.max_random_noops = max_random_noops\n\n obs_dims = self.environment.observation_space\n # Stores temporary observations used for pooling over two successive\n # frames.\n self.screen_buffer = [\n np.empty((obs_dims.shape[0], obs_dims.shape[1]), dtype=np.uint8),\n np.empty((obs_dims.shape[0], obs_dims.shape[1]), dtype=np.uint8),\n ]\n\n self.game_over = False\n self.lives = 0 # Will need to be set by reset().\n\n @property\n def observation_space(self):\n # Return the observation space adjusted to match the shape of the processed\n # observations.\n return Box(\n low=0,\n high=255,\n shape=(self.screen_size, self.screen_size, 1),\n dtype=np.uint8,\n )\n\n @property\n def action_space(self):\n return self.environment.action_space\n\n @property\n def reward_range(self):\n return self.environment.reward_range # type: ignore\n\n @property\n def metadata(self):\n return self.environment.metadata\n\n def close(self):\n return self.environment.close()\n\n def apply_random_noops(self):\n \"\"\"Steps self.environment with random no-ops.\"\"\"\n if self.max_random_noops <= 0:\n return\n # Other no-ops implementations actually always do at least 1 no-op. We\n # follow them.\n no_ops = self.environment.np_random.randint(1, self.max_random_noops + 1)\n for _ in range(no_ops):\n _, _, game_over, _, _ = self.environment.step(0)\n if game_over:\n self.environment.reset()\n\n def reset(self):\n \"\"\"Resets the environment.\n Returns:\n observation: numpy array, the initial observation emitted by the\n environment.\n \"\"\"\n self.environment.reset()\n self.apply_random_noops()\n\n self.lives = self.environment.unwrapped.ale.lives() # pytype:disable=attribute-error\n self._fetch_grayscale_observation(self.screen_buffer[0])\n self.screen_buffer[1].fill(0)\n return self._pool_and_resize()\n\n def render(self, mode):\n \"\"\"Renders the current screen, before preprocessing.\n This calls the Gym API's render() method.\n Args:\n mode: Mode argument for the environment's render() method.\n Valid values (str) are:\n 'rgb_array': returns the raw ALE image.\n 'human': renders to display via the Gym renderer.\n Returns:\n if mode='rgb_array': numpy array, the most recent screen.\n if mode='human': bool, whether the rendering was successful.\n \"\"\"\n return self.environment.render(mode) # pytype:disable=wrong-arg-count\n\n def step(self, action):\n \"\"\"Applies the given action in the environment.\n Remarks:\n * If a terminal state (from life loss or episode end) is reached, this may\n execute fewer than self.frame_skip steps in the environment.\n * Furthermore, in this case the returned observation may not contain valid\n image data and should be ignored.\n Args:\n action: The action to be executed.\n Returns:\n observation: numpy array, the observation following the action.\n reward: float, the reward following the action.\n is_terminal: bool, whether the environment has reached a terminal state.\n This is true when a life is lost and terminal_on_life_loss, or when the\n episode is over.\n info: Gym API's info data structure.\n \"\"\"\n accumulated_reward = 0.0\n\n for time_step in range(self.frame_skip):\n # We bypass the Gym observation altogether and directly fetch the\n # grayscale image from the ALE. This is a little faster.\n _, reward, game_over, _, info = self.environment.step(action)\n accumulated_reward += float(reward)\n\n if self.terminal_on_life_loss:\n new_lives = self.environment.unwrapped.ale.lives() # pytype:disable=attribute-error\n is_terminal = game_over or new_lives < self.lives\n self.lives = new_lives\n else:\n is_terminal = game_over\n\n if is_terminal:\n break\n # We max-pool over the last two frames, in grayscale.\n elif time_step >= self.frame_skip - 2:\n t = time_step - (self.frame_skip - 2)\n self._fetch_grayscale_observation(self.screen_buffer[t])\n\n # Pool the last two observations.\n observation = self._pool_and_resize()\n\n self.game_over = game_over\n return observation, accumulated_reward, is_terminal, info\n\n def _fetch_grayscale_observation(self, output):\n \"\"\"Returns the current observation in grayscale.\n The returned observation is stored in 'output'.\n Args:\n output: numpy array, screen buffer to hold the returned observation.\n Returns:\n observation: numpy array, the current observation in grayscale.\n \"\"\"\n self.environment.unwrapped.ale.getScreenGrayscale(output) # pytype:disable=attribute-error\n return output\n\n def _pool_and_resize(self):\n \"\"\"Transforms two frames into a Nature DQN observation.\n For efficiency, the transformation is done in-place in self.screen_buffer.\n Returns:\n transformed_screen: numpy array, pooled, resized screen.\n \"\"\"\n # Pool if there are enough screens to do so.\n if self.frame_skip > 1:\n np.maximum(\n self.screen_buffer[0],\n self.screen_buffer[1],\n out=self.screen_buffer[0],\n )\n\n transformed_image = cv2.resize(\n self.screen_buffer[0],\n (self.screen_size, self.screen_size),\n interpolation=cv2.INTER_LINEAR,\n )\n int_image = np.asarray(transformed_image, dtype=np.uint8)\n return np.expand_dims(int_image, axis=2)\n" }, { "path": "examples/ppo/test_episodes.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Test policy by playing a full Atari game.\"\"\"\n\nimport itertools\nfrom typing import Any\nfrom collections.abc import Callable\n\nimport flax\nimport numpy as np\n\nimport agent\nimport env_utils\n\n\ndef policy_test(\n n_episodes: int,\n apply_fn: Callable[..., Any],\n params: flax.core.frozen_dict.FrozenDict,\n game: str,\n):\n \"\"\"Perform a test of the policy in Atari environment.\n\n Args:\n n_episodes: number of full Atari episodes to test on\n apply_fn: the actor-critic apply function\n params: actor-critic model parameters, they define the policy being tested\n game: defines the Atari game to test on\n\n Returns:\n total_reward: obtained score\n \"\"\"\n test_env = env_utils.create_env(game, clip_rewards=False)\n for _ in range(n_episodes):\n obs = test_env.reset()\n state = obs[None, ...] # add batch dimension\n total_reward = 0.0\n for t in itertools.count():\n log_probs, _ = agent.policy_action(apply_fn, params, state)\n probs = np.exp(np.array(log_probs, dtype=np.float32))\n probabilities = probs[0] / probs[0].sum()\n action = np.random.choice(probs.shape[1], p=probabilities)\n obs, reward, done, _ = test_env.step(action)\n total_reward += reward\n next_state = obs[None, ...] if not done else None\n state = next_state\n if done:\n break\n return total_reward\n" }, { "path": "examples/seq2seq/README.md", "content": "## seq2seq addition\n\nThis example trains a simple LSTM on a sequence-to-sequence addition task using\nan encoder-decoder architecture. The data is generated on the fly.\n\nColab lets you edit the source files and interact with the model:\n\nhttps://colab.research.google.com/github/google/flax/blob/main/examples/seq2seq/seq2seq.ipynb\n\n### Example output\n\nFrom Colab run that also generated [tfhub.dev]\n\n```\nINFO:absl:[1800] accuracy=1.0, loss=0.0020284138154238462\nINFO:absl:DECODE: 14+381 = 395 (CORRECT)\nINFO:absl:DECODE: 68+91 = 159 (CORRECT)\nINFO:absl:DECODE: 0+807 = 707 (INCORRECT) correct=807\nINFO:absl:DECODE: 95+532 = 627 (CORRECT)\nINFO:absl:DECODE: 6+600 = 606 (CORRECT)\n```\n\n[tfhub.dev]: https://tensorboard.dev/experiment/TwvKVBqzTaKWgEbyebillw/#scalars&_smoothingWeight=0\n\n### How to run\n\n`python train.py`\n\nThe total runtime for 1200 steps on CPU (3.5GHz Intel Core i7, 16GB memory) is\nabout 4 minutes.\n" }, { "path": "examples/seq2seq/configs/default.py", "content": "# Copyright 2025 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Default Hyperparameter configuration.\"\"\"\n\nimport ml_collections\n\n\ndef get_config():\n \"\"\"Get the default hyperparameter configuration.\"\"\"\n config = ml_collections.ConfigDict()\n\n config.workdir = '/tmp/seq2seq'\n config.learning_rate = 0.003\n config.batch_size = 128\n config.hidden_size = 512\n config.num_train_steps = 10000\n config.decode_frequency = 200\n config.max_len_query_digit = 3\n\n return config\n" }, { "path": "examples/seq2seq/input_pipeline.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Input pipeline for seq2seq addition example.\"\"\"\n\nimport random\nfrom typing import Any, Dict, Optional, Tuple\nfrom collections.abc import Generator\n\nimport jax.numpy as jnp\nimport numpy as np\n\nArray = Any # pylint: disable=invalid-name\n\n\nclass CharacterTable:\n \"\"\"Encodes/decodes between strings and integer representations.\"\"\"\n\n def __init__(self, chars: str, max_len_query_digit: int = 3) -> None:\n self._chars = sorted(set(chars))\n self._char_indices = {ch: idx + 2 for idx, ch in enumerate(self._chars)}\n self._indices_char = {idx + 2: ch for idx, ch in enumerate(self._chars)}\n self._indices_char[self.pad_id] = '_'\n # Maximum length of a single input digit.\n self._max_len_query_digit = max_len_query_digit\n\n @property\n def pad_id(self) -> int:\n return 0\n\n @property\n def eos_id(self) -> int:\n return 1\n\n @property\n def vocab_size(self) -> int:\n # All characters + pad token and eos token.\n return len(self._chars) + 2\n\n @property\n def max_input_len(self) -> int:\n \"\"\"Returns the max length of an input sequence.\"\"\"\n # The input has the form \"digit1+digit2\", so the max input length is\n # the length of two digits plus two tokens for \"+\" and the EOS token.\n return self._max_len_query_digit * 2 + 2\n\n @property\n def max_output_len(self) -> int:\n \"\"\"Returns the max length of an output sequence.\"\"\"\n # The output has the form \"=digit\". If `digit` is the result of adding\n # two digits of max length x, then max length of `digit` is x+1.\n # Additionally, we require two more tokens for \"=\" and \" tuple[int, int, int]:\n return (1, self.max_input_len, self.vocab_size)\n\n @property\n def decoder_input_shape(self) -> tuple[int, int, int]:\n return (1, self.max_output_len, self.vocab_size)\n\n def encode(self, inputs: str) -> np.ndarray:\n \"\"\"Encodes from string to list of integers.\"\"\"\n return np.array(\n [self._char_indices[char] for char in inputs] + [self.eos_id]\n )\n\n def decode(self, inputs: Array) -> str:\n \"\"\"Decodes from list of integers to string.\"\"\"\n chars = []\n for elem in inputs.tolist():\n if elem == self.eos_id:\n break\n chars.append(self._indices_char[elem])\n return ''.join(chars)\n\n def one_hot(self, tokens: np.ndarray) -> np.ndarray:\n vecs = np.zeros((tokens.size, self.vocab_size), dtype=np.float32)\n vecs[np.arange(tokens.size), tokens] = 1\n return vecs\n\n def encode_onehot(\n self, batch_inputs: Array, max_len: int | None = None\n ) -> np.ndarray:\n \"\"\"One-hot encodes a string input.\"\"\"\n\n if max_len is None:\n max_len = self.max_input_len\n\n def encode_str(s):\n tokens = self.encode(s)\n unpadded_len = len(tokens)\n if unpadded_len > max_len:\n raise ValueError(f\"Sequence too long ({len(tokens)}>{max_len}): '{s}'\")\n tokens = np.pad(tokens, [(0, max_len - len(tokens))], mode='constant')\n return self.one_hot(tokens)\n\n return np.array([encode_str(inp) for inp in batch_inputs])\n\n def decode_onehot(self, batch_inputs: Array) -> np.ndarray:\n \"\"\"Decodes a batch of one-hot encoding to strings.\"\"\"\n decode_inputs = lambda inputs: self.decode(inputs.argmax(axis=-1))\n return np.array(list(map(decode_inputs, batch_inputs)))\n\n def generate_examples(\n self, num_examples: int\n ) -> Generator[tuple[str, str], None, None]:\n \"\"\"Yields `num_examples` examples.\"\"\"\n for _ in range(num_examples):\n max_digit = pow(10, self._max_len_query_digit) - 1\n # TODO(marcvanzee): Use jax.random here.\n key = tuple(sorted((random.randint(0, 99), random.randint(0, max_digit))))\n inputs = f'{key[0]}+{key[1]}'\n # Preprend output by the decoder's start token.\n outputs = '=' + str(key[0] + key[1])\n yield (inputs, outputs)\n\n def get_batch(self, batch_size: int) -> dict[str, np.ndarray]:\n \"\"\"Returns a batch of example of size @batch_size.\"\"\"\n inputs, outputs = zip(*self.generate_examples(batch_size))\n return {\n 'query': self.encode_onehot(inputs),\n 'answer': self.encode_onehot(outputs),\n }\n\n\ndef mask_sequences(sequence_batch: Array, lengths: Array) -> Array:\n \"\"\"Sets positions beyond the length of each sequence to 0.\"\"\"\n return sequence_batch * (\n lengths[:, np.newaxis] > np.arange(sequence_batch.shape[1])[np.newaxis]\n )\n\n\ndef get_sequence_lengths(sequence_batch: Array, eos_id: int) -> Array:\n \"\"\"Returns the length of each one-hot sequence, including the EOS token.\"\"\"\n # sequence_batch.shape = (batch_size, seq_length, vocab_size)\n eos_row = sequence_batch[:, :, eos_id]\n eos_idx = jnp.argmax(eos_row, axis=-1) # returns first occurrence\n # `eos_idx` is 0 if EOS is not present, so we use full length in that case.\n return jnp.where(\n eos_row[jnp.arange(eos_row.shape[0]), eos_idx],\n eos_idx + 1,\n sequence_batch.shape[1], # if there is no EOS, use full length\n )\n" }, { "path": "examples/seq2seq/main.py", "content": "# Copyright 2025 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Main script for seq2seq example.\"\"\"\n\nfrom absl import app\nfrom absl import flags\nfrom absl import logging\nimport train\nfrom ml_collections import config_flags\n\nFLAGS = flags.FLAGS\n\nconfig_flags.DEFINE_config_file(\n 'config',\n None,\n 'File path to the training hyperparameter configuration.',\n lock_config=True,\n)\n\n\ndef main(argv):\n del argv\n\n config = FLAGS.config\n\n # Set train.FLAGS values from config\n train.FLAGS.workdir = config.workdir\n train.FLAGS.learning_rate = config.learning_rate\n train.FLAGS.batch_size = config.batch_size\n train.FLAGS.hidden_size = config.hidden_size\n train.FLAGS.num_train_steps = config.num_train_steps\n train.FLAGS.decode_frequency = config.decode_frequency\n train.FLAGS.max_len_query_digit = config.max_len_query_digit\n\n logging.info('Starting training with config: %s', config)\n _ = train.train_and_evaluate(train.FLAGS.workdir)\n\n\nif __name__ == '__main__':\n app.run(main)\n" }, { "path": "examples/seq2seq/models.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"seq2seq example: Mode code.\"\"\"\n\n# See issue #620.\n# pytype: disable=wrong-keyword-args\n\nfrom typing import Tuple\n\nfrom flax import linen as nn\nimport jax\nimport jax.numpy as jnp\n\nArray = jax.Array\nPRNGKey = jax.Array\nLSTMCarry = tuple[Array, Array]\n\n\nclass DecoderLSTMCell(nn.RNNCellBase):\n \"\"\"DecoderLSTM Module wrapped in a lifted scan transform.\n\n Attributes:\n teacher_force: See docstring on Seq2seq module.\n vocab_size: Size of the vocabulary.\n \"\"\"\n\n features: int\n teacher_force: bool\n vocab_size: int\n\n @nn.compact\n def __call__(\n self, carry: tuple[LSTMCarry, Array], x: Array\n ) -> tuple[tuple[LSTMCarry, Array], tuple[Array, Array]]:\n \"\"\"Applies the DecoderLSTM model.\"\"\"\n lstm_state, last_prediction = carry\n if not self.teacher_force:\n x = last_prediction\n lstm_state, y = nn.LSTMCell(self.features)(lstm_state, x)\n logits = nn.Dense(features=self.vocab_size)(y)\n # Sample the predicted token using a categorical distribution over the\n # logits.\n categorical_rng = self.make_rng('lstm')\n predicted_token = jax.random.categorical(categorical_rng, logits)\n # Convert to one-hot encoding.\n prediction = jax.nn.one_hot(\n predicted_token, self.vocab_size, dtype=jnp.float32\n )\n\n return (lstm_state, prediction), (logits, prediction)\n\n @property\n def num_feature_axes(self) -> int:\n return 1\n\n\nclass Seq2seq(nn.Module):\n \"\"\"Sequence-to-sequence class using encoder/decoder architecture.\n\n Attributes:\n teacher_force: whether to use `decoder_inputs` as input to the decoder at\n every step. If False, only the first input (i.e., the \"=\" token) is used,\n followed by samples taken from the previous output logits.\n hidden_size: int, the number of hidden dimensions in the encoder and decoder\n LSTMs.\n vocab_size: the size of the vocabulary.\n eos_id: EOS id.\n \"\"\"\n\n teacher_force: bool\n hidden_size: int\n vocab_size: int\n eos_id: int = 1\n\n @nn.compact\n def __call__(\n self, encoder_inputs: Array, decoder_inputs: Array\n ) -> tuple[Array, Array]:\n \"\"\"Applies the seq2seq model.\n\n Args:\n encoder_inputs: [batch_size, max_input_length, vocab_size].\n padded batch of input sequences to encode.\n decoder_inputs: [batch_size, max_output_length, vocab_size].\n padded batch of expected decoded sequences for teacher forcing.\n When sampling (i.e., `teacher_force = False`), only the first token is\n input into the decoder (which is the token \"=\"), and samples are used\n for the following inputs. The second dimension of this tensor determines\n how many steps will be decoded, regardless of the value of\n `teacher_force`.\n\n Returns:\n Pair (logits, predictions), which are two arrays of length `batch_size`\n containing respectively decoded logits and predictions (in one hot\n encoding format).\n \"\"\"\n # Encode inputs.\n encoder = nn.RNN(\n nn.LSTMCell(self.hidden_size), return_carry=True, name='encoder'\n )\n decoder = nn.RNN(\n DecoderLSTMCell(\n decoder_inputs.shape[-1], self.teacher_force, self.vocab_size\n ),\n split_rngs={'params': False, 'lstm': True},\n name='decoder',\n )\n\n seq_lengths = self.get_seq_lengths(encoder_inputs)\n\n encoder_state, _ = encoder(encoder_inputs, seq_lengths=seq_lengths)\n logits, predictions = decoder(\n decoder_inputs[:, :-1],\n initial_carry=(encoder_state, decoder_inputs[:, 0]),\n )\n\n return logits, predictions\n\n def get_seq_lengths(self, inputs: Array) -> Array:\n \"\"\"Get segmentation mask for inputs.\"\"\"\n # undo one-hot encoding\n inputs = jnp.argmax(inputs, axis=-1)\n # calculate sequence lengths\n seq_lengths = jnp.argmax(inputs == self.eos_id, axis=-1)\n\n return seq_lengths\n" }, { "path": "examples/seq2seq/requirements.txt", "content": "absl-py==1.0.0\nclu==0.0.6\nflax==0.3.6\nnumpy==1.22.0\noptax==0.1.0\n" }, { "path": "examples/seq2seq/seq2seq.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"# Flax seq2seq Example\\n\",\n \"\\n\",\n \"\\\"Open\\n\",\n \"\\n\",\n \"Demonstration notebook for\\n\",\n \"https://github.com/google/flax/tree/main/examples/seq2seq\\n\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"The **Flax Notebook Workflow**:\\n\",\n \"\\n\",\n \"1. Run the entire notebook end-to-end and check out the outputs.\\n\",\n \" - This will open Python files in the right-hand editor!\\n\",\n \" - You'll be able to interactively explore metrics in TensorBoard.\\n\",\n \"2. Change some of the hyperparameters in the command-line flags in `train.py` for different hyperparameters. Check out the updated TensorBoard plots.\\n\",\n \"3. Update the code in `train.py`, `models.py`, and `input_pipeline.py`. \\n\",\n \" Thanks to `%autoreload`, any changes you make in the file will \\n\",\n \" automatically appear in the notebook. Some ideas to get you started:\\n\",\n \" - Change the model.\\n\",\n \" - Log some per-batch metrics during training.\\n\",\n \" - Add new hyperparameters to `models.py` and use them in `train.py`.\\n\",\n \" - Train on a different vocabulary by initializing `CharacterTable` with a\\n\",\n \" different character set.\\n\",\n \"4. At any time, feel free to paste code from the source code into the notebook\\n\",\n \" and modify it directly there!\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Setup\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {\n \"outputId\": \"4c0a705c-8d7e-44cc-d851-873a40ac115e\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"\\u001b[K |████████████████████████████████| 77 kB 3.1 MB/s \\n\",\n \"\\u001b[K |████████████████████████████████| 176 kB 30.2 MB/s \\n\",\n \"\\u001b[K |████████████████████████████████| 77 kB 5.2 MB/s \\n\",\n \"\\u001b[K |████████████████████████████████| 136 kB 45.5 MB/s \\n\",\n \"\\u001b[K |████████████████████████████████| 65 kB 2.8 MB/s \\n\",\n \"\\u001b[K |████████████████████████████████| 462 kB 44.3 MB/s \\n\",\n \"\\u001b[?25h Building wheel for ml-collections (setup.py) ... \\u001b[?25l\\u001b[?25hdone\\n\"\n ]\n }\n ],\n \"source\": [\n \"# Install CLU & Flax.\\n\",\n \"!pip install -q clu flax\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {\n \"tags\": []\n },\n \"outputs\": [],\n \"source\": [\n \"example_directory = 'examples/seq2seq'\\n\",\n \"editor_relpaths = ('train.py', 'input_pipeline.py', 'models.py')\\n\",\n \"\\n\",\n \"repo, branch = 'https://github.com/google/flax', 'main'\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {\n \"cellView\": \"form\",\n \"outputId\": \"4801432e-4090-4b13-f0f2-d99a3039ce47\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"Cloning into 'flaxrepo'...\\n\",\n \"remote: Enumerating objects: 349, done.\\u001b[K\\n\",\n \"remote: Counting objects: 100% (349/349), done.\\u001b[K\\n\",\n \"remote: Compressing objects: 100% (286/286), done.\\u001b[K\\n\",\n \"remote: Total 349 (delta 63), reused 220 (delta 51), pack-reused 0\\u001b[K\\n\",\n \"Receiving objects: 100% (349/349), 2.12 MiB | 13.39 MiB/s, done.\\n\",\n \"Resolving deltas: 100% (63/63), done.\\n\"\n ]\n },\n {\n \"data\": {\n \"text/html\": [\n \"

WARNING : Editing in VM - changes lost after reboot!!

\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"application/javascript\": [\n \"\\n\",\n \" ((filepath) => {{\\n\",\n \" if (!google.colab.kernel.accessAllowed) {{\\n\",\n \" return;\\n\",\n \" }}\\n\",\n \" google.colab.files.view(filepath);\\n\",\n \" }})(\\\"/content/examples/seq2seq/train.py\\\")\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"application/javascript\": [\n \"\\n\",\n \" ((filepath) => {{\\n\",\n \" if (!google.colab.kernel.accessAllowed) {{\\n\",\n \" return;\\n\",\n \" }}\\n\",\n \" google.colab.files.view(filepath);\\n\",\n \" }})(\\\"/content/examples/seq2seq/input_pipeline.py\\\")\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n },\n {\n \"data\": {\n \"application/javascript\": [\n \"\\n\",\n \" ((filepath) => {{\\n\",\n \" if (!google.colab.kernel.accessAllowed) {{\\n\",\n \" return;\\n\",\n \" }}\\n\",\n \" google.colab.files.view(filepath);\\n\",\n \" }})(\\\"/content/examples/seq2seq/models.py\\\")\"\n ],\n \"text/plain\": [\n \"\"\n ]\n },\n \"metadata\": {},\n \"output_type\": \"display_data\"\n }\n ],\n \"source\": [\n \"# (If you run this code in Jupyter[lab], then you're already in the\\n\",\n \"# example directory and nothing needs to be done.)\\n\",\n \"\\n\",\n \"#@markdown **Fetch newest Flax, copy example code**\\n\",\n \"#@markdown\\n\",\n \"#@markdown **If you select no** below, then the files will be stored on the\\n\",\n \"#@markdown *ephemeral* Colab VM. **After some time of inactivity, this VM will\\n\",\n \"#@markdown be restarted and any changes are lost**.\\n\",\n \"#@markdown\\n\",\n \"#@markdown **If you select yes** below, then you will be asked for your\\n\",\n \"#@markdown credentials to mount your personal Google Drive. In this case, all\\n\",\n \"#@markdown changes you make will be *persisted*, and even if you re-run the\\n\",\n \"#@markdown Colab later on, the files will still be the same (you can of course\\n\",\n \"#@markdown remove directories inside your Drive's `flax/` root if you want to\\n\",\n \"#@markdown manually revert these files).\\n\",\n \"\\n\",\n \"if 'google.colab' in str(get_ipython()):\\n\",\n \" import os\\n\",\n \" os.chdir('/content')\\n\",\n \" # Download Flax repo from Github.\\n\",\n \" if not os.path.isdir('flaxrepo'):\\n\",\n \" !git clone --depth=1 -b $branch $repo flaxrepo\\n\",\n \" # Copy example files & change directory.\\n\",\n \" mount_gdrive = 'no' #@param ['yes', 'no']\\n\",\n \" if mount_gdrive == 'yes':\\n\",\n \" DISCLAIMER = 'Note : Editing in your Google Drive, changes will persist.'\\n\",\n \" from google.colab import drive\\n\",\n \" drive.mount('/content/gdrive')\\n\",\n \" example_root_path = f'/content/gdrive/My Drive/flax/{example_directory}'\\n\",\n \" else:\\n\",\n \" DISCLAIMER = 'WARNING : Editing in VM - changes lost after reboot!!'\\n\",\n \" example_root_path = f'/content/{example_directory}'\\n\",\n \" from IPython import display\\n\",\n \" display.display(display.HTML(\\n\",\n \" f'

{DISCLAIMER}

'))\\n\",\n \" if not os.path.isdir(example_root_path):\\n\",\n \" os.makedirs(example_root_path)\\n\",\n \" !cp -r flaxrepo/$example_directory/* \\\"$example_root_path\\\"\\n\",\n \" os.chdir(example_root_path)\\n\",\n \" from google.colab import files\\n\",\n \" for relpath in editor_relpaths:\\n\",\n \" s = open(f'{example_root_path}/{relpath}').read()\\n\",\n \" open(f'{example_root_path}/{relpath}', 'w').write(\\n\",\n \" f'## {DISCLAIMER}\\\\n' + '#' * (len(DISCLAIMER) + 3) + '\\\\n\\\\n' + s)\\n\",\n \" files.view(f'{example_root_path}/{relpath}')\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {\n \"outputId\": \"a292a7a2-ae3c-4518-af28-9c2fa0ed2d7b\"\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"/content/examples/seq2seq\\n\"\n ]\n }\n ],\n \"source\": [\n \"# Note : In Colab, above cell changed the working directory.\\n\",\n \"!pwd\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Imports\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"from absl import app\\n\",\n \"app.parse_flags_with_usage(['seq2seq'])\\n\",\n \"\\n\",\n \"from absl import logging\\n\",\n \"logging.set_verbosity(logging.INFO)\\n\",\n \"\\n\",\n \"import jax\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {\n \"outputId\": \"7e1a29ce-9d8b-4715-ce60-9eae100a1df3\",\n \"tags\": []\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"The autoreload extension is already loaded. To reload it, use:\\n\",\n \" %reload_ext autoreload\\n\"\n ]\n }\n ],\n \"source\": [\n \"# Local imports from current directory - auto reload.\\n\",\n \"# Any changes you make to the three imported files will appear automatically.\\n\",\n \"%load_ext autoreload\\n\",\n \"%autoreload 2\\n\",\n \"import input_pipeline\\n\",\n \"import models\\n\",\n \"import train\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Dataset\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {\n \"outputId\": \"cb5f7f6e-1e6f-40ff-e0d6-5b428511d75b\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"[('72+789', '=861'),\\n\",\n \" ('58+858', '=916'),\\n\",\n \" ('77+358', '=435'),\\n\",\n \" ('99+264', '=363'),\\n\",\n \" ('94+115', '=209')]\"\n ]\n },\n \"execution_count\": 11,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"# Examples are generated on the fly.\\n\",\n \"ctable = input_pipeline.CharacterTable('0123456789+= ')\\n\",\n \"list(ctable.generate_examples(5))\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {\n \"outputId\": \"b58ea813-e757-4cc5-f3ba-3cb0f05d35a6\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"array([[0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\\n\",\n \" [0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\\n\",\n \" [0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0.],\\n\",\n \" [0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0.],\\n\",\n \" [0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0.],\\n\",\n \" [0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\\n\",\n \" [1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\\n\",\n \" [1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]],\\n\",\n \" dtype=float32)\"\n ]\n },\n \"execution_count\": 13,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"batch = ctable.get_batch(5)\\n\",\n \"# A single query (/answer) is one-hot encoded.\\n\",\n \"batch['query'][0]\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {\n \"outputId\": \"3b33e061-f0b5-42d7-ad49-5058e8fd3b90\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"array(['1+243'], dtype=' and Privacy Policy\\n\",\n \", and TensorBoard.dev's Terms of Service\\n\",\n \".\\n\",\n \"\\n\",\n \"This notice will not be shown again while you are logged into the uploader.\\n\",\n \"To log out, run `tensorboard dev auth revoke`.\\n\",\n \"\\n\",\n \"Continue? (yes/NO) yes\\n\",\n \"\\n\",\n \"Please visit this URL to authorize this application: https://accounts.google.com/o/oauth2/auth?response_type=code&client_id=373649185512-8v619h5kft38l4456nm2dj4ubeqsrvh6.apps.googleusercontent.com&redirect_uri=urn%3Aietf%3Awg%3Aoauth%3A2.0%3Aoob&scope=openid+https%3A%2F%2Fwww.googleapis.com%2Fauth%2Fuserinfo.email&state=IjociK9llsm6dSiC1TDFvFmJksFy49&prompt=consent&access_type=offline\\n\",\n \"Enter the authorization code: 4/1AX4XfWi6J9MqoDpbZ5Z_jd1AVheW7277VuUoTOEz5_8NRs_3oP9M7S4T81c\\n\",\n \"\\n\",\n \"\\n\",\n \"New experiment created. View your TensorBoard at: https://tensorboard.dev/experiment/pgfmdFlaQTy9odov72ZvVQ/\\n\",\n \"\\n\",\n \"\\u001b[1m[2022-02-25T09:34:21]\\u001b[0m Started scanning logdir.\\n\",\n \"\\u001b[1m[2022-02-25T09:34:22]\\u001b[0m Total uploaded: 38 scalars, 0 tensors, 0 binary objects\\n\",\n \"\\u001b[1m[2022-02-25T09:34:22]\\u001b[0m Done scanning logdir.\\n\",\n \"\\n\",\n \"\\n\",\n \"Done. View your TensorBoard at https://tensorboard.dev/experiment/pgfmdFlaQTy9odov72ZvVQ/\\n\"\n ]\n }\n ],\n \"source\": [\n \"if 'google.colab' in str(get_ipython()):\\n\",\n \" #@markdown You can upload the training results directly to https://tensorboard.dev\\n\",\n \" #@markdown\\n\",\n \" #@markdown Note that everybody with the link will be able to see the data.\\n\",\n \" upload_data = 'yes' #@param ['yes', 'no']\\n\",\n \" if upload_data == 'yes':\\n\",\n \" !tensorboard dev upload --one_shot --logdir ./workdirs --name 'Flax examples/seq2seq (Colab)'\"\n ]\n },\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"## Inference\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {\n \"outputId\": \"e22b7208-5413-4a63-abfb-b510af60f340\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"(1, 8, 15)\"\n ]\n },\n \"execution_count\": 23,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"inputs = ctable.encode_onehot(['2+40'])\\n\",\n \"# batch, max_length, vocab_size\\n\",\n \"inputs.shape\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# Using different random seeds generates different samples.\\n\",\n \"preds = train.decode(state.params, inputs, jax.random.key(0), ctable)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {\n \"outputId\": \"e5cdfd75-2c66-4165-8ab7-9fdecde5062a\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"array(['42'], dtype=' models.Seq2seq:\n return models.Seq2seq(\n teacher_force=teacher_force,\n hidden_size=FLAGS.hidden_size,\n eos_id=ctable.eos_id,\n vocab_size=ctable.vocab_size,\n )\n\n\ndef get_initial_params(\n model: models.Seq2seq, rng: PRNGKey, ctable: CTable\n) -> dict[str, Any]:\n \"\"\"Returns the initial parameters of a seq2seq model.\"\"\"\n rng1, rng2 = jax.random.split(rng)\n variables = model.init(\n {'params': rng1, 'lstm': rng2},\n jnp.ones(ctable.encoder_input_shape, jnp.float32),\n jnp.ones(ctable.decoder_input_shape, jnp.float32),\n )\n return variables['params']\n\n\ndef get_train_state(rng: PRNGKey, ctable: CTable) -> train_state.TrainState:\n \"\"\"Returns a train state.\"\"\"\n model = get_model(ctable)\n params = get_initial_params(model, rng, ctable)\n tx = optax.adam(FLAGS.learning_rate)\n state = train_state.TrainState.create(\n apply_fn=model.apply, params=params, tx=tx\n )\n return state\n\n\ndef cross_entropy_loss(\n logits: Array, labels: Array, lengths: Array\n) -> jax.Array:\n \"\"\"Returns cross-entropy loss.\"\"\"\n xe = jnp.sum(nn.log_softmax(logits) * labels, axis=-1)\n masked_xe = jnp.mean(mask_sequences(xe, lengths))\n return -masked_xe\n\n\ndef compute_metrics(\n logits: Array, labels: Array, eos_id: int\n) -> dict[str, jax.Array]:\n \"\"\"Computes metrics and returns them.\"\"\"\n lengths = get_sequence_lengths(labels, eos_id)\n loss = cross_entropy_loss(logits, labels, lengths)\n # Computes sequence accuracy, which is the same as the accuracy during\n # inference, since teacher forcing is irrelevant when all output are correct.\n token_accuracy = jnp.argmax(logits, -1) == jnp.argmax(labels, -1)\n sequence_accuracy = (\n jnp.sum(mask_sequences(token_accuracy, lengths), axis=-1) == lengths\n )\n accuracy = jnp.mean(sequence_accuracy)\n metrics = {\n 'loss': loss,\n 'accuracy': accuracy,\n }\n return metrics\n\n\n@jax.jit\ndef train_step(\n state: train_state.TrainState, batch: Array, lstm_rng: PRNGKey, eos_id: int\n) -> tuple[train_state.TrainState, dict[str, jax.Array]]:\n \"\"\"Trains one step.\"\"\"\n labels = batch['answer'][:, 1:]\n lstm_key = jax.random.fold_in(lstm_rng, state.step)\n\n def loss_fn(params):\n logits, _ = state.apply_fn(\n {'params': params},\n batch['query'],\n batch['answer'],\n rngs={'lstm': lstm_key},\n )\n loss = cross_entropy_loss(\n logits, labels, get_sequence_lengths(labels, eos_id)\n )\n return loss, logits\n\n grad_fn = jax.value_and_grad(loss_fn, has_aux=True)\n (_, logits), grads = grad_fn(state.params)\n state = state.apply_gradients(grads=grads)\n metrics = compute_metrics(logits, labels, eos_id)\n\n return state, metrics\n\n\ndef log_decode(question: str, inferred: str, golden: str):\n \"\"\"Logs the given question, inferred query, and correct query.\"\"\"\n suffix = (\n '(CORRECT)' if inferred == golden else (f'(INCORRECT) correct={golden}')\n )\n logging.info('DECODE: %s = %s %s', question, inferred, suffix)\n\n\n@functools.partial(jax.jit, static_argnums=3)\ndef decode(\n params: dict[str, Any], inputs: Array, decode_rng: PRNGKey, ctable: CTable\n) -> Array:\n \"\"\"Decodes inputs.\"\"\"\n init_decoder_input = ctable.one_hot(ctable.encode('=')[0:1])\n init_decoder_inputs = jnp.tile(\n init_decoder_input, (inputs.shape[0], ctable.max_output_len, 1)\n )\n model = get_model(ctable, teacher_force=False)\n _, predictions = model.apply(\n {'params': params}, inputs, init_decoder_inputs, rngs={'lstm': decode_rng}\n )\n return predictions\n\n\ndef decode_batch(\n state: train_state.TrainState,\n batch: dict[str, Array],\n decode_rng: PRNGKey,\n ctable: CTable,\n):\n \"\"\"Decodes and log results for a batch.\"\"\"\n inputs, outputs = batch['query'], batch['answer'][:, 1:]\n decode_rng = jax.random.fold_in(decode_rng, state.step)\n inferred = decode(state.params, inputs, decode_rng, ctable)\n questions = ctable.decode_onehot(inputs)\n infers = ctable.decode_onehot(inferred)\n goldens = ctable.decode_onehot(outputs)\n\n for question, inferred, golden in zip(questions, infers, goldens):\n log_decode(question, inferred, golden)\n\n\ndef train_and_evaluate(workdir: str) -> train_state.TrainState:\n \"\"\"Trains for a fixed number of steps and decode during training.\"\"\"\n\n # TODO(marcvanzee): Integrate ctable with train_state.\n ctable = CTable('0123456789+= ', FLAGS.max_len_query_digit)\n rng = jax.random.key(0)\n state = get_train_state(rng, ctable)\n\n writer = metric_writers.create_default_writer(workdir)\n for step in range(FLAGS.num_train_steps):\n batch = ctable.get_batch(FLAGS.batch_size)\n state, metrics = train_step(state, batch, rng, ctable.eos_id)\n if step and step % FLAGS.decode_frequency == 0:\n writer.write_scalars(step, metrics)\n batch = ctable.get_batch(5)\n decode_batch(state, batch, rng, ctable)\n\n return state\n\n\ndef main(_):\n _ = train_and_evaluate(FLAGS.workdir)\n\n\nif __name__ == '__main__':\n app.run(main)\n" }, { "path": "examples/seq2seq/train_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for flax.examples.seq2seq.train.\"\"\"\n\nimport functools\n\nfrom absl.testing import absltest\nfrom flax.training import train_state\nimport jax\nfrom jax import random\nimport numpy as np\nimport optax\n\nimport input_pipeline\nimport train\nimport models\n\njax.config.parse_flags_with_absl()\n\n\ndef create_ctable(chars='0123456789+= '):\n return input_pipeline.CharacterTable(chars)\n\n\ndef create_train_state(ctable):\n model = models.Seq2seq(\n teacher_force=False,\n hidden_size=train.FLAGS.hidden_size,\n vocab_size=ctable.vocab_size,\n )\n params = train.get_initial_params(model, jax.random.key(0), ctable)\n tx = optax.adam(train.FLAGS.learning_rate)\n state = train_state.TrainState.create(\n apply_fn=model.apply, params=params, tx=tx\n )\n return state\n\n\nclass TrainTest(absltest.TestCase):\n\n def test_character_table(self):\n ctable = create_ctable()\n text = '410+19'\n enc_text = ctable.encode(text)\n dec_text = ctable.decode(enc_text)\n # The text is possibly padded with whitespace, but the trimmed output should\n # be equal to the input.\n self.assertEqual(text, dec_text.strip())\n\n def test_mask_sequences(self):\n np.testing.assert_equal(\n input_pipeline.mask_sequences(\n np.arange(1, 13).reshape((4, 3)), np.array([3, 2, 1, 0])\n ),\n np.array([[1, 2, 3], [4, 5, 0], [7, 0, 0], [0, 0, 0]]),\n )\n\n def test_get_sequence_lengths(self):\n oh_sequence_batch = jax.vmap(\n functools.partial(jax.nn.one_hot, num_classes=4)\n )(np.array([[0, 1, 0], [1, 0, 2], [1, 2, 0], [1, 2, 3]]))\n np.testing.assert_equal(\n input_pipeline.get_sequence_lengths(oh_sequence_batch, eos_id=0),\n np.array([1, 2, 3, 3], np.int32),\n )\n np.testing.assert_equal(\n input_pipeline.get_sequence_lengths(oh_sequence_batch, eos_id=1),\n np.array([2, 1, 1, 1], np.int32),\n )\n np.testing.assert_equal(\n input_pipeline.get_sequence_lengths(oh_sequence_batch, eos_id=2),\n np.array([3, 3, 2, 2], np.int32),\n )\n\n def test_train_one_step(self):\n ctable = create_ctable()\n batch = ctable.get_batch(128)\n\n state = create_train_state(ctable)\n key = random.key(0)\n _, train_metrics = train.train_step(state, batch, key, ctable.eos_id)\n\n self.assertLessEqual(train_metrics['loss'], 5)\n self.assertGreaterEqual(train_metrics['accuracy'], 0)\n\n def test_decode_batch(self):\n ctable = create_ctable()\n batch = ctable.get_batch(5)\n key = random.key(0)\n state = create_train_state(ctable)\n train.decode_batch(state, batch, key, ctable)\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "examples/sst2/README.md", "content": "## SST-2 classification\n\nTrains a simple text classifier on the SST-2 sentiment classification dataset.\n\nYou can run this code and even modify it directly in Google Colab, no\ninstallation required:\n\nhttps://colab.research.google.com/github/google/flax/blob/main/examples/sst2/sst2.ipynb\n\n### Requirements\n* TensorFlow dataset `glue/sst2` will be downloaded and prepared automatically, if necessary.\n\n### Example output\n\n| Name | Platform | Epochs | Walltime | Accuracy | Metrics | Workdir |\n|:--------|:--------|--------:|:-----------|:-----------------|:----------------------------------------------------------------------------------------------------------------------|:-------------------------------------------------------------------------------------------------------------------------------|\n| default | TPU | 10 | 4.3m | 85.21% | [tensorboard.dev](https://tensorboard.dev/experiment/yTQjjRY9RlGRrZzg8h9PJw/) | |\n\n```\nINFO:absl:train epoch 010 loss 0.1918 accuracy 92.41\nINFO:absl:eval epoch 010 loss 0.4144 accuracy 85.21\n```\n\n### How to run\n\n```bash\npython main.py --workdir=/tmp/sst2 --config=configs/default.py`\n```\n\n#### Overriding Hyperparameter configurations\n\nThe SST2 example allows specifying a hyperparameter configuration by means of\nsetting the `--config` flag. The configuration flag is defined using\n[config_flags](https://github.com/google/ml_collections/tree/master#config-flags).\n`config_flags` allows overriding configuration fields. This can be done as\nfollows:\n\n```shell\npython main.py \\\n --workdir=/tmp/sst2 --config=configs/default.py \\\n --config.learning_rate=0.05 --config.num_epochs=5\n```\n" }, { "path": "examples/sst2/build_vocabulary.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"A vocabulary builder that generates vocab.txt to be used for training.\"\"\"\n\nimport time\nfrom collections.abc import Iterable, Sequence\n\nfrom absl import logging\nimport tensorflow as tf\nimport tensorflow_datasets as tfds\nimport tensorflow_text as tftext\n\nimport vocabulary\n\n\ndef get_tokenized_sequences(\n dataset: tf.data.Dataset,\n tokenizer: tftext.Tokenizer = tftext.WhitespaceTokenizer(),\n input_key: str = 'sentence',\n) -> Iterable[Sequence[bytes]]:\n \"\"\"Returns tokenized sequences for vocabulary building.\"\"\"\n dataset = dataset.map(\n lambda example: tokenizer.tokenize(example[input_key]),\n num_parallel_calls=tf.data.experimental.AUTOTUNE,\n )\n yield from tfds.as_numpy(dataset)\n\n\nif __name__ == '__main__':\n logging.set_verbosity(logging.INFO)\n start_time = time.time()\n\n # Loads the dataset to build the vocabulary from. We use the train split.\n dataset = tfds.load('glue/sst2', split='train')\n\n # Tokenizes the sequences in the dataset and keeps only those.\n tokenized_sequences = get_tokenized_sequences(dataset)\n\n # Builds the vocabulary from the tokenized sequences.\n # A token needs to appear at least 3 times to be in the vocabulary. You can\n # play with this. It is there to make sure we don't overfit on rare words.\n vocab = vocabulary.Vocabulary(\n tokenized_sequences=tokenized_sequences, min_freq=3\n )\n vocab.save('vocab.txt')\n\n logging.info('Total time elapsed: %f s', time.time() - start_time)\n" }, { "path": "examples/sst2/configs/default.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Default hyperparameter configuration for SST-2.\"\"\"\n\nimport ml_collections\n\n\ndef get_config():\n \"\"\"Get the default hyperparameter configuration.\"\"\"\n config = ml_collections.ConfigDict()\n\n config.embedding_size = 300\n config.hidden_size = 256\n config.vocab_size = None\n config.output_size = 1\n\n config.vocab_path = 'vocab.txt'\n config.max_input_length = 60\n\n config.dropout_rate = 0.5\n config.word_dropout_rate = 0.1\n config.unk_idx = 1\n\n config.learning_rate = 0.1\n config.momentum = 0.9\n config.weight_decay = 3e-6\n\n config.batch_size = 64\n config.bucket_size = 8\n config.num_epochs = 10\n\n config.seed = 0\n\n return config\n" }, { "path": "examples/sst2/input_pipeline.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"SST-2 input pipeline.\"\"\"\n\nimport sys\nfrom typing import Any, Dict, Optional\n\nfrom absl import logging\nimport numpy as np\nimport tensorflow as tf\nimport tensorflow_datasets as tfds\nif sys.version_info < (3, 13):\n import tensorflow_text as text\n\nimport vocabulary\n\n\nAUTOTUNE = tf.data.experimental.AUTOTUNE\nExample = dict[str, tf.Tensor]\n\n\ndef get_bucket_boundaries(bucket_size: int, max_size: int) -> np.ndarray:\n \"\"\"Bucket boundaries with `bucket_size` items per bucket, up to `max_size`.\n\n Example:\n ```\n get_bucket_boundaries(8, 24)\n [9, 17, 25]\n ```\n E.g., the first boundary covers items with sizes 0-8, the next boundary covers\n items with sizes 9-16, and the last bucket covers sizes 17-24. Each bucket\n covers 8 different sizes (e.g., sentence lengths).\n\n Args:\n bucket_size: The number of different items per bucket.\n max_size: The maximum size to be expected for a bucket.\n\n Returns:\n A list of (exclusive) bucket boundaries.\n \"\"\"\n return np.arange(bucket_size, max_size + bucket_size, bucket_size) + 1\n\n\ndef get_num_examples(dataset: tf.data.Dataset) -> int:\n \"\"\"Returns the number of examples in the dataset by iterating over it once.\"\"\"\n return dataset.reduce(np.int64(0), lambda x, _: x + 1).numpy()\n\n\ndef get_bucketed_batches(\n dataset: tf.data.Dataset,\n batch_size: int,\n bucket_size: int,\n max_length: int,\n padded_shapes: Any,\n example_size_fn: Any,\n shuffle: bool = False,\n shuffle_seed: int | None = None,\n drop_remainder: bool = False,\n) -> tf.data.Dataset:\n \"\"\"Returns padded batches of shuffled examples bucketed by length.\n\n This shuffles the examples randomly each epoch. The random order is\n deterministic and controlled by the seed.\n\n Batches are padded because sentences have different lengths.\n Sentences that are shorter in a batch will get 0s added at the end, until\n all sentences in the batch have the same length.\n\n For performance, examples of similar lengths are bucketed together. However,\n the contents of the buckets and their order is random each epoch, and\n controlled by the seed.\n\n Args:\n dataset: A TF Dataset with SST examples to be shuffled and batched.\n batch_size: The size of each batch. The remainder is dropped.\n bucket_size: How many different lengths go in each bucket.\n max_length: The maximum length to provide a bucket for.\n padded_shapes: A nested structure representing the shape to which the\n respective component of each input element should be padded prior to\n batching. See `tf.data.Dataset.padded_batch` for examples.\n example_size_fn: A TF function that returns the size of an example to\n determine in which bucket it goes. E.g., the sentence length.\n shuffle: Shuffle the dataset each epoch using seed.\n shuffle_seed: The seed that determines the shuffling order, with a\n different order each epoch.\n drop_remainder: Drop the last batch if it is not of size batch_size.\n\n Returns:\n A TF Dataset containing padded bucketed batches.\n \"\"\"\n if shuffle:\n assert shuffle_seed is not None, 'When shuffling you must provide a seed.'\n\n # For bucket_size 8 and max length 24, we get bucket boundaries [9, 17, 25].\n bucket_boundaries = get_bucket_boundaries(bucket_size, max_length)\n logging.info('Batching bucket boundaries: %r', bucket_boundaries)\n\n # One batch size for each bucket plus one additional one (per requirement).\n bucket_batch_sizes = [batch_size] * (len(bucket_boundaries) + 1)\n bucket_fn = tf.data.experimental.bucket_by_sequence_length(\n example_size_fn,\n bucket_boundaries,\n bucket_batch_sizes,\n padded_shapes=padded_shapes,\n pad_to_bucket_boundary=True,\n drop_remainder=drop_remainder,\n )\n\n if shuffle:\n # For shuffling we need to know how many training examples we have.\n num_examples = get_num_examples(dataset)\n num_batches = num_examples // batch_size\n return (\n dataset.shuffle(\n num_examples, seed=shuffle_seed, reshuffle_each_iteration=True\n )\n .apply(bucket_fn)\n .shuffle(num_batches, seed=shuffle_seed, reshuffle_each_iteration=True)\n .prefetch(tf.data.experimental.AUTOTUNE)\n )\n return dataset.apply(bucket_fn).prefetch(tf.data.experimental.AUTOTUNE)\n\n\ndef vocab_to_hashtable(\n vocab: vocabulary.Vocabulary, unk_idx: int\n) -> tf.lookup.StaticHashTable:\n \"\"\"Returns a TF lookup table (token -> ID) from a vocabulary.\"\"\"\n return tf.lookup.StaticHashTable(\n tf.lookup.KeyValueTensorInitializer(\n list(vocab.keys()), list(vocab.values())\n ),\n default_value=unk_idx,\n )\n\n\ndef vocab_to_inverse_hashtable(\n vocab: vocabulary.Vocabulary, unk_token: bytes\n) -> tf.lookup.StaticHashTable:\n \"\"\"Returns an inverse TF lookup table (ID -> token) from a vocabulary.\"\"\"\n return tf.lookup.StaticHashTable(\n tf.lookup.KeyValueTensorInitializer(\n list(vocab.values()),\n list(vocab.keys()),\n key_dtype=tf.int64,\n value_dtype=tf.string,\n ),\n default_value=unk_token,\n )\n\n\ndef _is_text_field(feature_name_and_type):\n \"\"\"Identifies a text field when given a feature (name, type) pair.\"\"\"\n _, feature_type = feature_name_and_type\n return isinstance(feature_type, tfds.features.Text)\n\n\ndef _is_class_label(feature_name_and_type):\n \"\"\"Identifies a class label field when given a feature (name, type) pair.\"\"\"\n _, feature_type = feature_name_and_type\n return isinstance(feature_type, tfds.features.ClassLabel)\n\n\nclass TextDataset:\n \"\"\"A text dataset with one sequence as input and a label.\"\"\"\n\n def __init__(\n self,\n tfds_name: str = 'glue/sst2',\n vocab_path: str = 'vocab.txt',\n tokenizer=None,\n split='train',\n ):\n \"\"\"Initializes the SST2 data source.\"\"\"\n self.dataset, self.info = tfds.load(tfds_name, split=split, with_info=True)\n\n # Look up the feature name of the text and label in the dataset.\n # We assume there is one text input and one label.\n text_fields = filter(_is_text_field, self.info.features.items())\n label_fields = filter(_is_class_label, self.info.features.items())\n self.text_feature_name, _ = next(text_fields)\n self.label_feature_name, _ = next(label_fields)\n\n # Load the vocabulary.\n self.vocab = vocabulary.Vocabulary(vocab_path=vocab_path)\n\n # Convert the sentences to sequences of token IDs and compute length.\n if tokenizer is None:\n tokenizer = text.WhitespaceTokenizer()\n self.tokenizer = tokenizer\n self.tf_vocab = vocab_to_hashtable(self.vocab, unk_idx=self.vocab.unk_idx)\n self.examples = self.dataset.map(\n self.prepare_example, num_parallel_calls=AUTOTUNE\n ).cache()\n\n @property\n def padded_shapes(self):\n \"\"\"The padded shapes used by batching functions.\"\"\"\n # None means variable length; pads to the longest sequence in the batch.\n return {'idx': [], 'token_ids': [None], 'label': [], 'length': []}\n\n def example_length_fn(self, example: Example) -> tf.Tensor:\n \"\"\"Returns the length of the example for the purpose of the bucketing.\"\"\"\n return tf.size(example['token_ids'])\n\n def add_bos_eos(self, sequence: tf.Tensor) -> tf.Tensor:\n \"\"\"Prepends BOS ID and appends EOS ID to a sequence of token IDs.\"\"\"\n return tf.concat([[self.vocab.bos_idx], sequence, [self.vocab.eos_idx]], 0)\n\n def prepare_example(self, example: Example) -> Example:\n \"\"\"Prepares an example by converting text to token IDs.\"\"\"\n tokens = self.tokenizer.tokenize(example[self.text_feature_name])\n label = example[self.label_feature_name]\n del example[self.text_feature_name]\n del example[self.label_feature_name]\n example['token_ids'] = self.add_bos_eos(self.tf_vocab.lookup(tokens))\n example['length'] = tf.size(example['token_ids'])\n example['label'] = label\n return example\n\n def get_batches(\n self,\n batch_size: int,\n drop_remainder: bool = False,\n shuffle: bool = False,\n shuffle_seed: int | None = None,\n fixed_pad_length: int | None = None,\n dataset: tf.data.Dataset | None = None,\n ):\n \"\"\"Returns an iterator with padded batches for the provided dataset.\"\"\"\n if dataset is None:\n dataset = self.examples\n if shuffle:\n buffer_size = get_num_examples(dataset)\n dataset = dataset.shuffle(\n buffer_size, seed=shuffle_seed, reshuffle_each_iteration=True\n )\n padded_shapes = {k: v for k, v in self.padded_shapes.items()}\n if fixed_pad_length is not None:\n padded_shapes['token_ids'] = fixed_pad_length\n return dataset.padded_batch(\n batch_size, padded_shapes=padded_shapes, drop_remainder=drop_remainder\n )\n\n def get_bucketed_batches(\n self,\n batch_size: int,\n bucket_size: int,\n max_input_length: int,\n drop_remainder: bool = False,\n shuffle: bool = False,\n shuffle_seed: int | None = None,\n dataset: tf.data.Dataset | None = None,\n ):\n \"\"\"Returns an iterator with bucketed batches for the provided dataset.\"\"\"\n if dataset is None:\n dataset = self.examples\n return get_bucketed_batches(\n dataset,\n batch_size,\n bucket_size,\n max_input_length,\n self.padded_shapes,\n self.example_length_fn,\n shuffle=shuffle,\n shuffle_seed=shuffle_seed,\n drop_remainder=drop_remainder,\n )\n" }, { "path": "examples/sst2/input_pipeline_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nimport pathlib\nimport sys\nimport tempfile\n\nfrom absl.testing import absltest\nimport tensorflow_datasets as tfds\n\nimport input_pipeline\nimport vocabulary\n\n\nclass InputPipelineTest(absltest.TestCase):\n\n def setUp(self):\n super().setUp()\n if sys.version_info >= (3, 13):\n self.skipTest('Test (and tensorflow-text) does not suport Python 3.13+')\n self.vocab_path = self._get_vocab_path()\n self.dataset = self._get_dataset(self.vocab_path)\n\n def _get_vocab_path(self):\n \"\"\"Prepares a mock vocabulary and returns a path to it.\"\"\"\n vocab_path = os.path.join(tempfile.mkdtemp(), 'vocab.txt')\n tokenized_sequences = (\n (b'this', b'is', b'a', b'test', b'sentence'),\n (b'this', b'is', b'a', b'test', b'sentence'),\n (b'this', b'is', b'a', b'test', b'sentence'),\n )\n vocab = vocabulary.Vocabulary(tokenized_sequences=tokenized_sequences)\n vocab.save(vocab_path)\n return vocab_path\n\n def _get_dataset(self, vocab_path):\n \"\"\"Uses mock data to create the dataset.\"\"\"\n # Go two directories up to the root of the flax directory.\n flax_root_dir = pathlib.Path(__file__).parents[2]\n data_dir = str(flax_root_dir) + '/.tfds/metadata' # pylint: disable=unused-variable\n with tfds.testing.mock_data(num_examples=128, data_dir=data_dir):\n return input_pipeline.TextDataset(vocab_path=vocab_path, split='train')\n\n def test_bucketed_dataset(self):\n \"\"\"Each batch should have a length that is a multiple of bucket_size.\"\"\"\n batch_size = 2\n bucket_size = 8\n for batch in self.dataset.get_bucketed_batches(\n batch_size=batch_size,\n bucket_size=bucket_size,\n max_input_length=60,\n shuffle=False,\n ).take(3):\n # Because of bucketing, sequence length must be multiple of bucket_size.\n length = batch['token_ids'].numpy().shape[-1]\n self.assertEqual(0, length % bucket_size)\n self.assertEqual((batch_size,), batch['length'].numpy().shape)\n self.assertEqual((batch_size,), batch['label'].numpy().shape)\n\n def test_batched_dataset(self):\n \"\"\"Tests that the length of a batch matches the longest sequence.\"\"\"\n batch_size = 2\n for batch in self.dataset.get_batches(\n batch_size=batch_size, shuffle=False\n ).take(1):\n # Each batch is padded to the maximum sentence length in the batch.\n max_length_in_batch = max(batch['length'].numpy())\n length = batch['token_ids'].numpy().shape[-1]\n self.assertEqual(max_length_in_batch, length)\n self.assertEqual((batch_size,), batch['length'].numpy().shape)\n self.assertEqual((batch_size,), batch['label'].numpy().shape)\n\n def test_batched_dataset_fixed_length(self):\n \"\"\"Tests that each batch has the fixed length.\"\"\"\n batch_size = 2\n fixed_pad_length = 77\n for batch in self.dataset.get_batches(\n batch_size=batch_size, shuffle=False, fixed_pad_length=fixed_pad_length\n ).take(1):\n length = batch['token_ids'].numpy().shape[-1]\n self.assertEqual(fixed_pad_length, length)\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "examples/sst2/main.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Main file for running the SST2 example.\nThis file is intentionally kept short. The majority for logic is in libraries\nthat can be easily tested and imported in Colab.\n\"\"\"\n\nfrom absl import app\nfrom absl import flags\nfrom absl import logging\nfrom clu import platform\nimport jax\nfrom ml_collections import config_flags\nimport tensorflow as tf\n\nimport train\n\n\nFLAGS = flags.FLAGS\n\nflags.DEFINE_string('workdir', None, 'Directory to store model data.')\nconfig_flags.DEFINE_config_file(\n 'config',\n None,\n 'File path to the training hyperparameter configuration.',\n lock_config=True,\n)\nflags.mark_flags_as_required(['config', 'workdir'])\n\n\ndef main(argv):\n if len(argv) > 1:\n raise app.UsageError('Too many command-line arguments.')\n\n # Hide any GPUs from TensorFlow. Otherwise TF might reserve memory and make\n # it unavailable to JAX.\n tf.config.experimental.set_visible_devices([], 'GPU')\n\n logging.info('JAX process: %d / %d', jax.process_index(), jax.process_count())\n logging.info('JAX local devices: %r', jax.local_devices())\n\n # Add a note so that we can tell which task is which JAX host.\n # (Depending on the platform task 0 is not guaranteed to be host 0)\n platform.work_unit().set_task_status(\n f'process_index: {jax.process_index()}, '\n f'process_count: {jax.process_count()}'\n )\n platform.work_unit().create_artifact(\n platform.ArtifactType.DIRECTORY, FLAGS.workdir, 'workdir'\n )\n\n train.train_and_evaluate(FLAGS.config, FLAGS.workdir)\n\n\nif __name__ == '__main__':\n app.run(main)\n" }, { "path": "examples/sst2/models.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"A text classification model.\"\"\"\n\nimport functools\nfrom typing import Any, Optional\nfrom collections.abc import Callable\n\nfrom flax import linen as nn\nimport jax\nfrom jax import numpy as jnp\n\nArray = jnp.ndarray\n\n\ndef sequence_mask(lengths: Array, max_length: int) -> Array:\n \"\"\"Computes a boolean mask over sequence positions for each given length.\n\n Example:\n ```\n sequence_mask([1, 2], 3)\n [[True, False, False],\n [True, True, False]]\n ```\n\n Args:\n lengths: The length of each sequence. [batch_size]\n max_length: The width of the boolean mask. Must be >= max(lengths).\n\n Returns:\n A mask with shape: [batch_size, max_length] indicating which\n positions are valid for each sequence.\n \"\"\"\n return jnp.arange(max_length)[None] < lengths[:, None]\n\n\n@jax.vmap\ndef flip_sequences(inputs: Array, lengths: Array) -> Array:\n \"\"\"Flips a sequence of inputs along the time dimension.\n\n This function can be used to prepare inputs for the reverse direction of a\n bidirectional LSTM. It solves the issue that, when naively flipping multiple\n padded sequences stored in a matrix, the first elements would be padding\n values for those sequences that were padded. This function keeps the padding\n at the end, while flipping the rest of the elements.\n\n Example:\n ```python\n inputs = [[1, 0, 0],\n [2, 3, 0]\n [4, 5, 6]]\n lengths = [1, 2, 3]\n flip_sequences(inputs, lengths) = [[1, 0, 0],\n [3, 2, 0],\n [6, 5, 4]]\n ```\n\n Args:\n inputs: An array of input IDs [batch_size, seq_length].\n lengths: The length of each sequence [batch_size].\n\n Returns:\n An ndarray with the flipped inputs.\n \"\"\"\n # Note: since this function is vmapped, the code below is effectively for\n # a single example.\n max_length = inputs.shape[0]\n return jnp.flip(jnp.roll(inputs, max_length - lengths, axis=0), axis=0)\n\n\nclass WordDropout(nn.Module):\n \"\"\"Applies word dropout to a batch of input IDs.\n\n This is basically the same as `nn.Dropout`, but allows specifying the\n value of dropped out items.\n \"\"\"\n\n dropout_rate: float\n unk_idx: int\n deterministic: bool | None = None\n\n @nn.compact\n def __call__(self, inputs: Array, deterministic: bool | None = None):\n deterministic = nn.module.merge_param(\n 'deterministic', self.deterministic, deterministic\n )\n if deterministic or self.dropout_rate == 0.0:\n return inputs\n rng = self.make_rng('dropout')\n mask = jax.random.bernoulli(rng, p=self.dropout_rate, shape=inputs.shape)\n return jnp.where(mask, jnp.array([self.unk_idx]), inputs)\n\n\nclass Embedder(nn.Module):\n \"\"\"Embeds batches of token IDs into feature space.\n\n Attributes:\n vocab_size: The size of the vocabulary (i.e., the number of embeddings).\n embedding_size: The dimensionality of the embeddings.\n embedding_init: The initializer used to initialize the embeddings.\n frozen: Freezes the embeddings table, keeping it fixed at initial values.\n dropout_rate: Percentage of units to drop after embedding the inputs.\n word_dropout_rate: Percentage of input words to replace with unk_idx.\n unk_idx: The index (integer) to use to replace inputs for word dropout.\n \"\"\"\n\n vocab_size: int\n embedding_size: int\n embedding_init: Callable[..., Array] = nn.initializers.normal(stddev=0.1)\n frozen: bool = False\n dropout_rate: float = 0.0\n word_dropout_rate: float = 0.0\n unk_idx: int | None = None\n deterministic: bool | None = None\n dtype: jnp.dtype = jnp.float32\n\n def setup(self):\n self.embedding = self.param(\n 'embedding',\n self.embedding_init,\n (self.vocab_size, self.embedding_size),\n self.dtype,\n )\n self.dropout_layer = nn.Dropout(rate=self.dropout_rate)\n self.word_dropout_layer = WordDropout(\n dropout_rate=self.word_dropout_rate, unk_idx=self.unk_idx\n )\n\n def __call__(\n self, inputs: Array, deterministic: bool | None = None\n ) -> Array:\n \"\"\"Embeds the input sequences and applies word dropout and dropout.\n\n Args:\n inputs: Batch of input token ID sequences [batch_size, seq_length].\n deterministic: Disables dropout when set to True.\n\n Returns:\n The embedded inputs, shape: [batch_size, seq_length,\n embedding_size].\n \"\"\"\n deterministic = nn.module.merge_param(\n 'deterministic', self.deterministic, deterministic\n )\n inputs = self.word_dropout_layer(inputs, deterministic=deterministic)\n embedded_inputs = self.embedding[inputs]\n\n # Keep the embeddings fixed at initial (e.g. pretrained) values.\n if self.frozen:\n embedded_inputs = jax.lax.stop_gradient(embedded_inputs)\n\n return self.dropout_layer(embedded_inputs, deterministic=deterministic)\n\n\nclass SimpleLSTM(nn.Module):\n \"\"\"A simple unidirectional LSTM.\"\"\"\n\n hidden_size: int\n\n @functools.partial(\n nn.transforms.scan,\n variable_broadcast='params',\n in_axes=1,\n out_axes=1,\n split_rngs={'params': False},\n )\n @nn.compact\n def __call__(self, carry, x):\n return nn.OptimizedLSTMCell(self.hidden_size)(carry, x)\n\n def initialize_carry(self, input_shape):\n # Use fixed random key since default state init fn is just zeros.\n return nn.OptimizedLSTMCell(self.hidden_size, parent=None).initialize_carry(\n jax.random.key(0), input_shape\n )\n\n\nclass SimpleBiLSTM(nn.Module):\n \"\"\"A simple bi-directional LSTM.\"\"\"\n\n hidden_size: int\n\n def setup(self):\n self.forward_lstm = SimpleLSTM(self.hidden_size)\n self.backward_lstm = SimpleLSTM(self.hidden_size)\n\n def __call__(self, embedded_inputs, lengths):\n # Forward LSTM.\n initial_state = self.forward_lstm.initialize_carry(\n embedded_inputs[:, 0].shape\n )\n _, forward_outputs = self.forward_lstm(initial_state, embedded_inputs)\n\n # Backward LSTM.\n reversed_inputs = flip_sequences(embedded_inputs, lengths)\n initial_state = self.backward_lstm.initialize_carry(\n reversed_inputs[:, 0].shape\n )\n _, backward_outputs = self.backward_lstm(initial_state, reversed_inputs)\n backward_outputs = flip_sequences(backward_outputs, lengths)\n\n # Concatenate the forward and backward representations.\n outputs = jnp.concatenate([forward_outputs, backward_outputs], -1)\n return outputs\n\n\nclass MLP(nn.Module):\n \"\"\"A simple Multilayer perceptron with 1 hidden layer.\n\n Attributes:\n hidden_size: The size of the hidden layer.\n output_size: The size of the output.\n activation: The activation function to apply to the hidden layer.\n dropout_rate: The dropout rate applied to the hidden layer.\n output_bias: If False, do not use a bias term in the last layer.\n deterministic: Disables dropout if set to True.\n \"\"\"\n\n hidden_size: int\n output_size: int\n activation: Callable[..., Any] = nn.tanh\n dropout_rate: float = 0.0\n output_bias: bool = False\n deterministic: bool | None = None\n\n def setup(self):\n self.intermediate_layer = nn.Dense(self.hidden_size)\n self.output_layer = nn.Dense(self.output_size, use_bias=self.output_bias)\n self.dropout_layer = nn.Dropout(rate=self.dropout_rate)\n\n def __call__(self, inputs: Array, deterministic: bool | None = None):\n \"\"\"Applies the MLP to the last dimension of the inputs.\n\n Args:\n inputs: [batch_size, ..., input_features].\n deterministic: Disables dropout when set to True.\n\n Returns:\n The MLP output [batch_size, ..., output_size]\n \"\"\"\n deterministic = nn.module.merge_param(\n 'deterministic', self.deterministic, deterministic\n )\n hidden = self.intermediate_layer(inputs)\n hidden = self.activation(hidden)\n hidden = self.dropout_layer(hidden, deterministic=deterministic)\n output = self.output_layer(hidden)\n return output\n\n\nclass KeysOnlyMlpAttention(nn.Module):\n \"\"\"Computes MLP-based attention scores based on keys alone, without a query.\n\n Attention scores are computed by feeding the keys through an MLP. This\n results in a single scalar per key, and for each sequence the attention\n scores are normalized using a softmax so that they sum to 1. Invalid key\n positions are ignored as indicated by the mask. This is also called\n \"Bahdanau attention\" and was originally proposed in:\n ```\n Bahdanau et al., 2015. Neural Machine Translation by Jointly Learning to\n Align and Translate. ICLR. https://arxiv.org/abs/1409.0473\n ```\n\n Attributes:\n hidden_size: The hidden size of the MLP that computes the attention score.\n \"\"\"\n\n hidden_size: int\n\n @nn.compact\n def __call__(self, keys: Array, mask: Array) -> Array:\n \"\"\"Applies model to the input keys and mask.\n\n Args:\n keys: The inputs for which to compute an attention score. Shape:\n [batch_size, seq_length, embeddings_size].\n mask: A mask that determines which values in `keys` are valid. Only\n values for which the mask is True will get non-zero attention scores.\n [batch_size, seq_length].\n\n Returns:\n The normalized attention scores. [batch_size, seq_length].\n \"\"\"\n hidden = nn.Dense(self.hidden_size, name='keys', use_bias=False)(keys)\n energy = nn.tanh(hidden)\n scores = nn.Dense(1, name='energy', use_bias=False)(energy)\n scores = scores.squeeze(-1) # New shape: [batch_size, seq_len].\n scores = jnp.where(mask, scores, -jnp.inf) # Using exp(-inf) = 0 below.\n scores = nn.softmax(scores, axis=-1)\n\n # Captures the scores if 'intermediates' is mutable, otherwise does nothing.\n self.sow('intermediates', 'attention', scores)\n\n return scores\n\n\nclass AttentionClassifier(nn.Module):\n \"\"\"A classifier that uses attention to summarize the inputs.\n\n Attributes:\n hidden_size: The hidden size of the MLP classifier.\n output_size: The number of output classes for the classifier.\n dropout_rate: The dropout rate applied over the encoded_inputs, the summary\n of the inputs, and inside the MLP. Applied when `deterministic` is False.\n deterministic: Disables dropout if True.\n \"\"\"\n\n hidden_size: int\n output_size: int\n dropout_rate: float = 0.0\n deterministic: bool | None = None\n\n def setup(self):\n self.dropout_layer = nn.Dropout(rate=self.dropout_rate)\n self.keys_only_mlp_attention = KeysOnlyMlpAttention(\n hidden_size=self.hidden_size\n )\n self.mlp = MLP(\n hidden_size=self.hidden_size,\n output_size=self.output_size,\n output_bias=False,\n dropout_rate=self.dropout_rate,\n )\n\n def __call__(\n self,\n encoded_inputs: Array,\n lengths: Array,\n deterministic: bool | None = None,\n ) -> Array:\n \"\"\"Applies model to the encoded inputs.\n\n Args:\n encoded_inputs: The inputs (e.g., sentences) that have already been\n encoded by some encoder, e.g., an LSTM. [batch_size,\n seq_length, encoded_inputs_size].\n lengths: The lengths of the inputs. [batch_size].\n deterministic: Disables dropout when set to True.\n\n Returns:\n An array of logits [batch_size, output_size].\n \"\"\"\n deterministic = nn.module.merge_param(\n 'deterministic', self.deterministic, deterministic\n )\n encoded_inputs = self.dropout_layer(\n encoded_inputs, deterministic=deterministic\n )\n\n # Compute attention. attention.shape: [batch_size, seq_len].\n mask = sequence_mask(lengths, encoded_inputs.shape[1])\n attention = self.keys_only_mlp_attention(encoded_inputs, mask)\n\n # Summarize the inputs by taking their weighted sum using attention scores.\n context = jnp.expand_dims(attention, 1) @ encoded_inputs\n context = context.squeeze(1) # [batch_size, encoded_inputs_size]\n context = self.dropout_layer(context, deterministic=deterministic)\n\n # Make the final prediction from the context vector (the summarized inputs).\n logits = self.mlp(context, deterministic=deterministic)\n return logits\n\n\nclass TextClassifier(nn.Module):\n \"\"\"A Text Classification model.\"\"\"\n\n embedding_size: int\n hidden_size: int\n vocab_size: int\n output_size: int\n\n dropout_rate: float\n word_dropout_rate: float\n unk_idx: int = 1\n deterministic: bool | None = None\n\n def setup(self):\n self.embedder = Embedder(\n vocab_size=self.vocab_size,\n embedding_size=self.embedding_size,\n dropout_rate=self.dropout_rate,\n word_dropout_rate=self.word_dropout_rate,\n unk_idx=self.unk_idx,\n )\n self.encoder = SimpleBiLSTM(hidden_size=self.hidden_size)\n self.classifier = AttentionClassifier(\n hidden_size=self.hidden_size,\n output_size=self.output_size,\n dropout_rate=self.dropout_rate,\n )\n\n def embed_token_ids(\n self, token_ids: Array, deterministic: bool | None = None\n ) -> Array:\n deterministic = nn.module.merge_param(\n 'deterministic', self.deterministic, deterministic\n )\n return self.embedder(token_ids, deterministic=deterministic)\n\n def logits_from_embedded_inputs(\n self,\n embedded_inputs: Array,\n lengths: Array,\n deterministic: bool | None = None,\n ) -> Array:\n deterministic = nn.module.merge_param(\n 'deterministic', self.deterministic, deterministic\n )\n encoded_inputs = self.encoder(embedded_inputs, lengths)\n return self.classifier(encoded_inputs, lengths, deterministic=deterministic)\n\n def __call__(\n self,\n token_ids: Array,\n lengths: Array,\n deterministic: bool | None = None,\n ) -> Array:\n \"\"\"Embeds the token IDs, encodes them, and classifies with attention.\"\"\"\n embedded_inputs = self.embed_token_ids(\n token_ids, deterministic=deterministic\n )\n logits = self.logits_from_embedded_inputs(\n embedded_inputs, lengths, deterministic=deterministic\n )\n return logits\n" }, { "path": "examples/sst2/models_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for sst2.models.\"\"\"\nfrom absl.testing import absltest\nfrom absl.testing import parameterized\nimport models\nimport jax\nimport jax.test_util\nimport numpy as np\n\n\n# Parse absl flags test_srcdir and test_tmpdir.\njax.config.parse_flags_with_absl()\n\n\nclass ModelTest(parameterized.TestCase):\n\n def test_embedder_returns_correct_output_shape(self):\n \"\"\"Tests if the embedder returns the correct shape.\"\"\"\n vocab_size = 5\n embedding_size = 3\n model = models.Embedder(\n vocab_size=vocab_size, embedding_size=embedding_size\n )\n rng = jax.random.key(0)\n token_ids = np.array([[2, 4, 3], [2, 6, 3]], dtype=np.int32)\n output, _ = model.init_with_output(rng, token_ids, deterministic=True)\n self.assertEqual((token_ids.shape) + (embedding_size,), output.shape)\n\n def test_lstm_returns_correct_output_shape(self):\n \"\"\"Tests if the simple LSTM returns the correct shape.\"\"\"\n batch_size = 2\n seq_len = 3\n embedding_size = 4\n hidden_size = 5\n model = models.SimpleLSTM(5)\n rng = jax.random.key(0)\n inputs = np.random.RandomState(0).normal(\n size=[batch_size, seq_len, embedding_size]\n )\n initial_state = model.initialize_carry(inputs[:, 0].shape)\n (_, output), _ = model.init_with_output(rng, initial_state, inputs)\n self.assertEqual((batch_size, seq_len, hidden_size), output.shape)\n\n def test_bilstm_returns_correct_output_shape(self):\n \"\"\"Tests if the simple BiLSTM returns the correct shape.\"\"\"\n batch_size = 2\n seq_len = 3\n embedding_size = 4\n hidden_size = 5\n model = models.SimpleBiLSTM(hidden_size=hidden_size)\n rng = jax.random.key(0)\n inputs = np.random.RandomState(0).normal(\n size=[batch_size, seq_len, embedding_size]\n )\n lengths = np.array([2, 3], dtype=np.int32)\n outputs, _ = model.init_with_output(rng, inputs, lengths)\n # We expect 2*hidden_size because we concatenate forward+backward LSTMs.\n self.assertEqual((batch_size, seq_len, 2 * hidden_size), outputs.shape)\n\n def test_text_classifier_returns_correct_output_shape(self):\n \"\"\"Tests if a TextClassifier model returns the correct shape.\"\"\"\n embedding_size = 3\n hidden_size = 7\n vocab_size = 5\n output_size = 3\n dropout_rate = 0.1\n word_dropout_rate = 0.2\n unk_idx = 1\n\n model = models.TextClassifier(\n embedding_size=embedding_size,\n hidden_size=hidden_size,\n vocab_size=vocab_size,\n output_size=output_size,\n dropout_rate=dropout_rate,\n word_dropout_rate=word_dropout_rate,\n unk_idx=unk_idx,\n deterministic=True,\n )\n\n rng = jax.random.key(0)\n token_ids = np.array([[2, 4, 3], [2, 6, 3]], dtype=np.int32)\n lengths = np.array([2, 3], dtype=np.int32)\n output, _ = model.init_with_output(rng, token_ids, lengths)\n batch_size = token_ids.shape[0]\n self.assertEqual((batch_size, output_size), output.shape)\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "examples/sst2/requirements.txt", "content": "absl-py==1.0.0\nclu==0.0.6\nflax==0.3.6\nml-collections==0.1.0\nnumpy==1.22.0\noptax==0.1.0\ntensorflow==2.11.1\ntensorflow-datasets==4.4.0\ntensorflow-text==2.7.0\n" }, { "path": "examples/sst2/sst2.ipynb", "content": "{\n \"cells\": [\n {\n \"id\": \"29fb3c7c\",\n \"cell_type\": \"markdown\",\n \"source\": [\n \"# Flax SST-2 Example\\n\",\n \"\\n\",\n \"\\u003ca href=\\\"https://colab.research.google.com/github/google/flax/blob/main/examples/sst2/sst2.ipynb\\\" \\u003e\\u003cimg src=\\\"https://colab.research.google.com/assets/colab-badge.svg\\\" alt=\\\"Open In Colab\\\"/\\u003e\\u003c/a\\u003e\\n\",\n \"\\n\",\n \"Demonstration notebook for\\n\",\n \"https://github.com/google/flax/tree/main/examples/sst2\"\n ],\n \"metadata\": {},\n \"execution_count\": null\n },\n {\n \"id\": \"5b526b04\",\n \"cell_type\": \"markdown\",\n \"source\": [\n \"**Before you start:** Select Runtime -\\u003e Change runtime type -\\u003e GPU.\"\n ],\n \"metadata\": {},\n \"execution_count\": null\n },\n {\n \"id\": \"89ec78c1\",\n \"cell_type\": \"markdown\",\n \"source\": [\n \"The **Flax Notebook Workflow**:\\n\",\n \"\\n\",\n \"1. Run the entire notebook end-to-end and check out the outputs.\\n\",\n \" - This will open Python files in the right-hand editor!\\n\",\n \" - You'll be able to interactively explore metrics in TensorBoard.\\n\",\n \"2. Change `config` and train for different hyperparameters. Check out the\\n\",\n \" updated TensorBoard plots.\\n\",\n \"3. Update the code in `train.py`. Thanks to `%autoreload`, any changes you\\n\",\n \" make in the file will automatically appear in the notebook. Some ideas to\\n\",\n \" get you started:\\n\",\n \" - Change the model.\\n\",\n \" - Log some per-batch metrics during training.\\n\",\n \" - Add new hyperparameters to `configs/default.py` and use them in\\n\",\n \" `train.py`.\\n\",\n \"4. At any time, feel free to paste code from `train.py` into the notebook\\n\",\n \" and modify it directly there!\"\n ],\n \"metadata\": {},\n \"execution_count\": null\n },\n {\n \"id\": \"7e4ba0dc\",\n \"cell_type\": \"markdown\",\n \"source\": [\n \"## Setup\"\n ],\n \"metadata\": {},\n \"execution_count\": null\n },\n {\n \"id\": \"ee8021b9\",\n \"cell_type\": \"code\",\n \"source\": [\n \"example_directory = 'examples/sst2'\\n\",\n \"editor_relpaths = ('configs/default.py', 'train.py', 'models.py')\"\n ],\n \"metadata\": {},\n \"execution_count\": null\n },\n {\n \"id\": \"36dab290\",\n \"cell_type\": \"code\",\n \"source\": [\n \"# (If you run this code in Jupyter[lab], then you're already in the\\n\",\n \"# example directory and nothing needs to be done.)\\n\",\n \"\\n\",\n \"#@markdown **Fetch newest Flax, copy example code**\\n\",\n \"#@markdown\\n\",\n \"#@markdown **If you select no** below, then the files will be stored on the\\n\",\n \"#@markdown *ephemeral* Colab VM. **After some time of inactivity, this VM will\\n\",\n \"#@markdown be restarted an any changes are lost**.\\n\",\n \"#@markdown\\n\",\n \"#@markdown **If you select yes** below, then you will be asked for your\\n\",\n \"#@markdown credentials to mount your personal Google Drive. In this case, all\\n\",\n \"#@markdown changes you make will be *persisted*, and even if you re-run the\\n\",\n \"#@markdown Colab later on, the files will still be the same (you can of course\\n\",\n \"#@markdown remove directories inside your Drive's `flax/` root if you want to\\n\",\n \"#@markdown manually revert these files).\\n\",\n \"\\n\",\n \"if 'google.colab' in str(get_ipython()):\\n\",\n \" import os\\n\",\n \" os.chdir('/content')\\n\",\n \" # Download Flax repo from Github.\\n\",\n \" if not os.path.isdir('flaxrepo'):\\n\",\n \" pass\\n\",\n \" !git clone --depth=1 https://github.com/google/flax flaxrepo\\n\",\n \" # Copy example files \\u0026 change directory.\\n\",\n \" mount_gdrive = 'no' #@param ['yes', 'no']\\n\",\n \" if mount_gdrive == 'yes':\\n\",\n \" DISCLAIMER = 'Note: Editing in your Google Drive, changes will persist.'\\n\",\n \" from google.colab import drive\\n\",\n \" drive.mount('/content/gdrive')\\n\",\n \" example_root_path = f'/content/gdrive/My Drive/flax/{example_directory}'\\n\",\n \" else:\\n\",\n \" DISCLAIMER = 'WARNING: Editing in VM - changes lost after reboot!!'\\n\",\n \" example_root_path = f'/content/{example_directory}'\\n\",\n \" from IPython import display\\n\",\n \" display.display(display.HTML(\\n\",\n \" f'\\u003ch1 style=\\\"color:red;\\\" class=\\\"blink\\\"\\u003e{DISCLAIMER}\\u003c/h1\\u003e'))\\n\",\n \" if not os.path.isdir(example_root_path):\\n\",\n \" os.makedirs(example_root_path)\\n\",\n \" !cp -r flaxrepo/$example_directory/* \\\"$example_root_path\\\"\\n\",\n \" os.chdir(example_root_path)\\n\",\n \" from google.colab import files\\n\",\n \" for relpath in editor_relpaths:\\n\",\n \" s = open(f'{example_root_path}/{relpath}').read()\\n\",\n \" open(f'{example_root_path}/{relpath}', 'w').write(\\n\",\n \" f'## {DISCLAIMER}\\\\n' + '#' * (len(DISCLAIMER) + 3) + '\\\\n\\\\n' + s)\\n\",\n \" files.view(f'{example_root_path}/{relpath}')\"\n ],\n \"metadata\": {},\n \"execution_count\": null\n },\n {\n \"id\": \"700d9428\",\n \"cell_type\": \"code\",\n \"source\": [\n \"# Note: In Colab, above cell changed the working directory.\\n\",\n \"!pwd\"\n ],\n \"metadata\": {},\n \"execution_count\": null\n },\n {\n \"id\": \"2fbc3e64\",\n \"cell_type\": \"code\",\n \"source\": [\n \"# Install SST-2 dependencies.\\n\",\n \"!pip install -q -r requirements.txt\"\n ],\n \"metadata\": {},\n \"execution_count\": null\n },\n {\n \"id\": \"e40c50cf\",\n \"cell_type\": \"markdown\",\n \"source\": [\n \"## Imports / Helpers\"\n ],\n \"metadata\": {},\n \"execution_count\": null\n },\n {\n \"id\": \"703f04fb\",\n \"cell_type\": \"code\",\n \"source\": [\n \"import os\\n\",\n \"os.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=8'\\n\",\n \"\\n\",\n \"import jax\\n\",\n \"jax.devices()\"\n ],\n \"metadata\": {},\n \"execution_count\": null\n },\n {\n \"id\": \"32e12a97\",\n \"cell_type\": \"code\",\n \"source\": [\n \"from absl import logging\\n\",\n \"import flax\\n\",\n \"import jax.numpy as jnp\\n\",\n \"import numpy as np\\n\",\n \"import tensorflow as tf\\n\",\n \"import tensorflow_datasets as tfds\\n\",\n \"import time\\n\",\n \"logging.set_verbosity(logging.INFO)\\n\",\n \"\\n\",\n \"# Make sure the GPU is for JAX, not for TF.\\n\",\n \"tf.config.experimental.set_visible_devices([], 'GPU')\"\n ],\n \"metadata\": {},\n \"execution_count\": null\n },\n {\n \"id\": \"94ece24d\",\n \"cell_type\": \"code\",\n \"source\": [\n \"# Local imports from current directory - auto reload.\\n\",\n \"# Any changes you make to train.py will appear automatically.\\n\",\n \"%load_ext autoreload\\n\",\n \"%autoreload 2\\n\",\n \"try:\\n\",\n \" import train\\n\",\n \" import models\\n\",\n \" import vocabulary\\n\",\n \" import input_pipeline\\n\",\n \" from configs import default as config_lib\\n\",\n \"except ModuleNotFoundError:\\n\",\n \" # Local imports may not be available in all contexts\\n\",\n \" pass\\n\",\n \"config = config_lib.get_config()\"\n ],\n \"metadata\": {\n \"tags\": []\n },\n \"execution_count\": null\n },\n {\n \"id\": \"c8a6ec00\",\n \"cell_type\": \"markdown\",\n \"source\": [\n \"## Dataset\"\n ],\n \"metadata\": {},\n \"execution_count\": null\n },\n {\n \"id\": \"5c615ca6\",\n \"cell_type\": \"code\",\n \"source\": [\n \"# Get datasets. \\n\",\n \"# If you get an error you need to install tensorflow_datasets from Github.\\n\",\n \"train_dataset = input_pipeline.TextDataset(split='train')\\n\",\n \"eval_dataset = input_pipeline.TextDataset(split='validation')\"\n ],\n \"metadata\": {\n \"tags\": []\n },\n \"execution_count\": null\n },\n {\n \"id\": \"7d7c55cd\",\n \"cell_type\": \"markdown\",\n \"source\": [\n \"## Training\"\n ],\n \"metadata\": {},\n \"execution_count\": null\n },\n {\n \"id\": \"df9d52ed\",\n \"cell_type\": \"code\",\n \"source\": [\n \"# Get a live update during training - use the \\\"refresh\\\" button!\\n\",\n \"# (In Jupyter[lab] start \\\"tensorboard\\\" in the local directory instead.)\\n\",\n \"if 'google.colab' in str(get_ipython()):\\n\",\n \" pass\\n\",\n \" %load_ext tensorboard\\n\",\n \" %tensorboard --logdir=.\"\n ],\n \"metadata\": {},\n \"execution_count\": null\n },\n {\n \"id\": \"f159d072\",\n \"cell_type\": \"code\",\n \"source\": [\n \"config.num_epochs = 10\\n\",\n \"model_name = 'bilstm'\\n\",\n \"start_time = time.time()\\n\",\n \"optimizer = train.train_and_evaluate(config, workdir=f'./models/{model_name}')\\n\",\n \"logging.info('Walltime: %f s', time.time() - start_time)\"\n ],\n \"metadata\": {\n \"tags\": []\n },\n \"execution_count\": null\n },\n {\n \"id\": \"b8e35c72\",\n \"cell_type\": \"code\",\n \"source\": [\n \"if 'google.colab' in str(get_ipython()):\\n\",\n \" #@markdown You can upload the training results directly to https://tensorboard.dev\\n\",\n \" #@markdown\\n\",\n \" #@markdown Note that everbody with the link will be able to see the data.\\n\",\n \" upload_data = 'yes' #@param ['yes', 'no']\\n\",\n \" if upload_data == 'yes':\\n\",\n \" pass\\n\",\n \" !tensorboard dev upload --one_shot --logdir ./models --name 'Flax examples/mnist'\"\n ],\n \"metadata\": {\n \"cellView\": \"form\",\n \"tags\": []\n },\n \"execution_count\": null\n }\n ],\n \"metadata\": {\n \"accelerator\": \"GPU\"\n },\n \"nbformat_minor\": 0,\n \"nbformat\": 4\n}" }, { "path": "examples/sst2/train.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Trains an SST2 text classifier.\"\"\"\nfrom typing import Any, Dict, Optional, Tuple, Union\nfrom collections.abc import Callable, Iterable, Sequence\n\nfrom absl import logging\nfrom flax import struct\nfrom flax.metrics import tensorboard\nfrom flax.training import train_state\nimport jax\nimport jax.numpy as jnp\nimport ml_collections\nimport numpy as np\nimport optax\nimport tensorflow as tf\n\nimport input_pipeline\nimport models\n\n\nArray = jnp.ndarray\nExample = dict[str, Array]\nTrainState = train_state.TrainState\n\n\nclass Metrics(struct.PyTreeNode):\n \"\"\"Computed metrics.\"\"\"\n\n loss: Array\n accuracy: Array\n count: int | None = None\n\n\n@jax.vmap\ndef sigmoid_cross_entropy_with_logits(*, labels: Array, logits: Array) -> Array:\n \"\"\"Sigmoid cross entropy loss.\"\"\"\n zeros = jnp.zeros_like(logits, dtype=logits.dtype)\n condition = logits >= zeros\n relu_logits = jnp.where(condition, logits, zeros)\n neg_abs_logits = jnp.where(condition, -logits, logits)\n return relu_logits - logits * labels + jnp.log1p(jnp.exp(neg_abs_logits))\n\n\ndef get_initial_params(rng, model):\n \"\"\"Returns randomly initialized parameters.\"\"\"\n token_ids = jnp.ones((2, 3), jnp.int32)\n lengths = jnp.ones((2,), dtype=jnp.int32)\n variables = model.init(rng, token_ids, lengths, deterministic=True)\n return variables['params']\n\n\ndef create_train_state(rng, config: ml_collections.ConfigDict, model):\n \"\"\"Create initial training state.\"\"\"\n params = get_initial_params(rng, model)\n tx = optax.chain(\n optax.sgd(learning_rate=config.learning_rate, momentum=config.momentum),\n optax.add_decayed_weights(weight_decay=config.weight_decay),\n )\n state = TrainState.create(apply_fn=model.apply, params=params, tx=tx)\n return state\n\n\ndef compute_metrics(*, labels: Array, logits: Array) -> Metrics:\n \"\"\"Computes the metrics, summed across the batch if a batch is provided.\"\"\"\n if labels.ndim == 1: # Prevent the labels from broadcasting over the logits.\n labels = jnp.expand_dims(labels, axis=1)\n loss = sigmoid_cross_entropy_with_logits(labels=labels, logits=logits)\n binary_predictions = logits >= 0.0\n binary_accuracy = jnp.equal(binary_predictions, labels)\n return Metrics(\n loss=jnp.sum(loss),\n accuracy=jnp.sum(binary_accuracy),\n count=logits.shape[0],\n )\n\n\ndef model_from_config(config: ml_collections.ConfigDict):\n \"\"\"Builds a text classification model from a config.\"\"\"\n model = models.TextClassifier(\n embedding_size=config.embedding_size,\n hidden_size=config.hidden_size,\n vocab_size=config.vocab_size,\n output_size=config.output_size,\n dropout_rate=config.dropout_rate,\n word_dropout_rate=config.word_dropout_rate,\n unk_idx=config.unk_idx,\n )\n return model\n\n\ndef train_step(\n state: TrainState,\n batch: dict[str, Array],\n rngs: dict[str, Any],\n) -> tuple[TrainState, Metrics]:\n \"\"\"Train for a single step.\"\"\"\n # Make sure to get a new RNG at every step.\n step = state.step\n rngs = {name: jax.random.fold_in(rng, step) for name, rng in rngs.items()}\n\n def loss_fn(params):\n variables = {'params': params}\n logits = state.apply_fn(\n variables,\n batch['token_ids'],\n batch['length'],\n deterministic=False,\n rngs=rngs,\n )\n\n labels = batch['label']\n if labels.ndim == 1:\n labels = jnp.expand_dims(labels, 1)\n loss = jnp.mean(\n sigmoid_cross_entropy_with_logits(labels=labels, logits=logits)\n )\n return loss, logits\n\n grad_fn = jax.value_and_grad(loss_fn, has_aux=True)\n value, grads = grad_fn(state.params)\n (_, logits) = value\n\n new_state = state.apply_gradients(grads=grads)\n metrics = compute_metrics(labels=batch['label'], logits=logits)\n return new_state, metrics\n\n\ndef eval_step(\n state: TrainState, batch: dict[str, Array], rngs: dict[str, Any]\n) -> Metrics:\n \"\"\"Evaluate for a single step. Model should be in deterministic mode.\"\"\"\n variables = {'params': state.params}\n logits = state.apply_fn(\n variables,\n batch['token_ids'],\n batch['length'],\n deterministic=True,\n rngs=rngs,\n )\n metrics = compute_metrics(labels=batch['label'], logits=logits)\n return metrics\n\n\ndef normalize_batch_metrics(batch_metrics: Sequence[Metrics]) -> Metrics:\n \"\"\"Consolidates and normalizes a list of per-batch metrics dicts.\"\"\"\n # Here we sum the metrics that were already summed per batch.\n total_loss = np.sum([metrics.loss for metrics in batch_metrics])\n total_accuracy = np.sum([metrics.accuracy for metrics in batch_metrics])\n total = np.sum([metrics.count for metrics in batch_metrics])\n # Divide each metric by the total number of items in the data set.\n return Metrics(\n loss=total_loss.item() / total, accuracy=total_accuracy.item() / total\n )\n\n\ndef batch_to_numpy(batch: dict[str, Array]) -> dict[str, Array]:\n \"\"\"Converts a batch with TF tensors to a batch of NumPy arrays.\"\"\"\n # _numpy() reuses memory, does not make a copy.\n # pylint: disable=protected-access\n return jax.tree_util.tree_map(lambda x: x._numpy(), batch)\n\n\ndef evaluate_model(\n eval_step_fn: Callable[..., Any],\n state: TrainState,\n batches: Iterable[Example] | tf.data.Dataset,\n epoch: int,\n rngs: dict[str, Any] | None = None,\n) -> Metrics:\n \"\"\"Evaluate a model on a dataset.\"\"\"\n batch_metrics = []\n for i, batch in enumerate(batches):\n batch = batch_to_numpy(batch)\n if rngs is not None: # New RNG for each step.\n rngs = {name: jax.random.fold_in(rng, i) for name, rng in rngs.items()}\n\n metrics = eval_step_fn(state, batch, rngs)\n batch_metrics.append(metrics)\n\n batch_metrics = jax.device_get(batch_metrics)\n metrics = normalize_batch_metrics(batch_metrics)\n logging.info(\n 'eval epoch %03d loss %.4f accuracy %.2f',\n epoch,\n metrics.loss,\n metrics.accuracy * 100,\n )\n return metrics\n\n\ndef train_epoch(\n train_step_fn: Callable[..., tuple[TrainState, Metrics]],\n state: TrainState,\n train_batches: tf.data.Dataset,\n epoch: int,\n rngs: dict[str, Any] | None = None,\n) -> tuple[TrainState, Metrics]:\n \"\"\"Train for a single epoch.\"\"\"\n batch_metrics = []\n for batch in train_batches:\n batch = batch_to_numpy(batch)\n state, metrics = train_step_fn(state, batch, rngs)\n batch_metrics.append(metrics)\n\n # Compute the metrics for this epoch.\n batch_metrics = jax.device_get(batch_metrics)\n metrics = normalize_batch_metrics(batch_metrics)\n\n logging.info(\n 'train epoch %03d loss %.4f accuracy %.2f',\n epoch,\n metrics.loss,\n metrics.accuracy * 100,\n )\n\n return state, metrics\n\n\ndef train_and_evaluate(\n config: ml_collections.ConfigDict, workdir: str\n) -> TrainState:\n \"\"\"Execute model training and evaluation loop.\n\n Args:\n config: Hyperparameter configuration for training and evaluation.\n workdir: Directory where the tensorboard summaries are written to.\n Returns:\n The final train state that includes the trained parameters.\n \"\"\"\n # Prepare datasets.\n train_dataset = input_pipeline.TextDataset(\n tfds_name='glue/sst2', split='train'\n )\n eval_dataset = input_pipeline.TextDataset(\n tfds_name='glue/sst2', split='validation'\n )\n train_batches = train_dataset.get_bucketed_batches(\n config.batch_size,\n config.bucket_size,\n max_input_length=config.max_input_length,\n drop_remainder=True,\n shuffle=True,\n shuffle_seed=config.seed,\n )\n eval_batches = eval_dataset.get_batches(batch_size=config.batch_size)\n\n # Keep track of vocab size in the config so that the embedder knows it.\n config.vocab_size = len(train_dataset.vocab)\n\n # Compile step functions.\n train_step_fn = jax.jit(train_step)\n eval_step_fn = jax.jit(eval_step)\n\n # Create model and a state that contains the parameters.\n rng = jax.random.key(config.seed)\n model = model_from_config(config)\n state = create_train_state(rng, config, model)\n\n summary_writer = tensorboard.SummaryWriter(workdir)\n summary_writer.hparams(dict(config))\n\n # Main training loop.\n logging.info('Starting training...')\n for epoch in range(1, config.num_epochs + 1):\n # Train for one epoch.\n rng, epoch_rng = jax.random.split(rng)\n rngs = {'dropout': epoch_rng}\n state, train_metrics = train_epoch(\n train_step_fn, state, train_batches, epoch, rngs\n )\n\n # Evaluate current model on the validation data.\n eval_metrics = evaluate_model(eval_step_fn, state, eval_batches, epoch)\n\n # Write metrics to TensorBoard.\n summary_writer.scalar('train_loss', train_metrics.loss, epoch)\n summary_writer.scalar('train_accuracy', train_metrics.accuracy * 100, epoch)\n summary_writer.scalar('eval_loss', eval_metrics.loss, epoch)\n summary_writer.scalar('eval_accuracy', eval_metrics.accuracy * 100, epoch)\n\n summary_writer.flush()\n return state\n" }, { "path": "examples/sst2/train_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for sst2.train.\"\"\"\nimport sys\n\nfrom absl.testing import absltest\nfrom absl.testing import parameterized\nimport jax\nimport jax.test_util\nimport numpy as np\n\nfrom configs import default as default_config\nimport train\n\n\n# Parse absl flags test_srcdir and test_tmpdir.\njax.config.parse_flags_with_absl()\n\n\nclass TrainTest(parameterized.TestCase):\n\n def test_train_step_updates_parameters(self):\n \"\"\"Tests if the train step updates the parameters in train state.\"\"\"\n # Create model and a state that contains the parameters.\n config = default_config.get_config()\n config.vocab_size = 13\n rng = jax.random.key(config.seed)\n model = train.model_from_config(config)\n state = train.create_train_state(rng, config, model)\n\n token_ids = np.array([[2, 4, 3], [2, 6, 3]], dtype=np.int32)\n lengths = np.array([2, 3], dtype=np.int32)\n labels = np.zeros_like(lengths)\n batch = {'token_ids': token_ids, 'length': lengths, 'label': labels}\n rngs = {'dropout': rng}\n train_step_fn = jax.jit(train.train_step)\n new_state, metrics = train_step_fn(state, batch, rngs)\n self.assertIsInstance(new_state, train.TrainState)\n self.assertIsInstance(metrics, train.Metrics)\n old_param_values = jax.tree_util.tree_leaves(state.params)\n new_param_values = jax.tree_util.tree_leaves(new_state.params)\n for old_array, new_array in zip(old_param_values, new_param_values):\n # Make sure parameters were updated.\n self.assertFalse(np.allclose(old_array, new_array))\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "examples/sst2/vocab.txt", "content": "\n\n\n\nthe\n,\na\nand\nof\n.\nto\n's\nis\nthat\nin\nit\nas\nwith\nan\nfilm\nits\nfor\nmovie\nthis\nyou\nbut\nbe\non\nn't\nby\nmore\n--\none\nat\nthan\nhas\nnot\nabout\nhis\nfrom\nare\nlike\nso\nor\nall\nhave\nmost\nstory\n'\ngood\n...\ninto\nout\ntoo\nwho\n)\nup\ncharacters\ni\nfunny\n(\ncomedy\nif\njust\nno\ndoes\nmuch\nwhat\ncan\neven\n`\nyour\ntheir\nwill\ntime\nsome\n``\nbad\nlittle\n''\nvery\nway\nwhich\nbest\nany\nlove\nbeen\nlife\nmake\nwork\nenough\nthere\nonly\nhe\nmakes\nus\nnew\nmovies\nnever\nsomething\ndo\nthey\nhumor\nthrough\nwas\nwell\naction\ngreat\nwould\nown\nmade\ndirector\nmany\nwe\nreally\nperformances\nplot\ndrama\nher\nhow\ncould\nfilms\nsense\nsee\nsuch\nbetter\nother\nfun\naudience\npeople\nevery\noff\ntwo\nwithout\ncast\nnothing\nfeel\nboth\nwhen\nbeing\nlook\ncharacter\nmay\nshould\nentertaining\nacting\never\nreal\noften\nperformance\nthem\nlong\n:\nwhile\nstill\nworld\nbecause\nscript\nalso\ninteresting\nanother\nheart\nkind\n're\nthose\ndialogue\nhollywood\nminutes\nwatch\nfirst\nscreen\ndown\nfew\nget\nbig\nover\nfar\nthriller\nmight\nhard\nless\nhuman\nmoments\nactors\ntale\ncompelling\nromantic\ncinema\nhad\nrather\nyear\nfamily\nalmost\nmaterial\nend\nwatching\nseen\n-\nworth\n've\nseem\nitself\npicture\noriginal\ntake\nbefore\nmy\ndocumentary\nseems\nwere\nemotional\nafter\nour\nquite\nfind\nold\nthese\nvisual\ncomes\nman\nthings\nback\nfascinating\nmoving\nsweet\nright\nworks\nbetween\nfeels\nhere\nscenes\nfull\ncome\npiece\ndirection\ncare\nyet\n;\nmusic\ngo\ndull\nme\ngoing\ntakes\nspecial\nultimately\nyears\nca\nyoung\nkeep\nmaking\n'll\nanything\nlaughs\namerican\ntimes\nwhy\nsmart\nworst\ngive\ncomic\nexperience\nenjoyable\nleast\ncinematic\nlot\npart\nwhere\nbeautiful\nentertainment\nhistory\nstyle\nsometimes\nthing\nthough\nart\nclever\nkids\naway\ngives\nagain\nhim\ntogether\nbit\nshe\nintelligence\ndark\ngets\nidea\namusing\nengaging\ntheater\npowerful\nsame\ngenre\nintelligent\nonce\nstar\nwomen\nenergy\nsubject\ndid\ncharming\nsurprisingly\nactually\nsummer\nanyone\ncharm\nscreenplay\nwant\npoint\nfilmmaking\nplace\nshort\nnarrative\nsolid\npretty\nflick\naround\nfeeling\nnearly\nfeature\nsilly\nsimply\nwhose\nmanages\nstrong\nface\npredictable\nenjoy\nthink\ntruly\nwar\nwit\noffers\nsay\nshow\ndeeply\ngoes\nknow\nperfect\nsatisfying\nthen\nfans\npower\nneed\nwhole\nalways\neffort\nbecomes\ndone\nspirit\nfresh\nbeautifully\npremise\ntrue\ntrying\nhalf\nquirky\nsince\nthree\nfilmmakers\nsuspense\ndramatic\nhilarious\nportrait\ntone\nhorror\nlast\nunder\nfine\ninterest\neffects\nflat\nrare\nhigh\nrich\nseries\nhours\nprobably\nchildren\neveryone\nromance\nideas\ntouching\n?\nfamiliar\nlooking\nmodern\nremarkable\nstudy\n'd\nespecially\nimagination\npleasure\nwonderful\nboring\nclassic\neasy\neverything\nsmall\nexercise\nleave\nset\ninstead\nlevel\ntitle\nhonest\nstuff\nculture\npast\ndumb\nintriguing\ntv\nwo\nfilmmaker\nlight\nvideo\nactor\nalready\nturn\naudiences\nsad\nstorytelling\nlack\nmatter\nrecent\nstories\nmind\nobvious\ndespite\nput\ntalent\nwritten\nending\nfrench\nimages\nmemorable\nproject\nterrific\nvisually\nserious\nadventure\ncompletely\nwoman\nbecome\nopera\nbeauty\ncamera\ngentle\nlikely\ntalented\nlooks\nemotionally\nmess\nfails\nday\nride\nslow\nsure\ncliches\ncold\nhaving\nhead\nhimself\nreason\nbeyond\ndirected\ngorgeous\ninside\njokes\nleft\nmr.\nbland\nmen\nmelodrama\nproves\nshot\nimpossible\nlow\nways\neasily\nrun\nabove\nhour\nstupid\nthoughtful\ncontrived\nexcellent\nmust\nsimple\ncomplex\ndebut\ndifferent\nelse\neyes\ntired\nugly\nde\nfairly\nlacks\notherwise\nviewer\nbelieve\nbrilliant\ncomedies\neach\nshows\nsort\nviewers\npassion\nwarmth\nattempt\nblack\ncertainly\nparticularly\nturns\nwriting\nplay\nviolence\nwelcome\nwrong\ncheap\nformula\nlost\nsocial\ngenuine\nsoap\nthemselves\nanimation\nbook\ncrime\ndelightful\neither\npersonal\nplays\nrole\nsequences\nthoroughly\nhero\nline\nbarely\nhistorical\nimpressive\nsex\nversion\nappealing\nfact\nhome\nnor\nseeing\nalong\ngags\nquality\nchange\nclichés\ngot\nworse\nadults\nfound\nlives\nmiddle\nold-fashioned\nsurprising\nambitious\ndeath\nengrossing\ngirl\nmessage\nnext\nrunning\ncreative\nfantasy\nimportant\nable\nlive\nnow\npretentious\nworthy\n!\n'm\ndecent\npsychological\nsequel\nends\nimagine\nnice\ntragedy\nwarm\nentirely\nmichael\nnone\nperfectly\nact\ncreepy\nwaste\ndeep\nremains\nsit\nconcept\ninventive\njob\njourney\nlaugh\npictures\nrock\nunfunny\nunsettling\ncool\ninsight\npainful\ntry\nusual\nvision\nwinning\nattention\nconvincing\njohn\nunique\nbring\nmoral\nneither\nmystery\nsatire\nstylish\nagainst\nbelievable\nknows\nleaves\nmaster\nnature\nside\nsuccess\nthin\nartist\nawful\nelements\nlacking\nlead\nreality\nseat\ntedious\nmood\nshallow\nworking\nappeal\nconsiderable\nepic\nfalls\nmoment\nperiod\nprovocative\nsituations\nview\ncreate\ndays\npolitical\nsentimental\ncinematography\nendearing\nhackneyed\nroad\nscene\nsensitive\nwatchable\ncolorful\nearnest\nfinally\ngetting\ngreatest\nhighly\nparents\ndifficult\nhelp\nloud\nuses\namount\ncall\ndelivers\ndry\ngame\npace\nproduction\nspy\nunderstand\nutterly\nwitty\nadaptation\nflaws\nfuture\nhold\nmostly\npossible\nrelationships\nscary\nsecond\nsharp\nstand\nthemes\ncliché\ndepth\nemotions\nkeeps\nmasterpiece\nodd\nsci-fi\nskill\nthroughout\nduring\nedge\ngiven\ngone\npoorly\nseveral\nsuccessful\nsustain\ntouch\ntries\nanimated\ncoming-of-age\ndrag\nexamination\nfour\nmindless\npointless\npromise\nuse\ndisney\nepisode\nflicks\nhit\ntell\ncaptures\ndefinitely\nexecution\nform\nmagic\nmajor\npieces\nscore\nsubtle\nsurprise\ntwists\nunexpected\nwhether\ncomplete\ndeserves\nensemble\ngrand\nintense\nones\npoor\npure\nspielberg\nacross\nactress\nboy\ndepressing\nintimate\nmemory\noscar\npacing\nproblem\nsomewhat\nstraight\nability\napproach\ncareer\ncouple\nexciting\nforced\nintrigue\nlovely\nsurprises\nthinking\ntold\nyork\ndelicate\ngoofy\nmanner\noffer\nrecommend\nstarts\ntaste\ntension\nthought\ntragic\naffection\nbrings\nfinal\ngrace\nmediocre\nsexy\ncontemporary\neye\nfemale\nheartfelt\nhumanity\noccasionally\novercome\nsets\nslightly\nsomeone\ntribute\nwords\nair\nallows\ncore\ncultural\nequally\nevil\nfully\ngeneric\nhand\nmanipulative\nnicely\npleasant\nplenty\npurpose\nremake\nstale\nstrange\ntoday\ntwisted\nuneven\nworthwhile\ncase\nevents\nfailure\nfelt\nframe\nfunnier\nmoney\nmonster\npainfully\nplayful\nsophisticated\nstriking\ntelling\nvivid\nweird\nwild\ncar\ncrafted\ndelight\ndoing\nhuge\nneeds\noffensive\nroom\nstoryline\nterrible\nclear\ncoming\nexpect\nlet\nliving\nmean\nmotion\nplaying\npoetry\nremarkably\nsingle\nsoul\nthanks\nurban\nacted\ncasting\nclarity\ndevoid\ninsightful\ninspiring\nsounds\nsupposed\nultimate\namong\nclass\nforget\nimaginative\njustice\nrarely\nclearly\ncompany\nenergetic\nexactly\nissues\nquiet\nsandler\nwants\nchilling\nclose\ndead\ndog\nfigure\nfunniest\ngenuinely\nparty\nrefreshing\nrespect\nruns\nsituation\nsmile\nspiritual\nstrangely\ntoward\nwildly\ncalled\ncountry\nempty\nextraordinary\nfestival\nguys\nminor\nnight\nproduct\nrealism\nschool\nsincere\nsloppy\nstep\ntiresome\napart\nchemistry\ncrowd\ncut\nmeditation\nmix\nshare\nsitting\ntaking\nteen\ntreat\ntruth\nviewing\nwasted\nwriter-director\nyourself\nago\nalone\nannoying\nbegins\neffective\ngripping\nms.\noriginality\npolitics\nridiculous\nwonderfully\nadd\nages\nbrain\ncute\ndream\nentire\nname\nproblems\nquestions\nsuperficial\ntarget\nuniversal\nvalue\nwise\n90\nchance\nconclusion\nderivative\nexcitement\nfare\nfire\nii\ninterested\nlame\nloss\nmerely\nplain\nsexual\nsick\ntwice\n10\nabsolutely\namateurish\namerica\ncreated\ndetailed\nextremely\nfirst-rate\nflawed\nimpact\nkid\nlatest\nliked\nputs\nshots\nslapstick\nsteven\nsucceeds\ntaken\nunnecessary\nwell-made\nattraction\nbattle\nbritish\nexcept\nhardly\nlots\nnatural\npopcorn\nrobert\nsappy\nstart\nsuffers\ntriumph\ndate\nfrustrating\nheartwarming\nhorrible\njoy\nmark\nmeandering\nplaces\nprevious\nread\nserved\nspectacle\ntypical\nusually\nalive\nattempts\nchan\ncity\ncredits\ndisaster\ndownright\ndreary\nhope\nincredibly\nmesmerizing\nnovel\nopportunity\nperformers\nseriously\ntreatment\nbehind\ncertain\nforgettable\nfrequently\ngirls\nholds\nindie\noverall\nrealistic\nrefreshingly\nrest\nstars\nsuspenseful\ntop\nused\nabsorbing\nallen\nbalance\nbrown\ncredit\ndelivery\nelegant\ngrant\nguilty\ninept\noddly\nprovides\nriveting\nstunning\ntrip\nweak\nwish\nyarn\navoid\nconventional\nconviction\ndazzling\ndoubt\nfinds\nfranchise\nperson\npop\nprocess\nreveals\nshame\nsuperb\nuntil\nvisuals\nwaiting\n2002\namazing\nbadly\nbased\ncrisp\ndavid\nediting\nguy\nlazy\nplease\npoetic\nthrills\ntrouble\n2\nachievement\nalthough\nbullock\ncollege\ncommercial\ncomplicated\ndamned\nfalse\nhate\nidentity\nincluding\nleads\nmaybe\nsubjects\nwonder\nadmirable\nb\nbehavior\nbore\nbright\ncomedic\ncreates\ndreams\nescape\nexquisite\nfashion\nfears\nlistless\nludicrous\nmatch\nnumber\nroutine\nsequels\nskin\nsubstance\nundeniably\ncleverly\nfavor\nlousy\nmeans\nmurphy\nnasty\npathetic\nrip-off\ntalents\nthought-provoking\nbizarre\ncause\nexcuse\nfiction\nglory\ngrief\ninspired\nmainstream\npay\npoignant\nquickly\nsentimentality\nsheer\nstandard\ntotally\nvarious\nwater\nwin\nwriter\nzone\nage\nauthentic\nbusiness\ncentral\ncheesy\ndamage\ndeal\ndie\nevocative\nfaith\ngay\nhappy\nhell\nholes\nphilosophical\npromising\nrepetitive\nsides\nstarring\nsuperior\ntender\nvoice\naspects\nbeneath\ncharisma\ncolor\ncreativity\ndevastating\nfall\nforgotten\nfriendship\nfront\ngifted\nglimpse\nheavy\nhumorous\nlargely\nmembers\nreach\nremember\nrevealing\nsequence\nsound\nsubtlety\ntheaters\nblue\ncruel\ndelightfully\ndeserve\ndesigned\ndirectors\nfast\nfilled\nfive\nflair\nfrank\ngeneral\ngross-out\nhaunting\nintended\ninvolving\nothers\nperspective\npoints\nraw\nunpleasant\nword\naccomplished\nconfusing\ndocument\nendless\nexhilarating\nfit\ngrowing\nintensity\nlaughter\nlevels\nparticular\nretread\nsign\nsongs\nspark\nstay\nstock\nstudio\nvirtually\nacts\nartificial\nbudget\nchallenging\ncliched\nconsistently\nentertain\nexploration\nfeel-good\nfinding\nforce\nformulaic\ngenerally\ngeorge\nhappens\njackson\nlines\nmodest\nmurder\nobservations\npossibly\nprovide\nquietly\nrhythms\nstereotypes\ntheme\nutter\nvibrant\naccessible\naffair\napparent\nbears\ndecades\nessentially\nfarce\nfatal\ngroup\nhappen\nhouse\nimax\nmarvelous\nnoir\noffering\npull\nrhythm\nsoundtrack\nstructure\nbelly\nbored\nbreak\nbuoyant\ncrazy\ndelights\ndeliver\nfake\nfavorite\nhip\nindustry\njoke\nlets\nmissing\nobnoxious\npersona\npotential\nsight\nslight\nsum\ntelevision\nunconvincing\nupon\naddition\nanswers\nchase\nclumsy\nenjoyed\nextreme\nflashy\nfurther\nhugely\nirritating\nloses\nlow-budget\nmain\nmasterful\nmention\nniro\nred\nrewarding\nsaw\nsouth\nspent\ntimely\nvehicle\nwashington\nachieves\nastonishing\nattractive\nbeginning\nbody\nbother\ncarries\ncommunity\ncomputer\ndeadly\ndisappointing\neloquent\nempathy\nera\nfrustration\ngradually\nhearts\nheaven\nlogic\nmixed\nnecessary\nolder\nplotting\nseason\nself-conscious\nsitcom\nstruggle\ntalking\ntells\nunimaginative\nvisceral\nwait\ncrap\ndirectorial\nexpected\ngreen\nhoffman\nincoherent\nking\nlaughing\nlived\nmarks\nmayhem\nmildly\nmove\nmoviegoers\nonto\npair\nrelationship\nrepresents\nsea\nslice\nsource\nterms\nundeniable\nunfocused\nunpredictable\nuplifting\nviolent\nwannabe\nwholly\nwife\n20\nabsurd\nadult\naffecting\nbond\ncartoon\nchristmas\ncrude\ndevelopment\nexploitative\nfan\nfat\ngangster\nholiday\nhowever\njoyous\nlarge\nmelodramatic\noffbeat\noutstanding\npoem\nresist\nsaturday\nsave\nsetting\nstandards\nsupporting\ntear\nwedding\namounts\nassured\nboys\nchoppy\ndeft\nemotion\nfantastic\ngames\ngrandeur\nhorrifying\nintoxicating\nlist\nlyrical\nmature\noffice\npleasures\nprecious\nreading\nroll\nsensual\nsmarter\nsomehow\nstuck\nsympathy\ntough\nvulgar\nweight\nwickedly\natmosphere\nbrought\ncaptivating\ncharmer\nconflict\ncontent\ncourage\ncrush\nedited\nendure\nkiddie\nlady\nlength\nmachine\nmother\npale\npassionate\npat\npersonality\nprice\nrealize\nrevelatory\nseagal\nshoot\nsoderbergh\ntears\ntrademark\nturned\nvapid\nwalk\nwell-crafted\nwry\n$\naverage\nbanal\ncheck\ncollection\ndetail\ndigital\ndirty\nendeavor\nguns\nhits\nhopeful\nintricate\nkiller\nlawrence\nmissed\nmurky\npack\npain\npositive\npotentially\nrealized\nrendered\nresonant\nservice\nsimplistic\nsink\nsour\nsplendid\nstage\nteenagers\nunlikable\nunoriginal\nvividly\nyes\nabsurdity\nadam\nbiting\nbittersweet\nblade\nbrothers\nbunch\ncharismatic\nconstructed\ndeftly\neventually\nfilmed\nfood\nfreshness\ngem\nheartbreaking\njapanese\nlate\nordinary\npoignancy\nportrayal\nsomewhere\nspirits\nstrength\nwars\namericans\nblend\ncombination\ncult\ndamn\ndvd\neffect\nemerges\nenthusiasm\nessence\nexample\nfear\nfootage\nfrightening\ngod\ngrow\ninstantly\ninvolved\nloved\nlush\nminds\nmixture\nplodding\npresent\nproceedings\nprofound\npulls\nresults\nricher\nsaid\nself-indulgent\nsell\nsimultaneously\nstrain\nstreet\nsurface\nterritory\ntouches\nunfolds\nwilling\n9\nartists\nband\nbold\nbreathtaking\ncapture\ncommentary\ncompassion\ndaring\ndeeper\ndignity\nexceptional\nfaster\nfight\nfriend\nhonestly\nknowledge\nleaving\nmeaningful\nnotice\npeter\nqueen\nrelease\nreturn\nrush\nsoon\nspooky\nsports\ntrite\nwhatever\nadventurous\nangst\nawkward\nbite\nbottom\ncapable\nchanging\nchannel\nchild\ncondition\ndesperate\ndetails\ndetermined\ndrawn\nendlessly\nfanciful\nfate\nimmensely\nincreasingly\nindian\ninsights\nkevin\nmagnificent\nmeasure\nmoore\nnightmare\nnowhere\nodds\nover-the-top\noverly\npaced\npic\npreachy\npredecessor\npresents\nquick\nryan\nsadly\nsane\nsmug\nsoulful\nspace\nspectacular\nstirring\nstring\nsweetly\nsweetness\nteam\ntiming\ntwist\nunforced\nunusual\nvalues\nwears\nwhite\nwilliams\nanimal\nb-movie\nbitter\nbroad\ncarried\ncomfortable\ncommitted\ndelivered\ndemands\ndish\nfeatures\nfriends\ngrab\nground\nharrowing\nheavy-handed\ninspiration\njames\nknowing\nlaughable\nliterary\nlosing\nmagical\nmovements\nneeded\npaper-thin\npatience\npolished\npretends\nprotagonist\nquestion\nrank\nrelentlessly\nshapeless\ntype\nunderstanding\nvital\nweek\nwhom\nwow\namused\nappalling\nartistic\nasks\nbigger\nburns\ncliche\nconfident\nconvinced\ncrass\ncredible\nearly\nfaithful\nfriday\ngrows\nhair\nhands\nhappening\nhomage\nitalian\nlonger\nmartin\nmonty\nnomination\npass\nprove\npublic\nreal-life\nrelatively\nrelief\nresourceful\nsadness\nscience\nscorsese\nsketchy\nstartling\nstrike\nsympathetic\nthird\ntiny\ntop-notch\nuncompromising\nunfortunately\nambition\napparently\nappear\naward\nbroken\ncelebration\ncreating\ndegree\ndespair\ndirect-to-video\ndisguise\nexploitation\nfair\nfinale\nfinish\nflimsy\nforgive\ngritty\nhappiness\nhip-hop\nhonesty\ninnocence\nkey\nlesson\nloves\nmusical\npleasing\npresence\nroles\nscreenwriter\nseek\nserves\nsigns\nsnow\nsparkling\nstate\nteacher\nticket\nuninspired\nwinner\nwitch\nwitless\nahead\naplomb\nappears\nbarbershop\nbelt\nbiggest\nbrutally\nbuy\ncharms\ncloying\nconcerned\ncop\ncorny\ncynical\ndepiction\ndeveloped\nearnestness\nenthusiastic\nescapism\nevent\nfights\nfolks\ngiant\nguessing\nhastily\nimagery\nimpression\ninevitable\nleaden\nlee\nmaudlin\nmerit\nmetaphor\nnote\noutrageous\nperverse\nrecycled\nrelies\nreport\nrooting\nsand\nscenery\nseconds\nshocking\nshooting\nskip\nslap\nstrikes\nstructured\nsuffering\ntechnical\nthrill\ntortured\ntreasure\ntrick\nwalked\nworked\n&\nabandon\naccept\nadmire\nalternative\nbenigni\nbreathtakingly\ncharacterizations\ncharmless\nchilly\ncohesive\ncomfort\ncommunal\nconvenient\ndelicious\ndelivering\ndesperately\ndoze\ndrive\ndumbed-down\nexecuted\nfeat\nfighting\ngiving\nherzog\ninsanely\nintentions\nirresistible\njuvenile\nlifeless\nlingering\nlively\nmale\nmanage\nmessages\nmiss\nmysterious\nnevertheless\nniche\noverwrought\npotent\npresented\nracial\nreasonably\nreferences\nsoulless\nspirited\nstaged\nterribly\nterrifying\ntestament\ntom\ntotal\nunfaithful\nurge\nvampire\nvariety\nverve\nwilde\nadventures\nastonishingly\nbelow\nbreaks\nbrutal\ncame\nclumsily\nconstant\ndeceptively\ndescribe\ndesire\ndisplay\ndubbed\nexceedingly\nforces\nfree\ngood-natured\nheard\nintegrity\ninternational\nlimited\nmoviemaking\nnational\node\npark\npile\npowers\npsychology\nquirks\nrises\nsake\nsecret\nshelf\nsilliness\nsmartly\nson\nspite\nsplendor\nstop\ntheatrical\nvillain\nvisit\nwarning\nalternately\nappreciate\nbasic\nbelongs\nblast\nchaotic\nchris\ncoherent\nconnection\ndemanding\ndirecting\ndistasteful\ndraws\ndullest\nearth\nembarrassment\nexotic\nfeces\ngently\ngraphic\nhalfway\nhelps\nhot\njason\nlabored\nlandscape\nlocations\nlooked\nmouth\noccasional\noverwhelming\nphony\nplayed\nplots\nplotted\npolice\nprivate\nproduced\nproud\nremain\nrent\nresult\nrevenge\nrising\nroger\nscratch\nscreenwriting\nsensational\nshamelessly\nsignificant\nsomber\nspare\nspell\nstraightforward\nsuck\nsweeping\ntasty\ntouched\ntraditional\ntrash\nturning\nunlikely\nwallace\nx\nyounger\nachingly\nanalyze\nblood\nbox\nburied\nchildhood\nchoices\nconsidered\ncontrol\ncourse\ndisappointed\ndiscovery\ndiversion\ndogs\ndrugs\ndue\ndying\neastwood\nembraces\nenthralling\nequivalent\nevening\nextended\nfactor\nfarts\nflow\nhandled\nimitation\nindeed\ningenious\ninherent\njones\nmotivated\nnation\nnear\nopen\nopenness\npaper\npolanski\nprofessional\nrace\nrichly\nrowdy\nsaccharine\nsafe\nscreaming\nscreenwriters\nshowtime\nshyamalan\nslack\nsociety\nsoggy\nsuited\nthumbs\ntradition\ntricks\nunlike\nwanting\nweirdly\nwooden\nafraid\nambiguous\narea\nartifice\nattract\nbeach\nbehold\nbroomfield\ncandy\ncapacity\ncheese\nchinese\nconsider\ndelicately\nefforts\nemerge\nengaged\nexpense\nexperiences\nfierce\nfish\nfollow\ngenerous\ngloriously\ngreek\nhandsome\nheld\ninane\nincompetent\ninspirational\nintermittently\njim\njunk\nlackluster\nlanguage\nmaintains\nmartha\nnonsensical\nonscreen\npath\npeculiar\nperhaps\npersonalities\npick\nplayers\npolitically\nquest\nrelevant\nrequired\nrewards\nsees\nsentiment\nstarted\nsubstantial\nthrows\ntragedies\ntranscends\nuncanny\nunflinching\nusing\nvague\nvitality\nwhimsy\nwindow\nwithin\nwondering\nwoody\n90-minute\nactresses\nbeat\nbringing\ncomplexity\ndestination\ndevices\nengage\nenglish\nenter\nexplore\nfantasies\nfather\nfill\nfinely\nfuzzy\ngarbage\ngeneration\ninsipid\ninspire\nintellectual\nintimacy\niranian\nkeeping\nkissinger\nkline\nknew\nlayers\nlocales\nmainly\nmakers\nmarriage\nmere\nmisery\nmorning\nmyself\nphotography\npossibilities\nranks\nrehash\nreturns\nrevolution\nrude\nrushed\nsarandon\nscream\nserve\nslapdash\nslip\nslowly\nsluggish\nstays\nstomach\nsubtly\nsumptuous\ntechnology\nterm\nthematic\ntopic\ntrailer\ntrappings\ntremendous\nuncommonly\nunderstands\nverbal\nvideos\nwent\nwon\n***\narts\naside\nastounding\nattitude\nawe\nbag\nbegin\nbreakthrough\nbreezy\nbuffs\ncelebrates\ncelluloid\nchoice\ncleverness\ncomplications\ncontext\nconvince\ncover\ndaily\ndealing\ndisgusting\ndistance\ndistinctive\ndistinguished\ndrags\ndramatically\ndysfunctional\nedgy\nelaborate\nembarrassed\nencouraging\nexistential\nfamilial\nfamous\nfiring\ngraceless\nholocaust\nhoping\nimprobable\nincisive\njack\nlikable\nloose\nlopez\nlow-key\nmean-spirited\nmild\noverlong\npacked\npaid\nparker\npleasurable\nportrays\nproblematic\nproper\nrefusal\nregret\nreleased\nreminds\nsecrets\nshaped\nshrewd\nsincerity\nsinging\nstands\nstation\nstereotypical\nsucceed\nsupremely\ntalk\ntext\nunbelievable\nunexpectedly\nvs.\nweeks\nweirdness\nwind\n/\naccents\nactual\nadolescent\naltogether\nask\nbrilliantly\nburn\ncamp\ncentury\nchances\nchills\ncount\ndancing\ndesperation\ndiverting\ndrawing\neager\neat\nelegantly\nenormously\nexcess\nexisted\nexit\nexpectation\nexposition\nfisher\nfollowing\ngenial\ngiddy\ngyllenhaal\nharry\nhokum\nhopkins\nhumorless\nhybrid\nignored\nincoherence\ninconsequential\ninfidelity\ninstincts\ninsulting\ninvention\nkicking\nkiss\nkong\nlends\nloads\nlovers\nmistake\nmovement\nneat\nnerve\nnewcomer\npathology\nplanet\nredundant\nreminiscent\nreplaced\nresolutely\nrespectable\nresponsible\nsatirical\nsavvy\nscenario\nseats\nsettles\nshake\nshocks\nsubculture\nsurely\ntalky\ntechnically\ntimeless\ntrees\ntrifle\nweb\nwhimsical\nadds\naims\nballot\nblair\nbloody\ncartoonish\nclueless\nconcern\nconfused\ncopy\ncreation\ncrucial\ncuriously\ndismissed\ndownbeat\nelevate\nenduring\nenormous\nexistence\nfinest\nfreaks\ngangs\nghetto\ngraceful\nhallmark\nhighest\nhistorically\nhunter\nhusband\nimpeccable\ninsult\nintolerable\njarring\njolt\nmartial\nmeaning\nmidnight\nnoise\nnotion\nobviously\npaul\npity\nponder\npopular\npredictably\nrandom\nrecord\nscattered\nsevere\nshining\nsimilarly\nsolondz\nsouls\nspoof\ntitles\nunfulfilling\nunusually\nuseless\nya-ya\n'70s\naccurate\naimless\nblind\nblockbuster\nbrooding\nchops\nclue\ncobbled\ncomparison\ncontinues\ncritics\ndenial\ndesign\ndickens\ndivine\ndomestic\ndragon\neasier\neddie\nenterprise\nexploiting\nflavor\nfluid\nformer\nfrailty\nhandful\nheroine\nimplausible\ninviting\njackie\nlearn\nmarginal\nmasterfully\nmiddle-aged\nmiracle\nmoved\nmuddled\nmythic\nnecessarily\nnoble\nnuanced\nobservation\nphenomenal\npianist\nponderous\npow\npreposterous\npresentation\nprotagonists\nproviding\nrating\nresonance\nrollicking\nrote\nsade\nsatisfies\nsaving\nseven\nsimilar\nsketch\nsong\nsordid\nsorvino\nspend\nstilted\nstretches\nteenage\nthinks\ntrapped\nunabashedly\nvanity\nwitherspoon\nwoo\nabsolute\nacademy\naccount\nadded\nallow\nand/or\nanybody\nartful\nblank\nbritney\nbroadway\nbuilt\ncalm\ncaper\ncares\ncaught\nchabrol\nclassics\ncollapses\ncombined\ncompetent\nconservative\nconstruct\ncreature\ncuriosity\ncurrent\ndangerous\ndeadpan\ndemented\ndisposable\ndrop\nelegance\nempowerment\nenjoyment\nequilibrium\nevelyn\nfancy\nfigures\nfrozen\nfrustratingly\nghost\ngorgeously\ngratuitous\ngrown-up\nhong\nhugh\nhype\nincongruous\ninhabit\ninnovative\nkill\nkung\nlump\nmalaise\nmanufactured\nmarvelously\nmorality\nmorally\nmushy\nobligatory\noutside\npitch-perfect\npostcard\npoverty\npretension\nprimarily\nradical\nrecipe\nrecovery\nremembering\nremote\nreplete\nscripts\nspider-man\nstinker\nsudden\nsurrounding\nswallow\ntasteful\ntepid\ntranslation\nvincent\nwave\n30\nappropriate\nawkwardness\nbill\nbits\nblandness\nbogs\nboredom\nbreath\ncaring\ncatch\ncautionary\ncharge\nclassical\ncommon\ncross\ndawns\ndawson\ndim-witted\ndoses\ndrown\nease\neffectively\nelement\nenjoyably\nequal\nesther\nevident\nexcesses\nexpressive\nfilling\nfluff\nfreeman\nharsh\nhideously\nhoot\nimpersonal\nimpressed\nincessant\ninconsistent\nirish\nironic\nkidman\nknown\nland\nlanguorous\nlarger\nletting\nlighting\nloneliness\nlurid\nmachismo\nmadonna\nmemories\nmisogyny\nmoves\nmusicians\nnotable\nnumbers\nobservant\npantheon\npertinent\npsyche\npulling\npulpy\nrecording\nredemption\nreel\nreign\nreputation\nrise\nruined\nseeking\nsensitivity\nshenanigans\nsleep\nsneaks\nsobering\nsomebody\nspectacularly\nspot\nstiff\nstore\nstunts\nthick\ntravel\ntruthful\nunsentimental\nvillains\nwell-written\nwet\nwillis\nwinds\nwonders\n1/2\n101\na-list\nache\nadmit\naging\nalbum\nalienation\nbaffling\nbecoming\nblown\nbone\nbound\nbracing\nbusy\ncarol\ncarvey\ncatches\nconcert\ncredibility\ncrimes\ndahmer\ndarker\ndemons\ndemonstrates\ndifference\neerie\nemptiness\nexceptionally\nexpectations\nexpress\nfamilies\nfiennes\nfills\nfits\nflashes\nfragile\ngodard\nheal\nhoward\nice\nidiotic\nimmediate\nimpressions\njennifer\nlonging\nmisfire\nmoody\nnaturally\nnonsense\npanache\nparody\nphones\npiano\npositively\npretensions\nproperly\npungent\nravishing\nrun-of-the-mill\nsanta\nscope\nself-consciously\nshameless\nshanghai\nshines\nshock\nslick\nsly\nsolidly\nsophomoric\nspin\nstatic\nsteal\nstylistic\ntable\ntextbook\ntoilet\ntorture\nunderbelly\nunderstated\nvintage\nwanted\nwhiny\nwrapped\n88-minute\nabuse\nadvantage\naffirming\nanger\nantwone\narnold\nartistry\nartsy\naspect\nauthority\nawfully\nbalanced\nbeast\nblow\ncapturing\ncarry\ncell\ncharacteristic\nchoose\nclaustrophobic\nclimactic\nclothes\nconcession\ncontradictory\ncontribution\ncontrivances\nconvey\ncost\ncostly\ndare\ndroll\ndud\nerotic\nethnic\nexaggerated\nfable\nfailing\nfallen\nfilming\nflamboyant\nflourishes\nformat\ngood-hearted\ngrowth\nhack\nhappened\nheroes\nhigh-concept\nhigher\nhitting\nhole\nhollow\nholly\nholm\nideal\nidiots\nilluminating\nincomprehensible\ninsistent\nirony\njustify\nkinda\nkorean\nlayered\nlessons\nliberating\nlie\nlifts\nlightweight\nload\nloving\nlust\nmarket\nmassive\nmeaningless\nmeat\nmile\nmomentum\nmoviegoing\nnose\nobjective\norganic\npayoff\npays\npit\npointed\npokemon\nprofile\npunch\nquarter\nraised\nremembered\nrohmer\nseemingly\nself-discovery\nsensibilities\nshockingly\nsisters\nsnipes\nspeak\nsplash\nstately\nstatus\nstruggled\nstunt\nsubversive\nsusan\nswinging\nteeth\ntest\nthousands\nthrillers\ntragically\ntreats\nuncertain\nuncomfortable\nuninteresting\nuntalented\nwholesome\nwide\nwondrous\n1\n15\narrive\nawkwardly\nbackground\nbeings\nbelieving\nblatant\nbodies\nbombastic\nbrother\nbruce\nbuilding\ncable\nchest\nchildlike\nchord\nchristian\nclichéd\ncloser\nconceits\nconsciousness\ncontains\nconveying\ncorner\ncriminal\ndentist\ndetermination\ndirects\ndisappointment\ndisjointed\ndocumentaries\ne.t.\nellen\nentertained\nexcruciating\nfoul\nfrance\ngenerate\ngenerates\nglaring\nhedonistic\nhilarity\nhorribly\nhumane\nimage\nimprovised\nincredible\ninformative\ninnocent\ninsane\ninsultingly\nkaufman\nkwan\nlater\nlatter\nlunacy\nmatters\nmonsters\nnarcissism\nnews\nnobody\nobscure\nopaque\nourselves\nouttakes\npaints\nparable\nparts\nperceptive\npitfalls\nplotless\npokes\npompeo\npredecessors\nproportions\npurely\nralph\nrecommendation\nrecommended\nrely\nreviews\nring\nsacrificing\nscores\nsearch\nself-satisfied\nsensibility\nseparate\nshowing\nshrill\nsimone\nsinks\nslo-mo\nspring\nsquanders\nsubplots\nsugar\nsurreal\nsurvival\ntediously\ntends\nthrilling\ntomorrow\ntones\ntorn\nunabashed\nunconventional\nunintentionally\nunnerving\nunsympathetic\nupper\nvictims\nwear\nwoefully\nwriter/director\naccent\naction-packed\nafrican-americans\nalabama\nalluring\namazingly\nanimals\nart-house\nassassin\natlantic\navary\navoids\nbar\nbarney\nbig-screen\nblame\nbleak\nbrand\nbridge\nbuilds\ncasual\ncategory\nchanges\nchina\ncircumstances\nclockstoppers\ncombines\nconfidence\nconspiracy\nconstantly\ncritical\ncry\ncurious\ncuteness\ncutting\ndares\ndarkly\ndefeated\ndefies\ndepartment\ndisappointingly\ndisguised\ndistant\ndreadful\ndrug\neccentric\neventual\nevokes\nexact\nexpert\nexplosion\nexpression\nexterior\nfail\nfearlessness\nfeathers\nfeaturing\nflaccid\nfocus\nglorious\nglossy\ngory\ngreater\nheads\nidealism\nimpenetrable\nimpressively\nimprovement\nincludes\nindependent\nindifference\nindulgent\ninnuendo\ninsomnia\ninsurance\njazzy\njordan\nkahlo\nleaky\nlife-affirming\nmeasured\nmike\nmillion\nminimalist\nnifty\nno.\nparade\nparanoia\nphotographed\npicture-perfect\npie\npleaser\nplight\npool\nprecision\nprison\nproduce\npryor\nqualities\nraise\nrambling\nrefuses\nrental\nresemble\nrevelation\nsacrifice\nschwarzenegger\nscratching\nsecretary\nself-deprecating\nself-important\nshoddy\nsignificance\nsimplicity\nsolutions\nsorry\nspeaking\nsuccumbs\nsuits\nsustains\ntap\ntaylor\ntexture\nties\ntour\ntrashy\nundercuts\nunrewarding\nupbeat\nvaluable\nwallop\nwebsite\nwell-acted\n20th\nabstract\naimed\nalien\nam\namiable\nanderson\narc\nbaggage\nbase\nboasts\nbrightly\nbrilliance\nbrooklyn\ncasts\nchuckle\nclinic\nclosing\ncomedian\ncomedy-drama\ncompellingly\nconsequences\ncontroversy\nconventions\ncostume\ndashing\ndavis\ndeaths\ndeceptions\ndecidedly\ndeniro\ndesired\ndishonest\ndisregard\ndiverse\ndose\ndramas\ndrunken\nduvall\nears\nemphasis\nepisodes\ner\nevans\nextraordinarily\nfathers\nfirst-time\nfix\nfly\nfoolish\nforeign\nfrida\ngave\ngifts\nglass\ngoals\ngosling\ngravity\ngrey\nguaranteed\nhalloween\nham-fisted\nharvard\nheated\nhundred\nhuston\ninnovation\ninterminable\nj.\njerry\njuicy\njumps\njunior\nlane\nlast-minute\nliberation\nlies\nlighthearted\nlimitations\nliterate\nmanhattan\nmeditative\nmonotonous\nnicholas\nnonjudgmental\nnostalgic\nnotorious\nnuance\nopens\npalma\npants\nparadiso\npopulated\npowerfully\nproducer\nprofane\nprofessionals\nprojects\npromised\npromises\npsychologically\npuerile\nputting\nradiant\nrefugees\nrichard\nromp\nrussian\nsandra\nsatisfyingly\nsavor\nsaying\nscare\nscott\nscreening\nshadows\nshoulders\nslightest\nspears\nstomach-turning\nstyles\nsublime\nsucked\nsun\nsundance\nsurprised\nswooping\nteens\ntense\ntight\ntool\ntosses\ntranscend\ntransforms\ntried\ntrust\nturf\nunassuming\nuncomfortably\nwallet\nwinners\nyearning\n3d\nacceptable\naccess\nacclaim\naffected\naliens\nangry\nappeals\nargentine\nasian\naspires\nassurance\nattal\naudacious\nawards\nbirthday\nboat\nbourne\ncage\ncelebrated\ncerebral\ncho\nclown\ncon\ncondescending\ncontinuity\ndeals\ndemonstrate\ndeserved\ndevelop\ndirectly\ndiscover\ndozen\ndragonfly\ndynamic\nego\neight\nencounter\nenvironmental\nextra\nfailed\nfeature-length\nferocity\nfirst-class\nfocused\nfollows\nforceful\ngang\ngel\ngrating\nhammy\nheights\nhey\nhilariously\nhoped\nhuppert\nidealistic\nidiosyncratic\nimmature\nimport\nimprove\nimprovisation\ninescapable\ninfuses\ninspires\nintensely\nirrelevant\njaglom\njapan\njell\njoyless\njudd\nknack\nkorea\nlarger-than-life\nlaugh-out-loud\nleading\nlegacy\nlesser\nlimp\nlit\nloaded\nlofty\nlosers\nmanhood\nmedium\nmental\nmind-numbingly\nmorvern\nmuddle\nmust-see\nnarratively\nneighborhood\nno-nonsense\nnonexistent\nobsession\nobsessive\noffend\noverbearing\noverexposed\nparticipants\npenetrating\npilot\nprime\nprocession\nprofoundly\npunching\npunk\npushes\nrapid-fire\nrecite\nrecognize\nreduced\nrejected\nremind\nremotely\nrolling\nsaga\nsalton\nscenic\nscience-fiction\nsecondhand\nsenses\nsinister\nsleepless\nsmack\nspecimen\nspontaneity\nsteve\nstrictly\nstrongest\nstrongly\nstudent\nstylized\nsuffer\nsuggests\nsystem\ntarantino\nteeming\nthematically\nthrowing\nthrown\ntongue-in-cheek\ntraditions\ntriangle\ntsai\nundoubtedly\nuneasy\nuneventful\nunforgettable\nunimaginable\nunless\nunsuspecting\nventure\nvictim\nwake\nwhale\nwins\nyiddish\nzero\n100\nachieve\nadrenaline\nafter-school\naggressively\nai\nalike\nallowed\nanthony\narbitrary\nartistes\nassembled\nattempted\nbard\nbarrel\nbasis\nbilked\nblockbusters\nboasting\nboldly\nborders\nbow\nbreadth\nbruckheimer\ncaine\ncallar\ncanny\ncenter\nchiller\ncoffee\ncollectively\ncollision\ncomposed\ncomposition\ncompromise\ncontest\nconvictions\ncrudup\ncultures\ncylinders\ndamning\ndance\ndecade\ndefinitive\ndetached\ndisappoint\ndong\ndreck\ndressed\ndrowned\neffortlessly\nempire\nentry\nexcited\nexplores\nfaces\nfairy\nfavors\nfield\nflat-out\nfrequent\ngear\ngenres\ngift\ngreatness\ngrin\ngross\nguess\nharmless\nhartley\nhas-been\nheady\nhelped\nhenry\nhewitt\nhost\nhurt\nilk\ninclination\ningredients\ninner\nintellect\nintelligently\nintentionally\ninventiveness\njewish\njoel\njumble\nkick\nlaudable\nlaughed\nlifestyle\nlisten\nloathsome\nloosely\nlunch\nmaintain\nmaintaining\nmanipulation\nmatinee\nmedical\nmediocrity\nmiserable\nmisguided\nmonstrous\nmoronic\nnarrator\nnary\nnicholson\nnijinsky\nnormally\nnot-very-funny\nnotes\nocean\nopening\novercooked\npathos\npatient\npeak\nphysical\npicked\npitch\npractically\npretend\npride\nprogram\nrandall\nrealizing\nrelative\nroberto\nroman\nrouge\nsack\nscrooge\nsegment\nsend\nshakespeare\nsisterhood\nsleeve\nsmash\nsquarely\nstrict\nstruggling\nsubstitute\nsubtitles\nsubtlest\nsuperficiality\nsupernatural\nsurrender\ntackling\ntensions\nthousand\ntook\ntraffic\ntrap\nundone\nvein\nwal-mart\nwarfare\nwarrior\nwarriors\nwearing\nweighty\nwell-meaning\nwhenever\nwilson\nwisdom\nwishing\nzany\n1958\n4ever\n51\nabroad\naccurately\nadmittedly\naimlessly\nallegory\nangels\nastute\nathletes\natrocious\nauteuil\nawakening\nawe-inspiring\nbelievability\nbenefit\nbiopic\nblack-and-white\nblaxploitation\nbrash\nbrave\ncanvas\ncat\ncgi\ncharacterization\nchristopher\ncircuit\nclub\nco-writer\ncomplexities\nconcoction\nconsistent\nconspicuous\ncrack\ncraft\ndaniel\ndialog\ndrew\ndrivel\ndriven\ndrooling\nduration\ndustin\nduty\nelsewhere\nemperor\nemphasizes\nenjoying\nentertainingly\neuropean\nexcels\nflailing\nflashbacks\nfreddy\nfreedom\ngasp\ngenerosity\ngoo\ngrim\ngun\nheaded\nhear\nhidden\nhitchcockian\nhotel\nhumanistic\nhumanly\nhumans\nhypnotic\niconoclastic\nimmediately\nindividual\ninexplicably\ninoffensive\ninvaluable\nironies\nirreparable\njealousy\njolts\nlapses\nlegendary\nleon\nlifetime\nlikeable\nliterally\nlola\nlongest\nlord\nlow-grade\nluck\nlunatic\nmaddeningly\nmelancholy\nmexican\nmira\nmolly\nmugging\nmundane\nnickleby\norder\npeek\nphilosophy\nplympton\nporridge\nprecisely\npredictability\nprefer\npuzzling\nramsay\nreefs\nrelaxed\nreliable\nremaining\nrestroom\nrewarded\nridiculousness\nrigor\nripe\nroots\nrowling\nschneider\nseams\nsensuality\nshowcases\nsickening\nsilver\nskillfully\nsloppily\nsmooth\nsmoother\nspike\nsportsmen\nsprightly\nsqueeze\nstealing\nstick\nstone\nstooping\nstorylines\nsubtext\nsuccessor\nsuccumb\nsucks\nsuggest\nsuit\nsuperbly\nsupport\nswear\nswim\ntales\ntheories\nthree-hour\ntoss\ntowards\ntrace\ntransformation\ntransformed\nturgid\nundead\nundercut\nunderneath\nuniquely\nuniverse\nunlimited\nvietnam\nwanders\nweaknesses\nwillingness\nwondrously\n11\nabsurdities\nadequate\nadmission\naffleck\nalert\nall-time\namuse\namy\nanatomical\nanswer\narrest\narrives\naudacity\nbackstage\nbaffled\nbands\nbeats\nbible\nbluster\nbogus\nbonus\nboobs\nbravery\nbreathes\nbride\nbrisk\nbrosnan\nbullets\nbuttons\nbuñuel\ncallow\nchaos\nchicago\nchick\nchosen\nclooney\ncolored\ncolors\ncoma\nconduct\nconjured\ncontrivance\ncram\ncrammed\ncranky\ncraven\ncreepiness\ncrisis\ncrossover\ndead-end\ndelectable\ndepressed\nderived\nderrida\ndevolves\ndiesel\ndiscipline\ndistanced\ndistraction\ndramatics\ndude\nechoes\nembrace\nembraced\nended\nendurance\nenigma\nenjoys\nestablished\nexcruciatingly\nextravaganza\nfabric\nfabulous\nfarcical\nfearless\nfeast\nfellow\nfewer\nfizz\nfollowed\nfool\nfulfill\nfury\ngender\ngerman\ngibson\nglorified\nglum\ngolden\ngracefully\ngrade\ngrew\ngrounded\nguide\nhates\nhopeless\nhopelessly\nhour-and-a-half\nhumming\nian\nicon\nidiocy\nillustrates\nimproved\ninimitable\ninsecure\ninsults\nintroduction\nisolation\njaw-dropping\njournalism\nkapur\nkilled\nknockout\nlavish\nlight-hearted\nlips\nlittered\nlock\nlodging\nlucky\nlurches\nmanaged\nmanners\nmarvel\nmasses\nmaze\nmel\nmetaphors\nmoralizing\nmountain\nmovie-going\nnames\nnash\nnaturalistic\nnotwithstanding\nnumbing\no\none-dimensional\noverused\npainted\npalette\npalpable\npar\npaxton\npedigree\nperformed\nperfunctory\nplotline\npraise\npraises\nprospect\nracist\nrandomness\nrange\nrapes\nreached\nreaching\nreacting\nready\nrecklessness\nrecycling\nrelentless\nreminder\nrescue\nresistance\nrevealed\nrevel\nrevelations\nrobin\nromances\nrotten\nsardonic\nscares\nsearching\nseas\nseeks\nshabby\nship\nshowcase\nskies\nskillful\nsoldiers\nsounding\nsources\nspaces\nspeaks\nspirituality\nspotlight\nstarting\nstarving\nstuart\nsubmarine\nsuitable\nswept\nthirty\nthoughtfulness\nthrow\ntimid\ntoback\ntranscendent\ntransparent\ntwo-dimensional\nurine\nvain\nvietnamese\nwell-intentioned\nwitnessed\nworry\nx.\nyesterday\nyouthful\n12-year-old\n13\nadams\nadding\nadequately\nafternoon\nanalysis\nanguish\nanime\nantonia\nappreciated\naspirations\nassayas\natmospheric\nbackdrops\nbackmasking\nbank\nbargain-basement\nbarry\nbeguiling\nbegun\nbucks\nburger\ncanon\ncapra\ncarrying\ncarved\nchateau\ncheeky\nchew\ncia\nclad\nclause\nclunky\nco-written\ncomics\ncompared\ncomplaint\nconcentrates\nconceptual\nconsolation\ncrafty\ncreations\ncrippled\ncumbersome\ncurves\ncuts\ndana\ndarkness\ndebate\ndeception\ndepraved\ndime-store\ndimension\ndisplays\ndistinct\ndistinguish\ndisturbing\nditsy\ndoofus\ndopey\ndozens\ndubious\nduke\ndupe\ndutiful\ndynamics\nedits\neffortless\nelizabeth\nengagingly\nenigmatic\nenlightening\nenticing\nerrors\nessential\nevade\nexamines\nexpecting\nexperienced\nexperiment\nexploit\nextravagant\nfade\nfascination\nfixating\nflawless\nfloor\nfluffy\nfolk\nformal\nforth\nfull-bodied\ngenius\ngrasp\ngrisly\ngrotesque\ngut-wrenching\nhall\nheat\nheightened\nhighs\nhill\nhistrionics\nhokey\nhot-button\ninappropriate\ninauthentic\nincorporates\nindescribably\ninexplicable\ninfectious\ninject\ninsensitivity\ninteraction\ninterestingly\nirresistibly\njackass\njagged\njake\njolly\njuliette\njumbo\nlanes\nlayer\nleather\nlegal\nlegged\nlevity\nlonely\nmad\nmajestic\nmakeup\nmaternal\nmechanical\nmelancholic\nmerits\nmidway\nmodesty\nmonths\nmysteries\nnair\nnegative\nnostalgia\nnoticeably\none-liners\nordeal\noutcome\noverrun\novertly\npackage\npaint-by-number\npandering\npastiche\npayami\npedestrian\nperform\nphone\npiercing\nplacement\npollution\nportray\npredicament\nprocedural\nproductions\npulse\npush\nraises\nreceived\nrelate\nreluctant\nrepellent\nrequisite\nresearch\nresonate\nrings\nritchie\nrival\nrobinson\nrocky\nsara\nschmaltz\nschmidt\nscreens\nscrewball\nself-caricature\nself-esteem\nsentence\nsequel-for-the-sake\nserenity\nsessions\nset-up\nshared\nshower\nshred\nshtick\nsignificantly\nsmiles\nsmugly\nso-called\nsober\nsorority\nsoufflé\nspice\nspontaneous\nsteers\nsticky\nstinks\nstoop\nstrangers\nstrengths\nstunningly\nsublimely\nsudsy\nsuspect\nsymbolic\ntart\ntattered\nteachers\ntechnique\ntempting\nterror\ntown\ntrek\nturd\nunbearably\nundercurrent\nunmemorable\nunsatisfying\nunwatchable\nurgently\nvan\nvastly\nvictimized\nvideotape\nweep\n1999\n40\n86\nabsurdly\naccomplish\nachievements\nachieving\nacute\nafterthought\nagent\nambiguity\nangel\nantics\napplies\narguments\natop\naustin\nball\nbanality\nben\nbile\nbinoche\nbiography\nbrainless\nbrains\nbrits\nbrushes\nbuild\nbump\nbuzz\ncalibrated\ncardboard\nchanneling\nchasing\ncleaner\nclear-eyed\nclimax\nclung-to\ncohesion\ncomprehension\ncomputer-generated\nconquer\nconversations\ncooper\ncostner\ncovers\ndamaged\ndarling\ndead-on\ndean\ndeliberately\ndepression\ndeserving\ndestined\ndirections\ndisquieting\ndowntown\ndownward\ndust\nearn\neastern\neducation\nencounters\nenervating\nepisodic\nevoking\nexperimental\nexpertly\nexplode\nexploitive\nexploring\nfacile\nfalling\nfascinate\nfast-paced\nfeet\nfest\nfondness\nfontaine\nforcefulness\nfun-for-fun\nguilt\ngutless\nguts\nhampered\nharris\nheart-pounding\nhitchcock\nhopes\nhypocritical\ninsignificance\ninstance\ninstitution\ninterpretation\ninterviews\njaded\njolie\nlaborious\nlobby\nloveable\nlows\nmadness\nmcdormand\nmechanics\nmember\nmerchant\nmercy\nmetropolitan\nmisanthropic\nmoods\nnamed\nnebrida\nnelson\nnourishing\nnumerous\none-of-a-kind\noozing\nopened\nopera-ish\noptimistic\norlando\npainting\nparent-child\npasses\npassing\npeculiarly\npeppered\nperils\npervasive\npileup\npitched\npompous\npossibility\npratfalls\npre-credit\nproficient\npulp\nradar\nrarity\nraunchy\nreasons\nrecognized\nregardless\nregion\nrequire\nreserved\nretaliation\nrhetoric\nrichness\nrubbo\nsanctimony\nscruffy\nsealed\nself-examination\nsensuous\nsermon\nsetups\nshape-shifting\nshapes\nshifting\nshine\nshrek\nsimpleminded\nsin\nslackers\nsleeper\nslippery\nsmartest\nsmirk\nsmoochy\nsnore\nsoccer\nspan\nspins\nsplit\nsquare\nsteady\nsteeped\nstepped\nsticks\nstill-inestimable\nsting\nstretched\nstroke\nsuffocating\nsuper\nsurefire\nsurfing\nsway\ntawdry\nthinly\nthinness\nthreadbare\nthree-dimensional\nticking\ntop-billed\ntraditionally\ntripe\ntug\ntuxedo\ntwisty\ntwohy\nugliness\nunbearable\nups\nvulnerable\nwalsh\nwastes\nwaters\nweave\nweepy\nwhat-if\nwitness\nworthless\nwound\n19th-century\nabomination\nabundant\nabyss\nactuary\nadorably\naftertaste\naids\nalfred\nall-enveloping\namateur\namorality\nanticipated\naplenty\narty\nattached\nattack\nawarded\naware\nbaby\nbang\nbarrymore\nbesson\nbetrayal\nblah\nbloated\nbogdanovich\nbogged\nbolt\nborn\nbrush\nburning\nburr\nbutt\ncandid\ncannes\ncarmen\ncarpenter\ncheaper\ncheated\nchekhov\nchelsea\nchiefly\nchoreographed\ncinderella\nclaim\ncliché-riddled\ncluelessness\ncoen\ncollapse\ncollateral\ncollective\ncompanionable\ncompetition\ncompulsively\nconnect-the-dots\nconsidering\ncrackle\ncrew\ncrowd-pleasing\ncutesy\ncutting-edge\ndad\ndeliberate\ndelinquent\ndeparture\ndesiccated\ndevastatingly\ndevelops\ndiaries\ndilemma\ndinner\ndiplomacy\ndips\ndisappointments\ndive\ndonovan\ndough\ndover\ndrains\ndriver\ndrowsy\ndual\ndullard\ndv\near\neducates\neerily\nefficient\neric\nethereal\nexpects\nexperimentation\nextremes\nfessenden\nfiercely\nflabby\nflash\nflesh\nforcefully\nforest\nformed\nfourth\nfrantic\nfresh-faced\nfreshly\ngasping\ngere\nglib\nglimpses\ngoofball\ngrinder\nhandsomely\nhaphazard\nhardened\nheartening\nheist\nhelm\nheroic\nhide\nhitler\nhole-ridden\nhooting\nhorrors\nice-t\nidemoto\nill-conceived\nill-fitting\nimaginatively\nimpersonation\nimproves\ninevitably\ninfomercial\ninformation\ninner-city\ninsightfully\nintact\ninterview\nintroducing\ninvolve\nkicks\nkilling\nkindness\nkosashvili\nlark\nlearning\nlibrarian\nlimits\nloyal\nlucy\nlulls\nlumbering\nm.\nmarching\nmasculine\nmeets\nmessiness\nmetaphorical\nmill\nmiller\nminute\nmisses\nmiyazaki\nmonumental\nnaughty\nneatly\nneil\nnon-stop\nnonstop\nnoticing\nnovels\nnuances\nobstacles\nodyssey\nofferings\noperates\nouting\noutsiders\noverproduced\noverwritten\npasse\npassions\npayne\npg-13\npinocchio\nplainly\npleasantly\npop-induced\npossess\npost\npotshots\npressed\nprevents\npropaganda\nprose\npushed\nqualify\nquestionable\nquivering\nr.\nrap\nrattling\nreaction\nrediscover\nredundancy\nrefused\nregard\nreleases\nreminded\nreruns\nrestrictive\nreveal\nriver\nrock-solid\nrollerball\nroot\nroutines\nrugrats\nrumination\nrut\nsatisfactory\nsaved\nscarcely\nscarier\nself-glorification\nsensationalism\nseptic\nshafer\nshaggy\nshape\nsoar\nsoars\nsouthern\nspader\nspecifically\nspeculation\nspeedy\nspellbinding\nspookiness\nstagy\nstaring\nstoked\nstrays\nstrokes\nstronger\nstruggles\nstupidity\nsubstances\nsucker\nsummertime\nsurfeit\nsurvive\nswan\nsweat\ntacky\ntactics\ntank\ntatters\nterminally\ntestimony\nthornberrys\ntitular\ntolerance\ntranscendence\ntranslate\ntraveler\ntrenchant\ntroubled\ntwaddle\ntwisting\nuh\nunamusing\nuncle\nunconditional\nuniformly\nunknown\nunpretentious\nunrelentingly\nvisible\nwan\nwander\nwaves\nweaving\nwelsh\nwendigo\nwesley\nwestern\nwhimper\nwhodunit\nwicked\nwilco\nwilliam\nwood\nwriters\nyu\n#\n1.8\n1970s\n21st\n6\n8\naccepting\naesthetic\naged\nallegedly\nallowing\nanachronistic\nanemic\nanimatronic\nanne\napparatus\nararat\nargue\narresting\narticulate\nassault\nassociated\nassumes\nastray\natmospherics\natrociously\naustralia\nawry\nbad-movie\nbankrupt\nbarrels\nbasketball\nbeacon\nbehave\nbielinsky\nbirot\nblanket\nblip\nbluescreen\nblurry\nboisterous\nbonds\nboozy\nborrowed\nborrows\nbouncy\nbowling\nbrawny\nbrief\nbroiling\nbrutality\nbudding\nbumbling\ncackles\ncalifornia\ncaptain\ncar-wreck\ncarefully\ncathartic\nchampion\nchaplin\ncheat\ncheesiest\nchief\nchopsocky\nchuckles\nchurn\ncodswallop\ncombat\ncomically\ncompanion\ncomplacency\ncomprehend\nconceit\nconception\ncondensed\nconfined\nconnect\nconned\nconvention\ncorniest\ncounterparts\ncourageous\ncriticism\ncrocodile\ncuban\ncynicism\ndanger\ndaughters\ndearth\ndecrepit\ndedication\ndeepest\ndefinition\ndeliciously\ndenied\ndenying\ndepths\ndetermine\ndiary\ndiscernible\ndiscovering\ndisgust\ndistracted\ndragons\ndredge\nedit\nefficiently\negregious\nelectric\nelephant\nelevates\nembracing\nenchanting\nengages\nenhances\nenters\nestablishes\nethical\nexamine\nexpressing\nexpressions\neyelids\nfabricated\nfantasma\nfellowship\nflatly\nflaw\nfondly\nfootnote\nford\nforgiven\nforgiveness\nfoul-mouthed\nfragmented\nfreighter\nfrenetic\nfused\nfuture-world\ngang-raped\ngimmicky\ngood-time\ngoodfellas\ngosford\ngratefully\ngrind\ngrown-ups\ngrownups\nguillen\ngunfight\ngutter\nhamming\nhawaiian\nhead-trip\nhealthy\nheart-wrenching\nherself\nhint\nholographic\nhopefully\nhorse\nhostile\nhumility\nhurry\nhysterics\nignore\nill\nill-considered\nill-timed\nimmediacy\nimmersed\nimpossibly\nimpulsive\ninarticulate\nincarnation\nindecent\nindians\nindie-heads\ninfatuated\ninfluenced\ninjects\nintricately\njaw-droppingly\njerking\njoie\njonah\njoys\nkalvert\nkids-and-family-oriented\nkinetic\nkinky\nkitchen\nlasting\nlaughably\nlaurice\nlengths\nlengthy\nlewis\nliberties\nlights\nlisping\nliteral\nlogical\nloopy\nlose\nlowbrow\nluckiest\nluridly\nmacabre\nmalkovich\nmamet\nmanic\nmedia\nmendes\nmentally\nmired\nmolested\nmorgan\nmother/daughter\nmtv\nmyth\nna\nnarration\nnatter\nneedlessly\nninth\nno-frills\nnonetheless\nnovak\nnoyce\nnumbingly\nnutty\nof-a-sequel\noliver\nought\nout-bad-act\noverdoing\noverflowing\noverladen\noverly-familiar\nowes\npairing\npaperbacks\nparanoid\npartly\nperdition\nperfection\npersonified\nphrase\npill\npink\npinnacle\npizza\nplayer\npreemptive\nprimitive\nprior\nprism\nprofanity\npuberty\npunny\nrapt\nraunch\nreactionary\nreceiving\nrecognizable\nredeeming\nrehashes\nrelic\nreminding\nrepeatedly\nrepetition\nrepresented\nreptilian\nresembling\nretains\nrevisionist\nriot\nritter\nroberts\nrotoscope\ns/m\nsag\nsalt\nsatisfy\nsavage\nsavagely\nsavory\nscale\nschaeffer\nschedule\nschumacher\nscreenings\nseattle\nsection\nseduce\nself-indulgence\nsemen\nseriousness\nsexually\nshadowy\nshop\nshortcomings\nshorter\nsillier\nsits\nsledgehammer\nsmoking\nsnl\nsolemn\nsoon-to-be-forgettable\nspider\nspiked\nspine\nsporadic\nsport\nsprung\nsquad\nstark\nstephen\nstopped\nstream\nstuffiest\nsub\nsubordinate\nsubtler\nsucceeded\nsuccessfully\nsupply\nsurvivors\nsuspects\nswooning\nt\ntapping\ntask\ntastelessness\ntea\ntelegraphed\ntendency\nthomas\ntissues\ntoilet-humor\ntolerate\ntravesty\ntrue-to-life\ntuned\nun-bear-able\nunable\nunappealing\nunderdone\nunderlying\nunimaginatively\nuninvolving\nunleashes\nunnoticed\nupdating\nupscale\nvibe\nviewed\nvignettes\nviolently\nvoice-over\nwasting\nwayward\nwell-shaped\nwest\nwhatsoever\nwhining\nwitnesses\nwork-in-progress\nwoven\nxerox\nxxx\nyouth\nzippy\n'60s\n1984\n2000\n3\n8-year-old\nabc\naccomplishes\naccomplishment\nactions\nadrenalin\nafrican\naisle\nalice\nalienate\nall-too-familiar\nanyway\nanywhere\napply\narduous\nareas\nargument\narms\nastoundingly\naudiard\nausterity\nauthentically\navarice\nbadness\nbalancing\nballistic\nbaroque\nbear\nbeaten\nbeliever\nbeloved\nbiased\nblasphemous\nbliss\nblonde\nbloodsucker\nbollywood\nbones\nbooks\nboomer\nbothers\nboundaries\nbowel\nbubba\ncampy\ncaptivates\ncareer-best\ncareers\ncartoons\ncatalyst\ncelebrity\ncharged\ncharles\nchasm\ncheering\nchurch\ncircle\nclaptrap\nclean\ncliche-ridden\nclimate\ncocktail\ncode\ncomments\nconceive\nconditioning\nconflicts\nconsigned\nconstraints\ncontemplation\ncontinuation\ncontrast\ncontroversial\ncops\ncounter-cultural\ncredulous\ncreek\ncreeping\ncruelty\ncrypt\ndamon\ndated\ndaydreams\ndecided\ndeeds\ndefine\ndepict\ndestiny\ndestroy\ndestruction\ndey\ndickensian\ndirect\ndisabilities\ndisintegrating\ndisorientated\ndon\ndoor\ndragged\ndreamed\ndroning\ndullness\ndurable\nearned\nebullient\nechelons\nedition\nedward\neffecting\nelbowed\nelevated\nemploys\nenchantment\nencourage\nenriched\nenvironment\netc.\nevaluate\nevenings\neverlasting\neveryday\nevolved\nexists\nexpands\nexplosions\nexuberance\nfairy-tale\nfashioned\nfees\nfeisty\nferrara\nfiascos\nfingered\nfires\nflagrantly\nflatter\nflickering\nflight\nflorid\nfoibles\nfolly\nforcing\nfriendships\nfuriously\nfuss\nfutile\ngaping\ngaze\ngenteel\ngeriatric\ngoers\ngoldmember\ngong\ngoods\ngreene\nguarantee\ngusto\nhalf-baked\nharder\nheart-warming\nheartache\nheritage\nhigh-wattage\nho-tep\nhomicide\nhonor\nhook\nhuman-scale\nhumanism\nimpart\nimportance\nimpress\nimproperly\ninclude\ninexperienced\ninfantile\ninfusion\ningenuity\ninherently\ninnovations\ninquisitive\nintent\nintrospective\ninuit\ninvasion\nironically\nissue\njagger\njet\nk-19\nkeen\nkieslowski\nkittenish\nknee-jerk\nlampoon\nlandscapes\nlate-twenty-somethings\nlaw\nlegally\nlegend\nletdown\nliability\nlift\nlikes\nlip-non-synching\nlittle-known\nlosses\nlustrous\nmajor-league\nmatthew\nmaturity\nmawkish\nmessenger\nmessy\nmetropolis\nmichel\nmine\nmisplaced\nmolestation\nmonsoon\nmorbidity\nmotions\nmotivations\nnails\nnaked\nnarc\nnavigate\nneck\nnerves\nnudity\noedekerk\noh-so\nopportunities\norgasm\norwell\noutdated\nover-blown\npacino\nparking\npencil\npickup\nplucking\nplus\npotato\npractice\nprecise\nprejudice\npreliminary\nprep-school\npress\nprevails\npreview\nproducers\npromisingly\nproposal\nproving\nprovoke\npseudo-educational\npsycho\npuff\npulled\npurposeless\npurposes\npushiness\nputrid\nquentin\nquirkily\nray\nreal-time\nrealistically\nrealm\nreassuring\nreceive\nredolent\nreflection\nreflective\nrefreshed\nrefreshes\nregalia\nreggio\nreinforcement\nrelegated\nreliance\nrepulsive\nretro\nrevenge-of-the-nerds\nreverent\nrevisiting\nriveted\nrobotically\nrobust\nrocks\nromancer\nsanctimonious\nsays\nscarifying\nschmaltzy\nscientists\nscrewed-up\nsearing\nsecond-guess\nseductive\nself-destructive\nsensation\nserials\nserving\nshadyac\nshakes\nsham\nsharpener\nsheerly\nshorts\nshout\nshowdown\nshowgirls\nsibling\nsignposts\nsing\nskewed\nslanted\nslap-happy\nslim\nsmacks\nsneak\nsneering\nsonny\nsorrow\nspiced\nspiffy\nspiteful\nsquandering\nstaggeringly\nstanley\nstitch\nstoops\nstops\nstrains\nstretch\nstubbornly\nstudios\nsuitably\nsummer-camp\nsurf\nsurvivable\ntaps\ntear-jerking\nteeth-clenching\ntheatre\ntheory\nthreatened\nthrown-together\ntosca\ntriteness\ntrope\ntropes\ntunney\nturkey\nunapologetic\nunder-inspired\nunderappreciated\nunderscore\nundisciplined\nunexplainable\nunintentional\nunrealized\nvaliantly\nvictories\nvolume\nvulnerability\nwalker\nwalls\nwater-born\nwaydowntown\nwinking\nwonderous\nword-of-mouth\nworn-out\nwrestling\nwryly\nyawning\nyoungsters\nzings\n'em\n146\n18-year-old\n80\naccumulates\naccuracy\nace\nacknowledges\nacumen\nadjective\nadmirably\naffirms\naficionados\naggressive\nalienating\nanchoring\nancient\nannals\nannoyed\nanomie\nantic\nantitrust\nappearing\nappointed\narray\nash\nashley\nassume\nastonish\nauditorium\nautobiographical\navailable\navant-garde\nbalances\nbanter\nbaran\nbarn-burningly\nbenefits\nbewitched\nblithe\nblood-curdling\nbodily\nbores\nbrainpower\nbravado\nbucked\nburrito\ncalls\ncamaraderie\ncapability\ncaptivated\ncard\ncatching\ncatholic\nchair\nchases\ncheapo\nchoosing\ncircus\nclash\ncleverest\nclosely\ncloyingly\nco-winner\ncompanionship\nconclusions\nconcocted\nconfidently\nconfront\nconfrontations\nconvoluted\ncorrupt\ncorruption\ncounts\ncourt\ncraftsmen\ncreatures\ncredited\ncrowd-pleaser\ncultivation\ncup\ncure\ncurse\ndash\ndead-undead\ndealer\ndear\ndecibel\ndefiance\ndeja\ndemand\ndemme\ndemonstrating\ndenis\nderailed\ndescribes\ndevotees\ndidactic\ndifferences\ndigest\ndimensions\ndimness\ndiscordant\ndiscussion\ndismiss\ndistinctly\ndisturbance\nditty\ndoing-it-for\ndope\ndrab\ndumped\nearlier\nearthly\neast\neating\necks\nelbows\neloquence\nelusive\nemployment\nencompassing\nenthusiasts\nentity\nenveloping\neven-flowing\neverybody\neverywhere\nevidence\nevolve\nexceeds\nexcites\nexecutive\nexist\nexuberant\nfacet\nfamily-friendly\nfascinates\nfatter\nfeeble\nfeelings\nfi\nfiend\nfilmgoing\nfishy\nflavorful\nflip-flop\nfluent\nfluke\nfoot\nfoul-natured\nfraction\nfreddie\nfrightfest\nfrissons\nfrontal\nfulfilling\ngary\ngeared\ngenerations\ngenial-rogue\nglass-shattering\nglimmer\nglow\ngoddammit\ngondry\ngoofiest\ngoofiness\ngraphics\ngratingly\ngrave\ngrueling\nguilty-pleasure\nguy-in-a-dress\nhalf-asleep\nhalf-hearted\nhallmarks\nhaneke\nharbor\nhawk\nhayek\nhelpings\nhigh-minded\nhighlight\nhindered\nhoary\nhodgepodge\nholding\nhungry-man\nhunk\nidentification\nillness\nimpassive\nimprobability\nin-depth\nindigestion\nineptly\ninjuries\ninstallment\nintegrated\nintro\ninvited\nirksome\nirrelevancy\nisrael\nitalicized\nj.k.\njane\njealous\njessica\njfk\njoe\njr.\njuice\njulia\njunkie\nkissing\nknock\nknockabout\nlaid\nlandmark\nlanguid\nlanguishing\nlaugh-a-minute\nleash\nleavened\nleftovers\nleguizamo\nlibertine\nlight-footed\nlilo\nlistening\nloser\nlow-rent\nlower\nluster\nlyricism\nmaelstrom\nmaneuvers\nmanifesto\nmann\nmannerisms\nmarkets\nmarveling\nmcgrath\nmemorial\nmeticulous\nmilestones\nmilquetoast\nminority\nmiraculous\nmiramax\nmiscalculation\nmission\nmist\nmistress\nmisuse\nmodern-day\nmonday\nmorph\nmultiplex\nnegligible\nnegotiate\nnet\nnewfangled\nnickelodeon-esque\nnorton\nnot-so-funny\nnoticeable\nnæss\nobject\nobjectivity\noily\nolivier\non-screen\none-trick\nopen-ended\norgy\nours\nout-to-change-the-world\nouter-space\noutright\noverwhelmed\noverworked\npadded\npains\npainterly\npassable\npassed\npawn\npearl\npeas\npenchant\npermeates\npetter\nphoenix\nphoney-feeling\nphoniness\npile-ups\npipeline\npony\npopularity\nportions\npower-lunchers\nprim\nprincipal\nprinciples\nprint\nprivy\nprobing\nproceeds\npubescent\npuddle\npumpkin\npunishment\npurposeful\npuzzle\nquandaries\nqueens\nquieter\nraindrop\nrally\nrant\nrate\nraucously\nrazzle-dazzle\nreadily\nreconstruction\nreflects\nrelying\nreno\nreprieve\nrerun\nresponsibility\nretitle\nrevitalize\nridicule\nrigid\nringing\nroughshod\nround\nrousing\nsadistic\nsan\nsassy\nschizophrenia\nsci\nscooped\nscooping\nscorpion\nselection\nself-aggrandizing\nself-aware\nself-awareness\nself-congratulation\nself-hatred\nself-righteousness\nsentiments\nserviceable\nsever\nsewage\nshakespearean\nshe-cute\nshown\nshrieky\nshudder\nsiegel\nsimple-minded\nsizzle\nskills\nslapping\nso-so\nsoft-core\nsolaris\nsorrowful\nsought\nspeed\nsquashed\nstages\nstalls\nstamp\nstand-up\nstates\nstellar\nstevens\nstifles\nstirs\nstores\nstrained\nstrangeness\nstreets\nstuffing\nstuffy\nstylists\nsubjugate\nsubzero\nsuccumbing\nsucking\nsuffered\nsuggested\nsung\nsuperhuman\nswedish\nsweetheart\nsymbolism\ntackles\ntagline\ntall\ntame\ntattoo\ntaut\ntavernier\ntax\nteach\ntearjerker\nteen-driven\nterrorism\nterrorists\ntestud\nthe-cash\nthreat\nthunderstorms\ntian\ntightrope\ntimelessness\ntoes\ntopless\ntopple\ntract\ntremble\ntwo-day\ntwo-hour-and-fifteen-minute\nunadorned\nunblinking\nuncreative\nundermining\nunfilmable\nunifying\nunited\nunmentionable\nunmistakable\nunrecoverable\nunsubtle\nuptight\nurgency\nurgent\nvacant\nvaguely\nvaliant\nvariation\nveers\nveracity\nvia\nvisionary\nvisualize\nvivre\nvon\nvu\nwaking\nwashed\nwaterlogged\nweary\nwednesday\nwell-balanced\nwell-deserved\nwell-developed\nwell-done\nwherever\nwide-smiling\nwig\nwildlife\nwinningly\nwore\nworking-class\nworkplace\nworld-renowned\nzingers\nzombie\n'30s\n'40s\n+\n22-year-old\n25\n3000\n4\n5\n84\n91-minute\nabove-average\nabsurdist\naccepts\naccident\naccused\naction-adventure\nadorable\nadvised\naesthetically\naffable\nagenda\nages-old\naggressiveness\naggrieved\nah\nairless\nalbeit\nall-star\nallusions\namateurishly\namazement\namish\namnesiac\namusement\nanciently\nangle\nanguished\nanonymous\nanymore\napocalypse\nappreciation\nargento\narmed\nartfully\nasia\nassassination\nastronomically\nawake\nawed\nbabak\nbackgrounds\nballplayer\nbanderas\nbartlett\nbathroom\nbest-sustained\nbiographical\nbludgeoning\nboilerplate\nbolero\nboost\nboots\nboundary-hopping\nbourgeois\nbratty\nbrim\nbrio\nbronx\nbug-eyed\nburst\nby-the-numbers\ncal\ncalculated\ncalculating\ncampus\ncardiac\ncaretakers\ncaricatures\ncarlin\ncastro\ncelebratory\nchallenge\nchallenges\ncharacteristically\ncharade\ncheapened\ncheatfully\nchecking\ncheer\ncheesier\ncherish\nchick-flicks\nchimes\nchump\ncleavage\ncletis\ncomedy/drama\ncommercialism\ncommunicates\ncompassionate\ncompetence\ncomplexly\nconcepts\nconditions\nconnected\nconstructs\nconsumed\ncontemptible\ncoolness\ncopout\ncostumes\ncotton\ncoupled\ncourageousness\ncraig\ncrane\ncrawl\ncreeps\ncricket\ncrossing-over\ncrossroads\ncrummy\ncrummy-looking\ncrushingly\ncuisine\ncurmudgeon\ndamaged-goods\ndamme\ndamon/bourne\ndangerously\ndaughter\ndavies\ndaytime\ndaytime-drama\ndebated\ndebt\ndecide\ndecision\ndefiantly\ndemise\ndemographic\ndenzel\ndetract\ndevastated\ndevelopments\ndevice\ndiabolical\ndiane\ndiner\ndirectionless\ndisastrous\ndisconnects\ndiscreet\ndisdain\ndismay\nditched\ndiversity\ndolls\ndoubts\ndown-to-earth\ndownhill\ndownsizing\ndraw\ndread\ndreaming\ndreamlike\ndrips\ndrowns\nduck\neagerness\nearplugs\necological\neconomical\necstasy\nedges\nelling\nenables\nencompasses\nensure\nepiphany\nerotically\nerror\nescort\nespionage\nestranged\nethics\nevaporates\nexasperating\nexhausted\nexpiration\nexpresses\nextent\nextreme-sports\neyre\nfacetious\nfactors\nfaltering\nfamiliarity\nfantasia\nfar-flung\nfast-moving\nfirst-timer\nflesh-and-blood\nfleshed-out\nflinching\nflippant\nfluidity\nfore\nforemost\nforgets\nforum\nfour-hour\nfreakshow\nfright\nfumbled\nfunction\nfundamentally\nfunniness\ngabbiest\ngaining\ngiant-screen\ngirlfriend\ngoodly\ngraduated\ngrandkids\ngrayish\ngremlins\ngrieving\ngriffiths\ngrips\ngrown\ngrumble\nguessed\nguffaw\ngullets\nhalf-step\nhandle\nhandling\nhard-hitting\nhaynes\nheart-felt\nheavily\nhennings\nhigh-octane\nhitch\nhorrendously\nhour-and-a-half-long\nhousehold\nhouses\nhug\nhughes\nhumanizing\nhumankind\nhungry\nhusband-and-wife\nignorant\niii\nills\nimpulses\ninadvertent\nincident\nincorporate\nindication\ninduce\ninhale\ninitial\ninsignificant\ninstinct\ninstructive\ninsufferable\ninterrogation\nintroduce\ninvestigate\nirresponsible\nirwin\nitalicizes\niwai\njacques\njean-claud\njesse\njoan\nkennedy\nkinetically-charged\nkosminsky\nkurt\nlabute\nladder\nlags\nlake\nlamer\nlan\nlarge-screen\nlarry\nlatin\nlearned\nliar\nlicense\nlife-changing\nlightness\nlilia\nlobotomy\nlong-faced\nlong-on-the-shelf\nlouis\nlucratively\nlumpen\nmacdowell\nmagnetic\nmakhmalbaf\nmall\nmargin\nmarquis\nmarshall\nmeal\nmeanderings\nmeara\nmenace\nmergers\nmid\nmiddling\nminded\nmormon\nmuddy\nmultilayered\nmumbo\nmutates\nmuzak\nmystical\nmétier\nnatured\nnaïveté\nnervy\nniches\nnights\nno-brainer\nnoon\nnoteworthy\nnumb\nobnoxiously\nobsessions\nodorous\noften-funny\nominous\non-camera\none-star\nops\norchestrated\noscar-winning\noscar-worthy\noverblown\novercomes\noverplayed\noverview\nozpetek\npact\npage\npanoramic\npaths\npersuasive\nphilip\nphlegmatic\nphoto\npillages\npitiful\nplacid\nplacing\nplane\nplateau\npleas\npleasuring\nploddingly\nplunging\npoet\npopulation\npregnant\nprevent\nprincipals\nprivileged\nproficiency\npropels\npsychotic\npub\npump\npunchy\npurer\nquick-witted\nrea\nreaches\nreasonable\nrecall\nreception\nreckless\nrecords\nredone\nreduces\nregain\nrelevance\nrembrandt\nrepeated\nresident\nretooling\nrevulsion\nripening\nrob\nrocket\nromanek\nrough-hewn\nrunteldat\nrustic\nsabotage\nsamira\nsatin\nsatisfied\nscandalous\nscathing\nschemes\nschlocky\nscotland\nscottish\nscreams\nscreenplays\nseamless\nseamy\nsedate\nseemed\nseeming\nsegal\nself-assured\nself-empowering\nself-preservation\nsexist\nsexuality\nshadow\nshainberg\nshaking\nshapable\nsharper\nsharply\nsheridan\nshockwaves\nshoe\nshoplifts\nshovel\nshu\nsignals\nsimulate\nsiuation\nsketched\nskins\nslob\nslog\nslovenly\nsmart-aleck\nsmarts\nsmeary\nsmith\nsmoothly\nsocially\nsoothing\nsoul-searching\nspecies\nspreads\nsprinkled\nsquareness\nstare\nsteals\nstewart\nstillborn\nstolid\nstomach-knotting\nstrings\nstud\nstyled\nsub-sophomoric\nsubgenre\nsubjected\nsuddenly\nsuperfluous\nsupermarket\nsurroundings\nsuspend\nsustenance\ntabloids\ntangled\ntargeted\ntechnological\nterrorist\ntheatrically\nthrowback\nthumbing\ntickets\ntickles\ntier\ntime-consuming\ntinny\ntitled\ntooth\ntoothless\ntopical\ntransfigures\ntransform\ntransgressive\ntravail\ntreasured\ntries-so-hard-to-be-cool\ntruck\ntumbleweeds\ntunes\nturbulent\ntwin\ntwitchy\ntwo-thirds\ntykwer\nunattractive\nundergrad\nunderrated\nundo\nunfakable\nungainly\nunhappy\nunholy\nuninhibited\nunknowable\nunmotivated\nunsurprising\nunturned\nupdate\nvacation\nveered\nveiling\nveneer\nverges\nversus\nveteran\nvicious\nvidgame\nvirulent\nvisions\nvolcano\nvowing\nwanderers\nwang\nwannabe-hip\nweaponry\nweasels\nwell-executed\nwell-lensed\nwell-thought\nwildean\nwindtalkers\nwiseman\nwistful\nwithered\nworm\nwounds\nwrenching\nyard\n10-year\n103-minute\n105\n110\n1937\n1975\n1989\n1993\n300\n451\n48\n50\n83\n94\n95-minute\nabandoned\nabhorrent\nably\naction-comedy\naddict\nadmirers\nadrift\nadvice\naerial\naffirm\nafterwards\nagreeably\nailments\naisles\nall-night\nall-woman\namaze\namble\nambrose\nanalytical\nanne-sophie\nanniversary\nannoyance\nanteing\nanti-semitism\nappeared\napproaches\narcane\nargentinian\narrangements\narrogant\nasset\nathlete\natrocities\nauteur\nautopsy\nawesome\nawfulness\nbale\nbang-up\nbare\nbarf\nbarlow\nbask\nbearing\nbeen-there\nbeen-told-a\nbehalf\nbela\nbender\nbenjamins\nbig-wave\nbio-pic\nbloodletting\nbloodstream\nblowing\nboard\nbomb\nbombs\nbone-chilling\nboom-bam\nbothered\nbottom-feeder\nbouquet\nbravura\nbreathing\nbrian\nbridget\nbrit\nbuddy\nbull\nbursts\ncaffeinated\ncafé\ncake\ncameos\ncamouflage\ncamouflaging\ncanned\ncaptured\ncaricature\ncat-and-mouse\ncatapulting\ncelebrate\ncentering\ncenturies\nceremonies\ncertified\nchafing\nchatter\nchen\ncherry\nchewy\nchiaroscuro\nchicken\nchocolate\nchoppiness\nclamoring\nclamorous\nclams\nclancy\nclashing\ncleansing\ncliché-laden\nclients\nclose-ups\nclubs\nco-operative\ncoal\ncockettes\ncollide\ncollinwood\ncolumbia\ncolumn\ncombustible\ncommanding\ncommonplace\ncompare\ncompatible\ncompendium\ncomposure\ncompressed\ncondescension\nconfronting\nconjures\nconsciously\nconsummate\ncontinue\ncontriving\ncoping\ncoppola\ncorners\ncorporate\ncorpse\ncourtship\ncows\ncrackles\ncrappy\ncreators\ncrescendo\ncringe\ncrises\ncrooks\ncrossed\nculkin\ndante\ndass\ndates\ndawdle\ndazzle\ndecay\ndecency\ndedicated\ndeep-seated\ndefensible\ndeflated\ndelay\ndelusional\ndemeaning\ndemographically\ndepressingly\ndescends\ndesert\ndesirable\ndespicable\ndestructive\ndetailing\ndetective\ndeuces\ndevastation\ndevote\ndevotion\ndigital-effects-heavy\ndigs\ndisbelief\ndiscarded\ndiscerned\ndiscord\ndisguising\ndisingenuous\ndispassionate\ndistaste\ndocumentary-making\ndocumented\ndocumenting\ndoles\ndoltish\ndouglas\ndour\ndowner\ndownfall\ndramatized\ndramaturgy\ndrink\ndrives\ndumbness\nduplicate\neats\ned\neileen\neloquently\nembarrassingly\nemerging\nencountering\nendangered\nendear\nenforced\nentertains\nentranced\nequals\nescapes\nevery-joke-has\nevocation\nevolution\nexasperated\nexcepting\nexistentialism\nexpose\nexposes\nextension\nexxon\neyeballs\nfaced\nfaked\nfame\nfamine\nfarewell-to-innocence\nfarrelly\nfaulty\nfeardotcom.com\nfell\nfemale-bonding\nfeminine\nfence\nferrera\nfervently\nfine-looking\nfinesse\nfingers\nfinished\nflowering\nflowers\nflowery\nfollies\nfootball\nforthright\nfoster\nfour-star\nfragment\nframing\nfreud\nfrighteningly\nfunky\nfussy\nfuturistic\ngag\ngap\ngaps\ngenerating\ngermanic\ngleaned\nglued\ngoggles\ngooding\ngrandparents\ngreasy\ngrievous\ngrit\ngrotesquely\ngroundbreaking\nguessable\nguilt-trip\nguitar\ngutsy\nhal\nhalfwit\nhallucinatory\nhang-ups\nhappily\nhard-to-predict\nhard-won\nhaunted\nhawn\nhazy\nhealing\nheavyweight\nhellish\nhelmer\nheralds\nhiatus\nhigh-profile\nhilary\nhispanic\nhistoric\nhomes\nhooks\nhooliganism\nhospital\nhostage\nhot-blooded\nhubert\nhumble\nhumdrum\nhurley\nhustler\ni.q.\nideally\nignoring\nillogical\nillusion\nillustrating\nimaginable\nimagined\nimitator\nimmaturity\ninadequately\ninc.\nincapable\nincinerates\nincrease\nindicative\nindictment\nindisputably\ninexcusable\ninfluence\ninfuriating\ningratiating\ninsufferably\ninterplay\ninvest\ninvigorating\ninvincible\niran\nirrepressible\nirreverent\nitch\nitems\nj\njagjit\njay\njimmy\njohnny\njonathan\njunior-high\njust-above-average\nkate\nkept\nkieran\nkinds\nkingsley\nknoxville\nkooky\nlabeled\nlabyrinthine\nlaid-back\nlaissez-passer\nlargest-ever\nlashing\nlazier\nlaziest\nlean\nleap\nlevy\nliberal\nlingers\nlogically\nloquacious\nlore\nlow-brow\nlugosi\nluminous\nlumpy\nmaddening\nmalarkey\nmannered\nmap\nmargarita\nmary\nmatron\nmcculloch\nmeander\nmeanspirited\nmeaty\nmegaplex\nmegaplexes\nmelodramas\nmemento\nmenacing\nmends\nmeow\nmeticulously\nmick\nmiddle-of-the-road\nmidlife\nmilk\nmisconceived\nmisconstrued\nmishmash\nmistaken\nmodeled\nmodernize\nmodestly\nmonotone\nmonster/science\nmoonlight\nmorrison\nmothers\nmothman\nmotives\nmournful\nmovie-star\nmusicals\nmusty\nmyopic\nnachtwey\nname-calling\nnamesake\nnationwide\nnative\nnazi\nnerds\nnincompoop\nnominated\nnot-at-all-good\nnotably\nnow-cliched\noff-beat\nold-fashioned-movie\nold-hat\noperative\nopposite\noppressively\noptimism\norchard\norchestrating\nout-sized\noutshined\noverstays\npained\npaint\npaintings\npanic\npared\nparent\npartisans\npassages\npasta-fagioli\npasty\npatchwork\npathetically\npatriot\npatronizing\npb\npbs\npeace\npedestal\npeploe\nperforming\npermits\nphilippe\nphillip\npique\npitifully\nplayed-out\npleased\npleasingly\npoint-and-shoot\npolemic\npoorly-constructed\npopulace\nporno\nportrayed\nprecocious\npredict\npredisposed\nprefabricated\nprepare\nprimer\nprince\npristine\nprobe\nprobes\nprojectile\nprologue\nprolonged\npronounced\nprotect\npythonesque\nquaid\nquibble\nquirkiness\nraces\nrachel\nracism\nradioactive\nram\nrapid\nrarest\nraymond\nreaffirming\nrealizes\nreceives\nrecreates\nredeems\nreductions\nrefresh\nrelayed\nreleasing\nrendering\nrenders\nrenowned\nreprehensible\nresentment\nresolution\nresponsibilities\nrestrained\nrestraint\nrewritten\nribcage\nrifkin\nright-thinking\nrighteousness\nriot-control\nripoff\nrisks\nrobustness\nrodriguez\nromanticization\nromped\nrug\nrules\nrunyon\nrural\nrussell\nsabotaged\nsaddled\nsandwich\nsanguine\nsaps\nsaucy\nsaves\nschizo\nschlock\nscrape\nscrutiny\nscummy\nself\nself-destructiveness\nself-hating\nsensitivities\nseparation\nsexpot\nsharing\nshedding\nshekhar\nshimmering\nshmear\nshock-you-into-laughter\nshoots\nshow-stoppingly\nsickeningly\nsights\nsimmer\nsinger-turned\nsix\nskunk\nsleekness\nslimed\nsmall-budget\nsmallest\nsmile-button\nsmoky\nsneers\nsociology\nsons\nsophistication\nsophomore\nspectator\nspeeds\nspicy\nspiffing\nspinning\nspliced\nspots\nsputters\nstab\nstacked\nstasis\nstimulating\nstitched\nstoryteller\nstraining\nstrands\nstrip\nstrives\nstruck\nstructures\nstrung-together\nstultifyingly\nsturdiest\nstylistically\nsuccesses\nsugarman\nsultry\nsun-drenched\nsurgeon\nsurrenders\nsurrounded\nsuspected\nsustained\nswashbucklers\nswings\nswitchblade\nsympathies\nszpilman\ntaboo\ntambor\ntaxi\ntaymor\nteeny-bopper\nten\ntenor\ntequila\nterry\nthankfully\ntherapy\nthrew\nthrough-line\nthrust\ntickled\ntiresomely\ntongue\ntopics\ntossed\ntout\ntransparently\ntransporting\ntreating\ntreatise\ntremors\ntriangles\ntrotting\ntrouble-in-the-ghetto\ntucked\ntwo-way\nunbridled\nuncompelling\nundeterminable\nundistinguished\nunentertaining\nunequivocally\nunflattering\nunfolding\nunfussily\nunimpressively\nuninflected\nunmentionables\nunrelated\nvampires\nversions\nvh1\nvibrance\nvile\nvin\nvirtues\nvirtuoso\nvirtuous\nvistas\nvolletta\nvulgarity\nwall-to-wall\nwarmed\nwasp\nwaster\nwatered\nwayne\nweighs\nwell-wrought\nwelled\nwertmuller\nwhimsicality\nwidow\nwindshield\nwithholds\nwithstand\nwladyslaw\nwoe-is-me\nworkable\nworlds\nwriter/directors\nye\nyearnings\nyelling\nyellow\nyorker\nyoung-guns\nyours\nzeal\n'80s\n12th\n1957\n2,500\n295\n90-plus\n93\nabrupt\naccuse\naction-movie\nad\nadapted\naddress\nadmitting\nadolescence\nadoring\nadorns\nadultery\naesthetics\naffectionate\naffectionately\naggravating\nagitprop\nagonizing\nally\nalterations\nambitions\namicable\namid\namok\namours\namp\nanew\nannex\nanti-erotic\napproached\napproaching\narchly\narnie\narrogance\nartworks\nas-nasty\nassumption\nassuredly\nathleticism\nattach\nattracting\nauspicious\nauthenticate\nauthor\nauto\nautopilot\navid\nbackdrop\nbackward\nbaked\nballistic-pyrotechnic\nbanged\nbarbed\nbartleby\nbeer-fueled\nbefallen\nbegging\nbelieved\nbesotted\nbewildered\nbewilderingly\nbias\nbig-hearted\nbigelow\nbike\nbinks\nbirmingham\nblacked\nblackout\nblanks\nblazingly\nblowout\nblueprint\nbodice-ripper\nbona\nboorishness\nbout\nbrain-deadening\nbrand-new\nbrazenly\nbread\nbreaking\nbreathless\nbroadcast\nbrockovich\nbrow\nbruised\nbrusqueness\nbrussels\nbubble\nbubbly\nbueller\nburlap\nbyzantine\nc.h.o.\ncalamity\ncampanella\ncanadian\ncannibal\ncannon\ncared\ncarnage\ncarpets\ncartons\ncastrated\ncensure\nchain\nchamber\ncharlie\nchatty\ncheapening\ncheckout\ncheery\nchildish\nchomps\nchooses\nchore\nchâteau\ncinemantic\ncivic\nclaims\nclips\nclock\nclosest\nco-stars\ncocoon\ncoheres\ncoldest\ncommerce\ncomparisons\ncompensate\nconan\nconfessions\nconflagration\nconfusion\nconniving\nconspiratorial\nconversion\ncopies\ncostuming\ncoughed\ncounty\ncracker\ncradles\ncraftsmanship\ncraziness\ncreaky\ncreep\ncritique\ncrosses\ncrucifixion\ncrudely\ncurling\ncushion\ndanny\ndarned\ndeadeningly\ndebrauwer\ndeceit\ndecorating\ndefensive\ndefined\ndefinite\ndemonstration\ndepleted\nderive\ndescribed\ndidacticism\ndifficulty\ndime\ndiminishing\ndimwits\ndirector/co-writer\ndisappearing/reappearing\ndiscomfort\ndiscourse\ndisease\ndissecting\ndisservice\ndistinguishing\ndivorce\ndog-paddle\ndogtown\ndollar\ndoomed\ndouble-pistoled\ndrastic\ndrawn-out\ndreaded\ndreadfulness\ndress\ndrill\ndrumline\nducts\ndull-witted\neagle\neccentricities\neclipses\neconomic\nefficiency\nelicit\nelicits\nelizabethans\nemergence\nempathizes\nemploy\nendeavors\nenergizes\nengross\nequalizer\nequation\nerin\nescaped\net\netc\neven-handedness\neventful\nexclamation\nexcursion\nexpressiveness\nextensive\nextra-dry\nextravagantly\nextremist\neye-popping\nfallibility\nfamily-film\nfart\nfast-edit\nfaults\nfeminist\nfiddle\nfide\nfidgeted\nfifteen-year-old\nfincher\nfireballs\nfireworks\nflavorless\nflop\nflopping\nfoundation\nframes\nframework\nfranc\nfreaky\nfringes\nfrontman\nfudged\nfull-length\nfumes\nfun-seeking\nfundamentals\nfussing\ngaiety\ngain\ngallery\ngallic\ngamble\ngantz\ngarnered\ngas\ngathering\ngawky\ngiggle\nglucose\ngoldbacher\ngoodwill\ngoofily\ngored\ngoth-vampire\ngrade-grubbers\ngrade-school\ngrandiloquent\ngrandness\ngray\ngrease\ngreatly\ngreenfingers\ngroan\ngroaners\ngroggy\ngrossly\nguest\ngulzar\ngut-busting\nhangover\nhankies\nhanks\nhaplessness\nhappenstance\nhard-bitten\nhard-hearted\nhard-pressed\nhard-sell\nhaul\nhead-turner\nheartwarmingly\nheathers\nheed\nherrings\nhigh-end\nhighlighted\nholofcener\nhonorable\nhonored\nhop\nhopped-up\nhussein\niconography\nidiotically\nidol\nillogic\nimage-mongering\nimbued\nimitative\nimplausibility\nimpudent\nimpulse\nincarnations\nincessantly\ninconceivable\nindefinitely\nindoor\ninferior\ninhalant\ninjustice\nintegrates\ninterdependence\ninterference\nintergalactic\nintoxication\ninvite\nisland\nitalics\njacket\njanice\njanine\njanklowicz-mann\njar\njeunet\njoseph\njournalist\njulianne\njunk-calorie\njuwanna\nkafka\nkeel\nkidnappings\nkids-cute\nkills\nkirkegaard\nknickknacks\nkrawczyk\nl.a.\nlabor\nladies\nlameness\nlanding\nlapping\nlaramie\nlarge-scale\nlast-place\nlaziness\nleague\nlear\nlegs\nlend\nleonine\nless-compelling\nlethargic\nletter\nlick\nlight-heartedness\nlika\nlimbs\nlion\nliterature\nll\nlovefest\nlow-wattage\nlower-wit\nlowered\nlowly\nlucia\nlulled\nmajority\nmale-ridden\nmanaging\nmarathons\nmarina\nmarked\nmarred\nmasochism\nmasquerade\nmattei\nmcdonald\nmelange\nmermaid\nmesmerize\nmessing\nmexico\nmgm\nmib\nmilked\nmimics\nmind-destroying\nmindset\nminutely\nminutiae\nmirren\nmisogynist\nmixed-up\nmoaning\nmoat\nmoderately\nmojo\nmonotony\nmopes\nmorbid\nmordantly\nmorsels\nmotherhood\nmotionless\nmoulin\nmounting\nmovie-of-the-week\nmuccino\nmug\nmulti-layered\nmulti-layers\nmultiple\nmusclefest\nmuttering\nmyers\nmyrtle\nnagging\nnaiveté\nnapoleon\nnegatives\nnervous\nnewcastle\nnewness\nnick\nnonconformist\nnone-too-funny\nnorth\nnuclear\no'fallon\nobscenely\noff-putting\nok\nold-school\noleander\none-hour\none-sided\noomph\nopen-hearted\noperational\noppressive\noprah\noption\nopts\norbits\nordered\norders\noscar-caliber\nouter\noutlandish\noutrage\noverhearing\noverkill\novermanipulative\noverripe\novershadows\noverweight\npablum\npackages\npadding\npapin\nparental\nparrots\nparticipate\npartnership\npatch\npaws\npeppering\nper\nperception\nperch\nperfected\npermission\nperversity\npiccoli\npizazz\nplague\nplagued\nplans\nplastic\nplatter\nplaylist\nplaywriting\nplethora\nplummer\nplunge\npoetics\npointing\npoo-poo\nporn\npornographic\nporous\npost-feminist\npostmodern\npotboiler\npractices\npreaching\nprecarious\npreciseness\npreordained\npresenting\npresident\npresiding\nprimal\nprofundity\nprojector\nprom\npromenade\nprotective\nproved\nprovokes\npun\npunctuated\npunctuation\npurity\npurposefully\npushing\nquick-buck\nr&d\nrae\nrah-rah\nramble\nransacked\nransom\nrat-a-tat\nraunch-fests\nreassuringly\nrecalls\nrecognizes\nreconciliation\nrecovers\nrecycle\nreeks\nreflect\nregarding\nreinvigorated\nrepetitious\nrepressed\nreputedly\nresembles\nrespite\nresulted\nresurrecting\nretelling\nretiring\nrevels\nrhapsodic\nrhapsodizes\nright-wing\nrip\nroad-trip\nrough-around-the-edges\nrow\nrueful\nruminations\nrumor\nrun-of-the-filth\nrunner\nruthlessly\ns.c.\nsacrifices\nsaddam\nsameness\nsands\nsap\nsatiric\nscant\nscherfig\nscooby\nscorn\nscrapbook\nseal\nseizures\nself-control\nself-mocking\nself-parody\nself-promoter\nself-promotion\nsenseless\nsentimentalizing\nserpent\nsetpieces\nsettle\nshades\nshameful\nshiner\nshiver-inducing\nshootings\nshreve\nshriek\nsignpost\nsillified\nsingh\nsitcom-worthy\nsketches\nskilfully\nslasher\nslender\nslopped\nslugfest\nslurs\nsmall-town\nsmallness\nsmeared\nsnail\nsnake\nsneaky\nsneeze\nsoap-opera\nsoliloquies\nsparks\nspeculative\nspends\nspiral\nspits\nsplendidly\nsprawling\nsquint\nsquirming\nstandbys\nstandoffish\nstandout\nstar-studded\nstature\nstaying\nsteaming\nstench\nsterling\nsticking\nstimulate\nstonehenge\nstones\nstreamed\nstreet-smart\nstripped\nstrong-minded\nstumbles\nstuttering\nsubstantive\nsuicidal\nsuper-dooper-adorability\nsuper-simple\nsuperstar\nsupposedly\nsurfacey\nsurrealism\nsurrendering\nswank\nswimfan\nswipe\nswitch\nsword\nswords\ntables\ntadpole\ntailor\ntapestry\ntargets\ntedium\nteenybopper\nteleprompter\ntenderness\ntens\nthesps\nthinkers\nthird-rate\nthoughtfully\nthrives\nthroat-singing\nthud\ntick\ntied\ntightened\ntime-switching\ntimeout\ntitus\ntolkien\ntonto\ntoo-conscientious\ntorrent\ntrain\ntransporter\ntransvestite\ntravails\ntreachery\ntreading\ntricky\ntriumphantly\ntrots\ntroubling\ntub-thumpingly\ntucker\ntundra\ntutorial\ntwenty\nu.n.\nultra-violent\nunacceptable\nunaffected\nunapologetically\nunbalanced\nuncertainties\nunchanged\nunclean\nundercover\nunderground\nundeserved\nundramatic\nunerring\nunexceptional\nunforgivingly\nunfortunate\nunhibited\nuni-dimensional\nunimpeachable\nuninventive\nunivac-like\nunpersuasive\nunprecedented\nunprovoked\nunrelenting\nunruly\nunsatisfied\nunshapely\nvagueness\nvainly\nvat\nveil\nvelocity\nvenice\nverisimilitude\nvices\nvictorious\nvirtuosic\nviscerally\nwacky\nwalking\nwalt\nwatery\nweakness\nweirdo\nwell-defined\nwell-paced\nwell-realized\nwelt\nwesterners\nwheedling\nwherein\nwide-eyed\nwild-and-woolly\nwizened\nwollter\nwoozy\nworldly\nwretchedly\nwrote\nyielded\nzaidan\nzealand\nzero-dimensional\n****\n10,000\n10th-grade\n13th\n163\n18\n1980s\n20-car\n21/2\n60\n66\n89\n95\n99\naaliyah\nabel\naccidental\naccountant\nachieved\naching\nacidic\nactivism\nactorly\nactress-producer\naddictive\naddressing\nadhering\nadjusting\nadvance\nafghani\nagile\naim\naladdin\nalas\nall-around\nall-inclusive\nall-male\nalleged\naloof\nambience\nambivalence\nample\namusements\nandie\nandroid\nanimations\nanticipation\nantiseptic\nantsy\nappealingly\napplying\nappreciates\nappropriately\nardent\nardently\narmageddon\narticulates\nartificiality\nartistically\nartless\nascends\nasleep\naspire\naspired\nastounds\nastringent\nautomatically\naverting\nb.s.\nbadder\nbags\nbalm\nbalto\nbanger\nbanter-filled\nbarbarism\nbarker\nbars\nbattlefield\nbeard\nbeer\nbefuddled\nbegley\nbenchmark\nbergmanesque\nberkley\nberry\nbitten\nbitterly\nblazing\nblight\nbling-bling\nblob\nblond\nbloodshed\nblues\nboatload\nboffo\nbogging\nbolstered\nboom\nbordering\nborderline\nbotched\nbracingly\nbreathe\nbrogue\nbrooms\nbuck\nbuffoons\nbuying\nbygone\ncad\ncagney\ncameo\ncanadians\ncancer\ncapped\ncardoso\ncareless\ncarl\ncars\ncassavetes\ncatch-22\ncatharsis\ncellophane-pop\nchaiken\ncheek\ncheered\ncheerfully\nchemically\nchicanery\nchillingly\nchin\nchoke\nchristianity\nchronicle\ncinematographer\nclassify\nclaustrophic\nclaw\nclayburgh\nclear-cut\nclerk\nclimb\nclone\nclouds\nclout\ncollaborative\ncollar\ncombining\ncomfortably\ncomfy\ncomical\ncommend\ncompassionately\nconcentration\nconclusive\nconscious\nconsiders\nconspicuously\nconstruction\nconsuming\ncontact\ncontender\ncontinuum\ncontradiction\ncontradicts\ncontrasting\ncontributions\ncope\ncorn\ncotswolds\ncox\ncrassly\ncrawls\ncreed\ncreepy-scary\ncribbing\ncrisper\ncrowds\ncrumb\ncube\ncue\ncussing\ncutoffs\ncyndi\nda\ndared\ndawn\ndeafening\ndebts\ndecisive\ndegraded\ndehumanizing\ndel\ndelusions\ndenouement\ndepartments\ndependable\nderek\nderring-do\ndesolate\ndevito\ndiaz\ndiggs\ndilutes\ndimming\ndip\ndisaffected-indie-film\ndisagree\ndisciplined\ndiscovered\ndisease-of-the-week\ndishonesty\ndisposition\ndisquietingly\ndissing\ndistances\ndistinctions\ndistracting\ndivisions\ndizzy\ndocumentarian\ndollars\ndolphin-gasm\ndone-that\ndorkier\ndouble-barreled\ndown-and-dirty\ndraft\ndraggy\ndreadfully\ndreamy\ndrifts\ndruggy\ndrying\ndug\ndumas\ndungpile\ndustbin\near-pleasing\nearmarks\nearns\neclair\nelegy\nelite\nelm\nembellished\nembroils\nenacted\nencumbers\nenergized\nenhanced\nensues\nensuing\nentered\nenthrall\nenthusiastically\nentrée\nenvirons\nerratic\nescapist\nesteemed\nethos\neuphoria\nevanescent\neverett\nevergreen\neviction\nevolves\nexalted\nexception\nexchanges\nexoticism\nexpanse\nexplains\nexplosive\nexpressionistic\nextra-large\nexudes\nfacial\nfaint\nfamily-oriented\nfanatics\nfantastically\nfar-fetched\nfavorably\nfeed\nfeel-bad\nfetishistic\nfighters\nfirmly\nfish-out-of-water\nfitting\nflakeball\nflame\nflatula\nflatulence\nflavors\nflck\nfleeing\nflexible\nflibbertigibbet\nflies\nfling\nfloating\nflourish\nfollowers\nfood-for-thought\nforbidden\nforming\nforster\nfragmentary\nfranklin\nfreak-outs\nfreakish\nfree-wheeling\nfreely\nfreezers\nfreshening\nfriggin\nfrom-television\nfrothy\nfrustrates\nfu\nfullness\nfumbles\nfunnybone\nfusion\ngamut\ngedeck\ngelati\ngender-provoking\ngeneralities\nghandi\nghosts\ngiants\ngibberish\ngigantic\ngiggles\ngirl-on-girl\ngive-and-take\nglee\nglides\nglobalizing\nglobe\nglobetrotters-generals\ngloom\ngloomy\ngloss\ngo-for-broke\ngobbler\ngods\ngon\ngood-bad\ngood-naturedly\ngoose\ngoosebumps\ngore\ngorefest\ngrad\ngrand-scale\ngraves\nhairs\nhallucinogenic\nhanded\nhandles\nhands-off\nhanging\nhardship\nhardy\nharm\nharrison\nhated\nhawke\nheart-string\nheartbreak\nhefty\nhell-bent\nhelluva\nhence\nhibernation\nhickenlooper\nhiding\nhigh-powered\nhigh-strung\nhighway\nhills\nhired\nhiss\nholden\nhollywood-action\nhomophobia\nhomosexual\nhoneys\nhorror/action\nhue\nhumanist\nhumidity\nhundreds\nhypocrisy\ni.e.\nick\niconic\nigby\nignite\nill-advised\nimbue\nimmaculate\nimperfect\nimposed\nimpostor\ninadequate\nincognito\nindia\nindifferent\nindoctrinated\nindomitability\nindulged\nindustrial-model\nineffective\ninelegant\ninescapably\ninfluential\ninnovators\ninnumerable\ninquiry\ninseparable\ninstilled\ninsufficiently\nintegrating\ninterests\ninterior\nintermezzo\nintrepid\nintroduces\nintrospection\ninvented\ninvites\ninvitingly\nirony-free\nirritates\nisraeli\nivan\njam\njar-jar\njersey\njokers\njokester\njostling\nkaige\nkeenest\nkenneth\nkid-vid\nkitsch\nl.\nlaboratory\nladen\nland-based\nlaptops\nlaser\nlaughingly\nlaundry\nlauper\nlax\nle\nleaps\nlectured\nlectures\nleigh\nlens\nless-is-more\nlibrary\nlieutenant\nlife-altering\nlife-embracing\nlightly\nliking\nlil\nlimit\nlinearity\nlocal\nloco\nlondon\nlong-suffering\nlong-winded\nlook-see\nlove-struck\nlower-class\nloyalty\nluis\nlurks\nluscious\nlynch\nlyne\nmacaroni\nmachines\nmade-for-movie\nmaid\nmaker\nmanager\nmarivaux\nmarking\nmarried\nmarveled\nmasterly\nmasterpieces\nmatrix\nmattel\nmattered\nmax\nmcbeal-style\nmeager\nmedicine\nmelancholia\nmelville\nmercilessly\nmessage-mongering\nmetal\nmetaphysical\nmidsection\nmikes\nmilder\nmind-numbing\nming-liang\nminimal\nmisdemeanor\nmissive\nmixing\nmode\nmodem\nmodus\nmoldering\nmom\nmoney-oriented\nmonitor\nmoralism\nmores\nmoron\nmorton\nmotivate\nmourns\nmuch-needed\nmuscles\nmuse\nmust-own\nmuy\nnaipaul\nnaive\nnarrow\nnba\nneedy\nneo-fascism\nneorealism\nnettelbeck\nnetwork\nneutral\nnew-agey\nnicest\nnicole\nnightmarish\nnohe\nnonchalant\nnonconformity\nnone-too-original\nnonfiction\nnorma\nnormal\nnosedive\nnot-nearly\nnot-so-bright\nnouvelle\nnuts\nobjects\nobligation\noddballs\noffended\noftentimes\nomission\nontiveros\nopportunists\norange\norigins\noscar-sweeping\nout-of-kilter\noutbursts\noutrageously\nover-dramatic\nover-indulgent\noverinflated\novertake\noveruse\nowed\nowen\npages\npainstaking\npalatable\nparochial\npasolini\npassably\npatiently\npatrolmen\npauly\npaunchy\npaying\npenn\nperennial\nperkiness\nperpetually\npessimism\npetri\nphilosophers\nphonce\npin-like\npiscopo\npissed\npithy\nplant\nplayfully\nplucks\nplummets\npocket\npoise\npokémon\npolite\npootie\npop-music\nportraits\nportraiture\nportuguese\nposition\nposturing\npot\npours\npowerpuff\npracticed\npranks\npreachy-keen\nprecedent\npreferably\nprepackaged\nprescient\nprescribed\npreserving\npretention\nprettiest\nprevention\npriceless\nprickly\nprimary\nprincess\nprincipled\nprocedure\nprolific\nproof\nproperties\nproposes\nproven\nproverbial\nprovided\nprovoked\npseudo-bio\npsychedelia\npsychopathic\npulpiness\npunchier\npunchlines\npuppets\npuppies\npurdy\npuzzled\nq\nquaint\nquaking\nquestioning\nquickie\nquietude\nrain\nrambles\nramblings\nrampant\nre-creations\nredeemed\nredundancies\nreeking\nrefined\nregistering\nrejigger\nreligion\nremoved\nreopens\nrepartee\nrepeating\nrepellantly\nrequiring\nresents\nreside\nresorting\nresorts\nrespond\nretail\nretard\nretro-refitting\nreverie\nreviewers\nrevolutionaries\nrife\nrivalry\nrode\nroiling\nrose\nrotting\nrough\nrumblings\nsaddest\nsaddle\nsags\nsalute\nsalvage\nsandbox\nsarah\nscandal\nscarily\nscheming\nschnitzler\nscientist\nscoring\nscripting\nscum\nscuzzy\nsean\nsecond-rate\nself-amused\nself-centered\nself-delusion\nself-glorified\nself-mutilation\nself-righteous\nselves\nsemi\nsense-spinning\nservants\nsewing\nshackles\nshakesperean\nshift\nshifted\nshindler\nshiri\nshoes\nshootout\nshore\nshouting\nshow-biz\nsica\nsiege\nsignature\nsilbersteins\nsimmering\nsincerely\nsinger\nsingles\nsitcomishly\nslaloming\nslam-dunk\nsleaziness\nsleazy\nsleek\nslope\nslumming\nslump\nsmashing\nsmear\nsmell\nsmokey\nsnail-like\nsnared\nsnickers\nsociological\nsolace\nsole\nsolidity\nsomeday\nsoothe\nsparking\nsparkles\nsparse\nspecific\nspill\nspirit-crushing\nspiritualism\nsplashed\nsplashy\nsplitting\nspoofy\nsprouts\nstable-full\nstagey\nstaggered\nstamina\nstandup\nstatements\nstein\nstepmother\nstills\nstockwell\nstorm\nstraight-ahead\nstraight-to-video\nstraight-up\nstrategic\nstrenuously\nstudied\nstunt-hungry\nsturdiness\nstylist\nsubstandard\nsubstitutable\nsuckers\nsugary\nsuggesting\nsuitcase\nsulky\nsunny\nsunshine\nsuper-powers\nsuperficially\nsuperhero\nsuppose\nsurehanded\nsurrealistic\nswaggers\nswaying\nswift\nswimming\ntang\ntanks\ntartly\ntaxicab\nteaming\ntech-geeks\nteeth-gnashing\ntelegrams\ntemple\nthe-night\nthirteen\nthoughts\nthree-minute\nthroat\nthrusts\ntics\ntodd\ntoddler\ntongue-tied\ntongues\ntoo-frosty\ntoo-long\ntorments\ntorpedo\ntoss-up\ntouchstone\ntoys\ntrade\ntreads\ntreated\ntried-and-true\ntrier\ntrifecta\ntrimmings\ntrio\ntriple\ntriumphant\ntube\ntuck\nturntablists\ntwinkle\nu.s.\nugly-looking\nuhf\nuhhh\nultra-loud\nunashamedly\nunburdened\nunclear\nundergo\nunderlay\nundernourished\nunderrehearsed\nunderstandable\nunderstatement\nundertaking\nundisputed\nunengaging\nunflappable\nunhidden\nuniversity\nunreligious\nunsettled\nunwieldy\nupfront\nuser-friendly\nutilizing\nvacuum\nvalentine\nvalid\nvalidated\nvalley\nvaporize\nvarying\nversatile\nvideo-game-based\nviewpoint\nvittorio\nvoices\nvoyeuristic\nwaldo\nwalled-off\nwallflower\nwarmest\nwashout\nwater-bound\nwatts\nweaves\nweighted\nwelcomed\nwell-characterized\nwell-timed\nwell-told\nwell-worn\nwerner\nwhiff\nwidowmaker\nwiggling\nwillies\nwind-in-the-hair\nwised-up\nwisely\nwittgenstein\nwonderment\nworkings\nworkshops\nworldly-wise\nworries\nwreaked\nwreck\nwreckage\nxmas\nyale\nyarn-spinner\nyou-are-there\nyourselves\nzelda\nzhang\nzipper\n'50s\n*\n10-course\n112-minute\n1920\n1938\n1952\n1995\n2001\n24/7\n2455\n80-minute\n=\nabject\nabsorb\nacademic\naccident-prone\naccumulated\nacid\nacquires\nactioners\nactorliness\nactory\nadage\nadherents\nadmitted\nadolescents\nadrenalized\naffluence\nafford\naffords\nafrican-american\nagency\nagitator\nahem\naimlessness\nairs\nalcatraz\nalexandre\nall-powerful\nallegiance\naltar\nambiguities\namerican-russian\namoral\nanchors\nangelina\nangling\naniston\nankle-deep\nanna\nante\nants\napartments\napex\napollo\nappetizing\napted\narbitrarily\narchitect\narchive\nardor\narguably\nargues\narmchair\narrow\nartwork\nas-it\nashamed\nasparagus\nasphalt\nassassins\nassaults\nassembles\nassets\nassigned\nate\nattackers\nattendant\natypically\nauschwitz\nawash\nbaaaaaaaaad\nbackhanded\nbadly-rendered\nbanquet\nbarrage\nbasically\nbeaches\nbean\nbecame\nbedside\nbefuddling\nbeginnings\nbegrudge\nberling\nbeseechingly\nbesides\nbest-known\nbestial\nbestowing\nbetters\nbeyond-lame\nbicentennial\nbid\nbidder\nbig-fisted\nbigger-name\nbind\nbisset\nbladerunner\nbled\nblender\nblessed\nbluff\nboardwalk\nbodacious\nboho\nboiling\nboils\nbolado\nbombshell\nbon\nboorish\nboot\nboss\nbottom-of-the-bill\nbouncing\nbouts\nbowser\nbrass\nbrats\nbraveheart\nbravo\nbrawn\nbrazen\nbrendan\nbrink\nbromides\nbryan\nbugsy\nbump-in\nbuoy\nburdened\nburnt\nburnt-out\nburstein\nbush\nbusinesses\nbusts\nbutterflies\nbutterworth\nbéart\nc.\ncacoyannis\ncalvin\ncampaign\ncampaign-trail\ncampfire\ncandidate\ncandy-coat\ncapably\ncareening\ncareer-defining\ncareful\ncarmichael\ncaruso\ncasings\ncatastrophic\ncaterer\ncatherine\ncaton-jones\ncattaneo\ncellular\ncentered\ncesspool\nchainsaw\nchampagne\ncharacterisations\ncharitable\ncharlotte\ncheap-looking\ncheeks\nchefs\ncherished\nchill\nchomp\nchords\nchronically\nchurlish\nchurns\ncircumstantial\nclare\nclassicism\nclassification\nclassified\ncliff-notes\nclinical\nclint\nclive\ncloak\nclownish\nclumsiness\ncockney\ncoda\ncoincidence\ncollaboration\ncollected\ncolumbine\ncoma-like\ncomatose\ncombine\ncoming-of-age/coming-out\ncommands\ncommunicate\ncomparatively\ncomplain\ncomplaining\ncompletion\ncompliment\ncomprehensible\ncompromising\nconcentrating\nconcoctions\nconfection\nconfuses\ncongratulate\nconjure\nconsiderably\ncontemplative\ncontenders\ncontentedly\nconversation\nconveyor\nconvolutions\ncookie-cutter\ncooler\ncoos\ncor-blimey-luv-a-duck\ncounterculture\ncourtesy\ncowardly\ncowering\ncrapulence\ncrawling\ncrazed\ncreatively\ncreepiest\ncriminals\ncritic-proof\ncronenberg\ncrudity\ncruelly\ncrystal\ncrystallize\ncumulative\ncuter\nd'etre\ndadaist\ndaft\ndancers\ndas\nday-to-day\ndeathly\ndebilitating\ndecipherable\ndecisions\ndecomposition\ndecorous\ndefend\ndefinitions\ndelighted\ndemocracies\ndenouements\ndenuded\ndepress\nderisive\ndevos\ndirect-to-video/dvd\ndisappoints\ndisarmingly\ndiscerning\ndisgracefully\ndisintegrates\ndislocation\ndistort\ndistract\ndistractions\ndoa\ndoc\ndocs\ndog-tag\ndone-to-death\ndowns\ndrain\ndreamworks\ndriver-esque\ndrops\ndullingly\ndummies\ndusty\ndwells\nego-destroying\neisenstein\nel\nelder\nelection\nem\nembellishment\nembody\neminently\nencomia\nencourages\nendorses\nengendering\nengland\nennui\nennui-hobbled\nenthusiasms\nentrapment\nequations\nequipment\nespn\neternal\nevasive\neve\never-growing\never-watchful\nexecutives\nexhaustion\nexpensive\nexplanation\nexpository\nexposure\nfabulously\nfabulousness\nfacing\nfahrenheit\nfailings\nfairytale\nfallible\nfalters\nfanatical\nfaraway\nfathom\nfaux-urban\nfax\nfertile\nfigured\nfilter\nfinery\nfinger\nfingering\nflames\nflashing\nfledgling\nflimsier\nflopped\nflops\nflower\nflower-power\nflush\nfluttering\nfluxing\nflying\nfolksy\nfood-spittingly\nforeground\nforsaken\nfortune\nfrancisco\nfranco\nfresher\nfreshman\nfreudianism\nfriel\nfrightful\nfrittered\nfrothing\nfrustrated\nfuddled\nfuel\nfuels\nfun-loving\nfunctions\nfuse\nfusty\ngadgets\ngall\nganesh\ngangster/crime\ngarage\ngarner\ngenerated\ngeniality\ngenitals\ngenre-busting\ngenre-curling\ngentlemen\ngestalt\nghastly\ngiggly\ngimmick\ngimmicks\ngirl-meets-girl\ngirlish\ngive-me-an-oscar\nglamorous\ngliding\nglinting\nglitter\ngodfrey\ngoldie\ngolf\ngood-naturedness\ngoodies\ngooeyness\ngovernment\ngovernments\ngraham\ngrain\ngrandfather\ngrandiosity\ngrandly\ngrasping\ngreenlight\ngrinds\ngrinning\ngrittily\ngroan-to-guffaw\ngroups\nguiding\nguilt-free\ngullible\ngum\ngunplay\nguru\ngut-clutching\nhalfhearted\nhandily\nhandy\nhang\nhannibal\nharanguing\nhard-driving\nhard-edged\nhard-to-swallow\nhardass\nhardhearted\nharmoniously\nharvesting\nhashiguchi\nhat-in-hand\nheartily\nheft\nhell-jaunt\nheyday\nhideous\nhijinks\nhippie-turned-yuppie\nholiday-season\nhomework\nhomogenized\nhorrified\nhudlin\nhuge-screen\nhurried\nhurts\nhyped\nhyper-realistic\nhypnotically\nhysteria\nichi\nicily\nicky\nilluminated\nimitations\nimmigrant\nimplies\nimponderably\nin-jokes\ninactive\ninconclusive\ninconsistencies\nindependent-community\nindignant\ninduces\ninexpressible\ninexpressive\ninfectiously\ninfrequently\ninfuse\ninfusing\ningredient\ninhabitants\ninquisitiveness\ninsider\ninsinuation\ninstallments\ninsulted\ninter-species\ninternalized\nintimidated\ninventing\ninvestigation\ninvulnerable\niota\nirrational\nirreconcilable\njabs\njams\njaw\njeff\njez\njoint\njon\njournalists\njulie\njump-in-your-seat\nkafka-inspired\nkaos\nkid-movie\nkid-pleasing\nkilt-wearing\nkiosks\nkitchen-sink\nklein\nknees\nknucklehead\nkoepp\nkouyate\nkraft\nla\nlack-of-attention\nlarson\nlate-inning\nlaugh-free\nlaugher\nlaugther\nlaurence\nlearns\nleatherbound\nlecture\nlesbian\nlet-down\nlewd\nlibretto\nlimpid\nlingerie\nlinking\nlisa\nlive-style\nlivelier\nliveliness\nliza\nlobbyists\nlong-lived\nlooney\nloopholes\nloosely-connected\nlothario\nlottery\nlovable-loser\nlucas\nlying\nlyrics\nm-16\nmacbeth\nmaggots\nmajors\nmake-believe\nmanipulating\nmankind\nmarine\nmarine/legal\nmarketable\nmask\nmasked\nmassoud\nmaximum\nmeant\nmeetings\nmelted\nmerrily\nmesmerised\nmidst\nmika\nmileage\nmillions\nminiseries\nmisanthropy\nmishandled\nmisty-eyed\nmisunderstood\nmixes\nmob\nmodels\nmoist\nmolina\nmoney-grubbing\nmonosyllabic\nmonument\nmoodiness\nmoralistic\nmorals\nmornings\nmouglalis\nmounted\nmourning\nmoviegoer\nmoviemakers\nmuck\nmulan\nmulti-character\nmumbo-jumbo\nmummy\nmurk\nmythology\nnalin\nnarcotized\nnatural-seeming\nnaturalism\nnausea\nnavel\nne\nnear-fatal\nneedless\nnegate\nneighbor\nnerve-rattling\nnewcomers\nnine-year-old\nno-surprise\nnoisy\nnolan\nnon-disney\nnon-exploitive\nnon-techies\nnorm\nnotches\nnurtured\nobscenity\noccur\noddity\nodor\nofficer\noft-brilliant\noh-so-hollywood\nol'\nold-time\nolympic\nomniscient\nonline\noperandi\noperatic\nopposed\nopposites\norganizing\noscar-nominated\notherworldly\nounce\noutlet\noutre\noutweighs\noverdue\noverstating\noverstimulated\noverture\noverwhelmingly\npaean\npageantry\npainless\npaint-by-numbers\npalate\npallid\nparalyzed\nparapsychological\nparis\nparlor\nparticipant\nparties\npassive-aggressive\npastry\npatchy\npathos-filled\npearce\npearls\npeels\npell-mell\npender\npent\npercentages\npercussion\nperoxide\nperry\npersonable\nphotos\npicaresque\npicpus\npiffle\npipe\npitted\npivotal\npixar\npixilated\nplayoff\npogue\npoint-of-view\npokey\npollute\npop-up\npopped\nportraying\nposing\nposterity\npreaches\npreciousness\npredominantly\nprejudices\nprescription\npresume\nprevalent\npreviously\nprobation\nproclaim\nprofiling\nprogresses\nprojection\nprominent\nprominently\npromotion\nprophet\npropriety-obsessed\nprostitute\nprovincial\nprovocations\nprowess\nprurient\npseudo-intellectual\npsychologizing\npublicists\npublicity\npunishable\npurists\npurr\npursuing\nputters\npyrotechnics\nquadrangle\nr\nracing\nraffish\nrage\nragged\nraison\nrake\nransacks\nraphael\nratio\nratliff\nravaging\nrazzie\nre-working\nreagan\nreassure\nreassures\nrebel\nrecharged\nrecompense\nrecovering\nrecruiting\nreduce\nreductive\nreef\nreek\nrefers\nrefracting\nrefuse\nregards\nregular\nrejiggering\nremarks\nremembrance\nreparations\nreplacing\nrepugnant\nresemblance\nrestage\nrestatement\nretaining\nreturning\nrevigorates\nrevives\nrevolting\nrevolutionary\nrice\nriled\nringside\nripping\nroach\nrobbed\nrobberies\nrodrigues\nrogue\nromanced\nron\nroster\nrounded\nruins\nrésumé\nsaddens\nsale\nsalle\nsally\nsalma\nsanctimoniousness\nsandlerian\nsanity\nsayles\nscariest\nscathingly\nscene-chewing\nschmucks\nschticky\nscientific\nscooby-doo\nscotches\nscreed\nscrewing\nscuttled\nseamstress\nsearches\nseasonal\nseizing\nself-absorbed\nself-absorption\nself-image\nself-knowledge\nself-revealing\nsemi-amusing\nsemimusical\nsent\nsermonize\nserrault\nsettings\nshaky\nshambles\nsheets\nshepard\nshivers\nshoot-em-up\nshrug\nsickly\nsidekicks\nsimpering\nsingular\nskeeved\nskeleton\nskidding\nskimpy\nskims\nskipping\nskirts\nskullduggery\nskulls\nslam-bang\nsleep-inducingly\nslickly\nslogans\nsloppiness\nslyly\nsmall-screen\nsmorgasbord\nsmothered\nsmutty\nsnagged\nsnap\nsnappy\nsnide\nsnooze\nso-five-minutes-ago\nsoft-porn\nsomebodies\nsopranos\nsorriest\nsorrowfully\nsouvlaki\nspanning\nspirals\nspiritless\nsplendid-looking\nsplit-screen\nspooks\nsports-movie\nspry\nsqueezed\nsquirm-inducing\nstabs\nstagings\nstains\nstar-power\nstardom\nsteam\nsteamy\nstepdad\nsterile\nstevenson\nstickiness\nstifling\nstiflingly\nstop-and-start\nstrangling\nstress-reducing\nstrolls\nstudents\nstudio-produced\nstylings\nsubconscious\nsubjective\nsuburban\nsuffocate\nsuggestion\nsupple\nsupplies\nsure-fire\nsurplus\nsurveillance\nsweet-tempered\nswill\nsword-and-sorcery\nswung\nsympathizing\nsytle\ntabloid\ntackled\ntactic\ntainted\ntalked\ntalks\ntardier\ntarzan\ntear-stained\ntearful\ntearing\ntechnologies\nteen-speak\ntelescope\ntemperamental\nterminal\nterrifically\ntestosterone-charged\ntheatrics\nthemed\ntherapy-dependent\nthirteen-year-old\nthorough\nthoughtless\nthrobbing\nthrowaway\nthurman\nticks\ntinged\ntinsel\ntirade\ntitillating\ntolerable\ntonight\ntons\ntops\ntotalitarian\ntougher\ntourists\ntoy\ntraces\ntraffics\ntransfixes\ntransparency\ntransported\ntraps\ntremendously\ntrend\ntrials\ntrudge\ntrue-crime\ntruest\ntruffaut\ntryingly\ntunnels\nturmoil\ntweener\ntyco\nunadulterated\nunblinkingly\nuncluttered\nuncouth\nund\nundemanding\nunderachiever\nunderconfident\nunderlies\nundermines\nunderscoring\nundertaken\nundertones\nunearth\nunendurable\nunfold\nunfulfilled\nunguarded\nunified\nuninspiring\nuniqueness\nunit\nunits\nunlikeable\nunpaid\nunparalleled\nunsaid\nunsalvageability\nunsatisfactorily\nunsuccessful\nunsung\nunveil\nunwary\nupdated\nupping\nuproarious\nurbane\nvariant\nvariations\nvaunted\nveering\nventura\nventuresome\nverdu\nverging\nvideo-cam\nvideodrome\nvigils\nvirtuosity\nvirulently\nvisitor\nvivacious\nvoid\nvonnegut\nvulgarities\nw.\nwafer-thin\nwaif\nwaited\nwalken\nwallowing\nwaltzed\nwar-movie\nwar-weary\nwarn\nwartime\nwatered-down\nweakly\nwelcomes\nwell-conceived\nwell-contructed\nwell-directed\nwell-edited\nwheezy\nwhippersnappers\nwhipping\nwhir\nwhirlwind\nwhite-knuckled\nwide-awake\nwidescreen\nwilder\nwimps\nwindows\nwishy-washy\nwispy\nwitlessness\nwizardry\nworkers\nworkmanlike\nworkout\nworkshop\nworried\nwould-be\nwrapping\nwriter-actor\nx-files\nyank\nzhuangzhuang\n129-minute\n140\n15-year\n1960s\n1998\n2-day\n3-d\n7th\n87\n9-11\n9/11\n?!?\na.c.\na.e.w.\nabrasive\nabsent\nacclaimed\naccompanies\naccompanying\naccomplishments\nacerbic\nacquire\nacting-workshop\nactivate\nactivities\nadaptations\nadhere\nadversity\nadvertised\nadvertisement\nadvises\naffability\naffectingly\naffinity\nafloat\nafterlife\naground\nahola\nakin\nal\naldrich\nall-out\nalmodovar\nalready-shallow\nalternating\namalgam\namusedly\nanarchist\nanchored\nanecdote\nannie\nanti-\napartheid\nappearance\nappetite\napple\naptitude\narchives\naristocracy\naristocrat\narliss\narrived\nassembly\nassert\nassociation\nassures\nastronaut\nattending\nattics\nattributable\naudience-pleaser\nautistic\navalanche\navert\nawakens\nawareness\nb-12\nbackbone\nbacked\nbacklash\nbai\nbalk\nballerinas\nballsy\nbalzac\nbang-bang\nbanking\nbarrie\nbarrow\nbasest\nbath\nbathtub\nbattered\nbatting\nbawdy\nbeat-the-clock\nbedfellows\nbedtime\nbeginners\nbehind-the-scenes\nbeliefs\nbelly-dancing\nbellyaching\nbenefited\nbenshan\nbernal\nbertrand\nbespeaks\nbestowed\nbetrayed\nbetting\nbeware\nbickle\nbikes\nbilling\nbinary\nbio\nbitchy\nblanchett\nbleed\nblended\nblips\nblock\nblockage\nblossom\nboll\nbolster\nbons\nborrow\nbottomlessly\nbought\nbounds\nbranagh\nbrat\nbrazil-like\nbreed\nbrimming\nbroader\nbrothers-style\nbruckheimeresque\nbuckaroo\nbucket\nbuild-up\nbuilt-in\nbullseye\nbumps\nbungle\nbushels\nbyatt\ncadence\ncalling\ncamerawork\ncandles\ncaptivatingly\ncarrey\ncautions\ncedar\ncelebi\ncharacterizes\nchoreography\nchouraqui\nchurch-wary\ncircular\ncivics\nclarke-williams\nclasses\nclassroom\nclassy\nclaude\nclaustrophobia\ncleaving\nclicking\ncling\nclocks\nclones\nclotted\nclunker\nclutch\nclyde\nco-writer/director\ncocky\ncogent\ncoherence\ncoke\ncold-hearted\ncombustion\ncomedically\ncommenting\ncommiserating\ncommitment\ncommon-man\ncommunicating\ncommunity-college\ncompels\nconceivable\nconcludes\nconcrete\nconducted\nconfigurations\nconnections\nconnoisseurs\nconrad\nconsistency\nconspirators\ncontained\ncontemporaries\ncontemptuous\ncontentious\ncontributed\nconundrum\nconverts\nconvincingly\nconvolution\ncooker\ncooly\ncop-flick\ncopenhagen\ncopycat\ncoral\ncorrectness\ncounter\ncounterproductive\ncourtroom\ncow\ncracked\ncracks\ncranked\ncrash-and-bash\ncrashing\ncrawlies\ncredulity\ncrispin\ncritically\ncriticizing\ncritiquing\ncross-country\ncrowdpleaser\nculture-clash\ncurtains\ncut-and-paste\ncutting-room\nd.w.\ndanish\ndarkest\ndate-night\ndave\nday-lewis\nday-old\ndeath-defying\ndecadent\ndecides\ndeckhand\ndecommissioned\ndeemed\ndefense\ndeferred\ndegrades\ndelineate\ndelving\ndemocracie\ndemonic\ndemonstrated\ndenlopp\ndeportment\nderisions\ndesires\ndespairing\ndesultory\ndevelopmentally\ndi\ndiamond\ndicaprio\ndiction\ndiego\ndim\ndirgelike\ndisabled\ndiscontent\ndiscoveries\ndismantle\ndismember\ndisneyland\ndisparate\ndistressing\ndistressingly\ndoctorate\ndoings\ndonald\ndoors\ndoorstep\ndouble-cross\ndoubling\ndoyle\ndozing\ndramatize\ndramatizing\ndrang\ndreamscape\ndropped\ndrug-induced\ndrunk\ndudsville\nduel\nduly\ndunno\ndunst\ndwindles\ndystopian\neducational\neee\neffectiveness\negocentricities\negypt\neh\nelemental\nelvis\nembalmed\nembedded\nemerged\nemphasizing\nemphatic\nemulates\nenergizing\nenforcement\nengine\nengorged\nenhance\nenhancing\nenlightenment\nentree\nentwined\nepics\neroded\neroti-comedy\nescapade\neudora\neurope\nevacuations\neven-toned\nex-girlfriend\nexamples\nexasperatingly\nexcellence\nexcessively\nexchange\nexercises\nexiled\nexpanded\nexpansion\nexpectant\nfabuleux\nfacades\nfamuyiwa\nfang-baring\nfantasti\nfarcically\nfashioning\nfastballs\nfatalism\nfateful\nfault\nfearlessly\nferal\nfictional\nfidel\nfiller\nfillers\nfillm\nfilmgoers\nfire-breathing\nfirth\nflabbergasting\nflag\nflag-waving\nfleeting\nflick-knife\nflinging\nflirts\nfloats\nflog\nflows\nfocuses\nfolds\nfollow-your-dream\nforefront\nforgettably\nformalist\nformidable\nforms\nforte\nfour-year-old\nfourth-rate\nfrankly\nfret\nfrosting\nfruit\nfull-fledged\nfutility\ngadzooks\ngai\ngainsbourg\ngalled\ngamesmanship\ngarcia\ngarcía\ngarde\ngarden\ngarish\ngates\ngay-niche\ngays\ngender-bending\ngenesis\ngestures\ngianni\ngifford\ngilliam\ngirls-behaving-badly\nglancing\nglasses\ngleefully\nglover\nglows\ngluing\ngold\ngoldberg\ngoldman\ngoliath\ngone-to-seed\ngoodall\ngoodness\ngordy\ngoth\ngotten\ngraces\ngracious\ngrandmother\ngrateful\ngratitude\ngrenade\ngriffith\ngroove\ngrossest\ngulpilil\ngushing\ngut\nguzman\nhail\nhalf-an-hour\nhalf-bad\nhallelujah\nhand-drawn\nhandbag-clutching\nhandheld\nhanna-barbera\nhanussen\nharangues\nhardware\nharmon\nharvey\nhatred\nheadline-fresh\nheap\nhears\nhearst\nheartbreakingly\nheck\nheller\nhelpful\nhelping\nhermocrates\nhews\nhibiscus\nhidebound\nhigh-spirited\nhigh-tech\nhighly-praised\nhipness\nhollowness\nholmes\nhonoring\nhopkins/rock\nhorizons\nhormonal\nhorrific\nhotter\nhounds\nhowlingly\nhubristic\nhumanize\nhumbling\nhumiliated\nhumor-seeking\nhuskies\nhustlers\nhyper-real\nhypothesis\nii-birkenau\nill-equipped\nill-wrought\nilluminates\nillumination\nillustrated\nimaginary\nimamura\nimmense\nimmersive\nimpacts\nimperious\nimpetus\nimpish\nimplodes\nimply\nin-your-face\nindignation\nindividuality\nindulgence\ninextricably\ninfamy\ninfatuation\ninflammatory\ninflate\ninfused\ninsecurity\ninsistence\ninsistently\ninstrument\nintacto\nintellectually\nintermediary\nintermingling\nintern\ninternal\nintersect\nintrusive\niris\nirritatingly\nitem\njanuary\njargon\njaunt\njazz\njie\njoin\njonze\njournalistic\njoyful\njugglers\njuiced\njuiceless\njustifies\njuxtaposition\njuxtapositions\nkathie\nkathy\nkatzenberg\nke\nkhan\nkibbitzes\nkid-empowerment\nkiddies\nkilt\nkim\nkingdom\nkirsten\nkitten\nknockaround\nknockoff\nlackadaisical\nlacked\nlad\nlagging\nlandbound\nlarded\nlarge-format\nlascivious-minded\nlathan\nlawn\nlaws\nlazily\nleaping\nleblanc\nlecherous\nlecter\nleontine\nleys\nlifelong\nlifting\nlily\nliman\nlimply\nlina\nliner\nlioness\nlip-gloss\nlip-reading\nliterarily\nlitmus\nlive-action\nlive-wire\nloathe\nlog\nlogistically\nlongtime\nlouiso\nlow-tech\nlukewarm\nlullaby\nlunar\nlushness\nmade-up\nmakeup-deep\nmakin\nmalapropisms\nmanchild\nmanoel\nmanual\nmarcken\nmarginally\nmargins\nmargolo\nmarilyn\nmarkedly\nmarketing\nmarry\nmason\nmasters\nmatches\nmaxim\nmcmullen\nmeddles\nmedia-soaked\nmeet\nmelanie\nmemoir\nmerchandised-to-the-max\nmerited\nmethodology\nmid-to-low\nmiddle-age\nmiike\nmilieu\nmillennial\nmillennium\nmillisecond\nmined\nminnie\nmiracles\nmiserably\nmisfiring\nmisunderstanding\nmockumentary\nmoldy-oldie\nmommy\nmonologues\nmonopoly\nmonth\nmopping\nmoratorium\nmordant\nmortal\nmother-daughter\nmotivation\nmount\nmountains\nmovie-esque\nmovie-making\nmush\nmuted\nmystic\nmélange\nnail\nnapoli\nnarcissistic\nnaval\nnavel-gazing\nnear-impossible\nnear-masterpiece\nneeson\nneglecting\nnemesis\nnicolas\nnighttime\nnimble\nno-bull\nnovelty\nnowheresville\nnubile\nnurses\nnurtures\nnymphette\noblivious\nobscured\nobserved\noccupation\nodoriferous\noff-kilter\nold-world\noliviera\nomits\nopen-mouthed\noppositions\nopulent\noral\nordinances\notherness\noutage\noutdoes\noutgag\noutpaces\noverachieving\novercoming-obstacles\noverdone\novereager\noverflows\noverlapping\noverlook\noverpowered\noversexed\noversimplification\noversized\noverstylized\npaeans\npainkillers\npaltrow\npan\nparables\nparadoxically\npardon\npasts\npatricio\npatriotic\npause\npeaks\npeevish\npellington\npenance\npeopled\nperceptiveness\nperiod-piece\npermitting\npersonnel\npet\npetty\nphenomena\nphenomenon\nphiladelphia\nphonograph\nphysician\npics\npiecing\npies\npiles\npillowcases\npiss\nplaced\nplaguing\nplastered\nplausible\nplaybook\nplug\nplumbing\nplumbs\npoking\npolitesse\npop-cyber\npop-influenced\npore\npornography\nposes\npositives\npossesses\npowerhouse\npowerment\nprank\npreceded\nprefeminist\npreoccupations\npreserves\npressure\nprey\nprisoners\npro\npro-wildlife\nprocessed\nprofessors\nprofundities\nprognosis\nprogressed\npropelled\nprophecies\npropulsive\nprovince\npublishing\npummel\npurge\npuréed\nqueasy\nqueasy-stomached\nquixotic\nquotations\nr&b\nrabbit-proof\nrainbows\nrampantly\nran\nrancorous\nranges\nrapport\nrash\nrat\nrates\nravel\nreader\nready-made\nrealities\nreams\nrecesses\nrecorded\nreeked\nreeses\nreference\nreferred\nrefusing\nregimen\nreincarnation\nrelays\nrelieved\nreligious\nreminders\nrenegade-cop\nrenner\nrepetitively\nreportedly\nreporting\nrepulse\nresolutions\nresources\nrespecting\nrespects\nrestoring\nretold\nretreats\nreturned\nrevelled\nrevived\nrewrite\nreynolds\nrhames\nribisi\nrick\nricture\nrigged\nrightly\nromanticism\nrooted\nrose-tinted\nroussillon\nrusi\ns1m0ne\nsalt-of-the-earth\nsalvation\nsappiness\nsatisfaction\nsc2\nscalds\nscared\nscarface\nscary-funny\nscented\nschool-age\nschoolers\nscooter\nscreenful\nscrutinize\nseater\nseedy\nseeping\nseesawing\nseinfeld\nseldhal\nself-critical\nself-defeatingly\nself-exploitation\nself-reflection\nself-satisfaction\nsend-up\nsending\nsends\nsense-of-humour\nserendipity\nserene\nsettled\nsexism\nshallows\nshapiro\nshear\nshell\nshirley\nshirt\nshocker\nshortness\nshriveled\nsidesplitting\nsidewalks\nsilence\nsilences\nsilent\nsilent-movie\nsilly-looking\nsimulation\nsingle-handedly\nsingle-minded\nsister\nsix-time\nsketchiest\nslapped\nslash-fest\nsleep-inducing\nsleeping\nslow-paced\nsmall-scale\nsmaller\nsnaps\nsnow-and-stuntwork\nsoft\nsoftheaded\nsolemnity\nsolo\nsomeplace\nsomnambulant\nsooooo\nsoporific\nsoul-stirring\nsoullessness\nspain\nspall\nspending\nspikes\nspiritually\nsplashing\nspoken\nspoofs\nsporting\nspotty\nspousal\nspring-break\nspringer\nsprings\nsquander\nstakes\nstammering\nstar-making\nsteinis\nstew\nstill-raw\nstodgy\nstranded\nstray\nstreaks\nstreamlined\nstudiously\nstuffs\nsturm\nstylishly\nsub-formulaic\nsub-par\nsub-tarantino\nsubliminally\nsubmerged\nsubplot\nsubtexts\nsufficient\nsufficiently\nsummery\nsunset\nsuper-serious\nsuper-stupid\nsuperlative\nsurface-effect\nsurrealist\nsuspecting\nswanson\nsweaty-palmed\nswiftly\nswirling\nswordfights\nsymbiotic\nsynergistic\nsyrup\nsystematically\ntacked\ntail\ntailor-made\ntalancón\ntamer\ntantamount\ntartakovsky\ntasteless\nteams\ntear-drenched\ntechniques\nteddy\nteen-exploitation\nteen-sleaze\nteetering\ntelevised\ntenacious\ntend\nterrified\ntested\ntesting\nthirst\nthornier\nthoroughfare\nthoughtlessly\nthreefold\nthrilled\nthrillingly\ntides\ntill\ntime-killer\ntiniest\ntits\nton\ntootsie\ntop-heavy\ntough-man\ntraced\ntrack\ntraits\ntrash-cinema\ntravel-agency\ntravis\ntrickery\ntrifling\ntrilogy\ntrims\ntriumphs\ntrombone\ntrove\ntruck-loving\ntruth-in-advertising\ntumult\ntumultuous\ntunisian\nturks\nturntablism\ntv-movie\ntwilight\ntwinkling\ntyped\ntypes\nugh\nuglier\numa\nunaware\nuncharismatically\nuncharted\nunconcerned\nuncool\nunder-10\nunderestimated\nundergraduate\nundertone\nundying\nunencouraging\nunlaughable\nunnamed\nunpleasantly\nunremittingly\nunstoppable\nupends\nupheaval\nuplifter\nurges\nuselessly\nvainglorious\nvary\nvast\nveins\nvibrantly\nvideo-shot\nvideologue\nviews\nvigorously\nvillainous\nving\nvolatile\nvolumes\nvulakoro\nwalking-dead\nwall\nwaltz\nwanes\nwarm-milk\nwarmed-over\nwarped\nwash\nwebcast\nweber\nweinstein\nweissman\nwell-behaved\nwell-formed\nwell-mounted\nwell-put-together\nwelty\nwewannour\nwhip-smart\nwhitewash\nwhoop\nwimmer\nwind-tunnel\nwinger\nwinter\nwise-beyond-her-years\nwisegirls\nwittier\nwoods\nwrap\nwrath\nwretched\nwritings\nyakusho\nyields\nyuen\nzhao\nzinger-filled\nzoning\nzzzzzzzzz\nélan\n!?\n'90s\n'til\n10-year-old\n120\n14-year-old\n1790\n1950s\n1959\n1960\n1962\n1973\n1997\n22\n3/4th\n6-year-old\n65-minute\n65-year-old\n71\n75\n88\n8th\n94-minute\n99-minute\na-bornin\nabbott\nabdul\nabiding\nabsence\nabsorbs\nabused\naccording\nachival\nacidity\naddresses\nadept\nadobo\nadroit\nadvancing\nadvocacy\naffectation\naffectation-free\naftermath\nagape\nagreed\nairy\nalagna\nalan\nall-over-the-map\nalternate\namélie\nanakin\nandrei\nangles\nanimated-movie\nanthropomorphic\nanti-adult\nanti-catholic\nanti-feminist\nantonio\nappetites\nappropriated\narchibald\narchival\nargot\naristocrats\narithmetic\narthur\nasiaphiles\nassaultive\nassuming\nattacks\nattentions\nattitudes\nattuned\naudrey\naustrian\nauthenticity\nauto-critique\navant\navoiding\naxel\nayatollah\nbabbitt\nbaboon\nback-stabbing\nbackseat\nbait-and-switch\nbaker\nballroom\nbanzai\nbarbarian\nbarking-mad\nbaseball\nbaseball-playing\nbatman\nbeachcombing\nbearable\nbeavis\nbecalmed\nbeers\nbeg\nbelgium\nbelittle\nbelong\nbermuda\nbernard\nbet\nbettany\nbettany/mcdowell\nbetter-focused\nbig-budget/all-star\nbilingual\nbilly\nbiologically\nbiscuit\nbjorkness\nblab\nblacklight\nblatantly\nblemishes\nbloodbath\nblown-out\nbluffs\nblustery\nbmw\nbo\nboiled\nbombay\nborder\nborscht\nbotches\nbotching\nbounces\nbowl\nbox-office\nbrady\nbray\nbrims\nbroke\nbronze\nbrooks\nbugged\nbullfighters\nbully\nbumper\nbundling\nburgeoning\nbutler\nbutter\nbutthead\nbyler\nbyron\nbytes\nbüttner\ncaesar\ncaliber\ncalories\ncanada\ncandor\ncantet\ncarré\ncash\ncassel\ncatcher\ncategorize\ncattle\ncaustic\nchabrolian\nchampionship\ncharacter-driven\ncharred\ncharting\nchastity\ncheap-shot\ncheapen\ncheats\nchips\nchou-chou\nchristian-themed\nchuckling\nchung\nchyna\ncinematically\nclarify\nclaws\nclean-cut\ncleaver\nclinch\nclinically\nclosed\nco-dependence\nco-writers\ncoaster\ncoinage\ncold-blooded\ncold-fish\ncollaborators\ncollect\ncollegiate\ncollie\ncoltish\ncolumbus\ncomeback\ncomforting\ncomic-book\ncommended\ncommunication\ncommunications\ncompass\ncompeting\ncomplicate\ncomprise\ncompromised\ncompulsive\nconceal\nconceited\nconceptions\nconcerning\nconcerns\nconduits\nconfuse\ncongrats\nconsoled\nconspiracies\nconsumerist\ncontagious\ncontours\ncontrol-alt-delete\ncontrolled\ncool-j\ncoordinated\ncorcuera\ncoupling\ncrave\ncream\ncreator\ncrematorium\ncringing\ncritic\ncriticizes\ncrooning\ncross-dressing\ncrudities\ncrushed\ncultivated\ncunning\ncurrently\ncurtsy\ncutes\ncyber\ndabbles\ndark-as-pitch\ndebris\ndecades-spanning\ndecent-enough\ndecline\ndeconstruction\ndeem\ndeficit\ndefines\ndegrading\ndelete\ndelhi\ndemi\ndemocracy\ndensely\ndeny\ndependent\ndepictions\ndepravity\ndescribing\ndestin\ndetractors\ndevils\ndevious\ndiatribes\ndictates\ndie-hard\ndies\ndipped\ndire\ndisadvantage\ndisagreeable\ndisgusted\ndislikable\ndismal\ndismissive\ndisoriented\ndispossessed\ndistinction\ndiver\ndivertissement\ndivided\ndocumentarians\ndodge\ndogme\ndolly\ndominate\ndosage\ndoting\ndoubles\ndoubtful\ndown-home\ndownplaying\ndrama/action\ndramatization\ndramedy\ndrawings\ndresses\ndridi\ndrinking\ndripping\ndriving\ndrudgery\ndrum\ndrumbeat\ndualistic\nduds\ndungeons\ndwarfs\ndylan\ndynamism\ndyslexia\neccentricity\neconomics\neditor\neditorial\neighth-grader\neinstein\nelectrocute\neludes\nemaciated\nemailed\nembodies\nemphasising\nemphasized\nenabling\nenact\nenactments\nencyclopedia\nendured\nenglish-language\nengulfed\nenlightened\nenlivens\nensures\nepitaph\nernest\nerrol\neternity\nexaggeration\nexalts\nexcite\nexemplify\nexhaustingly\nexhibitionism\nexhibits\nexisting\nexplicit\nexploratory\nextends\nextracting\nexuberantly\nexude\nfaceless\nfai\nfanboy\nfancies\nfarm\nfather-and-son\nfavored\nfeatured\nfeeds\nfencing\nfetishes\nfifty\nfiguring\nfiji\nfist\nfizzle\nflails\nflakiness\nflood\nfloundering\nflounders\nflurries\nfocusing\nfogging\nfollow-up\nforeman\nforged\nforgot\nformalism\nfortify\nforward\nfoundering\nfrat-boy\nfresh-squeezed\nfriendly\nfringe\nfrolic\nfruitful\nfrustrations\nfueled\nfuhrman\nfuneral\nfurious\nfuzziness\ng-rated\ngained\ngallo\ngarnish\ngaudy\ngaye\ngaza\ngenerational\ngesture\nglamour\nglosses\nglumly\ngodzilla\ngoose-pimple\ngore-free\ngovernmental\ngrabs\ngraced\ngraffiti\ngranted\ngraphically\ngratify\ngraze\ngroaner\nguarded\nguiltless\nguise\ngussied\nhackery\nhailed\nhalf-dozen\nhalf-lit\nham\nhandiwork\nhard-eyed\nhardest\nharshness\nhayao\nhead-on\nheadaches\nhearing\nheavyweights\nheels\nheidegger\nhellstenius\nhelms\nheroism\nhideousness\nhigh-adrenaline\nhigh-buffed\nhinges\nhistrionic\nhit-and-miss\nhit-man\nhitchens\nhobbled\nhogwash\nhollywood-itis\nhollywood-predictable\nholofcenter\nhomo-eroticism\nhoofing\nhooked\nhoopla\nhorrifyingly\nhorses\nhossein\nhotels\nhotsies\nhotter-two-years-ago\nhousing\nhowler\nhumbuggery\nhunger\nhunky\nhyper-artificiality\nhypermasculine\niced\nideology\nidling\nimbecilic\nimparted\nimpatient\nimplications\nimplicitly\nimplied\nimportantly\nimposter\nimpressionistic\nimprovise\nin-jokey\ninadvertently\nincurably\nindecipherable\nindistinct\nindoors\nindulge\nindulges\nineptitude\ninfantilized\ninfinitely\ninformed\ningenue\ningest\ninhuman\ninsatiable\ninside-show-biz\ninsouciance\ninstances\ninstruct\nintelligibility\ninterludes\ninterpreting\ninterweaves\nintolerant\nintroverted\ninvestment\ninvokes\niosseliani\niq\nirredeemably\nirvine\nirwins\nisabelle\nisraeli/palestinian\nitalian-language\nivory\njacobi\njacquot\njapanimator\njelly\njettisoned\njews\njoffé\njoined\njoins\njostles\njournalistically\njovial\njuan\njuliet/west\njury\njustifying\nka\nkahn\nkarim\nkazan\nkeenly\nkeg\nkendall\nkids-in-peril\nkinnear\nkitschy\nknocks\nkoury\nkubrick\nkumble\nkurys\nlab\nlabel\nlamentations\nlandau\nlarky\nlate-summer\nlawyers\nlds\nleader\nledger\nleery\nlemon\nleniency\nleroy\nlessen\nlethally\nli\nliberalism\nlife-at-arm\nlifted\nlighter\nlike-themed\nlin\nlinear\nlinger\nlint\nliotta\nlittle-remembered\nliu\nlizard\nlo\nlocation\nlolita\nlone\nloop\nlooseness\nlosin\nlovably\nlove-hate\nlugubrious\nlumps\nlungs\nlusty\nluther\nlynne\nmachinery\nmacho\nmacnaughton\nmadcap\nmagician\nmaintained\nmajidi\nmalcolm\nmalik\nmalleable\nmamá\nmandy\nmanhunter\nmapquest\nmarveilleux\nmarvin\nmarxian\nmasochistic\nmass\nmasterpeice\nmatched\nmateys\nmatinee-style\nmccann\nmckay\nmclaughlin\nmediterranean\nmedium-grade\nmelting\nmending\nmentioned\nmenu\nmerge\nmerges\nmeshes\nmeyjes\nmichele\nmichelle\nmicroscope\nmiddle-class\nmidwest\nmilitary\nmilking\nmilks\nminac\nmini\nminimum\nminor-league\nmire\nmiscalculates\nmiscalculations\nmisfortune\nmishandle\nmitch\nmitchell\nmoan\nmodel\nmodernizes\nmodicum\nmonkey\nmonologue\nmorlocks\nmorphs\nmorris\nmortality\nmovie-industry\nmovie-specific\nmugs\nmurder-on-campus\nmurderer\nmurders\nmurray\nmusketeer\nmyriad\nmystification\nmythologizing\nnada\nnail-biter\nnamely\nnaomi\nnarrated\nnatalie\nnaturalness\nnauseating\nneglects\nneo-nazism\nneophyte\nnesbitt\nneurotic\nneurotics\nneverland\nnietzsche-referencing\nninety\nnon-fan\nnon-narrative\nnon-porn\nnorthwest\nnorwegian\nnot-so-divine\nnot-so-small\nnotch\nnoticed\nnotting\nnumbered\nnyc\nobsessively\noccupied\noff-center\noffset\noften-hilarious\noh-so-important\nokay\nolives\nomitted\nomnibus\none-room\none-sidedness\noo\nooze\nopenly\nopportunism\nopting\noptions\norbit\norganized\norlean\noscars\noutline\noutrageousness\noveracted\noverboard\noverrides\noverstated\noverwhelms\np.t.\npa\npabulum\npage-turning\npageant\npalestinian\npamela\npander\npap\nparadigm\nparmentier\npartner\npartners\npartnerships\nparton\npaulette\npauline\npaved\npeanut\nperalta\nperil\nperiod-perfect\nperpetual\nperplexing\npersistence\npersistent\npersonally\npersonas\nperversely\npessimistic\nphantasms\nphantom\nphocion\nphotographic\npicnic\npiercingly\npig\npinheads\npinks\npinochet\nplasma\nplaystation\nplea\nplutonium\npoint-to-point\npointedly\npoke-mania\npolicy\npollak\npollyana\npomposity\npondering\npoof\npooper-scoopers\npopulating\nportent\npossession\nposter\nposter-boy\nposthumously\npouty-lipped\npowered\npraying\npre-teen\npre-wwii\npreteen\npretence\nprints\nproceed\nprod\nproduces\nproducing\nprone\nproperty\nprostituted\nprotestors\nprovocateur\npseudo-philosophic\npsychedelic\npublished\npunched\npunitive\npurgatory\npurportedly\npursued\npursuers\npère-fils\nq.\nquashed\nquasi-improvised\nquasi-shakespearean\nquick-cut\nquicksand\nquintet\nradiates\nrafael\nragbag\nraging\nrails\nramifications\nrandolph\nrape\nrapidly\nrapids\nrated\nraving\nraw-nerved\nre-hash\nre-invents\nreap\nrecently\nrecessive\nreclaiming\nrecognizably\nrecreation\nreenacting\nreflected\nregister\nrelating\nrelied\nrelish\nremove\nrender\nreplace\nresonances\nrestate\nresume\nresurrection\nresuscitate\nrethink\nretooled\nretrograde\nretrospective\nreunion\nreunions\nreversal\nreverse\nrevision\nrevisionism\nreward\nreworking\nrhapsodize\nrhino\nriding\nrigidly\nrigors\nringu\nrinzler\nrip-roaring\nripper\nrisky\nrivals\nroads\nroars\nrockumentary\nrodan\nroland\nroller\nromantics\nromeo\nromoli\nrooms\nrough-trade\nrouses\nroute\nroutes\nroy\nrubbish\nrudd\nruffle\nruggero\nruin\nruse\nrusted-out\nruthless\nrye\nsacrificed\nsamurai\nsanitised\nsat\nsatisfactorily\nscam\nscar\nschepisi\nschiffer\nschneidermeister\nschrader\nscoob\nscorcher\nscored\nscratches\nscreeching-metal\nseated\nsecondary\nselby\nseldom\nself-determination\nself-expression\nself-pitying\nself-reflexive\nself-styled\nsemi-stable\nseptember\nserbs\nserial\nserious-minded\nserry\nservicable\nsetpiece\nsetup\nseventy-minute\nseverely\nsexiness\nshaken\nshatters\nshaw\nshimmeringly\nshoulder\nsigning\nsirk\nsixth\nsixties\nsixties-style\nskid-row\nskilled\nskinny\nskyscraper\nskyscraper-trapeze\nslaps\nsleight-of-hand\nslickest\nslickness\nslide\nslivers\nslot\nslow-motion\nslow-moving\nsmokers\nsmuggling\nsnapping\nsniffle\nsnoozer\nsoaked\nsoapy\nsolomonic\nsommers\nsong-and-dance-man\nsooner\nspecial-interest\nspin-off\nsplatterfests\nspookily\nspot-on\nspouting\nspringboard\nspringing\nsquaddie\nsquirts\nstadium-seat\nstalker\nstalking\nstallone\nstandardized\nstanzas\nstartled\nstayed\nsteadfast\nstep-printing\nstereotype\nstiletto-stomps\nstimulus\nstoner\nstorytellers\nstraddle\nstraight-faced\nstraight-shooting\nstreetwise\nstrenuous\nstrikingly\nstrip-mined\nstroked\nstrutting\nstultifying\nstumble\nsubmerging\nsugar-coated\nsuggestive\nsullivan\nsummons\nsunbaked\nsuper-sized\nsuperman\nsupremes\nsurest\nsurge\nsurveys\nsusceptible\nsuspended\nswathe\nsweet-and-sour\nsweetest\nswims\nsydow\nt-shirt\ntad\ntaiwanese\ntakashi\ntalkiness\ntambién\ntangents\ntap-dancing\ntautou\ntechno-tripe\nteenaged\ntelanovela\ntelephone\ntemptingly\ntendencies\ntendentious\ntenderly\ntestosterone\nthank\ntheorist\ntherapeutic\ntherefore\nthewlis\nthievery\nthinks-it-is\nthousand-times\nthreatens\nthuds\ntick-tock\ntidal\ntie\ntim\ntipped\ntiring\ntitans\ntitle-bout\ntogetherness\ntolerable-to-adults\ntonal\ntoo-hot-for-tv\ntornatore\ntouchingly\ntoughest\ntradition-bound\ntrailer-trash\ntrained\ntrainspotting\ntranquil\ntransition\ntraveled\ntreacly\ntree\ntrembling\ntrial\ntrickster\ntrimmed\ntriviality\ntrivialize\ntron\ntroopers\ntropic\ntroubadour\ntrusted\ntruths\ntu\ntufano\ntv-insider\ntwelve\ntwenty-some\ntwentysomething\ntwinkly-eyed\nuk\nultra-cheesy\nunaccountable\nunafraid\nunbelievably\nunder-7\nunderdeveloped\nundermined\nunderutilized\nunderventilated\nundiminished\nunexplained\nunhappiness\nunhurried\nunleashed\nunorthodox\nunromantic\nunscathed\nunslick\nunspeakably\nunstable\nunsuccessfully\nunthinkable\nunwavering\nuseful\nushered\nusurp\nvacuous\nvardalos\nveggies\nvenality\nvenice/venice\nverdict\nverismo\nviolated\nvirgin\nvocalized\nvoting\nwai\nwalter\nwarned\nwarner\nwatched\nwaterboy\nwazoo\nwearisome\nweighed\nweightless\nwell-meant\nwell-structured\nwell-trod\nwheels\nwhet\nwhine\nwhiney\nwhirling\nwholesale\nwhoopee-cushion\nwidely\nwildcard\nwillie\nwillingly\nwince\nwire\nwireless\nwiser\nwives\nwobbly\nwondered\nwonderland\nwonton\nwrecks\nwrestler\nwrite\nwriter-producer-director\nwrites\nwyman\ny\nya-yas\nyawner\nyeah\nyorkers\nz-boys\nzap\nzemeckis\nzoe\n'53\n'n\n170\n1899\n1950\n1980\n1986\n4/5ths\n50-year-old\n50s\n52\n60-second\n72\n96\nabbass\nabridged\naccentuating\naccommodate\naccomodates\naccumulate\nacolytes\naction/thriller\nactorish\naddams\naddiction\nado\nadopts\nadrian\nadrien\nadvert\naffect\nafghan\nafterschool\nagendas\nagony\naid\naided\nal.\nalexander\nalias\nalienated\nallison\naltered\nalternatives\naltman\naltman-esque\namari\namaro\nambitiously\nambivalent\namc\name\namidst\namini\namir\namuses\nana\nanalgesic\nanarchy\nanchor\nangela\nanimaton\nanonymity\nanswered\nanti-war\nanyplace\nappetizer\naragorn\naranda\narch\narchetypal\narchitecture\nargentinean\narkansas\narwen\nasquith\nassess\nassign\nassignment\nastronauts\nattentive\nattracts\naudacious-impossible\naugustine\naurelie\naustralian\nauteil\nautocritique\naversion\naward-winning\nayres\nayurveda\nbabies\nbabysitter\nbackwater\nbagatelle\nbaird\nballerina\nballet\nbang-the-drum\nbarbs\nbarriers\nbathos\nbattles\nbe-all-end-all\nbe-bop\nbeast-within\nbeasts\nbeating\nbecker\nbeckons\nbed\nbedevilling\nbehan\nbeijing\nbelieves\nbelinsky\nbeloved-major\nbenevolent\nbergman\nberkeley\nberlin\nbetty\nbibbidy-bobbidi-bland\nblandly\nblarney\nblasting\nbless\nblimp\nbluer\nblush\nboarders\nbollywood/hollywood\nbombards\nbonding\nbooths\nboston\nbots\nbottom-rung\nboxes\nboy-meets-girl\nbrainy\nbreasts\nbreitbart\nbrent\nbristles\nbrittle\nbroadside\nbrody\nbubbles\nbuff\nbuffeted\nbug\nbug-eye\nbullwinkle\nburies\nburlesque\nbusby\nbuscemi\nbustling\nbypassing\ncabins\ncam\ncameras\ncapitalism\ncapitalize\ncapitalizes\ncaptions\ncaptive\ncaptives\ncaptors\ncarlito\ncarousel\ncascade\ncavorting\ncenters\nchai\nchapter\ncharacterize\nchilled\nchimney\nchimps\nchokes\nchristelle\ncimarron\ncirca\ncircles\ncivil\ncivilization\nclarissa\nclearasil\ncliffhanger\nclimbing\nclosure\nclothing\ncloudy\ncolin\ncolleagues\ncolorfully\ncom\ncommune\ncompelled\ncomplaints\ncomponents\ncompositions\ncomputerized\nconcealment\nconfessional\nconfirms\nconflicted\nconjuring\nconscience\nconsequence\nconsists\nconstricted\nconstrictive\nconsume\ncontain\ncontinually\ncontorting\ncontribute\nconveys\ncopious\ncorbett\ncorpus\ncorrectly\ncouch\ncounting\ncousin\ncovered\ncrassness\ncrazier\ncreepy-crawly\ncremaster\ncrimen\ncross-cultural\ncrushing\ncryin\ncuaron\ncub\ncurlers\ncurve\nd.\ndaddy\ndalloway\ndampened\ndanis\ndaredevils\ndazed\ndead-eye\ndearly\ndeblois\ndebuts\ndecasia\ndecisively\ndecoder\ndefeats\ndefiant\ndefuses\ndegenerates\ndeleting\ndeliberateness\ndenmark\ndense\ndependence\ndepending\ndepicted\nderanged\nderivativeness\ndesecrations\ndesplat\ndetachment\ndetention\ndeteriorates\ndevelopers\ndeviant\ndeviously\ndicey\ndichotomy\ndick\ndictator-madman\ndifferently\ndignified\ndimensional\ndisassociation\ndiscloses\ndisconcertingly\ndisconnection\ndiscursive\ndiscuss\ndiscussed\ndishes\ndismally\ndisney-style\ndisobedience\ndispatching\ndispel\ndisplacement\ndisposible\ndisrobed\ndistill\ndistinguishable\ndistorts\ndistracts\ndisturb\ndisturbed\ndiva\ndiverges\ndoctor\ndocu-drama\ndocumentary-like\ndodger\ndogma\ndognini\ndolby\ndominated\ndomineering\ndonna\ndoris\ndorm\ndoshas\ndot\ndotted\ndouble\ndoug\ndr.\ndrek\ndrinker\ndrumming\ndrung\ndry-eyed\nduking\ndumbo\ne-graveyard\neardrum-dicing\nearnhart\nedged\nelderly\nelect\nelectoral\nelectronic\neleven\nelicited\nelves\nelysian\nemergency\nenamored\nendgame\nendings\nenemies\nengineering\nentering\nentries\nenvelope\nenvelops\nenvironments\nepiphanies\neponymous\nequate\nerupt\nesoteric\nessayist\nestablish\nestablishing\nestablishment\nestela\netched\nethnography\nevenly\neverlyn\nexceeding\nexceptions\nexhibit\nexpand\nexplain\nexplodes\nexploits\nexpressively\nexpressly\nexquisitely\nextant\nextrusion\neye-boggling\neye-catching\nf\nf.\nfabian\nface-to-face\nfactory\nfacts\nfairies\nfamed\nfantasized\nfarenheit\nfeardotcom\nfetid\nfields\nfiery\nfifth\nfilm-culture\nfinch\nfinishing\nfirmer\nflashback\nflashbulb\nfleder\nfleet-footed\nfloyd\nforewarned\nforgoes\nfork\nforty\nforwards\nfosters\nfounders\nfrankenstein-monster\nfrankie\nfreaking\nfreewheeling\nfreight\nfrench-produced\nfrenzy\nfrighten\nfuses\ngaghan\ngaitskill\ngalvanize\ngambles\ngamely\ngandalf\ngaunt\ngaï\ngeeked\ngeneva\ngeorgia\ngesturing\ngheorghiu\ngidget\ngiggling\ngirlfriends\nglitz\ngobble\ngood-looking\ngoombah\ngorgeousness\ngovernance\ngrabowsky\ngreat-grandson\ngreengrass\ngrenoble\ngrip\ngrizzled\ngrouchy\ngrunge-pirate\ngrungy\ngurus\nhack-and-slash\nhades\nhagiographic\nhairdo\nhairier\nhalf-assed\nhalf-hour\nhalle\nhammer\nhammering\nhammers\nhammily\nhands-on\nhanley\nhaphazardness\nhappily-ever\nhard-core\nhardwood\nharlem\nharmed\nharps\nhart\nhaute\nheadbanger\nhealth\nheartland\nheartstrings\nhermitage\nherring\nhierarchy\nhighlander\nhighways\nhints\nhippopotamus\nhitler-study\nhjelje\nhobnail\nholistic\nholland\nholy\nhomeric\nhomosexuality\nhookers\nhorde\nhorrid\nhos\nhouseholds\nhowling\nhush\nhypertime\nhélène\ni-heard-a-joke\niben\nidiom\nignites\nimmortal\nimmune\nimparting\nimpatiently\nimpresses\ninauspicious\ninclusiveness\nindependence\nindispensable\nindividuals\nineffable\ninert\ninertia\ninfants\ninfinite\ninfirmity\ningeniously\ninsanity\ninsiders\ninsular\ninter-family\ninter-racial\ninterlocked\ninterminably\nintervention\ninterviewees\nintoxicatingly\nintractable\ninvested\nireland\nirreparably\nirreversible\nirrevocable\nirrigates\niteration\nitinerant\nj.r.r.\njackal\njacobson\njammies\njangle\njeremy\njesus\njia\njudgment\njumbled\njune\nkane\nkang\nkathryn\nkeener\nkiddie-oriented\nkidnapper\nknitting\nkubrick-meets-spielberg\nkuras\nkurds\nlabors\nlaced\nlacey\nladles\nlagaan\nlapaglia\nlasker\nlate-night\nlau\nleanest\nleaning\nlector\nled\nleers\nlending\nles\nless-than-thrilling\nlethargically\nletterman\nliana\nlimps\nlingered\nlionize\nlived-in\nliven\nliving-room\nloaf\nlocusts\nlohman\nlongley\nlouts\nlovingly\nlucid\nlucks\nluke\nlushly\nlynch-like\nmachinations\nmade-for-tv\nmafia\nmagi\nmaguire\nmaiden\nmail\nmaladjusted\nmalle\nmalls\nmanically\nmarch\nmarcus\nmargaritas\nmarker\nmarkers\nmaryam\nmassacres\nmctiernan\nmeanest\nmeatier\nmedia-constructed\nmethod\nmethodically\nmetro\nmiddle-earth\nmighty\nmind-bender\nmind-bending\nmisbegotten\nmischievous\nmisleading\nmobius\nmock\nmodern-office\nmolehill\nmonkeys\nmoon\nmora\nmorrissette\nmosque\nmovie-biz\nmrs.\nmuckraking\nmud\nmumbles\nmurderous\nmusset\nmuster\nnanook\nnaqoyqatsi\nnationalism\nnature/nurture\nnecessity\nneeding\nneedles\nnegated\nnegativity\nneo-augustinian\nneo-hitchcockianism\nneo-noir\nnerve-raked\nnewfoundland\nniblet\nnicks\nno-holds-barred\nnon-bondish\nnonbelievers\nnonethnic\nnotations\nnumbness\nobligations\nobserve\noccasion\nodd-couple\noddest\noff-puttingly\noff-screen\noh\noh-those-wacky-brits\nolympia\nonion\noodles\noozes\nopen-minded\nopinion\norleans\northodox\notar\nourside\noutnumber\noutward\noutweigh\nover-familiarity\novernight\noverwhelm\nowe\nownership\np.c.\npadre\npageants\npalm\npan-american\npandora\nparachutes\nparanormal\nparaphrase\nparrot\nparty-hearty\npasach\npassionately\npastel\npatent\npatting\npax\npay-off\npayback\npeaked\npedro\npeep\npercolating\npersecuted\npessimists\npete\npetite\nphifer\nplaintiveness\nplan\nplate\nplato\nplay-doh\nplayboy\nplaywright\npledge\nplod\nplods\nplunges\npluto\npolice-procedural\npoo\npopulist\npork\nportrayals\npost-camp\npost-production\npost-war\npotter\npotty-mouthed\npractical\npractitioners\npray\npre-9\npredators\npreferable\nprepared\npreposterousness\npresses\nprima\nprincesses\npro-serbian\nprofessor\nprofits\nprogression\npropensity\npros\nprosaic\nprotecting\nprozac\npseudo-rock-video\npsychodrama\npublicist\npulchritude\npunch-drunk\npunches\npunchline\npuritanical\npurports\npyschological\nqatsi\nquasi-documentary\nquinn\nrainbow\nrancid\nrape-payback\nrapturous\nre-fried\nre-imagining\nre-voiced\nreact\nreadings\nreaffirms\nrealization\nreassembled\nrecalling\nreconceptualize\nrecreating\nrecycles\nredeem\nredefinition\nrediscovers\nreel/real\nreferential\nreginald\nregrets\nreid\nreigen\nreinvention\nrejection\nrejects\nrekindles\nrelocated\nremained\nrendition\nrenewal\nreporters\nrepresent\nrequiem\nrequires\nresidences\nresidents\nrespectively\nresponses\nrestless\nrestored\nrests\nresumes\nretaliatory\nretrospectively\nreverberates\nreverence\nreyes\nriddle\nright-hand\nrisk\nrolls\nrom\nromance-novel\nromantic-comedy\nromanticized\nromero\nrosario\nrosemary\nrothman\nrubber-face\ns&m\nsafely\nsailor\nsalacious\nsalvaged\nsam\nsampi\nsandeman\nsanders\nsascha\nsatan\nsavaged\nscarpia\nschindler\nschools\nschtick\nscouse\nscrap\nscrewy\nscribe\nscripted\nscripters\nsecretions\nsecretly\nseeds\nseldahl\nself-flagellation\nself-inflicted\nself-sacrifice\nself-serious\nseller\nselling\nsemi-surrealist\nseparates\nsept.\nserenely\nserviceability\nshag\nshamefully\nshatner\nshattering\nshimizu\nshoe-loving\nshoestring\nshopping\nshort-story\nshrewdly\nshum\nsimpson\nsinuously\nsiren\nsituates\nskipped\nslather\nsleepwalk\nslides\nslights\nslop\nsloughs\nslowness\nsmashups\nsnapshot\nsnatch\nsniping\nsnobbery\nsober-minded\nsociopathy\nsodden\nsofia\nsolely\nsongbird\nsorcerer\nsorely\nspace-based\nspanish\nsparked\nspawned\nspecialized\nspecifics\nspied\nspinoff\nsplatter\nspouses\nspout\nspreading\nsquabbling\nstabbing\nstacks\nstacy\nstaggers\nstallion\nsteadily\nstepmom\nsteps\nstereo\nstinging\nstink\nstockings\nstomps\nstoppard\nstraddles\nstrangest\nstrategies\nstress\nstripe\nstupefying\nstyle-free\nsubsequent\nsubsided\nsubtitled\nsuicide\nsunk\nsuppression\nsurfer\nsurround\nsurrounds\nsven\nswingers\nsy\nsylvie\nsymbols\nt.\ntara\ntarkovsky\ntease\nteasing\ntelenovela\ntemporal\ntemptations\nten-year-old\ntenth\ntexan\ntheology\nthinly-conceived\nthree-to-one\nthroes\nthump\ntighter\ntinseltown\ntissue-thin\ntit-for-tat\ntnt\ntome\ntommy\ntonally\ntoolbags\ntools\ntormented\ntourism\ntov\ntowers\ntracking\ntranslating\ntrivializing\ntrot\ntrue-blue\ntrumpet\ntrusts\ntruth-telling\nturpin\ntwinkie\nugly-duckling\nultra-provincial\nunambitious\nunbreakable\nuncinematic\nuncommitted\nunder-rehearsed\nundergoing\nunderway\nunderwear\nunderwhelming\nundogmatic\nuneasily\nunembarrassing\nunemployment\nunevenly\nunfairly\nunhinged\nunreachable\nunrealistic\nunseemly\nunsophisticated\nuntrained\nunwillingness\nupdates\nupstaged\nusher\nut\nutilizes\nvariable\nvaudeville\nver\nverite\nvicarious\nvicente\nvideo-viewing\nvideo/dvd\nvillainess\nvilleneuve\nvirtue\nviva\nviveka\nvivi\nwade\nwaged\nwaits\nwalks\nwar-torn\nwarlord\nwash.\nwatercolor\nwatstein\nwattage\nweaver\nwedgie\nwell-rounded\nwelles\nwending\nwhiffle-ball\nwidget\nwiel\nwintry\nwishful\nwistfully\nwizard\nwoe\nwoodman\nwoolf\nworship\nwounded\nwounding\nwrists\nwrought\nyawn-provoking\nyear-end\nyoda\nyong\nyvan\nzeitgeist\nzen\nzips\nzwick\n" }, { "path": "examples/sst2/vocabulary.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"A vocabulary that represents the tokens in a dataset and maps them to indices.\"\"\"\n\nimport collections\nfrom typing import Optional\nfrom collections.abc import Iterable, Sequence\n\nfrom absl import logging\n\n\nclass Vocabulary:\n \"\"\"Represents a vocabulary that can be built from a dataset.\"\"\"\n\n def __init__(\n self,\n vocab_path: str | None = None,\n tokenized_sequences: Iterable[Sequence[bytes]] | None = None,\n min_freq: int = 1,\n pad_token: bytes = b'',\n unk_token: bytes = b'',\n bos_token: bytes = b'',\n eos_token: bytes = b'',\n ):\n \"\"\"Loads the vocab from disk (if `vocab_path` is given) or builds it from `tokenized_sequences`.\"\"\"\n self.pad_token = pad_token\n self.unk_token = unk_token\n self.bos_token = bos_token\n self.eos_token = eos_token\n self.special_tokens = (pad_token, unk_token, bos_token, eos_token)\n\n if vocab_path:\n self.load(vocab_path)\n elif tokenized_sequences is not None:\n self.build(tokenized_sequences, min_freq=min_freq)\n else:\n raise ValueError(\n (\n 'Vocabulary needs either `vocab_path` or `tokenized_sequences` to'\n ' be provided, got %r and %r.'\n )\n % (vocab_path, tokenized_sequences)\n )\n\n def build(\n self, tokenized_sequences: Iterable[Sequence[bytes]], min_freq: int = 1\n ):\n \"\"\"Builds a vocabulary over tokens with optional minimum frequency.\n\n Args:\n tokenized_sequences: Iterable of token sequences to build the vocabulary.\n min_freq: The minimum frequency of each token to be included. Default: 1.\n \"\"\"\n # Count all the tokens.\n counter = collections.Counter()\n for token_sequence in tokenized_sequences:\n counter.update(token_sequence)\n\n # Add special tokens to the start of vocab.\n vocab = collections.OrderedDict()\n for token in self.special_tokens:\n vocab[token] = len(vocab)\n\n # Add all other tokens to the vocab if their frequency is >= min_freq.\n for token, freq in sorted(\n # Sort by frequency (from high to low), and then by token string.\n # This makes sure high frequency tokens get a low token ID.\n counter.items(),\n key=lambda token_freq: (-token_freq[1], token_freq[0]),\n ):\n if freq >= min_freq:\n vocab[token] = len(vocab)\n\n logging.info('Number of unfiltered tokens: %d', len(counter))\n logging.info('Vocabulary size: %d', len(vocab))\n self.vocab = vocab\n\n def _getitem__(self, key: str):\n return self.vocab[key]\n\n def keys(self):\n return self.vocab.keys()\n\n def values(self):\n return self.vocab.values()\n\n def __len__(self):\n return len(self.vocab)\n\n @property\n def pad_idx(self):\n \"\"\"The index of the padding token.\"\"\"\n return self.vocab[self.pad_token]\n\n @property\n def unk_idx(self):\n \"\"\"The index of the unknown word token.\"\"\"\n return self.vocab[self.unk_token]\n\n @property\n def bos_idx(self):\n \"\"\"The index of the beginning-of-sequence token.\"\"\"\n return self.vocab[self.bos_token]\n\n @property\n def eos_idx(self):\n \"\"\"The index of the end-of-sequence token.\"\"\"\n return self.vocab[self.eos_token]\n\n def load(self, path: str) -> None:\n \"\"\"Loads a vocabulary (one token per line) from disk.\"\"\"\n vocab = collections.OrderedDict()\n with open(path, 'rb') as f:\n for i, token in enumerate(f):\n assert isinstance(token, bytes), 'Expected byte tokens.'\n vocab[token.rstrip()] = i\n logging.info('Loaded vocabulary, size: %d', len(vocab))\n self.vocab = vocab\n\n def save(self, path: str) -> None:\n \"\"\"Saves the vocabulary to disk.\"\"\"\n with open(path, 'wb') as f:\n for token in self.vocab:\n assert isinstance(token, bytes), 'Expected byte tokens.'\n f.write(token + b'\\n')\n logging.info('Saved vocabulary to %s.', path)\n" }, { "path": "examples/vae/README.md", "content": "# Basic VAE Example\n\nThis is an implementation of the paper [Auto-Encoding with Variational Bayes](http://arxiv.org/abs/1312.6114) by D.P.Kingma and M.Welling.\nThis code follows [pytorch/examples/vae](https://github.com/pytorch/examples/blob/master/vae/README.md).\n\n```bash\npip install -r requirements.txt\npython main.py --workdir=/tmp/mnist --config=configs/default.py\n```\n\n## Overriding Hyperparameter configurations\n\nThis VAE example allows specifying a hyperparameter configuration by the means of\nsetting `--config` flag. Configuration flag is defined using\n[config_flags](https://github.com/google/ml_collections/tree/master#config-flags).\n`config_flags` allows overriding configuration fields. This can be done as\nfollows:\n\n```shell\npython main.py \\\n--workdir=/tmp/mnist --config=configs/default.py \\\n--config.learning_rate=0.01 --config.num_epochs=10\n```\n\n\n## Examples\n\nIf you run the code by above command, you can get some generated images:\n\n![generated_mnist](./sample.png)\n\nand reconstructions of test set digits:\n\n![reconstruction_mnist](./reconstruction.png)\n\nThe test set loss after 10 epochs should be around `104`.\n" }, { "path": "examples/vae/configs/default.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Default Hyperparameter configuration.\"\"\"\n\nimport ml_collections\n\n\ndef get_config():\n \"\"\"Get the default hyperparameter configuration.\"\"\"\n config = ml_collections.ConfigDict()\n\n config.learning_rate = 0.001\n config.latents = 20\n config.batch_size = 128\n config.num_epochs = 30\n return config\n" }, { "path": "examples/vae/input_pipeline.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Input pipeline for VAE dataset.\"\"\"\n\nimport jax\nimport jax.numpy as jnp\nimport tensorflow as tf\nimport tensorflow_datasets as tfds\n\n\ndef build_train_set(batch_size, ds_builder):\n \"\"\"Builds train dataset.\"\"\"\n\n train_ds = ds_builder.as_dataset(split=tfds.Split.TRAIN)\n train_ds = train_ds.map(prepare_image)\n train_ds = train_ds.cache()\n train_ds = train_ds.repeat()\n train_ds = train_ds.shuffle(50000)\n train_ds = train_ds.batch(batch_size)\n train_ds = iter(tfds.as_numpy(train_ds))\n return train_ds\n\n\ndef build_test_set(ds_builder):\n \"\"\"Builds train dataset.\"\"\"\n test_ds = ds_builder.as_dataset(split=tfds.Split.TEST)\n test_ds = test_ds.map(prepare_image).batch(10000)\n test_ds = jnp.array(list(test_ds)[0])\n test_ds = jax.device_put(test_ds)\n return test_ds\n\n\ndef prepare_image(x):\n x = tf.cast(x['image'], tf.float32)\n x = tf.reshape(x, (-1,))\n return x\n" }, { "path": "examples/vae/main.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Main file for running the VAE example.\n\nThis file is intentionally kept short. The majority for logic is in libraries\nthat can be easily tested and imported in Colab.\n\"\"\"\n\nfrom absl import app\nfrom absl import flags\nfrom absl import logging\nfrom clu import platform\nimport jax\nfrom ml_collections import config_flags\nimport tensorflow as tf\n\nimport train\n\n\nFLAGS = flags.FLAGS\n\nflags.DEFINE_string('workdir', None, 'Directory to store model data.')\nconfig_flags.DEFINE_config_file(\n 'config',\n None,\n 'File path to the training hyperparameter configuration.',\n lock_config=True,\n)\nflags.mark_flags_as_required(['config', 'workdir'])\n\n\ndef main(argv):\n if len(argv) > 1:\n raise app.UsageError('Too many command-line arguments.')\n\n # Make sure tf does not allocate gpu memory.\n tf.config.experimental.set_visible_devices([], 'GPU')\n\n logging.info('JAX process: %d / %d', jax.process_index(), jax.process_count())\n logging.info('JAX local devices: %r', jax.local_devices())\n\n # Add a note so that we can tell which task is which JAX host.\n # (Depending on the platform task 0 is not guaranteed to be host 0)\n platform.work_unit().set_task_status(\n f'process_index: {jax.process_index()}, '\n f'process_count: {jax.process_count()}'\n )\n\n train.train_and_evaluate(FLAGS.config, FLAGS.workdir)\n\n\nif __name__ == '__main__':\n app.run(main)\n" }, { "path": "examples/vae/models.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"VAE model definitions.\"\"\"\n\nfrom flax import linen as nn\nfrom jax import random\nimport jax.numpy as jnp\n\n\nclass Encoder(nn.Module):\n \"\"\"VAE Encoder.\"\"\"\n\n latents: int\n\n @nn.compact\n def __call__(self, x):\n x = nn.Dense(500, name='fc1')(x)\n x = nn.relu(x)\n mean_x = nn.Dense(self.latents, name='fc2_mean')(x)\n logvar_x = nn.Dense(self.latents, name='fc2_logvar')(x)\n return mean_x, logvar_x\n\n\nclass Decoder(nn.Module):\n \"\"\"VAE Decoder.\"\"\"\n\n @nn.compact\n def __call__(self, z):\n z = nn.Dense(500, name='fc1')(z)\n z = nn.relu(z)\n z = nn.Dense(784, name='fc2')(z)\n return z\n\n\nclass VAE(nn.Module):\n \"\"\"Full VAE model.\"\"\"\n\n latents: int = 20\n\n def setup(self):\n self.encoder = Encoder(self.latents)\n self.decoder = Decoder()\n\n def __call__(self, x, z_rng):\n mean, logvar = self.encoder(x)\n z = reparameterize(z_rng, mean, logvar)\n recon_x = self.decoder(z)\n return recon_x, mean, logvar\n\n def generate(self, z):\n return nn.sigmoid(self.decoder(z))\n\n\ndef reparameterize(rng, mean, logvar):\n std = jnp.exp(0.5 * logvar)\n eps = random.normal(rng, logvar.shape)\n return mean + eps * std\n\n\ndef model(latents):\n return VAE(latents=latents)\n" }, { "path": "examples/vae/requirements.txt", "content": "absl-py==1.4.0\nflax==0.6.9\nnumpy==1.23.5\noptax==0.1.5\nPillow==10.2.0\ntensorflow==2.12.0\ntensorflow-datasets==4.9.2" }, { "path": "examples/vae/results/.gitignore", "content": "*.png\n" }, { "path": "examples/vae/train.py", "content": "# Copyright 2023 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"Training and evaluation logic.\"\"\"\nfrom typing import Any\n\nfrom absl import logging\nfrom flax import linen as nn\nimport input_pipeline\nimport models\nimport utils as vae_utils\nfrom flax.training import train_state\nimport jax\nfrom jax import random\nimport jax.numpy as jnp\nimport ml_collections\nimport optax\nimport tensorflow as tf\nimport tensorflow_datasets as tfds\n\n\n@jax.vmap\ndef kl_divergence(mean, logvar):\n return -0.5 * jnp.sum(1 + logvar - jnp.square(mean) - jnp.exp(logvar))\n\n\n@jax.vmap\ndef binary_cross_entropy_with_logits(logits, labels):\n logits = nn.log_sigmoid(logits)\n return -jnp.sum(\n labels * logits + (1.0 - labels) * jnp.log(-jnp.expm1(logits))\n )\n\n\ndef compute_metrics(recon_x, x, mean, logvar):\n bce_loss = binary_cross_entropy_with_logits(recon_x, x).mean()\n kld_loss = kl_divergence(mean, logvar).mean()\n return {'bce': bce_loss, 'kld': kld_loss, 'loss': bce_loss + kld_loss}\n\n\ndef train_step(state, batch, z_rng, latents):\n \"\"\"Train step.\"\"\"\n def loss_fn(params):\n recon_x, mean, logvar = models.model(latents).apply(\n {'params': params}, batch, z_rng\n )\n\n bce_loss = binary_cross_entropy_with_logits(recon_x, batch).mean()\n kld_loss = kl_divergence(mean, logvar).mean()\n loss = bce_loss + kld_loss\n return loss\n\n grads = jax.grad(loss_fn)(state.params)\n return state.apply_gradients(grads=grads)\n\n\ndef eval_f(params, images, z, z_rng, latents):\n \"\"\"Evaluation function.\"\"\"\n def eval_model(vae):\n recon_images, mean, logvar = vae(images, z_rng)\n comparison = jnp.concatenate([\n images[:8].reshape(-1, 28, 28, 1),\n recon_images[:8].reshape(-1, 28, 28, 1),\n ])\n\n generate_images = vae.generate(z)\n generate_images = generate_images.reshape(-1, 28, 28, 1)\n metrics = compute_metrics(recon_images, images, mean, logvar)\n return metrics, comparison, generate_images\n\n return nn.apply(eval_model, models.model(latents))({'params': params})\n\n\ndef train_and_evaluate(config: ml_collections.ConfigDict, workdir: str):\n \"\"\"Train and evaulate pipeline.\"\"\"\n tf.io.gfile.makedirs(workdir)\n\n rng = random.key(0)\n rng, key = random.split(rng)\n\n ds_builder = tfds.builder('binarized_mnist')\n ds_builder.download_and_prepare()\n\n logging.info('Initializing dataset.')\n train_ds = input_pipeline.build_train_set(config.batch_size, ds_builder)\n test_ds = input_pipeline.build_test_set(ds_builder)\n\n logging.info('Initializing model.')\n init_data = jnp.ones((config.batch_size, 784), jnp.float32)\n params = models.model(config.latents).init(key, init_data, rng)['params']\n\n state = train_state.TrainState.create(\n apply_fn=models.model(config.latents).apply,\n params=params,\n tx=optax.adam(config.learning_rate),\n )\n\n rng, z_key, eval_rng = random.split(rng, 3)\n z = random.normal(z_key, (64, config.latents))\n\n steps_per_epoch = (\n ds_builder.info.splits['train'].num_examples // config.batch_size\n )\n\n for epoch in range(config.num_epochs):\n for _ in range(steps_per_epoch):\n batch = next(train_ds)\n rng, key = random.split(rng)\n state = train_step(state, batch, key, config.latents)\n\n metrics, comparison, sample = eval_f(\n state.params, test_ds, z, eval_rng, config.latents\n )\n vae_utils.save_image(\n comparison, f'{workdir}/reconstruction_{epoch}.png', nrow=8\n )\n vae_utils.save_image(\n sample, f'{workdir}/sample_{epoch}.png', nrow=8\n )\n\n print(\n 'eval epoch: {}, loss: {:.4f}, BCE: {:.4f}, KLD: {:.4f}'.format(\n epoch + 1, metrics['loss'], metrics['bce'], metrics['kld']\n )\n )\n" }, { "path": "examples/vae/utils.py", "content": "# Copyright 2023 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nThis code is created with reference to torchvision/utils.py.\n\nModify: torch.tensor -> jax.numpy.DeviceArray\nIf you want to know about this file in detail, please visit the original code:\n https://github.com/pytorch/vision/blob/master/torchvision/utils.py\n\"\"\"\nimport math\n\nimport jax.numpy as jnp\nimport numpy as np\nfrom PIL import Image\n\n\ndef save_image(ndarray, fp, nrow=8, padding=2, pad_value=0.0, format_img=None):\n \"\"\"Make a grid of images and Save it into an image file.\n\n Args:\n ndarray (array_like): 4D mini-batch images of shape (B x H x W x C)\n fp: A filename(string) or file object\n nrow (int, optional): Number of images displayed in each row of the grid.\n The final grid size is ``(B / nrow, nrow)``. Default: ``8``.\n padding (int, optional): amount of padding. Default: ``2``.\n pad_value (float, optional): Value for the padded pixels. Default: ``0``.\n format_img(Optional): If omitted, the format to use is determined from the\n filename extension. If a file object was used instead of a filename,\n this parameter should always be used.\n \"\"\"\n\n if not (\n isinstance(ndarray, jnp.ndarray)\n or (\n isinstance(ndarray, list)\n and all(isinstance(t, jnp.ndarray) for t in ndarray)\n )\n ):\n raise TypeError(f'array_like of tensors expected, got {type(ndarray)}')\n\n ndarray = jnp.asarray(ndarray)\n\n if ndarray.ndim == 4 and ndarray.shape[-1] == 1: # single-channel images\n ndarray = jnp.concatenate((ndarray, ndarray, ndarray), -1)\n\n # make the mini-batch of images into a grid\n nmaps = ndarray.shape[0]\n xmaps = min(nrow, nmaps)\n ymaps = int(math.ceil(float(nmaps) / xmaps))\n height, width = (\n int(ndarray.shape[1] + padding),\n int(ndarray.shape[2] + padding),\n )\n num_channels = ndarray.shape[3]\n grid = jnp.full(\n (height * ymaps + padding, width * xmaps + padding, num_channels),\n pad_value,\n ).astype(jnp.float32)\n k = 0\n for y in range(ymaps):\n for x in range(xmaps):\n if k >= nmaps:\n break\n grid = grid.at[\n y * height + padding : (y + 1) * height,\n x * width + padding : (x + 1) * width,\n ].set(ndarray[k])\n k = k + 1\n\n # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer\n ndarr = np.array(jnp.clip(grid * 255.0 + 0.5, 0, 255).astype(jnp.uint8))\n im = Image.fromarray(ndarr.copy())\n im.save(fp, format=format_img)\n" }, { "path": "examples/wmt/README.md", "content": "## Machine Translation\n\nTrains a Transformer-based model (Vaswani *et al.*, 2017) on the WMT Machine\nTranslation en-de dataset.\n\nThis example uses linear learning rate warmup and inverse square root learning\nrate schedule.\n\nTable of contents:\n\n- [Requirements](#requirements)\n- [Example runs](#example-runs)\n- [Running on Cloud](#running-on-cloud)\n - [Preparing the dataset](#preparing-the-dataset)\n - [Google Cloud TPU](#google-cloud-tpu)\n\n### Requirements\n\n* TensorFlow datasets `wmt17_translate/de-en` and `wmt14_translate/de-en` need\n to be downloaded and prepared. A sentencepiece tokenizer vocabulary will be\n automatically generated and saved on each training run.\n* This example additionally depends on the `sentencepiece` and\n `tensorflow-text` packages.\n\n### Example runs\n\nYou should expect to get numbers similar to these:\n\n\nHardware | config | Training time | BLEU | TensorBoard.dev | Workdir\n-------- | ------- | ------------- | -------------- | ------------------------------------------------------------------------ | --------------------------------------------------------------------------------------------------------------------------\nTPU v3-8 | default | 24m
13h18m | 25.55
32.87 | [2021-08-04](https://tensorboard.dev/experiment/nnH7JNCxTgC1ROakWePTlg/) | [gs://flax_public/examples/wmt/default](https://console.cloud.google.com/storage/browser/flax_public/examples/wmt/default)\nTPU v3-32 | default | 3h1m | 32.45 | [2021-11-05](https://tensorboard.dev/experiment/7IKeXjoeRKiMtqysQlbqzw/) | [gs://flax_public/examples/wmt/default_v3-32](https://console.cloud.google.com/storage/browser/flax_public/examples/wmt/default_v3-32)\nGPU V100 x8 (Mixed Precision) | gpu_mixed_precision | 1h 58m | 25.69 | [2021-07-07](https://tensorboard.dev/experiment/9S2WuqNWRDemmBuQE8K6Ew/) | -\n\n\n### Running on Cloud\n\n#### Preparing the WMT Datasets\n\nWe recommend downloading and preparing the TFDS datasets beforehand. For Cloud\nTPUs, we recommend using a cheap standard instance and saving the prepared TFDS\ndata on a storage bucket, from where it can be loaded directly. Set the\n`TFDS_DATA_DIR` to your storage bucket path (`gs://`).\n\nYou can download and prepare any of the WMT datasets using TFDS directly:\n`python -m tensorflow_datasets.scripts.download_and_prepare\n--datasets=wmt17_translate/de-en`\n\nThe typical academic BLEU evaluation also uses the WMT 2014 Test set: `python -m\ntensorflow_datasets.scripts.download_and_prepare\n--datasets=wmt14_translate/de-en`\n\n#### Google Cloud TPU\n\nSee below for commands to set up a single VM with 8 TPUs attached\n(`--accelerator-type v3-8`), or for an entire TPU slice spanning multiple\nVMs (e.g. `--accelerator-type v3-32`). For more details about how to set up and\nuse TPUs, refer to Cloud docs for\n[single VM setup](https://cloud.google.com/tpu/docs/jax-quickstart-tpu-vm)\nand [pod slice setup](https://cloud.google.com/tpu/docs/jax-quickstart-tpu-vm).\n\nFirst create a single TPUv3-8 VM and connect to it:\n\n```\nZONE=us-central1-a\nTPU_TYPE=v3-8\nVM_NAME=wmt\n\ngcloud alpha compute tpus tpu-vm create $VM_NAME \\\n --zone $ZONE \\\n --accelerator-type $TPU_TYPE \\\n --version v2-alpha\n\ngcloud alpha compute tpus tpu-vm ssh $VM_NAME --zone $ZONE -- \\\n -L 6006:localhost:6006\n```\n\nWhen connected install JAX:\n\n```\npip install \"jax[tpu]>=0.2.21\" \\\n -f https://storage.googleapis.com/jax-releases/libtpu_releases.html\n```\n\nThen install Flax + the example dependencies:\n\n```\ngit clone --depth=1 --branch=main https://github.com/google/flax\ncd flax\npip install -e .\ncd examples/wmt\npip install -r requirements.txt\n```\n\nAnd finally start the training:\n\n```\npython3 main.py --workdir=$HOME/logs/wmt_256 \\\n --config.per_device_batch_size=32 \\\n --jax_backend_target=\"grpc://192.168.0.2:8470\"\n```\n\nNote that you might want to set `TFDS_DATA_DIR` as explained above. You probably\nalso want to start the long-running command above in a `tmux` session and start\nsome monitoring in a separate pane (note that we forwarded port 6006 locally\nabove):\n\n```\ntensorboard --logdir=$HOME/logs\n```\n\nWhen running on pod slices, after creating the TPU VM, there are different ways\nof running the training in SPMD fashion on the hosts connected to the TPUs that\nmake up the slice. We simply send the same installation/execution shell commands\nto all hosts in parallel with the command below. If anything fails it's\nusually a good idea to connect to a single host and execute the commands\ninteractively.\n\nFor convenience, the TPU creation commands are inlined below. Please note that\nwe define `GCS_TFDS_BUCKET` to where your data stands in your cloud bucket.\nAlso `YOUR_BUCKET` is the work directory you are experimenting in. You should\nchoose ZONE based on where your TPU and work directory is. [Here](https://cloud.google.com/tpu/docs/types-zones)\nhas some useful information on which zones you can have different types of TPUs.\n\n```shell\nVM_NAME=wmt\nREPO=https://github.com/google/flax\nBRANCH=main\nWORKDIR=gs://$YOUR_BUCKET/flax/examples/wmt/$(date +%Y%m%d_%H%M)\n\ngcloud alpha compute tpus tpu-vm create $VM_NAME \\\n --zone=$ZONE \\\n --version v2-alpha --accelerator-type v3-32\nFLAGS=\"--config.num_train_steps=$(( 100 * 1000 * 8/32 ))\n--config.warmup_steps=$(( 1000 * 8/32 ))\n--config.checkpoint_every_steps=$(( 10 * 1000 * 8/32 ))\"\n\ngcloud alpha compute tpus tpu-vm ssh $VM_NAME --zone $ZONE \\\n--worker=all --command \"\nset -x\npip install 'jax[tpu]>=0.2.21' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html &&\npip install --user git+$REPO.git &&\n(test -d flax || git clone --depth=1 -b $BRANCH $REPO) &&\ncd flax &&\npip install -e . &&\ncd examples/wmt &&\npip install -r requirements.txt &&\nexport TFDS_DATA_DIR=gs://$GCS_TFDS_BUCKET/datasets &&\npython3 main.py --workdir=$WORKDIR --config=configs/default.py $FLAGS\n\"\n```\nPlease don't forget to disconnect and delete your vm after you are done:\n\n```\ngcloud alpha compute tpus tpu-vm delete $VM_NAME \\\n --zone $ZONE\n```\n" }, { "path": "examples/wmt/bleu.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nr\"\"\"Parallel BLEU score calculation.\n\nThis version of BLEU calculation is derived from the MLPerf transformer reference.\nTries to match SacreBLEU metric reasonably well, but is not identical.\n\nRefs:\n tokenizer at:\n https://github.com/tensorflow/models/blob/master/official/transformer/utils/tokenizer.py\n original preprocessing tokenizer:\n https://github.com/moses-smt/mosesdecoder/blob/master/scripts/generic/mteval-v14.pl#L954-L983\n original t2t code:\n https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/utils/bleu_hook.py\n\nUsage:\n refs = '''food bar brown cow\n blee bloo dog sat\n or please take me out\n '''\n hyps = '''foo bar brown cow\n blee bloo dog sit\n please do take me out\n '''\n bleu_local(refs.split(\"\\n\"), hyps.split(\"\\n\")) # 39.65\n\"\"\"\n\nimport collections\nimport math\nimport re\nimport sys\nimport unicodedata\n\nimport numpy as np\n\n\nclass UnicodeRegex:\n \"\"\"Ad-hoc hack to recognize all punctuation and symbols.\"\"\"\n\n def __init__(self):\n punctuation = self.property_chars(\"P\")\n self.nondigit_punct_re = re.compile(r\"([^\\d])([\" + punctuation + r\"])\")\n self.punct_nondigit_re = re.compile(r\"([\" + punctuation + r\"])([^\\d])\")\n self.symbol_re = re.compile(\"([\" + self.property_chars(\"S\") + \"])\")\n\n def property_chars(self, prefix):\n return \"\".join(\n chr(x)\n for x in range(sys.maxunicode)\n if unicodedata.category(chr(x)).startswith(prefix)\n )\n\n\nuregex = UnicodeRegex()\n\n\ndef bleu_tokenize(string):\n r\"\"\"Tokenize a string following the official BLEU implementation.\n\n See https://github.com/moses-smt/mosesdecoder/'\n 'blob/master/scripts/generic/mteval-v14.pl#L954-L983\n In our case, the input string is expected to be just one line\n and no HTML entities de-escaping is needed.\n So we just tokenize on punctuation and symbols,\n except when a punctuation is preceded and followed by a digit\n (e.g. a comma/dot as a thousand/decimal separator).\n\n Note that a number (e.g. a year) followed by a dot at the end of sentence\n is NOT tokenized, i.e. the dot stays with the number because\n `s/(\\p{P})(\\P{N})/ $1 $2/g` does not match this case (unless we add a\n space after each sentence). However, this error is already in the\n original mteval-v14.pl and we want to be consistent with it.\n\n Args:\n string: the input string\n\n Returns:\n a list of tokens\n \"\"\"\n string = uregex.nondigit_punct_re.sub(r\"\\1 \\2 \", string)\n string = uregex.punct_nondigit_re.sub(r\" \\1 \\2\", string)\n string = uregex.symbol_re.sub(r\" \\1 \", string)\n return string.split()\n\n\ndef _get_ngrams(segment, max_order):\n \"\"\"Extracts all n-grams up to a given maximum order from an input segment.\n\n Args:\n segment: text segment from which n-grams will be extracted.\n max_order: maximum length in tokens of the n-grams returned by this methods.\n\n Returns:\n The Counter containing all n-grams up to max_order in segment\n with a count of how many times each n-gram occurred.\n \"\"\"\n ngram_counts = collections.Counter()\n for order in range(1, max_order + 1):\n for i in range(0, len(segment) - order + 1):\n ngram = tuple(segment[i : i + order])\n ngram_counts[ngram] += 1\n return ngram_counts\n\n\ndef compute_bleu_matches(reference_corpus, translation_corpus, max_order=4):\n \"\"\"Computes BLEU match stats of translations against one or more references.\n\n Args:\n reference_corpus: list of references for each translation. Each reference\n should be tokenized into a list of tokens.\n translation_corpus: list of translations to score. Each translation should\n be tokenized into a list of tokens.\n max_order: Maximum n-gram order to use when computing BLEU score.\n\n Returns:\n Aggregated n-gram stats for BLEU calculation.\n \"\"\"\n reference_length = 0\n translation_length = 0\n bp = 1.0\n geo_mean = 0\n\n matches_by_order = [0] * max_order\n possible_matches_by_order = [0] * max_order\n precisions = []\n\n for references, translations in zip(reference_corpus, translation_corpus):\n reference_length += len(references)\n translation_length += len(translations)\n ref_ngram_counts = _get_ngrams(references, max_order)\n translation_ngram_counts = _get_ngrams(translations, max_order)\n\n overlap = {\n ngram: min(count, translation_ngram_counts[ngram])\n for ngram, count in ref_ngram_counts.items()\n }\n\n for ngram in overlap:\n matches_by_order[len(ngram) - 1] += overlap[ngram]\n for ngram in translation_ngram_counts:\n possible_matches_by_order[len(ngram) - 1] += translation_ngram_counts[\n ngram\n ]\n\n return (\n np.array(matches_by_order),\n np.array(possible_matches_by_order),\n np.array(reference_length),\n np.array(translation_length),\n )\n\n\ndef bleu_partial(ref_lines, hyp_lines, case_sensitive=False):\n \"\"\"Compute n-gram statistics for two lists of references and translations.\"\"\"\n if len(ref_lines) != len(hyp_lines):\n raise ValueError(\n \"Reference and translation lists have different numbers of lines.\"\n )\n if not case_sensitive:\n ref_lines = [x.lower() for x in ref_lines]\n hyp_lines = [x.lower() for x in hyp_lines]\n ref_tokens = [bleu_tokenize(x) for x in ref_lines]\n hyp_tokens = [bleu_tokenize(x) for x in hyp_lines]\n return compute_bleu_matches(ref_tokens, hyp_tokens)\n\n\ndef complete_bleu(\n matches_by_order,\n possible_matches_by_order,\n reference_length,\n translation_length,\n max_order=4,\n use_bp=True,\n):\n \"\"\"Compute BLEU score from aggregated n-gram statistics.\"\"\"\n precisions = [0] * max_order\n smooth = 1.0\n geo_mean = 0.0\n for i in range(0, max_order):\n if possible_matches_by_order[i] > 0:\n precisions[i] = matches_by_order[i] / possible_matches_by_order[i]\n if matches_by_order[i] > 0:\n precisions[i] = matches_by_order[i] / possible_matches_by_order[i]\n else:\n smooth *= 2\n precisions[i] = 1.0 / (smooth * possible_matches_by_order[i])\n else:\n precisions[i] = 0.0\n\n if max(precisions) > 0:\n p_log_sum = sum(math.log(p) for p in precisions if p)\n geo_mean = math.exp(p_log_sum / max_order)\n\n if use_bp:\n if not reference_length:\n bp = 1.0\n else:\n ratio = translation_length / reference_length\n if ratio <= 0.0:\n bp = 0.0\n elif ratio >= 1.0:\n bp = 1.0\n else:\n bp = math.exp(1 - 1.0 / ratio)\n bleu = geo_mean * bp\n return float(bleu) * 100.0\n\n\ndef bleu_local(ref_lines, hyp_lines, case_sensitive=False):\n \"\"\"Compute BLEU for two lists of reference and hypothesis translations.\"\"\"\n stats = bleu_partial(ref_lines, hyp_lines, case_sensitive=case_sensitive)\n return complete_bleu(*stats) * 100\n" }, { "path": "examples/wmt/configs/default.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Default Hyperparameter configuration.\"\"\"\n\nimport ml_collections\n\n\ndef get_config():\n \"\"\"Get the default hyperparameter configuration.\"\"\"\n config = ml_collections.ConfigDict()\n\n # Path to load or store sentencepiece vocab file.\n config.vocab_path = None\n\n # Vocabulary size if `vocab_path` is not given.\n config.vocab_size = 32_000\n\n config.max_corpus_chars = 10**7\n\n # Name of TFDS translation dataset to use.\n config.dataset_name = 'wmt17_translate/de-en'\n\n # Optional name of TFDS translation dataset to use for evaluation.\n config.eval_dataset_name = 'wmt14_translate/de-en'\n config.eval_split = 'test'\n\n # Reverse the direction of translation.\n config.reverse_translation = False\n\n # Per device batch size for training.\n config.per_device_batch_size = 32\n\n # Beam size for inference.\n config.beam_size = 4\n\n config.num_train_steps = 100_000\n\n # Number of steps to take during evaluation.\n config.num_eval_steps = 20\n # Number of steps to generate predictions (used for BLEU score).\n # -1 will use the whole eval dataset.\n config.num_predict_steps = -1\n\n # Base learning rate.\n config.learning_rate = 0.002\n\n # Linear learning rate warmup.\n config.warmup_steps = 1000\n\n # Cross entropy loss label smoothing.\n config.label_smoothing = 0.1\n\n # Decay factor for AdamW style weight decay.\n config.weight_decay = 0.0\n\n # Maximum length cutoff for training examples.\n config.max_target_length = 256\n # Maximum length cutoff for eval examples.\n config.max_eval_target_length = 256\n # Maximum length cutoff for predicted tokens.\n config.max_predict_length = 256\n\n # Inputs and targets share embedding.\n config.share_embeddings = True\n\n # Final logit transform uses embedding matrix transpose.\n config.logits_via_embedding = True\n\n # Number of transformer layers.\n config.num_layers = 6\n\n # Size of query/key/value for attention.\n config.qkv_dim = 1024\n # Size of embeddings.\n config.emb_dim = 1024\n # Size of the MLP.\n config.mlp_dim = 4096\n\n # Number of attention heads.\n config.num_heads = 16\n\n # Dropout rate.\n config.dropout_rate = 0.1\n\n # Attention dropout rate.\n config.attention_dropout_rate = 0.1\n\n # Whether to save model checkpoints.\n config.save_checkpoints = True\n # Whether to restore from existing model checkpoints.\n config.restore_checkpoints = True\n\n # Save a checkpoint every these number of steps.\n config.checkpoint_every_steps = 10_000\n # Frequency of eval during training, e.g. every 1000 steps.\n config.eval_every_steps = 1_000\n\n # Use float16/bfloat16 (GPU/TPU) mixed precision training instead of float32.\n config.use_mixed_precision = True\n\n # Integer for PRNG random seed.\n config.seed = 0\n\n return config\n\n\ndef metrics():\n return [\n 'train_loss',\n 'eval_loss',\n 'bleu',\n 'eval_accuracy',\n 'train_accuracy',\n 'uptime',\n 'steps_per_sec',\n 'train_learning_rate',\n ]\n" }, { "path": "examples/wmt/decode.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Fast decoding routines for inference from a trained model.\"\"\"\n\nimport typing\n\nimport flax\nimport jax\nfrom jax import lax\nimport jax.numpy as jnp\nimport numpy as np\n\n\n# Constants\n# We assume the default End-of-Sentence token id is 2 (SentencePiece).\nEOS_ID = 2\n# \"Effective negative infinity\" constant for masking in beam search.\nNEG_INF = np.array(-1.0e7)\n\n\ndef brevity_penalty(alpha, length):\n \"\"\"Brevity penalty function for beam search penalizing short sequences.\n\n Args:\n alpha: float: brevity-penalty scaling parameter.\n length: int: length of considered sequence.\n\n Returns:\n Brevity penalty score as jax scalar.\n \"\"\"\n return jnp.power(((5.0 + length) / 6.0), alpha)\n\n\n# Beam handling utility functions:\n\n\ndef add_beam_dim(x, beam_size):\n \"\"\"Creates new beam dimension in non-scalar array and tiles into it.\"\"\"\n if x.ndim == 0: # ignore scalars (e.g. cache index)\n return x\n x = jnp.expand_dims(x, axis=1)\n tile_dims = [1] * x.ndim\n tile_dims[1] = beam_size\n return jnp.tile(x, tile_dims)\n\n\ndef flatten_beam_dim(x):\n \"\"\"Flattens the first two dimensions of a non-scalar array.\"\"\"\n if x.ndim == 0: # ignore scalars (e.g. cache index)\n return x\n return x.reshape((x.shape[0] * x.shape[1],) + x.shape[2:])\n\n\ndef unflatten_beam_dim(x, batch_size, beam_size):\n \"\"\"Unflattens the first, flat batch*beam dimension of a non-scalar array.\"\"\"\n if x.ndim == 0: # ignore scalars (e.g. cache index)\n return x\n assert batch_size * beam_size == x.shape[0]\n return x.reshape((batch_size, beam_size) + x.shape[1:])\n\n\ndef flat_batch_beam_expand(x, beam_size):\n \"\"\"Expands each batch item by beam_size in batch_dimension.\"\"\"\n return flatten_beam_dim(add_beam_dim(x, beam_size))\n\n\ndef gather_beams(nested, beam_indices, batch_size, new_beam_size):\n \"\"\"Gathers the beam slices indexed by beam_indices into new beam array.\n\n Args:\n nested: pytree of arrays or scalars (the latter ignored).\n beam_indices: array of beam_indices\n batch_size: int: size of batch.\n new_beam_size: int: size of _new_ beam dimension.\n\n Returns:\n New pytree with new beam arrays.\n [batch_size, old_beam_size, ...] --> [batch_size, new_beam_size, ...]\n \"\"\"\n batch_indices = jnp.reshape(\n jnp.arange(batch_size * new_beam_size) // new_beam_size,\n (batch_size, new_beam_size),\n )\n\n def gather_fn(x):\n if x.ndim == 0: # ignore scalars (e.g. cache index)\n return x\n else:\n return x[batch_indices, beam_indices]\n\n return jax.tree_util.tree_map(gather_fn, nested)\n\n\ndef gather_topk_beams(nested, score_or_log_prob, batch_size, new_beam_size):\n \"\"\"Gathers the top-k beam slices given by score_or_log_prob array.\n\n Args:\n nested: pytree of arrays or scalars (the latter ignored).\n score_or_log_prob: [batch_size, old_beam_size] array of values to sort by\n for top-k selection of beam slices.\n batch_size: int: size of batch.\n new_beam_size: int: size of _new_ top-k selected beam dimension\n\n Returns:\n New pytree with new beam arrays containing top k new_beam_size slices.\n [batch_size, old_beam_size, ...] --> [batch_size, new_beam_size, ...]\n \"\"\"\n _, topk_indices = lax.top_k(score_or_log_prob, k=new_beam_size)\n topk_indices = jnp.flip(topk_indices, axis=1)\n return gather_beams(nested, topk_indices, batch_size, new_beam_size)\n\n\n# Beam search state:\n\n\n@flax.struct.dataclass\nclass BeamState:\n \"\"\"Holds beam search state data.\"\"\"\n\n # The position of the decoding loop in the length dimension.\n cur_index: jax.Array # scalar int32: current decoded length index\n # The active sequence log probabilities and finished sequence scores.\n live_logprobs: jax.Array # float32: [batch_size, beam_size]\n finished_scores: jax.Array # float32: [batch_size, beam_size]\n # The current active-beam-searching and finished sequences.\n live_seqs: jax.Array # int32: [batch_size, beam_size, max_decode_len]\n finished_seqs: jax.Array # int32: [batch_size, beam_size,\n # max_decode_len]\n # Records which of the 'finished_seqs' is occupied and not a filler slot.\n finished_flags: jax.Array # bool: [batch_size, beam_size]\n # The current state of the autoregressive decoding caches.\n cache: typing.Any # Any pytree of arrays, e.g. flax attention Cache object\n\n\ndef beam_init(batch_size, beam_size, max_decode_len, cache):\n \"\"\"Initializes the beam search state data structure.\"\"\"\n cur_index0 = jnp.array(0)\n live_logprobs0 = jnp.tile(\n jnp.array([0.0] + [NEG_INF] * (beam_size - 1)), [batch_size, 1]\n )\n finished_scores0 = jnp.ones((batch_size, beam_size)) * NEG_INF\n live_seqs0 = jnp.zeros((batch_size, beam_size, max_decode_len), jnp.int32)\n finished_seqs0 = jnp.zeros((batch_size, beam_size, max_decode_len), jnp.int32)\n finished_flags0 = jnp.zeros((batch_size, beam_size), jnp.bool_)\n # add beam dimension to attention cache pytree elements\n beam_cache0 = jax.tree_util.tree_map(\n lambda x: add_beam_dim(x, beam_size), cache\n )\n return BeamState(\n cur_index=cur_index0,\n live_logprobs=live_logprobs0,\n finished_scores=finished_scores0,\n live_seqs=live_seqs0,\n finished_seqs=finished_seqs0,\n finished_flags=finished_flags0,\n cache=beam_cache0,\n )\n\n\n# Beam search routine:\n\n\ndef beam_search(\n inputs,\n cache,\n tokens_to_logits,\n beam_size=4,\n alpha=0.6,\n eos_id=EOS_ID,\n max_decode_len=None,\n):\n \"\"\"Beam search for transformer machine translation.\n\n Args:\n inputs: array: [batch_size, length] int32 sequence of tokens.\n cache: flax attention cache.\n tokens_to_logits: fast autoregressive decoder function taking single token\n slices and cache and returning next-token logits and updated cache.\n beam_size: int: number of beams to use in beam search.\n alpha: float: scaling factor for brevity penalty.\n eos_id: int: id of end-of-sentence token for target vocabulary.\n max_decode_len: int: maximum length of decoded translations.\n\n Returns:\n Tuple of:\n [batch_size, beam_size, max_decode_len] top-scoring sequences\n [batch_size, beam_size] beam-search scores.\n \"\"\"\n # We liberally annotate shape information for clarity below.\n\n batch_size = inputs.shape[0]\n if max_decode_len is None:\n max_decode_len = inputs.shape[1]\n end_marker = jnp.array(eos_id)\n\n # initialize beam search state\n beam_search_init_state = beam_init(\n batch_size, beam_size, max_decode_len, cache\n )\n\n def beam_search_loop_cond_fn(state):\n \"\"\"Beam search loop termination condition.\"\"\"\n # Have we reached max decoding length?\n not_at_end = state.cur_index < max_decode_len - 1\n\n # Is no further progress in the beam search possible?\n # Get the best possible scores from alive sequences.\n min_brevity_penalty = brevity_penalty(alpha, max_decode_len)\n best_live_scores = state.live_logprobs[:, -1:] / min_brevity_penalty\n # Get the worst scores from finished sequences.\n worst_finished_scores = jnp.min(\n state.finished_scores, axis=1, keepdims=True\n )\n # Mask out scores from slots without any actual finished sequences.\n worst_finished_scores = jnp.where(\n state.finished_flags, worst_finished_scores, NEG_INF\n )\n # If no best possible live score is better than current worst finished\n # scores, the search cannot improve the finished set further.\n search_terminated = jnp.all(worst_finished_scores > best_live_scores)\n\n # If we're not at the max decode length, and the search hasn't terminated,\n # continue looping.\n return not_at_end & (~search_terminated)\n\n def beam_search_loop_body_fn(state):\n \"\"\"Beam search loop state update function.\"\"\"\n # Collect the current position slice along length to feed the fast\n # autoregressive decoder model. Flatten the beam dimension into batch\n # dimension for feeding into the model.\n # --> [batch * beam, 1]\n flat_ids = flatten_beam_dim(\n lax.dynamic_slice(\n state.live_seqs, (0, 0, state.cur_index), (batch_size, beam_size, 1)\n )\n )\n # Flatten beam dimension into batch to be compatible with model.\n # {[batch, beam, ...], ...} --> {[batch * beam, ...], ...}\n flat_cache = jax.tree_util.tree_map(flatten_beam_dim, state.cache)\n\n # Call fast-decoder model on current tokens to get next-position logits.\n # --> [batch * beam, vocab]\n flat_logits, new_flat_cache = tokens_to_logits(flat_ids, flat_cache)\n\n # unflatten beam dimension\n # [batch * beam, vocab] --> [batch, beam, vocab]\n logits = unflatten_beam_dim(flat_logits, batch_size, beam_size)\n # Unflatten beam dimension in attention cache arrays\n # {[batch * beam, ...], ...} --> {[batch, beam, ...], ...}\n new_cache = jax.tree_util.tree_map(\n lambda x: unflatten_beam_dim(x, batch_size, beam_size), new_flat_cache\n )\n\n # Gather log probabilities from logits\n candidate_log_probs = jax.nn.log_softmax(logits)\n # Add new logprobs to existing prefix logprobs.\n # --> [batch, beam, vocab]\n log_probs = candidate_log_probs + jnp.expand_dims(\n state.live_logprobs, axis=2\n )\n\n # We'll need the vocab size, gather it from the log probability dimension.\n vocab_size = log_probs.shape[2]\n\n # Each item in batch has beam_size * vocab_size candidate sequences.\n # For each item, get the top 2*k candidates with the highest log-\n # probabilities. We gather the top 2*K beams here so that even if the best\n # K sequences reach EOS simultaneously, we have another K sequences\n # remaining to continue the live beam search.\n beams_to_keep = 2 * beam_size\n # Flatten beam and vocab dimensions.\n flat_log_probs = log_probs.reshape((batch_size, beam_size * vocab_size))\n # Gather the top 2*K scores from _all_ beams.\n # --> [batch, 2*beams], [batch, 2*beams]\n topk_log_probs, topk_indices = lax.top_k(flat_log_probs, k=beams_to_keep)\n # Recover the beam index by floor division.\n topk_beam_indices = topk_indices // vocab_size\n # Gather 2*k top beams.\n # --> [batch, 2*beams, length]\n topk_seq = gather_beams(\n state.live_seqs, topk_beam_indices, batch_size, beams_to_keep\n )\n\n # Append the most probable 2*K token IDs to the top 2*K sequences\n # Recover token id by modulo division and expand Id array for broadcasting.\n # --> [batch, 2*beams, 1]\n topk_ids = jnp.expand_dims(topk_indices % vocab_size, axis=2)\n # Update sequences for the 2*K top-k new sequences.\n # --> [batch, 2*beams, length]\n topk_seq = lax.dynamic_update_slice(\n topk_seq, topk_ids, (0, 0, state.cur_index + 1)\n )\n\n # Update LIVE (in-progress) sequences:\n # Did any of these sequences reach an end marker?\n # --> [batch, 2*beams]\n newly_finished = topk_seq[:, :, state.cur_index + 1] == end_marker\n # To prevent these newly finished sequences from being added to the LIVE\n # set of active beam search sequences, set their log probs to a very large\n # negative value.\n new_log_probs = topk_log_probs + newly_finished * NEG_INF\n # Determine the top k beam indices (from top 2*k beams) from log probs.\n # --> [batch, beams]\n _, new_topk_indices = lax.top_k(new_log_probs, k=beam_size)\n new_topk_indices = jnp.flip(new_topk_indices, axis=1)\n # Gather the top k beams (from top 2*k beams).\n # --> [batch, beams, length], [batch, beams]\n top_alive_seq, top_alive_log_probs = gather_beams(\n [topk_seq, new_log_probs], new_topk_indices, batch_size, beam_size\n )\n\n # Determine the top k beam indices from the original set of all beams.\n # --> [batch, beams]\n top_alive_indices = gather_beams(\n topk_beam_indices, new_topk_indices, batch_size, beam_size\n )\n # With these, gather the top k beam-associated caches.\n # --> {[batch, beams, ...], ...}\n top_alive_cache = gather_beams(\n new_cache, top_alive_indices, batch_size, beam_size\n )\n\n # Update FINISHED (reached end of sentence) sequences:\n # Calculate new seq scores from log probabilities.\n new_scores = topk_log_probs / brevity_penalty(alpha, state.cur_index + 1)\n # Mask out the still unfinished sequences by adding large negative value.\n # --> [batch, 2*beams]\n new_scores += (~newly_finished) * NEG_INF\n\n # Combine sequences, scores, and flags along the beam dimension and compare\n # new finished sequence scores to existing finished scores and select the\n # best from the new set of beams.\n finished_seqs = jnp.concatenate( # --> [batch, 3*beams, length]\n [state.finished_seqs, topk_seq], axis=1\n )\n finished_scores = jnp.concatenate( # --> [batch, 3*beams]\n [state.finished_scores, new_scores], axis=1\n )\n finished_flags = jnp.concatenate( # --> [batch, 3*beams]\n [state.finished_flags, newly_finished], axis=1\n )\n # --> [batch, beams, length], [batch, beams], [batch, beams]\n (\n top_finished_seq,\n top_finished_scores,\n top_finished_flags,\n ) = gather_topk_beams(\n [finished_seqs, finished_scores, finished_flags],\n finished_scores,\n batch_size,\n beam_size,\n )\n\n return BeamState(\n cur_index=state.cur_index + 1,\n live_logprobs=top_alive_log_probs,\n finished_scores=top_finished_scores,\n live_seqs=top_alive_seq,\n finished_seqs=top_finished_seq,\n finished_flags=top_finished_flags,\n cache=top_alive_cache,\n )\n\n # Run while loop and get final beam search state.\n final_state = lax.while_loop(\n beam_search_loop_cond_fn, beam_search_loop_body_fn, beam_search_init_state\n )\n\n # Account for the edge-case where there are no finished sequences for a\n # particular batch item. If so, return live sequences for that batch item.\n # --> [batch]\n none_finished = jnp.any(final_state.finished_flags, axis=1)\n # --> [batch, beams, length]\n finished_seqs = jnp.where(\n none_finished[:, None, None],\n final_state.finished_seqs,\n final_state.live_seqs,\n )\n # --> [batch, beams]\n finished_scores = jnp.where(\n none_finished[:, None],\n final_state.finished_scores,\n final_state.live_logprobs,\n )\n\n return finished_seqs, finished_scores\n" }, { "path": "examples/wmt/input_pipeline.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Input pipeline for a WMT dataset.\"\"\"\n\nimport os\nfrom typing import Dict, Optional, List, Union\n\nfrom clu import deterministic_data\nimport ml_collections\nimport tensorflow as tf\nimport tensorflow_datasets as tfds\n\nimport tokenizer\n\n\nAUTOTUNE = tf.data.AUTOTUNE\nFeatures = dict[str, tf.Tensor]\n\n\nclass NormalizeFeatureNamesOp:\n \"\"\"Normalizes feature names to 'inputs' and 'targets'.\"\"\"\n\n def __init__(self, ds_info: tfds.core.DatasetInfo, reverse_translation: bool):\n self.input_lang, self.target_lang = ds_info.supervised_keys\n if reverse_translation:\n self.input_lang, self.target_lang = self.target_lang, self.input_lang\n\n def __call__(self, features: Features) -> Features:\n features['inputs'] = features.pop(self.input_lang)\n features['targets'] = features.pop(self.target_lang)\n return features\n\n\ndef get_raw_dataset(\n dataset_builder: tfds.core.DatasetBuilder,\n split: str,\n *,\n reverse_translation: bool = False,\n) -> tf.data.Dataset:\n \"\"\"Loads a raw WMT dataset and normalizes feature keys.\n\n Args:\n dataset_builder: TFDS dataset builder that can build `slit`.\n split: Split to use. This must be the full split. We shard the split across\n multiple hosts and currently don't support sharding subsplits.\n reverse_translation: bool: whether to reverse the translation direction.\n e.g. for 'de-en' this translates from english to german.\n\n Returns:\n Dataset with source and target language features mapped to 'inputs' and\n 'targets'.\n \"\"\"\n num_examples = dataset_builder.info.splits[split].num_examples\n per_host_split = deterministic_data.get_read_instruction_for_host(\n split, num_examples, drop_remainder=False\n )\n ds = dataset_builder.as_dataset(split=per_host_split, shuffle_files=False)\n ds = ds.map(\n NormalizeFeatureNamesOp(\n dataset_builder.info, reverse_translation=reverse_translation\n ),\n num_parallel_calls=AUTOTUNE,\n )\n return ds\n\n\ndef pack_dataset(\n dataset: tf.data.Dataset,\n key2length: int | dict[str, int],\n keys: list[str] | None = None,\n) -> tf.data.Dataset:\n \"\"\"Creates a 'packed' version of a dataset on-the-fly.\n\n Adapted from the mesh-tf implementation.\n\n This is meant to replace the irritation of having to create a separate\n \"packed\" version of a dataset to train efficiently on TPU.\n Each example in the output dataset represents several examples in the\n input dataset.\n For each key in the input dataset, two additional keys are created:\n _segmentation: an int32 tensor identifying the parts\n representing the original example.\n _position: an int32 tensor identifying the position within the original\n example.\n Example:\n Two input examples get combined to form an output example.\n The input examples are:\n {\"inputs\": [8, 7, 1, 0], \"targets\":[4, 1, 0]}\n {\"inputs\": [2, 3, 4, 1], \"targets\":[5, 6, 1]}\n The output example is:\n {\n \"inputs\": [8, 7, 1, 2, 3, 4, 1, 0, 0, 0]\n \"inputs_segmentation\": [1, 1, 1, 2, 2, 2, 2, 0, 0, 0]\n \"inputs_position\": [0, 1, 2, 0, 1, 2, 3, 0, 0, 0]\n \"targets\": [4, 1, 5, 6, 1, 0, 0, 0, 0, 0]\n \"targets_segmentation\": [1, 1, 2, 2, 2, 0, 0, 0, 0, 0]\n \"targets_position\": [0, 1, 0, 1, 2, 0, 0, 0, 0, 0]\n }\n 0 represents padding in both the inputs and the outputs.\n Sequences in the incoming examples are truncated to length \"length\", and the\n sequences in the output examples all have fixed (padded) length \"length\".\n\n Args:\n dataset: a tf.data.Dataset\n key2length: an integer, or a dict from feature-key to integer\n keys: a list of strings (e.g. [\"inputs\", \"targets\"])\n\n Returns:\n a tf.data.Dataset\n \"\"\"\n shapes = tf.nest.map_structure(lambda spec: spec.shape, dataset.element_spec)\n if keys is None:\n keys = list(shapes.keys())\n for k in keys:\n if k not in shapes:\n raise ValueError(\n 'Key %s not found in dataset. Available keys are %s'\n % (k, shapes.keys())\n )\n if not shapes[k].is_compatible_with(tf.TensorShape([None])): # type: ignore[wrong-arg-types]\n raise ValueError('Tensors to be packed must be one-dimensional.')\n # make sure that the length dictionary contains all keys as well as the\n # keys suffixed by \"_segmentation\" and \"_position\"\n if isinstance(key2length, int):\n key2length = {k: key2length for k in keys}\n for k in keys:\n for suffix in ['_segmentation', '_position']:\n key2length[k + suffix] = key2length[k]\n\n # trim to length\n dataset = dataset.map(\n lambda x: {k: x[k][: key2length[k]] for k in keys},\n num_parallel_calls=AUTOTUNE,\n )\n # Setting batch_size=length ensures that the concatenated sequences (if they\n # have length >=1) are sufficient to fill at least one packed example.\n batch_size = max(key2length.values())\n dataset = dataset.padded_batch(\n batch_size, padded_shapes={k: [-1] for k in keys}\n )\n dataset = _pack_with_tf_ops(dataset, keys, key2length)\n\n # Set the Tensor shapes correctly since they get lost in the process.\n def my_fn(x):\n return {k: tf.reshape(v, [key2length[k]]) for k, v in x.items()}\n\n return dataset.map(my_fn, num_parallel_calls=AUTOTUNE)\n\n\ndef _pack_with_tf_ops(\n dataset: tf.data.Dataset, keys: list[str], key2length: dict[str, int]\n) -> tf.data.Dataset:\n \"\"\"Helper-function for packing a dataset which has already been batched.\n\n Helper for pack_dataset() Uses tf.while_loop.\n\n Args:\n dataset: a dataset containing padded batches of examples.\n keys: a list of strings\n key2length: an dict from feature-key to integer\n\n Returns:\n a dataset.\n \"\"\"\n empty_example = {}\n for k in keys:\n empty_example[k] = tf.zeros([0], dtype=tf.int32)\n empty_example[k + '_position'] = tf.zeros([0], dtype=tf.int32)\n keys_etc = empty_example.keys()\n\n def write_packed_example(partial, outputs):\n new_partial = empty_example.copy()\n new_outputs = {}\n for k in keys_etc:\n new_outputs[k] = outputs[k].write(\n outputs[k].size(),\n tf.pad(partial[k], [[0, key2length[k] - tf.size(partial[k])]]),\n )\n return new_partial, new_outputs\n\n def map_fn(x):\n \"\"\"Internal function to flat_map over.\n\n Consumes a batch of input examples and produces a variable number of output\n examples.\n Args:\n x: a single example\n\n Returns:\n a tf.data.Dataset\n \"\"\"\n partial = empty_example.copy()\n i = tf.zeros([], dtype=tf.int32)\n dynamic_batch_size = tf.shape(x[keys[0]])[0]\n outputs = {}\n for k in keys:\n outputs[k] = tf.TensorArray(\n tf.int32, size=0, dynamic_size=True, element_shape=[key2length[k]]\n )\n outputs[k + '_position'] = tf.TensorArray(\n tf.int32, size=0, dynamic_size=True, element_shape=[key2length[k]]\n )\n\n def body_fn(i, partial, outputs):\n \"\"\"Body function for while_loop.\n\n Args:\n i: integer scalar\n partial: dictionary of Tensor (partially-constructed example)\n outputs: dictionary of TensorArray\n\n Returns:\n A triple containing the new values of the inputs.\n \"\"\"\n can_append = True\n one_example = {}\n for k in keys:\n val = tf.cast(x[k][i], tf.int32)\n val = val[: tf.reduce_sum(tf.cast(tf.not_equal(val, 0), tf.int32))]\n one_example[k] = val\n for k in keys:\n can_append = tf.logical_and(\n can_append,\n tf.less_equal(\n tf.size(partial[k]) + tf.size(one_example[k]), key2length[k]\n ),\n )\n\n def false_fn():\n return write_packed_example(partial, outputs)\n\n def true_fn():\n return partial, outputs\n\n partial, outputs = tf.cond(can_append, true_fn, false_fn)\n new_partial = {}\n for k in keys:\n new_seq = one_example[k][: key2length[k]]\n new_seq_len = tf.size(new_seq)\n new_partial[k] = tf.concat([partial[k], new_seq], 0)\n new_partial[k + '_position'] = tf.concat(\n [partial[k + '_position'], tf.range(new_seq_len)], 0\n )\n partial = new_partial\n return i + 1, partial, outputs\n\n # For loop over all examples in the batch.\n i, partial, outputs = tf.while_loop(\n cond=lambda *_: True,\n body=body_fn,\n loop_vars=(i, partial, outputs),\n shape_invariants=(\n tf.TensorShape([]), # type: ignore[wrong-arg-types]\n {k: tf.TensorShape([None]) for k in keys_etc}, # type: ignore[wrong-arg-types]\n {k: tf.TensorShape(None) for k in keys_etc}, # type: ignore[wrong-arg-types]\n ),\n maximum_iterations=dynamic_batch_size,\n )\n _, outputs = write_packed_example(partial, outputs)\n packed = {k: outputs[k].stack() for k in keys_etc}\n for k in keys:\n packed[k + '_segmentation'] = tf.cumsum(\n tf.cast(tf.equal(packed[k + '_position'], 0), tf.int32), axis=1\n ) * tf.cast(tf.not_equal(packed[k], 0), tf.int32)\n return packed\n\n dataset = dataset.map(map_fn, num_parallel_calls=AUTOTUNE)\n return dataset.unbatch()\n\n\n# -----------------------------------------------------------------------------\n# Main dataset prep routines.\n# -----------------------------------------------------------------------------\ndef preprocess_wmt_data(\n dataset,\n shuffle: bool,\n num_epochs: int | None = 1,\n pack_examples: bool = True,\n shuffle_buffer_size: int = 1024,\n max_length: int = 512,\n batch_size: int = 256,\n drop_remainder: bool = True,\n prefetch_size: int = AUTOTUNE,\n):\n \"\"\"Shuffle and batch/pack the given dataset.\"\"\"\n\n def length_filter(max_len):\n def filter_fn(x):\n source, target = x['inputs'], x['targets']\n l = tf.maximum(tf.shape(source)[0], tf.shape(target)[0])\n return tf.less(l, max_len + 1)\n\n return filter_fn\n\n if max_length > 0:\n dataset = dataset.filter(length_filter(max_length))\n\n if shuffle:\n dataset = dataset.shuffle(shuffle_buffer_size)\n dataset = dataset.repeat(num_epochs)\n\n if pack_examples:\n dataset = pack_dataset(dataset, max_length)\n dataset = dataset.batch(batch_size, drop_remainder=drop_remainder)\n else: # simple (static-shape) padded batching\n dataset = dataset.padded_batch(\n batch_size,\n padded_shapes={'inputs': max_length, 'targets': max_length},\n padding_values={'inputs': 0, 'targets': 0},\n drop_remainder=drop_remainder,\n )\n\n if prefetch_size:\n dataset = dataset.prefetch(prefetch_size)\n\n return dataset\n\n\ndef get_wmt_datasets(\n config: ml_collections.ConfigDict,\n *,\n n_devices: int,\n reverse_translation: bool = True,\n vocab_path: str | None = None,\n):\n \"\"\"Load and return dataset of batched examples for use during training.\"\"\"\n if vocab_path is None:\n vocab_path = os.path.expanduser('~/wmt_sentencepiece_model')\n\n train_ds_builder = tfds.builder(config.dataset_name)\n train_data = get_raw_dataset(\n train_ds_builder, 'train', reverse_translation=reverse_translation\n )\n\n if config.eval_dataset_name:\n eval_ds_builder = tfds.builder(config.eval_dataset_name)\n else:\n eval_ds_builder = train_ds_builder\n eval_data = get_raw_dataset(\n eval_ds_builder,\n config.eval_split,\n reverse_translation=reverse_translation,\n )\n\n # Tokenize data.\n sp_tokenizer = tokenizer.load_or_train_tokenizer(\n train_data,\n vocab_path=vocab_path,\n vocab_size=config.vocab_size,\n max_corpus_chars=config.max_corpus_chars,\n )\n train_data = train_data.map(\n tokenizer.TokenizeOp(sp_tokenizer), num_parallel_calls=AUTOTUNE\n )\n eval_data = eval_data.map(\n tokenizer.TokenizeOp(sp_tokenizer), num_parallel_calls=AUTOTUNE\n )\n\n batch_size = config.per_device_batch_size * n_devices\n\n train_ds = preprocess_wmt_data(\n train_data,\n shuffle=True,\n num_epochs=None,\n pack_examples=True,\n batch_size=batch_size,\n max_length=config.max_target_length,\n )\n\n eval_ds = preprocess_wmt_data(\n eval_data,\n shuffle=False,\n pack_examples=False,\n batch_size=batch_size,\n max_length=config.max_eval_target_length,\n )\n\n predict_ds = preprocess_wmt_data(\n eval_data,\n shuffle=False,\n pack_examples=False,\n batch_size=batch_size,\n max_length=config.max_predict_length,\n drop_remainder=False,\n )\n\n return train_ds, eval_ds, predict_ds, sp_tokenizer\n" }, { "path": "examples/wmt/input_pipeline_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nimport pathlib\nimport sys\nimport tempfile\n\nfrom absl.testing import absltest\nimport tensorflow_datasets as tfds\n\nfrom configs import default\nimport input_pipeline\n\n\n# We just use different values here to verify that the input pipeline uses the\n# the correct value for the 3 different datasets.\n_TARGET_LENGTH = 32\n_EVAL_TARGET_LENGTH = 48\n_PREDICT_TARGET_LENGTH = 64\n\n\nclass InputPipelineTest(absltest.TestCase):\n\n def setUp(self):\n super().setUp()\n if sys.version_info >= (3, 13):\n self.skipTest('Test (and tensorflow-text) does not suport Python 3.13+')\n self.train_ds, self.eval_ds, self.predict_ds = self._get_datasets()\n\n def _get_datasets(self):\n config = default.get_config()\n config.per_device_batch_size = 1\n config.vocab_size = 32\n config.max_corpus_chars = 1000\n config.max_target_length = _TARGET_LENGTH\n config.max_eval_target_length = _EVAL_TARGET_LENGTH\n config.max_predict_length = _PREDICT_TARGET_LENGTH\n\n vocab_path = os.path.join(tempfile.mkdtemp(), 'sentencepiece_model')\n\n # Go two directories up to the root of the flax directory.\n flax_root_dir = pathlib.Path(__file__).parents[2]\n data_dir = str(flax_root_dir) + '/.tfds/metadata' # pylint: disable=unused-variable\n\n with tfds.testing.mock_data(num_examples=128, data_dir=data_dir):\n train_ds, eval_ds, predict_ds, _ = input_pipeline.get_wmt_datasets(\n n_devices=2, config=config, vocab_path=vocab_path\n )\n return train_ds, eval_ds, predict_ds\n\n def test_train_ds(self):\n expected_shape = [2, _TARGET_LENGTH] # 2 devices.\n # For training we pack multiple short examples in one example.\n # *_position and *_segmentation indicate the boundaries.\n for batch in self.train_ds.take(3):\n self.assertEqual(\n {k: v.shape.as_list() for k, v in batch.items()},\n {\n 'inputs': expected_shape,\n 'inputs_position': expected_shape,\n 'inputs_segmentation': expected_shape,\n 'targets': expected_shape,\n 'targets_position': expected_shape,\n 'targets_segmentation': expected_shape,\n },\n )\n\n def test_eval_ds(self):\n expected_shape = [2, _EVAL_TARGET_LENGTH] # 2 devices.\n for batch in self.eval_ds.take(3):\n self.assertEqual(\n {k: v.shape.as_list() for k, v in batch.items()},\n {\n 'inputs': expected_shape,\n 'targets': expected_shape,\n },\n )\n\n def test_predict_ds(self):\n expected_shape = [2, _PREDICT_TARGET_LENGTH] # 2 devices.\n for batch in self.predict_ds.take(3):\n self.assertEqual(\n {k: v.shape.as_list() for k, v in batch.items()},\n {\n 'inputs': expected_shape,\n 'targets': expected_shape,\n },\n )\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "examples/wmt/main.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Main file for running the WMT example.\n\nThis file is intentionally kept short. The majority for logic is in libraries\nthat can be easily tested and imported in Colab.\n\"\"\"\n\nfrom absl import app\nfrom absl import flags\nfrom absl import logging\nfrom clu import platform\nimport jax\nfrom ml_collections import config_flags\nimport tensorflow as tf\n\nimport train\n\n\nFLAGS = flags.FLAGS\n\nflags.DEFINE_string('workdir', None, 'Directory to store model data.')\nconfig_flags.DEFINE_config_file(\n 'config',\n 'configs/default.py',\n 'File path to the training hyperparameter configuration.',\n lock_config=True,\n)\nflags.mark_flags_as_required(['config', 'workdir'])\n\n\ndef main(argv):\n if len(argv) > 1:\n raise app.UsageError('Too many command-line arguments.')\n\n # Hide any GPUs from TensorFlow. Otherwise TF might reserve memory and make\n # it unavailable to JAX.\n tf.config.experimental.set_visible_devices([], 'GPU')\n\n logging.info('JAX process: %d / %d', jax.process_index(), jax.process_count())\n logging.info('JAX local devices: %r', jax.local_devices())\n\n # Add a note so that we can tell which task is which JAX host.\n # (Depending on the platform task 0 is not guaranteed to be host 0)\n platform.work_unit().set_task_status(\n f'process_index: {jax.process_index()}, '\n f'process_count: {jax.process_count()}'\n )\n platform.work_unit().create_artifact(\n platform.ArtifactType.DIRECTORY, FLAGS.workdir, 'workdir'\n )\n\n train.train_and_evaluate(FLAGS.config, FLAGS.workdir)\n\n\nif __name__ == '__main__':\n jax.config.config_with_absl()\n app.run(main)\n" }, { "path": "examples/wmt/models.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Transformer-based machine translation model.\"\"\"\n\n# pylint: disable=attribute-defined-outside-init,g-bare-generic\n# See issue #620.\n# pytype: disable=wrong-arg-count\n# pytype: disable=wrong-keyword-args\n# pytype: disable=attribute-error\n\nfrom collections.abc import Callable\nfrom typing import Any\n\nfrom flax import linen as nn\nfrom flax import struct\nfrom jax import lax\nimport jax.numpy as jnp\nimport numpy as np\n\n\n@struct.dataclass\nclass TransformerConfig:\n \"\"\"Global hyperparameters used to minimize obnoxious kwarg plumbing.\"\"\"\n\n vocab_size: int\n output_vocab_size: int\n share_embeddings: bool = False\n logits_via_embedding: bool = False\n dtype: Any = jnp.float32\n emb_dim: int = 512\n num_heads: int = 8\n num_layers: int = 6\n qkv_dim: int = 512\n mlp_dim: int = 2048\n max_len: int = 2048\n dropout_rate: float = 0.1\n attention_dropout_rate: float = 0.1\n deterministic: bool = False\n decode: bool = False\n kernel_init: Callable = nn.initializers.xavier_uniform()\n bias_init: Callable = nn.initializers.normal(stddev=1e-6)\n posemb_init: Callable | None = None\n\n\ndef shift_right(x, axis=1):\n \"\"\"Shift the input to the right by padding on axis 1.\"\"\"\n pad_widths = [(0, 0)] * len(x.shape)\n pad_widths[axis] = (1, 0)\n padded = jnp.pad(\n x, pad_widths, mode='constant', constant_values=x.dtype.type(0)\n )\n return padded[:, :-1]\n\n\ndef sinusoidal_init(max_len=2048, min_scale=1.0, max_scale=10000.0):\n \"\"\"1D Sinusoidal Position Embedding Initializer.\n\n Args:\n max_len: maximum possible length for the input.\n min_scale: float: minimum frequency-scale in sine grating.\n max_scale: float: maximum frequency-scale in sine grating.\n\n Returns:\n output: init function returning `(1, max_len, d_feature)`\n \"\"\"\n\n def init(key, shape, dtype=np.float32):\n \"\"\"Sinusoidal init.\"\"\"\n del key, dtype\n d_feature = shape[-1]\n pe = np.zeros((max_len, d_feature), dtype=np.float32)\n position = np.arange(0, max_len)[:, np.newaxis]\n scale_factor = -np.log(max_scale / min_scale) / (d_feature // 2 - 1)\n div_term = min_scale * np.exp(np.arange(0, d_feature // 2) * scale_factor)\n pe[:, : d_feature // 2] = np.sin(position * div_term)\n pe[:, d_feature // 2 : 2 * (d_feature // 2)] = np.cos(position * div_term)\n pe = pe[np.newaxis, :, :] # [1, max_len, d_feature]\n return jnp.array(pe)\n\n return init\n\n\nclass AddPositionEmbs(nn.Module):\n \"\"\"Adds (optionally learned) positional embeddings to the inputs.\n\n Attributes:\n config: TransformerConfig dataclass containing hyperparameters.\n decode: whether to run in single-position autoregressive mode.\n \"\"\"\n\n config: TransformerConfig\n decode: bool = False\n\n @nn.compact\n def __call__(self, inputs, inputs_positions=None):\n \"\"\"Applies AddPositionEmbs module.\n\n By default this layer uses a fixed sinusoidal embedding table. If a\n learned position embedding is desired, pass an initializer to\n posemb_init in the configuration.\n\n Args:\n inputs: input data.\n inputs_positions: input position indices for packed sequences.\n\n Returns:\n output: `(bs, timesteps, in_dim)`\n \"\"\"\n config = self.config\n # inputs.shape is (batch_size, seq_len, emb_dim)\n assert inputs.ndim == 3, (\n 'Number of dimensions should be 3, but it is: %d' % inputs.ndim\n )\n length = inputs.shape[1]\n pos_emb_shape = (1, config.max_len, inputs.shape[-1])\n if config.posemb_init is None:\n # Use a fixed (non-learned) sinusoidal position embedding.\n pos_embedding = sinusoidal_init(max_len=config.max_len)(\n None, pos_emb_shape, None\n )\n else:\n pos_embedding = self.param(\n 'pos_embedding', config.posemb_init, pos_emb_shape\n )\n pe = pos_embedding[:, :length, :]\n\n # We use a cache position index for tracking decoding position.\n if self.decode:\n is_initialized = self.has_variable('cache', 'cache_index')\n cache_index = self.variable(\n 'cache', 'cache_index', lambda: jnp.array(0, dtype=jnp.uint32)\n )\n if is_initialized:\n i = cache_index.value\n cache_index.value = i + 1\n _, _, df = pos_embedding.shape\n pe = lax.dynamic_slice(pos_embedding, jnp.array((0, i, 0)), (1, 1, df))\n if inputs_positions is None:\n # normal unpacked case:\n return inputs + pe\n else:\n # for packed data we need to use known position indices:\n return inputs + jnp.take(pe[0], inputs_positions, axis=0)\n\n\nclass MlpBlock(nn.Module):\n \"\"\"Transformer MLP / feed-forward block.\n\n Attributes:\n config: TransformerConfig dataclass containing hyperparameters.\n out_dim: optionally specify out dimension.\n \"\"\"\n\n config: TransformerConfig\n out_dim: int | None = None\n\n @nn.compact\n def __call__(self, inputs):\n \"\"\"Applies Transformer MlpBlock module.\"\"\"\n config = self.config\n actual_out_dim = inputs.shape[-1] if self.out_dim is None else self.out_dim\n x = nn.Dense(\n config.mlp_dim,\n dtype=config.dtype,\n kernel_init=config.kernel_init,\n bias_init=config.bias_init,\n )(inputs)\n x = nn.relu(x)\n x = nn.Dropout(rate=config.dropout_rate)(\n x, deterministic=config.deterministic\n )\n output = nn.Dense(\n actual_out_dim,\n dtype=config.dtype,\n kernel_init=config.kernel_init,\n bias_init=config.bias_init,\n )(x)\n output = nn.Dropout(rate=config.dropout_rate)(\n output, deterministic=config.deterministic\n )\n return output\n\n\nclass Encoder1DBlock(nn.Module):\n \"\"\"Transformer encoder layer.\n\n Attributes:\n config: TransformerConfig dataclass containing hyperparameters.\n \"\"\"\n\n config: TransformerConfig\n\n @nn.compact\n def __call__(self, inputs, encoder_mask=None):\n \"\"\"Applies Encoder1DBlock module.\n\n Args:\n inputs: input data.\n encoder_mask: encoder self-attention mask.\n\n Returns:\n output after transformer encoder block.\n \"\"\"\n config = self.config\n\n # Attention block.\n assert inputs.ndim == 3\n x = nn.LayerNorm(dtype=config.dtype)(inputs)\n x = nn.MultiHeadDotProductAttention(\n num_heads=config.num_heads,\n dtype=config.dtype,\n qkv_features=config.qkv_dim,\n kernel_init=config.kernel_init,\n bias_init=config.bias_init,\n use_bias=False,\n broadcast_dropout=False,\n dropout_rate=config.attention_dropout_rate,\n deterministic=config.deterministic,\n )(x, mask=encoder_mask)\n\n x = nn.Dropout(rate=config.dropout_rate)(\n x, deterministic=config.deterministic\n )\n x = x + inputs\n\n # MLP block.\n y = nn.LayerNorm(dtype=config.dtype)(x)\n y = MlpBlock(config=config)(y)\n\n return x + y\n\n\nclass EncoderDecoder1DBlock(nn.Module):\n \"\"\"Transformer encoder-decoder layer.\n\n Attributes:\n config: TransformerConfig dataclass containing hyperparameters.\n \"\"\"\n\n config: TransformerConfig\n\n @nn.compact\n def __call__(\n self, targets, encoded, decoder_mask=None, encoder_decoder_mask=None\n ):\n \"\"\"Applies EncoderDecoder1DBlock module.\n\n Args:\n targets: input data for decoder\n encoded: input data from encoder\n decoder_mask: decoder self-attention mask.\n encoder_decoder_mask: encoder-decoder attention mask.\n\n Returns:\n output after transformer encoder-decoder block.\n \"\"\"\n config = self.config\n\n # Decoder block.\n assert targets.ndim == 3\n x = nn.LayerNorm(dtype=config.dtype)(targets)\n x = nn.MultiHeadDotProductAttention(\n num_heads=config.num_heads,\n dtype=config.dtype,\n qkv_features=config.qkv_dim,\n kernel_init=config.kernel_init,\n bias_init=config.bias_init,\n use_bias=False,\n broadcast_dropout=False,\n dropout_rate=config.attention_dropout_rate,\n deterministic=config.deterministic,\n decode=config.decode,\n )(x, mask=decoder_mask)\n x = nn.Dropout(rate=config.dropout_rate)(\n x, deterministic=config.deterministic\n )\n x = x + targets\n\n # Encoder-Decoder block.\n y = nn.LayerNorm(dtype=config.dtype)(x)\n y = nn.MultiHeadDotProductAttention(\n num_heads=config.num_heads,\n dtype=config.dtype,\n qkv_features=config.qkv_dim,\n kernel_init=config.kernel_init,\n bias_init=config.bias_init,\n use_bias=False,\n broadcast_dropout=False,\n dropout_rate=config.attention_dropout_rate,\n deterministic=config.deterministic,\n )(y, encoded, mask=encoder_decoder_mask)\n\n y = nn.Dropout(rate=config.dropout_rate)(\n y, deterministic=config.deterministic\n )\n y = y + x\n\n # MLP block.\n z = nn.LayerNorm(dtype=config.dtype)(y)\n z = MlpBlock(config=config)(z)\n\n return y + z\n\n\nclass Encoder(nn.Module):\n \"\"\"Transformer Model Encoder for sequence to sequence translation.\n\n Attributes:\n config: TransformerConfig dataclass containing hyperparameters.\n shared_embedding: a shared embedding layer to use.\n \"\"\"\n\n config: TransformerConfig\n shared_embedding: Any = None\n\n @nn.compact\n def __call__(self, inputs, inputs_positions=None, encoder_mask=None):\n \"\"\"Applies Transformer model on the inputs.\n\n Args:\n inputs: input data\n inputs_positions: input subsequence positions for packed examples.\n encoder_mask: decoder self-attention mask.\n\n Returns:\n output of a transformer encoder.\n \"\"\"\n config = self.config\n assert inputs.ndim == 2 # (batch, len)\n\n # Input Embedding\n if self.shared_embedding is None:\n input_embed = nn.Embed(\n num_embeddings=config.vocab_size,\n features=config.emb_dim,\n embedding_init=nn.initializers.normal(stddev=1.0),\n )\n else:\n input_embed = self.shared_embedding\n x = inputs.astype('int32')\n x = input_embed(x)\n x = AddPositionEmbs(config=config, decode=False, name='posembed_input')(\n x, inputs_positions=inputs_positions\n )\n x = nn.Dropout(rate=config.dropout_rate)(\n x, deterministic=config.deterministic\n )\n\n x = x.astype(config.dtype)\n\n # Input Encoder\n for lyr in range(config.num_layers):\n x = Encoder1DBlock(config=config, name=f'encoderblock_{lyr}')(\n x, encoder_mask\n )\n\n encoded = nn.LayerNorm(dtype=config.dtype, name='encoder_norm')(x)\n\n return encoded\n\n\nclass Decoder(nn.Module):\n \"\"\"Transformer Model Decoder for sequence to sequence translation.\n\n Attributes:\n config: TransformerConfig dataclass containing hyperparameters.\n shared_embedding: a shared embedding layer to use.\n \"\"\"\n\n config: TransformerConfig\n shared_embedding: Any = None\n\n @nn.compact\n def __call__(\n self,\n encoded,\n targets,\n targets_positions=None,\n decoder_mask=None,\n encoder_decoder_mask=None,\n ):\n \"\"\"Applies Transformer model on the inputs.\n\n Args:\n encoded: encoded input data from encoder.\n targets: target inputs.\n targets_positions: input subsequence positions for packed examples.\n decoder_mask: decoder self-attention mask.\n encoder_decoder_mask: encoder-decoder attention mask.\n\n Returns:\n output of a transformer decoder.\n \"\"\"\n config = self.config\n\n assert encoded.ndim == 3 # (batch, len, depth)\n assert targets.ndim == 2 # (batch, len)\n\n # Target Embedding\n if self.shared_embedding is None:\n output_embed = nn.Embed(\n num_embeddings=config.output_vocab_size,\n features=config.emb_dim,\n embedding_init=nn.initializers.normal(stddev=1.0),\n )\n else:\n output_embed = self.shared_embedding\n\n y = targets.astype('int32')\n if not config.decode:\n y = shift_right(y)\n y = output_embed(y)\n y = AddPositionEmbs(\n config=config, decode=config.decode, name='posembed_output'\n )(y, inputs_positions=targets_positions)\n y = nn.Dropout(rate=config.dropout_rate)(\n y, deterministic=config.deterministic\n )\n\n y = y.astype(config.dtype)\n\n # Target-Input Decoder\n for lyr in range(config.num_layers):\n y = EncoderDecoder1DBlock(\n config=config, name=f'encoderdecoderblock_{lyr}'\n )(\n y,\n encoded,\n decoder_mask=decoder_mask,\n encoder_decoder_mask=encoder_decoder_mask,\n )\n y = nn.LayerNorm(dtype=config.dtype, name='encoderdecoder_norm')(y)\n\n # Decoded Logits\n if config.logits_via_embedding:\n # Use the transpose of embedding matrix for logit transform.\n logits = output_embed.attend(y.astype(jnp.float32))\n # Correctly normalize pre-softmax logits for this shared case.\n logits = logits / jnp.sqrt(y.shape[-1])\n else:\n logits = nn.Dense(\n config.output_vocab_size,\n dtype=config.dtype,\n kernel_init=config.kernel_init,\n bias_init=config.bias_init,\n name='logitdense',\n )(y)\n return logits\n\n\nclass Transformer(nn.Module):\n \"\"\"Transformer Model for sequence to sequence translation.\n\n Attributes:\n config: TransformerConfig dataclass containing hyperparameters.\n \"\"\"\n\n config: TransformerConfig\n\n def setup(self):\n config = self.config\n\n if config.share_embeddings:\n if config.output_vocab_size is not None:\n assert (\n config.output_vocab_size == config.vocab_size\n ), \"can't share embedding with different vocab sizes.\"\n self.shared_embedding = nn.Embed(\n num_embeddings=config.vocab_size,\n features=config.emb_dim,\n embedding_init=nn.initializers.normal(stddev=1.0),\n )\n else:\n self.shared_embedding = None\n\n self.encoder = Encoder(\n config=config, shared_embedding=self.shared_embedding\n )\n self.decoder = Decoder(\n config=config, shared_embedding=self.shared_embedding\n )\n\n def encode(self, inputs, inputs_positions=None, inputs_segmentation=None):\n \"\"\"Applies Transformer encoder-branch on the inputs.\n\n Args:\n inputs: input data.\n inputs_positions: input subsequence positions for packed examples.\n inputs_segmentation: input segmentation info for packed examples.\n\n Returns:\n encoded feature array from the transformer encoder.\n \"\"\"\n config = self.config\n # Make padding attention mask.\n encoder_mask = nn.make_attention_mask(\n inputs > 0, inputs > 0, dtype=config.dtype\n )\n # Add segmentation block-diagonal attention mask if using segmented data.\n if inputs_segmentation is not None:\n encoder_mask = nn.combine_masks(\n encoder_mask,\n nn.make_attention_mask(\n inputs_segmentation,\n inputs_segmentation,\n jnp.equal,\n dtype=config.dtype,\n ),\n )\n return self.encoder(\n inputs, inputs_positions=inputs_positions, encoder_mask=encoder_mask\n )\n\n def decode(\n self,\n encoded,\n inputs, # only needed for masks\n targets,\n targets_positions=None,\n inputs_segmentation=None,\n targets_segmentation=None,\n ):\n \"\"\"Applies Transformer decoder-branch on encoded-input and target.\n\n Args:\n encoded: encoded input data from encoder.\n inputs: input data (only needed for masking).\n targets: target data.\n targets_positions: target subsequence positions for packed examples.\n inputs_segmentation: input segmentation info for packed examples.\n targets_segmentation: target segmentation info for packed examples.\n\n Returns:\n logits array from transformer decoder.\n \"\"\"\n config = self.config\n\n # Make padding attention masks.\n if config.decode:\n # for fast autoregressive decoding only a special encoder-decoder mask is\n # used\n decoder_mask = None\n encoder_decoder_mask = nn.make_attention_mask(\n jnp.ones_like(targets) > 0, inputs > 0, dtype=config.dtype\n )\n else:\n decoder_mask = nn.combine_masks(\n nn.make_attention_mask(targets > 0, targets > 0, dtype=config.dtype),\n nn.make_causal_mask(targets, dtype=config.dtype),\n )\n encoder_decoder_mask = nn.make_attention_mask(\n targets > 0, inputs > 0, dtype=config.dtype\n )\n\n # Add segmentation block-diagonal attention masks if using segmented data.\n if inputs_segmentation is not None:\n decoder_mask = nn.combine_masks(\n decoder_mask,\n nn.make_attention_mask(\n targets_segmentation,\n targets_segmentation,\n jnp.equal,\n dtype=config.dtype,\n ),\n )\n encoder_decoder_mask = nn.combine_masks(\n encoder_decoder_mask,\n nn.make_attention_mask(\n targets_segmentation,\n inputs_segmentation,\n jnp.equal,\n dtype=config.dtype,\n ),\n )\n logits = self.decoder(\n encoded,\n targets,\n targets_positions=targets_positions,\n decoder_mask=decoder_mask,\n encoder_decoder_mask=encoder_decoder_mask,\n )\n return logits.astype(self.config.dtype)\n\n def __call__(\n self,\n inputs,\n targets,\n inputs_positions=None,\n targets_positions=None,\n inputs_segmentation=None,\n targets_segmentation=None,\n ):\n \"\"\"Applies Transformer model on the inputs.\n\n Args:\n inputs: input data.\n targets: target data.\n inputs_positions: input subsequence positions for packed examples.\n targets_positions: target subsequence positions for packed examples.\n inputs_segmentation: input segmentation info for packed examples.\n targets_segmentation: target segmentation info for packed examples.\n\n Returns:\n logits array from full transformer.\n \"\"\"\n encoded = self.encode(\n inputs,\n inputs_positions=inputs_positions,\n inputs_segmentation=inputs_segmentation,\n )\n\n return self.decode(\n encoded,\n inputs, # only used for masks\n targets,\n targets_positions=targets_positions,\n inputs_segmentation=inputs_segmentation,\n targets_segmentation=targets_segmentation,\n )\n" }, { "path": "examples/wmt/requirements.txt", "content": "absl-py==1.0.0\nclu==0.0.6\nflax==0.6.0\n-f https://storage.googleapis.com/jax-releases/libtpu_releases.html\n-f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\njax[cuda11_cudnn805]>=0.3.16 # change to jax[tpu] if running on tpus\nml-collections==0.1.0\nnumpy==1.22.0\noptax==0.1.0\nsentencepiece==0.1.96\nsix==1.15.0\ntensorflow==2.11.1\ntensorflow-datasets==4.4.0\ntensorflow-text==2.8.1\n" }, { "path": "examples/wmt/tokenizer.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Provides op for tokenizing a dataset.\"\"\"\n\nimport dataclasses\nimport os\nimport tempfile\nimport time\nfrom typing import Any, Dict, Tuple\nfrom collections.abc import Iterable\nimport sys\n\nfrom absl import logging\nimport jax\nimport tensorflow as tf\nif sys.version_info < (3, 13):\n import tensorflow_text as tftxt\nfrom sentencepiece import SentencePieceTrainer\n\nFeatures = dict[str, tf.Tensor]\n\n\ndef _dump_chars_to_textfile(\n dataset: tf.data.Dataset,\n maxchars: int = int(1e7),\n data_keys=('inputs', 'targets'),\n) -> tuple[str, int]:\n \"\"\"Write part of a TFDS sentence dataset to lines in a text file.\n\n Args:\n dataset: tf.dataset containing string-data.\n maxchars: int: approximate number of characters to save from dataset.\n data_keys: Tuple[str]: what keys in dataset to dump from.\n\n Returns:\n name of temp file with dataset bytes, exact number of characters dumped.\n \"\"\"\n char_count = 0\n ds_iter = dataset.as_numpy_iterator()\n with tempfile.NamedTemporaryFile(\n delete=False, prefix='/tmp/ds_chars'\n ) as outfp:\n while char_count < maxchars:\n example = next(ds_iter)\n for k in data_keys:\n line = example[k] + b'\\n'\n char_count += len(line)\n outfp.write(line)\n return outfp.name, char_count\n\n\ndef _train_sentencepiece(\n dataset: tf.data.Dataset,\n *,\n vocab_size: int,\n maxchars: int = int(1e7),\n model_path: str,\n model_type: str = 'unigram',\n character_coverage: float = 1.0,\n data_keys=('inputs', 'targets'),\n):\n \"\"\"Train SentencePiece tokenizer from subset of tf dataset.\n\n Args:\n dataset: tf.dataset\n vocab_size: int: size of vocab tokens to train.\n maxchars: int: number of characters to use for sentencepiece training.\n model_path: str: path of model file to save vocab model to.\n model_type: str: type of sentencepiece vocab to train.\n character_coverage: amount of characters covered by the model, good defaults\n are 0.9995 for languages with rich character set like Japanese or Chinese\n and 1.0 for other languages with small character set.\n data_keys: Tuple[str]: keys of dataset to use for training.\n\n Returns:\n path to the trained sentencepiece vocabulary model.\n \"\"\"\n if model_path.startswith('gs://'):\n abs_model_path = model_path\n else:\n abs_model_path = os.path.abspath(os.path.expanduser(model_path))\n fname, _ = _dump_chars_to_textfile(\n dataset, maxchars=maxchars, data_keys=data_keys\n )\n with tempfile.NamedTemporaryFile(\n delete=False, prefix='/tmp/sp_tmp'\n ) as model_fp:\n pass # we just want a prefix'd tmp-filename\n argstr = ' '.join([\n f'--input={fname}',\n f'--vocab_size={vocab_size}',\n f'--character_coverage={character_coverage}',\n f'--model_prefix={model_fp.name}',\n f'--model_type={model_type}',\n ])\n SentencePieceTrainer.Train(argstr)\n if jax.process_index() == 0:\n # Use an intermediate filename that is renamed to the target name to address\n # create and fill delays.\n copy_rename_path = abs_model_path + '.rntmp'\n tf.io.gfile.copy(model_fp.name + '.model', copy_rename_path, overwrite=True)\n tf.io.gfile.rename(copy_rename_path, abs_model_path, overwrite=True)\n logging.info('copied %s to %s', model_fp.name + '.model', abs_model_path)\n else:\n while not tf.io.gfile.exists(abs_model_path):\n time.sleep(1)\n time.sleep(1)\n return abs_model_path\n\n\ndef _load_sentencepiece_tokenizer(\n model_path: str,\n add_bos: bool = False,\n add_eos: bool = True,\n reverse: bool = False,\n):\n \"\"\"Load a tf-text SentencePiece tokenizer from given model filepath.\"\"\"\n with tf.io.gfile.GFile(model_path, 'rb') as model_fp:\n sp_model = model_fp.read()\n sp_tokenizer = tftxt.SentencepieceTokenizer(\n model=sp_model, add_bos=add_bos, add_eos=add_eos, reverse=reverse\n )\n return sp_tokenizer\n\n\ndef load_or_train_tokenizer(\n dataset: tf.data.Dataset,\n *,\n vocab_path: str,\n vocab_size: int,\n max_corpus_chars: int,\n data_keys: tuple[str, str] = ('inputs', 'targets'),\n):\n \"\"\"Loads the tokenizer at `vocab_path` or trains a one from `dataset`.\"\"\"\n try:\n return _load_sentencepiece_tokenizer(vocab_path)\n except tf.errors.NotFoundError:\n logging.info('SentencePiece vocab not found, building one from data.')\n vocab_path = _train_sentencepiece(\n dataset,\n vocab_size=vocab_size,\n maxchars=max_corpus_chars,\n model_path=vocab_path,\n data_keys=data_keys,\n )\n return _load_sentencepiece_tokenizer(vocab_path)\n\n\n@dataclasses.dataclass\nclass TokenizeOp:\n sp_tokenizer: Any\n data_keys: Iterable[str] = ('inputs', 'targets')\n\n def __call__(self, features: Features) -> Features:\n for k in self.data_keys:\n features[k] = self.sp_tokenizer.tokenize(features[k])\n return features\n" }, { "path": "examples/wmt/train.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Machine Translation example.\n\nThis script trains a Transformer on a WMT dataset.\n\"\"\"\n\n# pytype: disable=wrong-arg-count\n# pytype: disable=attribute-error\n\nimport collections\nimport functools\nimport os\n\nfrom absl import logging\nfrom clu import metric_writers\nfrom clu import periodic_actions\nfrom flax import jax_utils\nfrom flax import linen as nn\nfrom flax.training import checkpoints\nfrom flax.training import common_utils\nfrom flax.training import dynamic_scale as dynamic_scale_lib\nfrom flax.training import train_state\nimport jax\nimport jax.numpy as jnp\nimport ml_collections\nimport numpy as np\nimport optax\nimport orbax.checkpoint as ocp\nimport tensorflow as tf\n\nimport bleu\nimport decode\nimport input_pipeline\nimport models\n\n\nclass TrainState(train_state.TrainState):\n dynamic_scale: dynamic_scale_lib.DynamicScale\n\n\ndef rsqrt_schedule(\n init_value: float,\n shift: int = 0,\n):\n \"\"\"Applies a reverse square-root schedule.\n\n The reverse square root schedule is simply `lr = init_value / sqrt(step)`.\n\n Args:\n init_value: Base learning rate (before applying the rsqrt schedule).\n shift: How many steps the rsqrt should be shifted. Shifting the rsqrt\n schedule makes it less steep in the beginning (close to 0).\n\n Returns:\n A schedule `count -> learning_rate`.\n \"\"\"\n\n def schedule(count):\n return init_value * (count + shift) ** -0.5 * shift**0.5\n\n return schedule\n\n\ndef create_learning_rate_schedule(learning_rate: float, warmup_steps: int):\n \"\"\"Creates a rsqrt schedule with linear warmup.\"\"\"\n return optax.join_schedules(\n [\n optax.linear_schedule(\n init_value=0,\n end_value=learning_rate,\n transition_steps=warmup_steps,\n ),\n rsqrt_schedule(init_value=learning_rate, shift=warmup_steps),\n ],\n boundaries=[warmup_steps],\n )\n\n\ndef compute_weighted_cross_entropy(\n logits, targets, weights=None, label_smoothing=0.0\n):\n \"\"\"Compute weighted cross entropy and entropy for log probs and targets.\n\n Args:\n logits: [batch, length, num_classes] float array.\n targets: categorical targets [batch, length] int array.\n weights: None or array of shape [batch, length].\n label_smoothing: label smoothing constant, used to determine the on and off\n values.\n\n Returns:\n Tuple of scalar loss and batch normalizing factor.\n \"\"\"\n if logits.ndim != targets.ndim + 1:\n raise ValueError(\n \"Incorrect shapes. Got shape %s logits and %s targets\"\n % (str(logits.shape), str(targets.shape))\n )\n vocab_size = logits.shape[-1]\n confidence = 1.0 - label_smoothing\n low_confidence = (1.0 - confidence) / (vocab_size - 1)\n normalizing_constant = -(\n confidence * jnp.log(confidence)\n + (vocab_size - 1) * low_confidence * jnp.log(low_confidence + 1e-20)\n )\n soft_targets = common_utils.onehot(\n targets, vocab_size, on_value=confidence, off_value=low_confidence\n )\n\n loss = -jnp.sum(soft_targets * nn.log_softmax(logits), axis=-1)\n loss = loss - normalizing_constant\n\n normalizing_factor = np.prod(targets.shape)\n if weights is not None:\n loss = loss * weights\n normalizing_factor = weights.sum()\n\n return loss.sum(), normalizing_factor\n\n\ndef compute_weighted_accuracy(logits, targets, weights=None):\n \"\"\"Compute weighted accuracy for log probs and targets.\n\n Args:\n logits: [batch, length, num_classes] float array.\n targets: categorical targets [batch, length] int array.\n weights: None or array of shape [batch, length]\n\n Returns:\n Tuple of scalar loss and batch normalizing factor.\n \"\"\"\n if logits.ndim != targets.ndim + 1:\n raise ValueError(\n \"Incorrect shapes. Got shape %s logits and %s targets\"\n % (str(logits.shape), str(targets.shape))\n )\n loss = jnp.equal(jnp.argmax(logits, axis=-1), targets)\n normalizing_factor = np.prod(logits.shape[:-1])\n if weights is not None:\n loss = loss * weights\n normalizing_factor = weights.sum()\n\n return loss.sum(), normalizing_factor\n\n\ndef compute_metrics(logits, labels, weights, label_smoothing=0.0):\n \"\"\"Compute summary metrics.\"\"\"\n loss, weight_sum = compute_weighted_cross_entropy(\n logits, labels, weights, label_smoothing\n )\n acc, _ = compute_weighted_accuracy(logits, labels, weights)\n metrics = {\n \"loss\": loss,\n \"accuracy\": acc,\n \"denominator\": weight_sum,\n }\n metrics = jax.lax.psum(metrics, axis_name=\"batch\")\n return metrics\n\n\n# Primary training / eval / decode step functions.\n# -----------------------------------------------------------------------------\n\n\ndef train_step(\n state,\n batch,\n config,\n learning_rate_fn,\n label_smoothing=0.0,\n dropout_rng=None,\n):\n \"\"\"Perform a single training step.\"\"\"\n # X_position and X_segmentation are needed only when using \"packed examples\"\n # where multiple sequences are packed into the same example with this\n # metadata.\n # if such features are not present they are ignored and the example is treated\n # like a normal, unpacked sequence example.\n train_keys = [\n \"inputs\",\n \"targets\",\n \"inputs_position\",\n \"targets_position\",\n \"inputs_segmentation\",\n \"targets_segmentation\",\n ]\n (\n inputs,\n targets,\n inputs_positions,\n targets_positions,\n inputs_segmentation,\n targets_segmentation,\n ) = (batch.get(k, None) for k in train_keys)\n\n weights = jnp.where(targets > 0, 1, 0).astype(jnp.float32)\n\n dropout_rng = jax.random.fold_in(dropout_rng, state.step)\n\n def loss_fn(params):\n \"\"\"loss function used for training.\"\"\"\n logits = models.Transformer(config).apply(\n {\"params\": params},\n inputs,\n targets,\n inputs_positions=inputs_positions,\n targets_positions=targets_positions,\n inputs_segmentation=inputs_segmentation,\n targets_segmentation=targets_segmentation,\n rngs={\"dropout\": dropout_rng},\n )\n\n loss, weight_sum = compute_weighted_cross_entropy(\n logits, targets, weights, label_smoothing\n )\n mean_loss = loss / weight_sum\n return mean_loss, logits\n\n step = state.step\n\n if state.dynamic_scale:\n # dynamic scale takes care of averaging gradients across replicas\n grad_fn = state.dynamic_scale.value_and_grad(\n loss_fn, has_aux=True, axis_name=\"batch\"\n )\n dynamic_scale, is_fin, (_, logits), grads = grad_fn(state.params)\n state = state.replace(dynamic_scale=dynamic_scale)\n else:\n grad_fn = jax.value_and_grad(loss_fn, has_aux=True)\n (_, logits), grads = grad_fn(state.params)\n grads = jax.lax.pmean(grads, axis_name=\"batch\")\n\n new_state = state.apply_gradients(grads=grads)\n metrics = compute_metrics(logits, targets, weights)\n metrics[\"learning_rate\"] = learning_rate_fn(step)\n\n if state.dynamic_scale:\n # if is_fin == False the gradients contain Inf/NaNs and optimizer state and\n # params should be restored (= skip this step).\n select_fn = functools.partial(jnp.where, is_fin) # pylint: disable=undefined-variable\n new_state = new_state.replace(\n opt_state=jax.tree_util.tree_map(\n select_fn, new_state.opt_state, state.opt_state\n ),\n params=jax.tree_util.tree_map(\n select_fn, new_state.params, state.params\n ),\n )\n metrics[\"loss_scale\"] = dynamic_scale.scale * metrics[\"denominator\"] # pylint: disable=undefined-variable\n\n return new_state, metrics\n\n\ndef eval_step(params, batch, config, label_smoothing=0.0):\n \"\"\"Calculate evaluation metrics on a batch.\"\"\"\n inputs, targets = batch[\"inputs\"], batch[\"targets\"]\n weights = jnp.where(targets > 0, 1.0, 0.0)\n logits = models.Transformer(config).apply({\"params\": params}, inputs, targets)\n\n return compute_metrics(logits, targets, weights, label_smoothing)\n\n\ndef initialize_cache(inputs, max_decode_len, config):\n \"\"\"Initialize a cache for a given input shape and max decode length.\"\"\"\n target_shape = (inputs.shape[0], max_decode_len) + inputs.shape[2:]\n initial_variables = models.Transformer(config).init(\n jax.random.key(0),\n jnp.ones(inputs.shape, config.dtype),\n jnp.ones(target_shape, config.dtype),\n )\n return initial_variables[\"cache\"]\n\n\ndef predict_step(\n inputs, params, cache, eos_id, max_decode_len, config, beam_size=4\n):\n \"\"\"Predict translation with fast decoding beam search on a batch.\"\"\"\n # Prepare transformer fast-decoder call for beam search: for beam search, we\n # need to set up our decoder model to handle a batch size equal to\n # batch_size * beam_size, where each batch item's data is expanded in-place\n # rather than tiled.\n # i.e. if we denote each batch element subtensor as el[n]:\n # [el0, el1, el2] --> beamsize=2 --> [el0,el0,el1,el1,el2,el2]\n encoded_inputs = decode.flat_batch_beam_expand(\n models.Transformer(config).apply(\n {\"params\": params}, inputs, method=models.Transformer.encode\n ),\n beam_size,\n )\n raw_inputs = decode.flat_batch_beam_expand(inputs, beam_size)\n\n def tokens_ids_to_logits(flat_ids, flat_cache):\n \"\"\"Token slice to logits from decoder model.\"\"\"\n # --> [batch * beam, 1, vocab]\n flat_logits, new_vars = models.Transformer(config).apply(\n {\"params\": params, \"cache\": flat_cache},\n encoded_inputs,\n raw_inputs, # only needed for input padding mask\n flat_ids,\n mutable=[\"cache\"],\n method=models.Transformer.decode,\n )\n new_flat_cache = new_vars[\"cache\"]\n # Remove singleton sequence-length dimension:\n # [batch * beam, 1, vocab] --> [batch * beam, vocab]\n flat_logits = flat_logits.squeeze(axis=1)\n return flat_logits, new_flat_cache\n\n # Using the above-defined single-step decoder function, run a\n # beam search over possible sequences given input encoding.\n beam_seqs, _ = decode.beam_search(\n inputs,\n cache,\n tokens_ids_to_logits,\n beam_size=beam_size,\n alpha=0.6,\n eos_id=eos_id,\n max_decode_len=max_decode_len,\n )\n\n # Beam search returns [n_batch, n_beam, n_length + 1] with beam dimension\n # sorted in increasing order of log-probability.\n # Return the highest scoring beam sequence, drop first dummy 0 token.\n return beam_seqs[:, -1, 1:]\n\n\n# Utils for prediction and BLEU calculation\n# -----------------------------------------------------------------------------\n\n\ndef pad_examples(x, desired_batch_size):\n \"\"\"Expand batch to desired size by repeating last slice.\"\"\"\n batch_pad = desired_batch_size - x.shape[0]\n return np.concatenate([x, np.tile(x[-1], (batch_pad, 1))], axis=0)\n\n\ndef per_host_sum_pmap(in_tree):\n \"\"\"Execute psum on in_tree\"s leaves over one device per host.\"\"\"\n host2devices = collections.defaultdict(list)\n for d in jax.devices():\n host2devices[d.process_index].append(d)\n devices = [host2devices[k][0] for k in host2devices]\n host_psum = jax.pmap(lambda x: jax.lax.psum(x, \"i\"), \"i\", devices=devices)\n\n def pre_pmap(xs):\n return jax.tree_util.tree_map(\n lambda x: jnp.broadcast_to(x, (1,) + x.shape), xs\n )\n\n def post_pmap(xs):\n # Avoid degraded performance under the new jax.pmap. See\n # https://docs.jax.dev/en/latest/migrate_pmap.html#int-indexing-into-sharded-arrays.\n return jax.tree_util.tree_map(\n lambda x: x.addressable_shards[0].data.squeeze(0), xs\n )\n\n return post_pmap(host_psum(pre_pmap(in_tree)))\n\n\ndef tohost(x):\n \"\"\"Collect batches from all devices to host and flatten batch dimensions.\"\"\"\n n_device, n_batch, *remaining_dims = x.shape\n return np.array(x).reshape((n_device * n_batch,) + tuple(remaining_dims))\n\n\ndef evaluate(\n *, p_eval_step, params, eval_ds: tf.data.Dataset, num_eval_steps: int\n):\n \"\"\"Evaluate the params an return a dictionary with the metrics.\"\"\"\n logging.info(\"Gathering evaluation metrics.\")\n eval_metrics = []\n eval_iter = iter(eval_ds) # pytype: disable=wrong-arg-types\n for _, eval_batch in zip(range(num_eval_steps), eval_iter):\n eval_batch = jax.tree_util.tree_map(lambda x: x._numpy(), eval_batch) # pylint: disable=protected-access\n eval_batch = common_utils.shard(eval_batch)\n metrics = p_eval_step(params, eval_batch)\n eval_metrics.append(metrics)\n eval_metrics = common_utils.get_metrics(eval_metrics)\n eval_metrics_sums = jax.tree_util.tree_map(jnp.sum, eval_metrics)\n eval_denominator = eval_metrics_sums.pop(\"denominator\")\n eval_summary = jax.tree_util.tree_map(\n lambda x: x / eval_denominator, # pylint: disable=cell-var-from-loop\n eval_metrics_sums,\n )\n return eval_summary\n\n\ndef translate_and_calculate_bleu(\n *,\n p_pred_step,\n p_init_cache,\n params,\n predict_ds: tf.data.Dataset,\n decode_tokens,\n max_predict_length: int,\n):\n \"\"\"Translates the `predict_ds` and calculates the BLEU score.\"\"\"\n n_devices = jax.local_device_count()\n logging.info(\"Translating evaluation dataset.\")\n sources, references, predictions = [], [], []\n for pred_batch in predict_ds:\n pred_batch = jax.tree_util.tree_map(lambda x: x._numpy(), pred_batch) # pylint: disable=protected-access\n # Handle final odd-sized batch by padding instead of dropping it.\n cur_pred_batch_size = pred_batch[\"inputs\"].shape[0]\n if cur_pred_batch_size % n_devices:\n padded_size = int(np.ceil(cur_pred_batch_size / n_devices) * n_devices)\n pred_batch = jax.tree_util.tree_map(\n lambda x: pad_examples(x, padded_size), # pylint: disable=cell-var-from-loop\n pred_batch,\n )\n pred_batch = common_utils.shard(pred_batch)\n cache = p_init_cache(pred_batch[\"inputs\"])\n predicted = p_pred_step(\n pred_batch[\"inputs\"], params, cache, decode.EOS_ID, max_predict_length\n )\n predicted = tohost(predicted)\n inputs = tohost(pred_batch[\"inputs\"])\n targets = tohost(pred_batch[\"targets\"])\n # Iterate through non-padding examples of batch.\n for i, s in enumerate(predicted[:cur_pred_batch_size]):\n sources.append(decode_tokens(inputs[i]))\n references.append(decode_tokens(targets[i]))\n predictions.append(decode_tokens(s))\n logging.info(\n \"Translation: %d predictions %d references %d sources.\",\n len(predictions),\n len(references),\n len(sources),\n )\n\n # Calculate BLEU score for translated eval corpus against reference.\n bleu_matches = bleu.bleu_partial(references, predictions)\n all_bleu_matches = per_host_sum_pmap(bleu_matches)\n bleu_score = bleu.complete_bleu(*all_bleu_matches)\n # Save translation samples for tensorboard.\n exemplars = \"\"\n for n in np.random.choice(np.arange(len(predictions)), 8):\n exemplars += f\"{sources[n]}\\n\\n{references[n]}\\n\\n{predictions[n]}\\n\\n\"\n return exemplars, bleu_score\n\n\ndef preferred_dtype(config):\n platform = jax.local_devices()[0].platform\n if config.use_mixed_precision:\n if platform == \"tpu\":\n return jnp.bfloat16\n elif platform == \"gpu\":\n return jnp.float16\n return jnp.float32\n\n\ndef train_and_evaluate(config: ml_collections.ConfigDict, workdir: str):\n \"\"\"Runs a training and evaluation loop.\n\n Args:\n config: Configuration to use.\n workdir: Working directory for checkpoints and TF summaries. If this\n contains checkpoint training will be resumed from the latest checkpoint.\n \"\"\"\n tf.io.gfile.makedirs(workdir)\n\n vocab_path = config.vocab_path\n if vocab_path is None:\n vocab_path = os.path.join(workdir, \"sentencepiece_model\")\n config.vocab_path = vocab_path\n tf.io.gfile.makedirs(os.path.split(vocab_path)[0])\n\n # Load Dataset\n # ---------------------------------------------------------------------------\n logging.info(\"Initializing dataset.\")\n train_ds, eval_ds, predict_ds, encoder = input_pipeline.get_wmt_datasets(\n n_devices=jax.local_device_count(),\n config=config,\n reverse_translation=config.reverse_translation,\n vocab_path=vocab_path,\n )\n\n train_iter = iter(train_ds)\n vocab_size = int(encoder.vocab_size())\n eos_id = decode.EOS_ID # Default Sentencepiece EOS token.\n\n def decode_tokens(toks):\n valid_toks = toks[: np.argmax(toks == eos_id) + 1].astype(np.int32)\n return encoder.detokenize(valid_toks).numpy().decode(\"utf-8\")\n\n if config.num_predict_steps > 0:\n predict_ds = predict_ds.take(config.num_predict_steps)\n\n logging.info(\"Initializing model, optimizer, and step functions.\")\n\n dtype = preferred_dtype(config)\n\n # Build Model and Optimizer\n # ---------------------------------------------------------------------------\n train_config = models.TransformerConfig(\n vocab_size=vocab_size,\n output_vocab_size=vocab_size,\n share_embeddings=config.share_embeddings,\n logits_via_embedding=config.logits_via_embedding,\n dtype=dtype,\n emb_dim=config.emb_dim,\n num_heads=config.num_heads,\n num_layers=config.num_layers,\n qkv_dim=config.qkv_dim,\n mlp_dim=config.mlp_dim,\n max_len=max(config.max_target_length, config.max_eval_target_length),\n dropout_rate=config.dropout_rate,\n attention_dropout_rate=config.attention_dropout_rate,\n deterministic=False,\n decode=False,\n kernel_init=nn.initializers.xavier_uniform(),\n bias_init=nn.initializers.normal(stddev=1e-6),\n )\n eval_config = train_config.replace(deterministic=True)\n predict_config = train_config.replace(deterministic=True, decode=True)\n\n start_step = 0\n rng = jax.random.key(config.seed)\n rng, init_rng = jax.random.split(rng)\n input_shape = (config.per_device_batch_size, config.max_target_length)\n target_shape = (config.per_device_batch_size, config.max_target_length)\n\n m = models.Transformer(eval_config)\n initial_variables = jax.jit(m.init)(\n init_rng,\n jnp.ones(input_shape, jnp.float32),\n jnp.ones(target_shape, jnp.float32),\n )\n\n # Create train state with Adam optimizer and weight decay.\n learning_rate_fn = create_learning_rate_schedule(\n learning_rate=config.learning_rate, warmup_steps=config.warmup_steps\n )\n dynamic_scale = None\n if dtype == jnp.float16:\n dynamic_scale = dynamic_scale_lib.DynamicScale()\n state = TrainState.create(\n apply_fn=m.apply,\n params=initial_variables[\"params\"],\n tx=optax.adamw(\n learning_rate=learning_rate_fn,\n b1=0.9,\n b2=0.98,\n eps=1e-9,\n weight_decay=config.weight_decay,\n ),\n dynamic_scale=dynamic_scale,\n )\n\n # We access model params only via state.params\n del initial_variables\n\n if config.restore_checkpoints:\n # Restore unreplicated optimizer + model state from last checkpoint.\n state = checkpoints.restore_checkpoint(workdir, state)\n # Grab last step.\n start_step = int(state.step)\n\n writer = metric_writers.create_default_writer(\n workdir, just_logging=jax.process_index() > 0\n )\n if start_step == 0:\n writer.write_hparams(dict(config))\n\n # Replicate state.\n state = jax_utils.replicate(state)\n\n # compile multidevice versions of train/eval/predict step and cache init fn.\n p_train_step = jax.pmap(\n functools.partial(\n train_step,\n config=train_config,\n learning_rate_fn=learning_rate_fn,\n label_smoothing=config.label_smoothing,\n ),\n axis_name=\"batch\",\n donate_argnums=(0,),\n ) # pytype: disable=wrong-arg-types\n p_eval_step = jax.pmap(\n functools.partial(eval_step, config=eval_config), axis_name=\"batch\"\n )\n p_init_cache = jax.pmap(\n functools.partial(\n initialize_cache,\n max_decode_len=config.max_predict_length,\n config=predict_config,\n ),\n axis_name=\"batch\",\n )\n p_pred_step = jax.pmap(\n functools.partial(\n predict_step, config=predict_config, beam_size=config.beam_size\n ),\n axis_name=\"batch\",\n static_broadcasted_argnums=(3, 4),\n ) # eos token, max_length are constant\n\n # Main Train Loop\n # ---------------------------------------------------------------------------\n\n # We init the first set of dropout PRNG keys, but update it afterwards inside\n # the main pmap\"d training update for performance.\n dropout_rngs = jax.random.split(rng, jax.local_device_count())\n del rng\n\n logging.info(\"Starting training loop.\")\n hooks = []\n report_progress = periodic_actions.ReportProgress(\n num_train_steps=config.num_train_steps, writer=writer\n )\n if jax.process_index() == 0:\n hooks += [\n report_progress,\n periodic_actions.Profile(logdir=workdir, num_profile_steps=5),\n ]\n train_metrics = []\n with metric_writers.ensure_flushes(writer):\n for step in range(start_step, config.num_train_steps):\n is_last_step = step == config.num_train_steps - 1\n\n # Shard data to devices and do a training step.\n with jax.profiler.StepTraceAnnotation(\"train\", step_num=step):\n batch = common_utils.shard(\n jax.tree_util.tree_map(np.asarray, next(train_iter))\n )\n state, metrics = p_train_step(state, batch, dropout_rng=dropout_rngs)\n train_metrics.append(metrics)\n\n # Quick indication that training is happening.\n logging.log_first_n(logging.INFO, \"Finished training step %d.\", 5, step)\n for h in hooks:\n h(step)\n\n # Periodic metric handling.\n if step % config.eval_every_steps == 0 or is_last_step:\n with report_progress.timed(\"training_metrics\"):\n logging.info(\"Gathering training metrics.\")\n train_metrics = common_utils.get_metrics(train_metrics)\n lr = train_metrics.pop(\"learning_rate\").mean()\n metrics_sums = jax.tree_util.tree_map(jnp.sum, train_metrics)\n denominator = metrics_sums.pop(\"denominator\")\n summary = jax.tree_util.tree_map(\n lambda x: x / denominator, metrics_sums # pylint: disable=cell-var-from-loop\n )\n summary[\"learning_rate\"] = lr\n summary = {\"train_\" + k: v for k, v in summary.items()}\n writer.write_scalars(step, summary)\n train_metrics = []\n\n with report_progress.timed(\"eval\"):\n eval_results = evaluate(\n p_eval_step=p_eval_step,\n params=state.params,\n eval_ds=eval_ds,\n num_eval_steps=config.num_eval_steps,\n )\n writer.write_scalars(\n step, {\"eval_\" + k: v for k, v in eval_results.items()}\n )\n\n with report_progress.timed(\"translate_and_bleu\"):\n exemplars, bleu_score = translate_and_calculate_bleu(\n p_pred_step=p_pred_step,\n p_init_cache=p_init_cache,\n params=state.params,\n predict_ds=predict_ds,\n decode_tokens=decode_tokens,\n max_predict_length=config.max_predict_length,\n )\n writer.write_scalars(step, {\"bleu\": bleu_score})\n writer.write_texts(step, {\"samples\": exemplars})\n\n # Save a checkpoint on one host after every checkpoint_freq steps.\n save_checkpoint = (\n step % config.checkpoint_every_steps == 0 or is_last_step\n )\n if config.save_checkpoints and save_checkpoint:\n logging.info(\"Saving checkpoint step %d.\", step)\n\n # Orbax can not handle host local arrays from pmap.\n replicated_state = jax.tree_util.tree_map(\n ocp.utils.fully_replicated_host_local_array_to_global_array,\n state,\n )\n with report_progress.timed(\"checkpoint\"):\n checkpoints.save_checkpoint_multiprocess(\n workdir, replicated_state, step\n )\n" }, { "path": "examples/wmt/train_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport pathlib\nimport tempfile\nimport sys\n\nfrom absl import logging\nfrom absl.testing import absltest\nimport jax\nimport tensorflow as tf\nimport tensorflow_datasets as tfds\n\nfrom configs import default\nimport train\n\n\njax.config.update('jax_disable_most_optimizations', True)\n\n\nclass TrainTest(absltest.TestCase):\n \"\"\"Test cases for WMT library.\"\"\"\n\n def setUp(self):\n if sys.version_info >= (3, 13):\n self.skipTest('Test (and tensorflow-text) does not suport Python 3.13+')\n super().setUp()\n tf.config.experimental.set_visible_devices([], 'GPU')\n\n def test_train_and_evaluate(self):\n config = default.get_config()\n config.max_corpus_chars = 1000\n config.vocab_size = 32\n config.per_device_batch_size = 1\n config.num_train_steps = 1\n config.num_eval_steps = 1\n config.num_predict_steps = 1\n\n config.num_layers = 1\n config.qkv_dim = 128\n config.emb_dim = 128\n config.mlp_dim = 512\n config.num_heads = 2\n\n config.max_target_length = 32\n config.max_eval_target_length = 32\n config.max_predict_length = 32\n\n workdir = tempfile.mkdtemp()\n\n # Go two directories up to the root of the flax directory.\n flax_root_dir = pathlib.Path(__file__).parents[2]\n data_dir = str(flax_root_dir) + '/.tfds/metadata' # pylint: disable=unused-variable\n\n with tfds.testing.mock_data(num_examples=128, data_dir=data_dir):\n train.train_and_evaluate(config, workdir)\n logging.info('workdir content: %s', tf.io.gfile.listdir(workdir))\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "flax/__init__.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Flax API.\"\"\"\n\n# pylint: disable=g-import-not-at-top\n# pyformat: disable\n\nfrom flax import configurations\nconfig: configurations.Config = configurations.config\ndel configurations\n\nfrom flax import core\nfrom flax import jax_utils\nfrom flax import linen\nfrom flax import serialization\nfrom flax import traverse_util\n\nfrom flax import version\n__version__: str = version.__version__\ndel version\n\n# DO NOT REMOVE - Marker for internal deprecated API.\n\n# DO NOT REMOVE - Marker for internal logging.\n\n# pyformat: enable\n# pylint: enable=g-import-not-at-top\n" }, { "path": "flax/configurations.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Global configuration flags for Flax.\"\"\"\n\nimport os\nfrom contextlib import contextmanager\nfrom typing import Any, Generic, NoReturn, TypeVar, overload\n\n_T = TypeVar('_T')\n\n\nclass Config:\n flax_use_flaxlib: bool\n flax_array_ref: bool\n flax_pytree_module: bool\n flax_max_repr_depth: int | None\n flax_always_shard_variable: bool\n flax_hijax_variable: bool\n nnx_graph_mode: bool\n nnx_graph_updates: bool\n # See https://google.github.io/pytype/faq.html.\n _HAS_DYNAMIC_ATTRIBUTES = True\n\n def __init__(self):\n self._values = {}\n\n def _add_option(self, name, default):\n if name in self._values:\n raise RuntimeError(f'Config option {name} already defined')\n self._values[name] = default\n\n def _read(self, name):\n try:\n return self._values[name]\n except KeyError:\n raise LookupError(f'Unrecognized config option: {name}')\n\n @overload\n def update(self, name: str, value: Any, /) -> None:\n ...\n\n @overload\n def update(self, holder: 'FlagHolder[_T]', value: _T, /) -> None:\n ...\n\n def update(self, name_or_holder, value, /):\n \"\"\"Modify the value of a given flag.\n\n Args:\n name_or_holder: the name of the flag to modify or the corresponding\n flag holder object.\n value: new value to set.\n \"\"\"\n name = name_or_holder\n if isinstance(name_or_holder, FlagHolder):\n name = name_or_holder.name\n if name not in self._values:\n raise LookupError(f'Unrecognized config option: {name}')\n self._values[name] = value\n\n def __repr__(self):\n values_repr = ', '.join(f'\\n {k}={v!r}' for k, v in self._values.items())\n return f'Config({values_repr}\\n)'\n\n @contextmanager\n def temp_flip_flag(self, var_name: str, var_value: bool):\n \"\"\"Context manager to temporarily flip feature flags for test functions.\n\n Args:\n var_name: the config variable name (without the 'flax_' prefix)\n var_value: the boolean value to set var_name to temporarily\n \"\"\"\n old_value = getattr(self, f'flax_{var_name}')\n try:\n self.update(f'flax_{var_name}', var_value)\n yield\n finally:\n self.update(f'flax_{var_name}', old_value)\n\n\nconfig = Config()\n\n# Config parsing utils\n\n\nclass FlagHolder(Generic[_T]):\n def __init__(self, name, help):\n self.name = name\n self.__name__ = name[4:] if name.startswith('flax_') else name\n self.__doc__ = f'Flag holder for `{name}`.\\n\\n{help}'\n\n def __bool__(self) -> NoReturn:\n raise TypeError(\n \"bool() not supported for instances of type '{0}' \"\n \"(did you mean to use '{0}.value' instead?)\".format(type(self).__name__)\n )\n\n @property\n def value(self) -> _T:\n return config._read(self.name)\n\n\ndef bool_flag(name: str, *, default: bool, help: str) -> FlagHolder[bool]:\n \"\"\"Set up a boolean flag.\n\n Example::\n\n enable_foo = bool_flag(\n name='flax_enable_foo',\n default=False,\n help='Enable foo.',\n )\n\n Now the ``FLAX_ENABLE_FOO`` shell environment variable can be used to\n control the process-level value of the flag, in addition to using e.g.\n ``config.update(\"flax_enable_foo\", True)`` directly.\n\n Args:\n name: converted to lowercase to define the name of the flag. It is\n converted to uppercase to define the corresponding shell environment\n variable.\n default: a default value for the flag.\n help: used to populate the docstring of the returned flag holder object.\n\n Returns:\n A flag holder object for accessing the value of the flag.\n \"\"\"\n name = name.lower()\n config._add_option(name, static_bool_env(name.upper(), default))\n fh = FlagHolder[bool](name, help)\n setattr(Config, name, property(lambda _: fh.value, doc=help))\n return fh\n\n\ndef int_flag(name: str, *, default: int | None, help: str) -> FlagHolder[int]:\n \"\"\"Set up an integer flag.\n\n Example::\n\n num_foo = int_flag(\n name='flax_num_foo',\n default=42,\n help='Number of foo.',\n )\n\n Now the ``FLAX_NUM_FOO`` shell environment variable can be used to\n control the process-level value of the flag, in addition to using e.g.\n ``config.update(\"flax_num_foo\", 42)`` directly.\n\n Args:\n name: converted to lowercase to define the name of the flag. It is\n converted to uppercase to define the corresponding shell environment\n variable.\n default: a default value for the flag.\n help: used to populate the docstring of the returned flag holder object.\n\n Returns:\n A flag holder object for accessing the value of the flag.\n \"\"\"\n name = name.lower()\n config._add_option(name, static_int_env(name.upper(), default))\n fh = FlagHolder[int](name, help)\n setattr(Config, name, property(lambda _: fh.value, doc=help))\n return fh\n\n\ndef static_bool_env(varname: str, default: bool) -> bool:\n \"\"\"Read an environment variable and interpret it as a boolean.\n\n This is deprecated. Please use bool_flag() unless your flag\n will be used in a static method and does not require runtime updates.\n\n True values are (case insensitive): 'y', 'yes', 't', 'true', 'on', and '1';\n false values are 'n', 'no', 'f', 'false', 'off', and '0'.\n Args:\n varname: the name of the variable\n default: the default boolean value\n Returns:\n boolean return value derived from defaults and environment.\n Raises: ValueError if the environment variable is anything else.\n \"\"\"\n val = os.getenv(varname, str(default))\n val = val.lower()\n if val in ('y', 'yes', 't', 'true', 'on', '1'):\n return True\n elif val in ('n', 'no', 'f', 'false', 'off', '0'):\n return False\n else:\n raise ValueError(\n f'invalid truth value {val!r} for environment {varname!r}'\n )\n\n\ndef static_int_env(varname: str, default: int | None) -> int | None:\n \"\"\"Read an environment variable and interpret it as an integer.\n\n Args:\n varname: the name of the variable\n default: the default integer value\n Returns:\n integer return value derived from defaults and environment.\n Raises: ValueError if the environment variable is not an integer.\n \"\"\"\n val = os.getenv(varname)\n if val is None:\n return default\n try:\n return int(val)\n except ValueError:\n raise ValueError(\n f'invalid integer value {val!r} for environment {varname!r}'\n ) from None\n\n\n# Flax Global Configuration Variables:\n\nflax_filter_frames = bool_flag(\n name='flax_filter_frames',\n default=True,\n help='Whether to hide flax-internal stack frames from tracebacks.',\n)\n\nflax_profile = bool_flag(\n name='flax_profile',\n default=True,\n help='Whether to run Module methods under jax.named_scope for profiles.',\n)\n\nflax_use_orbax_checkpointing = bool_flag(\n name='flax_use_orbax_checkpointing',\n default=True,\n help='Whether to use Orbax to save checkpoints.',\n)\n\nflax_preserve_adopted_names = bool_flag(\n name='flax_preserve_adopted_names',\n default=False,\n help=\"When adopting outside modules, don't clobber existing names.\",\n)\n\n# TODO(marcuschiam): remove this feature flag once regular dict migration is complete\nflax_return_frozendict = bool_flag(\n name='flax_return_frozendict',\n default=False,\n help='Whether to return FrozenDicts when calling init or apply.',\n)\n\nflax_fix_rng = bool_flag(\n name='flax_fix_rng_separator',\n default=False,\n help=(\n 'Whether to add separator characters when folding in static data into'\n ' PRNG keys.'\n ),\n)\n\nflax_use_flaxlib = bool_flag(\n name='flax_use_flaxlib',\n default=False,\n help='Whether to use flaxlib for C++ acceleration.',\n)\nflax_array_ref = bool_flag(\n name='flax_array_ref',\n default=False,\n help='Whether to use array refs.',\n)\nflax_pytree_module = bool_flag(\n name='flax_pytree_module',\n default=True,\n help='Whether Modules are pytrees by default or not.',\n)\n\nflax_max_repr_depth = int_flag(\n name='flax_max_repr_depth',\n default=None,\n help='Maximum depth of reprs for nested flax objects. Default is None (no limit).',\n)\n\nflax_always_shard_variable = bool_flag(\n name='flax_always_shard_variable',\n default=True,\n help='Whether a `nnx.Variable` should always automatically be sharded if it contains sharding annotations.',\n)\nflax_hijax_variable = bool_flag(\n name='flax_hijax_variable',\n default=False,\n help='Whether to enable HiJAX support for `nnx.Variable`.',\n)\nnnx_graph_mode = bool_flag(\n name='nnx_graph_mode',\n default=True,\n help='Whether NNX APIs default to graph-mode (True) or tree-mode (False).',\n)\nnnx_graph_updates = bool_flag(\n name='nnx_graph_updates',\n default=True,\n help='Whether graph-mode uses dynamic (True) or simple (False) graph traversal.',\n)" }, { "path": "flax/core/__init__.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom .axes_scan import broadcast as broadcast\nfrom .frozen_dict import (\n FrozenDict as FrozenDict,\n copy as copy,\n freeze as freeze,\n pop as pop,\n pretty_repr as pretty_repr,\n unfreeze as unfreeze,\n)\nfrom .lift import (\n custom_vjp as custom_vjp,\n jit as jit,\n jvp as jvp,\n remat_scan as remat_scan,\n remat as remat,\n scan as scan,\n vjp as vjp,\n vmap as vmap,\n while_loop as while_loop,\n)\nfrom .meta import (\n AxisMetadata as AxisMetadata,\n map_axis_meta as map_axis_meta,\n unbox as unbox,\n)\nfrom .scope import (\n DenyList as DenyList,\n Scope as Scope,\n apply as apply,\n bind as bind,\n init as init,\n lazy_init as lazy_init,\n)\nfrom .tracers import (\n check_trace_level as check_trace_level,\n current_trace as current_trace,\n)\n\nfrom flax.typing import (\n Array as Array,\n)\n" }, { "path": "flax/core/axes_scan.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Wrapper around jax.lax.scan with in_axes/out_axes API.\"\"\"\nfrom collections.abc import Callable\nimport functools\nfrom typing import Any, Optional\n\nimport jax\nfrom jax import core\nfrom jax import lax\nfrom jax.extend import linear_util as lu\nfrom jax.interpreters import partial_eval as pe\nimport jax.numpy as jnp\nimport numpy as np\n\nScanAxis = Optional[int]\n\n\nclass _Broadcast:\n pass\n\n\nbroadcast = _Broadcast()\n\n\ndef build_shaped_array(x, batch_dim: bool = False) -> core.ShapedArray:\n \"\"\"Builds ShapedArray preserving as much information from x as possible.\"\"\"\n shape = jnp.shape(x)\n sharding = x.aval.sharding if hasattr(x, \"aval\") else None\n if batch_dim:\n shape = shape[1:]\n if sharding is not None:\n if sharding.spec[0] is not None:\n raise ValueError(\n \"Batch dimension in scan `xs` cannot be sharded.\"\n )\n sharding = sharding.update(\n spec=jax.sharding.PartitionSpec(*sharding.spec[1:]))\n return core.ShapedArray(\n shape=shape,\n dtype=jnp.result_type(x),\n sharding=sharding,\n **{k: getattr(x, k) for k in [\"weak_type\", \"manual_type\"]\n if hasattr(x, k)},\n )\n\n\ndef scan(\n fn: Callable[..., Any],\n in_axes: Any,\n out_axes: Any,\n length: int | None = None,\n reverse: bool = False,\n unroll: int = 1,\n _split_transpose: bool = False,\n check_constancy_invariants: bool = True,\n):\n \"\"\"A wrapper around `jax.lax.scan` with in_axes/out_axes api.\n\n Example::\n def body_fn(b, c, x):\n return b + 2, c + 1, 2 * x\n\n loop = scan(body_fn, in_axes=0, out_axes=0)\n broadcast_in = 1\n carry = 2\n xs = jnp.arange(3)\n broadcast_out, carry, ys = loop(broadcast_in, carry, xs)\n print(broadcast_out) # prints: 3\n print(carry) # prints: 5\n print(ys) # prints: [0, 2, 4]\n\n\n Args:\n fn: the body function of the scan loop of the form\n `(broadcast_in, carry, *args) -> (broadcast_out, carry, scan_out)`.\n the broadcast argument allows for loop independent inputs/outputs to\n be computed inside `fn`. `fn` will be called once to compute\n `broadcast_out`. The actual loop will receive `broadcast_out` as the new\n `broadcast_in`. This is useful for initializing values inside the loop.\n in_axes: specifies the axis along which arguments are scanned.\n Use `broadcast` to use the same value across iterations.\n out_axes: specifies the axis along which outputs are concatenated.\n Use `broadcast` if a return value should not be concatenated and\n is independent of the loop body.\n length: number of iterations. Only needs to be specified if there\n is no scan axis from which it can be derived.\n reverse: scan in reverse order from end to start.\n unroll: how many scan iterations to unroll within a single\n iteration of a loop (default: 1).\n _split_transpose: An experimental feature to split the transpose of scan\n into a scan and a map, backed by an experimental Jax lax.scan() feature.\n check_constancy_invariants: If true, the scan will verify that the\n broadcast constants are true loop invariants, and further supports\n broadcast function (non-carry) outputs. This requires an extra jax\n tracing step however, so setting to false can reduce trace time on larger\n models.\n Returns:\n the function that performs the scan of the form:\n (broadcast_in, carry_in, *args) -> (broadcast_out, carry_out, scan_out).\n \"\"\"\n\n def transpose_to_front(ax, xs):\n if ax is broadcast:\n return ()\n if ax == 0:\n return xs\n\n def trans(x):\n perm = tuple(range(x.ndim))\n perm = (ax,) + tuple(np.delete(perm, ax))\n return jnp.transpose(x, perm)\n\n return jax.tree_util.tree_map(trans, xs)\n\n def transpose_from_front(ax, xs):\n if ax is broadcast:\n return ()\n if ax == 0:\n return xs\n\n def trans(x):\n if ax < 0:\n pax = x.ndim + ax\n else:\n pax = ax\n assert pax < x.ndim\n perm = tuple(range(1, pax + 1)) + (0,) + tuple(range(pax + 1, x.ndim))\n return jnp.transpose(x, perm)\n\n return jax.tree_util.tree_map(trans, xs)\n\n def scan_fn(broadcast_in, init, *args):\n # Requires one extra tracing operation to test invariants:\n # Verifies that broadcast constants are true loop invariants, and further\n # supports broadcast function (non-carry) outputs.\n\n xs = jax.tree_util.tree_map(transpose_to_front, in_axes, args)\n\n def body_fn(c, xs, init_mode=False):\n # inject constants\n xs = jax.tree_util.tree_map(\n lambda ax, arg, x: (arg if ax is broadcast else x), in_axes, args, xs\n )\n broadcast_out, c, ys = fn(broadcast_in, c, *xs)\n\n if init_mode:\n ys = jax.tree_util.tree_map(\n lambda ax, y: (y if ax is broadcast else ()), out_axes, ys\n )\n return broadcast_out, ys\n else:\n ys = jax.tree_util.tree_map(\n lambda ax, y: (() if ax is broadcast else y), out_axes, ys\n )\n return c, ys\n\n broadcast_body = functools.partial(body_fn, init_mode=True)\n\n init_flat, carry_tree = jax.tree.flatten(init)\n xs_flat, scan_tree = jax.tree.flatten(xs)\n carry_avals = [build_shaped_array(x) for x in init_flat]\n scan_avals = [build_shaped_array(x, batch_dim=True) for x in xs_flat]\n in_avals = [*carry_avals, *scan_avals]\n in_tree = jax.tree_util.treedef_tuple((carry_tree, scan_tree))\n assert all(isinstance(a, core.AbstractValue) for a in in_avals), in_avals\n\n debug_info = jax.api_util.debug_info(\"flax scan\", broadcast_body,\n (in_tree,), {})\n f_flat, out_tree = jax.api_util.flatten_fun_nokwargs(\n lu.wrap_init(broadcast_body, debug_info=debug_info), in_tree\n )\n in_pvals = list(map(pe.PartialVal.unknown, in_avals))\n _, out_pvals, _ = pe.trace_to_jaxpr_nounits(f_flat, in_pvals)\n\n out_flat = []\n for pv, const in out_pvals:\n if pv is not None:\n raise ValueError(\n 'broadcasted variable has a data dependency on the scan body.'\n )\n out_flat.append(const)\n broadcast_in, constants_out = jax.tree_util.tree_unflatten(\n out_tree(), out_flat\n )\n\n if jax.version.__version_info__ > (0, 4, 25):\n c, ys = lax.scan(\n body_fn, init, xs, length=length, reverse=reverse, unroll=unroll,\n _split_transpose=_split_transpose\n )\n else:\n c, ys = lax.scan(\n body_fn, init, xs, length=length, reverse=reverse, unroll=unroll\n )\n ys = jax.tree_util.tree_map(transpose_from_front, out_axes, ys)\n ys = jax.tree_util.tree_map(\n lambda ax, const, y: (const if ax is broadcast else y),\n out_axes,\n constants_out,\n ys,\n )\n return broadcast_in, c, ys\n\n def simple_scan_fn(broadcast_in, init, *args):\n # Saves an extra tracing operation.\n # No verification of constancy, and no support for non-carry broadcast\n # function outputs.\n xs = jax.tree_util.tree_map(transpose_to_front, in_axes, args)\n\n if broadcast in jax.tree_util.tree_leaves(out_axes):\n raise ValueError(f\"nn.scan run with check_constancy_invariants=False \"\n f\"does not support broadcast non-carry function \"\n f\"outputs. out_axes was given as {out_axes}\")\n\n def body_fn(c, xs):\n # inject constants\n xs = jax.tree_util.tree_map(\n lambda ax, arg, x: (arg if ax is broadcast else x), in_axes, args, xs\n )\n _, c, ys = fn(broadcast_in, c, *xs)\n return c, ys\n\n if jax.version.__version_info__ > (0, 4, 25):\n c, ys = lax.scan(\n body_fn, init, xs, length=length, reverse=reverse, unroll=unroll,\n _split_transpose=_split_transpose\n )\n else:\n c, ys = lax.scan(\n body_fn, init, xs, length=length, reverse=reverse, unroll=unroll\n )\n ys = jax.tree_util.tree_map(transpose_from_front, out_axes, ys)\n return broadcast_in, c, ys\n\n if check_constancy_invariants:\n return scan_fn\n else:\n return simple_scan_fn\n" }, { "path": "flax/core/flax_functional_engine.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"code\",\n \"execution_count\": 1,\n \"metadata\": {\n \"colab_type\": \"code\"\n },\n \"outputs\": [],\n \"source\": [\n \"import functools\\n\",\n \"import jax\\n\",\n \"from jax import numpy as jnp, random, lax\\n\",\n \"import numpy as np\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 2,\n \"metadata\": {\n \"colab_type\": \"code\"\n },\n \"outputs\": [],\n \"source\": [\n \"from flax import linen as nn, struct\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 3,\n \"metadata\": {\n \"colab_type\": \"code\"\n },\n \"outputs\": [],\n \"source\": [\n \"from flax.core import Scope, init, apply, Array, lift, unfreeze\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 4,\n \"metadata\": {\n \"colab_type\": \"code\",\n \"outputId\": \"2558605e-e485-407e-b062-74d31cc49f1e\",\n \"tags\": []\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"FrozenDict({'params': {'kernel': DeviceArray([[ 0.15374057, -0.6807397 , -1.3350962 ],\\n\",\n \" [ 0.59940743, -0.69430196, -0.7663768 ]], dtype=float32), 'bias': DeviceArray([0., 0., 0.], dtype=float32)}})\\n\"\n ]\n },\n {\n \"data\": {\n \"text/plain\": [\n \"(DeviceArray([[0.17045607]], dtype=float32),\\n\",\n \" FrozenDict({'params': {'hidden': {'bias': DeviceArray([0., 0., 0.], dtype=float32), 'kernel': DeviceArray([[-0.22119394, 0.22075175, -0.0925657 ],\\n\",\n \" [ 0.40571952, 0.27750877, 1.0542233 ]], dtype=float32)}, 'out': {'kernel': DeviceArray([[ 0.21448377],\\n\",\n \" [-0.01530595],\\n\",\n \" [ 0.14402702]], dtype=float32), 'bias': DeviceArray([0.], dtype=float32)}}}))\"\n ]\n },\n \"execution_count\": 4,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"def dense(\\n\",\n \" scope: Scope,\\n\",\n \" inputs: Array,\\n\",\n \" features: int,\\n\",\n \" bias: bool = True,\\n\",\n \" kernel_init=nn.linear.default_kernel_init,\\n\",\n \" bias_init=nn.initializers.zeros_init(),\\n\",\n \"):\\n\",\n \" kernel = scope.param('kernel', kernel_init, (inputs.shape[-1], features))\\n\",\n \" y = jnp.dot(inputs, kernel)\\n\",\n \" if bias:\\n\",\n \" y += scope.param('bias', bias_init, (features,))\\n\",\n \" return y\\n\",\n \"\\n\",\n \"\\n\",\n \"model_fn = functools.partial(dense, features=3)\\n\",\n \"\\n\",\n \"x = jnp.ones((1, 2))\\n\",\n \"y, params = init(model_fn)(random.key(0), x)\\n\",\n \"print(params)\\n\",\n \"\\n\",\n \"\\n\",\n \"def mlp(scope: Scope, inputs: Array, features: int):\\n\",\n \" hidden = scope.child(dense, 'hidden')(inputs, features)\\n\",\n \" hidden = nn.relu(hidden)\\n\",\n \" return dense(scope.push('out'), hidden, 1)\\n\",\n \"\\n\",\n \"\\n\",\n \"init(mlp)(random.key(0), x, features=3)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 5,\n \"metadata\": {\n \"colab_type\": \"code\",\n \"outputId\": \"5790b763-df4f-47c8-9f4e-53fd1e1eb1fd\",\n \"tags\": []\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"[[ 0.11575121 -0.51936364 -1.113899 ]\\n\",\n \" [ 0.45569834 -0.5300623 -0.5873911 ]]\\n\",\n \"[ 0.45569834 -0.5300623 -0.5873911 ]\\n\",\n \"[[-1.5175114 -0.6617551]]\\n\"\n ]\n }\n ],\n \"source\": [\n \"@struct.dataclass\\n\",\n \"class Embedding:\\n\",\n \" table: np.ndarray\\n\",\n \"\\n\",\n \" def lookup(self, indices):\\n\",\n \" return self.table[indices]\\n\",\n \"\\n\",\n \" def attend(self, query):\\n\",\n \" return jnp.dot(query, self.table.T)\\n\",\n \"\\n\",\n \"\\n\",\n \"# all the embedding module does is provide a convenient initializers\\n\",\n \"\\n\",\n \"\\n\",\n \"def embedding(\\n\",\n \" scope: Scope,\\n\",\n \" num_embeddings: int,\\n\",\n \" features: int,\\n\",\n \" init_fn=nn.linear.default_embed_init,\\n\",\n \") -> Embedding:\\n\",\n \" table = scope.param('table', init_fn, (num_embeddings, features))\\n\",\n \" return Embedding(table)\\n\",\n \"\\n\",\n \"\\n\",\n \"embedding, _ = init(embedding)(random.key(0), num_embeddings=2, features=3)\\n\",\n \"print(embedding.table)\\n\",\n \"print(embedding.lookup(1))\\n\",\n \"print(\\n\",\n \" embedding.attend(\\n\",\n \" jnp.ones((\\n\",\n \" 1,\\n\",\n \" 3,\\n\",\n \" ))\\n\",\n \" )\\n\",\n \")\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 6,\n \"metadata\": {\n \"colab_type\": \"code\",\n \"outputId\": \"dd9c5079-10e7-4944-e09a-e9f65573a733\"\n },\n \"outputs\": [\n {\n \"data\": {\n \"text/plain\": [\n \"((((1, 3), (1, 3)), (1, 3)),\\n\",\n \" FrozenDict({'params': {'hf': {'bias': (3,), 'kernel': (3, 3)}, 'hg': {'bias': (3,), 'kernel': (3, 3)}, 'hi': {'bias': (3,), 'kernel': (3, 3)}, 'ho': {'bias': (3,), 'kernel': (3, 3)}, 'if': {'kernel': (2, 3)}, 'ig': {'kernel': (2, 3)}, 'ii': {'kernel': (2, 3)}, 'io': {'kernel': (2, 3)}}}))\"\n ]\n },\n \"execution_count\": 6,\n \"metadata\": {},\n \"output_type\": \"execute_result\"\n }\n ],\n \"source\": [\n \"def lstm(\\n\",\n \" scope,\\n\",\n \" carry,\\n\",\n \" inputs,\\n\",\n \" gate_fn=nn.activation.sigmoid,\\n\",\n \" activation_fn=nn.activation.tanh,\\n\",\n \" kernel_init=nn.linear.default_kernel_init,\\n\",\n \" recurrent_kernel_init=nn.initializers.orthogonal(),\\n\",\n \" bias_init=nn.initializers.zeros_init(),\\n\",\n \"):\\n\",\n \" r\\\"\\\"\\\"A long short-term memory (LSTM) cell.\\n\",\n \"\\n\",\n \" the mathematical definition of the cell is as follows\\n\",\n \" .. math::\\n\",\n \" \\\\begin{array}{ll}\\n\",\n \" i = \\\\sigma(W_{ii} x + W_{hi} h + b_{hi}) \\\\\\\\\\n\",\n \" f = \\\\sigma(W_{if} x + W_{hf} h + b_{hf}) \\\\\\\\\\n\",\n \" g = \\\\tanh(W_{ig} x + W_{hg} h + b_{hg}) \\\\\\\\\\n\",\n \" o = \\\\sigma(W_{io} x + W_{ho} h + b_{ho}) \\\\\\\\\\n\",\n \" c' = f * c + i * g \\\\\\\\\\n\",\n \" h' = o * \\\\tanh(c') \\\\\\\\\\n\",\n \" \\\\end{array}\\n\",\n \" where x is the input, h is the output of the previous time step, and c is\\n\",\n \" the memory.\\n\",\n \"\\n\",\n \" Args:\\n\",\n \" carry: the hidden state of the LSTM cell,\\n\",\n \" initialized using `LSTMCell.initialize_carry`.\\n\",\n \" inputs: an ndarray with the input for the current time step.\\n\",\n \" All dimensions except the final are considered batch dimensions.\\n\",\n \" gate_fn: activation function used for gates (default: sigmoid)\\n\",\n \" activation_fn: activation function used for output and memory update\\n\",\n \" (default: tanh).\\n\",\n \" kernel_init: initializer function for the kernels that transform\\n\",\n \" the input (default: lecun_normal).\\n\",\n \" recurrent_kernel_init: initializer function for the kernels that transform\\n\",\n \" the hidden state (default: orthogonal).\\n\",\n \" bias_init: initializer for the bias parameters (default: zeros_init())\\n\",\n \" Returns:\\n\",\n \" A tuple with the new carry and the output.\\n\",\n \" \\\"\\\"\\\"\\n\",\n \" c, h = carry\\n\",\n \" hidden_features = h.shape[-1]\\n\",\n \" # input and recurrent layers are summed so only one needs a bias.\\n\",\n \" dense_h = lambda name: scope.child(dense, name)(\\n\",\n \" h,\\n\",\n \" features=hidden_features,\\n\",\n \" bias=True,\\n\",\n \" kernel_init=recurrent_kernel_init,\\n\",\n \" bias_init=bias_init,\\n\",\n \" )\\n\",\n \" dense_i = lambda name: scope.child(dense, name)(\\n\",\n \" inputs, features=hidden_features, bias=False, kernel_init=kernel_init\\n\",\n \" )\\n\",\n \" i = gate_fn(dense_i(name='ii') + dense_h(name='hi'))\\n\",\n \" f = gate_fn(dense_i(name='if') + dense_h(name='hf'))\\n\",\n \" g = activation_fn(dense_i(name='ig') + dense_h(name='hg'))\\n\",\n \" o = gate_fn(dense_i(name='io') + dense_h(name='ho'))\\n\",\n \" new_c = f * c + i * g\\n\",\n \" new_h = o * activation_fn(new_c)\\n\",\n \" return (new_c, new_h), new_h\\n\",\n \"\\n\",\n \"\\n\",\n \"def lstm_init_carry(batch_dims, size, init_fn=jnp.zeros):\\n\",\n \" shape = batch_dims + (size,)\\n\",\n \" return init_fn(shape), init_fn(shape)\\n\",\n \"\\n\",\n \"\\n\",\n \"x = jnp.ones((1, 2))\\n\",\n \"carry = lstm_init_carry((1,), 3)\\n\",\n \"y, variables = init(lstm)(random.key(0), carry, x)\\n\",\n \"jax.tree_util.tree_map(np.shape, (y, variables))\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 9,\n \"metadata\": {\n \"tags\": []\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"initialized parameter shapes:\\n\",\n \" {'params': {'hf': {'bias': (2,), 'kernel': (2, 2)}, 'hg': {'bias': (2,), 'kernel': (2, 2)}, 'hi': {'bias': (2,), 'kernel': (2, 2)}, 'ho': {'bias': (2,), 'kernel': (2, 2)}, 'if': {'kernel': (2, 2)}, 'ig': {'kernel': (2, 2)}, 'ii': {'kernel': (2, 2)}, 'io': {'kernel': (2, 2)}}}\\n\"\n ]\n }\n ],\n \"source\": [\n \"def simple_scan(scope: Scope, xs):\\n\",\n \" init_carry = lstm_init_carry(xs.shape[:1], xs.shape[-1])\\n\",\n \" # cell = scope.child(lstm, 'cell')\\n\",\n \" # ys = []\\n\",\n \" # for i in range(xs.shape[1]):\\n\",\n \" # x = xs[:, i]\\n\",\n \" # init_carry, y = cell(init_carry, x)\\n\",\n \" # ys.append(y)\\n\",\n \" # return init_carry, ys\\n\",\n \" lstm_scan = lift.scan(\\n\",\n \" lstm,\\n\",\n \" in_axes=1,\\n\",\n \" out_axes=1,\\n\",\n \" variable_broadcast='params',\\n\",\n \" split_rngs={'params': False},\\n\",\n \" )\\n\",\n \" return lstm_scan(scope, init_carry, xs)\\n\",\n \"\\n\",\n \"\\n\",\n \"key1, key2 = random.split(random.key(0), 2)\\n\",\n \"xs = random.uniform(key1, (1, 5, 2))\\n\",\n \"\\n\",\n \"\\n\",\n \"y, init_variables = init(simple_scan)(key2, xs)\\n\",\n \"\\n\",\n \"print(\\n\",\n \" 'initialized parameter shapes:\\\\n',\\n\",\n \" jax.tree_util.tree_map(jnp.shape, unfreeze(init_variables)),\\n\",\n \")\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 10,\n \"metadata\": {\n \"tags\": []\n },\n \"outputs\": [\n {\n \"name\": \"stdout\",\n \"output_type\": \"stream\",\n \"text\": [\n \"output:\\n\",\n \" (DeviceArray([[-0.35626447, 0.25178757]], dtype=float32), DeviceArray([[-0.17885922, 0.13063088]], dtype=float32))\\n\"\n ]\n }\n ],\n \"source\": [\n \"y = apply(simple_scan)(init_variables, xs)[0]\\n\",\n \"print('output:\\\\n', y)\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": []\n }\n ],\n \"metadata\": {},\n \"nbformat\": 4,\n \"nbformat_minor\": 0\n}\n" }, { "path": "flax/core/frozen_dict.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Frozen Dictionary.\"\"\"\n\nimport collections\nfrom types import MappingProxyType\nfrom typing import Any, TypeVar\nfrom collections.abc import Hashable, Mapping\n\nimport jax\n\nfrom flax import serialization\n\n\nclass FrozenKeysView(collections.abc.KeysView):\n \"\"\"A wrapper for a more useful repr of the keys in a frozen dict.\"\"\"\n\n def __repr__(self):\n return f'frozen_dict_keys({list(self)})'\n\n\nclass FrozenValuesView(collections.abc.ValuesView):\n \"\"\"A wrapper for a more useful repr of the values in a frozen dict.\"\"\"\n\n def __repr__(self):\n return f'frozen_dict_values({list(self)})'\n\n\nK = TypeVar('K')\nV = TypeVar('V')\n\n\ndef _indent(x, num_spaces):\n indent_str = ' ' * num_spaces\n lines = x.split('\\n')\n assert not lines[-1]\n # skip the final line because it's empty and should not be indented.\n return '\\n'.join(indent_str + line for line in lines[:-1]) + '\\n'\n\n\n@jax.tree_util.register_pytree_with_keys_class\nclass FrozenDict(Mapping[K, V]):\n \"\"\"An immutable variant of the Python dict.\"\"\"\n\n __slots__ = ('_dict', '_hash')\n\n def __init__(self, *args, __unsafe_skip_copy__=False, **kwargs): # pylint: disable=invalid-name\n # make sure the dict is as\n xs = dict(*args, **kwargs)\n if __unsafe_skip_copy__:\n self._dict = xs\n else:\n self._dict = _prepare_freeze(xs)\n\n self._hash = None\n\n def __getitem__(self, key):\n v = self._dict[key]\n if isinstance(v, dict):\n return FrozenDict(v)\n return v\n\n def __setitem__(self, key, value):\n raise ValueError('FrozenDict is immutable.')\n\n def __contains__(self, key):\n return key in self._dict\n\n def __iter__(self):\n return iter(self._dict)\n\n def __len__(self):\n return len(self._dict)\n\n def __repr__(self):\n return self.pretty_repr()\n\n def __reduce__(self):\n return FrozenDict, (self.unfreeze(),)\n\n def get(self, key, default=None):\n \"\"\"Get an item from the FrozenDict.\"\"\"\n if key in self._dict:\n return self[key]\n return default\n\n def pretty_repr(self, num_spaces=4):\n \"\"\"Returns an indented representation of the nested dictionary.\"\"\"\n\n def pretty_dict(x):\n if not isinstance(x, dict):\n return repr(x)\n rep = ''\n for key, val in x.items():\n rep += f'{key}: {pretty_dict(val)},\\n'\n if rep:\n return '{\\n' + _indent(rep, num_spaces) + '}'\n else:\n return '{}'\n\n return f'FrozenDict({pretty_dict(self._dict)})'\n\n def __hash__(self):\n if self._hash is None:\n h = 0\n for key, value in self.items():\n h ^= hash((key, value))\n self._hash = h\n return self._hash\n\n def copy(\n self, add_or_replace: Mapping[K, V] = MappingProxyType({})\n ) -> 'FrozenDict[K, V]':\n \"\"\"Create a new FrozenDict with additional or replaced entries.\"\"\"\n return type(self)({**self, **unfreeze(add_or_replace)}) # type: ignore[arg-type]\n\n def keys(self):\n return FrozenKeysView(self)\n\n def values(self):\n return FrozenValuesView(self)\n\n def items(self):\n for key in self._dict:\n yield (key, self[key])\n\n def pop(self, key: K) -> tuple['FrozenDict[K, V]', V]:\n \"\"\"Create a new FrozenDict where one entry is removed.\n\n Example::\n\n >>> from flax.core import FrozenDict\n >>> variables = FrozenDict({'params': {...}, 'batch_stats': {...}})\n >>> new_variables, params = variables.pop('params')\n\n Args:\n key: the key to remove from the dict\n Returns:\n A pair with the new FrozenDict and the removed value.\n \"\"\"\n value = self[key]\n new_dict = dict(self._dict)\n new_dict.pop(key)\n new_self = type(self)(new_dict)\n return new_self, value\n\n def unfreeze(self) -> dict[K, V]:\n \"\"\"Unfreeze this FrozenDict.\n\n Returns:\n An unfrozen version of this FrozenDict instance.\n \"\"\"\n return unfreeze(self)\n\n def tree_flatten_with_keys(self) -> tuple[tuple[Any, ...], Hashable]:\n \"\"\"Flattens this FrozenDict.\n\n Returns:\n A flattened version of this FrozenDict instance.\n \"\"\"\n sorted_keys = sorted(self._dict)\n return tuple(\n [(jax.tree_util.DictKey(k), self._dict[k]) for k in sorted_keys]\n ), tuple(sorted_keys)\n\n @classmethod\n def tree_unflatten(cls, keys, values):\n # data is already deep copied due to tree map mechanism\n # we can skip the deep copy in the constructor\n return cls({k: v for k, v in zip(keys, values)}, __unsafe_skip_copy__=True)\n\n\ndef _prepare_freeze(xs: Any) -> Any:\n \"\"\"Deep copy unfrozen dicts to make the dictionary FrozenDict safe.\"\"\"\n if isinstance(xs, FrozenDict):\n # we can safely ref share the internal state of a FrozenDict\n # because it is immutable.\n return xs._dict # pylint: disable=protected-access\n if not isinstance(xs, dict):\n # return a leaf as is.\n return xs\n # recursively copy dictionary to avoid ref sharing\n return {key: _prepare_freeze(val) for key, val in xs.items()}\n\n\ndef freeze(xs: Mapping[Any, Any]) -> FrozenDict[Any, Any]:\n \"\"\"Freeze a nested dict.\n\n Makes a nested ``dict`` immutable by transforming it into ``FrozenDict``.\n\n Args:\n xs: Dictionary to freeze (a regualr Python dict).\n Returns:\n The frozen dictionary.\n \"\"\"\n return FrozenDict(xs)\n\n\ndef unfreeze(x: FrozenDict | dict[str, Any]) -> dict[Any, Any]:\n \"\"\"Unfreeze a FrozenDict.\n\n Makes a mutable copy of a ``FrozenDict`` mutable by transforming\n it into (nested) dict.\n\n Args:\n x: Frozen dictionary to unfreeze.\n Returns:\n The unfrozen dictionary (a regular Python dict).\n \"\"\"\n if isinstance(x, FrozenDict):\n # deep copy internal state of a FrozenDict\n # the dict branch would also work here but\n # it is much less performant because jax.tree_util.tree_map\n # uses an optimized C implementation.\n return jax.tree_util.tree_map(lambda y: y, x._dict) # type: ignore\n elif isinstance(x, dict):\n ys = {}\n for key, value in x.items():\n ys[key] = unfreeze(value)\n return ys\n else:\n return x\n\n\ndef copy(\n x: FrozenDict | dict[str, Any],\n add_or_replace: FrozenDict[str, Any] | dict[str, Any] = FrozenDict({}),\n) -> FrozenDict | dict[str, Any]:\n \"\"\"Create a new dict with additional and/or replaced entries. This is a utility\n function that can act on either a FrozenDict or regular dict and mimics the\n behavior of ``FrozenDict.copy``.\n\n Example::\n\n >>> from flax.core import FrozenDict, copy\n >>> variables = FrozenDict({'params': {...}, 'batch_stats': {...}})\n >>> new_variables = copy(variables, {'additional_entries': 1})\n\n Args:\n x: the dictionary to be copied and updated\n add_or_replace: dictionary of key-value pairs to add or replace in the dict x\n Returns:\n A new dict with the additional and/or replaced entries.\n \"\"\"\n\n if isinstance(x, FrozenDict):\n return x.copy(add_or_replace)\n elif isinstance(x, dict):\n new_dict = jax.tree_util.tree_map(\n lambda x: x, x\n ) # make a deep copy of dict x\n new_dict.update(add_or_replace)\n return new_dict\n raise TypeError(f'Expected FrozenDict or dict, got {type(x)}')\n\n\ndef pop(\n x: FrozenDict | dict[str, Any], key: str\n) -> tuple[FrozenDict | dict[str, Any], Any]:\n \"\"\"Create a new dict where one entry is removed. This is a utility\n function that can act on either a FrozenDict or regular dict and\n mimics the behavior of ``FrozenDict.pop``.\n\n Example::\n\n >>> from flax.core import FrozenDict, pop\n >>> variables = FrozenDict({'params': {...}, 'batch_stats': {...}})\n >>> new_variables, params = pop(variables, 'params')\n\n Args:\n x: the dictionary to remove the entry from\n key: the key to remove from the dict\n Returns:\n A pair with the new dict and the removed value.\n \"\"\"\n\n if isinstance(x, FrozenDict):\n return x.pop(key)\n elif isinstance(x, dict):\n new_dict = jax.tree_util.tree_map(\n lambda x: x, x\n ) # make a deep copy of dict x\n value = new_dict.pop(key)\n return new_dict, value\n raise TypeError(f'Expected FrozenDict or dict, got {type(x)}')\n\n\ndef pretty_repr(x: Any, num_spaces: int = 4) -> str:\n \"\"\"Returns an indented representation of the nested dictionary.\n This is a utility function that can act on either a FrozenDict or\n regular dict and mimics the behavior of ``FrozenDict.pretty_repr``.\n If x is any other dtype, this function will return ``repr(x)``.\n\n Args:\n x: the dictionary to be represented\n num_spaces: the number of space characters in each indentation level\n Returns:\n An indented string representation of the nested dictionary.\n \"\"\"\n\n if isinstance(x, FrozenDict):\n return x.pretty_repr()\n else:\n\n def pretty_dict(x):\n if not isinstance(x, dict):\n return repr(x)\n rep = ''\n for key, val in x.items():\n rep += f'{key}: {pretty_dict(val)},\\n'\n if rep:\n return '{\\n' + _indent(rep, num_spaces) + '}'\n else:\n return '{}'\n\n return pretty_dict(x)\n\n\ndef _frozen_dict_state_dict(xs):\n str_keys = {str(k) for k in xs.keys()}\n if len(str_keys) != len(xs):\n raise ValueError(\n 'Dict keys do not have a unique string representation: '\n f'{str_keys} vs given: {xs}'\n )\n return {str(key): serialization.to_state_dict(value) for key, value in xs.items()}\n\n\ndef _restore_frozen_dict(xs, states):\n diff = set(map(str, xs.keys())).difference(map(str, states.keys()))\n if diff:\n raise ValueError(\n 'The target dict keys and state dict keys do not match, target dict'\n f' contains keys {diff} which are not present in state dict at path'\n f' {serialization.current_path()}'\n )\n\n return FrozenDict(\n {\n key: serialization.from_state_dict(value, states[str(key)], name=key)\n for key, value in xs.items()\n }\n )\n\n\nserialization.register_serialization_state(\n FrozenDict, _frozen_dict_state_dict, _restore_frozen_dict\n)\n" }, { "path": "flax/core/lift.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Jax transform lifting.\"\"\"\n\nimport collections\nfrom collections.abc import Callable, Iterable, Mapping, Sequence\nimport contextlib\nimport dataclasses\nimport functools\nimport threading\nfrom typing import Any, Generic, TypeVar\nimport warnings\n\nfrom flax import traceback_util\nfrom flax import traverse_util\nfrom flax.typing import (\n In,\n InOutAxis,\n InOutScanAxis,\n Out,\n)\nimport jax\nfrom jax import random\n\nfrom . import axes_scan, meta\nfrom .frozen_dict import freeze, unfreeze\nfrom .scope import (\n CollectionFilter,\n DenyList, # pylint: disable=g-multiple-import\n Filter,\n LazyRng,\n PRNGSequenceFilter,\n Scope,\n group_collections,\n in_filter,\n intersect_filters,\n is_filter_empty,\n subtract_filters,\n union_filters,\n)\n\ntraceback_util.register_exclusion(__file__)\n\nA = TypeVar('A')\n\n\n@dataclasses.dataclass\nclass TransformContext(Generic[A], threading.local):\n \"\"\"Context for a transform.\"\"\"\n\n stack: list[A] = dataclasses.field(default_factory=list)\n\n @contextlib.contextmanager\n def push(self, a: A):\n self.stack.append(a)\n try:\n yield\n finally:\n self.stack.pop()\n\n def get(self) -> A:\n return self.stack[-1]\n\n\ndef tree_map_rngs(fn, tree):\n \"\"\"Needed for mapping JAX random.* functions over PRNGKey leaves.\"\"\"\n return jax.tree_util.tree_map(\n fn,\n tree,\n is_leaf=lambda x: hasattr(x, 'dtype')\n and jax.dtypes.issubdtype(x.dtype, jax.dtypes.prng_key),\n )\n\n\ndef _dedup_scopes(scopes):\n \"\"\"Deduplicated scopes.\"\"\"\n paths = []\n # must preseve insertion order for duplication to work correctly\n minimal_set = collections.OrderedDict((s, ()) for s in scopes)\n for leaf in scopes:\n scope = leaf.parent\n max_parent = leaf\n max_parent_path = ()\n path = [leaf.name]\n while scope is not None:\n if scope in minimal_set:\n max_parent = scope\n max_parent_path = tuple(reversed(path))\n path.append(scope.name)\n scope = scope.parent\n if max_parent is not leaf and leaf in minimal_set:\n del minimal_set[leaf]\n paths.append((max_parent, max_parent_path))\n return tuple(minimal_set), tuple(paths)\n\n\ndef _dup_scopes(orig_scopes, scopes, paths):\n \"\"\"Duplicated scopes.\"\"\"\n mapping = dict(zip(orig_scopes, scopes))\n scopes = []\n for root, path in paths:\n scope = mapping[root]\n for name in path:\n scope = scope.push(name, reuse=True)\n scopes.append(scope)\n return scopes\n\n\ndef _transpose(xs):\n return tuple(zip(*xs))\n\n\ndef _partial_pack(\n scope_tree: Scope,\n in_variable_filters: Sequence[CollectionFilter],\n out_variable_filters: Sequence[CollectionFilter],\n rng_filters: Sequence[PRNGSequenceFilter],\n name=None,\n) -> tuple[Callable[..., Any], Callable[..., Any], Any, Any, Callable[..., Any]]:\n \"\"\"Pack variables and rngs for functional transformations.\n\n The _partial_pack function is the building block for all other lifted transformations.\n\n Args:\n fn: The function to pack. `fn` has the signature\n in_variable_filters: Input variable filters.\n out_variable_filters: Output variable filters.\n rng_filters: RNG filters.\n name: The name of the packed scope.\n enable_kwargs: Whether to enable kwargs or not.\n Returns:\n `(scope_fn, repack_fn, variable_groups, rng_groups, publish_results_fn)`\n \"\"\"\n # pylint: disable=protected-access\n scopes, treedef = jax.tree_util.tree_flatten(scope_tree)\n scopes, paths = _dedup_scopes(scopes)\n\n variable_groups_xs = []\n\n for scope in scopes:\n scope._validate_trace_level()\n scope._populate_collections()\n variable_groups_xs.append(\n group_collections(scope._variables, in_variable_filters)\n )\n variable_groups_xs_t = _transpose(variable_groups_xs)\n\n # Make sure that in-only variable collections are frozen\n for variable_group_xs in variable_groups_xs_t:\n for variable_group in variable_group_xs:\n for col_name, collection in variable_group.items():\n col_in_out = any(\n in_filter(col_filter, col_name)\n for col_filter in out_variable_filters\n )\n if not col_in_out:\n variable_group[col_name] = freeze(collection)\n rng_groups_xs = []\n inner_rng_counters = []\n for scope in scopes:\n rng_counters = scope.rng_counters\n rng_groups = group_collections(scope.rngs, rng_filters)\n rng_groups_xs.append(rng_groups)\n inner_rng_counters.append(rng_counters)\n rng_groups_xs_t = _transpose(rng_groups_xs)\n\n def scope_fn(\n variable_groups_xs_t,\n rng_groups_xs_t,\n mutable_filter: CollectionFilter = True,\n ):\n inner_scopes = []\n mutable: Filter = False\n for out_filter in out_variable_filters:\n mutable = union_filters(mutable, out_filter)\n # could be () in the edge case where no rngs or variable_groups are lifted\n # in this case fallback to ((),) * len(scopes) to make sure the zip has\n # something to iterate over for each scope.\n variable_groups_xs = _transpose(variable_groups_xs_t) or ((),) * len(\n scopes\n )\n rng_groups_xs = _transpose(rng_groups_xs_t) or ((),) * len(scopes)\n assert len(variable_groups_xs) == len(scopes)\n assert len(rng_groups_xs) == len(scopes)\n for variable_groups, rng_groups, scope, rng_counters in zip(\n variable_groups_xs, rng_groups_xs, scopes, inner_rng_counters\n ):\n variables = {}\n rngs = {}\n for variable_group in variable_groups:\n variables.update(variable_group)\n for rng_group in rng_groups:\n rngs.update(rng_group)\n # make sure variable dicts are cloned and can't be manipulated by ref\n # sharing.\n variables = jax.tree_util.tree_map(lambda x: x, variables)\n scope_mutable = intersect_filters(\n intersect_filters(scope.mutable, mutable), mutable_filter\n )\n new_debug_path = scope.debug_path\n if name:\n if new_debug_path:\n new_debug_path = new_debug_path[:-1] + (\n f'{name}({new_debug_path[-1]})',\n )\n else:\n new_debug_path = (f'{name}()',)\n inner_scope = Scope(\n variables,\n name=scope.name,\n rngs=rngs,\n mutable=scope_mutable,\n parent=None,\n path=scope.path,\n debug_path=new_debug_path,\n flags=scope.flags,\n )\n inner_scope.rng_counters = rng_counters\n inner_scopes.append(inner_scope)\n inner_scopes = _dup_scopes(scopes, inner_scopes, paths)\n return treedef.unflatten(inner_scopes)\n\n def repack_fn(inner_scope_tree):\n inner_scopes = treedef.flatten_up_to(inner_scope_tree)\n inner_scopes, inner_paths = _dedup_scopes(inner_scopes)\n inner_scopes = list(inner_scopes)\n assert [p for _, p in paths] == [p for _, p in inner_paths]\n out_variable_groups_xs = []\n for inner_scope in inner_scopes:\n inner_scope.invalidate()\n inner_scope._validate_trace_level()\n mutable_variables = {\n key: val\n for key, val in inner_scope._variables.items()\n if in_filter(inner_scope.mutable, key)\n }\n out_variable_groups = group_collections(\n mutable_variables, tuple(out_variable_filters) + (True,)\n )\n remainder = tuple(out_variable_groups[-1].keys())\n if remainder:\n raise ValueError(f'unmapped output variables: {remainder}')\n out_variable_groups_xs.append(out_variable_groups[:-1])\n\n return _transpose(out_variable_groups_xs)\n\n def publish_results_fn(out_variable_groups_xs_t):\n out_variable_groups_xs = _transpose(out_variable_groups_xs_t)\n for scope, out_variable_groups, rng_counters in zip(\n scopes, out_variable_groups_xs, inner_rng_counters\n ):\n for out_variable_group in out_variable_groups:\n for col_name, collection in out_variable_group.items():\n if not scope.is_mutable_collection(col_name):\n # Some lifted transforms like scan return redundant variables.\n continue\n for var_name, value in collection.items():\n scope.put_variable(col_name, var_name, value)\n\n return (\n scope_fn,\n repack_fn,\n variable_groups_xs_t,\n rng_groups_xs_t,\n publish_results_fn,\n )\n\n\ndef pack(\n fn: Callable[..., Any],\n in_variable_filters: Sequence[CollectionFilter],\n out_variable_filters: Sequence[CollectionFilter],\n rng_filters: Sequence[PRNGSequenceFilter],\n name=None,\n enable_kwargs=False,\n) -> Callable[..., Any]:\n \"\"\"Pack variables and rngs for functional transformations.\n\n The pack function is the building block for all other lifted transformations.\n\n Args:\n fn: The function to pack. `fn` has the signature\n `(scope_fn, repack_fn, variable_groups, rng_groups, *args) ->\n (output, packed_variables)`.\n in_variable_filters: Input variable filters.\n out_variable_filters: Output variable filters.\n rng_filters: RNG filters.\n name: The name of the packed scope.\n enable_kwargs: Whether to enable kwargs or not.\n Returns:\n A callable which expects a scope as the first argument.\n \"\"\"\n\n @functools.wraps(fn)\n def wrapper(scope_tree: Scope, *args, **kwargs):\n if not enable_kwargs and kwargs:\n msg = 'kwargs are not supported in {}, so \"{}\" is(are) ignored'\n warnings.warn(msg.format(name, ', '.join(kwargs.keys())), RuntimeWarning)\n (\n scope_fn,\n repack_fn,\n variable_groups_xs_t,\n rng_groups_xs_t,\n publish_results_fn,\n ) = _partial_pack(scope_tree, in_variable_filters, out_variable_filters, rng_filters, name)\n if enable_kwargs:\n y, out_variable_groups_xs_t = fn(\n scope_fn,\n repack_fn,\n variable_groups_xs_t,\n rng_groups_xs_t,\n *args,\n **kwargs,\n )\n else:\n y, out_variable_groups_xs_t = fn(\n scope_fn, repack_fn, variable_groups_xs_t, rng_groups_xs_t, *args\n )\n publish_results_fn(out_variable_groups_xs_t)\n return y\n\n return wrapper\n\n\nid_fn = lambda x: x\n\n\ndef map_variables(\n fn: Callable[..., Any],\n mapped_collections: CollectionFilter,\n map_in_fn: Callable[..., Any] = id_fn,\n map_out_fn: Callable[..., Any] = id_fn,\n init: bool = False,\n mutable: bool = False,\n rngs: PRNGSequenceFilter = True,\n variables: CollectionFilter = True,\n) -> Callable[..., Any]:\n \"\"\"Map Variables inside a scope.\n\n Args:\n fn: the function to be transformed.\n mapped_collections: the collection(s) to be transformed.\n map_in_fn: creates a view of the target variables.\n map_out_fn: transforms the updated variables in the view after mutation.\n init: If True, variables are initialized before transformation.\n mutable: If True, the mapped variable collections will be mutable.\n rngs: PRNGSequences added to the transformed scope (default: all).\n variables: Additional Variable collections added to the transformed scope.\n Besides those specified by `target` (default: all).\n\n Returns:\n A callable expecting a scope as the first argument.\n \"\"\"\n is_target_out = mutable or init\n\n def wrapper(scope_fn, repack, variable_groups, rng_groups, *args, **kwargs):\n target, variables = variable_groups\n if init:\n scopes = scope_fn((target, variables), rng_groups)\n has_mutable_cols = any(\n not is_filter_empty(scope.mutable)\n for scope in jax.tree_util.tree_leaves(scopes)\n )\n if has_mutable_cols:\n fn(scopes, *args, **kwargs)\n target, _ = repack(scopes)\n target = tuple(map_out_fn(x) for x in target)\n target = tuple(map_in_fn(unfreeze(x)) for x in target)\n mfilter = True\n if not is_target_out:\n # mapped collections should not be mutable\n # unless the mapping supports it (by init=True or mutable=True)\n mfilter = subtract_filters(mfilter, mapped_collections)\n scopes = scope_fn((target, variables), rng_groups, mutable_filter=mfilter)\n y = fn(scopes, *args, **kwargs)\n out_target, out_vars = repack(scopes)\n if is_target_out:\n out_target = tuple(map_out_fn(x) for x in out_target)\n return y, (out_target, out_vars)\n\n in_vars = (mapped_collections, variables)\n out_vars = (\n in_vars\n if is_target_out\n else (False, subtract_filters(variables, mapped_collections))\n )\n return pack(\n wrapper,\n in_vars,\n out_vars,\n (rngs,),\n enable_kwargs=True,\n name='map_variables',\n )\n\n\ndef swap_collection(fn: Callable[..., Any], col_a: str, col_b: str):\n \"\"\"Swap two collections.\"\"\"\n\n def swap(target):\n a = target[col_a] if col_a in target else {}\n b = target[col_b] if col_b in target else {}\n target[col_b], target[col_a] = a, b\n return target\n\n return map_variables(fn, (col_a, col_b), swap, swap, mutable=True)\n\n\ndef _split_in_out_axes(xs: Mapping[CollectionFilter, Any]):\n unpack = lambda v: v.axis if isinstance(v, (In, Out)) else v\n in_axes = {k: unpack(v) for k, v in xs.items() if not isinstance(v, Out)}\n out_axes = {k: unpack(v) for k, v in xs.items() if not isinstance(v, In)}\n return in_axes, out_axes\n\n\ndef _bwd_wrapper(treedef, bwd_fn, tangent):\n vars_grad, *inputs_grad = bwd_fn(tangent)\n vars_grad = treedef.unflatten(vars_grad)\n return (vars_grad, *inputs_grad)\n\n\ndef vjp(\n fn: Callable[..., Any],\n scope: Scope,\n *primals,\n has_aux: bool = False,\n reduce_axes=(),\n vjp_variables: CollectionFilter = 'params',\n variables: CollectionFilter = True,\n rngs: PRNGSequenceFilter = True,\n) -> tuple[Any, Callable[..., Any]] | tuple[Any, Callable[..., Any], Any]:\n \"\"\"A lifted version of ``jax.vjp``.\n\n See ``jax.vjp`` for the unlifted vector-Jacobian product (backward gradient).\n\n Note that a gradient is returned for all variables in the collections\n specified by `vjp_variables`. However, the backward function only expects\n a cotangent for the return value of `fn`. If variables require a co-tangent\n as well they can be returned from `fn` using `scope.variables()`.\n\n Example::\n\n def learn_scale(scope, x, y):\n p = scope.param('scale', nn.initializers.zeros_init(), ())\n return p * x * y\n def f(scope, x, y):\n z, bwd = lift.vjp(learn_scale, scope, x, y)\n params_grad, x_grad, y_grad = bwd(jnp.ones(z.shape))\n return z, params_grad, x_grad, y_grad\n\n Args:\n fn: Function to be differentiated. Its arguments should be arrays, scalars,\n or standard Python containers of arrays or scalars. It should return an\n array, scalar, or standard Python container of arrays or scalars. It will\n receive the scope and primals as arguments.\n scope: The scope of which the variables will be differentiated.\n *primals: A sequence of primal values at which the Jacobian of ``fn``\n should be evaluated. The length of ``primals`` should be equal to the\n number of positional parameters to ``fn``. Each primal value should be a\n tuple of arrays, scalar, or standard Python containers thereof.\n has_aux: Optional, bool. Indicates whether ``fn`` returns a pair where the\n first element is considered the output of the mathematical function to be\n differentiated and the second element is auxiliary data. Default ``False``.\n vjp_variables: The vjpfun will return a cotangent vector for all\n variable collections specified by this filter.\n variables: other variables collections that are available inside `fn` but\n do not receive a cotangent.\n rngs: the prngs that are available inside `fn`.\n\n Returns:\n If ``has_aux`` is ``False``, returns a ``(primals_out, vjpfun)`` pair, where\n ``primals_out`` is ``fn(*primals)``.\n ``vjpfun`` is a function from a cotangent vector with the same shape as\n ``primals_out`` to a tuple of cotangent vectors with the same shape as\n ``primals``, representing the vector-Jacobian product of ``fn`` evaluated at\n ``primals``. If ``has_aux`` is ``True``, returns a\n ``(primals_out, vjpfun, aux)`` tuple where ``aux`` is the auxiliary data\n returned by ``fn``.\n \"\"\"\n if reduce_axes:\n raise NotImplementedError('reduce_axes argument to vjp is deprecated')\n del reduce_axes\n\n def inner(scope_fn, repack_fn, variable_groups, rng_groups, *args):\n vjp_vars, other_vars = variable_groups\n\n @functools.wraps(fn)\n def wrapper(vjp_vars, *args):\n variable_groups = (vjp_vars, other_vars)\n scope = scope_fn(variable_groups, rng_groups)\n if has_aux:\n y, aux = fn(scope, *args)\n else:\n y = fn(scope, *args)\n aux = ()\n return y, (aux, repack_fn(scope))\n\n y, bwd, (aux, out_vars) = jax.vjp(\n wrapper, vjp_vars, *args, has_aux=True\n )\n treedef = jax.tree_util.tree_structure(scope)\n bwd = jax.tree_util.Partial(functools.partial(_bwd_wrapper, treedef), bwd)\n if has_aux:\n return (y, bwd, aux), out_vars\n else:\n return (y, bwd), out_vars\n\n return pack(\n inner,\n (vjp_variables, variables),\n (variables,),\n (rngs,),\n name='vjp',\n enable_kwargs=False,\n )(scope, *primals)\n\n\ndef value_and_grad(\n fn: Callable[..., Any],\n scope: Scope,\n *primals,\n has_aux: bool = False,\n reduce_axes=(),\n variables: CollectionFilter = True,\n rngs: PRNGSequenceFilter = True,\n) -> tuple[Any, Callable[..., Any]] | tuple[Any, Callable[..., Any], Any]:\n \"\"\"A limited lifted version of ``jax.value_and_grad``.\n\n See ``jax.value_and_grad`` for the unlifted reverse mode gradient.\n\n Note that for this convenience function, gradients are only calculated for\n the function inputs (all function inputs), and not with respect to any scope\n variables. The target function must return a scalar-valued output.\n\n Example::\n\n def learn_scale(scope, x, y):\n p = scope.param('scale', nn.initializers.zeros_init(), ())\n return p * x * y\n def f(scope, x, y):\n z, x_grad, y_grad = lift.value_and_grad(learn_scale, scope, x, y)\n return z, x_grad, y_grad\n\n Args:\n fn: Function to be differentiated. Its arguments should be arrays, scalars,\n or standard Python containers of arrays or scalars. It should return an\n array, scalar, or standard Python container of arrays or scalars. It will\n receive the scope and primals as arguments.\n scope: The scope of which the variables will be differentiated.\n *primals: A sequence of primal values at which the Jacobian of ``fn``\n should be evaluated. The length of ``primals`` should be equal to the\n number of positional parameters to ``fn``. Each primal value should be a\n tuple of arrays, scalar, or standard Python containers thereof.\n has_aux: Optional, bool. Indicates whether ``fn`` returns a pair where the\n first element is considered the output of the mathematical function to be\n differentiated and the second element is auxiliary data. Default ``False``.\n variables: other variables collections that are available inside `fn` but\n do not receive a cotangent.\n rngs: the prngs that are available inside `fn`.\n\n Returns:\n If ``has_aux`` is ``False``, returns a ``(primals_out, grads)`` pair, where\n ``primals_out`` is ``fn(*primals)``.\n If ``has_aux`` is ``True``, returns a\n ``(primals_out, aux, grads)`` tuple where ``aux`` is the auxiliary data\n returned by ``fn``.\n \"\"\"\n if reduce_axes:\n raise NotImplementedError(\n 'reduce_axes argument to value_and_grad is deprecated')\n del reduce_axes\n\n def inner(scope_fn, repack_fn, variable_groups, rng_groups, *args):\n @functools.wraps(fn)\n def wrapper(*args):\n scope = scope_fn(variable_groups, rng_groups)\n if has_aux:\n y, aux = fn(scope, *args)\n else:\n y = fn(scope, *args)\n aux = ()\n return y, (aux, repack_fn(scope))\n\n y, bwd, (aux, out_vars) = jax.vjp(\n wrapper,\n *args,\n has_aux=True,\n )\n\n inputs_grad = bwd(jax.numpy.ones_like(y))\n\n if has_aux:\n return (y, aux, inputs_grad), out_vars\n else:\n return (y, inputs_grad), out_vars\n\n return pack(\n inner,\n (variables,),\n (variables,),\n (rngs,),\n name='value_and_grad',\n enable_kwargs=False,\n )(scope, *primals)\n\n\ndef jvp(\n fn: Callable[..., Any],\n scope: Scope,\n primals,\n tangents,\n variable_tangents,\n variables: CollectionFilter = True,\n rngs: PRNGSequenceFilter = True,\n) -> tuple[Any, Any]:\n \"\"\"A lifted version of ``jax.jvp``.\n\n See ``jax.jvp`` for the unlifted Jacobian-vector product (forward gradient).\n\n Note that no tangents are returned for variables. When variable tangents\n are required their value should be returned explicitly by `fn`\n using `scope.variables()`.\n\n Example::\n\n def learn_scale(scope, x):\n p = scope.param('scale', nn.initializers.zeros_init(), ())\n return p * x\n\n def f(scope, x):\n vars_t = jax.tree_util.tree_map(jnp.ones_like,\n scope.variables().get('params', {}))\n x, out_t = lift.jvp(\n learn_scale, scope, (x,), (jnp.zeros_like(x),),\n variable_tangents={'params': vars_t})\n return out_t\n\n Args:\n fn: The function to be transformed.\n scope: The scope(s) which should be lifted into the transform.\n primals: The primal values at which the Jacobian of ``fun`` should be\n evaluated. Should be either a tuple or a list of arguments,\n and its length should be equal to the number of positional parameters of\n ``fun``.\n tangents: The tangent vector for which the Jacobian-vector product should be\n evaluated. Should be either a tuple or a list of tangents, with the same\n tree structure and array shapes as ``primals``.\n variable_tangents: A dict or PyTree fo dicts with the same structure as\n scopes. Each entry in the dict specifies the tangents for a variable\n collection. Not specifying a collection in variable_tangents is\n equivalent to passing a zero vector as the tangent.\n variables: other variables collections that are available inside `fn` but\n do not receive a tangent.\n rngs: the prngs that are available inside `fn`.\n\n Returns:\n A ``(primals_out, tangents_out)`` pair, where ``primals_out`` is\n ``fun(*primals)``, and ``tangents_out`` is the Jacobian-vector product of\n ``function`` evaluated at ``primals`` with ``tangents``. The\n ``tangents_out`` value has the same Python tree structure and shapes as\n ``primals_out``.\n \"\"\"\n\n def inner(scope_fn, repack_fn, variable_groups, rng_groups, *args):\n jvp_vars, other_vars = variable_groups\n\n @functools.wraps(fn)\n def wrapper(vars_primals, args):\n variable_groups = (vars_primals, other_vars)\n scope = scope_fn(variable_groups, rng_groups)\n y = fn(scope, *args)\n return y, repack_fn(scope)\n\n (y, out_vars), out_tangents = jax.jvp(\n wrapper, (jvp_vars, args), (variable_tangents, tangents)\n )\n return (y, out_tangents[0]), out_vars\n\n # filter out empty tangent collections because JAX will error on non-equal\n # tree structure for example: {\"params\": {}} != {}.\n treedef = jax.tree_util.tree_structure(scope)\n\n variable_tangents = tuple(\n {k: v for k, v in vt.items() if v} # pylint: disable=g-complex-comprehension\n for vt in treedef.flatten_up_to(variable_tangents)\n )\n target = tuple(variable_tangents[0].keys())\n return pack(\n inner,\n (target, variables),\n (variables,),\n (rngs,),\n name='jvp',\n enable_kwargs=False,\n )(scope, *primals)\n\n\ndef vmap(\n fn: Callable[..., Any],\n variable_axes: Mapping[CollectionFilter, InOutAxis],\n split_rngs: Mapping[PRNGSequenceFilter, bool],\n in_axes=0,\n out_axes=0,\n axis_size: int | None = None,\n axis_name: str | None = None,\n spmd_axis_name: str | None = None,\n metadata_params: dict[Any, Any] = {},\n) -> Callable[..., Any]:\n \"\"\"A lifted version of ``jax.vmap``.\n\n See ``jax.vmap`` for the unlifted batch transform in Jax.\n\n ``vmap`` can be used to add a batch axis to a scope function.\n For example we could create a version of ``dense`` with\n a batch axis that does not share parameters::\n\n batch_dense = lift.vmap(\n nn.dense,\n in_axes=(0, None),\n variable_axes={'params': 0},\n split_rngs={'params': True})\n\n By using ``variable_axes={'params': 0}``, we indicate that the\n parameters themselves are mapped over and therefore not shared along\n the mapped axis. Consequently, we also split the 'params' RNG,\n otherwise the parameters would be initialized identically along\n the mapped axis.\n\n Similarly, ``vmap`` could be use to add a batch axis with parameter\n sharing::\n\n batch_foo = lift.vmap(\n foo,\n in_axes=0, out_axes=0,\n variable_axes={'params': None},\n split_rngs={'params': False})\n\n Here we use ``variable_axes={'params': None}`` to indicate the parameter\n variables are shared along the mapped axis. Consequently, the 'params'\n RNG must also be shared.\n\n Args:\n fn: the function to be transformed.\n variable_axes: the variable collections that are lifted into the batching\n transformation. Use `None` to indicate a broadcasted collection or an\n integer to map over an axis.\n split_rngs: Split PRNG sequences will be different for each index of the\n batch dimension. Unsplit PRNGs will be broadcasted.\n in_axes: Specifies the mapping of the input arguments (see `jax.vmap).\n out_axes: Specifies the mapping of the return value (see `jax.vmap).\n axis_size: Specifies the size of the batch axis. This only needs to be\n specified if it cannot be derived from the input arguments.\n axis_name: Specifies a name for the batch axis. Can be used together with\n parallel reduction primitives (e.g. `jax.lax.pmean`, `jax.lax.ppermute`,\n etc.). Note, this is only used for pmap and shmap. For SPMD jit, you do\n not need to manually synchronize. Just make sure that the axes are\n correctly annotated and XLA:SPMD will insert the necessary collectives.\n spmd_axis_name: Axis name added to any pjit sharding constraints appearing\n in `fn`. See also\n https://github.com/google/flax/blob/main/flax/linen/partitioning.py.\n metadata_params: arguments dict passed to AxisMetadata instances in the\n variable tree.\n\n Returns:\n A vectorized version of the input scope function.\n \"\"\"\n variable_in_axes, variable_out_axes = _split_in_out_axes(variable_axes)\n variable_in_groups, variable_in_axes = _unzip2(variable_in_axes.items())\n variable_out_groups, variable_out_axes = _unzip2(variable_out_axes.items())\n rng_groups, rng_splits = _unzip2(split_rngs.items())\n rng_axes = tuple(0 if rng_split else None for rng_split in rng_splits)\n\n def inner(scope_fn, repack_fn, variable_groups, rng_groups, *args):\n # optional user-defined variable transform on the way in\n new_variable_groups = []\n for var_group, axis in zip(variable_groups, variable_in_axes):\n if axis is not None:\n new_variable_groups.append(\n meta.remove_axis(var_group, axis, metadata_params)\n )\n else:\n new_variable_groups.append(var_group)\n variable_groups = tuple(new_variable_groups)\n\n # split rngs\n def find_axis_size(axis, x):\n if axis is not None:\n leaves = jax.tree_util.tree_leaves(x)\n if leaves:\n return leaves[0].shape[axis]\n return ()\n\n axis_sizes = jax.tree_util.tree_map(\n find_axis_size, (variable_in_axes, in_axes), (variable_groups, args),\n is_leaf=lambda x: x is None\n )\n axis_sizes = set(jax.tree_util.tree_leaves(axis_sizes))\n if axis_size is None and len(axis_sizes) == 1:\n (d_axis_size,) = axis_sizes\n elif len(axis_sizes) > 1:\n raise ValueError(f'Inconsistent batch axis sizes: {axis_sizes}')\n elif axis_size is None:\n raise ValueError('axis_size should be specified manually.')\n else:\n d_axis_size = axis_size\n\n def split_fn(rng):\n # random.clone is only available on Jax versions 0.4.26 or newer. See:\n # https://jax.readthedocs.io/en/latest/jax.experimental.key_reuse.htmls\n if hasattr(random, 'clone'):\n rng = random.clone(rng)\n rngs = random.split(rng, d_axis_size)\n if spmd_axis_name is not None:\n args_flat, _ = jax.tree.flatten(args)\n axes_flat = _broadcast_prefix_tree(in_axes, args)\n any_vmapped_axis_sharded = any(\n jax.typeof(x).sharding.spec[i] == spmd_axis_name\n for x, i in zip(args_flat, axes_flat)\n if i is not None\n )\n if any_vmapped_axis_sharded:\n rngs = jax.sharding.reshard(rngs, jax.P(spmd_axis_name))\n return rngs\n\n rng_groups = tuple(\n tree_map_rngs(split_fn, rng_group) if split else rng_group\n for rng_group, split in zip(rng_groups, rng_splits)\n )\n\n @functools.partial(\n jax.vmap,\n in_axes=(variable_in_axes, rng_axes, in_axes),\n out_axes=(out_axes, variable_out_axes),\n axis_name=axis_name,\n axis_size=axis_size,\n spmd_axis_name=spmd_axis_name,\n )\n @functools.wraps(fn)\n def mapped(variable_groups, rng_groups, args):\n scope = scope_fn(variable_groups, rng_groups)\n y = fn(scope, *args)\n return y, repack_fn(scope)\n\n # optional user-defined variable transform on the way out\n y, vars_out = mapped(variable_groups, rng_groups, args)\n new_vars_out = []\n for var_group, axis in zip(vars_out, variable_out_axes):\n if axis is not None:\n new_vars_out.append(meta.add_axis(var_group, axis, metadata_params))\n else:\n new_vars_out.append(var_group)\n vars_out = tuple(new_vars_out)\n return y, vars_out\n\n return pack(\n inner, variable_in_groups, variable_out_groups, rng_groups, name='vmap'\n )\n\n\ndef scan(\n fn: Callable[..., Any],\n variable_axes: Mapping[CollectionFilter, InOutScanAxis] = {},\n variable_broadcast: CollectionFilter = False,\n variable_carry: CollectionFilter = False,\n split_rngs: Mapping[PRNGSequenceFilter, bool] = {},\n in_axes=0,\n out_axes=0,\n length: int | None = None,\n reverse: bool = False,\n unroll: int = 1,\n _split_transpose: bool = False,\n data_transform: Callable[..., Any] | None = None,\n metadata_params: dict[Any, Any] = {},\n check_constancy_invariants: bool = True,\n) -> Callable[..., Any]:\n \"\"\"A lifted version of ``jax.lax.scan``.\n\n See ``jax.lax.scan`` for the unlifted scan in Jax.\n\n To improve consistency with ``vmap``, this version of scan\n uses ``in_axes`` and ``out_axes`` to determine which arguments\n are scanned over and along which axis.\n\n ``scan`` distinguishes between 3 different types of values inside the loop:\n\n 1. **scan**: a value that is iterated over in a loop. All scan values must\n have the same size in the axis they are scanned over. Scanned outputs\n will be stacked along the scan axis.\n 2. **carry**: A carried value is updated at each loop iteration. It must\n have the same shape and dtype throughout the loop.\n 3. **broadcast**: a value that is closed over by the loop. When a variable\n is broadcasted they are typically initialized inside the loop body but\n independent of the loop variables.\n\n The loop body should have the signature\n ``(scope, body, carry, *xs) -> (carry, ys)``, where ``xs`` and ``ys``\n are the scan values that go in and out of the loop.\n\n Example::\n\n scope.variable('counter', 'i', jnp.zeros, ())\n def body_fn(scope, c, x):\n counter = scope.variable('counter', 'i', jnp.zeros, ())\n counter.value += 1\n x = scope.child(nn.dense)(x, 1)\n return c, x\n\n _, ys = lift.scan(\n body_fn,\n variable_carry='counter',\n variable_broadcast='params',\n split_rngs={'params': False})(scope, (), xs)\n\n Args:\n fn: the function to be transformed.\n variable_axes: the variable collections that are scanned over.\n variable_broadcast: Specifies the broadcasted variable collections.\n A broadcasted variable should not depend on any computation that cannot b\n lifted out of the loop. This is typically used to define shared parameters\n inside the fn.\n variable_carry: Specifies the variable collections that are carried through\n the loop. Mutations to these variables are carried to the next iteration\n and will be preserved when the scan finishes.\n split_rngs: Split PRNG sequences will be different for each loop iterations.\n If split is False the PRNGs will be the same across iterations.\n in_axes: Specifies the axis to scan over for the arguments. Should be a\n prefix tree of the arguments. Use `flax.core.broadcast` to feed an entire\n input to each iteration of the scan body.\n out_axes: Specifies the axis to scan over for the return value. Should be a\n prefix tree of the return value.\n length: Specifies the number of loop iterations. This only needs\n to be specified if it cannot be derived from the scan arguments.\n reverse: If true, scan from end to start in reverse order.\n unroll: how many scan iterations to unroll within a single\n iteration of a loop (default: 1).\n _split_transpose: An experimental feature to split the transpose of a scan\n into a scan and a map, backed by an experimental Jax lax.scan() feature.\n data_transform: optional function to transform raw variable and rng groups,\n intended for inline SPMD annotations.\n metadata_params: arguments dict passed to AxisMetadata instances in the\n variable tree.\n check_constancy_invariants: If true, the scan will verify that the\n broadcast constants are true loop invariants, and further supports\n broadcast function (non-carry) outputs. This requires an extra jax\n tracing step however, so setting to false can reduce trace time on larger\n models.\n\n Returns:\n The scan function with the signature\n ``(scope, carry, *xxs) -> (carry, yys)``, where ``xxs`` and ``yys`` are the\n scan values that go in and out of the loop.\n \"\"\"\n variable_in_axes, variable_out_axes = _split_in_out_axes(variable_axes)\n variable_in_groups, variable_in_axes = _unzip2(variable_in_axes.items())\n variable_out_groups, variable_out_axes = _unzip2(variable_out_axes.items())\n assert all(isinstance(ax, int) for ax in variable_in_axes)\n assert all(isinstance(ax, int) for ax in variable_out_axes)\n rng_groups, rng_splits = _unzip2(split_rngs.items())\n rng_axes = tuple(\n 0 if rng_split else axes_scan.broadcast for rng_split in rng_splits\n )\n\n def inner(scope_fn, repack_fn, variable_groups, rng_groups, init, *args):\n def find_length(axis, x):\n if axis is not axes_scan.broadcast:\n leaves = jax.tree_util.tree_leaves(x)\n if leaves:\n return leaves[0].shape[axis]\n return ()\n\n # split rngs\n lengths = jax.tree_util.tree_map(find_length, in_axes, args)\n lengths = set(jax.tree_util.tree_leaves(lengths))\n if length is None and len(lengths) == 1:\n (d_length,) = lengths\n elif len(lengths) > 1:\n raise ValueError(f'Inconsistent scan lengths: {lengths}')\n elif length is None:\n raise ValueError('length should be specified manually.')\n else:\n d_length = length\n # random.clone is only available on Jax versions 0.4.26 or newer\n # see: https://jax.readthedocs.io/en/latest/jax.experimental.key_reuse.html\n if hasattr(random, 'clone'):\n split_fn = lambda rng: random.split(random.clone(rng), d_length)\n else:\n split_fn = lambda rng: random.split(rng, d_length)\n\n rng_groups = tuple(\n tree_map_rngs(split_fn, rng_group) if split else rng_group\n for rng_group, split in zip(rng_groups, rng_splits)\n )\n\n @functools.partial(\n axes_scan.scan,\n in_axes=(variable_in_axes, rng_axes, in_axes),\n out_axes=(out_axes, variable_out_axes),\n length=length,\n reverse=reverse,\n unroll=unroll,\n _split_transpose=_split_transpose,\n check_constancy_invariants=check_constancy_invariants,\n )\n def scanned(broadcast_vars, carry, scan_variable_groups, rng_groups, args):\n carry_vars, c = carry\n variable_groups = (broadcast_vars, carry_vars) + scan_variable_groups\n if data_transform is not None:\n variable_groups, rng_groups = data_transform(\n variable_groups, rng_groups\n )\n scope = scope_fn(variable_groups, rng_groups)\n c, y = fn(scope, c, *args)\n out_vars = repack_fn(scope)\n broadcast_vars_out = out_vars[0]\n carry_vars = out_vars[1]\n scan_vars = out_vars[2:]\n # add immutable broadcast vars back to broadcast output\n # otherwise they won't be fed to the actual scan body\n for in_group, out_group in zip(broadcast_vars, broadcast_vars_out):\n for col in in_group:\n if col not in out_group:\n out_group[col] = in_group[col]\n return broadcast_vars_out, (carry_vars, c), (y, scan_vars)\n\n broadcast_vars = variable_groups[0]\n carry_vars = variable_groups[1]\n scan_vars = variable_groups[2:]\n new_scan_vars = []\n for scan_group, axis in zip(scan_vars, variable_in_axes):\n new_scan_vars.append(meta.remove_axis(scan_group, axis, metadata_params))\n broadcast_vars, (carry_vars, c), (ys, scan_vars) = scanned(\n broadcast_vars,\n (carry_vars, init),\n tuple(new_scan_vars),\n rng_groups,\n args,\n )\n new_scan_vars = []\n for scan_group, axis in zip(scan_vars, variable_out_axes):\n new_scan_vars.append(meta.add_axis(scan_group, axis, metadata_params))\n scan_vars = tuple(new_scan_vars)\n out_vars = (\n broadcast_vars,\n carry_vars,\n ) + scan_vars\n return (c, ys), out_vars\n\n return pack(\n inner,\n (variable_broadcast, variable_carry) + variable_in_groups,\n (variable_broadcast, variable_carry) + variable_out_groups,\n rng_groups,\n name='scan',\n )\n\n\nC = TypeVar('C')\n\n\ndef while_loop(\n cond_fn: Callable[[Scope, C], bool],\n body_fn: Callable[[Scope, C], C],\n scope: Scope,\n init: C,\n carry_variables: CollectionFilter = False,\n broadcast_variables: CollectionFilter = True,\n split_rngs: Mapping[PRNGSequenceFilter, bool] = {},\n) -> C:\n \"\"\"Lifted version of jax.lax.while_loop.\n\n The lifted scope is passed to `cond_fn` and `body_fn`.\n Broadcasted variables are immutable. The carry variable are\n mutable but cannot change shape and dtype.\n This also means you cannot initialize variables inside\n the body. Consider calling `body_fn` once manually before\n calling `while_loop` if variable initialization is required.\n\n Example::\n\n def f(scope, x):\n def cond_fn(scope, c):\n return scope.get_variable('state', 'acc') < 10\n def body_fn(scope, c):\n acc = scope.variable('state', 'acc')\n acc += 1\n y = scope.child(nn.dense)(c, c.shape[-1])\n return y\n\n c = x\n c = body_fn(scope, c)\n return lift.while_loop(cond_fn, body_fn, scope, (),\n carry_variables='state')\n\n Args:\n cond_fn: Should return True as long as the loop should continue.\n body_fn: The body of the while loop.\n scope: The scope(s) which should be lifted into the loop.\n init: The initial state passed to the loop\n carry_variables: collections that are carried through the loop\n and are therefore mutable (default: none).\n broadcast_variables: collections that are closed over and are\n therefore read-only (default: all collections)\n split_rngs: Split PRNG sequences will be different for each loop iterations.\n If split is False the PRNGs will be the same across iterations.\n Returns:\n The final state after executing the while loop.\n \"\"\"\n rng_groups, rng_splits = _unzip2(split_rngs.items())\n\n def inner(scope_fn, repack_fn, variable_groups, rng_groups):\n carry_variables, broadcast_variables = variable_groups\n\n def make_loop_rngs(i):\n local_rng_groups = []\n for rng_group, rng_split in zip(rng_groups, rng_splits):\n if rng_split:\n rng_group = tree_map_rngs(\n lambda rng: random.fold_in(rng, i), rng_group\n )\n local_rng_groups.append(rng_group)\n return local_rng_groups\n\n def cond_wrapper(c):\n i, carry_variables, carry = c\n scope = scope_fn(\n (carry_variables, broadcast_variables),\n make_loop_rngs(-i),\n mutable_filter=False,\n )\n return cond_fn(scope, carry)\n\n def body_wrapper(c):\n i, carry_variables, carry = c\n scope = scope_fn(\n (carry_variables, broadcast_variables), make_loop_rngs(i)\n )\n carry = body_fn(scope, carry)\n (carry_variables,) = repack_fn(scope)\n return (i + 1, carry_variables, carry)\n\n c = (0, carry_variables, init)\n _, carry_variables, carry = jax.lax.while_loop(\n cond_wrapper, body_wrapper, c\n )\n return carry, (carry_variables,)\n\n return pack(\n inner,\n (carry_variables, broadcast_variables),\n (carry_variables,),\n rng_groups,\n name='while_loop',\n )(scope)\n\n\ndef cond(\n pred: Any,\n true_fun: Callable[..., C],\n false_fun: Callable[..., C],\n scope: Scope,\n *operands,\n variables: CollectionFilter = True,\n rngs: PRNGSequenceFilter = True,\n) -> C:\n \"\"\"Lifted version of ``jax.lax.cond``.\n\n The returned values from ``true_fun`` and ``false_fun``\n must have the same Pytree structure, shapes, and dtypes.\n The variables created or updated inside the\n branches must also have the same structure.\n Note that this constraint is violated when\n creating variables or submodules in only one branch.\n Because initializing variables in just one branch\n causes the parameter structure to be different.\n\n Example::\n\n def cond_example(scope, x, pred):\n scope.variable('state', 'true_count', lambda: 0)\n scope.variable('state', 'false_count', lambda: 0)\n def true_fn(scope, x):\n scope.variable('state', 'true_count').value += 1\n return scope.child(nn.dense)(x, 2)\n def false_fn(scope, x):\n scope.variable('state', 'false_count').value += 1\n return -scope.child(nn.dense)(x, 2)\n return lift.cond(pred, true_fn, false_fn, scope, x)\n\n\n Args:\n pred: determines if true_fun or false_fun is evaluated.\n true_fun: The function evalauted when ``pred`` is `True`.\n The signature is (Scope, *operands) -> T.\n false_fun: The function evalauted when ``pred`` is `False`.\n The signature is (Scope, *operands) -> T.\n scope: A Scope or Pytree of scopes to pass\n *operands: The arguments passed to ``true_fun`` and ``false_fun``\n variables: The variable collections passed to the conditional\n branches (default: all)\n rngs: The PRNG sequences passed to the conditionals (default: all)\n Returns:\n The result of the evaluated branch (``true_fun`` or ``false_fun``).\n \"\"\"\n branches = [true_fun, false_fun]\n\n def inner(scope_fn, repack_fn, variable_groups, rng_groups):\n def branch_wrapper(branch_fn, *operands):\n scope = scope_fn(variable_groups, rng_groups)\n y = branch_fn(scope, *operands)\n return y, repack_fn(scope)\n\n pure_branches = [\n functools.partial(branch_wrapper, branch_fn) for branch_fn in branches\n ]\n return jax.lax.cond(pred, pure_branches[0], pure_branches[1], *operands)\n\n return pack(inner, (variables,), (variables,), (rngs,), name='cond')(scope)\n\n\ndef switch(\n index: Any,\n branches: Sequence[Callable[..., C]],\n scope: Scope,\n *operands,\n variables: CollectionFilter = True,\n rngs: PRNGSequenceFilter = True,\n) -> C:\n \"\"\"Lifted version of ``jax.lax.switch``.\n\n The returned values from ``branches``\n must have the same Pytree structure, shapes, and dtypes.\n The variables created or updated inside the\n branches must also have the same structure.\n Note that this constraint is violated when\n creating variables or submodules in only one branch.\n Because initializing variables in just one branch\n causes the parameter structure to be different.\n\n Example::\n\n def switch_example(scope, x, index):\n scope.variable('state', 'a_count', lambda: 0)\n scope.variable('state', 'b_count', lambda: 0)\n scope.variable('state', 'c_count', lambda: 0)\n def a_fn(scope, x):\n scope.variable('state', 'a_count').value += 1\n return scope.child(nn.dense)(x, 2)\n def b_fn(scope, x):\n scope.variable('state', 'b_count').value += 1\n return -scope.child(nn.dense)(x, 2)\n def c_fn(scope, x):\n scope.variable('state', 'c_count').value += 1\n return scope.child(nn.dense)(x, 2)\n return lift.switch(index, [a_fn, b_fn, c_fn], scope, x)\n\n If you want to have a different parameter structure for each branch\n you should run all branch on initialization before calling switch::\n\n def multihead_switch_example(scope, x, index):\n def a_fn(scope, x):\n x = scope.child(nn.dense)(x, 10)\n x = scope.child(nn.dense)(x, 7)\n return scope.child(nn.dense)(x, 5)\n def b_fn(scope, x):\n x = scope.child(nn.dense)(x, 11)\n return scope.child(nn.dense)(x, 5)\n def c_fn(scope, x):\n return scope.child(nn.dense)(x, 5)\n\n branches = [a_fn, b_fn, c_fn]\n\n # run all branches on init\n if scope.is_mutable_collection('params'):\n for branch in branches:\n _ = branch(scope, x)\n\n return lift.switch(index, branches, scope, x)\n\n Args:\n index: Integer scalar type, indicating which branch function to apply.\n branches: Sequence of functions to be applied based on index.\n The signature of each function is (Scope, *operands) -> T.\n scope: A Scope or Pytree of scopes to pass\n *operands: The arguments passed to ``true_fun`` and ``false_fun``\n variables: The variable collections passed to the conditional\n branches (default: all)\n rngs: The PRNG sequences passed to the conditionals (default: all)\n Returns:\n The result of the evaluated branch.\n \"\"\"\n\n def inner(scope_fn, repack_fn, variable_groups, rng_groups):\n def branch_wrapper(branch_fn, *operands):\n scope = scope_fn(variable_groups, rng_groups)\n y = branch_fn(scope, *operands)\n return y, repack_fn(scope)\n\n pure_branches = [\n functools.partial(branch_wrapper, branch_fn) for branch_fn in branches\n ]\n return jax.lax.switch(index, pure_branches, *operands)\n\n return pack(inner, (variables,), (variables,), (rngs,), name='switch')(scope)\n\n\ndef custom_vjp(\n fn: Callable[..., Any],\n forward_fn: Callable[..., Any],\n backward_fn: Callable[..., Any],\n grad_vars: CollectionFilter = 'params',\n nondiff_argnums=(),\n):\n \"\"\"Lifted version of `jax.custom_vjp`.\n\n `forward_fn` and `backward_fn` together define a custom vjp for `fn`.\n The original `fn` will run in case a vjp (backward gradient) is not computed.\n\n The `forward_fn` receives the same arguments as `fn` but is expected to return\n a tuple containing the output of `fn(scope, *args)` and the residuals that are\n passed to `backward_fn`.\n\n The `backward_fn` receives the nondiff arguments, residuals, and the output\n tangents. It should return a tuple containing the variable and input tangents.\n\n Note that the vjp function returned by `lift.vjp` can be passed as residual\n and used in the `backward_fn`. The scope is unavailable during the backward\n pass. If the scope is required in `backward_fn`, a snapshot of the variables\n can be taken and returned as a residual in the `forward_fn`.\n\n Example::\n\n f = nn.dense\n\n def fwd(scope, x, features):\n y, vjp_fn = lift.vjp(partial(f, features=features), scope, x)\n return y, vjp_fn\n\n def bwd(features, vjp_fn, y_t):\n params_t, *inputs_t = vjp_fn(y_t)\n params_t = jax.tree_util.tree_map(jnp.sign, params_t)\n return (params_t, *inputs_t)\n\n dense_sign_grad = lift.custom_vjp(\n f, forward_fn=fwd, backward_fn=bwd, nondiff_argnums=(2,))\n\n Args:\n fn: The function to define a custom_vjp for. The first argument\n should be a ``Module`` instance.\n forward_fn: A function with the same arguments as `fn` returning an tuple\n with the original output and the residuals that will be passed to\n `backward_fn`.\n backward_fn: arguments are passed as (*nondiff_args, residuals, tangents)\n The function should return a tuple containing the tangents for the\n variable in the collections specified by `grad_vars` and the input\n arguments (except the scope and nondiff args).\n grad_vars: The collections for which a vjp will be computed\n (default: \"params\").\n nondiff_argnums: arguments for which no vjp is computed.\n Returns:\n A function with the same signature as `fn` with the custom vjp.\n \"\"\"\n\n def inner(scope_fn, repack_fn, variable_groups, rng_groups, *args):\n grad_variables, other_variables = variable_groups\n scopes_treedef = None\n\n def f(grad_variables, *args):\n scope = scope_fn((grad_variables, other_variables), rng_groups)\n y = fn(scope, *args)\n vars_out = repack_fn(scope)\n return y, vars_out\n\n f = jax.custom_vjp(f, nondiff_argnums=nondiff_argnums)\n\n def f_fwd(grad_variables, *args):\n nonlocal scopes_treedef\n scopes = scope_fn((grad_variables, other_variables), rng_groups)\n scopes_treedef = jax.tree_util.tree_structure(scopes)\n y, res = forward_fn(scopes, *args)\n vars_out = repack_fn(scopes)\n return (y, vars_out), res\n\n def f_bwd(*args):\n # the backward function does not pass a lifted scope to the user.\n # Currently, there is no way to have side effects flow out of backward\n # pass. Even without mutation variables would be ill-defined. For example,\n # would we take a snapshot of the variables before or after calling\n # `forward_fn`?\n nondiff_args = args[:-2]\n res, g = args[-2:] # pylint: disable=unbalanced-tuple-unpacking\n g_y, _ = g\n var_t, *inputs_t = backward_fn(*nondiff_args, res, g_y)\n assert scopes_treedef is not None, 'backward called before forward?!'\n var_t = tuple(scopes_treedef.flatten_up_to(var_t))\n return (var_t, *inputs_t)\n\n f.defvjp(f_fwd, f_bwd)\n\n return f(grad_variables, *args)\n\n variable_in_groups = (grad_vars, True)\n variable_out_groups = (grad_vars, True)\n rng_groups = (True,)\n return pack(\n inner,\n variable_in_groups,\n variable_out_groups,\n rng_groups,\n name='custom_vjp',\n )\n\n\ndef checkpoint(\n fn: Callable[..., Any],\n variables: CollectionFilter = True,\n rngs: PRNGSequenceFilter = True,\n concrete: bool = False,\n prevent_cse: bool = True,\n static_argnums: int | tuple[int, ...] = (),\n policy: Callable[..., bool] | None = None,\n) -> Callable[..., Any]:\n \"\"\"Lifted version of ``jax.checkpoint``.\n\n This function is aliased to ``lift.remat`` just like ``jax.remat``.\n\n Args:\n fn: scope function for which intermediate computations should be re-computed\n when computing gradients.\n variables: The variable collections that are lifted. By default all\n collections are lifted.\n rngs: The PRNG sequences that are lifted. By default all PRNG sequences are\n lifted.\n concrete: Optional, boolean indicating whether ``fun`` may involve\n value-dependent Python control flow (default ``False``). Support for such\n control flow is optional, and disabled by default, because in some\n edge-case compositions with :func:`jax.jit` it can lead to some extra\n computation.\n prevent_cse: Optional, boolean indicating whether to prevent common\n subexpression elimination (CSE) optimizations in the HLO generated from\n differentiation. This CSE prevention has costs because it can foil other\n optimizations, and because it can incur high overheads on some backends,\n especially GPU. The default is True because otherwise, under a ``jit`` or\n ``pmap``, CSE can defeat the purpose of this decorator. But in some\n settings, like when used inside a ``scan``, this CSE prevention mechanism\n is unnecessary, in which case ``prevent_cse`` can be set to False.\n static_argnums: Optional, int or sequence of ints, indicates which argument\n values on which to specialize for tracing and caching purposes. Specifying\n arguments as static can avoid ConcretizationTypeErrors when tracing, but\n at the cost of more retracing overheads.\n policy: Experimental checkpoint policy, see ``jax.checkpoint``.\n\n Returns:\n A wrapped version of ``fn``. When computing gradients intermediate\n computations will be re-computed when computing gradients.\n \"\"\"\n\n def inner(scope_fn, repack_fn, variable_groups, rng_groups, *args, **kwargs):\n # add 2 to each static_argnums because we add two initial arguments to rematted\n static_argnums_ = jax.tree_util.tree_map(lambda x: x + 2, static_argnums)\n\n # After JAX v0.3.16, concrete=False is a no-op and concrete=True raises\n # NotImplementedError. Starting in JAX v0.8.2, the concrete argument is\n # deprecated and will be removed in the future.\n if concrete:\n raise NotImplementedError(\n \"The concrete argument is deprecated. Use static_argnums instead.\"\n \" for more information, see\"\n \" https://docs.jax.dev/en/latest/jep/11830-new-remat-checkpoint.html\"\n )\n\n @functools.partial(\n jax.remat,\n static_argnums=static_argnums_,\n prevent_cse=prevent_cse,\n policy=policy,\n )\n @functools.wraps(fn)\n def rematted(variable_groups, rng_groups, *args, **kwargs):\n scope = scope_fn(variable_groups, rng_groups)\n y = fn(scope, *args, **kwargs)\n return y, repack_fn(scope)\n\n return rematted(variable_groups, rng_groups, *args, **kwargs)\n\n return pack(\n inner,\n (variables,),\n (variables,),\n (rngs,),\n name='remat',\n enable_kwargs=True,\n )\n\n\nremat = checkpoint\n\n\ndef _hashable_filter(x):\n \"\"\"Hashable version of CollectionFilter.\"\"\"\n if isinstance(x, str):\n return (x,)\n if isinstance(x, Iterable):\n return tuple(x) # convert un-hashable list & sets to tuple\n if isinstance(x, DenyList):\n return DenyList(\n _hashable_filter(x.deny)\n ) # convert inner filter recursively\n return x\n\n\nclass CountsHolder:\n\n def __init__(self, flat_d):\n self.flat_d = flat_d\n\n @classmethod\n def make(cls, d):\n flat_d = traverse_util.flatten_dict(d)\n flat_d = {k: v for k, v in flat_d.items()}\n return cls(flat_d)\n\n def sub(self, other):\n delta_flat_d = {}\n new_flat_d = collections.defaultdict(int, self.flat_d)\n old_flat_d = collections.defaultdict(int, other.flat_d)\n for k in new_flat_d:\n delta_flat_d[k] = new_flat_d[k] - old_flat_d[k]\n return CountsHolder(delta_flat_d)\n\n def add(self, other):\n delta_flat_d = {}\n new_flat_d = collections.defaultdict(int, self.flat_d)\n old_flat_d = collections.defaultdict(int, other.flat_d)\n for k in new_flat_d:\n delta_flat_d[k] = new_flat_d[k] + old_flat_d[k]\n return CountsHolder(delta_flat_d)\n\n def unflat(self):\n return traverse_util.unflatten_dict(self.flat_d)\n\n\ndef set_from_dict(original, updates):\n for k in updates:\n if k not in original:\n original[k] = updates[k]\n else:\n if isinstance(updates[k], dict):\n set_from_dict(original[k], updates[k])\n else:\n original[k] = updates[k]\n\n\nclass _SideEffectCache(threading.local):\n\n def __init__(self):\n self.cache = {}\n\n\n_side_effect_cache = _SideEffectCache()\n\n\ndef _restore_rng_counters(scopes, fingerprint, capture_old_counts):\n if fingerprint not in _side_effect_cache.cache:\n capture_new_counts = jax.tree.map(\n lambda s: CountsHolder.make(s.rng_counters), scopes\n )\n capture_delta_counts = jax.tree.map(\n lambda old, new: new.sub(old),\n capture_old_counts,\n capture_new_counts,\n )\n _side_effect_cache.cache[fingerprint] = capture_delta_counts\n else:\n updated_counts = jax.tree.map(\n lambda x, y: x.add(y).unflat(),\n _side_effect_cache.cache[fingerprint],\n capture_old_counts,\n )\n jax.tree.map(\n lambda s, u: set_from_dict(s.rng_counters, u),\n scopes,\n updated_counts,\n )\n\n\ndef jit(\n fn: Callable[..., Any],\n variables: CollectionFilter = True,\n rngs: PRNGSequenceFilter = True,\n static_argnums: int | Iterable[int] = (),\n static_argnames: str | Iterable[str] = (),\n donate_argnums: int | Iterable[int] = (),\n device=None,\n backend: str | None = None,\n) -> Callable[..., Any]:\n \"\"\"Lifted version of ``jax.jit``.\n\n Args:\n fn: Scope function to be jitted.\n variables: The variable collections that are lifted. By default all\n collections are lifted.\n rngs: The PRNG sequences that are lifted. By default all PRNG sequences\n are lifted.\n static_argnums: An int or collection of ints specifying which positional\n arguments to treat as static (compile-time constant). Operations that only\n depend on static arguments will be constant-folded in Python (during\n tracing), and so the corresponding argument values can be any Python\n object. Static arguments should be hashable, meaning both ``__hash__`` and\n ``__eq__`` are implemented, and immutable. Calling the jitted function\n with different values for these constants will trigger recompilation. If\n the jitted function is called with fewer positional arguments than\n indicated by ``static_argnums`` then an error is raised. Arguments that\n are not arrays or containers thereof must be marked as static.\n Defaults to ().\n static_argnames: An optional string or collection of strings specifying\n which named arguments to treat as static (compile-time constant). See the\n comment on ``static_argnums`` for details. If not\n provided but ``static_argnums`` is set, the default is based on calling\n ``inspect.signature(fun)`` to find corresponding named arguments.\n donate_argnums: Specify which arguments are \"donated\" to the computation.\n It is safe to donate arguments if you no longer need them once the\n computation has finished. In some cases XLA can make use of donated\n buffers to reduce the amount of memory needed to perform a computation,\n for example recycling one of your input buffers to store a result. You\n should not reuse buffers that you donate to a computation, JAX will raise\n an error if you try to.\n device: This is an experimental feature and the API is likely to change.\n Optional, the Device the jitted function will run on. (Available devices\n can be retrieved via :py:func:`jax.devices`.) The default is inherited\n from XLA's DeviceAssignment logic and is usually to use\n ``jax.devices()[0]``.\n backend: a string representing the XLA backend: ``'cpu'``, ``'gpu'``, or\n ``'tpu'``.\n\n Returns:\n A wrapped version of ``fn``, set up for just-in-time compilation.\n \"\"\"\n if not isinstance(static_argnums, Iterable):\n static_argnums = (static_argnums,)\n if not isinstance(donate_argnums, Iterable):\n donate_argnums = (donate_argnums,)\n # offset argnums by two because first argument in the original function is the\n # scope while jitted has 3 functions before the user arguments.\n static_argnums = (0,) + tuple(i + 2 for i in static_argnums if i > 0)\n donate_argnums = tuple(i + 2 for i in donate_argnums if i > 0)\n\n # Close over scope_fn & repack_fn to avoid recompilation\n # this is impure but we use the fingerprint arg to differentiate between cases\n # where scope_fn or repack_fn actually produce non-identical results.\n jit_context = TransformContext[tuple[Callable, Callable]]()\n\n @functools.partial(\n jax.jit,\n static_argnums=static_argnums,\n static_argnames=static_argnames,\n donate_argnums=donate_argnums,\n device=device,\n backend=backend,\n )\n @functools.wraps(fn)\n def jitted(fingerprint, variable_groups, rng_groups, *args, **kwargs):\n scope_fn, repack_fn = jit_context.get()\n hash_key = fingerprint[1]\n # fingerprint is only used to differentiate the cache signature\n # del fingerprint\n scope = scope_fn(variable_groups, rng_groups) # pylint: disable=not-callable\n y = fn(scope, hash_key, *args, **kwargs)\n return y, repack_fn(scope) # pylint: disable=not-callable\n\n def inner(\n scope_fn,\n repack_fn,\n variable_groups,\n rng_groups,\n module_hash_key,\n *args,\n **kwargs,\n ):\n with jit_context.push((scope_fn, repack_fn)):\n scopes: list[Scope] = jax.tree_util.tree_leaves(\n scope_fn(variable_groups, rng_groups)\n )\n mutable = tuple(_hashable_filter(scope.mutable) for scope in scopes)\n\n rng_groups = jax.tree.map(\n lambda x: x.clear_suffix() if isinstance(x, LazyRng) else x,\n rng_groups,\n is_leaf=lambda x: isinstance(x, LazyRng),\n )\n\n fingerprint = (mutable, module_hash_key)\n capture_old_counts = jax.tree.map(\n lambda s: CountsHolder.make(s.rng_counters), scopes\n )\n res = jitted(fingerprint, variable_groups, rng_groups, *args, **kwargs)\n _restore_rng_counters(scopes, fingerprint, capture_old_counts)\n return res\n\n return pack(\n inner, (variables,), (variables,), (rngs,), name='jit', enable_kwargs=True\n )\n\n\ndef remat_scan(\n body_fn: Callable[..., Any],\n lengths: Sequence[int],\n policy: Callable[..., bool] | None = None,\n variable_broadcast: CollectionFilter = False,\n variable_carry: CollectionFilter = False,\n variable_axes: Mapping[CollectionFilter, InOutScanAxis] = {True: 0},\n split_rngs: Mapping[PRNGSequenceFilter, bool] = {True: True},\n) -> Callable[..., Any]:\n \"\"\"Combines `lift.remat` and `lift.scan` for memory efficiency and constant time compilation.\n\n ``remat_scan`` allows for constant compile times and sublinear\n memory usage with respect to model depth. At a small constant\n penalty. This is typically beneficial for very deep models.\n\n Example::\n\n def body_fn(scope, x):\n return nn.dense(scope, x, features=x.shape[-1])\n # 100x dense with O(sqrt(N)) memory for gradient computation\n y = lift.remat_scan(body_fn, lengths=(10, 10))(scope, x)\n\n Args:\n body_fn: Scope function to be repeated using a (nested scan)\n lengths: number of loop iterations at the given level. The total number of\n iterations `n = prod(lengths)`. each loop is rematerialized. This way the\n memory consumption is proportional to `n^(1 / d)` where `d =\n len(lengths)`. Minimal memory consumptions requires tuning the lengths\n such that the same amount of memory is consumed at each level of the\n nested loop.\n policy: Experimental checkpoint policy, see ``jax.checkpoint``.\n variable_broadcast: Specifies the broadcasted variable collections. A\n broadcasted variable should not depend on any computation that cannot be\n lifted out of the loop. This is typically used to define shared parameters\n inside the fn.\n variable_carry: Specifies the variable collections that are carried through\n the loop. Mutations to these variables are carried to the next iteration\n and will be preserved when the scan finishes.\n variable_axes: the variable collections that are scanned over.\n split_rngs: Split PRNG sequences will be different for each loop iterations.\n If split is False the PRNGs will be the same across iterations.\n Returns:\n A wrapped version of ``body_fn`` that repeats itself prod(lengths) times.\n \"\"\"\n # TODO(jheek) should remat scan have scan inputs/outputs?\n scan_fn = functools.partial(\n scan,\n variable_broadcast=variable_broadcast,\n variable_carry=variable_carry,\n variable_axes=variable_axes,\n split_rngs=split_rngs,\n )\n if len(lengths) == 1:\n\n def wrapper(scope, carry):\n return body_fn(scope, carry), ()\n\n fn = lambda scope, c: scan_fn(wrapper, length=lengths[0])(scope, c)[0]\n else:\n\n @functools.partial(remat, policy=policy, prevent_cse=False)\n def inner_loop(scope, carry):\n carry = remat_scan(\n body_fn,\n lengths[1:],\n policy,\n variable_broadcast,\n variable_carry,\n variable_axes,\n split_rngs,\n )(scope, carry)\n return carry, ()\n\n fn = lambda scope, c: scan_fn(inner_loop, length=lengths[0])(scope, c)[0]\n return fn\n\n\ndef _unzip2(xs):\n ys = tuple(zip(*xs))\n return ys if ys else ((), ())\n\n\ndef _broadcast_prefix_tree(prefix_tree: Any, full_tree: Any) -> list[Any]:\n bcast_flat = []\n num_leaves_fn = lambda t: jax.tree.flatten(t)[1].num_leaves\n jax.tree.map(\n lambda x, subtree: bcast_flat.extend([x] * num_leaves_fn(subtree)),\n prefix_tree,\n full_tree,\n is_leaf=lambda x: x is None,\n )\n return bcast_flat\n\n\ndef fold_rngs(\n fn: Callable[..., Any],\n variables: CollectionFilter = True,\n rngs: PRNGSequenceFilter = True,\n) -> Callable[..., Any]:\n # Close over scope_fn & repack_fn to avoid recompilation\n # this is impure but we use the fingerprint arg to differentiate between cases\n # where scope_fn or repack_fn actually produce non-identical results.\n fold_rngs_context = TransformContext[tuple[Callable, Callable]]()\n\n @functools.wraps(fn)\n def wrapped_fold_rngs(fingerprint, variable_groups, rng_groups, *args, **kwargs):\n scope_fn, repack_fn = fold_rngs_context.get()\n hash_key = fingerprint[1]\n # fingerprint is only used to differentiate the cache signature\n # del fingerprint\n scope = scope_fn(variable_groups, rng_groups) # pylint: disable=not-callable\n y = fn(scope, hash_key, *args, **kwargs)\n return y, repack_fn(scope) # pylint: disable=not-callable\n\n def inner_fold_rngs(\n scope_fn,\n repack_fn,\n variable_groups,\n rng_groups,\n module_hash_key,\n *args,\n **kwargs,\n ):\n with fold_rngs_context.push((scope_fn, repack_fn)):\n scopes: list[Scope] = jax.tree_util.tree_leaves(\n scope_fn(variable_groups, rng_groups)\n )\n mutable = tuple(_hashable_filter(scope.mutable) for scope in scopes)\n\n rng_groups = jax.tree.map(\n lambda x: x.clear_suffix() if isinstance(x, LazyRng) else x,\n rng_groups,\n is_leaf=lambda x: isinstance(x, LazyRng),\n )\n\n fingerprint = (mutable, module_hash_key)\n capture_old_counts = jax.tree.map(\n lambda s: CountsHolder.make(s.rng_counters), scopes\n )\n res = wrapped_fold_rngs(\n fingerprint, variable_groups, rng_groups, *args, **kwargs\n )\n _restore_rng_counters(scopes, fingerprint, capture_old_counts)\n return res\n\n return pack(\n inner_fold_rngs,\n (variables,),\n (variables,),\n (rngs,),\n name='fold_rngs',\n enable_kwargs=True,\n )\n" }, { "path": "flax/core/meta.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Boxed Metadata API\n\nBoxed metadata enables tracking arbitrary metadata for linen variables\nthat is compatible with lifted transformations.\n\nSee ``Partitioned`` for a practical example on how to use this metadata\nto keep track of how variables should be partitioned with ``jax.pjit``.\n\"\"\"\n\nimport abc\nimport dataclasses\nimport functools\nfrom typing import Any, Generic, TypeVar\nfrom collections.abc import Callable\n\nfrom flax import errors, struct\nfrom flax.typing import LogicalNames\nimport jax\n\nA = TypeVar('A')\nB = TypeVar('B')\nTAxisMetadata = TypeVar('TAxisMetadata', bound='AxisMetadata[Any]')\n\n\nclass AxisMetadata(Generic[A], metaclass=abc.ABCMeta):\n \"\"\"Abstract base class for boxed Metadata.\n\n ``AxisMetadata`` enables arbitrary, per axis metadata for variables.\n By using ``unbox`` the metadata is stripped away to obtain the original\n variables. By using unboxing, most code handling variables does not need\n to handle ``AxisMetadata`` specifically, but can directly operate on the JAX\n arrays that they wrap.\n\n Additionally, ``AxisMetadata`` supports updating metadata whenever an axis\n is added or removed by a functional transformation\n (e.g.: ``nn.scan`` or ``nn.vmap``) using the ``add_axis`` and ``remove_axis``\n methods.\n\n By extending ``AxisMetadata``, custom metadata can be stored. See\n ``Partitioned`` for a specific implementation.\n \"\"\"\n\n @abc.abstractmethod\n def unbox(self) -> A:\n \"\"\"Returns the content of the AxisMetadata box.\n\n Note that unlike ``meta.unbox`` the unbox call should not recursively unbox\n metadata. It should simply return value that it wraps directly even\n if that value itself is an instance of AxisMetadata.\n\n In practise, AxisMetadata subclasses should be registered as PyTree nodes to\n support passing instances to JAX and Flax APIs. The leaves returned for this\n node should correspond to the value returned by unbox.\n\n Returns:\n The unboxed value.\n \"\"\"\n pass\n\n @abc.abstractmethod\n def replace_boxed(self, val: B) -> 'AxisMetadata[B]':\n \"\"\"Replaces the boxed value with the provided value.\n\n Args:\n val: The new value to be boxed by this AxisMetadata wrapper\n\n Returns:\n A new instance of the same type as self with `val` as the new ``unbox``\n content\n \"\"\"\n pass\n\n @abc.abstractmethod\n def add_axis(\n self: TAxisMetadata, index: int, params: dict[Any, Any]\n ) -> TAxisMetadata:\n \"\"\"Adds a new axis to the axis metadata.\n\n Note that add_axis and remove_axis should act as each other's inverse\n (meaning: ``x.add_axis(i, p).remove_axis(i, p) == x``)\n\n Args:\n index: The position at which the new axis will be inserted\n params: An arbitrary dictionary of parameters passed by the transformation\n that introduces the new axis (e.g.: ``nn.scan`` or ``nn.vmap``). The\n user passes this dictionary as the `metadata_param` argument to the\n transformation.\n\n Returns:\n A new instance of the same type as self and with the same ``unbox``\n content with updated axis metadata.\n \"\"\"\n pass\n\n @abc.abstractmethod\n def remove_axis(\n self: TAxisMetadata, index: int, params: dict[Any, Any]\n ) -> TAxisMetadata:\n \"\"\"Removes an axis from the axis metadata.\n\n Note that add_axis and remove_axis should act as each other's inverse\n (meaning: ``x.remove_axis(i, p).add_axis(i, p) == x``)\n\n Args:\n index: The position of the axis that is to be removed\n params: An arbitrary dictionary of parameters passed by the transformation\n that introduced the axis (e.g.: ``nn.scan`` or ``nn.vmap``). The user\n passes this dictionary as the `metadata_param` argument to the\n transformation.\n\n Returns:\n A new instance of the same type as self and with the same ``unbox``\n content with updated axis metadata.\n \"\"\"\n pass\n\n\ndef is_axis_metadata(val: Any) -> bool:\n \"\"\"Returns whether the argument is an instance of AxisMetadata.\"\"\"\n return isinstance(val, AxisMetadata)\n\n\ndef map_axis_meta(fn: Callable[[AxisMetadata[Any]], Any], tree: Any) -> Any:\n \"\"\"Maps over all PyTree nodes that are AxisMetadata instances.\"\"\"\n\n def wrapper(x):\n if isinstance(x, AxisMetadata):\n return fn(x)\n else:\n return x\n\n return jax.tree_util.tree_map(wrapper, tree, is_leaf=is_axis_metadata)\n\n\ndef add_axis(tree: Any, index: int, params: dict[Any, Any]) -> Any:\n \"\"\"Add an axis to each AxisMetadata node in a PyTree.\"\"\"\n return map_axis_meta(lambda x: x.add_axis(index, params), tree)\n\n\ndef remove_axis(tree: Any, index: int, params: dict[Any, Any]) -> Any:\n \"\"\"Remove an axis from each AxisMetadata node in a PyTree.\"\"\"\n return map_axis_meta(lambda x: x.remove_axis(index, params), tree)\n\n\ndef unbox(tree: Any) -> Any:\n \"\"\"Strips all AxisMetadata boxes from a PyTree.\"\"\"\n return map_axis_meta(lambda x: unbox(x.unbox()), tree)\n\n\ndef replace_boxed(tree: Any, updates: Any) -> Any:\n \"\"\"Updates all AxisMetadata boxes with the values in updates.\"\"\"\n\n def inner_update(c, v):\n if isinstance(c, AxisMetadata):\n return c.replace_boxed(replace_boxed(c.unbox(), v))\n else:\n return v\n\n return jax.tree_util.tree_map(\n inner_update, tree, updates, is_leaf=is_axis_metadata\n )\n\n\nPARTITION_NAME = 'partition_name'\n\n\ndef get_global_mesh() -> jax.sharding.AbstractMesh | jax.sharding.Mesh | None:\n mesh = jax.sharding.get_abstract_mesh()\n if mesh.empty:\n return None\n return mesh\n\n\ndef global_mesh_defined() -> bool:\n \"\"\"Checks if global mesh resource environment is defined.\"\"\"\n mesh = get_global_mesh()\n return mesh is not None\n\n\nclass Partitioned(struct.PyTreeNode, AxisMetadata[A]):\n \"\"\"Wrapper for partitioning metadata.\n\n ``Partitioned`` is used to extend variables with partitioning information\n required for ``jax.experimental.pjit``.\n\n The easiest way to define Partitioned variables is by using the\n ``with_partitioning`` wrapper around the variable initializer.\n\n Example::\n\n class MLP(nn.Module):\n hidden_size: int\n @nn.compact\n def __call__(self, x):\n ki = nn.linear.default_kernel_init\n h = nn.Dense(\n self.hidden_size,\n kernel_init=nn.with_partitioning(ki, ('data', 'model')))(x)\n h = nn.relu(h)\n return nn.Dense(\n x.shape[-1],\n kernel_init=nn.with_partitioning(ki, ('model', 'data')))(h)\n mlp = MLP(4096)\n x = jnp.ones((8 * 1024, 1024))\n # use eval_shape to get the Partitioned instances for the variables.\n # this way we can determine the PartitionSpecs for the init variables\n # before we call the init fn.\n var_spec = nn.get_partition_spec(\n jax.eval_shape(mlp.init, random.key(0), x))\n init_fn = mesh(pjit(mlp.init,\n (None, PartitionSpec(\"data\", \"model\")), var_spec))\n variables = init_fn(random.key(0), x)\n apply_fn = mesh(pjit(\n mlp.apply,\n (var_spec, PartitionSpec(\"data\", \"model\")),\n PartitionSpec(\"data\", \"model\")))\n apply_fn(variables, x)\n\n\n ``Partitioned`` values can gain additional axes when using transformations\n like ``nn.vmap`` and ``nn.scan``. In this case you can specify the name of\n the new axis with the `metadata_params` args in vmap/scan::\n\n class Model(nn.Module):\n @nn.compact\n def __call__(self, x):\n def body(mdl, c):\n c = MLP(4096)(c)\n return c, ()\n c, _ = nn.scan(\n body, variable_axes={\"params\": 0}, split_rngs={\"params\": 0}, length=8,\n metadata_params={nn.meta.PARTITION_NAME: \"layers\"})(self, x)\n return c\n \"\"\"\n\n value: Any\n names: LogicalNames = struct.field(pytree_node=False)\n mesh: jax.sharding.Mesh | None = struct.field(\n default=None, pytree_node=False\n )\n\n def unbox(self, apply_constraint=True) -> A:\n \"\"\"Returns the wrapped value with the partitioning applied as a sharding constraint.\"\"\"\n if apply_constraint and (global_mesh_defined() or self.mesh is not None):\n axis_resource = self.get_partition_spec()\n if self.mesh is not None:\n sharding = jax.sharding.NamedSharding(self.mesh, axis_resource)\n return jax.lax.with_sharding_constraint(self.value, sharding)\n return jax.lax.with_sharding_constraint(self.value, axis_resource)\n else:\n return self.value\n\n def replace_boxed(self, val: B) -> 'Partitioned[B]':\n return self.replace(value=val) # type: ignore\n\n def _get_partition_name(self, params: dict[Any, Any]) -> str:\n if PARTITION_NAME not in params:\n raise errors.PartitioningUnspecifiedError(self)\n return params[PARTITION_NAME]\n\n def add_axis(self, index: int, params: dict[Any, Any]) -> 'Partitioned[A]':\n axis_name = self._get_partition_name(params)\n names = list(self.names)\n while len(names) < index:\n names.append(None) # type: ignore\n names.insert(index, axis_name) # type: ignore\n return self.replace(names=tuple(names))\n\n def remove_axis(self, index: int, params: dict[Any, Any]) -> 'Partitioned[A]':\n axis_name = self._get_partition_name(params)\n names = list(self.names)\n assert names.pop(index) == axis_name\n return self.replace(names=tuple(names))\n\n def get_partition_spec(self) -> jax.sharding.PartitionSpec:\n \"\"\"Returns the ``Partitionspec`` for this partitioned value.\"\"\"\n return jax.sharding.PartitionSpec(*self.names)\n\n def get_sharding(self, mesh: jax.sharding.Mesh) -> jax.sharding.Sharding:\n \"\"\"Returns the ``NamedSharding`` for this partitioned value.\"\"\"\n return jax.sharding.NamedSharding(mesh, self.get_partition_spec())\n\n def to_nnx_metadata(self) -> dict[str, Any]:\n \"\"\"Return a dict of metadata that can translate into an `nnx.Variable`.\"\"\"\n metadata = dict(vars(self))\n metadata['out_sharding'] = metadata.pop('names')\n return metadata\n\n @classmethod\n def from_nnx_metadata(cls, metadata: dict[str, Any]):\n \"\"\"Given a dict of `nnx.Variable` format metadata, create a `nn.Partitioned`.\"\"\"\n metadata['names'] = metadata.pop('out_sharding')\n fields = {x.name for x in dataclasses.fields(cls)}\n return cls(**{k: v for k, v in metadata.items() if k in fields})\n\n\ndef with_partitioning(\n fn: Callable[..., Any],\n names: LogicalNames,\n mesh: jax.sharding.Mesh | None = None,\n) -> Callable[..., Partitioned[Any]]:\n \"\"\"Wraps a function's return value with Partitioned.\n\n Example::\n\n >>> import flax.linen as nn\n >>> kernel_init = nn.with_partitioning(\n ... nn.initializers.lecun_normal(), (None, \"data\"))\n >>> partitioned_dense = nn.Dense(features=3, kernel_init=kernel_init)\n\n Args:\n fn: The function to be wrapped. Typically this is an initializer.\n names: The logical axis passed to ``Partitioned``.\n mesh: The mesh to use for the partitioning. If None, the global mesh\n resource is used if available.\n\n Returns:\n A function wrapping ``fn`` that will return an instance of ``Partitioned``.\n \"\"\"\n\n @functools.wraps(fn)\n def wrapper(*args, **kwargs):\n return Partitioned(fn(*args, **kwargs), names, mesh=mesh)\n\n return wrapper\n\n\ndef _get_leaf_pspec(x: Any) -> jax.sharding.PartitionSpec | None:\n if hasattr(x, 'get_partition_spec'):\n return x.get_partition_spec()\n # Unboxed arrays, which should be replicated across all devices\n elif hasattr(x, 'shape'):\n return jax.sharding.PartitionSpec()\n else:\n return None\n\n\ndef get_partition_spec(tree: Any) -> Any:\n \"\"\"Extracts a PartitionSpec tree from a PyTree containing ``Partitioned`` values.\"\"\"\n return jax.tree_util.tree_map(\n _get_leaf_pspec, tree, is_leaf=lambda x: isinstance(x, AxisMetadata)\n )\n\n\ndef get_sharding(tree: Any, mesh: jax.sharding.Mesh) -> Any:\n \"\"\"Extracts a jax.sharding tree from a PyTree containing ``Partitioned`` values and a mesh.\"\"\"\n def f(x: Any) -> jax.sharding.Sharding | None:\n if hasattr(x, 'get_sharding'):\n return x.get_sharding(mesh)\n pspec = _get_leaf_pspec(x)\n if pspec is None:\n return None\n return jax.sharding.NamedSharding(mesh, pspec)\n\n return jax.tree_util.tree_map(\n f, tree, is_leaf=lambda x: isinstance(x, AxisMetadata)\n )\n" }, { "path": "flax/core/nn/__init__.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Flax Neural Network api.\"\"\"\n\n# pylint: disable=g-multiple-import\n# re-export commonly used modules and functions\nfrom flax.linen import activation as activation\nfrom flax.linen import initializers as initializers\nfrom flax.linen.activation import (\n celu as celu,\n elu as elu,\n gelu as gelu,\n glu as glu,\n leaky_relu as leaky_relu,\n log_sigmoid as log_sigmoid,\n log_softmax as log_softmax,\n relu as relu,\n sigmoid as sigmoid,\n silu as silu,\n soft_sign as soft_sign,\n softmax as softmax,\n softplus as softplus,\n swish as swish,\n tanh as tanh,\n)\nfrom flax.linen.pooling import (avg_pool as avg_pool, max_pool as max_pool)\nfrom .attention import (\n dot_product_attention as dot_product_attention,\n multi_head_dot_product_attention as multi_head_dot_product_attention,\n)\nfrom .linear import (\n Embedding as Embedding,\n conv_transpose as conv_transpose,\n conv as conv,\n dense_general as dense_general,\n dense as dense,\n embedding as embedding,\n)\nfrom .normalization import (\n batch_norm as batch_norm,\n group_norm as group_norm,\n layer_norm as layer_norm,\n)\nfrom .stochastic import dropout as dropout\n\n# pylint: enable=g-multiple-import\n" }, { "path": "flax/core/nn/attention.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Attention core modules for Flax.\"\"\"\n\nimport functools\nfrom collections.abc import Iterable # pylint: disable=g-importing-member\nfrom typing import Any\nfrom collections.abc import Callable\n\nimport jax\nimport jax.numpy as jnp\nimport numpy as np\nfrom jax import lax, random\n\nfrom flax import struct\nfrom flax.core import Scope\nfrom flax.linen import initializers\n\nfrom .linear import default_kernel_init, dense_general\n\n\ndef dot_product_attention(\n scope,\n query,\n key,\n value,\n dtype=jnp.float32,\n bias=None,\n axis=None,\n broadcast_dropout=True,\n dropout_rng=None,\n dropout_rate=0.0,\n deterministic=False,\n precision=None,\n):\n \"\"\"Computes dot-product attention given query, key, and value.\n\n This is the core function for applying attention based on\n https://arxiv.org/abs/1706.03762. It calculates the attention weights given\n query and key and combines the values using the attention weights. This\n function supports multi-dimensional inputs.\n\n\n Args:\n query: queries for calculating attention with shape of `[batch_size, dim1,\n dim2, ..., dimN, num_heads, mem_channels]`.\n key: keys for calculating attention with shape of `[batch_size, dim1, dim2,\n ..., dimN, num_heads, mem_channels]`.\n value: values to be used in attention with shape of `[batch_size, dim1,\n dim2,..., dimN, num_heads, value_channels]`.\n dtype: the dtype of the computation (default: float32)\n bias: bias for the attention weights. This can be used for incorporating\n autoregressive mask, padding mask, proximity bias.\n axis: axises over which the attention is applied.\n broadcast_dropout: bool: use a broadcasted dropout along batch dims.\n dropout_rng: JAX PRNGKey: to be used for dropout\n dropout_rate: dropout rate\n deterministic: bool, deterministic or not (to apply dropout)\n precision: numerical precision of the computation see `jax.lax.Precision`\n for details.\n\n Returns:\n Output of shape `[bs, dim1, dim2, ..., dimN,, num_heads, value_channels]`.\n \"\"\"\n assert key.shape[:-1] == value.shape[:-1]\n assert query.shape[0:1] == key.shape[0:1] and query.shape[-1] == key.shape[-1]\n\n if axis is None:\n axis = tuple(range(1, key.ndim - 2))\n if not isinstance(axis, Iterable):\n axis = (axis,)\n assert key.ndim == query.ndim\n assert key.ndim == value.ndim\n for ax in axis:\n if not (query.ndim >= 3 and 1 <= ax < query.ndim - 2):\n raise ValueError(\n 'Attention axis must be between the batch axis and the last-two axes.'\n )\n depth = query.shape[-1]\n n = key.ndim\n # batch_dims is , num_heads>\n batch_dims = tuple(np.delete(range(n), axis + (n - 1,)))\n # q & k -> (bs, , num_heads, , channels)\n qk_perm = batch_dims + axis + (n - 1,)\n key = key.transpose(qk_perm)\n query = query.transpose(qk_perm)\n # v -> (bs, , num_heads, channels, )\n v_perm = batch_dims + (n - 1,) + axis\n value = value.transpose(v_perm)\n\n query = query / jnp.sqrt(depth).astype(dtype)\n batch_dims_t = tuple(range(len(batch_dims)))\n attn_weights = lax.dot_general(\n query,\n key,\n (((n - 1,), (n - 1,)), (batch_dims_t, batch_dims_t)),\n precision=precision,\n )\n\n # apply attention bias: masking, droput, proximity bias, ect.\n if bias is not None:\n attn_weights = attn_weights + bias\n\n # normalize the attention weights\n norm_dims = tuple(range(attn_weights.ndim - len(axis), attn_weights.ndim))\n attn_weights = lax.exp(\n attn_weights\n - jax.scipy.special.logsumexp(attn_weights, axis=norm_dims, keepdims=True)\n )\n attn_weights = attn_weights.astype(dtype)\n\n # apply dropout\n if not deterministic and dropout_rate > 0.0:\n if dropout_rng is None:\n dropout_rng = scope.make_rng('dropout')\n keep_prob = 1.0 - dropout_rate\n if broadcast_dropout:\n # dropout is broadcast across the batch+head+non-attention dimension\n dropout_dims = attn_weights.shape[-(2 * len(axis)) :]\n dropout_shape = tuple([1] * len(batch_dims_t)) + dropout_dims\n keep = random.bernoulli(dropout_rng, keep_prob, dropout_shape)\n else:\n keep = random.bernoulli(dropout_rng, keep_prob, attn_weights.shape)\n multiplier = keep.astype(attn_weights.dtype) / jnp.asarray(\n keep_prob, dtype=dtype\n )\n attn_weights = attn_weights * multiplier\n\n # compute the new values given the attention weights\n wv_contracting_dims = (norm_dims, range(value.ndim - len(axis), value.ndim))\n y = lax.dot_general(\n attn_weights,\n value,\n (wv_contracting_dims, (batch_dims_t, batch_dims_t)),\n precision=precision,\n )\n\n # back to (bs, dim1, dim2, ..., dimN, num_heads, channels)\n perm_inv = _invert_perm(qk_perm)\n y = y.transpose(perm_inv)\n return y\n\n\ndef _invert_perm(perm):\n perm_inv = [0] * len(perm)\n for i, j in enumerate(perm):\n perm_inv[j] = i\n return tuple(perm_inv)\n\n\nclass CacheEntry(struct.PyTreeNode):\n key: np.ndarray\n value: np.ndarray\n i: np.ndarray\n\n\ndef multi_head_dot_product_attention(\n scope: Scope,\n inputs_q,\n inputs_kv,\n num_heads,\n dtype=jnp.float32,\n qkv_features=None,\n out_features=None,\n attention_axis=None,\n causal_mask=False,\n padding_mask=None,\n key_padding_mask=None,\n segmentation=None,\n key_segmentation=None,\n cache=False,\n broadcast_dropout=True,\n dropout_rng=None,\n dropout_rate=0.0,\n deterministic=False,\n precision=None,\n kernel_init=default_kernel_init,\n bias_init=initializers.zeros_init(),\n bias=True,\n attention_fn=dot_product_attention,\n):\n \"\"\"Applies multi-head dot product attention on the input data.\n\n Projects the inputs into multi-headed query, key, and value vectors,\n applies dot-product attention and project the results to an output vector.\n\n This can be used for encoder-decoder attention by specifying both `inputs_q`\n and `inputs_kv` orfor self-attention by only specifying `inputs_q` and\n setting `inputs_kv` to None.\n\n Args:\n inputs_q: input queries of shape `[bs, dim1, dim2, ..., dimN, features]`.\n inputs_kv: key/values of shape `[bs, dim1, dim2, ..., dimN, features]`\n or None for self-attention, inn which case key/values will be derived\n from inputs_q.\n num_heads: number of attention heads. Features (i.e. inputs_q.shape[-1])\n should be divisible by the number of heads.\n dtype: the dtype of the computation (default: float32)\n qkv_features: dimension of the key, query, and value.\n out_features: dimension of the last projection\n attention_axis: axes over which the attention is applied ( 'None' means\n attention over all axes, but batch, heads, and features).\n causal_mask: boolean specifying whether to apply a causal mask on the\n attention weights. If True, the output at timestep `t` will not depend\n on inputs at timesteps strictly greater than `t`.\n padding_mask: boolean specifying query tokens that are pad token.\n key_padding_mask: boolean specifying key-value tokens that are pad token.\n segmentation: segment indices for packed inputs_q data.\n key_segmentation: segment indices for packed inputs_kv data.\n cache: an instance of `flax.deprecated.nn.attention.Cache` used for\n efficient autoregressive decoding.\n broadcast_dropout: bool: use a broadcasted dropout along batch dims.\n dropout_rng: JAX PRNGKey: to be used for dropout\n dropout_rate: dropout rate\n deterministic: bool, deterministic or not (to apply dropout)\n precision: numerical precision of the computation see `jax.lax.Precision`\n for details.\n kernel_init: initializer for the kernel of the Dense layers.\n bias_init: initializer for the bias of the Dense layers.\n bias: bool: whether pointwise QKVO dense transforms use bias.\n attention_fn: dot_product_attention or compatible function. Accepts\n query, key, value, and returns output of shape\n `[bs, dim1, dim2, ..., dimN,, num_heads, value_channels]``\n\n Returns:\n output of shape `[bs, dim1, dim2, ..., dimN, features]`.\n \"\"\"\n\n assert (\n causal_mask or not cache\n ), 'Caching is only support for causal attention.'\n\n if inputs_kv is None:\n inputs_kv = inputs_q\n\n if attention_axis is None:\n attention_axis = tuple(range(1, inputs_q.ndim - 1))\n\n features = out_features or inputs_q.shape[-1]\n qkv_features = qkv_features or inputs_q.shape[-1]\n\n assert (\n qkv_features % num_heads == 0\n ), 'Memory dimension must be divisible by number of heads.'\n head_dim = qkv_features // num_heads\n\n dense = functools.partial(\n dense_general,\n axis=-1,\n dtype=dtype,\n features=(num_heads, head_dim),\n kernel_init=kernel_init,\n bias_init=bias_init,\n bias=bias,\n precision=precision,\n )\n # project inputs_q to multi-headed q/k/v\n # dimensions are then [bs, dims..., n_heads, n_features_per_head]\n query = scope.child(dense, 'query')(inputs_q)\n key = scope.child(dense, 'key')(inputs_kv)\n value = scope.child(dense, 'value')(inputs_kv)\n\n if cache:\n cache_entry: Callable[[Any], CacheEntry] | CacheEntry\n if not scope.has_variable('cache', 'entry'):\n ndim, tail_shape = (key.ndim, key.shape[-2:])\n\n def init_fn(shape, dtype=jnp.float32):\n full_shape = shape + tail_shape\n if len(full_shape) != ndim:\n raise ValueError(\n 'Shape should be a tuple with the shape of the batch'\n 'and attention dims.'\n )\n return CacheEntry(\n key=jnp.zeros(full_shape, dtype),\n value=jnp.zeros(full_shape, dtype),\n i=jnp.zeros((), jnp.uint32),\n )\n\n cache_entry = init_fn\n else:\n cache_entry = scope.get_variable('cache', 'entry')\n if not isinstance(cache_entry, CacheEntry):\n raise ValueError('Cache is not initialized.')\n\n expected_shape = list(cache_entry.key.shape[:-2])\n for attn_dim in attention_axis:\n expected_shape[attn_dim] = 1\n expected_shape = tuple(expected_shape) + inputs_q.shape[-1:]\n if expected_shape != inputs_q.shape:\n raise ValueError(\n 'Invalid shape provided, expected shape %s instead got %s.'\n % (expected_shape, inputs_q.shape)\n )\n\n cshape = cache_entry.key.shape\n indices = [0] * len(cshape)\n i = cache_entry.i\n attn_size = np.prod(np.take(cshape, attention_axis))\n for attn_dim in attention_axis:\n attn_size //= cshape[attn_dim]\n indices[attn_dim] = i // attn_size\n i = i % attn_size\n\n key = lax.dynamic_update_slice(cache_entry.key, key, indices) # type: ignore\n value = lax.dynamic_update_slice(cache_entry.value, value, indices) # type: ignore\n one = jnp.array(1, jnp.uint32)\n cache_entry = cache_entry.replace(\n i=cache_entry.i + one, key=key, value=value\n )\n\n # TODO(levskaya): verify this is still needed in translation decoding.\n key_padding_mask = jnp.broadcast_to(\n (jnp.arange(cshape[1]) < cache_entry.i), cshape[:2]\n )\n key_padding_mask = key_padding_mask.astype(jnp.float32)[..., None]\n scope.put_variable('cache', 'entry', cache_entry)\n\n # create attention masks\n mask_components = []\n\n if causal_mask:\n if cache and isinstance(cache_entry, CacheEntry):\n bias_pre_shape = (1,) * (key.ndim - 1)\n attn_shape = tuple(np.take(key.shape, attention_axis))\n attn_size = np.prod(attn_shape)\n ii = jnp.arange(attn_size, dtype=jnp.uint32)\n mask = ii < cache_entry.i\n mask_components.append(mask.reshape(bias_pre_shape + attn_shape))\n else:\n mask_components.append(_make_causal_mask(key, attention_axis))\n\n if padding_mask is not None:\n if key_padding_mask is None:\n key_padding_mask = padding_mask\n padding_mask = make_padding_mask(\n padding_mask_query=padding_mask,\n padding_mask_key=key_padding_mask,\n query_shape=query.shape,\n key_shape=key.shape,\n attention_axis=attention_axis,\n )\n mask_components.append(padding_mask)\n\n if segmentation is not None:\n if key_segmentation is None:\n key_segmentation = segmentation\n segmentation_mask = make_padding_mask(\n padding_mask_query=segmentation,\n padding_mask_key=key_segmentation,\n query_shape=query.shape,\n key_shape=key.shape,\n attention_axis=attention_axis,\n segmentation_mask=True,\n )\n mask_components.append(segmentation_mask)\n\n if mask_components:\n attention_mask = mask_components[0]\n for component in mask_components[1:]:\n attention_mask = jnp.logical_and(attention_mask, component)\n\n # attention mask in the form of attention bias\n attention_bias = lax.select(\n attention_mask > 0,\n jnp.full(attention_mask.shape, 0.0).astype(dtype),\n jnp.full(attention_mask.shape, -1e10).astype(dtype),\n )\n else:\n attention_bias = None\n\n # apply attention\n x = scope.child(attention_fn)(\n query,\n key,\n value,\n dtype=dtype,\n axis=attention_axis,\n bias=attention_bias,\n precision=precision,\n dropout_rng=dropout_rng,\n dropout_rate=dropout_rate,\n broadcast_dropout=broadcast_dropout,\n deterministic=deterministic,\n )\n\n # back to the original inputs dimensions\n out = scope.child(dense_general, name='out')(\n x,\n features=features,\n axis=(-2, -1),\n kernel_init=kernel_init,\n bias_init=bias_init,\n bias=bias,\n dtype=dtype,\n precision=precision,\n )\n\n return out\n\n\n# TODO(flax-dev): Consider refactoring MultiHeadDotProductAttention and moving\n# causal_mask and cache support into this class instead.\n# SelfAttention = MultiHeadDotProductAttention.partial(inputs_kv=None)\n\n\ndef make_padding_mask(\n padding_mask_query,\n padding_mask_key,\n query_shape,\n key_shape,\n attention_axis=None,\n segmentation_mask=False,\n):\n \"\"\"Makes padding mask for attention weights.\n\n In case of 1d inputs (i.e., `[bs, len, features]`, the attention weights will\n be `[bs, len, len]` and this function makes a square matrix [len, len].\n\n Args:\n padding_mask_query: padding mask of query \n padding_mask_key: padding mask of query \n query_shape: shape of the query\n key_shape: shape of the key, which is equal to the shape of value.\n attention_axis: axis over which attention is applied.\n segmentation_mask: bool: if true use equality on cartesian product rather\n than outer product for constructing segmentation masks.\n Returns:\n The padding mask for attention weights.\n \"\"\"\n assert query_shape[0] == key_shape[0]\n assert len(query_shape) == len(key_shape)\n\n ndim = len(key_shape)\n if attention_axis is None:\n attention_axis = tuple(range(1, ndim - 2))\n assert isinstance(attention_axis, tuple)\n for ax in attention_axis:\n if not (ndim >= 3 and 1 <= ax < ndim - 2):\n raise ValueError(\n 'Attention axis must be between the batch axis and the last-two axes.'\n )\n\n mask_shape_final = (query_shape[0], 1) # batch_size, 1 (for all heads)s\n for ax in attention_axis:\n mask_shape_final += (query_shape[ax],)\n for ax in attention_axis:\n mask_shape_final += (key_shape[ax],)\n\n padding_mask_query = padding_mask_query[..., None]\n padding_mask_key = padding_mask_key[..., None]\n perm = (0,) + tuple(np.flip(np.arange(padding_mask_key.ndim)))[:-1]\n if segmentation_mask:\n mask = jnp.equal(padding_mask_query, padding_mask_key.transpose(perm))\n else:\n mask = jnp.multiply(padding_mask_query, padding_mask_key.transpose(perm))\n\n mask = mask.reshape(mask_shape_final)\n mask = jax.lax.convert_element_type(mask, jnp.float32)\n return mask\n\n\ndef _make_causal_mask(key, attention_axis=None, self_mask=False):\n \"\"\"Makes a causal mask, to be used for masking out the future for attention.\n\n In case of 1d inputs (i.e., `[bs, len, features]`, the attention weights will\n be `[bs, len, len]` and this function makes a square matrix [len, len] with\n zeros in upper triangle and ones in lower triangle.\n\n Args:\n key: shape of the key, which is equal to the shape of value and is\n assumed to be equal to the shape of the query (since this is used in\n self-attention when decoding).\n attention_axis: axis over which attention is applied.\n self_mask: if mask out the diagonal or not.\n\n Returns:\n A causal mask to be used to mask out future positions.\n \"\"\"\n if attention_axis is None:\n attention_axis = tuple(range(1, key.ndim - 2))\n assert isinstance(attention_axis, tuple)\n for ax in attention_axis:\n if not (key.ndim >= 3 and 1 <= ax < key.ndim - 2):\n raise ValueError(\n 'Attention axis must be between the batch axis and the last-two axes.'\n )\n\n mask_shape = tuple([1] * (key.ndim - len(attention_axis) - 1))\n mask_shape_final = mask_shape\n for _ in range(2):\n flatten_dim = 1\n for ax in attention_axis:\n mask_shape_final += (key.shape[ax],)\n flatten_dim *= key.shape[ax]\n mask_shape += (flatten_dim,)\n\n def tri(n, m, k=0):\n # Tie in the key to avoid the mask becoming a constant.\n # This way XLA can construct the mask during computation and fuse it\n # with the attention ops.\n x = jnp.arange(n, dtype=jnp.int32)\n y = jnp.arange(m, dtype=jnp.int32)\n mask = lax.ge(\n (lax.broadcast_in_dim(x, shape=(n, m), broadcast_dimensions=(0,))) + k,\n lax.broadcast(y, [n]),\n )\n return mask\n\n k = -1 if self_mask else 0\n mask = tri(*mask_shape[-2:], k=k).reshape(mask_shape_final)\n return mask\n" }, { "path": "flax/core/nn/linear.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Linear modules.\"\"\"\n\nfrom collections.abc import Iterable # pylint: disable=g-importing-member\n\nimport jax.numpy as jnp\nimport numpy as np\nfrom jax import lax\n\nfrom flax import struct\nfrom flax.core import Scope\nfrom flax.linen import initializers\n\ndefault_kernel_init = initializers.lecun_normal()\n\n\ndef _normalize_axes(axes, ndim):\n # A tuple by convention. len(axes_tuple) then also gives the rank efficiently.\n return tuple(ax if ax >= 0 else ndim + ax for ax in axes)\n\n\ndef dense_general(\n scope,\n inputs,\n features,\n axis=-1,\n batch_dims=(),\n bias=True,\n dtype=jnp.float32,\n kernel_init=default_kernel_init,\n bias_init=initializers.zeros_init(),\n precision=None,\n):\n \"\"\"Applies a linear transformation to the inputs along multiple dimensions.\n\n Args:\n inputs: The nd-array to be transformed.\n features: tuple with numbers of output features.\n axis: tuple with axes to apply the transformation on.\n batch_dims: tuple with batch axes.\n bias: whether to add a bias to the output (default: True).\n dtype: the dtype of the computation (default: float32).\n kernel_init: initializer function for the weight matrix.\n bias_init: initializer function for the bias.\n precision: numerical precision of the computation see `jax.lax.Precision`\n for details.\n Returns:\n The transformed input.\n \"\"\"\n inputs = jnp.asarray(inputs, dtype)\n\n if not isinstance(features, Iterable):\n features = (features,)\n if not isinstance(axis, Iterable):\n axis = (axis,)\n if not isinstance(batch_dims, Iterable):\n batch_dims = (batch_dims,)\n features, axis, batch_dims = tuple(features), tuple(axis), tuple(batch_dims)\n\n if batch_dims:\n max_dim = np.max(batch_dims)\n if set(batch_dims) != set(range(max_dim + 1)):\n raise ValueError(\n 'batch_dims %s must be consecutive leading '\n 'dimensions starting from 0.' % str(batch_dims)\n )\n\n ndim = inputs.ndim\n n_batch_dims = len(batch_dims)\n axis = _normalize_axes(axis, ndim)\n batch_dims = _normalize_axes(batch_dims, ndim)\n n_axis, n_features = len(axis), len(features)\n\n def kernel_init_wrap(rng, shape, dtype=jnp.float32):\n size_batch_dims = np.prod(shape[:n_batch_dims], dtype=np.int32)\n flat_shape = (\n np.prod(shape[n_batch_dims : n_axis + n_batch_dims]),\n np.prod(shape[-n_features:]),\n )\n kernel = jnp.concatenate(\n [kernel_init(rng, flat_shape, dtype) for _ in range(size_batch_dims)],\n axis=0,\n )\n return jnp.reshape(kernel, shape)\n\n batch_shape = tuple(inputs.shape[ax] for ax in batch_dims)\n kernel_shape = tuple(inputs.shape[ax] for ax in axis) + features\n kernel = scope.param('kernel', kernel_init_wrap, batch_shape + kernel_shape)\n kernel = jnp.asarray(kernel, dtype)\n\n batch_ind = tuple(range(n_batch_dims))\n contract_ind = tuple(range(n_batch_dims, n_axis + n_batch_dims))\n out = lax.dot_general(\n inputs,\n kernel,\n ((axis, contract_ind), (batch_dims, batch_ind)),\n precision=precision,\n )\n if bias:\n\n def bias_init_wrap(rng, shape, dtype=jnp.float32):\n size_batch_dims = np.prod(shape[:n_batch_dims], dtype=np.int32)\n flat_shape = (np.prod(shape[-n_features:]),)\n bias = jnp.concatenate(\n [bias_init(rng, flat_shape, dtype) for _ in range(size_batch_dims)],\n axis=0,\n )\n return jnp.reshape(bias, shape)\n\n bias = scope.param('bias', bias_init_wrap, batch_shape + features)\n\n # Reshape bias for broadcast.\n expand_dims = sorted(set(range(inputs.ndim)) - set(axis) - set(batch_dims))\n for ax in expand_dims:\n bias = jnp.expand_dims(bias, ax)\n bias = jnp.asarray(bias, dtype)\n out = out + bias\n return out\n\n\ndef dense(\n scope,\n inputs,\n features,\n bias=True,\n dtype=jnp.float32,\n precision=None,\n kernel_init=default_kernel_init,\n bias_init=initializers.zeros_init(),\n):\n \"\"\"Applies a linear transformation to the inputs along the last dimension.\n\n Args:\n inputs: The nd-array to be transformed.\n features: the number of output features.\n bias: whether to add a bias to the output (default: True).\n dtype: the dtype of the computation (default: float32).\n precision: numerical precision of the computation see `jax.lax.Precision`\n for details.\n kernel_init: initializer function for the weight matrix.\n bias_init: initializer function for the bias.\n Returns:\n The transformed input.\n \"\"\"\n inputs = jnp.asarray(inputs, dtype)\n kernel = scope.param('kernel', kernel_init, (inputs.shape[-1], features))\n kernel = jnp.asarray(kernel, dtype)\n y = lax.dot_general(\n inputs,\n kernel,\n (((inputs.ndim - 1,), (0,)), ((), ())),\n precision=precision,\n )\n if bias:\n bias = scope.param('bias', bias_init, (features,))\n bias = jnp.asarray(bias, dtype)\n y += jnp.reshape(bias, (1,) * (y.ndim - 1) + (-1,))\n return y\n\n\ndef _conv_dimension_numbers(input_shape):\n \"\"\"Computes the dimension numbers based on the input shape.\"\"\"\n ndim = len(input_shape)\n lhs_spec = (0, ndim - 1) + tuple(range(1, ndim - 1))\n rhs_spec = (ndim - 1, ndim - 2) + tuple(range(0, ndim - 2))\n out_spec = lhs_spec\n return lax.ConvDimensionNumbers(lhs_spec, rhs_spec, out_spec)\n\n\ndef conv(\n scope,\n inputs,\n features,\n kernel_size,\n strides=None,\n padding='SAME',\n input_dilation=None,\n kernel_dilation=None,\n feature_group_count=1,\n bias=True,\n dtype=jnp.float32,\n precision=None,\n kernel_init=default_kernel_init,\n bias_init=initializers.zeros_init(),\n):\n \"\"\"Applies a convolution to the inputs.\n\n Args:\n inputs: input data with dimensions (batch, spatial_dims..., features).\n features: number of convolution filters.\n kernel_size: shape of the convolutional kernel.\n strides: a sequence of `n` integers, representing the inter-window\n strides.\n padding: either the string `'SAME'`, the string `'VALID'`, or a sequence\n of `n` `(low, high)` integer pairs that give the padding to apply before\n and after each spatial dimension.\n input_dilation: `None`, or a sequence of `n` integers, giving the\n dilation factor to apply in each spatial dimension of `inputs`.\n Convolution with input dilation `d` is equivalent to transposed\n convolution with stride `d`.\n kernel_dilation: `None`, or a sequence of `n` integers, giving the\n dilation factor to apply in each spatial dimension of the convolution\n kernel. Convolution with kernel dilation is also known as 'atrous\n convolution'.\n feature_group_count: integer, default 1. If specified divides the input\n features into groups.\n bias: whether to add a bias to the output (default: True).\n dtype: the dtype of the computation (default: float32).\n precision: numerical precision of the computation see `jax.lax.Precision`\n for details.\n kernel_init: initializer for the convolutional kernel.\n bias_init: initializer for the bias.\n Returns:\n The convolved data.\n \"\"\"\n\n inputs = jnp.asarray(inputs, dtype)\n\n if strides is None:\n strides = (1,) * (inputs.ndim - 2)\n\n in_features = inputs.shape[-1]\n assert in_features % feature_group_count == 0\n kernel_shape = kernel_size + (in_features // feature_group_count, features)\n kernel = scope.param('kernel', kernel_init, kernel_shape)\n kernel = jnp.asarray(kernel, dtype)\n\n dimension_numbers = _conv_dimension_numbers(inputs.shape)\n y = lax.conv_general_dilated(\n inputs,\n kernel,\n strides,\n padding,\n lhs_dilation=input_dilation,\n rhs_dilation=kernel_dilation,\n dimension_numbers=dimension_numbers,\n feature_group_count=feature_group_count,\n precision=precision,\n )\n\n if bias:\n bias = scope.param('bias', bias_init, (features,))\n bias = jnp.asarray(bias, dtype)\n y += jnp.reshape(bias, (1,) * (y.ndim - 1) + (-1,))\n return y\n\n\ndef conv_transpose(\n scope,\n inputs,\n features,\n kernel_size,\n strides=None,\n padding='SAME',\n kernel_dilation=None,\n bias=True,\n dtype=jnp.float32,\n precision=None,\n kernel_init=default_kernel_init,\n bias_init=initializers.zeros_init(),\n):\n \"\"\"Applies a transposed convolution to the inputs. Behaviour mirrors that of\n `jax.lax.conv_transpose`.\n\n Args:\n scope: functional scope.\n inputs: input data with dimensions (batch, spatial_dims..., features).\n features: number of convolution filters.\n kernel_size: shape of the convolutional kernel.\n strides: a sequence of `n` integers, representing the inter-window\n strides.\n padding: either the string `'SAME'`, the string `'VALID'`, or a sequence\n of `n` `(low, high)` integer pairs that give the padding to apply before\n and after each spatial dimension.\n kernel_dilation: `None`, or a sequence of `n` integers, giving the\n dilation factor to apply in each spatial dimension of the convolution\n kernel. Convolution with kernel dilation is also known as 'atrous\n convolution'.\n bias: whether to add a bias to the output (default: True).\n dtype: the dtype of the computation (default: float32).\n precision: numerical precision of the computation see `jax.lax.Precision`\n for details.\n kernel_init: initializer for the convolutional kernel.\n bias_init: initializer for the bias.\n Returns:\n The convolved data.\n \"\"\"\n inputs = jnp.asarray(inputs, dtype)\n strides = strides or (1,) * (inputs.ndim - 2)\n\n in_features = inputs.shape[-1]\n kernel_shape = kernel_size + (in_features, features)\n kernel = scope.param('kernel', kernel_init, kernel_shape)\n kernel = jnp.asarray(kernel, dtype)\n\n y = lax.conv_transpose(\n inputs,\n kernel,\n strides,\n padding,\n rhs_dilation=kernel_dilation,\n precision=precision,\n )\n\n if bias:\n bias = scope.param('bias', bias_init, (features,))\n bias = jnp.asarray(bias, dtype)\n y += jnp.reshape(bias, (1,) * (y.ndim - 1) + (-1,))\n return y\n\n\ndefault_embed_init = initializers.variance_scaling(\n 1.0, 'fan_in', 'normal', out_axis=0\n)\n\n\n@struct.dataclass\nclass Embedding:\n table: np.ndarray\n\n def lookup(self, indices):\n \"\"\"Embeds the inputs along the last dimension.\n\n Args:\n indices: input data, all dimensions are considered batch dimensions.\n\n Returns:\n Output which is embedded input data. The output shape follows the input,\n with an additional `features` dimension appended.\n \"\"\"\n if indices.dtype not in [jnp.int32, jnp.int64, jnp.uint32, jnp.uint64]:\n raise ValueError('Input type must be an integer or unsigned integer.')\n return self.table[indices]\n\n def attend(self, query):\n \"\"\"Attend over the embedding using a query array.\n\n Args:\n query: array with last dimension equal the feature depth `features` of the\n embedding.\n\n Returns:\n An array with final dim `num_embeddings` corresponding to the batched\n inner-product of the array of query vectors against each embedding.\n Commonly used for weight-sharing between embeddings and logit transform\n in NLP models.\n \"\"\"\n return jnp.dot(query, self.table.T)\n\n\ndef embedding(\n scope: Scope, num_embeddings: int, features: int, init_fn=default_embed_init\n) -> Embedding:\n \"\"\"Creates embedding dataclass.\n\n Args:\n num_embeddings: number of embeddings.\n features: Number of feature dimensions for each embedding.\n embedding_init: embedding initializer.\n\n Returns:\n Embedding dataclass with lookup and attend methods.\n \"\"\"\n table = scope.param('table', init_fn, (num_embeddings, features))\n return Embedding(table) # type: ignore\n" }, { "path": "flax/core/nn/normalization.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Normalization modules for Flax.\"\"\"\n\nimport jax.numpy as jnp\nfrom jax import lax\n\nfrom flax.core import Scope\nfrom flax.linen import initializers\n\n\ndef _absolute_dims(ndim, dims):\n return tuple(ndim + dim if dim < 0 else dim for dim in dims)\n\n\ndef batch_norm(\n scope: Scope,\n x,\n use_running_average=False,\n axis=-1,\n momentum=0.99,\n epsilon=1e-5,\n dtype=jnp.float32,\n bias=True,\n scale=True,\n bias_init=initializers.zeros_init(),\n scale_init=initializers.ones_init(),\n axis_name=None,\n axis_index_groups=None,\n kind='batch_stats',\n):\n x = jnp.asarray(x, jnp.float32)\n axis = axis if isinstance(axis, tuple) else (axis,)\n axis = _absolute_dims(x.ndim, axis)\n redux = tuple(i for i in range(x.ndim) if i not in axis)\n\n def pmean(x):\n m = jnp.mean(x, redux, keepdims=True)\n if axis_name is not None:\n m = lax.pmean(m, axis_name=axis_name, axis_index_groups=axis_index_groups)\n return m\n\n mean = pmean(x)\n squeeze_shape = jnp.squeeze(mean).shape\n mean2 = pmean(jnp.square(x))\n var = mean2 - jnp.square(mean)\n\n is_init = not scope.has_variable(kind, 'mean')\n ra_mean = scope.variable(kind, 'mean', jnp.zeros, squeeze_shape)\n ra_var = scope.variable(kind, 'var', jnp.ones, squeeze_shape)\n\n if use_running_average:\n # if ra_mean is not None:\n # raise ValueError('batch_stats should be provided if use_running_averages=True')\n mean = jnp.reshape(ra_mean.value, mean.shape)\n var = jnp.reshape(ra_var.value, var.shape)\n else:\n if not is_init:\n beta = 1.0 - momentum\n ra_mean.value += beta * (jnp.squeeze(mean) - ra_mean.value)\n ra_var.value += beta * (jnp.squeeze(var) - ra_var.value)\n y = x - mean\n mul = lax.rsqrt(var + epsilon)\n if scale:\n mul = mul * scope.param('scale', scale_init, squeeze_shape).reshape(\n mean.shape\n )\n y = y * mul\n if bias:\n y = y + scope.param('bias', bias_init, squeeze_shape).reshape(mean.shape)\n return jnp.asarray(y, dtype)\n\n\ndef layer_norm(\n scope: Scope,\n x,\n epsilon=1e-6,\n dtype=jnp.float32,\n bias=True,\n scale=True,\n bias_init=initializers.zeros_init(),\n scale_init=initializers.ones_init(),\n):\n \"\"\"Applies layer normalization on the input.\n It normalizes the activations of the layer for each given example in a\n batch independently, rather than across a batch like Batch Normalization.\n i.e. applies a transformation that maintains the mean activation within\n each example close to 0 and the activation standard deviation close to 1.\n Args:\n x: the inputs\n epsilon: A small float added to variance to avoid dividing by zero.\n dtype: the dtype of the computation (default: float32).\n bias: If True, bias (beta) is added.\n scale: If True, multiply by scale (gamma). When the next layer is linear\n (also e.g. nn.relu), this can be disabled since the scaling will be done\n by the next layer.\n bias_init: Initializer for bias, by default, zero.\n scale_init: Initializer for scale, by default, one.\n Returns:\n Normalized inputs (the same shape as inputs).\n \"\"\"\n features = x.shape[-1]\n mean = jnp.mean(x, axis=-1, keepdims=True)\n mean2 = jnp.mean(lax.square(x), axis=-1, keepdims=True)\n var = mean2 - lax.square(mean)\n mul = lax.rsqrt(var + epsilon)\n if scale:\n mul = mul * jnp.asarray(\n scope.param('scale', scale_init, (features,)), dtype\n )\n y = (x - mean) * mul\n if bias:\n y = y + jnp.asarray(scope.param('bias', bias_init, (features,)), dtype)\n return y\n\n\ndef group_norm(\n scope,\n x,\n num_groups=32,\n group_size=None,\n epsilon=1e-6,\n dtype=jnp.float32,\n bias=True,\n scale=True,\n bias_init=initializers.zeros_init(),\n scale_init=initializers.ones_init(),\n):\n \"\"\"Applies group normalization to the input (arxiv.org/abs/1803.08494).\n This op is similar to batch normalization, but statistics are shared across\n equally-sized groups of channels and not shared across batch dimension.\n Thus, group normalization does not depend on the batch composition and does\n not require maintaining internal state for storing statistics.\n The user should either specify the total number of channel groups or the\n number of channels per group.\n Args:\n x: the input of shape N...C, where N is a batch dimension and C is a\n channels dimensions. `...` represents an arbitrary number of extra\n dimensions that are used to accumulate statistics over.\n num_groups: the total number of channel groups. The default value of 32 is\n proposed by the original group normalization paper.\n group_size: the number of channels in a group.\n epsilon: A small float added to variance to avoid dividing by zero.\n dtype: the dtype of the computation (default: float32).\n bias: If True, bias (beta) is added.\n scale: If True, multiply by scale (gamma). When the next layer is linear\n (also e.g. nn.relu), this can be disabled since the scaling will be done\n by the next layer.\n bias_init: Initializer for bias, by default, zero.\n scale_init: Initializer for scale, by default, one.\n Returns:\n Normalized inputs (the same shape as inputs).\n \"\"\"\n x = jnp.asarray(x, jnp.float32)\n if (num_groups is None and group_size is None) or (\n num_groups is not None and group_size is not None\n ):\n raise ValueError(\n 'Either `num_groups` or `group_size` should be '\n 'specified, but not both of them.'\n )\n\n if group_size is not None:\n channels = x.shape[-1]\n if channels % group_size != 0:\n raise ValueError(\n 'Number of channels ({}) is not multiple of the '\n 'group size ({}).'.format(channels, group_size)\n )\n num_groups = channels // group_size\n\n input_shape = x.shape\n group_shape = x.shape[:-1] + (num_groups, x.shape[-1] // num_groups)\n\n x = x.reshape(group_shape)\n\n reduction_axis = list(range(1, x.ndim - 2)) + [x.ndim - 1]\n\n mean = jnp.mean(x, axis=reduction_axis, keepdims=True)\n mean_of_squares = jnp.mean(jnp.square(x), axis=reduction_axis, keepdims=True)\n var = mean_of_squares - jnp.square(mean)\n\n x = (x - mean) * lax.rsqrt(var + epsilon)\n\n x = x.reshape(input_shape)\n\n feature_shape = tuple([1 for d in input_shape[:-1]] + [input_shape[-1]])\n if scale:\n x = x * scope.param('scale', scale_init, feature_shape)\n if bias:\n x = x + scope.param('bias', bias_init, feature_shape)\n\n return x.astype(dtype)\n" }, { "path": "flax/core/nn/stochastic.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Stochastic modules.\"\"\"\n\nimport jax.numpy as jnp\nfrom jax import lax, random\n\n\ndef dropout(scope, inputs, rate, deterministic=False, rng=None):\n \"\"\"Applies a random dropout mask to the input.\n Args:\n inputs: the inputs that should be randomly masked.\n rate: the probablity of masking out a value.\n deterministic: if false the inputs are scaled by `1 / (1 - rate)` and\n masked, whereas if true, no mask is applied and the inputs are returned as\n is.\n rng: an optional `jax.random.PRNGKey`. By default `nn.make_rng()` will\n be used.\n Returns:\n The masked inputs.\n \"\"\"\n if rate == 0.0:\n return inputs\n keep_prob = 1.0 - rate\n\n if deterministic:\n return inputs\n else:\n if rng is None:\n rng = scope.make_rng('dropout')\n mask = random.bernoulli(rng, p=keep_prob, shape=inputs.shape)\n return lax.select(mask, inputs / keep_prob, jnp.zeros_like(inputs))\n" }, { "path": "flax/core/partial_eval.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport functools\nfrom typing import Any\n\nimport jax\nfrom jax import core\nfrom jax.extend import linear_util as lu\nfrom jax.interpreters import partial_eval as pe\n\nfrom flax import errors\n\n\ndef _maybe_unknown(x: Any) -> pe.PartialVal:\n if isinstance(x, jax.ShapeDtypeStruct):\n return pe.PartialVal.unknown(core.ShapedArray(x.shape, x.dtype))\n else:\n return pe.PartialVal.known(x)\n\n\ndef lazy_init(fn):\n \"\"\"Lazily evaluates a function by using the shapes of the inputs.\n\n The returned function accepts a combination of JAX values and\n ``jax.ShapeDtypeStruct`` instances for the inputs for which we\n don't need concrete values (only the shape and dtype).\n\n This API is used by ``core.lazy_init`` or ``Module.lazy_init``\n to initialize variables without doing any actual computation on the\n inputs.\n\n Args:\n fn: the function to be lazily evaluated.\n Returns:\n A new function that accepts a mix of concrete values and\n ``jax.ShapeDtypeStruct`` instances.\n \"\"\"\n\n @functools.wraps(fn)\n def wrapper(*args, **kwargs):\n # TODO(mattjj,jheek): use a public JAX API\n # flatten fn and prepare for internal JAX transform\n inputs_flat, in_tree = jax.tree_util.tree_flatten((args, kwargs))\n debug_info = jax.api_util.debug_info(\"lazy_init\", fn, (in_tree,), {})\n f_flat, out_tree = jax.api_util.flatten_fun(\n lu.wrap_init(fn, debug_info=debug_info), in_tree)\n # map inputs to PartialVal known/unknown\n # only the computations depending on knowns will be executed\n in_pvals = [_maybe_unknown(x) for x in inputs_flat]\n _, out_pvals, _ = pe.trace_to_jaxpr_nounits(f_flat, in_pvals)\n # all outputs should be knowns. If this fails\n # the user is creating variables that depend on a\n # argument that was passed as a ShapeDtypeStruct.\n out_flat = []\n for pv, const in out_pvals:\n if pv is None:\n # const is the actual value of the known output\n out_flat.append(const)\n else:\n raise errors.LazyInitError(pv)\n return jax.tree_util.tree_unflatten(out_tree(), out_flat)\n\n return wrapper\n" }, { "path": "flax/core/scope.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Flax functional core: Scopes.\"\"\"\n\nimport collections\nimport contextlib\nimport dataclasses\nimport functools\nimport hashlib\nimport typing\nfrom typing import (\n Any,\n Generic,\n Literal,\n Optional,\n TypeVar,\n Union,\n cast,\n overload,\n)\nfrom collections.abc import Callable, Iterable, Mapping, Sequence\n\nimport jax\nimport numpy as np\nfrom jax import numpy as jnp\nfrom jax import random, tree_util\n\nfrom flax import config as config\nfrom flax import errors, struct, traceback_util\nfrom flax.ids import uuid\nfrom flax.typing import (\n PRNGKey,\n Array,\n RNGSequences,\n Collection,\n MutableCollection,\n VariableDict,\n FrozenVariableDict as FrozenVariableDict,\n MutableVariableDict,\n PRNGFoldable,\n)\n\nfrom . import meta, partial_eval, tracers\nfrom .frozen_dict import FrozenDict, freeze, unfreeze\n\ntraceback_util.register_exclusion(__file__)\n\nT = TypeVar('T')\n\n\nFilter = Union[bool, str, typing.Collection[str], 'DenyList']\n\n# When conditioning on filters we require explicit boolean comparisons.\n# pylint: disable=g-bool-id-comparison\n\n\n@dataclasses.dataclass(frozen=True, eq=True)\nclass DenyList:\n \"\"\"DenyList represents an opt-out based mutability filter.\n DenyList can be used to make every collection mutable except the ones\n defined in the given filter.\n To for example make everything but the params collection mutable::\n nn.apply(fn, mutable=nn.DenyList([\"params\"]))\n Attributes:\n deny: The filter representing the collections that are not mutable.\n \"\"\"\n\n deny: Filter\n\n def __lt__(self, other):\n if isinstance(other, str):\n return False\n if isinstance(other, DenyList):\n return str(self.deny) < str(other.deny)\n return NotImplemented\n\n def __gt__(self, other):\n if isinstance(other, str):\n return True\n if isinstance(other, DenyList):\n return str(self.deny) > str(other.deny)\n return NotImplemented\n\n\nCollectionFilter = Filter\nPRNGSequenceFilter = Filter\n\n\nclass LazyRng(struct.PyTreeNode):\n \"\"\"Wrapper around JAX PRNGKey that lazily maintains a tuple of static data to be folded into the rng.\"\"\"\n\n rng: PRNGKey\n suffix: tuple[PRNGFoldable, ...] = struct.field(pytree_node=False)\n\n def as_jax_rng(self) -> PRNGKey:\n return _fold_in_static(self.rng, self.suffix)\n\n @staticmethod\n def create(\n rng: Union['LazyRng', PRNGKey], *suffix: PRNGFoldable\n ) -> 'LazyRng':\n if isinstance(rng, LazyRng):\n return LazyRng(rng.rng, rng.suffix + suffix)\n else:\n return LazyRng(rng, suffix)\n\n def clear_suffix(self):\n key = self.rng\n return LazyRng(key, ())\n\n\ndef _fold_in_static(\n rng: PRNGKey, data: typing.Collection[PRNGFoldable]\n) -> PRNGKey:\n \"\"\"Folds static data (strings & ints) into a jax.random.PRNGKey using its SHA-1 hash.\n\n This is faster than splitting an PRNGKey because it allows generating new PRNG\n keys in parallel that are independent of each other.\n\n Args:\n rng: the rng to fold the string into.\n data: the string to be folded in.\n\n Returns:\n The newly generated PRNG key.\n \"\"\"\n if not data:\n return rng\n m = hashlib.sha1()\n for x in data:\n if config.flax_fix_rng_separator:\n # encode seperate to avoid collisions like for example: (\"ab\", \"c\") and (\"a\", \"bc\")\n m.update(b'\\00')\n if isinstance(x, str):\n m.update(x.encode('utf-8'))\n elif isinstance(x, int):\n m.update(x.to_bytes((x.bit_length() + 7) // 8, byteorder='big'))\n else:\n raise ValueError(f'Expected int or string, got: {x}')\n d = m.digest()\n hash_int = int.from_bytes(d[:4], byteorder='big')\n return random.fold_in(rng, jnp.uint32(hash_int)) # type: ignore\n\n\ndef is_filter_empty(filter_like: Filter) -> bool:\n \"\"\"Returns True if `filter_like` is an empty filter.\n\n Args:\n filter_like: The filter to test.\n\n Returns:\n A filter is empty when it is an empty collection, it is a bool with value\n False, ir it is a DenyList that matches everything. A string filter is never\n empty.\n \"\"\"\n if isinstance(filter_like, str):\n return False\n if isinstance(filter_like, typing.Collection):\n return not filter_like\n if isinstance(filter_like, bool):\n return not filter_like\n if isinstance(filter_like, DenyList):\n # if any arbitrary collection is in the denylist it matches everything so\n # the filter is empty. This is checked with a stub.\n return in_filter(filter_like.deny, '__flax_internal_stub__')\n raise errors.InvalidFilterError(filter_like)\n\n\ndef in_filter(filter_like: Filter, col: str) -> bool:\n \"\"\"Checks whether a filter can be applied to a collection.\n\n Used for both collections and rng sequence filters.\n\n Args:\n filter_like: a filter (either a boolean, a string, or a list of strings) for\n a collection.\n col: a collection, which is a string identifying a dictionary of data, for\n instance \"params\" or \"batch_stats\".\n\n Returns:\n True if either `filter_like` is True, equal to `col`, or a sequence\n containing `col`.\n \"\"\"\n if isinstance(filter_like, str):\n return col == filter_like\n if isinstance(filter_like, typing.Collection):\n return col in filter_like\n if isinstance(filter_like, bool):\n return filter_like\n if isinstance(filter_like, DenyList):\n return not in_filter(filter_like.deny, col)\n raise errors.InvalidFilterError(filter_like)\n\n\ndef filter_to_set(x: Filter) -> set[str]:\n \"\"\"Converts a Filter into a set of collections, fails on the infinite set.\n\n Args:\n x: a filter (boolean, string, or list of strings).\n\n Returns:\n The input filter represented as a set of strings.\n \"\"\"\n assert x is not True and not isinstance(x, DenyList), 'Infinite set'\n if x is False:\n return set()\n if isinstance(x, str):\n return {x}\n if isinstance(x, typing.Collection):\n return set(x)\n raise errors.InvalidFilterError(x)\n\n\ndef union_filters(a: Filter, b: Filter) -> Filter:\n \"\"\"Takes the union of two filters (similar to a logical or).\n\n Args:\n a: a filter.\n b: a filter.\n\n Returns:\n The union of the two input filters. For instance,\n `union_filters('f1', ['f2']) = {'f1', 'f2'}`.\n \"\"\"\n if a is True or b is True:\n return True\n if isinstance(a, DenyList) and isinstance(b, DenyList):\n return DenyList(intersect_filters(a.deny, b.deny))\n if isinstance(b, DenyList):\n a, b = b, a\n if isinstance(a, DenyList):\n return DenyList(subtract_filters(a.deny, b))\n\n a = filter_to_set(a)\n b = filter_to_set(b)\n return a.union(b)\n\n\ndef subtract_filters(a: Filter, b: Filter) -> Filter:\n \"\"\"Returns the subtraction of b from a.\n\n Args:\n a: a filter.\n b: a filter.\n\n Returns:\n A filter matching with values in a that are not in b.\n \"\"\"\n if b is True:\n return False\n if a is True:\n return DenyList(b)\n if isinstance(a, DenyList) and isinstance(b, DenyList):\n return subtract_filters(b.deny, a.deny)\n if isinstance(a, DenyList):\n return DenyList(union_filters(a.deny, b))\n if isinstance(b, DenyList):\n return intersect_filters(a, b.deny)\n a = filter_to_set(a)\n b = filter_to_set(b)\n return a - b\n\n\ndef intersect_filters(a: Filter, b: Filter) -> Filter:\n \"\"\"Take the intersection of two filters (similar to a logical and).\n\n Args:\n a: a filter.\n b: a filter.\n\n Returns:\n The intersection of the two input filters. For instance,\n `intersect_filters('f1', ['f1', 'f2']) = {'f1'}`.\n \"\"\"\n if a is True:\n return b\n if b is True:\n return a\n if isinstance(a, DenyList) and isinstance(b, DenyList):\n return DenyList(union_filters(b.deny, a.deny))\n if isinstance(b, DenyList):\n b, a = a, b\n if isinstance(a, DenyList):\n return subtract_filters(b, a.deny)\n a = filter_to_set(a)\n b = filter_to_set(b)\n return a.intersection(b)\n\n\ndef group_collections(\n xs: VariableDict, col_filters: Sequence[CollectionFilter]\n) -> Sequence[MutableVariableDict]:\n \"\"\"Groups variables by collection filters.\n\n Iteratively applies the filters in `col_filters` to `xs`, and adds the result\n of applying each filter to the output sequence. Each key in `xs` is only added\n to the output once.\n\n Args:\n xs: a dictionary of variables, keyed by collections (strings).\n col_filters: a list of collection filters.\n\n Returns:\n A sequence S with `len(S) == len(col_filters)`. Each `S[i]` is the result of\n applying filter `col_filters[i]` to the remaining keys in `xs`.\n \"\"\"\n cols: Iterable[str]\n cols = xs.keys()\n groups = []\n for col_filter in col_filters:\n remaining_cols = []\n group = {}\n for col in cols:\n if in_filter(col_filter, col):\n group[col] = jax.tree_util.tree_map(lambda x: x, xs[col])\n else:\n remaining_cols.append(col)\n cols = remaining_cols\n groups.append(group)\n return tuple(groups)\n\n\nclass Variable(Generic[T]):\n \"\"\"A Variable object allows mutable access to a variable in a VariableDict.\n\n Variables are identified by a collection (e.g., \"batch_stats\") and a name\n (e.g., \"moving_mean\"). The value property gives access to the variable's\n content and can be assigned to for mutation.\n \"\"\"\n\n def __init__(self, scope: 'Scope', collection: str, name: str, unbox: bool):\n \"\"\"Initializes a variable.\n\n Args:\n scope: The scope in which the variable is stored.\n collection: The collection of the variable (e.g., \"params\").\n name: The name of the variable (e.g., \"dense\").\n unbox: Whether to unbox boxed values with metadata.\n \"\"\"\n self._id = uuid()\n self.scope = scope\n self.collection = collection\n self.name = name\n self.unbox = unbox\n\n @property\n def value(self) -> T:\n \"\"\"Returns the value of this Variable.\"\"\"\n v = self.scope.get_variable(self.collection, self.name)\n return meta.unbox(v) if self.unbox else v\n\n @value.setter\n def value(self, value: T):\n \"\"\"Updates the value of this Variable.\"\"\"\n if self.unbox:\n cur = self.scope.get_variable(self.collection, self.name)\n cur_struct = tree_util.tree_structure(cur, is_leaf=meta.is_axis_metadata)\n value_struct = tree_util.tree_structure(\n value, is_leaf=meta.is_axis_metadata\n )\n has_meta = any(map(meta.is_axis_metadata, cur_struct.flatten_up_to(cur)))\n if cur_struct == value_struct and has_meta: # type: ignore[operator]\n value = meta.replace_boxed(cur, value)\n\n self.scope.put_variable(self.collection, self.name, value)\n\n def is_mutable(self) -> bool:\n \"\"\"Checks if this Variable is mutable.\"\"\"\n return self.scope.is_mutable_collection(self.collection)\n\n\nclass _ChildRNGSentinel:\n pass\n\n\n# used to identify that an rng counter is meant for a child scope\nchild_rng_token = _ChildRNGSentinel()\n\n\nclass _DefaultSentinel:\n pass\n\n\n# used to denote no default flag value on scope\nno_flag = _DefaultSentinel()\n\n\n# Make sure reference sharing of child variable dictionaries isn't broken.\n# See https://github.com/google/flax/issues/2022 for more details.\ndef _put_variable(target, key, val):\n if (\n key in target\n and isinstance(target[key], dict)\n and isinstance(val, Mapping)\n ):\n for k, v in val.items():\n _put_variable(target[key], k, v)\n else:\n target[key] = val\n\n\nclass Scope:\n \"\"\"A Scope allows easy access to variables and manages RNGS of a neural network layer.\n\n Scopes are purely functional and encapsulated in\n :class:`flax.linen.module.Module`, so users writing neural network code\n usually generally do not interact with ``Scopes`` directly.\n\n See `core design tests\n `_\n for a number of examples using ``Scopes``.\n \"\"\"\n\n reservations: dict[str, set[str | None]]\n\n def __init__(\n self,\n variables: MutableVariableDict,\n rngs: RNGSequences | dict[str, LazyRng] | None = None,\n name: str | None = None,\n mutable: CollectionFilter = False,\n parent: Optional['Scope'] = None,\n path: Iterable[str] = (),\n debug_path: Iterable[str] = (),\n flags: Mapping | None = None,\n ):\n \"\"\"Initializes a Scope.\n\n Args:\n variables: VariableDict to initialize the Scope with.\n rngs: RNGs used in this scope or one of the child scopes.\n name: name of this scope.\n mutable: A CollectionFilter determining which variables are mutable.\n parent: The parent scope.\n path: The path in the variable tree from the root scope to this scope. It\n exactly matches the module path.\n debug_path: Similar to path but could contain transformation decorators.\n flags: internal flags.\n \"\"\"\n rngs = {k: LazyRng.create(v) for k, v in rngs.items()} if rngs else {}\n self._variables = variables\n self.parent = parent\n self.name = name\n self.path = tuple(path)\n self.debug_path = tuple(debug_path) or self.path\n self.rngs = rngs\n self.mutable = mutable\n self.flags = freeze({} if flags is None else flags)\n\n self._root = parent.root if parent else None\n self.trace_level = tracers.current_trace()\n\n self.rng_counters = {key: 0 for key in self.rngs}\n self.reservations = collections.defaultdict(set)\n\n self._invalid = False\n\n def __eq__(self, other: Any) -> bool:\n # If the root variable dict and path are the same, then two scopes behave\n # identically. Effectively, a scope is nothing more than a cursor into a\n # variable dict and an rng counter dict.\n if not isinstance(other, Scope):\n return False\n if self is other:\n return True\n return (\n self.root._variables is other.root._variables\n and self.path == other.path\n and self.rng_counters is other.rng_counters\n )\n\n def __hash__(self) -> int:\n # see __eq__\n return hash((id(self.root._variables), self.path, id(self.rng_counters)))\n\n @property\n def root(self) -> 'Scope':\n return self._root or self\n\n @property\n def path_text(self) -> str:\n \"\"\"Returns the debug path as a human readable string.\"\"\"\n return '/' + '/'.join(self.debug_path)\n\n @property\n def invalid(self) -> bool:\n \"\"\"Returns true if this scope is invalidated as a result of `Scope.temporary`.\"\"\"\n return self._invalid\n\n def _check_valid(self):\n if self._invalid:\n raise errors.InvalidScopeError(self.name)\n\n @contextlib.contextmanager\n def temporary(self):\n \"\"\"Returns a context manager that will invalidate this Scope when leaving the context.\"\"\"\n try:\n yield self\n finally:\n self.invalidate()\n\n def invalidate(self):\n \"\"\"Invalidates the Scope.\"\"\"\n self._invalid = True\n\n def mutable_variables(self) -> VariableDict | dict[str, Any]:\n \"\"\"Returns an immutable copy of the mutable variables belonging to this Scope.\"\"\"\n self._populate_collections()\n xs = {\n k: v for k, v in self._variables.items() if in_filter(self.mutable, k)\n }\n if config.flax_return_frozendict:\n return freeze(xs)\n return xs\n\n def variables(self) -> VariableDict | dict[str, Any]:\n \"\"\"Returns an immutable copy of the variables belonging to this Scope.\"\"\"\n self._populate_collections()\n if config.flax_return_frozendict:\n return freeze(self._variables)\n return self._variables\n\n def _validate_trace_level(self):\n tracers.check_trace_level(self.trace_level)\n\n def rewound(self, rewind_rngs: bool = False) -> 'Scope':\n \"\"\"Returns a rewound version of this Scope.\n\n Args:\n rewind_rngs: if true, reset the RNG counter of this scope.\n\n Returns:\n A rewound version of this scope, which means reservations are\n emptied, and the rng counter is optionally rewound.\n \"\"\"\n self._check_valid()\n scope = Scope(\n self._variables,\n self.rngs,\n self.name,\n self.mutable,\n self.parent,\n path=self.path,\n debug_path=self.debug_path,\n flags=self.flags,\n )\n if not rewind_rngs:\n scope.rng_counters = self.rng_counters\n return scope\n\n def name_reserved(self, name: str, col: str | None = None) -> bool:\n \"\"\"Checks whether a name for a child Scope or Variable is taken.\n\n Args:\n name: the name to check for collision.\n col: if a variable, the collection used.\n \"\"\"\n if name in self.reservations:\n # allow the same name for two variables in\n # different collections, otherwise raise error.\n if (\n None in self.reservations[name]\n or col is None\n or col in self.reservations[name]\n ):\n return True\n return False\n\n def reserve(self, name: str, col: str | None = None):\n \"\"\"Reserves a name for a child Scope or Variable.\n\n Throws an error if the name exists already.\n\n Args:\n name: the name to reserve.\n col: if a variable, the collection used.\n \"\"\"\n if not isinstance(name, str):\n raise TypeError(\n 'The type of scope \"{name}\" should be string but ' f'it is {type(name)}'\n )\n if self.name_reserved(name, col):\n raise ValueError(f'Duplicate use of scope name: \"{name}\"')\n self.reservations[name].add(col)\n\n def default_name(self, prefix: str) -> str:\n \"\"\"Generates an unreserved name with the given prefix.\n\n Args:\n prefix: prefix to use for generating an unreserved name.\n\n Returns:\n The generated name.\n \"\"\"\n i = 0\n while True:\n name = f'{prefix}{i}'\n if name not in self.reservations:\n return name\n i += 1\n\n def push(\n self, name: str | None = None, prefix: str = '', reuse=False\n ) -> 'Scope':\n \"\"\"Creates a child Scope.\n\n Args:\n name: optional name of the child.\n prefix: prefix used for generating the name if `name` is `None`.\n reuse: if True will return a pre-existing child scope with the given name\n instead of throwing an error.\n\n Returns:\n The child scope.\n \"\"\"\n self._check_valid()\n self._validate_trace_level()\n if name is None:\n name = self.default_name(prefix)\n if not reuse or name not in self.reservations:\n self.reserve(name)\n rngs = {key: LazyRng.create(rng, name) for key, rng in self.rngs.items()}\n rng_key = (child_rng_token, name)\n if rng_key in self.rng_counters:\n rng_counters = self.rng_counters.get(rng_key) # type: ignore\n else:\n rng_counters = {key: 0 for key in rngs}\n self.rng_counters[rng_key] = rng_counters # type: ignore\n scope = Scope(\n {},\n name=name,\n rngs=rngs,\n parent=self,\n mutable=self.mutable,\n path=self.path + (name,),\n debug_path=self.debug_path + (name,),\n flags=self.flags,\n )\n scope.rng_counters = rng_counters\n return scope\n\n def child(\n self,\n fn: Callable[..., Any],\n name: str | None = None,\n prefix: str | None = None,\n named_call: bool = True,\n **partial_kwargs,\n ) -> Callable[..., Any]:\n \"\"\"Partially applies a child scope to fn.\n\n When calling the returned function multiple times variables will be reused.\n\n Args:\n fn: the function to partially apply the child Scope to.\n name: optional name of the child.\n prefix: prefix used for generating name if it is `None`.\n named_call: if true, `fn` will be run under `jax.named_scope`. The XLA\n profiler will use this to name tag the computation.\n **partial_kwargs: additional kwargs partially applied to `fn`.\n\n Returns:\n The function with a partially applied scope.\n \"\"\"\n if name is None:\n if prefix is None:\n prefix = fn.__name__ + '_' if hasattr(fn, '__name__') else ''\n name = self.default_name(prefix)\n scope = self.push(name)\n\n @functools.wraps(fn)\n def wrapper(*args, **kwargs):\n kwargs = dict(partial_kwargs, **kwargs)\n if named_call:\n with jax.named_scope(name):\n res = fn(scope.rewound(), *args, **kwargs)\n else:\n res = fn(scope.rewound(), *args, **kwargs)\n return res\n\n return wrapper\n\n def is_mutable_collection(self, col: str) -> bool:\n \"\"\"Returns true if the collection `col` is mutable.\"\"\"\n return in_filter(self.mutable, col)\n\n def is_collection_empty(self, col: str) -> bool:\n \"\"\"Returns true if the collection is empty.\"\"\"\n if col in self.root._variables: # pylint: disable=protected-access\n return not self.root._variables[col] # pylint: disable=protected-access\n return True\n\n def _mutable_collection(self, col: str) -> MutableCollection:\n \"\"\"Returns the collection `col` as a mutable object.\"\"\"\n assert self.is_mutable_collection(col), f'Collection {col} is not mutable'\n\n # The actual variable dict is stored in the root scope only, and subscopes\n # hold references to subtrees relevant to them. This function ensures that\n # the collections are created in the top-level Scope and we return the\n # correct reference.\n if col not in self._variables:\n if not self.parent:\n # If this is the top-level Scope, just add an empty collection.\n self._variables[col] = {}\n else:\n assert self.name is not None # Only top-level Scope have name None.\n # Populate the parent collections recursively and obtain a reference to\n # the direct parent (which, by transitivity, is be a reference to a\n # dict in the root Scope).\n parent_col = self.parent._mutable_collection(col) # pylint: disable=protected-access\n if self.name not in parent_col:\n # If this Scope's name does not occur in the parent collection, add it\n # to the parent scope (updating the parent's variable dict).\n parent_col[self.name] = {}\n # Store a reference to the parent's scope collection for in this scope's\n # variable dict.\n self._variables[col] = parent_col[self.name]\n\n return self._variables[col]\n\n def _collection(self, col: str) -> Collection:\n \"\"\"Returns a collection of variables of collection `col`.\"\"\"\n if col not in self._variables:\n if self.parent:\n assert self.name is not None\n parent_col = self.parent._collection(col) # pylint: disable=protected-access\n if self.name not in parent_col:\n return FrozenDict()\n self._variables[col] = parent_col[self.name]\n else:\n return FrozenDict()\n return self._variables[col]\n\n def has_rng(self, name: str) -> bool:\n \"\"\"Returns true if a PRNGSequence with name `name` exists.\"\"\"\n return name in self.rngs\n\n def make_rng(self, name: str = 'params') -> PRNGKey:\n \"\"\"Generates A PRNGKey from a PRNGSequence with name `name`.\"\"\"\n if not self.has_rng(name):\n if self.has_rng('params'):\n name = 'params'\n else:\n raise errors.InvalidRngError(f'{self.name} needs PRNG for \"{name}\"')\n self._check_valid()\n self._validate_trace_level()\n self.rng_counters[name] += 1\n return LazyRng.create(self.rngs[name], self.rng_counters[name]).as_jax_rng()\n\n def get_variable(self, col: str, name: str, default: Any = None) -> Any:\n \"\"\"Retrieves the value of a Variable.\n\n Args:\n col: the variable collection.\n name: the name of the variable.\n default: the default value to return if the variable does not exist in\n this scope.\n\n Returns:\n The value of the input variable, of the default value if the variable\n doesn't exist in this scope.\n \"\"\"\n variables = self._collection(col)\n if name in variables:\n return variables[name]\n else:\n return default\n\n def has_variable(self, col: str, name: str) -> bool:\n \"\"\"Returns true if the given variable exists in this scope.\n\n Args:\n col: the collection of the variable.\n name: the name of the variable.\n \"\"\"\n variables = self._collection(col)\n return name in variables\n\n def put_variable(self, col: str, name: str, value: Any):\n \"\"\"Updates the value of the given variable if it is mutable, or an error otherwise.\n\n Args:\n col: the collection of the variable.\n name: the name of the variable.\n value: the new value of the given variable.\n \"\"\"\n self._check_valid()\n self._validate_trace_level()\n if not self.is_mutable_collection(col):\n raise errors.ModifyScopeVariableError(col, name, self.path_text)\n variables = self._mutable_collection(col)\n\n _put_variable(variables, name, value)\n\n @overload\n def variable(\n self,\n col: str,\n name: str,\n init_fn: Callable[..., T] | None = None,\n *init_args,\n ) -> Variable[T]:\n ...\n\n @overload\n def variable(\n self,\n col: str,\n name: str,\n init_fn: Callable[..., T] | None = None,\n *init_args,\n unbox: Literal[True],\n **init_kwargs,\n ) -> Variable[T]:\n ...\n\n @overload\n def variable(\n self,\n col: str,\n name: str,\n init_fn: Callable[..., T] | None = None,\n *init_args,\n unbox: Literal[False],\n **init_kwargs,\n ) -> Variable[meta.AxisMetadata[T]]:\n ...\n\n @overload\n def variable(\n self,\n col: str,\n name: str,\n init_fn: Callable[..., T] | None = None,\n *init_args,\n unbox: bool = True,\n **init_kwargs,\n ) -> Variable[T] | Variable[meta.AxisMetadata[T]]:\n ...\n\n def variable(\n self,\n col: str,\n name: str, # pylint: disable=keyword-arg-before-vararg\n init_fn: Callable[..., T] | None = None,\n *init_args,\n unbox: bool = True,\n **init_kwargs,\n ) -> Variable[T] | Variable[meta.AxisMetadata[T]]:\n \"\"\"Creates a variable if it doesn't exist yet in this scope and returns it.\n\n Args:\n col: the collection of the variable.\n name: the name of the variable.\n init_fn: a function taking a PRNGKey plus any other number of positional\n arguments. If None, the variable must already be initialized otherwise\n an error is raised.\n *init_args: the positional arguments to evaluate init_fn on lazily.\n unbox: If True, ``AxisMetadata`` instances are replaced by their unboxed\n value, see ``flax.nn.meta.unbox`` (default: True).\n **init_kwargs: the key-word arguments to evaluate init_fn on lazily.\n\n Returns:\n The variable. Throws an error if the variable exists already.\n \"\"\"\n self.reserve(name, col)\n if not self.has_variable(col, name):\n if not self.is_mutable_collection(col) or init_fn is None:\n if self.is_collection_empty(col):\n raise errors.ScopeCollectionNotFound(col, name, self.path_text)\n raise errors.ScopeVariableNotFoundError(name, col, self.path_text)\n init_value = init_fn(*init_args, **init_kwargs)\n self.put_variable(col, name, init_value)\n # cast to make static analyzers happy\n return cast(\n Union[Variable[T], Variable[meta.AxisMetadata[T]]],\n Variable(self, col, name, unbox=unbox),\n )\n\n @overload\n def param(\n self, name: str, init_fn: Callable[..., T], *init_args,\n ) -> T:\n ...\n\n @overload\n def param(\n self,\n name: str,\n init_fn: Callable[..., meta.AxisMetadata[T]] | Callable[..., T],\n *init_args,\n unbox: Literal[True],\n **init_kwargs,\n ) -> T:\n ...\n\n @overload\n def param(\n self,\n name: str,\n init_fn: Callable[..., T],\n *init_args,\n unbox: Literal[False],\n **init_kwargs,\n ) -> T:\n ...\n\n @overload\n def param(\n self,\n name: str,\n init_fn: Callable[..., T | meta.AxisMetadata[T]],\n *init_args,\n unbox: bool,\n **init_kwargs,\n ) -> T | meta.AxisMetadata[T]:\n ...\n\n def param(\n self,\n name: str,\n init_fn: Callable[..., T],\n *init_args,\n unbox: bool = True,\n **init_kwargs,\n ) -> T | meta.AxisMetadata[T]:\n \"\"\"Creates a parameter if it doesn't exist yet in this scope and returns it.\n\n If the parameter exists already, the existing value is simply returned.\n\n Args:\n name: the name of the parameter.\n init_fn: a function taking a PRNGKey plus any other number of positional\n arguments.\n *init_args: the positional arguments to evaluate init_fn on lazily.\n unbox: If True, ``AxisMetadata`` instances are replaced by their unboxed\n value, see ``flax.nn.meta.unbox`` (default: True).\n **init_kwargs: the key-word arguments to evaluate init_fn on lazily.\n\n Returns:\n The parameters. Throws an error if the params exist already.\n \"\"\"\n self.reserve(name, 'params')\n if self.has_variable('params', name):\n value = self.get_variable('params', name)\n if unbox:\n value = meta.unbox(value)\n # Validate that the shape of the init_fn output is the same as the shape\n # of the existing parameter. This is to make sure that the hparams set up\n # in a Flax Module match the shapes coming in during apply, and if not,\n # catch it with an error message.\n # NOTE: We could consider moving this to `self.`\n abs_value = jax.eval_shape(\n lambda: init_fn(random.key(0), *init_args, **init_kwargs)\n )\n abs_value_flat = jax.tree_util.tree_leaves(abs_value)\n value_flat = jax.tree_util.tree_leaves(value)\n for val, abs_val in zip(value_flat, abs_value_flat):\n # NOTE: We could check dtype consistency here as well but its usefulness\n # is less obvious. We might intentionally change the dtype for inference\n # to a half float type for example.\n if np.shape(val) != np.shape(abs_val):\n raise errors.ScopeParamShapeError(\n name, self.path_text, np.shape(val), np.shape(abs_val)\n )\n else:\n if not self.is_mutable_collection('params'):\n if self.is_collection_empty('params'):\n raise errors.ScopeCollectionNotFound('params', name, self.path_text)\n raise errors.ScopeParamNotFoundError(name, self.path_text)\n value = init_fn(self.make_rng('params'), *init_args, **init_kwargs)\n self.put_variable('params', name, value)\n if unbox:\n value = meta.unbox(value)\n return value\n\n def _populate_collections(self):\n collections = self.root._variables.keys() # pylint: disable=protected-access\n for col in collections:\n self._collection(col)\n\n def has_flag(self, key) -> bool:\n return key in self.flags\n\n def get_flag(self, key, default=no_flag) -> Any:\n if key not in self.flags and default is no_flag:\n return ValueError(f'Flag {key} not present on scope.')\n return self.flags.get(key, default)\n\n\ndef _unfreeze_variables(variables, mutable):\n new_variables = {}\n for key, value in variables.items():\n if in_filter(mutable, key):\n new_variables[key] = unfreeze(value)\n else:\n new_variables[key] = value\n return new_variables\n\n\ndef bind(\n variables: VariableDict,\n rngs: RNGSequences | None = None,\n mutable: CollectionFilter = False,\n flags: Mapping | None = None,\n):\n \"\"\"Binds variables and rngs to a new ``Scope``.\n\n bind provides a ``Scope`` instance without transforming a function with\n ``apply``. This is particularly useful for debugging and interactive use cases\n like notebooks where a function would limit the ability split up code into\n different cells.\n\n a ``Scope`` instance is a stateful object. Note that idiomatic JAX is\n functional and therefore a ``Scope` does not mix well well with vanilla JAX\n APIs. Therefore, we recommend using ``apply`` when code should be reusable and\n compatible across the JAX software ecosystem.\n\n Args:\n variables: Variable dictionary to bind.\n rngs: RNGs to bind.\n mutable: Which variable collections to treat as mutable.\n flags: internal flags.\n\n Returns:\n A new scope with the variables and rngs bound to it.\n \"\"\"\n if not _is_valid_variables(variables):\n raise errors.ApplyScopeInvalidVariablesTypeError()\n if rngs is not None and not _is_valid_rngs(rngs):\n raise errors.InvalidRngError(\n 'rngs should be a dictionary mapping strings to `jax.PRNGKey`.'\n )\n new_variables = _unfreeze_variables(variables, mutable)\n return Scope(new_variables, rngs=rngs, mutable=mutable, flags=flags)\n\n\ndef apply(\n fn: Callable[..., Any],\n mutable: CollectionFilter = False,\n flags: Mapping | None = None,\n) -> Callable[..., Any]:\n \"\"\"Functionalize a `Scope` function.\n\n Args:\n fn: a function taking a `Scope` as its first argument.\n mutable: the filter determining which variable collections are mutable.\n flags: internal flags.\n\n Returns:\n `fn` with the scope partially applied.\n \"\"\"\n\n @functools.wraps(fn)\n def wrapper(\n variables: VariableDict,\n *args,\n rngs: PRNGKey | RNGSequences | None = None,\n **kwargs,\n ) -> Any | tuple[Any, VariableDict | dict[str, Any]]:\n if rngs is not None:\n if not _is_valid_rng(rngs) and not _is_valid_rngs(rngs):\n raise ValueError(\n 'The ``rngs`` argument passed to an apply function should be a '\n '``jax.PRNGKey`` or a dictionary mapping strings to '\n '``jax.PRNGKey``.'\n )\n if not isinstance(rngs, (dict, FrozenDict)):\n rngs = {'params': rngs}\n\n # Try to detect if user accidentally passed {'params': {'params': ...}.\n if (\n 'params' in variables\n and isinstance(variables['params'], (dict, FrozenDict))\n and 'params' in variables['params']\n ):\n raise errors.ApplyScopeInvalidVariablesStructureError(variables)\n\n with bind(\n variables, rngs=rngs, mutable=mutable, flags=flags\n ).temporary() as root:\n y = fn(root, *args, **kwargs)\n if mutable is not False:\n return y, root.mutable_variables()\n else:\n return y\n\n return wrapper\n\n\ndef init(\n fn: Callable[..., Any],\n mutable: CollectionFilter = True,\n flags: Mapping | None = None,\n) -> Callable[..., Any]:\n \"\"\"Functionalize a `Scope` function for initialization.\n\n Args:\n fn: a function taking a `Scope` as its first argument.\n mutable: the filter determining which variable collections are mutable.\n flags: internal flags.\n\n Returns:\n `fn` with the scope partially applied.\n \"\"\"\n\n @functools.wraps(fn)\n def wrapper(rngs, *args, **kwargs) -> tuple[Any, VariableDict]:\n if not _is_valid_rng(rngs) and not _is_valid_rngs(rngs):\n raise ValueError(\n 'First argument passed to an init function should be a '\n '``jax.PRNGKey`` or a dictionary mapping strings to '\n '``jax.PRNGKey``.'\n )\n if not isinstance(rngs, (dict, FrozenDict)):\n rngs = {'params': rngs}\n init_flags = {**(flags if flags is not None else {}), 'initializing': True}\n return apply(fn, mutable=mutable, flags=init_flags)(\n {}, *args, rngs=rngs, **kwargs\n )\n\n return wrapper\n\n\ndef lazy_init(\n fn: Callable[..., Any],\n mutable: CollectionFilter = True,\n flags: Mapping | None = None,\n) -> Callable[..., Any]:\n \"\"\"Functionalizes a `Scope` function for lazy initialization.\n\n Similar to ``init`` except that the init function now accepts\n ``jax.ShapeDtypeStruct`` instances for arguments that do not\n affect the variable initialization (typically this is all the input data).\n\n Example::\n\n def f(scope, x):\n # the kernel init only uses the shape of x so we don't actually\n # need a value for x and can pass it as a ShapeDtypeStruct in lazy_init.\n k = scope.param(\"kernel\", nn.initializers.lecun_normal(), (x.shape[-1], x.shape[-1]))\n return x @ k\n init_fn = lazy_init(f)\n variables = init_fn(random.key(0), jax.ShapeDtypeStruct((1, 128), jnp.float32))\n\n\n Args:\n fn: a function taking a `Scope` as its first argument.\n mutable: the filter determining which variable collections are mutable.\n flags: internal flags.\n\n Returns:\n `fn` with the scope partially applied. Unlike ``init`` which returns a tuple of function\n output and variables, the lazy init function only returns the variables.\n \"\"\"\n return partial_eval.lazy_init(\n lambda *args, **kwargs: init(fn, mutable, flags)(*args, **kwargs)[1]\n )\n\n\ndef _is_valid_collection(col: VariableDict):\n if not isinstance(col, (FrozenDict, dict)):\n return False\n for name in col.keys():\n # Any value can be stored in a collection so only keys can be verified.\n if not isinstance(name, str):\n return False\n return True\n\n\ndef _is_valid_variables(variables: VariableDict) -> bool:\n \"\"\"Checks whether the given variable dict is valid.\n\n Args:\n variables: A variable dict.\n\n Returns:\n True if `variables` is a valid variable dict.\n \"\"\"\n for name, col in variables.items():\n if not isinstance(name, str):\n return False\n if not _is_valid_collection(col):\n return False\n return True\n\n\ndef _is_valid_rng(rng: Array):\n \"\"\"Checks whether rng is a valid JAX PRNGKey, also handling custom prngs.\"\"\"\n # Allow for user-provided LazyRng - useful for compatibility when refactoring.\n if isinstance(rng, LazyRng):\n return True\n\n # This check is valid for either new-style or old-style PRNG keys\n if not isinstance(rng, (np.ndarray, jnp.ndarray)):\n return False\n\n # Handle new-style typed PRNG keys\n if jax.dtypes.issubdtype(rng.dtype, jax.dtypes.prng_key):\n return rng.shape == ()\n\n # Handle old-style raw PRNG keys\n expected_rng = jax.eval_shape(\n lambda s: jax.random.key_data(jax.random.key(s)), 0\n )\n if (rng.shape, rng.dtype) != (expected_rng.shape, expected_rng.dtype):\n return False\n return True\n\n\ndef _is_valid_rngs(rngs: PRNGKey | RNGSequences):\n if not isinstance(rngs, (FrozenDict, dict)):\n return False\n for key, val in rngs.items():\n if not isinstance(key, str):\n return False\n if not _is_valid_rng(val):\n return False\n return True\n" }, { "path": "flax/core/spmd.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport contextlib\nimport dataclasses\nimport threading\n\nimport jax\nfrom jax.sharding import PartitionSpec, NamedSharding\nfrom flax.core import meta\nfrom jax.experimental.layout import Format\nfrom flax.typing import (\n LogicalRules,\n Sharding,\n)\n\ndef get_pspec(sharding, sharding_rules = None) -> PartitionSpec:\n \"\"\"Given an `nnx.Variable`, return its `PartitionSpec`.\"\"\"\n return apply_rules(sharding, sharding_rules) # type: ignore\n\ndef map_sharding(f, sharding):\n if isinstance(sharding, PartitionSpec) or isinstance(sharding, tuple):\n return PartitionSpec(*map(f, sharding))\n elif isinstance(sharding, NamedSharding):\n return NamedSharding(sharding.mesh, map_sharding(f, sharding.spec)) # type: ignore\n elif isinstance(sharding, Format):\n return Format(sharding.layout, map_sharding(f, sharding.sharding)) # type: ignore\n\ndef get_mesh(sharding):\n if isinstance(sharding, PartitionSpec) or isinstance(sharding, tuple):\n return None\n elif isinstance(sharding, NamedSharding):\n return sharding.mesh\n elif isinstance(sharding, Format):\n return get_mesh(sharding.sharding)\n\ndef apply_rules(sharding, sharding_rules):\n \"\"\"Rename the axes of a sharding specification (which can include `PartitionSpec`, `NamedSharding` or `Format` objects).\"\"\"\n if get_logical_axis_rules() or sharding_rules:\n context_rules = get_logical_axis_rules()\n rules = {alias: on_mesh for (alias, on_mesh) in composite_rules(context_rules, sharding_rules)}\n else:\n rules = {}\n return map_sharding(lambda a: rules.get(a, a), sharding)\n\ndef _apply_sharding(value, sharding, mesh):\n if isinstance(sharding, Format):\n return jax.lax.with_sharding_constraint(value, sharding)\n if mesh.are_all_axes_explicit:\n return jax.sharding.reshard(value, sharding)\n elif mesh.are_all_axes_auto:\n return jax.lax.with_sharding_constraint(value, sharding)\n else:\n raise ValueError(\n 'Mesh must have all axes as Explicit or all axes as Auto. '\n f'Got mixed axis types: {mesh.axis_types}')\n\ndef shard_value(value, out_sharding, sharding_rules, mesh):\n if not out_sharding:\n return value\n\n if mesh is None:\n mesh = meta.get_global_mesh()\n\n out_sharding = apply_rules(out_sharding, sharding_rules)\n\n sharding_mesh = get_mesh(out_sharding)\n\n if sharding_mesh:\n if mesh:\n assert mesh == out_sharding.mesh\n mesh = sharding_mesh\n\n if mesh is None:\n raise ValueError(\n 'An auto mesh context or metadata is required if creating a variable'\n f' with annotation {out_sharding=}. '\n 'For more guidance, see https://flax.readthedocs.io/en/latest/flip/4844-var-eager-sharding.html.')\n\n if isinstance(out_sharding, PartitionSpec):\n out_sharding = NamedSharding(mesh, out_sharding)\n\n return _apply_sharding(value, out_sharding, mesh)\n\n\n\n\n# Dynamic Axis Mapping Context\n# ------------------------------------------------------------------------------\n\n\n@dataclasses.dataclass\nclass _AxisRules(threading.local):\n \"\"\"Dynamic logical axis to mesh axis binding context.\"\"\"\n\n rules: LogicalRules = ()\n\n\n# Global axis binding context.\n_axis_rules = _AxisRules()\n\n\ndef set_logical_axis_rules(rules: LogicalRules):\n \"\"\"Sets the global logical axis to mesh axis binding.\"\"\"\n _axis_rules.rules = rules\n\n\ndef get_logical_axis_rules() -> LogicalRules:\n \"\"\"Returns the global logical axis to mesh axis binding.\"\"\"\n return _axis_rules.rules\n\n\n@contextlib.contextmanager\ndef logical_axis_rules(rules: LogicalRules):\n \"\"\"Context manager for setting the logical to mesh axis bindings.\"\"\"\n old_rules = _axis_rules.rules\n try:\n _axis_rules.rules = rules\n yield\n finally:\n _axis_rules.rules = old_rules\n\n\ndef composite_rules(rule1, rule2):\n if not rule1 and not rule2:\n return ()\n if rule1 and not rule2:\n return rule1\n if rule2 and not rule1:\n return rule2\n rules = {alias: value for alias, value in rule1}\n for alias, value in rule2:\n if alias in rules and rules[alias] != value:\n raise ValueError(\n f'Inconsistent logical axis annotations for {alias}: '\n f'{rules[alias]} vs {value}'\n )\n rules[alias] = value\n return tuple(rules.items())\n\n\ndef from_sharding_rules(\n sharding: Sharding, sharding_rules: LogicalRules\n) -> Sharding:\n rules = {alias: on_mesh for (alias, on_mesh) in sharding_rules}\n return tuple(\n rules[str(s)] if (s and str(s) in rules) else s for s in sharding\n )\n" }, { "path": "flax/core/tracers.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Functionality for inspecting jax tracers.\"\"\"\n\nimport jax\nimport jax.core\n\n\ndef current_trace():\n \"\"\"Returns the current JAX state tracer.\"\"\"\n if jax.__version_info__ <= (0, 4, 33):\n top = jax.core.find_top_trace(())\n if top:\n return top.level\n else:\n return float('-inf')\n\n return jax.core.get_opaque_trace_state(convention=\"flax\")\n\ndef check_trace_level(base_level):\n # TODO(cgarciae): skipping for now as it breaks\n # too many internal tests.\n # level = current_trace()\n # if level != base_level:\n # raise errors.JaxTransformError()\n pass\n" }, { "path": "flax/core/variables.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"A variable dict is a normal Python dictionary, which is a container for one\nor more \"variable collections\", each of which are nested dictionaries whose\nleaves are ``jax.numpy`` arrays.\n\nThe different variable collections share the same nested tree structure.\n\nFor example, consider the following variable dictionary::\n\n {\n \"params\": {\n \"Conv1\": { \"weight\": ..., \"bias\": ... },\n \"BatchNorm1\": { \"scale\": ..., \"mean\": ... },\n \"Conv2\": {...}\n },\n \"batch_stats\": {\n \"BatchNorm1\": { \"moving_mean\": ..., \"moving_average\": ...}\n }\n }\n\nIn this case, the ``\"BatchNorm1\"`` key lives in both the ``\"params\"`` and\n```\"batch_stats\"\"`` collections. This reflects the fact that the submodule\nnamed ``\"\"BatchNorm1\"\"`` has both trainable parameters (the ``\"params\"`` collection),\nas well as other non-trainable variables (the ``\"batch_stats\"`` collection)\n\nTODO: Make \"variable dict\" design note, and link to it from here.\n\"\"\"\n\nfrom .scope import Variable\n" }, { "path": "flax/cursor.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport dataclasses\nimport enum\nfrom typing import (\n Any,\n Generic,\n Protocol,\n TypeVar,\n runtime_checkable,\n)\nfrom collections.abc import Callable, Generator, Mapping\n\nfrom flax.core import FrozenDict\nfrom flax.errors import CursorFindError, TraverseTreeError\n\nA = TypeVar('A')\nKey = Any\n\n\n@runtime_checkable\nclass Indexable(Protocol):\n def __getitem__(self, key) -> Any:\n ...\n\n\nclass AccessType(enum.Enum):\n ITEM = enum.auto()\n ATTR = enum.auto()\n\n\n@dataclasses.dataclass\nclass ParentKey(Generic[A]):\n parent: 'Cursor[A]'\n key: Key\n access_type: AccessType\n\n\ndef is_named_tuple(obj):\n return (\n isinstance(obj, tuple)\n and hasattr(obj, '_fields')\n and hasattr(obj, '_asdict')\n and hasattr(obj, '_replace')\n )\n\n\ndef _traverse_tree(path, obj, *, update_fn=None, cond_fn=None):\n \"\"\"Helper function for ``Cursor.apply_update`` and ``Cursor.find_all``.\n Exactly one of ``update_fn`` and ``cond_fn`` must be not None.\n\n - If ``update_fn`` is not None, then ``Cursor.apply_update`` is calling\n this function and ``_traverse_tree`` will return a generator where\n each generated element is of type Tuple[Tuple[Union[str, int], AccessType], Any].\n The first element is a tuple of the key path and access type where the\n change was applied from the ``update_fn``, and the second element is\n the newly modified value. If the generator is non-empty, then the\n tuple key path will always be non-empty as well.\n - If ``cond_fn`` is not None, then ``Cursor.find_all`` is calling this\n function and ``_traverse_tree`` will return a generator where each\n generated element is of type Tuple[Union[str, int], AccessType]. The\n tuple contains the key path and access type where the object was found\n that fulfilled the conditions of the ``cond_fn``.\n \"\"\"\n if not (bool(update_fn) ^ bool(cond_fn)):\n raise TraverseTreeError(update_fn, cond_fn)\n\n if path:\n str_path = '/'.join(str(key) for key, _ in path)\n if update_fn:\n new_obj = update_fn(str_path, obj)\n if new_obj is not obj:\n yield path, new_obj\n return\n elif cond_fn(str_path, obj): # type: ignore\n yield path\n return\n\n if isinstance(obj, (FrozenDict, dict)):\n items = obj.items()\n access_type = AccessType.ITEM\n elif is_named_tuple(obj):\n items = ((name, getattr(obj, name)) for name in obj._fields) # type: ignore\n access_type = AccessType.ATTR\n elif isinstance(obj, (list, tuple)):\n items = enumerate(obj)\n access_type = AccessType.ITEM\n elif dataclasses.is_dataclass(obj):\n items = (\n (f.name, getattr(obj, f.name)) for f in dataclasses.fields(obj) if f.init\n )\n access_type = AccessType.ATTR\n else:\n return\n\n if update_fn:\n for key, value in items:\n yield from _traverse_tree(\n path + ((key, access_type),), value, update_fn=update_fn\n )\n else:\n for key, value in items:\n yield from _traverse_tree(\n path + ((key, access_type),), value, cond_fn=cond_fn\n )\n\n\nclass Cursor(Generic[A]):\n _obj: A\n _parent_key: ParentKey[A] | None\n _changes: dict[Any, 'Cursor[A]']\n\n def __init__(self, obj: A, parent_key: ParentKey[A] | None):\n # NOTE: we use `vars` here to avoid calling `__setattr__`\n # vars(self) = self.__dict__\n vars(self)['_obj'] = obj\n vars(self)['_parent_key'] = parent_key\n vars(self)['_changes'] = {}\n\n @property\n def _root(self) -> 'Cursor[A]':\n if self._parent_key is None:\n return self\n else:\n return self._parent_key.parent._root # type: ignore\n\n @property\n def _path(self) -> str:\n if self._parent_key is None:\n return ''\n if self._parent_key.access_type == AccessType.ITEM: # type: ignore\n if isinstance(self._parent_key.key, str): # type: ignore\n key = \"'\" + self._parent_key.key + \"'\" # type: ignore\n else:\n key = str(self._parent_key.key) # type: ignore\n return self._parent_key.parent._path + '[' + key + ']' # type: ignore\n # self.parent_key.access_type == AccessType.ATTR:\n return self._parent_key.parent._path + '.' + self._parent_key.key # type: ignore\n\n def __getitem__(self, key) -> 'Cursor[A]':\n if key in self._changes:\n return self._changes[key]\n\n if not isinstance(self._obj, Indexable):\n raise TypeError(f'Cannot index into {self._obj}')\n\n if isinstance(self._obj, Mapping) and key not in self._obj:\n raise KeyError(f'Key {key} not found in {self._obj}')\n\n if is_named_tuple(self._obj):\n return getattr(self, self._obj._fields[key]) # type: ignore\n\n child = Cursor(self._obj[key], ParentKey(self, key, AccessType.ITEM))\n self._changes[key] = child\n return child\n\n def __getattr__(self, name) -> 'Cursor[A]':\n if name in self._changes:\n return self._changes[name]\n\n if not hasattr(self._obj, name):\n raise AttributeError(f'Attribute {name} not found in {self._obj}')\n\n child = Cursor(\n getattr(self._obj, name), ParentKey(self, name, AccessType.ATTR)\n )\n self._changes[name] = child\n return child\n\n def __setitem__(self, key, value):\n if is_named_tuple(self._obj):\n return setattr(self, self._obj._fields[key], value) # type: ignore\n self._changes[key] = Cursor(value, ParentKey(self, key, AccessType.ITEM))\n\n def __setattr__(self, name, value):\n self._changes[name] = Cursor(value, ParentKey(self, name, AccessType.ATTR))\n\n def set(self, value) -> A:\n \"\"\"Set a new value for an attribute, property, element or entry\n in the Cursor object and return a copy of the original object,\n containing the new set value.\n\n Example::\n\n >>> from flax.cursor import cursor\n >>> from flax.training import train_state\n >>> import optax\n\n >>> dict_obj = {'a': 1, 'b': (2, 3), 'c': [4, 5]}\n >>> modified_dict_obj = cursor(dict_obj)['b'][0].set(10)\n >>> assert modified_dict_obj == {'a': 1, 'b': (10, 3), 'c': [4, 5]}\n\n >>> state = train_state.TrainState.create(\n ... apply_fn=lambda x: x,\n ... params=dict_obj,\n ... tx=optax.adam(1e-3),\n ... )\n >>> modified_state = cursor(state).params['b'][1].set(10)\n >>> assert modified_state.params == {'a': 1, 'b': (2, 10), 'c': [4, 5]}\n\n Args:\n value: the value used to set an attribute, property, element or entry in the Cursor object\n Returns:\n A copy of the original object with the new set value.\n \"\"\"\n if self._parent_key is None:\n return value\n parent, key = self._parent_key.parent, self._parent_key.key # type: ignore\n parent._changes[key] = Cursor(value, self._parent_key)\n return parent._root.build()\n\n def build(self) -> A:\n \"\"\"Create and return a copy of the original object with accumulated changes.\n This method is to be called after making changes to the Cursor object.\n\n .. note::\n The new object is built bottom-up, the changes will be first applied\n to the leaf nodes, and then its parent, all the way up to the root.\n\n Example::\n\n >>> from flax.cursor import cursor\n >>> from flax.training import train_state\n >>> import optax\n\n >>> dict_obj = {'a': 1, 'b': (2, 3), 'c': [4, 5]}\n >>> c = cursor(dict_obj)\n >>> c['b'][0] = 10\n >>> c['a'] = (100, 200)\n >>> modified_dict_obj = c.build()\n >>> assert modified_dict_obj == {'a': (100, 200), 'b': (10, 3), 'c': [4, 5]}\n\n >>> state = train_state.TrainState.create(\n ... apply_fn=lambda x: x,\n ... params=dict_obj,\n ... tx=optax.adam(1e-3),\n ... )\n >>> new_fn = lambda x: x + 1\n >>> c = cursor(state)\n >>> c.params['b'][1] = 10\n >>> c.apply_fn = new_fn\n >>> modified_state = c.build()\n >>> assert modified_state.params == {'a': 1, 'b': (2, 10), 'c': [4, 5]}\n >>> assert modified_state.apply_fn == new_fn\n\n Returns:\n A copy of the original object with the accumulated changes.\n \"\"\"\n changes = {\n key: child.build() if isinstance(child, Cursor) else child\n for key, child in self._changes.items()\n }\n if isinstance(self._obj, FrozenDict):\n obj = self._obj.copy(changes) # type: ignore\n elif isinstance(self._obj, (dict, list)):\n obj = self._obj.copy() # type: ignore\n for key, value in changes.items():\n obj[key] = value\n elif is_named_tuple(self._obj):\n obj = self._obj._replace(**changes) # type: ignore\n elif isinstance(self._obj, tuple):\n obj = list(self._obj) # type: ignore\n for key, value in changes.items():\n obj[key] = value\n obj = tuple(obj) # type: ignore\n elif dataclasses.is_dataclass(self._obj):\n obj = dataclasses.replace(self._obj, **changes) # type: ignore\n else:\n obj = self._obj # type: ignore\n return obj # type: ignore\n\n def apply_update(\n self,\n update_fn: Callable[[str, Any], Any],\n ) -> 'Cursor[A]':\n \"\"\"Traverse the Cursor object and record conditional changes recursively via an ``update_fn``.\n The changes are recorded in the Cursor object's ``._changes`` dictionary. To generate a copy\n of the original object with the accumulated changes, call the ``.build`` method after calling\n ``.apply_update``.\n\n The ``update_fn`` has a function signature of ``(str, Any) -> Any``:\n\n - The input arguments are the current key path (in the form of a string delimited\n by ``'/'``) and value at that current key path\n - The output is the new value (either modified by the ``update_fn`` or same as the\n input value if the condition wasn't fulfilled)\n\n .. note::\n - If the ``update_fn`` returns a modified value, this method will not recurse any further\n down that branch to record changes. For example, if we intend to replace an attribute that points\n to a dictionary with an int, we don't need to look for further changes inside the dictionary,\n since the dictionary will be replaced anyways.\n - The ``is`` operator is used to determine whether the return value is modified (by comparing it\n to the input value). Therefore if the ``update_fn`` modifies a mutable container (e.g. lists,\n dicts, etc.) and returns the same container, ``.apply_update`` will treat the returned value as\n unmodified as it contains the same ``id``. To avoid this, return a copy of the modified value.\n - ``.apply_update`` WILL NOT call the ``update_fn`` to the value at the top-most level of\n the pytree (i.e. the root node). The ``update_fn`` will first be called on the root node's\n children, and then the pytree traversal will continue recursively from there.\n\n Example::\n\n >>> import flax.linen as nn\n >>> from flax.cursor import cursor\n >>> import jax, jax.numpy as jnp\n\n >>> class Model(nn.Module):\n ... @nn.compact\n ... def __call__(self, x):\n ... x = nn.Dense(3)(x)\n ... x = nn.relu(x)\n ... x = nn.Dense(3)(x)\n ... x = nn.relu(x)\n ... x = nn.Dense(3)(x)\n ... x = nn.relu(x)\n ... return x\n\n >>> params = Model().init(jax.random.key(0), jnp.empty((1, 2)))['params']\n\n >>> def update_fn(path, value):\n ... '''Multiply all dense kernel params by 2 and add 1.\n ... Subtract the Dense_1 bias param by 1.'''\n ... if 'kernel' in path:\n ... return value * 2 + 1\n ... elif 'Dense_1' in path and 'bias' in path:\n ... return value - 1\n ... return value\n\n >>> c = cursor(params)\n >>> new_params = c.apply_update(update_fn).build()\n >>> for layer in ('Dense_0', 'Dense_1', 'Dense_2'):\n ... assert (new_params[layer]['kernel'] == 2 * params[layer]['kernel'] + 1).all()\n ... if layer == 'Dense_1':\n ... assert (new_params[layer]['bias'] == params[layer]['bias'] - 1).all()\n ... else:\n ... assert (new_params[layer]['bias'] == params[layer]['bias']).all()\n\n >>> assert jax.tree_util.tree_all(\n ... jax.tree_util.tree_map(\n ... lambda x, y: (x == y).all(),\n ... params,\n ... Model().init(jax.random.key(0), jnp.empty((1, 2)))[\n ... 'params'\n ... ],\n ... )\n ... ) # make sure original params are unchanged\n\n Args:\n update_fn: the function that will conditionally record changes to the Cursor object\n Returns:\n The current Cursor object with the recorded conditional changes specified by the\n ``update_fn``. To generate a copy of the original object with the accumulated\n changes, call the ``.build`` method after calling ``.apply_update``.\n \"\"\"\n for path, value in _traverse_tree((), self._obj, update_fn=update_fn):\n child = self\n for key, access_type in path[:-1]:\n if access_type is AccessType.ITEM:\n child = child[key]\n else: # access_type is AccessType.ATTR\n child = getattr(child, key)\n key, access_type = path[-1]\n if access_type is AccessType.ITEM:\n child[key] = value\n else: # access_type is AccessType.ATTR\n setattr(child, key, value)\n\n return self\n\n def find(self, cond_fn: Callable[[str, Any], bool]) -> 'Cursor[A]':\n \"\"\"Traverse the Cursor object and return a child Cursor object that fulfill the\n conditions in the ``cond_fn``. The ``cond_fn`` has a function signature of ``(str, Any) -> bool``:\n\n - The input arguments are the current key path (in the form of a string delimited\n by ``'/'``) and value at that current key path\n - The output is a boolean, denoting whether to return the child Cursor object at this path\n\n Raises a :meth:`CursorFindError ` if no object or more\n than one object is found that fulfills the condition of the ``cond_fn``. We raise an\n error because the user should always expect this method to return the only object whose\n corresponding key path and value fulfill the condition of the ``cond_fn``.\n\n .. note::\n - If the ``cond_fn`` evaluates to True at a particular key path, this method will not recurse\n any further down that branch; i.e. this method will find and return the \"earliest\" child node\n that fulfills the condition in ``cond_fn`` in a particular key path\n - ``.find`` WILL NOT search the the value at the top-most level of the pytree (i.e. the root\n node). The ``cond_fn`` will be evaluated recursively, starting at the root node's children.\n\n Example::\n\n >>> import flax.linen as nn\n >>> from flax.cursor import cursor\n >>> import jax, jax.numpy as jnp\n\n >>> class Model(nn.Module):\n ... @nn.compact\n ... def __call__(self, x):\n ... x = nn.Dense(3)(x)\n ... x = nn.relu(x)\n ... x = nn.Dense(3)(x)\n ... x = nn.relu(x)\n ... x = nn.Dense(3)(x)\n ... x = nn.relu(x)\n ... return x\n\n >>> params = Model().init(jax.random.PRNGKey(0), jnp.empty((1, 2)))['params']\n\n >>> def cond_fn(path, value):\n ... '''Find the second dense layer params.'''\n ... return 'Dense_1' in path\n\n >>> new_params = cursor(params).find(cond_fn)['bias'].set(params['Dense_1']['bias'] + 1)\n\n >>> for layer in ('Dense_0', 'Dense_1', 'Dense_2'):\n ... if layer == 'Dense_1':\n ... assert (new_params[layer]['bias'] == params[layer]['bias'] + 1).all()\n ... else:\n ... assert (new_params[layer]['bias'] == params[layer]['bias']).all()\n\n >>> c = cursor(params)\n >>> c2 = c.find(cond_fn)\n >>> c2['kernel'] += 2\n >>> c2['bias'] += 2\n >>> new_params = c.build()\n\n >>> for layer in ('Dense_0', 'Dense_1', 'Dense_2'):\n ... if layer == 'Dense_1':\n ... assert (new_params[layer]['kernel'] == params[layer]['kernel'] + 2).all()\n ... assert (new_params[layer]['bias'] == params[layer]['bias'] + 2).all()\n ... else:\n ... assert (new_params[layer]['kernel'] == params[layer]['kernel']).all()\n ... assert (new_params[layer]['bias'] == params[layer]['bias']).all()\n\n >>> assert jax.tree_util.tree_all(\n ... jax.tree_util.tree_map(\n ... lambda x, y: (x == y).all(),\n ... params,\n ... Model().init(jax.random.PRNGKey(0), jnp.empty((1, 2)))[\n ... 'params'\n ... ],\n ... )\n ... ) # make sure original params are unchanged\n\n Args:\n cond_fn: the function that will conditionally find child Cursor objects\n Returns:\n A child Cursor object that fulfills the condition in the ``cond_fn``.\n \"\"\"\n generator = self.find_all(cond_fn)\n try:\n cursor = next(generator)\n except StopIteration:\n raise CursorFindError()\n try:\n cursor2 = next(generator)\n raise CursorFindError(cursor, cursor2)\n except StopIteration:\n return cursor\n\n def find_all(\n self, cond_fn: Callable[[str, Any], bool]\n ) -> Generator['Cursor[A]', None, None]:\n \"\"\"Traverse the Cursor object and return a generator of child Cursor objects that fulfill the\n conditions in the ``cond_fn``. The ``cond_fn`` has a function signature of ``(str, Any) -> bool``:\n\n - The input arguments are the current key path (in the form of a string delimited\n by ``'/'``) and value at that current key path\n - The output is a boolean, denoting whether to return the child Cursor object at this path\n\n .. note::\n - If the ``cond_fn`` evaluates to True at a particular key path, this method will not recurse\n any further down that branch; i.e. this method will find and return the \"earliest\" child nodes\n that fulfill the condition in ``cond_fn`` in a particular key path\n - ``.find_all`` WILL NOT search the the value at the top-most level of the pytree (i.e. the root\n node). The ``cond_fn`` will be evaluated recursively, starting at the root node's children.\n\n Example::\n\n >>> import flax.linen as nn\n >>> from flax.cursor import cursor\n >>> import jax, jax.numpy as jnp\n\n >>> class Model(nn.Module):\n ... @nn.compact\n ... def __call__(self, x):\n ... x = nn.Dense(3)(x)\n ... x = nn.relu(x)\n ... x = nn.Dense(3)(x)\n ... x = nn.relu(x)\n ... x = nn.Dense(3)(x)\n ... x = nn.relu(x)\n ... return x\n\n >>> params = Model().init(jax.random.PRNGKey(0), jnp.empty((1, 2)))['params']\n\n >>> def cond_fn(path, value):\n ... '''Find all dense layer params.'''\n ... return 'Dense' in path\n\n >>> c = cursor(params)\n >>> for dense_params in c.find_all(cond_fn):\n ... dense_params['bias'] += 1\n >>> new_params = c.build()\n\n >>> for layer in ('Dense_0', 'Dense_1', 'Dense_2'):\n ... assert (new_params[layer]['bias'] == params[layer]['bias'] + 1).all()\n\n >>> assert jax.tree_util.tree_all(\n ... jax.tree_util.tree_map(\n ... lambda x, y: (x == y).all(),\n ... params,\n ... Model().init(jax.random.PRNGKey(0), jnp.empty((1, 2)))[\n ... 'params'\n ... ],\n ... )\n ... ) # make sure original params are unchanged\n\n Args:\n cond_fn: the function that will conditionally find child Cursor objects\n Returns:\n A generator of child Cursor objects that fulfill the condition in the ``cond_fn``.\n \"\"\"\n for path in _traverse_tree((), self._obj, cond_fn=cond_fn):\n child = self\n for key, access_type in path:\n if access_type is AccessType.ITEM:\n child = child[key]\n else: # access_type is AccessType.ATTR\n child = getattr(child, key)\n yield child\n\n def __str__(self):\n return str(self._obj)\n\n def __repr__(self):\n return self._pretty_repr()\n\n def _pretty_repr(self, indent=2, _prefix_indent=0):\n s = 'Cursor(\\n'\n obj_str = repr(self._obj).replace(\n '\\n', '\\n' + ' ' * (_prefix_indent + indent)\n )\n s += ' ' * (_prefix_indent + indent) + f'_obj={obj_str},\\n'\n s += ' ' * (_prefix_indent + indent) + '_changes={'\n if self._changes:\n s += '\\n'\n for key in self._changes:\n str_key = repr(key)\n prefix = ' ' * (_prefix_indent + 2 * indent) + str_key + ': '\n s += (\n prefix\n + self._changes[key]._pretty_repr(\n indent=indent, _prefix_indent=len(prefix)\n )\n + ',\\n'\n )\n s = s[\n :-2\n ] # remove comma and newline character for last element in self._changes\n s += '\\n' + ' ' * (_prefix_indent + indent) + '}\\n'\n else:\n s += '}\\n'\n s += ' ' * _prefix_indent + ')'\n return s\n\n def __len__(self):\n return len(self._obj)\n\n def __iter__(self):\n if isinstance(self._obj, (tuple, list)):\n return (self[i] for i in range(len(self._obj)))\n else:\n raise NotImplementedError(\n '__iter__ method only implemented for tuples and lists, not type'\n f' {type(self._obj)}'\n )\n\n def __reversed__(self):\n if isinstance(self._obj, (tuple, list)):\n return (self[i] for i in range(len(self._obj) - 1, -1, -1))\n else:\n raise NotImplementedError(\n '__reversed__ method only implemented for tuples and lists, not type'\n f' {type(self._obj)}'\n )\n\n def __add__(self, other):\n return self._obj + other\n\n def __sub__(self, other):\n return self._obj - other\n\n def __mul__(self, other):\n return self._obj * other\n\n def __matmul__(self, other):\n return self._obj @ other\n\n def __truediv__(self, other):\n return self._obj / other\n\n def __floordiv__(self, other):\n return self._obj // other\n\n def __mod__(self, other):\n return self._obj % other\n\n def __divmod__(self, other):\n return divmod(self._obj, other)\n\n def __pow__(self, other):\n return pow(self._obj, other)\n\n def __lshift__(self, other):\n return self._obj << other\n\n def __rshift__(self, other):\n return self._obj >> other\n\n def __and__(self, other):\n return self._obj & other\n\n def __xor__(self, other):\n return self._obj ^ other\n\n def __or__(self, other):\n return self._obj | other\n\n def __radd__(self, other):\n return other + self._obj\n\n def __rsub__(self, other):\n return other - self._obj\n\n def __rmul__(self, other):\n return other * self._obj\n\n def __rmatmul__(self, other):\n return other @ self._obj\n\n def __rtruediv__(self, other):\n return other / self._obj\n\n def __rfloordiv__(self, other):\n return other // self._obj\n\n def __rmod__(self, other):\n return other % self._obj\n\n def __rdivmod__(self, other):\n return divmod(other, self._obj)\n\n def __rpow__(self, other):\n return pow(other, self._obj)\n\n def __rlshift__(self, other):\n return other << self._obj\n\n def __rrshift__(self, other):\n return other >> self._obj\n\n def __rand__(self, other):\n return other & self._obj\n\n def __rxor__(self, other):\n return other ^ self._obj\n\n def __ror__(self, other):\n return other | self._obj\n\n def __neg__(self):\n return -self._obj\n\n def __pos__(self):\n return +self._obj\n\n def __abs__(self):\n return abs(self._obj)\n\n def __invert__(self):\n return ~self._obj\n\n def __round__(self, ndigits=None):\n return round(self._obj, ndigits)\n\n def __lt__(self, other):\n return self._obj < other\n\n def __le__(self, other):\n return self._obj <= other\n\n def __eq__(self, other):\n return self._obj == other\n\n def __ne__(self, other):\n return self._obj != other\n\n def __gt__(self, other):\n return self._obj > other\n\n def __ge__(self, other):\n return self._obj >= other\n\n\ndef cursor(obj: A) -> Cursor[A]:\n \"\"\"Wrap :class:`Cursor ` over ``obj`` and return it.\n Changes can then be applied to the Cursor object in the following ways:\n\n - single-line change via the ``.set`` method\n - multiple changes, and then calling the ``.build`` method\n - multiple changes conditioned on the pytree path and node value via the\n ``.apply_update`` method, and then calling the ``.build`` method\n\n ``.set`` example::\n\n >>> from flax.cursor import cursor\n\n >>> dict_obj = {'a': 1, 'b': (2, 3), 'c': [4, 5]}\n >>> modified_dict_obj = cursor(dict_obj)['b'][0].set(10)\n >>> assert modified_dict_obj == {'a': 1, 'b': (10, 3), 'c': [4, 5]}\n\n ``.build`` example::\n\n >>> from flax.cursor import cursor\n\n >>> dict_obj = {'a': 1, 'b': (2, 3), 'c': [4, 5]}\n >>> c = cursor(dict_obj)\n >>> c['b'][0] = 10\n >>> c['a'] = (100, 200)\n >>> modified_dict_obj = c.build()\n >>> assert modified_dict_obj == {'a': (100, 200), 'b': (10, 3), 'c': [4, 5]}\n\n ``.apply_update`` example::\n\n >>> from flax.cursor import cursor\n >>> from flax.training import train_state\n >>> import optax\n\n >>> def update_fn(path, value):\n ... '''Replace params with empty dictionary.'''\n ... if 'params' in path:\n ... return {}\n ... return value\n\n >>> state = train_state.TrainState.create(\n ... apply_fn=lambda x: x,\n ... params={'a': 1, 'b': 2},\n ... tx=optax.adam(1e-3),\n ... )\n >>> c = cursor(state)\n >>> state2 = c.apply_update(update_fn).build()\n >>> assert state2.params == {}\n >>> assert state.params == {'a': 1, 'b': 2} # make sure original params are unchanged\n\n If the underlying ``obj`` is a ``list`` or ``tuple``, iterating over the Cursor object\n to get the child Cursors is also possible::\n\n >>> from flax.cursor import cursor\n\n >>> c = cursor(((1, 2), (3, 4)))\n >>> for child_c in c:\n ... child_c[1] *= -1\n >>> assert c.build() == ((1, -2), (3, -4))\n\n View the docstrings for each method to see more examples of their usage.\n\n Args:\n obj: the object you want to wrap the Cursor in\n Returns:\n A Cursor object wrapped around obj.\n \"\"\"\n return Cursor(obj, None)\n" }, { "path": "flax/errors.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\"\"\" # Use an empty top-level docstring so Sphinx won't output the one below.\n\"\"\"Flax error classes.\n\n=== When to create a Flax error class?\n\nIf an error message requires more explanation than a one-liner, it is useful to\nadd it as a separate error class. This may lead to some duplication with\nexisting documentation or docstrings, but it will provide users with more help\nwhen they are debugging a problem. We can also point to existing documentation\nfrom the error docstring directly.\n\n=== How to name the error class?\n\n* If the error occurs when doing something, name the error\n Error\n\n For instance, if you want to raise an error when applying a module with an\n invalid method, the error can be: ApplyModuleInvalidMethodError.\n\n is optional, for instance if there is only one error when\n modifying a variable, the error can simply be: ModifyVariableError.\n\n* If there is no concrete action involved the only a description of the error is\n sufficient. For instance: InvalidFilterError, NameInUseError, etc.\n\n\n=== Copy/pastable template for new error messages:\n\nclass Template(FlaxError):\n \"\" \"\n\n \"\" \"\n def __init__(self):\n super().__init__(f'')\n\"\"\"\n\n\nclass FlaxError(Exception):\n def __init__(self, message):\n error_page = (\n 'https://flax.readthedocs.io/en/latest/api_reference/flax.errors.html'\n )\n module_name = self.__class__.__module__\n class_name = self.__class__.__name__\n if error_page not in message: # do not add a FlaxError link on unpickling\n message = f'{message} ({error_page}#{module_name}.{class_name})'\n super().__init__(message)\n\n def __reduce__(self):\n return (FlaxError, (str(self),))\n\n\n#################################################\n# NNX errors #\n#################################################\n\n\nclass TraceContextError(FlaxError):\n pass\n\n\n#################################################\n# lazy_init.py errors #\n#################################################\n\n\nclass LazyInitError(FlaxError):\n \"\"\"Lazy Init function has uncomputable return values.\n\n This happens when passing an argument to lazy_init with ``jax.ShapeDtypeStruct``\n that affects the initialized variables.\n Make sure the init function only uses the shape and dtype or pass an\n actual JAX array if this is impossible.\n\n Example::\n\n class Foo(nn.Module):\n @compact\n def __call__(self, x):\n # This parameter depends on the input x\n # this causes an error when using lazy_init.\n k = self.param(\"kernel\", lambda _: x)\n return x * k\n Foo().lazy_init(random.key(0), jax.ShapeDtypeStruct((8, 4), jnp.float32))\n \"\"\"\n\n def __init__(self, partial_val):\n super().__init__(\n 'Lazy init encountered a value that could with '\n f'the given inputs (shape: {partial_val}).'\n )\n\n\n#################################################\n# scope.py errors #\n#################################################\n\n\nclass InvalidRngError(FlaxError):\n \"\"\"All rngs used in a Module should be passed to\n\n :meth:`Module.init() ` and\n :meth:`Module.apply() ` appropriately. We explain\n both separately using the following example::\n\n class Bar(nn.Module):\n @nn.compact\n def __call__(self, x):\n some_param = self.param('some_param', nn.initializers.zeros_init(), (1, ))\n dropout_rng = self.make_rng('dropout')\n x = nn.Dense(features=4)(x)\n ...\n\n class Foo(nn.Module):\n @nn.compact\n def __call__(self, x):\n x = Bar()(x)\n ...\n\n **PRNGs for Module.init()**\n\n In this example, two rngs are used:\n\n * ``params`` is used for initializing the parameters of the model. This rng\n is used to initialize the ``some_params`` parameter, and for initializing\n the weights of the ``Dense`` Module used in ``Bar``.\n\n * ``dropout`` is used for the dropout rng that is used in ``Bar``.\n\n So, ``Foo`` is initialized as follows::\n\n init_rngs = {'params': random.key(0), 'dropout': random.key(1)}\n variables = Foo().init(init_rngs, init_inputs)\n\n If a Module only requires an rng for ``params``, you can use::\n\n SomeModule().init(rng, ...) # Shorthand for {'params': rng}\n\n\n **PRNGs for Module.apply()**\n\n When applying ``Foo``, only the rng for ``dropout`` is needed, because\n ``params`` is only used for initializing the Module parameters::\n\n Foo().apply(variables, inputs, rngs={'dropout': random.key(2)})\n\n If a Module only requires an rng for ``params``, you don't have to provide\n rngs for apply at all::\n\n SomeModule().apply(variables, inputs) # rngs=None\n \"\"\"\n\n def __init__(self, msg):\n # For this error message we pass the entire message, since there are various\n # different kinds of RNG errors and we want to be able to be more specific\n # in the error message, while always linking to the same documentation.\n super().__init__(msg)\n\n\nclass ApplyScopeInvalidVariablesTypeError(FlaxError):\n \"\"\"When calling :meth:`Module.apply() `, the first\n\n argument should be a variable dict. For more explanation on variable dicts,\n please see :mod:`flax.core.variables`.\n \"\"\"\n\n def __init__(self):\n super().__init__(\n 'The first argument passed to an apply function should be '\n 'a dictionary of collections. Each collection should be a '\n 'dictionary with string keys.'\n )\n\n\nclass ApplyScopeInvalidVariablesStructureError(FlaxError):\n \"\"\"This error is thrown when the dict passed as ``variables`` to apply() has an\n\n extra 'params' layer, i.e. {'params': {'params': ...}}.\n For more explanation on variable dicts, please see :mod:`flax.core.variables`.\n \"\"\"\n\n def __init__(self, variables):\n super().__init__(\n f'Expect the `variables` (first argument) passed to apply() '\n f'to be a dict with the structure {{\"params\": ...}}, but got a dict '\n f'with an extra params layer, i.e. {{\"params\": {{\"params\": ... }} }}. '\n f'You should instead pass in your dict\\'s [\"params\"].'\n )\n\n\nclass ScopeParamNotFoundError(FlaxError):\n \"\"\"This error is thrown when trying to access a parameter that does not exist.\n\n For instance, in the code below, the initialized embedding name 'embedding'\n does not match the apply name 'embed'::\n\n class Embed(nn.Module):\n num_embeddings: int\n features: int\n\n @nn.compact\n def __call__(self, inputs, embed_name='embedding'):\n inputs = inputs.astype('int32')\n embedding = self.param(embed_name,\n jax.nn.initializers.lecun_normal(),\n (self.num_embeddings, self.features))\n return embedding[inputs]\n\n model = Embed(4, 8)\n variables = model.init(random.key(0), jnp.ones((5, 5, 1)))\n _ = model.apply(variables, jnp.ones((5, 5, 1)), 'embed')\n \"\"\"\n\n def __init__(self, param_name, scope_path):\n super().__init__(\n f'Could not find parameter named \"{param_name}\" in scope '\n f'\"{scope_path}\".'\n )\n\n\nclass ScopeCollectionNotFound(FlaxError):\n \"\"\"This error is thrown when trying to access a variable from an empty\n collection.\n\n There are two common causes:\n\n 1. | The collection was not passed to ``apply`` correctly.\n | For example, you might have used ``module.apply(params, ...)`` instead\n | of ``module.apply({'params': params}, ...)``.\n 2. | The collection is empty because the variables need to be initialized.\n | In this case, you should have made the collection mutable during\n | apply (e.g.: ``module.apply(variables, ..., mutable=['state'])``.\n \"\"\"\n\n def __init__(self, col_name, var_name, scope_path):\n super().__init__(\n f'Tried to access \"{var_name}\" from collection \"{col_name}\" in '\n f'\"{scope_path}\" but the collection is empty.'\n )\n\n\nclass ScopeParamShapeError(FlaxError):\n \"\"\"This error is thrown when the shape of an existing parameter is different from\n\n the shape of the return value of the ``init_fn``. This can happen when the\n shape provided during :meth:`Module.apply() ` is\n different from the one used when initializing the module.\n\n For instance, the following code throws this error because the apply shape\n (``(5, 5, 1)``) is different from the init shape (``(5, 5``). As a result, the\n shape of the kernel during ``init`` is ``(1, 8)``, and the shape during\n ``apply`` is ``(5, 8)``, which results in this error.::\n\n class NoBiasDense(nn.Module):\n features: int = 8\n\n @nn.compact\n def __call__(self, x):\n kernel = self.param('kernel',\n lecun_normal(),\n (x.shape[-1], self.features)) # <--- ERROR\n y = lax.dot_general(x, kernel,\n (((x.ndim - 1,), (0,)), ((), ())))\n return y\n\n variables = NoBiasDense().init(random.key(0), jnp.ones((5, 5, 1)))\n _ = NoBiasDense().apply(variables, jnp.ones((5, 5)))\n \"\"\"\n\n def __init__(self, param_name, scope_path, value_shape, init_shape):\n super().__init__(\n f'For parameter \"{param_name}\" in \"{scope_path}\", the given '\n f'initializer is expected to generate shape {init_shape}, but the '\n f'existing parameter it received has shape {value_shape}.'\n )\n\n\nclass ScopeVariableNotFoundError(FlaxError):\n \"\"\"This error is thrown when trying to use a variable in a Scope in a collection\n\n that is immutable. In order to create this variable, mark the collection as\n mutable explicitly using the ``mutable`` keyword in\n :meth:`Module.apply() `.\n \"\"\"\n\n def __init__(self, name, col, scope_path):\n super().__init__(\n f'No Variable named \"{name}\" for collection \"{col}\" '\n f'exists in \"{scope_path}\".'\n )\n\n\nclass InvalidFilterError(FlaxError):\n \"\"\"A filter should be either a boolean, a string or a container object.\"\"\"\n\n def __init__(self, filter_like):\n super().__init__(f'Invalid Filter: \"{filter_like}\"')\n\n\nclass InvalidScopeError(FlaxError):\n \"\"\"A temporary Scope is only valid within the context in which it is created::\n\n with Scope(variables, rngs=rngs).temporary() as root:\n y = fn(root, *args, **kwargs)\n # Here root is valid.\n # Here root is invalid.\n \"\"\"\n\n def __init__(self, scope_name):\n super().__init__(f'The scope \"{scope_name}\" is no longer valid.')\n\n\nclass ModifyScopeVariableError(FlaxError):\n \"\"\"You cannot update a variable if the collection it belongs to is immutable.\n\n When you are applying a Module, you should specify which variable collections\n are mutable::\n\n class MyModule(nn.Module):\n @nn.compact\n def __call__(self, x):\n ...\n var = self.variable('batch_stats', 'mean', ...)\n var.value = ...\n ...\n\n v = MyModule.init(...)\n ...\n logits = MyModule.apply(v, batch) # This throws an error.\n logits = MyModule.apply(v, batch, mutable=['batch_stats']) # This works.\n \"\"\"\n\n def __init__(self, col, variable_name, scope_path):\n super().__init__(\n f'Cannot update variable \"{variable_name}\" in '\n f'\"{scope_path}\" because collection \"{col}\" is immutable.'\n )\n\n\nclass ImmutableVariableError(FlaxError):\n \"\"\"You cannot update a variable that is marked as immutable.\n\n This error occurs when attempting to modify a Variable that has been set to\n 'immutable' mode. Variables in immutable mode are read-only and cannot be\n changed after creation.\n\n To fix this error, either:\n 1. Use a different variable mode (e.g., 'qdd' or 'pytree')\n 2. Or ensure you're not trying to modify the variable's value\n \"\"\"\n\n def __init__(self, message):\n super().__init__(message)\n\n\nclass JaxTransformError(FlaxError):\n \"\"\"JAX transforms and Flax modules cannot be mixed.\n\n JAX's functional transformations expect pure function.\n When you want to use JAX transformations **inside** Flax models,\n you should make use of the Flax transformation wrappers\n (e.g.: ``flax.linen.vmap``, ``flax.linen.scan``, etc.).\n \"\"\"\n\n def __init__(self):\n super().__init__('Jax transforms and Flax models cannot be mixed.')\n\n\n#################################################\n# meta.py errors #\n#################################################\n\n\nclass PartitioningUnspecifiedError(FlaxError):\n \"\"\"This error is raised when trying to add an axis to a Partitioned variable by\n\n using a transformation (e.g.: ``scan``, ``vmap``) without specifying the\n \"partition_name\" in the ``metadata_params`` dict.\n \"\"\"\n\n def __init__(self, target):\n super().__init__(\n 'Trying to transform a Partitioned variable but \"partition_name\"'\n f' is not specified in metadata_params: {target}'\n )\n\n\n#################################################\n# module.py errors #\n#################################################\n\n\nclass NameInUseError(FlaxError):\n \"\"\"This error is raised when trying to create a submodule, param, or variable\n\n with an existing name. They are all considered to be in the same namespace.\n\n **Sharing Submodules**\n\n This is the wrong pattern for sharing submodules::\n\n y = nn.Dense(feature=3, name='bar')(x)\n z = nn.Dense(feature=3, name='bar')(x+epsilon)\n\n Instead, modules should be shared by instance::\n\n dense = nn.Dense(feature=3, name='bar')\n y = dense(x)\n z = dense(x+epsilon)\n\n If submodules are not provided with a name, a unique name will be given to\n them automatically::\n\n class MyModule(nn.Module):\n @nn.compact\n def __call__(self, x):\n x = MySubModule()(x)\n x = MySubModule()(x) # This is fine.\n return x\n\n **Parameters and Variables**\n\n A parameter name can collide with a submodule or variable, since they are all\n stored in the same variable dict::\n\n class Foo(nn.Module):\n @nn.compact\n def __call__(self, x):\n bar = self.param('bar', nn.initializers.zeros_init(), (1, ))\n embed = nn.Embed(num_embeddings=2, features=5, name='bar') # <-- ERROR!\n\n Variables should also have unique names, even if they have their own\n collection::\n\n class Foo(nn.Module):\n @nn.compact\n def __call__(self, inputs):\n _ = self.param('mean', initializers.lecun_normal(), (2, 2))\n _ = self.variable('stats', 'mean', initializers.zeros_init(), (2, 2))\n \"\"\"\n\n def __init__(self, key_type, value, module_name):\n # key_type is in {param, variable, submodule}.\n super().__init__(\n f'Could not create {key_type} \"{value}\" in Module '\n f'{module_name}: Name in use.'\n )\n\n\nclass AssignSubModuleError(FlaxError):\n \"\"\"You are only allowed to create submodules in two places:\n\n 1. If your Module is noncompact: inside\n :meth:`Module.setup() `.\n 2. If your Module is compact: inside the method wrapped in\n :meth:`nn.compact() `.\n\n For instance, the following code throws this error, because ``nn.Conv`` is\n created in ``__call__``, which is not marked as compact::\n\n class Foo(nn.Module):\n def setup(self):\n pass\n\n def __call__(self, x):\n conv = nn.Conv(features=3, kernel_size=3)\n\n Foo().init(random.key(0), jnp.zeros((1,)))\n\n Note that this error is also thrown if you partially defined a Module inside\n setup::\n\n class Foo(nn.Module):\n def setup(self):\n self.conv = functools.partial(nn.Conv, features=3)\n\n def __call__(self, x):\n x = self.conv(kernel_size=4)(x)\n return x\n\n Foo().init(random.key(0), jnp.zeros((1,)))\n\n In this case, ``self.conv(kernel_size=4)`` is called from ``__call__``, which\n is disallowed because it's neither within ``setup`` nor a method wrapped in\n x``nn.compact``.\n \"\"\"\n\n def __init__(self, cls):\n super().__init__(\n f'Submodule {cls} must be defined in `setup()` or in a '\n 'method wrapped in `@compact`'\n )\n\n\nclass SetAttributeInModuleSetupError(FlaxError):\n \"\"\"You are not allowed to modify Module class attributes in\n\n :meth:`Module.setup() `::\n\n class Foo(nn.Module):\n features: int = 6\n\n def setup(self):\n self.features = 3 # <-- ERROR\n\n def __call__(self, x):\n return nn.Dense(self.features)(x)\n\n variables = SomeModule().init(random.key(0), jnp.ones((1, )))\n\n Instead, these attributes should be set when initializing the Module::\n\n class Foo(nn.Module):\n features: int = 6\n\n @nn.compact\n def __call__(self, x):\n return nn.Dense(self.features)(x)\n\n variables = SomeModule(features=3).init(random.key(0), jnp.ones((1, )))\n\n TODO(marcvanzee): Link to a design note explaining why it's necessary for\n modules to stay frozen (otherwise we can't safely clone them, which we use for\n lifted transformations).\n \"\"\"\n\n def __init__(self):\n super().__init__(f'Module construction attributes are frozen.')\n\n\nclass SetAttributeFrozenModuleError(FlaxError):\n \"\"\"You can only assign Module attributes to ``self`` inside\n\n :meth:`Module.setup() `. Outside of that method, the\n Module instance is frozen (i.e., immutable). This behavior is similar to\n frozen Python dataclasses.\n\n For instance, this error is raised in the following case::\n\n class SomeModule(nn.Module):\n @nn.compact\n def __call__(self, x, num_features=10):\n self.num_features = num_features # <-- ERROR!\n x = nn.Dense(self.num_features)(x)\n return x\n\n s = SomeModule().init(random.key(0), jnp.ones((5, 5)))\n\n Similarly, the error is raised when trying to modify a submodule's attributes\n after constructing it, even if this is done in the ``setup()`` method of the\n parent module::\n\n class Foo(nn.Module):\n def setup(self):\n self.dense = nn.Dense(features=10)\n self.dense.features = 20 # <--- This is not allowed\n\n def __call__(self, x):\n return self.dense(x)\n \"\"\"\n\n def __init__(self, module_cls, attr_name, attr_val):\n super().__init__(\n f\"Can't set {attr_name}={attr_val} for Module of type \"\n f'{module_cls}: Module instance is frozen outside of '\n 'setup method.'\n )\n\n\nclass MultipleMethodsCompactError(FlaxError):\n \"\"\"The ``@compact`` decorator may only be added to at most one method in a Flax\n\n module. In order to resolve this, you can:\n\n * remove ``@compact`` and define submodules and variables using\n :meth:`Module.setup() `.\n * Use two separate modules that both have a unique ``@compact`` method.\n\n TODO(marcvanzee): Link to a design note explaining the motivation behind this.\n There is no need for an equivalent to ``hk.transparent`` and it makes submodules\n much more sane because there is no need to prefix the method names.\n \"\"\"\n\n def __init__(self):\n super().__init__(f'Only one method per class can be @compact')\n\n\nclass ReservedModuleAttributeError(FlaxError):\n \"\"\"This error is thrown when creating a Module that is using reserved attributes.\n\n The following attributes are reserved:\n\n * ``parent``: The parent Module of this Module.\n * ``name``: The name of this Module.\n \"\"\"\n\n def __init__(self, annotations):\n super().__init__(\n f'properties `parent` and `name` are reserved: {annotations}'\n )\n\n\nclass ApplyModuleInvalidMethodError(FlaxError):\n \"\"\"When calling :meth:`Module.apply() `, you can specify\n\n the method to apply using parameter ``method``. This error is thrown if the\n provided parameter is not a method in the Module and not a function with at\n least one argument.\n\n Learn more on the reference docs for\n :meth:`Module.apply() `.\n \"\"\"\n\n def __init__(self, method):\n super().__init__(\n f'Cannot call apply(): {method} is not a valid function for apply().'\n )\n\n\nclass CallCompactUnboundModuleError(FlaxError):\n \"\"\"This error occurs when you are trying to call a Module directly, rather than\n\n through :meth:`Module.apply() `. For instance, the\n error will be raised when trying to run this code::\n\n from flax import linen as nn\n import jax.numpy as jnp\n\n test_dense = nn.Dense(10)\n test_dense(jnp.ones((5,5)))\n\n Instead, you should pass the variables (parameters and other state) via\n :meth:`Module.apply() ` (or use\n :meth:`Module.init() ` to get initial variables)::\n\n from jax import random\n variables = test_dense.init(random.key(0), jnp.ones((5,5)))\n\n y = test_dense.apply(variables, jnp.ones((5,5)))\n \"\"\"\n\n def __init__(self):\n super().__init__(\"Can't call compact methods on unbound modules\")\n\n\nclass CallSetupUnboundModuleError(FlaxError):\n \"\"\"This error occurs when you are trying to call ``.setup()`` directly.\n\n For instance, the error will be raised when trying to run this code::\n\n from flax import linen as nn\n import jax.numpy as jnp\n\n class MyModule(nn.Module):\n def setup(self):\n self.submodule = MySubModule()\n\n module = MyModule()\n module.setup() # <-- ERROR!\n submodule = module.submodule\n\n In general you shouldn't call ``.setup()`` yourself, if you need to get access\n to a field or submodule defined inside ``setup`` you can instead create a\n function\n to extract it and pass it to ``nn.apply``::\n\n # setup() will be called automatically by ``nn.apply``\n def get_submodule(module):\n return module.submodule.clone() # avoid leaking the Scope\n\n empty_variables = {} # you can also use the real variables\n submodule = nn.apply(get_submodule, module)(empty_variables)\n \"\"\"\n\n def __init__(self):\n super().__init__(\"Can't call compact methods on unbound modules\")\n\n\nclass CallUnbindOnUnboundModuleError(FlaxError):\n \"\"\"This error occurs when you are trying to call ``.unbind()`` on an unbound\n Module. For instance, when you try running the following example,\n an error will be raised::\n\n from flax import linen as nn\n\n class MyModule(nn.Module):\n @nn.compact\n def __call__(self, x):\n return nn.Dense(features=10)(x)\n\n module = MyModule()\n module.unbind() # <-- ERROR!\n\n Instead, you should ``bind`` the Module to a variable collection before calling\n ``.unbind()``::\n\n bound_module = module.bind(variables)\n ... # do something with bound_module\n module = bound_module.unbind() # <-- OK!\n \"\"\"\n\n def __init__(self):\n super().__init__(\"Can't call `unbind()` on unbound modules\")\n\nclass CallShareScopeOnUnboundModuleError(FlaxError):\n \"\"\"This error occurs when you are trying to call ``nn.share_scope`` on an unbound\n Module. For instance, when you try to use ``nn.share_scope`` at the top-level::\n\n from flax import linen as nn\n\n class CustomDense(nn.Dense):\n def __call__(self, x):\n return super().__call__(x) + 1\n\n custom_dense = CustomDense(5)\n dense = nn.Dense(5) # has the parameters\n\n nn.share_scope(custom_dense, dense) # <-- ERROR!\n \"\"\"\n\n def __init__(self):\n super().__init__(\"Can't call `share_scope` on unbound modules\")\n\nclass InvalidInstanceModuleError(FlaxError):\n \"\"\"This error occurs when you are trying to call ``.init()``, ``.init_with_output()``, ``.apply()`` or ``.bind()``\n\n on the Module class itself, instead of an instance of the Module class.\n For example, the error will be raised when trying to run this code::\n\n class B(nn.Module):\n @nn.compact\n def __call__(self, x):\n return x\n\n k = random.key(0)\n x = random.uniform(random.key(1), (2,))\n B.init(k, x) # B is module class, not B() a module instance\n B.apply(vs, x) # similar issue with apply called on class instead of\n instance.\n \"\"\"\n\n def __init__(self):\n super().__init__(\n 'Can only call init, init_with_output or apply methods on an instance'\n ' of the Module class, not the Module class itself'\n )\n\n\nclass IncorrectPostInitOverrideError(FlaxError):\n \"\"\"This error occurs when you overrode ``.__post_init__()`` without calling ``super().__post_init__()``.\n\n For example, the error will be raised when trying to run this code::\n\n from flax import linen as nn\n import jax.numpy as jnp\n import jax\n class A(nn.Module):\n x: float\n def __post_init__(self):\n self.x_square = self.x ** 2\n # super().__post_init__() <-- forgot to add this line\n @nn.compact\n def __call__(self, input):\n return input + 3\n\n r = A(x=3)\n r.init(jax.random.key(2), jnp.ones(3))\n \"\"\"\n\n def __init__(self):\n super().__init__(\n 'Overrode `.__post_init__()` without calling `super().__post_init__()`'\n )\n\n\nclass DescriptorAttributeError(FlaxError):\n \"\"\"This error occurs when you are trying to access a property that is accessing a non-existent attribute.\n\n For example, the error will be raised when trying to run this code::\n\n class Foo(nn.Module):\n @property\n def prop(self):\n return self.non_existent_field # ERROR!\n def __call__(self, x):\n return self.prop\n\n foo = Foo()\n variables = foo.init(jax.random.key(0), jnp.ones(shape=(1, 8)))\n \"\"\"\n\n def __init__(self):\n super().__init__(\n 'Trying to access a property that is accessing a non-existent'\n ' attribute.'\n )\n\n\nclass InvalidCheckpointError(FlaxError):\n \"\"\"A checkpoint cannot be stored in a directory that already has\n\n a checkpoint at the current or a later step.\n\n You can pass ``overwrite=True`` to disable this behavior and\n overwrite existing checkpoints in the target directory.\n \"\"\"\n\n def __init__(self, path, step):\n super().__init__(\n f'Trying to save an outdated checkpoint at step: \"{step}\" and path:'\n f' \"{path}\".'\n )\n\n\nclass MPACheckpointingRequiredError(FlaxError):\n \"\"\"To optimally save and restore a multiprocess array (GDA or jax Array outputted from pjit), use GlobalAsyncCheckpointManager.\n\n You can create an GlobalAsyncCheckpointManager at top-level and pass it as\n argument::\n\n from jax.experimental.gda_serialization import serialization as gdas\n gda_manager = gdas.GlobalAsyncCheckpointManager()\n save_checkpoint(..., gda_manager=gda_manager)\n \"\"\"\n\n def __init__(self, path, step):\n super().__init__(\n f'Checkpoint failed at step: \"{step}\" and path: \"{path}\": Target '\n 'contains a multiprocess array should be saved/restored with a '\n 'GlobalAsyncCheckpointManager.'\n )\n\n\nclass MPARestoreTargetRequiredError(FlaxError):\n \"\"\"Provide a valid target when restoring a checkpoint with a multiprocess array.\n\n Multiprocess arrays need a sharding (global meshes and partition specs) to be\n initialized. Therefore, to restore a checkpoint that contains a multiprocess\n array, make sure the ``target`` you passed contains valid multiprocess arrays\n at the corresponding tree structure location. If you cannot provide a full\n valid ``target``, consider ``allow_partial_mpa_restoration=True``.\n \"\"\"\n\n def __init__(self, path, step, key=None):\n error_msg = (\n f'Restore checkpoint failed at step: \"{step}\" and path: \"{path}\": '\n 'Checkpoints containing a multiprocess array need to be restored with '\n 'a target with pre-created arrays. If you cannot provide a full valid '\n 'target, consider ``allow_partial_mpa_restoration=True``. '\n )\n if key:\n error_msg += f'This error fired when trying to restore array at {key}.'\n super().__init__(error_msg)\n\n\nclass MPARestoreDataCorruptedError(FlaxError):\n \"\"\"A multiprocess array stored in Google Cloud Storage doesn't contain a \"commit_success.txt\" file, which should be written at the end of the save.\n\n Failure of finding it could indicate a corruption of your saved GDA data.\n \"\"\"\n\n def __init__(self, step, path):\n super().__init__(\n f'Restore checkpoint failed at step: \"{step}\" on multiprocess array at '\n f' \"{path}\": No \"commit_success.txt\" found on this \"_gda\" directory. '\n 'Was its save halted before completion?'\n )\n\n\n#################################################\n# transforms.py errors #\n#################################################\n\n\nclass TransformedMethodReturnValueError(FlaxError):\n \"\"\"Transformed Module methods cannot return other Modules or Variables.\"\"\"\n\n def __init__(self, name):\n super().__init__(\n f'Transformed module method {name} cannot return Modules or Variables.'\n )\n\n\nclass TransformTargetError(FlaxError):\n \"\"\"Linen transformations must be applied to Modules classes or functions taking a Module instance as the first argument.\n\n This error occurs when passing an invalid target to a linen transform\n (nn.vmap, nn.scan, etc.). This occurs for example when trying to\n transform a Module instance::\n\n nn.vmap(nn.Dense(features))(x) # raises TransformTargetError\n\n You can transform the ``nn.Dense`` class directly instead::\n\n nn.vmap(nn.Dense)(features)(x)\n\n Or you can create a function that takes the module instance as the first\n argument::\n\n class BatchDense(nn.Module):\n @nn.compact\n def __call__(self, x):\n return nn.vmap(\n lambda mdl, x: mdl(x),\n variable_axes={'params': 0}, split_rngs={'params':\n True})(nn.Dense(3), x)\n \"\"\"\n\n def __init__(self, target):\n super().__init__(\n 'Linen transformations must be applied to Modules classes or'\n ' functions taking a Module instance as the first argument.'\n f' The provided target is not a Module class or callable: {target}'\n )\n\n\n#################################################\n# io.py errors #\n#################################################\n\n\nclass AlreadyExistsError(FlaxError):\n \"\"\"Attempting to overwrite a file via copy.\n\n You can pass ``overwrite=True`` to disable this behavior and overwrite\n existing files in.\n \"\"\"\n\n def __init__(self, path):\n super().__init__(f'Trying overwrite an existing file: \"{path}\".')\n\n\n#################################################\n# cursor.py errors #\n#################################################\n\n\nclass CursorFindError(FlaxError):\n \"\"\"Error when calling :meth:`Cursor.find() `.\n\n This error occurs if no object or more than one object is found, given\n the conditions of the ``cond_fn``.\n \"\"\"\n\n def __init__(self, cursor=None, cursor2=None):\n if cursor and cursor2:\n super().__init__(\n 'More than one object found given the conditions of the cond_fn. '\n 'The first two objects found have the following paths: '\n f'{cursor._path} and {cursor2._path}'\n )\n else:\n super().__init__('No object found given the conditions of the cond_fn.')\n\n\nclass TraverseTreeError(FlaxError):\n \"\"\"Error when calling ``Cursor._traverse_tree()``. This function has two\n modes:\n\n - if ``update_fn`` is not None, it will traverse the tree and return a\n generator of tuples containing the path where the ``update_fn`` was\n applied and the newly modified value.\n - if ``cond_fn`` is not None, it will traverse the tree and return a\n generator of tuple paths that fulfilled the conditions of the ``cond_fn``.\n\n This error occurs if either both ``update_fn`` and ``cond_fn`` are None,\n or both are not None.\n \"\"\"\n\n def __init__(self, update_fn, cond_fn):\n if update_fn is None and cond_fn is None:\n super().__init__(\n 'Both update_fn and cond_fn are None. Exactly one of them must be'\n ' None.'\n )\n else:\n super().__init__(\n 'Both update_fn and cond_fn are not None. Exactly one of them must be'\n ' not None.'\n )\n" }, { "path": "flax/experimental/__init__.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n" }, { "path": "flax/experimental/nnx.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom absl import logging\n\nfrom flax.nnx import *\n\n\nlogging.warning(\n \"Using 'flax.experimental.nnx' is deprecated. Please use 'flax.nnx' instead.\"\n)" }, { "path": "flax/ids.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"UUIDs for Flax internals.\"\"\"\n\nimport threading\n\n\nclass UUIDManager:\n \"\"\"Globally unique counter-based id manager.\n\n We need globally unique key ids for Module and Variable object instances\n to preserve and recreate sharing-by-reference relationship when lifting\n transforms and adopting outside Modules.\n - Use of id() is unacceptable because these identifiers are literally\n pointers which can be recycled, so we rely on a globally unique counter id\n instead.\n - We need to handle copy/deepcopy uniqueness via a wrapped type.\n \"\"\"\n\n def __init__(self):\n self._lock = threading.Lock()\n self._id = 0\n\n def __call__(self):\n with self._lock:\n self._id += 1\n return FlaxId(self._id)\n\n\nuuid = UUIDManager()\n\n\nclass FlaxId:\n \"\"\"Hashable wrapper for ids that handles uniqueness of copies.\"\"\"\n\n def __init__(self, rawid):\n self.id = rawid\n\n def __eq__(self, other):\n return isinstance(other, FlaxId) and other.id == self.id\n\n def __hash__(self):\n return hash(self.id)\n\n def __repr__(self):\n return f'FlaxId({self.id})'\n\n def __deepcopy__(self, memo):\n del memo\n return uuid()\n\n def __copy__(self):\n return uuid()\n" }, { "path": "flax/io.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"IO Abstraction Layer.\nThe sole purpose of this abstraction layer is to avoid requiring tensorflow\nas an open-source dependency solely for its tensorflow.io.gfile functions.\n\"\"\"\nimport contextlib\nimport glob as glob_module\nimport importlib\nimport os\nimport shutil\nfrom enum import Enum\n\nfrom absl import logging\n\nfrom . import errors\n\n# Global Modes and selective import of tensorflow.io gfile.\n\n\nclass BackendMode(Enum):\n DEFAULT = 0\n TF = 1\n\n\nio_mode = None\ngfile = None\n\nif importlib.util.find_spec('tensorflow'):\n from tensorflow.io import gfile # type: ignore\n\n io_mode = BackendMode.TF\nelse:\n logging.warning(\n 'Tensorflow library not found, tensorflow.io.gfile '\n 'operations will use native shim calls. '\n \"GCS paths (i.e. 'gs://...') cannot be accessed.\"\n )\n io_mode = BackendMode.DEFAULT\n\n\n# Constants and Exceptions\n\n\nif io_mode == BackendMode.TF:\n from tensorflow import errors as tf_errors # type: ignore\n\n NotFoundError = tf_errors.NotFoundError\nelse:\n NotFoundError = FileNotFoundError\n\n\n# Overrides for testing.\n\n\n@contextlib.contextmanager\ndef override_mode(override: BackendMode):\n # pylint: disable=g-doc-return-or-yield\n \"\"\"Returns a context manager that changes backend IO mode.\n Args:\n override: BackendMode enum value to set IO mode inside context.\n \"\"\"\n # pylint: enable=g-doc-return-or-yield\n global io_mode\n io_mode_prev = io_mode\n io_mode = override\n try:\n yield\n finally:\n io_mode = io_mode_prev\n\n\ndef set_mode(override: BackendMode):\n \"\"\"Sets global io mode.\n Args:\n override: BackendMode enum value to set for IO mode.\n \"\"\"\n global io_mode\n io_mode = override\n\n\n# tensorflow.io.gfile API shim functions.\n\n\ndef GFile(name, mode): # pylint: disable=invalid-name\n if io_mode == BackendMode.DEFAULT:\n if 'b' in mode:\n return open(name, mode) # pylint: disable=unspecified-encoding\n else:\n return open(name, mode, encoding='utf-8')\n elif io_mode == BackendMode.TF:\n return gfile.GFile(name, mode)\n else:\n raise ValueError('Unknown IO Backend Mode.')\n\n\ndef listdir(path):\n if io_mode == BackendMode.DEFAULT:\n return os.listdir(path=path)\n elif io_mode == BackendMode.TF:\n return gfile.listdir(path=path)\n else:\n raise ValueError('Unknown IO Backend Mode.')\n\n\ndef isdir(path):\n if io_mode == BackendMode.DEFAULT:\n return os.path.isdir(path)\n elif io_mode == BackendMode.TF:\n return gfile.isdir(path)\n else:\n raise ValueError('Unknown IO Backend Mode.')\n\n\ndef copy(src, dst, overwrite=False):\n if io_mode == BackendMode.DEFAULT:\n if os.path.exists(dst) and not overwrite:\n raise errors.AlreadyExistsError(dst)\n shutil.copy(src, dst)\n return\n elif io_mode == BackendMode.TF:\n return gfile.copy(src, dst, overwrite=overwrite)\n else:\n raise ValueError('Unknown IO Backend Mode.')\n\n\ndef rename(src, dst, overwrite=False):\n if io_mode == BackendMode.DEFAULT:\n if os.path.exists(dst) and not overwrite:\n raise errors.AlreadyExistsError(dst)\n return os.rename(src, dst)\n elif io_mode == BackendMode.TF:\n return gfile.rename(src, dst, overwrite=overwrite)\n else:\n raise ValueError('Unknown IO Backend Mode.')\n\n\ndef exists(path):\n if io_mode == BackendMode.DEFAULT:\n return os.path.exists(path)\n elif io_mode == BackendMode.TF:\n return gfile.exists(path)\n else:\n raise ValueError('Unknown IO Backend Mode.')\n\n\ndef makedirs(path):\n if io_mode == BackendMode.DEFAULT:\n return os.makedirs(path, exist_ok=True)\n elif io_mode == BackendMode.TF:\n return gfile.makedirs(path)\n else:\n raise ValueError('Unknown IO Backend Mode.')\n\n\ndef glob(pattern):\n if io_mode == BackendMode.DEFAULT:\n return [\n path.rstrip('/') for path in glob_module.glob(pattern, recursive=False)\n ]\n elif io_mode == BackendMode.TF:\n return gfile.glob(pattern)\n else:\n raise ValueError('Unknown IO Backend Mode.')\n\n\ndef remove(path):\n \"\"\"Remove the file at path. Might fail if used on a directory path.\"\"\"\n if io_mode == BackendMode.DEFAULT:\n return os.remove(path)\n elif io_mode == BackendMode.TF:\n return gfile.remove(path)\n else:\n raise ValueError('Unknown IO Backend Mode.')\n\n\ndef rmtree(path):\n \"\"\"Remove a directory and recursively all contents inside. Might fail if used on a file path.\"\"\"\n if io_mode == BackendMode.DEFAULT:\n return shutil.rmtree(path)\n elif io_mode == BackendMode.TF:\n return gfile.rmtree(path)\n else:\n raise ValueError('Unknown IO Backend Mode.')\n\n\ndef getsize(path):\n \"\"\"Return the size, in bytes, of path.\"\"\"\n if io_mode == BackendMode.DEFAULT:\n return os.path.getsize(path)\n elif io_mode == BackendMode.TF:\n return gfile.stat(path).length\n else:\n raise ValueError('Unknown IO Backend Mode.')\n" }, { "path": "flax/jax_utils.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Utilities we could consider upstreaming to Jax.\"\"\"\n\nimport collections\nimport itertools\nimport warnings\nfrom collections.abc import Iterable # pylint: disable=g-importing-member\n\nimport jax\nimport jax.numpy as jnp\nimport numpy as np\nfrom jax import core, lax\nfrom jax.extend import linear_util as lu\nfrom jax.interpreters import partial_eval as pe\n\n\ndef _pmap_device_order():\n return jax.local_devices()\n\n\ndef replicate(tree, devices=None):\n \"\"\"Replicates arrays to multiple devices.\n\n Args:\n tree: a pytree containing the arrays that should be replicated.\n devices: the devices the data is replicated to\n (default: same order as expected by ``jax.pmap()``).\n Returns:\n A new pytree containing the replicated arrays.\n \"\"\"\n devices = devices or _pmap_device_order()\n return jax.device_put_replicated(tree, devices)\n\n\ndef unreplicate(tree):\n \"\"\"Returns a single instance of a replicated array.\"\"\"\n def _unreplicate_one(x):\n # Avoid degraded performance under the new jax.pmap.\n # Handle 0-dimensional (scalar) arrays - cannot index into them\n if hasattr(x, 'ndim') and x.ndim == 0:\n return x\n if (\n not hasattr(x, 'sharding')\n or isinstance(x.sharding, jax.sharding.SingleDeviceSharding)\n or len(jax.local_devices()) == 1\n ):\n return x[0]\n if x.sharding.is_fully_replicated:\n return x.addressable_shards[0].data\n return x.addressable_shards[0].data.squeeze(0)\n\n return jax.tree_util.tree_map(_unreplicate_one, tree)\n\n\ndef pmean(xs, axis_name):\n warnings.warn('use jax.lax.pmean instead', DeprecationWarning)\n return lax.pmean(xs, axis_name)\n\n\ndef partial_eval_by_shape(fn, input_spec, *args, **kwargs):\n \"\"\"Lazily evaluate a function by using the shapes of the inputs.\n\n This function is similar to ``jax.eval_shape`` with the key difference that\n function outputs that can be computed without a concrete value of the\n inputs are returned as is instead of only the shape. See for example\n ``module.init_by_shape`` where this functionality is used to initialize a\n model without using input data lr computation.\n\n Args:\n fn: the function to be lazily evaluated.\n input_spec: an iterable of shapes or (shape, dtype) tuples specifying the\n shape and type of the inputs. If unspecified the dtype is float32.\n *args: other arguments passed to the module's apply function\n **kwargs: keyword arguments passed to the module's apply function\n Returns:\n A pair consisting of the model output and an instance of Model\n \"\"\"\n # output cannot be returned in lazy_create because jax.eval_shape will only\n # return the shape and dtype.\n # TODO(mattjj,jheek): use a public JAX API\n f = lambda *inputs: fn(*inputs, *args, **kwargs)\n input_structs = [_parse_spec(spec) for spec in input_spec]\n inputs_flat, in_tree = jax.tree_util.tree_flatten(input_structs)\n\n debug_info = jax.api_util.debug_info(\"flax partial_eval_by_shape\", f,\n (in_tree,), {})\n f_flat, out_tree = jax.api_util.flatten_fun_nokwargs(\n lu.wrap_init(f, debug_info=debug_info), in_tree)\n in_pvals = [\n pe.PartialVal.unknown(core.ShapedArray(x.shape, x.dtype))\n for x in inputs_flat\n ]\n _, out_pvals, _ = pe.trace_to_jaxpr_nounits(f_flat, in_pvals)\n out_flat = [\n const if pv is None else jax.ShapeDtypeStruct(pv.shape, pv.dtype)\n for pv, const in out_pvals\n ]\n return jax.tree_util.tree_unflatten(out_tree(), out_flat)\n\n\ndef _parse_spec(spec):\n \"\"\"Parse an input spec of the form (shape, dtype) or shape into a jax.ShapeDtypeStruct.\"\"\"\n spec = tuple(spec)\n if len(spec) == 2 and isinstance(spec[0], Iterable):\n return jax.ShapeDtypeStruct(tuple(spec[0]), spec[1])\n else:\n return jax.ShapeDtypeStruct(spec, jnp.float32)\n\n\ndef prefetch_to_device(iterator, size, devices=None):\n \"\"\"Shard and prefetch batches on device.\n\n This utility takes an iterator and returns a new iterator which fills an on\n device prefetch buffer. Eager prefetching can improve the performance of\n training loops significantly by overlapping compute and data transfer.\n\n This utility is mostly useful for GPUs, for TPUs and CPUs it should not be\n necessary -- the TPU & CPU memory allocators (normally) don't pick a memory\n location that isn't free yet so they don't block. Instead those allocators OOM.\n\n Args:\n iterator: an iterator that yields a pytree of ndarrays where the first\n dimension is sharded across devices.\n\n size: the size of the prefetch buffer.\n\n If you're training on GPUs, 2 is generally the best choice because this\n guarantees that you can overlap a training step on GPU with a data\n prefetch step on CPU.\n\n devices: the list of devices to which the arrays should be prefetched.\n\n Defaults to the order of devices expected by ``jax.pmap``.\n\n Yields:\n The original items from the iterator where each ndarray is now sharded to\n the specified devices.\n \"\"\"\n queue = collections.deque()\n devices = _pmap_device_order() if devices is None else devices\n\n def _prefetch(xs):\n return jax.device_put_sharded(list(xs), devices)\n\n def enqueue(n): # Enqueues *up to* `n` elements from the iterator.\n for data in itertools.islice(iterator, n):\n queue.append(jax.tree_util.tree_map(_prefetch, data))\n\n enqueue(size) # Fill up the buffer.\n while queue:\n yield queue.popleft()\n enqueue(1)\n\n\ndef _scan_nd(body_fn, init, xs, n=1, unroll=(1,)):\n \"\"\"Utility for performing an n-dimensional `lax.scan`.\n\n The n-d scan is simply recursive call of 1-d scan.\n Args:\n body_fn: the body of the loop of type (c, x) -> (c, y).\n init: initial value for the carry.\n xs: a pytree of tensors to scan over.\n n: number of dimensions to scan over (default: 1)\n Returns:\n A tuple of the final carry and the values returned by the body.\n \"\"\"\n if n == 1:\n return lax.scan(body_fn, init, xs, unroll=unroll[0])\n else:\n\n def scan_body(c, x):\n return _scan_nd(body_fn, c, x, n=n - 1, unroll=unroll[1:])\n\n return lax.scan(scan_body, init, xs, unroll=unroll[0])\n\n\ndef _invert_perm(perm):\n perm_inv = [0] * len(perm)\n for i, j in enumerate(perm):\n perm_inv[j] = i\n return tuple(perm_inv)\n\n\ndef scan_in_dim(body_fn, init, xs, axis=(0,), unroll=(1,), keepdims=False):\n \"\"\"utility for doing a scan along arbitrary dimensions.\n\n See `lax.scan` for details on how the scan operation works.\n\n Note on `unroll`: This argument gets left padded with ones to match the size\n of `axis`. Doing so allows unrolls to performed from the innermost loop first.\n For example, `scan_in_dim(..., axis=(1, 2, 3), unroll=5)` is equivalent to\n `scan_in_dim(..., axis=(1, 2, 3), unroll=(1, 1, 5))`.\n\n Args:\n body_fn: the body of the loop of type (c, x) -> (c, y).\n init: initial value for the carry.\n xs: a pytree of tensors to scan over.\n axis: the axis to scan over.\n keepdims: keep the dimensions that are scanned over.\n unroll: an optional positive integer, or tuple of positive integers\n showing how many iterations of the loop to be unrolled into a single\n iteration for each axis.\n Returns:\n A tuple of the final carry and the values returned by the body.\n \"\"\"\n if not isinstance(axis, Iterable):\n axis = (axis,)\n\n if not isinstance(unroll, Iterable):\n unroll = (unroll,)\n\n # Pad unroll with ones so we start unrolling from the innermost loop\n len_diff = len(axis) - len(unroll)\n unroll = (1,) * len_diff + unroll\n\n def transpose_in(x):\n perm = axis + tuple(np.delete(np.arange(x.ndim), axis))\n return x.transpose(perm)\n\n def transpose_out(x):\n perm = axis + tuple(np.delete(np.arange(x.ndim), axis))\n return x.transpose(_invert_perm(perm))\n\n def body_wrapper(c, xs):\n if keepdims:\n xs = jax.tree_util.tree_map(\n lambda x: x.reshape((1,) * len(axis) + x.shape), xs\n )\n xs = jax.tree_util.tree_map(transpose_out, xs)\n c, ys = body_fn(c, xs)\n if keepdims:\n ys = jax.tree_util.tree_map(transpose_in, ys)\n ys = jax.tree_util.tree_map(lambda x: x.reshape(x.shape[len(axis) :]), ys)\n return c, ys\n\n xs = jax.tree_util.tree_map(transpose_in, xs)\n c, ys = _scan_nd(body_wrapper, init, xs, n=len(axis), unroll=unroll)\n ys = jax.tree_util.tree_map(transpose_out, ys)\n return c, ys\n\n\n# Copied from https://github.com/google-research/big_vision\ndef pad_shard_unpad(\n wrapped, static_argnums=(0,), static_argnames=(), static_return=False\n):\n \"\"\"Wraps a function with code that pads, shards, then un-shards, un-pads.\n\n Args:\n wrapped: the function to be wrapped. Signature is ``params, *args, *kwargs``.\n static_argnums: indices of arguments to ``wrapped`` that should _not_ be\n padded and sharded, but instead be forwarded as-is. The default is (0,)\n because by far the most common use-case is to pass ``params`` first.\n static_argnames: names of kwargs to ``wrapped`` that should _not_ be padded\n and sharded, but instead be forwarded as-is.\n static_return: whether not to un-shard, and un-pad the return value; static\n return values are typically used with eval steps that compute metrics\n\n Returns:\n A new function that pads and shards its arguments before passing them to\n the wrapped function, and un-shards and un-pads the returned pytree.\n\n This is useful for calling a pmap'ed function with inputs that aren't\n divisible by the number of devices. A typical use is:\n @pad_shard_unpad\n @jax.pmap\n def forward(params, x): ...\n\n Notes:\n The padding is done in host-memory before being passed to the function, and\n the values returned by the function are transferred back to host memory.\n\n The returned function is augmented with a new keyword-only argument\n ``min_device_batch`` that, if specified, forces padding inputs to at least\n this size per device. This can be useful to avoid recompiles for the last\n batch and reduce memory fragmentation.\n\n For more information refer to https://flax.readthedocs.io/en/latest/guides/data_preprocessing/full_eval.html\n \"\"\"\n\n def pad_shard_unpad_wrapper(*args, min_device_batch=None, **kw):\n d = jax.local_device_count() # d = devices, b = batch\n batch_sizes = set()\n for i, a in enumerate(args):\n if i not in static_argnums:\n batch_sizes |= {t.shape[0] for t in jax.tree_util.tree_leaves(a)}\n for k, v in kw.items():\n if k not in static_argnames:\n batch_sizes |= {t.shape[0] for t in jax.tree_util.tree_leaves(v)}\n assert len(batch_sizes) == 1, f'Inconsistent batch-sizes: {batch_sizes}'\n b = batch_sizes.pop()\n\n def pad(x):\n _, *shape = x.shape\n db, rest = divmod(b, d)\n if rest:\n x = np.concatenate([x, np.zeros((d - rest, *shape), x.dtype)], axis=0)\n db += 1\n if min_device_batch and db < min_device_batch:\n x = np.concatenate(\n [x, np.zeros((d * (min_device_batch - db), *shape), x.dtype)]\n )\n db = min_device_batch\n return x.reshape(d, db, *shape)\n\n def maybe_pad(tree, actually_pad=True):\n if not actually_pad:\n return tree # For call-site convenience below.\n return jax.tree_util.tree_map(pad, tree)\n\n args = [maybe_pad(a, i not in static_argnums) for i, a in enumerate(args)]\n kw = {k: maybe_pad(v, k not in static_argnames) for k, v in kw.items()}\n out = wrapped(*args, **kw)\n\n def unpad(x):\n # Transfer back before cutting, to reduce on-device shape diversity.\n return jax.device_get(x).reshape([np.prod(x.shape[:2]), *x.shape[2:]])[:b]\n\n return out if static_return else jax.tree_util.tree_map(unpad, out)\n\n return pad_shard_unpad_wrapper\n" }, { "path": "flax/linen/README.md", "content": "# Linen: A comfortable evolution of Flax\n\nLinen is a neural network API developed based on learning from our users and the broader JAX community. Linen improves on much of the former `flax.nn` API (removed since v0.4.0), such as submodule sharing and better support for non-trainable variables.\nMoreover, Linen builds on a \"functional core\", enabling direct usage of JAX transformations such as `vmap`, `remat` or `scan` inside your modules.\n\nIn Linen, Modules behave much closer to vanilla Python objects, while still letting you opt-in to the concise single-method pattern many of our users love.\n\nThe Linen Module API is stable and currently recommended for new projects. We are already supporting users in the OSS community and within Google. Minor changes may come to the top-level `apply` and `init` patterns, which we will communicate clearly. We plan a few improvements, including writing up short design notes, adding more design tests (see last link below), and an API for interactive module instances.\n\nPlease open a [discussion](https://github.com/google/flax/discussions) if you have any questions or thoughts.\n\n**See the [Linen API reference docs](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/index.html)**, or take a look at our additional material:\n\n* 2-page intro to the [Linen Design Principles](https://docs.google.com/document/d/1ZlL_4bXCw5Xl0WstQw1GpnZqfb9JFOeUGAPcBVk-kn8)\n* [Slides from a talk to the JAX core team](https://docs.google.com/presentation/d/1ngKWUwsSqAwPRvATG8sAxMzu9ujv4N__cKsUofdNno0)\n* [Brief Intro to Linen](https://colab.research.google.com/github/google/flax/blob/main/docs/linen_intro.ipynb) in Colab\n* An [upgrade guide](https://docs.google.com/document/d/1hYavTVPaKVVe9Be8pCB7yW7r6dDv3RALVNit8NZca4c) + some additional questions we're considering\n* Ported [examples](https://github.com/google/flax/tree/main/examples)\n* \"Design tests\" used to ensure that our \"functional core\" supports [various advanced use-cases](https://github.com/google/flax/tree/main/tests/core/design), and that the mostly-syntactic-sugar Module abstraction\n [doesn't get in the way](https://github.com/google/flax/tree/main/examples/linen_design_test)\n" }, { "path": "flax/linen/__init__.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"The Flax Module system.\"\"\"\n\n\n# pylint: disable=g-multiple-import,useless-import-alias\n# re-export commonly used modules and functions\nfrom flax.core import (\n DenyList as DenyList,\n FrozenDict as FrozenDict,\n broadcast as broadcast,\n meta as meta,\n)\nfrom flax.core.meta import (\n PARTITION_NAME as PARTITION_NAME,\n Partitioned as Partitioned,\n get_partition_spec as get_partition_spec,\n get_sharding as get_sharding,\n unbox as unbox,\n with_partitioning as with_partitioning,\n)\nfrom flax.core.spmd import (\n get_logical_axis_rules as get_logical_axis_rules,\n logical_axis_rules as logical_axis_rules,\n set_logical_axis_rules as set_logical_axis_rules,\n)\nfrom .activation import (\n PReLU as PReLU,\n celu as celu,\n elu as elu,\n gelu as gelu,\n glu as glu,\n hard_sigmoid as hard_sigmoid,\n hard_silu as hard_silu,\n hard_swish as hard_swish,\n hard_tanh as hard_tanh,\n leaky_relu as leaky_relu,\n log_sigmoid as log_sigmoid,\n log_softmax as log_softmax,\n logsumexp as logsumexp,\n normalize as normalize,\n one_hot as one_hot,\n relu6 as relu6,\n relu as relu,\n selu as selu,\n sigmoid as sigmoid,\n silu as silu,\n soft_sign as soft_sign,\n softmax as softmax,\n softplus as softplus,\n standardize as standardize,\n swish as swish,\n tanh as tanh,\n)\nfrom .attention import (\n MultiHeadAttention as MultiHeadAttention,\n MultiHeadDotProductAttention as MultiHeadDotProductAttention,\n SelfAttention as SelfAttention,\n combine_masks as combine_masks,\n dot_product_attention_weights as dot_product_attention_weights,\n dot_product_attention as dot_product_attention,\n make_attention_mask as make_attention_mask,\n make_causal_mask as make_causal_mask,\n)\nfrom .batch_apply import BatchApply as BatchApply\nfrom .combinators import Sequential as Sequential\nfrom .fp8_ops import (\n Fp8DirectDotGeneralOp as Fp8DirectDotGeneralOp,\n Fp8DotGeneral as Fp8DotGeneral,\n Fp8DotGeneralOp as Fp8DotGeneralOp,\n Fp8Einsum as Fp8Einsum,\n NANOOFp8DotGeneralOp as NANOOFp8DotGeneralOp,\n)\nfrom .initializers import (\n ones_init as ones_init,\n ones as ones,\n zeros_init as zeros_init,\n zeros as zeros,\n)\nfrom .linear import (\n ConvLocal as ConvLocal,\n ConvTranspose as ConvTranspose,\n Conv as Conv,\n DenseGeneral as DenseGeneral,\n Dense as Dense,\n Einsum as Einsum,\n Embed as Embed,\n)\nfrom .module import (\n Module as Module,\n Variable as Variable,\n apply as apply,\n compact_name_scope as compact_name_scope,\n compact as compact,\n disable_named_call as disable_named_call,\n enable_named_call as enable_named_call,\n init_with_output as init_with_output,\n init as init,\n intercept_methods as intercept_methods,\n merge_param as merge_param,\n nowrap as nowrap,\n override_named_call as override_named_call,\n share_scope as share_scope,\n)\nfrom .normalization import (\n BatchNorm as BatchNorm,\n GroupNorm as GroupNorm,\n InstanceNorm as InstanceNorm,\n LayerNorm as LayerNorm,\n RMSNorm as RMSNorm,\n SpectralNorm as SpectralNorm,\n WeightNorm as WeightNorm,\n)\nfrom .pooling import (avg_pool as avg_pool, max_pool as max_pool, pool as pool)\nfrom .recurrent import (\n Bidirectional as Bidirectional,\n ConvLSTMCell as ConvLSTMCell,\n GRUCell as GRUCell,\n LSTMCell as LSTMCell,\n MGUCell as MGUCell,\n OptimizedLSTMCell as OptimizedLSTMCell,\n RNNCellBase as RNNCellBase,\n RNN as RNN,\n SimpleCell as SimpleCell,\n)\nfrom .spmd import (\n LogicallyPartitioned as LogicallyPartitioned,\n logical_to_mesh,\n logical_to_mesh_axes,\n logical_to_mesh_sharding,\n with_logical_constraint,\n with_logical_partitioning as with_logical_partitioning,\n)\nfrom .stochastic import Dropout as Dropout\nfrom .summary import tabulate\nfrom .transforms import (\n add_metadata_axis,\n checkpoint as checkpoint,\n cond as cond,\n custom_vjp as custom_vjp,\n fold_rngs as fold_rngs,\n grad as grad,\n jit as jit,\n jvp as jvp,\n map_variables as map_variables,\n named_call as named_call,\n remat_scan as remat_scan,\n remat as remat,\n scan as scan,\n switch as switch,\n value_and_grad as value_and_grad,\n vjp as vjp,\n vmap as vmap,\n while_loop as while_loop,\n)\n# pylint: enable=g-multiple-import\n" }, { "path": "flax/linen/activation.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Activation functions.\"\"\"\n\n# pylint: disable=unused-import\n# re-export activation functions from jax.nn\nfrom typing import Any, Optional\n\nfrom flax.linen.module import compact\nfrom flax.linen.module import Module\nfrom flax.linen.linear import Dense\nfrom flax.typing import Array, Dtype\n\nfrom jax.nn import celu\nfrom jax.nn import elu\nfrom jax.nn import gelu\nfrom jax.nn import glu\nfrom jax.nn import hard_sigmoid\nfrom jax.nn import hard_silu\nfrom jax.nn import hard_swish\nfrom jax.nn import hard_tanh\nfrom jax.nn import leaky_relu\nfrom jax.nn import log_sigmoid\nfrom jax.nn import log_softmax\nfrom jax.nn import logsumexp\nfrom jax.nn import one_hot\nfrom jax.nn import relu\nfrom jax.nn import relu6\nfrom jax.nn import selu\nfrom jax.nn import sigmoid\nfrom jax.nn import silu\nfrom jax.nn import soft_sign\nfrom jax.nn import softmax\nfrom jax.nn import softplus\nfrom jax.nn import standardize\nfrom jax.nn import swish\nimport jax.numpy as jnp\nfrom jax.numpy import tanh\n\n# Normalize is a deprecated alias of standardize\nnormalize = standardize\n\n# pylint: enable=unused-import\n\n\n\nclass PReLU(Module):\n \"\"\"Parametric Rectified Linear Unit (PReLU) activation function.\n\n Note that PReLU is a Flax layer and not a simple activation function, so\n it needs to be initialized before being called.\n\n Example usage::\n >>> import flax.linen as nn\n\n >>> class MLP(nn.Module):\n ... @nn.compact\n ... def __call__(self, x):\n ... x = nn.Dense(2)(x)\n ... x = nn.PReLU()(x) # initialized\n ... return x\n\n Attributes:\n param_dtype: the dtype passed to parameter initializers (default: float32).\n negative_slope_init: the value to initialize the negative slope\n (default 0.01).\n \"\"\"\n\n param_dtype: Dtype = jnp.float32\n negative_slope_init: float = 0.01\n\n @compact\n def __call__(self, inputs: Array) -> Array:\n \"\"\"Applies an activation to the inputs.\n\n Args:\n inputs: the nd-array to apply the activation function to.\n\n Returns:\n The transformed input.\n \"\"\"\n negative_slope = self.param(\n 'negative_slope',\n lambda k: jnp.asarray(self.negative_slope_init, self.param_dtype),\n )\n return jnp.where(\n inputs >= 0, inputs, jnp.asarray(negative_slope, inputs.dtype) * inputs\n )\n" }, { "path": "flax/linen/attention.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Attention core modules for Flax.\"\"\"\nfrom __future__ import annotations\n\nimport functools\nimport inspect\nimport warnings\nfrom typing import Any, overload\nfrom collections.abc import Callable\n\nimport jax\nimport jax.numpy as jnp\nfrom jax import lax, random\n\nfrom flax.linen import initializers\nfrom flax.linen.dtypes import promote_dtype\nfrom flax.linen.linear import (\n DenseGeneral,\n default_kernel_init,\n)\nfrom flax.linen.module import Module, compact, merge_param\nfrom flax.linen.normalization import LayerNorm\nfrom flax.typing import (\n Array,\n PRNGKey,\n Dtype,\n Shape as Shape,\n Initializer,\n PrecisionLike,\n DotGeneralT,\n)\n\n\ndef dot_product_attention_weights(\n query: Array,\n key: Array,\n bias: Array | None = None,\n mask: Array | None = None,\n broadcast_dropout: bool = True,\n dropout_rng: PRNGKey | None = None,\n dropout_rate: float = 0.0,\n deterministic: bool = False,\n dtype: Dtype | None = None,\n precision: PrecisionLike = None,\n module: Module | None = None,\n force_fp32_for_softmax: bool = False,\n einsum_dot_general: Callable[..., Array] | None = None,\n einsum: Callable[..., Array] | None = None,\n):\n \"\"\"Computes dot-product attention weights given query and key.\n\n Used by :func:`dot_product_attention`, which is what you'll most likely use.\n But if you want access to the attention weights for introspection, then\n you can directly call this function and call einsum yourself.\n\n Args:\n query: queries for calculating attention with shape of ``[batch...,\n q_length, num_heads, qk_depth_per_head]``.\n key: keys for calculating attention with shape of ``[batch..., kv_length,\n num_heads, qk_depth_per_head]``.\n bias: bias for the attention weights. This should be broadcastable to the\n shape ``[batch..., num_heads, q_length, kv_length]``. This can be used for\n incorporating causal masks, padding masks, proximity bias, etc.\n mask: mask for the attention weights. This should be broadcastable to the\n shape ``[batch..., num_heads, q_length, kv_length]``. This can be used for\n incorporating causal masks. Attention weights are masked out if their\n corresponding mask value is ``False``.\n broadcast_dropout: bool: use a broadcasted dropout along batch dims.\n dropout_rng: JAX PRNGKey: to be used for dropout\n dropout_rate: dropout rate\n deterministic: bool, deterministic or not (to apply dropout)\n dtype: the dtype of the computation (default: infer from inputs and params)\n precision: numerical precision of the computation see ``jax.lax.Precision``\n for details.\n module: the Module that will sow the attention weights into the\n 'intermediates' collection. Remember to mark 'intermediates' as mutable\n via ``mutable=['intermediates']`` in order to have that collection\n returned. If ``module`` is None, the attention weights will not be sowed.\n force_fp32_for_softmax: bool, whether to force the softmax to be computed in\n fp32. This is useful for mixed-precision training where higher precision\n is desired for numerical stability.\n einsum_dot_general: the dot_general to use in einsum.\n einsum: If unspecified, default `jnp.einsum` will be used. This argument is\n mutually exclusive with `precision` and `einsum_dot_general`.\n\n Raises:\n ValueError: if both `precision`/`einsum_dot_general` and `einsum` are\n specified.\n\n Returns:\n Output of shape ``[batch..., num_heads, q_length, kv_length]``.\n \"\"\"\n if (precision or einsum_dot_general) and einsum:\n raise ValueError(\n 'precision/einsum_dot_general and einsum are mutually exclusive. Please'\n ' specify only one of them.'\n )\n if not einsum:\n einsum = functools.partial(\n jnp.einsum,\n precision=precision,\n _dot_general=einsum_dot_general\n if einsum_dot_general\n else jax.lax.dot_general,\n )\n\n query, key = promote_dtype(query, key, dtype=dtype)\n dtype = query.dtype\n\n assert query.ndim == key.ndim, 'q, k must have same rank.'\n assert query.shape[:-3] == key.shape[:-3], 'q, k batch dims must match.'\n assert query.shape[-2] == key.shape[-2], 'q, k num_heads must match.'\n assert query.shape[-1] == key.shape[-1], 'q, k depths must match.'\n\n # calculate attention matrix\n depth = query.shape[-1]\n query = query / jnp.sqrt(depth).astype(dtype)\n # attn weight shape is (batch..., num_heads, q_length, kv_length)\n attn_weights = einsum('...qhd,...khd->...hqk', query, key)\n\n # apply attention bias: masking, dropout, proximity bias, etc.\n if bias is not None:\n attn_weights = attn_weights + bias\n # apply attention mask\n if mask is not None:\n big_neg = jnp.finfo(dtype).min\n attn_weights = jnp.where(mask, attn_weights, big_neg)\n\n # normalize the attention weights\n if force_fp32_for_softmax and dtype != jnp.float32:\n attn_weights = jax.nn.softmax(attn_weights.astype(jnp.float32))\n else:\n attn_weights = jax.nn.softmax(attn_weights).astype(dtype)\n\n if module:\n module.sow('intermediates', 'attention_weights', attn_weights)\n\n # apply attention dropout\n if not deterministic and dropout_rate > 0.0:\n keep_prob = 1.0 - dropout_rate\n if broadcast_dropout:\n # dropout is broadcast across the batch + head dimensions\n dropout_shape = tuple([1] * (key.ndim - 2)) + attn_weights.shape[-2:]\n keep = random.bernoulli(dropout_rng, keep_prob, dropout_shape) # type: ignore\n else:\n keep = random.bernoulli(dropout_rng, keep_prob, attn_weights.shape) # type: ignore\n multiplier = keep.astype(dtype) / jnp.asarray(keep_prob, dtype=dtype)\n attn_weights = attn_weights * multiplier\n\n return attn_weights\n\n\ndef dot_product_attention(\n query: Array,\n key: Array,\n value: Array,\n bias: Array | None = None,\n mask: Array | None = None,\n broadcast_dropout: bool = True,\n dropout_rng: PRNGKey | None = None,\n dropout_rate: float = 0.0,\n deterministic: bool = False,\n dtype: Dtype | None = None,\n precision: PrecisionLike = None,\n module: Module | None = None,\n force_fp32_for_softmax: bool = False,\n einsum_dot_general: Callable[..., Array] | None = None,\n qk_attn_weights_einsum: Callable[..., Array] | None = None,\n attn_weights_value_einsum: Callable[..., Array] | None = None,\n):\n \"\"\"Computes dot-product attention given query, key, and value.\n\n This is the core function for applying attention based on\n https://arxiv.org/abs/1706.03762. It calculates the attention weights given\n query and key and combines the values using the attention weights.\n\n .. note::\n ``query``, ``key``, ``value`` needn't have any batch dimensions.\n\n Args:\n query: queries for calculating attention with shape of ``[batch...,\n q_length, num_heads, qk_depth_per_head]``.\n key: keys for calculating attention with shape of ``[batch..., kv_length,\n num_heads, qk_depth_per_head]``.\n value: values to be used in attention with shape of ``[batch..., kv_length,\n num_heads, v_depth_per_head]``.\n bias: bias for the attention weights. This should be broadcastable to the\n shape ``[batch..., num_heads, q_length, kv_length]``. This can be used for\n incorporating causal masks, padding masks, proximity bias, etc.\n mask: mask for the attention weights. This should be broadcastable to the\n shape ``[batch..., num_heads, q_length, kv_length]``. This can be used for\n incorporating causal masks. Attention weights are masked out if their\n corresponding mask value is ``False``.\n broadcast_dropout: bool: use a broadcasted dropout along batch dims.\n dropout_rng: JAX PRNGKey: to be used for dropout\n dropout_rate: dropout rate\n deterministic: bool, deterministic or not (to apply dropout)\n dtype: the dtype of the computation (default: infer from inputs)\n precision: numerical precision of the computation see ``jax.lax.Precision`\n for details.\n module: the Module that will sow the attention weights into the\n 'intermediates' collection. Remember to mark 'intermediates' as mutable\n via ``mutable=['intermediates']`` in order to have that collection\n returned. If ``module`` is None, the attention weights will not be sowed.\n force_fp32_for_softmax: bool, whether to force the softmax to be computed in\n fp32. This is useful for mixed-precision training where higher precision\n is desired for numerical stability.\n einsum_dot_general: the dot_general to use in `jnp.einsum`.\n qk_attn_weights_einsum: the einsum for computing the attention weights. When\n unspecified, the default `jnp.einsum` will be used. This argument is\n mutually exclusive with `precision` and `einsum_dot_general`.\n attn_weights_value_einsum: the einsum for computing the product of the\n attention weights and the values. When unspecified, the default\n `jnp.einsum` will be used. This argument is mutually exclusive with\n `precision` and `einsum_dot_general`.\n\n Returns:\n Output of shape ``[batch..., q_length, num_heads, v_depth_per_head]``.\n\n Raises:\n ValueError: if both `precision`/`einsum_dot_general` and\n `qk_attn_weights_einsum`/`attn_weights_value_einsum` are\n specified.\n \"\"\"\n if (qk_attn_weights_einsum and not attn_weights_value_einsum) or (\n not qk_attn_weights_einsum and attn_weights_value_einsum\n ):\n raise ValueError(\n 'qk_attn_weights_einsum and attn_weights_value_einsum must be specified'\n ' together.'\n )\n if (precision or einsum_dot_general) and (\n qk_attn_weights_einsum or attn_weights_value_einsum\n ):\n raise ValueError(\n 'precision/einsum_dot_general and'\n ' qk_attn_weights_einsum/attn_weights_value_einsum are mutually'\n ' exclusive. Please specify only one of them.'\n )\n\n query, key, value = promote_dtype(query, key, value, dtype=dtype)\n dtype = query.dtype\n assert key.ndim == query.ndim == value.ndim, 'q, k, v must have same rank.'\n assert (\n query.shape[:-3] == key.shape[:-3] == value.shape[:-3]\n ), 'q, k, v batch dims must match.'\n assert (\n query.shape[-2] == key.shape[-2] == value.shape[-2]\n ), 'q, k, v num_heads must match.'\n assert key.shape[-3] == value.shape[-3], 'k, v lengths must match.'\n\n # compute attention weights\n attn_weights = dot_product_attention_weights(\n query,\n key,\n bias,\n mask,\n broadcast_dropout,\n dropout_rng,\n dropout_rate,\n deterministic,\n dtype,\n precision,\n module,\n force_fp32_for_softmax,\n einsum_dot_general=einsum_dot_general,\n einsum=qk_attn_weights_einsum,\n )\n if not attn_weights_value_einsum:\n attn_weights_value_einsum = functools.partial(\n jnp.einsum,\n precision=precision,\n _dot_general=einsum_dot_general\n if einsum_dot_general\n else jax.lax.dot_general,\n )\n # return weighted sum over values for each query position\n return attn_weights_value_einsum(\n '...hqk,...khd->...qhd',\n attn_weights,\n value,\n )\n\n\nclass MultiHeadDotProductAttention(Module):\n \"\"\"Multi-head dot-product attention.\n\n Example usage::\n\n >>> import flax.linen as nn\n >>> import jax\n\n >>> layer = nn.MultiHeadDotProductAttention(num_heads=8, qkv_features=16)\n >>> key1, key2, key3, key4, key5, key6 = jax.random.split(jax.random.key(0), 6)\n >>> shape = (4, 3, 2, 5)\n >>> q, k, v = jax.random.uniform(key1, shape), jax.random.uniform(key2, shape), jax.random.uniform(key3, shape)\n >>> variables = layer.init(jax.random.key(0), q)\n\n >>> # different inputs for inputs_q, inputs_k and inputs_v\n >>> out = layer.apply(variables, q, k, v)\n >>> # equivalent to layer.apply(variables, inputs_q=q, inputs_k=k, inputs_v=k)\n >>> out = layer.apply(variables, q, k)\n >>> # equivalent to layer.apply(variables, inputs_q=q, inputs_k=q) and layer.apply(variables, inputs_q=q, inputs_k=q, inputs_v=q)\n >>> out = layer.apply(variables, q)\n\n >>> attention_kwargs = dict(\n ... num_heads=8,\n ... qkv_features=16,\n ... kernel_init=nn.initializers.ones,\n ... bias_init=nn.initializers.zeros,\n ... dropout_rate=0.5,\n ... deterministic=False,\n ... )\n >>> class Module(nn.Module):\n ... attention_kwargs: dict\n ...\n ... @nn.compact\n ... def __call__(self, x, dropout_rng=None):\n ... out1 = nn.MultiHeadDotProductAttention(**self.attention_kwargs)(x, dropout_rng=dropout_rng)\n ... out2 = nn.MultiHeadDotProductAttention(**self.attention_kwargs)(x, dropout_rng=dropout_rng)\n ... return out1, out2\n >>> module = Module(attention_kwargs)\n >>> variables = module.init({'params': key1, 'dropout': key2}, q)\n\n >>> # out1 and out2 are different.\n >>> out1, out2 = module.apply(variables, q, rngs={'dropout': key3})\n >>> # out3 and out4 are different.\n >>> # out1 and out3 are different. out2 and out4 are different.\n >>> out3, out4 = module.apply(variables, q, rngs={'dropout': key4})\n >>> # out1 and out2 are the same.\n >>> out1, out2 = module.apply(variables, q, dropout_rng=key5)\n >>> # out1 and out2 are the same as out3 and out4.\n >>> # providing a `dropout_rng` arg will take precedence over the `rngs` arg in `.apply`\n >>> out3, out4 = module.apply(variables, q, rngs={'dropout': key6}, dropout_rng=key5)\n\n Attributes:\n num_heads: Number of attention heads. Features (i.e. inputs_q.shape[-1])\n should be divisible by the number of heads.\n dtype: The dtype of the computation (default: infer from inputs and params)\n param_dtype: The dtype passed to parameter initializers (default: float32)\n qkv_features: Dimension of the key, query, and value.\n out_features: Dimension of the last projection\n broadcast_dropout: Use a broadcasted dropout along batch dims.\n dropout_rate: Dropout rate.\n deterministic: If False, the attention weight is masked randomly using\n dropout, whereas if True, the attention weights are deterministic.\n precision: Numerical precision of the computation see ``jax.lax.Precision``\n for details.\n kernel_init: Initializer for the kernel of the Dense layers.\n out_kernel_init: Optional Initializer for the kernel of the output Dense layer,\n if None, ``kernel_init`` will be used.\n bias_init: Initializer for the bias of the Dense layers.\n out_bias_init: Optional Initializer for the bias of the output Dense layer,\n if None, ``bias_init`` will be used.\n use_bias: Whether pointwise QKVO dense transforms use bias.\n attention_fn: dot_product_attention or compatible function. Accepts query,\n key, value, and returns output of shape ``[bs, dim1, dim2, ..., dimN,,\n num_heads, value_channels]``\n decode: Whether to prepare and use an autoregressive cache.\n normalize_qk: Should QK normalization be applied (arxiv.org/abs/2302.05442).\n qk_attn_weights_einsum_cls: factory function to create the einsum for\n computing the attention weights.\n attn_weights_value_einsum_cls: factory function to create the einsum for\n computing the product of the attention weights and the values.\n \"\"\"\n\n num_heads: int\n dtype: Dtype | None = None\n param_dtype: Dtype = jnp.float32\n qkv_features: int | None = None\n out_features: int | None = None\n broadcast_dropout: bool = True\n dropout_rate: float = 0.0\n deterministic: bool | None = None\n precision: PrecisionLike = None\n kernel_init: Initializer = default_kernel_init\n out_kernel_init: Initializer | None = None\n bias_init: Initializer = initializers.zeros_init()\n out_bias_init: Initializer | None = None\n use_bias: bool = True\n attention_fn: Callable[..., Array] = dot_product_attention\n decode: bool = False\n normalize_qk: bool = False\n force_fp32_for_softmax: bool = False\n # Deprecated, will be removed.\n qkv_dot_general: DotGeneralT | None = None\n out_dot_general: DotGeneralT | None = None\n qkv_dot_general_cls: Any = None\n out_dot_general_cls: Any = None\n qk_attn_weights_einsum_cls: Callable[..., Callable[..., Array]] | None = None\n attn_weights_value_einsum_cls: Callable[..., Callable[..., Array]] | None = (\n None\n )\n\n @overload\n def __call__(\n self,\n inputs_q: Array,\n inputs_k: Array | None = None,\n inputs_v: Array | None = None,\n *,\n mask: Array | None = None,\n deterministic: bool | None = None,\n dropout_rng: PRNGKey | None = None,\n sow_weights: bool = False,\n ):\n ...\n\n @overload\n def __call__(\n self,\n inputs_q: Array,\n *,\n inputs_kv: Array | None = None,\n mask: Array | None = None,\n deterministic: bool | None = None,\n dropout_rng: PRNGKey | None = None,\n sow_weights: bool = False,\n ):\n ...\n\n @compact\n def __call__(\n self,\n inputs_q: Array,\n inputs_k: Array | None = None,\n inputs_v: Array | None = None,\n *,\n inputs_kv: Array | None = None,\n mask: Array | None = None,\n deterministic: bool | None = None,\n dropout_rng: PRNGKey | None = None,\n sow_weights: bool = False,\n ):\n \"\"\"Applies multi-head dot product attention on the input data.\n\n Projects the inputs into multi-headed query, key, and value vectors,\n applies dot-product attention and project the results to an output vector.\n\n If both inputs_k and inputs_v are None, they will both copy the value of\n inputs_q (self attention).\n If only inputs_v is None, it will copy the value of inputs_k.\n\n Args:\n inputs_q: input queries of shape ``[batch_sizes..., length, features]``.\n inputs_k: key of shape ``[batch_sizes..., length, features]``. If None,\n inputs_k will copy the value of inputs_q.\n inputs_v: values of shape ``[batch_sizes..., length, features]``. If None,\n inputs_v will copy the value of inputs_k.\n inputs_kv: key/values of shape ``[batch_sizes..., length, features]``. If\n None, inputs_kv will copy the value of inputs_q. This arg will be\n deprecated soon. Use inputs_k and inputs_v instead.\n mask: attention mask of shape ``[batch_sizes..., num_heads, query_length,\n key/value_length]``. Attention weights are masked out if their\n corresponding mask value is ``False``.\n deterministic: if false, the attention weight is masked randomly using\n dropout, whereas if true, the attention weights are deterministic.\n dropout_rng: optional rng key to pass to the attention layer's dropout\n mask. Otherwise, self.make_rng('dropout') is used instead.\n sow_weights: if ``True``, the attention weights are sowed into the\n 'intermediates' collection. Remember to mark 'intermediates' as\n mutable via ``mutable=['intermediates']`` in order to have that\n collection returned.\n\n Returns:\n output of shape ``[batch_sizes..., length, features]``.\n \"\"\"\n if inputs_kv is not None:\n if inputs_k is not None or inputs_v is not None:\n raise ValueError(\n 'If either `inputs_k` or `inputs_v` is not None, '\n '`inputs_kv` must be None. If `inputs_kv` is not None, both `inputs_k` '\n 'and `inputs_v` must be None. We recommend using `inputs_k` and '\n '`inputs_v` args, since `inputs_kv` will be deprecated soon. See '\n 'https://github.com/google/flax/discussions/3389 for more '\n 'information.'\n )\n inputs_k = inputs_v = inputs_kv\n warnings.warn(\n 'The inputs_kv arg will be deprecated soon. '\n 'Use inputs_k and inputs_v instead. See '\n 'https://github.com/google/flax/discussions/3389 '\n 'for more information.',\n DeprecationWarning,\n )\n else:\n if inputs_k is None:\n if inputs_v is not None:\n raise ValueError(\n '`inputs_k` cannot be None if `inputs_v` is not None. '\n 'To have both `inputs_k` and `inputs_v` be the same value, pass in the '\n 'value to `inputs_k` and leave `inputs_v` as None.'\n )\n inputs_k = inputs_q\n if inputs_v is None:\n inputs_v = inputs_k\n elif inputs_v.shape[-1] == inputs_v.shape[-2]:\n warnings.warn(\n f'You are passing an array of shape {inputs_v.shape} '\n 'to the `inputs_v` arg, when you may have intended '\n 'to pass it to the `mask` arg. As of Flax version '\n '0.7.4, the function signature of '\n \"MultiHeadDotProductAttention's `__call__` method \"\n 'has changed to `__call__(inputs_q, inputs_k=None, '\n 'inputs_v=None, *, inputs_kv=None, mask=None, '\n 'deterministic=None)`. Use the kwarg `mask` instead. '\n 'See https://github.com/google/flax/discussions/3389 '\n 'and read the docstring for more information.',\n DeprecationWarning,\n )\n\n features = self.out_features or inputs_q.shape[-1]\n qkv_features = self.qkv_features or inputs_q.shape[-1]\n assert qkv_features % self.num_heads == 0, (\n f'Memory dimension ({qkv_features}) must be divisible by number of'\n f' heads ({self.num_heads}).'\n )\n head_dim = qkv_features // self.num_heads\n\n dense = functools.partial(\n DenseGeneral,\n axis=-1,\n dtype=self.dtype,\n param_dtype=self.param_dtype,\n features=(self.num_heads, head_dim),\n kernel_init=self.kernel_init,\n bias_init=self.bias_init,\n use_bias=self.use_bias,\n precision=self.precision,\n dot_general=self.qkv_dot_general,\n dot_general_cls=self.qkv_dot_general_cls,\n )\n # project inputs_q to multi-headed q/k/v\n # dimensions are then [batch..., length, n_heads, n_features_per_head]\n query, key, value = (\n dense(name='query')(inputs_q),\n dense(name='key')(inputs_k),\n dense(name='value')(inputs_v),\n )\n\n if self.normalize_qk:\n # Normalizing query and key projections stabilizes training with higher\n # LR. See ViT-22B paper http://arxiv.org/abs/2302.05442 for analysis.\n query = LayerNorm(\n name='query_ln',\n use_bias=False,\n dtype=self.dtype,\n param_dtype=self.param_dtype,\n )(query) # type: ignore[call-arg]\n key = LayerNorm(\n name='key_ln',\n use_bias=False,\n dtype=self.dtype,\n param_dtype=self.param_dtype,\n )(key) # type: ignore[call-arg]\n\n # During fast autoregressive decoding, we feed one position at a time,\n # and cache the keys and values step by step.\n if self.decode:\n # detect if we're initializing by absence of existing cache data.\n is_initialized = self.has_variable('cache', 'cached_key')\n cached_key = self.variable(\n 'cache', 'cached_key', jnp.zeros, key.shape, key.dtype\n )\n cached_value = self.variable(\n 'cache', 'cached_value', jnp.zeros, value.shape, value.dtype\n )\n cache_index = self.variable(\n 'cache', 'cache_index', lambda: jnp.array(0, dtype=jnp.int32)\n )\n if is_initialized:\n (\n *batch_dims,\n max_length,\n num_heads,\n depth_per_head,\n ) = cached_key.value.shape\n # shape check of cached keys against query input\n expected_shape = tuple(batch_dims) + (1, num_heads, depth_per_head)\n if expected_shape != query.shape:\n raise ValueError(\n 'Autoregressive cache shape error, '\n 'expected query shape %s instead got %s.'\n % (expected_shape, query.shape)\n )\n # update key, value caches with our new 1d spatial slices\n cur_index = cache_index.value\n zero = jnp.array(0, dtype=lax.dtype(cur_index.dtype))\n indices: tuple[int | jax.Array, ...] = (zero,) * len(\n batch_dims\n ) + (\n cur_index,\n zero,\n zero,\n )\n key = lax.dynamic_update_slice(cached_key.value, key, indices)\n value = lax.dynamic_update_slice(cached_value.value, value, indices)\n cached_key.value = key\n cached_value.value = value\n cache_index.value = cache_index.value + 1\n # causal mask for cached decoder self-attention:\n # our single query position should only attend to those key\n # positions that have already been generated and cached,\n # not the remaining zero elements.\n mask = combine_masks(\n mask,\n jnp.broadcast_to(\n jnp.arange(max_length) <= cur_index,\n tuple(batch_dims) + (1, 1, max_length),\n ),\n )\n\n if (\n self.dropout_rate > 0.0\n ): # Require `deterministic` only if using dropout.\n m_deterministic = merge_param(\n 'deterministic', self.deterministic, deterministic\n )\n if not m_deterministic and dropout_rng is None:\n dropout_rng = self.make_rng('dropout')\n else:\n m_deterministic = True\n\n # `qk_attn_weights_einsum` and `attn_weights_value_einsum` are optional\n # arguments that can be used to override the default `jnp.einsum`. They\n # exist for quantized einsum support in AQT.\n qk_attn_weights_einsum = (\n self.qk_attn_weights_einsum_cls()\n if self.qk_attn_weights_einsum_cls\n else None\n )\n attn_weights_value_einsum = (\n self.attn_weights_value_einsum_cls()\n if self.attn_weights_value_einsum_cls\n else None\n )\n # apply attention\n attn_args = (query, key, value)\n # This kwargs list match the default nn.dot_product_attention.\n # For custom `attention_fn`s, invalid kwargs will be filtered.\n attn_kwargs = dict(\n mask=mask,\n dropout_rng=dropout_rng,\n dropout_rate=self.dropout_rate,\n broadcast_dropout=self.broadcast_dropout,\n deterministic=m_deterministic,\n dtype=self.dtype,\n precision=self.precision,\n force_fp32_for_softmax=self.force_fp32_for_softmax,\n qk_attn_weights_einsum=qk_attn_weights_einsum,\n attn_weights_value_einsum=attn_weights_value_einsum,\n )\n attn_kwargs = {\n k: v\n for k, v in attn_kwargs.items()\n if k in inspect.signature(self.attention_fn).parameters\n }\n if sow_weights:\n x = self.attention_fn(*attn_args, **attn_kwargs, module=self)\n else:\n x = self.attention_fn(*attn_args, **attn_kwargs)\n # back to the original inputs dimensions\n out = DenseGeneral(\n features=features,\n axis=(-2, -1),\n kernel_init=self.out_kernel_init or self.kernel_init,\n bias_init=self.out_bias_init or self.bias_init,\n use_bias=self.use_bias,\n dtype=self.dtype,\n param_dtype=self.param_dtype,\n precision=self.precision,\n dot_general=self.out_dot_general,\n dot_general_cls=self.out_dot_general_cls,\n name='out', # type: ignore[call-arg]\n )(x)\n return out\n\n\nclass MultiHeadAttention(MultiHeadDotProductAttention):\n \"\"\"Multi-head dot-product attention.\n Alias for ``MultiHeadDotProductAttention``.\n\n **NOTE**: ``MultiHeadAttention`` is a wrapper of ``MultiHeadDotProductAttention``,\n and so their implementations are identical. However ``MultiHeadAttention`` layers\n will, by default, be named ``MultiHeadAttention_{index}``, whereas ``MultiHeadDotProductAttention``\n will be named ``MultiHeadDotProductAttention_{index}``. Therefore, this could affect\n checkpointing, param collection names and RNG threading (since the layer name is\n used when generating new RNG's) within the module.\n\n Example usage::\n\n >>> import flax.linen as nn\n >>> import jax\n\n >>> layer = nn.MultiHeadAttention(num_heads=8, qkv_features=16)\n >>> key1, key2, key3, key4, key5, key6 = jax.random.split(jax.random.key(0), 6)\n >>> shape = (4, 3, 2, 5)\n >>> q, k, v = jax.random.uniform(key1, shape), jax.random.uniform(key2, shape), jax.random.uniform(key3, shape)\n >>> variables = layer.init(jax.random.key(0), q)\n\n >>> # different inputs for inputs_q, inputs_k and inputs_v\n >>> out = layer.apply(variables, q, k, v)\n >>> # equivalent to layer.apply(variables, inputs_q=q, inputs_k=k, inputs_v=k)\n >>> out = layer.apply(variables, q, k)\n >>> # equivalent to layer.apply(variables, inputs_q=q, inputs_k=q) and layer.apply(variables, inputs_q=q, inputs_k=q, inputs_v=q)\n >>> out = layer.apply(variables, q)\n\n >>> attention_kwargs = dict(\n ... num_heads=8,\n ... qkv_features=16,\n ... kernel_init=nn.initializers.ones,\n ... bias_init=nn.initializers.zeros,\n ... dropout_rate=0.5,\n ... deterministic=False,\n ... )\n >>> class Module(nn.Module):\n ... attention_kwargs: dict\n ...\n ... @nn.compact\n ... def __call__(self, x, dropout_rng=None):\n ... out1 = nn.MultiHeadAttention(**self.attention_kwargs)(x, dropout_rng=dropout_rng)\n ... out2 = nn.MultiHeadAttention(**self.attention_kwargs)(x, dropout_rng=dropout_rng)\n ... return out1, out2\n >>> module = Module(attention_kwargs)\n >>> variables = module.init({'params': key1, 'dropout': key2}, q)\n\n >>> # out1 and out2 are different.\n >>> out1, out2 = module.apply(variables, q, rngs={'dropout': key3})\n >>> # out3 and out4 are different.\n >>> # out1 and out3 are different. out2 and out4 are different.\n >>> out3, out4 = module.apply(variables, q, rngs={'dropout': key4})\n >>> # out1 and out2 are the same.\n >>> out1, out2 = module.apply(variables, q, dropout_rng=key5)\n >>> # out1 and out2 are the same as out3 and out4.\n >>> # providing a `dropout_rng` arg will take precedence over the `rngs` arg in `.apply`\n >>> out3, out4 = module.apply(variables, q, rngs={'dropout': key6}, dropout_rng=key5)\n\n Attributes:\n num_heads: number of attention heads. Features (i.e. inputs_q.shape[-1])\n should be divisible by the number of heads.\n dtype: the dtype of the computation (default: infer from inputs and params)\n param_dtype: the dtype passed to parameter initializers (default: float32)\n qkv_features: dimension of the key, query, and value.\n out_features: dimension of the last projection\n broadcast_dropout: bool: use a broadcasted dropout along batch dims.\n dropout_rate: dropout rate\n deterministic: if false, the attention weight is masked randomly using\n dropout, whereas if true, the attention weights are deterministic.\n precision: numerical precision of the computation see ``jax.lax.Precision``\n for details.\n kernel_init: initializer for the kernel of the Dense layers.\n bias_init: initializer for the bias of the Dense layers.\n use_bias: bool: whether pointwise QKVO dense transforms use bias.\n attention_fn: dot_product_attention or compatible function. Accepts query,\n key, value, and returns output of shape ``[bs, dim1, dim2, ..., dimN,,\n num_heads, value_channels]``\n decode: whether to prepare and use an autoregressive cache.\n normalize_qk: should QK normalization be applied (arxiv.org/abs/2302.05442).\n \"\"\"\n\n\nclass SelfAttention(MultiHeadDotProductAttention):\n \"\"\"Self-attention special case of multi-head dot-product attention.\n This layer is deprecated in favor of ``MultiHeadDotProductAttention``.\n\n Example usage::\n >>> import flax.linen as nn\n >>> import jax, jax.numpy as jnp\n >>> layer = nn.MultiHeadDotProductAttention(num_heads=8, qkv_features=16)\n >>> variables = layer.init(jax.random.key(0), jnp.ones((4, 3, 2, 5)))\n \"\"\"\n\n @compact\n def __call__( # type: ignore\n self,\n inputs_q: Array,\n mask: Array | None = None,\n deterministic: bool | None = None,\n dropout_rng: PRNGKey | None = None,\n sow_weights: bool = False,\n ):\n \"\"\"Applies multi-head dot product self-attention on the input data.\n\n Projects the inputs into multi-headed query, key, and value vectors,\n applies dot-product attention and project the results to an output vector.\n\n Args:\n inputs_q: input queries of shape ``[batch_sizes..., length, features]``.\n mask: attention mask of shape ``[batch_sizes..., num_heads, query_length,\n key/value_length]``. Attention weights are masked out if their\n corresponding mask value is ``False``.\n deterministic: if false, the attention weight is masked randomly using\n dropout, whereas if true, the attention weights are deterministic.\n\n Returns:\n output of shape ``[batch_sizes..., length, features]``.\n \"\"\"\n warnings.warn(\n 'SelfAttention will be deprecated soon. Use '\n '`MultiHeadDotProductAttention.__call__(inputs_q)` instead. '\n 'See https://github.com/google/flax/discussions/3389 '\n 'for more information.',\n DeprecationWarning,\n )\n return super().__call__(\n inputs_q,\n mask=mask,\n deterministic=deterministic,\n dropout_rng=dropout_rng,\n sow_weights=sow_weights,\n )\n\n\n# mask-making utility functions\n\n\ndef make_attention_mask(\n query_input: Array,\n key_input: Array,\n pairwise_fn: Callable[..., Any] = jnp.multiply,\n extra_batch_dims: int = 0,\n dtype: Dtype = jnp.float32,\n):\n \"\"\"Mask-making helper for attention weights.\n\n In case of 1d inputs (i.e., ``[batch..., len_q]``, ``[batch..., len_kv]``, the\n attention weights will be ``[batch..., heads, len_q, len_kv]`` and this\n function will produce ``[batch..., 1, len_q, len_kv]``.\n\n Args:\n query_input: a batched, flat input of query_length size\n key_input: a batched, flat input of key_length size\n pairwise_fn: broadcasting elementwise comparison function\n extra_batch_dims: number of extra batch dims to add singleton axes for, none\n by default\n dtype: mask return dtype\n\n Returns:\n A ``[batch..., 1, len_q, len_kv]`` shaped mask for 1d attention.\n \"\"\"\n mask = pairwise_fn(\n jnp.expand_dims(query_input, axis=-1), jnp.expand_dims(key_input, axis=-2)\n )\n mask = jnp.expand_dims(mask, axis=-3)\n mask = jnp.expand_dims(mask, axis=tuple(range(extra_batch_dims)))\n return mask.astype(dtype)\n\n\ndef make_causal_mask(\n x: Array, extra_batch_dims: int = 0, dtype: Dtype = jnp.float32\n) -> Array:\n \"\"\"Make a causal mask for self-attention.\n\n In case of 1d inputs (i.e., ``[batch..., len]``, the self-attention weights\n will be ``[batch..., heads, len, len]`` and this function will produce a\n causal mask of shape ``[batch..., 1, len, len]``.\n\n Args:\n x: input array of shape ``[batch..., len]``\n extra_batch_dims: number of batch dims to add singleton axes for, none by\n default\n dtype: mask return dtype\n\n Returns:\n A ``[batch..., 1, len, len]`` shaped causal mask for 1d attention.\n \"\"\"\n idxs = jnp.broadcast_to(jnp.arange(x.shape[-1], dtype=jnp.int32), x.shape)\n return make_attention_mask(\n idxs,\n idxs,\n jnp.greater_equal,\n extra_batch_dims=extra_batch_dims,\n dtype=dtype,\n )\n\n\ndef combine_masks(\n *masks: Array | None, dtype: Dtype = jnp.float32\n) -> Array | None:\n \"\"\"Combine attention masks.\n\n Args:\n *masks: set of attention mask arguments to combine, some can be None.\n dtype: dtype for the returned mask.\n\n Returns:\n Combined mask, reduced by logical and, returns None if no masks given.\n \"\"\"\n masks_list = [m for m in masks if m is not None]\n if not masks_list:\n return None\n assert all(\n map(lambda x: x.ndim == masks_list[0].ndim, masks_list)\n ), f'masks must have same rank: {tuple(map(lambda x: x.ndim, masks_list))}'\n mask, *other_masks = masks_list\n for other_mask in other_masks:\n mask = jnp.logical_and(mask, other_mask)\n return mask.astype(dtype)\n" }, { "path": "flax/linen/batch_apply.py", "content": "# Copyright 2023 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Batch apply.\"\"\"\n\nimport jax, jax.numpy as jnp\nimport numpy as np\n\n\ndef ndim_at_least(x, num_dims):\n if not (isinstance(x, jax.Array) or isinstance(x, np.ndarray)):\n x = jnp.asarray(x)\n return x.ndim >= num_dims\n\ndef arbitrary_mergeable_leaf(min_num_dims, args, kwargs):\n for a in jax.tree_util.tree_leaves(args):\n if ndim_at_least(a, min_num_dims):\n return a\n for k in jax.tree_util.tree_leaves(kwargs):\n if ndim_at_least(k, min_num_dims):\n return k\n # Couldn't find a satisfactory leaf.\n return None\n\ndef merge_leading_dims(x, num_dims):\n \"\"\"Merge leading dimensions.\"\"\"\n # Don't merge if there aren't dimensions to merge.\n if not ndim_at_least(x, num_dims):\n return x\n\n new_shape = (np.prod(x.shape[:num_dims]),) + x.shape[num_dims:]\n return x.reshape(new_shape)\n\ndef split_leading_dim(x, to_dim):\n new_shape = to_dim + x.shape[1:]\n return x.reshape(new_shape)\n\nclass BatchApply:\n r\"\"\"Temporarily merges leading dimensions of input tensors.\n\n Merges the leading dimensions of a tensor into a single dimension, runs the\n given callable, then splits the leading dimension of the result to match the\n input.\n\n Input arrays whose rank is smaller than the number of dimensions to collapse\n are passed unmodified.\n\n This may be useful for applying a module to each timestep of e.g. a\n ``[Time, Batch, ...]`` array.\n\n For some ``f``\\ s and platforms, this may be more efficient than\n :func:`jax.vmap`, especially when combined with other transformations like\n :func:`jax.grad`.\n\n Example usage::\n\n >>> import jax, jax.numpy as jnp\n\n >>> a = jax.random.normal(jax.random.key(0), [2, 3, 4])\n >>> b = jax.random.normal(jax.random.key(1), [4])\n\n >>> def raises(a, b):\n ... if len(a.shape) != 2:\n ... raise ValueError(\"a must be shape 2\")\n ... if len(b.shape) != 1:\n ... raise ValueError(\"b must be shape 1\")\n ... return jnp.dot(a, b)\n\n >>> out = BatchApply(raises)(a, b)\n >>> expected_merged_leading = raises(a.reshape(2*3, 4), b)\n >>> expected = expected_merged_leading.reshape((2, 3) + expected_merged_leading.shape[1:])\n >>> np.testing.assert_array_equal(out, expected)\n \"\"\"\n\n def __init__(self, f, num_dims=2):\n \"\"\"Constructs a :class:`BatchApply` module.\n\n Args:\n f: The callable to be applied to the reshaped array.\n num_dims: The number of dimensions to merge.\n \"\"\"\n self._f = f\n self.num_dims = num_dims\n\n def __call__(self, *args, **kwargs):\n example = arbitrary_mergeable_leaf(self.num_dims, args, kwargs)\n if example is None:\n raise ValueError(\n 'BatchApply requires at least one input with ndim >= '\n f'{self.num_dims}.'\n )\n\n merge = lambda x: merge_leading_dims(x, self.num_dims)\n split = lambda x: split_leading_dim(x, example.shape[:self.num_dims])\n args = jax.tree_util.tree_map(merge, args)\n kwargs = jax.tree_util.tree_map(merge, kwargs)\n outputs = self._f(*args, **kwargs)\n return jax.tree_util.tree_map(split, outputs)" }, { "path": "flax/linen/combinators.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Combinators of modules, such as a Sequential.\"\"\"\n\nfrom typing import Any\nfrom collections.abc import Callable, Sequence\n\nfrom flax.linen.module import Module, compact\n\n\nclass Sequential(Module):\n \"\"\"Applies a linear chain of Modules.\n\n Meant to be used only for the simple case of fusing together callables where\n the input of a particular module/op is the output of the previous one.\n\n Modules will be applied in the order that they are passed in the constructor.\n\n The ``__call__`` method of Sequential accepts any input and forwards it to the\n first module it contains. It chains the output sequentially to the input of\n the next module and returns the output of the final module.\n\n Example usage::\n\n >>> import flax.linen as nn\n\n >>> class Foo(nn.Module):\n ... @nn.compact\n ... def __call__(self, x):\n ... return nn.Sequential([nn.Dense(4),\n ... nn.relu,\n ... nn.Dense(2),\n ... nn.log_softmax])(x)\n\n Since `Sequential.__call__` is a `compact` method, you can also pass functions\n that construct Modules inline if you need shape inference::\n\n module = nn.Sequential([\n # << more layers\n lambda x: SomeModule(x.shape[-1])(x), # shape inference\n # << more layers\n ])\n\n This combinator supports also layers that return multiple outputs if returned\n as a tuple or a dictionary. If the output of a layer is a ``tuple`` it will be\n expanded as ``*args`` in the next layer, if its a ``dict`` it\n will be expanded as ``**kwargs``.\n\n Example usage::\n\n >>> class CrossAttentionBlock(nn.Module):\n ... num_heads: int = 2\n ... qkv_features: int = 16\n ...\n ... @nn.compact\n ... def __call__(self, query, key_value):\n ... output = nn.MultiHeadDotProductAttention(\n ... num_heads=self.num_heads, qkv_features=self.qkv_features)(query,\n ... key_value)\n ... output = nn.Dense(self.qkv_features)(output)\n ... return dict(query=output, key_value=key_value) # also works for tuples\n\n >>> from typing import Sequence\n >>> class CrossAttentionNetwork(nn.Module):\n ... num_layers: Sequence[int]\n ...\n ... @nn.compact\n ... def __call__(self, x):\n ... return nn.Sequential([CrossAttentionBlock() for _ in\n ... range(self.num_layers)])(query, key_value)\n\n\n Attributes:\n layers: A sequence of callables to be applied in order.\n\n Raises:\n ValueError: If layers is not a sequence.\n \"\"\"\n\n layers: Sequence[Callable[..., Any]]\n\n def __post_init__(self):\n if not isinstance(self.layers, Sequence):\n raise ValueError(\n f\"'layers' must be a sequence, got '{type(self.layers).__name__}'.\"\n )\n super().__post_init__()\n\n @compact\n def __call__(self, *args, **kwargs):\n if not self.layers:\n raise ValueError(f'Empty Sequential module {self.name}.')\n\n outputs = self.layers[0](*args, **kwargs)\n for layer in self.layers[1:]:\n if isinstance(outputs, tuple):\n outputs = layer(*outputs)\n elif isinstance(outputs, dict):\n outputs = layer(**outputs)\n else:\n outputs = layer(outputs)\n return outputs\n" }, { "path": "flax/linen/dtypes.py", "content": "# Copyright 2022 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"APIs for handling dtypes in Linen Modules.\"\"\"\n\nfrom typing import Any, TypeVar\nfrom flax.typing import Dtype\nfrom jax import numpy as jnp\n\nT = TypeVar('T', bound=tuple)\n\ndef canonicalize_dtype(\n *args, dtype: Dtype | None = None, inexact: bool = True\n) -> Dtype:\n \"\"\"Canonicalize an optional dtype to the definitive dtype.\n\n If the ``dtype`` is None this function will infer the dtype. If it is not\n None it will be returned unmodified or an exceptions is raised if the dtype\n is invalid.\n from the input arguments using ``jnp.result_type``.\n\n Args:\n *args: JAX array compatible values. None values\n are ignored.\n dtype: Optional dtype override. If specified the arguments are cast to\n the specified dtype instead and dtype inference is disabled.\n inexact: When True, the output dtype must be a subdtype\n of `jnp.inexact`. Inexact dtypes are real or complex floating points. This\n is useful when you want to apply operations that don't work directly on\n integers like taking a mean for example.\n Returns:\n The dtype that *args should be cast to.\n \"\"\"\n if dtype is None:\n args_filtered = [jnp.asarray(x) for x in args if x is not None]\n dtype = jnp.result_type(*args_filtered)\n if inexact and not jnp.issubdtype(dtype, jnp.inexact):\n dtype = jnp.promote_types(jnp.float32, dtype)\n if inexact and not jnp.issubdtype(dtype, jnp.inexact):\n raise ValueError(f'Dtype must be inexact: {dtype}')\n return dtype\n\n\ndef promote_dtype(*args, dtype=None, inexact=True) -> list[Any]:\n \"\"\" \"Promotes input arguments to a specified or inferred dtype.\n\n All args are cast to the same dtype. See ``canonicalize_dtype`` for how\n this dtype is determined.\n\n The behavior of promote_dtype is mostly a convinience wrapper around\n ``jax.numpy.promote_types``. The differences being that it automatically casts\n all input to the inferred dtypes, allows inference to be overridden by a\n forced dtype, and has an optional check to garantuee the resulting dtype is\n inexact.\n\n Args:\n *args: JAX array compatible values. None values are returned as is.\n dtype: Optional dtype override. If specified the arguments are cast to the\n specified dtype instead and dtype inference is disabled.\n inexact: When True, the output dtype must be a subdtype of `jnp.inexact`.\n Inexact dtypes are real or complex floating points. This is useful when\n you want to apply operations that don't work directly on integers like\n taking a mean for example.\n\n Returns:\n The arguments cast to arrays of the same dtype.\n \"\"\"\n dtype = canonicalize_dtype(*args, dtype=dtype, inexact=inexact)\n return [jnp.asarray(x, dtype) if x is not None else None for x in args]\n" }, { "path": "flax/linen/experimental/layers_with_named_axes.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Experimental layers with named axes for the partitioning API.\"\"\"\nimport dataclasses\nfrom typing import Any\nfrom collections.abc import Callable, Iterable, Sequence\n\nimport jax.numpy as jnp\nfrom jax import lax\n\nfrom flax import linen as nn\nfrom flax.linen import initializers\nfrom flax.linen.partitioning import param_with_axes, with_sharding_constraint\nfrom flax.typing import (\n Array,\n Dtype,\n Axes,\n Initializer,\n PrecisionLike,\n DotGeneralT,\n)\n\n# Type annotations\nActivation = Callable[..., Array]\n\n\ndefault_kernel_init = initializers.lecun_normal()\ndefault_embed_init = initializers.variance_scaling(\n 1.0, 'fan_in', 'normal', out_axis=0\n)\n\n\nclass Dense(nn.Module):\n \"\"\"A Dense layer with named axes for :meth:`jax.experimental.pjit.pjit`.\n\n .. warning:: This class is hightly EXPERIMENTAL and the API is likely to\n change. For regular (non-pjit) use, please use\n :class:`flax.linen.linear.Dense`.\n\n Attributes:\n features: the number of output features.\n use_bias: whether to add a bias to the output (default: True).\n dtype: the dtype of the computation (default: float32).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n precision: numerical precision of the computation see `jax.lax.Precision`\n for details.\n kernel_init: initializer function for the weight matrix.\n bias_init: initializer function for the bias.\n \"\"\"\n\n features: int\n use_bias: bool = True\n dtype: Dtype = jnp.float32\n param_dtype: Dtype = jnp.float32\n precision: PrecisionLike = None\n kernel_init: Initializer = default_kernel_init\n bias_init: Initializer = initializers.zeros_init()\n kernel_axes: tuple[str, ...] = ()\n # Deprecated. Will be removed.\n dot_general: DotGeneralT | None = None\n dot_general_cls: Any = None\n\n @nn.compact\n def __call__(self, inputs: Array) -> Array:\n \"\"\"Applies a linear transformation to the inputs along the last dimension.\n\n Args:\n inputs: The nd-array to be transformed.\n\n Returns:\n The transformed input.\n \"\"\"\n inputs = jnp.asarray(inputs, self.dtype)\n kernel = param_with_axes(\n 'kernel',\n self.kernel_init,\n (inputs.shape[-1], self.features),\n self.param_dtype,\n axes=self.kernel_axes,\n )\n kernel = jnp.asarray(kernel, self.dtype)\n\n if self.dot_general_cls is not None:\n dot_general = self.dot_general_cls()\n elif self.dot_general is not None:\n dot_general = self.dot_general\n else:\n dot_general = lax.dot_general\n y = dot_general(\n inputs,\n kernel,\n (((inputs.ndim - 1,), (0,)), ((), ())),\n precision=self.precision,\n )\n if self.use_bias:\n bias = param_with_axes(\n 'bias',\n self.bias_init,\n (self.features,),\n self.param_dtype,\n axes=(self.kernel_axes[-1],),\n )\n bias = jnp.asarray(bias, self.dtype)\n y += jnp.reshape(bias, (1,) * (y.ndim - 1) + (-1,))\n return y\n\n\nclass Embed(nn.Module):\n \"\"\"An embedding layer with named axes for :meth:`jax.experimental.pjit.pjit`.\n\n .. warning:: This class is hightly EXPERIMENTAL and the API is likely to\n change. For regular (non-pjit) use, please use\n :class:`flax.linen.linear.Embed`.\n\n Attributes:\n num_embeddings: number of embeddings.\n features: number of feature dimensions for each embedding.\n dtype: the dtype of the embedding vectors (default: float32).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n embedding_init: embedding initializer.\n one_hot: performs the gather with a one-hot contraction rather than a true\n gather. This is currently needed for SPMD partitioning.\n \"\"\"\n\n num_embeddings: int\n features: int\n cast_input_dtype: Dtype | None = None\n dtype: Dtype = jnp.float32\n param_dtype: Dtype = jnp.float32\n attend_dtype: Dtype | None = None\n embedding_init: Initializer = default_embed_init\n one_hot: bool = False\n embedding: Array = dataclasses.field(init=False)\n\n def setup(self):\n self.embedding = param_with_axes(\n 'embedding',\n self.embedding_init,\n (self.num_embeddings, self.features),\n self.param_dtype,\n axes=('vocab', 'embed'),\n )\n\n def __call__(self, inputs: Array) -> Array:\n \"\"\"Embeds the inputs along the last dimension.\n\n Args:\n inputs: input data, all dimensions are considered batch dimensions.\n\n Returns:\n Output which is embedded input data. The output shape follows the input,\n with an additional `features` dimension appended.\n \"\"\"\n if self.cast_input_dtype:\n inputs = inputs.astype(self.cast_input_dtype)\n if not jnp.issubdtype(inputs.dtype, jnp.integer):\n raise ValueError('Input type must be an integer or unsigned integer.')\n if self.one_hot:\n iota = lax.iota(jnp.int32, self.num_embeddings)\n one_hot = jnp.array(inputs[..., jnp.newaxis] == iota, dtype=self.dtype)\n output = jnp.dot(one_hot, jnp.asarray(self.embedding, self.dtype))\n else:\n output = jnp.asarray(self.embedding, self.dtype)[inputs]\n output = with_sharding_constraint(output, ('batch', 'length', 'embed'))\n return output\n\n def attend(self, query: Array) -> Array:\n \"\"\"Attend over the embedding using a query array.\n\n Args:\n query: array with last dimension equal the feature depth `features` of the\n embedding.\n\n Returns:\n An array with final dim `num_embeddings` corresponding to the batched\n inner-product of the array of query vectors against each embedding.\n Commonly used for weight-sharing between embeddings and logit transform\n in NLP models.\n \"\"\"\n dtype = self.attend_dtype if self.attend_dtype is not None else self.dtype\n return jnp.dot(query, jnp.asarray(self.embedding, dtype).T)\n\n\ndef _canonicalize_axes(rank: int, axes: Axes) -> Sequence[int]:\n \"\"\"Returns a tuple of deduplicated, sorted, and positive axes.\"\"\"\n if not isinstance(axes, Iterable):\n axes = (axes,)\n return tuple({rank + axis if axis < 0 else axis for axis in axes})\n\n\ndef _abs_sq(x):\n \"\"\"Computes the elementwise square of the absolute value |x|^2.\"\"\"\n if jnp.iscomplexobj(x):\n return lax.square(lax.real(x)) + lax.square(lax.imag(x))\n else:\n return lax.square(x)\n\n\ndef _compute_stats(x: Array, axes: Axes):\n \"\"\"Computes mean and variance statistics.\n\n This implementation takes care of a few important details:\n - Computes in float32 precision for half precision inputs\n - mean and variance is computable in a single XLA fusion,\n by using Var = E[|x|^2] - |E[x]|^2 instead of Var = E[|x - E[x]|^2]).\n - Clips negative variances to zero which can happen due to\n roundoff errors. This avoids downstream NaNs.\n - Supports averaging across a parallel axis and subgroups of a parallel axis\n with a single `lax.pmean` call to avoid latency.\n \"\"\"\n # promote x to at least float32, this avoids half precision computation\n # but preserves double or complex floating points\n x = jnp.asarray(x, jnp.promote_types(jnp.float32, jnp.result_type(x)))\n mean = jnp.mean(x, axes)\n mean2 = jnp.mean(_abs_sq(x), axes)\n # mean2 - _abs_sq(mean) is not guaranteed to be non-negative due\n # to floating point round-off errors.\n var = jnp.maximum(0.0, mean2 - _abs_sq(mean))\n return mean, var\n\n\ndef _normalize(\n mdl: nn.Module,\n x: Array,\n mean: Array,\n var: Array,\n reduction_axes: Axes,\n feature_axes: Axes,\n dtype: Dtype,\n param_dtype: Dtype,\n epsilon: float,\n use_bias: bool,\n use_scale: bool,\n bias_init: Initializer,\n scale_init: Initializer,\n):\n \"\"\" \"Normalizes the input of a normalization layer and optionally applies a learned scale and bias.\n\n A separate bias and scale is learned for each feature as specified by\n feature_axes.\n \"\"\"\n reduction_axes = _canonicalize_axes(x.ndim, reduction_axes)\n feature_axes = _canonicalize_axes(x.ndim, feature_axes)\n stats_shape = list(x.shape)\n for axis in reduction_axes:\n stats_shape[axis] = 1\n mean = mean.reshape(stats_shape)\n var = var.reshape(stats_shape)\n feature_shape = [1] * x.ndim\n reduced_feature_shape = []\n for ax in feature_axes:\n feature_shape[ax] = x.shape[ax]\n reduced_feature_shape.append(x.shape[ax])\n y = x - mean\n mul = lax.rsqrt(var + epsilon)\n if use_scale:\n scale = mdl.param_with_axes(\n 'scale', scale_init, reduced_feature_shape, param_dtype, axes=('embed',)\n ).reshape(feature_shape)\n mul *= scale\n y *= mul\n if use_bias:\n bias = mdl.param_with_axes(\n 'bias', bias_init, reduced_feature_shape, param_dtype, axes=('embed',)\n ).reshape(feature_shape)\n y += bias\n return jnp.asarray(y, dtype)\n\n\nclass LayerNorm(nn.Module):\n \"\"\"Layer normalization (https://arxiv.org/abs/1607.06450) with named axes for :meth:`jax.experimental.pjit.pjit`.\n\n .. warning:: This class is hightly EXPERIMENTAL and the API is likely to\n change. For regular (non-pjit) use, please use\n :class:`flax.linen.normalization.LayerNorm`.\n\n Operates on the last axis of the input data.\n\n It normalizes the activations of the layer for each given example in a\n batch independently, rather than across a batch like Batch Normalization.\n i.e. applies a transformation that maintains the mean activation within\n each example close to 0 and the activation standard deviation close to 1.\n\n Attributes:\n epsilon: A small float added to variance to avoid dividing by zero.\n dtype: the dtype of the computation (default: float32).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n use_bias: If True, bias (beta) is added.\n use_scale: If True, multiply by scale (gamma). When the next layer is linear\n (also e.g. nn.relu), this can be disabled since the scaling will be done\n by the next layer.\n bias_init: Initializer for bias, by default, zero.\n scale_init: Initializer for scale, by default, one.\n \"\"\"\n\n epsilon: float = 1e-6\n dtype: Any = jnp.float32\n param_dtype: Dtype = jnp.float32\n use_bias: bool = True\n use_scale: bool = True\n bias_init: Initializer = initializers.zeros_init()\n scale_init: Initializer = initializers.ones_init()\n\n @nn.compact\n def __call__(self, x):\n \"\"\"Applies layer normalization on the input.\n\n Args:\n x: the inputs\n\n Returns:\n Normalized inputs (the same shape as inputs).\n \"\"\"\n reduction_axes = (-1,)\n feature_axes = (-1,)\n\n mean, var = _compute_stats(x, reduction_axes)\n\n return _normalize(\n self,\n x,\n mean,\n var,\n reduction_axes,\n feature_axes,\n self.dtype,\n self.param_dtype,\n self.epsilon,\n self.use_bias,\n self.use_scale,\n self.bias_init,\n self.scale_init,\n )\n" }, { "path": "flax/linen/fp8_ops.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport dataclasses\nimport itertools\nimport numpy as np\nimport warnings\nfrom functools import partial\n\nfrom typing import Any\nDType = Any\n\nimport jax\nfrom jax import custom_jvp, custom_vjp, lax, random\nfrom jax import numpy as jnp\nfrom jax._src import core\nfrom jax._src import dtypes\nfrom jax.typing import DTypeLike\n\ntry:\n from jax._src import earray\n from jax._src.interpreters import pxla\n CAN_USE_EARRAY = True\nexcept (ModuleNotFoundError, ImportError):\n CAN_USE_EARRAY = False\n\nfrom flax.linen import initializers, module\n\nOVERWRITE_WITH_GRADIENT = '_overwrite_with_gradient'\n\n# Define a custom dtype for FP8 meta params.\nclass Fp8MetaTyRules:\n # tell JAX how to lower this dtype to an HLO dtype\n @staticmethod\n def physical_element_aval(dtype) -> core.ShapedArray:\n return core.ShapedArray((), dtype.float_dtype)\n\n if jax.__version_info__ < (0, 4, 29):\n @staticmethod\n def replicate_trailing_dims(ctx, val, aval):\n del ctx, aval\n return val\n\n @staticmethod\n def logical_sharding(aval, phys_sharding):\n return phys_sharding\n\n @staticmethod\n def physical_sharding(aval, sharding):\n return sharding # unlike KeyTyRules, assume same shape\n\n # allow conversions to and from the corresponding float type\n @staticmethod\n def convert_from(fp8_meta_dtype, other_dtype) -> bool:\n return fp8_meta_dtype.float_dtype == other_dtype\n\n @staticmethod\n def convert_to(other_dtype, fp8_meta_dtype) -> bool:\n return fp8_meta_dtype.float_dtype == other_dtype\n\n # define how autodiff should accumulate these values\n @staticmethod\n def add(dt, x, y):\n from_fp8_meta = partial(lax.convert_element_type, new_dtype=dt.float_dtype)\n to_fp8_meta = partial(lax.convert_element_type, new_dtype=dt)\n return to_fp8_meta(lax.max(from_fp8_meta(x), from_fp8_meta(y)))\n\n @staticmethod\n def zero(dt):\n neginf = np.array(-np.inf if dtypes.supports_inf(dt.float_dtype)\n else dtypes.finfo(dt.float_dtype).min, dt.float_dtype)\n return lax.convert_element_type(neginf, dt)\n\n @staticmethod\n def tangent_dtype(dtype):\n return dtype\n\n @staticmethod\n def full(shape, fill_value, dtype):\n fill_value = lax.convert_element_type(fill_value, dtype.float_dtype)\n out_raw = lax.full(shape, fill_value, dtype.float_dtype)\n return lax.convert_element_type(out_raw, dtype)\n\n @staticmethod\n def global_sharded_result_handler(aval, out_sharding, committed):\n if not CAN_USE_EARRAY:\n raise NotImplementedError(\"convert back under the jit\")\n\n phys_sharding = out_sharding # unlike KeyTyRules, assume same shape\n phys_aval = core.physical_aval(aval)\n phys_handler_maker = pxla.global_result_handlers[core.ShapedArray]\n phys_handler = phys_handler_maker(phys_aval, phys_sharding, committed)\n return lambda bufs: earray.EArray(aval, phys_handler(bufs))\n\n\n# class to use as second argument to jax.dtypes.issubdtype\nclass fp8_meta_dtype(dtypes.extended): pass\n\n# parameterized datatype for use in e.g. lax.convert_element_type\n@dataclasses.dataclass(frozen=True)\nclass fp8_meta_dtype_wrapper(dtypes.ExtendedDType):\n float_dtype: dtypes.DType\n _rules: type = Fp8MetaTyRules\n type: type = fp8_meta_dtype\n\n def __repr__(self) -> str:\n nbits = dtypes.finfo(self.float_dtype).bits\n return f'fp8_meta{nbits}'\n name = property(__repr__)\n\nfm32 = fp8_meta_dtype_wrapper(jnp.float32)\nfp32_max_grad = fp8_meta_dtype_wrapper(jnp.float32)\n\ndef get_fp8_max(fp8_dtype, out_dtype):\n assert fp8_dtype in (jnp.float8_e4m3fn, jnp.float8_e5m2,\n jnp.float8_e4m3fnuz, jnp.float8_e5m2fnuz)\n return jnp.finfo(fp8_dtype).max.astype(out_dtype)\n\ndef quantize(x, q_dtype, scale, compute_dtype):\n # Explicitly cast the max values to the compute dtype to avoid unnecessary\n # casting to FP32 during the subsequent math operations.\"\n dtype_max = get_fp8_max(q_dtype, compute_dtype)\n scaled_x = x / jnp.broadcast_to(scale.astype(compute_dtype), x.shape)\n clipped_x = jnp.clip(scaled_x, -dtype_max, dtype_max)\n return clipped_x.astype(q_dtype)\n\n\ndef dequantize(x, dq_dtype, scale):\n return x.astype(dq_dtype) * jnp.broadcast_to(scale.astype(dq_dtype), x.shape)\n\ndef qdq(x, q_dtype, scale, compute_dtype):\n qx = quantize(x, q_dtype, scale, compute_dtype)\n return dequantize(qx, x.dtype, scale)\n\n\ndef compute_scale(amax, scale, fp8_max, margin=0):\n # The algorithm for computing the new scale is sourced from\n # https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/api/jax.html#transformer_engine.jax.update_fp8_metas\n # wherein the `original_scale` corresponds to the reciprocal of the `scale`\n # passed in this function.\n scale = 1.0 / scale\n\n sf = (fp8_max / amax) / (2**margin)\n sf = jnp.where(amax > 0.0, sf, scale)\n sf = jnp.where(jnp.isfinite(amax), sf, scale)\n\n return 1.0 / sf\n\n\ndef compute_amax_history(x, amax_history):\n amax_update = jnp.max(jnp.abs(x)).astype(amax_history.dtype)\n new_history = jnp.roll(amax_history, shift=-1, axis=0).at[0].set(amax_update)\n return new_history\n\n\ndef update_fp8_meta(\n x, q_dtype, scale, amax_history\n):\n is_fmax32 = (scale.dtype == fm32 and amax_history.dtype == fm32)\n # convert fm32->f32 so we can do math\n if is_fmax32:\n amax_history = lax.convert_element_type(amax_history, jnp.float32)\n scale = lax.convert_element_type(scale, jnp.float32)\n\n # Update the fp8 meta\n dtype_max = get_fp8_max(q_dtype, jnp.float32)\n amax_from_history = jnp.max(amax_history, axis=0)\n\n new_scale = compute_scale(amax_from_history, scale, dtype_max)\n new_history = compute_amax_history(x, amax_history)\n\n if is_fmax32:\n new_history = lax.convert_element_type(new_history, fp32_max_grad)\n new_scale = lax.convert_element_type(new_scale, fp32_max_grad)\n return new_scale, new_history\n\ndef quantize_dequantize_update(x, q_dtype, scale, amax_history, compute_dtype):\n updated_scale, updated_history = update_fp8_meta(x, q_dtype, scale, amax_history)\n qdq_x = qdq(x, q_dtype, _fm32_to_float32(updated_scale), compute_dtype)\n return qdq_x, updated_scale, updated_history\n\ndef _fm32_to_float32(value):\n if value.dtype == fm32:\n return lax.convert_element_type(value, jnp.float32)\n return value\n\ndef dot_general_transpose_lhs(g, x, y, *, dimension_numbers, precision,\n preferred_element_type: DTypeLike | None,\n swap_ans=False):\n def _remaining(original, *removed_lists):\n removed = set(itertools.chain(*removed_lists))\n return [i for i in original if i not in removed]\n\n def _ranges_like(*xs):\n start = 0\n for x in xs:\n x_len = len(x)\n yield range(start, start + x_len)\n start += x_len\n\n (x_contract, y_contract), (x_batch, y_batch) = dimension_numbers\n x_ndim = x.aval.ndim\n x_kept = _remaining(range(x_ndim), x_contract, x_batch)\n y_kept = _remaining(range(np.ndim(y)), y_contract, y_batch)\n if swap_ans:\n ans_batch, ans_y, _ = _ranges_like(x_batch, y_kept, x_kept)\n else:\n ans_batch, _, ans_y = _ranges_like(x_batch, x_kept, y_kept)\n dims = ((ans_y, y_kept), (ans_batch, y_batch))\n x_contract_sorted_by_y = list(np.take(x_contract, np.argsort(y_contract)))\n out_axes = np.argsort(list(x_batch) + x_kept + x_contract_sorted_by_y)\n x_bar = lax.transpose(\n lax.dot_general(\n g, y, dims, precision=precision,\n preferred_element_type=preferred_element_type\n ),\n tuple(out_axes)\n )\n return x_bar\n\ndef dot_general_transpose_rhs(g, x, y, *, dimension_numbers, precision,\n preferred_element_type: DTypeLike | None):\n (x_contract, y_contract), (x_batch, y_batch) = dimension_numbers\n swapped_dimension_numbers = ((y_contract, x_contract), (y_batch, x_batch))\n y_bar = dot_general_transpose_lhs(\n g, y, x, dimension_numbers=swapped_dimension_numbers, precision=precision,\n preferred_element_type=preferred_element_type,\n swap_ans=True)\n return y_bar\n\n@partial(custom_vjp, nondiff_argnums=(0, 1))\ndef in_qdq(compute_dtype, q_dtype, inp, scale, amax_history):\n qin, _, _ = quantize_dequantize_update(\n inp, q_dtype, scale, amax_history, compute_dtype\n )\n return qin\n\n\ndef in_qdq_fwd(compute_dtype, q_dtype, inp, scale, amax_history):\n qin, new_scale, new_history = quantize_dequantize_update(\n inp, q_dtype, scale, amax_history, compute_dtype\n )\n return qin, (new_scale, new_history)\n\n\ndef in_qdq_bwd(compute_dtype, q_dtype, res, g):\n new_scale, new_history = res\n q_g = g\n return q_g, new_scale, new_history\n\n\nin_qdq.defvjp(in_qdq_fwd, in_qdq_bwd)\n\n\n@partial(custom_vjp, nondiff_argnums=(0, 1))\ndef out_qdq(compute_dtype, q_dtype, out, scale, amax_history):\n return out\n\n\ndef out_qdq_fwd(compute_dtype, q_dtype, out, scale, amax_history):\n return out, (scale, amax_history)\n\n\ndef out_qdq_bwd(compute_dtype, q_dtype, res, g):\n scale, amax_history = res\n q_g, new_scale, new_history = quantize_dequantize_update(\n g, q_dtype, scale, amax_history, compute_dtype\n )\n return q_g, new_scale, new_history\n\n\nout_qdq.defvjp(out_qdq_fwd, out_qdq_bwd)\n\n\n@partial(custom_vjp, nondiff_argnums=(0, 1))\ndef in_q(compute_dtype, q_dtype, inp, scale, amax_history):\n new_scale, _ = update_fp8_meta(inp, q_dtype, scale, amax_history)\n qin = quantize(inp, q_dtype, _fm32_to_float32(new_scale), compute_dtype)\n return qin, new_scale\n\ndef in_q_fwd(compute_dtype, q_dtype, inp, scale, amax_history):\n new_scale, new_history = update_fp8_meta(inp, q_dtype, scale, amax_history)\n qin = quantize(inp, q_dtype, _fm32_to_float32(new_scale), compute_dtype)\n return (qin, new_scale), (new_scale, new_history)\n\ndef in_q_bwd(compute_dtype, q_dtype, res, _):\n new_scale, new_history = res\n # We don't compute gradients for inp, scale and amax_history, but we pass through scale and history\n return None, new_scale, new_history\n\nin_q.defvjp(in_q_fwd, in_q_bwd)\n\n\n@partial(custom_vjp, nondiff_argnums=(0, ))\ndef out_dq(dq_type, lhs_scale, rhs_scale, out):\n q_out = dequantize(\n out,\n dq_type,\n _fm32_to_float32(lhs_scale) * _fm32_to_float32(rhs_scale)\n )\n return q_out\n\ndef out_dq_fwd(dq_type, lhs_scale, rhs_scale, out):\n return out_dq(dq_type, lhs_scale, rhs_scale, out), None\n\ndef out_dq_bwd(dq_type, _, g):\n return None, None, g\n\nout_dq.defvjp(out_dq_fwd, out_dq_bwd)\n\n@partial(custom_vjp, nondiff_argnums=(8, 9))\ndef quantized_dot(\n lhs,\n q_lhs,\n lhs_scale, # scale for this step\n rhs,\n q_rhs,\n rhs_scale, # scale for this step\n out_grad_scale, # scale from previous step\n out_grad_amax_history, # amax history from previous step\n dimension_numbers,\n preferred_element_type=None\n):\n return lax.dot_general(\n q_lhs,\n q_rhs,\n dimension_numbers,\n preferred_element_type=preferred_element_type,\n precision=lax.Precision.DEFAULT,\n )\n\ndef quantized_dot_fwd(\n lhs,\n q_lhs,\n lhs_scale,\n rhs,\n q_rhs,\n rhs_scale,\n out_grad_scale,\n out_grad_amax_history,\n dimension_numbers,\n preferred_element_type,\n):\n out = lax.dot_general(\n q_lhs,\n q_rhs,\n dimension_numbers,\n preferred_element_type=preferred_element_type,\n precision=lax.Precision.DEFAULT,\n )\n res = (\n lhs,\n q_lhs,\n lhs_scale,\n rhs,\n q_rhs,\n rhs_scale,\n out_grad_scale,\n out_grad_amax_history,\n )\n return out, res\n\ndef quantized_dot_bwd(\n dimension_numbers,\n preferred_element_type,\n res,\n g\n):\n (\n lhs,\n q_lhs,\n lhs_scale,\n rhs,\n q_rhs,\n rhs_scale,\n out_grad_scale,\n out_grad_amax_history,\n ) = res\n\n new_out_grad_scale, new_out_grad_amax_history = update_fp8_meta(\n g,\n jnp.float8_e5m2,\n out_grad_scale,\n out_grad_amax_history,\n )\n\n q_g = quantize(\n g,\n jnp.float8_e5m2,\n _fm32_to_float32(new_out_grad_scale),\n preferred_element_type\n )\n\n grad_lhs = dot_general_transpose_lhs(\n q_g,\n lhs,\n q_rhs,\n dimension_numbers=dimension_numbers,\n precision=lax.Precision.HIGHEST,\n preferred_element_type=preferred_element_type,\n )\n\n grad_lhs = dequantize(\n grad_lhs,\n preferred_element_type,\n _fm32_to_float32(rhs_scale) * _fm32_to_float32(new_out_grad_scale)\n )\n\n grad_rhs = dot_general_transpose_rhs(\n q_g,\n q_lhs,\n rhs,\n dimension_numbers=dimension_numbers,\n precision=lax.Precision.HIGHEST,\n preferred_element_type=preferred_element_type,\n )\n\n grad_rhs = dequantize(\n grad_rhs,\n preferred_element_type,\n _fm32_to_float32(lhs_scale) * _fm32_to_float32(new_out_grad_scale)\n )\n\n return (\n grad_lhs,\n None,\n None,\n grad_rhs,\n None,\n None,\n new_out_grad_scale,\n new_out_grad_amax_history,\n )\n\nquantized_dot.defvjp(quantized_dot_fwd, quantized_dot_bwd)\n\n# Wrapper function to achieve the same effect as the dot_general function\n# but with fp8 quantization and dequantization.\ndef fp8_scaled_dot_general(\n lhs,\n rhs,\n dimension_numbers,\n precision=None,\n preferred_element_type=None,\n *,\n lhs_scale=None,\n rhs_scale=None,\n grad_scale=None,\n lhs_amax_history=None,\n rhs_amax_history=None,\n grad_amax_history=None,\n quantize_compute_type=jnp.float32,\n):\n if precision != None:\n warnings.warn(\n 'The function fp8_scaled_dot_general will set the \"precision\" and '\n 'disregard any provided \"precision\" argument.'\n )\n q_lhs, new_lhs_scale = in_q(\n quantize_compute_type, jnp.float8_e4m3fn, lhs, lhs_scale, lhs_amax_history\n )\n q_rhs, new_rhs_scale = in_q(\n quantize_compute_type, jnp.float8_e4m3fn, rhs, rhs_scale, rhs_amax_history\n )\n y = quantized_dot(\n lhs,\n q_lhs,\n new_lhs_scale,\n rhs,\n q_rhs,\n new_rhs_scale,\n grad_scale,\n grad_amax_history,\n dimension_numbers,\n preferred_element_type\n )\n y = out_dq(\n dq_type=preferred_element_type,\n lhs_scale=new_lhs_scale,\n rhs_scale=new_rhs_scale,\n out=y\n )\n return y # type: ignore\n\n@partial(custom_jvp, nondiff_argnums=(2, 3, 4))\ndef dot_general_with_precision(\n lhs, rhs, dimension_numbers, precision=None, preferred_element_type=None\n):\n if precision != None or preferred_element_type != None:\n warnings.warn(\n 'The function dot_general_with_precision will set the '\n 'precision/preferred_element_type and disregard any provided '\n 'values.'\n )\n return lax.dot_general(\n lhs, rhs, dimension_numbers, precision=lax.Precision.DEFAULT\n )\n\n\n@dot_general_with_precision.defjvp\ndef dot_general_with_precision_jvp(\n dimension_numbers, precision, preferred_element_type, primals, tangents\n):\n lhs, rhs = primals\n lhs_dot, rhs_dot = tangents\n\n out = lax.dot_general(\n lhs, rhs, dimension_numbers, precision=lax.Precision.DEFAULT\n )\n grad_out = lax.dot_general(\n lhs_dot, rhs, dimension_numbers, precision=lax.Precision.HIGHEST\n ) + lax.dot_general(\n lhs, rhs_dot, dimension_numbers, precision=lax.Precision.HIGHEST\n )\n return out, grad_out\n\n\ndef _parse_dot_inputs(*args, **kwargs):\n assert len(args) == 3\n x = args[0]\n k = args[1]\n dimension_numbers = args[2]\n\n # Use the `k.dtype` since it aligns with the `dtype` of its layers,\n # namely, the computation data type.\n comp_dtype = k.dtype\n x = jnp.asarray(x, comp_dtype)\n return x, k, dimension_numbers, comp_dtype\n\n\nclass Fp8DotGeneralBase(module.Module):\n amax_history_length: int = 1024\n e4m3_dtype: DType = jnp.float8_e4m3fn\n e5m2_dtype: DType = jnp.float8_e5m2\n\n def setup(self) -> None:\n scale_args = (\n initializers.ones_init(),\n random.PRNGKey(0),\n (1,),\n jnp.float32,\n )\n amax_history_args = (\n initializers.zeros_init(),\n random.PRNGKey(0),\n (self.amax_history_length,),\n jnp.float32,\n )\n\n self.input_amax_history = self.variable(\n OVERWRITE_WITH_GRADIENT, 'input_amax_history', *amax_history_args\n )\n self.kernel_amax_history = self.variable(\n OVERWRITE_WITH_GRADIENT, 'kernel_amax_history', *amax_history_args\n )\n self.output_grad_amax_history = self.variable(\n OVERWRITE_WITH_GRADIENT, 'output_grad_amax_history', *amax_history_args\n )\n\n self.input_scale = self.variable(\n OVERWRITE_WITH_GRADIENT, 'input_scale', *scale_args\n )\n self.kernel_scale = self.variable(\n OVERWRITE_WITH_GRADIENT, 'kernel_scale', *scale_args\n )\n self.output_grad_scale = self.variable(\n OVERWRITE_WITH_GRADIENT, 'output_grad_scale', *scale_args\n )\n\n\nclass Fp8DotGeneralOp(Fp8DotGeneralBase):\n def __post_init__(self):\n super().__post_init__()\n if type(self) is Fp8DotGeneralOp:\n warnings.warn(\n 'The Fp8DotGeneralOp is deprecated. Use Fp8DirectDotGeneralOp or '\n 'Fp8Einsum instead.',\n DeprecationWarning,\n )\n\n def __call__(self, *args, **kwargs):\n x, k, dimension_numbers, comp_dtype = _parse_dot_inputs(\n *args, **kwargs\n )\n x_qdq = in_qdq(\n comp_dtype, self.e4m3_dtype, x, self.input_scale.value, self.input_amax_history.value\n )\n k_qdq = in_qdq(\n comp_dtype, self.e4m3_dtype, k, self.kernel_scale.value, self.kernel_amax_history.value\n )\n\n y_qdq = dot_general_with_precision(x_qdq, k_qdq, dimension_numbers) # type: ignore\n y = out_qdq(\n comp_dtype,\n self.e5m2_dtype,\n y_qdq,\n self.output_grad_scale.value,\n self.output_grad_amax_history.value,\n )\n\n return y # type: ignore\n\nclass Fp8DirectDotGeneralOp(Fp8DotGeneralBase):\n def __call__(self, *args, **kwargs):\n x, k, dimension_numbers, comp_dtype = _parse_dot_inputs(\n *args, **kwargs\n )\n\n y = fp8_scaled_dot_general(\n x,\n k,\n dimension_numbers,\n precision=None,\n preferred_element_type=x.dtype,\n lhs_scale=self.input_scale.value,\n rhs_scale=self.kernel_scale.value,\n grad_scale=self.output_grad_scale.value,\n lhs_amax_history=self.input_amax_history.value,\n rhs_amax_history=self.kernel_amax_history.value,\n grad_amax_history=self.output_grad_amax_history.value,\n quantize_compute_type=comp_dtype,\n )\n\n return y # type: ignore\n\nclass NANOOFp8DotGeneralOp(Fp8DotGeneralOp):\n e4m3_dtype: DType = jnp.float8_e4m3fnuz\n e5m2_dtype: DType = jnp.float8_e5m2fnuz\n\nclass Fp8Einsum(Fp8DotGeneralBase):\n\n def __call__(self, eqn, lhs: jnp.ndarray, rhs: jnp.ndarray,\n precision: lax.Precision | None = None,\n preferred_element_type: DTypeLike | None = None) -> jnp.ndarray:\n # Here we assume that the rhs is the weight and its dtype is the actual compute dtype (not storage dtype).\n # TODO(kaixih@nvidia): Better way to handle this?\n actual_compute_dtype = rhs.dtype\n lhs = lhs.astype(actual_compute_dtype)\n\n dot_general_fn = partial(\n fp8_scaled_dot_general,\n lhs_scale=self.input_scale.value,\n rhs_scale=self.kernel_scale.value,\n grad_scale=self.output_grad_scale.value,\n lhs_amax_history=self.input_amax_history.value,\n rhs_amax_history=self.kernel_amax_history.value,\n grad_amax_history=self.output_grad_amax_history.value,\n quantize_compute_type=actual_compute_dtype\n )\n out = jnp.einsum(eqn, lhs, rhs, precision=precision,\n preferred_element_type=preferred_element_type,\n _dot_general=dot_general_fn)\n return out\n\n# Alias for backward compatibility\nFp8DotGeneral = Fp8DirectDotGeneralOp\n" }, { "path": "flax/linen/initializers.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Initializers for Flax.\"\"\"\n\n# pylint: disable=unused-import\n# re-export initializer functions from jax.nn\nfrom jax.nn.initializers import constant as constant\nfrom jax.nn.initializers import delta_orthogonal as delta_orthogonal\nfrom jax.nn.initializers import glorot_normal as glorot_normal\nfrom jax.nn.initializers import glorot_uniform as glorot_uniform\nfrom jax.nn.initializers import he_normal as he_normal\nfrom jax.nn.initializers import he_uniform as he_uniform\nfrom jax.nn.initializers import kaiming_normal as kaiming_normal\nfrom jax.nn.initializers import kaiming_uniform as kaiming_uniform\nfrom jax.nn.initializers import lecun_normal as lecun_normal\nfrom jax.nn.initializers import lecun_uniform as lecun_uniform\nfrom jax.nn.initializers import normal as normal\nfrom jax.nn.initializers import ones as ones\nfrom jax.nn.initializers import orthogonal as orthogonal\nfrom jax.nn.initializers import truncated_normal as truncated_normal\nfrom jax.nn.initializers import uniform as uniform\nfrom jax.nn.initializers import variance_scaling as variance_scaling\nfrom jax.nn.initializers import xavier_normal as xavier_normal\nfrom jax.nn.initializers import xavier_uniform as xavier_uniform\nfrom jax.nn.initializers import zeros as zeros\nfrom flax.typing import Initializer as Initializer\n\n# pylint: enable=unused-import\n\n\ndef zeros_init() -> Initializer:\n \"\"\"Builds an initializer that returns a constant array full of zeros.\n\n >>> import jax, jax.numpy as jnp\n >>> from flax.linen.initializers import zeros_init\n >>> zeros_initializer = zeros_init()\n >>> zeros_initializer(jax.random.key(42), (2, 3), jnp.float32)\n Array([[0., 0., 0.],\n [0., 0., 0.]], dtype=float32)\n \"\"\"\n return zeros\n\n\ndef ones_init() -> Initializer:\n \"\"\"Builds an initializer that returns a constant array full of ones.\n\n >>> import jax, jax.numpy as jnp\n >>> from flax.linen.initializers import ones_init\n >>> ones_initializer = ones_init()\n >>> ones_initializer(jax.random.key(42), (3, 2), jnp.float32)\n Array([[1., 1.],\n [1., 1.],\n [1., 1.]], dtype=float32)\n \"\"\"\n return ones\n" }, { "path": "flax/linen/kw_only_dataclasses.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Support for keyword-only fields in dataclasses for Python versions <3.10.\n\nThis module provides wrappers for `dataclasses.dataclass` and\n`dataclasses.field` that simulate support for keyword-only fields for Python\nversions before 3.10 (which is the version where dataclasses added keyword-only\nfield support). If this module is imported in Python 3.10+, then\n`kw_only_dataclasses.dataclass` and `kw_only_dataclasses.field` will simply be\naliases for `dataclasses.dataclass` and `dataclasses.field`.\n\nFor earlier Python versions, when constructing a dataclass, any fields that have\nbeen marked as keyword-only (including inherited fields) will be moved to the\nend of the constructor's argument list. This makes it possible to have a base\nclass that defines a field with a default, and a subclass that defines a field\nwithout a default. E.g.:\n\n>>> from flax.linen import kw_only_dataclasses\n>>> @kw_only_dataclasses.dataclass\n... class Parent:\n... name: str = kw_only_dataclasses.field(default='', kw_only=True)\n\n>>> @kw_only_dataclasses.dataclass\n... class Child(Parent):\n... size: float # required.\n\n>>> import inspect\n>>> print(inspect.signature(Child.__init__))\n(self, size: float, name: str = '') -> None\n\n\n(If we used `dataclasses` rather than `kw_only_dataclasses` for the above\nexample, then it would have failed with TypeError \"non-default argument\n'size' follows default argument.\")\n\nWARNING: fields marked as keyword-only will not *actually* be turned into\nkeyword-only parameters in the constructor; they will only be moved to the\nend of the parameter list (after all non-keyword-only parameters).\n\"\"\"\n\nimport dataclasses\nimport functools\nimport inspect\nimport sys\nfrom types import MappingProxyType\nfrom typing import Any, TypeVar\n\nimport typing_extensions as tpe\n\nimport flax\n\nM = TypeVar('M', bound='flax.linen.Module')\nFieldName = str\nAnnotation = Any\nDefault = Any\n\n\nclass _KwOnlyType:\n \"\"\"Metadata tag used to tag keyword-only fields.\"\"\"\n\n def __repr__(self):\n return 'KW_ONLY'\n\n\nKW_ONLY = _KwOnlyType()\n\n\ndef field(*, metadata=None, kw_only=dataclasses.MISSING, **kwargs):\n \"\"\"Wrapper for dataclasses.field that adds support for kw_only fields.\n\n Args:\n metadata: A mapping or None, containing metadata for the field.\n kw_only: If true, the field will be moved to the end of `__init__`'s\n parameter list.\n **kwargs: Keyword arguments forwarded to `dataclasses.field`\n\n Returns:\n A `dataclasses.Field` object.\n \"\"\"\n if kw_only is not dataclasses.MISSING and kw_only:\n if (\n kwargs.get('default', dataclasses.MISSING) is dataclasses.MISSING\n and kwargs.get('default_factory', dataclasses.MISSING)\n is dataclasses.MISSING\n ):\n raise ValueError('Keyword-only fields with no default are not supported.')\n if metadata is None:\n metadata = {}\n metadata[KW_ONLY] = True\n return dataclasses.field(metadata=metadata, **kwargs)\n\n\n@tpe.dataclass_transform(field_specifiers=(field,)) # type: ignore[literal-required]\ndef dataclass(cls=None, extra_fields=None, **kwargs):\n \"\"\"Wrapper for dataclasses.dataclass that adds support for kw_only fields.\n\n Args:\n cls: The class to transform (or none to return a decorator).\n extra_fields: A list of `(name, type, Field)` tuples describing extra fields\n that should be added to the dataclass. This is necessary for linen's\n use-case of this module, since the base class (linen.Module) is *not* a\n dataclass. In particular, linen.Module class is used as the base for both\n frozen and non-frozen dataclass subclasses; but the frozen status of a\n dataclass must match the frozen status of any base dataclasses.\n **kwargs: Additional arguments for `dataclasses.dataclass`.\n\n Returns:\n `cls`.\n \"\"\"\n\n def wrap(cls):\n return _process_class(cls, extra_fields=extra_fields, **kwargs)\n\n return wrap if cls is None else wrap(cls)\n\n\ndef _process_class(cls: type[M], extra_fields=None, **kwargs):\n \"\"\"Transforms `cls` into a dataclass that supports kw_only fields.\"\"\"\n if sys.version_info < (3, 14) and '__annotations__' not in cls.__dict__:\n cls.__annotations__ = {}\n\n # The original __dataclass_fields__ dicts for all base classes. We will\n # modify these in-place before turning `cls` into a dataclass, and then\n # restore them to their original values.\n base_dataclass_fields = {} # dict[cls, cls.__dataclass_fields__.copy()]\n\n # The keyword only fields from `cls` or any of its base classes.\n kw_only_fields: dict[FieldName, tuple[Annotation, Default]] = {}\n\n # Scan for KW_ONLY marker.\n kw_only_name = None\n for name, annotation in cls.__annotations__.items():\n if annotation is KW_ONLY:\n if kw_only_name is not None:\n raise TypeError('Multiple KW_ONLY markers')\n kw_only_name = name\n elif kw_only_name is not None:\n if not hasattr(cls, name):\n raise ValueError(\n 'Keyword-only fields with no default are not supported.'\n )\n default = getattr(cls, name)\n if isinstance(default, dataclasses.Field):\n default.metadata = MappingProxyType({**default.metadata, KW_ONLY: True})\n else:\n default = field(default=default, kw_only=True)\n setattr(cls, name, default)\n if kw_only_name:\n del cls.__annotations__[kw_only_name]\n\n # Inject extra fields.\n if extra_fields:\n for name, annotation, default in extra_fields:\n if not (isinstance(name, str) and isinstance(default, dataclasses.Field)):\n raise ValueError(\n 'Expected extra_fields to a be a list of '\n '(name, type, Field) tuples.'\n )\n setattr(cls, name, default)\n cls.__annotations__[name] = annotation\n\n # Extract kw_only fields from base classes' __dataclass_fields__.\n for base in reversed(cls.__mro__[1:]):\n if not dataclasses.is_dataclass(base):\n continue\n if sys.version_info < (3, 14):\n base_annotations = base.__dict__.get('__annotations__', {})\n else:\n base_annotations = inspect.get_annotations(base)\n\n base_dataclass_fields[base] = dict(\n getattr(base, '__dataclass_fields__', {})\n )\n for base_field in list(dataclasses.fields(base)):\n field_name = base_field.name\n if base_field.metadata.get(KW_ONLY) or field_name in kw_only_fields:\n kw_only_fields[field_name] = (\n base_annotations.get(field_name),\n base_field,\n )\n del base.__dataclass_fields__[field_name]\n\n # Remove any keyword-only fields from this class.\n if sys.version_info < (3, 14):\n cls_annotations = cls.__dict__['__annotations__']\n else:\n cls_annotations = cls.__annotations__\n for name, annotation in list(cls_annotations.items()):\n value = getattr(cls, name, None)\n if (\n isinstance(value, dataclasses.Field) and value.metadata.get(KW_ONLY)\n ) or name in kw_only_fields:\n del cls_annotations[name]\n kw_only_fields[name] = (annotation, value)\n\n # Add keyword-only fields at the end of __annotations__, in the order they\n # were found in the base classes and in this class.\n for name, (annotation, default) in kw_only_fields.items():\n setattr(cls, name, default)\n cls_annotations.pop(name, None)\n cls_annotations[name] = annotation\n\n create_init = '__init__' not in vars(cls) and kwargs.get('init', True)\n\n # Apply the dataclass transform.\n transformed_cls: type[M] = dataclasses.dataclass(cls, **kwargs)\n\n # Restore the base classes' __dataclass_fields__.\n for _cls, fields in base_dataclass_fields.items():\n _cls.__dataclass_fields__ = fields\n\n if create_init:\n dataclass_init = transformed_cls.__init__\n # use sum to count the number of init fields that are not keyword-only\n expected_num_args = sum(\n f.init and not f.metadata.get(KW_ONLY, False)\n for f in dataclasses.fields(transformed_cls)\n )\n\n @functools.wraps(dataclass_init)\n def init_wrapper(self, *args, **kwargs):\n num_args = len(args)\n if num_args > expected_num_args:\n # we add + 1 to each to account for `self`, matching python's\n # default error message\n raise TypeError(\n f'__init__() takes {expected_num_args + 1} positional '\n f'arguments but {num_args + 1} were given'\n )\n\n dataclass_init(self, *args, **kwargs)\n\n init_wrapper.__signature__ = inspect.signature(dataclass_init) # type: ignore\n transformed_cls.__init__ = init_wrapper # type: ignore[method-assign]\n\n # Return the transformed dataclass\n return transformed_cls\n" }, { "path": "flax/linen/linear.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Linear modules.\"\"\"\n\nfrom collections.abc import Iterable, Sequence\nfrom typing import Any, Protocol\n\nfrom flax.core import meta\nfrom flax.linen import initializers\nfrom flax.linen import module\nfrom flax.linen.dtypes import promote_dtype\nfrom flax.linen.module import Module, compact\nfrom flax.typing import (\n Array,\n ConvGeneralDilatedT,\n DotGeneralT,\n Dtype,\n Initializer,\n LaxPadding,\n PRNGKey as PRNGKey,\n PaddingLike,\n PrecisionLike,\n Shape as Shape,\n)\nimport jax\nfrom jax import eval_shape, lax\nfrom jax.core import ShapedArray\nimport jax.numpy as jnp\nimport numpy as np\nimport opt_einsum\n\nclass PromoteDtypeFn(Protocol):\n def __call__(\n self, *args: jax.Array | None, dtype: Any = None, inexact: bool = True\n ) -> list[jax.Array | None]: ...\n\ndefault_kernel_init = initializers.lecun_normal()\n\n\ndef _normalize_axes(axes: tuple[int, ...], ndim: int) -> tuple[int, ...]:\n # A tuple by convention. len(axes_tuple) then also gives the rank efficiently.\n return tuple(sorted(ax if ax >= 0 else ndim + ax for ax in axes))\n\n\ndef _canonicalize_tuple(x: Sequence[int] | int) -> tuple[int, ...]:\n if isinstance(x, Iterable):\n return tuple(x)\n else:\n return (x,)\n\n\nclass DenseGeneral(Module):\n \"\"\"A linear transformation with flexible axes.\n\n Example usage::\n\n >>> import flax.linen as nn\n >>> import jax, jax.numpy as jnp\n\n >>> # equivalent to `nn.Dense(features=4)`\n >>> layer = nn.DenseGeneral(features=4)\n >>> # output features (4, 5)\n >>> layer = nn.DenseGeneral(features=(4, 5))\n >>> params = layer.init(jax.random.key(0), jnp.ones((1, 3)))\n >>> jax.tree_util.tree_map(jnp.shape, params)\n {'params': {'bias': (4, 5), 'kernel': (3, 4, 5)}}\n >>> # apply transformation on the the second and last axes\n >>> layer = nn.DenseGeneral(features=(4, 5), axis=(1, -1))\n >>> params = layer.init(jax.random.key(0), jnp.ones((1, 3, 6, 7)))\n >>> jax.tree_util.tree_map(jnp.shape, params)\n {'params': {'bias': (4, 5), 'kernel': (3, 7, 4, 5)}}\n\n Attributes:\n features: int or tuple with number of output features.\n axis: int or tuple with axes to apply the transformation on. For instance,\n (-2, -1) will apply the transformation to the last two axes.\n batch_dims: tuple with batch axes.\n use_bias: whether to add a bias to the output (default: True).\n dtype: the dtype of the computation (default: infer from input and params).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n kernel_init: initializer function for the weight matrix.\n bias_init: initializer function for the bias.\n precision: numerical precision of the computation see ``jax.lax.Precision``\n for details.\n promote_dtype: function to promote the dtype of the arrays to the desired\n dtype. The function should accept a tuple of ``(inputs, kernel, bias)``\n and a ``dtype`` keyword argument, and return a tuple of arrays with the\n promoted dtype.\n \"\"\"\n\n features: int | Sequence[int]\n axis: int | Sequence[int] = -1\n batch_dims: Sequence[int] = ()\n use_bias: bool = True\n dtype: Dtype | None = None\n param_dtype: Dtype = jnp.float32\n kernel_init: Initializer = default_kernel_init\n bias_init: Initializer = initializers.zeros_init()\n precision: PrecisionLike = None\n promote_dtype: PromoteDtypeFn = promote_dtype\n # Deprecated. Will be removed.\n dot_general: DotGeneralT | None = None\n dot_general_cls: Any = None\n\n @compact\n def __call__(self, inputs: Array) -> Array:\n \"\"\"Applies a linear transformation to the inputs along multiple dimensions.\n\n Args:\n inputs: The nd-array to be transformed.\n\n Returns:\n The transformed input.\n \"\"\"\n features = _canonicalize_tuple(self.features)\n axis = _canonicalize_tuple(self.axis)\n batch_dims = _canonicalize_tuple(self.batch_dims)\n if batch_dims:\n max_dim = np.max(batch_dims)\n if set(batch_dims) != set(range(max_dim + 1)):\n raise ValueError(\n 'batch_dims %s must be consecutive leading '\n 'dimensions starting from 0.' % str(batch_dims)\n )\n\n ndim = inputs.ndim\n n_batch_dims = len(batch_dims)\n axis = _normalize_axes(axis, ndim)\n batch_dims = _normalize_axes(batch_dims, ndim)\n n_axis, n_features = len(axis), len(features)\n\n def kernel_init_wrap(rng, shape, dtype=jnp.float32):\n flat_shape = (\n np.prod(shape[:n_batch_dims])\n * np.prod(shape[n_batch_dims : n_axis + n_batch_dims]),\n np.prod(shape[-n_features:]),\n )\n flat_shape = jax.tree_util.tree_map(int, flat_shape)\n kernel = self.kernel_init(rng, flat_shape, dtype)\n if isinstance(kernel, meta.AxisMetadata):\n return meta.replace_boxed(kernel, jnp.reshape(kernel.unbox(), shape))\n return jnp.reshape(kernel, shape)\n\n batch_shape = tuple(inputs.shape[ax] for ax in batch_dims)\n # batch and non-contracting dims of input with 1s for batch dims.\n expanded_batch_shape = tuple(\n inputs.shape[ax] if ax in batch_dims else 1\n for ax in range(inputs.ndim)\n if ax not in axis\n )\n kernel_shape = tuple(inputs.shape[ax] for ax in axis) + features\n kernel = self.param(\n 'kernel', kernel_init_wrap, batch_shape + kernel_shape, self.param_dtype\n )\n\n batch_ind = tuple(range(n_batch_dims))\n contract_ind = tuple(range(n_batch_dims, n_axis + n_batch_dims))\n\n if self.use_bias:\n\n def bias_init_wrap(rng, shape, dtype=jnp.float32):\n flat_shape = (\n np.prod(shape[:n_batch_dims]) * np.prod(shape[-n_features:]),\n )\n flat_shape = jax.tree_util.tree_map(int, flat_shape)\n bias = self.bias_init(rng, flat_shape, dtype)\n if isinstance(bias, meta.AxisMetadata):\n return meta.replace_boxed(bias, jnp.reshape(bias.unbox(), shape))\n return jnp.reshape(bias, shape)\n\n bias = self.param(\n 'bias', bias_init_wrap, batch_shape + features, self.param_dtype\n )\n else:\n bias = None\n\n inputs, kernel, bias = self.promote_dtype(\n inputs, kernel, bias, dtype=self.dtype\n )\n\n if self.dot_general_cls is not None:\n dot_general = self.dot_general_cls()\n elif self.dot_general is not None:\n dot_general = self.dot_general\n else:\n dot_general = lax.dot_general\n out = dot_general(\n inputs,\n kernel,\n ((axis, contract_ind), (batch_dims, batch_ind)),\n precision=self.precision,\n )\n # dot_general output has shape [batch_dims/group_dims] + [feature_dims]\n if self.use_bias:\n # expand bias shape to broadcast bias over batch dims.\n assert bias is not None\n bias = jnp.reshape(bias, expanded_batch_shape + features)\n out += bias\n return out\n\n\nclass Dense(Module):\n \"\"\"A linear transformation applied over the last dimension of the input.\n\n Example usage::\n\n >>> import flax.linen as nn\n >>> import jax, jax.numpy as jnp\n\n >>> layer = nn.Dense(features=4)\n >>> params = layer.init(jax.random.key(0), jnp.ones((1, 3)))\n >>> jax.tree_util.tree_map(jnp.shape, params)\n {'params': {'bias': (4,), 'kernel': (3, 4)}}\n\n Attributes:\n features: the number of output features.\n use_bias: whether to add a bias to the output (default: True).\n dtype: the dtype of the computation (default: infer from input and params).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n precision: numerical precision of the computation see ``jax.lax.Precision``\n for details.\n kernel_init: initializer function for the weight matrix.\n bias_init: initializer function for the bias.\n promote_dtype: function to promote the dtype of the arrays to the desired\n dtype. The function should accept a tuple of ``(inputs, kernel, bias)``\n and a ``dtype`` keyword argument, and return a tuple of arrays with the\n promoted dtype.\n \"\"\"\n\n features: int\n use_bias: bool = True\n dtype: Dtype | None = None\n param_dtype: Dtype = jnp.float32\n precision: PrecisionLike = None\n kernel_init: Initializer = default_kernel_init\n bias_init: Initializer = initializers.zeros_init()\n promote_dtype: PromoteDtypeFn = promote_dtype\n dot_general: DotGeneralT | None = None\n dot_general_cls: Any = None\n\n @compact\n def __call__(self, inputs: Array) -> Array:\n \"\"\"Applies a linear transformation to the inputs along the last dimension.\n\n Args:\n inputs: The nd-array to be transformed.\n\n Returns:\n The transformed input.\n \"\"\"\n kernel = self.param(\n 'kernel',\n self.kernel_init,\n (jnp.shape(inputs)[-1], self.features),\n self.param_dtype,\n )\n if self.use_bias:\n bias = self.param(\n 'bias', self.bias_init, (self.features,), self.param_dtype\n )\n else:\n bias = None\n inputs, kernel, bias = self.promote_dtype(\n inputs, kernel, bias, dtype=self.dtype\n )\n assert inputs is not None\n assert kernel is not None\n\n if self.dot_general_cls is not None:\n dot_general = self.dot_general_cls()\n elif self.dot_general is not None:\n dot_general = self.dot_general\n else:\n dot_general = lax.dot_general\n y = dot_general(\n inputs,\n kernel,\n (((inputs.ndim - 1,), (0,)), ((), ())),\n precision=self.precision,\n )\n if bias is not None:\n y += jnp.reshape(bias, (1,) * (y.ndim - 1) + (-1,))\n return y\n\n\nclass Einsum(Module):\n \"\"\"An einsum transformation with learnable kernel and bias.\n\n Example usage::\n\n >>> import flax.linen as nn\n >>> import jax, jax.numpy as jnp\n\n >>> layer = nn.Einsum((5, 6, 7), 'abc,cde->abde')\n >>> variables = layer.init(jax.random.key(0), jnp.ones((3, 4, 5)))\n >>> jax.tree_util.tree_map(jnp.shape, variables)\n {'params': {'bias': (6, 7), 'kernel': (5, 6, 7)}}\n\n Attributes:\n shape: the shape of the kernel.\n einsum_str: a string to denote the einsum equation. The equation must\n have exactly two operands, the lhs being the input passed in, and\n the rhs being the learnable kernel. Exactly one of ``einsum_str``\n in the constructor argument and call argument must be not None,\n while the other must be None.\n use_bias: whether to add a bias to the output (default: True).\n dtype: the dtype of the computation (default: infer from input and params).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n precision: numerical precision of the computation see ``jax.lax.Precision``\n for details.\n kernel_init: initializer function for the weight matrix.\n bias_init: initializer function for the bias.\n promote_dtype: function to promote the dtype of the arrays to the desired\n dtype. The function should accept a tuple of ``(inputs, kernel, bias)``\n and a ``dtype`` keyword argument, and return a tuple of arrays with the\n promoted dtype.\n \"\"\"\n\n shape: Shape\n einsum_str: str | None = None\n use_bias: bool = True\n dtype: Dtype | None = None\n param_dtype: Dtype = jnp.float32\n precision: PrecisionLike = None\n kernel_init: Initializer = default_kernel_init\n bias_init: Initializer = initializers.zeros_init()\n promote_dtype: PromoteDtypeFn = promote_dtype\n preferred_element_type: Dtype | None = None\n\n @compact\n def __call__(self, inputs: Array, einsum_str: str | None = None) -> Array:\n \"\"\"Applies a linear transformation to the inputs along the last dimension.\n\n Args:\n inputs: The nd-array to be transformed.\n einsum_str: a string to denote the einsum equation. The equation must\n have exactly two operands, the lhs being the input passed in, and\n the rhs being the learnable kernel. The ``einsum_str`` passed into\n the call method will take precedence over the ``einsum_str`` passed\n into the constructor.\n\n Returns:\n The transformed input.\n \"\"\"\n einsum_str = module.merge_param('einsum_str', self.einsum_str, einsum_str)\n\n einsum_str = einsum_str.replace(' ', '')\n if '->' not in einsum_str:\n raise ValueError(\n '`einsum_str` equation must be explicit and include \"->\".'\n )\n if einsum_str.count(',') != 1:\n raise ValueError(\n '`einsum_str` equation must have exactly two operands and '\n 'therefore, exactly one comma character, instead of '\n f'{einsum_str.count(\",\")}'\n )\n\n kernel = self.param(\n 'kernel',\n self.kernel_init,\n self.shape,\n self.param_dtype,\n )\n\n if self.use_bias:\n bias_shape, broadcasted_bias_shape = self._get_bias_shape(\n einsum_str, inputs, kernel\n )\n bias = self.param('bias', self.bias_init, bias_shape, self.param_dtype)\n else:\n bias = None\n\n inputs, kernel, bias = self.promote_dtype(\n inputs, kernel, bias, dtype=self.dtype\n )\n\n y = jnp.einsum(\n einsum_str,\n inputs,\n kernel,\n precision=self.precision,\n preferred_element_type=self.preferred_element_type,\n )\n\n if bias is not None:\n y += jnp.reshape(bias, broadcasted_bias_shape)\n return y\n\n def _get_bias_shape(self, einsum_str: str, lhs: Array, rhs: Array):\n \"\"\"Infer the bias shape and broadcasted bias shape given the ``einsum_str``,\n ``lhs`` and ``rhs`` arrays. This is needed for instantiating the bias\n parameter and adding the bias to the output during forward inference.\n\n This function first replaces all ellipses with actual letter characters,\n then computes the bias shape by checking to see which axes in the rhs\n array remain in the resulting array after einsumming. These axes are the\n embedding/feature dimensions, and all other axes in rhs are reduction axes.\n \"\"\"\n # More details on the parsing function: https://github.com/dgasmith/opt_einsum/blob/c826bb7df16f470a69f7bf90598fc27586209d11/opt_einsum/parser.py#L246\n # returns the einsum string representation of the operands and result, with\n # ellipsis replaced by actual letter characters\n operands_str, result_str, _ = opt_einsum.parser.parse_einsum_input(\n (einsum_str, lhs, rhs)\n )\n\n # rhs_dict is a dict{character:index} mapping that maps every character in\n # the rhs einsum string representation to its corresponding index position in the string\n rhs_dict = {c: i for i, c in enumerate(operands_str.split(',')[1])}\n assert len(rhs_dict) == len(self.shape)\n\n broadcasted_bias_shape = [1] * len(result_str)\n bias_shape = []\n for i, c in enumerate(result_str):\n if c in rhs_dict:\n broadcasted_bias_shape[i] = self.shape[rhs_dict[c]]\n bias_shape.append(self.shape[rhs_dict[c]])\n\n return bias_shape, broadcasted_bias_shape\n\n\ndef _conv_dimension_numbers(input_shape):\n \"\"\"Computes the dimension numbers based on the input shape.\"\"\"\n ndim = len(input_shape)\n lhs_spec = (0, ndim - 1) + tuple(range(1, ndim - 1))\n rhs_spec = (ndim - 1, ndim - 2) + tuple(range(0, ndim - 2))\n out_spec = lhs_spec\n return lax.ConvDimensionNumbers(lhs_spec, rhs_spec, out_spec)\n\n\ndef canonicalize_padding(padding: PaddingLike, rank: int) -> LaxPadding:\n \"\"\" \"Canonicalizes conv padding to a jax.lax supported format.\"\"\"\n if isinstance(padding, str):\n return padding\n if isinstance(padding, int):\n return [(padding, padding)] * rank\n if isinstance(padding, Sequence) and len(padding) == rank:\n new_pad = []\n for p in padding:\n if isinstance(p, int):\n new_pad.append((p, p))\n elif isinstance(p, tuple) and len(p) == 2:\n new_pad.append(p)\n else:\n break\n if len(new_pad) == rank:\n return new_pad\n raise ValueError(\n f'Invalid padding format: {padding}, should be str, int,'\n f' or a sequence of len {rank} where each element is an'\n ' int or pair of ints.'\n )\n\n\nclass _Conv(Module):\n \"\"\"Convolution Module wrapping ``lax.conv_general_dilated``.\n\n Attributes:\n features: number of convolution filters.\n kernel_size: shape of the convolutional kernel. An integer will be\n interpreted as a tuple of the single integer.\n strides: an integer or a sequence of `n` integers, representing the\n inter-window strides (default: 1).\n padding: either the string ``'SAME'``, the string ``'VALID'``, the string\n ``'CIRCULAR'`` (periodic boundary conditions), the string `'REFLECT'`\n (reflection across the padding boundary), or a sequence of ``n`` ``(low,\n high)`` integer pairs that give the padding to apply before and after each\n spatial dimension. A single int is interpreted as applying the same padding\n in all dims and assign a single int in a sequence causes the same padding\n to be used on both sides. ``'CAUSAL'`` padding for a 1D convolution will\n left-pad the convolution axis, resulting in same-sized output.\n input_dilation: an integer or a sequence of ``n`` integers, giving the\n dilation factor to apply in each spatial dimension of ``inputs``\n (default: 1). Convolution with input dilation ``d`` is equivalent to\n transposed convolution with stride ``d``.\n kernel_dilation: an integer or a sequence of ``n`` integers, giving the\n dilation factor to apply in each spatial dimension of the convolution\n kernel (default: 1). Convolution with kernel dilation\n is also known as 'atrous convolution'.\n feature_group_count: integer, default 1. If specified divides the input\n features into groups.\n use_bias: whether to add a bias to the output (default: True).\n mask: Optional mask for the weights during masked convolution. The mask must\n be the same shape as the convolution weight matrix.\n dtype: the dtype of the computation (default: infer from input and params).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n precision: numerical precision of the computation see ``jax.lax.Precision``\n for details.\n kernel_init: initializer for the convolutional kernel.\n bias_init: initializer for the bias.\n promote_dtype: function to promote the dtype of the arrays to the desired\n dtype. The function should accept a tuple of ``(inputs, kernel, bias)``\n and a ``dtype`` keyword argument, and return a tuple of arrays with the\n promoted dtype.\n \"\"\"\n\n features: int\n kernel_size: int | Sequence[int]\n strides: None | int | Sequence[int] = 1\n padding: PaddingLike = 'SAME'\n input_dilation: None | int | Sequence[int] = 1\n kernel_dilation: None | int | Sequence[int] = 1\n feature_group_count: int = 1\n use_bias: bool = True\n mask: Array | None = None\n dtype: Dtype | None = None\n param_dtype: Dtype = jnp.float32\n precision: PrecisionLike = None\n kernel_init: Initializer = default_kernel_init\n bias_init: Initializer = initializers.zeros_init()\n promote_dtype: PromoteDtypeFn = promote_dtype\n # Deprecated. Will be removed.\n conv_general_dilated: ConvGeneralDilatedT | None = None\n conv_general_dilated_cls: Any = None\n\n @property\n def shared_weights(self) -> bool: # type: ignore\n \"\"\"Defines whether weights are shared or not between different pixels.\n\n Returns:\n ``True`` to use shared weights in convolution (regular convolution).\n ``False`` to use different weights at different pixels, a.k.a.\n \"locally connected layer\", \"unshared convolution\", or \"local convolution\".\n\n \"\"\"\n ...\n\n @compact\n def __call__(self, inputs: Array) -> Array:\n \"\"\"Applies a (potentially unshared) convolution to the inputs.\n\n Args:\n inputs: input data with dimensions ``(*batch_dims, spatial_dims..., features)``.\n This is the channels-last convention, i.e. NHWC for a 2d convolution and\n NDHWC for a 3D convolution. Note: this is different from the input convention\n used by ``lax.conv_general_dilated``, which puts the spatial dimensions last.\n Note: If the input has more than 1 batch dimension, all batch dimensions\n are flattened into a single dimension for the convolution and restored\n before returning. In some cases directly vmap'ing the layer may yield\n better performance than this default flattening approach. If the input\n lacks a batch dimension it will be added for the convolution and removed\n n return, an allowance made to enable writing single-example code.\n\n Returns:\n The convolved data.\n \"\"\"\n\n kernel_size: Sequence[int]\n if isinstance(self.kernel_size, int):\n kernel_size = (self.kernel_size,)\n else:\n kernel_size = tuple(self.kernel_size)\n\n def maybe_broadcast(\n x: int | Sequence[int] | None,\n ) -> tuple[int, ...]:\n if x is None:\n # backward compatibility with using None as sentinel for\n # broadcast 1\n x = 1\n if isinstance(x, int):\n return (x,) * len(kernel_size)\n return tuple(x)\n\n # Combine all input batch dimensions into a single leading batch axis.\n num_batch_dimensions = inputs.ndim - (len(kernel_size) + 1)\n if num_batch_dimensions != 1:\n input_batch_shape = inputs.shape[:num_batch_dimensions]\n flat_input_shape = (-1,) + inputs.shape[\n num_batch_dimensions:\n ]\n inputs = jnp.reshape(inputs, flat_input_shape)\n\n # self.strides or (1,) * (inputs.ndim - 2)\n strides = maybe_broadcast(self.strides)\n input_dilation = maybe_broadcast(self.input_dilation)\n kernel_dilation = maybe_broadcast(self.kernel_dilation)\n\n padding_lax = canonicalize_padding(self.padding, len(kernel_size))\n if padding_lax in ('CIRCULAR', 'REFLECT'):\n assert isinstance(padding_lax, str)\n kernel_size_dilated = [\n (k - 1) * d + 1 for k, d in zip(kernel_size, kernel_dilation)\n ]\n zero_pad: list[tuple[int, int]] = [(0, 0)]\n pads = (\n zero_pad\n + [((k - 1) // 2, k // 2) for k in kernel_size_dilated]\n + [(0, 0)]\n )\n padding_mode = {'CIRCULAR': 'wrap', 'REFLECT': 'reflect'}[padding_lax]\n inputs = jnp.pad(inputs, pads, mode=padding_mode)\n padding_lax = 'VALID'\n elif padding_lax == 'CAUSAL':\n if len(kernel_size) != 1:\n raise ValueError(\n 'Causal padding is only implemented for 1D convolutions.'\n )\n left_pad = kernel_dilation[0] * (kernel_size[0] - 1)\n pads = [(0, 0), (left_pad, 0), (0, 0)]\n inputs = jnp.pad(inputs, pads)\n padding_lax = 'VALID'\n\n dimension_numbers = _conv_dimension_numbers(inputs.shape)\n in_features = jnp.shape(inputs)[-1]\n\n if self.shared_weights:\n # One shared convolutional kernel for all pixels in the output.\n assert in_features % self.feature_group_count == 0\n kernel_shape = kernel_size + (\n in_features // self.feature_group_count,\n self.features,\n )\n\n else:\n if self.feature_group_count != 1:\n raise NotImplementedError(\n '`lax.conv_general_dilated_local` does not support '\n f'`feature_group_count != 1`, got `{self.feature_group_count}`.'\n )\n\n # Need to know the spatial output shape of a standard convolution to\n # create the unshared convolution kernel.\n if self.conv_general_dilated_cls is not None:\n conv_general_dilated = self.conv_general_dilated_cls()\n elif self.conv_general_dilated is not None:\n conv_general_dilated = self.conv_general_dilated\n else:\n conv_general_dilated = lax.conv_general_dilated\n conv_output_shape = eval_shape(\n lambda lhs, rhs: conv_general_dilated( # pylint: disable=g-long-lambda\n lhs=lhs,\n rhs=rhs,\n window_strides=strides,\n padding=padding_lax,\n dimension_numbers=dimension_numbers,\n lhs_dilation=input_dilation,\n rhs_dilation=kernel_dilation,\n ),\n inputs,\n ShapedArray(kernel_size + (in_features, self.features), inputs.dtype),\n ).shape\n\n # One (unshared) convolutional kernel per each pixel in the output.\n kernel_shape = conv_output_shape[1:-1] + (\n np.prod(kernel_size) * in_features,\n self.features,\n )\n\n if self.mask is not None and self.mask.shape != kernel_shape:\n raise ValueError(\n 'Mask needs to have the same shape as weights. '\n f'Shapes are: {self.mask.shape}, {kernel_shape}'\n )\n\n kernel = self.param(\n 'kernel', self.kernel_init, kernel_shape, self.param_dtype\n )\n\n if self.mask is not None:\n kernel *= self.mask\n\n if self.use_bias:\n if self.shared_weights:\n # One bias weight per output channel, shared between pixels.\n bias_shape = (self.features,)\n else:\n # One bias weight per output entry, unshared betwen pixels.\n bias_shape = conv_output_shape[1:]\n\n bias = self.param('bias', self.bias_init, bias_shape, self.param_dtype)\n else:\n bias = None\n\n inputs, kernel, bias = self.promote_dtype(\n inputs, kernel, bias, dtype=self.dtype\n )\n assert inputs is not None\n assert kernel is not None\n\n if self.shared_weights:\n if self.conv_general_dilated_cls is not None:\n conv_general_dilated = self.conv_general_dilated_cls()\n elif self.conv_general_dilated is not None:\n conv_general_dilated = self.conv_general_dilated\n else:\n conv_general_dilated = lax.conv_general_dilated\n y = conv_general_dilated(\n inputs,\n kernel,\n strides,\n padding_lax,\n lhs_dilation=input_dilation,\n rhs_dilation=kernel_dilation,\n dimension_numbers=dimension_numbers,\n feature_group_count=self.feature_group_count,\n precision=self.precision,\n )\n else:\n y = lax.conv_general_dilated_local(\n lhs=inputs,\n rhs=kernel,\n window_strides=strides,\n padding=padding_lax,\n filter_shape=kernel_size,\n lhs_dilation=input_dilation,\n rhs_dilation=kernel_dilation,\n dimension_numbers=dimension_numbers,\n precision=self.precision,\n )\n\n if self.use_bias:\n bias = bias.reshape((1,) * (y.ndim - bias.ndim) + bias.shape) # type: ignore\n y += bias\n\n if num_batch_dimensions != 1:\n output_shape = input_batch_shape + y.shape[1:]\n y = jnp.reshape(y, output_shape)\n return y\n\n\nclass Conv(_Conv):\n \"\"\"Convolution Module wrapping ``lax.conv_general_dilated``.\n\n Example usage::\n\n >>> import flax.linen as nn\n >>> import jax, jax.numpy as jnp\n\n >>> # valid padding\n >>> layer = nn.Conv(features=4, kernel_size=(3,), padding='VALID')\n >>> out, variables = layer.init_with_output(jax.random.key(0), jnp.ones((1, 8, 3)))\n >>> jax.tree_util.tree_map(jnp.shape, variables)\n {'params': {'bias': (4,), 'kernel': (3, 3, 4)}}\n >>> out.shape\n (1, 6, 4)\n >>> # circular padding with stride 2\n >>> layer = nn.Conv(features=4, kernel_size=(3, 3), strides=2, padding='CIRCULAR')\n >>> out, variables = layer.init_with_output(jax.random.key(0), jnp.ones((1, 8, 3)))\n >>> jax.tree_util.tree_map(jnp.shape, variables)\n {'params': {'bias': (4,), 'kernel': (3, 3, 3, 4)}}\n >>> out.shape\n (1, 4, 4)\n >>> # apply lower triangle mask\n >>> mask = jnp.tril(jnp.ones((3, 3, 4)))\n >>> layer = nn.Conv(features=4, kernel_size=(3,), mask=mask, padding='VALID')\n >>> variables = layer.init(jax.random.key(0), jnp.ones((1, 8, 3)))\n\n Attributes:\n features: number of convolution filters.\n kernel_size: shape of the convolutional kernel. An integer will be\n interpreted as a tuple of the single integer.\n strides: an integer or a sequence of `n` integers, representing the\n inter-window strides (default: 1).\n padding: either the string ``'SAME'``, the string ``'VALID'``, the string\n ``'CIRCULAR'`` (periodic boundary conditions), or a sequence of ``n``\n ``(low, high)`` integer pairs that give the padding to apply before and\n after each spatial dimension. A single int is interpreted as applying the\n same padding in all dims and assign a single int in a sequence causes the\n same padding to be used on both sides. ``'CAUSAL'`` padding for a 1D\n convolution will left-pad the convolution axis, resulting in same-sized\n output.\n input_dilation: an integer or a sequence of ``n`` integers, giving the\n dilation factor to apply in each spatial dimension of ``inputs`` (default:\n 1). Convolution with input dilation ``d`` is equivalent to transposed\n convolution with stride ``d``.\n kernel_dilation: an integer or a sequence of ``n`` integers, giving the\n dilation factor to apply in each spatial dimension of the convolution\n kernel (default: 1). Convolution with kernel dilation is also known as\n 'atrous convolution'.\n feature_group_count: integer, default 1. If specified divides the input\n features into groups.\n use_bias: whether to add a bias to the output (default: True).\n mask: Optional mask for the weights during masked convolution. The mask must\n be the same shape as the convolution weight matrix.\n dtype: the dtype of the computation (default: infer from input and params).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n precision: numerical precision of the computation see ``jax.lax.Precision`\n for details.\n kernel_init: initializer for the convolutional kernel.\n bias_init: initializer for the bias.\n promote_dtype: function to promote the dtype of the arrays to the desired\n dtype. The function should accept a tuple of ``(inputs, kernel, bias)``\n and a ``dtype`` keyword argument, and return a tuple of arrays with the\n promoted dtype.\n \"\"\"\n\n @property\n def shared_weights(self) -> bool:\n return True\n\n\nclass ConvLocal(_Conv):\n \"\"\"Local convolution Module wrapping ``lax.conv_general_dilated_local``.\n\n Example usage::\n\n >>> import flax.linen as nn\n >>> import jax, jax.numpy as jnp\n\n >>> # valid padding\n >>> layer = nn.ConvLocal(features=4, kernel_size=(3,), padding='VALID')\n >>> out, variables = layer.init_with_output(jax.random.key(0), jnp.ones((1, 8, 3)))\n >>> jax.tree_util.tree_map(jnp.shape, variables)\n {'params': {'bias': (6, 4), 'kernel': (6, 9, 4)}}\n >>> out.shape\n (1, 6, 4)\n >>> # circular padding with stride 2\n >>> layer = nn.ConvLocal(features=4, kernel_size=(3, 3), strides=2, padding='CIRCULAR')\n >>> out, variables = layer.init_with_output(jax.random.key(0), jnp.ones((1, 8, 3)))\n >>> jax.tree_util.tree_map(jnp.shape, variables)\n {'params': {'bias': (1, 4, 4), 'kernel': (1, 4, 27, 4)}}\n >>> out.shape\n (1, 4, 4)\n >>> # apply lower triangle mask\n >>> mask = jnp.tril(jnp.ones((6, 9, 4)))\n >>> layer = nn.ConvLocal(features=4, kernel_size=(3,), mask=mask, padding='VALID')\n >>> variables = layer.init(jax.random.key(0), jnp.ones((1, 8, 3)))\n\n Attributes:\n features: number of convolution filters.\n kernel_size: shape of the convolutional kernel. An integer will be\n interpreted as a tuple of the single integer.\n strides: an integer or a sequence of `n` integers, representing the\n inter-window strides (default: 1).\n padding: either the string ``'SAME'``, the string ``'VALID'``, the string\n ``'CIRCULAR'`` (periodic boundary conditions), or a sequence of ``n``\n ``(low, high)`` integer pairs that give the padding to apply before and\n after each spatial dimension. A single int is interpreted as applying the\n same padding in all dims and assign a single int in a sequence causes the\n same padding to be used on both sides. ``'CAUSAL'`` padding for a 1D\n convolution will left-pad the convolution axis, resulting in same-sized\n output.\n input_dilation: an integer or a sequence of ``n`` integers, giving the\n dilation factor to apply in each spatial dimension of ``inputs`` (default:\n 1). Convolution with input dilation ``d`` is equivalent to transposed\n convolution with stride ``d``.\n kernel_dilation: an integer or a sequence of ``n`` integers, giving the\n dilation factor to apply in each spatial dimension of the convolution\n kernel (default: 1). Convolution with kernel dilation is also known as\n 'atrous convolution'.\n feature_group_count: integer, default 1. If specified divides the input\n features into groups.\n use_bias: whether to add a bias to the output (default: True).\n mask: Optional mask for the weights during masked convolution. The mask must\n be the same shape as the convolution weight matrix.\n dtype: the dtype of the computation (default: infer from input and params).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n precision: numerical precision of the computation see ``jax.lax.Precision``\n for details.\n kernel_init: initializer for the convolutional kernel.\n bias_init: initializer for the bias.\n promote_dtype: function to promote the dtype of the arrays to the desired\n dtype. The function should accept a tuple of ``(inputs, kernel, bias)``\n and a ``dtype`` keyword argument, and return a tuple of arrays with the\n promoted dtype.\n \"\"\"\n\n @property\n def shared_weights(self) -> bool:\n return False\n\n\nclass ConvTranspose(Module):\n \"\"\"Convolution Module wrapping ``lax.conv_transpose``.\n\n Example usage::\n\n >>> import flax.linen as nn\n >>> import jax, jax.numpy as jnp\n\n >>> # valid padding\n >>> layer = nn.ConvTranspose(features=4, kernel_size=(3,), padding='VALID')\n >>> out, variables = layer.init_with_output(jax.random.key(0), jnp.ones((1, 8, 3)))\n >>> jax.tree_util.tree_map(jnp.shape, variables)\n {'params': {'bias': (4,), 'kernel': (3, 3, 4)}}\n >>> out.shape\n (1, 10, 4)\n >>> # circular padding with stride 2\n >>> layer = nn.ConvTranspose(features=4, kernel_size=(6, 6), strides=(2, 2), padding='CIRCULAR', transpose_kernel=True)\n >>> out, variables = layer.init_with_output(jax.random.key(0), jnp.ones((1, 15, 15, 3)))\n >>> jax.tree_util.tree_map(jnp.shape, variables)\n {'params': {'bias': (4,), 'kernel': (6, 6, 4, 3)}}\n >>> out.shape\n (1, 30, 30, 4)\n >>> # apply lower triangle mask\n >>> mask = jnp.tril(jnp.ones((3, 3, 4)))\n >>> layer = nn.ConvTranspose(features=4, kernel_size=(3,), mask=mask, padding='VALID')\n >>> variables = layer.init(jax.random.key(0), jnp.ones((1, 8, 3)))\n\n Attributes:\n features: number of convolution filters.\n kernel_size: shape of the convolutional kernel. For 1D convolution, the\n kernel size can be passed as an integer, which will be interpreted as a\n tuple of the single integer. For all other cases, it must be a sequence of\n integers.\n strides: an integer or a sequence of `n` integers, representing the\n inter-window strides.\n padding: either the string `'SAME'`, the string `'VALID'`, the string\n `'CIRCULAR'` (periodic boundary conditions), or a sequence of `n` `(low,\n high)` integer pairs that give the padding to apply before and after each\n spatial dimension. A single int is interpreted as applying the same\n padding in all dims and assign a single int in a sequence causes the same\n padding to be used on both sides.\n kernel_dilation: ``None``, or an integer or a sequence of ``n`` integers,\n giving the dilation factor to apply in each spatial dimension of the convolution\n kernel. Convolution with kernel dilation is also known as 'atrous\n convolution'.\n use_bias: whether to add a bias to the output (default: True).\n mask: Optional mask for the weights during masked convolution. The mask must\n be the same shape as the convolution weight matrix.\n dtype: the dtype of the computation (default: infer from input and params).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n precision: numerical precision of the computation see ``jax.lax.Precision``\n for details.\n kernel_init: initializer for the convolutional kernel.\n bias_init: initializer for the bias.\n transpose_kernel: if ``True`` flips spatial axes and swaps the input/output\n channel axes of the kernel.\n promote_dtype: function to promote the dtype of the arrays to the desired\n dtype. The function should accept a tuple of ``(inputs, kernel, bias)``\n and a ``dtype`` keyword argument, and return a tuple of arrays with the\n promoted dtype.\n \"\"\"\n\n features: int\n kernel_size: int | Sequence[int]\n strides: Sequence[int] | None = None\n padding: PaddingLike = 'SAME'\n kernel_dilation: Sequence[int] | None = None\n use_bias: bool = True\n mask: Array | None = None\n dtype: Dtype | None = None\n param_dtype: Dtype = jnp.float32\n precision: PrecisionLike = None\n kernel_init: Initializer = default_kernel_init\n bias_init: Initializer = initializers.zeros_init()\n transpose_kernel: bool = False\n promote_dtype: PromoteDtypeFn = promote_dtype\n preferred_element_type: Dtype | None = None\n\n @compact\n def __call__(self, inputs: Array) -> Array:\n \"\"\"Applies a transposed convolution to the inputs.\n\n Behaviour mirrors of ``jax.lax.conv_transpose``.\n\n Args:\n inputs: input data with dimensions ``(*batch_dims, spatial_dims..., features).``\n This is the channels-last convention, i.e. NHWC for a 2d convolution and NDHWC\n for a 3D convolution. Note: this is different from the input convention used by\n ``lax.conv_general_dilated``, which puts the spatial dimensions last.\n Note: If the input has more than 1 batch dimension, all batch dimensions\n are flattened into a single dimension for the convolution and restored\n before returning. In some cases directly vmap'ing the layer may yield\n better performance than this default flattening approach. If the input\n lacks a batch dimension it will be added for the convolution and removed\n n return, an allowance made to enable writing single-example code.\n\n Returns:\n The convolved data.\n \"\"\"\n kernel_size: tuple[int, ...]\n if isinstance(self.kernel_size, int):\n kernel_size = (self.kernel_size,)\n else:\n kernel_size = tuple(self.kernel_size)\n\n def maybe_broadcast(\n x: int | Sequence[int] | None,\n ) -> tuple[int, ...]:\n if x is None:\n # backward compatibility with using None as sentinel for\n # broadcast 1\n x = 1\n if isinstance(x, int):\n return (x,) * len(kernel_size)\n return tuple(x)\n\n # Combine all input batch dimensions into a single leading batch axis.\n num_batch_dimensions = inputs.ndim - (len(kernel_size) + 1)\n if num_batch_dimensions != 1:\n input_batch_shape = inputs.shape[:num_batch_dimensions]\n flat_input_shape = (-1,) + inputs.shape[\n num_batch_dimensions:\n ]\n inputs = jnp.reshape(inputs, flat_input_shape)\n\n strides = maybe_broadcast(self.strides)\n kernel_dilation = maybe_broadcast(self.kernel_dilation)\n\n in_features = jnp.shape(inputs)[-1]\n if self.transpose_kernel:\n kernel_shape = kernel_size + (self.features, in_features)\n else:\n kernel_shape = kernel_size + (in_features, self.features)\n\n if self.mask is not None and self.mask.shape != kernel_shape:\n raise ValueError(\n 'Mask needs to have the same shape as weights. '\n f'Shapes are: {self.mask.shape}, {kernel_shape}'\n )\n\n kernel = self.param(\n 'kernel', self.kernel_init, kernel_shape, self.param_dtype\n )\n\n if self.mask is not None:\n kernel *= self.mask\n\n padding_lax = canonicalize_padding(self.padding, len(kernel_size))\n if padding_lax == 'CIRCULAR':\n padding_lax = 'VALID'\n\n if self.use_bias:\n bias = self.param(\n 'bias', self.bias_init, (self.features,), self.param_dtype\n )\n else:\n bias = None\n\n inputs, kernel, bias = self.promote_dtype(\n inputs, kernel, bias, dtype=self.dtype\n )\n assert inputs is not None\n assert kernel is not None\n\n y = lax.conv_transpose(\n inputs,\n kernel,\n strides,\n padding_lax,\n rhs_dilation=kernel_dilation,\n transpose_kernel=self.transpose_kernel,\n precision=self.precision,\n preferred_element_type=self.preferred_element_type,\n )\n\n if self.padding == 'CIRCULAR':\n # For circular padding, we need to identify the size of the final output\n # (\"period\") along each spatial dimension, pad each dimension to an\n # integer number of periods, and wrap the array periodically around each\n # dimension. Padding should be done in such a way that the start of the\n # original input data inside the padded array is located at integer\n # number of periods - otherwise the result would be circularly shifted.\n\n # Compute period along each spatial dimension - it's input size scaled\n # by the stride.\n scaled_x_dims = [\n x_dim * stride\n for x_dim, stride in zip(jnp.shape(inputs)[1:-1], strides)\n ]\n # Compute difference between the current size of y and the final output\n # size, and complement this difference to 2 * period - that gives how\n # much we need to pad.\n size_diffs = [\n -(y_dim - x_dim) % (2 * x_dim)\n for y_dim, x_dim in zip(y.shape[1:-1], scaled_x_dims)\n ]\n if self.transpose_kernel:\n # If the kernel is transposed, the \"+1\" is put on the right to\n # mirror the regular convolution. If the same kernel parameters are used\n # as for Conv, this layer then computes the proper transpose convolution.\n total_pad = [\n (size_diff // 2, (size_diff + 1) // 2) for size_diff in size_diffs\n ]\n else:\n # Divide the padding equally between left and right. The choice to put\n # \"+1\" on the left (and not on the right) represents a convention for\n # aligning even-sized kernels.\n total_pad = [\n ((size_diff + 1) // 2, size_diff // 2) for size_diff in size_diffs\n ]\n y = jnp.pad(y, [(0, 0)] + total_pad + [(0, 0)])\n # Wrap the result periodically around each spatial dimension,\n # one by one.\n for i in range(1, y.ndim - 1):\n y = y.reshape(\n y.shape[:i] + (-1, scaled_x_dims[i - 1]) + y.shape[i + 1 :]\n )\n y = y.sum(axis=i)\n\n if self.use_bias:\n y += jnp.reshape(bias, (1,) * (y.ndim - 1) + (-1,)) # type: ignore\n\n if num_batch_dimensions != 1:\n output_shape = input_batch_shape + y.shape[1:]\n y = jnp.reshape(y, output_shape)\n\n return y\n\n\ndefault_embed_init = initializers.variance_scaling(\n 1.0, 'fan_in', 'normal', out_axis=0\n)\n\n\nclass Embed(Module):\n \"\"\"Embedding Module.\n\n A parameterized function from integers [0, ``num_embeddings``) to\n ``features``-dimensional vectors. This ``Module`` will create an ``embedding``\n matrix with shape ``(num_embeddings, features)``. When calling this layer,\n the input values will be used to 0-index into the ``embedding`` matrix.\n Indexing on a value greater than or equal to ``num_embeddings`` will result\n in ``nan`` values. When ``num_embeddings`` equals to 1, it will\n broadcast the ``embedding`` matrix to input shape with ``features``\n dimension appended.\n\n Example usage::\n\n >>> import flax.linen as nn\n >>> import jax, jax.numpy as jnp\n\n >>> layer = nn.Embed(num_embeddings=5, features=3)\n >>> indices_input = jnp.array([[0, 1, 2], [-1, -2, -3]])\n >>> variables = layer.init(jax.random.key(0), indices_input)\n >>> variables\n {'params': {'embedding': Array([[ 0.04396089, -0.9328513 , -0.97328115],\n [ 0.41147125, 0.66334754, 0.49469155],\n [ 0.09719624, 0.49861377, 0.49519277],\n [-0.13316602, 0.6697022 , 0.3710195 ],\n [-0.5039532 , 0.287319 , 1.4369922 ]], dtype=float32)}}\n >>> # get the first three and last three embeddings\n >>> layer.apply(variables, indices_input)\n Array([[[ 0.04396089, -0.9328513 , -0.97328115],\n [ 0.41147125, 0.66334754, 0.49469155],\n [ 0.09719624, 0.49861377, 0.49519277]],\n \n [[-0.5039532 , 0.287319 , 1.4369922 ],\n [-0.13316602, 0.6697022 , 0.3710195 ],\n [ 0.09719624, 0.49861377, 0.49519277]]], dtype=float32)\n\n Attributes:\n num_embeddings: number of embeddings / vocab size.\n features: number of feature dimensions for each embedding.\n dtype: the dtype of the embedding vectors (default: same as embedding).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n embedding_init: embedding initializer.\n promote_dtype: function to promote the dtype of the arrays to the desired\n dtype. The function should accept a tuple of ``(embedding,)`` during ``__call__``\n or ``(query, embedding)`` during ``attend``, and a ``dtype`` keyword argument,\n and return a tuple of arrays with the promoted dtype.\n \"\"\"\n\n num_embeddings: int\n features: int\n dtype: Dtype | None = None\n param_dtype: Dtype = jnp.float32\n embedding_init: Initializer = default_embed_init\n promote_dtype: PromoteDtypeFn = promote_dtype\n\n def setup(self):\n self.embedding = self.param(\n 'embedding',\n self.embedding_init,\n (self.num_embeddings, self.features),\n self.param_dtype,\n )\n\n def __call__(self, inputs: Array) -> Array:\n \"\"\"Embeds the inputs along the last dimension.\n\n Args:\n inputs: input data, all dimensions are considered batch dimensions.\n Values in the input array must be integers.\n\n Returns:\n Output which is embedded input data. The output shape follows the input,\n with an additional ``features`` dimension appended.\n \"\"\"\n if not jnp.issubdtype(inputs.dtype, jnp.integer):\n raise ValueError('Input type must be an integer or unsigned integer.')\n # Use take because fancy indexing numpy arrays with JAX indices does not\n # work correctly.\n (embedding,) = self.promote_dtype(\n self.embedding, dtype=self.dtype, inexact=False\n )\n assert embedding is not None\n if self.num_embeddings == 1:\n return jnp.broadcast_to(embedding, inputs.shape + (self.features,))\n return jnp.take(embedding, inputs, axis=0)\n\n def attend(self, query: Array) -> Array:\n \"\"\"Attend over the embedding using a query array.\n\n Args:\n query: array with last dimension equal the feature depth ``features`` of the\n embedding.\n\n Returns:\n An array with final dim ``num_embeddings`` corresponding to the batched\n inner-product of the array of query vectors against each embedding.\n Commonly used for weight-sharing between embeddings and logit transform\n in NLP models.\n \"\"\"\n embedding: Array\n query, embedding = self.promote_dtype(\n query, self.embedding, dtype=self.dtype\n )\n assert query is not None\n assert embedding is not None\n return jnp.dot(query, embedding.T)\n" }, { "path": "flax/linen/module.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Flax Module.\"\"\"\n\nimport contextlib\nimport dataclasses\nimport enum\nimport functools\nimport inspect\nimport sys\nimport threading\nimport typing\nimport weakref\nfrom types import MappingProxyType\nfrom typing import (\n Any,\n Literal,\n Optional,\n TypeVar,\n Union,\n overload,\n)\nfrom collections.abc import Callable, Iterable, Iterator, Mapping\n\nimport jax\nimport jax.numpy as jnp\nimport typing_extensions as tpe\n\nimport flax\nimport flax.linen as nn\nfrom flax import (\n config,\n core,\n errors,\n serialization,\n traceback_util,\n traverse_util,\n)\nfrom flax.core import Scope, meta, partial_eval\nfrom flax.core.frozen_dict import FrozenDict\nfrom flax.core.scope import (\n CollectionFilter,\n DenyList,\n Variable,\n union_filters,\n)\nfrom flax.ids import FlaxId, uuid\nfrom flax.linen import kw_only_dataclasses\nfrom flax.typing import (\n RNGSequences,\n PRNGKey,\n FrozenVariableDict,\n VariableDict,\n)\n\ntraceback_util.register_exclusion(__file__)\n\n\nT = TypeVar('T')\nK = TypeVar('K')\nM = TypeVar('M', bound='Module')\n_CallableT = TypeVar('_CallableT', bound=Callable)\n\n\n# Used for abstractly testing module behavior.\nTestScope = type(\n 'TestScope',\n (Scope,),\n {'make_rng': lambda self, name: jax.random.key(0)},\n)\n\n\n# pylint: disable=protected-access,attribute-defined-outside-init\ndef _get_fn_name(fn):\n if isinstance(fn, functools.partial):\n return _get_fn_name(fn.func)\n return getattr(fn, '__name__', 'unnamed_function')\n\n\ndef _indent(x: str, num_spaces: int):\n indent_str = ' ' * num_spaces\n lines = x.split('\\n')\n # skip last line because it is always empty and should not be indented.\n assert not lines[-1]\n return '\\n'.join(indent_str + line for line in lines[:-1]) + '\\n'\n\n\ndef _attr_repr(value: Any):\n if callable(value) and (\n (isinstance(value, nn.Module) and value.__dict__.get('__name__', None))\n or (not isinstance(value, nn.Module) and getattr(value, '__name__', None))\n ):\n value_rep = value.__name__\n else:\n value_rep = repr(value)\n return value_rep\n\n\ndef _module_repr(module: 'Module', num_spaces: int = 4):\n \"\"\"Returns a pretty printed representation of the module.\"\"\"\n cls = type(module)\n try:\n fields = dataclasses.fields(cls)\n except TypeError:\n # Edge case with no fields e.g. module = nn.Module() causes error later.\n return object.__repr__(module)\n cls_name = cls.__name__\n rep = ''\n\n attributes = {\n f.name: f.type\n for f in fields\n if f.name not in ('parent', 'name') and f.repr\n }\n child_modules = {\n k: v\n for k, v in module._state.children.items() # pytype: disable=attribute-error\n if isinstance(v, Module)\n }\n if attributes:\n rep += '# attributes\\n'\n for attr in attributes.keys():\n # TODO(jheek): can we get a nice string representation of attribute types?\n value = module.__dict__.get(attr, None)\n value_rep = _attr_repr(value)\n rep += f'{attr} = {value_rep}\\n'\n if child_modules:\n rep += '# children\\n'\n for name, child in child_modules.items():\n child_rep = _module_repr(child, num_spaces)\n rep += f'{name} = {child_rep}\\n'\n if rep:\n return f'{cls_name}(\\n{_indent(rep, num_spaces)})'\n else:\n return f'{cls_name}()'\n\n\n# Tabulation utilities.\n# -----------------------------------------------------------------------------\n@dataclasses.dataclass\nclass _CallInfo:\n index: int\n path: tuple[str, ...]\n module: 'Module'\n rngs: dict[str, core.scope.PRNGKey | core.scope.LazyRng] | None\n mutable: bool\n method: str\n args: tuple[Any, ...]\n kwargs: dict[str, Any]\n outputs: Any\n\n\n@dataclasses.dataclass\nclass _CallInfoContext(threading.local):\n index: int\n calls: list[_CallInfo]\n\n def get_call_index(self) -> int:\n index = self.index\n self.index += 1\n return index\n\n\n@contextlib.contextmanager\ndef _tabulate_context():\n _context.call_info_stack.append(_CallInfoContext(0, []))\n try:\n yield\n finally:\n _context.call_info_stack.pop()\n\n\n# Track parent relationship across Modules.\n# -----------------------------------------------------------------------------\nclass _DynamicContext(threading.local):\n \"\"\"Dynamic context.\"\"\"\n\n # TODO(marcvanzee): switch to using contextvars once minimum python version is\n # 3.7\n\n def __init__(self):\n self.module_stack: list['Module' | None] = [\n None,\n ]\n self.capture_stack = []\n self.call_info_stack: list[_CallInfoContext] = []\n\n\n# The global context\n_context = _DynamicContext()\n\n\nclass _Sentinel:\n def __copy__(self):\n return self # Do not copy singleton sentinel.\n\n def __deepcopy__(self, memo):\n del memo\n return self # Do not copy singleton sentinel.\n\n def __reduce__(self):\n return _get_unspecified_parent, ()\n\n\ndef _get_unspecified_parent():\n return _unspecified_parent\n\n\n_unspecified_parent = _Sentinel()\n\n\n# Enable automatic named_call wrapping for labelling profile traces.\n# -----------------------------------------------------------------------------\n_use_named_call = config.flax_profile\n\n\ndef _derive_profiling_name(module, fn):\n fn_name = _get_fn_name(fn)\n method_suffix = f'.{fn_name}' if fn_name != '__call__' else ''\n module_name = module.name or module.__class__.__name__\n return f'{module_name}{method_suffix}'\n\n\ndef enable_named_call():\n \"\"\"Enables named call wrapping for labelling profile traces.\n\n When named call wrapping is enabled all JAX ops executed in a Module\n will be run under ``jax.named_scope``. The ``Module`` class name will\n show up around the operations belonging to that Module in the\n Tensorboard profiling UI, simplifying the profiling process.\n\n Note that ``jax.named_scope`` only works for\n compiled functions (e.g.: using jax.jit or jax.pmap).\n \"\"\"\n global _use_named_call\n _use_named_call = True\n\n\ndef disable_named_call():\n \"\"\"Disables named call wrapping.\n\n See ``enable_named_call``\n \"\"\"\n global _use_named_call\n _use_named_call = False\n\n\n@contextlib.contextmanager\ndef override_named_call(enable: bool = True):\n # pylint: disable=g-doc-return-or-yield\n \"\"\"Returns a context manager that enables/disables named call wrapping.\n\n Args:\n enable: If true, enables named call wrapping for labelling profile traces.\n (see ``enabled_named_call``).\n \"\"\"\n # pylint: enable=g-doc-return-or-yield\n global _use_named_call\n use_named_call_prev = _use_named_call\n _use_named_call = enable\n try:\n yield\n finally:\n _use_named_call = use_named_call_prev\n\n\n# Intercept module methods.\n# -----------------------------------------------------------------------------\n@dataclasses.dataclass(frozen=True)\nclass InterceptorContext:\n \"\"\"Read only state showing the calling context for method interceptors.\n\n Attributes:\n module: The Module instance whose method is being called.\n method_name: The name of the method being called on the module.\n orig_method: The original method defined on the module. Calling it will\n short circuit all other interceptors.\n \"\"\"\n\n module: 'Module'\n method_name: str\n orig_method: Callable[..., Any]\n\n\nclass ThreadLocalStack(threading.local):\n \"\"\"Thread-local stack.\"\"\"\n\n def __init__(self):\n self._storage = []\n\n def push(self, elem: Any) -> None:\n self._storage.append(elem)\n\n def pop(self) -> Any:\n return self._storage.pop()\n\n def __iter__(self) -> Iterator[Any]:\n return iter(reversed(self._storage))\n\n def __len__(self) -> int:\n return len(self._storage)\n\n def __repr__(self) -> str:\n return f'{self.__class__.__name__}({self._storage})'\n\n\nArgs = tuple[Any]\nKwargs = dict[str, Any]\nNextGetter = Callable[..., Any]\nInterceptor = Callable[[NextGetter, Args, Kwargs, InterceptorContext], Any]\n_global_interceptor_stack = ThreadLocalStack()\n\n\n@contextlib.contextmanager\ndef intercept_methods(interceptor: Interceptor):\n # pylint: disable=g-doc-return-or-yield\n r\"\"\"Registers a new method interceptor.\n\n Method interceptors allow you to (at a distance) intercept method calls to\n modules. It works similarly to decorators. You could modify args/kwargs before\n calling the underlying method and/or modify the result returning from calling\n the underlying method. Or you could completely skip calling the underlying\n method and decide to do something differently. For example::\n\n >>> import flax.linen as nn\n >>> import jax.numpy as jnp\n ...\n >>> class Foo(nn.Module):\n ... def __call__(self, x):\n ... return x\n ...\n >>> def my_interceptor1(next_fun, args, kwargs, context):\n ... print('calling my_interceptor1')\n ... return next_fun(*args, **kwargs)\n ...\n >>> foo = Foo()\n >>> with nn.intercept_methods(my_interceptor1):\n ... _ = foo(jnp.ones([1]))\n calling my_interceptor1\n\n You could also register multiple interceptors on the same method. Interceptors\n will run in order. For example::\n\n >>> def my_interceptor2(next_fun, args, kwargs, context):\n ... print('calling my_interceptor2')\n ... return next_fun(*args, **kwargs)\n ...\n >>> with nn.intercept_methods(my_interceptor1), \\\n ... nn.intercept_methods(my_interceptor2):\n ... _ = foo(jnp.ones([1]))\n calling my_interceptor1\n calling my_interceptor2\n\n You could skip other interceptors by directly calling the\n ``context.orig_method``. For example::\n\n >>> def my_interceptor3(next_fun, args, kwargs, context):\n ... print('calling my_interceptor3')\n ... return context.orig_method(*args, **kwargs)\n >>> with nn.intercept_methods(my_interceptor3), \\\n ... nn.intercept_methods(my_interceptor1), \\\n ... nn.intercept_methods(my_interceptor2):\n ... _ = foo(jnp.ones([1]))\n calling my_interceptor3\n\n The following methods couldn't be intercepted:\n\n 1. Methods decoratored with ``nn.nowrap``.\n 2. Dunder methods including ``__eq__``, ``__repr__``, ``__init__``, ``__hash__``, and ``__post_init__``.\n 3. Module dataclass fields.\n 4. Module descriptors.\n\n Args:\n interceptor: A method interceptor.\n \"\"\"\n _global_interceptor_stack.push(interceptor)\n try:\n yield\n finally:\n assert _global_interceptor_stack.pop() is interceptor\n\n\ndef run_interceptors(\n orig_method: Callable[..., Any],\n module: 'Module',\n *args,\n **kwargs,\n) -> Any:\n \"\"\"Runs method interceptors.\"\"\"\n method_name = _get_fn_name(orig_method)\n fun = functools.partial(orig_method, module)\n context = InterceptorContext(module, method_name, fun)\n\n def wrap_interceptor(interceptor, fun):\n \"\"\"Wraps `fun` with `interceptor`.\"\"\"\n\n @functools.wraps(fun)\n def wrapped(*args, **kwargs):\n return interceptor(fun, args, kwargs, context)\n\n return wrapped\n\n # Wraps interceptors around the original method. The innermost interceptor is\n # the last one added and directly wrapped around the original bound method.\n for interceptor in _global_interceptor_stack:\n fun = wrap_interceptor(interceptor, fun)\n return fun(*args, **kwargs)\n\n\n# Utilities for pytrees of Modules defined inside setup()\n# -----------------------------------------------------------------------------\n\n\ndef _sorted_items(x):\n \"\"\"Returns items of a dict ordered by keys.\"\"\"\n return sorted(x.items(), key=lambda x: x[0])\n\n\ndef _get_suffix_value_pairs(\n tree_or_leaf: Any,\n) -> list[tuple[str, type['Module']]]:\n \"\"\"Helper for naming pytrees of submodules.\"\"\"\n dict_or_leaf = serialization.to_state_dict(tree_or_leaf)\n if not isinstance(dict_or_leaf, dict) or not dict_or_leaf:\n return [('', tree_or_leaf)]\n else:\n flat_dict = traverse_util.flatten_dict(dict_or_leaf)\n return [('_' + '_'.join(k), v) for k, v in _sorted_items(flat_dict)]\n\n\ndef _map_over_modules_in_tree(fn, tree_or_leaf):\n \"\"\"Helper for mapping function over submodules.\"\"\"\n dict_or_leaf = serialization.to_state_dict(tree_or_leaf)\n if not isinstance(dict_or_leaf, dict) or not dict_or_leaf:\n return fn('', tree_or_leaf)\n else:\n flat_dict = traverse_util.flatten_dict(dict_or_leaf, keep_empty_nodes=True)\n mapped_flat_dict = {\n k: fn('_' + '_'.join(k), v) for k, v in _sorted_items(flat_dict)\n }\n return serialization.from_state_dict(\n tree_or_leaf, traverse_util.unflatten_dict(mapped_flat_dict)\n )\n\n\ndef _freeze_attr(val: Any) -> Any:\n \"\"\"Recursively wrap the given attribute `var` in ``FrozenDict``.\"\"\"\n if isinstance(val, (dict, FrozenDict)):\n return FrozenDict({k: _freeze_attr(v) for k, v in val.items()})\n elif isinstance(val, tuple):\n # Special case namedtuples and special JAX tuple structures otherwise they\n # would be downgraded to normal tuples.\n if hasattr(val, '_fields') or type(val).__name__ == 'PartitionSpec':\n return type(val)(*[_freeze_attr(v) for v in val])\n else:\n return tuple(_freeze_attr(v) for v in val)\n elif isinstance(val, list):\n return tuple(_freeze_attr(v) for v in val)\n else:\n return val\n\n\n# Method wrapping of \"compact methods\" and setup()\n# -----------------------------------------------------------------------------\ndef compact(fun: _CallableT) -> _CallableT:\n \"\"\"Marks the given module method allowing inlined submodules.\n\n Methods wrapped in @compact can define submodules directly within the method.\n\n For instance::\n\n >>> import flax.linen as nn\n\n >>> class Foo(nn.Module):\n ... @nn.compact\n ... def __call__(self, x, features):\n ... x = nn.Dense(features)(x)\n ... ...\n ... return x\n\n At most one method in each Module may be wrapped with @compact.\n\n Args:\n fun: The Module method to mark as compact.\n\n Returns:\n The given function ``fun`` marked as compact.\n \"\"\"\n fun.compact = True # type: ignore[attr-defined]\n return fun\n\n\ndef nowrap(fun: _CallableT) -> _CallableT:\n \"\"\"Marks the given module method as a helper method that needn't be wrapped.\n\n Methods wrapped in ``@nowrap`` are private helper methods that needn't be wrapped\n with the state handler or a separate named_call transform.\n\n This is needed in several concrete instances:\n - if you're subclassing a method like Module.param and don't want this\n overridden core function decorated with the state management wrapper.\n - If you want a method to be callable from an unbound Module (e.g.: a\n function of construction of arguments that doesn't depend on params/RNGs).\n If you want to learn more about how Flax Modules manage their state read the\n [The Flax Module lifecycle](https://flax.readthedocs.io/en/latest/developer_notes/module_lifecycle.html)\n guide.\n\n For instance::\n\n >>> import flax.linen as nn\n >>> import jax, jax.numpy as jnp\n\n >>> class Foo(nn.Module):\n ... num_features: int\n\n ... @nn.nowrap\n ... def _make_dense(self, num_features):\n ... return nn.Dense(num_features)\n\n ... @nn.compact\n ... def __call__(self, x):\n ... # now safe to use constructor helper even if using named_call\n ... dense = self._make_dense(self.num_features)\n ... return dense(x)\n\n Args:\n fun: The Module method to mark as nowrap.\n\n Returns:\n The given function ``fun`` marked as nowrap.\n \"\"\"\n fun.nowrap = True # type: ignore[attr-defined]\n return fun\n\n\ndef compact_name_scope(fun: _CallableT) -> _CallableT:\n \"\"\"Creates compact submodules from a method.\n\n This is a decorator that allows you to define compact submodules from a\n method. It's intention is to make it easier to port code Haiku code to Flax\n by providing the same functionality.\n\n Example::\n\n >>> import flax.linen as nn\n >>> import jax\n >>> import jax.numpy as jnp\n >>> from flax.core import pretty_repr\n ...\n >>> class Foo(nn.Module):\n ... @nn.compact_name_scope\n ... def up(self, x):\n ... return nn.Dense(3)(x)\n ...\n ... @nn.compact_name_scope\n ... def down(self, x):\n ... return nn.Dense(3)(x)\n ...\n ... def __call__(self, x):\n ... return self.up(x) + self.down(x)\n ...\n >>> module = Foo()\n >>> variables = module.init(jax.random.PRNGKey(0), jnp.ones((1, 2)))\n >>> params = variables['params']\n >>> print(pretty_repr(jax.tree_util.tree_map(jnp.shape, params)))\n {\n down: {\n Dense_0: {\n bias: (3,),\n kernel: (2, 3),\n },\n },\n up: {\n Dense_0: {\n bias: (3,),\n kernel: (2, 3),\n },\n },\n }\n\n You can also use ``compact_name_scope`` inside ``@compact`` methods or even\n other\n ``compact_name_scope`` methods. Methods that are decorated with\n ``compact_name_scope``\n can also be called directly from ``init`` or ``apply`` via the ``method``\n argument::\n\n >>> y_down = module.apply({'params': params}, jnp.ones((1, 2)), method='down')\n >>> y_down.shape\n (1, 3)\n\n Args:\n fun: The Module method to mark as compact_name_scope.\n\n Returns:\n The given function ``fun`` marked as compact_name_scope.\n \"\"\"\n\n @functools.wraps(fun)\n def compact_name_scope_wrapper(self: nn.Module, *args, **kwargs):\n name = fun.__name__\n if not hasattr(self, '_compact_name_scope_modules'):\n raise ValueError(\n f'Cannot call compact_name_scope method {name!r} on a Module that has not been '\n f'setup. This is likely because you are calling {name!r} '\n 'from outside of init or apply.'\n )\n module = self._compact_name_scope_modules[name]\n return module(*args, **kwargs)\n\n compact_name_scope_wrapper.compact_name_scope = True # type: ignore[attr-defined]\n compact_name_scope_wrapper.inner_fun = fun # type: ignore[attr-defined]\n compact_name_scope_wrapper.nowrap = True # type: ignore[attr-defined]\n return compact_name_scope_wrapper # type: ignore[return-value]\n\n\ndef _get_local_method_names(\n cls: Any, exclude: Iterable[str] = ()\n) -> tuple[str, ...]:\n \"\"\"Gets method names of a class, excluding class and static methods.\n\n Args:\n cls: The class to get method names for.\n exclude: Names to exclude from output.\n\n Returns:\n A list of method names.\n \"\"\"\n true_methods = set()\n for m in cls.__dict__:\n if callable(cls.__dict__[m]) and not inspect.isclass(\n cls.__dict__[m]\n ): # pytype: disable=not-supported-yet\n mtype = type(cls.__dict__[m])\n if mtype != staticmethod and mtype != classmethod:\n true_methods.add(m)\n return tuple(true_methods.difference(set(exclude)))\n\n\ndef _get_local_descriptor_names(\n cls: Any, exclude: Iterable[str] = ()\n) -> tuple[str, ...]:\n \"\"\"Gets descriptor names of a class.\n\n Args:\n cls: The class to get property names for.\n exclude: Names to exclude from output.\n\n Returns:\n A list of property names.\n \"\"\"\n true_properties = set()\n for m, attr in cls.__dict__.items():\n if not callable(attr) and (\n hasattr(attr, '__get__')\n or hasattr(attr, '__set__')\n or hasattr(attr, '__delete__')\n ):\n mtype = type(attr)\n if mtype != staticmethod and mtype != classmethod:\n true_properties.add(m)\n return tuple(true_properties.difference(set(exclude)))\n\n\ndef wrap_method_once(fun: Callable[..., Any]) -> Callable[..., Any]:\n \"\"\"Manages Module state for a given user-defined method.\n\n Args:\n fun: User-defined Module method to manage state for.\n\n Returns:\n Wrapped method.\n \"\"\"\n # Don't rewrap methods that have already had the state management wrapper\n # applied in the decorator stack. This wrapper should always be applied\n # before transformation wrappers.\n if hasattr(fun, 'method_handler_wrapped'):\n return fun\n\n @functools.wraps(fun)\n def wrapped_module_method(*args, **kwargs):\n # We might have incorrectly wrappped a callable\n # that is not a method. Check whether the first arg is self,\n # otherwise call the wrapped function as is.\n if args and isinstance(args[0], Module):\n self, args = args[0], args[1:]\n return self._call_wrapped_method(fun, args, kwargs)\n else:\n return fun(*args, **kwargs)\n\n wrapped_module_method.method_handler_wrapped = True # type: ignore[attr-defined]\n return wrapped_module_method\n\n\ndef wrap_descriptor_once(descriptor) -> 'DescriptorWrapper':\n \"\"\"Wraps a descriptor to give better error messages.\n\n Args:\n descriptor: User-defined Module attribute descriptor.\n\n Returns:\n Wrapped descriptor.\n \"\"\"\n # Don't rewrap descriptors.\n if isinstance(descriptor, DescriptorWrapper):\n return descriptor\n\n return create_descriptor_wrapper(descriptor)\n\n\ndef _wrap_hash(hash_fn: Callable[..., Any]) -> Callable[..., Any]:\n \"\"\"Wraps a hash function with some check for Flax Modules.\"\"\"\n\n @functools.wraps(hash_fn)\n def wrapped(self):\n if self.scope is not None:\n raise TypeError(\"Can't call __hash__ on modules that hold variables.\")\n try:\n hash_value = hash_fn(self)\n except TypeError as exc:\n raise TypeError(\n 'Failed to hash Flax Module. '\n 'The module probably contains unhashable attributes. '\n f'Module={self}'\n ) from exc\n return hash_value\n\n return wrapped\n\n\ndef _get_unbound_fn(method_or_fn: Callable[..., Any]) -> Callable[..., Any]:\n \"\"\"Returns an unbound function from a method that is possibly bound.\n\n This means that if the passed function belongs of an instance of a class, then\n the returned function does no longer depend on the instance, which is passed\n as the first argument to the function.\n\n Args:\n method_or_fn: A class method or function.\n\n Returns:\n An unbound version of input function.\n \"\"\"\n if inspect.ismethod(method_or_fn) and isinstance(\n method_or_fn.__self__, Module\n ): # pytype: disable=attribute-error\n method_or_fn = method_or_fn.__func__ # pytype: disable=attribute-error\n\n # The method should be callable, and it should have at least one argument\n # representing the class that is passed in.\n if (\n not callable(method_or_fn)\n or len(inspect.signature(method_or_fn).parameters) < 1\n ):\n raise errors.ApplyModuleInvalidMethodError(method_or_fn)\n\n return method_or_fn\n\n\ndef _map_submodules(fn: Callable[['Module'], Any], tree):\n \"\"\"Map a function over all submodules in a tree.\"\"\"\n g = lambda _, x: fn(x) if isinstance(x, Module) else x\n return _freeze_attr(_map_over_modules_in_tree(g, tree))\n\n\nclass SetupState(enum.IntEnum):\n # setup() has not been called.\n NEW = 0\n # setup() has been called outside a transform boundary.\n TRANSFORMED = 1\n # setup() has been called.\n DONE = 2\n\n\n@dataclasses.dataclass\nclass _ModuleInternalState:\n \"\"\"Ephemeral Module Evaluation State.\n\n For clarity, we collect all of the temporary flags and ephemeral state used by\n Modules for autonaming and error messages here, alongside the rules used\n to pass this ephemeral state across transform boundaries.\n \"\"\"\n\n in_compact_method: bool = False\n in_setup: bool = False\n setup_called: SetupState = SetupState.NEW\n is_initialized: bool = False\n autoname_cursor: dict[str, int] = dataclasses.field(default_factory=dict)\n children: dict[str, Union[str, 'Module']] = dataclasses.field(\n default_factory=dict\n )\n\n def reset(self) -> None:\n \"\"\"Resets transient state.\n\n This function is called after each module method, so only attributes that\n are method-dependent are reset.\n \"\"\"\n self.in_compact_method = False\n self.in_setup = False\n self.autoname_cursor = dict()\n\n def export(self) -> '_ModuleInternalState':\n \"\"\"Exports transform-preserved state across transform boundary.\"\"\"\n setup_state = (\n SetupState.TRANSFORMED if self.setup_called else SetupState.NEW\n )\n cloned = _ModuleInternalState(\n in_compact_method=self.in_compact_method,\n in_setup=self.in_setup,\n setup_called=setup_state,\n is_initialized=self.is_initialized,\n autoname_cursor=dict(self.autoname_cursor),\n )\n return cloned\n\n def reimport(self, other: '_ModuleInternalState') -> None:\n \"\"\"Re-imports transform-preserved state from across transform boundary.\"\"\"\n self.in_compact_method = other.in_compact_method\n self.in_setup = other.in_setup\n self.is_initialized = other.is_initialized\n self.autoname_cursor = dict(other.autoname_cursor)\n\n\n_uninitialized_module_internal_state = _ModuleInternalState()\n\n\n_UNDEFINED_COPY_PICKLE_METHODS = (\n '__getstate__',\n '__setstate__',\n '__getnewargs_ex__',\n '__reduce__',\n '__reduce_ex__',\n '__copy__',\n '__deepcopy__',\n)\n\n\n_caches: 'weakref.WeakKeyDictionary[Scope, weakref.WeakValueDictionary[FlaxId, Module]]' = weakref.WeakKeyDictionary()\n\n\ntuple_reduce = lambda xs, x: xs + (x,)\ntuple_init = lambda: ()\n\n\ncapture_call_intermediates = lambda _, method_name: method_name == '__call__'\n\n\nclass ParentDescriptor:\n \"\"\"Wraps parent module references in weak refs.\n\n This prevents reference cycles from forming via parent links which can lead\n to accidental OOMs in eager mode due to slow garbage collection as well as\n spurious tracer leaks during jit compilation.\n\n Note: \"descriptors\" are the underlying python mechanism for implementing\n dynamic @property decorators. We need to use a raw descriptor instead of the\n more common decorator in order to force that the appropriate getter/setter\n logic applies in subclasses even after various dataclass transforms.\n \"\"\"\n\n def __get__(self, obj, objtype=None):\n # check if obj is None, happens during %autoreload\n if obj is None:\n return None\n parent = object.__getattribute__(obj, '_parent_ref')\n return parent() if isinstance(parent, weakref.ReferenceType) else parent\n\n def __set__(self, obj, value):\n maybe_weak = weakref.ref(value) if isinstance(value, Module) else value\n object.__setattr__(obj, '_parent_ref', maybe_weak)\n\n\nclass Descriptor(tpe.Protocol):\n __isabstractmethod__: bool\n\n def __get__(self, obj, objtype=None) -> Any:\n ...\n\n def __set__(self, obj, value) -> None:\n ...\n\n def __delete__(self, obj) -> None:\n ...\n\n def __set_name__(self, owner, name) -> None:\n ...\n\n\nclass DescriptorWrapper:\n pass\n\n\ndef create_descriptor_wrapper(descriptor: Descriptor):\n \"\"\"Creates a descriptor wrapper that calls a get_fn on the descriptor.\"\"\"\n\n class _DescriptorWrapper(DescriptorWrapper):\n \"\"\"A descriptor that can wrap any descriptor.\"\"\"\n\n if hasattr(descriptor, '__isabstractmethod__'):\n __isabstractmethod__ = descriptor.__isabstractmethod__\n\n def __init__(self, wrapped: Descriptor):\n self.wrapped = wrapped\n\n # conditionally define descriptor methods\n if hasattr(descriptor, '__get__'):\n\n def __get__(self, *args, **kwargs):\n # here we will catch internal AttributeError and re-raise it as a\n # more informative and correct error message.\n try:\n return self.wrapped.__get__(*args, **kwargs)\n except AttributeError as e:\n raise errors.DescriptorAttributeError() from e\n\n if hasattr(descriptor, '__set__'):\n\n def __set__(self, *args, **kwargs):\n return self.wrapped.__set__(*args, **kwargs)\n\n if hasattr(descriptor, '__delete__'):\n\n def __delete__(self, *args, **kwargs):\n return self.wrapped.__delete__(*args, **kwargs)\n\n if hasattr(descriptor, '__set_name__'):\n\n def __set_name__(self, *args, **kwargs):\n self.wrapped.__set_name__(*args, **kwargs)\n\n def __getattr__(self, name):\n if 'wrapped' not in vars(self):\n raise AttributeError()\n return getattr(self.wrapped, name)\n\n return _DescriptorWrapper(descriptor)\n\n\n# Base Module definition.\n# -----------------------------------------------------------------------------\n\n\ndef module_field(*, kw_only: bool = False, default: Any | None = ...) -> Any:\n ...\n\n\n# The ModuleBase class is created only to make static analyzers happy\n# mainly pytype and pyright. Some notes:\n# * pyright (correctly) complains that Module itself is not a dataclass, even\n# though all its subclasses and intances ARE dataclasses. Because there is no\n# way to annotate this in a way that pyright understands, we create a\n# ModuleBase class decorated with `dataclass_transform` such that pyright\n# thinks Module is a dataclass (in reality only subclasses are instantiated\n# so this is fine).\n# * The `__dataclass_fields__` attribute is needed because pytype seems to\n# not understand the `dataclass_transform` decorator, therefore we need\n# to add the attribute manually.\n# * Other attributes are annotated for completeness. Because we are using\n# the `if typing.TYPE_CHECKING` pattern, these annotations are not present\n# at runtime so they don't affect the dataclass behavior.\n@tpe.dataclass_transform(field_specifiers=(module_field,)) # type: ignore[literal-required]\nclass ModuleBase:\n if typing.TYPE_CHECKING:\n scope: Scope | None\n _state: _ModuleInternalState\n _parent_ref: Union['Module', weakref.ReferenceType['Module'], None]\n __dataclass_fields__: dict[str, dataclasses.Field]\n\n\nclass Module(ModuleBase):\n \"\"\"Base class for all neural network modules.\n\n Layers and models should subclass this class.\n\n All Flax Modules are Python 3.7\n `dataclasses `_. Since\n dataclasses take over ``__init__``, you should instead override :meth:`setup`,\n which is automatically called to initialize the module.\n\n Modules can contain submodules, and in this way can be nested in a tree\n structure. Submodels can be assigned as regular attributes inside the\n :meth:`setup` method.\n\n You can define arbitrary \"forward pass\" methods on your Module subclass.\n While no methods are special-cased, ``__call__`` is a popular choice because\n it allows you to use module instances as if they are functions::\n\n >>> from flax import linen as nn\n >>> from typing import Tuple\n\n >>> class Module(nn.Module):\n ... features: Tuple[int, ...] = (16, 4)\n\n ... def setup(self):\n ... self.dense1 = nn.Dense(self.features[0])\n ... self.dense2 = nn.Dense(self.features[1])\n\n ... def __call__(self, x):\n ... return self.dense2(nn.relu(self.dense1(x)))\n\n Optionally, for more concise module implementations where submodules\n definitions are co-located with their usage, you can use the\n :meth:`compact` wrapper.\n \"\"\"\n\n if typing.TYPE_CHECKING:\n name: str | None = module_field(kw_only=True, default=None)\n parent: Union['Module', _Sentinel, None] = module_field(\n kw_only=True, default=None\n )\n\n def __init__(self, *args, **kwargs):\n # this stub makes sure pytype accepts constructor arguments.\n pass\n\n def __call__(self, *args, **kwargs) -> Any:\n # this stub allows pytype to accept Modules as Callables.\n pass\n\n @classmethod\n def __init_subclass__(cls, kw_only: bool = False, **kwargs: Any) -> None:\n \"\"\"Automatically initializes all subclasses as custom dataclasses.\"\"\"\n super().__init_subclass__(**kwargs)\n # All Flax Modules are dataclasses. We force this convention since\n # it encourages the stateless behavior needed to clone module instances for\n # functional transformation. Instead of using a python metaclass, we\n # automatically transform Modules into dataclasses at subclass creation\n # time, and we set the last dataclass arguments to `parent` and `name`.\n cls._customized_dataclass_transform(kw_only)\n # We wrap user-defined methods including setup and __call__ to enforce\n # a number of different checks and to provide clear error messages.\n cls._find_compact_name_scope_methods()\n cls._wrap_module_attributes()\n # Set empty class defaults.\n cls._state = _uninitialized_module_internal_state # type: ignore[attr-defined]\n cls.scope: Scope | None = None # type: ignore\n # Handles weak referencing of parent Modules to prevent reference cycles.\n cls._parent_ref = None # type: ignore[attr-defined]\n cls.parent = ParentDescriptor() # type: ignore[assignment]\n\n @classmethod\n def _customized_dataclass_transform(cls, kw_only: bool):\n \"\"\"Transforms `cls` into a dataclass, with custom additional behavior.\n\n 1. Inject `parent` and `name` fields. (If they are already present,\n then check that they have the expected types.)\n 2. Set compare, hash, and repr to False for non-init fields.\n 3. Generate a hash function (if not provided by cls).\n \"\"\"\n # Check reserved attributes have expected type annotations.\n if sys.version_info < (3, 14):\n annotations = dict(cls.__dict__.get('__annotations__', {}))\n else:\n annotations = inspect.get_annotations(cls)\n if annotations.get('parent', _ParentType) != _ParentType:\n raise errors.ReservedModuleAttributeError(annotations)\n if annotations.get('name', str) not in ('str', str, Optional[str]):\n raise errors.ReservedModuleAttributeError(annotations)\n\n # any non-init field will only be set in setup\n # During __hash__ and __eq__ the field is not set yet\n # so it should not be used in compare, hash or repr.\n for field in annotations:\n field_meta = getattr(cls, field, None)\n if isinstance(field_meta, dataclasses.Field) and not field_meta.init:\n field_meta.compare = False\n field_meta.hash = False\n field_meta.repr = False\n\n extra_fields = [\n (\n 'parent',\n _ParentType,\n kw_only_dataclasses.field(\n repr=False, default=_unspecified_parent, kw_only=True\n ),\n ),\n (\n 'name',\n Optional[str],\n kw_only_dataclasses.field(default=None, kw_only=True),\n ),\n ]\n\n if kw_only:\n if tuple(sys.version_info)[:3] >= (3, 10, 0):\n for (\n name,\n annotation, # pytype: disable=invalid-annotation\n default,\n ) in extra_fields:\n setattr(cls, name, default)\n cls.__annotations__[name] = annotation\n dataclasses.dataclass( # type: ignore[call-overload]\n unsafe_hash='__hash__' not in cls.__dict__,\n repr=False,\n kw_only=True,\n )(cls)\n else:\n raise TypeError('`kw_only` is not available before Py 3.10.')\n else:\n # Now apply dataclass transform (which operates in-place).\n # Do generate a hash function only if not provided by the class.\n kw_only_dataclasses.dataclass(\n cls,\n unsafe_hash='__hash__' not in cls.__dict__,\n repr=False,\n extra_fields=extra_fields,\n ) # pytype: disable=wrong-keyword-args\n\n cls.__hash__ = _wrap_hash(cls.__hash__) # type: ignore[method-assign]\n\n @classmethod\n def _find_compact_name_scope_methods(cls):\n \"\"\"Finds all compact_name_scope methods in the class.\"\"\"\n methods = [m[0] for m in inspect.getmembers(cls, predicate=callable)]\n compact_name_scope_fns = tuple(\n method_name\n for method_name in methods\n if hasattr(getattr(cls, method_name), 'compact_name_scope')\n )\n cls._compact_name_scope_methods = compact_name_scope_fns\n\n @classmethod\n def _wrap_module_attributes(cls):\n \"\"\"Wraps user-defined non-inherited methods and descriptors with state\n\n management functions.\n \"\"\"\n # wrap methods\n method_exclusions = [f.name for f in dataclasses.fields(cls)] + [\n '__eq__',\n '__repr__',\n '__init__',\n '__hash__',\n '__post_init__',\n ]\n for key in _get_local_method_names(cls, exclude=method_exclusions):\n method = getattr(cls, key)\n if hasattr(method, 'nowrap'):\n continue\n setattr(cls, key, wrap_method_once(method))\n\n # wrap descriptors\n descriptor_exclusions = [f.name for f in dataclasses.fields(cls)] + [\n 'parent',\n '__dict__',\n ]\n for key in _get_local_descriptor_names(cls, descriptor_exclusions):\n # don't use getattr here, since it will call the descriptor\n descriptor = cls.__dict__[key]\n if hasattr(descriptor, 'nowrap'):\n continue\n setattr(cls, key, wrap_descriptor_once(descriptor))\n return cls\n\n def _call_wrapped_method(self, fun, args, kwargs):\n \"\"\"Calls a wrapped method.\n\n This function is responsible for setting up the thread local state\n correctly before calling the method and cleaning up afterwards.\n This includes storing intermediates, setup of the compact scope,\n and making sure setup is called before any other method.\n\n Args:\n fun: The wrapped method.\n args: Named arguments passed to ``fun``.\n kwargs: Keyword arguments passed to ``fun``.\n\n Returns:\n The results of calling ``fun``.\n \"\"\"\n is_compact_method = hasattr(fun, 'compact')\n fun_name = _get_fn_name(fun)\n is_setup_method = fun_name == 'setup'\n add_call_info = not is_setup_method and len(_context.call_info_stack) > 0\n # We lazily call setup() only when needed.\n if is_setup_method:\n if self.scope is None:\n raise errors.CallSetupUnboundModuleError()\n is_recurrent = self._state.in_setup\n self._state.in_setup = True\n else:\n self._try_setup()\n\n if is_compact_method:\n if self.scope is None:\n raise errors.CallCompactUnboundModuleError()\n is_recurrent = self._state.in_compact_method\n self._state.in_compact_method = True\n _context.module_stack.append(self)\n try:\n # get call info\n if add_call_info:\n assert self.scope is not None\n call_index = _context.call_info_stack[-1].get_call_index()\n\n if _global_interceptor_stack:\n run_fun = functools.partial(run_interceptors, fun)\n else:\n run_fun = fun\n\n # call method\n if _use_named_call:\n with jax.named_scope(_derive_profiling_name(self, fun)):\n y = run_fun(self, *args, **kwargs)\n else:\n y = run_fun(self, *args, **kwargs)\n\n if _context.capture_stack:\n filter_fn = _context.capture_stack[-1]\n if filter_fn and filter_fn(self, fun_name):\n self.sow('intermediates', fun_name, y)\n if add_call_info:\n _args, _kwargs, _y = flax.linen.summary._represent_tree(\n (args, kwargs, y)\n )\n _context.call_info_stack[-1].calls.append(\n _CallInfo(\n call_index,\n self.path,\n self.clone(),\n self.scope.rngs,\n self.scope.mutable,\n fun.__name__,\n _args,\n _kwargs,\n _y,\n )\n )\n return y\n finally:\n _context.module_stack.pop()\n if is_compact_method:\n object.__setattr__(self, 'scope', self.scope.rewound())\n # setup or compact calls can be recurrent for example due to super calls\n # resetting the state would cause is compact/setup method\n # to be set to False prematurely.\n if (is_compact_method or is_setup_method) and not is_recurrent:\n self._state.reset()\n\n def __setattr__(self, name: str, val: Any):\n \"\"\"Sets an attribute on this Module.\n\n We overload setattr solely to support pythonic naming via assignment of\n submodules in the special :meth:`setup` function::\n\n self.submodule_name = MyModule(...)\n\n We also support lists and other general pytrees, e.g.::\n\n self.submodules = [MyModule0(..), MyModule1(..), ...]\n\n Args:\n name: Attribute to set.\n val: Value of the attribute.\n \"\"\"\n fields = self.__dataclass_fields__ # pytype: disable=attribute-error\n is_dataclass_attr = name in fields and fields[name].init\n\n if not self._state.in_setup:\n if not self._state.is_initialized:\n # Setting attributes before end of Module.__post_init__()\n object.__setattr__(self, name, val)\n return\n else:\n # If the attribute is a python special method, we allow setting it (this\n # is useful e.g. for IPython auto-reload).\n if name.startswith('__'):\n object.__setattr__(self, name, val)\n return\n # We're past all initialization and setup logic:\n # Raises a TypeError just like frozen python dataclasses.\n raise errors.SetAttributeFrozenModuleError(\n self.__class__.__name__, name, val\n )\n\n # We're inside the setup() method:\n if is_dataclass_attr:\n # These names are specified as dataclass fields. They should not be\n # initialized within the setup() method, but can be modified freely\n # before it.\n raise errors.SetAttributeInModuleSetupError()\n\n # Values (that may be variables or submodules) are being defined and\n # attached in setup(), we run some extra logic in that case.\n self._register_submodules(name, val)\n\n def __getattr__(self, name: str) -> Any:\n \"\"\"Call setup() before getting any setup-defined attributes.\"\"\"\n # We don't want to return anything for python copy / pickle methods.\n if name in _UNDEFINED_COPY_PICKLE_METHODS:\n raise AttributeError()\n self._try_setup()\n if name in self.__dict__:\n return self.__dict__[name]\n else:\n msg = f'\"{self.__class__.__name__}\" object has no attribute \"{name}\".'\n if self.scope is None:\n msg += (\n f' If \"{name}\" is defined in \\'.setup()\\', remember these fields '\n \"are only accessible from inside 'init' or 'apply'.\"\n )\n raise AttributeError(msg)\n\n def __dir__(self) -> list[str]:\n \"\"\"Call setup() before listing attributes.\"\"\"\n self._try_setup()\n return object.__dir__(self) # type: ignore\n\n def __post_init__(self) -> None:\n # DO NOT REMOVE - Marker for internal logging.\n # In dataclasses, __init__ is overridden to process dataclass arguments,\n # and __post_init__ is called immediately afterwards. Here, depending on the\n # type of `parent` passed to initialize the Module, we either defer\n # initialization, attach this Module as a submodule of a parent, or bind\n # this Module at the top-level to variables and rngs.\n\n object.__setattr__(self, '_id', uuid())\n object.__setattr__(self, '_state', _ModuleInternalState())\n\n # Typically we set the parent based on the dynamic module context.\n if self.parent is _unspecified_parent: # pytype: disable=attribute-error\n object.__setattr__(self, 'parent', _context.module_stack[-1])\n\n # Initialization is deferred for top level Modules or any other \"orphan\"\n # Modules until attachment by __setattr__ i.e. MyModule(..., parent=None)\n if self.parent is None:\n return\n\n # Register submodule on parent Module.\n if isinstance(self.parent, Module):\n # When initializing an unnamed Module inside setup()\n # initialization is deferred until attachment by __setattr__\n # i.e. self.mymodule = MyModule(...)\n self.name: str | None\n if (\n self.parent._state.in_setup and self.name is None\n ): # pytype: disable=attribute-error\n return\n if not self.parent._initialization_allowed:\n raise errors.AssignSubModuleError(self.__class__.__name__)\n # Autonaming of submodules.\n if self.name is None: # pytype: disable=attribute-error\n prefix = f'{self.__class__.__name__}'\n cursor = self.parent._state.autoname_cursor.get(prefix, 0)\n self.name = f'{prefix}_{cursor}'\n self.parent._state.autoname_cursor[prefix] = cursor + 1\n # Allow scope aliasing under transforms for submodules defined in setup.\n reuse_scopes = (\n self.parent._state.in_setup\n and self.parent._state.setup_called == SetupState.TRANSFORMED\n )\n # Perform name-collision check.\n if self.parent._name_taken(self.name, reuse_scopes=reuse_scopes):\n parent_class = self.parent.__class__.__name__\n raise errors.NameInUseError('submodule', self.name, parent_class)\n # Finalize attachment to parent and scope initialization.\n self.parent._state.children[self.name] = self\n assert self.parent.scope is not None\n object.__setattr__(\n self, 'scope', self.parent.scope.push(self.name, reuse=reuse_scopes)\n )\n\n # Top-level invocation with a functional Scope.\n elif isinstance(self.parent, Scope):\n object.__setattr__(self, 'scope', self.parent)\n else:\n raise ValueError('parent must be None, Module or Scope')\n\n # eagerly bind submodules if scope is available\n if self.scope is not None:\n for field in dataclasses.fields(self):\n if field.name not in ('parent', 'name') and field.init:\n self._register_submodules(field.name, getattr(self, field.name))\n\n self._state.is_initialized = True\n\n def __repr__(self) -> str:\n return _module_repr(self)\n\n def setup(self) -> None:\n \"\"\"Initializes a Module lazily (similar to a lazy ``__init__``).\n\n ``setup`` is called once lazily on a module instance when a module\n is bound, immediately before any other methods like ``__call__`` are\n invoked, or before a ``setup``-defined attribute on ``self`` is accessed.\n\n This can happen in three cases:\n\n 1. Immediately when invoking :meth:`apply`, :meth:`init` or\n :meth:`init_and_output`.\n\n 2. Once the module is given a name by being assigned to an attribute of\n another module inside the other module's ``setup`` method\n (see :meth:`__setattr__`)::\n\n >>> class MyModule(nn.Module):\n ... def setup(self):\n ... submodule = nn.Conv(...)\n\n ... # Accessing `submodule` attributes does not yet work here.\n\n ... # The following line invokes `self.__setattr__`, which gives\n ... # `submodule` the name \"conv1\".\n ... self.conv1 = submodule\n\n ... # Accessing `submodule` attributes or methods is now safe and\n ... # either causes setup() to be called once.\n\n 3. Once a module is constructed inside a method wrapped with\n :meth:`compact`, immediately before another method is called or\n ``setup`` defined attribute is accessed.\n \"\"\"\n pass\n\n def _register_submodules(self, name, val):\n \"\"\"Registers a submodule.\"\"\"\n assert self.scope, 'Trying to register submodules on unbound scope.'\n root = self.scope.root\n cache = _caches.get(root, weakref.WeakValueDictionary())\n _caches[root] = cache\n queue = []\n preserve_adopted_names = config.flax_preserve_adopted_names\n if hasattr(type(self), 'preserve_adopted_names'):\n preserve_adopted_names = type(self).preserve_adopted_names\n\n def adopt_attr_modules(cache, queue, suffix, subvalue):\n if isinstance(subvalue, Module):\n current_name = subvalue.name\n adopted_name = None\n if subvalue.parent is None:\n # Preserve sharing-by-reference relationships during adoption\n # via cache keyed on unique instance ids.\n key = subvalue._id\n # Module was passed from outside. It needs to be cloned.\n # Outside modules are named by attachment, not an outer name,\n # UNLESS we're using new adopted name policy, in which case an existing\n # name will be used, as is often supplied by config systems.\n if preserve_adopted_names:\n adopted_name = object.__getattribute__(subvalue, 'name')\n if key in cache:\n subvalue = cache[key]\n else:\n subvalue = subvalue.clone(name=None)\n cache[key] = subvalue\n if subvalue.name is None:\n object.__setattr__(subvalue, 'parent', self)\n if adopted_name is None:\n adopted_name = (\n f'{name}{suffix}'\n if not isinstance(subvalue, CompactNameScope)\n else current_name\n )\n object.__setattr__(subvalue, 'name', adopted_name)\n queue.append(subvalue)\n return subvalue\n\n val = _freeze_attr(\n _map_over_modules_in_tree(\n functools.partial(adopt_attr_modules, cache, queue), val\n )\n )\n object.__setattr__(self, name, val)\n for x in queue:\n x.__post_init__()\n\n def _try_setup(self, shallow: bool = False) -> None:\n \"\"\"Tries to setup module if scope is available and setup has not been called yet.\"\"\"\n if (\n self.scope\n and not self._state.in_setup\n and self._state.setup_called != SetupState.DONE\n ):\n try:\n self._state.in_setup = True\n # A shallow setup will only register attribute submodules but it does\n # not call the user's setup. This avoids running before a\n # transformation.\n for field in dataclasses.fields(self):\n if field.name not in ('parent', 'name') and field.init:\n self._register_submodules(field.name, getattr(self, field.name))\n if not shallow:\n self.setup()\n # create NonTransparent Modules\n self._compact_name_scope_modules = {\n name: CompactNameScope(\n getattr(type(self), name).inner_fun, lambda: self, name=name\n )\n for name in self._compact_name_scope_methods\n }\n\n # We run static checks abstractly once for setup before any transforms\n # to detect name collisions and other python errors.\n elif self._state.setup_called == SetupState.NEW:\n self._validate_setup()\n finally:\n self._state.in_setup = False\n if not shallow:\n self._state.setup_called = SetupState.DONE\n\n def _validate_setup(self) -> None:\n \"\"\"Abstractly evaluates setup only to run static checks.\"\"\"\n\n def run_setup_only(x):\n wrapped_id = wrap_method_once(lambda m, x: x)\n with TestScope({}, rngs={}, mutable=True).temporary() as root:\n return wrapped_id(self.clone(parent=root), x)\n\n _ = jax.eval_shape(run_setup_only, 0)\n\n def _name_taken(\n self,\n name: str,\n reuse_scopes: bool = False,\n collection: str | None = None,\n ) -> bool:\n assert self.scope is not None\n if reuse_scopes:\n return False\n return self.scope.name_reserved(name, collection)\n\n @property\n def _initialization_allowed(self):\n return (\n not self._state.is_initialized # allow eager attachment in post-init\n or self._state.in_setup\n or self._state.in_compact_method\n )\n\n @property\n def path(self):\n \"\"\"Get the path of this Module. Top-level root modules have an empty path ``()``.\n Note that this method can only be used on bound modules that have a valid scope.\n\n Example usage::\n\n >>> import flax.linen as nn\n >>> import jax, jax.numpy as jnp\n\n >>> class SubModel(nn.Module):\n ... @nn.compact\n ... def __call__(self, x):\n ... print(f'SubModel path: {self.path}')\n ... return x\n\n >>> class Model(nn.Module):\n ... @nn.compact\n ... def __call__(self, x):\n ... print(f'Model path: {self.path}')\n ... return SubModel()(x)\n\n >>> model = Model()\n >>> variables = model.init(jax.random.key(0), jnp.ones((1, 2)))\n Model path: ()\n SubModel path: ('SubModel_0',)\n \"\"\"\n\n if self.scope is None:\n raise ValueError(\"Can't access module paths on unbound modules.\")\n\n return self.scope.path\n\n def clone(\n self: M,\n *,\n parent: Union[Scope, 'Module', _Sentinel] | None = None,\n _deep_clone: bool | weakref.WeakValueDictionary = False,\n _reset_names: bool = False,\n **updates,\n ) -> M:\n \"\"\"Creates a clone of this Module, with optionally updated arguments.\n\n NOTE: end users are encouraged to use the ``copy`` method. ``clone`` is used\n primarily for internal routines, and ``copy`` offers simpler arguments and\n better defaults.\n\n Args:\n parent: The parent of the clone. The clone will have no parent if no\n explicit parent is specified.\n _deep_clone: A boolean or a weak value dictionary to control deep cloning\n of submodules. If True, submodules will be cloned recursively. If a weak\n value dictionary is passed, it will be used to cache cloned submodules.\n This flag is used by init/apply/bind to avoid scope leakage.\n _reset_names: If True, ``name=None`` is also passed to submodules when\n cloning. Resetting names in submodules is necessary when calling ``.unbind``.\n **updates: Attribute updates.\n\n Returns:\n A clone of the this Module with the updated attributes and parent.\n \"\"\"\n attrs = {\n f.name: getattr(self, f.name) for f in dataclasses.fields(self) if f.init\n }\n\n attrs.update(parent=parent, **updates)\n\n # Here we implement deep cloning of submodules, this is necessary to avoid scope leakage\n # from external submodules into init/apply/bind while preserving sharing-by-reference\n # relationships between submodules.\n if _deep_clone != False:\n # We use a weak value dictionary to cache cloned submodules. When a shared\n # submodule is cloned, its only cloned once else its fetched from the cache.\n cache = (\n weakref.WeakValueDictionary()\n if isinstance(_deep_clone, bool)\n else _deep_clone\n )\n\n def clone_fn(m: Module) -> Module:\n if hasattr(m, '_id'):\n key = m._id\n if key in cache:\n return cache[key]\n else:\n if _reset_names:\n clone = m.clone(\n _deep_clone=cache, _reset_names=_reset_names, name=None\n )\n else:\n clone = m.clone(_deep_clone=cache)\n cache[key] = clone\n return clone\n else:\n # If the module doesn't have an _id attribute it could be a mock object\n # so we return it as is.\n return m\n\n # _map_submodules will map over all submodules inside attrs\n # value here can be any pytree, non-module values are ignored\n for field_name, value in attrs.items():\n if field_name == 'parent':\n continue\n attrs[field_name] = _map_submodules(clone_fn, value)\n\n module = self.__class__(**attrs)\n\n return module\n\n def copy(\n self: M,\n *,\n parent: Union[Scope, 'Module', _Sentinel] | None = _unspecified_parent,\n name: str | None = None,\n **updates,\n ) -> M:\n \"\"\"Creates a copy of this Module, with optionally updated arguments.\n\n Args:\n parent: The parent of the copy. By default the current module is taken\n as parent if not explicitly specified.\n name: A new name for the copied Module, by default a new automatic name\n will be given.\n **updates: Attribute updates.\n\n Returns:\n A copy of the this Module with the updated name, parent, and attributes.\n \"\"\"\n return self.clone(\n parent=parent, name=name, _deep_clone=True, _reset_names=False, **updates\n )\n\n @overload\n def variable(\n self,\n col: str,\n name: str,\n init_fn: Callable[..., T] | None = None,\n *init_args,\n ) -> Variable[T]:\n ...\n\n @overload\n def variable(\n self,\n col: str,\n name: str,\n init_fn: Callable[..., T] | None = None,\n *init_args,\n unbox: Literal[True],\n **init_kwargs,\n ) -> Variable[T]:\n ...\n\n @overload\n def variable(\n self,\n col: str,\n name: str,\n init_fn: Callable[..., T] | None = None,\n *init_args,\n unbox: Literal[False],\n **init_kwargs,\n ) -> Variable[meta.AxisMetadata[T]]:\n ...\n\n @overload\n def variable(\n self,\n col: str,\n name: str,\n init_fn: Callable[..., T] | None = None,\n *init_args,\n unbox: bool = True,\n **init_kwargs,\n ) -> Variable[T] | Variable[meta.AxisMetadata[T]]:\n ...\n\n def variable(\n self,\n col: str,\n name: str,\n init_fn: Callable[..., T] | None = None,\n *init_args,\n unbox: bool = True,\n **init_kwargs,\n ) -> Variable[T] | Variable[meta.AxisMetadata[T]]:\n \"\"\"Declares and returns a variable in this Module.\n\n See :mod:`flax.core.variables` for more information. See also :meth:`param`\n for a shorthand way to define read-only variables in the \"params\"\n collection.\n\n Contrary to :meth:`param`, all arguments passing using ``init_fn`` should be\n passed on explicitly::\n\n >>> class Foo(nn.Module):\n ... @nn.compact\n ... def __call__(self, x):\n ... x = nn.Dense(4)(x)\n ... key = self.make_rng('stats')\n ... mean = self.variable('stats', 'mean', nn.initializers.lecun_normal(), key, x.shape)\n ... ...\n ... return x * mean.value\n >>> variables = Foo().init({'params': jax.random.key(0), 'stats': jax.random.key(1)}, jnp.ones((2, 3)))\n >>> jax.tree_util.tree_map(jnp.shape, variables)\n {'params': {'Dense_0': {'bias': (4,), 'kernel': (3, 4)}}, 'stats': {'mean': (2, 4)}}\n\n In the example above, the function ``lecun_normal`` expects two arguments:\n ``key`` and ``shape``, and both have to be passed on. The PRNG for ``stats``\n has to be provided explicitly when calling :meth:`init` and :meth:`apply`.\n\n Args:\n col: The variable collection name.\n name: The variable name.\n init_fn: The function that will be called to compute the initial value of\n this variable. This function will only be called the first time this\n variable is used in this module. If None, the variable must already be\n initialized otherwise an error is raised.\n *init_args: The positional arguments to pass to init_fn.\n unbox: If True, ``AxisMetadata`` instances are replaced by their unboxed\n value, see ``flax.nn.meta.unbox`` (default: True).\n **init_kwargs: The key-word arguments to pass to init_fn\n\n Returns:\n A :class:`flax.core.variables.Variable` that can be read or set via\n \".value\" attribute. Throws an error if the variable exists already.\n \"\"\"\n if not self._initialization_allowed:\n raise ValueError(\n 'Variables must be initialized in `setup()` or in a method '\n 'wrapped in `@compact`'\n )\n if self._name_taken(name, collection=col):\n raise errors.NameInUseError('variable', name, self.__class__.__name__)\n assert self.scope is not None\n v = self.scope.variable(\n col, name, init_fn, *init_args, unbox=unbox, **init_kwargs\n )\n self._state.children[name] = col\n return v\n\n @overload\n def param(\n self, name: str, init_fn: Callable[..., T], *init_args,\n ) -> T:\n ...\n\n @overload\n def param(\n self,\n name: str,\n init_fn: Callable[..., meta.AxisMetadata[T]] | Callable[..., T],\n *init_args,\n unbox: Literal[True],\n **init_kwargs,\n ) -> T:\n ...\n\n @overload\n def param(\n self,\n name: str,\n init_fn: Callable[..., T],\n *init_args,\n unbox: Literal[False],\n **init_kwargs,\n ) -> T:\n ...\n\n @overload\n def param(\n self,\n name: str,\n init_fn: Callable[..., T | meta.AxisMetadata[T]],\n *init_args,\n unbox: bool,\n **init_kwargs,\n ) -> T | meta.AxisMetadata[T]:\n ...\n def param(\n self,\n name: str,\n init_fn: Callable[..., T | meta.AxisMetadata[T]],\n *init_args,\n unbox: bool = True,\n **init_kwargs,\n ) -> T | meta.AxisMetadata[T]:\n \"\"\"Declares and returns a parameter in this Module.\n\n Parameters are read-only variables in the collection named \"params\". See\n :mod:`flax.core.variables` for more details on variables.\n\n The first argument of ``init_fn`` is assumed to be a PRNG key, which is\n provided automatically and does not have to be passed using ``init_args``\n or ``init_kwargs``::\n\n >>> class Foo(nn.Module):\n ... @nn.compact\n ... def __call__(self, x):\n ... x = nn.Dense(4)(x)\n ... mean = self.param('mean', nn.initializers.lecun_normal(), x.shape)\n ... ...\n ... return x * mean\n >>> variables = Foo().init({'params': jax.random.key(0), 'stats': jax.random.key(1)}, jnp.ones((2, 3)))\n >>> jax.tree_util.tree_map(jnp.shape, variables)\n {'params': {'Dense_0': {'bias': (4,), 'kernel': (3, 4)}, 'mean': (2, 4)}}\n\n In the example above, the function ``lecun_normal`` expects two arguments:\n ``key`` and ``shape``, but only ``shape`` has to be provided explicitly;\n ``key`` is set automatically using the PRNG for ``params`` that is passed\n when initializing the module using :meth:`init`.\n\n Args:\n name: The parameter name.\n init_fn: The function that will be called to compute the initial value of\n this variable. This function will only be called the first time this\n parameter is used in this module.\n *init_args: The positional arguments to pass to init_fn.\n unbox: If True, ``AxisMetadata`` instances are replaced by their unboxed\n value, see ``flax.nn.meta.unbox`` (default: True).\n **init_kwargs: The key-word arguments to pass to init_fn.\n\n Returns:\n The value of the initialized parameter. Throws an error if the parameter\n exists already.\n \"\"\"\n if not self._initialization_allowed:\n raise ValueError(\n 'Parameters must be initialized in `setup()` or in a method '\n 'wrapped in `@compact`'\n )\n if self._name_taken(name, collection='params'):\n raise errors.NameInUseError('param', name, self.__class__.__name__)\n assert self.scope is not None\n v: T | meta.AxisMetadata[T] = self.scope.param(\n name, init_fn, *init_args, unbox=unbox, **init_kwargs\n )\n self._state.children[name] = 'params'\n return v\n\n def has_variable(self, col: str, name: str) -> bool:\n \"\"\"Checks if a variable of given collection and name exists in this Module.\n\n See :mod:`flax.core.variables` for more explanation on variables and\n collections.\n\n Args:\n col: The variable collection name.\n name: The name of the variable.\n\n Returns:\n True if the variable exists.\n \"\"\"\n if self.scope is None:\n raise ValueError(\"Can't access variables on unbound modules\")\n return self.scope.has_variable(col, name)\n\n def is_mutable_collection(self, col: str) -> bool:\n \"\"\"Returns true if the collection ``col`` is mutable.\"\"\"\n if self.scope is None:\n raise ValueError(\"Can't check mutability on unbound modules\")\n return self.scope.is_mutable_collection(col)\n\n def has_rng(self, name: str) -> bool:\n \"\"\"Returns true if a PRNGSequence with name ``name`` exists.\"\"\"\n if self.scope is None:\n raise ValueError(\"Can't query for RNGs on unbound modules\")\n return self.scope.has_rng(name)\n\n def make_rng(self, name: str = 'params') -> PRNGKey:\n \"\"\"Returns a new RNG key from a given RNG sequence for this Module.\n\n The new RNG key is split from the previous one. Thus, every call to\n ``make_rng`` returns a new RNG key, while still guaranteeing full\n reproducibility.\n\n .. note::\n If an invalid name is passed (i.e. no RNG key was passed by\n the user in ``.init`` or ``.apply`` for this name), then ``name``\n will default to ``'params'``.\n\n Example::\n\n >>> import jax\n >>> import flax.linen as nn\n\n >>> class ParamsModule(nn.Module):\n ... def __call__(self):\n ... return self.make_rng('params')\n >>> class OtherModule(nn.Module):\n ... def __call__(self):\n ... return self.make_rng('other')\n\n >>> key = jax.random.key(0)\n >>> params_out, _ = ParamsModule().init_with_output({'params': key})\n >>> # self.make_rng('other') will default to using the 'params' RNG stream\n >>> other_out, _ = OtherModule().init_with_output({'params': key})\n >>> assert params_out == other_out\n\n Learn more about RNG's by reading the Flax RNG guide:\n https://flax.readthedocs.io/en/latest/guides/flax_fundamentals/rng_guide.html\n\n Args:\n name: The RNG sequence name.\n\n Returns:\n The newly generated RNG key.\n \"\"\"\n if self.scope is None:\n raise ValueError(\"Can't use RNGs on unbound modules\")\n return self.scope.make_rng(name)\n\n def is_initializing(self) -> bool:\n \"\"\"Returns True if running under self.init(...) or nn.init(...)().\n\n This is a helper method to handle the common case of simple initialization\n where we wish to have setup logic occur when only called under\n ``module.init`` or ``nn.init``. For more complicated multi-phase\n initialization scenarios it is better to test for the mutability of\n particular variable collections or for the presence of particular\n variables that potentially need to be initialized.\n \"\"\"\n if self.scope is None:\n raise ValueError(\"Can't check if running under init() on unbound modules\")\n return self.scope.get_flag('initializing', False)\n\n def _module_checks(self):\n \"\"\"Run standard runtime checks.\"\"\"\n\n if not isinstance(self, Module):\n raise errors.InvalidInstanceModuleError()\n\n overridden_post_init = self.__post_init__ != Module.__post_init__\n if overridden_post_init and not hasattr(self, '_id'):\n raise errors.IncorrectPostInitOverrideError()\n\n @traceback_util.api_boundary\n def bind(\n self: M,\n variables: VariableDict,\n *args,\n rngs: RNGSequences | None = None,\n mutable: CollectionFilter = False,\n ) -> M:\n \"\"\"Creates an interactive Module instance by binding variables and RNGs.\n\n ``bind`` provides an \"interactive\" instance of a Module directly without\n transforming a function with ``apply``. This is particularly useful for\n debugging and interactive use cases like notebooks where a function would\n limit the ability to split up code into different cells.\n\n Once the variables (and optionally RNGs) are bound to a ``Module`` it\n becomes a stateful object. Note that idiomatic JAX is functional and\n therefore an interactive instance does not mix well with vanilla JAX APIs.\n ``bind()`` should only be used for interactive experimentation, and in all\n other cases we strongly encourage users to use ``apply()`` instead.\n\n Example::\n\n >>> import jax\n >>> import jax.numpy as jnp\n >>> import flax.linen as nn\n\n >>> class AutoEncoder(nn.Module):\n ... def setup(self):\n ... self.encoder = nn.Dense(3)\n ... self.decoder = nn.Dense(5)\n ...\n ... def __call__(self, x):\n ... return self.decoder(self.encoder(x))\n\n >>> x = jnp.ones((16, 9))\n >>> ae = AutoEncoder()\n >>> variables = ae.init(jax.random.key(0), x)\n >>> model = ae.bind(variables)\n >>> z = model.encoder(x)\n >>> x_reconstructed = model.decoder(z)\n\n Args:\n variables: A dictionary containing variables keyed by variable\n collections. See :mod:`flax.core.variables` for more details about\n variables.\n *args: Named arguments (not used).\n rngs: a dict of PRNGKeys to initialize the PRNG sequences.\n mutable: Can be bool, str, or list. Specifies which collections should be\n treated as mutable: ``bool``: all/no collections are mutable. ``str``:\n The name of a single mutable collection. ``list``: A list of names of\n mutable collections.\n\n Returns:\n A copy of this instance with bound variables and RNGs.\n \"\"\"\n Module._module_checks(self)\n\n del args\n scope = core.bind(variables, rngs=rngs, mutable=mutable)\n return self.clone(parent=scope, _deep_clone=True)\n\n def unbind(self: M) -> tuple[M, VariableDict]:\n \"\"\"Returns an unbound copy of a Module and its variables.\n\n ``unbind`` helps create a stateless version of a bound Module.\n\n An example of a common use case: to extract a sub-Module defined inside\n ``setup()`` and its corresponding variables: 1) temporarily ``bind`` the\n parent Module; and then 2) ``unbind`` the desired sub-Module. (Recall that\n ``setup()`` is only called when the Module is bound.)::\n\n >>> class Encoder(nn.Module):\n ... @nn.compact\n ... def __call__(self, x):\n ... ...\n ... return nn.Dense(256)(x)\n\n >>> class Decoder(nn.Module):\n ... @nn.compact\n ... def __call__(self, x):\n ... ...\n ... return nn.Dense(784)(x)\n\n >>> class AutoEncoder(nn.Module):\n ... def setup(self):\n ... self.encoder = Encoder()\n ... self.decoder = Decoder()\n ...\n ... def __call__(self, x):\n ... return self.decoder(self.encoder(x))\n\n >>> module = AutoEncoder()\n >>> variables = module.init(jax.random.key(0), jnp.ones((1, 784)))\n\n >>> # Extract the Encoder sub-Module and its variables\n >>> encoder, encoder_vars = module.bind(variables).encoder.unbind()\n\n Returns:\n A tuple with an unbound copy of this Module and its variables.\n \"\"\"\n Module._module_checks(self)\n\n if self.scope is None:\n raise errors.CallUnbindOnUnboundModuleError()\n\n variables = self.variables\n module = self.clone(_deep_clone=True, _reset_names=True, name=None)\n return module, variables\n\n @traceback_util.api_boundary\n def apply(\n self,\n variables: VariableDict,\n *args,\n rngs: PRNGKey | RNGSequences | None = None,\n method: Callable[..., Any] | str | None = None,\n mutable: CollectionFilter = False,\n capture_intermediates: bool | Callable[['Module', str], bool] = False,\n **kwargs,\n ) -> Any | tuple[Any, FrozenVariableDict | dict[str, Any]]:\n \"\"\"Applies a module method to variables and returns output and modified variables.\n\n Note that ``method`` should be set if one would like to call ``apply`` on a\n different class method than ``__call__``. For instance, suppose a\n Transformer modules has a method called ``encode``, then the following calls\n ``apply`` on that method::\n\n >>> import flax.linen as nn\n >>> import jax, jax.numpy as jnp\n >>> import numpy as np\n\n >>> class Transformer(nn.Module):\n ... def encode(self, x):\n ... ...\n\n >>> x = jnp.ones((16, 9))\n >>> model = Transformer()\n >>> variables = model.init(jax.random.key(0), x, method=Transformer.encode)\n\n >>> encoded = model.apply(variables, x, method=Transformer.encode)\n\n If a function instance is provided, the unbound function is used. For\n instance, the example below is equivalent to the one above::\n\n >>> encoded = model.apply(variables, x, method=model.encode)\n\n You can also pass a string to a callable attribute of the module. For\n example, the previous can be written as::\n\n >>> encoded = model.apply(variables, x, method='encode')\n\n Note ``method`` can also be a function that is not defined in\n ``Transformer``. In that case, the function should have at least one\n argument representing an instance of the Module class::\n\n >>> def other_fn(instance, x):\n ... # instance.some_module_attr(...)\n ... instance.encode\n ... ...\n\n >>> model.apply(variables, x, method=other_fn)\n\n If you pass a single ``PRNGKey``, Flax will use it to feed the ``'params'``\n RNG stream. If you want to use a different RNG stream or need to use\n multiple streams, you can pass a dictionary mapping each RNG stream name\n to its corresponding ``PRNGKey`` to ``apply``. If ``self.make_rng(name)``\n is called on an RNG stream name that isn't passed by the user, it will\n default to using the ``'params'`` RNG stream.\n\n Example::\n\n >>> class Foo(nn.Module):\n ... @nn.compact\n ... def __call__(self, x, add_noise=False):\n ... x = nn.Dense(16)(x)\n ... x = nn.relu(x)\n ...\n ... if add_noise:\n ... # Add gaussian noise\n ... noise_key = self.make_rng('noise')\n ... x = x + jax.random.normal(noise_key, x.shape)\n ...\n ... return nn.Dense(1)(x)\n\n >>> x = jnp.empty((1, 7))\n >>> module = Foo()\n >>> rngs = {'params': jax.random.key(0), 'noise': jax.random.key(1)}\n >>> variables = module.init(rngs, x)\n >>> out0 = module.apply(variables, x, add_noise=True, rngs=rngs)\n\n >>> rngs['noise'] = jax.random.key(0)\n >>> out1 = module.apply(variables, x, add_noise=True, rngs=rngs)\n >>> # different output (key(1) vs key(0))\n >>> np.testing.assert_raises(AssertionError, np.testing.assert_allclose, out0, out1)\n\n >>> del rngs['noise']\n >>> # self.make_rng('noise') will default to using the 'params' RNG stream\n >>> out2 = module.apply(variables, x, add_noise=True, rngs=rngs)\n >>> # same output (key(0))\n >>> np.testing.assert_allclose(out1, out2)\n\n >>> # passing in a single key is equivalent to passing in {'params': key}\n >>> out3 = module.apply(variables, x, add_noise=True, rngs=jax.random.key(0))\n >>> # same output (key(0))\n >>> np.testing.assert_allclose(out2, out3)\n\n Args:\n variables: A dictionary containing variables keyed by variable\n collections. See :mod:`flax.core.variables` for more details about\n variables.\n *args: Named arguments passed to the specified apply method.\n rngs: a dict of PRNGKeys to initialize the PRNG sequences. The \"params\"\n PRNG sequence is used to initialize parameters.\n method: A function to call apply on. This is generally a function in the\n module. If provided, applies this method. If not provided, applies the\n ``__call__`` method of the module. A string can also be provided to\n specify a method by name.\n mutable: Can be bool, str, or list. Specifies which collections should be\n treated as mutable: ``bool``: all/no collections are mutable. ``str``:\n The name of a single mutable collection. ``list``: A list of names of\n mutable collections.\n capture_intermediates: If ``True``, captures intermediate return values of\n all Modules inside the \"intermediates\" collection. By default, only the\n return values of all ``__call__`` methods are stored. A function can be\n passed to change the filter behavior. The filter function takes the\n Module instance and method name and returns a bool indicating whether\n the output of that method invocation should be stored.\n **kwargs: Keyword arguments passed to the specified apply method.\n\n Returns:\n If ``mutable`` is False, returns output. If any collections are\n mutable, returns ``(output, vars)``, where ``vars`` are is a dict\n of the modified collections.\n \"\"\"\n Module._module_checks(self)\n\n if rngs is not None and not isinstance(rngs, dict):\n if not core.scope._is_valid_rng(rngs):\n raise errors.InvalidRngError(\n 'RNGs should be of shape (2,) or PRNGKey in Module '\n f'{self.__class__.__name__}, but rngs are: {rngs}'\n )\n rngs = {'params': rngs}\n\n if isinstance(method, str):\n attribute_name = method\n method = getattr(self, attribute_name)\n if not callable(method):\n class_name = type(self).__name__\n raise TypeError(\n f\"'{class_name}.{attribute_name}' must be a callable, got\"\n f' {type(method)}.'\n )\n # if the `method` string is a submodule, we create a lambda function\n # that calls the submodule, forwarding all arguments.\n if isinstance(method, Module):\n method = lambda self, *args, **kwargs: getattr(self, attribute_name)(\n *args, **kwargs\n )\n elif method is None:\n method = self.__call__\n method = _get_unbound_fn(method)\n return apply(\n method,\n self,\n mutable=mutable,\n capture_intermediates=capture_intermediates,\n )(variables, *args, **kwargs, rngs=rngs)\n\n @traceback_util.api_boundary\n def init_with_output(\n self,\n rngs: PRNGKey | RNGSequences,\n *args,\n method: Callable[..., Any] | str | None = None,\n mutable: CollectionFilter = DenyList('intermediates'),\n capture_intermediates: bool | Callable[['Module', str], bool] = False,\n **kwargs,\n ) -> tuple[Any, FrozenVariableDict | dict[str, Any]]:\n \"\"\"Initializes a module method with variables and returns output and modified variables.\n\n Args:\n rngs: The rngs for the variable collections.\n *args: Named arguments passed to the init function.\n method: An optional method. If provided, applies this method. If not\n provided, applies the ``__call__`` method. A string can also be\n provided to specify a method by name.\n mutable: Can be bool, str, or list. Specifies which collections should be\n treated as mutable: ``bool``: all/no collections are mutable. ``str``:\n The name of a single mutable collection. ``list``: A list of names of\n mutable collections. By default, all collections except \"intermediates\"\n are mutable.\n capture_intermediates: If ``True``, captures intermediate return values of\n all Modules inside the \"intermediates\" collection. By default only the\n return values of all ``__call__`` methods are stored. A function can be\n passed to change the filter behavior. The filter function takes the\n Module instance and method name and returns a bool indicating whether\n the output of that method invocation should be stored.\n **kwargs: Keyword arguments passed to the init function.\n\n Returns:\n ``(output, vars)``, where ``vars`` are is a dict of the modified\n collections.\n \"\"\"\n Module._module_checks(self)\n\n if not isinstance(rngs, dict):\n if not core.scope._is_valid_rng(rngs):\n raise errors.InvalidRngError(\n 'RNGs should be of shape (2,) or PRNGKey in Module '\n f'{self.__class__.__name__}, but rngs are: {rngs}'\n )\n rngs = {'params': rngs}\n\n if isinstance(method, str):\n attribute_name = method\n method = getattr(self, attribute_name)\n if not callable(method):\n class_name = type(self).__name__\n raise TypeError(\n f\"'{class_name}.{attribute_name}' must be a callable, got\"\n f' {type(method)}.'\n )\n elif method is None:\n method = self.__call__\n method = _get_unbound_fn(method)\n return init_with_output(\n method,\n self,\n mutable=mutable,\n capture_intermediates=capture_intermediates,\n )(rngs, *args, **kwargs)\n\n @traceback_util.api_boundary\n def init(\n self,\n rngs: PRNGKey | RNGSequences,\n *args,\n method: Callable[..., Any] | str | None = None,\n mutable: CollectionFilter = DenyList('intermediates'),\n capture_intermediates: bool | Callable[['Module', str], bool] = False,\n **kwargs,\n ) -> FrozenVariableDict | dict[str, Any]:\n \"\"\"Initializes a module method with variables and returns modified variables.\n\n ``init`` takes as first argument either a single ``PRNGKey``, or a\n dictionary mapping variable collections names to their ``PRNGKeys``, and\n will call ``method`` (which is the module's ``__call__`` function by\n default) passing ``*args`` and ``**kwargs``, and returns\n a dictionary of initialized variables.\n\n Example::\n\n >>> import flax.linen as nn\n >>> import jax, jax.numpy as jnp\n >>> import numpy as np\n\n >>> class Foo(nn.Module):\n ... @nn.compact\n ... def __call__(self, x, train):\n ... x = nn.Dense(16)(x)\n ... x = nn.BatchNorm(use_running_average=not train)(x)\n ... x = nn.relu(x)\n ... return nn.Dense(1)(x)\n\n >>> x = jnp.empty((1, 7))\n >>> module = Foo()\n >>> key = jax.random.key(0)\n >>> variables = module.init(key, x, train=False)\n\n If you pass a single ``PRNGKey``, Flax will use it to feed the ``'params'``\n RNG stream. If you want to use a different RNG stream or need to use\n multiple streams, you can pass a dictionary mapping each RNG stream name\n to its corresponding ``PRNGKey`` to ``init``. If ``self.make_rng(name)``\n is called on an RNG stream name that isn't passed by the user, it will\n default to using the ``'params'`` RNG stream.\n\n Example::\n\n >>> class Foo(nn.Module):\n ... @nn.compact\n ... def __call__(self, x):\n ... x = nn.Dense(16)(x)\n ... x = nn.relu(x)\n ...\n ... other_variable = self.variable(\n ... 'other_collection',\n ... 'other_variable',\n ... lambda x: jax.random.normal(self.make_rng('other_rng'), x.shape),\n ... x,\n ... )\n ... x = x + other_variable.value\n ...\n ... return nn.Dense(1)(x)\n\n >>> module = Foo()\n >>> rngs = {'params': jax.random.key(0), 'other_rng': jax.random.key(1)}\n >>> variables0 = module.init(rngs, x)\n\n >>> rngs['other_rng'] = jax.random.key(0)\n >>> variables1 = module.init(rngs, x)\n >>> # equivalent params (key(0))\n >>> _ = jax.tree_util.tree_map(\n ... np.testing.assert_allclose, variables0['params'], variables1['params']\n ... )\n >>> # different other_variable (key(1) vs key(0))\n >>> np.testing.assert_raises(\n ... AssertionError,\n ... np.testing.assert_allclose,\n ... variables0['other_collection']['other_variable'],\n ... variables1['other_collection']['other_variable'],\n ... )\n\n >>> del rngs['other_rng']\n >>> # self.make_rng('other_rng') will default to using the 'params' RNG stream\n >>> variables2 = module.init(rngs, x)\n >>> # equivalent params (key(0))\n >>> _ = jax.tree_util.tree_map(\n ... np.testing.assert_allclose, variables1['params'], variables2['params']\n ... )\n >>> # equivalent other_variable (key(0))\n >>> np.testing.assert_allclose(\n ... variables1['other_collection']['other_variable'],\n ... variables2['other_collection']['other_variable'],\n ... )\n\n >>> # passing in a single key is equivalent to passing in {'params': key}\n >>> variables3 = module.init(jax.random.key(0), x)\n >>> # equivalent params (key(0))\n >>> _ = jax.tree_util.tree_map(\n ... np.testing.assert_allclose, variables2['params'], variables3['params']\n ... )\n >>> # equivalent other_variable (key(0))\n >>> np.testing.assert_allclose(\n ... variables2['other_collection']['other_variable'],\n ... variables3['other_collection']['other_variable'],\n ... )\n\n Jitting ``init`` initializes a model lazily using only the shapes of the\n provided arguments, and avoids computing the forward pass with actual\n values. Example::\n\n >>> module = nn.Dense(1)\n >>> init_jit = jax.jit(module.init)\n >>> variables = init_jit(jax.random.key(0), x)\n\n ``init`` is a light wrapper over ``apply``, so other ``apply`` arguments\n like ``method``, ``mutable``, and ``capture_intermediates`` are also\n available.\n\n Args:\n rngs: The rngs for the variable collections.\n *args: Named arguments passed to the init function.\n method: An optional method. If provided, applies this method. If not\n provided, applies the ``__call__`` method. A string can also be provided\n to specify a method by name.\n mutable: Can be bool, str, or list. Specifies which collections should be\n treated as mutable: ``bool``: all/no collections are mutable. ``str``:\n The name of a single mutable collection. ``list``: A list of names of\n mutable collections. By default all collections except \"intermediates\"\n are mutable.\n capture_intermediates: If ``True``, captures intermediate return values of\n all Modules inside the \"intermediates\" collection. By default only the\n return values of all ``__call__`` methods are stored. A function can be\n passed to change the filter behavior. The filter function takes the\n Module instance and method name and returns a bool indicating whether\n the output of that method invocation should be stored.\n **kwargs: Keyword arguments passed to the init function.\n\n Returns:\n The initialized variable dict.\n \"\"\"\n Module._module_checks(self)\n\n _, v_out = self.init_with_output(\n rngs,\n *args,\n method=method,\n mutable=mutable,\n capture_intermediates=capture_intermediates,\n **kwargs,\n )\n return v_out\n\n @traceback_util.api_boundary\n def lazy_init(\n self,\n rngs: PRNGKey | RNGSequences,\n *args,\n method: Callable[..., Any] | None = None,\n mutable: CollectionFilter = DenyList('intermediates'),\n **kwargs,\n ) -> FrozenVariableDict:\n \"\"\"Initializes a module without computing on an actual input.\n\n lazy_init will initialize the variables without doing unnecessary compute.\n The input data should be passed as a ``jax.ShapeDtypeStruct`` which\n specifies the shape and dtype of the input but no concrete data.\n\n Example::\n\n >>> model = nn.Dense(features=256)\n >>> variables = model.lazy_init(\n ... jax.random.key(0), jax.ShapeDtypeStruct((1, 128), jnp.float32))\n\n The args and kwargs args passed to ``lazy_init`` can be a mix of\n concrete (jax arrays, scalars, bools) and abstract (ShapeDtypeStruct)\n values. Concrete values are only necessary for arguments that affect\n the initialization of variables. For example, the model might expect\n a keyword arg that enables/disables a subpart of the model.\n In this case, an explicit value (True/Flase) should be passed otherwise\n ``lazy_init`` cannot infer which variables should be initialized.\n\n Args:\n rngs: The rngs for the variable collections.\n *args: arguments passed to the init function.\n method: An optional method. If provided, applies this method. If not\n provided, applies the ``__call__`` method.\n mutable: Can be bool, str, or list. Specifies which collections should be\n treated as mutable: ``bool``: all/no collections are mutable. ``str``:\n The name of a single mutable collection. ``list``: A list of names of\n mutable collections. By default all collections except \"intermediates\"\n are mutable.\n **kwargs: Keyword arguments passed to the init function.\n\n Returns:\n The initialized variable dict.\n \"\"\"\n Module._module_checks(self)\n\n def lazy_wrapper(rngs, *args, **kwargs):\n return self.init(rngs, *args, method=method, mutable=mutable, **kwargs)\n\n return partial_eval.lazy_init(lazy_wrapper)(rngs, *args, **kwargs)\n\n @property\n def variables(self) -> VariableDict:\n \"\"\"Returns the variables in this module.\"\"\"\n if self.scope is None:\n raise ValueError(\"Can't access variables on unbound modules\")\n return self.scope.variables()\n\n def get_variable(self, col: str, name: str, default: T | None = None) -> T:\n \"\"\"Retrieves the value of a Variable.\n\n Args:\n col: the variable collection.\n name: the name of the variable.\n default: the default value to return if the variable does not exist in\n this scope.\n\n Returns:\n The value of the input variable, of the default value if the variable\n doesn't exist in this scope.\n \"\"\"\n if self.scope is None:\n raise ValueError(\"Can't access variables on unbound modules\")\n return self.scope.get_variable(col, name, default)\n\n def put_variable(self, col: str, name: str, value: Any):\n \"\"\"Updates the value of the given variable if it is mutable, or an error otherwise.\n\n Args:\n col: the variable collection.\n name: the name of the variable.\n value: the new value of the variable.\n \"\"\"\n if self.scope is None:\n raise ValueError(\"Can't access variables on unbound modules\")\n self.scope.put_variable(col, name, value)\n\n @overload\n def sow(self, col: str, name: str, value: Any) -> bool:\n ...\n\n @overload\n def sow(\n self,\n col: str,\n name: str,\n value: T,\n reduce_fn: Callable[[K, T], K] = tuple_reduce,\n init_fn: Callable[[], K] = tuple_init, # type: ignore\n ) -> bool:\n ...\n\n def sow(\n self,\n col: str,\n name: str,\n value: T,\n reduce_fn: Callable[[K, T], K] = tuple_reduce,\n init_fn: Callable[[], K] = tuple_init, # type: ignore\n ) -> bool:\n \"\"\"Stores a value in a collection.\n\n Collections can be used to collect intermediate values without\n the overhead of explicitly passing a container through each Module call.\n\n If the target collection is not mutable ``sow`` behaves like a no-op\n and returns ``False``.\n\n Example::\n\n >>> import jax\n >>> import jax.numpy as jnp\n >>> import flax.linen as nn\n\n >>> class Foo(nn.Module):\n ... @nn.compact\n ... def __call__(self, x):\n ... h = nn.Dense(4)(x)\n ... self.sow('intermediates', 'h', h)\n ... return nn.Dense(2)(h)\n\n >>> x = jnp.ones((16, 9))\n >>> model = Foo()\n >>> variables = model.init(jax.random.key(0), x)\n >>> y, state = model.apply(variables, x, mutable=['intermediates'])\n >>> jax.tree.map(jnp.shape, state['intermediates'])\n {'h': ((16, 4),)}\n\n By default the values are stored in a tuple and each stored value\n is appended at the end. This way all intermediates can be tracked when\n the same module is called multiple times. Alternatively, a custom\n init/reduce function can be passed::\n\n >>> class Foo2(nn.Module):\n ... @nn.compact\n ... def __call__(self, x):\n ... init_fn = lambda: 0\n ... reduce_fn = lambda a, b: a + b\n ... self.sow('intermediates', 'h', x,\n ... init_fn=init_fn, reduce_fn=reduce_fn)\n ... self.sow('intermediates', 'h', x * 2,\n ... init_fn=init_fn, reduce_fn=reduce_fn)\n ... return x\n\n >>> x = jnp.ones((1, 1))\n >>> model = Foo2()\n >>> variables = model.init(jax.random.key(0), x)\n >>> y, state = model.apply(\n ... variables, x, mutable=['intermediates'])\n >>> print(state['intermediates'])\n {'h': Array([[3.]], dtype=float32)}\n\n Args:\n col: The name of the variable collection.\n name: The name of the variable.\n value: The value of the variable.\n reduce_fn: The function used to combine the existing value with the new\n value. The default is to append the value to a tuple.\n init_fn: For the first value stored, ``reduce_fn`` will be passed the result\n of ``init_fn`` together with the value to be stored. The default is an\n empty tuple.\n\n Returns:\n ``True`` if the value has been stored successfully, ``False`` otherwise.\n \"\"\"\n if self.scope is None:\n raise ValueError(\"Can't store variables on unbound modules\")\n if not self.scope.is_mutable_collection(col):\n return False\n if self.scope.has_variable(col, name):\n xs = self.scope.get_variable(col, name)\n else:\n self.scope.reserve(name, col)\n self._state.children[name] = col\n xs = init_fn()\n xs = reduce_fn(xs, value)\n self.scope.put_variable(col, name, xs)\n return True\n\n def perturb(\n self, name: str, value: T, collection: str = 'perturbations'\n ) -> T:\n \"\"\"Add an zero-value variable ('perturbation') to the intermediate value.\n\n The gradient of ``value`` would be the same as the gradient of this\n perturbation variable. Therefore, if you define your loss function with\n both params and perturbations as standalone arguments, you can get the\n intermediate gradients of ``value`` by running ``jax.grad`` on the perturbation\n argument.\n\n .. note::\n This is an experimental API and may be tweaked later for better\n performance and usability.\n At its current stage, it creates extra dummy variables that occupies extra\n memory space. Use it only to debug gradients in training.\n\n Example::\n\n >>> class Foo(nn.Module):\n ... @nn.compact\n ... def __call__(self, x):\n ... x = nn.Dense(3)(x)\n ... x = self.perturb('dense3', x)\n ... return nn.Dense(2)(x)\n\n >>> def loss(variables, inputs, targets):\n ... preds = model.apply(variables, inputs)\n ... return jnp.square(preds - targets).mean()\n\n >>> x = jnp.ones((2, 9))\n >>> y = jnp.ones((2, 2))\n >>> model = Foo()\n >>> variables = model.init(jax.random.key(0), x)\n >>> intm_grads = jax.grad(loss, argnums=0)(variables, x, y)\n >>> print(intm_grads['perturbations']['dense3'])\n [[-0.04684732 0.06573904 -0.3194327 ]\n [-0.04684732 0.06573904 -0.3194327 ]]\n\n If perturbations are not passed to ``apply``, ``perturb`` behaves like a no-op\n so you can easily disable the behavior when not needed::\n\n >>> model.apply(variables, x) # works as expected\n Array([[-0.04579116, 0.50412744],\n [-0.04579116, 0.50412744]], dtype=float32)\n >>> model.apply({'params': variables['params']}, x) # behaves like a no-op\n Array([[-0.04579116, 0.50412744],\n [-0.04579116, 0.50412744]], dtype=float32)\n >>> intm_grads = jax.grad(loss, argnums=0)({'params': variables['params']}, x, y)\n >>> 'perturbations' not in intm_grads\n True\n \"\"\"\n if self.scope is None:\n raise ValueError(\"Can't store variables on unbound modules\")\n\n if self.is_mutable_collection(collection):\n if not self.scope.has_variable(collection, name):\n self.scope.reserve(name, collection)\n self._state.children[name] = collection\n zeros = jax.tree.map(jnp.zeros_like, value)\n self.scope.put_variable(collection, name, zeros) # type: ignore\n\n if collection in self.scope.root._variables:\n if self.scope.has_variable(collection, name):\n old_value = self.scope.get_variable(collection, name)\n value = jax.tree.map(jnp.add, value, old_value) # type: ignore\n else:\n raise ValueError(f\"Perturbation collection {collection} present, but \"\n f\"missing perturbation variable {name}\")\n\n return value\n\n def tabulate(\n self,\n rngs: PRNGKey | RNGSequences,\n *args,\n depth: int | None = None,\n show_repeated: bool = False,\n mutable: CollectionFilter = DenyList('intermediates'),\n console_kwargs: Mapping[str, Any] | None = None,\n table_kwargs: Mapping[str, Any] = MappingProxyType({}),\n column_kwargs: Mapping[str, Any] = MappingProxyType({}),\n compute_flops: bool = False,\n compute_vjp_flops: bool = False,\n **kwargs,\n ) -> str:\n \"\"\"Creates a summary of the Module represented as a table.\n\n This method has the same signature and internally calls ``Module.init``,\n but instead of returning the variables, it returns the string summarizing\n the Module in a table. ``tabulate`` uses ``jax.eval_shape`` to run the forward\n computation without consuming any FLOPs or allocating memory.\n\n Additional arguments can be passed into the ``console_kwargs`` argument, for\n example, ``{'width': 120}``. For a full list of ``console_kwargs`` arguments,\n see:\n https://rich.readthedocs.io/en/stable/reference/console.html#rich.console.Console\n\n Example::\n\n >>> import flax.linen as nn\n >>> import jax, jax.numpy as jnp\n\n >>> class Foo(nn.Module):\n ... @nn.compact\n ... def __call__(self, x):\n ... h = nn.Dense(4)(x)\n ... return nn.Dense(2)(h)\n\n >>> x = jnp.ones((16, 9))\n\n >>> # print(Foo().tabulate(\n >>> # jax.random.key(0), x, compute_flops=True, compute_vjp_flops=True))\n\n This gives the following output::\n\n Foo Summary\n ┏━━━━━━━━━┳━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━┳━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┓\n ┃ path ┃ module ┃ inputs ┃ outputs ┃ flops ┃ vjp_flops ┃ params ┃\n ┡━━━━━━━━━╇━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━╇━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━┩\n │ │ Foo │ float32[16,9] │ float32[16,2] │ 1504 │ 4460 │ │\n ├─────────┼────────┼───────────────┼───────────────┼───────┼───────────┼─────────────────┤\n │ Dense_0 │ Dense │ float32[16,9] │ float32[16,4] │ 1216 │ 3620 │ bias: │\n │ │ │ │ │ │ │ float32[4] │\n │ │ │ │ │ │ │ kernel: │\n │ │ │ │ │ │ │ float32[9,4] │\n │ │ │ │ │ │ │ │\n │ │ │ │ │ │ │ 40 (160 B) │\n ├─────────┼────────┼───────────────┼───────────────┼───────┼───────────┼─────────────────┤\n │ Dense_1 │ Dense │ float32[16,4] │ float32[16,2] │ 288 │ 840 │ bias: │\n │ │ │ │ │ │ │ float32[2] │\n │ │ │ │ │ │ │ kernel: │\n │ │ │ │ │ │ │ float32[4,2] │\n │ │ │ │ │ │ │ │\n │ │ │ │ │ │ │ 10 (40 B) │\n ├─────────┼────────┼───────────────┼───────────────┼───────┼───────────┼─────────────────┤\n │ │ │ │ │ │ Total │ 50 (200 B) │\n └─────────┴────────┴───────────────┴───────────────┴───────┴───────────┴─────────────────┘\n\n Total Parameters: 50 (200 B)\n\n **Note**: rows order in the table does not represent execution order,\n instead it aligns with the order of keys in ``variables`` which are sorted\n alphabetically.\n\n **Note**: ``vjp_flops`` returns ``0`` if the module is not differentiable.\n\n Args:\n rngs: The rngs for the variable collections as passed to ``Module.init``.\n *args: The arguments to the forward computation.\n depth: controls how many submodule deep the summary can go. By default,\n its ``None`` which means no limit. If a submodule is not shown because of\n the depth limit, its parameter count and bytes will be added to the row\n of its first shown ancestor such that the sum of all rows always adds\n up to the total number of parameters of the Module.\n show_repeated: If ``True``, repeated calls to the same module will be shown\n in the table, otherwise only the first call will be shown. Default is\n ``False``.\n mutable: Can be bool, str, or list. Specifies which collections should be\n treated as mutable: ``bool``: all/no collections are mutable. ``str``:\n The name of a single mutable collection. ``list``: A list of names of\n mutable collections. By default, all collections except 'intermediates'\n are mutable.\n console_kwargs: An optional dictionary with additional keyword arguments\n that are passed to ``rich.console.Console`` when rendering the table.\n Default arguments are ``{'force_terminal': True, 'force_jupyter':\n False}``.\n table_kwargs: An optional dictionary with additional keyword arguments\n that are passed to ``rich.table.Table`` constructor.\n column_kwargs: An optional dictionary with additional keyword arguments\n that are passed to ``rich.table.Table.add_column`` when adding columns to\n the table.\n compute_flops: whether to include a ``flops`` column in the table listing\n the estimated FLOPs cost of each module forward pass. Does incur actual\n on-device computation / compilation / memory allocation, but still\n introduces overhead for large modules (e.g. extra 20 seconds for a\n Stable Diffusion's UNet, whereas otherwise tabulation would finish in 5\n seconds).\n compute_vjp_flops: whether to include a ``vjp_flops`` column in the table\n listing the estimated FLOPs cost of each module backward pass.\n Introduces a compute overhead of about 2-3X of ``compute_flops``.\n **kwargs: keyword arguments to pass to the forward computation.\n\n Returns:\n A string summarizing the Module.\n \"\"\"\n from flax.linen import summary\n\n tabulate_fn = summary.tabulate(\n self,\n rngs,\n depth=depth,\n show_repeated=show_repeated,\n mutable=mutable,\n console_kwargs=console_kwargs,\n table_kwargs=table_kwargs,\n column_kwargs=column_kwargs,\n compute_flops=compute_flops,\n compute_vjp_flops=compute_vjp_flops,\n )\n return tabulate_fn(*args, **kwargs)\n\n def module_paths(\n self,\n rngs: PRNGKey | RNGSequences,\n *args,\n show_repeated: bool = False,\n mutable: CollectionFilter = DenyList('intermediates'),\n **kwargs,\n ) -> dict[str, 'Module']:\n \"\"\"Returns a dictionary mapping module paths to module instances.\n\n This method has the same signature and internally calls ``Module.init``,\n but instead of returning the variables, it returns a dictionary mapping\n module paths to unbounded copies of module instances that were used\n at runtime. ``module_paths`` uses ``jax.eval_shape`` to run the forward\n computation without consuming any FLOPs or allocating memory.\n\n Example::\n\n >>> import flax.linen as nn\n >>> import jax, jax.numpy as jnp\n\n >>> class Foo(nn.Module):\n ... @nn.compact\n ... def __call__(self, x):\n ... h = nn.Dense(4)(x)\n ... return nn.Dense(2)(h)\n\n >>> x = jnp.ones((16, 9))\n >>> modules = Foo().module_paths(jax.random.key(0), x)\n >>> print({\n ... p: type(m).__name__ for p, m in modules.items()\n ... })\n {'': 'Foo', 'Dense_0': 'Dense', 'Dense_1': 'Dense'}\n\n Args:\n rngs: The rngs for the variable collections as passed to ``Module.init``.\n *args: The arguments to the forward computation.\n show_repeated: If ``True``, repeated calls to the same module will be\n shown in the table, otherwise only the first call will be shown.\n Default is ``False``.\n mutable: Can be bool, str, or list. Specifies which collections should\n be treated as mutable: ``bool``: all/no collections are mutable.\n ``str``: The name of a single mutable collection. ``list``: A list of\n names of mutable collections. By default, all collections except\n 'intermediates' are mutable.\n **kwargs: keyword arguments to pass to the forward computation.\n\n Returns:\n A dict`ionary mapping module paths to module instances.\n \"\"\"\n from flax.linen import summary\n\n table = summary._get_module_table(\n module=self,\n depth=None,\n show_repeated=show_repeated,\n compute_flops=False,\n compute_vjp_flops=False,\n )(rngs, *args, **kwargs, mutable=mutable)\n\n return {'/'.join(row.path): row.module_copy for row in table}\n\n\n_ParentType = Union[Module, Scope, _Sentinel, None]\n\n\ndef merge_param(name: str, a: T | None, b: T | None) -> T:\n \"\"\"Merges construction- and call-time argument.\n\n This is a utility for supporting a pattern where a Module hyperparameter\n can be passed either to ``__init__`` or ``__call__``, and the value that is\n not ``None`` will be used.\n\n Example::\n\n >>> import flax.linen as nn\n >>> from typing import Optional\n\n >>> class Foo(nn.Module):\n ... train: Optional[bool] = None\n\n ... def __call__(self, train: Optional[bool] = None):\n ... train = nn.merge_param('train', self.train, train)\n\n An error is thrown when both arguments are ``None`` or both values are not\n ``None``.\n\n Args:\n name: the name of the parameter. Used for error messages.\n a: option a\n b: option b\n\n Returns:\n a or b whichever is not ``None``.\n \"\"\"\n if a is None and b is None:\n raise ValueError(\n f'Parameter \"{name}\" must be passed to the constructor or at call time.'\n )\n if a is not None and b is not None:\n raise ValueError(\n f'Parameter \"{name}\" was passed to the constructor and at call time.'\n ' Should be passed just once.'\n )\n if a is None:\n assert b is not None\n return b\n return a\n\n\n@traceback_util.api_boundary\ndef apply(\n fn: Callable[..., Any],\n module: Module,\n mutable: CollectionFilter = False,\n capture_intermediates: bool | Callable[[Module, str], bool] = False,\n) -> Callable[..., Any]:\n \"\"\"Creates an apply function to call ``fn`` with a bound module.\n\n Unlike ``Module.apply`` this function returns a new function with the\n signature ``(variables, *args, rngs=None, **kwargs) -> T`` where ``T`` is the\n return type of ``fn``. If ``mutable`` is not ``False`` the return type is a\n tuple where the second item is a ``FrozenDict`` with the mutated variables.\n\n The apply function that is returned can be directly composed with\n JAX transformations like ``jax.jit``::\n\n >>> class Foo(nn.Module):\n ... def encode(self, x):\n ... ...\n ... def decode(self, x):\n ... ...\n\n >>> def f(foo, x):\n ... z = foo.encode(x)\n ... y = foo.decode(z)\n ... # ...\n ... return y\n\n >>> variables = {}\n >>> foo = Foo()\n >>> f_jitted = jax.jit(nn.apply(f, foo))\n >>> f_jitted(variables, jnp.ones((1, 3)))\n\n Args:\n fn: The function that should be applied. The first argument passed will be\n a module instance of the ``module`` with variables and RNGs bound to it.\n module: The ``Module`` that will be used to bind variables and RNGs to. The\n ``Module`` passed as the first argument to ``fn`` will be a clone of\n module.\n mutable: Can be bool, str, or list. Specifies which collections should be\n treated as mutable: ``bool``: all/no collections are mutable. ``str``: The\n name of a single mutable collection. ``list``: A list of names of mutable\n collections.\n capture_intermediates: If ``True``, captures intermediate return values of all\n Modules inside the \"intermediates\" collection. By default, only the return\n values of all `__call__` methods are stored. A function can be passed to\n change the filter behavior. The filter function takes the Module instance\n and method name and returns a bool indicating whether the output of that\n method invocation should be stored.\n\n Returns:\n The apply function wrapping ``fn``.\n \"\"\"\n\n @functools.wraps(fn)\n def scope_fn(scope, *args, **kwargs):\n _context.capture_stack.append(capture_intermediates)\n try:\n return fn(module.clone(parent=scope, _deep_clone=True), *args, **kwargs)\n finally:\n _context.capture_stack.pop()\n\n if capture_intermediates is True: # pylint: disable=g-bool-id-comparison\n capture_intermediates = capture_call_intermediates\n if capture_intermediates:\n mutable = union_filters(mutable, 'intermediates')\n return core.apply(scope_fn, mutable=mutable)\n\n\n@traceback_util.api_boundary\ndef init_with_output(\n fn: Callable[..., Any],\n module: Module,\n mutable: CollectionFilter = DenyList('intermediates'),\n capture_intermediates: bool | Callable[[Module, str], bool] = False,\n) -> Callable[..., tuple[Any, FrozenVariableDict | dict[str, Any]]]:\n \"\"\"Creates an init function to call ``fn`` with a bound module that also returns the function outputs.\n\n Unlike ``Module.init_with_output`` this function returns a new function with\n the signature ``(rngs, *args, **kwargs) -> (T, variables)`` where ``T`` is the\n return type of ``fn``. The rngs can be a dict of PRNGKeys or a single\n ```PRNGKey`` which is equivalent to passing a dict with one PRNGKey with the\n name \"params\".\n\n The init function that is returned can be directly composed with\n JAX transformations like ``jax.jit``::\n\n >>> class Foo(nn.Module):\n ... def encode(self, x):\n ... ...\n ... def decode(self, x):\n ... ...\n\n >>> def f(foo, x):\n ... z = foo.encode(x)\n ... y = foo.decode(z)\n ... # ...\n ... return y\n\n >>> foo = Foo()\n >>> f_jitted = jax.jit(nn.init_with_output(f, foo))\n >>> y, variables = f_jitted(jax.random.key(0), jnp.ones((1, 3)))\n\n Args:\n fn: The function that should be applied. The first argument passed will be\n a module instance of the ``module`` with variables and RNGs bound to it.\n module: The ``Module`` that will be used to bind variables and RNGs to. The\n ``Module`` passed as the first argument to ``fn`` will be a clone of\n module.\n mutable: Can be bool, str, or list. Specifies which collections should be\n treated as mutable: ``bool``: all/no collections are mutable. ``str``: The\n name of a single mutable collection. ``list``: A list of names of mutable\n collections. By default, all collections except \"intermediates\" are\n mutable.\n capture_intermediates: If ``True``, captures intermediate return values of all\n Modules inside the \"intermediates\" collection. By default, only the return\n values of all `__call__` methods are stored. A function can be passed to\n change the filter behavior. The filter function takes the Module instance\n and method name and returns a bool indicating whether the output of that\n method invocation should be stored.\n\n Returns:\n The init function wrapping ``fn``.\n \"\"\"\n\n @functools.wraps(fn)\n def scope_fn(scope, *args, **kwargs):\n _context.capture_stack.append(capture_intermediates)\n try:\n return fn(module.clone(parent=scope, _deep_clone=True), *args, **kwargs)\n finally:\n _context.capture_stack.pop()\n\n if capture_intermediates is True: # pylint: disable=g-bool-id-comparison\n capture_intermediates = capture_call_intermediates\n if capture_intermediates:\n mutable = union_filters(mutable, 'intermediates')\n return core.init(scope_fn, mutable=mutable)\n\n\n@traceback_util.api_boundary\ndef init(\n fn: Callable[..., Any],\n module: Module,\n mutable: CollectionFilter = DenyList('intermediates'),\n capture_intermediates: bool | Callable[[Module, str], bool] = False,\n) -> Callable[..., FrozenVariableDict | dict[str, Any]]:\n \"\"\"Creates an init function to call ``fn`` with a bound module.\n\n Unlike ``Module.init`` this function returns a new function with the signature\n ``(rngs, *args, **kwargs) -> variables``.\n The rngs can be a dict of PRNGKeys or a single ```PRNGKey`` which is\n equivalent to passing a dict with one PRNGKey with the name \"params\".\n\n The init function that is returned can be directly composed with\n JAX transformations like ``jax.jit``::\n\n >>> class Foo(nn.Module):\n ... def encode(self, x):\n ... ...\n ... def decode(self, x):\n ... ...\n\n >>> def f(foo, x):\n ... z = foo.encode(x)\n ... y = foo.decode(z)\n ... # ...\n ... return y\n\n >>> foo = Foo()\n >>> f_jitted = jax.jit(nn.init(f, foo))\n >>> variables = f_jitted(jax.random.key(0), jnp.ones((1, 3)))\n\n Args:\n fn: The function that should be applied. The first argument passed will be\n a module instance of the ``module`` with variables and RNGs bound to it.\n module: The ``Module`` that will be used to bind variables and RNGs to. The\n ``Module`` passed as the first argument to ``fn`` will be a clone of\n module.\n mutable: Can be bool, str, or list. Specifies which collections should be\n treated as mutable: ``bool``: all/no collections are mutable. ``str``: The\n name of a single mutable collection. ``list``: A list of names of mutable\n collections. By default, all collections except \"intermediates\" are\n mutable.\n capture_intermediates: If `True`, captures intermediate return values of all\n Modules inside the \"intermediates\" collection. By default, only the return\n values of all `__call__` methods are stored. A function can be passed to\n change the filter behavior. The filter function takes the Module instance\n and method name and returns a bool indicating whether the output of that\n method invocation should be stored.\n\n Returns:\n The init function wrapping ``fn``.\n \"\"\"\n init_fn = init_with_output(fn, module, mutable, capture_intermediates)\n\n @functools.wraps(init_fn)\n def init_wrapper(*args, **kwargs):\n return init_fn(*args, **kwargs)[1]\n\n return init_wrapper\n\n\n# TODO(cgarciae): we are defining CompactNameScope just to\n# avoid a pytype bug with the Flax overlay. We should aim to\n# remove in the at some point as its not ergonomic.\nif not typing.TYPE_CHECKING:\n\n class CompactNameScope(Module):\n fn: Callable\n module_fn: Callable[[], Module]\n\n @compact\n def __call__(self, *args, **kwargs) -> Any:\n return self.fn(self.module_fn(), *args, **kwargs)\nelse:\n\n @dataclasses.dataclass\n class CompactNameScope:\n fn: Callable\n module_fn: Callable\n name: str\n\n def __call__(self, *args, **kwargs) -> Any:\n ...\n\n\ndef share_scope(module: Module, other: Module, /):\n \"\"\"Modifies one of the Modules such that they share the same scope. This is useful\n when you want to wrap a Module and extend its functionality without changing the\n parameter structure.\n\n ``share_scope`` takes two Modules, ``module`` and ``other``. ``module`` will use\n ``other``'s scope if ``other`` has a scope and its not a descendant of``module``'s\n scope::\n\n >>> import flax.linen as nn\n >>> import jax\n >>> from jax import numpy as jnp, random\n ...\n >>> class DenseLoRA(nn.Module):\n ... base: nn.Dense\n ... rank: int\n ...\n ... def setup(self):\n ... nn.share_scope(self, self.base)\n ...\n ... @nn.compact\n ... def __call__(self, x: jax.Array):\n ... din, dout = x.shape[-1], self.base.features\n ... A = self.param('A', nn.zeros_init(), (din, self.rank))\n ... B = self.param('B', nn.zeros_init(), (self.rank, dout))\n ... return self.base(x) + x @ A @ B\n ...\n >>> class Model(nn.Module):\n ... @nn.compact\n ... def __call__(self, x: jax.Array):\n ... dense = nn.Dense(10) # base scope\n ... return DenseLoRA(dense, rank=2)(x) # reuse the base scope\n ...\n >>> model = Model()\n ...\n >>> params = model.init(random.key(0), jnp.ones((1, 5)))['params']\n >>> list(params['Dense_0'].keys())\n ['A', 'B', 'kernel', 'bias']\n\n When ``other``'s scope is a descendant of ``module``'s scope then ``other``\n will use ``module``'s scope instead::\n\n >>> class DenseLoRA(nn.Module):\n ... features: int\n ... rank: int\n ...\n ... def setup(self):\n ... self.child = nn.Dense(self.features)\n ... nn.share_scope(self, self.child)\n ...\n ... @nn.compact\n ... def __call__(self, x: jax.Array):\n ... din, dout = x.shape[-1], self.features\n ... A = self.param('A', nn.zeros_init(), (din, self.rank))\n ... B = self.param('B', nn.zeros_init(), (self.rank, dout))\n ... return self.child(x) + x @ A @ B\n ...\n >>> class Model(nn.Module):\n ... @nn.compact\n ... def __call__(self, x: jax.Array):\n ... return DenseLoRA(10, rank=2)(x)\n ...\n >>> model = Model()\n ...\n >>> params = model.init(random.key(0), jnp.ones((1, 5)))['params']\n >>> list(params['DenseLoRA_0'].keys())\n ['A', 'B', 'kernel', 'bias']\n \"\"\"\n if module.scope is None or other.scope is None:\n raise errors.CallShareScopeOnUnboundModuleError()\n\n def _is_child_scope(scope: Scope, other: Scope) -> bool:\n target: Scope | None = other\n\n while target is not None:\n if target is scope:\n return True\n target = target.parent\n return False\n\n if _is_child_scope(module.scope, other.scope):\n # Child is a true child, overwrite its scope\n module_to_update = other\n new_scope = module.scope\n else:\n # Child has its own independent scope, overwrite\n # parent scope, so that we preserve the sharing\n module_to_update = module\n new_scope = other.scope\n\n old_scope = module_to_update.scope\n object.__setattr__(module_to_update, 'scope', new_scope)\n\n # Reattach all the children to the new scope as well.\n for m in module_to_update._state.children.values():\n if not isinstance(m, Module):\n continue\n # Should we go recursively to check if any of the ancestors point to the old\n # scope?\n if m.scope and m.scope.parent == old_scope:\n # Reserve the scope, so that if there is a conflict we can raise an error.\n if isinstance(m.scope.name, str):\n new_scope.reserve(m.scope.name)\n m.scope.parent = new_scope\n" }, { "path": "flax/linen/normalization.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Normalization modules for Flax.\"\"\"\n\nimport dataclasses\nimport functools\nfrom typing import Any\nfrom collections.abc import Iterable\n\nimport jax\nimport jax.numpy as jnp\nfrom jax import lax\nfrom jax.nn import initializers\n\nfrom flax.linen import dtypes, module, transforms\nfrom flax.typing import (\n Array,\n PRNGKey as PRNGKey,\n Dtype,\n Shape as Shape,\n Initializer,\n Axes,\n)\n\nfield = dataclasses.field\ncanonicalize_dtype = dtypes.canonicalize_dtype\ncompact = module.compact\nModule = module.Module\nmerge_param = module.merge_param\nmap_variables = transforms.map_variables\n\n\ndef _canonicalize_axes(rank: int, axes: Axes) -> tuple[int, ...]:\n \"\"\"Returns a tuple of deduplicated, sorted, and positive axes.\"\"\"\n if not isinstance(axes, Iterable):\n axes = (axes,)\n return tuple({rank + axis if axis < 0 else axis for axis in axes})\n\n\ndef _abs_sq(x):\n \"\"\"Computes the elementwise square of the absolute value |x|^2.\"\"\"\n if jnp.iscomplexobj(x):\n return lax.square(lax.real(x)) + lax.square(lax.imag(x))\n else:\n return lax.square(x)\n\n\ndef _compute_stats(\n x: Array,\n axes: Axes,\n dtype: Dtype | None,\n axis_name: str | None = None,\n axis_index_groups: Any = None,\n use_mean: bool = True,\n use_fast_variance: bool = True,\n mask: Array | None = None,\n force_float32_reductions=True,\n):\n \"\"\"Computes mean and variance statistics.\n\n This implementation takes care of a few important details:\n - By default, computes in float32 precision for stability\n in half precision training.\n - If `use_fast_variance` is `True`, mean and variance are computed using\n Var = E[|x|^2] - |E[x]|^2, instead of Var = E[|x - E[x]|^2]), in a single\n XLA fusion.\n - Clips negative variances to zero which can happen due to\n roundoff errors. This avoids downstream NaNs.\n - Supports averaging across a parallel axis and subgroups of a parallel axis\n with a single `lax.pmean` call to avoid latency.\n\n Arguments:\n x: Input array.\n axes: The axes in ``x`` to compute mean and variance statistics for.\n dtype: Optional dtype specifying the minimal precision. Statistics are\n always at least float32 for stability (default: dtype of x).\n axis_name: Optional name for the pmapped axis to compute mean over. Note,\n this is only used for pmap and shard map. For SPMD jit, you do not need to\n manually synchronize. Just make sure that the axes are correctly annotated\n and XLA:SPMD will insert the necessary collectives.\n axis_index_groups: Optional groups of indices within that named axis.\n use_mean: If true, calculate the mean from the input and use it when\n computing the variance. If false, set the mean to zero and compute the\n variance without subtracting the mean.\n use_fast_variance: If true, use a faster, but less numerically stable,\n calculation for the variance.\n mask: Binary array of shape broadcastable to `inputs` tensor, indicating the\n positions for which the mean and variance should be computed.\n force_float32_reductions: If false, this will skip float32 promotion and use\n the input dtype or inherited dtype from ``x``.\n\n Returns:\n A pair ``(mean, var)``.\n \"\"\"\n if dtype is None:\n dtype = jnp.result_type(x)\n # promote x to at least float32, this avoids half precision computation\n # but preserves double or complex floating points\n if force_float32_reductions:\n dtype = jnp.promote_types(dtype, jnp.float32)\n if isinstance(x, jax.Array):\n x = x.astype(dtype)\n else:\n x = jnp.asarray(x, dtype)\n axes = _canonicalize_axes(x.ndim, axes)\n\n def maybe_distributed_mean(*xs, mask=None):\n if mask is not None:\n mask = jnp.asarray(mask, dtype=bool)\n mus = tuple(x.mean(axes, where=mask) for x in xs)\n if axis_name is None:\n return mus if len(xs) > 1 else mus[0]\n else:\n # In the distributed case we stack multiple arrays to speed comms.\n if len(xs) > 1:\n reduced_mus = lax.pmean(\n jnp.stack(mus, axis=0),\n axis_name,\n axis_index_groups=axis_index_groups,\n )\n return tuple(reduced_mus[i] for i in range(len(xs)))\n else:\n return lax.pmean(mus[0], axis_name, axis_index_groups=axis_index_groups)\n\n if use_mean:\n if use_fast_variance:\n mu, mu2 = maybe_distributed_mean(x, _abs_sq(x), mask=mask)\n # mean2 - _abs_sq(mean) is not guaranteed to be non-negative due\n # to floating point round-off errors.\n var = jnp.maximum(0.0, mu2 - _abs_sq(mu))\n else:\n mu = maybe_distributed_mean(x, mask=mask)\n var = maybe_distributed_mean(\n _abs_sq(x - jnp.expand_dims(mu, axes)), mask=mask\n )\n else:\n var = maybe_distributed_mean(_abs_sq(x), mask=mask)\n mu = jnp.zeros_like(var)\n return mu, var\n\n\ndef _normalize(\n mdl: Module,\n x: Array,\n mean: Array,\n var: Array,\n reduction_axes: Axes,\n feature_axes: Axes,\n dtype: Dtype | None,\n param_dtype: Dtype,\n epsilon: float,\n use_bias: bool,\n use_scale: bool,\n bias_init: Initializer,\n scale_init: Initializer,\n force_float32_reductions: bool = True\n):\n \"\"\"Normalizes the input of a normalization layer and optionally applies a learned scale and bias.\n\n Arguments:\n mdl: Module to apply the normalization in (normalization params will reside\n in this module).\n x: The input.\n mean: Mean to use for normalization.\n var: Variance to use for normalization.\n reduction_axes: The axes in ``x`` to reduce.\n feature_axes: Axes containing features. A separate bias and scale is learned\n for each specified feature.\n dtype: The dtype of the result (default: infer from input and params).\n param_dtype: The dtype of the parameters.\n epsilon: Normalization epsilon.\n use_bias: If true, add a bias term to the output.\n use_scale: If true, scale the output.\n bias_init: Initialization function for the bias term.\n scale_init: Initialization function for the scaling function.\n force_float32_reductions: If false, the scale and bias parameters use the\n param_dtype. Otherwise, they will have at least float32 precision due to\n the mean and var being promoted to float32.\n\n Returns:\n The normalized input.\n \"\"\"\n reduction_axes = _canonicalize_axes(x.ndim, reduction_axes)\n feature_axes = _canonicalize_axes(x.ndim, feature_axes)\n feature_shape = [1] * x.ndim\n reduced_feature_shape = []\n for ax in feature_axes:\n feature_shape[ax] = x.shape[ax]\n reduced_feature_shape.append(x.shape[ax])\n\n mean = jnp.expand_dims(mean, reduction_axes)\n var = jnp.expand_dims(var, reduction_axes)\n y = x - mean\n mul = lax.rsqrt(var + epsilon)\n args = [x]\n if use_scale:\n scale = mdl.param(\n 'scale', scale_init, reduced_feature_shape, param_dtype\n ).reshape(feature_shape)\n if not force_float32_reductions:\n scale = jnp.asarray(scale, param_dtype)\n mul *= scale\n args.append(scale)\n y *= mul\n if use_bias:\n bias = mdl.param(\n 'bias', bias_init, reduced_feature_shape, param_dtype\n ).reshape(feature_shape)\n if not force_float32_reductions:\n bias = jnp.asarray(bias, param_dtype)\n y += bias\n args.append(bias)\n dtype = dtypes.canonicalize_dtype(*args, dtype=dtype)\n return jnp.asarray(y, dtype)\n\n\ndef _l2_normalize(x, axis=None, eps=1e-12):\n \"\"\"Normalizes along dimension `axis` using an L2 norm.\n\n This specialized function exists for numerical stability reasons.\n\n Args:\n x: An input ndarray.\n axis: Dimension along which to normalize, e.g. `1` to separately normalize\n vectors in a batch. Passing `None` views `t` as a flattened vector when\n calculating the norm (equivalent to Frobenius norm).\n eps: Epsilon to avoid dividing by zero.\n\n Returns:\n An array of the same shape as 'x' L2-normalized along 'axis'.\n \"\"\"\n return x * jax.lax.rsqrt((x * x).sum(axis=axis, keepdims=True) + eps)\n\n\nclass BatchNorm(Module):\n \"\"\"BatchNorm Module.\n\n Usage Note:\n If we define a model with BatchNorm, for example::\n\n >>> import flax.linen as nn\n >>> import jax, jax.numpy as jnp\n >>> BN = nn.BatchNorm(momentum=0.9, epsilon=1e-5, dtype=jnp.float32)\n\n The initialized variables dict will contain, in addition to a 'params'\n collection, a separate 'batch_stats' collection that will contain all the\n running statistics for all the BatchNorm layers in a model::\n\n >>> x = jax.random.normal(jax.random.key(0), (5, 6))\n >>> variables = BN.init(jax.random.key(1), x, use_running_average=False)\n >>> jax.tree_util.tree_map(jnp.shape, variables)\n {'batch_stats': {'mean': (6,), 'var': (6,)}, 'params': {'bias': (6,), 'scale': (6,)}}\n\n We then update the batch_stats during training by specifying that the\n ``batch_stats`` collection is mutable in the ``apply`` method for our\n module.::\n\n >>> y, new_batch_stats = BN.apply(variables, x, mutable=['batch_stats'], use_running_average=False)\n\n During eval we would define BN with ``use_running_average=True`` and use the\n batch_stats collection from training to set the statistics. In this case\n we are not mutating the batch statistics collection, and needn't mark it\n mutable::\n\n >>> y = BN.apply(variables, x, mutable=['batch_stats'], use_running_average=True)\n\n Attributes:\n use_running_average: if True, the statistics stored in batch_stats will be\n used instead of computing the batch statistics on the input.\n axis: the feature or non-batch axis of the input.\n momentum: decay rate for the exponential moving average of the batch\n statistics.\n epsilon: a small float added to variance to avoid dividing by zero.\n dtype: the dtype of the result (default: infer from input and params).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n use_bias: if True, bias (beta) is added.\n use_scale: if True, multiply by scale (gamma). When the next layer is linear\n (also e.g. nn.relu), this can be disabled since the scaling will be done\n by the next layer.\n bias_init: initializer for bias, by default, zero.\n scale_init: initializer for scale, by default, one.\n axis_name: the axis name used to combine batch statistics from multiple\n devices. See ``jax.pmap`` for a description of axis names (default: None).\n Note, this is only used for pmap and shard map. For SPMD jit, you do not\n need to manually synchronize. Just make sure that the axes are correctly\n annotated and XLA:SPMD will insert the necessary collectives.\n axis_index_groups: groups of axis indices within that named axis\n representing subsets of devices to reduce over (default: None). For\n example, ``[[0, 1], [2, 3]]`` would independently batch-normalize over the\n examples on the first two and last two devices. See ``jax.lax.psum`` for\n more details. This argument is currently not supported for SPMD jit.\n use_fast_variance: If true, use a faster, but less numerically stable,\n calculation for the variance.\n \"\"\"\n\n use_running_average: bool | None = None\n axis: int = -1\n momentum: float = 0.99\n epsilon: float = 1e-5\n dtype: Dtype | None = None\n param_dtype: Dtype = jnp.float32\n use_bias: bool = True\n use_scale: bool = True\n bias_init: Initializer = initializers.zeros\n scale_init: Initializer = initializers.ones\n axis_name: str | None = None\n axis_index_groups: Any = None\n use_fast_variance: bool = True\n force_float32_reductions: bool = True\n\n @compact\n def __call__(\n self,\n x,\n use_running_average: bool | None = None,\n *,\n mask: jax.Array | None = None,\n ):\n \"\"\"Normalizes the input using batch statistics.\n\n .. note::\n During initialization (when ``self.is_initializing()`` is ``True``) the running\n average of the batch statistics will not be updated. Therefore, the inputs\n fed during initialization don't need to match that of the actual input\n distribution and the reduction axis (set with ``axis_name``) does not have\n to exist.\n\n Args:\n x: the input to be normalized.\n use_running_average: if true, the statistics stored in batch_stats will be\n used instead of computing the batch statistics on the input.\n mask: Binary array of shape broadcastable to ``inputs`` tensor, indicating\n the positions for which the mean and variance should be computed.\n\n Returns:\n Normalized inputs (the same shape as inputs).\n \"\"\"\n\n use_running_average = module.merge_param(\n 'use_running_average', self.use_running_average, use_running_average\n )\n feature_axes = _canonicalize_axes(x.ndim, self.axis)\n reduction_axes = tuple(i for i in range(x.ndim) if i not in feature_axes)\n feature_shape = [x.shape[ax] for ax in feature_axes]\n\n ra_mean = self.variable(\n 'batch_stats',\n 'mean',\n lambda s: jnp.zeros(\n s,\n jnp.float32 if self.force_float32_reductions else self.param_dtype,\n ),\n feature_shape,\n )\n ra_var = self.variable(\n 'batch_stats',\n 'var',\n lambda s: jnp.ones(\n s,\n jnp.float32 if self.force_float32_reductions else self.param_dtype,\n ),\n feature_shape,\n )\n\n if use_running_average:\n mean = (\n ra_mean.value\n if self.force_float32_reductions\n else jnp.asarray(ra_mean.value, self.param_dtype)\n )\n var = (\n ra_var.value\n if self.force_float32_reductions\n else jnp.asarray(ra_var.value, self.param_dtype)\n )\n else:\n mean, var = _compute_stats(\n x,\n reduction_axes,\n dtype=self.dtype,\n axis_name=self.axis_name if not self.is_initializing() else None,\n axis_index_groups=self.axis_index_groups,\n use_fast_variance=self.use_fast_variance,\n mask=mask,\n force_float32_reductions=self.force_float32_reductions,\n )\n\n if not self.is_initializing():\n ra_mean.value = (\n self.momentum * ra_mean.value + (1 - self.momentum) * mean\n )\n ra_var.value = self.momentum * ra_var.value + (1 - self.momentum) * var\n\n return _normalize(\n self,\n x,\n mean,\n var,\n reduction_axes,\n feature_axes,\n self.dtype,\n self.param_dtype,\n self.epsilon,\n self.use_bias,\n self.use_scale,\n self.bias_init,\n self.scale_init,\n self.force_float32_reductions,\n )\n\n\nclass LayerNorm(Module):\n \"\"\"Layer normalization (https://arxiv.org/abs/1607.06450).\n\n LayerNorm normalizes the activations of the layer for each given example in a\n batch independently, rather than across a batch like Batch Normalization.\n i.e. applies a transformation that maintains the mean activation within\n each example close to 0 and the activation standard deviation close to 1.\n\n .. note::\n This normalization operation is identical to InstanceNorm and GroupNorm;\n the difference is simply which axes are reduced and the shape of the feature\n axes (i.e. the shape of the learnable scale and bias parameters).\n\n Example usage::\n\n >>> import flax.linen as nn\n >>> import jax\n >>> import numpy as np\n\n >>> x = jax.random.normal(jax.random.key(0), (3, 4, 5, 6))\n >>> layer = nn.LayerNorm()\n >>> variables = layer.init(jax.random.key(1), x)\n >>> variables\n {'params': {'scale': Array([1., 1., 1., 1., 1., 1.], dtype=float32), 'bias': Array([0., 0., 0., 0., 0., 0.], dtype=float32)}}\n >>> y = layer.apply(variables, x)\n\n >>> y = nn.LayerNorm(reduction_axes=(1, 2, 3)).apply(variables, x)\n >>> y2 = nn.GroupNorm(num_groups=1).apply(variables, x)\n >>> np.testing.assert_allclose(y, y2)\n\n >>> y = nn.LayerNorm(reduction_axes=(1, 2), feature_axes=-1).apply(variables, x)\n >>> y2 = nn.InstanceNorm(feature_axes=-1).apply(variables, x)\n >>> np.testing.assert_allclose(y, y2)\n\n Attributes:\n epsilon: A small float added to variance to avoid dividing by zero.\n dtype: the dtype of the result (default: infer from input and params).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n use_bias: If True, bias (beta) is added.\n use_scale: If True, multiply by scale (gamma). When the next layer is linear\n (also e.g. nn.relu), this can be disabled since the scaling will be done\n by the next layer.\n bias_init: Initializer for bias, by default, zero.\n scale_init: Initializer for scale, by default, one.\n reduction_axes: Axes for computing normalization statistics.\n feature_axes: Feature axes for learned bias and scaling.\n axis_name: the axis name used to combine batch statistics from multiple\n devices. See ``jax.pmap`` for a description of axis names (default: None).\n This is only needed if the model is subdivided across devices, i.e. the\n array being normalized is sharded across devices within a pmap or shard\n map. For SPMD jit, you do not need to manually synchronize. Just make sure\n that the axes are correctly annotated and XLA:SPMD will insert the\n necessary collectives.\n axis_index_groups: groups of axis indices within that named axis\n representing subsets of devices to reduce over (default: None). For\n example, ``[[0, 1], [2, 3]]`` would independently batch-normalize over the\n examples on the first two and last two devices. See ``jax.lax.psum`` for\n more details. This argument is currently not supported for SPMD jit.\n use_fast_variance: If true, use a faster, but less numerically stable,\n calculation for the variance.\n \"\"\"\n\n epsilon: float = 1e-6\n dtype: Dtype | None = None\n param_dtype: Dtype = jnp.float32\n use_bias: bool = True\n use_scale: bool = True\n bias_init: Initializer = initializers.zeros\n scale_init: Initializer = initializers.ones\n reduction_axes: Axes = -1\n feature_axes: Axes = -1\n axis_name: str | None = None\n axis_index_groups: Any = None\n use_fast_variance: bool = True\n force_float32_reductions: bool = True\n\n @compact\n def __call__(self, x, *, mask: jax.Array | None = None):\n \"\"\"Applies layer normalization on the input.\n\n Args:\n x: the inputs\n mask: Binary array of shape broadcastable to ``inputs`` tensor, indicating\n the positions for which the mean and variance should be computed.\n\n Returns:\n Normalized inputs (the same shape as inputs).\n \"\"\"\n mean, var = _compute_stats(\n x,\n self.reduction_axes,\n self.dtype,\n self.axis_name,\n self.axis_index_groups,\n use_fast_variance=self.use_fast_variance,\n mask=mask,\n force_float32_reductions=self.force_float32_reductions,\n )\n\n return _normalize(\n self,\n x,\n mean,\n var,\n self.reduction_axes,\n self.feature_axes,\n self.dtype,\n self.param_dtype,\n self.epsilon,\n self.use_bias,\n self.use_scale,\n self.bias_init,\n self.scale_init,\n self.force_float32_reductions,\n )\n\n\nclass RMSNorm(Module):\n \"\"\"RMS Layer normalization (https://arxiv.org/abs/1910.07467).\n\n RMSNorm normalizes the activations of the layer for each given example in a\n batch independently, rather than across a batch like Batch Normalization.\n Unlike LayerNorm which re-centers the mean to be 0 and normalizes by the\n standard deviation of the activations, RMSNorm does not re-center at all\n and instead normalizes by the root mean square of the activations.\n\n Example usage::\n\n >>> import flax.linen as nn\n >>> import jax\n\n >>> x = jax.random.normal(jax.random.key(0), (5, 6))\n >>> layer = nn.RMSNorm()\n >>> variables = layer.init(jax.random.key(1), x)\n >>> variables\n {'params': {'scale': Array([1., 1., 1., 1., 1., 1.], dtype=float32)}}\n >>> y = layer.apply(variables, x)\n\n Attributes:\n epsilon: A small float added to variance to avoid dividing by zero.\n dtype: the dtype of the result (default: infer from input and params).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n use_scale: If True, multiply by scale (gamma). When the next layer is linear\n (also e.g. nn.relu), this can be disabled since the scaling will be done\n by the next layer.\n scale_init: Initializer for scale, by default, one.\n reduction_axes: Axes for computing normalization statistics.\n feature_axes: Feature axes for learned bias and scaling.\n axis_name: the axis name used to combine batch statistics from multiple\n devices. See ``jax.pmap`` for a description of axis names (default: None).\n This is only needed if the model is subdivided across devices, i.e. the\n array being normalized is sharded across devices within a pmap or shard\n map. For SPMD jit, you do not need to manually synchronize. Just make sure\n that the axes are correctly annotated and XLA:SPMD will insert the\n necessary collectives.\n axis_index_groups: groups of axis indices within that named axis\n representing subsets of devices to reduce over (default: None). For\n example, ``[[0, 1], [2, 3]]`` would independently batch-normalize over the\n examples on the first two and last two devices. See ``jax.lax.psum`` for\n more details. This argument is currently not supported for SPMD jit.\n use_fast_variance: If true, use a faster, but less numerically stable,\n calculation for the variance.\n \"\"\"\n\n epsilon: float = 1e-6\n dtype: Dtype | None = None\n param_dtype: Dtype = jnp.float32\n use_scale: bool = True\n scale_init: Initializer = initializers.ones\n reduction_axes: Axes = -1\n feature_axes: Axes = -1\n axis_name: str | None = None\n axis_index_groups: Any = None\n use_fast_variance: bool = True\n force_float32_reductions: bool = True\n\n @compact\n def __call__(self, x, *, mask: jax.Array | None = None):\n \"\"\"Applies RMS layer normalization on the input.\n\n Args:\n x: the inputs\n mask: Binary array of shape broadcastable to ``inputs`` tensor, indicating\n the positions for which the mean and variance should be computed.\n\n Returns:\n Normalized inputs (the same shape as inputs).\n \"\"\"\n mean, var = _compute_stats(\n x,\n self.reduction_axes,\n self.dtype,\n self.axis_name,\n self.axis_index_groups,\n use_mean=False,\n use_fast_variance=self.use_fast_variance,\n mask=mask,\n force_float32_reductions=self.force_float32_reductions,\n )\n\n return _normalize(\n self,\n x,\n mean,\n var,\n self.reduction_axes,\n self.feature_axes,\n self.dtype,\n self.param_dtype,\n self.epsilon,\n False,\n self.use_scale,\n initializers.zeros,\n self.scale_init,\n self.force_float32_reductions,\n )\n\n\nclass GroupNorm(Module):\n \"\"\"Group normalization (arxiv.org/abs/1803.08494).\n\n This op is similar to batch normalization, but statistics are shared across\n equally-sized groups of channels and not shared across batch dimension.\n Thus, group normalization does not depend on the batch composition and does\n not require maintaining internal state for storing statistics.\n The user should either specify the total number of channel groups or the\n number of channels per group.\n\n .. note::\n LayerNorm is a special case of GroupNorm where ``num_groups=1``, and\n InstanceNorm is a special case of GroupNorm where ``group_size=1``.\n\n Example usage::\n\n >>> import flax.linen as nn\n >>> import jax\n >>> import numpy as np\n\n >>> x = jax.random.normal(jax.random.key(0), (3, 4, 5, 6))\n >>> layer = nn.GroupNorm(num_groups=3)\n >>> variables = layer.init(jax.random.key(1), x)\n >>> variables\n {'params': {'scale': Array([1., 1., 1., 1., 1., 1.], dtype=float32), 'bias': Array([0., 0., 0., 0., 0., 0.], dtype=float32)}}\n >>> y = layer.apply(variables, x)\n\n >>> y = nn.GroupNorm(num_groups=1).apply(variables, x)\n >>> y2 = nn.LayerNorm(reduction_axes=(1, 2, 3)).apply(variables, x)\n >>> np.testing.assert_allclose(y, y2)\n\n >>> y = nn.GroupNorm(num_groups=None, group_size=1).apply(variables, x)\n >>> y2 = nn.InstanceNorm(feature_axes=-1).apply(variables, x)\n >>> np.testing.assert_allclose(y, y2)\n\n Attributes:\n num_groups: the total number of channel groups. The default value of 32 is\n proposed by the original group normalization paper.\n group_size: the number of channels in a group.\n epsilon: A small float added to variance to avoid dividing by zero.\n dtype: the dtype of the result (default: infer from input and params).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n use_bias: If True, bias (beta) is added.\n use_scale: If True, multiply by scale (gamma). When the next layer is linear\n (also e.g. nn.relu), this can be disabled since the scaling will be done\n by the next layer.\n bias_init: Initializer for bias, by default, zero.\n scale_init: Initializer for scale, by default, one.\n reduction_axes: List of axes used for computing normalization statistics.\n This list must include the final dimension, which is assumed to be the\n feature axis. Furthermore, if the input used at call time has additional\n leading axes compared to the data used for initialisation, for example due\n to batching, then the reduction axes need to be defined explicitly.\n axis_name: the axis name used to combine batch statistics from multiple\n devices. See ``jax.pmap`` for a description of axis names (default: None).\n This is only needed if the model is subdivided across devices, i.e. the\n array being normalized is sharded across devices within a pmap or shard\n map. For SPMD jit, you do not need to manually synchronize. Just make sure\n that the axes are correctly annotated and XLA:SPMD will insert the\n necessary collectives.\n axis_index_groups: groups of axis indices within that named axis\n representing subsets of devices to reduce over (default: None). For\n example, ``[[0, 1], [2, 3]]`` would independently batch-normalize over the\n examples on the first two and last two devices. See ``jax.lax.psum`` for\n more details. This argument is currently not supported for SPMD jit.\n use_fast_variance: If true, use a faster, but less numerically stable,\n calculation for the variance.\n \"\"\"\n\n num_groups: int | None = 32\n group_size: int | None = None\n epsilon: float = 1e-6\n dtype: Dtype | None = None\n param_dtype: Dtype = jnp.float32\n use_bias: bool = True\n use_scale: bool = True\n bias_init: Initializer = initializers.zeros\n scale_init: Initializer = initializers.ones\n reduction_axes: Axes | None = None\n axis_name: str | None = None\n axis_index_groups: Any = None\n use_fast_variance: bool = True\n force_float32_reductions: bool = True\n\n @compact\n def __call__(self, x, *, mask: jax.Array | None = None):\n \"\"\"Applies group normalization to the input (arxiv.org/abs/1803.08494).\n\n Args:\n x: the input of shape ``...C`` where ``C`` is a channels dimension and ``...``\n represents an arbitrary number of extra dimensions that can be used to\n accumulate statistics over. If no reduction axes have been specified\n then all additional dimensions ``...`` will be used to accumulate\n statistics apart from the leading dimension which is assumed to\n represent the batch.\n mask: Binary array of shape broadcastable to ``inputs`` tensor, indicating\n the positions for which the mean and variance should be computed.\n\n Returns:\n Normalized inputs (the same shape as inputs).\n \"\"\"\n if self.reduction_axes is not None:\n reduction_axes = self.reduction_axes\n else:\n reduction_axes = list(range(1, x.ndim - 1)) + [-1]\n feature_axis = -1\n\n reduction_axes = _canonicalize_axes(x.ndim, reduction_axes)\n\n if reduction_axes[-1] != (feature_axis % x.ndim):\n raise ValueError(\n 'The reduction axes must include the final dimension '\n 'as this is assumed to be the feature axis.'\n )\n\n if (self.num_groups is None and self.group_size is None) or (\n self.num_groups is not None and self.group_size is not None\n ):\n raise ValueError(\n 'Either `num_groups` or `group_size` should be '\n 'specified. If `group_size` is to be specified, '\n 'pass `num_groups=None` as argument to override '\n 'the default `num_groups` value of 32.'\n )\n\n channels = x.shape[-1]\n if self.group_size is not None:\n if channels % self.group_size != 0:\n raise ValueError(\n 'Number of channels ({}) is not multiple of the '\n 'group size ({}).'.format(channels, self.group_size)\n )\n num_groups = channels // self.group_size\n else:\n num_groups = self.num_groups\n assert isinstance(num_groups, int)\n\n if num_groups <= 0 or channels % num_groups != 0:\n raise ValueError(\n 'Number of groups ({}) does not divide the number'\n ' of channels ({}).'.format(num_groups, channels)\n )\n\n group_size = x.shape[-1] // num_groups\n group_shape = x.shape[:-1] + (num_groups, group_size)\n\n if mask is not None:\n mask = mask.reshape(mask.shape[:-1] + (num_groups, group_size))\n\n mean, var = _compute_stats(\n x.reshape(group_shape),\n list(reduction_axes[:-1]) + [-1],\n self.dtype,\n self.axis_name,\n self.axis_index_groups,\n use_fast_variance=self.use_fast_variance,\n mask=mask,\n force_float32_reductions=self.force_float32_reductions,\n )\n mean = jnp.repeat(mean, group_size, axis=-1)\n var = jnp.repeat(var, group_size, axis=-1)\n\n return _normalize(\n self,\n x,\n mean,\n var,\n reduction_axes[:-1],\n (feature_axis,),\n self.dtype,\n self.param_dtype,\n self.epsilon,\n self.use_bias,\n self.use_scale,\n self.bias_init,\n self.scale_init,\n self.force_float32_reductions,\n )\n\n\nclass InstanceNorm(Module):\n \"\"\"Instance normalization (https://arxiv.org/abs/1607.08022v3).\n\n InstanceNorm normalizes the activations of the layer for each channel (rather\n than across all channels like Layer Normalization), and for each given example\n in a batch independently (rather than across an entire batch like Batch\n Normalization). i.e. applies a transformation that maintains the mean activation\n within each channel within each example close to 0 and the activation standard\n deviation close to 1.\n\n .. note::\n This normalization operation is identical to LayerNorm and GroupNorm; the\n difference is simply which axes are reduced and the shape of the feature axes\n (i.e. the shape of the learnable scale and bias parameters).\n\n Example usage::\n\n >>> import flax.linen as nn\n >>> import jax\n >>> import numpy as np\n\n >>> # dimensions: (batch, height, width, channel)\n >>> x = jax.random.normal(jax.random.key(0), (2, 3, 4, 5))\n >>> layer = nn.InstanceNorm()\n >>> variables = layer.init(jax.random.key(1), x)\n >>> variables\n {'params': {'scale': Array([1., 1., 1., 1., 1.], dtype=float32), 'bias': Array([0., 0., 0., 0., 0.], dtype=float32)}}\n >>> y = layer.apply(variables, x)\n\n >>> # having a channel_axis of -1 in InstanceNorm is identical to reducing all non-batch,\n >>> # non-channel axes and using the feature_axes as the feature_axes in LayerNorm\n >>> y2 = nn.LayerNorm(reduction_axes=[1, 2], feature_axes=-1).apply(variables, x)\n >>> np.testing.assert_allclose(y, y2, atol=1e-7)\n >>> y3 = nn.GroupNorm(num_groups=x.shape[-1]).apply(variables, x)\n >>> np.testing.assert_allclose(y, y3, atol=1e-7)\n\n Attributes:\n epsilon: A small float added to variance to avoid dividing by zero.\n dtype: the dtype of the result (default: infer from input and params).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n use_bias: If True, bias (beta) is added.\n use_scale: If True, multiply by scale (gamma). When the next layer is linear\n (also e.g. nn.relu), this can be disabled since the scaling will be done\n by the next layer.\n bias_init: Initializer for bias, by default, zero.\n scale_init: Initializer for scale, by default, one.\n feature_axes: Axes for features. The learned bias and scaling parameters will\n be in the shape defined by the feature axes. All other axes except the batch\n axes (which is assumed to be the leading axis) will be reduced.\n axis_name: the axis name used to combine batch statistics from multiple\n devices. See ``jax.pmap`` for a description of axis names (default: None).\n This is only needed if the model is subdivided across devices, i.e. the\n array being normalized is sharded across devices within a pmap or shard\n map. For SPMD jit, you do not need to manually synchronize. Just make sure\n that the axes are correctly annotated and XLA:SPMD will insert the\n necessary collectives.\n axis_index_groups: groups of axis indices within that named axis\n representing subsets of devices to reduce over (default: None). For\n example, ``[[0, 1], [2, 3]]`` would independently batch-normalize over the\n examples on the first two and last two devices. See ``jax.lax.psum`` for\n more details. This argument is currently not supported for SPMD jit.\n use_fast_variance: If true, use a faster, but less numerically stable,\n calculation for the variance.\n \"\"\"\n\n epsilon: float = 1e-6\n dtype: Dtype | None = None\n param_dtype: Dtype = jnp.float32\n use_bias: bool = True\n use_scale: bool = True\n bias_init: Initializer = initializers.zeros\n scale_init: Initializer = initializers.ones\n feature_axes: Axes = -1\n axis_name: str | None = None\n axis_index_groups: Any = None\n use_fast_variance: bool = True\n force_float32_reductions: bool = True\n\n @compact\n def __call__(self, x, *, mask: jax.Array | None = None):\n \"\"\"Applies instance normalization on the input.\n\n Args:\n x: the inputs\n mask: Binary array of shape broadcastable to ``inputs`` tensor, indicating\n the positions for which the mean and variance should be computed.\n\n Returns:\n Normalized inputs (the same shape as inputs).\n \"\"\"\n feature_axes = _canonicalize_axes(x.ndim, self.feature_axes)\n if 0 in feature_axes:\n raise ValueError('The channel axes cannot include the leading dimension '\n 'as this is assumed to be the batch axis.')\n reduction_axes = [i for i in range(1, x.ndim) if i not in feature_axes]\n\n mean, var = _compute_stats(\n x,\n reduction_axes,\n self.dtype,\n self.axis_name,\n self.axis_index_groups,\n use_fast_variance=self.use_fast_variance,\n mask=mask,\n force_float32_reductions=self.force_float32_reductions,\n )\n\n return _normalize(\n self,\n x,\n mean,\n var,\n reduction_axes,\n feature_axes,\n self.dtype,\n self.param_dtype,\n self.epsilon,\n self.use_bias,\n self.use_scale,\n self.bias_init,\n self.scale_init,\n self.force_float32_reductions,\n )\n\n\nclass SpectralNorm(Module):\n \"\"\"Spectral normalization.\n\n See:\n\n - https://arxiv.org/abs/1802.05957\n - https://arxiv.org/abs/1805.08318\n - https://arxiv.org/abs/1809.11096\n\n Spectral normalization normalizes the weight params so that the spectral\n norm of the matrix is equal to 1. This is implemented as a layer wrapper\n where each wrapped layer will have its params spectral normalized before\n computing its ``__call__`` output.\n\n .. note::\n The initialized variables dict will contain, in addition to a 'params'\n collection, a separate 'batch_stats' collection that will contain a\n ``u`` vector and ``sigma`` value, which are intermediate values used\n when performing spectral normalization. During training, we pass in\n ``update_stats=True`` and ``mutable=['batch_stats']`` so that ``u``\n and ``sigma`` are updated with the most recently computed values using\n power iteration. This will help the power iteration method approximate\n the true singular value more accurately over time. During eval, we pass\n in ``update_stats=False`` to ensure we get deterministic behavior from\n the model.\n\n Example usage::\n\n >>> import flax, flax.linen as nn\n >>> import jax, jax.numpy as jnp\n >>> import optax\n\n >>> class Foo(nn.Module):\n ... @nn.compact\n ... def __call__(self, x, train):\n ... x = nn.Dense(3)(x)\n ... # only spectral normalize the params of the second Dense layer\n ... x = nn.SpectralNorm(nn.Dense(4))(x, update_stats=train)\n ... x = nn.Dense(5)(x)\n ... return x\n\n >>> # init\n >>> x = jnp.ones((1, 2))\n >>> y = jnp.ones((1, 5))\n >>> model = Foo()\n >>> variables = model.init(jax.random.PRNGKey(0), x, train=False)\n >>> flax.core.freeze(jax.tree_util.tree_map(jnp.shape, variables))\n FrozenDict({\n batch_stats: {\n SpectralNorm_0: {\n Dense_1/kernel/sigma: (),\n Dense_1/kernel/u: (1, 4),\n },\n },\n params: {\n Dense_0: {\n bias: (3,),\n kernel: (2, 3),\n },\n Dense_1: {\n bias: (4,),\n kernel: (3, 4),\n },\n Dense_2: {\n bias: (5,),\n kernel: (4, 5),\n },\n },\n })\n\n >>> # train\n >>> def train_step(variables, x, y):\n ... def loss_fn(params):\n ... logits, updates = model.apply(\n ... {'params': params, 'batch_stats': variables['batch_stats']},\n ... x,\n ... train=True,\n ... mutable=['batch_stats'],\n ... )\n ... loss = jnp.mean(optax.l2_loss(predictions=logits, targets=y))\n ... return loss, updates\n ...\n ... (loss, updates), grads = jax.value_and_grad(loss_fn, has_aux=True)(\n ... variables['params']\n ... )\n ... return {\n ... 'params': jax.tree_util.tree_map(\n ... lambda p, g: p - 0.1 * g, variables['params'], grads\n ... ),\n ... 'batch_stats': updates['batch_stats'],\n ... }, loss\n >>> for _ in range(10):\n ... variables, loss = train_step(variables, x, y)\n\n >>> # inference / eval\n >>> out = model.apply(variables, x, train=False)\n\n Attributes:\n layer_instance: Module instance that is wrapped with SpectralNorm\n n_steps: How many steps of power iteration to perform to approximate the\n singular value of the weight params.\n epsilon: A small float added to l2-normalization to avoid dividing by zero.\n dtype: the dtype of the result (default: infer from input and params).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n error_on_non_matrix: Spectral normalization is only defined on matrices. By\n default, this module will return scalars unchanged and flatten\n higher-order tensors in their leading dimensions. Setting this flag to\n True will instead throw an error if a weight tensor with dimension greater\n than 2 is used by the layer.\n collection_name: Name of the collection to store intermediate values used\n when performing spectral normalization.\n \"\"\"\n\n layer_instance: Module\n n_steps: int = 1\n epsilon: float = 1e-12\n dtype: Dtype | None = None\n param_dtype: Dtype = jnp.float32\n error_on_non_matrix: bool = False\n collection_name: str = 'batch_stats'\n\n @compact\n def __call__(self, *args, update_stats: bool, **kwargs):\n \"\"\"Compute the largest singular value of the weights in ``self.layer_instance``\n using power iteration and normalize the weights using this value before\n computing the ``__call__`` output.\n\n Args:\n *args: positional arguments to be passed into the call method of the\n underlying layer instance in ``self.layer_instance``.\n update_stats: if True, update the internal ``u`` vector and ``sigma``\n value after computing their updated values using power iteration. This\n will help the power iteration method approximate the true singular value\n more accurately over time.\n **kwargs: keyword arguments to be passed into the call method of the\n underlying layer instance in ``self.layer_instance``.\n\n Returns:\n Output of the layer using spectral normalized weights.\n \"\"\"\n\n def layer_forward(layer_instance):\n return layer_instance(*args, **kwargs)\n\n return transforms.map_variables(\n layer_forward,\n trans_in_fn=lambda vs: jax.tree_util.tree_map_with_path(\n functools.partial(\n self._spectral_normalize,\n update_stats=update_stats,\n ),\n vs,\n ),\n init=self.is_initializing(),\n mutable=True,\n )(self.layer_instance)\n\n def _spectral_normalize(self, path, vs, update_stats):\n \"\"\"Compute the largest singular value using power iteration and normalize\n the variables ``vs`` using this value. This is intended to be a helper\n function used in this Module's ``__call__`` method in conjunction with\n ``nn.transforms.map_variables`` and ``jax.tree_util.tree_map_with_path``.\n\n Args:\n path: dict key path, used for naming the ``u`` and ``sigma`` variables\n vs: variables to be spectral normalized\n update_stats: if True, update the ``u`` vector and ``sigma`` variables\n after computing their updated values using power iteration. This will\n help the power iteration method approximate the true singular value\n more accurately over time.\n \"\"\"\n value = jnp.asarray(vs)\n value_shape = value.shape\n\n # Skip and return value if input is scalar, vector or if number of power\n # iterations is less than 1\n if value.ndim <= 1 or self.n_steps < 1:\n return value\n # Handle higher-order tensors.\n elif value.ndim > 2:\n if self.error_on_non_matrix:\n raise ValueError(\n f'Input is {value.ndim}D but error_on_non_matrix is True'\n )\n else:\n value = jnp.reshape(value, (-1, value.shape[-1]))\n\n u_var_name = (\n self.layer_instance.name\n + '/'\n + '/'.join(dict_key.key for dict_key in path[1:])\n + '/u'\n )\n u_var = self.variable(\n self.collection_name,\n u_var_name,\n jax.random.normal,\n self.make_rng('params')\n if not self.has_variable(self.collection_name, u_var_name)\n else None,\n (1, value.shape[-1]),\n self.param_dtype,\n )\n u0 = u_var.value\n sigma_var_name = (\n self.layer_instance.name\n + '/'\n + '/'.join(dict_key.key for dict_key in path[1:])\n + '/sigma'\n )\n sigma_var = self.variable(\n self.collection_name, sigma_var_name, jnp.ones, (), self.param_dtype\n )\n\n # Power iteration for the weight's singular value.\n for _ in range(self.n_steps):\n v0 = _l2_normalize(\n jnp.matmul(u0, value.transpose([1, 0])), eps=self.epsilon\n )\n u0 = _l2_normalize(jnp.matmul(v0, value), eps=self.epsilon)\n\n u0 = jax.lax.stop_gradient(u0)\n v0 = jax.lax.stop_gradient(v0)\n\n sigma = jnp.matmul(jnp.matmul(v0, value), jnp.transpose(u0))[0, 0]\n\n value /= jnp.where(sigma != 0, sigma, 1)\n value_bar = value.reshape(value_shape)\n\n if update_stats:\n u_var.value = u0\n sigma_var.value = sigma\n\n dtype = dtypes.canonicalize_dtype(vs, u0, v0, sigma, dtype=self.dtype)\n return jnp.asarray(value_bar, dtype)\n\n\nclass WeightNorm(Module):\n \"\"\"L2 weight normalization (https://arxiv.org/abs/1602.07868).\n\n Weight normalization normalizes the weight params so that the l2-norm of\n the matrix is equal to 1. This is implemented as a layer wrapper where\n each wrapped layer will have its params l2-normalized before computing\n its ``__call__`` output.\n\n Example usage::\n\n >>> import flax, flax.linen as nn\n >>> import jax, jax.numpy as jnp\n\n >>> class Baz(nn.Module):\n ... @nn.compact\n ... def __call__(self, x):\n ... return nn.Dense(2)(x)\n\n >>> class Bar(nn.Module):\n ... @nn.compact\n ... def __call__(self, x):\n ... x = Baz()(x)\n ... x = nn.Dense(3)(x)\n ... x = Baz()(x)\n ... x = nn.Dense(3)(x)\n ... return x\n\n >>> class Foo(nn.Module):\n ... @nn.compact\n ... def __call__(self, x):\n ... x = nn.Dense(3)(x)\n ... # l2-normalize all params of the second Dense layer\n ... x = nn.WeightNorm(nn.Dense(4), variable_filter=None)(x)\n ... x = nn.Dense(5)(x)\n ... # l2-normalize all kernels in the Bar submodule and all params in\n ... # the Baz submodule\n ... x = nn.WeightNorm(Bar(), variable_filter={'kernel', 'Baz'})(x)\n ... return x\n\n >>> # init\n >>> x = jnp.ones((1, 2))\n >>> model = Foo()\n >>> variables = model.init(jax.random.key(0), x)\n >>> flax.core.freeze(jax.tree_util.tree_map(jnp.shape, variables))\n FrozenDict({\n params: {\n Bar_0: {\n Baz_0: {\n Dense_0: {\n bias: (2,),\n kernel: (5, 2),\n },\n },\n Baz_1: {\n Dense_0: {\n bias: (2,),\n kernel: (3, 2),\n },\n },\n Dense_0: {\n bias: (3,),\n kernel: (2, 3),\n },\n Dense_1: {\n bias: (3,),\n kernel: (2, 3),\n },\n },\n Dense_0: {\n bias: (3,),\n kernel: (2, 3),\n },\n Dense_1: {\n bias: (4,),\n kernel: (3, 4),\n },\n Dense_2: {\n bias: (5,),\n kernel: (4, 5),\n },\n WeightNorm_0: {\n Dense_1/bias/scale: (4,),\n Dense_1/kernel/scale: (4,),\n },\n WeightNorm_1: {\n Bar_0/Baz_0/Dense_0/bias/scale: (2,),\n Bar_0/Baz_0/Dense_0/kernel/scale: (2,),\n Bar_0/Baz_1/Dense_0/bias/scale: (2,),\n Bar_0/Baz_1/Dense_0/kernel/scale: (2,),\n Bar_0/Dense_0/kernel/scale: (3,),\n Bar_0/Dense_1/kernel/scale: (3,),\n },\n },\n })\n\n Attributes:\n layer_instance: Module instance that is wrapped with WeightNorm\n epsilon: A small float added to l2-normalization to avoid dividing by zero.\n dtype: the dtype of the result (default: infer from input and params).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n use_scale: If True, creates a learnable variable ``scale`` that is\n multiplied to the ``layer_instance`` variables after l2-normalization.\n scale_init: Initialization function for the scaling function.\n feature_axes: The feature axes dimension(s). The l2-norm is calculated by\n reducing the ``layer_instance`` variables over the remaining (non-feature)\n axes. Therefore a separate l2-norm value is calculated and a separate\n scale (if ``use_scale=True``) is learned for each specified feature. By\n default, the trailing dimension is treated as the feature axis.\n variable_filter: An optional iterable that contains string items. The\n WeightNorm layer will selectively apply l2-normalization to the\n ``layer_instance`` variables whose key path (delimited by '/') has a match\n with ``variable_filter``. For example, ``variable_filter={'kernel'}`` will\n only apply l2-normalization to variables whose key path contains 'kernel'.\n By default, ``variable_filter={'kernel'}``.\n \"\"\"\n\n layer_instance: Module\n epsilon: float = 1e-12\n dtype: Dtype | None = None\n param_dtype: Dtype = jnp.float32\n use_scale: bool = True\n scale_init: Initializer = initializers.ones\n feature_axes: Axes | None = -1\n variable_filter: Iterable | None = dataclasses.field(\n default_factory=lambda: {'kernel'}\n )\n\n @compact\n def __call__(self, *args, **kwargs):\n \"\"\"Compute the l2-norm of the weights in ``self.layer_instance``\n and normalize the weights using this value before computing the\n ``__call__`` output.\n\n Args:\n *args: positional arguments to be passed into the call method of the\n underlying layer instance in ``self.layer_instance``.\n **kwargs: keyword arguments to be passed into the call method of the\n underlying layer instance in ``self.layer_instance``.\n\n Returns:\n Output of the layer using l2-normalized weights.\n \"\"\"\n\n def layer_forward(layer_instance):\n return layer_instance(*args, **kwargs)\n\n return transforms.map_variables(\n layer_forward,\n trans_in_fn=lambda vs: jax.tree_util.tree_map_with_path(\n self._l2_normalize,\n vs,\n ),\n init=self.is_initializing(),\n )(self.layer_instance)\n\n def _l2_normalize(self, path, vs):\n \"\"\"Compute the l2-norm and normalize the variables ``vs`` using this\n value. This is intended to be a helper function used in this Module's\n ``__call__`` method in conjunction with ``nn.transforms.map_variables``\n and ``jax.tree_util.tree_map_with_path``.\n\n Args:\n path: dict key path, used for naming the ``scale`` variable\n vs: variables to be l2-normalized\n \"\"\"\n value = jnp.asarray(vs)\n str_path = (\n self.layer_instance.name\n + '/'\n + jax.tree_util.keystr(path[1:], simple=True, separator='/')\n )\n if self.variable_filter:\n for variable_name in self.variable_filter:\n if variable_name in str_path:\n break\n else:\n return value\n\n if self.feature_axes is None:\n feature_axes = ()\n reduction_axes = tuple(i for i in range(value.ndim))\n else:\n feature_axes = _canonicalize_axes(value.ndim, self.feature_axes)\n reduction_axes = tuple(\n i for i in range(value.ndim) if i not in feature_axes\n )\n\n feature_shape = [1] * value.ndim\n reduced_feature_shape = []\n for ax in feature_axes:\n feature_shape[ax] = value.shape[ax]\n reduced_feature_shape.append(value.shape[ax])\n\n value_bar = _l2_normalize(value, axis=reduction_axes, eps=self.epsilon)\n\n args = [vs]\n if self.use_scale:\n scale = self.param(\n str_path + '/scale',\n self.scale_init,\n reduced_feature_shape,\n self.param_dtype,\n ).reshape(feature_shape)\n value_bar *= scale\n args.append(scale)\n\n dtype = dtypes.canonicalize_dtype(*args, dtype=self.dtype)\n return jnp.asarray(value_bar, dtype)\n" }, { "path": "flax/linen/partitioning.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Legacy utilities for working with pjit and partitioned models.\n\n**Experimental: please give feedback, and expect changes.**\n\nThis module introduces `axis_rules`, `logical_to_mesh_axes`,\n`with_sharding_constraint` for appyling pjit sharding constraints in terms of\n\"logical named axes\" rather than pjit's default mesh axes.\n\nAdditionally, flax linen methods `param_with_axes` and `variable_with_axes`\nare introduced alongside `get_axis_names` for defining variables and parameters\nand variables with logical axis name annotations that are managed as metadata.\n\nLastly, `*_with_axes` versions of `nn.scan` and `nn.vmap` are introduced to add\nlogical axis metadata to the underlying Lifted transformations.\n\"\"\"\n\nimport functools\nimport re\nfrom typing import (Any, Optional, Tuple)\nfrom collections.abc import Callable, Mapping\n\nimport flax\nfrom flax import linen as nn\nfrom flax import struct\nfrom flax.core.frozen_dict import freeze\nfrom flax.core.frozen_dict import unfreeze\nfrom flax.core.scope import (\n CollectionFilter as CollectionFilter,\n PRNGSequenceFilter as PRNGSequenceFilter,\n)\nfrom flax.core.spmd import logical_axis_rules as axis_rules # pylint: disable=unused-import\nfrom flax.core.spmd import set_logical_axis_rules as set_axis_rules # pylint: disable=unused-import\nfrom flax.core.spmd import get_logical_axis_rules as get_axis_rules # pylint: disable=unused-import\nfrom flax.linen.spmd import _is_logical_spec\nfrom flax.linen.spmd import _with_sharding_constraint # pylint: disable=unused-import\nfrom flax.linen.spmd import logical_to_mesh # pylint: disable=unused-import\nfrom flax.linen.spmd import logical_to_mesh_axes # pylint: disable=unused-import\nfrom flax.linen.spmd import RulesFallback\nfrom flax.linen.spmd import with_logical_constraint as with_sharding_constraint\nfrom flax.traverse_util import flatten_dict\nfrom flax.traverse_util import unflatten_dict\nfrom flax.typing import (\n Array,\n In as ScanIn, # pylint: disable=unused-import\n Out as ScanOut, # pylint: disable=unused-import\n InOutAxis,\n InOutScanAxis,\n LogicalRules, # pylint: disable=unused-import\n ArrayPytree, # pylint: disable=unused-import\n LogicalPartitionSpec, # pylint: disable=unused-import\n LogicalPartitionSpecPytree,\n PartitionSpecPytree, # pylint: disable=unused-import\n)\nimport jax\n\n\n# ------------------------------------------------------------------------------\n# NOTICE: This experimental partitioning utility API is deprecated\n#\n# We intend to continue supporting it indefinitely for those using it, but\n# we encourage new users to adopt the simpler metadata handling system found\n# in \"spmd.py\".\n# ------------------------------------------------------------------------------\n\n\n# Annotated parameters and Module axis metadata handling.\n# ------------------------------------------------------------------------------\n\n\n@struct.dataclass\nclass AxisMetadata:\n \"\"\"Contains a tuple of axis names, which is passed through FLAX.\"\"\"\n\n names: LogicalPartitionSpecPytree = struct.field(pytree_node=False)\n\n\ndef _param_with_axes_sow_reduce_fn(x, y):\n \"\"\"Reduction function for sow() calls.\n\n Args:\n x: Existing value, or () if there was none.\n y: New axis names sown.\n\n Returns:\n New axis names.\n\n Raises:\n TypeError: If the newly sown value is not an AxisMetadata.\n ValueError: If the newly sown axis names don't match previously sown axis\n names.\n AssertionError: If a previously sown value was truthy and not an\n AxisMetadata.\n \"\"\"\n if not isinstance(y, AxisMetadata):\n raise TypeError('Expected newly sown value to be an AxisMetadata')\n\n if isinstance(x, AxisMetadata):\n if x != y:\n raise ValueError(\n 'If axis names are sown twice, expected them to match. '\n f'Got {x} and {y}.'\n )\n elif x:\n # Shouldn't happen, so raise a fairly internal error.\n raise AssertionError(f'Non-initial-or-AxisMetadata value encountered: {x}')\n return y\n\n\ndef param_with_axes(\n name: str,\n init_fn,\n *init_args,\n axes: tuple[str, ...] | None = None,\n module: Optional['nn.Module'] = None,\n **init_kwargs,\n):\n \"\"\"Declares and returns a parameter with logical axes in the current Module.\n\n See :mod:`flax.linen.module.param` for original docstring.\n\n Args:\n name: The parameter name.\n init_fn: The function that will be called to compute the initial value\n of this variable. This function will only be called the first time\n this parameter is used in this module.\n *init_args: The positional arguments to pass to init_fn.\n axes: A tuple of axis names, must match the rank of the param array.\n module: Use an explicit module instead of deriving the most recent from\n dynamic module context.\n **init_kwargs: The key-word arguments to pass to init_fn.\n\n Returns:\n The value of the initialized parameter.\n\n Raises:\n TypeError: if axes specification is mal-formed.\n ValueError: if specified logical axes don't match parameter rank.\n \"\"\"\n # get current module if not explicitly provided\n if module is None:\n module = nn.module._context.module_stack[-1] # pylint: disable=protected-access\n assert module is not None\n # define/fetch parameter on that module\n module_param = module.param(name, init_fn, *init_args, **init_kwargs)\n if axes is not None:\n # apply logical axis constraint immediately\n module_param = with_sharding_constraint(\n module_param, jax.sharding.PartitionSpec(*axes)\n )\n # record logical axis constraint for global axis metadata\n module.sow(\n 'params_axes',\n f'{name}_axes',\n AxisMetadata(axes), # type: ignore\n reduce_fn=_param_with_axes_sow_reduce_fn,\n )\n return module_param\n\n\nclass PartitionedVariable(flax.core.scope.Variable):\n \"\"\"A PartitionedVariable object allows mutable access to a variable.\n\n PartitionedVariables are identified by a collection (e.g., \"batch_stats\") and\n a name (e.g., \"moving_mean\"). The value property gives access to the\n variable's content and can be assigned to for mutation. Additionally,\n PartitionedVariables enforce logical sharding constraints on both retrieval\n and assignment.\n \"\"\"\n\n def __init__(\n self,\n scope,\n collection: str,\n name: str,\n axes: tuple[str, ...] | None = None,\n fallback: RulesFallback = RulesFallback.AXIS_IS_UNSHARDED,\n ):\n \"\"\"Initializes a partitioned variable.\n\n Args:\n scope: The scope in which the variable is stored.\n collection: The collection of the variable (e.g., \"params\").\n name: The name of the variable (e.g., \"dense\").\n axes: logical axes name of variable.\n fallback: Fallback behavior if no matching rule is found.\n \"\"\"\n self.scope = scope\n self.collection = collection\n self.name = name\n self.axes = axes\n self.fallback = fallback\n\n @property\n def value(self):\n \"\"\"Returns the value of this Variable.\"\"\"\n value = self.scope.get_variable(self.collection, self.name)\n if self.axes is not None:\n value = with_sharding_constraint(value, self.axes, fallback=self.fallback)\n return value\n\n @value.setter\n def value(self, value):\n \"\"\"Updates the value of this Variable.\"\"\"\n if self.axes is not None:\n value = with_sharding_constraint(value, self.axes, fallback=self.fallback)\n self.scope.put_variable(self.collection, self.name, value)\n\n\ndef _core_variable_with_axes(\n scope,\n col: str,\n name: str,\n init_fn: Callable[..., Any],\n *init_args,\n axes: tuple[str, ...] | None = None,\n fallback: RulesFallback = RulesFallback.AXIS_IS_UNSHARDED,\n **init_kwargs,\n):\n \"\"\"Variant of flax core variable scope call with sharding constraints.\"\"\"\n scope.reserve(name)\n if not scope.has_variable(col, name):\n if not scope.is_mutable_collection(col):\n raise flax.errors.ScopeVariableNotFoundError(name, col, scope.path_text)\n init_value = init_fn(*init_args, **init_kwargs)\n if axes is not None:\n init_value = with_sharding_constraint(init_value, axes, fallback=fallback)\n scope.put_variable(col, name, init_value)\n return PartitionedVariable(scope, col, name, axes, fallback)\n\n\ndef variable_with_axes(\n collection: str,\n name: str,\n init_fn,\n *init_args,\n axes: tuple[str, ...] | None = None,\n module: Optional['nn.Module'] = None,\n fallback: RulesFallback = RulesFallback.AXIS_IS_UNSHARDED,\n **init_kwargs,\n):\n \"\"\"Declares and returns a variable with logical axes in the current Module.\n\n See :mod:`flax.linen.module.variable` for original docstring.\n\n Args:\n collection: The name of the variable collection.\n name: The variable name.\n init_fn: The function that will be called to compute the initial value\n of this variable. This function will only be called the first time\n this parameter is used in this module.\n *init_args: The positional arguments to pass to init_fn.\n axes: A tuple of axis names, must match the rank of the variable array.\n module: Use an explicit module instead of deriving the most recent from\n dynamic module context.\n fallback: How sharding should behave if there is no rule covering some axis.\n **init_kwargs: The key-word arguments to pass to init_fn.\n\n Returns:\n A flax `PartitionedVariable` object referencing the initialized variable\n array.\n\n Raises:\n TypeError: if axes specification is mal-formed.\n ValueError: if specified logical axes don't match parameter rank.\n \"\"\"\n # get current module if not explicitly provided\n if module is None:\n module = nn.module._context.module_stack[-1] # pylint: disable=protected-access\n assert module is not None\n module_var = _core_variable_with_axes(\n module.scope,\n collection,\n name,\n init_fn,\n *init_args,\n axes=axes,\n fallback=fallback,\n **init_kwargs,\n )\n if axes is not None:\n # record logical axis constraint for global axis metadata\n module.sow(\n f'{collection}_axes',\n f'{name}_axes',\n AxisMetadata(axes), # type: ignore\n reduce_fn=_param_with_axes_sow_reduce_fn,\n )\n return module_var\n\n\ndef get_axis_names(axes_metadata):\n \"\"\"Gets axis names for variables as logical PartitionSpecs.\n\n Args:\n axes_metadata: a single axes-metadata collection from a flax-initialized\n set of collections.\n\n Returns:\n Collection of Partitionspecs with logical axis names, with the \"_axes\"\n suffix on variable names removed to match original variable collection for\n annotations.\n \"\"\"\n\n def leaf_rewrite(x):\n return None if x is None else jax.sharding.PartitionSpec(*x)\n\n def rewrite(tree):\n return jax.tree_util.tree_map(leaf_rewrite, tree, is_leaf=_is_logical_spec)\n\n axes_metadata = unfreeze(axes_metadata) # pytype: disable=wrong-arg-types\n flat_dict = {\n re.sub(r'_axes$', '', '/'.join(k)): rewrite(v.names)\n for k, v in flatten_dict(axes_metadata).items()\n }\n return freeze(\n unflatten_dict({tuple(k.split('/')): v for k, v in flat_dict.items()})\n )\n\n\n# Metadata Aware Scan\n# -----------------------------------------------------------------------------\n\n\ndef _tree_map_axes(fn, tree):\n \"\"\"Only map over AxisMetadata leaves in pytree - identity for other leaves.\"\"\"\n safe_fn = lambda x: fn(x) if isinstance(x, AxisMetadata) else x\n return jax.tree_util.tree_map(\n safe_fn, tree, is_leaf=lambda x: isinstance(x, AxisMetadata)\n )\n\n\ndef _is_mutable(axis_col: str) -> bool:\n \"\"\"Determines whether a collection is mutable.\n\n For example, when a module is called with `module.apply(..., mutable=['z'])`,\n this function will return True for `axis_col='z'` and False otherwise.\n\n If there is no module in scope, this function will return True.\n\n Args:\n axis_col: Name of the collection in question.\n\n Returns:\n Whether it is currently mutable.\n \"\"\"\n last = nn.module._context.module_stack[-1] # pylint: disable=protected-access\n if last:\n return last.is_mutable_collection(axis_col)\n else:\n return True\n\n\n# uses this variable_transform to change 'params_axes' pytree as it bubbles\n# up / out from scan.\ndef _add_axis_to_metadata(fn, axis_pos, axis_name, axis_col='params_axes'):\n \"\"\"Insert a named axis to axes metadata.\"\"\"\n # Handle In() / Out() scan axis marker types.\n if hasattr(axis_pos, 'axis'):\n axis_pos = axis_pos.axis\n\n def insert_fn_leaf(names):\n if names is None:\n return names\n names = list(names)\n names.insert(axis_pos, axis_name)\n return tuple(names)\n\n def insert_fn(x):\n new_names = jax.tree_util.tree_map(\n insert_fn_leaf, x.names, is_leaf=_is_logical_spec\n )\n return x.replace(names=new_names)\n\n def remove_fn_leaf(names):\n if names is None:\n return names\n names = list(names)\n if names[axis_pos] != axis_name:\n raise ValueError(\n f'Expected axis {axis_name} at position {axis_pos} in '\n f'axis metadata {names}.'\n )\n names.pop(axis_pos)\n return tuple(names)\n\n def remove_fn(x):\n new_names = jax.tree_util.tree_map(\n remove_fn_leaf, x.names, is_leaf=_is_logical_spec\n )\n return x.replace(names=new_names)\n\n return nn.transforms.map_variables(\n fn,\n axis_col,\n mutable=_is_mutable(axis_col),\n trans_in_fn=lambda tree: _tree_map_axes(remove_fn, tree),\n trans_out_fn=lambda tree: _tree_map_axes(insert_fn, tree),\n )\n\n\n# pylint: disable=dangerous-default-value\ndef scan_with_axes(\n target: 'flax.linen.transforms.Target',\n variable_axes: Mapping[\n CollectionFilter, InOutScanAxis\n ] = {},\n variable_broadcast: CollectionFilter = False,\n variable_carry: CollectionFilter = False,\n split_rngs: Mapping[PRNGSequenceFilter, bool] = {},\n in_axes=0,\n out_axes=0,\n length: int | None = None,\n reverse: bool = False,\n unroll: int = 1,\n axis_name: str = 'layers',\n axes_collections: tuple[str, ...] = ('params',),\n data_transform: Callable[..., Any] | None = None,\n methods=None,\n) -> 'flax.linen.transforms.Target':\n \"\"\"Wrapped version of nn.scan that handles logical axis metadata.\"\"\"\n\n # we broadcast the static metadata collections.\n axes_filters = tuple(f'{col}_axes' for col in axes_collections)\n variable_broadcast = flax.core.scope.union_filters(\n variable_broadcast, axes_filters\n )\n\n # perform usual lifted scan\n scanned = flax.linen.transforms.lift_transform(\n flax.core.lift.scan,\n target,\n variable_axes=variable_axes,\n variable_broadcast=variable_broadcast,\n variable_carry=variable_carry,\n split_rngs=split_rngs,\n in_axes=in_axes,\n out_axes=out_axes,\n length=length,\n reverse=reverse,\n unroll=unroll,\n data_transform=data_transform,\n methods=methods,\n )\n\n # add scan axis to logical axes metadata\n for col in axes_collections:\n if col in variable_axes:\n scanned = _add_axis_to_metadata(\n scanned,\n axis_pos=variable_axes[col],\n axis_name=axis_name,\n axis_col=f'{col}_axes',\n )\n return scanned\n\n\n# pylint: disable=dangerous-default-value\ndef vmap_with_axes(\n target: 'flax.linen.transforms.Target',\n variable_axes: Mapping[\n CollectionFilter, InOutAxis\n ],\n split_rngs: Mapping[PRNGSequenceFilter, bool] = {},\n in_axes=0,\n out_axes=0,\n axis_size: int | None = None,\n axis_name: str | None = None,\n partitioning_axis_names: Mapping[Any, str] = {},\n spmd_axis_name: str | None = None,\n methods=None,\n) -> 'flax.linen.transforms.Target':\n \"\"\"Wrapped version of nn.vmap that handles logical axis metadata.\"\"\"\n\n # tell normal vmap to broadcast axis metadata.\n variable_axes = dict(variable_axes) # shallow copy\n for name in partitioning_axis_names:\n variable_axes[f'{name}_axes'] = None\n\n # perform usual lifted vmap\n vmapped = flax.linen.transforms.lift_transform(\n flax.core.lift.vmap,\n target,\n variable_axes=variable_axes,\n split_rngs=split_rngs,\n in_axes=in_axes,\n out_axes=out_axes,\n axis_size=axis_size,\n axis_name=axis_name,\n spmd_axis_name=spmd_axis_name,\n methods=methods,\n )\n\n for collection_name, axis in variable_axes.items():\n if collection_name in partitioning_axis_names:\n vmapped = _add_axis_to_metadata( # pylint: disable=protected-access\n vmapped,\n axis_pos=axis,\n axis_name=partitioning_axis_names[collection_name],\n axis_col=f'{collection_name}_axes',\n )\n\n return vmapped\n\n\n# Remat abstraction bug hotfix\n# ------------------------------------------------------------------------------\n# TODO(levskaya): upstream this fix into main flax.core.lift.remat.\n# Workaround a scan(remat(...)) abstraction bug by manually implementing a\n# static_argnums behavior for flax remat via closure before applying jax remat.\n\n\ndef core_remat_static(\n fn,\n variables=True,\n rngs=True,\n concrete=False,\n prevent_cse=True,\n static_argnums=(),\n policy=None,\n):\n \"\"\"Flax functional core remat version with static_argnums.\"\"\"\n\n static_argnums = tuple(sorted(static_argnums))\n\n def _repack_remat_args(dyn_args, static_args):\n \"\"\"Remake arg list from static and dynamic args given static_argnums.\"\"\"\n args = []\n s_cnt, d_cnt = 0, 0\n for i in range(len(dyn_args) + len(static_args)):\n if i in static_argnums:\n args.append(static_args[s_cnt])\n s_cnt += 1\n else:\n args.append(dyn_args[d_cnt])\n d_cnt += 1\n return tuple(args)\n\n def inner(scope_fn, repack_fn, variable_groups, rng_groups, *args):\n static_args = tuple(x for i, x in enumerate(args) if i in static_argnums)\n dyn_args = tuple(x for i, x in enumerate(args) if i not in static_argnums)\n\n # After JAX v0.3.16, concrete=False is a no-op and concrete=True raises\n # NotImplementedError. Starting in JAX v0.8.2, the concrete argument is\n # deprecated and will be removed in the future.\n if concrete:\n raise NotImplementedError(\n \"The concrete argument is deprecated. Use static_argnums instead.\"\n \" for more information, see\"\n \" https://docs.jax.dev/en/latest/jep/11830-new-remat-checkpoint.html\"\n )\n\n @functools.partial(\n jax.remat, prevent_cse=prevent_cse, policy=policy\n )\n @functools.wraps(fn)\n def rematted(variable_groups, rng_groups, *dyn_args):\n args = _repack_remat_args(dyn_args, static_args)\n scope = scope_fn(variable_groups, rng_groups)\n y = fn(scope, *args)\n return y, repack_fn(scope)\n\n return rematted(variable_groups, rng_groups, *dyn_args)\n\n return flax.core.lift.pack(\n inner, (variables,), (variables,), (rngs,), name='remat'\n )\n\n\ndef remat(\n target,\n variables=True,\n rngs=True,\n concrete=False,\n prevent_cse=True,\n static_argnums=(),\n policy=None,\n methods=None,\n):\n \"\"\"Flax lifted remat that supports static_argnums.\"\"\"\n return flax.linen.transforms.lift_transform(\n core_remat_static,\n target,\n variables=variables,\n rngs=rngs,\n concrete=concrete,\n prevent_cse=prevent_cse,\n static_argnums=static_argnums,\n policy=policy,\n methods=methods,\n )\n" }, { "path": "flax/linen/pooling.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Pooling modules.\"\"\"\n\nimport jax.numpy as jnp\nimport numpy as np\nfrom jax import lax\n\n\ndef pool(inputs, init, reduce_fn, window_shape, strides, padding):\n \"\"\"Helper function to define pooling functions.\n\n Pooling functions are implemented using the ReduceWindow XLA op.\n\n .. note::\n Be aware that pooling is not generally differentiable.\n That means providing a reduce_fn that is differentiable does not imply that\n pool is differentiable.\n\n Args:\n inputs: input data with dimensions (batch, window dims..., features).\n init: the initial value for the reduction\n reduce_fn: a reduce function of the form ``(T, T) -> T``.\n window_shape: a shape tuple defining the window to reduce over.\n strides: a sequence of ``n`` integers, representing the inter-window\n strides (default: ``(1, ..., 1)``).\n padding: either the string ``'SAME'``, the string ``'VALID'``, or a sequence\n of ``n`` ``(low, high)`` integer pairs that give the padding to apply before\n and after each spatial dimension.\n Returns:\n The output of the reduction for each window slice.\n \"\"\"\n num_batch_dims = inputs.ndim - (len(window_shape) + 1)\n strides = strides or (1,) * len(window_shape)\n assert len(window_shape) == len(\n strides\n ), f'len({window_shape}) must equal len({strides})'\n strides = (1,) * num_batch_dims + strides + (1,)\n dims = (1,) * num_batch_dims + window_shape + (1,)\n\n is_single_input = False\n if num_batch_dims == 0:\n # add singleton batch dimension because lax.reduce_window always\n # needs a batch dimension.\n inputs = inputs[None]\n strides = (1,) + strides\n dims = (1,) + dims\n is_single_input = True\n\n assert inputs.ndim == len(dims), f'len({inputs.shape}) != len({dims})'\n if not isinstance(padding, str):\n padding = tuple(map(tuple, padding))\n assert len(padding) == len(window_shape), (\n f'padding {padding} must specify pads for same number of dims as '\n f'window_shape {window_shape}'\n )\n assert all(\n [len(x) == 2 for x in padding]\n ), f'each entry in padding {padding} must be length 2'\n padding = ((0, 0),) + padding + ((0, 0),)\n y = lax.reduce_window(inputs, init, reduce_fn, dims, strides, padding)\n if is_single_input:\n y = jnp.squeeze(y, axis=0)\n return y\n\n\ndef avg_pool(\n inputs, window_shape, strides=None, padding='VALID', count_include_pad=True\n):\n \"\"\"Pools the input by taking the average over a window.\n\n Args:\n inputs: input data with dimensions (batch, window dims..., features).\n window_shape: a shape tuple defining the window to reduce over.\n strides: a sequence of ``n`` integers, representing the inter-window\n strides (default: ``(1, ..., 1)``).\n padding: either the string ``'SAME'``, the string ``'VALID'``, or a sequence\n of ``n`` ``(low, high)`` integer pairs that give the padding to apply before\n and after each spatial dimension (default: ``'VALID'``).\n count_include_pad: a boolean whether to include padded tokens\n in the average calculation (default: ``True``).\n Returns:\n The average for each window slice.\n \"\"\"\n y = pool(inputs, 0.0, lax.add, window_shape, strides, padding)\n if count_include_pad:\n y = y / np.prod(window_shape)\n else:\n div_shape = inputs.shape[:-1] + (1,)\n if len(div_shape) - 2 == len(window_shape):\n div_shape = (1,) + div_shape[1:]\n y = y / pool(\n jnp.ones(div_shape), 0.0, lax.add, window_shape, strides, padding\n )\n return y\n\n\ndef max_pool(inputs, window_shape, strides=None, padding='VALID'):\n \"\"\"Pools the input by taking the maximum of a window slice.\n\n Args:\n inputs: input data with dimensions (batch, window dims..., features).\n window_shape: a shape tuple defining the window to reduce over.\n strides: a sequence of ``n`` integers, representing the inter-window\n strides (default: ``(1, ..., 1)``).\n padding: either the string ``'SAME'``, the string ``'VALID'``, or a sequence\n of ``n`` ``(low, high)`` integer pairs that give the padding to apply before\n and after each spatial dimension (default: ``'VALID'``).\n Returns:\n The maximum for each window slice.\n \"\"\"\n y = pool(inputs, -jnp.inf, lax.max, window_shape, strides, padding)\n return y\n\n\ndef min_pool(inputs, window_shape, strides=None, padding='VALID'):\n \"\"\"Pools the input by taking the minimum of a window slice.\n\n Args:\n inputs: Input data with dimensions (batch, window dims..., features).\n window_shape: A shape tuple defining the window to reduce over.\n strides: A sequence of ``n`` integers, representing the inter-window strides\n (default: ``(1, ..., 1)``).\n padding: Either the string ``'SAME'``, the string ``'VALID'``, or a sequence of\n ``n`` ``(low, high)`` integer pairs that give the padding to apply before and\n after each spatial dimension (default: ``'VALID'``).\n\n Returns:\n The minimum for each window slice.\n \"\"\"\n return pool(inputs, jnp.inf, lax.min, window_shape, strides, padding)\n" }, { "path": "flax/linen/recurrent.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Recurrent neural network modules.\n\nTHe RNNCell modules can be scanned using lifted transforms. For more information\nsee: https://flax.readthedocs.io/en/latest/developer_notes/lift.html.\n\"\"\"\n\nfrom functools import partial # pylint: disable=g-importing-member\nfrom typing import (\n Any,\n TypeVar,\n)\nfrom collections.abc import Callable, Mapping, Sequence\n\nimport jax\nimport numpy as np\nfrom absl import logging\nfrom jax import numpy as jnp\nfrom jax import random\nfrom typing_extensions import Protocol\n\nfrom flax.core.frozen_dict import FrozenDict\nfrom flax.core.scope import CollectionFilter, PRNGSequenceFilter\nfrom flax.linen import initializers, transforms\nfrom flax.linen.activation import sigmoid, tanh\nfrom flax.linen.dtypes import promote_dtype\nfrom flax.linen.linear import Conv, Dense, default_kernel_init\nfrom flax.linen.module import Module, compact, nowrap\nfrom flax.typing import (\n Array,\n PRNGKey,\n Dtype,\n InOutScanAxis,\n Initializer,\n PrecisionLike,\n)\n\nA = TypeVar('A')\nCarry = Any\nCarryHistory = Any\nOutput = Any\n\n\nclass RNNCellBase(Module):\n \"\"\"RNN cell base class.\"\"\"\n\n @nowrap\n def initialize_carry(\n self, rng: PRNGKey, input_shape: tuple[int, ...]\n ) -> Carry:\n \"\"\"Initialize the RNN cell carry.\n\n Args:\n rng: random number generator passed to the init_fn.\n input_shape: a tuple providing the shape of the input to the cell.\n\n Returns:\n An initialized carry for the given RNN cell.\n \"\"\"\n raise NotImplementedError\n\n @property\n def num_feature_axes(self) -> int:\n \"\"\"Returns the number of feature axes of the RNN cell.\"\"\"\n raise NotImplementedError\n\n\nclass LSTMCell(RNNCellBase):\n r\"\"\"LSTM cell.\n\n The mathematical definition of the cell is as follows\n\n .. math::\n \\begin{array}{ll}\n i = \\sigma(W_{ii} x + W_{hi} h + b_{hi}) \\\\\n f = \\sigma(W_{if} x + W_{hf} h + b_{hf}) \\\\\n g = \\tanh(W_{ig} x + W_{hg} h + b_{hg}) \\\\\n o = \\sigma(W_{io} x + W_{ho} h + b_{ho}) \\\\\n c' = f * c + i * g \\\\\n h' = o * \\tanh(c') \\\\\n \\end{array}\n\n where x is the input, h is the output of the previous time step, and c is\n the memory.\n\n Example usage::\n\n >>> import flax.linen as nn\n >>> import jax, jax.numpy as jnp\n\n >>> x = jax.random.normal(jax.random.key(0), (2, 3))\n >>> layer = nn.LSTMCell(features=4)\n >>> carry = layer.initialize_carry(jax.random.key(1), x.shape)\n >>> variables = layer.init(jax.random.key(2), carry, x)\n >>> new_carry, out = layer.apply(variables, carry, x)\n\n Attributes:\n features: number of output features.\n gate_fn: activation function used for gates (default: sigmoid).\n activation_fn: activation function used for output and memory update\n (default: tanh).\n kernel_init: initializer function for the kernels that transform\n the input (default: lecun_normal).\n recurrent_kernel_init: initializer function for the kernels that transform\n the hidden state (default: initializers.orthogonal()).\n bias_init: initializer for the bias parameters (default: initializers.zeros_init())\n dtype: the dtype of the computation (default: infer from inputs and params).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n \"\"\"\n\n features: int\n gate_fn: Callable[..., Any] = sigmoid\n activation_fn: Callable[..., Any] = tanh\n kernel_init: Initializer = default_kernel_init\n recurrent_kernel_init: Initializer = initializers.orthogonal()\n bias_init: Initializer = initializers.zeros_init()\n dtype: Dtype | None = None\n param_dtype: Dtype = jnp.float32\n carry_init: Initializer = initializers.zeros_init()\n\n @compact\n def __call__(self, carry, inputs):\n r\"\"\"A long short-term memory (LSTM) cell.\n\n Args:\n carry: the hidden state of the LSTM cell,\n initialized using ``LSTMCell.initialize_carry``.\n inputs: an ndarray with the input for the current time step.\n All dimensions except the final are considered batch dimensions.\n\n Returns:\n A tuple with the new carry and the output.\n \"\"\"\n c, h = carry\n hidden_features = h.shape[-1]\n # input and recurrent layers are summed so only one needs a bias.\n dense_h = partial(\n Dense,\n features=hidden_features,\n use_bias=True,\n kernel_init=self.recurrent_kernel_init,\n bias_init=self.bias_init,\n dtype=self.dtype,\n param_dtype=self.param_dtype,\n )\n dense_i = partial(\n Dense,\n features=hidden_features,\n use_bias=False,\n kernel_init=self.kernel_init,\n dtype=self.dtype,\n param_dtype=self.param_dtype,\n )\n i = self.gate_fn(dense_i(name='ii')(inputs) + dense_h(name='hi')(h))\n f = self.gate_fn(dense_i(name='if')(inputs) + dense_h(name='hf')(h))\n g = self.activation_fn(dense_i(name='ig')(inputs) + dense_h(name='hg')(h))\n o = self.gate_fn(dense_i(name='io')(inputs) + dense_h(name='ho')(h))\n new_c = f * c + i * g\n new_h = o * self.activation_fn(new_c)\n return (new_c, new_h), new_h\n\n @nowrap\n def initialize_carry(\n self, rng: PRNGKey, input_shape: tuple[int, ...]\n ) -> tuple[Array, Array]:\n \"\"\"Initialize the RNN cell carry.\n\n Args:\n rng: random number generator passed to the init_fn.\n input_shape: a tuple providing the shape of the input to the cell.\n Returns:\n An initialized carry for the given RNN cell.\n \"\"\"\n batch_dims = input_shape[:-1]\n key1, key2 = random.split(rng)\n mem_shape = batch_dims + (self.features,)\n c = self.carry_init(key1, mem_shape, self.param_dtype)\n h = self.carry_init(key2, mem_shape, self.param_dtype)\n return (c, h)\n\n @property\n def num_feature_axes(self) -> int:\n return 1\n\n\nclass DenseParams(Module):\n \"\"\"Dummy module for creating parameters matching ``flax.linen.Dense``.\"\"\"\n\n features: int\n use_bias: bool = True\n param_dtype: Dtype = jnp.float32\n precision: PrecisionLike = None\n kernel_init: Initializer = default_kernel_init\n bias_init: Initializer = initializers.zeros_init()\n\n @compact\n def __call__(self, inputs: Array) -> tuple[Array, Array | None]:\n k = self.param(\n 'kernel',\n self.kernel_init,\n (inputs.shape[-1], self.features),\n self.param_dtype,\n )\n if self.use_bias:\n b = self.param('bias', self.bias_init, (self.features,), self.param_dtype)\n else:\n b = None\n return k, b\n\n\nclass OptimizedLSTMCell(RNNCellBase):\n r\"\"\"More efficient LSTM Cell that concatenates state components before matmul.\n\n The parameters are compatible with ``LSTMCell``. Note that this cell is often\n faster than ``LSTMCell`` as long as the hidden size is roughly <= 2048 units.\n\n The mathematical definition of the cell is the same as ``LSTMCell`` and as\n follows\n\n .. math::\n\n \\begin{array}{ll}\n i = \\sigma(W_{ii} x + W_{hi} h + b_{hi}) \\\\\n f = \\sigma(W_{if} x + W_{hf} h + b_{hf}) \\\\\n g = \\tanh(W_{ig} x + W_{hg} h + b_{hg}) \\\\\n o = \\sigma(W_{io} x + W_{ho} h + b_{ho}) \\\\\n c' = f * c + i * g \\\\\n h' = o * \\tanh(c') \\\\\n \\end{array}\n\n where x is the input, h is the output of the previous time step, and c is\n the memory.\n\n Example usage::\n\n >>> import flax.linen as nn\n >>> import jax, jax.numpy as jnp\n\n >>> x = jax.random.normal(jax.random.key(0), (2, 3))\n >>> layer = nn.OptimizedLSTMCell(features=4)\n >>> carry = layer.initialize_carry(jax.random.key(1), x.shape)\n >>> variables = layer.init(jax.random.key(2), carry, x)\n >>> new_carry, out = layer.apply(variables, carry, x)\n\n Attributes:\n gate_fn: activation function used for gates (default: sigmoid).\n activation_fn: activation function used for output and memory update\n (default: tanh).\n kernel_init: initializer function for the kernels that transform\n the input (default: lecun_normal).\n recurrent_kernel_init: initializer function for the kernels that transform\n the hidden state (default: initializers.orthogonal()).\n bias_init: initializer for the bias parameters (default: initializers.zeros_init()).\n dtype: the dtype of the computation (default: infer from inputs and params).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n \"\"\"\n\n features: int\n gate_fn: Callable[..., Any] = sigmoid\n activation_fn: Callable[..., Any] = tanh\n kernel_init: Initializer = default_kernel_init\n recurrent_kernel_init: Initializer = initializers.orthogonal()\n bias_init: Initializer = initializers.zeros_init()\n dtype: Dtype | None = None\n param_dtype: Dtype = jnp.float32\n carry_init: Initializer = initializers.zeros_init()\n\n @compact\n def __call__(\n self, carry: tuple[Array, Array], inputs: Array\n ) -> tuple[tuple[Array, Array], Array]:\n r\"\"\"An optimized long short-term memory (LSTM) cell.\n\n Args:\n carry: the hidden state of the LSTM cell, initialized using\n ``LSTMCell.initialize_carry``.\n inputs: an ndarray with the input for the current time step. All\n dimensions except the final are considered batch dimensions.\n\n Returns:\n A tuple with the new carry and the output.\n \"\"\"\n c, h = carry\n hidden_features = h.shape[-1]\n\n def _concat_dense(\n inputs: Array,\n params: Mapping[str, tuple[Array, Array | None]],\n use_bias: bool = True,\n ) -> dict[str, Array]:\n # Concatenates the individual kernels and biases, given in params, into a\n # single kernel and single bias for efficiency before applying them using\n # dot_general.\n kernels = [kernel for kernel, _ in params.values()]\n kernel = jnp.concatenate(kernels, axis=-1)\n if use_bias:\n biases = []\n for _, bias in params.values():\n if bias is None:\n raise ValueError('bias is None but use_bias is True.')\n biases.append(bias)\n bias = jnp.concatenate(biases, axis=-1)\n else:\n bias = None\n inputs, kernel, bias = promote_dtype(\n inputs, kernel, bias, dtype=self.dtype\n )\n y = jnp.dot(inputs, kernel)\n if use_bias:\n # This assert is here since mypy can't infer that bias cannot be None\n assert bias is not None\n y += jnp.reshape(bias, (1,) * (y.ndim - 1) + (-1,))\n\n # Split the result back into individual (i, f, g, o) outputs.\n split_indices = np.cumsum([kernel.shape[-1] for kernel in kernels[:-1]])\n ys = jnp.split(y, split_indices, axis=-1)\n return dict(zip(params.keys(), ys))\n\n # Create params with the same names/shapes as `LSTMCell` for compatibility.\n dense_params_h = {}\n dense_params_i = {}\n for component in ['i', 'f', 'g', 'o']:\n dense_params_i[component] = DenseParams(\n features=hidden_features,\n use_bias=False,\n param_dtype=self.param_dtype,\n kernel_init=self.kernel_init,\n bias_init=self.bias_init,\n name=f'i{component}', # type: ignore[call-arg]\n )(inputs)\n dense_params_h[component] = DenseParams(\n features=hidden_features,\n use_bias=True,\n param_dtype=self.param_dtype,\n kernel_init=self.recurrent_kernel_init,\n bias_init=self.bias_init,\n name=f'h{component}', # type: ignore[call-arg]\n )(h)\n dense_h = _concat_dense(h, dense_params_h, use_bias=True)\n dense_i = _concat_dense(inputs, dense_params_i, use_bias=False)\n\n i = self.gate_fn(dense_h['i'] + dense_i['i'])\n f = self.gate_fn(dense_h['f'] + dense_i['f'])\n g = self.activation_fn(dense_h['g'] + dense_i['g'])\n o = self.gate_fn(dense_h['o'] + dense_i['o'])\n\n new_c = f * c + i * g\n new_h = o * self.activation_fn(new_c)\n return (new_c, new_h), new_h\n\n @nowrap\n def initialize_carry(\n self, rng: PRNGKey, input_shape: tuple[int, ...]\n ) -> tuple[Array, Array]:\n \"\"\"Initialize the RNN cell carry.\n\n Args:\n rng: random number generator passed to the init_fn.\n input_shape: a tuple providing the shape of the input to the cell.\n\n Returns:\n An initialized carry for the given RNN cell.\n \"\"\"\n batch_dims = input_shape[:-1]\n key1, key2 = random.split(rng)\n mem_shape = batch_dims + (self.features,)\n c = self.carry_init(key1, mem_shape, self.param_dtype)\n h = self.carry_init(key2, mem_shape, self.param_dtype)\n return c, h\n\n @property\n def num_feature_axes(self) -> int:\n return 1\n\n\nclass SimpleCell(RNNCellBase):\n r\"\"\"Simple cell.\n\n The mathematical definition of the cell is as follows\n\n .. math::\n\n \\begin{array}{ll}\n h' = \\tanh(W_i x + b_i + W_h h)\n \\end{array}\n\n where x is the input and h is the output of the previous time step.\n\n If `residual` is `True`,\n\n .. math::\n\n \\begin{array}{ll}\n h' = \\tanh(W_i x + b_i + W_h h + h)\n \\end{array}\n\n Example usage::\n\n >>> import flax.linen as nn\n >>> import jax, jax.numpy as jnp\n\n >>> x = jax.random.normal(jax.random.key(0), (2, 3))\n >>> layer = nn.SimpleCell(features=4)\n >>> carry = layer.initialize_carry(jax.random.key(1), x.shape)\n >>> variables = layer.init(jax.random.key(2), carry, x)\n >>> new_carry, out = layer.apply(variables, carry, x)\n\n Attributes:\n features: number of output features.\n activation_fn: activation function used for output and memory update\n (default: tanh).\n kernel_init: initializer function for the kernels that transform\n the input (default: lecun_normal).\n recurrent_kernel_init: initializer function for the kernels that transform\n the hidden state (default: initializers.orthogonal()).\n bias_init: initializer for the bias parameters (default: initializers.zeros_init())\n dtype: the dtype of the computation (default: None).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n residual: pre-activation residual connection (https://arxiv.org/abs/1801.06105).\n \"\"\"\n\n features: int\n activation_fn: Callable[..., Any] = tanh\n kernel_init: Initializer = default_kernel_init\n recurrent_kernel_init: Initializer = initializers.orthogonal()\n bias_init: Initializer = initializers.zeros_init()\n dtype: Dtype | None = None\n param_dtype: Dtype = jnp.float32\n carry_init: Initializer = initializers.zeros_init()\n residual: bool = False\n\n @compact\n def __call__(self, carry, inputs):\n \"\"\"Simple cell.\n\n Args:\n carry: the hidden state of the Simple cell,\n initialized using ``SimpleCell.initialize_carry``.\n inputs: an ndarray with the input for the current time step.\n All dimensions except the final are considered batch dimensions.\n\n Returns:\n A tuple with the new carry and the output.\n \"\"\"\n hidden_features = carry.shape[-1]\n # input and recurrent layers are summed so only one needs a bias.\n dense_h = partial(\n Dense,\n features=hidden_features,\n use_bias=False,\n dtype=self.dtype,\n param_dtype=self.param_dtype,\n kernel_init=self.recurrent_kernel_init,\n )\n dense_i = partial(\n Dense,\n features=hidden_features,\n use_bias=True,\n dtype=self.dtype,\n param_dtype=self.param_dtype,\n kernel_init=self.kernel_init,\n bias_init=self.bias_init,\n )\n new_carry = dense_i(name='i')(inputs) + dense_h(name='h')(carry)\n if self.residual:\n new_carry += carry\n new_carry = self.activation_fn(new_carry)\n return new_carry, new_carry\n\n @nowrap\n def initialize_carry(self, rng: PRNGKey, input_shape: tuple[int, ...]):\n \"\"\"Initialize the RNN cell carry.\n\n Args:\n rng: random number generator passed to the init_fn.\n input_shape: a tuple providing the shape of the input to the cell.\n\n Returns:\n An initialized carry for the given RNN cell.\n \"\"\"\n batch_dims = input_shape[:-1]\n mem_shape = batch_dims + (self.features,)\n return self.carry_init(rng, mem_shape, self.param_dtype)\n\n @property\n def num_feature_axes(self) -> int:\n return 1\n\n\nclass GRUCell(RNNCellBase):\n r\"\"\"GRU cell.\n\n The mathematical definition of the cell is as follows\n\n .. math::\n\n \\begin{array}{ll}\n r = \\sigma(W_{ir} x + b_{ir} + W_{hr} h) \\\\\n z = \\sigma(W_{iz} x + b_{iz} + W_{hz} h) \\\\\n n = \\tanh(W_{in} x + b_{in} + r * (W_{hn} h + b_{hn})) \\\\\n h' = (1 - z) * n + z * h \\\\\n \\end{array}\n\n where x is the input and h is the output of the previous time step.\n\n Example usage::\n\n >>> import flax.linen as nn\n >>> import jax, jax.numpy as jnp\n\n >>> x = jax.random.normal(jax.random.key(0), (2, 3))\n >>> layer = nn.GRUCell(features=4)\n >>> carry = layer.initialize_carry(jax.random.key(1), x.shape)\n >>> variables = layer.init(jax.random.key(2), carry, x)\n >>> new_carry, out = layer.apply(variables, carry, x)\n\n Attributes:\n features: number of output features.\n gate_fn: activation function used for gates (default: sigmoid).\n activation_fn: activation function used for output and memory update\n (default: tanh).\n kernel_init: initializer function for the kernels that transform\n the input (default: lecun_normal).\n recurrent_kernel_init: initializer function for the kernels that transform\n the hidden state (default: initializers.orthogonal()).\n bias_init: initializer for the bias parameters (default: initializers.zeros_init())\n dtype: the dtype of the computation (default: None).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n \"\"\"\n\n features: int\n gate_fn: Callable[..., Any] = sigmoid\n activation_fn: Callable[..., Any] = tanh\n kernel_init: Initializer = default_kernel_init\n recurrent_kernel_init: Initializer = initializers.orthogonal()\n bias_init: Initializer = initializers.zeros_init()\n dtype: Dtype | None = None\n param_dtype: Dtype = jnp.float32\n carry_init: Initializer = initializers.zeros_init()\n\n @compact\n def __call__(self, carry, inputs):\n \"\"\"Gated recurrent unit (GRU) cell.\n\n Args:\n carry: the hidden state of the GRU cell,\n initialized using ``GRUCell.initialize_carry``.\n inputs: an ndarray with the input for the current time step.\n All dimensions except the final are considered batch dimensions.\n\n Returns:\n A tuple with the new carry and the output.\n \"\"\"\n h = carry\n hidden_features = h.shape[-1]\n # input and recurrent layers are summed so only one needs a bias.\n dense_h = partial(\n Dense,\n features=hidden_features,\n use_bias=False,\n dtype=self.dtype,\n param_dtype=self.param_dtype,\n kernel_init=self.recurrent_kernel_init,\n bias_init=self.bias_init,\n )\n dense_i = partial(\n Dense,\n features=hidden_features,\n use_bias=True,\n dtype=self.dtype,\n param_dtype=self.param_dtype,\n kernel_init=self.kernel_init,\n bias_init=self.bias_init,\n )\n r = self.gate_fn(dense_i(name='ir')(inputs) + dense_h(name='hr')(h))\n z = self.gate_fn(dense_i(name='iz')(inputs) + dense_h(name='hz')(h))\n # add bias because the linear transformations aren't directly summed.\n n = self.activation_fn(\n dense_i(name='in')(inputs) + r * dense_h(name='hn', use_bias=True)(h)\n )\n new_h = (1.0 - z) * n + z * h\n return new_h, new_h\n\n @nowrap\n def initialize_carry(self, rng: PRNGKey, input_shape: tuple[int, ...]):\n \"\"\"Initialize the RNN cell carry.\n\n Args:\n rng: random number generator passed to the init_fn.\n input_shape: a tuple providing the shape of the input to the cell.\n\n Returns:\n An initialized carry for the given RNN cell.\n \"\"\"\n batch_dims = input_shape[:-1]\n mem_shape = batch_dims + (self.features,)\n return self.carry_init(rng, mem_shape, self.param_dtype)\n\n @property\n def num_feature_axes(self) -> int:\n return 1\n\n\nclass MGUCell(RNNCellBase):\n r\"\"\"MGU cell (https://arxiv.org/pdf/1603.09420.pdf).\n\n The mathematical definition of the cell is as follows\n\n .. math::\n\n \\begin{array}{ll}\n f = \\sigma(W_{if} x + b_{if} + W_{hf} h) \\\\\n n = \\tanh(W_{in} x + b_{in} + f * (W_{hn} h + b_{hn})) \\\\\n h' = (1 - f) * n + f * h \\\\\n \\end{array}\n\n where x is the input and h is the output of the previous time step.\n\n If ``reset_gate`` is false, the above becomes\n\n .. math::\n\n \\begin{array}{ll}\n f = \\sigma(W_{if} x + b_{if} + W_{hf} h) \\\\\n n = \\tanh(W_{in} x + b_{in} + W_{hn} h) \\\\\n h' = (1 - f) * n + f * h \\\\\n \\end{array}\n\n Example usage::\n\n >>> import flax.linen as nn\n >>> import jax, jax.numpy as jnp\n\n >>> x = jax.random.normal(jax.random.key(0), (2, 3))\n >>> layer = nn.MGUCell(features=4)\n >>> carry = layer.initialize_carry(jax.random.key(1), x.shape)\n >>> variables = layer.init(jax.random.key(2), carry, x)\n >>> new_carry, out = layer.apply(variables, carry, x)\n\n Attributes:\n features: number of output features.\n gate_fn: activation function used for gates (default: sigmoid).\n activation_fn: activation function used for output and memory update\n (default: tanh).\n kernel_init: initializer function for the kernels that transform\n the input (default: lecun_normal).\n recurrent_kernel_init: initializer function for the kernels that transform\n the hidden state (default: initializers.orthogonal()).\n forget_bias_init: initializer for the bias parameters of the forget gate.\n The default is set to initializers.ones_init() because this prevents\n vanishing gradients. See https://proceedings.mlr.press/v37/jozefowicz15.pdf,\n section 2.2 for more details.\n activation_bias_init: initializer for the bias parameters of the activation\n output (default: initializers.zeros_init()).\n dtype: the dtype of the computation (default: None).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n reset_gate: flag for applying reset gating.\n \"\"\"\n\n features: int\n gate_fn: Callable[..., Any] = sigmoid\n activation_fn: Callable[..., Any] = tanh\n kernel_init: Initializer = default_kernel_init\n recurrent_kernel_init: Initializer = initializers.orthogonal()\n forget_bias_init: Initializer = initializers.ones_init()\n activation_bias_init: Initializer = initializers.zeros_init()\n dtype: Dtype | None = None\n param_dtype: Dtype = jnp.float32\n carry_init: Initializer = initializers.zeros_init()\n reset_gate: bool = True\n\n @compact\n def __call__(self, carry, inputs):\n \"\"\"Minimal gated unit (MGU) cell.\n\n Args:\n carry: the hidden state of the MGU cell,\n initialized using ``MGUCell.initialize_carry``.\n inputs: an ndarray with the input for the current time step.\n All dimensions except the final are considered batch dimensions.\n\n Returns:\n A tuple with the new carry and the output.\n \"\"\"\n h = carry\n hidden_features = h.shape[-1]\n # input and recurrent layers are summed so only one needs a bias.\n dense_h = partial(\n Dense,\n features=hidden_features,\n use_bias=False,\n dtype=self.dtype,\n param_dtype=self.param_dtype,\n kernel_init=self.recurrent_kernel_init,\n bias_init=self.activation_bias_init,\n )\n dense_i = partial(\n Dense,\n features=hidden_features,\n use_bias=True,\n dtype=self.dtype,\n param_dtype=self.param_dtype,\n kernel_init=self.kernel_init,\n )\n f = self.gate_fn(\n dense_i(name='if', bias_init=self.forget_bias_init)(inputs)\n + dense_h(name='hf')(h)\n )\n # add bias when the linear transformations aren't directly summed.\n x = dense_h(name=\"hn\", use_bias=self.reset_gate)(h)\n if self.reset_gate:\n x *= f\n n = self.activation_fn(\n dense_i(name=\"in\", bias_init=self.activation_bias_init)(inputs) + x\n )\n new_h = (1.0 - f) * n + f * h\n return new_h, new_h\n\n @nowrap\n def initialize_carry(self, rng: PRNGKey, input_shape: tuple[int, ...]):\n \"\"\"Initialize the RNN cell carry.\n\n Args:\n rng: random number generator passed to the init_fn.\n input_shape: a tuple providing the shape of the input to the cell.\n\n Returns:\n An initialized carry for the given RNN cell.\n \"\"\"\n batch_dims = input_shape[:-1]\n mem_shape = batch_dims + (self.features,)\n return self.carry_init(rng, mem_shape, self.param_dtype)\n\n @property\n def num_feature_axes(self) -> int:\n return 1\n\n\nclass ConvLSTMCell(RNNCellBase):\n r\"\"\"A convolutional LSTM cell.\n\n The implementation is based on xingjian2015convolutional.\n Given x_t and the previous state (h_{t-1}, c_{t-1})\n the core computes\n\n .. math::\n\n \\begin{array}{ll}\n i_t = \\sigma(W_{ii} * x_t + W_{hi} * h_{t-1} + b_i) \\\\\n f_t = \\sigma(W_{if} * x_t + W_{hf} * h_{t-1} + b_f) \\\\\n g_t = \\tanh(W_{ig} * x_t + W_{hg} * h_{t-1} + b_g) \\\\\n o_t = \\sigma(W_{io} * x_t + W_{ho} * h_{t-1} + b_o) \\\\\n c_t = f_t c_{t-1} + i_t g_t \\\\\n h_t = o_t \\tanh(c_t)\n \\end{array}\n\n where * denotes the convolution operator;\n i_t, f_t, o_t are input, forget and output gate activations,\n and g_t is a vector of cell updates.\n\n .. note::\n Forget gate initialization:\n Following jozefowicz2015empirical we add 1.0 to b_f\n after initialization in order to reduce the scale of forgetting in\n the beginning of the training.\n\n Example usage::\n\n >>> import flax.linen as nn\n >>> import jax, jax.numpy as jnp\n\n >>> x = jax.random.normal(jax.random.key(0), (3, 5, 5))\n >>> layer = nn.ConvLSTMCell(features=4, kernel_size=(2, 2))\n >>> carry = layer.initialize_carry(jax.random.key(1), x.shape)\n >>> variables = layer.init(jax.random.key(2), carry, x)\n >>> new_carry, out = layer.apply(variables, carry, x)\n\n Attributes:\n features: number of convolution filters.\n kernel_size: shape of the convolutional kernel.\n strides: a sequence of ``n`` integers, representing the inter-window\n strides.\n padding: either the string ``'SAME'``, the string ``'VALID'``, or a sequence\n of ``n`` ``(low, high)`` integer pairs that give the padding to apply before\n and after each spatial dimension.\n bias: whether to add a bias to the output (default: True).\n dtype: the dtype of the computation (default: None).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n \"\"\"\n\n features: int\n kernel_size: Sequence[int]\n strides: Sequence[int] | None = None\n padding: str | Sequence[tuple[int, int]] = 'SAME'\n use_bias: bool = True\n dtype: Dtype | None = None\n param_dtype: Dtype = jnp.float32\n carry_init: Initializer = initializers.zeros_init()\n\n @compact\n def __call__(self, carry, inputs):\n \"\"\"Constructs a convolutional LSTM.\n\n Args:\n carry: the hidden state of the Conv2DLSTM cell,\n initialized using ``Conv2DLSTM.initialize_carry``.\n inputs: input data with dimensions (batch, spatial_dims..., features).\n Returns:\n A tuple with the new carry and the output.\n \"\"\"\n c, h = carry\n input_to_hidden = partial(\n Conv,\n features=4 * self.features,\n kernel_size=self.kernel_size,\n strides=self.strides,\n padding=self.padding,\n use_bias=self.use_bias,\n dtype=self.dtype,\n param_dtype=self.param_dtype,\n name='ih',\n )\n\n hidden_to_hidden = partial(\n Conv,\n features=4 * self.features,\n kernel_size=self.kernel_size,\n strides=self.strides,\n padding=self.padding,\n use_bias=self.use_bias,\n dtype=self.dtype,\n param_dtype=self.param_dtype,\n name='hh',\n )\n\n gates = input_to_hidden()(inputs) + hidden_to_hidden()(h)\n i, g, f, o = jnp.split(gates, indices_or_sections=4, axis=-1)\n\n f = sigmoid(f + 1)\n new_c = f * c + sigmoid(i) * jnp.tanh(g)\n new_h = sigmoid(o) * jnp.tanh(new_c)\n return (new_c, new_h), new_h\n\n @nowrap\n def initialize_carry(self, rng: PRNGKey, input_shape: tuple[int, ...]):\n \"\"\"Initialize the RNN cell carry.\n\n Args:\n rng: random number generator passed to the init_fn.\n input_shape: a tuple providing the shape of the input to the cell.\n\n Returns:\n An initialized carry for the given RNN cell.\n \"\"\"\n # (*batch_dims, *signal_dims, features)\n signal_dims = input_shape[-self.num_feature_axes : -1]\n batch_dims = input_shape[: -self.num_feature_axes]\n key1, key2 = random.split(rng)\n mem_shape = batch_dims + signal_dims + (self.features,)\n c = self.carry_init(key1, mem_shape, self.param_dtype)\n h = self.carry_init(key2, mem_shape, self.param_dtype)\n return c, h\n\n @property\n def num_feature_axes(self) -> int:\n return len(self.kernel_size) + 1\n\n\nclass RNN(Module):\n \"\"\"The ``RNN`` module takes any :class:`RNNCellBase` instance and applies it over a sequence\n\n using :func:`flax.linen.scan`.\n\n Example::\n\n >>> import jax.numpy as jnp\n >>> import jax\n >>> import flax.linen as nn\n\n >>> x = jnp.ones((10, 50, 32)) # (batch, time, features)\n >>> lstm = nn.RNN(nn.LSTMCell(64))\n >>> variables = lstm.init(jax.random.key(0), x)\n >>> y = lstm.apply(variables, x)\n >>> y.shape # (batch, time, cell_size)\n (10, 50, 64)\n\n As shown above, RNN uses the ``cell_size`` argument to set the ``size``\n argument for the cell's\n ``initialize_carry`` method, in practice this is typically the number of\n hidden units you want\n for the cell. However, this may vary depending on the cell you are using, for\n example the\n :class:`ConvLSTMCell` requires a ``size`` argument of the form\n ``(kernel_height, kernel_width, features)``::\n\n >>> x = jnp.ones((10, 50, 32, 32, 3)) # (batch, time, height, width, features)\n >>> conv_lstm = nn.RNN(nn.ConvLSTMCell(64, kernel_size=(3, 3)))\n >>> y, variables = conv_lstm.init_with_output(jax.random.key(0), x)\n >>> y.shape # (batch, time, height, width, features)\n (10, 50, 32, 32, 64)\n\n By default RNN expect the time dimension after the batch dimension (``(*batch,\n time, *features)``),\n if you set ``time_major=True`` RNN will instead expect the time dimension to be\n at the beginning\n (``(time, *batch, *features)``)::\n\n >>> x = jnp.ones((50, 10, 32)) # (time, batch, features)\n >>> lstm = nn.RNN(nn.LSTMCell(64), time_major=True)\n >>> variables = lstm.init(jax.random.key(0), x)\n >>> y = lstm.apply(variables, x)\n >>> y.shape # (time, batch, cell_size)\n (50, 10, 64)\n\n The output is an array of shape ``(*batch, time, *cell_size)`` by default\n (typically), however\n if you set ``return_carry=True`` it will instead return a tuple of the final\n carry and the output::\n\n >>> x = jnp.ones((10, 50, 32)) # (batch, time, features)\n >>> lstm = nn.RNN(nn.LSTMCell(64), return_carry=True)\n >>> variables = lstm.init(jax.random.key(0), x)\n >>> carry, y = lstm.apply(variables, x)\n >>> jax.tree_util.tree_map(jnp.shape, carry) # ((batch, cell_size), (batch, cell_size))\n ((10, 64), (10, 64))\n >>> y.shape # (batch, time, cell_size)\n (10, 50, 64)\n\n To support variable length sequences, you can pass a ``seq_lengths`` which is\n an integer\n array of shape ``(*batch)`` where each element is the length of the sequence\n in the batch.\n For example::\n\n >>> seq_lengths = jnp.array([3, 2, 5])\n\n The output elements corresponding to padding elements are NOT zeroed out. If\n ``return_carry``\n is set to ``True`` the carry will be the state of the last valid element of\n each sequence.\n\n RNN also accepts some of the arguments of :func:`flax.linen.scan`, by default\n they are set to\n work with cells like :class:`LSTMCell` and :class:`GRUCell` but they can be\n overridden as needed.\n Overriding default values to scan looks like this::\n\n >>> lstm = nn.RNN(\n ... nn.LSTMCell(64),\n ... unroll=1, variable_axes={}, variable_broadcast='params',\n ... variable_carry=False, split_rngs={'params': False})\n\n Attributes:\n cell: an instance of :class:`RNNCellBase`.\n time_major: if ``time_major=False`` (default) it will expect inputs with\n shape ``(*batch, time, *features)``, else it will expect inputs with shape\n ``(time, *batch, *features)``.\n return_carry: if ``return_carry=False`` (default) only the output sequence\n is returned, else it will return a tuple of the final carry and the output\n sequence.\n reverse: if ``reverse=False`` (default) the sequence is processed from left\n to right and returned in the original order, else it will be processed\n from right to left, and returned in reverse order. If ``seq_lengths`` is\n passed, padding will always remain at the end of the sequence.\n keep_order: if ``keep_order=True``, when ``reverse=True`` the output will be\n reversed back to the original order after processing, this is useful to\n align sequences in bidirectional RNNs. If ``keep_order=False`` (default),\n the output will remain in the order specified by ``reverse``.\n unroll: how many scan iterations to unroll within a single iteration of a\n loop, defaults to 1. This argument will be passed to ``nn.scan``.\n variable_axes: a dictionary mapping each collection to either an integer\n ``i`` (meaning we scan over dimension ``i``) or ``None`` (replicate rather\n than scan). This argument is forwarded to ``nn.scan``.\n variable_broadcast: Specifies the broadcasted variable collections. A\n broadcasted variable should not depend on any computation that cannot be\n lifted out of the loop. This is typically used to define shared parameters\n inside the fn. This argument is forwarded to ``nn.scan``.\n variable_carry: Specifies the variable collections that are carried through\n the loop. Mutations to these variables are carried to the next iteration\n and will be preserved when the scan finishes. This argument is forwarded\n to ``nn.scan``.\n split_rngs: a mapping from PRNGSequenceFilter to bool specifying whether a\n collection's PRNG key should be split such that its values are different\n at each step, or replicated such that its values remain the same at each\n step. This argument is forwarded to ``nn.scan``.\n \"\"\"\n\n cell: RNNCellBase\n time_major: bool = False\n return_carry: bool = False\n reverse: bool = False\n keep_order: bool = False\n unroll: int = 1\n variable_axes: Mapping[\n CollectionFilter, InOutScanAxis\n ] = FrozenDict()\n variable_broadcast: CollectionFilter = 'params'\n variable_carry: CollectionFilter = False\n split_rngs: Mapping[PRNGSequenceFilter, bool] = FrozenDict(\n {'params': False}\n )\n\n def __call__(\n self,\n inputs: jax.Array,\n *,\n initial_carry: Carry | None = None,\n init_key: PRNGKey | None = None,\n seq_lengths: Array | None = None,\n return_carry: bool | None = None,\n time_major: bool | None = None,\n reverse: bool | None = None,\n keep_order: bool | None = None,\n ) -> Output | tuple[Carry, Output]:\n \"\"\"\n Applies the RNN to the inputs.\n\n ``__call__`` allows you to optionally override some attributes like ``return_carry``\n and ``time_major`` defined in the constructor.\n\n Arguments:\n inputs: the input sequence.\n initial_carry: the initial carry, if not provided it will be initialized\n using the cell's :meth:`RNNCellBase.initialize_carry` method.\n init_key: a PRNG key used to initialize the carry, if not provided\n ``jax.random.key(0)`` will be used. Most cells will ignore this\n argument.\n seq_lengths: an optional integer array of shape ``(*batch)`` indicating\n the length of each sequence, elements whose index in the time dimension\n is greater than the corresponding length will be considered padding and\n will be ignored.\n return_carry: if ``return_carry=False`` (default) only the output sequence is returned,\n else it will return a tuple of the final carry and the output sequence.\n time_major: if ``time_major=False`` (default) it will expect inputs with shape\n ``(*batch, time, *features)``, else it will expect inputs with shape ``(time, *batch, *features)``.\n reverse: overrides the ``reverse`` attribute, if ``reverse=False`` (default) the sequence is\n processed from left to right and returned in the original order, else it will be processed\n from right to left, and returned in reverse order. If ``seq_lengths`` is passed,\n padding will always remain at the end of the sequence.\n keep_order: overrides the ``keep_order`` attribute, if ``keep_order=True``, when ``reverse=True``\n the output will be reversed back to the original order after processing, this is\n useful to align sequences in bidirectional RNNs. If ``keep_order=False`` (default),\n the output will remain in the order specified by ``reverse``.\n Returns:\n if ``return_carry=False`` (default) only the output sequence is returned,\n else it will return a tuple of the final carry and the output sequence.\n \"\"\"\n\n if return_carry is None:\n return_carry = self.return_carry\n if time_major is None:\n time_major = self.time_major\n if reverse is None:\n reverse = self.reverse\n if keep_order is None:\n keep_order = self.keep_order\n\n # Infer the number of batch dimensions from the input shape.\n # Cells like ConvLSTM have additional spatial dimensions.\n time_axis = (\n 0 if time_major else inputs.ndim - (self.cell.num_feature_axes + 1)\n )\n\n # make time_axis positive\n if time_axis < 0:\n time_axis += inputs.ndim\n\n if time_major:\n # we add +1 because we moved the time axis to the front\n batch_dims = inputs.shape[1 : -self.cell.num_feature_axes]\n else:\n batch_dims = inputs.shape[:time_axis]\n\n # maybe reverse the sequence\n if reverse:\n inputs = jax.tree_util.tree_map(\n lambda x: flip_sequences(\n x,\n seq_lengths,\n num_batch_dims=len(batch_dims),\n time_major=time_major, # type: ignore\n ),\n inputs,\n )\n\n carry: Carry\n if initial_carry is None:\n if init_key is None:\n init_key = random.key(0)\n\n input_shape = inputs.shape[:time_axis] + inputs.shape[time_axis + 1 :]\n carry = self.cell.initialize_carry(init_key, input_shape)\n else:\n carry = initial_carry\n\n slice_carry = seq_lengths is not None and return_carry\n\n def scan_fn(\n cell: RNNCellBase, carry: Carry, x: Array\n ) -> tuple[Carry, Array] | tuple[Carry, tuple[Carry, Array]]:\n carry, y = cell(carry, x)\n # When we have a segmentation mask we return the carry as an output\n # so that we can select the last carry for each sequence later.\n # This uses more memory but is faster than using jnp.where at each\n # iteration. As a small optimization do this when we really need it.\n if slice_carry:\n return carry, (carry, y)\n else:\n return carry, y\n\n scan = transforms.scan(\n scan_fn,\n in_axes=time_axis,\n out_axes=(0, time_axis) if slice_carry else time_axis,\n unroll=self.unroll,\n variable_axes=self.variable_axes,\n variable_broadcast=self.variable_broadcast,\n variable_carry=self.variable_carry,\n split_rngs=self.split_rngs,\n )\n\n scan_output = scan(self.cell, carry, inputs)\n\n # Next we select the final carry. If a segmentation mask was provided and\n # return_carry is True we slice the carry history and select the last valid\n # carry for each sequence. Otherwise we just use the last carry.\n if slice_carry:\n assert seq_lengths is not None\n _, (carries, outputs) = scan_output\n # seq_lengths[None] expands the shape of the mask to match the\n # number of dimensions of the carry.\n carry = _select_last_carry(carries, seq_lengths)\n else:\n carry, outputs = scan_output\n\n if reverse and keep_order:\n outputs = jax.tree_util.tree_map(\n lambda x: flip_sequences(\n x,\n seq_lengths,\n num_batch_dims=len(batch_dims),\n time_major=time_major, # type: ignore\n ),\n outputs,\n )\n\n if return_carry:\n return carry, outputs\n else:\n return outputs\n\n\ndef _select_last_carry(sequence: A, seq_lengths: jnp.ndarray) -> A:\n last_idx = seq_lengths - 1\n\n def _slice_array(x: jnp.ndarray):\n return x[last_idx, jnp.arange(x.shape[1])]\n\n return jax.tree_util.tree_map(_slice_array, sequence)\n\n\ndef _expand_dims_like(x, target):\n \"\"\"Expands the shape of `x` to match `target`'s shape by adding singleton dimensions.\"\"\"\n return x.reshape(list(x.shape) + [1] * (target.ndim - x.ndim))\n\n\ndef flip_sequences(\n inputs: Array,\n seq_lengths: Array | None,\n num_batch_dims: int,\n time_major: bool,\n) -> Array:\n \"\"\"Flips a sequence of inputs along the time axis.\n\n This function can be used to prepare inputs for the reverse direction of a\n bidirectional LSTM. It solves the issue that, when naively flipping multiple\n padded sequences stored in a matrix, the first elements would be padding\n values for those sequences that were padded. This function keeps the padding\n at the end, while flipping the rest of the elements.\n\n Example:\n ```python\n inputs = [[1, 0, 0],\n [2, 3, 0]\n [4, 5, 6]]\n lengths = [1, 2, 3]\n flip_sequences(inputs, lengths) = [[1, 0, 0],\n [3, 2, 0],\n [6, 5, 4]]\n ```\n\n Args:\n inputs: An array of input IDs [batch_size, seq_length].\n lengths: The length of each sequence [batch_size].\n\n Returns:\n An ndarray with the flipped inputs.\n \"\"\"\n # Compute the indices to put the inputs in flipped order as per above example.\n time_axis = 0 if time_major else num_batch_dims\n max_steps = inputs.shape[time_axis]\n\n if seq_lengths is None:\n # reverse inputs and return\n inputs = jnp.flip(inputs, axis=time_axis)\n return inputs\n\n seq_lengths = jnp.expand_dims(seq_lengths, axis=time_axis)\n\n # create indexes\n idxs = jnp.arange(max_steps - 1, -1, -1) # [max_steps]\n if time_major:\n idxs = jnp.reshape(idxs, [max_steps] + [1] * num_batch_dims)\n else:\n idxs = jnp.reshape(\n idxs, [1] * num_batch_dims + [max_steps]\n ) # [1, ..., max_steps]\n idxs = (idxs + seq_lengths) % max_steps # [*batch, max_steps]\n idxs = _expand_dims_like(\n idxs, target=inputs\n ) # [*batch, max_steps, *features]\n # Select the inputs in flipped order.\n outputs = jnp.take_along_axis(inputs, idxs, axis=time_axis)\n\n return outputs\n\n\ndef _concatenate(a: Array, b: Array) -> Array:\n \"\"\"Concatenates two arrays along the last dimension.\"\"\"\n return jnp.concatenate([a, b], axis=-1)\n\n\nclass RNNBase(Protocol):\n def __call__(\n self,\n inputs: jax.Array,\n *,\n initial_carry: Carry | None = None,\n init_key: PRNGKey | None = None,\n seq_lengths: Array | None = None,\n return_carry: bool | None = None,\n time_major: bool | None = None,\n reverse: bool | None = None,\n keep_order: bool | None = None,\n ) -> Output | tuple[Carry, Output]:\n ...\n\n\nclass Bidirectional(Module):\n \"\"\"Processes the input in both directions and merges the results.\n\n Example usage::\n\n >>> import flax.linen as nn\n >>> import jax, jax.numpy as jnp\n\n >>> layer = nn.Bidirectional(nn.RNN(nn.GRUCell(4)), nn.RNN(nn.GRUCell(4)))\n >>> x = jnp.ones((2, 3))\n >>> variables = layer.init(jax.random.key(0), x)\n >>> out = layer.apply(variables, x)\n \"\"\"\n\n forward_rnn: RNNBase\n backward_rnn: RNNBase\n merge_fn: Callable[[Array, Array], Array] = _concatenate\n time_major: bool = False\n return_carry: bool = False\n\n def __call__(\n self,\n inputs: jax.Array,\n *,\n initial_carry: Carry | None = None,\n init_key: PRNGKey | None = None,\n seq_lengths: Array | None = None,\n return_carry: bool | None = None,\n time_major: bool | None = None,\n reverse: bool | None = None,\n keep_order: bool | None = None,\n ) -> Output | tuple[Carry, Output]:\n if time_major is None:\n time_major = self.time_major\n if return_carry is None:\n return_carry = self.return_carry\n if init_key is not None:\n key_forward, key_backward = random.split(init_key)\n else:\n key_forward = key_backward = None\n if initial_carry is not None:\n initial_carry_forward, initial_carry_backward = initial_carry\n else:\n initial_carry_forward = initial_carry_backward = None\n # Throw a warning in case the user accidentally re-uses the forward RNN\n # for the backward pass and does not intend for them to share parameters.\n if self.forward_rnn is self.backward_rnn:\n logging.warning(\n 'forward_rnn and backward_rnn is the same object, so '\n 'they will share parameters.'\n )\n\n # Encode in the forward direction.\n carry_forward, outputs_forward = self.forward_rnn(\n inputs,\n initial_carry=initial_carry_forward,\n init_key=key_forward,\n seq_lengths=seq_lengths,\n return_carry=True,\n time_major=time_major,\n reverse=False,\n )\n\n carry_backward, outputs_backward = self.backward_rnn(\n inputs,\n initial_carry=initial_carry_backward,\n init_key=key_backward,\n seq_lengths=seq_lengths,\n return_carry=True,\n time_major=time_major,\n reverse=True,\n keep_order=True,\n )\n\n carry = (carry_forward, carry_backward)\n outputs = jax.tree_util.tree_map(\n self.merge_fn, outputs_forward, outputs_backward\n )\n\n if return_carry:\n return carry, outputs\n else:\n return outputs\n" }, { "path": "flax/linen/spmd.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Utilities for working with jit and partitioned models.\n\nThis module introduces ``axis_rules``, ``logical_to_mesh_axes``,\n``logical_to_mesh``, ``with_logical_constraint`` for appyling jit sharding\nconstraints in terms of \"logical named axes\" rather than jit's default mesh\naxes.\n\nAdditionally the ``LogicallyPartitioned`` metadata wrapper is defined as\nwell as the initializer function wrapper ``with_logical_partitioning``for\nintroducing logical axis metadata into a model's variables.\n\"\"\"\n\nimport collections\nimport dataclasses\nimport enum\nimport functools\nfrom typing import Any\nfrom collections.abc import Callable, Sequence\n\nimport jax\nfrom jax import lax\n\nfrom flax import struct\nfrom flax.core import meta\nfrom flax.core.spmd import (\n get_logical_axis_rules,\n)\nfrom flax.typing import (\n Array,\n LogicalNames,\n LogicalRules,\n ArrayPytree, # pylint: disable=invalid-name\n LogicalPartitionSpec, # pylint: disable=unused-import\n LogicalPartitionSpecPytree, # pylint: disable=invalid-name\n )\n\n\nclass _UnassignedAxis:\n \"\"\"Sentinel class for unassigned logical axis name.\"\"\"\n\n def __repr__(self):\n return 'UnassignedAxis'\n\n def __bool__(self):\n return False\n\n\n_unassigned_axis = _UnassignedAxis()\n\n\ndef _mesh_assignment_free(new_assignment, existing_assignments):\n \"\"\"Determines if a given mesh axis has already been assigned.\"\"\"\n new = set(jax.tree_util.tree_leaves(new_assignment))\n existing = set(jax.tree_util.tree_leaves(existing_assignments))\n new.discard(jax.sharding.PartitionSpec.UNCONSTRAINED)\n new.discard(None)\n if existing.intersection(new):\n return False\n return True\n\n\ndef _logical_to_mesh_axes(\n array_dim_names: Sequence[str | None] | None,\n rules: LogicalRules | None = None,\n) -> list[_UnassignedAxis | None | str | tuple[str, ...]] | None:\n \"\"\"Same as logical_to_mesh_axes, but doesn't fill in _unassigned_axis.\"\"\"\n if array_dim_names is None:\n return None\n if rules is None:\n rules = get_logical_axis_rules()\n axis_name_counts = collections.Counter(array_dim_names)\n # None and special values such as PartitionSpec.UNCONSTRAINED can appear more\n # then once.\n dups = tuple(\n k for k, v in axis_name_counts.items() if v > 1 and isinstance(k, str)\n )\n if dups:\n raise ValueError(\n f'Unsupported: Dimensions {dups} occur more than once in array names.'\n )\n if not isinstance(rules, (tuple, list)):\n raise ValueError('Unknown axis rule specification type.')\n # We assign mesh axes using a priority based ruleset over logical axis names.\n result: list[_UnassignedAxis | None | str | tuple[str, ...]]\n result = [\n (_unassigned_axis if isinstance(name, str) else name)\n for name in array_dim_names\n ]\n for rule_model_name, rule_mesh_names in rules:\n if rule_model_name in array_dim_names:\n pos = array_dim_names.index(rule_model_name)\n if (\n _mesh_assignment_free(rule_mesh_names, result)\n and result[pos] == _unassigned_axis\n ):\n result[pos] = rule_mesh_names\n return result\n\n\ndef logical_to_mesh_axes(\n array_dim_names: Sequence[str | None] | None,\n rules: LogicalRules | None = None,\n) -> jax.sharding.PartitionSpec | None:\n \"\"\"Compute layout for an array.\n\n The rules are in order of precedence, and consist of pairs:\n ``(ArrayDimensionName, MeshDimensionName)``, meaning that the given array\n dimension (if present and unused) should be sharded across the given\n mesh dimension (if present and unused).\n\n A Layout of an Array is expressed as a tuple with one element for each\n dimension in the Array. The element is either None, or is the name of a\n mesh-dimension, meaning that this dimension of the array is sharded across\n this dimension of the mesh.\n\n For example, given an array with::\n\n array_dim_names = ('batch', 'length', 'heads', 'features')\n\n and the layout rules are::\n\n rules = (('batch', 'X'),\n ('features', 'X'),\n ('heads', 'Y'),\n ('batch', 'Z'))\n\n then this function will return::\n\n PartitionSpec('X', None, 'Y', None)\n\n Args:\n array_dim_names: Tuple of array dimension names or None.\n rules: Optional logical to mesh rules override. Defaults to using the\n rules defined in the dynamic context set from the ``axis_rules`` function.\n\n Returns:\n PartitionSpec for the parameter.\n \"\"\"\n result = _logical_to_mesh_axes(array_dim_names, rules)\n if result is None:\n return None\n # We default to None - ie unsharded along the dimension.\n result = [None if x is _unassigned_axis else x for x in result]\n return jax.sharding.PartitionSpec(*result)\n\n\ndef logical_to_mesh(tree: Any, rules: LogicalRules | None = None) -> Any:\n \"\"\"Applies logical_to_mesh_axes to pytrees of logical PartitionSpecs.\"\"\"\n return jax.tree_util.tree_map(\n lambda x: logical_to_mesh_axes(x, rules),\n tree,\n is_leaf=lambda x: isinstance(x, jax.sharding.PartitionSpec),\n )\n\n\ndef logical_to_mesh_sharding(\n tree: Any,\n mesh: jax.sharding.Mesh,\n rules: LogicalRules | None = None,\n) -> Any:\n \"\"\"Convert pytrees of logical PartitionSpecs to shardings.\"\"\"\n return jax.tree_util.tree_map(\n lambda x: jax.sharding.NamedSharding(mesh, x),\n logical_to_mesh(tree, rules),\n is_leaf=lambda x: isinstance(x, jax.sharding.PartitionSpec),\n )\n\n\nclass RulesFallback(enum.Enum):\n \"\"\"How a sharding constraint should behave when no matching rule is found.\"\"\"\n\n AXIS_IS_UNSHARDED = 'axis_is_unsharded'\n RAISE_ERROR = 'raise_error'\n NO_CONSTRAINT = 'no_constraint'\n\n\ndef _with_sharding_constraint(\n x: Array,\n axis_resources: jax.sharding.PartitionSpec | None,\n mesh: jax.sharding.Mesh | None = None,\n):\n \"\"\"Wrapper for lax.with_sharding_constraint, no-op on cpu or outside jit.\"\"\"\n if not meta.global_mesh_defined() and mesh is None:\n return x\n else:\n if mesh is not None and axis_resources is not None:\n sharding = jax.sharding.NamedSharding(mesh, axis_resources)\n return lax.with_sharding_constraint(x, sharding)\n return lax.with_sharding_constraint(x, axis_resources)\n\n\ndef _with_sharding_constraint_one_fallback(\n axis_resources: LogicalPartitionSpec,\n x: Array,\n fallback: RulesFallback = RulesFallback.AXIS_IS_UNSHARDED,\n rules: LogicalRules | None = None,\n mesh: jax.sharding.Mesh | None = None,\n):\n \"\"\"Either imposes a sharding constraint or applies fallback.\"\"\"\n mesh_axes = _logical_to_mesh_axes(axis_resources, rules)\n if mesh_axes is None:\n return _with_sharding_constraint(x, None, mesh=mesh)\n\n if fallback == RulesFallback.AXIS_IS_UNSHARDED:\n mesh_axes = [None if x is _unassigned_axis else x for x in mesh_axes]\n else:\n if any(x is _unassigned_axis for x in mesh_axes):\n if fallback == RulesFallback.RAISE_ERROR:\n raise ValueError(f'Axis names {axis_resources} did not match a rule')\n else:\n return x\n return _with_sharding_constraint(\n x, jax.sharding.PartitionSpec(*mesh_axes), mesh=mesh\n )\n\n\ndef _is_axis_spec(x):\n return (\n isinstance(x, str)\n or x is jax.sharding.PartitionSpec.UNCONSTRAINED\n or x is None\n )\n\n\ndef _is_logical_spec(x):\n return x is None or (\n isinstance(x, tuple) and all(_is_axis_spec(e) for e in x)\n )\n\n\ndef with_logical_constraint(\n x: ArrayPytree,\n logical_axis_resources: LogicalPartitionSpecPytree,\n rules: LogicalRules | None = None,\n mesh: jax.sharding.Mesh | None = None,\n fallback: RulesFallback = RulesFallback.AXIS_IS_UNSHARDED,\n):\n \"\"\"Version of jit's with_sharding_constraint that uses logical axis names.\"\"\"\n # If no axis binding is set, this is a no-op.\n if rules is None:\n rules = get_logical_axis_rules()\n if not rules or logical_axis_resources is None:\n return x\n # Translate logical names to mesh assignments.\n return jax.tree_util.tree_map(\n functools.partial(\n _with_sharding_constraint_one_fallback,\n fallback=fallback,\n rules=rules,\n mesh=mesh,\n ),\n logical_axis_resources,\n x,\n is_leaf=_is_logical_spec,\n )\n\n\n# Logical Partitioning Axis Metadata\n# ------------------------------------------------------------------------------\n\n\nclass LogicallyPartitioned(meta.Partitioned):\n rules: LogicalRules | None = struct.field(default=None, pytree_node=False)\n\n # Directly comparing members of the type `jax.Array` can throw an error:\n # \"The truth value of an array with more than one element is ambiguous.\"\n # So we bring back an explicit implementation of __eq__ like it was prior to\n # Python 3.13 in order work around this possibility.\n def __eq__(self, other):\n if self is other:\n return True\n if other.__class__ is self.__class__:\n return (self.value,) == (other.value,)\n return NotImplemented\n\n def unbox(self, apply_constraint=True) -> Any:\n \"\"\"Returns the wrapped value with the partitioning constraint applied.\"\"\"\n if apply_constraint and (meta.global_mesh_defined() or self.mesh is not None):\n return with_logical_constraint(\n self.value,\n self.get_partition_spec(),\n rules=self.rules,\n mesh=self.mesh,\n )\n else:\n return self.value\n\n def to_nnx_metadata(self) -> dict[str, Any]:\n \"\"\"Return a dict of metadata that can translate into an `nnx.Variable`.\"\"\"\n metadata = vars(self)\n if 'names' in metadata:\n metadata['out_sharding'] = metadata.pop('names')\n if 'rules' in metadata:\n metadata['sharding_rules'] = metadata.pop('rules')\n return metadata\n\n @classmethod\n def from_nnx_metadata(cls, metadata: dict[str, Any]):\n \"\"\"Given a dict of `nnx.Variable` format metadata, create a `nn.LogicallyPartitioned`.\"\"\"\n metadata['names'] = metadata.pop('out_sharding')\n metadata['rules'] = metadata.pop('sharding_rules')\n fields = {x.name for x in dataclasses.fields(cls)}\n return cls(**{k: v for k, v in metadata.items() if k in fields})\n\n\ndef with_logical_partitioning(\n fn: Callable[..., Any],\n names: LogicalNames,\n mesh: jax.sharding.Mesh | None = None,\n rules: LogicalRules | None = None,\n) -> Callable[..., LogicallyPartitioned]:\n \"\"\"Wraps a function's return value with LogicallyPartitioned.\n\n Example::\n\n >>> import flax.linen as nn\n >>> kernel_init = nn.with_logical_partitioning(\n ... nn.initializers.lecun_normal(), (None, \"data\"))\n >>> partitioned_dense = nn.Dense(features=3, kernel_init=kernel_init)\n\n Args:\n fn: The function to be wrapped. Typically this is an initializer.\n names: The logical axis passed to ``LogicallyPartitioned``.\n mesh: The mesh to use for the partitioning. If None, the global mesh\n resource is used if available.\n rules: Optional logical to mesh rules use. If None, the global rules\n are used if available.\n Returns:\n A function wrapping ``fn`` that will return an instance of\n ``LogicallyPartitioned``.\n \"\"\"\n\n @functools.wraps(fn)\n def wrapper(*args, **kwargs):\n return LogicallyPartitioned(\n fn(*args, **kwargs), names, mesh=mesh, rules=rules\n ) # pytype: disable=wrong-keyword-args\n\n return wrapper\n" }, { "path": "flax/linen/stochastic.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Stochastic modules.\"\"\"\n\nfrom collections.abc import Sequence\n\nimport jax.numpy as jnp\nfrom jax import lax, random\n\nfrom flax.linen.module import Module, compact, merge_param\nfrom flax.typing import PRNGKey\n\n\nclass Dropout(Module):\n \"\"\"Create a dropout layer.\n\n .. note::\n When using :meth:`Module.apply() `, make sure\n to include an RNG seed named ``'dropout'``. Dropout isn't necessary for\n variable initialization.\n\n Example usage::\n\n >>> import flax.linen as nn\n >>> import jax, jax.numpy as jnp\n\n >>> class MLP(nn.Module):\n ... @nn.compact\n ... def __call__(self, x, train):\n ... x = nn.Dense(4)(x)\n ... x = nn.Dropout(0.5, deterministic=not train)(x)\n ... return x\n\n >>> model = MLP()\n >>> x = jnp.ones((1, 3))\n >>> variables = model.init(jax.random.key(0), x, train=False) # don't use dropout\n >>> model.apply(variables, x, train=False) # don't use dropout\n Array([[-0.17875527, 1.6255447 , -1.2431065 , -0.02554005]], dtype=float32)\n >>> model.apply(variables, x, train=True, rngs={'dropout': jax.random.key(1)}) # use dropout\n Array([[-0.35751054, 3.2510893 , 0. , 0. ]], dtype=float32)\n\n Attributes:\n rate: the dropout probability. (_not_ the keep rate!)\n broadcast_dims: dimensions that will share the same dropout mask\n deterministic: if false the inputs are scaled by ``1 / (1 - rate)`` and\n masked, whereas if true, no mask is applied and the inputs are returned as\n is.\n rng_collection: the rng collection name to use when requesting an rng key.\n \"\"\"\n\n rate: float\n broadcast_dims: Sequence[int] = ()\n deterministic: bool | None = None\n rng_collection: str = 'dropout'\n\n @compact\n def __call__(\n self,\n inputs,\n deterministic: bool | None = None,\n rng: PRNGKey | None = None,\n ):\n \"\"\"Applies a random dropout mask to the input.\n\n Args:\n inputs: the inputs that should be randomly masked.\n deterministic: if false the inputs are scaled by ``1 / (1 - rate)`` and\n masked, whereas if true, no mask is applied and the inputs are returned\n as is.\n rng: an optional PRNGKey used as the random key, if not specified, one\n will be generated using ``make_rng`` with the ``rng_collection`` name.\n\n Returns:\n The masked inputs reweighted to preserve mean.\n \"\"\"\n deterministic = merge_param(\n 'deterministic', self.deterministic, deterministic\n )\n\n if (self.rate == 0.0) or deterministic:\n return inputs\n\n # Prevent gradient NaNs in 1.0 edge-case.\n if self.rate == 1.0:\n return jnp.zeros_like(inputs)\n\n keep_prob = 1.0 - self.rate\n if rng is None:\n rng = self.make_rng(self.rng_collection)\n broadcast_shape = list(inputs.shape)\n for dim in self.broadcast_dims:\n broadcast_shape[dim] = 1\n mask = random.bernoulli(rng, p=keep_prob, shape=broadcast_shape)\n mask = jnp.broadcast_to(mask, inputs.shape)\n return lax.select(mask, inputs / keep_prob, jnp.zeros_like(inputs))\n" }, { "path": "flax/linen/summary.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Flax Module summary library.\"\"\"\n\nimport dataclasses\nimport enum\nimport io\nfrom abc import ABC, abstractmethod\nfrom types import MappingProxyType\nfrom typing import (\n Any,\n)\nfrom collections.abc import Callable, Iterable, Mapping, Sequence\n\nimport jax\nimport jax.numpy as jnp\nimport numpy as np\nimport rich.console\nimport rich.table\nimport rich.text\nimport yaml\n\nimport flax.linen.module as module_lib\nfrom flax.core import meta, unfreeze\nfrom flax.core.scope import (\n CollectionFilter,\n DenyList,\n LazyRng,\n)\nfrom flax.typing import (\n Array,\n PRNGKey,\n RNGSequences,\n FrozenVariableDict,\n MutableVariableDict,\n LogicalNames,\n)\n\n\nclass _ValueRepresentation(ABC):\n \"\"\"A class that represents a value in the summary table.\"\"\"\n\n @abstractmethod\n def render(self) -> str:\n ...\n\n\n@dataclasses.dataclass\nclass _ArrayRepresentation(_ValueRepresentation):\n shape: tuple[int, ...]\n dtype: Any\n\n @classmethod\n def from_array(cls, x: Array) -> '_ArrayRepresentation':\n return cls(jnp.shape(x), jnp.result_type(x))\n\n @classmethod\n def render_array(cls, x) -> str:\n return cls.from_array(x).render()\n\n def render(self):\n shape_repr = ','.join(str(x) for x in self.shape)\n return f'[dim]{self.dtype}[/dim][{shape_repr}]'\n\n\n@dataclasses.dataclass\nclass _PartitionedArrayRepresentation(_ValueRepresentation):\n array_representation: _ArrayRepresentation\n names: LogicalNames\n\n @classmethod\n def from_partitioned(\n cls, partitioned: meta.Partitioned\n ) -> '_PartitionedArrayRepresentation':\n return cls(\n _ArrayRepresentation.from_array(partitioned.value), partitioned.names\n )\n\n def render(self):\n return self.array_representation.render() + f' [dim]P[/dim]{self.names}'\n\n\n@dataclasses.dataclass\nclass _ObjectRepresentation(_ValueRepresentation):\n obj: Any\n\n def render(self):\n return repr(self.obj)\n\n\n@dataclasses.dataclass\nclass Row:\n \"\"\"Contains the information about a single row in the summary table.\n\n Attributes:\n path: A tuple of strings that represents the path to the module.\n module_copy: A copy of the module being summarized.\n method: method of the module called.\n inputs: inputs to the module.\n outputs: Output of the Module as reported by `capture_intermediates`.\n module_variables: Dictionary of variables in the module (no submodules\n included).\n counted_variables: Dictionary of variables that should be counted for this\n row, if no summarization is done (e.g. `depth=None` in `module_summary`)\n then this field is the same as `module_variables`, however if a\n summarization is done then this dictionary potentially contains parameters\n from submodules depending on the depth of the Module in question.\n flops: FLOPs cost of calling the module method.\n vjp_flops: FLOPs cost of calling the VJP of the module method.\n \"\"\"\n\n path: tuple[str, ...]\n module_copy: module_lib.Module\n method: str\n inputs: Any\n outputs: Any\n module_variables: dict[str, dict[str, Any]]\n counted_variables: dict[str, dict[str, Any]]\n flops: int\n vjp_flops: int\n\n def __post_init__(self):\n self.inputs = self.inputs\n self.outputs = self.outputs\n self.module_variables = self.module_variables\n self.counted_variables = self.counted_variables\n\n def size_and_bytes(\n self, collections: Iterable[str]\n ) -> dict[str, tuple[int, int]]:\n return {\n col: (\n _size_and_bytes(self.counted_variables[col])\n if col in self.counted_variables\n else (0, 0)\n )\n for col in collections\n }\n\n\nclass Table(list[Row]):\n \"\"\"A list of Row objects.\n\n Table inherits from `List[Row]` so it has all the methods of a list, however\n it also contains some additional fields:\n\n * `module`: the module that this table is summarizing\n * `collections`: a list containing the parameter collections (e.g. 'params', 'batch_stats', etc)\n \"\"\"\n\n def __init__(\n self,\n module: module_lib.Module,\n collections: Sequence[str],\n rows: Iterable[Row],\n ):\n super().__init__(rows)\n self.module = module\n self.collections = collections\n\n\ndef tabulate(\n module: module_lib.Module,\n rngs: PRNGKey | RNGSequences,\n depth: int | None = None,\n show_repeated: bool = False,\n mutable: CollectionFilter = DenyList('intermediates'),\n console_kwargs: Mapping[str, Any] | None = None,\n table_kwargs: Mapping[str, Any] = MappingProxyType({}),\n column_kwargs: Mapping[str, Any] = MappingProxyType({}),\n compute_flops: bool = False,\n compute_vjp_flops: bool = False,\n **kwargs,\n) -> Callable[..., str]:\n \"\"\"Returns a function that creates a summary of the Module represented as a table.\n\n This function accepts most of the same arguments and internally calls\n `Module.init`, except that it returns a function of the form\n `(*args, **kwargs) -> str` where `*args` and `**kwargs` are passed to\n `method` (e.g. `__call__`) during the forward pass.\n\n `tabulate` uses `jax.eval_shape` under the hood to run the forward computation\n without consuming any FLOPs or allocating memory.\n\n Additional arguments can be passed into the `console_kwargs` argument, for\n example, `{'width': 120}`. For a full list of `console_kwargs` arguments, see:\n https://rich.readthedocs.io/en/stable/reference/console.html#rich.console.Console\n\n Example::\n\n >>> import flax.linen as nn\n >>> import jax, jax.numpy as jnp\n\n >>> class Foo(nn.Module):\n ... @nn.compact\n ... def __call__(self, x):\n ... h = nn.Dense(4)(x)\n ... return nn.Dense(2)(h)\n\n >>> x = jnp.ones((16, 9))\n >>> tabulate_fn = nn.tabulate(\n ... Foo(), jax.random.key(0), compute_flops=True, compute_vjp_flops=True)\n\n >>> # print(tabulate_fn(x))\n\n This gives the following output::\n\n Foo Summary\n ┏━━━━━━━━━┳━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━┳━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┓\n ┃ path ┃ module ┃ inputs ┃ outputs ┃ flops ┃ vjp_flops ┃ params ┃\n ┡━━━━━━━━━╇━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━╇━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━┩\n │ │ Foo │ float32[16,9] │ float32[16,2] │ 1504 │ 4460 │ │\n ├─────────┼────────┼───────────────┼───────────────┼───────┼───────────┼─────────────────┤\n │ Dense_0 │ Dense │ float32[16,9] │ float32[16,4] │ 1216 │ 3620 │ bias: │\n │ │ │ │ │ │ │ float32[4] │\n │ │ │ │ │ │ │ kernel: │\n │ │ │ │ │ │ │ float32[9,4] │\n │ │ │ │ │ │ │ │\n │ │ │ │ │ │ │ 40 (160 B) │\n ├─────────┼────────┼───────────────┼───────────────┼───────┼───────────┼─────────────────┤\n │ Dense_1 │ Dense │ float32[16,4] │ float32[16,2] │ 288 │ 840 │ bias: │\n │ │ │ │ │ │ │ float32[2] │\n │ │ │ │ │ │ │ kernel: │\n │ │ │ │ │ │ │ float32[4,2] │\n │ │ │ │ │ │ │ │\n │ │ │ │ │ │ │ 10 (40 B) │\n ├─────────┼────────┼───────────────┼───────────────┼───────┼───────────┼─────────────────┤\n │ │ │ │ │ │ Total │ 50 (200 B) │\n └─────────┴────────┴───────────────┴───────────────┴───────┴───────────┴─────────────────┘\n\n Total Parameters: 50 (200 B)\n\n\n **Note**: rows order in the table does not represent execution order,\n instead it aligns with the order of keys in `variables` which are sorted\n alphabetically.\n\n **Note**: `vjp_flops` returns `0` if the module is not differentiable.\n\n Args:\n module: The module to tabulate.\n rngs: The rngs for the variable collections as passed to `Module.init`.\n depth: controls how many submodule deep the summary can go. By default its\n `None` which means no limit. If a submodule is not shown because of the\n depth limit, its parameter count and bytes will be added to the row of its\n first shown ancestor such that the sum of all rows always adds up to the\n total number of parameters of the Module.\n show_repeated: If `True`, repeated calls to the same module will be shown\n in the table, otherwise only the first call will be shown. Default is\n `False`.\n mutable: Can be bool, str, or list. Specifies which collections should be\n treated as mutable: ``bool``: all/no collections are mutable. ``str``: The\n name of a single mutable collection. ``list``: A list of names of mutable\n collections. By default all collections except 'intermediates' are\n mutable.\n console_kwargs: An optional dictionary with additional keyword arguments\n that are passed to `rich.console.Console` when rendering the table.\n Default arguments are `{'force_terminal': True, 'force_jupyter': False}`.\n table_kwargs: An optional dictionary with additional keyword arguments that\n are passed to `rich.table.Table` constructor.\n column_kwargs: An optional dictionary with additional keyword arguments that\n are passed to `rich.table.Table.add_column` when adding columns to the\n table.\n compute_flops: whether to include a `flops` column in the table listing the\n estimated FLOPs cost of each module forward pass. Does incur actual\n on-device computation / compilation / memory allocation, but still\n introduces overhead for large modules (e.g. extra 20 seconds for a\n Stable Diffusion's UNet, whereas otherwise tabulation would finish in 5\n seconds).\n compute_vjp_flops: whether to include a `vjp_flops` column in the table\n listing the estimated FLOPs cost of each module backward pass. Introduces\n a compute overhead of about 2-3X of `compute_flops`.\n **kwargs: Additional arguments passed to `Module.init`.\n\n Returns:\n A function that accepts the same `*args` and `**kwargs` of the forward pass\n (`method`) and returns a string with a tabular representation of the\n Modules.\n \"\"\"\n # add non-default arguments to kwargs, this prevents some issue we overloading init\n # see: https://github.com/google/flax/issues/3299\n if mutable != DenyList('intermediates'):\n kwargs['mutable'] = mutable\n\n def _tabulate_fn(*fn_args, **fn_kwargs):\n table_fn = _get_module_table(\n module,\n depth=depth,\n show_repeated=show_repeated,\n compute_flops=compute_flops,\n compute_vjp_flops=compute_vjp_flops,\n )\n\n table = table_fn(rngs, *fn_args, **fn_kwargs, **kwargs)\n\n non_param_cols = [\n 'path',\n 'module',\n 'inputs',\n 'outputs',\n ]\n\n if compute_flops:\n non_param_cols.append('flops')\n if compute_vjp_flops:\n non_param_cols.append('vjp_flops')\n\n return _render_table(\n table, console_kwargs, table_kwargs, column_kwargs, non_param_cols\n )\n\n return _tabulate_fn\n\n\ndef _get_flops(fn, *args, **kwargs):\n e = jax.jit(fn).lower(*args, **kwargs)\n cost = e.cost_analysis()\n if cost is None:\n return 0\n flops = int(cost['flops']) if 'flops' in cost else 0\n return flops\n\n\ndef _get_call_flops(\n c: module_lib._CallInfo,\n compute_flops: bool,\n compute_vjp_flops: bool,\n) -> tuple[int, int]:\n \"\"\"Return the FLOPs of executing the call `c` in the call stack.\n\n Does not perform actual computation / compilation / memory allocation, but\n still introduces overhead for large modules.\n\n Args:\n c: ``_CallInfo``.\n compute_flops: whether to compute forward pass FLOPs. Return `-1` otherwise.\n compute_vjp_flops: whether to compute backward pass FLOPs. Return `-1`\n otherwise.\n\n Returns:\n FLOPs of executing forward pass of `c`, and its VJP.\n \"\"\"\n\n if not compute_flops and not compute_vjp_flops:\n return -1, -1\n\n rngs = jax.tree_util.tree_map(\n lambda x: x.rng, c.rngs, is_leaf=lambda x: isinstance(x, LazyRng)\n )\n\n args = jax.tree_util.tree_map(_from_value_representation, c.args)\n kwargs = jax.tree_util.tree_map(_from_value_representation, c.kwargs)\n\n leaves, treedef = jax.tree_util.tree_flatten((args, kwargs))\n dynamic_leaves = []\n dynamic_idxs = []\n for i, arg in enumerate(leaves):\n if isinstance(arg, jax.ShapeDtypeStruct):\n dynamic_leaves.append(arg)\n dynamic_idxs.append(i)\n\n def _get_inputs(dynamic_leaves):\n new_leaves: list[Any] = leaves.copy()\n for i, arg in zip(dynamic_idxs, dynamic_leaves):\n new_leaves[i] = arg\n return treedef.unflatten(new_leaves)\n\n def init(rngs, dynamic_leaves):\n \"\"\"`c.module.init` closed over static keyword arguments.\"\"\"\n args, kwargs = _get_inputs(dynamic_leaves)\n return c.module.init(\n rngs,\n *args,\n method=c.method,\n mutable=c.mutable,\n **kwargs,\n )\n\n variables = jax.eval_shape(init, rngs, dynamic_leaves)\n\n def apply(variables, rngs, dynamic_leaves):\n \"\"\"`c.module.apply` closed over static keyword arguments.\"\"\"\n args, kwargs = _get_inputs(dynamic_leaves)\n return c.module.apply(\n variables,\n *args,\n rngs=rngs,\n method=c.method,\n mutable=c.mutable,\n **kwargs,\n )\n\n # Forward pass FLOPs\n if compute_flops:\n flops = _get_flops(apply, variables, rngs, dynamic_leaves)\n else:\n flops = -1\n\n if compute_vjp_flops:\n # Backward pass FLOPs\n def apply_vjp(variables, rngs, dynamic_leaves):\n \"\"\"VJP of `c.module.apply` closed over static keyword arguments.\"\"\"\n out, vjp_fn = jax.vjp(apply, variables, rngs, dynamic_leaves)\n return vjp_fn(out)\n\n vjp_flops = _get_flops(apply_vjp, variables, rngs, dynamic_leaves)\n else:\n vjp_flops = -1\n\n return flops, vjp_flops\n\n\ndef _get_module_table(\n module: module_lib.Module,\n depth: int | None,\n show_repeated: bool,\n compute_flops: bool,\n compute_vjp_flops: bool,\n) -> Callable[..., Table]:\n \"\"\"A function that takes a Module and returns function with the same signature\n as `init` but returns the Table representation of the Module.\"\"\"\n\n def _get_table_fn(*args, **kwargs):\n with module_lib._tabulate_context():\n\n def _get_variables():\n return module.init(*args, **kwargs)\n # TODO(cgarciae): is it possible to avoid leaking tracers for summaries?\n with jax.check_tracer_leaks(False):\n variables = jax.eval_shape(_get_variables)\n calls = module_lib._context.call_info_stack[-1].calls\n calls.sort(key=lambda c: c.index)\n\n collections: set[str] = set(variables.keys())\n rows = []\n all_paths: set[tuple[str, ...]] = {call.path for call in calls}\n visited_paths: set[tuple[str, ...]] = set()\n\n for c in calls:\n call_depth = len(c.path)\n inputs = _process_inputs(c.args, c.kwargs)\n\n if c.path in visited_paths or not hasattr(c.module, c.method):\n if not show_repeated:\n continue\n module_vars = {}\n counted_vars = {}\n elif depth is not None:\n if call_depth > depth:\n continue\n module_vars, _ = _get_module_variables(c.path, variables, all_paths)\n if call_depth == depth:\n counted_vars = _get_path_variables(c.path, variables)\n else:\n counted_vars = module_vars\n else:\n module_vars, _ = _get_module_variables(c.path, variables, all_paths)\n counted_vars = module_vars\n\n visited_paths.add(c.path)\n rows.append(\n Row(\n c.path,\n c.module.copy(parent=None),\n c.method,\n inputs,\n c.outputs,\n module_vars,\n counted_vars,\n *_get_call_flops(c, compute_flops, compute_vjp_flops),\n )\n )\n\n return Table(module, tuple(collections), rows)\n\n return _get_table_fn\n\n\ndef _get_module_variables(\n path: tuple[str, ...],\n variables: FrozenVariableDict,\n all_paths: set[tuple[str, ...]],\n) -> tuple[MutableVariableDict, Any]:\n \"\"\"A function that takes a path and variables structure and returns a\n\n (module_variables, submodule_variables) tuple for that path.\n _get_module_variables\n uses the `all_paths` set to determine if a variable belongs to a submodule or\n not.\n \"\"\"\n module_variables = _get_path_variables(path, variables)\n submodule_variables: Any = {collection: {} for collection in module_variables}\n all_keys = {\n key for collection in module_variables.values() for key in collection\n }\n\n for key in all_keys:\n submodule_path = path + (key,)\n if submodule_path in all_paths:\n for collection in module_variables:\n if key in module_variables[collection]:\n submodule_variables[collection][key] = module_variables[\n collection\n ].pop(key)\n\n return module_variables, submodule_variables\n\n\ndef _get_path_variables(\n path: tuple[str, ...], variables: FrozenVariableDict\n) -> MutableVariableDict:\n \"\"\"A function that takes a path and a variables structure and returns the\n variable structure at that path.\n \"\"\"\n path_variables = {}\n\n for collection in variables:\n collection_variables = variables[collection]\n for name in path:\n if name not in collection_variables:\n collection_variables = None\n break\n collection_variables = collection_variables[name]\n\n if collection_variables is not None:\n path_variables[collection] = unfreeze(collection_variables)\n\n return path_variables\n\n\ndef _process_inputs(args, kwargs) -> Any:\n \"\"\"A function that normalizes the representation of the ``args`` and\n ``kwargs`` for the ``inputs`` column.\n \"\"\"\n if args and kwargs:\n input_values = (*args, kwargs)\n elif args and not kwargs:\n input_values = args[0] if len(args) == 1 else args\n elif kwargs and not args:\n input_values = kwargs\n else:\n input_values = ()\n\n return input_values\n\n\ndef _render_table(\n table: Table,\n console_extras: Mapping[str, Any] | None,\n table_kwargs: Mapping[str, Any],\n column_kwargs: Mapping[str, Any],\n non_params_cols: list[str],\n) -> str:\n \"\"\"A function that renders a Table to a string representation using rich.\"\"\"\n console_kwargs = {'force_terminal': True, 'force_jupyter': False}\n if console_extras is not None:\n console_kwargs.update(console_extras)\n\n rich_table = rich.table.Table(\n show_header=True,\n show_lines=True,\n show_footer=True,\n title=f'{table.module.__class__.__name__} Summary',\n **table_kwargs,\n )\n\n for c in non_params_cols:\n rich_table.add_column(c, **column_kwargs)\n\n for col in table.collections:\n rich_table.add_column(col, **column_kwargs)\n\n for row in table:\n collections_size_repr = []\n\n for collection, size_bytes in row.size_and_bytes(table.collections).items():\n col_repr = ''\n\n if collection in row.module_variables:\n module_variables = _represent_tree(row.module_variables[collection])\n module_variables = _normalize_structure(module_variables)\n col_repr += _as_yaml_str(\n _summary_tree_map(_maybe_render, module_variables)\n )\n if col_repr:\n col_repr += '\\n\\n'\n\n col_repr += f'[bold]{_size_and_bytes_repr(*size_bytes)}[/bold]'\n collections_size_repr.append(col_repr)\n\n no_show_methods = {'__call__', ''}\n path_repr = '/'.join(row.path)\n method_repr = (\n f' [dim]({row.method})[/dim]' if row.method not in no_show_methods else ''\n )\n rich_table.add_row(\n path_repr,\n type(row.module_copy).__name__ + method_repr,\n *(\n _as_yaml_str(\n _summary_tree_map(\n _maybe_render, _normalize_structure(getattr(row, c))\n )\n )\n for c in non_params_cols[2:]\n ),\n *collections_size_repr,\n )\n\n # add footer with totals\n n_non_params_cols = len(non_params_cols)\n rich_table.columns[n_non_params_cols - 1].footer = rich.text.Text.from_markup(\n 'Total', justify='right'\n )\n\n # get collection totals\n collection_total = {col: (0, 0) for col in table.collections}\n for row in table:\n for col, size_bytes in row.size_and_bytes(table.collections).items():\n collection_total[col] = (\n collection_total[col][0] + size_bytes[0],\n collection_total[col][1] + size_bytes[1],\n )\n\n # add totals to footer\n for i, col in enumerate(table.collections):\n rich_table.columns[n_non_params_cols + i].footer = _size_and_bytes_repr(\n *collection_total[col]\n )\n\n # add final totals to caption\n caption_totals = (0, 0)\n for size, num_bytes in collection_total.values():\n caption_totals = (\n caption_totals[0] + size,\n caption_totals[1] + num_bytes,\n )\n\n rich_table.caption_style = 'bold'\n rich_table.caption = (\n f'\\nTotal Parameters: {_size_and_bytes_repr(*caption_totals)}'\n )\n\n return '\\n' + _get_rich_repr(rich_table, console_kwargs) + '\\n'\n\n\ndef _summary_tree_map(f, tree, *rest):\n return jax.tree_util.tree_map(f, tree, *rest, is_leaf=lambda x: x is None)\n\n\ndef _size_and_bytes_repr(size: int, num_bytes: int) -> str:\n if not size:\n return ''\n bytes_repr = _bytes_repr(num_bytes)\n return f'{size:,} [dim]({bytes_repr})[/dim]'\n\n\ndef _size_and_bytes(pytree: Any) -> tuple[int, int]:\n leaves = jax.tree_util.tree_leaves(pytree)\n size = sum(x.size for x in leaves if hasattr(x, 'size'))\n num_bytes = sum(\n x.size * x.dtype.itemsize for x in leaves if hasattr(x, 'size')\n )\n return size, num_bytes\n\n\ndef _get_rich_repr(obj, console_kwargs):\n f = io.StringIO()\n console = rich.console.Console(file=f, **console_kwargs)\n console.print(obj)\n return f.getvalue()\n\n\ndef _as_yaml_str(value) -> str:\n if (hasattr(value, '__len__') and len(value) == 0) or value is None:\n return ''\n\n file = io.StringIO()\n yaml.safe_dump(\n value,\n file,\n default_flow_style=False,\n indent=2,\n sort_keys=False,\n explicit_end=False,\n )\n return file.getvalue().replace('\\n...', '').replace(\"'\", '').strip()\n\n\ndef _normalize_structure(obj):\n if isinstance(obj, _ValueRepresentation):\n return obj\n if isinstance(obj, (tuple, list)):\n return tuple(map(_normalize_structure, obj))\n elif isinstance(obj, Mapping):\n return {\n _normalize_structure(k): _normalize_structure(v) for k, v in obj.items()\n }\n elif dataclasses.is_dataclass(obj):\n return {\n f.name: _normalize_structure(getattr(obj, f.name))\n for f in dataclasses.fields(obj)\n }\n elif isinstance(obj, enum.Enum):\n # `yaml.safe_dump` does not support Enum key types so extract the underlying value\n return obj.value\n else:\n return obj\n\n\ndef _bytes_repr(num_bytes):\n count, units = (\n (f'{num_bytes / 1e9 :,.1f}', 'GB')\n if num_bytes > 1e9\n else (f'{num_bytes / 1e6 :,.1f}', 'MB')\n if num_bytes > 1e6\n else (f'{num_bytes / 1e3 :,.1f}', 'KB')\n if num_bytes > 1e3\n else (f'{num_bytes:,}', 'B')\n )\n\n return f'{count} {units}'\n\n\ndef _get_value_representation(x: Any) -> _ValueRepresentation:\n if isinstance(x, (int, float, bool, type(None))) or (\n isinstance(x, np.ndarray) and np.isscalar(x)\n ):\n return _ObjectRepresentation(x)\n elif isinstance(x, meta.Partitioned):\n return _PartitionedArrayRepresentation.from_partitioned(x)\n try:\n return _ArrayRepresentation.from_array(x)\n except:\n return _ObjectRepresentation(x)\n\n\ndef _from_value_representation(x: _ValueRepresentation) -> Any:\n if isinstance(x, _ArrayRepresentation):\n return jax.ShapeDtypeStruct(x.shape, x.dtype)\n\n elif isinstance(x, _PartitionedArrayRepresentation):\n return jax.ShapeDtypeStruct(\n x.array_representation.shape, x.array_representation.dtype\n )\n\n elif isinstance(x, _ObjectRepresentation):\n return x.obj\n\n raise TypeError(x, type(x))\n\n\ndef _represent_tree(x):\n \"\"\"Returns a tree with the same structure as `x` but with each leaf replaced\n by a `_ValueRepresentation` object.\"\"\"\n return jax.tree_util.tree_map(\n _get_value_representation,\n x,\n is_leaf=lambda x: x is None or isinstance(x, meta.Partitioned),\n )\n\n\ndef _maybe_render(x):\n return x.render() if hasattr(x, 'render') else repr(x)\n" }, { "path": "flax/linen/transforms.py", "content": "# Copyright 2023 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"JAX transformations on Modules.\n\nJax functional transformations operate on pure functions.\nFlax extends these transformations to also operate on Module's which\nhave stateful variables and PRNG sequences. We refer to these extended\nversions as \"lifted transformations\".\n\nA lifted transformation can be applied to a ``Module`` class or a\nfunction that takes a ``Module`` instance as its first argument.\n\"\"\"\n\nfrom collections.abc import Callable, Iterable, Mapping, Sequence\nimport contextlib\nimport dataclasses\nimport functools\nimport inspect\nfrom typing import (\n Any,\n TypeVar,\n Union,\n)\nimport weakref\n\nfrom flax import core\nfrom flax import errors, struct, traceback_util\nfrom flax import serialization\nfrom flax.core import Scope, lift, meta\nfrom flax.core.frozen_dict import FrozenDict\nfrom flax.core.scope import (\n CollectionFilter,\n LazyRng,\n PRNGSequenceFilter,\n)\nfrom flax.ids import FlaxId\nfrom flax.linen import module as linen_module\nfrom flax.linen.module import (\n Module,\n Variable,\n _derive_profiling_name,\n _get_unbound_fn,\n wrap_method_once,\n)\nfrom flax.typing import (\n InOutAxis,\n InOutScanAxis,\n)\nimport jax\n\ntraceback_util.register_exclusion(__file__)\n\n# pylint: disable=protected-access,dangerous-default-value\n\n\n# Utils\n# -----------------------------------------------------------------------------\ndef clean_clone(x):\n \"\"\"Remove scopes and tracers from children.\"\"\"\n if isinstance(x, Module):\n object.__setattr__(\n x, 'children', {k: clean_clone(v) for k, v in x.children.items()}\n )\n object.__setattr__(x, 'scope', None)\n return x\n\n\n@struct.dataclass\nclass VariablePlaceholder:\n \"\"\"Used to mark Variables in a JAX-compatible way when lifting arguments.\"\"\"\n\n collection: str = struct.field(pytree_node=False)\n name: str = struct.field(pytree_node=False)\n unbox: bool = struct.field(pytree_node=False)\n id: int = struct.field(pytree_node=False)\n\n\n@struct.dataclass\nclass InstancePlaceholder:\n \"\"\"Marks module instances in a JAX-compatible way when lifting arguments.\"\"\"\n\n cls: type[Any] = struct.field(pytree_node=False)\n attrs: dict[Any, Any] = struct.field(pytree_node=False)\n id: int = struct.field(pytree_node=False)\n\n\ndef _memoize_by_id(fn, refs):\n \"\"\"Memoization by module/variable id to handle aliasing in traversal.\"\"\"\n\n @functools.wraps(fn)\n def wrapped_fn(x):\n nonlocal refs\n if isinstance(x, (VariablePlaceholder, InstancePlaceholder)):\n x_id = x.id\n elif isinstance(x, (Variable, Module)):\n x_id = x._id\n else:\n return fn(x)\n if x_id not in refs:\n refs[x_id] = fn(x)\n return refs[x_id]\n\n return wrapped_fn\n\n\ndef get_module_scopes(module, args=None, kwargs=None):\n \"\"\"Get all scopes on module, including constructor Module arguments.\n\n To properly functionalize a Module that has other bound Modules passed in\n \"from the outside\" as dataclass attributes, we need to traverse all dataclass\n fields to find the Scopes associated with the Module. Additionally, because\n we allow Modules to be passed inside pytrees on the dataclass attributes, we\n must traverse all dataclass attributes as pytrees to find all Modules. We\n additionally handle lifting Variables (which are just references to data in\n particular scopes) and Module instances that are passed as arguments to\n methods.\n\n Args:\n module: a bound flax Module.\n args: an *args list possibly containing Variables or Module instances\n referencing a scope.\n kwargs: a **kwargs dict possibly containing Variables or Module instances\n referencing a scope.\n\n Returns:\n A list of all functional-core Scopes bound on self and inside dataclass\n fields as well as any Scopes passed via argument Variables, an updated args\n list, and an updated kwargs dict that have both had Variables replaced with\n VariablePlaceholders and Module instances replaced with InstancePlaceholders\n that are compatible with jax functions.\n \"\"\"\n scopes: list[Scope] = []\n refs = {}\n\n # Gather scopes associated with Variables and Module instances passed as\n # positional and keyword arguments.\n @functools.partial(_memoize_by_id, refs=refs)\n def get_arg_scope(x):\n nonlocal scopes\n if isinstance(x, Variable) and isinstance(x.scope, Scope):\n scopes.append(x.scope)\n return VariablePlaceholder(x.collection, x.name, x.unbox, x._id)\n elif isinstance(x, Module) and isinstance(x.scope, Scope):\n x._try_setup(shallow=True)\n scopes.append(x.scope)\n attrs = {\n f.name: getattr(x, f.name)\n for f in dataclasses.fields(x)\n if f.name != 'parent' and f.init\n }\n attrs = jax.tree_util.tree_map(get_arg_scope, attrs)\n return InstancePlaceholder(x.__class__, attrs, x._id)\n return x\n\n new_args, new_kwargs = jax.tree_util.tree_map(get_arg_scope, (args, kwargs))\n\n # Gather scopes in Variables and Submodules passed as Module attributes.\n @functools.partial(_memoize_by_id, refs=refs)\n def get_scopes(module):\n nonlocal scopes\n module._try_setup(shallow=True)\n\n def get_scopes_inner(x):\n nonlocal scopes\n if isinstance(x, Module) and isinstance(x.scope, Scope):\n get_scopes(x)\n elif isinstance(x, Variable) and isinstance(x.scope, Scope):\n scopes.append(x.scope)\n\n attrs = {\n f.name: getattr(module, f.name)\n for f in dataclasses.fields(module)\n if f.name != 'parent' and f.init\n }\n for leaf in jax.tree_util.tree_leaves(attrs):\n get_scopes_inner(leaf)\n scopes.append(module.scope)\n\n get_scopes(module)\n return scopes, new_args, new_kwargs\n\n\ndef set_module_scopes(module, args, kwargs, scopes):\n \"\"\"Set all scopes on module, including those on Modules in dataclass fields.\n\n To properly functionalize a Module we must also \"rehydrate\" it with Scopes\n from `get_module_scopes`. We need to set scopes not just on the Module but\n also on any Module living inside dataclass attributes or even pytrees in its\n dataclass attributes. We additionally handle restoring Variables and Module\n instances from their placeholders in the method positional and keyword\n arguments. The order of traversal through this method is the same as in\n `get_module_scopes`, guaranteeing the correct Scopes are applied to each\n Module.\n\n Args:\n module: a flax Module.\n args: an *args list possibly containing VariablePlaceholder or\n InstancePlaceholder members.\n kwargs: a **kwargs dict possibly containing VariablePlaceholder or\n InstancePlaceholder members.\n scopes: a list of Scopes corresponding to this Module and its arguments that\n was created by the `get_module_scopes` function.\n\n Returns:\n A copy of the module with it and its attributes bound to the scopes passed\n to this function, an updated args list, and an updated kwargs dict with\n updated Variable and Module instance references.\n \"\"\"\n idx = 0\n refs = {}\n\n # Set scopes associated with Variables and Module instances passed as\n # positional and keyword arguments.\n @functools.partial(_memoize_by_id, refs=refs)\n def set_arg_scope(x):\n nonlocal idx\n if isinstance(x, VariablePlaceholder):\n new_x = Variable(\n scope=scopes[idx], collection=x.collection, name=x.name, unbox=x.unbox\n )\n idx += 1\n return new_x\n elif isinstance(x, InstancePlaceholder):\n instance_scope = scopes[idx]\n idx += 1\n instance_attrs = jax.tree_util.tree_map(set_arg_scope, x.attrs)\n return x.cls(parent=instance_scope, **instance_attrs)\n return x\n\n def is_placeholder(x):\n return isinstance(x, (VariablePlaceholder, InstancePlaceholder))\n\n new_args, new_kwargs = jax.tree_util.tree_map(\n set_arg_scope, (args, kwargs), is_leaf=is_placeholder\n )\n\n # set scopes in Variables and Submodules passed as Module attributes\n @functools.partial(_memoize_by_id, refs=refs)\n def set_scopes(module):\n nonlocal idx\n\n def set_scopes_inner(x):\n nonlocal idx\n if isinstance(x, Module) and isinstance(x.scope, Scope):\n return set_scopes(x)\n elif isinstance(x, Variable) and isinstance(x.scope, Scope):\n new_x = Variable(\n scope=scopes[idx],\n collection=x.collection,\n name=x.name,\n unbox=x.unbox,\n )\n idx += 1\n return new_x\n else:\n return x\n\n attrs = {\n f.name: getattr(module, f.name)\n for f in dataclasses.fields(module)\n if f.name != 'parent' and f.init\n }\n new_attrs = jax.tree_util.tree_map(set_scopes_inner, attrs)\n new_module = module.clone(parent=scopes[idx], **new_attrs)\n idx += 1\n return new_module\n\n new_module = set_scopes(module)\n assert len(scopes) == idx, f'scope list mismatch {len(scopes)} != {idx}'\n return new_module, new_args, new_kwargs\n\n\ndef _test_transformed_return_values(tree, method_name):\n \"\"\"Tests whether the return value contains any Modules or Variables.\"\"\"\n impure = any(\n map(\n lambda x: isinstance(x, (Module, Variable)),\n jax.tree_util.tree_leaves(tree),\n )\n )\n if impure:\n raise errors.TransformedMethodReturnValueError(method_name)\n\n\n# Class lifting\n# -----------------------------------------------------------------------------\ndef module_class_lift_transform(\n transform, module_class, *trafo_args, methods=None, **trafo_kwargs\n):\n \"\"\"Module class lift transform.\"\"\"\n # TODO(marcvanzee): Improve docstrings (#1977).\n # TODO(levskaya): find nicer argument convention for multi-method case?\n\n # Prepare per-method transform args, kwargs.\n if methods is None:\n # Default case, just transform __call__\n class_trafo_args = {'__call__': (trafo_args, trafo_kwargs)}\n elif isinstance(methods, (list, tuple)):\n # Transform every method in methods with given args, kwargs.\n class_trafo_args = {m: (trafo_args, trafo_kwargs) for m in methods}\n elif isinstance(methods, dict):\n # Pass different trafo args per each method.\n class_trafo_args = {k: ((), v) for k, v in methods.items()}\n else:\n raise ValueError(\n 'transform methods argument must be None, tuple, list, or dict.'\n )\n\n # Handle partially initialized module class constructors.\n if isinstance(module_class, functools.partial) and issubclass(\n module_class.func, Module\n ):\n partial_object = module_class\n module_class = module_class.func\n else:\n partial_object = None\n\n def create_trans_fn(fn_name, fn_trafo_args):\n # get existing unbound method from class\n fn = getattr(module_class, fn_name)\n trafo_args, trafo_kwargs = fn_trafo_args\n\n # we need to create a scope-function from our class for the given method\n @functools.wraps(fn)\n def wrapped_fn(self, *args, **kwargs):\n state = self._state.export()\n\n # make a scope-function to transform\n def core_fn(scopes, *args, **kwargs):\n # make a clone of self using its arguments\n attrs = {\n f.name: getattr(self, f.name)\n for f in dataclasses.fields(self)\n if f.name != 'parent' and f.init\n }\n # we reference module_class, not self.__class__ to avoid infinite loop\n cloned = module_class(parent=None, **attrs)\n cloned, args, kwargs = set_module_scopes(cloned, args, kwargs, scopes)\n object.__setattr__(cloned, '_state', state.export())\n res = fn(cloned, *args, **kwargs)\n self._state.reimport(cloned._state)\n _test_transformed_return_values(res, fn_name)\n return res\n\n # here we apply the given lifting transform to the scope-ingesting fn\n trafo_fn = transform(core_fn, *trafo_args, **trafo_kwargs)\n module_scopes, args, kwargs = get_module_scopes(self, args, kwargs)\n ret = trafo_fn(module_scopes, *args, **kwargs)\n return ret\n\n return wrapped_fn\n\n transformed_fns = {\n fn_name: create_trans_fn(fn_name, fn_trafo_args)\n for fn_name, fn_trafo_args in class_trafo_args.items()\n }\n # construct new dynamic class w. transformed methods\n transformed_cls = type(\n transform.__name__.capitalize() + module_class.__name__,\n (module_class,),\n transformed_fns,\n )\n # Handle partially initialized module class constructors.\n if partial_object is not None:\n transformed_cls = functools.partial(\n transformed_cls, *partial_object.args, **partial_object.keywords\n )\n return transformed_cls\n\n\n# Function lifting as decorator on methods __inside__ class definition.\n# -----------------------------------------------------------------------------\ndef decorator_lift_transform(\n transform, class_fn, *trafo_args, multi_scope=True, **trafo_kwargs\n):\n \"\"\"Decorator for lifted transform.\"\"\"\n # TODO(marcvanzee): Improve docstrings (#1977).\n # Due to the ordering of method decorators, we must wrap the class_fn\n # with the module state management wrapper first to maintain Module state\n # correctly.\n if isinstance(class_fn, tuple):\n class_fns = class_fn\n else:\n class_fns = (class_fn,)\n prewrapped_fns = [wrap_method_once(class_fn) for class_fn in class_fns]\n\n @functools.wraps(prewrapped_fns[0])\n def wrapped_fn(self: Module, *args, **kwargs):\n state = self._state.export()\n\n # make a scope-function to transform\n def core_fn(prewrapped_fn, class_fn, scopes, *args, **kwargs):\n if not multi_scope:\n scopes = [scopes]\n cloned, args, kwargs = set_module_scopes(self, args, kwargs, scopes)\n object.__setattr__(cloned, '_state', state.export())\n res = prewrapped_fn(cloned, *args, **kwargs)\n self._state.reimport(cloned._state)\n _test_transformed_return_values(res, getattr(class_fn, '__name__', None))\n return res\n\n core_fns = [\n functools.partial(core_fn, prewrapped_fn, class_fn)\n for prewrapped_fn, class_fn in zip(prewrapped_fns, class_fns)\n ]\n # here we apply the given lifting transform to the scope-ingesting fn\n trafo_fn = transform(*core_fns, *trafo_args, **trafo_kwargs)\n module_scopes, args, kwargs = get_module_scopes(self, args, kwargs)\n if not multi_scope:\n if len(module_scopes) != 1:\n # TODO(levskaya): transforms like jvp & vjp have args that follow the\n # pytree structure of scopes. The user doesn't explicitly control shared\n # modules passed as arguments to methods or as attributes to Module\n # constructors. Therefore, there is no obvious API for specifying\n # arguments per lifted Module.\n raise NotImplementedError(\n 'This transform does not yet support'\n ' Modules that include other Modules passed as arguments.'\n )\n module_scopes = module_scopes[0]\n return trafo_fn(module_scopes, *args, **kwargs)\n\n return wrapped_fn\n\n\n@dataclasses.dataclass(frozen=True)\nclass _HashableProxy:\n \"\"\"A hashable proxy object that is use to define a hash for Modules.\n\n The hash produced by _HashableProxy is useful for nn.jit to decide if a\n function should be retraced or not\n \"\"\"\n\n module_ref: weakref.ref\n hash_key: int\n\n @classmethod\n def from_module(cls, module: Module) -> '_HashableProxy':\n fingerprint = _module_fingerprint(module)\n hash_key = hash(fingerprint)\n return cls(weakref.ref(module), hash_key)\n\n def __hash__(self):\n return self.hash_key\n\n def __eq__(self, other):\n return isinstance(other, _HashableProxy) and self.hash_key == other.hash_key\n\n @property\n def module(self):\n return self.module_ref()\n\n\ndef _module_fingerprint(module: Module) -> tuple[type[Any], Any]:\n return _fingerprint_recursive(module, (), {})\n\n\ndef _fingerprint_recursive(\n obj: Any, path: tuple[str, ...], seen_modules: dict[FlaxId, int]\n) -> Any:\n \"\"\"Creates a hashable representation for a Module by traversing its structure recursively.\"\"\"\n\n def _get_fingerprint(name: str, value: Any) -> tuple[str, Any]:\n return name, _fingerprint_recursive(value, (*path, name), seen_modules)\n\n if isinstance(obj, str):\n return obj\n elif hasattr(obj, '__fn_or_cls__'): # support PaxConfig objects\n return _fingerprint_recursive(obj.__fn_or_cls__, path, seen_modules)\n elif isinstance(obj, Module):\n fingerprint: Any\n if obj._id in seen_modules:\n # if we have already seen the module we just use the index\n # as its static component\n fingerprint = seen_modules[obj._id]\n return type(obj), fingerprint\n else:\n # if its a new module we add it to the cache and give it\n # a new index\n seen_modules[obj._id] = len(seen_modules)\n # TODO(cgarciae): define a way for the user of nn.jit to define\n # what fields it wants to ignore per Module instance.\n fingerprints = []\n for field in dataclasses.fields(obj):\n if not hasattr(obj, field.name):\n continue\n if field.name not in ('parent', 'name'):\n value = getattr(obj, field.name)\n fingerprints.append(_get_fingerprint(field.name, value))\n # add state fingerprint\n state_fingerprint = (\n _get_fingerprint('in_compact_method', obj._state.in_compact_method),\n _get_fingerprint('in_setup', obj._state.in_setup),\n _get_fingerprint('setup_called', obj._state.setup_called),\n _get_fingerprint('is_initialized', obj._state.is_initialized),\n _get_fingerprint('autoname_cursor', obj._state.autoname_cursor),\n )\n fingerprints.append(('_state', state_fingerprint))\n # add scope fingerprint\n scope = obj.scope\n if scope is not None:\n static_scope = (\n _get_fingerprint('mutable', scope.mutable),\n _get_fingerprint('flags', scope.flags),\n _get_fingerprint('rng_counts', scope.rng_counters),\n _get_fingerprint('reservations', scope.reservations),\n )\n _check_field_is_hashable((*path, 'scope'), static_scope)\n fingerprints.append(('scope', static_scope))\n fingerprint = tuple(fingerprints)\n return type(obj), fingerprint\n elif dataclasses.is_dataclass(obj):\n fingerprints = []\n for field in dataclasses.fields(obj):\n if not hasattr(obj, field.name):\n continue\n value = getattr(obj, field.name)\n value_fingerprint = _get_fingerprint(field.name, value)\n fingerprints.append((field.name, value_fingerprint))\n return type(obj), tuple(fingerprints)\n elif isinstance(obj, core.DenyList):\n return type(obj), _get_fingerprint('deny', obj.deny)\n elif isinstance(obj, dict):\n fingerprint = tuple((k, _get_fingerprint(k, v)) for k, v in obj.items())\n return fingerprint\n elif serialization.is_serializable(obj):\n state = serialization.to_state_dict(obj)\n fingerprint = _fingerprint_recursive(state, path, seen_modules)\n return type(obj), fingerprint\n elif isinstance(obj, Mapping):\n return tuple((k, _get_fingerprint(k, v)) for k, v in obj.items())\n elif isinstance(obj, Iterable):\n return tuple(_get_fingerprint(str(i), v) for i, v in enumerate(obj))\n else:\n _check_field_is_hashable(path, obj)\n return obj\n\n\ndef _check_field_is_hashable(path: tuple[str, ...], x: Any):\n \"\"\"Checks if a field is hashable.\"\"\"\n try:\n hash(x)\n except Exception as e:\n path_name = '/'.join(path)\n raise ValueError(f\"Value at '{path_name}' is not hashable: {e}\") from e\n\n\ndef decorator_lift_transform_cached(transform, class_fn, **trafo_kwargs):\n \"\"\"Decorator for lifted transform.\n\n Similar to `decorator_lift_transform` but specialized for `jit`, it reuses the\n previous transform when available to avoid retracing.\n \"\"\"\n # TODO(marcvanzee): Improve docstrings (#1977).\n # Due to the ordering of method decorators, we must wrap the class_fn\n # with the module state management wrapper first to maintain Module state\n # correctly.\n multi_scope = True\n\n if isinstance(class_fn, tuple):\n class_fns = class_fn\n else:\n class_fns = (class_fn,)\n prewrapped_fns = [wrap_method_once(class_fn) for class_fn in class_fns]\n trafo_fn = None\n\n @functools.wraps(prewrapped_fns[0])\n def wrapped_fn(self: Module, *args, **kwargs):\n nonlocal trafo_fn\n state = self._state.export()\n\n # increment rng counters for all rngs in scope\n with fork_rngs(self):\n # make a scope-function to transform\n def core_fn(\n prewrapped_fn,\n class_fn,\n scopes,\n module_hash,\n *args,\n **kwargs,\n ):\n # self = hash_key.obj\n self: Module = module_hash.module\n if not multi_scope:\n scopes = [scopes]\n cloned, args, kwargs = set_module_scopes(self, args, kwargs, scopes)\n object.__setattr__(cloned, '_state', state.export())\n res = prewrapped_fn(cloned, *args, **kwargs)\n self._state.reimport(cloned._state)\n _test_transformed_return_values(\n res, getattr(class_fn, '__name__', None)\n )\n return res\n\n core_fns = [\n functools.wraps(class_fn)(\n functools.partial(core_fn, prewrapped_fn, class_fn)\n )\n for prewrapped_fn, class_fn in zip(prewrapped_fns, class_fns)\n ]\n\n # here we apply the given lifting transform to the scope-ingesting fn\n if trafo_fn is None:\n trafo_fn = transform(*core_fns, **trafo_kwargs)\n\n module_scopes, args, kwargs = get_module_scopes(self, args, kwargs)\n\n if not multi_scope:\n if len(module_scopes) != 1:\n # TODO(levskaya): transforms like jvp & vjp have args that follow the\n # pytree structure of scopes. The user doesn't explicitly control shared\n # modules passed as arguments to methods or as attributes to Module\n # constructors. Therefore, there is no obvious API for specifying\n # arguments per lifted Module.\n raise NotImplementedError(\n 'This transform does not yet support'\n ' Modules that include other Modules passed as arguments.'\n )\n module_scopes = module_scopes[0]\n\n # get a hashable proxy object for the Module\n hash_key = _HashableProxy.from_module(self)\n\n return trafo_fn(module_scopes, hash_key, *args, **kwargs)\n\n return wrapped_fn\n\n\n@contextlib.contextmanager\ndef fork_rngs(module: Module):\n \"\"\"Context manager to fork rngs in a module.\"\"\"\n if module.scope is None:\n yield\n return\n\n current_rngs = module.scope.rngs.copy()\n module.scope.rngs = {\n name: LazyRng.create(module.make_rng(name)) for name in current_rngs\n }\n\n try:\n yield\n finally:\n module.scope.rngs = current_rngs\n\n\ndef module_class_lift_transform_cached(\n transform, module_class, methods=None, **trafo_kwargs\n):\n \"\"\"Module class lift transform.\"\"\"\n # TODO(marcvanzee): Improve docstrings (#1977).\n # TODO(levskaya): find nicer argument convention for multi-method case?\n trafo_args = ()\n\n # Prepare per-method transform args, kwargs.\n if methods is None:\n # Default case, just transform __call__\n class_trafo_args = {'__call__': (trafo_args, trafo_kwargs)}\n elif isinstance(methods, (list, tuple)):\n # Transform every method in methods with given args, kwargs.\n class_trafo_args = {m: (trafo_args, trafo_kwargs) for m in methods}\n elif isinstance(methods, dict):\n # Pass different trafo args per each method.\n class_trafo_args = {k: ((), v) for k, v in methods.items()}\n else:\n raise ValueError(\n 'transform methods argument must be None, tuple, list, or dict.'\n )\n\n # Handle partially initialized module class constructors.\n if isinstance(module_class, functools.partial) and issubclass(\n module_class.func, Module\n ):\n partial_object = module_class\n module_class = module_class.func\n else:\n partial_object = None\n\n def create_trans_fn(fn_name, fn_trafo_args):\n # get existing unbound method from class\n fn = getattr(module_class, fn_name)\n trafo_args, trafo_kwargs = fn_trafo_args\n trafo_fn = None\n\n # we need to create a scope-function from our class for the given method\n @functools.wraps(fn)\n def wrapped_fn(self: Module, *args, **kwargs):\n assert self.scope is not None\n nonlocal trafo_fn\n state = self._state.export()\n\n # increment rng counters for all rngs in scope\n with fork_rngs(self):\n # make a scope-function to transform\n def core_fn(scopes, module_hash, *args, **kwargs):\n self: Module = module_hash.module\n # make a clone of self using its arguments\n attrs = {\n f.name: getattr(self, f.name)\n for f in dataclasses.fields(self)\n if f.name != 'parent' and f.init\n }\n # we reference module_class, not self.__class__ to avoid infinite loop\n cloned = module_class(parent=None, **attrs)\n cloned, args, kwargs = set_module_scopes(cloned, args, kwargs, scopes)\n object.__setattr__(cloned, '_state', state.export())\n res = fn(cloned, *args, **kwargs)\n self._state.reimport(cloned._state)\n _test_transformed_return_values(res, fn_name)\n return res\n\n # here we apply the given lifting transform to the scope-ingesting fn\n trafo_fn = trafo_fn or transform(core_fn, *trafo_args, **trafo_kwargs)\n module_scopes, args, kwargs = get_module_scopes(self, args, kwargs)\n\n # get a hash for the Module by using its repr as a proxy\n hash_key = _HashableProxy.from_module(self)\n\n ret = trafo_fn(module_scopes, hash_key, *args, **kwargs)\n return ret\n\n return wrapped_fn\n\n transformed_fns = {\n fn_name: create_trans_fn(fn_name, fn_trafo_args)\n for fn_name, fn_trafo_args in class_trafo_args.items()\n }\n # construct new dynamic class w. transformed methods\n transformed_cls = type(\n transform.__name__.capitalize() + module_class.__name__,\n (module_class,),\n transformed_fns,\n )\n # Handle partially initialized module class constructors.\n if partial_object is not None:\n transformed_cls = functools.partial(\n transformed_cls, *partial_object.args, **partial_object.keywords\n )\n return transformed_cls\n\n\n# Utility to wrap a class or to use as decorator in def of class method.\n# -----------------------------------------------------------------------------\n\nTransformTarget = Union[type[Module], Callable[..., Any]]\nTarget = TypeVar('Target', bound=TransformTarget)\n\n\ndef _is_module_class(target: TransformTarget) -> bool:\n return (\n inspect.isclass(target)\n and issubclass(target, Module)\n or (isinstance(target, functools.partial))\n and _is_module_class(target.func)\n )\n\n\ndef lift_transform(\n transform, target, *trafo_args, methods=None, **trafo_kwargs\n):\n \"\"\"Applies to class or as a decorator on class fns.\"\"\"\n # TODO(marcvanzee): Improve docstrings (#1977).\n if _is_module_class(target):\n return module_class_lift_transform(\n transform, target, *trafo_args, methods=methods, **trafo_kwargs\n )\n # we presume this is being used as a function decorator in class definition\n elif callable(target) and not isinstance(target, Module):\n return decorator_lift_transform(\n transform, target, *trafo_args, **trafo_kwargs\n )\n else:\n raise errors.TransformTargetError(target)\n\n\ndef lift_transform_cached(\n transform, target, *trafo_args, methods=None, **trafo_kwargs\n):\n \"\"\"Applies to class or as a decorator on class fns.\"\"\"\n # TODO(marcvanzee): Improve docstrings (#1977).\n if _is_module_class(target):\n return module_class_lift_transform_cached(\n transform, target, *trafo_args, methods=methods, **trafo_kwargs\n )\n # we presume this is being used as a function decorator in class definition\n elif callable(target) and not isinstance(target, Module):\n return decorator_lift_transform_cached(\n transform, target, *trafo_args, **trafo_kwargs\n )\n else:\n raise errors.TransformTargetError(target)\n\n\ndef lift_direct_transform(\n transform: Callable[..., Any],\n targets: tuple[Callable[..., Any], ...],\n mdl: Module,\n *args,\n multi_scope=True,\n **kwargs,\n):\n \"\"\"Lift direct transform.\"\"\"\n # TODO(marcvanzee): Improve docstrings (#1977).\n for target in targets:\n if _is_module_class(target):\n raise ValueError(\n f'The {transform.__name__} transform can only be applied on a Module'\n ' method. That is function that takes a Module instance as its first'\n ' arg.'\n )\n elif not callable(target):\n raise ValueError('transform target must be callable')\n # normalize self.foo bound methods to class.foo unbound methods.\n targets = tuple(_get_unbound_fn(target) for target in targets)\n aug_transform = lambda *fns: functools.partial(transform, *fns)\n return decorator_lift_transform(\n aug_transform, targets, multi_scope=multi_scope\n )(mdl, *args, **kwargs)\n\n\ndef vmap(\n target: Target,\n variable_axes: Mapping[CollectionFilter, InOutAxis] = FrozenDict(),\n split_rngs: Mapping[PRNGSequenceFilter, bool] = FrozenDict(),\n in_axes=0,\n out_axes=0,\n axis_size: int | None = None,\n axis_name: str | None = None,\n spmd_axis_name: str | None = None,\n metadata_params: Mapping[Any, Any] = {},\n methods=None,\n) -> Target:\n \"\"\"A lifted version of ``jax.vmap``.\n\n See ``jax.vmap`` for the unlifted batch transform in Jax.\n\n ``vmap`` can be used to add a batch axis to a ``Module``.\n For example we could create a version of ``Dense`` with\n a batch axis that does not share parameters::\n\n >>> import flax.linen as nn\n >>> BatchDense = nn.vmap(\n ... nn.Dense,\n ... in_axes=0, out_axes=0,\n ... variable_axes={'params': 0},\n ... split_rngs={'params': True})\n\n By using ``variable_axes={'params': 0}``, we indicate that the\n parameters themselves are mapped over and therefore not shared along\n the mapped axis. Consequently, we also split the 'params' RNG,\n otherwise the parameters would be initialized identically along\n the mapped axis.\n\n Similarly, ``vmap`` could be used to add a batch axis with parameter\n sharing::\n\n >>> import flax.linen as nn\n >>> BatchDense = nn.vmap(\n ... nn.Dense,\n ... in_axes=0, out_axes=0,\n ... variable_axes={'params': None},\n ... split_rngs={'params': False})\n\n Here we use ``variable_axes={'params': None}`` to indicate the parameter\n variables are shared along the mapped axis. Consequently, the 'params'\n RNG must also be shared.\n\n Args:\n target: a ``Module`` or a function taking a ``Module`` as its first\n argument.\n variable_axes: the variable collections that are lifted into the batching\n transformation. Use ``None`` to indicate a broadcasted collection or an\n integer to map over an axis. For example, passing in\n ``variable_axes={'params': None}`` will indicate that the\n parameter variables should be shared along the mapped axis.\n split_rngs: Split PRNG sequences will be different for each index of the\n batch dimension. Unsplit PRNGs will be broadcasted.\n in_axes: Specifies the mapping of the input arguments (see ``jax.vmap``).\n out_axes: Specifies the mapping of the return value (see ``jax.vmap``).\n axis_size: Specifies the size of the batch axis. This only needs to be\n specified if it cannot be derived from the input arguments.\n axis_name: Specifies a name for the batch axis. Can be used together with\n parallel reduction primitives (e.g. ``jax.lax.pmean``, ``jax.lax.ppermute``,\n etc.). Note, this is only used for pmap and shard map. For SPMD jit, you\n do not need to manually synchronize. Just make sure that the axes are\n correctly annotated and XLA:SPMD will insert the necessary collectives.\n methods: If ``target`` is a ``Module``, the methods of ``Module`` to vmap over.\n spmd_axis_name: Axis name added to any pjit sharding constraints appearing\n in ``fn``. See also\n https://github.com/google/flax/blob/main/flax/linen/partitioning.py.\n metadata_params: arguments dict passed to AxisMetadata instances in the\n variable tree.\n\n Returns:\n A batched/vectorized version of ``target``, with the same arguments but with\n extra axes at positions indicated by ``in_axes``, and the same return value,\n but with extra axes at positions indicated by ``out_axes``.\n \"\"\"\n return lift_transform(\n lift.vmap,\n target,\n variable_axes,\n split_rngs,\n methods=methods,\n in_axes=in_axes,\n out_axes=out_axes,\n axis_size=axis_size,\n axis_name=axis_name,\n metadata_params=metadata_params,\n spmd_axis_name=spmd_axis_name,\n )\n\n\ndef jit(\n target: Target,\n variables: CollectionFilter = True,\n rngs: PRNGSequenceFilter = True,\n static_argnums: int | Iterable[int] = (),\n static_argnames: str | Iterable[str] = (),\n donate_argnums: int | Iterable[int] = (),\n device=None,\n backend: str | None = None,\n methods=None,\n) -> Target:\n \"\"\"Lifted version of ``jax.jit``.\n\n Args:\n target: a ``Module`` or a function taking a ``Module`` as its first\n argument.\n variables: The variable collections that are lifted. By default all\n collections are lifted.\n rngs: The PRNG sequences that are lifted. By default all PRNG sequences are\n lifted.\n static_argnums: An int or collection of ints specifying which positional\n arguments to treat as static (compile-time constant). Operations that only\n depend on static arguments will be constant-folded in Python (during\n tracing), and so the corresponding argument values can be any Python\n object. Static arguments should be hashable, meaning both ``__hash__`` and\n ``__eq__`` are implemented, and immutable. Calling the jitted function\n with different values for these constants will trigger recompilation. If\n the jitted function is called with fewer positional arguments than\n indicated by ``static_argnums`` then an error is raised. Arguments that\n are not arrays or containers thereof must be marked as static. Defaults to\n ().\n static_argnames: An optional string or collection of strings specifying\n which named arguments to treat as static (compile-time constant). See the\n comment on ``static_argnums`` for details. If not provided but\n ``static_argnums`` is set, the default is based on calling\n ``inspect.signature(fun)`` to find corresponding named arguments.\n donate_argnums: Specify which arguments are \"donated\" to the computation. It\n is safe to donate arguments if you no longer need them once the\n computation has finished. In some cases XLA can make use of donated\n buffers to reduce the amount of memory needed to perform a computation,\n for example recycling one of your input buffers to store a result. You\n should not reuse buffers that you donate to a computation, JAX will raise\n an error if you try to.\n device: This is an experimental feature and the API is likely to change.\n Optional, the Device the jitted function will run on. (Available devices\n can be retrieved via :py:func:`jax.devices`.) The default is inherited\n from XLA's DeviceAssignment logic and is usually to use\n ``jax.devices()[0]``.\n backend: a string representing the XLA backend: ``'cpu'``, ``'gpu'``, or\n ``'tpu'``.\n methods: If ``target`` is a ``Module``, the methods of ``Module`` to jit.\n\n Returns:\n A wrapped version of target, set up for just-in-time compilation.\n \"\"\"\n # TODO(marcvanzee): Improve docstrings (#1977).\n return lift_transform_cached(\n lift.jit,\n target,\n variables=variables,\n rngs=rngs,\n static_argnums=static_argnums,\n static_argnames=static_argnames,\n donate_argnums=donate_argnums,\n device=device,\n backend=backend,\n methods=methods,\n )\n\n\ndef checkpoint(\n target: Target,\n variables: CollectionFilter = True,\n rngs: PRNGSequenceFilter = True,\n concrete: bool = False,\n prevent_cse: bool = True,\n static_argnums: int | tuple[int, ...] = (),\n policy: Callable[..., bool] | None = None,\n methods=None,\n) -> Target:\n \"\"\"Lifted version of ``jax.checkpoint``.\n\n Checkpointing is a technique for reducing memory usage by recomputing\n activations during backpropagation. When training large models, it can be\n helpful to checkpoint parts of the model to trade off memory usage for\n additional computation.\n\n Example::\n\n >>> import jax\n >>> import jax.numpy as jnp\n >>> import flax.linen as nn\n ...\n >>> class CheckpointedMLP(nn.Module):\n ... @nn.checkpoint\n ... @nn.compact\n ... def __call__(self, x):\n ... x = nn.Dense(128)(x)\n ... x = nn.relu(x)\n ... x = nn.Dense(1)(x)\n ... return x\n ...\n >>> model = CheckpointedMLP()\n >>> variables = model.init(jax.random.key(0), jnp.ones((1, 16)))\n\n This function is aliased to ``remat`` just like ``jax.remat``.\n\n Args:\n target: a ``Module`` or a function taking a ``Module``\n as its first argument. intermediate computations will be\n re-computed when computing gradients for the target.\n variables: The variable collections that are lifted. By default all\n collections are lifted.\n rngs: The PRNG sequences that are lifted. By default all PRNG sequences\n are lifted.\n concrete: Optional, boolean indicating whether ``fun`` may involve\n value-dependent Python control flow (default ``False``). Support for such\n control flow is optional, and disabled by default, because in some\n edge-case compositions with :func:`jax.jit` it can lead to some extra\n computation.\n prevent_cse: Optional, boolean indicating whether to prevent common\n subexpression elimination (CSE) optimizations in the HLO generated from\n differentiation. This CSE prevention has costs because it can foil other\n optimizations, and because it can incur high overheads on some backends,\n especially GPU. The default is True because otherwise, under a ``jit`` or\n ``pmap``, CSE can defeat the purpose of this decorator. But in some\n settings, like when used inside a ``scan``, this CSE prevention mechanism\n is unnecessary, in which case ``prevent_cse`` should be set to False.\n static_argnums: Optional, int or sequence of ints, indicates which argument\n values on which to specialize for tracing and caching purposes. Specifying\n arguments as static can avoid ConcretizationTypeErrors when tracing, but\n at the cost of more retracing overheads.\n policy: Experimental checkpoint policy, see ``jax.checkpoint``.\n methods: An optional list of method names that will be lifted, if ``methods``\n is None (default) only the ``__call__`` method will be lifted. If``target``\n is a function, ``methods`` is ignored.\n\n Returns:\n A wrapped version of ``target``. When computing gradients intermediate\n computations will be re-computed on the backward pass.\n \"\"\"\n # subtract 1 from each static_argnums because 'self' is not passed to the\n # lifted function\n static_argnums = jax.tree_util.tree_map(lambda x: x - 1, static_argnums)\n return lift_transform(\n lift.checkpoint,\n target,\n variables=variables,\n rngs=rngs,\n concrete=concrete,\n static_argnums=static_argnums,\n prevent_cse=prevent_cse,\n policy=policy,\n methods=methods,\n )\n\n\nremat = checkpoint\n\n\ndef remat_scan(\n target: Target,\n lengths: Sequence[int] | None = (),\n policy: Callable[..., bool] | None = None,\n variable_broadcast: CollectionFilter = False,\n variable_carry: CollectionFilter = False,\n variable_axes: Mapping[CollectionFilter, InOutScanAxis] = FrozenDict(\n {True: 0}\n ),\n split_rngs: Mapping[PRNGSequenceFilter, bool] = FrozenDict({True: True}),\n) -> Target:\n \"\"\"Combines remat and scan for memory efficiency and constant time compilation.\n\n ``remat_scan`` allows for constant compile times and sublinear\n memory usage with respect to model depth. At a small constant\n penalty. This is typically beneficial for very deep models.\n\n Example::\n\n >>> import flax.linen as nn\n\n >>> class BigModel(nn.Module):\n ... @nn.compact\n ... def __call__(self, x):\n ... DenseStack = nn.remat_scan(nn.Dense, lengths=(10, 10))\n ... # 100x dense with O(sqrt(N)) memory for gradient computation\n ... return DenseStack(8, name=\"dense_stack\")(x)\n\n Args:\n target: a ``Module`` or a function taking a ``Module`` as its first\n argument.\n lengths: number of loop iterations at the given level. The total number of\n iterations ``n = prod(lengths)``. each loop is rematerialized. This way the\n memory consumption is proportional to ``n^(1 / d)`` where ``d =\n len(lengths)``. Minimal memory consumptions requires tuning the lengths\n such that the same amount of memory is consumed at each level of the\n nested loop.\n policy: Experimental checkpoint policy, see ``jax.checkpoint``.\n variable_broadcast: Specifies the broadcasted variable collections. A\n broadcasted variable should not depend on any computation that cannot be\n lifted out of the loop. This is typically used to define shared parameters\n inside the fn.\n variable_carry: Specifies the variable collections that are carried through\n the loop. Mutations to these variables are carried to the next iteration\n and will be preserved when the scan finishes.\n variable_axes: the variable collections that are scanned over. Defaults to\n ``{True: 0}``.\n split_rngs: Split PRNG sequences will be different for each loop iterations.\n If split is False the PRNGs will be the same across iterations. Defaults\n to ``{True: True}``.\n\n Returns:\n A wrapped version of ``target`` that repeats itself prod(lengths) times.\n \"\"\"\n return lift_transform(\n lift.remat_scan,\n target,\n lengths=lengths,\n variable_broadcast=variable_broadcast,\n variable_carry=variable_carry,\n variable_axes=variable_axes,\n split_rngs=split_rngs,\n policy=policy,\n )\n\n\ndef scan(\n target: Target,\n variable_axes: Mapping[CollectionFilter, InOutScanAxis] = FrozenDict(),\n variable_broadcast: CollectionFilter = False,\n variable_carry: CollectionFilter = False,\n split_rngs: Mapping[PRNGSequenceFilter, bool] = FrozenDict(),\n in_axes=0,\n out_axes=0,\n length: int | None = None,\n reverse: bool = False,\n unroll: int = 1,\n data_transform: Callable[..., Any] | None = None,\n metadata_params: Mapping[Any, Any] = {},\n methods=None,\n _split_transpose: bool = False,\n check_constancy_invariants: bool = True,\n) -> Target:\n \"\"\"A lifted version of ``jax.lax.scan``.\n\n See ``jax.lax.scan`` for the unlifted scan in Jax.\n\n To improve consistency with ``vmap``, this version of scan\n uses ``in_axes`` and ``out_axes`` to determine which arguments\n are scanned over and along which axis.\n\n ``scan`` distinguishes between 3 different types of values inside the loop:\n\n #. **scan**: a value that is iterated over in a loop. All scan values must\n have the same size in the axis they are scanned over. Scanned outputs\n will be stacked along the scan axis.\n\n #. **carry**: A carried value is updated at each loop iteration. It must\n have the same shape and dtype throughout the loop.\n\n #. **broadcast**: a value that is closed over by the loop. When a variable\n is broadcasted they are typically initialized inside the loop body but\n independent of the loop variables.\n\n The ``target`` should have the signature\n ``(module, carry, *xs) -> (carry, ys)``, where ``xs`` and ``ys``\n are the scan values that go in and out of the loop.\n\n Example::\n\n >>> import flax.linen as nn\n >>> import jax\n >>> import jax.numpy as jnp\n ...\n >>> class LSTM(nn.Module):\n ... features: int\n ...\n ... @nn.compact\n ... def __call__(self, x):\n ... ScanLSTM = nn.scan(\n ... nn.LSTMCell, variable_broadcast=\"params\",\n ... split_rngs={\"params\": False}, in_axes=1, out_axes=1)\n ...\n ... lstm = ScanLSTM(self.features)\n ... input_shape = x[:, 0].shape\n ... carry = lstm.initialize_carry(jax.random.key(0), input_shape)\n ... carry, x = lstm(carry, x)\n ... return x\n ...\n >>> x = jnp.ones((4, 12, 7))\n >>> module = LSTM(features=32)\n >>> y, variables = module.init_with_output(jax.random.key(0), x)\n\n Note that when providing a function to ``nn.scan``, the scanning happens over\n all arguments starting from the third argument, as specified by ``in_axes``.\n The previous example could also be written using the functional form as::\n\n >>> class LSTM(nn.Module):\n ... features: int\n ...\n ... @nn.compact\n ... def __call__(self, x):\n ...\n ... cell = nn.LSTMCell(self.features)\n ... def body_fn(cell, carry, x):\n ... carry, y = cell(carry, x)\n ... return carry, y\n ... scan = nn.scan(\n ... body_fn, variable_broadcast=\"params\",\n ... split_rngs={\"params\": False}, in_axes=1, out_axes=1)\n ...\n ... input_shape = x[:, 0].shape\n ... carry = cell.initialize_carry(\n ... jax.random.key(0), input_shape)\n ... carry, x = scan(cell, carry, x)\n ... return x\n ...\n >>> module = LSTM(features=32)\n >>> variables = module.init(jax.random.key(0), jnp.ones((4, 12, 7)))\n\n You can also use ``scan`` to reduce the compilation time of your JAX program\n by merging multiple layers into a single scan loop, you can do this when\n you have a sequence of identical layers that you want to apply iteratively\n to an input. For example::\n\n >>> class ResidualMLPBlock(nn.Module):\n ... @nn.compact\n ... def __call__(self, x, _):\n ... h = nn.Dense(features=2)(x)\n ... h = nn.relu(h)\n ... return x + h, None\n ...\n >>> class ResidualMLP(nn.Module):\n ... n_layers: int = 4\n ...\n ... @nn.compact\n ... def __call__(self, x):\n ... ScanMLP = nn.scan(\n ... ResidualMLPBlock, variable_axes={'params': 0},\n ... variable_broadcast=False, split_rngs={'params': True},\n ... length=self.n_layers)\n ... x, _ = ScanMLP()(x, None)\n ... return x\n ...\n >>> model = ResidualMLP(n_layers=4)\n >>> variables = model.init(jax.random.key(42), jnp.ones((1, 2)))\n\n To reduce both compilation and memory usage, you can use :func:`remat_scan`\n which will in addition checkpoint each layer in the scan loop.\n\n Args:\n target: a ``Module`` or a function taking a ``Module`` as its first\n argument.\n variable_axes: the variable collections that are scanned over.\n variable_broadcast: Specifies the broadcasted variable collections. A\n broadcasted variable should not depend on any computation that cannot be\n lifted out of the loop. This is typically used to define shared parameters\n inside the fn.\n variable_carry: Specifies the variable collections that are carried through\n the loop. Mutations to these variables are carried to the next iteration\n and will be preserved when the scan finishes.\n split_rngs: Split PRNG sequences will be different for each loop iterations.\n If split is False the PRNGs will be the same across iterations.\n in_axes: Specifies the axis to scan over for the arguments. Should be a\n prefix tree of the arguments. Use ``flax.core.broadcast`` to feed an entire\n input to each iteration of the scan body.\n out_axes: Specifies the axis to scan over for the return value. Should be a\n prefix tree of the return value.\n length: Specifies the number of loop iterations. This only needs to be\n specified if it cannot be derived from the scan arguments.\n reverse: If true, scan from end to start in reverse order.\n unroll: how many scan iterations to unroll within a single iteration of a\n loop (default: 1).\n data_transform: optional function to transform raw functional-core variable\n and rng groups inside lifted scan body_fn, intended for inline SPMD\n annotations.\n metadata_params: arguments dict passed to AxisMetadata instances in the\n variable tree.\n methods: If ``target`` is a ``Module``, the methods of ``Module`` to scan over.\n _split_transpose: An experimental feature to split the transpose of a scan\n into a scan and a map, backed by an experimental Jax lax.scan() feature.\n check_constancy_invariants: If true, the scan will verify that the\n broadcast constants are true loop invariants, and further supports\n broadcast function (non-carry) outputs. This requires an extra jax\n tracing step however, so setting to false can reduce trace time on larger\n models.\n\n Returns:\n The scan function with the signature ``(module, carry, *xs) -> (carry,\n ys)``, where ``xs`` and ``ys`` are the scan values that go in and out of\n the loop.\n \"\"\"\n return lift_transform(\n lift.scan,\n target,\n variable_axes=variable_axes,\n variable_broadcast=variable_broadcast,\n variable_carry=variable_carry,\n split_rngs=split_rngs,\n in_axes=in_axes,\n out_axes=out_axes,\n length=length,\n reverse=reverse,\n unroll=unroll,\n _split_transpose=_split_transpose,\n data_transform=data_transform,\n metadata_params=metadata_params,\n methods=methods,\n check_constancy_invariants=check_constancy_invariants,\n )\n\n\ndef map_variables(\n target: Target,\n mapped_collections: CollectionFilter = True,\n trans_in_fn: Callable[..., Any] = lift.id_fn,\n trans_out_fn: Callable[..., Any] = lift.id_fn,\n init: bool = False,\n mutable: bool = False,\n rngs: PRNGSequenceFilter = True,\n variables: CollectionFilter = True,\n methods=None,\n) -> Target:\n \"\"\"Map Variables inside a module.\n\n ``map_variables`` can be used to transform the variables inside a module\n both before and after the module is applied. This is useful among other\n things for masking the weights of a module without having to modify the\n module itself.\n\n Example::\n\n >>> import jax\n >>> import jax.numpy as jnp\n >>> import flax.linen as nn\n ...\n >>> class CausalDense(nn.Module):\n ... '''A dense layer that masks the weights such that the output is\n ... causal, i.e. output i only depends on input <= i.\n ... '''\n ... features: int\n ...\n ... def apply_mask(self, variables):\n ... return (jax.tree_util.tree_map(jnp.triu, variables)\n ... if not self.is_initializing() else variables)\n ...\n ... def setup(self):\n ... # temporary class\n ... _CausalDense = nn.map_variables(\n ... nn.Dense, 'params', self.apply_mask, init=self.is_initializing())\n ... self.dense = _CausalDense(features=self.features, use_bias=False)\n ...\n ... def __call__(self, x):\n ... return self.dense(x)\n ...\n >>> module = CausalDense(features=5)\n >>> variables = module.init(jax.random.key(0), jnp.ones((1, 5)))\n\n Args:\n target: the module or function to be transformed.\n mapped_collections: the collection(s) to be transformed.\n trans_in_fn: modifies the variables before applying the module or function.\n trans_out_fn: modifies the variables after applying the module or function,\n it is only applied if either ``init`` or ``mutable`` are not False.\n init: If True, variables are initialized before transformation.\n mutable: If True, the mapped variable collections will be mutable.\n rngs: PRNGSequences added to the transformed scope (default: all).\n variables: Additional Variable collections added to the transformed scope.\n Besides those specified by ``target`` (default: all).\n methods: If ``target`` is a ``Module``, the methods of ``Module`` to map\n variables for.\n\n Returns:\n a wrapped version of ``target`` that will map the specified collections.\n \"\"\"\n\n return lift_transform(\n lift.map_variables,\n target,\n mapped_collections,\n trans_in_fn,\n trans_out_fn,\n init,\n mutable,\n rngs,\n variables,\n methods=methods,\n )\n\n\ndef vjp(\n fn: Callable[..., Any],\n mdl: Module,\n *primals,\n has_aux: bool = False,\n reduce_axes=(),\n vjp_variables: CollectionFilter = 'params',\n variables: CollectionFilter = True,\n rngs: PRNGSequenceFilter = True,\n multi_scope: bool = False,\n):\n \"\"\"A lifted version of ``jax.vjp``.\n\n See ``jax.vjp`` for the unlifted vector-Jacobian product (backward gradient).\n\n Note that a gradient is returned for all variables in the collections\n specified by ``vjp_variables``. However, the backward function only expects\n a cotangent for the return value of ``fn``. If variables require a co-tangent\n as well they can be returned from ``fn`` using ``Module.variables``.\n\n Example::\n\n >>> import flax.linen as nn\n >>> import jax.numpy as jnp\n\n >>> class LearnScale(nn.Module):\n ... @nn.compact\n ... def __call__(self, x, y):\n ... p = self.param('scale', nn.initializers.zeros_init(), ())\n ... return p * x * y\n\n >>> class Foo(nn.Module):\n ... @nn.compact\n ... def __call__(self, x, y):\n ... z, bwd = nn.vjp(lambda mdl, x, y: mdl(x, y), LearnScale(), x, y)\n ... params_grad, x_grad, y_grad = bwd(jnp.ones(z.shape))\n ... return z, params_grad, x_grad, y_grad\n\n Args:\n fn: Function to be differentiated. Its arguments should be arrays, scalars,\n or standard Python containers of arrays or scalars. It should return an\n array, scalar, or standard Python container of arrays or scalars. It will\n receive the scope and primals as arguments.\n mdl: The module of which the variables will be differentiated.\n *primals: A sequence of primal values at which the Jacobian of ``fn``\n should be evaluated. The length of ``primals`` should be equal to the\n number of positional parameters to ``fn``. Each primal value should be a\n tuple of arrays, scalar, or standard Python containers thereof.\n has_aux: Optional, bool. Indicates whether ``fn`` returns a pair where the\n first element is considered the output of the mathematical function to be\n differentiated and the second element is auxiliary data. Default ``False``.\n vjp_variables: The vjpfun will return a cotangent vector for all\n variable collections specified by this filter.\n variables: other variables collections that are available inside ``fn`` but\n do not receive a cotangent.\n rngs: the prngs that are available inside ``fn``.\n multi_scope: for Modules containing multiple scopes from outside modules passed in,\n allow for variable gradients to be returned for multiple scopes instead of erroring.\n Returns:\n If ``has_aux`` is ``False``, returns a ``(primals_out, vjpfun)`` pair, where\n ``primals_out`` is ``fn(*primals)``.\n ``vjpfun`` is a function from a cotangent vector with the same shape as\n ``primals_out`` to a tuple of cotangent vectors with the same shape as\n ``primals``, representing the vector-Jacobian product of ``fn`` evaluated at\n ``primals``. If ``has_aux`` is ``True``, returns a\n ``(primals_out, vjpfun, aux)`` tuple where ``aux`` is the auxiliary data\n returned by ``fn``.\n \"\"\"\n if reduce_axes:\n raise NotImplementedError('reduce_axes argument to vjp is deprecated')\n del reduce_axes\n\n return lift_direct_transform(\n lift.vjp,\n (fn,),\n mdl,\n *primals,\n multi_scope=multi_scope,\n has_aux=has_aux,\n vjp_variables=vjp_variables,\n variables=variables,\n rngs=rngs,\n )\n\n\ndef value_and_grad(\n fn: Callable[..., Any],\n mdl: Module,\n *primals,\n has_aux: bool = False,\n reduce_axes=(),\n variables: CollectionFilter = True,\n rngs: PRNGSequenceFilter = True,\n):\n \"\"\"A limited, lifted equivalent of ``jax.value_and_grad``.\n\n Note that for this convenience function, gradients are only calculated for\n the function inputs, and not with respect to any module variables. The\n target function must return a scalar-valued output. For a more general\n lifted vjp, see ``nn.vjp`` for the lifted vector-Jacobian product.\n\n Example::\n\n class LearnScale(nn.Module):\n @nn.compact\n def __call__(self, x, y):\n p = self.param('scale', nn.initializers.zeros_init(), ())\n return p * x * y\n\n class Foo(nn.Module):\n @nn.compact\n def __call__(self, x, y):\n z, (x_grad, y_grad) = nn.value_and_grad(\n lambda mdl, x, y: mdl(x, y), LearnScale(), x, y)\n return z, x_grad, y_grad\n\n Args:\n fn: Function to be differentiated. Its arguments should be arrays, scalars,\n or standard Python containers of arrays or scalars. It should return an\n array, scalar, or standard Python container of arrays or scalars. It will\n receive the scope and primals as arguments.\n mdl: The module of which the variables will be differentiated.\n *primals: A sequence of primal values at which the Jacobian of ``fn``\n should be evaluated. The length of ``primals`` should be equal to the\n number of positional parameters to ``fn``. Each primal value should be a\n tuple of arrays, scalar, or standard Python containers thereof.\n has_aux: Optional, bool. Indicates whether ``fn`` returns a pair where the\n first element is considered the output of the mathematical function to be\n differentiated and the second element is auxiliary data. Default ``False``.\n variables: variables collections that are available inside ``fn`` but\n do not receive a cotangent.\n rngs: the prngs that are available inside ``fn``.\n Returns:\n If ``has_aux`` is ``False``, returns a ``primals_out, grads`` pair, where\n ``primals_out`` is ``fn(*primals)``. ``grads`` are the gradients for the\n corresponding primals and do not include the gradients for module variables.\n If ``has_aux`` is ``True``, returns a\n ``(primals_out, aux), grads`` tuple where ``aux`` is the auxiliary data\n returned by ``fn``.\n \"\"\"\n if reduce_axes:\n raise NotImplementedError(\n 'reduce_axes argument to value_and_grad is deprecated')\n del reduce_axes\n\n grad_partial = functools.partial(\n lift_direct_transform,\n lift.value_and_grad,\n (fn,),\n mdl,\n *primals,\n has_aux=has_aux,\n variables=variables,\n rngs=rngs,\n )\n\n if has_aux:\n out, aux, argument_grads = grad_partial()\n if out.shape != ():\n raise ValueError(\n 'grad can only work on functions with '\n f'scalar-valued outputs. out shape={out.shape}'\n )\n return (out, aux), argument_grads\n else:\n out, argument_grads = grad_partial()\n if out.shape != ():\n raise ValueError(\n 'grad can only work on functions with '\n f'scalar-valued outputs. out shape={out.shape}'\n )\n return out, argument_grads\n\n\ndef grad(\n fn: Callable[..., Any],\n mdl: Module,\n *primals,\n has_aux: bool = False,\n reduce_axes=(),\n variables: CollectionFilter = True,\n rngs: PRNGSequenceFilter = True,\n):\n \"\"\"A limited, lifted equivalent of ``jax.grad``.\n\n Note that for this convenience function, gradients are only calculated for\n the function inputs, and not with respect to any module variables. The\n target function must return a scalar-valued output. For a more general\n lifted vjp, see ``nn.vjp`` for the lifted vector-Jacobian product.\n\n Example::\n\n class LearnScale(nn.Module):\n @nn.compact\n def __call__(self, x, y):\n p = self.param('scale', nn.initializers.zeros_init(), ())\n return p * x * y\n\n class Foo(nn.Module):\n @nn.compact\n def __call__(self, x, y):\n x_grad, y_grad = nn.grad(\n lambda mdl, x, y: mdl(x, y), LearnScale(), x, y)\n return x_grad, y_grad\n\n Args:\n fn: Function to be differentiated. Its arguments should be arrays, scalars,\n or standard Python containers of arrays or scalars. It should return an\n array, scalar, or standard Python container of arrays or scalars. It will\n receive the scope and primals as arguments.\n mdl: The module of which the variables will be differentiated.\n *primals: A sequence of primal values at which the Jacobian of ``fn``\n should be evaluated. The length of ``primals`` should be equal to the\n number of positional parameters to ``fn``. Each primal value should be a\n tuple of arrays, scalar, or standard Python containers thereof.\n has_aux: Optional, bool. Indicates whether ``fn`` returns a pair where the\n first element is considered the output of the mathematical function to be\n differentiated and the second element is auxiliary data. Default ``False``.\n variables: variables collections that are available inside ``fn`` but\n do not receive a cotangent.\n rngs: the prngs that are available inside ``fn``.\n Returns:\n If ``has_aux`` is ``False``, returns ``grads``, where ``grads`` are the\n gradients for the corresponding primals and do not include the gradients\n for module variables.\n If ``has_aux`` is ``True``, returns a\n ``(grads, aux)`` tuple where ``aux`` is the auxiliary data\n returned by ``fn``.\n \"\"\"\n if reduce_axes:\n raise NotImplementedError('reduce_axes argument to grad is deprecated')\n del reduce_axes\n\n value_and_grad_partial = functools.partial(\n value_and_grad,\n fn,\n mdl,\n *primals,\n has_aux=has_aux,\n variables=variables,\n rngs=rngs,\n )\n\n if has_aux:\n (_, aux), argument_grads = value_and_grad_partial()\n return argument_grads, aux\n else:\n _, argument_grads = value_and_grad_partial()\n return argument_grads\n\n\ndef jvp(\n fn: Callable[..., Any],\n mdl: Module,\n primals,\n tangents,\n variable_tangents,\n variables: CollectionFilter = True,\n rngs: PRNGSequenceFilter = True,\n) -> tuple[Any, Callable[..., Any]] | tuple[Any, Callable[..., Any], Any]:\n \"\"\"A lifted version of ``jax.jvp``.\n\n See ``jax.jvp`` for the unlifted Jacobian-vector product (forward gradient).\n\n Note that no tangents are returned for variables. When variable tangents\n are required their value should be returned explicitly by ``fn``\n using ``Module.variables``::\n\n >>> import flax.linen as nn\n >>> import jax.numpy as jnp\n\n >>> class LearnScale(nn.Module):\n ... @nn.compact\n ... def __call__(self, x):\n ... p = self.param('test', nn.initializers._init(), ())\n ... return p * x\n\n >>> class Foo(nn.Module):\n ... @nn.compact\n ... def __call__(self, x):\n ... scale = LearnScale()\n ... vars_t = jax.tree_util.tree_map(jnp.ones_like,\n ... scale.variables.get('params', {}))\n ... _, out_t = nn.jvp(\n ... lambda mdl, x: mdl(x), scale, (x,), (jnp.zeros_like(x),),\n ... variable_tangents={'params': vars_t})\n ... return out_t\n\n Example::\n\n >>> def learn_scale(scope, x):\n ... p = scope.param('scale', nn.initializers.zeros_init(), ())\n ... return p * x\n\n >>> def f(scope, x):\n ... vars_t = jax.tree_util.tree_map(jnp.ones_like, scope.variables().get('params', {}))\n ... x, out_t = lift.jvp(\n ... learn_scale, scope, (x,), (jnp.zeros_like(x),),\n ... variable_tangents={'params': vars_t})\n ... return out_t\n\n Args:\n fn: Function to be differentiated. Its arguments should be arrays, scalars,\n or standard Python containers of arrays or scalars. It should return an\n array, scalar, or standard Python container of arrays or scalars. It will\n receive the scope and primals as arguments.\n mdl: The module of which the variables will be differentiated.\n primals: The primal values at which the Jacobian of ``fun`` should be\n evaluated. Should be either a tuple or a list of arguments,\n and its length should be equal to the number of positional parameters of\n ``fun``.\n tangents: The tangent vector for which the Jacobian-vector product should be\n evaluated. Should be either a tuple or a list of tangents, with the same\n tree structure and array shapes as ``primals``.\n variable_tangents: A dict or PyTree fo dicts with the same structure as\n scopes. Each entry in the dict specifies the tangents for a variable\n collection. Not specifying a collection in variable_tangents is\n equivalent to passing a zero vector as the tangent.\n variables: other variables collections that are available in ``fn`` but\n do not receive a tangent.\n rngs: the prngs that are available inside ``fn``.\n\n Returns:\n A ``(primals_out, tangents_out)`` pair, where ``primals_out`` is\n ``fun(*primals)``, and ``tangents_out`` is the Jacobian-vector product of\n ``function`` evaluated at ``primals`` with ``tangents``. The\n ``tangents_out`` value has the same Python tree structure and shapes as\n ``primals_out``.\n \"\"\"\n return lift_direct_transform(\n lift.jvp,\n (fn,),\n mdl,\n primals,\n tangents,\n variable_tangents,\n multi_scope=False,\n variables=variables,\n rngs=rngs,\n )\n\n\nModuleT = TypeVar('ModuleT', bound=Module)\nC = TypeVar('C')\n\n\ndef while_loop(\n cond_fn: Callable[[ModuleT, C], bool],\n body_fn: Callable[[ModuleT, C], C],\n mdl: ModuleT,\n init: C,\n carry_variables: CollectionFilter = False,\n broadcast_variables: CollectionFilter = True,\n split_rngs: Mapping[PRNGSequenceFilter, bool] = FrozenDict(),\n) -> C:\n \"\"\"Lifted version of jax.lax.while_loop.\n\n The lifted scope is passed to ``cond_fn`` and ``body_fn``.\n Broadcasted variables are immutable. The carry variable are\n mutable but cannot change shape and dtype.\n This also means you cannot initialize variables inside\n the body. Consider calling ``body_fn`` once manually before\n calling ``while_loop`` if variable initialization is required.\n\n Example::\n\n >>> import flax.linen as nn\n >>> import jax, jax.numpy as jnp\n\n >>> class WhileLoopExample(nn.Module):\n ... @nn.compact\n ... def __call__(self, x):\n ... def cond_fn(mdl, c):\n ... return mdl.variables['state']['acc'] < 10\n ... def body_fn(mdl, c):\n ... acc = mdl.variable('state', 'acc', lambda: jnp.array(0))\n ... acc.value += 1\n ... y = nn.Dense(c.shape[-1])(c)\n ... return y\n ... c = x\n ... if self.is_mutable_collection('params'):\n ... return body_fn(self, c)\n ... else:\n ... return nn.while_loop(cond_fn, body_fn, self, c,\n ... carry_variables='state')\n\n >>> k = jax.random.key(0)\n >>> x = jnp.ones((2, 2))\n >>> initial_vars = WhileLoopExample().init(k, x)\n >>> result, state = WhileLoopExample().apply(initial_vars, x, mutable=['state'])\n\n Args:\n cond_fn: Should return True as long as the loop should continue.\n body_fn: The body of the while loop.\n mdl: The Module which should be lifted into the loop.\n init: The initial state passed to the loop\n carry_variables: collections that are carried through the loop\n and are therefore mutable (default: none).\n broadcast_variables: collections that are closed over and are\n therefore read-only (default: all collections)\n split_rngs: Split PRNG sequences will be different for each loop iterations.\n If split is False the PRNGs will be the same across iterations.\n Returns:\n The final state after executing the while loop.\n \"\"\"\n return lift_direct_transform(\n lift.while_loop,\n (cond_fn, body_fn),\n mdl,\n init,\n carry_variables,\n broadcast_variables,\n split_rngs,\n )\n\n\ndef _cond_wrapper(t_fn, f_fn, scope, pred, *ops, variables, rngs):\n return lift.cond(\n pred, t_fn, f_fn, scope, *ops, variables=variables, rngs=rngs\n )\n\n\ndef cond(\n pred: Any,\n true_fun: Callable[..., C],\n false_fun: Callable[..., C],\n mdl: Module,\n *operands,\n variables: CollectionFilter = True,\n rngs: PRNGSequenceFilter = True,\n) -> C:\n \"\"\"Lifted version of ``jax.lax.cond``.\n\n The returned values from ``true_fun`` and ``false_fun``\n must have the same Pytree structure, shapes, and dtypes.\n The variables created or updated inside the\n branches must also have the same structure.\n Note that this constraint is violated when\n creating variables or submodules in only one branch.\n Because initializing variables in just one branch\n causes the parameter structure to be different.\n\n Example::\n\n >>> import flax.linen as nn\n\n >>> class CondExample(nn.Module):\n ... @nn.compact\n ... def __call__(self, x, pred):\n ... self.variable('state', 'true_count', lambda: 0)\n ... self.variable('state', 'false_count', lambda: 0)\n ... def true_fn(mdl, x):\n ... mdl.variable('state', 'true_count').value += 1\n ... return nn.Dense(2, name='dense')(x)\n ... def false_fn(mdl, x):\n ... mdl.variable('state', 'false_count').value += 1\n ... return -nn.Dense(2, name='dense')(x)\n ... return nn.cond(pred, true_fn, false_fn, self, x)\n\n Args:\n pred: determines if true_fun or false_fun is evaluated.\n true_fun: The function evaluated when ``pred`` is ``True``.\n The signature is (module, *operands) -> T.\n false_fun: The function evaluated when ``pred`` is ``False``.\n The signature is (module, *operands) -> T.\n mdl: A Module target to pass.\n *operands: The arguments passed to ``true_fun`` and ``false_fun``\n variables: The variable collections passed to the conditional\n branches (default: all)\n rngs: The PRNG sequences passed to the conditionals (default: all)\n Returns:\n The result of the evaluated branch (``true_fun`` or ``false_fun``).\n \"\"\"\n return lift_direct_transform(\n _cond_wrapper,\n (true_fun, false_fun),\n mdl,\n pred,\n *operands,\n variables=variables,\n rngs=rngs,\n )\n\n\ndef _switch_wrapper(*args, variables, rngs, n_branches):\n # first n_branches arguments are branches.\n # then scope, index, and the rest are *operands\n branches = args[:n_branches]\n scope, index, *operands = args[n_branches:]\n return lift.switch(\n index, branches, scope, *operands, variables=variables, rngs=rngs\n )\n\n\ndef switch(\n index: Any,\n branches: Sequence[Callable[..., C]],\n mdl: Module,\n *operands,\n variables: CollectionFilter = True,\n rngs: PRNGSequenceFilter = True,\n) -> C:\n \"\"\"Lifted version of ``jax.lax.switch``.\n\n The returned values from ``branches``\n must have the same Pytree structure, shapes, and dtypes.\n The variables created or updated inside the\n branches must also have the same structure.\n Note that this constraint is violated when\n creating variables or submodules in only one branch.\n Because initializing variables in just one branch\n causes the parameter structure to be different.\n\n Example::\n\n >>> import flax.linen as nn\n\n >>> class SwitchExample(nn.Module):\n ... @nn.compact\n ... def __call__(self, x, index):\n ... self.variable('state', 'a_count', lambda: 0)\n ... self.variable('state', 'b_count', lambda: 0)\n ... self.variable('state', 'c_count', lambda: 0)\n ... def a_fn(mdl, x):\n ... mdl.variable('state', 'a_count').value += 1\n ... return nn.Dense(2, name='dense')(x)\n ... def b_fn(mdl, x):\n ... mdl.variable('state', 'b_count').value += 1\n ... return -nn.Dense(2, name='dense')(x)\n ... def c_fn(mdl, x):\n ... mdl.variable('state', 'c_count').value += 1\n ... return nn.Dense(2, name='dense')(x)\n ... return nn.switch(index, [a_fn, b_fn, c_fn], self, x)\n\n If you want to have a different parameter structure for each branch\n you should run all branches on initialization before calling switch::\n\n >>> class MultiHeadSwitchExample(nn.Module):\n ... def setup(self) -> None:\n ... self.heads = [\n ... nn.Sequential([nn.Dense(10), nn.Dense(7), nn.Dense(5)]),\n ... nn.Sequential([nn.Dense(11), nn.Dense(5)]),\n ... nn.Dense(5),\n ... ]\n ...\n ... @nn.compact\n ... def __call__(self, x, index):\n ... def head_fn(i):\n ... return lambda mdl, x: mdl.heads[i](x)\n ... branches = [head_fn(i) for i in range(len(self.heads))]\n ...\n ... # run all branches on init\n ... if self.is_mutable_collection('params'):\n ... for branch in branches:\n ... _ = branch(self, x)\n ...\n ... return nn.switch(index, branches, self, x)\n\n Args:\n index: Integer scalar type, indicating which branch function to apply.\n branches: Sequence of functions to be applied based on index.\n The signature of each function is (module, *operands) -> T.\n mdl: A Module target to pass.\n *operands: The arguments passed to the branches.\n variables: The variable collections passed to the conditional\n branches (default: all)\n rngs: The PRNG sequences passed to the conditionals (default: all)\n Returns:\n The result of the evaluated branch.\n \"\"\"\n return lift_direct_transform(\n _switch_wrapper,\n tuple(branches),\n mdl,\n index,\n *operands,\n variables=variables,\n rngs=rngs,\n n_branches=len(branches),\n )\n\n\n# a version of lift.custom_vjp with a single scope function\n# this avoids having to lift multiple functions in\n# lift_transform.\ndef _custom_vjp_single_scope_fn(\n fn: Callable[..., Any],\n backward_fn: Callable[..., Any],\n grad_vars: CollectionFilter = 'params',\n nondiff_argnums=(),\n):\n nodiff_fn = functools.partial(fn, needs_residual=False)\n forward_fn = functools.partial(fn, needs_residual=True)\n return lift.custom_vjp(\n nodiff_fn, forward_fn, backward_fn, grad_vars, nondiff_argnums\n )\n\n\ndef custom_vjp(\n fn: Callable[..., Any],\n forward_fn: Callable[..., Any],\n backward_fn: Callable[..., Any],\n grad_vars: CollectionFilter = 'params',\n nondiff_argnums=(),\n):\n \"\"\"Lifted version of ``jax.custom_vjp``.\n\n ``forward_fn`` and ``backward_fn`` together define a custom vjp for ``fn``.\n The original ``fn`` will run in case a vjp (backward gradient) is not computed.\n\n The ``forward_fn`` receives the same arguments as ``fn`` but is expected to return\n a tuple containing the output of ``fn(mdl, *args)`` and the residuals that are\n passed to ``backward_fn``.\n\n The ``backward_fn`` receives the nondiff arguments, residuals, and the output\n tangents. It should return a tuple containing the variable and input tangents.\n\n Note that the vjp function returned by ``nn.vjp`` can be passed as residual and\n used in the ``backward_fn``. The scope is unavailable during the backward pass.\n If the module is required in ``backward_fn``, a snapshot of the variables can\n be taken and returned as a residual in the ``forward_fn``.\n\n Example::\n\n >>> import flax.linen as nn\n >>> import jax, jax.numpy as jnp\n\n >>> class Foo(nn.Module):\n ... @nn.compact\n ... def __call__(self, x):\n ... def f(mdl, x):\n ... return mdl(x)\n ...\n ... def fwd(mdl, x):\n ... return nn.vjp(f, mdl, x)\n ...\n ... def bwd(vjp_fn, y_t):\n ... params_t, *inputs_t = vjp_fn(y_t)\n ... params_t = jax.tree_util.tree_map(jnp.sign, params_t)\n ... return (params_t, *inputs_t)\n ...\n ... sign_grad = nn.custom_vjp(\n ... f, forward_fn=fwd, backward_fn=bwd)\n ... return sign_grad(nn.Dense(1), x).reshape(())\n\n >>> x = jnp.ones((2,))\n >>> variables = Foo().init(jax.random.key(0), x)\n >>> grad = jax.grad(Foo().apply)(variables, x)\n\n Args:\n fn: The function to define a custom_vjp for.\n forward_fn: A function with the same arguments as ``fn`` returning an tuple\n with the original output and the residuals that will be passed to\n ``backward_fn``.\n backward_fn: arguments are passed as\n ``(*nondiff_args, residuals, tangents)`` The function should return a\n tuple containing the tangents for the variable in the collections\n specified by ``grad_vars`` and the input arguments (except the module and\n nondiff args).\n grad_vars: The collections for which a vjp will be computed\n (default: \"params\").\n nondiff_argnums: arguments for which no vjp is computed.\n Returns:\n A function with the same signature as ``fn`` with the custom vjp.\n \"\"\"\n\n def shared_forward_fn(*args, needs_residual, **kwargs):\n if needs_residual:\n return forward_fn(*args, **kwargs)\n else:\n return fn(*args, **kwargs)\n\n return decorator_lift_transform(\n _custom_vjp_single_scope_fn,\n shared_forward_fn,\n backward_fn=backward_fn,\n grad_vars=grad_vars,\n nondiff_argnums=nondiff_argnums,\n multi_scope=False,\n )\n\n\ndef named_call(class_fn, force=True):\n \"\"\"Labels a method for labelled traces in profiles.\n\n Note that it is better to use the `jax.named_scope` context manager directly\n to add names to JAX's metadata name stack.\n\n Args:\n class_fn: The class method to label.\n force: If True, the named_call transform is applied even if it is globally\n disabled. (e.g.: by calling `flax.linen.disable_named_call()`)\n Returns:\n A wrapped version of ``class_fn`` that is labeled.\n \"\"\"\n\n # We use JAX's dynamic name-stack named_call. No transform boundary needed!\n @functools.wraps(class_fn)\n def wrapped_fn(self, *args, **kwargs):\n if (not force and not linen_module._use_named_call) or self._state.in_setup: # pylint: disable=protected-access # pylint: disable=protected-access\n return class_fn(self, *args, **kwargs)\n full_name = _derive_profiling_name(self, class_fn)\n return jax.named_call(class_fn, name=full_name)(self, *args, **kwargs)\n\n return wrapped_fn\n\n\ndef add_metadata_axis(\n target: Target,\n variable_axes: Mapping[CollectionFilter, InOutAxis] = FrozenDict(),\n metadata_params: dict[Any, Any] = {},\n) -> Target:\n \"\"\"A helper to manipulate boxed axis metadata.\n\n This is a helper to manipulate the *metadata* in boxed variables, similar\n to how lifted ``vmap`` and ``scan`` will handle the introduction and stripping\n of the new metadata axis across a transform boundary.\n\n Args:\n target: a ``Module`` or a function taking a ``Module``\n as its first argument.\n variable_axes: the variable collections whose axis metadata is being\n transformed. Use `None` to indicate a broadcasted collection or an integer\n to specify an axis index for an introduced axis.\n methods: If `target` is a `Module`, the methods of `Module` to vmap over.\n metadata_params: arguments dict passed to AxisMetadata instances in the\n variable tree.\n Returns:\n A transformed version of ``target`` that performs a transform of the\n axis metadata on its variables.\n \"\"\"\n\n def add_fn(axis):\n return lambda x: meta.add_axis(x, axis, metadata_params)\n\n def remove_fn(axis):\n return lambda x: meta.remove_axis(x, axis, metadata_params)\n\n for col_name, axis in variable_axes.items():\n target = map_variables(\n target,\n col_name,\n trans_in_fn=remove_fn(axis),\n trans_out_fn=add_fn(axis),\n mutable=True,\n )\n return target\n\n\ndef fold_rngs(\n target: Target,\n variables: CollectionFilter = True,\n rngs: PRNGSequenceFilter = True,\n) -> Target:\n return lift_transform_cached(\n lift.fold_rngs,\n target,\n variables=variables,\n rngs=rngs,\n )\n" }, { "path": "flax/metrics/__init__.py", "content": "# Copyright 2023 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n" }, { "path": "flax/metrics/tensorboard.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Write Summaries from JAX for use with Tensorboard.\"\"\"\n\nimport contextlib\nimport functools\nimport os\n\nimport numpy as np\nimport tensorflow as tf # pytype: disable=import-error\nfrom tensorboard.plugins.hparams import api as hparams_api\n\n# pylint: disable=g-import-not-at-top\nfrom flax import io\n\n\ndef _flatten_dict(input_dict, parent_key='', sep='.'):\n \"\"\"Flattens and simplifies dict such that it can be used by hparams.\n\n Args:\n input_dict: Input dict, e.g., from ConfigDict.\n parent_key: String used in recursion.\n sep: String used to separate parent and child keys.\n\n Returns:\n Flattened dict.\n \"\"\"\n items = []\n for k, v in input_dict.items():\n new_key = parent_key + sep + k if parent_key else k\n\n # Valid types according to https://github.com/tensorflow/tensorboard/blob/1204566da5437af55109f7a4af18f9f8b7c4f864/tensorboard/plugins/hparams/summary_v2.py\n valid_types = (\n bool,\n int,\n float,\n str,\n np.bool_,\n np.integer,\n np.floating,\n np.character,\n )\n\n if isinstance(v, dict):\n # Recursively flatten the dict.\n items.extend(_flatten_dict(v, new_key, sep=sep).items())\n continue\n elif not isinstance(v, valid_types):\n # Cast any incompatible values as strings such that they can be handled by hparams\n v = str(v)\n items.append((new_key, v))\n return dict(items)\n\n\n@contextlib.contextmanager\ndef _as_default(summary_writer: tf.summary.SummaryWriter, auto_flush: bool):\n \"\"\"No-flush variation of summary_writer.as_default().\"\"\"\n context_manager = summary_writer.as_default()\n try:\n context_manager.__enter__()\n yield summary_writer\n finally:\n old_flush = summary_writer.flush\n new_flush = old_flush if auto_flush else lambda: None\n summary_writer.flush = new_flush\n context_manager.__exit__(None, None, None)\n summary_writer.flush = old_flush\n\n\nclass SummaryWriter:\n \"\"\"Saves data in event and summary protos for tensorboard.\"\"\"\n\n def __init__(self, log_dir, auto_flush=True):\n \"\"\"Create a new SummaryWriter.\n\n Args:\n log_dir: path to record tfevents files in.\n auto_flush: if true, flush after every reported metric.\n \"\"\"\n log_dir = os.fspath(log_dir)\n\n # If needed, create log_dir directory as well as missing parent directories.\n if not io.isdir(log_dir):\n io.makedirs(log_dir)\n\n self._event_writer = tf.summary.create_file_writer(log_dir)\n self._as_default = functools.partial(_as_default, auto_flush=auto_flush)\n self._closed = False\n\n def close(self):\n \"\"\"Close SummaryWriter. Final!\"\"\"\n if not self._closed:\n self._event_writer.close()\n self._closed = True\n del self._event_writer\n\n def flush(self):\n self._event_writer.flush()\n\n def scalar(self, tag, value, step):\n \"\"\"Saves scalar value.\n\n Args:\n tag: str: label for this data\n value: int/float: number to log\n step: int: training step\n \"\"\"\n value = float(np.array(value))\n with self._as_default(self._event_writer):\n tf.summary.scalar(name=tag, data=value, step=step)\n\n def image(self, tag, image, step, max_outputs=3):\n \"\"\"Saves RGB image summary from np.ndarray [H,W], [H,W,1], or [H,W,3].\n\n Args:\n tag: str: label for this data\n image: ndarray: [H,W], [H,W,1], [H,W,3], [K,H,W], [K,H,W,1], [K,H,W,3]\n Save image in greyscale or colors.\n Pixel values could be either uint8 or float.\n Floating point values should be in range [0, 1).\n step: int: training step\n max_outputs: At most this many images will be emitted at each step.\n Defaults to 3.\n \"\"\"\n image = np.array(image)\n # tf.summary.image expects image to have shape [k, h, w, c] where,\n # k = number of samples, h = height, w = width, c = number of channels.\n if len(np.shape(image)) == 2:\n image = image[np.newaxis, :, :, np.newaxis]\n elif len(np.shape(image)) == 3:\n # this could be either [k, h, w] or [h, w, c]\n if np.shape(image)[-1] in (1, 3):\n image = image[np.newaxis, :, :, :]\n else:\n image = image[:, :, :, np.newaxis]\n if np.shape(image)[-1] == 1:\n image = np.repeat(image, 3, axis=-1)\n\n # Convert to tensor value as tf.summary.image expects data to be a tensor.\n image = tf.convert_to_tensor(image)\n with self._as_default(self._event_writer):\n tf.summary.image(name=tag, data=image, step=step, max_outputs=max_outputs)\n\n def audio(self, tag, audiodata, step, sample_rate=44100, max_outputs=3):\n \"\"\"Saves audio as wave.\n\n NB: single channel only right now.\n\n Args:\n tag: str: label for this data\n audiodata: ndarray [Nsamples, Nframes, Nchannels]: audio data to\n be saved as wave. The data will be clipped to [-1.0, 1.0].\n step: int: training step\n sample_rate: sample rate of passed in audio buffer\n max_outputs: At most this many audio clips will be emitted at each\n step. Defaults to 3.\n \"\"\"\n # tf.summary.audio expects the audio data to have floating values in\n # [-1.0, 1.0].\n audiodata = np.clip(np.array(audiodata), -1, 1)\n\n # Convert to tensor value as tf.summary.audio expects data to be a tensor.\n audio = tf.convert_to_tensor(audiodata, dtype=tf.float32)\n with self._as_default(self._event_writer):\n tf.summary.audio(\n name=tag,\n data=audio,\n sample_rate=sample_rate,\n step=step,\n max_outputs=max_outputs,\n encoding='wav',\n )\n\n def histogram(self, tag, values, step, bins=None):\n \"\"\"Saves histogram of values.\n\n Args:\n tag: str: label for this data\n values: ndarray: will be flattened by this routine\n step: int: training step\n bins: number of bins in histogram\n \"\"\"\n values = np.array(values)\n values = np.reshape(values, -1)\n with self._as_default(self._event_writer):\n tf.summary.histogram(name=tag, data=values, step=step, buckets=bins)\n\n def text(self, tag, textdata, step):\n \"\"\"Saves a text summary.\n\n Args:\n tag: str: label for this data\n textdata: string\n step: int: training step\n Note: markdown formatting is rendered by tensorboard.\n \"\"\"\n if not isinstance(textdata, (str, bytes)):\n raise ValueError('`textdata` should be of the type `str` or `bytes`.')\n with self._as_default(self._event_writer):\n tf.summary.text(name=tag, data=tf.constant(textdata), step=step)\n\n def write(self, tag, tensor, step, metadata=None):\n \"\"\"Saves an arbitrary tensor summary.\n\n Useful when working with custom plugins or constructing a summary directly.\n\n Args:\n tag: str: label for this data\n tensor: ndarray: tensor data to save.\n step: int: training step\n metadata: Optional SummaryMetadata, as a proto or serialized bytes.\n Note: markdown formatting is rendered by tensorboard.\n \"\"\"\n with self._as_default(self._event_writer):\n tf.summary.write(tag=tag, tensor=tensor, step=step, metadata=metadata)\n\n def hparams(self, hparams):\n \"\"\"Saves hyper parameters.\n\n Args:\n hparams: Flat mapping from hyper parameter name to value.\n \"\"\"\n\n with self._as_default(self._event_writer):\n hparams_api.hparams(hparams=_flatten_dict(hparams))\n" }, { "path": "flax/nnx/README.md", "content": "[![codecov](https://codecov.io/gh/cgarciae/nnx/branch/main/graph/badge.svg?token=VqJjL474Z7)](https://codecov.io/gh/cgarciae/nnx)\n\n# NNX\n\n_**N**eural **N**etworks for JA**X**_ - | [docs](https://flax.readthedocs.io/en/latest/nnx/index.html) |\n\nNNX is a JAX-based neural network library that focuses on providing the best development experience to make\nbuilding and experimenting with neural networks as easy and intuitive as possible.\n\n* **Pythonic**: Modules are standard Python classes, promoting ease of use and a more familiar\n development experience.\n* **Easy-to-use**: NNX provides a set of transforms that take care of state management, allowing\n users to focus on building their models and training loops.\n* **Expressive**: NNX allows fine-grained over the Module state with lifted transforms, enabling\n users to define complex architectures.\n* **Compatible**: NNX allows functionalizing Module state, making it possible to directly use JAX\n transformations when needed.\n\n## What does NNX look like?\n\nNNX removes most of the friction from building and training neural networks in JAX. It provides\na Module system that uses standard Python classes, and a set of transforms that extend\nJAX to handle objects.\n\n```python\nfrom flax import nnx\nimport optax\n\nclass Model(nnx.Module):\n def __init__(self, din, dmid, dout, rngs: nnx.Rngs):\n self.linear = nnx.Linear(din, dmid, rngs=rngs)\n self.bn = nnx.BatchNorm(dmid, rngs=rngs)\n self.dropout = nnx.Dropout(0.2, rngs=rngs)\n self.linear_out = nnx.Linear(dmid, dout, rngs=rngs)\n\n def __call__(self, x):\n x = nnx.relu(self.dropout(self.bn(self.linear(x))))\n return self.linear_out(x)\n\n\nmodel = Model(2, 64, 3, rngs=nnx.Rngs(0)) # eager initialization\noptimizer = nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param) # reference sharing\n\n@nnx.jit # automatic state management\ndef train_step(model, optimizer, x, y):\n def loss_fn(model):\n y_pred = model(x) # call methods directly\n return ((y_pred - y) ** 2).mean()\n\n loss, grads = nnx.value_and_grad(loss_fn)(model)\n optimizer.update(model, grads) # inplace updates\n\n return loss\n```\n\nTo learn more about the `Module` abstraction, check out our [NNX Basics](https://flax.readthedocs.io/en/latest/nnx/nnx_basics.html#) guide.\n\n## Installation\n\nTo get started with `nnx`, install Flax from GitHub:\n```\npip install git+https://github.com/google/flax.git\n```\n\n### Examples\n\n* [LM1B](https://github.com/google/flax/tree/main/examples/lm1b_nnx): A language model trained on the 1 Billion Word Benchmark dataset.\n\n#### Toy Examples\n* [Basic Example](https://github.com/google/flax/tree/main/examples/nnx_toy_examples/02_lifted_transforms.py): Shows how to train a simple model using NNX.\n* [Using the Functional API](https://github.com/google/flax/tree/main/examples/nnx_toy_examples/01_functional_api.py): Shows how to train a simple model using the functional API.\n* [Training a VAE](https://github.com/google/flax/tree/main/examples/nnx_toy_examples/05_vae.py): Shows how to train a VAE on the binarized MNIST dataset.\n* [Scan over layers](https://github.com/google/flax/tree/main/examples/nnx_toy_examples/06_scan_over_layers.py): An contrived example that implements scan over layers with dropout and a share BatcNorm layer to showcase how lifted transforms can be implemented. It uses the functional API along with `jax.vmap` and `jax.lax.scan`.\n" }, { "path": "flax/nnx/__init__.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom flax.core.spmd import logical_axis_rules as logical_axis_rules\nfrom flax.linen.pooling import avg_pool as avg_pool\nfrom flax.linen.pooling import max_pool as max_pool\nfrom flax.linen.pooling import min_pool as min_pool\nfrom flax.linen.pooling import pool as pool\nfrom flax.typing import Initializer as Initializer\n\nfrom .bridge import wrappers as wrappers\nfrom .filterlib import WithTag as WithTag\nfrom .filterlib import PathContains as PathContains\nfrom .filterlib import OfType as OfType\nfrom .filterlib import Any as Any\nfrom .filterlib import All as All\nfrom .filterlib import Not as Not\nfrom .filterlib import Everything as Everything\nfrom .filterlib import Nothing as Nothing\nfrom .graphlib import GraphDef as GraphDef\nfrom .graphlib import GraphState as GraphState\nfrom .graphlib import PureState as PureState\nfrom . import pytreelib as object\nfrom .pytreelib import Pytree as Pytree\nfrom .pytreelib import Object as Object\nfrom .pytreelib import Data as Data\nfrom .pytreelib import Static as Static\nfrom .pytreelib import dataclass as dataclass\nfrom .pytreelib import data as data\nfrom .pytreelib import static as static\nfrom .pytreelib import register_data_type as register_data_type\nfrom .pytreelib import is_data as is_data\nfrom .pytreelib import has_data as has_data\nfrom .pytreelib import check_pytree as check_pytree\nfrom .helpers import Dict as Dict\nfrom .helpers import List as List\nfrom .helpers import Sequential as Sequential\nfrom .helpers import TrainState as TrainState\nfrom .module import M as M\nfrom .module import Module as Module\nfrom .module import capture as capture\nfrom .module import view as view\nfrom .module import view_info as view_info\nfrom .module import with_attributes as with_attributes\nfrom .module import iter_children as iter_children, iter_modules as iter_modules\nfrom .graphlib import merge as merge\nfrom .graphlib import UpdateContext as UpdateContext\nfrom .graphlib import update_context as update_context\nfrom .graphlib import current_update_context as current_update_context\nfrom .graphlib import split as split\nfrom .graphlib import update as update\nfrom .graphlib import clone as clone\nfrom .graphlib import pop as pop\nfrom .graphlib import state as state\nfrom .graphlib import graphdef as graphdef\nfrom .graphlib import iter_graph as iter_graph\nfrom .graphlib import recursive_map as recursive_map\nfrom .graphlib import find_duplicates as find_duplicates\nfrom .graphlib import map as map\nfrom .graphlib import call as call\nfrom .graphlib import set_metadata as set_metadata\nfrom .graphlib import SplitContext as SplitContext\nfrom .graphlib import split_context as split_context\nfrom .graphlib import MergeContext as MergeContext\nfrom .graphlib import merge_context as merge_context\nfrom .graphlib import variables as variables\nfrom .graphlib import vars_as as vars_as\nfrom .graphlib import pure as pure\nfrom .graphlib import cached_partial as cached_partial\nfrom .graphlib import flatten as flatten\nfrom .graphlib import unflatten as unflatten\nfrom .graphlib import set_graph_mode as set_graph_mode\nfrom .graphlib import set_graph_updates as set_graph_updates\nfrom .nn import initializers as initializers\nfrom .nn.activations import celu as celu\nfrom .nn.activations import elu as elu\nfrom .nn.activations import gelu as gelu\nfrom .nn.activations import glu as glu\nfrom .nn.activations import hard_sigmoid as hard_sigmoid\nfrom .nn.activations import hard_silu as hard_silu\nfrom .nn.activations import hard_swish as hard_swish\nfrom .nn.activations import hard_tanh as hard_tanh\nfrom .nn.activations import leaky_relu as leaky_relu\nfrom .nn.activations import log_sigmoid as log_sigmoid\nfrom .nn.activations import log_softmax as log_softmax\nfrom .nn.activations import logsumexp as logsumexp\nfrom .nn.activations import one_hot as one_hot\nfrom .nn.activations import relu as relu\nfrom .nn.activations import relu6 as relu6\nfrom .nn.activations import selu as selu\nfrom .nn.activations import sigmoid as sigmoid\nfrom .nn.activations import identity as identity\nfrom .nn.activations import silu as silu\nfrom .nn.activations import soft_sign as soft_sign\nfrom .nn.activations import softmax as softmax\nfrom .nn.activations import softplus as softplus\nfrom .nn.activations import standardize as standardize\nfrom .nn.activations import swish as swish\nfrom .nn.activations import tanh as tanh\nfrom .nn.activations import PReLU as PReLU\nfrom .nn.attention import MultiHeadAttention as MultiHeadAttention\nfrom .nn.attention import combine_masks as combine_masks\nfrom .nn.attention import dot_product_attention as dot_product_attention\nfrom .nn.attention import make_attention_mask as make_attention_mask\nfrom .nn.attention import make_causal_mask as make_causal_mask\nfrom .nn.recurrent import RNNCellBase as RNNCellBase\nfrom .nn.recurrent import LSTMCell as LSTMCell\nfrom .nn.recurrent import GRUCell as GRUCell\nfrom .nn.recurrent import OptimizedLSTMCell as OptimizedLSTMCell\nfrom .nn.recurrent import SimpleCell as SimpleCell\nfrom .nn.recurrent import RNN as RNN\nfrom .nn.recurrent import Bidirectional as Bidirectional\nfrom .nn.linear import Conv as Conv\nfrom .nn.linear import ConvTranspose as ConvTranspose\nfrom .nn.linear import Embed as Embed\nfrom .nn.linear import Linear as Linear\nfrom .nn.linear import LinearGeneral as LinearGeneral\nfrom .nn.linear import Einsum as Einsum\nfrom .nn.lora import LoRA as LoRA\nfrom .nn.lora import LoRALinear as LoRALinear\nfrom .nn.lora import LoRAParam as LoRAParam\nfrom .nn.normalization import BatchNorm as BatchNorm\nfrom .nn.normalization import LayerNorm as LayerNorm\nfrom .nn.normalization import RMSNorm as RMSNorm\nfrom .nn.normalization import GroupNorm as GroupNorm\nfrom .nn.normalization import InstanceNorm as InstanceNorm\nfrom .nn.normalization import SpectralNorm as SpectralNorm\nfrom .nn.normalization import WeightNorm as WeightNorm\nfrom .nn.stochastic import Dropout as Dropout\nfrom .rnglib import Rngs as Rngs\nfrom .rnglib import RngStream as RngStream\nfrom .rnglib import RngState as RngState\nfrom .rnglib import RngKey as RngKey\nfrom .rnglib import RngCount as RngCount\nfrom .rnglib import fork_rngs as fork_rngs\nfrom .rnglib import reseed as reseed\nfrom .rnglib import split_rngs as split_rngs\nfrom .rnglib import restore_rngs as restore_rngs\nfrom .spmd import PARTITION_NAME as PARTITION_NAME\nfrom .spmd import get_partition_spec as get_partition_spec\nfrom .spmd import get_named_sharding as get_named_sharding\nfrom .spmd import with_partitioning as with_partitioning\nfrom .spmd import get_abstract_model as get_abstract_model\nfrom .spmd import abstract_with_sharding as abstract_with_sharding\nfrom .statelib import FlatState as FlatState\nfrom .statelib import State as State\nfrom .statelib import to_flat_state as to_flat_state\nfrom .statelib import from_flat_state as from_flat_state\nfrom .statelib import to_pure_dict as to_pure_dict\nfrom .statelib import replace_by_pure_dict as replace_by_pure_dict\nfrom .statelib import restore_int_paths as restore_int_paths\nfrom .statelib import filter_state as filter_state\nfrom .statelib import merge_state as merge_state\nfrom .statelib import split_state as split_state\nfrom .statelib import map_state as map_state\nfrom .training import metrics as metrics\nfrom .variablelib import Param as Param\n# this needs to be imported before optimizer to prevent circular import\nfrom .training import optimizer as optimizer\nfrom .training.metrics import Metric as Metric\nfrom .training.metrics import MultiMetric as MultiMetric\nfrom .training.optimizer import OptState as OptState\nfrom .training.optimizer import OptArray as OptArray\nfrom .training.optimizer import OptVariable as OptVariable\nfrom .training.optimizer import Optimizer as Optimizer\nfrom .training.optimizer import ModelAndOptimizer as ModelAndOptimizer\nfrom .training.optimizer import OptState as OptState\nfrom .transforms.autodiff import DiffState as DiffState\nfrom .transforms.autodiff import grad as grad\nfrom .transforms.autodiff import value_and_grad as value_and_grad\nfrom .transforms.autodiff import custom_vjp as custom_vjp\nfrom .transforms.autodiff import vjp as vjp\nfrom .transforms.autodiff import jvp as jvp\nfrom .transforms.autodiff import remat as remat\nfrom .transforms.compilation import jit as jit\nfrom .transforms.compilation import jit_partial as jit_partial\nfrom .transforms.compilation import shard_map as shard_map\nfrom .transforms.compilation import StateSharding as StateSharding\nfrom .transforms.iteration import Carry as Carry\nfrom .transforms.iteration import scan as scan\nfrom .transforms.iteration import vmap as vmap\nfrom .transforms.iteration import pmap as pmap\nfrom .transforms.transforms import eval_shape as eval_shape\nfrom .transforms.transforms import cond as cond\nfrom .transforms.transforms import switch as switch\nfrom .transforms.transforms import checkify as checkify\nfrom .transforms.iteration import while_loop as while_loop\nfrom .transforms.iteration import fori_loop as fori_loop\nfrom .transforms.iteration import StateAxes as StateAxes\nfrom .transforms.iteration import transform_metadata as transform_metadata\nfrom .variablelib import A as A\nfrom .variablelib import BatchStat as BatchStat\nfrom .variablelib import Cache as Cache\nfrom .variablelib import Intermediate as Intermediate\nfrom .variablelib import Perturbation as Perturbation\nfrom .variablelib import Variable as Variable\nfrom .variablelib import VariableMetadata as VariableMetadata\nfrom .variablelib import with_metadata as with_metadata\nfrom .variablelib import variable_type_from_name as variable_type_from_name\nfrom .variablelib import variable_name_from_type as variable_name_from_type\nfrom .variablelib import register_variable_name as register_variable_name\nfrom .variablelib import use_eager_sharding as use_eager_sharding, using_eager_sharding as using_eager_sharding, var_defaults as var_defaults\nfrom .visualization import display as display\nfrom .extract import to_tree as to_tree\nfrom .extract import from_tree as from_tree\nfrom .extract import NodeStates as NodeStates\nfrom .summary import tabulate as tabulate\nfrom . import traversals as traversals\nfrom . import graphlib as graphlib\n# import last to prevent potential import cycles\nfrom . import graph as graph\nfrom . import compat as compat\n\nimport typing as _tp\n\nif _tp.TYPE_CHECKING:\n VariableState = Variable\nelse:\n def __getattr__(name):\n if name == \"VariableState\":\n import warnings\n warnings.warn(\n \"'VariableState' was removed, this is just an alias to 'Variable'. \"\n \"Plase use 'Variable' directly instead.\",\n DeprecationWarning,\n stacklevel=2,\n )\n return Variable\n raise AttributeError(f\"Module {__name__} has no attribute '{name}'\")\n" }, { "path": "flax/nnx/bridge/__init__.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nfrom .wrappers import functional as functional\nfrom .wrappers import Functional as Functional\nfrom .wrappers import ToNNX as ToNNX\nfrom .wrappers import lazy_init as lazy_init\nfrom .wrappers import ToLinen as ToLinen\nfrom .wrappers import to_linen as to_linen\nfrom .variables import NNXMeta as NNXMeta\nfrom .variables import with_partitioning as with_partitioning\nfrom .module import Module as Module\nfrom .module import Scope as Scope\nfrom .module import AttrPriority as AttrPriority\nfrom .module import compact as compact\nfrom .module import current_context as current_context\nfrom .module import current_module as current_module\nfrom .interop import nnx_in_bridge_mdl as nnx_in_bridge_mdl\nfrom .interop import linen_in_bridge_mdl as linen_in_bridge_mdl\nfrom flax.nnx.nn import initializers as initializers\n" }, { "path": "flax/nnx/bridge/interop.py", "content": "# Copyright 2025 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport typing as tp\n\nfrom flax.linen import module as nn_module\nfrom flax.nnx import graphlib, rnglib\nfrom flax.nnx.bridge import wrappers\nfrom flax.nnx.bridge import module as bdg_module\nimport flax.nnx.module as nnx_module\nfrom flax.nnx.transforms.transforms import eval_shape as nnx_eval_shape\nfrom flax.nnx.transforms.compilation import jit as nnx_jit\n\n\ndef nnx_in_bridge_mdl(factory: tp.Callable[[rnglib.Rngs], nnx_module.Module],\n name: str | None = None) -> nnx_module.Module:\n \"\"\"Make pure NNX modules a submodule of a bridge module.\n\n Create module at init time, or make abstract module and let parent bind\n it with its state.\n Use current bridge module scope for RNG generation.\n\n Args:\n factory: a function that takes an `nnx.Rngs` arg and returns an NNX module.\n name: the name of the module. Only used during `bridge.compact` functions;\n in setup() function the user will set it to an attribute explicitly.\n Returns:\n A submodule (`nnx.Module`) of the bridge module.\n \"\"\"\n parent_ctx, parent = bdg_module.current_context(), bdg_module.current_module()\n assert parent_ctx is not None and parent is not None, 'nnx_in_bridge_mdl() only needed inside bridge Module'\n parent = parent_ctx.module\n assert parent.scope is not None\n\n if parent.is_initializing():\n module = factory(parent.scope.rngs)\n else:\n rngs = parent.scope.rngs if parent.scope.rngs else rnglib.Rngs(7) # dummy\n module = nnx_eval_shape(factory, rngs, graph=True)\n\n @nnx_jit\n def rng_state(rngs):\n return graphlib.state(factory(rngs), rnglib.RngState, graph=True)\n\n # Make sure the internal rng state is not abstract - other vars shall be\n if parent.scope.rngs:\n graphlib.update(module, rng_state(parent.scope.rngs))\n\n # Automatically set the attribute if compact. If setup, user is responsible\n # for setting the attribute of the superlayer.\n if parent_ctx.in_compact:\n if name is None:\n name = bdg_module._auto_submodule_name(parent_ctx, type(module))\n setattr(parent, name, module)\n return module\n\n\ndef linen_in_bridge_mdl(linen_module: nn_module.Module,\n name: str | None = None) -> nnx_module.Module:\n \"\"\"Make Linen modules a submodule of a bridge module using wrappers.ToNNX().\n\n Args:\n linen_module: the underlying Linen module instance.\n name: the name of the module. Only used during `bridge.compact` functions;\n in setup() function the user will set it to an attribute explicitly.\n Returns:\n A submodule (`nnx.Module`) of the bridge module.\n \"\"\"\n parent_ctx, parent = bdg_module.current_context(), bdg_module.current_module()\n assert parent_ctx is not None and parent is not None, 'linen_in_bridge_mdl() only needed inside bridge Module'\n assert parent.scope is not None\n module = wrappers.ToNNX(linen_module, parent.scope.rngs)\n wrappers._set_initializing(module, parent.is_initializing())\n if parent_ctx.in_compact:\n if name is None:\n name = bdg_module._auto_submodule_name(parent_ctx, type(linen_module))\n setattr(parent, name, module)\n return module\n" }, { "path": "flax/nnx/bridge/module.py", "content": "# Copyright 2025 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom __future__ import annotations\n\nfrom collections import defaultdict\nimport dataclasses\nimport enum\nimport functools\nimport inspect\nimport threading\nimport typing as tp\nimport jax\nimport typing_extensions as tpe\n\nfrom flax import errors\nfrom flax.core import meta\nfrom flax.core.scope import CollectionFilter\nfrom flax.core.frozen_dict import FrozenDict\nfrom flax.nnx import graphlib, rnglib, statelib, traversals\nfrom flax.nnx import variablelib\nimport flax.nnx.module as nnx_module\nfrom flax.nnx.pytreelib import Pytree\nfrom flax.nnx import variablelib\nfrom flax.nnx.bridge import variables as bridge_variables\nimport numpy as np\n\nA = tp.TypeVar('A')\nM = tp.TypeVar('M', bound='Module')\nF = tp.TypeVar('F', bound=tp.Callable[..., tp.Any])\n\n\n@dataclasses.dataclass\nclass ModuleStackEntry:\n module: Module\n in_compact: bool\n type_counter: defaultdict[type, int] = dataclasses.field(\n default_factory=lambda: defaultdict(int)\n )\n\n\n@dataclasses.dataclass\nclass ModuleContext(threading.local):\n module_stack: list[ModuleStackEntry | None] = dataclasses.field(\n default_factory=lambda: [None]\n )\n\n\nMODULE_CONTEXT = ModuleContext()\n\n\nclass ModuleState(statelib.State):\n pass\n\nfrom flax.nnx.pytreelib import register_data_type\nregister_data_type(ModuleState)\n\n\nclass Scope(Pytree):\n def __init__(self, rngs: rnglib.Rngs, mutable: CollectionFilter):\n self.rngs = rngs\n self.mutable = mutable\n\n def copy(self):\n return Scope(self.rngs, self.mutable)\n\n\nclass _HasSetup(tp.Protocol):\n def setup(self) -> None: ...\n\n\ndef has_setup(x: tp.Any) -> tp.TypeGuard[_HasSetup]:\n return hasattr(x, 'setup')\n\n\ndef _maybe_call_setup(module: Module):\n if (\n has_setup(module)\n and isinstance(module, Module)\n and not module._pytree__state.is_setup\n ):\n # void parent context\n MODULE_CONTEXT.module_stack.append(\n ModuleStackEntry(module, in_compact=False)\n )\n try:\n module.setup() # type: ignore[attribute-error]\n module._pytree__state._is_setup = True\n finally:\n MODULE_CONTEXT.module_stack.pop()\n\n\ndef _bind_module(parent: Module, module: Module) -> Module:\n assert parent.scope is not None\n\n for _, value in reversed(list(graphlib.iter_graph(module, graph=True))):\n if isinstance(value, Module):\n if module.scope is None:\n value.scope = parent.scope.copy() # type: ignore[attribute-error]\n _maybe_call_setup(value)\n return module\n\n\ndef current_context() -> ModuleStackEntry | None:\n return MODULE_CONTEXT.module_stack[-1]\n\n\ndef current_module() -> Module | None:\n \"\"\"A quick util to get the current bridge module.\"\"\"\n ctx = current_context()\n if ctx is None:\n return None\n return ctx.module\n\n\ndef _auto_submodule_name(parent_ctx, cls):\n \"\"\"Increment type count and generate a new submodule name.\"\"\"\n type_index = parent_ctx.type_counter[cls]\n parent_ctx.type_counter[cls] += 1\n return f'{cls.__name__}_{type_index}'\n\n\nclass ModuleMeta(nnx_module.ModuleMeta):\n\n def _pytree_meta_construct(cls, self, *args, **kwargs):\n vars(self)['scope'] = None\n super()._pytree_meta_construct(self, *args, **kwargs)\n\n\ndef _module_meta_call(cls: type[M], *args, **kwargs) -> M:\n # compact behavior\n parent_ctx = MODULE_CONTEXT.module_stack[-1]\n parent = None\n module: M\n\n name = None\n if parent_ctx is not None:\n if 'parent' in kwargs:\n parent = kwargs.pop('parent')\n if parent_ctx.in_compact and parent is not None:\n raise ValueError(\n f\"'parent' can only be set to None, got {type(parent).__name__}\"\n )\n else:\n parent = parent_ctx.module\n\n if 'name' in kwargs:\n name = kwargs['name']\n if not 'name' in inspect.get_annotations(cls):\n kwargs.pop('name')\n if not isinstance(name, str):\n raise ValueError(f\"'name' must be a 'str', got {type(name).__name__}\")\n elif parent_ctx.in_compact:\n name = _auto_submodule_name(parent_ctx, cls)\n\n module = nnx_module.ModuleMeta.__call__(cls, *args, **kwargs)\n module.scope = None\n module.attr_priorities = {}\n\n if parent is not None:\n assert parent.scope is not None\n # compact, or setup if `name` exists\n if name is not None:\n setattr(parent, name, module)\n parent.set_attr_priority(name, AttrPriority.INIT_PARENT)\n\n return module # type: ignore\n\n# set ModuleMeta.__call__ like this because pytype doesn't understand\n# the use of TYPE_CHECKING conditionals for metaclass methods\nModuleMeta.__call__ = _module_meta_call # type: ignore\n\nclass AttrPriority(enum.IntEnum):\n HIGH = 0\n INIT_PARENT = 20\n DEFAULT = 50\n LOW = 100\n\n\nclass PriorityStr(str):\n _priority: AttrPriority\n\n def __new__(cls, priority: AttrPriority, value: str):\n obj = super().__new__(cls, value)\n obj._priority = priority\n return obj\n\n def _check_and_get_priority(self, other) -> AttrPriority:\n if not isinstance(other, (str, PriorityStr)):\n raise NotImplementedError(\n f'Cannot compare {type(self)} with {type(other)}'\n )\n if isinstance(other, PriorityStr):\n return other._priority\n return AttrPriority.DEFAULT\n\n def __lt__(self, other) -> bool:\n other_priority = self._check_and_get_priority(other)\n if self._priority == other_priority:\n return super().__lt__(other)\n return self._priority < other_priority\n\n def __gt__(self, other) -> bool:\n other_priority = self._check_and_get_priority(other)\n if self._priority == other_priority:\n return super().__gt__(other)\n return self._priority > other_priority\n\nclass ModuleBase:\n if tp.TYPE_CHECKING:\n scope: Scope | None\n attr_priorities: dict[str, AttrPriority]\n\n\n@tpe.dataclass_transform(field_specifiers=(dataclasses.field,)) # type: ignore[not-supported-yet]\nclass Module(nnx_module.Module, ModuleBase, metaclass=ModuleMeta):\n def __init_subclass__(cls) -> None:\n super().__init_subclass__(pytree=False)\n\n cls = dataclasses.dataclass(repr=False)(cls)\n cls.__hash__ = object.__hash__ # type: ignore[method-assign]\n\n def __getattribute__(self, name: str):\n return type(self)._getattr(self, name)\n\n def _getattr(self, name: str) -> tp.Any:\n value = super().__getattribute__(name)\n if isinstance(value, ModuleState):\n raise AttributeError\n return value\n\n def _setattr(self, name: str, value: tp.Any) -> None:\n if self.scope is not None:\n if name in vars(self) and isinstance(\n state := vars(self)[name], ModuleState\n ):\n graphlib.update(value, state)\n for leaf in jax.tree.leaves(value, is_leaf=graphlib.is_graph_node):\n if isinstance(leaf, Module):\n leaf._pytree__state._initializing = self.is_initializing()\n _bind_module(self, leaf)\n super()._setattr(name, value)\n\n def _graph_node_flatten(self):\n nodes = vars(self).copy()\n keys = (\n PriorityStr(self.attr_priorities.get(k, AttrPriority.DEFAULT), k)\n for k in nodes.keys()\n )\n sorted_nodes = list((k, nodes[k]) for k in sorted(keys))\n return sorted_nodes, type(self)\n\n def set_attr_priority(self, name: str, value: AttrPriority):\n self.attr_priorities[name] = value\n\n def make_rng(self, name: str = 'default') -> jax.Array:\n if self.scope is None:\n raise ValueError(\"Can't use RNGs on unbound modules\")\n return self.scope.rngs[name]() # type: ignore[attribute-error]\n\n def param( # type: ignore[invalid-annotation]\n self,\n name: str,\n init_fn: tp.Callable[..., A],\n *init_args,\n unbox: bool = True,\n **init_kwargs,\n ) -> variablelib.Param[A]:\n # TODO(cgarciae): implement same condition as linen\n if self.scope is None:\n raise ValueError(\n 'Parameters must be initialized in `setup()` or in a method '\n 'wrapped in `@compact`'\n )\n if hasattr(self, name):\n value = getattr(self, name)\n # TODO(cgarciae): implement reservations\n # if self._name_taken(name):\n # raise errors.NameInUseError('param', name, self.__class__.__name__)\n\n if isinstance(value, variablelib.Variable):\n if not isinstance(value, variablelib.Param):\n raise ValueError(\n f\"Expected '{name}' to be a Param, got {type(value).__name__}\"\n )\n return value\n\n abs_value = jax.eval_shape(\n lambda: init_fn(jax.random.key(0), *init_args, **init_kwargs)\n )\n abs_value_flat = jax.tree.leaves(abs_value)\n value_flat = jax.tree.leaves(value)\n for val, abs_val in zip(value_flat, abs_value_flat):\n if np.shape(val) != np.shape(abs_val):\n raise errors.ScopeParamShapeError(\n name, '', np.shape(abs_val), np.shape(val)\n )\n\n if isinstance(abs_value, variablelib.VariableMetadata):\n abs_value.raw_value = value\n value = abs_value\n\n variable = variablelib.Param(value)\n else:\n value = init_fn(self.make_rng('params'), *init_args, **init_kwargs)\n variable = variablelib.Param(value)\n\n setattr(self, name, variable)\n return variable\n\n def variable( # type: ignore[invalid-annotation]\n self,\n col: str,\n name: str,\n init_fn: tp.Callable[..., A] | None = None,\n *init_args,\n unbox: bool = True,\n **init_kwargs,\n ) -> variablelib.Variable[A]:\n variable_type = variablelib.variable_type_from_name(\n col, allow_register=True\n )\n if self.scope is None:\n raise ValueError(\n 'Variables must be initialized in `setup()` or in a method '\n 'wrapped in `@compact`'\n )\n\n if hasattr(self, name):\n value = getattr(self, name)\n # TODO(cgarciae): implement reservations\n # if self._name_taken(name):\n # raise errors.NameInUseError('param', name, self.__class__.__name__)\n\n if isinstance(value, variablelib.Variable):\n return value\n\n if init_fn is None:\n raise ValueError(f\"Expected 'init_fn' to be a callable, got None\")\n\n abs_value = jax.eval_shape(lambda: init_fn(*init_args, **init_kwargs))\n abs_value_flat = jax.tree.leaves(abs_value)\n value_flat = jax.tree.leaves(value)\n for val, abs_val in zip(value_flat, abs_value_flat):\n if np.shape(val) != np.shape(abs_val):\n raise errors.ScopeParamShapeError(\n name, '', np.shape(abs_val), np.shape(val)\n )\n\n if isinstance(abs_value, variablelib.VariableMetadata):\n abs_value.raw_value = value\n value = abs_value\n\n variable = variable_type(value)\n else:\n if init_fn is None:\n raise ValueError(f\"Expected 'init_fn' to be a callable, got None\")\n\n value = init_fn(*init_args, **init_kwargs)\n variable = variable_type(value)\n\n setattr(self, name, variable)\n return variable\n\n def _get_variables(self) -> tp.Mapping:\n state = graphlib.state(self, graph=True)\n _variables: dict = {}\n\n variable: variablelib.Variable\n for path, variable in statelib.to_flat_state(state):\n if isinstance(variable, rnglib.RngState):\n # Don't return RNG states, since Linen doesn't have them.\n continue\n\n try:\n collection = variablelib.variable_name_from_type(type(variable))\n except ValueError:\n collection = type(variable).__name__\n\n if collection not in _variables:\n _variables[collection] = {}\n\n if isinstance(\n variable, variablelib.Variable\n ) and bridge_variables.is_vanilla_variable(variable):\n leaf = variable.get_value()\n else:\n leaf = bridge_variables.to_linen_var(variable)\n\n _variables[collection][path] = leaf\n\n _variables = {\n collection: traversals.unflatten_mapping(flat_state)\n for collection, flat_state in _variables.items()\n }\n\n return _variables\n\n @property\n def variables(self):\n _variables = FrozenDict(self._get_variables())\n return _variables\n\n def apply(\n self,\n variables: dict[str, tp.Mapping],\n *args,\n rngs: int | jax.Array | dict[str, jax.Array] | rnglib.Rngs | None = None,\n method: tp.Callable[..., tp.Any] | str = '__call__',\n mutable: CollectionFilter = False,\n _initialize: bool = False,\n **kwargs,\n ) -> tp.Any:\n module = graphlib.clone(self, graph=True)\n\n # create variables\n real_variables = dict(variables)\n for col_name, linen_collection in variables.items():\n\n def to_variable(value):\n return bridge_variables.to_nnx_var(col_name, value)\n\n linen_collection = jax.tree.map(\n to_variable,\n linen_collection,\n is_leaf=lambda x: isinstance(x, meta.AxisMetadata),\n )\n real_variables[col_name] = linen_collection\n\n states = ({},) if not real_variables else real_variables.values()\n state = statelib.merge_state(*states, cls=ModuleState)\n graphlib.update(module, state)\n\n if rngs is None:\n rngs = rnglib.Rngs()\n elif isinstance(rngs, jax.Array | int):\n rngs = rnglib.Rngs(rngs)\n elif isinstance(rngs, dict):\n rngs = rnglib.Rngs(**rngs)\n\n # get method\n _method: tp.Callable[..., tp.Any]\n if isinstance(method, str):\n attribute_name = method\n _method = getattr(module, attribute_name)\n if not callable(_method):\n class_name = type(module).__name__\n raise TypeError(\n f\"'{class_name}.{attribute_name}' must be a callable, got\"\n f' {type(_method)}.'\n )\n # if the `method` string is a submodule, we create a lambda function\n # that calls the submodule, forwarding all arguments.\n if isinstance(_method, Module):\n _method = lambda module, *args, **kwargs: getattr(\n module, attribute_name\n )(*args, **kwargs)\n else:\n _method = method\n _method = _get_unbound_fn(_method)\n\n # set temporary state\n for _, value in graphlib.iter_graph(module, graph=True):\n if isinstance(value, Pytree):\n value._pytree__state._initializing = _initialize\n if isinstance(value, Module):\n value.scope = Scope(rngs, mutable)\n _maybe_call_setup(value)\n\n MODULE_CONTEXT.module_stack.append(\n ModuleStackEntry(module, in_compact=False)\n )\n try:\n out = _method(module, *args, **kwargs)\n finally:\n MODULE_CONTEXT.module_stack.pop()\n # reset temporary state\n for _, value in graphlib.iter_graph(module, graph=True):\n if isinstance(value, Pytree):\n value._pytree__state._initializing = False\n if isinstance(value, Module):\n value.scope = None\n\n _variables: tp.Mapping = module._get_variables()\n\n if mutable is False:\n return out\n else:\n return out, _variables\n\n def init(\n self,\n rngs: int | jax.Array | dict[str, jax.Array] | rnglib.Rngs | None = None,\n *args,\n method: tp.Callable[..., tp.Any] | str = '__call__',\n **kwargs,\n ):\n out, variables = self.apply(\n {},\n *args,\n _initialize=True,\n mutable=True,\n rngs=rngs,\n method=method,\n **kwargs,\n )\n return variables\n\n def init_with_output(\n self,\n rngs: int | jax.Array | dict[str, jax.Array] | rnglib.Rngs | None = None,\n *args,\n method: tp.Callable[..., tp.Any] | str = '__call__',\n mutable: tp.Any = False,\n # capture_intermediates: bool | Callable[['Module', str], bool] = False,\n **kwargs,\n ) -> tuple[tp.Any, dict[str, tp.Mapping]]:\n return self.apply(\n {},\n *args,\n rngs=rngs,\n method=method,\n mutable=True,\n _initialize=True,\n **kwargs,\n )\n\n def is_initializing(self) -> bool:\n return self._pytree__state._initializing\n\n\ndef compact(f: F) -> F:\n @functools.wraps(f)\n def compact_wrapper(self, *args, **kwargs):\n if not isinstance(self, Module):\n raise ValueError(\n f\"Expected 'self' to be a nnx.bridge.Module, got {type(self).__name__}\"\n )\n\n MODULE_CONTEXT.module_stack.append(ModuleStackEntry(self, in_compact=True))\n\n try:\n return f(self, *args, **kwargs)\n finally:\n MODULE_CONTEXT.module_stack.pop()\n\n return compact_wrapper # type: ignore\n\n\ndef _get_unbound_fn(method_or_fn: tp.Callable) -> tp.Callable:\n if inspect.ismethod(method_or_fn) and isinstance(\n method_or_fn.__self__, Module\n ): # pytype: disable=attribute-error\n method_or_fn = method_or_fn.__func__ # pytype: disable=attribute-error\n if (\n not callable(method_or_fn)\n or len(inspect.signature(method_or_fn).parameters) < 1\n ):\n raise errors.ApplyModuleInvalidMethodError(method_or_fn)\n\n return method_or_fn\n\n" }, { "path": "flax/nnx/bridge/variables.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom typing import Any, TypeVar\nimport typing as tp\n\nimport jax\nfrom flax import struct\nfrom flax.core import meta\nfrom flax.nnx import spmd\nfrom flax.nnx import traversals\nfrom flax.nnx import variablelib\nfrom flax.typing import LogicalNames\n\n\nA = TypeVar('A')\nB = TypeVar('B')\n\n\ndef sort_variable_types(types: tp.Iterable[type]):\n def _variable_parents_count(t: type):\n return sum(1 for p in t.mro() if issubclass(p, variablelib.Variable))\n parent_count = {t: _variable_parents_count(t) for t in types}\n return sorted(types, key=lambda t: -parent_count[t])\n\n\n#############################################\n### NNX Variable <-> Linen metadata boxes ###\n#############################################\n\n\nclass NNXMeta(struct.PyTreeNode, meta.AxisMetadata[A]):\n \"\"\"Default Flax metadata class for `nnx.Variable`.\"\"\"\n\n var_type: type[variablelib.Variable[tp.Any]] = struct.field(pytree_node=False)\n value: Any = struct.field(pytree_node=True)\n metadata: dict[str, tp.Any] = struct.field(pytree_node=False)\n\n def unbox(self) -> A:\n return self.value\n\n def replace_boxed(self, val: B) -> 'NNXMeta[B]':\n return self.replace(value=val) # type: ignore\n\n def add_axis(self, index: int, params: dict[Any, Any]) -> 'NNXMeta[A]':\n # TODO: implement this, supporting hooks\n return self\n\n def remove_axis(self, index: int, params: dict[Any, Any]) -> 'NNXMeta[A]':\n # TODO: implement this, supporting hooks\n return self\n\n def get_partition_spec(self) -> jax.sharding.PartitionSpec:\n \"\"\"Returns the ``Partitionspec`` for this partitioned value.\"\"\"\n nnx_var = self.to_nnx_variable()\n spec = spmd.get_partition_spec(nnx_var).get_raw_value()\n assert isinstance(spec, jax.sharding.PartitionSpec)\n return spec\n\n def to_nnx_variable(self) -> variablelib.Variable:\n return self.var_type(self.value, **self.metadata)\n\n\ndef is_vanilla_variable(vs: variablelib.Variable) -> bool:\n \"\"\"A variable is vanilla if its metadata is essentially blank.\n\n Returns False only if it has non-empty hooks or any non-built-in attribute.\n \"\"\"\n for key, value in vs.get_metadata().items():\n if key in variablelib.Variable.required_metadata:\n continue\n if key.endswith('_hooks') and value == ():\n continue\n return False\n return True\n\n\ndef to_linen_var(vs: variablelib.Variable) -> meta.AxisMetadata:\n metadata = vs.get_metadata()\n if 'linen_meta_type' in metadata:\n linen_type = metadata['linen_meta_type']\n if hasattr(linen_type, 'from_nnx_metadata'):\n return linen_type.from_nnx_metadata({'value': vs.get_value(), **metadata})\n return linen_type(vs.get_value(), **metadata)\n if is_vanilla_variable(vs):\n return vs.get_value()\n return NNXMeta(type(vs), vs.get_value(), metadata)\n\n\ndef get_col_name(keypath: tp.Sequence[Any]) -> str:\n \"\"\"Given the keypath of a Flax variable type, return its Linen collection name.\"\"\"\n # Infer variable type from the leaf's path, which contains its Linen collection name\n assert isinstance(keypath[0], jax.tree_util.DictKey)\n return str(keypath[0].key)\n\n\ndef to_nnx_var(col: str, x: meta.AxisMetadata | Any) -> variablelib.Variable:\n \"\"\"Convert a Linen variable to an NNX variable.\"\"\"\n vtype = variablelib.variable_type_from_name(col, allow_register=True)\n if isinstance(x, NNXMeta):\n assert vtype == x.var_type, f'Type stored in NNXMeta {x.var_type} != type inferred from collection name {vtype}'\n return x.to_nnx_variable()\n if isinstance(x, meta.AxisMetadata):\n x_metadata = vars(x)\n if hasattr(x, 'to_nnx_metadata'):\n x_metadata = x.to_nnx_metadata()\n assert hasattr(x, 'value')\n return vtype(**x_metadata, linen_meta_type=type(x))\n return vtype(x)\n\n\ndef _recursive_merge(dict1, dict2):\n \"\"\"Recursively merge two dicts.\"\"\"\n flat_map = traversals.flatten_mapping(dict1)\n flat_map |= traversals.flatten_mapping(dict2)\n return traversals.unflatten_mapping(flat_map)\n\n\ndef linen_vars_to_nnx_attrs(variables: tp.Mapping[str, Any]) -> dict[str, Any]:\n \"\"\"Convert a dict of Linen-style variables to NNX variables.\"\"\"\n nnx_vars = jax.tree_util.tree_map_with_path(\n lambda kp, x: to_nnx_var(get_col_name(kp), x),\n variables,\n is_leaf=lambda x: not isinstance(x, dict),\n )\n flat_paths: dict[tuple, tp.Any] = {}\n for col_name, col_variables in nnx_vars.items(): # pylint: disable=unused-variable\n for path, variable in traversals.flatten_mapping(col_variables).items():\n if path in flat_paths:\n raise ValueError(\n f\"Found duplicate variable path {path} with variables \"\n f\"{flat_paths[path]} and {variable}. \"\n \"This is not allowed in NNX.\"\n )\n flat_paths[path] = variable\n\n nnx_vars = traversals.unflatten_mapping(flat_paths)\n return nnx_vars\n\n\ndef nnx_attrs_to_linen_vars(nnx_attrs: dict) -> dict:\n \"\"\"Convert a dict of NNX variables to Linen-style variables.\"\"\"\n linen_structured = {}\n for kp, v in traversals.flatten_mapping(nnx_attrs).items():\n if isinstance(v, variablelib.Variable):\n col_name = variablelib.variable_name_from_type(type(v))\n v = to_linen_var(v.to_state())\n else:\n raise ValueError(f'Cannot infer collection name from value: {v}')\n linen_structured[(col_name, *kp)] = v\n variables = traversals.unflatten_mapping(linen_structured)\n return variables\n\n\n\ndef with_partitioning(\n fn: tp.Callable[..., tp.Any],\n names: LogicalNames,\n mesh: jax.sharding.Mesh | None = None,\n) -> tp.Callable[..., meta.Partitioned[tp.Any]]:\n \"\"\"Same interface as Linen, but calls NNX `with_partitioning` within.\"\"\"\n return spmd.with_partitioning(fn, names, mesh,\n linen_meta_type=meta.Partitioned)" }, { "path": "flax/nnx/bridge/wrappers.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom functools import partial\nimport typing as tp\nfrom typing import Any\nimport warnings\nimport dataclasses\n\nfrom flax import linen\nfrom flax import core\nfrom flax import nnx\nfrom flax.core import FrozenDict\nfrom flax.core import meta\nfrom flax.nnx import graphlib\nfrom flax.nnx import variablelib\nfrom flax.nnx.bridge import variables as bv\nfrom flax.nnx.bridge import module as bdg_module\nfrom flax.nnx.module import Module\nfrom flax.nnx.statelib import State\nfrom flax.nnx.pytreelib import Pytree\nfrom flax.nnx.rnglib import Rngs\nimport jax\nfrom jax import tree_util as jtu\n\nM = tp.TypeVar('M', bound=Module)\n\n# Flax-like style is NNX\n@dataclasses.dataclass\nclass Functional(tp.Generic[M]):\n module_type: tp.Type[M]\n graphdef: tp.Optional[graphlib.GraphDef[M]]\n args: tuple[tp.Any, ...]\n kwargs: dict[str, tp.Any]\n\n def init(self, *, rngs: tp.Optional[Rngs] = None) -> State:\n kwargs = {}\n if rngs is not None:\n kwargs['rngs'] = rngs\n module = self.module_type(*self.args, **self.kwargs, **kwargs)\n graphdef, state = nnx.split(module)\n self.graphdef = graphdef\n return state # type: ignore\n\n def apply(self, *states: tp.Any):\n assert self.graphdef is not None\n return self.graphdef.apply(*states)\n\n\ndef functional(cls: tp.Type[M]) -> tp.Callable[..., Functional[M]]:\n def _functional_constructor(*args: tp.Any, **kwargs: tp.Any) -> Functional[M]:\n return Functional(cls, None, args, kwargs)\n\n return _functional_constructor\n\n\ndef _set_initializing(module: Module, initializing: bool):\n for _, value in graphlib.iter_graph(module, graph=True):\n if isinstance(value, Pytree):\n value._pytree__state._initializing = initializing\n\n\ndef lazy_init(fn: Module | tp.Callable[..., tp.Any], *args, **kwargs):\n \"\"\"To run through an arbitrary nnx.Module method and initialize all its needed state.\n\n Here used to trigger initialization of all `LinenToNNX` module variables.\"\"\"\n if isinstance(fn, Module):\n module = fn\n assert callable(fn)\n else:\n if not (hasattr(fn, '__self__') and isinstance(fn.__self__, Module)):\n raise ValueError(f'{fn = } needs to be a method of an NNX Module.')\n module = fn.__self__\n _set_initializing(module, True)\n try:\n _ = fn(*args, **kwargs)\n finally:\n _set_initializing(module, False)\n return fn\n\ndef current_linen_module() -> linen.Module | None:\n \"\"\"Get the current Linen module from the Linen context.\"\"\"\n if linen.module._context.module_stack: # pylint: disable=W0212\n return linen.module._context.module_stack[-1] # pylint: disable=W0212\n return None\n\nclass ToNNX(Module):\n \"\"\"A wrapper to turn any Linen module into an NNX module.\n\n The result NNX module can be used standalone with all NNX APIs, or as a submodule of\n another NNX module.\n\n Since Linen module initialization requires a sample input, you need to call `lazy_init`\n with an argument to initialize the variables.\n\n Example::\n\n >>> from flax import linen as nn, nnx\n >>> import jax\n >>> linen_module = nn.Dense(features=64)\n >>> x = jax.numpy.ones((1, 32))\n >>> # Like Linen init(), initialize with a sample input\n >>> model = nnx.bridge.ToNNX(linen_module, rngs=nnx.Rngs(0)).lazy_init(x)\n >>> # Like Linen apply(), but using NNX's direct call method\n >>> y = model(x)\n >>> model.kernel.shape\n (32, 64)\n\n Args:\n module: The Linen Module instance.\n rngs: The `nnx.Rngs` instance being passed to any NNX module.\n\n Returns:\n A stateful NNX module that behaves the same as the wrapped Linen module.\n \"\"\"\n\n def __init__(\n self,\n module: linen.Module,\n rngs: Rngs | jax.Array | None = None,\n ):\n self.to_nnx__module = module\n\n self.to_nnx__rngs: Rngs | None\n if isinstance(rngs, jax.Array):\n self.to_nnx__rngs = Rngs(params=rngs)\n elif isinstance(rngs, nnx.Rngs):\n self.to_nnx__rngs = rngs.fork()\n else:\n self.to_nnx__rngs = rngs\n\n @property\n def rngs(self) -> Rngs | None:\n warnings.warn(\n '`ToNNX.rngs` is deprecated. Please use `to_nnx__rngs` instead.',\n DeprecationWarning,\n )\n return self.to_nnx__rngs\n\n @property\n def module(self) -> linen.Module:\n warnings.warn(\n '`ToNNX.module` is deprecated. Please use `to_nnx__module` instead.',\n DeprecationWarning,\n )\n return self.to_nnx__module\n\n def _setattr(self, name, value):\n if not nnx.is_data(value) and nnx.has_data(value):\n value = nnx.data(value)\n super()._setattr(name, value)\n\n def lazy_init(self, *args, **kwargs):\n \"\"\"A shortcut of calling `nnx.bridge.lazy_init()` upon this module.\"\"\"\n return lazy_init(self, *args, **kwargs)\n\n def __getattr__(self, name: str):\n if hasattr(super(), name):\n return super().__getattribute__(name)\n maybe_method = getattr(type(self.to_nnx__module), name, None)\n if callable(maybe_method):\n method = partial(self.__call__, method=maybe_method)\n method.__self__ = self\n return method\n return super().__getattribute__(name)\n\n def __call__(\n self,\n *args: Any,\n rngs: Rngs | jax.Array | None = None,\n method: tp.Callable[..., Any] | str | None = None,\n mutable: tp.Any = None,\n **kwargs: Any,\n ) -> Any:\n # Shape-based lazy init of the flax variables\n if rngs is None:\n rngs = self.to_nnx__rngs\n if isinstance(rngs, nnx.Rngs):\n _rngs = {name: stream() for name, stream in rngs.items()}\n elif isinstance(rngs, jax.Array):\n _rngs = {'params': rngs}\n else:\n _rngs = {}\n # rename default to params\n if 'params' not in _rngs and 'default' in _rngs:\n _rngs['params'] = _rngs.pop('default')\n if self._pytree__state.initializing:\n out, updates = self.to_nnx__module.init_with_output(_rngs, *args, method=method, **kwargs)\n else:\n nnx_attrs = {\n k: v\n for k, v in vars(self).items()\n if not k.startswith('to_nnx__') and not k.startswith('_pytree__')\n }\n variables = bv.nnx_attrs_to_linen_vars(nnx_attrs)\n\n # Get `mutable` from top level bridge.Module context if any\n if mutable is not None:\n pass\n elif (m := bdg_module.current_module()) is not None: # type: ignore[assignment]\n assert m.scope is not None\n mutable = m.scope.mutable\n elif (m := current_linen_module()) is not None: # type: ignore[assignment]\n assert m.scope is not None\n mutable = m.scope.mutable\n else:\n mutable = False\n\n out = self.to_nnx__module.apply(\n variables, *args, rngs=_rngs, method=method, mutable=mutable, **kwargs\n )\n\n # Split out the updates if `mutable` is passed into the Flax module\n if mutable is not False:\n out, updates = out\n else:\n updates = None\n\n # Split out the updates if `mutable` is passed into the Flax module\n if updates:\n nnx_attrs = bv.linen_vars_to_nnx_attrs(updates)\n # nnx.update(self, nnx_attrs)\n # TODO(cgarciae): ideally we just do an update but currently dictionaries don't allow\n # insertion of new keys, we need to enable this in NNX to simplify the code bellow\n # to the simple nnx.update(self, nnx_attrs) above.\n for attr_name, value in nnx_attrs.items():\n if hasattr(self, attr_name) and isinstance(value, dict):\n original_value = getattr(self, attr_name)\n new_values = bv._recursive_merge(original_value, value)\n setattr(self, attr_name, nnx.data(new_values))\n else:\n setattr(self, attr_name, nnx.data(value))\n\n return out\n\n\ndef linen_rngs_dict(linen_module: linen.Module, add_default: bool = False):\n \"\"\"Given a module, split out one of its every active RNG key collections.\"\"\"\n assert linen_module.scope is not None, 'linen_rngs_dict() must be called inside a Linen module.'\n rngs: dict[str, tp.Any] = {\n name: linen_module.make_rng(name)\n for name in linen_module.scope.rngs.keys()\n }\n if add_default and 'default' not in rngs:\n rngs['default'] = 0\n return rngs\n\ndef _get_module_method(module, method: tp.Callable[..., Any] | str | None):\n \"\"\"Get a callable method from the module, or raise TypeError.\"\"\"\n if method is None:\n method = '__call__'\n\n if isinstance(method, str):\n attribute_name = method\n method = getattr(type(module), attribute_name)\n if not callable(method):\n class_name = type(module).__name__\n raise TypeError(\n f\"'{class_name}.{attribute_name}' must be a callable, got\"\n f' {type(method)}.'\n )\n if not callable(method):\n class_name = type(module).__name__\n raise TypeError(\n f\"'{method}' must be a callable, got {type(method)}.\"\n )\n\n return method\n\nclass ToLinen(linen.Module):\n \"\"\"A wrapper to turn any NNX module into a Linen module.\n\n The result Linen module can be used standalone with all Linen APIs, or as a\n submodule of\n another Linen module.\n\n Since NNX modules are stateful and owns the state, we only create it once\n during init\n time, and will track its state and static data as separate variables.\n\n Example::\n\n >>> from flax import linen as nn, nnx\n >>> import jax\n >>> model = nnx.bridge.ToLinen(nnx.Linear, args=(32, 64))\n >>> x = jax.numpy.ones((1, 32))\n >>> y, variables = model.init_with_output(jax.random.key(0), x)\n >>> y.shape\n (1, 64)\n >>> variables['params']['kernel'].shape\n (32, 64)\n >>> # The static GraphDef of the underlying NNX module\n >>> variables.keys()\n dict_keys(['params'])\n\n Args:\n nnx_class: The NNX Module class (not instance!).\n args: The arguments that normally would be passed in to create the NNX\n module.\n kwargs: The keyword arguments that normally would be passed in to create the\n NNX module.\n skip_rng: True if this NNX module doesn't need `rngs` arg during\n initialization (not common).\n\n Returns:\n A stateful NNX module that behaves the same as the wrapped Linen module.\n \"\"\"\n\n nnx_class: tp.Callable[..., Module]\n args: tp.Sequence = ()\n kwargs: tp.Mapping[str, tp.Any] = FrozenDict({})\n skip_rng: bool = False\n metadata_fn: tp.Callable[[variablelib.Variable], tp.Any] | None = bv.to_linen_var\n\n @linen.compact\n def __call__(\n self, *args, nnx_method: tp.Callable[..., Any] | str | None = None, **kwargs\n ):\n def _module_kwargs():\n maybe_add_default = not self.is_initializing()\n module_kwargs = dict(self.kwargs)\n if not self.skip_rng:\n module_kwargs['rngs'] = nnx.Rngs(\n **linen_rngs_dict(self, add_default=maybe_add_default)\n )\n return module_kwargs\n\n # init codepath\n if self.is_initializing():\n module = self.nnx_class(*self.args, **_module_kwargs())\n # TODO: add lazy_init here in case there's an `ToNNX` submodule under `module`.\n # update linen variables before call module to save initial state\n self._update_variables(module)\n method_fn = _get_module_method(module, nnx_method)\n out = method_fn(module, *args, **kwargs)\n return out\n\n # create the nnx module\n module = self.nnx_class(*self.args, **_module_kwargs())\n\n # update nnx module from linen variables\n def maybe_unbox(x):\n if isinstance(x, meta.AxisMetadata):\n return x.unbox()\n return x\n states = jtu.tree_map(\n maybe_unbox,\n list(core.unfreeze(self.variables).values()), # type: ignore[wrong-arg-types, arg-type]\n is_leaf=lambda x: isinstance(x, meta.AxisMetadata),\n )\n if not states:\n states = ({},)\n\n new_state = nnx.merge_state(*states)\n new_state_flat = nnx.traversals.flatten_mapping(new_state)\n current_state_flat = nnx.traversals.flatten_mapping(nnx.state(module))\n unknown_state_flat = {path: v for path, v in new_state_flat.items() if path not in current_state_flat}\n\n if unknown_state_flat:\n paths_str = \"\"\n for path, _ in unknown_state_flat.items():\n paths_str += f\"\\n - {'/'.join(map(str, path))}\"\n\n warnings.warn(f\"Found unknown module paths in incoming state:{paths_str}\")\n\n nnx.update(module, new_state)\n\n method_fn = _get_module_method(module, nnx_method)\n out = method_fn(module, *args, **kwargs)\n self._update_variables(module)\n return out\n\n def __getattr__(self, name: str):\n if hasattr(super(), name):\n return super().__getattribute__(name)\n if name in self.kwargs:\n return self.kwargs[name]\n maybe_method = getattr(self.nnx_class, name, None)\n if callable(maybe_method):\n method = partial(self.__call__, nnx_method=maybe_method)\n method.__self__ = self\n return method\n return super().__getattribute__(name)\n\n def _update_variables(self, module):\n \"\"\"Store the NNX module's graph def and state inside Linen module variables.\"\"\"\n state = nnx.state(module, nnx.Not(nnx.RngState))\n\n collection_flat_state: dict[str, list[tuple[tuple[tp.Any, ...], tp.Any]]] = {}\n\n # group state by collection\n for path, leaf in nnx.to_flat_state(state):\n type_ = type(leaf)\n collection = variablelib.variable_name_from_type(\n type_, allow_register=True\n )\n if collection not in collection_flat_state:\n collection_flat_state[collection] = []\n collection_flat_state[collection].append((path, leaf))\n\n # update linen variables\n for collection, flat_state in collection_flat_state.items():\n if self.is_mutable_collection(collection):\n\n def _to_linen_var(x):\n if isinstance(x, nnx.Variable):\n if self.metadata_fn is not None:\n return self.metadata_fn(x) # pylint: disable=too-many-function-args\n else:\n return x.get_value()\n return x\n\n collection_state = nnx.traversals.unflatten_mapping(flat_state)\n collection_state = jax.tree.map(\n _to_linen_var,\n collection_state,\n is_leaf=lambda x: isinstance(x, nnx.Variable),\n )\n for k, v in collection_state.items():\n self.put_variable(collection, k, v)\n\n\n\nclass _Missing:\n ...\n\n\n_MISSING = _Missing()\n\n\ndef to_linen(\n nnx_class: tp.Callable[..., Module],\n *args,\n metadata_fn: (\n tp.Callable[[variablelib.Variable], tp.Any] | None\n ) = bv.to_linen_var,\n name: str | None = None,\n skip_rng: bool = False,\n abstract_init: bool = True,\n **kwargs,\n):\n \"\"\"Shortcut of `nnx.bridge.ToLinen` if user is not changing any of its default fields.\"\"\"\n return ToLinen(\n nnx_class,\n args=args,\n kwargs=FrozenDict(kwargs),\n metadata_fn=metadata_fn,\n skip_rng=skip_rng,\n name=name,\n )\n\ndef to_linen_class(\n base_nnx_class: type[M],\n base_metadata_fn: tp.Callable[[variablelib.VariableState], tp.Any] | None = bv.to_linen_var,\n base_skip_rng: bool = False,\n **partial_kwargs: tp.Any,\n) -> type[ToLinen]:\n \"\"\"Dynamically wraps an NNX module class into a Flax Linen module class.\"\"\"\n\n class ToLinenPartial(ToLinen):\n \"\"\"A dynamically created Linen Module that wraps a specific NNX Module.\n\n This class is not meant to be used directly. Instead, it is created and\n returned by the `to_linen_class` function. It acts as a \"partially applied\"\n version of the `ToLinen` wrapper, where the NNX module to be wrapped and\n its default arguments are pre-configured.\n\n When you instantiate this class, it behaves like a standard Linen module.\n The arguments you provide during instantiation can override the defaults\n that were set when this class was created by `to_linen_class`.\n\n For example:\n >>> from flax import linen as nn, nnx\n >>> from maxtext.src.maxtext.layers import linears\n >>> # Create a specialized Linen wrapper for linears.DenseGeneral\n >>> LinenDenseGeneral = to_linen_class(linears.DenseGeneral)\n >>> # Now, LinenDenseGeneral can be used like a regular Linen module\n >>> class MyModel(nn.Module):\n ... def setup(self):\n ... # Instantiate the wrapped linears.DenseGeneral with its arguments\n ... self.dense = LinenDenseGeneral(\n ... in_features_shape=10, out_features_shape=5\n ... )\n ... def __call__(self, x):\n ... return self.dense(x)\n\n Attributes:\n (The attributes are dynamically set by the `ToLinen` parent class based\n on the arguments provided during instantiation.)\n \"\"\"\n\n def __init_subclass__(cls, **kwargs):\n super().__init_subclass__(**kwargs)\n\n def __init__(\n self,\n args=None,\n kwargs=None,\n nnx_class=None,\n skip_rng=None,\n metadata_fn=None,\n name=_MISSING,\n parent=_MISSING,\n **other_kwargs,\n ):\n linen_kwargs = {}\n if not isinstance(parent, _Missing):\n linen_kwargs[\"parent\"] = parent\n if not isinstance(name, _Missing):\n linen_kwargs[\"name\"] = name\n ToLinen.__init__(\n self,\n nnx_class=nnx_class or base_nnx_class,\n args=args or (),\n metadata_fn=metadata_fn or base_metadata_fn,\n skip_rng=skip_rng or base_skip_rng,\n kwargs=FrozenDict({**partial_kwargs, **(kwargs or {}), **other_kwargs}),\n **linen_kwargs,\n )\n\n cls.__init__ = __init__\n\n return ToLinenPartial\n" }, { "path": "flax/nnx/compat.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"Compat API.\n\nCompatibility wrappers for NNX APIs. Each function in this module mirrors the\ncorresponding ``nnx.*`` API but enforces ``graph=True`` (and\n``graph_updates=True`` for transforms), preserving the pre-tree-mode behavior.\n\"\"\"\n\nimport functools\n\nfrom flax.nnx import graphlib as _graphlib\nfrom flax.nnx import module as _module\nfrom flax.nnx import rnglib as _rnglib\nfrom flax.nnx.transforms import autodiff as _autodiff\nfrom flax.nnx.transforms import compilation as _compilation\nfrom flax.nnx.transforms import iteration as _iteration\nfrom flax.nnx.transforms import transforms as _transforms\nfrom flax.nnx import spmd as _spmd\n\n# graphlib\nsplit = functools.partial(_graphlib.split, graph=True)\nstate = functools.partial(_graphlib.state, graph=True)\nclone = functools.partial(_graphlib.clone, graph=True)\ngraphdef = functools.partial(_graphlib.graphdef, graph=True)\nflatten = functools.partial(_graphlib.flatten, graph=True)\niter_graph = functools.partial(_graphlib.iter_graph, graph=True)\nrecursive_map = functools.partial(_graphlib.recursive_map, graph=True)\n\n# module\nview = functools.partial(_module.view, graph=True)\nview_info = functools.partial(_module.view_info, graph=True)\niter_modules = functools.partial(_module.iter_modules, graph=True)\niter_children = functools.partial(_module.iter_children, graph=True) # type: ignore[has-type]\n\n# rnglib\nsplit_rngs = functools.partial(_rnglib.split_rngs, graph=True)\nfork_rngs = functools.partial(_rnglib.fork_rngs, graph=True)\nreseed = functools.partial(_rnglib.reseed, graph=True)\nbackup_keys = functools.partial(_rnglib.backup_keys, graph=True)\n\n# transforms - compilation\njit = functools.partial(_compilation.jit, graph=True, graph_updates=True)\nshard_map = functools.partial(\n _compilation.shard_map, graph=True, graph_updates=True\n)\n\n# transforms - autodiff\ngrad = functools.partial(_autodiff.grad, graph=True, graph_updates=True)\nvalue_and_grad = functools.partial(\n _autodiff.value_and_grad, graph=True, graph_updates=True\n)\ncustom_vjp = functools.partial(\n _autodiff.custom_vjp, graph=True, graph_updates=True\n)\nvjp = functools.partial(_autodiff.vjp, graph=True, graph_updates=True)\njvp = functools.partial(_autodiff.jvp, graph=True, graph_updates=True)\nremat = functools.partial(_autodiff.remat, graph=True, graph_updates=True)\n\n# transforms - iteration\nvmap = functools.partial(_iteration.vmap, graph=True, graph_updates=True)\nscan = functools.partial(_iteration.scan, graph=True, graph_updates=True)\npmap = functools.partial(_iteration.pmap, graph=True, graph_updates=True)\nwhile_loop = functools.partial(\n _iteration.while_loop, graph=True, graph_updates=True\n)\nfori_loop = functools.partial(\n _iteration.fori_loop, graph=True, graph_updates=True\n)\n\n# transforms - general\neval_shape = functools.partial(\n _transforms.eval_shape, graph=True, graph_updates=True\n)\ncond = functools.partial(_transforms.cond, graph=True, graph_updates=True)\nswitch = functools.partial(_transforms.switch, graph=True, graph_updates=True)\ncheckify = functools.partial(\n _transforms.checkify, graph=True, graph_updates=True\n)\n\n# spmd\nget_abstract_model = functools.partial(_spmd.get_abstract_model, graph=True)\n" }, { "path": "flax/nnx/extract.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport abc\nimport functools\nimport typing as tp\nfrom collections import namedtuple\n\nimport jax\n\nfrom flax import struct\nfrom flax import typing\nfrom flax.nnx.pytreelib import Pytree\nfrom flax.typing import Missing, PathParts\nfrom flax.nnx import graphlib, variablelib\n\n\nA = tp.TypeVar('A')\nIndex = int\nKeyPath = tuple[tp.Hashable, ...]\nPrefix = tp.Any\nLeaf = tp.Any\n\n\nclass PrefixMapping(abc.ABC):\n @abc.abstractmethod\n def map_prefix(\n self,\n path: typing.PathParts,\n variable: variablelib.Variable,\n /,\n ) -> tp.Any: ...\n\ndef check_consistent_aliasing(\n node: tp.Any,\n prefix: tp.Any,\n /,\n *,\n node_prefixes: dict[int, list[tuple[PathParts, tp.Any]]] | None = None,\n):\n \"\"\"Check for consistent aliasing of nodes when extracting graph.\"\"\"\n if node_prefixes is None:\n node_prefixes = {}\n\n # Store variable references for error messages\n node_id_to_variable: dict[int, tp.Any] = {}\n\n # collect all paths and prefixes for each node\n for path, value in graphlib.iter_graph(node, graph=True):\n if graphlib.is_graph_node(value) or isinstance(value, graphlib.Variable):\n if isinstance(value, Pytree):\n value._check_valid_context(\n lambda: f'Trying to extract graph node from different trace level, got {value!r}'\n )\n if isinstance(value, graphlib.Variable):\n if not value._can_update:\n raise ValueError(\n f'Cannot extract graph node from different trace level, got {value!r}'\n )\n if isinstance(prefix, PrefixMapping):\n variable_prefix = prefix.map_prefix(path, value)\n else:\n variable_prefix = prefix\n\n value_id = id(value)\n node_id_to_variable[value_id] = value\n if value_id in node_prefixes:\n paths_prefixes = node_prefixes[value_id]\n paths_prefixes.append((path, variable_prefix))\n else:\n node_prefixes[value_id] = [(path, variable_prefix)]\n\n # check for inconsistent aliasing\n node_msgs = []\n for node_id, paths_prefixes in node_prefixes.items():\n unique_prefixes = {prefix for _, prefix in paths_prefixes}\n if len(unique_prefixes) > 1:\n path_prefix_repr = '\\n'.join(\n f' {\"/\".join(map(str,path)) if path else \"\"}: {prefix}'\n for path, prefix in paths_prefixes\n )\n # Get the variable type name if available\n if node_id in node_id_to_variable:\n variable = node_id_to_variable[node_id]\n node_type_name = type(variable).__name__\n else:\n node_type_name = f'Node ID: {node_id}'\n\n nodes_msg = f'Node: {node_type_name}\\n{path_prefix_repr}'\n node_msgs.append(nodes_msg)\n\n if node_msgs:\n raise ValueError(\n 'Inconsistent aliasing detected. The following nodes have different prefixes:\\n'\n + '\\n'.join(node_msgs)\n )\n\n\ndef check_consistent_aliasing2(\n node: tp.Any,\n prefix: tp.Any,\n /,\n *,\n base_path: tuple[tp.Any, ...] = (),\n node_prefixes: dict[int, list[tuple[PathParts, tp.Any]]] | None = None,\n):\n if node_prefixes is None:\n node_prefixes = {}\n\n node_id_to_variable: dict[int, tp.Any] = {}\n\n for path, value in graphlib.iter_graph(node, graph=True):\n path = base_path + path\n if graphlib.is_graph_node(value) or isinstance(value, graphlib.Variable):\n value_id = id(value)\n node_id_to_variable[value_id] = value\n if value_id in node_prefixes:\n node_prefixes[value_id].append((path, prefix))\n else:\n node_prefixes[value_id] = [(path, prefix)]\n\n node_msgs = []\n for node_id, paths_prefixes in node_prefixes.items():\n unique_prefixes = {p for _, p in paths_prefixes}\n if len(unique_prefixes) > 1:\n path_prefix_repr = '\\n'.join(\n f' {\"/\".join(map(str,path)) if path else \"\"}: {p}'\n for path, p in paths_prefixes\n )\n if node_id in node_id_to_variable:\n variable = node_id_to_variable[node_id]\n node_type_name = type(variable).__name__\n else:\n node_type_name = f'Node ID: {node_id}'\n\n node_msgs.append(f'Node: {node_type_name}\\n{path_prefix_repr}')\n\n if node_msgs:\n raise ValueError(\n 'Inconsistent aliasing detected. The following nodes have different prefixes:\\n'\n + '\\n'.join(node_msgs)\n )\n\n# -----------------------------\n# to_tree/from_tree\n# -----------------------------\n\ndef broadcast_prefix(\n prefix_tree: tp.Any,\n full_tree: tp.Any,\n prefix_is_leaf: tp.Callable[[tp.Any], bool] | None = None,\n tree_is_leaf: tp.Callable[[tp.Any], bool] | None = None,\n) -> list[tp.Any]:\n # If prefix_tree is not a tree prefix of full_tree, this code can raise a\n # ValueError; use prefix_errors to find disagreements and raise more precise\n # error messages.\n result = []\n num_leaves = lambda t: jax.tree_util.tree_structure(\n t, is_leaf=tree_is_leaf\n ).num_leaves\n add_leaves = lambda x, subtree: result.extend([x] * num_leaves(subtree))\n jax.tree.map(\n add_leaves,\n prefix_tree,\n full_tree,\n is_leaf=lambda x: isinstance(x, variablelib.Variable)\n or graphlib.is_graph_node(x)\n or (prefix_is_leaf is not None and prefix_is_leaf(x)),\n )\n return result\n\ndef broadcast_prefix2(\n prefix_tree: tp.Any,\n full_tree: tp.Any,\n is_leaf: tp.Callable[[tp.Any], bool] | None = None,\n) -> tuple[list[KeyPath], list[tp.Any]]:\n paths: list[KeyPath] = []\n leaves: list[tp.Any] = []\n num_leaves = lambda t: jax.tree_util.tree_structure(t).num_leaves\n def add_leaves(path, x, subtree):\n n = num_leaves(subtree)\n paths.extend([path] * n)\n leaves.extend([x] * n)\n jax.tree.map_with_path(add_leaves, prefix_tree, full_tree, is_leaf=is_leaf)\n return paths, leaves\n\ndef broadcast_prefix_map(\n f: tp.Callable[..., tp.Any],\n prefix_tree: tp.Any,\n full_tree: tp.Any,\n *rest: tp.Any,\n is_leaf: tp.Callable[[tp.Any], bool] | None = None,\n) -> tp.Any:\n paths, prefix_leaves = broadcast_prefix2(prefix_tree, full_tree, is_leaf=is_leaf)\n leaves, treedef = jax.tree_util.tree_flatten(full_tree, is_leaf=is_leaf)\n full_prefix_tree = treedef.unflatten(prefix_leaves)\n return jax.tree.map_with_path(f, full_prefix_tree, full_tree, *rest, is_leaf=is_leaf)\n\n\nclass GraphDefState(struct.PyTreeNode):\n graphdef: graphlib.GraphDef[tp.Any] = struct.field(pytree_node=False)\n state: graphlib.GraphState = struct.field(pytree_node=True)\n\nS = tp.TypeVar(\n 'S', bound=graphlib.GraphState | graphlib.GraphFlatState | list[tp.Any]\n)\n\nclass NodeStates(struct.PyTreeNode):\n _graphdef: graphlib.GraphDef[tp.Any] | None\n states: tuple[tp.Any, ...]\n metadata: tp.Any = struct.field(pytree_node=False)\n\n @property\n def graphdef(self) -> graphlib.GraphDef[tp.Any]:\n if self._graphdef is None:\n raise ValueError('No graphdef available')\n return self._graphdef\n\n @property\n def state(self) -> tp.Any:\n if len(self.states) != 1:\n raise ValueError(\n f'Expected exactly one GraphDefState, got {len(self.states)}'\n )\n return self.states[0]\n\n @classmethod\n def from_split(\n cls,\n graphdef: graphlib.GraphDef[tp.Any] | None,\n state: tp.Any,\n /,\n *states: tp.Any,\n metadata: tp.Any = None,\n ):\n return cls(_graphdef=graphdef, states=(state, *states), metadata=metadata)\n\n @classmethod\n def from_states(\n cls,\n state: tp.Any,\n *states: tp.Any,\n ):\n return cls(_graphdef=None, states=(state, *states), metadata=None)\n\n @classmethod\n def from_prefixes(\n cls,\n prefixes: tp.Iterable[tp.Any],\n /,\n *,\n metadata: tp.Any = None,\n ):\n return cls(_graphdef=None, states=tuple(prefixes), metadata=metadata)\n\n\ndef default_split_fn(\n ctx: graphlib.SplitContext, path: KeyPath, prefix: Prefix, leaf: Leaf\n) -> tp.Any:\n return NodeStates.from_split(*ctx.split(leaf))\n\n\ndef to_tree(\n tree,\n /,\n *,\n prefix: tp.Any = Missing,\n split_fn: tp.Callable[\n [graphlib.SplitContext, KeyPath, Prefix, Leaf], tp.Any\n ] = default_split_fn,\n map_non_graph_nodes: bool = False,\n ctxtag: tp.Hashable | None = None,\n check_aliasing: bool = True,\n) -> tp.Any:\n if prefix is Missing or prefix is None:\n # fast path, no need for prefix broadcasting or consistent aliasing checks\n with graphlib.split_context(ctxtag) as split_ctx:\n return jax.tree.map(\n lambda x: split_fn(split_ctx, (), prefix, x)\n if map_non_graph_nodes\n or graphlib.is_graph_node(x)\n or isinstance(x, variablelib.Variable)\n else x,\n tree,\n is_leaf=lambda x: isinstance(x, variablelib.Variable)\n or graphlib.is_graph_node(x),\n )\n leaf_prefixes = broadcast_prefix(\n prefix,\n tree,\n prefix_is_leaf=lambda x: x is None\n or isinstance(x, variablelib.Variable)\n or graphlib.is_graph_node(x),\n tree_is_leaf=lambda x: isinstance(x, variablelib.Variable)\n or graphlib.is_graph_node(x),\n )\n leaf_keys, treedef = jax.tree_util.tree_flatten_with_path(\n tree,\n is_leaf=lambda x: isinstance(x, variablelib.Variable)\n or graphlib.is_graph_node(x),\n )\n\n assert len(leaf_keys) == len(leaf_prefixes)\n leaves_out = []\n node_prefixes: dict[int, list[tuple[PathParts, tp.Any]]] = {}\n\n with graphlib.split_context(ctxtag) as split_ctx:\n for (keypath, leaf), leaf_prefix in zip(leaf_keys, leaf_prefixes):\n if graphlib.is_graph_node(leaf) or isinstance(leaf, variablelib.Variable):\n if check_aliasing:\n check_consistent_aliasing(\n leaf, leaf_prefix, node_prefixes=node_prefixes\n )\n tree_node = split_fn(split_ctx, keypath, leaf_prefix, leaf)\n leaves_out.append(tree_node)\n else:\n if map_non_graph_nodes:\n leaf = split_fn(split_ctx, keypath, leaf_prefix, leaf)\n leaves_out.append(leaf)\n\n pytree_out = jax.tree.unflatten(treedef, leaves_out)\n return pytree_out\n\n\ndef to_tree2(\n tree,\n /,\n *,\n prefix: tp.Any = Missing,\n check_aliasing: bool = True,\n) -> tp.Any:\n \"\"\"to_tree2 has two main tasks:\n\n 1. Convert all graph nodes to NodeStates (a tree representation).\n 2. Check all Variables are aliased consistently given the prefix tree,\n e.g. vmap's in/out_axes arguments.\n\n Each NodeState contains the `GraphDef` and State for each object, these\n are generated using `graphlib.flatten`. `extract.broadcast_prefix` is used\n to calculate the prefix for each node, `check_consistent_aliasing2` traverses\n the nodes subgraph and checks for Variable aliasing.\n \"\"\"\n ref_index: graphlib.RefMap = graphlib.RefMap()\n\n def _to_node_states(leaf):\n if not (graphlib.is_graph_node(leaf) or isinstance(leaf, variablelib.Variable)):\n return leaf\n graphdef, flat_state = graphlib.flatten(\n leaf, ref_index=ref_index, graph=True\n )\n (state,) = graphlib._to_nested_state(graphdef, (flat_state,))\n return NodeStates.from_split(graphdef, state)\n\n is_leaf = lambda x: (\n isinstance(x, variablelib.Variable) or graphlib.is_graph_node(x)\n )\n\n if prefix is Missing or prefix is None:\n return jax.tree.map(_to_node_states, tree, is_leaf=is_leaf)\n\n leaf_prefixes = broadcast_prefix(\n prefix,\n tree,\n prefix_is_leaf=lambda x: x is None or is_leaf(x),\n tree_is_leaf=is_leaf,\n )\n leaf_paths, treedef = jax.tree_util.tree_flatten_with_path(tree, is_leaf=is_leaf)\n\n assert len(leaf_paths) == len(leaf_prefixes)\n leaves_out = []\n node_prefixes: dict[int, list[tuple[PathParts, tp.Any]]] = {}\n\n for (keypath, leaf), leaf_prefix in zip(leaf_paths, leaf_prefixes):\n if is_leaf(leaf):\n if check_aliasing:\n base_path = graphlib.jax_to_nnx_path(keypath)\n check_consistent_aliasing2(\n leaf, leaf_prefix, base_path=base_path, node_prefixes=node_prefixes\n )\n leaves_out.append(_to_node_states(leaf))\n else:\n leaves_out.append(leaf)\n\n return jax.tree.unflatten(treedef, leaves_out)\n\n\ndef from_tree2(tree: tp.Any, /) -> tp.Any:\n index_ref = graphlib.IndexMap()\n\n def _from_node_states(x):\n if not isinstance(x, NodeStates):\n return x\n state = graphlib._merge_to_flat_state(x.states)\n return graphlib.unflatten(\n x.graphdef, state, index_ref=index_ref,\n )\n\n return jax.tree.map(\n _from_node_states,\n tree,\n is_leaf=lambda x: (\n isinstance(x, NodeStates)\n or graphlib.is_graph_node(x)\n or isinstance(x, variablelib.Variable)\n ),\n )\n\n\ndef merge_tree_node(\n ctx: graphlib.MergeContext, path: KeyPath, prefix: Prefix, leaf: Leaf\n) -> tp.Any:\n if not isinstance(leaf, NodeStates):\n raise ValueError(f'Expected TreeNode, got {type(leaf)} at path {path}')\n return ctx.merge(leaf.graphdef, *leaf.states)\n\n\ndef is_tree_node(x):\n return isinstance(x, NodeStates)\n\n\ndef from_tree(\n tree: tp.Any,\n /,\n *,\n prefix: tp.Any = Missing,\n merge_fn: tp.Callable[\n [graphlib.MergeContext, KeyPath, Prefix, Leaf], tp.Any\n ] = merge_tree_node,\n is_node_leaf: tp.Callable[[Leaf], bool] = is_tree_node,\n is_leaf: tp.Callable[[Leaf], bool] = is_tree_node,\n map_non_graph_nodes: bool = False,\n is_inner: bool | None = None,\n ctxtag: tp.Hashable | None = None,\n) -> tp.Any:\n if prefix is Missing or prefix is None:\n # fast path, no need for prefix broadcasting or consistent aliasing checks\n with graphlib.merge_context(ctxtag, is_inner) as merge_ctx:\n\n def maybe_split(x):\n if (\n map_non_graph_nodes\n or is_node_leaf(x)\n or isinstance(x, variablelib.Variable)\n ):\n return merge_fn(merge_ctx, (), prefix, x)\n return x\n\n return jax.tree.map(maybe_split, tree, is_leaf=is_leaf)\n leaf_prefixes = broadcast_prefix(\n prefix,\n tree,\n prefix_is_leaf=lambda x: x is None or is_leaf(x),\n tree_is_leaf=is_leaf,\n )\n leaf_keys, treedef = jax.tree_util.tree_flatten_with_path(\n tree, is_leaf=is_leaf\n )\n assert len(leaf_keys) == len(leaf_prefixes)\n leaves_out = []\n\n with graphlib.merge_context(ctxtag, is_inner) as merge_ctx:\n for (keypath, leaf), leaf_prefix in zip(leaf_keys, leaf_prefixes):\n if (\n map_non_graph_nodes\n or is_node_leaf(leaf)\n or isinstance(leaf, variablelib.Variable)\n ):\n leaf = merge_fn(merge_ctx, keypath, leaf_prefix, leaf)\n leaves_out.append(leaf)\n\n pytree_out = jax.tree.unflatten(treedef, leaves_out)\n return pytree_out\n\ndef clear_non_graph_nodes(tree):\n return jax.tree.map(\n lambda x: x\n if graphlib.is_graph_node(x) or isinstance(x, variablelib.Variable)\n else None,\n tree,\n is_leaf=lambda x: isinstance(x, variablelib.Variable)\n or graphlib.is_graph_node(x),\n )\n\nclass Mask(tp.NamedTuple):\n pass\n\ndef mask_at(t: tuple, index: int | None) -> tuple:\n if index is None:\n return t\n return tuple(\n Mask() if i == index else x\n for i, x in enumerate(t)\n )\n\ndef replace_at(t: tuple, index: int, value: tp.Any) -> tuple:\n return tuple(\n value if i == index else x\n for i, x in enumerate(t)\n )\n\ndef updates_and_snapshot(args: A) -> tuple[A, A]:\n is_leaf = lambda x: isinstance(x, variablelib.Variable)\n leaves, treedef = jax.tree.flatten(args, is_leaf=is_leaf)\n updates_leaves: list[variablelib.Variable | Mask] = []\n snapshot_leaves: list[variablelib.Variable | Mask] = []\n for leaf in leaves:\n if isinstance(leaf, variablelib.Variable):\n updates_leaves.append(leaf)\n snapshot_leaves.append(leaf.copy())\n else:\n updates_leaves.append(Mask())\n snapshot_leaves.append(Mask())\n updates = jax.tree.unflatten(treedef, updates_leaves)\n snapshot = jax.tree.unflatten(treedef, snapshot_leaves)\n return updates, snapshot\n\n\ndef check_no_aliases(fn_name: str, /, **kwargs):\n Attrs = namedtuple('Attrs', kwargs.keys()) # type: ignore[misc]\n container = Attrs(**kwargs)\n is_leaf = lambda x: isinstance(x, variablelib.Variable)\n seen: dict[int, jax.tree_util.KeyPath] = {}\n for path, leaf in jax.tree.leaves_with_path(\n container, is_leaf=is_leaf\n ):\n if not isinstance(leaf, variablelib.Variable):\n continue\n var_id = id(leaf)\n if var_id in seen:\n path_str = jax.tree_util.keystr(path)\n seen_path_str = jax.tree_util.keystr(seen[var_id])\n raise ValueError(\n f'Duplicate {leaf}\\nfound at paths:\\n\\n'\n f' - {seen_path_str}\\n'\n f' - {path_str}\\n\\n'\n f'nnx.{fn_name} with graph_updates=False does not support '\n 'returning input Variables as outputs. '\n f'Consider the following options:\\n\\n'\n f'1. Remove the duplicate Variables.\\n'\n f'2. Create new Variables via nnx.clone() and use those instead.\\n'\n f'3. Enable graph mode and graph updates by passing graph=True and '\n f'graph_updates=True to {fn_name}\\n\\n'\n f' nnx.{fn_name}(..., graph=True, graph_updates=True)\\n\\n'\n f'4. Use nnx.compat.{fn_name} (sets graph and graph_updates to True '\n f'automatically)\\n\\n'\n f' nnx.compat.{fn_name}(...)'\n )\n seen[var_id] = path\n\n\ndef check_prefix(prefix: tp.Any, prefix_name: str, fn_name: str):\n def _check(path, leaf):\n if graphlib.is_graph_node(leaf) or isinstance(leaf, variablelib.Variable):\n raise ValueError(\n f'Found graph node or Variable of type {type(leaf).__name__} '\n f'at path {jax.tree_util.keystr(path)} in `{prefix_name}` '\n f'for nnx.{fn_name}. Graph nodes and Variables are not allowed '\n f'as prefixes when graph=True and graph_updates=False'\n )\n jax.tree.map_with_path(\n _check, prefix,\n is_leaf=lambda x: isinstance(x, variablelib.Variable)\n or graphlib.is_graph_node(x),\n )\n\n\ndef variable_changed(post: variablelib.Variable, pre: variablelib.Variable) -> bool:\n post_leaves, post_td = jax.tree.flatten(post)\n pre_leaves, pre_td = jax.tree.flatten(pre)\n return post_td != pre_td or any( # type: ignore[operator]\n a is not b for a, b in zip(post_leaves, pre_leaves)\n )\n\nKeepFn = tp.Callable[\n [PathParts, tp.Any, variablelib.Variable, variablelib.Variable], bool\n]\n\ndef mask_variable_updates(\n current_tree: A,\n snapshot_tree: A,\n *,\n prefix: tp.Any = Missing,\n keep_fn: KeepFn | None = None,\n) -> A:\n if keep_fn is None:\n keep_fn = lambda _, _pfx, cur, snap: variable_changed(cur, snap)\n\n def _mask_updates(path, prefix_leaf, current, snapshot):\n if isinstance(current, variablelib.Variable):\n if current.hijax or current.ref:\n return Mask()\n if keep_fn(path, prefix_leaf, current, snapshot):\n return current\n return Mask()\n is_leaf = lambda x: isinstance(x, variablelib.Variable) or x is None\n if prefix is Missing:\n return jax.tree.map_with_path(\n lambda path, cur, snap: _mask_updates(path, None, cur, snap),\n current_tree, snapshot_tree, is_leaf=is_leaf,\n )\n return broadcast_prefix_map(\n _mask_updates, prefix, current_tree, snapshot_tree, is_leaf=is_leaf,\n )\n\n\ndef apply_variable_updates(args_tree: A, updates_tree: A):\n is_leaf = lambda x: isinstance(x, variablelib.Variable) or isinstance(x, Mask)\n args_leaves = jax.tree.leaves(args_tree, is_leaf=is_leaf)\n _, treedef = jax.tree.flatten(args_tree, is_leaf=is_leaf)\n updates_leaves = treedef.flatten_up_to(updates_tree)\n for variable, update in zip(args_leaves, updates_leaves, strict=True):\n if isinstance(update, variablelib.Variable):\n assert isinstance(variable, variablelib.Variable)\n variable.update_from_state(update)\n\n\ndef treemap_copy_args(f):\n @functools.wraps(f)\n def wrapper(*args, **kwargs):\n args, kwargs = jax.tree.map(lambda x: x, (args, kwargs))\n return f(*args, **kwargs)\n return wrapper\n\n\ndef check_same_variables(inputs, outputs, transform_name: str = ''):\n def _check(in_leaf, out_leaf):\n if isinstance(in_leaf, variablelib.Variable) and in_leaf is not out_leaf:\n raise ValueError(\n f'{transform_name} Variable identity must be preserved '\n 'across iterations.'\n )\n is_leaf = lambda x: isinstance(x, (Mask, variablelib.Variable))\n jax.tree.map(\n _check, inputs, outputs,\n is_leaf=is_leaf,\n )\n\n\ndef update_carry_variables(init_val, val_out):\n def _update(in_leaf, out_leaf):\n if isinstance(in_leaf, variablelib.Variable):\n in_leaf.update_from_state(out_leaf)\n return in_leaf\n return out_leaf\n\n return jax.tree.map(\n _update, init_val, val_out,\n is_leaf=lambda x: isinstance(x, variablelib.Variable),\n )\n" }, { "path": "flax/nnx/filterlib.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport builtins\nimport dataclasses\nfrom flax.typing import Key, PathParts\nimport typing as tp\n\nif tp.TYPE_CHECKING:\n ellipsis = builtins.ellipsis\nelse:\n ellipsis = tp.Any\n\nPredicate = tp.Callable[[PathParts, tp.Any], bool]\n\nFilterLiteral = tp.Union[type, str, Predicate, bool, ellipsis, None]\nFilter = tp.Union[FilterLiteral, tuple['Filter', ...], list['Filter']]\n\n\n\ndef to_predicate(filter: Filter) -> Predicate:\n \"\"\"Converts a Filter to a predicate function.\n See `Using Filters `__.\n \"\"\"\n\n if isinstance(filter, str):\n return WithTag(filter)\n elif isinstance(filter, type):\n return OfType(filter)\n elif isinstance(filter, bool):\n if filter:\n return Everything()\n else:\n return Nothing()\n elif filter is Ellipsis:\n return Everything()\n elif filter is None:\n return Nothing()\n elif callable(filter):\n return filter\n elif isinstance(filter, (list, tuple)):\n return Any(*filter)\n else:\n raise TypeError(f'Invalid collection filter: {filter:!r}. ')\n\ndef filters_to_predicates(\n filters: tp.Sequence[Filter],\n) -> tuple[Predicate, ...]:\n for i, filter_ in enumerate(filters):\n if filter_ in (..., True) and i != len(filters) - 1:\n remaining_filters = filters[i + 1 :]\n if not all(f in (..., True) for f in remaining_filters):\n raise ValueError(\n '`...` or `True` can only be used as the last filters, '\n f'got {filter_} it at index {i}.'\n )\n return tuple(map(to_predicate, filters))\n\n\nclass HasTag(tp.Protocol):\n tag: str\n\n\ndef _has_tag(x: tp.Any) -> tp.TypeGuard[HasTag]:\n return hasattr(x, 'tag')\n\n\n@dataclasses.dataclass(frozen=True)\nclass WithTag:\n tag: str\n\n def __call__(self, path: PathParts, x: tp.Any):\n return _has_tag(x) and x.tag == self.tag\n\n def __repr__(self):\n return f'WithTag({self.tag!r})'\n\n\n@dataclasses.dataclass(frozen=True)\nclass PathContains:\n key: Key | str\n exact: bool = True\n\n def __call__(self, path: PathParts, x: tp.Any):\n if self.exact:\n return self.key in path\n return any(str(self.key) in str(part) for part in path)\n\n def __repr__(self):\n return f'PathContains({self.key!r}, exact={self.exact})'\n\n\nclass PathIn:\n def __init__(self, *paths: PathParts):\n self.paths = frozenset(paths)\n\n def __call__(self, path: PathParts, x: tp.Any):\n return path in self.paths\n\n def __repr__(self):\n paths_repr = ','.join(map(repr, self.paths))\n return f'PathIn({paths_repr})'\n\n def __eq__(self, other):\n return isinstance(other, PathIn) and self.paths == other.paths\n\n def __hash__(self):\n return hash(self.paths)\n\n\n@dataclasses.dataclass(frozen=True)\nclass OfType:\n type: type\n\n def __call__(self, path: PathParts, x: tp.Any):\n return isinstance(x, self.type)\n\n def __repr__(self):\n return f'OfType({self.type!r})'\n\n\nclass Any:\n def __init__(self, *filters: Filter):\n self.predicates = tuple(\n to_predicate(collection_filter) for collection_filter in filters\n )\n\n def __call__(self, path: PathParts, x: tp.Any):\n return any(predicate(path, x) for predicate in self.predicates)\n\n def __repr__(self):\n return f'Any({\", \".join(map(repr, self.predicates))})'\n\n def __eq__(self, other):\n return isinstance(other, Any) and self.predicates == other.predicates\n\n def __hash__(self):\n return hash(self.predicates)\n\n\nclass All:\n def __init__(self, *filters: Filter):\n self.predicates = tuple(\n to_predicate(collection_filter) for collection_filter in filters\n )\n\n def __call__(self, path: PathParts, x: tp.Any):\n return all(predicate(path, x) for predicate in self.predicates)\n\n def __repr__(self):\n return f'All({\", \".join(map(repr, self.predicates))})'\n\n def __eq__(self, other):\n return isinstance(other, All) and self.predicates == other.predicates\n\n def __hash__(self):\n return hash(self.predicates)\n\n\nclass Not:\n def __init__(self, collection_filter: Filter, /):\n self.predicate = to_predicate(collection_filter)\n\n def __call__(self, path: PathParts, x: tp.Any):\n return not self.predicate(path, x)\n\n def __repr__(self):\n return f'Not({self.predicate!r})'\n\n def __eq__(self, other):\n return isinstance(other, Not) and self.predicate == other.predicate\n\n def __hash__(self):\n return hash(self.predicate)\n\n\nclass Everything:\n def __call__(self, path: PathParts, x: tp.Any):\n return True\n\n def __repr__(self):\n return 'Everything()'\n\n def __eq__(self, other):\n return isinstance(other, Everything)\n\n def __hash__(self):\n return hash(Everything)\n\n\nclass Nothing:\n def __call__(self, path: PathParts, x: tp.Any):\n return False\n\n def __repr__(self):\n return 'Nothing()'\n\n def __eq__(self, other):\n return isinstance(other, Nothing)\n\n def __hash__(self):\n return hash(Nothing)\n" }, { "path": "flax/nnx/graph.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"Graph module.\n\nRe-exports APIs from ``flax.nnx.graphlib``. This module is kept for backward\ncompatibility with code that imports from ``flax.nnx.graph``.\n\"\"\"\n\nfrom flax.nnx.graphlib import *\n\n" }, { "path": "flax/nnx/graphlib.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom __future__ import annotations\n\nimport contextlib\nimport dataclasses\nimport functools\nimport threading\nimport typing as tp\nimport builtins\n\nimport jax.core\n\nfrom flax import config\nfrom flax.nnx import filterlib, reprlib, traversals, variablelib\nfrom flax.nnx import statelib\nfrom flax.nnx.proxy_caller import (\n ApplyCaller,\n CallableProxy,\n DelayedAccessor,\n)\nfrom flax.nnx.statelib import FlatState, State, map_state\nfrom flax.nnx.variablelib import Variable, is_array_ref, V\nfrom flax.typing import BaseConfigContext, HashableMapping, Key, PathParts, is_key_like\nimport jax\nimport numpy as np\nimport treescope # type: ignore[import-not-found,import-untyped]\nimport typing_extensions as tpe\n\nA = tp.TypeVar('A')\nB = tp.TypeVar('B')\nC = tp.TypeVar('C')\nF = tp.TypeVar('F', bound=tp.Callable)\n\nHA = tp.TypeVar('HA', bound=tp.Hashable)\nHB = tp.TypeVar('HB', bound=tp.Hashable)\nKeyT = tp.TypeVar('KeyT', bound=Key)\n\nIndex = int\n\ndef _tree_mode_suggestion_api(fn_name: str) -> str:\n return (\n f'Consider the following options:\\n\\n'\n '1. Remove the duplicates and guarantee a tree structure.\\n'\n f'2. Enable graph mode by passing graph=True to {fn_name} e.g.\\n\\n'\n f' nnx.{fn_name}(..., graph=True)\\n\\n'\n f'3. Use nnx.compat.{fn_name} instead e.g.\\n\\n'\n f' nnx.compat.{fn_name}(...)'\n )\n\ndef _tree_mode_suggestion_transform(fn_name: str) -> str:\n return (\n f'Consider the following options:\\n\\n'\n '1. Remove the duplicates.\\n'\n f'2. Enable graph mode and graph updates by passing graph=True and '\n f'graph_updates=True to {fn_name} e.g.\\n\\n'\n f' nnx.{fn_name}(..., graph=True, graph_updates=True)\\n\\n'\n f'3. Use nnx.compat.{fn_name} instead e.g.\\n\\n'\n f' nnx.compat.{fn_name}(...)'\n )\n\ndef _check_valid_pytree(\n node: tp.Any, fn_name: str, path: str = '',\n) -> None:\n from flax.nnx import pytreelib\n if (\n isinstance(node, pytreelib.Pytree)\n and not node._pytree__is_pytree\n ):\n msg = (\n f\"Cannot use '{fn_name}' with graph=False on a \"\n f\"'{type(node).__name__}' instance that has pytree=False. \"\n )\n if path:\n msg += f\"Found at path: {path}. \"\n msg += (\n f\"Pytree subclasses with pytree=False are not registered as \"\n f\"JAX pytrees and cannot be used in tree-mode. \"\n + _tree_mode_suggestion_api(fn_name)\n )\n raise ValueError(msg)\n\nNames = tp.Sequence[int]\nNode = tp.TypeVar('Node')\nLeaf = tp.TypeVar('Leaf')\nAuxData = tp.TypeVar('AuxData')\n\n\n@jax.tree_util.register_static\n@dataclasses.dataclass(frozen=True, slots=True)\nclass NoUpdate: ...\n\n\nNO_UPDATE = NoUpdate()\n\n\n@jax.tree_util.register_static\n@dataclasses.dataclass(frozen=True, slots=True)\nclass Repeated: ...\n\n\nREPEATED = Repeated()\n\n\n@jax.tree_util.register_dataclass\n@dataclasses.dataclass(frozen=True, slots=True, repr=False)\nclass ArrayRefOutput(reprlib.Representable):\n value: jax.Array\n\n def __nnx_repr__(self):\n yield reprlib.Object(type=type(self))\n yield reprlib.Attr('value', self.value)\n\n def __treescope_repr__(self, path, subtree_renderer):\n return treescope.repr_lib.render_object_constructor(\n object_type=type(self),\n attributes={\n 'value': self.value,\n },\n path=path,\n subtree_renderer=subtree_renderer,\n )\n\n\nLeafType = tp.Union[\n Variable,\n jax.Array,\n np.ndarray,\n variablelib.Ref,\n ArrayRefOutput,\n NoUpdate,\n]\nGraphState = State[Key, LeafType]\nGraphFlatState = FlatState[LeafType]\n\n\ndef is_node_leaf(x: tp.Any) -> tpe.TypeGuard[LeafType]:\n return isinstance(x, LeafType) or variablelib.is_array_ref(x) # type: ignore[misc, arg-type]\n\n\nclass IndexMap(dict[Index, tp.Any]):\n @staticmethod\n def from_refmap(refmap: RefMap) -> IndexMap:\n return IndexMap((index, value) for value, index in refmap.items())\n\n\nif config.flax_use_flaxlib:\n import flaxlib # type: ignore[import]\n\n globals()['IndexMap'] = flaxlib.IndexMap\n\n\n# RefMap = dict\nclass RefMap(tp.MutableMapping[tp.Any, int], reprlib.MappingReprMixin):\n \"\"\"A mapping that hashes keys by their identity.\"\"\"\n\n def __init__(\n self,\n mapping: tp.Mapping[tp.Any, int]\n | tp.Iterable[tuple[tp.Any, int]]\n | None = None,\n /,\n ):\n self._mapping: dict[int, tuple[tp.Any, int]] = dict()\n if mapping is not None:\n self.update(mapping)\n\n @staticmethod\n def from_indexmap(indexmap: IndexMap) -> RefMap:\n refmap = RefMap()\n refmap.update((value, index) for index, value in indexmap.items())\n return refmap\n\n def get(self, key: tp.Any, default: int | None = None) -> int | None: # type: ignore[override]\n return self._mapping.get(id(key), (None, default))[1]\n\n def __getitem__(self, key: tp.Any) -> int:\n return self._mapping[id(key)][1]\n\n def __setitem__(self, key: tp.Any, value: int):\n self._mapping[id(key)] = (key, value)\n\n def __delitem__(self, key: tp.Any):\n del self._mapping[id(key)]\n\n def __len__(self) -> int:\n return len(self._mapping)\n\n def __contains__(self, key: tp.Any) -> bool:\n return id(key) in self._mapping\n\n def __iter__(self) -> tp.Iterator[tp.Any]:\n for key, _ in self._mapping.values():\n yield key\n\n def items(self) -> tp.ItemsView[tp.Any, int]:\n return self._mapping.values() # type: ignore\n\n\n# save python version\nPythonRefMap = RefMap\n\nif config.flax_use_flaxlib:\n import flaxlib # type: ignore[import]\n\n globals()['RefMap'] = flaxlib.RefMap\n\n\n@dataclasses.dataclass(frozen=True, slots=True)\nclass NodeImplBase(tp.Generic[Node, Leaf, AuxData]):\n type: type[Node]\n flatten: tp.Callable[[Node], tuple[tp.Sequence[tuple[Key, Leaf]], AuxData]]\n\n def node_dict(self, node: Node) -> dict[Key, tp.Any]:\n nodes, _ = self.flatten(node)\n return {\n key: value.value if isinstance(value, DataElem | StaticElem) else value\n for key, value in nodes\n }\n\n\n@dataclasses.dataclass(frozen=True, slots=True)\nclass GraphNodeImpl(NodeImplBase[Node, Leaf, AuxData]):\n set_key: tp.Callable[[Node, Key, Leaf], None]\n pop_key: tp.Callable[[Node, Key], Leaf]\n create_empty: tp.Callable[[AuxData], Node]\n clear: tp.Callable[[Node], None]\n init: tp.Callable[[Node, tp.Iterable[tuple[Key, Leaf]]], None]\n\n\n@dataclasses.dataclass(frozen=True, slots=True)\nclass PytreeNodeImpl(NodeImplBase[Node, Leaf, AuxData]):\n unflatten: tp.Callable[[tp.Sequence[tuple[Key, Leaf]], AuxData], Node]\n set_key: tp.Callable[[Node, Key, Leaf], None] | None\n pop_key: tp.Callable[[Node, Key], Leaf] | None\n\n\nNodeImpl = tp.Union[\n GraphNodeImpl[Node, Leaf, AuxData], PytreeNodeImpl[Node, Leaf, AuxData]\n]\n\n\nGRAPH_REGISTRY: dict[type, NodeImpl[tp.Any, tp.Any, tp.Any]] = {}\nPYTREE_REGISTRY: dict[type, PytreeNodeImpl[tp.Any, tp.Any, tp.Any]] = {}\n\n\ndef register_graph_node_type(\n type: type,\n flatten: tp.Callable[[Node], tuple[tp.Sequence[tuple[Key, Leaf]], AuxData]],\n set_key: tp.Callable[[Node, Key, Leaf], None],\n pop_key: tp.Callable[[Node, Key], Leaf],\n create_empty: tp.Callable[[AuxData], Node],\n clear: tp.Callable[[Node], None],\n init: tp.Callable[[Node, tp.Iterable[tuple[Key, Leaf]]], None],\n):\n if type in GRAPH_REGISTRY:\n raise ValueError(f'Node type {type} is already registered.')\n\n GRAPH_REGISTRY[type] = GraphNodeImpl(\n type=type,\n flatten=flatten,\n set_key=set_key,\n pop_key=pop_key,\n create_empty=create_empty,\n clear=clear,\n init=init,\n )\n\n\ndef register_pytree_node_type(\n type: type,\n flatten: tp.Callable[[Node], tuple[tp.Sequence[tuple[Key, Leaf]], AuxData]],\n unflatten: tp.Callable[[tp.Sequence[tuple[Key, Leaf]], AuxData], Node],\n *,\n set_key: tp.Callable[[Node, Key, Leaf], None] | None = None,\n pop_key: tp.Callable[[Node, Key], Leaf] | None = None,\n):\n if type in PYTREE_REGISTRY:\n raise ValueError(f'Node type {type} is already registered.')\n\n PYTREE_REGISTRY[type] = PytreeNodeImpl(\n type=type,\n flatten=flatten,\n unflatten=unflatten,\n set_key=set_key,\n pop_key=pop_key,\n )\n\n\ndef is_node(x: tp.Any) -> bool:\n if isinstance(x, Variable):\n return False\n if type(x) in GRAPH_REGISTRY:\n return True\n return is_pytree_node(x)\n\n\ndef is_graph_node(x: tp.Any) -> bool:\n return (\n type(x) in GRAPH_REGISTRY\n or variablelib.is_array_ref(x)\n or isinstance(x, Variable)\n )\n\n\ndef is_node_type(x: type[tp.Any]) -> bool:\n return x in GRAPH_REGISTRY or x in PYTREE_REGISTRY or x is GenericPytree\n\n\ndef get_node_impl(x: Node) -> NodeImpl[Node, tp.Any, tp.Any] | None:\n if isinstance(x, Variable):\n return None\n\n node_type = type(x)\n\n if node_type in GRAPH_REGISTRY:\n return GRAPH_REGISTRY[node_type]\n elif node_type in PYTREE_REGISTRY:\n return PYTREE_REGISTRY[node_type]\n elif node_type in JAX_PYTREE_REGISTRY or issubclass(node_type, tuple):\n return PYTREE_NODE_IMPL # type: ignore\n else:\n return None\n\n\ndef get_node_impl_for_type(\n x: type[Node],\n) -> NodeImpl[Node, tp.Any, tp.Any] | None:\n if x is GenericPytree:\n return PYTREE_NODE_IMPL # type: ignore\n elif x in PYTREE_REGISTRY:\n return PYTREE_REGISTRY[x]\n elif x in GRAPH_REGISTRY:\n return GRAPH_REGISTRY[x]\n else:\n return None\n\n# use type-aware sorting to support int keys\ndef _type_aware_sort(item: tuple[tp.Any, tp.Any]) -> tuple[int, tp.Any]:\n key, _ = item\n if isinstance(key, int):\n return (0, key)\n elif isinstance(key, str):\n return (1, key)\n else:\n raise ValueError(f'Unsupported key type: {type(key)!r}')\n\n\n@jax.tree_util.register_static\n@dataclasses.dataclass(frozen=True, repr=False)\nclass NodeRef(tp.Generic[Node], reprlib.Representable):\n index: int\n\n def __nnx_repr__(self):\n yield reprlib.Object(type=type(self))\n yield reprlib.Attr('index', self.index)\n\n def __treescope_repr__(self, path, subtree_renderer):\n return treescope.repr_lib.render_object_constructor(\n object_type=type(self),\n attributes={'index': self.index},\n path=path,\n subtree_renderer=subtree_renderer,\n )\n\n\nif config.flax_use_flaxlib:\n import flaxlib # type: ignore[import]\n\n jax.tree_util.register_static(flaxlib.NodeRef)\n globals()['NodeRef'] = flaxlib.NodeRef\n\n\n@dataclasses.dataclass(frozen=True, repr=False)\nclass VariableDef(reprlib.Representable, tp.Generic[Node]):\n type: type[Node]\n index: int\n outer_index: int | None\n metadata: HashableMapping[str, tp.Any]\n array_refdef: ArrayRefDef | NodeRef | None\n\n def with_no_outer_index(self) -> VariableDef:\n return VariableDef(\n type=self.type,\n index=self.index,\n outer_index=None,\n metadata=self.metadata,\n array_refdef=self.array_refdef.with_no_outer_index()\n if isinstance(self.array_refdef, ArrayRefDef)\n else self.array_refdef,\n )\n\n def with_same_outer_index(self) -> VariableDef:\n return VariableDef(\n type=self.type,\n index=self.index,\n outer_index=self.index,\n metadata=self.metadata,\n array_refdef=self.array_refdef.with_same_outer_index()\n if isinstance(self.array_refdef, ArrayRefDef)\n else self.array_refdef,\n )\n\n def with_matching_outer_index(self, other) -> VariableDef:\n return VariableDef(\n type=self.type,\n index=self.index,\n outer_index=other.outer_index,\n metadata=self.metadata,\n array_refdef=self.array_refdef.with_matching_outer_index(other.array_refdef)\n if isinstance(self.array_refdef, ArrayRefDef)\n else self.array_refdef\n )\n\n def __nnx_repr__(self):\n yield reprlib.Object(type=type(self))\n yield reprlib.Attr('type', self.type.__name__)\n yield reprlib.Attr('index', self.index)\n yield reprlib.Attr('outer_index', self.outer_index)\n yield reprlib.Attr('metadata', reprlib.PrettyMapping(self.metadata))\n\n def __treescope_repr__(self, path, subtree_renderer):\n return treescope.repr_lib.render_object_constructor(\n object_type=type(self),\n attributes={\n 'type': self.type,\n 'index': self.index,\n 'outer_index': self.outer_index,\n 'metadata': self.metadata,\n },\n path=path,\n subtree_renderer=subtree_renderer,\n )\n\n\nif config.flax_use_flaxlib:\n import flaxlib # type: ignore[import]\n\n jax.tree_util.register_static(flaxlib.VariableDef)\n globals()['VariableDef'] = flaxlib.VariableDef\n\n\n@dataclasses.dataclass(frozen=True, repr=False)\nclass ArrayRefDef(reprlib.Representable):\n index: int\n outer_index: int | None\n\n def with_no_outer_index(self):\n return ArrayRefDef(\n index=self.index,\n outer_index=None,\n )\n\n def with_same_outer_index(self):\n return ArrayRefDef(\n index=self.index,\n outer_index=self.index,\n )\n\n def with_matching_outer_index(self, other):\n return ArrayRefDef(\n index=self.index,\n outer_index=other.outer_index,\n )\n\n def __nnx_repr__(self):\n yield reprlib.Object(type=type(self))\n yield reprlib.Attr('index', self.index)\n yield reprlib.Attr('outer_index', self.outer_index)\n\n def __treescope_repr__(self, path, subtree_renderer):\n return treescope.repr_lib.render_object_constructor(\n object_type=type(self),\n attributes={\n 'index': self.index,\n 'outer_index': self.outer_index,\n },\n path=path,\n subtree_renderer=subtree_renderer,\n )\n\n\n@jax.tree_util.register_static\n@dataclasses.dataclass(frozen=True, repr=False, slots=True)\nclass NodeDef(tp.Generic[Node], reprlib.Representable):\n \"\"\"A dataclass that denotes the tree structure of a\n :class:`Module`. A ``GraphDef`` can be generated by either\n calling :func:`split` or :func:`graphdef` on the :class:`Module`.\"\"\"\n\n type: tp.Type[Node]\n index: int | None\n outer_index: int | None\n num_attributes: int\n metadata: tp.Any\n\n def with_no_outer_index(self) -> NodeDef[Node]:\n return NodeDef(\n type=self.type,\n index=self.index,\n outer_index=None,\n num_attributes=self.num_attributes,\n metadata=self.metadata,\n )\n\n def with_same_outer_index(self) -> NodeDef[Node]:\n return NodeDef(\n type=self.type,\n index=self.index,\n outer_index=self.index,\n num_attributes=self.num_attributes,\n metadata=self.metadata,\n )\n\n def with_matching_outer_index(self, other) -> NodeDef[Node]:\n return NodeDef(\n type=self.type,\n index=self.index,\n outer_index=other.outer_index,\n num_attributes=self.num_attributes,\n metadata=self.metadata,\n )\n\n def __nnx_repr__(self):\n yield reprlib.Object(type=type(self))\n\n yield reprlib.Attr('type', self.type.__name__)\n yield reprlib.Attr('index', self.index)\n yield reprlib.Attr('outer_index', self.outer_index)\n yield reprlib.Attr('num_attributes', self.num_attributes)\n yield reprlib.Attr('metadata', self.metadata)\n\n def __treescope_repr__(self, path, subtree_renderer):\n return treescope.repr_lib.render_object_constructor(\n object_type=type(self),\n attributes={\n 'type': self.type,\n 'index': self.index,\n 'outer_index': self.outer_index,\n 'num_attributes': self.num_attributes,\n 'metadata': self.metadata,\n },\n path=path,\n subtree_renderer=subtree_renderer,\n )\n\n\nif config.flax_use_flaxlib:\n import flaxlib # type: ignore[import]\n\n jax.tree_util.register_static(flaxlib.NodeDef)\n globals()['NodeDef'] = flaxlib.NodeDef\n\n@jax.tree_util.register_static\n@dataclasses.dataclass(frozen=True, slots=True)\nclass TreeNodeDef(tp.Generic[Node]):\n type: tp.Type[Node]\n treedef: jax.tree_util.PyTreeDef\n path_index: tuple[tuple[PathParts, int], ...]\n\n def with_no_outer_index(self) -> TreeNodeDef[Node]:\n return self\n\n def with_same_outer_index(self) -> TreeNodeDef[Node]:\n return self\n\n def with_matching_outer_index(self, other) -> TreeNodeDef[Node]:\n return self\n\nNodeDefType = tp.Union[\n NodeDef[Node],\n NodeRef[Node],\n VariableDef[Node],\n ArrayRefDef,\n TreeNodeDef[Node],\n]\n\n\n@dataclasses.dataclass(frozen=True, slots=True)\nclass NodeAttr:\n pass\n\n\nNODE_ATTR = NodeAttr()\n\n@dataclasses.dataclass(frozen=True, slots=True)\nclass LeafAttr:\n pass\n\nLEAF_ATTR = LeafAttr()\n\nAttrType = tp.Union[\n NodeAttr,\n LeafAttr,\n 'Static[tp.Any]',\n]\n\n\n# GraphDef = tp.Union[NodeDef[Node], NodeRef[Node], VariableDef[Node]]\n@jax.tree_util.register_static\n@dataclasses.dataclass(frozen=True, slots=True)\nclass GraphDef(tp.Generic[Node]):\n nodes: list[NodeDefType[tp.Any]]\n attributes: list[tuple[Key, AttrType]]\n num_leaves: int\n\n def __hash__(self) -> int:\n return hash((tuple(self.nodes), tuple(self.attributes)))\n\n def with_no_outer_index(self) -> GraphDef[Node]:\n return GraphDef(\n nodes=[\n node.with_no_outer_index() if not isinstance(node, NodeRef) else node\n for node in self.nodes\n ],\n attributes=self.attributes,\n num_leaves=self.num_leaves,\n )\n\n def with_matching_outer_index(self, other) -> GraphDef[Node]:\n return GraphDef(\n nodes=[\n node.with_matching_outer_index(other_node) if not isinstance(node, NodeRef) else node\n for node, other_node in zip(self.nodes, other.nodes)\n ],\n attributes=self.attributes,\n num_leaves=self.num_leaves,\n )\n\n def with_same_outer_index(self) -> GraphDef[Node]:\n return GraphDef(\n nodes=[\n node.with_same_outer_index() if not isinstance(node, NodeRef) else node\n for node in self.nodes\n ],\n attributes=self.attributes,\n num_leaves=self.num_leaves,\n )\n\n # TODO(cgarciae): remove this method\n def apply(\n self, state: GraphState, *states: GraphState,\n graph: bool | None = None,\n ) -> ApplyCaller[tuple[GraphDef[Node], GraphState]]:\n accessor = DelayedAccessor()\n\n def _apply(\n accessor: DelayedAccessor, *args, **kwargs\n ) -> tuple[tp.Any, tuple[GraphDef[Node], GraphState]]:\n module = merge(self, state, *states)\n fn = accessor(module)\n out = fn(*args, **kwargs)\n if graph is None:\n _graph = config.nnx_graph_mode\n else:\n _graph = graph\n graphdef, flat_state = flatten(module, graph=_graph)\n state_ = statelib.from_flat_state(flat_state)\n return out, (graphdef, state_)\n\n return CallableProxy(_apply, accessor) # type: ignore\n\n\nPureState = tuple[GraphDef[Node], GraphState]\n\n\ndef _tree_flatten(\n node: tp.Any,\n nodes: list[NodeDefType[tp.Any]],\n leaves: list[tp.Any],\n paths: list[PathParts] | None,\n) -> None:\n seen_variables: dict[int, str] = {}\n seen_refs: dict[int, str] = {}\n def _is_leaf(path, x):\n if isinstance(x, Variable):\n var_id = id(x)\n str_path = jax.tree_util.keystr(path)\n if var_id in seen_variables:\n raise ValueError(\n f'Duplicate {x}\\nfound at paths:\\n\\n'\n f' - {seen_variables[var_id]}\\n'\n f' - {str_path}\\n\\n'\n 'Tree mode (graph=False) does not support shared references. '\n + _tree_mode_suggestion_api('split')\n )\n seen_variables[var_id] = str_path\n return True\n if variablelib.is_array_ref(x):\n ref_id = id(x)\n str_path = jax.tree_util.keystr(path)\n if ref_id in seen_refs:\n raise ValueError(\n f'Duplicate {x}\\nfound at paths:\\n\\n'\n f' - {seen_refs[ref_id]}\\n'\n f' - {str_path}\\n\\n'\n 'Tree mode (graph=False) does not support shared references. '\n + _tree_mode_suggestion_api('split')\n )\n seen_refs[ref_id] = str_path\n _check_valid_pytree(x, 'flatten', jax.tree_util.keystr(path))\n return False\n jax_leaves, treedef = jax.tree_util.tree_flatten_with_path(\n node, is_leaf=_is_leaf, is_leaf_takes_path=True\n )\n nnx_paths_and_leaves: list[tuple[PathParts, tp.Any]] = [\n (jax_to_nnx_path(jax_path), value) for jax_path, value in jax_leaves\n ]\n original_indices = {p: i for i, (p, _) in enumerate(nnx_paths_and_leaves)}\n nnx_paths_and_leaves.sort()\n path_index = tuple(\n (p, original_indices[p]) for p, _ in nnx_paths_and_leaves\n )\n\n tree_nodedef: TreeNodeDef[tp.Any] = TreeNodeDef(\n type=type(node),\n treedef=treedef,\n path_index=path_index,\n )\n nodes.append(tree_nodedef)\n\n sorted_leaf_index = 0\n for nnx_path, value in nnx_paths_and_leaves:\n if isinstance(value, Variable):\n nodes.append(VariableDef(\n type=value.var_type,\n index=sorted_leaf_index,\n outer_index=None,\n metadata=HashableMapping(value.get_metadata()),\n array_refdef=None,\n ))\n leaves.append(value)\n if paths is not None:\n paths.append(nnx_path)\n sorted_leaf_index += 1\n\n\n@tp.overload\ndef flatten( # type: ignore[invalid-annotation]\n node: Node,\n /,\n *,\n ref_index: RefMap | None = ...,\n ref_outer_index: RefMap | None = ...,\n graph: bool = ...,\n) -> tuple[GraphDef[Node], FlatState[tp.Any]]: ...\n@tp.overload\ndef flatten( # type: ignore[invalid-annotation]\n node: Node,\n /,\n *,\n with_paths: tp.Literal[True],\n ref_index: RefMap | None = ...,\n ref_outer_index: RefMap | None = ...,\n graph: bool = ...,\n) -> tuple[\n GraphDef[Node],\n FlatState[tp.Any],\n]: ...\n@tp.overload\ndef flatten( # type: ignore[invalid-annotation]\n node: Node,\n /,\n *,\n with_paths: tp.Literal[False],\n ref_index: RefMap | None = ...,\n ref_outer_index: RefMap | None = ...,\n graph: bool = ...,\n) -> tuple[\n GraphDef[Node],\n list[tp.Any],\n]: ...\n@tp.overload\ndef flatten( # type: ignore[invalid-annotation]\n node: Node,\n /,\n *,\n with_paths: bool,\n ref_index: RefMap | None = ...,\n ref_outer_index: RefMap | None = ...,\n graph: bool = ...,\n) -> tuple[\n GraphDef[Node],\n FlatState[tp.Any] | list[tp.Any],\n]: ...\ndef flatten( # type: ignore[invalid-annotation]\n node: Node,\n /,\n *,\n with_paths: bool = True,\n ref_index: RefMap | None = None,\n ref_outer_index: RefMap | None = None,\n graph: bool | None = None,\n) -> tuple[\n GraphDef[Node],\n FlatState[tp.Any] | list[tp.Any],\n]:\n \"\"\"Flattens a graph node into a (graphdef, state) pair.\n\n Args:\n x: A graph node.\n ref_index: A mapping from nodes to indexes, defaults to None. If not provided, a new\n empty dictionary is created. This argument can be used to flatten a sequence of graph\n nodes that share references.\n with_paths: A boolean that indicates whether to return a FlatState object that includes\n the paths, or just a list of the Variable's inner values.\n graph: If ``True`` (default), uses graph-mode which supports the full\n NNX feature set including shared references. If ``False``, uses\n tree-mode which treats Modules as regular JAX pytrees, avoiding\n the overhead of the graph protocol.\n \"\"\"\n if graph is None:\n graph = set_graph_mode.current_value()\n if ref_index is None:\n ref_index = RefMap()\n leaves: list[tp.Any] = []\n path: list[Key] | None = [] if with_paths else None\n paths: list[PathParts] | None = [] if with_paths else None\n nodes: list[NodeDefType[tp.Any]] = []\n attributes: list[tuple[Key, AttrType]] = []\n if graph:\n node_impl = get_node_impl(node)\n _graph_flatten(\n node,\n node_impl,\n path,\n ref_index,\n ref_outer_index,\n nodes,\n attributes,\n leaves,\n paths,\n )\n else:\n _tree_flatten(\n node,\n nodes,\n leaves,\n paths,\n )\n graphdef: GraphDef = GraphDef(\n nodes=nodes, attributes=attributes, num_leaves=len(leaves)\n )\n\n if paths is not None:\n return graphdef, FlatState.from_sorted_keys_values(tuple(paths), leaves) # type: ignore[return-value]\n else:\n return graphdef, leaves\n\n@dataclasses.dataclass(frozen=True, slots=True)\nclass DataElem:\n value: tp.Any\n\n\n@dataclasses.dataclass(frozen=True, slots=True)\nclass StaticElem:\n value: tp.Any\n\ndef _graph_flatten(\n node: Node,\n node_impl: NodeImpl[Node, Leaf, AuxData] | None,\n path: list[Key] | None,\n ref_index: RefMap,\n ref_outer_index: RefMap | None,\n nodes: list[NodeDefType[tp.Any]],\n attributes: list[tuple[Key, AttrType]],\n leaves: list[tp.Any],\n paths: list[PathParts] | None,\n) -> None:\n is_pytree_node_ = type(node_impl) is PytreeNodeImpl\n\n index: int | None\n if not is_pytree_node_ and node in ref_index:\n nodes.append(NodeRef(index := ref_index[node]))\n return\n\n is_graph_node_ = type(node_impl) is GraphNodeImpl\n is_variable = isinstance(node, Variable)\n is_array_ref = variablelib.is_array_ref(node)\n\n # only cache graph nodes, we don't add array refs here\n # as they are added in the make_mutable_arraydef function\n if is_graph_node_ or is_variable:\n index = len(ref_index)\n ref_index[node] = index\n else:\n index = None\n\n def make_mutable_arraydef(value: variablelib.Ref):\n if value in ref_index:\n index = ref_index[value]\n return NodeRef(index), REPEATED\n else:\n index = len(ref_index)\n ref_index[value] = index\n output_value: NoUpdate | ArrayRefOutput | variablelib.Ref\n if ref_outer_index is not None:\n if value in ref_outer_index:\n outer_index = ref_outer_index[value]\n output_value = NO_UPDATE\n array_refdef = ArrayRefDef(index=index, outer_index=outer_index)\n else:\n output_value = ArrayRefOutput(value[...])\n array_refdef = ArrayRefDef(index=index, outer_index=None)\n else:\n output_value = value\n array_refdef = ArrayRefDef(index=index, outer_index=None)\n return array_refdef, output_value\n\n if is_variable:\n assert isinstance(node, Variable)\n assert index is not None\n prev_inner_value = node.get_raw_value()\n if variablelib.is_array_ref(prev_inner_value):\n array_refdef, inner_value = make_mutable_arraydef(prev_inner_value)\n else:\n array_refdef = None\n inner_value = prev_inner_value\n if path is None:\n leaf = inner_value\n else:\n leaf = node # type: ignore[assignment]\n if inner_value is not prev_inner_value:\n leaf.set_raw_value(inner_value)\n\n variabledef = VariableDef(\n type=node.var_type, # type: ignore\n index=index,\n outer_index=ref_outer_index.get(node, None) if ref_outer_index else None,\n metadata=HashableMapping(node.get_metadata()),\n array_refdef=array_refdef,\n )\n if type(inner_value) is not Repeated:\n assert not isinstance(leaf, Repeated)\n leaves.append(leaf)\n if path is not None:\n assert paths is not None\n paths.append(tuple(path))\n nodes.append(variabledef)\n return\n elif is_array_ref:\n array_refdef, leaf = make_mutable_arraydef(node) # type: ignore[arg-type]\n if not isinstance(leaf, Repeated):\n leaves.append(leaf)\n if path is not None:\n assert paths is not None\n paths.append(tuple(path))\n nodes.append(array_refdef)\n return\n elif not is_pytree_node_ and not is_graph_node_:\n # unkown leaf\n leaves.append(node)\n if path is not None:\n assert paths is not None\n paths.append(tuple(path))\n return\n\n if node_impl is None:\n raise RuntimeError(f'Unsupported type: {type(node)}, this is a bug.')\n\n values, metadata = node_impl.flatten(node)\n num_attributes = len(values)\n nodedef = NodeDef(\n node_impl.type,\n index,\n ref_outer_index[node]\n if is_graph_node_ and ref_outer_index and node in ref_outer_index\n else None,\n num_attributes,\n metadata,\n )\n nodes.append(nodedef)\n\n for key, value in values:\n is_data = None\n if isinstance(value, DataElem):\n value = value.value\n is_data = True\n elif isinstance(value, StaticElem):\n value = value.value\n is_data = False\n\n if is_data is False:\n attributes.append((key, Static(value)))\n continue\n\n value_node_impl = get_node_impl(value)\n if path is not None:\n path.append(key)\n if value_node_impl is not None or isinstance(value, Variable):\n attributes.append((key, NODE_ATTR))\n _graph_flatten(\n value,\n value_node_impl,\n path,\n ref_index,\n ref_outer_index,\n nodes,\n attributes,\n leaves,\n paths,\n )\n elif variablelib.is_array_ref(value):\n attributes.append((key, NODE_ATTR))\n array_refdef, leaf = make_mutable_arraydef(value)\n if not isinstance(leaf, Repeated):\n leaves.append(leaf)\n if paths is not None:\n paths.append(tuple(path)) # type: ignore\n nodes.append(array_refdef)\n elif isinstance(value, (jax.Array, np.ndarray)) or is_data:\n attributes.append((key, LEAF_ATTR))\n if paths is not None:\n paths.append(tuple(path)) # type: ignore\n leaves.append(value)\n else:\n attributes.append((key, Static(value)))\n\n if path is not None:\n path.pop()\n\n return\n\n\ndef _get_sorted_leaves(\n xs: tp.Mapping[tp.Any, tp.Any],\n) -> list[tp.Any]:\n if not isinstance(xs, tp.Mapping): # type: ignore\n raise TypeError(f'expected Mapping; got {type(xs).__qualname__}')\n leaves: list[tp.Any] = []\n\n def _flatten(xs):\n if not isinstance(xs, tp.Mapping):\n leaves.append(xs)\n else:\n for _, value in sorted(xs.items()):\n _flatten(value)\n\n _flatten(xs)\n return leaves\n\n\ndef _tree_unflatten(\n graphdef: GraphDef[tp.Any],\n leaves: list[tp.Any],\n copy_variables: bool,\n) -> tp.Any:\n tree_nodedef = graphdef.nodes[0]\n assert isinstance(tree_nodedef, TreeNodeDef)\n variable_defs_iter = iter(\n node for node in graphdef.nodes[1:] if isinstance(node, VariableDef)\n )\n variabledef = next(variable_defs_iter, None)\n\n original_leaves: list[tp.Any] = [None] * len(leaves)\n for i, (path, original_index) in enumerate(tree_nodedef.path_index):\n leaf = leaves[i]\n if variabledef is not None and variabledef.index == i:\n if isinstance(leaf, Variable):\n if copy_variables:\n leaf = leaf.copy()\n else:\n leaf = variabledef.type.from_metadata(\n leaf, dict(variabledef.metadata)\n )\n variabledef = next(variable_defs_iter, None)\n original_leaves[original_index] = leaf\n\n return tree_nodedef.treedef.unflatten(original_leaves)\n\n\ndef unflatten( # type: ignore[invalid-annotation]\n graphdef: GraphDef[Node],\n state: State[Key, tp.Any] | FlatState[tp.Any] | list[tp.Any],\n /,\n *,\n index_ref: IndexMap | None = None,\n outer_index_outer_ref: IndexMap | None = None,\n copy_variables: bool = False,\n) -> Node:\n \"\"\"Unflattens a graphdef into a node with the given state.\n\n Args:\n graphdef: A GraphDef instance.\n state: A State instance.\n index_ref: A mapping from indexes to nodes references found during the graph\n traversal, defaults to None. If not provided, a new empty dictionary is\n created. This argument can be used to unflatten a sequence of (graphdef,\n state) pairs that share the same index space.\n index_ref_cache: A mapping from indexes to existing nodes that can be\n reused. When an reference is reused, ``GraphNodeImpl.clear`` is called to\n leave the object in an empty state and then filled by the unflatten\n process, as a result existing graph nodes are mutated to have the new\n content/topology specified by the graphdef.\n copy_variables: If True variables in the state will be copied onto the new\n new structure, else variables will be shared. Default is False.\n \"\"\"\n if isinstance(state, (State, dict)):\n leaves = _get_sorted_leaves(state)\n elif isinstance(state, FlatState):\n leaves = state.leaves\n elif isinstance(state, list): # type: ignore\n leaves = state\n else:\n raise ValueError(f'Unsupported state type: {type(state)}')\n\n if len(leaves) != graphdef.num_leaves:\n raise ValueError(\n f'Incorrect number of leaves, expected {graphdef.num_leaves} leaves, but got {len(leaves)}.'\n )\n\n if graphdef.nodes and isinstance(graphdef.nodes[0], TreeNodeDef):\n return _tree_unflatten(graphdef, leaves, copy_variables)\n\n if index_ref is None:\n index_ref = IndexMap()\n\n if len(graphdef.nodes) == 0:\n return leaves[0]\n elif isinstance(nodedef := graphdef.nodes[0], NodeRef):\n node = index_ref[nodedef.index]\n else:\n node_iter = iter(graphdef.nodes)\n attribute_iter = iter(graphdef.attributes)\n leaves_iter = iter(leaves)\n nodedef = next(node_iter)\n assert not isinstance(nodedef, NodeRef)\n if isinstance(nodedef, ArrayRefDef):\n node_impl = None\n else:\n node_impl = get_node_impl_for_type(nodedef.type)\n node = _graph_unflatten(\n nodedef,\n node_impl,\n node_iter,\n attribute_iter,\n leaves_iter,\n index_ref,\n outer_index_outer_ref,\n copy_variables,\n )\n\n try:\n next(leaves_iter)\n except StopIteration:\n pass\n else:\n raise ValueError('Incorrect number of leaves in state.')\n\n return node\n\n\ndef _graph_unflatten(\n nodedef: NodeDefType[Node],\n node_impl: NodeImpl[Node, Leaf, AuxData] | None,\n node_iter: tp.Iterator[NodeDefType[Node]],\n attribute_iter: tp.Iterator[tuple[Key, AttrType]],\n leaves_iter: tp.Iterator[tp.Any],\n index_ref: IndexMap,\n outer_index_outer_ref: IndexMap | None,\n copy_variables: bool,\n) -> Node:\n \"\"\"Recursive helper for graph_unflatten.\n\n Args:\n nodedef: A GraphDef instance or an index to a node in the cache.\n state: A mapping from attribute names to variables or subgraphs.\n index_ref: A mapping from indexes to nodes that have been traversed.\n If a node is already in the cache, it won't be traversed again.\n outer_index_outer_ref: A mapping from indexes to existing nodes that can be reused.\n When an reference is reused, ``GraphNodeImpl.clear`` is called to leave the\n object in an empty state and then filled by the unflatten process, as a result\n existing graph nodes are mutated to have the new content/topology\n specified by the nodedef.\n \"\"\"\n\n def get_mutable_array(array_refdef: ArrayRefDef, leaf):\n assert type(array_refdef) is ArrayRefDef\n if (\n outer_index_outer_ref is not None\n and array_refdef.outer_index is not None\n and array_refdef.outer_index in outer_index_outer_ref\n ):\n # if array ref exists, update it\n array_ref = outer_index_outer_ref[array_refdef.outer_index]\n if not variablelib.is_array_ref(array_ref):\n raise RuntimeError(f'Expected a ArrayRef type but got {array_ref}.')\n if type(leaf) is not NoUpdate:\n raise RuntimeError(f'Expected a no update for ArrayRef but got {leaf}.')\n elif type(leaf) in (NoUpdate, Repeated):\n raise ValueError(\n f\"Expected a ArrayRefOutput type but got '{leaf}.'\"\n )\n elif type(leaf) is ArrayRefOutput:\n array_ref = jax.new_ref(leaf.value)\n elif variablelib.is_array_ref(leaf):\n array_ref = leaf\n else:\n # here we allow merging frozen arrays and will not create a new array ref\n array_ref = leaf\n\n index_ref[array_refdef.index] = array_ref\n return array_ref\n\n if type(nodedef) is NodeRef:\n return index_ref[nodedef.index]\n\n if type(nodedef) is VariableDef:\n variabledef = tp.cast(VariableDef[Variable], nodedef)\n # its a unseen variable, create a new one\n\n if variabledef.array_refdef is not None:\n if type(variabledef.array_refdef) is NodeRef:\n value = index_ref[variabledef.array_refdef.index]\n else:\n value = next(leaves_iter)\n assert type(variabledef.array_refdef) is ArrayRefDef\n if isinstance(value, Variable):\n copy_ref = not isinstance(\n value.get_raw_value(), (NoUpdate, Repeated, ArrayRefOutput)\n )\n value = value.copy(_copy_ref=copy_ref) if copy_variables else value\n inner_value = value.get_raw_value()\n array_ref = get_mutable_array(variabledef.array_refdef, inner_value)\n if array_ref is not inner_value:\n value.set_raw_value(array_ref)\n else:\n # if value is an array or array ref, we need call get_mutable_array\n # to register it in the index_ref\n value = get_mutable_array(variabledef.array_refdef, value)\n else:\n value = next(leaves_iter)\n if isinstance(value, Variable) and copy_variables:\n copy_ref = not isinstance(\n value.get_raw_value(), (NoUpdate, Repeated, ArrayRefOutput)\n )\n value = value.copy(_copy_ref=copy_ref)\n\n # when idxmap is present, check if the Varable exists there\n # and update existing variables if it does\n if (\n outer_index_outer_ref is not None\n and variabledef.outer_index is not None\n and variabledef.outer_index in outer_index_outer_ref\n ):\n # if variable exists, update it\n variable = outer_index_outer_ref[variabledef.outer_index]\n if not isinstance(variable, Variable):\n raise ValueError(f'Expected a Variable type but got {type(variable)}.')\n elif isinstance(value, Variable):\n variable.update_from_state(value)\n else:\n variable.set_raw_value(value)\n else: # variabledef.index not in index_ref_cache\n # variable reference does not exist outside, create a new one\n if isinstance(value, Variable):\n variable = value\n else:\n variable = variabledef.type.from_metadata(\n value, dict(variabledef.metadata)\n )\n index_ref[variabledef.index] = variable\n return variable # type: ignore[return-value]\n\n if type(nodedef) is ArrayRefDef:\n leaf = next(leaves_iter)\n array_ref = get_mutable_array(nodedef, leaf)\n return array_ref # type: ignore[return-value]\n\n assert type(nodedef) is NodeDef\n if node_impl is None:\n raise RuntimeError(f'Unsupported type: {nodedef.type}, this is a bug.')\n if nodedef.index is not None and nodedef.index in index_ref:\n raise RuntimeError(f'GraphDef index {nodedef.index} already used.')\n\n def _get_children() -> list[tuple[Key, tp.Any]]:\n children: list[tuple[Key, LeafType | Node]] = [] # type: ignore[invalid-annotation]\n\n assert type(nodedef) is NodeDef\n for _ in range(nodedef.num_attributes):\n key, value = next(attribute_iter)\n if type(value) is Static:\n children.append((key, value.value)) # type: ignore[attribute-error]\n elif type(value) is LeafAttr:\n leaf = next(leaves_iter)\n children.append((key, leaf))\n elif type(value) is NodeAttr:\n node_def = next(node_iter)\n if isinstance(node_def, NodeRef):\n node = index_ref[node_def.index]\n elif isinstance(node_def, ArrayRefDef):\n leaf = next(leaves_iter)\n node = get_mutable_array(node_def, leaf)\n elif isinstance(node_def, NodeDef | VariableDef):\n value_node_impl = get_node_impl_for_type(node_def.type)\n node = _graph_unflatten(\n node_def,\n value_node_impl,\n node_iter,\n attribute_iter,\n leaves_iter,\n index_ref,\n outer_index_outer_ref,\n copy_variables,\n )\n else:\n raise RuntimeError(f'Unknown node definition: {node_def!r}')\n children.append((key, node))\n elif type(value) is NodeRef:\n children.append((key, index_ref[value.index])) # type: ignore[attribute-error]\n else:\n raise RuntimeError(f'Unknown static field: {key!r}')\n\n return children\n\n if isinstance(node_impl, GraphNodeImpl):\n # we create an empty node first and add it to the index\n # this avoids infinite recursion when there is a reference cycle\n assert type(nodedef) is NodeDef\n if (\n outer_index_outer_ref is not None\n and nodedef.outer_index is not None\n and nodedef.outer_index in outer_index_outer_ref\n ):\n node = outer_index_outer_ref[nodedef.outer_index]\n if type(node) != nodedef.type:\n raise ValueError(\n f'Expected a node of type {nodedef.type} for index '\n f'{nodedef.index}, but got a node of type {type(node)}.'\n )\n node_impl.clear(node)\n else:\n node = node_impl.create_empty(nodedef.metadata)\n assert nodedef.index is not None\n index_ref[nodedef.index] = node\n node_impl.init(node, _get_children())\n else:\n # if the node type does not support the creation of an empty object it means\n # that it cannot reference itself, so we can create its children first\n node = node_impl.unflatten(_get_children(), nodedef.metadata)\n\n return node\n\n\ndef graph_pop(\n node: tp.Any,\n filters: tuple[filterlib.Filter, ...],\n) -> tuple[GraphState, ...]:\n id_to_index: dict[int, Index] = {}\n path_parts: PathParts = ()\n predicates = tuple(filterlib.to_predicate(filter) for filter in filters)\n flat_states: tuple[dict[PathParts, LeafType], ...] = tuple(\n {} for _ in predicates\n )\n _graph_pop(node, id_to_index, path_parts, flat_states, predicates)\n return tuple(\n statelib.from_flat_state(flat_state) for flat_state in flat_states\n )\n\n\ndef _graph_pop(\n node: tp.Any,\n id_to_index: dict[int, Index],\n path_parts: PathParts,\n flat_states: tuple[dict[PathParts, LeafType], ...],\n predicates: tuple[filterlib.Predicate, ...],\n) -> None:\n if not is_node(node):\n raise RuntimeError(f'Unsupported type: {type(node)}, this is a bug.')\n\n if id(node) in id_to_index:\n return\n\n id_to_index[id(node)] = len(id_to_index)\n node_impl = get_node_impl(node)\n if node_impl is None:\n raise TypeError(f'Unknown node type: {type(node)}')\n node_dict = node_impl.node_dict(node)\n\n for name, value in node_dict.items():\n if is_node(value):\n _graph_pop(\n node=value,\n id_to_index=id_to_index,\n path_parts=(*path_parts, name),\n flat_states=flat_states,\n predicates=predicates,\n )\n continue\n elif not is_node_leaf(value):\n continue\n elif id(value) in id_to_index:\n continue\n\n node_path = (*path_parts, name)\n node_impl = get_node_impl(node)\n if node_impl is None:\n raise TypeError(f'Unknown node type: {type(node)}')\n\n for state, predicate in zip(flat_states, predicates):\n if predicate(node_path, value):\n if node_impl.pop_key is None:\n raise ValueError(\n f'Cannot pop key {name!r} from node of type {type(node).__name__}'\n )\n id_to_index[id(value)] = len(id_to_index)\n node_impl.pop_key(node, name)\n if isinstance(value, Variable):\n value = value\n state[node_path] = value # type: ignore[index] # mypy is wrong here?\n break\n else:\n # NOTE: should we raise an error here?\n pass\n\n\ndef _graph_update_dynamic(node: tp.Any, state: tp.Mapping[KeyT, tp.Any]):\n def _update_variable(node: Variable, value):\n if isinstance(value, Variable):\n # updated from Variable\n node.update_from_state(value)\n else:\n # updated from raw value\n if isinstance(value, State) and not value:\n # NOTE: this is a special case when trying to update a Variable from state\n # created when flattening into a NodeRef, which creates an empty State. This\n # can happen when using standalone Variables with `grad`\n pass\n else:\n if is_array_ref(node.get_raw_value()) and (\n isinstance(value, jax.Array) or is_array_ref(value)\n ):\n node[...] = value[...]\n else:\n node.set_raw_value(value, _unsafe_bypass_check=True)\n\n if isinstance(node, Variable):\n _update_variable(node, state)\n return\n\n if not is_node(node):\n raise RuntimeError(f'Unsupported type: {type(node)}')\n\n node_impl = get_node_impl(node)\n if node_impl is None:\n raise TypeError(f'Unknown node type: {type(node)}')\n node_dict = node_impl.node_dict(node)\n for key, value in state.items():\n # case 1: new state is being added\n if key not in node_dict:\n if node_impl.set_key is None:\n raise ValueError(\n f'Cannot set key {key!r} on immutable node of '\n f'type {type(node).__name__}'\n )\n if isinstance(value, Variable):\n copy_ref = not isinstance(\n value.get_raw_value(), (NoUpdate, Repeated, ArrayRefOutput)\n )\n value = value.copy(_copy_ref=copy_ref)\n node_impl.set_key(node, key, value)\n continue\n\n current_value = node_dict[key]\n\n # case 2: subgraph is being updated\n if is_array_ref(current_value):\n current_value[...] = value\n elif is_node(current_value):\n if is_node_leaf(value):\n raise ValueError(f'Expected a subgraph for {key!r}, but got: {value!r}')\n _graph_update_dynamic(current_value, value)\n elif isinstance(current_value, Variable):\n _update_variable(current_value, value)\n elif node_impl.set_key is not None:\n node_impl.set_key(node, key, value)\n else:\n raise ValueError(\n f'Cannot set key {key!r} on immutable node of '\n f'type {type(node).__name__}'\n )\n\n\n# --------------------------------------------------------\n# UpdateContext\n# --------------------------------------------------------\n\n\nclass StaticCache(tp.NamedTuple):\n graphdef: GraphDef[tp.Any]\n final_graphdef: GraphDef[tp.Any]\n paths: tuple[PathParts, ...]\n variables: list[Variable[tp.Any]]\n new_ref_index: RefMap\n new_index_ref: IndexMap\n\n @staticmethod\n def create(\n graphdef: GraphDef[tp.Any],\n paths: tuple[PathParts, ...],\n variables: list[Variable[tp.Any]],\n new_ref_index: RefMap,\n ):\n new_index_ref = IndexMap.from_refmap(new_ref_index)\n final_graphdef: GraphDef[tp.Any]\n final_graphdef = graphdef.with_same_outer_index()\n return StaticCache(\n graphdef=graphdef,\n final_graphdef=final_graphdef,\n paths=paths,\n variables=variables,\n new_ref_index=new_ref_index,\n new_index_ref=new_index_ref,\n )\n\n\n@dataclasses.dataclass\nclass GraphContext(threading.local):\n update_context_stacks: dict[tp.Hashable, list[UpdateContext]] = (\n dataclasses.field(default_factory=dict)\n )\n ref_index_stack: list[SplitContext] = dataclasses.field(default_factory=list)\n index_ref_stack: list[MergeContext] = dataclasses.field(default_factory=list)\n tmp_static_cache: tp.MutableMapping[tp.Any, StaticCache] | None = None\n caching: bool = False\n graph_mode_stack: list[bool] = dataclasses.field(default_factory=list)\n graph_updates_stack: list[bool] = dataclasses.field(default_factory=list)\n\n\nGRAPH_CONTEXT = GraphContext()\n\n\nclass set_graph_mode(BaseConfigContext):\n get_default = classmethod(lambda cls: config.nnx_graph_mode)\n get_stack = classmethod(lambda cls: GRAPH_CONTEXT.graph_mode_stack)\n\n\nclass set_graph_updates(BaseConfigContext):\n get_default = classmethod(lambda cls: config.nnx_graph_updates)\n get_stack = classmethod(lambda cls: GRAPH_CONTEXT.graph_updates_stack)\n\n\n@contextlib.contextmanager\ndef static_cache(static_cache: tp.MutableMapping[tp.Any, StaticCache]):\n if GRAPH_CONTEXT.caching:\n yield\n return\n\n GRAPH_CONTEXT.tmp_static_cache = static_cache\n\n try:\n yield\n finally:\n if GRAPH_CONTEXT.tmp_static_cache is not None:\n raise ValueError(\n 'GRAPH_CONTEXT.tmp_static_cache should be None, no context consumed it.'\n )\n\n\ndef _cached_partial(f: tp.Callable[..., tp.Any], *cached_args, graph: bool | None = None):\n \"\"\"Create a partial from a NNX transformed function alog with some cached input arguments\n and reduces the python overhead by caching the traversal of NNX graph nodes. This is useful\n for speed up function that are called repeatedly with the same subset of inputs e.g. a\n ``train_step`` with a ``model`` and ``optimizer``::\n\n >>> from flax import nnx\n >>> import jax.numpy as jnp\n >>> import optax\n ...\n >>> model = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n >>> optimizer = nnx.Optimizer(model, optax.adamw(1e-3), wrt=nnx.Param)\n ...\n >>> @nnx.jit\n ... def train_step(model, optimizer, x, y):\n ... def loss_fn(model):\n ... return jnp.mean((model(x) - y) ** 2)\n ...\n ... loss, grads = nnx.value_and_grad(loss_fn)(model)\n ... optimizer.update(model, grads)\n ... return loss\n ...\n >>> cached_train_step = nnx.cached_partial(train_step, model, optimizer)\n ...\n >>> for step in range(total_steps:=2):\n ... x, y = jnp.ones((10, 2)), jnp.ones((10, 3))\n ... # loss = train_step(model, optimizer, x, y)\n ... loss = cached_train_step(x, y)\n ... print(f'Step {step}: loss={loss:.3f}')\n Step 0: loss=2.669\n Step 1: loss=2.660\n\n Note that ``cached_partial`` will clone all cached graph nodes to gurantee the validity\n of the cache, and these clones will contain references to the same Variable objects\n which guarantees that state is propagated correctly back to the original graph nodes.\n Because of the previous, the final structure of all graph nodes must be the same\n after each call to the cached function, otherwise an error will be raised. Temporary\n mutations are allowed (e.g. the use of ``Module.sow``) as long as they are cleaned up before\n the function returns (e.g. via ``nnx.pop``).\n\n Args:\n f: A function to cache.\n *cached_args: A subset of the input arguments containing the graph nodes to cache.\n\n Returns:\n A partial function expecting the remaining arguments to the original function.\n \"\"\"\n if graph is None:\n graph = set_graph_mode.current_value()\n if not graph:\n raise ValueError(\n 'cached_partial is a graph-mode-only API and does not support '\n 'tree-mode (graph=False).'\n )\n cache: tp.MutableMapping[tp.Any, StaticCache] = PythonRefMap() # type: ignore\n original_ref_index: RefMap = RefMap()\n index_ref: IndexMap = IndexMap()\n cached_ref_index: RefMap = RefMap()\n\n def create_static_cache(x):\n # TODO(cgarciae): support Array attribute updates for graph nodes\n if is_graph_node(x) or isinstance(x, Variable):\n graphdef, flat_state = flatten(\n x, with_paths=True, ref_index=original_ref_index, graph=True\n )\n paths = flat_state.paths\n variables = flat_state.leaves\n # clone but keep the same variable references\n node_cache = unflatten(\n graphdef, flat_state, index_ref=index_ref, copy_variables=False,\n )\n start_index = len(cached_ref_index)\n flatten(\n node_cache,\n ref_index=cached_ref_index,\n with_paths=False,\n graph=True,\n )\n cached_new_ref_index = RefMap(\n (key, value)\n for key, value in cached_ref_index.items()\n if value >= start_index\n )\n cache[node_cache] = StaticCache.create(\n graphdef, paths, variables, cached_new_ref_index\n )\n return node_cache\n return x\n\n cached_args = jax.tree.map(\n create_static_cache,\n cached_args,\n is_leaf=lambda x: is_graph_node(x) or isinstance(x, Variable),\n )\n\n @functools.wraps(f)\n def cache_args_wrapper(*args, **kwargs):\n with static_cache(cache):\n return f(*cached_args, *args, **kwargs)\n\n return cache_args_wrapper\n\n\nif tp.TYPE_CHECKING:\n cached_partial = functools.partial\nelse:\n cached_partial = _cached_partial\n\n\n@dataclasses.dataclass\nclass SplitContext:\n ctxtag: tp.Hashable | None\n ref_index: RefMap\n is_inner: bool | None\n\n @tp.overload\n def split(self, graph_node: A, /) -> tuple[GraphDef[A], GraphState]: ... # type: ignore[invalid-annotation]\n\n @tp.overload\n def split( # type: ignore[invalid-annotation]\n self, graph_node: A, first: filterlib.Filter, /\n ) -> tuple[GraphDef[A], GraphState]: ...\n\n @tp.overload\n def split(\n self,\n graph_node: A,\n first: filterlib.Filter,\n second: filterlib.Filter,\n /,\n *filters: filterlib.Filter,\n ) -> tuple[GraphDef[A], GraphState, tpe.Unpack[tuple[GraphState, ...]]]: ... # type: ignore[not-supported-yet]\n\n def split(\n self, node: A, *filters: filterlib.Filter\n ) -> tuple[GraphDef[A], tpe.Unpack[tuple[GraphState, ...]]]: # type: ignore[not-supported-yet]\n ctx = (\n current_update_context(self.ctxtag) if self.ctxtag is not None else None\n )\n inner_ref_outer_index = (\n ctx.inner_ref_outer_index if ctx and ctx.inner_ref_outer_index else None\n )\n graphdef, flat_state = flatten(\n node, ref_index=self.ref_index, ref_outer_index=inner_ref_outer_index, graph=True\n )\n flat_states = _split_state(flat_state, filters)\n states = _to_nested_state(graphdef, flat_states)\n\n return graphdef, *states\n\n @tp.overload\n def flatten( # type: ignore[invalid-annotation]\n self,\n graph_node: A,\n /,\n *,\n with_paths: tp.Literal[False],\n ) -> tuple[GraphDef[A], list[tp.Any]]: ...\n\n @tp.overload\n def flatten( # type: ignore[invalid-annotation]\n self,\n graph_node: A,\n /,\n ) -> tuple[GraphDef[A], FlatState[tp.Any]]: ...\n\n @tp.overload\n def flatten( # type: ignore[invalid-annotation]\n self,\n graph_node: A,\n first: filterlib.Filter,\n /,\n ) -> tuple[GraphDef[A], FlatState[tp.Any]]: ...\n\n @tp.overload\n def flatten( # type: ignore[invalid-annotation]\n self,\n graph_node: A,\n first: filterlib.Filter,\n second: filterlib.Filter,\n /,\n *filters: filterlib.Filter,\n ) -> tuple[\n GraphDef[A],\n FlatState[tp.Any],\n tpe.Unpack[tuple[FlatState[tp.Any], ...]],\n ]: ...\n\n def flatten( # type: ignore[invalid-annotation]\n self,\n node: A,\n *filters: filterlib.Filter,\n with_paths: bool = True,\n ) -> tuple[\n GraphDef[A],\n FlatState[tp.Any] | list[tp.Any],\n tpe.Unpack[tuple[FlatState[tp.Any], ...]],\n ]:\n if not with_paths and filters:\n raise ValueError('Cannot use filters with with_paths=False')\n\n ctx = (\n current_update_context(self.ctxtag) if self.ctxtag is not None else None\n )\n static_cache = (\n ctx.static_cache if ctx is not None and self.is_inner is False else None\n )\n ref_outer_index = (\n ctx.inner_ref_outer_index if ctx and ctx.inner_ref_outer_index else None\n )\n flat_state: FlatState[tp.Any] | list[tp.Any]\n leaves: list[tp.Any]\n if node in self.ref_index:\n # node is already in the ref_index, call flatten which will return a NodeRef\n graphdef, flat_state = flatten(\n node,\n ref_index=self.ref_index,\n ref_outer_index=ref_outer_index,\n with_paths=with_paths,\n graph=True,\n )\n if with_paths:\n assert isinstance(flat_state, FlatState)\n paths = flat_state.paths\n leaves = flat_state.leaves\n else:\n assert isinstance(flat_state, list)\n paths = None\n leaves = flat_state\n elif static_cache is not None and node in static_cache:\n node_static_cache = static_cache[node]\n graphdef = node_static_cache.graphdef\n # add the new references to the ref_index\n self.ref_index.update(node_static_cache.new_ref_index)\n\n if with_paths:\n paths = node_static_cache.paths\n leaves = node_static_cache.variables\n else:\n paths = None\n leaves = [\n variable.get_raw_value() for variable in node_static_cache.variables\n ]\n else:\n graphdef, flat_state = flatten(\n node,\n ref_index=self.ref_index,\n ref_outer_index=ref_outer_index,\n with_paths=with_paths,\n graph=True,\n )\n if with_paths:\n assert isinstance(flat_state, FlatState)\n paths = flat_state.paths\n leaves = flat_state.leaves\n else:\n assert isinstance(flat_state, list)\n paths = None\n leaves = flat_state\n\n if with_paths:\n assert paths is not None\n flat_state = FlatState.from_sorted_keys_values(paths, leaves)\n flat_states = _split_state(flat_state, filters)\n return graphdef, *flat_states # type: ignore[bad-return-type]\n else:\n return graphdef, leaves\n\n\n@contextlib.contextmanager\ndef split_context(ctxtag: tp.Hashable | None = None):\n ctx = current_update_context(ctxtag) if ctxtag is not None else None\n is_inner = ctx.outer_ref_outer_index is not None if ctx is not None else None\n GRAPH_CONTEXT.ref_index_stack.append(SplitContext(ctxtag, RefMap(), is_inner))\n\n try:\n yield GRAPH_CONTEXT.ref_index_stack[-1]\n finally:\n flatten_ctx = GRAPH_CONTEXT.ref_index_stack.pop()\n if ctxtag is not None:\n ctx = current_update_context(ctxtag)\n ctx.flatten_end(flatten_ctx.ref_index)\n del flatten_ctx.ref_index\n del flatten_ctx.ctxtag\n\n\n@dataclasses.dataclass\nclass MergeContext:\n ctxtag: tp.Hashable | None\n index_ref: IndexMap\n is_inner: bool | None\n\n def merge( # type: ignore[invalid-annotation]\n self,\n graphdef: GraphDef[A],\n state: GraphState,\n /,\n *states: GraphState,\n ) -> A:\n ctx = (\n current_update_context(self.ctxtag) if self.ctxtag is not None else None\n )\n outer_index_outer_ref = (\n ctx.outer_index_outer_ref if ctx and ctx.outer_index_outer_ref else None\n )\n\n _state = _merge_to_flat_state((state, *states))\n node = unflatten(\n graphdef,\n _state,\n index_ref=self.index_ref,\n outer_index_outer_ref=outer_index_outer_ref,\n copy_variables=True,\n )\n return node\n\n def unflatten( # type: ignore[invalid-annotation]\n self,\n graphdef: GraphDef[A],\n flat_state: GraphFlatState | list[tp.Any],\n /,\n *flat_states: GraphFlatState,\n ) -> A:\n ctx = (\n current_update_context(self.ctxtag) if self.ctxtag is not None else None\n )\n static_cache = (\n ctx.static_cache if ctx is not None and self.is_inner is False else None\n )\n state: FlatState[tp.Any] | list[tp.Any]\n if type(flat_state) is list:\n if flat_states:\n raise ValueError(\n 'Cannot use multiple flat_states when flat_state is a list, '\n f'got flat_state: {flat_state!r}, flat_states: {flat_states!r}'\n )\n state = flat_state\n else:\n state = FlatState.merge(flat_state, *flat_states)\n\n if type(graphdef.nodes[0]) is NodeRef:\n node = unflatten(\n graphdef,\n state,\n index_ref=self.index_ref,\n )\n\n elif static_cache is not None:\n assert isinstance(graphdef.nodes[0], NodeDef) or isinstance(graphdef.nodes[0], VariableDef)\n assert ctx is not None\n if (outer_index := graphdef.nodes[0].outer_index) is not None:\n outer_index_outer_ref = ctx.outer_index_outer_ref\n assert outer_index_outer_ref is not None\n node = outer_index_outer_ref[outer_index]\n\n if node in static_cache:\n static_cache_node = static_cache[node]\n if static_cache_node.final_graphdef != graphdef:\n raise ValueError(\n 'The graph structure of a node added to cached_partial was mutated inside the transformation, '\n f'this is not allowed.\\nNode: {node}\\nOuput graphdef: {graphdef}\\nExpected graphdef: {static_cache_node.final_graphdef}'\n )\n if type(state) is list:\n leaves = state\n elif type(state) is FlatState:\n leaves = state.leaves\n else:\n raise ValueError(f'Unsupported state type: {type(state)}')\n\n if len(leaves) != len(static_cache_node.variables):\n raise ValueError(\n f'Incorrect number of leaves: expected {len(static_cache_node.variables)} '\n f'leaves in the state, got {len(leaves)}'\n )\n for variable, leaf in zip(static_cache_node.variables, leaves):\n if isinstance(leaf, Variable):\n variable.update_from_state(leaf)\n else:\n variable.set_raw_value(leaf)\n self.index_ref.update(static_cache_node.new_index_ref)\n else:\n # uncached node, create it\n node = unflatten(\n graphdef,\n state,\n index_ref=self.index_ref,\n outer_index_outer_ref=outer_index_outer_ref,\n )\n else: # graphdef.outer_index is None\n # its a new node, create it\n node = unflatten(\n graphdef,\n state,\n index_ref=self.index_ref,\n )\n else:\n outer_index_outer_ref = (\n ctx.outer_index_outer_ref if ctx and ctx.outer_index_outer_ref else None\n )\n node = unflatten(\n graphdef,\n state,\n index_ref=self.index_ref,\n outer_index_outer_ref=outer_index_outer_ref,\n )\n return node\n\n\n@tp.overload\n@contextlib.contextmanager\ndef merge_context() -> tp.Generator[MergeContext, None, None]: ... # type: ignore[bad-return-type]\n@tp.overload\n@contextlib.contextmanager\ndef merge_context(\n ctxtag: tp.Hashable | None, inner: bool | None\n) -> tp.Generator[MergeContext, None, None]: ... # type: ignore[bad-return-type]\n@contextlib.contextmanager\ndef merge_context(ctxtag: tp.Hashable | None = None, inner: bool | None = None):\n GRAPH_CONTEXT.index_ref_stack.append(MergeContext(ctxtag, IndexMap(), inner))\n\n try:\n yield GRAPH_CONTEXT.index_ref_stack[-1]\n finally:\n unflatten_ctx = GRAPH_CONTEXT.index_ref_stack.pop()\n index_ref = unflatten_ctx.index_ref\n if ctxtag is not None:\n if inner is None:\n raise ValueError('inner_merge must be specified when using ctxtag')\n ctx = current_update_context(ctxtag)\n ctx.unflatten_end(index_ref, inner)\n del unflatten_ctx.index_ref\n del unflatten_ctx.ctxtag\n\n\n@jax.tree_util.register_static\n@dataclasses.dataclass\nclass UpdateContext:\n \"\"\"A context manager for handling complex state updates.\"\"\"\n\n tag: tp.Hashable\n outer_ref_outer_index: RefMap | None\n outer_index_inner_ref: IndexMap | None\n # reverse caches\n outer_index_outer_ref: IndexMap | None\n inner_ref_outer_index: RefMap | None\n static_cache: tp.MutableMapping[tp.Any, StaticCache] | None\n\n # define hash and eq to make this an opaque object\n def __hash__(self):\n return 0\n\n def __eq__(self, other):\n return isinstance(other, UpdateContext)\n\n def flatten_end(self, ref_index: RefMap):\n if self.outer_ref_outer_index is None:\n # outer split (1), store the references\n self.outer_ref_outer_index = ref_index\n self.outer_index_outer_ref = IndexMap.from_refmap(\n self.outer_ref_outer_index\n )\n else:\n # inner split (3), clear index_ref\n self.outer_index_inner_ref = None\n self.inner_ref_outer_index = None\n\n def unflatten_end(self, index_ref: IndexMap, inner_merge: bool):\n if inner_merge:\n # inner merge (2)\n self.outer_index_inner_ref = index_ref\n self.inner_ref_outer_index = RefMap.from_indexmap(index_ref)\n\n\n@dataclasses.dataclass\nclass UpdateContextManager:\n tag: tp.Hashable\n\n def __enter__(self):\n if GRAPH_CONTEXT.tmp_static_cache is not None:\n # take current static cache\n static_cache = GRAPH_CONTEXT.tmp_static_cache\n GRAPH_CONTEXT.tmp_static_cache = None\n else:\n static_cache = None\n ctx = UpdateContext(\n tag=self.tag,\n outer_ref_outer_index=None,\n outer_index_inner_ref=None,\n outer_index_outer_ref=None,\n inner_ref_outer_index=None,\n static_cache=static_cache,\n )\n if self.tag not in GRAPH_CONTEXT.update_context_stacks:\n GRAPH_CONTEXT.update_context_stacks[self.tag] = [ctx]\n else:\n GRAPH_CONTEXT.update_context_stacks[self.tag].append(ctx)\n return ctx\n\n def __exit__(self, *args):\n if self.tag not in GRAPH_CONTEXT.update_context_stacks:\n raise RuntimeError(\n f'No update context found for tag {self.tag!r}, this is a bug.'\n )\n stack = GRAPH_CONTEXT.update_context_stacks[self.tag]\n\n ctx = stack.pop()\n # clear references\n del ctx.outer_ref_outer_index\n del ctx.outer_index_inner_ref\n del ctx.outer_index_outer_ref\n del ctx.inner_ref_outer_index\n\n if not stack:\n del GRAPH_CONTEXT.update_context_stacks[self.tag]\n\n def __call__(self, f: F) -> F:\n @functools.wraps(f)\n def update_context_manager_wrapper(*args, **kwargs):\n with self:\n return f(*args, **kwargs)\n\n return update_context_manager_wrapper # type: ignore\n\n\ndef update_context(tag: tp.Hashable):\n \"\"\"Creates an :class:`UpdateContext` context manager which can be used to handle\n more complex state updates beyond what ``nnx.update`` can handle, including\n updates to static properties and graph structure.\n\n UpdateContext exposes a ``split`` and ``merge`` API with the same\n signature as ``nnx.split`` / ``nnx.merge`` but performs some bookkeeping\n to have the necessary information in order to perfectly update the input\n objects based on the changes made inside the transform. The UpdateContext\n must call split and merge a total of 4 times, the first\n and last calls happen outside the transform and the second and third calls\n happen inside the transform as shown in the diagram below::\n\n\n idxmap\n (2) merge ─────────────────────────────► split (3)\n ▲ │\n │ inside │\n │. . . . . . . . . . . . . . . . . . │ index_mapping\n │ outside │\n │ ▼\n (1) split──────────────────────────────► merge (4)\n refmap\n\n\n The first call to split ``(1)`` creates a ``refmap`` which keeps track of the\n outer references, and the first call to merge ``(2)`` creates an ``idxmap`` which\n keeps track of the inner references. The second call to split ``(3)`` combines\n the refmap and idxmap to produce the ``index_mapping`` which indicates\n how the outer references map to the inner references. Finally, the last call to\n merge ``(4)`` uses the index_mapping and the refmap to reconstruct the\n output of the transform while reusing/updating the inner references. To avoid\n memory leaks, the idxmap is cleared after ``(3)`` and the refmap is\n cleared after ``(4)``, and both are cleared after the context manager exits.\n\n Here is a simple example showing the use of ``update_context``::\n\n >>> from flax import nnx\n ...\n >>> class Foo(nnx.Module): pass\n ...\n >>> m1 = Foo()\n >>> with nnx.update_context('example'):\n ... with nnx.split_context('example') as ctx:\n ... graphdef, state = ctx.split(m1)\n ... @jax.jit\n ... def f(graphdef, state):\n ... with nnx.merge_context('example', inner=True) as ctx:\n ... m2 = ctx.merge(graphdef, state)\n ... m2.a = 1\n ... m2.ref = m2 # create a reference cycle\n ... with nnx.split_context('example') as ctx:\n ... return ctx.split(m2)\n ... graphdef_out, state_out = f(graphdef, state)\n ... with nnx.merge_context('example', inner=False) as ctx:\n ... m3 = ctx.merge(graphdef_out, state_out)\n ...\n >>> assert m1 is m3\n >>> assert m1.a == 1\n >>> assert m1.ref is m1\n\n Note that ``update_context`` takes in a ``tag`` argument which is used\n primarily as a safety mechanism reduce the risk of accidentally using the\n wrong UpdateContext when using :func:`current_update_context` to access the\n current active context. ``update_context`` can also be used as a\n decorator that creates/activates an UpdateContext context for the\n duration of the function::\n\n >>> from flax import nnx\n ...\n >>> class Foo(nnx.Module): pass\n ...\n >>> m1 = Foo()\n >>> @jax.jit\n ... def f(graphdef, state):\n ... with nnx.merge_context('example', inner=True) as ctx:\n ... m2 = ctx.merge(graphdef, state)\n ... m2.a = 1 # insert static attribute\n ... m2.ref = m2 # create a reference cycle\n ... with nnx.split_context('example') as ctx:\n ... return ctx.split(m2)\n ...\n >>> @nnx.update_context('example')\n ... def g(m1):\n ... with nnx.split_context('example') as ctx:\n ... graphdef, state = ctx.split(m1)\n ... graphdef_out, state_out = f(graphdef, state)\n ... with nnx.merge_context('example', inner=False) as ctx:\n ... return ctx.merge(graphdef_out, state_out)\n ...\n >>> m3 = g(m1)\n >>> assert m1 is m3\n >>> assert m1.a == 1\n >>> assert m1.ref is m1\n\n The context can be accessed using :func:`current_update_context`.\n\n Args:\n tag: A string tag to identify the context.\n \"\"\"\n return UpdateContextManager(tag=tag)\n\n\ndef current_update_context(tag: tp.Hashable) -> UpdateContext:\n \"\"\"Returns the current active :class:`UpdateContext` for the given tag.\"\"\"\n if tag not in GRAPH_CONTEXT.update_context_stacks:\n raise ValueError(f'No update context found for tag {tag!r}.')\n return GRAPH_CONTEXT.update_context_stacks[tag][-1]\n\n\n# --------------------------------------------------------\n# Functional API\n# --------------------------------------------------------\n\n\ndef _split_state(\n state: FlatState[tp.Any],\n filters: tuple[filterlib.Filter, ...],\n) -> tuple[FlatState[tp.Any], tpe.Unpack[tuple[FlatState[tp.Any], ...]]]:\n if not filters:\n return (state,) # type: ignore[bad-return-type]\n states = state.split(*filters)\n if not isinstance(states, tuple):\n return (states,) # type: ignore[bad-return-type]\n assert len(states) > 0\n return states # type: ignore[return-value]\n\n\n@tp.overload\ndef split( # type: ignore[invalid-annotation]\n graph_node: A, /, *, graph: bool | None = None,\n) -> tuple[GraphDef[A], GraphState]: ...\n@tp.overload\ndef split( # type: ignore[invalid-annotation]\n graph_node: A, first: filterlib.Filter, /, *, graph: bool | None = None,\n) -> tuple[GraphDef[A], GraphState]: ...\n@tp.overload\ndef split( # type: ignore[invalid-annotation]\n graph_node: A,\n first: filterlib.Filter,\n second: filterlib.Filter,\n /,\n *filters: filterlib.Filter,\n graph: bool | None = None,\n) -> tuple[\n GraphDef[A],\n GraphState,\n tpe.Unpack[tuple[GraphState, ...]],\n]: ...\ndef split( # type: ignore[invalid-annotation]\n node: A, *filters: filterlib.Filter, graph: bool | None = None,\n) -> tuple[\n GraphDef[A],\n GraphState,\n tpe.Unpack[tuple[GraphState, ...]],\n]:\n \"\"\"Split a graph node into a :class:`GraphDef` and one or more :class:`State`s. State is\n a ``Mapping`` from strings or integers to ``Variables``, Arrays or nested States. GraphDef\n contains all the static information needed to reconstruct a ``Module`` graph, it is analogous\n to JAX's ``PyTreeDef``. :func:`split` is used in conjunction with :func:`merge` to switch\n seamlessly between stateful and stateless representations of the graph.\n\n Example usage::\n\n >>> from flax import nnx\n >>> import jax, jax.numpy as jnp\n ...\n >>> class Foo(nnx.Module):\n ... def __init__(self, rngs):\n ... self.batch_norm = nnx.BatchNorm(2, rngs=rngs)\n ... self.linear = nnx.Linear(2, 3, rngs=rngs)\n ...\n >>> node = Foo(nnx.Rngs(0))\n >>> graphdef, params, batch_stats = nnx.split(node, nnx.Param, nnx.BatchStat)\n ...\n >>> jax.tree.map(jnp.shape, params)\n State({\n 'batch_norm': {\n 'bias': Param(\n value=(2,)\n ),\n 'scale': Param(\n value=(2,)\n )\n },\n 'linear': {\n 'bias': Param(\n value=(3,)\n ),\n 'kernel': Param(\n value=(2, 3)\n )\n }\n })\n >>> jax.tree.map(jnp.shape, batch_stats)\n State({\n 'batch_norm': {\n 'mean': BatchStat(\n value=(2,)\n ),\n 'var': BatchStat(\n value=(2,)\n )\n }\n })\n\n :func:`split` and :func:`merge` are primarily used to interact directly with JAX\n transformations, see\n `Functional API `__\n for more information.\n\n Arguments:\n node: graph node to split.\n *filters: some optional filters to group the state into mutually exclusive substates.\n graph: If ``True`` (default), uses graph-mode which supports the full\n NNX feature set including shared references. If ``False``, uses\n tree-mode which treats Modules as regular JAX pytrees, avoiding\n the overhead of the graph protocol.\n Returns:\n ``GraphDef`` and one or more ``States`` equal to the number of filters passed. If no\n filters are passed, a single ``State`` is returned.\n \"\"\"\n if graph is None:\n graph = set_graph_mode.current_value()\n graphdef, flat_state = flatten(node, graph=graph)\n flat_states = _split_state(flat_state, filters)\n states = _to_nested_state(graphdef, flat_states)\n return graphdef, *states # type: ignore[return-value]\n\n\ndef _to_nested_state(\n graphdef: GraphDef[A], flat_states: tp.Iterable[tp.Any]\n) -> tuple[tp.Any, ...]:\n def _nested_or_leaf(flat_state):\n if not flat_state:\n return State({})\n if len(flat_state) == 1 and flat_state[0][0] == ():\n return flat_state[0][1]\n return statelib.from_flat_state(flat_state)\n states = tuple(\n _nested_or_leaf(flat_state) for flat_state in flat_states\n )\n return states\n\n\ndef _merge_to_flat_state(states: tp.Iterable[tp.Any]):\n flat_state: list[tuple[PathParts, tp.Any]] = []\n\n for state in states:\n if isinstance(state, dict | State):\n flat_state.extend(traversals.flatten_to_sequence(state))\n elif isinstance(state, FlatState):\n flat_state.extend(state)\n else:\n flat_state.append(((), state))\n\n flat_state.sort()\n return [value for _, value in flat_state]\n\n\ndef merge( # type: ignore[invalid-annotation]\n graphdef: GraphDef[A],\n state: tp.Any,\n /,\n *states: tp.Any,\n copy: bool = False,\n) -> A:\n \"\"\"The inverse of :func:`flax.nnx.split`.\n\n ``nnx.merge`` takes a :class:`flax.nnx.GraphDef` and one or more :class:`flax.nnx.State`'s\n and creates a new node with the same structure as the original node.\n\n Recall: :func:`flax.nnx.split` is used to represent a :class:`flax.nnx.Module`\n by: 1) a static ``nnx.GraphDef`` that captures its Pythonic static information;\n and 2) one or more :class:`flax.nnx.Variable` ``nnx.State``'(s) that capture\n its ``jax.Array``'s in the form of JAX pytrees.\n\n ``nnx.merge`` is used in conjunction with ``nnx.split`` to switch seamlessly\n between stateful and stateless representations of the graph.\n\n Example usage::\n\n >>> from flax import nnx\n >>> import jax, jax.numpy as jnp\n ...\n >>> class Foo(nnx.Module):\n ... def __init__(self, rngs):\n ... self.batch_norm = nnx.BatchNorm(2, rngs=rngs)\n ... self.linear = nnx.Linear(2, 3, rngs=rngs)\n ...\n >>> node = Foo(nnx.Rngs(0))\n >>> graphdef, params, batch_stats = nnx.split(node, nnx.Param, nnx.BatchStat)\n ...\n >>> new_node = nnx.merge(graphdef, params, batch_stats)\n >>> assert isinstance(new_node, Foo)\n >>> assert isinstance(new_node.batch_norm, nnx.BatchNorm)\n >>> assert isinstance(new_node.linear, nnx.Linear)\n\n ``nnx.split`` and ``nnx.merge`` are primarily used to interact directly with JAX\n transformations (refer to\n `Functional API `__\n for more information.\n\n Args:\n graphdef: A :class:`flax.nnx.GraphDef` object.\n state: A :class:`flax.nnx.State` object.\n *states: Additional :class:`flax.nnx.State` objects.\n copy: Whether to create new copies of the Variables in the states, defaults to ``False``.\n Returns:\n The merged :class:`flax.nnx.Module`.\n \"\"\"\n if isinstance(state, list):\n if len(states) != 0:\n raise ValueError(f'Only one state can be passed as a list.')\n _state = state\n else:\n _state = _merge_to_flat_state((state, *states))\n node = unflatten(graphdef, _state, copy_variables=copy)\n return node\n\n\ndef update(node, state: tp.Any, /, *states: tp.Any) -> None:\n \"\"\"Update the given graph node with a new state(s) in-place.\n\n Example usage::\n\n >>> from flax import nnx\n >>> import jax, jax.numpy as jnp\n\n >>> x = jnp.ones((1, 2))\n >>> y = jnp.ones((1, 3))\n >>> model = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n\n >>> def loss_fn(model, x, y):\n ... return jnp.mean((y - model(x))**2)\n >>> prev_loss = loss_fn(model, x, y)\n\n >>> grads = nnx.grad(loss_fn)(model, x, y)\n >>> new_state = jax.tree.map(lambda p, g: p - 0.1*g, nnx.state(model), grads)\n >>> nnx.update(model, new_state)\n >>> assert loss_fn(model, x, y) < prev_loss\n\n Args:\n node: A graph node to update.\n state: A :class:`State` object.\n *states: Additional :class:`State` objects.\n \"\"\"\n if states:\n if isinstance(node, Variable):\n non_empty_states = [\n _state\n for _state in (state, *states)\n if not isinstance(_state, tp.Mapping) or _state\n ]\n if len(non_empty_states) != 1:\n all_states = (state, *states)\n raise ValueError(\n f'Expected exactly one non-empty state, got: {all_states!r}'\n )\n state = non_empty_states[0]\n else:\n state = statelib.merge_state(state, *states)\n _graph_update_dynamic(node, state)\n\n\n@tp.overload\ndef state(node, /, *, graph: bool | None = None) -> GraphState: ...\n@tp.overload\ndef state(node, first: filterlib.Filter, /, *, graph: bool | None = None) -> GraphState: ...\n@tp.overload\ndef state(\n node,\n first: filterlib.Filter,\n second: filterlib.Filter,\n /,\n *filters: filterlib.Filter,\n graph: bool | None = None,\n) -> tuple[GraphState, ...]: ...\ndef state(\n node,\n *filters: filterlib.Filter,\n graph: bool | None = None,\n) -> tp.Union[GraphState, tuple[GraphState, ...]]:\n \"\"\"Similar to :func:`split` but only returns the :class:`State`'s indicated by the filters.\n\n Example usage::\n\n >>> from flax import nnx\n\n >>> class Model(nnx.Module):\n ... def __init__(self, rngs):\n ... self.batch_norm = nnx.BatchNorm(2, rngs=rngs)\n ... self.linear = nnx.Linear(2, 3, rngs=rngs)\n ... def __call__(self, x):\n ... return self.linear(self.batch_norm(x))\n\n >>> model = Model(rngs=nnx.Rngs(0))\n >>> # get the learnable parameters from the batch norm and linear layer\n >>> params = nnx.state(model, nnx.Param)\n >>> # get the batch statistics from the batch norm layer\n >>> batch_stats = nnx.state(model, nnx.BatchStat)\n >>> # get them separately\n >>> params, batch_stats = nnx.state(model, nnx.Param, nnx.BatchStat)\n >>> # get them together\n >>> state = nnx.state(model)\n\n Args:\n node: A graph node object.\n *filters: One or more :class:`Variable` objects to filter by.\n graph: If ``True`` (default), uses graph-mode which supports the full\n NNX feature set including shared references. If ``False``, uses\n tree-mode which treats Modules as regular JAX pytrees, avoiding\n the overhead of the graph protocol.\n Returns:\n One or more :class:`State` mappings.\n \"\"\"\n if graph is None:\n graph = set_graph_mode.current_value()\n _, flat_state = flatten(node, graph=graph)\n state = flat_state.to_nested_state()\n\n states: GraphState | tuple[GraphState, ...]\n if len(filters) == 0:\n states = state # type: ignore[assignment]\n elif len(filters) == 1:\n states = statelib.filter_state(state, filters[0])\n else:\n states = statelib.filter_state(state, filters[0], filters[1], *filters[2:])\n\n return states\n\n\nvariables = state\n\n\ndef map(\n f: tp.Callable[[tuple, tp.Any], tp.Any],\n node: A,\n /,\n *,\n graph: bool | None = None,\n) -> A:\n \"\"\"Map a function over the state of a graph node.\n\n ``map`` extracts the state from ``node`` using :func:`split`, applies ``f``\n to every ``(path, value)`` pair using :func:`map_state`, and returns a\n new node with the mapped values merged back into the original structure.\n Note that the leaves in the state are :class:`Variable` objects, so ``f``\n should handle them accordingly.\n\n Example usage::\n\n >>> from flax import nnx\n >>> import jax.numpy as jnp\n\n >>> model = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n >>> new_model = nnx.map(lambda path, v: v.replace(jnp.zeros_like(v)), model)\n >>> assert jnp.all(new_model.kernel[...] == 0)\n >>> assert jnp.all(new_model.bias[...] == 0)\n\n Args:\n f: A callable ``(path, value) -> new_value`` applied to each leaf in the\n state. ``path`` is a tuple of path parts and ``value`` is the\n corresponding leaf (typically a :class:`Variable`).\n node: A graph node object.\n graph: If ``True``, uses graph-mode which supports the full\n NNX feature set including shared references. If ``False``, uses\n tree-mode which treats Modules as regular JAX pytrees, avoiding\n the overhead of the graph protocol.\n Returns:\n A :class:`State` with the mapped values.\n \"\"\"\n graphdef, state = split(node, graph=graph)\n state = statelib.map_state(f, state)\n return merge(graphdef, state)\n\n\ndef graphdef(\n node: tp.Any, /, *, graph: bool | None = None,\n) -> GraphDef[tp.Any]:\n \"\"\"Get the :class:`GraphDef` of the given graph node.\n\n Example usage::\n\n >>> from flax import nnx\n\n >>> model = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n >>> graphdef, _ = nnx.split(model)\n >>> assert graphdef == nnx.graphdef(model)\n\n Args:\n node: A graph node object.\n graph: If ``True`` (default), uses graph-mode which supports the full\n NNX feature set including shared references. If ``False``, uses\n tree-mode which treats Modules as regular JAX pytrees, avoiding\n the overhead of the graph protocol.\n Returns:\n The :class:`GraphDef` of the :class:`Module` object.\n \"\"\"\n if graph is None:\n graph = set_graph_mode.current_value()\n graphdef, _ = flatten(node, graph=graph)\n return graphdef\n\n\n@tp.overload\ndef pop(\n node,\n filter: filterlib.Filter,\n /,\n) -> GraphState: ...\n\n\n@tp.overload\ndef pop(\n node,\n filter: filterlib.Filter,\n filter2: filterlib.Filter,\n /,\n *filters: filterlib.Filter,\n) -> tuple[GraphState, ...]: ...\n\n\ndef pop(\n node, *filters: filterlib.Filter\n) -> tp.Union[GraphState, tuple[GraphState, ...]]:\n \"\"\"Pop one or more :class:`Variable` types from the graph node.\n\n Example usage::\n\n >>> from flax import nnx\n >>> import jax.numpy as jnp\n\n >>> class Model(nnx.Module):\n ... def __init__(self, rngs):\n ... self.linear1 = nnx.Linear(2, 3, rngs=rngs)\n ... self.linear2 = nnx.Linear(3, 4, rngs=rngs)\n ... def __call__(self, x):\n ... x = self.linear1(x)\n ... self.i = nnx.Intermediate(x)\n ... x = self.linear2(x)\n ... return x\n\n >>> x = jnp.ones((1, 2))\n >>> model = Model(rngs=nnx.Rngs(0))\n >>> assert not hasattr(model, 'i')\n >>> y = model(x)\n >>> assert hasattr(model, 'i')\n\n >>> intermediates = nnx.pop(model, nnx.Intermediate)\n >>> assert intermediates['i'].shape == (1, 3)\n >>> assert not hasattr(model, 'i')\n\n Args:\n node: A graph node object.\n *filters: One or more :class:`Variable` objects to filter by.\n Returns:\n The popped :class:`State` containing the :class:`Variable`\n objects that were filtered for.\n \"\"\"\n if len(filters) == 0:\n raise ValueError('Expected at least one filter')\n\n id_to_index: dict[int, Index] = {}\n path_parts: PathParts = ()\n predicates = tuple(filterlib.to_predicate(filter) for filter in filters)\n flat_states: tuple[dict[PathParts, LeafType], ...] = tuple(\n {} for _ in predicates\n )\n _graph_pop(\n node=node,\n id_to_index=id_to_index,\n path_parts=path_parts,\n flat_states=flat_states,\n predicates=predicates,\n )\n states = tuple(\n statelib.from_flat_state(flat_state) for flat_state in flat_states\n )\n\n if len(states) == 1:\n return states[0]\n else:\n return states\n\n\ndef clone(node: Node, variables: bool = True, *, graph: bool | None = None) -> Node:\n \"\"\"Create a deep copy of the given graph node.\n\n Example usage::\n\n >>> from flax import nnx\n\n >>> model = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n >>> cloned_model = nnx.clone(model)\n >>> model.bias[...] += 1\n >>> assert (model.bias[...] != cloned_model.bias[...]).all()\n\n Args:\n node: A graph node object.\n variables: If ``True`` (default) copies of the :class:`Variable` objects are created,\n otherwise the Variables are shared between the original and cloned node.\n graph: If ``True`` (default), uses graph-mode which supports the full\n NNX feature set including shared references. If ``False``, uses\n tree-mode which treats Modules as regular JAX pytrees, avoiding\n the overhead of the graph protocol.\n Returns:\n A deep copy of the :class:`Module` object.\n \"\"\"\n graphdef, state = split(node, graph=graph)\n return merge(graphdef, state, copy=variables)\n\n\ndef vars_as(\n node: A,\n /,\n *,\n hijax: bool | None = None,\n ref: bool | None = None,\n mutable: bool | None = None,\n only: filterlib.Filter = ...,\n allow_duplicates: bool = False,\n) -> A:\n \"\"\" \"\"\"\n new_attrs: dict[str, bool] = {}\n if hijax is not None:\n new_attrs['hijax'] = hijax\n if ref is not None:\n new_attrs['ref'] = ref\n if mutable is not None:\n new_attrs['mutable'] = mutable\n\n def _different_vars(path, x):\n return isinstance(x, Variable) and any(\n getattr(x, attr) != value for attr, value in new_attrs.items()\n )\n\n only = filterlib.All(_different_vars, only)\n predicate = filterlib.to_predicate(only)\n\n if not allow_duplicates and (\n all_duplicates := find_duplicates(node, only=only)\n ):\n duplicates_strs = '\\n ---'\n for node_duplicates in all_duplicates:\n for path in node_duplicates:\n path_str = '/'.join(builtins.map(str, path))\n duplicates_strs += f'\\n {path_str}'\n duplicates_strs += '\\n ---'\n raise ValueError(f'Found duplicate at paths:{duplicates_strs}')\n\n def _to_refs(jax_path, x):\n if predicate(jax_to_nnx_path(jax_path), x):\n assert isinstance(x, Variable)\n variable = x.copy(**new_attrs)\n return variable\n return x\n\n node = jax.tree.map_with_path(\n _to_refs, node, is_leaf=lambda x: isinstance(x, Variable)\n )\n return node\n\n\ndef pure(tree: A) -> A:\n \"\"\"Returns a new tree with all ``Variable`` objects replaced with inner values.\n\n This can be used to remove Variable metadata when its is not needed for tasks like\n serialization or exporting.\n\n Example::\n\n >>> from flax import nnx\n >>> import jax\n >>> import jax.numpy as jnp\n ...\n >>> model = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n >>> graphdef, state = nnx.split(model)\n >>> jax.tree.map(jnp.shape, state)\n State({\n 'bias': Param(\n value=(3,)\n ),\n 'kernel': Param(\n value=(2, 3)\n )\n })\n >>> pure_state = nnx.pure(state)\n >>> jax.tree.map(jnp.shape, pure_state)\n State({\n 'bias': (3,),\n 'kernel': (2, 3)\n })\n\n Args:\n tree: A pytree potentially containing ``Variable`` objects.\n Returns:\n A new pytree with all ``Variable`` objects replaced with their\n inner values.\n \"\"\"\n\n def _pure_fn(x):\n if isinstance(x, Variable):\n return pure(x.get_raw_value())\n elif variablelib.is_array_ref(x):\n return x[...]\n return x\n\n return jax.tree.map(\n _pure_fn,\n tree,\n is_leaf=lambda x: isinstance(x, Variable),\n )\n\n\ndef call(\n graphdef_state: tuple[GraphDef[A], GraphState], /\n) -> ApplyCaller[tuple[GraphDef[A], GraphState]]:\n \"\"\"Calls a method underlying graph node defined by a (GraphDef, State) pair.\n\n ``call`` takes a ``(GraphDef, State)`` pair and creates a proxy object that can be\n used to call methods on the underlying graph node. When a method is called, the\n output is returned along with a new (GraphDef, State) pair that represents the\n updated state of the graph node. ``call`` is equivalent to :func:`merge` > ``method``\n > :func:`split` but is more convenient to use in pure JAX functions.\n\n Example::\n\n >>> from flax import nnx\n >>> import jax\n >>> import jax.numpy as jnp\n ...\n >>> class StatefulLinear(nnx.Module):\n ... def __init__(self, din, dout, rngs):\n ... self.w = nnx.Param(jax.random.uniform(rngs(), (din, dout)))\n ... self.b = nnx.Param(jnp.zeros((dout,)))\n ... self.count = Variable(jnp.array(0, dtype=jnp.uint32))\n ...\n ... def increment(self):\n ... self.count[...] += 1\n ...\n ... def __call__(self, x):\n ... self.increment()\n ... return x @ self.w + self.b\n ...\n >>> linear = StatefulLinear(3, 2, nnx.Rngs(0))\n >>> linear_state = nnx.split(linear)\n ...\n >>> @jax.jit\n ... def forward(x, linear_state):\n ... y, linear_state = nnx.call(linear_state)(x)\n ... return y, linear_state\n ...\n >>> x = jnp.ones((1, 3))\n >>> y, linear_state = forward(x, linear_state)\n >>> y, linear_state = forward(x, linear_state)\n ...\n >>> linear = nnx.merge(*linear_state)\n >>> linear.count[...]\n Array(2, dtype=uint32)\n\n The proxy object returned by ``call`` supports indexing and attribute access\n to access nested methods. In the example below, the ``increment`` method indexing\n is used to call the ``increment`` method of the ``StatefulLinear`` module\n at the ``b`` key of a ``nodes`` dictionary.\n\n >>> class StatefulLinear(nnx.Module):\n ... def __init__(self, din, dout, rngs):\n ... self.w = nnx.Param(jax.random.uniform(rngs(), (din, dout)))\n ... self.b = nnx.Param(jnp.zeros((dout,)))\n ... self.count = nnx.Variable(jnp.array(0, dtype=jnp.uint32))\n ...\n ... def increment(self):\n ... self.count[...] += 1\n ...\n ... def __call__(self, x):\n ... self.increment()\n ... return x @ self.w + self.b\n ...\n >>> rngs = nnx.Rngs(0)\n >>> nodes = dict(\n ... a=StatefulLinear(3, 2, rngs),\n ... b=StatefulLinear(2, 1, rngs),\n ... )\n ...\n >>> node_state = nnx.split(nodes)\n >>> # use attribute access\n >>> _, node_state = nnx.call(node_state)['b'].increment()\n ...\n >>> nodes = nnx.merge(*node_state)\n >>> nodes['a'].count[...]\n Array(0, dtype=uint32)\n >>> nodes['b'].count[...]\n Array(1, dtype=uint32)\n \"\"\"\n\n def pure_caller(accessor: DelayedAccessor, *args, **kwargs):\n node = merge(*graphdef_state)\n method = accessor(node)\n out = method(*args, **kwargs)\n return out, split(node)\n\n return CallableProxy(pure_caller) # type: ignore\n\n\ndef set_metadata(\n node: tp.Any, /, *, only: filterlib.Filter = Variable, **metadata: tp.Any\n) -> None:\n \"\"\"Sets the metadata of all :class:`Variable` objects in the given graph node in-place.\n\n Example::\n\n >>> from flax import nnx\n >>> import jax, jax.numpy as jnp\n ...\n >>> class Foo(nnx.Module):\n ... def __init__(self):\n ... self.param = nnx.Param(0.0)\n ... self.variable = nnx.Variable(0.0)\n ...\n >>> node = Foo()\n ...\n >>> # set differentiable to False for all nnx.Param objects\n >>> nnx.set_metadata(node, differentiable=False, only=nnx.Param)\n ...\n >>> # check that only the nnx.Param was updated\n >>> assert node.param.get_metadata('differentiable') is False\n\n Args:\n node: A graph node object.\n only: A Filter to specify which :class:`Variable` objects to set metadata for.\n metadata: Key-value pairs to set as metadata on the :class:`Variable` objects.\n \"\"\"\n def _set_metadata(path: PathParts, variable: V) -> None:\n del path # unused\n if isinstance(variable, Variable):\n variable.set_metadata(**metadata)\n\n # inplace update of variable_state metadata\n map_state(_set_metadata, state(node, only))\n\n\ndef iter_graph(\n node: tp.Any, /, *, graph: bool | None = None,\n) -> tp.Iterator[tuple[PathParts, tp.Any]]:\n \"\"\"Iterates over all nested nodes and leaves of the given graph node, including the current node.\n\n ``iter_graph`` creates a generator that yields path and value pairs, where the\n path is a tuple of strings or integers representing the path to the value from\n the root. Repeated nodes are visited only once. Leaves include static values.\n\n Example::\n\n >>> from flax import nnx\n >>> import jax.numpy as jnp\n ...\n >>> class Linear(nnx.Module):\n ... def __init__(self, din, dout, *, rngs: nnx.Rngs):\n ... self.din, self.dout = din, dout\n ... self.w = nnx.Param(jax.random.uniform(rngs.next(), (din, dout)))\n ... self.b = nnx.Param(jnp.zeros((dout,)))\n ...\n >>> module = Linear(3, 4, rngs=nnx.Rngs(0))\n >>> graph = [module, module]\n ...\n >>> for path, value in nnx.iter_graph(graph):\n ... print(path, type(value).__name__)\n ...\n (0, '_pytree__nodes') HashableMapping\n (0, '_pytree__state') PytreeState\n (0, 'b') Param\n (0, 'din') int\n (0, 'dout') int\n (0, 'w') Param\n (0,) Linear\n () list\n\n Args:\n node: A graph node object.\n graph: If ``True`` (default), uses graph-mode which supports the full\n NNX feature set including shared references. If ``False``, uses\n tree-mode which treats Modules as regular JAX pytrees, avoiding\n the overhead of the graph protocol.\n \"\"\"\n if graph is None:\n graph = set_graph_mode.current_value()\n if graph:\n return _iter_graph(node)\n else:\n return _iter_tree(node)\n\n\ndef _iter_graph(node: tp.Any, /) -> tp.Iterator[tuple[PathParts, tp.Any]]:\n visited: set[int] = set()\n stack: list[tuple[PathParts, tp.Any, bool]] = [((), node, False)]\n while stack:\n path_parts, node, traversed = stack.pop(-1)\n if traversed or not (is_node(node) or isinstance(node, Variable)):\n yield path_parts, node\n continue\n\n if id(node) in visited:\n continue\n visited.add(id(node))\n\n if (node_impl := get_node_impl(node)) is None:\n yield path_parts, node\n continue\n\n stack.append((path_parts, node, True))\n for key, child in reversed(node_impl.node_dict(node).items()):\n stack.append(((*path_parts, key), child, False))\n\n\ndef _iter_tree(node: tp.Any, /) -> tp.Iterator[tuple[PathParts, tp.Any]]:\n in_progress: dict[int, str] = {}\n seen_refs: dict[int, str] = {}\n stack: list[tuple[PathParts, tp.Any, bool]] = [((), node, False)]\n while stack:\n path, current, traversed = stack.pop()\n\n if traversed:\n in_progress.pop(id(current), None)\n yield path, current\n continue\n\n if not is_pytree_node(current, check_graph_registry=False):\n _check_valid_pytree(current, 'iter_graph', '/'.join(builtins.map(str, path)))\n if isinstance(current, Variable) or variablelib.is_array_ref(current):\n obj_id = id(current)\n str_path = '/'.join(builtins.map(str, path))\n if obj_id in seen_refs:\n raise ValueError(\n f'Duplicate {current}\\nfound at paths:\\n\\n'\n f' - {seen_refs[obj_id]}\\n'\n f' - {str_path}\\n\\n'\n 'Tree mode (graph=False) does not support shared references. '\n + _tree_mode_suggestion_api('iter_graph')\n )\n seen_refs[obj_id] = str_path\n yield path, current\n continue\n\n obj_id = id(current)\n str_path = '/'.join(builtins.map(str, path))\n if obj_id in in_progress:\n raise ValueError(\n f'Cycle detected for {type(current).__name__}\\nfound at paths:\\n\\n'\n f' - {in_progress[obj_id]}\\n'\n f' - {str_path}\\n\\n'\n 'Cycles are not supported with graph=False. '\n + _tree_mode_suggestion_api('iter_graph')\n )\n in_progress[obj_id] = str_path\n\n stack.append((path, current, True))\n children, _ = jax.tree_util.tree_flatten_with_path(\n current, is_leaf=lambda x: x is not current\n )\n for jax_key_path, child in reversed(children):\n key = _key_path_to_key(jax_key_path[0])\n stack.append(((*path, key), child, False))\n\n\ndef iter_children(\n node: tp.Any, /, *, graph: bool | None = None,\n) -> tp.Iterator[tuple[Key, tp.Any]]:\n \"\"\"Iterates over all immediate child nodes of a given node. This\n function is similar to :func:`iter_graph`, except it only iterates over the\n immediate children, and does not recurse further down.\n\n Specifically, this function creates a generator that yields the key and the\n child node instance, where the key is a string representing the attribute\n name to access the corresponding child.\n\n Example::\n\n >>> from flax import nnx\n ...\n >>> class SubModule(nnx.Module):\n ... def __init__(self, din, dout, rngs):\n ... self.linear1 = nnx.Linear(din, dout, rngs=rngs)\n ... self.linear2 = nnx.Linear(din, dout, rngs=rngs)\n ...\n >>> class Block(nnx.Module):\n ... def __init__(self, din, dout, *, rngs: nnx.Rngs):\n ... self.linear = nnx.Linear(din, dout, rngs=rngs)\n ... self.submodule = SubModule(din, dout, rngs=rngs)\n ... self.dropout = nnx.Dropout(0.5)\n ... self.batch_norm = nnx.BatchNorm(10, rngs=rngs)\n ...\n >>> model = Block(2, 5, rngs=nnx.Rngs(0))\n >>> for path, module in nnx.iter_children(model):\n ... print(path, type(module).__name__)\n ...\n batch_norm BatchNorm\n dropout Dropout\n linear Linear\n submodule SubModule\n\n Args:\n node: A graph node object.\n graph: If ``True`` (default), uses graph-mode which supports the full\n NNX feature set including shared references. If ``False``, uses\n tree-mode which treats Modules as regular JAX pytrees, avoiding\n the overhead of the graph protocol.\n \"\"\"\n if graph is None:\n graph = set_graph_mode.current_value()\n if graph:\n node_impl = get_node_impl(node)\n if node_impl is None:\n raise ValueError(\n f'Expected a graph node, got {type(node).__name__}. '\n 'If this is a regular pytree, use graph=False.'\n )\n node_dict = node_impl.node_dict(node)\n for key, value in node_dict.items():\n if is_graph_node(value):\n yield key, value\n else:\n _check_valid_pytree(node, 'iter_children')\n if not is_pytree_node(node, check_graph_registry=False):\n raise ValueError(\n f'Expected a pytree node, got {type(node).__name__}. '\n 'If this is a graph node, use graph=True.'\n )\n children, _ = jax.tree_util.tree_flatten_with_path(\n node, is_leaf=lambda x: x is not node\n )\n for jax_key_path, child in children:\n if is_graph_node(child):\n key = _key_path_to_key(jax_key_path[0])\n yield key, child\n\n\ndef recursive_map(\n f: tp.Callable[[PathParts, tp.Any], tp.Any],\n node: tp.Any,\n /,\n *,\n graph: bool | None = None,\n):\n \"\"\"Recursively applies a function to all nodes and leaves of the given graph node.\n\n Example::\n\n >>> from flax import nnx\n >>> class MyModule(nnx.Module):\n ... def __init__(self, *, rngs: nnx.Rngs):\n ... self.lin = nnx.Linear(16, 16, rngs=rngs)\n ... self.conv = nnx.Conv(16, 3, 1, 1, rngs=rngs)\n ...\n >>> def print_modules(path, node):\n ... if isinstance(node, nnx.Module):\n ... s = \".\" + \".\".join(path)\n ... print(f\"Path = {s:<10}{node.__class__.__name__}\")\n ... return node\n ...\n >>> model = MyModule(rngs=nnx.Rngs(0))\n >>> new_model = nnx.recursive_map(print_modules, model)\n ...\n Path = .conv Conv\n Path = .lin Linear\n Path = . MyModule\n\n Args:\n f: A function that takes a path and a node and returns a new node.\n node: A graph node object.\n graph: If ``True`` (default), uses graph-mode which supports the full\n NNX feature set including shared references. If ``False``, uses\n tree-mode which treats Modules as regular JAX pytrees, avoiding\n the overhead of the graph protocol.\n \"\"\"\n if graph is None:\n graph = set_graph_mode.current_value()\n if graph:\n node = clone(node, variables=False, graph=True)\n path_parts: PathParts = ()\n visited: set[int] = set()\n results: dict[int, tp.Any] = {}\n return _recursive_map_graph(f, node, path_parts, visited, results)\n else:\n return _recursive_map_tree(f, node)\n\n\ndef _recursive_map_graph(\n f: tp.Callable[[PathParts, tp.Any], tp.Any],\n node: tp.Any,\n path: PathParts,\n visited: set[int],\n results: dict[int, tp.Any],\n) -> tp.Any:\n node_id = id(node)\n if node_id in visited:\n if node_id in results:\n return results[node_id]\n path_str = '/'.join(builtins.map(str, path))\n raise ValueError(\n f\"Found cycle in the graph at path '{path_str}'. Node of type\"\n f' {type(node)} has already been visited but has not been returned yet.'\n )\n node_impl = get_node_impl(node)\n if (\n type(node_impl) is GraphNodeImpl\n or isinstance(node, Variable)\n or is_array_ref(node)\n ):\n visited.add(node_id)\n if node_impl is not None:\n for key, value in node_impl.node_dict(node).items():\n new_value = _recursive_map_graph(f, value, (*path, key), visited, results)\n if new_value is not value:\n if node_impl.set_key is not None and value is not new_value:\n node_impl.set_key(node, key, new_value)\n else:\n raise ValueError(\n f\"Cannot update key '{key}' for node of type '{type(node)}'\"\n ' because the node does not support mutation.'\n )\n\n new_node = f(path, node)\n results[node_id] = new_node\n return new_node\n\n\ndef _recursive_map_tree(\n f: tp.Callable[[PathParts, tp.Any], tp.Any],\n node: tp.Any,\n) -> tp.Any:\n in_progress: dict[int, str] = {}\n seen_refs: dict[int, str] = {}\n\n def _recurse(path: PathParts, current: tp.Any) -> tp.Any:\n if not is_pytree_node(current, check_graph_registry=False):\n _check_valid_pytree(current, 'recursive_map', '/'.join(builtins.map(str, path)))\n if isinstance(current, Variable) or is_array_ref(current):\n obj_id = id(current)\n str_path = '/'.join(builtins.map(str, path))\n if obj_id in seen_refs:\n raise ValueError(\n f'Duplicate {current}\\nfound at paths:\\n\\n'\n f' - {seen_refs[obj_id]}\\n'\n f' - {str_path}\\n\\n'\n 'Tree mode (graph=False) does not support shared references. '\n + _tree_mode_suggestion_api('recursive_map')\n )\n seen_refs[obj_id] = str_path\n return f(path, current)\n\n obj_id = id(current)\n str_path = '/'.join(builtins.map(str, path))\n if obj_id in in_progress:\n raise ValueError(\n f'Cycle detected for {type(current).__name__}\\nfound at paths:\\n\\n'\n f' - {in_progress[obj_id]}\\n'\n f' - {str_path}\\n\\n'\n 'Cycles are not supported with graph=False. '\n + _tree_mode_suggestion_api('recursive_map')\n )\n in_progress[obj_id] = str_path\n\n children_with_path, treedef = jax.tree_util.tree_flatten_with_path(\n current, is_leaf=lambda x: x is not current\n )\n new_children = []\n for jax_key_path, child in children_with_path:\n key = _key_path_to_key(jax_key_path[0])\n new_child = _recurse((*path, key), child)\n new_children.append(new_child)\n\n new_node = treedef.unflatten(new_children)\n result = f(path, new_node)\n\n in_progress.pop(obj_id, None)\n return result\n\n return _recurse((), node)\n\n\ndef find_duplicates(node: tp.Any, /, *, only: filterlib.Filter = ...) -> list[list[PathParts]]:\n \"\"\"Finds duplicate nodes or node leaves in the given node.\n\n This function traverses the graph node and collects paths to nodes and leaves\n that have the same identity. It returns a list of lists, where each inner list\n contains paths to nodes or leaves that are duplicates.\n\n Example::\n\n >>> from flax import nnx\n >>> import jax.numpy as jnp\n ...\n >>> class SharedVariables(nnx.Module):\n ... def __init__(self):\n ... self.a = nnx.Param(jnp.array(1.0))\n ... self.b = nnx.Param(jnp.array(2.0))\n ... self.c = self.b # shared Variable\n ...\n >>> model = SharedVariables()\n >>> duplicates = nnx.find_duplicates(model)\n >>> len(duplicates)\n 1\n >>> for path in duplicates[0]:\n ... print(path)\n ('b',)\n ('c',)\n\n ``find_duplicates`` will also find duplicates nodes such as Modules that are\n referenced multiple times in the graph::\n\n >>> class SharedModules(nnx.Module):\n ... def __init__(self, rngs: nnx.Rngs):\n ... self.a = nnx.Linear(1, 1, rngs=rngs)\n ... self.b = nnx.Linear(1, 1, rngs=rngs)\n ... self.c = self.a # shared Module\n ...\n >>> model = SharedModules(nnx.Rngs(0))\n >>> for duplicate_paths in nnx.find_duplicates(model):\n ... print(duplicate_paths)\n [('a',), ('c',)]\n\n Args:\n node: A graph node object.\n only: A Filter to specify which nodes or leaves to consider for duplicates.\n Returns:\n A list of lists, where each inner list contains the different paths for a\n for a duplicate node or leaf.\n \"\"\"\n node_paths: dict[int, list[PathParts]] = {}\n duplicate_candidate = filterlib.to_predicate(only)\n _node_paths(node, node_paths, (), duplicate_candidate)\n _duplicates = [paths for paths in node_paths.values() if len(paths) > 1]\n return _duplicates\n\n\ndef _node_paths(\n node: tp.Any,\n node_paths: dict[int, list[PathParts]],\n path: PathParts,\n duplicate_candidate: filterlib.Predicate,\n /,\n):\n _is_graph_node = is_graph_node(node)\n _is_pytree_node = is_pytree_node(node)\n _is_node_leaf = is_node_leaf(node)\n\n if _is_graph_node or _is_pytree_node or _is_node_leaf:\n node_id = id(node)\n if node_id in node_paths:\n if (_is_graph_node or _is_node_leaf) and duplicate_candidate(path, node):\n node_paths[node_id].append(path)\n return\n if _is_graph_node or _is_node_leaf:\n node_paths[node_id] = [path]\n node_impl = get_node_impl(node)\n if node_impl is None:\n return\n node_dict = node_impl.node_dict(node)\n for key, value in node_dict.items():\n _node_paths(value, node_paths, (*path, key), duplicate_candidate)\n\n\n@jax.tree_util.register_static\n@dataclasses.dataclass(frozen=True, slots=True)\nclass Static(tp.Generic[A]):\n \"\"\"An empty pytree node that treats its inner value as static.\n ``value`` must define ``__eq__`` and ``__hash__``.\n \"\"\"\n\n value: A\n\n\n# ---------------------------------------------------------\n# Pytree\n# ---------------------------------------------------------\nclass GenericPytree: ...\n\n\nfrom jax._src.tree_util import _registry as JAX_PYTREE_REGISTRY\n\n\ndef is_pytree_node(\n x: tp.Any, *, check_graph_registry: bool = True,\n) -> bool:\n if check_graph_registry and type(x) in GRAPH_REGISTRY:\n return False\n elif isinstance(x, Variable):\n return False\n elif type(x) in JAX_PYTREE_REGISTRY:\n return True\n elif isinstance(x, tuple):\n return True\n else:\n return False\n\n\ndef _key_path_to_key(key: tp.Any) -> Key:\n if isinstance(key, jax.tree_util.SequenceKey):\n return key.idx\n elif isinstance(\n key, (jax.tree_util.DictKey, jax.tree_util.FlattenedIndexKey)\n ):\n if not is_key_like(key.key): # type: ignore[not-supported-yet]\n raise ValueError(\n f'Invalid key: {key.key}. May be due to its type not being hashable or comparable.'\n )\n return key.key\n elif isinstance(key, jax.tree_util.GetAttrKey):\n return key.name\n else:\n return str(key)\n\n\ndef jax_to_nnx_path(jax_path: tuple, /):\n return tuple(_key_path_to_key(part) for part in jax_path)\n\n\nclass IndexesPytreeDef(tp.NamedTuple):\n key_index: HashableMapping[Key, int]\n treedef: jax.tree_util.PyTreeDef\n\n\ndef _flatten_pytree(pytree: tp.Any):\n leaves, treedef = jax.tree_util.tree_flatten_with_path(\n pytree, is_leaf=lambda x: x is not pytree\n )\n nodes = [(_key_path_to_key(path[0]), value) for path, value in leaves]\n key_index = HashableMapping(\n {key: i for i, (key, _) in enumerate(nodes)}, copy=False\n )\n # Sort by key to match the path-sorted order used by _merge_to_flat_state.\n # key_index records the original jax tree_flatten order so _unflatten_pytree\n # can restore it before calling treedef.unflatten.\n nodes.sort()\n return nodes, IndexesPytreeDef(key_index, treedef)\n\n\ndef _unflatten_pytree(\n nodes: tuple[tuple[Key, tp.Any], ...], metadata: IndexesPytreeDef\n):\n # sort to original order\n sorted_nodes = sorted(nodes, key=lambda x: metadata.key_index[x[0]])\n pytree = metadata.treedef.unflatten(value for _, value in sorted_nodes)\n return pytree\n\n\nPYTREE_NODE_IMPL = PytreeNodeImpl(\n type=GenericPytree,\n flatten=_flatten_pytree,\n unflatten=_unflatten_pytree, # type: ignore\n set_key=None,\n pop_key=None,\n)\ndef _list_set_key(x: list[tp.Any], key: int, value: tp.Any):\n x[key] = value\n\n# common pytrees\n# list\nregister_pytree_node_type(\n list,\n flatten=lambda x: (list(enumerate(x)), None),\n unflatten=lambda nodes, _: [value for _, value in nodes], # type: ignore\n set_key=_list_set_key, # type: ignore\n)\n# tuple\nregister_pytree_node_type(\n tuple,\n flatten=lambda x: (list(enumerate(x)), None),\n unflatten=lambda nodes, _: tuple(value for _, value in nodes), # type: ignore\n)\n\n\ndef _mutable_mapping_set_key(\n x: tp.MutableMapping[Key, tp.Any], key: Key, value: tp.Any\n):\n x[key] = value\n\n\ndef _mutable_mapping_pop_key(x: tp.MutableMapping[Key, tp.Any], key: Key):\n x.pop(key)\n\n\n# dict\nregister_pytree_node_type(\n dict,\n flatten=lambda x: (sorted(x.items()), None),\n unflatten=lambda nodes, _: dict(nodes), # type: ignore\n set_key=_mutable_mapping_set_key,\n pop_key=_mutable_mapping_pop_key,\n)\n# State\nregister_pytree_node_type(\n State,\n flatten=lambda x: (sorted(x.raw_mapping.items()), None),\n unflatten=lambda nodes, _: State(nodes), # type: ignore\n set_key=_mutable_mapping_set_key,\n pop_key=_mutable_mapping_pop_key,\n)\n# None\nregister_pytree_node_type(\n type(None),\n flatten=lambda x: ([], None),\n unflatten=lambda _, __: None, # type: ignore\n)\n" }, { "path": "flax/nnx/helpers.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nimport inspect\nimport typing as tp\n\nimport jax\nimport jax.numpy as jnp\nimport optax\n\nfrom flax.nnx import graphlib, reprlib\nfrom flax.nnx.graphlib import GraphDef\nfrom flax.nnx.module import Module\nfrom flax.nnx.proxy_caller import ApplyCaller\nfrom flax.nnx.rnglib import Rngs\nfrom flax.nnx.statelib import State\nfrom flax.training.train_state import struct\nfrom flax.nnx.variablelib import Variable\n\nA = tp.TypeVar('A')\nM = tp.TypeVar('M', bound=Module)\nTS = tp.TypeVar('TS', bound='TrainState')\n\nclass Dict(reprlib.MappingReprMixin, Module, tp.MutableMapping[str, A]):\n \"\"\"A Module that implements a mutable mapping.\n\n This class provides a way to store and manipulate a mapping of keys to values\n contained a mixed set of data (e.g. Array, Variables, Modules) and static\n (e.g. functions, strings) types.\n\n Example:\n >>> from flax import nnx\n ...\n >>> rngs = nnx.Rngs(0)\n >>> layers = nnx.Dict({\n ... 'kernel1': nnx.Linear(1, 32, rngs=rngs), # data\n ... 'activation1': nnx.relu, # static\n ... 'kernel2': nnx.Linear(32, 1, rngs=rngs), # data\n ... })\n \"\"\"\n @tp.overload\n def __init__(self, iterable: tp.Iterable[tp.Tuple[str, A]], /): ...\n\n @tp.overload\n def __init__(\n self, mapping: tp.Optional[tp.Mapping[str, A]] = None, /, **kwargs: A\n ): ...\n\n def __init__(self, *args, **kwargs):\n for name, value in dict(*args, **kwargs).items():\n setattr(self, name, value)\n\n def __getitem__(self, key) -> A:\n try:\n return getattr(self, key)\n except AttributeError as e:\n raise KeyError(key) from e\n\n def __setitem__(self, key, value):\n setattr(self, key, value)\n\n def __iter__(self) -> tp.Iterator[str]:\n return (k for k in vars(self) if not k.startswith('_pytree__'))\n\n def __len__(self) -> int:\n length = 0\n for _ in self:\n length += 1\n return length\n\n def __hash__(self) -> int:\n return id(self)\n\n def __delitem__(self, key: str) -> None:\n try:\n delattr(self, key)\n except AttributeError as e:\n raise KeyError(key) from e\n\n if tp.TYPE_CHECKING:\n def __getattr__(self, key: str) -> A:\n ...\n def __setattr__(self, key: str, value: A) -> None:\n ...\n\n\nclass List(reprlib.SequenceReprMixin, Module, tp.MutableSequence[A]):\n \"\"\"A Module that implements a mutable sequence.\n\n This class provides a way to store and manipulate a sequence of values\n contained a mixed set of data (e.g. Array, Variables, Modules) and static\n (e.g. functions, strings) types.\n\n Example:\n >>> from flax import nnx\n ...\n >>> rngs = nnx.Rngs(0)\n >>> layers = nnx.List([\n ... nnx.Linear(1, 32, rngs=rngs), # data\n ... nnx.relu, # static\n ... nnx.Linear(32, 1, rngs=rngs), # data\n ... ])\n \"\"\"\n def __init__(self, it: tp.Iterable[A] | None = None, /):\n \"\"\"\n Args:\n it: An iterable of values to initialize the list.\n \"\"\"\n self._length: int = 0\n if it is not None:\n for value in it:\n self.append(value)\n\n def _get_elem(self, key: int) -> A:\n return getattr(self, str(key))\n\n def _set_elem(self, key: int, value: A) -> None:\n setattr(self, str(key), value)\n\n def _del_elem(self, key: int) -> None:\n delattr(self, str(key))\n\n def __len__(self) -> int:\n return self._length\n\n def append(self, value: A) -> None:\n self._set_elem(self._length, value)\n self._length += 1\n\n def insert(self, index: int, value: A) -> None:\n if index < 0:\n index += self._length\n if index < 0:\n index = 0\n if index > self._length:\n index = self._length\n\n # Shift elements to the right\n for i in range(self._length, index, -1):\n self._set_elem(i, self._get_elem(i - 1))\n\n # Insert the new value\n self._set_elem(index, value)\n self._length += 1\n\n def __iter__(self) -> tp.Iterator[A]:\n for i in range(self._length):\n yield self._get_elem(i)\n\n @tp.overload\n def __getitem__(self, index: int) -> A: ...\n @tp.overload\n def __getitem__(self, index: slice) -> tp.List[A]: ...\n def __getitem__(self, index: int | slice) -> A | tp.List[A]:\n if isinstance(index, int):\n if index < 0:\n index += self._length\n if index < 0 or index >= self._length:\n raise IndexError('Index out of bounds')\n return self._get_elem(index)\n elif isinstance(index, slice):\n idxs = list(range(self._length))[index]\n return [self._get_elem(i) for i in idxs]\n else:\n raise TypeError('Invalid index type')\n\n def __setitem__(self, index: int | slice, value: A | tp.Iterable[A]) -> None:\n if isinstance(index, int):\n if index < 0:\n index += self._length\n if index < 0 or index >= self._length:\n raise IndexError('Index out of bounds')\n self._set_elem(index, value) # type: ignore[arg-type]\n elif isinstance(index, slice):\n if not isinstance(value, tp.Iterable):\n raise TypeError('Expected an iterable')\n values = list(value)\n idxs = list(range(self._length))[index]\n if len(idxs) != len(values):\n raise ValueError('Length mismatch')\n for i, v in zip(idxs, values):\n self._set_elem(i, v)\n else:\n raise TypeError('Invalid index type')\n\n def _graph_node_set_key(self, key: str, value: tp.Any):\n if not isinstance(key, int):\n raise KeyError(f'Invalid key: {key}')\n elif key < len(self):\n if isinstance(variable := self[key], Variable) and isinstance(value, Variable):\n variable.update_from_state(value)\n else:\n self[key] = value\n else:\n self.insert(key, value)\n\n def __delitem__(self, index: int | slice) -> None:\n if isinstance(index, int):\n if index < 0:\n index += self._length\n if index < 0 or index >= self._length:\n raise IndexError('Index out of bounds')\n self._del_elem(index)\n for i in range(index + 1, self._length):\n self._set_elem(i - 1, self._get_elem(i))\n self._length -= 1\n elif isinstance(index, slice):\n idxs = list(range(self._length))[index]\n for i in reversed(idxs):\n # implement recursively\n del self[i]\n else:\n raise TypeError('Invalid index type')\n\n _pytree__has_int_keys = True\n\n\nclass Sequential(Module):\n \"\"\"A Module that applies a sequence of callables.\n\n This class provides a way to store and manipulate a sequence of callables\n (e.g. layers, activation functions) and apply them in order.\n\n Example::\n\n >>> from flax import nnx\n >>> import jax.numpy as jnp\n ...\n >>> rngs = nnx.Rngs(0)\n >>> model = nnx.Sequential(\n ... nnx.Linear(1, 4, rngs=rngs), # data\n ... nnx.relu, # static\n ... nnx.Linear(4, 2, rngs=rngs), # data\n ... )\n >>> x = jnp.ones((1, 1))\n >>> y = model(x)\n >>> y.shape\n (1, 2)\n \"\"\"\n\n def __init__(self, *fns: tp.Callable[..., tp.Any]):\n \"\"\"\n Args:\n *fns: A sequence of callables to apply.\n \"\"\"\n self.layers = List(fns)\n\n def __call__(self, *args, rngs: tp.Optional[Rngs] = None, **kwargs) -> tp.Any:\n if len(self.layers) == 0:\n if len(args) == 1:\n return args[0]\n elif len(args) > 0:\n return args\n elif len(kwargs) > 0:\n return kwargs\n else:\n return None\n\n output: tp.Any = None\n\n for i, f in enumerate(self.layers):\n if not callable(f):\n raise TypeError(f'Sequence[{i}] is not callable: {f}')\n if i > 0:\n if isinstance(output, tuple):\n args = output\n kwargs = {}\n elif isinstance(output, dict):\n args = ()\n kwargs = output\n else:\n args = (output,)\n kwargs = {}\n if rngs is not None and has_keyword_arg(f, 'rngs'):\n kwargs['rngs'] = rngs\n\n output = f(*args, **kwargs)\n\n return output\n\n\nclass ModuleDefApply(tp.Protocol, tp.Generic[M]):\n def __call__(\n self, state: State, *states: State\n ) -> ApplyCaller[tuple[State, GraphDef[M]]]: ...\n\n\nclass TrainState(tp.Generic[M], struct.PyTreeNode):\n graphdef: graphlib.GraphDef[M]\n params: State\n opt_state: optax.OptState\n step: jax.Array\n tx: optax.GradientTransformation = struct.field(pytree_node=False)\n\n @classmethod\n def create(\n cls,\n graphdef: graphlib.GraphDef[M],\n *,\n params: State,\n tx: optax.GradientTransformation,\n step: int = 0,\n **kwargs,\n ):\n return cls(\n graphdef=graphdef,\n params=params,\n opt_state=tx.init(params),\n step=jnp.asarray(step),\n tx=tx,\n **kwargs,\n )\n\n if tp.TYPE_CHECKING:\n\n def __getattr__(self, key: str) -> tp.Any: ...\n\n def apply(\n self, state: tp.Union[State, str], *states: tp.Union[State, str]\n ) -> ApplyCaller[tuple[GraphDef[M], State]]:\n states = (state, *states)\n\n _states: list[State] = []\n\n for _state in states:\n if isinstance(_state, str):\n _state_key = _state\n _state = getattr(self, _state_key)\n if not isinstance(_state, State):\n raise TypeError(\n f'Expected {self.__class__.__name__}.{_state_key} to be a State, got {type(_state)}'\n )\n _states.append(_state)\n\n return self.graphdef.apply(*_states)\n\n def apply_gradients(self: TS, grads: State, **kwargs) -> TS:\n updates, opt_state = self.tx.update(grads, self.opt_state, self.params)\n params = optax.apply_updates(self.params, updates) # type: ignore\n step = self.step + 1\n return self.replace(\n params=params,\n opt_state=opt_state,\n step=step,\n **kwargs,\n )\n\n\ndef has_keyword_arg(func: tp.Callable[..., tp.Any], name: str) -> bool:\n \"\"\"Return True if func has keyword-only arguments with the given name.\"\"\"\n return any(\n param.name == name\n and param.kind in (param.KEYWORD_ONLY, param.POSITIONAL_OR_KEYWORD)\n for param in inspect.signature(func).parameters.values()\n )\n" }, { "path": "flax/nnx/ids.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"UUIDs for Flax internals.\"\"\"\n\nimport threading\n\n\nclass UUIDManager:\n \"\"\"Globally unique counter-based id manager.\n\n We need globally unique key ids for Module and Variable object instances\n to preserve and recreate sharing-by-reference relationship when lifting\n transforms and adopting outside Modules.\n - Use of id() is unacceptable because these identifiers are literally\n pointers which can be recycled, so we rely on a globally unique counter id\n instead.\n - We need to handle copy/deepcopy uniqueness via a wrapped type.\n \"\"\"\n\n def __init__(self):\n self._lock = threading.Lock()\n self._id = 0\n\n def __call__(self):\n with self._lock:\n self._id += 1\n return UUID(self._id)\n\n\nuuid = UUIDManager()\n\n\nclass UUID:\n \"\"\"Hashable wrapper for ids that handles uniqueness of copies.\"\"\"\n\n def __init__(self, rawid):\n self.id = rawid\n\n def __eq__(self, other):\n return isinstance(other, UUID) and other.id == self.id\n\n def __hash__(self):\n return hash(self.id)\n\n def __repr__(self):\n return f'UUID({self.id})'\n\n def __deepcopy__(self, memo):\n del memo\n return uuid()\n\n def __copy__(self):\n return uuid()\n" }, { "path": "flax/nnx/module.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom __future__ import annotations\n\nimport inspect\nimport typing as tp\n\nimport jax\nimport jax.numpy as jnp\n\nfrom flax.nnx import (\n filterlib,\n graphlib,\n pytreelib,\n)\nfrom flax.nnx import variablelib as variableslib\nfrom flax.nnx.pytreelib import Pytree, PytreeMeta\nfrom flax.nnx.graphlib import GraphState\nfrom flax.nnx.statelib import split_state, State\nimport functools as ft\nfrom flax.typing import Key, Path, PathParts\nfrom collections.abc import MutableMapping\nimport warnings\n\nA = tp.TypeVar('A')\nB = tp.TypeVar('B')\nM = tp.TypeVar('M', bound='Module')\nS = tp.TypeVar('S', bound=tp.Union[GraphState, tuple[GraphState, ...]])\nV = tp.TypeVar('V', bound=variableslib.Variable[tp.Any])\nF = tp.TypeVar('F', bound=tp.Callable[..., tp.Any])\n\nStateMapping = tp.Mapping[Path, tp.Any]\ntuple_reduce = lambda xs, x: xs + (x,)\ntuple_init = lambda: ()\n\n\nclass ModuleMeta(PytreeMeta):\n # we keep a trivial derived class just in case we need to\n # add more functionality in the future\n pass\n\n\nclass Module(Pytree, metaclass=ModuleMeta):\n \"\"\"Base class for all neural network modules.\n\n Layers and models should subclass this class.\n\n ``Module``'s can contain submodules, and in this way can be nested in a tree\n structure. Submodules can be assigned as regular attributes inside the\n ``__init__`` method.\n\n You can define arbitrary \"forward pass\" methods on your ``Module`` subclass.\n While no methods are special-cased, ``__call__`` is a popular choice since\n you can call the ``Module`` directly::\n\n >>> from flax import nnx\n >>> import jax.numpy as jnp\n\n >>> class Model(nnx.Module):\n ... def __init__(self, rngs):\n ... self.linear1 = nnx.Linear(2, 3, rngs=rngs)\n ... self.linear2 = nnx.Linear(3, 4, rngs=rngs)\n ... def __call__(self, x):\n ... x = self.linear1(x)\n ... x = nnx.relu(x)\n ... x = self.linear2(x)\n ... return x\n\n >>> x = jnp.ones((1, 2))\n >>> model = Model(rngs=nnx.Rngs(0))\n >>> y = model(x)\n \"\"\"\n\n def sow(\n self,\n variable_type: type[variableslib.Variable[B]] | str,\n name: str,\n value: A,\n reduce_fn: tp.Callable[[B, A], B] = tuple_reduce,\n init_fn: tp.Callable[[], B] = tuple_init, # type: ignore\n ) -> bool:\n \"\"\"Store intermediate values during module execution for later extraction.\n\n Used with :func:`nnx.capture` decorator to collect intermediate values without\n explicitly passing containers through module calls. Values are stored under\n the specified ``name`` in a collection associated with ``variable_type``.\n\n By default, values are appended to a tuple, allowing multiple values to be\n tracked when the same module is called multiple times.\n\n Example usage::\n\n >>> from flax import nnx\n >>> import jax.numpy as jnp\n\n >>> class Model(nnx.Module):\n ... def __init__(self, rngs):\n ... self.linear1 = nnx.Linear(2, 3, rngs=rngs)\n ... self.linear2 = nnx.Linear(3, 4, rngs=rngs)\n ... def __call__(self, x):\n ... x = self.linear1(x)\n ... self.sow(nnx.Intermediate, 'features', x)\n ... x = self.linear2(x)\n ... return x\n\n >>> # With the capture decorator, sow returns intermediates\n >>> model = Model(rngs=nnx.Rngs(0))\n >>> @nnx.capture(nnx.Intermediate)\n ... def forward(model, x):\n ... return model(x)\n >>> result, intermediates = forward(model, jnp.ones(2))\n >>> assert 'features' in intermediates\n\n Custom init/reduce functions can be passed to control accumulation::\n\n >>> class Model(nnx.Module):\n ... def __init__(self, rngs):\n ... self.linear = nnx.Linear(2, 3, rngs=rngs)\n ... def __call__(self, x):\n ... x = self.linear(x)\n ... self.sow(nnx.Intermediate, 'sum', x,\n ... init_fn=lambda: 0,\n ... reduce_fn=lambda prev, curr: prev+curr)\n ... return x\n\n Args:\n variable_type: The :class:`Variable` type for the stored value.\n Typically :class:`Intermediate` or a subclass is used.\n name: A string key for storing the value in the collection.\n value: The value to be stored.\n reduce_fn: Function to combine existing and new values. Default appends\n to a tuple.\n init_fn: Function providing initial value for first ``reduce_fn`` call.\n Default is an empty tuple.\n \"\"\"\n if isinstance(variable_type, str):\n variable_type = variableslib.variable_type_from_name(\n variable_type, allow_register=True\n )\n\n if hasattr(self, '__captures__'):\n for var in self.__captures__:\n if type(var) == variable_type:\n if name in var:\n var[name] = reduce_fn(var[name], value)\n else:\n var[name] = reduce_fn(init_fn(), value)\n return True\n else:\n return False\n elif hasattr(self, name):\n variable = getattr(self, name)\n if not isinstance(variable, variableslib.Variable):\n raise ValueError(\n f\"Expected '{name}' to be a Variable, got {type(variable).__name__}\"\n )\n elif type(variable) != variable_type:\n raise ValueError(\n f\"Expected '{name}' to be of type '{variable_type.__name__}', \"\n f\"got '{type(variable).__name__}'\"\n )\n variable.set_value(reduce_fn(variable.get_value(), value))\n else:\n reduced_value = reduce_fn(init_fn(), value)\n setattr(self, name, variable_type(reduced_value))\n warnings.warn(\n \"\"\"Using 'Module.sow()' outside of 'nnx.capture()' is deprecated; see\n https://flax.readthedocs.io/en/stable/capturing_intermediates.html for more information.\n \"\"\",\n DeprecationWarning,\n stacklevel=2,\n )\n return True\n\n def perturb(\n self,\n name: str,\n value: tp.Any,\n variable_type: (\n str | type[variableslib.Variable[tp.Any]]\n ) = variableslib.Perturbation,\n ):\n \"\"\"Extract gradients of intermediate values during training.\n\n Used with :func:`nnx.capture` to record intermediate values in the forward pass\n and their gradients in the backward pass. Returns the value plus whatever perturbation\n is stored under ``name`` in the current capture context, allowing gradient computation via ``nnx.grad``.\n\n The workflow has four steps:\n 1. Initialize perturbations with ``nnx.capture(model, nnx.Perturbation)``\n 2. Run model with ``nnx.capture(model, nnx.Intermediate, init=perturbations)``\n 3. Take gradients with respect to perturbations using ``nnx.grad``\n 4. Combine results with ``nnx.merge_state(perturb_grads, intermediates)``\n\n .. note::\n This creates extra variables of the same size as ``value``, thus\n occupies more memory. Use it only to debug gradients in training.\n\n Example usage::\n\n >>> from flax import nnx\n >>> import jax.numpy as jnp\n\n >>> class Model(nnx.Module):\n ... def __call__(self, x):\n ... x2 = self.perturb('grad_of_x', x)\n ... return 3 * x2\n\n >>> model = Model()\n >>> x = 1.0\n\n >>> # Step 1: Initialize perturbations\n >>> forward = nnx.capture(model, nnx.Perturbation)\n >>> _, perturbations = forward(x)\n\n >>> # Steps 2-4: Capture gradients\n >>> def train_step(model, perturbations, x):\n ... def loss(model, perturbations, x):\n ... return nnx.capture(model, nnx.Intermediate, init=perturbations)(x)\n ... (grads, perturb_grads), sowed = nnx.grad(loss, argnums=(0, 1), has_aux=True)(model, perturbations, x)\n ... return nnx.merge_state(perturb_grads, sowed)\n\n >>> metrics = train_step(model, perturbations, x)\n >>> # metrics contains gradients of intermediate values\n\n Args:\n name: A string key for storing the perturbation value.\n value: The intermediate value to capture gradients for. You must use\n the returned value (not the original) for gradient capturing to work.\n variable_type: The :class:`Variable` type for the stored perturbation.\n Default is :class:`nnx.Perturbation`.\n \"\"\"\n if isinstance(variable_type, str):\n variable_type = variableslib.variable_type_from_name(\n variable_type, allow_register=True\n )\n\n if hasattr(self, '__captures__'):\n for var in self.__captures__:\n if type(var) == variable_type:\n if name not in var:\n zeros = jax.tree.map(jnp.zeros_like, value)\n var[name] = zeros\n old_value = var[name]\n return old_value + value\n else:\n return value\n elif hasattr(self, name):\n var = getattr(self, name)\n if not isinstance(var, variable_type):\n raise ValueError(\n f\"Expected '{name}' to be of type '{variable_type.__name__}', \"\n f\"got '{type(var).__name__}'\"\n )\n old_value = var.get_value()\n else:\n old_value = jax.tree.map(jnp.zeros_like, value)\n setattr(self, name, variable_type(old_value))\n warnings.warn(\"\"\"\n Using 'Module.perturb()' outside of 'nnx.capture()' is deprecated; see\n https://flax.readthedocs.io/en/stable/capturing_intermediates.html for more information.\n \"\"\",\n DeprecationWarning,\n stacklevel=2)\n return old_value + value\n\n def iter_modules(self) -> tp.Iterator[tuple[PathParts, Module]]:\n \"\"\"\n Warning: this method is method is deprecated; use :func:`iter_modules` instead.\n\n Recursively iterates over all nested :class:`Module`'s of the current Module, including\n the current Module. Alias of :func:`iter_modules`.\n \"\"\"\n warnings.warn(\n \"The 'm.iter_modules()' method is deprecated; use the 'nnx.iter_modules(m)' function instead.\",\n DeprecationWarning,\n stacklevel=2,\n )\n yield from iter_modules(self)\n\n def iter_children(self) -> tp.Iterator[tuple[Key, Module]]:\n \"\"\"\n Warning: this method is method is deprecated; use :func:`iter_children` instead.\n\n Iterates over all children :class:`Module`'s of the current Module. This\n method is similar to :func:`iter_modules`, except it only iterates over the\n immediate children, and does not recurse further down. Alias of :func:`iter_children`.\n \"\"\"\n warnings.warn(\n \"The 'm.iter_children()' method is deprecated; use the 'nnx.iter_children(m)' function instead.\",\n DeprecationWarning,\n stacklevel=2,\n )\n yield from iter_children(self)\n\n def set_attributes(\n self,\n *filters: filterlib.Filter,\n raise_if_not_found: bool = True,\n graph: bool | None = None,\n **attributes: tp.Any,\n ) -> None:\n \"\"\"Sets the attributes of nested Modules including the current Module.\n If the attribute is not found in the Module, it is ignored.\n\n Example::\n\n >>> from flax import nnx\n ...\n >>> class Block(nnx.Module):\n ... def __init__(self, din, dout, *, rngs: nnx.Rngs):\n ... self.linear = nnx.Linear(din, dout, rngs=rngs)\n ... self.dropout = nnx.Dropout(0.5, deterministic=False)\n ... self.batch_norm = nnx.BatchNorm(10, use_running_average=False, rngs=rngs)\n ...\n >>> block = Block(2, 5, rngs=nnx.Rngs(0))\n >>> block.dropout.deterministic, block.batch_norm.use_running_average\n (False, False)\n >>> block.set_attributes(deterministic=True, use_running_average=True)\n >>> block.dropout.deterministic, block.batch_norm.use_running_average\n (True, True)\n\n ``Filter``'s can be used to set the attributes of specific Modules::\n\n >>> block = Block(2, 5, rngs=nnx.Rngs(0))\n >>> block.set_attributes(nnx.Dropout, deterministic=True)\n >>> # Only the dropout will be modified\n >>> block.dropout.deterministic, block.batch_norm.use_running_average\n (True, False)\n\n Args:\n *filters: Filters to select the Modules to set the attributes of.\n raise_if_not_found: If True (default), raises a ValueError if at least one attribute\n instance is not found in one of the selected Modules.\n **attributes: The attributes to set.\n \"\"\"\n remaining_attributes = set(attributes.keys())\n if not filters:\n filters = (True,)\n predicates = tuple(map(filterlib.to_predicate, filters))\n for path, module in iter_modules(self, graph=graph):\n for predicate in predicates:\n if predicate(path, module):\n for name, value in attributes.items():\n if hasattr(module, name):\n if name in remaining_attributes:\n remaining_attributes.remove(name)\n setattr(module, name, value)\n break\n\n if remaining_attributes and raise_if_not_found:\n raise ValueError(\n 'Could not find at least one instance of the following'\n f' attributes: {sorted(remaining_attributes)}'\n )\n\n def train(self, **attributes):\n \"\"\"Sets the Module to training mode.\n\n ``train`` uses ``set_attributes`` to recursively set attributes ``deterministic=False``\n and ``use_running_average=False`` of all nested Modules that have these attributes.\n Its primarily used to control the runtime behavior of the ``Dropout`` and ``BatchNorm``\n Modules.\n\n Example::\n\n >>> from flax import nnx\n ...\n >>> class Block(nnx.Module):\n ... def __init__(self, din, dout, *, rngs: nnx.Rngs):\n ... self.linear = nnx.Linear(din, dout, rngs=rngs)\n ... # initialize Dropout and BatchNorm in eval mode\n ... self.dropout = nnx.Dropout(0.5, deterministic=True)\n ... self.batch_norm = nnx.BatchNorm(10, use_running_average=True, rngs=rngs)\n ...\n >>> block = Block(2, 5, rngs=nnx.Rngs(0))\n >>> block.dropout.deterministic, block.batch_norm.use_running_average\n (True, True)\n >>> block.train()\n >>> block.dropout.deterministic, block.batch_norm.use_running_average\n (False, False)\n\n Args:\n **attributes: additional attributes passed to ``set_attributes``.\n \"\"\"\n return self.set_attributes(\n deterministic=False,\n use_running_average=False,\n **attributes,\n raise_if_not_found=False,\n )\n\n def eval(self, **attributes):\n \"\"\"Sets the Module to evaluation mode.\n\n ``eval`` uses ``set_attributes`` to recursively set attributes ``deterministic=True``\n and ``use_running_average=True`` of all nested Modules that have these attributes.\n Its primarily used to control the runtime behavior of the ``Dropout`` and ``BatchNorm``\n Modules.\n\n Example::\n\n >>> from flax import nnx\n ...\n >>> class Block(nnx.Module):\n ... def __init__(self, din, dout, *, rngs: nnx.Rngs):\n ... self.linear = nnx.Linear(din, dout, rngs=rngs)\n ... self.dropout = nnx.Dropout(0.5)\n ... self.batch_norm = nnx.BatchNorm(10, rngs=rngs)\n ...\n >>> block = Block(2, 5, rngs=nnx.Rngs(0))\n >>> block.dropout.deterministic, block.batch_norm.use_running_average\n (False, False)\n >>> block.eval()\n >>> block.dropout.deterministic, block.batch_norm.use_running_average\n (True, True)\n\n Args:\n **attributes: additional attributes passed to ``set_attributes``.\n \"\"\"\n return self.set_attributes(\n deterministic=True,\n use_running_average=True,\n **attributes,\n raise_if_not_found=False,\n )\n\ndef view(node: A, /, *, only: filterlib.Filter = ..., raise_if_not_found: bool = True, graph: bool | None = None, **kwargs) -> A:\n \"\"\"Creates a new node with static attributes updated according to ``**kwargs``.\n\n The new node contains references to jax arrays in the original node. If a\n kwarg is not found in any module, this method raises a ValueError. Uses the\n ``set_view`` class method in nnx.Modules. ``set_view`` class methods should\n return any unused kwargs.\n\n Example::\n >>> from flax import nnx\n ...\n >>> class Block(nnx.Module):\n ... def __init__(self, din, dout, *, rngs: nnx.Rngs):\n ... self.linear = nnx.Linear(din, dout, rngs=rngs)\n ... self.dropout = nnx.Dropout(0.5, deterministic=False)\n ... self.batch_norm = nnx.BatchNorm(10, use_running_average=False, rngs=rngs)\n ...\n >>> block = Block(2, 5, rngs=nnx.Rngs(0))\n >>> block.dropout.deterministic, block.batch_norm.use_running_average\n (False, False)\n >>> new_block = nnx.view(block, deterministic=True, use_running_average=True)\n >>> new_block.dropout.deterministic, new_block.batch_norm.use_running_average\n (True, True)\n\n ``Filter``'s can be used to set the attributes of specific Modules::\n >>> block = Block(2, 5, rngs=nnx.Rngs(0))\n >>> new_block = nnx.view(block, only=nnx.Dropout, deterministic=True)\n >>> # Only the dropout will be modified\n >>> new_block.dropout.deterministic, new_block.batch_norm.use_running_average\n (True, False)\n\n Args:\n node: the object to create a copy of.\n only: Filters to select the Modules to set the attributes of.\n graph: If ``True`` (default), uses graph-mode which supports the full\n NNX feature set including shared references. If ``False``, uses\n tree-mode which treats Modules as regular JAX pytrees, avoiding\n the overhead of the graph protocol.\n **kwargs: The attributes to set.\n \"\"\"\n predicate = filterlib.to_predicate(only)\n\n remaining = set(kwargs)\n\n def _set_mode_fn(path, node):\n if hasattr(node, 'set_view') and predicate(path, node):\n sig = inspect.signature(node.set_view)\n has_var_keyword = any(\n p.kind == inspect.Parameter.VAR_KEYWORD\n for p in sig.parameters.values()\n )\n if has_var_keyword:\n node.set_view(**kwargs)\n remaining.clear()\n else:\n named_params = {\n name\n for name, p in sig.parameters.items()\n if p.kind\n in (\n inspect.Parameter.POSITIONAL_OR_KEYWORD,\n inspect.Parameter.KEYWORD_ONLY,\n )\n }\n filtered_kwargs = {\n k: v for k, v in kwargs.items() if k in named_params\n }\n node.set_view(**filtered_kwargs)\n remaining.difference_update(named_params)\n return node\n\n out = graphlib.recursive_map(_set_mode_fn, node, graph=graph)\n\n if raise_if_not_found and remaining:\n raise ValueError(f\"Unused keys found in nnx.view: {sorted(remaining)}\")\n\n return out\n\ndef with_attributes(\n node: A,\n /,\n *,\n only: filterlib.Filter = ...,\n raise_if_not_found: bool = True,\n graph: bool | None = None,\n **attributes: tp.Any,\n) -> A:\n \"\"\"Creates a new node with attributes updated according to ``**attributes``.\n\n The new node contains references to jax arrays in the original node. Unlike\n ``set_attributes``, this function does not modify the original node.\n\n Example::\n >>> from flax import nnx\n ...\n >>> class Block(nnx.Module):\n ... def __init__(self, din, dout, *, rngs: nnx.Rngs):\n ... self.linear = nnx.Linear(din, dout, rngs=rngs)\n ... self.dropout = nnx.Dropout(0.5, deterministic=False)\n ... self.batch_norm = nnx.BatchNorm(10, use_running_average=False, rngs=rngs)\n ...\n >>> block = Block(2, 5, rngs=nnx.Rngs(0))\n >>> block.dropout.deterministic, block.batch_norm.use_running_average\n (False, False)\n >>> new_block = nnx.with_attributes(block, deterministic=True, use_running_average=True)\n >>> new_block.dropout.deterministic, new_block.batch_norm.use_running_average\n (True, True)\n >>> block.dropout.deterministic, block.batch_norm.use_running_average\n (False, False)\n\n ``Filter``'s can be used to set the attributes of specific Modules::\n >>> block = Block(2, 5, rngs=nnx.Rngs(0))\n >>> new_block = nnx.with_attributes(block, only=nnx.Dropout, deterministic=True)\n >>> # Only the dropout will be modified\n >>> new_block.dropout.deterministic, new_block.batch_norm.use_running_average\n (True, False)\n\n Args:\n node: the object to create a copy of.\n only: Filters to select the Modules to set the attributes of.\n raise_if_not_found: If True (default), raises a ValueError if at least one\n attribute instance is not found in one of the selected Modules.\n graph: If ``True`` (default), uses graph-mode which supports the full\n NNX feature set including shared references. If ``False``, uses\n tree-mode which treats Modules as regular JAX pytrees, avoiding\n the overhead of the graph protocol.\n **attributes: The attributes to set.\n \"\"\"\n predicate = filterlib.to_predicate(only)\n remaining_attributes = set(attributes.keys())\n\n def _set_attributes_fn(path, node):\n if isinstance(node, Module) and predicate(path, node):\n for name, value in attributes.items():\n if hasattr(node, name):\n setattr(node, name, value)\n remaining_attributes.discard(name)\n return node\n\n out = graphlib.recursive_map(_set_attributes_fn, node, graph=graph)\n\n if remaining_attributes and raise_if_not_found:\n raise ValueError(\n 'Could not find at least one instance of the '\n f'following attributes: {sorted(remaining_attributes)}'\n )\n\n return out\n\ndef _parse_docstring_args(doc_str: str) -> dict[str, str]:\n \"\"\"Parses parameters from `Args:` section of a function docstring.\n Assumes Google style docstrings. Returns a dictionary with\n keys representing argument names and values representing descriptions.\n Each description has lines starting with 4 spaces.\n \"\"\"\n lines = doc_str.split(\"\\n\")\n\n # Get lines with the parameter names\n inds = [i for i, l in enumerate(lines) if l.startswith(\" \") and not l.startswith(\" \")]\n inds.append(len(lines))\n out = dict()\n\n # Parse each argument\n for i in range(len(inds)-1):\n start, end = inds[i], inds[i+1]\n\n # Process first line for the description\n first_colon = lines[start].find(\":\")\n name = lines[start][:first_colon].strip()\n desc = [\" \"*4 + lines[start][first_colon+1:].strip()]\n\n # Append remaining description lines\n for j in range(start+1, end):\n desc.append(lines[j])\n\n out[name] = \"\\n\".join(desc)\n return out\n\n\n\ndef view_info(node: Module, /, *, only: filterlib.Filter = ..., graph: bool | None = None) -> str:\n \"\"\"Provides information about the ``view`` arguments for a module and all\n submodules. If no docstring is provided for a module's `set_view`, this function\n puts the `set_view` signature below the function.\n\n Example::\n >>> from flax import nnx\n ...\n >>> class CustomModel(nnx.Module):\n ... def __init__(self, *, rngs):\n ... self.mha = nnx.MultiHeadAttention(4, 8, 32, rngs=rngs)\n ... self.drop = nnx.Dropout(0.5, rngs=rngs)\n ... self.bn = nnx.BatchNorm(32, rngs=rngs)\n ...\n >>> model = CustomModel(rngs=nnx.Rngs(0))\n >>> print(nnx.view_info(model))\n BatchNorm:\n use_running_average: bool | None = None\n if True, the stored batch statistics will be\n used instead of computing the batch statistics on the input.\n Dropout:\n deterministic: bool | None = None\n if True, disables dropout masking.\n MultiHeadAttention:\n deterministic: bool | None = None\n if True, the module is set to deterministic mode.\n decode: bool | None = None\n if True, the module is set to decode mode.\n batch_size: int | Shape | None = None\n the batch size to use for the cache.\n max_length: int | None = None\n the max length to use for the cache.\n\n Args:\n node: the object to display ``view`` information for.\n only: Filters to select the Modules to display information for.\n graph: If ``True`` (default), uses graph-mode which supports the full\n NNX feature set including shared references. If ``False``, uses\n tree-mode which treats Modules as regular JAX pytrees, avoiding\n the overhead of the graph protocol.\n \"\"\"\n predicate = filterlib.to_predicate(only)\n classes: set[Module] = set()\n\n def _set_mode_info_fn(path, node):\n if hasattr(node, 'set_view') and predicate(path, node):\n classes.add(node.__class__)\n return node\n\n graphlib.recursive_map(_set_mode_info_fn, node, graph=graph)\n\n class_list = sorted(list(classes), key=lambda x: x.__qualname__)\n out_str = []\n for c in class_list:\n out_str.append(f\"{c.__qualname__}:\")\n sig = inspect.signature(c.set_view)\n doc = inspect.getdoc(c.set_view)\n\n # Parse docstring\n if isinstance(doc, str):\n start, end = doc.find(\"Args:\\n\"), doc.find(\"Returns:\\n\")\n if end == -1:\n end = len(doc)\n doc = doc[start+6:end]\n parsed_docstring = _parse_docstring_args(doc)\n\n # Generate output from signature and docstring\n skip_names = {\"self\", \"args\", \"kwargs\"}\n for name, param in sig.parameters.items():\n if name in skip_names:\n continue\n\n if param.default is inspect.Parameter.empty:\n out_str.append(f\" {name}: {param.annotation}\")\n else:\n out_str.append(f\" {name}: {param.annotation} = {param.default}\")\n out_str.append(parsed_docstring[name])\n else:\n out_str.append(f\" set_view{sig}\")\n\n\n return \"\\n\".join(out_str)\n\ndef first_from(*args: tp.Optional[A], error_msg: str) -> A:\n \"\"\"Return the first non-None argument.\n\n If all arguments are None, raise a ValueError with the given error message.\n\n Args:\n *args: the arguments to check\n error_msg: the error message to raise if all arguments are None\n Returns:\n The first non-None argument.\n \"\"\"\n for arg in args:\n if arg is not None:\n return arg\n raise ValueError(error_msg)\n\ndef iter_modules(\n module: Module, /, *, graph: bool | None = None,\n) -> tp.Iterator[tuple[PathParts, Module]]:\n \"\"\"Recursively iterates over all nested :class:`Module`'s of the given Module, including\n the argument.\n\n Specifically, this function creates a generator that yields the path and the Module instance, where\n the path is a tuple of strings or integers representing the path to the Module from the\n root Module.\n\n Example::\n\n >>> from flax import nnx\n ...\n >>> class SubModule(nnx.Module):\n ... def __init__(self, din, dout, rngs):\n ... self.linear1 = nnx.Linear(din, dout, rngs=rngs)\n ... self.linear2 = nnx.Linear(din, dout, rngs=rngs)\n ...\n >>> class Block(nnx.Module):\n ... def __init__(self, din, dout, *, rngs: nnx.Rngs):\n ... self.linear = nnx.Linear(din, dout, rngs=rngs)\n ... self.submodule = SubModule(din, dout, rngs=rngs)\n ... self.dropout = nnx.Dropout(0.5)\n ... self.batch_norm = nnx.BatchNorm(10, rngs=rngs)\n ...\n >>> model = Block(2, 5, rngs=nnx.Rngs(0))\n >>> for path, module in nnx.iter_modules(model):\n ... print(path, type(module).__name__)\n ...\n ('batch_norm',) BatchNorm\n ('dropout',) Dropout\n ('linear',) Linear\n ('submodule', 'linear1') Linear\n ('submodule', 'linear2') Linear\n ('submodule',) SubModule\n () Block\n\n Args:\n module: A :class:`Module` object.\n graph: If ``True`` (default), uses graph-mode which supports the full\n NNX feature set including shared references. If ``False``, uses\n tree-mode which treats Modules as regular JAX pytrees, avoiding\n the overhead of the graph protocol.\n \"\"\"\n for path, value in graphlib.iter_graph(module, graph=graph):\n if isinstance(value, Module):\n yield path, value\n\niter_children = graphlib.iter_children\n\nP = tp.ParamSpec(\"P\")\nR = tp.TypeVar(\"R\")\n\n@tp.overload\ndef capture(\n fn: tp.Callable[P, R],\n *var_types: type[variableslib.Variable],\n init: tp.Optional[State] = None,\n method_outputs: tp.Optional[type[variableslib.Variable]] = None\n) -> tp.Callable[P, tuple[R, State]]: ...\n\n@tp.overload\ndef capture(\n fn: type[variableslib.Variable],\n *var_types: type[variableslib.Variable],\n init: tp.Optional[State] = None,\n method_outputs: tp.Optional[type[variableslib.Variable]] = None\n) -> tp.Callable[[tp.Callable[P, R]], tp.Callable[P, tuple[R, State]]]: ...\n\ndef capture(fn: tp.Callable[P, R] | type[variableslib.Variable], *var_types: type[variableslib.Variable],\n init : tp.Optional[State] = None,\n method_outputs : tp.Optional[type[variableslib.Variable]] = None\n) -> tp.Callable[P, tuple[R, State]] | tp.Callable[[tp.Callable[P, R]], tp.Callable[P, tuple[R, State]]]:\n \"\"\"Wraps a function to capture intermediate values from a module during execution.\n\n This function wraps a `Callable`, executing it while collecting intermediate values that were stored using\n ``Module.sow()`` or ``Module.perturb()``.\n\n The `fn` argument can be either a function, a Module instance, or a bound method.\n If `fn` is a function, its first argument should be the module in which intermediate values are to be recorded.\n If `fn` is a bound method, the module used for storage is inferred from the instance.\n If `fn` is a Module, its `__call__` method will be wrapped.\n\n Args:\n fn: The `Callable` to wrap.\n var_types: Variable types to capture. If None, defaults to [].\n init: MutableMapping used to initialize perturbation values. This is useful for gradient extraction.\n method_outputs: If provided, automatically sows the output of each method\n in the module and its submodules using this variable type.\n\n Returns:\n A wrapped function that returns\n a tuple of (result, *intermediates) where result is the output of the function\n and each intermediate is a State containing the captured values with the corresponding type in `var_types`.\n\n Example with manual sowing::\n\n class Foo(nnx.Module):\n def __call__(self, x):\n self.sow(nnx.Intermediate, 'features', x)\n return x\n\n model = Foo(rngs=nnx.Rngs(0))\n forward = nnx.capture(model, nnx.Intermediate)\n result, intermediates = forward(x)\n # intermediates['features'] contains the sowed value\n\n Example with method outputs::\n\n class Foo(nnx.Module):\n def features(self, x):\n return x\n def classifier(self, x):\n return x\n def __call__(self, x):\n return self.classifier(self.features(x))\n\n model = Foo(rngs=nnx.Rngs(0))\n result, intermediates = nnx.capture(\n model, method_output_type=nnx.Intermediate)(x)\n # intermediates contains outputs of features(), classifier(), and __call__()\n\n Example with gradient extraction::\n\n class Model(nnx.Module):\n def __call__(self, x):\n x2 = self.perturb('grad_of_x', x)\n return 3 * x2\n\n model = Model()\n forward = nnx.capture(lambda model, x: model(x), nnx.Perturbation) # Initialize perturbations\n _, perturbations = forward_capture(model, x)\n\n # Compute gradients with respect to perturbations\n loss = nnx.capture(forward, init=perturbations)\n grads, sowed = nnx.grad(loss, has_aux=True)(model, perturbations, x)\n \"\"\"\n\n # Handle partial evaluation when first arg is a Variable type\n if isinstance(fn, type) and issubclass(fn, variableslib.Variable):\n # Partial application: return a function that waits for the actual fn\n all_var_types = (fn,) + var_types\n def partial_capture(actual_fn: tp.Callable[P, R] | Module) -> tp.Callable[P, tuple[R, State]]:\n return capture(actual_fn, *all_var_types, init=init, method_outputs=method_outputs)\n return partial_capture\n\n # Handle bound methods and callable Modules\n module_instance = None\n if inspect.ismethod(fn) and isinstance(fn.__self__, Module):\n module_instance = fn.__self__\n elif isinstance(fn, Module):\n module_instance = fn\n\n ft.wraps(fn)\n def wrapper(*fn_args, **kwargs):\n if module_instance is None:\n module = fn_args[0]\n else:\n module = module_instance\n\n # Extract initial values from state\n state_by_path = _collect_state_by_path(init) if init else {}\n\n # Initialize __captures__ as a tuple of Variables (one per type)\n for path, m in iter_modules(module):\n # Create initial dicts for each variable type\n initial_dicts = {}\n for var_type in var_types:\n initial_dicts[var_type] = {}\n\n # Populate from state if available\n if path in state_by_path:\n for name, var in state_by_path[path].items():\n var_type = type(var)\n if var_type not in initial_dicts:\n initial_dicts[var_type] = {}\n initial_dicts[var_type][name] = var.get_value()\n\n # Create the captures tuple\n captures_tuple = tuple(k(v) for (k,v) in initial_dicts.items())\n m.__captures__ = pytreelib.data(captures_tuple)\n\n # Wrap methods with capturing if required\n if method_outputs:\n for _, m in iter_modules(module):\n _add_capturing(type(m), method_outputs)\n\n try:\n result = fn(*fn_args, **kwargs)\n finally:\n\n # Undo method sowing modification\n for _, m in iter_modules(module):\n _remove_capturing(type(m))\n\n # Extract intermediates manually from __captures__\n interms = State({})\n _extract_captures(module, interms, set(var_types))\n if len(var_types) == 0:\n return result\n split_states = split_state(interms, *var_types)\n if len(var_types) == 1:\n return result, split_states\n else:\n return (result, *split_states)\n\n return wrapper\n\ndef _collect_state_by_path(state):\n \"\"\"Build a mapping from module path to state Variables.\"\"\"\n state_by_path = {}\n\n def collect(s, path_parts):\n if isinstance(s, MutableMapping):\n for key, value in s.items():\n if isinstance(value, variableslib.Variable):\n path_tuple = tuple(path_parts)\n if path_tuple not in state_by_path:\n state_by_path[path_tuple] = {}\n state_by_path[path_tuple][key] = value\n elif isinstance(value, MutableMapping):\n collect(value, path_parts + [key])\n\n collect(state, [])\n return state_by_path\n\ndef _navigate_to_path(state, path):\n \"\"\"Navigate to a nested path in state, creating dicts as needed.\"\"\"\n current = state\n for part in path:\n if part not in current:\n current[part] = State({})\n current = current[part]\n return current\n\ndef _extract_captures(module, state, var_types):\n \"\"\"Extract intermediates from __captures__ tuple into state dict.\"\"\"\n for path, mod in iter_modules(module):\n if hasattr(mod, '__captures__'):\n captures_tuple = mod.__captures__\n for var in captures_tuple:\n if not type(var) in var_types:\n continue\n current = _navigate_to_path(state, path)\n for key, value in var.items():\n current[key] = type(var)(value)\n delattr(mod, '__captures__')\n\n\ndef _add_capturing(cls, variable_type):\n \"\"\"Adds capturing to methods of a Module.\n Does not instrument superclass methods.\"\"\"\n for name, method in cls.__dict__.items():\n if callable(method) and (not name.startswith('_') or name == '__call__'):\n if not hasattr(method, '_does_capturing'):\n def closure(name, method): # Necessary to make 'name' immutable during iteration\n @ft.wraps(method)\n def wrapper(self, *args, **kwargs):\n result = method(self, *args, **kwargs)\n self.sow(variable_type, name, result)\n return result\n wrapper._does_capturing = True\n setattr(cls, name, wrapper)\n closure(name, method)\n return cls\n\ndef _remove_capturing(cls):\n \"\"\"Remove capturing methods from a Module.\"\"\"\n for name, method in cls.__dict__.items():\n if hasattr(method, '_does_capturing'):\n setattr(cls, name, method.__wrapped__)\n return cls\n" }, { "path": "flax/nnx/nn/__init__.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n" }, { "path": "flax/nnx/nn/activations.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport typing as tp\nfrom types import MappingProxyType\nfrom jax.nn import (\n celu,\n elu,\n gelu,\n glu,\n hard_sigmoid,\n hard_silu,\n hard_swish,\n hard_tanh,\n leaky_relu,\n log_sigmoid,\n log_softmax,\n logsumexp,\n one_hot,\n relu,\n identity,\n relu6,\n selu,\n sigmoid,\n silu,\n soft_sign,\n softmax,\n softplus,\n standardize,\n swish,\n)\nimport jax.numpy as jnp\nfrom jax.numpy import tanh\n\nfrom flax import nnx\nfrom flax.nnx.nn import dtypes\nfrom flax.typing import Array, Dtype, PromoteDtypeFn\n\n\n__all__ = [\n 'celu',\n 'elu',\n 'gelu',\n 'glu',\n 'hard_sigmoid',\n 'hard_silu',\n 'hard_swish',\n 'hard_tanh',\n 'leaky_relu',\n 'log_sigmoid',\n 'log_softmax',\n 'logsumexp',\n 'one_hot',\n 'relu',\n 'identity',\n 'relu6',\n 'selu',\n 'sigmoid',\n 'silu',\n 'soft_sign',\n 'softmax',\n 'softplus',\n 'standardize',\n 'swish',\n 'tanh',\n 'PReLU',\n]\n\n\nclass PReLU(nnx.Module):\n \"\"\"Parametric Rectified Linear Unit (PReLU) activation function.\n\n Note that PReLU is a Flax layer and not a simple activation function, so\n it needs to be initialized before being called.\n\n Example::\n\n >>> import flax.nnx as nnx\n\n >>> class MLP(nnx.Module):\n ... def __init__(self):\n ... self.linear = nnx.Linear(3, 2)\n ... self.act = nnx.PReLU(negative_slope_init=0.1)\n ...\n ... def __call__(self, x):\n ... x = self.linear(x)\n ... x = self.act(x)\n ... return x\n\n Args:\n negative_slope_init: the value to initialize the negative slope (default 0.01).\n dtype: the dtype of the computation (default: infer from input and params).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n promote_dtype: function to promote the dtype of all input array arguments\n (including Variables accessed through ``self``) to the desired dtype. The\n function should accept a tuple of ``(inputs, negative_slope)`` and a ``dtype``\n keyword argument, and return a tuple of arrays with the promoted dtype.\n negative_slope_metadata: Optional metadata dictionary to set when initializing\n the negative slope.\n \"\"\"\n def __init__(\n self,\n negative_slope_init: float = 0.01,\n *,\n dtype: Dtype | None = None,\n param_dtype: Dtype = jnp.float32,\n promote_dtype: PromoteDtypeFn = dtypes.promote_dtype,\n negative_slope_metadata: tp.Mapping[str, tp.Any] = MappingProxyType({}),\n ):\n self.negative_slope = nnx.Param(\n jnp.asarray(negative_slope_init, dtype=param_dtype), **negative_slope_metadata\n )\n self.dtype = dtype\n self.param_dtype = param_dtype\n self.promote_dtype = promote_dtype\n\n def __call__(self, inputs: Array) -> Array:\n negative_slope = self.negative_slope[...]\n if self.dtype is not None:\n inputs, negative_slope = self.promote_dtype(\n (inputs, negative_slope), dtype=self.dtype\n )\n else:\n # Match Linen behavior: cast parameter to input dtype\n negative_slope = jnp.asarray(negative_slope, inputs.dtype)\n\n return jnp.where(\n inputs >= 0,\n inputs,\n negative_slope * inputs,\n )\n" }, { "path": "flax/nnx/nn/attention.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Attention core modules for Flax.\"\"\"\n\nfrom __future__ import annotations\n\nimport functools\nfrom typing import Any\nfrom collections.abc import Mapping\nfrom types import MappingProxyType\nfrom collections.abc import Callable\nimport math\n\nimport jax\nimport jax.numpy as jnp\nfrom jax import lax, random\n\nfrom flax import nnx\nfrom flax.nnx import rnglib\nfrom flax.nnx.module import Module, first_from\nfrom flax.nnx.nn import initializers\nfrom flax.nnx.nn import dtypes\nfrom flax.nnx.nn.linear import (\n LinearGeneral,\n default_kernel_init,\n)\nfrom flax.nnx.nn.normalization import LayerNorm\nfrom flax.typing import (\n Dtype,\n PromoteDtypeFn,\n Shape,\n Initializer,\n PrecisionLike,\n DotGeneralT,\n)\n\nArray = jax.Array\n\n\ndef dot_product_attention_weights(\n query: Array,\n key: Array,\n bias: Array | None = None,\n mask: Array | None = None,\n broadcast_dropout: bool = True,\n dropout_rng: Array | None = None,\n dropout_rate: float = 0.0,\n deterministic: bool = False,\n dtype: Dtype | None = None,\n precision: PrecisionLike = None,\n module: Module | None = None,\n promote_dtype: PromoteDtypeFn = dtypes.promote_dtype,\n is_causal: bool = False,\n):\n \"\"\"Computes dot-product attention weights given query and key.\n\n Used by :func:`dot_product_attention`, which is what you'll most likely use.\n But if you want access to the attention weights for introspection, then\n you can directly call this function and call einsum yourself.\n\n Args:\n query: queries for calculating attention with shape of `[batch..., q_length,\n num_heads, qk_depth_per_head]`.\n key: keys for calculating attention with shape of `[batch..., kv_length,\n num_heads, qk_depth_per_head]`.\n bias: bias for the attention weights. This should be broadcastable to the\n shape `[batch..., num_heads, q_length, kv_length]`. This can be used for\n incorporating causal masks, padding masks, proximity bias, etc.\n mask: mask for the attention weights. This should be broadcastable to the\n shape `[batch..., num_heads, q_length, kv_length]`. This can be used for\n incorporating causal masks. Attention weights are masked out if their\n corresponding mask value is `False`.\n broadcast_dropout: bool: use a broadcasted dropout along batch dims.\n dropout_rng: JAX PRNGKey: to be used for dropout\n dropout_rate: dropout rate\n deterministic: bool, deterministic or not (to apply dropout)\n dtype: the dtype of the computation (default: infer from inputs and params)\n precision: numerical precision of the computation see `jax.lax.Precision`\n for details.\n module: the Module that will sow the attention weights into the\n ``nnx.Intermediate`` collection. If ``module`` is None, the attention\n weights will not be sowed.\n promote_dtype: function to promote the dtype of the arrays to the desired\n dtype. The function should accept a tuple of ``(query, key)`` and a ``dtype``\n keyword argument, and return a tuple of arrays with the promoted dtype.\n is_causal: If true, causal attention will be applied. Note, some\n implementations like xla will generate a mask tensor and apply it to\n the logits to mask out the non-causal parts of the attention matrix,\n but other implementations like cudnn will avoid computing the\n non-causal regions, providing speedups.\n\n Returns:\n Output of shape `[batch..., num_heads, q_length, kv_length]`.\n \"\"\"\n query, key = promote_dtype((query, key), dtype=dtype) # type: ignore[bad-unpacking]\n dtype = query.dtype\n\n assert query.ndim == key.ndim, 'q, k must have same rank.'\n assert query.shape[:-3] == key.shape[:-3], 'q, k batch dims must match.'\n assert query.shape[-1] == key.shape[-1], 'q, k depths must match.'\n\n # check if we need to broadcast Key heads to match Query heads\n is_gqa = False\n if query.shape[-2] != key.shape[-2]:\n q_heads = query.shape[-2]\n k_heads = key.shape[-2]\n\n if q_heads % k_heads != 0:\n raise ValueError(\n f\"Query heads ({q_heads}) must be multiple of \"\n f\"Key heads ({k_heads}) for Grouped Query Attention.\"\n )\n\n n_rep = q_heads // k_heads\n is_gqa = True\n # Reshape Query: [..., Q, H_k * n_rep, D] -> [..., Q, H_k, n_rep, D]\n query = query.reshape(query.shape[:-2] + (k_heads, n_rep, query.shape[-1]))\n # Expand Key: [..., K, H_k, D] -> [..., K, H_k, 1, D]\n key = jnp.expand_dims(key, axis=-2)\n\n # Contract: q(h)gd, k(h)1d -> hgqk (h=H_k, g=n_rep)\n einsum_str = '...qhgd,...kh1d->...hgqk'\n else:\n q_heads = query.shape[-2]\n einsum_str = '...qhd,...khd->...hqk'\n assert query.shape[-2] == key.shape[-2], 'q, k num_heads must match.'\n\n # calculate attention matrix\n depth = query.shape[-1]\n query = query / jnp.sqrt(depth).astype(dtype)\n\n # attn weight shape is (batch..., num_heads, q_length, kv_length)\n attn_weights = jnp.einsum(einsum_str, query, key, precision=precision)\n\n if is_gqa:\n attn_weights = attn_weights.reshape(attn_weights.shape[:-4] + (q_heads, attn_weights.shape[-2], attn_weights.shape[-1]))\n\n # apply attention bias: masking, dropout, proximity bias, etc.\n if bias is not None:\n attn_weights = attn_weights + bias\n # apply attention mask\n if mask is not None or is_causal:\n big_neg = jnp.finfo(dtype).min\n masks = [m for m in [mask] if m is not None]\n if is_causal:\n T, S = attn_weights.shape[-2:]\n causal_mask = jnp.tril(jnp.ones((T, S), dtype=dtype))\n target_shape = mask.shape if mask is not None else attn_weights.shape\n masks.append(jnp.broadcast_to(causal_mask, target_shape))\n combined_mask = combine_masks(*masks, dtype=dtype)\n assert combined_mask is not None\n attn_weights = jnp.where(combined_mask, attn_weights, big_neg)\n\n # normalize the attention weights\n attn_weights = jax.nn.softmax(attn_weights).astype(dtype)\n\n if module:\n module.sow(nnx.Intermediate, 'attention_weights', attn_weights)\n\n # apply attention dropout\n if not deterministic and dropout_rate > 0.0:\n keep_prob = 1.0 - dropout_rate\n # use original key.ndim because we might have expanded key dim\n ndim_base = key.ndim - 1 if is_gqa else key.ndim\n\n if broadcast_dropout:\n # dropout is broadcast across the batch + head dimensions\n dropout_shape = tuple([1] * (ndim_base - 2)) + attn_weights.shape[-2:]\n keep = random.bernoulli(dropout_rng, keep_prob, dropout_shape) # type: ignore\n else:\n keep = random.bernoulli(dropout_rng, keep_prob, attn_weights.shape) # type: ignore\n multiplier = keep.astype(dtype) / jnp.asarray(keep_prob, dtype=dtype)\n attn_weights = attn_weights * multiplier\n\n return attn_weights\n\n\ndef dot_product_attention(\n query: Array,\n key: Array,\n value: Array,\n bias: Array | None = None,\n mask: Array | None = None,\n broadcast_dropout: bool = True,\n dropout_rng: Array | None = None,\n dropout_rate: float = 0.0,\n deterministic: bool = False,\n dtype: Dtype | None = None,\n precision: PrecisionLike = None,\n module: Module | None = None,\n promote_dtype: PromoteDtypeFn = dtypes.promote_dtype,\n is_causal: bool = False,\n):\n \"\"\"Computes dot-product attention given query, key, and value.\n\n This is the core function for applying attention based on\n https://arxiv.org/abs/1706.03762. It calculates the attention weights given\n query and key and combines the values using the attention weights.\n\n Will use the more optimized `jax.nn.dot_product_attention` if dropout is\n not activated and `module=None`.\n\n .. note::\n ``query``, ``key``, ``value`` needn't have any batch dimensions.\n\n Args:\n query: queries for calculating attention with shape of ``[batch..., q_length,\n num_heads, qk_depth_per_head]``.\n key: keys for calculating attention with shape of ``[batch..., kv_length,\n num_heads, qk_depth_per_head]``.\n value: values to be used in attention with shape of ``[batch..., kv_length,\n num_heads, v_depth_per_head]``.\n bias: bias for the attention weights. This should be broadcastable to the\n shape `[batch..., num_heads, q_length, kv_length]`. This can be used for\n incorporating causal masks, padding masks, proximity bias, etc.\n mask: mask for the attention weights. This should be broadcastable to the\n shape `[batch..., num_heads, q_length, kv_length]`. This can be used for\n incorporating causal masks. Attention weights are masked out if their\n corresponding mask value is `False`.\n broadcast_dropout: bool: use a broadcasted dropout along batch dims.\n dropout_rng: JAX PRNGKey: to be used for dropout\n dropout_rate: dropout rate\n deterministic: bool, deterministic or not (to apply dropout)\n dtype: the dtype of the computation (default: infer from inputs)\n precision: numerical precision of the computation see `jax.lax.Precision`\n for details.\n module: the Module that will sow the attention weights into the\n ``nnx.Intermediate`` collection. If ``module`` is None, the attention\n weights will not be sowed.\n promote_dtype: function to promote the dtype of the arrays to the desired\n dtype. The function should accept a tuple of ``(query, key, value)`` and a\n ``dtype`` keyword argument, and return a tuple of arrays with the promoted\n dtype.\n is_causal: If true, causal attention will be applied. Note, some\n implementations like xla will generate a mask tensor and apply it to\n the logits to mask out the non-causal parts of the attention matrix,\n but other implementations like cudnn will avoid computing the\n non-causal regions, providing speedups.\n\n Returns:\n Output of shape `[batch..., q_length, num_heads, v_depth_per_head]`.\n \"\"\"\n query, key, value = promote_dtype((query, key, value), dtype=dtype) # type: ignore[bad-unpacking]\n dtype = query.dtype\n\n assert key.ndim == query.ndim == value.ndim, 'q, k, v must have same rank.'\n assert (\n query.shape[:-3] == key.shape[:-3] == value.shape[:-3]\n ), 'q, k, v batch dims must match.'\n assert key.shape[-3] == value.shape[-3], 'k, v lengths must match.'\n\n # Criteria that invoke the more optimized dot product attention\n if dropout_rate == 0.0 and module is None:\n # make sure qkv batch are compressed to one dim\n query_shape = query.shape\n if len(query_shape) > 4:\n def reshape_4d(x):\n return jnp.reshape(x, (math.prod(x.shape[:-3]), *x.shape[-3:]))\n query, key, value, bias, mask = jax.tree.map(\n reshape_4d, (query, key, value, bias, mask))\n if mask is not None:\n mask = mask.astype(jnp.bool)\n out = jax.nn.dot_product_attention(query, key, value, bias, mask, is_causal=is_causal)\n if len(query_shape) > 4:\n out = jnp.reshape(out, query_shape)\n return out\n\n # compute attention weights\n attn_weights = dot_product_attention_weights(\n query,\n key,\n bias,\n mask,\n broadcast_dropout,\n dropout_rng,\n dropout_rate,\n deterministic,\n dtype,\n precision,\n module,\n promote_dtype,\n is_causal,\n )\n\n # return weighted sum over values for each query position\n # check if need to broadcast Value heads to match Query heads (GQA)\n if attn_weights.shape[-3] != value.shape[-2]:\n q_heads = attn_weights.shape[-3]\n v_heads = value.shape[-2]\n if q_heads % v_heads != 0:\n raise ValueError(f\"Query heads ({q_heads}) must be multiple of Value heads ({v_heads})\")\n\n n_rep = q_heads // v_heads\n # Reshape weights: [..., H_v, n_rep, Q, K]\n attn_weights = attn_weights.reshape(attn_weights.shape[:-3] + (v_heads, n_rep) + attn_weights.shape[-2:])\n # Expand Value: [..., K, H_v, 1, D]\n value = jnp.expand_dims(value, axis=-2)\n # Contract: hgqk, kh1d -> qhgd (h=H_v, g=n_rep)\n out = jnp.einsum('...hgqk,...kh1d->...qhgd', attn_weights, value, precision=precision)\n # Flatten: [..., Q, H_q, D]\n out = out.reshape(out.shape[:-3] + (q_heads, out.shape[-1]))\n else:\n out = jnp.einsum(\n '...hqk,...khd->...qhd', attn_weights, value, precision=precision\n )\n\n return out\n\n\nclass MultiHeadAttention(Module):\n \"\"\"Multi-head attention.\n\n Example usage::\n\n >>> from flax import nnx\n >>> import jax\n\n >>> layer = nnx.MultiHeadAttention(num_heads=8, in_features=5, qkv_features=16,\n ... decode=False, rngs=nnx.Rngs(0))\n >>> key1, key2, key3 = jax.random.split(jax.random.key(0), 3)\n >>> shape = (4, 3, 2, 5)\n >>> q, k, v = (\n ... jax.random.uniform(key1, shape),\n ... jax.random.uniform(key2, shape),\n ... jax.random.uniform(key3, shape),\n ... )\n\n >>> # different inputs for inputs_q, inputs_k and inputs_v\n >>> out = layer(q, k, v)\n >>> # equivalent output when inferring v\n >>> assert (layer(q, k) == layer(q, k, k)).all()\n >>> # equivalent output when inferring k and v\n >>> assert (layer(q) == layer(q, q)).all()\n >>> assert (layer(q) == layer(q, q, q)).all()\n\n Args:\n num_heads: number of attention heads. Features (i.e. inputs_q.shape[-1])\n should be divisible by the number of heads.\n in_features: int or tuple with number of input features.\n qkv_features: dimension of the key, query, and value.\n out_features: dimension of the last projection.\n in_kv_features: number of input features for computing key and value.\n num_kv_heads: number of key and value heads. If None, it defaults to\n ``num_heads``. If set to a value smaller than ``num_heads``, Grouped Query\n Attention (GQA) is used. ``num_heads`` must be divisible by\n ``num_kv_heads``.\n dtype: the dtype of the computation (default: infer from inputs and params)\n param_dtype: the dtype passed to parameter initializers (default: float32)\n broadcast_dropout: bool: use a broadcasted dropout along batch dims.\n dropout_rate: dropout rate\n deterministic: if false, the attention weight is masked randomly using\n dropout, whereas if true, the attention weights are deterministic.\n precision: numerical precision of the computation see `jax.lax.Precision`\n for details.\n kernel_init: initializer for the kernel of the Dense layers.\n out_kernel_init: optional initializer for the kernel of the output Dense layer,\n if None, the kernel_init is used.\n bias_init: initializer for the bias of the Dense layers.\n out_bias_init: optional initializer for the bias of the output Dense layer,\n if None, the bias_init is used.\n use_bias: bool: whether pointwise QKVO dense transforms use bias.\n attention_fn: dot_product_attention or compatible function. Accepts query,\n key, value, and returns output of shape `[bs, dim1, dim2, ..., dimN,,\n num_heads, value_channels]``\n decode: whether to prepare and use an autoregressive cache.\n normalize_qk: should QK normalization be applied (arxiv.org/abs/2302.05442).\n qkv_promote_dtype: function to promote the dtype of all input array arguments\n (including Variables accessed through ``self``) to the desired dtype for the\n query, key, and value LinearGeneral submodules.\n out_promote_dtype: function to promote the dtype of all input array arguments\n (including Variables accessed through ``self``) to the desired dtype for the\n output LinearGeneral submodule.\n ln_promote_dtype: function to promote the dtype of all input array arguments\n (including Variables accessed through ``self``) to the desired dtype for the\n LayerNorm submodules (query_ln and key_ln) when normalize_qk=True.\n rngs: rng key.\n keep_rngs: whether to store the input rngs as attribute (i.e. `self.rngs = rngs`)\n (default: True). If rngs is stored, we should split the module as\n `graphdef, params, nondiff = nnx.split(module, nnx.Param, ...)` where `nondiff`\n contains RNG object associated with stored `self.rngs`.\n kernel_metadata: Optional metadata dictionary to set when initializing\n the Dense layers.\n out_kernel_metadata: Optional metadata dictionary to set when initializing\n the output Dense layers. If None, the kernel_metadata is used.\n bias_metadata: Optional metadata dictionary to set when initializing\n the bias of the Dense layers.\n out_bias_metadata: Optional metadata dictionary to set when initializing\n the bias of the output Dense layers. If None, the bias_metadata is used.\n query_ln_scale_metadata: Optional metadata dictionary to set when initializing\n the scale of the query layer norm layer.\n key_ln_scale_metadata: Optional metadata dictionary to set when initializing\n the scale of the key layer norm layer.\n \"\"\"\n\n def __init__(\n self,\n num_heads: int,\n in_features: int,\n qkv_features: int | None = None,\n out_features: int | None = None,\n num_kv_heads: int | None = None,\n in_kv_features: int | None = None,\n *,\n dtype: Dtype | None = None,\n param_dtype: Dtype = jnp.float32,\n broadcast_dropout: bool = True,\n dropout_rate: float = 0.0,\n deterministic: bool | None = None,\n precision: PrecisionLike = None,\n kernel_init: Initializer = default_kernel_init,\n out_kernel_init: Initializer | None = None,\n bias_init: Initializer = initializers.zeros_init(),\n out_bias_init: Initializer | None = None,\n use_bias: bool = True,\n attention_fn: Callable[..., Array] = dot_product_attention,\n decode: bool | None = None,\n normalize_qk: bool = False,\n qkv_promote_dtype: PromoteDtypeFn = dtypes.promote_dtype,\n out_promote_dtype: PromoteDtypeFn = dtypes.promote_dtype,\n ln_promote_dtype: PromoteDtypeFn = dtypes.promote_dtype,\n # Deprecated, will be removed.\n qkv_dot_general: DotGeneralT | None = None,\n out_dot_general: DotGeneralT | None = None,\n qkv_dot_general_cls: Any = None,\n out_dot_general_cls: Any = None,\n rngs: rnglib.Rngs,\n keep_rngs: bool = True,\n kernel_metadata: Mapping[str, Any] = MappingProxyType({}),\n out_kernel_metadata: Mapping[str, Any] = MappingProxyType({}),\n bias_metadata: Mapping[str, Any] = MappingProxyType({}),\n out_bias_metadata: Mapping[str, Any] = MappingProxyType({}),\n query_ln_scale_metadata: Mapping[str, Any] = MappingProxyType({}),\n key_ln_scale_metadata: Mapping[str, Any] = MappingProxyType({}),\n ):\n self.num_heads = num_heads\n self.in_features = in_features\n self.qkv_features = (\n qkv_features if qkv_features is not None else in_features\n )\n self.out_features = (\n out_features if out_features is not None else in_features\n )\n self.in_kv_features = (\n in_kv_features if in_kv_features is not None else in_features\n )\n self.num_kv_heads = (\n num_kv_heads if num_kv_heads is not None else num_heads\n )\n\n if self.num_heads % self.num_kv_heads != 0:\n raise ValueError(\n f\"num_heads ({self.num_heads}) must be divisible by \"\n f\"num_kv_heads ({self.num_kv_heads}).\"\n )\n\n self.dtype = dtype\n self.param_dtype = param_dtype\n self.broadcast_dropout = broadcast_dropout\n self.dropout_rate = dropout_rate\n self.deterministic = deterministic\n self.precision = precision\n self.use_bias = use_bias\n self.attention_fn = attention_fn\n self.decode = decode\n self.normalize_qk = normalize_qk\n self.qkv_promote_dtype = qkv_promote_dtype\n self.out_promote_dtype = out_promote_dtype\n self.ln_promote_dtype = ln_promote_dtype\n self.qkv_dot_general = qkv_dot_general\n self.out_dot_general = out_dot_general\n self.qkv_dot_general_cls = qkv_dot_general_cls\n self.out_dot_general_cls = out_dot_general_cls\n\n if self.qkv_features % self.num_heads != 0:\n raise ValueError(\n f'Memory dimension ({self.qkv_features}) must be divisible by '\n f\"'num_heads' heads ({self.num_heads}).\"\n )\n\n self.head_dim = self.qkv_features // self.num_heads\n\n linear_general = functools.partial(\n LinearGeneral,\n dtype=self.dtype,\n param_dtype=self.param_dtype,\n kernel_init=kernel_init,\n bias_init=bias_init,\n use_bias=self.use_bias,\n precision=self.precision,\n promote_dtype=self.qkv_promote_dtype,\n dot_general=self.qkv_dot_general,\n dot_general_cls=self.qkv_dot_general_cls,\n kernel_metadata=kernel_metadata,\n bias_metadata=bias_metadata,\n )\n\n # project inputs_q to multi-headed q/k/v\n # dimensions are then [batch..., length, n_heads, n_features_per_head]\n self.query = linear_general(\n self.in_features,\n out_features=(self.num_heads, self.head_dim),\n rngs=rngs\n )\n self.key = linear_general(\n self.in_kv_features,\n out_features=(self.num_kv_heads, self.head_dim),\n rngs=rngs\n )\n self.value = linear_general(\n self.in_kv_features,\n out_features=(self.num_kv_heads, self.head_dim),\n rngs=rngs\n )\n\n self.query_ln: LayerNorm | None\n self.key_ln: LayerNorm | None\n if self.normalize_qk:\n # Normalizing query and key projections stabilizes training with higher\n # LR. See ViT-22B paper http://arxiv.org/abs/2302.05442 for analysis.\n self.query_ln = LayerNorm(\n self.head_dim,\n use_bias=False,\n dtype=self.dtype,\n param_dtype=self.param_dtype,\n promote_dtype=self.ln_promote_dtype,\n rngs=rngs,\n scale_metadata=query_ln_scale_metadata,\n )\n self.key_ln = LayerNorm(\n self.head_dim,\n use_bias=False,\n dtype=self.dtype,\n param_dtype=self.param_dtype,\n promote_dtype=self.ln_promote_dtype,\n rngs=rngs,\n scale_metadata=key_ln_scale_metadata,\n )\n else:\n self.query_ln = nnx.data(None)\n self.key_ln = nnx.data(None)\n\n self.out = LinearGeneral(\n in_features=(self.num_heads, self.head_dim),\n out_features=self.out_features,\n axis=(-2, -1),\n kernel_init=out_kernel_init or kernel_init,\n bias_init=out_bias_init or bias_init,\n use_bias=self.use_bias,\n dtype=self.dtype,\n param_dtype=self.param_dtype,\n precision=self.precision,\n promote_dtype=self.out_promote_dtype,\n dot_general=self.out_dot_general,\n dot_general_cls=self.out_dot_general_cls,\n rngs=rngs,\n kernel_metadata=out_kernel_metadata or kernel_metadata,\n bias_metadata=out_bias_metadata or bias_metadata,\n )\n self.rngs = rngs.dropout.fork() if keep_rngs and dropout_rate > 0 else None\n\n self.cached_key: nnx.Cache[Array] | None = nnx.data(None)\n self.cached_value: nnx.Cache[Array] | None = nnx.data(None)\n self.cache_index: nnx.Cache[Array] | None = nnx.data(None)\n\n def __call__(\n self,\n inputs_q: Array,\n inputs_k: Array | None = None,\n inputs_v: Array | None = None,\n *,\n mask: Array | None = None,\n deterministic: bool | None = None,\n rngs: rnglib.Rngs | rnglib.RngStream | None = None,\n sow_weights: bool = False,\n decode: bool | None = None,\n ):\n \"\"\"Applies multi-head dot product attention on the input data.\n\n Projects the inputs into multi-headed query, key, and value vectors,\n applies dot-product attention and project the results to an output vector.\n\n If both inputs_k and inputs_v are None, they will both copy the value of\n inputs_q (self attention).\n If only inputs_v is None, it will copy the value of inputs_k.\n\n Args:\n inputs_q: input queries of shape `[batch_sizes..., length, features]`.\n inputs_k: key of shape `[batch_sizes..., length, features]`. If None,\n inputs_k will copy the value of inputs_q.\n inputs_v: values of shape `[batch_sizes..., length, features]`. If None,\n inputs_v will copy the value of inputs_k.\n mask: attention mask of shape `[batch_sizes..., num_heads, query_length,\n key/value_length]`. Attention weights are masked out if their\n corresponding mask value is `False`.\n deterministic: if false, the attention weight is masked randomly using\n dropout, whereas if true, the attention weights are deterministic. The\n ``deterministic`` flag passed into the call method will take precedence\n over the ``deterministic`` flag passed into the constructor.\n rngs: rng key. The rng key passed into the call method will take\n precedence over the rng key passed into the constructor.\n sow_weights: if ``True``, the attention weights are sowed into the\n 'intermediates' collection.\n decode: whether to prepare and use an autoregressive cache. The ``decode``\n flag passed into the call method will take precedence over the ``decode``\n flag passed into the constructor.\n\n Returns:\n output of shape `[batch_sizes..., length, features]`.\n \"\"\"\n if rngs is None:\n rngs = self.rngs\n elif isinstance(rngs, rnglib.Rngs):\n rngs = rngs.dropout\n\n if inputs_k is None:\n if inputs_v is not None:\n raise ValueError(\n '`inputs_k` cannot be None if `inputs_v` is not None. '\n 'To have both `inputs_k` and `inputs_v` be the same value, pass in the '\n 'value to `inputs_k` and leave `inputs_v` as None.'\n )\n inputs_k = inputs_q\n if inputs_v is None:\n inputs_v = inputs_k\n\n if inputs_q.shape[-1] != self.in_features:\n raise ValueError(\n f'Incompatible input dimension, got {inputs_q.shape[-1]} '\n f'but module expects {self.in_features}.'\n )\n\n query = self.query(inputs_q)\n key = self.key(inputs_k)\n value = self.value(inputs_v)\n\n if self.normalize_qk:\n assert self.query_ln is not None and self.key_ln is not None\n # Normalizing query and key projections stabilizes training with higher\n # LR. See ViT-22B paper http://arxiv.org/abs/2302.05442 for analysis.\n query = self.query_ln(query)\n key = self.key_ln(key)\n\n # During fast autoregressive decoding, we feed one position at a time,\n # and cache the keys and values step by step.\n decode = first_from(\n decode,\n self.decode,\n error_msg=\"\"\"No `decode` argument was provided to MultiHeadAttention\n as either a __call__ argument, class attribute, or nnx.flag.\"\"\",\n )\n\n if decode:\n if (\n self.cached_key is None\n or self.cached_value is None\n or self.cache_index is None\n ):\n raise ValueError(\n 'Autoregressive cache not initialized, call ``init_cache`` first.'\n )\n (\n *batch_dims,\n max_length,\n num_kv_heads,\n depth_per_head,\n ) = self.cached_key.shape\n # shape check of cached keys against key input\n expected_shape = tuple(batch_dims) + (1, num_kv_heads, depth_per_head)\n if expected_shape != key.shape:\n raise ValueError(\n 'Autoregressive cache shape error, '\n f'expected key shape {expected_shape} instead got {key.shape}.'\n )\n # update key, value caches with our new 1d spatial slices\n cur_index = self.cache_index[...]\n zero = jnp.array(0, dtype=lax.dtype(cur_index.dtype))\n indices = (zero,) * len(batch_dims) + (cur_index, zero, zero)\n key = lax.dynamic_update_slice(self.cached_key[...], key, indices)\n value = lax.dynamic_update_slice(self.cached_value[...], value, indices)\n self.cached_key[...] = key\n self.cached_value[...] = value\n self.cache_index[...] += 1\n # causal mask for cached decoder self-attention:\n # our single query position should only attend to those key\n # positions that have already been generated and cached,\n # not the remaining zero elements.\n mask = combine_masks(\n mask,\n jnp.broadcast_to(\n jnp.arange(max_length) <= cur_index,\n tuple(batch_dims) + (1, 1, max_length),\n ),\n )\n\n if (\n self.dropout_rate > 0.0\n ): # Require `deterministic` only if using dropout.\n deterministic = first_from(\n deterministic,\n self.deterministic,\n error_msg=\"\"\"No `deterministic` argument was provided to MultiHeadAttention\n as either a __call__ argument, class attribute, or nnx.flag.\"\"\",\n )\n if not deterministic:\n if rngs is None:\n raise ValueError(\n \"'rngs' must be provided to __call__ method if \"\n \"MultiHeadAttention instance is defined with keep_rngs=False.\"\n )\n dropout_rng = rngs()\n else:\n dropout_rng = None\n else:\n deterministic = True\n dropout_rng = None\n\n # apply attention\n x = self.attention_fn(\n query,\n key,\n value,\n mask=mask,\n dropout_rng=dropout_rng,\n dropout_rate=self.dropout_rate,\n broadcast_dropout=self.broadcast_dropout,\n deterministic=deterministic,\n dtype=self.dtype,\n precision=self.precision,\n module=self if sow_weights else None,\n )\n # back to the original inputs dimensions\n out = self.out(x)\n return out\n\n def init_cache(self, input_shape: Shape, dtype: Dtype = jnp.float32):\n \"\"\"Initializes cache for fast autoregressive decoding. When\n ``decode=True``, this method must be called first before performing\n forward inference. When in decode mode, only one token must be passed\n at a time.\n\n Example usage::\n\n >>> from flax import nnx\n >>> import jax.numpy as jnp\n ...\n >>> batch_size = 5\n >>> embed_dim = 3\n >>> x = jnp.ones((batch_size, 1, embed_dim)) # single token\n ...\n >>> model_nnx = nnx.MultiHeadAttention(\n ... num_heads=2,\n ... in_features=3,\n ... qkv_features=6,\n ... out_features=6,\n ... decode=True,\n ... rngs=nnx.Rngs(42),\n ... )\n ...\n >>> # out_nnx = model_nnx(x) <-- throws an error because cache isn't initialized\n ...\n >>> model_nnx.init_cache(x.shape)\n >>> out_nnx = model_nnx(x)\n \"\"\"\n cache_shape = (*input_shape[:-1], self.num_kv_heads, self.head_dim)\n self.cached_key = nnx.Cache(jnp.zeros(cache_shape, dtype))\n self.cached_value = nnx.Cache(jnp.zeros(cache_shape, dtype))\n self.cache_index = nnx.Cache(jnp.array(0, dtype=jnp.int32))\n\n def set_view(\n self,\n deterministic: bool | None = None,\n decode: bool | None = None,\n batch_size: int | Shape | None = None,\n max_length: int | None = None,\n ):\n \"\"\"Class method used by ``nnx.view``.\n\n Args:\n train: if True, the module is set to training mode.\n deterministic: if True, the module is set to deterministic mode.\n decode: if True, the module is set to decode mode.\n batch_size: the batch size to use for the cache.\n max_length: the max length to use for the cache.\n \"\"\"\n if deterministic is not None:\n self.deterministic = deterministic\n\n if decode is not None:\n self.decode = decode\n if (\n not hasattr(self, 'cached_key')\n or not hasattr(self, 'cached_value')\n or not hasattr(self, 'cache_index')\n ):\n if batch_size is None:\n raise TypeError(\n \"'batch_size' must be provided when initializing cache.\"\n )\n if max_length is None:\n raise TypeError(\n \"'max_length' must be provided when initializing cache.\"\n )\n if isinstance(batch_size, int):\n batch_size = (batch_size,)\n\n # initialize cache\n cache_shape = (*batch_size, max_length, self.num_kv_heads, self.head_dim)\n self.cached_key = nnx.Cache(jnp.zeros(cache_shape, self.dtype))\n self.cached_value = nnx.Cache(jnp.zeros(cache_shape, self.dtype))\n self.cache_index = nnx.Cache(jnp.array(0, dtype=jnp.int32))\n\n\n# mask-making utility functions\n\n\ndef make_attention_mask(\n query_input: Array,\n key_input: Array,\n pairwise_fn: Callable[..., Any] = jnp.multiply,\n extra_batch_dims: int = 0,\n dtype: Dtype = jnp.float32,\n):\n \"\"\"Mask-making helper for attention weights.\n\n In case of 1d inputs (i.e., `[batch..., len_q]`, `[batch..., len_kv]`, the\n attention weights will be `[batch..., heads, len_q, len_kv]` and this\n function will produce `[batch..., 1, len_q, len_kv]`.\n\n Args:\n query_input: a batched, flat input of query_length size\n key_input: a batched, flat input of key_length size\n pairwise_fn: broadcasting elementwise comparison function\n extra_batch_dims: number of extra batch dims to add singleton axes for, none\n by default\n dtype: mask return dtype\n\n Returns:\n A `[batch..., 1, len_q, len_kv]` shaped mask for 1d attention.\n \"\"\"\n mask = pairwise_fn(\n jnp.expand_dims(query_input, axis=-1), jnp.expand_dims(key_input, axis=-2)\n )\n mask = jnp.expand_dims(mask, axis=-3)\n mask = jnp.expand_dims(mask, axis=tuple(range(extra_batch_dims)))\n return mask.astype(dtype)\n\n\ndef make_causal_mask(\n x: Array, extra_batch_dims: int = 0, dtype: Dtype = jnp.float32\n) -> Array:\n \"\"\"Make a causal mask for self-attention.\n\n In case of 1d inputs (i.e., `[batch..., len]`, the self-attention weights\n will be `[batch..., heads, len, len]` and this function will produce a\n causal mask of shape `[batch..., 1, len, len]`.\n\n Args:\n x: input array of shape `[batch..., len]`\n extra_batch_dims: number of batch dims to add singleton axes for, none by\n default\n dtype: mask return dtype\n\n Returns:\n A `[batch..., 1, len, len]` shaped causal mask for 1d attention.\n \"\"\"\n idxs = jnp.broadcast_to(jnp.arange(x.shape[-1], dtype=jnp.int32), x.shape)\n return make_attention_mask(\n idxs,\n idxs,\n jnp.greater_equal,\n extra_batch_dims=extra_batch_dims,\n dtype=dtype,\n )\n\n\ndef combine_masks(\n *masks: Array | None, dtype: Dtype = jnp.float32\n) -> Array | None:\n \"\"\"Combine attention masks.\n\n Args:\n *masks: set of attention mask arguments to combine, some can be None.\n dtype: dtype for the returned mask.\n\n Returns:\n Combined mask, reduced by logical and, returns None if no masks given.\n \"\"\"\n masks_list = [m for m in masks if m is not None]\n if not masks_list:\n return None\n assert all(\n map(lambda x: x.ndim == masks_list[0].ndim, masks_list)\n ), f'masks must have same rank: {tuple(map(lambda x: x.ndim, masks_list))}'\n mask, *other_masks = masks_list\n for other_mask in other_masks:\n mask = jnp.logical_and(mask, other_mask)\n return mask.astype(dtype)\n" }, { "path": "flax/nnx/nn/dtypes.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport typing as tp\nfrom flax.typing import Dtype\nfrom jax import numpy as jnp\n\nT = tp.TypeVar('T', bound=tuple)\n\n\ndef canonicalize_dtype(\n *args, dtype: Dtype | None = None, inexact: bool = True\n) -> Dtype:\n \"\"\"Canonicalize an optional dtype to the definitive dtype.\n\n If the ``dtype`` is None this function will infer the dtype. If it is not\n None it will be returned unmodified or an exceptions is raised if the dtype\n is invalid.\n from the input arguments using ``jnp.result_type``.\n\n Args:\n *args: JAX array compatible values. None values\n are ignored.\n dtype: Optional dtype override. If specified the arguments are cast to\n the specified dtype instead and dtype inference is disabled.\n inexact: When True, the output dtype must be a subdtype\n of `jnp.inexact`. Inexact dtypes are real or complex floating points. This\n is useful when you want to apply operations that don't work directly on\n integers like taking a mean for example.\n Returns:\n The dtype that *args should be cast to.\n \"\"\"\n if dtype is None:\n args_filtered = [jnp.asarray(x) for x in args if x is not None]\n dtype = jnp.result_type(*args_filtered)\n if inexact and not jnp.issubdtype(dtype, jnp.inexact):\n dtype = jnp.promote_types(jnp.float32, dtype)\n if inexact and not jnp.issubdtype(dtype, jnp.inexact):\n raise ValueError(f'Dtype must be inexact: {dtype}')\n return dtype\n\n\ndef promote_dtype(args: T, /, *, dtype=None, inexact=True) -> T:\n \"\"\" \"Promotes input arguments to a specified or inferred dtype.\n\n All args are cast to the same dtype. See ``canonicalize_dtype`` for how\n this dtype is determined.\n\n The behavior of promote_dtype is mostly a convinience wrapper around\n ``jax.numpy.promote_types``. The differences being that it automatically casts\n all input to the inferred dtypes, allows inference to be overridden by a\n forced dtype, and has an optional check to garantuee the resulting dtype is\n inexact.\n\n Args:\n *args: JAX array compatible values. None values\n are returned as is.\n dtype: Optional dtype override. If specified the arguments are cast to\n the specified dtype instead and dtype inference is disabled.\n inexact: When True, the output dtype must be a subdtype\n of `jnp.inexact`. Inexact dtypes are real or complex floating points. This\n is useful when you want to apply operations that don't work directly on\n integers like taking a mean for example.\n Returns:\n The arguments cast to arrays of the same dtype.\n \"\"\"\n dtype = canonicalize_dtype(*args, dtype=dtype, inexact=inexact)\n arrays = tuple(jnp.asarray(x, dtype) if x is not None else None for x in args)\n return arrays # type: ignore[return-value]\n" }, { "path": "flax/nnx/nn/initializers.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport typing as tp\n\nfrom jax.nn.initializers import constant as constant\nfrom jax.nn.initializers import delta_orthogonal as delta_orthogonal\nfrom jax.nn.initializers import glorot_normal as glorot_normal\nfrom jax.nn.initializers import glorot_uniform as glorot_uniform\nfrom jax.nn.initializers import he_normal as he_normal\nfrom jax.nn.initializers import he_uniform as he_uniform\nfrom jax.nn.initializers import kaiming_normal as kaiming_normal\nfrom jax.nn.initializers import kaiming_uniform as kaiming_uniform\nfrom jax.nn.initializers import lecun_normal as lecun_normal\nfrom jax.nn.initializers import lecun_uniform as lecun_uniform\nfrom jax.nn.initializers import normal as normal\nfrom jax.nn.initializers import ones as ones\nfrom jax.nn.initializers import orthogonal as orthogonal\nfrom jax.nn.initializers import truncated_normal as truncated_normal\nfrom jax.nn.initializers import uniform as uniform\nfrom jax.nn.initializers import variance_scaling as variance_scaling\nfrom jax.nn.initializers import xavier_normal as xavier_normal\nfrom jax.nn.initializers import xavier_uniform as xavier_uniform\nfrom jax.nn.initializers import zeros as zeros\nfrom flax.typing import Initializer\n\nDtypeLikeInexact = tp.Any\n\n\ndef zeros_init() -> Initializer:\n \"\"\"Builds an initializer that returns a constant array full of zeros.\n\n >>> import jax, jax.numpy as jnp\n >>> from flax.nnx import initializers\n >>> zeros_initializer = initializers.zeros_init()\n >>> zeros_initializer(jax.random.key(42), (2, 3), jnp.float32)\n Array([[0., 0., 0.],\n [0., 0., 0.]], dtype=float32)\n \"\"\"\n return zeros\n\n\ndef ones_init() -> Initializer:\n \"\"\"Builds an initializer that returns a constant array full of ones.\n\n >>> import jax, jax.numpy as jnp\n >>> from flax.nnx import initializers\n >>> ones_initializer = initializers.ones_init()\n >>> ones_initializer(jax.random.key(42), (3, 2), jnp.float32)\n Array([[1., 1.],\n [1., 1.],\n [1., 1.]], dtype=float32)\n \"\"\"\n return ones\n" }, { "path": "flax/nnx/nn/linear.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nimport typing as tp\nfrom types import MappingProxyType\n\nimport jax\nimport jax.numpy as jnp\nimport numpy as np\nfrom jax import lax\nimport opt_einsum\n\nfrom flax.core.frozen_dict import FrozenDict\nfrom flax import nnx\nfrom flax.nnx import rnglib, variablelib\nfrom flax.nnx.module import Module, first_from\nfrom flax.nnx.nn import dtypes, initializers\nfrom flax.typing import (\n Dtype,\n Shape,\n Initializer,\n PrecisionLike,\n DotGeneralT,\n ConvGeneralDilatedT,\n PaddingLike,\n LaxPadding,\n PromoteDtypeFn,\n EinsumT,\n)\n\nArray = jax.Array\nAxis = int\nSize = int\n\n\ndefault_kernel_init = initializers.lecun_normal()\ndefault_bias_init = initializers.zeros_init()\n\n\ndef canonicalize_padding(padding: PaddingLike, rank: int) -> LaxPadding:\n \"\"\" \"Canonicalizes conv padding to a jax.lax supported format.\"\"\"\n if isinstance(padding, str):\n return padding\n if isinstance(padding, int):\n return [(padding, padding)] * rank\n if isinstance(padding, tp.Sequence) and len(padding) == rank:\n new_pad = []\n for p in padding:\n if isinstance(p, int):\n new_pad.append((p, p))\n elif isinstance(p, tuple) and len(p) == 2:\n new_pad.append(p)\n else:\n break\n if len(new_pad) == rank:\n return new_pad\n raise ValueError(\n f'Invalid padding format: {padding}, should be str, int,'\n f' or a sequence of len {rank} where each element is an'\n ' int or pair of ints.'\n )\n\n\ndef _conv_dimension_numbers(input_shape):\n \"\"\"Computes the dimension numbers based on the input shape.\"\"\"\n ndim = len(input_shape)\n lhs_spec = (0, ndim - 1) + tuple(range(1, ndim - 1))\n rhs_spec = (ndim - 1, ndim - 2) + tuple(range(0, ndim - 2))\n out_spec = lhs_spec\n return lax.ConvDimensionNumbers(lhs_spec, rhs_spec, out_spec)\n\n\ndef _normalize_axes(axes: tuple[int, ...], ndim: int) -> tuple[int, ...]:\n # A tuple by convention. len(axes_tuple) then also gives the rank efficiently.\n return tuple(sorted(ax if ax >= 0 else ndim + ax for ax in axes))\n\n\ndef _canonicalize_tuple(x: tp.Sequence[int] | int) -> tuple[int, ...]:\n if isinstance(x, tp.Iterable):\n return tuple(x)\n else:\n return (x,)\n\n\nclass LinearGeneral(Module):\n \"\"\"A linear transformation with flexible axes.\n\n Example usage::\n\n >>> from flax import nnx\n >>> import jax, jax.numpy as jnp\n ...\n >>> # equivalent to `nnx.Linear(2, 4)`\n >>> layer = nnx.LinearGeneral(2, 4, rngs=nnx.Rngs(0))\n >>> layer.kernel.shape\n (2, 4)\n >>> # output features (4, 5)\n >>> layer = nnx.LinearGeneral(2, (4, 5), rngs=nnx.Rngs(0))\n >>> layer.kernel.shape\n (2, 4, 5)\n >>> layer.bias.shape\n (4, 5)\n >>> # apply transformation on the the second and last axes\n >>> layer = nnx.LinearGeneral((2, 3), (4, 5), axis=(1, -1), rngs=nnx.Rngs(0))\n >>> layer.kernel.shape\n (2, 3, 4, 5)\n >>> layer.bias.shape\n (4, 5)\n >>> y = layer(jnp.ones((16, 2, 3)))\n >>> y.shape\n (16, 4, 5)\n\n Args:\n in_features: int or tuple with number of input features.\n out_features: int or tuple with number of output features.\n axis: int or tuple with axes to apply the transformation on. For instance,\n (-2, -1) will apply the transformation to the last two axes.\n batch_axis: mapping of batch axis indices to axis size.\n use_bias: whether to add a bias to the output (default: True).\n dtype: the dtype of the computation (default: infer from input and params).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n kernel_init: initializer function for the weight matrix.\n bias_init: initializer function for the bias.\n dot_general: dot product function (default: None). If neither this nor\n ``dot_general_cls`` are provided, ``jax.lax.dot_general`` is used.\n dot_general_cls: dot product function class to instantiate a dot product function\n as ``dot_general = dot_general_cls()`` (default: None).\n precision: numerical precision of the computation see ``jax.lax.Precision``\n for details.\n promote_dtype: function to promote the dtype of the arrays to the desired\n dtype. The function should accept a tuple of ``(inputs, kernel, bias)``\n and a ``dtype`` keyword argument, and return a tuple of arrays with the\n promoted dtype.\n preferred_element_type: Optional parameter controls the data type output by\n the dot product. This argument is passed to ``dot_general`` function.\n See ``jax.lax.dot`` for details.\n rngs: rng key.\n kernel_metadata: Optional metadata dictionary to set when initializing\n the weight matrix.\n bias_metadata: Optional metadata dictionary to set when initializing\n the bias.\n \"\"\"\n\n def __init__(\n self,\n in_features: Size | tp.Sequence[Size],\n out_features: Size | tp.Sequence[Size],\n *,\n axis: Axis | tp.Sequence[Axis] = -1,\n batch_axis: tp.Mapping[Axis, Size] = FrozenDict({}),\n use_bias: bool = True,\n dtype: Dtype | None = None,\n param_dtype: Dtype = jnp.float32,\n kernel_init: Initializer = default_kernel_init,\n bias_init: Initializer = default_bias_init,\n precision: PrecisionLike = None,\n promote_dtype: PromoteDtypeFn = dtypes.promote_dtype,\n dot_general: DotGeneralT | None = None,\n dot_general_cls: tp.Any = None,\n preferred_element_type: Dtype | None = None,\n rngs: rnglib.Rngs,\n kernel_metadata: tp.Mapping[str, tp.Any] = MappingProxyType({}),\n bias_metadata: tp.Mapping[str, tp.Any] = MappingProxyType({}),\n ):\n self.in_features = _canonicalize_tuple(in_features)\n self.out_features = _canonicalize_tuple(out_features)\n self.axis = _canonicalize_tuple(axis)\n self.batch_axis = FrozenDict[Axis, Size](batch_axis)\n self.use_bias = use_bias\n self.dtype = dtype\n self.param_dtype = param_dtype\n self.precision = precision\n self.dot_general = dot_general\n self.dot_general_cls = dot_general_cls\n self.promote_dtype = promote_dtype\n self.preferred_element_type = preferred_element_type\n\n if len(self.in_features) != len(self.axis):\n raise ValueError(\n 'in_features and axis must have the same length. '\n f'Got {self.in_features} and {self.axis}.'\n )\n\n if batch_axis:\n batch_dims = tuple(batch_axis.keys())\n max_dim = np.max(batch_dims)\n if set(batch_dims) != set(range(max_dim + 1)):\n raise ValueError(\n 'batch_dims %s must be consecutive leading '\n 'dimensions starting from 0.' % str(batch_dims)\n )\n\n n_batch_axis = len(self.batch_axis)\n n_in_features = len(self.in_features)\n n_out_features = len(self.out_features)\n\n def kernel_init_wrap(rng, shape, dtype):\n flat_shape = (\n np.prod(shape[:n_batch_axis])\n * np.prod(shape[n_batch_axis : n_in_features + n_batch_axis]),\n np.prod(shape[-n_out_features:]),\n )\n flat_shape = jax.tree.map(int, flat_shape)\n kernel = kernel_init(rng, flat_shape, dtype)\n if isinstance(kernel, variablelib.VariableMetadata):\n kernel.raw_value = jnp.reshape(kernel.raw_value, shape)\n else:\n kernel = jnp.reshape(kernel, shape)\n\n return kernel\n\n batch_shape = tuple(self.batch_axis.values())\n kernel_shape = (\n *batch_shape,\n *self.in_features,\n *self.out_features,\n )\n self.kernel = nnx.Param(\n kernel_init_wrap(rngs.params(), kernel_shape, self.param_dtype), **kernel_metadata\n )\n\n self.bias: nnx.Param[jax.Array] | None\n if self.use_bias:\n def bias_init_wrap(rng, shape, dtype):\n flat_shape = (int(np.prod(shape)),)\n bias = bias_init(rng, flat_shape, dtype)\n if isinstance(bias, variablelib.VariableMetadata):\n bias.raw_value = jnp.reshape(bias.raw_value, shape)\n else:\n bias = jnp.reshape(bias, shape)\n return bias\n\n bias_shape = (*batch_shape, *self.out_features)\n self.bias = nnx.Param(\n bias_init_wrap(rngs.params(), bias_shape, self.param_dtype),\n **bias_metadata,\n )\n else:\n self.bias = nnx.data(None)\n\n def __call__(self, inputs: Array, out_sharding = None) -> Array:\n \"\"\"Applies a linear transformation to the inputs along multiple dimensions.\n\n Args:\n inputs: The nd-array to be transformed.\n\n Returns:\n The transformed input.\n \"\"\"\n\n ndim = inputs.ndim\n n_batch_dims = len(self.batch_axis)\n axis = _normalize_axes(self.axis, ndim)\n batch_axis = _normalize_axes(tuple(self.batch_axis.keys()), ndim)\n n_axis = len(axis)\n\n # batch and non-contracting dims of input with 1s for batch dims.\n expanded_batch_shape = tuple(\n inputs.shape[ax] if ax in batch_axis else 1\n for ax in range(inputs.ndim)\n if ax not in axis\n )\n kernel = self.kernel[...]\n bias = self.bias[...] if self.bias is not None else None\n\n batch_ind = tuple(range(n_batch_dims))\n contract_ind = tuple(range(n_batch_dims, n_axis + n_batch_dims))\n\n inputs, kernel, bias = self.promote_dtype(\n (inputs, kernel, bias), dtype=self.dtype\n )\n\n if self.dot_general_cls is not None:\n dot_general = self.dot_general_cls()\n elif self.dot_general is not None:\n dot_general = self.dot_general\n else:\n dot_general = lax.dot_general\n # We use dot_general_kwargs for BC compatibility with\n # user custom dot_general/dot_general_cls which may not have\n # preferred_element_type argument to avoid breaking\n # existing code\n dot_general_kwargs = {'out_sharding': out_sharding}\n if self.preferred_element_type is not None:\n dot_general_kwargs[\"preferred_element_type\"] = self.preferred_element_type\n out = dot_general(\n inputs,\n kernel,\n ((axis, contract_ind), (batch_axis, batch_ind)),\n precision=self.precision,\n **dot_general_kwargs,\n )\n # dot_general output has shape [batch_dims/group_dims] + [feature_dims]\n if bias is not None:\n # expand bias shape to broadcast bias over batch dims.\n bias = jnp.reshape(bias, (*expanded_batch_shape, *self.out_features))\n out += bias\n return out\n\n\nclass Linear(Module):\n \"\"\"A linear transformation applied over the last dimension of the input.\n\n Example usage::\n\n >>> from flax import nnx\n >>> import jax, jax.numpy as jnp\n\n >>> layer = nnx.Linear(in_features=3, out_features=4, rngs=nnx.Rngs(0))\n >>> jax.tree.map(jnp.shape, nnx.state(layer))\n State({\n 'bias': Param(\n value=(4,)\n ),\n 'kernel': Param(\n value=(3, 4)\n )\n })\n\n Args:\n in_features: the number of input features.\n out_features: the number of output features.\n use_bias: whether to add a bias to the output (default: True).\n dtype: the dtype of the computation (default: infer from input and params).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n precision: numerical precision of the computation see ``jax.lax.Precision``\n for details.\n kernel_init: initializer function for the weight matrix.\n bias_init: initializer function for the bias.\n dot_general: dot product function.\n promote_dtype: function to promote the dtype of the arrays to the desired\n dtype. The function should accept a tuple of ``(inputs, kernel, bias)``\n and a ``dtype`` keyword argument, and return a tuple of arrays with the\n promoted dtype.\n preferred_element_type: Optional parameter controls the data type output by\n the dot product. This argument is passed to ``dot_general`` function.\n See ``jax.lax.dot`` for details.\n rngs: rng key.\n kernel_metadata: Optional metadata dictionary to set when initializing\n the weight matrix.\n bias_metadata: Optional metadata dictionary to set when initializing\n the bias.\n \"\"\"\n\n def __init__(\n self,\n in_features: int,\n out_features: int,\n *,\n use_bias: bool = True,\n dtype: tp.Optional[Dtype] = None,\n param_dtype: Dtype = jnp.float32,\n precision: PrecisionLike = None,\n kernel_init: Initializer = default_kernel_init,\n bias_init: Initializer = default_bias_init,\n dot_general: DotGeneralT = lax.dot_general,\n promote_dtype: PromoteDtypeFn = dtypes.promote_dtype,\n preferred_element_type: Dtype | None = None,\n rngs: rnglib.Rngs,\n kernel_metadata: tp.Mapping[str, tp.Any] = MappingProxyType({}),\n bias_metadata: tp.Mapping[str, tp.Any] = MappingProxyType({}),\n ):\n kernel_key = rngs.params()\n self.kernel = nnx.Param(\n kernel_init(kernel_key, (in_features, out_features), param_dtype),\n **kernel_metadata,\n )\n self.bias: nnx.Param[jax.Array] | None\n if use_bias:\n bias_key = rngs.params()\n self.bias = nnx.Param(\n bias_init(bias_key, (out_features,), param_dtype),\n **bias_metadata,\n )\n else:\n self.bias = nnx.data(None)\n\n self.in_features = in_features\n self.out_features = out_features\n self.use_bias = use_bias\n self.dtype = dtype\n self.param_dtype = param_dtype\n self.precision = precision\n self.dot_general = dot_general\n self.promote_dtype = promote_dtype\n self.preferred_element_type = preferred_element_type\n\n def __call__(self, inputs: Array, out_sharding = None) -> Array:\n \"\"\"Applies a linear transformation to the inputs along the last dimension.\n\n Args:\n inputs: The nd-array to be transformed.\n\n Returns:\n The transformed input.\n \"\"\"\n kernel = self.kernel[...]\n bias = self.bias[...] if self.bias is not None else None\n\n inputs, kernel, bias = self.promote_dtype(\n (inputs, kernel, bias), dtype=self.dtype\n )\n # We use dot_general_kwargs for BC compatibility with\n # user custom self.dot_general method which may not have\n # preferred_element_type argument to avoid breaking\n # existing code\n dot_general_kwargs = {'out_sharding': out_sharding}\n if self.preferred_element_type is not None:\n dot_general_kwargs[\"preferred_element_type\"] = self.preferred_element_type\n y = self.dot_general(\n inputs,\n kernel,\n (((inputs.ndim - 1,), (0,)), ((), ())),\n precision=self.precision,\n **dot_general_kwargs,\n )\n assert self.use_bias == (bias is not None)\n if bias is not None:\n y += jnp.reshape(bias, (1,) * (y.ndim - 1) + (-1,))\n return y\n\n\nclass Einsum(Module):\n \"\"\"An einsum transformation with learnable kernel and bias.\n\n Example usage::\n\n >>> from flax import nnx\n >>> import jax.numpy as jnp\n ...\n >>> layer = nnx.Einsum('nta,hab->nthb', (8, 2, 4), (8, 4), rngs=nnx.Rngs(0))\n >>> layer.kernel.shape\n (8, 2, 4)\n >>> layer.bias.shape\n (8, 4)\n >>> y = layer(jnp.ones((16, 11, 2)))\n >>> y.shape\n (16, 11, 8, 4)\n\n Args:\n einsum_str: a string to denote the einsum equation. The equation must\n have exactly two operands, the lhs being the input passed in, and\n the rhs being the learnable kernel. Exactly one of ``einsum_str``\n in the constructor argument and call argument must be not None,\n while the other must be None.\n kernel_shape: the shape of the kernel.\n bias_shape: the shape of the bias. If this is None, a bias won't be used.\n dtype: the dtype of the computation (default: infer from input and params).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n precision: numerical precision of the computation see ``jax.lax.Precision``\n for details.\n kernel_init: initializer function for the weight matrix.\n bias_init: initializer function for the bias.\n promote_dtype: function to promote the dtype of the arrays to the desired\n dtype. The function should accept a tuple of ``(inputs, kernel, bias)``\n and a ``dtype`` keyword argument, and return a tuple of arrays with the\n promoted dtype.\n einsum_op: An injectable alternative of `jnp.einsum` to do the computation.\n Should support same signature as `jnp.einsum`.\n preferred_element_type: Optional parameter controls the data type output by\n the dot product. This argument is passed to ``dot_general`` function.\n See ``jax.lax.dot`` for details.\n rngs: rng key.\n kernel_metadata: Optional metadata dictionary to set when initializing\n the weight matrix.\n bias_metadata: Optional metadata dictionary to set when initializing\n the bias.\n \"\"\"\n\n def __init__(\n self,\n einsum_str: str,\n kernel_shape: Shape,\n bias_shape: tp.Optional[Shape] = None,\n *,\n dtype: tp.Optional[Dtype] = None,\n param_dtype: Dtype = jnp.float32,\n precision: PrecisionLike = None,\n kernel_init: Initializer = default_kernel_init,\n bias_init: Initializer = default_bias_init,\n promote_dtype: PromoteDtypeFn = dtypes.promote_dtype,\n einsum_op: EinsumT = jnp.einsum,\n preferred_element_type: Dtype | None = None,\n rngs: rnglib.Rngs,\n kernel_metadata: tp.Mapping[str, tp.Any] = MappingProxyType({}),\n bias_metadata: tp.Mapping[str, tp.Any] = MappingProxyType({}),\n ):\n einsum_str = einsum_str.replace(' ', '')\n self._einsum_str_check(einsum_str)\n\n kernel_key = rngs.params()\n self.kernel = nnx.Param(\n kernel_init(kernel_key, kernel_shape, param_dtype), **kernel_metadata\n )\n\n self.bias: nnx.Param | None\n if bias_shape is not None:\n bias_key = rngs.params()\n self.bias = nnx.Param(\n bias_init(bias_key, bias_shape, param_dtype), **bias_metadata\n )\n else:\n self.bias = nnx.data(None)\n\n self.einsum_str = einsum_str\n self.kernel_shape = kernel_shape\n self.bias_shape = bias_shape\n self.dtype = dtype\n self.param_dtype = param_dtype\n self.precision = precision\n self.promote_dtype = promote_dtype\n self.einsum_op = einsum_op\n self.preferred_element_type = preferred_element_type\n\n def __call__(\n self, inputs: Array, einsum_str: tp.Optional[str] = None, out_sharding = None\n ) -> Array:\n \"\"\"Applies a linear transformation to the inputs along the last dimension.\n\n Args:\n inputs: The nd-array to be transformed.\n einsum_str: a string to denote the einsum equation. The equation must\n have exactly two operands, the lhs being the input passed in, and\n the rhs being the learnable kernel. Exactly one of ``einsum_str``\n in the constructor argument and call argument must be not None,\n while the other must be None.\n\n Returns:\n The transformed input.\n \"\"\"\n einsum_str = first_from(\n einsum_str,\n self.einsum_str,\n error_msg=\"\"\"No `einsum_str` argument was provided to Einsum\n as either a __call__ argument, or class attribute.\"\"\",\n )\n einsum_str = einsum_str.replace(' ', '')\n self._einsum_str_check(einsum_str)\n\n inputs, kernel, bias = self.promote_dtype(\n (\n inputs,\n self.kernel[...],\n self.bias[...] if self.bias is not None else self.bias,\n ),\n dtype=self.dtype,\n )\n # We use einsum_op_kwargs for BC compatibility with\n # user custom self.einsum_op method which may not have\n # preferred_element_type argument to avoid breaking\n # existing code\n einsum_op_kwargs = {'out_sharding': out_sharding}\n if self.preferred_element_type is not None:\n einsum_op_kwargs[\"preferred_element_type\"] = self.preferred_element_type\n\n y = self.einsum_op(\n einsum_str, inputs, kernel, precision=self.precision, **einsum_op_kwargs\n )\n\n if bias is not None:\n broadcasted_bias_shape = self._infer_broadcasted_bias_shape(\n einsum_str, inputs, kernel\n )\n y += jnp.reshape(bias, broadcasted_bias_shape)\n return y\n\n def _infer_broadcasted_bias_shape(\n self, einsum_str: str, lhs: Array, rhs: Array\n ):\n \"\"\"Infer the broadcasted bias shape given the ``einsum_str``, ``lhs``\n and ``rhs`` arrays. This is needed reshaping the bias and it to the\n output during forward inference.\n\n This function first replaces all ellipses with actual letter characters,\n then computes the broadcasted bias shape by checking to see which axes in\n the rhs array remain in the resulting array after einsumming. These axes\n are the embedding/feature dimensions, and all other axes in rhs are\n reduction axes.\n \"\"\"\n # More details on the parsing function: https://github.com/dgasmith/opt_einsum/blob/c826bb7df16f470a69f7bf90598fc27586209d11/opt_einsum/parser.py#L246\n # returns the einsum string representation of the operands and result, with\n # ellipsis replaced by actual letter characters\n operands_str, result_str, _ = opt_einsum.parser.parse_einsum_input(\n (einsum_str, lhs, rhs)\n )\n\n # rhs_dict is a dict{character:index} mapping that maps every character in\n # the rhs einsum string representation to its corresponding index position in the string\n rhs_dict = {c: i for i, c in enumerate(operands_str.split(',')[1])}\n assert len(rhs_dict) == len(self.kernel_shape)\n\n broadcasted_bias_shape = [1] * len(result_str)\n for i, c in enumerate(result_str):\n if c in rhs_dict:\n broadcasted_bias_shape[i] = self.kernel_shape[rhs_dict[c]]\n\n return broadcasted_bias_shape\n\n def _einsum_str_check(self, einsum_str):\n if '->' not in einsum_str:\n raise ValueError(\n '`einsum_str` equation must be explicit and include \"->\".'\n )\n if einsum_str.count(',') != 1:\n raise ValueError(\n '`einsum_str` equation must have exactly two operands and '\n 'therefore, exactly one comma character, instead of '\n f'{einsum_str.count(\",\")}'\n )\n\n\nclass Conv(Module):\n \"\"\"Convolution Module wrapping ``lax.conv_general_dilated``.\n\n Example usage::\n\n >>> from flax import nnx\n >>> import jax.numpy as jnp\n\n >>> rngs = nnx.Rngs(0)\n >>> x = jnp.ones((1, 8, 3))\n\n >>> # valid padding\n >>> layer = nnx.Conv(in_features=3, out_features=4, kernel_size=(3,),\n ... padding='VALID', rngs=rngs)\n >>> layer.kernel.shape\n (3, 3, 4)\n >>> layer.bias.shape\n (4,)\n >>> out = layer(x)\n >>> out.shape\n (1, 6, 4)\n\n >>> # circular padding with stride 2\n >>> layer = nnx.Conv(in_features=3, out_features=4, kernel_size=(3, 3),\n ... strides=2, padding='CIRCULAR', rngs=rngs)\n >>> layer.kernel.shape\n (3, 3, 3, 4)\n >>> layer.bias.shape\n (4,)\n >>> out = layer(x)\n >>> out.shape\n (1, 4, 4)\n\n >>> # apply lower triangle mask\n >>> mask = jnp.tril(jnp.ones((3, 3, 4)))\n >>> layer = nnx.Conv(in_features=3, out_features=4, kernel_size=(3,),\n ... mask=mask, padding='VALID', rngs=rngs)\n >>> out = layer(x)\n\n Args:\n in_features: int or tuple with number of input features.\n out_features: int or tuple with number of output features.\n kernel_size: shape of the convolutional kernel. For 1D convolution,\n the kernel size can be passed as an integer, which will be interpreted\n as a tuple of the single integer. For all other cases, it must be a\n sequence of integers.\n strides: an integer or a sequence of ``n`` integers, representing the\n inter-window strides (default: 1).\n padding: either the string ``'SAME'``, the string ``'VALID'``, the string\n ``'CIRCULAR'`` (periodic boundary conditions), the string `'REFLECT'`\n (reflection across the padding boundary), or a sequence of ``n``\n ``(low, high)`` integer pairs that give the padding to apply before and after each\n spatial dimension. A single int is interpreted as applying the same padding\n in all dims and passing a single int in a sequence causes the same padding\n to be used on both sides. ``'CAUSAL'`` padding for a 1D convolution will\n left-pad the convolution axis, resulting in same-sized output.\n input_dilation: an integer or a sequence of ``n`` integers, giving the\n dilation factor to apply in each spatial dimension of ``inputs``\n (default: 1). Convolution with input dilation ``d`` is equivalent to\n transposed convolution with stride ``d``.\n kernel_dilation: an integer or a sequence of ``n`` integers, giving the\n dilation factor to apply in each spatial dimension of the convolution\n kernel (default: 1). Convolution with kernel dilation\n is also known as 'atrous convolution'.\n feature_group_count: integer, default 1. If specified divides the input\n features into groups.\n use_bias: whether to add a bias to the output (default: True).\n mask: Optional mask for the weights during masked convolution. The mask must\n be the same shape as the convolution weight matrix.\n dtype: the dtype of the computation (default: infer from input and params).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n conv_general_dilated: the convolution function to use (default:\n ``jax.lax.conv_general_dilated``).\n precision: numerical precision of the computation see ``jax.lax.Precision``\n for details.\n kernel_init: initializer for the convolutional kernel.\n bias_init: initializer for the bias.\n promote_dtype: function to promote the dtype of the arrays to the desired\n dtype. The function should accept a tuple of ``(inputs, kernel, bias)``\n and a ``dtype`` keyword argument, and return a tuple of arrays with the\n promoted dtype.\n preferred_element_type: Optional parameter controls the data type output by\n the convolution. This argument is passed to ``conv_general_dilated``\n function. See ``jax.lax.conv_general_dilated`` for details.\n rngs: rng key.\n kernel_metadata: Optional metadata dictionary to set when initializing\n the weight matrix.\n bias_metadata: Optional metadata dictionary to set when initializing\n the bias.\n \"\"\"\n\n def __init__(\n self,\n in_features: int,\n out_features: int,\n kernel_size: int | tp.Sequence[int],\n strides: tp.Union[None, int, tp.Sequence[int]] = 1,\n *,\n padding: PaddingLike = 'SAME',\n input_dilation: tp.Union[None, int, tp.Sequence[int]] = 1,\n kernel_dilation: tp.Union[None, int, tp.Sequence[int]] = 1,\n feature_group_count: int = 1,\n use_bias: bool = True,\n mask: tp.Optional[Array] = None,\n dtype: tp.Optional[Dtype] = None,\n param_dtype: Dtype = jnp.float32,\n precision: PrecisionLike = None,\n kernel_init: Initializer = default_kernel_init,\n bias_init: Initializer = default_bias_init,\n conv_general_dilated: ConvGeneralDilatedT = lax.conv_general_dilated,\n promote_dtype: PromoteDtypeFn = dtypes.promote_dtype,\n preferred_element_type: Dtype | None = None,\n rngs: rnglib.Rngs,\n kernel_metadata: tp.Mapping[str, tp.Any] = MappingProxyType({}),\n bias_metadata: tp.Mapping[str, tp.Any] = MappingProxyType({}),\n ):\n if isinstance(kernel_size, int):\n kernel_size = (kernel_size,)\n else:\n kernel_size = tuple(kernel_size)\n\n kernel_shape = kernel_size + (\n in_features // feature_group_count,\n out_features,\n )\n kernel_key = rngs.params()\n self.kernel_shape = kernel_shape\n self.kernel = nnx.Param(\n kernel_init(kernel_key, kernel_shape, param_dtype), **kernel_metadata\n )\n\n self.bias: nnx.Param[jax.Array] | None\n if use_bias:\n bias_shape = (out_features,)\n bias_key = rngs.params()\n self.bias = nnx.Param(\n bias_init(bias_key, bias_shape, param_dtype), **bias_metadata\n )\n else:\n self.bias = nnx.data(None)\n\n self.in_features = in_features\n self.out_features = out_features\n self.kernel_size = kernel_size\n self.strides = strides\n self.padding = padding\n self.input_dilation = input_dilation\n self.kernel_dilation = kernel_dilation\n self.feature_group_count = feature_group_count\n self.use_bias = use_bias\n self.mask = mask\n self.dtype = dtype\n self.param_dtype = param_dtype\n self.precision = precision\n self.conv_general_dilated = conv_general_dilated\n self.promote_dtype = promote_dtype\n self.preferred_element_type = preferred_element_type\n\n def __call__(self, inputs: Array, out_sharding=None) -> Array:\n \"\"\"Applies a (potentially unshared) convolution to the inputs.\n\n Args:\n inputs: input data with dimensions ``(*batch_dims, spatial_dims..., features)``.\n This is the channels-last convention, i.e. NHWC for a 2d convolution and\n NDHWC for a 3D convolution. Note: this is different from the input convention\n used by ``lax.conv_general_dilated``, which puts the spatial dimensions last.\n Note: If the input has more than 1 batch dimension, all batch dimensions\n are flattened into a single dimension for the convolution and restored\n before returning. In some cases directly vmap'ing the layer may yield\n better performance than this default flattening approach. If the input\n lacks a batch dimension it will be added for the convolution and removed\n n return, an allowance made to enable writing single-example code.\n out_sharding: Optional sharding specification (e.g.,\n ``jax.sharding.PartitionSpec``) for the output array. When using JAX's\n explicit sharding mode with a mesh context with ``AxisType.Explicit``.\n If ``None`` (default), the compiler automatically determines output\n sharding.\n\n Returns:\n The convolved data.\n \"\"\"\n\n assert isinstance(self.kernel_size, tuple)\n kernel_size = self.kernel_size\n\n def maybe_broadcast(\n x: tp.Optional[tp.Union[int, tp.Sequence[int]]],\n ) -> tuple[int, ...]:\n if x is None:\n # backward compatibility with using None as sentinel for\n # broadcast 1\n x = 1\n if isinstance(x, int):\n return (x,) * len(kernel_size)\n return tuple(x)\n\n # Combine all input batch dimensions into a single leading batch axis.\n num_batch_dimensions = inputs.ndim - (len(kernel_size) + 1)\n if num_batch_dimensions != 1:\n input_batch_shape = inputs.shape[:num_batch_dimensions]\n flat_input_shape = (-1,) + inputs.shape[\n num_batch_dimensions:\n ]\n inputs = jnp.reshape(inputs, flat_input_shape)\n\n # self.strides or (1,) * (inputs.ndim - 2)\n strides = maybe_broadcast(self.strides)\n input_dilation = maybe_broadcast(self.input_dilation)\n kernel_dilation = maybe_broadcast(self.kernel_dilation)\n\n padding_lax = canonicalize_padding(self.padding, len(kernel_size))\n if padding_lax in ('CIRCULAR', 'REFLECT'):\n assert isinstance(padding_lax, str)\n kernel_size_dilated = [\n (k - 1) * d + 1 for k, d in zip(kernel_size, kernel_dilation)\n ]\n zero_pad: list[tuple[int, int]] = [(0, 0)]\n pads = (\n zero_pad\n + [((k - 1) // 2, k // 2) for k in kernel_size_dilated]\n + [(0, 0)]\n )\n padding_mode = {'CIRCULAR': 'wrap', 'REFLECT': 'reflect'}[padding_lax]\n inputs = jnp.pad(inputs, pads, mode=padding_mode)\n padding_lax = 'VALID'\n elif padding_lax == 'CAUSAL':\n if len(kernel_size) != 1:\n raise ValueError(\n 'Causal padding is only implemented for 1D convolutions.'\n )\n left_pad = kernel_dilation[0] * (kernel_size[0] - 1)\n pads = [(0, 0), (left_pad, 0), (0, 0)]\n inputs = jnp.pad(inputs, pads)\n padding_lax = 'VALID'\n\n dimension_numbers = _conv_dimension_numbers(inputs.shape)\n\n # One shared convolutional kernel for all pixels in the output.\n assert self.in_features % self.feature_group_count == 0\n\n if self.mask is not None and self.mask.shape != self.kernel_shape:\n raise ValueError(\n 'Mask needs to have the same shape as weights. '\n f'Shapes are: {self.mask.shape}, {self.kernel_shape}'\n )\n\n kernel = self.kernel[...]\n\n if self.mask is not None:\n kernel *= self.mask\n\n bias = self.bias[...] if self.bias is not None else None\n\n inputs, kernel, bias = self.promote_dtype(\n (inputs, kernel, bias), dtype=self.dtype\n )\n\n # We use conv_kwargs for BC compatibility with\n # user custom self.conv_general_dilated method which may not have\n # preferred_element_type argument to avoid breaking\n # existing code\n conv_kwargs = {'out_sharding': out_sharding}\n if self.preferred_element_type is not None:\n conv_kwargs[\"preferred_element_type\"] = self.preferred_element_type\n\n y = self.conv_general_dilated(\n inputs,\n kernel,\n strides,\n padding_lax,\n lhs_dilation=input_dilation,\n rhs_dilation=kernel_dilation,\n dimension_numbers=dimension_numbers,\n feature_group_count=self.feature_group_count,\n precision=self.precision,\n **conv_kwargs,\n )\n\n if self.use_bias:\n bias = bias.reshape((1,) * (y.ndim - bias.ndim) + bias.shape) # type: ignore\n y += bias\n\n if num_batch_dimensions != 1:\n output_shape = input_batch_shape + y.shape[1:]\n y = jnp.reshape(y, output_shape)\n return y\n\n\nclass ConvTranspose(Module):\n \"\"\"Convolution Module wrapping ``lax.conv_transpose``.\n\n **Note:** The `padding` argument behaves differently from PyTorch; see the argument description below.\n\n Example usage::\n\n >>> from flax import nnx\n >>> import jax.numpy as jnp\n\n >>> rngs = nnx.Rngs(0)\n >>> x = jnp.ones((1, 8, 3))\n\n >>> # valid padding\n >>> layer = nnx.ConvTranspose(in_features=3, out_features=4, kernel_size=(3,),\n ... padding='VALID', rngs=rngs)\n >>> layer.kernel.shape\n (3, 3, 4)\n >>> layer.bias.shape\n (4,)\n >>> out = layer(x)\n >>> out.shape\n (1, 10, 4)\n\n >>> # circular padding with stride 2\n >>> layer = nnx.ConvTranspose(in_features=3, out_features=4, kernel_size=(6, 6),\n ... strides=(2, 2), padding='CIRCULAR',\n ... transpose_kernel=True, rngs=rngs)\n >>> layer.kernel.shape\n (6, 6, 4, 3)\n >>> layer.bias.shape\n (4,)\n >>> out = layer(jnp.ones((1, 15, 15, 3)))\n >>> out.shape\n (1, 30, 30, 4)\n\n >>> # apply lower triangle mask\n >>> mask = jnp.tril(jnp.ones((3, 3, 4)))\n >>> layer = nnx.Conv(in_features=3, out_features=4, kernel_size=(3,),\n ... mask=mask, padding='VALID', rngs=rngs)\n >>> out = layer(x)\n\n Args:\n in_features: int or tuple with number of input features.\n out_features: int or tuple with number of output features.\n kernel_size: shape of the convolutional kernel. For 1D convolution,\n the kernel size can be passed as an integer, which will be interpreted\n as a tuple of the single integer. For all other cases, it must be a\n sequence of integers.\n strides: an integer or a sequence of ``n`` integers, representing the\n inter-window strides (default: 1).\n padding: either a string indicating a specialized padding mode,\n or a sequence of ``n`` ``(low, high)`` integer pairs that give the padding to apply before and after each\n spatial dimension. A single int is interpreted as applying the same padding\n in all dims and a single int in a sequence causes the same padding\n to be used on both sides.\n\n **Note that this behavior is different from\n PyTorch**. In PyTorch, the padding argument effectively adds ``dilation * (kernel_size - 1) - padding``\n amount of zero padding to the input instead. This is set so that when ``torch.Conv2d`` and ``torch.ConvTranspose2d``\n are initialized with the same parameters, they are inverses of each other in regard to the input and output shapes.\n ``nnx.Conv`` and ``nnx.ConvTranspose`` do *not* have this behavior; if you want a ``nnx.ConvTranspose`` layer\n to invert the shape change produced by a ``nnx.Conv`` layer with a given padding and dilation, you should explicitly pass\n ``dilation * (kernel_size - 1) - padding`` as the `padding` argument to the ``nnx.ConvTranspose`` layer.\n\n Strings for specifying padding modes can be one of the following:\n\n - ``VALID`` adds ``dilation * (kernel_size - 1)`` padding to all dimensions. This is set so that a\n ``nnx.Conv`` layer with ``VALID`` padding would produce the inverse shape transformation.\n\n - ``SAME`` pads the input so that the output shape is the same as the input shape.\n\n - ``CIRCULAR`` pads the input with periodic boundary conditions.\n\n - ``CAUSAL`` padding for a 1D convolution will left-pad the convolution axis, resulting in same-sized output.\n\n kernel_dilation: an integer or a sequence of ``n`` integers, giving the\n dilation factor to apply in each spatial dimension of the convolution\n kernel (default: 1). Convolution with kernel dilation\n is also known as 'atrous convolution'.\n use_bias: whether to add a bias to the output (default: True).\n mask: Optional mask for the weights during masked convolution. The mask must\n be the same shape as the convolution weight matrix.\n dtype: the dtype of the computation (default: infer from input and params).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n precision: numerical precision of the computation see ``jax.lax.Precision``\n for details.\n kernel_init: initializer for the convolutional kernel.\n bias_init: initializer for the bias.\n transpose_kernel: if ``True`` flips spatial axes and swaps the input/output\n channel axes of the kernel.\n promote_dtype: function to promote the dtype of the arrays to the desired\n dtype. The function should accept a tuple of ``(inputs, kernel, bias)``\n and a ``dtype`` keyword argument, and return a tuple of arrays with the\n promoted dtype.\n preferred_element_type: Optional parameter controls the data type output by\n the transposed convolution. This argument is passed to\n ``jax.lax.conv_transpose`` function. See ``jax.lax.conv_transpose``\n for details.\n rngs: rng key.\n kernel_metadata: Optional metadata dictionary to set when initializing\n the weight matrix.\n bias_metadata: Optional metadata dictionary to set when initializing\n the bias.\n \"\"\"\n\n def __init__(\n self,\n in_features: int,\n out_features: int,\n kernel_size: int | tp.Sequence[int],\n strides: int | tp.Sequence[int] | None = None,\n *,\n padding: PaddingLike = 'SAME',\n kernel_dilation: int | tp.Sequence[int] | None = None,\n use_bias: bool = True,\n mask: Array | None = None,\n dtype: Dtype | None = None,\n param_dtype: Dtype = jnp.float32,\n precision: PrecisionLike | None = None,\n kernel_init: Initializer = default_kernel_init,\n bias_init: Initializer = default_bias_init,\n transpose_kernel: bool = False,\n promote_dtype: PromoteDtypeFn = dtypes.promote_dtype,\n preferred_element_type: Dtype | None = None,\n rngs: rnglib.Rngs,\n kernel_metadata: tp.Mapping[str, tp.Any] = MappingProxyType({}),\n bias_metadata: tp.Mapping[str, tp.Any] = MappingProxyType({}),\n ):\n if isinstance(kernel_size, int):\n kernel_size = (kernel_size,)\n else:\n kernel_size = tuple(kernel_size)\n\n self.kernel_size = kernel_size\n self.in_features = in_features\n self.out_features = out_features\n self.strides = strides\n self.padding = padding\n self.kernel_dilation = kernel_dilation\n self.use_bias = use_bias\n self.mask = mask\n self.dtype = dtype\n self.param_dtype = param_dtype\n self.precision = precision\n self.transpose_kernel = transpose_kernel\n self.promote_dtype = promote_dtype\n self.preferred_element_type = preferred_element_type\n\n if self.transpose_kernel:\n kernel_shape = kernel_size + (self.out_features, in_features)\n else:\n kernel_shape = kernel_size + (in_features, self.out_features)\n\n self.kernel_shape = kernel_shape\n self.kernel = nnx.Param(\n kernel_init(rngs.params(), kernel_shape, self.param_dtype), **kernel_metadata\n )\n\n self.bias: nnx.Param | None\n if self.use_bias:\n self.bias = nnx.Param(\n bias_init(rngs.params(), (self.out_features,), self.param_dtype), **bias_metadata\n )\n else:\n self.bias = nnx.data(None)\n\n def __call__(self, inputs: Array) -> Array:\n \"\"\"Applies a transposed convolution to the inputs.\n\n Behaviour mirrors of ``jax.lax.conv_transpose``.\n\n Args:\n inputs: input data with dimensions ``(*batch_dims, spatial_dims..., features).\n This is the channels-last convention, i.e. NHWC for a 2d convolution and NDHWC\n for a 3D convolution. Note: this is different from the input convention used by\n ``lax.conv_general_dilated``, which puts the spatial dimensions last.\n Note: If the input has more than 1 batch dimension, all batch dimensions\n are flattened into a single dimension for the convolution and restored\n before returning. In some cases directly vmap'ing the layer may yield\n better performance than this default flattening approach. If the input\n lacks a batch dimension it will be added for the convolution and removed\n n return, an allowance made to enable writing single-example code.\n\n Returns:\n The convolved data.\n \"\"\"\n kernel_size = self.kernel_size\n\n def maybe_broadcast(\n x: tp.Optional[tp.Union[int, tp.Sequence[int]]],\n ) -> tuple[int, ...]:\n if x is None:\n # backward compatibility with using None as sentinel for\n # broadcast 1\n x = 1\n if isinstance(x, int):\n return (x,) * len(kernel_size)\n return tuple(x)\n\n # Combine all input batch dimensions into a single leading batch axis.\n num_batch_dimensions = inputs.ndim - (len(kernel_size) + 1)\n if num_batch_dimensions != 1:\n input_batch_shape = inputs.shape[:num_batch_dimensions]\n flat_input_shape = (-1,) + inputs.shape[\n num_batch_dimensions:\n ]\n inputs = jnp.reshape(inputs, flat_input_shape)\n\n strides = maybe_broadcast(self.strides)\n kernel_dilation = maybe_broadcast(self.kernel_dilation)\n\n kernel_shape = self.kernel_shape\n\n if self.mask is not None and self.mask.shape != kernel_shape:\n raise ValueError(\n 'Mask needs to have the same shape as weights. '\n f'Shapes are: {self.mask.shape}, {kernel_shape}'\n )\n\n kernel = self.kernel[...]\n\n if self.mask is not None:\n kernel *= self.mask\n\n padding_lax = canonicalize_padding(self.padding, len(kernel_size))\n if padding_lax == 'CIRCULAR':\n padding_lax = 'VALID'\n\n bias = self.bias[...] if self.bias is not None else None\n\n inputs, kernel, bias = self.promote_dtype(\n (inputs, kernel, bias), dtype=self.dtype\n )\n\n y = lax.conv_transpose(\n inputs,\n kernel,\n strides,\n padding_lax,\n rhs_dilation=kernel_dilation,\n transpose_kernel=self.transpose_kernel,\n precision=self.precision,\n preferred_element_type=self.preferred_element_type\n )\n\n if self.padding == 'CIRCULAR':\n # For circular padding, we need to identify the size of the final output\n # (\"period\") along each spatial dimension, pad each dimension to an\n # integer number of periods, and wrap the array periodically around each\n # dimension. Padding should be done in such a way that the start of the\n # original input data inside the padded array is located at integer\n # number of periods - otherwise the result would be circularly shifted.\n\n # Compute period along each spatial dimension - it's input size scaled\n # by the stride.\n scaled_x_dims = [\n x_dim * stride\n for x_dim, stride in zip(jnp.shape(inputs)[1:-1], strides)\n ]\n # Compute difference between the current size of y and the final output\n # size, and complement this difference to 2 * period - that gives how\n # much we need to pad.\n size_diffs = [\n -(y_dim - x_dim) % (2 * x_dim)\n for y_dim, x_dim in zip(y.shape[1:-1], scaled_x_dims)\n ]\n if self.transpose_kernel:\n # If the kernel is transposed, the \"+1\" is put on the right to\n # mirror the regular convolution. If the same kernel parameters are used\n # as for Conv, this layer then computes the proper transpose convolution.\n total_pad = [\n (size_diff // 2, (size_diff + 1) // 2) for size_diff in size_diffs\n ]\n else:\n # Divide the padding equally between left and right. The choice to put\n # \"+1\" on the left (and not on the right) represents a convention for\n # aligning even-sized kernels.\n total_pad = [\n ((size_diff + 1) // 2, size_diff // 2) for size_diff in size_diffs\n ]\n y = jnp.pad(y, [(0, 0)] + total_pad + [(0, 0)])\n # Wrap the result periodically around each spatial dimension,\n # one by one.\n for i in range(1, y.ndim - 1):\n y = y.reshape(\n y.shape[:i] + (-1, scaled_x_dims[i - 1]) + y.shape[i + 1 :]\n )\n y = y.sum(axis=i)\n\n if self.use_bias:\n y += jnp.reshape(bias, (1,) * (y.ndim - 1) + (-1,)) # type: ignore\n\n if num_batch_dimensions != 1:\n output_shape = input_batch_shape + y.shape[1:]\n y = jnp.reshape(y, output_shape)\n\n return y\n\n\ndefault_embed_init = initializers.variance_scaling(\n 1.0, 'fan_in', 'normal', out_axis=0\n)\n\n\nclass Embed(Module):\n \"\"\"Embedding Module.\n\n Example usage::\n\n >>> from flax import nnx\n >>> import jax.numpy as jnp\n\n >>> layer = nnx.Embed(num_embeddings=5, features=3, rngs=nnx.Rngs(0))\n >>> nnx.state(layer)\n State({\n 'embedding': Param( # 15 (60 B)\n value=Array([[ 0.57966787, -0.523274 , -0.43195742],\n [-0.676289 , -0.50300646, 0.33996582],\n [ 0.41796115, -0.59212935, 0.95934135],\n [-1.0917838 , -0.7441663 , 0.07713798],\n [-0.66570747, 0.13815777, 1.007365 ]], dtype=float32)\n )\n })\n >>> # get the first three and last three embeddings\n >>> indices_input = jnp.array([[0, 1, 2], [-1, -2, -3]])\n >>> layer(indices_input)\n Array([[[ 0.57966787, -0.523274 , -0.43195742],\n [-0.676289 , -0.50300646, 0.33996582],\n [ 0.41796115, -0.59212935, 0.95934135]],\n \n [[-0.66570747, 0.13815777, 1.007365 ],\n [-1.0917838 , -0.7441663 , 0.07713798],\n [ 0.41796115, -0.59212935, 0.95934135]]], dtype=float32)\n\n A parameterized function from integers [0, ``num_embeddings``) to\n ``features``-dimensional vectors. This ``Module`` will create an ``embedding``\n matrix with shape ``(num_embeddings, features)``. When calling this layer,\n the input values will be used to 0-index into the ``embedding`` matrix.\n Indexing on a value greater than or equal to ``num_embeddings`` will result\n in ``nan`` values. When ``num_embeddings`` equals to 1, it will\n broadcast the ``embedding`` matrix to input shape with ``features``\n dimension appended.\n\n Args:\n num_embeddings: number of embeddings / vocab size.\n features: number of feature dimensions for each embedding.\n dtype: the dtype of the embedding vectors (default: same as embedding).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n embedding_init: embedding initializer.\n promote_dtype: function to promote the dtype of the arrays to the desired\n dtype. The function should accept a tuple of ``(embedding,)`` during ``__call__``\n or ``(query, embedding)`` during ``attend``, and a ``dtype`` keyword argument,\n and return a tuple of arrays with the promoted dtype.\n rngs: rng key.\n embedding_metadata: Optional metadata dictionary to set when initializing\n the embedding matrix.\n \"\"\"\n\n def __init__(\n self,\n num_embeddings: int,\n features: int,\n *,\n dtype: tp.Optional[Dtype] = None,\n param_dtype: Dtype = jnp.float32,\n embedding_init: Initializer = default_embed_init,\n promote_dtype: PromoteDtypeFn = dtypes.promote_dtype,\n rngs: rnglib.Rngs,\n embedding_metadata: tp.Mapping[str, tp.Any] = MappingProxyType({}),\n ):\n self.embedding = nnx.Param(\n embedding_init(rngs.params(), (num_embeddings, features), param_dtype),\n **embedding_metadata,\n )\n\n self.num_embeddings = num_embeddings\n self.features = features\n self.dtype = dtype or self.embedding.dtype\n self.param_dtype = param_dtype\n self.promote_dtype = promote_dtype\n\n def __call__(self, inputs: Array, out_sharding=None) -> Array:\n \"\"\"Embeds the inputs along the last dimension.\n\n Args:\n inputs: input data, all dimensions are considered batch dimensions.\n Values in the input array must be integers.\n\n Returns:\n Output which is embedded input data. The output shape follows the input,\n with an additional ``features`` dimension appended.\n \"\"\"\n if not jnp.issubdtype(inputs.dtype, jnp.integer):\n raise ValueError('Input type must be an integer or unsigned integer.')\n # Use take because fancy indexing numpy arrays with JAX indices does not\n # work correctly.\n (embedding,) = self.promote_dtype(\n (self.embedding[...],), dtype=self.dtype, inexact=False\n )\n if self.num_embeddings == 1:\n return jnp.broadcast_to(embedding, inputs.shape + (self.features,))\n if out_sharding is not None:\n # Use auto_axes to handle out_sharding as jnp.take does not support it.\n take_fn = lambda embedding, inputs: jnp.take(embedding, inputs, axis=0)\n sharded_take = jax.sharding.auto_axes(out_sharding=out_sharding)(take_fn)\n return sharded_take(embedding, inputs)\n return jnp.take(embedding, inputs, axis=0)\n\n def attend(self, query: Array, out_sharding=None) -> Array:\n \"\"\"Attend over the embedding using a query array.\n\n Args:\n query: array with last dimension equal the feature depth ``features`` of the\n embedding.\n out_sharding: Optional sharding specification (e.g.,\n ``jax.sharding.PartitionSpec``) for the output array. When using JAX's\n explicit sharding mode with a mesh context with ``AxisType.Explicit``.\n If ``None`` (default), the compiler automatically determines output\n sharding.\n\n Returns:\n An array with final dim ``num_embeddings`` corresponding to the batched\n inner-product of the array of query vectors against each embedding.\n Commonly used for weight-sharing between embeddings and logit transform\n in NLP models.\n \"\"\"\n query, embedding = self.promote_dtype(\n (query, self.embedding[...]), dtype=self.dtype\n )\n return jnp.dot(query, embedding.T, out_sharding=out_sharding)\n" }, { "path": "flax/nnx/nn/lora.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nimport typing as tp\nfrom types import MappingProxyType\n\nfrom flax.nnx import rnglib, variablelib\nfrom flax.nnx.module import Module\nfrom flax.nnx.nn import initializers, dtypes\nfrom flax.nnx.nn.linear import Linear\nfrom flax.typing import Dtype, Initializer, PromoteDtypeFn\nimport jax\nimport jax.numpy as jnp\n\nArray = jax.Array\nAxis = int\nSize = int\nA = tp.TypeVar('A')\n\ndefault_a_initializer = initializers.he_uniform()\ndefault_b_initializer = initializers.zeros\n\n\nclass LoRAParam(variablelib.Param[A]):\n pass\n\n\nclass LoRA(Module):\n \"\"\"A standalone LoRA layer.\n\n Example usage::\n\n >>> from flax import nnx\n >>> import jax, jax.numpy as jnp\n >>> layer = nnx.LoRA(3, 2, 4, rngs=nnx.Rngs(0))\n >>> layer.lora_a.shape\n (3, 2)\n >>> layer.lora_b.shape\n (2, 4)\n >>> # Wrap around existing layer\n >>> linear = nnx.Linear(3, 4, rngs=nnx.Rngs(0))\n >>> wrapper = nnx.LoRA(3, 2, 4, base_module=linear, rngs=nnx.Rngs(1))\n >>> assert wrapper.base_module == linear\n >>> wrapper.lora_a.shape\n (3, 2)\n >>> layer.lora_b.shape\n (2, 4)\n >>> y = layer(jnp.ones((16, 3)))\n >>> y.shape\n (16, 4)\n\n Args:\n in_features: the number of input features.\n lora_rank: the rank of the LoRA dimension.\n out_features: the number of output features.\n base_module: a base module to call and substitute, if possible.\n dtype: the dtype of the computation (default: infer from input and params).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n precision: numerical precision of the computation see `jax.lax.Precision`\n for details.\n a_initializer: initializer function for the fan-in matrices. Default to\n `he_uniform`.\n b_initializer: initializer function for the fan-out matrices. Default to\n `zero initializer`.\n lora_param_type: the type of the LoRA params.\n promote_dtype: function to promote the dtype of all input array arguments\n (including Variables accessed through ``self``) to the desired dtype. The\n function should accept a tuple of ``(inputs, lora_a, lora_b)`` and a ``dtype``\n keyword argument, and return a tuple of arrays with the promoted dtype.\n rngs: rng key.\n a_metadata: Optional metadata dictionary to set when initializing\n the fan-in matrices.\n b_metadata: Optional metadata dictionary to set when initializing\n the fan-out matrices.\n \"\"\"\n\n def __init__(\n self,\n in_features: int,\n lora_rank: int,\n out_features: int,\n *,\n base_module: tp.Optional[Module] = None,\n dtype: tp.Optional[Dtype] = None,\n param_dtype: Dtype = jnp.float32,\n a_initializer: Initializer = default_a_initializer,\n b_initializer: Initializer = default_b_initializer,\n lora_param_type: tp.Type[variablelib.Variable] = LoRAParam,\n promote_dtype: PromoteDtypeFn = dtypes.promote_dtype,\n rngs: rnglib.Rngs,\n a_metadata: tp.Mapping[str, tp.Any] = MappingProxyType({}),\n b_metadata: tp.Mapping[str, tp.Any] = MappingProxyType({}),\n ):\n self.in_features = in_features\n self.out_features = out_features\n self.dtype = dtype\n self.param_dtype = param_dtype\n self.lora_param_type = lora_param_type\n self.base_module = base_module\n self.promote_dtype = promote_dtype\n\n self.lora_a = lora_param_type(\n a_initializer(rngs.params(), (in_features, lora_rank), param_dtype),\n **a_metadata,\n )\n self.lora_b = lora_param_type(\n b_initializer(rngs.params(), (lora_rank, out_features), param_dtype),\n **b_metadata,\n )\n\n def __call__(self, x: jax.Array):\n x, lora_a, lora_b = self.promote_dtype(\n (x, self.lora_a[...], self.lora_b[...]), dtype=self.dtype\n )\n out = x @ lora_a @ lora_b\n if self.base_module is not None:\n if not callable(self.base_module):\n raise ValueError('`self.base_module` must be callable.')\n out += self.base_module(x)\n return out\n\n\nclass LoRALinear(Linear):\n \"\"\"An `nnx.Linear` layer in which the output will be LoRAified.\n\n The model state structure will be compatible with that of Linear.\n\n Example usage::\n\n >>> from flax import nnx\n >>> import jax, jax.numpy as jnp\n >>> linear = nnx.Linear(3, 4, rngs=nnx.Rngs(0))\n >>> lora_linear = nnx.LoRALinear(3, 4, lora_rank=2, rngs=nnx.Rngs(0))\n >>> linear.kernel.shape\n (3, 4)\n >>> lora_linear.kernel.shape\n (3, 4)\n >>> lora_linear.lora.lora_a.shape\n (3, 2)\n >>> jnp.allclose(linear.kernel[...], lora_linear.kernel[...])\n Array(True, dtype=bool)\n >>> y = lora_linear(jnp.ones((16, 3)))\n >>> y.shape\n (16, 4)\n\n Args:\n in_features: the number of input features.\n out_features: the number of output features.\n lora_rank: the rank of the LoRA dimension.\n base_module: a base module to call and substitute, if possible.\n dtype: the dtype of the computation (default: infer from input and params).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n precision: numerical precision of the computation see `jax.lax.Precision`\n for details.\n a_initializer: initializer function for the fan-in matrices. Default to\n `he_uniform`.\n b_initializer: initializer function for the fan-out matrices. Default to\n `zero initializer`.\n lora_param_type: the type of the LoRA params.\n lora_promote_dtype: function to promote the dtype for the LoRA submodule.\n a_metadata: Optional metadata dictionary to set when initializing\n the fan-in matrices.\n b_metadata: Optional metadata dictionary to set when initializing\n the fan-out matrices.\n \"\"\"\n\n def __init__(\n self,\n in_features: int,\n out_features: int,\n *,\n lora_rank: int,\n lora_dtype: tp.Optional[Dtype] = None,\n lora_param_dtype: Dtype = jnp.float32,\n a_initializer: Initializer = default_a_initializer,\n b_initializer: Initializer = default_b_initializer,\n lora_param_type: tp.Type[variablelib.Variable] = LoRAParam,\n lora_promote_dtype: PromoteDtypeFn = dtypes.promote_dtype,\n rngs: rnglib.Rngs,\n a_metadata: tp.Mapping[str, tp.Any] = MappingProxyType({}),\n b_metadata: tp.Mapping[str, tp.Any] = MappingProxyType({}),\n **kwargs,\n ):\n super().__init__(in_features, out_features, rngs=rngs, **kwargs)\n self.lora = LoRA(\n in_features,\n lora_rank,\n out_features,\n dtype=lora_dtype,\n param_dtype=lora_param_dtype,\n a_initializer=a_initializer,\n b_initializer=b_initializer,\n lora_param_type=lora_param_type,\n promote_dtype=lora_promote_dtype,\n rngs=rngs,\n a_metadata=a_metadata,\n b_metadata=b_metadata,\n )\n\n def __call__(self, x: jax.Array, out_sharding = None):\n y = super().__call__(x, out_sharding=out_sharding)\n y += self.lora(x)\n return y\n" }, { "path": "flax/nnx/nn/normalization.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport typing as tp\nfrom types import MappingProxyType\n\nimport jax\nimport jax.numpy as jnp\nfrom jax import lax\n\nfrom flax import nnx\nfrom flax.nnx import rnglib\nfrom flax.nnx.module import Module, first_from\nfrom flax.nnx.nn import dtypes, initializers\nfrom flax.typing import (\n Array,\n Dtype,\n Initializer,\n Axes,\n PromoteDtypeFn,\n)\n\n\ndef _canonicalize_axes(rank: int, axes: Axes) -> tp.Tuple[int, ...]:\n \"\"\"Returns a tuple of deduplicated, sorted, and positive axes.\"\"\"\n if not isinstance(axes, tp.Iterable):\n axes = (axes,)\n return tuple({rank + axis if axis < 0 else axis for axis in axes})\n\n\ndef _abs_sq(x):\n \"\"\"Computes the elementwise square of the absolute value |x|^2.\"\"\"\n if jnp.iscomplexobj(x):\n return lax.square(lax.real(x)) + lax.square(lax.imag(x))\n else:\n return lax.square(x)\n\n\ndef _compute_stats(\n x: Array,\n axes: Axes,\n dtype: tp.Optional[Dtype],\n axis_name: tp.Optional[str] = None,\n axis_index_groups: tp.Any = None,\n use_mean: bool = True,\n use_fast_variance: bool = True,\n mask: tp.Optional[Array] = None,\n) -> tuple[Array, Array]:\n \"\"\"Computes mean and variance statistics.\n\n This implementation takes care of a few important details:\n - Computes in float32 precision for stability in half precision training.\n - If ``use_fast_variance`` is ``True``, mean and variance are computed using\n Var = E[|x|^2] - |E[x]|^2, instead of Var = E[|x - E[x]|^2]), in a single\n XLA fusion.\n - Clips negative variances to zero which can happen due to\n roundoff errors. This avoids downstream NaNs.\n - Supports averaging across a parallel axis and subgroups of a parallel axis\n with a single ``lax.pmean`` call to avoid latency.\n\n Arguments:\n x: Input array.\n axes: The axes in ``x`` to compute mean and variance statistics for.\n dtype: Optional dtype specifying the minimal precision. Statistics are\n always at least float32 for stability (default: dtype of x).\n axis_name: Optional name for the pmapped axis to compute mean over. Note,\n this is only used for pmap and shard map. For SPMD jit, you do not need to\n manually synchronize. Just make sure that the axes are correctly annotated\n and XLA:SPMD will insert the necessary collectives.\n axis_index_groups: Optional groups of indices within that named axis.\n use_mean: If true, calculate the mean from the input and use it when\n computing the variance. If false, set the mean to zero and compute the\n variance without subtracting the mean.\n use_fast_variance: If true, use a faster, but less numerically stable,\n calculation for the variance.\n mask: Binary array of shape broadcastable to ``inputs`` tensor, indicating\n the positions for which the mean and variance should be computed.\n\n Returns:\n A pair ``(mean, var)``.\n \"\"\"\n if dtype is None:\n dtype = jnp.result_type(x)\n # promote x to at least float32, this avoids half precision computation\n # but preserves double or complex floating points\n dtype = jnp.promote_types(dtype, jnp.float32)\n x = jnp.asarray(x, dtype)\n axes = _canonicalize_axes(x.ndim, axes)\n\n def maybe_distributed_mean(*xs, mask=None):\n mus = tuple(x.mean(axes, where=mask) for x in xs)\n if axis_name is None:\n return mus if len(xs) > 1 else mus[0]\n else:\n # In the distributed case we stack multiple arrays to speed comms.\n if len(xs) > 1:\n reduced_mus = lax.pmean(\n jnp.stack(mus, axis=0),\n axis_name,\n axis_index_groups=axis_index_groups,\n )\n return tuple(reduced_mus[i] for i in range(len(xs)))\n else:\n return lax.pmean(mus[0], axis_name, axis_index_groups=axis_index_groups)\n\n if use_mean:\n if use_fast_variance:\n mu, mu2 = maybe_distributed_mean(x, _abs_sq(x), mask=mask)\n # mean2 - _abs_sq(mean) is not guaranteed to be non-negative due\n # to floating point round-off errors.\n var = jnp.maximum(0.0, mu2 - _abs_sq(mu))\n else:\n mu = maybe_distributed_mean(x, mask=mask)\n var = maybe_distributed_mean(\n _abs_sq(x - jnp.expand_dims(mu, axes)), mask=mask\n )\n else:\n var = maybe_distributed_mean(_abs_sq(x), mask=mask)\n mu = jnp.zeros_like(var)\n return mu, var\n\n\ndef _normalize(\n x: Array,\n mean: Array,\n var: Array,\n scale: tp.Optional[Array],\n bias: tp.Optional[Array],\n reduction_axes: Axes,\n feature_axes: Axes,\n dtype: tp.Optional[Dtype],\n epsilon: float,\n):\n \"\"\" \"Normalizes the input of a normalization layer and optionally applies a learned scale and bias.\n\n Arguments:\n x: The input.\n mean: Mean to use for normalization.\n var: Variance to use for normalization.\n reduction_axes: The axes in ``x`` to reduce.\n feature_axes: Axes containing features. A separate bias and scale is learned\n for each specified feature.\n dtype: The dtype of the result (default: infer from input and params).\n epsilon: Normalization epsilon.\n\n Returns:\n The normalized input.\n \"\"\"\n reduction_axes = _canonicalize_axes(x.ndim, reduction_axes)\n feature_axes = _canonicalize_axes(x.ndim, feature_axes)\n stats_shape = list(x.shape)\n for axis in reduction_axes:\n stats_shape[axis] = 1\n mean = mean.reshape(stats_shape)\n var = var.reshape(stats_shape)\n feature_shape = [1] * x.ndim\n for ax in feature_axes:\n feature_shape[ax] = x.shape[ax]\n y = x - mean\n mul = lax.rsqrt(var + epsilon)\n args = [x]\n if scale is not None:\n scale = scale.reshape(feature_shape)\n mul *= scale\n args.append(scale)\n y *= mul\n if bias is not None:\n bias = bias.reshape(feature_shape)\n y += bias\n args.append(bias)\n dtype = dtypes.canonicalize_dtype(*args, dtype=dtype)\n return jnp.asarray(y, dtype)\n\n\ndef _l2_normalize(x, axis=None, eps=1e-12):\n \"\"\"Normalizes along dimension `axis` using an L2 norm.\n This specialized function exists for numerical stability reasons.\n Args:\n x: An input ndarray.\n axis: Dimension along which to normalize, e.g. `1` to separately normalize\n vectors in a batch. Passing `None` views `t` as a flattened vector when\n calculating the norm (equivalent to Frobenius norm).\n eps: Epsilon to avoid dividing by zero.\n Returns:\n An array of the same shape as 'x' L2-normalized along 'axis'.\n \"\"\"\n return x * jax.lax.rsqrt((x * x).sum(axis=axis, keepdims=True) + eps)\n\n\nclass BatchNorm(Module):\n \"\"\"BatchNorm Module.\n\n To calculate the batch norm on the input and update the batch statistics,\n call the :func:`train` method (or pass in ``use_running_average=False`` in\n the constructor or during call time).\n\n To use the stored batch statistics' running average, call the :func:`eval`\n method (or pass in ``use_running_average=True`` in the constructor or\n during call time).\n\n Example usage::\n\n >>> from flax import nnx\n >>> import jax, jax.numpy as jnp\n\n >>> x = jax.random.normal(jax.random.key(0), (5, 6))\n >>> layer = nnx.BatchNorm(num_features=6, momentum=0.9, epsilon=1e-5,\n ... dtype=jnp.float32, rngs=nnx.Rngs(0))\n >>> jax.tree.map(jnp.shape, nnx.state(layer))\n State({\n 'bias': Param(\n value=(6,)\n ),\n 'mean': BatchStat(\n value=(6,)\n ),\n 'scale': Param(\n value=(6,)\n ),\n 'var': BatchStat(\n value=(6,)\n )\n })\n\n >>> # calculate batch norm on input and update batch statistics\n >>> layer.train()\n >>> y = layer(x)\n >>> batch_stats1 = nnx.clone(nnx.state(layer, nnx.BatchStat)) # keep a copy\n >>> y = layer(x)\n >>> batch_stats2 = nnx.state(layer, nnx.BatchStat)\n >>> assert (batch_stats1['mean'][...] != batch_stats2['mean'][...]).all()\n >>> assert (batch_stats1['var'][...] != batch_stats2['var'][...]).all()\n\n >>> # use stored batch statistics' running average\n >>> layer.eval()\n >>> y = layer(x)\n >>> batch_stats3 = nnx.state(layer, nnx.BatchStat)\n >>> assert (batch_stats2['mean'][...] == batch_stats3['mean'][...]).all()\n >>> assert (batch_stats2['var'][...] == batch_stats3['var'][...]).all()\n\n Args:\n num_features: the number of input features.\n use_running_average: if True, the stored batch statistics will be\n used instead of computing the batch statistics on the input.\n axis: the feature or non-batch axis of the input.\n momentum: decay rate for the exponential moving average of\n the batch statistics.\n epsilon: a small float added to variance to avoid dividing by zero.\n dtype: the dtype of the result (default: infer from input and params).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n use_bias: if True, bias (beta) is added.\n use_scale: if True, multiply by scale (gamma).\n When the next layer is linear (also e.g. nn.relu), this can be disabled\n since the scaling will be done by the next layer.\n bias_init: initializer for bias, by default, zero.\n scale_init: initializer for scale, by default, one.\n axis_name: the axis name used to combine batch statistics from multiple\n devices. See ``jax.pmap`` for a description of axis names (default: None).\n axis_index_groups: groups of axis indices within that named axis\n representing subsets of devices to reduce over (default: None). For\n example, ``[[0, 1], [2, 3]]`` would independently batch-normalize over\n the examples on the first two and last two devices. See ``jax.lax.psum``\n for more details. This argument is currently not supported for SPMD jit.\n use_fast_variance: If true, use a faster, but less numerically stable,\n calculation for the variance.\n promote_dtype: function to promote the dtype of all input array arguments\n (including Variables accessed through ``self``) to the desired dtype. The\n function should accept a tuple of ``(inputs, mean, var, scale, bias)`` and\n a ``dtype`` keyword argument, and return a tuple of arrays with the promoted\n dtype.\n rngs: rng key.\n bias_metadata: Optional metadata dictionary to set when initializing\n the bias.\n scale_metadata: Optional metadata dictionary to set when initializing\n the scale.\n \"\"\"\n\n def __init__(\n self,\n num_features: int,\n *,\n use_running_average: bool = False,\n axis: int = -1,\n momentum: float = 0.99,\n epsilon: float = 1e-5,\n dtype: tp.Optional[Dtype] = None,\n param_dtype: Dtype = jnp.float32,\n use_bias: bool = True,\n use_scale: bool = True,\n bias_init: Initializer = initializers.zeros_init(),\n scale_init: Initializer = initializers.ones_init(),\n axis_name: tp.Optional[str] = None,\n axis_index_groups: tp.Any = None,\n use_fast_variance: bool = True,\n promote_dtype: PromoteDtypeFn = dtypes.promote_dtype,\n rngs: rnglib.Rngs,\n bias_metadata: tp.Mapping[str, tp.Any] = MappingProxyType({}),\n scale_metadata: tp.Mapping[str, tp.Any] = MappingProxyType({}),\n ):\n feature_shape = (num_features,)\n self.mean = nnx.BatchStat(jnp.zeros(feature_shape, jnp.float32))\n self.var = nnx.BatchStat(jnp.ones(feature_shape, jnp.float32))\n\n self.scale: nnx.Param[jax.Array] | None\n if use_scale:\n key = rngs.params()\n self.scale = nnx.Param(scale_init(key, feature_shape, param_dtype), **scale_metadata)\n else:\n self.scale = nnx.data(None)\n\n self.bias: nnx.Param[jax.Array] | None\n if use_bias:\n key = rngs.params()\n self.bias = nnx.Param(bias_init(key, feature_shape, param_dtype), **bias_metadata)\n else:\n self.bias = nnx.data(None)\n\n self.num_features = num_features\n self.use_running_average = use_running_average\n self.axis = axis\n self.momentum = momentum\n self.epsilon = epsilon\n self.dtype = dtype\n self.param_dtype = param_dtype\n self.use_bias = use_bias\n self.use_scale = use_scale\n self.axis_name = axis_name\n self.axis_index_groups = axis_index_groups\n self.use_fast_variance = use_fast_variance\n self.promote_dtype = promote_dtype\n\n def __call__(\n self,\n x,\n use_running_average: tp.Optional[bool] = None,\n *,\n mask: tp.Optional[jax.Array] = None,\n ):\n \"\"\"Normalizes the input using batch statistics.\n\n Args:\n x: the input to be normalized.\n use_running_average: if true, the stored batch statistics will be\n used instead of computing the batch statistics on the input. The\n ``use_running_average`` flag passed into the call method will take\n precedence over the ``use_running_average`` flag passed into the\n constructor.\n\n Returns:\n Normalized inputs (the same shape as inputs).\n \"\"\"\n\n use_running_average = first_from(\n use_running_average,\n self.use_running_average,\n error_msg=\"\"\"No `use_running_average` argument was provided to BatchNorm\n as either a __call__ argument, class attribute, or nnx.flag.\"\"\",\n )\n feature_axes = _canonicalize_axes(x.ndim, self.axis)\n reduction_axes = tuple(i for i in range(x.ndim) if i not in feature_axes)\n\n # Promote dtypes for input and all Variables\n scale = self.scale[...] if self.scale is not None else None\n bias = self.bias[...] if self.bias is not None else None\n x, mean, var, scale, bias = self.promote_dtype(\n (x, self.mean[...], self.var[...], scale, bias), dtype=self.dtype\n )\n if not use_running_average:\n mean, var = _compute_stats(\n x,\n reduction_axes,\n dtype=self.dtype,\n axis_name=self.axis_name,\n axis_index_groups=self.axis_index_groups,\n use_fast_variance=self.use_fast_variance,\n mask=mask,\n )\n # stop_gradient only for flax_array_ref\n if self.mean._can_update or self.var._can_update:\n stop_gradient = jax.lax.stop_gradient\n else:\n stop_gradient = lambda x: x\n\n self.mean[...] = stop_gradient(\n self.momentum * self.mean[...] + (1 - self.momentum) * mean\n )\n self.var[...] = stop_gradient(\n self.momentum * self.var[...] + (1 - self.momentum) * var\n )\n\n return _normalize(\n x,\n mean,\n var,\n scale,\n bias,\n reduction_axes,\n feature_axes,\n self.dtype,\n self.epsilon,\n )\n\n def set_view(\n self,\n use_running_average: bool | None = None,\n ):\n \"\"\"Class method used by ``nnx.view``.\n\n Args:\n use_running_average: if True, the stored batch statistics will be\n used instead of computing the batch statistics on the input.\n \"\"\"\n if use_running_average is not None:\n self.use_running_average = use_running_average\n\n\nclass LayerNorm(Module):\n \"\"\"Layer normalization (https://arxiv.org/abs/1607.06450).\n\n LayerNorm normalizes the activations of the layer for each given example in a\n batch independently, rather than across a batch like Batch Normalization.\n i.e. applies a transformation that maintains the mean activation within\n each example close to 0 and the activation standard deviation close to 1.\n\n Example usage::\n\n >>> from flax import nnx\n >>> import jax\n\n >>> x = jax.random.normal(jax.random.key(0), (3, 4, 5, 6))\n >>> layer = nnx.LayerNorm(num_features=6, rngs=nnx.Rngs(0))\n\n >>> nnx.state(layer)\n State({\n 'bias': Param( # 6 (24 B)\n value=Array([0., 0., 0., 0., 0., 0.], dtype=float32)\n ),\n 'scale': Param( # 6 (24 B)\n value=Array([1., 1., 1., 1., 1., 1.], dtype=float32)\n )\n })\n\n >>> y = layer(x)\n\n Args:\n num_features: the number of input features.\n epsilon: A small float added to variance to avoid dividing by zero.\n dtype: the dtype of the result (default: infer from input and params).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n use_bias: If True, bias (beta) is added.\n use_scale: If True, multiply by scale (gamma). When the next layer is linear\n (also e.g. nnx.relu), this can be disabled since the scaling will be done\n by the next layer.\n bias_init: Initializer for bias, by default, zero.\n scale_init: Initializer for scale, by default, one.\n reduction_axes: Axes for computing normalization statistics.\n feature_axes: Feature axes for learned bias and scaling.\n axis_name: the axis name used to combine batch statistics from multiple\n devices. See ``jax.pmap`` for a description of axis names (default: None).\n This is only needed if the model is subdivided across devices, i.e. the\n array being normalized is sharded across devices within a pmap.\n axis_index_groups: groups of axis indices within that named axis\n representing subsets of devices to reduce over (default: None). For\n example, ``[[0, 1], [2, 3]]`` would independently batch-normalize over\n the examples on the first two and last two devices. See ``jax.lax.psum``\n for more details. This argument is currently not supported for SPMD jit.\n use_fast_variance: If true, use a faster, but less numerically stable,\n calculation for the variance.\n promote_dtype: function to promote the dtype of all input array arguments\n (including Variables accessed through ``self``) to the desired dtype. The\n function should accept a tuple of ``(inputs, scale, bias)`` and a ``dtype``\n keyword argument, and return a tuple of arrays with the promoted dtype.\n rngs: rng key.\n bias_metadata: Optional metadata dictionary to set when initializing\n the bias.\n scale_metadata: Optional metadata dictionary to set when initializing\n the scale.\n \"\"\"\n\n def __init__(\n self,\n num_features: int,\n *,\n epsilon: float = 1e-6,\n dtype: tp.Optional[Dtype] = None,\n param_dtype: Dtype = jnp.float32,\n use_bias: bool = True,\n use_scale: bool = True,\n bias_init: Initializer = initializers.zeros_init(),\n scale_init: Initializer = initializers.ones_init(),\n reduction_axes: Axes = -1,\n feature_axes: Axes = -1,\n axis_name: tp.Optional[str] = None,\n axis_index_groups: tp.Any = None,\n use_fast_variance: bool = True,\n promote_dtype: PromoteDtypeFn = dtypes.promote_dtype,\n rngs: rnglib.Rngs,\n bias_metadata: tp.Mapping[str, tp.Any] = MappingProxyType({}),\n scale_metadata: tp.Mapping[str, tp.Any] = MappingProxyType({}),\n ):\n feature_shape = (num_features,)\n\n self.scale: nnx.Param[jax.Array] | None\n if use_scale:\n key = rngs.params()\n self.scale = nnx.Param(scale_init(key, feature_shape, param_dtype), **scale_metadata)\n else:\n self.scale = nnx.data(None)\n\n self.bias: nnx.Param[jax.Array] | None\n if use_bias:\n key = rngs.params()\n self.bias = nnx.Param(bias_init(key, feature_shape, param_dtype), **bias_metadata)\n else:\n self.bias = nnx.data(None)\n\n self.num_features = num_features\n self.epsilon = epsilon\n self.dtype = dtype\n self.param_dtype = param_dtype\n self.use_bias = use_bias\n self.use_scale = use_scale\n self.reduction_axes = reduction_axes\n self.feature_axes = feature_axes\n self.axis_name = axis_name\n self.axis_index_groups = axis_index_groups\n self.use_fast_variance = use_fast_variance\n self.promote_dtype = promote_dtype\n\n def __call__(self, x, *, mask: tp.Optional[jax.Array] = None):\n \"\"\"Applies layer normalization on the input.\n\n Args:\n x: the inputs\n\n Returns:\n Normalized inputs (the same shape as inputs).\n \"\"\"\n # Promote dtypes for input and all Variables\n scale = self.scale[...] if self.scale else None\n bias = self.bias[...] if self.bias else None\n x, scale, bias = self.promote_dtype(\n (x, scale, bias), dtype=self.dtype\n )\n mean, var = _compute_stats(\n x,\n self.reduction_axes,\n self.dtype,\n self.axis_name,\n self.axis_index_groups,\n use_fast_variance=self.use_fast_variance,\n mask=mask,\n )\n\n return _normalize(\n x,\n mean,\n var,\n scale,\n bias,\n self.reduction_axes,\n self.feature_axes,\n self.dtype,\n self.epsilon,\n )\n\n\nclass RMSNorm(Module):\n \"\"\"RMS Layer normalization (https://arxiv.org/abs/1910.07467).\n\n RMSNorm normalizes the activations of the layer for each given example in a\n batch independently, rather than across a batch like Batch Normalization.\n Unlike LayerNorm which re-centers the mean to be 0 and normalizes by the\n standard deviation of the activations, RMSNorm does not re-center at all\n and instead normalizes by the root mean square of the activations.\n\n Example usage::\n\n >>> from flax import nnx\n >>> import jax\n\n >>> x = jax.random.normal(jax.random.key(0), (5, 6))\n >>> layer = nnx.RMSNorm(num_features=6, rngs=nnx.Rngs(0))\n\n >>> nnx.state(layer)\n State({\n 'scale': Param( # 6 (24 B)\n value=Array([1., 1., 1., 1., 1., 1.], dtype=float32)\n )\n })\n\n >>> y = layer(x)\n\n Args:\n num_features: the number of input features.\n epsilon: A small float added to variance to avoid dividing by zero.\n dtype: the dtype of the result (default: infer from input and params).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n use_scale: If True, multiply by scale (gamma). When the next layer is linear\n (also e.g. nn.relu), this can be disabled since the scaling will be done\n by the next layer.\n scale_init: Initializer for scale, by default, one.\n reduction_axes: Axes for computing normalization statistics.\n feature_axes: Feature axes for learned bias and scaling.\n axis_name: the axis name used to combine batch statistics from multiple\n devices. See ``jax.pmap`` for a description of axis names (default: None).\n This is only needed if the model is subdivided across devices, i.e. the\n array being normalized is sharded across devices within a pmap.\n axis_index_groups: groups of axis indices within that named axis\n representing subsets of devices to reduce over (default: None). For\n example, ``[[0, 1], [2, 3]]`` would independently batch-normalize over\n the examples on the first two and last two devices. See ``jax.lax.psum``\n for more details. This argument is currently not supported for SPMD jit.\n use_fast_variance: If true, use a faster, but less numerically stable,\n calculation for the variance.\n promote_dtype: function to promote the dtype of all input array arguments\n (including Variables accessed through ``self``) to the desired dtype. The\n function should accept a tuple of ``(inputs, scale)`` and a ``dtype``\n keyword argument, and return a tuple of arrays with the promoted dtype.\n rngs: rng key.\n scale_metadata: Optional metadata dictionary to set when initializing\n the scale.\n \"\"\"\n\n def __init__(\n self,\n num_features: int,\n *,\n epsilon: float = 1e-6,\n dtype: tp.Optional[Dtype] = None,\n param_dtype: Dtype = jnp.float32,\n use_scale: bool = True,\n scale_init: Initializer = initializers.ones,\n reduction_axes: Axes = -1,\n feature_axes: Axes = -1,\n axis_name: tp.Optional[str] = None,\n axis_index_groups: tp.Any = None,\n use_fast_variance: bool = True,\n promote_dtype: PromoteDtypeFn = dtypes.promote_dtype,\n rngs: rnglib.Rngs,\n scale_metadata: tp.Mapping[str, tp.Any] = MappingProxyType({}),\n ):\n feature_shape = (num_features,)\n\n self.scale: nnx.Param[jax.Array] | None\n if use_scale:\n key = rngs.params()\n self.scale = nnx.Param(scale_init(key, feature_shape, param_dtype), **scale_metadata)\n else:\n self.scale = nnx.data(None)\n\n self.num_features = num_features\n self.epsilon = epsilon\n self.dtype = dtype\n self.param_dtype = param_dtype\n self.use_scale = use_scale\n self.reduction_axes = reduction_axes\n self.feature_axes = feature_axes\n self.axis_name = axis_name\n self.axis_index_groups = axis_index_groups\n self.use_fast_variance = use_fast_variance\n self.promote_dtype = promote_dtype\n\n def __call__(self, x, mask: tp.Optional[jax.Array] = None):\n \"\"\"Applies layer normalization on the input.\n\n Args:\n x: the inputs\n\n Returns:\n Normalized inputs (the same shape as inputs).\n \"\"\"\n # Promote dtypes for input and all Variables\n scale = self.scale[...] if self.scale else None\n x, scale = self.promote_dtype(\n (x, scale), dtype=self.dtype\n )\n mean, var = _compute_stats(\n x,\n self.reduction_axes,\n self.dtype,\n self.axis_name,\n self.axis_index_groups,\n use_mean=False,\n use_fast_variance=self.use_fast_variance,\n mask=mask,\n )\n\n return _normalize(\n x,\n mean,\n var,\n scale,\n None,\n self.reduction_axes,\n self.feature_axes,\n self.dtype,\n self.epsilon,\n )\n\n\nclass GroupNorm(Module):\n \"\"\"Group normalization (arxiv.org/abs/1803.08494).\n\n This op is similar to batch normalization, but statistics are shared across\n equally-sized groups of channels and not shared across batch dimension.\n Thus, group normalization does not depend on the batch composition and does\n not require maintaining internal state for storing statistics.\n The user should either specify the total number of channel groups or the\n number of channels per group.\n\n .. note::\n LayerNorm is a special case of GroupNorm where ``num_groups=1``.\n\n Example usage::\n\n >>> from flax import nnx\n >>> import jax\n >>> import numpy as np\n ...\n >>> x = jax.random.normal(jax.random.key(0), (3, 4, 5, 6))\n >>> layer = nnx.GroupNorm(num_features=6, num_groups=3, rngs=nnx.Rngs(0))\n >>> nnx.state(layer)\n State({\n 'bias': Param( # 6 (24 B)\n value=Array([0., 0., 0., 0., 0., 0.], dtype=float32)\n ),\n 'scale': Param( # 6 (24 B)\n value=Array([1., 1., 1., 1., 1., 1.], dtype=float32)\n )\n })\n >>> y = layer(x)\n ...\n >>> y = nnx.GroupNorm(num_features=6, num_groups=1, rngs=nnx.Rngs(0))(x)\n >>> y2 = nnx.LayerNorm(num_features=6, reduction_axes=(1, 2, 3), rngs=nnx.Rngs(0))(x)\n >>> np.testing.assert_allclose(y, y2)\n\n Args:\n num_features: the number of input features/channels.\n num_groups: the total number of channel groups. The default value of 32 is\n proposed by the original group normalization paper.\n group_size: the number of channels in a group.\n epsilon: A small float added to variance to avoid dividing by zero.\n dtype: the dtype of the result (default: infer from input and params).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n use_bias: If True, bias (beta) is added.\n use_scale: If True, multiply by scale (gamma). When the next layer is linear\n (also e.g. nn.relu), this can be disabled since the scaling will be done\n by the next layer.\n bias_init: Initializer for bias, by default, zero.\n scale_init: Initializer for scale, by default, one.\n reduction_axes: List of axes used for computing normalization statistics.\n This list must include the final dimension, which is assumed to be the\n feature axis. Furthermore, if the input used at call time has additional\n leading axes compared to the data used for initialisation, for example due\n to batching, then the reduction axes need to be defined explicitly.\n axis_name: the axis name used to combine batch statistics from multiple\n devices. See ``jax.pmap`` for a description of axis names (default: None).\n This is only needed if the model is subdivided across devices, i.e. the\n array being normalized is sharded across devices within a pmap or shard\n map. For SPMD jit, you do not need to manually synchronize. Just make sure\n that the axes are correctly annotated and XLA:SPMD will insert the\n necessary collectives.\n axis_index_groups: groups of axis indices within that named axis\n representing subsets of devices to reduce over (default: None). For\n example, ``[[0, 1], [2, 3]]`` would independently batch-normalize over the\n examples on the first two and last two devices. See ``jax.lax.psum`` for\n more details. This argument is currently not supported for SPMD jit.\n use_fast_variance: If true, use a faster, but less numerically stable,\n calculation for the variance.\n promote_dtype: function to promote the dtype of all input array arguments\n (including Variables accessed through ``self``) to the desired dtype. The\n function should accept a tuple of ``(inputs, scale, bias)`` and a ``dtype``\n keyword argument, and return a tuple of arrays with the promoted dtype.\n rngs: rng key.\n bias_metadata: Optional metadata dictionary to set when initializing\n the bias.\n scale_metadata: Optional metadata dictionary to set when initializing\n the scale.\n \"\"\"\n\n def __init__(\n self,\n num_features: int,\n num_groups: tp.Optional[int] = 32,\n group_size: tp.Optional[int] = None,\n *,\n epsilon: float = 1e-6,\n dtype: tp.Optional[Dtype] = None,\n param_dtype: Dtype = jnp.float32,\n use_bias: bool = True,\n use_scale: bool = True,\n bias_init: Initializer = initializers.zeros_init(),\n scale_init: Initializer = initializers.ones_init(),\n reduction_axes: tp.Optional[Axes] = None,\n axis_name: tp.Optional[str] = None,\n axis_index_groups: tp.Any = None,\n use_fast_variance: bool = True,\n promote_dtype: PromoteDtypeFn = dtypes.promote_dtype,\n rngs: rnglib.Rngs,\n bias_metadata: tp.Mapping[str, tp.Any] = MappingProxyType({}),\n scale_metadata: tp.Mapping[str, tp.Any] = MappingProxyType({}),\n ):\n self.feature_axis = -1\n\n if (num_groups is None and group_size is None) or (\n num_groups is not None and group_size is not None\n ):\n raise ValueError(\n 'Either `num_groups` or `group_size` should be '\n 'specified. If `group_size` is to be specified, '\n 'pass `num_groups=None` as argument to override '\n 'the default `num_groups` value of 32.'\n )\n\n if group_size is not None:\n if num_features % group_size != 0:\n raise ValueError(\n 'Number of features ({}) is not multiple of the '\n 'group size ({}).'.format(num_features, group_size)\n )\n self.num_groups = num_features // group_size\n self.group_size = group_size\n else:\n if not isinstance(num_groups, int) or num_groups <= 0 or (\n num_features % num_groups != 0\n ):\n raise ValueError(\n 'Number of groups ({}) does not divide the number'\n ' of channels ({}).'.format(num_groups, num_features)\n )\n self.num_groups = num_groups\n self.group_size = num_features // num_groups\n\n feature_shape = (num_features,)\n self.scale: nnx.Param[jax.Array] | None\n if use_scale:\n key = rngs.params()\n self.scale = nnx.Param(scale_init(key, feature_shape, param_dtype), **scale_metadata)\n else:\n self.scale = nnx.data(None)\n\n self.bias: nnx.Param[jax.Array] | None\n if use_bias:\n key = rngs.params()\n self.bias = nnx.Param(bias_init(key, feature_shape, param_dtype), **bias_metadata)\n else:\n self.bias = nnx.data(None)\n\n self.epsilon = epsilon\n self.dtype = dtype\n self.param_dtype = param_dtype\n self.use_bias = use_bias\n self.use_scale = use_scale\n self.reduction_axes = reduction_axes\n self.axis_name = axis_name\n self.axis_index_groups = axis_index_groups\n self.use_fast_variance = use_fast_variance\n self.promote_dtype = promote_dtype\n\n def __call__(self, x, *, mask: tp.Optional[jax.Array] = None):\n \"\"\"Applies group normalization to the input (arxiv.org/abs/1803.08494).\n\n Args:\n x: the input of shape ``...self.num_features`` where ``self.num_features``\n is a channels dimension and ``...`` represents an arbitrary number of\n extra dimensions that can be used to accumulate statistics over. If no\n reduction axes have been specified then all additional dimensions ``...``\n will be used to accumulate statistics apart from the leading dimension\n which is assumed to represent the batch.\n mask: Binary array of shape broadcastable to ``inputs`` tensor, indicating\n the positions for which the mean and variance should be computed.\n\n Returns:\n Normalized inputs (the same shape as inputs).\n \"\"\"\n if self.reduction_axes is not None:\n reduction_axes = self.reduction_axes\n else:\n reduction_axes = list(range(1, x.ndim - 1)) + [-1]\n reduction_axes = _canonicalize_axes(x.ndim, reduction_axes)\n\n group_shape = x.shape[:-1] + (self.num_groups, self.group_size)\n if mask is not None:\n mask = mask.reshape(mask.shape[:-1] + (self.num_groups, self.group_size))\n\n # Promote dtypes for input and all Variables\n scale = self.scale[...] if self.scale else None\n bias = self.bias[...] if self.bias else None\n x, scale, bias = self.promote_dtype(\n (x, scale, bias), dtype=self.dtype\n )\n mean, var = _compute_stats(\n x.reshape(group_shape),\n list(reduction_axes[:-1]) + [-1],\n self.dtype,\n self.axis_name,\n self.axis_index_groups,\n use_fast_variance=self.use_fast_variance,\n mask=mask,\n )\n mean = jnp.repeat(mean, self.group_size, axis=1)\n var = jnp.repeat(var, self.group_size, axis=1)\n\n return _normalize(\n x,\n mean,\n var,\n scale,\n bias,\n reduction_axes[:-1],\n (self.feature_axis,),\n self.dtype,\n self.epsilon,\n )\n\nclass WeightNorm(nnx.Module):\n \"\"\"L2 weight normalization (https://arxiv.org/abs/1602.07868).\n\n Weight normalization normalizes the weight params so that the l2-norm of\n the matrix is equal to 1. This is implemented as a layer wrapper where\n each wrapped layer will have its params l2-normalized before computing\n its ``__call__`` output.\n\n Example usage::\n\n >>> import jax\n >>> import numpy as np\n >>> from flax import nnx\n\n >>> class Foo(nnx.Module):\n ... def __init__(self, rngs: nnx.Rngs):\n ... self.normed_linear = nnx.WeightNorm(\n ... nnx.Linear(8, 4, rngs=rngs),\n ... variable_filter=nnx.PathContains('kernel'),\n ... rngs=rngs,\n ... )\n ...\n ... def __call__(self, x: jax.Array) -> jax.Array:\n ... return self.normed_linear(x)\n\n >>> rng = jax.random.key(42)\n >>> model = Foo(rngs=nnx.Rngs(rng))\n\n >>> x = jax.random.normal(rng, (5, 8))\n >>> y = model(x)\n >>> y.shape\n (5, 4)\n\n >>> w = model.normed_linear.layer_instance.kernel[...]\n >>> col_norms = np.linalg.norm(np.array(w), axis=0)\n >>> np.testing.assert_allclose(col_norms, np.ones(4))\n\n Args:\n layer_instance: The layer instance to wrap.\n feature_axes: The axes to normalize.\n use_scale: Whether to use a scale parameter.\n scale_init: The initializer for the scale parameter, by default ones.\n epsilon: The epsilon value for the normalization, by default 1e-12.\n dtype: The dtype of the result, by default infer from input and params.\n param_dtype: The dtype of the parameters, by default float32.\n variable_filter: The variable filter, by default ``nnx.PathContains('kernel')``.\n promote_dtype: function to promote the dtype of all input array arguments\n (including Variables accessed through ``self``) to the desired dtype. This\n is used internally by WeightNorm when normalizing weights.\n rngs: The rng key.\n \"\"\"\n def __init__(\n self,\n layer_instance: nnx.Module,\n *,\n feature_axes: Axes | None = -1,\n use_scale: bool = True,\n scale_init: Initializer = initializers.ones,\n epsilon: float = 1e-12,\n dtype: tp.Optional[Dtype] = None,\n param_dtype: Dtype = jnp.float32,\n variable_filter: nnx.filterlib.Filter = nnx.PathContains('kernel'),\n promote_dtype: PromoteDtypeFn = dtypes.promote_dtype,\n rngs: rnglib.Rngs,\n ):\n self.layer_instance = layer_instance\n self.feature_axes = () if feature_axes is None else feature_axes\n self.use_scale = use_scale\n self.scale_init = scale_init\n self.epsilon = epsilon\n self.dtype = dtype\n self.param_dtype = param_dtype\n self.variable_filter = nnx.filterlib.to_predicate(variable_filter)\n self.promote_dtype = promote_dtype\n self.scales : tp.Optional[dict] = None\n\n if use_scale:\n state = nnx.state(self.layer_instance, nnx.Param)\n def init_scales(param):\n feature_axes = _canonicalize_axes(param.ndim, self.feature_axes)\n scale_shape = tuple(param.shape[ax] for ax in feature_axes)\n return scale_init(rngs['params'], scale_shape)\n self.scales = nnx.data({\n path: init_scales(param) for path, param in nnx.to_flat_state(state)\n if self.variable_filter(path, param)})\n\n def _weightnorm_inplace(self, path, param):\n if not self.variable_filter(path, param):\n return\n\n if self.feature_axes is None:\n feature_axes = ()\n reduction_axes = tuple(range(param.ndim))\n else:\n feature_axes = _canonicalize_axes(param.ndim, self.feature_axes)\n reduction_axes = tuple(\n i for i in range(param.ndim) if i not in feature_axes\n )\n\n value_bar = _l2_normalize(param, axis=reduction_axes, eps=self.epsilon)\n\n if self.use_scale:\n if path not in self.scales:\n raise RuntimeError(\n f'Could not find the scale corresponding to the param {path} '\n 'in scales dict. Parameters of the layer_instance should not change!'\n )\n scale_value = self.scales[path]\n\n if len(feature_axes) < param.ndim:\n broadcast_shape = [1] * param.ndim\n for ax in feature_axes:\n broadcast_shape[ax] = param.shape[ax]\n scale_value = scale_value.reshape(broadcast_shape)\n value_bar = value_bar * scale_value\n\n cast_args = [param]\n if self.use_scale:\n cast_args.append(scale_value)\n\n final_dtype = dtypes.canonicalize_dtype(*cast_args, dtype=self.dtype)\n param.set_value(jnp.asarray(value_bar, final_dtype))\n\n def __call__(self, x: Array, *args, **kwargs) -> Array:\n \"\"\"Compute the l2-norm of the weights in ``self.layer_instance``\n and normalize the weights using this value before computing the\n ``__call__`` output.\n\n Args:\n *args: positional arguments to be passed into the call method of the\n underlying layer instance in ``self.layer_instance``.\n **kwargs: keyword arguments to be passed into the call method of the\n underlying layer instance in ``self.layer_instance``.\n\n Returns:\n Output of the layer using l2-normalized weights.\n \"\"\"\n state = nnx.state(self.layer_instance)\n for path, param in nnx.to_flat_state(state):\n self._weightnorm_inplace(path, param)\n\n return self.layer_instance(x, *args, **kwargs) # type: ignore\n\nclass InstanceNorm(Module):\n \"\"\"Instance normalization (https://arxiv.org/abs/1607.08022v3).\n InstanceNorm normalizes the activations of the layer for each channel (rather\n than across all channels like Layer Normalization), and for each given example\n in a batch independently (rather than across an entire batch like Batch\n Normalization). i.e. applies a transformation that maintains the mean activation\n within each channel within each example close to 0 and the activation standard\n deviation close to 1.\n\n .. note::\n This normalization operation is identical to LayerNorm and GroupNorm; the\n difference is simply which axes are reduced and the shape of the feature axes\n (i.e. the shape of the learnable scale and bias parameters).\n\n Example usage::\n\n >>> from flax import nnx\n >>> import jax\n >>> import numpy as np\n >>> # dimensions: (batch, height, width, channel)\n >>> x = jax.random.normal(jax.random.key(0), (2, 3, 4, 5))\n >>> layer = nnx.InstanceNorm(5, rngs=nnx.Rngs(0))\n >>> nnx.state(layer, nnx.Param)\n State({\n 'bias': Param( # 5 (20 B)\n value=Array([0., 0., 0., 0., 0.], dtype=float32)\n ),\n 'scale': Param( # 5 (20 B)\n value=Array([1., 1., 1., 1., 1.], dtype=float32)\n )\n })\n >>> y = layer(x)\n >>> # having a channel_axis of -1 in InstanceNorm is identical to reducing all non-batch,\n >>> # non-channel axes and using the feature_axes as the feature_axes in LayerNorm\n >>> y2 = nnx.LayerNorm(5, reduction_axes=[1, 2], feature_axes=-1, rngs=nnx.Rngs(0))(x)\n >>> np.testing.assert_allclose(y, y2, atol=1e-7)\n >>> y3 = nnx.GroupNorm(5, num_groups=x.shape[-1], rngs=nnx.Rngs(0))(x)\n >>> np.testing.assert_allclose(y, y3, atol=1e-7)\n\n Args:\n num_features: the number of input features/channels.\n epsilon: A small float added to variance to avoid dividing by zero.\n dtype: the dtype of the result (default: infer from input and params).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n use_bias: If True, bias (beta) is added.\n use_scale: If True, multiply by scale (gamma). When the next layer is linear\n (also e.g. nn.relu), this can be disabled since the scaling will be done\n by the next layer.\n bias_init: Initializer for bias, by default, zero.\n scale_init: Initializer for scale, by default, one.\n feature_axes: Axes for features. The learned bias and scaling parameters will\n be in the shape defined by the feature axes. All other axes except the batch\n axes (which is assumed to be the leading axis) will be reduced.\n axis_name: the axis name used to combine batch statistics from multiple\n devices. See ``jax.pmap`` for a description of axis names (default: None).\n This is only needed if the model is subdivided across devices, i.e. the\n array being normalized is sharded across devices within a pmap or shard\n map. For SPMD jit, you do not need to manually synchronize. Just make sure\n that the axes are correctly annotated and XLA:SPMD will insert the\n necessary collectives.\n axis_index_groups: groups of axis indices within that named axis\n representing subsets of devices to reduce over (default: None). For\n example, ``[[0, 1], [2, 3]]`` would independently batch-normalize over the\n examples on the first two and last two devices. See ``jax.lax.psum`` for\n more details. This argument is currently not supported for SPMD jit.\n use_fast_variance: If true, use a faster, but less numerically stable,\n calculation for the variance.\n promote_dtype: function to promote the dtype of all input array arguments\n (including Variables accessed through ``self``) to the desired dtype. The\n function should accept a tuple of ``(inputs, scale, bias)`` and a ``dtype``\n keyword argument, and return a tuple of arrays with the promoted dtype.\n rngs: The rng key.\n bias_metadata: Optional metadata dictionary to set when initializing\n the bias.\n scale_metadata: Optional metadata dictionary to set when initializing\n the scale.\n \"\"\"\n\n def __init__(\n self,\n num_features: int,\n *,\n epsilon: float = 1e-6,\n dtype: tp.Optional[Dtype] = None,\n param_dtype: Dtype = jnp.float32,\n use_bias: bool = True,\n use_scale: bool = True,\n bias_init: Initializer = initializers.zeros,\n scale_init: Initializer = initializers.ones,\n feature_axes: Axes = -1,\n axis_name: tp.Optional[str] = None,\n axis_index_groups: tp.Any = None,\n use_fast_variance: bool = True,\n promote_dtype: PromoteDtypeFn = dtypes.promote_dtype,\n rngs: rnglib.Rngs,\n bias_metadata: tp.Mapping[str, tp.Any] = MappingProxyType({}),\n scale_metadata: tp.Mapping[str, tp.Any] = MappingProxyType({}),\n ):\n feature_shape = (num_features,)\n self.scale: nnx.Param[jax.Array] | None\n if use_scale:\n key = rngs.params()\n self.scale = nnx.Param(scale_init(key, feature_shape, param_dtype), **scale_metadata)\n else:\n self.scale = None\n\n self.bias: nnx.Param[jax.Array] | None\n if use_bias:\n key = rngs.params()\n self.bias = nnx.Param(bias_init(key, feature_shape, param_dtype), **bias_metadata)\n else:\n self.bias = None\n\n self.num_features = num_features\n self.epsilon = epsilon\n self.dtype = dtype\n self.param_dtype = param_dtype\n self.use_bias = use_bias\n self.use_scale = use_scale\n self.feature_axes = feature_axes\n self.axis_name = axis_name\n self.axis_index_groups = axis_index_groups\n self.use_fast_variance = use_fast_variance\n self.promote_dtype = promote_dtype\n\n def __call__(self, x, *, mask: tp.Optional[jax.Array] = None):\n \"\"\"Applies instance normalization on the input.\n\n Args:\n x: the inputs\n mask: Binary array of shape broadcastable to ``inputs`` array, indicating\n the positions for which the mean and variance should be computed.\n\n Returns:\n Normalized inputs (the same shape as inputs).\n \"\"\"\n feature_axes = _canonicalize_axes(x.ndim, self.feature_axes)\n if 0 in feature_axes:\n raise ValueError('The channel axes cannot include the leading dimension '\n 'as this is assumed to be the batch axis.')\n reduction_axes = [i for i in range(1, x.ndim) if i not in feature_axes]\n\n # Promote dtypes for input and all Variables\n scale = self.scale[...] if self.scale else None\n bias = self.bias[...] if self.bias else None\n x, scale, bias = self.promote_dtype(\n (x, scale, bias), dtype=self.dtype\n )\n mean, var = _compute_stats(\n x,\n reduction_axes,\n self.dtype,\n self.axis_name,\n self.axis_index_groups,\n use_fast_variance=self.use_fast_variance,\n mask=mask,\n )\n\n return _normalize(\n x,\n mean,\n var,\n scale,\n bias,\n reduction_axes,\n feature_axes,\n self.dtype,\n self.epsilon,\n )\n\n\nclass SpectralNorm(Module):\n \"\"\"Spectral normalization.\n\n See:\n\n - https://arxiv.org/abs/1802.05957\n - https://arxiv.org/abs/1805.08318\n - https://arxiv.org/abs/1809.11096\n\n Spectral normalization normalizes the weight params so that the spectral\n norm of the matrix is equal to 1. This is implemented as a layer wrapper\n where each wrapped layer will have its params spectral normalized before\n computing its ``__call__`` output.\n\n .. note::\n The initialized variables dict will contain, in addition to a 'params'\n collection, a separate 'batch_stats' collection that will contain a\n ``u`` vector and ``sigma`` value, which are intermediate values used\n when performing spectral normalization. During training, we pass in\n ``update_stats=True`` so that ``u`` and ``sigma`` are updated with\n the most recently computed values using power iteration. This will\n help the power iteration method approximate the true singular value\n more accurately over time. During eval, we pass in ``update_stats=False``\n to ensure we get deterministic behavior from the model.\n\n Example usage::\n\n >>> from flax import nnx\n >>> import jax\n >>> rngs = nnx.Rngs(0)\n >>> x = jax.random.normal(jax.random.key(0), (3, 4))\n >>> layer = nnx.SpectralNorm(nnx.Linear(4, 5, rngs=rngs), rngs=rngs)\n >>> jax.tree.map(jax.numpy.shape, nnx.state(layer, nnx.Param))\n State({\n 'layer_instance': {\n 'bias': Param(\n value=(5,)\n ),\n 'kernel': Param(\n value=(4, 5)\n )\n }\n })\n >>> y = layer(x, update_stats=True)\n\n Args:\n layer_instance: Module instance that is wrapped with SpectralNorm\n n_steps: How many steps of power iteration to perform to approximate the\n singular value of the weight params.\n epsilon: A small float added to l2-normalization to avoid dividing by zero.\n dtype: the dtype of the result (default: infer from input and params).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n error_on_non_matrix: Spectral normalization is only defined on matrices. By\n default, this module will return scalars unchanged and flatten\n higher-order tensors in their leading dimensions. Setting this flag to\n True will instead throw an error if a weight tensor with dimension greater\n than 2 is used by the layer.\n update_stats: if True, the stored batch statistics will be\n used instead of computing the batch statistics on the input.\n rngs: rng key.\n \"\"\"\n\n def __init__(\n self,\n layer_instance: Module,\n *,\n n_steps: int = 1,\n epsilon: float = 1e-12,\n dtype: tp.Optional[Dtype] = None,\n param_dtype: Dtype = jnp.float32,\n error_on_non_matrix: bool = False,\n update_stats: bool = True,\n rngs: rnglib.Rngs,\n ):\n self.layer_instance = layer_instance\n self.n_steps = n_steps\n self.epsilon = epsilon\n self.dtype = dtype\n self.param_dtype = param_dtype\n self.error_on_non_matrix = error_on_non_matrix\n\n # We define here self.use_running_average attribute to make\n # .train() and .eval() work. These methods internally change\n # self.use_running_average to False in train mode and\n # to True in eval mode\n # update_stats flag has the opposite logic.\n self.use_running_average = not update_stats\n\n # Initialize batch stat variables:\n state = nnx.state(self.layer_instance, nnx.Param)\n\n def init_batch_stats(path, param):\n if param.ndim <= 1 or self.n_steps < 1:\n return None\n elif param.ndim > 2:\n if self.error_on_non_matrix:\n raise ValueError(\n f'Layer instance parameter is {param.ndim}D but error_on_non_matrix is True'\n )\n else:\n param = jnp.reshape(param, (-1, param.shape[-1]))\n\n path_u = path + (\"u\", )\n path_sigma = path + (\"sigma\", )\n\n key = rngs.params()\n return [\n (\n path_u,\n nnx.BatchStat(\n initializers.normal()(key, (1, param.shape[-1]), self.param_dtype))\n ),\n (\n path_sigma,\n nnx.BatchStat(initializers.ones(key, (), self.param_dtype)),\n )\n ]\n\n batch_stats: dict[tuple, nnx.BatchStat] = {}\n for path, param in nnx.to_flat_state(state):\n batch_stats_per_param = init_batch_stats(path, param)\n if batch_stats_per_param is None:\n continue\n for new_path, bstat in batch_stats_per_param:\n batch_stats[new_path] = bstat\n\n self.batch_stats = nnx.data(batch_stats)\n\n def __call__(\n self,\n x,\n update_stats: tp.Optional[bool] = None,\n ):\n \"\"\"Compute the largest singular value of the weights in ``self.layer_instance``\n using power iteration and normalize the weights using this value before\n computing the ``__call__`` output.\n\n Args:\n x: the input array of the nested layer\n update_stats: if True, update the internal ``u`` vector and ``sigma``\n value after computing their updated values using power iteration. This\n will help the power iteration method approximate the true singular value\n more accurately over time.\n\n Returns:\n Output of the layer using spectral normalized weights.\n \"\"\"\n update_stats = first_from(\n update_stats,\n not self.use_running_average,\n error_msg=\"\"\"No `update_stats` argument was provided to SpectralNorm\n as either a __call__ argument or __init__ argument.\"\"\",\n )\n state = nnx.state(self.layer_instance, nnx.Param)\n for path, param in nnx.to_flat_state(state):\n self._spectral_normalize_inplace(path, param, update_stats=update_stats)\n return self.layer_instance(x)\n\n def _spectral_normalize_inplace(self, path, orig_param, update_stats):\n param = orig_param\n param_shape = param.shape\n if param.ndim <= 1 or self.n_steps < 1:\n return\n elif param.ndim > 2:\n if self.error_on_non_matrix:\n raise ValueError(\n f'Layer instance parameter is {param.ndim}D but error_on_non_matrix is True'\n )\n else:\n param = jnp.reshape(param, (-1, param.shape[-1]))\n\n path_u = path + (\"u\", )\n path_sigma = path + (\"sigma\", )\n\n if path_u not in self.batch_stats:\n raise RuntimeError(\n f\"Could not find the path for u batch stat corresponding to the param {path} \"\n \"in the batch stats dict. Parameters of the layer_instance should not change!\"\n )\n if path_sigma not in self.batch_stats:\n raise RuntimeError(\n f\"Could not find the path for sigma batch stat corresponding to the param {path} \"\n \"in the batch stats dict. Parameters of the layer_instance should not change!\"\n )\n\n u = self.batch_stats[path_u][...]\n\n for _ in range(self.n_steps):\n v = _l2_normalize(jnp.matmul(u, param.T), eps=self.epsilon)\n u = _l2_normalize(jnp.matmul(v, param), eps=self.epsilon)\n\n u = lax.stop_gradient(u)\n v = lax.stop_gradient(v)\n\n sigma = jnp.matmul(jnp.matmul(v, param), u.T)[0, 0]\n param = param / jnp.where(sigma != 0, sigma, 1)\n param = param.reshape(param_shape)\n\n if update_stats:\n self.batch_stats[path_u][...] = u\n self.batch_stats[path_sigma][...] = sigma\n\n dtype = dtypes.canonicalize_dtype(param, u, v, sigma, dtype=self.dtype)\n orig_param[...] = jnp.asarray(param, dtype)\n" }, { "path": "flax/nnx/nn/recurrent.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"RNN modules for Flax.\"\"\"\nimport warnings\nfrom typing import Any, TypeVar\nfrom collections.abc import Mapping\nfrom types import MappingProxyType\nfrom collections.abc import Mapping\nfrom collections.abc import Callable\nfrom functools import partial\nfrom typing_extensions import Protocol\nfrom absl import logging\n\nimport jax\nimport jax.numpy as jnp\n\nfrom flax import nnx\nfrom flax.nnx import filterlib, rnglib\nfrom flax.nnx.module import Module\nfrom flax.nnx.nn import initializers, dtypes\nfrom flax.nnx.nn.linear import Linear\nfrom flax.nnx.nn.activations import sigmoid\nfrom flax.nnx.nn.activations import tanh\nfrom flax.nnx.transforms import iteration\nfrom flax.typing import Dtype, Initializer, PromoteDtypeFn, Shape\n\ndefault_kernel_init = initializers.lecun_normal()\ndefault_bias_init = initializers.zeros_init()\n\nA = TypeVar(\"A\")\nArray = jax.Array\nOutput = Any\nCarry = Any\n\nclass RNNCellBase(Module):\n \"\"\"RNN cell base class.\"\"\"\n\n def initialize_carry(\n self,\n input_shape: tuple[int, ...],\n rngs: rnglib.Rngs | rnglib.RngStream | None = None,\n carry_init: Initializer | None = None,\n ) -> Carry:\n \"\"\"Initialize the RNN cell carry.\n\n Args:\n input_shape: a tuple providing the shape of the input to the cell.\n rngs: random number generator passed to the init_fn.\n carry_init: optional carry initializer.\n\n Returns:\n An initialized carry for the given RNN cell.\n \"\"\"\n raise NotImplementedError\n\n def __call__(\n self,\n carry: Carry,\n inputs: Array\n ) -> tuple[Carry, Array]:\n \"\"\"Run the RNN cell.\n\n Args:\n carry: the hidden state of the RNN cell.\n inputs: an ndarray with the input for the current time step.\n All dimensions except the final are considered batch dimensions.\n\n Returns:\n A tuple with the new carry and the output.\n \"\"\"\n raise NotImplementedError\n\n @property\n def num_feature_axes(self) -> int:\n \"\"\"Returns the number of feature axes of the RNN cell.\"\"\"\n raise NotImplementedError\n\ndef modified_orthogonal(key: Array, shape: Shape, dtype: Dtype = jnp.float32) -> Array:\n \"\"\"Modified orthogonal initializer for compatibility with half precision.\"\"\"\n initializer = initializers.orthogonal()\n return initializer(key, shape).astype(dtype)\n\nclass LSTMCell(RNNCellBase):\n r\"\"\"LSTM cell.\n\n The mathematical definition of the cell is as follows\n\n .. math::\n \\begin{array}{ll}\n i = \\sigma(W_{ii} x + W_{hi} h + b_{hi}) \\\\\n f = \\sigma(W_{if} x + W_{hf} h + b_{hf}) \\\\\n g = \\tanh(W_{ig} x + W_{hg} h + b_{hg}) \\\\\n o = \\sigma(W_{io} x + W_{ho} h + b_{ho}) \\\\\n c' = f * c + i * g \\\\\n h' = o * \\tanh(c') \\\\\n \\end{array}\n\n where x is the input, h is the output of the previous time step, and c is\n the memory.\n \"\"\"\n\n def __init__(\n self,\n in_features: int,\n hidden_features: int,\n *,\n gate_fn: Callable[..., Any] = sigmoid,\n activation_fn: Callable[..., Any] = tanh,\n kernel_init: Initializer = default_kernel_init,\n recurrent_kernel_init: Initializer = modified_orthogonal,\n bias_init: Initializer = initializers.zeros_init(),\n dtype: Dtype | None = None,\n param_dtype: Dtype = jnp.float32,\n carry_init: Initializer | None = None,\n promote_dtype: PromoteDtypeFn = dtypes.promote_dtype,\n keep_rngs: bool = False,\n rngs: rnglib.Rngs,\n kernel_metadata: Mapping[str, Any] = MappingProxyType({}),\n recurrent_kernel_metadata: Mapping[str, Any] = MappingProxyType({}),\n bias_metadata: Mapping[str, Any] = MappingProxyType({}),\n ):\n self.in_features = in_features\n self.hidden_features = hidden_features\n self.gate_fn = gate_fn\n self.activation_fn = activation_fn\n self.dtype = dtype\n self.param_dtype = param_dtype\n self.promote_dtype = promote_dtype\n self.rngs: rnglib.RngStream | None\n if keep_rngs:\n self.rngs = rngs.carry.fork()\n else:\n self.rngs = nnx.data(None)\n\n # input and recurrent layers are summed so only one needs a bias.\n dense_i = partial(\n Linear,\n in_features=in_features,\n out_features=hidden_features,\n use_bias=False,\n kernel_init=kernel_init,\n dtype=self.dtype,\n param_dtype=self.param_dtype,\n promote_dtype=self.promote_dtype,\n rngs=rngs,\n kernel_metadata=kernel_metadata,\n )\n\n dense_h = partial(\n Linear,\n in_features=hidden_features,\n out_features=hidden_features,\n use_bias=True,\n kernel_init=recurrent_kernel_init,\n bias_init=bias_init,\n dtype=self.dtype,\n param_dtype=self.param_dtype,\n promote_dtype=self.promote_dtype,\n rngs=rngs,\n kernel_metadata=recurrent_kernel_metadata,\n bias_metadata=bias_metadata,\n )\n\n self.ii = dense_i()\n self.if_ = dense_i()\n self.ig = dense_i()\n self.io = dense_i()\n self.hi = dense_h()\n self.hf = dense_h()\n self.hg = dense_h()\n self.ho = dense_h()\n\n if carry_init:\n warnings.warn(\n \"carry_init is provided in __init__. \"\n \"Please, use carry_init argument in `initialize_carry` method instead to initialize the carry. \"\n \"Otherwise, two instances with same configuration but different carry_init \"\n \"functions will have different graphdefs.\"\n )\n\n self.carry_init = carry_init\n\n def __call__(\n self, carry: tuple[Array, Array], inputs: Array\n ) -> tuple[tuple[Array, Array], Array]: # type: ignore[override]\n r\"\"\"A long short-term memory (LSTM) cell.\n\n Args:\n carry: the hidden state of the LSTM cell,\n initialized using ``LSTMCell.initialize_carry``.\n inputs: an ndarray with the input for the current time step.\n All dimensions except the final are considered batch dimensions.\n\n Returns:\n A tuple with the new carry and the output.\n \"\"\"\n c, h = carry\n i = self.gate_fn(self.ii(inputs) + self.hi(h))\n f = self.gate_fn(self.if_(inputs) + self.hf(h))\n g = self.activation_fn(self.ig(inputs) + self.hg(h))\n o = self.gate_fn(self.io(inputs) + self.ho(h))\n new_c = f * c + i * g\n new_h = o * self.activation_fn(new_c)\n return (new_c, new_h), new_h\n\n def initialize_carry(\n self,\n input_shape: tuple[int, ...],\n rngs: rnglib.Rngs | rnglib.RngStream | None = None,\n carry_init: Initializer | None = None,\n ) -> tuple[Array, Array]: # type: ignore[override]\n \"\"\"Initialize the RNN cell carry.\n\n Args:\n rng: random number generator passed to the init_fn.\n input_shape: a tuple providing the shape of the input to the cell.\n Returns:\n An initialized carry for the given RNN cell.\n \"\"\"\n batch_dims = input_shape[:-1]\n if rngs is None:\n rngs = self.rngs\n if isinstance(rngs, rnglib.Rngs):\n rngs = rngs.carry\n if rngs is None:\n raise ValueError('RNGs must be provided to initialize the cell carry.')\n\n if self.carry_init is None and carry_init is None:\n carry_init = initializers.zeros_init()\n elif carry_init is None:\n carry_init = self.carry_init\n assert carry_init is not None # just to please mypy\n\n mem_shape = batch_dims + (self.hidden_features,)\n c = carry_init(rngs(), mem_shape, self.param_dtype)\n h = carry_init(rngs(), mem_shape, self.param_dtype)\n return (c, h)\n\n @property\n def num_feature_axes(self) -> int:\n return 1\n\n\nclass OptimizedLSTMCell(RNNCellBase):\n r\"\"\"More efficient LSTM Cell that concatenates state components before matmul.\n\n The parameters are compatible with ``LSTMCell``. Note that this cell is often\n faster than ``LSTMCell`` as long as the hidden size is roughly <= 2048 units.\n\n The mathematical definition of the cell is the same as ``LSTMCell`` and as\n follows:\n\n .. math::\n\n \\begin{array}{ll}\n i = \\sigma(W_{ii} x + W_{hi} h + b_{hi}) \\\\\n f = \\sigma(W_{if} x + W_{hf} h + b_{hf}) \\\\\n g = \\tanh(W_{ig} x + W_{hg} h + b_{hg}) \\\\\n o = \\sigma(W_{io} x + W_{ho} h + b_{ho}) \\\\\n c' = f * c + i * g \\\\\n h' = o * \\tanh(c') \\\\\n \\end{array}\n\n where x is the input, h is the output of the previous time step, and c is\n the memory.\n\n Args:\n gate_fn: activation function used for gates (default: sigmoid).\n activation_fn: activation function used for output and memory update\n (default: tanh).\n kernel_init: initializer function for the kernels that transform\n the input (default: lecun_normal).\n recurrent_kernel_init: initializer function for the kernels that transform\n the hidden state (default: initializers.orthogonal()).\n bias_init: initializer for the bias parameters (default: initializers.zeros_init()).\n dtype: the dtype of the computation (default: infer from inputs and params).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n keep_rngs: whether to store the input rngs as attribute (i.e. `self.rngs = rngs`)\n (default: True). If rngs is stored, we should split the module as\n `graphdef, params, nondiff = nnx.split(module, nnx.Param, ...)` where `nondiff`\n contains RNG object associated with stored `self.rngs`.\n rngs: rng key.\n kernel_metadata: Optional metadata dictionary to set when initializing\n the kernels that transform the input.\n recurrent_kernel_metadata: Optional metadata dictionary to set when initializing\n the kernels that transform the hidden state.\n bias_metadata: Optional metadata dictionary to set when initializing\n the bias of layers that transform the hidden state.\n \"\"\"\n\n def __init__(\n self,\n in_features: int,\n hidden_features: int,\n *,\n gate_fn: Callable[..., Any] = sigmoid,\n activation_fn: Callable[..., Any] = tanh,\n kernel_init: Initializer = default_kernel_init,\n recurrent_kernel_init: Initializer = initializers.orthogonal(),\n bias_init: Initializer = initializers.zeros_init(),\n dtype: Dtype | None = None,\n param_dtype: Dtype = jnp.float32,\n carry_init: Initializer | None = None,\n promote_dtype: PromoteDtypeFn = dtypes.promote_dtype,\n keep_rngs: bool = False,\n rngs: rnglib.Rngs,\n kernel_metadata: Mapping[str, Any] = MappingProxyType({}),\n recurrent_kernel_metadata: Mapping[str, Any] = MappingProxyType({}),\n bias_metadata: Mapping[str, Any] = MappingProxyType({}),\n ):\n self.in_features = in_features\n self.hidden_features = hidden_features\n self.gate_fn = gate_fn\n self.activation_fn = activation_fn\n self.dtype = dtype\n self.param_dtype = param_dtype\n self.promote_dtype = promote_dtype\n self.rngs: rnglib.RngStream | None\n if keep_rngs:\n self.rngs = rngs.carry.fork()\n else:\n self.rngs = nnx.data(None)\n\n # input and recurrent layers are summed so only one needs a bias.\n self.dense_i = Linear(\n in_features=in_features,\n out_features=4 * hidden_features,\n use_bias=False,\n kernel_init=kernel_init,\n dtype=self.dtype,\n param_dtype=self.param_dtype,\n promote_dtype=self.promote_dtype,\n rngs=rngs,\n kernel_metadata=kernel_metadata,\n )\n\n self.dense_h = Linear(\n in_features=hidden_features,\n out_features=4 * hidden_features,\n use_bias=True,\n kernel_init=recurrent_kernel_init,\n bias_init=bias_init,\n dtype=self.dtype,\n param_dtype=self.param_dtype,\n promote_dtype=self.promote_dtype,\n rngs=rngs,\n kernel_metadata=recurrent_kernel_metadata,\n bias_metadata=bias_metadata,\n )\n\n if carry_init:\n warnings.warn(\n \"carry_init is provided in __init__. \"\n \"Please, use carry_init argument in `initialize_carry` method instead to initialize the carry. \"\n \"Otherwise, two instances with same configuration but different carry_init \"\n \"functions will have different graphdefs.\"\n )\n\n self.carry_init = carry_init\n\n def __call__(\n self, carry: tuple[Array, Array], inputs: Array\n ) -> tuple[tuple[Array, Array], Array]: # type: ignore[override]\n r\"\"\"An optimized long short-term memory (LSTM) cell.\n\n Args:\n carry: the hidden state of the LSTM cell, initialized using\n ``LSTMCell.initialize_carry``.\n inputs: an ndarray with the input for the current time step.\n All dimensions except the final are considered batch dimensions.\n\n Returns:\n A tuple with the new carry and the output.\n \"\"\"\n c, h = carry\n\n # Compute combined transformations for inputs and hidden state\n y = self.dense_i(inputs) + self.dense_h(h)\n\n # Split the combined transformations into individual gates\n i, f, g, o = jnp.split(y, indices_or_sections=4, axis=-1)\n\n # Apply gate activations\n i = self.gate_fn(i)\n f = self.gate_fn(f)\n g = self.activation_fn(g)\n o = self.gate_fn(o)\n\n # Update cell state and hidden state\n new_c = f * c + i * g\n new_h = o * self.activation_fn(new_c)\n return (new_c, new_h), new_h\n\n def initialize_carry(\n self,\n input_shape: tuple[int, ...],\n rngs: rnglib.Rngs | rnglib.RngStream | None = None,\n carry_init: Initializer | None = None,\n ) -> tuple[Array, Array]: # type: ignore[override]\n \"\"\"Initialize the RNN cell carry.\n\n Args:\n rngs: random number generator passed to the init_fn.\n input_shape: a tuple providing the shape of the input to the cell.\n\n Returns:\n An initialized carry for the given RNN cell.\n \"\"\"\n batch_dims = input_shape[:-1]\n if rngs is None:\n rngs = self.rngs\n if isinstance(rngs, rnglib.Rngs):\n rngs = rngs.carry\n if rngs is None:\n raise ValueError('RNGs must be provided to initialize the cell carry.')\n mem_shape = batch_dims + (self.hidden_features,)\n\n if self.carry_init is None and carry_init is None:\n carry_init = initializers.zeros_init()\n elif carry_init is None:\n carry_init = self.carry_init\n assert carry_init is not None # just to please mypy\n\n c = carry_init(rngs(), mem_shape, self.param_dtype)\n h = carry_init(rngs(), mem_shape, self.param_dtype)\n return (c, h)\n\n @property\n def num_feature_axes(self) -> int:\n return 1\n\n\nclass SimpleCell(RNNCellBase):\n r\"\"\"Simple cell.\n\n The mathematical definition of the cell is as follows\n\n .. math::\n\n \\begin{array}{ll}\n h' = \\tanh(W_i x + b_i + W_h h)\n \\end{array}\n\n where x is the input and h is the output of the previous time step.\n\n If `residual` is `True`,\n\n .. math::\n\n \\begin{array}{ll}\n h' = \\tanh(W_i x + b_i + W_h h + h)\n \\end{array}\n \"\"\"\n\n def __init__(\n self,\n in_features: int,\n hidden_features: int, # not inferred from carry for now\n *,\n dtype: Dtype = jnp.float32,\n param_dtype: Dtype = jnp.float32,\n carry_init: Initializer | None = None,\n residual: bool = False,\n activation_fn: Callable[..., Any] = tanh,\n kernel_init: Initializer = initializers.lecun_normal(),\n recurrent_kernel_init: Initializer = initializers.orthogonal(),\n bias_init: Initializer = initializers.zeros_init(),\n promote_dtype: PromoteDtypeFn = dtypes.promote_dtype,\n keep_rngs: bool = False,\n rngs: rnglib.Rngs,\n kernel_metadata: Mapping[str, Any] = MappingProxyType({}),\n recurrent_kernel_metadata: Mapping[str, Any] = MappingProxyType({}),\n bias_metadata: Mapping[str, Any] = MappingProxyType({}),\n ):\n self.in_features = in_features\n self.hidden_features = hidden_features\n self.dtype = dtype\n self.param_dtype = param_dtype\n self.residual = residual\n self.activation_fn = activation_fn\n self.promote_dtype = promote_dtype\n self.rngs: rnglib.RngStream | None\n if keep_rngs:\n self.rngs = rngs.carry.fork()\n else:\n self.rngs = nnx.data(None)\n\n # self.hidden_features = carry.shape[-1]\n # input and recurrent layers are summed so only one needs a bias.\n self.dense_h = Linear(\n in_features=self.hidden_features,\n out_features=self.hidden_features,\n use_bias=False,\n dtype=self.dtype,\n param_dtype=self.param_dtype,\n kernel_init=recurrent_kernel_init,\n promote_dtype=self.promote_dtype,\n rngs=rngs,\n kernel_metadata=recurrent_kernel_metadata,\n )\n self.dense_i = Linear(\n in_features=self.in_features,\n out_features=self.hidden_features,\n use_bias=True,\n dtype=self.dtype,\n param_dtype=self.param_dtype,\n kernel_init=kernel_init,\n bias_init=bias_init,\n promote_dtype=self.promote_dtype,\n rngs=rngs,\n kernel_metadata=kernel_metadata,\n bias_metadata=bias_metadata,\n )\n\n if carry_init:\n warnings.warn(\n \"carry_init is provided in __init__. \"\n \"Please, use carry_init argument in `initialize_carry` method instead to initialize the carry. \"\n \"Otherwise, two instances with same configuration but different carry_init \"\n \"functions will have different graphdefs.\"\n )\n\n self.carry_init = carry_init\n\n def __call__(self, carry: Array, inputs: Array) -> tuple[Array, Array]: # type: ignore[override]\n new_carry = self.dense_i(inputs) + self.dense_h(carry)\n if self.residual:\n new_carry += carry\n new_carry = self.activation_fn(new_carry)\n return new_carry, new_carry\n\n def initialize_carry(\n self,\n input_shape: tuple[int, ...],\n rngs: rnglib.Rngs | rnglib.RngStream | None = None,\n carry_init: Initializer | None = None,\n ) -> Array: # type: ignore[override]\n \"\"\"Initialize the RNN cell carry.\n\n Args:\n rng: random number generator passed to the init_fn.\n input_shape: a tuple providing the shape of the input to the cell.\n\n Returns:\n An initialized carry for the given RNN cell.\n \"\"\"\n if rngs is None:\n rngs = self.rngs\n if isinstance(rngs, rnglib.Rngs):\n rngs = rngs.carry\n if rngs is None:\n raise ValueError('RNGs must be provided to initialize the cell carry.')\n\n batch_dims = input_shape[:-1]\n mem_shape = batch_dims + (self.hidden_features,)\n\n if self.carry_init is None and carry_init is None:\n carry_init = initializers.zeros_init()\n elif carry_init is None:\n carry_init = self.carry_init\n assert carry_init is not None # just to please mypy\n\n return carry_init(rngs(), mem_shape, self.param_dtype)\n\n @property\n def num_feature_axes(self) -> int:\n return 1\n\n\nclass GRUCell(RNNCellBase):\n r\"\"\"GRU cell.\n\n The mathematical definition of the cell is as follows\n\n .. math::\n\n \\begin{array}{ll}\n r = \\sigma(W_{ir} x + b_{ir} + W_{hr} h) \\\\\n z = \\sigma(W_{iz} x + b_{iz} + W_{hz} h) \\\\\n n = \\tanh(W_{in} x + b_{in} + r * (W_{hn} h + b_{hn})) \\\\\n h' = (1 - z) * n + z * h \\\\\n \\end{array}\n\n where x is the input and h is the output of the previous time step.\n\n Args:\n in_features: number of input features.\n hidden_features: number of output features.\n gate_fn: activation function used for gates (default: sigmoid).\n activation_fn: activation function used for output and memory update\n (default: tanh).\n kernel_init: initializer function for the kernels that transform\n the input (default: lecun_normal).\n recurrent_kernel_init: initializer function for the kernels that transform\n the hidden state (default: initializers.orthogonal()).\n bias_init: initializer for the bias parameters (default: initializers.zeros_init()).\n dtype: the dtype of the computation (default: None).\n param_dtype: the dtype passed to parameter initializers (default: float32).\n keep_rngs: whether to store the input rngs as attribute (i.e. `self.rngs = rngs`)\n (default: True). If rngs is stored, we should split the module as\n `graphdef, params, nondiff = nnx.split(module, nnx.Param, ...)` where `nondiff`\n contains RNG object associated with stored `self.rngs`.\n rngs: rng key.\n kernel_metadata: Optional metadata dictionary to set when initializing\n the kernels that transform the input.\n recurrent_kernel_metadata: Optional metadata dictionary to set when initializing\n the kernels that transform the hidden state.\n bias_metadata: Optional metadata dictionary to set when initializing\n the bias of layers that transform the input.\n \"\"\"\n\n def __init__(\n self,\n in_features: int,\n hidden_features: int,\n *,\n gate_fn: Callable[..., Any] = sigmoid,\n activation_fn: Callable[..., Any] = tanh,\n kernel_init: Initializer = default_kernel_init,\n recurrent_kernel_init: Initializer = initializers.orthogonal(),\n bias_init: Initializer = initializers.zeros_init(),\n dtype: Dtype | None = None,\n param_dtype: Dtype = jnp.float32,\n carry_init: Initializer | None = None,\n promote_dtype: PromoteDtypeFn = dtypes.promote_dtype,\n keep_rngs: bool = False,\n rngs: rnglib.Rngs,\n kernel_metadata: Mapping[str, Any] = MappingProxyType({}),\n recurrent_kernel_metadata: Mapping[str, Any] = MappingProxyType({}),\n bias_metadata: Mapping[str, Any] = MappingProxyType({}),\n ):\n self.in_features = in_features\n self.hidden_features = hidden_features\n self.gate_fn = gate_fn\n self.activation_fn = activation_fn\n self.dtype = dtype\n self.param_dtype = param_dtype\n self.promote_dtype = promote_dtype\n self.rngs: rnglib.RngStream | None\n if keep_rngs:\n self.rngs = rngs.carry.fork()\n else:\n self.rngs = nnx.data(None)\n\n # Combine input transformations into a single linear layer\n self.dense_i = Linear(\n in_features=in_features,\n out_features=3 * hidden_features, # r, z, n\n use_bias=True,\n kernel_init=kernel_init,\n bias_init=bias_init,\n dtype=self.dtype,\n param_dtype=self.param_dtype,\n promote_dtype=self.promote_dtype,\n rngs=rngs,\n kernel_metadata=kernel_metadata,\n bias_metadata=bias_metadata,\n )\n\n self.dense_h = Linear(\n in_features=hidden_features,\n out_features=3 * hidden_features, # r, z, n\n use_bias=False,\n kernel_init=recurrent_kernel_init,\n dtype=self.dtype,\n param_dtype=self.param_dtype,\n promote_dtype=self.promote_dtype,\n rngs=rngs,\n kernel_metadata=recurrent_kernel_metadata,\n )\n\n if carry_init:\n warnings.warn(\n \"carry_init is provided in __init__. \"\n \"Please, use carry_init argument in `initialize_carry` method instead to initialize the carry. \"\n \"Otherwise, two instances with same configuration but different carry_init \"\n \"functions will have different graphdefs.\"\n )\n\n self.carry_init = carry_init\n\n def __call__(self, carry: Array, inputs: Array) -> tuple[Array, Array]: # type: ignore[override]\n \"\"\"Gated recurrent unit (GRU) cell.\n\n Args:\n carry: the hidden state of the GRU cell,\n initialized using ``GRUCell.initialize_carry``.\n inputs: an ndarray with the input for the current time step.\n All dimensions except the final are considered batch dimensions.\n\n Returns:\n A tuple with the new carry and the output.\n \"\"\"\n h = carry\n\n # Compute combined transformations for inputs and hidden state\n x_transformed = self.dense_i(inputs)\n h_transformed = self.dense_h(h)\n\n # Split the combined transformations into individual components\n xi_r, xi_z, xi_n = jnp.split(x_transformed, 3, axis=-1)\n hh_r, hh_z, hh_n = jnp.split(h_transformed, 3, axis=-1)\n\n # Compute gates\n r = self.gate_fn(xi_r + hh_r)\n z = self.gate_fn(xi_z + hh_z)\n\n # Compute n with an additional linear transformation on h\n n = self.activation_fn(xi_n + r * hh_n)\n\n # Update hidden state\n new_h = (1.0 - z) * n + z * h\n return new_h, new_h\n\n def initialize_carry(\n self,\n input_shape: tuple[int, ...],\n rngs: rnglib.Rngs | rnglib.RngStream | None = None,\n carry_init: Initializer | None = None,\n ) -> Array: # type: ignore[override]\n \"\"\"Initialize the RNN cell carry.\n\n Args:\n rngs: random number generator passed to the init_fn.\n input_shape: a tuple providing the shape of the input to the cell.\n\n Returns:\n An initialized carry for the given RNN cell.\n \"\"\"\n batch_dims = input_shape[:-1]\n if rngs is None:\n rngs = self.rngs\n if isinstance(rngs, rnglib.Rngs):\n rngs = rngs.carry\n if rngs is None:\n raise ValueError('RNGs must be provided to initialize the cell carry.')\n\n mem_shape = batch_dims + (self.hidden_features,)\n\n if self.carry_init is None and carry_init is None:\n carry_init = initializers.zeros_init()\n elif carry_init is None:\n carry_init = self.carry_init\n assert carry_init is not None # just to please mypy\n\n h = carry_init(rngs(), mem_shape, self.param_dtype)\n return h\n\n @property\n def num_feature_axes(self) -> int:\n return 1\n\n\nclass RNN(Module):\n \"\"\"The ``RNN`` module takes any :class:`RNNCellBase` instance and applies it over a sequence\n\n using :func:`flax.nnx.scan`.\n \"\"\"\n\n state_axes: dict[str, int | type[iteration.Carry] | None]\n\n def __init__(\n self,\n cell: RNNCellBase,\n *,\n time_major: bool = False,\n return_carry: bool = False,\n reverse: bool = False,\n keep_order: bool = False,\n unroll: int = 1,\n state_axes: Mapping[str, int | type[iteration.Carry] | None] | None = None,\n broadcast_rngs: filterlib.Filter = None,\n rngs: rnglib.Rngs | rnglib.RngStream | bool = True,\n ):\n self.cell = cell\n self.time_major = time_major\n self.return_carry = return_carry\n self.reverse = reverse\n self.keep_order = keep_order\n self.unroll = unroll\n self.rngs: rnglib.RngStream | None\n if rngs is True:\n self.rngs = rnglib.RngStream(0, tag='carry')\n elif isinstance(rngs, rnglib.Rngs):\n self.rngs = rngs.carry.fork()\n elif rngs is False:\n self.rngs = nnx.data(None)\n else:\n raise ValueError(\n 'Expected rngs to be a jax.Array, int, Rngs, or bool. '\n f'Got {type(rngs)}.'\n )\n self.state_axes = state_axes or nnx.StateAxes({...: iteration.Carry}) # type: ignore\n self.broadcast_rngs = broadcast_rngs\n\n def __call__(\n self,\n inputs: Array,\n *,\n initial_carry: Carry | None = None,\n seq_lengths: Array | None = None,\n return_carry: bool | None = None,\n time_major: bool | None = None,\n reverse: bool | None = None,\n keep_order: bool | None = None,\n rngs: rnglib.Rngs | rnglib.RngStream | None = None,\n ):\n if return_carry is None:\n return_carry = self.return_carry\n if time_major is None:\n time_major = self.time_major\n if reverse is None:\n reverse = self.reverse\n if keep_order is None:\n keep_order = self.keep_order\n\n # Infer the number of batch dimensions from the input shape.\n # Cells like ConvLSTM have additional spatial dimensions.\n time_axis = 0 if time_major else inputs.ndim - (self.cell.num_feature_axes + 1)\n\n # make time_axis positive\n if time_axis < 0:\n time_axis += inputs.ndim\n\n if time_major:\n # we add +1 because we moved the time axis to the front\n batch_dims = inputs.shape[1 : -self.cell.num_feature_axes]\n else:\n batch_dims = inputs.shape[:time_axis]\n\n # maybe reverse the sequence\n if reverse:\n inputs = jax.tree_util.tree_map(\n lambda x: flip_sequences(\n x,\n seq_lengths,\n num_batch_dims=len(batch_dims),\n time_major=time_major, # type: ignore\n ),\n inputs,\n )\n if rngs is None:\n rngs = self.rngs\n if isinstance(rngs, rnglib.Rngs):\n rngs = rngs.carry.fork()\n\n carry: Carry = (\n self.cell.initialize_carry(\n inputs.shape[:time_axis] + inputs.shape[time_axis + 1 :], rngs\n )\n if initial_carry is None\n else initial_carry\n )\n\n slice_carry = seq_lengths is not None and return_carry\n broadcast_rngs = nnx.All(nnx.RngState, self.broadcast_rngs)\n state_axes = iteration.StateAxes({broadcast_rngs: None, **self.state_axes}) # type: ignore[misc]\n\n # we use split_rngs with splits=1 and squeeze=True to get unique rngs\n # every time RNN is called\n @nnx.split_rngs(splits=1, only=self.broadcast_rngs, squeeze=True)\n @nnx.scan(\n in_axes=(state_axes, iteration.Carry, time_axis),\n out_axes=(iteration.Carry, (0, time_axis))\n if slice_carry\n else (iteration.Carry, time_axis),\n unroll=self.unroll,\n )\n def scan_fn(\n cell: RNNCellBase, carry: Carry, x: Array\n ) -> tuple[Carry, Array] | tuple[Carry, tuple[Carry, Array]]:\n carry, y = cell(carry, x)\n if slice_carry:\n return carry, (carry, y)\n return carry, y\n\n scan_output = scan_fn(self.cell, carry, inputs)\n\n # Next we select the final carry. If a segmentation mask was provided and\n # return_carry is True we slice the carry history and select the last valid\n # carry for each sequence. Otherwise we just use the last carry.\n if slice_carry:\n assert seq_lengths is not None\n _, (carries, outputs) = scan_output\n # seq_lengths[None] expands the shape of the mask to match the\n # number of dimensions of the carry.\n carry = _select_last_carry(carries, seq_lengths)\n else:\n carry, outputs = scan_output\n\n if reverse and keep_order:\n outputs = jax.tree_util.tree_map(\n lambda x: flip_sequences(\n x,\n seq_lengths,\n num_batch_dims=len(batch_dims),\n time_major=time_major, # type: ignore\n ),\n outputs,\n )\n\n if return_carry:\n return carry, outputs\n else:\n return outputs\n\n\ndef _select_last_carry(sequence: A, seq_lengths: jnp.ndarray) -> A:\n last_idx = seq_lengths - 1\n\n def _slice_array(x: jnp.ndarray):\n return x[last_idx, jnp.arange(x.shape[1])]\n\n return jax.tree_util.tree_map(_slice_array, sequence)\n\n\ndef _expand_dims_like(x, target):\n \"\"\"Expands the shape of `x` to match `target`'s shape by adding singleton dimensions.\"\"\"\n return x.reshape(list(x.shape) + [1] * (target.ndim - x.ndim))\n\n\ndef flip_sequences(\n inputs: Array,\n seq_lengths: Array | None,\n num_batch_dims: int,\n time_major: bool,\n) -> Array:\n \"\"\"Flips a sequence of inputs along the time axis.\n\n This function can be used to prepare inputs for the reverse direction of a\n bidirectional LSTM. It solves the issue that, when naively flipping multiple\n padded sequences stored in a matrix, the first elements would be padding\n values for those sequences that were padded. This function keeps the padding\n at the end, while flipping the rest of the elements.\n\n Example::\n\n >>> from flax.nnx.nn.recurrent import flip_sequences\n >>> from jax import numpy as jnp\n >>> inputs = jnp.array([[1, 0, 0], [2, 3, 0], [4, 5, 6]])\n >>> lengths = jnp.array([1, 2, 3])\n >>> flip_sequences(inputs, lengths, 1, False)\n Array([[1, 0, 0],\n [3, 2, 0],\n [6, 5, 4]], dtype=int32)\n\n\n Args:\n inputs: An array of input IDs [batch_size, seq_length].\n lengths: The length of each sequence [batch_size].\n\n Returns:\n An ndarray with the flipped inputs.\n \"\"\"\n # Compute the indices to put the inputs in flipped order as per above example.\n time_axis = 0 if time_major else num_batch_dims\n max_steps = inputs.shape[time_axis]\n\n if seq_lengths is None:\n # reverse inputs and return\n inputs = jnp.flip(inputs, axis=time_axis)\n return inputs\n\n seq_lengths = jnp.expand_dims(seq_lengths, axis=time_axis)\n\n # create indexes\n idxs = jnp.arange(max_steps - 1, -1, -1) # [max_steps]\n if time_major:\n idxs = jnp.reshape(idxs, [max_steps] + [1] * num_batch_dims)\n else:\n idxs = jnp.reshape(\n idxs, [1] * num_batch_dims + [max_steps]\n ) # [1, ..., max_steps]\n idxs = (idxs + seq_lengths) % max_steps # [*batch, max_steps]\n idxs = _expand_dims_like(idxs, target=inputs) # [*batch, max_steps, *features]\n # Select the inputs in flipped order.\n outputs = jnp.take_along_axis(inputs, idxs, axis=time_axis)\n\n return outputs\n\n\ndef _concatenate(a: Array, b: Array) -> Array:\n \"\"\"Concatenates two arrays along the last dimension.\"\"\"\n return jnp.concatenate([a, b], axis=-1)\n\n\nclass RNNBase(Protocol):\n def __call__(\n self,\n inputs: Array,\n *,\n initial_carry: Carry | None = None,\n rngs: rnglib.Rngs | rnglib.RngStream | None = None,\n seq_lengths: Array | None = None,\n return_carry: bool | None = None,\n time_major: bool | None = None,\n reverse: bool | None = None,\n keep_order: bool | None = None,\n ) -> Output | tuple[Carry, Output]: ...\n\n\nclass Bidirectional(Module):\n \"\"\"Processes the input in both directions and merges the results.\n\n Example usage::\n\n >>> from flax import nnx\n >>> import jax\n >>> import jax.numpy as jnp\n\n >>> # Define forward and backward RNNs\n >>> forward_rnn = RNN(GRUCell(in_features=3, hidden_features=4, rngs=nnx.Rngs(0)))\n >>> backward_rnn = RNN(GRUCell(in_features=3, hidden_features=4, rngs=nnx.Rngs(0)))\n\n >>> # Create Bidirectional layer\n >>> layer = Bidirectional(forward_rnn=forward_rnn, backward_rnn=backward_rnn)\n\n >>> # Input data\n >>> x = jnp.ones((2, 3, 3))\n\n >>> # Apply the layer\n >>> out = layer(x)\n >>> print(out.shape)\n (2, 3, 8)\n\n \"\"\"\n\n forward_rnn: RNNBase\n backward_rnn: RNNBase\n merge_fn: Callable[[Array, Array], Array] = _concatenate\n time_major: bool = False\n return_carry: bool = False\n\n def __init__(\n self,\n forward_rnn: RNNBase,\n backward_rnn: RNNBase,\n *,\n merge_fn: Callable[[Array, Array], Array] = _concatenate,\n time_major: bool = False,\n return_carry: bool = False,\n rngs: rnglib.Rngs | rnglib.RngStream | bool = True,\n ):\n self.forward_rnn = forward_rnn\n self.backward_rnn = backward_rnn\n self.merge_fn = merge_fn\n self.time_major = time_major\n self.return_carry = return_carry\n self.rngs: rnglib.RngStream | None\n if rngs is True:\n self.rngs = rnglib.RngStream(0, tag='carry')\n elif rngs is False:\n self.rngs = None\n elif isinstance(rngs, rnglib.Rngs):\n self.rngs = rngs.carry.fork()\n elif isinstance(rngs, rnglib.RngStream):\n self.rngs = rngs\n else:\n raise TypeError(\n f'rngs must be a Rngs, jax.Array, int, or bool, but got {type(rngs)}.'\n )\n\n def __call__(\n self,\n inputs: Array,\n *,\n initial_carry: tuple[Carry, Carry] | None = None,\n rngs: rnglib.Rngs | rnglib.RngStream | None = None,\n seq_lengths: Array | None = None,\n return_carry: bool | None = None,\n time_major: bool | None = None,\n reverse: bool | None = None, # unused\n keep_order: bool | None = None, # unused\n ) -> Output | tuple[tuple[Carry, Carry], Output]:\n if time_major is None:\n time_major = self.time_major\n if return_carry is None:\n return_carry = self.return_carry\n if rngs is None:\n rngs = self.rngs\n if isinstance(rngs, rnglib.Rngs):\n rngs = rngs.carry\n\n if initial_carry is not None:\n initial_carry_forward, initial_carry_backward = initial_carry\n else:\n initial_carry_forward = None\n initial_carry_backward = None\n # Throw a warning in case the user accidentally re-uses the forward RNN\n # for the backward pass and does not intend for them to share parameters.\n if self.forward_rnn is self.backward_rnn:\n logging.warning(\n 'forward_rnn and backward_rnn is the same object, so '\n 'they will share parameters.'\n )\n\n # Encode in the forward direction.\n carry_forward, outputs_forward = self.forward_rnn(\n inputs,\n initial_carry=initial_carry_forward,\n rngs=rngs,\n seq_lengths=seq_lengths,\n return_carry=True,\n time_major=time_major,\n reverse=False,\n )\n\n # Encode in the backward direction.\n carry_backward, outputs_backward = self.backward_rnn(\n inputs,\n initial_carry=initial_carry_backward,\n rngs=rngs,\n seq_lengths=seq_lengths,\n return_carry=True,\n time_major=time_major,\n reverse=True,\n keep_order=True,\n )\n\n carry = (carry_forward, carry_backward) if return_carry else None\n outputs = jax.tree_util.tree_map(\n self.merge_fn, outputs_forward, outputs_backward\n )\n\n if return_carry:\n return carry, outputs\n else:\n return outputs\n" }, { "path": "flax/nnx/nn/stochastic.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nfrom collections.abc import Sequence\n\nimport jax\nimport jax.numpy as jnp\nfrom jax import lax, random\n\nfrom flax.nnx import rnglib\nfrom flax.nnx.module import Module, first_from\nfrom flax import nnx\n\n\nclass Dropout(Module):\n \"\"\"Create a dropout layer.\n\n To use dropout, call the :func:`train` method (or pass in\n ``deterministic=False`` in the constructor or during call time).\n\n To disable dropout, call the :func:`eval` method (or pass in\n ``deterministic=True`` in the constructor or during call time).\n\n Example usage::\n\n >>> from flax import nnx\n >>> import jax.numpy as jnp\n\n >>> class MLP(nnx.Module):\n ... def __init__(self, rngs):\n ... self.linear = nnx.Linear(in_features=3, out_features=4, rngs=rngs)\n ... self.dropout = nnx.Dropout(0.5, rngs=rngs)\n ... def __call__(self, x):\n ... x = self.linear(x)\n ... x = self.dropout(x)\n ... return x\n\n >>> model = MLP(rngs=nnx.Rngs(0))\n >>> x = jnp.ones((1, 3))\n\n >>> model.train() # use dropout\n >>> model(x)\n Array([[ 2.1067007, -2.5359864, -1.592019 , -2.5238838]], dtype=float32)\n\n >>> model.eval() # don't use dropout\n >>> model(x)\n Array([[ 1.0533503, -1.2679932, -0.7960095, -1.2619419]], dtype=float32)\n\n Args:\n rate: the dropout probability. (_not_ the keep rate!)\n broadcast_dims: dimensions that will share the same dropout mask\n deterministic: if false the inputs are scaled by ``1 / (1 - rate)`` and\n masked, whereas if true, no mask is applied and the inputs are returned\n as is.\n rng_collection: the rng collection name to use when requesting an rng key.\n rngs: rng key.\n \"\"\"\n\n def __init__(\n self,\n rate: float,\n *,\n broadcast_dims: Sequence[int] = (),\n deterministic: bool = False,\n rng_collection: str = 'dropout',\n rngs: rnglib.Rngs | rnglib.RngStream | None = None,\n ):\n self.rate = rate\n self.broadcast_dims = broadcast_dims\n self.deterministic = deterministic\n self.rng_collection = rng_collection\n\n if isinstance(rngs, rnglib.Rngs):\n self.rngs = rngs[self.rng_collection].fork()\n elif isinstance(rngs, rnglib.RngStream):\n self.rngs = rngs.fork()\n elif rngs is None:\n self.rngs = nnx.data(None)\n else:\n raise TypeError(\n f'rngs must be a Rngs, RngStream or None, but got {type(rngs)}.'\n )\n\n def __call__(\n self,\n inputs,\n *,\n deterministic: bool | None = None,\n rngs: rnglib.Rngs | rnglib.RngStream | jax.Array | None = None,\n ) -> jax.Array:\n \"\"\"Applies a random dropout mask to the input.\n\n Args:\n inputs: the inputs that should be randomly masked.\n deterministic: if false the inputs are scaled by ``1 / (1 - rate)`` and\n masked, whereas if true, no mask is applied and the inputs are returned\n as is. The ``deterministic`` flag passed into the call method will take\n precedence over the ``deterministic`` flag passed into the constructor.\n rngs: an optional key, RngStream, or Rngs object used to generate the dropout mask.\n If given it will take precedence over the rngs passed into the constructor.\n\n Returns:\n The masked inputs reweighted to preserve mean.\n \"\"\"\n deterministic = first_from(\n deterministic,\n self.deterministic,\n error_msg=\"\"\"No `deterministic` argument was provided to Dropout\n as either a __call__ argument or class attribute\"\"\",\n )\n\n if (self.rate == 0.0) or deterministic:\n return inputs\n\n # Prevent gradient NaNs in 1.0 edge-case.\n if self.rate == 1.0:\n return jnp.zeros_like(inputs)\n\n rngs = first_from( # type: ignore[assignment]\n rngs,\n self.rngs,\n error_msg=\"\"\"`deterministic` is False, but no `rngs` argument was provided to Dropout\n as either a __call__ argument or class attribute.\"\"\",\n )\n\n if isinstance(rngs, rnglib.Rngs):\n key = rngs[self.rng_collection]()\n elif isinstance(rngs, rnglib.RngStream):\n key = rngs()\n elif isinstance(rngs, jax.Array):\n key = rngs\n else:\n raise TypeError(\n f'rngs must be a Rngs, RngStream or jax.Array, but got {type(rngs)}.'\n )\n\n keep_prob = 1.0 - self.rate\n broadcast_shape = list(inputs.shape)\n for dim in self.broadcast_dims:\n broadcast_shape[dim] = 1\n mask = random.bernoulli(\n key, p=keep_prob, shape=broadcast_shape, out_sharding=jax.typeof(inputs).sharding\n )\n mask = jnp.broadcast_to(mask, inputs.shape)\n return lax.select(mask, inputs / keep_prob, jnp.zeros_like(inputs))\n\n def set_view(\n self,\n deterministic: bool | None = None,\n ):\n \"\"\"Class method used by ``nnx.view``.\n\n Args:\n deterministic: if True, disables dropout masking.\n \"\"\"\n if deterministic is not None:\n self.deterministic = deterministic\n" }, { "path": "flax/nnx/proxy_caller.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nimport dataclasses\nimport typing as tp\n\nimport jax\n\n\nA = tp.TypeVar('A', covariant=True) # type: ignore[not-supported-yet]\n\n\ndef _identity(x):\n return x\n\n@dataclasses.dataclass(frozen=True)\nclass GetItem:\n key: tp.Any\n\n\n@dataclasses.dataclass(frozen=True)\nclass GetAttr:\n name: str\n\n\n@dataclasses.dataclass(frozen=True)\nclass DelayedAccessor:\n actions: tuple[GetItem | GetAttr, ...] = ()\n\n def __call__(self, x):\n for action in self.actions:\n if isinstance(action, GetItem):\n x = x[action.key]\n elif isinstance(action, GetAttr):\n x = getattr(x, action.name)\n return x\n\n def __getattr__(self, name):\n return DelayedAccessor(self.actions + (GetAttr(name),))\n\n def __getitem__(self, key):\n return DelayedAccessor(self.actions + (GetItem(key),))\n\njax.tree_util.register_static(DelayedAccessor)\n\n\nclass _AccessorCall(tp.Protocol):\n def __call__(self, accessor: DelayedAccessor, /, *args, **kwargs) -> tp.Any:\n ...\n\nclass CallableProxy:\n def __init__(\n self, callable: _AccessorCall, accessor: DelayedAccessor | None = None\n ):\n self._callable = callable\n self._accessor = DelayedAccessor() if accessor is None else accessor\n\n def __call__(self, *args, **kwargs):\n return self._callable(self._accessor, *args, **kwargs)\n\n def __getattr__(self, name) -> CallableProxy:\n return CallableProxy(self._callable, getattr(self._accessor, name))\n\n def __getitem__(self, key) -> CallableProxy:\n return CallableProxy(self._callable, self._accessor[key])\n\n\nclass ApplyCaller(tp.Protocol, tp.Generic[A]):\n def __getattr__(self, __name) -> ApplyCaller[A]:\n ...\n\n def __getitem__(self, __name) -> ApplyCaller[A]:\n ...\n\n def __call__(self, *args, **kwargs) -> tuple[tp.Any, A]:\n ...\n" }, { "path": "flax/nnx/pytreelib.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom __future__ import annotations\n\nimport dataclasses\nimport inspect\nimport os\nimport threading\nimport typing as tp\nfrom abc import ABCMeta\nfrom copy import deepcopy\nimport warnings\n\nfrom flax.nnx import variablelib\nimport jax\nimport numpy as np\nimport treescope # type: ignore[import-untyped]\nfrom treescope import rendering_parts\n\nfrom flax import errors, nnx\nfrom flax.nnx import (\n graphlib,\n reprlib,\n tracers,\n visualization,\n)\nfrom flax import config\nfrom flax.nnx.variablelib import Variable\nfrom flax.typing import MISSING, Missing, SizeBytes\n\nBUILDING_DOCS = 'FLAX_DOC_BUILD' in os.environ\n\nA = tp.TypeVar('A')\nP = tp.TypeVar('P', bound='Pytree')\nT = tp.TypeVar('T', bound=type)\n\nDataAnnotation = '__data__'\nData = tp.Annotated[A, DataAnnotation]\nData.__doc__ = \"\"\"Data marks attributes of a class as pytree data using type annotations.\n\nData annotations must be used at the class level and will apply to all instances.\nThe usage of Data is recommended when type annotations are used already present\nor required e.g. for dataclasses.\n\"\"\"\nDATA_REGISTRY: set[type] = set()\n\n@tp.overload\ndef data(value: A, /) -> A: ...\n@tp.overload\ndef data(\n *,\n default: A = dataclasses.MISSING, # type: ignore[assignment]\n default_factory: tp.Callable[[], A] | None = None, # type: ignore[assignment]\n init: bool = True,\n repr: bool = True,\n hash: bool | None = None,\n compare: bool = True,\n metadata: tp.Mapping[str, tp.Any] | None = None,\n kw_only: bool = False,\n) -> tp.Any: ...\ndef data(value: tp.Any = MISSING, /, **kwargs) -> tp.Any:\n \"\"\"Annotates a an attribute as pytree data.\n\n The return value from `data` must be directly assigned to an Object attribute\n which will be registered as a pytree data attribute.\n\n Example::\n\n from flax import nnx\n import jax\n\n class Foo(nnx.Pytree):\n def __init__(self):\n self.data_attr = nnx.data(42) # pytree data\n self.static_attr = \"hello\" # static attribute\n\n foo = Foo()\n\n assert jax.tree.leaves(foo) == [42]\n\n Args:\n value: The value to annotate as data.\n\n Returns:\n A value which will register the attribute as data on assignment.\n\n \"\"\"\n if not isinstance(value, Missing) and kwargs:\n raise TypeError(\n 'nnx.data() accepts either a single positional argument or keyword'\n ' arguments, but not both.'\n )\n metadata = {'nnx_value': value}\n if 'metadata' in kwargs and kwargs['metadata'] is not None:\n if 'static' in kwargs['metadata']:\n raise ValueError(\n \"Cannot use 'static' key in metadata argument for nnx.data.\"\n )\n metadata.update(kwargs.pop('metadata'))\n metadata['static'] = False\n return dataclasses.field(**kwargs, metadata=metadata) # type: ignore[return-value]\n\n\ndef register_data_type(type_: T, /) -> T:\n \"\"\"Registers a type as pytree data type recognized by Object.\n\n Custom types registered as data will be automatically recognized\n as data attributes when assigned to an Object attribute. This means\n that values of this type do not need to be wrapped in `nnx.data(...)`\n for Object to mark the attribute its being assigned to as data.\n\n Example::\n\n from flax import nnx\n from dataclasses import dataclass\n\n @dataclass(frozen=True)\n class MyType:\n value: int\n\n nnx.register_data_type(MyType)\n\n class Foo(nnx.Pytree):\n def __init__(self, a):\n self.a = MyType(a) # Automatically registered as data\n self.b = \"hello\" # str not registered as data\n\n foo = Foo(42)\n\n assert nnx.is_data(foo.a) # True\n assert jax.tree.leaves(foo) == [MyType(value=42)]\n\n Can also be used as a decorator::\n\n @nnx.register_data_type\n @dataclass(frozen=True)\n class MyType:\n value: int\n \"\"\"\n DATA_REGISTRY.add(type_)\n return type_\n\n\ndef is_data(value: tp.Any, /) -> bool:\n \"\"\"Checks if a value is a registered data type.\n\n This function checks a the value is registered data type, which means it is\n automatically recognized as data when assigned a :class:`flax.nnx.Pytree` attribute.\n\n Data types are:\n\n - ``jax.Array``\n - ``np.ndarray``\n - ``ArrayRef``\n - Variables (:class:`flax.nnx.Param`, :class:`flax.nnx.BatchStat`, `nnx.RngState`, etc.)\n - All graph nodes (:class:`flax.nnx.Object`, :class:`flax.nnx.Module`, :class:`flax.nnx.Rngs`, etc.)\n - Any type registered with :func:`flax.nnx.register_data_type`\n - Any pytree that contains at least one node or leaf element of the above\n\n\n Example::\n\n >>> from flax import nnx\n >>> import jax.numpy as jnp\n ... # ------ DATA ------------\n >>> assert nnx.is_data( jnp.array(0) ) # Arrays\n >>> assert nnx.is_data( nnx.Param(1) ) # Variables\n >>> assert nnx.is_data( nnx.Rngs(2) ) # nnx.Pytrees\n >>> assert nnx.is_data( nnx.Linear(1, 1,rngs=nnx.Rngs(0)) ) # Modules\n ... # ------ STATIC ------------\n >>> assert not nnx.is_data( 'hello' ) # strings, arbitrary objects\n >>> assert not nnx.is_data( 42 ) # int, float, bool, complex, etc.\n >>> assert not nnx.is_data( [1, 2.0, 3j, jnp.array(1)] ) # list, dict, tuple, pytrees\n\n\n Args:\n value: The value to check.\n\n Returns:\n A string representing the attribute status.\n \"\"\"\n return (\n graphlib.is_node_leaf(value)\n or graphlib.is_graph_node(value)\n or type(value) in DATA_REGISTRY\n )\n\n\ndef has_data(value: tp.Any, /) -> list[tp.Any]:\n visited: set[int] = set()\n def _is_leaf(x):\n if id(x) in visited:\n return True\n visited.add(id(x))\n return is_data(x)\n leaves = jax.tree.leaves(value, is_leaf=_is_leaf)\n return [leaf for leaf in leaves if is_data(leaf)]\n\n\nStaticAnnotation = '__static__'\nStatic = tp.Annotated[A, StaticAnnotation]\nStatic.__doc__ = \"\"\"Static marks attributes of a class as static using type annotations.\nStatic annotations must be used at the class level and will apply to all instances.\nThe usage of Static is recommended when type annotations are used already present\nor required e.g. for dataclasses.\n\"\"\"\n\n@tp.overload\ndef static(value: A, /) -> A: ...\n@tp.overload\ndef static(\n *,\n default: A = dataclasses.MISSING, # type: ignore[assignment]\n default_factory: tp.Callable[[], A] | None = None,\n init: bool = True,\n repr: bool = True,\n hash: bool | None = None,\n compare: bool = True,\n metadata: tp.Mapping[str, tp.Any] | None = None,\n kw_only: bool = False,\n) -> tp.Any: ...\ndef static(value: tp.Any = MISSING, /, **kwargs) -> tp.Any:\n \"\"\"Annotates a an attribute as static.\n\n The return value from `static` must be directly assigned to an Object\n attribute\n which will be registered as static attribute.\n\n Example::\n\n from flax import nnx\n\n class Foo(nnx.Pytree):\n def __init__(self, a, b):\n self.a = nnx.static(a) # pytree metadata\n self.b = nnx.data(b) # pytree data\n\n foo = Foo(\"one\", \"two\")\n\n assert jax.tree.leaves(foo) == [\"two\"]\n\n By default ``nnx.Pytree`` will ...\n \"\"\"\n if not isinstance(value, Missing) and kwargs:\n raise TypeError(\n 'nnx.static() accepts either a single positional argument or keyword'\n ' arguments, but not both.'\n )\n metadata = {'nnx_value': value}\n if 'metadata' in kwargs and kwargs['metadata'] is not None:\n if 'static' in kwargs['metadata']:\n raise ValueError(\n \"Cannot use 'static' key in metadata argument for nnx.static.\"\n )\n metadata.update(kwargs.pop('metadata'))\n metadata['static'] = True\n return dataclasses.field(**kwargs, metadata=metadata) # type: ignore[return-value]\n\n@tp.overload\ndef dataclass(cls: type[A], /) -> type[A]: ...\n@tp.overload\ndef dataclass(\n *,\n init: bool = True,\n eq: bool = True,\n order: bool = False,\n unsafe_hash: bool = False,\n match_args: bool = True,\n kw_only: bool = False,\n slots: bool = False,\n weakref_slot: bool = False,\n) -> tp.Callable[[type[A]], type[A]]: ...\n\n@tp.dataclass_transform(field_specifiers=(dataclasses.field, data, static))\ndef dataclass(\n cls=None,\n /,\n *,\n init: bool = True,\n eq: bool = True,\n order: bool = False,\n unsafe_hash: bool = False,\n match_args: bool = True,\n kw_only: bool = False,\n slots: bool = False,\n weakref_slot: bool = False,\n) -> tp.Any:\n return dataclasses.dataclass(\n cls,\n init=init,\n eq=eq,\n order=order,\n unsafe_hash=unsafe_hash,\n match_args=match_args,\n kw_only=kw_only,\n slots=slots,\n weakref_slot=weakref_slot,\n )\n\ndef _collect_stats(\n node: tp.Any, node_stats: dict[int, dict[type[Variable], SizeBytes]]\n):\n if not graphlib.is_node(node) and not isinstance(node, Variable):\n raise ValueError(f'Expected a graph node or Variable, got {type(node)!r}.')\n\n if id(node) in node_stats:\n return\n\n stats: dict[type[Variable], SizeBytes] = {}\n node_stats[id(node)] = stats\n\n if isinstance(node, Variable):\n var_type = node.var_type\n if issubclass(var_type, nnx.RngState):\n var_type = nnx.RngState\n size_bytes = SizeBytes.from_any(node.get_raw_value())\n if size_bytes:\n stats[var_type] = size_bytes\n\n else:\n node_impl = graphlib.get_node_impl(node)\n assert node_impl is not None\n node_dict = node_impl.node_dict(node)\n for key, value in node_dict.items():\n if id(value) in node_stats:\n continue\n if graphlib.is_node(value) or isinstance(value, Variable):\n _collect_stats(value, node_stats)\n child_stats = node_stats[id(value)]\n for var_type, size_bytes in child_stats.items():\n if var_type in stats:\n stats[var_type] += size_bytes\n else:\n stats[var_type] = size_bytes\n\n\n@dataclasses.dataclass\nclass ObjectContext(threading.local):\n seen_modules_repr: set[int] | None = None\n node_stats: dict[int, dict[type[Variable], SizeBytes]] | None = None\n\n\nOBJECT_CONTEXT = ObjectContext()\n\n\nclass PytreeState(reprlib.Representable):\n __slots__ = ('_trace_state', '_initializing', '_is_setup')\n\n def __init__(self, initializing: bool = False, is_setup: bool = False):\n self._trace_state = tracers.TraceState()\n self._initializing = initializing\n self._is_setup = is_setup\n\n @property\n def trace_state(self) -> tracers.TraceState:\n return self._trace_state\n\n @property\n def initializing(self) -> bool:\n return self._initializing\n\n @property\n def is_setup(self) -> bool:\n return self._is_setup\n\n def __nnx_repr__(self):\n yield reprlib.Object(type(self))\n yield reprlib.Attr('trace_state', self._trace_state)\n\n def __treescope_repr__(self, path, subtree_renderer):\n return visualization.render_object_constructor(\n object_type=type(self),\n attributes={'trace_state': self._trace_state},\n path=path,\n subtree_renderer=subtree_renderer,\n )\n\n\ndef _flatten_pytree_state(state: PytreeState):\n return (), (state.initializing, state.is_setup)\n\n\ndef _unflatten_pytree_state(static: tuple[bool, bool], _):\n initializing, setup = static\n return PytreeState(initializing, setup)\n\n\njax.tree_util.register_pytree_node(\n PytreeState,\n _flatten_pytree_state,\n _unflatten_pytree_state,\n)\n\n\ndef check_pytree(pytree):\n \"\"\"Checks if a pytree is valid.\"\"\"\n if not isinstance(pytree, Pytree):\n raise TypeError(f'Expected a Pytree, got {type(pytree)}.')\n\n for name, value in vars(pytree).items():\n pytree._check_value(name, value, new_status=None)\n\n\nclass PytreeMeta(ABCMeta):\n if not tp.TYPE_CHECKING:\n\n def __call__(cls, *args: Any, **kwargs: Any) -> Any:\n return _graph_node_meta_call(cls, *args, **kwargs)\n\n def _pytree_meta_construct(cls, self, *args, **kwargs):\n self.__init__(*args, **kwargs)\n\nObjectMeta = PytreeMeta\n\ndef _graph_node_meta_call(cls: tp.Type[P], *args, **kwargs) -> P:\n node = cls.__new__(cls, *args, **kwargs)\n object.__setattr__(node, '_pytree__state', PytreeState())\n object.__setattr__(node, '_pytree__nodes', cls._pytree__nodes)\n cls._pytree_meta_construct(node, *args, **kwargs)\n if cls._pytree__is_pytree:\n missing: dict[str, bool] = {}\n for name, value in vars(node).items():\n if name not in node._pytree__nodes:\n missing[name] = is_data(value)\n if missing:\n object.__setattr__(\n node, '_pytree__nodes', node._pytree__nodes.update(missing)\n )\n check_pytree(node)\n\n return node\n\n\n@dataclasses.dataclass(frozen=True, repr=False)\nclass ArrayRepr(reprlib.Representable):\n shape: tp.Tuple[int, ...]\n dtype: tp.Any\n\n @staticmethod\n def from_array(array: jax.Array | np.ndarray) -> ArrayRepr:\n return ArrayRepr(array.shape, array.dtype)\n\n def __nnx_repr__(self):\n yield reprlib.Object(type='Array', same_line=True)\n yield reprlib.Attr('shape', self.shape)\n yield reprlib.Attr('dtype', self.dtype)\n\n@dataclasses.dataclass(frozen=True, repr=False)\nclass VariableRepr(reprlib.Representable):\n var_type: type[Variable]\n value: tp.Any\n metadata: dict[str, tp.Any]\n\n def __nnx_repr__(self):\n variable = self.var_type._new(self.value, self.metadata)\n yield from variable.__nnx_repr__()\n\n\n@dataclasses.dataclass(frozen=True, repr=False)\nclass MutableArrayRepr(reprlib.Representable):\n shape: tp.Tuple[int, ...]\n dtype: tp.Any\n\n @staticmethod\n def from_array(array: jax.Array | np.ndarray) -> MutableArrayRepr:\n return MutableArrayRepr(array.shape, array.dtype)\n\n def __nnx_repr__(self):\n yield reprlib.Object(type='ArrayRef', same_line=True)\n yield reprlib.Attr('shape', self.shape)\n yield reprlib.Attr('dtype', self.dtype)\n\ndef _to_shape_dtype(x):\n if isinstance(x, Variable):\n value = x.get_raw_value()\n metadata = x.get_metadata()\n value = jax.tree.map(_to_shape_dtype, value)\n return VariableRepr(x.var_type, value, metadata)\n elif variablelib.is_array_ref(x) and np.prod(x.shape) > 1:\n return MutableArrayRepr(x.shape, x.dtype)\n elif (\n isinstance(x, (np.ndarray, jax.Array))\n and np.prod(x.shape) > 1\n ):\n return ArrayRepr(x.shape, x.dtype)\n return x\n\nclass AttributeStatus(tp.NamedTuple):\n is_data: bool\n explicit: bool\n\n\nclass Pytree(reprlib.Representable, metaclass=PytreeMeta):\n \"\"\"Base class for all NNX objects.\"\"\"\n\n if tp.TYPE_CHECKING:\n _pytree__nodes: graphlib.HashableMapping[tp.Any, bool]\n _pytree__state: PytreeState\n _pytree__is_pytree: bool\n\n def __init_subclass__(\n cls,\n *,\n pytree: bool = config.flax_pytree_module,\n **kwargs,\n ) -> None:\n super().__init_subclass__(**kwargs)\n if slots := getattr(cls, '__slots__', ()):\n raise TypeError(\n 'Pytree currently does not support __slots__, '\n f\"found __slots__={slots} in '{cls.__name__}'.\"\n )\n cls._pytree__is_pytree = pytree\n\n graphlib.register_graph_node_type(\n type=cls,\n flatten=cls._graph_node_flatten,\n set_key=cls._graph_node_set_key, # type: ignore\n pop_key=cls._graph_node_pop_key, # type: ignore\n create_empty=cls._graph_node_create_empty,\n clear=cls._graph_node_clear,\n init=cls._graph_node_init, # type: ignore\n )\n\n nodes: dict[str, bool] = dict(getattr(cls, '_pytree__nodes', ()))\n nodes['_pytree__state'] = True\n try:\n type_hints = tp.get_type_hints(\n cls, globals(), {cls.__name__: cls}, include_extras=True\n )\n except NameError:\n type_hints = cls.__annotations__\n # add annotation attributes\n for name, type_ in type_hints.items():\n if isinstance(type_, str):\n if type_.startswith('nnx.Data'):\n warnings.warn(\n f\"'Data' is deprecated, please replace:\\n\\n\"\n ' some_field: nnx.Data[SomeType]\\n\\n'\n f'with:\\n\\n'\n ' some_field: SomeType = nnx.data()\\n\\n',\n DeprecationWarning,\n stacklevel=2,\n )\n nodes[name] = True\n elif type_.startswith('nnx.Static'):\n warnings.warn(\n f\"'Static' is deprecated, please replace:\\n\\n\"\n ' some_field: nnx.Static[SomeType]\\n\\n'\n f'with:\\n\\n'\n ' some_field: SomeType = nnx.static()\\n\\n',\n DeprecationWarning,\n stacklevel=2,\n )\n nodes[name] = False\n else:\n type_metadata = getattr(type_, '__metadata__', ())\n if DataAnnotation in type_metadata:\n warnings.warn(\n f\"'Data' is deprecated, please replace:\\n\\n\"\n ' some_field: nnx.Data[SomeType]\\n\\n'\n f'with:\\n\\n'\n ' some_field: SomeType = nnx.data()\\n\\n',\n DeprecationWarning,\n stacklevel=2,\n )\n nodes[name] = True\n elif StaticAnnotation in type_metadata:\n warnings.warn(\n f\"'Static' is deprecated, please replace:\\n\\n\"\n ' some_field: nnx.Static[SomeType]\\n\\n'\n f'with:\\n\\n'\n ' some_field: SomeType = nnx.static()\\n\\n',\n DeprecationWarning,\n stacklevel=2,\n )\n nodes[name] = False\n\n for name, value in vars(cls).items():\n if isinstance(value, dataclasses.Field) and 'static' in value.metadata:\n if not isinstance(value.metadata['static'], bool):\n raise ValueError(\n f\"Invalid 'static' metadata for attribute\"\n f\" '{cls.__name__}.{name}': expected bool, got\"\n f' {type(value.metadata[\"static\"]).__name__}.'\n )\n is_node = not value.metadata['static']\n if name in nodes and nodes[name] != is_node:\n raise ValueError(\n f'Conflicting pytree annotation for attribute'\n f\" '{cls.__name__}.{name}': previously registered as\"\n f' {\"data\" if nodes[name] else \"static\"}, but found'\n f' nnx.{\"data\" if is_node else \"static\"}(...) annotation.'\n )\n nodes[name] = is_node\n\n cls._pytree__nodes = graphlib.HashableMapping(nodes, copy=False)\n\n if pytree:\n jax.tree_util.register_pytree_with_keys(\n cls,\n flatten_with_keys=cls._pytree__flatten_with_paths,\n unflatten_func=cls._pytree__unflatten,\n flatten_func=cls._pytree__flatten,\n )\n\n if BUILDING_DOCS:\n # set correct signature for sphinx\n cls.__signature__ = inspect.signature(cls.__init__)\n\n # Backward compatibility with PR #4863\n @property\n def _object__nodes(self):\n warnings.warn(\n \"'_object__nodes' is deprecated, use '_pytree__nodes' instead.\",\n DeprecationWarning,\n stacklevel=2,\n )\n return self._pytree__nodes\n\n @property\n def _object__state(self):\n warnings.warn(\n \"'_object__state' is deprecated, use '_pytree__state' instead.\",\n DeprecationWarning,\n stacklevel=2,\n )\n return self._pytree__state\n\n if not tp.TYPE_CHECKING:\n\n def __setattr__(self, name: str, value: Any) -> None:\n self._setattr(name, value)\n\n def _setattr(self, name, value: tp.Any) -> None:\n self._check_valid_context(\n lambda: f\"Cannot mutate '{type(self).__name__}' from different trace level\"\n )\n data: bool = False\n explicit: bool = False\n if isinstance(value, dataclasses.Field) and 'nnx_value' in value.metadata:\n is_static = value.metadata['static']\n value = value.metadata['nnx_value']\n if self._pytree__is_pytree:\n data = not is_static\n explicit = True\n elif self._pytree__is_pytree:\n data = is_data(value)\n if self._pytree__is_pytree:\n self._check_value(name, value, AttributeStatus(data, explicit))\n if name not in self._pytree__nodes or (\n explicit and self._pytree__nodes[name] != data\n ):\n object.__setattr__(\n self, '_pytree__nodes', self._pytree__nodes.update({name: data})\n )\n object.__setattr__(self, name, value)\n\n def _check_value(self, key, value, new_status: AttributeStatus | None):\n\n def _get_annotations(leaves):\n return {\n 'static' if leaf.metadata['static'] else 'data'\n for leaf in leaves\n if isinstance(leaf, dataclasses.Field) and 'nnx_value' in leaf.metadata\n }\n\n def _has_visited(x):\n if id(x) in visited:\n return True\n visited.add(id(x))\n return False\n\n current_is_data = (\n self._pytree__nodes[key] if key in self._pytree__nodes else False\n )\n existing_attr = key in vars(self)\n if (\n new_status is not None\n and not new_status.explicit\n and new_status.is_data\n and existing_attr\n and not current_is_data\n ):\n raise ValueError(\n f\"Cannot assign data value of type '{type(value)}' to static\"\n f\" attribute '{key}' of Pytree type '{type(self)}'. To override the\"\n ' status explicitly wrap the value with nnx.data on assignment:\\n\\n '\n f' _.{key} = nnx.data(...)\\n\\n'\n )\n\n visited: set[int] = set()\n leaves = jax.tree.leaves(\n value, is_leaf=lambda x: _has_visited(x) or is_data(x)\n )\n data_leaves = [leaf for leaf in leaves if is_data(leaf)]\n if data_leaves:\n # check no data in nnx.static assignments\n if new_status is not None:\n if not new_status.is_data and new_status.explicit:\n raise ValueError(\n f\"Found data in value of type '{type(value)}' annotated with \"\n f\"nnx.static(...) when setting attribute '{key}' of Pytree type \"\n f\"'{type(self)}'.\"\n )\n if not new_status.is_data and not current_is_data:\n base_pytree_type = Pytree\n for t in type(self).mro()[1:]:\n if issubclass(t, nnx.Pytree):\n base_pytree_type = t\n break\n data_leaves_type_names = {type(leaf).__name__ for leaf in data_leaves}\n raise ValueError(\n f\"Found data on value of type '{type(value)}' assigned to\"\n f\" static attribute '{key}' of Pytree type '{type(self)}'. Static\"\n ' attributes should not contain data values. Consider one of'\n ' the following options:\\n\\n1. If the attribute is meant to be'\n f' static, remove the data values of type(s):\\n\\n {\", \".join(data_leaves_type_names)}'\n f'\\n\\n2. If the attribute is meant to be data, wrap the value with nnx.data '\n f' on assignment:\\n\\n _.{key} = nnx.data(...)\\n\\n3. Alternatively,'\n ' annotate the class attribute with nnx.data:\\n\\n class'\n f' {type(self).__name__}({base_pytree_type.__name__}):\\n '\n f' {key}: {type(value).__name__} = nnx.data()\\n\\n4. If the'\n ' container is a list or dict, try using nnx.List(...) or'\n ' nnx.Dict(...) instead.\\n\\n5. Disable pytree'\n ' for this class:\\n\\n class'\n f' {type(self).__name__}({base_pytree_type.__name__},'\n ' pytree=False):\\n\\n'\n )\n # check no data in static attributes after __init__\n elif not current_is_data:\n base_pytree_type = Pytree\n for t in type(self).mro()[1:]:\n if issubclass(t, nnx.Pytree):\n base_pytree_type = t\n break\n raise ValueError(\n f'Found unexpected data on value of type {type(value)} in static'\n f\" attribute '{key}' of Pytree type '{type(self)}'. This is an\"\n ' error starting from Flax version 0.12.0.\\nConsider one of the'\n ' following options:\\n\\n1. If the attribute is meant to be static,'\n ' either remove the data value or wrap it in a static'\n ' container.\\n2. Wrap the value with nnx.data on'\n f' assignment:\\n\\n _.{key} = nnx.data(...)\\n\\n3. Annotate the'\n ' class attribute with nnx.data:\\n\\n class'\n f' {type(self).__name__}({base_pytree_type.__name__}):\\n {key}:'\n f' {type(value).__name__} = nnx.data()\\n\\n4. If the container is a'\n ' list or dict, try using nnx.List(...) or nnx.Dict(...)'\n ' instead.\\n5. Disable pytree for this class:\\n\\n class'\n f' {type(self).__name__}({base_pytree_type.__name__},'\n f' pytree=False):\\n\\n'\n )\n if tags := _get_annotations(leaves):\n raise ValueError(\n f'Found unexpected tags {tags} on attribute'\n f\" '{type(self).__name__}.{key}'. Values from nnx.data(...)\"\n ' and\\nnnx.static(...) should be assigned to nnx.Pytree attributes'\n ' directly, they should not be inside other structures. Got value of'\n f\" type '{type(value)}' on Pytree of type '{type(self)}'.\"\n )\n\n def _check_valid_context(self, error_msg: tp.Callable[[], str]) -> None:\n if not self._pytree__state.trace_state.is_valid():\n raise errors.TraceContextError(error_msg())\n\n def __deepcopy__(self: P, memo=None) -> P:\n graphdef, state = graphlib.split(self, graph=True)\n graphdef = deepcopy(graphdef)\n state = deepcopy(state)\n return graphlib.merge(graphdef, state)\n\n def __nnx_repr__(self):\n if OBJECT_CONTEXT.node_stats is None or id(self) not in OBJECT_CONTEXT.node_stats:\n node_stats: dict[int, dict[type[Variable], SizeBytes]] = {}\n _collect_stats(self, node_stats)\n OBJECT_CONTEXT.node_stats = node_stats\n stats = node_stats[id(self)]\n clear_node_stats = True\n else:\n stats = OBJECT_CONTEXT.node_stats[id(self)]\n clear_node_stats = False\n\n if OBJECT_CONTEXT.seen_modules_repr is None:\n OBJECT_CONTEXT.seen_modules_repr = set()\n clear_seen = True\n else:\n clear_seen = False\n\n if id(self) in OBJECT_CONTEXT.seen_modules_repr:\n yield reprlib.Object(type=type(self), empty_repr='...')\n return\n\n try:\n if stats:\n stats_repr = ' # ' + ', '.join(\n f'{var_type.__name__}: {size_bytes}'\n for var_type, size_bytes in stats.items()\n )\n if len(stats) > 1:\n total_bytes = sum(stats.values(), SizeBytes(0, 0))\n stats_repr += f', Total: {total_bytes}'\n else:\n stats_repr = ''\n\n yield reprlib.Object(type=type(self), comment=stats_repr)\n OBJECT_CONTEXT.seen_modules_repr.add(id(self))\n\n for name, value in vars(self).items():\n if str(name).startswith('_pytree__'):\n continue\n if str(name).startswith('_') and not self._pytree__nodes.get(str(name), False):\n continue\n\n value = jax.tree.map(_to_shape_dtype, value, is_leaf=graphlib.is_graph_node)\n yield reprlib.Attr(name, value)\n finally:\n if clear_seen:\n OBJECT_CONTEXT.seen_modules_repr = None\n if clear_node_stats:\n OBJECT_CONTEXT.node_stats = None\n\n def __treescope_repr__(self, path, subtree_renderer):\n from flax import nnx\n\n if OBJECT_CONTEXT.node_stats is None or id(self) not in OBJECT_CONTEXT.node_stats:\n node_stats: dict[int, dict[type[Variable], SizeBytes]] = {}\n _collect_stats(self, node_stats)\n OBJECT_CONTEXT.node_stats = node_stats\n stats = node_stats[id(self)]\n clear_node_stats = True\n else:\n stats = OBJECT_CONTEXT.node_stats[id(self)]\n clear_node_stats = False\n\n try:\n if stats:\n stats_repr = ' # ' + ', '.join(\n f'{var_type.__name__}: {size_bytes}'\n for var_type, size_bytes in stats.items()\n )\n if len(stats) > 1:\n total_bytes = sum(stats.values(), SizeBytes(0, 0))\n stats_repr += f', Total: {total_bytes}'\n\n first_line_annotation = rendering_parts.comment_color(\n rendering_parts.text(f'{stats_repr}')\n )\n else:\n first_line_annotation = None\n children = {}\n for name, value in vars(self).items():\n if str(name).startswith('_pytree__'):\n continue\n if str(name).startswith('_') and not self._pytree__nodes.get(str(name), False):\n continue\n children[name] = value\n\n if isinstance(self, nnx.Module):\n color = treescope.formatting_util.color_from_string(\n type(self).__qualname__\n )\n else:\n color = None\n return visualization.render_object_constructor(\n object_type=type(self),\n attributes=children,\n path=path,\n subtree_renderer=subtree_renderer,\n first_line_annotation=first_line_annotation,\n color=color,\n )\n finally:\n if clear_node_stats:\n OBJECT_CONTEXT.node_stats = None\n\n # pickle support\n def __getstate__(self):\n return vars(self).copy()\n\n def __setstate__(self, state):\n vars(self).update(state)\n\n # -------------------------\n # Pytree Definition\n # -------------------------\n _pytree__has_int_keys: bool = False\n\n def _pytree__flatten_with_paths(self):\n obj_items = vars(self).items()\n if self._pytree__has_int_keys:\n obj_items = ((_maybe_int(name), value) for name, value in obj_items)\n key_fn = graphlib._type_aware_sort\n else:\n key_fn = None\n node_attributes = self._pytree__nodes\n node_keys: list[str | int] = []\n node_attrs: list[tuple[tp.Any, tp.Any]] = []\n static_keys: list[str | int] = []\n static_attrs: list[tp.Any] = []\n for key, value in sorted(obj_items, key=key_fn):\n # get string representation of the key because\n # node_attributes keys are strings\n key_str = _get_str(key)\n if key_str in node_attributes and node_attributes[key_str]:\n node_keys.append(key)\n node_attrs.append((\n jax.tree_util.GetAttrKey(key)\n if isinstance(key, str)\n else jax.tree_util.SequenceKey(key),\n value,\n ))\n else:\n static_keys.append(key)\n static_attrs.append(value)\n\n return (\n node_attrs,\n (tuple(node_keys), tuple(static_keys), tuple(static_attrs)),\n )\n\n def _pytree__flatten(self):\n obj_items = vars(self).items()\n if self._pytree__has_int_keys:\n obj_items = ((_maybe_int(name), value) for name, value in obj_items)\n key_fn = graphlib._type_aware_sort\n else:\n key_fn = None\n node_attributes = self._pytree__nodes\n node_keys: list[str | int] = []\n node_attrs: list[tp.Any] = []\n static_keys: list[str | int] = []\n static_attrs: list[tp.Any] = []\n for key, value in sorted(obj_items, key=key_fn):\n # get string representation of the key because\n # node_attributes keys are strings\n key_str = _get_str(key)\n if key_str in node_attributes and node_attributes[key_str]:\n node_keys.append(key)\n node_attrs.append(value)\n else:\n static_keys.append(key)\n static_attrs.append(value)\n\n return (\n node_attrs,\n (tuple(node_keys), tuple(static_keys), tuple(static_attrs)),\n )\n\n @classmethod\n def _pytree__unflatten(\n cls,\n static: tuple[tp.Iterable[str | int], tp.Iterable[str | int], tp.Iterable[tp.Any]],\n node_attrs: tp.Iterable[tp.Any],\n ):\n node_keys, static_keys, static_attrs = static\n obj = object.__new__(cls)\n for name, value in zip(node_keys, node_attrs, strict=True):\n object.__setattr__(obj, _get_str(name), value)\n for name, value in zip(static_keys, static_attrs, strict=True):\n object.__setattr__(obj, _get_str(name), value)\n return obj\n\n # -------------------------\n # Graph Definition\n # -------------------------\n def _graph_node_flatten(self):\n obj_items = vars(self).items()\n if self._pytree__is_pytree:\n pytree_nodes = self._pytree__nodes\n obj_items = (\n (\n name,\n nnx.graphlib.DataElem(value)\n if name in pytree_nodes and pytree_nodes[name]\n else nnx.graphlib.StaticElem(value),\n )\n for name, value in obj_items\n )\n if self._pytree__has_int_keys:\n obj_items = ((_maybe_int(name), value) for name, value in obj_items)\n key_fn = graphlib._type_aware_sort\n else:\n key_fn = None\n nodes = sorted(obj_items, key=key_fn)\n return nodes, type(self)\n\n def _graph_node_set_key(self, key, value: tp.Any):\n if self._pytree__has_int_keys and isinstance(key, int):\n key = str(key)\n if not isinstance(key, str):\n raise KeyError(f'Invalid key: {key!r}')\n elif (\n hasattr(self, key)\n and isinstance(variable := getattr(self, key), Variable)\n and isinstance(value, Variable)\n ):\n variable.update_from_state(value)\n else:\n setattr(self, key, value)\n\n def _graph_node_pop_key(self, key):\n if self._pytree__has_int_keys and isinstance(key, int):\n key = str(key)\n value = getattr(self, key)\n delattr(self, key)\n return value\n\n def __delattr__(self, name: str) -> None:\n if name in self._pytree__nodes:\n mapping = {k: v for k, v in self._pytree__nodes.items() if k != name}\n object.__setattr__(\n self, '_pytree__nodes', graphlib.HashableMapping(mapping, copy=False)\n )\n\n super().__delattr__(name)\n\n @staticmethod\n def _graph_node_create_empty(node_type: tp.Type[P]) -> P:\n node = object.__new__(node_type)\n return node\n\n def _graph_node_clear(self):\n vars(self).clear()\n\n def _graph_node_init(self, attributes: tp.Iterable[tuple[str | int, tp.Any]]):\n for name, value in attributes:\n object.__setattr__(self, _get_str(name), value)\n\n if tp.TYPE_CHECKING:\n def __call__(self, *args: tp.Any, **kwargs: tp.Any) -> tp.Any: ...\n\n\nclass Object(Pytree, pytree=False):\n \"\"\"Base class for NNX objects that are not pytrees.\"\"\"\n\n def __init_subclass__(cls, **kwargs):\n pytree = kwargs.pop('pytree', False)\n if pytree is not False:\n raise ValueError(\n \"Object is not a pytree, but 'pytree' was explicitly set to \"\n f'{pytree!r} for type {cls}.'\n )\n super().__init_subclass__(pytree=pytree, **kwargs)\n\ndef _maybe_int(x):\n try:\n return int(x)\n except (ValueError, TypeError):\n return x\n\ndef _get_str(x):\n return x if isinstance(x, str) else str(x)" }, { "path": "flax/nnx/reprlib.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport dataclasses\nimport os\nimport sys\nimport threading\nimport typing as tp\nfrom flax import config as flax_config\n\nA = tp.TypeVar('A')\nB = tp.TypeVar('B')\n\n\ndef supports_color() -> bool:\n \"\"\"\n Returns True if the running system's terminal supports color, and False otherwise.\n \"\"\"\n try:\n from IPython import get_ipython\n\n ipython_available = get_ipython() is not None\n except ImportError:\n ipython_available = False\n\n supported_platform = sys.platform != 'win32' or 'ANSICON' in os.environ\n is_a_tty = hasattr(sys.stdout, 'isatty') and sys.stdout.isatty()\n return (supported_platform and is_a_tty) or ipython_available\n\n\nclass Color(tp.NamedTuple):\n TYPE: str\n ATTRIBUTE: str\n SEP: str\n PAREN: str\n COMMENT: str\n INT: str\n STRING: str\n FLOAT: str\n BOOL: str\n NONE: str\n END: str\n\n\nNO_COLOR = Color(\n TYPE='',\n ATTRIBUTE='',\n SEP='',\n PAREN='',\n COMMENT='',\n INT='',\n STRING='',\n FLOAT='',\n BOOL='',\n NONE='',\n END='',\n)\n\n\n# Use python vscode theme colors\nif supports_color():\n COLOR = Color(\n TYPE='\\x1b[38;2;79;201;177m',\n ATTRIBUTE='\\033[38;2;156;220;254m',\n SEP='\\x1b[38;2;212;212;212m',\n PAREN='\\x1b[38;2;255;213;3m',\n # COMMENT='\\033[38;2;87;166;74m',\n COMMENT='\\033[38;2;105;105;105m', # Dark gray\n INT='\\x1b[38;2;182;207;169m',\n STRING='\\x1b[38;2;207;144;120m',\n FLOAT='\\x1b[38;2;182;207;169m',\n BOOL='\\x1b[38;2;86;156;214m',\n NONE='\\x1b[38;2;86;156;214m',\n END='\\x1b[0m',\n )\nelse:\n COLOR = NO_COLOR\n\n\n@dataclasses.dataclass\nclass ReprContext(threading.local):\n current_color: Color = COLOR\n depth: int = 0\n\n\nREPR_CONTEXT = ReprContext()\n\n\ndef colorized(x, /):\n c = REPR_CONTEXT.current_color\n if isinstance(x, list):\n return f'{c.PAREN}[{c.END}{\", \".join(map(lambda i: colorized(i), x))}{c.PAREN}]{c.END}'\n elif isinstance(x, tuple):\n if len(x) == 1:\n return f'{c.PAREN}({c.END}{colorized(x[0])},{c.PAREN}){c.END}'\n return f'{c.PAREN}({c.END}{\", \".join(map(lambda i: colorized(i), x))}{c.PAREN}){c.END}'\n elif isinstance(x, dict):\n open, close = '{', '}'\n return f'{c.PAREN}{open}{c.END}{\", \".join(f\"{c.STRING}{k!r}{c.END}: {colorized(v)}\" for k, v in x.items())}{c.PAREN}{close}{c.END}'\n elif isinstance(x, set):\n open, close = '{', '}'\n return f'{c.PAREN}{open}{c.END}{\", \".join(map(lambda i: colorized(i), x))}{c.PAREN}{close}{c.END}'\n elif isinstance(x, type):\n return f'{c.TYPE}{x.__name__}{c.END}'\n elif isinstance(x, bool):\n return f'{c.BOOL}{x}{c.END}'\n elif isinstance(x, int):\n return f'{c.INT}{x}{c.END}'\n elif isinstance(x, str):\n return f'{c.STRING}{x!r}{c.END}'\n elif isinstance(x, float):\n return f'{c.FLOAT}{x}{c.END}'\n elif x is None:\n return f'{c.NONE}{x}{c.END}'\n elif isinstance(x, Representable):\n return get_repr(x)\n else:\n return repr(x)\n\n\n@dataclasses.dataclass\nclass Object:\n type: tp.Union[str, type]\n start: str = '('\n end: str = ')'\n kv_sep: str = '='\n indent: str = ' '\n empty_repr: str = ''\n comment: str = ''\n same_line: bool = False\n\n @property\n def elem_sep(self):\n return ', ' if self.same_line else ',\\n'\n\n\n@dataclasses.dataclass\nclass Attr:\n key: str\n value: tp.Union[str, tp.Any]\n start: str = ''\n end: str = ''\n use_raw_value: bool = False\n use_raw_key: bool = False\n\n\nclass Representable:\n __slots__ = ()\n\n def __nnx_repr__(self) -> tp.Iterator[tp.Union[Object, Attr]]:\n raise NotImplementedError\n\n def __repr__(self) -> str:\n current_color = REPR_CONTEXT.current_color\n REPR_CONTEXT.current_color = NO_COLOR\n try:\n return get_repr(self)\n finally:\n REPR_CONTEXT.current_color = current_color\n\n def __str__(self) -> str:\n return get_repr(self)\n\n\ndef get_repr(obj: Representable) -> str:\n REPR_CONTEXT.depth += 1\n try:\n if not isinstance(obj, Representable):\n raise TypeError(f'Object {obj!r} is not representable')\n\n c = REPR_CONTEXT.current_color\n iterator = obj.__nnx_repr__()\n config = next(iterator)\n\n if not isinstance(config, Object):\n raise TypeError(f'First item must be Config, got {type(config).__name__}')\n\n kv_sep = f'{c.SEP}{config.kv_sep}{c.END}'\n\n def _repr_elem(elem: tp.Any) -> str:\n if not isinstance(elem, Attr):\n raise TypeError(f'Item must be Elem, got {type(elem).__name__}')\n\n value_repr = elem.value if elem.use_raw_value else colorized(elem.value)\n value_repr = value_repr.replace('\\n', '\\n' + config.indent)\n key = elem.key if elem.use_raw_key else f'{c.ATTRIBUTE}{elem.key}{c.END}'\n indent = '' if config.same_line else config.indent\n\n return f'{indent}{elem.start}{key}{kv_sep}{value_repr}{elem.end}'\n\n max_depth_reached = (\n flax_config.flax_max_repr_depth is not None\n and REPR_CONTEXT.depth > flax_config.flax_max_repr_depth\n )\n if max_depth_reached:\n elems = '...'\n else:\n elems = config.elem_sep.join(map(_repr_elem, iterator))\n\n if elems:\n if config.same_line or max_depth_reached:\n elems_repr = elems\n comment = ''\n else:\n elems_repr = '\\n' + elems + '\\n'\n comment = f'{c.COMMENT}{config.comment}{c.END}'\n else:\n elems_repr = config.empty_repr\n comment = ''\n\n type_repr = (\n config.type if isinstance(config.type, str) else config.type.__name__\n )\n type_repr = f'{c.TYPE}{type_repr}{c.END}' if type_repr else ''\n start = f'{c.PAREN}{config.start}{c.END}' if config.start else ''\n end = f'{c.PAREN}{config.end}{c.END}' if config.end else ''\n\n out = f'{type_repr}{start}{comment}{elems_repr}{end}'\n return out\n finally:\n REPR_CONTEXT.depth -= 1\n\nclass MappingReprMixin(Representable):\n def __nnx_repr__(self):\n yield Object(type=type(self), kv_sep=': ', start='({', end='})')\n\n for key, value in self.items(): # type: ignore\n yield Attr(colorized(key), value, use_raw_key=True)\n\n@dataclasses.dataclass(repr=False)\nclass PrettyMapping(Representable):\n mapping: tp.Mapping\n\n def __nnx_repr__(self):\n yield Object(type=type(self), kv_sep=': ', start='({', end='})')\n\n for key, value in self.mapping.items():\n yield Attr(colorized(key), value, use_raw_key=True)\n\n\n@dataclasses.dataclass(repr=False)\nclass SequenceReprMixin(Representable):\n def __nnx_repr__(self):\n yield Object(type=type(self), kv_sep='', start='([', end='])')\n\n for value in self: # type: ignore\n yield Attr('', value, use_raw_key=True)\n\n\n@dataclasses.dataclass(repr=False)\nclass PrettySequence(Representable):\n sequence: tp.Sequence\n\n def __nnx_repr__(self):\n yield Object(type=type(self), kv_sep='', start='([', end='])')\n\n for value in self.sequence:\n yield Attr('', value, use_raw_key=True)" }, { "path": "flax/nnx/rnglib.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nimport functools\nimport typing as tp\n\nimport jax\nfrom jax import random\nimport jax.numpy as jnp\n\nfrom flax import struct\nfrom flax import typing\nfrom flax.nnx import graphlib\nfrom flax.nnx.nn import initializers\nfrom flax.nnx.variablelib import Variable\nfrom flax.nnx import filterlib\nfrom flax.nnx.pytreelib import Pytree\nfrom flax.typing import MISSING, Key, Missing\nimport warnings\n\nF = tp.TypeVar('F', bound=tp.Callable[..., tp.Any])\nA = tp.TypeVar('A')\nCounts = list[int]\nAxesValue = tp.Union[int, None]\nSplitPattern = tp.Union[AxesValue, tuple[AxesValue, ...]]\nOutShardingType: tp.TypeAlias = (\n jax.sharding.NamedSharding | jax.sharding.PartitionSpec | None\n)\nFargs = tp.ParamSpec('Fargs')\n\n\n@tp.runtime_checkable\nclass KeylessInitializer(tp.Protocol):\n def __call__(\n self,\n shape: typing.Shape,\n dtype: tp.Any | None = None,\n out_sharding: OutShardingType = None,\n ) -> jax.Array:\n raise NotImplementedError\n\n\ndef _to_keyless(\n initializer_constructor: tp.Callable[Fargs, jax.nn.initializers.Initializer],\n) -> tp.Callable[Fargs, KeylessInitializer]:\n raise NotImplementedError\n\n\ndef _function_to_method(random_f):\n @functools.wraps(random_f)\n def rngs_random_method(self: Rngs | RngStream, *args, **kwargs) -> jax.Array:\n return random_f(self(), *args, **kwargs)\n\n return rngs_random_method\n\n\ndef _initializer_to_method(\n initializer_constructor: tp.Callable[Fargs, jax.nn.initializers.Initializer],\n):\n def rngs_initializer_method(\n self: Rngs | RngStream, *args: Fargs.args, **kwargs: Fargs.kwargs\n ) -> KeylessInitializer:\n init_fn = initializer_constructor(*args, **kwargs)\n\n def rngs_keyless_initializer(*init_args, **init_kwargs):\n return init_fn(self(), *init_args, **init_kwargs)\n\n return rngs_keyless_initializer\n\n return rngs_initializer_method\n\n\nclass RngState(Variable[jax.Array]):\n tag: str\n\n\nclass RngCount(RngState): ...\n\n\nclass RngKey(RngState): ...\n\n\nNotKey = filterlib.All(RngState, filterlib.Not(RngKey))\n\n\nclass RngStream(Pytree):\n\n def __init__(\n self,\n key: jax.Array | int,\n *,\n tag: str,\n ):\n if isinstance(key, int):\n key = random.key(key)\n elif isinstance(key, jax.Array) and key.dtype == jnp.uint32:\n key = random.wrap_key_data(key)\n\n if not isinstance(key, jax.Array) or not jnp.issubdtype(key.dtype, jax.dtypes.prng_key):\n raise ValueError(f'Invalid rng value: {key}, expected a '\n f'jax.Array of jax.dtypes.prng_key sub-dtype')\n\n count = jnp.zeros(key.shape, dtype=jnp.uint32)\n self.tag = tag\n self.key = RngKey(key, tag=tag)\n self.count = RngCount(count, tag=tag)\n\n def __call__(self) -> jax.Array:\n self.count._check_can_update()\n key = random.fold_in(self.key[...], self.count[...])\n self.count[...] += 1\n return key\n\n def split(self, k: int | tuple[int, ...]):\n key = random.split(self(), k)\n return type(self)(key, tag=self.tag)\n\n def fork(self, *, split: int | tuple[int, ...] | None = None):\n if split is not None:\n warnings.warn(\n \"The 'split' argument of 'fork' is deprecated; use the 'split' method instead.\",\n DeprecationWarning,\n stacklevel=2,\n )\n key = self()\n if split is not None:\n key = random.split(key, split)\n return type(self)(key, tag=self.tag)\n\n # ----------------------------------------------------------\n # random functions\n # ----------------------------------------------------------\n if tp.TYPE_CHECKING:\n bits = staticmethod(functools.partial(random.bits, random.key(0)))\n uniform = staticmethod(\n functools.partial(random.uniform, random.key(0))\n )\n randint = staticmethod(\n functools.partial(random.randint, random.key(0))\n )\n permutation = staticmethod(\n functools.partial(random.permutation, random.key(0))\n )\n choice = staticmethod(functools.partial(random.choice, random.key(0)))\n normal = staticmethod(functools.partial(random.normal, random.key(0)))\n multivariate_normal = staticmethod(\n functools.partial(random.multivariate_normal, random.key(0))\n )\n truncated_normal = staticmethod(\n functools.partial(random.truncated_normal, random.key(0))\n )\n bernoulli = staticmethod(\n functools.partial(random.bernoulli, random.key(0))\n )\n beta = staticmethod(functools.partial(random.beta, random.key(0)))\n cauchy = staticmethod(functools.partial(random.cauchy, random.key(0)))\n dirichlet = staticmethod(\n functools.partial(random.dirichlet, random.key(0))\n )\n exponential = staticmethod(\n functools.partial(random.exponential, random.key(0))\n )\n gamma = staticmethod(functools.partial(random.gamma, random.key(0)))\n loggamma = staticmethod(\n functools.partial(random.loggamma, random.key(0))\n )\n poisson = staticmethod(\n functools.partial(random.poisson, random.key(0))\n )\n gumbel = staticmethod(functools.partial(random.gumbel, random.key(0)))\n categorical = staticmethod(\n functools.partial(random.categorical, random.key(0))\n )\n laplace = staticmethod(\n functools.partial(random.laplace, random.key(0))\n )\n logistic = staticmethod(\n functools.partial(random.logistic, random.key(0))\n )\n pareto = staticmethod(functools.partial(random.pareto, random.key(0)))\n t = staticmethod(functools.partial(random.t, random.key(0)))\n chisquare = staticmethod(\n functools.partial(random.chisquare, random.key(0))\n )\n f = staticmethod(functools.partial(random.f, random.key(0)))\n rademacher = staticmethod(\n functools.partial(random.rademacher, random.key(0))\n )\n maxwell = staticmethod(\n functools.partial(random.maxwell, random.key(0))\n )\n double_sided_maxwell = staticmethod(\n functools.partial(random.double_sided_maxwell, random.key(0))\n )\n weibull_min = staticmethod(\n functools.partial(random.weibull_min, random.key(0))\n )\n orthogonal = staticmethod(\n functools.partial(random.orthogonal, random.key(0))\n )\n generalized_normal = staticmethod(\n functools.partial(random.generalized_normal, random.key(0))\n )\n ball = staticmethod(functools.partial(random.ball, random.key(0)))\n rayleigh = staticmethod(\n functools.partial(random.rayleigh, random.key(0))\n )\n wald = staticmethod(functools.partial(random.wald, random.key(0)))\n geometric = staticmethod(\n functools.partial(random.geometric, random.key(0))\n )\n triangular = staticmethod(\n functools.partial(random.triangular, random.key(0))\n )\n lognormal = staticmethod(\n functools.partial(random.lognormal, random.key(0))\n )\n binomial = staticmethod(\n functools.partial(random.binomial, random.key(0))\n )\n multinomial = staticmethod(\n functools.partial(random.multinomial, random.key(0))\n )\n else:\n bits = _function_to_method(random.bits)\n uniform = _function_to_method(random.uniform)\n randint = _function_to_method(random.randint)\n permutation = _function_to_method(random.permutation)\n choice = _function_to_method(random.choice)\n normal = _function_to_method(random.normal)\n multivariate_normal = _function_to_method(random.multivariate_normal)\n truncated_normal = _function_to_method(random.truncated_normal)\n bernoulli = _function_to_method(random.bernoulli)\n beta = _function_to_method(random.beta)\n cauchy = _function_to_method(random.cauchy)\n dirichlet = _function_to_method(random.dirichlet)\n exponential = _function_to_method(random.exponential)\n gamma = _function_to_method(random.gamma)\n loggamma = _function_to_method(random.loggamma)\n poisson = _function_to_method(random.poisson)\n gumbel = _function_to_method(random.gumbel)\n categorical = _function_to_method(random.categorical)\n laplace = _function_to_method(random.laplace)\n logistic = _function_to_method(random.logistic)\n pareto = _function_to_method(random.pareto)\n t = _function_to_method(random.t)\n chisquare = _function_to_method(random.chisquare)\n f = _function_to_method(random.f)\n rademacher = _function_to_method(random.rademacher)\n maxwell = _function_to_method(random.maxwell)\n double_sided_maxwell = _function_to_method(random.double_sided_maxwell)\n weibull_min = _function_to_method(random.weibull_min)\n orthogonal = _function_to_method(random.orthogonal)\n generalized_normal = _function_to_method(random.generalized_normal)\n ball = _function_to_method(random.ball)\n rayleigh = _function_to_method(random.rayleigh)\n wald = _function_to_method(random.wald)\n geometric = _function_to_method(random.geometric)\n triangular = _function_to_method(random.triangular)\n lognormal = _function_to_method(random.lognormal)\n binomial = _function_to_method(random.binomial)\n multinomial = _function_to_method(random.multinomial)\n\n # ----------------------------------------------------------\n # initializers\n # ----------------------------------------------------------\n if tp.TYPE_CHECKING:\n # skip constant\n delta_orthogonal = staticmethod(_to_keyless(initializers.delta_orthogonal))\n glorot_normal = staticmethod(_to_keyless(initializers.glorot_normal))\n glorot_uniform = staticmethod(_to_keyless(initializers.glorot_uniform))\n he_normal = staticmethod(_to_keyless(initializers.he_normal))\n he_uniform = staticmethod(_to_keyless(initializers.he_uniform))\n kaiming_normal = staticmethod(_to_keyless(initializers.kaiming_normal))\n kaiming_uniform = staticmethod(_to_keyless(initializers.kaiming_uniform))\n lecun_normal = staticmethod(_to_keyless(initializers.lecun_normal))\n lecun_uniform = staticmethod(_to_keyless(initializers.lecun_uniform))\n # skip normal as it conflicts with jax.random.normal\n # skip ones\n # skip orthogonal as it conflicts with jax.random.orthogonal\n # skip truncated_normal as it conflicts with jax.random.truncated_normal\n # skip uniform as it conflicts with jax.random.uniform\n variance_scaling = staticmethod(_to_keyless(initializers.variance_scaling))\n xavier_normal = staticmethod(_to_keyless(initializers.xavier_normal))\n xavier_uniform = staticmethod(_to_keyless(initializers.xavier_uniform))\n # skip zeros\n else:\n # skip constant\n delta_orthogonal = _initializer_to_method(initializers.delta_orthogonal)\n glorot_normal = _initializer_to_method(initializers.glorot_normal)\n glorot_uniform = _initializer_to_method(initializers.glorot_uniform)\n he_normal = _initializer_to_method(initializers.he_normal)\n he_uniform = _initializer_to_method(initializers.he_uniform)\n kaiming_normal = _initializer_to_method(initializers.kaiming_normal)\n kaiming_uniform = _initializer_to_method(initializers.kaiming_uniform)\n lecun_normal = _initializer_to_method(initializers.lecun_normal)\n lecun_uniform = _initializer_to_method(initializers.lecun_uniform)\n # skip normal as it conflicts with jax.random.normal\n # skip ones\n # skip orthogonal as it conflicts with jax.random.orthogonal\n # skip truncated_normal as it conflicts with jax.random.truncated_normal\n # skip uniform as it conflicts with jax.random.uniform\n variance_scaling = _initializer_to_method(initializers.variance_scaling)\n xavier_normal = _initializer_to_method(initializers.xavier_normal)\n xavier_uniform = _initializer_to_method(initializers.xavier_uniform)\n # skip zeros\n\n\nRngValue = tp.Union[int, jax.Array]\n\nclass Rngs(Pytree):\n \"\"\"A small abstraction to manage RNG state.\n\n ``Rngs`` allows the creation of ``RngStream`` which are used to easily generate new unique\n random keys on demand. An ``RngStream`` is a wrapper around a JAX random ``key``, and a\n ``counter``. Every time a key is requested, the counter is incremented and the key is\n generated from the seed key and the counter by using ``jax.random.fold_in``.\n\n To create an ``Rngs`` pass in an integer or ``jax.random.key`` to the\n constructor as a keyword argument with the name of the stream. The key will be used as the\n starting seed for the stream, and the counter will be initialized to zero. Then call the\n stream to get a key::\n\n >>> from flax import nnx\n >>> import jax, jax.numpy as jnp\n\n >>> rngs = nnx.Rngs(params=0, dropout=1)\n\n >>> param_key1 = rngs.params()\n >>> param_key2 = rngs.params()\n >>> dropout_key1 = rngs.dropout()\n >>> dropout_key2 = rngs.dropout()\n ...\n >>> assert param_key1 != dropout_key1\n\n Trying to generate a key for a stream that was not specified during construction\n will result in an error being raised::\n\n >>> rngs = nnx.Rngs(params=0, dropout=1)\n >>> try:\n ... key = rngs.unkown_stream()\n ... except AttributeError as e:\n ... print(e)\n No RngStream named 'unkown_stream' found in Rngs.\n\n The ``default`` stream can be created by passing in a key to the constructor without\n specifying a stream name. When the ``default`` stream is set the ``rngs`` object can be\n called directly to get a key, and calling streams that were not specified during\n construction will fallback to ``default``::\n\n >>> rngs = nnx.Rngs(0, params=1)\n ...\n >>> key1 = rngs.default() # uses 'default'\n >>> key2 = rngs() # uses 'default'\n >>> key3 = rngs.params() # uses 'params'\n >>> key4 = rngs.dropout() # uses 'default'\n >>> key5 = rngs.unkown_stream() # uses 'default'\n \"\"\"\n\n def __init__(\n self,\n default: RngValue\n | RngStream\n | tp.Mapping[str, RngValue | RngStream]\n | None = None,\n **rngs: RngValue | RngStream,\n ):\n \"\"\"\n Args:\n default: the starting seed for the ``default`` stream, defaults to None.\n **rngs: keyword arguments specifying the starting seed for each stream.\n The key can be an integer or a ``jax.random.key``.\n \"\"\"\n if default is not None:\n if isinstance(default, tp.Mapping):\n rngs = {**default, **rngs}\n else:\n rngs['default'] = default\n\n for tag, key in rngs.items():\n if isinstance(key, RngStream):\n key = key.key.get_value()\n stream = RngStream(\n key=key,\n tag=tag,\n )\n setattr(self, tag, stream)\n\n def _get_stream(self, name: str, error_type: type[Exception]) -> RngStream:\n stream_vars = vars(self)\n if name not in stream_vars:\n if 'default' not in stream_vars:\n raise error_type(f\"No RngStream named '{name}' found in Rngs.\")\n stream = stream_vars['default']\n else:\n stream = stream_vars[name]\n\n return stream\n\n def __getitem__(self, name: str):\n return self._get_stream(name, KeyError)\n\n def __getattr__(self, name: str):\n return self._get_stream(name, AttributeError)\n\n def __call__(self):\n return self.default()\n\n def __iter__(self) -> tp.Iterator[str]:\n for name, stream in vars(self).items():\n if isinstance(stream, RngStream):\n yield name\n\n def __len__(self) -> int:\n return sum(\n 1 for stream in vars(self).values() if isinstance(stream, RngStream)\n )\n\n def __contains__(self, name: tp.Any) -> bool:\n return name in vars(self)\n\n def items(self):\n for name, stream in vars(self).items():\n if isinstance(stream, RngStream):\n yield name, stream\n\n def split(self, k: tp.Mapping[filterlib.Filter, int | tuple[int, ...]] | int | tuple[int, ...]):\n \"\"\"\n Splits the keys of the newly created ``Rngs`` object.\n\n Example::\n >>> from flax import nnx\n ...\n >>> rngs = nnx.Rngs(params=1, dropout=2)\n >>> new_rngs = rngs.split(5)\n ...\n >>> assert new_rngs.params.key.shape == (5,)\n >>> assert new_rngs.dropout.key.shape == (5,)\n\n ``split`` also accepts a mapping of\n `Filters `__ to\n split sizes or None to control which streams are split and how they are split::\n\n >>> rngs = nnx.Rngs(params=1, dropout=2, noise=3)\n >>> new_rngs = rngs.split({\n ... 'params': 5, # split params into 5 keys\n ... 'dropout': None, # don't split dropout\n ... ...: (2, 5), # split anything else into 2x5 keys\n ... })\n ...\n >>> assert new_rngs.params.key.shape == (5,)\n >>> assert new_rngs.dropout.key.shape == ()\n >>> assert new_rngs.noise.key.shape == (2, 5)\n \"\"\"\n if isinstance(k, int):\n k = {...: k}\n elif isinstance(k, tuple):\n k = {...: k}\n\n split_predicates = {filterlib.to_predicate(k): v for k, v in k.items()}\n keys: dict[str, RngStream] = {}\n for name, stream in self.items():\n for predicate, num_splits in split_predicates.items():\n if predicate((), stream):\n if num_splits is None:\n keys[name] = stream\n else:\n keys[name] = stream.split(num_splits)\n break\n else:\n keys[name] = stream\n\n return Rngs(**keys)\n\n def fork(\n self,\n /,\n *,\n split: tp.Mapping[filterlib.Filter, int | tuple[int, ...]]\n | int\n | tuple[int, ...]\n | None = None,\n ):\n \"\"\"Returns a new Rngs object with new unique RNG keys.\n\n Example::\n >>> from flax import nnx\n ...\n >>> rngs = nnx.Rngs(params=1, dropout=2)\n >>> new_rngs = rngs.fork()\n ...\n >>> assert rngs.params() != new_rngs.params()\n \"\"\"\n if split is not None:\n warnings.warn(\n \"The 'split' argument of 'fork' is deprecated; use the 'split' method instead.\",\n DeprecationWarning,\n stacklevel=2,\n )\n\n if split is None:\n split = {}\n elif isinstance(split, int):\n split = {...: split}\n elif isinstance(split, tuple):\n split = {...: split}\n\n split_predicates = {filterlib.to_predicate(k): v for k, v in split.items()}\n keys: dict[str, RngStream] = {}\n for name, stream in self.items():\n for predicate, num_splits in split_predicates.items():\n if predicate((), stream):\n keys[name] = stream.fork(split=num_splits)\n break\n else:\n keys[name] = stream.fork()\n\n return Rngs(**keys)\n\n # ----------------------------------------------------------\n # random functions\n # ----------------------------------------------------------\n if tp.TYPE_CHECKING:\n bits = staticmethod(functools.partial(random.bits, random.key(0)))\n uniform = staticmethod(\n functools.partial(random.uniform, random.key(0))\n )\n randint = staticmethod(\n functools.partial(random.randint, random.key(0))\n )\n permutation = staticmethod(\n functools.partial(random.permutation, random.key(0))\n )\n choice = staticmethod(functools.partial(random.choice, random.key(0)))\n normal = staticmethod(functools.partial(random.normal, random.key(0)))\n multivariate_normal = staticmethod(\n functools.partial(random.multivariate_normal, random.key(0))\n )\n truncated_normal = staticmethod(\n functools.partial(random.truncated_normal, random.key(0))\n )\n bernoulli = staticmethod(\n functools.partial(random.bernoulli, random.key(0))\n )\n beta = staticmethod(functools.partial(random.beta, random.key(0)))\n cauchy = staticmethod(functools.partial(random.cauchy, random.key(0)))\n dirichlet = staticmethod(\n functools.partial(random.dirichlet, random.key(0))\n )\n exponential = staticmethod(\n functools.partial(random.exponential, random.key(0))\n )\n gamma = staticmethod(functools.partial(random.gamma, random.key(0)))\n loggamma = staticmethod(\n functools.partial(random.loggamma, random.key(0))\n )\n poisson = staticmethod(\n functools.partial(random.poisson, random.key(0))\n )\n gumbel = staticmethod(functools.partial(random.gumbel, random.key(0)))\n categorical = staticmethod(\n functools.partial(random.categorical, random.key(0))\n )\n laplace = staticmethod(\n functools.partial(random.laplace, random.key(0))\n )\n logistic = staticmethod(\n functools.partial(random.logistic, random.key(0))\n )\n pareto = staticmethod(functools.partial(random.pareto, random.key(0)))\n t = staticmethod(functools.partial(random.t, random.key(0)))\n chisquare = staticmethod(\n functools.partial(random.chisquare, random.key(0))\n )\n f = staticmethod(functools.partial(random.f, random.key(0)))\n rademacher = staticmethod(\n functools.partial(random.rademacher, random.key(0))\n )\n maxwell = staticmethod(\n functools.partial(random.maxwell, random.key(0))\n )\n double_sided_maxwell = staticmethod(\n functools.partial(random.double_sided_maxwell, random.key(0))\n )\n weibull_min = staticmethod(\n functools.partial(random.weibull_min, random.key(0))\n )\n orthogonal = staticmethod(\n functools.partial(random.orthogonal, random.key(0))\n )\n generalized_normal = staticmethod(\n functools.partial(random.generalized_normal, random.key(0))\n )\n ball = staticmethod(functools.partial(random.ball, random.key(0)))\n rayleigh = staticmethod(\n functools.partial(random.rayleigh, random.key(0))\n )\n wald = staticmethod(functools.partial(random.wald, random.key(0)))\n geometric = staticmethod(\n functools.partial(random.geometric, random.key(0))\n )\n triangular = staticmethod(\n functools.partial(random.triangular, random.key(0))\n )\n lognormal = staticmethod(\n functools.partial(random.lognormal, random.key(0))\n )\n binomial = staticmethod(\n functools.partial(random.binomial, random.key(0))\n )\n multinomial = staticmethod(\n functools.partial(random.multinomial, random.key(0))\n )\n else:\n bits = _function_to_method(random.bits)\n uniform = _function_to_method(random.uniform)\n randint = _function_to_method(random.randint)\n permutation = _function_to_method(random.permutation)\n choice = _function_to_method(random.choice)\n normal = _function_to_method(random.normal)\n multivariate_normal = _function_to_method(random.multivariate_normal)\n truncated_normal = _function_to_method(random.truncated_normal)\n bernoulli = _function_to_method(random.bernoulli)\n beta = _function_to_method(random.beta)\n cauchy = _function_to_method(random.cauchy)\n dirichlet = _function_to_method(random.dirichlet)\n exponential = _function_to_method(random.exponential)\n gamma = _function_to_method(random.gamma)\n loggamma = _function_to_method(random.loggamma)\n poisson = _function_to_method(random.poisson)\n gumbel = _function_to_method(random.gumbel)\n categorical = _function_to_method(random.categorical)\n laplace = _function_to_method(random.laplace)\n logistic = _function_to_method(random.logistic)\n pareto = _function_to_method(random.pareto)\n t = _function_to_method(random.t)\n chisquare = _function_to_method(random.chisquare)\n f = _function_to_method(random.f)\n rademacher = _function_to_method(random.rademacher)\n maxwell = _function_to_method(random.maxwell)\n double_sided_maxwell = _function_to_method(random.double_sided_maxwell)\n weibull_min = _function_to_method(random.weibull_min)\n orthogonal = _function_to_method(random.orthogonal)\n generalized_normal = _function_to_method(random.generalized_normal)\n ball = _function_to_method(random.ball)\n rayleigh = _function_to_method(random.rayleigh)\n wald = _function_to_method(random.wald)\n geometric = _function_to_method(random.geometric)\n triangular = _function_to_method(random.triangular)\n lognormal = _function_to_method(random.lognormal)\n binomial = _function_to_method(random.binomial)\n multinomial = _function_to_method(random.multinomial)\n\n # ----------------------------------------------------------\n # initializers\n # ----------------------------------------------------------\n if tp.TYPE_CHECKING:\n # skip constant\n delta_orthogonal = staticmethod(_to_keyless(initializers.delta_orthogonal))\n glorot_normal = staticmethod(_to_keyless(initializers.glorot_normal))\n glorot_uniform = staticmethod(_to_keyless(initializers.glorot_uniform))\n he_normal = staticmethod(_to_keyless(initializers.he_normal))\n he_uniform = staticmethod(_to_keyless(initializers.he_uniform))\n kaiming_normal = staticmethod(_to_keyless(initializers.kaiming_normal))\n kaiming_uniform = staticmethod(_to_keyless(initializers.kaiming_uniform))\n lecun_normal = staticmethod(_to_keyless(initializers.lecun_normal))\n lecun_uniform = staticmethod(_to_keyless(initializers.lecun_uniform))\n # skip normal as it conflicts with jax.random.normal\n # skip ones\n # skip orthogonal as it conflicts with jax.random.orthogonal\n # skip truncated_normal as it conflicts with jax.random.truncated_normal\n # skip uniform as it conflicts with jax.random.uniform\n variance_scaling = staticmethod(_to_keyless(initializers.variance_scaling))\n xavier_normal = staticmethod(_to_keyless(initializers.xavier_normal))\n xavier_uniform = staticmethod(_to_keyless(initializers.xavier_uniform))\n # skip zeros\n else:\n # skip constant\n delta_orthogonal = _initializer_to_method(initializers.delta_orthogonal)\n glorot_normal = _initializer_to_method(initializers.glorot_normal)\n glorot_uniform = _initializer_to_method(initializers.glorot_uniform)\n he_normal = _initializer_to_method(initializers.he_normal)\n he_uniform = _initializer_to_method(initializers.he_uniform)\n kaiming_normal = _initializer_to_method(initializers.kaiming_normal)\n kaiming_uniform = _initializer_to_method(initializers.kaiming_uniform)\n lecun_normal = _initializer_to_method(initializers.lecun_normal)\n lecun_uniform = _initializer_to_method(initializers.lecun_uniform)\n # skip normal as it conflicts with jax.random.normal\n # skip ones\n # skip orthogonal as it conflicts with jax.random.orthogonal\n # skip truncated_normal as it conflicts with jax.random.truncated_normal\n # skip uniform as it conflicts with jax.random.uniform\n variance_scaling = _initializer_to_method(initializers.variance_scaling)\n xavier_normal = _initializer_to_method(initializers.xavier_normal)\n xavier_uniform = _initializer_to_method(initializers.xavier_uniform)\n # skip zeros\n\n\nStreamBackup = (\n tuple[RngStream, jax.Array, jax.Array] | tuple[RngStream, jax.Array]\n)\n\n\nclass SplitBackups(struct.PyTreeNode, tp.Iterable[StreamBackup]):\n backups: list[StreamBackup]\n\n def __iter__(self) -> tp.Iterator[StreamBackup]:\n return iter(self.backups)\n\n def __enter__(self):\n return self\n\n def __exit__(self, *args):\n restore_rngs(self)\n\n\n@tp.overload\ndef split_rngs(\n node: tp.Any,\n /,\n *,\n splits: int | tuple[int, ...],\n only: filterlib.Filter = ...,\n squeeze: bool = False,\n graph: tp.Literal[True] | None = None,\n) -> SplitBackups: ...\n@tp.overload\ndef split_rngs(\n node: A,\n /,\n *,\n splits: int | tuple[int, ...],\n only: filterlib.Filter = ...,\n squeeze: bool = False,\n graph: tp.Literal[False],\n) -> A: ...\n@tp.overload\ndef split_rngs(\n *,\n splits: int | tuple[int, ...],\n only: filterlib.Filter = ...,\n squeeze: bool = False,\n graph: bool | None = None,\n) -> tp.Callable[[F], F]: ...\ndef split_rngs(\n node: tp.Any = MISSING,\n /,\n *,\n splits: int | tuple[int, ...],\n only: filterlib.Filter = ...,\n squeeze: bool = False,\n graph: bool | None = None,\n) -> SplitBackups | tp.Any | tp.Callable[[F], F]:\n \"\"\"Splits the (nested) Rng states of the given node.\n\n Args:\n node: the base node containing the rng states to split.\n splits: an integer or tuple of integers specifying the\n shape of the split rng keys.\n only: a Filter selecting which rng states to split.\n graph: If ``True`` (default), uses graph-mode which supports the full\n NNX feature set including shared references. If ``False``, uses\n tree-mode which treats Modules as regular JAX pytrees, avoiding\n the overhead of the graph protocol.\n\n Returns:\n A SplitBackups iterable if ``node`` is provided, otherwise a\n decorator that splits the rng states of the inputs to the\n decorated function.\n\n Example::\n\n >>> from flax import nnx\n ...\n >>> rngs = nnx.Rngs(params=0, dropout=1)\n >>> _ = nnx.split_rngs(rngs, splits=5)\n >>> rngs.params.key.shape, rngs.dropout.key.shape\n ((5,), (5,))\n\n >>> rngs = nnx.Rngs(params=0, dropout=1)\n >>> _ = nnx.split_rngs(rngs, splits=(2, 5))\n >>> rngs.params.key.shape, rngs.dropout.key.shape\n ((2, 5), (2, 5))\n\n\n >>> rngs = nnx.Rngs(params=0, dropout=1)\n >>> _ = nnx.split_rngs(rngs, splits=5, only='params')\n >>> rngs.params.key.shape, rngs.dropout.key.shape\n ((5,), ())\n\n Once split, random state can be used with transforms like :func:`nnx.vmap`::\n\n >>> class Model(nnx.Module):\n ... def __init__(self, rngs):\n ... self.linear = nnx.Linear(2, 3, rngs=rngs)\n ... self.dropout = nnx.Dropout(0.5, rngs=rngs)\n ...\n >>> rngs = nnx.Rngs(params=0, dropout=1)\n >>> _ = nnx.split_rngs(rngs, splits=5, only='params')\n ...\n >>> state_axes = nnx.StateAxes({(nnx.Param, 'params'): 0, ...: None})\n ...\n >>> @nnx.vmap(in_axes=(state_axes,), out_axes=state_axes)\n ... def create_model(rngs):\n ... return Model(rngs)\n ...\n >>> model = create_model(rngs)\n >>> model.dropout.rngs.key.shape\n ()\n\n ``split_rngs`` returns a SplitBackups object that can be used to restore the\n original unsplit rng states using :func:`nnx.restore_rngs`, this is useful\n when you only want to split the rng states temporarily::\n\n >>> rngs = nnx.Rngs(params=0, dropout=1)\n ...\n >>> backups = nnx.split_rngs(rngs, splits=5, only='params')\n >>> model = create_model(rngs)\n >>> nnx.restore_rngs(backups)\n ...\n >>> model.dropout.rngs.key.shape\n ()\n\n SplitBackups can also be used as a context manager to automatically restore\n the rng states when exiting the context::\n\n >>> rngs = nnx.Rngs(params=0, dropout=1)\n ...\n >>> with nnx.split_rngs(rngs, splits=5, only='params'):\n ... model = create_model(rngs)\n ...\n >>> model.dropout.rngs.key.shape\n ()\n\n >>> state_axes = nnx.StateAxes({(nnx.Param, 'params'): 0, ...: None})\n ...\n >>> @nnx.split_rngs(splits=5, only='params')\n ... @nnx.vmap(in_axes=(state_axes,), out_axes=state_axes)\n ... def create_model(rngs):\n ... return Model(rngs)\n ...\n >>> rngs = nnx.Rngs(params=0, dropout=1)\n >>> model = create_model(rngs)\n >>> model.dropout.rngs.key.shape\n ()\n\n\n \"\"\"\n if graph is None:\n graph = graphlib.set_graph_mode.current_value()\n\n if isinstance(node, Missing):\n\n def split_rngs_decorator(f: F) -> F:\n @functools.wraps(f)\n def split_rngs_wrapper(*args, **kwargs):\n if graph:\n with split_rngs(\n (args, kwargs), splits=splits, only=only, squeeze=squeeze,\n graph=True,\n ):\n return f(*args, **kwargs)\n else:\n args, kwargs = split_rngs(\n (args, kwargs), splits=splits, only=only, squeeze=squeeze,\n graph=False,\n )\n return f(*args, **kwargs)\n\n return tp.cast(F, split_rngs_wrapper)\n\n return split_rngs_decorator # type: ignore[bad-return-type]\n\n if squeeze and splits != 1:\n raise ValueError('squeeze=True is only supported for splits=1')\n\n if graph:\n return _graph_split_rngs(\n node, splits=splits, only=only, squeeze=squeeze,\n )\n else:\n return _tree_split_rngs(\n node, splits=splits, only=only, squeeze=squeeze,\n )\n\n\ndef _graph_split_rngs(\n node: tp.Any,\n /,\n *,\n splits: int | tuple[int, ...],\n only: filterlib.Filter = ...,\n squeeze: bool = False,\n) -> SplitBackups:\n predicate = filterlib.to_predicate(only)\n backups: list[StreamBackup] = []\n for path, stream in graphlib.iter_graph(node, graph=True):\n if (\n isinstance(stream, RngStream)\n and predicate((*path, 'key'), stream.key)\n and predicate((*path, 'count'), stream.count)\n ):\n key = stream()\n backups.append((stream, stream.key[...], stream.count[...]))\n key = random.split(key, splits)\n if squeeze:\n key = key[0]\n stream.key.set_value(key)\n if squeeze:\n counts_shape = stream.count.shape\n elif isinstance(splits, int):\n counts_shape = (splits, *stream.count.shape)\n else:\n counts_shape = (*splits, *stream.count.shape)\n\n stream.count.set_value(jnp.zeros(counts_shape, dtype=jnp.uint32))\n\n return SplitBackups(backups)\n\n\ndef _tree_split_rngs(\n node: tp.Any,\n /,\n *,\n splits: int | tuple[int, ...],\n only: filterlib.Filter = ...,\n squeeze: bool = False,\n) -> tp.Any:\n predicate = filterlib.to_predicate(only)\n\n def _split_stream(path, node):\n if (\n isinstance(node, RngStream)\n and predicate((*path, 'key'), node.key)\n and predicate((*path, 'count'), node.count)\n ):\n key = random.split(node(), splits)\n if squeeze:\n key = key[0]\n if squeeze:\n counts_shape = node.count.shape\n elif isinstance(splits, int):\n counts_shape = (splits, *node.count.shape)\n else:\n counts_shape = (*splits, *node.count.shape)\n\n node.key = RngKey(key, tag=node.tag)\n node.count = RngCount(\n jnp.zeros(counts_shape, dtype=jnp.uint32), tag=node.tag\n )\n return node\n\n return graphlib.recursive_map(_split_stream, node, graph=False)\n\n@tp.overload\ndef fork_rngs(\n node: tp.Any,\n /,\n *,\n split: tp.Mapping[filterlib.Filter, int | tuple[int, ...] | None]\n | int\n | None = None,\n graph: bool | None = None,\n) -> SplitBackups: ...\n@tp.overload\ndef fork_rngs(\n *,\n split: tp.Mapping[filterlib.Filter, int | tuple[int, ...] | None]\n | int\n | None = None,\n graph: bool | None = None,\n) -> tp.Callable[[F], F]: ...\ndef fork_rngs(\n node: tp.Any = MISSING,\n /,\n *,\n split: tp.Mapping[filterlib.Filter, int | tuple[int, ...] | None]\n | int\n | None = None,\n graph: bool | None = None,\n) -> SplitBackups | tp.Callable[[F], F]:\n \"\"\"Forks the (nested) Rng states of the given node.\n\n Args:\n node: the base node containing the rng states to fork.\n graph: If ``True`` (default), uses graph-mode which supports the full\n NNX feature set including shared references. If ``False``, uses\n tree-mode which treats Modules as regular JAX pytrees, avoiding\n the overhead of the graph protocol.\n\n Returns:\n A SplitBackups iterable if ``node`` is provided, otherwise a\n decorator that forks the rng states of the inputs to the\n decorated function.\n\n Example::\n\n >>> from flax import nnx\n ...\n >>> rngs = nnx.Rngs(params=0, dropout=1)\n >>> _ = nnx.fork_rngs(rngs)\n\n ``fork_rngs`` returns a SplitBackups object that can be used to restore the\n original unforked rng states using :func:`nnx.restore_rngs`, this is useful\n when you only want to fork the rng states temporarily::\n\n >>> rngs = nnx.Rngs(params=0, dropout=1)\n ...\n >>> backups = nnx.fork_rngs(rngs)\n >>> model = nnx.Linear(2, 3, rngs=rngs)\n >>> nnx.restore_rngs(backups)\n ...\n\n SplitBackups can also be used as a context manager to automatically restore\n the rng states when exiting the context::\n\n >>> rngs = nnx.Rngs(params=0, dropout=1)\n ...\n >>> with nnx.fork_rngs(rngs):\n ... model = nnx.Linear(2, 3, rngs=rngs)\n\n \"\"\"\n if isinstance(node, Missing):\n\n def fork_rngs_decorator(f: F) -> F:\n @functools.wraps(f)\n def fork_rngs_wrapper(*args, **kwargs):\n with fork_rngs((args, kwargs), split=split):\n return f(*args, **kwargs)\n\n return tp.cast(F, fork_rngs_wrapper)\n\n return fork_rngs_decorator # type: ignore[bad-return-type]\n\n if split is None:\n split = {...: None}\n elif isinstance(split, int | tuple):\n split = {...: split}\n\n predicate_splits = {\n filterlib.to_predicate(k): v for k, v in split.items()\n }\n backups: list[StreamBackup] = []\n for path, stream in graphlib.iter_graph(node, graph=graph):\n for predicate, splits in predicate_splits.items():\n if (\n isinstance(stream, RngStream)\n and predicate((*path, 'key'), stream.key)\n and predicate((*path, 'count'), stream.count)\n ):\n forked_stream = stream.fork(split=splits)\n # backup the original stream state\n backups.append((stream, stream.key[...], stream.count[...]))\n # apply the forked key and count to the original stream\n stream.key.set_value(forked_stream.key.get_value())\n stream.count.set_value(forked_stream.count.get_value())\n\n return SplitBackups(backups)\n\n\ndef backup_keys(node: tp.Any, /, *, graph: bool | None = None):\n backups: list[StreamBackup] = []\n for _, stream in graphlib.iter_graph(node, graph=graph):\n if isinstance(stream, RngStream):\n backups.append((stream, stream.key[...]))\n return backups\n\ndef _scalars_only(\n path: tuple[Key, ...], scalar_key: jax.Array, target_shape: tuple[int, ...]\n) -> jax.Array:\n if target_shape != ():\n raise ValueError(\n f'Cannot reseed stream at path {path!r} becuase it has a non-scalar key, '\n f'found key with shape {target_shape}. If all your multi-dimensional '\n 'keys have unique values on all dimensions, set policy=\"match_shape\", '\n 'else provide a custom reseed policy.'\n )\n return scalar_key\n\n\ndef _match_shape(\n path: tuple[Key, ...], scalar_key: jax.Array, target_shape: tuple[int, ...]\n) -> jax.Array:\n if target_shape == ():\n return scalar_key\n return random.split(scalar_key, target_shape)\n\n\ndef reseed(\n node,\n /,\n *,\n graph: bool | None = None,\n policy: tp.Literal['scalars_only', 'match_shape']\n | tp.Callable[\n [tuple, jax.Array, tuple[int, ...]], jax.Array\n ] = 'scalars_only',\n **stream_keys: RngValue,\n):\n \"\"\"Update the keys of the specified RNG streams with new keys.\n\n Args:\n node: the node to reseed the RNG streams in.\n graph: If ``True`` (default), uses graph-mode which supports the full\n NNX feature set including shared references. If ``False``, uses\n tree-mode which treats Modules as regular JAX pytrees, avoiding\n the overhead of the graph protocol.\n policy: defines how the the new scalar key is for each RngStream is used to\n reseed the stream. If ``'scalars_only'`` is given (the default), an error is raised\n if the target stream key is not a scalar. If ``'match_shape'`` is given, the new\n scalar key is split to match the shape of the target stream key. A callable\n of the form ``(path, scalar_key, target_shape) -> new_key`` can be passed to\n define a custom reseeding policy.\n **stream_keys: a mapping of stream names to new keys. The keys can be\n either integers or ``jax.random.key``.\n\n Example::\n\n >>> from flax import nnx\n >>> import jax.numpy as jnp\n ...\n >>> class Model(nnx.Module):\n ... def __init__(self, rngs):\n ... self.linear = nnx.Linear(2, 3, rngs=rngs)\n ... self.dropout = nnx.Dropout(0.5, rngs=rngs)\n ... def __call__(self, x):\n ... return self.dropout(self.linear(x))\n ...\n >>> model = Model(nnx.Rngs(params=0, dropout=42))\n >>> x = jnp.ones((1, 2))\n ...\n >>> y1 = model(x)\n ...\n >>> # reset the ``dropout`` stream key to 42\n >>> nnx.reseed(model, dropout=42)\n >>> y2 = model(x)\n ...\n >>> jnp.allclose(y1, y2)\n Array(True, dtype=bool)\n \"\"\"\n if policy == 'scalars_only':\n policy = _scalars_only\n elif policy == 'match_shape':\n policy = _match_shape\n elif not callable(policy):\n raise ValueError(\n f'policy must be \"scalars_only\", \"match_shape\" or a callable, '\n f'got {policy!r}'\n )\n rngs = Rngs(**stream_keys)\n for path, stream in graphlib.iter_graph(node, graph=graph):\n if isinstance(stream, RngStream):\n if stream.key.tag in stream_keys:\n key = rngs[stream.key.tag]()\n key = policy(path, key, stream.key.shape)\n stream.key.set_value(key)\n stream.count.set_value(jnp.zeros(key.shape, dtype=jnp.uint32))\n\n\ndef restore_rngs(backups: tp.Iterable[StreamBackup], /):\n for backup in backups:\n stream = backup[0]\n stream.key.set_value(backup[1])\n if len(backup) == 3:\n stream.count.set_value(backup[2]) # count\n" }, { "path": "flax/nnx/scripts/requirements.txt", "content": "datasets>=2.12.0\n" }, { "path": "flax/nnx/scripts/run-all-examples.bash", "content": "set -e\n\nsource .venv/bin/activate\n\nfor f in $(find examples/nnx_toy_examples -name \"*.py\" -maxdepth 1); do\n echo -e \"\\n---------------------------------\"\n echo \"$f\"\n echo \"---------------------------------\"\n MPLBACKEND=Agg python \"$f\"\ndone\n" }, { "path": "flax/nnx/spmd.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport typing as tp\n\nimport flax.core.spmd as core_spmd\nfrom flax.nnx import variablelib, graphlib\nfrom flax.nnx.transforms.transforms import eval_shape\nfrom flax.typing import (\n Sharding,\n)\nimport jax\nfrom jax.sharding import PartitionSpec\n\nA = tp.TypeVar('A')\nF = tp.TypeVar('F', bound=tp.Callable[..., tp.Any])\nPARTITION_NAME = 'partition_name'\n\n\n# Transform axis change helpers\n# ------------------------------------------------------------------------------\n\n\ndef add_axis(tree: A, index: int, transform_metadata: tp.Mapping) -> A:\n axis_name, other_meta = _get_partition_name_and_metadata(transform_metadata)\n\n def insert_field(fields, index, value):\n iterable = list(fields)\n while len(iterable) < index:\n iterable.append(None)\n iterable.insert(index, value)\n return tuple(iterable)\n\n def _add_axis(x: tp.Any):\n if isinstance(x, variablelib.Variable):\n metadata = x.get_metadata()\n if 'out_sharding' in metadata and metadata['out_sharding']:\n sharding = metadata['out_sharding']\n x.set_metadata(out_sharding=insert_field(sharding, index, axis_name))\n\n for k, v in other_meta.items():\n if hasattr(x, k) and (t := getattr(x, k)) and isinstance(t, tuple):\n x.set_metadata(k, insert_field(t, index, v))\n\n assert isinstance(x, variablelib.Variable)\n x.add_axis(index, axis_name)\n return x\n\n return jax.tree.map(\n _add_axis, tree, is_leaf=lambda x: isinstance(x, variablelib.Variable)\n )\n\n\ndef remove_axis(\n tree: A, index: int, transform_metadata: tp.Mapping[tp.Any, tp.Any]\n) -> A:\n axis_name, other_meta = _get_partition_name_and_metadata(transform_metadata)\n\n def remove_field(fields, index, value):\n iterable = list(fields)\n removed = iterable.pop(index)\n if removed != value:\n raise ValueError(\n f'Expected to remove {value!r} at index {index} from '\n f'{fields!r}, but found {removed!r}.'\n )\n return tuple(iterable)\n\n def _remove_axis(x: tp.Any):\n if isinstance(x, variablelib.Variable):\n if hasattr(x, 'out_sharding') and x.out_sharding is not None:\n x.set_metadata(\n out_sharding=remove_field(x.out_sharding, index, axis_name)\n )\n\n for k, v in other_meta.items():\n if hasattr(x, k) and (t := getattr(x, k)) and isinstance(t, tuple):\n x.set_metadata(k, remove_field(t, index, v))\n\n x.remove_axis(index, axis_name)\n return x\n\n return jax.tree.map(\n _remove_axis,\n tree,\n is_leaf=lambda x: isinstance(x, variablelib.Variable),\n )\n\n\ndef _get_partition_name_and_metadata(\n transform_metadata: tp.Mapping[tp.Any, tp.Any],\n) -> tuple[str, tp.Mapping[tp.Any, tp.Any]]:\n if PARTITION_NAME not in transform_metadata:\n raise ValueError(\n 'Trying to transform a Partitioned variable but \"partition_name\" '\n f'is not specified in transform_metadata: {transform_metadata}'\n )\n other_meta = dict(transform_metadata) # shallow copy\n other_meta.pop(PARTITION_NAME)\n return transform_metadata[PARTITION_NAME], other_meta\n\n\n# Annotation handling\n# ------------------------------------------------------------------------------\n\n\ndef with_partitioning(\n initializer: F,\n sharding: Sharding,\n mesh: tp.Optional[jax.sharding.Mesh] = None,\n **metadata: tp.Any,\n) -> F:\n \"\"\"A wrapper over any initializer to add sharding annotation data to a `Variable`.\"\"\"\n return variablelib.with_metadata(\n initializer,\n out_sharding=sharding,\n mesh=mesh,\n **metadata,\n )\n\n\ndef get_var_pspec(v: variablelib.Variable) -> PartitionSpec | None:\n \"\"\"Given an `nnx.Variable`, return its `PartitionSpec`.\"\"\"\n metadata = v.get_metadata()\n if 'out_sharding' in metadata and metadata['out_sharding']:\n sharding = metadata['out_sharding']\n if core_spmd.get_logical_axis_rules() or 'sharding_rules' in metadata:\n context_rules = core_spmd.get_logical_axis_rules()\n local_rules = metadata.get('sharding_rules', ())\n rules = core_spmd.composite_rules(context_rules, local_rules)\n return PartitionSpec(*core_spmd.from_sharding_rules(sharding, rules))\n return PartitionSpec(*sharding)\n elif hasattr(v, 'shape'):\n return PartitionSpec()\n return None\n\n\ndef get_partition_spec(tree: A) -> A:\n \"\"\"Extracts a PartitionSpec tree from a PyTree containing ``Variable`` values.\"\"\"\n\n def f(x):\n if isinstance(x, variablelib.Variable):\n return x.replace(get_var_pspec(x))\n elif hasattr(x, 'shape'):\n return PartitionSpec()\n return None\n\n return jax.tree.map(\n f, tree, is_leaf=lambda x: isinstance(x, variablelib.Variable)\n )\n\n\ndef get_named_sharding(tree: A, mesh: jax.sharding.Mesh) -> A:\n spec = get_partition_spec(tree)\n sharding = jax.tree.map(lambda p: jax.sharding.NamedSharding(mesh, p), spec)\n return sharding\n\n\n# Other utilities\n# ------------------------------------------------------------------------------\n\n\ndef get_abstract_model(init_fn, mesh, *, graph: bool | None = None):\n with jax.set_mesh(mesh):\n abs_model = eval_shape(init_fn, graph=graph)\n gdef, abs_state = graphlib.split(abs_model, graph=graph)\n abs_state = jax.tree.map(\n lambda a, s: jax.ShapeDtypeStruct(a.shape, a.dtype, sharding=s),\n abs_state, get_named_sharding(abs_state, mesh)\n )\n return gdef, abs_state\n\n\ndef abstract_with_sharding(\n tree: A, graph: bool | None = None\n) -> A:\n \"\"\"Add sharding information to abstract Variables.\n\n When creating models with :func:`eval_shape`, Variables are abstract\n (backed by ``jax.ShapeDtypeStruct``) and may not carry sharding\n information, especially when using meshes with\n :attr:`jax.sharding.AxisType.Auto` axes. ``abstract_with_sharding`` inspects each\n Variable in ``tree`` and, if it has ``out_sharding`` metadata but no\n sharding already set, attaches a :class:`jax.sharding.NamedSharding`\n derived from the Variable's ``out_sharding`` and either its ``mesh``\n metadata or the current abstract mesh (``jax.sharding.get_abstract_mesh``).\n\n Example usage::\n\n from flax import nnx\n import jax\n\n mesh = jax.make_mesh((2, 2), ('a', 'b'),\n axis_types=(jax.sharding.AxisType.Auto,) * 2)\n with jax.set_mesh(mesh):\n abs_model = nnx.eval_shape(\n lambda: nnx.Linear(4, 8, rngs=nnx.Rngs(0),\n kernel_metadata={'out_sharding': ('a', 'b')}))\n abs_model = nnx.abstract_with_sharding(abs_model)\n assert abs_model.kernel.sharding.spec == jax.P('a', 'b')\n\n Args:\n tree: A graph node (e.g. an :class:`nnx.Module`) whose Variables should\n be annotated with sharding (via ``out_sharding`` metadata).\n graph: Forwarded to :func:`nnx.map`. If ``True``, uses graph-mode;\n if ``False``, uses tree-mode.\n Returns:\n A tree with sharding-annotated ShapeDtypeStruct values inside Variables.\n \"\"\"\n def add_sharding(_path, x):\n if (\n isinstance(x, variablelib.Variable)\n and hasattr(value := x.get_value(), 'shape')\n and hasattr(value, 'dtype')\n and getattr(value, 'sharding', None) is None\n and x.has_metadata('out_sharding')\n ):\n if x.has_metadata('mesh'):\n mesh = x.get_metadata('mesh')\n else:\n mesh = jax.sharding.get_abstract_mesh()\n specs = get_var_pspec(x)\n sharding = jax.sharding.NamedSharding(mesh, specs)\n abs_var = x.replace(\n jax.ShapeDtypeStruct(value.shape, value.dtype, sharding=sharding)\n )\n return abs_var\n return x\n return graphlib.map(add_sharding, tree, graph=graph)\n" }, { "path": "flax/nnx/statelib.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# pytype: skip-file\nfrom __future__ import annotations\n\nimport typing as tp\nfrom collections.abc import MutableMapping\nfrom functools import partial\nimport warnings\n\nimport jax\nimport jax.tree_util as jtu\nimport treescope # type: ignore[import-not-found,import-untyped]\n\nfrom flax.nnx import filterlib, reprlib, traversals, variablelib\nfrom flax.typing import Key, PathParts\n\nA = tp.TypeVar('A')\nK = tp.TypeVar('K', bound=tp.Hashable)\nS = tp.TypeVar('S', bound='State')\nV = tp.TypeVar('V')\n\nExtractValueFn = tp.Callable[[tp.Any], tp.Any]\nSetValueFn = tp.Callable[[V, tp.Any], V]\n\n\nclass NestedStateRepr(reprlib.Representable):\n def __init__(self, state: State):\n self.state = state\n\n def __nnx_repr__(self):\n yield reprlib.Object('', kv_sep=': ', start='{', end='}')\n\n for r in self.state.__nnx_repr__():\n if isinstance(r, reprlib.Object):\n continue\n yield r\n\n def __treescope_repr__(self, path, subtree_renderer):\n children = {}\n for k, v in self.state.items():\n if isinstance(v, State):\n v = NestedStateRepr(v)\n children[k] = v\n # Render as the dictionary itself at the same path.\n return subtree_renderer(children, path=path)\n\nclass FlatState(tp.Sequence[tuple[PathParts, V]], reprlib.Representable):\n __slots__ = ('_keys', '_values')\n\n _keys: tuple[PathParts, ...]\n _values: list[V]\n\n def __init__(self, items: tp.Iterable[tuple[PathParts, V]], /, *, sort: bool):\n keys, values = [], []\n if sort:\n items = sorted(items)\n for key, value in items:\n keys.append(key)\n values.append(value)\n self._keys = tuple(keys)\n self._values = values\n\n @staticmethod\n def from_sorted_keys_values(\n keys: tuple[PathParts, ...], values: list[V], /\n ) -> FlatState[V]:\n flat_state = object.__new__(FlatState)\n flat_state._keys = keys\n flat_state._values = values\n return flat_state\n\n @property\n def paths(self) -> tp.Tuple[PathParts, ...]:\n return self._keys\n\n @property\n def leaves(self) -> list[V]:\n return self._values\n\n def __nnx_repr__(self):\n yield reprlib.Object(type='FlatState', kv_sep='', start='([', end='])')\n\n for value in self:\n yield reprlib.Attr('', value)\n\n @tp.overload\n def __getitem__(self, index: int) -> tuple[PathParts, V]: ...\n @tp.overload\n def __getitem__(self, index: slice) -> FlatState[V]: ...\n def __getitem__(\n self, index: int | slice\n ) -> tuple[PathParts, V] | FlatState[V]:\n if isinstance(index, int):\n return self._keys[index], self._values[index]\n return FlatState(zip(self._keys[index], self._values[index]), sort=False)\n\n def __len__(self) -> int:\n return len(self._keys)\n\n def __iter__(self) -> tp.Iterator[tuple[PathParts, V]]:\n return iter(zip(self._keys, self._values))\n\n def to_nested_state(self) -> State[Key, V]:\n return from_flat_state(self)\n\n @tp.overload\n def split(self, first: filterlib.Filter, /) -> FlatState[V]: ...\n\n @tp.overload\n def split(\n self,\n first: filterlib.Filter,\n second: filterlib.Filter,\n /,\n *filters: filterlib.Filter,\n ) -> tuple[FlatState[V], ...]: ...\n\n @tp.overload\n def split(\n self, /, *filters: filterlib.Filter\n ) -> tp.Union[FlatState[V], tuple[FlatState[V], ...]]: ...\n\n def split( # type: ignore[misc]\n self, first: filterlib.Filter, /, *filters: filterlib.Filter\n ) -> tp.Union[FlatState[V], tuple[FlatState[V], ...]]:\n filters = (first, *filters)\n *flat_states_, rest = _split_state(self, *filters)\n\n if rest:\n raise ValueError(\n 'Non-exhaustive filters, got a non-empty remainder: '\n f'{rest}.\\nUse `...` to match all remaining elements.'\n )\n\n flat_states: FlatState[V] | tuple[FlatState[V], ...]\n if len(flat_states_) == 1:\n flat_states = flat_states_[0]\n else:\n flat_states = tuple(flat_states_)\n return flat_states # type: ignore\n\n @tp.overload\n def filter(self, first: filterlib.Filter, /) -> FlatState[V]: ...\n\n @tp.overload\n def filter(\n self,\n first: filterlib.Filter,\n second: filterlib.Filter,\n /,\n *filters: filterlib.Filter,\n ) -> tuple[FlatState[V], ...]: ...\n\n def filter(\n self,\n first: filterlib.Filter,\n /,\n *filters: filterlib.Filter,\n ) -> tp.Union[FlatState[V], tuple[FlatState[V], ...]]:\n *flat_states_, _rest = _split_state(self, first, *filters)\n\n assert len(flat_states_) == len(filters) + 1\n\n flat_states: FlatState[V] | tuple[FlatState[V], ...]\n if len(flat_states_) == 1:\n flat_states = flat_states_[0]\n else:\n flat_states = tuple(flat_states_)\n\n return flat_states # type: ignore\n\n @staticmethod\n def merge(\n flat_state: tp.Iterable[tuple[PathParts, V]],\n /,\n *flat_states: tp.Iterable[tuple[PathParts, V]],\n ) -> FlatState[V]:\n if not flat_states:\n if isinstance(flat_state, FlatState):\n return flat_state\n return FlatState(flat_state, sort=True)\n flat_states = (flat_state, *flat_states)\n\n return FlatState(\n (elem for flat_state in flat_states for elem in flat_state), sort=True\n )\n\n\ndef _flat_state_pytree_flatten(x: FlatState[V]):\n return x._values, x._keys\n\n\ndef _flat_state_pytree_unflatten(\n keys: tuple[PathParts, ...], values: list[V]\n) -> FlatState[V]:\n flat_state = object.__new__(FlatState)\n flat_state._keys = keys\n flat_state._values = values\n return flat_state\n\n\njax.tree_util.register_pytree_node(\n FlatState,\n _flat_state_pytree_flatten,\n _flat_state_pytree_unflatten,\n)\n\n\nclass State(MutableMapping[K, V], reprlib.Representable):\n \"\"\"A pytree-like ``Mapping`` with hashable and comparable keys.\n \"\"\"\n\n def __init__(\n self,\n mapping: tp.Union[\n tp.Mapping[K, tp.Mapping | V],\n tp.Iterator[tuple[K, tp.Mapping | V]],\n ],\n /,\n *,\n _copy: bool = True,\n ):\n if _copy:\n _mapping = dict(mapping)\n else:\n if not isinstance(mapping, dict):\n raise ValueError(\n 'Expected a dictionary when `_copy=False`, '\n f'got {type(mapping)} instead.'\n )\n _mapping = mapping\n\n if tp.TYPE_CHECKING:\n self._mapping = _mapping\n else:\n super().__setattr__('_mapping', _mapping)\n\n @property\n def raw_mapping(self) -> dict[K, tp.Mapping[K, tp.Any] | V]:\n return self._mapping # type: ignore\n\n def __contains__(self, key) -> bool:\n return key in self._mapping\n\n def __getitem__(self, key: K) -> State | V: # type: ignore\n value = self._mapping[key]\n if isinstance(value, dict):\n return type(self)(value, _copy=False)\n return value # type: ignore[return-value]\n\n def __getattr__(self, key: K) -> State | V: # type: ignore[misc]\n if '_mapping' not in vars(self) or key not in self._mapping:\n raise AttributeError(f\"No attribute '{key}' in State\")\n return self[key]\n\n def __setitem__(self, key: K, value: State | V) -> None:\n if key == '__orig_class__':\n object.__setattr__(self, key, value) # type: ignore\n elif isinstance(value, State):\n self._mapping[key] = value._mapping\n else:\n self._mapping[key] = value\n\n __setattr__ = __setitem__ # type: ignore\n\n def __delitem__(self, key: K) -> None:\n del self._mapping[key]\n\n def __iter__(self) -> tp.Iterator[K]:\n return iter(self._mapping)\n\n def __len__(self) -> int:\n return len(self._mapping)\n\n def __nnx_repr__(self):\n yield reprlib.Object(type(self), kv_sep=': ', start='({', end='})')\n\n for k, v in self.items():\n if isinstance(v, State):\n v = NestedStateRepr(v)\n yield reprlib.Attr(repr(k), v)\n\n def __treescope_repr__(self, path, subtree_renderer):\n children = {}\n for k, v in self.items():\n if isinstance(v, State):\n v = NestedStateRepr(v)\n children[k] = v\n return treescope.repr_lib.render_dictionary_wrapper(\n object_type=type(self),\n wrapped_dict=children,\n path=path,\n subtree_renderer=subtree_renderer,\n )\n\n def map(self, f: tp.Callable[[tuple, V], V]) -> State[K, V]:\n warnings.warn(\n '`flax.nnx.State` will be deprecated and be replaced by the built-in '\n 'Python dict. Please use the equivalent `nnx.map_state` instead.',\n DeprecationWarning,\n )\n return map_state(f, self)\n\n def flat_state(self) -> FlatState[V]:\n warnings.warn(\n '`flax.nnx.State` will be deprecated and be replaced by the built-in '\n 'Python dict. Please use the equivalent `nnx.to_flat_state` instead.',\n DeprecationWarning,\n )\n return to_flat_state(self)\n\n @classmethod\n def from_flat_path(\n cls,\n flat_state: tp.Mapping[PathParts, V] | tp.Iterable[tuple[PathParts, V]],\n /,\n ):\n warnings.warn(\n '`flax.nnx.State` will be deprecated and be replaced by the built-in '\n 'Python dict. Please use the equivalent `nnx.from_flat_state` instead.',\n DeprecationWarning,\n )\n return from_flat_state(flat_state, cls=cls)\n\n def to_pure_dict(self,\n extract_fn: ExtractValueFn | None = None\n ) -> dict[str, tp.Any]:\n warnings.warn(\n '`flax.nnx.State` will be deprecated and be replaced by the built-in '\n 'Python dict. Please use the equivalent `nnx.to_pure_dict` instead.',\n DeprecationWarning,\n )\n return to_pure_dict(self, extract_fn)\n\n def replace_by_pure_dict(self,\n pure_dict: dict[str, tp.Any],\n replace_fn: SetValueFn | None = None):\n warnings.warn(\n '`flax.nnx.State` will be deprecated and be replaced by the built-in '\n 'Python dict. Please use the equivalent `nnx.replace_by_pure_dict` '\n 'instead.',\n DeprecationWarning,\n )\n return replace_by_pure_dict(self, pure_dict, replace_fn)\n\n @tp.overload\n def split(self, first: filterlib.Filter, /) -> State[K, V]: ...\n\n @tp.overload\n def split(\n self,\n first: filterlib.Filter,\n second: filterlib.Filter,\n /,\n *filters: filterlib.Filter,\n ) -> tuple[State[K, V], ...]: ...\n\n @tp.overload\n def split(\n self, /, *filters: filterlib.Filter\n ) -> tp.Union[State[K, V], tuple[State[K, V], ...]]: ...\n\n def split( # type: ignore[misc]\n self, first: filterlib.Filter, /, *filters: filterlib.Filter\n ) -> tp.Union[State[K, V], tuple[State[K, V], ...]]:\n warnings.warn(\n '`flax.nnx.State` will be deprecated and be replaced by the built-in '\n 'Python dict. Please use the equivalent `nnx.split_state` instead.',\n DeprecationWarning,\n )\n return split_state(self, first, *filters)\n\n @tp.overload\n def filter(\n self,\n first: filterlib.Filter,\n /,\n ) -> State[K, V]: ...\n\n @tp.overload\n def filter(\n self,\n first: filterlib.Filter,\n second: filterlib.Filter,\n /,\n *filters: filterlib.Filter,\n ) -> tuple[State[K, V], ...]: ...\n\n def filter(\n self,\n first: filterlib.Filter,\n /,\n *filters: filterlib.Filter,\n ) -> tp.Union[State[K, V], tuple[State[K, V], ...]]:\n warnings.warn(\n '`flax.nnx.State` will be deprecated and be replaced by the built-in '\n 'Python dict. Please use the equivalent `nnx.filter_state` instead.',\n DeprecationWarning,\n )\n return filter_state(self, first, *filters)\n\n @classmethod\n def merge(cls, state: tp.Mapping[K, V], /, *states: tp.Mapping[K, V]):\n warnings.warn(\n '`flax.nnx.State` will be deprecated and be replaced by the built-in '\n 'Python dict. Please use the equivalent `nnx.merge_state` instead.',\n DeprecationWarning,\n )\n return merge_state(state, *states)\n\n def __or__(self, other: State[K, V]) -> State[K, V]:\n if not other:\n return self\n return merge_state(self, other)\n\n def __sub__(self, other: State[K, V]) -> State[K, V]:\n warnings.warn(\n '`flax.nnx.State` will be deprecated and be replaced by the built-in '\n 'Python dict. Please use the equivalent `nnx.diff` instead.',\n DeprecationWarning,\n )\n return diff(self, other)\n\n def __init_subclass__(cls) -> None:\n super().__init_subclass__()\n\n jax.tree_util.register_pytree_with_keys(\n cls,\n _state_flatten_with_keys,\n partial(_state_unflatten, cls), # type: ignore[arg-type]\n )\n\n\n\ndef _state_flatten_with_keys(x: State):\n items = sorted(x._mapping.items())\n children = tuple((jtu.DictKey(key), value) for key, value in items)\n return children, tuple(key for key, _ in items)\n\n\ndef _state_unflatten(\n cls: type[S],\n static: tuple[K, ...],\n leaves: tuple[V, ...] | tuple[dict[K, V]],\n):\n return cls(zip(static, leaves))\n\n\njax.tree_util.register_pytree_with_keys(\n State,\n _state_flatten_with_keys,\n partial(_state_unflatten, State), # type: ignore[arg-type]\n)\n\ndef map_state(f: tp.Callable[[tuple, tp.Any], tp.Any], state: State) -> State:\n \"\"\"Map ``f`` over :class:`State` object.\n\n Arguments:\n f: A function to be mapped\n state: A :class:`State` object.\n Returns:\n New state :class:`State`.\n \"\"\"\n flat_state = to_flat_state(state)\n result = [\n (path, f(path, variable_state)) for path, variable_state in flat_state\n ]\n return from_flat_state(result)\n\n\ndef to_flat_state(state: State) -> FlatState:\n \"\"\"Convert state into flat state\n\n Arguments:\n state: A :class:`State` object.\n Returns:\n Flat state :class:`FlatState`\n \"\"\"\n return FlatState(traversals.flatten_to_sequence(state._mapping), sort=True)\n\n\ndef from_flat_state(\n flat_state: tp.Mapping[PathParts, V] | tp.Iterable[tuple[PathParts, V]],\n *, cls = State, # for compatibility with State subclasses\n) -> State:\n \"\"\"Convert flat state object into :class:`State` object.\n\n Arguments:\n flat_state: A :class:`FlatState` object.\n Returns:\n State :class:`State` object.\n \"\"\"\n if not isinstance(flat_state, tp.Mapping):\n flat_state = dict(flat_state)\n nested_state = traversals.unflatten_mapping(flat_state)\n return cls(nested_state)\n\n\ndef to_pure_dict(\n state: State, extract_fn: ExtractValueFn | None = None\n) -> dict[str, tp.Any]:\n \"\"\"Convert :class:`State` object into pure dictionary state.\n\n Arguments:\n state: A :class:`State` object.\n extract_fn: optional extraction function.\n Returns:\n Pure dictionary.\n \"\"\"\n # Works for nnx.Variable\n if extract_fn is None:\n extract_fn = lambda x: x.get_value() if isinstance(x, variablelib.Variable) else x\n flat_values = {k: extract_fn(x) for k, x in to_flat_state(state)}\n return traversals.unflatten_mapping(flat_values)\n\n\ndef restore_int_paths(pure_dict: dict[str, tp.Any]):\n \"\"\"Restore integer paths from string value in the dict.\n This method can be helpful when restoring the state from a checkpoint as\n pure dictionary:\n\n Example::\n\n >>> from flax import nnx\n >>> import orbax.checkpoint as ocp\n >>> import tempfile\n ...\n >>> model = nnx.List([nnx.Linear(10, 10, rngs=nnx.Rngs(0)) for _ in range(2)])\n >>> pure_dict_state = nnx.to_pure_dict(nnx.state(model))\n >>> list(pure_dict_state.keys())\n [0, 1]\n >>> checkpointer = ocp.StandardCheckpointer()\n >>> with tempfile.TemporaryDirectory() as tmpdir:\n ... checkpointer.save(f'{tmpdir}/ckpt', pure_dict_state)\n ... restored_pure_dict = checkpointer.restore(f'{tmpdir}/ckpt')\n ... list(restored_pure_dict.keys())\n ['0', '1']\n >>> restored_pure_dict = nnx.restore_int_paths(restored_pure_dict)\n >>> list(restored_pure_dict.keys())\n [0, 1]\n\n Arguments:\n pure_dict: state as pure dictionary\n Returns:\n state as pure dictionary with restored integers paths\n \"\"\"\n def try_convert_int(x):\n try:\n return int(x)\n except ValueError:\n return x\n\n fixed = {\n tuple(map(try_convert_int, path)): value\n for path, value in traversals.flatten_mapping(pure_dict).items()\n }\n return traversals.unflatten_mapping(fixed)\n\ndef replace_by_pure_dict(\n state: State, pure_dict: dict[str, tp.Any], replace_fn: SetValueFn | None = None\n):\n \"\"\"Replace input ``state`` values with ``pure_dict`` values.\n\n Arguments:\n state: A :class:`State` object.\n pure_dict: pure dictionary with values to be used for replacement.\n replace_fn: optional replace function.\n \"\"\"\n def try_convert_int(x):\n try:\n return int(x)\n except ValueError:\n return x\n\n # Works for nnx.Variable\n if replace_fn is None:\n replace_fn = lambda x, v: x.replace(v) if hasattr(x, 'replace') else v\n current_flat = dict(to_flat_state(state))\n for kp, v in traversals.flatten_mapping(pure_dict).items():\n # Try exact match first, then integer conversion\n if kp in current_flat:\n matched_key = kp\n else:\n int_kp = tuple(map(try_convert_int, kp))\n if int_kp in current_flat:\n matched_key = int_kp\n else:\n raise ValueError(f'key in pure_dict not available in state: {kp}')\n current_flat[matched_key] = replace_fn(current_flat[matched_key], v)\n state.update(traversals.unflatten_mapping(current_flat))\n\n\n@tp.overload\ndef split_state(state: State, first: filterlib.Filter, /) -> State: ...\n\n\n@tp.overload\ndef split_state(\n state: State,\n first: filterlib.Filter,\n second: filterlib.Filter,\n /,\n *filters: filterlib.Filter,\n) -> tuple[State, ...]: ...\n\n\n@tp.overload\ndef split_state(\n state: State, /, *filters: filterlib.Filter\n) -> tp.Union[State, tuple[State, ...]]: ...\n\n\ndef split_state( # type: ignore[misc]\n state: State, first: filterlib.Filter, /, *filters: filterlib.Filter\n) -> tp.Union[State, tuple[State, ...]]:\n \"\"\"Split a :class:`State` into one or more :class:`State`'s. The\n user must pass at least one ``Filter`` (i.e. :class:`Variable`),\n and the filters must be exhaustive (i.e. they must cover all\n :class:`Variable` types in the ``State``).\n\n Example usage::\n\n >>> from flax import nnx\n\n >>> class Model(nnx.Module):\n ... def __init__(self, rngs):\n ... self.batchnorm = nnx.BatchNorm(2, rngs=rngs)\n ... self.linear = nnx.Linear(2, 3, rngs=rngs)\n ... def __call__(self, x):\n ... return self.linear(self.batchnorm(x))\n\n >>> model = Model(rngs=nnx.Rngs(0))\n >>> state = nnx.state(model)\n >>> param, batch_stats = nnx.split_state(state, nnx.Param, nnx.BatchStat)\n\n Arguments:\n first: The first filter\n *filters: The optional, additional filters to group the state into mutually exclusive substates.\n Returns:\n One or more ``States`` equal to the number of filters passed.\n \"\"\"\n filters = (first, *filters)\n flat_states = _split_state(to_flat_state(state), *filters)\n *states_, rest = (state.to_nested_state() for state in flat_states)\n\n if rest:\n raise ValueError(\n 'Non-exhaustive filters, got a non-empty remainder: '\n f'{rest}.\\nUse `...` to match all remaining elements.'\n )\n\n states: State | tuple[State, ...]\n if len(states_) == 1:\n states = states_[0]\n else:\n states = tuple(states_)\n return states # type: ignore\n\n\n\n@tp.overload\ndef filter_state(\n state: State,\n first: filterlib.Filter,\n /,\n) -> State: ...\n\n\n@tp.overload\ndef filter_state(\n state: State,\n first: filterlib.Filter,\n second: filterlib.Filter,\n /,\n *filters: filterlib.Filter,\n) -> tuple[State, ...]: ...\n\n\ndef filter_state(\n state: State,\n first: filterlib.Filter,\n /,\n *filters: filterlib.Filter,\n) -> tp.Union[State, tuple[State, ...]]:\n \"\"\"Filter a ``State`` into one or more ``State``'s. The\n user must pass at least one ``Filter`` (i.e. :class:`Variable`).\n This method is similar to :meth:`split() `,\n except the filters can be non-exhaustive.\n\n Example usage::\n\n >>> from flax import nnx\n\n >>> class Model(nnx.Module):\n ... def __init__(self, rngs):\n ... self.batchnorm = nnx.BatchNorm(2, rngs=rngs)\n ... self.linear = nnx.Linear(2, 3, rngs=rngs)\n ... def __call__(self, x):\n ... return self.linear(self.batchnorm(x))\n\n >>> model = Model(rngs=nnx.Rngs(0))\n >>> state = nnx.state(model)\n >>> param = nnx.filter_state(state, nnx.Param)\n >>> batch_stats = nnx.filter_state(state, nnx.BatchStat)\n >>> param, batch_stats = nnx.filter_state(state, nnx.Param, nnx.BatchStat)\n\n Arguments:\n first: The first filter\n *filters: The optional, additional filters to group the state into mutually exclusive substates.\n Returns:\n One or more ``States`` equal to the number of filters passed.\n \"\"\"\n flat_states = _split_state(to_flat_state(state), first, *filters)\n *states_, _rest = (state.to_nested_state() for state in flat_states)\n\n assert len(states_) == len(filters) + 1\n\n states: State | tuple[State, ...]\n if len(states_) == 1:\n states = states_[0]\n else:\n states = tuple(states_)\n\n return states # type: ignore\n\n\ndef merge_state(state: tp.Mapping, /, *states: tp.Mapping,\n cls = State # for compatibility with State subclasses\n ) -> State:\n \"\"\"The inverse of :meth:`split() `.\n\n ``merge`` takes one or more :class:`State`'s and creates\n a new :class:`State`.\n\n Example usage::\n\n >>> from flax import nnx\n >>> import jax.numpy as jnp\n\n >>> class Model(nnx.Module):\n ... def __init__(self, rngs):\n ... self.batchnorm = nnx.BatchNorm(2, rngs=rngs)\n ... self.linear = nnx.Linear(2, 3, rngs=rngs)\n ... def __call__(self, x):\n ... return self.linear(self.batchnorm(x))\n\n >>> model = Model(rngs=nnx.Rngs(0))\n >>> params, batch_stats = nnx.state(model, nnx.Param, nnx.BatchStat)\n >>> params['linear']['bias'][...] += 1\n\n >>> state = nnx.merge_state(params, batch_stats)\n >>> nnx.update(model, state)\n >>> assert (model.linear.bias[...] == jnp.array([1, 1, 1])).all()\n\n Args:\n state: A :class:`State` object.\n *states: Additional :class:`State` objects.\n Returns:\n The merged :class:`State`.\n \"\"\"\n if not states:\n if isinstance(state, cls):\n return state\n return cls(state)\n\n states = (state, *states)\n\n new_state: dict[PathParts, tp.Any] = {}\n\n for state in states:\n new_state.update(traversals.flatten_mapping(state)) # type: ignore[attribute-error] # pytype is wrong here\n\n return from_flat_state(new_state, cls=cls)\n\n\ndef diff(state: State, other: State) -> State:\n if not other:\n return state\n\n self_flat = to_flat_state(state)\n other_flat = to_flat_state(other)\n diff = {k: v for k, v in self_flat if k not in other_flat.paths}\n\n return from_flat_state(diff)\n\n\ndef _split_state(\n flat_state: FlatState[V],\n *filters: filterlib.Filter,\n) -> tuple[FlatState[V], ...]:\n for i, filter_ in enumerate(filters):\n if filter_ in (..., True) and i != len(filters) - 1:\n remaining_filters = filters[i + 1 :]\n if not all(f in (..., True) for f in remaining_filters):\n raise ValueError(\n '`...` or `True` can only be used as the last filters, '\n f'got {filter_} it at index {i}.'\n )\n predicates = tuple(map(filterlib.to_predicate, filters))\n\n # we have n + 1 states, where n is the number of predicates\n # the last state is for values that don't match any predicate\n flat_states: tuple[list[tuple[PathParts, V]], ...] = tuple(\n [] for _ in range(len(predicates) + 1)\n )\n\n for path, value in flat_state:\n for i, predicate in enumerate(predicates):\n if predicate(path, value):\n flat_states[i].append((path, value)) # type: ignore[index] # mypy is wrong here?\n break\n else:\n # if we didn't break, set leaf to last state\n flat_states[-1].append((path, value)) # type: ignore[index] # mypy is wrong here?\n\n return tuple(FlatState(flat_state, sort=False) for flat_state in flat_states)\n\n\ndef create_path_filters(state: State):\n flat_state = to_flat_state(state)\n value_paths: dict[tp.Any, set[PathParts]] = {}\n for path, value in flat_state:\n if isinstance(value, variablelib.Variable):\n value = value.get_value()\n value_paths.setdefault(value, set()).add(path)\n return {filterlib.PathIn(*value_paths[value]): value for value in value_paths}" }, { "path": "flax/nnx/summary.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nfrom collections import defaultdict\nimport dataclasses\nimport inspect\nimport io\nimport typing as tp\nfrom types import MappingProxyType\nimport functools\nimport itertools\n\nimport jax\nimport numpy as np\nimport rich.console\nimport rich.table\nimport rich.text\nimport yaml\nimport jax.numpy as jnp\n\nfrom flax import nnx\nfrom flax import typing\nfrom flax.nnx import graphlib, statelib, variablelib\n\nfrom functools import wraps\n\n\ntry:\n from IPython import get_ipython\n\n in_ipython = get_ipython() is not None\nexcept ImportError:\n in_ipython = False\n\n# Custom YAML dumper to represent None as 'None' string (not YAML 'null') for clarity\nclass NoneDumper(yaml.SafeDumper):\n pass\n\nNoneDumper.add_representer(\n type(None),\n lambda dumper, data: dumper.represent_scalar('tag:yaml.org,2002:str', 'None'),\n)\n\nclass SizeBytes(typing.SizeBytes):\n def __repr__(self) -> str:\n bytes_repr = _bytes_repr(self.bytes)\n return f'{self.size:,} [dim]({bytes_repr})[/dim]'\n\nclass ObjectInfo(tp.NamedTuple):\n path: statelib.PathParts\n stats: dict[type[variablelib.Variable], SizeBytes]\n variable_groups: defaultdict[\n type[variablelib.Variable], defaultdict[typing.Key, variablelib.Variable]\n ]\n\nNodeStats = dict[int, ObjectInfo | None]\n\ndef _collect_stats(\n path: statelib.PathParts,\n node: tp.Any,\n node_stats: NodeStats,\n object_types: set[type],\n):\n if not graphlib.is_node(node) and not isinstance(node, variablelib.Variable):\n raise ValueError(f'Expected a graph node or Variable, got {type(node)!r}.')\n\n if id(node) in node_stats:\n return\n\n stats: dict[type[variablelib.Variable], SizeBytes] = {}\n variable_groups: defaultdict[\n type[variablelib.Variable], defaultdict[typing.Key, variablelib.Variable]\n ] = defaultdict(lambda: defaultdict())\n node_stats[id(node)] = ObjectInfo(path, stats, variable_groups)\n\n if isinstance(node, nnx.Pytree):\n node._nnx_tabulate_id = id(node) # type: ignore\n object_types.add(type(node))\n\n node_impl = graphlib.get_node_impl(node)\n assert node_impl is not None\n node_dict = node_impl.node_dict(node)\n for key, value in node_dict.items():\n if id(value) in node_stats:\n continue\n elif isinstance(value, variablelib.Variable):\n var_type = type(value)\n if issubclass(var_type, nnx.RngState):\n var_type = nnx.RngState\n size_bytes = SizeBytes.from_any(value.get_value())\n if var_type in stats:\n stats[var_type] += size_bytes\n else:\n stats[var_type] = size_bytes\n variable_groups[var_type][key] = value\n node_stats[id(value)] = None\n elif graphlib.is_node(value):\n _collect_stats((*path, key), value, node_stats, object_types)\n # accumulate stats from children\n child_info = node_stats[id(value)]\n assert child_info is not None\n for var_type, size_bytes in child_info.stats.items():\n if var_type in stats:\n stats[var_type] += size_bytes\n else:\n stats[var_type] = size_bytes\n\n@dataclasses.dataclass(frozen=True, repr=False)\nclass ArrayRepr:\n shape: tuple[int, ...]\n dtype: tp.Any\n\n @classmethod\n def from_array(cls, x: jax.Array | np.ndarray):\n return cls(jnp.shape(x), jnp.result_type(x))\n\n def __str__(self):\n shape_repr = ','.join(str(x) for x in self.shape)\n return f'[dim]{self.dtype}[/dim][{shape_repr}]'\n\n\n@dataclasses.dataclass\nclass CallInfo:\n call_order: int\n object_id: int\n type: type\n path: statelib.PathParts\n inputs_repr: str\n outputs: tp.Any\n flops: int | None\n vjp_flops: int | None\n\nclass SimpleObjectRepr:\n def __init__(self, obj: tp.Any):\n self.type = type(obj)\n\n def __str__(self):\n return f'{self.type.__name__}(...)'\n\n def __repr__(self):\n return f'{self.type.__name__}(...)'\n\n\ndef _to_dummy_array(x):\n if isinstance(x,jax.ShapeDtypeStruct):\n return ArrayRepr(x.shape, x.dtype)\n elif isinstance(x, jax.Array | np.ndarray):\n return ArrayRepr.from_array(x)\n elif graphlib.is_graph_node(x):\n return SimpleObjectRepr(x)\n else:\n return x\n\ndef _pure_nnx_vjp(f, model, *args, **kwargs):\n \"Wrap nnx functional api around jax.vjp. Only handles pure method calls.\"\n graphdef, state = nnx.split(model)\n def inner(state, *args, **kwargs):\n model = nnx.merge(graphdef, state)\n return f(model, *args, **kwargs)\n return jax.vjp(inner, state, *args, **kwargs)\n\ndef filter_rng_streams(row: CallInfo):\n return not issubclass(row.type, nnx.RngStream)\n\ndef _create_obj_env(object_types):\n \"Turn a set of object types into a dictionary mapping (type, method name) pairs to methods\"\n result = {}\n for obj_type in object_types:\n for name, top_method in inspect.getmembers(obj_type, inspect.isfunction):\n if not name.startswith('_') or name == '__call__':\n result[(obj_type, name)] = top_method\n return result\n\ndef _get_inputs_repr(args, kwargs):\n input_args, input_kwargs = jax.tree.map(\n _to_dummy_array, (args, kwargs)\n )\n inputs_repr = ''\n if input_args:\n if len(input_args) == 1 and not input_kwargs:\n inputs_repr += _as_yaml_str(input_args[0])\n else:\n inputs_repr += _as_yaml_str(input_args)\n if input_kwargs:\n inputs_repr += '\\n'\n if input_kwargs:\n inputs_repr += _as_yaml_str(input_kwargs)\n return inputs_repr\n\ndef _save_call_info(counter, tracer_args, f, node_stats, compute_flops, compute_vjp_flops, seen):\n \"Wrap a function to save its arguments\"\n\n # Used when computing vjp flops\n def do_vjp(*args, **kwargs):\n primals, f_vjp = jax.vjp(f, *args, **kwargs)\n return f_vjp(primals)\n\n method_name = f.__name__\n\n @functools.partial(jax.jit)\n def jit_f(graphdef, state):\n args, kwargs = nnx.merge(graphdef, state)\n return f(*args, **kwargs)\n\n @wraps(f)\n def wrapper(obj, *args, **kwargs):\n inputs_repr = _get_inputs_repr(args, kwargs)\n object_id = getattr(obj, '_nnx_tabulate_id')\n node_info = node_stats[object_id]\n path = node_info.path\n if method_name != '__call__':\n path = (*path, method_name)\n identifier = (inputs_repr, object_id)\n counter_val = next(counter)\n graphdef, state = nnx.split(((obj, *args), kwargs))\n if compute_flops:\n lowered = jit_f.lower(graphdef, state)\n flops = _get_flops(lowered)\n outputs = lowered.out_info\n else:\n flops = None\n outputs = jit_f(graphdef, state)\n if identifier not in seen:\n seen.add(identifier)\n output_repr = jax.tree.map(_to_dummy_array, outputs)\n vjp_flops = _get_flops(jax.jit(do_vjp).lower(\n obj, *args, **kwargs)) if compute_vjp_flops else None\n tracer_args.append(\n CallInfo(counter_val, object_id, type(obj), path, inputs_repr,\n output_repr, flops, vjp_flops))\n return jit_f(graphdef, state)\n return wrapper\n\ndef _overwrite_methods(env):\n \"Overwrite methods with functions from an environment\"\n for (obj_type, name), f in env.items():\n setattr(obj_type, name, f)\n\ndef _get_flops(e) -> int:\n cost = e.cost_analysis() or e.compile().cost_analysis()\n return 0 if cost is None or 'flops' not in cost else int(cost['flops'])\n\ndef tabulate(\n obj,\n *input_args,\n depth: int | None = None,\n method: str = '__call__',\n row_filter: tp.Callable[[CallInfo], bool] = filter_rng_streams,\n table_kwargs: tp.Mapping[str, tp.Any] = MappingProxyType({}),\n column_kwargs: tp.Mapping[str, tp.Any] = MappingProxyType({}),\n console_kwargs: tp.Mapping[str, tp.Any] = MappingProxyType({}),\n compute_flops: bool = False,\n compute_vjp_flops: bool = False,\n **input_kwargs,\n) -> str:\n \"\"\"Creates a summary of the graph object represented as a table.\n\n The table summarizes the object's state and metadata. The table is\n structured as follows:\n\n - The first column represents the path of the object in the graph.\n - The second column represents the type of the object.\n - The third column represents the input arguments passed to the object's\n method.\n - The fourth column represents the output of the object's method.\n - The following columns provide information about the object's state,\n grouped by Variable types.\n\n Example::\n\n >>> from flax import nnx\n ...\n >>> class Block(nnx.Module):\n ... def __init__(self, din, dout, rngs: nnx.Rngs):\n ... self.linear = nnx.Linear(din, dout, rngs=rngs)\n ... self.bn = nnx.BatchNorm(dout, rngs=rngs)\n ... self.dropout = nnx.Dropout(0.2, rngs=rngs)\n ...\n ... def __call__(self, x):\n ... return nnx.relu(self.dropout(self.bn(self.linear(x))))\n ...\n >>> class Foo(nnx.Module):\n ... def __init__(self, rngs: nnx.Rngs):\n ... self.block1 = Block(32, 128, rngs=rngs)\n ... self.block2 = Block(128, 10, rngs=rngs)\n ...\n ... def __call__(self, x):\n ... return self.block2(self.block1(x))\n ...\n >>> foo = Foo(nnx.Rngs(0))\n >>> # print(nnx.tabulate(foo, jnp.ones((1, 32))))\n\n Foo Summary\n ┏━━━━━━━━━━━━━━━━┳━━━━━━━━━━━┳━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━┓\n ┃ path ┃ type ┃ inputs ┃ outputs ┃ BatchStat ┃ Param ┃ RngState ┃\n ┡━━━━━━━━━━━━━━━━╇━━━━━━━━━━━╇━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━┩\n │ │ Foo │ float32[1,32] │ float32[1,10] │ 276 (1.1 KB) │ 5,790 (23.2 KB) │ 2 (12 B) │\n ├────────────────┼───────────┼────────────────┼────────────────┼────────────────────┼─────────────────────────┼──────────┤\n │ block1 │ Block │ float32[1,32] │ float32[1,128] │ 256 (1.0 KB) │ 4,480 (17.9 KB) │ 2 (12 B) │\n ├────────────────┼───────────┼────────────────┼────────────────┼────────────────────┼─────────────────────────┼──────────┤\n │ block1/linear │ Linear │ float32[1,32] │ float32[1,128] │ │ bias: float32[128] │ │\n │ │ │ │ │ │ kernel: float32[32,128] │ │\n │ │ │ │ │ │ │ │\n │ │ │ │ │ │ 4,224 (16.9 KB) │ │\n ├────────────────┼───────────┼────────────────┼────────────────┼────────────────────┼─────────────────────────┼──────────┤\n │ block1/bn │ BatchNorm │ float32[1,128] │ float32[1,128] │ mean: float32[128] │ bias: float32[128] │ │\n │ │ │ │ │ var: float32[128] │ scale: float32[128] │ │\n │ │ │ │ │ │ │ │\n │ │ │ │ │ 256 (1.0 KB) │ 256 (1.0 KB) │ │\n ├────────────────┼───────────┼────────────────┼────────────────┼────────────────────┼─────────────────────────┼──────────┤\n │ block1/dropout │ Dropout │ float32[1,128] │ float32[1,128] │ │ │ 2 (12 B) │\n ├────────────────┼───────────┼────────────────┼────────────────┼────────────────────┼─────────────────────────┼──────────┤\n │ block2 │ Block │ float32[1,128] │ float32[1,10] │ 20 (80 B) │ 1,310 (5.2 KB) │ │\n ├────────────────┼───────────┼────────────────┼────────────────┼────────────────────┼─────────────────────────┼──────────┤\n │ block2/linear │ Linear │ float32[1,128] │ float32[1,10] │ │ bias: float32[10] │ │\n │ │ │ │ │ │ kernel: float32[128,10] │ │\n │ │ │ │ │ │ │ │\n │ │ │ │ │ │ 1,290 (5.2 KB) │ │\n ├────────────────┼───────────┼────────────────┼────────────────┼────────────────────┼─────────────────────────┼──────────┤\n │ block2/bn │ BatchNorm │ float32[1,10] │ float32[1,10] │ mean: float32[10] │ bias: float32[10] │ │\n │ │ │ │ │ var: float32[10] │ scale: float32[10] │ │\n │ │ │ │ │ │ │ │\n │ │ │ │ │ 20 (80 B) │ 20 (80 B) │ │\n ├────────────────┼───────────┼────────────────┼────────────────┼────────────────────┼─────────────────────────┼──────────┤\n │ block2/dropout │ Dropout │ float32[1,10] │ float32[1,10] │ │ │ │\n ├────────────────┼───────────┼────────────────┼────────────────┼────────────────────┼─────────────────────────┼──────────┤\n │ │ │ │ Total │ 276 (1.1 KB) │ 5,790 (23.2 KB) │ 2 (12 B) │\n └────────────────┴───────────┴────────────────┴────────────────┴────────────────────┴─────────────────────────┴──────────┘\n\n Total Parameters: 6,068 (24.3 KB)\n\n Note that ``block2/dropout`` is not shown in the table because it shares the\n same ``RngState`` with ``block1/dropout``.\n\n Args:\n obj: A object to summarize. It can a pytree or a graph objects\n such as nnx.Module or nnx.Optimizer.\n *input_args: Positional arguments passed to the object's method.\n **input_kwargs: Keyword arguments passed to the object's method.\n depth: The depth of the table.\n method: The method to call on the object. Default is ``'__call__'``.\n row_filter: A callable that filters the rows to be displayed in the table.\n By default, it filters out rows with type ``nnx.RngStream``.\n table_kwargs: An optional dictionary with additional keyword arguments\n that are passed to ``rich.table.Table`` constructor.\n column_kwargs: An optional dictionary with additional keyword arguments\n that are passed to ``rich.table.Table.add_column`` when adding columns to\n the table.\n console_kwargs: An optional dictionary with additional keyword arguments\n that are passed to `rich.console.Console` when rendering the table.\n Default arguments are ``'force_terminal': True``, and ``'force_jupyter'``\n is set to ``True`` if the code is running in a Jupyter notebook, otherwise\n it is set to ``False``.\n compute_flops: whether to include a `flops` column in the table listing the\n estimated FLOPs cost of each module forward pass. Does incur actual\n on-device computation / compilation / memory allocation, but still\n introduces overhead for large modules (e.g. extra 20 seconds for a\n Stable Diffusion's UNet, whereas otherwise tabulation would finish in 5\n seconds).\n compute_vjp_flops: whether to include a `vjp_flops` column in the table\n listing the estimated FLOPs cost of each module backward pass. Introduces\n a compute overhead of about 2-3X of `compute_flops`.\n\n Returns:\n A string summarizing the object.\n \"\"\"\n _console_kwargs = {'force_terminal': True, 'force_jupyter': in_ipython}\n _console_kwargs.update(console_kwargs)\n\n obj = graphlib.clone(obj) # create copy to avoid side effects\n node_stats: NodeStats = {}\n object_types: set[type] = set()\n _collect_stats((), obj, node_stats, object_types)\n _variable_types: set[type] = {\n nnx.RngState # type: ignore[misc]\n if isinstance(leaf, nnx.RngState)\n else type(leaf)\n for _, leaf in nnx.to_flat_state(nnx.state(obj))\n }\n variable_types: list[type] = sorted(_variable_types, key=lambda t: t.__name__)\n\n # Create a dictionary-version of the object's class. This makes\n # iteration over methods easier.\n env = _create_obj_env(object_types)\n\n # Information is recorded in post-order, but should be presented as a pre-order traversal.\n # This keeps track of the order of calls.\n counter = itertools.count(0)\n\n # Modify all the object's methods to save their lowered JIT representations.\n rows : list[CallInfo] = []\n seen : set = set()\n jits = {k: _save_call_info(counter, rows, v, node_stats, compute_flops, compute_vjp_flops, seen)\n for k, v in env.items()}\n _overwrite_methods(jits)\n\n # Trace the top function (which indirectly traces all the others)\n jits[(type(obj), method)](obj, *input_args, **input_kwargs)\n\n # Sort call info in pre-order traversal order\n rows.sort(key=lambda x: x.call_order)\n\n # Restore the object's original methods\n _overwrite_methods(env)\n\n if depth is not None:\n rows = [row for row in rows if len(row.path) <= depth and row_filter(row)]\n else:\n rows = [row for row in rows if row_filter(row)]\n\n rich_table = rich.table.Table(\n show_header=True,\n show_lines=True,\n show_footer=True,\n title=f'{type(obj).__name__} Summary',\n **table_kwargs,\n )\n\n rich_table.add_column('path', **column_kwargs)\n rich_table.add_column('type', **column_kwargs)\n rich_table.add_column('inputs', **column_kwargs)\n rich_table.add_column('outputs', **column_kwargs)\n if compute_flops:\n rich_table.add_column('flops', **column_kwargs)\n if compute_vjp_flops:\n rich_table.add_column('vjp_flops', **column_kwargs)\n\n for var_type in variable_types:\n rich_table.add_column(var_type.__name__, **column_kwargs)\n\n for row in rows:\n node_info = node_stats[row.object_id]\n assert node_info is not None\n col_reprs: list[str] = []\n path_str = '/'.join(map(str, row.path))\n col_reprs.append(path_str)\n col_reprs.append(row.type.__name__)\n col_reprs.append(row.inputs_repr)\n col_reprs.append(_as_yaml_str(row.outputs))\n if compute_flops:\n col_reprs.append(str(row.flops))\n if compute_vjp_flops:\n col_reprs.append(str(row.vjp_flops))\n\n for var_type in variable_types:\n attributes = {}\n variable: variablelib.Variable\n for name, variable in node_info.variable_groups[var_type].items():\n value = variable.get_value()\n value_repr = _render_array(value) if _has_shape_dtype(value) else ''\n metadata = variable.get_metadata()\n for required_key in var_type.required_metadata:\n metadata.pop(required_key, None)\n if metadata:\n attributes[name] = {\n 'value': value_repr,\n **metadata,\n }\n elif value_repr:\n attributes[name] = value_repr # type: ignore[assignment]\n\n if attributes:\n col_repr = _as_yaml_str(attributes) + '\\n\\n'\n else:\n col_repr = ''\n\n size_bytes = node_info.stats.get(var_type) # type: ignore[call-overload]\n if size_bytes:\n col_repr += f'[bold]{size_bytes}[/bold]'\n col_reprs.append(col_repr)\n\n rich_table.add_row(*col_reprs)\n\n total_offset = 3 + int(compute_flops) + int(compute_vjp_flops)\n rich_table.columns[total_offset].footer = rich.text.Text.from_markup(\n 'Total', justify='right'\n )\n node_info = node_stats[id(obj)]\n assert node_info is not None\n for i, var_type in enumerate(variable_types):\n size_bytes = node_info.stats[var_type]\n rich_table.columns[i + total_offset + 1].footer = str(size_bytes)\n\n rich_table.caption_style = 'bold'\n total_size = sum(node_info.stats.values(), SizeBytes(0, 0))\n rich_table.caption = f'\\nTotal Parameters: {total_size}'\n\n return _get_rich_repr(rich_table, _console_kwargs)\n\n\ndef _get_rich_repr(obj, console_kwargs):\n f = io.StringIO()\n console = rich.console.Console(file=f, **console_kwargs)\n console.print(obj)\n return f.getvalue()\n\n\ndef _size_and_bytes(pytree: tp.Any) -> tuple[int, int]:\n leaves = jax.tree.leaves(pytree)\n size = sum(x.size for x in leaves if hasattr(x, 'size'))\n num_bytes = sum(\n x.size * x.dtype.itemsize for x in leaves if hasattr(x, 'size')\n )\n return size, num_bytes\n\n\ndef _size_and_bytes_repr(size: int, num_bytes: int) -> str:\n if not size:\n return ''\n bytes_repr = _bytes_repr(num_bytes)\n return f'{size:,} [dim]({bytes_repr})[/dim]'\n\n\ndef _bytes_repr(num_bytes):\n count, units = (\n (f'{num_bytes / 1e9:,.1f}', 'GB')\n if num_bytes > 1e9\n else (f'{num_bytes / 1e6:,.1f}', 'MB')\n if num_bytes > 1e6\n else (f'{num_bytes / 1e3:,.1f}', 'KB')\n if num_bytes > 1e3\n else (f'{num_bytes:,}', 'B')\n )\n\n return f'{count} {units}'\n\n\ndef _has_shape_dtype(value):\n return hasattr(value, 'shape') and hasattr(value, 'dtype')\n\n\ndef _normalize_values(x):\n if isinstance(x, type):\n return f'type[{x.__name__}]'\n elif isinstance(x, ArrayRepr | SimpleObjectRepr):\n return str(x)\n else:\n return repr(x)\n\ndef _maybe_pytree_to_dict(pytree: tp.Any):\n path_leaves = jax.tree_util.tree_flatten_with_path(pytree)[0]\n path_leaves = [\n (tuple(map(graphlib._key_path_to_key, path)), value)\n for path, value in path_leaves\n ]\n if len(path_leaves) < 1:\n return pytree\n elif len(path_leaves) == 1 and path_leaves[0][0] == ():\n return pytree\n else:\n return _unflatten_to_simple_structure(path_leaves, original=pytree)\n\n\ndef _unflatten_to_simple_structure(\n xs: list[tuple[tuple[tp.Any, ...], tp.Any]], *, original: tp.Any\n):\n \"\"\"Rebuild a simple Python structure from path/value leaves.\n\n This variant is aware of the original object so it can:\n - Preserve empty containers that were elided by JAX flattening.\n - Pad trailing missing list/tuple items using the original length.\n - Distinguish placeholders for empty dict/list vs None.\n \"\"\"\n\n def _get_by_path(x, path: tuple[tp.Any, ...]):\n cur = x\n for k in path:\n cur = cur[k]\n return cur\n\n def _to_simple(x):\n # Convert to display-friendly simple structures\n if isinstance(x, (list, tuple)):\n return [_to_simple(e) for e in x]\n if isinstance(x, dict):\n return {k: _to_simple(v) for k, v in x.items()}\n return x\n\n result: list | dict = (\n [] if len(xs) > 0 and isinstance(xs[0][0][0], int) else {}\n )\n for path, value in xs:\n cursor = result\n for i, key in enumerate(path[:-1]):\n if isinstance(cursor, list):\n # Ensure list has slot for current key; infer placeholder from original\n while len(cursor) <= key:\n # path to the slot we are about to create\n slot_path = path[:i] + (len(cursor),)\n try:\n orig_slot = _get_by_path(original, slot_path)\n except Exception:\n orig_slot = None\n if isinstance(orig_slot, (list, tuple)):\n cursor.append([])\n elif isinstance(orig_slot, dict):\n cursor.append({})\n else:\n cursor.append(None)\n else:\n if key not in cursor:\n next_key = path[i + 1]\n if isinstance(next_key, int):\n cursor[key] = []\n else:\n cursor[key] = {}\n cursor = cursor[key]\n if isinstance(cursor, list):\n # Handle gaps in indices caused by JAX flattening eliding empty containers\n while len(cursor) <= path[-1]:\n slot_path = path[:-1] + (len(cursor),)\n try:\n orig_slot = _get_by_path(original, slot_path)\n except Exception:\n orig_slot = None\n if isinstance(orig_slot, (list, tuple)):\n cursor.append([])\n elif isinstance(orig_slot, dict):\n cursor.append({})\n else:\n cursor.append(None)\n cursor[path[-1]] = value\n else:\n assert isinstance(cursor, dict)\n cursor[path[-1]] = value\n # If original is a sequence and result is a list, pad trailing items\n if isinstance(original, (list, tuple)) and isinstance(result, list):\n for i in range(len(result), len(original)):\n slot = original[i]\n result.append(_to_simple(slot))\n return result\n\ndef _as_yaml_str(value) -> str:\n if (hasattr(value, '__len__') and len(value) == 0) or value is None:\n return ''\n\n value = jax.tree.map(_normalize_values, value)\n value = _maybe_pytree_to_dict(value)\n\n file = io.StringIO()\n yaml.dump(\n value,\n file,\n Dumper=NoneDumper,\n default_flow_style=False,\n indent=2,\n sort_keys=False,\n explicit_end=False,\n )\n return file.getvalue().replace('\\n...', '').replace(\"'\", '').strip()\n\n\ndef _render_array(x):\n shape, dtype = jnp.shape(x), jnp.result_type(x)\n shape_repr = ','.join(str(x) for x in shape)\n return f'[dim]{dtype}[/dim][{shape_repr}]'\n\n\ndef _sort_variable_types(types: tp.Iterable[type]) -> list[type]:\n def _variable_parents_count(t: type):\n return sum(1 for p in t.mro() if issubclass(p, variablelib.Variable))\n\n type_sort_key = {t: (-_variable_parents_count(t), t.__name__) for t in types}\n return sorted(types, key=lambda t: type_sort_key[t])\n" }, { "path": "flax/nnx/tracers.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# Taken from flax/core/tracer.py 🏴‍☠️\n\n\nimport jax\nimport jax.core\nimport treescope # type: ignore[import-not-found,import-untyped]\n\nfrom flax.nnx import reprlib\n\n\ndef current_jax_trace():\n \"\"\"Returns the Jax tracing state.\"\"\"\n if jax.__version_info__ <= (0, 4, 33):\n return jax.core.thread_local_state.trace_state.trace_stack.dynamic\n return jax.core.get_opaque_trace_state(convention=\"nnx\")\n\n\nclass TraceState(reprlib.Representable):\n __slots__ = ['_jax_trace']\n\n def __init__(self):\n self._jax_trace = current_jax_trace()\n\n @property\n def jax_trace(self):\n return self._jax_trace\n\n def is_valid(self) -> bool:\n return self._jax_trace == current_jax_trace()\n\n def __nnx_repr__(self):\n yield reprlib.Object(f'{type(self).__name__}')\n yield reprlib.Attr('jax_trace', self._jax_trace)\n\n def __treescope_repr__(self, path, subtree_renderer):\n return treescope.repr_lib.render_object_constructor(\n object_type=type(self),\n attributes={'jax_trace': self._jax_trace},\n path=path,\n subtree_renderer=subtree_renderer,\n )\n\n def __eq__(self, other):\n if jax.__version_info__ <= (0, 4, 33):\n return isinstance(other, TraceState) and self._jax_trace is other._jax_trace\n\n return isinstance(other, TraceState) and self._jax_trace == other._jax_trace\n\n # pickle support\n def __getstate__(self):\n return {}\n\n def __setstate__(self, state):\n self._jax_trace = current_jax_trace()\n\ndef _flatten_trace_state(trace_state: TraceState):\n return (), None\n\n\ndef _unflatten_trace_state(_1, _2):\n return TraceState()\n\n\njax.tree_util.register_pytree_node(\n TraceState,\n _flatten_trace_state,\n _unflatten_trace_state,\n)" }, { "path": "flax/nnx/training/__init__.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n" }, { "path": "flax/nnx/training/metrics.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nimport typing as tp\n\nimport numpy as np\n\nfrom flax import struct\nfrom flax.nnx import filterlib, graphlib\nfrom flax.nnx.pytreelib import Pytree\nfrom flax.nnx.variablelib import Variable\nimport jax, jax.numpy as jnp\n\n# TODO: add tests and docstrings\n\n\nclass MetricState(Variable):\n \"\"\"Wrapper class for Metric Variables.\"\"\"\n\n pass\n\n\nclass Metric(Pytree):\n \"\"\"Base class for metrics. Any class that subclasses ``Metric`` should\n implement a ``compute``, ``reset`` and ``update`` method.\"\"\"\n\n def __init__(self):\n raise NotImplementedError('Must override `__init__()` method.')\n\n def reset(self) -> None:\n \"\"\"In-place reset the ``Metric``.\"\"\"\n raise NotImplementedError('Must override `reset()` method.')\n\n def update(self, **kwargs) -> None:\n \"\"\"In-place update the ``Metric``.\"\"\"\n raise NotImplementedError('Must override `update()` method.')\n\n def compute(self):\n \"\"\"Compute and return the value of the ``Metric``.\"\"\"\n raise NotImplementedError('Must override `compute()` method.')\n\n def split(self, *filters: filterlib.Filter):\n return graphlib.split(self, *filters)\n\n\nclass Average(Metric):\n \"\"\"Average metric.\n\n Example usage::\n\n >>> import jax.numpy as jnp\n >>> from flax import nnx\n\n >>> batch_loss = jnp.array([1, 2, 3, 4])\n >>> batch_loss2 = jnp.array([3, 2, 1, 0])\n\n >>> metrics = nnx.metrics.Average()\n >>> metrics.compute()\n Array(nan, dtype=float32)\n >>> metrics.update(values=batch_loss)\n >>> metrics.compute()\n Array(2.5, dtype=float32)\n >>> metrics.update(values=batch_loss2)\n >>> metrics.compute()\n Array(2., dtype=float32)\n >>> metrics.reset()\n >>> metrics.compute()\n Array(nan, dtype=float32)\n \"\"\"\n\n def __init__(self, argname: str = 'values'):\n \"\"\"Pass in a string denoting the key-word argument that :func:`update` will use to derive the new value.\n For example, constructing the metric as ``avg = Average('test')`` would allow you to make updates with\n ``avg.update(test=new_value)``.\n\n Args:\n argname: an optional string denoting the key-word argument that\n :func:`update` will use to derive the new value. Defaults to\n ``'values'``.\n \"\"\"\n self.argname = argname\n self.total = MetricState(jnp.array(0, dtype=jnp.float32))\n self.count = MetricState(jnp.array(0, dtype=jnp.int32))\n\n def reset(self) -> None:\n \"\"\"Reset this ``Metric``.\"\"\"\n self.total[...] = jnp.array(0, dtype=jnp.float32)\n self.count[...] = jnp.array(0, dtype=jnp.int32)\n\n def update(self, mask: jax.Array | None = None, **kwargs) -> None:\n \"\"\"In-place update this ``Metric``. This method will use the value from\n ``kwargs[self.argname]`` to update the metric, where ``self.argname`` is\n defined on construction.\n\n Args:\n mask: optional mask to ignore the values from the computation.\n We use `values * mask` rule and mask shape should be broadcastable\n to the shape of values array.\n **kwargs: the key-word arguments that contains a ``self.argname``\n entry that maps to the value we want to use to update this metric.\n \"\"\"\n if self.argname not in kwargs:\n raise TypeError(f\"Expected keyword argument '{self.argname}'\")\n values: tp.Union[int, float, jax.Array] = kwargs[self.argname]\n if mask is not None and isinstance(values, (int, float)):\n raise ValueError(f\"If mask is provided, {self.argname} should be a jax array\")\n if isinstance(values, (int, float)):\n self.total[...] += values\n self.count[...] += 1\n elif mask is None:\n self.total[...] += values.sum()\n self.count[...] += values.size\n else:\n self.total[...] += (values * mask.astype(values.dtype)).sum()\n self.count[...] += mask.sum().astype(self.count.dtype)\n\n def compute(self) -> jax.Array:\n \"\"\"Compute and return the average.\"\"\"\n return self.total / self.count\n\n\n@struct.dataclass\nclass Statistics:\n mean: jnp.float32\n standard_error_of_mean: jnp.float32\n standard_deviation: jnp.float32\n\n\nclass Welford(Metric):\n \"\"\"Uses Welford's algorithm to compute the mean and variance of a stream of data.\n\n Example usage::\n\n >>> import jax.numpy as jnp\n >>> from flax import nnx\n\n >>> batch_loss = jnp.array([1, 2, 3, 4])\n >>> batch_loss2 = jnp.array([3, 2, 1, 0])\n\n >>> metrics = nnx.metrics.Welford()\n >>> metrics.compute()\n Statistics(mean=Array(0., dtype=float32), standard_error_of_mean=Array(nan, dtype=float32), standard_deviation=Array(nan, dtype=float32))\n >>> metrics.update(values=batch_loss)\n >>> metrics.compute()\n Statistics(mean=Array(2.5, dtype=float32), standard_error_of_mean=Array(0.559017, dtype=float32), standard_deviation=Array(1.118034, dtype=float32))\n >>> metrics.update(values=batch_loss2)\n >>> metrics.compute()\n Statistics(mean=Array(2., dtype=float32), standard_error_of_mean=Array(0.43301272, dtype=float32), standard_deviation=Array(1.2247449, dtype=float32))\n >>> metrics.reset()\n >>> metrics.compute()\n Statistics(mean=Array(0., dtype=float32), standard_error_of_mean=Array(nan, dtype=float32), standard_deviation=Array(nan, dtype=float32))\n \"\"\"\n\n def __init__(self, argname: str = 'values'):\n \"\"\"Pass in a string denoting the key-word argument that :func:`update` will use to derive the new value.\n For example, constructing the metric as ``wf = Welford('test')`` would allow you to make updates with\n ``wf.update(test=new_value)``.\n\n Args:\n argname: an optional string denoting the key-word argument that\n :func:`update` will use to derive the new value. Defaults to\n ``'values'``.\n \"\"\"\n self.argname = argname\n self.count = MetricState(jnp.array(0, dtype=jnp.int32))\n self.mean = MetricState(jnp.array(0, dtype=jnp.float32))\n self.m2 = MetricState(jnp.array(0, dtype=jnp.float32))\n\n def reset(self) -> None:\n \"\"\"Reset this ``Metric``.\"\"\"\n self.count[...] = jnp.array(0, dtype=jnp.uint32)\n self.mean[...] = jnp.array(0, dtype=jnp.float32)\n self.m2[...] = jnp.array(0, dtype=jnp.float32)\n\n def update(self, **kwargs) -> None:\n \"\"\"In-place update this ``Metric``. This method will use the value from\n ``kwargs[self.argname]`` to update the metric, where ``self.argname`` is\n defined on construction.\n\n Args:\n **kwargs: the key-word arguments that contains a ``self.argname``\n entry that maps to the value we want to use to update this metric.\n \"\"\"\n if self.argname not in kwargs:\n raise TypeError(f\"Expected keyword argument '{self.argname}'\")\n values: tp.Union[int, float, jax.Array] = kwargs[self.argname]\n count = 1 if isinstance(values, (int, float)) else values.size\n original_count = self.count[...]\n self.count[...] += count\n delta = (\n values if isinstance(values, (int, float)) else values.mean()\n ) - self.mean\n self.mean[...] += delta * count / self.count\n m2 = 0.0 if isinstance(values, (int, float)) else values.var() * count\n self.m2[...] += m2 + delta * delta * count * original_count / self.count\n\n def compute(self) -> Statistics:\n \"\"\"Compute and return the mean and variance statistics in a\n ``Statistics`` dataclass object.\n \"\"\"\n variance = self.m2 / self.count\n standard_deviation = variance**0.5\n sem = standard_deviation / (self.count**0.5)\n return Statistics(\n mean=self.mean[...],\n standard_error_of_mean=sem,\n standard_deviation=standard_deviation,\n )\n\n\nclass Accuracy(Average):\n \"\"\"Accuracy metric. This metric subclasses :class:`Average`,\n and so they share the same ``reset`` and ``compute`` method\n implementations. Unlike :class:`Average`, no string needs to\n be passed to ``Accuracy`` during construction.\n\n Example usage::\n\n >>> from flax import nnx\n >>> import jax, jax.numpy as jnp\n\n >>> logits = jax.random.normal(jax.random.key(0), (5, 2))\n >>> labels = jnp.array([0, 1, 1, 1, 0])\n >>> logits2 = jax.random.normal(jax.random.key(1), (5, 2))\n >>> labels2 = jnp.array([0, 1, 1, 1, 1])\n\n >>> metrics = nnx.metrics.Accuracy()\n >>> metrics.compute()\n Array(nan, dtype=float32)\n >>> metrics.update(logits=logits, labels=labels)\n >>> metrics.compute()\n Array(0.6, dtype=float32)\n >>> metrics.update(logits=logits2, labels=labels2)\n >>> metrics.compute()\n Array(0.4, dtype=float32)\n >>> metrics.reset()\n >>> metrics.compute()\n Array(nan, dtype=float32)\n\n >>> logits3 = jax.random.normal(jax.random.key(2), (5,))\n >>> labels3 = jnp.array([0, 1, 0, 1, 1])\n >>> accuracy = nnx.metrics.Accuracy(threshold=0.5)\n >>> accuracy.update(logits=logits3, labels=labels3)\n >>> accuracy.compute()\n Array(0.8, dtype=float32)\n \"\"\"\n\n def __init__(self, threshold: float | None = None, *args, **kwargs):\n \"\"\"For binary classification, pass in a float denoting a threshold to determine if a\n prediction is positive. For example, constructing the metric as\n ``acc = Accuracy(threshold=0.5)`` would cause any logit greater than or equal to 0.5\n to be interpreted as a positive classification. For multi-class classification, do\n not pass in a threshold.\n\n Args:\n threshold: for binary classification, determines if a prediction is\n positive. Defaults to None.\n \"\"\"\n if (threshold is not None) and (not isinstance(threshold, float)):\n raise TypeError(f'Expected threshold to be a float, got {type(threshold)}')\n\n self.threshold = threshold\n super().__init__(*args, **kwargs)\n\n def update( # type: ignore[override]\n self, *, logits: jax.Array, labels: jax.Array, mask: jax.Array | None = None, **_\n ) -> None:\n \"\"\"In-place update this ``Metric``.\n\n Args:\n logits: the outputted predicted activations. For multi-class\n classification, these values are argmax-ed (on the trailing\n dimension), before comparing them to the labels. For binary\n classification, these values are compared to the labels directly.\n labels: the ground truth integer labels.\n mask: optional mask to ignore the logits and labels values from the computation.\n We use `array * mask` rule and mask shape should be broadcastable\n to the shape of labels array.\n \"\"\"\n if self.threshold is not None: # Binary classification case\n if logits.ndim != labels.ndim:\n raise ValueError(\n 'For binary classification, expected logits.ndim==labels.ndim, got '\n f'{logits.ndim} and {labels.ndim}'\n )\n elif logits.ndim != labels.ndim + 1: # Multi-class classification case\n raise ValueError(\n 'For multi-class classification, expected logits.ndim==labels.ndim+1, '\n f'got {logits.ndim} and {labels.ndim}'\n )\n\n if labels.dtype in (jnp.int64, np.int32, np.int64):\n labels = jnp.astype(labels, jnp.int32)\n elif labels.dtype != jnp.int32:\n raise ValueError(f'Expected labels.dtype==jnp.int32, got {labels.dtype}')\n\n if self.threshold is not None: # Binary classification case\n super().update(values=((logits >= self.threshold) == (labels > 0)), mask=mask)\n return\n\n # Multi-class classification case\n super().update(values=(logits.argmax(axis=-1) == labels), mask=mask)\n\n\nclass MultiMetric(Metric):\n \"\"\"MultiMetric class to store multiple metrics and update them in a single call.\n\n Example usage::\n\n >>> from flax import nnx\n >>> import jax, jax.numpy as jnp\n\n >>> metrics = nnx.MultiMetric(\n ... accuracy=nnx.metrics.Accuracy(), loss=nnx.metrics.Average()\n ... )\n\n >>> metrics\n MultiMetric( # MetricState: 4 (16 B)\n accuracy=Accuracy( # MetricState: 2 (8 B)\n threshold=None,\n argname='values',\n total=MetricState( # 1 (4 B)\n value=Array(0., dtype=float32)\n ),\n count=MetricState( # 1 (4 B)\n value=Array(0, dtype=int32)\n )\n ),\n loss=Average( # MetricState: 2 (8 B)\n argname='values',\n total=MetricState( # 1 (4 B)\n value=Array(0., dtype=float32)\n ),\n count=MetricState( # 1 (4 B)\n value=Array(0, dtype=int32)\n )\n )\n )\n\n >>> metrics.accuracy\n Accuracy( # MetricState: 2 (8 B)\n threshold=None,\n argname='values',\n total=MetricState( # 1 (4 B)\n value=Array(0., dtype=float32)\n ),\n count=MetricState( # 1 (4 B)\n value=Array(0, dtype=int32)\n )\n )\n\n >>> metrics.loss\n Average( # MetricState: 2 (8 B)\n argname='values',\n total=MetricState( # 1 (4 B)\n value=Array(0., dtype=float32)\n ),\n count=MetricState( # 1 (4 B)\n value=Array(0, dtype=int32)\n )\n )\n\n >>> logits = jax.random.normal(jax.random.key(0), (5, 2))\n >>> labels = jnp.array([0, 1, 1, 1, 0])\n >>> logits2 = jax.random.normal(jax.random.key(1), (5, 2))\n >>> labels2 = jnp.array([0, 1, 1, 1, 1])\n\n >>> batch_loss = jnp.array([1, 2, 3, 4])\n >>> batch_loss2 = jnp.array([3, 2, 1, 0])\n\n >>> metrics.compute()\n {'accuracy': Array(nan, dtype=float32), 'loss': Array(nan, dtype=float32)}\n >>> metrics.update(logits=logits, labels=labels, values=batch_loss)\n >>> metrics.compute()\n {'accuracy': Array(0.6, dtype=float32), 'loss': Array(2.5, dtype=float32)}\n >>> metrics.update(logits=logits2, labels=labels2, values=batch_loss2)\n >>> metrics.compute()\n {'accuracy': Array(0.4, dtype=float32), 'loss': Array(2., dtype=float32)}\n >>> metrics.reset()\n >>> metrics.compute()\n {'accuracy': Array(nan, dtype=float32), 'loss': Array(nan, dtype=float32)}\n \"\"\"\n\n def __init__(self, **metrics):\n \"\"\"Pass in key-word arguments to the constructor, e.g.\n ``MultiMetric(keyword1=Average(), keyword2=Accuracy(), ...)``.\n\n Args:\n **metrics: the key-word arguments that will be used to access\n the corresponding ``Metric``.\n \"\"\"\n # TODO: raise error if a kwarg is passed that is in ('reset', 'update', 'compute'), since these names are reserved for methods\n self._metric_names = []\n for metric_name, metric in metrics.items():\n self._metric_names.append(metric_name)\n setattr(self, metric_name, metric)\n\n def reset(self) -> None:\n \"\"\"Reset all underlying ``Metric``'s.\"\"\"\n for metric_name in self._metric_names:\n getattr(self, metric_name).reset()\n\n def update(self, **updates) -> None:\n \"\"\"In-place update all underlying ``Metric``'s in this ``MultiMetric``. All\n ``**updates`` will be passed to the ``update`` method of all underlying\n ``Metric``'s.\n\n Args:\n **updates: the key-word arguments that will be passed to the underlying ``Metric``'s\n ``update`` method. It can contain ``mask`` argument as mask to be passed to the underlying metrics.\n ``mask`` can be a ``jax.Array` and it will be passed to all the metrics. ``mask`` can be a dictionary\n with metric name as keys and ``jax.Array`` as values, e.g ``{metric_name1: metric_mask1, ...}``\n \"\"\"\n # TODO: should we give the option of updating only some of the metrics and not all? e.g. if for some kwargs==None, don't do update\n # TODO: should we raise an error if a kwarg is passed into **updates that has no match with any underlying metric? e.g. user typo\n mask = updates.pop(\"mask\", None)\n for metric_name in self._metric_names:\n metric_mask_kwarg = {}\n metric_mask = mask.get(metric_name, None) if isinstance(mask, dict) else mask\n if metric_mask is not None:\n metric_mask_kwarg = {\"mask\": metric_mask}\n getattr(self, metric_name).update(**(updates | metric_mask_kwarg))\n\n def compute(self) -> dict[str, tp.Any]:\n \"\"\"Compute and return the value of all underlying ``Metric``'s. This method\n will return a dictionary, mapping strings (defined by the key-word arguments\n ``**metrics`` passed to the constructor) to the corresponding metric value.\n \"\"\"\n return {\n f'{metric_name}': getattr(self, metric_name).compute()\n for metric_name in self._metric_names\n }\n" }, { "path": "flax/nnx/training/optimizer.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nimport functools\nimport typing as tp\n\nimport jax\nimport jax.numpy as jnp\nimport optax\n\nfrom flax import nnx\nfrom flax.nnx import filterlib\nfrom flax.nnx.pytreelib import Pytree\nfrom flax.nnx.variablelib import Variable\n\nM = tp.TypeVar('M', bound=nnx.Module)\nF = tp.TypeVar('F', bound=tp.Callable[..., tp.Any])\n\nclass OptState(Variable):\n \"\"\"Any optimizer state\"\"\"\n\n pass\n\n\nclass OptArray(OptState):\n \"\"\"Optimizer state for an array.\"\"\"\n\n pass\n\n\nclass OptVariable(OptState):\n \"\"\"Optimizer state for a Variable.\"\"\"\n\n pass\n\n\ndef to_opt_state(tree):\n def _to_opt_state(x):\n if isinstance(x, Variable):\n opt_state = OptVariable(x.get_value(), **x.get_metadata()) # type: ignore\n else:\n opt_state = OptArray(x)\n return opt_state\n\n tree = jax.tree.map(\n _to_opt_state,\n tree,\n is_leaf=lambda x: isinstance(x, Variable),\n )\n return tree\n\nclass _Missing:\n pass\n\nMISSING = _Missing()\n\ndef _check_grads_arg_passed(f: F) -> F:\n @functools.wraps(f)\n def _check_grads_wrapper(self, model, grads=MISSING, **kwargs):\n if isinstance(grads, _Missing):\n raise TypeError(\n 'Missing required argument `grads`. As of Flax 0.11.0 update requires both (model, grads) arguments '\n 'to be passed. If you want to keep the previous use nnx.ModelAndOptimizer instead of nnx.Optimizer.'\n )\n return f(self, model, grads, **kwargs)\n return _check_grads_wrapper # type: ignore\n\ndef _check_wrt_arg_passed(f: F) -> F:\n @functools.wraps(f)\n def _check_wrt_wrapper(*args, wrt=MISSING, **kwargs):\n if isinstance(wrt, _Missing):\n raise TypeError(\n 'Missing required argument `wrt`. As of Flax 0.11.0 the `wrt` argument is required, '\n 'if you want to keep the previous use nnx.ModelAndOptimizer instead of nnx.Optimizer.'\n )\n return f(*args, wrt=wrt, **kwargs)\n return _check_wrt_wrapper # type: ignore\n\nclass Optimizer(Pytree, tp.Generic[M]):\n \"\"\"Simple train state for the common case with a single Optax optimizer.\n\n Example usage::\n\n >>> import jax, jax.numpy as jnp\n >>> from flax import nnx\n >>> import optax\n ...\n >>> class Model(nnx.Module):\n ... def __init__(self, rngs):\n ... self.linear1 = nnx.Linear(2, 3, rngs=rngs)\n ... self.linear2 = nnx.Linear(3, 4, rngs=rngs)\n ... def __call__(self, x):\n ... return self.linear2(self.linear1(x))\n ...\n >>> x = jax.random.normal(jax.random.key(0), (1, 2))\n >>> y = jnp.ones((1, 4))\n ...\n >>> model = Model(nnx.Rngs(0))\n >>> tx = optax.adam(1e-3)\n >>> optimizer = nnx.Optimizer(model, tx, wrt=nnx.Param)\n ...\n >>> loss_fn = lambda model: ((model(x) - y) ** 2).mean()\n >>> loss_fn(model)\n Array(2.3359997, dtype=float32)\n >>> grads = nnx.grad(loss_fn)(model)\n >>> optimizer.update(model, grads)\n >>> loss_fn(model)\n Array(2.310461, dtype=float32)\n\n Attributes:\n step: An ``OptState`` :class:`Variable` that tracks the step count.\n tx: An Optax gradient transformation.\n opt_state: The Optax optimizer state.\n \"\"\"\n\n @_check_wrt_arg_passed\n def __init__(\n self,\n model: M,\n tx: optax.GradientTransformation,\n *,\n wrt: filterlib.Filter, # type: ignore\n ):\n \"\"\"\n Instantiate the class and wrap the :class:`Module` and Optax gradient\n transformation. Instantiate the optimizer state to keep track of\n :class:`Variable` types specified in ``wrt``. Set the step count to 0.\n\n Args:\n model: An NNX Module.\n tx: An Optax gradient transformation.\n wrt: filter to specify for which :class:`Variable`'s to keep\n track of in the optimizer state. These should be the :class:`Variable`'s\n that you plan on updating; i.e. this argument value should match the\n ``wrt`` argument passed to the ``nnx.grad`` call that will generate the\n gradients that will be passed into the ``grads`` argument of the\n :func:`update` method. The filter should match the filter used in nnx.grad.\n \"\"\"\n self.step = OptState(jnp.array(0, dtype=jnp.uint32))\n self.tx = tx\n self.opt_state = nnx.data(\n to_opt_state(tx.init(nnx.state(model, wrt)))\n )\n self.wrt = wrt\n\n if not tp.TYPE_CHECKING:\n def __getattribute__(self, name: str) -> tp.Any:\n if name == 'model' and name not in vars(self):\n raise AttributeError(\n f\"{type(self).__name__} does not have attribute 'model' since Flax 0.11.0. \"\n \"To keep the previous behavior, use nnx.ModelAndOptimizer instead of nnx.Optimizer.\"\n )\n return super().__getattribute__(name)\n\n @_check_grads_arg_passed\n def update(self, model: M, grads, /, **kwargs):\n \"\"\"Updates the optimizer state and model parameters given the gradients.\n\n Example::\n\n >>> from flax import nnx\n >>> import jax, jax.numpy as jnp\n >>> import optax\n ...\n >>> class Model(nnx.Module):\n ... def __init__(self, rngs):\n ... self.linear = nnx.Linear(2, 3, rngs=rngs)\n ... self.count = nnx.Variable(jnp.array(0))\n ...\n ... def __call__(self, x):\n ... self.count[...] += 1\n ... return self.linear(x)\n ...\n >>> model = Model(rngs=nnx.Rngs(0))\n ...\n >>> loss_fn = lambda model, x, y: ((model(x) - y) ** 2).mean()\n >>> optimizer = nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param)\n >>> grads = nnx.grad(loss_fn, argnums=nnx.DiffState(0, nnx.Param))(\n ... model, jnp.ones((1, 2)), jnp.ones((1, 3))\n ... )\n >>> optimizer.update(model, grads)\n\n Note that internally this function calls ``.tx.update()`` followed by a call\n to ``optax.apply_updates()`` to update ``params`` and ``opt_state``.\n\n Args:\n grads: the gradients derived from ``nnx.grad``.\n **kwargs: additional keyword arguments passed to the tx.update, to support\n ``GradientTransformationExtraArgs``, such as ``optax.scale_by_backtracking_linesearch``.\n \"\"\"\n param_arrays = nnx.pure(nnx.state(model, self.wrt))\n grad_arrays = nnx.pure(nnx.state(grads, self.wrt))\n opt_state_arrays = nnx.pure(self.opt_state)\n kwargs_arrays = nnx.pure(kwargs)\n\n updates, new_opt_state = self.tx.update(\n grad_arrays, opt_state_arrays, param_arrays, **kwargs_arrays\n )\n new_params = optax.apply_updates(param_arrays, updates)\n\n nnx.update(model, new_params)\n nnx.update(self.opt_state, nnx.state(new_opt_state))\n self.step[...] += 1\n\nclass ModelAndOptimizer(Optimizer[M]):\n \"\"\"A convenience class that combines a model and an optimizer.\n\n This class is deprecated and will be removed in a future release.\n Use :class:`Optimizer` instead.\n \"\"\"\n\n def __init__(self, model: M, tx: optax.GradientTransformation, *, wrt: filterlib.Filter = nnx.Param):\n super().__init__(model, tx, wrt=wrt)\n self.model = model\n\n def update(self, grads, /, **kwargs): # type: ignore\n return super().update(self.model, grads, **kwargs)\n" }, { "path": "flax/nnx/transforms/__init__.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n" }, { "path": "flax/nnx/transforms/autodiff.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom collections import deque\nimport dataclasses\nimport functools\nimport typing as tp\n\n\nfrom flax import struct\nfrom flax.nnx import (\n extract,\n filterlib,\n graphlib,\n variablelib,\n)\nfrom flax.nnx.statelib import State\nimport jax\n\nfrom flax.nnx.transforms import general\nfrom flax.nnx.transforms.transforms import (\n resolve_kwargs,\n _resolve_bound_callable,\n _raise_bound_method_error,\n)\nfrom flax.typing import MISSING, Missing\n\n\nA = tp.TypeVar('A')\n# C = tp.TypeVar('C')\n# B = tp.TypeVar('B')\nF = tp.TypeVar('F', bound=tp.Callable[..., tp.Any])\n# G = tp.TypeVar('G', bound=tp.Callable[..., tp.Any])\n# M = tp.TypeVar('M', bound=Module)\n# MA = tp.TypeVar('MA', bound=Module)\n# N = tp.TypeVar('N', bound=Module)\n# StrInt = tp.TypeVar('StrInt', str, int)\nAxisName = tp.Hashable\n# Leaves = tp.List[Leaf]\n# Index = int\n\n\n# -------------------------------\n# grad\n# -------------------------------\n\n\n@dataclasses.dataclass(frozen=True)\nclass DiffState:\n argnum: int\n filter: filterlib.Filter\n\n\n@dataclasses.dataclass(eq=False)\nclass SimpleGradFn:\n f: tp.Callable[..., tp.Any]\n has_aux: bool\n graph: bool\n\n def __post_init__(self):\n functools.update_wrapper(self, self.f, updated=())\n\n @extract.treemap_copy_args\n def __call__(self, *args, **kwargs):\n updates, snapshot = extract.updates_and_snapshot((args, kwargs))\n if self.graph:\n args, kwargs = extract.from_tree2((args, kwargs))\n out = self.f(*args, **kwargs)\n if self.graph:\n out = extract.to_tree2(out)\n extract.check_no_aliases('grad', args=updates[0], kwargs=updates[1], out=out)\n updates = extract.mask_variable_updates(updates, snapshot)\n\n if self.has_aux:\n loss, aux = out\n return loss, (updates, aux)\n else:\n return out, updates\n\n\n@dataclasses.dataclass(eq=False)\nclass GradFn:\n f: tp.Callable[..., tp.Any]\n has_aux: bool\n nondiff_states: deque[State | None]\n\n def __post_init__(self):\n functools.update_wrapper(self, self.f)\n\n def __call__(self, *pure_args):\n # rebuild diff_state from substates in args\n\n def _grad_merge_fn(\n ctx: graphlib.MergeContext, path, prefix, value: extract.NodeStates\n ):\n nondiff = self.nondiff_states.popleft()\n if nondiff is None:\n return ctx.merge(value.graphdef, value.state)\n else:\n return ctx.merge(value.graphdef, value.state, nondiff)\n\n args = extract.from_tree(\n pure_args, merge_fn=_grad_merge_fn, ctxtag='grad', is_inner=True\n )\n\n out = self.f(*args)\n\n args_out = extract.clear_non_graph_nodes(args)\n pure_args_out, pure_out = extract.to_tree((args_out, out), ctxtag='grad')\n\n if self.has_aux:\n loss, pure_aux = pure_out\n fn_out = (loss, (pure_args_out, pure_aux))\n else:\n loss = pure_out\n fn_out = (loss, pure_args_out)\n\n return fn_out\n\n\ndef _grad_general(\n f: tp.Callable[..., tp.Any],\n argnums: int | DiffState | tp.Sequence[int | DiffState],\n has_aux: bool,\n holomorphic: bool,\n allow_int: bool,\n return_value: bool,\n graph: bool,\n graph_updates: bool,\n) -> tp.Callable[..., tp.Any]:\n\n transform = jax.value_and_grad if return_value else jax.grad\n\n if not graph or not graph_updates:\n if any(isinstance(x, DiffState) for x in jax.tree.leaves(argnums)):\n raise ValueError(\n '`argnums` cannot contain `DiffState` objects '\n 'when `graph=False`. '\n + graphlib._tree_mode_suggestion_transform('grad')\n )\n\n gradded_fn = transform(\n SimpleGradFn(f, has_aux, graph=graph),\n argnums=argnums, # type: ignore[arg-type]\n has_aux=True,\n holomorphic=holomorphic,\n allow_int=allow_int,\n )\n\n def tree_grad_wrapper(*args, **kwargs):\n if graph:\n diff_argnums = (argnums,) if isinstance(argnums, int) else argnums\n args_prefix = tuple(\n i in diff_argnums for i in range(len(args))\n )\n args, kwargs = extract.to_tree2(\n (args, kwargs), prefix=(args_prefix, False),\n )\n\n extract.check_no_aliases('grad', args=args, kwargs=kwargs)\n\n fn_out = gradded_fn(*args, **kwargs)\n\n if return_value:\n if has_aux:\n (loss, (updates, aux)), grads = fn_out\n if graph: grads, aux = extract.from_tree2((grads, aux))\n result = (loss, aux), grads\n else:\n (loss, updates), grads = fn_out\n if graph: grads = extract.from_tree2(grads)\n result = loss, grads\n else:\n if has_aux:\n grads, (updates, aux) = fn_out\n if graph: grads, aux = extract.from_tree2((grads, aux))\n result = grads, aux\n else:\n grads, updates = fn_out\n if graph: grads = extract.from_tree2(grads)\n result = grads\n\n extract.apply_variable_updates((args, kwargs), updates)\n return result\n\n return tree_grad_wrapper\n\n jax_argnums: int | tuple[int, ...]\n if isinstance(argnums, (int, DiffState)):\n jax_argnums = argnums.argnum if isinstance(argnums, DiffState) else argnums\n else:\n jax_argnums = tuple(\n x.argnum if isinstance(x, DiffState) else x for x in argnums\n )\n\n _argnums = (argnums,) if isinstance(argnums, (int, DiffState)) else argnums\n index_filter: dict[int, DiffState] = {}\n for argnum in _argnums:\n index = argnum.argnum if isinstance(argnum, DiffState) else argnum\n if index in index_filter:\n raise ValueError(f'argnum {index} is repeated in argnums')\n index_filter[index] = (\n dataclasses.replace(argnum, argnum=-1)\n if isinstance(argnum, DiffState)\n else DiffState(-1, variablelib.Param)\n )\n\n @graphlib.update_context('grad')\n def grad_wrapper(*args, **kwargs):\n args = resolve_kwargs(f, args, kwargs)\n del kwargs\n nondiff_states: deque[State | variablelib.Variable | None] = deque()\n\n def _grad_split_fn(\n ctx: graphlib.SplitContext, path, prefix: DiffState | None, value\n ):\n if prefix is None or (prefix.argnum == -1 and isinstance(value, variablelib.Variable)):\n nondiff_states.append(None)\n return extract.NodeStates.from_split(*ctx.split(value))\n else:\n graphdef, diff, nondiff = ctx.split(value, prefix.filter, ...) # type: ignore[misc]\n nondiff_states.append(nondiff) # type: ignore[container-type-mismatch]\n return extract.NodeStates.from_split(graphdef, diff)\n\n arg_filters = tuple(index_filter.get(i) for i in range(len(args)))\n pure_args = extract.to_tree(\n args, prefix=arg_filters, split_fn=_grad_split_fn, ctxtag='grad'\n )\n\n gradded_fn = transform(\n GradFn(f, has_aux, nondiff_states),\n argnums=jax_argnums,\n has_aux=True,\n holomorphic=holomorphic,\n allow_int=allow_int,\n )\n\n fn_out = gradded_fn(*pure_args)\n\n def process_grads(grads):\n return jax.tree.map(\n lambda x: x.state if isinstance(x, extract.NodeStates) else x,\n grads,\n is_leaf=lambda x: isinstance(x, extract.NodeStates),\n )\n\n def process_out(pure_out: A, /) -> A:\n return extract.from_tree(pure_out, ctxtag='grad', is_inner=False)\n\n if return_value:\n # unpack value_and_grad output\n if has_aux:\n (loss, (pure_args_out, pure_aux)), grads = fn_out\n grads = process_grads(grads)\n _args_out, aux = process_out((pure_args_out, pure_aux))\n return (loss, aux), grads\n else:\n (loss, pure_args_out), grads = fn_out\n grads = process_grads(grads)\n _args_out = process_out(pure_args_out)\n return loss, grads\n else:\n # unpack grad output\n if has_aux:\n grads, (pure_args_out, pure_aux) = fn_out\n grads = process_grads(grads)\n _args_out, aux = process_out((pure_args_out, pure_aux))\n return grads, aux\n else:\n grads, pure_args_out = fn_out\n grads = process_grads(grads)\n _args_out = process_out(pure_args_out)\n return grads\n\n return grad_wrapper\n\n\n@tp.overload\ndef grad(\n f: tp.Callable[..., tp.Any],\n *,\n argnums: int | DiffState | tp.Sequence[int | DiffState] = 0,\n has_aux: bool = False,\n holomorphic: bool = False,\n allow_int: bool = False,\n reduce_axes: tp.Sequence[AxisName] = (),\n graph: bool | None = None,\n graph_updates: bool | None = None,\n) -> tp.Callable[..., tp.Any]: ...\n@tp.overload\ndef grad(\n *,\n argnums: int | DiffState | tp.Sequence[int | DiffState] = 0,\n has_aux: bool = False,\n holomorphic: bool = False,\n allow_int: bool = False,\n reduce_axes: tp.Sequence[AxisName] = (),\n graph: bool | None = None,\n graph_updates: bool | None = None,\n) -> tp.Callable[[tp.Callable[..., tp.Any]], tp.Callable[..., tp.Any]]: ...\ndef grad(\n f: tp.Callable[..., tp.Any] | Missing = MISSING,\n *,\n argnums: int | DiffState | tp.Sequence[int | DiffState] = 0,\n has_aux: bool = False,\n holomorphic: bool = False,\n allow_int: bool = False,\n reduce_axes: tp.Sequence[AxisName] = (),\n graph: bool | None = None,\n graph_updates: bool | None = None,\n) -> (\n tp.Callable[..., tp.Any]\n | tp.Callable[[tp.Callable[..., tp.Any]], tp.Callable[..., tp.Any]]\n):\n \"\"\"Object-aware version of ``jax.grad`` that can handle Modules / graph nodes as\n arguments.\n\n Example::\n\n >>> from flax import nnx\n >>> import jax.numpy as jnp\n ...\n >>> m = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n >>> x = jnp.ones((1, 2))\n >>> y = jnp.ones((1, 3))\n ...\n >>> loss_fn = lambda m, x, y: jnp.mean((m(x) - y) ** 2)\n >>> grad_fn = nnx.grad(loss_fn)\n ...\n >>> grads = grad_fn(m, x, y)\n >>> jax.tree.map(jnp.shape, grads)\n State({\n 'bias': Param(\n value=(3,)\n ),\n 'kernel': Param(\n value=(2, 3)\n )\n })\n\n By default, NNX objects are differentiated with respect to all their ``nnx.Param``\n Variables. You can specify which substates are differentiable by passing a ``DiffState``\n object to the ``argnums`` argument. For example, if you want to differentiate only the\n ``kernel`` attribute of the ``Linear`` class, you can use the ``PathContains`` filter::\n\n >>> m = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n ...\n >>> kernel_attribute = nnx.PathContains('kernel')\n >>> diff_state = nnx.DiffState(0, kernel_attribute)\n ...\n >>> loss_fn = lambda m, x, y: jnp.mean((m(x) - y) ** 2)\n >>> grad_fn = nnx.grad(loss_fn, argnums=diff_state)\n ...\n >>> grads = grad_fn(m, x, y)\n >>> jax.tree.map(jnp.shape, grads)\n State({\n 'kernel': Param(\n value=(2, 3)\n )\n })\n\n For more information on how to create custom filters, see\n `Using Filters `__\n guide.\n\n Args:\n fun: Function to be differentiated. Its arguments at positions specified by\n ``argnums`` should be arrays, scalars, graph nodes or standard Python\n containers. Argument arrays in the positions specified by ``argnums`` must\n be of inexact (i.e., floating-point or complex) type. It should return a\n scalar (which includes arrays with shape ``()`` but not arrays with shape\n ``(1,)`` etc.)\n argnums: Optional, integer or sequence of integers. Specifies which\n positional argument(s) to differentiate with respect to (default 0).\n has_aux: Optional, bool. Indicates whether ``fun`` returns a pair where the\n first element is considered the output of the mathematical function to be\n differentiated and the second element is auxiliary data. Default False.\n holomorphic: Optional, bool. Indicates whether ``fun`` is promised to be\n holomorphic. If True, inputs and outputs must be complex. Default False.\n allow_int: Optional, bool. Whether to allow differentiating with\n respect to integer valued inputs. The gradient of an integer input will\n have a trivial vector-space dtype (float0). Default False.\n graph: If ``True`` (default), uses graph-mode which supports the full\n NNX feature set including shared references and reference semantics.\n If ``False``, uses tree-mode which treats Modules as regular JAX\n pytrees, avoiding the overhead of the graph protocol. Tree-mode does\n not support ``DiffState`` or shared ``Variable`` references.\n graph_updates: If ``True``, propagates updates on graph structure\n that happen inside the transform to the input graphs, has no\n effect when ``graph=False``. When ``False``, using ``DiffState``\n is not supported.\n \"\"\"\n if graph is None:\n graph = graphlib.set_graph_mode.current_value()\n if graph_updates is None:\n graph_updates = graphlib.set_graph_updates.current_value()\n if reduce_axes:\n raise NotImplementedError('reduce_axes argument to grad is deprecated')\n del reduce_axes\n\n if isinstance(f, Missing):\n return functools.partial(\n grad,\n argnums=argnums,\n has_aux=has_aux,\n holomorphic=holomorphic,\n allow_int=allow_int,\n graph=graph,\n graph_updates=graph_updates,\n )\n # Detect bound nnx.Module methods and raise error.\n f_unbound, _, was_bound = _resolve_bound_callable(f)\n\n if was_bound:\n _raise_bound_method_error('grad')\n\n return _grad_general(\n f_unbound,\n argnums,\n has_aux,\n holomorphic,\n allow_int,\n return_value=False,\n graph=graph,\n graph_updates=graph_updates,\n )\n\n\n@tp.overload\ndef value_and_grad(\n f: tp.Callable[..., tp.Any],\n *,\n argnums: int | DiffState | tp.Sequence[int | DiffState] = 0,\n has_aux: bool = False,\n holomorphic: bool = False,\n allow_int: bool = False,\n reduce_axes: tp.Sequence[AxisName] = (),\n graph: bool | None = None,\n graph_updates: bool | None = None,\n) -> tp.Callable[..., tp.Any]: ...\n@tp.overload\ndef value_and_grad(\n *,\n argnums: int | DiffState | tp.Sequence[int | DiffState] = 0,\n has_aux: bool = False,\n holomorphic: bool = False,\n allow_int: bool = False,\n reduce_axes: tp.Sequence[AxisName] = (),\n graph: bool | None = None,\n graph_updates: bool | None = None,\n) -> tp.Callable[[tp.Callable[..., tp.Any]], tp.Callable[..., tp.Any]]: ...\ndef value_and_grad(\n f: tp.Callable[..., tp.Any] | type[Missing] = Missing,\n *,\n argnums: int | DiffState | tp.Sequence[int | DiffState] = 0,\n has_aux: bool = False,\n holomorphic: bool = False,\n allow_int: bool = False,\n reduce_axes: tp.Sequence[AxisName] = (),\n graph: bool | None = None,\n graph_updates: bool | None = None,\n) -> (\n tp.Callable[..., tp.Any]\n | tp.Callable[[tp.Callable[..., tp.Any]], tp.Callable[..., tp.Any]]\n):\n \"\"\"Object-aware version of ``jax.value_and_grad``.\n\n Like :func:`grad`, but returns both the value and the gradient of ``f``.\n\n Args:\n f: Function to be differentiated. Its arguments at positions specified by\n ``argnums`` should be arrays, scalars, graph nodes or standard Python\n containers. Argument arrays in the positions specified by ``argnums`` must\n be of inexact (i.e., floating-point or complex) type. It should return a\n scalar (which includes arrays with shape ``()`` but not arrays with shape\n ``(1,)`` etc.)\n argnums: Optional, integer or sequence of integers. Specifies which\n positional argument(s) to differentiate with respect to (default 0).\n has_aux: Optional, bool. Indicates whether ``f`` returns a pair where the\n first element is considered the output of the mathematical function to be\n differentiated and the second element is auxiliary data. Default False.\n holomorphic: Optional, bool. Indicates whether ``f`` is promised to be\n holomorphic. If True, inputs and outputs must be complex. Default False.\n allow_int: Optional, bool. Whether to allow differentiating with\n respect to integer valued inputs. The gradient of an integer input will\n have a trivial vector-space dtype (float0). Default False.\n graph: If ``True`` (default), uses graph-mode which supports the full\n NNX feature set including shared references and reference semantics.\n If ``False``, uses tree-mode which treats Modules as regular JAX\n pytrees, avoiding the overhead of the graph protocol. Tree-mode does\n not support ``DiffState`` or shared ``Variable`` references.\n graph_updates: If ``True``, propagates updates on graph structure\n that happen inside the transform to the input graphs, has no\n effect when ``graph=False``. When ``False``, using ``DiffState``\n is not supported.\n\n Returns:\n A function with the same arguments as ``f`` that evaluates both ``f``\n and the gradient of ``f`` and returns them as a pair (a two-element\n tuple). If ``argnums`` is an integer then the gradient has the same shape\n and type as the positional argument indicated by that integer. If argnums is\n a sequence of integers, the gradient is a tuple of values with the same\n shapes and types as the corresponding arguments. If ``has_aux`` is True\n then a tuple of ((value, auxiliary_data), gradient) is returned.\n \"\"\"\n if graph is None:\n graph = graphlib.set_graph_mode.current_value()\n if graph_updates is None:\n graph_updates = graphlib.set_graph_updates.current_value()\n if reduce_axes:\n raise NotImplementedError(\n 'reduce_axes argument to value_and_grad is deprecated')\n del reduce_axes\n\n if f is Missing:\n return functools.partial(\n value_and_grad,\n argnums=argnums,\n has_aux=has_aux,\n holomorphic=holomorphic,\n allow_int=allow_int,\n graph=graph,\n graph_updates=graph_updates,\n )\n # Detect bound nnx.Module methods and raise error.\n f_unbound, _, was_bound = _resolve_bound_callable(f)\n\n if was_bound:\n _raise_bound_method_error('value_and_grad')\n\n return _grad_general(\n f_unbound,\n argnums,\n has_aux,\n holomorphic,\n allow_int,\n return_value=True,\n graph=graph,\n graph_updates=graph_updates,\n )\n\n# -----------------------------------------------\n# vjp\n# -----------------------------------------------\n\n\n@dataclasses.dataclass(eq=False)\nclass SimpleVjpFn:\n f: tp.Callable[..., tp.Any]\n has_aux: bool\n graph: bool\n\n def __post_init__(self):\n functools.update_wrapper(self, self.f, updated=())\n\n @extract.treemap_copy_args\n def __call__(self, *args):\n updates, snapshot = extract.updates_and_snapshot(args)\n if self.graph:\n args = extract.from_tree2(args)\n out = self.f(*args)\n if self.graph:\n out = extract.to_tree2(out)\n extract.check_no_aliases('vjp', args=updates, out=out)\n updates = extract.mask_variable_updates(updates, snapshot)\n if self.has_aux:\n primals_out, aux = out\n return primals_out, (updates, aux)\n else:\n return out, updates\n\n\n@tp.overload\ndef vjp(\n f: tp.Callable[..., tp.Any],\n *primals: tp.Any,\n has_aux: bool = False,\n reduce_axes: tp.Sequence[AxisName] = (),\n graph: bool | None = None,\n graph_updates: bool | None = None,\n) -> tuple[tp.Any, tp.Callable] | tuple[tp.Any, tp.Callable, tp.Any]: ...\n@tp.overload\ndef vjp(\n *,\n has_aux: bool = False,\n reduce_axes: tp.Sequence[AxisName] = (),\n graph: bool | None = None,\n graph_updates: bool | None = None,\n) -> tp.Callable[[tp.Callable[..., tp.Any]], tp.Callable[..., tp.Any]]: ...\ndef vjp(\n f: tp.Callable[..., tp.Any] | Missing = MISSING,\n *primals: tp.Any,\n has_aux: bool = False,\n reduce_axes: tp.Sequence[AxisName] = (),\n graph: bool | None = None,\n graph_updates: bool | None = None,\n) -> (\n tuple[tp.Any, tp.Callable]\n | tuple[tp.Any, tp.Callable, tp.Any]\n | tp.Callable[[tp.Callable[..., tp.Any]], tp.Callable[..., tp.Any]]\n):\n \"\"\"Stateful version of ``jax.vjp`` that propagates NNX Variable updates.\n\n Example::\n\n >>> from flax import nnx\n >>> import jax.numpy as jnp\n ...\n >>> m = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n >>> x = jnp.ones((1, 2))\n ...\n >>> def loss_fn(m, x):\n ... return jnp.sum(m(x))\n ...\n >>> primals_out, vjp_fn = nnx.vjp(loss_fn, m, x, graph=False)\n >>> m_grad, x_grad = vjp_fn(jnp.ones_like(primals_out))\n\n Can also be used as a decorator::\n\n >>> @nnx.vjp(graph=False)\n ... def f(m, x):\n ... return jnp.sum(m(x))\n ...\n >>> primals_out, vjp_fn = f(m, x)\n\n Args:\n f: Function to be differentiated. Its arguments can be arrays, scalars,\n or pytrees containing arrays and NNX Variables.\n *primals: A sequence of primal values at which the Jacobian of ``f``\n should be evaluated.\n has_aux: Optional, bool. Indicates whether ``f`` returns a pair where the\n first element is considered the output of the mathematical function to be\n differentiated and the second element is auxiliary data. Default False.\n reduce_axes: Deprecated, do not use.\n graph: If ``True`` (default), uses graph-mode which supports the full\n NNX feature set including shared references and reference semantics.\n If ``False``, uses tree-mode which treats Modules as regular JAX\n pytrees, avoiding the overhead of the graph protocol.\n graph_updates: If ``True``, propagates updates on graph structure\n that happen inside the transform to the input graphs, has no\n effect when ``graph=False``.\n\n Returns:\n If ``has_aux`` is False, returns a ``(primals_out, vjp_fn)`` pair.\n ``vjp_fn`` takes a cotangent with the same structure as ``primals_out``\n and returns gradients for each primal argument.\n If ``has_aux`` is True, returns ``(primals_out, vjp_fn, aux)``.\n \"\"\"\n if graph is None:\n graph = graphlib.set_graph_mode.current_value()\n if graph_updates is None:\n graph_updates = graphlib.set_graph_updates.current_value()\n if graph and graph_updates:\n raise NotImplementedError(\n 'graph-mode with graph_updates is not supported for nnx.vjp. '\n 'Set graph=False or graph_updates=False.'\n )\n if reduce_axes:\n raise NotImplementedError('reduce_axes argument to vjp is deprecated')\n del reduce_axes\n\n if isinstance(f, Missing):\n return functools.partial( # type: ignore[return-value]\n vjp,\n has_aux=has_aux,\n graph=graph,\n graph_updates=graph_updates,\n )\n\n f_unbound, _, was_bound = _resolve_bound_callable(f)\n if was_bound:\n _raise_bound_method_error('vjp')\n\n if not primals:\n return functools.partial( # type: ignore[return-value]\n vjp,\n f,\n has_aux=has_aux,\n graph=graph,\n graph_updates=graph_updates,\n )\n\n if graph:\n primals = extract.to_tree2(primals)\n extract.check_no_aliases('vjp', primals=primals)\n primals_out, vjp_fn, aux = jax.vjp(\n SimpleVjpFn(f_unbound, has_aux=has_aux, graph=graph),\n *primals,\n has_aux=True,\n )\n if has_aux:\n updates, user_aux = aux\n else:\n updates = aux\n user_aux = None\n if graph:\n primals_out = extract.from_tree2(primals_out)\n raw_vjp_fn = vjp_fn\n def vjp_fn(g):\n return extract.from_tree2(raw_vjp_fn(g))\n extract.apply_variable_updates(primals, updates)\n if has_aux:\n return primals_out, vjp_fn, user_aux\n else:\n return primals_out, vjp_fn\n\n\n# -----------------------------------------------\n# jvp\n# -----------------------------------------------\n\n\n@dataclasses.dataclass(eq=False)\nclass SimpleJvpFn:\n f: tp.Callable[..., tp.Any]\n has_aux: bool\n graph: bool\n\n def __post_init__(self):\n functools.update_wrapper(self, self.f, updated=())\n\n @extract.treemap_copy_args\n def __call__(self, *args):\n updates, snapshot = extract.updates_and_snapshot(args)\n if self.graph:\n args = extract.from_tree2(args)\n out = self.f(*args)\n if self.graph:\n out = extract.to_tree2(out)\n extract.check_no_aliases('jvp', args=updates, out=out)\n updates = extract.mask_variable_updates(updates, snapshot)\n if self.has_aux:\n primals_out, aux = out\n return (primals_out, updates), aux\n else:\n return out, updates\n\n\n@tp.overload\ndef jvp(\n f: tp.Callable[..., tp.Any],\n primals: tuple[tp.Any, ...],\n tangents: tuple[tp.Any, ...],\n *,\n has_aux: bool = False,\n graph: bool | None = None,\n graph_updates: bool | None = None,\n) -> tuple[tp.Any, ...]: ...\n@tp.overload\ndef jvp(\n *,\n has_aux: bool = False,\n graph: bool | None = None,\n graph_updates: bool | None = None,\n) -> tp.Callable[[tp.Callable[..., tp.Any]], tp.Callable[..., tp.Any]]: ...\n@tp.overload\ndef jvp(\n f: tp.Callable[..., tp.Any],\n *,\n has_aux: bool = False,\n graph: bool | None = None,\n graph_updates: bool | None = None,\n) -> tp.Callable[..., tp.Any]: ...\ndef jvp(\n f: tp.Callable[..., tp.Any] | Missing = MISSING,\n primals: tuple[tp.Any, ...] | Missing = MISSING,\n tangents: tuple[tp.Any, ...] | Missing = MISSING,\n *,\n has_aux: bool = False,\n graph: bool | None = None,\n graph_updates: bool | None = None,\n) -> (\n tuple[tp.Any, ...]\n | tp.Callable[..., tp.Any]\n | tp.Callable[[tp.Callable[..., tp.Any]], tp.Callable[..., tp.Any]]\n):\n \"\"\"Stateful version of ``jax.jvp`` that propagates NNX Variable updates.\n\n Example::\n\n >>> from flax import nnx\n >>> import jax.numpy as jnp\n ...\n >>> m = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n >>> x = jnp.ones((1, 2))\n ...\n >>> def f(m, x):\n ... return jnp.sum(m(x))\n ...\n >>> m_tangent = jax.tree.map(jnp.zeros_like, m)\n >>> x_tangent = jnp.ones_like(x)\n >>> primals_out, tangent_out = nnx.jvp(\n ... f, (m, x), (m_tangent, x_tangent), graph=False\n ... )\n\n Can also be used as a decorator::\n\n >>> @nnx.jvp(graph=False)\n ... def f(m, x):\n ... return jnp.sum(m(x))\n ...\n >>> primals_out, tangent_out = f((m, x), (m_tangent, x_tangent))\n\n Args:\n f: Function to be differentiated. Its arguments can be arrays, scalars,\n or pytrees containing arrays and NNX Variables.\n primals: A tuple of primal values at which the Jacobian of ``f``\n should be evaluated.\n tangents: A tuple of tangent vectors, with the same structure as\n ``primals``.\n has_aux: Optional, bool. Indicates whether ``f`` returns a pair where the\n first element is considered the output of the mathematical function to be\n differentiated and the second element is auxiliary data. Default False.\n graph: If ``True`` (default), uses graph-mode which supports the full\n NNX feature set including shared references and reference semantics.\n If ``False``, uses tree-mode which treats Modules as regular JAX\n pytrees, avoiding the overhead of the graph protocol.\n graph_updates: If ``True``, propagates updates on graph structure\n that happen inside the transform to the input graphs, has no\n effect when ``graph=False``.\n\n Returns:\n If ``has_aux`` is False, returns ``(primals_out, tangent_out)``.\n If ``has_aux`` is True, returns ``(primals_out, tangent_out, aux)``.\n \"\"\"\n if graph is None:\n graph = graphlib.set_graph_mode.current_value()\n if graph_updates is None:\n graph_updates = graphlib.set_graph_updates.current_value()\n if graph and graph_updates:\n raise NotImplementedError(\n 'graph-mode with graph_updates is not supported for nnx.jvp. '\n 'Set graph=False or graph_updates=False.'\n )\n\n if isinstance(f, Missing):\n return functools.partial(\n jvp,\n has_aux=has_aux,\n graph=graph,\n graph_updates=graph_updates,\n )\n\n f_unbound, _, was_bound = _resolve_bound_callable(f)\n if was_bound:\n _raise_bound_method_error('jvp')\n\n if isinstance(primals, Missing) or isinstance(tangents, Missing):\n return functools.partial(\n jvp,\n f,\n has_aux=has_aux,\n graph=graph,\n graph_updates=graph_updates,\n )\n\n if graph:\n primals = extract.to_tree2(primals)\n tangents = extract.to_tree2(tangents)\n extract.check_no_aliases('jvp', primals=primals)\n extract.check_no_aliases('jvp', tangents=tangents)\n if has_aux:\n (primals_out, updates), (tangent_out, _updates_tangent), aux = jax.jvp(\n SimpleJvpFn(f_unbound, has_aux=True, graph=graph),\n primals,\n tangents,\n has_aux=True,\n )\n else:\n (primals_out, updates), (tangent_out, _updates_tangent) = jax.jvp(\n SimpleJvpFn(f_unbound, has_aux=False, graph=graph),\n primals,\n tangents,\n )\n if graph:\n primals_out = extract.from_tree2(primals_out)\n tangent_out = extract.from_tree2(tangent_out)\n extract.apply_variable_updates(primals, updates)\n if has_aux:\n return primals_out, tangent_out, aux\n else:\n return primals_out, tangent_out\n\n\n# -----------------------------------------------\n# custom_vjp\n# -----------------------------------------------\n\n\n@dataclasses.dataclass(eq=False)\nclass SimpleCustomVjpFn:\n f: tp.Callable[..., tp.Any]\n graph: bool\n nondiff_argnums: tuple[int, ...]\n\n def __post_init__(self):\n functools.update_wrapper(self, self.f, updated=())\n\n @extract.treemap_copy_args\n def __call__(self, *args):\n updates, snapshot = extract.updates_and_snapshot(args)\n if self.graph:\n args = extract.from_tree2(args)\n out = self.f(*args)\n if self.graph:\n out = extract.to_tree2(out)\n extract.check_no_aliases('custom_vjp', args=updates, out=out)\n diff_prefix = tuple(\n i not in self.nondiff_argnums for i in range(len(args))\n )\n def keep_fn(path, diff_arg, cur, snap):\n assert isinstance(diff_arg, bool)\n changed = extract.variable_changed(cur, snap)\n if diff_arg and changed:\n raise ValueError(\n f'Variables in differentiable argument were mutated inside '\n f'custom_vjp at {jax.tree_util.keystr(path)}.\\n'\n f'This is not supported when '\n f'graph_updates=False because the gradient for the Variable '\n f'updates would be silently dropped. Move the Variable mutation '\n f'to a non-differentiable argument, or use graph_updates=True.'\n )\n return changed\n updates = extract.mask_variable_updates(\n updates, snapshot, prefix=diff_prefix, keep_fn=keep_fn,\n )\n return out, updates\n\n\n@dataclasses.dataclass(eq=False)\nclass SimpleFwdFn:\n fwd: tp.Callable[..., tp.Any]\n graph: bool\n\n def __post_init__(self):\n functools.update_wrapper(self, self.fwd, updated=())\n\n @extract.treemap_copy_args\n def __call__(self, *args):\n updates, snapshot = extract.updates_and_snapshot(args)\n if self.graph:\n args = extract.from_tree2(args)\n out, residual = self.fwd(*args)\n if self.graph:\n out = extract.to_tree2(out)\n residual = extract.to_tree2(residual)\n extract.check_no_aliases('custom_vjp', args=updates, out=out)\n updates = extract.mask_variable_updates(updates, snapshot)\n return (out, updates), residual\n\n\n@dataclasses.dataclass(eq=False)\nclass SimpleBwdFn:\n bwd: tp.Callable[..., tp.Any]\n graph: bool\n\n def __post_init__(self):\n functools.update_wrapper(self, self.bwd, updated=())\n\n @extract.treemap_copy_args\n def __call__(self, *args):\n *nondiff, residual, (out_g, _updates_g) = args\n if self.graph:\n nondiff = extract.from_tree2(nondiff)\n residual = extract.from_tree2(residual)\n result = self.bwd(*nondiff, residual, out_g)\n if self.graph:\n result = extract.to_tree2(result)\n return result\n\n\nclass SimpleCustomVjp(tp.Generic[A]):\n def __init__(\n self,\n fun: tp.Callable[..., A],\n nondiff_argnums: tuple[int, ...],\n graph: bool,\n ):\n functools.update_wrapper(self, fun)\n self.fun = fun\n self.nondiff_argnums = nondiff_argnums\n self.graph = graph\n self.custom_vjp_fn = jax.custom_vjp(\n fun=SimpleCustomVjpFn(fun, graph=graph, nondiff_argnums=nondiff_argnums),\n nondiff_argnums=nondiff_argnums,\n )\n\n def __call__(\n self, *args: tp.Any, **kwargs: tp.Any\n ) -> A:\n args = resolve_kwargs(self.fun, args, kwargs)\n del kwargs\n if self.graph:\n prefix = tuple(\n i not in self.nondiff_argnums for i in range(len(args))\n )\n args = extract.to_tree2(args, prefix=prefix)\n extract.check_no_aliases('custom_vjp', args=args)\n (out, updates) = self.custom_vjp_fn(*args)\n if self.graph:\n out = extract.from_tree2(out)\n extract.apply_variable_updates(args, updates)\n return out\n\n def defvjp(\n self,\n fwd: tp.Callable[..., tuple[A, tp.Any]],\n bwd: tp.Callable[..., tuple[tp.Any, ...]],\n symbolic_zeros: bool = False,\n ) -> None:\n self.fwd = fwd\n self.bwd = bwd\n self.symbolic_zeros = symbolic_zeros\n self.custom_vjp_fn.defvjp(\n fwd=SimpleFwdFn(fwd, graph=self.graph),\n bwd=SimpleBwdFn(bwd, graph=self.graph),\n symbolic_zeros=symbolic_zeros,\n )\n\n\n# custom_vjp is one of the most complicated transforms as it requires\n# to handle 4 different functions:\n# 1. CustomVJP: the main object that runs the outer logic, converts input graph nodes\n# to pytrees and output pytrees to graph nodes.\n# 2. CustomVjpFnWrapper: function that wraps the user's function, it converts\n# its input pytrees to graph nodes and output graph nodes to pytrees.\n# 3. FwdFn: wraps the user's fwd function, it converts its input pytrees to graph nodes\n# and output graph nodes to pytrees. Since it might run by itself in a separate context,\n# it needs to be aware if the update_context is active or not in order to update the outer\n# referenes.\n# 4. BwdFn: wraps the user's bwd function, it converts its input pytrees to graph nodes\n# and output graph nodes to pytrees. It doesn't need to be aware of the outer context\n# since it will never update the outer references as it runs during the backward pass.\n\ndef _custom_vjp_merge_fn(\n ctx: graphlib.MergeContext,\n path,\n prefix: bool | DiffState,\n value: extract.NodeStates,\n *,\n nondiff_states: deque[extract.GraphDefState],\n):\n nondiff = nondiff_states.popleft()\n return ctx.merge(nondiff.graphdef, value.state, nondiff.state)\n\n\ndef _custom_vjp_split_fn(\n ctx: graphlib.SplitContext,\n path,\n prefix: bool | DiffState,\n value,\n *,\n nondiff_states: list[extract.GraphDefState],\n):\n broadcast: graphlib.GraphState\n if prefix is False:\n # pure non-differentiable arg, not supported\n raise TypeError(\n 'Passing integers to nondiff_argnums for graph nodes arguments in custom_vjp is not supported. '\n f'Got {prefix} at path {jax.tree_util.keystr(path)} for value {value}'\n )\n elif prefix is True:\n # pure differentiable arg, we pass all the state through\n # but we return a TreeNode.from_states which doesn't have a graphdef\n # in order to keep the gradients clean from any metadata\n graphdef, passed = ctx.split(value)\n broadcast = State({})\n nondiff_states.append(extract.GraphDefState(graphdef, broadcast))\n return extract.NodeStates.from_states(passed)\n else:\n # differentiable arg with DiffState filter, we use the filter to split the state\n # as before we return a TreeNode.from_states to keep the gradients clean\n # from any metadata, the non-differentiable state is stored in a deque\n # which is broadcasted during the forward pass\n graphdef, passed, broadcast = ctx.split(value, prefix.filter, ...) # type: ignore[misc]\n nondiff_states.append(extract.GraphDefState(graphdef, broadcast))\n return extract.NodeStates.from_states(passed)\n\n\n nondiff_argnums: tuple[int, ...] = struct.field(pytree_node=False)\n tangent_tree_node_args: tuple[tp.Any, ...] = struct.field(pytree_node=False)\n\ndef _extract_nodedefs(x, *, nodedefs: deque[graphlib.GraphDef]):\n if isinstance(x, graphlib.GraphDef):\n nodedefs.append(x)\n return x.with_no_outer_index()\n return x\n\n@dataclasses.dataclass(eq=False)\nclass CustomVjpFnWrapper:\n f: tp.Callable[..., tp.Any]\n jax_nondiff_argnums: tuple[int, ...]\n ctxtag: str\n nondiff_states: list[extract.GraphDefState]\n nodedefs: deque[graphlib.GraphDef]\n\n def __post_init__(self):\n functools.update_wrapper(self, self.f)\n\n def __call__(self, *pure_args):\n nondiff_states = deque(self.nondiff_states)\n args = extract.from_tree(\n pure_args,\n merge_fn=functools.partial(\n _custom_vjp_merge_fn, nondiff_states=nondiff_states\n ),\n ctxtag=self.ctxtag,\n is_inner=True,\n )\n\n out = self.f(*args)\n\n # remove nondiff from pure_args_out_g\n args_out = tuple(\n x for i, x in enumerate(args) if i not in self.jax_nondiff_argnums\n )\n args_out = extract.clear_non_graph_nodes(args_out)\n pure_args_out, pure_out = extract.to_tree(\n (args_out, out), ctxtag=self.ctxtag\n )\n # remove outer_index from GraphDef's but store them in global context\n\n pure_args_out, pure_out = jax.tree.map(\n functools.partial(_extract_nodedefs, nodedefs=self.nodedefs),\n (pure_args_out, pure_out),\n is_leaf=lambda x: isinstance(x, graphlib.GraphDef),\n )\n\n return pure_args_out, pure_out\n\n\n@dataclasses.dataclass(eq=False)\nclass FwdFn:\n fwd: tp.Callable[..., tp.Any]\n nondiff_argnums: tuple[int, ...]\n ctxtag: str\n nondiff_states: list[extract.GraphDefState]\n nodedefs: deque[graphlib.GraphDef]\n\n def __post_init__(self):\n functools.update_wrapper(self, self.fwd)\n\n def __call__(self, *pure_args):\n # here we need to be aware if the update_context is active or not\n # when its not active, index_mappings will be None\n # when its active, we will remove the index_mappings from the GraphDef's and store them\n # in the index_mappings deque created by CustomVjp\n update_context_active = (\n self.ctxtag in graphlib.GRAPH_CONTEXT.update_context_stacks\n )\n nondiff_states = deque(self.nondiff_states)\n args = extract.from_tree(\n pure_args,\n merge_fn=functools.partial(\n _custom_vjp_merge_fn, nondiff_states=nondiff_states\n ),\n ctxtag=self.ctxtag if update_context_active else None,\n is_inner=True,\n )\n\n out, residual = self.fwd(*args)\n\n # remove nondiff from pure_args_out_g\n args_out = tuple(\n x for i, x in enumerate(args) if i not in self.nondiff_argnums\n )\n args_out = extract.clear_non_graph_nodes(args_out)\n pure_args_out, pure_out = extract.to_tree(\n (args_out, out),\n ctxtag=self.ctxtag if update_context_active else None,\n )\n pure_residual = extract.to_tree(residual)\n\n if update_context_active:\n # remove outer_index from GraphDef's but store them in global context\n pure_args_out, pure_out = jax.tree.map(\n functools.partial(_extract_nodedefs, nodedefs=self.nodedefs),\n (pure_args_out, pure_out),\n is_leaf=lambda x: isinstance(x, graphlib.GraphDef),\n )\n\n return (pure_args_out, pure_out), pure_residual\n\n\n@dataclasses.dataclass(eq=False)\nclass BwdFn:\n bwd: tp.Callable[..., tp.Any]\n tree_node_args: tuple[tp.Any, ...]\n\n def __post_init__(self):\n functools.update_wrapper(self, self.bwd)\n\n def __call__(self, *args):\n *nondiff, pure_residual, (pure_args_out_g, pure_out_g) = args\n residual = extract.from_tree(pure_residual, is_inner=True)\n (pure_args_out_g, pure_out_g) = jax.tree.map(\n lambda x: x.state if isinstance(x, extract.NodeStates) else x,\n (pure_args_out_g, pure_out_g),\n is_leaf=lambda x: isinstance(x, extract.NodeStates),\n )\n\n tangent = self.bwd(*nondiff, residual, (pure_args_out_g, pure_out_g))\n\n def state_to_node_states(is_differentiable: bool, x):\n if is_differentiable:\n if isinstance(x, jax.Array):\n return x\n elif not isinstance(x, State | variablelib.Variable):\n raise ValueError(f'Expected State or Variable, got {type(x)}')\n return extract.NodeStates.from_states(x)\n return x\n\n pure_tangent = jax.tree.map(\n state_to_node_states,\n self.tree_node_args,\n tangent,\n is_leaf=lambda x: isinstance(x, State | variablelib.Variable),\n )\n return pure_tangent\n\n\nclass CustomVjp(tp.Generic[A]):\n def __init__(\n self,\n fun: tp.Callable[..., A],\n nondiff_argnums: tuple[int | DiffState, ...],\n ):\n functools.update_wrapper(self, fun)\n # first argument is metadata\n self.jax_nondiff_argnums = tuple(\n x for x in nondiff_argnums if isinstance(x, int)\n )\n self.ctxtag = f'custom_vjp_{fun.__name__}_{id(fun)}'\n self.fun = fun\n self.fwd: tp.Callable | None = None\n self.bwd: tp.Callable | None = None\n self.symbolic_zeros: bool | None = None\n self.nondiff_argnums = nondiff_argnums\n self.diff_filter: dict[int, tp.Literal[False] | DiffState] = {}\n for argnum in self.nondiff_argnums:\n index = argnum.argnum if isinstance(argnum, DiffState) else argnum\n if index in self.diff_filter:\n raise ValueError(f'argnum {index} is repeated in nondiff_argnums')\n self.diff_filter[index] = (\n dataclasses.replace(argnum, argnum=-1)\n if isinstance(argnum, DiffState)\n else False\n )\n\n # def __getattr__(self, name: str) -> tp.Any:\n # if not hasattr(self.custom_vjp_fn, name):\n # raise AttributeError(f'{type(self).__name__} has no attribute {name}')\n # return getattr(self.custom_vjp_fn, name)\n\n def __call__(\n self, *args: tp.Any, **kwargs: tp.Any\n ) -> A: # pytype: disable=invalid-annotation\n with graphlib.update_context(self.ctxtag):\n args = resolve_kwargs(self.fun, args, kwargs)\n del kwargs\n nondiff_states: list[extract.GraphDefState] = []\n arg_filters = tuple(\n self.diff_filter.get(i, True) for i in range(len(args))\n )\n pure_args = extract.to_tree(\n args,\n prefix=arg_filters,\n split_fn=functools.partial(\n _custom_vjp_split_fn, nondiff_states=nondiff_states\n ),\n ctxtag=self.ctxtag,\n )\n tree_node_args = jax.tree.map(\n lambda x: isinstance(x, extract.NodeStates),\n pure_args,\n is_leaf=lambda x: isinstance(x, extract.NodeStates),\n )\n tree_node_args = tuple(\n x\n for i, x in enumerate(tree_node_args)\n if i not in self.jax_nondiff_argnums\n )\n nodedefs: deque[graphlib.GraphDef] = deque()\n if self.fwd is None or self.bwd is None or self.symbolic_zeros is None:\n raise ValueError()\n\n custom_vjp_fn = jax.custom_vjp(\n fun=CustomVjpFnWrapper(\n f=self.fun,\n jax_nondiff_argnums=self.jax_nondiff_argnums,\n ctxtag=self.ctxtag,\n nondiff_states=nondiff_states,\n nodedefs=nodedefs,\n ),\n nondiff_argnums=self.jax_nondiff_argnums,\n )\n custom_vjp_fn.defvjp(\n fwd=FwdFn(\n fwd=self.fwd,\n nondiff_argnums=self.jax_nondiff_argnums,\n ctxtag=self.ctxtag,\n nondiff_states=nondiff_states,\n nodedefs=nodedefs,\n ),\n bwd=BwdFn(\n bwd=self.bwd,\n tree_node_args=tree_node_args,\n ),\n symbolic_zeros=self.symbolic_zeros,\n )\n pure_args_out, pure_out = custom_vjp_fn(*pure_args)\n\n # insert index_mappings\n def _insert_index_mappings(x):\n if isinstance(x, graphlib.GraphDef):\n nodedef = nodedefs.popleft()\n return nodedef\n return x\n\n pure_args_out, pure_out = jax.tree_util.tree_map(\n _insert_index_mappings,\n (pure_args_out, pure_out),\n is_leaf=lambda x: isinstance(x, graphlib.GraphDef),\n )\n\n args_out, out = extract.from_tree(\n (pure_args_out, pure_out), ctxtag=self.ctxtag, is_inner=False\n )\n\n return out\n\n def defvjp(\n self,\n fwd: tp.Callable[..., tuple[A, tp.Any]],\n bwd: tp.Callable[..., tuple[tp.Any, ...]],\n symbolic_zeros: bool = False,\n ) -> None:\n self.fwd = fwd\n self.bwd = bwd\n self.symbolic_zeros = symbolic_zeros\n\n\n@tp.overload\ndef custom_vjp(\n fun: tp.Callable[..., A],\n *,\n nondiff_argnums: tuple[int | DiffState, ...] = (),\n graph: bool | None = None,\n graph_updates: bool | None = None,\n) -> CustomVjp[A] | SimpleCustomVjp[A]: ...\n@tp.overload\ndef custom_vjp(\n *,\n nondiff_argnums: tuple[int | DiffState, ...] = (),\n graph: bool | None = None,\n graph_updates: bool | None = None,\n) -> tp.Callable[[tp.Callable[..., A]], CustomVjp[A] | SimpleCustomVjp[A]]: ...\ndef custom_vjp(\n fun: tp.Callable[..., A] | Missing = MISSING,\n *,\n nondiff_argnums: tuple[int | DiffState, ...] = (),\n graph: bool | None = None,\n graph_updates: bool | None = None,\n) -> CustomVjp[A] | SimpleCustomVjp[A] | tp.Callable[[tp.Callable[..., A]], CustomVjp[A] | SimpleCustomVjp[A]]:\n \"\"\"Reference aware version of\n `jax.custom_vjp `__.\n\n ``nnx.custom_vjp`` accepts Modules and other Flax NNX objects as arguments. The main difference\n with the JAX version is that, because Modules follow reference semantics, they propagate the State\n updates for the inputs as auxiliary outputs. This means that the incoming gradients in the ``bwd`` function\n will have the form ``(input_updates_g, out_g)`` where ``input_updates_g`` is the gradient updated state of\n the inputs w.r.t. to the inputs. All Module terms on the inputs will an associated ``State`` term in\n ``input_updates_g``, while all non-Module terms will appear as None. The shape of the tangent will be\n expected to have the same shape as the input, with ``State`` terms in place of the corresponding Module terms.\n\n Example::\n\n >>> import jax\n >>> import jax.numpy as jnp\n >>> from flax import nnx\n ...\n >>> class Foo(nnx.Module):\n ... def __init__(self, x, y):\n ... self.x = nnx.Param(x)\n ... self.y = nnx.Param(y)\n ...\n >>> @nnx.custom_vjp\n ... def f(m: Foo):\n ... return jnp.sin(m.x) * m.y\n ...\n >>> def f_fwd(m: Foo):\n ... return f(m), (jnp.cos(m.x), jnp.sin(m.x), m)\n ...\n >>> def f_bwd(res, g):\n ... input_updates_g, out_g = g\n ... cos_x, sin_x, m = res\n ... (m_updates_g,) = input_updates_g\n ... m_g = jax.tree.map(lambda x: x, m_updates_g) # create copy\n ...\n ... m_g['x'][...] = cos_x * out_g * m.y\n ... m_g['y'][...] = sin_x * out_g\n ... return (m_g,)\n ...\n >>> f.defvjp(f_fwd, f_bwd)\n ...\n >>> m = Foo(x=jnp.array(1.), y=jnp.array(2.))\n >>> grads = nnx.grad(f)(m)\n ...\n >>> jax.tree.map(jnp.shape, grads)\n State({\n 'x': Param(\n value=()\n ),\n 'y': Param(\n value=()\n )\n })\n\n Note that the State objects that represent Module terms on ``input_updates_g`` have the\n same shape as the State objects expected in the output tanget. This means that you can\n usually just copy them from ``input_updates_g`` and update them with their corresponding\n gradient values.\n\n You can select which substates are differentiable (have a tangent) for Modules and other\n graph nodes by passing a ``DiffState`` to ``nondiff_argnums``. For example, if you want to\n differentiate only the ``x`` attribute of the ``Foo`` class, you can do the following::\n\n >>> x_attribute = nnx.PathContains('x')\n >>> diff_state = nnx.DiffState(0, x_attribute)\n ...\n >>> @nnx.custom_vjp(nondiff_argnums=(diff_state,))\n ... def f(m: Foo):\n ... return jnp.sin(m.x) * m.y # type: ignore\n\n >>> def f_fwd(m: Foo):\n ... y = f(m)\n ... res = (jnp.cos(m.x), m) # type: ignore\n ... return y, res\n ...\n >>> def f_bwd(res, g):\n ... input_updates_g, out_g = g\n ... cos_x, m = res\n ... (m_updates_g,) = input_updates_g\n ... m_g = jax.tree.map(lambda x: x, m_updates_g) # create copy\n ...\n ... m_g.x[...] = cos_x * out_g * m.y\n ... del m_g['y'] # y is not differentiable\n ... return (m_g,)\n\n >>> f.defvjp(f_fwd, f_bwd)\n ...\n >>> m = Foo(x=jnp.array(1.), y=jnp.array(2.))\n >>> grad = nnx.grad(f, argnums=nnx.DiffState(0, x_attribute))(m)\n ...\n >>> jax.tree.map(jnp.shape, grad)\n State({\n 'x': Param(\n value=()\n )\n })\n\n Note that ``grad`` cannot calculate gradients for states that don't have a tangent\n defined by ``custom_vjp``, in the example above we reuse the same ``x_attribute``\n filter to keep ``custom_vjp`` and ``grad`` in sync.\n\n **graph_updates=False**\n\n When ``graph_updates=False`` or ``graph=False``, the behavior is closer to\n ``jax.custom_vjp``: the ``bwd`` function receives ``out_g`` directly, and\n tangent types are the same as the input types, this means the tangent for a\n Module is a Module instance with gradient values set on its attributes.\n This mode does not support ``DiffState`` in ``nondiff_argnums``. Additionally,\n Variables in differentiable arguments cannot be mutated inside ``f``. If\n mutations are needed, pass the relevant Variables through a non-differentiable\n argument instead.\n\n Example::\n\n >>> @nnx.custom_vjp(graph_updates=False)\n ... def f(m: Foo):\n ... return jnp.sin(m.x) * m.y\n ...\n >>> def f_fwd(m: Foo):\n ... return f(m), (jnp.cos(m.x), jnp.sin(m.x), m)\n ...\n >>> def f_bwd(res, g):\n ... cos_x, sin_x, m = res\n ... m_g = nnx.clone(m)\n ... m_g.x[...] = cos_x * g * m.y\n ... m_g.y[...] = sin_x * g\n ... return (m_g,)\n ...\n >>> f.defvjp(f_fwd, f_bwd)\n\n Args:\n fun: Callable base function.\n nondiff_argnums: Tuple of integers or DiffState objects specifying the\n argument indices that are not differentiated. By default all arguments are\n differentiated. Integers cannot be used to mark graph nodes such as Modules\n as non-differentiable, in this case use a DiffState object. DiffState objects\n define the set of differentiable substates, contrary to what the name of this\n argument suggests, this is done for compatibility with ``grad``.\n graph: If ``True`` (default), uses graph-mode which supports the full\n NNX feature set including shared references and reference semantics.\n If ``False``, uses tree-mode which treats Modules as regular JAX\n pytrees, avoiding the overhead of the graph protocol. Tree-mode does\n not support ``DiffState`` in ``nondiff_argnums``.\n graph_updates: If ``True``, propagates updates on graph structure\n that happen inside the transform to the input graphs, has no\n effect when ``graph=False``. When ``False``, using ``DiffState``\n is not supported.\n\n \"\"\"\n if graph is None:\n graph = graphlib.set_graph_mode.current_value()\n if graph_updates is None:\n graph_updates = graphlib.set_graph_updates.current_value()\n if isinstance(fun, Missing):\n return functools.partial(\n custom_vjp, nondiff_argnums=nondiff_argnums, graph=graph,\n graph_updates=graph_updates,\n )\n\n # Detect bound nnx.Module methods and raise error.\n fun_unbound, _, was_bound = _resolve_bound_callable(fun)\n if was_bound:\n _raise_bound_method_error('custom_vjp')\n\n if not graph or not graph_updates:\n if any(isinstance(x, DiffState) for x in nondiff_argnums):\n raise ValueError(\n '`nondiff_argnums` cannot contain `DiffState` objects '\n 'when `graph=False`. '\n + graphlib._tree_mode_suggestion_transform('custom_vjp')\n )\n return SimpleCustomVjp(fun_unbound, nondiff_argnums, graph=graph) # type: ignore[arg-type]\n\n return CustomVjp(fun_unbound, nondiff_argnums)\n\n\n# -------------------------------\n# remat\n# -------------------------------\n\n\n@dataclasses.dataclass(eq=False)\nclass SimpleRematFn:\n f: tp.Callable[..., tp.Any]\n graph: bool\n\n def __post_init__(self):\n functools.update_wrapper(self, self.f, updated=())\n\n @extract.treemap_copy_args\n def __call__(self, *args, **kwargs):\n updates, snapshot = extract.updates_and_snapshot((args, kwargs))\n if self.graph:\n args, kwargs = extract.from_tree2((args, kwargs))\n out = self.f(*args, **kwargs)\n if self.graph:\n out = extract.to_tree2(out)\n extract.check_no_aliases('remat', args=updates[0], kwargs=updates[1], out=out)\n updates = extract.mask_variable_updates(updates, snapshot)\n return out, updates\n\n@tp.overload\ndef remat(\n *,\n prevent_cse: bool = True,\n static_argnums: int | tuple[int, ...] = (),\n policy: tp.Callable[..., bool] | None = None,\n graph: bool | None = None,\n graph_updates: bool | None = None,\n) -> tp.Callable[[F], F]: ...\n@tp.overload\ndef remat(\n f: F,\n *,\n prevent_cse: bool = True,\n static_argnums: int | tuple[int, ...] = (),\n policy: tp.Callable[..., bool] | None = None,\n graph: bool | None = None,\n graph_updates: bool | None = None,\n) -> F: ...\ndef remat(\n f: F | Missing = MISSING,\n *,\n prevent_cse: bool = True,\n static_argnums: int | tuple[int, ...] = (),\n policy: tp.Callable[..., bool] | None = None,\n graph: bool | None = None,\n graph_updates: bool | None = None,\n) -> F | tp.Callable[[F], F]:\n \"\"\"A 'lifted' version of the\n `jax.checkpoint `__\n (a.k.a. ``jax.remat``).\n\n ``flax.nnx.remat``, similar to ``jax.checkpoint`` can provide control over, for\n example, how ``flax.nnx.grad`` values are computed and saved during the forward pass versus\n how they are recomputed during the backward pass, trading off memory and FLOPs.\n\n Learn more in `Flax NNX vs JAX Transformations `_.\n\n To learn about ``jax.remat``, go to JAX's\n `fundamentals of jax.checkpoint `_\n and `practical notes `_.\n\n Args:\n f: Function to be rematerialized.\n prevent_cse: Optional, bool. If True, prevents common subexpression\n elimination. Default True.\n static_argnums: Optional, int or tuple of ints. Specifies which\n positional arguments to treat as static.\n policy: Optional, callable. A policy for which intermediates to save\n during the forward pass.\n graph: If ``True`` (default), uses graph-mode which supports the full\n NNX feature set including shared references and reference semantics.\n If ``False``, uses tree-mode which treats Modules as regular JAX\n pytrees, avoiding the overhead of the graph protocol. Tree-mode does\n not support shared ``Variable`` references.\n graph_updates: If ``True``, propagates updates on graph structure\n that happen inside the transform to the input graphs, has no\n effect when ``graph=False``.\n \"\"\"\n if graph is None:\n graph = graphlib.set_graph_mode.current_value()\n if graph_updates is None:\n graph_updates = graphlib.set_graph_updates.current_value()\n if isinstance(f, Missing):\n return functools.partial(\n remat,\n prevent_cse=prevent_cse,\n static_argnums=static_argnums,\n policy=policy,\n graph=graph,\n graph_updates=graph_updates,\n ) # type: ignore[return-value]\n\n f_unbound, _, was_bound = _resolve_bound_callable(f)\n\n if was_bound:\n _raise_bound_method_error('remat')\n\n if not graph or not graph_updates:\n checkpointed_fn = jax.checkpoint(\n SimpleRematFn(f_unbound, graph=graph),\n prevent_cse=prevent_cse,\n static_argnums=static_argnums,\n policy=policy,\n )\n\n @functools.wraps(f_unbound)\n def simple_remat_wrapper(*args, **kwargs):\n if graph:\n args, kwargs = extract.to_tree2((args, kwargs))\n extract.check_no_aliases('remat', args=args, kwargs=kwargs)\n out, updates = checkpointed_fn(*args, **kwargs)\n if graph:\n out = extract.from_tree2(out)\n extract.apply_variable_updates((args, kwargs), updates)\n return out\n\n return simple_remat_wrapper # type: ignore[return-value]\n\n # Unbound function path: preserve the concise composition used in NNX.\n return resolve_kwargs()( # type: ignore[return-value]\n graphlib.update_context('remat')(\n general.split_inputs(\n jax.checkpoint(\n general.merge_inputs(f_unbound, ctxtag='remat'),\n prevent_cse=prevent_cse,\n static_argnums=static_argnums,\n policy=policy,\n ),\n ctxtag='remat',\n ),\n )\n )\n" }, { "path": "flax/nnx/transforms/compilation.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# pytype: skip-file\nfrom __future__ import annotations\n\nimport dataclasses\nimport functools\nimport inspect\nimport operator\nimport typing as tp\n\nimport jax\nfrom jax.sharding import AbstractMesh, Mesh, PartitionSpec\n\nfrom flax.nnx import (\n extract,\n filterlib,\n graphlib,\n statelib,\n variablelib,\n)\nfrom flax.nnx.transforms.transforms import (\n _resolve_bound_callable,\n _raise_bound_method_error,\n)\nfrom flax.typing import MISSING, Missing, PathParts\n\nF = tp.TypeVar('F', bound=tp.Callable[..., tp.Any])\nP = tp.ParamSpec('P')\nR = tp.TypeVar('R')\nSpecs = tp.Any\nAxisName = tp.Hashable\n\n# -------------------------------\n# jit\n# -------------------------------\n\n\nclass StateSharding(extract.PrefixMapping):\n def __init__(\n self,\n filter_sharding: statelib.State\n | tp.Mapping[filterlib.Filter, tp.Any]\n | tp.Iterable[tuple[filterlib.Filter, tp.Any]],\n /,\n ):\n if isinstance(filter_sharding, statelib.State):\n filter_sharding = statelib.create_path_filters(filter_sharding) # type: ignore\n\n iterable = tuple(\n filter_sharding.items()\n if isinstance(filter_sharding, tp.Mapping)\n else filter_sharding\n )\n self._filters = tuple(filter for filter, _ in iterable)\n self._shardings = tuple(axis for _, axis in iterable)\n\n @property\n def filters(self) -> tuple[filterlib.Filter, ...]:\n return self._filters\n\n @property\n def shardings(self) -> tuple[tp.Any, ...]:\n return self._shardings\n\n def map_prefix(\n self, path: PathParts, variable: variablelib.Variable\n ) -> tp.Any:\n for filter, sharding in zip(self.filters, self.shardings):\n predicate = filterlib.to_predicate(filter)\n if predicate(path, variable):\n return sharding\n raise ValueError(f'No axis found for {path=}, {variable=}')\n\n def __repr__(self):\n return f'StateSharding({dict(zip(self.filters, self.shardings))})'\n\n def __eq__(self, other):\n return (\n isinstance(other, StateSharding)\n and self.filters == other.filters\n and self.shardings == other.shardings\n )\n\n def __hash__(self):\n return hash((self.filters, self.shardings))\n\n\ndef _jit_split_fn(ctx: graphlib.SplitContext, path, prefix, x):\n if isinstance(prefix, StateSharding):\n graphdef, *states = ctx.flatten(x, *prefix.filters)\n return extract.NodeStates.from_split(graphdef, *states, metadata=prefix)\n return extract.NodeStates.from_split(*ctx.flatten(x, with_paths=False))\n\n\ndef _jit_merge_fn(ctx: graphlib.MergeContext, path, prefix, leaf) -> tp.Any:\n if not isinstance(leaf, extract.NodeStates):\n raise ValueError(f'Expected TreeNode, got {type(leaf)} at path {path}')\n return ctx.unflatten(leaf.graphdef, *leaf.states)\n\n\n@dataclasses.dataclass(eq=False)\nclass JitFn:\n f: tp.Callable[..., tp.Any]\n in_shardings: tp.Any\n out_shardings: tp.Any\n kwarg_shardings: tp.Any\n ctxtag: tp.Hashable\n\n def __post_init__(self):\n # Prevent overwriting our ctxtag info with the child function's\n orig_ctxtag = self.ctxtag\n functools.update_wrapper(self, self.f, updated=())\n self.ctxtag = orig_ctxtag\n\n def __call__(self, *pure_args, **pure_kwargs):\n args, kwargs = extract.from_tree(\n (pure_args, pure_kwargs),\n merge_fn=_jit_merge_fn,\n ctxtag=self.ctxtag,\n is_inner=True,\n )\n\n out = self.f(*args, **kwargs)\n\n args_out, kwargs_out = extract.clear_non_graph_nodes((args, kwargs))\n pure_args_out, pure_kwargs_out, pure_out = extract.to_tree(\n (args_out, kwargs_out, out),\n prefix=(self.in_shardings, self.kwarg_shardings, self.out_shardings),\n ctxtag=self.ctxtag,\n split_fn=_jit_split_fn,\n )\n\n return pure_args_out, pure_kwargs_out, pure_out\n\n\n@tp.overload\ndef jit(\n *,\n in_shardings: tp.Any = None,\n out_shardings: tp.Any = None,\n static_argnums: int | tp.Sequence[int] | None = None,\n static_argnames: str | tp.Iterable[str] | None = None,\n donate_argnums: int | tp.Sequence[int] | None = None,\n donate_argnames: str | tp.Iterable[str] | None = None,\n keep_unused: bool = False,\n device: tp.Optional[jax.Device] = None,\n backend: tp.Optional[str] = None,\n inline: bool = False,\n graph: bool | None = None,\n graph_updates: bool | None = None,\n) -> tp.Callable[[tp.Callable[P, R]], JitWrapped[P, R]]: ...\n@tp.overload\ndef jit(\n fun: tp.Callable[P, R],\n *,\n in_shardings: tp.Any = None,\n out_shardings: tp.Any = None,\n static_argnums: int | tp.Sequence[int] | None = None,\n static_argnames: str | tp.Iterable[str] | None = None,\n donate_argnums: int | tp.Sequence[int] | None = None,\n donate_argnames: str | tp.Iterable[str] | None = None,\n keep_unused: bool = False,\n device: tp.Optional[jax.Device] = None,\n backend: tp.Optional[str] = None,\n inline: bool = False,\n graph: bool | None = None,\n graph_updates: bool | None = None,\n) -> JitWrapped[P, R]: ...\ndef jit(\n fun: tp.Callable[P, R] | Missing = MISSING,\n *,\n in_shardings: tp.Any = None,\n out_shardings: tp.Any = None,\n static_argnums: int | tp.Sequence[int] | None = None,\n static_argnames: str | tp.Iterable[str] | None = None,\n donate_argnums: int | tp.Sequence[int] | None = None,\n donate_argnames: str | tp.Iterable[str] | None = None,\n keep_unused: bool = False,\n device: tp.Optional[jax.Device] = None,\n backend: tp.Optional[str] = None,\n inline: bool = False,\n graph: bool | None = None,\n graph_updates: bool | None = None,\n) -> JitWrapped[P, R] | tp.Callable[[tp.Callable[P, R]], JitWrapped[P, R]]:\n \"\"\"\n Lifted version of ``jax.jit`` that can handle Modules / graph nodes as\n arguments.\n\n .. note::\n If jitted function has a model and an optimizer as inputs, we can\n reduce accelerator's memory usage if we specify them in\n ``donate_argnums`` or ``donate_argnames``:\n\n >>> from flax import nnx\n >>>\n >>> @nnx.jit(donate_argnames=(\"model\", \"optimizer\"))\n ... def func(model: nnx.Module, optimizer: nnx.Optimizer, other_args):\n ... pass\n\n For details please see `this discussion `_.\n\n Args:\n fun: Function to be jitted. ``fun`` should be a pure function, as\n side-effects may only be executed once.\n\n The arguments and return value of ``fun`` should be arrays,\n scalars, or (nested) standard Python containers (tuple/list/dict) thereof.\n Positional arguments indicated by ``static_argnums`` can be anything at\n all, provided they are hashable and have an equality operation defined.\n Static arguments are included as part of a compilation cache key, which is\n why hash and equality operators must be defined.\n\n JAX keeps a weak reference to ``fun`` for use as a compilation cache key,\n so the object ``fun`` must be weakly-referenceable. Most :class:`Callable`\n objects will already satisfy this requirement.\n\n .. note::\n Bound methods (e.g., ``module.method``) are not supported. Use the\n decorator form ``@nnx.jit`` on the method definition or call\n ``nnx.jit(MyClass.method)(instance, ...)`` with the unbound method.\n in_shardings: Pytree of structure matching that of arguments to ``fun``,\n with all actual arguments replaced by resource assignment specifications.\n It is also valid to specify a pytree prefix (e.g. one value in place of a\n whole subtree), in which case the leaves get broadcast to all values in\n that subtree.\n\n The ``in_shardings`` argument is optional. JAX will infer the shardings\n from the input :py:class:`jax.Array`'s and defaults to replicating the input\n if the sharding cannot be inferred.\n\n The valid resource assignment specifications are:\n - :py:class:`Sharding`, which will decide how the value\n will be partitioned. With this, using a mesh context manager is not\n required.\n - :py:obj:`None`, will give JAX the freedom to choose whatever sharding\n it wants.\n For in_shardings, JAX will mark is as replicated but this behavior\n can change in the future.\n For out_shardings, we will rely on the XLA GSPMD partitioner to\n determine the output shardings.\n\n The size of every dimension has to be a multiple of the total number of\n resources assigned to it. This is similar to pjit's in_shardings.\n out_shardings: Like ``in_shardings``, but specifies resource\n assignment for function outputs. This is similar to pjit's\n out_shardings.\n\n The ``out_shardings`` argument is optional. If not specified, :py:func:`jax.jit`\n will use GSPMD's sharding propagation to figure out what the sharding of the\n output(s) should be.\n static_argnums: An optional int or collection of ints that specify which\n positional arguments to treat as static (compile-time constant).\n Operations that only depend on static arguments will be constant-folded in\n Python (during tracing), and so the corresponding argument values can be\n any Python object.\n\n Static arguments should be hashable, meaning both ``__hash__`` and\n ``__eq__`` are implemented, and immutable. Calling the jitted function\n with different values for these constants will trigger recompilation.\n Arguments that are not arrays or containers thereof must be marked as\n static.\n\n If neither ``static_argnums`` nor ``static_argnames`` is provided, no\n arguments are treated as static. If ``static_argnums`` is not provided but\n ``static_argnames`` is, or vice versa, JAX uses\n :code:`inspect.signature(fun)` to find any positional arguments that\n correspond to ``static_argnames``\n (or vice versa). If both ``static_argnums`` and ``static_argnames`` are\n provided, ``inspect.signature`` is not used, and only actual\n parameters listed in either ``static_argnums`` or ``static_argnames`` will\n be treated as static.\n static_argnames: An optional string or collection of strings specifying\n which named arguments to treat as static (compile-time constant). See the\n comment on ``static_argnums`` for details. If not\n provided but ``static_argnums`` is set, the default is based on calling\n ``inspect.signature(fun)`` to find corresponding named arguments.\n donate_argnums: Specify which positional argument buffers are \"donated\" to\n the computation. It is safe to donate argument buffers if you no longer\n need them once the computation has finished. In some cases XLA can make\n use of donated buffers to reduce the amount of memory needed to perform a\n computation, for example recycling one of your input buffers to store a\n result. You should not reuse buffers that you donate to a computation, JAX\n will raise an error if you try to. By default, no argument buffers are\n donated.\n\n If neither ``donate_argnums`` nor ``donate_argnames`` is provided, no\n arguments are donated. If ``donate_argnums`` is not provided but\n ``donate_argnames`` is, or vice versa, JAX uses\n :code:`inspect.signature(fun)` to find any positional arguments that\n correspond to ``donate_argnames``\n (or vice versa). If both ``donate_argnums`` and ``donate_argnames`` are\n provided, ``inspect.signature`` is not used, and only actual\n parameters listed in either ``donate_argnums`` or ``donate_argnames`` will\n be donated.\n\n For more details on buffer donation see the\n `FAQ `_.\n donate_argnames: An optional string or collection of strings specifying\n which named arguments are donated to the computation. See the\n comment on ``donate_argnums`` for details. If not\n provided but ``donate_argnums`` is set, the default is based on calling\n ``inspect.signature(fun)`` to find corresponding named arguments.\n keep_unused: If `False` (the default), arguments that JAX determines to be\n unused by `fun` *may* be dropped from resulting compiled XLA executables.\n Such arguments will not be transferred to the device nor provided to the\n underlying executable. If `True`, unused arguments will not be pruned.\n device: This is an experimental feature and the API is likely to change.\n Optional, the Device the jitted function will run on. (Available devices\n can be retrieved via :py:func:`jax.devices`.) The default is inherited\n from XLA's DeviceAssignment logic and is usually to use\n ``jax.devices()[0]``.\n backend: This is an experimental feature and the API is likely to change.\n Optional, a string representing the XLA backend: ``'cpu'``, ``'gpu'``, or\n ``'tpu'``.\n inline: Specify whether this function should be inlined into enclosing\n jaxprs (rather than being represented as an application of the xla_call\n primitive with its own subjaxpr). Default False.\n graph: If ``True`` (default), uses graph-mode which supports the full\n NNX feature set including shared references, reference semantics, and\n structural changes to Modules inside the jitted function. If ``False``,\n uses tree-mode which treats Modules as regular JAX pytrees, avoiding\n the overhead of the graph protocol. Tree-mode is faster but does not\n support shared ``Variable`` references or returning mutable array\n references from the jitted function.\n\n Returns:\n A wrapped version of ``fun``, set up for just-in-time compilation.\n \"\"\"\n\n if graph is None:\n graph = graphlib.set_graph_mode.current_value()\n if graph_updates is None:\n graph_updates = graphlib.set_graph_updates.current_value()\n if isinstance(fun, Missing):\n return functools.partial(\n jit,\n in_shardings=in_shardings,\n out_shardings=out_shardings,\n static_argnums=static_argnums,\n static_argnames=static_argnames,\n donate_argnums=donate_argnums,\n donate_argnames=donate_argnames,\n keep_unused=keep_unused,\n device=device,\n backend=backend,\n inline=inline,\n graph=graph,\n graph_updates=graph_updates,\n ) # type: ignore[return-value]\n fun_unbound, _, was_bound = _resolve_bound_callable(fun)\n if was_bound:\n _raise_bound_method_error('jit')\n\n if not graph:\n if any(isinstance(x, StateSharding) for x in jax.tree.leaves(in_shardings)):\n raise ValueError(\n '`in_shardings` cannot contain `StateSharding` objects '\n 'when `graph=False`. '\n + graphlib._tree_mode_suggestion_transform('jit')\n )\n if any(isinstance(x, StateSharding) for x in jax.tree.leaves(out_shardings)):\n raise ValueError(\n '`out_shardings` cannot contain `StateSharding` objects '\n 'when `graph=False`. '\n + graphlib._tree_mode_suggestion_transform('jit')\n )\n\n if graph and not graph_updates:\n if in_shardings is not None:\n extract.check_prefix(in_shardings, 'in_shardings', 'jit')\n if out_shardings is not None:\n extract.check_prefix(out_shardings, 'out_shardings', 'jit')\n\n wrapped_cls: tp.Any\n if graph and graph_updates:\n wrapped_cls = JitWrapped\n else:\n wrapped_cls = functools.partial(SimpleJitWrapped, graph=graph)\n return wrapped_cls(\n fun_unbound,\n in_shardings=in_shardings,\n out_shardings=out_shardings,\n static_argnums=static_argnums,\n static_argnames=static_argnames,\n donate_argnums=donate_argnums,\n donate_argnames=donate_argnames,\n keep_unused=keep_unused,\n device=device,\n backend=backend,\n inline=inline,\n )\n\n\n@dataclasses.dataclass(frozen=True, slots=True)\nclass PartialState:\n \"\"\"Container for a pre-flattened partial argument.\n\n Stores the pytree structure (``treedef``) as static metadata and the\n flattened leaves as dynamic data. Variables within the original argument\n are kept as leaves so their values can change between calls without\n triggering recompilation.\n \"\"\"\n treedef: jax.tree_util.PyTreeDef\n leaves: list[tp.Any]\n\njax.tree_util.register_dataclass(\n PartialState,\n data_fields=['leaves'],\n meta_fields=['treedef'],\n)\n\n\ndef _flatten_to_partial_state(\n arg: tp.Any,\n ref_index: graphlib.RefMap | None,\n) -> PartialState:\n if ref_index is not None:\n graphdef, flat_state = graphlib.flatten(arg, ref_index=ref_index, graph=True)\n return PartialState(treedef=graphdef, leaves=flat_state.leaves)\n is_leaf = lambda x: isinstance(x, variablelib.Variable)\n leaves, treedef = jax.tree.flatten(arg, is_leaf=is_leaf)\n return PartialState(treedef=treedef, leaves=leaves)\n\n\n\n\n@dataclasses.dataclass(eq=False)\nclass SimpleJitFn:\n f: tp.Callable[..., tp.Any]\n out_shardings: tp.Any\n donate_argnums: frozenset[int]\n donate_argnames: frozenset[str]\n graph: bool\n\n def __post_init__(self):\n functools.update_wrapper(self, self.f, updated=())\n\n @extract.treemap_copy_args\n def __call__(self, *args, **kwargs):\n updates, snapshot = extract.updates_and_snapshot((args, kwargs))\n args_updates, kwargs_updates = updates\n args_snapshot, kwargs_snapshot = snapshot\n if self.graph:\n args, kwargs = extract.from_tree2((args, kwargs))\n out = self.f(*args, **kwargs)\n if self.graph:\n out = extract.to_tree2(out, prefix=self.out_shardings)\n extract.check_no_aliases('jit', args=args_updates, kwargs=kwargs_updates, out=out)\n def donated_arg(jax_path, prefix, c, s):\n path = graphlib.jax_to_nnx_path(jax_path)\n return path[0] in self.donate_argnums or extract.variable_changed(c, s)\n args_updates = extract.mask_variable_updates(\n args_updates, args_snapshot, keep_fn=donated_arg)\n def donated_kwarg(jax_path, prefix, c, s):\n path = graphlib.jax_to_nnx_path(jax_path)\n return path[0] in self.donate_argnames or extract.variable_changed(c, s)\n kwargs_updates = extract.mask_variable_updates(\n kwargs_updates, kwargs_snapshot, keep_fn=donated_kwarg)\n return out, (args_updates, kwargs_updates)\n\n\nclass SimpleJitWrapped(tp.Generic[P, R]):\n\n def __init__(\n self,\n fun: tp.Callable[P, R],\n in_shardings: tp.Any,\n out_shardings: tp.Any,\n static_argnums: int | tp.Sequence[int] | None = None,\n static_argnames: str | tp.Iterable[str] | None = None,\n donate_argnums: int | tp.Sequence[int] | None = None,\n donate_argnames: str | tp.Iterable[str] | None = None,\n keep_unused: bool = False,\n device: tp.Optional[jax.Device] = None,\n backend: tp.Optional[str] = None,\n inline: bool = False,\n partial_args: tuple[PartialState, ...] = (),\n graph: bool = True,\n ):\n functools.update_wrapper(self, fun)\n self.fun: tp.Callable[P, R] = fun\n self.out_shardings = out_shardings\n self.partial_args = partial_args\n self.graph = graph\n\n if in_shardings is not None and isinstance(in_shardings, (tuple, list)) and (\n static_argnums or static_argnames\n ):\n resolved = _resolve_argnums(fun, static_argnums, static_argnames)\n expanded = list(in_shardings)\n for i in sorted(resolved):\n expanded.insert(i, None)\n self.in_shardings = tuple(expanded)\n else:\n self.in_shardings = in_shardings\n\n jit_out_shardings: tp.Any\n if in_shardings is not None or out_shardings is not None:\n if isinstance(in_shardings, (tuple, list)) and (\n static_argnums or static_argnames\n ):\n resolved = _resolve_argnums(fun, static_argnums, static_argnames)\n expanded = list(in_shardings)\n for i in sorted(resolved):\n expanded.insert(i, None)\n out_in_shardings = tuple(expanded)\n else:\n out_in_shardings = in_shardings\n jit_out_shardings = (out_shardings, (out_in_shardings, None))\n else:\n jit_out_shardings = None\n\n donate_argnums_set = frozenset(\n (donate_argnums,) if isinstance(donate_argnums, int)\n else donate_argnums or ()\n )\n donate_argnames_set = frozenset(\n (donate_argnames,) if isinstance(donate_argnames, str)\n else donate_argnames or ()\n )\n self.jitted_fn = jax.jit(\n SimpleJitFn(fun, out_shardings, donate_argnums_set, donate_argnames_set, graph),\n in_shardings=in_shardings,\n out_shardings=jit_out_shardings,\n static_argnums=static_argnums,\n static_argnames=static_argnames,\n donate_argnums=donate_argnums,\n donate_argnames=donate_argnames,\n keep_unused=keep_unused,\n device=device,\n backend=backend,\n inline=inline,\n )\n\n def _maybe_to_tree(self, args, kwargs):\n if self.graph:\n args, kwargs = extract.to_tree2(\n (args, kwargs),\n prefix=(self.in_shardings, None)\n if self.in_shardings is not None\n else None,\n check_aliasing=self.in_shardings is not None,\n )\n return args, kwargs\n\n def _maybe_from_tree(self, out):\n if self.graph:\n out = extract.from_tree2(out)\n return out\n\n def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R:\n args = (*self.partial_args, *args) # type: ignore[assignment]\n args, kwargs = self._maybe_to_tree(args, kwargs)\n if not self.graph: # skip check for graph mode\n extract.check_no_aliases('jit', args=args, kwargs=kwargs)\n out, updates = self.jitted_fn(*args, **kwargs)\n extract.apply_variable_updates((args, kwargs), updates)\n return self._maybe_from_tree(out)\n\n def __get__(self, obj, objtype=None):\n if obj is None:\n return self\n return functools.partial(self, obj)\n\n def eval_shape(self, *args, **kwargs):\n args = (*self.partial_args, *args)\n args, kwargs = self._maybe_to_tree(args, kwargs)\n if not self.graph:\n extract.check_no_aliases('jit', args=args, kwargs=kwargs)\n out, updates = self.jitted_fn.eval_shape(*args, **kwargs)\n return self._maybe_from_tree(out)\n\n def trace(self, *args, **kwargs):\n args = (*self.partial_args, *args)\n args, kwargs = self._maybe_to_tree(args, kwargs)\n if not self.graph:\n extract.check_no_aliases('jit', args=args, kwargs=kwargs)\n traced = self.jitted_fn.trace(*args, **kwargs)\n return SimpleTraced(traced, self)\n\n def lower(self, *args, **kwargs):\n args = (*self.partial_args, *args)\n args, kwargs = self._maybe_to_tree(args, kwargs)\n if not self.graph:\n extract.check_no_aliases('jit', args=args, kwargs=kwargs)\n lowered = self.jitted_fn.lower(*args, **kwargs)\n return SimpleLowered(lowered, self)\ndef jit_partial(\n fun: tp.Callable[..., R],\n *partial_args: tp.Any,\n in_shardings: tp.Any = None,\n out_shardings: tp.Any = None,\n donate_argnums: int | tp.Sequence[int] | None = None,\n donate_argnames: str | tp.Iterable[str] | None = None,\n keep_unused: bool = False,\n device: tp.Optional[jax.Device] = None,\n backend: tp.Optional[str] = None,\n inline: bool = False,\n graph: bool | None = None,\n graph_updates: bool | None = None,\n) -> SimpleJitWrapped[..., R]:\n \"\"\"JIT-compile ``fun`` with pre-flattened partial arguments.\n\n Similar to ``nnx.cached_partial`` but designed for tree-mode\n (``graph=False``). Each ``partial_arg`` is flattened into a\n ``PartialState`` whose pytree structure is fixed at construction time.\n Variable values inside partial arguments can still change between calls\n without triggering recompilation, and any mutations to Variables are\n propagated back to the originals after each call.\n\n Example usage::\n\n >>> from flax import nnx\n >>> import jax.numpy as jnp\n >>> import optax\n ...\n >>> x, y = jnp.ones((4, 2)), jnp.ones((4, 3))\n >>> model = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n >>> optimizer = nnx.Optimizer(model, optax.adamw(1e-3), wrt=nnx.Param)\n ...\n >>> def train_step(model, optimizer, x, y):\n ... def loss_fn(model):\n ... return jnp.mean((model(x) - y) ** 2)\n ... loss, grads = nnx.value_and_grad(loss_fn)(model)\n ... optimizer.update(model, grads)\n ... return loss\n ...\n >>> train_step_fn = nnx.jit_partial(train_step, model, optimizer, graph=False)\n ...\n >>> loss = train_step_fn(x, y)\n\n Args:\n fun: The function to JIT-compile.\n *partial_args: Arguments to be pre-flattened and bound. These must\n appear as the first positional arguments of ``fun``.\n in_shardings: Sharding specification for inputs. When a tuple/list,\n the first ``len(partial_args)`` entries correspond to partial\n arguments and are broadcast against their original pytree\n structure. A non-tuple value (e.g. a single ``PartitionSpec``)\n is passed through directly to ``jax.jit`` and broadcast across\n all arguments uniformly.\n out_shardings: Like ``in_shardings``, but for function outputs.\n donate_argnums: Positional argument indices whose buffers may be\n donated to the computation.\n donate_argnames: Named arguments whose buffers may be donated.\n keep_unused: If ``True``, unused arguments are not pruned.\n device: Optional device to run on.\n backend: Optional backend to use.\n inline: If ``True``, inline the function.\n graph: If ``None``, uses the ``nnx_graph_mode`` config value.\n graph_updates: If ``True``, propagates updates on graph structure\n that happen inside the transform to the input graphs, has no\n effect when ``graph=False``. When ``False``, using\n ``StateSharding`` is not supported.\n\n Returns:\n A callable expecting the remaining (runtime)\n arguments.\n \"\"\"\n if graph is None:\n graph = graphlib.set_graph_mode.current_value()\n if graph_updates is None:\n graph_updates = graphlib.set_graph_updates.current_value()\n if graph_updates and graph:\n raise ValueError(\n '`graph_updates` not supported by `jit_partial`'\n )\n if any(isinstance(x, StateSharding) for x in jax.tree.leaves(in_shardings)):\n raise ValueError(\n '`in_shardings` cannot contain `StateSharding` objects '\n 'in `jit_partial`'\n )\n if any(isinstance(x, StateSharding) for x in jax.tree.leaves(out_shardings)):\n raise ValueError(\n '`out_shardings` cannot contain `StateSharding` objects '\n 'in `jit_partial`'\n )\n\n is_variable = lambda x: isinstance(x, variablelib.Variable)\n ref_index = graphlib.RefMap() if graph else None\n flat_partial_args = tuple(\n _flatten_to_partial_state(arg, ref_index=ref_index)\n for arg in partial_args\n )\n\n jit_in_shardings: tp.Any = None\n if in_shardings is not None and isinstance(in_shardings, (tuple, list)) and not graph:\n num_partial = len(partial_args)\n partial_shardings = in_shardings[:num_partial]\n runtime_shardings = in_shardings[num_partial:]\n\n flat_partial_shardings = []\n for flat_arg, orig_arg, sharding in zip(\n flat_partial_args, partial_args, partial_shardings):\n broadcasted = extract.broadcast_prefix(\n sharding, orig_arg,\n prefix_is_leaf=lambda x: x is None\n or isinstance(x, variablelib.Variable),\n tree_is_leaf=is_variable,\n )\n flat_partial_shardings.append(\n PartialState(treedef=flat_arg.treedef, leaves=broadcasted)\n )\n jit_in_shardings = (*flat_partial_shardings, *runtime_shardings)\n else:\n jit_in_shardings = in_shardings\n\n @functools.wraps(fun)\n def wrapped_fun(*args, **kwargs):\n index_ref = graphlib.IndexMap() if graph else None\n def _unflatten(arg):\n if not isinstance(arg, PartialState):\n return arg\n elif graph:\n return graphlib.unflatten(\n arg.treedef, arg.leaves, index_ref=index_ref,\n copy_variables=False,\n )\n else:\n return jax.tree.unflatten(arg.treedef, arg.leaves)\n args = (_unflatten(a) for a in args)\n return fun(*args, **kwargs)\n\n return SimpleJitWrapped(\n wrapped_fun,\n in_shardings=jit_in_shardings,\n out_shardings=out_shardings,\n donate_argnums=donate_argnums,\n donate_argnames=donate_argnames,\n keep_unused=keep_unused,\n device=device,\n backend=backend,\n inline=inline,\n partial_args=flat_partial_args,\n graph=graph,\n )\n\n\nclass JitWrapped(tp.Generic[P, R]):\n \"\"\"A function ready to be traced, lowered, and compiled.\n\n This protocol reflects the output of functions such as\n ``jax.jit``. Calling it results in JIT (just-in-time) lowering,\n compilation, and execution. It can also be explicitly lowered prior\n to compilation, and the result compiled prior to execution.\n \"\"\"\n\n def __init__(\n self,\n fun: tp.Callable[P, R],\n in_shardings: tp.Any,\n out_shardings: tp.Any,\n static_argnums: int | tp.Sequence[int] | None = None,\n static_argnames: str | tp.Iterable[str] | None = None,\n donate_argnums: int | tp.Sequence[int] | None = None,\n donate_argnames: str | tp.Iterable[str] | None = None,\n keep_unused: bool = False,\n device: tp.Optional[jax.Device] = None,\n backend: tp.Optional[str] = None,\n inline: bool = False,\n ):\n functools.update_wrapper(self, fun)\n self.fun: tp.Callable[P, R] = fun\n kwarg_shardings = None\n self.jax_in_shardings = jax.tree.map(\n lambda x: extract.NodeStates.from_prefixes(x.shardings, metadata=x)\n if isinstance(x, StateSharding)\n else x,\n in_shardings,\n )\n self.jax_out_shardings = jax.tree.map(\n lambda x: extract.NodeStates.from_prefixes(x.shardings, metadata=x)\n if isinstance(x, StateSharding)\n else x,\n out_shardings,\n )\n\n if isinstance(in_shardings, (tuple, list)) and (static_argnums or static_argnames):\n # We should reintroduce None values into in_shardings corresponding to static arguments\n static_argnums = _resolve_argnums(fun, static_argnums, static_argnames)\n in_shardings = list(in_shardings)\n for static_arg_index in sorted(static_argnums):\n in_shardings.insert(static_arg_index, None)\n in_shardings = tuple(in_shardings)\n\n jax_out_in_shardings = jax.tree.map(\n lambda x: extract.NodeStates.from_prefixes(x.shardings, metadata=x)\n if isinstance(x, StateSharding)\n else x,\n in_shardings,\n )\n\n self.jitted_fn = jax.jit(\n JitFn(fun, in_shardings, out_shardings, kwarg_shardings, self),\n in_shardings=self.jax_in_shardings,\n out_shardings=(jax_out_in_shardings, kwarg_shardings, self.jax_out_shardings),\n static_argnums=static_argnums,\n static_argnames=static_argnames,\n donate_argnums=donate_argnums,\n donate_argnames=donate_argnames,\n keep_unused=keep_unused,\n device=device,\n backend=backend,\n inline=inline,\n )\n self.in_shardings = in_shardings\n self.out_shardings = out_shardings\n self.kwarg_shardings = kwarg_shardings\n self.static_argnums = static_argnums\n\n # implement descriptor protocol so that we can use this as a method\n def __get__(self, obj, objtype=None):\n if obj is None:\n return self\n return functools.partial(self, obj)\n\n def _get_pure_args_kwargs(self, args, kwargs):\n pure_args, pure_kwargs = extract.to_tree(\n (args, kwargs),\n prefix=(self.in_shardings, self.kwarg_shardings)\n if self.in_shardings is not None or self.kwarg_shardings is not None\n else None,\n split_fn=_jit_split_fn,\n check_aliasing=self.in_shardings is not None\n or self.kwarg_shardings is not None,\n ctxtag=self,\n )\n return pure_args, pure_kwargs\n\n def _get_non_pure_out(self, pure_args_out, pure_kwargs_out, pure_out, /):\n _args_out, _kwargs_out, out = extract.from_tree(\n (pure_args_out, pure_kwargs_out, pure_out),\n merge_fn=_jit_merge_fn,\n is_inner=False,\n ctxtag=self,\n )\n return out\n\n def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R:\n # run dynamic_cache_context before update_context\n with graphlib.update_context(self):\n pure_args, pure_kwargs = self._get_pure_args_kwargs(args, kwargs)\n pure_args_out, pure_kwargs_out, pure_out = self.jitted_fn(\n *pure_args, **pure_kwargs\n )\n out = self._get_non_pure_out(pure_args_out, pure_kwargs_out, pure_out)\n return out\n\n def eval_shape(self, *args, **kwargs):\n \"\"\"See ``jax.eval_shape``.\"\"\"\n args, kwargs = graphlib.clone((args, kwargs))\n with graphlib.update_context(self):\n pure_args, pure_kwargs = self._get_pure_args_kwargs(args, kwargs)\n pure_args_out, pure_kwargs_out, pure_out = self.jitted_fn.eval_shape(\n *pure_args, **pure_kwargs\n )\n out = self._get_non_pure_out(pure_args_out, pure_kwargs_out, pure_out)\n return out\n\n def trace(self, *args, **kwargs) -> Traced:\n \"\"\"Trace this function explicitly for the given arguments.\n\n A traced function is staged out of Python and translated to a jaxpr. It is\n ready for lowering but not yet lowered.\n\n Returns:\n A ``Traced`` instance representing the tracing.\n \"\"\"\n with graphlib.update_context(self):\n pure_args, pure_kwargs = self._get_pure_args_kwargs(args, kwargs)\n traced = self.jitted_fn.trace(*pure_args, **pure_kwargs)\n return Traced(traced, self)\n\n def lower(self, *args, **kwargs) -> Lowered:\n \"\"\"Lower this function explicitly for the given arguments.\n\n This is a shortcut for ``self.trace(*args, **kwargs).lower()``.\n\n A lowered function is staged out of Python and translated to a\n compiler's input language, possibly in a backend-dependent\n manner. It is ready for compilation but not yet compiled.\n\n Returns:\n A ``Lowered`` instance representing the lowering.\n \"\"\"\n with graphlib.update_context(self):\n pure_args, pure_kwargs = self._get_pure_args_kwargs(args, kwargs)\n lowered = self.jitted_fn.lower(*pure_args, **pure_kwargs)\n return Lowered(lowered, self)\n\n\nclass Stage:\n args_info: tp.Any # PyTree of ArgInfo\n\n @property\n def _inner_obj(self) -> tp.Any:\n raise NotImplementedError\n\n @property\n def in_tree(self) -> jax.tree_util.PyTreeDef:\n return self._inner_obj.in_tree\n\n @property\n def in_avals(self):\n return self._inner_obj.in_avals\n\n @property\n def donate_argnums(self):\n return self._inner_obj.donate_argnums\n\n@dataclasses.dataclass(frozen=True, slots=True)\nclass Compiled(Stage):\n \"\"\"Compiled representation of a function specialized to types/values.\n\n A compiled computation is associated with an executable and the\n remaining information needed to execute it. It also provides a\n common API for querying properties of compiled computations across\n JAX's various compilation paths and backends.\n \"\"\"\n\n compiled: jax.stages.Compiled\n jit_wrapped: JitWrapped\n\n @property\n def _inner_obj(self):\n return self.compiled\n\n @property\n def args_info(self) -> tp.Any: # PyTree of ArgInfo\n raise self.compiled.args_info\n\n @staticmethod\n def call(*args, **kwargs):\n raise NotImplementedError\n\n def __call__(self, *args, **kwargs):\n with graphlib.update_context(self.jit_wrapped):\n pure_args, pure_kwargs = self.jit_wrapped._get_pure_args_kwargs(\n args, kwargs\n )\n pure_args_out, pure_kwargs_out, pure_out = self.compiled(\n *pure_args, **pure_kwargs\n )\n out = self.jit_wrapped._get_non_pure_out(\n pure_args_out, pure_kwargs_out, pure_out\n )\n return out\n\n @property\n def out_tree(self) -> jax.tree_util.PyTreeDef:\n return self.compiled.out_tree\n\n def as_text(self) -> str | None:\n \"\"\"A human-readable text representation of this executable.\n\n Intended for visualization and debugging purposes. This is not a valid nor\n reliable serialization.\n\n Returns ``None`` if unavailable, e.g. based on backend, compiler, or\n runtime.\n \"\"\"\n return self.compiled.as_text()\n\n def cost_analysis(self) -> tp.Any | None:\n \"\"\"A summary of execution cost estimates.\n\n Intended for visualization and debugging purposes. The object output by\n this is some simple data structure that can easily be printed or serialized\n (e.g. nested dicts, lists, and tuples with numeric leaves). However, its\n structure can be arbitrary: it may be inconsistent across versions of JAX\n and jaxlib, or even across invocations.\n\n Returns ``None`` if unavailable, e.g. based on backend, compiler, or\n runtime.\n \"\"\"\n return self.compiled.cost_analysis()\n\n def memory_analysis(self) -> tp.Any | None:\n \"\"\"A summary of estimated memory requirements.\n\n Intended for visualization and debugging purposes. The object output by\n this is some simple data structure that can easily be printed or serialized\n (e.g. nested dicts, lists, and tuples with numeric leaves). However, its\n structure can be arbitrary: it may be inconsistent across versions of JAX\n and jaxlib, or even across invocations.\n\n Returns ``None`` if unavailable, e.g. based on backend, compiler, or\n runtime.\n \"\"\"\n return self.compiled.memory_analysis()\n\n def runtime_executable(self) -> tp.Any | None:\n \"\"\"An arbitrary object representation of this executable.\n\n Intended for debugging purposes. This is not valid nor reliable\n serialization. The output has no guarantee of consistency across\n invocations.\n\n Returns ``None`` if unavailable, e.g. based on backend, compiler, or\n runtime.\n \"\"\"\n return self.compiled.runtime_executable()\n\n @property\n def input_shardings(self): # PyTree[sharding.Sharding]\n return self.compiled.input_shardings\n\n @property\n def output_shardings(self): # PyTree[sharding.Sharding]\n return self.compiled.output_shardings\n\n @property\n def input_layouts(self):\n return self.compiled.input_formats\n\n\n@dataclasses.dataclass(frozen=True, slots=True)\nclass Lowered(Stage):\n \"\"\"Lowering of a function specialized to argument types and values.\n\n A lowering is a computation ready for compilation. This class\n carries a lowering together with the remaining information needed to\n later compile and execute it. It also provides a common API for\n querying properties of lowered computations across JAX's various\n lowering paths (:func:`~jax.jit`, :func:`~jax.pmap`, etc.).\n \"\"\"\n\n lowered: jax.stages.Lowered\n jit_wrapped: JitWrapped\n\n @property\n def _inner_obj(self):\n return self.lowered\n\n @property\n def args_info(self) -> tp.Any: # PyTree of ArgInfo\n return self.lowered.args_info\n\n @property\n def out_tree(self):\n return self.lowered.out_tree\n\n @classmethod\n def from_flat_info(\n cls,\n lowering: tp.Any, # type: ignore[name-defined]\n in_tree: jax.tree_util.PyTreeDef,\n in_avals,\n donate_argnums: tuple[int, ...],\n out_tree: jax.tree_util.PyTreeDef,\n no_kwargs: bool = False,\n ):\n raise NotImplementedError\n\n def compile(\n self, compiler_options: jax.stages.CompilerOptions | None = None\n ) -> Compiled:\n \"\"\"Compile, returning a corresponding ``Compiled`` instance.\"\"\"\n compiled = self.lowered.compile(compiler_options)\n return Compiled(compiled, self.jit_wrapped)\n\n def as_text(\n self, dialect: str | None = None, *, debug_info: bool = False\n ) -> str:\n \"\"\"A human-readable text representation of this lowering.\n\n Intended for visualization and debugging purposes. This need not be a valid\n nor reliable serialization.\n Use `jax.export` if you want reliable and portable serialization.\n\n Args:\n dialect: Optional string specifying a lowering dialect (e.g. \"stablehlo\",\n or \"hlo\").\n debug_info: Whether to include debugging information,\n e.g., source location.\n \"\"\"\n return self.lowered.as_text(dialect=dialect, debug_info=debug_info)\n\n def compiler_ir(self, dialect: str | None = None) -> tp.Any | None:\n \"\"\"An arbitrary object representation of this lowering.\n\n Intended for debugging purposes. This is not a valid nor reliable\n serialization. The output has no guarantee of consistency across\n invocations.\n Use `jax.export` if you want reliable and portable serialization.\n\n Returns ``None`` if unavailable, e.g. based on backend, compiler, or\n runtime.\n\n Args:\n dialect: Optional string specifying a lowering dialect (e.g. \"stablehlo\",\n or \"hlo\").\n \"\"\"\n return self.lowered.compiler_ir(dialect=dialect)\n\n def cost_analysis(self) -> tp.Any | None:\n \"\"\"A summary of execution cost estimates.\n\n Intended for visualization and debugging purposes. The object output by\n this is some simple data structure that can easily be printed or serialized\n (e.g. nested dicts, lists, and tuples with numeric leaves). However, its\n structure can be arbitrary: it may be inconsistent across versions of JAX\n and jaxlib, or even across invocations.\n\n Returns ``None`` if unavailable, e.g. based on backend, compiler, or\n runtime.\n \"\"\"\n return self.lowered.cost_analysis()\n\n@dataclasses.dataclass(frozen=True, slots=True)\nclass Traced(Stage):\n \"\"\"Traced form of a function specialized to argument types and values.\n\n A traced computation is ready for lowering. This class carries the\n traced representation with the remaining information needed to later\n lower, compile, and execute it.\n \"\"\"\n\n traced: jax.stages.Traced\n jit_wrapped: JitWrapped\n\n @property\n def _inner_obj(self):\n return self.traced\n\n @property\n def out_info(self):\n return self.traced.out_info\n\n def lower(\n self, *, lowering_platforms: tuple[str, ...] | None = None\n ) -> Lowered:\n \"\"\"Lower to compiler input, returning a ``Lowered`` instance.\"\"\"\n lowered = self.traced.lower(lowering_platforms=lowering_platforms)\n return Lowered(lowered, self.jit_wrapped)\n\n\n@dataclasses.dataclass(frozen=True, slots=True)\nclass SimpleCompiled(Stage):\n compiled: jax.stages.Compiled\n jit_wrapped: SimpleJitWrapped\n\n @property\n def _inner_obj(self):\n return self.compiled\n\n @property\n def args_info(self) -> tp.Any:\n raise self.compiled.args_info\n\n @staticmethod\n def call(*args, **kwargs):\n raise NotImplementedError\n\n def __call__(self, *args, **kwargs):\n args = (*self.jit_wrapped.partial_args, *args)\n args, kwargs = self.jit_wrapped._maybe_to_tree(args, kwargs)\n if not self.jit_wrapped.graph:\n extract.check_no_aliases('jit', args=args, kwargs=kwargs)\n out, updates = self.compiled(*args, **kwargs)\n extract.apply_variable_updates((args, kwargs), updates)\n return self.jit_wrapped._maybe_from_tree(out)\n\n @property\n def out_tree(self) -> jax.tree_util.PyTreeDef:\n return self.compiled.out_tree\n\n def as_text(self) -> str | None:\n return self.compiled.as_text()\n\n def cost_analysis(self) -> tp.Any | None:\n return self.compiled.cost_analysis()\n\n def memory_analysis(self) -> tp.Any | None:\n return self.compiled.memory_analysis()\n\n def runtime_executable(self) -> tp.Any | None:\n return self.compiled.runtime_executable()\n\n @property\n def input_shardings(self):\n return self.compiled.input_shardings\n\n @property\n def output_shardings(self):\n return self.compiled.output_shardings\n\n @property\n def input_layouts(self):\n return self.compiled.input_formats\n\n\n@dataclasses.dataclass(frozen=True, slots=True)\nclass SimpleLowered(Stage):\n lowered: jax.stages.Lowered\n jit_wrapped: SimpleJitWrapped\n\n @property\n def _inner_obj(self):\n return self.lowered\n\n @property\n def args_info(self) -> tp.Any:\n return self.lowered.args_info\n\n @property\n def out_tree(self):\n return self.lowered.out_tree\n\n def compile(\n self, compiler_options: jax.stages.CompilerOptions | None = None\n ) -> SimpleCompiled:\n compiled = self.lowered.compile(compiler_options)\n return SimpleCompiled(compiled, self.jit_wrapped)\n\n def as_text(\n self, dialect: str | None = None, *, debug_info: bool = False\n ) -> str:\n return self.lowered.as_text(dialect=dialect, debug_info=debug_info)\n\n def compiler_ir(self, dialect: str | None = None) -> tp.Any | None:\n return self.lowered.compiler_ir(dialect=dialect)\n\n def cost_analysis(self) -> tp.Any | None:\n return self.lowered.cost_analysis()\n\n\n@dataclasses.dataclass(frozen=True, slots=True)\nclass SimpleTraced(Stage):\n traced: jax.stages.Traced\n jit_wrapped: SimpleJitWrapped\n\n @property\n def _inner_obj(self):\n return self.traced\n\n @property\n def out_info(self):\n return self.traced.out_info\n\n def lower(\n self, *, lowering_platforms: tuple[str, ...] | None = None\n ) -> SimpleLowered:\n lowered = self.traced.lower(lowering_platforms=lowering_platforms)\n return SimpleLowered(lowered, self.jit_wrapped)\n# -------------------------------\n# shard_map\n# -------------------------------\n\n# TODO: create StateSpec and consider enabling a mode that does\n# not use filters during split for performance. Overall there might\n# be performance limitations for using shard_map at a top-level\n\n\n@dataclasses.dataclass(eq=False)\nclass SimpleShardMapFn:\n f: tp.Callable[..., tp.Any]\n graph: bool\n out_specs: tp.Any\n\n def __post_init__(self):\n functools.update_wrapper(self, self.f, updated=())\n\n @extract.treemap_copy_args\n def __call__(self, *args):\n updates, snapshot = extract.updates_and_snapshot(args)\n if self.graph:\n args = extract.from_tree2(args)\n out = self.f(*args)\n if self.graph:\n out = extract.to_tree2(out, prefix=self.out_specs)\n extract.check_no_aliases('shard_map', args=updates, out=out)\n updates = extract.mask_variable_updates(updates, snapshot)\n return out, updates\n\n\n@dataclasses.dataclass(eq=False)\nclass ShardMapFn:\n f: tp.Callable[..., tp.Any]\n in_specs: tp.Any\n out_specs: tp.Any\n kwarg_specs: tp.Any\n ctxtag: tp.Hashable\n\n def __post_init__(self):\n functools.update_wrapper(self, self.f)\n\n def __call__(self, *pure_args, **pure_kwargs):\n args, kwargs = extract.from_tree(\n (pure_args, pure_kwargs),\n merge_fn=_jit_merge_fn,\n ctxtag=self.ctxtag,\n is_inner=True,\n )\n\n out = self.f(*args, **kwargs)\n\n args_out, kwargs_out = extract.clear_non_graph_nodes((args, kwargs))\n pure_args_out, pure_kwargs_out, pure_out = extract.to_tree(\n (args_out, kwargs_out, out),\n prefix=(self.in_specs, self.kwarg_specs, self.out_specs),\n ctxtag=self.ctxtag,\n split_fn=_jit_split_fn,\n )\n\n return pure_args_out, pure_kwargs_out, pure_out\n\n\n@tp.overload\ndef shard_map(\n f: F,\n *,\n mesh: Mesh | AbstractMesh,\n in_specs: Specs,\n out_specs: Specs,\n axis_names: tp.AbstractSet[AxisName] = frozenset(),\n check_vma: bool = True,\n graph: bool | None = None,\n graph_updates: bool | None = None,\n) -> F:\n ...\n\n\n@tp.overload\ndef shard_map(\n *,\n mesh: Mesh | AbstractMesh,\n in_specs: Specs,\n out_specs: Specs,\n axis_names: tp.AbstractSet[AxisName] = frozenset(),\n check_vma: bool = True,\n graph: bool | None = None,\n graph_updates: bool | None = None,\n) -> tp.Callable[[F], F]:\n ...\n\n\ndef shard_map(\n f: F | type[Missing] = Missing,\n *,\n mesh: Mesh | AbstractMesh,\n in_specs: Specs,\n out_specs: Specs,\n axis_names: tp.AbstractSet[AxisName] = frozenset(),\n check_vma: bool = True,\n graph: bool | None = None,\n graph_updates: bool | None = None,\n) -> F | tp.Callable[[F], F]:\n \"\"\"\n Lifted version of\n `jax.shard_map `_\n that can handle Modules / graph nodes as arguments.\n\n Simple data parallel example::\n\n import jax\n import jax.numpy as jnp\n from flax import nnx\n from jax.sharding import PartitionSpec as P\n\n mesh = jax.sharding.Mesh(jax.local_devices(), ('data',))\n\n m = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n x = jnp.ones((32, 2))\n\n @nnx.shard_map(\n mesh=mesh, in_specs=(P(None), P('data')), out_specs=P('data')\n )\n def f(m, x):\n return m(x)\n\n y = f(m, x)\n\n jax.debug.visualize_array_sharding(y)\n\n Notice that here we simply used some ``PartitionSpec`` to define the spec\n the the whole model and data. This works for simple cases but if we need\n to assign different ``PartitionSpec`` to different parts of the model we\n need to use ``StateSharding`` and create some filters that allow us to target\n specific parts of the model. Here's an example of how to do tensor parallelism\n for a simple MLP block using ``StateSharding`` and filters::\n\n mesh = jax.sharding.Mesh(jax.local_devices(), ('model',))\n\n class MLP(nnx.Module):\n def __init__(self, din, dhidden, dout, *, rngs: nnx.Rngs):\n self.linear1 = nnx.Linear(din, dhidden, use_bias=False, rngs=rngs)\n self.linear2 = nnx.Linear(dhidden, dout, use_bias=False, rngs=rngs)\n\n def __call__(self, x):\n return self.linear2(jax.nn.relu(self.linear1(x)))\n\n m = MLP(2, 64, 3, rngs=nnx.Rngs(0))\n x = jnp.ones((32, 2))\n\n def path_ends_with(*path_suffix): # custom filter\n return lambda path, value: path[-len(path_suffix):] == path_suffix\n\n model_spec = nnx.StateSharding({\n path_ends_with('linear1', 'kernel'): P(None, 'model'),\n path_ends_with('linear2', 'kernel'): P('model', None),\n })\n\n @nnx.shard_map(mesh=mesh, in_specs=(model_spec, P(None)), out_specs=P(None))\n def f(m, x):\n y = m(x)\n return jax.lax.psum(y, 'model')\n\n y = f(m, x)\n\n jax.debug.visualize_array_sharding(m.linear1.kernel[...])\n jax.debug.visualize_array_sharding(m.linear2.kernel[...])\n\n\n Alternatively, a ``State`` object with the exact PartitionSpec for each\n state then you can be passed to ``StateSharding``::\n\n mesh = jax.sharding.Mesh(jax.local_devices(), ('model',))\n\n class MLP(nnx.Module):\n def __init__(self, din, dhidden, dout, *, rngs: nnx.Rngs):\n self.linear1 = nnx.Linear(din, dhidden, use_bias=False, rngs=rngs)\n self.linear2 = nnx.Linear(dhidden, dout, use_bias=False, rngs=rngs)\n\n def __call__(self, x):\n return self.linear2(jax.nn.relu(self.linear1(x)))\n\n m = MLP(2, 64, 3, rngs=nnx.Rngs(0))\n x = jnp.ones((32, 2))\n\n model_spec = nnx.State(\n {\n 'linear1': {'kernel': P(None, 'model')},\n 'linear2': {'kernel': P('model', None)},\n }\n )\n\n @nnx.shard_map(\n mesh=mesh,\n in_specs=(nnx.StateSharding(model_spec), P(None)),\n out_specs=P(None),\n )\n def f(m, x):\n y = m(x)\n return jax.lax.psum(y, 'model')\n\n y = f(m, x)\n\n jax.debug.visualize_array_sharding(m.linear1.kernel[...])\n jax.debug.visualize_array_sharding(m.linear2.kernel[...])\n\n Here ``model_spec`` was created manually but you can also automate\n this process by using ``nnx.get_partition_spec`` to automatically\n create it for you (see\n `Scale up on multiple devices `_\n ).\n\n Args:\n f: callable to be mapped. Each application of ``f``, or \"instance\" of ``f``,\n takes as input a shard of the mapped-over arguments and produces a shard\n of the output.\n mesh: a ``jax.sharding.Mesh`` representing the array of devices over which\n to shard the data and on which to execute instances of ``f``. The names of\n the ``Mesh`` can be used in collective communication operations in ``f``.\n This is typically created by a utility function like\n :func:`jax.experimental.mesh_utils.create_device_mesh`.\n in_specs: a pytree with ``jax.sharding.PartitionSpec``or ``nnx.StateSharding``\n (mapping substates to ``PartitionSpec``s) instances as leaves,\n with a tree structure that is a tree prefix of the\n args tuple to be mapped over. Similar to ``jax.sharding.NamedSharding``,\n each ``PartitionSpec`` represents how the corresponding argument (or subtree\n of arguments) should be sharded along the named axes of ``mesh``. In each\n ``PartitionSpec``, mentioning a ``mesh`` axis name at a position expresses sharding\n the corresponding argument array axis along that positional axis; not\n mentioning an axis name expresses replication. If an argument, or argument\n subtree, has a corresponding spec of None, that argument is not sharded.\n out_specs: a pytree with ``jax.sharding.PartitionSpec`` or ``nnx.StateSharding``\n (mapping substates to ``PartitionSpec``s) instances as leaves, with a tree structure\n that is a tree prefix of the output of ``f``.\n Each ``PartitionSpec`` represents how the corresponding output shards should be\n concatenated. In each ``PartitionSpec``, metioning a ``mesh`` axis name at\n a position expresses concatenation of that mesh axis's shards along the\n corresponding positional axis. Not mentioning a ``mesh`` axis name\n expresses a promise that the output values are equal along that mesh axis,\n and that rather than concatenating only a single value should be produced.\n axis_names: optional set of axis names from ``mesh`` over which the function ``f``\n is manual. If empty, ``f``, is manual over all mesh axes.\n check_vma: optional boolean representing whether to enable additional validity\n checks and automatic differentiation optimizations. The validity checks concern\n whether any mesh axis names not mentioned in ``out_specs`` are consistent with\n how the outputs of ``f`` are replicated.\n graph: If ``True`` (default), uses graph-mode which supports the full\n NNX feature set including shared references and reference semantics.\n If ``False``, uses tree-mode which treats Modules as regular JAX\n pytrees, avoiding the overhead of the graph protocol. Tree-mode does\n not support ``StateSharding`` or shared ``Variable`` references.\n graph_updates: If ``True``, propagates updates on graph structure\n that happen inside the transform to the input graphs, has no\n effect when ``graph=False``. When ``False``, using\n ``StateSharding`` is not supported.\n\n Returns:\n A callable that applies the input function ``f`` across data sharded according to\n the ``mesh`` and ``in_specs``.\n \"\"\"\n if graph is None:\n graph = graphlib.set_graph_mode.current_value()\n if graph_updates is None:\n graph_updates = graphlib.set_graph_updates.current_value()\n if f is Missing:\n return functools.partial(\n shard_map,\n mesh=mesh,\n in_specs=in_specs,\n out_specs=out_specs,\n axis_names=axis_names,\n check_vma=check_vma,\n graph=graph,\n graph_updates=graph_updates,\n ) # type: ignore[return-value]\n assert not isinstance(f, type)\n\n f_unbound, _, was_bound = _resolve_bound_callable(f)\n if was_bound:\n _raise_bound_method_error('shard_map')\n\n if not graph or not graph_updates:\n if any(isinstance(x, StateSharding) for x in jax.tree.leaves(in_specs)):\n raise ValueError(\n '`in_specs` cannot contain `StateSharding` objects '\n 'when `graph=False`'\n )\n if any(isinstance(x, StateSharding) for x in jax.tree.leaves(out_specs)):\n raise ValueError(\n '`out_specs` cannot contain `StateSharding` objects '\n 'when `graph=False`'\n )\n if graph:\n extract.check_prefix(in_specs, 'in_specs', 'shard_map')\n extract.check_prefix(out_specs, 'out_specs', 'shard_map')\n\n shard_map_fn = jax.shard_map(\n SimpleShardMapFn(f_unbound, graph=graph, out_specs=out_specs),\n mesh=mesh,\n in_specs=in_specs,\n out_specs=(out_specs, in_specs),\n axis_names=axis_names,\n check_vma=check_vma,\n )\n\n @functools.wraps(f)\n def shard_map_wrapper(*args, **kwargs):\n if graph:\n args = extract.to_tree2(\n args,\n prefix=in_specs,\n check_aliasing=in_specs is not None,\n )\n extract.check_no_aliases('shard_map', args=args)\n out, updates = shard_map_fn(*args, **kwargs)\n extract.apply_variable_updates(args, updates)\n if graph:\n out = extract.from_tree2(out)\n return out\n\n shard_map_wrapper.inner = shard_map_fn # type: ignore\n return shard_map_wrapper # type: ignore\n\n kwarg_specs = PartitionSpec()\n jax_in_specs = jax.tree.map(\n lambda x: extract.NodeStates(\n _graphdef=PartitionSpec(), # type: ignore[arg-type]\n states=x.shardings,\n metadata=x,\n )\n if isinstance(x, StateSharding)\n else x,\n in_specs,\n )\n jax_out_specs = jax.tree.map(\n lambda x: extract.NodeStates(\n _graphdef=PartitionSpec(), # type: ignore[arg-type]\n states=x.shardings,\n metadata=x,\n )\n if isinstance(x, StateSharding)\n else x,\n out_specs,\n )\n\n @functools.wraps(f) # type: ignore[no-redef]\n def shard_map_wrapper(*args, **kwargs):\n with graphlib.update_context(shard_map_wrapper):\n pure_args, pure_kwargs = extract.to_tree(\n (args, kwargs),\n prefix=(in_specs, kwarg_specs)\n if in_specs is not None or kwarg_specs is not None\n else None,\n split_fn=_jit_split_fn,\n check_aliasing=in_specs is not None or kwarg_specs is not None,\n ctxtag=shard_map_wrapper,\n )\n pure_args_out, pure_kwargs_out, pure_out = shard_map_fn(\n *pure_args, **pure_kwargs\n )\n _args_out, _kwargs_out, out = extract.from_tree(\n (pure_args_out, pure_kwargs_out, pure_out),\n merge_fn=_jit_merge_fn,\n is_inner=False,\n ctxtag=shard_map_wrapper,\n )\n return out\n\n shard_map_fn = jax.shard_map(\n ShardMapFn(f_unbound, in_specs, out_specs, kwarg_specs, shard_map_wrapper),\n mesh=mesh,\n in_specs=jax_in_specs,\n out_specs=(jax_in_specs, kwarg_specs, jax_out_specs), # type: ignore\n axis_names=axis_names,\n check_vma=check_vma,\n )\n\n shard_map_wrapper.inner = shard_map_fn # type: ignore\n\n return shard_map_wrapper # type: ignore\n\n\n# We can't use private methods from jax._src.api_util\n# We copy the function: api_util.fun_signature\ndef _fun_signature(fun: tp.Callable) -> inspect.Signature | None:\n try:\n return inspect.signature(fun)\n except (ValueError, TypeError):\n return None\n\n# Adapted copy of private jax function from api_util: fun_signature\ndef _resolve_argnums(\n fun: tp.Callable,\n static_argnums: int | tp.Sequence[int] | None,\n static_argnames: str | tp.Iterable[str] | None,\n) -> tuple[int, ...]:\n def _ensure_index_tuple(x: tp.Any) -> tuple[int, ...]:\n \"\"\"Convert x to a tuple of indices.\"\"\"\n try:\n return (operator.index(x),)\n except TypeError:\n return tuple(map(operator.index, x))\n\n def _ensure_str(x: str) -> str:\n if not isinstance(x, str):\n raise TypeError(f\"argument is not a string: {x}\")\n return x\n\n def _ensure_str_tuple(x: str | tp.Iterable[str]) -> tuple[str, ...]:\n \"\"\"Convert x to a tuple of strings.\"\"\"\n if isinstance(x, str):\n return (x,)\n else:\n return tuple(map(_ensure_str, x))\n\n signature = _fun_signature(fun)\n if signature is None:\n # Some built-in functions don't support signature.\n # See: https://github.com/python/cpython/issues/73485\n # In this case no validation is done\n static_argnums = () if static_argnums is None else _ensure_index_tuple(\n static_argnums)\n else:\n # Infer argnums and argnames according to docstring\n # If nums is None and names is not None, then nums are inferred from the\n # names and vice-versa.\n _POSITIONAL_OR_KEYWORD = inspect.Parameter.POSITIONAL_OR_KEYWORD\n _POSITIONAL_ARGUMENTS = (\n inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD\n )\n\n def infer_argnums_and_argnames(\n sig: inspect.Signature,\n argnums: int | tp.Iterable[int] | None,\n argnames: str | tp.Iterable[str] | None,\n ) -> tuple[tuple[int, ...], tuple[str, ...]]:\n \"\"\"Infer missing argnums and argnames for a function with inspect.\"\"\"\n if argnums is None and argnames is None:\n return (), ()\n\n if argnums is not None and argnames is not None:\n argnums = _ensure_index_tuple(argnums)\n argnames = _ensure_str_tuple(argnames)\n return argnums, argnames\n\n parameters = sig.parameters\n if argnums is None:\n assert argnames is not None\n argnames = _ensure_str_tuple(argnames)\n argnums = tuple(\n i for i, (k, param) in enumerate(parameters.items())\n if param.kind == _POSITIONAL_OR_KEYWORD and k in argnames\n )\n else:\n argnums = _ensure_index_tuple(argnums)\n argnames = tuple(\n k for i, (k, param) in enumerate(parameters.items())\n if param.kind == _POSITIONAL_OR_KEYWORD and i in argnums\n )\n return argnums, argnames\n\n def _validate_argnums(sig: inspect.Signature, argnums: tuple[int, ...], argnums_name: str) -> None:\n n_pos_args = 0\n for param in sig.parameters.values():\n if param.kind in _POSITIONAL_ARGUMENTS:\n n_pos_args += 1\n\n elif param.kind is inspect.Parameter.VAR_POSITIONAL:\n # We can have any number of positional arguments\n return\n\n if argnums and (-min(argnums) > n_pos_args or max(argnums) >= n_pos_args):\n raise ValueError(f\"Jitted function has {argnums_name}={argnums}, \"\n f\"but only accepts {n_pos_args} positional arguments.\")\n\n static_argnums, static_argnames = infer_argnums_and_argnames(\n signature, static_argnums, static_argnames)\n\n # Validation\n _validate_argnums(signature, static_argnums, \"static_argnums\")\n\n return static_argnums\n" }, { "path": "flax/nnx/transforms/general.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport functools\nimport typing as tp\n\nfrom flax.nnx import (\n extract,\n graphlib,\n)\nfrom flax.typing import MISSING, Missing\n\nA = tp.TypeVar('A')\nF = tp.TypeVar('F', bound=tp.Callable[..., tp.Any])\n\n\n# -------------------------------\n# (split|merge)_inputs\n# -------------------------------\n\n\n@tp.overload\ndef split_inputs(\n *,\n ctxtag: str = 'split_merge_inputs',\n) -> tp.Callable[[F], F]: ...\n@tp.overload\ndef split_inputs(\n f: F,\n *,\n ctxtag: str = 'split_merge_inputs',\n) -> F: ...\ndef split_inputs(\n f: F | Missing = MISSING,\n *,\n ctxtag: str = 'split_merge_inputs',\n) -> F | tp.Callable[[F], F]:\n \"\"\"Takes in a function that contains graph nodes in the inputs and outputs, and\n returns a function that replaces the graph nodes with some jax-compatible data\n structures. Must be used in conjunction with :func:`merge_inputs`.\n\n Args:\n f: The function to be transformed.\n ctxtag: The context tag to be used for the transformation. Defaults to\n 'split_merge_inputs'.\n\n Returns:\n The transformed function.\n\n ``split_inputs`` and ``merge_inputs`` can be used to lift functions that operate\n on jax datastructures (pytrees) to functions that operate on graph nodes. ``split_inputs``\n will take graph nodes in the inputs and outputs and replace them with jax-compatible data\n structures, usually before calling into the transformed function, while ``merge_inputs``\n will convert the jax-compatible data structures back to graph nodes, usually inside the\n transformed function. For common transforms like ``jax.jit`` and ``jax.vmap`` NNX will\n provide a version that works with graph nodes, but for other transforms you can use\n ``split_inputs`` and ``merge_inputs`` to manually lift the function.\n\n The following example demonstrates how to use ``split_inputs`` and ``merge_inputs`` to\n lift ``jax.jit`` to work over a silly function has a stateful operation that zeros out\n the kernel of a linear layer::\n\n >>> from flax import nnx\n >>> import jax.numpy as jnp\n >>> import jax\n ...\n >>> @split_inputs\n ... @jax.jit\n ... @merge_inputs\n ... def forward_and_zero(model: nnx.Linear, x: jax.Array):\n ... y = model(x)\n ... model.kernel[...] *= 0 # zero out the kernel\n ... return y\n ...\n >>> model = nnx.Linear(2, 2, rngs=nnx.Rngs(0))\n >>> y = forward_and_zero(model, jnp.ones((1, 2)))\n >>> y.shape\n (1, 2)\n >>> assert jnp.allclose(model.kernel, 0)\n\n As shown above, not only does the lifted function work with graph nodes, but it also\n propagates the side effects of the original function. **Note**: in practice use ``nnx.jit``\n instead.\n\n Splitting and merging can also be applied to multiple functions in a pipeline. The following\n example show how to lift ``jax.lax.cond`` by using ``split_inputs`` over ``cond`` and\n ``merge_inputs`` over the branches::\n\n >>> model = nnx.Linear(2, 2, rngs=nnx.Rngs(0))\n >>> x = jnp.ones((1, 2))\n ...\n >>> true_fn = lambda m, x: m(x)\n >>> false_fn = lambda m, x: x + 1\n ...\n >>> y = split_inputs(jax.lax.cond)(\n ... False,\n ... merge_inputs(true_fn),\n ... merge_inputs(false_fn), # <== gets called\n ... model,\n ... x,\n ... )\n >>> assert jnp.allclose(y, 2)\n\n **Lifting functions with output semantics**\n\n ``merge_inputs`` internally returns a ``(inputs, output)`` tuple, where ``inputs`` is the\n tuple of the input arguments with non-graph node leaves set to ``None``, and ``output`` is\n the output of the function. This is done to propage all the state changes in the function\n to the graph nodes outside the function. If the transform function has output semantics\n like e.g. ``jax.vmap``'s ``out_axes``, you must account for this in the by configuring\n the arguments accordingly::\n\n >>> from functools import partial\n ...\n >>> model = nnx.Linear(2, 2, rngs=nnx.Rngs(0))\n ...\n >>> in_axes = (None, 0)\n >>> out_axes = (in_axes, 0) # <== internal output arrangement\n ...\n >>> @split_inputs\n ... @partial(jax.vmap, in_axes=in_axes, out_axes=out_axes)\n ... @merge_inputs\n ... def forward(model: nnx.Linear, x: jax.Array):\n ... return model(x)\n ...\n >>> x = jnp.ones((10, 2))\n >>> y = forward(model, x)\n >>> y.shape\n (10, 2)\n\n .. note::\n If the transform has a rigid output structure like ``jax.grad`` or ``jax.lax.scan``\n then ``split_inputs`` and ``merge_inputs`` will not work. In this case, use the\n `Functional API `__.\n \"\"\"\n if isinstance(f, Missing):\n return functools.partial(split_inputs, ctxtag=ctxtag) # type: ignore[return-value]\n\n @graphlib.update_context(ctxtag)\n @functools.wraps(f)\n def split_inputs_wrapper(*args):\n pure_args = extract.to_tree(args, ctxtag=ctxtag)\n pure_args_out, pure_out = f(*pure_args)\n args_out, out = extract.from_tree(\n (pure_args_out, pure_out), ctxtag=ctxtag, is_inner=False\n )\n return out\n\n return split_inputs_wrapper # type: ignore\n\n@tp.overload\ndef merge_inputs(\n *,\n ctxtag: str = 'split_merge_inputs',\n) -> tp.Callable[[F], F]: ...\n@tp.overload\ndef merge_inputs(\n f: F,\n *,\n ctxtag: str = 'split_merge_inputs',\n) -> F: ...\ndef merge_inputs(\n f: F | Missing = MISSING,\n *,\n ctxtag: str = 'split_merge_inputs',\n) -> F | tp.Callable[[F], F]:\n \"\"\"Takes in a function that contains jax-compatible data structures in the\n inputs and outputs, and returns a function that replaces the jax-compatible\n data structures the corresponding graph nodes. Must be used in conjunction\n with :func:`split_inputs`.\n\n Args:\n f: The function to be transformed.\n ctxtag: The context tag to be used for the transformation. Defaults to\n 'split_merge_inputs'.\n\n Returns:\n The transformed function.\n\n For more information and examples, see :func:`split_inputs`.\n \"\"\"\n if isinstance(f, Missing):\n return functools.partial(merge_inputs, ctxtag=ctxtag) # type: ignore[return-value]\n\n @functools.wraps(f)\n def merge_inputs_wrapper(*pure_args):\n args = extract.from_tree(pure_args, ctxtag=ctxtag, is_inner=True)\n out = f(*args)\n args_out = extract.clear_non_graph_nodes(args)\n pure_args_out, pure_out = extract.to_tree((args_out, out), ctxtag=ctxtag)\n return pure_args_out, pure_out\n\n return merge_inputs_wrapper # type: ignore\n" }, { "path": "flax/nnx/transforms/iteration.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# pytype: skip-file\n\nfrom collections import deque\nimport dataclasses\nimport functools\nimport typing as tp\n\n\nfrom flax import struct\nfrom flax import typing\nfrom flax.core.frozen_dict import FrozenDict\nfrom flax.nnx import extract, filterlib, graphlib, spmd, variablelib\nfrom flax.nnx import statelib\nfrom flax.nnx.module import Module\nfrom flax.nnx.statelib import State\nfrom flax.nnx.transforms.transforms import (\n resolve_kwargs,\n _resolve_bound_callable,\n _raise_bound_method_error,\n)\nfrom flax.typing import Leaf, Missing, PytreeDeque\nimport jax\nimport jax.core\nimport jax.numpy as jnp\nimport jax.stages\nimport numpy as np\n\nA = tp.TypeVar('A')\nC = tp.TypeVar('C')\nB = tp.TypeVar('B')\nF = tp.TypeVar('F', bound=tp.Callable[..., tp.Any])\nG = tp.TypeVar('G', bound=tp.Callable[..., tp.Any])\nM = tp.TypeVar('M', bound=Module)\nMA = tp.TypeVar('MA', bound=Module)\nN = tp.TypeVar('N', bound=Module)\nT = tp.TypeVar('T')\nStrInt = tp.TypeVar('StrInt', str, int)\nAxisName = tp.Hashable\nLeaves = list[Leaf]\nIndex = int\n\n\nclass Carry:\n \"\"\"Helper class for :func:`flax.nnx.scan` function to mark input and output axis as carry.\n \"\"\"\n pass\n\n\n# -------------------------------\n# transform_metadata\n# -------------------------------\ndef _apply_axis_fn(\n tree: tp.Any,\n axes: tp.Any,\n metadata: tp.Mapping[str, tp.Any],\n axis_fn: tp.Callable[..., tp.Any],\n) -> None:\n is_leaf = lambda x: x is None or isinstance(x, variablelib.Variable)\n _, per_leaf_axes = extract.broadcast_prefix2(axes, tree, is_leaf=is_leaf)\n leaves = jax.tree_util.tree_leaves(tree, is_leaf=is_leaf)\n for leaf, axis in zip(leaves, per_leaf_axes):\n if (axis is None or isinstance(axis, int)) and isinstance(\n leaf, variablelib.Variable\n ):\n axis_fn(leaf, axis, metadata)\n\n\n@tp.overload\ndef transform_metadata(\n *,\n in_axes: tp.Any = 0,\n out_axes: tp.Any = 0,\n partition: str,\n graph: bool | None = None,\n) -> tp.Callable[[F], F]:\n ...\n\n\n@tp.overload\ndef transform_metadata(\n f: F,\n *,\n in_axes: tp.Any = 0,\n out_axes: tp.Any = 0,\n graph: bool | None = None,\n partition: str,\n) -> F:\n ...\n\n\ndef transform_metadata(\n f: F | type[Missing] = Missing,\n *,\n in_axes: tp.Any = 0,\n out_axes: tp.Any = 0,\n graph: bool | None = None,\n partition: str,\n) -> F | tp.Callable[[F], F]:\n if f is Missing:\n return functools.partial(\n transform_metadata,\n in_axes=in_axes,\n out_axes=out_axes,\n partition=partition,\n graph=graph,\n ) # type: ignore[return-value]\n\n if graph is None:\n graph = graphlib.set_graph_mode.current_value()\n\n metadata: tp.Mapping[str, tp.Any] = {\n spmd.PARTITION_NAME: partition,\n }\n if graph:\n extract.check_prefix(in_axes, 'in_axes', 'transform_metadata')\n extract.check_prefix(out_axes, 'out_axes', 'transform_metadata')\n\n @functools.wraps(f)\n def wrapper(*in_args, **in_kwargs):\n in_args = resolve_kwargs(f, in_args, in_kwargs)\n if graph:\n in_args = extract.to_tree2(in_args, prefix=in_axes)\n extract.check_no_aliases('transform_metadata', args=in_args)\n args = graphlib.clone(in_args, graph=graph)\n _apply_axis_fn(args, in_axes, metadata, spmd.remove_axis)\n updates, snapshot = extract.updates_and_snapshot(args)\n if graph:\n args = extract.from_tree2(args)\n out = f(*args)\n if graph:\n out = extract.to_tree2(out, prefix=out_axes)\n extract.check_no_aliases('transform_metadata', args=updates, out=out)\n _apply_axis_fn(args, in_axes, metadata, spmd.add_axis)\n _apply_axis_fn(out, out_axes, metadata, spmd.add_axis)\n updates = extract.mask_variable_updates(updates, snapshot)\n extract.apply_variable_updates(in_args, updates)\n if graph:\n out = extract.from_tree2(out)\n return out\n\n return wrapper # type: ignore[return-value]\n\n\n# -------------------------------\n# vmap\n# -------------------------------\n\n\nclass StateAxes(extract.PrefixMapping, tp.Mapping):\n\n def __init__(\n self,\n filter_axes: (\n statelib.State\n | tp.Mapping[filterlib.Filter, Index | type[Carry] | None]\n | tp.Iterable[tuple[filterlib.Filter, Index | type[Carry] | None]]\n ),\n /,\n ):\n if isinstance(filter_axes, statelib.State):\n filter_axes = statelib.create_path_filters(filter_axes) # type: ignore\n\n iterable = tuple(\n filter_axes.items()\n if isinstance(filter_axes, tp.Mapping)\n else filter_axes\n )\n self._filters = tuple(filter for filter, _ in iterable)\n self._axes = tuple(axis for _, axis in iterable)\n\n @property\n def filters(self) -> tuple[filterlib.Filter, ...]:\n return self._filters\n\n @property\n def axes(self) -> tuple[Index | type[Carry] | None, ...]:\n return self._axes\n\n def map_prefix(\n self, path: typing.PathParts, variable: variablelib.Variable\n ) -> tp.Any:\n for filter, axis in zip(self.filters, self.axes):\n predicate = filterlib.to_predicate(filter)\n if predicate(path, variable):\n return axis\n raise ValueError(f'No axis found for {path=}, {variable=}')\n\n def __repr__(self):\n return f'StateAxes({dict(self.items())})'\n\n def items(self):\n return zip(self.filters, self.axes)\n\n def __getitem__(self, key):\n return self.axes[self.filters.index(key)]\n\n def __iter__(self):\n return iter(self.filters)\n\n def __len__(self):\n return len(self.filters)\n\n def __eq__(self, other):\n return (\n isinstance(other, StateAxes)\n and self.filters == other.filters\n and self.axes == other.axes\n )\n\n def __hash__(self):\n return hash((self.filters, self.axes))\n\n\nAxisFn = tp.Callable[\n [graphlib.GraphState | variablelib.Variable, int, tp.Mapping],\n graphlib.GraphState | variablelib.Variable,\n]\n\n\ndef _update_variable_sharding_metadata(\n tree, transform_metadata, axis_fn: AxisFn\n):\n def _update_axes_fn(node_states):\n if isinstance(node_states, extract.NodeStates) and isinstance(\n node_states.metadata, (StateAxes, int)\n ):\n if isinstance(node_states.metadata, int):\n state = node_states.state\n assert isinstance(state, State | variablelib.Variable)\n state = axis_fn(state, node_states.metadata, transform_metadata)\n return node_states.replace(states=(state,))\n else:\n states_out: list[graphlib.GraphState | variablelib.Variable] = []\n for state, axis in zip(node_states.states, node_states.metadata.axes):\n assert isinstance(state, graphlib.State | variablelib.Variable)\n if isinstance(axis, int):\n state = axis_fn(state, axis, transform_metadata)\n states_out.append(state)\n return node_states.replace(states=tuple(states_out))\n return node_states\n\n return jax.tree.map(\n _update_axes_fn, tree, is_leaf=lambda x: isinstance(x, extract.NodeStates)\n )\n\n\ndef _vmap_split_fn(ctx: graphlib.SplitContext, path, prefix, x):\n if isinstance(prefix, StateAxes):\n return extract.NodeStates.from_split(\n *ctx.split(x, *prefix.filters), metadata=prefix\n )\n return extract.NodeStates.from_split(*ctx.split(x), metadata=prefix)\n\n\n@dataclasses.dataclass(eq=False)\nclass SimpleVmapFn:\n f: tp.Callable[..., tp.Any]\n graph: bool\n out_axes: tp.Any\n\n def __post_init__(self):\n functools.update_wrapper(self, self.f, updated=())\n\n @extract.treemap_copy_args\n def __call__(self, *args, **kwargs):\n updates, snapshot = extract.updates_and_snapshot((args, kwargs))\n if self.graph:\n args, kwargs = extract.from_tree2((args, kwargs))\n out = self.f(*args, **kwargs)\n if self.graph:\n out = extract.to_tree2(out, prefix=self.out_axes)\n extract.check_no_aliases('vmap', args=updates[0], kwargs=updates[1], out=out)\n updates = extract.mask_variable_updates(updates, snapshot)\n return out, updates\n\n\n@dataclasses.dataclass(eq=False)\nclass SimplePmapFn:\n f: tp.Callable[..., tp.Any]\n graph: bool\n out_axes: tp.Any\n\n def __post_init__(self):\n functools.update_wrapper(self, self.f, updated=())\n\n @extract.treemap_copy_args\n def __call__(self, *args, **kwargs):\n updates, snapshot = extract.updates_and_snapshot((args, kwargs))\n if self.graph:\n args, kwargs = extract.from_tree2((args, kwargs))\n out = self.f(*args, **kwargs)\n if self.graph:\n out = extract.to_tree2(out, prefix=self.out_axes)\n extract.check_no_aliases('pmap', args=updates[0], kwargs=updates[1], out=out)\n updates = extract.mask_variable_updates(updates, snapshot)\n return out, updates\n\n\n@dataclasses.dataclass(eq=False)\nclass VmapFn:\n f: tp.Callable[..., tp.Any]\n transform_metadata: tp.Mapping[str, tp.Any]\n in_axes: tp.Any\n out_axes: tp.Any\n\n def __post_init__(self):\n functools.update_wrapper(self, self.f)\n\n def __call__(self, *pure_args: tuple[tp.Any, ...]):\n if spmd.PARTITION_NAME in self.transform_metadata:\n pure_args = _update_variable_sharding_metadata(\n pure_args, self.transform_metadata, spmd.remove_axis\n )\n args = extract.from_tree(pure_args, ctxtag='vmap', is_inner=True)\n\n out = self.f(*args)\n\n args_out = extract.clear_non_graph_nodes(args)\n pure_args_out, pure_out = extract.to_tree(\n (args_out, out),\n prefix=(self.in_axes, self.out_axes),\n split_fn=_vmap_split_fn,\n ctxtag='vmap',\n )\n if spmd.PARTITION_NAME in self.transform_metadata:\n pure_args_out, pure_out = _update_variable_sharding_metadata(\n (pure_args_out, pure_out), self.transform_metadata, spmd.add_axis\n )\n return pure_args_out, pure_out\n\n\n@tp.overload\ndef vmap(\n *,\n in_axes: int | None | tp.Sequence[tp.Any] = 0,\n out_axes: tp.Any = 0,\n axis_name: AxisName | None = None,\n axis_size: int | None = None,\n spmd_axis_name: AxisName | tuple[AxisName, ...] | None = None,\n # nnx specific\n transform_metadata: tp.Mapping[str, tp.Any] = FrozenDict({}),\n graph: bool | None = None,\n graph_updates: bool | None = None,\n) -> tp.Callable[[F], F]:\n ...\n\n\n@tp.overload\ndef vmap(\n f: F,\n *,\n in_axes: int | None | tp.Sequence[tp.Any] = 0,\n out_axes: tp.Any = 0,\n axis_name: AxisName | None = None,\n axis_size: int | None = None,\n spmd_axis_name: AxisName | tuple[AxisName, ...] | None = None,\n # nnx specific\n transform_metadata: tp.Mapping[str, tp.Any] = FrozenDict({}),\n graph: bool | None = None,\n graph_updates: bool | None = None,\n) -> F:\n ...\n\n\ndef vmap(\n f: F | type[Missing] = Missing,\n *,\n in_axes: int | None | tp.Sequence[tp.Any] = 0,\n out_axes: tp.Any = 0,\n axis_name: AxisName | None = None,\n axis_size: int | None = None,\n spmd_axis_name: AxisName | tuple[AxisName, ...] | None = None,\n # nnx specific\n transform_metadata: tp.Mapping[str, tp.Any] = FrozenDict({}),\n graph: bool | None = None,\n graph_updates: bool | None = None,\n) -> F | tp.Callable[[F], F]:\n \"\"\"Reference-aware version of `jax.vmap `__.\n\n Args:\n f: Function to be mapped over additional axes.\n in_axes: An integer, None, or sequence of values specifying which input\n array axes to map over (see `jax.vmap\n `__). In\n addition to integers and None, :class:`StateAxes` can be used to control\n how graph nodes like Modules are vectorized by specifying the axes to be\n applied to substates of the graph node given a `Filter\n `__.\n out_axes: An integer, None, or pytree indicating where the mapped axis\n should appear in the output (see `jax.vmap\n `__).\n axis_name: Optional, a hashable Python object used to identify the mapped\n axis so that parallel collectives can be applied.\n axis_size: Optional, an integer indicating the size of the axis to be\n mapped. If not provided, the mapped axis size is inferred from arguments.\n graph: If ``True`` (default), uses graph-mode which supports the full\n NNX feature set including shared references and reference semantics.\n If ``False``, uses tree-mode which treats Modules as regular JAX\n pytrees, avoiding the overhead of the graph protocol. Tree-mode does\n not support ``StateAxes`` or shared ``Variable`` references.\n graph_updates: If ``True``, propagates updates on graph structure\n that happen inside the transform to the input graphs, has no\n effect when ``graph=False``. When ``False``, using ``StateAxes``\n is not supported.\n\n Returns:\n Batched/vectorized version of ``f`` with arguments that correspond to\n those of ``f``, but with extra array axes at positions indicated by\n ``in_axes``, and a return value that corresponds to that of ``f``, but\n with extra array axes at positions indicated by ``out_axes``.\n\n Example::\n\n >>> from flax import nnx\n >>> from jax import random, numpy as jnp\n ...\n >>> model = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n >>> x = jnp.ones((5, 2))\n ...\n >>> @nnx.vmap(in_axes=(None, 0), out_axes=0)\n ... def forward(model, x):\n ... return model(x)\n ...\n >>> y = forward(model, x)\n >>> y.shape\n (5, 3)\n\n >>> class LinearEnsemble(nnx.Module):\n ... def __init__(self, num, rngs):\n ... self.w = nnx.Param(jax.random.uniform(rngs(), (num, 2, 3)))\n ...\n >>> model = LinearEnsemble(5, rngs=nnx.Rngs(0))\n >>> x = jnp.ones((2,))\n ...\n >>> @nnx.vmap(in_axes=(0, None), out_axes=0)\n ... def forward(model, x):\n ... return x @ model.w\n ...\n >>> y = forward(model, x)\n >>> y.shape\n (5, 3)\n\n To control control how graph node substates are vectorized, ``StateAxes``\n can be passed to ``in_axes`` and ``out_axes`` specifying the axes to be\n applied to each substate given a filter. The following example shows how to\n share the parameters between the ensemble members which keeping different\n batch statistics and dropout random state::\n\n >>> class Foo(nnx.Module):\n ... def __init__(self):\n ... self.a = nnx.Param(jnp.arange(4))\n ... self.b = nnx.BatchStat(jnp.arange(4))\n ...\n >>> state_axes = nnx.StateAxes({nnx.Param: 0, nnx.BatchStat: None})\n >>> @nnx.vmap(in_axes=(state_axes,), out_axes=0)\n ... def mul(foo):\n ... return foo.a * foo.b\n ...\n >>> foo = Foo()\n >>> y = mul(foo)\n >>> y\n Array([[0, 0, 0, 0],\n [0, 1, 2, 3],\n [0, 2, 4, 6],\n [0, 3, 6, 9]], dtype=int32)\n \"\"\"\n if graph is None:\n graph = graphlib.set_graph_mode.current_value()\n if graph_updates is None:\n graph_updates = graphlib.set_graph_updates.current_value()\n if f is Missing:\n return functools.partial(\n vmap,\n in_axes=in_axes,\n out_axes=out_axes,\n axis_name=axis_name,\n axis_size=axis_size,\n spmd_axis_name=spmd_axis_name,\n transform_metadata=transform_metadata,\n graph=graph,\n graph_updates=graph_updates,\n ) # type: ignore[return-value]\n\n f_unbound, _, was_bound = _resolve_bound_callable(f)\n if was_bound:\n _raise_bound_method_error('vmap')\n\n if not graph or not graph_updates:\n if any(isinstance(x, StateAxes) for x in jax.tree.leaves(in_axes)):\n raise ValueError(\n '`in_axes` cannot contain `StateAxes` objects '\n 'when `graph=False`. '\n + graphlib._tree_mode_suggestion_transform('vmap')\n )\n if any(isinstance(x, StateAxes) for x in jax.tree.leaves(out_axes)):\n raise ValueError(\n '`out_axes` cannot contain `StateAxes` objects '\n 'when `graph=False`. '\n + graphlib._tree_mode_suggestion_transform('vmap')\n )\n if graph:\n extract.check_prefix(in_axes, 'in_axes', 'vmap')\n extract.check_prefix(out_axes, 'out_axes', 'vmap')\n\n vmapped_fn = jax.vmap(\n SimpleVmapFn(f_unbound, graph=graph, out_axes=out_axes),\n in_axes=in_axes,\n out_axes=(out_axes, (in_axes, 0)),\n axis_name=axis_name,\n axis_size=axis_size,\n spmd_axis_name=spmd_axis_name,\n )\n\n @functools.wraps(f_unbound)\n def simple_vmap_wrapper(*args, **kwargs):\n if graph:\n args, kwargs = extract.to_tree2(\n (args, kwargs),\n prefix=(in_axes, None)\n if in_axes is not None\n else None,\n check_aliasing=in_axes is not None,\n )\n extract.check_no_aliases('vmap', args=args, kwargs=kwargs)\n out, updates = vmapped_fn(*args, **kwargs)\n extract.apply_variable_updates((args, kwargs), updates)\n if graph:\n out = extract.from_tree2(out)\n return out\n\n return simple_vmap_wrapper # type: ignore[return-value]\n\n\n jax_in_axes = jax.tree.map(\n lambda x: extract.NodeStates.from_prefixes(x.axes, metadata=x)\n if isinstance(x, StateAxes)\n else x,\n in_axes,\n )\n jax_out_axes = jax.tree.map(\n lambda x: extract.NodeStates.from_prefixes(x.axes, metadata=x)\n if isinstance(x, StateAxes)\n else x,\n out_axes,\n )\n vmapped_fn = jax.vmap( # type: ignore[assignment]\n VmapFn(f_unbound, transform_metadata, in_axes, out_axes),\n in_axes=jax_in_axes,\n out_axes=(jax_in_axes, jax_out_axes),\n axis_name=axis_name,\n axis_size=axis_size,\n spmd_axis_name=spmd_axis_name,\n )\n\n @functools.wraps(f)\n @graphlib.update_context('vmap')\n def vmap_wrapper(*args, **kwargs):\n args = resolve_kwargs(f, args, kwargs)\n pure_args = extract.to_tree(\n args, prefix=in_axes, split_fn=_vmap_split_fn, ctxtag='vmap'\n )\n pure_args_out, pure_out = vmapped_fn(*pure_args)\n _args_out, out = extract.from_tree(\n (pure_args_out, pure_out), ctxtag='vmap', is_inner=False\n )\n return out\n\n return vmap_wrapper # type: ignore\n\n\n# -------------------------------\n# pmap\n# -------------------------------\n\n\n@dataclasses.dataclass(eq=False)\nclass PmapFn:\n f: tp.Callable[..., tp.Any]\n transform_metadata: tp.Mapping[str, tp.Any]\n in_axes: tp.Any\n out_axes: tp.Any\n\n def __post_init__(self):\n functools.update_wrapper(self, self.f)\n\n def __call__(self, *pure_args: tuple[tp.Any, ...]):\n if spmd.PARTITION_NAME in self.transform_metadata:\n pure_args = _update_variable_sharding_metadata(\n pure_args, self.transform_metadata, spmd.remove_axis\n )\n args = extract.from_tree(pure_args, ctxtag='pmap', is_inner=True)\n\n out = self.f(*args)\n\n args_out = extract.clear_non_graph_nodes(args)\n pure_args_out, pure_out = extract.to_tree(\n (args_out, out),\n prefix=(self.in_axes, self.out_axes),\n split_fn=_vmap_split_fn,\n ctxtag='pmap',\n )\n if spmd.PARTITION_NAME in self.transform_metadata:\n pure_args_out, pure_out = _update_variable_sharding_metadata(\n (pure_args_out, pure_out), self.transform_metadata, spmd.add_axis\n )\n return pure_args_out, pure_out\n\n\n@tp.overload\ndef pmap(\n *,\n axis_name: AxisName | None = None,\n in_axes: tp.Any = 0,\n out_axes: tp.Any = 0,\n static_broadcasted_argnums: int | tp.Iterable[int] = (),\n devices: tp.Sequence[jax.Device] | None = None, # noqa: F811\n backend: str | None = None,\n axis_size: int | None = None,\n donate_argnums: int | tp.Iterable[int] = (),\n # nnx specific\n transform_metadata: tp.Mapping[str, tp.Any] = FrozenDict({}),\n graph: bool | None = None,\n graph_updates: bool | None = None,\n) -> tp.Callable[[F], F]:\n ...\n\n\n@tp.overload\ndef pmap(\n f: F,\n *,\n axis_name: AxisName | None = None,\n in_axes: tp.Any = 0,\n out_axes: tp.Any = 0,\n static_broadcasted_argnums: int | tp.Iterable[int] = (),\n devices: tp.Sequence[jax.Device] | None = None, # noqa: F811\n backend: str | None = None,\n axis_size: int | None = None,\n donate_argnums: int | tp.Iterable[int] = (),\n # nnx specific\n transform_metadata: tp.Mapping[str, tp.Any] = FrozenDict({}),\n graph: bool | None = None,\n graph_updates: bool | None = None,\n) -> F:\n ...\n\n\ndef pmap(\n f: F | type[Missing] = Missing,\n *,\n axis_name: AxisName | None = None,\n in_axes: tp.Any = 0,\n out_axes: tp.Any = 0,\n static_broadcasted_argnums: int | tp.Iterable[int] = (),\n devices: tp.Sequence[jax.Device] | None = None, # noqa: F811\n backend: str | None = None,\n axis_size: int | None = None,\n donate_argnums: int | tp.Iterable[int] = (),\n # nnx specific\n transform_metadata: tp.Mapping[str, tp.Any] = FrozenDict({}),\n graph: bool | None = None,\n graph_updates: bool | None = None,\n) -> F | tp.Callable[[F], F]:\n \"\"\"Reference-aware version of `jax.pmap `__.\n\n Args:\n f: Function to be mapped over argument axes. Its arguments and return\n value should be arrays, scalars, graph nodes, or (nested) standard\n Python containers (tuple/list/dict) thereof. Positional arguments\n indicated by ``static_broadcasted_argnums`` can be anything at all,\n provided they are hashable and have an equality operation defined.\n axis_name: Optional, a hashable Python object used to identify the mapped\n axis so that parallel collectives can be applied.\n in_axes: A non-negative integer, None, or nested Python container thereof\n that specifies which axes of positional arguments to map over. Arguments\n passed as keywords are always mapped over their leading axis (i.e. axis\n index 0). In addition to integers and None, :class:`StateAxes` can be\n used to control how graph nodes like Modules are vectorized by specifying\n the axes to be applied to substates of the graph node given a `Filter\n `__.\n out_axes: A non-negative integer, None, or nested Python container thereof\n indicating where the mapped axis should appear in the output. All outputs\n with a mapped axis must have a non-None ``out_axes`` specification.\n static_broadcasted_argnums: An int or collection of ints specifying which\n positional arguments to treat as static (compile-time constant).\n Operations that only depend on static arguments will be constant-folded.\n Calling the pmapped function with different values for these constants\n will trigger recompilation. If the pmapped function is called with fewer\n positional arguments than indicated by ``static_broadcasted_argnums`` then\n an error is raised. Each of the static arguments will be broadcasted to\n all devices. Arguments that are not arrays or containers thereof must be\n marked as static. Defaults to ().\n Static arguments must be hashable, meaning both ``__hash__`` and\n ``__eq__`` are implemented, and should be immutable.\n devices: This is an experimental feature and the API is likely to change.\n Optional, a sequence of Devices to map over. (Available devices can be\n retrieved via jax.devices()). Must be given identically for each process\n in multi-process settings (and will therefore include devices across\n processes). If specified, the size of the mapped axis must be equal to\n the number of devices in the sequence local to the given process. Nested\n ``pmap`` s with ``devices`` specified in either the inner or outer\n ``pmap`` are not yet supported.\n backend: This is an experimental feature and the API is likely to change.\n Optional, a string representing the XLA backend. 'cpu', 'gpu', or 'tpu'.\n axis_size: Optional; the size of the mapped axis.\n donate_argnums: Specify which positional argument buffers are \"donated\" to\n the computation. It is safe to donate argument buffers if you no longer\n need them once the computation has finished. In some cases XLA can make\n use of donated buffers to reduce the amount of memory needed to perform a\n computation, for example recycling one of your input buffers to store a\n result. You should not reuse buffers that you donate to a computation,\n JAX will raise an error if you try to. Note that donate_argnums only\n work for positional arguments, and keyword arguments will not be donated.\n transform_metadata: Optional mapping of metadata for the transform.\n graph: if True, use graph-mode (default). If False, use tree-mode.\n If None, uses the value of ``nnx_graph_mode`` config.\n graph_updates: If ``True``, propagates updates on graph structure\n that happen inside the transform to the input graphs, has no\n effect when ``graph=False``. When ``False``, using ``StateAxes``\n is not supported.\n\n Returns:\n A parallelized version of ``f`` with arguments that correspond to those of\n ``f`` but with extra array axes at positions indicated by ``in_axes`` and\n with output that has an additional leading array axis (with the same size).\n \"\"\"\n if graph is None:\n graph = graphlib.set_graph_mode.current_value()\n if graph_updates is None:\n graph_updates = graphlib.set_graph_updates.current_value()\n if f is Missing:\n return functools.partial(\n pmap,\n axis_name=axis_name,\n in_axes=in_axes,\n out_axes=out_axes,\n static_broadcasted_argnums=static_broadcasted_argnums,\n devices=devices,\n backend=backend,\n axis_size=axis_size,\n donate_argnums=donate_argnums,\n transform_metadata=transform_metadata,\n graph=graph,\n graph_updates=graph_updates,\n ) # type: ignore[return-value]\n\n f_unbound, _, was_bound = _resolve_bound_callable(f)\n if was_bound:\n _raise_bound_method_error('pmap')\n\n if not graph or not graph_updates:\n if any(isinstance(x, StateAxes) for x in jax.tree.leaves(in_axes)):\n raise ValueError(\n '`in_axes` cannot contain `StateAxes` objects '\n 'when `graph=False`. '\n + graphlib._tree_mode_suggestion_transform('pmap')\n )\n if any(isinstance(x, StateAxes) for x in jax.tree.leaves(out_axes)):\n raise ValueError(\n '`out_axes` cannot contain `StateAxes` objects '\n 'when `graph=False`. '\n + graphlib._tree_mode_suggestion_transform('pmap')\n )\n if graph:\n extract.check_prefix(in_axes, 'in_axes', 'pmap')\n extract.check_prefix(out_axes, 'out_axes', 'pmap')\n\n pmapped_fn = jax.pmap(\n SimplePmapFn(f_unbound, graph=graph, out_axes=out_axes),\n axis_name=axis_name,\n in_axes=in_axes,\n out_axes=(out_axes, (in_axes, 0)),\n static_broadcasted_argnums=static_broadcasted_argnums,\n devices=devices,\n backend=backend,\n axis_size=axis_size,\n donate_argnums=donate_argnums,\n )\n\n @functools.wraps(f_unbound)\n def simple_pmap_wrapper(*args, **kwargs):\n if graph:\n args, kwargs = extract.to_tree2(\n (args, kwargs),\n prefix=(in_axes, None)\n if in_axes is not None\n else None,\n check_aliasing=in_axes is not None,\n )\n extract.check_no_aliases('pmap', args=args, kwargs=kwargs)\n out, updates = pmapped_fn(*args, **kwargs)\n extract.apply_variable_updates((args, kwargs), updates)\n if graph:\n out = extract.from_tree2(out)\n return out\n\n return simple_pmap_wrapper # type: ignore[return-value]\n\n jax_in_axes = jax.tree.map(\n lambda x: extract.NodeStates.from_prefixes(x.axes, metadata=x)\n if isinstance(x, StateAxes)\n else x,\n in_axes,\n )\n jax_out_axes = jax.tree.map(\n lambda x: extract.NodeStates.from_prefixes(x.axes, metadata=x)\n if isinstance(x, StateAxes)\n else x,\n out_axes,\n )\n pmapped_fn = jax.pmap(\n PmapFn(f_unbound, transform_metadata, in_axes, out_axes),\n axis_name=axis_name,\n in_axes=jax_in_axes,\n out_axes=(jax_in_axes, jax_out_axes),\n static_broadcasted_argnums=static_broadcasted_argnums,\n devices=devices,\n backend=backend,\n axis_size=axis_size,\n donate_argnums=donate_argnums,\n )\n\n @functools.wraps(f)\n @graphlib.update_context('pmap')\n def vmap_wrapper(*args):\n pure_args = extract.to_tree(\n args, prefix=in_axes, split_fn=_vmap_split_fn, ctxtag='pmap'\n )\n pure_args_out, pure_out = pmapped_fn(*pure_args)\n _args_out, out = extract.from_tree(\n (pure_args_out, pure_out), ctxtag='pmap', is_inner=False\n )\n return out\n\n return vmap_wrapper # type: ignore\n\n\n# -------------------------------\n# scan\n# -------------------------------\n\n\nclass Broadcasted(struct.PyTreeNode):\n data: tp.Any\n\ndef _get_carry_argnum(axes, is_in_axes: bool):\n if axes is Carry:\n return 'all'\n elif isinstance(axes, int) or axes is None:\n return None\n\n obj_repr = 'in_axes' if is_in_axes else 'out_axes'\n carry_argnum: int | None = None\n prev_key: tp.Any = None\n for key, x in jax.tree_util.tree_leaves_with_path(axes):\n if x is not Carry:\n continue\n assert isinstance(key[0], jax.tree_util.SequenceKey)\n i = key[0].idx\n if len(key) >= 2:\n raise ValueError(\n f'Carry must at the top-level, it cannot be nested. Found {axes=}'\n )\n if carry_argnum is not None:\n raise ValueError(\n f'Found multiple Carry definitions at '\n f'{obj_repr}{jax.tree_util.keystr(prev_key)} and '\n f'{obj_repr}{jax.tree_util.keystr(key)}'\n )\n carry_argnum = i\n prev_key = key\n\n return carry_argnum\n\n\ndef _check_out_axes(out_axes):\n for key, x in jax.tree_util.tree_leaves_with_path(\n out_axes, is_leaf=lambda x: x is None\n ):\n if x is None:\n raise ValueError(\n f'Cannot broadcast output state. '\n f'Got out_axes=None at: out_axes{jax.tree_util.keystr(key)}'\n )\n elif isinstance(x, StateAxes):\n for filter, value in x.items():\n if value is None:\n raise ValueError(\n f'Cannot broadcast output state. '\n f'Got StateAxes({{{filter}: None}}) at: out_axes'\n f'{jax.tree_util.keystr(key)}'\n )\n elif value is Carry:\n raise ValueError(\n f'Cannot carry output state. '\n f'Got StateAxes({{{filter}: Carry}}) at: out_axes'\n f'{jax.tree_util.keystr(key)}'\n )\ndef _check_carry_same_references(carry_arg, carry_arg_out):\n def check_carry_same_references(key_path, arg, out):\n if (\n not isinstance(arg, jax.Array) or not isinstance(out, jax.Array)\n ) and arg is not out:\n raise ValueError(\n 'Carry references must be the same between iterations. '\n f'Got {arg=} with id={id(arg)} and {out=} with id={id(out)} '\n f'at carry{jax.tree_util.keystr(key_path)}'\n )\n\n jax.tree_util.tree_map_with_path(\n check_carry_same_references,\n carry_arg,\n carry_arg_out,\n is_leaf=lambda x: graphlib.is_graph_node(x)\n and not isinstance(x, variablelib.Variable),\n )\n\n\ndef _scan_split_in(\n carry_deque: PytreeDeque[list[State | variablelib.Variable]],\n graphdefs_deque: PytreeDeque[graphlib.GraphDef],\n broadcast_deque: PytreeDeque[list[State | variablelib.Variable]],\n broadcast_arrays: PytreeDeque[Broadcasted],\n /,\n ctx: graphlib.SplitContext,\n path,\n prefix,\n x,\n):\n if graphlib.is_graph_node(x) or isinstance(x, variablelib.Variable):\n vectorized_states: list[State | variablelib.Variable] = []\n carry_states: list[State | variablelib.Variable] = []\n broadcast_states: list[State | variablelib.Variable] = []\n if isinstance(prefix, StateAxes):\n graphdef, *states = ctx.split(x, *prefix.filters)\n\n for state, axis in zip(states, prefix.axes):\n if axis is None:\n broadcast_states.append(state)\n elif isinstance(axis, int):\n if axis != 0:\n state = jax.tree.map(lambda x: jnp.moveaxis(x, axis, 0), state)\n vectorized_states.append(state)\n else: # axis is Carry\n carry_states.append(state)\n if not vectorized_states:\n vectorized_states.append(State({}))\n carry_deque.append(carry_states)\n graphdefs_deque.append(graphdef)\n broadcast_deque.append(broadcast_states)\n return extract.NodeStates.from_split(\n None, *vectorized_states, metadata=prefix\n )\n elif isinstance(prefix, int):\n graphdef, state = ctx.split(x)\n if prefix != 0:\n state = jax.tree.map(lambda x: jnp.moveaxis(x, prefix, 0), state)\n vectorized_states.append(state)\n elif prefix is None:\n graphdef, state = ctx.split(x)\n broadcast_states.append(state)\n vectorized_states.append(State({}))\n elif prefix is Carry:\n graphdef, state = ctx.split(x)\n carry_states.append(state)\n vectorized_states.append(State({}))\n else:\n raise ValueError(\n f'Invalid axes {prefix} args{jax.tree_util.keystr(path)}'\n )\n\n if not vectorized_states:\n vectorized_states.append(State({}))\n carry_deque.append(carry_states)\n graphdefs_deque.append(graphdef)\n broadcast_deque.append(broadcast_states)\n return extract.NodeStates.from_split(\n None, *vectorized_states, metadata=prefix\n )\n else:\n if isinstance(prefix, StateAxes):\n raise ValueError(\n 'Cannot use StateAxes on non-graph nodes, '\n f'found {prefix} args{jax.tree_util.keystr(path)}'\n )\n elif prefix is Carry:\n return x\n elif prefix is None:\n broadcast_arrays.append(Broadcasted(x))\n return Broadcasted(None)\n elif isinstance(prefix, int):\n if not isinstance(x, (jax.Array, np.ndarray)):\n raise ValueError(\n f'Expected an array, got {type(x).__name__} args'\n f'{jax.tree_util.keystr(path)}'\n )\n if prefix != 0:\n x = jnp.moveaxis(x, prefix, 0)\n return x\n else:\n raise ValueError(\n f'Invalid axes {prefix} args{jax.tree_util.keystr(path)}'\n )\n\n\ndef _scan_split_out(\n carry_deque: PytreeDeque[list[State | variablelib.Variable]],\n graphdefs_deque: PytreeDeque[graphlib.GraphDef],\n broadcast_deque: PytreeDeque[list[State | variablelib.Variable]],\n /,\n ctx: graphlib.SplitContext,\n path: extract.KeyPath,\n prefix,\n x,\n):\n assert isinstance(path[0], jax.tree_util.SequenceKey)\n is_input_arg = path[0].idx == 0\n\n if graphlib.is_graph_node(x) or isinstance(x, variablelib.Variable):\n vectorized_states: list[State | variablelib.Variable] = []\n carry_states: list[State | variablelib.Variable] = []\n broadcast_states: list[State | variablelib.Variable] = []\n if isinstance(prefix, StateAxes):\n graphdef, *states = ctx.split(x, *prefix.filters)\n\n for state, filter, axis in zip(states, prefix.filters, prefix.axes):\n if axis is None:\n assert is_input_arg # validated by _check_out_axes\n broadcast_states.append(state)\n elif isinstance(axis, int):\n vectorized_states.append(state)\n elif axis is Carry:\n assert is_input_arg # validated by _check_out_axes\n carry_states.append(state)\n else:\n obj_repr = 'args' if is_input_arg else 'out'\n raise ValueError(\n f'Invalid axes {axis} for filter {filter} at '\n f'{obj_repr}{jax.tree_util.keystr(path)}'\n )\n\n if not vectorized_states:\n vectorized_states.append(State({}))\n if is_input_arg:\n carry_deque.append(carry_states)\n broadcast_deque.append(broadcast_states)\n graphdefs_deque.append(graphdef)\n return extract.NodeStates.from_split(\n None, *vectorized_states, metadata=prefix\n )\n elif isinstance(prefix, int):\n graphdef, state = ctx.split(x)\n vectorized_states.append(state)\n elif prefix is None:\n assert is_input_arg # validated by _check_out_axes\n graphdef, state = ctx.split(x)\n broadcast_states.append(state)\n vectorized_states.append(State({}))\n elif prefix is Carry:\n assert is_input_arg # validated by _check_out_axes\n graphdef, state = ctx.split(x)\n carry_states.append(state)\n vectorized_states.append(State({}))\n else:\n obj_repr = 'args' if is_input_arg else 'out'\n raise ValueError(\n f'Invalid axes {prefix} at {obj_repr}{jax.tree_util.keystr(path)}'\n )\n if not vectorized_states:\n vectorized_states.append(State({}))\n if is_input_arg:\n carry_deque.append(carry_states)\n broadcast_deque.append(broadcast_states)\n graphdefs_deque.append(graphdef)\n return extract.NodeStates.from_split(\n None, *vectorized_states, metadata=prefix\n )\n else:\n if isinstance(prefix, StateAxes):\n obj_repr = 'args' if is_input_arg else 'out'\n raise ValueError(\n 'Cannot use StateAxes on non-graph nodes, '\n f'found {prefix} at {obj_repr}{jax.tree_util.keystr(path)}'\n )\n elif prefix is Carry:\n return x\n elif prefix is None:\n assert not is_input_arg # validated by _check_out_axes\n return Broadcasted(None)\n elif isinstance(prefix, int):\n return x\n else:\n obj_repr = 'args' if is_input_arg else 'out'\n raise ValueError(\n f'Invalid axes {prefix} at {obj_repr}{jax.tree_util.keystr(path)}'\n )\n\n\ndef _scan_merge_in(\n carry_deque: PytreeDeque[list[State]],\n graphdefs_deque: PytreeDeque[graphlib.GraphDef],\n broadcast_deque: PytreeDeque[list[State]],\n broadcast_arrays: PytreeDeque[Broadcasted],\n /,\n ctx: graphlib.MergeContext,\n path,\n prefix,\n x,\n):\n if isinstance(x, extract.NodeStates):\n carry_states = carry_deque.popleft()\n broadcast_states = broadcast_deque.popleft()\n graphdef = graphdefs_deque.popleft()\n return ctx.merge(graphdef, *x.states, *carry_states, *broadcast_states)\n elif isinstance(x, Broadcasted):\n assert x.data is None\n return broadcast_arrays.popleft().data\n else:\n return x\n\n\ndef _scan_merge_out(\n carry_deque: PytreeDeque[list[State]],\n graphdefs_deque: PytreeDeque[graphlib.GraphDef],\n broadcast_deque: PytreeDeque[list[State]],\n /,\n ctx: graphlib.MergeContext,\n path,\n prefix,\n x,\n):\n assert isinstance(path[0], jax.tree_util.SequenceKey)\n is_input_arg = path[0].idx == 0\n\n if isinstance(x, extract.NodeStates):\n states: list[State] = []\n graphdef = graphdefs_deque.popleft()\n if is_input_arg:\n carry_states = deque(carry_deque.popleft())\n broadcast_states = deque(broadcast_deque.popleft())\n else:\n carry_states = deque[State]()\n broadcast_states = deque[State]()\n if isinstance(prefix, StateAxes):\n vectorized_states = deque(x.states)\n for axis in prefix.axes:\n if isinstance(axis, int):\n state = vectorized_states.popleft()\n state = jax.tree.map(\n lambda x: jnp.moveaxis(x, 0, axis) if axis != 0 else x,\n state,\n )\n states.append(state)\n elif axis is None:\n states.append(broadcast_states.popleft())\n else: # axis is Carry\n states.append(carry_states.popleft())\n assert not carry_states and not broadcast_states\n assert not vectorized_states or (\n len(vectorized_states) == 1 and not vectorized_states[0]\n )\n elif isinstance(prefix, int):\n state = jax.tree.map(\n lambda x: jnp.moveaxis(x, 0, prefix) if prefix != 0 else x, x.state\n )\n states.extend((state, *carry_states, *broadcast_states))\n elif prefix is None:\n assert is_input_arg\n states.extend(broadcast_states)\n elif prefix is Carry:\n assert is_input_arg\n states.extend(carry_states)\n else:\n obj_repr = 'args' if is_input_arg else 'out'\n raise ValueError(\n f'Invalid axes {prefix} at {obj_repr}{jax.tree_util.keystr(path)}'\n )\n return ctx.merge(graphdef, *states)\n else:\n if isinstance(prefix, StateAxes):\n obj_repr = 'args' if is_input_arg else 'out'\n raise ValueError(\n 'Cannot use StateAxes on non-graph nodes, '\n f'found {prefix} at {obj_repr}{jax.tree_util.keystr(path)}'\n )\n elif prefix is Carry:\n return x\n elif prefix is None:\n return x\n elif isinstance(prefix, int):\n if not isinstance(x, (jax.Array, np.ndarray)):\n obj_repr = 'args' if is_input_arg else 'out'\n raise ValueError(\n f'Expected an array, got {type(x).__name__} at '\n f'{obj_repr}{jax.tree_util.keystr(path)}'\n )\n if prefix != 0:\n x = jnp.moveaxis(x, 0, prefix)\n return x\n else:\n obj_repr = 'args' if is_input_arg else 'out'\n raise ValueError(\n f'Invalid axes {prefix} at {obj_repr}{jax.tree_util.keystr(path)}'\n )\n\n\n@dataclasses.dataclass(eq=False)\nclass ScanFn:\n f: tp.Callable[..., tp.Any]\n input_carry_argnum: int | None | tp.Literal['all']\n output_carry_argnum: int | None | tp.Literal['all']\n in_axes: tp.Any\n out_axes: tp.Any\n transform_metadata: tp.Mapping[str, tp.Any]\n\n def __post_init__(self):\n functools.update_wrapper(self, self.f)\n\n def __call__(\n self,\n carry: tuple[\n tp.Any, # carry_arg\n PytreeDeque[list[State]], # carry_deque\n PytreeDeque[list[State]], # broadcast_deque\n PytreeDeque[Broadcasted], # broadcast_arrays\n ],\n scan_in: tuple[tp.Any, ...],\n ):\n pure_carry_arg, carry_deque, broadcast_deque, broadcast_arrays = carry\n broadcast_deque_out = PytreeDeque(broadcast_deque)\n broadcast_arrays_out = PytreeDeque(broadcast_arrays)\n graphdefs_deque, pure_args = scan_in\n\n if self.input_carry_argnum == 'all':\n assert pure_args == ()\n pure_args = (pure_carry_arg,)\n elif isinstance(self.input_carry_argnum, int):\n assert pure_args[self.input_carry_argnum] is None\n pure_args = extract.replace_at(pure_args, self.input_carry_argnum, pure_carry_arg)\n else:\n assert self.input_carry_argnum is None\n assert pure_carry_arg is None\n\n if spmd.PARTITION_NAME in self.transform_metadata:\n pure_args = _update_variable_sharding_metadata(\n pure_args, self.transform_metadata, spmd.remove_axis\n )\n\n args: tuple = extract.from_tree(\n pure_args,\n prefix=self.in_axes,\n merge_fn=functools.partial(\n _scan_merge_in, carry_deque, graphdefs_deque, broadcast_deque, broadcast_arrays\n ),\n is_leaf=lambda x: isinstance(x, (extract.NodeStates, Broadcasted)),\n map_non_graph_nodes=True,\n ctxtag='scan',\n is_inner=True,\n )\n assert not carry_deque and not broadcast_deque and not broadcast_arrays\n\n out = self.f(*args)\n\n # extract the carry from the args\n if self.input_carry_argnum == 'all':\n carry_arg = args[0]\n elif isinstance(self.input_carry_argnum, int):\n carry_arg = args[self.input_carry_argnum]\n else:\n assert self.input_carry_argnum is None\n carry_arg = None\n\n # extract the carry from the output\n if self.output_carry_argnum == 'all':\n carry_arg_out = out\n out = None\n elif isinstance(self.output_carry_argnum, int):\n assert isinstance(out, tuple)\n carry_arg_out = out[self.output_carry_argnum]\n out = extract.replace_at(out, self.output_carry_argnum, None)\n else:\n assert self.output_carry_argnum is None\n carry_arg_out = None\n\n # TODO(cgarciae): allowing new references might lead to inconsistencies with\n # scan's looping semantics and we would also need to propagate the input\n _check_carry_same_references(carry_arg, carry_arg_out)\n\n args_out: tuple = extract.clear_non_graph_nodes(args)\n\n # replace the carry from the input args with the carry from the output\n if self.input_carry_argnum == 'all':\n args_out = (carry_arg_out,)\n elif isinstance(self.input_carry_argnum, int):\n args_out = extract.replace_at(args_out, self.input_carry_argnum, carry_arg_out)\n else:\n assert self.input_carry_argnum is None\n assert carry_arg_out is None\n\n carry_deque_out = PytreeDeque[list[State | variablelib.Variable]]()\n graphdefs_out = PytreeDeque[graphlib.GraphDef]()\n _broadcast_deque_out_tmp = PytreeDeque[\n list[State | variablelib.Variable]\n ]() # discarded\n pure_args_out: tuple\n pure_args_out, pure_out = extract.to_tree(\n (args_out, out),\n prefix=(self.in_axes, self.out_axes),\n split_fn=functools.partial(\n _scan_split_out, carry_deque_out, graphdefs_out, _broadcast_deque_out_tmp\n ),\n map_non_graph_nodes=True,\n ctxtag='scan',\n )\n if spmd.PARTITION_NAME in self.transform_metadata:\n pure_args_out, pure_out = _update_variable_sharding_metadata(\n (pure_args_out, pure_out),\n self.transform_metadata,\n spmd.add_axis,\n )\n\n # extract the pure carry from the pure args\n if self.input_carry_argnum == 'all':\n pure_carry_arg_out = pure_args_out[0]\n pure_args_out = ()\n elif isinstance(self.input_carry_argnum, int):\n pure_carry_arg_out = pure_args_out[self.input_carry_argnum]\n pure_args_out = extract.replace_at(pure_args_out, self.input_carry_argnum, None)\n else:\n assert self.input_carry_argnum is None\n pure_carry_arg_out = None\n\n carry_arg_out = (\n pure_carry_arg_out,\n carry_deque_out,\n broadcast_deque_out,\n broadcast_arrays_out,\n )\n scan_out = (\n graphdefs_out,\n pure_args_out,\n pure_out,\n )\n return carry_arg_out, scan_out\n\n\n@dataclasses.dataclass(eq=False)\nclass SimpleScanFn:\n f: tp.Callable[..., tp.Any]\n graph: bool\n in_axes: tp.Any\n out_axes: tp.Any\n out_is_tuple: bool\n carry_arg_index: int | None\n carry_out_index: int | None\n\n def __post_init__(self):\n functools.update_wrapper(self, self.f, updated=())\n\n @extract.treemap_copy_args\n def __call__(self, *args):\n updates, snapshot = extract.updates_and_snapshot(args)\n updates = extract.mask_at(updates, self.carry_arg_index)\n snapshot = extract.mask_at(snapshot, self.carry_arg_index)\n if self.carry_arg_index is not None:\n carry_in = args[self.carry_arg_index]\n else:\n carry_in = None\n if self.graph:\n args = extract.from_tree2(args)\n\n out = self.f(*args)\n\n if self.graph:\n out = extract.to_tree2(out, prefix=self.out_axes)\n\n if self.carry_out_index is not None:\n carry_out = out[self.carry_out_index] if self.out_is_tuple else out\n extract.check_same_variables(carry_in, carry_out, 'scan')\n\n def keep_fn(path, prefix, cur, snap):\n changed = extract.variable_changed(cur, snap)\n if prefix is None and changed:\n raise ValueError(\n f'Broadcast (None axis) Variable at {jax.tree_util.keystr(path)} '\n 'was mutated during scan. Only Carry and scanned Variables can be '\n 'updated.'\n )\n return changed\n\n extract.check_no_aliases('scan', args=updates, out=out)\n updates = extract.mask_variable_updates(\n updates, snapshot, prefix=self.in_axes, keep_fn=keep_fn,\n )\n if self.out_is_tuple:\n return (*out, updates)\n return (out, updates)\n\n\n@tp.overload\ndef scan(\n *,\n length: int | None = None,\n reverse: bool = False,\n unroll: int | bool = 1,\n _split_transpose: bool = False,\n # extended api\n in_axes: int | None | type[Carry] | tuple[tp.Any, ...] = (Carry, 0),\n out_axes: tp.Any = (Carry, 0),\n # nnx specific\n transform_metadata: tp.Mapping[str, tp.Any] = FrozenDict({}),\n graph: bool | None = None,\n graph_updates: bool | None = None,\n) -> tp.Callable[[F], F]:\n ...\n\n\n@tp.overload\ndef scan(\n f: F,\n *,\n length: int | None = None,\n reverse: bool = False,\n unroll: int | bool = 1,\n _split_transpose: bool = False,\n # extended api\n in_axes: int | None | type[Carry] | tuple[tp.Any, ...] = (Carry, 0),\n out_axes: tp.Any = (Carry, 0),\n # nnx specific\n transform_metadata: tp.Mapping[str, tp.Any] = FrozenDict({}),\n graph: bool | None = None,\n graph_updates: bool | None = None,\n) -> F:\n ...\n\n\ndef scan(\n f: F | type[Missing] = Missing,\n *,\n length: int | None = None,\n reverse: bool = False,\n unroll: int | bool = 1,\n _split_transpose: bool = False,\n # extended api\n in_axes: int | None | type[Carry] | tuple[tp.Any, ...] = (Carry, 0),\n out_axes: tp.Any = (Carry, 0),\n # nnx specific\n transform_metadata: tp.Mapping[str, tp.Any] = FrozenDict({}),\n graph: bool | None = None,\n graph_updates: bool | None = None,\n) -> F | tp.Callable[[F], F]:\n \"\"\"A Flax NNX transformation of `jax.lax.scan`_.\n\n Example::\n\n import jax\n from flax import nnx\n\n class Block(nnx.Module):\n def __init__(self, input_dim, features, *, rngs):\n self.linear = nnx.Linear(input_dim, features, rngs=rngs)\n self.dropout = nnx.Dropout(0.1, rngs=rngs)\n\n def __call__(self, x: jax.Array):\n x = self.linear(x)\n x = self.dropout(x)\n x = jax.nn.relu(x)\n return x\n\n class Model(nnx.Module):\n def __init__(self, num_layers, features, *, rngs):\n # In this model implementation we create\n # multiple blocks using vmap\n\n # As Block contains dropout op, we prefer\n # to split RNG into num_layers of RNGs\n # using @nnx.split_rngs decorator.\n # Next, nnx.vmap creates a vectorized version of Block.\n # in_axes and out_axes define vectorization axis\n # of the input splitted rngs and the output Block instance.\n # Both axes should be 0.\n @nnx.split_rngs(splits=num_layers)\n @nnx.vmap(in_axes=(0,), out_axes=0)\n def create_block(rngs: nnx.Rngs):\n return Block(features, features, rngs=rngs)\n\n self.blocks = create_block(rngs)\n self.num_layers = num_layers\n\n def __call__(self, x):\n # Forward pass method implementation\n\n # We use nnx.scan to apply sequentially the blocks\n # on the input, for example with num_layers=3\n # output = block[0](x)\n # output = block[1](output)\n # output = block[2](output)\n #\n # In `forward` function defined below:\n # - x represents the loop carry value\n # - model is the data to scan along the leading axis\n # nnx.scan args:\n # - in_axes marks the inputs: x is marked as carry\n # and the model is to scan along the axis 0\n # - out_axes marks the output as carry\n @nnx.scan(in_axes=(nnx.Carry, 0), out_axes=nnx.Carry)\n def forward(x, model):\n x = model(x)\n return x\n\n return forward(x, self.blocks)\n\n # Alternatively, we can also decorate `self.__call__` method\n # @nnx.scan(in_axes=(0, nnx.Carry), out_axes=nnx.Carry)\n # def __call__(self, x):\n # return self.blocks(x)\n\n model = Model(2, 4, rngs=nnx.Rngs(0))\n _, params, _ = nnx.split(model, nnx.Param, ...)\n print(params) # kernel of shape: (2, 4, 4)\n\n x = jnp.arange(5 * 4, dtype=\"float32\").reshape((5, 4))\n y = model(x)\n print(y.shape) # shape: (5, 4)\n\n Args:\n f: a Python function to be scanned\n length: optional integer specifying the number of loop iterations\n reverse: optional boolean specifying whether to run the scan iteration\n forward (the default) or in reverse\n unroll: optional positive int or bool specifying, in the underlying\n operation of the scan primitive, how many scan iterations to unroll\n within a single iteration of a loop.\n in_axes: integer, None, :class:`flax.nnx.Carry` or sequence of values specifying\n the kind of input args. Integer value would specify the axis of corresponding\n input data to scan along. :class:`flax.nnx.Carry` marks the input data as\n loop carry value. None marks the input data as auxiliary input.\n out_axes: integer, None, :class:`flax.nnx.Carry` or sequence of values specifying\n the kind of output args. See ``in_axes`` for details. Note that If ``in_axes``\n contains :class:`flax.nnx.Carry` then ``out_axes`` must also contain :class:`flax.nnx.Carry`.\n graph_updates: If ``True``, propagates updates on graph structure\n that happen inside the transform to the input graphs, has no\n effect when ``graph=False``. When ``False``, using ``StateAxes``\n is not supported.\n\n .. _jax.lax.scan: https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.scan.html>\n \"\"\"\n if f is Missing:\n return functools.partial(\n scan,\n length=length,\n reverse=reverse,\n unroll=unroll,\n _split_transpose=_split_transpose,\n in_axes=in_axes,\n out_axes=out_axes,\n transform_metadata=transform_metadata,\n graph=graph,\n graph_updates=graph_updates,\n ) # type: ignore[return-value]\n\n f_unbound, _, was_bound = _resolve_bound_callable(f)\n if was_bound:\n _raise_bound_method_error('scan')\n\n if graph is None:\n graph = graphlib.set_graph_mode.current_value()\n if graph_updates is None:\n graph_updates = graphlib.set_graph_updates.current_value()\n if not graph or not graph_updates:\n return _simple_scan(\n f, f_unbound, graph=graph,\n in_axes=in_axes, out_axes=out_axes,\n length=length, reverse=reverse, unroll=unroll,\n _split_transpose=_split_transpose,\n )\n\n return _graph_updates_scan(\n f, f_unbound,\n in_axes=in_axes, out_axes=out_axes,\n length=length, reverse=reverse, unroll=unroll,\n _split_transpose=_split_transpose,\n transform_metadata=transform_metadata,\n )\n\ndef _simple_scan(\n f, f_unbound, *,\n graph, in_axes, out_axes,\n length, reverse, unroll, _split_transpose,\n):\n if any(isinstance(x, StateAxes) for x in jax.tree.leaves(in_axes)):\n raise ValueError(\n '`in_axes` cannot contain `StateAxes` objects '\n 'when `graph=False`. '\n + graphlib._tree_mode_suggestion_transform('scan')\n )\n if any(isinstance(x, StateAxes) for x in jax.tree.leaves(out_axes)):\n raise ValueError(\n '`out_axes` cannot contain `StateAxes` objects '\n 'when `graph=False`. '\n + graphlib._tree_mode_suggestion_transform('scan')\n )\n if graph:\n extract.check_prefix(in_axes, 'in_axes', 'scan')\n extract.check_prefix(out_axes, 'out_axes', 'scan')\n\n out_is_tuple = isinstance(out_axes, tuple)\n if in_axes is Carry:\n in_axes = (Carry,)\n if isinstance(in_axes, tuple):\n carry_arg_index = next(\n (i for i, ax in enumerate(in_axes) if ax is Carry), None\n )\n updates_out_axes = extract.mask_at(in_axes, carry_arg_index)\n else:\n carry_arg_index = None\n updates_out_axes = in_axes\n\n if isinstance(out_axes, tuple):\n carry_out_index = next(\n (i for i, ax in enumerate(out_axes) if ax is Carry), None\n )\n else:\n carry_out_index = None\n\n simple_scan_fn = SimpleScanFn(\n f_unbound, graph=graph,\n in_axes=in_axes, out_axes=out_axes,\n out_is_tuple=out_is_tuple,\n carry_arg_index=carry_arg_index,\n carry_out_index=carry_out_index,\n )\n\n if out_is_tuple:\n augmented_out_axes = (*out_axes, updates_out_axes)\n else:\n augmented_out_axes = (out_axes, updates_out_axes)\n\n @functools.wraps(f)\n def simple_scan_wrapper(*args):\n args = resolve_kwargs(f, args, {})\n if graph:\n args = extract.to_tree2(args, prefix=in_axes)\n\n extract.check_no_aliases('scan', args=args)\n\n result = pure_jax_fancy_scan(\n simple_scan_fn,\n *args,\n length=length,\n reverse=reverse,\n unroll=unroll,\n _split_transpose=_split_transpose,\n in_axes=in_axes,\n out_axes=augmented_out_axes,\n )\n\n if out_is_tuple:\n n = len(out_axes)\n out = result[:n]\n updates = result[n]\n else:\n out, updates = result\n\n masked_args = extract.mask_at(args, carry_arg_index)\n extract.apply_variable_updates(masked_args, updates)\n\n if carry_arg_index is not None:\n carry_in = args[carry_arg_index]\n carry_out = (\n out[carry_out_index] if out_is_tuple else out\n )\n extract.update_carry_variables(carry_in, carry_out)\n if out_is_tuple:\n out = extract.replace_at(out, carry_out_index, carry_in)\n else:\n out = carry_in\n\n if graph:\n out = extract.from_tree2(out)\n\n return out\n\n return simple_scan_wrapper\n\n\ndef _graph_updates_scan(\n f, f_unbound, *,\n in_axes, out_axes,\n length, reverse, unroll, _split_transpose,\n transform_metadata,\n):\n _check_out_axes(out_axes)\n\n input_carry_argnum = _get_carry_argnum(in_axes, is_in_axes=True)\n output_carry_argnum = _get_carry_argnum(out_axes, is_in_axes=False)\n\n if (input_carry_argnum is None and output_carry_argnum is not None) or (\n input_carry_argnum is not None and output_carry_argnum is None\n ):\n raise ValueError(\n 'If one of in_axes or out_axes has Carry, the other must also have Carry. '\n f'Got {in_axes=!r} and {out_axes=!r}'\n )\n\n scan_fn = ScanFn(\n f_unbound,\n input_carry_argnum,\n output_carry_argnum,\n in_axes,\n out_axes,\n transform_metadata,\n )\n\n @functools.wraps(f)\n @graphlib.update_context('scan')\n def scan_wrapper(*args, **kwargs):\n args = resolve_kwargs(f, args, kwargs)\n\n if in_axes is Carry and len(args) != 1:\n raise ValueError(\n f'When in_axes=Carry, the function must take exactly one argument, '\n f'got {len(args)} arguments.'\n )\n\n graphdefs_deque = PytreeDeque()\n carry_deque = PytreeDeque()\n broadcast_deque = PytreeDeque()\n broadcast_arrays = PytreeDeque()\n pure_args: tuple = extract.to_tree(\n args,\n prefix=in_axes,\n split_fn=functools.partial(\n _scan_split_in, carry_deque, graphdefs_deque, broadcast_deque, broadcast_arrays\n ),\n map_non_graph_nodes=True,\n ctxtag='scan',\n )\n if isinstance(input_carry_argnum, int):\n pure_carry_arg = pure_args[input_carry_argnum]\n pure_args = extract.replace_at(pure_args, input_carry_argnum, None)\n elif input_carry_argnum == 'all':\n pure_carry_arg = pure_args[0]\n pure_args = ()\n else:\n assert input_carry_argnum is None\n pure_carry_arg = None\n\n carry = (pure_carry_arg, carry_deque, broadcast_deque, broadcast_arrays)\n scan_in = (graphdefs_deque, pure_args)\n\n carry_out, scan_out = jax.lax.scan(\n scan_fn,\n carry,\n scan_in,\n length=length,\n reverse=reverse,\n unroll=unroll,\n _split_transpose=_split_transpose,\n )\n (\n pure_carry_arg_out,\n carry_deque_out,\n broadcast_deque_out,\n broadcast_arrays_out,\n ) = carry_out\n (\n graphdefs_out,\n pure_args_out,\n pure_out,\n ) = scan_out\n\n if input_carry_argnum == 'all':\n pure_args_out = (pure_carry_arg_out,)\n elif isinstance(input_carry_argnum, int):\n pure_args_out = extract.replace_at(pure_args_out, input_carry_argnum, pure_carry_arg_out)\n else:\n assert input_carry_argnum is None\n assert pure_carry_arg_out is None\n\n args_out, out = extract.from_tree(\n (pure_args_out, pure_out),\n prefix=(in_axes, out_axes),\n merge_fn=functools.partial(\n _scan_merge_out, carry_deque_out, graphdefs_out, broadcast_deque_out\n ),\n is_leaf=lambda x: isinstance(x, (extract.NodeStates, Broadcasted)),\n map_non_graph_nodes=True,\n ctxtag='scan',\n is_inner=False,\n )\n\n if input_carry_argnum == 'all':\n carry_arg = args_out[0]\n elif isinstance(input_carry_argnum, int):\n carry_arg = args_out[input_carry_argnum]\n else:\n assert input_carry_argnum is None\n carry_arg = None\n\n if output_carry_argnum == 'all':\n out = carry_arg\n elif isinstance(output_carry_argnum, int):\n out = extract.replace_at(out, output_carry_argnum, carry_arg)\n else:\n assert output_carry_argnum is None\n assert carry_arg is None\n\n return out\n\n return scan_wrapper\n\n\ndef pure_jax_fancy_scan(\n f,\n *args,\n length: int | None = None,\n reverse: bool = False,\n unroll: int | bool = 1,\n _split_transpose: bool = False,\n in_axes: tp.Any = (Carry, 0),\n out_axes: tp.Any = (Carry, 0),\n):\n if in_axes is Carry:\n in_axes = (Carry,)\n is_axis_leaf = lambda x: x is None or x is Carry\n\n if isinstance(in_axes, tuple):\n for i, ax in enumerate(in_axes):\n if ax is Carry or ax is None or isinstance(ax, int):\n continue\n for leaf in jax.tree.leaves(ax, is_leaf=is_axis_leaf):\n if leaf is Carry:\n raise ValueError(\n 'Carry must be a top-level argument, it cannot be nested. '\n f'Found Carry inside in_axes[{i}]={ax}'\n )\n\n if isinstance(out_axes, tuple):\n for i, ax in enumerate(out_axes):\n if ax is Carry or ax is None or isinstance(ax, int):\n continue\n for path, leaf in jax.tree_util.tree_leaves_with_path(\n ax, is_leaf=is_axis_leaf,\n ):\n if leaf is Carry:\n raise ValueError(\n 'Carry must be a top-level argument, it cannot be nested. '\n f'Found Carry at out_axes[{i}]{jax.tree_util.keystr(path)}'\n )\n\n in_has_carry = in_axes is Carry or (\n isinstance(in_axes, tuple) and Carry in in_axes\n )\n out_has_carry = out_axes is Carry or (\n isinstance(out_axes, tuple) and Carry in out_axes\n )\n if in_has_carry != out_has_carry:\n raise ValueError(\n 'If one of in_axes or out_axes has Carry, the other must also '\n f'have Carry. Got {in_axes=}, {out_axes=}'\n )\n\n\n args_flat, args_treedef = jax.tree.flatten(args)\n _, in_axes_flat = extract.broadcast_prefix2(\n in_axes, args, is_leaf=is_axis_leaf,\n )\n\n carry_indices: list[int] = []\n broadcast_indices: list[int] = []\n scan_indices: list[int] = []\n scan_in_axes: list[int] = []\n\n carry_leaves: list[tp.Any] = []\n broadcast_leaves: list[tp.Any] = []\n scan_leaves: list[tp.Any] = []\n\n for i, (leaf, ax) in enumerate(zip(args_flat, in_axes_flat, strict=True)):\n if ax is Carry:\n carry_indices.append(i)\n carry_leaves.append(leaf)\n elif ax is None:\n broadcast_indices.append(i)\n broadcast_leaves.append(leaf)\n elif isinstance(ax, int):\n scan_indices.append(i)\n scan_in_axes.append(ax)\n if ax != 0:\n leaf = jnp.moveaxis(leaf, ax, 0)\n scan_leaves.append(leaf)\n else:\n raise ValueError(f'Invalid in_axes leaf value: {ax}')\n\n n_in = len(args_flat)\n out_info: list[tuple[\n jax.tree_util.PyTreeDef, list[int], list[int], list[int],\n ]] = []\n\n in_broadcast = jax.tree.map(lambda x: x, broadcast_leaves)\n\n def body_fn(carry_state, scan_x):\n flat = [None] * n_in\n for idx, j in enumerate(carry_indices):\n flat[j] = carry_state[idx]\n for idx, j in enumerate(broadcast_indices):\n flat[j] = in_broadcast[idx]\n if scan_x is not None:\n for idx, j in enumerate(scan_indices):\n flat[j] = scan_x[idx]\n\n reconstructed = args_treedef.unflatten(flat)\n out = f(*reconstructed)\n\n out_flat, out_treedef = jax.tree.flatten(out)\n out_axes_paths, out_axes_flat = extract.broadcast_prefix2(\n out_axes, out, is_leaf=is_axis_leaf,\n )\n\n if not out_info:\n out_carry_idx = []\n out_scan_idx = []\n out_scan_axes = []\n out_broadcast_idx = []\n for j, oax in enumerate(out_axes_flat):\n if oax is Carry:\n out_carry_idx.append(j)\n elif oax is None:\n out_broadcast_idx.append(j)\n elif isinstance(oax, int):\n out_scan_idx.append(j)\n out_scan_axes.append(oax)\n else:\n raise ValueError(f'Invalid out_axes leaf value: {oax}')\n if out_broadcast_idx:\n broadcast_paths = [\n jax.tree_util.keystr(out_axes_paths[j]) for j in out_broadcast_idx\n ]\n broadcast_str = \"\\n\\n \".join(broadcast_paths)\n raise ValueError(\n 'Scan does not support broadcast outputs (None axis). The following '\n f'output leaves are broadcast:\\n\\n {broadcast_str}\\n'\n )\n out_info.append(\n (out_treedef, out_carry_idx, out_scan_idx, out_scan_axes),\n )\n\n oci = out_info[0][1]\n osi = out_info[0][2]\n new_carry = [out_flat[j] for j in oci]\n new_ys = [out_flat[j] for j in osi]\n\n return new_carry, new_ys\n\n final_carry, stacked_ys = jax.lax.scan(\n body_fn,\n carry_leaves,\n scan_leaves if scan_leaves else None,\n length=length,\n reverse=reverse,\n unroll=unroll,\n _split_transpose=_split_transpose,\n )\n\n out_treedef, out_carry_idx, out_scan_idx, out_scan_axes = (\n out_info[0]\n )\n n_out = out_treedef.num_leaves\n out_flat: list[tp.Any] = [None] * n_out\n for idx, j in enumerate(out_carry_idx):\n out_flat[j] = final_carry[idx]\n for idx, j in enumerate(out_scan_idx):\n y = stacked_ys[idx]\n ax = out_scan_axes[idx]\n if ax != 0:\n y = jnp.moveaxis(y, 0, ax)\n out_flat[j] = y\n\n return out_treedef.unflatten(out_flat)\n\n\n# -------------------------------\n# while_loop\n# -------------------------------\n\n\n@dataclasses.dataclass(eq=False)\nclass SimpleWhileLoopBodyFn:\n f: tp.Callable[..., tp.Any]\n graph: bool\n\n def __post_init__(self):\n functools.update_wrapper(self, self.f, updated=())\n\n @extract.treemap_copy_args\n def __call__(self, val):\n val_variables, _ = extract.updates_and_snapshot(val)\n if self.graph:\n val = extract.from_tree2(val)\n out = self.f(val)\n if self.graph:\n out = extract.to_tree2(out)\n extract.check_same_variables(val_variables, out, 'while_loop')\n return out\n\n\n@dataclasses.dataclass(eq=False)\nclass SimpleWhileLoopCondFn:\n f: tp.Callable[..., tp.Any]\n graph: bool\n\n def __post_init__(self):\n functools.update_wrapper(self, self.f)\n\n def __call__(self, val):\n if self.graph:\n val = extract.from_tree2(val)\n return self.f(val)\n\n\n@dataclasses.dataclass(eq=False)\nclass WhileLoopCondFn:\n f: tp.Callable[..., tp.Any]\n\n def __post_init__(self):\n functools.update_wrapper(self, self.f)\n\n def __call__(self, pure_val):\n val = extract.from_tree(pure_val)\n out = self.f(val)\n return out\n\n\ndef _reconsile_index_mapping(tree_to_fix, example_tree):\n def f(a, b):\n if not isinstance(a, extract.NodeStates) or not isinstance(\n a._graphdef, graphlib.GraphDef\n ):\n return a\n return dataclasses.replace(\n a, _graphdef=a._graphdef.with_matching_outer_index(b._graphdef)\n )\n\n return jax.tree.map(f, tree_to_fix, example_tree,\n is_leaf=lambda x: isinstance(x, extract.NodeStates))\n\ndef _add_fake_index_mapping(tree: tp.Any):\n def per_node_state(node_state: extract.NodeStates | tp.Any):\n if not isinstance(node_state, extract.NodeStates) or not isinstance(\n node_state._graphdef, graphlib.GraphDef\n ):\n return node_state\n\n return dataclasses.replace(\n node_state, _graphdef=node_state._graphdef.with_same_outer_index()\n )\n\n return jax.tree.map(per_node_state, tree,\n is_leaf=lambda x: isinstance(x, extract.NodeStates))\n\n\ndef _remove_index_mapping(tree: tp.Any):\n \"\"\"Remove a fake outer_index for the input to match that of the output.\"\"\"\n\n def per_node_state(node_state: extract.NodeStates | tp.Any):\n if not isinstance(node_state, extract.NodeStates) or not isinstance(\n node_state._graphdef, graphlib.GraphDef\n ):\n return node_state\n assert isinstance(node_state._graphdef, graphlib.GraphDef)\n node_state = dataclasses.replace(\n node_state, _graphdef=node_state._graphdef.with_no_outer_index()\n )\n return node_state\n\n return jax.tree.map(per_node_state, tree,\n is_leaf=lambda x: isinstance(x, extract.NodeStates))\n\n\n@dataclasses.dataclass(eq=False)\nclass WhileLoopBodyFn:\n f: tp.Callable[..., tp.Any]\n\n def __post_init__(self):\n functools.update_wrapper(self, self.f)\n\n @graphlib.update_context('while_loop_body')\n def __call__(self, pure_val):\n # Removing the dummy index mapping being added outside of body function.\n pure_val_in = _remove_index_mapping(pure_val)\n\n val = extract.from_tree(\n pure_val_in, ctxtag='while_loop_body', is_inner=True\n )\n out = self.f(val)\n pure_out = extract.to_tree(out, ctxtag='while_loop_body')\n\n try:\n jax.tree.map(lambda a, b: None, pure_val, pure_out)\n except ValueError as e:\n msg = (\n \"nnx.while_loop requires body function's input and output to \"\n 'have the same reference and pytree structure, but they differ. '\n 'If the mismatch comes from `outer_index` field, you might '\n 'have modified reference structure within the body function, '\n 'which is not allowed.'\n f'Detail of the mismatch: \\n {str(e)}'\n )\n raise ValueError(msg)\n\n return pure_out\n\n\n@graphlib.update_context('while_loop')\ndef while_loop(cond_fun: tp.Callable[[T], tp.Any],\n body_fun: tp.Callable[[T], T],\n init_val: T,\n *,\n graph: bool | None = None,\n graph_updates: bool | None = None) -> T:\n \"\"\"A Flax NNX transformation of `jax.lax.while_loop `_.\n\n Caution: for the NNX internal reference tracing mechanism to work, you cannot\n change the variable reference structure of ``init_val`` inside ``body_fun``.\n\n Example::\n\n >>> import jax\n >>> from flax import nnx\n >>> def fwd_fn(input):\n ... module, x, count = input\n ... return module, module(x), count - 1.0\n\n >>> module = nnx.Linear(10, 10, rngs=nnx.Rngs(0))\n >>> x = jax.random.normal(jax.random.key(0), (10,))\n >>> # `module` will be called three times\n >>> _, y, _ = nnx.while_loop(\n ... lambda input: input[-1] > 0, fwd_fn, (module, x, 3.0))\n\n\n Args:\n cond_fun: A function for the continue condition of the while loop, taking a\n single input of type ``T`` and outputting a boolean.\n body_fun: A function that takes an input of type ``T`` and outputs an ``T``.\n Note that both data and modules of ``T`` must have the same reference\n structure between inputs and outputs.\n init_val: The initial input for ``cond_fun`` and ``body_fun``. Must be of type ``T``.\n graph: if True, use graph-mode (default). If False, use tree-mode.\n If None, uses the value of ``nnx_graph_mode`` config.\n graph_updates: If ``True``, propagates updates on graph structure\n that happen inside the transform to the input graphs, has no\n effect when ``graph=False``.\n\n \"\"\"\n if graph is None:\n graph = graphlib.set_graph_mode.current_value()\n if graph_updates is None:\n graph_updates = graphlib.set_graph_updates.current_value()\n if not graph or not graph_updates:\n simple_body_fn = SimpleWhileLoopBodyFn(body_fun, graph=graph)\n simple_cond_fn = SimpleWhileLoopCondFn(cond_fun, graph=graph)\n\n if graph:\n init_val = extract.to_tree2(init_val)\n val_out = jax.lax.while_loop(simple_cond_fn, simple_body_fn, init_val)\n val_out = extract.update_carry_variables(init_val, val_out)\n if graph:\n val_out = extract.from_tree2(val_out)\n return val_out\n\n pure_init_val = extract.to_tree(init_val, ctxtag='while_loop')\n\n pure_init_val = _add_fake_index_mapping(pure_init_val)\n\n pure_out = jax.lax.while_loop(\n WhileLoopCondFn(cond_fun),\n WhileLoopBodyFn(body_fun),\n pure_init_val,\n )\n out = extract.from_tree(pure_out, ctxtag='while_loop', is_inner=False)\n return out\n\n\n@dataclasses.dataclass(eq=False)\nclass SimpleForiLoopBodyFn:\n f: tp.Callable[..., tp.Any]\n graph: bool\n\n def __post_init__(self):\n functools.update_wrapper(self, self.f, updated=())\n\n @extract.treemap_copy_args\n def __call__(self, i, val):\n val_variables, _ = extract.updates_and_snapshot(val)\n if self.graph:\n val = extract.from_tree2(val)\n out = self.f(i, val)\n if self.graph:\n out = extract.to_tree2(out)\n extract.check_same_variables(val_variables, out, 'fori_loop')\n return out\n\n\n@dataclasses.dataclass(eq=False)\nclass ForiLoopBodyFn:\n f: tp.Callable[..., tp.Any]\n\n def __post_init__(self):\n functools.update_wrapper(self, self.f)\n\n @graphlib.update_context('fori_loop_body')\n def __call__(self, i, pure_val_in):\n val = extract.from_tree(pure_val_in, ctxtag='fori_loop_body', is_inner=True)\n out = self.f(i, val)\n pure_out = extract.to_tree(out, ctxtag='fori_loop_body')\n return pure_out\n\n\n@graphlib.update_context('fori_loop')\ndef fori_loop(lower: int, upper: int,\n body_fun: tp.Callable[[int, T], T],\n init_val: T,\n *,\n unroll: int | bool | None = None,\n graph: bool | None = None,\n graph_updates: bool | None = None) -> T:\n \"\"\"A Flax NNX transformation of `jax.lax.fori_loop `_.\n\n Caution: for the NNX internal reference tracing mechanism to work, you cannot\n change the variable reference structure of `init_val` inside `body_fun`.\n\n Example::\n\n >>> import jax\n >>> from flax import nnx\n\n >>> def fwd_fn(i, input):\n ... m, x = input\n ... m.kernel[...] = jnp.identity(10) * i\n ... return m, m(x)\n\n >>> module = nnx.Linear(10, 10, rngs=nnx.Rngs(0))\n >>> x = jax.random.normal(jax.random.key(0), (10,))\n >>> _, y = nnx.fori_loop(2, 4, fwd_fn, (module, x))\n >>> np.testing.assert_array_equal(y, x * 2 * 3)\n\n\n Args:\n lower: An integer representing the loop index lower bound (inclusive).\n upper: An integer representing the loop index upper bound (exclusive).\n body_fun: a function that takes an input of type ``T`` and outputs an ``T``.\n Note that both data and modules of ``T`` must have the same reference\n structure between inputs and outputs.\n init_val: the initial input for body_fun. Must be of type ``T``.\n unroll: An optional integer or boolean that determines how much to unroll\n the loop. If an integer is provided, it determines how many unrolled\n loop iterations to run within a single rolled iteration of the loop. If a\n boolean is provided, it will determine if the loop is competely unrolled\n (i.e. ``unroll=True``) or left completely unrolled (i.e. ``unroll=False``).\n This argument is only applicable if the loop bounds are statically known.\n graph: if True, use graph-mode (default). If False, use tree-mode.\n If None, uses the value of ``nnx_graph_mode`` config.\n graph_updates: If ``True``, propagates updates on graph structure\n that happen inside the transform to the input graphs, has no\n effect when ``graph=False``.\n\n Returns:\n A loop value from the final iteration, of type ``T``.\n\n \"\"\"\n if graph is None:\n graph = graphlib.set_graph_mode.current_value()\n if graph_updates is None:\n graph_updates = graphlib.set_graph_updates.current_value()\n if not graph or not graph_updates:\n simple_body_fn = SimpleForiLoopBodyFn(body_fun, graph=graph)\n\n if graph:\n init_val = extract.to_tree2(init_val)\n val_out = jax.lax.fori_loop(\n lower, upper,\n simple_body_fn,\n init_val,\n unroll=unroll,\n )\n val_out = extract.update_carry_variables(init_val, val_out)\n if graph:\n val_out = extract.from_tree2(val_out)\n return val_out\n\n pure_init_val = extract.to_tree(init_val, ctxtag='fori_loop')\n body = ForiLoopBodyFn(body_fun)\n pure_out = jax.eval_shape(body, lower, pure_init_val)\n pure_init_val = _reconsile_index_mapping(pure_init_val, pure_out)\n pure_out = jax.lax.fori_loop(lower, upper,\n body, pure_init_val,\n unroll=unroll)\n out = extract.from_tree(pure_out, ctxtag='fori_loop', is_inner=False)\n return out\n" }, { "path": "flax/nnx/transforms/transforms.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# pytype: skip-file\nfrom __future__ import annotations\n\nfrom abc import abstractmethod\nimport dataclasses\nimport functools\nimport inspect\nimport typing as tp\n\nfrom jax._src import checkify as checkify_lib\n\nfrom flax.nnx import (\n extract,\n graphlib,\n variablelib,\n)\nfrom flax.nnx.module import Module\nfrom flax.nnx.proxy_caller import (\n CallableProxy,\n DelayedAccessor,\n)\nfrom flax.nnx.transforms import general\nfrom flax.typing import MISSING, Leaf, Missing\nimport jax\nimport jax.core\nimport jax.stages\n\nA = tp.TypeVar('A')\nC = tp.TypeVar('C')\nB = tp.TypeVar('B')\nF = tp.TypeVar('F', bound=tp.Callable[..., tp.Any])\nG = tp.TypeVar('G', bound=tp.Callable[..., tp.Any])\nM = tp.TypeVar('M', bound=Module)\nMA = tp.TypeVar('MA', bound=Module)\nN = tp.TypeVar('N', bound=Module)\nStrInt = tp.TypeVar('StrInt', str, int)\nAxisName = tp.Hashable\nLeaves = list[Leaf]\nIndex = int\n\n\n@tp.overload\ndef resolve_kwargs(\n fun: tp.Callable[..., tp.Any],\n args: tuple,\n kwargs: dict[str, tp.Any],\n) -> tuple: ...\n@tp.overload\ndef resolve_kwargs() -> tp.Callable[[F], F]: ...\ndef resolve_kwargs(\n fun: tp.Callable[..., tp.Any] | Missing = MISSING,\n args: tuple | Missing = MISSING,\n kwargs: dict[str, tp.Any] | Missing = MISSING,\n) -> tuple | tp.Callable[[F], F]:\n if isinstance(fun, Missing):\n\n def resolve_kwargs_decorator(f):\n @functools.wraps(f)\n def resolve_kwargs_wrapper(*args, **kwargs):\n args = resolve_kwargs(f, args, kwargs)\n return f(*args)\n\n return resolve_kwargs_wrapper\n\n return resolve_kwargs_decorator # type: ignore\n\n if isinstance(args, Missing):\n raise ValueError('args must be provided')\n if isinstance(kwargs, Missing):\n raise ValueError('kwargs must be provided')\n\n if isinstance(fun, functools.partial):\n # functools.partial should have an opaque signature.\n fun = lambda *args, **kwargs: None\n ba = inspect.signature(fun).bind(*args, **kwargs)\n ba.apply_defaults()\n if ba.kwargs:\n raise TypeError('keyword arguments could not be resolved to positions')\n else:\n return ba.args\n\n\n\n# -------------------------------\n# helper utilities for bound methods & indices\n# -------------------------------\n\ndef _resolve_bound_callable(\n f: tp.Callable[..., tp.Any],\n) -> tuple[tp.Callable[..., tp.Any], tp.Any | None, bool]:\n \"\"\"Detects and extracts bound methods from NNX Module callables.\n\n This function unwraps functools.partial layers to reach the underlying\n callable before checking if it's a bound method of an NNX Module.\n\n Args:\n f: A callable that may be a bound method of an NNX Module, potentially\n wrapped in functools.partial.\n\n Returns:\n A tuple of (unbound_fn, bound_self, was_bound) where:\n - unbound_fn: The unbound function (or original if not bound)\n - bound_self: The Module instance if f was bound, None otherwise\n - was_bound: True if f was a bound method, False otherwise\n\n Note:\n Preserves functools.partial wrappers around the callable and follows\n the same detection pattern as _get_unbound_fn in bridge/module.py.\n Detection occurs before any argnum shifting or index normalization.\n \"\"\"\n # Unwrap functools.partial layers to reach the underlying callable.\n partials: list[tuple[tuple[tp.Any, ...], dict[str, tp.Any] | None]] = []\n g = f\n while isinstance(g, functools.partial): # type: ignore[arg-type]\n partials.append((g.args or (), g.keywords)) # type: ignore[attr-defined]\n g = g.func # type: ignore[attr-defined]\n\n bound_self = getattr(g, \"__self__\", None)\n was_bound = bool(inspect.ismethod(g) and isinstance(bound_self, Module))\n if was_bound:\n g = g.__func__ # type: ignore[attr-defined]\n\n # Reapply partials in reverse unwrap order.\n for args, kwargs in reversed(partials):\n kwargs = {} if kwargs is None else kwargs\n g = functools.partial(g, *args, **kwargs)\n\n return g, (bound_self if was_bound else None), was_bound\n\n\ndef _raise_bound_method_error(transform_name: str):\n \"\"\"Raises a standardized error for bound method usage with NNX transforms.\n\n Args:\n transform_name: Name of the transform (e.g., 'grad', 'jit', 'remat').\n \"\"\"\n raise ValueError(\n f\"nnx.{transform_name} does not support bound methods. \"\n f\"Use the decorator form @nnx.{transform_name} or call \"\n f\"nnx.{transform_name}(MyClass.method)(instance, ...) with the unbound method.\"\n )\n\n\nclass LiftedModule(tp.Generic[M], Module): # type: ignore[ignored-abstractmethod]\n @abstractmethod\n def _call(self, accessor: DelayedAccessor, *args, **kwargs) -> tp.Any:\n pass\n\n @property\n @abstractmethod\n def _submodule(self) -> M:\n pass # type: ignore[bad-return-type] # why pytype?\n\n def __call__(self, *args, **kwargs) -> tp.Any:\n return self.call(*args, **kwargs) # type: ignore\n\n @property\n def call(self) -> tp.Any:\n module = self\n\n def check_and_call(accessor: DelayedAccessor, *args, **kwargs):\n return self._call(accessor, *args, **kwargs)\n\n proxy = CallableProxy(check_and_call) # type: ignore[arg-type]\n\n while isinstance(module._submodule, LiftedModule):\n module = module._submodule\n proxy = proxy.call\n\n return proxy # type: ignore\n\n\n# -------------------------------\n# simple transforms\n# -------------------------------\n@dataclasses.dataclass(frozen=True)\nclass ValueMetadata:\n var_type: type[variablelib.Variable]\n value: tp.Any\n metadata: dict[str, tp.Any]\n\n\ndef _flatten_value_metadata(\n value_metadata: tp.Union[tp.Any, ValueMetadata],\n):\n metadata = tuple(sorted(value_metadata.metadata.items()))\n return (value_metadata.value,), (value_metadata.var_type, metadata)\n\n\ndef _unflatten_value_metadata(aux_data, children):\n var_type, metadata_items = aux_data\n metadata = dict(metadata_items)\n return ValueMetadata(var_type=var_type, value=children[0], metadata=metadata)\n\n\njax.tree_util.register_pytree_node(\n ValueMetadata,\n _flatten_value_metadata,\n _unflatten_value_metadata,\n)\n\n\ndef _to_value_metadata(node):\n def to_value_metadata(x):\n if isinstance(x, variablelib.Variable):\n value = x.get_raw_value()\n if variablelib.is_array_ref(value):\n value = value[...]\n metadata = x.get_metadata()\n return ValueMetadata(var_type=x.var_type, value=value, metadata=metadata)\n return x\n\n return jax.tree.map(\n to_value_metadata,\n node,\n is_leaf=lambda x: isinstance(x, variablelib.Variable),\n )\n\n\ndef _to_variable(node):\n # import here to avoid circular imports\n from flax.nnx.spmd import get_var_pspec\n\n def to_variable(x):\n if isinstance(x, ValueMetadata):\n var = x.var_type._new(x.value, x.metadata)\n\n global_mesh = jax.sharding.get_abstract_mesh()\n if global_mesh.axis_sizes == ():\n global_mesh = None\n mesh = var.get_metadata(\"mesh\", None) or global_mesh\n if mesh is not None and (not hasattr(var, 'sharding') or var.sharding is None):\n pspec = get_var_pspec(var)\n sharding = jax.sharding.NamedSharding(mesh=mesh, spec=pspec)\n var.set_value(jax.ShapeDtypeStruct(shape=var.shape, dtype=var.dtype, sharding=sharding))\n return var\n return x\n\n return jax.tree.map(\n to_variable, node, is_leaf=lambda x: isinstance(x, ValueMetadata)\n )\n\n\n@dataclasses.dataclass(eq=False)\nclass SimpleEvalShapeFn:\n f: tp.Callable[..., tp.Any]\n graph: bool\n\n def __post_init__(self):\n functools.update_wrapper(self, self.f, updated=())\n\n @extract.treemap_copy_args\n def __call__(self, *args, **kwargs):\n if self.graph:\n args, kwargs = extract.from_tree2((args, kwargs))\n out = self.f(*args, **kwargs)\n if self.graph:\n out = extract.to_tree2(out)\n extract.check_no_aliases('eval_shape', args=args, kwargs=kwargs, out=out)\n return out\n\n\ndef eval_shape(\n f: tp.Callable[..., A],\n *args: tp.Any,\n graph: bool | None = None,\n graph_updates: bool | None = None,\n **kwargs: tp.Any,\n) -> A:\n \"\"\"A \\\"lifted\\\" version of `jax.eval_shape `_\n that can handle `flax.nnx.Module `_\n / graph nodes as arguments.\n\n Similar to ``jax.eval_shape``, it computes the shape/dtype of a function `f` without\n performing any floating point operations (FLOPs) which can be expensive. This can be\n useful for performing shape inference, for example. Unlike `jax.eval_shape`,\n `nnx.eval_shape` will automatically compute the expected sharding based on Flax sharding metadata\n for all Variables not using explicit sharding.\n\n Args:\n f: the function to evaluate.\n *args: positional arguments to ``f``.\n graph: If ``True`` (default), uses graph-mode which supports the full\n NNX feature set including shared references and reference semantics.\n If ``False``, uses tree-mode which treats Modules as regular JAX\n pytrees, avoiding the overhead of the graph protocol.\n graph_updates: If ``True``, propagates updates on graph structure\n that happen inside the transform to the input graphs, has no\n effect when ``graph=False``.\n **kwargs: keyword arguments to ``f``.\n\"\"\"\n f_call, _, was_bound = _resolve_bound_callable(f)\n\n if was_bound:\n _raise_bound_method_error('eval_shape')\n\n if graph is None:\n graph = graphlib.set_graph_mode.current_value()\n if graph_updates is None:\n graph_updates = graphlib.set_graph_updates.current_value()\n if not graph or not graph_updates:\n if graph:\n args, kwargs = extract.to_tree2((args, kwargs))\n extract.check_no_aliases('eval_shape', args=args, kwargs=kwargs)\n out = jax.eval_shape(\n SimpleEvalShapeFn(f_call, graph=graph), *args, **kwargs\n )\n if graph:\n out = extract.from_tree2(out)\n return out\n\n args, kwargs = extract.to_tree((args, kwargs))\n\n @functools.wraps(f)\n def _eval_shape_fn(*args, **kwargs):\n args, kwargs = extract.from_tree((args, kwargs))\n out = f_call(*args, **kwargs)\n return _to_value_metadata(extract.to_tree(out))\n\n out = jax.eval_shape(_eval_shape_fn, *args, **kwargs)\n return extract.from_tree(_to_variable(out))\n\n@dataclasses.dataclass(eq=False)\nclass CheckifyFn:\n f: tp.Callable[..., tp.Any]\n\n def __post_init__(self):\n functools.update_wrapper(self, self.f)\n\n def __call__(self, *pure_args, **pure_kwargs):\n args, kwargs = extract.from_tree(\n (pure_args, pure_kwargs), ctxtag='checkify', is_inner=True\n )\n out = self.f(*args, **kwargs)\n\n args_out, kwargs_out = extract.clear_non_graph_nodes((args, kwargs))\n pure_args_out, pure_kwargs_out, pure_out = extract.to_tree(\n (args, kwargs, out), ctxtag='checkify'\n )\n return pure_args_out, pure_kwargs_out, pure_out\n\n\n@dataclasses.dataclass(eq=False)\nclass SimpleCheckifyFn:\n f: tp.Callable[..., tp.Any]\n graph: bool\n\n def __post_init__(self):\n functools.update_wrapper(self, self.f, updated=())\n\n @extract.treemap_copy_args\n def __call__(self, *args):\n updates, snapshot = extract.updates_and_snapshot(args)\n if self.graph:\n args = extract.from_tree2(args)\n out = self.f(*args)\n if self.graph:\n out = extract.to_tree2(out)\n extract.check_no_aliases('checkify', args=updates, out=out)\n updates = extract.mask_variable_updates(updates, snapshot)\n return out, updates\n\ndef checkify(\n f: tp.Callable[..., checkify_lib.Out],\n errors: frozenset[type[checkify_lib.JaxException]] = checkify_lib.user_checks, # type: ignore\n graph: bool | None = None,\n graph_updates: bool | None = None,\n) -> tp.Callable[..., tuple[checkify_lib.Error, checkify_lib.Out]]:\n \"\"\"Reference-aware version of `jax.experimental.checkify\n `_.\n\n Example::\n\n >>> import jax\n >>> import jax.numpy as jnp\n >>> from jax.experimental import checkify\n >>> import dataclasses\n >>> from flax import nnx\n ...\n >>> class Foo(nnx.Module):\n ... def __init__(self, a):\n ... self.a = nnx.Param(a)\n ...\n >>> @nnx.jit\n ... def f(m):\n ... y = jnp.sin(m.a) # error\n ... return m.a + y\n ...\n >>> m = Foo(a=jnp.inf)\n >>> err, out = nnx.checkify(f, errors=checkify.float_checks)(m)\n >>> # err.throw()\n >>> print(err)\n Error(nan generated by primitive: sin.)\n\n Args:\n f: the function to checkify.\n errors: the set of error checks to enable.\n graph: If ``True`` (default), uses graph-mode which supports the full\n NNX feature set including shared references and reference semantics.\n If ``False``, uses tree-mode which treats Modules as regular JAX\n pytrees, avoiding the overhead of the graph protocol.\n graph_updates: If ``True``, propagates updates on graph structure\n that happen inside the transform to the input graphs, has no\n effect when ``graph=False``.\n \"\"\"\n f_call, _, was_bound = _resolve_bound_callable(f)\n\n if was_bound:\n _raise_bound_method_error('checkify')\n\n if graph is None:\n graph = graphlib.set_graph_mode.current_value()\n if graph_updates is None:\n graph_updates = graphlib.set_graph_updates.current_value()\n if not graph or not graph_updates:\n checkify_fn = checkify_lib.checkify(\n SimpleCheckifyFn(f_call, graph=graph), errors,\n )\n\n @functools.wraps(f)\n def simple_checkify_wrapper(*args):\n if graph:\n args = extract.to_tree2(args)\n extract.check_no_aliases('checkify', args=args)\n error, (out, updates) = checkify_fn(*args)\n if graph:\n out = extract.from_tree2(out)\n extract.apply_variable_updates(args, updates)\n return error, out\n\n return simple_checkify_wrapper # type: ignore\n\n checkify_fn = checkify_lib.checkify(CheckifyFn(f_call), errors)\n @functools.wraps(f)\n @graphlib.update_context('checkify')\n def checkify_wrapper(*args, **kwargs):\n pure_args, pure_kwargs = extract.to_tree(\n (args, kwargs),\n ctxtag='checkify',\n )\n error, (pure_args_out, pure_kwargs_out, pure_out) = checkify_fn(\n *pure_args, **pure_kwargs\n )\n\n args_out, kwargs_out, out = extract.from_tree(\n (pure_args_out, pure_kwargs_out, pure_out),\n ctxtag='checkify',\n is_inner=False,\n )\n\n return error, out\n\n return checkify_wrapper # type: ignore\n\n\n@dataclasses.dataclass(eq=False)\nclass SimpleCondFn:\n f: tp.Callable[..., tp.Any]\n graph: bool\n\n def __post_init__(self):\n functools.update_wrapper(self, self.f, updated=())\n\n @extract.treemap_copy_args\n def __call__(self, *args):\n updates, _snapshot = extract.updates_and_snapshot(args)\n if self.graph:\n args = extract.from_tree2(args)\n out = self.f(*args)\n if self.graph:\n out = extract.to_tree2(out)\n extract.check_no_aliases('switch', args=updates, out=out)\n return out, updates\n\n\ndef cond(\n pred,\n true_fun: tp.Callable[..., A],\n false_fun: tp.Callable[..., A],\n *operands,\n graph: bool | None = None,\n graph_updates: bool | None = None,\n) -> A:\n \"\"\"Conditionally apply ``true_fun`` or ``false_fun``.\n\n Wraps `jax.lax.cond `__\n to support Flax NNX modules and variables.\n\n Args:\n pred: boolean scalar. If True, ``true_fun`` is applied, otherwise\n ``false_fun``.\n true_fun: function to apply if ``pred`` is True.\n false_fun: function to apply if ``pred`` is False.\n *operands: operands passed to whichever branch is selected.\n graph: If ``True`` (default), uses graph-mode which supports the full\n NNX feature set including shared references and reference semantics.\n If ``False``, uses tree-mode which treats Modules as regular JAX\n pytrees, avoiding the overhead of the graph protocol.\n graph_updates: If ``True``, propagates updates on graph structure\n that happen inside the transform to the input graphs, has no\n effect when ``graph=False``.\n \"\"\"\n if graph is None:\n graph = graphlib.set_graph_mode.current_value()\n if graph_updates is None:\n graph_updates = graphlib.set_graph_updates.current_value()\n if not graph or not graph_updates:\n if graph:\n operands = extract.to_tree2(operands)\n extract.check_no_aliases('cond', operands=operands)\n out, updates = jax.lax.cond(\n pred,\n SimpleCondFn(true_fun, graph=graph),\n SimpleCondFn(false_fun, graph=graph),\n *operands,\n )\n if graph:\n out = extract.from_tree2(out)\n extract.apply_variable_updates(operands, updates)\n return out\n\n @general.split_inputs(ctxtag='cond')\n def _cond(pred, true_fun, false_fun, *operands):\n return jax.lax.cond(\n pred,\n general.merge_inputs(true_fun, ctxtag='cond'),\n general.merge_inputs(false_fun, ctxtag='cond'),\n *operands,\n )\n\n return _cond(pred, true_fun, false_fun, *operands)\n\n\ndef switch(\n index,\n branches: tp.Sequence[tp.Callable[..., A]],\n *operands,\n graph: bool | None = None,\n graph_updates: bool | None = None,\n) -> A:\n \"\"\"Select and apply one of ``branches`` based on ``index``.\n\n Wraps `jax.lax.switch `__\n to support Flax NNX modules and variables.\n\n Args:\n index: integer scalar indicating which branch to apply.\n branches: sequence of functions to select from.\n *operands: operands passed to the selected branch.\n graph: If ``True`` (default), uses graph-mode which supports the full\n NNX feature set including shared references and reference semantics.\n If ``False``, uses tree-mode which treats Modules as regular JAX\n pytrees, avoiding the overhead of the graph protocol.\n graph_updates: If ``True``, propagates updates on graph structure\n that happen inside the transform to the input graphs, has no\n effect when ``graph=False``.\n \"\"\"\n if graph is None:\n graph = graphlib.set_graph_mode.current_value()\n if graph_updates is None:\n graph_updates = graphlib.set_graph_updates.current_value()\n if not graph or not graph_updates:\n if graph:\n operands = extract.to_tree2(operands)\n extract.check_no_aliases('switch', operands=operands)\n out, updates = jax.lax.switch(\n index,\n [SimpleCondFn(f, graph=graph) for f in branches],\n *operands,\n )\n if graph:\n out = extract.from_tree2(out)\n extract.apply_variable_updates(operands, updates)\n return out\n\n @general.split_inputs(ctxtag='switch')\n def _switch(index, branches, *operands):\n return jax.lax.switch(\n index,\n [general.merge_inputs(f, ctxtag='switch') for f in branches],\n *operands,\n )\n\n return _switch(index, branches, *operands)\n" }, { "path": "flax/nnx/traversals.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Utilities for flattening and unflattening mappings.\n\"\"\"\nfrom __future__ import annotations\n\nfrom collections.abc import Callable, Mapping\nfrom collections.abc import Iterable, Sequence\nfrom typing import Any, overload\n\nfrom flax import struct\n\n\n# the empty node is a struct.dataclass to be compatible with JAX.\n@struct.dataclass\nclass _EmptyNode:\n pass\n\n\nempty_node = _EmptyNode()\n\n\n# TODO: In Python 3.10, use TypeAlias.\nIsLeafCallable = Callable[[tuple[Any, ...], Mapping[Any, Any]], bool]\n\n\n@overload\ndef flatten_mapping(xs: Mapping[Any, Any],\n /,\n *,\n keep_empty_nodes: bool = False,\n is_leaf: None | IsLeafCallable = None,\n sep: None = None\n ) -> dict[tuple[Any, ...], Any]:\n ...\n\n@overload\ndef flatten_mapping(xs: Mapping[Any, Any],\n /,\n *,\n keep_empty_nodes: bool = False,\n is_leaf: None | IsLeafCallable = None,\n sep: str,\n ) -> dict[str, Any]:\n ...\n\ndef flatten_mapping(xs: Mapping[Any, Any],\n /,\n *,\n keep_empty_nodes: bool = False,\n is_leaf: None | IsLeafCallable = None,\n sep: None | str = None\n ) -> dict[Any, Any]:\n \"\"\"Flatten a nested mapping.\n\n The nested keys are flattened to a tuple. See ``unflatten_mapping`` on how to\n restore the nested mapping.\n\n Example::\n\n >>> from flax import nnx\n >>> xs = {'foo': 1, 'bar': {'a': 2, 'b': {}}}\n >>> flat_xs = nnx.traversals.flatten_mapping(xs)\n >>> flat_xs\n {('foo',): 1, ('bar', 'a'): 2}\n\n Note that empty mappings are ignored and will not be restored by\n ``unflatten_mapping``.\n\n Args:\n xs: a nested mapping\n keep_empty_nodes: replaces empty mappings with\n ``traverse_util.empty_node``.\n is_leaf: an optional function that takes the next nested mapping and nested\n keys and returns True if the nested mapping is a leaf (i.e., should not be\n flattened further).\n sep: if specified, then the keys of the returned mapping will be\n ``sep``-joined strings (if ``None``, then keys will be tuples).\n Returns:\n The flattened mapping.\n \"\"\"\n assert isinstance(\n xs, Mapping\n ), f'expected Mapping; got {type(xs).__qualname__}'\n\n def _key(path: tuple[Any, ...]) -> tuple[Any, ...] | str:\n if sep is None:\n return path\n return sep.join(path)\n\n def _flatten(xs: Any, prefix: tuple[Any, ...]) -> dict[Any, Any]:\n if not isinstance(xs, Mapping) or (\n is_leaf and is_leaf(prefix, xs)\n ):\n return {_key(prefix): xs}\n result = {}\n is_empty = True\n for key, value in xs.items():\n is_empty = False\n path = prefix + (key,)\n result.update(_flatten(value, path))\n if keep_empty_nodes and is_empty:\n if prefix == (): # when the whole input is empty\n return {}\n return {_key(prefix): empty_node}\n return result\n\n return _flatten(xs, ())\n\ndef flatten_to_sequence(\n xs: Mapping[Any, Any],\n /,\n *,\n is_leaf: IsLeafCallable | None = None,\n) -> list[tuple[Any, Any]]:\n \"\"\"Flatten a nested mapping.\n\n The nested keys are flattened to a tuple. See ``unflatten_mapping`` on how to\n restore the nested mapping.\n\n Example::\n\n >>> from flax import nnx\n >>> xs = {'foo': 1, 'bar': {'a': 2, 'b': {}}}\n >>> flat_xs = nnx.traversals.flatten_to_sequence(xs)\n >>> flat_xs\n [(('foo',), 1), (('bar', 'a'), 2)]\n\n Note that empty mappings are ignored and will not be restored by\n ``unflatten_mapping``.\n\n Args:\n xs: a nested mapping\n is_leaf: an optional function that takes the next nested mapping and nested\n keys and returns True if the nested mapping is a leaf (i.e., should not be\n flattened further).\n\n Returns:\n The flattened mapping.\n \"\"\"\n assert isinstance(\n xs, Mapping\n ), f'expected Mapping; got {type(xs).__qualname__}'\n result = []\n\n def _flatten(xs: Any, prefix: tuple[Any, ...]):\n if not isinstance(xs, Mapping) or (is_leaf and is_leaf(prefix, xs)):\n result.append((prefix, xs))\n else:\n for key, value in xs.items():\n _flatten(value, (*prefix, key))\n\n _flatten(xs, ())\n return result\n\n\n@overload\ndef unflatten_mapping(\n xs: Sequence[tuple[tuple[Any, ...], Any]], /, *, sep: None = None\n) -> dict[Any, Any]:\n ...\n\n\n@overload\ndef unflatten_mapping(\n xs: Mapping[tuple[Any, ...], Any], /, *, sep: None = None\n) -> dict[Any, Any]:\n ...\n\n\n@overload\ndef unflatten_mapping(xs: Mapping[str, Any], /, *, sep: str) -> dict[Any, Any]:\n ...\n\n\ndef unflatten_mapping(xs: Any, /, *, sep: str | None = None) -> dict[Any, Any]:\n \"\"\"Unflatten a mapping.\n\n See ``flatten_mapping``\n\n Example::\n\n >>> from flax import nnx\n >>> flat_xs = {\n ... ('foo',): 1,\n ... ('bar', 'a'): 2,\n ... }\n >>> xs = nnx.traversals.unflatten_mapping(flat_xs)\n >>> xs\n {'foo': 1, 'bar': {'a': 2}}\n\n Args:\n xs: a flattened mapping.\n sep: separator (same as used with ``flatten_mapping()``).\n Returns:\n The nested mapping.\n \"\"\"\n if isinstance(xs, Mapping):\n xs = xs.items()\n\n if not isinstance(xs, Iterable):\n raise TypeError(\n f'expected Mapping or Iterable; got {type(xs).__qualname__}'\n )\n result: dict[Any, Any] = {}\n for path, value in xs:\n if sep is not None:\n path = path.split(sep) # type: ignore\n if value is empty_node:\n value = {}\n cursor = result\n for key in path[:-1]:\n if key not in cursor:\n cursor[key] = {}\n cursor = cursor[key]\n cursor[path[-1]] = value\n return result\n" }, { "path": "flax/nnx/variablelib.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# pytype: skip-file\nfrom __future__ import annotations\n\nimport dataclasses\nimport functools\nfrom functools import partial\nimport itertools as it\nimport threading\nimport typing as tp\nfrom typing import Any\nimport warnings\n\nfrom flax import config\nfrom flax import errors\nfrom flax.core import spmd as core_spmd\nfrom flax.nnx import reprlib, tracers, visualization\nfrom flax.typing import BaseConfigContext, MISSING, Missing, SizeBytes\nimport jax\nfrom jax._src.state.types import AbstractRef\nimport jax.experimental\nfrom jax.experimental import hijax as hjx\nimport jax.tree_util as jtu\nimport treescope # type: ignore[import-untyped]\n\nA = tp.TypeVar('A')\nB = tp.TypeVar('B')\nC = tp.TypeVar('C')\nF = tp.TypeVar('F', bound=tp.Callable[..., tp.Any])\nP = tp.TypeVar('P', bound=property)\nV = tp.TypeVar('V', bound='Variable[Any]')\nGetValueHook = tp.Callable[['Variable[A]', A], A]\nSetValueHook = tp.Callable[['Variable[A]', A], A]\nCreateValueHook = tp.Callable[['Variable[A]', A], A]\nAxisName = str\nAxisIndex = int\nAddAxisHook = tp.Callable[[V, AxisIndex, AxisName | None], None]\nRemoveAxisHook = tp.Callable[[V, AxisIndex, AxisName | None], None]\n\n# JAX array refs were renamed a few times between JAX v0.7.0 and v0.8.0.\n# The following ensures we avoid an ImportError or DeprecationWarning.\nif hasattr(jax, 'new_ref') and hasattr(jax, 'Ref'):\n # JAX v0.7.2 or newer\n from jax import Ref\nelif hasattr(jax, 'array_ref') and hasattr(jax, 'ArrayRef'):\n # JAX v0.7.1\n from jax import ArrayRef as Ref # type: ignore[import-untyped,no-redef]\nelse:\n # JAX v0.7.0 or older\n from jax.experimental import MutableArray as Ref # type: ignore[no-redef]\n\n\n@dataclasses.dataclass\nclass VariableContext(threading.local):\n variable_hijax_stack: list[bool] = dataclasses.field(default_factory=list)\n variable_ref_stack: list[bool] = dataclasses.field(default_factory=list)\n eager_shard_stack: list[bool] = dataclasses.field(default_factory=list)\n\n\nVARIABLE_CONTEXT = VariableContext()\n\n\nclass use_eager_sharding(BaseConfigContext):\n \"\"\"Sets whether Variables should use eager sharding by default or not.\n\n Example usage::\n\n >>> from flax import nnx\n >>> # Use eager sharding by default\n >>> nnx.use_eager_sharding(True)\n <...>\n >>> # Variable will now use eager sharding\n >>> nnx.using_eager_sharding()\n True\n\n It can also be used as a context manager to temporarily\n change the default behavior for a block of code::\n\n >>> nnx.use_eager_sharding(False)\n <...>\n >>> with nnx.use_eager_sharding(True):\n ... nnx.using_eager_sharding()\n True\n >>> # it will reset outside\n >>> v = nnx.Variable(jax.numpy.ones((2, 3)))\n >>> nnx.using_eager_sharding()\n False\n\n Args:\n value: A boolean indicating if Variables should use eager sharding by default.\n\n Returns:\n A context manager that resets the context to the previous value.\n \"\"\"\n get_default = classmethod(lambda cls: config.flax_always_shard_variable)\n get_stack = classmethod(lambda cls: VARIABLE_CONTEXT.eager_shard_stack)\n\n\ndef using_eager_sharding() -> bool:\n \"\"\"Returns whether Variables are using eager sharding by default.\n\n Example::\n\n >>> from flax import nnx\n >>> nnx.use_eager_sharding(True)\n <...>\n >>> nnx.using_eager_sharding()\n True\n >>> nnx.use_eager_sharding(False)\n <...>\n >>> nnx.using_eager_sharding()\n False\n\n\n Returns:\n A boolean indicating if Variables are using eager sharding by default.\n \"\"\"\n return use_eager_sharding.current_value()\n\n\n@dataclasses.dataclass(frozen=True)\nclass VarDefaults(tp.Mapping[str, tp.Any]):\n hijax: bool\n ref: bool\n\n def __getitem__(self, key: str) -> tp.Any:\n return getattr(self, key)\n\n def __iter__(self) -> tp.Iterator[str]:\n return iter(dataclasses.asdict(self))\n\n def __len__(self) -> int:\n return len(dataclasses.fields(self))\n\n\n@tp.overload\ndef var_defaults() -> VarDefaults: ...\n\n\n@tp.overload\ndef var_defaults(\n *, hijax: bool | None = None, ref: bool | None = None\n) -> VarDefaultsContext: ...\n\n\ndef var_defaults(\n *, hijax: bool | None = None, ref: bool | None = None\n) -> VarDefaultsContext | VarDefaults:\n if hijax is None and ref is None:\n return VarDefaults(\n hijax=VARIABLE_CONTEXT.variable_hijax_stack[-1]\n if VARIABLE_CONTEXT.variable_hijax_stack\n else config.flax_hijax_variable,\n ref=VARIABLE_CONTEXT.variable_ref_stack[-1]\n if VARIABLE_CONTEXT.variable_ref_stack\n else False,\n )\n\n hijax_prev = None\n if hijax is not None:\n if VARIABLE_CONTEXT.variable_hijax_stack:\n hijax_prev = VARIABLE_CONTEXT.variable_hijax_stack[-1]\n VARIABLE_CONTEXT.variable_hijax_stack[-1] = hijax\n else:\n VARIABLE_CONTEXT.variable_hijax_stack.append(hijax)\n\n ref_prev = None\n if ref is not None:\n if VARIABLE_CONTEXT.variable_ref_stack:\n ref_prev = VARIABLE_CONTEXT.variable_ref_stack[-1]\n VARIABLE_CONTEXT.variable_ref_stack[-1] = ref\n else:\n VARIABLE_CONTEXT.variable_ref_stack.append(ref)\n\n return VarDefaultsContext(\n hijax_prev=hijax_prev,\n hijax_new=hijax,\n ref_prev=ref_prev,\n ref_new=ref,\n )\n\n\nclass VarDefaultsContext:\n def __init__(\n self,\n *,\n hijax_prev: bool | None,\n hijax_new: bool | None,\n ref_prev: bool | None,\n ref_new: bool | None,\n ):\n self.hijax_prev = hijax_prev\n self.hijax_new = hijax_new\n self.ref_prev = ref_prev\n self.ref_new = ref_new\n\n def __enter__(self):\n if self.hijax_new is not None and self.hijax_prev is not None:\n VARIABLE_CONTEXT.variable_hijax_stack.insert(-1, self.hijax_prev)\n if self.ref_new is not None and self.ref_prev is not None:\n VARIABLE_CONTEXT.variable_ref_stack.insert(-1, self.ref_prev)\n\n def __exit__(self, exc_type, exc_value, traceback):\n if self.hijax_new is not None:\n VARIABLE_CONTEXT.variable_hijax_stack.pop()\n if self.ref_new is not None:\n VARIABLE_CONTEXT.variable_ref_stack.pop()\n\n def __call__(self, f: F) -> F:\n # undo stack change for decorator usage\n if self.hijax_new is not None:\n VARIABLE_CONTEXT.variable_hijax_stack.pop()\n if self.hijax_prev is not None:\n VARIABLE_CONTEXT.variable_hijax_stack.append(self.hijax_prev)\n\n if self.ref_new is not None:\n VARIABLE_CONTEXT.variable_ref_stack.pop()\n if self.ref_prev is not None:\n VARIABLE_CONTEXT.variable_ref_stack.append(self.ref_prev)\n\n @functools.wraps(f)\n def var_defaults_wrapper(*args, **kwargs):\n if self.hijax_new is not None:\n VARIABLE_CONTEXT.variable_hijax_stack.append(self.hijax_new)\n if self.ref_new is not None:\n VARIABLE_CONTEXT.variable_ref_stack.append(self.ref_new)\n try:\n return f(*args, **kwargs)\n finally:\n if self.hijax_new is not None:\n VARIABLE_CONTEXT.variable_hijax_stack.pop()\n if self.ref_new is not None:\n VARIABLE_CONTEXT.variable_ref_stack.pop()\n\n return var_defaults_wrapper # type: ignore[return-value]\n\n\ndef is_array_ref(x) -> tp.TypeGuard[Ref]:\n return isinstance(x, jax.Array | AbstractRef | Ref) and isinstance(\n jax.typeof(x), AbstractRef | Ref\n )\n\n\n@dataclasses.dataclass\nclass VariableMetadata(tp.Generic[A]):\n raw_value: A\n set_value_hooks: tuple[SetValueHook[A], ...] = ()\n get_value_hooks: tuple[GetValueHook[A], ...] = ()\n create_value_hooks: tuple[CreateValueHook[A], ...] = ()\n add_axis_hooks: tuple[AddAxisHook[Variable[A]], ...] = ()\n remove_axis_hooks: tuple[RemoveAxisHook[Variable[A]], ...] = ()\n metadata: tp.Mapping[str, tp.Any] = dataclasses.field(default_factory=dict)\n\n\nPyTreeDef = tp.Any\nLeaf = tp.Any\n\n# ---------------------------------\n# hijax\n# ---------------------------------\n\n\n@dataclasses.dataclass(frozen=True)\nclass VariableQDD:\n leaf_avals: tuple[hjx.AbstractValue, ...]\n treedef: PyTreeDef\n var_type: type[Variable[Any]]\n\n def to_tangent_qdd(self):\n leaf_avals = tuple(a.to_tangent_aval() for a in self.leaf_avals)\n return VariableQDD(leaf_avals, self.treedef, self.var_type)\n\n def normalize(self):\n leaf_types = tuple(a.normalize() for a in self.leaf_avals)\n return VariableQDD(leaf_types, self.treedef, self.var_type)\n\nclass VariableEffect(jax.core.Effect): ...\n\n\nvariable_effect = VariableEffect()\nhjx.control_flow_allowed_effects.add_type(VariableEffect)\n\n\ndef _bind_new_variable(\n *leaves, treedef, var_type, has_qdd, ref\n) -> HijaxVariable:\n \"\"\"Binds new_variable_p after instantiating any Zero tangents.\"\"\"\n leaves = tuple(hjx.instantiate_zeros(leaf) for leaf in leaves)\n return new_variable_p.bind(\n *leaves,\n treedef=treedef,\n var_type=var_type,\n has_qdd=has_qdd,\n ref=ref,\n )\n\n\ndef _new_hijax_from_variable(variable: Variable) -> HijaxVariable:\n has_qdd = not variable.ref\n leaves, treedef = jax.tree.flatten(variable)\n var_type = type(variable)\n hijax_var = _bind_new_variable(\n *leaves,\n treedef=treedef,\n var_type=var_type,\n has_qdd=has_qdd,\n ref=variable.ref,\n )\n return hijax_var\n\n\nclass NewVariable(hjx.HiPrimitive):\n def is_high(self, *leaves, treedef, var_type, has_qdd, ref) -> bool:\n return True # type: ignore\n\n def impl(self, *leaves, treedef, var_type, has_qdd, ref):\n return HijaxVariable._new(\n leaves, treedef, var_type, has_qdd, ref=ref\n )\n\n def abstract_eval(self, *leaves, treedef, var_type, has_qdd, ref):\n aval = AbstractVariable(\n var_type, treedef, leaves, has_qdd, ref=ref\n )\n if has_qdd:\n qdd = VariableQDD(tuple(leaves), treedef, var_type)\n aval_qdd = hjx.AvalQDD(aval, qdd) # type: ignore\n return aval_qdd, {variable_effect}\n else:\n return aval, set()\n\n def to_lojax(self, *leaves, treedef, var_type, has_qdd, ref):\n return HijaxVariable._new(leaves, treedef, var_type, has_qdd, ref=ref)\n\n def jvp(_, primals, tangents, *, treedef, var_type, has_qdd, ref):\n if has_qdd:\n raise NotImplementedError(\n \"jvp not implemented for 'new_variable' with QDD\"\n )\n primal_hijax_var = _bind_new_variable(\n *primals, treedef=treedef, var_type=var_type, has_qdd=has_qdd, ref=ref\n )\n tangent_hijax_var = _bind_new_variable(\n *tangents, treedef=treedef, var_type=var_type, has_qdd=has_qdd, ref=ref\n )\n return primal_hijax_var, tangent_hijax_var\n\n def transpose(\n _, out_var: HijaxVariable, *input_leaves, treedef, var_type, has_qdd, ref\n ):\n if has_qdd:\n raise NotImplementedError(\n \"transpose not implemented for 'new_variable' with QDD\"\n )\n avals = tuple(\n map(\n lambda x: x.aval if hjx.is_undefined_primal(x) else jax.typeof(x),\n input_leaves,\n )\n )\n leaves_dot = get_variable_p.bind(\n out_var,\n treedef=treedef,\n avals=avals,\n var_type=var_type,\n has_qdd=has_qdd,\n )\n return leaves_dot\n\n\nnew_variable_p = NewVariable(f'new_variable')\n\n\ndef _set_hijax_state(hijax_var, variable: Variable):\n leaves, treedef = jax.tree.flatten(variable)\n set_variable_p.bind(\n hijax_var, *leaves, treedef=treedef, var_type=type(variable)\n )\n\n\nclass SetVariable(hjx.HiPrimitive):\n multiple_results = True\n\n def is_high(_, *leaf_avals, treedef, var_type) -> bool:\n return True # type: ignore\n\n # TODO: upstream this to Box\n def impl(_, hijax_var: HijaxVariable, *leaves, treedef, var_type):\n if not hijax_var.has_qdd:\n raise errors.ImmutableVariableError(\n \"Trying to update Variable with 'has_qdd=False'.\"\n )\n assert var_type is hijax_var._var_type\n object.__setattr__(hijax_var, '_leaves', leaves)\n object.__setattr__(hijax_var, '_treedef', treedef)\n return []\n\n def abstract_eval(\n _, aval_mutable_qdd: hjx.AvalMutableQDD, *leaf_avals, treedef, var_type\n ):\n hijax_var: AbstractVariable = aval_mutable_qdd.aval # type: ignore\n assert isinstance(hijax_var, AbstractVariable)\n if not hijax_var.has_qdd:\n raise errors.ImmutableVariableError(\n \"Trying to update Variable with 'has_qdd=False'.\"\n )\n assert var_type is hijax_var._var_type\n aval_mutable_qdd.mutable_qdd.update(\n VariableQDD(leaf_avals, treedef, var_type)\n )\n effects = {variable_effect} if hijax_var.has_qdd else set()\n return [], effects # TODO better typechecking...\n\n def to_lojax(_, hijax_var: HijaxVariable, *leaves, treedef, var_type):\n if not hijax_var.has_qdd:\n raise errors.ImmutableVariableError(\n \"Trying to update Variable with 'has_qdd=False'.\"\n )\n assert var_type is hijax_var._var_type\n object.__setattr__(hijax_var, '_leaves', leaves)\n object.__setattr__(hijax_var, '_treedef', treedef)\n return []\n\n def jvp(_, primals, tangents, *, treedef, var_type):\n variable: Variable\n variable, *vals = primals\n variable_dot: Variable\n variable_dot, *val_dots = tangents\n if type(variable_dot._raw_value) is hjx.Zero:\n raise Exception(\n \"can't differentiate Variable._set operation, \"\n 'did you forget jax.lax.stop_gradient?'\n )\n set_variable_p.bind(\n variable, *vals, treedef=treedef, var_type=type(variable)\n )\n set_variable_p.bind(\n variable_dot, *val_dots, treedef=treedef, var_type=type(variable_dot)\n )\n return [], []\n\n def transpose(_, *args, treedef, var_type):\n raise NotImplementedError('transpose not implemented for SetHijaxVariable')\n\n\nset_variable_p = SetVariable(f'set_variable')\n\n\ndef _get_hijax_state(hijax_var: HijaxVariable | AbstractVariable) -> Variable:\n if hijax_var.has_qdd:\n tys: VariableQDD = jax.experimental.cur_qdd(hijax_var)\n leaf_vals = get_variable_p.bind(\n hijax_var,\n treedef=tys.treedef,\n avals=tuple(tys.leaf_avals),\n var_type=hijax_var._var_type,\n has_qdd=hijax_var.has_qdd,\n )\n variable = jax.tree.unflatten(tys.treedef, leaf_vals)\n else:\n assert hijax_var._treedef is not None\n assert hijax_var._leaves is not None\n if isinstance(hijax_var, (jax.core.Tracer, AbstractVariable)):\n leaf_avals = hijax_var._leaves\n else:\n leaf_avals = tuple(map(jax.typeof, hijax_var._leaves))\n leaf_vals = get_variable_p.bind(\n hijax_var,\n treedef=hijax_var._treedef,\n avals=leaf_avals,\n var_type=hijax_var._var_type,\n has_qdd=hijax_var.has_qdd,\n )\n variable = jax.tree.unflatten(hijax_var._treedef, leaf_vals)\n\n return variable\n\n\nclass GetVariable(hjx.HiPrimitive):\n multiple_results = True\n\n def impl(\n self, hijax_var: HijaxVariable, *, treedef, avals, var_type, has_qdd\n ):\n return hijax_var._leaves\n\n def abstract_eval(self, abstract_var, *, treedef, avals, var_type, has_qdd):\n if has_qdd:\n return avals, {variable_effect}\n else:\n return avals, set()\n\n def to_lojax(\n _, hijax_var: HijaxVariable, *, treedef, avals, var_type, has_qdd\n ):\n return hijax_var._leaves\n\n def jvp(_, primals, tangents, *, treedef, avals, var_type, has_qdd):\n if has_qdd:\n raise NotImplementedError(\n \"jvp not implemented for 'get_variable' with QDD\"\n )\n (hijax_var,), (hijax_var_dot,) = primals, tangents\n return (\n get_variable_p.bind(\n hijax_var,\n treedef=treedef,\n avals=avals,\n var_type=var_type,\n has_qdd=has_qdd,\n ),\n get_variable_p.bind(\n hijax_var_dot,\n treedef=treedef,\n avals=tuple(a.to_tangent_aval() for a in avals),\n var_type=var_type,\n has_qdd=has_qdd,\n ),\n )\n\n def transpose(_, out, hijax_var, *, treedef, avals, var_type, has_qdd):\n if has_qdd:\n raise NotImplementedError(\n \"transpose not implemented for 'get_variable' with QDD\"\n )\n abstract_var: AbstractVariable = (\n hijax_var.aval\n if hjx.is_undefined_primal(hijax_var)\n else jax.typeof(hijax_var)\n )\n hijax_var_dot = _bind_new_variable(\n *out,\n treedef=abstract_var._treedef,\n var_type=var_type,\n has_qdd=has_qdd,\n ref=abstract_var.ref,\n )\n return (hijax_var_dot,)\n\n\nget_variable_p = GetVariable(f'get_variable')\n\n\n# ---------------------------------\n# HijaxVariable\n# ---------------------------------\ndef _variable_has_changed(old: Variable, new: Variable) -> bool:\n old_structure = jax.tree.structure(old)\n new_structure = jax.tree.structure(new)\n if old_structure != new_structure: # type: ignore[operator]\n return True\n old_leaves = jax.tree.leaves(old)\n new_leaves = jax.tree.leaves(new)\n return any(o is not n for o, n in zip(old_leaves, new_leaves))\n\n\ndef _as_hijax_property(name: str, *, get: bool, set: bool) -> property:\n \"\"\"Creates a property that operates on the hijax type.\"\"\"\n\n def _getter_wrapper(hijax_var):\n variable = _get_hijax_state(hijax_var)\n old_state = jax.tree.map(lambda x: x, variable)\n out = getattr(variable, name)\n if _variable_has_changed(old_state, variable):\n _set_hijax_state(hijax_var, variable)\n return out\n\n def _setter_wrapper(hijax_var, value):\n variable = _get_hijax_state(hijax_var)\n setattr(variable, name, value)\n _set_hijax_state(hijax_var, variable)\n\n _hijax_property = property(\n fget=_getter_wrapper if get else None,\n fset=_setter_wrapper if set else None,\n )\n return _hijax_property # type: ignore[return]\n\n\ndef _as_aval_property(p: property) -> hjx.aval_property:\n \"\"\"Wraps a property `p` operate on the aval type.\"\"\"\n _aval_property = hjx.aval_property(fget=p.fget)\n return _aval_property # type: ignore[return]\n\n\ndef _as_hijax_attribute(name: str) -> property:\n \"\"\"Creates a property that operates on the hijax type.\"\"\"\n\n def _getter_wrapper(hijax_var):\n variable = _get_hijax_state(hijax_var)\n old_state = jax.tree.map(lambda x: x, variable)\n out = getattr(variable, name)\n if _variable_has_changed(old_state, variable):\n _set_hijax_state(hijax_var, variable)\n return out\n\n _getter_wrapper.__name__ = name\n _hijax_property = property(fget=_getter_wrapper)\n\n return _hijax_property # type: ignore[return]\n\n\ndef _as_hijax_method(name: str) -> tp.Any:\n \"\"\"Creates a method that operates on the hijax type.\"\"\"\n\n def hijax_method_wrapper(hijax_var, *args, **kwargs):\n variable = _get_hijax_state(hijax_var)\n old_state = jax.tree.map(lambda x: x, variable)\n method = getattr(variable, name)\n out = method(*args, **kwargs)\n if _variable_has_changed(old_state, variable):\n _set_hijax_state(hijax_var, variable)\n return out\n\n hijax_method_wrapper.__name__ = name\n\n return hijax_method_wrapper\n\n\ndef _as_tracer_method(name: str):\n def op(self, hijax_var, *args, **kwargs):\n variable = _get_hijax_state(hijax_var)\n old_state = jax.tree.map(lambda x: x, variable)\n out = getattr(variable, name)(*args, **kwargs)\n if _variable_has_changed(old_state, variable):\n _set_hijax_state(hijax_var, variable)\n return out\n\n op.__name__ = name\n return op\n\ndef _not_an_attribute_property(name: str):\n def _op(self):\n raise AttributeError(\n f\"'{type(self).__name__}' object has no attribute '{name}'\"\n )\n\n return property(_op)\n\nclass HijaxVariableMeta(type):\n def __instancecheck__(self, instance):\n if super().__instancecheck__(instance):\n return True\n\n if isinstance(instance, jax.core.Tracer):\n ty = jax.typeof(instance)\n return isinstance(ty, AbstractVariable)\n return False\n\n\nclass HijaxVariable(\n tp.Generic[A], reprlib.Representable, metaclass=HijaxVariableMeta\n): # type: ignore\n __slots__ = ('_treedef', '_leaves', '_var_type', 'has_qdd', '_ref')\n _treedef: PyTreeDef\n _leaves: tuple[Leaf, ...]\n _var_type: type[Variable[tp.Any]]\n has_qdd: bool\n _ref: bool\n\n @classmethod\n def _new(\n cls,\n leaves: tuple[Leaf, ...],\n treedef: PyTreeDef,\n var_type: type[Variable[A]],\n has_qdd: bool,\n *,\n ref: bool = False,\n ):\n hijax_var = object.__new__(cls)\n object.__setattr__(hijax_var, '_treedef', treedef)\n object.__setattr__(hijax_var, '_leaves', leaves)\n object.__setattr__(hijax_var, '_var_type', var_type)\n object.__setattr__(hijax_var, 'has_qdd', has_qdd)\n object.__setattr__(hijax_var, '_ref', ref)\n return hijax_var\n\n __init__ = _as_hijax_method('__init__')\n\n @property\n def value(self) -> A:\n raise NotImplementedError(\n 'HijaxVariable.value property is not implemented. For Variable[Array] instances use:\\n\\n'\n ' variable[...]\\n\\n'\n 'For other Variable types use:\\n\\n'\n ' variable.get_value()\\n'\n )\n\n @value.setter\n def value(self, new_value: A):\n raise NotImplementedError(\n 'HijaxVariable.value property is not implemented. For Variable[Array] instances use:\\n\\n'\n ' variable[...] = new_value\\n\\n'\n 'For other Variable types use:\\n\\n'\n ' variable.set_value(new_value)\\n'\n )\n\n @property\n def var_type(self) -> type[Variable[A]]:\n return self._var_type\n\n _trace_state = _as_hijax_property('_trace_state', get=True, set=False)\n _can_update = _as_hijax_property('_can_update', get=True, set=False)\n _check_can_update = _as_hijax_method('_check_can_update')\n __getattr__ = _as_hijax_method('__getattr__')\n __setattr__ = _as_hijax_method('__setattr__')\n __delattr__ = _as_hijax_method('__delattr__')\n type = _as_hijax_property('type', get=True, set=False)\n type = _as_hijax_property('type', get=True, set=False)\n hijax = _as_hijax_property('hijax', get=True, set=False)\n\n @property\n def ref(self) -> bool:\n return self._ref\n\n get_metadata = _as_hijax_method('get_metadata')\n set_metadata = _as_hijax_method('set_metadata')\n\n def copy_from(self, other: Variable[A] | HijaxVariable[A]) -> None:\n if isinstance(other, HijaxVariable):\n other = _get_hijax_state(other)\n variable = _get_hijax_state(self)\n variable.copy_from(other) # type: ignore[arg-type]\n _set_hijax_state(self, variable)\n\n def update_from_state(self, variable_state: Variable[A] | HijaxVariable[A]):\n if isinstance(variable_state, HijaxVariable):\n variable_state = _get_hijax_state(variable_state)\n variable = _get_hijax_state(self)\n variable.update_from_state(variable_state) # type: ignore[arg-type]\n _set_hijax_state(self, variable)\n\n get_raw_value = _as_hijax_method('get_raw_value')\n set_raw_value = _as_hijax_method('set_raw_value')\n set_value = _as_hijax_method('set_value')\n get_value = _as_hijax_method('get_value')\n create_value = _as_hijax_method('create_value')\n set_raw_value = _as_hijax_method('set_raw_value')\n add_axis = _as_hijax_method('add_axis')\n remove_axis = _as_hijax_method('remove_axis')\n copy = _as_hijax_method('copy')\n replace = _as_hijax_method('replace')\n to_state = _as_hijax_method('to_state')\n\n @classmethod\n def from_metadata(cls, value: A, metadata: dict[str, tp.Any]):\n return cls._var_type.from_metadata(value, metadata) # type: ignore[misc]\n\n __nnx_repr__ = _as_hijax_method('__nnx_repr__')\n __treescope_repr__ = _as_hijax_method('__treescope_repr__')\n\n # --------------------------------------------\n # proxy methods\n # --------------------------------------------\n __jax_array__ = _as_hijax_method('__jax_array__')\n __getitem__ = _as_hijax_method('__getitem__')\n __setitem__ = _as_hijax_method('__setitem__')\n __delitem__ = _as_hijax_method('__delitem__')\n __call__ = _as_hijax_method('__call__')\n __len__ = _as_hijax_method('__len__')\n __iter__ = _as_hijax_method('__iter__')\n __contains__ = _as_hijax_method('__contains__')\n __add__ = _as_hijax_method('__add__')\n __sub__ = _as_hijax_method('__sub__')\n __mul__ = _as_hijax_method('__mul__')\n __matmul__ = _as_hijax_method('__matmul__')\n __truediv__ = _as_hijax_method('__truediv__')\n __floordiv__ = _as_hijax_method('__floordiv__')\n __mod__ = _as_hijax_method('__mod__')\n __divmod__ = _as_hijax_method('__divmod__')\n __pow__ = _as_hijax_method('__pow__')\n __lshift__ = _as_hijax_method('__lshift__')\n __rshift__ = _as_hijax_method('__rshift__')\n __and__ = _as_hijax_method('__and__')\n __xor__ = _as_hijax_method('__xor__')\n __or__ = _as_hijax_method('__or__')\n __radd__ = _as_hijax_method('__radd__')\n __rsub__ = _as_hijax_method('__rsub__')\n __rmul__ = _as_hijax_method('__rmul__')\n __rmatmul__ = _as_hijax_method('__rmatmul__')\n __rtruediv__ = _as_hijax_method('__rtruediv__')\n __rfloordiv__ = _as_hijax_method('__rfloordiv__')\n __rmod__ = _as_hijax_method('__rmod__')\n __rdivmod__ = _as_hijax_method('__rdivmod__')\n __rpow__ = _as_hijax_method('__rpow__')\n __rlshift__ = _as_hijax_method('__rlshift__')\n __rrshift__ = _as_hijax_method('__rrshift__')\n __rand__ = _as_hijax_method('__rand__')\n __rxor__ = _as_hijax_method('__rxor__')\n __ror__ = _as_hijax_method('__ror__')\n __iadd__ = _as_hijax_method('__iadd__')\n __isub__ = _as_hijax_method('__isub__')\n __imul__ = _as_hijax_method('__imul__')\n __imatmul__ = _as_hijax_method('__imatmul__')\n __itruediv__ = _as_hijax_method('__itruediv__')\n __ifloordiv__ = _as_hijax_method('__ifloordiv__')\n __imod__ = _as_hijax_method('__imod__')\n __ipow__ = _as_hijax_method('__ipow__')\n __ilshift__ = _as_hijax_method('__ilshift__')\n __irshift__ = _as_hijax_method('__irshift__')\n __iand__ = _as_hijax_method('__iand__')\n __ixor__ = _as_hijax_method('__ixor__')\n __ior__ = _as_hijax_method('__ior__')\n __neg__ = _as_hijax_method('__neg__')\n __pos__ = _as_hijax_method('__pos__')\n __abs__ = _as_hijax_method('__abs__')\n __invert__ = _as_hijax_method('__invert__')\n __complex__ = _as_hijax_method('__complex__')\n __int__ = _as_hijax_method('__int__')\n __float__ = _as_hijax_method('__float__')\n __index__ = _as_hijax_method('__index__')\n __round__ = _as_hijax_method('__round__')\n __trunc__ = _as_hijax_method('__trunc__')\n __floor__ = _as_hijax_method('__floor__')\n __ceil__ = _as_hijax_method('__ceil__')\n\n # --------------------------------------------\n # hijax interface\n # --------------------------------------------\n\n def cur_qdd(self):\n return self.type_state()\n\n def type_state(self):\n leaf_avals = tuple(map(jax.typeof, self._leaves))\n return VariableQDD(leaf_avals, self._treedef, self._var_type)\n\n\ndef _to_abstract_variable(hijax_var: HijaxVariable):\n if hijax_var.has_qdd:\n treedef = None\n leaves = None\n else:\n leaves = tuple(map(jax.typeof, hijax_var._leaves))\n treedef = hijax_var._treedef\n return AbstractVariable(\n hijax_var._var_type,\n treedef,\n leaves,\n hijax_var.has_qdd,\n ref=hijax_var.ref,\n )\n\n\nhjx.register_hitype(HijaxVariable, _to_abstract_variable)\n\n\n# ---------------------------------\n# AbstractVariable\n# ---------------------------------\nclass AbstractVariable(tp.Generic[A], hjx.MutableHiType):\n __slots__ = ['_var_type', '_treedef', '_leaves', 'has_qdd', '_ref']\n _var_type: type[Variable[A]]\n _treedef: PyTreeDef | None\n _leaves: tuple[hjx.AbstractValue, ...] | None\n has_qdd: bool\n _ref: bool\n\n @property\n def ref(self) -> bool:\n return self._ref\n\n @property\n def hijax(self):\n return True\n\n _check_can_update = hjx.aval_method(HijaxVariable._check_can_update)\n\n def __init__(\n self,\n var_type: type[Variable[A]],\n treedef: PyTreeDef | None,\n leaves: tuple[hjx.AbstractValue, ...] | None,\n has_qdd: bool,\n *,\n ref: bool = False,\n ):\n if (treedef is None) ^ (leaves is None):\n raise ValueError('treedef and leaves must be both provided or both None')\n object.__setattr__(self, '_treedef', treedef)\n object.__setattr__(self, '_leaves', leaves)\n object.__setattr__(self, '_var_type', var_type)\n object.__setattr__(self, 'has_qdd', has_qdd)\n object.__setattr__(self, '_ref', ref)\n\n @property\n def dtype(self):\n raise AttributeError\n\n @property\n def ndim(self):\n raise AttributeError\n\n @property\n def size(self):\n raise AttributeError\n\n @property\n def shape(self):\n raise AttributeError\n\n def __getattr__(self, name: str):\n # Forward unknown attributes to the value\n if hasattr(AbstractVariable, name):\n raise AttributeError(\n f\"'{type(self).__name__}' object has no attribute '{name}'\"\n )\n if name.startswith('_'):\n raise AttributeError(\n f\"'{type(self).__name__}' object has no attribute '{name}'\"\n )\n return _as_aval_property(_as_hijax_attribute(name))\n\n # __setattr__ supported via __getattr__\n # __delattr__ CURRENTLY NOT SUPPORTED\n type = _as_aval_property(HijaxVariable.type)\n get_metadata = hjx.aval_method(HijaxVariable.get_metadata)\n set_metadata = hjx.aval_method(HijaxVariable.set_metadata)\n copy_from = hjx.aval_method(HijaxVariable.copy_from)\n update_from_state = hjx.aval_method(HijaxVariable.update_from_state)\n get_raw_value = hjx.aval_method(HijaxVariable.get_raw_value)\n set_raw_value = hjx.aval_method(HijaxVariable.set_raw_value)\n set_value = hjx.aval_method(HijaxVariable.set_value)\n get_value = hjx.aval_method(HijaxVariable.get_value)\n create_value = hjx.aval_method(HijaxVariable.create_value)\n set_raw_value = hjx.aval_method(HijaxVariable.set_raw_value)\n add_axis = hjx.aval_method(HijaxVariable.add_axis)\n remove_axis = hjx.aval_method(HijaxVariable.remove_axis)\n replace = hjx.aval_method(HijaxVariable.replace)\n\n @hjx.aval_method\n def from_metadata(self, value, metadata: dict[str, tp.Any]):\n aval: AbstractVariable = self.aval # type: ignore\n variable = aval._var_type.from_metadata(value, metadata)\n return variable\n\n copy = hjx.aval_method(HijaxVariable.copy)\n replace = hjx.aval_method(HijaxVariable.replace)\n to_state = hjx.aval_method(HijaxVariable.to_state)\n\n def __str__(self):\n return f'{self._var_type.__name__}()'\n\n def __repr__(self):\n return f'{self._var_type.__name__}()'\n\n @hjx.aval_method\n def __treescope_repr__(self, path, subtree_renderer):\n raise NotImplementedError\n\n # ---------------------------------\n # proxy methods\n # ---------------------------------\n __jax_array__ = hjx.aval_method(HijaxVariable.__jax_array__)\n _getitem = _as_tracer_method('__getitem__')\n _setitem = _as_tracer_method('__setitem__')\n # __delitem__ CURRENTLY NOT SUPPORTED\n # __call__ CURRENTLY NOT SUPPORTED\n _len = _as_tracer_method('__len__')\n _iter = _as_tracer_method('__iter__')\n # __contains__ CURRENTLY NOT SUPPORTED\n _add = _as_tracer_method('__add__')\n _sub = _as_tracer_method('__sub__')\n _mul = _as_tracer_method('__mul__')\n _matmul = _as_tracer_method('__matmul__')\n _truediv = _as_tracer_method('__truediv__')\n _floordiv = _as_tracer_method('__floordiv__')\n _mod = _as_tracer_method('__mod__')\n _divmod = _as_tracer_method('__divmod__')\n _pow = _as_tracer_method('__pow__')\n _lshift = _as_tracer_method('__lshift__')\n _rshift = _as_tracer_method('__rshift__')\n _and = _as_tracer_method('__and__')\n _xor = _as_tracer_method('__xor__')\n _or = _as_tracer_method('__or__')\n _radd = _as_tracer_method('__radd__')\n _rsub = _as_tracer_method('__rsub__')\n _rmul = _as_tracer_method('__rmul__')\n _rmatmul = _as_tracer_method('__rmatmul__')\n _rtruediv = _as_tracer_method('__rtruediv__')\n _rfloordiv = _as_tracer_method('__rfloordiv__')\n _rmod = _as_tracer_method('__rmod__')\n _rdivmod = _as_tracer_method('__rdivmod__')\n _rpow = _as_tracer_method('__rpow__')\n _rlshift = _as_tracer_method('__rlshift__')\n _rrshift = _as_tracer_method('__rrshift__')\n _rand = _as_tracer_method('__rand__')\n _rxor = _as_tracer_method('__rxor__')\n _ror = _as_tracer_method('__ror__')\n # _iadd CURRENTLY NOT SUPPORTED\n # _isub CURRENTLY NOT SUPPORTED\n # _imul CURRENTLY NOT SUPPORTED\n # _imatmul CURRENTLY NOT SUPPORTED\n # _itruediv CURRENTLY NOT SUPPORTED\n # _ifloordiv CURRENTLY NOT SUPPORTED\n # _imod CURRENTLY NOT SUPPORTED\n # _ipow CURRENTLY NOT SUPPORTED\n # _ilshift CURRENTLY NOT SUPPORTED\n # _irshift CURRENTLY NOT SUPPORTED\n # _iand CURRENTLY NOT SUPPORTED\n # _ixor CURRENTLY NOT SUPPORTED\n # _ior CURRENTLY NOT SUPPORTED\n _neg = _as_tracer_method('__neg__')\n _pos = _as_tracer_method('__pos__')\n _abs = _as_tracer_method('__abs__')\n _invert = _as_tracer_method('__invert__')\n _complex = _as_tracer_method('__complex__')\n _int = _as_tracer_method('__int__')\n _float = _as_tracer_method('__float__')\n _index = _as_tracer_method('__index__')\n _round = _as_tracer_method('__round__')\n _trunc = _as_tracer_method('__trunc__')\n _floor = _as_tracer_method('__floor__')\n _ceil = _as_tracer_method('__ceil__')\n\n # --------------------------------\n # hijax interface\n # --------------------------------\n cur_qdd = _not_an_attribute_property('cur_qdd')\n\n def __hash__(self):\n if self._leaves is not None and self._treedef is not None:\n return hash(\n (AbstractVariable, self._var_type, self._treedef, self._leaves)\n )\n else:\n assert self._leaves is None and self._treedef is None\n return hash((AbstractVariable, self._var_type))\n\n def __eq__(self, other):\n return (\n isinstance(other, AbstractVariable) and self._var_type == other._var_type\n )\n\n def str_short(self, short_dtypes=False, **_) -> str: # type: ignore\n return f'{self._var_type.__name__}()'\n\n # mutable interface\n def lo_ty_qdd(self, variable_state: VariableQDD) -> list: # type: ignore\n return [lo_ty for t in variable_state.leaf_avals for lo_ty in t.lo_ty()]\n\n def new_from_loval( # type: ignore[override]\n self, variable_state: VariableQDD, *lo_vals\n ) -> HijaxVariable:\n lo_vals_ = iter(lo_vals)\n hi_vals = [\n hi_ty.raise_val(*it.islice(lo_vals_, len(hi_ty.lo_ty()))) # type: ignore\n for hi_ty in variable_state.leaf_avals\n ]\n assert next(lo_vals_, None) is None\n variable: Variable = jax.tree.unflatten(variable_state.treedef, hi_vals)\n return HijaxVariable._new(\n hi_vals,\n variable_state.treedef,\n self._var_type,\n has_qdd=self.has_qdd,\n ref=self.ref,\n ) # will be mutated\n\n def read_loval(self, variable_state: VariableQDD, variable) -> list: # type: ignore\n leaf_vals, treedef = jax.tree.flatten(_get_hijax_state(variable))\n assert treedef == variable_state.treedef\n return [\n lo_val\n for hi_ty, hi_val in zip(variable_state.leaf_avals, leaf_vals)\n for lo_val in hi_ty.lower_val(hi_val)\n ] # type: ignore\n\n def update_from_loval( # type: ignore[override]\n self, box_state: VariableQDD, variable, *lo_vals\n ) -> None:\n lo_vals_ = iter(lo_vals)\n hi_vals = [\n hi_ty.raise_val(*it.islice(lo_vals_, len(hi_ty.lo_ty()))) # type: ignore\n for hi_ty in box_state.leaf_avals\n ]\n assert next(lo_vals_, None) is None\n _set_hijax_state(variable, jax.tree.unflatten(box_state.treedef, hi_vals))\n\n def to_tangent_aval(self):\n return AbstractVariable(\n self._var_type,\n self._treedef,\n self._leaves,\n self.has_qdd,\n ref=self.ref,\n )\n\n# --------------------------------------------\n# Variable\n# --------------------------------------------\n\n\ndef _remap_sharding_metadata(metadata: dict[str, tp.Any]) -> None:\n if 'sharding' in metadata:\n warnings.warn(\n \"'sharding' is deprecated, use 'out_sharding' instead.\",\n DeprecationWarning,\n stacklevel=3,\n )\n metadata['out_sharding'] = metadata.pop('sharding')\n if 'sharding_names' in metadata:\n warnings.warn(\n \"'sharding_names' is deprecated, use 'out_sharding' instead.\",\n DeprecationWarning,\n stacklevel=3,\n )\n metadata['out_sharding'] = metadata.pop('sharding_names')\n\n\ndef _variable_operator(name: str) -> tp.Callable[[Variable[A], tp.Any], A]:\n def variable_operator_method(self, other):\n value = self.get_value()\n if isinstance(other, Variable):\n other = other.get_value()\n return getattr(value, name)(other)\n\n variable_operator_method.__name__ = name\n return variable_operator_method\n\n\ndef _variable_unary_operator(name: str) -> tp.Callable[[Variable[A]], A]:\n def variable_unary_operator_method(self):\n value = self.get_value()\n return getattr(value, name)()\n\n variable_unary_operator_method.__name__ = name\n return variable_unary_operator_method\n\n\nclass VariableMeta(type):\n def __new__(cls, cls_name, bases, attrs):\n if '__slots__' not in attrs:\n attrs['__slots__'] = ()\n return super().__new__(cls, cls_name, bases, attrs)\n\n def __instancecheck__(self, instance):\n if super().__instancecheck__(instance):\n return True\n\n if isinstance(instance, jax.core.Tracer):\n ty = jax.typeof(instance)\n if isinstance(ty, AbstractVariable):\n return issubclass(ty._var_type, self)\n if isinstance(instance, HijaxVariable):\n return issubclass(instance._var_type, self)\n return False\n\n if not tp.TYPE_CHECKING:\n\n def __call__(cls, *args, **kwargs):\n return cls._variable_meta_call(*args, **kwargs)\n\n def _variable_meta_call(cls, *args, **kwargs):\n variable = super().__call__(*args, **kwargs)\n if variable.hijax:\n return _new_hijax_from_variable(variable)\n return variable\n\n\nclass Variable(tp.Generic[A], reprlib.Representable, metaclass=VariableMeta):\n \"\"\"The base class for all ``Variable`` types. Create custom ``Variable``\n types by subclassing this class. Numerous NNX graph functions can filter\n for specific ``Variable`` types, for example, :func:`split`, :func:`state`,\n :func:`pop`, and :func:`State.filter`.\n\n Example usage::\n\n >>> from flax import nnx\n >>> import jax, jax.numpy as jnp\n\n >>> class CustomVariable(nnx.Variable):\n ... pass\n\n >>> class Model(nnx.Module):\n ... def __init__(self, rngs):\n ... self.linear = nnx.Linear(2, 3, rngs=rngs)\n ... self.custom_variable = CustomVariable(jnp.ones((1, 3)))\n ... def __call__(self, x):\n ... return self.linear(x) + self.custom_variable\n >>> model = Model(rngs=nnx.Rngs(0))\n\n >>> linear_variables = nnx.state(model, nnx.Param)\n >>> jax.tree.map(jnp.shape, linear_variables)\n State({\n 'linear': {\n 'bias': Param(\n value=(3,)\n ),\n 'kernel': Param(\n value=(2, 3)\n )\n }\n })\n\n >>> custom_variable = nnx.state(model, CustomVariable)\n >>> jax.tree.map(jnp.shape, custom_variable)\n State({\n 'custom_variable': CustomVariable(\n value=(1, 3)\n )\n })\n\n >>> variables = nnx.state(model)\n >>> jax.tree.map(jnp.shape, variables)\n State({\n 'custom_variable': CustomVariable(\n value=(1, 3)\n ),\n 'linear': {\n 'bias': Param(\n value=(3,)\n ),\n 'kernel': Param(\n value=(2, 3)\n )\n }\n })\n \"\"\"\n\n __slots__ = ('_raw_value', '_trace_state', '_var_metadata')\n _raw_value: A\n _trace_state: tracers.TraceState\n _var_metadata: dict[str, tp.Any]\n required_metadata = frozenset(\n ['hijax', 'ref', 'eager_sharding']\n )\n\n @property\n def var_type(self):\n return type(self)\n\n @property\n def hijax(self) -> bool:\n return self._var_metadata['hijax']\n\n @property\n def ref(self) -> bool:\n return self._var_metadata['ref']\n\n @property\n def shape(self: Variable[jax.Array]) -> tuple[int, ...]:\n return self.get_value().shape\n\n @property\n def sharding_names(self):\n warnings.warn(\n \"'sharding_names' is deprecated, use 'out_sharding' instead.\",\n DeprecationWarning,\n stacklevel=2,\n )\n return self.get_metadata('out_sharding', None)\n\n def __init__(\n self,\n value: A | VariableMetadata[A],\n *,\n hijax: bool | None = None,\n ref: bool | None = None,\n eager_sharding: bool | None = None,\n **metadata: tp.Any,\n ):\n var_t = type(self)\n\n if isinstance(value, VariableMetadata):\n aux_metadata = dict(value.metadata)\n if 'hijax' in aux_metadata:\n if hijax is not None and hijax != aux_metadata['hijax']:\n raise ValueError(\n 'Cannot specify hijax both in VariableMetadata and as an '\n 'argument to Variable constructor.'\n )\n hijax = aux_metadata.pop('hijax')\n if 'ref' in aux_metadata:\n if ref is not None and ref != aux_metadata['ref']:\n raise ValueError(\n 'Cannot specify ref both in VariableMetadata and as an '\n 'argument to Variable constructor.'\n )\n ref = aux_metadata.pop('ref')\n if 'eager_sharding' in aux_metadata:\n if (\n eager_sharding is not None\n and eager_sharding != aux_metadata['eager_sharding']\n ):\n raise ValueError(\n 'Cannot specify eager_sharding both in VariableMetadata and as '\n 'an argument to Variable constructor.'\n )\n eager_sharding = aux_metadata['eager_sharding']\n metadata.update(aux_metadata)\n value = tp.cast(A, value.raw_value)\n\n if hijax is None:\n hijax = var_defaults().hijax\n\n if ref is None:\n ref = var_defaults().ref\n\n if eager_sharding is None:\n eager_sharding = using_eager_sharding()\n\n if any(is_array_ref(v) for v in jax.tree.leaves(value)):\n raise ValueError('Cannot pass a Ref directly into Variable constructor.')\n\n metadata['hijax'] = hijax\n metadata['ref'] = ref\n metadata['eager_sharding'] = eager_sharding\n object.__setattr__(self, '_trace_state', tracers.TraceState())\n object.__setattr__(self, '_var_metadata', metadata)\n object.__setattr__(self, '_raw_value', value)\n\n if hasattr(var_t, 'on_get_value') and 'on_get_value' not in metadata:\n metadata['on_get_value'] = var_t.on_get_value\n\n if hasattr(var_t, 'on_set_value') and 'on_set_value' not in metadata:\n metadata['on_set_value'] = var_t.on_set_value\n\n if hasattr(var_t, 'on_create_value') and 'on_create_value' not in metadata:\n metadata['on_create_value'] = var_t.on_create_value\n\n if hasattr(var_t, 'on_add_axis') and 'on_add_axis' not in metadata:\n metadata['on_add_axis'] = var_t.on_add_axis\n\n if hasattr(var_t, 'on_remove_axis') and 'on_remove_axis' not in metadata:\n metadata['on_remove_axis'] = var_t.on_remove_axis\n\n _remap_sharding_metadata(metadata)\n\n # run create_value hooks\n if 'on_create_value' in metadata:\n value = metadata['on_create_value'](self, value)\n\n object.__setattr__(self, '_raw_value', value)\n # run create_value hook\n value = self.create_value(value) # type: ignore\n # shard the _value if applicable\n if eager_sharding and 'out_sharding' in metadata:\n value = core_spmd.shard_value(\n value,\n metadata['out_sharding'],\n metadata.get('sharding_rules', None),\n metadata.get('mesh', None),\n )\n if ref:\n value = jax.new_ref(value) # type: ignore\n object.__setattr__(self, '_raw_value', value)\n\n @property\n def _can_update(self) -> bool:\n \"\"\"Whether the Variable can be updated in-place in the current trace context.\"\"\"\n if self.hijax:\n return True\n else:\n return self._trace_state.is_valid()\n\n def _check_can_update(self):\n if not self.hijax and not self._trace_state.is_valid():\n raise errors.TraceContextError(\n f'Cannot mutate {type(self).__name__} from a different trace level'\n )\n\n def __getattr__(self, name: str) -> tp.Any:\n if name in object.__getattribute__(self, '_var_metadata'):\n return self._var_metadata[name]\n return getattr(object.__getattribute__(self, '_raw_value'), name)\n\n def __setattr__(self, name: str, value: tp.Any):\n self._check_can_update()\n try:\n object.__setattr__(self, name, value)\n except AttributeError as e:\n raise AttributeError(\n f'Cannot set attribute {name}. '\n f'To set Variable metadata use either:\\n\\n'\n f' variable.set_metadata({name}=value)\\n\\nor\\n\\n'\n f\" variable.set_metadata('{name}', value)\"\n ) from e\n\n def __delattr__(self, name: str):\n self._check_can_update()\n try:\n object.__delattr__(self, name)\n except AttributeError as e:\n raise AttributeError(\n f'Cannot delete attribute {name}. '\n f'To delete Variable metadata use:\\n\\n'\n f\" variable.del_metadata('{name}')\"\n ) from e\n\n # NOTE(cgarciae): adding this for backward compatibility with VariableState\n @property\n def type(self):\n \"\"\"The type of the variable.\"\"\"\n\n return type(self)\n\n @tp.overload\n def get_metadata(\n self, *, exclude_required: bool = False\n ) -> dict[str, tp.Any]: ...\n\n @tp.overload\n def get_metadata(self, name: str, default: tp.Any = MISSING) -> tp.Any: ...\n\n def get_metadata(\n self,\n name: str | None = None,\n default: tp.Any = MISSING,\n *,\n exclude_required: bool | None = None,\n ) -> tp.Any:\n \"\"\"Get metadata for the Variable.\n\n Args:\n name: The key of the metadata element to get. If not provided, returns\n the full metadata dictionary.\n default: The default value to return if the metadata key is not found. If\n not provided and the key is not found, raises a KeyError.\n \"\"\"\n if name is not None and exclude_required is not None:\n raise TypeError(\n \"Cannot specify both 'name' and 'exclude_required' arguments.\"\n )\n metadata = self._var_metadata.copy()\n if name is None:\n if not isinstance(default, Missing):\n raise TypeError(\n \"Cannot provide a default value when 'name' is not provided. \"\n f'Got default={default}'\n )\n if exclude_required:\n for key in self.required_metadata:\n metadata.pop(key, None)\n return metadata\n if name not in metadata and not isinstance(default, Missing):\n return default\n return metadata[name]\n\n @tp.overload\n def set_metadata(self, metadata: dict[str, tp.Any], /) -> None: ...\n\n @tp.overload\n def set_metadata(self, name: str, value: tp.Any, /) -> None: ...\n\n @tp.overload\n def set_metadata(self, **metadata: tp.Any) -> None: ...\n\n def set_metadata(self, *args, **kwargs) -> None:\n \"\"\"Set metadata for the Variable.\n\n `set_metadata` can be called in 3 ways:\n\n 1. By passing a dictionary of metadata as the first argument, this will replace\n the entire Variable's metadata.\n 2. By passing a name and value as the first two arguments, this will set\n the metadata entry for the given name to the given value.\n 3. By using keyword arguments, this will update the Variable's metadata\n with the provided key-value pairs.\n \"\"\"\n self._check_can_update()\n if args and kwargs:\n raise TypeError(\n 'Cannot mix positional and keyword arguments in set_metadata'\n )\n if len(args) == 1:\n metadata = dict(args[0])\n _remap_sharding_metadata(metadata)\n if 'hijax' not in metadata:\n metadata['hijax'] = self.hijax\n if metadata['hijax'] != self.hijax:\n raise ValueError(\n f'Cannot change `hijax` metadata, expected {self.hijax}, '\n f'got {metadata[\"hijax\"]}'\n )\n if 'ref' not in metadata:\n metadata['ref'] = self.ref\n if metadata['ref'] != self.ref:\n raise ValueError(\n f'Cannot change `ref` metadata, expected {self.ref}, '\n f'got {metadata[\"ref\"]}'\n )\n if 'eager_sharding' not in metadata:\n metadata['eager_sharding'] = self.eager_sharding\n if metadata['eager_sharding'] != self.eager_sharding:\n raise ValueError(\n f'Cannot change `eager_sharding` metadata, expected '\n f'{self.eager_sharding}, got {metadata[\"eager_sharding\"]}'\n )\n self._var_metadata = metadata\n elif len(args) == 2:\n name, value = args\n if name == 'sharding_names':\n warnings.warn(\n \"'sharding_names' is deprecated, use 'out_sharding' instead.\",\n DeprecationWarning,\n stacklevel=2,\n )\n name = 'out_sharding'\n elif name == 'sharding':\n warnings.warn(\n \"'sharding' is deprecated, use 'out_sharding' instead.\",\n DeprecationWarning,\n stacklevel=2,\n )\n name = 'out_sharding'\n if name == 'hijax' and value != self.hijax:\n raise ValueError(\n f'Cannot change `hijax` metadata, expected {self.hijax}, got {value}'\n )\n if name == 'ref' and value != self.ref:\n raise ValueError(\n f'Cannot change `ref` metadata, expected {self.ref}, got {value}'\n )\n self._var_metadata[name] = value\n elif kwargs:\n _remap_sharding_metadata(kwargs)\n if 'hijax' in kwargs and kwargs['hijax'] != self.hijax:\n raise ValueError(\n f'Cannot change `hijax` metadata, expected {self.hijax}, '\n f'got {kwargs[\"hijax\"]}'\n )\n if 'ref' in kwargs and kwargs['ref'] != self.ref:\n raise ValueError(\n f'Cannot change `ref` metadata, expected {self.ref}, '\n f'got {kwargs[\"ref\"]}'\n )\n self._var_metadata.update(kwargs)\n else:\n raise TypeError(\n f'set_metadata takes either 1 or 2 arguments, or at least 1 keyword argument, '\n f'got args={args}, kwargs={kwargs}'\n )\n\n def has_metadata(self, name: str) -> bool:\n \"\"\"Check if the Variable has a metadata entry for the given name.\n\n Args:\n name: The key of the metadata element to check.\n Returns:\n True if the metadata entry exists, False otherwise.\n \"\"\"\n return name in self._var_metadata\n\n def del_metadata(self, name: str) -> None:\n \"\"\"Delete a metadata entry for the Variable.\n\n Args:\n name: The key of the metadata element to delete.\n \"\"\"\n self._check_can_update()\n if name in ('hijax', 'ref'):\n raise ValueError(f'Cannot delete `{name}` metadata')\n del self._var_metadata[name]\n\n def copy_from(self, other: Variable[A]) -> None:\n if type(self) is not type(other):\n raise ValueError(\n f'Cannot copy from incompatible container, '\n f'expected {type(self).__name__}, got {type(other).__name__}'\n )\n if self is other:\n return\n self._raw_value = other._raw_value\n self._var_metadata.clear()\n self._var_metadata.update(other.get_metadata())\n\n def update_from_state(self, variable_state: Variable[A]):\n self._raw_value = variable_state._raw_value\n\n if self._var_metadata != variable_state._var_metadata:\n metadata = variable_state.get_metadata()\n metadata['hijax'] = self.hijax\n metadata['ref'] = self.ref\n self._var_metadata = metadata\n\n @tp.final\n def get_raw_value(self) -> A:\n return self._raw_value\n\n # @tp.final\n def set_raw_value(self, value: A, *, _unsafe_bypass_check: bool = False):\n if not _unsafe_bypass_check:\n self._check_can_update()\n self._raw_value = value\n\n @property\n def raw_value(self) -> A:\n warnings.warn(\n \"'.raw_value' access is now deprecated. Use:\\n\\n\"\n ' variable.get_raw_value()\\n',\n DeprecationWarning,\n stacklevel=2,\n )\n return self.get_raw_value()\n\n @raw_value.setter\n def raw_value(self, value: A):\n warnings.warn(\n \"'.raw_value' setter is now deprecated. Use:\\n\\n\"\n ' variable.set_raw_value(value)\\n',\n DeprecationWarning,\n stacklevel=2,\n )\n self.set_raw_value(value)\n\n @property\n def value(self) -> A:\n warnings.warn(\n \"'.value' access is now deprecated. For Variable[Array] instances use:\\n\\n\"\n ' variable[...]\\n\\n'\n 'For other Variable types use:\\n\\n'\n ' variable.get_value()\\n',\n DeprecationWarning,\n stacklevel=2,\n )\n return self.get_value()\n\n @value.setter\n def value(self, value: A):\n warnings.warn(\n \"'.value' setter is now deprecated. For Variable[Array] instances use:\\n\\n\"\n ' variable[...] = value\\n\\n'\n 'For other Variable types use:\\n\\n'\n ' variable.set_value(value)\\n',\n DeprecationWarning,\n stacklevel=2,\n )\n self.set_value(value)\n\n def create_value(self, value: A):\n return value\n\n def get_value(self, *, index: tp.Any = MISSING) -> A:\n value = jax.tree.map(lambda x: x, self._raw_value) # make a copy\n if not isinstance(index, Missing):\n if is_array_ref(value):\n value = value[index]\n elif isinstance(value, jax.Array) and index is ...:\n pass # skip trivial access\n else:\n value = value[index]\n elif is_array_ref(value):\n value = value[...]\n if 'on_get_value' in self._var_metadata:\n value = self._var_metadata['on_get_value'](self, value)\n return value # type: ignore\n\n def set_value(self, value: A, *, index: tp.Any = MISSING):\n value = jax.tree.map(lambda x: x, value) # make a copy\n if 'on_set_value' in self._var_metadata:\n value = self._var_metadata['on_set_value'](self, value)\n # update _raw_value\n if is_array_ref(self._raw_value):\n if isinstance(index, Missing):\n self._raw_value[...] = value\n else:\n self._raw_value[index] = value\n elif isinstance(self._raw_value, jax.Array) and (\n not isinstance(index, Missing)\n ):\n # check if its a full replace to av\n if (\n index == ...\n and isinstance(value, jax.Array)\n and value.shape == self._raw_value[index].shape\n and value.dtype == self._raw_value.dtype\n and (\n getattr(value, 'sharding', None)\n == getattr(self._raw_value, 'sharding', None)\n )\n ):\n self._raw_value = value\n else:\n self._raw_value = self._raw_value.at[index].set(value) # type: ignore\n else:\n if isinstance(index, Missing):\n self._raw_value = value\n else:\n self._raw_value[index] = value # type: ignore\n\n def add_axis(self, axis_index: AxisIndex, axis_name: AxisName | None):\n if 'on_add_axis' in self._var_metadata:\n self._var_metadata['on_add_axis'](self, axis_index, axis_name)\n\n def remove_axis(self, axis_index: AxisIndex, axis_name: AxisName | None):\n if 'on_remove_axis' in self._var_metadata:\n self._var_metadata['on_remove_axis'](self, axis_index, axis_name)\n\n @tp.overload\n def copy(self, value: B, **kwargs) -> Variable[B]: ...\n\n @tp.overload\n def copy(self, **kwargs) -> Variable[A]: ...\n\n def copy(\n self,\n value: tp.Any = MISSING,\n *,\n _copy_ref: bool = True,\n **updates,\n ) -> Variable[tp.Any]:\n assert 'raw_value' not in updates\n\n new_metadata = self.get_metadata() | updates\n\n if not isinstance(value, Missing):\n pass\n elif 'value' in updates:\n value = updates.pop('value')\n else:\n value = self.get_raw_value()\n if _copy_ref and is_array_ref(value):\n value = value[...]\n\n if _copy_ref and new_metadata['ref']:\n value = jax.new_ref(value)\n new_metadata['ref'] = True\n\n value = jax.tree.map(lambda x: x, value) # make a copy\n obj = self.from_metadata(value, new_metadata)\n return obj\n\n @classmethod\n def _new(\n cls,\n value: A,\n metadata: dict[str, tp.Any],\n ) -> Variable[A]:\n obj = object.__new__(cls)\n # skip __setattr__ for trace_state initialization\n object.__setattr__(obj, '_trace_state', tracers.TraceState())\n object.__setattr__(obj, '_var_metadata', metadata)\n object.__setattr__(obj, '_raw_value', value)\n return obj\n\n @classmethod\n def from_metadata(\n cls,\n value: A,\n attributes: dict[str, tp.Any],\n ) -> Variable[A]:\n variable = cls._new(value, dict(attributes))\n if attributes['hijax']:\n variable = _new_hijax_from_variable(variable) # type: ignore[assignment]\n return variable # type: ignore[return-value]\n\n replace = copy\n to_state = copy\n\n def __nnx_repr__(self):\n stats = SizeBytes.from_any(self._raw_value)\n if stats:\n comment = f' # {stats}'\n else:\n comment = ''\n\n yield reprlib.Object(type=type(self).__name__, comment=comment)\n yield reprlib.Attr('value', self.get_value())\n for name, value in self._var_metadata.items():\n if name == 'hijax' and value == config.flax_hijax_variable:\n continue\n if name == 'ref' and not value:\n continue\n if name == 'eager_sharding' and value == config.flax_always_shard_variable:\n continue\n yield reprlib.Attr(name, value)\n\n def __treescope_repr__(self, path, subtree_renderer):\n size_bytes = SizeBytes.from_any(self.get_value())\n if size_bytes:\n stats_repr = f' # {size_bytes}'\n first_line_annotation = treescope.rendering_parts.comment_color(\n treescope.rendering_parts.text(f'{stats_repr}')\n )\n else:\n first_line_annotation = None\n\n metadata = {\n name: value\n for name, value in self._var_metadata.items()\n if not (name == 'hijax' and value == config.flax_hijax_variable)\n and not (name == 'ref' and not value)\n and not (name == 'eager_sharding' and value == config.flax_always_shard_variable)\n }\n children = {'value': self.get_value(), **metadata}\n return visualization.render_object_constructor(\n object_type=type(self),\n attributes=children,\n path=path,\n subtree_renderer=subtree_renderer,\n first_line_annotation=first_line_annotation,\n )\n\n # hooks API\n if tp.TYPE_CHECKING:\n\n def on_get_value(self, value: A) -> A: ...\n\n def on_set_value(self, value: A) -> A: ...\n\n def on_create_value(self, value: A) -> A: ...\n\n def on_add_axis(\n self: V, axis_index: AxisIndex, axis_name: AxisName | None\n ) -> V: ...\n\n def on_remove_axis(\n self: V, axis_index: AxisIndex, axis_name: AxisName | None\n ) -> V: ...\n\n def __jax_array__(self):\n return self.get_value()\n\n # pickle support\n def __getstate__(self):\n return {\n '_raw_value': self._raw_value,\n '_trace_state': self._trace_state,\n '_var_metadata': self._var_metadata,\n }\n\n def __setstate__(self, state):\n # skip __setattr__ for trace_state initialization\n object.__setattr__(self, '_trace_state', state['_trace_state'])\n object.__setattr__(self, '_var_metadata', state['_var_metadata'])\n object.__setattr__(self, '_raw_value', state['_raw_value'])\n\n # --------------------------------------------\n # proxy methods\n # --------------------------------------------\n @tp.overload\n def __getitem__(self: Variable[jax.Array], key) -> jax.Array:\n ...\n\n @tp.overload\n def __getitem__(self: Variable[dict[tp.Any, B]], key) -> B:\n ...\n\n @tp.overload\n def __getitem__(self: Variable[list[B]], key: int) -> B:\n ...\n\n @tp.overload\n def __getitem__(self: Variable[tuple[B, ...]], key: int) -> B:\n ...\n\n @tp.overload\n def __getitem__(self, key) -> tp.Any:\n ...\n\n def __getitem__(self, key):\n return self.get_value(index=key)\n\n def __setitem__(self, key, value) -> None:\n self.set_value(value, index=key)\n\n def __delitem__(self, key) -> None:\n value = self.get_value()\n del value[key] # type: ignore\n self.set_value(value) # type: ignore\n\n def __call__(self, *args, **kwargs) -> tp.Any:\n return self.get_value()(*args, **kwargs) # type: ignore\n\n def __len__(self) -> int:\n return len(self.get_value()) # type: ignore\n\n def __iter__(self) -> tp.Iterator:\n return iter(self.get_value()) # type: ignore\n\n def __contains__(self, item) -> bool:\n return item in self.get_value() # type: ignore\n\n __add__ = _variable_operator('__add__')\n __sub__ = _variable_operator('__sub__')\n __mul__ = _variable_operator('__mul__')\n __matmul__ = _variable_operator('__matmul__')\n __truediv__ = _variable_operator('__truediv__')\n __floordiv__ = _variable_operator('__floordiv__')\n __mod__ = _variable_operator('__mod__')\n __pow__ = _variable_operator('__pow__')\n __lshift__ = _variable_operator('__lshift__')\n __rshift__ = _variable_operator('__rshift__')\n __and__ = _variable_operator('__and__')\n __xor__ = _variable_operator('__xor__')\n __or__ = _variable_operator('__or__')\n __radd__ = _variable_operator('__radd__')\n __rsub__ = _variable_operator('__rsub__')\n __rmul__ = _variable_operator('__rmul__')\n __rmatmul__ = _variable_operator('__rmatmul__')\n __rtruediv__ = _variable_operator('__rtruediv__')\n __rfloordiv__ = _variable_operator('__rfloordiv__')\n __rmod__ = _variable_operator('__rmod__')\n __rpow__ = _variable_operator('__rpow__')\n __rlshift__ = _variable_operator('__rlshift__')\n __rrshift__ = _variable_operator('__rrshift__')\n __rand__ = _variable_operator('__rand__')\n __rxor__ = _variable_operator('__rxor__')\n __ror__ = _variable_operator('__ror__')\n\n def __eq__(self, other) -> bool:\n if isinstance(other, Variable):\n other = other.get_value()\n return self.get_value().__eq__(other) # type: ignore\n\n def __iadd__(self: V, other) -> V:\n raise NotImplementedError(\n 'In-place operations are no longer supported for Variable.\\n'\n 'Use `variable[...] += x` instead.'\n )\n\n def __isub__(self: V, other) -> V:\n raise NotImplementedError(\n 'In-place operations are no longer supported for Variable.\\n'\n 'Use `variable[...] -= x` instead.'\n )\n\n def __imul__(self: V, other) -> V:\n raise NotImplementedError(\n 'In-place operations are no longer supported for Variable.\\n'\n 'Use `variable[...] *= x` instead.'\n )\n\n def __imatmul__(self: V, other) -> V:\n raise NotImplementedError(\n 'In-place operations are no longer supported for Variable.\\n'\n 'Use `variable.value @= x` instead.'\n )\n\n def __itruediv__(self: V, other) -> V:\n raise NotImplementedError(\n 'In-place operations are no longer supported for Variable.\\n'\n 'Use `variable.value /= x` instead.'\n )\n\n def __ifloordiv__(self: V, other) -> V:\n raise NotImplementedError(\n 'In-place operations are no longer supported for Variable.\\n'\n 'Use `variable.value //= x`` instead.'\n )\n\n def __imod__(self: V, other) -> V:\n raise NotImplementedError(\n 'In-place operations are no longer supported for Variable.\\n'\n 'Use `variable.value %= x` instead.'\n )\n\n def __ipow__(self: V, other) -> V:\n raise NotImplementedError(\n 'In-place operations are no longer supported for Variable.\\n'\n 'Use `variable.value **= x`` instead.'\n )\n\n def __ilshift__(self: V, other) -> V:\n raise NotImplementedError(\n 'In-place operations are no longer supported for Variable.\\n'\n 'Use `variable.value <<= x`` instead.'\n )\n\n def __irshift__(self: V, other) -> V:\n raise NotImplementedError(\n 'In-place operations are no longer supported for Variable.\\n'\n 'Use `variable.value >>= x`` instead.'\n )\n\n def __iand__(self: V, other) -> V:\n raise NotImplementedError(\n 'In-place operations are no longer supported for Variable.\\n'\n 'Use `variable.value &= x` instead.'\n )\n\n def __ixor__(self: V, other) -> V:\n raise NotImplementedError(\n 'In-place operations are no longer supported for Variable.\\n'\n 'Use `variable.value ^= x` instead.'\n )\n\n def __ior__(self: V, other) -> V:\n raise NotImplementedError(\n 'In-place operations are no longer supported for Variable.\\n'\n 'Use `variable.value |= x` instead.'\n )\n\n __neg__ = _variable_unary_operator('__neg__')\n __pos__ = _variable_unary_operator('__pos__')\n __abs__ = _variable_unary_operator('__abs__')\n __invert__ = _variable_unary_operator('__invert__')\n __complex__ = _variable_unary_operator('__complex__')\n __int__ = _variable_unary_operator('__int__')\n __float__ = _variable_unary_operator('__float__')\n __index__ = _variable_unary_operator('__index__')\n __trunc__ = _variable_unary_operator('__trunc__')\n __floor__ = _variable_unary_operator('__floor__')\n __ceil__ = _variable_unary_operator('__ceil__')\n\n def __round__(self, ndigits: int = 0) -> A:\n return self.get_value().__round__(ndigits) # type: ignore\n\n # --------------------------------------------\n\n def __init_subclass__(cls) -> None:\n if '__slots__' not in vars(cls):\n cls.__slots__ = () # type: ignore[assignment]\n super().__init_subclass__()\n jax.tree_util.register_pytree_with_keys(\n cls,\n flatten_with_keys=_variable_flatten_with_keys,\n unflatten_func=partial(_variable_unflatten, cls), # type: ignore\n flatten_func=_variable_flatten,\n )\n\n\ndef _variable_flatten_with_keys(x: Variable[tp.Any]):\n metadata = tuple(sorted(x._var_metadata.items()))\n node = (jtu.GetAttrKey('value'), x._raw_value)\n return (node,), metadata\n\n\ndef _variable_flatten(x: Variable[tp.Any]):\n metadata = tuple(sorted(x._var_metadata.items()))\n return (x._raw_value,), metadata\n\n\ndef _variable_unflatten(\n cls: type[Variable[tp.Any]],\n static: tuple[tuple[str, tp.Any], ...],\n children: tuple[tp.Any],\n):\n return cls._new(children[0], dict(static))\n\n\njax.tree_util.register_pytree_with_keys(\n Variable,\n flatten_with_keys=_variable_flatten_with_keys,\n unflatten_func=partial(_variable_unflatten, Variable), # type: ignore\n flatten_func=_variable_flatten,\n)\n\nVariableState = Variable\n\n\nclass Param(Variable[A]):\n \"\"\"The canonical learnable parameter. All learnable parameters\n in NNX layer modules will have the ``Param`` :class:`Variable`\n type::\n\n >>> from flax import nnx\n >>> import jax, jax.numpy as jnp\n\n >>> layer = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n >>> jax.tree.map(jnp.shape, nnx.state(layer))\n State({\n 'bias': Param(\n value=(3,)\n ),\n 'kernel': Param(\n value=(2, 3)\n )\n })\n \"\"\"\n\n pass\n\n\nclass BatchStat(Variable[A]):\n \"\"\"The mean and variance batch statistics stored in\n the :class:`BatchNorm` layer. Note, these are not the\n learnable scale and bias parameters, but rather the\n running average statistics that are typically used\n during post-training inference::\n\n >>> from flax import nnx\n >>> import jax, jax.numpy as jnp\n\n >>> layer = nnx.BatchNorm(3, rngs=nnx.Rngs(0))\n >>> jax.tree.map(jnp.shape, nnx.state(layer))\n State({\n 'bias': Param(\n value=(3,)\n ),\n 'mean': BatchStat(\n value=(3,)\n ),\n 'scale': Param(\n value=(3,)\n ),\n 'var': BatchStat(\n value=(3,)\n )\n })\n \"\"\"\n\n pass\n\n\nclass Cache(Variable[A]):\n \"\"\"Autoregressive cache in :class:`MultiHeadAttention`::\n\n >>> from flax import nnx\n >>> import jax, jax.numpy as jnp\n\n >>> layer = nnx.MultiHeadAttention(\n ... num_heads=2,\n ... in_features=3,\n ... qkv_features=6,\n ... out_features=6,\n ... decode=True,\n ... rngs=nnx.Rngs(0),\n ... )\n\n >>> layer.init_cache((1, 3))\n >>> jax.tree.map(jnp.shape, nnx.state(layer, nnx.Cache))\n State({\n 'cache_index': Cache(\n value=()\n ),\n 'cached_key': Cache(\n value=(1, 2, 3)\n ),\n 'cached_value': Cache(\n value=(1, 2, 3)\n )\n })\n \"\"\"\n\n pass\n\n\nclass Intermediate(Variable[A]):\n \"\"\":class:`Variable` type that is typically used for\n :func:`Module.sow`. Use :func:`nnx.capture` to retrieve\n the sowed values::\n\n >>> from flax import nnx\n >>> import jax, jax.numpy as jnp\n\n >>> class Model(nnx.Module):\n ... def __init__(self, rngs):\n ... self.linear1 = nnx.Linear(2, 3, rngs=rngs)\n ... self.linear2 = nnx.Linear(3, 4, rngs=rngs)\n ... def __call__(self, x):\n ... x = self.linear1(x)\n ... self.sow(nnx.Intermediate, 'i', x)\n ... x = self.linear2(x)\n ... return x\n >>> model = Model(rngs=nnx.Rngs(0))\n\n >>> x = jnp.ones((1, 2))\n >>> y, intms = nnx.capture(model, nnx.Intermediate)(x)\n >>> jax.tree.map(jnp.shape, intms)\n State({\n 'i': Intermediate(\n value=((1, 3),)\n )\n })\n \"\"\"\n\n pass\n\n\nclass Perturbation(Intermediate[A]):\n \"\"\":class:`Variable` type that is typically used for\n :func:`Module.perturb`. Use :func:`nnx.capture` to retrieve\n the perturbation values::\n\n >>> from flax import nnx\n >>> import jax, jax.numpy as jnp\n\n >>> class Model(nnx.Module):\n ... def __init__(self, rngs):\n ... self.linear1 = nnx.Linear(2, 3, rngs=rngs)\n ... self.linear2 = nnx.Linear(3, 4, rngs=rngs)\n ... def __call__(self, x):\n ... x = self.linear1(x)\n ... x = self.perturb('i', x)\n ... x = self.linear2(x)\n ... return x\n >>> model = Model(rngs=nnx.Rngs(0))\n\n >>> x = jnp.ones((1, 2))\n >>> y, perturbations = nnx.capture(model, nnx.Perturbation)(x)\n >>> jax.tree.map(jnp.shape, perturbations)\n State({\n 'i': Perturbation(\n value=(1, 3)\n )\n })\n \"\"\"\n\n pass\n\n\ndef with_metadata(\n initializer: F,\n set_value_hooks: tp.Union[SetValueHook[A], tp.Sequence[SetValueHook[A]]] = (),\n get_value_hooks: tp.Union[SetValueHook[A], tp.Sequence[SetValueHook[A]]] = (),\n create_value_hooks: tp.Union[\n CreateValueHook[A], tp.Sequence[CreateValueHook[A]]\n ] = (),\n add_axis_hooks: tp.Union[\n AddAxisHook[Variable[A]], tp.Sequence[AddAxisHook[Variable[A]]]\n ] = (),\n remove_axis_hooks: tp.Union[\n RemoveAxisHook[Variable[A]],\n tp.Sequence[RemoveAxisHook[Variable[A]]],\n ] = (),\n **metadata: tp.Any,\n) -> F:\n if set_value_hooks:\n if callable(set_value_hooks):\n set_value_hooks = (set_value_hooks,)\n else:\n set_value_hooks = tuple(set_value_hooks)\n else:\n set_value_hooks = ()\n\n if get_value_hooks:\n if callable(get_value_hooks):\n get_value_hooks = (get_value_hooks,)\n else:\n get_value_hooks = tuple(get_value_hooks)\n else:\n get_value_hooks = ()\n\n if create_value_hooks:\n if callable(create_value_hooks):\n create_value_hooks = (create_value_hooks,)\n else:\n create_value_hooks = tuple(create_value_hooks)\n else:\n create_value_hooks = ()\n\n if add_axis_hooks:\n if callable(add_axis_hooks):\n add_axis_hooks = (add_axis_hooks,)\n else:\n add_axis_hooks = tuple(add_axis_hooks)\n else:\n add_axis_hooks = ()\n\n if remove_axis_hooks:\n if callable(remove_axis_hooks):\n remove_axis_hooks = (remove_axis_hooks,)\n else:\n remove_axis_hooks = tuple(remove_axis_hooks)\n else:\n remove_axis_hooks = ()\n\n @functools.wraps(initializer)\n def wrapper(*args):\n return VariableMetadata(\n initializer(*args),\n set_value_hooks=set_value_hooks,\n get_value_hooks=get_value_hooks,\n create_value_hooks=create_value_hooks,\n add_axis_hooks=add_axis_hooks,\n remove_axis_hooks=remove_axis_hooks,\n metadata=metadata,\n )\n\n return wrapper # type: ignore\n\n\n###################################################\n### Variable type/class <-> string name mapping ###\n###################################################\n# Assumption: the mapping is 1-1 and unique.\n\nVariableTypeCache: dict[str, tp.Type[Variable[tp.Any]]] = {}\n\n\ndef variable_type_from_name(\n name: str,\n /,\n *,\n base: type[Variable[tp.Any]] = Variable,\n allow_register: bool = False,\n) -> tp.Type[Variable[tp.Any]]:\n \"\"\"Given a Linen-style collection name, get or create its NNX Variable class.\"\"\"\n if name not in VariableTypeCache:\n if not allow_register:\n raise ValueError(\n f'Name {name} is not registered in the registry. '\n 'To register a new name, use register_variable_name() '\n 'or set allow_register=True.'\n )\n VariableTypeCache[name] = type(name, (base,), {})\n return VariableTypeCache[name]\n\n\ndef variable_name_from_type(\n typ: tp.Type[Variable[tp.Any]], /, *, allow_register: bool = False\n) -> str:\n \"\"\"Given an NNX Variable type, get its Linen-style collection name.\n\n Should output the exact inversed result of `variable_type_from_name()`.\"\"\"\n for name, t in VariableTypeCache.items():\n if typ == t:\n return name\n\n if not allow_register:\n raise ValueError(\n f'Type {typ} is not registered in the registry. '\n 'To register a new type, use register_variable_name() '\n 'or set allow_register=True.'\n )\n name = typ.__name__\n if name in VariableTypeCache:\n raise ValueError(\n 'Name {name} is already registered in the registry as {VariableTypeCache[name]}. '\n 'It cannot be linked with this type {typ}.'\n )\n register_variable_name(name, typ)\n return name\n\n\n@tp.overload\ndef register_variable_name(\n name: str,\n typ: type[Variable[tp.Any]],\n *,\n overwrite: bool = False,\n) -> type[Variable[tp.Any]]: ...\n\n\n@tp.overload\ndef register_variable_name(\n name: str,\n *,\n overwrite: bool = False,\n) -> tp.Callable[[type[Variable[tp.Any]]], type[Variable[tp.Any]]]: ...\n\n\ndef register_variable_name(\n name: str,\n typ: type[Variable[A]] | Missing = MISSING,\n *,\n overwrite=False,\n) -> type[Variable[A]] | tp.Callable[[type[Variable[A]]], type[Variable[A]]]:\n \"\"\"Register a pair of Linen collection name and its NNX type.\"\"\"\n if isinstance(typ, Missing):\n return partial(register_variable_name, name, overwrite=overwrite)\n typ = tp.cast(type[Variable[A]], typ)\n if not overwrite and name in VariableTypeCache:\n raise ValueError(\n f'Name {name} already mapped to type {VariableTypeCache[name]}. '\n 'To overwrite, call register_variable_name() with `overwrite=True`.'\n )\n VariableTypeCache[name] = typ\n return typ\n\n\n# add known variable type names\nregister_variable_name('params', Param)\nregister_variable_name('batch_stats', BatchStat)\nregister_variable_name('cache', Cache)\nregister_variable_name('intermediates', Intermediate)\nregister_variable_name('perturbations', Perturbation)\n" }, { "path": "flax/nnx/visualization.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport typing as tp\n\nimport treescope # type: ignore[import-untyped]\nfrom treescope import rendering_parts, renderers\n\ntry:\n from IPython import get_ipython\n\n in_ipython = get_ipython() is not None\nexcept ImportError:\n in_ipython = False\n\n\ndef display(*args):\n \"\"\"Display the given objects using the Treescope pretty-printer.\n\n If treescope is not installed or the code is not running in IPython,\n ``display`` will print the objects instead.\n \"\"\"\n if not in_ipython:\n for x in args:\n print(x)\n return\n\n for x in args:\n treescope.display(x, ignore_exceptions=True, autovisualize=True)\n\n\ndef render_object_constructor(\n object_type: type[tp.Any],\n attributes: tp.Mapping[str, tp.Any],\n path: str | None,\n subtree_renderer: renderers.TreescopeSubtreeRenderer,\n roundtrippable: bool = False,\n color: str | None = None,\n first_line_annotation: rendering_parts.RenderableTreePart | None = None,\n) -> rendering_parts.Rendering:\n \"\"\"Renders an object in \"constructor format\", similar to a dataclass.\n\n This produces a rendering like `Foo(bar=1, baz=2)`, where Foo identifies the\n type of the object, and bar and baz are the names of the attributes of the\n object. It is a *requirement* that these are the actual attributes of the\n object, which can be accessed via `obj.bar` or similar; otherwise, the\n path renderings will break.\n\n This can be used from within a `__treescope_repr__` implementation via ::\n\n def __treescope_repr__(self, path, subtree_renderer):\n return repr_lib.render_object_constructor(\n object_type=type(self),\n attributes=,\n path=path,\n subtree_renderer=subtree_renderer,\n )\n\n Args:\n object_type: The type of the object.\n attributes: The attributes of the object, which will be rendered as keyword\n arguments to the constructor.\n path: The path to the object. When `render_object_constructor` is called\n from `__treescope_repr__`, this should come from the `path` argument to\n `__treescope_repr__`.\n subtree_renderer: The renderer to use to render subtrees. When\n `render_object_constructor` is called from `__treescope_repr__`, this\n should come from the `subtree_renderer` argument to `__treescope_repr__`.\n roundtrippable: Whether evaluating the rendering as Python code will produce\n an object that is equal to the original object. This implies that the\n keyword arguments are actually the keyword arguments to the constructor,\n and not some other attributes of the object.\n color: The background color to use for the object rendering. If None, does\n not use a background color. A utility for assigning a random color based\n on a string key is given in `treescope.formatting_util`.\n first_line_annotation: An annotation for the first line of the node when it\n is expanded.\n\n Returns:\n A rendering of the object, suitable for returning from `__treescope_repr__`.\n \"\"\"\n if roundtrippable:\n constructor = rendering_parts.siblings(\n rendering_parts.maybe_qualified_type_name(object_type), '('\n )\n closing_suffix = rendering_parts.text(')')\n else:\n constructor = rendering_parts.siblings(\n rendering_parts.roundtrip_condition(roundtrip=rendering_parts.text('<')),\n rendering_parts.maybe_qualified_type_name(object_type),\n '(',\n )\n closing_suffix = rendering_parts.siblings(\n ')',\n rendering_parts.roundtrip_condition(roundtrip=rendering_parts.text('>')),\n )\n\n children = []\n for i, (name, value) in enumerate(attributes.items()):\n child_path = None if path is None else f'{path}.{name}'\n\n if i < len(attributes) - 1:\n # Not the last child. Always show a comma, and add a space when\n # collapsed.\n comma_after = rendering_parts.siblings(\n ',',\n rendering_parts.fold_condition(collapsed=rendering_parts.text(' ')),\n )\n else:\n # Last child: only show the comma when the node is expanded.\n comma_after = rendering_parts.fold_condition(\n expanded=rendering_parts.text(',')\n )\n\n child_line = rendering_parts.build_full_line_with_annotations(\n rendering_parts.siblings_with_annotations(\n f'{name}=',\n subtree_renderer(value, path=child_path),\n ),\n comma_after,\n )\n children.append(child_line)\n\n return rendering_parts.build_foldable_tree_node_from_children(\n prefix=constructor,\n children=children,\n suffix=closing_suffix,\n path=path,\n background_color=color,\n first_line_annotation=first_line_annotation,\n )" }, { "path": "flax/oss/ .git-blame-ignore-revs", "content": "# .git-blame-ignore-revs\n#\n# These commits will be ignored by the github blame view.\n# The git blame CLI can ignore them as well by doing:\n# git blame --ignore-revs-file .git-blame-ignore-revs \n# or via global config:\n# git config --global blame.ignoreRevsFile .git-blame-ignore-revs\n# see the blame.markIgnoredLines and blame.markUnblamableLines options.\n#\n# remove all trailing whitespaces\n442df07ca1a90f04c685cdae9f8e488bbffc2f83\n" }, { "path": "flax/py.typed", "content": "# Marker file for PEP 561. The package uses inline types.\n" }, { "path": "flax/serialization.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Serialization utilities for Jax.\n\nAll Flax classes that carry state (e.g., Optimizer) can be turned into a\nstate dict of numpy arrays for easy serialization.\n\"\"\"\nimport enum\nimport threading\nfrom contextlib import contextmanager\nfrom typing import Any\n\nimport jax\nimport msgpack\nimport numpy as np\n\n_STATE_DICT_REGISTRY: dict[Any, Any] = {}\n\n\nclass _ErrorContext(threading.local):\n \"\"\"Context for deserialization error messages.\"\"\"\n\n def __init__(self):\n self.path = []\n\n\n_error_context = _ErrorContext()\n\n\n@contextmanager\ndef _record_path(name):\n try:\n _error_context.path.append(name)\n yield\n finally:\n _error_context.path.pop()\n\n\ndef current_path():\n \"\"\"Current state_dict path during deserialization for error messages.\"\"\"\n return '/'.join(_error_context.path)\n\n\nclass _NamedTuple:\n \"\"\"Fake type marker for namedtuple for registry.\"\"\"\n\n pass\n\n\ndef _is_namedtuple(x):\n \"\"\"Duck typing test for namedtuple factory-generated objects.\"\"\"\n return isinstance(x, tuple) and hasattr(x, '_fields')\n\n\ndef from_state_dict(target, state: dict[str, Any], name: str = '.'):\n \"\"\"Restores the state of the given target using a state dict.\n\n This function takes the current target as an argument. This\n lets us know the exact structure of the target,\n as well as lets us add assertions that shapes and dtypes don't change.\n\n In practice, none of the leaf values in ``target`` are actually\n used. Only the tree structure, shapes and dtypes.\n\n Args:\n target: the object of which the state should be restored.\n state: a dictionary generated by ``to_state_dict`` with the desired new\n state for ``target``.\n name: name of branch taken, used to improve deserialization error messages.\n Returns:\n A copy of the object with the restored state.\n \"\"\"\n if _is_namedtuple(target):\n ty = _NamedTuple\n else:\n ty = type(target)\n if ty not in _STATE_DICT_REGISTRY:\n return state\n ty_from_state_dict = _STATE_DICT_REGISTRY[ty][1]\n with _record_path(name):\n return ty_from_state_dict(target, state)\n\n\ndef to_state_dict(target) -> dict[str, Any]:\n \"\"\"Returns a dictionary with the state of the given target.\"\"\"\n if _is_namedtuple(target):\n ty = _NamedTuple\n else:\n ty = type(target)\n if ty not in _STATE_DICT_REGISTRY:\n return target\n\n ty_to_state_dict = _STATE_DICT_REGISTRY[ty][0]\n state_dict = ty_to_state_dict(target)\n if isinstance(state_dict, dict):\n for key in state_dict.keys():\n assert isinstance(key, str), 'A state dict must only have string keys.'\n return state_dict\n\n\ndef is_serializable(target):\n if not isinstance(target, type):\n target = type(target)\n return target in _STATE_DICT_REGISTRY\n\n\ndef register_serialization_state(\n ty, ty_to_state_dict, ty_from_state_dict, override=False\n):\n \"\"\"Register a type for serialization.\n\n Args:\n ty: the type to be registered\n ty_to_state_dict: a function that takes an instance of ty and\n returns its state as a dictionary.\n ty_from_state_dict: a function that takes an instance of ty and\n a state dict, and returns a copy of the instance with the restored state.\n override: override a previously registered serialization handler\n (default: False).\n \"\"\"\n if ty in _STATE_DICT_REGISTRY and not override:\n raise ValueError(\n f'a serialization handler for \"{ty.__name__}\" is already registered'\n )\n _STATE_DICT_REGISTRY[ty] = (ty_to_state_dict, ty_from_state_dict)\n\n\ndef _list_state_dict(xs: list[Any]) -> dict[str, Any]:\n return {str(i): to_state_dict(x) for i, x in enumerate(xs)}\n\n\ndef _restore_list(xs, state_dict: dict[str, Any]) -> list[Any]:\n if len(state_dict) != len(xs):\n raise ValueError(\n 'The size of the list and the state dict do not match,'\n f' got {len(xs)} and {len(state_dict)} '\n f'at path {current_path()}'\n )\n ys = []\n for i in range(len(state_dict)):\n y = from_state_dict(xs[i], state_dict[str(i)], name=str(i))\n ys.append(y)\n return ys\n\n\ndef _dict_state_dict(xs: dict[str, Any]) -> dict[str, Any]:\n str_keys = {str(k) for k in xs.keys()}\n if len(str_keys) != len(xs):\n raise ValueError(\n 'Dict keys do not have a unique string representation: '\n f'{str_keys} vs given: {xs}'\n )\n return {str(key): to_state_dict(value) for key, value in xs.items()}\n\n\ndef _restore_dict(xs, states: dict[str, Any]) -> dict[str, Any]:\n diff = set(map(str, xs.keys())).difference(states.keys())\n if diff:\n raise ValueError(\n 'The target dict keys and state dict keys do not match, target dict'\n f' contains keys {diff} which are not present in state dict at path'\n f' {current_path()}'\n )\n\n return {\n key: from_state_dict(value, states[str(key)], name=str(key))\n for key, value in xs.items()\n }\n\n\ndef _namedtuple_state_dict(nt) -> dict[str, Any]:\n return {key: to_state_dict(getattr(nt, key)) for key in nt._fields}\n\n\ndef _restore_namedtuple(xs, state_dict: dict[str, Any]):\n \"\"\"Rebuild namedtuple from serialized dict.\"\"\"\n if set(state_dict.keys()) == {'name', 'fields', 'values'}:\n # TODO(jheek): remove backward compatible named tuple restoration early 2022\n state_dict = {\n state_dict['fields'][str(i)]: state_dict['values'][str(i)]\n for i in range(len(state_dict['fields']))\n }\n\n sd_keys = set(state_dict.keys())\n nt_keys = set(xs._fields)\n\n if sd_keys != nt_keys:\n raise ValueError(\n 'The field names of the state dict and the named tuple do not match,'\n f' got {sd_keys} and {nt_keys} at path {current_path()}'\n )\n fields = {\n k: from_state_dict(getattr(xs, k), v, name=k) for k, v in state_dict.items()\n }\n return type(xs)(**fields)\n\n\nregister_serialization_state(dict, _dict_state_dict, _restore_dict)\nregister_serialization_state(list, _list_state_dict, _restore_list)\nregister_serialization_state(\n tuple,\n _list_state_dict,\n lambda xs, state_dict: tuple(_restore_list(list(xs), state_dict)),\n)\nregister_serialization_state(\n _NamedTuple, _namedtuple_state_dict, _restore_namedtuple\n)\n\nregister_serialization_state(\n jax.tree_util.Partial,\n lambda x: (\n {\n 'args': to_state_dict(x.args),\n 'keywords': to_state_dict(x.keywords),\n }\n ),\n lambda x, sd: jax.tree_util.Partial(\n x.func,\n *from_state_dict(x.args, sd['args']),\n **from_state_dict(x.keywords, sd['keywords']),\n ),\n)\n\n# On-the-wire / disk serialization format\n\n# We encode state-dicts via msgpack, using its custom type extension.\n# https://github.com/msgpack/msgpack/blob/master/spec.md\n#\n# - ndarrays and DeviceArrays are serialized to nested msgpack-encoded string\n# of (shape-tuple, dtype-name (e.g. 'float32'), row-major array-bytes).\n# Note: only simple ndarray types are supported, no objects or fields.\n#\n# - native complex scalars are converted to nested msgpack-encoded tuples\n# (real, imag).\n\n\ndef _ndarray_to_bytes(arr) -> bytes:\n \"\"\"Save ndarray to simple msgpack encoding.\"\"\"\n if isinstance(arr, jax.Array):\n arr = np.array(arr)\n if arr.dtype.hasobject or arr.dtype.isalignedstruct:\n raise ValueError(\n 'Object and structured dtypes not supported '\n 'for serialization of ndarrays.'\n )\n tpl = (arr.shape, arr.dtype.name, arr.tobytes('C'))\n return msgpack.packb(tpl, use_bin_type=True)\n\n\ndef _dtype_from_name(name: str):\n \"\"\"Handle JAX bfloat16 dtype correctly.\"\"\"\n if name == b'bfloat16':\n return jax.numpy.bfloat16\n else:\n return np.dtype(name)\n\n\ndef _ndarray_from_bytes(data: bytes) -> np.ndarray:\n \"\"\"Load ndarray from simple msgpack encoding.\"\"\"\n shape, dtype_name, buffer = msgpack.unpackb(data, raw=True)\n return np.frombuffer(\n buffer, dtype=_dtype_from_name(dtype_name), count=-1, offset=0\n ).reshape(shape, order='C')\n\n\nclass _MsgpackExtType(enum.IntEnum):\n \"\"\"Messagepack custom type ids.\"\"\"\n\n ndarray = 1\n native_complex = 2\n npscalar = 3\n\n\ndef _msgpack_ext_pack(x):\n \"\"\"Messagepack encoders for custom types.\"\"\"\n # TODO(flax-dev): Array here only work when they are fully addressable.\n # If they are not fully addressable, use the GDA path for checkpointing.\n if isinstance(x, (np.ndarray, jax.Array)):\n return msgpack.ExtType(_MsgpackExtType.ndarray, _ndarray_to_bytes(x))\n if isinstance(x, np.generic):\n # pack scalar as ndarray\n return msgpack.ExtType(\n _MsgpackExtType.npscalar, _ndarray_to_bytes(np.asarray(x))\n )\n elif isinstance(x, complex):\n return msgpack.ExtType(\n _MsgpackExtType.native_complex, msgpack.packb((x.real, x.imag))\n )\n return x\n\n\ndef _msgpack_ext_unpack(code, data):\n \"\"\"Messagepack decoders for custom types.\"\"\"\n if code == _MsgpackExtType.ndarray:\n return _ndarray_from_bytes(data)\n elif code == _MsgpackExtType.native_complex:\n complex_tuple = msgpack.unpackb(data)\n return complex(complex_tuple[0], complex_tuple[1])\n elif code == _MsgpackExtType.npscalar:\n ar = _ndarray_from_bytes(data)\n return ar[()] # unpack ndarray to scalar\n return msgpack.ExtType(code, data)\n\n\n# Chunking array leaves\n\n# msgpack has a hard limit of 2**31 - 1 bytes per object leaf. To circumvent\n# this limit for giant arrays (e.g. embedding tables), we traverse the tree\n# and break up arrays near the limit into flattened array chunks.\n\n# True limit is 2**31 - 1, but leave a margin for encoding padding.\nMAX_CHUNK_SIZE = 2**30\n\n\ndef _np_convert_in_place(d):\n \"\"\"Convert any jax devicearray leaves to numpy arrays in place.\"\"\"\n if isinstance(d, dict):\n for k, v in d.items():\n if isinstance(v, jax.Array):\n d[k] = np.array(v)\n elif isinstance(v, dict):\n _np_convert_in_place(v)\n elif isinstance(d, jax.Array):\n return np.array(d)\n return d\n\n\n_tuple_to_dict = lambda tpl: {str(x): y for x, y in enumerate(tpl)}\n_dict_to_tuple = lambda dct: tuple(dct[str(i)] for i in range(len(dct)))\n\n\ndef _chunk(arr) -> dict[str, Any]:\n \"\"\"Convert array to a canonical dictionary of chunked arrays.\"\"\"\n chunksize = max(1, int(MAX_CHUNK_SIZE / arr.dtype.itemsize))\n data = {'__msgpack_chunked_array__': True, 'shape': _tuple_to_dict(arr.shape)}\n flatarr = arr.reshape(-1)\n chunks = [\n flatarr[i : i + chunksize] for i in range(0, flatarr.size, chunksize)\n ]\n data['chunks'] = _tuple_to_dict(chunks)\n return data\n\n\ndef _unchunk(data: dict[str, Any]):\n \"\"\"Convert canonical dictionary of chunked arrays back into array.\"\"\"\n assert '__msgpack_chunked_array__' in data\n shape = _dict_to_tuple(data['shape'])\n flatarr = np.concatenate(_dict_to_tuple(data['chunks']))\n return flatarr.reshape(shape)\n\n\ndef _chunk_array_leaves_in_place(d):\n \"\"\"Convert oversized array leaves to safe chunked form in place.\"\"\"\n if isinstance(d, dict):\n for k, v in d.items():\n if isinstance(v, np.ndarray):\n if v.size * v.dtype.itemsize > MAX_CHUNK_SIZE:\n d[k] = _chunk(v)\n elif isinstance(v, dict):\n _chunk_array_leaves_in_place(v)\n elif isinstance(d, np.ndarray):\n if d.size * d.dtype.itemsize > MAX_CHUNK_SIZE:\n return _chunk(d)\n return d\n\n\ndef _unchunk_array_leaves_in_place(d):\n \"\"\"Convert chunked array leaves back into array leaves, in place.\"\"\"\n if isinstance(d, dict):\n if '__msgpack_chunked_array__' in d:\n return _unchunk(d)\n else:\n for k, v in d.items():\n if isinstance(v, dict) and '__msgpack_chunked_array__' in v:\n d[k] = _unchunk(v)\n elif isinstance(v, dict):\n _unchunk_array_leaves_in_place(v)\n return d\n\n\n# User-facing API calls:\n\n\ndef msgpack_serialize(pytree, in_place: bool = False) -> bytes:\n \"\"\"Save data structure to bytes in msgpack format.\n\n Low-level function that only supports python trees with array leaves,\n for custom objects use ``to_bytes``. It splits arrays above MAX_CHUNK_SIZE into\n multiple chunks.\n\n Args:\n pytree: python tree of dict, list, tuple with python primitives\n and array leaves.\n in_place: boolean specifying if pytree should be modified in place.\n\n Returns:\n msgpack-encoded bytes of pytree.\n \"\"\"\n if not in_place:\n pytree = jax.tree_util.tree_map(lambda x: x, pytree)\n pytree = _np_convert_in_place(pytree)\n pytree = _chunk_array_leaves_in_place(pytree)\n return msgpack.packb(pytree, default=_msgpack_ext_pack, strict_types=True)\n\n\ndef msgpack_restore(encoded_pytree: bytes):\n \"\"\"Restore data structure from bytes in msgpack format.\n\n Low-level function that only supports python trees with array leaves,\n for custom objects use ``from_bytes``.\n\n Args:\n encoded_pytree: msgpack-encoded bytes of python tree.\n\n Returns:\n Python tree of dict, list, tuple with python primitive\n and array leaves.\n \"\"\"\n state_dict = msgpack.unpackb(\n encoded_pytree, ext_hook=_msgpack_ext_unpack, raw=False\n )\n return _unchunk_array_leaves_in_place(state_dict)\n\n\ndef from_bytes(target, encoded_bytes: bytes):\n \"\"\"Restore optimizer or other object from msgpack-serialized state-dict.\n\n Args:\n target: template object with state-dict registrations that matches\n the structure being deserialized from ``encoded_bytes``.\n encoded_bytes: msgpack serialized object structurally isomorphic to\n ``target``. Typically a flax model or optimizer.\n\n Returns:\n A new object structurally isomorphic to ``target`` containing the updated\n leaf data from saved data.\n \"\"\"\n state_dict = msgpack_restore(encoded_bytes)\n return from_state_dict(target, state_dict)\n\n\ndef to_bytes(target) -> bytes:\n \"\"\"Save optimizer or other object as msgpack-serialized state-dict.\n\n Args:\n target: template object with state-dict registrations to be\n serialized to msgpack format. Typically a flax model or optimizer.\n\n Returns:\n Bytes of msgpack-encoded state-dict of ``target`` object.\n \"\"\"\n state_dict = to_state_dict(target)\n return msgpack_serialize(state_dict, in_place=True)\n" }, { "path": "flax/struct.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Utilities for defining custom classes that can be used with jax transformations.\"\"\"\n\nfrom collections.abc import Callable\nimport dataclasses\nimport functools\nfrom typing import TypeVar, overload\n\nimport jax\nfrom typing_extensions import (\n dataclass_transform, # pytype: disable=not-supported-yet\n)\n\nfrom . import serialization\n\n_T = TypeVar('_T')\n\n\ndef field(pytree_node=True, *, metadata=None, **kwargs):\n return dataclasses.field(metadata=(metadata or {}) | {'pytree_node': pytree_node},\n **kwargs)\n\n\n@dataclass_transform(field_specifiers=(field,)) # type: ignore[literal-required]\n@overload\ndef dataclass(clz: _T, **kwargs) -> _T:\n ...\n\n\n@dataclass_transform(field_specifiers=(field,)) # type: ignore[literal-required]\n@overload\ndef dataclass(**kwargs) -> Callable[[_T], _T]:\n ...\n\n\n@dataclass_transform(field_specifiers=(field,)) # type: ignore[literal-required]\ndef dataclass(\n clz: _T | None = None,\n **kwargs,\n) -> _T | Callable[[_T], _T]:\n \"\"\"Create a class which can be passed to functional transformations.\n\n .. note::\n Inherit from ``PyTreeNode`` instead to avoid type checking issues when\n using PyType.\n\n Jax transformations such as ``jax.jit`` and ``jax.grad`` require objects that are\n immutable and can be mapped over using the ``jax.tree_util`` methods.\n The ``dataclass`` decorator makes it easy to define custom classes that can be\n passed safely to Jax. Define JAX data as normal attribute fields, and use\n ``pytree_node=False`` to define static metadata.\n\n See example::\n\n >>> from flax import struct\n >>> import jax\n >>> from typing import Any, Callable\n\n >>> @struct.dataclass\n ... class Model:\n ... params: Any\n ... # use pytree_node=False to indicate an attribute should not be touched\n ... # by Jax transformations.\n ... apply_fn: Callable = struct.field(pytree_node=False)\n\n ... def __apply__(self, *args):\n ... return self.apply_fn(*args)\n\n >>> params = {}\n >>> params_b = {}\n >>> apply_fn = lambda v, x: x\n >>> model = Model(params, apply_fn)\n\n >>> # model.params = params_b # Model is immutable. This will raise an error.\n >>> model_b = model.replace(params=params_b) # Use the replace method instead.\n\n >>> # This class can now be used safely in Jax to compute gradients w.r.t. the\n >>> # parameters.\n >>> model = Model(params, apply_fn)\n >>> loss_fn = lambda model: 3.\n >>> model_grad = jax.grad(loss_fn)(model)\n\n Note that dataclasses have an auto-generated ``__init__`` where\n the arguments of the constructor and the attributes of the created\n instance match 1:1. If you desire a \"smart constructor\", for example to\n optionally derive some of the attributes from others,\n make an additional static or class method. Consider the following example::\n\n >>> @struct.dataclass\n ... class DirectionAndScaleKernel:\n ... direction: jax.Array\n ... scale: jax.Array\n\n ... @classmethod\n ... def create(cls, kernel):\n ... scale = jax.numpy.linalg.norm(kernel, axis=0, keepdims=True)\n ... direction = direction / scale\n ... return cls(direction, scale)\n\n Args:\n clz: the class that will be transformed by the decorator.\n **kwargs: arguments to pass to the dataclass constructor.\n\n Returns:\n The new class.\n \"\"\"\n # Support passing arguments to the decorator (e.g. @dataclass(kw_only=True))\n if clz is None:\n return functools.partial(dataclass, **kwargs) # type: ignore[bad-return-type]\n\n # check if already a flax dataclass\n if '_flax_dataclass' in clz.__dict__:\n return clz\n\n if 'frozen' not in kwargs.keys():\n kwargs['frozen'] = True\n data_clz = dataclasses.dataclass(**kwargs)(clz) # type: ignore\n meta_fields = []\n data_fields = []\n for field_info in dataclasses.fields(data_clz):\n is_pytree_node = field_info.metadata.get('pytree_node', True)\n if is_pytree_node:\n data_fields.append(field_info.name)\n else:\n meta_fields.append(field_info.name)\n\n def replace(self, **updates):\n \"\"\"Returns a new object replacing the specified fields with new values.\"\"\"\n return dataclasses.replace(self, **updates)\n\n data_clz.replace = replace\n\n jax.tree_util.register_dataclass(data_clz, data_fields, meta_fields)\n\n def to_state_dict(x):\n state_dict = {\n name: serialization.to_state_dict(getattr(x, name))\n for name in data_fields\n }\n return state_dict\n\n def from_state_dict(x, state):\n \"\"\"Restore the state of a data class.\"\"\"\n state = state.copy() # copy the state so we can pop the restored fields.\n updates = {}\n for name in data_fields:\n if name not in state:\n raise ValueError(\n f'Missing field {name} in state dict while restoring'\n f' an instance of {clz.__name__},'\n f' at path {serialization.current_path()}'\n )\n value = getattr(x, name)\n value_state = state.pop(name)\n updates[name] = serialization.from_state_dict(\n value, value_state, name=name\n )\n if state:\n names = ','.join(state.keys())\n raise ValueError(\n f'Unknown field(s) \"{names}\" in state dict while'\n f' restoring an instance of {clz.__name__}'\n f' at path {serialization.current_path()}'\n )\n return x.replace(**updates)\n\n serialization.register_serialization_state(\n data_clz, to_state_dict, from_state_dict\n )\n\n # add a _flax_dataclass flag to distinguish from regular dataclasses\n data_clz._flax_dataclass = True # type: ignore[attr-defined]\n\n return data_clz # type: ignore\n\n\nTNode = TypeVar('TNode', bound='PyTreeNode')\n\n\n@dataclass_transform(field_specifiers=(field,)) # type: ignore[literal-required]\nclass PyTreeNode:\n \"\"\"Base class for dataclasses that should act like a JAX pytree node.\n\n See ``flax.struct.dataclass`` for the ``jax.tree_util`` behavior.\n This base class additionally avoids type checking errors when using PyType.\n\n Example::\n\n >>> from flax import struct\n >>> import jax\n >>> from typing import Any, Callable\n\n >>> class Model(struct.PyTreeNode):\n ... params: Any\n ... # use pytree_node=False to indicate an attribute should not be touched\n ... # by Jax transformations.\n ... apply_fn: Callable = struct.field(pytree_node=False)\n\n ... def __apply__(self, *args):\n ... return self.apply_fn(*args)\n\n >>> params = {}\n >>> params_b = {}\n >>> apply_fn = lambda v, x: x\n >>> model = Model(params, apply_fn)\n\n >>> # model.params = params_b # Model is immutable. This will raise an error.\n >>> model_b = model.replace(params=params_b) # Use the replace method instead.\n\n >>> # This class can now be used safely in Jax to compute gradients w.r.t. the\n >>> # parameters.\n >>> model = Model(params, apply_fn)\n >>> loss_fn = lambda model: 3.\n >>> model_grad = jax.grad(loss_fn)(model)\n \"\"\"\n\n def __init_subclass__(cls, **kwargs):\n super().__init_subclass__()\n dataclass(cls, **kwargs) # pytype: disable=wrong-arg-types\n\n def __init__(self, *args, **kwargs):\n # stub for pytype\n raise NotImplementedError\n\n def replace(self: TNode, **overrides) -> TNode:\n # stub for pytype\n raise NotImplementedError\n" }, { "path": "flax/testing/__init__.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\n\"\"\"Flax testing utilities.\"\"\"\n\nfrom .benchmark import Benchmark\n" }, { "path": "flax/testing/benchmark.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Benchmark class for Flax regression and integration testing.\n\nThis file defines utility functions for collecting model training results\nfrom TensorBoard summaries, and for reporting benchmarks in a JSON format\nfor pickup by continuous integration / monitoring frameworks.\n\nWhen the `benchmark_output_dir` is provided, the benchmark results are\nsaved in this directory in the JSON format with a single file per\nbenchmark.\n\"\"\"\n\nimport functools\nimport inspect\nimport itertools\nimport json\nimport os\nimport tempfile\n\nfrom absl import flags, logging\nfrom absl.testing import absltest\nfrom tensorboard.backend.event_processing import (\n directory_watcher,\n event_file_loader,\n io_wrapper,\n)\nfrom tensorboard.summary import v1 as summary_lib\nfrom tensorboard.util import tensor_util\n\nfrom flax import io\n\n_BENCHMARK_OUTPUT_DIR = flags.DEFINE_string(\n 'benchmark_output_dir', default=None, help='Benchmark output directory.'\n)\n\n_SCALAR_PLUGIN_NAME = (\n summary_lib.scalar_pb('', 0).value[0].metadata.plugin_data.plugin_name\n)\n\n\ndef _make_events_generator(path):\n \"\"\"Makes a generator yielding TensorBoard events from files in `path`.\"\"\"\n return directory_watcher.DirectoryWatcher(\n path, event_file_loader.EventFileLoader, io_wrapper.IsSummaryEventsFile\n ).Load()\n\n\ndef _is_scalar_value(value):\n if value.HasField('metadata') and value.metadata.HasField('plugin_data'):\n plugin_data = value.metadata.plugin_data\n return plugin_data.plugin_name == _SCALAR_PLUGIN_NAME\n\n return False\n\n\ndef _process_event(event):\n \"\"\"Parse TensorBoard scalars into a (tag, wall_time, step, scalar) tuple.\"\"\"\n for value in event.summary.value:\n if not _is_scalar_value(value):\n continue\n\n if value.HasField('tensor'):\n yield (\n value.tag,\n event.wall_time,\n event.step,\n tensor_util.make_ndarray(value.tensor).item(),\n )\n\n\ndef _get_tensorboard_scalars(path):\n \"\"\"Read and parse scalar TensorBoard summaries.\n\n Args:\n path: str. Path containing TensorBoard event files.\n\n Returns:\n Dictionary mapping summary tags (str) to lists of\n (wall_time, step, scalar) tuples.\n \"\"\"\n gen = _make_events_generator(path)\n data = filter(lambda x: x.HasField('summary'), gen)\n data = itertools.chain.from_iterable(map(_process_event, data))\n\n data_by_key = {}\n for tag, wall_time, step, value in data:\n if not tag in data_by_key:\n data_by_key[tag] = []\n data_by_key[tag].append((wall_time, step, value))\n return data_by_key\n\n\nclass Benchmark(absltest.TestCase):\n \"\"\"Benchmark class for Flax examples.\n\n This class overrides the behaviour of `self.assert*` methods to be\n deferred instead of failing immediately. This allows for using absltest\n assert methods for checking benchmark target metrics. This is also\n necessary for correctly reporting benchmark results and determining its\n success.\n \"\"\"\n\n def __init__(self, *args, **kwargs):\n \"\"\"Wrap test methods in a try-except decorator to delay exceptions.\"\"\"\n super().__init__(*args, **kwargs)\n for func_name in dir(self):\n if func_name.startswith('assert'):\n func = getattr(self, func_name)\n patched_func = functools.partial(self._collect_assert_wrapper, fn=func)\n setattr(self, func_name, patched_func)\n\n self._benchmark_output_dir = _BENCHMARK_OUTPUT_DIR.value\n # Create target directory if defined.\n if self._benchmark_output_dir and not io.exists(self._benchmark_output_dir):\n io.makedirs(self._benchmark_output_dir)\n\n # pylint: disable=invalid-name\n def _collect_assert_wrapper(self, *args, fn=None, **kwargs):\n \"\"\"Wrapper around assert methods that caputres and collects failures.\"\"\"\n try:\n return fn(*args, **kwargs)\n except self.failureException as ex:\n self._outstanding_fails.append(ex)\n\n def setUp(self):\n \"\"\"Setup ran before each test.\"\"\"\n super().setUp()\n self._reported_name = None\n self._reported_wall_time = None\n self._reported_metrics = {}\n self._reported_extras = {}\n self._outstanding_fails = []\n\n def tearDown(self):\n \"\"\"Tear down after each test.\"\"\"\n super().tearDown()\n self._report_benchmark_results()\n for message in self._outstanding_fails:\n raise self.failureException(message)\n\n def get_tmp_model_dir(self):\n \"\"\"Returns an unique temporary directory for storing model data.\n\n Returns path by appending Classname.testname to `benchmark_output_dir` flag\n if defined else uses a temporary directory. This helps to export summary\n files to tensorboard as multiple separate runs for each test method.\n \"\"\"\n if self._benchmark_output_dir:\n model_dir = self._benchmark_output_dir\n else:\n model_dir = tempfile.mkdtemp()\n model_dir_path = os.path.join(\n model_dir, self._reported_name or self._get_test_name()\n )\n # Create directories if they don't exist.\n if not io.exists(model_dir_path):\n io.makedirs(model_dir_path)\n return model_dir_path\n\n def has_outstanding_fails(self):\n \"\"\"Determine whether the benchmark failed, but the error is deferred.\"\"\"\n return len(self._outstanding_fails) > 0\n\n def read_summaries(self, path):\n \"\"\"Read TensorBoard summaries.\"\"\"\n return _get_tensorboard_scalars(path)\n\n def report_wall_time(self, wall_time: float):\n \"\"\"Report wall time for the benchmark.\"\"\"\n self._update_reported_name()\n self._reported_wall_time = wall_time\n\n def report_metrics(self, metrics: dict[str, float]):\n \"\"\"Report metrics for the benchmark.\"\"\"\n self._update_reported_name()\n self._reported_metrics.update(metrics)\n\n def report_metric(self, name: str, value: float):\n \"\"\"Report a single metric for the benchmark.\"\"\"\n self.report_metrics({name: value})\n\n def report_extras(self, extras: dict[str, str]):\n \"\"\"Report extras for the benchmark.\"\"\"\n self._update_reported_name()\n self._reported_extras.update(extras)\n\n def report_extra(self, name: str, value: str):\n \"\"\"Report a single extra for the benchmark.\"\"\"\n self.report_extras({name: value})\n\n def _get_test_name(self, prefix='test_'):\n \"\"\"Returns full name of test class and method calling report_benchmark.\n\n The name is based on the *outermost* Benchmark class in the class stack.\n Based on tensorflow/python/platform/benchmark.py\n\n Args:\n prefix: str. Prefix that the caller method must have.\n\n Returns:\n Resolved test name as `ClassName.test_name`.\n \"\"\"\n # Find the caller method (outermost Benchmark class).\n stack = inspect.stack()\n calling_class, name = None, None\n for frame_info in stack[::-1]:\n f_locals = frame_info.frame.f_locals\n f_self = f_locals.get('self', None)\n if isinstance(f_self, Benchmark):\n name = frame_info.function\n if name.startswith(prefix):\n calling_class = f_self\n break\n if calling_class is None:\n raise ValueError('Unable to determine the calling Benchmark class.')\n\n # Prefix the name with the class name.\n class_name = type(calling_class).__name__\n name = f'{class_name}.{name}'\n return name\n\n def _update_reported_name(self):\n \"\"\"Record / update test name for the benchmark.\"\"\"\n self._reported_name = self._reported_name or self._get_test_name()\n\n def _report_benchmark_results(self):\n \"\"\"Produce benchmark results report.\n\n Results are reported as a JSON string with the following schema:\n ```\n {\n \"name\": \n \"succeeded\": true / false\n \"wall_time\": float (containing wall-time for the benchmark)\n \"metrics\": {\n \"string\" -> float map of other performance metrics\n }\n \"extras\": {\n \"string\" -> \"string\" map containing anything else of interest\n }\n }\n ```\n \"\"\"\n name = self._reported_name\n if not name:\n raise ValueError(\n 'Unable to determine test name for reporting '\n 'benchmark results. Make sure you are using '\n '`self.report_*` methods.'\n )\n\n succeeded = not self.has_outstanding_fails()\n results = {\n 'name': name,\n 'succeeded': succeeded,\n 'metrics': self._reported_metrics,\n 'extras': self._reported_extras,\n }\n if self._reported_wall_time is not None:\n results['wall_time'] = self._reported_wall_time\n if not succeeded:\n msg = '\\n'.join([str(fail) for fail in self._outstanding_fails])\n results['extras']['failed_assertions'] = msg\n\n results_str = json.dumps(results)\n logging.info(results_str)\n\n # Maybe save results as a file for pickup by CI / monitoring frameworks.\n if self._benchmark_output_dir:\n filename = os.path.join(self._benchmark_output_dir, name + '.json')\n with io.GFile(filename, 'w') as fout:\n fout.write(results_str)\n" }, { "path": "flax/traceback_util.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Flax specific traceback_util functions.\"\"\"\n\nfrom jax._src import traceback_util as jax_traceback_util\nfrom jax.extend import source_info_util\n\nfrom flax import config\n\n# pylint: disable=protected-access\n\n# Globals:\n# Whether to filter flax frames from traceback.\n_flax_filter_tracebacks = config.flax_filter_frames\n# Flax specific set of paths to exclude from tracebacks.\n_flax_exclusions = set()\n\n\n# re-import JAX symbol for convenience.\napi_boundary = jax_traceback_util.api_boundary\n\n\ndef register_exclusion(path):\n \"\"\"Marks a Flax source file for exclusion.\"\"\"\n global _flax_exclusions, _flax_filter_tracebacks\n # Record flax exclusions so we can dynamically add and remove them.\n _flax_exclusions.add(path)\n if _flax_filter_tracebacks:\n jax_traceback_util.register_exclusion(path)\n source_info_util.register_exclusion(path)\n\ndef hide_flax_in_tracebacks():\n \"\"\"Hides Flax internal stack frames in tracebacks.\"\"\"\n global _flax_exclusions, _flax_filter_tracebacks\n _flax_filter_tracebacks = True\n for exclusion in _flax_exclusions:\n if exclusion not in jax_traceback_util._exclude_paths:\n jax_traceback_util._exclude_paths.append(exclusion)\n\n\ndef show_flax_in_tracebacks():\n \"\"\"Shows Flax internal stack frames in tracebacks.\"\"\"\n global _flax_exclusions, _flax_filter_tracebacks\n _flax_filter_tracebacks = False\n for exclusion in _flax_exclusions:\n if exclusion in jax_traceback_util._exclude_paths:\n jax_traceback_util._exclude_paths.remove(exclusion)\n" }, { "path": "flax/training/__init__.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Flax training utilities.\"\"\"\n" }, { "path": "flax/training/checkpoints.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Checkpointing helper functions.\n\nHandles saving and restoring optimizer checkpoints based on step-number or\nother numerical metric in filename. Cleans up older / worse-performing\ncheckpoint files.\n\"\"\"\n\nimport functools\nimport os\nimport pathlib\nimport re\nimport time\nimport warnings\nfrom concurrent.futures import thread\nfrom typing import (\n Any,\n)\nfrom collections.abc import Callable, Iterable\n\nimport jax\nimport orbax.checkpoint as ocp\nfrom absl import logging\nfrom jax import monitoring, process_index\nfrom jax import tree_util as jtu\nfrom jax.experimental.array_serialization.serialization import (\n GlobalAsyncCheckpointManager,\n get_tensorstore_spec,\n)\nfrom jax.experimental.multihost_utils import sync_global_devices\n\nfrom flax import config, core, errors, io, serialization, traverse_util\nfrom flax.training import orbax_utils\n\n\n_READ_CHECKPOINT_EVENT: str = '/jax/checkpoint/read/durations_sec'\n_WRITE_CHECKPOINT_EVENT: str = '/jax/checkpoint/write/durations_sec'\n\n\n# Single-group reg-exps for int or float numerical substrings.\n# captures sign:\nSIGNED_FLOAT_RE = re.compile(r'([-+]?(?:\\d+(?:\\.\\d*)?|\\.\\d+)(?:[eE][-+]?\\d+)?)')\n# does not capture sign:\nUNSIGNED_FLOAT_RE = re.compile(\n r'[-+]?((?:\\d+(?:\\.\\d*)?|\\.\\d+)(?:[eE][-+]?\\d+)?)'\n)\n# Module name followed by number.\nMODULE_NUM_RE = re.compile(r'(.*)_\\d+$')\n# Alternative schemes handled by `gfile`, e.g. on Google Cloud Storage (GCS).\nSCHEME_RE = re.compile('^(?P[a-z][a-z0-9.+-]+://)?(?P.*)', re.I)\n\n# Multiprocess arrays (GlobalDeviceArray, or JAX array with multiprocess\n# sharding) is across processes and will be stored in directories with this\n# postfix, separated from the non-distributed data (e.g. the larger pytree)\nMP_ARRAY_POSTFIX = '_gda'\n# Occurrences of multiprocess arrays in the target pytree will be\n# replaced by this string placeholder.\nMP_ARRAY_PH = '//GDAPlaceholder:'\n\n# Add a copy-success file to a distributed array directory to indicate the\n# array save is complete.\n# We need this for GCS because GCS's directory move is not atomic.\nCOMMIT_SUCCESS_FILE = 'commit_success.txt'\n\n# Orbax main checkpoint file name.\nORBAX_CKPT_FILENAME = 'checkpoint'\nORBAX_MANIFEST_OCDBT = 'manifest.ocdbt'\nORBAX_METADATA_FILENAME = '_METADATA'\n\nPyTree = Any\n\n# TODO(flax-dev): Remove this once flax is using the latest jax release\n# containing jax.Array attribute.\nMultiprocessArrayType = Any\n\n\ndef _is_multiprocess_array(value: Any) -> bool:\n \"\"\"Use GlobalAsyncCheckpointManager to save the array if it's only partially available on this host.\"\"\"\n if isinstance(value, jax.Array):\n return not value.is_fully_addressable\n return False\n\n\ndef _checkpoint_path(\n ckpt_dir: str, step: int | float | str, prefix: str = 'checkpoint_'\n) -> str:\n return os.path.join(ckpt_dir, f'{prefix}{step}')\n\n\ndef _checkpoint_path_step(path: str) -> float | None:\n \"\"\"Returns the step number of a checkpoint path.\"\"\"\n for s in SIGNED_FLOAT_RE.split(path)[::-1]:\n if SIGNED_FLOAT_RE.match(s):\n return float(s)\n return None\n\n\ndef _allowempty_listdir(path: str):\n try:\n return io.listdir(path)\n except io.NotFoundError:\n return []\n\n\ndef _safe_remove(path: str):\n \"\"\"Identify whether a path is a dir or list and choose the correct remove method.\"\"\"\n if io.isdir(path):\n io.rmtree(path)\n else:\n io.remove(path)\n\n\ndef _is_orbax_checkpoint(path: str) -> bool:\n return (\n io.exists(os.path.join(path, ORBAX_CKPT_FILENAME))\n or io.exists(os.path.join(path, ORBAX_METADATA_FILENAME))\n or io.exists(os.path.join(path, ORBAX_MANIFEST_OCDBT))\n )\n\n\nclass AsyncManager:\n \"\"\"A simple object to track async checkpointing.\n\n How to use: create an instance and pass to save_checkpoint() calls:\n am = AsyncManager()\n save_checkpoint(..., async_manager=am)\n \"\"\"\n\n def __init__(self, max_workers: int = 1):\n self.executor = thread.ThreadPoolExecutor(max_workers=max_workers)\n self.save_future = None\n\n def wait_previous_save(self):\n \"\"\"Block until the previous save finishes, to keep files' consistency.\"\"\"\n if self.save_future and not self.save_future.done():\n logging.warning(\n 'The previous async save_checkpoint has not finished yet. Waiting '\n 'for it to complete before the next save.'\n )\n self.save_future.result()\n\n def save_async(self, task: Callable[[], Any]):\n \"\"\"Run a task async.\n\n The future will be tracked as self.save_future.\n\n Args:\n task: The callable to be executed asynchronously.\n \"\"\"\n self.wait_previous_save()\n self.save_future = self.executor.submit(task) # type: ignore\n\n\ndef _split_mp_arrays(\n target: dict[str, Any]\n) -> tuple[dict[str, Any], list[tuple[MultiprocessArrayType, str]]]:\n \"\"\"Split out the multiprocess arrays from the target pytree to save.\"\"\"\n # When target is a single leaf instead of a pytree dict.\n if not isinstance(target, (core.FrozenDict, dict)):\n if _is_multiprocess_array(target):\n return MP_ARRAY_PH, [(target, '')]\n return target, []\n # Traverse the target and handle distributed arrays.\n flattened = traverse_util.flatten_dict(target, keep_empty_nodes=True)\n mpa_targets = []\n for key, value in flattened.items():\n if _is_multiprocess_array(value):\n subpath = '/'.join(key)\n mpa_targets.append((value, subpath))\n flattened[key] = MP_ARRAY_PH + subpath\n target = traverse_util.unflatten_dict(flattened)\n return target, mpa_targets\n\n\ndef _make_mpa_dirs(\n mpa_targets: list[tuple[MultiprocessArrayType, str]], tmp_path: str\n):\n # Temporary array path is not used in GCS.\n if tmp_path.startswith('gs://'):\n return\n mpa_tmp_path = tmp_path + MP_ARRAY_POSTFIX\n # Clean up the previous MPA dir, in case some leftover from last preemption\n # lingers.\n if io.exists(mpa_tmp_path):\n logging.info('Removing outdated MPA temporary files at %s', mpa_tmp_path)\n io.rmtree(mpa_tmp_path)\n _, mpa_subpaths = zip(*mpa_targets)\n for subpath in mpa_subpaths:\n io.makedirs(os.path.join(mpa_tmp_path, subpath))\n\n\ndef _save_mpas(\n gda_manager,\n mpa_targets: list[tuple[MultiprocessArrayType, str]],\n tmp_path: str,\n final_path: str,\n base_path: str,\n keep: int,\n overwrite: bool,\n keep_every_n_steps: int | None,\n ckpt_start_time: float,\n async_manager: AsyncManager | None = None,\n):\n \"\"\"Save the multiprocess arrays given the paths.\"\"\"\n mpa_list, mpa_subpaths = zip(*mpa_targets)\n mpa_tmp_path, mpa_final_path = (\n tmp_path + MP_ARRAY_POSTFIX,\n final_path + MP_ARRAY_POSTFIX,\n )\n write_commit_success = False\n # If the checkpoint directory is a GCS directory, then keep the final\n # checkpoint directory as the temporary checkpoint directory. This is because\n # renames are not atomic on GCS. When restoring check for the existence of a\n # success file.\n # TODO: figure out a way to unit-test the behavior.\n if tmp_path.startswith('gs://'):\n mpa_tmp_path = mpa_final_path\n write_commit_success = True\n mpa_paths = [os.path.join(mpa_tmp_path, x) for x in mpa_subpaths]\n ts_specs = [get_tensorstore_spec(x) for x in mpa_paths]\n gda_manager.serialize(\n list(mpa_list),\n ts_specs,\n on_commit_callback=functools.partial(\n _save_commit,\n tmp_path,\n final_path,\n base_path,\n keep,\n overwrite,\n keep_every_n_steps,\n ckpt_start_time,\n has_mpa=True,\n write_commit_success=write_commit_success,\n async_manager=async_manager,\n ),\n )\n\n\ndef _restore_mpas(\n state_dict,\n target: Any | None,\n ckpt_path: str,\n step: int | float | None,\n gda_manager: GlobalAsyncCheckpointManager | None,\n allow_partial: bool = False,\n):\n \"\"\"Restore the multiprocess arrays given the target structure and type.\"\"\"\n\n def _check_mpa_errors():\n if not gda_manager:\n raise errors.MPACheckpointingRequiredError(ckpt_path, step)\n if not target and not allow_partial:\n raise errors.MPARestoreTargetRequiredError(ckpt_path, step)\n\n def _safe_deserialize(\n target_mpas: list[tuple[tuple[Any, ...], MultiprocessArrayType, str]],\n gda_manager: Any,\n ) -> list[MultiprocessArrayType]:\n gda_manager.wait_until_finished()\n\n # Check if reading from GCS and the array dir is potentially corrupted.\n if ckpt_path.startswith('gs://') and not io.exists(\n os.path.join(ckpt_path + MP_ARRAY_POSTFIX, COMMIT_SUCCESS_FILE)\n ):\n raise errors.MPARestoreDataCorruptedError(step, ckpt_path)\n\n # Check if the given target array types are valid.\n shardings = []\n for _, arr, path in target_mpas:\n if isinstance(arr, jax.Array):\n shardings.append(arr.sharding)\n\n # Restore the arrays.\n ts_specs = [get_tensorstore_spec(path) for _, _, path in target_mpas]\n return gda_manager.deserialize(shardings, ts_specs)\n\n # When target is a single leaf instead of a pytree dict.\n if not isinstance(state_dict, (core.FrozenDict, dict)):\n if (\n _is_multiprocess_array(target)\n and isinstance(state_dict, str)\n and state_dict.startswith(MP_ARRAY_PH)\n ):\n _check_mpa_errors()\n return _safe_deserialize(\n [((), target, ckpt_path + MP_ARRAY_POSTFIX)], gda_manager\n )[0]\n return state_dict\n\n # Go through the restored checkpoint pytree for all MPAs\n flattened = traverse_util.flatten_dict(state_dict, keep_empty_nodes=True)\n target_flattened = {}\n if target:\n target_flattened = traverse_util.flatten_dict(\n serialization.to_state_dict(target), keep_empty_nodes=True\n )\n # A list of (state_dict_key, target_array, array_file_path) for every array\n # to be restored\n target_mpas = []\n for key, value in flattened.items():\n if isinstance(value, str) and value.startswith(MP_ARRAY_PH):\n _check_mpa_errors()\n if (\n not target\n or (key not in target_flattened)\n or (not _is_multiprocess_array(target_flattened[key]))\n ):\n if allow_partial:\n logging.warning(\n 'Multiprocess array %s could not be restored because a valid'\n ' array is not found in target at the corresponding location.'\n ' Proceed to restore other arrays because'\n ' allow_partial_restoration=True',\n key,\n )\n else:\n raise errors.MPARestoreTargetRequiredError(ckpt_path, step, key)\n else:\n mpa_path = os.path.join(\n ckpt_path + MP_ARRAY_POSTFIX, value[len(MP_ARRAY_PH) :]\n )\n target_mpas.append((key, target_flattened[key], mpa_path))\n\n # If any MPA needs to be restored, call deserialize\n if target_mpas:\n mpa_list = _safe_deserialize(target_mpas, gda_manager)\n for mpa, (key, _, _) in zip(mpa_list, target_mpas):\n flattened[key] = mpa\n state_dict = traverse_util.unflatten_dict(flattened)\n return state_dict\n\n\ndef natural_sort(file_list: Iterable[str], signed: bool = True) -> list[str]:\n \"\"\"Natural sort for filenames with numerical substrings.\n\n Args:\n file_list: list of paths to sort containing numerical substrings.\n signed: bool: if leading '-' (or '+') signs should be included in numerical\n substrings as a sign or treated as a separator.\n\n Returns:\n List of filenames sorted 'naturally', not lexicographically: any\n integer substrings are used to subsort numerically. e.g.\n file_1, file_10, file_2 --> file_1, file_2, file_10\n file_0.1, file_-0.2, file_2.0 --> file_-0.2, file_0.1, file_2.0\n \"\"\"\n float_re = SIGNED_FLOAT_RE if signed else UNSIGNED_FLOAT_RE\n\n def maybe_num(s):\n if float_re.match(s):\n return float(s)\n else:\n return s\n\n def split_keys(s):\n return [maybe_num(c) for c in float_re.split(s)]\n\n return sorted(file_list, key=split_keys)\n\n\ndef safe_normpath(path: str) -> str:\n \"\"\"Normalizes path safely to get around `io.glob()` limitations.\"\"\"\n match = SCHEME_RE.match(path)\n assert match is not None\n d = match.groupdict()\n return (d['scheme'] or '') + os.path.normpath(d['path'])\n\n\ndef _remove_invalid_ckpts(\n ckpt_path: str,\n base_path: str,\n keep: int,\n overwrite: bool,\n keep_every_n_steps: int | None,\n has_mpa: bool,\n) -> None:\n \"\"\"Clean up the checkpoint space according to `overwrite`, `keep`, and `keep_every_n_steps` parameters.\"\"\"\n dir_path, prefix = os.path.split(base_path)\n checkpoint_files: list[Any] = [\n pathlib.PurePath(c) for c in _allowempty_listdir(dir_path)\n ]\n checkpoint_files = [\n os.path.join(dir_path, c)\n for c in checkpoint_files\n if c.match(f'{prefix}*') and not c.match(f'*{MP_ARRAY_POSTFIX}')\n ]\n checkpoint_files = natural_sort(checkpoint_files)\n\n # Remove newer checkpoints\n if overwrite and ckpt_path in checkpoint_files:\n ind = checkpoint_files.index(ckpt_path) + 1\n newer_ckpts = checkpoint_files[ind:]\n checkpoint_files = checkpoint_files[:ind]\n for path in newer_ckpts:\n logging.info('Removing checkpoint at %s', path)\n if has_mpa:\n # MPA might be removed already but the main ckpt is still there. This\n # can happen if the job is previously preempted after deleting the MPA\n # checkpoint folder and before deleting the main checkpoint.\n if io.exists(path + MP_ARRAY_POSTFIX):\n io.rmtree(path + MP_ARRAY_POSTFIX)\n _safe_remove(path)\n\n # Remove old checkpoint files.\n last_kept = -float('inf')\n if len(checkpoint_files) > keep:\n old_ckpts = checkpoint_files[:-keep]\n # Note: old_ckpts is sorted from oldest to newest.\n for path in old_ckpts:\n if keep_every_n_steps:\n step_number = _checkpoint_path_step(path)\n if step_number and (step_number - last_kept) >= keep_every_n_steps:\n logging.debug(\n 'Not deleting %s, because last_kept=%f and keeping '\n 'every %d steps.',\n path,\n last_kept,\n keep_every_n_steps,\n )\n last_kept = step_number\n continue\n logging.info('Removing checkpoint at %s', path)\n if has_mpa:\n # MPA might be removed already but the main ckpt is still there.\n if io.exists(path + MP_ARRAY_POSTFIX):\n io.rmtree(path + MP_ARRAY_POSTFIX)\n _safe_remove(path)\n\n\ndef _save_commit(\n ckpt_tmp_path: str,\n ckpt_path: str,\n base_path: str,\n keep: int,\n overwrite: bool,\n keep_every_n_steps: int | None,\n ckpt_start_time: float,\n has_mpa: bool,\n write_commit_success: bool,\n async_manager: AsyncManager | None = None,\n) -> None:\n \"\"\"Commit changes after saving checkpoints to disk.\n\n This function does the following, sequentially:\n 1. Make sure all ckpt writing finishes, and rename them from temp path to\n the final path.\n 2. Remove newer checkpoints (files that ordered larger than this save) if\n `overwrite=True`.\n 3. Remove old checkpoint files based on `keep` and `keep_every_n_steps`.\n 4. Record program duration saved by this checkpoint.\n \"\"\"\n mpa_ckpt_tmp_path, mpa_ckpt_path = (\n ckpt_tmp_path + MP_ARRAY_POSTFIX,\n ckpt_path + MP_ARRAY_POSTFIX,\n )\n # Rename the multiprocess array path once serialization and writing finished.\n if has_mpa:\n if write_commit_success:\n commit_success_path = os.path.join(mpa_ckpt_path, COMMIT_SUCCESS_FILE)\n with io.GFile(commit_success_path, 'w') as f:\n f.write(f'Checkpoint commit was successful to {mpa_ckpt_path}')\n else:\n # Commits are a two stage process (renaming the array folder and renaming\n # the main ckpt file in sequential order). We always try to overwrite\n # here because the array ckpt might be already renamed in a previously\n # interrupted commit. NOTE: io.rename does not support overwriting\n # directories via `rename` so we manually overwrite it.\n if io.exists(mpa_ckpt_path):\n logging.info('Removing outdated checkpoint at %s', mpa_ckpt_path)\n io.rmtree(mpa_ckpt_path)\n io.rename(mpa_ckpt_tmp_path, mpa_ckpt_path)\n # Commit the main checkpoint file after arrays (if any) are committed\n if async_manager:\n async_manager.wait_previous_save()\n io.rename(ckpt_tmp_path, ckpt_path, overwrite=overwrite)\n logging.info('Saved checkpoint at %s', ckpt_path)\n\n # Remove newer and older invalid checkpoints.\n _remove_invalid_ckpts(\n ckpt_path, base_path, keep, overwrite, keep_every_n_steps, has_mpa\n )\n # Record checkpoint-related metrics.\n ocp.utils.record_saved_duration(ckpt_start_time)\n if async_manager:\n jax.monitoring.record_event_duration_secs(\n '/jax/checkpoint/write/async/total_duration_secs',\n time.time() - ckpt_start_time,\n )\n\n\ndef _check_overwrite_error(\n ckpt_tmp_path: str, ckpt_path: str, base_path: str, step: int\n):\n \"\"\"Throw error if a ckpt file of this step or higher already exists.\"\"\"\n dir_path, prefix = os.path.split(base_path)\n checkpoint_files: list[Any] = [\n pathlib.PurePath(c) for c in _allowempty_listdir(dir_path)\n ]\n checkpoint_files = [\n os.path.join(dir_path, c)\n for c in checkpoint_files\n if c.match(f'{prefix}*') and not c.match(f'*{MP_ARRAY_POSTFIX}')\n ]\n if ckpt_path in checkpoint_files:\n raise errors.InvalidCheckpointError(ckpt_path, step)\n checkpoint_files.append(ckpt_path)\n\n checkpoint_files = natural_sort(checkpoint_files)\n # Handle the case if the job was preempted after the temporary checkpoint\n # was written, but before it was renamed to the final checkpoint name\n if checkpoint_files[-1] == ckpt_tmp_path:\n checkpoint_files.pop()\n if ckpt_path != checkpoint_files[-1]:\n raise errors.InvalidCheckpointError(ckpt_path, step)\n\n\ndef _save_main_ckpt_file(\n target: bytes,\n has_mpa: bool,\n paths: tuple[str, str],\n base_path: str,\n step: int,\n keep: int,\n overwrite: bool,\n keep_every_n_steps: int | None,\n ckpt_start_time: float,\n):\n \"\"\"Save the main checkpoint file via file system.\"\"\"\n ckpt_tmp_path, ckpt_path = paths\n io.makedirs(os.path.dirname(ckpt_path))\n\n with io.GFile(ckpt_tmp_path, 'wb') as fp:\n fp.write(target)\n\n # Postpone the commitment of checkpoint to after MPA writes are done.\n if not has_mpa:\n _save_commit(\n ckpt_tmp_path,\n ckpt_path,\n base_path,\n keep,\n overwrite,\n keep_every_n_steps,\n ckpt_start_time,\n has_mpa=False,\n write_commit_success=False,\n )\n\n\ndef _get_checkpoint_paths(\n ckpt_dir: str | os.PathLike,\n step: int | float,\n prefix: str = 'checkpoint_',\n) -> tuple[str, str, str]:\n \"\"\"Generate the checkpoint paths used in this save operation.\"\"\"\n ckpt_dir = os.fspath(ckpt_dir) # Pathlib -> str\n logging.info('Saving checkpoint at step: %s', step)\n # normalize path because io.glob() can modify path './', '//' ...\n ckpt_dir = safe_normpath(ckpt_dir)\n ckpt_tmp_path = _checkpoint_path(ckpt_dir, 'tmp', prefix)\n ckpt_path = _checkpoint_path(ckpt_dir, step, prefix)\n base_path = os.path.join(ckpt_dir, prefix)\n return ckpt_path, ckpt_tmp_path, base_path\n\n\ndef save_checkpoint(\n ckpt_dir: str | os.PathLike,\n target: PyTree,\n step: int | float,\n prefix: str = 'checkpoint_',\n keep: int = 1,\n overwrite: bool = False,\n keep_every_n_steps: int | None = None,\n async_manager: AsyncManager | None = None,\n orbax_checkpointer: ocp.Checkpointer | None = None,\n) -> str:\n \"\"\"Save a checkpoint of the model. Suitable for single-host.\n\n In this method, every JAX process saves the checkpoint on its own. Do not\n use it if you have multiple processes and you intend for them to save data\n to a common directory (e.g., a GCloud bucket). To save multi-process\n checkpoints to a shared storage or to save ``GlobalDeviceArray``s, use\n ``save_checkpoint_multiprocess()`` instead.\n\n Pre-emption safe by writing to temporary before a final rename and cleanup\n of past files. However, if async_manager is used, the final\n commit will happen inside an async callback, which can be explicitly waited\n by calling ``async_manager.wait_previous_save()``.\n\n Example usage::\n\n >>> from flax.training import checkpoints\n >>> import jax.numpy as jnp\n >>> import tempfile\n\n >>> with tempfile.TemporaryDirectory() as dir_path:\n ... test_object = {\n ... 'a': jnp.array([1, 2, 3], jnp.int32),\n ... 'b': jnp.array([1, 1, 1], jnp.int32),\n ... }\n ... file_path = checkpoints.save_checkpoint(\n ... dir_path, target=test_object, step=0, prefix='test_', keep=1\n ... )\n ... restored_object = checkpoints.restore_checkpoint(\n ... file_path, target=None\n ... )\n >>> restored_object\n {'a': Array([1, 2, 3], dtype=int32), 'b': Array([1, 1, 1], dtype=int32)}\n\n Args:\n ckpt_dir: str or pathlib-like path to store checkpoint files in.\n target: serializable flax object, usually a flax optimizer.\n step: int or float: training step number or other metric number.\n prefix: str: checkpoint file name prefix.\n keep: number of past checkpoint files to keep.\n overwrite: overwrite existing checkpoint files if a checkpoint at the\n current or a later step already exits (default: False).\n keep_every_n_steps: if defined, keep every checkpoints every n steps (in\n addition to keeping the last 'keep' checkpoints).\n async_manager: if defined, the save will run without blocking the main\n thread. Only works for single host. Note that an ongoing save will still\n block subsequent saves, to make sure overwrite/keep logic works correctly.\n orbax_checkpointer: if defined, the save will be done by ocp. In the future,\n all Flax checkpointing features will be migrated to Orbax, and starting to\n use an ``orbax_checkpointer`` is recommended. Please check out the\n checkpointing guide\n (https://flax.readthedocs.io/en/latest/guides/training_techniques/use_checkpointing.html#save-checkpoints)\n for how to use Orbax checkpointers.\n\n Returns:\n Filename of saved checkpoint.\n \"\"\"\n jax.monitoring.record_event('/jax/flax/checkpoint/save')\n start_time = time.time()\n # Make sure all saves are finished before the logic of checking and removing\n # outdated checkpoints happens.\n if async_manager:\n async_manager.wait_previous_save()\n\n ckpt_path, ckpt_tmp_path, base_path = _get_checkpoint_paths(\n ckpt_dir, step, prefix\n )\n\n if config.flax_use_orbax_checkpointing or orbax_checkpointer:\n logging.info(\n 'Using Orbax as backend to save Flax checkpoints. For potential'\n ' troubleshooting see:'\n ' https://flax.readthedocs.io/en/latest/guides/training_techniques/use_checkpointing.html#orbax-as-backend-troubleshooting'\n )\n if jax.process_count() > 1:\n logging.warning(\n 'Multiple JAX processes detected when calling single-process'\n ' `save_checkpoint`. Your devices will HANG if this function is only'\n ' called on process 0! Troubleshoot at:'\n ' https://flax.readthedocs.io/en/latest/guides/training_techniques/use_checkpointing.html#if-your-devices-hang-when-writing-checkpoints'\n )\n\n # Make sure any previous work is done before making file changes.\n if orbax_checkpointer and isinstance(\n orbax_checkpointer, ocp.AsyncCheckpointer\n ):\n orbax_checkpointer.wait_until_finished()\n # If no checkpointer provided, save synchronously with default setting.\n if not orbax_checkpointer:\n orbax_checkpointer = ocp.Checkpointer(\n ocp.PyTreeCheckpointHandler()\n )\n # Check singular target.\n if jtu.treedef_is_leaf(jtu.tree_structure(target)) and not isinstance(\n orbax_checkpointer._handler,\n ocp.ArrayCheckpointHandler, # pylint: disable=protected-access\n ):\n raise ValueError(\n 'Orbax backend only accept pytree as save target. To save singular'\n ' objects like numbers or Numpy arrays, checkout'\n ' https://flax.readthedocs.io/en/latest/guides/training_techniques/use_checkpointing.html#if-you-don-t-save-pytrees'\n )\n\n orbax_checkpointer.save(\n ckpt_path, target, force=overwrite\n )\n # Do a process check here in case people call this for multihost.\n if process_index() == 0:\n _remove_invalid_ckpts(\n ckpt_path, base_path, keep, overwrite, keep_every_n_steps, True\n )\n end_time = time.time()\n monitoring.record_event_duration_secs(\n _WRITE_CHECKPOINT_EVENT, end_time - start_time\n )\n return ckpt_path\n\n warnings.warn(\n (\n 'Flax Checkpointing will soon be deprecated in favor of Orbax'\n ' (https://github.com/google/orbax). Please refer to the Checkpoint'\n ' Upgrade Guide'\n ' (https://flax.readthedocs.io/en/latest/guides/converting_and_upgrading/orbax_upgrade_guide.html)'\n ' to self-migrate your code to ocp.'\n ),\n DeprecationWarning,\n )\n if not overwrite:\n _check_overwrite_error(ckpt_tmp_path, ckpt_path, base_path, step) # type: ignore\n\n target = serialization.to_bytes(target)\n\n # Save the files via I/O sync or async.\n def save_main_ckpt_task():\n jax.monitoring.record_event('/jax/flax/checkpoint/save_main_ckpt_task')\n return _save_main_ckpt_file(\n target,\n False,\n (ckpt_tmp_path, ckpt_path),\n base_path,\n step,\n keep,\n overwrite,\n keep_every_n_steps,\n start_time,\n )\n\n if async_manager:\n async_manager.save_async(save_main_ckpt_task)\n else:\n save_main_ckpt_task()\n end_time = time.time()\n monitoring.record_event_duration_secs(\n _WRITE_CHECKPOINT_EVENT, end_time - start_time\n )\n return ckpt_path\n\n\ndef save_checkpoint_multiprocess(\n ckpt_dir: str | os.PathLike,\n target: PyTree,\n step: int | float,\n prefix: str = 'checkpoint_',\n keep: int = 1,\n overwrite: bool = False,\n keep_every_n_steps: int | None = None,\n async_manager: AsyncManager | None = None,\n gda_manager: GlobalAsyncCheckpointManager | None = None,\n orbax_checkpointer: ocp.Checkpointer | None = None,\n) -> str:\n \"\"\"Save a checkpoint of the model in multi-process environment.\n\n Use this method to save ``GlobalDeviceArray``s, or to save data to a\n common directory. Only process 0 will save the main checkpoint file and\n remove old checkpoint files.\n\n Pre-emption safe by writing to temporary before a final rename and cleanup\n of past files. However, if async_manager or gda_manager is used, the final\n commit will happen inside an async callback, which can be explicitly waited\n by calling ``async_manager.wait_previous_save()`` or\n ``gda_manager.wait_until_finished()``.\n\n Args:\n ckpt_dir: str or pathlib-like path to store checkpoint files in.\n target: serializable flax object, usually a flax optimizer.\n step: int or float: training step number or other metric number.\n prefix: str: checkpoint file name prefix.\n keep: number of past checkpoint files to keep.\n overwrite: overwrite existing checkpoint files if a checkpoint at the\n current or a later step already exits (default: False).\n keep_every_n_steps: if defined, keep every checkpoints every n steps (in\n addition to keeping the last 'keep' checkpoints).\n async_manager: if defined, the save will run without blocking the main\n thread. Only works for single host. Note that an ongoing save will still\n block subsequent saves, to make sure overwrite/keep logic works correctly.\n gda_manager: required if target contains a JAX GlobalDeviceArray. Will save\n the GDAs to a separate subdirectory with postfix \"_gda\" asynchronously.\n Same as async_manager, this will block subsequent saves.\n orbax_checkpointer: if defined, the save will be done by Orbax. In the\n future, all Flax checkpointing features will be migrated to Orbax, and\n starting to use an ``orbax_checkpointer`` is recommended. Please check out\n the checkpointing guide\n (https://flax.readthedocs.io/en/latest/guides/training_techniques/use_checkpointing.html#save-checkpoints)\n for how to use Orbax checkpointers.\n\n Returns:\n Filename of saved checkpoint.\n \"\"\"\n jax.monitoring.record_event('/jax/flax/checkpoint/save')\n start_time = time.time()\n # Make sure all saves are finished before the logic of checking and removing\n # outdated checkpoints happens.\n sync_global_devices('Flax:Checkpoint:StartSave')\n if async_manager:\n async_manager.wait_previous_save()\n if gda_manager:\n gda_manager.wait_until_finished()\n sync_global_devices('Flax:Checkpoint:WaitLastSaveDone')\n\n ckpt_path, ckpt_tmp_path, base_path = _get_checkpoint_paths(\n ckpt_dir, step, prefix\n )\n\n if config.flax_use_orbax_checkpointing or orbax_checkpointer:\n logging.info(\n 'Using Orbax as backend to save Flax checkpoints. For potential'\n ' troubleshooting see:'\n ' https://flax.readthedocs.io/en/latest/guides/training_techniques/use_checkpointing.html#orbax-as-backend-troubleshooting'\n )\n # Make sure any previous work is done before making file changes.\n if orbax_checkpointer and isinstance(\n orbax_checkpointer, ocp.AsyncCheckpointer\n ):\n orbax_checkpointer.wait_until_finished()\n\n # If no checkpointer provided, save synchronously with default setting.\n if not orbax_checkpointer:\n orbax_checkpointer = ocp.Checkpointer(\n ocp.PyTreeCheckpointHandler()\n )\n # Check singular target.\n if jtu.treedef_is_leaf(jtu.tree_structure(target)) and not isinstance(\n orbax_checkpointer._handler,\n ocp.ArrayCheckpointHandler, # pylint: disable=protected-access\n ):\n raise ValueError(\n 'Orbax backend only accept pytree as save target. To save singular'\n ' objects like numbers or Numpy arrays, checkout'\n ' https://flax.readthedocs.io/en/latest/guides/training_techniques/use_checkpointing.html#if-you-don-t-save-pytrees'\n )\n\n if process_index() == 0:\n _remove_invalid_ckpts(\n ckpt_path, base_path, keep, overwrite, keep_every_n_steps, True\n )\n orbax_checkpointer.save(\n ckpt_path, target, force=overwrite\n )\n end_time = time.time()\n monitoring.record_event_duration_secs(\n _WRITE_CHECKPOINT_EVENT, end_time - start_time\n )\n return ckpt_path\n\n warnings.warn(\n (\n 'Flax Checkpointing will soon be deprecated in favor of Orbax'\n ' (https://github.com/google/orbax). Please refer to the Checkpoint'\n ' Upgrade Guide'\n ' (https://flax.readthedocs.io/en/latest/guides/converting_and_upgrading/orbax_upgrade_guide.html)'\n ' to self-migrate your code to ocp.'\n ),\n DeprecationWarning,\n )\n\n target = serialization.to_state_dict(target)\n target, mpa_targets = _split_mp_arrays(target)\n target = serialization.msgpack_serialize(target)\n has_mpa = bool(mpa_targets)\n\n if not overwrite:\n _check_overwrite_error(ckpt_tmp_path, ckpt_path, base_path, step) # type: ignore\n sync_global_devices('Flax:Checkpoint:CheckOverwriteBeforeSave')\n\n # Save the files via I/O sync or async.\n def save_main_ckpt_task():\n jax.monitoring.record_event('/jax/flax/checkpoint/save_main_ckpt_task')\n return _save_main_ckpt_file(\n target,\n has_mpa,\n (ckpt_tmp_path, ckpt_path),\n base_path,\n step,\n keep,\n overwrite,\n keep_every_n_steps,\n start_time,\n )\n\n # Write the main checkpoint file only via process 0, to avoid race condition.\n if process_index() == 0:\n if async_manager:\n async_manager.save_async(save_main_ckpt_task)\n else:\n save_main_ckpt_task()\n\n if has_mpa:\n if not gda_manager:\n raise errors.MPACheckpointingRequiredError(ckpt_path, step)\n # Creating the directory containing GDAs explicitly. This should happen only\n # on process 0 and before any worker starts to write GDA data.\n if process_index() == 0:\n _make_mpa_dirs(mpa_targets, ckpt_tmp_path)\n sync_global_devices('Flax:Checkpoint:AfterCreateMPADir')\n _save_mpas(\n gda_manager,\n mpa_targets,\n ckpt_tmp_path,\n ckpt_path,\n base_path,\n keep,\n overwrite,\n keep_every_n_steps,\n start_time,\n async_manager,\n )\n\n end_time = time.time()\n monitoring.record_event_duration_secs(\n _WRITE_CHECKPOINT_EVENT, end_time - start_time\n )\n return ckpt_path\n\n\ndef _all_checkpoints(\n ckpt_dir: str | os.PathLike, prefix: str = 'checkpoint_'\n) -> list[str]:\n \"\"\"Retrieve all checkpoint paths in directory.\n\n Args:\n ckpt_dir: str: directory of checkpoints to restore from.\n prefix: str: name prefix of checkpoint files.\n\n Returns:\n Sorted list of checkpoint paths or empty list if no checkpoints were found.\n \"\"\"\n ckpt_dir = os.fspath(ckpt_dir) # Pathlib -> str\n checkpoint_files: list[Any] = [\n pathlib.PurePath(c) for c in _allowempty_listdir(ckpt_dir)\n ]\n checkpoint_files = [\n os.path.join(ckpt_dir, c)\n for c in checkpoint_files\n if c.match(f'{prefix}*')\n and not c.match(f'{prefix}tmp')\n and not c.match(f'*{MP_ARRAY_POSTFIX}')\n and not c.match(f'*{ocp.utils.TMP_DIR_SUFFIX}*')\n ]\n checkpoint_files = natural_sort(checkpoint_files)\n if checkpoint_files:\n return checkpoint_files\n else:\n return []\n\n\ndef latest_checkpoint(\n ckpt_dir: str | os.PathLike, prefix: str = 'checkpoint_'\n) -> str | None:\n \"\"\"Retrieve the path of the latest checkpoint in a directory.\n\n Args:\n ckpt_dir: str: directory of checkpoints to restore from.\n prefix: str: name prefix of checkpoint files.\n\n Returns:\n The latest checkpoint path or None if no checkpoints were found.\n \"\"\"\n checkpoint_files = _all_checkpoints(ckpt_dir, prefix)\n if checkpoint_files:\n return checkpoint_files[-1]\n else:\n return None\n\n\ndef available_steps(\n ckpt_dir: str | os.PathLike,\n prefix: str = 'checkpoint_',\n step_type: type = int,\n) -> list[int | float]:\n \"\"\"Return step numbers of available checkpoints in a directory.\n\n\n Args:\n ckpt_dir: str: directory of checkpoints to restore from.\n prefix: str: name prefix of checkpoint files.\n step_type: type: type for steps, int (default) or float.\n\n Returns:\n Sorted list of available steps or empty list if no checkpoints were found.\n \"\"\"\n checkpoint_files = _all_checkpoints(ckpt_dir, prefix)\n\n checkpoint_steps = []\n\n for file in checkpoint_files:\n prefix_idx = file.rfind(prefix)\n checkpoint_steps += [step_type(file[prefix_idx + len(prefix) :])]\n\n return checkpoint_steps\n\n\ndef restore_checkpoint(\n ckpt_dir: str | os.PathLike,\n target: Any | None,\n step: int | float | None = None,\n prefix: str = 'checkpoint_',\n parallel: bool = True,\n gda_manager: GlobalAsyncCheckpointManager | None = None,\n allow_partial_mpa_restoration: bool = False,\n orbax_checkpointer: ocp.Checkpointer | None = None,\n orbax_transforms: dict | None = None,\n) -> PyTree:\n \"\"\"Restore last/best checkpoint from checkpoints in path.\n\n Sorts the checkpoint files naturally, returning the highest-valued\n file, e.g.:\n\n * ``ckpt_1, ckpt_2, ckpt_3 --> ckpt_3``\n\n * ``ckpt_0.01, ckpt_0.1, ckpt_0.001 --> ckpt_0.1``\n\n * ``ckpt_-1.0, ckpt_1.0, ckpt_1e5 --> ckpt_1e5``\n\n Example usage::\n\n >>> from flax.training import checkpoints\n >>> import jax.numpy as jnp\n >>> import tempfile\n ...\n >>> with tempfile.TemporaryDirectory() as dir_path:\n ... test_object = {\n ... 'a': jnp.array([1, 2, 3], jnp.int32),\n ... 'b': jnp.array([1, 1, 1], jnp.int32),\n ... }\n ... file_path = checkpoints.save_checkpoint(\n ... dir_path, target=test_object, step=0, prefix='test_', keep=1\n ... )\n ... restored_object = checkpoints.restore_checkpoint(\n ... file_path, target=None\n ... )\n >>> restored_object\n {'a': Array([1, 2, 3], dtype=int32), 'b': Array([1, 1, 1], dtype=int32)}\n\n Args:\n ckpt_dir: str: checkpoint file or directory of checkpoints to restore from.\n target: matching object to rebuild via deserialized state-dict. If None, the\n deserialized state-dict is returned as-is.\n step: int or float: step number to load or None to load latest. If\n specified, ckpt_dir must be a directory.\n prefix: str: name prefix of checkpoint files.\n parallel: bool: whether to load seekable checkpoints in parallel, for speed.\n gda_manager: required if checkpoint contains a multiprocess array\n (GlobalDeviceArray or jax Array from pjit). Will read the arrays from the\n separate subdirectory with postfix \"_gda\".\n allow_partial_mpa_restoration: If true, the given ``target`` doesn't have to\n contain all valid multiprocess arrays. As a result, the restored Pytree\n may have some MPAs not restored correctly. Use this if you cannot provide\n a fully valid ``target`` and don't need all the MPAs in the checkpoint to\n be restored.\n orbax_checkpointer: the ``ocp.Checkpointer`` that handles the underlying\n restore, if the given checkpoint is saved with ocp.\n orbax_transforms: the Orbax transformations that will be passed into\n ``orbax_checkpointer.restore()`` call.\n\n Returns:\n Restored ``target`` updated from checkpoint file, or if no step specified\n and\n no checkpoint files present, returns the passed-in ``target`` unchanged.\n If a file path is specified and is not found, the passed-in ``target`` will\n be\n returned. This is to match the behavior of the case where a directory path\n is specified but the directory has not yet been created.\n \"\"\"\n jax.monitoring.record_event('/jax/flax/checkpoint/restore')\n start_time = time.time()\n # Make sure any previous work is done before checking files.\n if orbax_checkpointer and isinstance(\n orbax_checkpointer, ocp.AsyncCheckpointer\n ):\n orbax_checkpointer.wait_until_finished()\n\n ckpt_dir = os.fspath(ckpt_dir) # Pathlib -> str\n ckpt_dir = safe_normpath(ckpt_dir)\n if step is not None:\n ckpt_path = _checkpoint_path(ckpt_dir, step, prefix)\n if not io.exists(ckpt_path):\n raise ValueError(f'Matching checkpoint not found: {ckpt_path}')\n else:\n if not io.exists(ckpt_dir):\n logging.info('Found no checkpoint directory at %s', ckpt_dir)\n return target\n if io.isdir(ckpt_dir):\n # This means the given dir is an orbax checkpoint.\n if _is_orbax_checkpoint(ckpt_dir):\n ckpt_path = ckpt_dir\n else:\n ckpt_path = latest_checkpoint(ckpt_dir, prefix) # type: ignore\n if not ckpt_path:\n logging.info(\n 'Found no checkpoint files in %s with prefix %s', ckpt_dir, prefix\n )\n return target\n else:\n ckpt_path = ckpt_dir\n\n # Restore the checkpoint with Orbax if needed.\n is_orbax = _is_orbax_checkpoint(ckpt_path)\n ckpt_type = 'orbax' if is_orbax else 'legacy Flax'\n logging.info(f'Restoring {ckpt_type} checkpoint from {ckpt_path}')\n if is_orbax:\n if not orbax_checkpointer:\n orbax_checkpointer = ocp.Checkpointer(\n ocp.PyTreeCheckpointHandler()\n )\n\n restore_kwargs = {}\n if target is not None:\n restore_kwargs['restore_args'] = orbax_utils.restore_args_from_target(\n target\n )\n if isinstance(orbax_checkpointer._handler, ocp.PyTreeCheckpointHandler): # pylint: disable=protected-access\n restore_kwargs[\n 'transforms'\n ] = orbax_utils.maybe_construct_transformations(\n target, orbax_transforms\n )\n restored = orbax_checkpointer.restore(\n ckpt_path, item=target, **restore_kwargs\n )\n restored = serialization.to_state_dict(restored)\n if target is not None:\n restored = serialization.from_state_dict(target, restored)\n end_time = time.time()\n monitoring.record_event_duration_secs(\n _READ_CHECKPOINT_EVENT, end_time - start_time\n )\n return restored\n\n # Legacy Flax checkpoint restoration.\n ckpt_size = io.getsize(ckpt_path)\n with io.GFile(ckpt_path, 'rb') as fp:\n if parallel and fp.seekable():\n buf_size = 128 << 20 # 128M buffer.\n num_bufs = ckpt_size / buf_size\n logging.debug('num_bufs: %d', num_bufs)\n checkpoint_contents = bytearray(ckpt_size)\n\n def read_chunk(i):\n # NOTE: We have to re-open the file to read each chunk, otherwise the\n # parallelism has no effect. But we could reuse the file pointers\n # within each thread.\n with io.GFile(ckpt_path, 'rb') as f:\n f.seek(i * buf_size)\n buf = f.read(buf_size)\n if buf:\n checkpoint_contents[i * buf_size : i * buf_size + len(buf)] = buf\n return len(buf) / buf_size\n\n pool_size = 32\n pool = thread.ThreadPoolExecutor(pool_size)\n results = pool.map(read_chunk, range(int(num_bufs) + 1))\n pool.shutdown(wait=False)\n logging.debug(f'results: {list(results)}')\n else:\n checkpoint_contents = fp.read()\n\n state_dict = serialization.msgpack_restore(checkpoint_contents)\n state_dict = _restore_mpas(\n state_dict,\n target,\n ckpt_path,\n step,\n gda_manager,\n allow_partial_mpa_restoration,\n )\n\n if target is None:\n restored_checkpoint = state_dict\n else:\n restored_checkpoint = serialization.from_state_dict(target, state_dict)\n\n end_time = time.time()\n monitoring.record_event_duration_secs(\n _READ_CHECKPOINT_EVENT, end_time - start_time\n )\n\n return restored_checkpoint\n\n\ndef convert_pre_linen(params: PyTree) -> PyTree:\n \"\"\"Converts a pre-Linen parameter pytree.\n\n In pre-Linen API submodules were numbered incrementally, independent of the\n submodule class. With Linen this behavior has changed to keep separate\n submodule counts per module class.\n\n Consider the following module::\n\n class Model(nn.Module):\n @nn.compact\n def __call__(self, x):\n x = nn.Conv(1, 1)(x)\n x = nn.Dense(1)(x)\n return x\n\n In pre-Linen the resulting params would have had the structure:\n\n ``{'Conv_0': { ... }, 'Dense_1': { ... } }``\n\n With Linen the resulting params would instead have had the structure:\n\n ``{'Conv_0': { ... }, 'Dense_0': { ... } }``\n\n To convert from pre-Linen format to Linen simply call::\n\n params = convert_pre_linen(pre_linen_params)\n\n Note that you can also use this utility to convert pre-Linen collections\n because they're following the same module naming. Note though that collections\n were \"flat\" in pre-Linen and first need to be unflattened before they can be\n used with this function::\n\n batch_stats = convert_pre_linen(flax.traverse_util.unflatten_dict({\n tuple(k.split('/')[1:]): v\n for k, v in pre_linen_model_state.as_dict().items()\n }))\n\n Then Linen variables can be defined from these converted collections::\n\n variables = {'params': params, 'batch_stats': batch_stats}\n\n Args:\n params: Parameter pytree in pre-Linen format. If the pytree is already in\n Linen format, then the returned pytree is unchanged (i.e. this function\n can safely be called on any loaded checkpoint for use with Linen).\n\n Returns:\n Parameter pytree with Linen submodule naming.\n \"\"\"\n if not isinstance(params, (dict, core.FrozenDict)):\n return params\n params_renamed = {}\n counts: dict[Any, Any] = {}\n names = natural_sort(params.keys())\n for name in names:\n value = params[name]\n match = MODULE_NUM_RE.match(name)\n if match:\n module = match.group(1)\n num = counts.get(module, 0)\n name = f'{module}_{num}'\n counts[module] = num + 1\n params_renamed[name] = convert_pre_linen(value)\n\n if isinstance(params, core.FrozenDict):\n params_renamed = core.freeze(params_renamed) # type: ignore\n return params_renamed\n" }, { "path": "flax/training/common_utils.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Common utility functions used in data-parallel Flax examples.\n\nThis module is a historical grab-bag of utility functions primarily concerned\nwith helping write pmap-based data-parallel training loops.\n\"\"\"\n\nimport jax\nimport jax.numpy as jnp\nfrom jax import lax\n\n\ndef shard(xs):\n \"\"\"Helper for pmap to shard a pytree of arrays by local_device_count.\n\n Args:\n xs: a pytree of arrays.\n Returns:\n A matching pytree with arrays' leading dimensions sharded by the\n local device count.\n \"\"\"\n local_device_count = jax.local_device_count()\n return jax.tree_util.tree_map(\n lambda x: x.reshape((local_device_count, -1) + x.shape[1:]), xs\n )\n\n\ndef shard_prng_key(prng_key):\n \"\"\"Helper to shard (aka split) a PRNGKey for use with pmap'd functions.\n\n PRNG keys can be used at train time to drive stochastic modules\n e.g. Dropout. We would like a different PRNG key for each local\n device so that we end up with different random numbers on each one,\n hence we split our PRNG key.\n\n Args:\n prng_key: JAX PRNGKey\n Returns:\n A new array of PRNGKeys with leading dimension equal to local device count.\n \"\"\"\n return jax.random.split(prng_key, num=jax.local_device_count())\n\n\ndef stack_forest(forest):\n \"\"\"Helper function to stack the leaves of a sequence of pytrees.\n\n Args:\n forest: a sequence of pytrees (e.g tuple or list) of matching structure\n whose leaves are arrays with individually matching shapes.\n Returns:\n A single pytree of the same structure whose leaves are individually\n stacked arrays.\n \"\"\"\n stack_args = lambda *args: jnp.stack(args)\n return jax.tree_util.tree_map(stack_args, *forest)\n\n\ndef get_metrics(device_metrics):\n \"\"\"Helper utility for pmap, gathering replicated timeseries metric data.\n\n Args:\n device_metrics: replicated, device-resident pytree of metric data,\n whose leaves are presumed to be a sequence of arrays recorded over time.\n Returns:\n A pytree of unreplicated, host-resident, stacked-over-time arrays useful for\n computing host-local statistics and logging.\n \"\"\"\n # We select the first element of x in order to get a single copy of a\n # device-replicated metric.\n device_metrics = jax.tree_util.tree_map(\n lambda x: x.addressable_shards[0].data.squeeze(0), device_metrics\n )\n metrics_np = jax.device_get(device_metrics)\n return stack_forest(metrics_np)\n\n\ndef onehot(labels, num_classes, on_value=1.0, off_value=0.0):\n \"\"\"Create a dense one-hot version of an indexed array.\n\n NB: consider using the more standard ``jax.nn.one_hot`` instead.\n\n Args:\n labels: an n-dim JAX array whose last dimension contains integer indices.\n num_classes: the maximum possible index.\n on_value: the \"on\" value for the one-hot array, defaults to 1.0.\n off_value: the \"off\" value for the one-hot array, defaults to 0.0.\n Returns:\n A (n+1)-dim array whose last dimension contains one-hot vectors of length\n num_classes.\n \"\"\"\n x = labels[..., None] == jnp.arange(num_classes).reshape(\n (1,) * labels.ndim + (-1,)\n )\n x = lax.select(x, jnp.full(x.shape, on_value), jnp.full(x.shape, off_value))\n return x.astype(jnp.float32)\n" }, { "path": "flax/training/dynamic_scale.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Dynamic loss scaling for mixed precision gradients.\"\"\"\n\nimport functools\nfrom typing import Any, NamedTuple\nfrom collections.abc import Callable, Sequence\n\nimport jax\nimport jax.numpy as jnp\nfrom jax import lax\n\nfrom flax import struct\nfrom flax.typing import Array\n\n\nclass DynamicScaleResult(NamedTuple):\n dynamic_scale: 'DynamicScale'\n finite: Array\n aux: Any\n grad: Any\n\n\nclass DynamicScale(struct.PyTreeNode):\n \"\"\"Dynamic loss scaling for mixed precision gradients.\n\n For many models gradient computations in float16 will result in numerical\n issues because small/large gradients being flushed to zero/infinity.\n Dynamic loss scaling is an algorithm that aims to find the largest scalar\n multiple for which the gradient does not overflow. This way the risk of\n underflow is minimized.\n\n the `value_and_grad` method mimicks `jax.value_and_grad`. Beside the loss\n and gradients it also ouputs and updated `DynamicScale` instance with the\n current loss scale factor. This method also returns a boolean value indicating\n whether the gradients are finite.\n\n Example::\n\n from flax.training.dynamic_scale import DynamicScale\n\n def loss_fn(p):\n return jnp.asarray(p, jnp.float16) ** 2\n p = jnp.array(1., jnp.float32)\n\n dyn_scale = DynamicScale(growth_interval=10)\n compute_grad = jax.jit(lambda ds, p: ds.value_and_grad(loss_fn)(p))\n for _ in range(100):\n dyn_scale, is_fin, loss, grad = compute_grad(dyn_scale, p)\n p += jnp.where(is_fin, 0.01 * grad, 0.)\n print(loss)\n\n Jax currently cannot execute conditionals efficiently on GPUs therefore we\n selectively ignore the gradient update using `jax.numpy.where` in case of\n non-finite gradients.\n\n Attributes:\n growth_factor: how much to grow the scalar after a period of finite\n gradients (default: 2.).\n backoff_factor: how much to shrink the scalar after a non-finite gradient\n (default: 0.5).\n growth_interval: after how many steps of finite gradients the scale should\n be increased (default: 2000).\n fin_steps: indicates how many gradient steps in a row have been finite.\n scale: the current scale by which the loss is multiplied.\n minimum_scale: the minimum value that the scale can take (default: the\n smallest positive number representable in floating point).\n \"\"\"\n\n growth_factor: float = struct.field(pytree_node=False, default=2.0)\n backoff_factor: float = struct.field(pytree_node=False, default=0.5)\n growth_interval: int = struct.field(pytree_node=False, default=2000)\n fin_steps: int = 0\n scale: float = 65536.0\n minimum_scale: float | None = struct.field(\n pytree_node=False, default=jnp.finfo(jnp.float32).tiny\n )\n\n def value_and_grad(\n self,\n fun: Callable[..., Any],\n argnums: int | Sequence[int] = 0,\n has_aux: bool = False,\n axis_name: str | None = None,\n ) -> Callable[..., DynamicScaleResult]:\n \"\"\"Wrapper around `jax.value_and_grad`.\n\n Args:\n fun: Function to be differentiated. Its arguments at positions specified\n by ``argnums`` should be arrays, scalars, or standard Python containers.\n It should return a scalar (which includes arrays with shape ``()`` but\n not arrays with shape ``(1,)`` etc.)\n argnums: Optional, integer or sequence of integers. Specifies which\n positional argument(s) to differentiate with respect to (default 0).\n has_aux: Optional, bool. Indicates whether ``fun`` returns a pair where\n the first element is considered the output of the mathematical function\n to be differentiated and the second element is auxiliary data. Default\n False.\n axis_name: If an axis is given the gradients will be averaged across\n replicas (default: None). Note, this is only used for pmap and shard\n map. For SPMD jit, you do not need to manually synchronize. Just make\n sure that the axes are correctly annotated and XLA:SPMD will insert the\n necessary collectives.\n\n Returns:\n A function that takes the same arguments as `fun` and\n returns a DynamicScaleResult\n \"\"\"\n\n @functools.wraps(fun)\n def loss_wrapper(*args):\n aux = fun(*args)\n if has_aux:\n return (self.scale * aux[0], aux[1])\n else:\n return self.scale * aux\n\n grad_fn = jax.value_and_grad(loss_wrapper, argnums, has_aux)\n\n def grad_fn_wrapper(*args):\n aux, grad = grad_fn(*args)\n aux = (aux[0] / self.scale, aux[1]) if has_aux else aux / self.scale\n\n grad = jax.tree_util.tree_map(\n lambda g: jnp.asarray(g, jnp.float32) / self.scale, grad\n )\n if axis_name is not None:\n grad = lax.pmean(grad, axis_name)\n\n finite = jnp.array(True)\n for g in jax.tree_util.tree_leaves(grad):\n finite &= jnp.all(lax.is_finite(g))\n\n grow = self.fin_steps == self.growth_interval\n fin_scale = jnp.where(\n grow & finite,\n jnp.minimum(\n self.scale * self.growth_factor, jnp.finfo(jnp.float32).max\n ),\n self.scale,\n )\n inf_scale = self.scale * self.backoff_factor\n if self.minimum_scale is not None:\n inf_scale = jnp.maximum(inf_scale, self.minimum_scale)\n new_scale = jnp.where(finite, fin_scale, inf_scale)\n new_fin_steps = jnp.where(grow | (~finite), 0, self.fin_steps + 1)\n\n new_self = self.replace(fin_steps=new_fin_steps, scale=new_scale)\n return DynamicScaleResult(new_self, finite, aux, grad)\n\n return grad_fn_wrapper\n" }, { "path": "flax/training/early_stopping.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Early stopping.\"\"\"\n\nimport math\n\nfrom flax import struct\n\n\nclass EarlyStopping(struct.PyTreeNode):\n \"\"\"Early stopping to avoid overfitting during training.\n\n The following example stops training early if the difference between losses\n recorded in the current epoch and previous epoch is less than 1e-3\n consecutively for 2 times::\n\n >>> from flax.training.early_stopping import EarlyStopping\n\n >>> def train_epoch(optimizer, train_ds, batch_size, epoch, input_rng):\n ... ...\n ... loss = [4, 3, 3, 3, 2, 2, 2, 2, 1, 1][epoch]\n ... return None, {'loss': loss}\n\n >>> early_stop = EarlyStopping(min_delta=1e-3, patience=2)\n >>> optimizer = None\n >>> for epoch in range(10):\n ... optimizer, train_metrics = train_epoch(\n ... optimizer=optimizer, train_ds=None, batch_size=None, epoch=epoch, input_rng=None)\n ... early_stop = early_stop.update(train_metrics['loss'])\n ... if early_stop.should_stop:\n ... print(f'Met early stopping criteria, breaking at epoch {epoch}')\n ... break\n Met early stopping criteria, breaking at epoch 7\n\n Attributes:\n min_delta: Minimum delta between updates to be considered an\n improvement.\n patience: Number of steps of no improvement before stopping.\n best_metric: Current best metric value.\n patience_count: Number of steps since last improving update.\n should_stop: Whether the training loop should stop to avoid\n overfitting.\n has_improved: Whether the metric has improved greater or\n equal to the min_delta in the last ``.update`` call.\n \"\"\"\n\n min_delta: float = 0\n patience: int = 0\n best_metric: float = float('inf')\n patience_count: int = 0\n should_stop: bool = False\n has_improved: bool = False\n\n def reset(self):\n return self.replace(\n best_metric=float('inf'),\n patience_count=0,\n should_stop=False,\n has_improved=False,\n )\n\n def update(self, metric):\n \"\"\"Update the state based on metric.\n\n Returns:\n The updated EarlyStopping class. The ``.has_improved`` attribute is True\n when there was an improvement greater than ``min_delta`` from the previous\n ``best_metric``.\n \"\"\"\n\n if (\n math.isinf(self.best_metric) or self.best_metric - metric > self.min_delta\n ):\n return self.replace(\n best_metric=metric, patience_count=0, has_improved=True\n )\n else:\n should_stop = self.patience_count >= self.patience or self.should_stop\n return self.replace(\n patience_count=self.patience_count + 1,\n should_stop=should_stop,\n has_improved=False,\n )\n" }, { "path": "flax/training/lr_schedule.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Learning rate schedules used in FLAX image classification examples.\n\nNote that with `FLIP #1009`_ learning rate schedules in ``flax.training`` are\n**effectively deprecated** in favor of Optax_ schedules. Please refer to\n`Optimizer Schedules`_ for more information.\n\n.. _FLIP #1009: https://github.com/google/flax/blob/main/docs/flip/1009-optimizer-api.md\n.. _Optax: https://github.com/deepmind/optax\n.. _Optimizer Schedules: https://optax.readthedocs.io/en/latest/api.html#optimizer-schedules\n\"\"\"\n\nimport jax.numpy as jnp\nimport numpy as np\nfrom absl import logging\n\n\ndef _piecewise_constant(boundaries, values, t):\n index = jnp.sum(boundaries < t)\n return jnp.take(values, index)\n\n\ndef create_constant_learning_rate_schedule(\n base_learning_rate, steps_per_epoch, warmup_length=0.0\n):\n \"\"\"Create a constant learning rate schedule with optional warmup.\n\n Note that with `FLIP #1009`_ learning rate schedules in ``flax.training`` are\n **effectively deprecated** in favor of Optax_ schedules. Please refer to\n `Optimizer Schedules`_ for more information.\n\n .. _FLIP #1009: https://github.com/google/flax/blob/main/docs/flip/1009-optimizer-api.md\n .. _Optax: https://github.com/deepmind/optax\n .. _Optimizer Schedules: https://optax.readthedocs.io/en/latest/api.html#optimizer-schedules\n\n Holds the learning rate constant. This function also offers a learing rate\n warmup as per https://arxiv.org/abs/1706.02677, for the purpose of training\n with large mini-batches.\n\n Args:\n base_learning_rate: the base learning rate\n steps_per_epoch: the number of iterations per epoch\n warmup_length: if > 0, the learning rate will be modulated by a warmup\n factor that will linearly ramp-up from 0 to 1 over the first\n ``warmup_length`` epochs\n\n Returns:\n Function ``f(step) -> lr`` that computes the learning rate for a given step.\n \"\"\"\n logging.warning(\n 'Learning rate schedules in ``flax.training`` are effectively deprecated '\n 'in favor of Optax schedules. Please refer to '\n 'https://optax.readthedocs.io/en/latest/api.html#optimizer-schedules'\n ' for alternatives.'\n )\n\n def learning_rate_fn(step):\n lr = base_learning_rate\n if warmup_length > 0.0:\n lr = lr * jnp.minimum(1.0, step / float(warmup_length) / steps_per_epoch)\n return lr\n\n return learning_rate_fn\n\n\ndef create_stepped_learning_rate_schedule(\n base_learning_rate, steps_per_epoch, lr_sched_steps, warmup_length=0.0\n):\n \"\"\"Create a stepped learning rate schedule with optional warmup.\n\n Note that with `FLIP #1009`_ learning rate schedules in ``flax.training`` are\n **effectively deprecated** in favor of Optax_ schedules. Please refer to\n `Optimizer Schedules`_ for more information.\n\n .. _FLIP #1009: https://github.com/google/flax/blob/main/docs/flip/1009-optimizer-api.md\n .. _Optax: https://github.com/deepmind/optax\n .. _Optimizer Schedules: https://optax.readthedocs.io/en/latest/api.html#optimizer-schedules\n\n A stepped learning rate schedule decreases the learning rate\n by specified amounts at specified epochs. The steps are given as\n the ``lr_sched_steps`` parameter. A common ImageNet schedule decays the\n learning rate by a factor of 0.1 at epochs 30, 60 and 80. This would be\n specified as::\n\n [\n [30, 0.1],\n [60, 0.01],\n [80, 0.001]\n ]\n\n This function also offers a learing rate warmup as per\n https://arxiv.org/abs/1706.02677, for the purpose of training with large\n mini-batches.\n\n Args:\n base_learning_rate: the base learning rate\n steps_per_epoch: the number of iterations per epoch\n lr_sched_steps: the schedule as a list of steps, each of which is\n a ``[epoch, lr_factor]`` pair; the step occurs at epoch ``epoch`` and\n sets the learning rate to ``base_learning_rage * lr_factor``\n warmup_length: if > 0, the learning rate will be modulated by a warmup\n factor that will linearly ramp-up from 0 to 1 over the first\n ``warmup_length`` epochs\n\n Returns:\n Function ``f(step) -> lr`` that computes the learning rate for a given step.\n \"\"\"\n logging.warning(\n 'Learning rate schedules in ``flax.training`` are effectively deprecated '\n 'in favor of Optax schedules. Please refer to '\n 'https://optax.readthedocs.io/en/latest/api.html#optimizer-schedules'\n ' for alternatives.'\n )\n boundaries = [step[0] for step in lr_sched_steps]\n decays = [step[1] for step in lr_sched_steps]\n boundaries = np.array(boundaries) * steps_per_epoch\n boundaries = np.round(boundaries).astype(int)\n values = np.array([1.0] + decays) * base_learning_rate\n\n def learning_rate_fn(step):\n lr = _piecewise_constant(boundaries, values, step)\n if warmup_length > 0.0:\n lr = lr * jnp.minimum(1.0, step / float(warmup_length) / steps_per_epoch)\n return lr\n\n return learning_rate_fn\n\n\ndef create_cosine_learning_rate_schedule(\n base_learning_rate, steps_per_epoch, halfcos_epochs, warmup_length=0.0\n):\n \"\"\"Create a cosine learning rate schedule with optional warmup.\n\n Note that with `FLIP #1009`_ learning rate schedules in ``flax.training`` are\n **effectively deprecated** in favor of Optax_ schedules. Please refer to\n `Optimizer Schedules`_ for more information.\n\n .. _FLIP #1009: https://github.com/google/flax/blob/main/docs/flip/1009-optimizer-api.md\n .. _Optax: https://github.com/deepmind/optax\n .. _Optimizer Schedules: https://optax.readthedocs.io/en/latest/api.html#optimizer-schedules\n\n A cosine learning rate schedule modules the learning rate with\n half a cosine wave, gradually scaling it to 0 at the end of training.\n\n This function also offers a learing rate warmup as per\n https://arxiv.org/abs/1706.02677, for the purpose of training with large\n mini-batches.\n\n Args:\n base_learning_rate: the base learning rate\n steps_per_epoch: the number of iterations per epoch\n halfcos_epochs: the number of epochs to complete half a cosine wave;\n normally the number of epochs used for training\n warmup_length: if > 0, the learning rate will be modulated by a warmup\n factor that will linearly ramp-up from 0 to 1 over the first\n ``warmup_length`` epochs\n\n Returns:\n Function ``f(step) -> lr`` that computes the learning rate for a given step.\n \"\"\"\n logging.warning(\n 'Learning rate schedules in ``flax.training`` are effectively deprecated '\n 'in favor of Optax schedules. Please refer to '\n 'https://optax.readthedocs.io/en/latest/api.html#optimizer-schedules'\n ' for alternatives.'\n )\n halfwavelength_steps = halfcos_epochs * steps_per_epoch\n\n def learning_rate_fn(step):\n scale_factor = jnp.cos(step * jnp.pi / halfwavelength_steps) * 0.5 + 0.5\n lr = base_learning_rate * scale_factor\n if warmup_length > 0.0:\n lr = lr * jnp.minimum(1.0, step / float(warmup_length) / steps_per_epoch)\n return lr\n\n return learning_rate_fn\n" }, { "path": "flax/training/orbax_utils.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Utils for Orbax Checkpointing, available even after Flax Checkpointing is deprecated.\"\"\"\n\nimport warnings\nfrom typing import Any\n\nimport jax\nimport numpy as np\nimport orbax.checkpoint as ocp\nfrom jax.sharding import Mesh\n\nPyTree = Any\n\n\ndef is_multi_device_array(value: Any) -> bool:\n \"\"\"Instruct Orbax to save this array with Tensorstore instead of msgpack.\"\"\"\n if isinstance(value, jax.Array):\n return not value.is_fully_replicated\n return False\n\n\ndef save_args_from_target(target: Any) -> Any:\n return jax.tree_util.tree_map(lambda _: ocp.SaveArgs(), target)\n\n\ndef maybe_construct_transformations(\n target: Any, transforms: Any | None\n) -> Any:\n if transforms is not None:\n return transforms\n flat_transforms = {}\n flat_target = ocp.utils.to_flat_dict(target, sep='/', keep_empty_nodes=True)\n for k, v in flat_target.items():\n if v is None:\n flat_transforms[k] = ocp.Transform(use_fallback=True)\n return flat_transforms\n\n\ndef restore_args_from_target(target: Any, mesh: Mesh | None = None) -> Any:\n \"\"\"Creates Orbax `restore_args` given a target Pytree.\n\n Args:\n target: The Pytree that has the same structure as the checkpoint. The arrays\n restored from checkpoint will have the same `sharding` as the target\n Pytree's corresponding arrays.\n mesh: DEPRECATED ARG. Please simply use your mesh to create the arrays\n in your `target`, no need to pass it here.\n\n Returns:\n A Pytree of Orbax `RestoreArgs` or `ArrayRestoreArgs`\n \"\"\"\n\n def find_sharding(x):\n if hasattr(x, 'sharding'):\n return x.sharding\n return None\n\n # Simpler case: no JAX arrays\n if not any(\n jax.tree_util.tree_flatten(jax.tree_util.tree_map(find_sharding, target))[\n 0\n ]\n ):\n return jax.tree_util.tree_map(\n lambda x: ocp.RestoreArgs(restore_type=np.ndarray), target\n )\n\n # JAX arrays: find sharding from the given target and create RestoreArgs\n sharding_tree = jax.tree_util.tree_map(find_sharding, target)\n if mesh is not None:\n warnings.warn(\n (\n 'restore_args_from_target(): `mesh` arg is deprecated. Simply'\n ' calling the function with target pytree should suffice.'\n ),\n DeprecationWarning,\n )\n def substitute_embedding(s):\n return jax.sharding.NamedSharding(mesh, s.spec)\n sharding_tree = jax.tree_util.tree_map(substitute_embedding, sharding_tree)\n restore_args = ocp.checkpoint_utils.construct_restore_args(\n target, sharding_tree, set_global_shape=False\n )\n return restore_args\n" }, { "path": "flax/training/prefetch_iterator.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Utility for constructing an iterator which prefetches data asynchronously.\"\"\"\n\nimport threading\nimport warnings\n\n\nclass PrefetchIterator:\n \"\"\"Wraps an iterator to provide async prefetching.\n\n DEPRECATION WARNING:\n TensorFlow datasets no longer require manual prefetching.\n\n Previously this class was used to make data loading using TensorFlow datasets\n more efficient. Now TF data handles prefetching with NumPy iterators\n correctly.\n\n Example::\n\n tf_iter = dataset.as_numpy_iterator() # only loads data while calling next\n tf_iter = PrefetchIterator(tf_iter) # prefetches data in the background\n\n \"\"\"\n\n def __init__(self, data_iter, buffer_size=1):\n \"\"\"Construct a PrefetchIterator.\n\n Args:\n data_iter: the Iterator that should be prefetched.\n buffer_size: how many items to prefetch (default: 1).\n \"\"\"\n warnings.warn(\n 'PrefetchIterator is deprecated. Use the standard `tf.data`'\n ' prefetch method instead',\n DeprecationWarning,\n )\n\n self._data_iter = data_iter\n self.buffer_size = buffer_size\n self._cond = threading.Condition()\n self._buffer = []\n self._active = True\n self._thread = threading.Thread(target=self._prefetch_loop, daemon=True)\n self._thread.start()\n self._error = None\n\n def __iter__(self):\n return self\n\n def __next__(self):\n with self._cond:\n self._cond.wait_for(lambda: self._buffer or not self._active)\n if self._buffer:\n item = self._buffer.pop(0)\n self._cond.notify_all()\n return item\n if self._error:\n raise self._error # pylint: disable=raising-bad-type\n assert not self._active\n raise StopIteration()\n\n def close(self):\n with self._cond:\n self._active = False\n self._cond.notify_all()\n\n def _prefetch_loop(self):\n \"\"\"Prefetch loop that prefetches a tf dataset.\"\"\"\n\n def _predicate():\n return len(self._buffer) < self.buffer_size or not self._active\n\n while True:\n try:\n item = next(self._data_iter)\n with self._cond:\n self._buffer.append(item)\n self._cond.notify_all()\n self._cond.wait_for(_predicate)\n if not self._active:\n return\n except Exception as e: # pylint: disable=broad-except\n with self._cond:\n self._error = e\n self._active = False\n self._cond.notify_all()\n return\n" }, { "path": "flax/training/train_state.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom typing import Any\nfrom collections.abc import Callable\n\nimport optax\n\nimport jax\nfrom flax import core, struct\nfrom flax.linen.fp8_ops import OVERWRITE_WITH_GRADIENT\n\n\nclass TrainState(struct.PyTreeNode):\n \"\"\"Simple train state for the common case with a single Optax optimizer.\n\n Example usage::\n\n >>> import flax.linen as nn\n >>> from flax.training.train_state import TrainState\n >>> import jax, jax.numpy as jnp\n >>> import optax\n\n >>> x = jnp.ones((1, 2))\n >>> y = jnp.ones((1, 2))\n >>> model = nn.Dense(2)\n >>> variables = model.init(jax.random.key(0), x)\n >>> tx = optax.adam(1e-3)\n\n >>> state = TrainState.create(\n ... apply_fn=model.apply,\n ... params=variables['params'],\n ... tx=tx)\n\n >>> def loss_fn(params, x, y):\n ... predictions = state.apply_fn({'params': params}, x)\n ... loss = optax.l2_loss(predictions=predictions, targets=y).mean()\n ... return loss\n >>> loss_fn(state.params, x, y)\n Array(1.8136346, dtype=float32)\n\n >>> grads = jax.grad(loss_fn)(state.params, x, y)\n >>> state = state.apply_gradients(grads=grads)\n >>> loss_fn(state.params, x, y)\n Array(1.8079796, dtype=float32)\n\n Note that you can easily extend this dataclass by subclassing it for storing\n additional data (e.g. additional variable collections).\n\n For more exotic usecases (e.g. multiple optimizers) it's probably best to\n fork the class and modify it.\n\n Args:\n step: Counter starts at 0 and is incremented by every call to\n ``.apply_gradients()``.\n apply_fn: Usually set to ``model.apply()``. Kept in this dataclass for\n convenience to have a shorter params list for the ``train_step()`` function\n in your training loop.\n params: The parameters to be updated by ``tx`` and used by ``apply_fn``.\n tx: An Optax gradient transformation.\n opt_state: The state for ``tx``.\n \"\"\"\n\n step: int | jax.Array\n apply_fn: Callable = struct.field(pytree_node=False)\n params: core.FrozenDict[str, Any] = struct.field(pytree_node=True)\n tx: optax.GradientTransformation = struct.field(pytree_node=False)\n opt_state: optax.OptState = struct.field(pytree_node=True)\n\n def apply_gradients(self, *, grads, **kwargs):\n \"\"\"Updates ``step``, ``params``, ``opt_state`` and ``**kwargs`` in return value.\n\n Note that internally this function calls ``.tx.update()`` followed by a call\n to ``optax.apply_updates()`` to update ``params`` and ``opt_state``.\n\n Args:\n grads: Gradients that have the same pytree structure as ``.params``.\n **kwargs: Additional dataclass attributes that should be ``.replace()``-ed.\n\n Returns:\n An updated instance of ``self`` with ``step`` incremented by one, ``params``\n and ``opt_state`` updated by applying ``grads``, and additional attributes\n replaced as specified by ``kwargs``.\n \"\"\"\n if OVERWRITE_WITH_GRADIENT in grads:\n grads_with_opt = grads['params']\n params_with_opt = self.params['params']\n else:\n grads_with_opt = grads\n params_with_opt = self.params\n\n updates, new_opt_state = self.tx.update(\n grads_with_opt, self.opt_state, params_with_opt\n )\n new_params_with_opt = optax.apply_updates(params_with_opt, updates)\n\n # As implied by the OWG name, the gradients are used directly to update the\n # parameters.\n if OVERWRITE_WITH_GRADIENT in grads:\n new_params = {\n 'params': new_params_with_opt,\n OVERWRITE_WITH_GRADIENT: grads[OVERWRITE_WITH_GRADIENT],\n }\n else:\n new_params = new_params_with_opt\n return self.replace(\n step=self.step + 1,\n params=new_params,\n opt_state=new_opt_state,\n **kwargs,\n )\n\n @classmethod\n def create(cls, *, apply_fn, params, tx, **kwargs):\n \"\"\"Creates a new instance with ``step=0`` and initialized ``opt_state``.\"\"\"\n # We exclude OWG params when present because they do not need opt states.\n params_with_opt = (\n params['params'] if OVERWRITE_WITH_GRADIENT in params else params\n )\n opt_state = tx.init(params_with_opt)\n return cls(\n step=0,\n apply_fn=apply_fn,\n params=params,\n tx=tx,\n opt_state=opt_state,\n **kwargs,\n )\n" }, { "path": "flax/traverse_util.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"A utility for traversing immutable datastructures.\n\nA Traversal can be used to iterate and update complex data structures.\nTraversals take in an object and return a subset of its contents.\nFor example, a Traversal could select an attribute of an object::\n\n >>> from flax import traverse_util\n >>> import dataclasses\n\n >>> @dataclasses.dataclass\n ... class Foo:\n ... foo: int = 0\n ... bar: int = 0\n ...\n >>> x = Foo(foo=1)\n >>> iterator = traverse_util.TraverseAttr('foo').iterate(x)\n >>> list(iterator)\n [1]\n\nMore complex traversals can be constructed using composition.\nIt is often useful to start from the identity traversal and use a method chain\nto construct the intended Traversal::\n\n >>> data = [{'foo': 1, 'bar': 2}, {'foo': 3, 'bar': 4}]\n >>> traversal = traverse_util.t_identity.each()['foo']\n >>> iterator = traversal.iterate(data)\n >>> list(iterator)\n [1, 3]\n\nTraversals can also be used to make changes using the ``update`` method::\n\n >>> data = {'foo': Foo(bar=2)}\n >>> traversal = traverse_util.t_identity['foo'].bar\n >>> data = traversal.update(lambda x: x + x, data)\n >>> data\n {'foo': Foo(foo=0, bar=4)}\n\nTraversals never mutate the original data. Therefore, an update essentially\nreturns a copy of the data including the provided updates.\n\"\"\"\n\nimport abc\nimport copy\nimport dataclasses\nimport warnings\nfrom typing import Any\nfrom collections.abc import Callable\n\nimport jax\n\nimport flax\nfrom flax.core.scope import VariableDict\nfrom flax.typing import PathParts\n\nfrom . import struct\n\n\n# the empty node is a struct.dataclass to be compatible with JAX.\n@struct.dataclass\nclass _EmptyNode:\n pass\n\n\nempty_node = _EmptyNode()\n\n\ndef _flatten(xs, prefix, keep_empty_nodes, is_leaf, sep):\n def _key(path):\n if sep is None:\n return path\n return sep.join(path)\n\n if not isinstance(xs, (flax.core.FrozenDict, dict)) or (\n is_leaf and is_leaf(prefix, xs)\n ):\n return {_key(prefix): xs}\n result = {}\n is_empty = True\n for key, value in xs.items():\n is_empty = False\n path = prefix + (key,)\n result.update(_flatten(value, path, keep_empty_nodes, is_leaf, sep))\n if keep_empty_nodes and is_empty:\n if prefix == (): # when the whole input is empty\n return {}\n return {_key(prefix): empty_node}\n return result\n\n\ndef flatten_dict(xs, keep_empty_nodes=False, is_leaf=None, sep=None):\n \"\"\"Flatten a nested dictionary.\n\n The nested keys are flattened to a tuple.\n See ``unflatten_dict`` on how to restore the\n nested dictionary structure.\n\n Example::\n\n >>> from flax.traverse_util import flatten_dict\n\n >>> xs = {'foo': 1, 'bar': {'a': 2, 'b': {}}}\n >>> flat_xs = flatten_dict(xs)\n >>> flat_xs\n {('foo',): 1, ('bar', 'a'): 2}\n\n Note that empty dictionaries are ignored and\n will not be restored by ``unflatten_dict``.\n\n Args:\n xs: a nested dictionary\n keep_empty_nodes: replaces empty dictionaries\n with ``traverse_util.empty_node``.\n is_leaf: an optional function that takes the\n next nested dictionary and nested keys and\n returns True if the nested dictionary is a\n leaf (i.e., should not be flattened further).\n sep: if specified, then the keys of the returned\n dictionary will be ``sep``-joined strings (if\n ``None``, then keys will be tuples).\n Returns:\n The flattened dictionary.\n \"\"\"\n assert isinstance(\n xs, (flax.core.FrozenDict, dict)\n ), f'expected (frozen)dict; got {type(xs)}'\n\n return _flatten(xs, (), keep_empty_nodes, is_leaf, sep)\n\n\ndef unflatten_dict(xs, sep=None):\n \"\"\"Unflatten a dictionary.\n\n See ``flatten_dict``\n\n Example::\n\n >>> flat_xs = {\n ... ('foo',): 1,\n ... ('bar', 'a'): 2,\n ... }\n >>> xs = unflatten_dict(flat_xs)\n >>> xs\n {'foo': 1, 'bar': {'a': 2}}\n\n Args:\n xs: a flattened dictionary\n sep: separator (same as used with ``flatten_dict()``).\n Returns:\n The nested dictionary.\n \"\"\"\n assert isinstance(xs, dict), f'input is not a dict; it is a {type(xs)}'\n result = {}\n for path, value in xs.items():\n if sep is not None:\n path = path.split(sep)\n if value is empty_node:\n value = {}\n cursor = result\n for key in path[:-1]:\n if key not in cursor:\n cursor[key] = {}\n cursor = cursor[key]\n cursor[path[-1]] = value\n return result\n\n\ndef path_aware_map(\n f: Callable[[PathParts, Any], Any], nested_dict: VariableDict\n) -> VariableDict:\n \"\"\"A map function that operates over nested dictionary structures while taking\n the path to each leaf into account.\n\n Example::\n\n >>> import jax.numpy as jnp\n >>> from flax import traverse_util\n\n >>> params = {'a': {'x': 10, 'y': 3}, 'b': {'x': 20}}\n >>> f = lambda path, x: x + 5 if 'x' in path else -x\n >>> traverse_util.path_aware_map(f, params)\n {'a': {'x': 15, 'y': -3}, 'b': {'x': 25}}\n\n Args:\n f: A callable that takes in ``(path, value)`` arguments and maps them\n to a new value. Here ``path`` is a tuple of strings.\n nested_dict: A nested dictionary structure.\n\n Returns:\n A new nested dictionary structure with the mapped values.\n \"\"\"\n flat = flatten_dict(nested_dict, keep_empty_nodes=True)\n return unflatten_dict(\n {k: f(k, v) if v is not empty_node else v for k, v in flat.items()}\n )\n\n\nclass Traversal(abc.ABC):\n \"\"\"Base class for all traversals.\"\"\"\n\n def __new__(cls, *args, **kwargs):\n # Must override __new__ instead of __init__ since this is an ABC\n warnings.warn(\n '`flax.traverse_util.Traversal` will be deprecated. If you are using '\n 'it for `flax.optim`, use `optax` instead. Refer to the update guide '\n 'https://flax.readthedocs.io/en/latest/guides/converting_and_upgrading/optax_update_guide.html '\n 'for detailed instructions.',\n DeprecationWarning,\n )\n return super().__new__(cls)\n\n @abc.abstractmethod\n def update(self, fn, inputs):\n \"\"\"Update the focused items.\n\n Args:\n fn: the callback function that maps each traversed item\n to its updated value.\n inputs: the object that should be traversed.\n Returns:\n A new object with the updated values.\n \"\"\"\n pass\n\n @abc.abstractmethod\n def iterate(self, inputs):\n \"\"\"Iterate over the values selected by this ``Traversal``.\n\n Args:\n inputs: the object that should be traversed.\n Returns:\n An iterator over the traversed values.\n \"\"\"\n pass\n\n def set(self, values, inputs):\n \"\"\"Overrides the values selected by the ``Traversal``.\n\n Args:\n values: a list containing the new values.\n inputs: the object that should be traversed.\n Returns:\n A new object with the updated values.\n \"\"\"\n\n def update_fn(_):\n if not values:\n raise ValueError('Not enough values provided')\n return values.pop(0)\n\n y = self.update(update_fn, inputs)\n if values:\n raise ValueError('Too many values provided')\n return y\n\n def compose(self, other):\n \"\"\"Compose two traversals.\"\"\"\n return TraverseCompose(self, other)\n\n def merge(self, *traversals):\n \"\"\"Compose an arbitrary number of traversals and merge the results.\"\"\"\n return self.compose(TraverseMerge(*traversals))\n\n def each(self):\n \"\"\"Traverse each item in the selected containers.\"\"\"\n return self.compose(TraverseEach())\n\n def tree(self):\n \"\"\"Traverse each item in a pytree.\"\"\"\n return self.compose(TraverseTree())\n\n def filter(self, fn):\n \"\"\"Filter the selected values.\"\"\"\n return self.compose(TraverseFilter(fn))\n\n def __getattr__(self, attr):\n return self.compose(TraverseAttr(attr))\n\n def __getitem__(self, key):\n return self.compose(TraverseItem(key))\n\n\nclass TraverseId(Traversal):\n \"\"\"The identity Traversal.\"\"\"\n\n def update(self, fn, inputs):\n return fn(inputs)\n\n def iterate(self, inputs):\n yield inputs\n\n\nwith warnings.catch_warnings():\n warnings.simplefilter('ignore', DeprecationWarning)\n t_identity = TraverseId()\n\n\nclass TraverseMerge(Traversal):\n \"\"\"Merges the selection from a set of traversals.\"\"\"\n\n def __init__(self, *traversals):\n self._traversals = traversals\n\n def update(self, fn, inputs):\n for traversal in self._traversals:\n inputs = traversal.update(fn, inputs)\n return inputs\n\n def iterate(self, inputs):\n for traversal in self._traversals:\n yield from traversal.iterate(inputs)\n\n\nclass TraverseCompose(Traversal):\n \"\"\"Compose two traversals.\"\"\"\n\n def __init__(self, x, y):\n self._x = x\n self._y = y\n\n def update(self, fn, inputs):\n def update_fn(x):\n return self._y.update(fn, x)\n\n return self._x.update(update_fn, inputs)\n\n def iterate(self, inputs):\n for x in self._x.iterate(inputs):\n yield from self._y.iterate(x)\n\n\nclass TraverseFilter(Traversal):\n \"\"\"Filter selected values based on a predicate.\"\"\"\n\n def __init__(self, fn):\n self._fn = fn\n\n def update(self, fn, inputs):\n if self._fn(inputs):\n return fn(inputs)\n else:\n return inputs\n\n def iterate(self, inputs):\n if self._fn(inputs):\n yield inputs\n\n\ndef _is_namedtuple(t):\n return issubclass(t, tuple) and hasattr(t, '_fields')\n\n\nclass TraverseAttr(Traversal):\n \"\"\"Traverse the attribute of an object.\"\"\"\n\n def __init__(self, attr):\n self._attr = attr\n\n def update(self, fn, inputs):\n value = fn(getattr(inputs, self._attr))\n if _is_namedtuple(type(inputs)):\n return inputs._replace(**{self._attr: value})\n elif dataclasses.is_dataclass(inputs):\n return dataclasses.replace(inputs, **{self._attr: value})\n else:\n inputs = copy.copy(inputs)\n setattr(inputs, self._attr, value)\n return inputs\n\n def iterate(self, inputs):\n yield getattr(inputs, self._attr)\n\n\nclass TraverseItem(Traversal):\n \"\"\"Traverse the item of an object.\"\"\"\n\n def __init__(self, key):\n self._key = key\n\n def update(self, fn, inputs):\n if isinstance(inputs, tuple):\n ty = type(inputs)\n if isinstance(self._key, slice):\n sl = self._key\n else:\n sl = slice(self._key, self._key + 1)\n indices = set(range(*sl.indices(len(inputs))))\n\n args = [\n fn(inputs[i]) if i in indices else inputs[i] for i in range(len(inputs))\n ]\n if _is_namedtuple(ty):\n return ty(*args)\n else:\n return ty(args)\n else:\n xs = copy.copy(inputs)\n xs[self._key] = fn(xs[self._key])\n return xs\n\n def iterate(self, inputs):\n if isinstance(self._key, slice):\n yield from inputs[self._key]\n else:\n yield inputs[self._key]\n\n\nclass TraverseEach(Traversal):\n \"\"\"Traverse each item of a container.\"\"\"\n\n def update(self, fn, inputs):\n ty = type(inputs)\n if ty is dict:\n return {key: fn(val) for key, val in inputs.items()}\n if ty not in {list, tuple}:\n raise ValueError('Only the entries of a list or tuple can be traversed.')\n return ty(fn(x) for x in inputs)\n\n def iterate(self, inputs):\n if isinstance(inputs, dict):\n yield from inputs.values()\n else:\n yield from inputs\n\n\nclass TraverseTree(Traversal):\n \"\"\"Traverse every item in a pytree.\"\"\"\n\n def update(self, fn, inputs):\n return jax.tree_util.tree_map(fn, inputs)\n\n def iterate(self, inputs):\n yield from jax.tree_util.tree_leaves(inputs)\n\n\ndef _get_params_dict(inputs):\n if isinstance(inputs, (dict, flax.core.FrozenDict)):\n return flax.core.unfreeze(inputs)\n else:\n raise ValueError(\n 'Can only traverse a flax Model instance or a nested dict, not '\n f'{type(inputs)}'\n )\n\n\ndef _sorted_items(x):\n \"\"\"Returns items of a dict ordered by keys.\"\"\"\n return sorted(x.items(), key=lambda x: x[0])\n\n\nclass ModelParamTraversal(Traversal):\n \"\"\"Select model parameters using a name filter.\n\n This traversal operates on a nested dictionary of parameters and selects a\n subset based on the ``filter_fn`` argument.\n\n See :class:`flax.optim.MultiOptimizer` for an example of how to use\n :class:`ModelParamTraversal` to update subsets of the parameter tree with a\n specific optimizer.\n \"\"\"\n\n def __init__(self, filter_fn):\n \"\"\"Constructor a new ModelParamTraversal.\n\n Args:\n filter_fn: a function that takes a parameter's full name and its value and\n returns whether this parameter should be selected or not. The name of a\n parameter is determined by the module hierarchy and the parameter name\n (for example: '/module/sub_module/parameter_name').\n \"\"\"\n self._filter_fn = filter_fn\n\n def iterate(self, inputs):\n params = _get_params_dict(inputs)\n flat_dict = flatten_dict(params)\n for key, value in _sorted_items(flat_dict):\n path = '/' + '/'.join(key)\n if self._filter_fn(path, value):\n yield value\n\n def update(self, fn, inputs):\n params = _get_params_dict(inputs)\n flat_dict = flatten_dict(params, keep_empty_nodes=True)\n new_dict = {}\n for key, value in _sorted_items(flat_dict):\n # empty_node is not an actual leave. It's just a stub for empty nodes\n # in the nested dict.\n if value is not empty_node:\n path = '/' + '/'.join(key)\n if self._filter_fn(path, value):\n value = fn(value)\n new_dict[key] = value\n new_params = unflatten_dict(new_dict)\n if isinstance(inputs, flax.core.FrozenDict):\n return flax.core.FrozenDict(new_params)\n else:\n return new_params\n" }, { "path": "flax/typing.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nimport abc\nfrom collections import deque\nimport functools\nfrom functools import partial\nfrom typing import (\n Any,\n Generic,\n Optional,\n Protocol,\n TypeGuard,\n TypeVar,\n Union,\n)\nfrom collections.abc import Iterator\nfrom collections.abc import Callable, Hashable, Mapping, Sequence\n\nimport jax\nimport jax.numpy as jnp\nimport numpy as np\nfrom flax.core import FrozenDict\n\nimport dataclasses\nimport jax.tree_util as jtu\n\n\n# General\n\nArray = Union[jax.Array, Any]\nPRNGKey = jax.Array\nRNGSequences = dict[str, PRNGKey]\nDtype = Union[jax.typing.DTypeLike, Any]\nShape = Sequence[int]\nK = TypeVar('K')\n\nclass Key(Hashable, Protocol):\n def __lt__(self: K, value: K, /) -> bool:\n ...\n\ndef is_key_like(x: Any) -> TypeGuard[Key]:\n return hasattr(x, '__hash__') and hasattr(x, '__lt__')\n\nPath = str\nPathParts = tuple[Key, ...]\n\nLeaf = Any\n\n\n# Linear\n\nPrecisionLike = Union[\n None,\n str,\n jax.lax.Precision,\n tuple[str, str],\n tuple[jax.lax.Precision, jax.lax.Precision],\n]\nDotGeneralT = Callable[..., Array]\nConvGeneralDilatedT = Callable[..., Array]\nEinsumT = Callable[..., Array]\n\nPaddingLike = Union[str, int, Sequence[Union[int, tuple[int, int]]]]\nLaxPadding = Union[str, Sequence[tuple[int, int]]]\n\n\n# Initializers\n\nInitializer = Union[jax.nn.initializers.Initializer, Callable[..., Any]]\n\n\n# Collections\n\nCollection = Mapping[str, Any]\nMutableCollection = dict[str, Any]\n\n\n# Dicts\n\nVariableDict = Mapping[str, Collection]\nFrozenVariableDict = FrozenDict[str, Collection]\nMutableVariableDict = dict[str, MutableCollection]\n\nPRNGFoldable = Union[int, str]\n\n\n# Axes\n\nT = TypeVar('T')\n\n@dataclasses.dataclass(frozen=True)\nclass In(Generic[T]):\n \"\"\"Specifies a variable collection should only be lifted as input.\"\"\"\n\n axis: T\n\n@dataclasses.dataclass(frozen=True)\nclass Out(Generic[T]):\n \"\"\"Specifies a variable collection should only be lifted as output.\"\"\"\n\n axis: T\n\nAxis = Optional[int]\nInOutAxis = Union[Axis, In[Axis], Out[Axis]]\n\nScanAxis = int\nInOutScanAxis = Union[ScanAxis, In[ScanAxis], Out[ScanAxis]]\n\nAxes = Union[int, Sequence[int]]\n\n\n# SPMD\n\nLogicalNames = tuple[Union[str, None], ...]\nAxisName = str | tuple[str, ...] | None\n\n# Maps each logical axis to physical mesh, can be either None (replicated),\n# one physical axis or a tuple of physical axes.\nLogicalRules = Sequence[tuple[str, AxisName]]\nArrayPytree = Any # pylint: disable=invalid-name\nLogicalPartitionSpec = Any # pylint: disable=invalid-name\nLogicalPartitionSpecPytree = Any # pylint: disable=invalid-name\nPartitionSpecPytree = Any # pylint: disable=invalid-name\n\nSharding = tuple[AxisName, ...]\n\nA = TypeVar('A')\nHA = TypeVar('HA', bound=Hashable)\nHB = TypeVar('HB')\n\n\nclass PytreeDeque(deque[A]):\n pass\n\n\ndef _pytree_deque_flatten(xs: PytreeDeque, *, with_path: bool):\n if with_path:\n nodes = tuple((jtu.SequenceKey(i), x) for i, x in enumerate(xs))\n return nodes, ()\n else:\n return xs, ()\n\n\ndef _pytree_deque_unflatten(_, nodes):\n return PytreeDeque(nodes)\n\n\njtu.register_pytree_with_keys(\n PytreeDeque,\n partial(_pytree_deque_flatten, with_path=True),\n _pytree_deque_unflatten,\n flatten_func=partial(_pytree_deque_flatten, with_path=False),\n)\n\nclass Missing:\n pass\n\n\nMISSING = Missing()\n\n\ndef _bytes_repr(num_bytes):\n count, units = (\n (f'{num_bytes / 1e9:,.1f}', 'GB')\n if num_bytes > 1e9\n else (f'{num_bytes / 1e6:,.1f}', 'MB')\n if num_bytes > 1e6\n else (f'{num_bytes / 1e3:,.1f}', 'KB')\n if num_bytes > 1e3\n else (f'{num_bytes:,}', 'B')\n )\n\n return f'{count} {units}'\n\n\nclass ShapeDtype(Protocol):\n shape: Shape\n dtype: Dtype\n\n\ndef has_shape_dtype(x: Any) -> TypeGuard[ShapeDtype]:\n return hasattr(x, 'shape') and hasattr(x, 'dtype')\n\n\n@dataclasses.dataclass(frozen=True, slots=True)\nclass SizeBytes: # type: ignore[misc]\n size: int\n bytes: int\n\n @classmethod\n def from_array(cls, x: ShapeDtype):\n size = int(np.prod(x.shape))\n dtype: jnp.dtype\n if isinstance(x.dtype, str):\n dtype = jnp.dtype(x.dtype)\n else:\n dtype = x.dtype # type: ignore\n bytes = size * dtype.itemsize # type: ignore\n return cls(size, bytes)\n\n def __add__(self, other: SizeBytes):\n return type(self)(self.size + other.size, self.bytes + other.bytes)\n\n def __bool__(self) -> bool:\n return bool(self.size)\n\n def __repr__(self) -> str:\n bytes_repr = _bytes_repr(self.bytes)\n return f'{self.size:,} ({bytes_repr})'\n\n @classmethod\n def from_any(cls, x):\n leaves = jax.tree.leaves(x)\n size_bytes = cls(0, 0)\n for leaf in leaves:\n if has_shape_dtype(leaf):\n size_bytes += cls.from_array(leaf)\n\n return size_bytes\n\n\nTupleArg = TypeVar('TupleArg', bound=tuple)\n\n\nclass PromoteDtypeFn(Protocol):\n def __call__(\n self, args: TupleArg, /, *, dtype: Any = None, inexact: bool = True\n ) -> TupleArg: ...\n\n\nclass HashableMapping(Mapping[HA, HB], Hashable):\n _mapping: dict[HA, HB] | Mapping[HA, HB]\n\n def __init__(self, mapping: Mapping[HA, HB], copy: bool = True):\n self._mapping = dict(mapping) if copy else mapping\n\n def __contains__(self, key: object) -> bool:\n return key in self._mapping\n\n def __getitem__(self, key: HA) -> HB:\n return self._mapping[key]\n\n def __iter__(self) -> Iterator[HA]:\n return iter(self._mapping)\n\n def __len__(self) -> int:\n return len(self._mapping)\n\n def __hash__(self) -> int:\n # use type-aware sorting to support int keys\n def _pytree__key_sort_fn(item: tuple[Any, Any]) -> tuple[int, Any]:\n key, _ = item\n if isinstance(key, int):\n return (0, key)\n elif isinstance(key, str):\n return (1, key)\n else:\n raise ValueError(f'Unsupported key type: {type(key)!r}')\n\n return hash(tuple(sorted(self._mapping.items(), key=_pytree__key_sort_fn)))\n\n def __eq__(self, other: Any) -> bool:\n return (\n isinstance(other, HashableMapping) and self._mapping == other._mapping\n )\n\n def __repr__(self) -> str:\n return repr(self._mapping)\n\n def update(self, other: Mapping[HA, HB]) -> HashableMapping[HA, HB]:\n \"\"\"Updates the mapping with another mapping.\"\"\"\n mapping = dict(self._mapping)\n mapping.update(other)\n return HashableMapping(mapping, copy=False)\n\n\nF = TypeVar('F', bound=Callable[..., Any])\n\n\nclass BaseConfigContext(abc.ABC):\n @classmethod\n @abc.abstractmethod\n def get_default(cls):\n ...\n\n @classmethod\n @abc.abstractmethod\n def get_stack(cls) -> list:\n ...\n\n def __init__(self, value, /):\n stack = self.get_stack()\n if stack:\n self.prev_value = stack[-1]\n stack[-1] = value\n else:\n self.prev_value = None\n stack.append(value)\n self.new_value = value\n\n @classmethod\n def current_value(cls):\n stack = cls.get_stack()\n if stack:\n return stack[-1]\n return cls.get_default()\n\n def __enter__(self):\n if self.prev_value is not None:\n self.get_stack().insert(-1, self.prev_value)\n\n def __exit__(self, exc_type, exc_value, traceback):\n self.get_stack().pop()\n\n def __call__(self, f: F) -> F:\n self.get_stack().pop()\n if self.prev_value is not None:\n self.get_stack().append(self.prev_value)\n\n @functools.wraps(f)\n def wrapper(*args, **kwargs):\n self.get_stack().append(self.new_value)\n try:\n return f(*args, **kwargs)\n finally:\n self.get_stack().pop()\n\n return wrapper # type: ignore[return-value]\n" }, { "path": "flax/version.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Current Flax version at head on Github.\"\"\"\n__version__ = '0.12.6'\n" }, { "path": "flaxlib_src/.gitignore", "content": "/target\n\n# Byte-compiled / optimized / DLL files\n__pycache__/\n.pytest_cache/\n*.py[cod]\n\n# C extensions\n*.so\n\n# Distribution / packaging\n.Python\n.venv/\nenv/\nbin/\nbuild/\ndevelop-eggs/\ndist/\neggs/\nlib/\nlib64/\nparts/\nsdist/\nvar/\ninclude/\nman/\nvenv/\n*.egg-info/\n.installed.cfg\n*.egg\n\n# Installer logs\npip-log.txt\npip-delete-this-directory.txt\npip-selfcheck.json\n\n# Unit test / coverage reports\nhtmlcov/\n.tox/\n.coverage\n.cache\nnosetests.xml\ncoverage.xml\n\n# Translations\n*.mo\n\n# Mr Developer\n.mr.developer.cfg\n.project\n.pydevproject\n\n# Rope\n.ropeproject\n\n# Django stuff:\n*.log\n*.pot\n\n.DS_Store\n\n# Sphinx documentation\ndocs/_build/\n\n# PyCharm\n.idea/\n\n# VSCode\n.vscode/\n\n# Pyenv\n.python-version\n\n# cibuildwheel\n/wheelhouse" }, { "path": "flaxlib_src/CMakeLists.txt", "content": "# Set the minimum CMake version and policies for highest tested version\ncmake_minimum_required(VERSION 3.15...3.27)\n\n# Set up the project and ensure there is a working C++ compiler\nproject(flaxlib LANGUAGES CXX)\n\n# Warn if the user invokes CMake directly\nif (NOT SKBUILD)\n message(WARNING \"\\\n This CMake file is meant to be executed using 'scikit-build-core'.\n Running it directly will almost certainly not produce the desired\n result. If you are a user trying to install this package, use the\n command below, which will install all necessary build dependencies,\n compile the package in an isolated environment, and then install it.\n =====================================================================\n $ pip install .\n =====================================================================\n If you are a software developer, and this is your own package, then\n it is usually much more efficient to install the build dependencies\n in your environment once and use the following command that avoids\n a costly creation of a new virtual environment at every compilation:\n =====================================================================\n $ pip install nanobind scikit-build-core[pyproject]\n $ pip install --no-build-isolation -ve .\n =====================================================================\n You may optionally add -Ceditable.rebuild=true to auto-rebuild when\n the package is imported. Otherwise, you need to rerun the above\n after editing C++ files.\")\nendif()\n\n# Try to import all Python components potentially needed by nanobind\nfind_package(Python 3.8\n REQUIRED COMPONENTS Interpreter Development.Module\n OPTIONAL_COMPONENTS Development.SABIModule)\n\n# Import nanobind through CMake's find_package mechanism\nfind_package(nanobind CONFIG REQUIRED)\n\n# We are now ready to compile the actual extension module\nnanobind_add_module(\n # Name of the extension\n flaxlib_cpp\n\n # Target the stable ABI for Python 3.12+, which reduces\n # the number of binary wheels that must be built. This\n # does nothing on older Python versions\n STABLE_ABI\n\n # Source code goes here\n src/lib.cc\n)\n\n# Install directive for scikit-build-core\ninstall(TARGETS flaxlib_cpp LIBRARY DESTINATION flaxlib)" }, { "path": "flaxlib_src/Cargo.toml", "content": "[package]\nname = \"flaxlib\"\nversion = \"0.0.1-a1\"\nedition = \"2021\"\n\n# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html\n[lib]\nname = \"flaxlib\"\ncrate-type = [\"cdylib\"]\n\n[dependencies]\npyo3 = \"0.21.2\"\n" }, { "path": "flaxlib_src/LICENSE", "content": " Version 2.0, January 2004\n http://www.apache.org/licenses/\n\n TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION\n\n 1. Definitions.\n\n \"License\" shall mean the terms and conditions for use, reproduction,\n and distribution as defined by Sections 1 through 9 of this document.\n\n \"Licensor\" shall mean the copyright owner or entity authorized by\n the copyright owner that is granting the License.\n\n \"Legal Entity\" shall mean the union of the acting entity and all\n other entities that control, are controlled by, or are under common\n control with that entity. For the purposes of this definition,\n \"control\" means (i) the power, direct or indirect, to cause the\n direction or management of such entity, whether by contract or\n otherwise, or (ii) ownership of fifty percent (50%) or more of the\n outstanding shares, or (iii) beneficial ownership of such entity.\n\n \"You\" (or \"Your\") shall mean an individual or Legal Entity\n exercising permissions granted by this License.\n\n \"Source\" form shall mean the preferred form for making modifications,\n including but not limited to software source code, documentation\n source, and configuration files.\n\n \"Object\" form shall mean any form resulting from mechanical\n transformation or translation of a Source form, including but\n not limited to compiled object code, generated documentation,\n and conversions to other media types.\n\n \"Work\" shall mean the work of authorship, whether in Source or\n Object form, made available under the License, as indicated by a\n copyright notice that is included in or attached to the work\n (an example is provided in the Appendix below).\n\n \"Derivative Works\" shall mean any work, whether in Source or Object\n form, that is based on (or derived from) the Work and for which the\n editorial revisions, annotations, elaborations, or other modifications\n represent, as a whole, an original work of authorship. For the purposes\n of this License, Derivative Works shall not include works that remain\n separable from, or merely link (or bind by name) to the interfaces of,\n the Work and Derivative Works thereof.\n\n \"Contribution\" shall mean any work of authorship, including\n the original version of the Work and any modifications or additions\n to that Work or Derivative Works thereof, that is intentionally\n submitted to Licensor for inclusion in the Work by the copyright owner\n or by an individual or Legal Entity authorized to submit on behalf of\n the copyright owner. For the purposes of this definition, \"submitted\"\n means any form of electronic, verbal, or written communication sent\n to the Licensor or its representatives, including but not limited to\n communication on electronic mailing lists, source code control systems,\n and issue tracking systems that are managed by, or on behalf of, the\n Licensor for the purpose of discussing and improving the Work, but\n excluding communication that is conspicuously marked or otherwise\n designated in writing by the copyright owner as \"Not a Contribution.\"\n\n \"Contributor\" shall mean Licensor and any individual or Legal Entity\n on behalf of whom a Contribution has been received by Licensor and\n subsequently incorporated within the Work.\n\n 2. Grant of Copyright License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n copyright license to reproduce, prepare Derivative Works of,\n publicly display, publicly perform, sublicense, and distribute the\n Work and such Derivative Works in Source or Object form.\n\n 3. Grant of Patent License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n (except as stated in this section) patent license to make, have made,\n use, offer to sell, sell, import, and otherwise transfer the Work,\n where such license applies only to those patent claims licensable\n by such Contributor that are necessarily infringed by their\n Contribution(s) alone or by combination of their Contribution(s)\n with the Work to which such Contribution(s) was submitted. If You\n institute patent litigation against any entity (including a\n cross-claim or counterclaim in a lawsuit) alleging that the Work\n or a Contribution incorporated within the Work constitutes direct\n or contributory patent infringement, then any patent licenses\n granted to You under this License for that Work shall terminate\n as of the date such litigation is filed.\n\n 4. Redistribution. You may reproduce and distribute copies of the\n Work or Derivative Works thereof in any medium, with or without\n modifications, and in Source or Object form, provided that You\n meet the following conditions:\n\n (a) You must give any other recipients of the Work or\n Derivative Works a copy of this License; and\n\n (b) You must cause any modified files to carry prominent notices\n stating that You changed the files; and\n\n (c) You must retain, in the Source form of any Derivative Works\n that You distribute, all copyright, patent, trademark, and\n attribution notices from the Source form of the Work,\n excluding those notices that do not pertain to any part of\n the Derivative Works; and\n\n (d) If the Work includes a \"NOTICE\" text file as part of its\n distribution, then any Derivative Works that You distribute must\n include a readable copy of the attribution notices contained\n within such NOTICE file, excluding those notices that do not\n pertain to any part of the Derivative Works, in at least one\n of the following places: within a NOTICE text file distributed\n as part of the Derivative Works; within the Source form or\n documentation, if provided along with the Derivative Works; or,\n within a display generated by the Derivative Works, if and\n wherever such third-party notices normally appear. The contents\n of the NOTICE file are for informational purposes only and\n do not modify the License. You may add Your own attribution\n notices within Derivative Works that You distribute, alongside\n or as an addendum to the NOTICE text from the Work, provided\n that such additional attribution notices cannot be construed\n as modifying the License.\n\n You may add Your own copyright statement to Your modifications and\n may provide additional or different license terms and conditions\n for use, reproduction, or distribution of Your modifications, or\n for any such Derivative Works as a whole, provided Your use,\n reproduction, and distribution of the Work otherwise complies with\n the conditions stated in this License.\n\n 5. Submission of Contributions. Unless You explicitly state otherwise,\n any Contribution intentionally submitted for inclusion in the Work\n by You to the Licensor shall be under the terms and conditions of\n this License, without any additional terms or conditions.\n Notwithstanding the above, nothing herein shall supersede or modify\n the terms of any separate license agreement you may have executed\n with Licensor regarding such Contributions.\n\n 6. Trademarks. This License does not grant permission to use the trade\n names, trademarks, service marks, or product names of the Licensor,\n except as required for reasonable and customary use in describing the\n origin of the Work and reproducing the content of the NOTICE file.\n\n 7. Disclaimer of Warranty. Unless required by applicable law or\n agreed to in writing, Licensor provides the Work (and each\n Contributor provides its Contributions) on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or\n implied, including, without limitation, any warranties or conditions\n of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A\n PARTICULAR PURPOSE. You are solely responsible for determining the\n appropriateness of using or redistributing the Work and assume any\n risks associated with Your exercise of permissions under this License.\n\n 8. Limitation of Liability. In no event and under no legal theory,\n whether in tort (including negligence), contract, or otherwise,\n unless required by applicable law (such as deliberate and grossly\n negligent acts) or agreed to in writing, shall any Contributor be\n liable to You for damages, including any direct, indirect, special,\n incidental, or consequential damages of any character arising as a\n result of this License or out of the use or inability to use the\n Work (including but not limited to damages for loss of goodwill,\n work stoppage, computer failure or malfunction, or any and all\n other commercial damages or losses), even if such Contributor\n has been advised of the possibility of such damages.\n\n 9. Accepting Warranty or Additional Liability. While redistributing\n the Work or Derivative Works thereof, You may choose to offer,\n and charge a fee for, acceptance of support, warranty, indemnity,\n or other liability obligations and/or rights consistent with this\n License. However, in accepting such obligations, You may act only\n on Your own behalf and on Your sole responsibility, not on behalf\n of any other Contributor, and only if You agree to indemnify,\n defend, and hold each Contributor harmless for any liability\n incurred by, or claims asserted against, such Contributor by reason\n of your accepting any such warranty or additional liability.\n\n END OF TERMS AND CONDITIONS\n\n APPENDIX: How to apply the Apache License to your work.\n\n To apply the Apache License to your work, attach the following\n boilerplate notice, with the fields enclosed by brackets \"[]\"\n replaced with your own identifying information. (Don't include\n the brackets!) The text should be enclosed in the appropriate\n comment syntax for the file format. We also recommend that a\n file or class name and description of purpose be included on the\n same \"printed page\" as the copyright notice for easier\n identification within third-party archives.\n\n Copyright [yyyy] [name of copyright owner]\n\n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n\n http://www.apache.org/licenses/LICENSE-2.0\n\n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n" }, { "path": "flaxlib_src/README.md", "content": "# flaxlib\n\n## Build flaxlib from source\n\nInstall necessary dependencies to build the C++ based package.\n\n```shell\npip install meson-python ninja build\n```\n\nClone the Flax repository, navigate to the flaxlib source directory.\n\n```shell\ngit clone git@github.com:google/flax.git\ncd flax/flaxlib_src\n```\n\nConfigure the build.\n\n```shell\nmkdir -p subprojects\nmeson wrap install robin-map\nmeson wrap install nanobind\nmeson setup builddir\n```\n\nCompile the code. You'll need to run this repeatedly if you modify the source\ncode. Note that the actual wheel name will differ depending on your system.\n\n```shell\nmeson compile -C builddir\npython -m build . -w\npip install dist/flaxlib-0.0.1-cp311-cp311-macosx_14_0_arm64.whl --force-reinstall\n```\n" }, { "path": "flaxlib_src/pyproject.toml", "content": "[build-system]\nrequires = [\"scikit-build-core >=0.4.3\", \"nanobind >=1.3.2\"]\nbuild-backend = \"scikit_build_core.build\"\n\n[project]\nname = \"flaxlib\"\nversion = \"0.0.1\"\nrequires-python = \">=3.10\"\nclassifiers = [\n \"Programming Language :: C++\",\n \"Programming Language :: Python :: Implementation :: CPython\",\n \"Programming Language :: Python :: Implementation :: PyPy\",\n]\n\n[project.optional-dependencies]\ntests = [\n \"pytest\",\n]\n\n[tool.scikit-build]\n# Protect the configuration against future changes in scikit-build-core\nminimum-version = \"0.4\"\n\n# Setuptools-style build caching in a local directory\nbuild-dir = \"build/{wheel_tag}\"\n\n# Build stable ABI wheels for CPython 3.12+\nwheel.py-api = \"cp312\"" }, { "path": "flaxlib_src/src/flaxlib/__init__.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom .flaxlib_cpp import RefMap as RefMap\nfrom .flaxlib_cpp import IndexMap as IndexMap\nfrom .flaxlib_cpp import NodeDef as NodeDef\nfrom .flaxlib_cpp import VariableDef as VariableDef\nfrom .flaxlib_cpp import NodeRef as NodeRef\n" }, { "path": "flaxlib_src/src/flaxlib/flaxlib_cpp.pyi", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport typing as tp\n\nRefMap = tp.MutableMapping[tp.Any, int]\nIndexMap = dict[int, tp.Any]\n\nclass NodeDef:\n type: type\n index: int | None\n outer_index: int | None\n num_attributes: int\n metadata: tp.Any\n\n def with_no_outer_index(self) -> NodeDef: ...\n def with_same_outer_index(self) -> NodeDef: ...\n def __eq__(self, other: tp.Any) -> bool: ...\n def __hash__(self) -> int: ...\n def __getstate__(\n self,\n ) -> tuple[tp.Any, tp.Any, tp.Any, tp.Any, tp.Any]: ...\n @staticmethod\n def __setstate__(\n nodedef: NodeDef, state: tuple[tp.Any, tp.Any, tp.Any, tp.Any, tp.Any]\n ) -> None: ...\n\nclass VariableDef:\n type: type\n index: int\n outer_index: int | None\n metadata: tp.Any\n\n def with_no_outer_index(self) -> VariableDef: ...\n def with_same_outer_index(self) -> VariableDef: ...\n def __eq__(self, other: tp.Any) -> bool: ...\n def __hash__(self) -> int: ...\n def __getstate__(\n self,\n ) -> tuple[tp.Any, int, tp.Any, tp.Any]: ...\n @staticmethod\n def __setstate__(\n variabledef: 'VariableDef', state: tuple[tp.Any, int, tp.Any, tp.Any]\n ) -> None: ...\n\nclass NodeRef:\n index: int\n\n def __eq__(self, other: tp.Any) -> bool: ...\n def __hash__(self) -> int: ...\n def __getstate__(self) -> tuple[int]: ...\n @staticmethod\n def __setstate__(noderef: NodeRef, state: tuple[int]) -> None: ...\n\n" }, { "path": "flaxlib_src/src/lib.cc", "content": "// Copyright 2024 The Flax Authors.\n//\n// Licensed under the Apache License, Version 2.0 (the \"License\");\n// you may not use this file except in compliance with the License.\n// You may obtain a copy of the License at\n//\n// http://www.apache.org/licenses/LICENSE-2.0\n//\n// Unless required by applicable law or agreed to in writing, software\n// distributed under the License is distributed on an \"AS IS\" BASIS,\n// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n// See the License for the specific language governing permissions and\n// limitations under the License.\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nnamespace nb = nanobind;\nusing namespace nb::literals;\n\n// -----------------------------------\n// helper functions\n// -----------------------------------\nintptr_t nb_id(const nb::object &obj)\n{\n // Get the object ID\n return reinterpret_cast(obj.ptr());\n}\n\nnb::tuple vector_to_tuple(const std::vector &vec)\n{\n\n if (vec.empty())\n {\n return nb::tuple();\n }\n else\n {\n return nb::tuple(nb::cast(vec));\n }\n}\n\n// 1. Hash function for nb::object\nstruct NbObjectHash\n{\n std::size_t operator()(const nb::object &obj) const\n {\n return nb::hash(obj);\n }\n};\n\n// 2. Equality function for nb::object (Important!)\nstruct NbObjectEqual\n{\n bool operator()(const nb::object &a, const nb::object &b) const\n {\n return a.equal(b);\n }\n};\n\nnamespace flaxlib\n{\n //---------------------------------------------------------------\n // NNXContext\n //---------------------------------------------------------------\n\n struct PythonContext\n {\n nb::object nnx;\n nb::object graph;\n nb::object jax;\n nb::object np;\n nb::object jax_Array;\n nb::object np_ndarray;\n nb::type_object GraphNodeImpl;\n nb::type_object PytreeNodeImpl;\n nb::type_object Object;\n nb::type_object Variable;\n nb::object get_node_impl;\n\n PythonContext()\n {\n nnx = nb::module_::import_(\"flax.nnx\");\n graph = nb::module_::import_(\"flax.nnx.graph\");\n jax = nb::module_::import_(\"jax\");\n np = nb::module_::import_(\"numpy\");\n jax_Array = jax.attr(\"Array\");\n np_ndarray = np.attr(\"ndarray\");\n GraphNodeImpl = graph.attr(\"GraphNodeImpl\");\n PytreeNodeImpl = graph.attr(\"PytreeNodeImpl\");\n Object = nnx.attr(\"Object\");\n Variable = graph.attr(\"Variable\");\n get_node_impl = graph.attr(\"get_node_impl\");\n }\n\n ~PythonContext()\n {\n graph.release();\n jax.release();\n np.release();\n jax_Array.release();\n np_ndarray.release();\n GraphNodeImpl.release();\n PytreeNodeImpl.release();\n Variable.release();\n get_node_impl.release();\n }\n };\n\n static std::optional _python_context;\n\n PythonContext &get_python_context()\n {\n if (!_python_context)\n {\n _python_context.emplace();\n }\n return *_python_context;\n }\n\n //---------------------------------------------------------------\n // IndexMap\n //---------------------------------------------------------------\n\n struct IndexMap : public std::unordered_map\n {\n };\n\n //---------------------------------------------------------------\n // RefMap\n //---------------------------------------------------------------\n\n struct RefMapKeysIterator\n {\n std::unordered_map>::iterator it;\n std::unordered_map>::iterator end;\n\n RefMapKeysIterator(std::unordered_map>::iterator it, std::unordered_map>::iterator end)\n : it(it), end(end) {}\n\n nb::object __next__()\n {\n if (it == end)\n {\n throw nb::stop_iteration();\n }\n auto elem = it->second;\n ++it;\n return std::get<0>(elem);\n }\n };\n\n struct RefMapItemsIterator\n {\n std::unordered_map>::iterator it;\n std::unordered_map>::iterator end;\n\n RefMapItemsIterator(std::unordered_map>::iterator it, std::unordered_map>::iterator end)\n : it(it), end(end) {}\n\n RefMapItemsIterator __iter__()\n {\n return *this;\n }\n\n nb::tuple __next__()\n {\n if (it == end)\n {\n throw nb::stop_iteration();\n }\n auto elem = it->second;\n ++it;\n return nb::make_tuple(std::get<0>(elem), std::get<1>(elem));\n }\n };\n\n struct RefMap\n {\n std::unordered_map> mapping;\n\n RefMap() {}\n\n RefMap(const nb::object &iterable) : RefMap()\n {\n for (auto item : iterable)\n {\n nb::object obj = item[0];\n auto value = nb::cast(item[1]);\n mapping[nb_id(obj)] = {obj, value};\n }\n }\n\n void update(const RefMap &other)\n {\n for (const auto &[key_id, value_tuple] : other.mapping)\n {\n mapping[key_id] = value_tuple;\n }\n }\n\n int __getitem__(const nb::object &key)\n {\n return std::get<1>(mapping[nb_id(key)]);\n }\n\n void __setitem__(const nb::object &key, int value)\n {\n mapping[nb_id(key)] = std::make_tuple(key, value);\n }\n\n int __len__() const\n {\n return mapping.size();\n }\n\n bool __contains__(const nb::object &key) const\n {\n return mapping.find(nb_id(key)) != mapping.end();\n }\n\n RefMapKeysIterator __iter__()\n {\n return RefMapKeysIterator(mapping.begin(), mapping.end());\n };\n\n RefMapItemsIterator items()\n {\n return RefMapItemsIterator(mapping.begin(), mapping.end());\n }\n\n std::optional get(const nb::object &key, std::optional default_value = std::nullopt)\n {\n auto it = mapping.find(nb_id(key));\n if (it != mapping.end())\n {\n return std::get<1>(it->second);\n }\n return default_value;\n }\n };\n\n static IndexMap indexmap_from_refmap(const RefMap &refmap)\n {\n IndexMap indexmap;\n for (const auto &[_, value_index] : refmap.mapping)\n {\n nb::object value = std::get<0>(value_index);\n int index = std::get<1>(value_index);\n indexmap[index] = value;\n }\n return indexmap;\n };\n\n static RefMap refmap_from_indexmap(const IndexMap &indexmap)\n {\n RefMap refmap;\n for (const auto &[index, value] : indexmap)\n {\n refmap.mapping[nb_id(value)] = std::make_tuple(value, index);\n }\n return refmap;\n };\n\n //---------------------------------------------------------------\n // NodeDef\n //---------------------------------------------------------------\n\n struct NodeDef\n {\n nb::object type;\n std::optional index;\n std::optional outer_index;\n int num_attributes;\n nb::object metadata;\n\n NodeDef(nb::object type, std::optional index, std::optional outer_index, int num_attributes, nb::object metadata)\n : type(type), index(index), outer_index(outer_index), num_attributes(num_attributes), metadata(metadata) {}\n\n NodeDef with_no_outer_index() const\n {\n return NodeDef(type, index, std::nullopt, num_attributes, metadata);\n }\n\n NodeDef with_same_outer_index() const\n {\n return NodeDef(type, index, index, num_attributes, metadata);\n }\n\n bool __eq__(const nb::object &other_obj) const\n {\n if (!nb::isinstance(other_obj))\n {\n return false;\n }\n NodeDef other = nb::cast(other_obj);\n return type.equal(other.type) && index == other.index && outer_index == other.outer_index && num_attributes == other.num_attributes && metadata.equal(other.metadata);\n }\n\n int __hash__() const\n {\n // return nb::hash(type) ^ nb::hash(nb::cast(index)) ^ nb::hash(nb::cast(outer_index)) ^ nb::hash(nb::cast(num_attributes)) ^ nb::hash(metadata);\n return nb::hash(nb::make_tuple(type, index, outer_index, num_attributes, metadata));\n }\n\n nb::tuple __getstate__() const\n {\n return nb::make_tuple(type, index, outer_index, num_attributes, metadata);\n }\n\n static void __setstate__(NodeDef &nodedef, nb::tuple &t)\n {\n new (&nodedef) NodeDef(t[0], nb::cast>(t[1]), nb::cast>(t[2]), nb::cast(t[3]), t[4]);\n }\n };\n\n //---------------------------------------------------------------\n // VariableDef\n //---------------------------------------------------------------\n\n struct VariableDef\n {\n nb::object type;\n int index;\n std::optional outer_index;\n nb::object metadata;\n\n VariableDef(nb::object type, int index, std::optional outer_index, nb::object metadata)\n : type(type), index(index), outer_index(outer_index), metadata(metadata) {}\n\n VariableDef with_no_outer_index() const\n {\n return VariableDef(type, index, std::nullopt, metadata);\n }\n\n VariableDef with_same_outer_index() const\n {\n return VariableDef(type, index, index, metadata);\n }\n\n bool __eq__(const nb::object &other_obj) const\n {\n if (!nb::isinstance(other_obj))\n {\n return false;\n }\n VariableDef other = nb::cast(other_obj);\n return type.equal(other.type) && index == other.index && outer_index == other.outer_index && metadata.equal(other.metadata);\n }\n\n int __hash__() const\n {\n // return nb::hash(type) ^ nb::hash(nb::cast(index)) ^ nb::hash(nb::cast(outer_index)) ^ nb::hash(metadata);\n return nb::hash(nb::make_tuple(type, index, outer_index, metadata));\n }\n\n nb::tuple __getstate__() const\n {\n return nb::make_tuple(type, index, outer_index, metadata);\n }\n\n static void __setstate__(VariableDef &variabledef, nb::tuple &t)\n {\n new (&variabledef) VariableDef(t[0], nb::cast(t[1]), nb::cast>(t[2]), t[3]);\n }\n };\n\n //---------------------------------------------------------------\n // NodeRef\n //---------------------------------------------------------------\n\n struct NodeRef\n {\n int index;\n\n NodeRef(int index)\n : index(index) {}\n\n bool __eq__(const nb::object &other_obj) const\n {\n if (!nb::isinstance(other_obj))\n {\n return false;\n }\n NodeRef other = nb::cast(other_obj);\n return index == other.index;\n }\n\n int __hash__() const\n {\n return nb::hash(nb::cast(index));\n }\n\n nb::tuple __getstate__() const\n {\n return nb::make_tuple(index);\n }\n\n static void __setstate__(NodeRef &noderef, nb::tuple &t)\n {\n new (&noderef) NodeRef(nb::cast(t[0]));\n }\n };\n\n NB_MODULE(flaxlib_cpp, m)\n {\n nb::bind_map(m, \"IndexMap\")\n .def_static(\"from_refmap\", &indexmap_from_refmap);\n nb::class_(m, \"RefMapKeysIterator\")\n .def(\"__next__\", &flaxlib::RefMapKeysIterator::__next__);\n\n nb::class_(m, \"RefMapItemsIterator\")\n .def(\"__iter__\", &flaxlib::RefMapItemsIterator::__iter__)\n .def(\"__next__\", &flaxlib::RefMapItemsIterator::__next__);\n\n nb::class_(m, \"RefMap\")\n .def(nb::init<>())\n .def(nb::init(), nb::arg(\"iterable\"))\n .def(\"update\", &flaxlib::RefMap::update)\n .def_static(\"from_indexmap\", &refmap_from_indexmap)\n .def(\"__getitem__\", &flaxlib::RefMap::__getitem__, nb::arg(\"key\").none())\n .def(\"__setitem__\", &flaxlib::RefMap::__setitem__, nb::arg(\"key\").none(), nb::arg(\"value\"))\n .def(\"__len__\", &flaxlib::RefMap::__len__)\n .def(\"__contains__\", &flaxlib::RefMap::__contains__, nb::arg(\"key\").none())\n .def(\"__iter__\", &flaxlib::RefMap::__iter__)\n .def(\"items\", &flaxlib::RefMap::items)\n .def(\"get\", &flaxlib::RefMap::get, nb::arg(\"key\").none(), nb::arg(\"default_value\").none());\n\n nb::class_(m, \"NodeDef\")\n .def(nb::init, std::optional, int, nb::object>(),\n nb::arg(\"type\"), nb::arg(\"index\").none(), nb::arg(\"outer_index\").none(), nb::arg(\"num_attributes\"), nb::arg(\"metadata\").none())\n .def_prop_ro(\"type\", [](const flaxlib::NodeDef &n)\n { return n.type; })\n .def_prop_ro(\"index\", [](const flaxlib::NodeDef &n)\n { return n.index; })\n .def_prop_ro(\"outer_index\", [](const flaxlib::NodeDef &n)\n { return n.outer_index; })\n .def_prop_ro(\"num_attributes\", [](const flaxlib::NodeDef &n)\n { return n.num_attributes; })\n .def_prop_ro(\"metadata\", [](const flaxlib::NodeDef &n)\n { return n.metadata; })\n .def(\"with_no_outer_index\", &flaxlib::NodeDef::with_no_outer_index)\n .def(\"with_same_outer_index\", &flaxlib::NodeDef::with_same_outer_index)\n .def(\"__eq__\", &flaxlib::NodeDef::__eq__, nb::arg().none())\n .def(\"__hash__\", &flaxlib::NodeDef::__hash__)\n .def(\"__getstate__\", &flaxlib::NodeDef::__getstate__)\n .def(\"__setstate__\", &flaxlib::NodeDef::__setstate__);\n\n nb::class_(m, \"VariableDef\")\n .def(nb::init, nb::object>(),\n nb::arg(\"type\"), nb::arg(\"index\"), nb::arg(\"outer_index\").none(), nb::arg(\"metadata\").none())\n .def_prop_ro(\"type\", [](const flaxlib::VariableDef &n)\n { return n.type; })\n .def_prop_ro(\"index\", [](const flaxlib::VariableDef &n)\n { return n.index; })\n .def_prop_ro(\"outer_index\", [](const flaxlib::VariableDef &n)\n { return n.outer_index; })\n .def_prop_ro(\"metadata\", [](const flaxlib::VariableDef &n)\n { return n.metadata; })\n .def(\"with_no_outer_index\", &flaxlib::VariableDef::with_no_outer_index)\n .def(\"with_same_outer_index\", &flaxlib::VariableDef::with_same_outer_index)\n .def(\"__eq__\", &flaxlib::VariableDef::__eq__, nb::arg().none())\n .def(\"__hash__\", &flaxlib::VariableDef::__hash__)\n .def(\"__getstate__\", &flaxlib::VariableDef::__getstate__)\n .def(\"__setstate__\", &flaxlib::VariableDef::__setstate__);\n\n nb::class_(m, \"NodeRef\")\n .def(nb::init(),\n nb::arg(\"index\"))\n .def_prop_ro(\"index\", [](const flaxlib::NodeRef &n)\n { return n.index; })\n .def(\"__eq__\", &flaxlib::NodeRef::__eq__, nb::arg().none())\n .def(\"__hash__\", &flaxlib::NodeRef::__hash__)\n .def(\"__getstate__\", &flaxlib::NodeRef::__getstate__)\n .def(\"__setstate__\", &flaxlib::NodeRef::__setstate__);\n }\n} // namespace flaxlib" }, { "path": "nnx.py", "content": "# Copyright 2026 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Standalone nnx module shim.\n\nThis module allows importing nnx directly as a standalone module:\n\n import nnx\n\ninstead of:\n\n from flax import nnx\n\nBoth imports refer to the exact same module, ensuring that\n`nnx.Module` and `flax.nnx.Module` are the same class in memory.\n\"\"\"\n\nfrom flax.nnx import *\nfrom flax.version import __version__ as __version__\nfrom flax import nnx as _nnx\n\n# Re-export the module's metadata\n__all__: list[str] = _nnx.__all__ if hasattr(_nnx, '__all__') else []" }, { "path": "pylintrc", "content": "# This Pylint rcfile contains a best-effort configuration to uphold the\n# best-practices and style described in the Google Python style guide:\n# https://google.github.io/styleguide/pyguide.html\n#\n# Its canonical open-source location is:\n# https://google.github.io/styleguide/pylintrc\n\n[MASTER]\n\n# Add files or directories to the blacklist. They should be base names, not\n# paths.\nignore=third_party\n\n# Add files or directories matching the regex patterns to the blacklist. The\n# regex matches against base names, not paths.\nignore-patterns=\n\n# Pickle collected data for later comparisons.\npersistent=no\n\n# List of plugins (as comma separated values of python modules names) to load,\n# usually to register additional checkers.\nload-plugins=\n\n# Use multiple processes to speed up Pylint.\njobs=4\n\n# Allow loading of arbitrary C extensions. Extensions are imported into the\n# active Python interpreter and may run arbitrary code.\nunsafe-load-any-extension=no\n\n# A comma-separated list of package or module names from where C extensions may\n# be loaded. Extensions are loading into the active Python interpreter and may\n# run arbitrary code\nextension-pkg-whitelist=\n\n\n[MESSAGES CONTROL]\n\n# Only show warnings with the listed confidence levels. Leave empty to show\n# all. Valid levels: HIGH, INFERENCE, INFERENCE_FAILURE, UNDEFINED\nconfidence=\n\n# Enable the message, report, category or checker with the given id(s). You can\n# either give multiple identifier separated by comma (,) or put this option\n# multiple time (only on the command line, not in the configuration file where\n# it should appear only once). See also the \"--disable\" option for examples.\n#enable=\n\n# Disable the message, report, category or checker with the given id(s). You\n# can either give multiple identifiers separated by comma (,) or put this\n# option multiple times (only on the command line, not in the configuration\n# file where it should appear only once).You can also use \"--disable=all\" to\n# disable everything first and then reenable specific checks. For example, if\n# you want to run only the similarities checker, you can use \"--disable=all\n# --enable=similarities\". If you want to run only the classes checker, but have\n# no Warning level messages displayed, use\"--disable=all --enable=classes\n# --disable=W\"\ndisable=apply-builtin,\n backtick,\n bad-option-value,\n buffer-builtin,\n c-extension-no-member,\n cmp-builtin,\n cmp-method,\n coerce-builtin,\n coerce-method,\n delslice-method,\n div-method,\n duplicate-code,\n eq-without-hash,\n execfile-builtin,\n file-builtin,\n filter-builtin-not-iterating,\n fixme,\n getslice-method,\n global-statement,\n hex-method,\n idiv-method,\n implicit-str-concat-in-sequence,\n import-error,\n import-self,\n import-star-module-level,\n input-builtin,\n intern-builtin,\n invalid-str-codec,\n locally-disabled,\n long-builtin,\n long-suffix,\n map-builtin-not-iterating,\n metaclass-assignment,\n next-method-called,\n next-method-defined,\n no-absolute-import,\n no-else-break,\n no-else-continue,\n no-else-raise,\n no-else-return,\n no-member,\n no-self-use,\n nonzero-method,\n oct-method,\n old-division,\n old-ne-operator,\n old-octal-literal,\n old-raise-syntax,\n parameter-unpacking,\n print-statement,\n raising-string,\n range-builtin-not-iterating,\n raw_input-builtin,\n rdiv-method,\n reduce-builtin,\n relative-import,\n reload-builtin,\n round-builtin,\n setslice-method,\n signature-differs,\n standarderror-builtin,\n suppressed-message,\n sys-max-int,\n too-few-public-methods,\n too-many-ancestors,\n too-many-arguments,\n too-many-boolean-expressions,\n too-many-branches,\n too-many-instance-attributes,\n too-many-locals,\n too-many-public-methods,\n too-many-return-statements,\n too-many-statements,\n trailing-newlines,\n unichr-builtin,\n unicode-builtin,\n unpacking-in-except,\n useless-else-on-loop,\n useless-suppression,\n using-cmp-argument,\n xrange-builtin,\n zip-builtin-not-iterating,\n\n\n[REPORTS]\n\n# Set the output format. Available formats are text, parseable, colorized, msvs\n# (visual studio) and html. You can also give a reporter class, eg\n# mypackage.mymodule.MyReporterClass.\noutput-format=text\n\n# Put messages in a separate file for each module / package specified on the\n# command line instead of printing them on stdout. Reports (if any) will be\n# written in a file name \"pylint_global.[txt|html]\". This option is deprecated\n# and it will be removed in Pylint 2.0.\nfiles-output=no\n\n# Tells whether to display a full report or only the messages\nreports=no\n\n# Python expression which should return a note less than 10 (10 is the highest\n# note). You have access to the variables errors warning, statement which\n# respectively contain the number of errors / warnings messages and the total\n# number of statements analyzed. This is used by the global evaluation report\n# (RP0004).\nevaluation=10.0 - ((float(5 * error + warning + refactor + convention) / statement) * 10)\n\n# Template used to display messages. This is a python new-style format string\n# used to format the message information. See doc for all details\n#msg-template=\n\n\n[BASIC]\n\n# Good variable names which should always be accepted, separated by a comma\ngood-names=main,_\n\n# Bad variable names which should always be refused, separated by a comma\nbad-names=\n\n# Colon-delimited sets of names that determine each other's naming style when\n# the name regexes allow several styles.\nname-group=\n\n# Include a hint for the correct naming format with invalid-name\ninclude-naming-hint=no\n\n# List of decorators that produce properties, such as abc.abstractproperty. Add\n# to this list to register other decorators that produce valid properties.\nproperty-classes=abc.abstractproperty,cached_property.cached_property,cached_property.threaded_cached_property,cached_property.cached_property_with_ttl,cached_property.threaded_cached_property_with_ttl\n\n# Regular expression matching correct function names\nfunction-rgx=^(?:(?PsetUp|tearDown|setUpModule|tearDownModule)|(?P_?[A-Z][a-zA-Z0-9]*)|(?P_?[a-z][a-z0-9_]*))$\n\n# Regular expression matching correct variable names\nvariable-rgx=^[a-z][a-z0-9_]*$\n\n# Regular expression matching correct constant names\nconst-rgx=^(_?[A-Z][A-Z0-9_]*|__[a-z0-9_]+__|_?[a-z][a-z0-9_]*)$\n\n# Regular expression matching correct attribute names\nattr-rgx=^_{0,2}[a-z][a-z0-9_]*$\n\n# Regular expression matching correct argument names\nargument-rgx=^[a-z][a-z0-9_]*$\n\n# Regular expression matching correct class attribute names\nclass-attribute-rgx=^(_?[A-Z][A-Z0-9_]*|__[a-z0-9_]+__|_?[a-z][a-z0-9_]*)$\n\n# Regular expression matching correct inline iteration names\ninlinevar-rgx=^[a-z][a-z0-9_]*$\n\n# Regular expression matching correct class names\nclass-rgx=^_?[A-Z][a-zA-Z0-9]*$\n\n# Regular expression matching correct module names\nmodule-rgx=^(_?[a-z][a-z0-9_]*|__init__)$\n\n# Regular expression matching correct method names\nmethod-rgx=(?x)^(?:(?P_[a-z0-9_]+__|runTest|setUp|tearDown|setUpTestCase|tearDownTestCase|setupSelf|tearDownClass|setUpClass|(test|assert)_*[A-Z0-9][a-zA-Z0-9_]*|next)|(?P_{0,2}[A-Z][a-zA-Z0-9_]*)|(?P_{0,2}[a-z][a-z0-9_]*))$\n\n# Regular expression which should only match function or class names that do\n# not require a docstring.\nno-docstring-rgx=(__.*__|main|test.*|.*test|.*Test)$\n\n# Minimum line length for functions/classes that require docstrings, shorter\n# ones are exempt.\ndocstring-min-length=10\n\n\n[TYPECHECK]\n\n# List of decorators that produce context managers, such as\n# contextlib.contextmanager. Add to this list to register other decorators that\n# produce valid context managers.\ncontextmanager-decorators=contextlib.contextmanager,contextlib2.contextmanager\n\n# Tells whether missing members accessed in mixin class should be ignored. A\n# mixin class is detected if its name ends with \"mixin\" (case insensitive).\nignore-mixin-members=yes\n\n# List of module names for which member attributes should not be checked\n# (useful for modules/projects where namespaces are manipulated during runtime\n# and thus existing member attributes cannot be deduced by static analysis. It\n# supports qualified module names, as well as Unix pattern matching.\nignored-modules=\n\n# List of class names for which member attributes should not be checked (useful\n# for classes with dynamically set attributes). This supports the use of\n# qualified names.\nignored-classes=optparse.Values,thread._local,_thread._local\n\n# List of members which are set dynamically and missed by pylint inference\n# system, and so shouldn't trigger E1101 when accessed. Python regular\n# expressions are accepted.\ngenerated-members=\n\n\n[FORMAT]\n\n# Maximum number of characters on a single line.\nmax-line-length=80\n\n# TODO(https://github.com/PyCQA/pylint/issues/3352): Direct pylint to exempt\n# lines made too long by directives to pytype.\n\n# Regexp for a line that is allowed to be longer than the limit.\nignore-long-lines=(?x)(\n ^\\s*(\\#\\ )??$|\n ^\\s*(from\\s+\\S+\\s+)?import\\s+.+$)\n\n# Allow the body of an if to be on the same line as the test if there is no\n# else.\nsingle-line-if-stmt=yes\n\n# List of optional constructs for which whitespace checking is disabled. `dict-\n# separator` is used to allow tabulation in dicts, etc.: {1 : 1,\\n222: 2}.\n# `trailing-comma` allows a space between comma and closing bracket: (a, ).\n# `empty-line` allows space-only lines.\nno-space-check=\n\n# Maximum number of lines in a module\nmax-module-lines=99999\n\n# String used as indentation unit. The internal Google style guide mandates 2\n# spaces. Google's externaly-published style guide says 4, consistent with\n# PEP 8. Here, we use 2 spaces, for conformity with many open-sourced Google\n# projects (like TensorFlow).\nindent-string=' '\n\n# Number of spaces of indent required inside a hanging or continued line.\nindent-after-paren=4\n\n# Expected format of line ending, e.g. empty (any line ending), LF or CRLF.\nexpected-line-ending-format=\n\n\n[MISCELLANEOUS]\n\n# List of note tags to take in consideration, separated by a comma.\nnotes=TODO\n\n\n[VARIABLES]\n\n# Tells whether we should check for unused import in __init__ files.\ninit-import=no\n\n# A regular expression matching the name of dummy variables (i.e. expectedly\n# not used).\ndummy-variables-rgx=^\\*{0,2}(_$|unused_|dummy_)\n\n# List of additional names supposed to be defined in builtins. Remember that\n# you should avoid to define new builtins when possible.\nadditional-builtins=\n\n# List of strings which can identify a callback function by name. A callback\n# name must start or end with one of those strings.\ncallbacks=cb_,_cb\n\n# List of qualified module names which can have objects that can redefine\n# builtins.\nredefining-builtins-modules=six,six.moves,past.builtins,future.builtins,functools\n\n\n[LOGGING]\n\n# Logging modules to check that the string format arguments are in logging\n# function parameter format\nlogging-modules=logging,absl.logging\n\n\n[SIMILARITIES]\n\n# Minimum lines number of a similarity.\nmin-similarity-lines=4\n\n# Ignore comments when computing similarities.\nignore-comments=yes\n\n# Ignore docstrings when computing similarities.\nignore-docstrings=yes\n\n# Ignore imports when computing similarities.\nignore-imports=no\n\n\n[SPELLING]\n\n# Spelling dictionary name. Available dictionaries: none. To make it working\n# install python-enchant package.\nspelling-dict=\n\n# List of comma separated words that should not be checked.\nspelling-ignore-words=\n\n# A path to a file that contains private dictionary; one word per line.\nspelling-private-dict-file=\n\n# Tells whether to store unknown words to indicated private dictionary in\n# --spelling-private-dict-file option instead of raising a message.\nspelling-store-unknown-words=no\n\n\n[IMPORTS]\n\n# Deprecated modules which should not be used, separated by a comma\ndeprecated-modules=regsub,\n TERMIOS,\n Bastion,\n rexec,\n sets\n\n# Create a graph of every (i.e. internal and external) dependencies in the\n# given file (report RP0402 must not be disabled)\nimport-graph=\n\n# Create a graph of external dependencies in the given file (report RP0402 must\n# not be disabled)\next-import-graph=\n\n# Create a graph of internal dependencies in the given file (report RP0402 must\n# not be disabled)\nint-import-graph=\n\n# Force import order to recognize a module as part of the standard\n# compatibility libraries.\nknown-standard-library=\n\n# Force import order to recognize a module as part of a third party library.\nknown-third-party=enchant, absl\n\n# Analyse import fallback blocks. This can be used to support both Python 2 and\n# 3 compatible code, which means that the block might have code that exists\n# only in one or another interpreter, leading to false positives when analysed.\nanalyse-fallback-blocks=no\n\n\n[CLASSES]\n\n# List of method names used to declare (i.e. assign) instance attributes.\ndefining-attr-methods=__init__,\n __new__,\n setUp\n\n# List of member names, which should be excluded from the protected access\n# warning.\nexclude-protected=_asdict,\n _fields,\n _replace,\n _source,\n _make\n\n# List of valid names for the first argument in a class method.\nvalid-classmethod-first-arg=cls,\n class_\n\n# List of valid names for the first argument in a metaclass class method.\nvalid-metaclass-classmethod-first-arg=mcs\n\n\n[EXCEPTIONS]\n\n# Exceptions that will emit a warning when being caught. Defaults to\n# \"Exception\"\novergeneral-exceptions=StandardError,\n Exception,\n BaseException\n" }, { "path": "pyproject.toml", "content": "[build-system]\nrequires = [\"setuptools\", \"setuptools-scm\"]\nbuild-backend = \"setuptools.build_meta\"\n\n[project]\nname = \"flax\"\nrequires-python = \">=3.11\"\ndescription = \"Flax: A neural network library for JAX designed for flexibility\"\nkeywords = []\nauthors = [\n {name = \"Flax team\", email = \"flax-dev@google.com\"},\n]\ndependencies = [\n \"numpy>=1.23.2\",\n # keep in sync with jax-version in .github/workflows/build.yml\n \"jax>=0.8.1\",\n \"msgpack\",\n \"optax\",\n \"orbax-checkpoint\",\n \"tensorstore\",\n \"rich>=11.1\",\n \"typing_extensions>=4.2\",\n \"PyYAML>=5.4.1\",\n \"treescope>=0.1.7\",\n \"orbax-export>=0.0.8\",\n]\nclassifiers = [\n \"Development Status :: 3 - Alpha\",\n \"Intended Audience :: Developers\",\n \"Intended Audience :: Science/Research\",\n \"License :: OSI Approved :: Apache Software License\",\n \"Programming Language :: Python :: 3.11\",\n \"Programming Language :: Python :: 3.12\",\n \"Topic :: Scientific/Engineering :: Artificial Intelligence\",\n]\ndynamic = [\"version\", \"readme\"]\n\n[project.optional-dependencies]\ntesting = [\n \"clu\",\n \"clu<=0.0.9; python_version<'3.10'\",\n \"einops\",\n \"gymnasium[atari]; python_version<'3.14'\",\n \"jaxlib\",\n \"jaxtyping\",\n \"jraph>=0.0.6dev0\",\n \"ml-collections\",\n \"mypy\",\n \"opencv-python\",\n # Set protobuf version for python 3.13+ to prevent error in\n # examples/mnist/train_test.py::TrainTest::test_train_and_evaluate\n # Failed to construct dataset \"mnist\", builder_kwargs \"{}\": Value out of range: 11594722\n \"protobuf<6; python_version>='3.13'\",\n \"pytest\",\n \"pytest-cov\",\n \"pytest-custom_exit_code\",\n \"pytest-xdist\",\n \"pytype\",\n # WMT/LM1B examples\n \"sentencepiece==0.2.0\", # Segfault bug in 0.2.1\n \"tensorflow_text>=2.11.0; platform_system!='Darwin' and python_version < '3.13'\",\n \"tensorflow_datasets\",\n \"tensorflow>=2.12.0; python_version<'3.13'\", # to fix Numpy np.bool8 deprecation error\n \"tensorflow>=2.20.0; python_version>='3.13' and python_version<'3.14'\",\n # Temporary fix for https://github.com/google/flax/issues/5143\n \"keras<3.13\",\n \"torch\",\n \"treescope>=0.1.1; python_version>='3.10'\",\n \"cloudpickle>=3.0.0\",\n \"ale-py>=0.10.2; python_version<'3.14'\",\n]\ndocs = [\n \"sphinx==6.2.1\",\n \"sphinx-book-theme\",\n \"Pygments>=2.6.1\",\n \"ipykernel\",\n \"tqdm==4.67.1\",\n \"myst_nb\",\n \"nbstripout\",\n \"recommonmark\",\n \"ipython_genutils\",\n \"sphinx-design\",\n \"jupytext==1.13.8\",\n \"dm-haiku>=0.0.14\",\n \"docutils\",\n # The next packages are for notebooks.\n \"matplotlib\",\n \"scikit-learn\",\n # The next packages are used in testcode blocks.\n \"ml_collections\",\n # notebooks\n \"einops\",\n \"kagglehub>=0.3.3\",\n \"ipywidgets>=8.1.5\",\n]\ndev = [\n \"nanobind>=2.5.0\",\n \"pre-commit>=3.8.0\",\n \"scikit-build-core[pyproject]>=0.11.0\",\n]\n\n[project.urls]\nhomepage = \"https://github.com/google/flax\"\n\n[tool.setuptools.dynamic]\nreadme = {file = [\"README.md\"], content-type = \"text/markdown\"}\nversion = {attr = \"flax.version.__version__\"}\n\n[tool.setuptools.packages.find]\ninclude = [\"flax*\"]\n\n[tool.setuptools.package-data]\nflax = [\"*py.typed\"]\n\n[tool.yapf]\nbased_on_style = \"yapf\"\n\n[tool.pytype]\n# TODO(levskaya): figure out why we get pyi-error from flax's root __init__.py\n# could be a pytype bug.\ndisable = \"pyi-error\"\n\n[tool.mypy]\nshow_error_codes = true\nno_implicit_optional = true\ndisable_error_code = \"attr-defined\"\n\n[[tool.mypy.overrides]]\nmodule = [\n \"tensorflow.*\",\n \"tensorboard.*\",\n \"absl.*\",\n \"jax.*\",\n \"rich.*\",\n \"flax.*\",\n \"jaxlib.cuda.*\",\n \"jaxlib.cpu.*\",\n \"msgpack\",\n \"numpy.*\",\n \"optax.*\",\n \"orbax.*\",\n \"opt_einsum.*\",\n \"scipy.*\",\n \"libtpu.*\",\n \"jaxlib.mlir.*\",\n \"yaml\",\n]\nignore_missing_imports = true\ndisable_error_code = \"annotation-unchecked\"\n# exclude nnx examples\n[[tool.mypy.overrides]]\nmodule = \"flax.nnx.examples.*\"\nignore_errors = true\n\n[tool.pytest.ini_options]\nfilterwarnings = [\n # By default error out on any warnings.\n \"error\",\n # Jax warning when no gpu/tpu found.\n\t\"ignore:No GPU/TPU found, falling back to CPU.*:UserWarning\",\n # traverse_util.Traversal will be removed soon.\n \"ignore:`flax.traverse_util.Traversal` will be deprecated.*:DeprecationWarning\",\n # Deprecated legacy checkpoint - just want to keep the tests running for a while\n \"ignore:Flax Checkpointing will soon be deprecated in favor of Orbax.*:DeprecationWarning\",\n # DeprecationWarning: The inputs_kv arg will be deprecated soon. Use inputs_k and inputs_v instead.\n \"ignore:.*The inputs_kv arg will be deprecated soon. Use inputs_k and inputs_v instead.*:DeprecationWarning\",\n # DeprecationWarning: the function signature of MultiHeadDotProductAttention's `__call__` method has changed\n \"ignore:.*the function signature of MultiHeadDotProductAttention's `__call__` method has changed.*:DeprecationWarning\",\n # DeprecationWarning: ml_dtypes.float8_e4m3b11 is deprecated.\n \"ignore:.*ml_dtypes.float8_e4m3b11 is deprecated.*:DeprecationWarning\",\n # pytest-cov uses a deprecated feature of pytest-xdist. (2023-11-06)\n \"ignore:The --rsyncdir command line argument and rsyncdirs config variable are deprecated.:DeprecationWarning\",\n # DeprecationWarning: jax.random.KeyArray is deprecated.\n \"ignore:.*jax.random.KeyArray is deprecated.*:DeprecationWarning\",\n # DeprecationWarning: jax.core.Shape is deprecated.\n \"ignore:.*jax.core.Shape is deprecated.*:DeprecationWarning\",\n # DeprecationWarning: pkg_resources is deprecated as an API.\n \"ignore:.*pkg_resources is deprecated as an API.*:DeprecationWarning\",\n # DeprecationWarning: Deprecated call to `pkg_resources.declare_namespace('google')`.\n \"ignore:.*Deprecated call to.*pkg_resources.declare_namespace.*:DeprecationWarning\",\n # jax.xla_computation is deprecated but TF still uses it.\n \"ignore:.*jax.xla_computation is deprecated.*:DeprecationWarning\",\n # Orbax warnings inside deprecated `flax.training` package.\n \"ignore:.*Couldn't find sharding info under RestoreArgs.*:UserWarning\",\n # RuntimeWarning: invalid value encountered in cast\n \"ignore:.*invalid value encountered in cast.*:RuntimeWarning\",\n # RuntimeWarning: divide by zero encountered in equal/not_equal\n \"ignore:.*divide by zero encountered in.*:RuntimeWarning\",\n # DeprecationWarning: numpy.core is deprecated\n \"ignore:.*numpy.core is deprecated.*:DeprecationWarning\",\n # DeprecationWarning: shape requires ndarray or scalar arguments\n \"ignore:.*shape requires ndarray or scalar arguments.*:DeprecationWarning\",\n # UserWarning: Sharding info not provided when restoring\n \"ignore:.*Sharding info not provided when restoring.*:UserWarning\",\n # UserWarning: pkg_resources is deprecated as an API.\n \"ignore:.*pkg_resources is deprecated as an API.*:UserWarning\",\n # DeprecationWarning: 'Data' is deprecated, please replace:\n \"ignore:.*[Data|Static]' is deprecated, please replace.*:DeprecationWarning\",\n # DeprecationWarning: Implicit conversion of an array to a dtype is deprecated; rather than dtype=arr use dtype=arr.dtype.\n \"ignore:.*Implicit conversion of an array to a dtype is deprecated; rather than dtype=arr use dtype=arr.dtype.*:DeprecationWarning\",\n # DeprecationWarning: Setting `jax_pmap_shmap_merge` is deprecated in JAX v0.9.0 and will be removed in JAX v0.10.0\n \"ignore:.*Setting `jax_pmap_shmap_merge` is deprecated in JAX.*:DeprecationWarning\",\n]\n\n[tool.coverage.report]\nexclude_lines = [\n \"@abc.abstractmethod\",\n \"raise NotImplementedError\",\n]\n\n[tool.pyink]\npyink-indentation = 2\npyink-use-majority-quotes = true\nline-length = 80\npreview = true\n\n[tool.ruff]\n# Exclude a variety of commonly ignored directories.\nexclude = [\n \"__init__.py\",\n \"activation.py\",\n \"partitioning.py\",\n \"flax/core/variables.py\",\n \"examples/\",\n]\n\nline-length = 80\nindent-width = 2\n\n[tool.ruff.lint]\n# Enable Pyflakes (`F`) and a subset of the pycodestyle (`E`) codes by default.\nselect = [\"F401\"]\nignore = []\n# Allow fix for all enabled rules (when `--fix`) is provided.\n# Full list of rules: https://docs.astral.sh/ruff/rules/\nfixable = [\"ALL\"]\nunfixable = []\n\n[tool.ruff.format]\nindent-style = \"space\"\nquote-style = \"single\"\n\n\n[tool.uv]\n# Ignore uv.lock and always upgrade the package to the latest\nupgrade-package = [\"jax\", \"jaxlib\", \"orbax-checkpoint\"]\n" }, { "path": "tests/checkpoints_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for flax.training.checkpoints.\"\"\"\n\nimport copy\nimport os\nimport pathlib\nfrom typing import Any\n\nimport jax\nimport numpy as np\nimport orbax.checkpoint as orbax\nfrom absl.testing import absltest, parameterized\nfrom jax import numpy as jnp\n\nfrom flax import config, core, errors, io, struct\nfrom flax import linen as nn\nfrom flax.training import checkpoints\n\n# Parse absl flags test_srcdir and test_tmpdir.\njax.config.parse_flags_with_absl()\n\n\nPyTree = Any\n\n\ndef check_eq(xs, ys):\n return jax.tree_util.tree_all(\n jax.tree_util.tree_map(np.testing.assert_allclose, xs, ys)\n )\n\n\ndef shuffle(l):\n \"\"\"Functional shuffle.\"\"\"\n l = copy.copy(l)\n np.random.shuffle(l)\n return l\n\n\nclass Inner(nn.Module):\n \"\"\"Inner class based on nn.\"\"\"\n\n @nn.compact\n def __call__(self, x):\n x = nn.Conv(10, (2, 2))(x)\n x = nn.normalization.BatchNorm(True)(x)\n return x\n\n\nclass Model(nn.Module):\n \"\"\"Simple model based on nn.\"\"\"\n\n @nn.compact\n def __call__(self, inputs):\n x = nn.Conv(10, (2, 2))(inputs)\n x = Inner()(x)\n x = x.reshape([x.shape[0], -1])\n x = nn.normalization.BatchNorm(True)(x)\n x = nn.Dense(10)(x)\n x = nn.log_softmax(x)\n return x\n\n\n@struct.dataclass\nclass CustomDC:\n foo: Any\n bar: Any\n\n\nclass CheckpointsTest(parameterized.TestCase):\n def setUp(self):\n super().setUp()\n config.update('flax_use_orbax_checkpointing', False) # default value\n\n def test_naturalsort(self):\n np.random.seed(0)\n tests = [\n ['file_1', 'file_2', 'file_10', 'file_11', 'file_21'],\n ['file_0.001', 'file_0.01', 'file_0.1', 'file_1'],\n ['file_-3.0', 'file_-2', 'file_-1', 'file_0.0'],\n ['file_1e1', 'file_1.0e2', 'file_1e3', 'file_1.0e4'],\n ['file_1', 'file_2', 'file_9', 'file_1.0e1', 'file_11'],\n ]\n for test in tests:\n self.assertEqual(test, checkpoints.natural_sort(shuffle(test)))\n\n def test_safe_normpath(self):\n tests = ['./a/b/c', '/a//b/c', '/a/../b/c', 'a/b/./c', 'gs://a//b/c']\n expected = ['a/b/c', '/a/b/c', '/b/c', 'a/b/c', 'gs://a/b/c']\n for test, expect in zip(tests, expected):\n self.assertEqual(expect, checkpoints.safe_normpath(test))\n\n @parameterized.parameters({'use_orbax': True}, {'use_orbax': False})\n def test_save_restore_checkpoints(self, use_orbax):\n config.update('flax_use_orbax_checkpointing', use_orbax)\n tmp_dir = pathlib.Path(self.create_tempdir().full_path)\n test_object0 = {\n 'a': np.array([0, 0, 0], np.int32),\n 'b': np.array([0, 0, 0], np.int32),\n }\n test_object1 = {\n 'a': np.array([1, 2, 3], np.int32),\n 'b': np.array([1, 1, 1], np.int32),\n }\n test_object2 = {\n 'a': np.array([4, 5, 6], np.int32),\n 'b': np.array([2, 2, 2], np.int32),\n }\n new_object = checkpoints.restore_checkpoint(\n tmp_dir, test_object0, prefix='test_'\n )\n check_eq(new_object, test_object0)\n checkpoints.save_checkpoint(\n tmp_dir, test_object1, 0, prefix='test_', keep=1\n )\n self.assertIn('test_0', os.listdir(tmp_dir))\n new_object = checkpoints.restore_checkpoint(\n tmp_dir, test_object0, prefix='test_'\n )\n check_eq(new_object, test_object1)\n checkpoints.save_checkpoint(\n tmp_dir, test_object1, 1, prefix='test_', keep=1\n )\n checkpoints.save_checkpoint(\n tmp_dir, test_object2, 2, prefix='test_', keep=1\n )\n new_object = checkpoints.restore_checkpoint(\n tmp_dir, test_object0, prefix='test_'\n )\n check_eq(new_object, test_object2)\n checkpoints.save_checkpoint(\n tmp_dir, test_object2, 3, prefix='test_', keep=2\n )\n checkpoints.save_checkpoint(\n tmp_dir, test_object1, 4, prefix='test_', keep=2\n )\n new_object = checkpoints.restore_checkpoint(\n tmp_dir, test_object0, prefix='test_'\n )\n check_eq(new_object, test_object1)\n new_object = checkpoints.restore_checkpoint(\n tmp_dir, test_object0, step=3, prefix='test_'\n )\n check_eq(new_object, test_object2)\n\n # Restore with a specific checkpoint path, not the directory path.\n new_object = checkpoints.restore_checkpoint(\n os.path.join(tmp_dir, 'test_3'), test_object0\n )\n check_eq(new_object, test_object2)\n # If a specific path is specified, but it does not exist, the same behavior\n # as when a directory is empty should apply: the target is returned\n # unchanged.\n new_object = checkpoints.restore_checkpoint(\n os.path.join(tmp_dir, 'test_not_there'), test_object0\n )\n check_eq(new_object, test_object0)\n with self.assertRaises(ValueError):\n checkpoints.restore_checkpoint(\n tmp_dir, test_object0, step=5, prefix='test_'\n )\n\n @parameterized.parameters({'use_orbax': True}, {'use_orbax': False})\n def test_overwrite_checkpoints(self, use_orbax):\n config.update('flax_use_orbax_checkpointing', use_orbax)\n overwrite_error = ValueError if use_orbax else errors.InvalidCheckpointError\n tmp_dir = self.create_tempdir().full_path\n test_object0 = {'a': np.array([0, 0, 0], np.int32)}\n test_object = {'a': np.array([1, 2, 3], np.int32)}\n\n checkpoints.save_checkpoint(tmp_dir, test_object0, 0, keep=1)\n with self.assertRaises(overwrite_error):\n checkpoints.save_checkpoint(tmp_dir, test_object, 0, keep=1)\n checkpoints.save_checkpoint(tmp_dir, test_object, 0, keep=1, overwrite=True)\n new_object = checkpoints.restore_checkpoint(tmp_dir, test_object0)\n check_eq(new_object, test_object)\n\n non_norm_dir_path = tmp_dir + '//'\n checkpoints.save_checkpoint(non_norm_dir_path, test_object, 4, keep=1)\n new_object = checkpoints.restore_checkpoint(non_norm_dir_path, test_object0)\n check_eq(new_object, test_object)\n\n @parameterized.parameters(\n {'use_orbax': True, 'keep_every_n_steps': None},\n {'use_orbax': False, 'keep_every_n_steps': 7},\n )\n def test_keep(self, use_orbax, keep_every_n_steps):\n config.update('flax_use_orbax_checkpointing', use_orbax)\n tmp_dir = self.create_tempdir().full_path\n test_object = {'a': np.array([1, 2, 3], np.int32)}\n steps_start = 17\n steps_end = 37\n keep = 3\n increment = 5\n\n for step in range(steps_start, steps_end, increment):\n checkpoints.save_checkpoint(\n tmp_dir,\n test_object,\n step=step,\n keep=keep,\n keep_every_n_steps=keep_every_n_steps,\n )\n\n last_checkpoint = -float('inf')\n for step in range(steps_start, steps_end, increment):\n if ((steps_end - step) / increment <= keep) or (\n keep_every_n_steps and (step - last_checkpoint) >= keep_every_n_steps\n ):\n restored = checkpoints.restore_checkpoint(\n tmp_dir, target=None, step=step\n )\n check_eq(restored, test_object)\n last_checkpoint = step\n else:\n with self.assertRaises(ValueError):\n checkpoints.restore_checkpoint(tmp_dir, target=None, step=step)\n\n @parameterized.parameters({'use_orbax': True}, {'use_orbax': False})\n def test_save_restore_checkpoints_w_float_steps(self, use_orbax):\n config.update('flax_use_orbax_checkpointing', use_orbax)\n tmp_dir = self.create_tempdir().full_path\n test_object0 = {\n 'a': np.array([0, 0, 0], np.int32),\n 'b': np.array([0, 0, 0], np.int32),\n }\n test_object1 = {\n 'a': np.array([1, 2, 3], np.int32),\n 'b': np.array([1, 1, 1], np.int32),\n }\n test_object2 = {\n 'a': np.array([4, 5, 6], np.int32),\n 'b': np.array([2, 2, 2], np.int32),\n }\n checkpoints.save_checkpoint(\n tmp_dir, test_object1, 0.0, prefix='test_', keep=1\n )\n self.assertIn('test_0.0', os.listdir(tmp_dir))\n new_object = checkpoints.restore_checkpoint(\n tmp_dir, test_object0, prefix='test_'\n )\n check_eq(new_object, test_object1)\n checkpoints.save_checkpoint(\n tmp_dir, test_object1, 2.0, prefix='test_', keep=1\n )\n checkpoints.save_checkpoint(\n tmp_dir, test_object2, 3.0, prefix='test_', keep=2\n )\n self.assertIn('test_3.0', os.listdir(tmp_dir))\n self.assertIn('test_2.0', os.listdir(tmp_dir))\n check_eq(new_object, test_object1)\n\n @parameterized.parameters({'use_orbax': True}, {'use_orbax': False})\n def test_save_restore_checkpoints_target_none(self, use_orbax):\n config.update('flax_use_orbax_checkpointing', use_orbax)\n tmp_dir = self.create_tempdir().full_path\n test_object0 = {\n 'a': np.array([0, 0, 0], np.int32),\n 'b': np.array([0, 0, 0], np.int32),\n }\n # Target pytree is a dictionary, so it's equal to a restored state_dict.\n checkpoints.save_checkpoint(tmp_dir, test_object0, 0)\n new_object = checkpoints.restore_checkpoint(tmp_dir, target=None)\n check_eq(new_object, test_object0)\n # Target pytree it's a tuple, check the expected state_dict is recovered.\n test_object1 = (\n np.array([0, 0, 0], np.int32),\n np.array([1, 1, 1], np.int32),\n )\n checkpoints.save_checkpoint(tmp_dir, test_object1, 1)\n new_object = checkpoints.restore_checkpoint(tmp_dir, target=None)\n expected_new_object = {str(k): v for k, v in enumerate(test_object1)}\n check_eq(new_object, expected_new_object)\n\n @parameterized.parameters({'use_orbax': True}, {'use_orbax': False})\n def test_save_restore_checkpoints_target_singular(self, use_orbax):\n config.update('flax_use_orbax_checkpointing', use_orbax)\n tmp_dir = self.create_tempdir().full_path\n test_object0 = np.array([0, 0, 0], np.int32)\n test_object1 = np.array([1, 1, 1], np.int32)\n # Orbax backend returns error if target is singular. Orbax user need to use\n # ArrayCheckpointHandler instead.\n if use_orbax:\n with self.assertRaises(ValueError):\n checkpoints.save_checkpoint(tmp_dir, test_object1, 0)\n else:\n checkpoints.save_checkpoint(tmp_dir, test_object1, 0)\n new_object = checkpoints.restore_checkpoint(tmp_dir, target=None)\n check_eq(new_object, test_object1)\n checkpoints.save_checkpoint(tmp_dir, test_object0, 1)\n new_object = checkpoints.restore_checkpoint(tmp_dir, target=test_object1)\n check_eq(new_object, test_object0)\n\n @parameterized.parameters({'use_orbax': True}, {'use_orbax': False})\n def test_save_restore_checkpoints_target_empty(self, use_orbax):\n config.update('flax_use_orbax_checkpointing', use_orbax)\n tmp_dir = self.create_tempdir().full_path\n test_object0 = {}\n test_object1 = []\n # Orbax returns ValueError if the target is empty, but legacy Flax doesn't.\n if use_orbax:\n with self.assertRaises(ValueError):\n checkpoints.save_checkpoint(tmp_dir, test_object1, 0)\n else:\n checkpoints.save_checkpoint(tmp_dir, test_object1, 0)\n new_object = checkpoints.restore_checkpoint(tmp_dir, target=None)\n check_eq(new_object, test_object0)\n checkpoints.save_checkpoint(tmp_dir, test_object0, 1)\n new_object = checkpoints.restore_checkpoint(tmp_dir, target=test_object1)\n check_eq(new_object, test_object1)\n\n def test_async_save_checkpoints(self):\n tmp_dir = pathlib.Path(self.create_tempdir().full_path)\n test_object0 = {\n 'a': np.array([0, 0, 0], np.int32),\n 'b': np.array([0, 0, 0], np.int32),\n }\n test_object1 = {\n 'a': np.random.normal(size=(1000, 1000)),\n 'b': np.random.normal(size=(1000, 1000)),\n }\n test_object2 = {\n 'a': np.random.normal(size=(1000, 1000)),\n 'b': np.random.normal(size=(1000, 1000)),\n }\n test_object3 = {\n 'a': np.random.normal(size=(1000, 1000)),\n 'b': np.random.normal(size=(1000, 1000)),\n }\n am = checkpoints.AsyncManager()\n checkpoints.save_checkpoint(\n tmp_dir, test_object1, 0, prefix='test_', keep=1, async_manager=am\n )\n # Hard-wait the write to be done, then check its content.\n am.save_future.result()\n self.assertIn('test_0', os.listdir(tmp_dir))\n new_object = checkpoints.restore_checkpoint(\n tmp_dir, test_object1, prefix='test_'\n )\n check_eq(new_object, test_object1)\n # Check two consecutive saves happen in the right order.\n checkpoints.save_checkpoint(\n tmp_dir, test_object2, 1, prefix='test_', keep=1, async_manager=am\n )\n checkpoints.save_checkpoint(\n tmp_dir, test_object3, 2, prefix='test_', keep=1, async_manager=am\n )\n am.save_future.result()\n self.assertIn('test_2', os.listdir(tmp_dir))\n new_object = checkpoints.restore_checkpoint(\n tmp_dir, test_object1, prefix='test_'\n )\n check_eq(new_object, test_object3)\n\n def test_last_checkpoint(self):\n tmp_dir = pathlib.Path(self.create_tempdir().full_path)\n with io.GFile(os.path.join(tmp_dir, 'test_tmp'), 'w') as f:\n f.write('test_tmp')\n io.makedirs(os.path.join(tmp_dir, 'test_tmp_gda'))\n self.assertEqual(checkpoints.latest_checkpoint(tmp_dir, 'test_'), None)\n\n with io.GFile(os.path.join(tmp_dir, 'test_0'), 'w') as f:\n f.write('test_0')\n io.makedirs(os.path.join(tmp_dir, 'test_0_gda'))\n self.assertEqual(\n checkpoints.latest_checkpoint(tmp_dir, 'test_'),\n os.path.join(tmp_dir, 'test_0'),\n )\n\n with io.GFile(os.path.join(tmp_dir, 'test_10'), 'w') as f:\n f.write('test_10')\n self.assertEqual(\n checkpoints.latest_checkpoint(tmp_dir, 'test_'),\n os.path.join(tmp_dir, 'test_10'),\n )\n self.assertEqual(checkpoints.latest_checkpoint(tmp_dir, 'ckpt_'), None)\n\n path = f'orbaxtest_{orbax.utils.TMP_DIR_SUFFIX}_10'\n with io.GFile(os.path.join(tmp_dir, path), 'w') as f:\n f.write('orbaxtest_10')\n self.assertIsNone(checkpoints.latest_checkpoint(tmp_dir, 'orbaxtest_'))\n\n @parameterized.parameters(\n {'step_type': int, 'steps': [1, 5, 112]},\n {'step_type': float, 'steps': [1.0, 4.5, 5.6]},\n )\n def test_available_steps(self, step_type, steps):\n tmp_dir = pathlib.Path(self.create_tempdir().full_path)\n with io.GFile(os.path.join(tmp_dir, 'test_tmp'), 'w') as f:\n f.write('test_tmp')\n io.makedirs(os.path.join(tmp_dir, 'test_tmp_gda'))\n\n for step in steps:\n with io.GFile(os.path.join(tmp_dir, 'test_' + str(step)), 'w') as f:\n f.write('test_' + str(step))\n io.makedirs(os.path.join(tmp_dir, 'test_' + str(step) + '_gda'))\n\n self.assertEqual(\n checkpoints.available_steps(tmp_dir, 'test_', step_type=step_type),\n steps,\n )\n\n @parameterized.parameters({'use_orbax': True}, {'use_orbax': False})\n def test_complex_pytree(self, use_orbax):\n config.update('flax_use_orbax_checkpointing', use_orbax)\n tmp_dir = self.create_tempdir().full_path\n to_save = [\n CustomDC(foo=12, bar=core.freeze({'x': jnp.array((1, 4))})),\n np.array((2, 3)),\n ]\n target = [\n CustomDC(foo=0, bar=core.freeze({'x': jnp.array((0, 0))})),\n np.array((0, 0)),\n ]\n checkpoints.save_checkpoint(tmp_dir, to_save, 0)\n restored = checkpoints.restore_checkpoint(tmp_dir, target=target)\n check_eq(restored, to_save)\n\n # restore_checkpoint can automatically restore either orbax or legacy files.\n def test_auto_restore(self):\n tmp_dir = self.create_tempdir().full_path\n to_save = [CustomDC(foo=12, bar={'x': jnp.array((1, 4))}), np.array((2, 3))]\n target = [CustomDC(foo=0, bar={'x': jnp.array((0, 0))}), np.array((0, 0))]\n # Store an orbax ckpt\n config.update('flax_use_orbax_checkpointing', True)\n checkpoints.save_checkpoint(tmp_dir, to_save, 0, prefix='test_')\n # And a legacy ckpt\n config.update('flax_use_orbax_checkpointing', False)\n checkpoints.save_checkpoint(tmp_dir, to_save, 1, prefix='test_', keep=2)\n\n # Both gets restored with same API.\n restored = checkpoints.restore_checkpoint(\n os.path.join(tmp_dir, 'test_0'), target=target\n )\n check_eq(restored, to_save)\n restored = checkpoints.restore_checkpoint(\n os.path.join(tmp_dir, 'test_1'), target=target\n )\n check_eq(restored, to_save)\n\n @parameterized.parameters({'use_orbax': True}, {'use_orbax': False})\n def test_smaller_target(self, use_orbax):\n config.update('flax_use_orbax_checkpointing', use_orbax)\n tmp_dir = self.create_tempdir().full_path\n to_save = {'a': jnp.ones((16, 256, 1024))}\n target = {'a': jnp.zeros((2, 3))}\n\n checkpoints.save_checkpoint(tmp_dir, to_save, 0, keep=1)\n new_object = checkpoints.restore_checkpoint(tmp_dir, target)\n check_eq(new_object, to_save)\n\n def test_convert_pre_linen(self):\n params = checkpoints.convert_pre_linen(\n {\n 'mod_0': {\n 'submod1_0': {},\n 'submod2_1': {},\n 'submod1_2': {},\n },\n 'mod2_2': {'submod2_2_0': {}},\n 'mod2_11': {'submod2_11_0': {}},\n 'mod2_1': {'submod2_1_0': {}},\n }\n )\n self.assertDictEqual(\n core.unfreeze(params),\n {\n 'mod_0': {\n 'submod1_0': {},\n 'submod1_1': {},\n 'submod2_0': {},\n },\n 'mod2_0': {'submod2_1_0': {}},\n 'mod2_1': {'submod2_2_0': {}},\n 'mod2_2': {'submod2_11_0': {}},\n },\n )\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/colab_tpu_jax_version.ipynb", "content": "{\n \"cells\": [\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# JAX/jaxlib should be both 0.3.25\\n\",\n \"# because newer JAX versions are *not* supported on TPU runtimes\\n\",\n \"# Flax should be included in a ƒresh kernel.\\n\",\n \"!pip freeze | egrep 'jax|flax'\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# should show 8 TPU devices\\n\",\n \"import jax, jax.tools.colab_tpu\\n\",\n \"jax.tools.colab_tpu.setup_tpu()\\n\",\n \"jax.devices()\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# sometimes it's necessary to install additional packages; but we need to keep\\n\",\n \"# JAX/jaxlib versions pinned to what is supported by the TPU runtime...\\n\",\n \"!pip install jax==0.3.25 jaxlib==0.3.25 flax\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# in case JAX version has changed after the '!pip install`, below command should\\n\",\n \"# show the offending packages\\n\",\n \"!pip install -qq pipdeptree\\n\",\n \"!pipdeptree -w silence -r -p jax\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# it's possible to get dependency tree without installing packages, but this\\n\",\n \"# usually takes some 2-3 minutes...\\n\",\n \"!pip install -qq pipgrip\\n\",\n \"!pipgrip --tree flax==0.6.4\"\n ]\n }\n ],\n \"metadata\": {\n \"accelerator\": \"TPU\",\n \"gpuClass\": \"standard\",\n \"language_info\": {\n \"name\": \"python\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 0\n}\n" }, { "path": "tests/configurations_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom unittest import mock\n\nfrom absl.testing import absltest\n\nfrom flax.configurations import bool_flag, config\n\n\nclass MyTestCase(absltest.TestCase):\n def setUp(self):\n super().setUp()\n self.enter_context(mock.patch.object(config, '_values', {}))\n self._flag = bool_flag('flax_test', default=False, help='Just a test flag.')\n\n def test_duplicate_flag(self):\n with self.assertRaisesRegex(RuntimeError, 'already defined'):\n bool_flag(self._flag.name, default=False, help='Another test flag.')\n\n def test_default(self):\n self.assertFalse(self._flag.value)\n self.assertFalse(config.flax_test)\n\n def test_typed_update(self):\n config.update(self._flag, True)\n self.assertTrue(self._flag.value)\n self.assertTrue(config.flax_test)\n\n def test_untyped_update(self):\n config.update(self._flag.name, True)\n self.assertTrue(self._flag.value)\n self.assertTrue(config.flax_test)\n\n def test_update_unknown_flag(self):\n with self.assertRaisesRegex(LookupError, 'Unrecognized config option'):\n config.update('unknown', True)\n\n def test_temp_flip(self):\n self.assertFalse(self._flag.value)\n with config.temp_flip_flag('test', True):\n self.assertTrue(self._flag.value)\n self.assertFalse(self._flag.value)\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/core/core_frozen_dict_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport jax\nfrom absl.testing import absltest, parameterized\n\nfrom flax.core import FrozenDict, copy, freeze, pop, pretty_repr, unfreeze\n\n\nclass FrozenDictTest(parameterized.TestCase):\n def test_frozen_dict_copies(self):\n xs = {'a': 1, 'b': {'c': 2}}\n frozen = freeze(xs)\n xs['a'] += 1\n xs['b']['c'] += 1\n self.assertEqual(unfreeze(frozen), {'a': 1, 'b': {'c': 2}})\n\n def test_frozen_dict_maps(self):\n xs = {'a': 1, 'b': {'c': 2}}\n frozen = FrozenDict(xs)\n frozen2 = jax.tree_util.tree_map(lambda x: x + x, frozen)\n self.assertEqual(unfreeze(frozen2), {'a': 2, 'b': {'c': 4}})\n\n def test_frozen_dict_pop(self):\n xs = {'a': 1, 'b': {'c': 2}}\n b, a = FrozenDict(xs).pop('a')\n self.assertEqual(a, 1)\n self.assertEqual(unfreeze(b), {'b': {'c': 2}})\n\n def test_frozen_dict_partially_maps(self):\n x = jax.tree_util.tree_map(\n lambda a, b: (a, b), freeze({'a': 2}), freeze({'a': {'b': 1}})\n )\n self.assertEqual(unfreeze(x), {'a': (2, {'b': 1})})\n\n def test_frozen_dict_hash(self):\n xs = {'a': 1, 'b': {'c': 2}}\n ys = {'a': 1, 'b': {'c': 3}}\n self.assertNotEqual(hash(freeze(xs)), hash(freeze(ys)))\n\n def test_frozen_items(self):\n xs = {'a': 1, 'b': {'c': 2}}\n items = list(freeze(xs).items())\n\n self.assertEqual(items, [('a', 1), ('b', freeze(xs['b']))])\n\n def test_frozen_dict_repr(self):\n expected = \"\"\"FrozenDict({\n a: 1,\n b: {\n c: 2,\n d: {},\n },\n})\"\"\"\n\n xs = FrozenDict({'a': 1, 'b': {'c': 2, 'd': {}}})\n self.assertEqual(repr(xs), expected)\n self.assertEqual(repr(FrozenDict()), 'FrozenDict({})')\n\n def test_frozen_dict_reduce(self):\n before = FrozenDict(a=FrozenDict(b=1, c=2))\n cl, data = before.__reduce__()\n after = cl(*data)\n self.assertIsNot(before, after)\n self.assertEqual(before, after)\n self.assertEqual(after, {'a': {'b': 1, 'c': 2}})\n\n def test_frozen_dict_copy_reserved_name(self):\n result = FrozenDict({'a': 1}).copy({'cls': 2})\n self.assertEqual(result, {'a': 1, 'cls': 2})\n\n @parameterized.parameters(\n {\n 'x': {'a': 1, 'b': {'c': 2}},\n 'key': 'b',\n 'actual_new_x': {'a': 1},\n 'actual_value': {'c': 2},\n },\n {\n 'x': FrozenDict({'a': 1, 'b': {'c': 2}}),\n 'key': 'b',\n 'actual_new_x': FrozenDict({'a': 1}),\n 'actual_value': FrozenDict({'c': 2}),\n },\n )\n def test_utility_pop(self, x, key, actual_new_x, actual_value):\n new_x, value = pop(x, key)\n self.assertTrue(\n new_x == actual_new_x and isinstance(new_x, type(actual_new_x))\n )\n self.assertTrue(\n value == actual_value and isinstance(value, type(actual_value))\n )\n\n @parameterized.parameters(\n {\n 'x': {'a': 1, 'b': {'c': 2}},\n 'add_or_replace': {'b': {'c': -1, 'd': 3}},\n 'actual_new_x': {'a': 1, 'b': {'c': -1, 'd': 3}},\n },\n {\n 'x': FrozenDict({'a': 1, 'b': {'c': 2}}),\n 'add_or_replace': FrozenDict({'b': {'c': -1, 'd': 3}}),\n 'actual_new_x': FrozenDict({'a': 1, 'b': {'c': -1, 'd': 3}}),\n },\n )\n def test_utility_copy(self, x, add_or_replace, actual_new_x):\n new_x = copy(x, add_or_replace=add_or_replace)\n self.assertTrue(\n new_x == actual_new_x and isinstance(new_x, type(actual_new_x))\n )\n\n @parameterized.parameters(\n {\n 'x': {'a': 1, 'b': {'c': 2}},\n },\n {\n 'x': FrozenDict({'a': 1, 'b': {'c': 2}}),\n },\n )\n def test_utility_copy_singlearg(self, x):\n new_x = copy(x)\n self.assertTrue(new_x == x and isinstance(new_x, type(x)))\n\n @parameterized.parameters(\n {\n 'x': {'a': 1, 'b': {'c': 2}},\n 'pretty_str': '{\\n a: 1,\\n b: {\\n c: 2,\\n },\\n}',\n },\n {\n 'x': FrozenDict({'a': 1, 'b': {'c': 2}}),\n 'pretty_str': (\n 'FrozenDict({\\n a: 1,\\n b: {\\n c: 2,\\n },\\n})'\n ),\n },\n {\n 'x': 345,\n 'pretty_str': '345',\n },\n )\n def test_utility_pretty_repr(self, x, pretty_str):\n self.assertEqual(pretty_repr(x), pretty_str)\n\n def test_flatten(self):\n frozen = freeze({'c': 1, 'b': {'a': 2}})\n flat_leaves, tdef = jax.tree_util.tree_flatten(frozen)\n self.assertEqual(flat_leaves, [2, 1])\n self.assertEqual(\n jax.tree_util.tree_unflatten(tdef, flat_leaves),\n frozen,\n )\n flat_path_leaves, tdef = jax.tree_util.tree_flatten_with_path(frozen)\n self.assertEqual(\n flat_path_leaves,\n [\n ((jax.tree_util.DictKey('b'), jax.tree_util.DictKey('a')), 2),\n ((jax.tree_util.DictKey('c'),), 1),\n ],\n )\n self.assertEqual(\n jax.tree_util.tree_unflatten(tdef, [l for _, l in flat_path_leaves]),\n frozen,\n )\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/core/core_lift_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport jax\nimport numpy as np\nfrom absl.testing import absltest\nfrom jax import numpy as jnp\nfrom jax import random\n\nfrom flax import errors\nfrom flax.core import FrozenDict, apply, copy, init, lift, nn\n\n\nclass LiftTest(absltest.TestCase):\n def test_aliasing(self):\n def f(scope):\n a = scope.push('a')\n\n def g(scopes, _):\n scope, a = scopes\n self.assertEqual(a.parent, scope)\n\n lift.vmap(g, variable_axes={}, split_rngs={})((scope, a), jnp.ones((1,)))\n\n init(f)(random.key(0))\n\n def test_undefined_param(self):\n def f(scope):\n dense = lift.vmap(\n nn.dense,\n in_axes=(0, None),\n out_axes=0,\n variable_axes={'params': 0},\n split_rngs={'params': True},\n )\n dense(scope.push('dense'), np.ones((3, 2)), 2)\n\n msg = r'Could not find parameter named \"kernel\" in scope \"/vmap\\(dense\\)\".'\n with self.assertRaisesRegex(errors.ScopeParamNotFoundError, msg):\n apply(f)({'params': {'dense': {'abc': np.ones((3, 3))}}})\n\n def test_jit_cache(self):\n compiles = 0\n\n @lift.jit\n def f(scope, _module_hash, x):\n nonlocal compiles\n compiles += 1\n if scope.is_mutable_collection(\n 'intermediates'\n ) and not scope.is_mutable_collection('params'):\n scope.put_variable('intermediates', 'x', x + 1)\n return nn.dense(scope, x, 1)\n\n x = np.ones((3, 2))\n module_hash = 1\n _, params = init(f)(random.key(0), module_hash, x)\n init(f)(random.key(0), module_hash, x)\n self.assertEqual(compiles, 1)\n apply(f)(params, module_hash, x)\n self.assertEqual(compiles, 2) # apply should cause a compile\n apply(f)(params, module_hash, x)\n self.assertEqual(compiles, 2) # applying again should not\n # edge case where only the implicit return of the jitted functions changes.\n # this should not use the previously cached apply.\n _, state = apply(f, mutable='intermediates')(params, module_hash, x)\n self.assertEqual(compiles, 3) # applying again should not\n self.assertEqual(state['intermediates']['x'].sum(), 3 * 2 * 2)\n\n def test_vjp(self):\n def g(scope, x, y):\n p = scope.param('test', nn.initializers.constant(0.5), ())\n scope.variable('state', 'counter', lambda: 0)\n return p * x * y\n\n def f(scope, x, y):\n z, bwd = lift.vjp(g, scope, x, y)\n return bwd(jnp.ones(y.shape))\n\n x = jnp.array([1.0, 2.0, 3.0])\n y = jnp.array([4.0, 5.0, 6.0])\n _, params = init(f)(random.key(0), x, y)\n params_grad, x_grad, y_grad = apply(f)(params, x, y)\n self.assertEqual(\n params_grad,\n {\n 'params': FrozenDict({'test': 32.0}),\n },\n )\n np.testing.assert_allclose(x_grad, [2.0, 2.5, 3.0])\n np.testing.assert_allclose(y_grad, [0.5, 1.0, 1.5])\n\n def test_jvp(self):\n def g(scope, x):\n p = scope.param('test', nn.initializers.zeros_init(), ())\n scope.variable('state', 'counter', lambda: 0)\n return p * x\n\n def f(scope, x):\n vars_t = jax.tree_util.tree_map(\n jnp.ones_like, scope.variables().get('params', {})\n )\n _, out_t = lift.jvp(\n g, scope, (x,), (jnp.zeros_like(x),), {'params': vars_t}\n )\n return out_t\n\n x = jnp.ones((3,))\n _, params = init(f)(random.key(0), x)\n y_t = apply(f)(params, x)\n np.testing.assert_allclose(y_t, jnp.ones_like(x))\n\n def test_while_loop(self):\n\n def f(scope, x):\n key_zero = random.key(0)\n key_zero = jnp.broadcast_to(key_zero, (2, *key_zero.shape))\n scope.param('inc', lambda _: 1)\n scope.put_variable('state', 'acc', 0)\n scope.put_variable('state', 'rng_params', key_zero)\n scope.put_variable('state', 'rng_loop', jax.random.clone(key_zero))\n\n def cond_fn(scope, c):\n acc = scope.get_variable('state', 'acc')\n return acc < x\n\n def body_fn(scope, c):\n i = scope.get_variable('state', 'acc')\n p_rng = scope.make_rng('params')\n l_rng = scope.make_rng('loop')\n scope.put_variable(\n 'state',\n 'rng_params',\n scope.get_variable('state', 'rng_params').at[i].set(p_rng),\n )\n scope.put_variable(\n 'state',\n 'rng_loop',\n scope.get_variable('state', 'rng_loop').at[i].set(l_rng),\n )\n inc = scope.get_variable('params', 'inc')\n scope.put_variable('state', 'acc', i + inc)\n return c + 2\n\n return lift.while_loop(\n cond_fn,\n body_fn,\n scope,\n 0,\n carry_variables='state',\n split_rngs={'params': False, 'loop': True},\n )\n\n x = 2\n c, vars = apply(f, mutable=True)(\n {}, x, rngs={'params': random.key(1), 'loop': random.key(2)}\n )\n self.assertEqual(vars['state']['acc'], x)\n self.assertEqual(c, 2 * x)\n self.assertEqual(\n vars['state']['rng_params'][0], vars['state']['rng_params'][1]\n )\n with jax.debug_key_reuse(False):\n self.assertNotEqual(\n vars['state']['rng_loop'][0],\n vars['state']['rng_loop'][1],\n )\n\n def test_cond(self):\n def f(scope, x, pred):\n scope.variable('state', 'true_count', lambda: 0)\n scope.variable('state', 'false_count', lambda: 0)\n\n def true_fn(scope, x):\n scope.variable('state', 'true_count').value += 1\n return scope.child(nn.dense)(x, 2)\n\n def false_fn(scope, x):\n scope.variable('state', 'false_count').value += 1\n return -scope.child(nn.dense)(x, 2)\n\n return lift.cond(pred, true_fn, false_fn, scope, x)\n\n x = jnp.ones((1, 3))\n y1, vars = init(f)(random.key(0), x, True)\n self.assertEqual(vars['state'], {'true_count': 1, 'false_count': 0})\n y2, vars = apply(f, mutable='state')(vars, x, False)\n self.assertEqual(vars['state'], {'true_count': 1, 'false_count': 1})\n np.testing.assert_allclose(y1, -y2)\n\n def test_switch(self):\n def f(scope, x, index):\n scope.variable('state', 'a_count', lambda: 0)\n scope.variable('state', 'b_count', lambda: 0)\n scope.variable('state', 'c_count', lambda: 0)\n\n def a_fn(scope, x):\n scope.variable('state', 'a_count').value += 1\n return scope.child(nn.dense)(x, 2)\n\n def b_fn(scope, x):\n scope.variable('state', 'b_count').value += 1\n return -scope.child(nn.dense)(x, 2)\n\n def c_fn(scope, x):\n scope.variable('state', 'c_count').value += 1\n return scope.child(nn.dense)(x, 2)\n\n return lift.switch(index, [a_fn, b_fn, c_fn], scope, x)\n\n x = jnp.ones((1, 3))\n y1, vars = init(f)(random.key(0), x, 0)\n self.assertEqual(vars['state'], {'a_count': 1, 'b_count': 0, 'c_count': 0})\n y2, updates = apply(f, mutable='state')(vars, x, 1)\n vars = copy(vars, updates)\n self.assertEqual(vars['state'], {'a_count': 1, 'b_count': 1, 'c_count': 0})\n np.testing.assert_allclose(y1, -y2)\n y3, updates = apply(f, mutable='state')(vars, x, 2)\n vars = copy(vars, updates)\n self.assertEqual(vars['state'], {'a_count': 1, 'b_count': 1, 'c_count': 1})\n np.testing.assert_allclose(y1, y3)\n\n def test_subscope_var_aliasing(self):\n def test(scope, x):\n subscope = scope.push(name='a')\n subscope.put_variable('state', 'x', 0.0)\n _ = lift.while_loop(\n lambda scope, x: False,\n lambda scope, x: x,\n scope,\n jnp.array(0, jnp.int32),\n carry_variables=['state'],\n )\n subscope.put_variable('state', 'x', 1.0)\n val0 = scope.variables()['state']['a']['x']\n val1 = subscope.variables()['state']['x']\n self.assertEqual(val0, val1)\n return x\n\n init(test)(random.key(0), 1.0)\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/core/core_meta_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport jax\nfrom absl.testing import absltest\nfrom jax import numpy as jnp\nfrom jax import random, sharding\nfrom jax.experimental import mesh_utils\n\nfrom flax import errors\nfrom flax.core import init, lift, meta, nn\n\n\nclass MetaTest(absltest.TestCase):\n def test_boxed_param(self):\n def f(scope, xs):\n def g(scope, x):\n kernel_init = meta.with_partitioning(\n nn.initializers.ones_init(), ('in', 'out')\n )\n kernel = scope.param('kernel', kernel_init, (x.shape[-1], 2))\n kernel_box = scope.get_variable('params', 'kernel')\n self.assertIsInstance(kernel_box, meta.Partitioned)\n self.assertEqual(kernel_box.names, ('in', 'out'))\n return x @ kernel\n\n lift.vmap(\n g,\n in_axes=0,\n variable_axes={'params': 0},\n split_rngs={'params': True},\n metadata_params={meta.PARTITION_NAME: 'batch'},\n )(scope, xs)\n\n _, variables = init(f)(random.key(0), jnp.zeros((8, 3)))\n self.assertEqual(\n variables['params']['kernel'].names, ('batch', 'in', 'out')\n )\n\n def test_boxed_variable(self):\n def f(scope, xs):\n def g(scope, x):\n kernel_init = meta.with_partitioning(\n nn.initializers.ones_init(), ('in', 'out')\n )\n kernel = scope.variable(\n 'params',\n 'kernel',\n kernel_init,\n scope.make_rng('params'),\n (x.shape[-1], 2),\n )\n kernel.value += 1.0\n self.assertEqual(kernel.value.sum(), kernel.value.size * 2)\n kernel_box = scope.get_variable('params', 'kernel')\n self.assertIsInstance(kernel_box, meta.Partitioned)\n self.assertEqual(kernel_box.names, ('in', 'out'))\n return x @ kernel.value\n\n lift.vmap(\n g,\n in_axes=0,\n variable_axes={'params': 0},\n split_rngs={'params': True},\n metadata_params={meta.PARTITION_NAME: 'batch'},\n )(scope, xs)\n\n _, variables = init(f)(random.key(0), jnp.zeros((8, 3)))\n self.assertEqual(\n variables['params']['kernel'].names, ('batch', 'in', 'out')\n )\n\n def test_partition_axis_unspecified(self):\n def f(scope, xs):\n def g(scope, x):\n kernel_init = meta.with_partitioning(\n nn.initializers.ones_init(), ('in', 'out')\n )\n scope.param('kernel', kernel_init, (3, 2))\n return x\n\n with self.assertRaises(errors.PartitioningUnspecifiedError):\n lift.vmap(\n g,\n in_axes=0,\n variable_axes={'params': 0},\n split_rngs={'params': True},\n metadata_params={},\n )(scope, xs)\n\n init(f)(random.key(0), jnp.zeros((8, 3)))\n\n def test_unbox(self):\n xs = {\n 'kernel': meta.Partitioned(jnp.zeros((3, 2)), ('in', 'out')),\n 'complex': meta.Partitioned(\n {'K': jnp.zeros((3, 2)), 'b': jnp.zeros((3,))}, ('data',)\n ),\n }\n unboxed = meta.unbox(xs)\n unboxed_shapes = jax.tree_util.tree_map(jnp.shape, unboxed)\n self.assertEqual(\n unboxed_shapes,\n {\n 'kernel': (3, 2),\n 'complex': {\n 'K': (3, 2),\n 'b': (3,),\n },\n },\n )\n\n def test_scan_over_layers(self):\n def f(scope, x):\n def body(scope, x):\n kernel_init = meta.with_partitioning(\n nn.initializers.ones_init(), ('in', 'out')\n )\n y = nn.dense(scope, x, 3, kernel_init=kernel_init)\n return y, ()\n\n c, _ = lift.scan(\n body,\n variable_axes={'params': 0},\n split_rngs={'params': True},\n length=8,\n metadata_params={meta.PARTITION_NAME: 'layers'},\n )(scope, x)\n return c\n\n _, variables = init(f)(random.key(0), jnp.zeros((8, 3)))\n boxed_shapes = jax.tree_util.tree_map(jnp.shape, variables['params'])\n self.assertEqual(\n boxed_shapes,\n {\n 'kernel': meta.Partitioned((8, 3, 3), ('layers', 'in', 'out')),\n 'bias': (8, 3),\n },\n )\n\n def test_get_partition_spec(self):\n xs = {\n 'kernel': meta.Partitioned(jnp.zeros((8, 3, 3)), ('layers', 'in', 'out')),\n 'bias': jnp.zeros((8, 3)),\n 'step': jnp.array(100),\n }\n ps = meta.get_partition_spec(xs)\n self.assertEqual(\n ps,\n {\n 'kernel': jax.sharding.PartitionSpec('layers', 'in', 'out'),\n 'bias': jax.sharding.PartitionSpec(),\n 'step': jax.sharding.PartitionSpec(),\n },\n )\n\n def test_get_sharding(self):\n devices = mesh_utils.create_device_mesh((jax.local_device_count(), 1))\n mesh = sharding.Mesh(devices, ('in', 'out'))\n xs = {\n 'kernel': meta.Partitioned(jnp.zeros((8, 3)), ('in', 'out')),\n 'bias': jnp.zeros((8, 3)),\n 'step': jnp.array(100),\n }\n ps = meta.get_sharding(xs, mesh)\n self.assertEqual(\n ps,\n {\n 'kernel': jax.sharding.NamedSharding(\n mesh, jax.sharding.PartitionSpec('in', 'out')\n ),\n 'bias': jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec()),\n 'step': jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec()),\n },\n )\n\n def test_boxed_param_with_mesh(self):\n devices = mesh_utils.create_device_mesh((jax.local_device_count(), 1))\n mesh = sharding.Mesh(devices, ('in', 'out'))\n\n def f(scope, x):\n kernel_init = meta.with_partitioning(\n nn.initializers.ones_init(), ('in', 'out'), mesh=mesh\n )\n kernel = scope.param('kernel', kernel_init, (x.shape[-1], 2))\n kernel_box = scope.get_variable('params', 'kernel')\n self.assertIsInstance(kernel_box, meta.Partitioned)\n self.assertEqual(kernel_box.names, ('in', 'out'))\n return x @ kernel\n\n @jax.jit\n def create_state():\n y, variables = init(f)(random.key(0), jnp.zeros((8, 4)))\n spec = meta.get_partition_spec(variables)\n shardings = jax.tree_util.tree_map(\n lambda s: sharding.NamedSharding(mesh, s), spec\n )\n variables = jax.lax.with_sharding_constraint(variables, shardings)\n return variables\n\n variables = create_state()\n self.assertEqual(variables['params']['kernel'].names, ('in', 'out'))\n self.assertIs(variables['params']['kernel'].mesh, mesh)\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/core/core_scope_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport jax\nimport numpy as np\nfrom absl.testing import absltest\nfrom jax import numpy as jnp\nfrom jax import random\n\nfrom flax import errors\nfrom flax import config\nfrom flax.core import Scope, apply, freeze, init, lazy_init, nn, scope\nfrom flax.core.scope import LazyRng\n\n\nclass ScopeTest(absltest.TestCase):\n def test_rng(self):\n def f(scope):\n self.assertTrue(scope.has_rng('params'))\n self.assertFalse(scope.has_rng('dropout'))\n rng = scope.make_rng('params')\n self.assertTrue(\n np.all(rng == LazyRng.create(random.key(0), 1).as_jax_rng())\n )\n\n init(f)(random.key(0))\n\n def test_in_filter(self):\n filter_true = lambda x, y: self.assertTrue(scope.in_filter(x, y))\n filter_false = lambda x, y: self.assertFalse(scope.in_filter(x, y))\n\n filter_true(True, 'any_string1')\n filter_false(False, 'any_string2')\n filter_true('exact_match', 'exact_match')\n filter_false('no_match1', 'no_match2')\n filter_true(['one', 'two'], 'one')\n filter_false(['one', 'two'], 'three')\n filter_false([], 'one')\n filter_false([], None)\n\n def test_union_filter(self):\n def union_check(a, b, ans):\n self.assertEqual(scope.union_filters(a, b), ans)\n self.assertEqual(scope.union_filters(b, a), ans)\n\n union_check(['a', 'b'], ['b', 'c'], {'a', 'b', 'c'})\n union_check(True, False, True)\n union_check(False, False, set())\n union_check(True, True, True)\n union_check(\n scope.DenyList(['a', 'b']),\n scope.DenyList(['b', 'c']),\n scope.DenyList({'b'}),\n )\n union_check(\n scope.DenyList(['a', 'b']), ['b', 'c'], scope.DenyList({'a'})\n )\n\n def test_intersect_filter(self):\n def intersect_check(a, b, ans):\n self.assertEqual(scope.intersect_filters(a, b), ans)\n self.assertEqual(scope.intersect_filters(b, a), ans)\n\n intersect_check(['a', 'b'], ['b', 'c'], {'b'})\n intersect_check(True, False, False)\n intersect_check(False, False, set())\n intersect_check(True, True, True)\n intersect_check(\n scope.DenyList(['a', 'b']),\n scope.DenyList(['b', 'c']),\n scope.DenyList({'a', 'b', 'c'}),\n )\n intersect_check(scope.DenyList(['a', 'b']), ['b', 'c'], {'c'})\n\n def test_subtract_filter(self):\n def subtract_check(a, b, ans):\n self.assertEqual(scope.subtract_filters(a, b), ans)\n\n subtract_check(['a', 'b'], ['b', 'c'], {'a'})\n subtract_check(True, False, scope.DenyList(False))\n subtract_check(False, False, set())\n subtract_check(True, True, False)\n subtract_check(True, 'a', scope.DenyList('a'))\n subtract_check(\n scope.DenyList(['a', 'b']), scope.DenyList(['b', 'c']), {'c'}\n )\n subtract_check(\n scope.DenyList(['a', 'b']),\n ['b', 'c'],\n scope.DenyList({'a', 'b', 'c'}),\n )\n\n def test_group_collections(self):\n params = {'dense1': {'x': [10, 20]}}\n batch_stats = {'dense1': {'ema': 5}}\n xs = {'params': params, 'batch_stats': batch_stats}\n\n # Retrieve all keys only once.\n group = scope.group_collections(xs, ['params', 'params'])\n self.assertEqual(group, ({'params': params}, {}))\n\n # Ignore non-existing keys.\n self.assertEqual(scope.group_collections(xs, ['vars']), ({},))\n\n # False gets nothing and True retrieves all keys once.\n self.assertEqual(\n scope.group_collections(xs, [False, True, True]), ({}, xs, {})\n )\n\n def test_inconsistent_param_shapes(self):\n def f(scope):\n scope.param('test', nn.initializers.ones_init(), (4,))\n\n msg = (\n r'For parameter \"test\" in \"/\", the given initializer is expected to'\n r' generate shape \\(4,\\), but the existing parameter it received has'\n r' shape \\(2,\\).'\n )\n with self.assertRaisesRegex(errors.ScopeParamShapeError, msg):\n apply(f)(freeze({'params': {'test': np.ones((2,))}}))\n\n def test_apply_variables_bad_pytree(self):\n def f(scope):\n scope.param('kernel', nn.initializers.ones_init(), (4,))\n\n params = freeze(\n {\n 'params': {\n 'kernel': np.ones((4,)),\n },\n }\n )\n apply(f)(params) # Valid.\n msg = 'but got a dict with an extra params layer'\n with self.assertRaisesRegex(\n errors.ApplyScopeInvalidVariablesStructureError, msg\n ):\n apply(f)({'params': params})\n\n def test_mutate_undefined_collection(self):\n def f(scope):\n scope.put_variable('state', 'test', 123)\n\n msg = (\n r'Cannot update variable \"test\" in \"/\" because collection \"state\" is'\n r' immutable.'\n )\n with self.assertRaisesRegex(errors.ModifyScopeVariableError, msg):\n init(f, mutable='params')(random.key(0))\n\n def test_undefined_param(self):\n def f(scope):\n nn.dense(scope.push('dense'), np.ones((1, 2)), 2)\n\n msg = r'Could not find parameter named \"kernel\" in scope \"/dense\".'\n with self.assertRaisesRegex(errors.ScopeParamNotFoundError, msg):\n apply(f)({'params': {'abc': 1}})\n\n def test_variable_is_mutable(self):\n def f(scope, should_be_mutable):\n test = scope.variable('state', 'test', lambda: 1)\n self.assertEqual(test.is_mutable(), should_be_mutable)\n\n _, variables = apply(f, mutable='state')({}, True)\n apply(f, mutable=False)(variables, False)\n\n def test_rngs_check_w_frozen_dict(self):\n def f(scope, x):\n return x\n\n _ = apply(f)({}, np.array([0.0]), rngs=freeze({'a': random.key(0)}))\n\n def test_rng_check_w_old_and_new_keys(self):\n # random.key always returns a new-style typed PRNG key.\n key = random.key(0)\n self.assertTrue(scope._is_valid_rng(key))\n self.assertFalse(scope._is_valid_rng(random.split(key)))\n\n # random.PRNGKey returns an old-style uint32 key by default.\n old_key = random.PRNGKey(0)\n self.assertTrue(scope._is_valid_rng(old_key))\n self.assertFalse(scope._is_valid_rng(random.split(old_key)))\n\n # Also explicitly test raw key data, because the jax_enable_custom_prng\n # flag can make PRNGKey return new-style keys.\n raw_key = random.key_data(key)\n self.assertTrue(scope._is_valid_rng(raw_key))\n self.assertFalse(scope._is_valid_rng(random.split(raw_key)))\n\n def test_rng_check_w_lazy_rng(self):\n key = random.key(0)\n self.assertTrue(scope._is_valid_rng(scope.LazyRng.create(key, 1)))\n\n def test_jax_leak_detector(self):\n with jax.check_tracer_leaks(True):\n\n def f(scope):\n def g(scope):\n pass\n\n scope.child(g)()\n\n jax.jit(init(f))(random.key(0))\n\n def test_rng_counter_reuse(self):\n root = Scope({}, {'dropout': random.key(0)})\n\n def f(scope):\n return scope.make_rng('dropout')\n\n a = root.child(f)()\n root = root.rewound()\n b = root.child(f)()\n self.assertFalse(jnp.allclose(a, b))\n\n def test_empty_col_error(self):\n root = Scope({})\n with self.assertRaises(errors.ScopeCollectionNotFound):\n root.param('test', nn.initializers.zeros_init(), ())\n root = Scope({'params': {}})\n with self.assertRaises(errors.ScopeCollectionNotFound):\n root.param('test', nn.initializers.zeros_init(), ())\n\n root = Scope({'params': {'abc': 1}})\n with self.assertRaises(errors.ScopeCollectionNotFound):\n root.variable('state', 'test', jnp.zeros, ())\n root = Scope({'state': {}})\n with self.assertRaises(errors.ScopeCollectionNotFound):\n root.variable('state', 'test', jnp.zeros, ())\n\n def test_variable_no_init(self):\n root = Scope({}, mutable='state')\n with self.assertRaises(errors.ScopeCollectionNotFound):\n root.variable('state', 'test')\n root = Scope({'state': {'abc': 1}}, mutable='state')\n abc = root.variable('state', 'abc')\n self.assertEqual(abc.value, 1)\n with self.assertRaises(errors.ScopeVariableNotFoundError):\n root.variable('state', 'test')\n\n def test_variable_alias(self):\n scope = Scope({}, mutable='state')\n subscope = scope.push(name='a')\n subscope.put_variable('state', 'x', 0.0)\n scope.put_variable('state', 'a', {'x': jnp.array(1.0, jnp.float32)})\n self.assertEqual(\n scope.variables()['state']['a']['x'], subscope.variables()['state']['x']\n )\n\n def test_lazy_init(self):\n def f(scope, x):\n k = scope.param(\n 'kernel', nn.initializers.lecun_normal(), (x.shape[-1], x.shape[-1])\n )\n return x @ k\n\n init_fn = lazy_init(f)\n # provide a massive input message which would OOM if any compute ops were actually executed\n variables = init_fn(\n random.key(0),\n jax.ShapeDtypeStruct((1024 * 1024 * 1024, 128), jnp.float32),\n )\n self.assertEqual(variables['params']['kernel'].shape, (128, 128))\n\n def test_lazy_init_fails_on_data_dependence(self):\n def f(scope, x):\n # kernel is initialized with x so params are now dependent on the input\n k = scope.param('kernel', lambda _: x)\n return x * k\n\n init_fn = lazy_init(f)\n with self.assertRaises(errors.LazyInitError):\n init_fn(random.key(0), jax.ShapeDtypeStruct((8, 4), jnp.float32))\n\n @config.temp_flip_flag('fix_rng_separator', True)\n def test_fold_in_static_seperator(self):\n x = LazyRng(random.key(0), ('ab', 'c'))\n y = LazyRng(random.key(0), ('a', 'bc'))\n self.assertFalse(np.all(x.as_jax_rng() == y.as_jax_rng()))\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/core/design/core_attention_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom functools import partial\nfrom collections.abc import Callable, Sequence\n\nimport jax\nfrom absl.testing import absltest\nfrom jax import lax, random\nfrom jax import numpy as jnp\n\nfrom flax.core import Array, Scope, init, lift, nn, unfreeze\n\n\ndef softmax_attn(scope: Scope, weights: Array):\n del scope\n norm_dims = tuple(range(weights.ndim // 2, weights.ndim))\n log_norms = jax.scipy.special.logsumexp(\n weights, axis=norm_dims, keepdims=True\n )\n return jnp.exp(weights - log_norms)\n\n\ndef with_dropout(fn, rate: float, deterministic: bool = False):\n def attn_fn(scope: Scope, weights: Array):\n attn_weights = fn(scope, weights)\n return nn.dropout(\n scope, attn_weights, deterministic=deterministic, rate=rate\n )\n\n return attn_fn\n\n\ndef _dot_product_attention(\n scope: Scope,\n query: Array,\n key: Array,\n value: Array,\n bias: Array | None = None,\n attn_fn: Callable = softmax_attn,\n dtype=jnp.float32,\n):\n assert key.ndim == query.ndim\n assert key.ndim == value.ndim\n\n n = query.ndim\n attn_weights = lax.dot_general(query, key, (((n - 1,), (n - 1,)), ((), ())))\n if bias is not None:\n attn_weights += bias\n attn_weights = attn_fn(scope, attn_weights)\n attn_weights = attn_weights.astype(dtype)\n\n contract_dims = (\n tuple(range(n - 1, attn_weights.ndim)),\n tuple(range(0, n - 1)),\n )\n y = lax.dot_general(attn_weights, value, (contract_dims, ((), ())))\n return y\n\n\ndef dot_product_attention(\n scope: Scope,\n inputs_q: Array,\n inputs_kv: Array,\n bias: Array | None = None,\n qkv_features: int | None = None,\n out_features: int | None = None,\n attn_fn: Callable = softmax_attn,\n dtype=jnp.float32,\n):\n if qkv_features is None:\n qkv_features = inputs_q.shape[-1]\n if out_features is None:\n out_features = inputs_q.shape[-1]\n dense = partial(nn.dense, features=qkv_features, bias=False, dtype=dtype)\n\n query = scope.child(dense, 'query')(inputs_q)\n key = scope.child(dense, 'key')(inputs_kv)\n value = scope.child(dense, 'value')(inputs_kv)\n\n y = _dot_product_attention(\n scope, query, key, value, bias=bias, attn_fn=attn_fn, dtype=dtype\n )\n\n return scope.child(nn.dense, 'out')(y, features=out_features, dtype=dtype)\n\n\ndef multi_head_dot_product_attention(\n scope: Scope,\n inputs_q: Array,\n inputs_kv: Array,\n bias: Array | None = None,\n qkv_features: int | None = None,\n out_features: int | None = None,\n attn_fn: Callable = softmax_attn,\n batch_axes: Sequence[int] = (0,),\n num_heads: int = 1,\n dtype=jnp.float32,\n broadcast_dropout=False,\n):\n if qkv_features is None:\n qkv_features = inputs_q.shape[-1]\n if out_features is None:\n out_features = inputs_q.shape[-1]\n\n attn_fn = partial(\n dot_product_attention,\n attn_fn=attn_fn,\n qkv_features=qkv_features // num_heads,\n out_features=out_features,\n dtype=dtype,\n )\n attn_fn = lift.vmap(\n attn_fn,\n in_axes=(None, None, None),\n out_axes=-2,\n axis_size=num_heads,\n variable_axes={'params': 0},\n split_rngs={'params': True, 'dropout': not broadcast_dropout},\n )\n for axis in reversed(sorted(batch_axes)):\n attn_fn = lift.vmap(\n attn_fn,\n in_axes=(axis, axis, axis),\n out_axes=axis,\n variable_axes={'params': None},\n split_rngs={'params': False, 'dropout': not broadcast_dropout},\n )\n\n y = attn_fn(scope, inputs_q, inputs_kv, bias)\n return y.mean(axis=-2)\n\n\nclass AttentionTest(absltest.TestCase):\n def test_attention(self):\n inputs = jnp.ones((2, 7, 16))\n model = partial(\n multi_head_dot_product_attention,\n num_heads=2,\n batch_axes=(0,),\n attn_fn=with_dropout(softmax_attn, 0.1, deterministic=False),\n )\n\n rngs = {'params': random.key(0), 'dropout': random.key(1)}\n y, variables = jax.jit(init(model))(rngs, inputs, inputs)\n variable_shapes = jax.tree_util.tree_map(jnp.shape, variables['params'])\n self.assertEqual(y.shape, (2, 7, 16))\n self.assertEqual(\n unfreeze(variable_shapes),\n {\n 'key': {'kernel': (2, 16, 8)},\n 'value': {'kernel': (2, 16, 8)},\n 'query': {'kernel': (2, 16, 8)},\n 'out': {'bias': (2, 16), 'kernel': (2, 8, 16)},\n },\n )\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/core/design/core_auto_encoder_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom dataclasses import dataclass\nfrom collections.abc import Callable\n\nimport jax\nfrom absl.testing import absltest\nfrom jax import numpy as jnp\nfrom jax import random\n\nfrom flax.core import Array, Scope, init, nn, unfreeze\n\n\ndef mlp(scope: Scope, x: Array, hidden: int, out: int):\n x = scope.child(nn.dense, 'hidden')(x, hidden)\n x = nn.relu(x)\n return scope.child(nn.dense, 'out')(x, out)\n\n\n@dataclass\nclass AutoEncoder:\n latents: int\n features: int\n hidden: int\n\n def __call__(self, scope, x):\n z = self.encode(scope, x)\n return self.decode(scope, z)\n\n def encode(self, scope, x):\n return scope.child(mlp, 'encoder')(x, self.hidden, self.latents)\n\n def decode(self, scope, z):\n return scope.child(mlp, 'decoder')(z, self.hidden, self.features)\n\n\ndef module_method(fn, name=None):\n if name is None:\n name = fn.__name__ if hasattr(fn, '__name__') else None\n\n def wrapper(self, *args, **kwargs):\n scope = self.scope.rewound()\n mod_fn = lambda scope: fn(self, scope, *args, **kwargs)\n return scope.child(mod_fn, name)()\n\n return wrapper\n\n\n@dataclass\nclass AutoEncoder2:\n scope: Scope\n latents: int\n features: int\n hidden: int\n\n def __call__(self, x):\n z = self.encode(x)\n return self.decode(z)\n\n @module_method\n def encode(self, scope, x):\n return mlp(scope, x, self.hidden, self.latents)\n\n @module_method\n def decode(self, scope, z):\n return mlp(scope, z, self.hidden, self.features)\n\n\n@dataclass\nclass AutoEncoder3:\n encode: Callable\n decode: Callable\n\n @staticmethod\n def create(scope, hidden: int, latents: int, features: int):\n enc = scope.child(mlp, 'encode', hidden=hidden, out=latents)\n dec = scope.child(mlp, 'decode', hidden=hidden, out=features)\n return AutoEncoder3(enc, dec)\n\n def __call__(self, x):\n z = self.encode(x)\n return self.decode(z)\n\n\nclass AutoEncoderTest(absltest.TestCase):\n def test_auto_encoder_hp_struct(self):\n ae = AutoEncoder(latents=2, features=4, hidden=3)\n x = jnp.ones((1, 4))\n x_r, variables = init(ae)(random.key(0), x)\n self.assertEqual(x.shape, x_r.shape)\n variable_shapes = unfreeze(\n jax.tree_util.tree_map(jnp.shape, variables['params'])\n )\n self.assertEqual(\n variable_shapes,\n {\n 'encoder': {\n 'hidden': {'kernel': (4, 3), 'bias': (3,)},\n 'out': {'kernel': (3, 2), 'bias': (2,)},\n },\n 'decoder': {\n 'hidden': {'kernel': (2, 3), 'bias': (3,)},\n 'out': {'kernel': (3, 4), 'bias': (4,)},\n },\n },\n )\n\n def test_auto_encoder_with_scope(self):\n ae = lambda scope, x: AutoEncoder2(scope, latents=2, features=4, hidden=3)(\n x\n )\n x = jnp.ones((1, 4))\n\n x_r, variables = init(ae)(random.key(0), x)\n self.assertEqual(x.shape, x_r.shape)\n variable_shapes = unfreeze(\n jax.tree_util.tree_map(jnp.shape, variables['params'])\n )\n self.assertEqual(\n variable_shapes,\n {\n 'encode': {\n 'hidden': {'kernel': (4, 3), 'bias': (3,)},\n 'out': {'kernel': (3, 2), 'bias': (2,)},\n },\n 'decode': {\n 'hidden': {'kernel': (2, 3), 'bias': (3,)},\n 'out': {'kernel': (3, 4), 'bias': (4,)},\n },\n },\n )\n\n def test_auto_encoder_bind_method(self):\n ae = lambda scope, x: AutoEncoder3.create(\n scope, latents=2, features=4, hidden=3\n )(x)\n x = jnp.ones((1, 4))\n\n x_r, variables = init(ae)(random.key(0), x)\n self.assertEqual(x.shape, x_r.shape)\n variable_shapes = unfreeze(\n jax.tree_util.tree_map(jnp.shape, variables['params'])\n )\n self.assertEqual(\n variable_shapes,\n {\n 'encode': {\n 'hidden': {'kernel': (4, 3), 'bias': (3,)},\n 'out': {'kernel': (3, 2), 'bias': (2,)},\n },\n 'decode': {\n 'hidden': {'kernel': (2, 3), 'bias': (3,)},\n 'out': {'kernel': (3, 4), 'bias': (4,)},\n },\n },\n )\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/core/design/core_big_resnets_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom functools import partial\n\nimport jax\nimport numpy as np\nfrom absl.testing import absltest\nfrom jax import numpy as jnp\nfrom jax import random\n\nfrom flax.core import Array, Scope, init, lift, nn, unfreeze\n\ndefault_norm = partial(nn.batch_norm)\n\n\ndef residual_block(scope: Scope, x: Array, conv, norm, act, features: int):\n residual = x\n x = scope.child(conv, 'conv_1')(x, features, (3, 3))\n x = scope.child(norm, 'bn_1')(x)\n x = act(x)\n x = scope.child(conv, 'conv_2')(x, features, (3, 3))\n x = scope.child(norm, 'bn_2')(x)\n return act(residual + x)\n\n\ndef big_resnet(\n scope: Scope,\n x,\n blocks=(10, 5),\n dtype=jnp.float32,\n norm=default_norm,\n act=nn.relu,\n):\n conv = partial(nn.conv, bias=False, dtype=dtype)\n norm = partial(norm, dtype=dtype)\n\n # a two stage resnet where inner blocks are rematerialized to make sure\n # memory consumtion grows as O(sqrt(N)) and compute is O(N) where N is the number of blocks..\n # we use a double scan such that the compiled binary is of size O(1).\n print('total residual blocks:', np.prod(blocks))\n\n def body_fn(scope, x):\n return residual_block(scope, x, conv, norm, act, features=x.shape[-1])\n\n return lift.remat_scan(\n body_fn,\n lengths=blocks,\n variable_axes={'params': 0, 'batch_stats': 0},\n split_rngs={'params': True},\n policy=None,\n )(scope, x)\n\n\nclass BigResnetTest(absltest.TestCase):\n def test_big_resnet(self):\n x = random.normal(random.key(0), (1, 8, 8, 8))\n y, variables = init(big_resnet)(random.key(1), x)\n self.assertEqual(y.shape, (1, 8, 8, 8))\n param_shapes = unfreeze(\n jax.tree_util.tree_map(jnp.shape, variables['params'])\n )\n batch_stats_shapes = unfreeze(\n jax.tree_util.tree_map(jnp.shape, variables['batch_stats'])\n )\n self.assertEqual(\n param_shapes,\n {\n 'conv_1': {'kernel': (10, 5, 3, 3, 8, 8)},\n 'conv_2': {'kernel': (10, 5, 3, 3, 8, 8)},\n 'bn_1': {'scale': (10, 5, 8), 'bias': (10, 5, 8)},\n 'bn_2': {'scale': (10, 5, 8), 'bias': (10, 5, 8)},\n },\n )\n self.assertEqual(\n batch_stats_shapes,\n {\n 'bn_1': {'var': (10, 5, 8), 'mean': (10, 5, 8)},\n 'bn_2': {'var': (10, 5, 8), 'mean': (10, 5, 8)},\n },\n )\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/core/design/core_custom_vjp_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom functools import partial\nfrom collections.abc import Callable, Sequence\n\nimport jax\nimport numpy as np\nfrom absl.testing import absltest\nfrom jax import numpy as jnp\nfrom jax import random\n\nfrom flax.core import Array, Scope, apply, init, lift, nn, unfreeze\n\n\ndef mlp_custom_grad(\n scope: Scope,\n x: Array,\n sizes: Sequence[int] = (8, 1),\n act_fn: Callable[[Array], Array] = nn.relu,\n):\n f = nn.dense\n\n def fwd(scope, x, features):\n y, vjp_fn = lift.vjp(partial(f, features=features), scope, x)\n return y, vjp_fn\n\n def bwd(features, res, y_t):\n del features\n vjp_fn = res\n params_t, *input_t = vjp_fn(y_t)\n params_t = jax.tree_util.tree_map(jnp.sign, params_t)\n return (params_t, *input_t)\n\n dense_custom_grad = lift.custom_vjp(\n f, forward_fn=fwd, backward_fn=bwd, nondiff_argnums=(2,)\n )\n\n # hidden layers\n for size in sizes[:-1]:\n x = scope.child(dense_custom_grad, prefix='hidden_')(x, size)\n x = act_fn(x)\n\n # output layer\n return scope.child(dense_custom_grad, 'out')(x, sizes[-1])\n\n\nclass CustomVJPTest(absltest.TestCase):\n def test_custom_vjp(self):\n x = random.normal(random.key(0), (1, 4))\n y, variables = init(mlp_custom_grad)(random.key(1), x)\n param_shapes = unfreeze(\n jax.tree_util.tree_map(jnp.shape, variables['params'])\n )\n loss_fn = lambda p, x: jnp.mean(apply(mlp_custom_grad)(p, x) ** 2)\n grad = jax.grad(loss_fn)(variables, x)\n grad_shapes = unfreeze(jax.tree_util.tree_map(jnp.shape, grad['params']))\n self.assertEqual(y.shape, (1, 1))\n expected_param_shapes = {\n 'hidden_0': {'kernel': (4, 8), 'bias': (8,)},\n 'out': {'kernel': (8, 1), 'bias': (1,)},\n }\n self.assertEqual(param_shapes, expected_param_shapes)\n self.assertEqual(grad_shapes, expected_param_shapes)\n for g in jax.tree_util.tree_leaves(grad):\n self.assertTrue(np.all(g == np.sign(g)))\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/core/design/core_dense_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom dataclasses import dataclass\nfrom typing import Any\n\nimport jax\nfrom absl.testing import absltest\nfrom jax import numpy as jnp\nfrom jax import random\n\nfrom flax import struct\nfrom flax.core import Array, init, nn, unfreeze\n\n\n@dataclass\nclass Dense:\n features: int\n bias: bool = True\n kernel_init: Any = nn.linear.default_kernel_init\n bias_init: Any = nn.initializers.zeros_init()\n\n def __call__(self, scope, x):\n kernel = scope.param(\n 'kernel', self.kernel_init, (x.shape[-1], self.features)\n )\n y = x @ kernel\n if self.bias:\n bias = scope.param('bias', self.bias_init, (self.features,))\n y += bias.reshape((1,) * (y.ndim - 1) + (-1,))\n return y\n\n\n@struct.dataclass\nclass ExplicitDense:\n kernel: Array\n bias: Array | None\n\n # a fully explicit \"scope free\" version\n @staticmethod\n def create(\n rng,\n in_size,\n out_size,\n bias=True,\n kernel_init=nn.linear.default_kernel_init,\n bias_init=nn.initializers.zeros_init(),\n ):\n k1, k2 = random.split(rng, 2)\n kernel = kernel_init(k1, (in_size, out_size))\n if bias:\n bias = bias_init(k2, (out_size,))\n else:\n bias = None\n return ExplicitDense(kernel, bias)\n\n # a semi-explicit version where a scope is used to create explicit params\n @staticmethod\n def create_in_scope(\n scope,\n in_size,\n out_size,\n bias=True,\n kernel_init=nn.linear.default_kernel_init,\n bias_init=nn.initializers.zeros_init(),\n ):\n kernel = scope.param('kernel', kernel_init, (in_size, out_size))\n if bias:\n bias = scope.param('bias', bias_init, (out_size,))\n else:\n bias = None\n return ExplicitDense(kernel, bias)\n\n def __call__(self, x):\n y = x @ self.kernel\n if self.bias is not None:\n y += self.bias.reshape((1,) * (y.ndim - 1) + (-1,))\n return y\n\n\ndef explicit_mlp(scope, x, sizes=(3, 1)):\n for i, size in enumerate(sizes):\n dense = scope.param(f'dense_{i}', ExplicitDense.create, x.shape[-1], size)\n x = dense(x)\n if i + 1 < len(sizes):\n x = nn.relu(x)\n return x\n\n\ndef semi_explicit_mlp(scope, x, sizes=(3, 1)):\n for i, size in enumerate(sizes):\n dense = scope.child(ExplicitDense.create_in_scope, prefix='dense_')(\n x.shape[-1], size\n )\n x = dense(x)\n if i + 1 < len(sizes):\n x = nn.relu(x)\n return x\n\n\nclass DenseTest(absltest.TestCase):\n def test_dense(self):\n model = Dense(features=4)\n x = jnp.ones((1, 3))\n y, variables = init(model)(random.key(0), x)\n param_shapes = unfreeze(\n jax.tree_util.tree_map(jnp.shape, variables['params'])\n )\n self.assertEqual(y.shape, (1, 4))\n self.assertEqual(\n param_shapes,\n {\n 'kernel': (3, 4),\n 'bias': (4,),\n },\n )\n\n def test_explicit_dense(self):\n x = jnp.ones((1, 3))\n y, variables = init(explicit_mlp)(random.key(0), x)\n param_shapes = unfreeze(\n jax.tree_util.tree_map(jnp.shape, variables['params'])\n )\n self.assertEqual(y.shape, (1, 4))\n self.assertEqual(\n param_shapes,\n {\n 'kernel': (3, 4),\n 'bias': (4,),\n },\n )\n\n def test_explicit_dense(self):\n x = jnp.ones((1, 4))\n y, variables = init(explicit_mlp)(random.key(0), x)\n param_shapes = unfreeze(\n jax.tree_util.tree_map(jnp.shape, variables['params'])\n )\n self.assertEqual(y.shape, (1, 1))\n self.assertEqual(\n param_shapes,\n {\n 'dense_0': ExplicitDense((4, 3), (3,)),\n 'dense_1': ExplicitDense((3, 1), (1,)),\n },\n )\n\n def test_semi_explicit_dense(self):\n x = jnp.ones((1, 4))\n y, variables = init(semi_explicit_mlp)(random.key(0), x)\n param_shapes = unfreeze(\n jax.tree_util.tree_map(jnp.shape, variables['params'])\n )\n self.assertEqual(y.shape, (1, 1))\n self.assertEqual(\n param_shapes,\n {\n 'dense_0': {'kernel': (4, 3), 'bias': (3,)},\n 'dense_1': {'kernel': (3, 1), 'bias': (1,)},\n },\n )\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/core/design/core_flow_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom dataclasses import dataclass\nfrom typing import Any\nfrom collections.abc import Sequence\n\nimport jax\nfrom absl.testing import absltest\nfrom jax import numpy as jnp\nfrom jax import random\nfrom jax.scipy.linalg import expm\n\nfrom flax.core import Array, Scope, apply, init, nn, unfreeze\n\nInitializer = Any\nFlow = Any\n\n\n@dataclass\nclass DenseFlow:\n kernel_init: Initializer = nn.linear.default_kernel_init\n bias_init: Initializer = nn.initializers.zeros_init()\n\n def params(self, scope: Scope, features: int):\n kernel = scope.param('kernel', self.kernel_init, (features, features))\n bias = scope.param('bias', self.bias_init, (features,))\n return kernel, bias\n\n def forward(self, scope: Scope, x: Array):\n kernel, bias = self.params(scope, x.shape[-1])\n return jnp.dot(x, expm(kernel)) + bias.reshape((1,) * (x.ndim - 1) + (-1,))\n\n def backward(self, scope: Scope, y: Array):\n kernel, bias = self.params(scope, y.shape[-1])\n return jnp.dot(y - bias.reshape((1,) * (y.ndim - 1) + (-1,)), expm(-kernel))\n\n\n@dataclass\nclass StackFlow:\n flows: Sequence[Flow]\n\n def forward(self, scope: Scope, x: Array):\n for i, f in enumerate(self.flows):\n x = scope.child(f.forward, name=str(i))(x)\n return x\n\n def backward(self, scope: Scope, x: Array):\n for i, f in reversed(tuple(enumerate(self.flows))):\n x = scope.child(f.backward, name=str(i))(x)\n return x\n\n\nclass FlowTest(absltest.TestCase):\n def test_flow(self):\n x = jnp.ones((1, 3))\n flow = StackFlow((DenseFlow(),) * 3)\n y, variables = init(flow.forward)(random.key(0), x)\n param_shapes = unfreeze(\n jax.tree_util.tree_map(jnp.shape, variables['params'])\n )\n self.assertEqual(y.shape, (1, 3))\n self.assertEqual(\n param_shapes,\n {\n '0': {'kernel': (3, 3), 'bias': (3,)},\n '1': {'kernel': (3, 3), 'bias': (3,)},\n '2': {'kernel': (3, 3), 'bias': (3,)},\n },\n )\n x_restored = apply(flow.backward)(variables, y)\n self.assertTrue(jnp.allclose(x, x_restored))\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/core/design/core_resnet_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom functools import partial\n\nimport jax\nfrom absl.testing import absltest\nfrom jax import numpy as jnp\nfrom jax import random\n\nfrom flax.core import Array, Scope, init, nn, unfreeze\n\ndefault_norm = partial(nn.batch_norm)\n\n\ndef residual_block(\n scope: Scope, x: Array, conv, norm, act, features: int, strides=(1, 1)\n):\n residual = x\n x = scope.child(conv, 'conv_1')(x, features, (1, 1))\n x = scope.child(norm, 'bn_1')(x)\n x = act(x)\n x = scope.child(conv, 'conv_2')(x, 4 * features, (3, 3), strides=strides)\n x = scope.child(norm, 'bn_2')(x)\n x = act(x)\n x = scope.child(conv, 'conv_3')(x, 4 * features, (1, 1))\n x = scope.child(norm, 'bn_3')(x)\n\n if x.shape != residual.shape:\n residual = scope.child(conv, 'proj_conv')(\n residual, 4 * features, (1, 1), strides=strides\n )\n residual = scope.child(norm, 'proj_bn')(residual)\n\n return act(residual + x)\n\n\ndef resnet(\n scope: Scope,\n x,\n block_sizes=(3, 4, 6, 3),\n features=16,\n num_classes=1000,\n dtype=jnp.float32,\n norm=default_norm,\n act=nn.relu,\n):\n conv = partial(nn.conv, bias=False, dtype=dtype)\n norm = partial(norm, dtype=dtype)\n\n x = scope.child(conv, 'init_conv')(x, 16, (7, 7), padding=((3, 3), (3, 3)))\n x = scope.child(norm, 'init_bn')(x)\n x = act(x)\n x = nn.max_pool(x, (2, 2), (2, 2), 'SAME')\n\n for i, size in enumerate(block_sizes):\n for j in range(size):\n strides = (1, 1)\n if i > 0 and j == 0:\n strides = (2, 2)\n block_features = features * 2**i\n block_scope = scope.push(f'block_{i}_{j}')\n x = residual_block(\n block_scope, x, conv, norm, act, block_features, strides\n )\n # we can access parameters of the sub module by operating on the scope\n # Example:\n # block_scope.get_kind('params')['conv_1']['kernel']\n x = jnp.mean(x, (1, 2))\n x = scope.child(nn.dense, 'out')(x, num_classes)\n return x\n\n\nclass ResNetTest(absltest.TestCase):\n def test_resnet(self):\n block_sizes = (2, 2)\n x = random.normal(random.key(0), (1, 64, 64, 3))\n y, variables = init(resnet)(\n random.key(1), x, block_sizes=block_sizes, features=16\n )\n param_shapes = unfreeze(\n jax.tree_util.tree_map(jnp.shape, variables['params'])\n )\n self.assertEqual(y.shape, (1, 1000))\n\n self.assertEqual(\n param_shapes,\n {\n 'init_conv': {'kernel': (7, 7, 3, 16)},\n 'init_bn': {'bias': (16,), 'scale': (16,)},\n 'out': {'kernel': (128, 1000), 'bias': (1000,)},\n 'block_0_0': {\n 'conv_1': {'kernel': (1, 1, 16, 16)},\n 'conv_2': {'kernel': (3, 3, 16, 64)},\n 'conv_3': {'kernel': (1, 1, 64, 64)},\n 'bn_1': {'bias': (16,), 'scale': (16,)},\n 'bn_2': {'bias': (64,), 'scale': (64,)},\n 'bn_3': {'bias': (64,), 'scale': (64,)},\n 'proj_conv': {'kernel': (1, 1, 16, 64)},\n 'proj_bn': {'bias': (64,), 'scale': (64,)},\n },\n 'block_0_1': {\n 'conv_1': {'kernel': (1, 1, 64, 16)},\n 'conv_2': {'kernel': (3, 3, 16, 64)},\n 'conv_3': {'kernel': (1, 1, 64, 64)},\n 'bn_1': {'bias': (16,), 'scale': (16,)},\n 'bn_2': {'bias': (64,), 'scale': (64,)},\n 'bn_3': {'bias': (64,), 'scale': (64,)},\n },\n 'block_1_0': {\n 'conv_1': {'kernel': (1, 1, 64, 32)},\n 'conv_2': {'kernel': (3, 3, 32, 128)},\n 'conv_3': {'kernel': (1, 1, 128, 128)},\n 'bn_1': {'bias': (32,), 'scale': (32,)},\n 'bn_2': {'bias': (128,), 'scale': (128,)},\n 'bn_3': {'bias': (128,), 'scale': (128,)},\n 'proj_conv': {'kernel': (1, 1, 64, 128)},\n 'proj_bn': {'bias': (128,), 'scale': (128,)},\n },\n 'block_1_1': {\n 'conv_1': {'kernel': (1, 1, 128, 32)},\n 'conv_2': {'kernel': (3, 3, 32, 128)},\n 'conv_3': {'kernel': (1, 1, 128, 128)},\n 'bn_1': {'bias': (32,), 'scale': (32,)},\n 'bn_2': {'bias': (128,), 'scale': (128,)},\n 'bn_3': {'bias': (128,), 'scale': (128,)},\n },\n },\n )\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/core/design/core_scan_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport jax\nfrom absl.testing import absltest\nfrom jax import numpy as jnp\nfrom jax import random\n\nfrom flax.core import Array, Scope, init, lift, nn, unfreeze\n\n\ndef mlp_scan(scope: Scope, xs: Array, share_params: bool = False):\n scope.variable('counter', 'i', jnp.zeros, ())\n\n def body_fn(scope, c, x):\n counter = scope.variable('counter', 'i', jnp.zeros, ())\n counter.value += 1\n x = scope.child(nn.dense)(x, 1)\n return c, x\n\n if share_params:\n _, ys = lift.scan(\n body_fn,\n variable_carry='counter',\n variable_broadcast='params',\n split_rngs={'params': False},\n )(scope, (), xs)\n else:\n _, ys = lift.scan(\n body_fn,\n variable_carry='counter',\n variable_axes={'params': 0},\n split_rngs={'params': True},\n )(scope, (), xs)\n\n # output layer\n return ys\n\n\nclass ScanTest(absltest.TestCase):\n def test_scan_unshared_params(self):\n x = random.normal(random.key(0), (1, 4))\n x = jnp.concatenate([x, x], 0)\n y, variables = init(mlp_scan)(random.key(1), x, share_params=False)\n\n param_shapes = unfreeze(\n jax.tree_util.tree_map(jnp.shape, variables['params'])\n )\n self.assertEqual(variables['counter']['i'], 2)\n self.assertEqual(\n param_shapes,\n {\n 'dense_0': {'kernel': (2, 4, 1), 'bias': (2, 1)},\n },\n )\n\n self.assertNotEqual(y[0], y[1])\n k1, k2 = variables['params']['dense_0']['kernel']\n self.assertFalse(jnp.allclose(k1, k2))\n\n def test_scan_shared_params(self):\n x = random.normal(random.key(0), (1, 4))\n x = jnp.concatenate([x, x], 0)\n y, variables = init(mlp_scan)(random.key(1), x, share_params=True)\n\n param_shapes = unfreeze(\n jax.tree_util.tree_map(jnp.shape, variables['params'])\n )\n self.assertEqual(variables['counter']['i'], 2)\n self.assertEqual(\n param_shapes,\n {\n 'dense_0': {'kernel': (4, 1), 'bias': (1,)},\n },\n )\n\n self.assertEqual(y[0], y[1])\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/core/design/core_tied_autoencoder_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom dataclasses import dataclass\n\nimport jax\nfrom absl.testing import absltest\nfrom jax import numpy as jnp\nfrom jax import random\n\nfrom flax.core import init, lift, nn, unfreeze\n\n\ndef transpose(fn):\n def trans(variables):\n return jax.tree_util.tree_map(lambda x: x.T, variables)\n\n return lift.map_variables(\n fn, 'params', map_in_fn=trans, map_out_fn=trans, mutable=True\n )\n\n\n@dataclass\nclass TiedAutoEncoder:\n latents: int\n features: int\n\n def __call__(self, scope, x):\n z = self.encode(scope, x)\n return self.decode(scope, z)\n\n def encode(self, scope, x):\n return nn.dense(scope, x, self.latents, bias=False)\n\n def decode(self, scope, z):\n return transpose(nn.dense)(scope, z, self.features, bias=False)\n\n\nclass TiedAutoEncoderTest(absltest.TestCase):\n def test_tied_auto_encoder(self):\n ae = TiedAutoEncoder(latents=2, features=4)\n x = jnp.ones((1, ae.features))\n x_r, variables = init(ae)(random.key(0), x)\n\n param_shapes = unfreeze(\n jax.tree_util.tree_map(jnp.shape, variables['params'])\n )\n self.assertEqual(\n param_shapes,\n {\n 'kernel': (4, 2),\n },\n )\n self.assertEqual(x.shape, x_r.shape)\n\n def test_init_from_decoder(self):\n ae = TiedAutoEncoder(latents=2, features=4)\n z = jnp.ones((1, ae.latents))\n x_r, variables = init(ae.decode)(random.key(0), z)\n\n param_shapes = unfreeze(\n jax.tree_util.tree_map(jnp.shape, variables['params'])\n )\n self.assertEqual(\n param_shapes,\n {\n 'kernel': (4, 2),\n },\n )\n self.assertEqual(x_r.shape, (1, 4))\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/core/design/core_vmap_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom collections.abc import Callable, Sequence\n\nimport jax\nfrom absl.testing import absltest\nfrom jax import numpy as jnp\nfrom jax import random\n\nfrom flax.core import Array, Scope, init, lift, nn, unfreeze\n\n\ndef mlp_vmap(\n scope: Scope,\n x: Array,\n sizes: Sequence[int] = (8, 1),\n act_fn: Callable[[Array], Array] = nn.relu,\n share_params: bool = False,\n):\n if share_params:\n dense_vmap = lift.vmap(\n nn.dense,\n in_axes=(0, None),\n variable_axes={'params': None},\n split_rngs={'params': False},\n )\n else:\n dense_vmap = lift.vmap(\n nn.dense,\n in_axes=(0, None),\n variable_axes={'params': 0},\n split_rngs={'params': True},\n )\n\n # hidden layers\n for size in sizes[:-1]:\n x = scope.child(dense_vmap, prefix='hidden_')(x, size)\n x = act_fn(x)\n\n # output layer\n return scope.child(dense_vmap, 'out')(x, sizes[-1])\n\n\nclass VMapTest(absltest.TestCase):\n def test_vmap_shared(self):\n x = random.normal(random.key(0), (1, 4))\n x = jnp.concatenate([x, x], 0)\n\n y, variables = init(mlp_vmap)(random.key(1), x, share_params=True)\n\n param_shapes = unfreeze(\n jax.tree_util.tree_map(jnp.shape, variables['params'])\n )\n self.assertEqual(\n param_shapes,\n {\n 'hidden_0': {'kernel': (4, 8), 'bias': (8,)},\n 'out': {'kernel': (8, 1), 'bias': (1,)},\n },\n )\n self.assertEqual(y.shape, (2, 1))\n self.assertTrue(jnp.allclose(y[0], y[1]))\n\n def test_vmap_unshared(self):\n x = random.normal(random.key(0), (1, 4))\n x = jnp.concatenate([x, x], 0)\n\n y, variables = init(mlp_vmap)(random.key(1), x, share_params=False)\n\n param_shapes = unfreeze(\n jax.tree_util.tree_map(jnp.shape, variables['params'])\n )\n self.assertEqual(\n param_shapes,\n {\n 'hidden_0': {'kernel': (2, 4, 8), 'bias': (2, 8)},\n 'out': {'kernel': (2, 8, 1), 'bias': (2, 1)},\n },\n )\n self.assertEqual(y.shape, (2, 1))\n self.assertFalse(jnp.allclose(y[0], y[1]))\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/core/design/core_weight_std_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom functools import partial\nfrom collections.abc import Sequence\n\nimport jax\nfrom absl.testing import absltest\nfrom jax import numpy as jnp\nfrom jax import random\n\nfrom flax.core import Array, Scope, apply, init, lift, nn, unfreeze\n\n\ndef weight_std(fn, kernel_name='kernel', eps=1e-8):\n def std(variables):\n params = variables['params']\n assert kernel_name in params\n kernel = params[kernel_name]\n redux = tuple(range(kernel.ndim - 1))\n norm = jnp.square(kernel).sum(redux, keepdims=True)\n std_kernel = kernel / jnp.sqrt(norm + eps)\n params[kernel_name] = std_kernel\n return variables\n\n # map_variables handles a few of nasty edge cases here...\n # the transformed kind will be immutable inside fn\n # this way we avoid lost mutations to param\n # map_variables also avoids accidental reuse of rngs\n # and it makes sure that other state is updated correctly (not twice during init!)\n return lift.map_variables(fn, 'params', std, init=True)\n\n\ndef mlp(scope: Scope, x: Array, sizes: Sequence[int] = (8, 1)):\n std_dense = weight_std(\n partial(nn.dense, kernel_init=nn.initializers.normal(stddev=1e5))\n )\n for size in sizes[:-1]:\n x = scope.child(std_dense, prefix='hidden_')(x, size)\n return scope.child(nn.dense, 'out')(x, sizes[-1])\n\n\nclass WeightStdTest(absltest.TestCase):\n def test_weight_std(self):\n x = random.normal(\n random.key(0),\n (\n 1,\n 4,\n ),\n )\n y, variables = init(mlp)(random.key(1), x)\n\n param_shapes = unfreeze(\n jax.tree_util.tree_map(jnp.shape, variables['params'])\n )\n self.assertEqual(\n param_shapes,\n {\n 'hidden_0': {'kernel': (4, 8), 'bias': (8,)},\n 'out': {'kernel': (8, 1), 'bias': (1,)},\n },\n )\n self.assertEqual(y.shape, (1, 1))\n self.assertTrue(y.ravel() < 1.0)\n\n y2 = apply(mlp)(variables, x)\n self.assertTrue(jnp.allclose(y, y2))\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/cursor_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for flax.struct.\"\"\"\n\n\nimport dataclasses\nfrom typing import Any, NamedTuple\n\nimport jax\nimport jax.numpy as jnp\nimport optax\nfrom absl.testing import absltest\n\nimport flax\nimport flax.linen as nn\nfrom flax.core import freeze\nfrom flax.cursor import AccessType, _traverse_tree, cursor\nfrom flax.errors import CursorFindError, TraverseTreeError\nfrom flax.training import train_state\n\n# Parse absl flags test_srcdir and test_tmpdir.\njax.config.parse_flags_with_absl()\n\n\nclass GenericTuple(NamedTuple):\n x: Any\n y: Any = None\n z: Any = None\n\n\n@dataclasses.dataclass\nclass GenericDataClass:\n x: Any\n y: Any = None\n z: Any = None\n\n\nclass CursorTest(absltest.TestCase):\n def test_repr(self):\n g = GenericTuple(1, 'a', (2, 'b'))\n c = cursor(\n {'a': {1: {(2, 3): 'z', 4: g, '6': (7, 8)}, 'b': [1, 2, 3]}, 'z': -1}\n )\n self.assertEqual(\n repr(c),\n \"\"\"Cursor(\n _obj={'a': {1: {(2, 3): 'z', 4: GenericTuple(x=1, y='a', z=(2, 'b')), '6': (7, 8)}, 'b': [1, 2, 3]}, 'z': -1},\n _changes={}\n)\"\"\",\n )\n\n # test overwriting\n c['z'] = -2\n c['z'] = -3\n c['a']['b'][1] = -2\n c['a']['b'] = None\n\n # test deep mutation\n c['a'][1][4].x = (2, 4, 6)\n c['a'][1][4].z[0] = flax.core.freeze({'a': 1, 'b': {'c': 2, 'd': 3}})\n\n self.assertEqual(\n repr(c),\n \"\"\"Cursor(\n _obj={'a': {1: {(2, 3): 'z', 4: GenericTuple(x=1, y='a', z=(2, 'b')), '6': (7, 8)}, 'b': [1, 2, 3]}, 'z': -1},\n _changes={\n 'z': Cursor(\n _obj=-3,\n _changes={}\n ),\n 'a': Cursor(\n _obj={1: {(2, 3): 'z', 4: GenericTuple(x=1, y='a', z=(2, 'b')), '6': (7, 8)}, 'b': [1, 2, 3]},\n _changes={\n 'b': Cursor(\n _obj=None,\n _changes={}\n ),\n 1: Cursor(\n _obj={(2, 3): 'z', 4: GenericTuple(x=1, y='a', z=(2, 'b')), '6': (7, 8)},\n _changes={\n 4: Cursor(\n _obj=GenericTuple(x=1, y='a', z=(2, 'b')),\n _changes={\n 'x': Cursor(\n _obj=(2, 4, 6),\n _changes={}\n ),\n 'z': Cursor(\n _obj=(2, 'b'),\n _changes={\n 0: Cursor(\n _obj=FrozenDict({\n a: 1,\n b: {\n c: 2,\n d: 3,\n },\n }),\n _changes={}\n )\n }\n )\n }\n )\n }\n )\n }\n )\n }\n)\"\"\",\n )\n\n def test_magic_methods(self):\n def same_value(v1, v2):\n if isinstance(v1, tuple):\n return all([\n jnp.all(jax.tree_util.tree_map(lambda x, y: x == y, e1, e2))\n for e1, e2 in zip(v1, v2)\n ])\n return jnp.all(jax.tree_util.tree_map(lambda x, y: x == y, v1, v2))\n\n list_obj = [(1, 2), (3, 4)]\n for l, tuple_wrap in ((list_obj, lambda x: x), (tuple(list_obj), tuple)):\n c = cursor(l)\n # test __len__\n self.assertTrue(same_value(len(c), len(l)))\n # test __iter__\n for i, child_c in enumerate(c):\n child_c[1] += i + 1\n self.assertEqual(c.build(), tuple_wrap([(1, 3), (3, 6)]))\n # test __reversed__\n for i, child_c in enumerate(reversed(c)):\n child_c[1] += i + 1\n self.assertEqual(c.build(), tuple_wrap([(1, 5), (3, 7)]))\n # test __iter__ error\n with self.assertRaisesRegex(\n NotImplementedError,\n '__iter__ method only implemented for tuples and lists, not type \",\n ):\n c = cursor({'a': 1, 'b': 2})\n for key in c:\n c[key] *= -1\n # test __iter__ error\n with self.assertRaisesRegex(\n NotImplementedError,\n '__reversed__ method only implemented for tuples and lists, not type'\n \" \",\n ):\n c = cursor({'a': 1, 'b': 2})\n for key in reversed(c):\n c[key] *= -1\n\n for obj_value in (2, jnp.array([[1, -2], [3, 4]])):\n for c in (\n cursor(obj_value),\n cursor([obj_value])[0],\n cursor((obj_value,))[0],\n cursor({0: obj_value})[0],\n cursor(flax.core.freeze({0: obj_value}))[0],\n cursor(GenericTuple(x=obj_value)).x,\n cursor(GenericDataClass(x=obj_value)).x,\n ):\n # test __neg__\n self.assertTrue(same_value(-c, -obj_value))\n # test __pos__\n self.assertTrue(same_value(+c, +obj_value))\n # test __abs__\n self.assertTrue(same_value(abs(-c), abs(-obj_value)))\n # test __invert__\n self.assertTrue(same_value(~c, ~obj_value))\n # test __round__\n self.assertTrue(same_value(round(c + 0.123), round(obj_value + 0.123)))\n self.assertTrue(\n same_value(round(c + 0.123, 2), round(obj_value + 0.123, 2))\n )\n\n for other_value in (3, jnp.array([[5, 6], [7, 8]])):\n # test __add__\n self.assertTrue(same_value(c + other_value, obj_value + other_value))\n # test __radd__\n self.assertTrue(same_value(other_value + c, other_value + obj_value))\n # test __sub__\n self.assertTrue(same_value(c - other_value, obj_value - other_value))\n # test __rsub__\n self.assertTrue(same_value(other_value - c, other_value - obj_value))\n # test __mul__\n self.assertTrue(same_value(c * other_value, obj_value * other_value))\n # test __rmul__\n self.assertTrue(same_value(other_value * c, other_value * obj_value))\n # test __truediv__\n self.assertTrue(same_value(c / other_value, obj_value / other_value))\n # test __rtruediv__\n self.assertTrue(same_value(other_value / c, other_value / obj_value))\n # test __floordiv__\n self.assertTrue(\n same_value(c // other_value, obj_value // other_value)\n )\n # test __rfloordiv__\n self.assertTrue(\n same_value(other_value // c, other_value // obj_value)\n )\n # test __mod__\n self.assertTrue(same_value(c % other_value, obj_value % other_value))\n # test __rmod__\n self.assertTrue(same_value(other_value % c, other_value % obj_value))\n # test __divmod__\n self.assertTrue(\n same_value(divmod(c, other_value), divmod(obj_value, other_value))\n )\n # test __rdivmod__\n self.assertTrue(\n same_value(divmod(other_value, c), divmod(other_value, obj_value))\n )\n # test __pow__\n self.assertTrue(\n same_value(pow(c, other_value), pow(obj_value, other_value))\n )\n # test __rpow__\n self.assertTrue(\n same_value(pow(other_value, c), pow(other_value, obj_value))\n )\n # test __lshift__\n self.assertTrue(\n same_value(c << other_value, obj_value << other_value)\n )\n # test __rlshift__\n self.assertTrue(\n same_value(other_value << c, other_value << obj_value)\n )\n # test __rshift__\n self.assertTrue(\n same_value(c >> other_value, obj_value >> other_value)\n )\n # test __rrshift__\n self.assertTrue(\n same_value(other_value >> c, other_value >> obj_value)\n )\n # test __and__\n self.assertTrue(same_value(c & other_value, obj_value & other_value))\n # test __rand__\n self.assertTrue(same_value(other_value & c, other_value & obj_value))\n # test __xor__\n self.assertTrue(same_value(c ^ other_value, obj_value ^ other_value))\n # test __rxor__\n self.assertTrue(same_value(other_value ^ c, other_value ^ obj_value))\n # test __or__\n self.assertTrue(same_value(c | other_value, obj_value | other_value))\n # test __ror__\n self.assertTrue(same_value(other_value | c, other_value | obj_value))\n\n if isinstance(obj_value, jax.Array) and isinstance(\n other_value, jax.Array\n ):\n # test __matmul__\n self.assertTrue(\n same_value(c @ other_value, obj_value @ other_value)\n )\n # test __rmatmul__\n self.assertTrue(\n same_value(other_value @ c, other_value @ obj_value)\n )\n\n # test __lt__\n self.assertTrue(same_value(c < other_value, obj_value < other_value))\n self.assertTrue(same_value(other_value < c, other_value < obj_value))\n # test __le__\n self.assertTrue(\n same_value(c <= other_value, obj_value <= other_value)\n )\n self.assertTrue(\n same_value(other_value <= c, other_value <= obj_value)\n )\n # test __eq__\n self.assertTrue(\n same_value(c == other_value, obj_value == other_value)\n )\n self.assertTrue(\n same_value(other_value == c, other_value == obj_value)\n )\n # test __ne__\n self.assertTrue(\n same_value(c != other_value, obj_value != other_value)\n )\n self.assertTrue(\n same_value(other_value != c, other_value != obj_value)\n )\n # test __gt__\n self.assertTrue(same_value(c > other_value, obj_value > other_value))\n self.assertTrue(same_value(other_value > c, other_value > obj_value))\n # test __ge__\n self.assertTrue(\n same_value(c >= other_value, obj_value >= other_value)\n )\n self.assertTrue(\n same_value(other_value >= c, other_value >= obj_value)\n )\n\n def test_path(self):\n c = cursor(\n GenericTuple(\n x=[\n 0,\n {'a': 1, 'b': (2, 3), ('c', 'd'): [4, 5]},\n (100, 200),\n [3, 4, 5],\n ],\n y=train_state.TrainState.create(\n apply_fn=lambda x: x,\n params=freeze({'a': 1, 'b': (2, 3), 'c': [4, 5]}),\n tx=optax.adam(1e-3),\n ),\n )\n )\n self.assertEqual(c.x[1][('c', 'd')][0]._path, \".x[1][('c', 'd')][0]\")\n self.assertEqual(c.x[2][1]._path, '.x[2][1]')\n self.assertEqual(c.y.params['b'][1]._path, \".y.params['b'][1]\")\n\n # test path when first access type is item access\n c = cursor([1, GenericTuple('a', 2), (3, 4)])\n self.assertEqual(c[1].x._path, '[1].x')\n self.assertEqual(c[2][0]._path, '[2][0]')\n\n def test_traverse_tree(self):\n c = cursor(\n GenericTuple(\n x=[\n 0,\n {'a': 1, 'b': (2, 3), ('c', 'd'): [4, 5]},\n (100, 200),\n [3, 4, 5],\n ],\n y=3,\n )\n )\n\n def update_fn(path, value):\n if value == 4:\n return -4\n return value\n\n def cond_fn(path, value):\n return value == 3\n\n with self.assertRaisesRegex(\n TraverseTreeError,\n 'Both update_fn and cond_fn are None. Exactly one of them must be'\n ' None.',\n ):\n next(_traverse_tree((), c._obj))\n with self.assertRaisesRegex(\n TraverseTreeError,\n 'Both update_fn and cond_fn are not None. Exactly one of them must be'\n ' not None.',\n ):\n next(_traverse_tree((), c._obj, update_fn=update_fn, cond_fn=cond_fn))\n\n (p, v), (p2, v2) = _traverse_tree((), c._obj, update_fn=update_fn)\n self.assertEqual(\n p,\n (\n ('x', AccessType.ATTR),\n (1, AccessType.ITEM),\n (('c', 'd'), AccessType.ITEM),\n (0, AccessType.ITEM),\n ),\n )\n self.assertEqual(v, -4)\n self.assertEqual(\n p2, (('x', AccessType.ATTR), (3, AccessType.ITEM), (1, AccessType.ITEM))\n )\n self.assertEqual(v2, -4)\n\n p, p2, p3 = _traverse_tree((), c._obj, cond_fn=cond_fn)\n self.assertEqual(\n p,\n (\n ('x', AccessType.ATTR),\n (1, AccessType.ITEM),\n ('b', AccessType.ITEM),\n (1, AccessType.ITEM),\n ),\n )\n self.assertEqual(\n p2, (('x', AccessType.ATTR), (3, AccessType.ITEM), (0, AccessType.ITEM))\n )\n self.assertEqual(p3, (('y', AccessType.ATTR),))\n\n def test_set_and_build(self):\n # test regular dict and FrozenDict\n dict_obj = {'a': 1, 'b': (2, 3), 'c': [4, 5]}\n for d, freeze_wrap in ((dict_obj, lambda x: x), (freeze(dict_obj), freeze)):\n # set API\n self.assertEqual(\n cursor(d)['b'][0].set(10),\n freeze_wrap({'a': 1, 'b': (10, 3), 'c': [4, 5]}),\n )\n # build API\n c = cursor(d)\n c['b'][0] = 20\n c['a'] = (100, 200)\n d2 = c.build()\n self.assertEqual(\n d2, freeze_wrap({'a': (100, 200), 'b': (20, 3), 'c': [4, 5]})\n )\n self.assertEqual(\n dict_obj, {'a': 1, 'b': (2, 3), 'c': [4, 5]}\n ) # make sure original object is unchanged\n\n # test list and tuple\n list_obj = [0, dict_obj, (1, 2), [3, 4, 5]]\n for l, tuple_wrap in ((list_obj, lambda x: x), (tuple(list_obj), tuple)):\n # set API\n self.assertEqual(\n cursor(l)[1]['b'][0].set(10),\n tuple_wrap([0, {'a': 1, 'b': (10, 3), 'c': [4, 5]}, (1, 2), [3, 4, 5]]),\n )\n # build API\n c = cursor(l)\n c[1]['b'][0] = 20\n c[2] = (100, 200)\n l2 = c.build()\n self.assertEqual(\n l2,\n tuple_wrap(\n [0, {'a': 1, 'b': (20, 3), 'c': [4, 5]}, (100, 200), [3, 4, 5]]\n ),\n )\n self.assertEqual(\n list_obj, [0, {'a': 1, 'b': (2, 3), 'c': [4, 5]}, (1, 2), [3, 4, 5]]\n ) # make sure original object is unchanged\n\n # test TrainState\n state = train_state.TrainState.create(\n apply_fn=lambda x: x,\n params=dict_obj,\n tx=optax.adam(1e-3),\n )\n # set API\n self.assertEqual(\n cursor(state).params['b'][0].set(10).params,\n {'a': 1, 'b': (10, 3), 'c': [4, 5]},\n )\n # build API\n new_fn = lambda x: x + 1\n c = cursor(state)\n c.apply_fn = new_fn\n c.params['b'][0] = 20\n c.params['a'] = (100, 200)\n state2 = c.build()\n self.assertEqual(state2.apply_fn, new_fn)\n self.assertEqual(\n state2.params, {'a': (100, 200), 'b': (20, 3), 'c': [4, 5]}\n )\n\n self.assertEqual(\n dict_obj, {'a': 1, 'b': (2, 3), 'c': [4, 5]}\n ) # make sure original object is unchanged\n\n # test NamedTuple\n # set API\n t = GenericTuple(GenericTuple(0))\n self.assertEqual(cursor(t).x.x.set(1), GenericTuple(GenericTuple(1)))\n\n # build API\n c = cursor(t)\n c.x.x = 2\n c.x.y = 3\n c.y = 4\n t2 = c.build()\n self.assertEqual(t2, GenericTuple(GenericTuple(2, 3), 4))\n\n self.assertEqual(\n t, GenericTuple(GenericTuple(0))\n ) # make sure original object is unchanged\n\n def test_apply_update(self):\n # test list and tuple\n def update_fn(path, value):\n \"\"\"Multiply the first element of all leaf nodes of the pytree by -1.\"\"\"\n if path[-1] == '0' and isinstance(value, int):\n return value * -1\n return value\n\n for tuple_wrap in (lambda x: x, tuple):\n l = tuple_wrap([tuple_wrap([1, 2]), tuple_wrap([3, 4])])\n c = cursor(l)\n l2 = c.apply_update(update_fn).build()\n self.assertEqual(\n l2, tuple_wrap([tuple_wrap([-1, 2]), tuple_wrap([-3, 4])])\n )\n self.assertEqual(\n l, tuple_wrap([tuple_wrap([1, 2]), tuple_wrap([3, 4])])\n ) # make sure the original object is unchanged\n\n # test regular dict and FrozenDict\n def update_fn(path, value):\n \"\"\"Multiply all dense kernel params by 2 and add 1.\n Subtract the Dense_1 bias param by 1.\"\"\"\n if 'kernel' in path:\n return value * 2 + 1\n elif 'Dense_1' in path and 'bias' in path:\n return value - 1\n return value\n\n class Model(nn.Module):\n @nn.compact\n def __call__(self, x):\n x = nn.Dense(3)(x)\n x = nn.relu(x)\n x = nn.Dense(3)(x)\n x = nn.relu(x)\n x = nn.Dense(3)(x)\n x = nn.relu(x)\n return x\n\n for freeze_wrap in (lambda x: x, freeze):\n params = freeze_wrap(\n Model().init(jax.random.key(0), jnp.empty((1, 2)))['params']\n )\n\n c = cursor(params)\n params2 = c.apply_update(update_fn).build()\n for layer in ('Dense_0', 'Dense_1', 'Dense_2'):\n self.assertTrue(\n (params2[layer]['kernel'] == 2 * params[layer]['kernel'] + 1).all()\n )\n if layer == 'Dense_1':\n self.assertTrue(\n (params2[layer]['bias'] == jnp.array([-1, -1, -1])).all()\n )\n else:\n self.assertTrue(\n (params2[layer]['bias'] == params[layer]['bias']).all()\n )\n self.assertTrue(\n jax.tree_util.tree_all(\n jax.tree_util.tree_map(\n lambda x, y: (x == y).all(),\n params,\n freeze_wrap(\n Model().init(jax.random.key(0), jnp.empty((1, 2)))['params']\n ),\n )\n )\n ) # make sure original params are unchanged\n\n # test TrainState\n def update_fn(path, value):\n \"\"\"Replace params with empty dictionary.\"\"\"\n if 'params' in path:\n return {}\n return value\n\n state = train_state.TrainState.create(\n apply_fn=lambda x: x,\n params={'a': 1, 'b': 2},\n tx=optax.adam(1e-3),\n )\n c = cursor(state)\n state2 = c.apply_update(update_fn).build()\n self.assertEqual(state2.params, {})\n self.assertEqual(\n state.params, {'a': 1, 'b': 2}\n ) # make sure original params are unchanged\n\n # test NamedTuple\n def update_fn(path, value):\n \"\"\"Add 5 to all x-attribute values that are ints.\"\"\"\n if path[-1] == 'x' and isinstance(value, int):\n return value + 5\n return value\n\n t = GenericTuple(GenericTuple(0, 1), GenericTuple(2, 3))\n c = cursor(t)\n t2 = c.apply_update(update_fn).build()\n self.assertEqual(t2, GenericTuple(GenericTuple(5, 1), GenericTuple(7, 3)))\n self.assertEqual(\n t, GenericTuple(GenericTuple(0, 1), GenericTuple(2, 3))\n ) # make sure original object is unchanged\n\n def test_apply_update_root_node_unmodified(self):\n def update_fn(path, value):\n if isinstance(value, list):\n value = value.copy()\n value.append(-1)\n return value\n\n l = [[1, 2], [3, 4], 5]\n l2 = cursor(l).apply_update(update_fn).build()\n self.assertEqual(l2, [[1, 2, -1], [3, 4, -1], 5])\n\n def test_multi_modify(self):\n d = {'a': 1, 'b': (2, 3), 'c': [4, 5]}\n c = cursor(d)\n # test multiple changes on same element\n c['b'][0] = 6\n c['b'][0] = 7\n d2 = c.build()\n self.assertEqual(d2, {'a': 1, 'b': (7, 3), 'c': [4, 5]})\n # test nested changes\n c['a'] = (100, 200)\n c['a'][1] = -1\n d3 = c.build()\n self.assertEqual(d3, {'a': (100, -1), 'b': (7, 3), 'c': [4, 5]})\n\n def test_hidden_change(self):\n # test list\n l = [1, 2]\n c = cursor(l)\n c[0] = 100\n l[1] = -1\n l2 = c.build()\n self.assertEqual(\n l2, [100, -1]\n ) # change in l affects l2 (this is expected behavior)\n self.assertEqual(l, [1, -1])\n\n # test regular dict\n d = {'a': 1, 'b': 2}\n c = cursor(d)\n c['a'] = 100\n d['b'] = -1\n d2 = c.build()\n self.assertEqual(\n d2, {'a': 100, 'b': -1}\n ) # change in d affects d2 (this is expected behavior)\n self.assertEqual(d, {'a': 1, 'b': -1})\n\n # test TrainState\n params = {'a': 1, 'b': 2}\n state = train_state.TrainState.create(\n apply_fn=lambda x: x,\n params=params,\n tx=optax.adam(1e-3),\n )\n c = cursor(state)\n c.params['a'] = 100\n params['b'] = -1\n state2 = c.build()\n self.assertEqual(\n state2.params, {'a': 100, 'b': -1}\n ) # change in state affects state2 (this is expected behavior)\n self.assertEqual(state.params, {'a': 1, 'b': -1})\n\n def test_named_tuple_multi_access(self):\n t = GenericTuple(GenericTuple(0, 1), GenericTuple(2, 3))\n c = cursor(t)\n c.x.x = 4\n c[0].y = 5\n c.y.x = 6\n c.y[1] = 7\n self.assertEqual(\n c.build(), GenericTuple(GenericTuple(4, 5), GenericTuple(6, 7))\n )\n\n c[0][1] = -5\n self.assertEqual(\n c.build(), GenericTuple(GenericTuple(4, -5), GenericTuple(6, 7))\n )\n c.x[1] = -6\n self.assertEqual(\n c.build(), GenericTuple(GenericTuple(4, -6), GenericTuple(6, 7))\n )\n c.x.y = -7\n self.assertEqual(\n c.build(), GenericTuple(GenericTuple(4, -7), GenericTuple(6, 7))\n )\n c[0].y = -8\n self.assertEqual(\n c.build(), GenericTuple(GenericTuple(4, -8), GenericTuple(6, 7))\n )\n\n def test_find(self):\n c = cursor(\n GenericTuple(\n x=[\n 0,\n {'a': 1, 'b': (2, 3), ('c', 'd'): [4, 5]},\n (100, 200),\n [3, 4, 5],\n ],\n y=train_state.TrainState.create(\n apply_fn=lambda x: x,\n params=freeze({'a': 1, 'b': (2, 3), 'c': [4, 5]}),\n tx=optax.adam(1e-3),\n ),\n )\n )\n\n with self.assertRaisesRegex(\n CursorFindError,\n 'More than one object found given the conditions of the cond_fn\\\\. '\n 'The first two objects found have the following paths: '\n \"\\\\.x\\\\[1]\\\\['b'] and \\\\.y\\\\.params\\\\['b'] \",\n ):\n c.find(lambda path, value: 'b' in path and isinstance(value, tuple))\n with self.assertRaisesRegex(\n CursorFindError,\n 'No object found given the conditions of the cond_fn\\\\.',\n ):\n c.find(lambda path, value: 'b' in path and isinstance(value, str))\n\n self.assertEqual(\n c.find(lambda path, value: path.endswith('params/b'))[1].set(30).y.params,\n freeze({'a': 1, 'b': (2, 30), 'c': [4, 5]}),\n )\n\n def test_find_all(self):\n # test list and tuple\n def cond_fn(path, value):\n \"\"\"Get all lists that are not the first element in its parent.\"\"\"\n return path[-1] != '0' and isinstance(value, (tuple, list))\n\n for tuple_wrap in (lambda x: x, tuple):\n l = tuple_wrap(\n [tuple_wrap([1, 2]), tuple_wrap([3, 4]), tuple_wrap([5, 6])]\n )\n c = cursor(l)\n c2, c3 = c.find_all(cond_fn)\n c2[0] *= -1\n c3[1] *= -2\n self.assertEqual(\n c.build(),\n tuple_wrap(\n [tuple_wrap([1, 2]), tuple_wrap([-3, 4]), tuple_wrap([5, -12])]\n ),\n )\n self.assertEqual(\n l,\n tuple_wrap(\n [tuple_wrap([1, 2]), tuple_wrap([3, 4]), tuple_wrap([5, 6])]\n ),\n ) # make sure the original object is unchanged\n\n # test regular dict and FrozenDict\n def cond_fn(path, value):\n \"\"\"Get the second and third dense params.\"\"\"\n return 'Dense_1' in path or 'Dense_2' in path\n\n class Model(nn.Module):\n @nn.compact\n def __call__(self, x):\n x = nn.Dense(3)(x)\n x = nn.relu(x)\n x = nn.Dense(3)(x)\n x = nn.relu(x)\n x = nn.Dense(3)(x)\n x = nn.relu(x)\n return x\n\n for freeze_wrap in (lambda x: x, freeze):\n params = freeze_wrap(\n Model().init(jax.random.PRNGKey(0), jnp.empty((1, 2)))['params']\n )\n c = cursor(params)\n for i, c2 in enumerate(c.find_all(cond_fn)):\n self.assertEqual(\n c2['kernel'].set(123)[f'Dense_{i+1}'],\n freeze_wrap({'kernel': 123, 'bias': params[f'Dense_{i+1}']['bias']}),\n )\n self.assertTrue(\n jax.tree_util.tree_all(\n jax.tree_util.tree_map(\n lambda x, y: (x == y).all(),\n params,\n freeze_wrap(\n Model().init(jax.random.PRNGKey(0), jnp.empty((1, 2)))['params']\n ),\n )\n )\n ) # make sure original params are unchanged\n\n # test TrainState\n def cond_fn(path, value):\n \"\"\"Find TrainState params.\"\"\"\n return 'params' in path\n\n state = train_state.TrainState.create(\n apply_fn=lambda x: x,\n params={'a': 1, 'b': 2},\n tx=optax.adam(1e-3),\n )\n c = cursor(state)\n c2 = list(c.find_all(cond_fn))\n self.assertEqual(len(c2), 1)\n c2 = c2[0]\n self.assertEqual(c2['b'].set(-1).params, {'a': 1, 'b': -1})\n self.assertEqual(\n state.params, {'a': 1, 'b': 2}\n ) # make sure original params are unchanged\n\n # test NamedTuple\n def cond_fn(path, value):\n \"\"\"Get all GenericTuples that have int x-attribute values.\"\"\"\n return isinstance(value, GenericTuple) and isinstance(value.x, int)\n\n t = GenericTuple(\n GenericTuple(0, 'a'), GenericTuple(1, 'b'), GenericTuple('c', 2)\n )\n c = cursor(t)\n c2, c3 = c.find_all(cond_fn)\n c2.x += 5\n c3.x += 6\n self.assertEqual(\n c.build(),\n GenericTuple(\n GenericTuple(5, 'a'), GenericTuple(7, 'b'), GenericTuple('c', 2)\n ),\n )\n self.assertEqual(\n t,\n GenericTuple(\n GenericTuple(0, 'a'), GenericTuple(1, 'b'), GenericTuple('c', 2)\n ),\n ) # make sure original object is unchanged\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/download_dataset_metadata.sh", "content": "#!/bin/bash\n# If you get an error like:\n# Cloning into 'datasets'...\n# fatal: cannot change to 'https://github.com/tensorflow/datasets/': No such file or directory\n# error: failed to initialize sparse-checkout\n# This mean your git version is outdated. Just update it.\n\n\nset -e\n\n# Download TFDS metadata to flax/.tfds/metadata directory.\n# This allows the tests to specify the `data_dir` when using tfds.testing.mock_data().\ncd \"$( dirname \"$0\" )\"\n\nif [ -d \"../.tfds/metadata\" ]; then\n echo 'TFDS metadata already exists.';\nelse\n echo 'TFDS metadata does not exist. Downloading...';\n git clone --branch v4.8.2 --depth 3 --filter=blob:none --sparse https://github.com/tensorflow/datasets/\n cd datasets\n git sparse-checkout set tensorflow_datasets/testing/metadata\n mkdir ../../.tfds\n mv tensorflow_datasets/testing/metadata/ ../../.tfds/metadata/\n cd ..\n rm -rf datasets\nfi\n" }, { "path": "tests/early_stopping_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for flax.training.early_stopping.\"\"\"\n\n\nimport jax\nfrom absl.testing import absltest\n\nfrom flax.training import early_stopping\n\n# Parse absl flags test_srcdir and test_tmpdir.\njax.config.parse_flags_with_absl()\n\n\nclass EarlyStoppingTests(absltest.TestCase):\n def test_update(self):\n es = early_stopping.EarlyStopping(min_delta=0, patience=0)\n\n for i in range(2):\n improve_steps = 0\n for step in range(10):\n metric = 1.0\n es = es.update(metric)\n if not es.has_improved:\n improve_steps += 1\n if es.should_stop:\n break\n\n self.assertEqual(improve_steps, 1)\n self.assertEqual(step, 1)\n\n es = es.reset() # ensure object is reusable if reset.\n\n def test_patience(self):\n es = early_stopping.EarlyStopping(min_delta=0, patience=0)\n patient_es = early_stopping.EarlyStopping(min_delta=0, patience=6)\n for step in range(10):\n metric = 1.0\n es = es.update(metric)\n if es.should_stop:\n break\n\n self.assertEqual(step, 1)\n\n for patient_step in range(10):\n metric = 1.0\n patient_es = patient_es.update(metric)\n if patient_es.should_stop:\n break\n\n self.assertEqual(patient_step, 7)\n\n def test_delta(self):\n es = early_stopping.EarlyStopping(min_delta=0, patience=0)\n delta_es = early_stopping.EarlyStopping(min_delta=1e-3, patience=0)\n delta_patient_es = early_stopping.EarlyStopping(min_delta=1e-3, patience=1)\n metric = 1.0\n for step in range(100):\n metric -= 1e-4\n es = es.update(metric)\n if es.should_stop:\n break\n\n self.assertEqual(step, 99)\n\n metric = 1.0\n for step in range(100):\n metric -= 1e-4\n delta_es = delta_es.update(metric)\n if delta_es.should_stop:\n break\n\n self.assertEqual(step, 1)\n\n metrics = [\n 0.01,\n 0.005,\n 0.0033,\n 0.0025,\n 0.002,\n 0.0017,\n 0.0014,\n 0.0012,\n 0.0011,\n 0.001,\n ]\n improvement_steps = 0\n for step in range(10):\n metric = metrics[step]\n delta_patient_es = delta_patient_es.update(metric)\n if delta_patient_es.has_improved:\n improvement_steps += 1\n if delta_patient_es.should_stop:\n break\n\n self.assertEqual(improvement_steps, 4) # steps 0, 1, 2, 4\n self.assertEqual(step, 6)\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/flaxlib_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\n# TODO: Re-enable this test after setting up CI build for flaxlib CC.\n\n# from absl.testing import absltest\n# import flaxlib\n\n\n# class TestFlaxlib(absltest.TestCase):\n\n# def test_flaxlib(self):\n# self.assertEqual(flaxlib.sum_as_string(1, 2), '3')\n" }, { "path": "tests/import_test.ipynb", "content": "{\n \"cells\": [\n {\n \"attachments\": {},\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"# Test Import in Colab\\n\",\n \"\\n\",\n \"\\\"Run all\\\" to test that all the Flax imports work in head.\\n\",\n \"\\n\",\n \"Change runtime type as needed.\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# Colab runtimes are pre-built with JAX/Flax:\\n\",\n \"!pip freeze | egrep 'jax|flax'\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {\n \"tags\": [\n \"skip-execution\"\n ]\n },\n \"outputs\": [],\n \"source\": [\n \"# Install from head\\n\",\n \"!pip install git+https://github.com/google/flax.git\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": null,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# Check versions after installing Flax from Github:\\n\",\n \"!pip freeze | egrep 'jax|flax'\"\n ]\n },\n {\n \"cell_type\": \"code\",\n \"execution_count\": 9,\n \"metadata\": {},\n \"outputs\": [],\n \"source\": [\n \"# Verify we can import everything.\\n\",\n \"import flax\\n\",\n \"from flax.training import (checkpoints, dynamic_scale, early_stopping, lr_schedule,\\n\",\n \" orbax_utils, prefetch_iterator, train_state, common_utils)\\n\",\n \"from flax.metrics import tensorboard\"\n ]\n }\n ],\n \"metadata\": {\n \"language_info\": {\n \"codemirror_mode\": {\n \"name\": \"ipython\",\n \"version\": 3\n },\n \"file_extension\": \".py\",\n \"mimetype\": \"text/x-python\",\n \"name\": \"python\",\n \"nbconvert_exporter\": \"python\",\n \"pygments_lexer\": \"ipython3\",\n \"version\": \"3.9.15\"\n }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 2\n}\n" }, { "path": "tests/io_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for flax.io.\"\"\"\n\nimport os\nimport tempfile\n\nimport jax\nimport tensorflow as tf\nfrom absl.testing import absltest, parameterized\n\nfrom flax import errors, io\n\n# Parse absl flags test_srcdir and test_tmpdir.\njax.config.parse_flags_with_absl()\n\n\nclass IOTest(parameterized.TestCase):\n @parameterized.parameters(\n {'backend_mode': io.BackendMode.DEFAULT},\n {'backend_mode': io.BackendMode.TF},\n )\n def test_override(self, backend_mode):\n with io.override_mode(backend_mode):\n self.assertEqual(io.io_mode, backend_mode)\n\n @parameterized.parameters(\n {'write_mode': io.BackendMode.DEFAULT, 'read_mode': io.BackendMode.TF},\n {'write_mode': io.BackendMode.TF, 'read_mode': io.BackendMode.DEFAULT},\n )\n def test_GFile(self, write_mode, read_mode):\n test_string = b'testing write and read'\n with tempfile.TemporaryDirectory() as temp_dir_path:\n test_path = os.path.join(temp_dir_path, 'test')\n\n with io.override_mode(write_mode):\n with io.GFile(test_path, 'wb') as file:\n file.write(test_string)\n\n with io.override_mode(read_mode):\n with io.GFile(test_path, 'rb') as file:\n self.assertEqual(file.read(), test_string)\n\n def test_listdir(self):\n with tempfile.TemporaryDirectory() as temp_dir_path:\n os.mkdir(os.path.join(temp_dir_path, 'a'))\n os.mkdir(os.path.join(temp_dir_path, 'as'))\n os.mkdir(os.path.join(temp_dir_path, 'af'))\n os.mkdir(os.path.join(temp_dir_path, 'test'))\n os.mkdir(os.path.join(temp_dir_path, 'at'))\n\n with io.override_mode(io.BackendMode.DEFAULT):\n default_dir_set = set(io.listdir(temp_dir_path))\n\n with io.override_mode(io.BackendMode.TF):\n tf_dir_set = set(io.listdir(temp_dir_path))\n\n self.assertEqual(default_dir_set, tf_dir_set)\n\n @parameterized.parameters(\n {'create_temp_fn': tempfile.TemporaryDirectory},\n {'create_temp_fn': tempfile.NamedTemporaryFile},\n )\n def test_isdir(self, create_temp_fn):\n with create_temp_fn() as temp:\n path = temp.name if hasattr(temp, 'name') else temp\n\n with io.override_mode(io.BackendMode.DEFAULT):\n default_isdir = io.isdir(path)\n\n with io.override_mode(io.BackendMode.TF):\n tf_isdir = io.isdir(path)\n\n self.assertEqual(default_isdir, tf_isdir)\n\n def test_copy(self):\n test_string = b'testing copy'\n with tempfile.TemporaryDirectory() as temp_dir_path:\n test_path = os.path.join(temp_dir_path, 'test')\n copy1_path = os.path.join(temp_dir_path, 'copy1')\n copy2_path = os.path.join(temp_dir_path, 'copy2')\n\n with io.GFile(test_path, 'wb') as file:\n file.write(test_string)\n\n with io.override_mode(io.BackendMode.DEFAULT):\n io.copy(test_path, copy1_path)\n\n with io.override_mode(io.BackendMode.TF):\n io.copy(copy1_path, copy2_path)\n\n with io.GFile(copy2_path, 'rb') as file:\n self.assertEqual(file.read(), test_string)\n\n @parameterized.parameters(\n {\n 'backend_mode': io.BackendMode.DEFAULT,\n 'error_type': errors.AlreadyExistsError,\n },\n {\n 'backend_mode': io.BackendMode.TF,\n 'error_type': tf.errors.AlreadyExistsError,\n },\n )\n def test_copy_raises_error(self, backend_mode, error_type):\n with tempfile.NamedTemporaryFile() as temp_file:\n with io.override_mode(backend_mode):\n with self.assertRaises(error_type):\n io.copy(temp_file.name, temp_file.name)\n\n def test_rename(self):\n with tempfile.TemporaryDirectory() as temp_dir_path:\n test_path = os.path.join(temp_dir_path, 'test')\n rename1_path = os.path.join(temp_dir_path, 'rename1')\n rename2_path = os.path.join(temp_dir_path, 'rename2')\n\n with io.GFile(test_path, 'wb') as file:\n file.write(b'placeholder text')\n\n with io.override_mode(io.BackendMode.DEFAULT):\n io.rename(test_path, rename1_path)\n\n with io.override_mode(io.BackendMode.TF):\n io.rename(rename1_path, rename2_path)\n\n with io.GFile(rename2_path, 'rb') as file:\n self.assertTrue(os.path.exists(rename2_path))\n\n @parameterized.parameters(\n {\n 'backend_mode': io.BackendMode.DEFAULT,\n 'error_type': errors.AlreadyExistsError,\n },\n {\n 'backend_mode': io.BackendMode.TF,\n 'error_type': tf.errors.AlreadyExistsError,\n },\n )\n def test_rename_raises_error(self, backend_mode, error_type):\n with tempfile.NamedTemporaryFile() as temp_file:\n with io.override_mode(backend_mode):\n with self.assertRaises(error_type):\n io.rename(temp_file.name, temp_file.name)\n\n def test_exists(self):\n with tempfile.NamedTemporaryFile() as temp_file:\n with io.override_mode(io.BackendMode.DEFAULT):\n default_exists = io.exists(temp_file.name)\n\n with io.override_mode(io.BackendMode.TF):\n tf_exists = io.exists(temp_file.name)\n\n self.assertEqual(default_exists, tf_exists)\n\n @parameterized.parameters(\n {'backend_mode': io.BackendMode.DEFAULT},\n {'backend_mode': io.BackendMode.TF},\n )\n def test_makedirs(self, backend_mode):\n with tempfile.TemporaryDirectory() as temp_dir_path:\n test_dir_path = os.path.join(temp_dir_path, 'test_dir')\n\n with io.override_mode(backend_mode):\n io.makedirs(test_dir_path)\n self.assertTrue(\n os.path.exists(test_dir_path) and (os.path.isdir(test_dir_path))\n )\n\n def test_glob(self):\n with tempfile.TemporaryDirectory() as temp_dir_path:\n os.mkdir(os.path.join(temp_dir_path, 'a'))\n os.mkdir(os.path.join(temp_dir_path, 'as'))\n os.mkdir(os.path.join(temp_dir_path, 'af'))\n os.mkdir(os.path.join(temp_dir_path, 'test'))\n os.mkdir(os.path.join(temp_dir_path, 'at'))\n\n with io.override_mode(io.BackendMode.DEFAULT):\n default_glob_set = set(io.glob('a*/'))\n\n with io.override_mode(io.BackendMode.TF):\n tf_glob_set = set(io.glob('a*/'))\n\n self.assertEqual(default_glob_set, tf_glob_set)\n\n @parameterized.parameters(\n {'backend_mode': io.BackendMode.DEFAULT},\n {'backend_mode': io.BackendMode.TF},\n )\n def test_remove(self, backend_mode):\n with tempfile.TemporaryDirectory() as temp_dir_path:\n test_path = os.path.join(temp_dir_path, 'test')\n\n with io.GFile(test_path, 'wb') as file:\n file.write(b'placeholder text')\n\n with io.override_mode(backend_mode):\n io.remove(test_path)\n\n self.assertTrue(not os.path.exists(test_path))\n\n @parameterized.parameters(\n {'backend_mode': io.BackendMode.DEFAULT},\n {'backend_mode': io.BackendMode.TF},\n )\n def test_rmtree(self, backend_mode):\n with tempfile.TemporaryDirectory() as temp_dir_path:\n dir0_path = os.path.join(temp_dir_path, 'dir0')\n\n os.mkdir(dir0_path)\n os.mkdir(os.path.join(dir0_path, 'dir1'))\n os.mkdir(os.path.join(dir0_path, 'dir1', 'dir2'))\n os.mkdir(os.path.join(dir0_path, 'dir1', 'dir3'))\n os.mkdir(os.path.join(dir0_path, 'dir4'))\n os.mkdir(os.path.join(dir0_path, 'dir4', 'dir5'))\n os.mkdir(os.path.join(dir0_path, 'dir6'))\n\n with io.override_mode(backend_mode):\n io.rmtree(dir0_path)\n\n self.assertTrue(not os.path.exists(dir0_path))\n\n @parameterized.parameters(\n {'backend_mode': io.BackendMode.DEFAULT},\n {'backend_mode': io.BackendMode.TF},\n )\n def test_getsize(self, backend_mode):\n with tempfile.TemporaryDirectory() as temp_dir_path:\n test_path = os.path.join(temp_dir_path, 'test')\n\n content = b'placeholder text'\n with io.GFile(test_path, 'wb') as file:\n file.write(content)\n\n with io.override_mode(backend_mode):\n size = io.getsize(test_path)\n\n self.assertEqual(size, len(content))\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/jax_utils_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for flax.jax_utils.\"\"\"\n\nfrom functools import partial\nimport os\nos.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=4'\n\nfrom absl.testing import absltest\nfrom absl.testing import parameterized\nfrom flax import jax_utils\nimport jax\nimport jax.numpy as jnp\nimport numpy as np\n\nNDEV = 4\n\n\ndef assert_max_traces(n):\n \"\"\"Decorator to assert that a function is traced at most n times.\"\"\"\n from functools import wraps\n def decorator(fn):\n trace_count = {'count': 0}\n\n @wraps(fn)\n def wrapped(*args, **kwargs):\n trace_count['count'] += 1\n if trace_count['count'] > n:\n raise AssertionError(\n f\"Function was traced {trace_count['count']} times, \"\n f\"expected at most {n} traces\"\n )\n return fn(*args, **kwargs)\n\n wrapped.trace_count = trace_count\n return wrapped\n return decorator\n\n\nclass PadShardUnpadTest(parameterized.TestCase):\n BATCH_SIZES = [NDEV, NDEV + 1, NDEV - 1, 5 * NDEV, 5 * NDEV + 1, 5 * NDEV - 1]\n DTYPES = [np.float32, np.uint8, jax.numpy.bfloat16, np.int32]\n\n\n\n @parameterized.product(dtype=DTYPES, bs=BATCH_SIZES)\n def test_basics(self, dtype, bs):\n # Just tests that basic calling works without exploring caveats.\n @partial(jax_utils.pad_shard_unpad, static_argnums=())\n def add(a, b):\n b = jnp.asarray(b, dtype=dtype)\n return a + b\n\n x = np.arange(bs, dtype=dtype)\n y = add(x, 10 * x)\n self.assertEqual(y.dtype, x.dtype)\n np.testing.assert_allclose(np.float64(y), np.float64(x + 10 * x))\n\n @parameterized.product(dtype=DTYPES, bs=BATCH_SIZES)\n def test_trees(self, dtype, bs):\n # Just tests that basic calling works without exploring caveats.\n @partial(jax_utils.pad_shard_unpad, static_argnums=())\n def add(a, b):\n return a['a'] + b[0]\n\n x = jnp.arange(bs, dtype=dtype)\n y = add(dict(a=x), (10 * x,))\n self.assertEqual(y.dtype, x.dtype)\n np.testing.assert_allclose(np.float64(y), np.float64(x + 10 * x))\n\n @parameterized.parameters(DTYPES)\n def test_min_device_batch_avoids_recompile(self, dtype):\n @partial(jax_utils.pad_shard_unpad, static_argnums=())\n @jax.jit\n @assert_max_traces(n=1)\n def add(a, b):\n b = jnp.asarray(b, dtype=dtype)\n return a + b\n\n for bs in self.BATCH_SIZES:\n x = jnp.arange(bs, dtype=dtype)\n y = add(x, 10 * x, min_device_batch=9) # pylint: disable=unexpected-keyword-arg\n self.assertEqual(y.dtype, x.dtype)\n np.testing.assert_allclose(np.float64(y), np.float64(x + 10 * x))\n\n @parameterized.product(dtype=DTYPES, bs=BATCH_SIZES)\n def test_static_argnum(self, dtype, bs):\n @partial(jax_utils.pad_shard_unpad, static_argnums=(1,))\n def add(a, b):\n return a + jnp.asarray(b, dtype=dtype)\n\n x = jnp.arange(bs, dtype=dtype)\n y = add(x, 10)\n self.assertEqual(y.dtype, x.dtype)\n np.testing.assert_allclose(np.float64(y), np.float64(x + 10))\n\n @parameterized.product(dtype=DTYPES, bs=BATCH_SIZES)\n def test_static_argnames(self, dtype, bs):\n # In this test, leave static_argnums at the default value too, in order to\n # test the default/most canonical path where `params` are the first arg.\n @partial(jax_utils.pad_shard_unpad, static_argnames=('b',))\n def add(params, a, *, b):\n params = jnp.asarray(params, dtype=dtype)\n b = jnp.asarray(b, dtype=dtype)\n return params * a + b\n\n x = jnp.arange(bs, dtype=dtype)\n y = add(5, x, b=10)\n self.assertEqual(y.dtype, x.dtype)\n np.testing.assert_allclose(np.float64(y), np.float64(5 * x + 10))\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/linen/initializers_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for flax.linen.initializers.\"\"\"\n\nimport jax\nimport jax.numpy as jnp\nimport numpy as np\nfrom absl.testing import absltest, parameterized\nfrom jax import random\n\nfrom flax import linen as nn\nfrom flax.linen import initializers\n\n# Parse absl flags test_srcdir and test_tmpdir.\njax.config.parse_flags_with_absl()\n\n\nclass InitializersTest(parameterized.TestCase):\n @parameterized.parameters(\n {\n 'builder_fn': initializers.zeros_init,\n 'params_shape': (2, 3),\n 'expected_params': jnp.zeros((2, 3)),\n },\n {\n 'builder_fn': initializers.ones_init,\n 'params_shape': (3, 2),\n 'expected_params': jnp.ones((3, 2)),\n },\n )\n def test_call_builder(self, builder_fn, params_shape, expected_params):\n params = builder_fn()(random.key(42), params_shape, jnp.float32)\n np.testing.assert_allclose(params, expected_params)\n\n @parameterized.parameters(\n {\n 'builder_fn': initializers.zeros_init,\n 'expected_params': jnp.zeros((2, 5)),\n },\n {\n 'builder_fn': initializers.ones_init,\n 'expected_params': jnp.ones((2, 5)),\n },\n )\n def test_kernel_builder(self, builder_fn, expected_params):\n layer = nn.Dense(5, kernel_init=builder_fn())\n params = layer.init(random.key(42), jnp.empty((3, 2)))['params']\n np.testing.assert_allclose(params['kernel'], expected_params)\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/linen/kw_only_dataclasses_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for kw_only_dataclasses.\"\"\"\n\nimport dataclasses\nimport inspect\n\nfrom absl.testing import absltest\n\nfrom flax.linen import kw_only_dataclasses\n\n\nclass KwOnlyDataclassesTest(absltest.TestCase):\n def test_kwonly_args_moved_to_end(self):\n @kw_only_dataclasses.dataclass\n class TestClass:\n a: int = 1\n b: int = kw_only_dataclasses.field(default=2, kw_only=True)\n c: int = 3\n\n params = inspect.signature(TestClass.__init__).parameters\n self.assertEqual(list(params), ['self', 'a', 'c', 'b'])\n self.assertEqual(params['a'].default, 1)\n self.assertEqual(params['b'].default, 2)\n self.assertEqual(params['c'].default, 3)\n\n v1 = TestClass()\n self.assertDictEqual(dataclasses.asdict(v1), dict(a=1, b=2, c=3))\n\n v2 = TestClass(b=20)\n self.assertDictEqual(dataclasses.asdict(v2), dict(a=1, b=20, c=3))\n\n v3 = TestClass(1, 30)\n self.assertDictEqual(dataclasses.asdict(v3), dict(a=1, b=2, c=30))\n\n def test_base_optional_subclass_required(self):\n @kw_only_dataclasses.dataclass\n class Parent:\n a: int = kw_only_dataclasses.field(default=2, kw_only=True)\n\n @kw_only_dataclasses.dataclass\n class Child(Parent):\n b: int\n\n child_params = inspect.signature(Child.__init__).parameters\n self.assertEqual(list(child_params), ['self', 'b', 'a'])\n self.assertEqual(child_params['a'].default, 2)\n self.assertEqual(child_params['b'].default, inspect.Parameter.empty)\n\n v1 = Child(4)\n self.assertDictEqual(dataclasses.asdict(v1), dict(a=2, b=4))\n\n v2 = Child(4, a=5) # pylint: disable=too-many-function-args\n self.assertDictEqual(dataclasses.asdict(v2), dict(a=5, b=4))\n\n def test_subclass_overrides_base(self):\n # Note: if a base class declares a field as keyword-only, then\n # subclasses don't need to also declare it as keyword-only.\n\n @kw_only_dataclasses.dataclass\n class A:\n x: int = kw_only_dataclasses.field(default=1, kw_only=True)\n\n @kw_only_dataclasses.dataclass\n class B(A):\n size: float\n y: int = kw_only_dataclasses.field(default=3, kw_only=True)\n x: int = 2\n\n @kw_only_dataclasses.dataclass\n class C(B):\n name: str\n\n a_params = inspect.signature(A.__init__).parameters\n b_params = inspect.signature(B.__init__).parameters\n c_params = inspect.signature(C.__init__).parameters\n\n self.assertEqual(list(a_params), ['self', 'x'])\n self.assertEqual(list(b_params), ['self', 'size', 'x', 'y'])\n self.assertEqual(list(c_params), ['self', 'size', 'name', 'x', 'y'])\n\n self.assertEqual(a_params['x'].default, 1)\n self.assertEqual(b_params['x'].default, 2)\n self.assertEqual(b_params['y'].default, 3)\n self.assertEqual(b_params['size'].default, inspect.Parameter.empty)\n self.assertEqual(c_params['x'].default, 2)\n self.assertEqual(c_params['y'].default, 3)\n self.assertEqual(c_params['name'].default, inspect.Parameter.empty)\n self.assertEqual(c_params['size'].default, inspect.Parameter.empty)\n\n value = C(4, 'foo') # pylint: disable=too-many-function-args\n self.assertDictEqual(\n dataclasses.asdict(value), dict(name='foo', size=4, x=2, y=3)\n )\n\n def test_kwonly_marker(self):\n @kw_only_dataclasses.dataclass\n class A:\n x: float\n _: kw_only_dataclasses.KW_ONLY\n a: int = 5\n b: int = kw_only_dataclasses.field(default=2)\n c: int = kw_only_dataclasses.field(default=2, kw_only=True)\n\n @kw_only_dataclasses.dataclass\n class B(A):\n z: str\n\n a_params = inspect.signature(A.__init__).parameters\n b_params = inspect.signature(B.__init__).parameters\n self.assertEqual(list(a_params), ['self', 'x', 'a', 'b', 'c'])\n self.assertEqual(list(b_params), ['self', 'x', 'z', 'a', 'b', 'c'])\n\n def test_whatever(self):\n import abc\n from collections.abc import Iterator, Iterable\n import io\n from typing import Protocol, TypeVar\n\n T = TypeVar(\"T\")\n\n class CheckpointableIterator(Iterator[T], Protocol[T]):\n pass\n\n isinstance(io.TextIOBase, Iterable)\n\n from flax import linen as nn\n\n class Steppable(metaclass=abc.ABCMeta):\n path: str\n\n class SequenceLayer(nn.Module, Steppable):\n pass\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/linen/linen_activation_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for flax.linen.activation.\"\"\"\n\nfrom absl.testing import absltest\nfrom flax import linen as nn\nimport jax\nfrom jax import random\nimport jax.numpy as jnp\nimport numpy as np\n\n\n# Parse absl flags test_srcdir and test_tmpdir.\njax.config.parse_flags_with_absl()\n\n\nclass ActivationTest(absltest.TestCase):\n\n def test_prelu(self):\n rng = random.key(0)\n key, skey_1, skey_2 = jax.random.split(rng, 3)\n x = jax.random.uniform(skey_1, (4, 6, 5)) - 0.5\n act = nn.PReLU()\n y, params = act.init_with_output(skey_2, x)\n expected_y = jnp.where(x < 0, x * act.negative_slope_init, x)\n init_negative_slope = params['params']['negative_slope']\n expected_negative_slope = jnp.array(\n act.negative_slope_init, dtype=jnp.float32\n )\n\n self.assertEqual(y.shape, x.shape)\n np.testing.assert_array_almost_equal(expected_y, y)\n np.testing.assert_array_equal(init_negative_slope, expected_negative_slope)\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/linen/linen_attention_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for flax.linen.attention.\"\"\"\n\nfrom absl.testing import absltest, parameterized\nfrom flax import errors, jax_utils\nfrom flax import linen as nn\nfrom flax.core import pop\nimport jax\nfrom jax import lax, random\nfrom jax.nn import initializers\nimport jax.numpy as jnp\nimport numpy as np\n\ntry:\n # JAX v0.8.0 and newer\n from jax import enable_x64\nexcept ImportError:\n from jax.experimental import enable_x64\n\n# Parse absl flags test_srcdir and test_tmpdir.\njax.config.parse_flags_with_absl()\n\n\nclass AttentionTest(parameterized.TestCase):\n def test_multihead_self_attention(self):\n rng = random.key(0)\n x = jnp.ones((4, 6, 5))\n sa_module = nn.MultiHeadDotProductAttention(\n num_heads=8,\n qkv_features=16,\n kernel_init=initializers.ones,\n bias_init=initializers.zeros,\n deterministic=False,\n )\n y, _ = sa_module.init_with_output(rng, x)\n self.assertEqual(y.shape, x.shape)\n self.assertEqual(y.dtype, jnp.float32)\n\n def test_dtype_infer(self):\n rng = random.key(0)\n x = jnp.ones((4, 6, 5), jnp.complex64)\n sa_module = nn.MultiHeadDotProductAttention(\n num_heads=8,\n qkv_features=16,\n kernel_init=initializers.ones,\n bias_init=initializers.zeros,\n deterministic=False,\n )\n y, _ = sa_module.init_with_output(rng, x)\n self.assertEqual(y.shape, x.shape)\n self.assertEqual(y.dtype, jnp.complex64)\n\n def test_multihead_encoder_decoder_attention(self):\n rng = random.key(0)\n q = jnp.ones((4, 2, 3, 5))\n sa_module = nn.MultiHeadDotProductAttention(\n num_heads=8,\n qkv_features=16,\n kernel_init=initializers.ones,\n bias_init=initializers.zeros,\n deterministic=False,\n )\n y, _ = sa_module.init_with_output(rng, q)\n self.assertEqual(y.shape, q.shape)\n\n def test_mha_out_initializers(self):\n rng = random.key(0)\n q = jnp.ones((4, 2, 3, 5))\n sa_module = nn.MultiHeadDotProductAttention(\n num_heads=8,\n qkv_features=16,\n kernel_init=initializers.ones,\n out_kernel_init=initializers.zeros,\n bias_init=initializers.zeros,\n out_bias_init=initializers.ones,\n deterministic=False,\n )\n variables = sa_module.init(rng, q)\n params = variables['params']\n # test kernels\n np.testing.assert_allclose(params['query']['kernel'], 1.0)\n np.testing.assert_allclose(params['key']['kernel'], 1.0)\n np.testing.assert_allclose(params['value']['kernel'], 1.0)\n np.testing.assert_allclose(params['out']['kernel'], 0.0)\n # test biases\n np.testing.assert_allclose(params['query']['bias'], 0.0)\n np.testing.assert_allclose(params['key']['bias'], 0.0)\n np.testing.assert_allclose(params['value']['bias'], 0.0)\n np.testing.assert_allclose(params['out']['bias'], 1.0)\n\n def test_multihead_self_attention_w_dropout(self):\n rng = random.key(0)\n x = jnp.ones((4, 2, 3, 5))\n sa_module = nn.MultiHeadDotProductAttention(\n num_heads=8,\n qkv_features=16,\n kernel_init=initializers.ones,\n bias_init=initializers.zeros,\n dropout_rate=0.1,\n deterministic=False,\n )\n rng1, rng2 = random.split(rng)\n rngs = {'params': rng1, 'dropout': rng2}\n y, _ = sa_module.init_with_output(rngs, x)\n self.assertEqual(y.shape, x.shape)\n\n def test_multihead_self_attention_explicit_dropout(self):\n def clone(key):\n return jax.tree.map(jax.random.clone, key)\n\n class Foo(nn.Module):\n attention_kwargs: dict\n\n @nn.compact\n def __call__(self, x, dropout_rng=None):\n a = nn.MultiHeadDotProductAttention(**self.attention_kwargs)(\n x, x, dropout_rng=dropout_rng\n )\n if dropout_rng is not None:\n dropout_rng = clone(dropout_rng)\n b = nn.MultiHeadDotProductAttention(**self.attention_kwargs)(\n x, x, dropout_rng=dropout_rng\n )\n return a, b\n\n module = Foo(\n dict(\n num_heads=8,\n qkv_features=16,\n kernel_init=initializers.ones,\n bias_init=initializers.zeros,\n dropout_rate=0.5,\n deterministic=False,\n )\n )\n rng1, rng2, rng3, rng4 = random.split(random.key(0), 4)\n x = jnp.ones((4, 2, 3, 5))\n rngs = {'params': rng1, 'dropout': rng2}\n v = module.init(rngs, x)\n a, b = module.apply(v, x, rngs=clone(rngs))\n c, d = module.apply(v, x, rngs={'dropout': clone(rng2)})\n e, f = module.apply(v, x, rngs={'dropout': rng3})\n self.assertFalse((a == b).all())\n self.assertTrue((a == c).all())\n self.assertTrue((b == d).all())\n self.assertFalse((a == e).all())\n self.assertFalse((b == f).all())\n a, b = module.apply(v, x, rngs=clone(rngs), dropout_rng=rng4)\n self.assertTrue((a == b).all())\n a, b = module.apply(v, x, dropout_rng=clone(rng4))\n self.assertTrue((a == b).all())\n self.assertTrue(a.shape == b.shape == x.shape)\n\n def test_multihead_self_attention_w_dropout_disabled(self):\n rng = random.key(0)\n x = jnp.ones((4, 2, 3, 5))\n sa_module0 = nn.MultiHeadDotProductAttention(\n num_heads=8,\n qkv_features=16,\n kernel_init=initializers.ones,\n bias_init=initializers.zeros,\n dropout_rate=0.0,\n deterministic=True,\n )\n rng1, rng2, rng3, rng4 = random.split(rng, 4)\n rngs1 = {'params': rng1, 'dropout': rng2}\n rngs2 = {'params': rng3, 'dropout': rng4}\n y1, vs = sa_module0.init_with_output(rngs1, x)\n y2, _ = sa_module0.init_with_output(rngs2, x)\n np.testing.assert_allclose(y1, y2)\n y3 = sa_module0.apply(vs, x, rngs=rngs1)\n y4 = sa_module0.apply(vs, x, rngs=rngs2)\n np.testing.assert_allclose(y3, y4)\n sa_module1 = nn.MultiHeadDotProductAttention(\n num_heads=8,\n qkv_features=16,\n kernel_init=initializers.ones,\n bias_init=initializers.zeros,\n dropout_rate=0.0,\n )\n y5 = sa_module1.apply(vs, x, deterministic=True, rngs=rngs1)\n y6 = sa_module1.apply(vs, x, deterministic=True, rngs=rngs2)\n np.testing.assert_allclose(y5, y6)\n sa_module2 = nn.MultiHeadDotProductAttention(\n num_heads=8,\n qkv_features=16,\n kernel_init=initializers.ones,\n bias_init=initializers.zeros,\n dropout_rate=0.5,\n )\n y7 = sa_module2.apply(vs, x, deterministic=True, rngs=rngs1)\n y8 = sa_module2.apply(vs, x, deterministic=True, rngs=rngs2)\n np.testing.assert_allclose(y7, y8)\n\n def test_causal_mask_1d(self):\n \"\"\"Tests autoregressive masking for 1d attention.\"\"\"\n x = jnp.ones((3, 16)) # (bs1, length)\n mask_1d = nn.attention.make_causal_mask(x)\n ts = np.arange(16)\n mask_1d_simple = (ts[:, None] >= ts[None, :])[None, None, :, :]\n mask_1d_simple = jnp.broadcast_to(mask_1d_simple, (3, 1, 16, 16))\n np.testing.assert_allclose(\n mask_1d,\n mask_1d_simple,\n )\n\n @parameterized.parameters([((5,), (1,)), ((6, 5), (2,))])\n def test_decoding(self, spatial_shape, attn_dims):\n bs = 2\n num_heads = 3\n num_features = 4\n rng = random.key(0)\n key1, key2 = random.split(rng)\n inputs = random.normal(\n key1, (bs,) + spatial_shape + (num_heads * num_features,)\n )\n module = nn.MultiHeadDotProductAttention(\n num_heads=num_heads,\n qkv_features=num_heads * num_features,\n precision=lax.Precision.HIGHEST,\n deterministic=False,\n decode=False,\n )\n decode_module = module.clone(decode=True)\n\n initial_vars = decode_module.init(key2, inputs)\n state, params = pop(initial_vars, 'params')\n causal_mask = nn.attention.make_causal_mask(jnp.ones((bs,) + spatial_shape))\n y_ref = jax.jit(lambda x, y: module.apply(initial_vars, x, mask=y))(\n inputs, causal_mask\n )\n\n # feed the inputs sequentially to simulate decoding\n def body_fn(state, x):\n y, state = decode_module.apply(\n {'params': params, **state}, x, mutable=['cache']\n )\n return state, y\n\n # scan_in_dim supports scanning multiple dims\n _, y = jax_utils.scan_in_dim(\n body_fn, state, inputs, axis=attn_dims, keepdims=True\n )\n\n np.testing.assert_allclose(y_ref, y, atol=1e-5)\n\n def test_autoregressive_receptive_field_1d(self):\n \"\"\"Tests the autoregressive self-attention receptive field.\"\"\"\n rng = random.key(0)\n rng1, rng2 = random.split(rng, num=2)\n\n length = 10\n dim = 1\n num_heads = 1\n input_shape = (1, length, dim)\n inputs = random.normal(rng2, input_shape)\n\n module = nn.MultiHeadDotProductAttention(\n num_heads=num_heads,\n kernel_init=jax.nn.initializers.ones,\n deterministic=False,\n )\n\n initial_vars = module.init(rng1, inputs)\n causal_mask = nn.attention.make_causal_mask(jnp.ones(input_shape[:-1]))\n\n def model_loss(inputs, pos):\n out = module.apply(initial_vars, inputs, mask=causal_mask)\n assert out.shape == input_shape\n assert len(out.shape) == 3\n return out[0, pos, :].sum()\n\n grad_fn = jax.jit(jax.grad(model_loss))\n\n def get_receptive_field_1d(pos):\n g = grad_fn(inputs, pos)[0, :, :]\n return jnp.any((jnp.abs(g) > 1e-5).astype(jnp.uint32), axis=-1)\n\n for i in range(length):\n deps = get_receptive_field_1d(i)\n assert (deps[:i] == 1).all(), (\n 'Receptive Field Error: Some of the '\n 'previous postions are not reachable '\n 'in autoregressive self-attention.'\n )\n if i != length - 1:\n k = i + 1\n assert (deps[k:] == 0).all(), (\n 'Receptive Field Error: Some of the '\n 'future postions are reachable in '\n 'autoregressive self-attention.'\n )\n\n def test_multihead_kv_args(self):\n key1, key2 = random.split(random.key(0), 2)\n query = random.uniform(key1, (3, 5))\n key_value = random.uniform(key2, (9, 5))\n module = nn.MultiHeadDotProductAttention(\n num_heads=8,\n qkv_features=16,\n kernel_init=initializers.ones,\n bias_init=initializers.zeros,\n deterministic=False,\n )\n key = lambda: random.key(43279)\n y0, v0 = module.init_with_output(\n key(), query, inputs_k=key_value, inputs_v=key_value\n )\n y1, v1 = module.init_with_output(key(), query, inputs_k=key_value)\n with self.assertWarnsRegex(\n DeprecationWarning, 'The inputs_kv arg will be deprecated soon.'\n ):\n y2, v2 = module.init_with_output(key(), query, inputs_kv=key_value)\n self.assertTrue((y0 == y1).all() and (y1 == y2).all())\n self.assertTrue(\n jax.tree_util.tree_all(\n jax.tree_util.tree_map(\n lambda x, y, z: (x == y).all() and (y == z).all(), v0, v1, v2\n )\n )\n )\n\n with self.assertRaisesRegex(\n ValueError, '`inputs_k` cannot be None if `inputs_v` is not None.'\n ):\n y3, v3 = module.init_with_output(key(), query, inputs_v=key_value)\n with self.assertRaisesRegex(\n ValueError,\n 'If either `inputs_k` or `inputs_v` is not None, `inputs_kv` must be None.',\n ):\n y3, v3 = module.init_with_output(\n key(), query, inputs_kv=key_value, inputs_v=key_value\n )\n with self.assertRaisesRegex(\n ValueError,\n 'If either `inputs_k` or `inputs_v` is not None, `inputs_kv` must be None.',\n ):\n y3, v3 = module.init_with_output(\n key(), query, key_value, key_value, inputs_kv=key_value\n )\n\n def test_multihead_mask_warning(self):\n rng = random.key(0)\n rng1, rng2 = random.split(rng, num=2)\n\n length = 10\n dim = 1\n num_heads = 1\n input_shape = (1, length, dim)\n query = key = random.normal(rng2, input_shape)\n\n module = nn.MultiHeadDotProductAttention(\n num_heads=num_heads,\n kernel_init=jax.nn.initializers.ones,\n deterministic=False,\n )\n\n initial_vars = module.init(rng1, query, key)\n causal_mask = nn.attention.make_causal_mask(jnp.ones(input_shape[:-1]))\n\n module.apply(initial_vars, query, key, mask=causal_mask)\n with self.assertWarnsRegex(\n DeprecationWarning,\n \"the function signature of MultiHeadDotProductAttention's `__call__` method has changed\",\n ):\n with self.assertRaises(errors.ScopeParamShapeError):\n module.apply(initial_vars, query, key, causal_mask)\n\n def test_multihead_sow_attention_weights(self):\n rng = random.key(0)\n x = jnp.ones((4, 6, 5))\n\n class Model(nn.Module):\n attention_kwargs: dict\n\n @nn.compact\n def __call__(self, x, sow_weights=False):\n x = nn.MultiHeadDotProductAttention(**self.attention_kwargs)(\n x, sow_weights=sow_weights\n )\n x = nn.MultiHeadDotProductAttention(**self.attention_kwargs)(x)\n x = nn.MultiHeadDotProductAttention(**self.attention_kwargs)(\n x, sow_weights=sow_weights\n )\n return x\n\n module = Model(\n dict(\n num_heads=8,\n qkv_features=16,\n kernel_init=initializers.ones,\n bias_init=initializers.zeros,\n deterministic=False,\n )\n )\n v = module.init(rng, x)\n _, intermediates = module.apply(\n v, x, mutable=['intermediates'], sow_weights=True\n )\n self.assertEqual(\n intermediates['intermediates']['MultiHeadDotProductAttention_0'][\n 'attention_weights'\n ][0].shape,\n (4, 8, 6, 6),\n )\n self.assertNotIn(\n 'MultiHeadDotProductAttention_1', intermediates['intermediates']\n )\n self.assertEqual(\n intermediates['intermediates']['MultiHeadDotProductAttention_2'][\n 'attention_weights'\n ][0].shape,\n (4, 8, 6, 6),\n )\n _, intermediates = module.apply(\n v, x, mutable=['intermediates'], sow_weights=False\n )\n self.assertNotIn('intermediates', intermediates)\n\n def test_autoregressive_decode_with_x64(self):\n with enable_x64():\n x = jnp.ones((1, 4, 4))\n module = nn.MultiHeadDotProductAttention(\n num_heads=2,\n qkv_features=4,\n decode=True\n )\n\n rng = random.PRNGKey(0)\n variables = module.init(rng, x, x, x)\n params, cache = variables['params'], variables['cache']\n y1, updates = module.apply(\n { 'params': params, 'cache': cache },\n x[:, :1, :],\n mutable=['cache']\n )\n cache = updates['cache']\n y2, updates = module.apply(\n { 'params': params, 'cache': cache },\n x[:, 1:2, :],\n mutable=['cache']\n )\n assert y1.shape == (1, 1, 4)\n assert y2.shape == (1, 1, 4)\n\n def test_attention_alias_equivalence(self):\n key1, key2 = random.split(random.key(0), 2)\n query = random.uniform(key1, (3, 5))\n key_value = random.uniform(key2, (9, 5))\n attention_kwargs = dict(\n num_heads=8,\n qkv_features=16,\n kernel_init=initializers.lecun_normal(),\n bias_init=initializers.uniform(),\n deterministic=False,\n )\n module1 = nn.MultiHeadDotProductAttention(**attention_kwargs)\n module2 = nn.MultiHeadAttention(**attention_kwargs)\n key = lambda: random.key(43279)\n out1, v1 = module1.init_with_output(key(), query, key_value)\n out2, v2 = module2.init_with_output(key(), query, key_value, key_value)\n self.assertTrue((out1 == out2).all())\n self.assertTrue(\n jax.tree_util.tree_all(\n jax.tree_util.tree_map(lambda x, y: (x == y).all(), v1, v2)\n )\n )\n\n def test_attention_alias_submodule(self):\n key1, key2 = random.split(random.key(0), 2)\n query = random.uniform(key1, (3, 5))\n key_value = random.uniform(key2, (9, 5))\n attention_kwargs = dict(\n num_heads=8,\n qkv_features=16,\n kernel_init=initializers.lecun_normal(),\n bias_init=initializers.uniform(),\n deterministic=False,\n )\n\n class Foo1(nn.Module):\n attention_kwargs: dict\n\n @nn.compact\n def __call__(self, query, key):\n return nn.MultiHeadDotProductAttention(**self.attention_kwargs)(\n query, key\n )\n\n class Foo2(nn.Module):\n attention_kwargs: dict\n\n @nn.compact\n def __call__(self, query, key, value):\n return nn.MultiHeadAttention(**self.attention_kwargs)(query, key, value)\n\n key = lambda: random.key(5478392)\n module1 = Foo1(attention_kwargs)\n module2 = Foo2(attention_kwargs)\n out1, v1 = module1.init_with_output(key(), query, key_value)\n out2, v2 = module2.init_with_output(key(), query, key_value, key_value)\n\n # test different output and variables if layer names are different\n self.assertTrue((out1 != out2).all())\n v2['params']['MultiHeadDotProductAttention_0'] = v2['params'][\n 'MultiHeadAttention_0'\n ]\n del v2['params']['MultiHeadAttention_0']\n self.assertTrue(\n jax.tree_util.tree_all(\n jax.tree_util.tree_map(lambda x, y: (x != y).all(), v1, v2)\n )\n )\n\n # test same output if variables are the same\n v2 = jax.tree_util.tree_map(lambda x: x, v1)\n v2['params']['MultiHeadAttention_0'] = v2['params'][\n 'MultiHeadDotProductAttention_0'\n ]\n del v2['params']['MultiHeadDotProductAttention_0']\n out2 = module2.apply(v2, query, key_value, key_value)\n self.assertTrue((out1 == out2).all())\n\n # test same output and variables if names are the same\n class Foo2(nn.Module):\n attention_kwargs: dict\n\n @nn.compact\n def __call__(self, query, key, value):\n return nn.MultiHeadAttention(\n **self.attention_kwargs, name='MultiHeadDotProductAttention_0'\n )(query, key, value)\n\n module2 = Foo2(attention_kwargs)\n out2, v2 = module2.init_with_output(key(), query, key_value, key_value)\n self.assertTrue((out1 == out2).all())\n self.assertTrue(\n jax.tree_util.tree_all(\n jax.tree_util.tree_map(lambda x, y: (x == y).all(), v1, v2)\n )\n )\n\n @parameterized.parameters(\n {'force_fp32': True, 'attn_weights_dtype': jnp.float32},\n {'force_fp32': False, 'attn_weights_dtype': jnp.bfloat16},\n )\n def test_mixed_precision_multihead_attention(\n self, force_fp32, attn_weights_dtype\n ):\n input_key, params_key, dropout_key = random.split(random.key(0), 3)\n x = random.uniform(input_key, (2, 4))\n attention_kwargs = dict(\n num_heads=2,\n qkv_features=4,\n kernel_init=initializers.lecun_normal(),\n bias_init=initializers.uniform(),\n force_fp32_for_softmax=force_fp32,\n deterministic=False,\n dtype=jnp.bfloat16,\n )\n mha = nn.MultiHeadDotProductAttention(**attention_kwargs)\n init_vars = mha.init({'params': params_key, 'dropout': dropout_key}, x)\n _, updated_vars = mha.apply(\n init_vars, x, mutable=['intermediates'], sow_weights=True\n )\n self.assertEqual(\n updated_vars['intermediates']['attention_weights'][0].dtype,\n attn_weights_dtype,\n )\n\n @parameterized.parameters(\n (lax.Precision.DEFAULT, None),\n (None, jax.lax.dot_general),\n )\n def test_dot_product_attention_precision_and_einsum_override(\n self, precision, einsum_dot_general\n ):\n # Test that we raise a ValueError if the user specifies both\n # `precision` and/or `einsum_dot_general` and `qk_attn_weights_einsum`.\n einsum_cls = lambda: jnp.einsum\n self.assertRaises(\n ValueError,\n nn.dot_product_attention,\n query=jnp.ones((1, 4, 2)),\n key=jnp.ones((1, 4, 2)),\n value=jnp.ones((1, 4, 2)),\n precision=precision,\n einsum_dot_general=einsum_dot_general,\n qk_attn_weights_einsum=einsum_cls,\n attn_weights_value_einsum=einsum_cls,\n )\n\n @parameterized.parameters(\n (lambda: jax.lax.dot_general, None),\n (None, lambda: jax.lax.dot_general),\n )\n def test_dot_product_attention_specify_einsums_together(\n self, qk_attn_weights_einsum, attn_weights_value_einsum\n ):\n # Test that we raise a ValueError if the user specifies only one of\n # `qk_attn_weights_einsum` and `attn_weights_value_einsum`.\n self.assertRaises(\n ValueError,\n nn.dot_product_attention,\n query=jnp.ones((1, 4, 2)),\n key=jnp.ones((1, 4, 2)),\n value=jnp.ones((1, 4, 2)),\n qk_attn_weights_einsum=qk_attn_weights_einsum,\n attn_weights_value_einsum=attn_weights_value_einsum,\n )\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/linen/linen_batch_apply_test.py", "content": "# Copyright 2023 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for flax.linen.batch_apply.\"\"\"\n\nimport jax\nimport jax.numpy as jnp\nimport numpy as np\nfrom absl.testing import absltest, parameterized\n\nfrom flax import linen as nn\n\n# Parse absl flags test_srcdir and test_tmpdir.\njax.config.parse_flags_with_absl()\n\n\nclass BatchApplyTest(parameterized.TestCase):\n @parameterized.parameters(\n {'fn': lambda a, b: a + b.reshape(1, -1)},\n {'fn': lambda a, b: jnp.dot(a, b)},\n )\n def test_batchapply(self, fn):\n a = jax.random.normal(jax.random.key(0), [2, 3, 4])\n b = jax.random.normal(jax.random.key(1), [4])\n\n def raises(a, b):\n if len(a.shape) != 2:\n raise ValueError('a must be shape 2')\n if len(b.shape) != 1:\n raise ValueError('b must be shape 1')\n return fn(a, b)\n\n out = nn.BatchApply(raises)(a, b)\n expected_merged_leading = raises(a.reshape(2 * 3, 4), b)\n expected = expected_merged_leading.reshape(\n (2, 3) + expected_merged_leading.shape[1:]\n )\n np.testing.assert_array_equal(out, expected)\n\n def test_batchapply_accepts_float(self):\n def raises(a, b):\n if len(a.shape) != 2:\n raise ValueError('a must be shape 2')\n return a + b\n\n out = nn.BatchApply(raises)(jnp.ones([2, 3, 4]), 2.0)\n np.testing.assert_array_equal(out, 3 * jnp.ones([2, 3, 4]))\n\n def test_batchapply_accepts_none(self):\n def raises(a, b):\n if a is not None:\n raise ValueError('a must be None.')\n if len(b.shape) != 2:\n raise ValueError('b must be shape 2')\n return 3 * b\n\n out = nn.BatchApply(raises)(None, jnp.ones([2, 3, 4]))\n np.testing.assert_array_equal(out, 3 * jnp.ones([2, 3, 4]))\n\n def test_batchapply_raises(self):\n with self.assertRaisesRegex(ValueError, 'requires at least one input'):\n nn.BatchApply(lambda: 1)()\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/linen/linen_combinators_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for flax.linen.combinators.\"\"\"\n\nfrom typing import Any\nfrom collections.abc import Sequence\n\nimport jax\nimport numpy as np\nfrom absl.testing import absltest\nfrom jax import numpy as jnp\nfrom jax import random\n\nfrom flax import linen as nn\n\n# Parse absl flags test_srcdir and test_tmpdir.\njax.config.parse_flags_with_absl()\n\n\nclass MLP(nn.Module):\n layer_sizes: Sequence[int]\n activation: Any | None = None\n activation_final: Any | None = None\n\n @nn.compact\n def __call__(self, inputs):\n x = inputs\n for layer_size in self.layer_sizes[:-1]:\n x = nn.Dense(\n features=layer_size, kernel_init=nn.initializers.ones_init()\n )(x)\n if self.activation is not None:\n x = self.activation(x)\n x = nn.Dense(\n features=self.layer_sizes[-1], kernel_init=nn.initializers.ones_init()\n )(x)\n if self.activation_final is None:\n return x\n return self.activation_final(x)\n\n\nclass AttentionTuple(nn.Module):\n num_heads: int = 2\n qkv_features: int = 16\n\n @nn.compact\n def __call__(self, query, key_value):\n output = nn.MultiHeadDotProductAttention(\n num_heads=self.num_heads, qkv_features=self.qkv_features\n )(query, key_value)\n return output, key_value\n\n\nclass AttentionDict(nn.Module):\n num_heads: int = 2\n qkv_features: int = 16\n\n @nn.compact\n def __call__(self, query, key_value):\n output = nn.MultiHeadDotProductAttention(\n num_heads=self.num_heads, qkv_features=self.qkv_features\n )(query, key_value)\n return dict(query=output, key_value=key_value)\n\n\nclass SequentialTest(absltest.TestCase):\n def test_construction(self):\n sequential = nn.Sequential([nn.Dense(4), nn.Dense(2)])\n key1, key2 = random.split(random.key(0), 2)\n x = random.uniform(key1, (3, 1, 5))\n params = sequential.init(key2, x)\n output = sequential.apply(params, x)\n self.assertEqual(output.shape, (3, 1, 2))\n\n def test_fails_if_layers_empty(self):\n sequential = nn.Sequential([])\n with self.assertRaisesRegex(ValueError, 'Empty Sequential module'):\n sequential.init(random.key(42), jnp.ones((3, 5)))\n\n def test_same_output_as_mlp(self):\n sequential = nn.Sequential(\n [\n nn.Dense(4, kernel_init=nn.initializers.ones_init()),\n nn.Dense(8, kernel_init=nn.initializers.ones_init()),\n nn.Dense(2, kernel_init=nn.initializers.ones_init()),\n ]\n )\n mlp = MLP(layer_sizes=[4, 8, 2])\n\n key1, key2 = random.split(random.key(0), 2)\n x = random.uniform(key1, (3, 5))\n params_1 = sequential.init(key2, x)\n params_2 = mlp.init(key2, x)\n\n output_1 = sequential.apply(params_1, x)\n output_2 = mlp.apply(params_2, x)\n np.testing.assert_array_equal(output_1, output_2)\n\n def test_same_output_as_mlp_with_activation(self):\n sequential = nn.Sequential(\n [\n nn.Dense(4, kernel_init=nn.initializers.ones_init()),\n nn.relu,\n nn.Dense(8, kernel_init=nn.initializers.ones_init()),\n nn.relu,\n nn.Dense(2, kernel_init=nn.initializers.ones_init()),\n nn.log_softmax,\n ]\n )\n\n mlp = MLP(\n layer_sizes=[4, 8, 2],\n activation=nn.relu,\n activation_final=nn.log_softmax,\n )\n\n key1, key2 = random.split(random.key(0), 2)\n x = random.uniform(key1, (3, 5))\n params_1 = sequential.init(key2, x)\n params_2 = mlp.init(key2, x)\n\n output_1 = sequential.apply(params_1, x)\n output_2 = mlp.apply(params_2, x)\n np.testing.assert_array_equal(output_1, output_2)\n\n def test_tuple_output(self):\n sequential = nn.Sequential(\n [\n AttentionTuple(),\n AttentionTuple(),\n ]\n )\n\n key1, key2, key3 = random.split(random.key(0), 3)\n query = random.uniform(key1, (3, 5))\n key_value = random.uniform(key2, (9, 5))\n params_1 = sequential.init(key3, query, key_value)\n outputs = sequential.apply(params_1, query, key_value)\n np.testing.assert_equal(len(outputs), 2)\n out_query, out_key_value = outputs\n np.testing.assert_equal(out_query.shape, (3, 5))\n np.testing.assert_equal(out_key_value.shape, (9, 5))\n\n def test_dict_output(self):\n sequential = nn.Sequential(\n [\n AttentionDict(),\n AttentionDict(),\n ]\n )\n\n key1, key2, key3 = random.split(random.key(0), 3)\n query = random.uniform(key1, (3, 5))\n key_value = random.uniform(key2, (9, 5))\n params_1 = sequential.init(key3, query, key_value)\n outputs = sequential.apply(params_1, query, key_value)\n np.testing.assert_equal(len(outputs), 2)\n out_query, out_key_value = outputs['query'], outputs['key_value']\n np.testing.assert_equal(out_query.shape, (3, 5))\n np.testing.assert_equal(out_key_value.shape, (9, 5))\n\n def test_sequential_compact(self):\n\n mlp = nn.Sequential([\n lambda x: nn.Dense(x.shape[-1])(x),\n nn.relu,\n lambda x: nn.Dense(x.shape[-1])(x),\n nn.relu,\n lambda x: nn.Dense(x.shape[-1])(x),\n ])\n\n params = mlp.init(random.key(0), jnp.ones((3, 5)))['params']\n\n self.assertIn('Dense_0', params)\n self.assertIn('Dense_1', params)\n self.assertIn('Dense_2', params)\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/linen/linen_dtypes_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for flax.linen.dtypes.\"\"\"\n\n\nfrom absl.testing import absltest\nfrom jax import numpy as jnp\n\nfrom flax.linen import dtypes\n\ntry:\n # JAX v0.8.0 and newer\n from jax import enable_x64\nexcept ImportError:\n from jax.experimental import enable_x64\n\ndefault_float_dtype = jnp.result_type(1.0)\n\n\nclass DtypesTest(absltest.TestCase):\n def test_no_inexact_dtype(self):\n i32 = jnp.int32(1.0)\n self.assertEqual(dtypes.canonicalize_dtype(i32, inexact=False), jnp.int32)\n\n def test_inexact_dtype(self):\n with enable_x64():\n i64 = jnp.int64(1)\n self.assertEqual(dtypes.canonicalize_dtype(i64), jnp.float32)\n i32 = jnp.int32(1)\n self.assertEqual(dtypes.canonicalize_dtype(i32), jnp.float32)\n i16 = jnp.int16(1.0)\n self.assertEqual(dtypes.canonicalize_dtype(i16), jnp.float32)\n\n def test_explicit_downcast(self):\n f32 = jnp.float32(1.0)\n (x,) = dtypes.promote_dtype(f32, dtype=jnp.float16)\n self.assertEqual(x.dtype, jnp.float16)\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/linen/linen_linear_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for flax.linen.linear.\"\"\"\n\nimport functools\n\nimport jax\nimport jax.numpy as jnp\nimport numpy as np\nfrom absl.testing import absltest, parameterized\nfrom jax import random\nfrom jax.nn import initializers\n\nfrom flax import linen as nn\n\n# Parse absl flags test_srcdir and test_tmpdir.\njax.config.parse_flags_with_absl()\n\n\nclass LinearTest(parameterized.TestCase):\n def test_dense(self):\n rng = dict(params=random.key(0))\n x = jnp.ones((1, 3))\n dense_module = nn.Dense(\n features=4,\n kernel_init=initializers.ones,\n bias_init=initializers.ones,\n )\n y, _ = dense_module.init_with_output(rng, x)\n self.assertEqual(y.shape, (1, 4))\n self.assertEqual(y.dtype, jnp.float32)\n np.testing.assert_allclose(y, np.full((1, 4), 4.0))\n\n def test_dense_extra_batch_dims(self):\n rng = dict(params=random.key(0))\n x = jnp.ones((1, 2, 3))\n dense_module = nn.Dense(\n features=4,\n kernel_init=initializers.ones,\n bias_init=initializers.ones,\n )\n y, _ = dense_module.init_with_output(rng, x)\n np.testing.assert_allclose(y, np.full((1, 2, 4), 4.0))\n\n def test_dense_no_bias(self):\n rng = dict(params=random.key(0))\n x = jnp.ones((1, 3))\n dense_module = nn.Dense(\n features=4,\n use_bias=False,\n kernel_init=initializers.ones,\n )\n y, _ = dense_module.init_with_output(rng, x)\n np.testing.assert_allclose(y, np.full((1, 4), 3.0))\n\n def test_dense_is_dense_general(self):\n x = jax.random.normal(random.key(0), (5, 3))\n dense_module = nn.Dense(\n features=4,\n use_bias=True,\n bias_init=initializers.normal(),\n )\n y1, _ = dense_module.init_with_output(dict(params=random.key(1)), x)\n dg_module = nn.DenseGeneral(\n features=4,\n use_bias=True,\n bias_init=initializers.normal(),\n )\n y2, _ = dg_module.init_with_output(dict(params=random.key(1)), x)\n\n np.testing.assert_allclose(y1, y2)\n\n def test_dense_general_batch_dim_raises(self):\n rng = dict(params=random.key(0))\n x = jnp.ones((1, 3, 2, 5))\n with self.assertRaises(ValueError):\n dg_module = nn.DenseGeneral(\n features=4,\n batch_dims=(0, 2),\n kernel_init=initializers.ones,\n bias_init=initializers.ones,\n )\n dg_module.init_with_output(rng, x)\n\n def test_dense_general_two_out(self):\n rng = dict(params=random.key(0))\n x = jnp.ones((1, 3))\n dg_module = nn.DenseGeneral(\n features=(2, 2),\n kernel_init=initializers.ones,\n bias_init=initializers.ones,\n )\n y, _ = dg_module.init_with_output(rng, x)\n np.testing.assert_allclose(y, np.full((1, 2, 2), 4.0))\n\n def test_dense_general_two_in(self):\n rng = dict(params=random.key(0))\n x = jnp.ones((1, 2, 2))\n dg_module = nn.DenseGeneral(\n features=3,\n axis=(-2, 2),\n kernel_init=initializers.ones,\n bias_init=initializers.ones,\n )\n y, _ = dg_module.init_with_output(rng, x)\n np.testing.assert_allclose(y, np.full((1, 3), 5.0))\n\n def test_dense_general_batch_dim(self):\n rng = dict(params=random.key(0))\n x = jnp.ones((2, 1, 3, 5))\n\n state = {'counter': 0.0}\n\n def _counter_init(rng, shape, dtype, state):\n del rng, dtype\n state['counter'] += 1.0\n return jnp.full(shape, state['counter'])\n\n counter_init = functools.partial(_counter_init, state=state)\n\n dg_module = nn.DenseGeneral(\n features=7,\n axis=(3, -2),\n batch_dims=0,\n bias_init=initializers.ones,\n kernel_init=counter_init,\n )\n y, _ = dg_module.init_with_output(rng, x)\n target = np.full((2, 1, 7), 16.0)\n np.testing.assert_allclose(y, target)\n\n @parameterized.parameters(\n [\n ((-2, 3), (), 'bijk,jklm->bilm'),\n ((3, -2), (), 'bijk,jklm->bilm'),\n ((-2, 3), (0,), 'bijk,bjklm->bilm'),\n ]\n )\n def test_dense_general_vs_numpy(self, axis, batch_dims, einsum_expr):\n rng = dict(params=random.key(0))\n x = jnp.ones((16, 8, 9, 10))\n\n dg_module = nn.DenseGeneral(\n features=(11, 12),\n axis=axis,\n batch_dims=batch_dims,\n bias_init=initializers.ones,\n kernel_init=initializers.normal(),\n )\n y, initial_params = dg_module.init_with_output(rng, x)\n target = np.einsum(einsum_expr, x, initial_params['params']['kernel']) + 1.0\n np.testing.assert_allclose(y, target, atol=1e-6)\n\n def test_complex_params_dense(self):\n dense = nn.Dense(features=2, param_dtype=jnp.complex64)\n x = jnp.ones((1, 2), jnp.float32)\n variables = dense.init(random.key(0), x)\n self.assertEqual(variables['params']['kernel'].dtype, jnp.complex64)\n self.assertEqual(variables['params']['bias'].dtype, jnp.complex64)\n y = dense.apply(variables, x)\n self.assertEqual(y.dtype, jnp.complex64)\n\n def test_complex_input_dense(self):\n dense = nn.Dense(features=2)\n x = jnp.ones((1, 2), jnp.complex64)\n variables = dense.init(random.key(0), x)\n self.assertEqual(variables['params']['kernel'].dtype, jnp.float32)\n self.assertEqual(variables['params']['bias'].dtype, jnp.float32)\n y = dense.apply(variables, x)\n self.assertEqual(y.dtype, jnp.complex64)\n\n @parameterized.parameters(\n [\n (\n 'abc,cde->abde',\n (3, 4, 5),\n (5, 6, 7),\n (3, 4, 6, 7),\n (6, 7),\n (1, 1, 6, 7),\n ),\n (\n 'abcd,abcd->abcd',\n (3, 4, 5, 6),\n (3, 4, 5, 6),\n (3, 4, 5, 6),\n (3, 4, 5, 6),\n (3, 4, 5, 6),\n ),\n (\n 'abcd,abcd->abd',\n (3, 4, 5, 6),\n (3, 4, 5, 6),\n (3, 4, 6),\n (3, 4, 6),\n (3, 4, 6),\n ),\n (\n 'abcd,cdef->abef',\n (3, 4, 5, 6),\n (5, 6, 7, 8),\n (3, 4, 7, 8),\n (7, 8),\n (1, 1, 7, 8),\n ),\n (\n 'abcd,eafc->bfed',\n (3, 4, 5, 6),\n (7, 3, 8, 5),\n (4, 8, 7, 6),\n (8, 7),\n (1, 8, 7, 1),\n ),\n (\n 'abcd,cbedf->abfe',\n (3, 4, 5, 6),\n (5, 4, 7, 6, 8),\n (3, 4, 8, 7),\n (4, 8, 7),\n (1, 4, 8, 7),\n ),\n (\n 'ab...,bc...->ac...',\n (3, 4, 6),\n (4, 5, 6),\n (3, 5, 6),\n (5, 6),\n (1, 5, 6),\n ),\n (\n 'd...ab,bc...->ad...c',\n (8, 6, 7, 3, 4),\n (4, 5, 6, 7),\n (3, 8, 6, 7, 5),\n (6, 7, 5),\n (1, 1, 6, 7, 5),\n ),\n (\n 'd...ab,bc...->adc',\n (8, 6, 7, 3, 4),\n (4, 5, 6, 7),\n (3, 8, 5),\n (5,),\n (1, 1, 5),\n ),\n (\n 'abd...,bc...->ac...',\n (3, 4, 6),\n (4, 5, 6),\n (3, 5, 6),\n (5, 6),\n (1, 5, 6),\n ),\n (\n 'a...d,ej...f->adef',\n (3, 4, 5, 6),\n (7, 4, 5, 8),\n (3, 6, 7, 8),\n (7, 8),\n (1, 1, 7, 8),\n ),\n (\n 'ab...d,ej...f->ad...f',\n (3, 4, 5, 6),\n (7, 4, 5, 8),\n (3, 6, 5, 8),\n (5, 8),\n (1, 1, 5, 8),\n ),\n ]\n )\n def test_einsum_init_apply(\n self,\n einsum_str,\n lhs_shape,\n rhs_shape,\n expected_result_shape,\n expected_bias_shape,\n bias_broadcast_shape,\n ):\n layer = nn.Einsum(rhs_shape, einsum_str, bias_init=nn.initializers.normal())\n x = jax.random.normal(jax.random.key(0), lhs_shape)\n\n v = layer.init(jax.random.key(1), x)\n self.assertEqual(rhs_shape, v['params']['kernel'].shape)\n self.assertEqual(expected_bias_shape, v['params']['bias'].shape)\n\n out = layer.apply(v, x)\n self.assertEqual(out.shape, expected_result_shape)\n expected_out = jnp.einsum(einsum_str, x, v['params']['kernel']) + v[\n 'params'\n ]['bias'].reshape(bias_broadcast_shape)\n np.testing.assert_allclose(out, expected_out)\n\n @parameterized.parameters(\n [\n (\n ('abd,bce->ace', 'abd,bc...->ac...', 'abd...,bc...->ac...'),\n (3, 4, 6),\n (4, 5, 6),\n ),\n (\n (\n 'abcd,ejcf->adef',\n 'ab...d,ej...f->adef',\n 'ab...d,e...f->adef',\n 'a...d,ej...f->adef',\n ),\n (3, 4, 5, 6),\n (7, 4, 5, 8),\n ),\n (\n ('abcd,ejcf->adcf', 'ab...d,ej...f->ad...f'),\n (3, 4, 5, 6),\n (7, 4, 5, 8),\n ),\n ]\n )\n def test_einsum_ellipsis_equivalence(\n self, einsum_str_list, lhs_shape, rhs_shape\n ):\n x = jax.random.uniform(jax.random.key(0), lhs_shape)\n layer = nn.Einsum(\n rhs_shape, einsum_str_list[0], bias_init=nn.initializers.normal()\n )\n v = layer.init(jax.random.key(1), x)\n out = layer.apply(v, x)\n\n for einsum_str in einsum_str_list[1:]:\n layer2 = nn.Einsum(\n rhs_shape, einsum_str, bias_init=nn.initializers.normal()\n )\n v2 = layer2.init(jax.random.key(1), x)\n np.testing.assert_allclose(v['params']['kernel'], v2['params']['kernel'])\n np.testing.assert_allclose(v['params']['bias'], v2['params']['bias'])\n np.testing.assert_allclose(out, layer2.apply(v2, x))\n\n def test_einsum_str_arg(self):\n einsum_str = 'abc,cde->abde'\n x = jax.random.normal(jax.random.key(0), (3, 4, 5))\n\n constructed_layer = nn.Einsum(\n (5, 6, 7), einsum_str, bias_init=nn.initializers.normal()\n )\n constructed_v = constructed_layer.init(jax.random.key(1), x)\n constructed_out = constructed_layer.apply(constructed_v, x)\n\n called_layer = nn.Einsum((5, 6, 7), bias_init=nn.initializers.normal())\n called_v = called_layer.init(jax.random.key(1), x, einsum_str)\n called_out = called_layer.apply(called_v, x, einsum_str)\n\n np.testing.assert_allclose(\n constructed_v['params']['kernel'], called_v['params']['kernel']\n )\n np.testing.assert_allclose(\n constructed_v['params']['bias'], called_v['params']['bias']\n )\n np.testing.assert_allclose(constructed_out, called_out)\n\n with self.assertRaisesWithLiteralMatch(\n ValueError,\n 'Parameter \"einsum_str\" was passed to the constructor and at call time. Should be passed just once.',\n ):\n constructed_layer.init(jax.random.key(1), x, einsum_str)\n with self.assertRaisesWithLiteralMatch(\n ValueError,\n 'Parameter \"einsum_str\" must be passed to the constructor or at call time.',\n ):\n called_layer.init(jax.random.key(1), x)\n\n def test_einsum_space_str(self):\n x = jax.random.normal(jax.random.key(0), (3, 4, 5))\n\n layer1 = nn.Einsum(\n (5, 6, 7), 'abc,cde->abde', bias_init=nn.initializers.normal()\n )\n v1 = layer1.init(jax.random.key(1), x)\n out1 = layer1.apply(v1, x)\n\n layer2 = nn.Einsum(\n (5, 6, 7),\n ' ab c , c d e - >a b d e ',\n bias_init=nn.initializers.normal(),\n )\n v2 = layer2.init(jax.random.key(1), x)\n out2 = layer2.apply(v1, x)\n\n np.testing.assert_allclose(v1['params']['kernel'], v2['params']['kernel'])\n np.testing.assert_allclose(v1['params']['bias'], v2['params']['bias'])\n np.testing.assert_allclose(out1, out2)\n\n @parameterized.parameters(\n [\n ('abc,cde', '`einsum_str` equation must be explicit and include \"->\".'),\n (\n 'abc->',\n f'`einsum_str` equation must have exactly two operands and therefore, exactly one comma character, instead of 0',\n ),\n (\n 'abc,cde,efg->abdfg',\n f'`einsum_str` equation must have exactly two operands and therefore, exactly one comma character, instead of 2',\n ),\n ]\n )\n def test_einsum_error(self, einsum_str, error_msg):\n x = jax.random.normal(jax.random.key(0), (3, 4, 5))\n\n layer = nn.Einsum((5, 6, 7), einsum_str)\n with self.assertRaisesRegex(ValueError, error_msg):\n layer.init(jax.random.key(1), x)\n\n @parameterized.product(use_bias=(True, False))\n def test_conv(self, use_bias):\n rng = dict(params=random.key(0))\n x = jnp.ones((1, 8, 3))\n conv_module = nn.Conv(\n features=4,\n use_bias=use_bias,\n kernel_size=(3,),\n padding='VALID',\n kernel_init=initializers.ones,\n bias_init=initializers.ones,\n )\n y, initial_params = conv_module.init_with_output(rng, x)\n self.assertEqual(initial_params['params']['kernel'].shape, (3, 3, 4))\n expected = 10.0 if use_bias else 9.0\n np.testing.assert_allclose(y, np.full((1, 6, 4), expected))\n\n @parameterized.product(use_bias=(True, False))\n def test_multibatch_input_conv(self, use_bias):\n rng = dict(params=random.key(0))\n x = jnp.ones((2, 5, 8, 3))\n conv_module = nn.Conv(\n features=4,\n use_bias=use_bias,\n kernel_size=(3,),\n padding='VALID',\n kernel_init=initializers.ones,\n bias_init=initializers.ones,\n )\n y, initial_params = conv_module.init_with_output(rng, x)\n self.assertEqual(initial_params['params']['kernel'].shape, (3, 3, 4))\n expected = 10.0 if use_bias else 9.0\n np.testing.assert_allclose(y, np.full((2, 5, 6, 4), expected))\n\n def test_conv_local(self):\n rng = dict(params=random.key(0))\n x = jnp.ones((1, 8, 2))\n conv_module = nn.ConvLocal(\n features=4,\n kernel_size=(3,),\n padding='VALID',\n kernel_init=initializers.ones,\n bias_init=initializers.ones,\n )\n y, initial_params = conv_module.init_with_output(rng, x)\n self.assertEqual(initial_params['params']['kernel'].shape, (6, 3 * 2, 4))\n np.testing.assert_allclose(y, np.full((1, 6, 4), 7.0))\n\n def test_single_input_conv(self):\n rng = dict(params=random.key(0))\n x = jnp.ones((8, 3))\n conv_module = nn.Conv(\n features=4,\n kernel_size=(3,),\n padding='VALID',\n kernel_init=initializers.ones,\n bias_init=initializers.ones,\n )\n y, initial_params = conv_module.init_with_output(rng, x)\n self.assertEqual(initial_params['params']['kernel'].shape, (3, 3, 4))\n np.testing.assert_allclose(y, np.full((6, 4), 10.0))\n\n def test_single_input_masked_conv(self):\n rng = dict(params=random.key(0))\n x = jnp.ones((8, 3))\n m = jnp.tril(jnp.ones((3, 3, 4)))\n conv_module = nn.Conv(\n features=4,\n kernel_size=(3,),\n padding='VALID',\n mask=m,\n kernel_init=initializers.ones,\n bias_init=initializers.ones,\n )\n expected = jnp.array(\n [\n [10.0, 7.0, 4.0, 1.0],\n [10.0, 7.0, 4.0, 1.0],\n [10.0, 7.0, 4.0, 1.0],\n [10.0, 7.0, 4.0, 1.0],\n [10.0, 7.0, 4.0, 1.0],\n [10.0, 7.0, 4.0, 1.0],\n ]\n )\n y, initial_params = conv_module.init_with_output(rng, x)\n self.assertEqual(initial_params['params']['kernel'].shape, (3, 3, 4))\n np.testing.assert_allclose(y, expected)\n\n def test_single_input_conv_local(self):\n rng = dict(params=random.key(0))\n x = jnp.ones((8, 2))\n conv_module = nn.ConvLocal(\n features=4,\n kernel_size=(3,),\n padding='VALID',\n kernel_init=initializers.ones,\n bias_init=initializers.ones,\n )\n y, initial_params = conv_module.init_with_output(rng, x)\n self.assertEqual(initial_params['params']['kernel'].shape, (6, 3 * 2, 4))\n np.testing.assert_allclose(y, np.full((6, 4), 7.0))\n\n def test_group_conv(self):\n rng = dict(params=random.key(0))\n x = jnp.ones((1, 8, 4))\n conv_module = nn.Conv(\n features=4,\n kernel_size=(3,),\n feature_group_count=2,\n padding='VALID',\n kernel_init=initializers.ones,\n bias_init=initializers.ones,\n )\n y, initial_params = conv_module.init_with_output(rng, x)\n self.assertEqual(initial_params['params']['kernel'].shape, (3, 2, 4))\n np.testing.assert_allclose(y, np.full((1, 6, 4), 7.0))\n\n @parameterized.product(\n n_batch=(1, 3),\n n_features=(1, 2),\n kernel_size=(1, 2, 3, 9),\n n_input_features=(1, 3),\n input_size=(1, 8, 16),\n module=(nn.Conv, nn.ConvLocal),\n )\n def test_circular_conv_1d_constant(\n self,\n n_batch,\n n_features,\n kernel_size,\n n_input_features,\n input_size,\n module,\n ):\n \"\"\"\n Test 1D convolution with circular padding: filter with all elements equal\n to 1 applied on an input with all elements equal to 1.\n Result should have the same shape as input (except for the feature\n dimension) and have all elements equal to\n `n_input_features * kernel_lin_size`.\n \"\"\"\n rng = dict(params=random.key(0))\n x = jnp.ones((n_batch, input_size, n_input_features))\n conv_module = module(\n features=n_features,\n kernel_size=(kernel_size,),\n padding='CIRCULAR',\n kernel_init=initializers.ones,\n bias_init=initializers.zeros,\n )\n y, initial_params = conv_module.init_with_output(rng, x)\n\n kernel_shape = self._get_kernel_shape(\n x.shape, (kernel_size,), module, n_features\n )\n\n self.assertEqual(\n initial_params['params']['kernel'].shape,\n kernel_shape,\n )\n correct_ans = np.full(\n (n_batch, input_size, n_features), kernel_size * n_input_features\n )\n np.testing.assert_allclose(y, correct_ans)\n\n def _get_kernel_shape(self, input_shape, kernel_size, module, n_features):\n if module == nn.Conv:\n kernel_shape = kernel_size + (input_shape[-1], n_features)\n elif module == nn.ConvLocal:\n kernel_shape = input_shape[1:-1] + (\n input_shape[-1] * np.prod(kernel_size),\n n_features,\n )\n else:\n raise ValueError(module)\n return kernel_shape\n\n @parameterized.product(\n n_batch=(1, 3),\n n_features=(1, 2, 10),\n kernel_lin_size=(1, 2, 3, 9),\n n_input_features=(1, 5),\n input_x_size=(14,),\n input_y_size=(5, 10),\n module=(nn.Conv, nn.ConvLocal),\n )\n def test_circular_conv_2d_constant(\n self,\n n_batch,\n n_features,\n kernel_lin_size,\n n_input_features,\n input_x_size,\n input_y_size,\n module,\n ):\n \"\"\"\n Test 2D convolution with circular padding: square filter with all elements\n equal to 1 applied on an input with all elements equal to 1.\n Result should have the same shape as input (except for the feature\n dimension) and have all elements equal to\n `n_input_features * kernel_lin_size^2`.\n \"\"\"\n rng = dict(params=random.key(0))\n x = jnp.ones((n_batch, input_x_size, input_y_size, n_input_features))\n kernel_size = (kernel_lin_size, kernel_lin_size)\n conv_module = module(\n features=n_features,\n kernel_size=kernel_size,\n padding='CIRCULAR',\n kernel_init=initializers.ones,\n bias_init=initializers.zeros,\n )\n y, initial_params = conv_module.init_with_output(rng, x)\n\n kernel_shape = self._get_kernel_shape(\n x.shape, kernel_size, module, n_features\n )\n\n self.assertEqual(\n initial_params['params']['kernel'].shape,\n kernel_shape,\n )\n correct_ans = np.full(\n (n_batch, input_x_size, input_y_size, n_features),\n kernel_lin_size * kernel_lin_size * n_input_features,\n )\n np.testing.assert_allclose(y, correct_ans)\n\n def test_circular_conv_1d_custom(self):\n \"\"\"Test 1d convolution with circular padding and a stride.\"\"\"\n rng = dict(params=random.key(0))\n x = np.arange(1, 6)\n x = np.expand_dims(x, (0, 2))\n kernel = np.array((1, 2, 1))\n kernel = np.expand_dims(kernel, (1, 2))\n\n conv_module = nn.Conv(\n features=1,\n kernel_size=(3,),\n strides=(3,),\n padding='CIRCULAR',\n kernel_init=lambda *_: kernel,\n bias_init=initializers.zeros,\n )\n y, initial_params = conv_module.init_with_output(rng, x)\n\n self.assertEqual(initial_params['params']['kernel'].shape, (3, 1, 1))\n # Compare with manually computed convolution\n correct_ans = np.array((5 + 2 * 1 + 2, 3 + 2 * 4 + 5))\n correct_ans = np.expand_dims(correct_ans, (0, 2))\n np.testing.assert_allclose(y, correct_ans)\n\n def test_circular_conv_local_1d_custom(self):\n \"\"\"\n Test 1d local convolution with circular padding and a stride\n \"\"\"\n rng = dict(params=random.key(0))\n x = np.arange(1, 6)\n x = np.expand_dims(x, (0, 2))\n kernel = np.array(((-1, 2, 3), (4, 5, 6)))\n kernel = np.expand_dims(kernel, (2,))\n conv_module = nn.ConvLocal(\n features=1,\n kernel_size=(3,),\n strides=(3,),\n padding='CIRCULAR',\n kernel_init=lambda *_: kernel,\n bias_init=initializers.zeros,\n )\n y, initial_params = conv_module.init_with_output(rng, x)\n\n self.assertEqual(initial_params['params']['kernel'].shape, (2, 3, 1))\n # Compare with manually computed convolution\n correct_ans = np.array((-1 * 5 + 2 * 1 + 3 * 2, 4 * 3 + 5 * 4 + 6 * 5))\n correct_ans = np.expand_dims(correct_ans, (0, 2))\n np.testing.assert_allclose(y, correct_ans)\n\n def test_circular_conv_1d_dilation(self):\n \"\"\"Test 1d convolution with circular padding and kernel dilation.\"\"\"\n rng = dict(params=random.key(0))\n x = np.arange(1, 6)\n x = np.expand_dims(x, (0, 2))\n kernel = np.array((1, 2, 1))\n kernel = np.expand_dims(kernel, (1, 2))\n\n conv_module = nn.Conv(\n features=1,\n kernel_size=(3,),\n padding='CIRCULAR',\n kernel_init=lambda *_: kernel,\n bias_init=initializers.zeros,\n kernel_dilation=(3,),\n )\n y, initial_params = conv_module.init_with_output(rng, x)\n\n self.assertEqual(initial_params['params']['kernel'].shape, (3, 1, 1))\n # Compare with manually computed convolution\n correct_ans = np.array(\n (\n 3 + 2 * 1 + 4,\n 4 + 2 * 2 + 5,\n 5 + 2 * 3 + 1,\n 1 + 2 * 4 + 2,\n 2 + 2 * 5 + 3,\n )\n )\n correct_ans = np.expand_dims(correct_ans, (0, 2))\n np.testing.assert_allclose(y, correct_ans)\n\n def test_circular_conv_local_1d_dilation(self):\n \"\"\"\n Test 1d local convolution with circular padding and kernel dilation\n \"\"\"\n rng = dict(params=random.key(0))\n x = np.arange(1, 6)\n x = np.expand_dims(x, (0, 2))\n kernel = np.array(\n ((1, 2, 1), (3, 4, 5), (-1, 1, 2), (2, 3, 4), (-1, -2, -3))\n )\n kernel = np.expand_dims(kernel, (2,))\n\n conv_module = nn.ConvLocal(\n features=1,\n kernel_size=(3,),\n padding='CIRCULAR',\n kernel_init=lambda *_: kernel,\n bias_init=initializers.zeros,\n kernel_dilation=(3,),\n )\n y, initial_params = conv_module.init_with_output(rng, x)\n\n self.assertEqual(initial_params['params']['kernel'].shape, (5, 3, 1))\n # Compare with manually computed convolution\n correct_ans = np.array(\n (\n 1 * 3 + 2 * 1 + 1 * 4,\n 3 * 4 + 4 * 2 + 5 * 5,\n -1 * 5 + 1 * 3 + 2 * 1,\n 2 * 1 + 3 * 4 + 4 * 2,\n -1 * 2 + -2 * 5 + -3 * 3,\n )\n )\n correct_ans = np.expand_dims(correct_ans, (0, 2))\n np.testing.assert_allclose(y, correct_ans)\n\n def test_circular_conv_2d_custom(self):\n \"\"\"Test 2d convolution with circular padding on a 3x3 example.\"\"\"\n rng = dict(params=random.key(0))\n x = np.array(((1, 2, 3), (4, 5, 6), (7, 8, 9)))\n x = np.expand_dims(x, (0, 3))\n kernel = np.array(((0, 1, 0), (1, 2, 1), (0, 1, 0)))\n kernel = np.expand_dims(kernel, (2, 3))\n\n conv_module = nn.Conv(\n features=1,\n kernel_size=(3, 3),\n padding='CIRCULAR',\n kernel_init=lambda *_: kernel,\n bias_init=initializers.zeros,\n )\n y, initial_params = conv_module.init_with_output(rng, x)\n\n self.assertEqual(initial_params['params']['kernel'].shape, (3, 3, 1, 1))\n # Compare with manually computed convolution\n correct_ans = np.array(\n (\n (2 * 1 + 7 + 2 + 4 + 3, 2 * 2 + 8 + 3 + 5 + 1, 2 * 3 + 9 + 1 + 6 + 2),\n (2 * 4 + 1 + 5 + 7 + 6, 2 * 5 + 2 + 6 + 8 + 4, 2 * 6 + 3 + 4 + 9 + 5),\n (2 * 7 + 4 + 8 + 1 + 9, 2 * 8 + 5 + 9 + 2 + 7, 2 * 9 + 6 + 7 + 3 + 8),\n )\n )\n correct_ans = np.expand_dims(correct_ans, (0, 3))\n np.testing.assert_allclose(y, correct_ans)\n\n def test_circular_conv_local_2d_custom(self):\n \"\"\"\n Test 2d local convolution with circular padding on a 3x3 example\n \"\"\"\n rng = dict(params=random.key(0))\n x = np.array(((1, 2, 3), (4, 5, 6), (7, 8, 9)))\n x = np.expand_dims(x, (0, 3))\n kernel = np.array(\n (\n (\n ((0, 1, 0), (1, 2, 1), (0, 1, 0)),\n ((0, 1, 0), (1, 3, 1), (0, 1, 0)),\n ((0, 1, 0), (1, 4, 1), (0, 1, 0)),\n ),\n (\n ((0, 1, 0), (1, 5, 1), (0, 1, 0)),\n ((0, 1, 0), (1, 6, 1), (0, 1, 0)),\n ((0, 1, 0), (1, 7, 1), (0, 1, 0)),\n ),\n (\n ((0, 1, 0), (1, 8, 1), (0, 1, 0)),\n ((0, 1, 0), (1, 9, 1), (0, 1, 0)),\n ((0, 1, 0), (1, 10, 1), (0, 1, 0)),\n ),\n )\n )\n kernel = np.expand_dims(kernel, (3,))\n kernel = np.reshape(kernel, (3, 3, 9, 1))\n\n conv_module = nn.ConvLocal(\n features=1,\n kernel_size=(3, 3),\n padding='CIRCULAR',\n kernel_init=lambda *_: kernel,\n bias_init=initializers.zeros,\n )\n y, initial_params = conv_module.init_with_output(rng, x)\n\n self.assertEqual(initial_params['params']['kernel'].shape, (3, 3, 9, 1))\n # Compare with manually computed convolution\n correct_ans = np.array(\n (\n (2 * 1 + 7 + 2 + 4 + 3, 3 * 2 + 8 + 3 + 5 + 1, 4 * 3 + 9 + 1 + 6 + 2),\n (5 * 4 + 1 + 5 + 7 + 6, 6 * 5 + 2 + 6 + 8 + 4, 7 * 6 + 3 + 4 + 9 + 5),\n (8 * 7 + 4 + 8 + 1 + 9, 9 * 8 + 5 + 9 + 2 + 7, 10 * 9 + 6 + 7 + 3 + 8),\n )\n )\n correct_ans = np.expand_dims(correct_ans, (0, 3))\n np.testing.assert_allclose(y, correct_ans)\n\n def test_reflect_conv_1d_custom(self):\n \"\"\"Test 1d convolution with reflection padding and a stride.\"\"\"\n rng = dict(params=random.key(0))\n x = np.arange(1, 6)\n x = np.expand_dims(x, (0, 2))\n kernel = np.array((1, 2, 1))\n kernel = np.expand_dims(kernel, (1, 2))\n\n conv_module = nn.Conv(\n features=1,\n kernel_size=(3,),\n strides=(2,),\n padding='REFLECT',\n kernel_init=lambda *_: kernel,\n bias_init=initializers.zeros,\n )\n y, initial_params = conv_module.init_with_output(rng, x)\n\n self.assertEqual(initial_params['params']['kernel'].shape, (3, 1, 1))\n # Compare with manually computed convolution\n correct_ans = np.array((2 + 2 * 1 + 2, 2 + 2 * 3 + 4, 4 + 2 * 5 + 4))\n correct_ans = np.expand_dims(correct_ans, (0, 2))\n np.testing.assert_allclose(y, correct_ans)\n\n def test_reflect_conv_2d_custom(self):\n \"\"\"Test 2d convolution with reflect padding on a 3x3 example.\"\"\"\n rng = dict(params=random.key(0))\n x = np.array(((1, 2, 3), (4, 5, 6), (7, 8, 9)))\n x = np.expand_dims(x, (0, 3))\n kernel = np.array(((0, 1, 0), (1, 2, 1), (0, 1, 0)))\n kernel = np.expand_dims(kernel, (2, 3))\n\n conv_module = nn.Conv(\n features=1,\n kernel_size=(3, 3),\n padding='REFLECT',\n kernel_init=lambda *_: kernel,\n bias_init=initializers.zeros,\n )\n y, initial_params = conv_module.init_with_output(rng, x)\n\n self.assertEqual(initial_params['params']['kernel'].shape, (3, 3, 1, 1))\n # Compare with manually computed convolution\n correct_ans = np.array(\n (\n (2 * 1 + 4 + 2 + 4 + 2, 2 * 2 + 5 + 3 + 5 + 1, 2 * 3 + 6 + 2 + 6 + 2),\n (2 * 4 + 1 + 5 + 7 + 5, 2 * 5 + 2 + 6 + 8 + 4, 2 * 6 + 3 + 5 + 9 + 5),\n (2 * 7 + 4 + 8 + 8 + 4, 2 * 8 + 5 + 9 + 5 + 7, 2 * 9 + 6 + 8 + 6 + 8),\n )\n )\n correct_ans = np.expand_dims(correct_ans, (0, 3))\n np.testing.assert_allclose(y, correct_ans)\n\n def test_causal_conv1d(self):\n rng = dict(params=random.key(0))\n x = jnp.ones((1, 8, 4))\n conv_module = nn.Conv(\n features=4,\n kernel_size=(3,),\n padding='CAUSAL',\n kernel_init=initializers.ones,\n bias_init=initializers.ones,\n )\n y, _ = conv_module.init_with_output(rng, x)\n correct_ans = np.array(\n [\n [\n [5.0, 5.0, 5.0, 5.0],\n [9.0, 9.0, 9.0, 9.0],\n [13.0, 13.0, 13.0, 13.0],\n [13.0, 13.0, 13.0, 13.0],\n [13.0, 13.0, 13.0, 13.0],\n [13.0, 13.0, 13.0, 13.0],\n [13.0, 13.0, 13.0, 13.0],\n [13.0, 13.0, 13.0, 13.0],\n ]\n ]\n )\n np.testing.assert_allclose(y, correct_ans)\n np.testing.assert_array_equal(correct_ans.shape, y.shape)\n\n @parameterized.product(\n use_bias=(True, False),\n )\n def test_conv_transpose(self, use_bias):\n rng = dict(params=random.key(0))\n x = jnp.ones((1, 8, 3))\n conv_transpose_module = nn.ConvTranspose(\n features=4,\n use_bias=use_bias,\n kernel_size=(3,),\n padding='VALID',\n kernel_init=initializers.ones,\n bias_init=initializers.ones,\n )\n y, initial_params = conv_transpose_module.init_with_output(rng, x)\n self.assertEqual(initial_params['params']['kernel'].shape, (3, 3, 4))\n correct_ans = np.array(\n [\n [\n [4.0, 4.0, 4.0, 4.0],\n [7.0, 7.0, 7.0, 7.0],\n [10.0, 10.0, 10.0, 10.0],\n [10.0, 10.0, 10.0, 10.0],\n [10.0, 10.0, 10.0, 10.0],\n [10.0, 10.0, 10.0, 10.0],\n [10.0, 10.0, 10.0, 10.0],\n [10.0, 10.0, 10.0, 10.0],\n [7.0, 7.0, 7.0, 7.0],\n [4.0, 4.0, 4.0, 4.0],\n ]\n ]\n )\n if not use_bias:\n correct_ans -= 1.0\n np.testing.assert_allclose(y, correct_ans)\n\n @parameterized.product(\n use_bias=(True, False),\n )\n def test_multibatch_input_conv_transpose(self, use_bias):\n rng = dict(params=random.key(0))\n x = jnp.ones((2, 5, 8, 3))\n conv_transpose_module = nn.ConvTranspose(\n features=4,\n use_bias=use_bias,\n kernel_size=(3,),\n padding='VALID',\n kernel_init=initializers.ones,\n bias_init=initializers.ones,\n )\n y, initial_params = conv_transpose_module.init_with_output(rng, x)\n self.assertEqual(initial_params['params']['kernel'].shape, (3, 3, 4))\n correct_ans = np.array(\n [\n [\n [4.0, 4.0, 4.0, 4.0],\n [7.0, 7.0, 7.0, 7.0],\n [10.0, 10.0, 10.0, 10.0],\n [10.0, 10.0, 10.0, 10.0],\n [10.0, 10.0, 10.0, 10.0],\n [10.0, 10.0, 10.0, 10.0],\n [10.0, 10.0, 10.0, 10.0],\n [10.0, 10.0, 10.0, 10.0],\n [7.0, 7.0, 7.0, 7.0],\n [4.0, 4.0, 4.0, 4.0],\n ]\n ]\n )\n correct_ans = np.repeat(correct_ans[None], repeats=2, axis=0)\n correct_ans = np.repeat(correct_ans, repeats=5, axis=1)\n if not use_bias:\n correct_ans -= 1.0\n np.testing.assert_allclose(y, correct_ans)\n\n def test_single_input_conv_transpose(self):\n rng = dict(params=random.key(0))\n x = jnp.ones((8, 3))\n conv_transpose_module = nn.ConvTranspose(\n features=4,\n kernel_size=(3,),\n padding='VALID',\n kernel_init=initializers.ones,\n bias_init=initializers.ones,\n )\n y, initial_params = conv_transpose_module.init_with_output(rng, x)\n self.assertEqual(initial_params['params']['kernel'].shape, (3, 3, 4))\n correct_ans = np.array(\n [\n [4.0, 4.0, 4.0, 4.0],\n [7.0, 7.0, 7.0, 7.0],\n [10.0, 10.0, 10.0, 10.0],\n [10.0, 10.0, 10.0, 10.0],\n [10.0, 10.0, 10.0, 10.0],\n [10.0, 10.0, 10.0, 10.0],\n [10.0, 10.0, 10.0, 10.0],\n [10.0, 10.0, 10.0, 10.0],\n [7.0, 7.0, 7.0, 7.0],\n [4.0, 4.0, 4.0, 4.0],\n ]\n )\n np.testing.assert_allclose(y, correct_ans)\n\n def test_single_input_masked_conv_transpose(self):\n rng = dict(params=random.key(0))\n x = jnp.ones((8, 3))\n m = jnp.tril(jnp.ones((3, 3, 4)))\n conv_transpose_module = nn.ConvTranspose(\n features=4,\n kernel_size=(3,),\n padding='VALID',\n mask=m,\n kernel_init=initializers.ones,\n bias_init=initializers.ones,\n )\n y, initial_params = conv_transpose_module.init_with_output(rng, x)\n self.assertEqual(initial_params['params']['kernel'].shape, (3, 3, 4))\n correct_ans = np.array(\n [\n [4.0, 3.0, 2.0, 1.0],\n [7.0, 5.0, 3.0, 1.0],\n [10.0, 7.0, 4.0, 1.0],\n [10.0, 7.0, 4.0, 1.0],\n [10.0, 7.0, 4.0, 1.0],\n [10.0, 7.0, 4.0, 1.0],\n [10.0, 7.0, 4.0, 1.0],\n [10.0, 7.0, 4.0, 1.0],\n [7.0, 5.0, 3.0, 1.0],\n [4.0, 3.0, 2.0, 1.0],\n ]\n )\n np.testing.assert_allclose(y, correct_ans)\n\n @parameterized.product(\n n_batch=(1, 3),\n n_features=(1, 2),\n kernel_size=(1, 2, 3, 9),\n n_input_features=(1, 3),\n input_size=(1, 8, 16),\n )\n def test_circular_conv_transpose_1d_constant(\n self, n_batch, n_features, kernel_size, n_input_features, input_size\n ):\n \"\"\"\n Test 1D transposed convolution with circular padding: filter with all\n elements equal to 1 applied on an input with all elements equal to 1.\n Result should have the same shape as input (except for the feature\n dimension) and have all elements equal to\n `n_input_features * kernel_lin_size`.\n \"\"\"\n rng = dict(params=random.key(0))\n x = jnp.ones((n_batch, input_size, n_input_features))\n conv_module = nn.ConvTranspose(\n features=n_features,\n kernel_size=(kernel_size,),\n padding='CIRCULAR',\n kernel_init=initializers.ones,\n bias_init=initializers.zeros,\n )\n y, initial_params = conv_module.init_with_output(rng, x)\n\n self.assertEqual(\n initial_params['params']['kernel'].shape,\n (kernel_size, n_input_features, n_features),\n )\n correct_ans = np.full(\n (n_batch, input_size, n_features), kernel_size * n_input_features\n )\n np.testing.assert_allclose(y, correct_ans)\n\n @parameterized.product(\n n_batch=(1, 3),\n n_features=(1, 2, 10),\n kernel_lin_size=(1, 2, 3, 9),\n n_input_features=(1, 5),\n input_x_size=(14,),\n input_y_size=(5, 10),\n )\n def test_circular_conv_transpose_2d_constant(\n self,\n n_batch,\n n_features,\n kernel_lin_size,\n n_input_features,\n input_x_size,\n input_y_size,\n ):\n \"\"\"\n Test 2D transposed convolution with circular padding: square filter with all\n elements equal to 1 applied on an input with all elements equal to 1.\n Result should have the same shape as input (except for the feature\n dimension) and have all elements equal to\n `n_input_features * kernel_lin_size^2`.\n \"\"\"\n rng = dict(params=random.key(0))\n x = jnp.ones((n_batch, input_x_size, input_y_size, n_input_features))\n conv_module = nn.ConvTranspose(\n features=n_features,\n kernel_size=(kernel_lin_size, kernel_lin_size),\n padding='CIRCULAR',\n kernel_init=initializers.ones,\n bias_init=initializers.zeros,\n )\n y, initial_params = conv_module.init_with_output(rng, x)\n\n self.assertEqual(\n initial_params['params']['kernel'].shape,\n (kernel_lin_size, kernel_lin_size, n_input_features, n_features),\n )\n correct_ans = np.full(\n (n_batch, input_x_size, input_y_size, n_features),\n kernel_lin_size * kernel_lin_size * n_input_features,\n )\n np.testing.assert_allclose(y, correct_ans)\n\n def test_circular_conv_transpose_2d_with_vmap(self):\n layer = nn.ConvTranspose(features=5, kernel_size=(3,), padding='CIRCULAR')\n\n # this is ok\n sample_input = jnp.ones((1, 32, 2))\n out, vars = layer.init_with_output(jax.random.key(0), sample_input)\n self.assertEqual(out.shape, (1, 32, 5))\n\n batch_input = jnp.ones((8, 32, 2))\n batch_apply = jax.vmap(layer.apply, in_axes=(None, 0))\n\n # this breaks with the error provided\n batch_out = batch_apply(vars, batch_input)\n self.assertEqual(batch_out.shape, (8, 32, 5))\n\n def test_circular_conv_transpose_1d_custom(self):\n \"\"\"Test 1d transposed convolution with circular padding and a stride.\"\"\"\n rng = dict(params=random.key(0))\n x = np.arange(1, 6)\n x = np.expand_dims(x, (0, 2))\n kernel = np.array((1, 2, 1))\n kernel = np.expand_dims(kernel, (1, 2))\n\n conv_module = nn.ConvTranspose(\n features=1,\n kernel_size=(3,),\n strides=(3,),\n padding='CIRCULAR',\n kernel_init=lambda *_: kernel,\n bias_init=initializers.zeros,\n )\n y, initial_params = conv_module.init_with_output(rng, x)\n\n self.assertEqual(initial_params['params']['kernel'].shape, (3, 1, 1))\n # Compare with manually computed convolution\n correct_ans = np.array(\n ( # pyformat: disable\n 1 * 1,\n 1 * 2,\n 1 * 1,\n 2 * 1,\n 2 * 2,\n 2 * 1,\n 3 * 1,\n 3 * 2,\n 3 * 1,\n 4 * 1,\n 4 * 2,\n 4 * 1,\n 5 * 1,\n 5 * 2,\n 5 * 1,\n )\n )\n correct_ans = np.expand_dims(correct_ans, (0, 2))\n np.testing.assert_allclose(y, correct_ans)\n\n def test_circular_conv_transpose_2d_custom(self):\n \"\"\"Test 2d transposed convolution with circular padding on a 3x3 example.\"\"\"\n rng = dict(params=random.key(0))\n x = np.array(\n (\n (1, 2, 3),\n (4, 5, 6),\n (7, 8, 9),\n )\n )\n x = np.expand_dims(x, (0, 3))\n kernel = np.array(((0, 1, 0), (1, 2, 1), (0, 1, 0)))\n kernel = np.expand_dims(kernel, (2, 3))\n\n conv_module = nn.ConvTranspose(\n features=1,\n kernel_size=(3, 3),\n padding='CIRCULAR',\n kernel_init=lambda *_: kernel,\n bias_init=initializers.zeros,\n )\n y, initial_params = conv_module.init_with_output(rng, x)\n\n self.assertEqual(initial_params['params']['kernel'].shape, (3, 3, 1, 1))\n # Compare with manually computed convolution\n correct_ans = np.array(\n (\n (18, 21, 24),\n (27, 30, 33),\n (36, 39, 42),\n )\n )\n correct_ans = np.expand_dims(correct_ans, (0, 3))\n np.testing.assert_allclose(y, correct_ans)\n\n def test_circular_conv_transpose_2d_custom_bias(self):\n \"\"\"Test 2d transposed convolution with circular padding on a 2x2 example with bias.\"\"\"\n rng = dict(params=random.key(0))\n x = np.array(((1, 2), (3, 4)))\n x = np.expand_dims(x, (0, 3))\n kernel = np.array(\n (\n (1, 2),\n (3, 4),\n )\n )\n kernel = np.expand_dims(kernel, (2, 3))\n\n conv_module = nn.ConvTranspose(\n features=1,\n kernel_size=(2, 2),\n padding='CIRCULAR',\n kernel_init=lambda *_: kernel,\n bias_init=initializers.ones,\n )\n y, initial_params = conv_module.init_with_output(rng, x)\n\n self.assertEqual(initial_params['params']['kernel'].shape, (2, 2, 1, 1))\n # Compare with manually computed convolution\n correct_ans = np.array(\n (\n (21, 23),\n (29, 31),\n )\n )\n correct_ans = np.expand_dims(correct_ans, (0, 3))\n np.testing.assert_allclose(y, correct_ans)\n\n @parameterized.product(use_bias=(True, False))\n def test_transpose_kernel_conv_transpose(self, use_bias):\n rng = dict(params=random.key(0))\n x = jnp.ones((1, 15, 15, 3))\n conv_module = nn.ConvTranspose(\n features=4,\n use_bias=use_bias,\n strides=2,\n kernel_size=(6, 6),\n padding='CIRCULAR',\n transpose_kernel=True,\n )\n y, initial_params = conv_module.init_with_output(rng, x)\n self.assertEqual(initial_params['params']['kernel'].shape, (6, 6, 4, 3))\n self.assertEqual(y.shape, (1, 30, 30, 4))\n\n @parameterized.product(module=(nn.Conv, nn.ConvLocal, nn.ConvTranspose))\n def test_int_kernel_equality(self, module):\n conv_int = module(features=4, kernel_size=3)\n conv_seq = module(features=4, kernel_size=(3,))\n x = jnp.ones((8, 3))\n self.assertTrue(\n jax.tree_util.tree_all(\n jax.tree_util.tree_map(\n lambda x, y: (x == y).all(),\n conv_int.init(random.key(0), x),\n conv_seq.init(random.key(0), x),\n )\n )\n )\n\n def test_embed(self):\n rng = dict(params=random.key(0))\n x = jnp.arange(4)[None]\n dummy_embedding = jnp.broadcast_to(jnp.arange(4)[..., None], (4, 3)).astype(\n jnp.float32\n )\n embed_module = nn.Embed(\n num_embeddings=4,\n features=3,\n embedding_init=lambda rng, shape, dtype: dummy_embedding,\n )\n y, initial_params = embed_module.init_with_output(rng, x)\n np.testing.assert_allclose(y, dummy_embedding[None])\n z = embed_module.apply(\n initial_params, jnp.ones((3,)), method=embed_module.attend\n )\n np.testing.assert_allclose(z, 3.0 * jnp.arange(4))\n\n def test_embed_numpy(self):\n rng = dict(params=random.key(0))\n x = jnp.arange(4)[None]\n dummy_embedding = np.broadcast_to(np.arange(4)[..., None], (4, 3)).astype(\n np.float32\n )\n embed_module = nn.Embed(\n num_embeddings=4,\n features=3,\n embedding_init=lambda rng, shape, dtype: dummy_embedding,\n )\n y, initial_params = embed_module.init_with_output(rng, x)\n np.testing.assert_allclose(y, dummy_embedding[None])\n z = embed_module.apply(\n initial_params, jnp.ones((3,)), method=embed_module.attend\n )\n np.testing.assert_allclose(z, 3.0 * jnp.arange(4))\n\n def test_embed_hash(self):\n self.assertEqual(hash(nn.Embed(2, 3)), hash(nn.Embed(2, 3)))\n self.assertNotEqual(hash(nn.Embed(3, 4)), hash(nn.Embed(2, 3)))\n\n def test_non_final_axis(self):\n class Foo(nn.Module):\n @nn.compact\n def __call__(self, x):\n return nn.DenseGeneral(features=6, axis=1, name='dense')(x)\n\n x = jnp.ones((2, 4, 8))\n y, variables = Foo().init_with_output(random.key(0), x)\n self.assertEqual(\n jax.tree_util.tree_map(jnp.shape, variables['params']),\n {'dense': {'kernel': (4, 6), 'bias': (6,)}},\n )\n self.assertEqual(y.shape, (2, 8, 6))\n\n def test_non_final_axes(self):\n class Foo(nn.Module):\n @nn.compact\n def __call__(self, x):\n return nn.DenseGeneral(features=6, axis=(0, 1), name='dense')(x)\n\n x = jnp.ones((2, 4, 8))\n y, variables = Foo().init_with_output(random.key(0), x)\n self.assertEqual(\n jax.tree_util.tree_map(jnp.shape, variables['params']),\n {'dense': {'kernel': (2, 4, 6), 'bias': (6,)}},\n )\n self.assertEqual(y.shape, (8, 6))\n\n def test_canonicalize_padding(self):\n def test_pad(pad, rank, expected=None):\n if expected is None:\n with self.assertRaises(ValueError):\n nn.linear.canonicalize_padding(pad, rank)\n else:\n self.assertEqual(nn.linear.canonicalize_padding(pad, rank), expected)\n\n test_pad('SAME', 2, 'SAME')\n test_pad(2, 3, [(2, 2), (2, 2), (2, 2)])\n test_pad((2, 2), 3)\n test_pad((2, 2), 1)\n test_pad([1, (2, 3)], 2, [(1, 1), (2, 3)])\n test_pad([None, (1, 2)], 2)\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/linen/linen_meta_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for linen_meta.\"\"\"\n\nimport jax\nfrom absl.testing import absltest\nfrom jax import numpy as jnp\nfrom jax import random\nfrom jax.experimental import mesh_utils\nfrom jax.sharding import Mesh, PartitionSpec\n\nfrom flax import linen as nn\n\n\nclass LinenMetaTest(absltest.TestCase):\n def test_boxed_param(self):\n class Bar(nn.Module):\n @nn.compact\n def __call__(mdl_self, x): # pylint: disable=no-self-argument\n kernel_init = nn.with_partitioning(\n nn.initializers.ones_init(), ('in', 'out')\n )\n kernel = mdl_self.param('kernel', kernel_init, (x.shape[-1], 2))\n kernel_box = mdl_self.get_variable('params', 'kernel')\n self.assertIsInstance(kernel_box, nn.Partitioned)\n self.assertEqual(kernel_box.names, ('in', 'out'))\n return x @ kernel\n\n class Foo(nn.Module):\n @nn.compact\n def __call__(self, xs):\n return nn.vmap(\n Bar,\n in_axes=0,\n variable_axes={'params': 0},\n split_rngs={'params': True},\n metadata_params={nn.PARTITION_NAME: 'batch'},\n )(name='bar')(xs)\n\n m = Foo()\n variables = m.init(random.key(0), jnp.zeros((8, 3)))\n self.assertEqual(\n variables['params']['bar']['kernel'].names, ('batch', 'in', 'out')\n )\n\n def test_boxed_variable(self):\n class Bar(nn.Module):\n @nn.compact\n def __call__(mdl_self, x): # pylint: disable=no-self-argument\n kernel_init = nn.with_partitioning(\n nn.initializers.ones_init(), ('in', 'out')\n )\n kernel = mdl_self.variable(\n 'params',\n 'kernel',\n kernel_init,\n mdl_self.make_rng('params'),\n (x.shape[-1], 2),\n )\n kernel.value += 1.0\n self.assertEqual(kernel.value.sum(), kernel.value.size * 2)\n kernel_box = mdl_self.get_variable('params', 'kernel')\n self.assertIsInstance(kernel_box, nn.Partitioned)\n self.assertEqual(kernel_box.names, ('in', 'out'))\n return x @ kernel.value\n\n class Foo(nn.Module):\n @nn.compact\n def __call__(self, xs):\n return nn.vmap(\n Bar,\n in_axes=0,\n variable_axes={'params': 0},\n split_rngs={'params': True},\n metadata_params={nn.PARTITION_NAME: 'batch'},\n )(name='bar')(xs)\n\n m = Foo()\n variables = m.init(random.key(0), jnp.zeros((8, 3)))\n self.assertEqual(\n variables['params']['bar']['kernel'].names, ('batch', 'in', 'out')\n )\n\n # def test_boxed_variable(self):\n # def f(scope, xs):\n # def g(scope, x):\n # kernel_init = nn.with_partitioning(nn.initializers.ones_init(),\n # ('in', 'out'))\n # kernel = scope.variable('params', 'kernel', kernel_init,\n # scope.make_rng('params'), (x.shape[-1], 2))\n # kernel.value += 1.\n # self.assertEqual(kernel.value.sum(), kernel.value.size * 2)\n # kernel_box = scope.get_variable('params', 'kernel')\n # self.assertIsInstance(kernel_box, nn.Partitioned)\n # self.assertEqual(kernel_box.names, ('in', 'out'))\n # return x @ kernel.value\n\n # nn.vmap(\n # g, in_axes=0,\n # variable_axes={'params': 0}, split_rngs={'params': True},\n # metadata_params={nn.PARTITION_NAME: 'batch'})(scope, xs)\n\n # _, variables = init(f)(random.key(0), jnp.zeros((8, 3)))\n # self.assertEqual(variables['params']['kernel'].names,\n # ('batch', 'in', 'out'))\n\n def test_pjit_scan_over_layers(self):\n class MLP(nn.Module):\n hidden_size: int\n\n @nn.compact\n def __call__(self, x):\n ki = nn.linear.default_kernel_init\n h = nn.Dense(\n self.hidden_size,\n kernel_init=nn.with_partitioning(ki, ('data', 'model')),\n )(x)\n h = nn.relu(h)\n return nn.Dense(\n x.shape[-1], kernel_init=nn.with_partitioning(ki, ('model', 'data'))\n )(h)\n\n class Model(nn.Module):\n @nn.compact\n def __call__(self, x):\n def body(_, c):\n c = MLP(512)(c)\n return c, ()\n\n c, _ = nn.scan(\n body,\n variable_axes={'params': 0},\n split_rngs={'params': 0},\n length=8,\n metadata_params={nn.PARTITION_NAME: None},\n )(self, x)\n return c\n\n devs = mesh_utils.create_device_mesh((jax.device_count(), 1))\n mesh = Mesh(devs, ['data', 'model'])\n model = Model()\n x = jnp.ones((8, 128))\n spec = nn.get_partition_spec(jax.eval_shape(model.init, random.key(0), x))\n self.assertEqual(\n spec,\n {\n 'params': {\n 'MLP_0': {\n 'Dense_0': {\n 'bias': PartitionSpec(),\n 'kernel': PartitionSpec(None, 'data', 'model'),\n },\n 'Dense_1': {\n 'bias': PartitionSpec(),\n 'kernel': PartitionSpec(None, 'model', 'data'),\n },\n },\n },\n },\n )\n x_spec = PartitionSpec('data', 'model')\n f = lambda x: jax.sharding.NamedSharding(mesh, x)\n key_spec = PartitionSpec()\n init_fn = jax.jit(\n model.init,\n in_shardings=jax.tree_util.tree_map(f, (key_spec, x_spec)),\n out_shardings=jax.tree_util.tree_map(f, spec),\n )\n variables = init_fn(random.key(0), x)\n apply_fn = jax.jit(\n model.apply,\n in_shardings=jax.tree_util.tree_map(f, (spec, x_spec)),\n out_shardings=jax.tree_util.tree_map(f, x_spec),\n )\n y = apply_fn(variables, x)\n self.assertEqual(y.shape, (8, 128))\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/linen/linen_module_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for flax.linen.\"\"\"\n\nimport contextlib\nimport copy\nimport dataclasses\nimport enum\nimport functools\nimport gc\nimport inspect\nimport operator\nimport sys\nfrom tempfile import TemporaryDirectory\nfrom typing import (\n Any,\n Generic,\n NamedTuple,\n TypeVar,\n get_type_hints,\n)\nfrom collections.abc import Callable, Mapping, Sequence\nfrom unittest.mock import patch\n\nimport jax\nimport jax.numpy as jnp\nimport numpy as np\nfrom absl.testing import absltest\nfrom jax import random\nfrom jax.nn import initializers\n\nfrom flax import config, errors, struct\nfrom flax import linen as nn\nfrom flax.core import FrozenDict, Scope, freeze\nfrom flax.linen import compact\n\n# Parse absl flags test_srcdir and test_tmpdir.\njax.config.parse_flags_with_absl()\n\n\ndef tree_equals(x, y):\n return jax.tree_util.tree_all(jax.tree_util.tree_map(operator.eq, x, y))\n\n\n@contextlib.contextmanager\ndef set_config(option: str, value: bool):\n old_value = getattr(config, option)\n try:\n config.update(option, value)\n yield None\n finally:\n config.update(option, old_value)\n\n\nclass DummyModule(nn.Module):\n @compact\n def __call__(self, x):\n bias = self.param('bias', initializers.ones, x.shape)\n return x + bias\n\n\nclass Dense(nn.Module):\n features: int\n\n @compact\n def __call__(self, x):\n kernel = self.param(\n 'kernel', initializers.lecun_normal(), (x.shape[-1], self.features)\n )\n y = jnp.dot(x, kernel)\n return y\n\n\nclass IdentityModule(nn.Module):\n def __call__(self, x):\n return x\n\n\nclass RaisesModule(nn.Module):\n def __call__(self):\n assert False\n\n\nclass ModuleTest(absltest.TestCase):\n def test_init_module(self):\n rngkey = jax.random.key(0)\n x = jnp.array([1.0])\n scope = Scope({}, {'params': rngkey}, mutable=['params'])\n y = DummyModule(parent=scope)(x)\n params = scope.variables()['params']\n y2 = DummyModule(parent=scope.rewound())(x)\n np.testing.assert_allclose(y, y2)\n np.testing.assert_allclose(y, jnp.array([2.0]))\n self.assertEqual(params, {'bias': jnp.array([1.0])})\n\n def test_lazy_init(self):\n class Foo(nn.Module):\n @compact\n def __call__(self, x):\n k = self.param(\n 'kernel', nn.initializers.lecun_normal(), (x.shape[-1], x.shape[-1])\n )\n return x @ k\n\n # provide a massive input message which would OOM if any compute ops were actually executed\n variables = Foo().lazy_init(\n random.key(0),\n jax.ShapeDtypeStruct((1024 * 1024 * 1024, 128), jnp.float32),\n )\n self.assertEqual(variables['params']['kernel'].shape, (128, 128))\n\n def test_lazy_init_fails_on_data_dependence(self):\n class Foo(nn.Module):\n @compact\n def __call__(self, x):\n k = self.param('kernel', lambda _: x)\n return x * k\n\n with self.assertRaises(errors.LazyInitError):\n Foo().lazy_init(random.key(0), jax.ShapeDtypeStruct((8, 4), jnp.float32))\n\n def test_arg_module(self):\n rngkey = jax.random.key(0)\n x = jnp.ones((10,))\n scope = Scope({}, {'params': rngkey}, mutable=['params'])\n y = Dense(3, parent=scope)(x)\n params = scope.variables()['params']\n y2 = Dense(3, parent=scope.rewound())(x)\n np.testing.assert_allclose(y, y2)\n self.assertEqual(params['kernel'].shape, (10, 3))\n\n def test_util_fun(self):\n rngkey = jax.random.key(0)\n\n class MLP(nn.Module):\n @compact\n def __call__(self, x):\n x = self._mydense(x)\n x = self._mydense(x)\n return x\n\n def _mydense(self, x):\n return Dense(3)(x)\n\n x = jnp.ones((10,))\n scope = Scope({}, {'params': rngkey}, mutable=['params'])\n y = MLP(parent=scope)(x)\n params = scope.variables()['params']\n y2 = MLP(parent=scope.rewound())(x)\n np.testing.assert_allclose(y, y2)\n param_shape = jax.tree_util.tree_map(jnp.shape, params)\n self.assertEqual(\n param_shape,\n {'Dense_0': {'kernel': (10, 3)}, 'Dense_1': {'kernel': (3, 3)}},\n )\n\n def test_nested_module_reuse(self):\n rngkey = jax.random.key(0)\n\n class MLP(nn.Module):\n @compact\n def __call__(self, x):\n x = self._mydense(x)\n x = self._mydense(x)\n return x\n\n def _mydense(self, x):\n return Dense(3)(x)\n\n class Top(nn.Module):\n @compact\n def __call__(self, x):\n mlp = MLP()\n y = mlp(x)\n z = mlp(x)\n return y + z\n\n x = jnp.ones((10,))\n scope = Scope({}, {'params': rngkey}, mutable=['params'])\n y = Top(parent=scope)(x)\n params = scope.variables()['params']\n y2 = Top(parent=scope.rewound())(x)\n np.testing.assert_allclose(y, y2)\n param_shape = jax.tree_util.tree_map(jnp.shape, params)\n self.assertEqual(\n param_shape,\n {\n 'MLP_0': {\n 'Dense_0': {'kernel': (10, 3)},\n 'Dense_1': {'kernel': (3, 3)},\n }\n },\n )\n\n def test_setup_dict_assignment(self):\n rngkey = jax.random.key(0)\n\n class MLP(nn.Module):\n def setup(self):\n self.lyrs1 = {\n 'a': Dense(3),\n 'b': Dense(3),\n }\n self.lyrs2 = [Dense(3), Dense(3)]\n\n def __call__(self, x):\n y = self.lyrs1['a'](x)\n z = self.lyrs1['b'](y)\n return z\n\n x = jnp.ones((10,))\n scope = Scope({}, {'params': rngkey}, mutable=['params'])\n y = MLP(parent=scope)(x)\n params = scope.variables()['params']\n y2 = MLP(parent=scope.rewound())(x)\n np.testing.assert_allclose(y, y2)\n param_shape = jax.tree_util.tree_map(jnp.shape, params)\n self.assertEqual(\n param_shape,\n {'lyrs1_a': {'kernel': (10, 3)}, 'lyrs1_b': {'kernel': (3, 3)}},\n )\n\n def test_setup_dict_nonstring_keys(self):\n class Foo(nn.Module):\n def setup(self):\n self.a = {(1, 2): nn.Dense(2)} # Tuple as key.\n\n @nn.compact\n def __call__(self, x):\n return self.a[(1, 2)](x)\n\n foo = Foo()\n x = jnp.ones(shape=(1, 3))\n params = foo.init(random.key(0), x)['params']\n param_shape = jax.tree_util.tree_map(jnp.shape, params)\n self.assertEqual(\n param_shape, {'a_(1, 2)': {'kernel': (3, 2), 'bias': (2,)}}\n )\n\n def test_setup_frozen_dict_nonstring_keys(self):\n a = {1: 2}\n\n class Foo(nn.Module):\n def setup(self):\n self.a = FrozenDict(a) # int as key.\n\n @nn.compact\n def __call__(self, x):\n return self.a[x]\n\n foo = Foo()\n x = 1\n params = foo.init(random.key(0), x)\n assert foo.apply(params, x) == a[x]\n\n def test_setup_dict_nonstring_keys_in_state(self):\n class Foo(nn.Module):\n a: dict[int, int] # int as key.\n\n @nn.compact\n def __call__(self, x):\n return self.a[x]\n\n a = {1: 2}\n foo = Foo(a)\n x = 1\n params = foo.init(random.key(0), x)\n assert foo.apply(params, x) == a[x]\n\n def test_setup_cloning(self):\n class MLP(nn.Module):\n def setup(self):\n self.dense = Dense(3)\n\n scope = Scope({})\n unused_clone = MLP(parent=scope).clone()\n\n def test_submodule_attr(self):\n rngkey = jax.random.key(0)\n\n class Inner(nn.Module):\n @compact\n def __call__(self):\n self.param('x', lambda rng: 40)\n\n class Outer(nn.Module):\n inner: nn.Module\n\n def __call__(self):\n return self.inner()\n\n class Wrapper(nn.Module):\n def setup(self):\n self.inner = Inner()\n self.outer = Outer(self.inner)\n\n def __call__(self):\n return self.outer()\n\n scope = Scope({'params': {}}, rngs={'params': rngkey}, mutable=['params'])\n # Make sure this doesn't raise \"Can't attach to remote parent\"\n wrapper = Wrapper(parent=scope)\n wrapper()\n\n # Make sure that variables are registered at the level of the\n # Wrapper submodule, not the Outer submodule.\n self.assertEqual(40, scope.variables()['params']['inner']['x'])\n\n def test_param_in_setup(self):\n rngkey = jax.random.key(0)\n\n class DummyModuleWithoutCompact(nn.Module):\n xshape: tuple[int, ...]\n\n def setup(self):\n self.bias = self.param('bias', initializers.ones, self.xshape)\n\n def __call__(self, x):\n return x + self.bias\n\n x = jnp.array([1.0])\n scope = Scope({}, {'params': rngkey}, mutable=['params'])\n y = DummyModuleWithoutCompact(x.shape, parent=scope)(x)\n params = scope.variables()['params']\n y2 = DummyModuleWithoutCompact(x.shape, parent=scope.rewound())(x)\n np.testing.assert_allclose(y, y2)\n np.testing.assert_allclose(y, jnp.array([2.0]))\n self.assertEqual(params, {'bias': jnp.array([1.0])})\n\n def test_init_outside_setup_without_compact(self):\n rngkey = jax.random.key(0)\n\n class DummyModuleWithoutCompact(nn.Module):\n def __call__(self, x):\n bias = self.param('bias', initializers.ones, x.shape)\n return x + bias\n\n x = jnp.array([1.0])\n scope = Scope({}, {'params': rngkey}, mutable=['params'])\n with self.assertRaisesRegex(ValueError, 'must be initialized.*setup'):\n unused_y = DummyModuleWithoutCompact(parent=scope)(x)\n\n def test_init_outside_call(self):\n rngkey = jax.random.key(0)\n\n class Dummy(nn.Module):\n @compact\n def __call__(self, x):\n bias = self.param('bias', initializers.ones, x.shape)\n return x + bias\n\n def foo(self, x):\n bias = self.param('bias', initializers.ones, x.shape)\n return x + bias\n\n x = jnp.array([1.0])\n scope = Scope({}, {'params': rngkey}, mutable=['params'])\n with self.assertRaisesRegex(ValueError, 'must be initialized.*setup'):\n unused_y = Dummy(parent=scope).foo(x)\n\n def test_setup_call_var_collision(self):\n rngkey = jax.random.key(0)\n\n class Dummy(nn.Module):\n xshape: tuple[int, ...]\n\n def setup(self):\n self.bias = self.param('bias', initializers.ones, self.xshape)\n\n @compact\n def __call__(self, x):\n unused_bias = self.param('bias', initializers.ones, x.shape)\n return x + self.bias\n\n x = jnp.array([1.0])\n scope = Scope({}, {'params': rngkey}, mutable=['params'])\n msg = 'Could not create param \"bias\" in Module Dummy: Name in use'\n with self.assertRaisesRegex(errors.NameInUseError, msg):\n unused_y = Dummy(x.shape, parent=scope)(x)\n\n def test_call_var_collision(self):\n rngkey = jax.random.key(0)\n\n class Dummy(nn.Module):\n xshape: tuple[int, ...]\n\n @compact\n def __call__(self, x):\n bias = self.param('bias', initializers.ones, self.xshape)\n bias = self.param('bias', initializers.ones, self.xshape)\n return x + bias\n\n x = jnp.array([1.0])\n scope = Scope({}, {'params': rngkey}, mutable=['params'])\n msg = 'Could not create param \"bias\" in Module Dummy: Name in use'\n with self.assertRaisesRegex(errors.NameInUseError, msg):\n unused_y = Dummy(x.shape, parent=scope)(x)\n\n def test_setup_var_collision(self):\n rngkey = jax.random.key(0)\n\n class Dummy(nn.Module):\n xshape: tuple[int, ...]\n\n def setup(self):\n self.bias = self.param('bias', initializers.ones, self.xshape)\n self.bias = self.param('bias', initializers.ones, self.xshape)\n\n def __call__(self, x):\n return x + self.bias\n\n x = jnp.array([1.0])\n scope = Scope({}, {'params': rngkey}, mutable=['params'])\n msg = 'Could not create param \"bias\" in Module Dummy: Name in use'\n with self.assertRaisesRegex(errors.NameInUseError, msg):\n unused_y = Dummy(x.shape, parent=scope)(x)\n\n def test_setattr_name_var_disagreement_allowed_in_lists(self):\n rngkey = jax.random.key(0)\n\n class Dummy(nn.Module):\n xshape: tuple[int, ...]\n\n def setup(self):\n self.biases = [\n self.param(f'bias_{i}', initializers.ones, self.xshape)\n for i in range(4)\n ]\n\n def __call__(self, x):\n return x + self.biases[0]\n\n x = jnp.array([1.0])\n scope = Scope({}, {'params': rngkey}, mutable=['params'])\n y = Dummy(x.shape, parent=scope)(x)\n self.assertEqual(y, jnp.array([2.0]))\n\n def test_setattr_name_var_disagreement_allowed_in_dicts(self):\n rngkey = jax.random.key(0)\n\n class Dummy(nn.Module):\n xshape: tuple[int, ...]\n\n def setup(self):\n self.biases = {\n # NOTE: Keys still must be strings. This is to make a possible\n # future transition to automatically derived parameter names when\n # assigned as a dict easier (like we currently have with\n # submodules). See a bit of discussion here:\n # https://github.com/google/flax/issues/705#issuecomment-738761853\n str(i): self.param(f'bias_{i}', initializers.ones, self.xshape)\n for i in range(4)\n }\n\n def __call__(self, x):\n return x + self.biases['0']\n\n x = jnp.array([1.0])\n scope = Scope({}, {'params': rngkey}, mutable=['params'])\n y = Dummy(x.shape, parent=scope)(x)\n self.assertEqual(y, jnp.array([2.0]))\n\n def test_submodule_var_collision_with_scope(self):\n rngkey = jax.random.key(0)\n\n class Dummy(nn.Module):\n xshape: tuple[int, ...]\n\n def setup(self):\n self.bias = self.param('bias', initializers.ones, self.xshape)\n self.bias = DummyModule()\n\n def __call__(self, x):\n return x + self.bias\n\n x = jnp.array([1.0])\n scope = Scope({}, {'params': rngkey}, mutable=['params'])\n\n with self.assertRaises(errors.NameInUseError):\n unused_y = Dummy(x.shape, parent=scope)(x)\n\n def test_submodule_var_collision_with_submodule(self):\n rngkey = jax.random.key(0)\n\n class Dummy(nn.Module):\n xshape: tuple[int, ...]\n\n def setup(self):\n self.bias = self.param('bias', initializers.ones, self.xshape)\n\n @compact\n def __call__(self, x):\n unused_bias = DummyModule(name='bias')\n return x + self.bias\n\n x = jnp.array([1.0])\n scope = Scope({}, {'params': rngkey}, mutable=['params'])\n\n msg = 'Could not create submodule \"bias\" in Module Dummy: Name in use'\n with self.assertRaisesRegex(errors.NameInUseError, msg):\n unused_y = Dummy(x.shape, parent=scope)(x)\n\n def test_submodule_var_collision_with_params(self):\n rngkey = jax.random.key(0)\n\n class Dummy(nn.Module):\n xshape: tuple[int, ...]\n\n def setup(self):\n self.bias = DummyModule()\n\n @compact\n def __call__(self, x):\n unused_bias = self.param('bias', initializers.ones, self.xshape)\n return x + self.bias\n\n x = jnp.array([1.0])\n scope = Scope({}, {'params': rngkey}, mutable=['params'])\n\n msg = 'Could not create param \"bias\" in Module Dummy: Name in use'\n with self.assertRaisesRegex(errors.NameInUseError, msg):\n unused_y = Dummy(x.shape, parent=scope)(x)\n\n def test_attr_empty_container(self):\n class Foo(nn.Module):\n bar: Mapping[str, Any]\n\n @compact\n def __call__(self):\n pass\n\n Foo({'a': ()}).apply({})\n\n def test_multiple_compact_methods(self):\n \"\"\"Test that multiple methods with the @compact decorator can be used.\n\n NOTE: in the near future we might want to have compact methods reset the\n autoname_cursor such that Dense would be reused in the second method.\n \"\"\"\n\n class MultipleCompactMethods(nn.Module):\n @compact\n def __call__(self, x):\n x = nn.Dense(1)(x)\n return self.method(x)\n\n @compact\n def method(self, x):\n x = nn.Dense(1)(x)\n return x\n\n m = MultipleCompactMethods()\n variables = m.init(random.key(0), jnp.ones((1, 1)))\n params = variables['params']\n self.assertIn('Dense_0', params)\n self.assertIn('Dense_1', params)\n\n def test_only_one_compact_method_subclass(self):\n class Dummy(nn.Module):\n @nn.compact\n def __call__(self):\n pass\n\n class SubDummy(Dummy):\n @nn.compact\n def __call__(self):\n super().__call__()\n\n scope = Scope(variables={})\n\n subdummy = SubDummy(parent=scope)\n # Make sure the @compact annotation is valid on both base class and\n # subclass, as long as its on the same method.\n subdummy()\n\n def test_forgotten_compact_annotation(self):\n class Bar(nn.Module):\n # user forgot to add @compact\n def __call__(self, x):\n return nn.Dense(1)(x)\n\n class Foo(nn.Module):\n @nn.compact\n def __call__(self, x):\n bar = Bar()\n x = bar(x)\n x = bar(x)\n return x\n\n msg = (\n r'Submodule Dense must be defined in `setup\\(\\)` or in a method '\n 'wrapped in `@compact`'\n )\n with self.assertRaisesRegex(errors.AssignSubModuleError, msg):\n Foo().init(random.key(0), jnp.ones((1, 3)))\n\n def test_forgotten_compact_annotation_with_explicit_parent(self):\n class Bar(nn.Module):\n def __call__(self, x):\n return nn.Dense(1, parent=self)(x)\n\n class Foo(nn.Module):\n @nn.compact\n def __call__(self, x):\n bar = Bar()\n x = bar(x)\n x = bar(x)\n return x\n\n msg = (\n r'Submodule Dense must be defined in `setup\\(\\)` or in a method '\n 'wrapped in `@compact`'\n )\n with self.assertRaisesRegex(errors.AssignSubModuleError, msg):\n Foo().init(random.key(0), jnp.ones((1, 3)))\n\n def test_numpy_array_shape_class_args(self):\n class MLP(nn.Module):\n widths: Sequence[int]\n\n @nn.compact\n def __call__(self, x):\n for width in self.widths[:-1]:\n x = nn.relu(nn.Dense(width)(x))\n return nn.Dense(self.widths[-1])(x)\n\n test = MLP(np.array([3, 3], np.int32))\n params = test.init({'params': random.key(42)}, jnp.ones((3, 3)))\n _ = test.apply(params, jnp.ones((3, 3)))\n\n def test_get_local_methods(self):\n class Base:\n @staticmethod\n def bar(x):\n return x\n\n @classmethod\n def baz(cls, x):\n return x\n\n def bleep(self, x):\n return x\n\n class Derived1(Base):\n @staticmethod\n def bar2(x):\n return x\n\n @classmethod\n def baz2(cls, x):\n return x\n\n def bloop(self, x):\n return x\n\n class Derived2(Derived1):\n pass\n\n self.assertEqual(nn.module._get_local_method_names(Base), ('bleep',))\n self.assertEqual(nn.module._get_local_method_names(Derived1), ('bloop',))\n self.assertEqual(\n nn.module._get_local_method_names(Derived1, exclude=('bloop',)), ()\n )\n self.assertEqual(nn.module._get_local_method_names(Derived2), ())\n\n def test_inheritance_dataclass_attribs(self):\n class Test(nn.Module):\n bar: int\n\n def __call__(self, x):\n return x\n\n class Test2(Test):\n baz: int\n\n def __call__(self, x):\n return x\n\n class Test3(Test):\n baz: int\n\n def __call__(self, x):\n return x\n\n class Test4(Test2):\n def __call__(self, x):\n return x\n\n key = random.key(0)\n x = jnp.ones((5,))\n test1 = Test(bar=4)\n test2 = Test2(bar=4, baz=2)\n test3 = Test3(bar=4, baz=2)\n test4 = Test4(bar=5, baz=3)\n self.assertEqual(test1.init_with_output(key, x), (x, freeze({})))\n self.assertEqual(test2.init_with_output(key, x), (x, freeze({})))\n self.assertEqual(test3.init_with_output(key, x), (x, freeze({})))\n self.assertEqual(test4.init_with_output(key, x), (x, freeze({})))\n self.assertTrue(hasattr(test1, 'bar'))\n self.assertTrue(hasattr(test1, 'name'))\n self.assertTrue(hasattr(test1, 'parent'))\n self.assertTrue(hasattr(test2, 'bar'))\n self.assertTrue(hasattr(test2, 'baz'))\n self.assertTrue(hasattr(test2, 'name'))\n self.assertTrue(hasattr(test2, 'parent'))\n self.assertTrue(hasattr(test3, 'bar'))\n self.assertTrue(hasattr(test3, 'baz'))\n self.assertTrue(hasattr(test3, 'name'))\n self.assertTrue(hasattr(test3, 'parent'))\n self.assertTrue(hasattr(test4, 'bar'))\n self.assertTrue(hasattr(test4, 'baz'))\n self.assertTrue(hasattr(test4, 'name'))\n self.assertTrue(hasattr(test4, 'parent'))\n self.assertEqual(\n list(Test.__dataclass_fields__.keys()), ['bar', 'parent', 'name']\n )\n self.assertEqual(\n list(Test2.__dataclass_fields__.keys()),\n ['bar', 'baz', 'parent', 'name'],\n )\n self.assertEqual(\n list(Test3.__dataclass_fields__.keys()),\n ['bar', 'baz', 'parent', 'name'],\n )\n self.assertEqual(\n list(Test4.__dataclass_fields__.keys()),\n ['bar', 'baz', 'parent', 'name'],\n )\n\n def test_get_suffix_value_pairs(self):\n for x in [(), [], {}, None, 0, set()]:\n self.assertEqual(nn.module._get_suffix_value_pairs(x), [('', x)])\n self.assertEqual(\n nn.module._get_suffix_value_pairs({'a': 1, 'b': 2}),\n [('_a', 1), ('_b', 2)],\n )\n self.assertEqual(\n nn.module._get_suffix_value_pairs([1, 2, 3]),\n [('_0', 1), ('_1', 2), ('_2', 3)],\n )\n x1 = [nn.Dense(10), nn.relu, nn.Dense(10)]\n y1 = nn.module._get_suffix_value_pairs(x1)\n self.assertEqual(y1, [('_0', x1[0]), ('_1', x1[1]), ('_2', x1[2])])\n x2 = {'a': 1, 'b': {'c': nn.Dense(10), 'd': nn.relu}}\n y2 = nn.module._get_suffix_value_pairs(x2)\n self.assertEqual(\n y2, [('_a', 1), ('_b_c', x2['b']['c']), ('_b_d', x2['b']['d'])]\n )\n\n def test_mixed_list_assignment_in_setup(self):\n class Test(nn.Module):\n def setup(self):\n self.layers = [nn.Dense(10), nn.relu, nn.Dense(10)]\n\n def __call__(self, x):\n for lyr in self.layers:\n x = lyr(x)\n return x\n\n x = random.uniform(random.key(0), (5, 5))\n variables = Test().init(random.key(0), jnp.ones((5, 5)))\n y = Test().apply(variables, x)\n m0 = variables['params']['layers_0']['kernel']\n m1 = variables['params']['layers_2']['kernel']\n self.assertTrue(jnp.all(y == jnp.dot(nn.relu(jnp.dot(x, m0)), m1)))\n\n def test_module_is_hashable(self):\n module_a = nn.Dense(10)\n module_a_2 = nn.Dense(10)\n module_b = nn.Dense(5)\n self.assertEqual(hash(module_a), hash(module_a_2))\n self.assertNotEqual(hash(module_a), hash(module_b))\n\n def test_module_custom_hash(self):\n class Test(nn.Module):\n x: int = 3\n y: int = 5\n\n def __hash__(self):\n return 42 + self.x\n\n module_a = Test(1, 2)\n module_a_2 = Test(1, 5)\n module_b = Test(2, 2)\n self.assertEqual(hash(module_a), hash(module_a_2))\n self.assertNotEqual(hash(module_a), hash(module_b))\n\n def test_module_with_scope_is_not_hashable(self):\n module_a = nn.Dense(10, parent=Scope({}))\n msg = \"Can't call __hash__ on modules that hold variables.\"\n with self.assertRaisesWithLiteralMatch(TypeError, msg):\n hash(module_a)\n\n def test_module_trace(self):\n class MLP(nn.Module):\n act: Callable = nn.relu\n sizes: Sequence[int] = (3, 2)\n\n @nn.compact\n def __call__(self, x):\n for size in self.sizes:\n x = nn.Dense(size)(x)\n x = self.act(x)\n return repr(self)\n\n mlp = MLP()\n expected_trace = \"\"\"MLP(\n # attributes\n act = relu\n sizes = (3, 2)\n # children\n Dense_0 = Dense(\n # attributes\n features = 3\n use_bias = True\n dtype = None\n param_dtype = float32\n precision = None\n kernel_init = init\n bias_init = zeros\n promote_dtype = promote_dtype\n dot_general = None\n dot_general_cls = None\n )\n Dense_1 = Dense(\n # attributes\n features = 2\n use_bias = True\n dtype = None\n param_dtype = float32\n precision = None\n kernel_init = init\n bias_init = zeros\n promote_dtype = promote_dtype\n dot_general = None\n dot_general_cls = None\n )\n)\"\"\"\n x = jnp.ones((1, 2))\n trace, variables = mlp.init_with_output(random.key(0), x)\n self.assertEqual(trace, expected_trace)\n trace = mlp.apply(variables, x)\n self.assertEqual(trace, expected_trace)\n\n def test_default_params_rng_equivalence(self):\n class Model(nn.Module):\n @nn.compact\n def __call__(self, x, add_dropout=False, add_noise=False):\n x = nn.Dense(16)(x)\n x = nn.Dropout(0.5)(x, deterministic=not add_dropout)\n if add_noise:\n x += jax.random.normal(self.make_rng('params'))\n return x\n\n model = Model()\n key0, key1, key2 = jax.random.split(jax.random.key(0), 3)\n x = jax.random.normal(key0, (10, 8))\n\n with self.assertRaisesRegex(\n ValueError,\n 'First argument passed to an init function should be a ``jax.PRNGKey``',\n ):\n model.init({'params': 'test'}, x)\n with self.assertRaisesRegex(\n errors.InvalidRngError,\n 'RNGs should be of shape \\\\(2,\\\\) or PRNGKey in Module Model, but rngs are: test',\n ):\n model.init('test', x)\n # should not throw an error, since nn.Dropout will get an RNG key from the 'params' stream\n model.init(key1, x, add_dropout=True)\n\n v = model.init({'params': key1}, x)\n v2 = model.init(key1, x)\n jax.tree_util.tree_map(np.testing.assert_allclose, v, v2)\n\n for add_dropout, add_noise in [[True, False], [False, True], [True, True]]:\n out = model.apply(\n v,\n x,\n add_dropout=add_dropout,\n add_noise=add_noise,\n rngs={'params': key2},\n )\n out2 = model.apply(\n v, x, add_dropout=add_dropout, add_noise=add_noise, rngs=key2\n )\n np.testing.assert_allclose(out, out2)\n\n with self.assertRaisesRegex(\n ValueError,\n 'The ``rngs`` argument passed to an apply function should be a ``jax.PRNGKey``',\n ):\n model.apply(v, x, rngs={'params': 'test'})\n with self.assertRaisesRegex(\n errors.InvalidRngError,\n 'RNGs should be of shape \\\\(2,\\\\) or PRNGKey in Module Model, but rngs are: test',\n ):\n model.apply(v, x, rngs='test')\n\n def test_module_apply_method(self):\n class Foo(nn.Module):\n not_callable: int = 1\n\n @nn.compact\n def __call__(self):\n pass\n\n def test(self):\n pass\n\n # We can use both instance and class methods in apply.\n Foo().apply({}, method=Foo.test)\n Foo().apply({}, method=Foo().test)\n\n # We also use a function that is not in the provided Module, although it\n # should have a first argument representing an instance of the Module (Foo\n # in this case).\n x = Foo().apply({}, method=lambda foo_instance: foo_instance)\n self.assertEqual(type(x), type(Foo()))\n\n # This is not allowed.\n msg = 'Cannot call apply()'\n with self.assertRaisesRegex(errors.ApplyModuleInvalidMethodError, msg):\n Foo().apply({}, method=lambda: True)\n\n # string method names are also allowed.\n Foo().apply({}, method='test')\n # test same for init.\n Foo().init({}, method='test')\n\n # non-existent attribute names will yield AttributeError.\n with self.assertRaisesRegex(AttributeError, 'allowed_apply_fn'):\n Foo().apply({}, method='allowed_apply_fn')\n # test same for init.\n Foo().init({}, method='allowed_apply_fn')\n\n # attributes which are not callables yield TypeError.\n with self.assertRaisesRegex(\n TypeError, \"'Foo.not_callable' must be a callable\"\n ):\n Foo().apply({}, method='not_callable')\n # test same for init.\n Foo().init({}, method='not_callable')\n\n def test_module_apply_method_submodule(self):\n class Foo(nn.Module):\n bar: nn.Module\n\n @nn.compact\n def __call__(self, x):\n return self.bar(x)\n\n foo = Foo(nn.Dense(3))\n variables = foo.init(jax.random.PRNGKey(0), jnp.zeros(3))\n\n foo.apply(variables, jnp.zeros(3), method='bar')\n\n def test_call_unbound_compact_module_methods(self):\n dense = Dense(3)\n msg = r'Can\\'t call compact methods on unbound modules'\n with self.assertRaisesRegex(errors.CallCompactUnboundModuleError, msg):\n dense(jnp.ones((1,)))\n\n def test_call_unbound_has_variable(self):\n class EmptyModule(nn.Module):\n def foo(self):\n self.has_variable('bar', 'baz')\n\n empty = EmptyModule()\n with self.assertRaisesRegex(ValueError, 'variable.*unbound module'):\n empty.foo()\n\n def test_call_unbound_make_rng(self):\n class EmptyModule(nn.Module):\n def foo(self):\n self.make_rng('bar')\n\n empty = EmptyModule()\n with self.assertRaisesRegex(ValueError, 'RNGs.*unbound module'):\n empty.foo()\n\n def test_call_unbound_variables(self):\n class EmptyModule(nn.Module):\n def foo(self):\n self.variables\n\n empty = EmptyModule()\n with self.assertRaisesRegex(ValueError, 'variables.*unbound module'):\n empty.foo()\n\n def test_call_unbound_noncompact_module_methods(self):\n class EmptyModule(nn.Module):\n foo: int = 3\n\n def bar(self):\n return self.foo\n\n empty = EmptyModule()\n # It's fine to call methods of unbound methods that don't depend on\n # attributes defined during `setup`.\n self.assertEqual(empty.bar(), 3)\n\n def test_call_unbound_noncompact_module_methods_depending_on_setup(self):\n class EmptyModule(nn.Module):\n def setup(self):\n self.foo = 2\n\n def bar(self):\n return self.foo\n\n empty = EmptyModule()\n msg = r'\"EmptyModule\" object has no attribute \"foo\"'\n with self.assertRaisesRegex(AttributeError, msg):\n empty.bar()\n\n def test_module_with_attrs(self):\n class Foo(nn.Module):\n bar: nn.Dense = dataclasses.field(init=False)\n\n def setup(self):\n self.bar = nn.Dense(3)\n\n def __call__(self, x):\n return self.bar(x)\n\n foo = Foo()\n x = jnp.ones((2,))\n variables = foo.init(random.key(0), x)\n self.assertEqual(variables['params']['bar']['kernel'].shape, (2, 3))\n\n def test_noncompact_module_frozen(self):\n class Foo(nn.Module):\n def setup(self):\n self.i = 1 # This is allowed (for assigning submodules).\n\n def __call__(self):\n self.i = 2 # This is not allowed.\n\n msg = (\n \"Can't set i=2 for Module of type Foo: Module instance is frozen \"\n 'outside of setup method.'\n )\n with self.assertRaisesRegex(errors.SetAttributeFrozenModuleError, msg):\n Foo().init(random.key(0))\n\n def test_compact_module_frozen(self):\n class Foo(nn.Module):\n @nn.compact\n def __call__(self):\n self.i = 2\n\n msg = (\n \"Can't set i=2 for Module of type Foo: Module instance is frozen \"\n 'outside of setup method.'\n )\n with self.assertRaisesRegex(errors.SetAttributeFrozenModuleError, msg):\n Foo().init(random.key(0))\n\n def test_submodule_frozen(self):\n class Foo(nn.Module):\n @nn.compact\n def __call__(self):\n dense = nn.Dense(10)\n dense.features = 20 # <--- This is not allowed\n\n msg = (\n \"Can't set features=20 for Module of type Dense: Module instance \"\n 'is frozen outside of setup method.'\n )\n with self.assertRaisesRegex(errors.SetAttributeFrozenModuleError, msg):\n Foo().init(random.key(0))\n\n def test_module_call_not_implemented(self):\n class Foo(nn.Module):\n pass\n\n msg = '\"Foo\" object has no attribute \"__call__\"'\n with self.assertRaisesRegex(AttributeError, msg):\n Foo().init(random.key(0))\n\n def test_is_mutable_collection(self):\n class EmptyModule(nn.Module):\n def __call__(self):\n return self.is_mutable_collection('test')\n\n empty = EmptyModule()\n self.assertTrue(empty.apply({}, mutable=['test'])[0])\n self.assertFalse(empty.apply({}, mutable=False))\n\n def test_module_lazy_getattr_setup(self):\n class A(nn.Module):\n def setup(self):\n self.d = nn.Dense(2)\n\n def __call__(self, x):\n return self.d(x)\n\n class B(nn.Module):\n def setup(self):\n self.a = A()\n\n def __call__(self, x):\n y1 = self.a.d(x)\n y2 = self.a(x)\n return y1, y2\n\n key = random.key(0)\n x = jnp.ones((2,))\n\n (y1, y2), unused_vars = B().init_with_output(key, x)\n np.testing.assert_array_equal(y1, y2)\n\n def test_module_lazy_dir_setup(self):\n class A(nn.Module):\n def setup(self):\n self.d = nn.Dense(2)\n\n def __call__(self, x):\n return self.d(x)\n\n class B(nn.Module):\n def setup(self):\n self.a = A()\n\n def __call__(self, x):\n assert 'd' in dir(self.a)\n y1 = self.a.d(x)\n y2 = self.a(x)\n return y1, y2\n\n key = random.key(0)\n x = jnp.ones((2,))\n _ = B().init_with_output(key, x)\n\n def test_module_unbound_getattr(self):\n class A(nn.Module):\n def setup(self):\n b = B()\n b.c # B is unbound because it is not yet assigned to an attribute.\n self.b = b\n\n def __call__(self):\n pass\n\n class B(nn.Module):\n def setup(self):\n self.c = nn.Dense(2)\n\n msg = '\"B\" object has no attribute \"c\"'\n with self.assertRaisesRegex(AttributeError, msg):\n A().init(random.key(0))\n\n def test_unbound_setup_call(self):\n setup_called = False\n\n class A(nn.Module):\n def setup(self):\n nonlocal setup_called\n setup_called = True\n\n def test(self):\n pass\n\n A().test()\n self.assertFalse(setup_called)\n\n def test_module_pass_as_attr(self):\n class A(nn.Module):\n def setup(self):\n self.b = B(nn.Dense(2))\n\n def __call__(self, x):\n return self.b(x)\n\n class B(nn.Module):\n foo: Any\n\n def __call__(self, x):\n return self.foo(x)\n\n variables = A().init(random.key(0), jnp.ones((1,)))\n var_shapes = jax.tree_util.tree_map(jnp.shape, variables)\n ref_var_shapes = {\n 'params': {\n 'b': {\n 'foo': {\n 'bias': (2,),\n 'kernel': (1, 2),\n }\n },\n },\n }\n self.assertTrue(tree_equals(var_shapes, ref_var_shapes))\n\n def test_module_pass_in_closure(self):\n a = nn.Dense(2)\n\n class B(nn.Module):\n def setup(self):\n self.foo = a\n\n def __call__(self, x):\n return self.foo(x)\n\n variables = B().init(random.key(0), jnp.ones((1,)))\n var_shapes = jax.tree_util.tree_map(jnp.shape, variables)\n ref_var_shapes = {\n 'params': {\n 'foo': {\n 'bias': (2,),\n 'kernel': (1, 2),\n }\n },\n }\n self.assertTrue(tree_equals(var_shapes, ref_var_shapes))\n self.assertIsNone(a.name)\n\n def test_toplevel_submodule_adoption(self):\n class Encoder(nn.Module):\n n_layers: int\n ch: int\n\n def setup(self):\n self.layers = [nn.Dense(self.ch) for _ in range(self.n_layers)]\n\n def __call__(self, x):\n for layer in self.layers:\n x = layer(x)\n x = nn.relu(x)\n return x\n\n class Model(nn.Module):\n encoder: nn.Module\n n_out: int\n\n def setup(self):\n self.dense_out = nn.Dense(self.n_out)\n\n def __call__(self, x):\n x = self.encoder(x)\n return self.dense_out(x)\n\n # Define model.\n encoder = Encoder(n_layers=1, ch=8)\n model = Model(encoder=encoder, n_out=5)\n\n # Initialize.\n key = jax.random.key(0)\n x = random.uniform(key, (4, 4))\n\n variables = model.init(key, x)\n y = model.apply(variables, x)\n self.assertEqual(y.shape, (4, 5))\n\n var_shapes = jax.tree_util.tree_map(jnp.shape, variables)\n ref_var_shapes = {\n 'params': {\n 'dense_out': {\n 'bias': (5,),\n 'kernel': (8, 5),\n },\n 'encoder': {\n 'layers_0': {\n 'bias': (8,),\n 'kernel': (4, 8),\n },\n },\n },\n }\n self.assertTrue(tree_equals(var_shapes, ref_var_shapes))\n\n def test_toplevel_submodule_adoption_pytree(self):\n class A(nn.Module):\n @nn.compact\n def __call__(self, c, x):\n counter = self.variable('counter', 'i', jnp.zeros, ())\n counter.value += 1\n x = nn.Dense(1)(x)\n return c, x\n\n class B(nn.Module):\n A: Any\n\n @nn.compact\n def __call__(self, c, x):\n return self.A['foo'](*self.A['bar'](c, x))\n\n unused_a = A()\n a_pytree = {'foo': A(), 'bar': A()}\n b = B(a_pytree)\n\n key = random.key(0)\n x = jnp.ones((2, 2))\n\n params = B(a_pytree).init(key, x, x)\n unused_y, counters = b.apply(params, x, x, mutable='counter')\n ref_counters = {\n 'counter': {\n 'A_bar': {\n 'i': jnp.array(2.0),\n },\n 'A_foo': {\n 'i': jnp.array(2.0),\n },\n },\n }\n self.assertTrue(\n jax.tree_util.tree_all(\n jax.tree_util.tree_map(\n lambda x, y: np.testing.assert_allclose(x, y, atol=1e-7),\n counters,\n ref_counters,\n )\n )\n )\n\n def test_toplevel_submodule_adoption_sharing(self):\n dense = functools.partial(nn.Dense, use_bias=False)\n\n class A(nn.Module):\n @nn.compact\n def __call__(self, x):\n return dense(2)(x)\n\n class B(nn.Module):\n a: nn.Module\n\n @nn.compact\n def __call__(self, x):\n return dense(2)(x) + self.a(x)\n\n class C(nn.Module):\n a: nn.Module\n b: nn.Module\n\n @nn.compact\n def __call__(self, x):\n return dense(2)(x) + self.b(x) + self.a(x)\n\n key = random.key(0)\n x = jnp.ones((2, 2))\n a = A()\n b = B(a)\n c = C(a, b)\n p = c.init(key, x)\n var_shapes = jax.tree_util.tree_map(jnp.shape, p)\n ref_var_shapes = {\n 'params': {\n 'Dense_0': {\n 'kernel': (2, 2),\n },\n 'a': {\n 'Dense_0': {\n 'kernel': (2, 2),\n },\n },\n 'b': {\n 'Dense_0': {\n 'kernel': (2, 2),\n },\n },\n },\n }\n self.assertTrue(tree_equals(var_shapes, ref_var_shapes))\n\n def test_toplevel_named_submodule_adoption(self):\n dense = functools.partial(nn.Dense, use_bias=False)\n\n class A(nn.Module):\n def setup(self):\n self.dense = dense(4)\n\n def __call__(self, x):\n return self.dense(x)\n\n class B(nn.Module):\n a: A\n\n def setup(self):\n self.proj = dense(6)\n\n def __call__(self, x):\n return self.proj(self.a(x))\n\n a = A(name='foo')\n b = B(a=a)\n k = jax.random.key(0)\n x = jnp.zeros((5, 5))\n init_vars = b.init(k, x)\n var_shapes = jax.tree_util.tree_map(jnp.shape, init_vars)\n if config.flax_preserve_adopted_names:\n ref_var_shapes = {\n 'params': {\n 'foo': {\n 'dense': {\n 'kernel': (5, 4),\n },\n },\n 'proj': {\n 'kernel': (4, 6),\n },\n },\n }\n else:\n ref_var_shapes = {\n 'params': {\n 'a': {\n 'dense': {\n 'kernel': (5, 4),\n },\n },\n 'proj': {\n 'kernel': (4, 6),\n },\n },\n }\n self.assertTrue(tree_equals(var_shapes, ref_var_shapes))\n\n def test_toplevel_submodule_pytree_adoption_sharing(self):\n class A(nn.Module):\n @nn.compact\n def __call__(self, x):\n counter = self.variable('counter', 'i', jnp.zeros, ())\n counter.value += 1\n x = nn.Dense(1)(x)\n return x\n\n class B(nn.Module):\n A: Any\n\n @nn.compact\n def __call__(self, x):\n return self.A['foo'](x) + self.A['bar'](x) + self.A['baz'](x)\n\n key = random.key(0)\n x = jnp.ones((2, 2))\n\n a = A()\n a_pytree = {'foo': a, 'bar': a, 'baz': a}\n b = B(a_pytree)\n\n params = b.init(key, x)\n _, counters = b.apply(params, x, mutable='counter')\n ref_counters = {\n 'counter': {\n 'A_bar': {\n 'i': jnp.array(6.0),\n },\n },\n }\n self.assertTrue(tree_equals(counters, ref_counters))\n\n def test_inner_class_def(self):\n class X(nn.Module):\n class Hyper(struct.PyTreeNode):\n a: int\n\n hyper: Hyper\n\n @nn.compact\n def __call__(self, x):\n return x + 1\n\n self.assertIsInstance(X.Hyper(a=1), X.Hyper)\n\n def test_sow(self):\n class Foo(nn.Module):\n @nn.compact\n def __call__(self, x, **sow_args):\n self.sow('intermediates', 'h', x, **sow_args)\n self.sow('intermediates', 'h', 2 * x, **sow_args)\n return 3 * x\n\n variables = Foo().init(random.key(0), 1)\n # During init we should not collect intermediates by default...\n self.assertNotIn('intermediates', variables)\n # ...unless we override mutable.\n variables = Foo().init(random.key(0), 1, mutable=True)\n self.assertEqual(variables, {'intermediates': {'h': (1, 2)}})\n\n _, state = Foo().apply({}, 1, mutable=['intermediates'])\n self.assertEqual(state, {'intermediates': {'h': (1, 2)}})\n _, state = Foo().apply(\n {},\n 1,\n init_fn=lambda: 0,\n reduce_fn=lambda a, b: a + b,\n mutable=['intermediates'],\n )\n self.assertEqual(state, {'intermediates': {'h': 3}})\n self.assertEqual(Foo().apply({}, 1), 3)\n\n def test_capture_intermediates(self):\n class Bar(nn.Module):\n def test(self, x):\n return x + 1\n\n class Foo(nn.Module):\n @nn.compact\n def __call__(self, x):\n return Bar().test(x) + 1\n\n _, state = Foo().apply({}, 1, capture_intermediates=True)\n self.assertEqual(state, {'intermediates': {'__call__': (3,)}})\n fn = lambda mdl, _: isinstance(mdl, Bar)\n _, state = Foo().apply({}, 1, capture_intermediates=fn)\n self.assertEqual(state, {'intermediates': {'Bar_0': {'test': (2,)}}})\n\n def test_perturb(self):\n class Foo(nn.Module):\n @nn.compact\n def __call__(self, x):\n x = nn.Dense(10)(x)\n x = self.perturb('before_multiply', x)\n x = 4 * x\n x = self.perturb('after_multiply', x)\n return x\n\n def loss(params, perturbations, inputs, targets):\n variables = {'params': params, 'perturbations': perturbations}\n preds = Foo().apply(variables, inputs)\n return jnp.square(preds - targets).mean()\n\n x = jax.random.uniform(jax.random.key(1), shape=(10,))\n y = jax.random.uniform(jax.random.key(2), shape=(10,))\n variables = Foo().init(jax.random.key(0), x)\n intm_grads = jax.grad(loss, argnums=1)(\n variables['params'], variables['perturbations'], x, y\n )\n # activation * 4 so reverse gradient also * 4\n self.assertTrue(\n all(intm_grads['after_multiply'] * 4 == intm_grads['before_multiply'])\n )\n\n def test_perturb_setup(self):\n class Foo(nn.Module):\n def setup(self):\n self.a = nn.Dense(10)\n\n def __call__(self, x):\n x = self.a(x)\n x = self.perturb('before_multiply', x)\n x = 4 * x\n x = self.perturb('after_multiply', x)\n return x\n\n def loss(params, perturbations, inputs, targets):\n variables = {'params': params, 'perturbations': perturbations}\n preds = Foo().apply(variables, inputs)\n return jnp.square(preds - targets).mean()\n\n x = jax.random.uniform(jax.random.key(1), shape=(10,))\n y = jax.random.uniform(jax.random.key(2), shape=(10,))\n variables = Foo().init(jax.random.key(0), x)\n intm_grads = jax.grad(loss, argnums=1)(\n variables['params'], variables['perturbations'], x, y\n )\n # activation * 4 so reverse gradient also * 4\n self.assertTrue(\n all(intm_grads['after_multiply'] * 4 == intm_grads['before_multiply'])\n )\n\n def test_perturb_noop(self):\n class Foo(nn.Module):\n @nn.compact\n def __call__(self, x):\n x = nn.Dense(10)(x)\n x = self.perturb('before_multiply', x)\n x = 4 * x\n x = self.perturb('after_multiply', x)\n return x\n\n x = jax.random.uniform(jax.random.key(1), shape=(10,))\n module = Foo()\n variables = module.init(jax.random.key(0), x)\n params = variables['params']\n perturbations = variables['perturbations']\n\n # check no error if perturbations is not passed\n module.apply({'params': params}, x)\n\n # check errors if perturbations is passed but empty\n with self.assertRaisesRegex(ValueError, 'Perturbation collection'):\n module.apply({'params': params, 'perturbations': {}}, x)\n\n # check no error if perturbations is passed and not empty\n module.apply({'params': params, 'perturbations': perturbations}, x)\n\n def test_functional_apply(self):\n class Foo(nn.Module):\n def setup(self):\n self.a = nn.Dense(3)\n self.b = nn.Dense(1)\n\n def f(foo, x):\n x = foo.a(x)\n return foo.b(x)\n\n foo = Foo()\n x = jnp.ones((4,))\n f_init = nn.init_with_output(f, foo)\n f_apply = nn.apply(f, foo)\n y1, variables = f_init(random.key(0), x)\n y2 = f_apply(variables, x)\n self.assertEqual(y1, y2)\n\n def test_bind(self):\n class Foo(nn.Module):\n def setup(self):\n self.a = nn.Dense(3)\n self.b = nn.Dense(1)\n\n def f(foo, x):\n x = foo.a(x)\n return foo.b(x)\n\n foo = Foo()\n x = jnp.ones((4,))\n f_init = nn.init_with_output(f, foo)\n y1, variables = f_init(random.key(0), x)\n y2 = f(foo.bind(variables), x)\n self.assertEqual(y1, y2)\n\n def test_bind_stateful(self):\n class Foo(nn.Module):\n def setup(self):\n self.a = nn.Dense(3)\n self.bn = nn.BatchNorm()\n self.b = nn.Dense(1)\n\n def f(foo, x):\n x = foo.a(x)\n x = foo.bn(x, use_running_average=False)\n return foo.b(x)\n\n foo = Foo()\n x = jnp.ones((4,))\n f_init = nn.init_with_output(f, foo)\n y1, variables = f_init(random.key(0), x)\n foo_b = foo.bind(variables, mutable='batch_stats')\n y2 = f(foo_b, x)\n y3, new_state = nn.apply(f, foo, mutable='batch_stats')(variables, x)\n self.assertEqual(y1, y2)\n self.assertEqual(y2, y3)\n bs_1 = new_state['batch_stats']\n bs_2 = foo_b.variables['batch_stats']\n for x, y in zip(\n jax.tree_util.tree_leaves(bs_1), jax.tree_util.tree_leaves(bs_2)\n ):\n np.testing.assert_allclose(x, y)\n\n def test_unbind(self):\n class Foo(nn.Module):\n def setup(self):\n self.encoder = nn.Dense(4)\n self.decoder = nn.Dense(2)\n\n def __call__(self, x):\n x = self.encoder(x)\n return self.decoder(x)\n\n foo = Foo()\n x = jnp.ones((2,))\n\n variables = foo.init(random.key(0), x)\n encoder, encoder_vars = foo.bind(variables).encoder.unbind()\n decoder, decoder_vars = foo.bind(variables).decoder.unbind()\n\n self.assertIsInstance(encoder, nn.Dense)\n self.assertEqual(encoder.features, 4)\n self.assertIsInstance(decoder, nn.Dense)\n self.assertEqual(decoder.features, 2)\n\n self.assertTrue(\n jax.tree_util.tree_all(\n jax.tree_util.tree_map(\n lambda v1, v2: (v1 == v2).all(),\n variables['params']['encoder'],\n encoder_vars['params'],\n )\n )\n )\n self.assertTrue(\n jax.tree_util.tree_all(\n jax.tree_util.tree_map(\n lambda v1, v2: (v1 == v2).all(),\n variables['params']['decoder'],\n decoder_vars['params'],\n )\n )\n )\n\n def test_bind_unbind_equality(self):\n class Foo(nn.Module):\n sub_module: Any\n\n @nn.compact\n def __call__(self, x):\n x = nn.Dense(2)(x)\n return self.sub_module(x)\n\n sub_module = Foo(nn.Dense(3))\n module = Foo(sub_module)\n x = jnp.ones((1, 2))\n variables = module.init(jax.random.PRNGKey(0), x)\n\n bound_module = module.bind(variables)\n self.assertTrue((module.apply(variables, x) == bound_module(x)).all())\n new_module, new_variables = bound_module.unbind()\n self.assertTrue(\n jax.tree_util.tree_all(\n jax.tree_util.tree_map(\n lambda v1, v2: (v1 == v2).all(), variables, new_variables\n )\n )\n )\n self.assertEqual(module, new_module)\n\n def test_passing_mutable_variables(self):\n class Foo(nn.Module):\n @nn.compact\n def __call__(self, x):\n return nn.Dense(2)(x)\n\n x = jnp.ones((3,))\n variables = Foo().init(random.key(0), x)\n y = Foo().apply(variables, x)\n self.assertEqual(y.shape, (2,))\n\n def test_super_compact(self):\n class Foo(nn.Module):\n @nn.compact\n def __call__(self, x):\n return nn.Dense(4)(x)\n\n class Bar(Foo):\n @nn.compact\n def __call__(self, x):\n y = super().__call__(x)\n return nn.Dense(3)(y)\n\n k = random.key(0)\n x = jnp.ones((4, 7))\n\n variables = Bar().init(k, x)\n shapes = jax.tree_util.tree_map(np.shape, variables['params'])\n self.assertEqual(\n shapes,\n {\n 'Dense_0': {'kernel': (7, 4), 'bias': (4,)},\n 'Dense_1': {'kernel': (4, 3), 'bias': (3,)},\n },\n )\n y = Bar().apply(variables, x)\n self.assertEqual(y.shape, (4, 3))\n\n def test_super_setup(self):\n class Foo(nn.Module):\n def setup(self):\n self.a = nn.Dense(4)\n\n class Bar(Foo):\n def setup(self):\n super().setup()\n self.b = nn.Dense(3)\n\n def __call__(self, x):\n y = self.a(x)\n return self.b(y)\n\n k = random.key(0)\n x = jnp.ones((4, 7))\n\n variables = Bar().init(k, x)\n y = Bar().apply(variables, x)\n self.assertEqual(y.shape, (4, 3))\n\n def test_freeze_attr(self):\n class Foo(NamedTuple):\n a: int\n b: int\n\n self.assertEqual(nn.module._freeze_attr([1, 2]), (1, 2))\n xs = nn.module._freeze_attr(Foo(1, 2))\n self.assertEqual(xs, (1, 2))\n self.assertEqual(\n type(xs), Foo\n ) # equality test for NamedTuple doesn't check class!\n\n def test_generic_multiple_inheritance(self):\n T = TypeVar('T')\n\n class MyComponent(nn.Module, Generic[T]):\n pass\n\n class MyModule(nn.Module):\n submodule: MyComponent[jnp.ndarray]\n\n class MyComponent2(Generic[T], nn.Module):\n pass\n\n class MyModule2(nn.Module):\n submodule: MyComponent2[jnp.ndarray]\n\n def test_jit_rng_equivalance(self):\n model = nn.fold_rngs(nn.Dense)(1, use_bias=False)\n jit_model = nn.jit(nn.Dense)(1, use_bias=False)\n param = model.init(random.key(0), np.ones((1, 1)))['params']['kernel']\n param_2 = jit_model.init(random.key(0), np.ones((1, 1)))['params']['kernel']\n self.assertEqual(param, param_2)\n\n def test_rng_reuse_after_rewind(self):\n class C(nn.Module):\n @nn.compact\n def __call__(self):\n # Some module that has dropouts in it, in general,\n # it does more than just dropout!\n return self.make_rng('dropout')\n\n class A(nn.Module):\n @nn.compact\n def __call__(self):\n # Some module that has dropouts in it, in general,\n # it does more than just dropout!\n return C()()\n\n class B(nn.Module):\n @nn.compact\n def __call__(self):\n a = A()\n x0 = a()\n x1 = a()\n return jnp.all(x0 == x1)\n\n k = random.key(0)\n rng_equals = B().apply({}, rngs={'dropout': k})\n self.assertFalse(rng_equals)\n\n def test_module_get_put_has_variable(self):\n class A(nn.Module):\n @nn.compact\n def __call__(self, x):\n self.put_variable('test_col', 'a', x)\n assert self.has_variable('test_col', 'a')\n return self.get_variable('test_col', 'a')\n\n class B(nn.Module):\n def __call__(self, x):\n self.put_variable('test_col', 'a', x)\n assert self.has_variable('test_col', 'a')\n return self.get_variable('test_col', 'a')\n\n class C(nn.Module):\n def setup(self):\n self.put_variable(\n 'test_col',\n 'a',\n jnp.ones(\n 2,\n ),\n )\n assert self.has_variable('test_col', 'a')\n\n def __call__(self):\n return self.get_variable('test_col', 'a')\n\n x = jnp.ones((2,))\n\n y, vs = A().apply({}, x, mutable=['test_col'])\n np.testing.assert_array_equal(x, y)\n np.testing.assert_array_equal(x, vs['test_col']['a'])\n\n y, vs = B().apply({}, x, mutable=['test_col'])\n np.testing.assert_array_equal(x, y)\n np.testing.assert_array_equal(x, vs['test_col']['a'])\n\n y, vs = C().apply({}, mutable=['test_col'])\n np.testing.assert_array_equal(y, jnp.ones((2,)))\n np.testing.assert_array_equal(y, vs['test_col']['a'])\n\n def test_generic_module(self):\n # See https://github.com/google/flax/issues/1899\n T = TypeVar('T')\n\n class C(nn.Module, Generic[T]):\n def f(self, t: T) -> T:\n return t\n\n class D(nn.Module):\n def setup(self):\n unused_c = C[Any]()\n\n def __call__(self) -> None:\n pass\n\n rngs = {}\n D().init(rngs)\n\n def test_modifying_attribs_in_post_init(self):\n class Foo(nn.Module):\n love: int = 99\n\n def __post_init__(self):\n self.hate = 100 - self.love\n super().__post_init__()\n\n foo = Foo()\n self.assertEqual(foo.love, 99)\n self.assertEqual(foo.hate, 1)\n\n class Bar(nn.Module):\n love: int = 99\n\n def __post_init__(self):\n self.love = 101\n super().__post_init__()\n\n bar = Bar()\n self.assertEqual(bar.love, 101)\n\n def test_has_rng(self):\n class Foo(nn.Module):\n def __call__(self):\n return self.has_rng('bar')\n\n foo = Foo()\n with self.assertRaisesRegex(ValueError, 'RNGs.*unbound module'):\n foo()\n k = random.key(0)\n self.assertTrue(foo.apply({}, rngs={'bar': k}))\n self.assertFalse(foo.apply({}, rngs={'baz': k}))\n\n def test_is_initializing(self):\n class Foo(nn.Module):\n def __call__(self):\n return self.is_initializing()\n\n foo = Foo()\n k = random.key(0)\n self.assertTrue(foo.init_with_output(k)[0])\n self.assertFalse(foo.apply({}))\n\n def test_throws_invalid_instance_module_error(self):\n class B(nn.Module):\n @nn.compact\n def __call__(self, x):\n return x\n\n k = random.key(0)\n x = random.uniform(random.key(1), (2,))\n\n with self.assertRaises(errors.InvalidInstanceModuleError):\n B.init(k, x) # B is module class, not B() a module instance\n with self.assertRaises(errors.InvalidInstanceModuleError):\n B.init_with_output(k, x)\n with self.assertRaises(errors.InvalidInstanceModuleError):\n B.apply(\n {}, x\n ) # similar issue w. apply called on class instead of instance.\n with self.assertRaises(errors.InvalidInstanceModuleError):\n B.bind(\n {}, x\n ) # similar issue w. apply called on class instead of instance.\n\n def test_throws_incorrect_post_init_override_error(self):\n class A(nn.Module):\n x: float\n\n def __post_init__(self):\n self.x_square = self.x**2\n\n @nn.compact\n def __call__(self, input):\n return input + 3\n\n r = A(x=3)\n\n with self.assertRaises(errors.IncorrectPostInitOverrideError):\n r.init(jax.random.key(2), jnp.ones(3))\n\n def test_deepcopy_unspecified_parent(self):\n parent_parameter = inspect.signature(DummyModule).parameters['parent']\n unspecified_parent = parent_parameter.default\n\n self.assertIs(unspecified_parent, copy.copy(unspecified_parent))\n\n self.assertIs(unspecified_parent, copy.deepcopy(unspecified_parent))\n\n def test_type_hints(self):\n class Network(nn.Module):\n layers: int\n\n type_hints = get_type_hints(Network)\n self.assertEqual(type_hints['layers'], int)\n\n def test_incorrect_property(self):\n class Foo(nn.Module):\n @property\n def prop(self):\n return self.non_existent\n\n def __call__(self):\n return self.prop\n\n foo = Foo()\n with self.assertRaisesRegex(\n errors.DescriptorAttributeError, 'Trying to access a property that'\n ):\n foo.apply({})\n\n def test_custom_descriptor(self):\n class Descriptor:\n def __get__(self, obj, objtype=None):\n return 10\n\n class Foo(nn.Module):\n prop = Descriptor()\n\n def __call__(self):\n return self.prop\n\n foo = Foo()\n res = foo.apply({})\n self.assertEqual(res, 10)\n\n def test_custom_descriptor_error(self):\n class Descriptor:\n def __get__(self, obj, objtype=None):\n return obj.non_existent\n\n class Foo(nn.Module):\n prop = Descriptor()\n\n def __call__(self):\n return self.prop\n\n foo = Foo()\n with self.assertRaisesRegex(\n errors.DescriptorAttributeError, 'Trying to access a property that'\n ):\n foo.apply({})\n\n def test_nested_external_modules(self):\n class Baz(nn.Module):\n a: int\n\n def setup(self):\n self.b = self.param('b', lambda k: 2)\n\n def __call__(self, x):\n return x + self.a * self.b\n\n class Bar(nn.Module):\n baz: Baz\n\n def __call__(self, x):\n return self.baz(x)\n\n class Foo(nn.Module):\n def setup(self):\n self.bar = Bar(baz=Baz(a=1))\n\n def __call__(self, x):\n return self.bar.baz(x)\n\n module = Foo()\n y, variables = module.init_with_output(jax.random.key(0), 1)\n self.assertEqual(y, 3)\n\n def test_getattribute_triggers_setup(self):\n class B(nn.Module):\n def setup(self):\n self.p1 = self.param('p1', lambda k: jnp.ones((2,)))\n\n def fn1(self, x):\n return self.p1 + x\n\n class A(nn.Module):\n b: nn.Module\n\n def __call__(self, x):\n return self.b.fn1(x)\n\n a = A(b=B())\n k = random.key(0)\n x = jnp.zeros((2,))\n vs = nn.init(lambda a, x: a(x), a)(k, x)\n y = nn.apply(lambda a, x: a.b.fn1(x), a)(vs, x)\n np.testing.assert_array_equal(y, jnp.ones((2,)))\n\n def test_nested_sequential_in_call(self):\n class Foo(nn.Module):\n def setup(self):\n self.seq = nn.Sequential([nn.Dense(10) for i in range(10)])\n\n def __call__(self, x):\n # try calling only the first layer\n return self.seq.layers[0](x)\n\n module = Foo()\n variables = module.init(jax.random.key(0), jnp.ones((1, 10)))\n\n def test_setup_called_bounded_submodules(self):\n module = nn.Sequential(\n [\n nn.Sequential(\n [\n nn.Dense(2),\n nn.relu,\n nn.Dense(2),\n ]\n ),\n nn.relu,\n nn.Dense(2),\n ]\n )\n x = jnp.ones((1, 3))\n variables = module.init(jax.random.key(0), x)\n bound_module = module.bind(variables)\n\n self.assertIsNotNone(bound_module.layers[0].layers[0].scope)\n self.assertIsNotNone(bound_module.layers[0].layers[2].scope)\n self.assertIsNotNone(bound_module.layers[2].scope)\n\n def test_call_bounded_toplevel_mutable(self):\n class Bar(nn.Module):\n a: int\n\n def setup(self):\n self.b = self.param('b', lambda k: 1)\n\n def __call__(self, x):\n return x + self.a * self.b\n\n class Foo(nn.Module):\n bars: Sequence[Bar]\n\n def __call__(self, x):\n for bar in self.bars:\n x = bar(x)\n return x\n\n module = Foo(bars=[])\n module.bars = [Bar(a=1)]\n\n variables = module.init(jax.random.key(0), jnp.ones(()))\n bound_module = module.bind(variables)\n\n bar1 = bound_module.bars[0]\n self.assertIsNotNone(bar1.scope)\n\n def test_nested_init(self):\n class Baz(nn.Module):\n a: int\n\n def setup(self):\n self.b = self.param('b', lambda k: jnp.ones(()))\n\n def __call__(self, x):\n return x + self.a * self.b\n\n class Bar(nn.Module):\n baz: Baz\n\n def setup(self):\n a = 1\n\n def __call__(self, x):\n return self.baz(x)\n\n class Foo(nn.Module):\n def setup(self):\n self.bar: Bar = Bar(baz=Baz(a=1))\n\n def __call__(self, x):\n # y = self.bar(x)\n y, bar_vars = self.bar.init_with_output(jax.random.key(0), x)\n return y, bar_vars\n\n # create foo\n module = Foo()\n\n # run foo\n (y, bar_vars), variables = module.init_with_output(\n jax.random.key(0), jnp.ones(())\n )\n\n self.assertIn('params', bar_vars)\n\n def test_nested_shared(self):\n class Shared(nn.Module):\n @nn.compact\n def __call__(self, x):\n return nn.Dense(1)(x)\n\n class Unshared(nn.Module):\n shared: nn.Module\n\n def __call__(self, x):\n return self.shared(x)\n\n class Super(nn.Module):\n a: nn.Module\n b: nn.Module\n\n def run_a(self, x):\n return self.a(x)\n\n def run_b(self, x):\n return self.b(x)\n\n def __call__(self, x):\n return self.a(x) + self.b(x)\n\n sh = Shared()\n a = Unshared(shared=sh)\n b = Unshared(shared=sh)\n module = Super(a=a, b=b)\n\n rng = jax.random.key(0)\n params = module.init(rng, jnp.ones(1))['params']\n\n module.apply({'params': params}, jnp.ones(1)) # works as expected\n module.apply(\n {'params': params}, jnp.ones(1), method='run_a'\n ) # works as expected\n module.apply(\n {'params': params}, jnp.ones(1), method='run_b'\n ) # ScopeParamNotFoundError: Could not find parameter named \"kernel\" in scope \"/b/shared/Dense_0\"\n\n def test_repr(self):\n class Base1(nn.Module):\n a: int\n\n class Base2(nn.Module):\n b: str\n\n class Foo(Base2, Base1):\n c: float\n\n module = Foo(a=1, b='ok', c=3.0)\n str_rep = repr(module)\n\n self.assertIn('a = 1', str_rep)\n self.assertIn(\"b = 'ok'\", str_rep)\n self.assertIn('c = 3.0', str_rep)\n\n def test_repr_should_not_cause_setup(self):\n class MLP(nn.Module):\n @nn.compact\n def __call__(self, x):\n x = nn.Dense(1)(x)\n return repr(self)\n\n class Foo(nn.Module):\n a: float\n b: MLP\n\n scope = Scope({})\n module = Foo(parent=scope, a=1, b=MLP(parent=scope))\n str_rep = repr(module)\n self.assertIn('a = 1', str_rep)\n\n self.assertEqual(module._state.setup_called, nn.module.SetupState.NEW)\n # repr() on a module should not cause inadvertent setup of submodules\n # i.e. module.b._state.setup_called should remain nn.module.SetupState.NEW\n # and not nn.module.SetupState.DONE\n self.assertEqual(module.b._state.setup_called, nn.module.SetupState.NEW)\n\n def test_kw_only(self):\n def create_kw_layers():\n class BaseLayer(nn.Module, kw_only=True):\n base_multiplier: int | None = -1\n\n class ChildLayer(BaseLayer):\n child_multiplier: int # Don't want to have to set a default argument!\n\n def __call__(self, x):\n return x * self.child_multiplier * self.base_multiplier\n\n return BaseLayer, ChildLayer\n\n if tuple(sys.version_info)[:3] < (3, 10, 0):\n with self.assertRaisesRegex(TypeError, 'not available before Py 3.10'):\n BaseLayer, ChildLayer = create_kw_layers()\n else:\n BaseLayer, ChildLayer = create_kw_layers()\n with self.assertRaisesRegex(TypeError, 'positional argument'):\n _ = BaseLayer(2)\n # Like in Python dataclass, `kw_only` is not inherited, so ChildLayer can\n # take positional arg. It takes BaseLayer's default kwargs though.\n np.testing.assert_equal(ChildLayer(8)(np.ones(10)), -8 * np.ones(10))\n\n def test_positional_cannot_be_kw_only(self):\n class Foo(nn.Module):\n a: int\n\n Foo(1) # ok\n Foo(a=1) # ok\n with self.assertRaisesRegex(\n TypeError, r'takes 2 positional arguments but 3 were'\n ):\n Foo(1, None)\n Foo(a=1, parent=None) # type: ignore[call-arg]\n\n def test_module_path_empty(self):\n rngkey = jax.random.key(0)\n scope = Scope({}, {'params': rngkey}, mutable=['params'])\n m1 = DummyModule(parent=scope)\n\n self.assertEqual(m1.path, ())\n\n scope = Scope({}, {'params': rngkey}, mutable=['params'], path=['root'])\n m2 = DummyModule(parent=scope)\n\n self.assertEqual(m2.path, ('root',))\n\n m3 = DummyModule(parent=scope.rewound())\n\n self.assertEqual(m3.path, ('root',))\n\n def test_module_path_unbound_module_error(self):\n m1 = DummyModule()\n with self.assertRaisesRegex(ValueError, 'unbound module'):\n _ = m1.path\n\n def test_module_path_in_nested_module(self):\n module_paths = []\n debug_paths = []\n\n class A(nn.Module):\n def setup(self):\n self.b1 = B()\n self.b2 = B()\n self.c1 = C()\n\n module_paths.append(self.path)\n debug_paths.append(self.scope.debug_path)\n\n def __call__(self, x):\n return self.b1(x) + self.b2(x) + self.c1(x)\n\n class B(nn.Module):\n def setup(self):\n self.c1 = nn.remat(nn.remat(C))()\n self.c2 = C()\n\n module_paths.append(self.path)\n debug_paths.append(self.scope.debug_path)\n\n def __call__(self, x):\n return self.c1(x) + self.c2(x)\n\n class C(nn.Module):\n def setup(self):\n super().setup()\n if self.scope.__class__.__name__ != 'TestScope':\n module_paths.append(self.path)\n debug_paths.append(self.scope.debug_path)\n\n @nn.compact\n def __call__(self, x):\n return x\n\n a = A()\n k = random.key(0)\n x = random.uniform(random.key(42), (2,))\n _ = a.init(k, x)\n expected_module_paths = [\n (),\n ('b1',),\n ('b1', 'c1'),\n ('b1', 'c2'),\n ('b2',),\n ('b2', 'c1'),\n ('b2', 'c2'),\n ('c1',),\n ]\n expected_debug_paths = [\n (),\n ('b1',),\n ('b1', 'remat(remat(c1))'),\n ('b1', 'c2'),\n ('b2',),\n ('b2', 'remat(remat(c1))'),\n ('b2', 'c2'),\n ('c1',),\n ]\n\n self.assertEqual(module_paths, expected_module_paths)\n self.assertEqual(debug_paths, expected_debug_paths)\n\n def test_intercept_methods(self):\n mod = IdentityModule(parent=None)\n x = jnp.ones([])\n call_count = []\n\n def add_one_interceptor(f, args, kwargs, context):\n call_count.append(None)\n self.assertLen(dataclasses.fields(context), 3)\n self.assertIs(context.module, mod)\n self.assertEqual(context.method_name, '__call__')\n self.assertEqual(context.orig_method(3), 3)\n self.assertEqual(args, (x,))\n self.assertEmpty(kwargs)\n y = f(*args, **kwargs)\n return y + 1\n\n y1 = mod(x)\n with nn.intercept_methods(add_one_interceptor):\n y2 = mod(x)\n y3 = mod(x)\n\n self.assertLen(call_count, 1)\n self.assertEqual(y1, 1)\n self.assertEqual(y2, 2)\n self.assertEqual(y3, 1)\n\n def test_intercept_methods_compact(self):\n class CompactModule(nn.Module):\n @compact\n def __call__(self, x):\n return nn.Dense(2)(x)\n\n mod = CompactModule()\n x = jnp.ones(shape=(1, 3))\n variables = mod.init(jax.random.key(0), x)\n call_modules = []\n\n def log_interceptor(f, args, kwargs, context):\n call_modules.append(context.module)\n self.assertLen(dataclasses.fields(context), 3)\n self.assertEqual(context.method_name, '__call__')\n self.assertEqual(args, (x,))\n self.assertEmpty(kwargs)\n return f(*args, **kwargs)\n\n with nn.intercept_methods(log_interceptor):\n _ = mod.apply(variables, x)\n\n self.assertLen(call_modules, 2)\n self.assertIsInstance(call_modules[0], CompactModule)\n self.assertIsInstance(call_modules[1], nn.Dense)\n\n def test_intercept_methods_setup(self):\n class SetupModule(nn.Module):\n def setup(self):\n self.layer = nn.Dense(2)\n\n def __call__(self, x):\n return self.layer(x)\n\n mod = SetupModule()\n x = jnp.ones(shape=(1, 3))\n variables = mod.init(jax.random.key(0), x)\n call_modules = []\n log = []\n\n def log_interceptor(f, args, kwargs, context):\n call_modules.append(context.module)\n log.append((context.method_name, args, kwargs))\n return f(*args, **kwargs)\n\n with nn.intercept_methods(log_interceptor):\n _ = mod.apply(variables, x)\n\n self.assertLen(call_modules, 3)\n self.assertIsInstance(call_modules[0], SetupModule)\n self.assertIsInstance(call_modules[1], SetupModule)\n self.assertIsInstance(call_modules[2], nn.Dense)\n self.assertEqual(\n log, [('setup', (), {}), ('__call__', (x,), {}), ('__call__', (x,), {})]\n )\n\n def test_intercept_methods_calling_underlying_optional(self):\n def do_nothing_interceptor(f, args, kwargs, context):\n del f, context\n self.assertEmpty(args)\n self.assertEmpty(kwargs)\n\n m = RaisesModule()\n with nn.intercept_methods(do_nothing_interceptor):\n m()\n\n with self.assertRaises(AssertionError):\n m()\n\n with nn.intercept_methods(do_nothing_interceptor):\n m()\n\n def test_intercept_methods_run_in_lifo_order(self):\n def op_interceptor(op):\n def _interceptor(f, args, kwargs, context):\n del context\n y = f(*args, **kwargs)\n return op(y)\n\n return _interceptor\n\n mod = IdentityModule(parent=None)\n x = 7\n with (\n nn.intercept_methods(op_interceptor(lambda a: a + 1)),\n nn.intercept_methods(op_interceptor(lambda a: a**2)),\n ):\n y = mod(x)\n\n self.assertEqual(y, (x**2) + 1)\n\n with (\n nn.intercept_methods(op_interceptor(lambda a: a**2)),\n nn.intercept_methods(op_interceptor(lambda a: a + 1)),\n ):\n y = mod(x)\n\n self.assertEqual(y, (x + 1) ** 2)\n\n def test_intercept_methods_subclasses(self):\n class Foo(IdentityModule):\n def __call__(self, x): # pylint: disable=useless-parent-delegation\n return super().__call__(x)\n\n class Bar(Foo):\n def __call__(self, x): # pylint: disable=useless-parent-delegation\n return super().__call__(x)\n\n bar = Bar(parent=None)\n x = jnp.ones([])\n called = []\n\n def record_interceptor(f, args, kwargs, context):\n called.append(None)\n self.assertIs(context.module, bar)\n self.assertEqual(context.method_name, '__call__')\n self.assertEqual(args, (x,))\n self.assertEmpty(kwargs)\n return f(*args, **kwargs)\n\n with nn.intercept_methods(record_interceptor):\n bar(x)\n\n # Bar.__call__, Foo.__call__ and IdenityModule.__call__\n self.assertLen(called, 3)\n\n def test_intercept_methods_nested_module(self):\n class Foo(nn.Module):\n def __call__(self, x):\n return x\n\n class Bar(nn.Module):\n sub: nn.Module\n\n def __call__(self, x):\n return self.sub(x)\n\n foo = Foo()\n bar = Bar(sub=foo)\n x = jnp.ones([])\n called = []\n\n def record_interceptor(f, args, kwargs, context):\n called.append(context.module)\n self.assertEqual(context.method_name, '__call__')\n self.assertEqual(args, (x,))\n self.assertEmpty(kwargs)\n return f(*args, **kwargs)\n\n with nn.intercept_methods(record_interceptor):\n bar(x)\n\n # bar.__call__ and foo.__call__\n self.assertLen(called, 2)\n self.assertIs(called[0], bar)\n self.assertIs(called[1], foo)\n\n def test_cloudpickle_class(self):\n import cloudpickle\n\n class MyModule(nn.Module):\n pass\n\n a = MyModule()\n\n UnpickledMyModule = cloudpickle.loads(cloudpickle.dumps(MyModule))\n b = UnpickledMyModule()\n\n def test_cloudpickle_module(self):\n from cloudpickle import cloudpickle_fast\n\n class NNModuleWithProperty(nn.Module):\n a: int\n b: str\n\n @property\n def my_property(self):\n return self.b * self.a\n\n m = NNModuleWithProperty(a=2, b='ok')\n\n with TemporaryDirectory() as tmpdir:\n filename = f'{tmpdir}/module.pkl'\n with open(filename, 'wb') as f:\n cloudpickle_fast.dump(m, f)\n\n with open(filename, 'rb') as f:\n obj_loaded = cloudpickle_fast.load(f)\n\n self.assertEqual(obj_loaded.a, 2)\n self.assertEqual(obj_loaded.b, 'ok')\n self.assertEqual(obj_loaded.my_property, 'okok')\n\n def test_module_paths(self):\n class Bar(nn.Module):\n @nn.compact\n def __call__(self, x):\n x = nn.Dense(3)(x)\n x = nn.Dense(4)(x)\n return x\n\n class Foo(nn.Module):\n @nn.compact\n def __call__(self, x):\n x = Bar()(x)\n x = nn.Dense(5)(x)\n return x\n\n x = jnp.ones((1, 2))\n m = Foo()\n module_paths = m.module_paths(random.key(0), x)\n\n # assert all module are unbounded\n for module in module_paths.values():\n self.assertIsNone(module.scope)\n\n # test paths\n self.assertIn('', module_paths)\n self.assertEqual(type(module_paths['']), Foo)\n self.assertIn('Dense_0', module_paths)\n self.assertEqual(type(module_paths['Dense_0']), nn.Dense)\n self.assertIn('Bar_0', module_paths)\n self.assertEqual(type(module_paths['Bar_0']), Bar)\n self.assertIn('Bar_0/Dense_0', module_paths)\n self.assertEqual(type(module_paths['Bar_0/Dense_0']), nn.Dense)\n self.assertIn('Bar_0/Dense_1', module_paths)\n self.assertEqual(type(module_paths['Bar_0/Dense_1']), nn.Dense)\n\n def test_init_apply_default_rng(self):\n class SubModel(nn.Module):\n @nn.compact\n def __call__(self, x, apply_dropout):\n x = nn.Dense(8)(x)\n x = nn.Dropout(0.8)(x, deterministic=not apply_dropout)\n p = self.param(\n 'parameter', lambda key, shape: jax.random.normal(key, shape), x.shape\n )\n noise = jax.random.normal(self.make_rng('noise'), x.shape)\n return x * p + noise\n\n class Model(nn.Module):\n @nn.compact\n def __call__(self, x, apply_dropout):\n x = nn.Dense(16)(x)\n x = SubModel()(x, apply_dropout)\n x = nn.Dropout(0.5)(x, deterministic=not apply_dropout)\n v = self.variable(\n 'var_collection',\n 'variable',\n lambda shape: jax.random.normal(self.make_rng('var_rng'), shape),\n x.shape,\n )\n noise = jax.random.normal(self.make_rng('noise'), x.shape)\n return x * v.value + noise\n\n key0, key1, key2 = jax.random.split(jax.random.key(0), 3)\n x = jax.random.normal(key0, (10, 4))\n model = Model()\n\n # test init equality\n default_variables = model.init({'params': key1}, x, apply_dropout=False)\n rngs = {'params': key1, 'var_rng': key1, 'noise': key1}\n explicit_variables = model.init(rngs, x, apply_dropout=False)\n self.assertTrue(\n jax.tree_util.tree_all(\n jax.tree_util.tree_map(\n lambda v1, v2: (v1 == v2).all(),\n default_variables,\n explicit_variables,\n )\n )\n )\n\n # test init inequality\n for rng_name in ('params', 'var_rng'):\n rngs[rng_name] = key2\n explicit_variables = model.init(rngs, x, apply_dropout=False)\n self.assertFalse(\n jax.tree_util.tree_all(\n jax.tree_util.tree_map(\n lambda v1, v2: (v1 == v2).all(),\n default_variables,\n explicit_variables,\n )\n )\n )\n rngs[rng_name] = key1\n\n # test apply equality\n default_out = model.apply(\n default_variables, x, apply_dropout=True, rngs={'params': key1}\n )\n rngs = {'dropout': key1, 'noise': key1}\n explicit_out = model.apply(\n default_variables, x, apply_dropout=True, rngs=rngs\n )\n np.testing.assert_allclose(default_out, explicit_out)\n\n # test apply inequality\n for rng_name in ('dropout', 'noise'):\n rngs[rng_name] = key2\n explicit_out = model.apply(\n default_variables, x, apply_dropout=True, rngs=rngs\n )\n with self.assertRaises(AssertionError):\n np.testing.assert_allclose(default_out, explicit_out, atol=1e-1)\n rngs[rng_name] = key1\n\n def test_default_make_rng(self):\n class SubModel(nn.Module):\n @nn.compact\n def __call__(self, x):\n noise = jax.random.normal(self.make_rng(), x.shape)\n return x + noise\n\n class Model(nn.Module):\n @nn.compact\n def __call__(self, x):\n x = SubModel()(x)\n noise = jax.random.normal(self.make_rng(), x.shape)\n return x + noise\n\n key0, key1 = jax.random.split(jax.random.key(0), 2)\n x = jax.random.normal(key0, (10, 4))\n default_out = Model().apply({}, x, rngs={'params': key1})\n\n class SubModel(nn.Module):\n @nn.compact\n def __call__(self, x):\n noise = jax.random.normal(self.make_rng('params'), x.shape)\n return x + noise\n\n class Model(nn.Module):\n @nn.compact\n def __call__(self, x):\n x = SubModel()(x)\n noise = jax.random.normal(self.make_rng('params'), x.shape)\n return x + noise\n\n explicit_out = Model().apply({}, x, rngs={'params': key1})\n np.testing.assert_allclose(default_out, explicit_out)\n\n def test_default_rng_error(self):\n class Model(nn.Module):\n @nn.compact\n def __call__(self, x):\n return nn.Dense(2)(x)\n\n model = Model()\n with self.assertRaisesRegex(\n errors.InvalidRngError, 'Dense_0 needs PRNG for \"params\"'\n ):\n model.init({'other_rng_stream': jax.random.key(0)}, jnp.ones((1, 3)))\n\n class Model(nn.Module):\n @nn.compact\n def __call__(self, x):\n return x + jax.random.normal(self.make_rng(), x.shape)\n\n model = Model()\n with self.assertRaisesRegex(\n errors.InvalidRngError, 'None needs PRNG for \"params\"'\n ):\n model.init({'other_rng_stream': jax.random.key(0)}, jnp.ones((1, 3)))\n\n def test_compact_name_scope(self):\n class Foo(nn.Module):\n @nn.compact_name_scope\n def up(self, x):\n return nn.Dense(3)(x)\n\n @nn.compact_name_scope\n def down(self, x):\n return nn.Dense(3)(x)\n\n @nn.compact\n def __call__(self, x):\n return self.up(x) + self.down(x) + nn.Dense(3)(x)\n\n m = Foo()\n x = jnp.ones((1, 2))\n\n self.assertEqual(set(m._compact_name_scope_methods), {'up', 'down'})\n\n variables = m.init(random.key(0), x)\n params = variables['params']\n\n self.assertIn('Dense_0', params)\n self.assertIn('down', params)\n self.assertIn('up', params)\n self.assertIn('Dense_0', params['down'])\n self.assertIn('Dense_0', params['up'])\n\n y = m.apply(variables, x)\n y_up = m.apply(variables, x, method='up')\n y_down = m.apply(variables, x, method='down')\n\n assert y.shape == (1, 3)\n assert y_up.shape == (1, 3)\n assert y_down.shape == (1, 3)\n\n def test_compact_name_scope_outside_compact(self):\n class Foo(nn.Module):\n @nn.compact_name_scope\n def up(self, x):\n return nn.Dense(3)(x)\n\n @nn.compact_name_scope\n def down(self, x):\n return nn.Dense(3)(x)\n\n def __call__(self, x):\n return self.up(x) + self.down(x)\n\n m = Foo()\n x = jnp.ones((1, 2))\n\n self.assertEqual(set(m._compact_name_scope_methods), {'up', 'down'})\n\n variables = m.init(random.key(0), x)\n params = variables['params']\n\n self.assertIn('down', params)\n self.assertIn('up', params)\n self.assertIn('Dense_0', params['down'])\n self.assertIn('Dense_0', params['up'])\n\n y = m.apply(variables, x)\n y_up = m.apply(variables, x, method='up')\n y_down = m.apply(variables, x, method='down')\n\n assert y.shape == (1, 3)\n assert y_up.shape == (1, 3)\n assert y_down.shape == (1, 3)\n\n\nclass LeakTests(absltest.TestCase):\n def test_tracer_leaks(self):\n model = nn.Sequential([nn.Dense(50)])\n\n @jax.jit\n @functools.partial(jax.vmap, in_axes=(0, None))\n def sample_from_prior(rng, inp):\n params = model.init(rng, np.zeros((10, 50)))\n out = model.apply(params, inp)\n del params\n return out\n\n # disable manual gc.collect call in jax leak checker\n # so that we can test tracer leaks in ref-cycles. This is a\n # reasonable proxy for transiently leaked memory during\n # eager execution.\n with patch.object(gc, 'collect', return_value=0):\n with jax.checking_leaks():\n for i in range(5):\n rngs = jax.random.split(jax.random.key(23), 100)\n out = sample_from_prior(rngs, np.ones((4, 50)))\n out.block_until_ready()\n del out, rngs\n\n\nclass RelaxedNamingTests(absltest.TestCase):\n def test_relaxed_adoption(self):\n class Foo(nn.Module):\n @nn.compact\n def __call__(self, x):\n p = self.param('p', nn.initializers.zeros, x.shape)\n return x + p\n\n class Bar(nn.Module):\n sub: nn.Module\n\n def __call__(self, x):\n return self.sub(x)\n\n with set_config('flax_preserve_adopted_names', True):\n foo = Foo(name='foo')\n bar = Bar(sub=foo)\n k = random.key(0)\n x = jnp.zeros((1,))\n vs = bar.init(k, x)\n self.assertTrue('foo' in vs['params'], 'relaxed naming failure')\n y = bar.apply(vs, x)\n\n with set_config('flax_preserve_adopted_names', False):\n foo = Foo(name='foo')\n bar = Bar(sub=foo)\n k = random.key(0)\n x = jnp.zeros((1,))\n vs = bar.init(k, x)\n self.assertTrue('sub' in vs['params'], 'old policy naming failure')\n y = bar.apply(vs, x)\n\n def test_class_optional_adoption_name_preservation(self):\n class Foo(nn.Module):\n @nn.compact\n def __call__(self, x):\n p = self.param('p', nn.initializers.zeros, x.shape)\n return x + p\n\n class Bar1(nn.Module):\n sub: nn.Module\n preserve_adopted_names = True\n\n def __call__(self, x):\n return self.sub(x)\n\n class Bar2(nn.Module):\n sub: nn.Module\n preserve_adopted_names = False\n\n def __call__(self, x):\n return self.sub(x)\n\n with set_config('flax_preserve_adopted_names', False):\n foo = Foo(name='foo')\n bar = Bar1(sub=foo)\n k = random.key(0)\n x = jnp.zeros((1,))\n vs = bar.init(k, x)\n self.assertTrue('foo' in vs['params'], 'adoption naming failure')\n y = bar.apply(vs, x)\n\n with set_config('flax_preserve_adopted_names', True):\n foo = Foo(name='foo')\n bar = Bar2(sub=foo)\n k = random.key(0)\n x = jnp.zeros((1,))\n vs = bar.init(k, x)\n self.assertTrue('sub' in vs['params'], 'adoption naming failure')\n y = bar.apply(vs, x)\n\n def test_nested_class_optional_adoption_name_preservation(self):\n class Foo(nn.Module):\n @nn.compact\n def __call__(self, x):\n p = self.param('p', nn.initializers.zeros, x.shape)\n return x + p\n\n class Bar(nn.Module):\n sub: nn.Module\n preserve_adopted_names = True\n\n def __call__(self, x):\n return self.sub(x)\n\n class Baz(nn.Module):\n sub: nn.Module\n preserve_adopted_names = True\n\n def __call__(self, x):\n return self.sub(x)\n\n with set_config('flax_preserve_adopted_names', False):\n foo = Foo(name='foo')\n bar = Bar(sub=foo, name='bar')\n baz = Baz(sub=bar)\n k = random.key(0)\n x = jnp.zeros((1,))\n vs = baz.init(k, x)\n self.assertTrue('bar' in vs['params'], 'adoption naming failure')\n self.assertTrue('foo' in vs['params']['bar'], 'adoption naming failure')\n y = baz.apply(vs, x)\n\n def test_relaxed_adoption_still_conflict_checks(self):\n class Foo(nn.Module):\n @nn.compact\n def __call__(self, x):\n p = self.param('p', nn.initializers.zeros, x.shape)\n return x + p\n\n class Bar(nn.Module):\n sub1: nn.Module\n sub2: nn.Module\n\n def __call__(self, x):\n return self.sub(x)\n\n with set_config('flax_preserve_adopted_names', True):\n foo1 = Foo(name='foo')\n foo2 = Foo(name='foo')\n bar = Bar(sub1=foo1, sub2=foo2)\n k = random.key(0)\n x = jnp.zeros((1,))\n with self.assertRaises(errors.NameInUseError):\n vs = bar.init(k, x)\n\n def test_relaxed_adoption_unnamed_adoptee(self):\n class Foo(nn.Module):\n @nn.compact\n def __call__(self, x):\n p = self.param('p', nn.initializers.zeros, x.shape)\n return x + p\n\n class Bar(nn.Module):\n sub: nn.Module\n\n def __call__(self, x):\n return self.sub(x)\n\n with set_config('flax_preserve_adopted_names', True):\n foo = Foo(name=None)\n bar = Bar(sub=foo)\n k = random.key(0)\n x = jnp.zeros((1,))\n vs = bar.init(k, x)\n self.assertTrue('sub' in vs['params'], 'relaxed naming failure')\n y = bar.apply(vs, x)\n\n with set_config('flax_preserve_adopted_names', False):\n foo = Foo(name='foo')\n bar = Bar(sub=foo)\n k = random.key(0)\n x = jnp.zeros((1,))\n vs = bar.init(k, x)\n self.assertTrue('sub' in vs['params'], 'old policy naming failure')\n y = bar.apply(vs, x)\n\n def test_relaxed_python_conflict(self):\n class Foo(nn.Module):\n dummy = 0\n\n @nn.compact\n def __call__(self, x):\n p = self.param('dummy', nn.initializers.zeros, x.shape)\n return x + p\n\n foo = Foo(name='foo')\n k = random.key(0)\n x = jnp.zeros((1,))\n vs = foo.init(k, x)\n\n def test_relaxed_intercollection_conflict(self):\n class Foo(nn.Module):\n @nn.compact\n def __call__(self, x):\n v1 = self.variable('col1', 'v', lambda x: jnp.zeros(x), x.shape)\n v2 = self.variable('col2', 'v', lambda x: jnp.zeros(x), x.shape)\n return x + v1.value + v2.value\n\n foo = Foo(name='foo')\n k = random.key(0)\n x = jnp.zeros((1,))\n vs = foo.init(k, x)\n\n def test_relaxed_intercollection_conflict_set(self):\n class Foo(nn.Module):\n @nn.compact\n def __call__(self, x):\n v1 = self.variable('col1', 'v', lambda x: jnp.zeros(x), x.shape)\n v2 = self.variable('col2', 'v', lambda x: jnp.zeros(x), x.shape)\n v3 = self.variable('col1', 'v', lambda x: jnp.zeros(x), x.shape)\n return x + v1.value + v2.value + v3.value\n\n foo = Foo(name='foo')\n k = random.key(0)\n x = jnp.zeros((1,))\n with self.assertRaises(errors.NameInUseError):\n vs = foo.init(k, x)\n\n def test_internal_deep_clone(self):\n class Child(nn.Module):\n @nn.compact\n def __call__(self, x):\n w = self.param('w', nn.initializers.zeros, (5, x.shape[1]))\n return x @ w\n\n class Parent(nn.Module):\n num_layers: int\n child_template: Child\n\n @nn.compact\n def __call__(self, x):\n for i in range(self.num_layers):\n x = self.child_template.clone(\n parent=self, _deep_clone=True, name=None\n )(x)\n return x\n\n model = Parent(num_layers=2, child_template=Child())\n x = jnp.ones((32, 5))\n variables = model.init(jax.random.key(0), x)\n output = model.apply(variables, x)\n self.assertTrue(\n variables['params']['Child_0']['w'].shape\n == variables['params']['Child_1']['w'].shape\n )\n\n def test_copy_method(self):\n class Parent(nn.Module):\n @nn.compact\n def __call__(self, x):\n child = nn.Dense(\n 2,\n )\n x = child(x)\n x = child.copy()(x)\n return x\n\n model = Parent()\n x = jnp.ones((2, 2))\n variables = model.init(jax.random.key(0), x)\n output = model.apply(variables, x)\n self.assertTrue(\n variables['params']['Dense_0']['kernel'].shape\n == variables['params']['Dense_1']['kernel'].shape\n )\n\n def test_copy_from_template(self):\n class Child(nn.Module):\n @nn.compact\n def __call__(self, x):\n w = self.param('w', nn.initializers.zeros, (5, x.shape[1]))\n return x @ w\n\n class Parent(nn.Module):\n num_layers: int\n child_template: Child\n\n @nn.compact\n def __call__(self, x):\n for i in range(self.num_layers):\n x = self.child_template.copy()(x)\n for i in range(self.num_layers):\n x = self.child_template.copy(name=f'next_layer_{i}')(x)\n return x\n\n model = Parent(num_layers=2, child_template=Child())\n x = jnp.ones((32, 5))\n variables = model.init(jax.random.key(0), x)\n output = model.apply(variables, x)\n self.assertTrue(\n variables['params']['Child_0']['w'].shape\n == variables['params']['Child_1']['w'].shape\n )\n self.assertIn('Child_0', variables['params'])\n self.assertIn('Child_1', variables['params'])\n self.assertIn('next_layer_0', variables['params'])\n self.assertIn('next_layer_1', variables['params'])\n self.assertNotIn('child_template', variables['params'])\n\n def test_nonstring_keys_in_dict_on_module(self):\n class MyEnum(str, enum.Enum):\n a = 'a'\n b = 'b'\n\n class MyModule(nn.Module):\n config: dict[MyEnum, int]\n\n def __call__(self, inputs):\n return inputs\n\n module = MyModule(config={MyEnum.a: 1, MyEnum.b: 2})\n variables = module.init(jax.random.key(0), jnp.zeros([0]))\n\n\nclass FrozenDictTests(absltest.TestCase):\n def test_frozendict_flag(self):\n with set_config('flax_return_frozendict', True):\n x = jnp.zeros((2, 3))\n layer = nn.Dense(5)\n params = layer.init(random.key(0), x)\n self.assertTrue(isinstance(params, FrozenDict))\n\n with set_config('flax_return_frozendict', False):\n x = jnp.zeros((2, 3))\n layer = nn.Dense(5)\n params = layer.init(random.key(0), x)\n self.assertTrue(isinstance(params, dict))\n\n\nclass ShareScopeTest(absltest.TestCase):\n def test_basic(self):\n class DenseLoRA(nn.Module):\n inner: nn.Dense\n rank: int\n\n def setup(self):\n nn.share_scope(self, self.inner)\n\n @nn.compact\n def __call__(self, x: jax.Array):\n din, dout = x.shape[-1], self.inner.features\n A = self.param('A', nn.zeros_init(), (din, self.rank))\n B = self.param('B', nn.zeros_init(), (self.rank, dout))\n return self.inner(x) + x @ A @ B\n\n dense_lora = DenseLoRA(nn.Dense(10), rank=2)\n\n params = dense_lora.init(random.key(0), jnp.ones((1, 5)))['params']\n\n self.assertIn('kernel', params)\n self.assertIn('bias', params)\n self.assertIn('A', params)\n self.assertIn('B', params)\n\n def test_child_scope(self):\n class DenseLoRA(nn.Module):\n rank: int\n\n def setup(self):\n self.child = nn.Dense(10)\n nn.share_scope(self, self.child)\n\n @nn.compact\n def __call__(self, x: jax.Array):\n din, dout = x.shape[-1], self.child.features\n A = self.param('A', nn.zeros_init(), (din, self.rank))\n B = self.param('B', nn.zeros_init(), (self.rank, dout))\n return self.child(x) + x @ A @ B\n\n dense_lora = DenseLoRA(rank=2)\n\n params = dense_lora.init(random.key(0), jnp.ones((1, 5)))['params']\n\n self.assertIn('kernel', params)\n self.assertIn('bias', params)\n self.assertIn('A', params)\n self.assertIn('B', params)\n\n def test_in_compact(self):\n class DenseLoRA(nn.Module):\n rank: int\n\n def setup(self):\n self.child = nn.Dense(10)\n nn.share_scope(self, self.child)\n\n @nn.compact\n def __call__(self, x: jax.Array):\n din, dout = x.shape[-1], self.child.features\n A = self.param('A', nn.zeros_init(), (din, self.rank))\n B = self.param('B', nn.zeros_init(), (self.rank, dout))\n return self.child(x) + x @ A @ B\n\n class Model(nn.Module):\n @nn.compact\n def __call__(self, x: jax.Array):\n return DenseLoRA(rank=2)(x)\n\n model = Model()\n\n params = model.init(random.key(0), jnp.ones((1, 5)))['params']\n\n self.assertIn('kernel', params['DenseLoRA_0'])\n self.assertIn('bias', params['DenseLoRA_0'])\n self.assertIn('A', params['DenseLoRA_0'])\n self.assertIn('B', params['DenseLoRA_0'])\n\n def test_adopt_child_name(self):\n class DenseLoRA(nn.Module):\n inner: nn.Dense\n rank: int\n\n def setup(self):\n nn.share_scope(self, self.inner)\n\n @nn.compact\n def __call__(self, x: jax.Array):\n din, dout = x.shape[-1], self.inner.features\n A = self.param('A', nn.zeros_init(), (din, self.rank))\n B = self.param('B', nn.zeros_init(), (self.rank, dout))\n return self.inner(x) + x @ A @ B\n\n class Model(nn.Module):\n @nn.compact\n def __call__(self, x: jax.Array):\n return DenseLoRA(nn.Dense(10), rank=2)(x)\n\n model = Model()\n\n params = model.init(random.key(0), jnp.ones((1, 5)))['params']\n\n self.assertIn('kernel', params['Dense_0'])\n self.assertIn('bias', params['Dense_0'])\n self.assertIn('A', params['Dense_0'])\n self.assertIn('B', params['Dense_0'])\n\n def test_other_scope_is_none(self):\n class DenseLoRA(nn.Module):\n inner: nn.Dense\n rank: int\n\n def setup(self):\n nn.share_scope(self, self.inner)\n\n @nn.compact\n def __call__(self, x: jax.Array):\n din, dout = x.shape[-1], self.inner.features\n A = self.param('A', nn.zeros_init(), (din, self.rank))\n B = self.param('B', nn.zeros_init(), (self.rank, dout))\n return self.inner(x) + x @ A @ B\n\n class Model(nn.Module):\n def setup(self):\n # here Dense doesn't have a scope yet\n self.dense_lora = DenseLoRA(nn.Dense(10), rank=2)\n\n @nn.compact\n def __call__(self, x: jax.Array):\n return self.dense_lora(x)\n\n model = Model()\n\n params = model.init(random.key(0), jnp.ones((1, 5)))['params']\n\n self.assertIn('kernel', params['dense_lora'])\n self.assertIn('bias', params['dense_lora'])\n self.assertIn('A', params['dense_lora'])\n self.assertIn('B', params['dense_lora'])\n\n def test_external_grandchild_scope_correct(self):\n class GrandChild(nn.Module):\n @nn.compact\n def __call__(self):\n return nn.Dense(50)(jnp.zeros(10))\n\n class Child(nn.Module):\n child: GrandChild\n\n @nn.compact\n def __call__(self, *args: Any, **kwargs: Any) -> Any:\n return self.child(*args, **kwargs)\n\n class Parent(nn.Module):\n main_child: Child\n\n def setup(self):\n nn.share_scope(self, self.main_child)\n\n @nn.compact\n def __call__(self, *args: Any, **kwargs: Any) -> Any:\n nn.Dense(10)(jnp.zeros(10))\n r = self.main_child(*args, **kwargs)\n return r\n\n params = Parent(Child(GrandChild())).init(jax.random.key(0))\n self.assertNotIn('main_child', params['params'])\n self.assertIn('child', params['params'])\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/linen/linen_recurrent_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Recurrent tests.\"\"\"\n\n\nimport jax\nimport jax.numpy as jnp\nimport numpy as np\nfrom absl.testing import absltest\n\nfrom flax import linen as nn\nfrom flax.linen.recurrent import flip_sequences\n\n# Parse absl flags test_srcdir and test_tmpdir.\njax.config.parse_flags_with_absl()\n\n\nclass RNNTest(absltest.TestCase):\n def test_rnn_basic_forward(self):\n batch_size = 10\n seq_len = 40\n channels_in = 5\n channels_out = 15\n\n rnn = nn.RNN(nn.LSTMCell(channels_out), return_carry=True)\n\n xs = jnp.ones((batch_size, seq_len, channels_in))\n variables = rnn.init(jax.random.key(0), xs)\n ys: jnp.ndarray\n carry, ys = rnn.apply(variables, xs)\n\n self.assertEqual(ys.shape, (batch_size, seq_len, channels_out))\n\n for layer_params in variables['params']['cell'].values():\n if 'bias' in layer_params:\n self.assertEqual(layer_params['bias'].shape, (channels_out,))\n self.assertIn(\n layer_params['kernel'].shape[0], [channels_in, channels_out]\n )\n self.assertEqual(layer_params['kernel'].shape[1], channels_out)\n\n def test_rnn_multiple_batch_dims(self):\n batch_dims = (10, 11)\n seq_len = 40\n channels_in = 5\n channels_out = 15\n\n rnn = nn.RNN(nn.LSTMCell(channels_out), return_carry=True)\n\n xs = jnp.ones((*batch_dims, seq_len, channels_in))\n variables = rnn.init(jax.random.key(0), xs)\n ys: jnp.ndarray\n carry, ys = rnn.apply(variables, xs)\n\n self.assertEqual(ys.shape, (*batch_dims, seq_len, channels_out))\n\n for layer_params in variables['params']['cell'].values():\n if 'bias' in layer_params:\n self.assertEqual(layer_params['bias'].shape, (channels_out,))\n self.assertIn(\n layer_params['kernel'].shape[0], [channels_in, channels_out]\n )\n self.assertEqual(layer_params['kernel'].shape[1], channels_out)\n\n def test_rnn_unroll(self):\n batch_size = 10\n seq_len = 40\n channels_in = 5\n channels_out = 15\n\n rnn = nn.RNN(nn.LSTMCell(channels_out), unroll=10, return_carry=True)\n\n xs = jnp.ones((batch_size, seq_len, channels_in))\n variables = rnn.init(jax.random.key(0), xs)\n ys: jnp.ndarray\n carry, ys = rnn.apply(variables, xs)\n\n self.assertEqual(ys.shape, (batch_size, seq_len, channels_out))\n\n for layer_params in variables['params']['cell'].values():\n if 'bias' in layer_params:\n self.assertEqual(layer_params['bias'].shape, (channels_out,))\n self.assertIn(\n layer_params['kernel'].shape[0], [channels_in, channels_out]\n )\n self.assertEqual(layer_params['kernel'].shape[1], channels_out)\n\n def test_rnn_time_major(self):\n seq_len = 40\n batch_size = 10\n channels_in = 5\n channels_out = 15\n\n rnn = nn.RNN(nn.LSTMCell(channels_out), time_major=True, return_carry=True)\n\n xs = jnp.ones((seq_len, batch_size, channels_in))\n variables = rnn.init(jax.random.key(0), xs)\n\n ys: jnp.ndarray\n carry, ys = rnn.apply(variables, xs)\n\n # carry state should not be zeros after apply\n for leaf in jax.tree_util.tree_leaves(carry):\n assert not np.allclose(leaf, jnp.zeros_like(leaf))\n self.assertEqual(leaf.shape, (batch_size, channels_out))\n\n self.assertEqual(ys.shape, (seq_len, batch_size, channels_out))\n\n for layer_params in variables['params']['cell'].values():\n if 'bias' in layer_params:\n self.assertEqual(layer_params['bias'].shape, (channels_out,))\n self.assertIn(\n layer_params['kernel'].shape[0], [channels_in, channels_out]\n )\n self.assertEqual(layer_params['kernel'].shape[1], channels_out)\n\n def test_rnn_with_spatial_dimensions(self):\n batch_size = 10\n seq_len = 40\n kernel_size = (3, 3)\n image_size = (32, 32)\n channels_in = 5\n channels_out = 15\n\n rnn = nn.RNN(\n nn.ConvLSTMCell(channels_out, kernel_size),\n )\n\n xs = jnp.ones((batch_size, seq_len, *image_size, channels_in))\n variables = rnn.init(jax.random.key(0), xs)\n\n ys: jnp.ndarray\n carry, ys = rnn.apply(variables, xs, return_carry=True)\n\n # carry state should not be zeros after apply\n for leaf in jax.tree_util.tree_leaves(carry):\n assert not np.allclose(leaf, jnp.zeros_like(leaf))\n self.assertEqual(leaf.shape[:-1], (batch_size, *image_size))\n self.assertIn(leaf.shape[-1], [channels_in, channels_out])\n\n self.assertEqual(ys.shape, (batch_size, seq_len, *image_size, channels_out))\n\n for layer_params in variables['params']['cell'].values():\n if 'bias' in layer_params:\n self.assertEqual(layer_params['bias'].shape, (channels_out * 4,))\n self.assertIn(\n layer_params['kernel'].shape[2],\n [channels_in, channels_out, channels_out * 4],\n )\n self.assertEqual(layer_params['kernel'].shape[3], channels_out * 4)\n\n def test_numerical_equivalence(self):\n batch_size = 3\n seq_len = 4\n channels_in = 5\n channels_out = 6\n\n rnn = nn.RNN(nn.LSTMCell(channels_out), return_carry=True)\n\n xs = jnp.ones((batch_size, seq_len, channels_in))\n ys: jnp.ndarray\n (carry, ys), variables = rnn.init_with_output(jax.random.key(0), xs)\n\n cell_carry = rnn.cell.initialize_carry(jax.random.key(0), xs[:, 0].shape)\n cell_params = variables['params']['cell']\n\n for i in range(seq_len):\n cell_carry, y = rnn.cell.apply(\n {'params': cell_params}, cell_carry, xs[:, i, :]\n )\n np.testing.assert_allclose(y, ys[:, i, :], rtol=1e-5)\n\n np.testing.assert_allclose(cell_carry, carry, rtol=1e-5)\n\n def test_numerical_equivalence_with_mask(self):\n batch_size = 3\n seq_len = 4\n channels_in = 5\n channels_out = 6\n\n key = jax.random.key(0)\n seq_lengths = jax.random.randint(\n key, (batch_size,), minval=1, maxval=seq_len + 1\n )\n\n rnn = nn.RNN(nn.LSTMCell(channels_out), return_carry=True)\n\n xs = jnp.ones((batch_size, seq_len, channels_in))\n ys: jnp.ndarray\n (carry, ys), variables = rnn.init_with_output(\n jax.random.key(0), xs, seq_lengths=seq_lengths\n )\n\n cell_carry = rnn.cell.initialize_carry(jax.random.key(0), xs[:, 0].shape)\n cell_params = variables['params']['cell']\n carries = []\n\n for i in range(seq_len):\n cell_carry, y = rnn.cell.apply(\n {'params': cell_params}, cell_carry, xs[:, i, :]\n )\n np.testing.assert_allclose(y, ys[:, i, :], rtol=1e-5)\n carries.append(cell_carry)\n\n for batch_idx, length in enumerate(seq_lengths):\n t = int(length) - 1\n for carries_t_, carry_ in zip(carries[t], carry):\n np.testing.assert_allclose(\n carries_t_[batch_idx], carry_[batch_idx], rtol=1e-5\n )\n\n def test_numerical_equivalence_single_batch(self):\n batch_size = 3\n seq_len = 4\n channels_in = 5\n channels_out = 6\n\n rnn = nn.RNN(nn.LSTMCell(channels_out), return_carry=True)\n\n xs = jnp.ones((batch_size, seq_len, channels_in))\n ys: jnp.ndarray\n (carry, ys), variables = rnn.init_with_output(jax.random.key(0), xs)\n\n cell_params = variables['params']['cell']\n\n for batch_idx in range(batch_size):\n cell_carry = rnn.cell.initialize_carry(jax.random.key(0), xs[:1, 0].shape)\n\n for i in range(seq_len):\n cell_carry, y = rnn.cell.apply(\n {'params': cell_params}, cell_carry, xs[batch_idx, i, :][None]\n )\n np.testing.assert_allclose(y[0], ys[batch_idx, i, :], rtol=1e-6)\n\n carry_i = jax.tree_util.tree_map(\n lambda x: x[batch_idx : batch_idx + 1], carry\n )\n np.testing.assert_allclose(cell_carry, carry_i, rtol=1e-6)\n\n def test_numerical_equivalence_single_batch_nn_scan(self):\n batch_size = 3\n seq_len = 4\n channels_in = 5\n channels_out = 6\n\n cell: nn.LSTMCell = nn.LSTMCell(channels_out)\n rnn: nn.LSTMCell = nn.scan(\n nn.LSTMCell,\n in_axes=1,\n out_axes=1,\n variable_broadcast='params',\n split_rngs={'params': False},\n )(channels_out)\n\n xs = jnp.ones((batch_size, seq_len, channels_in))\n carry = rnn.initialize_carry(jax.random.key(0), xs[:, 0].shape)\n ys: jnp.ndarray\n (carry, ys), variables = rnn.init_with_output(jax.random.key(0), carry, xs)\n\n cell_params = variables['params']\n\n for batch_idx in range(batch_size):\n cell_carry = cell.initialize_carry(jax.random.key(0), xs[:1, 0].shape)\n\n for i in range(seq_len):\n cell_carry, y = cell.apply(\n {'params': cell_params},\n cell_carry,\n xs[batch_idx : batch_idx + 1, i, :],\n )\n np.testing.assert_allclose(y[0], ys[batch_idx, i, :], rtol=1e-5)\n\n carry_i = jax.tree_util.tree_map(\n lambda x: x[batch_idx : batch_idx + 1], carry\n )\n np.testing.assert_allclose(cell_carry, carry_i, rtol=1e-5)\n\n def test_numerical_equivalence_single_batch_jax_scan(self):\n batch_size = 3\n seq_len = 4\n channels_in = 5\n channels_out = 6\n\n xs = jax.random.uniform(\n jax.random.key(0), (batch_size, seq_len, channels_in)\n )\n cell: nn.LSTMCell = nn.LSTMCell(channels_out)\n carry = cell.initialize_carry(jax.random.key(0), xs[:, 0].shape)\n variables = cell.init(jax.random.key(0), carry, xs[:, 0])\n cell_params = variables['params']\n\n def scan_fn(carry, x):\n return cell.apply({'params': cell_params}, carry, x)\n\n ys: jnp.ndarray\n carry, ys = jax.lax.scan(scan_fn, carry, xs.swapaxes(0, 1))\n ys = ys.swapaxes(0, 1)\n\n cell_carry = cell.initialize_carry(jax.random.key(0), xs[:, 0].shape)\n\n for i in range(seq_len):\n cell_carry, y = cell.apply(\n {'params': cell_params}, cell_carry, xs[:, i, :]\n )\n np.testing.assert_allclose(y, ys[:, i, :], rtol=1e-4)\n\n np.testing.assert_allclose(cell_carry, carry, rtol=1e-4)\n\n def test_reverse(self):\n batch_size = 3\n seq_len = 4\n channels_in = 5\n channels_out = 6\n\n rnn = nn.RNN(nn.LSTMCell(channels_out), return_carry=True, reverse=True)\n\n xs = jnp.ones((batch_size, seq_len, channels_in))\n ys: jnp.ndarray\n (carry, ys), variables = rnn.init_with_output(jax.random.key(0), xs)\n\n cell_params = variables['params']['cell']\n\n for batch_idx in range(batch_size):\n cell_carry = rnn.cell.initialize_carry(jax.random.key(0), xs[:1, 0].shape)\n\n for i in range(seq_len):\n cell_carry, y = rnn.cell.apply(\n {'params': cell_params},\n cell_carry,\n xs[batch_idx, seq_len - i - 1, :][None],\n )\n np.testing.assert_allclose(y[0], ys[batch_idx, i, :], rtol=1e-5)\n\n np.testing.assert_allclose(\n cell_carry,\n jax.tree_util.tree_map(lambda x: x[batch_idx : batch_idx + 1], carry),\n rtol=1e-5,\n )\n\n def test_reverse_but_keep_order(self):\n batch_size = 3\n seq_len = 4\n channels_in = 5\n channels_out = 6\n\n rnn = nn.RNN(\n nn.LSTMCell(channels_out),\n return_carry=True,\n reverse=True,\n keep_order=True,\n )\n\n xs = jnp.ones((batch_size, seq_len, channels_in))\n ys: jnp.ndarray\n (carry, ys), variables = rnn.init_with_output(jax.random.key(0), xs)\n\n cell_params = variables['params']['cell']\n\n for batch_idx in range(batch_size):\n cell_carry = rnn.cell.initialize_carry(jax.random.key(0), xs[:1, 0].shape)\n\n for i in range(seq_len):\n cell_carry, y = rnn.cell.apply(\n {'params': cell_params},\n cell_carry,\n xs[batch_idx, seq_len - i - 1, :][None],\n )\n np.testing.assert_allclose(\n y[0], ys[batch_idx, seq_len - i - 1, :], rtol=1e-5\n )\n\n np.testing.assert_allclose(\n cell_carry,\n jax.tree_util.tree_map(lambda x: x[batch_idx : batch_idx + 1], carry),\n rtol=1e-5,\n )\n\n def test_flip_sequence(self):\n x = jnp.arange(2 * 5).reshape((2, 5))\n seq_lengths = jnp.array([4, 2])\n\n flipped = flip_sequences(x, seq_lengths, num_batch_dims=1, time_major=False)\n\n self.assertEqual(flipped.shape, (2, 5))\n np.testing.assert_allclose(flipped[0, :4], [3, 2, 1, 0])\n np.testing.assert_allclose(flipped[1, :2], [6, 5])\n\n def test_flip_sequence_more_feature_dims(self):\n x = jnp.arange(2 * 5 * 3).reshape((2, 5, 3))\n seq_lengths = jnp.array([4, 2])\n\n flipped = flip_sequences(x, seq_lengths, num_batch_dims=1, time_major=False)\n\n self.assertEqual(flipped.shape, (2, 5, 3))\n np.testing.assert_allclose(flipped[0, :4], x[0, :4][::-1])\n np.testing.assert_allclose(flipped[1, :2], x[1, :2][::-1])\n\n def test_flip_sequence_time_major(self):\n x = jnp.arange(2 * 5).reshape((5, 2))\n seq_lengths = jnp.array([4, 2])\n\n flipped = flip_sequences(x, seq_lengths, num_batch_dims=1, time_major=True)\n\n self.assertEqual(flipped.shape, (5, 2))\n np.testing.assert_allclose(flipped[:4, 0], x[:4, 0][::-1])\n np.testing.assert_allclose(flipped[:2, 1], x[:2, 1][::-1])\n\n def test_flip_sequence_time_major_more_feature_dims(self):\n x = jnp.arange(2 * 5 * 3).reshape((5, 2, 3))\n seq_lengths = jnp.array([4, 2])\n\n flipped = flip_sequences(x, seq_lengths, num_batch_dims=1, time_major=True)\n\n self.assertEqual(flipped.shape, (5, 2, 3))\n np.testing.assert_allclose(flipped[:4, 0], x[:4, 0][::-1])\n np.testing.assert_allclose(flipped[:2, 1], x[:2, 1][::-1])\n\n def test_basic_seq_lengths(self):\n x = jnp.ones((2, 10, 6))\n lstm = nn.RNN(nn.LSTMCell(265))\n variables = lstm.init(jax.random.key(0), x)\n y = lstm.apply(variables, x, seq_lengths=jnp.array([5, 5]))\n\n\nclass BidirectionalTest(absltest.TestCase):\n def test_bidirectional(self):\n batch_size = 3\n seq_len = 4\n channels_in = 5\n channels_out = 6\n\n bdirectional = nn.Bidirectional(\n nn.RNN(nn.LSTMCell(channels_out)), nn.RNN(nn.LSTMCell(channels_out))\n )\n\n xs = jnp.ones((batch_size, seq_len, channels_in))\n ys: jnp.ndarray\n ys, variables = bdirectional.init_with_output(jax.random.key(0), xs)\n\n self.assertEqual(ys.shape, (batch_size, seq_len, channels_out * 2))\n\n def test_shared_cell(self):\n batch_size = 3\n seq_len = 4\n channels_in = 5\n channels_out = 6\n\n cell = nn.LSTMCell(channels_out)\n bdirectional = nn.Bidirectional(nn.RNN(cell), nn.RNN(cell))\n\n xs = jnp.ones((batch_size, seq_len, channels_in))\n ys: jnp.ndarray\n ys, variables = bdirectional.init_with_output(jax.random.key(0), xs)\n\n self.assertEqual(ys.shape, (batch_size, seq_len, channels_out * 2))\n\n def test_custom_merge_fn(self):\n batch_size = 3\n seq_len = 4\n channels_in = 5\n channels_out = 6\n\n bdirectional = nn.Bidirectional(\n nn.RNN(nn.LSTMCell(channels_out)),\n nn.RNN(nn.LSTMCell(channels_out)),\n merge_fn=lambda x, y: x + y,\n )\n\n xs = jnp.ones((batch_size, seq_len, channels_in))\n ys: jnp.ndarray\n ys, variables = bdirectional.init_with_output(jax.random.key(0), xs)\n\n self.assertEqual(ys.shape, (batch_size, seq_len, channels_out))\n\n def test_return_carry(self):\n batch_size = 3\n seq_len = 4\n channels_in = 5\n channels_out = 6\n\n bdirectional = nn.Bidirectional(\n nn.RNN(nn.LSTMCell(channels_out)),\n nn.RNN(nn.LSTMCell(channels_out)),\n return_carry=True,\n )\n\n xs = jnp.ones((batch_size, seq_len, channels_in))\n ys: jnp.ndarray\n (carry, ys), variables = bdirectional.init_with_output(\n jax.random.key(0), xs\n )\n carry_forward, carry_backward = carry\n\n self.assertEqual(ys.shape, (batch_size, seq_len, channels_out * 2))\n self.assertEqual(\n jax.tree_util.tree_map(jnp.shape, carry_forward),\n ((batch_size, channels_out), (batch_size, channels_out)),\n )\n self.assertEqual(\n jax.tree_util.tree_map(jnp.shape, carry_backward),\n ((batch_size, channels_out), (batch_size, channels_out)),\n )\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/linen/linen_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for flax.linen.\"\"\"\n\nimport copy\nimport functools\nfrom typing import Any\n\nfrom absl.testing import absltest, parameterized\nfrom flax import ids\nfrom flax import linen as nn\nfrom flax.linen import fp8_ops\nfrom flax.training import train_state\nimport jax\nfrom jax import random\nimport jax.numpy as jnp\nimport numpy as np\nimport optax\n\n# Parse absl flags test_srcdir and test_tmpdir.\njax.config.parse_flags_with_absl()\n\n\ndef check_eq(xs, ys):\n return jax.tree_util.tree_all(\n jax.tree_util.tree_map(np.testing.assert_allclose, xs, ys)\n )\n\n\nclass PoolTest(parameterized.TestCase):\n def test_pool_custom_reduce(self):\n x = jnp.full((1, 3, 3, 1), 2.0)\n mul_reduce = lambda x, y: x * y\n y = nn.pooling.pool(x, 1.0, mul_reduce, (2, 2), (1, 1), 'VALID')\n np.testing.assert_allclose(y, np.full((1, 2, 2, 1), 2.0**4))\n\n @parameterized.parameters(\n {'count_include_pad': True}, {'count_include_pad': False}\n )\n def test_avg_pool(self, count_include_pad):\n x = jnp.full((1, 3, 3, 1), 2.0)\n pool = lambda x: nn.avg_pool(x, (2, 2), count_include_pad=count_include_pad)\n y = pool(x)\n np.testing.assert_allclose(y, np.full((1, 2, 2, 1), 2.0))\n y_grad = jax.grad(lambda x: pool(x).sum())(x)\n expected_grad = jnp.array(\n [\n [0.25, 0.5, 0.25],\n [0.5, 1.0, 0.5],\n [0.25, 0.5, 0.25],\n ]\n ).reshape((1, 3, 3, 1))\n np.testing.assert_allclose(y_grad, expected_grad)\n\n @parameterized.parameters(\n {'count_include_pad': True}, {'count_include_pad': False}\n )\n def test_avg_pool_no_batch(self, count_include_pad):\n x = jnp.full((3, 3, 1), 2.0)\n pool = lambda x: nn.avg_pool(x, (2, 2), count_include_pad=count_include_pad)\n y = pool(x)\n np.testing.assert_allclose(y, np.full((2, 2, 1), 2.0))\n y_grad = jax.grad(lambda x: pool(x).sum())(x)\n expected_grad = jnp.array(\n [\n [0.25, 0.5, 0.25],\n [0.5, 1.0, 0.5],\n [0.25, 0.5, 0.25],\n ]\n ).reshape((3, 3, 1))\n np.testing.assert_allclose(y_grad, expected_grad)\n\n def test_max_pool(self):\n x = jnp.arange(9).reshape((1, 3, 3, 1)).astype(jnp.float32)\n pool = lambda x: nn.max_pool(x, (2, 2))\n expected_y = jnp.array(\n [\n [4.0, 5.0],\n [7.0, 8.0],\n ]\n ).reshape((1, 2, 2, 1))\n y = pool(x)\n np.testing.assert_allclose(y, expected_y)\n y_grad = jax.grad(lambda x: pool(x).sum())(x)\n expected_grad = jnp.array(\n [\n [0.0, 0.0, 0.0],\n [0.0, 1.0, 1.0],\n [0.0, 1.0, 1.0],\n ]\n ).reshape((1, 3, 3, 1))\n np.testing.assert_allclose(y_grad, expected_grad)\n\n @parameterized.parameters(\n {'count_include_pad': True}, {'count_include_pad': False}\n )\n def test_avg_pool_padding_same(self, count_include_pad):\n x = jnp.array([1.0, 2.0, 3.0, 4.0]).reshape((1, 2, 2, 1))\n pool = lambda x: nn.avg_pool(\n x, (2, 2), padding='SAME', count_include_pad=count_include_pad\n )\n y = pool(x)\n if count_include_pad:\n expected_y = jnp.array([10.0 / 4, 6.0 / 4, 7.0 / 4, 4.0 / 4]).reshape(\n (1, 2, 2, 1)\n )\n else:\n expected_y = jnp.array([10.0 / 4, 6.0 / 2, 7.0 / 2, 4.0 / 1]).reshape(\n (1, 2, 2, 1)\n )\n np.testing.assert_allclose(y, expected_y)\n\n def test_pooling_variable_batch_dims(self):\n x = jnp.zeros((1, 8, 32, 32, 3), dtype=jnp.float32)\n y = nn.max_pool(x, (2, 2), (2, 2))\n\n assert y.shape == (1, 8, 16, 16, 3)\n\n def test_pooling_no_batch_dims(self):\n x = jnp.zeros((32, 32, 3), dtype=jnp.float32)\n y = nn.max_pool(x, (2, 2), (2, 2))\n\n assert y.shape == (16, 16, 3)\n\n\nclass NormalizationTest(parameterized.TestCase):\n def test_layer_norm_mask(self):\n key = random.key(0)\n keys = random.split(key)\n x = random.normal(keys[0], (3, 4, 5))\n m = random.choice(keys[1], 2, x.shape).astype(bool)\n m = m.at[..., :2].set(True) # guarantee at least 2 elements\n x = jnp.where(m, x, jnp.nan)\n\n module = nn.LayerNorm()\n y, w = module.init_with_output(key, x, mask=m)\n\n z = y.mean(-1, where=m)\n np.testing.assert_allclose(z, 0, atol=1e-4)\n\n z = y.var(-1, where=m)\n np.testing.assert_allclose(z, 1, atol=1e-4)\n\n def test_rms_norm_mask(self):\n key = random.key(0)\n keys = random.split(key)\n x = random.normal(keys[0], (3, 4, 5))\n m = random.choice(keys[1], 2, x.shape).astype(bool)\n m = m.at[..., :1].set(True) # guarantee at least 1 element\n x = jnp.where(m, x, jnp.nan)\n\n module = nn.RMSNorm()\n y, w = module.init_with_output(key, x, mask=m)\n\n z = np.square(y).mean(-1, where=m)\n np.testing.assert_allclose(z, 1, atol=1e-4)\n\n def test_group_norm_mask(self):\n key = random.key(0)\n keys = random.split(key)\n x = random.normal(keys[0], (13, 3, 5, 7 * 11))\n m = random.choice(keys[1], 2, x.shape).astype(bool)\n m = m.at[..., :2].set(True) # guarantee at least 2 elements\n x = jnp.where(m, x, jnp.nan)\n\n module = nn.GroupNorm(7, use_bias=False, use_scale=False)\n y, w = module.init_with_output(key, x, mask=m)\n\n yr = y.reshape((13, 3, 5, 7, 11))\n mr = m.reshape((13, 3, 5, 7, 11))\n\n axes = list(range(1, x.ndim - 1)) + [-1]\n\n z = yr.mean(axes, where=mr)\n np.testing.assert_allclose(z, 0, atol=1e-4)\n\n z = yr.var(axes, where=mr)\n np.testing.assert_allclose(z, 1, atol=1e-4)\n\n @parameterized.parameters({'test_mask': True}, {'test_mask': False})\n def test_batch_norm(self, test_mask):\n rng = random.key(0)\n key1, key2, key3 = random.split(rng, 3)\n x = random.normal(key1, (4, 3, 2))\n if test_mask:\n m = random.randint(\n key2, (4, 3, 1), minval=0, maxval=2, dtype=jnp.int32\n ).astype(jnp.bool_)\n x = jnp.where(m, x, jnp.ones_like(x) * jnp.nan)\n else:\n m = None\n model_cls = nn.BatchNorm(momentum=0.9, use_running_average=False)\n y, initial_params = model_cls.init_with_output(key3, x, mask=m)\n\n mean = y.mean((0, 1), where=m)\n var = y.var((0, 1), where=m)\n np.testing.assert_allclose(mean, np.array([0.0, 0.0]), atol=1e-4)\n np.testing.assert_allclose(var, np.array([1.0, 1.0]), rtol=1e-4)\n _, vars_out = model_cls.apply(\n initial_params, x, mutable=['batch_stats'], mask=m\n )\n\n ema = vars_out['batch_stats']\n np.testing.assert_allclose(\n ema['mean'], 0.1 * x.mean((0, 1), keepdims=False, where=m), atol=1e-4\n )\n np.testing.assert_allclose(\n ema['var'],\n 0.9 + 0.1 * x.var((0, 1), keepdims=False, where=m),\n rtol=1e-4,\n )\n\n @parameterized.parameters({'test_mask': True}, {'test_mask': False})\n def test_batch_norm_complex(self, test_mask):\n rng = random.key(0)\n key1, key2, key3 = random.split(rng, 3)\n x = random.normal(key1, (4, 3, 2), dtype=jnp.complex64)\n if test_mask:\n m = random.randint(\n key2, (4, 3, 1), minval=0, maxval=2, dtype=jnp.int32\n ).astype(jnp.bool_)\n x = jnp.where(m, x, jnp.ones_like(x) * jnp.nan)\n else:\n m = None\n model_cls = nn.BatchNorm(\n momentum=0.9, use_running_average=False, dtype=jnp.complex64\n )\n y, initial_params = model_cls.init_with_output(key3, x, mask=m)\n\n mean = y.mean((0, 1), where=m)\n var = y.var((0, 1), where=m)\n np.testing.assert_allclose(mean, np.array([0.0, 0.0]), atol=1e-4)\n np.testing.assert_allclose(var, np.array([1.0, 1.0]), rtol=1e-4)\n self.assertEqual(mean.dtype, jnp.complex64)\n\n _, vars_out = model_cls.apply(\n initial_params, x, mutable=['batch_stats'], mask=m\n )\n\n ema = vars_out['batch_stats']\n np.testing.assert_allclose(\n ema['mean'], 0.1 * x.mean((0, 1), keepdims=False, where=m), atol=1e-4\n )\n np.testing.assert_allclose(\n ema['var'],\n 0.9 + 0.1 * x.var((0, 1), keepdims=False, where=m),\n rtol=1e-4,\n )\n\n @parameterized.parameters(\n {'reduction_axes': -1},\n {'reduction_axes': 1},\n {'reduction_axes': (1, 2)},\n {'reduction_axes': (0, 1, 2)},\n {'reduction_axes': -1, 'use_fast_variance': False},\n )\n def test_layer_norm(self, reduction_axes, use_fast_variance=True):\n rng = random.key(0)\n key1, key2 = random.split(rng)\n e = 1e-5\n x = random.normal(key1, (2, 3, 4))\n if not use_fast_variance:\n x += 1e6 # This blows up fast variance, but should work otherwise.\n model_cls = nn.LayerNorm(\n use_bias=False,\n use_scale=False,\n epsilon=e,\n reduction_axes=reduction_axes,\n use_fast_variance=use_fast_variance,\n )\n y, _ = model_cls.init_with_output(key2, x)\n self.assertEqual(x.dtype, y.dtype)\n self.assertEqual(x.shape, y.shape)\n y_one_liner = (\n x - x.mean(axis=reduction_axes, keepdims=True)\n ) * jax.lax.rsqrt(x.var(axis=reduction_axes, keepdims=True) + e)\n np.testing.assert_allclose(y_one_liner, y, atol=1e-3, rtol=1e-3)\n\n @parameterized.parameters(\n {'reduction_axes': -1}, {'reduction_axes': 1}, {'reduction_axes': (1, 2)}\n )\n def test_rms_norm(self, reduction_axes):\n rng = random.key(0)\n key1, key2 = random.split(rng)\n e = 1e-5\n x = random.normal(key1, (2, 3, 4))\n model_cls = nn.RMSNorm(\n use_scale=False, epsilon=e, reduction_axes=reduction_axes\n )\n y, _ = model_cls.init_with_output(key2, x)\n self.assertEqual(x.dtype, y.dtype)\n self.assertEqual(x.shape, y.shape)\n y_one_liner = x * jax.lax.rsqrt(\n jnp.mean(jax.lax.square(x), axis=reduction_axes, keepdims=True) + e\n )\n np.testing.assert_allclose(y_one_liner, y, atol=1e-4)\n\n def test_group_norm(self):\n rng = random.key(0)\n key1, key2 = random.split(rng)\n e = 1e-5\n x = random.normal(key1, (2, 5, 4, 4, 32))\n model_cls = nn.GroupNorm(\n num_groups=2, use_bias=False, use_scale=False, epsilon=e\n )\n\n y, _ = model_cls.init_with_output(key2, x)\n self.assertEqual(x.dtype, y.dtype)\n self.assertEqual(x.shape, y.shape)\n\n x_gr = x.reshape([2, 5, 4, 4, 2, 16])\n y_test = (\n x_gr - x_gr.mean(axis=[1, 2, 3, 5], keepdims=True)\n ) * jax.lax.rsqrt(x_gr.var(axis=[1, 2, 3, 5], keepdims=True) + e)\n y_test = y_test.reshape([2, 5, 4, 4, 32])\n\n np.testing.assert_allclose(y_test, y, atol=1e-4)\n\n def test_group_norm_unbatched(self):\n rng = random.key(0)\n key1, key2 = random.split(rng)\n e = 1e-5\n x = random.normal(key1, (2, 5, 4, 4, 32))\n model_cls = nn.GroupNorm(\n num_groups=2,\n use_bias=False,\n use_scale=False,\n epsilon=e,\n reduction_axes=(0, 1, 3, 4),\n )\n\n y, _ = model_cls.init_with_output(key2, x)\n self.assertEqual(x.dtype, y.dtype)\n self.assertEqual(x.shape, y.shape)\n\n x_gr = x.reshape([2, 5, 4, 4, 2, 16])\n y_test = (\n x_gr - x_gr.mean(axis=[0, 1, 3, 5], keepdims=True)\n ) * jax.lax.rsqrt(x_gr.var(axis=[0, 1, 3, 5], keepdims=True) + e)\n y_test = y_test.reshape([2, 5, 4, 4, 32])\n\n np.testing.assert_allclose(y_test, y, atol=1e-4)\n\n def test_group_norm_batched(self):\n rng = random.key(0)\n key1, key2 = random.split(rng)\n e = 1e-5\n x = random.normal(key1, (3, 4, 32))\n model_cls = nn.GroupNorm(\n num_groups=2,\n use_bias=False,\n use_scale=False,\n epsilon=e,\n reduction_axes=(-3, -2, -1),\n )\n\n y, _ = model_cls.init_with_output(key2, x)\n self.assertEqual(x.dtype, y.dtype)\n self.assertEqual(x.shape, y.shape)\n\n x_stacked = jnp.stack([x] * 5)\n y_stacked = model_cls.apply({}, x_stacked)\n\n np.testing.assert_allclose(y, y_stacked[0, ...], atol=1e-4)\n\n x_gr = x_stacked.reshape([5, 3, 4, 2, 16])\n y_test = (x_gr - x_gr.mean(axis=[1, 2, 4], keepdims=True)) * jax.lax.rsqrt(\n x_gr.var(axis=[1, 2, 4], keepdims=True) + e\n )\n y_test = y_test.reshape([5, 3, 4, 32])\n\n np.testing.assert_allclose(y_test, y_stacked, atol=1e-4)\n\n def test_group_norm_raises(self):\n rng = random.key(0)\n key1, key2 = random.split(rng)\n e = 1e-5\n x = random.normal(key1, (2, 5, 4, 4, 32))\n model_cls = nn.GroupNorm(\n num_groups=3, use_bias=False, use_scale=False, epsilon=e\n )\n\n with self.assertRaises(ValueError):\n model_cls.init_with_output(key2, x)\n\n def test_group_norm_raises_incorrect_reduction_axes(self):\n rng = random.key(0)\n key1, key2 = random.split(rng)\n e = 1e-5\n x = random.normal(key1, (2, 5, 4, 4, 32))\n model_cls = nn.GroupNorm(\n num_groups=3,\n use_bias=False,\n use_scale=False,\n epsilon=e,\n reduction_axes=(0, 1, 2, 3),\n )\n\n with self.assertRaises(ValueError):\n model_cls.init_with_output(key2, x)\n\n def test_batch_norm_multi_init(self):\n class Foo(nn.Module):\n @nn.compact\n def __call__(self, x):\n norm = nn.BatchNorm(\n name='norm',\n use_running_average=False,\n axis_name='batch',\n )\n x = norm(x)\n return x, norm(x)\n\n key = random.key(0)\n model = Foo()\n x = random.normal(random.key(1), (2, 4))\n (y1, y2), _ = model.init_with_output(key, x)\n np.testing.assert_allclose(y1, y2, rtol=0.005)\n\n @parameterized.parameters(\n {'feature_axes': -1},\n {'feature_axes': (1, 2)},\n {'feature_axes': (1, 2, 3)},\n {'feature_axes': -1, 'use_fast_variance': False},\n )\n def test_instance_norm(self, feature_axes, use_fast_variance=True):\n rng = random.key(0)\n key1, key2 = random.split(rng)\n e = 1e-5\n x = random.normal(key1, (2, 3, 4, 5))\n if not use_fast_variance:\n x += 1e4 # This blows up fast variance, but should work otherwise.\n model_cls = nn.InstanceNorm(\n use_bias=False,\n use_scale=False,\n epsilon=e,\n feature_axes=feature_axes,\n use_fast_variance=use_fast_variance,\n )\n y, _ = model_cls.init_with_output(key2, x)\n self.assertEqual(x.dtype, y.dtype)\n self.assertEqual(x.shape, y.shape)\n\n canonicalized_feature_axes = [\n i if i >= 0 else (x.ndim + i)\n for i in (\n feature_axes if isinstance(feature_axes, tuple) else (feature_axes,)\n )\n ]\n reduction_axes = [\n i for i in range(1, x.ndim) if i not in canonicalized_feature_axes\n ]\n y_one_liner = (\n x - x.mean(axis=reduction_axes, keepdims=True)\n ) * jax.lax.rsqrt(x.var(axis=reduction_axes, keepdims=True) + e)\n\n np.testing.assert_allclose(y_one_liner, y, atol=1e-6)\n\n @parameterized.parameters(\n {'feature_axes': 0},\n {'feature_axes': -4},\n {'feature_axes': (0, 3)},\n {'feature_axes': (2, -4)},\n )\n def test_instance_norm_raise_error(self, feature_axes):\n with self.assertRaisesRegex(\n ValueError,\n 'The channel axes cannot include the leading dimension '\n 'as this is assumed to be the batch axis.',\n ):\n x = jax.random.normal(jax.random.key(0), (2, 3, 4, 5))\n layer = nn.InstanceNorm(feature_axes=feature_axes)\n _ = layer.init(jax.random.key(1), x)\n\n @parameterized.parameters(\n {\n 'layer1': nn.LayerNorm(feature_axes=(1, 2)),\n 'layer2': nn.InstanceNorm(feature_axes=(1, 2)),\n },\n {\n 'layer1': nn.LayerNorm(reduction_axes=(1, 2), feature_axes=-1),\n 'layer2': nn.InstanceNorm(feature_axes=-1),\n },\n {\n 'layer1': nn.LayerNorm(\n reduction_axes=-2,\n feature_axes=(1, 3),\n bias_init=nn.initializers.uniform(),\n scale_init=nn.initializers.uniform(),\n ),\n 'layer2': nn.InstanceNorm(\n feature_axes=(1, -1),\n bias_init=nn.initializers.uniform(),\n scale_init=nn.initializers.uniform(),\n ),\n },\n {\n 'layer1': nn.LayerNorm(\n reduction_axes=(1, 2, 3),\n bias_init=nn.initializers.uniform(),\n scale_init=nn.initializers.uniform(),\n ),\n 'layer2': nn.GroupNorm(\n num_groups=1,\n bias_init=nn.initializers.uniform(),\n scale_init=nn.initializers.uniform()\n ),\n },\n {\n 'layer1': nn.InstanceNorm(\n bias_init=nn.initializers.uniform(),\n scale_init=nn.initializers.uniform(),\n ),\n 'layer2': nn.GroupNorm(\n num_groups=None,\n group_size=1,\n bias_init=nn.initializers.uniform(),\n scale_init=nn.initializers.uniform(),\n ),\n },\n )\n def test_normalization_equivalence(self, layer1, layer2):\n x = jax.random.normal(jax.random.key(0), (2, 3, 4, 5))\n layer1_variables = layer1.init(jax.random.key(1), x)\n layer2_variables = layer2.init(jax.random.key(1), x)\n self.assertTrue(\n jax.tree_util.tree_all(\n jax.tree_util.tree_map(\n lambda v1, v2: (v1 == v2).all(),\n layer1_variables,\n layer2_variables,\n )\n )\n )\n\n layer1_y = layer1.apply(layer1_variables, x)\n layer2_y = layer2.apply(layer2_variables, x)\n np.testing.assert_allclose(layer1_y, layer2_y, atol=1e-7)\n\n @parameterized.parameters(\n {\n 'model_index': 0,\n 'key_paths': {'Dense_1/kernel/u', 'Dense_1/kernel/sigma'},\n },\n {\n 'model_index': 1,\n 'key_paths': {'Conv_0/kernel/u', 'Conv_0/kernel/sigma'},\n },\n {\n 'model_index': 2,\n 'key_paths': {\n 'MultiHeadDotProductAttention_0/key/bias/u',\n 'MultiHeadDotProductAttention_0/key/kernel/u',\n 'MultiHeadDotProductAttention_0/out/kernel/u',\n 'MultiHeadDotProductAttention_0/query/bias/u',\n 'MultiHeadDotProductAttention_0/query/kernel/u',\n 'MultiHeadDotProductAttention_0/value/bias/u',\n 'MultiHeadDotProductAttention_0/value/kernel/u',\n 'MultiHeadDotProductAttention_0/key/bias/sigma',\n 'MultiHeadDotProductAttention_0/key/kernel/sigma',\n 'MultiHeadDotProductAttention_0/out/kernel/sigma',\n 'MultiHeadDotProductAttention_0/query/bias/sigma',\n 'MultiHeadDotProductAttention_0/query/kernel/sigma',\n 'MultiHeadDotProductAttention_0/value/bias/sigma',\n 'MultiHeadDotProductAttention_0/value/kernel/sigma',\n },\n },\n )\n def test_spectral_norm_train(self, model_index, key_paths):\n class FooDense(nn.Module):\n @nn.compact\n def __call__(self, x, train):\n x = nn.Dense(8)(x)\n x = nn.SpectralNorm(nn.Dense(6))(x, update_stats=train)\n x = nn.Dense(4)(x)\n return x\n\n class FooConv(nn.Module):\n @nn.compact\n def __call__(self, x, train):\n x = nn.Dense(9)(x)\n x = x.reshape((1, 3, 3))\n x = nn.SpectralNorm(nn.Conv(2, kernel_size=(2, 2)))(\n x, update_stats=train\n )\n x = x.reshape(1, -1)\n x = nn.Dense(4)(x)\n return x\n\n class FooAttention(nn.Module):\n @nn.compact\n def __call__(self, x, train):\n a = nn.Dense(4)(x)\n b = nn.Dense(4)(x)\n x = nn.SpectralNorm(nn.attention.MultiHeadDotProductAttention(4))(\n a, b, update_stats=train\n )\n x = nn.Dense(4)(x)\n return x\n\n key1, key2, key3 = random.split(random.PRNGKey(0), 3)\n x = random.normal(key1, (1, 4))\n y = random.normal(key2, (1, 4))\n\n model_cls = (FooDense, FooConv, FooAttention)[model_index]\n variables = model_cls().init(key3, x, train=False)\n params, batch_stats = variables['params'], variables['batch_stats']\n self.assertEqual(key_paths, batch_stats['SpectralNorm_0'].keys())\n\n class TrainState(train_state.TrainState):\n batch_stats: Any\n\n state = TrainState.create(\n apply_fn=model_cls().apply,\n params=params,\n batch_stats=batch_stats,\n tx=optax.adam(1e-3),\n )\n\n @jax.jit\n def train_step(state, batch):\n def loss_fn(params):\n logits, updates = state.apply_fn(\n {'params': params, 'batch_stats': state.batch_stats},\n x=batch['image'],\n train=True,\n mutable=['batch_stats'],\n )\n loss = jnp.mean(\n optax.l2_loss(predictions=logits, targets=batch['label'])\n )\n return loss, updates\n\n grad_fn = jax.value_and_grad(loss_fn, has_aux=True)\n (loss, updates), grads = grad_fn(state.params)\n state = state.apply_gradients(grads=grads)\n state = state.replace(batch_stats=updates['batch_stats'])\n return state, loss\n\n prev_loss = float('inf')\n for _ in range(10):\n state, loss = train_step(state, {'image': x, 'label': y})\n self.assertLess(loss, prev_loss)\n prev_loss = loss\n\n @parameterized.parameters(\n {'n_steps': 1, 'update_stats': True, 'result': 4.0},\n {'n_steps': 3, 'update_stats': True, 'result': 4.0},\n {'n_steps': 10, 'update_stats': True, 'result': 4.0},\n {'n_steps': 1, 'update_stats': False, 'result': 1.0},\n )\n def test_spectral_norm_sigma(self, n_steps, update_stats, result):\n class Foo(nn.Module):\n @nn.compact\n def __call__(self, x, train):\n x = nn.SpectralNorm(nn.Dense(8, use_bias=False), n_steps=n_steps)(\n x, update_stats=train\n )\n return x\n\n x = jnp.ones((1, 8))\n model_cls = Foo()\n variables = model_cls.init(random.PRNGKey(0), x, train=False)\n params, batch_stats = variables['params'], variables['batch_stats']\n params = jax.tree_util.tree_map(lambda x: 4 * jnp.eye(*x.shape), params)\n _, updates = model_cls.apply(\n {'params': params, 'batch_stats': batch_stats},\n x=x,\n train=update_stats,\n mutable=True,\n )\n np.testing.assert_allclose(\n updates['batch_stats']['SpectralNorm_0']['Dense_0/kernel/sigma'],\n result,\n atol=1e-3,\n )\n\n @parameterized.parameters(\n {'error_on_non_matrix': True}, {'error_on_non_matrix': False}\n )\n def test_spectral_norm_3d_tensor(self, error_on_non_matrix):\n class Foo(nn.Module):\n @nn.compact\n def __call__(self, x, train):\n x = nn.SpectralNorm(\n nn.DenseGeneral((3, 4), use_bias=False),\n error_on_non_matrix=error_on_non_matrix,\n )(x, update_stats=train)\n return x\n\n x = jnp.ones((1, 2))\n model_cls = Foo()\n\n if error_on_non_matrix:\n with self.assertRaisesRegex(\n ValueError, 'Input is 3D but error_on_non_matrix is True'\n ):\n _ = model_cls.init(random.PRNGKey(0), x, train=False)\n else:\n _ = model_cls.init(random.PRNGKey(0), x, train=False)\n\n @parameterized.parameters(\n {'feature_axes': -1, 'reduction_axes': 0, 'variable_filter': {'kernel'}},\n {'feature_axes': 0, 'reduction_axes': 1, 'variable_filter': {'kernel'}},\n {\n 'feature_axes': (0, 1),\n 'reduction_axes': (),\n 'variable_filter': {'kernel'},\n },\n {\n 'feature_axes': (),\n 'reduction_axes': (0, 1),\n 'variable_filter': {'kernel'},\n },\n {\n 'feature_axes': None,\n 'reduction_axes': (0, 1),\n 'variable_filter': {'kernel'},\n },\n {'feature_axes': 0, 'reduction_axes': (), 'variable_filter': {'bias'}},\n {'feature_axes': (), 'reduction_axes': -1, 'variable_filter': {'bias'}},\n )\n def test_manual_weight_norm(\n self, feature_axes, reduction_axes, variable_filter\n ):\n class Foo(nn.Module):\n @nn.compact\n def __call__(self, x):\n return nn.WeightNorm(\n nn.Dense(2, bias_init=nn.initializers.normal()),\n feature_axes=feature_axes,\n variable_filter=variable_filter,\n )(x)\n\n key1, key2 = jax.random.split(jax.random.key(1))\n x = jax.random.normal(key1, (1, 3))\n module = Foo()\n v = module.init(key2, x)\n v = jax.tree_util.tree_map(lambda x: x + 0.5, v)\n out = module.apply(v, x)\n\n kernel = v['params']['Dense_0']['kernel']\n if 'kernel' in variable_filter:\n kernel /= jnp.sqrt(jnp.sum(kernel**2, axis=reduction_axes, keepdims=True))\n kernel_scale = jnp.expand_dims(\n v['params']['WeightNorm_0']['Dense_0/kernel/scale'],\n axis=reduction_axes,\n )\n else:\n kernel_scale = 1\n bias = v['params']['Dense_0']['bias']\n if 'bias' in variable_filter:\n bias /= jnp.sqrt(jnp.sum(bias**2, axis=reduction_axes, keepdims=True))\n bias_scale = jnp.expand_dims(\n v['params']['WeightNorm_0']['Dense_0/bias/scale'], axis=reduction_axes\n )\n else:\n bias_scale = 1\n manual_out = jnp.dot(x, kernel_scale * kernel) + (\n bias_scale * bias\n ).reshape(1, -1)\n\n self.assertTrue(jnp.allclose(out, manual_out))\n\n @parameterized.parameters(\n {\n 'variable_filters': ({}, None, {'kernel', 'bias'}, {'Bar'}),\n 'key_paths': {\n 'Bar_0/Baz_0/Dense_0/kernel/scale',\n 'Bar_0/Baz_0/Dense_0/bias/scale',\n 'Bar_0/Dense_0/kernel/scale',\n 'Bar_0/Dense_0/bias/scale',\n 'Bar_0/Baz_1/Dense_0/kernel/scale',\n 'Bar_0/Baz_1/Dense_0/bias/scale',\n 'Bar_0/Dense_1/kernel/scale',\n 'Bar_0/Dense_1/bias/scale',\n },\n },\n {\n 'variable_filters': ({'kernel'},),\n 'key_paths': {\n 'Bar_0/Baz_0/Dense_0/kernel/scale',\n 'Bar_0/Dense_0/kernel/scale',\n 'Bar_0/Baz_1/Dense_0/kernel/scale',\n 'Bar_0/Dense_1/kernel/scale',\n },\n },\n {\n 'variable_filters': ({'Baz', 'kernel'},),\n 'key_paths': {\n 'Bar_0/Baz_0/Dense_0/kernel/scale',\n 'Bar_0/Baz_0/Dense_0/bias/scale',\n 'Bar_0/Dense_0/kernel/scale',\n 'Bar_0/Baz_1/Dense_0/kernel/scale',\n 'Bar_0/Baz_1/Dense_0/bias/scale',\n 'Bar_0/Dense_1/kernel/scale',\n },\n },\n )\n def test_weight_norm_variable_filter(self, variable_filters, key_paths):\n class Baz(nn.Module):\n @nn.compact\n def __call__(self, x):\n return nn.Dense(2)(x)\n\n class Bar(nn.Module):\n @nn.compact\n def __call__(self, x):\n x = Baz()(x)\n x = nn.Dense(3)(x)\n x = Baz()(x)\n x = nn.Dense(3)(x)\n return x\n\n for variable_filter in variable_filters:\n\n class Foo(nn.Module):\n @nn.compact\n def __call__(self, x):\n return nn.WeightNorm(Bar(), variable_filter=variable_filter)(x)\n\n v = Foo().init(jax.random.key(0), jnp.ones((1, 4)))\n self.assertEqual(key_paths, v['params']['WeightNorm_0'].keys())\n\n @parameterized.parameters(\n {'model_index': 0, 'key_paths': {'Dense_1/kernel/scale'}},\n {'model_index': 1, 'key_paths': {'Conv_0/kernel/scale'}},\n {\n 'model_index': 2,\n 'key_paths': {\n 'MultiHeadDotProductAttention_0/key/kernel/scale',\n 'MultiHeadDotProductAttention_0/out/kernel/scale',\n 'MultiHeadDotProductAttention_0/query/kernel/scale',\n 'MultiHeadDotProductAttention_0/value/kernel/scale',\n },\n },\n )\n def test_weight_norm_train(self, model_index, key_paths):\n class FooDense(nn.Module):\n @nn.compact\n def __call__(\n self,\n x,\n ):\n x = nn.Dense(8)(x)\n x = nn.WeightNorm(nn.Dense(6))(x)\n x = nn.Dense(4)(x)\n return x\n\n class FooConv(nn.Module):\n @nn.compact\n def __call__(\n self,\n x,\n ):\n x = nn.Dense(9)(x)\n x = x.reshape((1, 3, 3))\n x = nn.WeightNorm(nn.Conv(2, kernel_size=(2, 2)))(x)\n x = x.reshape(1, -1)\n x = nn.Dense(4)(x)\n return x\n\n class FooAttention(nn.Module):\n @nn.compact\n def __call__(self, x):\n a = nn.Dense(4)(x)\n b = nn.Dense(4)(x)\n x = nn.WeightNorm(nn.attention.MultiHeadDotProductAttention(4))(a, b)\n x = nn.Dense(4)(x)\n return x\n\n key1, key2, key3 = random.split(random.PRNGKey(0), 3)\n x = random.normal(key1, (1, 4))\n y = random.normal(key2, (1, 4))\n\n model_cls = (FooDense, FooConv, FooAttention)[model_index]\n params = model_cls().init(key3, x)['params']\n self.assertEqual(key_paths, params['WeightNorm_0'].keys())\n\n state = train_state.TrainState.create(\n apply_fn=model_cls().apply,\n params=params,\n tx=optax.adam(1e-3),\n )\n\n @jax.jit\n def train_step(state, batch):\n def loss_fn(params):\n logits = state.apply_fn(\n {'params': params},\n x=batch['image'],\n )\n loss = jnp.mean(\n optax.l2_loss(predictions=logits, targets=batch['label'])\n )\n return loss\n\n grad_fn = jax.value_and_grad(loss_fn)\n loss, grads = grad_fn(state.params)\n state = state.apply_gradients(grads=grads)\n return state, loss\n\n prev_loss = float('inf')\n for _ in range(10):\n state, loss = train_step(state, {'image': x, 'label': y})\n self.assertLess(loss, prev_loss)\n prev_loss = loss\n\n def test_weight_norm_compatibility_with_partitioning(self):\n replicated_module = nn.WeightNorm(nn.Dense(10))\n\n @jax.jit\n def _init_replicated(x):\n return replicated_module.init(jax.random.key(0), x)\n\n expected = _init_replicated(jnp.ones((10, 10)))\n\n sharded_module = nn.WeightNorm(\n nn.Dense(\n 10,\n kernel_init=nn.with_partitioning(\n nn.initializers.lecun_normal(), ('x',)\n ),\n )\n )\n\n @jax.jit\n def _init_sharded(x):\n return sharded_module.init(jax.random.key(0), x)\n\n got = _init_sharded(jnp.ones((10, 10)))\n\n def _strip_partitioning(x):\n if isinstance(x, nn.Partitioned):\n return x.value\n return x\n\n got = jax.tree.map(\n _strip_partitioning,\n got,\n is_leaf=lambda x: isinstance(x, nn.Partitioned),\n )\n\n expected_scale = expected['params'].pop('layer_instance/kernel/scale')\n # NOTE: `value` is from `nn.Partitioned.value`\n got_scale = got['params'].pop('layer_instance/kernel/value/scale')\n np.testing.assert_array_equal(got_scale, expected_scale)\n\n # Compares the rest of PyTree nodes\n jax.tree.map(lambda x, y: np.testing.assert_allclose(x, y), expected, got)\n\n\nclass StochasticTest(parameterized.TestCase):\n def test_dropout(self):\n rng = random.key(0)\n key1, key2 = random.split(rng)\n module = nn.Dropout(rate=0.5)\n y1 = module.apply(\n {}, jnp.ones((20, 20)), deterministic=False, rngs={'dropout': key1}\n )\n y2 = module.apply(\n {}, jnp.ones((20, 20)), deterministic=False, rngs={'dropout': key2}\n )\n self.assertFalse(np.all(y1 == y2))\n\n y1 = module.apply(\n {}, jnp.ones((20, 20)), deterministic=True, rngs={'dropout': key1}\n )\n y2 = module.apply(\n {}, jnp.ones((20, 20)), deterministic=True, rngs={'dropout': key2}\n )\n self.assertTrue(np.all(y1 == y2))\n\n def test_dropout_rate_stats(self):\n rootkey = random.key(0)\n for rate in np.arange(0.1, 1.0, 0.1):\n rootkey, subkey = random.split(rootkey)\n module = nn.Dropout(rate=rate)\n n_trials = 10\n nonzero_counts = 0\n for key in random.split(subkey, n_trials):\n y = module.apply(\n {}, jnp.ones((100, 100)), deterministic=False, rngs={'dropout': key}\n )\n nonzero_counts += np.sum(y > 0.0)\n all_counts = np.prod((100, 100, n_trials))\n frac = np.sum(nonzero_counts) / all_counts\n keep_rate = 1.0 - rate\n # just check within 4 sigma.\n delta = 4 * np.sqrt(rate * keep_rate) / np.sqrt(all_counts)\n self.assertTrue(keep_rate - delta < frac < keep_rate + delta)\n\n def test_dropout_rate_limits(self):\n rng = random.key(0)\n key1, key2, key3 = random.split(rng, 3)\n inputs = jnp.ones((20, 20))\n d0 = nn.Dropout(rate=0.0)\n y1 = d0.apply({}, inputs, deterministic=False, rngs={'dropout': key1})\n np.testing.assert_array_equal(y1, inputs)\n d1 = nn.Dropout(rate=1.0)\n y2 = d1.apply({}, inputs, deterministic=False, rngs={'dropout': key2})\n np.testing.assert_array_equal(y2, np.zeros_like(inputs))\n # ensure gradient of rate==1.0 case is non-NaN\n fn = lambda x, k: d1.apply({}, x, rngs={'dropout': k}, deterministic=False)\n res = jax.grad(lambda x, k: jnp.sum(fn(x, k)))(inputs, key3)\n self.assertFalse(np.isnan(res).any())\n\n @parameterized.parameters(\n {\n 'num_dims': 2,\n 'broadcast_dims': (1,),\n 'slice_fn': lambda out, i: out[i, :],\n 'summed_total': 2 * 10,\n },\n {\n 'num_dims': 2,\n 'broadcast_dims': (0,),\n 'slice_fn': lambda out, i: out[:, i],\n 'summed_total': 2 * 10,\n },\n {\n 'num_dims': 3,\n 'broadcast_dims': (1, 2),\n 'slice_fn': lambda out, i: out[i, :, :],\n 'summed_total': 2 * 10 * 10,\n },\n {\n 'num_dims': 3,\n 'broadcast_dims': (1,),\n 'slice_fn': lambda out, i, j: out[i, :, j],\n 'summed_total': 2 * 10,\n },\n {\n 'num_dims': 4,\n 'broadcast_dims': (0, 2, 3),\n 'slice_fn': lambda out, i: out[:, i, :, :],\n 'summed_total': 2 * 10 * 10 * 10,\n },\n {\n 'num_dims': 4,\n 'broadcast_dims': (0, 1),\n 'slice_fn': lambda out, i, j: out[:, :, i, j],\n 'summed_total': 2 * 10 * 10,\n },\n {\n 'num_dims': 4,\n 'broadcast_dims': (3,),\n 'slice_fn': lambda out, i, j, k: out[i, j, k, :],\n 'summed_total': 2 * 10,\n },\n )\n def test_dropout_broadcast(\n self, num_dims, broadcast_dims, slice_fn, summed_total\n ):\n module = nn.Dropout(\n rate=0.5, broadcast_dims=broadcast_dims, deterministic=False\n )\n x = jnp.ones((10,) * num_dims)\n out = module.apply({}, x, rngs={'dropout': random.key(0)})\n\n for i in range(10):\n if num_dims - len(broadcast_dims) >= 2:\n for j in range(10):\n if num_dims - len(broadcast_dims) >= 3:\n for k in range(10):\n self.assertTrue(slice_fn(out, i, j, k).sum() in (0, summed_total))\n else:\n self.assertTrue(slice_fn(out, i, j).sum() in (0, summed_total))\n else:\n self.assertTrue(slice_fn(out, i).sum() in (0, summed_total))\n\n def test_dropout_manual_rng(self):\n class Foo(nn.Module):\n @nn.compact\n def __call__(self, x):\n key = self.make_rng('dropout')\n x1 = nn.Dropout(rate=0.5, deterministic=False)(x, rng=key)\n x2 = nn.Dropout(rate=0.5, deterministic=False)(x, rng=jax.random.clone(key))\n return x1, x2\n\n module = Foo()\n x1, x2 = module.apply(\n {}, jnp.ones((20, 20)), rngs={'dropout': random.key(0)}\n )\n\n np.testing.assert_array_equal(x1, x2)\n\n\n# TODO(flax-dev): add integration tests for RNN cells\nclass RecurrentTest(parameterized.TestCase):\n def test_lstm(self):\n lstm = nn.LSTMCell(features=4)\n rng = random.key(0)\n rng, key1, key2 = random.split(rng, 3)\n x = random.normal(key1, (2, 3))\n c0, h0 = lstm.initialize_carry(rng, x.shape)\n self.assertEqual(c0.shape, (2, 4))\n self.assertEqual(h0.shape, (2, 4))\n (carry, y), initial_params = lstm.init_with_output(key2, (c0, h0), x)\n self.assertEqual(carry[0].shape, (2, 4))\n self.assertEqual(carry[1].shape, (2, 4))\n np.testing.assert_allclose(y, carry[1])\n param_shapes = jax.tree_util.tree_map(np.shape, initial_params['params'])\n self.assertEqual(\n param_shapes,\n {\n 'ii': {'kernel': (3, 4)},\n 'if': {'kernel': (3, 4)},\n 'ig': {'kernel': (3, 4)},\n 'io': {'kernel': (3, 4)},\n 'hi': {'kernel': (4, 4), 'bias': (4,)},\n 'hf': {'kernel': (4, 4), 'bias': (4,)},\n 'hg': {'kernel': (4, 4), 'bias': (4,)},\n 'ho': {'kernel': (4, 4), 'bias': (4,)},\n },\n )\n\n @parameterized.parameters(\n {\n 'module_cls': nn.SimpleCell,\n 'expected_param_shapes': {\n 'i': {'kernel': (3, 4), 'bias': (4,)},\n 'h': {'kernel': (4, 4)},\n },\n },\n {\n 'module_cls': nn.GRUCell,\n 'expected_param_shapes': {\n 'ir': {'kernel': (3, 4), 'bias': (4,)},\n 'iz': {'kernel': (3, 4), 'bias': (4,)},\n 'in': {'kernel': (3, 4), 'bias': (4,)},\n 'hr': {'kernel': (4, 4)},\n 'hz': {'kernel': (4, 4)},\n 'hn': {'kernel': (4, 4), 'bias': (4,)},\n },\n },\n {\n 'module_cls': nn.MGUCell,\n 'expected_param_shapes': {\n 'if': {'kernel': (3, 4), 'bias': (4,)},\n 'in': {'kernel': (3, 4), 'bias': (4,)},\n 'hf': {'kernel': (4, 4)},\n 'hn': {'kernel': (4, 4), 'bias': (4,)},\n },\n },\n )\n def test_gated_units(self, module_cls, expected_param_shapes):\n module = module_cls(features=4)\n rng = random.key(0)\n rng, key1, key2 = random.split(rng, 3)\n x = random.normal(key1, (2, 3))\n carry0 = module.initialize_carry(rng, x.shape)\n self.assertEqual(carry0.shape, (2, 4))\n (carry, y), initial_params = module.init_with_output(key2, carry0, x)\n self.assertEqual(carry.shape, (2, 4))\n np.testing.assert_allclose(y, carry)\n param_shapes = jax.tree_util.tree_map(np.shape, initial_params['params'])\n self.assertEqual(\n param_shapes,\n expected_param_shapes,\n )\n if module_cls == nn.MGUCell:\n self.assertTrue(\n (initial_params['params']['if']['bias'] == jnp.ones((4,))).all()\n )\n self.assertTrue(\n (initial_params['params']['in']['bias'] == jnp.zeros((4,))).all()\n )\n self.assertTrue(\n (initial_params['params']['hn']['bias'] == jnp.zeros((4,))).all()\n )\n\n @parameterized.parameters(\n {'module_cls': nn.SimpleCell},\n {'module_cls': nn.GRUCell},\n {'module_cls': nn.MGUCell},\n )\n def test_complex_input_gated_units(self, module_cls):\n module_instance = module_cls(features=4)\n rng = random.key(0)\n rng, key1, key2 = random.split(rng, 3)\n x = random.normal(key1, (2, 3), dtype=jnp.complex64)\n carry0 = module_instance.initialize_carry(rng, x.shape)\n self.assertEqual(carry0.shape, (2, 4))\n (carry, y), _ = module_instance.init_with_output(key2, carry0, x)\n self.assertEqual(carry.dtype, jnp.complex64)\n self.assertEqual(y.dtype, jnp.complex64)\n\n def test_convlstm(self):\n lstm = nn.ConvLSTMCell(features=6, kernel_size=(3, 3))\n rng = random.key(0)\n rng, key1, key2 = random.split(rng, 3)\n x = random.normal(key1, (2, 4, 4, 3))\n c0, h0 = lstm.initialize_carry(rng, x.shape)\n self.assertEqual(c0.shape, (2, 4, 4, 6))\n self.assertEqual(h0.shape, (2, 4, 4, 6))\n (carry, y), initial_params = lstm.init_with_output(key2, (c0, h0), x)\n self.assertEqual(carry[0].shape, (2, 4, 4, 6))\n self.assertEqual(carry[1].shape, (2, 4, 4, 6))\n np.testing.assert_allclose(y, carry[1])\n param_shapes = jax.tree_util.tree_map(np.shape, initial_params['params'])\n self.assertEqual(\n param_shapes,\n {\n 'hh': {'bias': (6 * 4,), 'kernel': (3, 3, 6, 6 * 4)},\n 'ih': {'bias': (6 * 4,), 'kernel': (3, 3, 3, 6 * 4)},\n },\n )\n\n def test_optimized_lstm_cell_matches_regular(self):\n # Create regular LSTMCell.\n lstm = nn.LSTMCell(features=4)\n rng = random.key(0)\n rng, key1, key2 = random.split(rng, 3)\n x = random.normal(key1, (2, 3))\n c0, h0 = lstm.initialize_carry(rng, x.shape)\n self.assertEqual(c0.shape, (2, 4))\n self.assertEqual(h0.shape, (2, 4))\n (_, y), lstm_params = lstm.init_with_output(key2, (c0, h0), x)\n\n # Create OptimizedLSTMCell.\n lstm_opt = nn.OptimizedLSTMCell(features=4)\n rng = random.key(0)\n rng, key1, key2 = random.split(rng, 3)\n x = random.normal(key1, (2, 3))\n c0, h0 = lstm_opt.initialize_carry(rng, x.shape)\n self.assertEqual(c0.shape, (2, 4))\n self.assertEqual(h0.shape, (2, 4))\n (_, y_opt), lstm_opt_params = lstm_opt.init_with_output(key2, (c0, h0), x)\n\n np.testing.assert_allclose(y, y_opt, rtol=1e-6)\n check_eq(lstm_params, lstm_opt_params)\n\n def test_mgu_reset_gate(self):\n module = nn.MGUCell(features=4, reset_gate=False)\n rng = random.key(0)\n rng, key1, key2 = random.split(rng, 3)\n x = random.normal(key1, (2, 3))\n carry0 = module.initialize_carry(rng, x.shape)\n (carry, y), v = module.init_with_output(key2, carry0, x)\n\n self.assertIn('kernel', v['params']['hn'])\n self.assertNotIn('bias', v['params']['hn'])\n\n f = jax.nn.sigmoid(\n jnp.dot(x, v['params']['if']['kernel'])\n + v['params']['if']['bias'].reshape(1, -1)\n + jnp.dot(carry0, v['params']['hf']['kernel'])\n )\n n = jax.nn.tanh(\n jnp.dot(x, v['params']['in']['kernel'])\n + v['params']['in']['bias'].reshape(1, -1)\n + jnp.dot(carry0, v['params']['hn']['kernel'])\n )\n expected_out = (1 - f) * n + f * carry0\n np.testing.assert_allclose(y, expected_out)\n\n\nclass IdsTest(absltest.TestCase):\n def test_hashable(self):\n id1 = ids.uuid()\n id2 = ids.uuid()\n self.assertEqual(id1, id1)\n self.assertNotEqual(id1, id2)\n self.assertNotEqual(hash(id1), hash(id2))\n id1c = copy.copy(id1)\n id1dc = copy.deepcopy(id1)\n self.assertNotEqual(hash(id1), hash(id1c))\n self.assertNotEqual(hash(id1), hash(id1dc))\n\ndef get_fp8_dtypes(fp8_genre):\n assert fp8_genre in ('OCP', 'NANOO')\n if fp8_genre == 'OCP':\n e4m3_dtype = jnp.float8_e4m3fn\n e5m2_dtype = jnp.float8_e5m2\n else: # fp8_genre == 'NANOO'\n e4m3_dtype = jnp.float8_e4m3fnuz\n e5m2_dtype = jnp.float8_e5m2fnuz\n return e4m3_dtype, e5m2_dtype\n\nclass Fp8Test(parameterized.TestCase):\n @parameterized.parameters(\n {'x_shape': (16, 32), 'y_shape': (32, 64), 'g_shape': (16, 64), 'eqn': 'mk,kn->mn'},\n {'x_shape': (2, 3, 32), 'y_shape': (64, 32), 'g_shape': (2, 3, 64), 'eqn': '...k,nk->...n'},\n {'x_shape': (2, 3, 64), 'y_shape': (64, 32), 'g_shape': (2, 3, 32), 'eqn': '...k,kn->...n'},\n )\n def test_fp8_einsum(self, x_shape, y_shape, g_shape, eqn):\n rng, key1, key2, key3 = random.split(random.key(42), 4)\n x = random.normal(key1, x_shape)\n y = random.normal(key2, y_shape)\n g = random.normal(key3, g_shape)\n e4m3_dtype = jnp.float8_e4m3fn\n e5m2_dtype = jnp.float8_e5m2\n cast_to_representable = functools.partial(\n fp8_ops.qdq,\n scale=jnp.ones((1,)),\n compute_dtype=jnp.float32,\n )\n\n x = cast_to_representable(x, e4m3_dtype)\n y = cast_to_representable(y, e4m3_dtype)\n g = cast_to_representable(g, e5m2_dtype)\n\n p = nn.Fp8Einsum()\n vars = p.init(rng, eqn, x, y)\n def loss_fn(vars, x, y):\n out = p.apply(vars, eqn, x, y)\n return jnp.sum(out * g.astype(out.dtype))\n step_fn = jax.value_and_grad(loss_fn, argnums=[1, 2])\n out, grads = jax.jit(step_fn)(vars, x, y)\n\n def loss_fn_ref(x, y):\n out = jnp.einsum(eqn, x, y)\n return jnp.sum(out * g.astype(out.dtype))\n step_fn_ref = jax.value_and_grad(loss_fn_ref, argnums=[0, 1])\n out_ref, grads_ref = jax.jit(step_fn_ref)(x, y)\n\n np.testing.assert_allclose(out, out_ref, atol=1e-02, rtol=1e-02)\n np.testing.assert_allclose(grads[0], grads_ref[0], atol=1e-02, rtol=1e-02)\n np.testing.assert_allclose(grads[1], grads_ref[1], atol=1e-02, rtol=1e-02)\n\n\n @parameterized.parameters(\n {'fp8_genre': 'OCP'}, {'fp8_genre': 'NANOO'}\n )\n def test_fp8_dot_general_injection(self, fp8_genre):\n # Used to cast the inputs to be representable in FP8, so that the difference\n # of the results from the original gemm and fp8 gemm is small.\n cast_to_representable = functools.partial(\n fp8_ops.qdq,\n scale=jnp.ones((1,)),\n compute_dtype=jnp.float32,\n )\n\n e4m3_dtype, e5m2_dtype = get_fp8_dtypes(fp8_genre)\n\n init_key, random_key = random.split(random.PRNGKey(seed=123), 2)\n x = cast_to_representable(\n random.uniform(random_key, (16, 32)), e4m3_dtype\n )\n dy = cast_to_representable(\n random.uniform(random_key, (16, 64)), e5m2_dtype\n )\n\n quant_cls = nn.Fp8DotGeneral if fp8_genre == 'OCP' else nn.NANOOFp8DotGeneralOp\n\n def run(fp8_injection, expected_shapes):\n p = nn.DenseGeneral(features=64, name='dense')\n\n if fp8_injection:\n p.dot_general_cls = quant_cls\n\n init_fn = jax.jit(p.init_with_output)\n y, initial_vars = init_fn(init_key, x)\n var_shapes = jax.tree_util.tree_map(jnp.shape, initial_vars)\n self.assertEqual(var_shapes, expected_shapes)\n\n def _train(variables, x):\n y = p.apply(variables, x)\n loss = y * dy\n return jnp.mean(loss)\n\n train_fn = jax.jit(jax.value_and_grad(_train, argnums=[0, 1]))\n outputs, grads = train_fn(initial_vars, x)\n return outputs, grads\n\n expected_shapes_original = {\n 'params': {'kernel': (32, 64), 'bias': (64,)},\n }\n\n expected_shapes_new = {\n 'params': {'kernel': (32, 64), 'bias': (64,)},\n fp8_ops.OVERWRITE_WITH_GRADIENT: {\n f'{quant_cls.__name__}_0': {\n 'input_amax_history': (1024,),\n 'kernel_amax_history': (1024,),\n 'output_grad_amax_history': (1024,),\n 'input_scale': (1,),\n 'kernel_scale': (1,),\n 'output_grad_scale': (1,),\n }\n },\n }\n output1a, output1b = run(False, expected_shapes_original)\n output2a, output2b = run(True, expected_shapes_new)\n dw1, dw2 = output1b[0]['params']['kernel'], output2b[0]['params']['kernel']\n dx1, dx2 = output1b[1], output2b[1]\n\n np.testing.assert_allclose(output1a, output2a, atol=1e-02)\n np.testing.assert_allclose(dw1, dw2, atol=1e-04)\n np.testing.assert_allclose(dx1, dx2, atol=1e-04)\n\n @parameterized.parameters(\n {'fp8_genre': 'OCP'}, {'fp8_genre': 'NANOO'}\n )\n def test_fp8_train_state(self, fp8_genre):\n key, init_key, random_key = random.split(random.PRNGKey(seed=123), 3)\n x = random.uniform(random_key, (16, 16), dtype=jnp.float32)\n\n quant_cls = nn.Fp8DotGeneral if fp8_genre == 'OCP' else nn.NANOOFp8DotGeneralOp\n dense = nn.DenseGeneral(\n features=32, use_bias=True, dot_general_cls=quant_cls\n )\n\n init_fn = jax.jit(dense.init)\n variables = init_fn(init_key, x)\n opt = optax.adam(learning_rate=0.1)\n state = train_state.TrainState.create(\n params=variables, tx=opt, apply_fn=dense.apply\n )\n\n def _roll_and_update(amax_h, update):\n return jnp.roll(amax_h, shift=-1, axis=0).at[0].set(update)\n\n def _train_loss(state, x, dy):\n def loss_fn(vars):\n y = state.apply_fn(vars, x)\n loss = y * dy.astype(y.dtype)\n return jnp.sum(loss)\n\n grad_fn = jax.grad(loss_fn)\n grads = grad_fn(state.params)\n state = state.apply_gradients(grads=grads)\n return state\n\n train_fn = jax.jit(_train_loss)\n\n scale_x, amax_history_x = jnp.ones(()), jnp.zeros((1024,))\n scale_k, amax_history_k = jnp.ones(()), jnp.zeros((1024,))\n scale_g, amax_history_g = jnp.ones(()), jnp.zeros((1024,))\n e4m3_dtype, e5m2_dtype = get_fp8_dtypes(fp8_genre)\n e4m3_max = jnp.finfo(e4m3_dtype).max.astype(jnp.float32)\n e5m2_max = jnp.finfo(e5m2_dtype).max.astype(jnp.float32)\n\n for _ in range(5):\n key, random_key = random.split(key, 2)\n x = random.normal(random_key, (16, 16), dtype=jnp.float32)\n g = random.normal(random_key, (16, 32), dtype=jnp.float32)\n k = state.params['params']['kernel']\n\n # Manually compute the expected amax history and scaling factors.\n amax_from_history_x = jnp.max(amax_history_x, axis=0)\n amax_from_history_k = jnp.max(amax_history_k, axis=0)\n amax_from_history_g = jnp.max(amax_history_g, axis=0)\n scale_x = fp8_ops.compute_scale(amax_from_history_x, scale_x, e4m3_max)\n scale_k = fp8_ops.compute_scale(amax_from_history_k, scale_k, e4m3_max)\n scale_g = fp8_ops.compute_scale(amax_from_history_g, scale_g, e5m2_max)\n amax_history_x = _roll_and_update(amax_history_x, jnp.max(jnp.abs(x)))\n amax_history_k = _roll_and_update(amax_history_k, jnp.max(jnp.abs(k)))\n amax_history_g = _roll_and_update(amax_history_g, jnp.max(jnp.abs(g)))\n\n state = train_fn(state, x, g)\n\n rtol, atol = 0.001, 0.001\n fp8_vars = state.params[fp8_ops.OVERWRITE_WITH_GRADIENT][\n f'{quant_cls.__name__}_0'\n ]\n np.testing.assert_allclose(\n fp8_vars['input_amax_history'],\n amax_history_x,\n rtol=rtol,\n atol=atol,\n )\n np.testing.assert_allclose(\n fp8_vars['kernel_amax_history'],\n amax_history_k,\n rtol=rtol,\n atol=atol,\n )\n np.testing.assert_allclose(\n fp8_vars['output_grad_amax_history'],\n amax_history_g,\n rtol=rtol,\n atol=atol,\n )\n\n np.testing.assert_allclose(fp8_vars['input_scale'][0], scale_x)\n np.testing.assert_allclose(fp8_vars['kernel_scale'][0], scale_k)\n np.testing.assert_allclose(fp8_vars['output_grad_scale'][0], scale_g)\n\n @parameterized.parameters(\n {'fp8_genre': 'OCP', 'use_jit': True},\n {'fp8_genre': 'OCP', 'use_jit': False},\n {'fp8_genre': 'NANOO', 'use_jit': True},\n {'fp8_genre': 'NANOO', 'use_jit': False}\n )\n def test_fp8_meta_dtype(self, fp8_genre, use_jit):\n if not use_jit and not fp8_ops.CAN_USE_EARRAY:\n self.skipTest(\"TODO: requires newer jax that has earray\")\n f32 = jnp.dtype('float32')\n fmax32 = fp8_ops.fp32_max_grad\n e4m3_dtype, _ = get_fp8_dtypes(fp8_genre)\n e4m3_max = 448 if fp8_genre == 'OCP' else 240\n\n # Create a scan loop with reused ah_f32 and sf_f32. So, the autograd will\n # accumulate the grads of them. We expect the max op (rather than add op)\n # for the accumulation by converting them to fmax32 dtype.\n def outer(x, ah_f32, sf_f32):\n ah_fmax32 = jax.lax.convert_element_type(ah_f32, fmax32)\n sf_fmax32 = jax.lax.convert_element_type(sf_f32, fmax32)\n array_x = jnp.array([x], f32)\n def body_fun(carry, _):\n carry = fp8_ops.in_qdq(f32, e4m3_dtype, carry, sf_fmax32, ah_fmax32)\n return carry, None\n array_x, _ = jax.lax.scan(body_fun, array_x, None, length=3)\n return array_x[0]\n\n outer_fn = jax.grad(outer, (0, 1, 2))\n if use_jit:\n outer_fn = jax.jit(outer_fn)\n ah = jnp.array([0., 0., 0.], f32)\n sf = jnp.array([1.], f32)\n # 1st iteration\n grads, new_ah, new_sf = outer_fn(2.0, ah, sf)\n np.testing.assert_allclose(new_ah, [2., 0., 0.])\n np.testing.assert_allclose(new_sf, [1.])\n # 2nd iteration\n grads, new_ah, new_sf = outer_fn(3., new_ah, new_sf)\n np.testing.assert_allclose(new_ah, [3., 0., 2.])\n np.testing.assert_allclose(new_sf, [2. / e4m3_max])\n # 3rd iteration\n grads, new_ah, new_sf = outer_fn(4., new_ah, new_sf)\n np.testing.assert_allclose(new_ah, [4., 2., 3.])\n np.testing.assert_allclose(new_sf, [3. / e4m3_max])\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/linen/linen_transforms_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Transforms tests.\"\"\"\n\nfrom functools import partial\nimport operator\nfrom typing import Any\nfrom collections.abc import Callable, Sequence\nimport unittest\n\nfrom absl.testing import absltest, parameterized\nfrom flax import errors\nfrom flax import linen as nn\nfrom flax import serialization\nfrom flax import struct\nfrom flax.core import copy, freeze, AxisMetadata\nfrom flax.linen.transforms import _HashableProxy\nimport jax\nfrom jax import random\nimport jax.numpy as jnp\nimport numpy as np\n\n# Parse absl flags test_srcdir and test_tmpdir.\njax.config.parse_flags_with_absl()\n\n# pylint: disable=attribute-defined-outside-init,unused-variable,g-wrong-blank-lines,g-bare-generic\n\n\ndef tree_equals(x, y):\n return jax.tree_util.tree_all(jax.tree_util.tree_map(operator.eq, x, y))\n\n\ndef tree_allclose(x, y):\n return jax.tree_util.tree_all(\n jax.tree_util.tree_map(lambda x, y: np.all(np.isclose(x, y)), x, y)\n )\n\n\nid_fn = lambda x: x\n\n\nclass TransformedMLP(nn.Module):\n features: Sequence[int]\n transform: Callable = id_fn\n\n @nn.compact\n def __call__(self, inputs):\n x = inputs\n for i, feat in enumerate(self.features):\n # JIT the Module (it's __call__ fn by default.)\n x = self.transform(nn.Dense)(feat, name=f'layers_{i}')(x)\n if i != len(self.features) - 1:\n x = nn.relu(x)\n return x\n\n\ndef decorated_MLP(transform: Callable = id_fn):\n class MLP(nn.Module):\n features: Sequence[int]\n\n @transform\n @nn.compact\n def __call__(self, inputs):\n x = inputs\n for i, feat in enumerate(self.features):\n # JIT the Module (it's __call__ fn by default.)\n x = nn.Dense(feat, name=f'layers_{i}')(x)\n if i != len(self.features) - 1:\n x = nn.relu(x)\n return x\n\n return MLP\n\n\nclass TransformTest(parameterized.TestCase):\n\n def assert_keys_equal(self, key1, key2):\n self.assertEqual(key1.dtype, key2.dtype)\n np.testing.assert_array_equal(random.key_data(key1), random.key_data(key2))\n\n def test_jit(self):\n key1, key2 = random.split(random.key(3), 2)\n x = random.uniform(key1, (4, 4))\n\n normal_model = TransformedMLP(features=[3, 4, 5])\n jit_model = TransformedMLP(features=[3, 4, 5], transform=nn.jit)\n init_variables = normal_model.init(key2, x)\n y1 = normal_model.apply(init_variables, x)\n y2 = jit_model.apply(init_variables, x)\n\n self.assertTrue(np.all(y1 == y2))\n\n def test_jit_decorated(self):\n key1, key2 = random.split(random.key(3), 2)\n x = random.uniform(key1, (4, 4))\n\n normal_model = decorated_MLP()(features=[3, 4, 5])\n jit_model = decorated_MLP(nn.jit)(features=[3, 4, 5])\n init_variables = normal_model.init(key2, x)\n y1 = normal_model.apply(init_variables, x)\n y2 = jit_model.apply(init_variables, x)\n\n self.assertTrue(np.all(y1 == y2))\n\n def test_jit_init_fn(self):\n class Foo(nn.Module):\n @nn.compact\n def __call__(self, x):\n return nn.Dense(2)(x)\n\n @nn.jit\n def init_with_output(self, rngs, *args, **kwargs):\n return super().init_with_output(rngs, *args, **kwargs)\n\n Foo().init_with_output(random.key(0), jnp.ones((2, 3)))\n\n\n def test_remat(self):\n key1, key2 = random.split(random.key(3), 2)\n x = random.uniform(key1, (4, 4))\n\n normal_model = TransformedMLP(features=[3, 4, 5])\n remat_model = TransformedMLP(features=[3, 4, 5], transform=nn.remat)\n init_variables = normal_model.init(key2, x)\n y1 = normal_model.apply(init_variables, x)\n y2 = remat_model.apply(init_variables, x)\n\n self.assertTrue(np.all(y1 == y2))\n\n def test_remat_decorated(self):\n key1, key2 = random.split(random.key(3), 2)\n x = random.uniform(key1, (4, 4))\n\n normal_model = decorated_MLP()(features=[3, 4, 5])\n remat_model = decorated_MLP(nn.remat)(features=[3, 4, 5])\n init_variables = normal_model.init(key2, x)\n y1 = normal_model.apply(init_variables, x)\n y2 = remat_model.apply(init_variables, x)\n\n self.assertTrue(np.all(y1 == y2))\n\n def test_remat_kwargs(self):\n raise unittest.SkipTest('test breaks with grad')\n\n class ConditionalReLU(nn.Module):\n\n @nn.compact\n def __call__(self, input, apply_relu: bool = False):\n return nn.relu(input) if apply_relu else input\n\n key = random.key(0)\n x = jnp.ones((4, 4)) * -1\n remat_model = nn.remat(ConditionalReLU)()\n p = remat_model.init(key, x)\n y = remat_model.apply(p, x, apply_relu=True)\n\n self.assertTrue(np.all(y == jnp.zeros_like(x)))\n\n # This next line crashes with a concretization error\n _ = jax.grad(lambda x: remat_model.apply(p, x, apply_relu=True))(x)\n\n def test_remat_static_argnums(self):\n test = self\n\n class Foo(nn.Module):\n train_is_static: bool\n\n @nn.compact\n def __call__(self, inputs, train: bool):\n if self.train_is_static:\n test.assertTrue(isinstance(train, bool))\n else:\n test.assertTrue(isinstance(train, jnp.ndarray))\n\n return nn.Dense(3, use_bias=False)(inputs)\n\n # set train as a static argument\n FooRemat = nn.remat(Foo, static_argnums=(2,))\n foo = FooRemat(train_is_static=True)\n\n x = jnp.empty((1, 2))\n variables = foo.init(random.key(0), x, True)\n y = foo.apply(variables, x, False)\n self.assertEqual(y.shape, (1, 3))\n\n # set train as a non-static arguments\n FooRemat = nn.remat(Foo, static_argnums=())\n foo = FooRemat(train_is_static=False)\n\n variables = foo.init(random.key(0), x, True)\n y = foo.apply(variables, x, False)\n self.assertEqual(y.shape, (1, 3))\n\n def test_remat_decorator_static_argnums(self):\n test = self\n\n class FooTrainStatic(nn.Module):\n\n @partial(nn.remat, static_argnums=(2,))\n @nn.compact\n def __call__(self, inputs, train: bool):\n test.assertTrue(isinstance(train, bool))\n\n return nn.Dense(3, use_bias=False)(inputs)\n\n # set train as a static argument\n foo = FooTrainStatic()\n\n x = jnp.empty((1, 2))\n variables = foo.init(random.key(0), x, True)\n y = foo.apply(variables, x, False)\n self.assertEqual(y.shape, (1, 3))\n\n class FooTrainDynamic(nn.Module):\n\n @partial(nn.remat, static_argnums=())\n @nn.compact\n def __call__(self, inputs, train: bool):\n test.assertTrue(isinstance(train, jnp.ndarray))\n\n return nn.Dense(3, use_bias=False)(inputs)\n\n # set train as a non-static arguments\n foo = FooTrainDynamic()\n\n variables = foo.init(random.key(0), x, True)\n y = foo.apply(variables, x, False)\n self.assertEqual(y.shape, (1, 3))\n\n def test_vmap(self):\n key1, key2, key3 = random.split(random.key(3), 3)\n x = random.uniform(key1, (4, 4))\n x2 = random.uniform(key2, (5, 4, 4))\n\n def vmap(cls):\n return nn.vmap(\n cls,\n in_axes=(0,),\n variable_axes={'params': None},\n split_rngs={'params': False},\n )\n\n normal_model = TransformedMLP(features=[3, 4, 5])\n vmap_model = TransformedMLP(features=[3, 4, 5], transform=vmap)\n init_variables = normal_model.init(key3, x)\n # simulate vmap in python for comparison:\n y1 = jnp.vstack([\n normal_model.apply(init_variables, x2[i])[None, ...]\n for i in np.arange(x2.shape[0])\n ])\n y2 = vmap_model.apply(init_variables, x2)\n np.testing.assert_allclose(y1, y2, atol=1e-6)\n\n def test_vmap_decorated(self):\n key1, key2, key3 = random.split(random.key(3), 3)\n x = random.uniform(key1, (4, 4))\n x2 = random.uniform(key2, (5, 4, 4))\n\n def vmap(fn):\n return nn.vmap(\n fn,\n in_axes=(0,),\n variable_axes={'params': None},\n split_rngs={'params': False},\n )\n\n normal_model = decorated_MLP()(features=[3, 4, 5])\n vmap_model = decorated_MLP(vmap)(features=[3, 4, 5])\n init_variables = normal_model.init(key3, x)\n # simulate vmap in python for comparison:\n y1 = jnp.vstack([\n normal_model.apply(init_variables, x2[i])[None, ...]\n for i in np.arange(x2.shape[0])\n ])\n y2 = vmap_model.apply(init_variables, x2)\n np.testing.assert_allclose(y1, y2, atol=1e-6)\n\n def test_vmap_batchnorm(self):\n key1, key2, key3 = random.split(random.key(3), 3)\n x = random.uniform(key1, (4, 4))\n x2 = random.uniform(key2, (5, 4, 4))\n\n def vmap(cls):\n return nn.vmap(\n cls,\n in_axes=(0,),\n variable_axes={'params': None, 'batch_stats': None},\n split_rngs={'params': False},\n axis_name='batch',\n )\n\n class MlpBn(nn.Module):\n axis_name: Any = None\n\n @nn.compact\n def __call__(self, x):\n x = nn.Dense(3)(x)\n x = nn.BatchNorm(axis_name=self.axis_name, use_running_average=False)(x)\n return x\n\n normal_model = MlpBn()\n vmap_model = vmap(MlpBn)(axis_name='batch')\n init_variables = normal_model.init(key3, x)\n y1 = normal_model.apply(\n init_variables, x2.reshape((-1, 4)), mutable=['batch_stats']\n )[0]\n y1 = y1.reshape((5, 4, 3))\n y2 = vmap_model.apply(init_variables, x2, mutable=['batch_stats'])[0]\n np.testing.assert_allclose(y1, y2, atol=1e-5)\n\n def test_scan(self):\n class SimpleScan(nn.Module):\n features: int\n\n @nn.compact\n def __call__(self, c, xs):\n LSTM = nn.scan(\n nn.LSTMCell,\n variable_broadcast='params',\n split_rngs={'params': False},\n )\n return LSTM(self.features, name='lstm_cell')(c, xs)\n\n key1, key2 = random.split(random.key(0), 2)\n xs = random.uniform(key1, (5, 3, 2))\n dummy_rng = random.key(0)\n init_carry = nn.LSTMCell(2).initialize_carry(dummy_rng, xs[0].shape)\n model = SimpleScan(2)\n init_variables = model.init(key2, init_carry, xs)\n # simulate scan in python for comparison:\n c = init_carry\n ys = []\n lstmcell_variables = freeze(\n {'params': init_variables['params']['lstm_cell']}\n )\n for i in range(xs.shape[0]):\n c, y = nn.LSTMCell(2).apply(lstmcell_variables, c, xs[i])\n ys.append(y[None, ...])\n y1 = jnp.vstack(ys)\n\n c2, y2 = model.apply(init_variables, init_carry, xs)\n np.testing.assert_allclose(y1, y2, atol=1e-7)\n np.testing.assert_allclose(c[0], c2[0], atol=1e-7)\n np.testing.assert_allclose(c[1], c2[1], atol=1e-7)\n\n def test_scan_decorated(self):\n class SimpleScan(nn.Module):\n features: int\n\n @partial(\n nn.scan,\n variable_broadcast='params',\n in_axes=(nn.broadcast, 0),\n split_rngs={'params': False},\n )\n @nn.compact\n def __call__(self, c, b, xs):\n assert b.shape == (4,)\n return nn.LSTMCell(self.features, name='lstm_cell')(c, xs)\n\n key1, key2 = random.split(random.key(0), 2)\n xs = random.uniform(key1, (4, 3, 2))\n b = jnp.ones((4,))\n dummy_rng = random.key(0)\n init_carry = nn.LSTMCell(2).initialize_carry(dummy_rng, xs[0].shape)\n model = SimpleScan(2)\n init_variables = model.init(key2, init_carry, b, xs)\n # simulate scan in python for comparison:\n c = init_carry\n ys = []\n lstmcell_variables = freeze(\n {'params': init_variables['params']['lstm_cell']}\n )\n for i in range(xs.shape[0]):\n c, y = nn.LSTMCell(2).apply(lstmcell_variables, c, xs[i])\n ys.append(y[None, ...])\n y1 = jnp.vstack(ys)\n\n c2, y2 = model.apply(init_variables, init_carry, b, xs)\n np.testing.assert_allclose(y1, y2, atol=1e-7)\n np.testing.assert_allclose(c[0], c2[0], atol=1e-7)\n np.testing.assert_allclose(c[1], c2[1], atol=1e-7)\n\n def test_scan_negative_axes(self):\n class Foo(nn.Module):\n\n @nn.compact\n def __call__(self, _, x):\n x = nn.Dense(4)(x)\n return None, x\n\n class Bar(nn.Module):\n\n @nn.compact\n def __call__(self, x):\n _, x = nn.scan(\n Foo,\n variable_broadcast='params',\n split_rngs=dict(params=False),\n in_axes=1,\n out_axes=-1,\n )()(None, x)\n return x\n\n y, variables = Bar().init_with_output(\n {'params': jax.random.PRNGKey(0)},\n jax.random.normal(jax.random.PRNGKey(1), shape=[1, 2, 3]),\n )\n params = variables['params']\n\n self.assertEqual(y.shape, (1, 4, 2))\n self.assertEqual(params['ScanFoo_0']['Dense_0']['kernel'].shape, (3, 4))\n self.assertEqual(params['ScanFoo_0']['Dense_0']['bias'].shape, (4,))\n\n def test_multiscope_lifting_simple(self):\n class Counter(nn.Module):\n\n @nn.compact\n def __call__(self):\n v = self.variable('counter', 'foo', lambda: jnp.array([0]))\n v.value += jnp.array([1])\n return v.value\n\n class Outer(nn.Module):\n\n @nn.compact\n def __call__(self, x):\n cntr = nn.jit(Counter)(name='cntr')()\n return x\n\n class Inner(nn.Module):\n outer_module: nn.Module\n\n @nn.compact\n def __call__(self, x):\n return self.outer_module(x)\n\n class Test(nn.Module):\n\n @nn.compact\n def __call__(self, x):\n outer_dense = nn.jit(Outer)(name='outer')\n # we share stateful outer module as arg to two different, transformed modules:\n inner = nn.jit(Inner)(outer_dense, name='inner1')\n inner2 = nn.jit(Inner)(outer_dense, name='inner2')\n res = inner(x) + inner2(x)\n return res\n\n x = jnp.ones((10, 10))\n rngs = random.key(0)\n init_vars = Test(parent=None).init(rngs, x)\n _, new_vars = Test(parent=None).apply(init_vars, x, mutable=['counter'])\n self.assertEqual(\n init_vars['counter']['outer']['cntr']['foo'], jnp.array([2], jnp.int32)\n )\n self.assertEqual(\n new_vars['counter']['outer']['cntr']['foo'], jnp.array([4], jnp.int32)\n )\n\n def test_multiscope_lifting_simple_decorator(self):\n class Counter(nn.Module):\n\n @nn.jit\n @nn.compact\n def __call__(self):\n v = self.variable('counter', 'foo', lambda: jnp.array([0]))\n v.value += jnp.array([1])\n return v.value\n\n class Outer(nn.Module):\n\n @nn.jit\n @nn.compact\n def __call__(self, x):\n cntr = Counter(name='cntr')()\n return x\n\n class Inner(nn.Module):\n outer_module: nn.Module\n\n @nn.jit\n @nn.compact\n def __call__(self, x):\n return self.outer_module(x)\n\n class Test(nn.Module):\n\n @nn.compact\n def __call__(self, x):\n outer_dense = Outer(name='outer')\n # we share stateful outer module as arg to two different, transformed modules:\n inner = Inner(outer_dense, name='inner1')\n inner2 = Inner(outer_dense, name='inner2')\n res = inner(x) + inner2(x)\n return res\n\n x = jnp.ones((1, 1))\n rngs = random.key(0)\n init_vars = Test(parent=None).init(rngs, x)\n _, new_vars = Test(parent=None).apply(init_vars, x, mutable=['counter'])\n self.assertEqual(\n init_vars['counter']['outer']['cntr']['foo'], jnp.array([2], jnp.int32)\n )\n self.assertEqual(\n new_vars['counter']['outer']['cntr']['foo'], jnp.array([4], jnp.int32)\n )\n\n def test_multiscope_lifting_argtree(self):\n class Counter(nn.Module):\n\n @nn.compact\n def __call__(self):\n v = self.variable('counter', 'foo', lambda: jnp.array([0]))\n v.value += jnp.array([1])\n return v.value\n\n class Outer(nn.Module):\n\n @nn.compact\n def __call__(self, x):\n cntr = nn.jit(Counter)(name='cntr')()\n return x\n\n class Inner(nn.Module):\n outer_module: Sequence[nn.Module]\n\n @nn.compact\n def __call__(self, x):\n return self.outer_module[0](x) + self.outer_module[1](x)\n\n class Test(nn.Module):\n\n @nn.compact\n def __call__(self, x):\n outer_dense1 = nn.jit(Outer)(name='outer1')\n outer_dense2 = nn.jit(Outer)(name='outer2')\n # we share stateful outer module as arg to two different, transformed modules:\n inner1 = nn.jit(Inner)((outer_dense1, outer_dense2), name='inner1')\n inner2 = nn.jit(Inner)((outer_dense1, outer_dense2), name='inner2')\n res = inner1(x) + inner2(x)\n return res\n\n x = jnp.ones((1, 1))\n rngs = random.key(0)\n init_vars = Test(parent=None).init(rngs, x)\n _, new_vars = Test(parent=None).apply(init_vars, x, mutable=['counter'])\n self.assertEqual(\n init_vars['counter']['outer1']['cntr']['foo'], jnp.array([2], jnp.int32)\n )\n self.assertEqual(\n new_vars['counter']['outer1']['cntr']['foo'], jnp.array([4], jnp.int32)\n )\n self.assertEqual(\n init_vars['counter']['outer2']['cntr']['foo'], jnp.array([2], jnp.int32)\n )\n self.assertEqual(\n new_vars['counter']['outer2']['cntr']['foo'], jnp.array([4], jnp.int32)\n )\n\n def test_multiscope_lifting_argtree_decorator(self):\n class Counter(nn.Module):\n\n @nn.jit\n @nn.compact\n def __call__(self):\n v = self.variable('counter', 'foo', lambda: jnp.array([0]))\n v.value += jnp.array([1])\n return v.value\n\n class Outer(nn.Module):\n\n @nn.jit\n @nn.compact\n def __call__(self, x):\n cntr = nn.jit(Counter)(name='cntr')()\n return x\n\n class Inner(nn.Module):\n outer_module: Sequence[nn.Module]\n\n @nn.jit\n @nn.compact\n def __call__(self, x):\n return self.outer_module[0](x) + self.outer_module[1](x)\n\n class Test(nn.Module):\n\n @nn.compact\n def __call__(self, x):\n outer_dense1 = Outer(name='outer1')\n outer_dense2 = Outer(name='outer2')\n # we share stateful outer module as arg to two different, transformed modules:\n inner1 = Inner((outer_dense1, outer_dense2), name='inner1')\n inner2 = Inner((outer_dense1, outer_dense2), name='inner2')\n res = inner1(x) + inner2(x)\n return res\n\n x = jnp.ones((1, 1))\n rngs = random.key(0)\n init_vars = Test(parent=None).init(rngs, x)\n _, new_vars = Test(parent=None).apply(init_vars, x, mutable=['counter'])\n self.assertEqual(\n init_vars['counter']['outer1']['cntr']['foo'], jnp.array([2], jnp.int32)\n )\n self.assertEqual(\n new_vars['counter']['outer1']['cntr']['foo'], jnp.array([4], jnp.int32)\n )\n self.assertEqual(\n init_vars['counter']['outer2']['cntr']['foo'], jnp.array([2], jnp.int32)\n )\n self.assertEqual(\n new_vars['counter']['outer2']['cntr']['foo'], jnp.array([4], jnp.int32)\n )\n\n def test_multiscope_lifting_simple_decorator_w_jit(self):\n # TODO: actually test jaxpr on a simpler module.\n class Counter(nn.Module):\n\n @nn.jit\n @nn.compact\n def __call__(self):\n v = self.variable('counter', 'foo', lambda: jnp.array([0]))\n v.value += jnp.array([1])\n return v.value\n\n class Outer(nn.Module):\n\n @nn.jit\n @nn.compact\n def __call__(self, x):\n cntr = Counter(name='cntr')()\n return x\n\n class Inner(nn.Module):\n outer_module: nn.Module\n\n @nn.jit\n @nn.compact\n def __call__(self, x):\n return self.outer_module(x)\n\n class Test(nn.Module):\n\n @nn.jit\n @nn.compact\n def __call__(self, x):\n outer_dense = Outer(name='outer')\n # we share stateful outer module as arg to two different, transformed modules:\n inner = Inner(outer_dense, name='inner1')\n inner2 = Inner(outer_dense, name='inner2')\n res = inner(x) + inner2(x)\n return res\n\n x = jnp.ones((1, 1))\n rngs = random.key(0)\n init_vars = Test(parent=None).init(rngs, x)\n _, new_vars = Test(parent=None).apply(init_vars, x, mutable=['counter'])\n self.assertEqual(\n init_vars['counter']['outer']['cntr']['foo'], jnp.array([2], jnp.int32)\n )\n self.assertEqual(\n new_vars['counter']['outer']['cntr']['foo'], jnp.array([4], jnp.int32)\n )\n\n def test_vmapped_outer_module(self):\n class Outer(nn.Module):\n\n @nn.jit\n @nn.compact\n def __call__(self, x):\n return nn.Dense(5)(x)\n\n class Inner(nn.Module):\n outer_module: nn.Module\n\n @partial(\n nn.vmap,\n in_axes=(0,),\n variable_axes={'params': 0},\n split_rngs={'params': True},\n )\n @nn.jit\n @nn.compact\n def __call__(self, x):\n return self.outer_module(x)\n\n class Test(nn.Module):\n\n @nn.compact\n def __call__(self, x):\n outer_dense = Outer(name='outer')\n inner = Inner(outer_dense, name='inner1')\n inner2 = Inner(outer_dense, name='inner2')\n res = inner(x) + inner2(x)\n return res\n\n x = jnp.ones((3, 1, 2))\n rngs = random.key(0)\n init_vars = Test(parent=None).init(rngs, x)\n y = Test(parent=None).apply(init_vars, x)\n self.assertEqual(\n init_vars['params']['outer']['Dense_0']['kernel'].shape, (3, 2, 5)\n )\n self.assertEqual(\n init_vars['params']['outer']['Dense_0']['bias'].shape, (3, 5)\n )\n self.assertEqual(y.shape, (3, 1, 5))\n\n def test_module_transform_with_setup(self):\n class Foo(nn.Module):\n\n def setup(self):\n self.test = self.param('test', nn.initializers.ones_init(), ())\n\n def __call__(self, x):\n return x * self.test\n\n FooVmap = nn.vmap(\n Foo,\n in_axes=0,\n out_axes=0,\n variable_axes={'params': 0},\n split_rngs={'params': True},\n )\n variables = FooVmap().init(random.key(0), jnp.ones((4,)))\n self.assertEqual(variables['params']['test'].shape, (4,))\n\n def test_nested_module_args_vmap(self):\n class A(nn.Module):\n\n @nn.compact\n def __call__(self, x):\n return nn.Dense(3)(x)\n\n class B(nn.Module):\n A: nn.Module\n\n @nn.compact\n def __call__(self, x):\n return self.A(x)\n\n class C(nn.Module):\n B: nn.Module\n\n @partial(\n nn.vmap, variable_axes={'params': 0}, split_rngs={'params': True}\n )\n @nn.compact\n def __call__(self, x):\n return self.B(x)\n\n class D(nn.Module):\n\n @nn.compact\n def __call__(self, x):\n a = A()\n b = B(a)\n c = C(b)\n return c(x)\n\n key = random.key(0)\n x = jnp.ones((10, 10))\n p = D().init(key, x)\n\n variable_shapes = jax.tree_util.tree_map(jnp.shape, p)\n self.assertEqual(\n variable_shapes['params']['A_0']['Dense_0']['kernel'], (10, 10, 3)\n )\n self.assertEqual(\n variable_shapes['params']['A_0']['Dense_0']['bias'], (10, 3)\n )\n\n def test_nested_module_args_vmap_2(self):\n class A(nn.Module):\n\n @nn.compact\n def __call__(self, x):\n return nn.Dense(3)(x)\n\n class B(nn.Module):\n A: nn.Module\n\n @nn.compact\n def __call__(self, x):\n return self.A(x)\n\n class C(nn.Module):\n A: nn.Module\n B: nn.Module\n\n @partial(\n nn.vmap, variable_axes={'params': 0}, split_rngs={'params': True}\n )\n @nn.compact\n def __call__(self, x):\n return self.B(x) + self.A(x)\n\n class D(nn.Module):\n\n @nn.compact\n def __call__(self, x):\n a1 = A()\n a2 = A()\n b = B(a1)\n c = C(a2, b)\n return c(x)\n\n key = random.key(0)\n x = jnp.ones((10, 10))\n p = D().init(key, x)\n\n variable_shapes = jax.tree_util.tree_map(jnp.shape, p)\n self.assertEqual(\n variable_shapes['params']['A_0']['Dense_0']['kernel'], (10, 10, 3)\n )\n self.assertEqual(\n variable_shapes['params']['A_0']['Dense_0']['bias'], (10, 3)\n )\n self.assertEqual(\n variable_shapes['params']['A_1']['Dense_0']['kernel'], (10, 10, 3)\n )\n self.assertEqual(\n variable_shapes['params']['A_1']['Dense_0']['bias'], (10, 3)\n )\n\n def test_nested_setup_calls_count(self):\n D = 3\n N = 4\n setup_cntr = 0\n call_cntr = 0\n\n class Repeat(nn.Module):\n mdl_def: Any\n\n def setup(self):\n self.lyrs = [self.mdl_def() for _ in range(N)]\n\n @nn.remat # we just use remat as a convenient test of transform logic\n def __call__(self, x):\n for lyr in self.lyrs:\n lyr(x)\n return x\n\n class Counter(nn.Module):\n\n def setup(self):\n nonlocal setup_cntr\n setup_cntr += 1\n self.dense = nn.Dense(2, use_bias=False)\n\n @nn.remat\n def __call__(self, x):\n nonlocal call_cntr\n call_cntr += 1\n return self.dense(x)\n\n def nested_repeat(mdl):\n for _ in range(D):\n mdl = partial(Repeat, mdl)\n return mdl()\n\n _ = nested_repeat(Counter).init(random.key(0), jnp.ones((2,)))\n # setup_cntr == 128 due to 1 call in Counter.setup by _validate_setup\n # and 1 further \"real\" call.\n self.assertEqual(setup_cntr, 128)\n self.assertEqual(call_cntr, 64)\n\n def test_multimethod_setup_calls(self):\n cntr = 0\n\n class A(nn.Module):\n\n def setup(self):\n nonlocal cntr\n cntr += 1\n self.d = nn.Dense(2)\n\n @nn.remat\n def foo(self, x):\n return self.d(x)\n\n @nn.remat\n def bar(self, x):\n return self.d(x)\n\n class B(nn.Module):\n\n def setup(self):\n self.a = A()\n\n def __call__(self, x):\n y1 = self.a.foo(x)\n y2 = self.a.bar(x)\n return y1, y2\n\n key = random.key(0)\n x = jnp.ones((2,))\n (y1, y2), _ = B().init_with_output(key, x)\n np.testing.assert_array_equal(y1, y2)\n # cntr == 4 due to:\n # 1 call by _validate_setup\n # 1 call for the setup() outside transform boundary\n # and two further \"real\" calls in transform boundaries\n self.assertEqual(cntr, 4)\n\n def test_toplevel_submodule_adoption_transform(self):\n class A(nn.Module):\n\n @nn.compact\n def __call__(self, x):\n return nn.Dense(3)(x)\n\n class B(nn.Module):\n A: nn.Module\n\n @nn.compact\n def __call__(self, x):\n return self.A(x)\n\n class C(nn.Module):\n A: nn.Module\n B: nn.Module\n\n @partial(\n nn.vmap, variable_axes={'params': 0}, split_rngs={'params': True}\n )\n @nn.compact\n def __call__(self, x):\n return self.B(x) + self.A(x)\n\n class Csimple(nn.Module):\n A: nn.Module\n B: nn.Module\n\n @nn.compact\n def __call__(self, x):\n return self.B(x) + self.A(x)\n\n class D(nn.Module):\n\n @nn.compact\n def __call__(self, x):\n a1 = A()\n a2 = A()\n b = B(a1)\n c = C(a2, b)\n return c(x)\n\n key = random.key(0)\n x = jnp.ones((10, 10))\n p1 = D().init(key, x)\n y1 = D().apply(p1, x)\n\n a1 = A()\n a2 = A()\n b = B(a1)\n p2 = freeze({\n 'params': {\n 'A': p1['params']['A_0'],\n 'B': {\n 'A': p1['params']['A_1'],\n },\n }\n })\n\n # Test method wrapper transform.\n y2 = C(a2, b).apply(p2, x)\n np.testing.assert_allclose(y1, y2, atol=1e-7)\n # Test class transform.\n Ctrafo = nn.vmap(\n Csimple, variable_axes={'params': 0}, split_rngs={'params': True}\n )\n\n y3 = Ctrafo(a2, b).apply(p2, x)\n np.testing.assert_allclose(y1, y3, atol=1e-7)\n\n def test_toplevel_submodule_adoption_pytree_transform(self):\n class A(nn.Module):\n\n @nn.compact\n def __call__(self, c, x):\n counter = self.variable('counter', 'i', jnp.zeros, ())\n counter.value += 1\n x = nn.Dense(1)(x)\n return c, x\n\n class B(nn.Module):\n A: Any\n\n @nn.compact\n def __call__(self, c, x):\n return self.A['foo'](*self.A['bar'](c, x))\n\n a = A()\n As = {'foo': A(), 'bar': A()}\n b = nn.scan(\n B,\n in_axes=0,\n variable_carry='counter',\n variable_broadcast='params',\n split_rngs={'params': False},\n )(As)\n\n key = random.key(0)\n x = jnp.ones((10, 2))\n\n p = B(As).init(key, x, x)\n y, cntrs = b.apply(p, x, x, mutable='counter')\n ref_cntrs = {\n 'counter': {\n 'A_bar': {\n 'i': jnp.array(11.0),\n },\n 'A_foo': {\n 'i': jnp.array(11.0),\n },\n },\n }\n self.assertTrue(\n jax.tree_util.tree_all(\n jax.tree_util.tree_map(\n lambda x, y: np.testing.assert_allclose(x, y, atol=1e-7),\n cntrs,\n ref_cntrs,\n )\n )\n )\n\n def test_partially_applied_module_constructor_transform(self):\n k = random.key(0)\n x = jnp.ones((3, 4, 4))\n dense = partial(nn.Dense, use_bias=False)\n vmap_dense = nn.vmap(\n dense, variable_axes={'params': 0}, split_rngs={'params': True}\n )(4)\n init_vars = vmap_dense.init(k, x)\n init_vars_shapes = jax.tree_util.tree_map(jnp.shape, init_vars)\n ref_var_shapes = {\n 'params': {\n 'kernel': (3, 4, 4),\n },\n }\n self.assertTrue(tree_equals(init_vars_shapes, ref_var_shapes))\n\n def test_partial_module_method(self):\n k = random.key(0)\n x = jnp.ones((3, 4, 4))\n\n class Foo(nn.Module):\n\n @nn.compact\n def inner(self, x):\n return nn.Dense(2, use_bias=False)(x)\n\n def __call__(self, x):\n return nn.vmap(\n partial(Foo.inner),\n variable_axes={'params': 0},\n split_rngs={'params': True},\n )(self, x)\n\n init_vars = Foo().init(k, x)\n init_vars_shapes = jax.tree_util.tree_map(jnp.shape, init_vars)\n ref_var_shapes = {\n 'params': {'Dense_0': {'kernel': (3, 4, 2)}},\n }\n self.assertTrue(tree_equals(init_vars_shapes, ref_var_shapes))\n\n def test_variable_in_args_transform(self):\n class Test(nn.Module):\n\n @nn.jit\n @nn.compact\n def __call__(self, x):\n baz = self.variable('test', 'baz', jnp.zeros, x.shape)\n y = self.mutate_variable_in_method(x, baz)\n return y\n\n @nn.jit\n def mutate_variable_in_method(self, x, baz):\n baz.value += x\n return baz.value\n\n k = random.key(0)\n x = jnp.ones((1,))\n variables = Test().init(k, x)\n np.testing.assert_allclose(\n variables['test']['baz'],\n jnp.array([\n 1.0,\n ]),\n atol=1e-7,\n )\n y, variables = Test().apply(variables, x, mutable=['test'])\n np.testing.assert_allclose(\n variables['test']['baz'],\n jnp.array([\n 2.0,\n ]),\n atol=1e-7,\n )\n\n def test_module_instance_in_args_transform(self):\n class Inner(nn.Module):\n\n @nn.jit\n @nn.compact\n def __call__(self, x):\n baz = self.variable('test', 'baz', jnp.zeros, x.shape)\n baz.value += x\n return baz.value\n\n class Test(nn.Module):\n\n @nn.jit\n @nn.compact\n def __call__(self, x):\n inner = Inner(name='inner')\n y = self.call_instance_arg_in_method(x, inner)\n return y\n\n @nn.jit\n def call_instance_arg_in_method(self, x, inner):\n return inner(x)\n\n k = random.key(0)\n x = jnp.ones((1,))\n variables = Test().init(k, x)\n np.testing.assert_allclose(\n variables['test']['inner']['baz'],\n jnp.array([\n 1.0,\n ]),\n atol=1e-7,\n )\n y, variables = Test().apply(variables, x, mutable=['test'])\n np.testing.assert_allclose(\n variables['test']['inner']['baz'],\n jnp.array([\n 2.0,\n ]),\n atol=1e-7,\n )\n\n def test_module_instance_in_args_transform_nested(self):\n class Inner(nn.Module):\n\n @nn.jit\n @nn.compact\n def __call__(self, x):\n baz = self.variable('test', 'baz', jnp.zeros, x.shape)\n baz.value += x\n return baz.value\n\n class Outer(nn.Module):\n\n @nn.jit\n @nn.compact\n def __call__(self, inner, x):\n y = self.call_instance_arg_in_method(x, inner)\n return y\n\n @nn.jit\n def call_instance_arg_in_method(self, x, inner):\n return inner(x)\n\n class Test(nn.Module):\n\n @nn.jit\n @nn.compact\n def __call__(self, x):\n inner = Inner(name='inner')\n outer = Outer(name='outer')\n return outer(inner, x)\n\n k = random.key(0)\n x = jnp.ones((1,))\n variables = Test().init(k, x)\n np.testing.assert_allclose(\n variables['test']['inner']['baz'],\n jnp.array([\n 1.0,\n ]),\n atol=1e-7,\n )\n y, variables = Test().apply(variables, x, mutable=['test'])\n np.testing.assert_allclose(\n variables['test']['inner']['baz'],\n jnp.array([\n 2.0,\n ]),\n atol=1e-7,\n )\n\n def test_nested_variable_passing(self):\n class NestedVarUser(nn.Module):\n somevar: nn.Variable\n\n @nn.jit\n @nn.compact\n def __call__(self, x):\n self.somevar.value += x\n return x\n\n class VarUser(nn.Module):\n somevar: nn.Variable\n\n @nn.jit\n @nn.compact\n def __call__(self, x):\n return NestedVarUser(self.somevar)(x)\n\n class VarPasser(nn.Module):\n\n @nn.jit\n @nn.compact\n def __call__(self, x):\n baz = self.variable('test', 'baz', jnp.zeros, x.shape)\n y = VarUser(baz)(x)\n return y\n\n k = random.key(0)\n x = jnp.ones((1,))\n variables = VarPasser().init(k, x)\n np.testing.assert_allclose(\n variables['test']['baz'],\n jnp.array([\n 1.0,\n ]),\n atol=1e-7,\n )\n y, variables = VarPasser().apply(variables, x, mutable=['test'])\n np.testing.assert_allclose(\n variables['test']['baz'],\n jnp.array([\n 2.0,\n ]),\n atol=1e-7,\n )\n\n def test_returned_module_warning(self):\n class Foo(nn.Module):\n\n @nn.compact\n def __call__(self, x):\n return x\n\n class Bar(nn.Module):\n\n @nn.compact\n def __call__(self, x):\n f = self._helper()\n return f(x)\n\n @nn.jit\n def _helper(self):\n return Foo()\n\n b = Bar()\n with self.assertRaises(errors.TransformedMethodReturnValueError):\n b.apply({}, jnp.ones(2))\n\n def test_returned_variable_warning(self):\n class Bar(nn.Module):\n\n @nn.compact\n def __call__(self, x):\n f = self._helper()\n return f(x)\n\n @nn.jit\n def _helper(self):\n return nn.Variable(None, None, None, False)\n\n b = Bar()\n with self.assertRaises(errors.TransformedMethodReturnValueError):\n b.apply({}, jnp.ones(2))\n\n def test_nowrap(self):\n class Bar(nn.Module):\n\n @nn.compact\n def __call__(self, x):\n return self._helper(x)\n\n @nn.nowrap\n def _helper(self, x):\n if len(nn.module._context.module_stack) > 2: # pylint: disable=protected-access\n raise ValueError('Module stack too deep.')\n return x\n\n b = Bar()\n b.apply({}, jnp.ones(2))\n\n def test_map_variables_tied_autoencoder(self):\n def trans(variables):\n return jax.tree_util.tree_map(lambda x: x.T, variables)\n\n class TiedAutencoder(nn.Module):\n features: int\n latents: int\n\n @nn.compact\n def _call(self, x, decode):\n def f(self):\n return nn.Dense(\n self.features if decode else self.latents, use_bias=False\n )(x)\n\n if decode:\n map_fn = trans\n else:\n map_fn = lambda x: x\n return nn.map_variables(f, 'params', map_fn, map_fn, mutable=True)(self)\n\n def encode(self, x):\n return self._call(x, False)\n\n def decode(self, x):\n return self._call(x, True)\n\n def __call__(self, x):\n return self.decode(self.encode(x))\n\n x = jnp.ones((2, 4))\n ae = TiedAutencoder(4, 5)\n variables = ae.init(random.key(0), x)\n param_shapes = jax.tree_util.tree_map(jnp.shape, variables['params'])\n self.assertEqual(param_shapes, {'Dense_0': {'kernel': (4, 5)}})\n\n def test_map_variables_bit_weights(self):\n class BitWeights(nn.Module):\n\n @nn.compact\n def __call__(self, x):\n def sign(x):\n return jax.tree_util.tree_map(jnp.sign, x)\n\n BitDense = nn.map_variables(nn.Dense, 'params', sign, init=True)\n return BitDense(4)(x)\n\n bw = BitWeights()\n x = jnp.ones((2, 4))\n y, variables = bw.init_with_output(random.key(0), x)\n y_2 = bw.apply(variables, x)\n np.testing.assert_allclose(y, y_2)\n\n def test_remat_scan(self):\n class BigModel(nn.Module):\n\n @nn.compact\n def __call__(self, x):\n DenseStack = nn.remat_scan(nn.Dense, lengths=(100,))\n return DenseStack(8, name='dense_stack')(x)\n\n x = jnp.ones((2, 8))\n model = BigModel()\n variables = model.init(random.key(0), x)\n param_shapes = jax.tree_util.tree_map(jnp.shape, variables['params'])\n self.assertEqual(param_shapes['dense_stack']['kernel'], (100, 8, 8))\n self.assertEqual(param_shapes['dense_stack']['bias'], (100, 8))\n y = model.apply(variables, x)\n self.assertEqual(y.shape, (2, 8))\n\n def test_vjp(self):\n class Bar(nn.Module):\n\n @nn.compact\n def __call__(self, x, y):\n p = self.param('test', nn.initializers.constant(0.5), ())\n self.variable('state', 'counter', lambda: 0)\n return p * x * y\n\n class Foo(nn.Module):\n\n @nn.compact\n def __call__(self, x, y):\n z, bwd = nn.vjp(Bar.__call__, Bar(), x, y)\n return bwd(jnp.ones(z.shape))\n\n x = jnp.array([1.0, 2.0, 3.0])\n y = jnp.array([4.0, 5.0, 6.0])\n params = Foo().init(random.key(0), x, y)\n params_grad, x_grad, y_grad = Foo().apply(params, x, y)\n self.assertEqual(\n params_grad,\n {\n 'params': nn.FrozenDict({'test': 32.0}),\n },\n )\n np.testing.assert_allclose(x_grad, [2.0, 2.5, 3.0])\n np.testing.assert_allclose(y_grad, [0.5, 1.0, 1.5])\n\n def test_jvp(self):\n class Bar(nn.Module):\n\n @nn.compact\n def __call__(self, x):\n p = self.param('test', nn.initializers.zeros, ())\n self.variable('state', 'counter', lambda: 0)\n return p * x\n\n class Foo(nn.Module):\n\n @nn.compact\n def __call__(self, x):\n bar = Bar()\n vars_t = jax.tree_util.tree_map(\n jnp.ones_like, bar.variables.get('params', {})\n )\n _, out_t = nn.jvp(\n Bar.__call__, bar, (x,), (jnp.zeros_like(x),), {'params': vars_t}\n )\n return out_t\n\n x = jnp.ones((3,))\n params = Foo().init(random.key(0), x)\n y_t = Foo().apply(params, x)\n np.testing.assert_allclose(y_t, jnp.ones_like(x))\n\n def test_complicated_alias_mutation(self):\n class A(nn.Module):\n b: nn.Module\n\n @nn.jit\n @nn.compact\n def __call__(self, x):\n return self.b(x)\n\n class B(nn.Module):\n c: nn.Module\n\n @nn.jit\n @nn.compact\n def __call__(self, x):\n y = C(name='outer_c')(x)\n z = self.c(x)\n return z\n\n class C(nn.Module):\n\n @nn.jit\n @nn.compact\n def __call__(self, x):\n initialized = self.has_variable('muts', 'v')\n v = self.variable('muts', 'v', lambda: jnp.zeros_like(x))\n if initialized:\n v.value += x\n return x\n\n a = A(b=B(c=C()))\n k = random.key(0)\n x = jnp.ones((1,), jnp.float32)\n vs = a.init(k, x)\n y, vs_new = a.apply(\n vs,\n x,\n mutable=[\n 'muts',\n ],\n )\n np.testing.assert_array_equal(\n vs_new['muts']['b']['c']['v'], jnp.array([1.0], jnp.float32)\n )\n np.testing.assert_array_equal(\n vs_new['muts']['b']['outer_c']['v'], jnp.array([1.0], jnp.float32)\n )\n\n def test_custom_vjp(self):\n class Foo(nn.Module):\n\n @nn.compact\n def __call__(self, x):\n def f(mdl, x):\n return mdl(x)\n\n def fwd(mdl, x):\n return nn.vjp(f, mdl, x)\n\n def bwd(vjp_fn, y_t):\n params_t, input_t = vjp_fn(y_t)\n params_t = jax.tree_util.tree_map(jnp.sign, params_t)\n return params_t, input_t\n\n sign_grad = nn.custom_vjp(f, forward_fn=fwd, backward_fn=bwd)\n return sign_grad(nn.Dense(1), x).reshape(())\n\n x = jnp.ones((2,))\n variables = Foo().init(random.key(0), x)\n grad = jax.grad(Foo().apply)(variables, x)\n for grad_leaf in jax.tree_util.tree_leaves(grad):\n self.assertTrue(jnp.all(jnp.abs(grad_leaf) == 1.0))\n\n def test_transform_with_setup_and_methods_on_submodules(self):\n # This is the archetypal example motivating the introduction of\n # SetupState as a triple-enum to handle multiple setup() calls\n # across transform boundaries and scope reuse.\n class Foo(nn.Module):\n\n def setup(self):\n self.inner = nn.Dense(2)\n\n def helper(self, x, m):\n return m(x)\n\n def __call__(self, x):\n return self.helper(x, self.inner)\n\n k = random.key(0)\n x = jnp.ones((2,))\n vs_foo = Foo().init(k, x)\n\n class Bar(nn.Module):\n\n def setup(self):\n self.inner = nn.Dense(2)\n\n @nn.jit\n def helper(self, x, m):\n return m(x)\n\n @nn.jit\n def __call__(self, x):\n return self.helper(x, self.inner)\n\n vs_bar = Bar().init(k, x)\n self.assertTrue(\n tree_equals(\n jax.tree_util.tree_map(jnp.shape, vs_foo),\n jax.tree_util.tree_map(jnp.shape, vs_bar),\n )\n )\n\n def test_transform_methods_on_submodules_still_reserve_names(self):\n class Foo(nn.Module):\n\n @nn.jit\n def helper(self, x, m):\n conflicting_a = nn.Dense(2, name='a')\n return m(x)\n\n @nn.jit\n @nn.compact\n def __call__(self, x):\n a = nn.Dense(2, name='a')\n return self.helper(x, a)\n\n k = random.key(0)\n x = jnp.ones((2,))\n with self.assertRaises(errors.NameInUseError):\n vs = Foo().init(k, x)\n\n def test_transform_setup_still_reserve_names(self):\n class Identity(nn.Module):\n\n @nn.compact\n def __call__(self, x):\n return x\n\n class Test(nn.Module):\n\n def setup(self):\n self.sub = Identity()\n self.sub = Identity()\n\n @nn.jit\n def __call__(self, x):\n return x\n\n k = random.key(0)\n x = jnp.array([1.0])\n\n with self.assertRaises(errors.NameInUseError):\n y = Test().init(k, x)\n\n def test_transform_with_setup_and_methods_on_submodule_pytrees(self):\n class Foo(nn.Module):\n\n def setup(self):\n self.inners = [nn.Dense(2), nn.Dense(2)]\n\n def helper(self, x, ms):\n return ms[0](x) + ms[1](x)\n\n @nn.fold_rngs\n def __call__(self, x):\n return self.helper(x, self.inners)\n\n class JitFoo(nn.Module):\n\n def setup(self):\n self.inners = [nn.Dense(2), nn.Dense(2)]\n\n def helper(self, x, ms):\n return ms[0](x) + ms[1](x)\n\n @nn.jit\n def __call__(self, x):\n return self.helper(x, self.inners)\n\n k = random.key(0)\n x = jnp.ones((2,))\n\n vs_0 = Foo().init(k, x)\n vs_1 = JitFoo().init(k, x)\n\n self.assertTrue(tree_allclose(vs_0, vs_1))\n\n def test_transform_setup_still_reserve_names_pytrees(self):\n class Identity(nn.Module):\n\n @nn.compact\n def __call__(self, x):\n return x\n\n class Test(nn.Module):\n\n def setup(self):\n self.subs = [Identity(), Identity()]\n self.subs = [Identity(), Identity()]\n\n @nn.jit\n def __call__(self, x):\n return x\n\n k = random.key(0)\n x = jnp.array([1.0])\n\n msg = r'Could not create submodule \"subs_0\".*'\n with self.assertRaisesRegex(errors.NameInUseError, msg):\n y = Test().init(k, x)\n\n def test_scan_of_setup_parameter(self):\n class Body(nn.Module):\n\n def setup(self):\n self.dense = nn.Dense(1)\n self.p = self.param('p', lambda k: jnp.ones((1,)))\n\n def __call__(self, x):\n return self.dense(x) + self.p, None\n\n scanbody = nn.scan(\n Body, variable_axes={'params': 0}, split_rngs={'params': True}, length=2\n )\n k = random.key(0)\n x = jnp.ones((1,))\n vs = scanbody().init(k, x)\n y = scanbody().apply(vs, x)\n\n def test_multi_method_class_transform(self):\n class Foo(nn.Module):\n\n def setup(self):\n self.dense0 = nn.Dense(2)\n self.dense1 = nn.Dense(2)\n\n def method_0(self, x):\n return self.dense0(x), x\n\n def method_1(self, x, y):\n return self.dense1(x) + y, None\n\n class Bar(nn.Module):\n\n @nn.compact\n def __call__(self, x):\n ScanFoo = nn.scan(\n Foo,\n methods={\n 'method_0': dict(\n variable_axes={'params': 0},\n split_rngs={'params': True},\n in_axes=nn.broadcast,\n out_axes=0,\n length=3,\n ),\n 'method_1': dict(\n variable_axes={'params': 0},\n split_rngs={'params': True},\n in_axes=0,\n length=3,\n ),\n },\n )\n sf = ScanFoo()\n y, ys = sf.method_0(x)\n z, _ = sf.method_1(y, ys)\n return z\n\n k = random.key(0)\n x = random.uniform(random.key(1), (2, 2))\n vs = Bar().init(k, x)\n y = Bar().apply(vs, x)\n\n def test_compact_aliasing_collision(self):\n class Foo(nn.Module):\n m1: nn.Module\n m2: nn.Module\n\n @nn.compact\n def __call__(self, x):\n x = self.m2(self.m1(x))\n return x\n\n class Bar(nn.Module):\n\n @nn.compact\n def __call__(self, x):\n dense = nn.Dense(2)\n x = nn.jit(Foo)(dense, dense)(x)\n return x\n\n k = random.key(0)\n x = jnp.zeros((2, 2))\n _ = Bar().init(k, x)\n\n def test_compact_aliasing_collision_arg_and_attrib(self):\n class Foo(nn.Module):\n m1: nn.Module\n\n @nn.compact\n def __call__(self, x, m2):\n x = m2(self.m1(x))\n return x\n\n class Bar(nn.Module):\n\n @nn.compact\n def __call__(self, x):\n dense = nn.Dense(2)\n x = nn.jit(Foo)(dense)(x, dense)\n return x\n\n k = random.key(0)\n x = jnp.zeros((2, 2))\n _ = Bar().init(k, x)\n\n def test_jit_with_setup_helpers(self):\n class Foo(nn.Module):\n\n def setup(self):\n self.a = nn.Dense(2)\n self.setup_helper()\n\n def setup_helper(self):\n self.b = nn.Dense(2)\n\n @nn.fold_rngs\n def __call__(self, x):\n return self.b(self.a(x))\n\n class JitFoo(nn.Module):\n\n def setup(self):\n self.a = nn.Dense(2)\n self.setup_helper()\n\n def setup_helper(self):\n self.b = nn.Dense(2)\n\n @nn.jit\n def __call__(self, x):\n return self.b(self.a(x))\n\n k = random.key(0)\n x = jnp.ones((2, 2))\n vs = JitFoo().init(k, x)\n y0 = JitFoo().apply(vs, x)\n vs = Foo().init(k, x)\n y1 = Foo().apply(vs, x)\n np.testing.assert_array_equal(y0, y1)\n\n def test_jit_kwargs(self):\n class Foo(nn.Module):\n\n @nn.jit\n def __call__(self, a: jax.Array, b: jax.Array):\n return a + b\n\n m = Foo()\n y = m.apply({}, jnp.array(1.0), b=jnp.array(2.0))\n\n np.testing.assert_array_equal(y, jnp.array(3.0))\n\n def test_jit_static_argnames(self):\n s = None\n\n class Foo(nn.Module):\n\n @partial(nn.jit, static_argnames=['b'])\n def __call__(self, a: jax.Array, b: str):\n nonlocal s\n s = b\n return a\n\n m = Foo()\n y = m.apply({}, jnp.array(1.0), b='hi')\n\n self.assertEqual(s, 'hi')\n np.testing.assert_array_equal(y, jnp.array(1.0))\n\n def test_jit_and_sow(self):\n class Inner(nn.Module):\n\n @nn.compact\n def __call__(self, x):\n self.sow('intermediates', 'loss', jnp.sum(x))\n return x + 1\n\n class Outer(nn.Module):\n\n def setup(self):\n self.inner = Inner()\n\n @nn.jit\n def __call__(self, x):\n return self.inner(x)\n\n m = Outer()\n x = jnp.ones((2, 2))\n vs = m.init(random.key(0), x)\n y, updates = m.apply(vs, x, mutable=['intermediates'])\n np.testing.assert_array_equal(\n updates['intermediates']['inner']['loss'], 4.0\n )\n np.testing.assert_array_equal(y, 2)\n\n def test_fold_rngs(self):\n class Foo(nn.Module):\n\n def __call__(self, use_jit: bool):\n def f(foo: Foo):\n return foo.make_rng('params')\n\n if use_jit:\n key = nn.jit(f)(self)\n else:\n key = nn.fold_rngs(f)(self)\n\n return key\n\n foo = Foo()\n key_jit = foo.apply({}, True, rngs={'params': random.key(0)})\n key_fold_rngs = foo.apply({}, False, rngs={'params': random.key(0)})\n\n self.assert_keys_equal(key_jit, key_fold_rngs)\n\n def test_same_key(self):\n\n class Block(nn.Module):\n\n @nn.jit\n @nn.compact\n def __call__(self, carry, inputs):\n # dump_rng_info(self)\n key = self.make_rng('params')\n # y = jax.random.uniform(self.make_rng('params'), (2,))\n return carry, key\n\n class Transformer(nn.Module):\n\n @nn.compact\n def __call__(self):\n num_blocks = 10\n carry, key = nn.scan(\n Block,\n variable_axes={'params': 0},\n split_rngs={'params': True},\n # length=num_blocks,\n )()(None, jnp.arange(num_blocks))\n return key\n\n model = Transformer()\n keys1, _ = model.init_with_output(jax.random.key(1))\n keys2, _ = model.init_with_output(jax.random.key(1))\n keys3, _ = model.init_with_output(jax.random.key(1))\n keys4, _ = model.init_with_output(jax.random.key(1))\n\n self.assert_keys_equal(keys1, keys2)\n self.assert_keys_equal(keys2, keys3)\n self.assert_keys_equal(keys2, keys3)\n\n def test_jit_repr_hash(self):\n n = 0\n\n @partial(jax.jit, static_argnums=0)\n def f(obj):\n nonlocal n\n n += 1\n return None\n\n f(_HashableProxy.from_module(nn.Dense(10)))\n self.assertEqual(n, 1)\n f(_HashableProxy.from_module(nn.Dense(10)))\n self.assertEqual(n, 1)\n\n f(_HashableProxy.from_module(nn.Dense(20)))\n self.assertEqual(n, 2)\n f(_HashableProxy.from_module(nn.Dense(20)))\n self.assertEqual(n, 2)\n\n def test_jit_reuse(self):\n n = 0\n\n class Foo(nn.Module):\n\n @nn.jit\n def __call__(self, x):\n nonlocal n\n n += 1\n return x\n\n x = jnp.array(1.0)\n m = Foo()\n\n self.assertEqual(n, 0)\n\n y = m.apply({}, x)\n self.assertEqual(n, 1)\n y = m.apply({}, x)\n self.assertEqual(n, 1)\n\n def test_jit_recursive(self):\n n = 0\n\n class Foo(nn.Module):\n\n @partial(nn.jit, static_argnames='recurse_once')\n def __call__(self, x, *, recurse_once: bool = True):\n nonlocal n\n n += 1\n if recurse_once:\n x = self(x, recurse_once=False)\n return x + 1\n\n x = jnp.array(1.0)\n m = Foo()\n\n self.assertEqual(n, 0)\n\n y = m.apply({}, x)\n self.assertEqual(n, 2)\n y = m.apply({}, x)\n self.assertEqual(n, 2)\n\n @parameterized.named_parameters(('class', True), ('method', False))\n def test_jit_reuse_hash(self, jit_class: bool):\n n = 0\n\n class Foo(nn.Module):\n key: int\n\n def __call__(self, x):\n nonlocal n\n n += 1\n return x\n\n if jit_class:\n Foo = nn.jit(Foo)\n else: # jit method\n Foo.__call__ = nn.jit(Foo.__call__)\n\n x = jnp.array(1.0)\n self.assertEqual(n, 0)\n\n y = Foo(1).apply({}, x)\n self.assertEqual(n, 1)\n y = Foo(1).apply({}, x)\n self.assertEqual(n, 1)\n y = Foo(2).apply({}, x)\n self.assertEqual(n, 2)\n y = Foo(2).apply({}, x)\n self.assertEqual(n, 2)\n\n @parameterized.named_parameters(('class', True), ('method', False))\n def test_jit_reuse_submodules(self, jit_class: bool):\n test = self\n n = 0\n key = None\n name = None\n\n class Foo(nn.Module):\n key: int\n\n def __call__(self, x):\n nonlocal n, key, name\n n += 1\n key = self.key\n name = self.name\n return x\n\n if jit_class:\n Foo = nn.jit(Foo)\n else: # jit method\n Foo.__call__ = nn.jit(Foo.__call__)\n\n class Parent(nn.Module):\n\n @nn.compact\n def __call__(self, x):\n\n for i in range(3):\n m = Foo(i)\n y = m(x)\n test.assertEqual(key, i)\n test.assertEndsWith(name, f'Foo_{i}')\n test.assertEqual(n, i + 1)\n\n x = jnp.array(1.0)\n self.assertEqual(n, 0)\n\n y = Parent().apply({}, x)\n\n @parameterized.named_parameters(('class', True), ('method', False))\n def test_jit_stateful_submodules(self, jit_class: bool):\n n = 0\n\n class Foo(nn.Module):\n key: int\n\n @nn.compact\n def __call__(self, x):\n nonlocal n\n n += 1\n count = self.variable('counts', 'count', lambda: 0)\n if not self.is_initializing():\n count.value += 1\n return x\n\n if jit_class:\n Foo = nn.jit(Foo)\n else: # jit method\n Foo.__call__ = nn.jit(Foo.__call__)\n\n class Parent(nn.Module):\n @nn.compact\n def __call__(self, x):\n for _ in range(3):\n m = Foo(0)\n x = m(x)\n return x\n\n m = Parent()\n x = jnp.array(1.0)\n counts = m.init({}, x)['counts']\n self.assertEqual(n, 1)\n\n y, updates = m.apply({'counts': counts}, x, mutable=['counts'])\n counts = updates['counts']\n self.assertEqual(n, 2)\n\n for count in jax.tree.leaves(counts):\n self.assertEqual(count, 1)\n\n y, updates = m.apply({'counts': counts}, x, mutable=['counts'])\n counts = updates['counts']\n self.assertEqual(n, 2)\n\n for count in jax.tree.leaves(counts):\n self.assertEqual(count, 2)\n\n @parameterized.named_parameters(('class', True), ('method', False))\n def test_jit_reuse_nested_submodules(self, jit_class: bool):\n test = self\n n = 0\n\n class Foo(nn.Module):\n key: int\n\n def __call__(self, x):\n return x\n\n class Parent(nn.Module):\n submodules: list[Foo]\n\n @nn.compact\n def __call__(self, x):\n nonlocal n\n n += 1\n for i, m in enumerate(self.submodules):\n x = m(x)\n\n return x\n\n if jit_class:\n Parent = nn.jit(Parent)\n else: # jit method\n Parent.__call__ = nn.jit(Parent.__call__)\n\n x = jnp.array(1.0)\n self.assertEqual(n, 0)\n\n y = Parent([Foo(1), Foo(2)]).apply({}, x)\n self.assertEqual(n, 1)\n y = Parent([Foo(1), Foo(2)]).apply({}, x)\n self.assertEqual(n, 1)\n y = Parent([Foo(3), Foo(4)]).apply({}, x)\n self.assertEqual(n, 2)\n y = Parent([Foo(3), Foo(4)]).apply({}, x)\n self.assertEqual(n, 2)\n\n def test_jit_hashes_serializable_types(self):\n class Node:\n\n def __init__(self, a: int):\n self.a = a\n\n def __hash__(self):\n # test object is not being passed as static\n raise Exception('this should not be called')\n\n def __eq__(self, __value, /):\n raise Exception('this should not be called')\n\n def to_dict(node: Node):\n return {'a': node.a}\n\n def from_dict(node: Node, d: dict[str, Any]):\n node.a = d['a']\n return node\n\n serialization.register_serialization_state(Node, to_dict, from_dict)\n\n try:\n n = 0\n\n class Foo(nn.Module):\n node: Node\n\n @nn.jit\n @nn.compact\n def __call__(self, x):\n nonlocal n\n n += 1\n return self.node.a + nn.Dense(2)(x)\n\n m = Foo(Node(1))\n m.init_with_output(random.key(0), jnp.ones((2, 2)))\n self.assertEqual(n, 1)\n m.init_with_output(random.key(0), jnp.ones((2, 2)))\n self.assertEqual(n, 1)\n finally:\n del serialization._STATE_DICT_REGISTRY[Node]\n\n def test_while_loop(self):\n class Foo(nn.Module):\n\n @nn.compact\n def __call__(self, x):\n key_zero = random.key(0)\n key_zero = jnp.broadcast_to(key_zero, (2, *key_zero.shape))\n self.param('inc', lambda _: 1)\n self.put_variable('state', 'acc', 0)\n self.put_variable('state', 'rng_params', key_zero)\n self.put_variable('state', 'rng_loop', key_zero)\n\n def cond_fn(mdl, c):\n acc = mdl.get_variable('state', 'acc')\n return acc < x\n\n def body_fn(mdl, c):\n i = mdl.get_variable('state', 'acc')\n p_rng = mdl.make_rng('params')\n l_rng = mdl.make_rng('loop')\n mdl.put_variable(\n 'state',\n 'rng_params',\n mdl.get_variable('state', 'rng_params').at[i].set(p_rng),\n )\n mdl.put_variable(\n 'state',\n 'rng_loop',\n mdl.get_variable('state', 'rng_loop').at[i].set(l_rng),\n )\n inc = mdl.get_variable('params', 'inc')\n mdl.put_variable('state', 'acc', i + inc)\n return c\n\n return nn.while_loop(\n cond_fn,\n body_fn,\n self,\n (),\n carry_variables='state',\n split_rngs={'params': False, 'loop': True},\n )\n\n x = 2\n mdl = Foo()\n _, vars = mdl.apply(\n {},\n x,\n mutable=True,\n rngs={'params': random.key(1), 'loop': random.key(2)},\n )\n self.assertEqual(vars['state']['acc'], x)\n self.assertTrue(\n jnp.equal(vars['state']['rng_params'][0], vars['state']['rng_params'][1])\n )\n with jax.debug_key_reuse(False):\n self.assertFalse(\n jnp.equal(\n vars['state']['rng_loop'][0],\n vars['state']['rng_loop'][1],\n )\n )\n\n def test_while_loop_denylist_split_rngs(self):\n def cond_fn(module, carry):\n del module\n return carry < 10\n\n def body_fn(module, carry):\n rng = module.make_rng('random')\n return carry + 1\n\n class Foo(nn.Module):\n\n def setup(self):\n pass\n\n f = Foo().bind({}, rngs={'random': jax.random.PRNGKey(0)})\n init = jnp.zeros(())\n\n result = nn.while_loop(\n cond_fn,\n body_fn,\n f,\n init,\n split_rngs={\n 'params': False,\n nn.DenyList(('params',)): True,\n },\n )\n np.testing.assert_array_equal(result, jnp.array(10.0))\n\n def test_cond(self):\n class Foo(nn.Module):\n\n @nn.compact\n def __call__(self, x, pred):\n self.variable('state', 'true_count', lambda: 0)\n self.variable('state', 'false_count', lambda: 0)\n\n def true_fn(mdl, x):\n mdl.variable('state', 'true_count').value += 1\n return nn.Dense(2, name='dense')(x)\n\n def false_fn(mdl, x):\n mdl.variable('state', 'false_count').value += 1\n return -nn.Dense(2, name='dense')(x)\n\n return nn.cond(pred, true_fn, false_fn, self, x)\n\n def test_switch(self):\n class Foo(nn.Module):\n\n @nn.compact\n def __call__(self, x, pred):\n self.variable('state', 'a_count', lambda: 0)\n self.variable('state', 'b_count', lambda: 0)\n self.variable('state', 'c_count', lambda: 0)\n\n def a_fn(mdl, x):\n mdl.variable('state', 'a_count').value += 1\n return nn.Dense(2, name='dense')(x)\n\n def b_fn(mdl, x):\n mdl.variable('state', 'b_count').value += 1\n return -nn.Dense(2, name='dense')(x)\n\n def c_fn(mdl, x):\n mdl.variable('state', 'c_count').value += 1\n return nn.Dense(2, name='dense')(x)\n\n return nn.switch(pred, [a_fn, b_fn, c_fn], self, x)\n\n x = jnp.ones((1, 3))\n foo = Foo()\n y1, vars = foo.init_with_output(random.key(0), x, 0)\n self.assertEqual(vars['state'], {'a_count': 1, 'b_count': 0, 'c_count': 0})\n y2, updates = foo.apply(vars, x, 1, mutable='state')\n vars = copy(vars, updates)\n self.assertEqual(vars['state'], {'a_count': 1, 'b_count': 1, 'c_count': 0})\n np.testing.assert_allclose(y1, -y2)\n y3, updates = foo.apply(vars, x, 2, mutable='state')\n vars = copy(vars, updates)\n self.assertEqual(vars['state'], {'a_count': 1, 'b_count': 1, 'c_count': 1})\n np.testing.assert_allclose(y1, y3)\n\n def test_switch_multihead(self):\n class Foo(nn.Module):\n\n def setup(self) -> None:\n self.heads = [\n nn.Sequential([nn.Dense(10), nn.Dense(7), nn.Dense(5)]),\n nn.Sequential([nn.Dense(11), nn.Dense(5)]),\n nn.Dense(5),\n ]\n\n @nn.compact\n def __call__(self, x, index):\n def head_fn(i):\n def fn(mdl, x):\n mdl.variable('state', f'{i}_count', lambda: -1).value += 1\n return mdl.heads[i](x)\n\n return fn\n\n branches = [head_fn(i) for i in range(len(self.heads))]\n\n if self.is_mutable_collection('params'):\n for branch in branches:\n _ = branch(self, x)\n\n return nn.switch(index, branches, self, x)\n\n x = jnp.ones((1, 3))\n foo = Foo()\n y1, vars = foo.init_with_output(random.key(0), x, 0)\n self.assertEqual(vars['state'], {'0_count': 1, '1_count': 0, '2_count': 0})\n y2, updates = foo.apply(vars, x, 1, mutable='state')\n vars = copy(vars, updates)\n self.assertEqual(vars['state'], {'0_count': 1, '1_count': 1, '2_count': 0})\n y3, updates = foo.apply(vars, x, 2, mutable='state')\n vars = copy(vars, updates)\n self.assertEqual(vars['state'], {'0_count': 1, '1_count': 1, '2_count': 1})\n\n self.assertEqual(\n vars['params']['heads_0']['layers_0']['kernel'].shape, (3, 10)\n )\n self.assertEqual(vars['params']['heads_0']['layers_0']['bias'].shape, (10,))\n self.assertEqual(\n vars['params']['heads_0']['layers_1']['kernel'].shape, (10, 7)\n )\n self.assertEqual(vars['params']['heads_0']['layers_1']['bias'].shape, (7,))\n self.assertEqual(\n vars['params']['heads_0']['layers_2']['kernel'].shape, (7, 5)\n )\n self.assertEqual(vars['params']['heads_0']['layers_2']['bias'].shape, (5,))\n\n self.assertEqual(\n vars['params']['heads_1']['layers_0']['kernel'].shape, (3, 11)\n )\n self.assertEqual(vars['params']['heads_1']['layers_0']['bias'].shape, (11,))\n self.assertEqual(\n vars['params']['heads_1']['layers_1']['kernel'].shape, (11, 5)\n )\n self.assertEqual(vars['params']['heads_1']['layers_1']['bias'].shape, (5,))\n\n self.assertEqual(vars['params']['heads_2']['kernel'].shape, (3, 5))\n self.assertEqual(vars['params']['heads_2']['bias'].shape, (5,))\n\n def test_lift_instance_error(self):\n class Foo(nn.Module):\n\n @nn.compact\n def __call__(self, x):\n return nn.checkpoint(nn.Dense(2))(x)\n\n with self.assertRaises(errors.TransformTargetError):\n Foo().init(random.key(0), jnp.zeros((2, 3)))\n\n def test_scan_compact_count(self):\n class Foo(nn.Module):\n num_layers: int = 5\n\n @nn.compact\n def __call__(self, x):\n def body_fn(mdl, x):\n return nn.Dense(features=x.shape[-1])(x), ()\n\n x, _ = nn.scan(\n body_fn,\n length=self.num_layers,\n variable_axes={'params': 0},\n split_rngs={'params': True},\n )(self, x)\n return x\n\n m = Foo()\n x = jnp.ones((3,))\n v = m.init(jax.random.key(0), x)\n self.assertEqual(v['params']['Dense_0']['kernel'].shape, (5, 3, 3))\n m.apply(v, x)\n\n def test_bound_methods_in_direct_transforms(self):\n class CondModel(nn.Module):\n\n def setup(self):\n self.dense = nn.Dense(3)\n\n def f1(self, arr):\n arr = self.dense(arr)\n return arr\n\n def f2(self, arr):\n _ = self.dense(arr)\n return arr\n\n def __call__(self, x):\n return nn.cond(x.sum() > 0, self.f1, self.f2, self, x)\n\n cond_model = CondModel()\n\n output, init_params = jax.jit(cond_model.init_with_output)(\n jax.random.key(0), x=jnp.ones(3)\n )\n\n def test_add_metadata_axis(self):\n vars_copy = None\n\n class Foo(nn.Module):\n\n @nn.compact\n def __call__(self, x):\n nonlocal vars_copy\n kernel_init = nn.with_partitioning(\n nn.initializers.lecun_normal(), ('foo', 'bar')\n )\n vars_copy = self.variables\n return nn.Dense(\n 4, kernel_init=kernel_init, use_bias=False, name='dense'\n )(x)\n\n class Test(nn.Module):\n\n @partial(\n nn.add_metadata_axis,\n variable_axes={'params': 0},\n metadata_params={nn.PARTITION_NAME: 'baz'},\n )\n @nn.compact\n def __call__(self, x):\n return Foo(name='foo')(x)\n\n k = random.key(0)\n x = jnp.ones((4, 4), dtype=jnp.float32)\n vs = Test().init(k, x)\n y = Test().apply(vs, x)\n outer_expect = jax.tree_util.tree_map(\n jnp.shape,\n freeze({\n 'params': {\n 'foo': {\n 'dense': {\n 'kernel': nn.Partitioned(\n jnp.ones((4, 4)), names=('baz', 'foo', 'bar')\n )\n }\n }\n }\n }),\n )\n inner_expect = jax.tree_util.tree_map(\n jnp.shape,\n freeze({\n 'params': {\n 'dense': {\n 'kernel': nn.Partitioned(\n jnp.ones((4, 4)), names=('foo', 'bar')\n )\n }\n }\n }),\n )\n self.assertEqual(jax.tree_util.tree_map(jnp.shape, vs), outer_expect)\n self.assertEqual(jax.tree_util.tree_map(jnp.shape, vars_copy), inner_expect)\n\n def test_outer_setup_called_with_sharing_across_transforms(self):\n class A(nn.Module):\n\n def setup(self):\n self.foo = self.param('foo', nn.initializers.zeros, (2, 2), jnp.float32)\n\n def __call__(self, x):\n return self.foo\n\n class B(nn.Module):\n a: Any\n\n @nn.compact\n def __call__(self, x):\n return self.a(x)\n\n class C(nn.Module):\n\n def setup(self):\n self.a = A()\n self.b = nn.jit(B)(self.a)\n\n def __call__(self, x):\n b = self.b(x)\n a = self.a(x)\n return a + b\n\n k = random.key(0)\n x = random.randint(k, (2, 2), minval=0, maxval=10)\n vs = C().init(k, x)\n y = C().apply(vs, x)\n outer_expect = jax.tree_util.tree_map(\n jnp.shape, freeze({'params': {'a': {'foo': jnp.zeros((2, 2))}}})\n )\n self.assertEqual(jax.tree_util.tree_map(jnp.shape, vs), outer_expect)\n\n def test_grad_simple(self):\n class LearnScale(nn.Module):\n\n @nn.compact\n def __call__(self, x, y):\n p = self.param('scale', nn.initializers.ones_init(), ())\n return jnp.sum(p * x * y)\n\n class Foo(nn.Module):\n\n @nn.compact\n def __call__(self, x, y):\n x_grad, y_grad = nn.grad(\n lambda mdl, x, y: mdl(x, y), LearnScale(), x, y\n )\n return x_grad, y_grad\n\n x = random.uniform(random.key(1), (4,))\n y = random.uniform(random.key(2), (4,))\n vs = Foo().init(random.key(0), x, y)\n\n x_grad, y_grad = Foo().apply(vs, x, y)\n self.assertTrue(tree_allclose(x_grad, y))\n self.assertTrue(tree_allclose(y_grad, x))\n\n def test_grad_simple_with_aux(self):\n class LearnScale(nn.Module):\n\n @nn.compact\n def __call__(self, x, y):\n p = self.param('scale', nn.initializers.ones_init(), ())\n return jnp.sum(p * x * y), p\n\n class Foo(nn.Module):\n\n @nn.compact\n def __call__(self, x, y):\n (x_grad, y_grad), aux = nn.grad(\n lambda mdl, x, y: mdl(x, y), LearnScale(), x, y, has_aux=True\n )\n return aux, x_grad, y_grad\n\n x = random.uniform(random.key(1), (4,))\n y = random.uniform(random.key(2), (4,))\n vs = Foo().init(random.key(0), x, y)\n\n aux, x_grad, y_grad = Foo().apply(vs, x, y)\n self.assertTrue(tree_allclose(x_grad, y))\n self.assertTrue(tree_allclose(y_grad, x))\n self.assertTrue(tree_allclose(aux, vs['params']['LearnScale_0']['scale']))\n\n def test_value_and_grad_simple(self):\n class LearnScale(nn.Module):\n\n @nn.compact\n def __call__(self, x, y):\n p = self.param('scale', nn.initializers.ones_init(), ())\n return jnp.sum(p * x * y)\n\n class Foo(nn.Module):\n\n @nn.compact\n def __call__(self, x, y):\n z, (x_grad, y_grad) = nn.value_and_grad(\n lambda mdl, x, y: mdl(x, y), LearnScale(), x, y\n )\n return z, x_grad, y_grad\n\n x = random.uniform(random.key(1), (4,))\n y = random.uniform(random.key(2), (4,))\n vs = Foo().init(random.key(0), x, y)\n\n z, x_grad, y_grad = Foo().apply(vs, x, y)\n self.assertTrue(tree_allclose(x_grad, y))\n self.assertTrue(tree_allclose(y_grad, x))\n\n def test_value_and_grad_simple_with_aux(self):\n class LearnScale(nn.Module):\n\n @nn.compact\n def __call__(self, x, y):\n p = self.param('scale', nn.initializers.ones_init(), ())\n return jnp.sum(p * x * y), p\n\n class Foo(nn.Module):\n\n @nn.compact\n def __call__(self, x, y):\n (z, aux), (x_grad, y_grad) = nn.value_and_grad(\n lambda mdl, x, y: mdl(x, y), LearnScale(), x, y, has_aux=True\n )\n return z, aux, x_grad, y_grad\n\n x = random.uniform(random.key(1), (4,))\n y = random.uniform(random.key(2), (4,))\n vs = Foo().init(random.key(0), x, y)\n\n z, aux, x_grad, y_grad = Foo().apply(vs, x, y)\n self.assertTrue(tree_allclose(x_grad, y))\n self.assertTrue(tree_allclose(y_grad, x))\n self.assertTrue(tree_allclose(aux, vs['params']['LearnScale_0']['scale']))\n\n def test_value_and_grad_multiscope(self):\n class Foo(nn.Module):\n bar: nn.Module\n\n @nn.compact\n def __call__(self, x, y):\n def fn(self, x, y):\n qup = nn.Dense(y.shape[-1])\n delta = y - self.bar(qup(x))\n return jnp.sum(delta**2)\n\n z, (x_grad, y_grad) = nn.value_and_grad(fn, self, x, y)\n return z, x_grad, y_grad\n\n class Baz(nn.Module):\n\n @nn.compact\n def __call__(self, x, y):\n bar = nn.Dense(y.shape[-1])\n return Foo(bar=bar)(x, y)\n\n x = random.uniform(random.key(1), (4,))\n y = random.uniform(random.key(2), (4,))\n vs = Baz().init(random.key(0), x, y)\n z, x_grad, y_grad = Baz().apply(vs, x, y)\n\n def comparison_fn(x, y):\n w1 = vs['params']['Foo_0']['Dense_0']['kernel']\n w2 = vs['params']['Dense_0']['kernel']\n delta = y - jnp.dot(jnp.dot(x, w1), w2)\n return jnp.sum(delta**2)\n\n self.assertTrue(tree_allclose(comparison_fn(x, y), z))\n self.assertTrue(tree_allclose(jax.grad(comparison_fn, 0)(x, y), x_grad))\n self.assertTrue(tree_allclose(jax.grad(comparison_fn, 1)(x, y), y_grad))\n\n def test_value_and_grad_multiscope_adopted(self):\n class Foo(nn.Module):\n bar: nn.Module\n qup: nn.Module\n\n @nn.compact\n def __call__(self, x, y):\n def fn(self, x, y):\n delta = y - self.bar(self.qup(x))\n return jnp.sum(delta**2)\n\n z, (x_grad, y_grad) = nn.value_and_grad(fn, self, x, y)\n return z, x_grad, y_grad\n\n x = random.uniform(random.key(1), (4,))\n y = random.uniform(random.key(2), (4,))\n vs = Foo(bar=nn.Dense(4), qup=nn.Dense(4)).init(random.key(0), x, y)\n z, x_grad, y_grad = Foo(bar=nn.Dense(4), qup=nn.Dense(4)).apply(vs, x, y)\n\n def comparison_fn(x, y):\n w1 = vs['params']['qup']['kernel']\n w2 = vs['params']['bar']['kernel']\n delta = y - jnp.dot(jnp.dot(x, w1), w2)\n return jnp.sum(delta**2)\n\n self.assertTrue(tree_allclose(comparison_fn(x, y), z))\n self.assertTrue(tree_allclose(jax.grad(comparison_fn, 0)(x, y), x_grad))\n self.assertTrue(tree_allclose(jax.grad(comparison_fn, 1)(x, y), y_grad))\n\n def test_vmap_add_remove_axis_transforms(self):\n class BoxedData(struct.PyTreeNode, AxisMetadata):\n value: Any\n def unbox(self):\n return self.value\n def replace_boxed(self, val):\n return self.replace(value=val)\n def add_axis(self, index: int, params: dict[Any, Any]):\n value = jnp.mean(self.value, axis=index)\n return self.replace(value=value)\n def remove_axis(self, index: int, params: dict[Any, Any]):\n value_shape = list(self.value.shape)\n value_shape.insert(index, params['axis_size'])\n value = jnp.broadcast_to(self.value, value_shape)\n return self.replace(value=value)\n\n class Top(nn.Module):\n @nn.compact\n def __call__(self, x):\n VFoo = nn.vmap(\n Foo, in_axes=0, out_axes=0, variable_axes={'params':0, 'aux': 0},\n metadata_params={'axis_size': x.shape[0]},\n )\n vfoo = VFoo(name=\"vfoo\")\n y = vfoo(x)\n y = vfoo(x)\n assert vfoo.variables['aux']['v'].value.shape == ()\n return y\n\n class Foo(nn.Module):\n @nn.compact\n def __call__(self, x):\n if self.has_variable('aux', 'v'):\n assert self.variables['aux']['v'].value.shape == ()\n boxed_v = self.variable('aux', 'v', lambda: BoxedData(jnp.ones(())))\n assert self.variables['aux']['v'].value.shape == ()\n return x\n\n vs = Top().init(random.key(0), jnp.ones((2,5)))\n y = Top().apply(vs, jnp.ones((2, 5)))\n assert vs['aux']['vfoo']['v'].value.shape == ()\n\n def test_vjp_tracer_leak(self):\n class LearnScale(nn.Module):\n @nn.compact\n def __call__(self, x):\n p = self.param('scale', nn.initializers.zeros, ())\n return p * x\n class Foo(nn.Module):\n @nn.compact\n def __call__(self, x):\n y, bwd = nn.vjp(lambda mdl, x: mdl(x), LearnScale(), x)\n params_grad, x_grad = bwd(jnp.ones(y.shape))\n return y, params_grad, x_grad\n key = jax.random.PRNGKey(0)\n x = jnp.ones((2, 3))\n foo = Foo()\n with jax.check_tracer_leaks():\n params = foo.init(key, x)\n foo.apply(params, x)\n\n @parameterized.named_parameters(\n ('retracing scan', True), ('simple scan', False)\n )\n def test_jit_scan_retracing(self, retracing_scan: bool):\n num_blocks = 4\n num_patterns = 4\n features = 4\n trace_counts = [0, 0]\n\n class Block(nn.Module):\n def setup(self):\n self.dense = nn.Dense(features, use_bias=False)\n @nn.jit\n def __call__(self, x):\n nonlocal trace_counts\n trace_counts[1] += 1\n return self.dense(x)\n\n class BlockSequence(nn.Module):\n def setup(self):\n self.blocks = [Block() for _ in range(num_blocks)]\n @nn.jit\n def __call__(self, carry, inputs):\n nonlocal trace_counts\n trace_counts[0] += 1\n for block in self.blocks:\n carry = block(carry)\n return carry, inputs\n\n class Transformer(nn.Module):\n retracing_scan: bool = True\n def setup(self):\n self.scan = nn.scan(\n BlockSequence,\n variable_axes={'params': 0},\n split_rngs={'params': False},\n length=num_patterns,\n check_constancy_invariants=retracing_scan,\n )()\n def __call__(self, inputs):\n return self.scan(jnp.zeros_like(inputs), inputs)\n\n model = Transformer(retracing_scan=retracing_scan)\n _ = model.init(random.key(0), jnp.ones((num_patterns, features,)))\n self.assertEqual(trace_counts[0], 2 if retracing_scan else 1)\n self.assertEqual(trace_counts[1], 2 if retracing_scan else 1)\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/linen/partitioning_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for flax.linen.partitioning.\"\"\"\n\nimport jax\nimport jax.numpy as jnp\nfrom absl.testing import absltest, parameterized\nfrom jax import random, sharding\nfrom jax.experimental import mesh_utils\n\nfrom flax import linen as nn\nfrom flax.core import freeze, unfreeze\nfrom flax.linen import partitioning\n\nmock = absltest.mock\n\n# Parse absl flags test_srcdir and test_tmpdir.\njax.config.parse_flags_with_absl()\n\n# Testing constants.\nAXIS_RULES_1 = (('foo', 'data'), ('bar', 'model'), ('baz', None))\nAXIS_RULES_2 = (('foo', 'model'), ('bar', None), ('baz', 'data'))\n\n\nclass PartitioningTest(parameterized.TestCase):\n def test_axis_rules(self):\n self.assertEqual(nn.spmd.get_logical_axis_rules(), ())\n partitioning.set_axis_rules(AXIS_RULES_1)\n self.assertEqual(nn.spmd.get_logical_axis_rules(), AXIS_RULES_1)\n self.assertEqual(partitioning.get_axis_rules(), AXIS_RULES_1)\n partitioning.set_axis_rules(())\n\n def test_axis_rules_context(self):\n partitioning.set_axis_rules(AXIS_RULES_1)\n self.assertEqual(partitioning.get_axis_rules(), AXIS_RULES_1)\n with partitioning.axis_rules(AXIS_RULES_2):\n self.assertEqual(partitioning.get_axis_rules(), AXIS_RULES_2)\n self.assertEqual(partitioning.get_axis_rules(), AXIS_RULES_1)\n\n def test_logical_to_mesh_axes_resolves_to_none_or_unconstrained(self):\n unconstrained = jax.sharding.PartitionSpec.UNCONSTRAINED\n rules = (\n ('foo', None),\n ('bad', None),\n ('bar', unconstrained),\n ('baz', unconstrained),\n )\n self.assertEqual(\n partitioning.logical_to_mesh_axes(\n ('foo', 'bad', 'bar', 'baz'), rules=rules\n ),\n (None, None, unconstrained, unconstrained),\n )\n\n def test_logical_to_mesh_axes(self):\n axes_0 = ('foo', 'bar')\n # direct rule assignment\n self.assertEqual(\n partitioning.logical_to_mesh_axes(axes_0, rules=AXIS_RULES_1),\n ('data', 'model'),\n )\n\n # Repeated None and Unconstrained\n unconstrained = jax.sharding.PartitionSpec.UNCONSTRAINED\n axes_repeated = ('foo', unconstrained, unconstrained, None, None)\n self.assertEqual(\n partitioning.logical_to_mesh_axes(axes_repeated, rules=AXIS_RULES_1),\n ('data', unconstrained, unconstrained, None, None),\n )\n # axis rules context\n with partitioning.axis_rules(AXIS_RULES_1):\n self.assertEqual(\n partitioning.logical_to_mesh_axes(axes_0), ('data', 'model')\n )\n # nested context\n with partitioning.axis_rules(AXIS_RULES_2):\n self.assertEqual(\n partitioning.logical_to_mesh_axes(axes_0), ('model', None)\n )\n # duplicated logical names\n with partitioning.axis_rules(AXIS_RULES_1):\n with self.assertRaises(ValueError):\n partitioning.logical_to_mesh_axes(('foo', 'foo', 'baz'))\n\n def test_logical_to_mesh_axes_priorities(self):\n p_rules = (('foo', 'model'), ('bar', 'model'), ('baz', 'data'))\n with partitioning.axis_rules(p_rules):\n self.assertEqual(\n partitioning.logical_to_mesh_axes(('foo', 'bar', 'baz')),\n ('model', None, 'data'),\n )\n self.assertEqual(\n partitioning.logical_to_mesh_axes(('bar', 'foo', 'baz')),\n (None, 'model', 'data'),\n )\n self.assertEqual(\n partitioning.logical_to_mesh_axes(('baz', 'bar', 'foo')),\n ('data', None, 'model'),\n )\n self.assertEqual(\n partitioning.logical_to_mesh_axes(('baz', 'bar', 'foo', 'unassigned')),\n ('data', None, 'model', None),\n )\n\n @parameterized.parameters(\n dict(\n rules=(('a', ('model', 'data')), ('b', 'data')),\n axes=('a', 'b'),\n expected=(('model', 'data'), None),\n ),\n dict(\n rules=(('a', ('model', 'replica')), ('b', 'data')),\n axes=('a', 'b'),\n expected=(('model', 'replica'), 'data'),\n ),\n dict(\n rules=(('a', ('model', 'replica')), ('b', ('data', 'model'))),\n axes=('a', 'b'),\n expected=(('model', 'replica'), None),\n ),\n dict(\n rules=(('a', ('model', 'replica')), ('b', 'model')),\n axes=('a', 'b', 'c'),\n expected=(('model', 'replica'), None, None),\n ),\n dict(rules=(), axes=('a', 'b', 'c'), expected=(None, None, None)),\n dict(\n rules=(('a', None), ('a', 'model')),\n axes=('a', 'b'),\n expected=(None, None),\n ),\n dict(\n rules=(\n ('baz', 'data'),\n ('bar', None),\n ('foo', 'model'),\n ('foo', 'data'),\n ),\n axes=('baz', 'bar', 'foo'),\n expected=('data', None, 'model'),\n ),\n dict(\n rules=(('baz', 'data'), ('foo', ('model', 'emb'))),\n axes=('baz', jax.sharding.PartitionSpec.UNCONSTRAINED, 'foo'),\n expected=(\n 'data',\n jax.sharding.PartitionSpec.UNCONSTRAINED,\n ('model', 'emb'),\n ),\n ),\n )\n def test_logical_to_mesh_axes_cases(self, rules, axes, expected):\n with partitioning.axis_rules(rules):\n result = partitioning.logical_to_mesh_axes(axes)\n self.assertEqual(result, expected)\n\n @mock.patch('flax.linen.spmd._with_sharding_constraint')\n def test_with_sharding_constraint(self, wsc_fn):\n unconstrained = jax.sharding.PartitionSpec.UNCONSTRAINED\n arr = jnp.ones((2, 2))\n axes = ('foo', 'bar')\n partitioning.set_axis_rules(())\n _ = partitioning.with_sharding_constraint(arr, axes)\n wsc_fn.assert_not_called()\n with partitioning.axis_rules(AXIS_RULES_1):\n _ = partitioning.with_sharding_constraint(arr, None)\n wsc_fn.assert_not_called()\n _ = partitioning.with_sharding_constraint(arr, axes)\n wsc_fn.assert_called_with(\n arr, jax.sharding.PartitionSpec('data', 'model'), mesh=None\n )\n _ = partitioning.with_sharding_constraint(arr, ('foo', unconstrained))\n wsc_fn.assert_called_with(\n arr, jax.sharding.PartitionSpec('data', unconstrained), mesh=None\n )\n\n @mock.patch('flax.linen.spmd._with_sharding_constraint')\n def test_with_sharding_constraint_fallback(self, wsc_fn):\n arr = jnp.ones((2, 2))\n with partitioning.axis_rules(AXIS_RULES_1):\n _ = partitioning.with_sharding_constraint(arr, ('foo', 'not_recognized'))\n wsc_fn.assert_called_with(\n arr, jax.sharding.PartitionSpec('data', None), mesh=None\n )\n wsc_fn.reset_mock()\n _ = partitioning.with_sharding_constraint(\n arr,\n ('foo', 'not_recognized'),\n fallback=partitioning.RulesFallback.AXIS_IS_UNSHARDED,\n )\n wsc_fn.assert_called_with(\n arr, jax.sharding.PartitionSpec('data', None), mesh=None\n )\n wsc_fn.reset_mock()\n with self.assertRaises(ValueError):\n _ = partitioning.with_sharding_constraint(\n arr,\n ('foo', 'not_recognized'),\n fallback=partitioning.RulesFallback.RAISE_ERROR,\n )\n wsc_fn.assert_not_called()\n _ = partitioning.with_sharding_constraint(\n arr,\n ('foo', 'not_recognized'),\n fallback=partitioning.RulesFallback.NO_CONSTRAINT,\n )\n wsc_fn.assert_not_called()\n\n @parameterized.parameters(dict(axes_spec=None), dict(axes_spec=()))\n def test_param_with_axes_no_axes(self, axes_spec):\n class ParamTest(nn.Module):\n @nn.compact\n def __call__(self, x):\n foo = partitioning.param_with_axes(\n 'foo',\n lambda k, s, d: jnp.zeros(s, d),\n (2, 2),\n x.dtype,\n axes=axes_spec,\n )\n return x + foo\n\n k = random.key(0)\n x = jnp.ones((2, 2))\n _ = ParamTest().init(k, x)\n\n def test_param_with_axes(self):\n class ParamTest(nn.Module):\n @nn.compact\n def __call__(self, x):\n foo = partitioning.param_with_axes(\n 'foo',\n lambda k, s, d: jnp.zeros(s, d),\n (2, 2),\n x.dtype,\n axes=('foo', 'bar'),\n )\n return x + foo\n\n p_rules = (('foo', 'model'), ('bar', 'data'), ('baz', None))\n k = random.key(0)\n x = jnp.ones((2, 2))\n with partitioning.axis_rules(p_rules):\n variables = ParamTest().init(k, x)\n self.assertIn('params', variables)\n self.assertIn('params_axes', variables)\n self.assertEqual(\n variables['params_axes']['foo_axes'],\n partitioning.AxisMetadata(names=('foo', 'bar')),\n )\n logical_axis_names = partitioning.get_axis_names(variables['params_axes'])\n self.assertEqual(\n logical_axis_names, {'foo': jax.sharding.PartitionSpec('foo', 'bar')}\n )\n\n def test_param_pytree_with_axes(self):\n def init_fn(k, s, d):\n del k\n return {'a': jnp.zeros(s, d), 'b': (jnp.zeros(s, d), jnp.zeros(s, d))}\n\n axes = {'a': ('foo', 'bar'), 'b': (('foo', 'bar'), ('bar', 'foo'))}\n\n class ParamTest(nn.Module):\n @nn.compact\n def __call__(self, x):\n foo = partitioning.param_with_axes(\n 'foo', init_fn, (2, 2), x.dtype, axes=axes\n )\n return x + foo['a']\n\n p_rules = (('foo', 'model'), ('bar', 'data'), ('baz', None))\n k = random.key(0)\n x = jnp.ones((2, 2))\n with partitioning.axis_rules(p_rules):\n variables = ParamTest().init(k, x)\n self.assertIn('params', variables)\n self.assertIn('params_axes', variables)\n self.assertEqual(\n variables['params_axes']['foo_axes'],\n partitioning.AxisMetadata(names=axes),\n )\n logical_axis_names = partitioning.get_axis_names(variables['params_axes'])\n expected = freeze(\n {\n 'foo': {\n 'a': jax.sharding.PartitionSpec('foo', 'bar'),\n 'b': (\n jax.sharding.PartitionSpec('foo', 'bar'),\n jax.sharding.PartitionSpec('bar', 'foo'),\n ),\n }\n }\n )\n self.assertEqual(logical_axis_names, expected)\n\n @parameterized.parameters(dict(axes_spec=None), dict(axes_spec=()))\n def test_variable_with_axes_no_axes(self, axes_spec):\n class VarTest(nn.Module):\n @nn.compact\n def __call__(self, x):\n foo = partitioning.variable_with_axes(\n 'test', 'foo', jnp.zeros, (2, 2), x.dtype, axes=axes_spec\n )\n return x + foo.value\n\n k = random.key(0)\n x = jnp.ones((2, 2))\n _ = VarTest().init(k, x)\n\n def test_variable_with_empty_tuple_has_empty_axes(self):\n class VarTest(nn.Module):\n @nn.compact\n def __call__(self, x):\n foo = partitioning.variable_with_axes(\n 'test', 'foo', jnp.zeros, (2, 2), x.dtype, axes=()\n )\n return x + foo.value\n\n k = random.key(0)\n x = jnp.ones((2, 2))\n variables = VarTest().init(k, x)\n logical_axis_names = partitioning.get_axis_names(variables['test_axes'])\n self.assertEqual(logical_axis_names, {'foo': jax.sharding.PartitionSpec()})\n\n def test_variable_with_axes(self):\n class VarTest(nn.Module):\n @nn.compact\n def __call__(self, x):\n foo = partitioning.variable_with_axes(\n 'test', 'foo', jnp.zeros, (2, 2), x.dtype, axes=('foo', 'bar')\n )\n return x + foo.value\n\n p_rules = (('foo', 'model'), ('bar', 'data'), ('baz', None))\n k = random.key(0)\n x = jnp.ones((2, 2))\n with partitioning.axis_rules(p_rules):\n variables = VarTest().init(k, x)\n self.assertIn('test', variables)\n self.assertIn('test_axes', variables)\n self.assertEqual(\n variables['test_axes']['foo_axes'],\n partitioning.AxisMetadata(names=('foo', 'bar')),\n )\n logical_axis_names = partitioning.get_axis_names(variables['test_axes'])\n self.assertEqual(\n logical_axis_names, {'foo': jax.sharding.PartitionSpec('foo', 'bar')}\n )\n\n @mock.patch('flax.linen.partitioning._with_sharding_constraint')\n def test_variable_with_axes_fallback(self, wsc_fn):\n class VarTest(nn.Module):\n @nn.compact\n def __call__(self, x):\n foo = partitioning.variable_with_axes(\n 'test',\n 'foo',\n jnp.zeros,\n (2, 2),\n x.dtype,\n axes=('foo', 'bar'),\n fallback=partitioning.RulesFallback.NO_CONSTRAINT,\n )\n return x + foo.value\n\n p_rules = (\n # No rule for 'foo':\n ('bar', 'data'),\n ('baz', None),\n )\n k = random.key(0)\n x = jnp.ones((2, 2))\n with partitioning.axis_rules(p_rules):\n variables = VarTest().init(k, x)\n\n wsc_fn.assert_not_called()\n self.assertIn('test', variables)\n self.assertIn('test_axes', variables)\n self.assertEqual(\n variables['test_axes']['foo_axes'],\n partitioning.AxisMetadata(names=('foo', 'bar')),\n )\n logical_axis_names = partitioning.get_axis_names(variables['test_axes'])\n self.assertEqual(\n logical_axis_names, {'foo': jax.sharding.PartitionSpec('foo', 'bar')}\n )\n\n def test_scan_with_axes(self):\n # MLP Hparams\n B, L, E = 8, 4, 32 # pylint: disable=invalid-name\n # fake inputs\n x = jnp.ones((B, E))\n k = random.key(0)\n\n class SinDot(nn.Module):\n depth: int\n\n @nn.compact\n def __call__(self, x):\n W1 = partitioning.param_with_axes( # pylint: disable=invalid-name\n 'W1',\n nn.initializers.xavier_normal(),\n (x.shape[-1], self.depth),\n axes=('emb', 'mlp'),\n )\n W2 = partitioning.param_with_axes( # pylint: disable=invalid-name\n 'W2',\n nn.initializers.xavier_normal(),\n (self.depth, x.shape[-1]),\n axes=('mlp', 'emb'),\n )\n y = jnp.dot(jnp.sin(jnp.dot(x, W1)), W2)\n _ = partitioning.variable_with_axes(\n 'stats', 'y_st', lambda: y, axes=('batch', 'emb')\n )\n # scan expects a (carry, out) return signature.\n return y, None\n\n class Scanned(nn.Module):\n num_layers: int\n depth: int\n\n @nn.compact\n def __call__(self, x):\n scanned_sindot = partitioning.scan_with_axes(\n SinDot,\n in_axes=(),\n variable_axes={'params': 0, 'stats': 1},\n split_rngs={'params': True},\n axis_name='layer',\n axes_collections=('params', 'stats'),\n length=self.num_layers,\n )(self.depth, name='scanned_layer')\n y, _ = scanned_sindot(x)\n # test calling again to test metadata compatibility across calls\n _, _ = scanned_sindot(x)\n return y\n\n p_rules = (('emb', 'data'), ('mlp', 'model'), ('batch', 'data'))\n with partitioning.axis_rules(p_rules):\n variables = Scanned(L, E).init(k, x)\n\n # Ensure that the module can be called when 'params_axes' is not mutable.\n Scanned(L, E).apply(variables, x)\n self.assertIn('params', variables)\n self.assertIn('params_axes', variables)\n self.assertIn('stats', variables)\n self.assertIn('stats_axes', variables)\n self.assertEqual(\n variables['params_axes']['scanned_layer']['W1_axes'],\n partitioning.AxisMetadata(names=('layer', 'emb', 'mlp')),\n )\n logical_axis_names = partitioning.get_axis_names(variables['params_axes'])\n self.assertEqual(\n logical_axis_names,\n {\n 'scanned_layer': {\n 'W1': jax.sharding.PartitionSpec('layer', 'emb', 'mlp'),\n 'W2': jax.sharding.PartitionSpec('layer', 'mlp', 'emb'),\n }\n },\n )\n logical_axis_names = partitioning.get_axis_names(variables['stats_axes'])\n self.assertEqual(\n logical_axis_names,\n {\n 'scanned_layer': {\n 'y_st': jax.sharding.PartitionSpec('batch', 'layer', 'emb')\n }\n },\n )\n\n def test_vmap_with_axes(self):\n class Foo(nn.Module):\n @nn.compact\n def __call__(self, x):\n return (\n partitioning.param_with_axes(\n 'w', jax.nn.initializers.uniform(), [4, 3], axes=('out', 'in')\n )\n @ x\n )\n\n class Vmapped(nn.Module):\n @nn.compact\n def __call__(self, x):\n FooVmapped = partitioning.vmap_with_axes( # pylint: disable=invalid-name\n Foo,\n variable_axes={\n 'params': 1,\n },\n split_rngs={'params': True},\n partitioning_axis_names={'params': 'vmap_axis'},\n )\n return FooVmapped(name='foo_vmapped')(x)\n\n p_rules = (('out', None), ('in', 'data'), ('vmap_axis', 'model'))\n\n # check that regular Food module is correct\n with partitioning.axis_rules(p_rules):\n variables = Foo().init(jax.random.key(0), jnp.array([1, 2, 3]))\n variables = unfreeze(variables)\n variables['params'] = jax.tree_util.tree_map(\n lambda x: x.shape, variables['params']\n )\n self.assertDictEqual(\n variables,\n {\n 'params': {'w': (4, 3)},\n 'params_axes': {\n 'w_axes': partitioning.AxisMetadata(names=('out', 'in'))\n },\n },\n )\n\n # check that FooVmapped adds 'vmap_axis' to axis 1\n with partitioning.axis_rules(p_rules):\n variables = Vmapped().init(\n jax.random.key(0), jnp.array([[1, 2, 3], [4, 5, 6]])\n )\n variables = unfreeze(variables)\n variables['params'] = jax.tree_util.tree_map(\n lambda x: x.shape, variables['params']\n )\n self.assertDictEqual(\n variables,\n {\n 'params': {'foo_vmapped': {'w': (4, 2, 3)}},\n 'params_axes': {\n 'foo_vmapped': {\n 'w_axes': partitioning.AxisMetadata(\n names=('out', 'vmap_axis', 'in')\n )\n }\n },\n },\n )\n\n def test_logical_with_mesh_and_rules(self):\n devices = mesh_utils.create_device_mesh((jax.local_device_count(), 1))\n mesh = sharding.Mesh(devices, ('in', 'out'))\n test = self\n rules = (('a', 'in'), ('b', 'out'))\n\n class Foo(nn.Module):\n @nn.compact\n def __call__(self, x):\n kernel_init = nn.with_logical_partitioning(\n nn.initializers.ones_init(), ('a', 'b'), mesh=mesh, rules=rules\n )\n kernel = self.param('kernel', kernel_init, (x.shape[-1], 2))\n kernel_box = self.get_variable('params', 'kernel')\n test.assertIsInstance(kernel_box, nn.Partitioned)\n test.assertEqual(kernel_box.names, ('a', 'b'))\n return x @ kernel\n\n @jax.jit\n def create_state():\n module = Foo()\n variables = module.init(random.key(0), jnp.zeros((8, 4)))\n logical_spec = nn.get_partition_spec(variables)\n shardings = nn.logical_to_mesh_sharding(logical_spec, mesh, rules)\n variables = jax.lax.with_sharding_constraint(variables, shardings)\n return variables\n\n variables = create_state()\n self.assertEqual(variables['params']['kernel'].names, ('a', 'b'))\n self.assertIs(variables['params']['kernel'].mesh, mesh)\n self.assertEqual(variables['params']['kernel'].rules, rules)\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/linen/summary_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport enum\n\nimport jax\nimport jax.numpy as jnp\nimport numpy as np\nfrom absl.testing import absltest\nfrom jax import random\n\nfrom flax import linen as nn\nfrom flax import struct\nfrom flax.core.scope import Array\nfrom flax.linen import summary\n\n# Parse absl flags test_srcdir and test_tmpdir.\njax.config.parse_flags_with_absl()\n\nCONSOLE_TEST_KWARGS = dict(force_terminal=False, no_color=True, width=10_000)\n\n\ndef _get_shapes(pytree):\n return jax.tree_util.tree_map(\n lambda x: x.shape if hasattr(x, 'shape') else x, pytree\n )\n\n\ndef _get_obj_repr_value(x):\n if isinstance(x, summary._ObjectRepresentation):\n return x.obj\n return x\n\n\nclass ConvBlock(nn.Module):\n features: int\n kernel_size: list[int]\n test_sow: bool\n\n def setup(self) -> None:\n self.conv = nn.Conv(self.features, self.kernel_size)\n self.bn = nn.BatchNorm()\n self.dropout = nn.Dropout(0.5)\n\n def block_method(self, x: Array, training: bool) -> Array:\n x = self.conv(x)\n\n if self.test_sow:\n self.sow('intermediates', 'INTERM', x)\n\n x = self.bn(x, use_running_average=not training)\n x = self.dropout(x, deterministic=not training)\n x = nn.relu(x)\n return x\n\n def __call__(self, x: Array, training: bool) -> Array:\n x = self.conv(x)\n\n if self.test_sow:\n self.sow('intermediates', 'INTERM', x)\n\n x = self.bn(x, use_running_average=not training)\n x = self.dropout(x, deterministic=not training)\n x = nn.relu(x)\n return x\n\n\nclass CNN(nn.Module):\n test_sow: bool\n\n def setup(self) -> None:\n self.block1 = ConvBlock(32, [3, 3], test_sow=self.test_sow)\n self.block2 = ConvBlock(64, [3, 3], test_sow=self.test_sow)\n self.dense = nn.Dense(10)\n\n def cnn_method(self, x: Array, training: bool) -> Array:\n x = self.block1.block_method(x, training=training)\n x = self.block2.block_method(x, training=training)\n x = x.mean(axis=(1, 2))\n\n if self.test_sow:\n self.sow('intermediates', 'INTERM', x)\n\n x = self.dense(x)\n\n return x, dict(a=x, b=x + 1.0)\n\n def __call__(self, x: Array, training: bool) -> Array:\n x = self.block1.block_method(x, training=training)\n x = self.block2.block_method(x, training=training)\n x = x.mean(axis=(1, 2))\n\n if self.test_sow:\n self.sow('intermediates', 'INTERM', x)\n\n x = self.dense(x)\n\n return x, dict(a=x, b=x + 1.0)\n\n\nclass SummaryTest(absltest.TestCase):\n def test_module_summary(self):\n \"\"\"\n This test creates a Table using `module_summary` and checks that it\n matches the expected output given the CNN model defined in\n `_get_tabulate_cnn`.\n \"\"\"\n\n batch_size = 32\n\n x = jnp.ones((batch_size, 28, 28, 1))\n module = CNN(test_sow=False)\n\n table = summary._get_module_table(\n module,\n depth=None,\n show_repeated=True,\n compute_flops=True,\n compute_vjp_flops=True,\n )(\n {'dropout': random.key(0), 'params': random.key(1)},\n x,\n training=True,\n mutable=True,\n )\n # get values for inputs and outputs from their _ValueRepresentation\n for row in table:\n row.inputs = jax.tree_util.tree_map(_get_obj_repr_value, row.inputs)\n row.outputs = jax.tree_util.tree_map(_get_obj_repr_value, row.outputs)\n\n # 10 rows = 1 CNN + 4 ConvBlock_0 + 4 ConvBlock_1 + 1 Dense_0\n self.assertLen(table, 10)\n\n # check paths\n self.assertEqual(table[0].path, ())\n\n self.assertEqual(table[1].path, ('block1',))\n self.assertEqual(table[2].path, ('block1', 'conv'))\n self.assertEqual(table[3].path, ('block1', 'bn'))\n self.assertEqual(table[4].path, ('block1', 'dropout'))\n\n self.assertEqual(table[5].path, ('block2',))\n self.assertEqual(table[6].path, ('block2', 'conv'))\n self.assertEqual(table[7].path, ('block2', 'bn'))\n self.assertEqual(table[8].path, ('block2', 'dropout'))\n\n self.assertEqual(table[9].path, ('dense',))\n\n # check outputs shapes\n self.assertEqual(\n (table[0].inputs[0].shape, table[0].inputs[1]),\n (x.shape, dict(training=True)),\n )\n self.assertEqual(\n _get_shapes(table[0].outputs),\n ((batch_size, 10), dict(a=(batch_size, 10), b=(batch_size, 10))),\n )\n\n self.assertEqual(\n _get_shapes(table[1].inputs),\n ((batch_size, 28, 28, 1), {'training': True}),\n )\n self.assertEqual(table[1].outputs.shape, (batch_size, 28, 28, 32))\n self.assertEqual(table[2].inputs.shape, (batch_size, 28, 28, 1))\n self.assertEqual(table[2].outputs.shape, (batch_size, 28, 28, 32))\n self.assertEqual(\n _get_shapes(table[3].inputs),\n ((batch_size, 28, 28, 32), {'use_running_average': False}),\n )\n self.assertEqual(table[3].outputs.shape, (batch_size, 28, 28, 32))\n self.assertEqual(\n _get_shapes(table[4].inputs),\n ((batch_size, 28, 28, 32), {'deterministic': False}),\n )\n self.assertEqual(table[4].outputs.shape, (batch_size, 28, 28, 32))\n\n self.assertEqual(\n _get_shapes(table[5].inputs),\n ((batch_size, 28, 28, 32), {'training': True}),\n )\n self.assertEqual(table[5].outputs.shape, (batch_size, 28, 28, 64))\n self.assertEqual(table[6].inputs.shape, (batch_size, 28, 28, 32))\n self.assertEqual(table[6].outputs.shape, (batch_size, 28, 28, 64))\n self.assertEqual(\n _get_shapes(table[7].inputs),\n ((batch_size, 28, 28, 64), {'use_running_average': False}),\n )\n self.assertEqual(table[7].outputs.shape, (batch_size, 28, 28, 64))\n self.assertEqual(\n _get_shapes(table[8].inputs),\n ((batch_size, 28, 28, 64), {'deterministic': False}),\n )\n self.assertEqual(table[8].outputs.shape, (batch_size, 28, 28, 64))\n\n self.assertEqual(table[9].inputs.shape, (batch_size, 64))\n self.assertEqual(table[9].outputs.shape, (batch_size, 10))\n\n # check no summary is performed\n for row in table:\n self.assertEqual(\n row.module_variables,\n row.counted_variables,\n )\n\n # Each module FLOPs >= sum of its submodule FLOPs.\n # Can be greater due to ops like `nn.relu` not belonging to any submodule.\n for r in table:\n flops, vjp_flops = r.flops, r.vjp_flops\n submodule_flops, submodule_vjp_flops = 0, 0\n for s in table:\n if len(s.path) == len(r.path) + 1 and s.path[: len(r.path)] == r.path:\n submodule_flops += s.flops\n submodule_vjp_flops += s.vjp_flops\n\n self.assertGreaterEqual(flops, submodule_flops)\n self.assertGreaterEqual(vjp_flops, submodule_vjp_flops)\n\n def test_module_summary_with_depth(self):\n \"\"\"\n This test creates a Table using `module_summary` set the `depth` argument\n to `1`, table should have fewer rows as a consequence.\n \"\"\"\n batch_size = 32\n\n x = jnp.ones((batch_size, 28, 28, 1))\n module = CNN(test_sow=False)\n\n table = summary._get_module_table(\n module,\n depth=1,\n show_repeated=True,\n compute_flops=True,\n compute_vjp_flops=True,\n )(\n {'dropout': random.key(0), 'params': random.key(1)},\n x,\n training=True,\n mutable=True,\n )\n # get values for inputs and outputs from their _ValueRepresentation\n\n for row in table:\n row.inputs = jax.tree_util.tree_map(_get_obj_repr_value, row.inputs)\n row.outputs = jax.tree_util.tree_map(_get_obj_repr_value, row.outputs)\n\n # 4 rows = 1 CNN + 1 ConvBlock_0 + 1 ConvBlock_1 + 1 Dense_0\n self.assertLen(table, 4)\n\n # check paths\n self.assertEqual(table[0].path, ())\n\n self.assertEqual(table[1].path, ('block1',))\n self.assertEqual(table[2].path, ('block2',))\n self.assertEqual(table[3].path, ('dense',))\n\n # check outputs shapes\n self.assertEqual(\n (table[0].inputs[0].shape, table[0].inputs[1]),\n (x.shape, dict(training=True)),\n )\n self.assertEqual(\n _get_shapes(table[0].outputs),\n ((batch_size, 10), dict(a=(batch_size, 10), b=(batch_size, 10))),\n )\n\n self.assertEqual(\n _get_shapes(table[1].inputs),\n ((batch_size, 28, 28, 1), {'training': True}),\n )\n self.assertEqual(table[1].outputs.shape, (batch_size, 28, 28, 32))\n\n self.assertEqual(\n _get_shapes(table[2].inputs),\n ((batch_size, 28, 28, 32), {'training': True}),\n )\n self.assertEqual(table[2].outputs.shape, (batch_size, 28, 28, 64))\n\n self.assertEqual(table[3].inputs.shape, (batch_size, 64))\n self.assertEqual(table[3].outputs.shape, (batch_size, 10))\n\n # check ConvBlock_0 and ConvBlock_1 are summarized\n self.assertNotEqual(table[1].module_variables, table[1].counted_variables)\n self.assertNotEqual(table[2].module_variables, table[2].counted_variables)\n\n # check CNN and Dense_0 output are not summarized\n self.assertEqual(table[0].module_variables, table[0].counted_variables)\n self.assertEqual(table[3].module_variables, table[3].counted_variables)\n\n # Top level FLOPs > sum of listed submodule FLOPs, since not all are listed.\n self.assertGreater(table[0].flops, sum(r.flops for r in table[1:]))\n self.assertGreater(table[0].vjp_flops, sum(r.vjp_flops for r in table[1:]))\n\n def test_tabulate(self):\n \"\"\"\n This test creates a string representation of a Module using `Module.tabulate`\n and checks that it matches the expected output given the CNN model defined\n in `_get_tabulate_cnn`.\n \"\"\"\n batch_size = 32\n\n x = jnp.ones((batch_size, 28, 28, 1))\n module = CNN(test_sow=False)\n\n module_repr = module.tabulate(\n {'dropout': random.key(0), 'params': random.key(1)},\n x,\n training=True,\n console_kwargs=CONSOLE_TEST_KWARGS,\n compute_flops=True,\n compute_vjp_flops=True,\n )\n\n # NOTE: it's tricky to validate the content of lines\n # because it seems to be shell-dependent, so we will\n # just check lines that won't change between environments\n lines = module_repr.split('\\n')\n\n # check title\n module_name = module.__class__.__name__\n self.assertIn(f'{module_name} Summary', lines[1])\n\n # check headers are correct\n self.assertIn('path', lines[3])\n self.assertIn('module', lines[3])\n self.assertIn('inputs', lines[3])\n self.assertIn('outputs', lines[3])\n self.assertIn('params', lines[3])\n self.assertIn('flops', lines[3])\n self.assertIn('vjp_flops', lines[3])\n self.assertIn('batch_stats', lines[3])\n\n # collection counts\n self.assertIn('Total', lines[-6])\n self.assertIn('192', lines[-6])\n self.assertIn('768 B', lines[-6])\n self.assertIn('19,658', lines[-6])\n self.assertIn('78.6 KB', lines[-6])\n\n # total counts\n self.assertIn('Total Parameters', lines[-3])\n self.assertIn('19,850', lines[-3])\n self.assertIn('79.4 KB', lines[-3])\n\n def test_tabulate_with_sow(self):\n batch_size = 32\n\n x = jnp.ones((batch_size, 28, 28, 1))\n module = CNN(test_sow=True)\n\n module_repr = module.tabulate(\n {'dropout': random.key(0), 'params': random.key(1)},\n x,\n training=True,\n mutable=True,\n console_kwargs=CONSOLE_TEST_KWARGS,\n compute_flops=True,\n compute_vjp_flops=True,\n )\n\n self.assertIn('intermediates', module_repr)\n self.assertIn('INTERM', module_repr)\n\n def test_tabulate_with_method(self):\n batch_size = 32\n\n x = jnp.ones((batch_size, 28, 28, 1))\n module = CNN(test_sow=False)\n\n module_repr = module.tabulate(\n {'dropout': random.key(0), 'params': random.key(1)},\n x,\n training=True,\n method=CNN.cnn_method,\n console_kwargs=CONSOLE_TEST_KWARGS,\n compute_flops=True,\n compute_vjp_flops=True,\n )\n\n self.assertIn('(block_method)', module_repr)\n self.assertIn('(cnn_method)', module_repr)\n\n def test_tabulate_function(self):\n \"\"\"\n This test creates a string representation of a Module using\n `Module.tabulate` and checks that it matches the expected output given the\n CNN model defined in `_get_tabulate_cnn`.\n \"\"\"\n batch_size = 32\n\n x = jnp.ones((batch_size, 28, 28, 1))\n module = CNN(test_sow=False)\n\n module_repr = nn.tabulate(\n module,\n {'dropout': random.key(0), 'params': random.key(1)},\n console_kwargs=CONSOLE_TEST_KWARGS,\n compute_flops=True,\n compute_vjp_flops=True,\n )(x, training=True)\n\n lines = module_repr.split('\\n')\n\n # check title\n module_name = module.__class__.__name__\n self.assertIn(f'{module_name} Summary', lines[1])\n\n # check headers are correct\n self.assertIn('path', lines[3])\n self.assertIn('module', lines[3])\n self.assertIn('inputs', lines[3])\n self.assertIn('outputs', lines[3])\n self.assertIn('params', lines[3])\n self.assertIn('flops', lines[3])\n self.assertIn('batch_stats', lines[3])\n\n # collection counts\n self.assertIn('Total', lines[-6])\n self.assertIn('192', lines[-6])\n self.assertIn('768 B', lines[-6])\n self.assertIn('19,658', lines[-6])\n self.assertIn('78.6 KB', lines[-6])\n\n # total counts\n self.assertIn('Total Parameters', lines[-3])\n self.assertIn('19,850', lines[-3])\n self.assertIn('79.4 KB', lines[-3])\n\n def test_lifted_transform(self):\n class LSTM(nn.Module):\n features: int\n\n @nn.compact\n def __call__(self, x):\n carry = nn.LSTMCell(self.features).initialize_carry(\n random.key(0), x[:, 0].shape\n )\n ScanLSTM = nn.scan(\n nn.LSTMCell,\n variable_broadcast='params',\n split_rngs={'params': False},\n in_axes=1,\n out_axes=1,\n )\n return ScanLSTM(self.features, name='ScanLSTM')(carry, x)\n\n lstm = LSTM(features=128)\n\n with jax.check_tracer_leaks(True):\n module_repr = lstm.tabulate(\n random.key(0),\n x=jnp.ones((32, 128, 64)),\n console_kwargs=CONSOLE_TEST_KWARGS,\n compute_flops=True,\n compute_vjp_flops=True,\n )\n\n lines = module_repr.splitlines()\n\n self.assertIn('LSTM', lines[5])\n self.assertIn('ScanLSTM', lines[9])\n self.assertIn('LSTMCell', lines[9])\n self.assertIn('ScanLSTM/ii', lines[13])\n self.assertIn('Dense', lines[13])\n\n def test_lifted_transform_no_rename(self):\n class LSTM(nn.Module):\n features: int\n\n @nn.compact\n def __call__(self, x):\n carry = nn.LSTMCell(self.features).initialize_carry(\n random.key(0), x[:, 0].shape\n )\n ScanLSTM = nn.scan(\n nn.LSTMCell,\n variable_broadcast='params',\n split_rngs={'params': False},\n in_axes=1,\n out_axes=1,\n )\n return ScanLSTM(self.features)(carry, x)\n\n lstm = LSTM(features=128)\n\n with jax.check_tracer_leaks(True):\n module_repr = lstm.tabulate(\n random.key(0),\n x=jnp.ones((32, 128, 64)),\n console_kwargs=CONSOLE_TEST_KWARGS,\n compute_flops=True,\n compute_vjp_flops=True,\n )\n\n lines = module_repr.splitlines()\n\n self.assertIn('LSTM', lines[5])\n self.assertIn('ScanLSTMCell_0', lines[9])\n self.assertIn('LSTMCell', lines[9])\n self.assertIn('ScanLSTMCell_0/ii', lines[13])\n self.assertIn('Dense', lines[13])\n\n def test_module_reuse(self):\n class ConvBlock(nn.Module):\n @nn.compact\n def __call__(self, x):\n x = nn.Conv(32, [3, 3])(x)\n x = nn.BatchNorm(use_running_average=True)(x)\n x = nn.Dropout(0.5, deterministic=True)(x)\n x = nn.relu(x)\n return x\n\n class CNN(nn.Module):\n @nn.compact\n def __call__(self, x):\n block = ConvBlock()\n x = block(x)\n x = block(x)\n x = block(x)\n return x\n\n x = jnp.ones((4, 28, 28, 32))\n module_repr = CNN().tabulate(\n jax.random.key(0),\n x=x,\n show_repeated=True,\n console_kwargs=CONSOLE_TEST_KWARGS,\n compute_flops=True,\n compute_vjp_flops=True,\n )\n lines = module_repr.splitlines()\n\n # first call\n self.assertIn('ConvBlock_0/Conv_0', lines[9])\n self.assertIn('bias', lines[9])\n self.assertIn('ConvBlock_0/BatchNorm_0', lines[14])\n self.assertIn('mean', lines[14])\n self.assertIn('bias', lines[14])\n self.assertIn('ConvBlock_0/Dropout_0', lines[19])\n\n # second call\n self.assertIn('ConvBlock_0/Conv_0', lines[23])\n self.assertNotIn('bias', lines[23])\n self.assertIn('ConvBlock_0/BatchNorm_0', lines[25])\n self.assertNotIn('mean', lines[25])\n self.assertNotIn('bias', lines[25])\n self.assertIn('ConvBlock_0/Dropout_0', lines[27])\n\n # third call\n self.assertIn('ConvBlock_0/Conv_0', lines[31])\n self.assertNotIn('bias', lines[31])\n self.assertIn('ConvBlock_0/BatchNorm_0', lines[33])\n self.assertNotIn('mean', lines[33])\n self.assertNotIn('bias', lines[33])\n self.assertIn('ConvBlock_0/Dropout_0', lines[35])\n\n # Test that CNN FLOPs are 3x ConvBlock FLOPs.\n args = ({'dropout': random.key(0), 'params': random.key(1)}, x)\n cnn = summary._get_module_table(\n CNN(),\n depth=1,\n show_repeated=True,\n compute_flops=True,\n compute_vjp_flops=True,\n )(*args, mutable=True)\n\n block = summary._get_module_table(\n ConvBlock(),\n depth=1,\n show_repeated=True,\n compute_flops=True,\n compute_vjp_flops=True,\n )(*args, mutable=True)\n\n # Total forward/backward FLOPs equal to their sums of sub blocks.\n self.assertEqual(cnn[0].flops, sum(r.flops for r in cnn[1:]))\n self.assertEqual(cnn[0].vjp_flops, sum(r.vjp_flops for r in cnn[1:]))\n\n # Each sub block has cost equal to ConvBlock instantiated separately.\n for r in cnn[1:]:\n self.assertEqual(r.flops, block[0].flops)\n self.assertEqual(r.vjp_flops, block[0].vjp_flops)\n\n def test_empty_input(self):\n class EmptyInput(nn.Module):\n @nn.compact\n def __call__(self):\n return 1\n\n module = EmptyInput()\n module_repr = module.tabulate(\n {},\n console_kwargs=CONSOLE_TEST_KWARGS,\n compute_flops=True,\n compute_vjp_flops=True,\n )\n lines = module_repr.splitlines()\n\n # 1 output and 0 forward / backward FLOPs.\n self.assertRegex(\n lines[5], r'│\\s*│\\s*EmptyInput\\s*│\\s*│\\s*1\\s*│\\s*0\\s*│\\s*0\\s*│'\n )\n\n def test_numpy_scalar(self):\n class Submodule(nn.Module):\n def __call__(self, x):\n return x + 1\n\n class EmptyInput(nn.Module):\n @nn.compact\n def __call__(self):\n return Submodule()(x=np.pi)\n\n module = EmptyInput()\n module_repr = module.tabulate(\n {},\n console_kwargs=CONSOLE_TEST_KWARGS,\n compute_flops=True,\n compute_vjp_flops=True,\n )\n lines = module_repr.splitlines()\n\n self.assertIn('4.141592', lines[5])\n self.assertIn('x: 3.141592', lines[7])\n self.assertIn('4.141592', lines[7])\n\n # 0 forward / backward FLOPs due to precomputed output values.\n self.assertIn('│ 0 │ 0', lines[5])\n self.assertIn('│ 0 │ 0', lines[7])\n\n def test_partitioned_params(self):\n class Classifier(nn.Module):\n @nn.compact\n def __call__(self, x):\n hidden = nn.Dense(\n features=1024,\n kernel_init=nn.with_partitioning(\n nn.initializers.lecun_normal(), (None, 'data')\n ),\n bias_init=nn.with_partitioning(nn.initializers.zeros, (None,)),\n name='hidden',\n )\n x = x / 255.0\n x = x.reshape((x.shape[0], -1)) # flatten\n x = nn.relu(hidden(x))\n x = nn.Dense(features=10, name='head')(x)\n return x\n\n module = Classifier()\n lines = module.tabulate(\n jax.random.key(0),\n jnp.empty((1, 28, 28, 1)),\n console_kwargs=CONSOLE_TEST_KWARGS,\n compute_flops=True,\n compute_vjp_flops=True,\n ).splitlines()\n self.assertIn('P(None,)', lines[7])\n self.assertIn('P(None, data)', lines[8])\n\n # Per-layer forward FLOPs:\n self.assertIn('1606656', lines[7]) # 1 * (28 * 28 * 1) * 1024 * 2 + 1024\n self.assertIn('20490', lines[12]) # 1 * (1024 ) * 10 * 2 + 10\n\n # Total forward FLOPs: input division + ReLU + two dense layers above.\n # (1 * 28 * 28 * 1) + (1 * 1024) + 1606656 + 20490.\n self.assertIn('1628954', lines[5])\n\n # Per-layer backward FLOPs:\n\n # [3x MMs: forward, input cotangent, weight cotangent] 1024 * 784 * 2 * 3\n # + [forward bias addition] 1024\n # + [`mutable=True`: weight and bias sizes] 1024 * 784 + 1024\n self.assertIn('5621760', lines[7])\n\n # [3x matmuls: forward, input cotangent, weight cotangent] 1024 * 10 * 2 * 3\n # + [forward bias addition] 10\n # + [`mutable=True`: weight and bias sizes] 1024 * 10 + 10\n self.assertIn('71700', lines[12])\n\n # Total backward FLOPs: input division + ReLU + two dense layers above.\n # 2 * (1 * 28 * 28 * 1) + 3 * (1 * 1024) + 5621760 + 71700.\n self.assertIn('5698100', lines[5])\n\n def test_non_array_variables(self):\n class Metadata(struct.PyTreeNode):\n names: tuple = struct.field(pytree_node=False)\n\n class Foo(nn.Module):\n @nn.compact\n def __call__(self):\n self.sow('foo', 'bar', Metadata(('baz', 'qux')))\n\n module = Foo()\n lines = module.tabulate(\n {},\n console_kwargs=CONSOLE_TEST_KWARGS,\n compute_flops=True,\n compute_vjp_flops=True,\n ).splitlines()\n self.assertIn('names', lines[6])\n self.assertIn('baz', lines[7])\n self.assertIn('qux', lines[8])\n\n # 0 forward and backward FLOPs.\n self.assertIn('│ 0 │ 0', lines[5])\n\n def test_tabulate_param_count_and_flops(self):\n class Foo(nn.Module):\n @nn.compact\n def __call__(self, x):\n h = nn.Dense(4)(x)\n return nn.Dense(2)(h)\n\n module = Foo()\n rng = jax.random.key(0)\n x = jnp.ones((16, 9))\n\n rep = module.tabulate(\n rng,\n x,\n console_kwargs=CONSOLE_TEST_KWARGS,\n compute_flops=True,\n compute_vjp_flops=True,\n )\n lines = rep.splitlines()\n self.assertIn('Total Parameters: 50', lines[-2])\n\n def test_tabulate_enum(self):\n class Net(nn.Module):\n @nn.compact\n def __call__(self, inputs):\n x = inputs['x']\n x = nn.Dense(features=2)(x)\n return jnp.sum(x)\n\n class InputEnum(str, enum.Enum):\n x = 'x'\n\n inputs = {InputEnum.x: jnp.ones((1, 1))}\n # test args\n lines = Net().tabulate(jax.random.key(0), inputs).split('\\n')\n self.assertIn('x: \\x1b[2mfloat32\\x1b[0m[1,1]', lines[5])\n # test kwargs\n lines = Net().tabulate(jax.random.key(0), inputs=inputs).split('\\n')\n self.assertIn('inputs:', lines[5])\n self.assertIn('x: \\x1b[2mfloat32\\x1b[0m[1,1]', lines[6])\n\n def test_tabulate_norm_wrapper(self):\n class SubModel(nn.Module):\n @nn.compact\n def __call__(self, x):\n x = nn.SpectralNorm(nn.Dense(5))(x, update_stats=False)\n x = nn.Dense(6)(x)\n x = nn.WeightNorm(nn.Dense(7))(x)\n return x\n\n class Model(nn.Module):\n @nn.compact\n def __call__(self, x):\n x = nn.WeightNorm(nn.Dense(3))(x)\n x = nn.Dense(4)(x)\n x = SubModel()(x)\n x = nn.Dense(8)(x)\n x = nn.SpectralNorm(nn.Dense(9))(x, update_stats=False)\n return x\n\n x = jnp.ones((1, 2))\n key = jax.random.key(0)\n model = Model()\n\n lines = model.tabulate(\n key,\n x,\n console_kwargs=CONSOLE_TEST_KWARGS,\n compute_flops=True,\n compute_vjp_flops=True,\n ).splitlines()\n\n self.assertIn('Model', lines[5])\n self.assertIn('WeightNorm_0', lines[7])\n self.assertIn('Dense_0/kernel/scale', lines[7])\n self.assertIn('Dense_0', lines[11])\n self.assertIn('Dense_1', lines[16])\n self.assertIn('SubModel_0', lines[21])\n self.assertIn('SubModel_0/SpectralNorm_0', lines[23])\n self.assertIn('Dense_0/kernel/sigma', lines[23])\n self.assertIn('Dense_0/kernel/u', lines[24])\n self.assertIn('SubModel_0/Dense_0', lines[28])\n self.assertIn('SubModel_0/Dense_1', lines[33])\n self.assertIn('SubModel_0/WeightNorm_0', lines[38])\n self.assertIn('Dense_2/kernel/scale', lines[38])\n self.assertIn('SubModel_0/Dense_2', lines[42])\n self.assertIn('Dense_2', lines[47])\n self.assertIn('SpectralNorm_0', lines[52])\n self.assertIn('Dense_3/kernel/sigma', lines[52])\n self.assertIn('Dense_3/kernel/u', lines[53])\n self.assertIn('Dense_3', lines[57])\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/nnx/__init__.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n" }, { "path": "tests/nnx/bridge/module_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nfrom typing import Any\n\nfrom flax.linen.dtypes import promote_dtype\n\nos.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=4'\n\nimport jax\nimport jax.numpy as jnp\nfrom absl.testing import absltest\nimport numpy as np\n\nfrom flax import linen as nn\nfrom flax import nnx\nfrom flax.nnx import bridge\nfrom flax.nnx.bridge.module import MODULE_CONTEXT\n\n\nclass TestBridgeModule(absltest.TestCase):\n def test_update(self):\n class Foo(bridge.Module):\n a: int\n\n foo = Foo(1)\n state = {'b': {'c': nnx.Param(jnp.array(2))}}\n nnx.update(foo, state)\n\n def test_module_stack(self):\n \"\"\"Test that apply set the module stack correctly.\"\"\"\n test = self\n\n class Foo(bridge.Module):\n def setup(self):\n current_ctx = MODULE_CONTEXT.module_stack[-1]\n test.assertIs(current_ctx.module, self)\n test.assertFalse(current_ctx.in_compact)\n\n def __call__(self):\n current_ctx = MODULE_CONTEXT.module_stack[-1]\n test.assertIs(current_ctx.module, self)\n test.assertFalse(current_ctx.in_compact)\n\n foo = Foo()\n foo.apply({})\n\n def test_compact_basic(self):\n test = self\n class Linear(bridge.Module):\n dout: int\n\n @bridge.compact\n def __call__(self, x):\n w = self.param(\n 'w', nnx.initializers.uniform(), (x.shape[-1], self.dout)\n )\n b = self.param('b', nn.initializers.zeros_init(), (self.dout,))\n return x @ w + b[None]\n\n class Foo(bridge.Module):\n dout: int\n\n @bridge.compact\n def __call__(self, x):\n din = x.shape[-1]\n self.linear = Linear(self.dout)\n x = self.linear(x)\n\n # NNX\n graphdef, state = nnx.split(self)\n test.assertIn('Linear_0', state)\n test.assertIn('w', state['Linear_0'])\n test.assertIn('b', state['Linear_0'])\n\n return x\n\n foo = Foo(5)\n x = jnp.ones((3, 2))\n\n self.assertIsInstance(foo, nnx.Module)\n\n variables = foo.init(0, x)\n params = variables['params']\n\n self.assertIn('Linear_0', params)\n self.assertIn('w', params['Linear_0'])\n self.assertIn('b', params['Linear_0'])\n self.assertEqual(params['Linear_0']['w'].shape, (2, 5))\n self.assertEqual(params['Linear_0']['b'].shape, (5,))\n\n y: jax.Array = foo.apply(variables, x)\n\n self.assertEqual(y.shape, (3, 5))\n\n def test_mutable_state(self):\n class FooLinen(nn.Module):\n @nn.compact\n def __call__(self):\n count = self.variable(\n 'counts', 'count', lambda: jnp.zeros((), jnp.int32)\n )\n count.value += 1\n\n model_linen = FooLinen()\n initial_vars_linen = model_linen.init({})\n _, vars_linen = model_linen.apply(initial_vars_linen, mutable='counts')\n\n class FooNNX(bridge.Module):\n @bridge.compact\n def __call__(self):\n count = self.variable(\n 'counts', 'count', lambda: jnp.zeros((), jnp.int32)\n )\n count[...] += 1\n\n model_nnx = FooNNX()\n\n initial_vars_nnx = model_nnx.init({})\n _, vars_nnx = model_nnx.apply(initial_vars_nnx, mutable='counts')\n\n self.assertEqual(\n initial_vars_linen['counts']['count'], initial_vars_nnx['counts']['count']\n )\n self.assertEqual(vars_linen['counts']['count'], vars_nnx['counts']['count'])\n\n def test_compact_parent_none(self):\n class Foo(bridge.Module):\n pass\n\n class Bar(bridge.Module):\n @bridge.compact\n def __call__(self):\n return Foo().scope\n\n bar = Bar()\n scope = bar.apply({}, rngs=1)\n self.assertIsNone(bar.scope)\n\n self.assertEqual(scope.rngs.default.key[...], jax.random.key(1))\n self.assertEqual(scope.rngs.default.count[...], 0)\n\n class Baz(bridge.Module):\n @bridge.compact\n def __call__(self):\n return Foo(parent=None).scope\n\n baz = Baz()\n scope = baz.apply({}, rngs=1)\n self.assertIsNone(scope)\n\n def test_dense_port(self):\n class Dense(bridge.Module):\n features: int\n use_bias: bool = True\n dtype: Any = None\n param_dtype: Any = jnp.float32\n precision: Any = None\n kernel_init: Any = nnx.initializers.lecun_normal()\n bias_init: Any = nnx.initializers.zeros_init()\n # Deprecated. Will be removed.\n dot_general: Any | None = None\n dot_general_cls: Any = None\n\n @bridge.compact\n def __call__(self, inputs: jax.Array) -> jax.Array:\n kernel = self.param(\n 'kernel',\n self.kernel_init,\n (jnp.shape(inputs)[-1], self.features),\n self.param_dtype,\n )\n if self.use_bias:\n bias = self.param(\n 'bias', self.bias_init, (self.features,), self.param_dtype\n )\n else:\n bias = None\n inputs, kernel, bias = promote_dtype(\n inputs, kernel, bias, dtype=self.dtype\n )\n\n if self.dot_general_cls is not None:\n dot_general = self.dot_general_cls()\n elif self.dot_general is not None:\n dot_general = self.dot_general\n else:\n dot_general = jax.lax.dot_general\n y = dot_general(\n inputs,\n kernel,\n (((inputs.ndim - 1,), (0,)), ((), ())),\n precision=self.precision,\n )\n if bias is not None:\n y += jnp.reshape(bias, (1,) * (y.ndim - 1) + (-1,))\n return y\n\n m = Dense(3)\n x = jnp.ones((1, 10, 2))\n y, variables = m.init_with_output(0, x)\n\n self.assertEqual(y.shape, (1, 10, 3))\n self.assertEqual(variables['params']['kernel'].shape, (2, 3))\n self.assertEqual(variables['params']['bias'].shape, (3,))\n\n y = m.apply(variables, x)\n\n self.assertEqual(y.shape, (1, 10, 3))\n self.assertEqual(variables['params']['kernel'].shape, (2, 3))\n self.assertEqual(variables['params']['bias'].shape, (3,))\n\n @jax.jit\n def train_step(params, x, y):\n def loss_fn(params):\n y_pred = m.apply({'params': params}, x)\n return jnp.mean((y - y_pred) ** 2)\n\n grads = jax.grad(loss_fn)(params)\n\n params = jax.tree.map(lambda p, g: p - 0.01 * g, params, grads)\n\n return params\n\n params = variables['params']\n x = jnp.ones((1, 10, 2))\n y = jnp.ones((1, 10, 3))\n\n params = train_step(params, x, y)\n\n def test_metadata(self):\n class Linear(bridge.Module):\n dout: int\n\n @bridge.compact\n def __call__(self, x):\n w = self.param(\n 'w', bridge.with_partitioning(nnx.initializers.uniform(), ('in', 'out')),\n (x.shape[-1], self.dout)\n )\n b = self.param('b', nnx.initializers.zeros_init(), (self.dout,))\n return x @ w + b[None]\n\n foo = Linear(6)\n x = jnp.ones((4, 2))\n\n mesh = jax.make_mesh(\n (2, 2),\n ('in', 'out'),\n axis_types=(jax.sharding.AxisType.Auto,) * len(('in', 'out')),\n )\n with jax.set_mesh(mesh):\n variables = foo.init(0, x)\n y: jax.Array = foo.apply(variables, x)\n\n params = variables['params']\n self.assertIsInstance(params['w'], nn.Partitioned)\n self.assertEqual(params['w'].value.shape, (2, 6))\n self.assertEqual(params['w'].names, ('in', 'out'))\n self.assertEqual(nn.get_partition_spec(variables)['params']['w'],\n jax.sharding.PartitionSpec('in', 'out'))\n self.assertIsInstance(params['b'], jax.Array)\n self.assertEqual(params['b'].shape, (6,))\n\n self.assertEqual(y.shape, (4, 6))\n\n def test_pure_nnx_submodule(self):\n class NNXLayer(nnx.Module):\n def __init__(self, dim, dropout, rngs):\n self.linear = nnx.Linear(dim, dim, use_bias=False, rngs=rngs)\n self.dropout = nnx.Dropout(dropout, rngs=rngs)\n self.count = nnx.Intermediate(jnp.array([0.]))\n def __call__(self, x):\n # Required check to avoid state update in `init()`. Can this be avoided?\n if not bridge.current_module().is_initializing():\n self.count[...] += 1\n x = self.linear(x)\n x = self.dropout(x)\n return x\n\n class BridgeMLP(bridge.Module):\n @bridge.compact\n def __call__(self, x):\n x = bridge.nnx_in_bridge_mdl(lambda r: NNXLayer(8, 0.3, rngs=r))(x)\n x = bridge.nnx_in_bridge_mdl(\n lambda r: NNXLayer(8, 0.3, rngs=r), name='another')(x)\n return x\n\n model = BridgeMLP()\n x = jax.random.normal(jax.random.key(0), (4, 8))\n variables = model.init(jax.random.key(1), x)\n self.assertSameElements(variables['params'].keys(),\n ['NNXLayer_0', 'another'])\n self.assertFalse(jnp.array_equal(\n variables['params']['NNXLayer_0']['linear']['kernel'],\n variables['params']['another']['linear']['kernel'], ))\n self.assertEqual(variables['intermediates']['NNXLayer_0']['count'], 0)\n\n k1, k2, k3 = jax.random.split(jax.random.key(0), 3)\n y1 = model.apply(variables, x, rngs={'params': k1, 'dropout': k2})\n y2 = model.apply(variables, x, rngs={'params': k1, 'dropout': k3})\n assert not jnp.array_equal(y1, y2)\n\n _, updates = model.apply(variables, x, rngs={'params': k1, 'dropout': k3},\n mutable=True)\n self.assertEqual(updates['intermediates']['NNXLayer_0']['count'], 1)\n\n class BridgeMLPSetup(bridge.Module):\n def setup(self):\n self.layer = bridge.nnx_in_bridge_mdl(\n lambda r: NNXLayer(8, 0.3, rngs=r))\n def __call__(self, x):\n return self.layer(x)\n\n model = BridgeMLPSetup()\n variables = model.init(jax.random.key(1), x)\n self.assertSameElements(variables['params'].keys(), ['layer'])\n y1 = model.apply(variables, x, rngs={'params': k1, 'dropout': k2})\n y2 = model.apply(variables, x, rngs={'params': k1, 'dropout': k3})\n assert not jnp.array_equal(y1, y2)\n\n def test_pure_nnx_submodule_modified_rng(self):\n class FooStack(nnx.Module):\n def __init__(self, in_dim, key):\n keys = jax.random.split(key, in_dim)\n self.rngs = nnx.Rngs(keys)\n def __call__(self, x):\n @nnx.vmap\n def generate_weights(r):\n return jax.random.normal(r.default(), (2,))\n w = generate_weights(self.rngs)\n return x @ w\n\n class BridgeFoo(bridge.Module):\n @bridge.compact\n def __call__(self, x):\n x = bridge.nnx_in_bridge_mdl(lambda r: FooStack(4, r.default()))(x)\n return x\n\n model = BridgeFoo()\n v = model.init(jax.random.key(1), jnp.ones((1, 4)))\n y = model.apply(v, jnp.ones((1, 4)), rngs=jax.random.key(1))\n\n def test_linen_submodule(self):\n class LinenLayer(nn.Module):\n dim: int\n dropout_rate: float\n def setup(self):\n self.linear = nn.Dense(self.dim, use_bias=False)\n self.dropout = nn.Dropout(self.dropout_rate, deterministic=False)\n\n def __call__(self, x):\n if not self.is_initializing():\n self.sow('intermediates', 'count', 1,\n init_fn=lambda: 0, reduce_fn=lambda a, b: a + b)\n x = self.linear(x)\n x = self.dropout(x)\n return x\n\n class BridgeMLP(bridge.Module):\n @bridge.compact\n def __call__(self, x):\n x = bridge.linen_in_bridge_mdl(LinenLayer(8, 0.3))(x)\n x = bridge.linen_in_bridge_mdl(LinenLayer(8, 0.3), name='another')(x)\n return x\n\n model = BridgeMLP()\n x = jax.random.normal(jax.random.key(0), (4, 8))\n variables = model.init(jax.random.key(1), x)\n self.assertFalse(jnp.array_equal(\n variables['params']['LinenLayer_0']['linear']['kernel'],\n variables['params']['another']['linear']['kernel'], ))\n\n k1, k2, k3 = jax.random.split(jax.random.key(0), 3)\n y1 = model.apply(variables, x, rngs={'params': k1, 'dropout': k2})\n y2 = model.apply(variables, x, rngs={'params': k1, 'dropout': k3})\n assert not jnp.array_equal(y1, y2)\n\n _, updates = model.apply(variables, x, rngs={'params': k1, 'dropout': k3},\n mutable=True)\n self.assertEqual(updates['intermediates']['LinenLayer_0']['count'], 1)\n\n class BridgeMLPSetup(bridge.Module):\n def setup(self):\n self.layer = bridge.linen_in_bridge_mdl(LinenLayer(8, 0.3))\n def __call__(self, x):\n return self.layer(x)\n\n model = BridgeMLPSetup()\n variables = model.init(jax.random.key(1), x)\n self.assertSameElements(variables['params'].keys(), ['layer'])\n y1 = model.apply(variables, x, rngs={'params': k1, 'dropout': k2})\n y2 = model.apply(variables, x, rngs={'params': k1, 'dropout': k3})\n assert not jnp.array_equal(y1, y2)\n\n def test_name(self):\n class Foo(bridge.Module):\n name: str\n\n class Bar(bridge.Module):\n @bridge.compact\n def __call__(self):\n f = Foo(name='f')\n assert f.name == 'f'\n assert self.f == f\n\n Bar().init()\n\n def test_transforms(self):\n class Dense(bridge.Module):\n dout: int\n @bridge.compact\n def __call__(self, x: jax.Array) -> jax.Array:\n return x @ self.param('w', nn.initializers.normal(),\n (x.shape[-1], self.dout))\n\n class MLP(bridge.Module):\n dim: int\n num_layers: int\n def setup(self):\n @nnx.split_rngs(splits=self.num_layers)\n @nnx.vmap(\n in_axes=(nnx.StateAxes({nnx.RngState: 0, ...: None}),),\n axis_size=self.num_layers,\n transform_metadata={nnx.PARTITION_NAME: None},\n )\n def create_block(parent):\n block = Dense(self.dim)\n parent.block = block\n create_block(self)\n\n def __call__(self, x):\n @nnx.split_rngs(splits=self.num_layers)\n @nnx.scan(\n in_axes=(0, nnx.Carry),\n out_axes=nnx.Carry,\n transform_metadata={nnx.PARTITION_NAME: None},\n )\n def forward_block(model, x):\n return model(x)\n x = forward_block(self.block, x)\n return x\n\n model = MLP(dim=32, num_layers=2)\n x = jnp.ones((4, 32))\n variables = model.init(jax.random.key(0), x)\n y = model.apply(variables, x)\n w = variables['params']['block']['w']\n np.testing.assert_array_equal(y, x @ w[0] @ w[1])\n\n def test_shared_modules(self):\n class Dense(bridge.Module):\n dout: int\n @bridge.compact\n def __call__(self, x):\n return x @ self.param('w', nn.initializers.normal(),\n (x.shape[-1], self.dout))\n\n class Bottom(bridge.Module):\n def setup(self):\n self.layer = Dense(4)\n\n def __call__(self, x):\n return self.layer(x)\n\n class Top(bridge.Module):\n def setup(self):\n self.zzz = Bottom()\n self.set_attr_priority('zzz', bridge.AttrPriority.HIGH)\n self.aaa = self.zzz # another reference\n self.dense = self.aaa.layer # and another reference\n\n def forward(self, x):\n return self.aaa(x)\n\n def __call__(self, x):\n forward = nnx.remat(self.__class__.forward)\n return forward(self, x)\n\n model = Top()\n x = jnp.ones((4, 32))\n params = model.init(jax.random.key(0), x)['params']\n self.assertSameElements(['zzz'], params.keys())\n\n def test_linen_layer_naming(self):\n class Dense(bridge.Module):\n dout: int\n @bridge.compact\n def __call__(self, x):\n return x @ self.param('w', lambda _: jnp.ones((x.shape[-1], self.dout)))\n\n class MLP(bridge.Module):\n nlayers: int\n def setup(self):\n self.layers = [Dense(4, name=f'layer_{i}') for i in range(self.nlayers)]\n def __call__(self, x):\n for layer in self.layers:\n x = layer(x)\n return x\n\n model = MLP(nlayers=3)\n x = jnp.ones((2, 4))\n params = model.init(jax.random.key(0), x)['params']\n self.assertSameElements([f'layer_{i}' for i in range(3)], params.keys())\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/nnx/bridge/wrappers_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nos.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=4'\n\nimport jax\nimport jax.numpy as jnp\nimport numpy as np\nfrom absl.testing import absltest\n\nimport flax\nfrom flax import linen as nn\nfrom flax import nnx\nfrom flax.nnx import bridge\n\n\nclass TestCompatibility(absltest.TestCase):\n def setUp(self):\n super().setUp()\n dim1 = max(jax.device_count() // 2, 1)\n device_mesh = np.array(jax.devices()).reshape(dim1, jax.device_count() // dim1)\n self.mesh = jax.sharding.Mesh(devices=device_mesh, axis_names=('in', 'out'))\n\n def test_functional(self):\n # Functional API for NNX Modules\n functional = bridge.functional(nnx.Linear)(32, 64)\n state = functional.init(rngs=nnx.Rngs(0))\n x = jax.numpy.ones((1, 32))\n y, updates = functional.apply(state)(x)\n\n ##################\n ### LinenToNNX ###\n ##################\n\n def test_linen_to_nnx(self):\n ## Wrapper API for Linen Modules\n linen_module = nn.Dense(features=64)\n x = jax.numpy.ones((1, 32))\n model = bridge.ToNNX(linen_module, rngs=nnx.Rngs(0)).lazy_init(x) # like linen init\n y = model(x) # like linen apply\n assert y.shape == (1, 64)\n self.assertIsInstance(model.kernel, nnx.Variable)\n # NNX automatically adds metadata box regardless of original Linen module.\n linen_vars = {\n 'params': {'kernel': model.kernel[...], 'bias': model.bias[...]}\n }\n linen_y = linen_module.apply(linen_vars, x)\n np.testing.assert_array_equal(y, linen_y)\n\n def test_linen_to_nnx_submodule(self):\n class NNXOuter(nnx.Module):\n def __init__(self, dout: int, *, rngs: nnx.Rngs):\n self.nn_dense1 = bridge.ToNNX(nn.Dense(dout, use_bias=False), rngs=rngs)\n self.b = nnx.Param(jax.random.uniform(rngs.params(), (1, dout,)))\n self.batchnorm = bridge.ToNNX(nn.BatchNorm(use_running_average=True), rngs=rngs)\n self.rngs = rngs\n\n def __call__(self, x):\n x = self.nn_dense1(x) + self.b\n return self.batchnorm(x)\n\n x = jax.random.normal(jax.random.key(0), (2, 4))\n model = NNXOuter(3, rngs=nnx.Rngs(0))\n gdef_before_lazy_init, _ = nnx.split(model)\n bridge.lazy_init(model, x)\n gdef_full, state = nnx.split(model)\n assert gdef_before_lazy_init != gdef_full\n assert 'nn_dense1' in state\n assert 'batchnorm' in state\n assert 'kernel' in state['nn_dense1']\n y = model(x)\n k, b = state['nn_dense1']['kernel'][...], state['b'][...]\n np.testing.assert_allclose(y, x @ k + b, rtol=1e-5)\n assert gdef_full == nnx.graphdef(model) # static data is stable now\n\n def test_linen_to_nnx_noncall_method(self):\n class Foo(nn.Module):\n @nn.compact\n def __call__(self, x):\n b = self.param('b', nn.zeros_init(), (1, 3,))\n return self.dot(x) + b\n\n @nn.compact\n def dot(self, x):\n w = self.param('w', nn.initializers.lecun_normal(), (4, 3))\n return x @ w\n\n def rngs(self):\n raise ValueError('This should not be called because ToNNX has .rngs')\n\n x = jax.random.normal(jax.random.key(0), (2, 4))\n model = bridge.ToNNX(Foo(), rngs=nnx.Rngs(0))\n bridge.lazy_init(model.dot, x)\n y = model.dot(x)\n np.testing.assert_allclose(y, x @ nnx.state(model)['w'][...])\n # lazy_init only initialized param w inside dot(), so calling __call__ should fail\n with self.assertRaises(flax.errors.ScopeParamNotFoundError):\n y = model(x)\n assert isinstance(model.to_nnx__rngs, nnx.Rngs)\n\n def test_linen_to_nnx_mutable(self):\n class Foo(nn.Module):\n def setup(self):\n self.count = self.variable('counter', 'count', lambda: jnp.zeros((), jnp.int32))\n\n def __call__(self, x):\n if not self.is_initializing():\n self.count.value += 1\n return x\n\n x = lambda: jnp.zeros((), jnp.int32)\n model = bridge.ToNNX(Foo(), rngs=nnx.Rngs(0)).lazy_init(x)\n self.assertEqual(nnx.state(model)['count'][...], 0)\n y = model(x, mutable=True)\n self.assertEqual(nnx.state(model)['count'][...], 1)\n\n def test_linen_to_nnx_transform(self):\n class NNXOuter(nnx.Module):\n def __init__(self, dout: int, rngs: nnx.Rngs):\n self.inner = nnx.bridge.ToNNX(nn.Dense(dout), rngs=rngs)\n self.rngs = rngs\n\n def __call__(self, x):\n\n @nnx.split_rngs(splits=5)\n @nnx.vmap(in_axes=(0, None), axis_size=5)\n def vmap_fn(inner, x):\n return inner(x)\n\n return vmap_fn(self.inner, x)\n\n x = jax.random.normal(jax.random.key(0), (2, 4))\n model = NNXOuter(3, rngs=nnx.Rngs(0))\n nnx.bridge.lazy_init(model, x)\n\n self.assertEqual(model.inner.kernel.shape, (5, 4, 3))\n self.assertEqual(model.inner.bias.shape, (5, 3))\n\n def test_linen_to_nnx_metadata(self):\n linen_module = nn.Dense(\n features=64,\n kernel_init=nn.with_partitioning(nn.initializers.lecun_normal(),\n ('in', 'out')),\n bias_init=nn.with_logical_partitioning(nn.initializers.zeros_init(),\n ('out-alias',),\n rules=(('out-alias', 'out'),)),\n )\n x = jax.numpy.ones((1, 32))\n linen_vars = linen_module.init(jax.random.key(0), x)\n\n @nnx.jit\n def create_sharded_nnx_module(x):\n model = bridge.lazy_init(bridge.ToNNX(linen_module, rngs=nnx.Rngs(0)), x)\n state = nnx.state(model)\n sharded_state = jax.lax.with_sharding_constraint(state, nnx.get_partition_spec(state))\n nnx.update(model, sharded_state)\n return model\n with jax.set_mesh(self.mesh):\n nnx_model = create_sharded_nnx_module(x)\n\n # nn.Partitioned metadata boxes translated into valid nnx.Variable boxes.\n self.assertIsInstance(linen_vars['params']['kernel'], nn.Partitioned)\n self.assertIsInstance(linen_vars['params']['bias'], nn.LogicallyPartitioned)\n self.assertIsInstance(nnx_model.kernel, nnx.Variable)\n assert nnx_model.kernel.out_sharding == ('in', 'out')\n assert nnx_model.kernel[...].sharding.is_equivalent_to(\n jax.sharding.NamedSharding(\n self.mesh, jax.sharding.PartitionSpec('in', 'out')\n ),\n ndim=2,\n ), f'{nnx_model.kernel[...].sharding = }'\n\n assert nnx_model.bias.out_sharding == ('out-alias',)\n assert nnx_model.bias.sharding_rules == (('out-alias', 'out'),)\n assert nnx_model.bias[...].sharding.is_equivalent_to(\n jax.sharding.NamedSharding(self.mesh, jax.sharding.PartitionSpec('out')),\n ndim=1,\n )\n\n def test_linen_to_nnx_state_structure_consistency(self):\n class LinenInner(nn.Module):\n dout: int\n @nn.compact\n def __call__(self, x):\n w = self.param('w', nn.initializers.lecun_normal(), (x.shape[-1], self.dout))\n return nn.Dropout(rate=0.5, deterministic=False)(x @ w)\n\n class LinenMiddle(nn.Module):\n dout: int\n @nn.compact\n def __call__(self, x):\n dot = LinenInner(self.dout, name='dot')\n b = self.variable('bias', 'b', nn.initializers.zeros_init(), None, (1, self.dout))\n return dot(x) + b.value\n\n @nnx.register_variable_name('bias')\n class Bias(nnx.Variable):\n pass\n\n class NNXMiddle(nnx.Module):\n def __init__(self, dout: int, *, rngs: nnx.Rngs):\n self.dot = bridge.ToNNX(LinenInner(dout), rngs=rngs)\n self.b = Bias(nnx.initializers.zeros_init()(rngs.params(), (1, dout)))\n def __call__(self, x):\n return self.dot(x) + self.b\n\n x = jax.random.normal(jax.random.key(42), (2, 4))\n from_top = bridge.lazy_init(\n bridge.ToNNX(LinenMiddle(dout=3), rngs=nnx.Rngs(0, dropout=1)), x)\n from_middle = bridge.lazy_init(\n NNXMiddle(dout=3, rngs=nnx.Rngs(0, dropout=1)), x)\n\n # Remove the NNX-module-local RNG states, which will be different\n # because the NNX modules are on different level\n def get_weights(model):\n return nnx.split(model, nnx.RngCount, nnx.RngKey, ...)[3]\n from_top_weights = get_weights(from_top)\n from_middle_weights = get_weights(from_middle)\n\n # Confirm the rest of the state has the same structure.\n self.assertEqual(jax.tree.structure(from_top_weights),\n jax.tree.structure(from_middle_weights))\n\n def test_adding_new_attributes(self):\n class LinenModule(nn.Module):\n @nn.compact\n def __call__(self):\n if self.is_initializing() and self.is_mutable_collection('cache'):\n self.put_variable('cache', 'x', 0)\n res = self.get_variable('cache', 'x')\n return res\n\n class NNXModule(nnx.Module):\n def __init__(self):\n self.module = nnx.bridge.ToNNX(LinenModule()).lazy_init()\n\n def __call__(self):\n result1 = self.module(mutable=['cache'])\n assert result1 == 0\n result2 = self.module()\n assert result2 == 0, result2 # fails: result2 is None\n\n module = NNXModule()\n module()\n\n ##################\n ### NNXToLinen ###\n ##################\n\n def test_nnx_to_linen(self):\n model = bridge.to_linen(nnx.Linear, 32, out_features=64)\n x = jax.numpy.ones((1, 32))\n y, variables = model.init_with_output(jax.random.key(0), x)\n assert y.shape == (1, 64)\n np.testing.assert_allclose(y, x @ variables['params']['kernel'])\n\n def test_nnx_to_linen_multiple_rngs(self):\n class NNXInner(nnx.Module):\n def __init__(self, din, dout, rngs):\n self.w = nnx.Param(nnx.initializers.lecun_normal()(rngs.params(), (din, dout)))\n self.dropout = nnx.Dropout(rate=0.5, rngs=rngs)\n def __call__(self, x):\n return self.dropout(x @ self.w[...])\n\n class LinenOuter(nn.Module):\n @nn.compact\n def __call__(self, x):\n inner = bridge.to_linen(NNXInner, 4, 3)\n return inner(x)\n\n xkey, pkey, dkey1, dkey2 = jax.random.split(jax.random.key(0), 4)\n x = jax.random.normal(xkey, (2, 4))\n model = LinenOuter()\n y1, var = model.init_with_output({'params': pkey, 'dropout': dkey1}, x)\n y2 = model.apply(var, x, rngs={'dropout': dkey2})\n assert not jnp.allclose(y1, y2) # dropout keys are different\n\n def test_nnx_to_linen_multiple_collections(self):\n class NNXInner(nnx.Module):\n def __init__(self, din, dout, rngs):\n self.w = nnx.Param(nnx.initializers.lecun_normal()(rngs.params(), (din, dout)))\n self.bn = nnx.BatchNorm(dout, use_running_average=False, rngs=rngs)\n self.lora = nnx.LoRA(din, 3, dout, rngs=rngs)\n\n def __call__(self, x):\n return self.bn(x @ self.w[...]) + self.lora(x)\n\n xkey, pkey, dkey = jax.random.split(jax.random.key(0), 3)\n x = jax.random.normal(xkey, (2, 4))\n model = bridge.to_linen(NNXInner, 4, 3)\n var = model.init({'params': pkey, 'dropout': dkey}, x)\n self.assertSameElements(var.keys(), ['LoRAParam', 'params', 'batch_stats'])\n y = model.apply(var, x)\n assert y.shape == (2, 3)\n\n def test_nnx_to_linen_mutable(self):\n\n @nnx.register_variable_name('Count', overwrite=True)\n class Count(nnx.Variable):\n pass\n\n class Counter(nnx.Module):\n def __init__(self):\n self.count = Count(jnp.array(0))\n def __call__(self):\n self.count[...] += 1\n\n model = bridge.ToLinen(Counter, skip_rng=True)\n variables = model.init(jax.random.key(0))\n assert variables['Count']['count'] == 0\n\n _, updates = model.apply(variables, mutable='Count')\n assert updates['Count']['count'] == 1\n _ = model.apply(variables | updates)\n\n def test_to_linen_method_call(self):\n class Foo(nn.Module):\n def setup(self):\n self.embedding = nnx.bridge.to_linen(nnx.Embed, 2, 3)\n\n def __call__(self, x):\n return self.embedding(x)\n\n def attend(self, x):\n return self.embedding.attend(x)\n\n module = Foo()\n x = jnp.ones((1,), dtype=jnp.int32)\n z = jnp.ones((1, 3))\n y , params = module.init_with_output(jax.random.key(0), x)\n assert y.shape == (1, 3)\n assert params['params']['embedding']['embedding'].shape == (2, 3)\n x_out = module.apply(params, z, method='attend')\n assert x_out.shape == (1, 2)\n\n def test_to_linen_nnx_method_arg(self):\n module = nnx.bridge.to_linen(nnx.Embed, 2, 3)\n x = jnp.ones((1,), dtype=jnp.int32)\n z = jnp.ones((1, 3))\n y , params = module.init_with_output(jax.random.key(0), x)\n assert y.shape == (1, 3)\n assert params['params']['embedding'].shape == (2, 3)\n x_out = module.apply(params, z, nnx_method='attend')\n assert x_out.shape == (1, 2)\n\n\n\n def test_nnx_to_linen_mutated_static_data(self):\n\n @nnx.register_variable_name('Count', overwrite=True)\n class Count(nnx.Variable):\n pass\n\n class Counter(nnx.Module):\n def __init__(self):\n self.count = Count(jnp.array(0))\n def __call__(self):\n self.count[...] += 1\n self.count_nonzero = nnx.Intermediate(jnp.array(1))\n\n model = bridge.ToLinen(Counter, skip_rng=True)\n variables = model.init(jax.random.key(0))\n assert variables['Count']['count'] == 0\n\n _, updates = model.apply(variables, mutable=['Count', 'intermediates'])\n assert updates['Count']['count'] == 1\n assert updates['intermediates']['count_nonzero'] == 1\n del updates['intermediates']\n _ = model.apply(variables | updates)\n\n def test_nnx_to_linen_transforms(self):\n class LinenOuter(nn.Module):\n dout: int\n @nn.compact\n def __call__(self, x):\n inner = nn.vmap(\n bridge.ToLinen,\n variable_axes={'params': 0},\n split_rngs={'params': True},\n )(nnx.Linear, args=(x.shape[-1], self.dout))\n return inner(x)\n\n xkey, pkey, _ = jax.random.split(jax.random.key(0), 3)\n x = jax.random.normal(xkey, (2, 4))\n model = LinenOuter(dout=3)\n y, var = model.init_with_output(pkey, x)\n k = var['params']['VmapToLinen_0']['kernel']\n assert k.shape == (2, 4, 3)\n np.testing.assert_allclose(y, jnp.einsum('ab,abc->ac', x, k))\n\n def test_nnx_to_linen_metadata(self):\n model = bridge.to_linen(\n nnx.Linear, 32, 64,\n kernel_init=nnx.with_partitioning(nnx.initializers.lecun_normal(), ('in', 'out')))\n x = jax.numpy.ones((1, 32))\n with jax.set_mesh(self.mesh):\n y, variables = model.init_with_output(jax.random.key(0), x)\n pspec_tree = nn.get_partition_spec(variables)\n assert y.shape == (1, 64)\n self.assertIsInstance(variables['params']['kernel'], nnx.bridge.NNXMeta)\n assert variables['params']['kernel'].metadata['out_sharding'] == ('in', 'out')\n self.assertEqual(pspec_tree['params']['kernel'],\n jax.sharding.PartitionSpec('in', 'out'))\n np.testing.assert_allclose(y, x @ variables['params']['kernel'].value)\n\n def test_nnx_to_linen_metadata_transform(self):\n # TODO: add support and testing after axis add/remove in transform is fixed.\n pass\n\n def test_nnx_to_linen_pytree_structure_consistency(self):\n class NNXInner(nnx.Module):\n def __init__(self, din: int, dout: int, *, rngs: nnx.Rngs):\n self.w = nnx.Param(nnx.initializers.lecun_normal()(rngs.params(), (din, dout)))\n self.dropout = nnx.Dropout(rate=0.5, rngs=rngs)\n def __call__(self, x):\n return self.dropout(x @ self.w)\n\n @nnx.register_variable_name('bias', overwrite=True)\n class Bias(nnx.Variable):\n pass\n\n class NNXMiddle(nnx.Module):\n def __init__(self, din: int, dout: int, *, rngs: nnx.Rngs):\n self.dot = NNXInner(din, dout, rngs=rngs)\n self.b = Bias(nnx.initializers.zeros_init()(rngs.params(), (1, dout)))\n def __call__(self, x):\n return self.dot(x) + self.b\n\n class LinenMiddle(nn.Module):\n dout: int\n @nn.compact\n def __call__(self, x):\n dot = bridge.to_linen(NNXInner, x.shape[-1], self.dout, name='dot')\n b = self.variable('bias', 'b', nn.initializers.zeros_init(), None, (1, self.dout))\n return dot(x) + b.value\n\n x = jax.random.normal(jax.random.key(42), (2, 4))\n keys = {'params': jax.random.key(0), 'dropout': jax.random.key(1)}\n from_top = bridge.to_linen(NNXMiddle, din=4, dout=3).init(keys, x)\n from_middle = LinenMiddle(dout=3).init(keys, x)\n\n # Remove the NNX-module-local RNG states, which will be different\n # because the NNX modules are on different level\n def get_weights(variables):\n non_rngs = {}\n for kp, v in flax.traverse_util.flatten_dict(variables).items():\n if 'rngs' not in kp:\n non_rngs[kp] = v\n return flax.traverse_util.unflatten_dict(non_rngs)\n from_top_weights = get_weights(from_top)\n from_middle_weights = get_weights(from_middle)\n\n # Confirm the rest of the state has the same structure.\n self.assertEqual(jax.tree.structure(from_top_weights),\n jax.tree.structure(from_middle_weights))\n\n ############################\n ### Hybrid mix-and-match ###\n ############################\n\n def test_nnx_linen_nnx(self):\n class NNXInner(nnx.Module):\n def __init__(self, din, dout, dropout_rate, rngs):\n self.w = nnx.Param(\n nnx.with_partitioning(nnx.initializers.lecun_normal(), sharding=('in', 'out')\n )(rngs.params(), (din, dout)))\n self.dropout = nnx.Dropout(rate=dropout_rate, rngs=rngs)\n def __call__(self, x):\n return self.dropout(x @ self.w)\n\n class LinenMiddle(nn.Module):\n dout: int\n dropout_rate: float\n @nn.compact\n def __call__(self, x):\n dot = bridge.to_linen(NNXInner, x.shape[-1], self.dout, self.dropout_rate, name='dot')\n logical_init = nn.with_logical_partitioning(\n nn.initializers.lecun_normal(), ('out-alias',), rules=(('out-alias', 'out'),))\n b = self.param('b', logical_init, (2, self.dout))\n return dot(x) + b\n\n class NNXOuter(nnx.Module):\n def __init__(self, dout: int, dropout_rate: float, *, rngs: nnx.Rngs):\n self.inner = bridge.ToNNX(LinenMiddle(dout, dropout_rate), rngs=rngs)\n self.rngs = rngs\n def __call__(self, x):\n return self.inner(x)\n\n x = jax.random.normal(jax.random.key(0), (2, 4))\n\n # Test the RNG\n with jax.set_mesh(self.mesh):\n model = bridge.lazy_init(NNXOuter(dout=6, dropout_rate=0.5,\n rngs=nnx.Rngs(default=1, dropout=2)), x)\n nnx.reseed(model, dropout=2)\n y1, y2 = model(x), model(x)\n # The dropout key of lowest NNX level still changes over stateful calls\n assert not jnp.allclose(y1, y2)\n # Another reseed resets the RNG key back\n nnx.reseed(model, dropout=2)\n np.testing.assert_array_equal(y1, model(x))\n\n # Test the param value with disabled dropout\n with jax.set_mesh(self.mesh):\n model = bridge.lazy_init(NNXOuter(dout=6, dropout_rate=0.,\n rngs=nnx.Rngs(default=1, dropout=2)), x)\n w, b = model.inner.dot['w'], model.inner.b\n np.testing.assert_allclose(model(x), x @ w + b)\n self.assertIsInstance(w, nnx.Param)\n assert hasattr(w, 'out_sharding') and w.out_sharding == ('in', 'out')\n assert hasattr(b, 'out_sharding') and b.out_sharding == ('out-alias', )\n\n def test_linen_nnx_linen(self):\n # TODO: add when we can safely `lazy_init` the NNX module inside `ToLinen` without\n # messing up the stateful part of the NNX module.\n pass\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/nnx/containers_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nfrom flax import nnx\nfrom absl.testing import absltest\nimport jax.numpy as jnp\n\n\nclass TestContainers(absltest.TestCase):\n def test_unbox(self):\n x = nnx.Param(\n jnp.array(1),\n on_get_value=lambda c, x: x + 3, # type: ignore\n )\n\n assert x[...] == 4\n\n def test_on_set_value(self):\n x = nnx.Param(\n jnp.array(1), # type: ignore\n on_set_value=lambda c, x: x + 7, # type: ignore\n )\n x[...] = 5\n\n assert x.get_raw_value() == 12\n\n def test_module_unbox(self):\n class Foo(nnx.Module):\n def __init__(self) -> None:\n self.x = nnx.Param(1, on_get_value=lambda c, x: x + 3)\n\n module = Foo()\n\n assert module.x.get_value() == 4\n assert vars(module)['x'].get_raw_value() == 1\n\n def test_module_box(self):\n class Foo(nnx.Module):\n def __init__(self) -> None:\n self.x = nnx.Param(\n jnp.array(1),\n on_set_value=lambda c, x: x + 7, # type: ignore\n )\n\n module = Foo()\n module.x[...] = 5\n\n assert module.x[...] == 12\n assert vars(module)['x'][...] == 12\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/nnx/filters_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom absl.testing import absltest\n\nfrom flax import nnx\n\n\nclass TestFilters(absltest.TestCase):\n def test_path_contains(self):\n class Model(nnx.Module):\n def __init__(self, rngs):\n self.backbone1 = nnx.Linear(2, 3, rngs=rngs)\n self.backbone2 = nnx.Linear(3, 3, rngs=rngs)\n self.head = nnx.Linear(3, 10, rngs=rngs)\n\n model = Model(nnx.Rngs(0))\n\n head_state = nnx.state(model, nnx.PathContains('head'))\n backbones_state = nnx.state(model, nnx.PathContains('backbone', exact=False))\n\n self.assertIn('head', head_state)\n self.assertNotIn('backbone', head_state)\n self.assertIn('backbone1', backbones_state)\n self.assertIn('backbone2', backbones_state)\n self.assertNotIn('head', backbones_state)\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/nnx/graph_utils_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport dataclasses\nfrom functools import partial\nfrom threading import Thread\nfrom typing import Any\n\nfrom absl.testing import absltest, parameterized\nimport numpy as np\nfrom flax import linen, nnx, struct\nimport jax\nimport jax.numpy as jnp\n\n\nclass StatefulLinear(nnx.Module):\n def __init__(self, din, dout, rngs):\n self.w = nnx.Param(jax.random.uniform(rngs(), (din, dout)))\n self.b = nnx.Param(jnp.zeros((dout,)))\n self.count = nnx.Variable(jnp.array(0, dtype=jnp.uint32))\n\n def increment(self):\n self.count[...] += 1\n\n def __call__(self, x):\n self.count[...] += 1\n return x @ self.w + self.b[None]\n\n\nclass TestGraphUtils(parameterized.TestCase):\n def test_flatten(self):\n a = {'a': 1, 'b': nnx.Param(2)}\n g = [a, 3, a, nnx.Param(4)]\n\n refmap = nnx.graphlib.RefMap()\n graphdef, flat_state = nnx.graphlib.flatten(g, ref_index=refmap, graph=True)\n\n assert flat_state[0][1].get_value() == 2\n assert flat_state[1][1].get_value() == 4\n\n assert len(refmap) == 2 # 2 Variables\n assert a['b'] in refmap\n assert g[3] in refmap\n\n def test_flatten_no_paths(self):\n a = {'a': 1, 'b': nnx.Param(jnp.array(2))}\n g = [a, 3, a, nnx.Param(jnp.array(4))]\n\n refmap = nnx.graphlib.RefMap()\n graphdef, flat_state = nnx.graphlib.flatten(\n g, ref_index=refmap, with_paths=False, graph=True\n )\n\n assert flat_state[0][...] == 2\n assert flat_state[1][...] == 4\n\n assert len(refmap) == 2 # 2 Variables\n assert a['b'] in refmap\n assert g[3] in refmap\n\n def test_unflatten(self):\n a = nnx.Dict(a=1, b=nnx.Param(2))\n g = nnx.List([a, 3, a, nnx.Param(4)])\n\n graphdef, state = nnx.split(g)\n g = nnx.merge(graphdef, state)\n\n assert g[0] is g[2]\n\n @parameterized.parameters(True, False)\n def test_flatten_unflatten_unkown_leaves(self, graph):\n x = jnp.array(1.0)\n graphdef, flat_state = nnx.graphlib.flatten(x, graph=graph)\n\n self.assertIs(flat_state[0][1], x)\n\n x1 = nnx.merge(graphdef, flat_state)\n self.assertIs(x1, x)\n\n @parameterized.parameters(True, False)\n def test_split_merge_unkown_leaves(self, graph):\n x = jnp.array(1.0)\n graphdef, state = nnx.graphlib.split(x, graph=graph)\n\n self.assertIs(state, x)\n\n x1 = nnx.merge(graphdef, state)\n self.assertIs(x1, x)\n\n @parameterized.parameters(True, False)\n def test_split_merge_unkown_leaves_with_filters(self, graph):\n x = jnp.array(1.0)\n graphdef, state, rest = nnx.graphlib.split(x, jax.Array, ..., graph=graph)\n\n self.assertIs(state, x)\n\n x1 = nnx.merge(graphdef, state, rest)\n self.assertIs(x1, x)\n\n def test_unflatten_pure_dict(self):\n a = nnx.Dict(a=1, b=nnx.Param(2))\n g = nnx.List([a, 3, a, nnx.Param(4)])\n\n graphdef, state = nnx.split(g)\n pure_state = nnx.to_pure_dict(state)\n\n g = nnx.merge(graphdef, pure_state)\n\n assert g[0] is g[2]\n\n def test_unflatten_pytree(self):\n a = {'a': 1, 'b': nnx.Param(2)}\n g = [a, 3, a, nnx.Param(4)]\n\n graphdef, state = nnx.split(g)\n g = nnx.merge(graphdef, state)\n\n assert g[0] is not g[2]\n\n def test_unflatten_empty(self):\n a = nnx.Dict({'a': 1, 'b': nnx.Param(2)})\n g = nnx.List([a, 3, a, nnx.Param(4)])\n\n graphdef, state = nnx.split(g)\n\n with self.assertRaisesRegex(ValueError, 'Incorrect number of leaves'):\n nnx.graphlib.unflatten(graphdef, nnx.State({}))\n\n def test_unflatten_return_variables(self):\n a = nnx.Dict({'a': 1, 'b': nnx.Param(2)})\n g = nnx.List([a, 3, a, nnx.Param(4)])\n\n graphdef, state = nnx.graphlib.flatten(\n g, with_paths=True, graph=True\n )\n\n self.assertLen(state, 2)\n self.assertIsInstance(state, nnx.graphlib.FlatState)\n self.assertIsInstance(state[0][1], nnx.Param)\n self.assertIsInstance(state[1][1], nnx.Param)\n\n def test_update_dynamic(self):\n a = {'a': 1, 'b': nnx.Param(jnp.array(2))}\n g = [a, 3, a, nnx.Param(jnp.array(4))]\n\n graphdef, state = nnx.split(g)\n\n state[0]['b'][...] = 3\n nnx.update(g, state)\n\n assert g[0]['b'][...] == 3\n assert g[2]['b'][...] == 3\n\n def test_update_from_pure_dict(self):\n a = {'a': 1, 'b': nnx.Param(jnp.array(2))}\n g = [a, 3, a, nnx.Param(jnp.array(4))]\n\n graphdef, state = nnx.split(g)\n pure_state = nnx.to_pure_dict(state)\n\n pure_state[0]['b'] = jnp.array(3)\n nnx.update(g, pure_state)\n\n assert g[0]['b'][...] == 3\n assert g[2]['b'][...] == 3\n\n @parameterized.parameters(True, False)\n def test_module_list(self, graph):\n rngs = nnx.Rngs(0)\n ls = [\n nnx.Linear(2, 2, rngs=rngs),\n nnx.BatchNorm(2, rngs=rngs),\n ]\n\n graphdef, state = nnx.split(ls, graph=graph)\n\n assert state[0]['kernel'].shape == (2, 2)\n assert state[0]['bias'].shape == (2,)\n assert state[1]['scale'].shape == (2,)\n assert state[1]['bias'].shape == (2,)\n assert state[1]['mean'].shape == (2,)\n assert state[1]['var'].shape == (2,)\n\n def test_shared_variables(self):\n v = nnx.Param(1)\n g = [v, v]\n\n graphdef, state = nnx.split(g)\n\n assert len(nnx.to_flat_state(state)) == 1\n\n g2 = nnx.merge(graphdef, state)\n\n assert g2[0] is g2[1]\n\n def test_tied_weights(self):\n class Foo(nnx.Module):\n def __init__(self, *, rngs: nnx.Rngs) -> None:\n self.bar = nnx.Linear(2, 2, rngs=rngs)\n self.baz = nnx.Linear(2, 2, rngs=rngs)\n\n # tie the weights\n self.baz.kernel = self.bar.kernel\n\n node = Foo(rngs=nnx.Rngs(0))\n graphdef, state = nnx.split(node)\n\n assert len(nnx.to_flat_state(state)) == 3 # 2 bias + 1 kernel\n\n node2 = nnx.merge(graphdef, state)\n\n assert node2.bar.kernel is node2.baz.kernel\n\n def test_tied_weights_example(self):\n class LinearTranspose(nnx.Module):\n def __init__(self, dout: int, din: int, *, rngs: nnx.Rngs) -> None:\n self.kernel = nnx.Param(\n nnx.initializers.lecun_normal()(rngs(), (dout, din))\n )\n\n def __call__(self, x):\n return x @ self.kernel.T\n\n class Encoder(nnx.Module):\n def __init__(self, *, rngs: nnx.Rngs) -> None:\n self.embed = nnx.Embed(10, 2, rngs=rngs)\n ...\n self.linear_out = LinearTranspose(10, 2, rngs=rngs)\n\n # tie the weights\n self.linear_out.kernel = self.embed.embedding\n\n def __call__(self, x):\n x = self.embed(x)\n ...\n return self.linear_out(x)\n\n model = Encoder(rngs=nnx.Rngs(0))\n graphdef, state = nnx.split(model)\n\n assert len(nnx.to_flat_state(state)) == 1\n\n x = jax.random.randint(jax.random.key(0), (2,), 0, 10)\n y = model(x)\n\n assert y.shape == (2, 10)\n\n def test_state_variables_shared_with_graph(self):\n class Foo(nnx.Module):\n def __init__(self):\n self.a = nnx.Param(jnp.array(1))\n\n m = Foo()\n graphdef, state = nnx.split(m)\n\n assert isinstance(m.a, nnx.Param)\n assert isinstance(state['a'], nnx.Param)\n assert m.a is state['a']\n assert m.a[...] == state['a'][...]\n\n m2 = nnx.merge(graphdef, state)\n\n assert isinstance(m2.a, nnx.Param)\n assert isinstance(state['a'], nnx.Param)\n assert m2.a is state['a']\n assert m2.a[...] == state['a'][...]\n\n def test_shared_state_variables_shared_with_graph(self):\n class Foo(nnx.Module):\n def __init__(self):\n p = nnx.Param(jnp.array(1))\n self.a = p\n self.b = p\n\n m = Foo()\n graphdef, state = nnx.split(m)\n\n assert isinstance(m.a, nnx.Param)\n assert isinstance(m.b, nnx.Param)\n assert isinstance(state['a'], nnx.Param)\n assert 'b' not in state\n assert m.a is state['a']\n assert m.b is state['a']\n assert m.a[...] == state['a'][...]\n assert m.b[...] == state['a'][...]\n\n m2 = nnx.merge(graphdef, state)\n\n assert isinstance(m2.a, nnx.Param)\n assert isinstance(m2.b, nnx.Param)\n assert isinstance(state['a'], nnx.Param)\n assert m2.a is state['a']\n assert m2.b is state['a']\n assert m2.a[...] == state['a'][...]\n assert m2.b[...] == state['a'][...]\n assert m2.a is m2.b\n\n def test_pytree_flatten(self):\n @struct.dataclass\n class Tree:\n a: int\n b: str = struct.field(pytree_node=False)\n\n p = Tree(1, 'a')\n\n leaves, treedef = nnx.graphlib._flatten_pytree(p)\n fields = dict(leaves)\n\n assert 'a' in fields\n assert 'b' not in fields\n assert fields['a'] == 1\n\n p2 = nnx.graphlib._unflatten_pytree(leaves, treedef)\n\n assert isinstance(p2, Tree)\n assert p2.a == 1\n\n def test_pytree_node(self):\n @struct.dataclass\n class Tree:\n a: nnx.Param[int]\n b: str = struct.field(pytree_node=False)\n\n class Foo(nnx.Module):\n def __init__(self):\n self.tree = nnx.data(Tree(nnx.Param(1), 'a'))\n\n m = Foo()\n\n graphdef, state = nnx.split(m)\n\n assert 'tree' in state\n assert 'a' in state['tree']\n\n m2 = nnx.merge(graphdef, state)\n\n assert isinstance(m2.tree, Tree)\n assert m2.tree.a.get_value() == 1\n assert m2.tree.b == 'a'\n assert m2.tree.a is m.tree.a\n assert m2.tree is not m.tree\n\n def test_cached_unflatten(self):\n class Foo(nnx.Module):\n def __init__(self, *, rngs: nnx.Rngs):\n self.a = nnx.Linear(2, 2, rngs=rngs)\n self.b = nnx.BatchNorm(2, rngs=rngs)\n\n def f(m: Foo):\n m.a, m.b = m.b, m.a # type: ignore\n\n m = Foo(rngs=nnx.Rngs(0))\n a = m.a\n b = m.b\n\n ref_out_idx_out = nnx.graphlib.RefMap()\n graphdef: nnx.graphlib.GraphDef[Foo]\n graphdef, state = nnx.graphlib.flatten(m, ref_index=ref_out_idx_out, graph=True)\n state = state.to_nested_state()\n\n @partial(jax.jit, static_argnums=(0,))\n def f_pure(graphdef: nnx.graphlib.GraphDef[Foo], state):\n idx_out_ref_in = nnx.graphlib.IndexMap()\n m = nnx.graphlib.unflatten(graphdef, state, index_ref=idx_out_ref_in)\n ref_in_idx_out = nnx.graphlib.RefMap.from_indexmap(idx_out_ref_in)\n f(m)\n ref_in_idx_in = nnx.graphlib.RefMap()\n graphdef, state = nnx.graphlib.flatten(\n m, ref_index=ref_in_idx_in, ref_outer_index=ref_in_idx_out, graph=True\n )\n state = state.to_nested_state()\n return state, graphdef\n\n state, graphdef_out = f_pure(graphdef, state)\n idx_out_ref_out = nnx.graphlib.IndexMap.from_refmap(ref_out_idx_out)\n m2 = nnx.graphlib.unflatten(\n graphdef_out, state, outer_index_outer_ref=idx_out_ref_out\n )\n assert m2 is m\n assert m2.a is b\n assert m2.b is a\n\n def test_cached_unflatten_swap_variables(self):\n class Foo(nnx.Module):\n def __init__(self):\n self.a = nnx.Param(1)\n self.b = nnx.Param(2)\n\n def f(m: Foo):\n m.a, m.b = m.b, m.a\n\n m = Foo()\n a = m.a\n b = m.b\n\n ref_out_idx_out = nnx.graphlib.RefMap()\n graphdef: nnx.graphlib.GraphDef[Foo]\n graphdef, state = nnx.graphlib.flatten(m, ref_index=ref_out_idx_out, graph=True)\n idx_out_ref_out = {v: k for k, v in ref_out_idx_out.items()}\n state = state.to_nested_state()\n\n @partial(jax.jit, static_argnums=(0,))\n def f_pure(graphdef: nnx.graphlib.GraphDef[Foo], state):\n idx_out_ref_in = nnx.graphlib.IndexMap()\n m = nnx.graphlib.unflatten(graphdef, state, index_ref=idx_out_ref_in)\n ref_in_idx_out = nnx.graphlib.RefMap.from_indexmap(idx_out_ref_in)\n f(m)\n ref_in_idx_in = nnx.graphlib.RefMap()\n graphdef, state = nnx.graphlib.flatten(\n m, ref_index=ref_in_idx_in, ref_outer_index=ref_in_idx_out, graph=True\n )\n state = state.to_nested_state()\n return state, graphdef\n\n state, graphdef = f_pure(graphdef, state)\n m2 = nnx.graphlib.unflatten(\n graphdef, state, outer_index_outer_ref=idx_out_ref_out\n )\n assert m2 is m\n assert m2.a is b\n assert m2.b is a\n\n def test_cached_unflatten_add_self_reference(self):\n class Foo(nnx.Module):\n def __init__(self):\n self.ref = nnx.data(None)\n\n def f(m: Foo):\n m.ref = m\n\n m = Foo()\n\n ref_out_idx_out = nnx.graphlib.RefMap()\n graphdef: nnx.graphlib.GraphDef[Foo]\n graphdef, state = nnx.graphlib.flatten(m, ref_index=ref_out_idx_out, graph=True)\n idx_out_ref_out = nnx.graphlib.IndexMap.from_refmap(ref_out_idx_out)\n state = state.to_nested_state()\n\n @partial(jax.jit, static_argnums=(0,))\n def f_pure(graphdef: nnx.graphlib.GraphDef[Foo], state):\n idx_out_ref_in = nnx.graphlib.IndexMap()\n m = nnx.graphlib.unflatten(graphdef, state, index_ref=idx_out_ref_in)\n ref_in_idx_out = nnx.graphlib.RefMap.from_indexmap(idx_out_ref_in)\n f(m)\n ref_in_idx_in = nnx.graphlib.RefMap()\n graphdef, state = nnx.graphlib.flatten(\n m, ref_index=ref_in_idx_in, ref_outer_index=ref_in_idx_out, graph=True\n )\n state = state.to_nested_state()\n return state, graphdef\n\n state, graphdef_out = f_pure(graphdef, state)\n m2 = nnx.graphlib.unflatten(\n graphdef_out, state, outer_index_outer_ref=idx_out_ref_out\n )\n assert m2 is m\n assert m2.ref is m2\n\n def test_call_jit_update(self):\n class Counter(nnx.Module):\n def __init__(self):\n self.count = nnx.Param(jnp.zeros(()))\n\n def inc(self):\n self.count[...] += 1\n return 1\n\n graph_state = nnx.split(Counter())\n\n @jax.jit\n def update(graph_state: nnx.PureState[Counter]):\n out, graph_state = nnx.call(graph_state).inc()\n self.assertEqual(out, 1)\n return graph_state\n\n graph_state = update(graph_state)\n graph_state = update(graph_state)\n\n counter = nnx.merge(*graph_state)\n\n self.assertEqual(counter.count[...], 2)\n\n def test_stateful_linear(self):\n linear = StatefulLinear(3, 2, nnx.Rngs(0))\n linear_state = nnx.split(linear)\n\n @jax.jit\n def forward(x, pure_linear: nnx.PureState[StatefulLinear]):\n y, pure_linear = nnx.call(pure_linear)(x)\n return y, pure_linear\n\n x = jnp.ones((1, 3))\n y, linear_state = forward(x, linear_state)\n y, linear_state = forward(x, linear_state)\n\n self.assertEqual(linear.count[...], 0)\n new_linear = nnx.merge(*linear_state)\n self.assertEqual(new_linear.count[...], 2)\n\n def test_getitem(self):\n rngs = nnx.Rngs(0)\n nodes = dict(\n a=StatefulLinear(3, 2, rngs),\n b=StatefulLinear(2, 1, rngs),\n )\n node_state = nnx.split(nodes)\n _, node_state = nnx.call(node_state)['b'].increment()\n\n nodes = nnx.merge(*node_state)\n\n self.assertEqual(nodes['a'].count[...], 0)\n self.assertEqual(nodes['b'].count[...], 1)\n\n def test_object_state_propagation(self):\n test = self\n\n class Foo(nnx.Module):\n def __call__(self):\n test.assertTrue(self._pytree__state.initializing)\n self = nnx.merge(*nnx.split(self))\n test.assertTrue(self._pytree__state.initializing)\n\n module = Foo()\n nnx.bridge.lazy_init(module)\n\n def test_object_state_propagation_nested(self):\n class NNXOuter(nnx.Module):\n def __init__(self, dout: int, rngs: nnx.Rngs):\n self.inner = nnx.bridge.ToNNX(linen.Dense(dout), rngs=rngs)\n self.rngs = rngs\n\n def __call__(self, x):\n @nnx.split_rngs(splits=5)\n @nnx.vmap(in_axes=(0, None), axis_size=5)\n def vmap_fn(inner, x):\n return inner(x)\n\n return vmap_fn(self.inner, x)\n\n x = jax.random.normal(jax.random.key(0), (2, 4))\n model = NNXOuter(3, rngs=nnx.Rngs(0))\n nnx.bridge.lazy_init(model, x)\n\n self.assertEqual(model.inner.kernel.shape, (5, 4, 3))\n self.assertEqual(model.inner.bias.shape, (5, 3))\n\n def test_split_merge_context(self):\n m = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n with nnx.graphlib.split_context() as ctx:\n graphdef1, state1 = ctx.split(m)\n graphdef2, state2 = ctx.split(m)\n\n self.assertFalse(hasattr(ctx, 'ref_index'))\n self.assertFalse(hasattr(ctx, 'ctxtag'))\n self.assertIsInstance(graphdef1.nodes[0], nnx.graphlib.NodeDef)\n self.assertIsInstance(graphdef2.nodes[0], nnx.graphlib.NodeRef)\n self.assertLen(nnx.to_flat_state(state1), 2)\n self.assertLen(nnx.to_flat_state(state2), 0)\n\n with nnx.graphlib.merge_context() as ctx:\n m1 = ctx.merge(graphdef1, state1)\n m2 = ctx.merge(graphdef2, state2)\n\n self.assertIs(m1, m2)\n self.assertFalse(hasattr(ctx, 'index_ref'))\n self.assertFalse(hasattr(ctx, 'ctxtag'))\n\n def test_split_merge_context_example(self):\n m1 = nnx.Dict({})\n with nnx.update_context('example'):\n with nnx.split_context('example') as ctx:\n graphdef, state = ctx.split(m1)\n\n @jax.jit\n def f(graphdef, state):\n with nnx.merge_context('example', True) as ctx:\n m2 = ctx.merge(graphdef, state)\n m2.a = 1\n m2.ref = m2 # create a reference cycle\n with nnx.split_context('example') as ctx:\n return ctx.split(m2)\n\n graphdef_out, state_out = f(graphdef, state)\n with nnx.merge_context('example', False) as ctx:\n m3 = ctx.merge(graphdef_out, state_out)\n\n def test_split_merge_context_nested(self):\n m2 = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n m1 = nnx.Sequential(m2)\n with nnx.graphlib.split_context() as ctx:\n graphdef1, state1 = ctx.split(m1)\n graphdef2, state2 = ctx.split(m2)\n\n self.assertIsInstance(graphdef1.nodes[0], nnx.graphlib.NodeDef)\n self.assertIsInstance(graphdef2.nodes[0], nnx.graphlib.NodeRef)\n self.assertLen(nnx.to_flat_state(state1), 2)\n self.assertLen(nnx.to_flat_state(state2), 0)\n\n with nnx.graphlib.merge_context() as ctx:\n m1 = ctx.merge(graphdef1, state1)\n m2 = ctx.merge(graphdef2, state2)\n\n self.assertIs(m2, m1.layers[0])\n self.assertFalse(hasattr(ctx, 'index_ref'))\n self.assertFalse(hasattr(ctx, 'ctxtag'))\n\n def test_split_merge_update_context(self):\n class Foo(nnx.Module):\n def __init__(self):\n self.a = nnx.Param(1)\n self.b = nnx.data(2)\n\n m = Foo()\n ctxtag = 'test'\n\n with nnx.update_context(ctxtag):\n with nnx.graphlib.split_context(ctxtag) as ctx:\n graphdef1, state1 = ctx.split(m)\n graphdef2, state2 = ctx.split(m)\n\n self.assertFalse(hasattr(ctx, 'ref_index'))\n self.assertFalse(hasattr(ctx, 'ctxtag'))\n self.assertIsInstance(graphdef1.nodes[0], nnx.graphlib.NodeDef)\n self.assertIsInstance(graphdef2.nodes[0], nnx.graphlib.NodeRef)\n self.assertLen(nnx.to_flat_state(state1), 2)\n self.assertLen(nnx.to_flat_state(state2), 0)\n\n @jax.jit\n def f(graphdef1, state1, graphdef2, state2):\n with nnx.graphlib.merge_context(ctxtag, True) as ctx:\n m1 = ctx.merge(graphdef1, state1)\n m2 = ctx.merge(graphdef2, state2)\n\n self.assertIs(m1, m2)\n self.assertFalse(hasattr(ctx, 'index_ref'))\n self.assertFalse(hasattr(ctx, 'ctxtag'))\n\n # swap a and b\n m1.a, m1.b = m1.b, m1.a\n\n with nnx.graphlib.split_context(ctxtag) as ctx:\n graphdef1, state1 = ctx.split(m1)\n graphdef2, state2 = ctx.split(m2)\n\n return graphdef1, state1, graphdef2, state2\n\n graphdef1, state1, graphdef2, state2 = f(\n graphdef1, state1, graphdef2, state2\n )\n\n with nnx.graphlib.merge_context(ctxtag, False) as ctx:\n m1_out = ctx.merge(graphdef1, state1)\n m2_out = ctx.merge(graphdef2, state2)\n\n self.assertIs(m, m1_out)\n self.assertIs(m, m2_out)\n self.assertEqual(m.a, 2)\n self.assertEqual(m.b[...], 1) # type: ignore\n\n self.assertFalse(hasattr(ctx, 'index_ref'))\n self.assertFalse(hasattr(ctx, 'ctxtag'))\n\n def test_to_tree_simple(self):\n m = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n impure_tree = (m, 1, {'b': m})\n\n pure_tree = nnx.to_tree(impure_tree)\n\n t1 = pure_tree[0]\n t2 = pure_tree[2]['b']\n\n self.assertEqual(pure_tree[1], 1)\n self.assertIsInstance(t1, nnx.NodeStates)\n assert isinstance(t1, nnx.NodeStates)\n self.assertIsInstance(t2, nnx.NodeStates)\n assert isinstance(t2, nnx.NodeStates)\n self.assertIsInstance(t1.graphdef.nodes[0], nnx.graphlib.NodeDef)\n self.assertIsInstance(t2.graphdef.nodes[0], nnx.graphlib.NodeRef)\n self.assertLen(nnx.to_flat_state(t1.states[0]), 2)\n self.assertLen(nnx.to_flat_state(t2.states[0]), 0)\n\n impure_tree2 = nnx.from_tree(pure_tree)\n\n m1_out = impure_tree2[0]\n m2_out = impure_tree2[2]['b']\n\n self.assertIs(m1_out, m2_out)\n self.assertEqual(impure_tree2[1], 1)\n\n def test_to_tree_update_context(self):\n class Foo(nnx.Module):\n def __init__(self):\n self.a = nnx.Param(1)\n self.b = nnx.data(2)\n\n m = Foo()\n impure_tree = (m, 1, {'b': m})\n ctxtag = 'test'\n\n with nnx.update_context(ctxtag):\n pure_tree = nnx.to_tree(impure_tree, ctxtag=ctxtag)\n\n t1 = pure_tree[0]\n t2 = pure_tree[2]['b']\n\n self.assertEqual(pure_tree[1], 1)\n self.assertIsInstance(t1, nnx.NodeStates)\n assert isinstance(t1, nnx.NodeStates)\n self.assertIsInstance(t2, nnx.NodeStates)\n assert isinstance(t2, nnx.NodeStates)\n self.assertIsInstance(t1.graphdef.nodes[0], nnx.graphlib.NodeDef)\n self.assertIsInstance(t2.graphdef.nodes[0], nnx.graphlib.NodeRef)\n self.assertLen(nnx.to_flat_state(t1.states[0]), 2)\n self.assertLen(nnx.to_flat_state(t2.states[0]), 0)\n\n @jax.jit\n def f(pure_tree):\n impure_tree2 = nnx.from_tree(pure_tree, ctxtag=ctxtag, is_inner=True)\n m1_out = impure_tree2[0]\n m2_out = impure_tree2[2]['b']\n\n self.assertIs(m1_out, m2_out)\n # self.assertEqual(impure_tree2[1], 1)\n\n # swap a and b\n m1_out.a, m1_out.b = m1_out.b, m1_out.a\n\n pure_tree2 = nnx.to_tree(impure_tree2, ctxtag=ctxtag)\n\n t1 = pure_tree2[0]\n t2 = pure_tree2[2]['b']\n\n # self.assertEqual(pure_tree2[1], 1)\n self.assertIsInstance(t1, nnx.NodeStates)\n assert isinstance(t1, nnx.NodeStates)\n self.assertIsInstance(t2, nnx.NodeStates)\n assert isinstance(t2, nnx.NodeStates)\n self.assertIsInstance(t1.graphdef.nodes[0], nnx.graphlib.NodeDef)\n self.assertIsInstance(t2.graphdef.nodes[0], nnx.graphlib.NodeRef)\n self.assertLen(nnx.to_flat_state(t1.states[0]), 2)\n self.assertLen(nnx.to_flat_state(t2.states[0]), 0)\n\n return pure_tree2\n\n pure_tree2 = f(pure_tree)\n\n impure_tree2 = nnx.from_tree(pure_tree2, ctxtag=ctxtag, is_inner=False)\n\n m1_out = impure_tree2[0]\n m2_out = impure_tree2[2]['b']\n\n self.assertIs(m, m1_out)\n self.assertIs(m, m2_out)\n self.assertEqual(m.a, 2)\n self.assertEqual(m.b[...], 1) # type: ignore\n self.assertEqual(impure_tree2[1], 1)\n\n def test_graph_flatten_with_data_wrapper(self):\n class Foo(nnx.Pytree):\n def __init__(self, data, static):\n self.data = nnx.data(data)\n self.static = nnx.static(static)\n\n tree = Foo(1, 2)\n state = nnx.state(tree)\n\n self.assertIn('data', state)\n self.assertIsInstance(state['data'], int)\n self.assertEqual(state['data'], 1)\n self.assertNotIn('static', state)\n\n def test_to_tree_consistent_prefix(self):\n m = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n impure_tree = (m, 1, {'b': m})\n prefix = (0, None, 0)\n pure_tree = nnx.to_tree(impure_tree, prefix=prefix)\n\n prefix = (0, None, 1)\n with self.assertRaisesRegex(ValueError, 'Inconsistent aliasing detected'):\n nnx.to_tree(impure_tree, prefix=prefix)\n\n def test_simple_vmap(self):\n @dataclasses.dataclass(frozen=True)\n class StateAxes:\n params: Any\n batch_stats: Any\n\n class Foo(nnx.Module):\n def __init__(self, a, b):\n self.a = nnx.Param(a)\n self.b = nnx.BatchStat(b)\n\n ctxtag = 'test'\n with nnx.update_context(ctxtag):\n m1 = Foo(a=jnp.array(0), b=jnp.arange(5))\n m2 = Foo(a=jnp.array(1), b=jnp.array(2))\n\n args = (m1, m2, {'b': m1})\n m1_axes = StateAxes(None, 0)\n in_axes = (m1_axes, None, {'b': m1_axes})\n jax_in_axes = jax.tree.map(\n lambda x: nnx.NodeStates.from_prefixes((x.params, x.batch_stats))\n if isinstance(x, StateAxes)\n else x,\n in_axes,\n )\n out_axes = 0\n\n def split_fn(ctx: nnx.SplitContext, path, prefix, x):\n if isinstance(prefix, StateAxes):\n return nnx.NodeStates.from_split(\n *ctx.split(x, nnx.Param, nnx.BatchStat)\n )\n return nnx.NodeStates.from_split(*ctx.split(x))\n\n pure_args = nnx.to_tree(\n args, ctxtag=ctxtag, prefix=in_axes, split_fn=split_fn\n )\n\n @partial(jax.vmap, in_axes=jax_in_axes, out_axes=(jax_in_axes, out_axes))\n def f(*pure_args):\n args = nnx.from_tree(pure_args, ctxtag=ctxtag, is_inner=True)\n\n y = 0\n\n self.assertIs(args[0], args[2]['b'])\n for path, m in nnx.iter_graph(args):\n if isinstance(m, Foo):\n self.assertEqual(m.a.shape, ())\n self.assertEqual(m.b.shape, ())\n y += m.a + m.b\n\n args_out = nnx.extract.clear_non_graph_nodes(args)\n\n pure_args_out, y = nnx.to_tree(\n (args_out, y),\n prefix=(in_axes, out_axes),\n ctxtag=ctxtag,\n split_fn=split_fn,\n )\n return pure_args_out, y\n\n pure_args_out, y = f(*pure_args)\n\n args_out, y = nnx.from_tree(\n (pure_args_out, y), ctxtag=ctxtag, is_inner=False\n )\n\n self.assertEqual(y.shape, (5,))\n self.assertGreater(y.sum(), 5)\n self.assertIs(m1, args_out[0])\n self.assertIs(m1, args_out[2]['b'])\n self.assertIs(m2, args_out[1])\n\n @parameterized.parameters(True, False)\n def test_split_variable(self, graph):\n v = nnx.Param(1)\n graphdef, state = nnx.split(v, graph=graph)\n\n expected_type = nnx.graphlib.VariableDef if graph else nnx.graphlib.TreeNodeDef\n self.assertIsInstance(graphdef.nodes[0], expected_type)\n self.assertIsInstance(state, nnx.Variable)\n\n v2 = nnx.merge(graphdef, state)\n self.assertIsInstance(v2, nnx.Param)\n\n @parameterized.parameters(True, False)\n def test_split_filter_variable(self, graph):\n v = nnx.Param(1)\n graphdef, batch_stats, params, rest = nnx.split(\n v, nnx.BatchStat, nnx.Param, ..., graph=graph\n )\n\n expected_type = nnx.graphlib.VariableDef if graph else nnx.graphlib.TreeNodeDef\n self.assertIsInstance(graphdef.nodes[0], expected_type)\n self.assertIsInstance(params, nnx.Variable)\n self.assertIsInstance(batch_stats, nnx.State)\n self.assertEmpty(batch_stats)\n self.assertIsInstance(rest, nnx.State)\n self.assertEmpty(rest)\n\n v2 = nnx.merge(graphdef, batch_stats, params, rest)\n self.assertIsInstance(v2, nnx.Param)\n\n @parameterized.parameters(True, False)\n def test_split_update_variable(self, graph):\n v = nnx.Param(jnp.array(1))\n graphdef, state = nnx.split(v, graph=graph)\n\n expected_type = nnx.graphlib.VariableDef if graph else nnx.graphlib.TreeNodeDef\n self.assertIsInstance(graphdef.nodes[0], expected_type)\n self.assertIsInstance(state, nnx.Variable)\n\n state[...] = 2\n nnx.update(v, state)\n\n self.assertEqual(v[...], 2)\n\n @parameterized.parameters(True, False)\n def test_split_update_filter_variable(self, graph):\n v = nnx.Param(jnp.array(1))\n graphdef, batch_stats, params, rest = nnx.split(\n v, nnx.BatchStat, nnx.Param, ..., graph=graph\n )\n\n expected_type = nnx.graphlib.VariableDef if graph else nnx.graphlib.TreeNodeDef\n self.assertIsInstance(graphdef.nodes[0], expected_type)\n self.assertIsInstance(params, nnx.Variable)\n self.assertIsInstance(batch_stats, nnx.State)\n self.assertEmpty(batch_stats)\n self.assertIsInstance(rest, nnx.State)\n self.assertEmpty(rest)\n\n params[...] = 2\n nnx.update(v, batch_stats, params, rest)\n\n self.assertEqual(v[...], 2)\n\n @parameterized.parameters(\n (lambda: nnx.Param(1),),\n (lambda: 42,),\n (lambda: jnp.array([1, 2, 3]),),\n )\n def test_split_leaf(self, leaf_fn):\n leaf = leaf_fn()\n graphdef, state = nnx.split(leaf)\n\n out = nnx.merge(graphdef, state)\n self.assertIs(out, state)\n\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_jit_variable(self, graph, graph_updates):\n v = nnx.Param(1)\n\n @nnx.jit(graph=graph, graph_updates=graph_updates)\n def f(v):\n v[...] += 1\n\n f(v)\n\n np.testing.assert_allclose(v[...], 2)\n\n def test_jit_pytree_of_variables(self):\n v1 = nnx.Param(jnp.array(1))\n v2 = nnx.Param(jnp.array(2))\n vs = [v1, v1, v2]\n\n @nnx.jit\n def f(vs):\n self.assertIs(vs[0], vs[1])\n self.assertIsNot(vs[0], vs[2])\n vs[0][...] += 10\n\n f(vs)\n\n self.assertIs(vs[0], vs[1])\n self.assertIsNot(vs[0], vs[2])\n np.testing.assert_allclose(vs[0][...], 11)\n np.testing.assert_allclose(vs[2][...], 2)\n\n def test_variable_reference_in_module(self):\n class Foo(nnx.Module):\n def __init__(self, var):\n self.var = var\n\n var = nnx.Param(1)\n foo = Foo(var)\n\n @nnx.jit\n def increment_var(var, foo):\n self.assertIs(var, foo.var)\n var[...] += 1\n\n increment_var(var, foo)\n self.assertEqual(foo.var[...], 2)\n\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_variables_example(self, graph, graph_updates):\n def stateful_linear_init(din: int, dout: int, rngs: nnx.Rngs):\n w = nnx.Param(jax.random.normal(rngs(), (din, dout)))\n b = nnx.Param(jnp.zeros((dout,)))\n count = nnx.Variable(jnp.array(0))\n return w, b, count\n\n rngs = nnx.Rngs(0)\n w, b, count = stateful_linear_init(2, 3, rngs=rngs)\n\n @nnx.jit(graph=graph, graph_updates=graph_updates)\n def stateful_linear(w, b, count, x):\n count[...] += 1\n return x @ w + b[None]\n\n x = jax.random.normal(rngs(), (1, 2))\n y = stateful_linear(w, b, count, x)\n self.assertEqual(count[...], 1)\n\n y = stateful_linear(w, b, count, x)\n self.assertEqual(count[...], 2)\n self.assertEqual(y.shape, (1, 3))\n\n @parameterized.parameters(True, False)\n def test_array_attributes(self, graph):\n class Foo(nnx.Module):\n def __init__(self):\n self.a = jnp.array(1)\n self.b = 'yes'\n\n m = Foo()\n\n graphdef, state = nnx.split(m, graph=graph)\n\n self.assertLen(state, 1)\n self.assertIsInstance(state['a'], jax.Array)\n\n m2 = nnx.merge(graphdef, state)\n\n self.assertIsInstance(m2.a, jax.Array)\n self.assertEqual(m2.a, 1)\n self.assertEqual(m2.b, 'yes')\n\n def test_transform_array_attributes(self):\n class Foo(nnx.Module):\n def __init__(self):\n self.a = jnp.array(1)\n self.b = 'yes'\n\n m = Foo()\n\n @nnx.jit\n def f(m):\n m.a += 1\n self.assertEqual(m.b, 'yes')\n\n f(m)\n\n self.assertEqual(m.a, 2)\n\n def test_data_after_init(self):\n test = self\n class Foo(nnx.Module):\n def __init__(self):\n self.ls = []\n self.ls.append(jnp.array(1))\n\n with self.assertRaisesRegex(\n ValueError, 'Found unexpected data on value of type'\n ):\n m = Foo()\n\n def test_update_dict(self):\n node = {\n 'a': {\n 'b': 1,\n 'c': nnx.Param(jnp.array(2)),\n 'd': 3,\n },\n }\n\n updates = {\n 'a': {\n 'b': 4,\n 'c': jnp.array(10),\n },\n }\n\n nnx.update(node, updates)\n\n self.assertEqual(node['a']['b'], 4)\n self.assertEqual(node['a']['c'][...], 10)\n self.assertEqual(node['a']['d'], 3)\n\n def test_pop_dict(self):\n node = {\n 'a': {\n 'b': jnp.array(1),\n 'c': nnx.Param(jnp.array(2)),\n 'd': jnp.array(3.0),\n },\n }\n lt_2 = lambda _, x: x < 2\n popped = nnx.pop(node, (nnx.Param, lt_2))\n\n self.assertEqual(popped['a']['b'], 1)\n self.assertEqual(popped['a']['c'][...], 2)\n self.assertEqual(node['a']['d'], 3.0)\n self.assertLen(jax.tree.leaves(node), 1)\n self.assertLen(jax.tree.leaves(popped), 2)\n\n def test_iter_graph(self):\n arr0 = jnp.zeros(1)\n arr1 = jnp.zeros(1)\n var0 = nnx.Variable(jnp.zeros(1))\n var1 = nnx.Variable(jnp.zeros(1))\n\n child = nnx.Module()\n child.a = var0\n child.b = arr0\n child.c = var1\n child.d = var0\n child.e = arr1\n child.f = arr0\n\n root = nnx.Module()\n root.a = child\n root.b = var1\n root.c = arr1\n root.d = var1\n root.e = child\n root.f = var0\n root.g = arr1\n\n nodes = [node for _, node in nnx.iter_graph(root)]\n count = lambda e: sum(node is e for node in nodes)\n\n # All internal nodes must be visited exactly once.\n self.assertEqual(count(var0), 1)\n self.assertEqual(count(var1), 1)\n self.assertEqual(count(child), 1)\n self.assertEqual(count(root), 1)\n\n # Arrays must not be deduplicated.\n self.assertEqual(count(arr0), 2)\n self.assertEqual(count(arr1), 3)\n\n # Nodes must be yielded in DFS order.\n expected = [var0, arr0, var1, arr1, arr0, child, arr1, arr1, root]\n unique = (arr0, arr1, var0, var1, child, root)\n index = lambda e: next(node is e for node in unique)\n actual = [node for node in nodes if any(node is e for e in unique)]\n self.assertEqual(list(map(index, actual)), list(map(index, expected)))\n\n def test_cached_partial_docstring_example(self):\n import optax\n\n model = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n optimizer = nnx.Optimizer(model, optax.adamw(1e-3), wrt=nnx.Param)\n\n @nnx.jit\n def train_step(model, optimizer, x, y):\n def loss_fn(model):\n return jnp.mean((model(x) - y) ** 2)\n\n loss, grads = nnx.value_and_grad(loss_fn)(model)\n optimizer.update(model, grads)\n return loss\n\n cached_train_step = nnx.cached_partial(train_step, model, optimizer)\n\n for step in range(2):\n x, y = jnp.ones((10, 2)), jnp.ones((10, 3))\n loss = cached_train_step(x, y)\n self.assertIsInstance(loss, jax.Array)\n\n def test_find_duplicates(self):\n class SharedModules(nnx.Module):\n def __init__(self, rngs: nnx.Rngs):\n self.a = nnx.Linear(1, 1, rngs=rngs)\n self.b = nnx.Linear(1, 1, rngs=rngs)\n self.c = self.a # shared Module\n\n model = SharedModules(nnx.Rngs(0))\n duplicates = nnx.find_duplicates(model)\n\n self.assertLen(duplicates, 1)\n self.assertEqual(duplicates[0], [('a',), ('c',)])\n\n def test_resursive_map(self):\n class Foo(nnx.Pytree):\n def __init__(self, d):\n self.d = d\n\n foo1 = Foo(10)\n foo2 = Foo(20)\n bar = [foo1, foo2, foo1]\n n = 0\n\n def inc_d(path, node):\n nonlocal n\n if isinstance(node, Foo):\n n += 1\n node.d += 1\n return node\n\n bar2 = nnx.recursive_map(inc_d, bar)\n self.assertIs(bar2[0], bar2[2])\n self.assertEqual(bar2[0].d, 11)\n self.assertEqual(bar2[1].d, 21)\n self.assertEqual(n, 2)\n\n def test_resursive_map_replace(self):\n class Foo(nnx.Pytree):\n def __init__(self, d):\n self.d = d\n\n foo1 = Foo(10)\n foo2 = Foo(20)\n bar = [foo1, foo2, foo1]\n n = 0\n\n def swap(path, node):\n nonlocal n\n if isinstance(node, Foo):\n n += 1\n node = Foo(-node.d)\n return node\n\n bar2 = nnx.recursive_map(swap, bar)\n self.assertIs(bar2[0], bar2[2])\n self.assertEqual(bar2[0].d, -10)\n self.assertEqual(bar2[1].d, -20)\n self.assertEqual(n, 2)\n\n @parameterized.parameters(True, False)\n def test_recursive_map_with_list(self, graph):\n rngs = nnx.Rngs(0)\n model = nnx.Sequential(nnx.Linear(2, 3, rngs=rngs), nnx.relu, nnx.Linear(3, 4, rngs=rngs))\n\n def add_rank2_lora(_, node):\n if isinstance(node, nnx.Linear):\n return nnx.LoRA(node.in_features, 2, node.out_features, base_module=node, rngs=rngs)\n return node\n\n self.assertEqual(len(nnx.recursive_map(add_rank2_lora, model, graph=graph).layers), 3)\n\n def test_graphdef_hash_with_sequential(self):\n rngs = nnx.Rngs(0)\n net = nnx.Sequential(\n nnx.Linear(2, 1, rngs=rngs),\n )\n hash(nnx.graphdef(net))\n\n @nnx.set_graph_mode(False)\n def test_split_graph_error(self):\n v = nnx.Variable(jnp.array(1.0))\n with self.assertRaisesRegex(\n ValueError, 'found at paths'\n ):\n graphdef, state = nnx.split((v, v))\n\nclass SimpleModule(nnx.Module):\n pass\n\n\nclass TestThreading(parameterized.TestCase):\n def test_threading(self):\n x = SimpleModule()\n\n class MyThread(Thread):\n def run(self) -> None:\n nnx.graphlib.split(x)\n\n thread = MyThread()\n thread.start()\n thread.join()\n\n\nclass TestTreeFlatten(parameterized.TestCase):\n def test_tree_flatten_unflatten(self):\n a = {'a': 1, 'b': nnx.Param(jnp.array(2))}\n b = {'a': 5, 'b': nnx.Param(jnp.array(6))}\n g = [a, 3, b, nnx.Param(jnp.array(4))]\n\n graphdef, flat_state = nnx.graphlib.flatten(g, graph=False)\n\n self.assertIsInstance(graphdef.nodes[0], nnx.graphlib.TreeNodeDef)\n\n g2 = nnx.graphlib.unflatten(graphdef, flat_state)\n self.assertIsInstance(g2, list)\n self.assertLen(g2, 4)\n self.assertEqual(g2[0]['a'], 1)\n self.assertEqual(g2[1], 3)\n self.assertIsInstance(g2[0]['b'], nnx.Param)\n self.assertIsInstance(g2[3], nnx.Param)\n self.assertEqual(g2[0]['b'][...], 2)\n self.assertEqual(g2[3][...], 4)\n\n def test_tree_flatten_no_paths(self):\n a = {'a': 1, 'b': nnx.Param(jnp.array(2))}\n b = {'a': 5, 'b': nnx.Param(jnp.array(6))}\n g = [a, 3, b, nnx.Param(jnp.array(4))]\n\n graphdef, leaves = nnx.graphlib.flatten(g, with_paths=False, graph=False)\n self.assertIsInstance(graphdef.nodes[0], nnx.graphlib.TreeNodeDef)\n self.assertIsInstance(leaves, list)\n\n def test_tree_split_merge(self):\n a = {'a': 1, 'b': nnx.Param(jnp.array(2))}\n b = {'a': 5, 'b': nnx.Param(jnp.array(6))}\n g = [a, 3, b, nnx.Param(jnp.array(4))]\n\n graphdef, state = nnx.split(g, graph=False)\n g2 = nnx.merge(graphdef, state)\n\n self.assertIsInstance(g2, list)\n self.assertEqual(g2[0]['a'], 1)\n self.assertEqual(g2[1], 3)\n self.assertIsInstance(g2[0]['b'], nnx.Param)\n self.assertEqual(g2[0]['b'][...], 2)\n self.assertEqual(g2[3][...], 4)\n\n def test_tree_split_merge_module(self):\n m = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n\n graphdef, state = nnx.split(m, graph=False)\n self.assertIsInstance(graphdef.nodes[0], nnx.graphlib.TreeNodeDef)\n\n m2 = nnx.merge(graphdef, state)\n self.assertIsInstance(m2, nnx.Linear)\n self.assertEqual(m2.kernel.shape, (2, 3))\n self.assertEqual(m2.bias.shape, (3,))\n\n def test_tree_shared_variables_raises(self):\n v = nnx.Param(jnp.array(1))\n g = [v, v]\n\n with self.assertRaises(ValueError):\n nnx.split(g, graph=False)\n\n def test_tree_shared_refs_raises(self):\n ref = jax.new_ref(jnp.array(1.0))\n g = [ref, ref]\n\n with self.assertRaises(ValueError):\n nnx.split(g, graph=False)\n\n def test_tree_shared_variables_state_raises(self):\n v = nnx.Param(jnp.array(1))\n g = [v, v]\n\n with self.assertRaises(ValueError):\n nnx.state(g, graph=False)\n\n def test_tree_shared_variables_graphdef_raises(self):\n v = nnx.Param(jnp.array(1))\n g = [v, v]\n\n with self.assertRaises(ValueError):\n nnx.graphdef(g, graph=False)\n\n def test_tree_shared_variables_clone_raises(self):\n v = nnx.Param(jnp.array(1))\n g = [v, v]\n\n with self.assertRaises(ValueError):\n nnx.clone(g, graph=False)\n\n def test_tree_flatten_unflatten_ordering(self):\n m = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n graphdef, state = nnx.split(m, graph=False)\n\n tree_nodedef = graphdef.nodes[0]\n self.assertIsInstance(tree_nodedef, nnx.graphlib.TreeNodeDef)\n paths = [p for p, _ in tree_nodedef.path_index]\n self.assertEqual(paths, sorted(paths))\n\n m2 = nnx.merge(graphdef, state)\n np.testing.assert_array_equal(m.kernel[...], m2.kernel[...])\n np.testing.assert_array_equal(m.bias[...], m2.bias[...])\n\n def test_tree_flatten_dict(self):\n g = {'z': nnx.Param(jnp.array(1)), 'a': jnp.array(2)}\n graphdef, state = nnx.split(g, graph=False)\n g2 = nnx.merge(graphdef, state)\n self.assertEqual(g2['z'][...], 1)\n np.testing.assert_array_equal(g2['a'], jnp.array(2))\n\n def test_tree_flatten_tuple(self):\n g = (nnx.Param(jnp.array(1)), jnp.array(2), 3)\n graphdef, state = nnx.split(g, graph=False)\n g2 = nnx.merge(graphdef, state)\n self.assertIsInstance(g2, tuple)\n self.assertEqual(g2[0][...], 1)\n np.testing.assert_array_equal(g2[1], jnp.array(2))\n self.assertEqual(g2[2], 3)\n\n def test_tree_flatten_namedtuple(self):\n import collections\n Point = collections.namedtuple('Point', ['y', 'x'])\n g = Point(\n y=nnx.Param(jnp.array(1.0)),\n x=nnx.Param(jnp.array(2.0)),\n )\n graphdef, flat_state = nnx.graphlib.flatten(g, graph=False)\n self.assertLen(flat_state, 2)\n path, value = flat_state[0]\n self.assertEqual(path, ('x',))\n self.assertEqual(value, 2.0)\n path, value = flat_state[1]\n self.assertEqual(path, ('y',))\n self.assertEqual(value, 1.0)\n\n g2 = nnx.graphlib.unflatten(graphdef, flat_state)\n self.assertIsInstance(g2, Point)\n self.assertEqual(g2.y[...], 1.0)\n self.assertEqual(g2.x[...], 2.0)\n\n def test_tree_flatten_registered_dataclass(self):\n @jax.tree_util.register_dataclass\n @dataclasses.dataclass\n class MyData:\n z_param: Any\n a_value: Any\n\n g = MyData(\n z_param=nnx.Param(jnp.array(10.0)),\n a_value=jnp.array(20.0),\n )\n graphdef, flat_state = nnx.graphlib.flatten(g, graph=False)\n self.assertLen(flat_state, 2)\n path, value = flat_state[0]\n self.assertEqual(path, ('a_value',))\n np.testing.assert_array_equal(value, 20.0)\n path, value = flat_state[1]\n self.assertEqual(path, ('z_param',))\n self.assertEqual(value, 10.0)\n\n g2 = nnx.graphlib.unflatten(graphdef, flat_state)\n self.assertIsInstance(g2, MyData)\n self.assertEqual(g2.z_param[...], 10.0)\n np.testing.assert_array_equal(g2.a_value, jnp.array(20.0))\n\n def test_tree_flatten_nested_mixed(self):\n g = {\n 'b': [nnx.Param(jnp.array(1)), jnp.array(2)],\n 'a': (nnx.Param(jnp.array(3)), 4),\n }\n graphdef, flat_state = nnx.graphlib.flatten(g, graph=False)\n self.assertLen(flat_state, 4)\n path, value = flat_state[0]\n self.assertEqual(path, ('a', 0))\n self.assertEqual(value, 3)\n path, value = flat_state[1]\n self.assertEqual(path, ('a', 1))\n self.assertEqual(value, 4)\n path, value = flat_state[2]\n self.assertEqual(path, ('b', 0))\n self.assertEqual(value, 1)\n path, value = flat_state[3]\n self.assertEqual(path, ('b', 1))\n np.testing.assert_array_equal(value, 2)\n\n g2 = nnx.graphlib.unflatten(graphdef, flat_state)\n self.assertIsInstance(g2, dict)\n self.assertEqual(g2['b'][0][...], 1)\n np.testing.assert_array_equal(g2['b'][1], jnp.array(2))\n self.assertEqual(g2['a'][0][...], 3)\n self.assertEqual(g2['a'][1], 4)\n\n @parameterized.parameters(True, False)\n def test_iter_graph(self, graph):\n var0 = nnx.Variable(jnp.zeros(1))\n var1 = nnx.Variable(jnp.zeros(1))\n arr0 = jnp.zeros(1)\n\n child = nnx.Module()\n child.a = var0\n child.b = arr0\n child.c = var1\n\n root = nnx.Module()\n root.x = child\n root.y = jnp.ones(2)\n\n node_ids = [id(node) for _, node in nnx.iter_graph(root, graph=graph)]\n self.assertIn(id(var0), node_ids)\n self.assertIn(id(var1), node_ids)\n self.assertIn(id(child), node_ids)\n self.assertIn(id(root), node_ids)\n\n def test_iter_graph_tree_mode_shared_variable_raises(self):\n var = nnx.Variable(jnp.zeros(1))\n root = nnx.Module()\n root.a = var\n root.b = var\n\n with self.assertRaisesRegex(\n ValueError, 'found at paths'\n ):\n list(nnx.iter_graph(root, graph=False))\n\n def test_iter_graph_tree_mode_cycle_raises(self):\n a = nnx.List([1])\n b = nnx.List([2, a])\n a.append(b)\n\n with self.assertRaisesRegex(\n ValueError, 'found at paths'\n ):\n list(nnx.iter_graph(a, graph=False))\n\n @parameterized.parameters(True, False)\n def test_iter_modules(self, graph):\n model = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n\n modules = list(nnx.iter_modules(model, graph=graph))\n self.assertLen(modules, 1)\n path, m = modules[0]\n self.assertEqual(path, ())\n self.assertIs(m, model)\n\n @parameterized.parameters(True, False)\n def test_iter_modules_nested(self, graph):\n class Block(nnx.Module):\n def __init__(self, rngs):\n self.linear = nnx.Linear(2, 3, rngs=rngs)\n self.dropout = nnx.Dropout(0.5)\n\n model = Block(nnx.Rngs(0))\n modules = list(nnx.iter_modules(model, graph=graph))\n module_types = [type(m).__name__ for _, m in modules]\n self.assertIn('Block', module_types)\n self.assertIn('Linear', module_types)\n self.assertIn('Dropout', module_types)\n self.assertLen(modules, 3)\n\n def test_recursive_map_tree_mode(self):\n class Foo(nnx.Pytree):\n def __init__(self, d):\n self.d = d\n\n foo1 = Foo(10)\n foo2 = Foo(20)\n bar = [foo1, foo2]\n n = 0\n\n def inc_d(path, node):\n nonlocal n\n if isinstance(node, Foo):\n n += 1\n node.d += 1\n return node\n\n bar2 = nnx.recursive_map(inc_d, bar, graph=False)\n self.assertEqual(bar2[0].d, 11)\n self.assertEqual(bar2[1].d, 21)\n self.assertEqual(n, 2)\n\n def test_recursive_map_tree_mode_replace(self):\n class Foo(nnx.Pytree):\n def __init__(self, d):\n self.d = d\n\n foo1 = Foo(10)\n foo2 = Foo(20)\n bar = [foo1, foo2]\n n = 0\n\n def swap(path, node):\n nonlocal n\n if isinstance(node, Foo):\n n += 1\n node = Foo(-node.d)\n return node\n\n bar2 = nnx.recursive_map(swap, bar, graph=False)\n self.assertEqual(bar2[0].d, -10)\n self.assertEqual(bar2[1].d, -20)\n self.assertEqual(n, 2)\n\n def test_recursive_map_tree_mode_with_list(self):\n rngs = nnx.Rngs(0)\n model = nnx.Sequential(\n nnx.Linear(2, 3, rngs=rngs), nnx.relu, nnx.Linear(3, 4, rngs=rngs)\n )\n\n def add_rank2_lora(_, node):\n if isinstance(node, nnx.Linear):\n return nnx.LoRA(\n node.in_features, 2, node.out_features,\n base_module=node, rngs=rngs,\n )\n return node\n\n result = nnx.recursive_map(add_rank2_lora, model, graph=False)\n self.assertLen(result.layers, 3)\n\n def test_recursive_map_tree_mode_shared_variable_raises(self):\n v = nnx.Param(jnp.array(1))\n g = [v, v]\n\n with self.assertRaisesRegex(\n ValueError, 'found at paths'\n ):\n nnx.recursive_map(lambda path, node: node, g, graph=False)\n\n def test_recursive_map_tree_mode_cycle_raises(self):\n a = nnx.List([1])\n b = nnx.List([2, a])\n a.append(b)\n\n with self.assertRaisesRegex(\n ValueError, 'found at paths'\n ):\n nnx.recursive_map(lambda path, node: node, a, graph=False)\n\n def test_check_valid_pytree_flatten(self):\n class NotAPytree(nnx.Pytree, pytree=False):\n def __init__(self):\n self.x = 1\n\n node = [NotAPytree()]\n with self.assertRaisesRegex(\n ValueError, \"pytree=False.*Found at path\"\n ):\n nnx.graphlib.flatten(node, graph=False)\n\n def test_check_valid_pytree_iter_graph(self):\n class NotAPytree(nnx.Pytree, pytree=False):\n def __init__(self):\n self.x = 1\n\n node = nnx.List([NotAPytree()])\n with self.assertRaisesRegex(\n ValueError, \"pytree=False.*Found at path\"\n ):\n list(nnx.iter_graph(node, graph=False))\n\n def test_check_valid_pytree_iter_children(self):\n class NotAPytree(nnx.Pytree, pytree=False):\n def __init__(self):\n self.x = 1\n\n node = NotAPytree()\n with self.assertRaisesRegex(\n ValueError, \"pytree=False\"\n ):\n list(nnx.iter_children(node, graph=False))\n\n def test_check_valid_pytree_recursive_map(self):\n class NotAPytree(nnx.Pytree, pytree=False):\n def __init__(self):\n self.x = 1\n\n node = nnx.List([NotAPytree()])\n with self.assertRaisesRegex(\n ValueError, \"pytree=False.*Found at path\"\n ):\n nnx.recursive_map(lambda path, node: node, node, graph=False)\n\n @parameterized.parameters(True, False)\n def test_map(self, graph):\n model = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n new_model = nnx.map(lambda path, x: jnp.zeros_like(x), model, graph=graph)\n\n self.assertTrue(hasattr(new_model, 'kernel'))\n self.assertTrue(hasattr(new_model, 'bias'))\n np.testing.assert_array_equal(new_model.kernel, jnp.zeros((2, 3)))\n np.testing.assert_array_equal(new_model.bias, jnp.zeros((3,)))\n\n def test_map_with_path(self):\n model = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n paths_seen = []\n\n def record_path(path, x):\n paths_seen.append(path)\n return x\n\n nnx.map(record_path, model)\n self.assertLen(paths_seen, 2)\n path_last_parts = sorted(p[-1] for p in paths_seen)\n self.assertEqual(path_last_parts, ['bias', 'kernel'])\n\n def test_map_nested(self):\n class Model(nnx.Module):\n def __init__(self, rngs):\n self.linear = nnx.Linear(2, 3, rngs=rngs)\n\n model = Model(rngs=nnx.Rngs(0))\n new_model = nnx.map(lambda path, x: jnp.ones_like(x), model)\n\n self.assertTrue(hasattr(new_model, 'linear'))\n np.testing.assert_array_equal(new_model.linear.kernel, jnp.ones((2, 3)))\n np.testing.assert_array_equal(new_model.linear.bias, jnp.ones((3,)))\n\n def test_map_replace(self):\n model = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n new_model = nnx.map(\n lambda path, v: v.replace(jnp.zeros_like(v)), model\n )\n\n self.assertTrue(hasattr(new_model, 'kernel'))\n self.assertTrue(hasattr(new_model, 'bias'))\n self.assertIsInstance(new_model.kernel, nnx.Param)\n np.testing.assert_array_equal(new_model.kernel[...], jnp.zeros((2, 3)))\n np.testing.assert_array_equal(new_model.bias[...], jnp.zeros((3,)))\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/nnx/helpers_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport jax\nimport jax.numpy as jnp\nimport optax\n\nfrom absl.testing import absltest\nimport numpy as np\n\nfrom flax import linen\nfrom flax import nnx\n\n\nclass TrainState(nnx.TrainState):\n batch_stats: nnx.State\n\n\nclass TestHelpers(absltest.TestCase):\n def test_train_state(self):\n m = nnx.Dict(a=nnx.Param(1), b=nnx.BatchStat(2))\n\n graphdef, params, batch_stats = nnx.split(m, nnx.Param, nnx.BatchStat)\n\n state = TrainState.create(\n graphdef,\n params=params,\n tx=optax.sgd(1.0),\n batch_stats=batch_stats,\n )\n\n leaves = jax.tree_util.tree_leaves(state)\n\n def test_train_state_methods(self):\n class Foo(nnx.Module):\n def __init__(self, *, rngs: nnx.Rngs):\n self.linear = nnx.Linear(2, 4, rngs=rngs)\n self.batch_norm = nnx.BatchNorm(4, rngs=rngs)\n\n def __call__(self, x: jax.Array, train: bool) -> jax.Array:\n x = self.linear(x)\n x = self.batch_norm(x, use_running_average=not train)\n return x\n\n module = Foo(rngs=nnx.Rngs(0))\n graphdef, params, batch_stats = nnx.split(module, nnx.Param, nnx.BatchStat)\n\n state = TrainState.create(\n graphdef,\n params=params,\n tx=optax.sgd(1.0),\n batch_stats=batch_stats,\n )\n\n x = jax.numpy.ones((1, 2))\n y, _updates = state.apply('params', 'batch_stats')(x, train=True)\n\n assert y.shape == (1, 4)\n\n # fake gradient\n grads = jax.tree.map(jnp.ones_like, state.params)\n # test apply_gradients\n state = state.apply_gradients(grads)\n\n def test_nnx_linen_sequential_equivalence(self):\n key1, key2 = jax.random.split(jax.random.key(0), 2)\n rngs = nnx.Rngs(0)\n x = jax.random.uniform(key1, (3, 1, 5))\n\n model_nnx = nnx.Sequential(\n nnx.Linear(5, 4, rngs=rngs), nnx.Linear(4, 2, rngs=rngs)\n )\n model = linen.Sequential([linen.Dense(4), linen.Dense(2)])\n\n variables = model.init(key2, x)\n for layer_index in range(2):\n for param in ('kernel', 'bias'):\n variables['params'][f'layers_{layer_index}'][param] = getattr(\n model_nnx.layers[layer_index], param\n )[...]\n out_nnx = model_nnx(x)\n out = model.apply(variables, x)\n np.testing.assert_array_equal(out, out_nnx)\n\n variables = model.init(key2, x)\n for layer_index in range(2):\n for param in ('kernel', 'bias'):\n getattr(model_nnx.layers[layer_index], param)[...] = variables[\n 'params'\n ][f'layers_{layer_index}'][param]\n out_nnx = model_nnx(x)\n out = model.apply(variables, x)\n np.testing.assert_array_equal(out, out_nnx)\n\n def test_nnx_empty_sequential_is_identity(self):\n iden = nnx.Sequential()\n assert iden(12) == 12\n assert iden(12, 23) == (12, 23)\n assert iden() is None\n assert iden(k=2) == {'k': 2}\n\n def test_dict_mutable_mapping(self):\n d = nnx.Dict({'a': 1, 'b': 2})\n self.assertEqual(d['a'], 1)\n self.assertEqual(d['b'], 2)\n self.assertEqual(len(d), 2)\n\n d['c'] = 3\n self.assertEqual(d['c'], 3)\n self.assertEqual(len(d), 3)\n\n del d['a']\n self.assertEqual(len(d), 2)\n with self.assertRaises(KeyError):\n _ = d['a']\n\n self.assertSetEqual(set(d), {'b', 'c'})\n\n def test_dict_setdefault(self):\n d = nnx.Dict({'a': 1, 'b': 2})\n self.assertEqual(d.setdefault('a', 10), 1)\n self.assertEqual(d['a'], 1)\n\n self.assertEqual(d.setdefault('c', 3), 3)\n self.assertEqual(d['c'], 3)\n self.assertEqual(len(d), 3)\n\n def test_dict_contains(self):\n d = nnx.Dict({'a': 1, 'b': 2})\n self.assertIn('a', d)\n self.assertIn('b', d)\n self.assertNotIn('c', d)\n\n d['c'] = 3\n self.assertIn('c', d)\n\n del d['a']\n self.assertNotIn('a', d)\n\n def test_list_mutable_sequence(self):\n l = nnx.List([1, 2, 3])\n self.assertEqual(len(l), 3)\n self.assertEqual(l[0], 1)\n self.assertEqual(l[1], 2)\n self.assertEqual(l[2], 3)\n\n l.append(4)\n self.assertEqual(len(l), 4)\n self.assertEqual(l[3], 4)\n\n l.insert(1, 5)\n self.assertEqual(len(l), 5)\n self.assertEqual(l[0], 1)\n self.assertEqual(l[1], 5)\n self.assertEqual(l[2], 2)\n self.assertEqual(l[3], 3)\n self.assertEqual(l[4], 4)\n\n del l[2]\n self.assertEqual(len(l), 4)\n self.assertEqual(l[0], 1)\n self.assertEqual(l[1], 5)\n self.assertEqual(l[2], 3)\n self.assertEqual(l[3], 4)\n\n l[1:3] = [6, 7]\n self.assertEqual(l[1], 6)\n self.assertEqual(l[2], 7)\n\n self.assertEqual(l[1:3], [6, 7])\n\n def test_list_fori_loop(self):\n class Foo(nnx.Module):\n def __init__(self):\n self.layers = nnx.List([\n nnx.Linear(1, 1, rngs=nnx.Rngs(0)),\n nnx.Linear(1, 1, rngs=nnx.Rngs(0)),\n ])\n\n def batch_loop_body(i, carry):\n return carry\n\n net = Foo()\n jax.lax.fori_loop(0, 2, batch_loop_body, net)\n\n def test_list_pytree_default_behavior(self):\n ls = nnx.List([jnp.array(1), jnp.array(2), jnp.array(3)])\n leaves = jax.tree_util.tree_leaves(ls)\n self.assertLen(leaves, 3)\n np.testing.assert_array_equal(leaves[0], jnp.array(1))\n np.testing.assert_array_equal(leaves[1], jnp.array(2))\n np.testing.assert_array_equal(leaves[2], jnp.array(3))\n\n def test_list_pytree_static_elements(self):\n ls = nnx.List([nnx.static(10), nnx.static(20), nnx.static(30)])\n leaves = jax.tree_util.tree_leaves(ls)\n self.assertEmpty(leaves)\n\n def test_list_pytree_data_elements(self):\n ls = nnx.List([nnx.data(1), nnx.data(2), nnx.data(3)])\n leaves = jax.tree_util.tree_leaves(ls)\n self.assertLen(leaves, 3)\n self.assertEqual(leaves[0], 1)\n self.assertEqual(leaves[1], 2)\n self.assertEqual(leaves[2], 3)\n\n def test_list_pytree_mixed_static_data(self):\n ls = nnx.List([\n nnx.data(jnp.array(1)),\n nnx.static(100),\n nnx.data(jnp.array(2)),\n nnx.static(200),\n ])\n leaves = jax.tree_util.tree_leaves(ls)\n self.assertLen(leaves, 2)\n np.testing.assert_array_equal(leaves[0], jnp.array(1))\n np.testing.assert_array_equal(leaves[1], jnp.array(2))\n\n def test_list_pytree_flatten_unflatten(self):\n ls = nnx.List([nnx.data(10), nnx.static('hello'), nnx.data(20)])\n leaves, treedef = jax.tree_util.tree_flatten(ls)\n self.assertLen(leaves, 2)\n self.assertEqual(leaves[0], 10)\n self.assertEqual(leaves[1], 20)\n\n new_leaves = [x * 2 for x in leaves]\n new_ls = jax.tree_util.tree_unflatten(treedef, new_leaves)\n self.assertEqual(new_ls[0], 20)\n self.assertEqual(new_ls[1], 'hello')\n self.assertEqual(new_ls[2], 40)\n\n def test_list_pytree_jit(self):\n ls = nnx.List([nnx.data(jnp.array(1.0)), nnx.static(999)])\n\n @jax.jit\n def double(ls):\n return jax.tree.map(lambda x: x * 2, ls)\n\n result = double(ls)\n np.testing.assert_array_equal(result[0], jnp.array(2.0))\n self.assertEqual(result[1], 999)\n\n\nif __name__ == '__main__':\n absltest.main()\n\n" }, { "path": "tests/nnx/ids_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport copy\n\nfrom absl.testing import absltest\nfrom flax.nnx import ids\n\n\nclass TestIds(absltest.TestCase):\n def test_hashable(self):\n id1 = ids.uuid()\n id2 = ids.uuid()\n assert id1 == id1\n assert id1 != id2\n assert hash(id1) != hash(id2)\n id1c = copy.copy(id1)\n id1dc = copy.deepcopy(id1)\n assert hash(id1) != hash(id1c)\n assert hash(id1) != hash(id1dc)\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/nnx/integration_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport tempfile\nimport typing as tp\n\nfrom absl.testing import absltest\nfrom absl.testing import parameterized\nimport jax\nimport jax.numpy as jnp\nimport numpy as np\nimport orbax.checkpoint as ocp\nimport optax\n\nfrom flax import nnx\n\nA = tp.TypeVar('A')\n\n\nclass TestIntegration(parameterized.TestCase):\n\n @parameterized.parameters(True, False)\n def test_basic_view_example(self, graph_mode):\n class Model(nnx.Module):\n\n def __init__(self, din, dmid, dout, rngs: nnx.Rngs):\n self.linear = nnx.Linear(din, dmid, rngs=rngs)\n self.bn = nnx.BatchNorm(dmid, rngs=rngs)\n self.dropout = nnx.Dropout(0.2, rngs=rngs)\n self.linear_out = nnx.Linear(dmid, dout, rngs=rngs)\n\n def __call__(self, x):\n x = nnx.relu(self.dropout(self.bn(self.linear(x))))\n return self.linear_out(x)\n\n model = Model(2, 64, 3, rngs=nnx.Rngs(0)) # eager initialization\n train_model = nnx.view(\n model, deterministic=False, use_running_average=False\n )\n eval_model = nnx.view(\n model, deterministic=True, use_running_average=True\n )\n optimizer = nnx.Optimizer(train_model, optax.adam(1e-3), wrt=nnx.Param)\n\n self.assertEqual(train_model.dropout.deterministic, False)\n self.assertEqual(train_model.bn.use_running_average, False)\n self.assertEqual(eval_model.dropout.deterministic, True)\n self.assertEqual(eval_model.bn.use_running_average, True)\n self.assertIs(train_model.dropout.rngs.count, eval_model.dropout.rngs.count)\n\n @nnx.jit(graph=graph_mode) # automatic state management for JAX transforms\n def train_step(model, optimizer, x, y):\n def loss_fn(model):\n y_pred = model(x)\n return jnp.mean((y_pred - y) ** 2)\n\n loss, grads = nnx.value_and_grad(loss_fn)(model)\n optimizer.update(model, grads) # in-place updates\n\n return loss\n\n @nnx.jit(graph=graph_mode)\n def eval_step(model, x, y):\n y_pred = model(x)\n return jnp.mean((y_pred - y) ** 2)\n\n x = jax.random.normal(jax.random.key(0), (8, 2))\n y = jax.random.normal(jax.random.key(1), (8, 3))\n\n train_step(train_model, optimizer, x, y)\n self.assertEqual(train_model.dropout.rngs.count[...], 1)\n eval_step(eval_model, x, y)\n self.assertEqual(train_model.dropout.rngs.count[...], 1)\n\n def test_shared_modules(self):\n class Block(nnx.Module):\n def __init__(self, linear: nnx.Linear, *, rngs):\n self.linear = linear\n self.bn = nnx.BatchNorm(2, rngs=rngs)\n\n def __call__(self, x):\n x = self.linear(x)\n x = self.bn(x)\n return nnx.relu(x)\n\n class Model(nnx.Module):\n def __init__(self, *, rngs):\n shared = nnx.Linear(2, 2, rngs=rngs)\n self.block1 = Block(shared, rngs=rngs)\n self.block2 = Block(shared, rngs=rngs)\n\n def __call__(self, x):\n x = self.block1(x)\n x = self.block2(x)\n return x\n\n @nnx.jit\n def train_step(model: Model, x, y):\n @nnx.grad\n def loss_fn(model: Model):\n y_pred = model(x)\n return jnp.mean((y - y_pred) ** 2)\n\n grads = loss_fn(model)\n nnx.update(\n model,\n jax.tree.map(\n lambda w, g: w - 0.1 * g, nnx.state(model, nnx.Param), grads\n ),\n )\n\n model = Model(rngs=nnx.Rngs(0))\n\n x = np.random.uniform(size=(4, 2))\n y = np.random.uniform(size=(4, 2))\n model.set_attributes(use_running_average=False, graph=True)\n\n for _i in range(3):\n train_step(model, x, y)\n\n assert model.block1.linear is model.block2.linear\n assert model.block1.linear.bias is not None\n assert model.block1.bn is not model.block2.bn\n\n def test_shared_modules_view(self):\n class Block(nnx.Module):\n def __init__(self, linear: nnx.Linear, *, rngs):\n self.linear = linear\n self.bn = nnx.BatchNorm(2, rngs=rngs)\n\n def __call__(self, x):\n x = self.linear(x)\n x = self.bn(x)\n return nnx.relu(x)\n\n class Model(nnx.Module):\n def __init__(self, *, rngs):\n shared = nnx.Linear(2, 2, rngs=rngs)\n self.block1 = Block(shared, rngs=rngs)\n self.block2 = Block(shared, rngs=rngs)\n\n def __call__(self, x):\n x = self.block1(x)\n x = self.block2(x)\n return x\n\n @nnx.jit\n def train_step(model: Model, x, y):\n @nnx.grad\n def loss_fn(model: Model):\n y_pred = model(x)\n return jnp.mean((y - y_pred) ** 2)\n\n grads = loss_fn(model)\n nnx.update(\n model,\n jax.tree.map(\n lambda w, g: w - 0.1 * g, nnx.state(model, nnx.Param), grads\n ),\n )\n\n model = Model(rngs=nnx.Rngs(0))\n\n x = np.random.uniform(size=(4, 2))\n y = np.random.uniform(size=(4, 2))\n new_model = nnx.view(model, use_running_average=False)\n\n for _i in range(3):\n train_step(model, x, y)\n\n assert new_model.block1.linear is new_model.block2.linear\n assert new_model.block1.linear.bias is not None\n assert new_model.block1.bn is not new_model.block2.bn\n\n def test_shared_modules_pure(self):\n class Block(nnx.Module):\n def __init__(self, linear: nnx.Linear, *, rngs: nnx.Rngs):\n self.linear = linear\n self.bn = nnx.BatchNorm(2, rngs=rngs)\n\n def __call__(self, x):\n x = self.linear(x)\n x = self.bn(x)\n return nnx.relu(x)\n\n class Model(nnx.Module):\n def __init__(self, *, rngs: nnx.Rngs):\n shared = nnx.Linear(2, 2, rngs=rngs)\n self.block1 = Block(shared, rngs=rngs)\n self.block2 = Block(shared, rngs=rngs)\n\n def __call__(self, x):\n x = self.block1(x)\n x = self.block2(x)\n return x\n\n @jax.jit\n def train_step(state: nnx.State, graphdef: nnx.GraphDef[Model], x, y):\n model = nnx.merge(graphdef, state)\n model.set_attributes(use_running_average=False, graph=True)\n\n @nnx.grad\n def loss_fn(model: Model):\n y_pred = model(x)\n return jnp.mean((y - y_pred) ** 2)\n\n grads = loss_fn(model)\n nnx.update(\n model,\n jax.tree.map(\n lambda w, g: w - 0.1 * g, nnx.state(model, nnx.Param), grads\n ),\n )\n\n return nnx.split(model)\n\n graphdef: nnx.GraphDef[Model]\n graphdef, state = nnx.split(Model(rngs=nnx.Rngs(0)))\n\n x = np.random.uniform(size=(4, 2))\n y = np.random.uniform(size=(4, 2))\n\n for _i in range(3):\n graphdef, state = train_step(state, graphdef, x, y)\n\n model = nnx.merge(graphdef, state)\n\n assert model.block1.linear.bias is not None\n assert model.block2.linear.bias is not None\n assert model.block1.linear.kernel is model.block2.linear.kernel\n assert model.block1.linear.bias is model.block2.linear.bias\n assert model.block1.bn is not model.block2.bn\n\n def test_shared_modules_pure_view(self):\n class Block(nnx.Module):\n def __init__(self, linear: nnx.Linear, *, rngs: nnx.Rngs):\n self.linear = linear\n self.bn = nnx.BatchNorm(2, rngs=rngs)\n\n def __call__(self, x):\n x = self.linear(x)\n x = self.bn(x)\n return nnx.relu(x)\n\n class Model(nnx.Module):\n def __init__(self, *, rngs: nnx.Rngs):\n shared = nnx.Linear(2, 2, rngs=rngs)\n self.block1 = Block(shared, rngs=rngs)\n self.block2 = Block(shared, rngs=rngs)\n\n def __call__(self, x):\n x = self.block1(x)\n x = self.block2(x)\n return x\n\n @jax.jit\n def train_step(state: nnx.State, graphdef: nnx.GraphDef[Model], x, y):\n model = nnx.merge(graphdef, state)\n new_model = nnx.view(model, use_running_average=False, graph=True)\n\n @nnx.grad\n def loss_fn(model: Model):\n y_pred = model(x)\n return jnp.mean((y - y_pred) ** 2)\n\n grads = loss_fn(new_model)\n nnx.update(\n new_model,\n jax.tree.map(\n lambda w, g: w - 0.1 * g, nnx.state(new_model, nnx.Param), grads\n ),\n )\n\n return nnx.split(new_model)\n\n graphdef: nnx.GraphDef[Model]\n graphdef, state = nnx.split(Model(rngs=nnx.Rngs(0)))\n\n x = np.random.uniform(size=(4, 2))\n y = np.random.uniform(size=(4, 2))\n\n for _ in range(3):\n graphdef, state = train_step(state, graphdef, x, y)\n\n model = nnx.merge(graphdef, state)\n\n assert model.block1.linear.bias is not None\n assert model.block2.linear.bias is not None\n assert model.block1.linear.kernel is model.block2.linear.kernel\n assert model.block1.linear.bias is model.block2.linear.bias\n assert model.block1.bn is not model.block2.bn\n\n @parameterized.parameters(True, False)\n def test_stateful_example(self, graph_mode):\n class State(nnx.Variable[A]):\n pass\n\n class Linear(nnx.Module):\n def __init__(self, din: int, dout: int, *, rngs: nnx.Rngs):\n key = rngs.params()\n self.w = nnx.Param(jax.random.uniform(key, (din, dout)))\n self.b = nnx.Param(jnp.zeros((dout,)))\n self.count = State(jnp.array(0))\n\n def __call__(self, x):\n self.count[...] += 1\n return x @ self.w + self.b[None]\n\n model = Linear(din=12, dout=2, rngs=nnx.Rngs(0))\n # forward pass\n x = jnp.ones((8, 12))\n y = model(x)\n assert model.count[...] == 1\n\n @nnx.jit(graph=graph_mode)\n def train_step(model, x, y):\n def loss_fn(model):\n y_pred = model(x)\n return jax.numpy.mean((y_pred - y) ** 2)\n\n # compute gradient\n grads: nnx.State = nnx.grad(loss_fn)(model)\n # SGD update\n nnx.update(\n model,\n jax.tree.map(\n lambda w, g: w - 0.1 * g, nnx.state(model, nnx.Param), grads\n ),\n )\n\n # execute the training step\n train_step(model, x, y)\n assert model.count[...] == 2\n\n def test_functional_example(self):\n class Count(nnx.Variable[A]):\n pass\n\n class Linear(nnx.Module):\n def __init__(self, din: int, dout: int, *, rngs: nnx.Rngs):\n key = rngs.params()\n self.w = nnx.Param(jax.random.uniform(key, (din, dout)))\n self.b = nnx.Param(jnp.zeros((dout,)))\n self.count = Count(jnp.array(0))\n\n def __call__(self, x):\n self.count[...] += 1\n return x @ self.w + self.b[None]\n\n model = Linear(din=12, dout=2, rngs=nnx.Rngs(0))\n # forward pass\n x = jnp.ones((8, 12))\n y = model(x)\n assert model.count[...] == 1\n\n graphdef, params, counts = nnx.split(model, nnx.Param, Count)\n\n @jax.jit\n def train_step(params, counts, x, y):\n def loss_fn(params):\n model = nnx.merge(graphdef, params, counts, copy=True)\n loss = jax.numpy.mean((model(x) - y) ** 2)\n return loss, nnx.state(model, Count)\n\n # compute gradient\n grads, counts = jax.grad(loss_fn, has_aux=True)(params)\n # SGD update\n params = jax.tree.map(lambda w, g: w - 0.1 * g, params, grads)\n\n return params, counts\n\n # execute the training step\n params, counts = train_step(params, counts, x, y)\n model = nnx.merge(graphdef, params, counts)\n assert model.count[...] == 2\n\n def test_intermediates_example(self):\n class Linear(nnx.Module):\n def __init__(self, din: int, dout: int, *, rngs: nnx.Rngs):\n key = rngs.params()\n self.w = nnx.Param(jax.random.uniform(key, (din, dout)))\n self.b = nnx.Param(jnp.zeros((dout,)))\n\n def __call__(self, x):\n y = x @ self.w + self.b[None]\n self.y = nnx.Intermediate(y)\n return y\n\n model = Linear(12, 2, rngs=nnx.Rngs(0))\n\n y = model(jnp.ones((8, 12)))\n\n intermediates = nnx.pop(model, nnx.Intermediate)\n\n assert 'y' in intermediates\n\n def test_intermediates_example_functional(self):\n class Linear(nnx.Module):\n def __init__(self, din: int, dout: int, *, rngs: nnx.Rngs):\n key = rngs.params()\n self.w = nnx.Param(jax.random.uniform(key, (din, dout)))\n self.b = nnx.Param(jnp.zeros((dout,)))\n\n def __call__(self, x):\n y = x @ self.w + self.b[None]\n self.y = nnx.Intermediate(y)\n return y\n\n model = Linear(12, 2, rngs=nnx.Rngs(0))\n\n graphdef, state = nnx.split(model)\n\n y, (_, state) = graphdef.apply(state)(jnp.ones((8, 12)))\n\n intermediates, state = nnx.split_state(state, nnx.Intermediate, ...)\n\n assert 'y' in intermediates\n\n def test_replace_by_pure_dict(self):\n class MLPs(nnx.Module):\n def __init__(self, dim, rngs: nnx.Rngs):\n self.layers = nnx.List()\n for _ in range(4):\n self.layers.append(nnx.Linear(dim, dim, rngs=rngs, use_bias=False))\n\n def __call__(self, x):\n for layer in self.layers:\n x = layer(x)\n return x\n\n model = MLPs(4, rngs=nnx.Rngs(0))\n x = jax.random.normal(jax.random.key(42), (3, 4))\n assert model(x).shape == (3, 4)\n\n _, state = nnx.split(model)\n pure_dict_state = nnx.to_pure_dict(state)\n nnx.display(pure_dict_state)\n\n with tempfile.TemporaryDirectory() as tmpdir:\n ckpt_dir = ocp.test_utils.erase_and_create_empty(\n tmpdir + '/my-checkpoints/'\n )\n checkpointer = ocp.StandardCheckpointer()\n # checkpointer.save(ckpt_dir / 'state', state)\n checkpointer.save(ckpt_dir / 'pure_dict', pure_dict_state)\n\n # Restore as a pure dictionary.\n restored_pure_dict = checkpointer.restore(ckpt_dir / 'pure_dict')\n restored_pure_dict = nnx.statelib.restore_int_paths(restored_pure_dict)\n\n model = nnx.eval_shape(lambda: MLPs(4, rngs=nnx.Rngs(0)))\n nnx.update(model, restored_pure_dict)\n assert model(x).shape == (3, 4) # The model still works!\n\n @nnx.var_defaults(hijax=True)\n def test_example_mutable_arrays(self):\n class Model(nnx.Module):\n def __init__(self, din, dmid, dout, rngs: nnx.Rngs):\n self.linear = nnx.Linear(din, dmid, rngs=rngs)\n self.bn = nnx.BatchNorm(dmid, rngs=rngs)\n self.dropout = nnx.Dropout(0.2, rngs=rngs)\n self.linear_out = nnx.Linear(dmid, dout, rngs=rngs)\n\n def __call__(self, x):\n x = nnx.relu(self.dropout(self.bn(self.linear(x))))\n return self.linear_out(x)\n\n model = Model(2, 64, 3, rngs=nnx.Rngs(0)) # eager initialization\n optimizer = nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param)\n\n @jax.jit # automatic state management for JAX transforms\n def train_step(x, y):\n graphdef, params, nondiff = nnx.split(model, nnx.Param, ...)\n def loss_fn(params):\n model = nnx.merge(graphdef, params, nondiff)\n return ((model(x) - y) ** 2).mean() # call methods directly\n\n loss, grads = jax.value_and_grad(loss_fn)(\n nnx.vars_as(params, hijax=False)\n )\n optimizer.update(model, grads) # in-place updates\n\n return loss\n\n x = jax.random.normal(jax.random.key(0), (8, 2))\n y = jax.random.normal(jax.random.key(1), (8, 3))\n\n train_step(x, y)\n\n def test_tree_mode_train_step(self):\n model = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n optimizer = nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param)\n\n @nnx.jit(graph=False)\n def train_step(model, optimizer, x, y):\n def loss_fn(model):\n y_pred = model(x)\n return jnp.mean((y_pred - y) ** 2)\n\n loss, grads = nnx.value_and_grad(loss_fn, graph=False)(model)\n optimizer.update(model, grads)\n return loss\n\n x = jax.random.normal(jax.random.key(0), (4, 2))\n y = jax.random.normal(jax.random.key(1), (4, 3))\n\n loss0 = train_step(model, optimizer, x, y)\n loss1 = train_step(model, optimizer, x, y)\n self.assertLess(loss1, loss0)\n\n def test_tree_mode_multi_module(self):\n class Block(nnx.Module):\n def __init__(self, din, dout, *, rngs):\n self.linear = nnx.Linear(din, dout, rngs=rngs)\n self.bn = nnx.BatchNorm(dout, rngs=rngs)\n\n def __call__(self, x):\n x = self.linear(x)\n x = self.bn(x)\n return nnx.relu(x)\n\n class Model(nnx.Module):\n def __init__(self, *, rngs):\n self.block1 = Block(2, 2, rngs=rngs)\n self.block2 = Block(2, 2, rngs=rngs)\n\n def __call__(self, x):\n x = self.block1(x)\n x = self.block2(x)\n return x\n\n @nnx.jit(graph=False)\n def train_step(model: Model, x, y):\n @nnx.grad(graph=False)\n def loss_fn(model: Model):\n y_pred = model(x)\n return jnp.mean((y - y_pred) ** 2)\n\n grads = loss_fn(model)\n model = jax.tree.map(\n lambda w, g: w - 0.1 * g, model, grads,\n )\n return model\n\n model = Model(rngs=nnx.Rngs(0))\n\n x = np.random.uniform(size=(4, 2))\n y = np.random.uniform(size=(4, 2))\n model.set_attributes(use_running_average=False)\n\n for _i in range(3):\n model = train_step(model, x, y)\n\n def test_tree_mode_stateful(self):\n class State(nnx.Variable[A]):\n pass\n\n class Linear(nnx.Module):\n def __init__(self, din: int, dout: int, *, rngs: nnx.Rngs):\n key = rngs.params()\n self.w = nnx.Param(jax.random.uniform(key, (din, dout)))\n self.b = nnx.Param(jnp.zeros((dout,)))\n self.count = State(jnp.array(0))\n\n def __call__(self, x):\n self.count[...] = self.count[...] + 1\n return x @ self.w + self.b[None]\n\n model = Linear(din=12, dout=2, rngs=nnx.Rngs(0))\n x = jnp.ones((8, 12))\n y = model(x)\n assert model.count[...] == 1\n\n @nnx.jit(graph=False)\n def train_step(model, x, y):\n def loss_fn(model):\n y_pred = model(x)\n return jax.numpy.mean((y_pred - y) ** 2)\n\n grads = nnx.grad(loss_fn, graph=False, allow_int=True)(model)\n params = jax.tree.map(\n lambda w, g: w - 0.1 * g,\n nnx.state(model, nnx.Param),\n nnx.state(grads, nnx.Param),\n )\n nnx.update(model, params)\n\n train_step(model, x, y)\n assert model.count[...] == 2\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/nnx/metrics_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom absl.testing import absltest\nfrom absl.testing import parameterized\nfrom flax import nnx\nimport jax\nimport jax.numpy as jnp\nimport numpy as np\n\n\nclass TestMetrics(parameterized.TestCase):\n\n def test_split_merge(self):\n logits = jnp.array(\n [[-1.0, 1.0], [-1.0, 1.0], [-1.0, 1.0], [1.0, -1.0], [1.0, -1.0]]\n )\n labels = jnp.array([1, 1, 1, 1, 1])\n logits2 = jnp.array(\n [[-1.0, 1.0], [-1.0, 1.0], [-1.0, 1.0], [1.0, -1.0], [1.0, -1.0]]\n )\n labels2 = jnp.array([1, 1, 1, 1, 0])\n\n accuracy = nnx.metrics.Accuracy()\n accuracy.update(logits=logits, labels=labels)\n graphdef, state = accuracy.split()\n accuracy = nnx.merge(graphdef, state)\n self.assertEqual(accuracy.compute(), 0.6)\n accuracy.update(logits=logits2, labels=labels2)\n self.assertEqual(accuracy.compute(), 0.7)\n\n def test_welford(self):\n values = jax.random.normal(jax.random.key(0), (5, 2))\n\n welford = nnx.metrics.Welford()\n welford.update(values=values)\n graphdef, state = welford.split()\n welford = nnx.merge(graphdef, state)\n expected = nnx.metrics.Statistics(\n mean=values.mean(),\n standard_deviation=values.std(),\n standard_error_of_mean=values.std() / jnp.sqrt(values.size),\n )\n computed = welford.compute()\n self.assertAlmostEqual(computed.mean, expected.mean, )\n self.assertAlmostEqual(computed.standard_deviation, expected.standard_deviation)\n self.assertAlmostEqual(\n computed.standard_error_of_mean, expected.standard_error_of_mean\n )\n\n def test_welford_large(self):\n values = jax.random.normal(jax.random.key(0), (5, 2)) + 1e16\n\n welford = nnx.metrics.Welford()\n welford.update(values=values)\n graphdef, state = welford.split()\n welford = nnx.merge(graphdef, state)\n expected = nnx.metrics.Statistics(\n mean=values.mean(),\n standard_deviation=values.std(),\n standard_error_of_mean=values.std() / jnp.sqrt(values.size),\n )\n computed = welford.compute()\n self.assertAlmostEqual(computed.mean, expected.mean)\n self.assertAlmostEqual(computed.standard_deviation, expected.standard_deviation)\n self.assertAlmostEqual(\n computed.standard_error_of_mean, expected.standard_error_of_mean\n )\n\n def test_welford_many(self):\n values = jax.random.normal(jax.random.key(0), (50_000,))\n\n welford = nnx.metrics.Welford()\n welford.update(values=values)\n graphdef, state = welford.split()\n welford = nnx.merge(graphdef, state)\n computed = welford.compute()\n self.assertAlmostEqual(\n computed.mean, 0.0, delta=3 * computed.standard_error_of_mean\n )\n self.assertAlmostEqual(computed.standard_deviation, 1.0, places=2)\n\n @parameterized.product(with_mask=[True, False])\n def test_multimetric(self, with_mask):\n logits = jnp.array(\n [[-1.0, 1.0], [-1.0, 1.0], [-1.0, 1.0], [1.0, -1.0], [1.0, -1.0]]\n )\n labels = jnp.array([1, 1, 0, 1, 0])\n mask = labels > 0 if with_mask else None\n logits2 = jnp.array(\n [[-1.0, 1.0], [-1.0, 1.0], [-1.0, 1.0], [1.0, -1.0], [1.0, -1.0]]\n )\n labels2 = jnp.array([0, 1, 1, 1, 1])\n mask2 = labels2 > 0 if with_mask else None\n batch_loss = jnp.array([1, 2, 3, 4])\n batch_loss2 = jnp.array([3, 2, 1, 0])\n loss_mask2 = batch_loss2 > 1 if with_mask else None\n\n metrics = nnx.MultiMetric(\n accuracy=nnx.metrics.Accuracy(), loss=nnx.metrics.Average()\n )\n values = metrics.compute()\n self.assertTrue(jnp.isnan(values['accuracy']))\n self.assertTrue(jnp.isnan(values['loss']))\n\n metrics.update(logits=logits, labels=labels, values=batch_loss, mask={\"accuracy\": mask})\n values = metrics.compute()\n self.assertEqual(values['accuracy'], 0.6 if not with_mask else 2 / 3)\n self.assertEqual(values['loss'], 2.5)\n\n metrics.update(logits=logits2, labels=labels2, values=batch_loss2, mask={\"accuracy\": mask2, \"loss\": loss_mask2})\n values = metrics.compute()\n self.assertEqual(values['accuracy'], 0.5 if not with_mask else (2 + 2) / (3 + 4))\n self.assertEqual(values['loss'], 2.0 if not with_mask else 15 / 6)\n\n metrics.reset()\n values = metrics.compute()\n self.assertTrue(jnp.isnan(values['accuracy']))\n self.assertTrue(jnp.isnan(values['loss']))\n\n @parameterized.product(with_mask=[True, False])\n def test_multimetric_with_custom_metric(self, with_mask):\n class CustomAccuracy(nnx.metrics.Accuracy):\n # we use unused values arg is on purpose instead of using **kwargs\n # to reproduce the error if mask arg is injected to the custom metric without mask arg\n def update(self, y_preds, y_true, values):\n super().update(logits=y_preds, labels=y_true)\n\n logits = jnp.array(\n [[-1.0, 1.0], [-1.0, 1.0], [-1.0, 1.0], [1.0, -1.0], [1.0, -1.0]]\n )\n labels = jnp.array([1, 1, 0, 1, 0])\n batch_loss = jnp.array([1, 2, 3, 4])\n mask = batch_loss > 1 if with_mask else None\n\n metrics = nnx.MultiMetric(\n accuracy=CustomAccuracy(), loss=nnx.metrics.Average()\n )\n metrics.update(y_preds=logits, y_true=labels, values=batch_loss, mask={\"loss\": mask})\n values = metrics.compute()\n self.assertEqual(values['accuracy'], 0.6)\n self.assertEqual(values['loss'], 3.0 if with_mask else 2.5)\n\n def test_binary_classification_accuracy(self):\n logits = jnp.array([0.4, 0.7, 0.2, 0.6])\n labels = jnp.array([0, 1, 1, 1])\n logits2 = jnp.array([0.1, 0.9, 0.8, 0.3])\n labels2 = jnp.array([0, 1, 1, 0])\n accuracy = nnx.metrics.Accuracy(threshold=0.5)\n accuracy.update(logits=logits, labels=labels)\n self.assertEqual(accuracy.compute(), 0.75)\n accuracy.update(logits=logits2, labels=labels2)\n self.assertEqual(accuracy.compute(), 0.875)\n\n @parameterized.parameters(\n {\n 'logits': np.array([[[0.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0]]]),\n 'labels': np.array([0, 0, 0, 0]),\n 'threshold': None,\n 'error_msg': 'For multi-class classification'\n },\n {\n 'logits': np.array([0.0, 0.0, 0.0, 0.0]),\n 'labels': np.array([[0, 0], [0, 0]]),\n 'threshold': 0.5,\n 'error_msg': 'For binary classification'\n }\n )\n def test_accuracy_dims(self, logits, labels, threshold, error_msg):\n accuracy = nnx.metrics.Accuracy(threshold=threshold)\n with self.assertRaisesRegex(ValueError, error_msg):\n accuracy.update(logits=logits, labels=labels)\n\n @parameterized.product(\n with_mask=[True, False],\n scalar_values=[True, False],\n )\n def test_average(self, with_mask, scalar_values):\n average = nnx.metrics.Average()\n value1 = 2 if scalar_values else jnp.arange(5 * 3, dtype=jnp.int32).reshape(5, 3)\n value2 = 3 if scalar_values else jnp.arange(5 * 3 * 4, dtype=jnp.float32).reshape(5, 3, 4)\n if with_mask:\n list_masks = [\n jnp.ones(5) if scalar_values else jnp.where(value1 > 4, 1.0, 0.0),\n jnp.ones(5) if scalar_values else value2 > 10,\n ]\n else:\n list_masks = [None, None]\n\n list_values = [value1, value2]\n\n for mask, values in zip(list_masks, list_values):\n if with_mask and scalar_values:\n with self.assertRaisesRegex(ValueError, \"should be a jax array\"):\n average.update(mask=mask, values=values)\n else:\n average.update(mask=mask, values=values)\n\n if with_mask and scalar_values:\n return\n\n self.assertEqual(\n average.count,\n (\n len(list_values) if scalar_values\n else sum(m.sum() for m in list_masks) if with_mask\n else sum(v.size for v in list_values)\n )\n )\n self.assertEqual(\n average.total,\n (\n sum(list_values) if scalar_values\n else sum((v * m).sum() for m, v in zip(list_masks, list_values)) if with_mask\n else sum(v.sum() for v in list_values)\n )\n )\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/nnx/module_test.py", "content": "# Copyright 2024 The Flax Authors.\n\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom copy import deepcopy\nimport dataclasses\nimport pickle\nimport tempfile\nfrom typing import TypeVar\n\nfrom absl.testing import absltest\nfrom absl.testing import parameterized\nimport cloudpickle\nfrom flax import errors, nnx\nimport jax\nimport jax.numpy as jnp\nimport numpy as np\nimport flax\n\nA = TypeVar('A')\n\nfrom contextlib import contextmanager\n\n@contextmanager\ndef set_graph_mode(mode):\n old_mode = flax.config._read('nnx_graph_mode')\n try:\n flax.config.update('nnx_graph_mode', mode)\n yield\n finally:\n flax.config.update('nnx_graph_mode', old_mode)\n\n\nclass PytreeTest(absltest.TestCase):\n def test_pytree(self):\n class Foo(nnx.Pytree):\n def __init__(self, a, b):\n self.a = nnx.data(a)\n self.b = nnx.static(b)\n\n foo = Foo(a=1, b=2)\n\n self.assertEqual(jax.tree.leaves(foo), [1])\n\n def test_sequential_map(self):\n model = nnx.Sequential(nnx.Linear(2,8, rngs=nnx.Rngs(0)))\n jax.tree.map(lambda x: x + 1, model) # shouldn't error\n\n def test_sequential_has_leaves(self):\n model = nnx.Sequential(nnx.Linear(2,8, rngs=nnx.Rngs(0)))\n self.assertLen(jax.tree.leaves(model), 2)\n\n def test_consistent_attrs(self):\n class Foo(nnx.Pytree):\n def __init__(self, a, b, c):\n self.a = nnx.data(a)\n self.b = nnx.static(b)\n self.c = c\n\n foo = Foo(a=1, b=2, c=jnp.array(3))\n self.assertLen(jax.tree.leaves(foo), 2)\n\n foo.a = 3\n self.assertLen(jax.tree.leaves(foo), 2)\n\n foo.a = nnx.static(3)\n self.assertLen(jax.tree.leaves(foo), 1)\n\n foo.b = 4 # ok\n self.assertLen(jax.tree.leaves(foo), 1)\n\n foo.b = nnx.data(4)\n self.assertLen(jax.tree.leaves(foo), 2)\n\n foo.c = jnp.array(5) # ok\n self.assertLen(jax.tree.leaves(foo), 2)\n\n foo.c = nnx.static(5)\n self.assertLen(jax.tree.leaves(foo), 1)\n\n with self.assertRaisesRegex(\n ValueError,\n 'Found data on value of type',\n ):\n foo.a = ['hi', jnp.array(6)]\n\n with self.assertRaisesRegex(\n ValueError,\n 'Found data in value of type',\n ):\n foo.b = nnx.static(jnp.array(4))\n\n def test_assing_pytree_with_data(self):\n class Foo(nnx.Pytree):\n pass\n\n foo = Foo()\n with self.assertRaisesRegex(\n ValueError,\n 'Found data on value of type',\n ):\n foo.a = [nnx.Variable(1)]\n\n def test_consistent_attrs_frozen_dataclass(self):\n @nnx.dataclass\n class Foo(nnx.Pytree):\n a: int = nnx.data()\n b: int = nnx.static()\n c: jax.Array\n\n foo = Foo(a=1, b=2, c=jnp.array(3))\n self.assertLen(jax.tree.leaves(foo), 2)\n\n def test_consistent_attrs_dataclass_annotations(self):\n @dataclasses.dataclass\n class Foo(nnx.Pytree):\n a: nnx.Data[int]\n b: nnx.Static[int]\n c: jax.Array\n\n foo = Foo(a=1, b=2, c=jnp.array(3))\n self.assertLen(jax.tree.leaves(foo), 2)\n\n foo.a = 3\n self.assertLen(jax.tree.leaves(foo), 2)\n\n foo.a = nnx.static(3)\n self.assertLen(jax.tree.leaves(foo), 1)\n\n foo.b = 4 # ok\n self.assertLen(jax.tree.leaves(foo), 1)\n\n foo.b = nnx.data(4)\n self.assertLen(jax.tree.leaves(foo), 2)\n\n foo.c = jnp.array(5) # ok\n self.assertLen(jax.tree.leaves(foo), 2)\n\n foo.c = nnx.static(5)\n self.assertLen(jax.tree.leaves(foo), 1)\n\n with self.assertRaisesRegex(\n ValueError,\n 'Found data on value of type',\n ):\n foo.a = ['hi', jnp.array(6)]\n\n with self.assertRaisesRegex(\n ValueError,\n 'Found data in value of type',\n ):\n foo.b = nnx.static(jnp.array(4))\n\n def test_explicit_dont_change(self):\n class Foo(nnx.Pytree):\n def __init__(self):\n self.b = nnx.data(2)\n\n foo = Foo()\n self.assertEqual(jax.tree.leaves(foo), [2])\n foo.b = \"hello\"\n self.assertEqual(jax.tree.leaves(foo), [\"hello\"])\n\n def test_no_data_in_static(self):\n class Foo(nnx.Pytree):\n def __init__(self):\n self.a = nnx.static(jnp.array(1))\n\n with self.assertRaisesRegex(\n ValueError,\n 'Found data in value of type',\n ):\n foo = Foo()\n\nclass TestCapture(parameterized.TestCase):\n\n def test_vmap(self):\n\n class Foo(nnx.Module):\n def __init__(self, dim):\n self.w = nnx.Param(jax.random.normal(jax.random.key(0), dim))\n\n def __call__(self, x):\n x = self.perturb('grad_of_x', x)\n y = jnp.dot(x, self.w)\n self.sow(nnx.Intermediate, 'y', y)\n return y\n\n def pre_run(self, x):\n graphdef, intms, params = nnx.split(model, nnx.Intermediate, nnx.Param)\n def run(intms, params, x):\n return nnx.merge(graphdef, intms, params)(x)\n nnx.vmap(run, in_axes=(0, None, 0))(intms, params, x)\n\n @nnx.jit\n def train_step(model, perturbations, x):\n def loss_grad(model, perturbations, x):\n def loss(model, perturbations, x):\n loss, interms = nnx.capture(model, nnx.Intermediate, init=perturbations)(x)\n return loss, interms\n (grads, perturb_grads), sowed = nnx.grad(loss, argnums=(0, 1), has_aux=True)(model, perturbations, x)\n return grads, nnx.merge_state(perturb_grads, sowed)\n return nnx.vmap(loss_grad, in_axes=(None, 0, 0))(model, perturbations,x)\n\n model, x = Foo(4), jnp.ones((3, 4))\n\n pre_run_capture = nnx.capture(model.pre_run, nnx.Perturbation)\n _, perturbations = pre_run_capture(x)\n _, intermediates = train_step(model, perturbations, x)\n np.testing.assert_allclose(intermediates['grad_of_x'].get_value(),\n jnp.broadcast_to(model.w.get_value()[None, :], (3, 4)))\n self.assertEqual(intermediates['y'].get_value()[0].shape, (3,))\n\n\n @parameterized.parameters(True, False)\n def test_fwd_bwd(self, graph_mode):\n with set_graph_mode(graph_mode):\n\n class Foo(nnx.Module):\n @nnx.jit\n def __call__(self, x):\n x = self.perturb('grad_of_x', x)\n y = 3 * x\n self.sow(nnx.Intermediate, 'y', y)\n return y\n\n model = Foo()\n\n @nnx.jit\n def train_step(model, perturbations, x):\n def loss(model, perturbations, x):\n loss, interms = nnx.capture(model, nnx.Intermediate, init=perturbations)(x)\n return loss, interms\n (grads, perturb_grads), sowed = nnx.grad(loss, argnums=(0, 1), has_aux=True)(model, perturbations, x)\n return grads, nnx.merge_state(perturb_grads, sowed)\n\n x = 1.0\n\n forward = nnx.capture(model, nnx.Perturbation)\n _, perturbations = forward(x)\n grads, intermediates = train_step(model, perturbations, x)\n self.assertEqual(intermediates['grad_of_x'], 3)\n self.assertEqual(intermediates['y'][0], 3)\n\n @parameterized.parameters(True, False)\n def test_nested_modules(self, graph_mode):\n with set_graph_mode(graph_mode):\n\n class Foo(nnx.Module):\n def __call__(self, x):\n x = self.perturb('grad_of_x', x)\n y = 3 * x\n self.sow(nnx.Intermediate, 'y', y)\n return y\n class Bar(nnx.Module):\n def __init__(self):\n self.foos = nnx.data([Foo() for _ in range(3)])\n def __call__(self, x):\n for block in self.foos:\n x = block(x)\n return x\n\n model = Bar()\n\n @nnx.jit\n def train_step(model, perturbations, x):\n def loss(model, perturbations, x):\n loss, interms = nnx.capture(model, nnx.Intermediate, init=perturbations)(x)\n return loss, interms\n (grads, perturb_grads), sowed = nnx.grad(loss, argnums=(0, 1), has_aux=True)(model, perturbations, x)\n return grads, nnx.merge_state(perturb_grads, sowed)\n\n x = 1.0\n\n forward = nnx.capture(model, nnx.Perturbation)\n _, perturbations = forward(x)\n _, intermediates = train_step(model, perturbations, x)\n for i in range(3):\n self.assertEqual(intermediates['foos'][i]['grad_of_x'], 3**(3-i))\n self.assertEqual(intermediates['foos'][i]['y'][0], 3**(i+1))\n\n def test_method_outputs_single_module(self):\n class Foo(nnx.Module):\n def __init__(self, dim):\n self.w = nnx.Param(jax.random.normal(jax.random.key(0), (dim, dim)))\n def __call__(self, x):\n return x @ self.w\n def helper(self, x):\n return jnp.sin(x)\n def run(self, x):\n y = self(x)\n z = self.helper(y)\n return (y, z)\n\n model = Foo(8)\n x = jnp.ones((4, 8))\n\n run_with_capture = nnx.capture(\n model.run,\n nnx.Intermediate, method_outputs=nnx.Intermediate\n )\n (y, z), intms = run_with_capture(x)\n\n self.assertIn('__call__', intms)\n self.assertIn('helper', intms)\n np.testing.assert_allclose(intms['__call__'][0], y)\n np.testing.assert_allclose(intms['helper'][0], z)\n\n def test_method_outputs_nested_modules(self):\n class Inner(nnx.Module):\n def __init__(self, dim, rngs):\n self.w = nnx.Param(jax.random.normal(rngs.params(), (dim, dim)))\n def __call__(self, x):\n return x @ self.w\n def process(self, x):\n return jnp.sin(x)\n\n class Outer(nnx.Module):\n def __init__(self, rngs):\n self.inner1 = Inner(8, rngs)\n self.inner2 = Inner(8, rngs)\n def __call__(self, x):\n x = self.inner1(x)\n x = self.inner2.process(x)\n return x\n\n model = Outer(nnx.Rngs(0))\n x = jnp.ones((4, 8))\n\n forward = nnx.capture(\n model,\n nnx.Intermediate, method_outputs=nnx.Intermediate\n )\n y, intms = forward(x)\n\n self.assertIn('__call__', intms)\n self.assertIn('inner1', intms)\n self.assertIn('process', intms['inner2'])\n self.assertEqual(intms['inner1']['__call__'][0].shape, (4, 8))\n self.assertEqual(intms['inner2']['process'][0].shape, (4, 8))\n\n def test_method_outputs_mixed_with_sow(self):\n class Foo(nnx.Module):\n def __init__(self, dim):\n self.w = nnx.Param(jax.random.normal(jax.random.key(0), (dim, dim)))\n def __call__(self, x):\n x = x @ self.w\n self.sow(nnx.Intermediate, 'intermediate', x)\n return jnp.sin(x)\n\n model = Foo(8)\n x = jnp.ones((4, 8))\n\n forward = nnx.capture(\n model,\n nnx.Intermediate, method_outputs=nnx.Intermediate\n )\n y, intms = forward(x)\n\n self.assertIn('__call__', intms)\n self.assertIn('intermediate', intms)\n np.testing.assert_allclose(intms['__call__'][0], y)\n np.testing.assert_allclose(jnp.sin(intms['intermediate'][0]), y)\n\nclass SowMod(nnx.Module):\n def __init__(self, rngs: nnx.Rngs):\n self.linear = nnx.Linear(4, 4, rngs=rngs)\n\n def __call__(self, x):\n y = self.linear(x)\n self.sow(nnx.Intermediate, \"my_summary\", y.mean())\n return y * 2\n\nclass TestModule(parameterized.TestCase):\n def test_has_module_state(self):\n class Foo(nnx.Module): ...\n\n foo = Foo()\n\n assert hasattr(foo, '_pytree__state')\n\n def test_trace_level(self):\n m = nnx.Dict(a=nnx.Param(1))\n\n @jax.jit\n def f():\n with self.assertRaisesRegex(\n errors.TraceContextError,\n \"Cannot mutate 'Dict' from different trace level\",\n ):\n m.a = 2\n\n f()\n\n def test_tree_map(self):\n m = nnx.Dict(a=nnx.Param(1))\n\n graphdef, state = nnx.split(m)\n\n state = jax.tree.map(lambda x: x + 1, state)\n\n def test_split_2(self):\n m = nnx.Dict(a=nnx.Param(1))\n\n graphdef, empty, some = nnx.split(m, None, ...)\n\n some = jax.tree.map(lambda x: x + 1, some)\n\n def test_split_merge(self):\n m = nnx.Dict(a=nnx.Param(1))\n\n @jax.jit\n def g(graphdef: nnx.GraphDef[nnx.Dict], state: nnx.State):\n m = nnx.merge(graphdef, state)\n m.a = 2\n return nnx.split(m)\n\n graphdef, state = g(*nnx.split(m))\n m2 = nnx.merge(graphdef, state)\n\n assert m2.a == 2\n\n def test_call(self):\n class Foo(nnx.Module):\n def __init__(self, c: float, *, rngs: nnx.Rngs):\n key = rngs.params()\n self.w = nnx.Param(jax.random.uniform(key, ()))\n self.c = c\n\n def __call__(self, x, *, rngs: nnx.Rngs):\n return self.w * x + rngs.e.normal(()) + self.c\n\n foo = Foo(c=1.0, rngs=nnx.Rngs(0))\n\n y = foo(x=2.0, rngs=nnx.Rngs(e=1))\n\n assert isinstance(y, jax.Array)\n\n def test_shared_module(self):\n m1 = nnx.Dict(a=nnx.Param(1), b=nnx.Param(2))\n m2 = nnx.Dict(x=m1, y=m1, z=nnx.Param(3))\n\n m3 = nnx.merge(*nnx.split(m2))\n\n assert m3['x'] is m3['y']\n assert m3['x']['a'] is m3['y']['a']\n assert m3['x']['b'] is m3['y']['b']\n\n def test_module_graph(self):\n class Foo(nnx.Module):\n def __init__(self):\n self.a = nnx.Param(1)\n self.sub = self\n\n m = Foo()\n\n graphdef, state = nnx.split(m)\n assert len(state) == 1\n\n m2 = nnx.merge(graphdef, state)\n assert m2 is m2.sub\n\n def test_deref_through_jit(self):\n r1 = nnx.Variable(1)\n r2 = nnx.Variable(2)\n\n m = m0 = nnx.Dict({'a': nnx.List([r1, r2]), 'b': r1})\n\n @jax.jit\n def f(graphdef: nnx.GraphDef[nnx.Dict], state: nnx.State):\n m = nnx.merge(graphdef, state)\n\n assert m['a'][0] is m['b']\n assert m['a'][1] is not m['b']\n\n return nnx.split(m)\n\n graphdef, state = f(*nnx.split(m))\n m = nnx.merge(graphdef, state)\n\n assert m['a'][0] is m['b']\n assert m['a'][1] is not m['b']\n\n # compare with original\n assert m['a'][0] is not m0['a'][0]\n assert m['a'][1] is not m0['a'][1]\n assert m['b'] is not m0['b']\n\n def test_cross_barrier(self):\n m = nnx.Dict(a=nnx.Param(jnp.array(1)))\n\n @jax.jit\n def g(graphdef: nnx.GraphDef[nnx.Dict], state: nnx.State):\n m = nnx.merge(graphdef, state)\n m.a[...] += 1\n return nnx.split(m)\n\n graphdef, state = g(*nnx.split(m))\n m2 = nnx.merge(graphdef, state)\n assert m2 is not m\n assert m.a[...] == 1\n assert m2.a[...] == 2\n\n def test_no_rejit(self):\n n = 0\n m = nnx.Dict(a=nnx.Param(jnp.array(1)))\n\n @jax.jit\n def g(state_and_def):\n nonlocal n\n n += 1\n m = nnx.merge(*state_and_def)\n m.a[...] += 1\n return nnx.split(m)\n\n m2 = nnx.merge(*g(nnx.split(m)))\n\n assert n == 1\n assert m2 is not m\n assert m.a[...] == 1\n assert m2.a[...] == 2\n\n g(nnx.split(m))\n assert n == 1\n\n g(nnx.split(m2))\n assert n == 1\n\n m2.b = nnx.Param(10)\n g(nnx.split(m2))\n\n assert n == 2\n\n def test_deref_number_of_fields(self):\n r1 = nnx.Variable(1)\n r2 = nnx.Variable(2)\n v1 = 3\n m = nnx.Dict(\n {\n 'a': nnx.List([r1, r2, v1]),\n 'b': nnx.Dict({'c': r1, 'd': r2}),\n }\n )\n\n graphdef, p = nnx.split(m)\n assert len(nnx.to_flat_state(p)) == 2\n assert len(jax.tree_util.tree_leaves(p)) == 2\n\n def test_clone(self):\n m = nnx.Dict(\n a=nnx.List([nnx.Param(1), nnx.Param(2), 3]),\n b=nnx.Dict(c=nnx.Param(1), d=nnx.Param(2)),\n )\n\n m2 = nnx.clone(m)\n\n assert m is not m2\n assert m2.a[0].get_value() == m2.b.c.get_value()\n assert m2.a[1].get_value() == m2.b.d.get_value()\n\n assert m.a[0].get_value() == m2.a[0].get_value()\n assert m.a[1].get_value() == m2.a[1].get_value()\n assert m.b.c.get_value() == m2.b.c.get_value()\n assert m.b.d.get_value() == m2.b.d.get_value()\n\n def test_sow_existing_non_variable_field(self):\n class Foo(nnx.Module):\n def __init__(self) -> None:\n self.y = 10\n\n def __call__(self, x):\n y = x + 1\n self.sow(nnx.Intermediate, 'y', y)\n return y\n\n m = Foo()\n\n with self.assertRaisesRegex(ValueError, 'to be a Variable, got'):\n m(2)\n\n def test_sow_wrong_collection(self):\n class Foo(nnx.Module):\n def __init__(self) -> None:\n self.y = nnx.Param(10)\n\n def __call__(self, x):\n y = x + 1\n self.sow(nnx.Intermediate, 'y', y)\n return y\n\n m = Foo()\n\n with self.assertRaisesRegex(ValueError, 'to be of type'):\n m(2)\n\n def test_sow_pop(self):\n x = jnp.ones((2, 4))\n model = SowMod(nnx.Rngs(42))\n out, intermediates = nnx.capture(model, nnx.Intermediate)(x)\n attr_names = set(model._pytree__nodes)\n assert 'my_summary' not in attr_names\n\n def test_cached_partial(self):\n model = SowMod(nnx.Rngs(42))\n x = jnp.ones((2, 4))\n\n @nnx.jit\n def train_step(model, x):\n out, intermediates = nnx.capture(model, nnx.Intermediate)(x)\n return out, intermediates\n\n train_step_fn = nnx.cached_partial(train_step, model)\n train_step_fn(x)\n\n def test_update_static_state_submodules(self):\n class Bar(nnx.Module):\n def __init__(self) -> None:\n self.x = 1\n\n def add_field(self):\n self.y = 2\n\n class Foo(nnx.Module):\n def __init__(self) -> None:\n self.a = Bar()\n self.b = self.a\n\n m1 = Foo()\n with nnx.update_context('test'):\n with nnx.split_context('test') as ctx:\n graphdef, state = ctx.split(m1)\n with nnx.merge_context('test', inner=True) as ctx:\n m2 = ctx.merge(graphdef, state)\n m2.a.add_field()\n with nnx.split_context('test') as ctx:\n new_graphdef, state = ctx.split(m2)\n\n with nnx.merge_context('test', inner=False) as ctx:\n m3 = ctx.merge(new_graphdef, state)\n\n assert m3 is m1\n assert m1.a.x == 1\n assert m1.a.y == 2\n assert m1.b.x == 1\n assert m1.b.y == 2\n\n def test_update_new_submodule(self):\n class Bar(nnx.Module):\n def __init__(self) -> None:\n self.x = 1\n\n class Foo(nnx.Module):\n def __init__(self) -> None:\n self.a = Bar()\n\n def add_module(self):\n self.b = Bar()\n\n m1 = Foo()\n with nnx.update_context('test'):\n with nnx.split_context('test') as ctx:\n graphdef, state = ctx.split(m1)\n with nnx.merge_context('test', inner=True) as ctx:\n m2 = ctx.merge(graphdef, state)\n m2.add_module()\n with nnx.split_context('test') as ctx:\n new_graphdef, state = ctx.split(m2)\n\n with nnx.merge_context('test', inner=False) as ctx:\n m3 = ctx.merge(new_graphdef, state)\n\n assert m3 is m1\n assert m1.a.x == 1\n assert m1.b.x == 1\n\n def test_update_update_submodule(self):\n class Bar(nnx.Module):\n def __init__(self) -> None:\n self.x = 1\n\n class Foo(nnx.Module):\n def __init__(self) -> None:\n self.a = Bar()\n self.b = self.a\n\n m1 = Foo()\n with nnx.update_context('test'):\n with nnx.split_context('test') as ctx:\n graphdef, state = ctx.split(m1)\n with nnx.merge_context('test', inner=True) as ctx:\n m2 = ctx.merge(graphdef, state)\n m2.a.x = 2\n with nnx.split_context('test') as ctx:\n new_graphdef, state = ctx.split(m2)\n with nnx.merge_context('test', inner=False) as ctx:\n m3 = ctx.merge(new_graphdef, state)\n\n assert m3 is m1\n assert m1.a.x == 2\n assert m1.b.x == 2\n\n def test_update_add_shared(self):\n class Bar(nnx.Module):\n def __init__(self) -> None:\n self.x = 1\n\n class Foo(nnx.Module):\n def __init__(self) -> None:\n self.a = Bar()\n self.b = self.a\n\n def add_submodule(self):\n self.c = self.a\n\n m1 = Foo()\n with nnx.update_context('test'):\n with nnx.split_context('test') as ctx:\n graphdef, state = ctx.split(m1)\n with nnx.merge_context('test', inner=True) as ctx:\n m2 = ctx.merge(graphdef, state)\n m2.add_submodule()\n with nnx.split_context('test') as ctx:\n new_graphdef, state = ctx.split(m2)\n with nnx.merge_context('test', inner=False) as ctx:\n m3 = ctx.merge(new_graphdef, state)\n\n assert m3 is m1\n assert hasattr(m1, 'c')\n\n def test_create_abstract(self):\n linear = nnx.eval_shape(lambda: nnx.Linear(2, 3, rngs=nnx.Rngs(0)))\n\n assert linear.kernel.get_value() == jax.ShapeDtypeStruct((2, 3), jnp.float32)\n assert linear.bias.get_value() == jax.ShapeDtypeStruct((3,), jnp.float32)\n\n def test_create_abstract_stateful(self):\n linear = nnx.eval_shape(lambda: nnx.Dropout(0.5, rngs=nnx.Rngs(0)))\n\n assert linear.rngs.key.get_value() == jax.ShapeDtypeStruct(\n (), jax.random.key(0).dtype\n )\n\n def test_partial_init(self):\n linear = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n state = nnx.state(linear)\n\n del state['bias']\n\n @nnx.jit\n def partial_init(state: nnx.State):\n m = nnx.Linear(\n 2, 3, bias_init=nnx.initializers.ones_init(), rngs=nnx.Rngs(1)\n )\n nnx.update(m, state)\n return m\n\n linear2 = partial_init(state)\n\n np.testing.assert_allclose(linear.kernel[...], linear2.kernel[...])\n np.testing.assert_allclose(linear.bias[...], 0)\n np.testing.assert_allclose(linear2.bias[...], 1)\n\n def test_deepcopy(self):\n class Foo(nnx.Module):\n def __init__(self) -> None:\n self.a = nnx.Param(jnp.array(1))\n self.b = [1, 2, 3]\n self.c = nnx.Param(jnp.array([1.0]))\n self.self = self\n\n m1 = Foo()\n m2 = deepcopy(m1)\n\n assert m1.a[...] == m2.a[...]\n assert vars(m1)['a'] is not vars(m2)['a']\n assert m1.b is not m2.b\n assert m1.c is not m2.c\n assert m1.self is m1\n\n def test_set_attributes(self):\n class Block(nnx.Module):\n def __init__(self, din, dout, *, rngs: nnx.Rngs):\n self.linear = nnx.Linear(din, dout, rngs=rngs)\n self.dropout = nnx.Dropout(0.5, deterministic=False)\n self.batch_norm = nnx.BatchNorm(\n 10, use_running_average=False, rngs=rngs\n )\n\n block = Block(2, 5, rngs=nnx.Rngs(0))\n assert block.dropout.deterministic == False\n assert block.batch_norm.use_running_average == False\n\n block.set_attributes(deterministic=True, use_running_average=True)\n assert block.dropout.deterministic == True\n assert block.batch_norm.use_running_average == True\n\n block = Block(2, 5, rngs=nnx.Rngs(0))\n block.set_attributes(nnx.Dropout, deterministic=True)\n # Only the dropout will be modified\n assert block.dropout.deterministic == True\n assert block.batch_norm.use_running_average == False\n\n def test_set_attribute_error(self):\n class Block(nnx.Module):\n def __init__(self, din, dout, *, rngs: nnx.Rngs):\n self.linear = nnx.Linear(din, dout, rngs=rngs)\n self.dropout = nnx.Dropout(0.5, deterministic=False)\n self.batch_norm = nnx.BatchNorm(\n 10, use_running_average=False, rngs=rngs\n )\n\n block = Block(2, 5, rngs=nnx.Rngs(0))\n\n with self.assertRaisesRegex(\n ValueError,\n (\n 'Could not find at least one instance of the following attributes:'\n \" \\\\['unknown'\\\\]\"\n ),\n ):\n block.set_attributes(\n deterministic=True, use_running_average=True, unknown=True\n )\n\n block.set_attributes(\n deterministic=True,\n use_running_average=True,\n unknown=True,\n raise_if_not_found=False,\n )\n\n @parameterized.parameters(True, False)\n def test_view(self, graph):\n class Block(nnx.Module):\n def __init__(self, din, dout, *, rngs: nnx.Rngs):\n self.linear = nnx.Linear(din, dout, rngs=rngs)\n self.dropout = nnx.Dropout(0.5, deterministic=False, rngs=rngs)\n self.batch_norm = nnx.BatchNorm(\n 10, use_running_average=False, rngs=rngs\n )\n\n block = Block(2, 5, rngs=nnx.Rngs(0))\n assert block.dropout.deterministic == False\n assert block.batch_norm.use_running_average == False\n\n new_block = nnx.view(block, deterministic=True, use_running_average=True, graph=graph)\n assert new_block.dropout.deterministic == True\n assert new_block.batch_norm.use_running_average == True\n assert new_block.linear.kernel is block.linear.kernel\n\n block = Block(2, 5, rngs=nnx.Rngs(0))\n new_block = nnx.view(block, only=nnx.Dropout, deterministic=True, graph=graph)\n assert new_block.dropout.deterministic == True\n assert new_block.batch_norm.use_running_average == False\n\n @parameterized.parameters(True, False)\n def test_with_attributes(self, graph):\n class Block(nnx.Module):\n def __init__(self, din, dout, *, rngs: nnx.Rngs):\n self.linear = nnx.Linear(din, dout, rngs=rngs)\n self.dropout = nnx.Dropout(0.5, deterministic=False)\n self.batch_norm = nnx.BatchNorm(\n 10, use_running_average=False, rngs=rngs\n )\n\n block = Block(2, 5, rngs=nnx.Rngs(0))\n assert block.dropout.deterministic == False\n assert block.batch_norm.use_running_average == False\n\n new_block = nnx.with_attributes(\n block, deterministic=True, use_running_average=True, graph=graph\n )\n assert new_block.dropout.deterministic == True\n assert new_block.batch_norm.use_running_average == True\n assert new_block.linear.kernel is block.linear.kernel\n assert block.dropout.deterministic == False\n assert block.batch_norm.use_running_average == False\n\n @parameterized.parameters(True, False)\n def test_with_attributes_filter(self, graph):\n class Block(nnx.Module):\n def __init__(self, din, dout, *, rngs: nnx.Rngs):\n self.linear = nnx.Linear(din, dout, rngs=rngs)\n self.dropout = nnx.Dropout(0.5, deterministic=False)\n self.batch_norm = nnx.BatchNorm(\n 10, use_running_average=False, rngs=rngs\n )\n\n block = Block(2, 5, rngs=nnx.Rngs(0))\n new_block = nnx.with_attributes(\n block, only=nnx.Dropout, deterministic=True, graph=graph\n )\n assert new_block.dropout.deterministic == True\n assert new_block.batch_norm.use_running_average == False\n\n @parameterized.parameters(True, False)\n def test_with_attributes_error(self, graph):\n class Block(nnx.Module):\n def __init__(self, din, dout, *, rngs: nnx.Rngs):\n self.linear = nnx.Linear(din, dout, rngs=rngs)\n self.dropout = nnx.Dropout(0.5, deterministic=False)\n self.batch_norm = nnx.BatchNorm(\n 10, use_running_average=False, rngs=rngs\n )\n\n block = Block(2, 5, rngs=nnx.Rngs(0))\n with self.assertRaisesRegex(\n ValueError,\n (\n 'Could not find at least one instance of the following attributes:'\n \" \\\\['unknown'\\\\]\"\n ),\n ):\n nnx.with_attributes(\n block, deterministic=True, use_running_average=True,\n unknown=True, graph=graph,\n )\n\n new_block = nnx.with_attributes(\n block, deterministic=True, use_running_average=True,\n unknown=True, raise_if_not_found=False, graph=graph,\n )\n assert new_block.dropout.deterministic == True\n assert new_block.batch_norm.use_running_average == True\n\n @parameterized.parameters(True, False)\n def test_view_error(self, graph):\n class Block(nnx.Module):\n def __init__(self, din, dout, *, rngs: nnx.Rngs):\n self.linear = nnx.Linear(din, dout, rngs=rngs)\n self.dropout = nnx.Dropout(0.5, deterministic=False, rngs=rngs)\n self.batch_norm = nnx.BatchNorm(\n 10, use_running_average=False, rngs=rngs\n )\n block = Block(2, 5, rngs=nnx.Rngs(0))\n\n with self.assertRaisesRegex(\n ValueError,\n (\n \"Unused keys found in nnx.view: \\\\['unknown'\\\\]\"\n ),\n ):\n nnx.view(block, deterministic=True, use_running_average=True, unknown=True, graph=graph)\n\n def test_cloud_pickle(self):\n import platform\n if platform.python_version().startswith('3.11'):\n self.skipTest(\"Cloudpickle cannot pickle PRNGKeyArray on python 3.11\")\n class Model(nnx.Module):\n def __init__(self, din, dmid, dout, rngs: nnx.Rngs):\n self.linear = nnx.Linear(din, dmid, rngs=rngs)\n self.bn = nnx.BatchNorm(dmid, rngs=rngs)\n self.dropout = nnx.Dropout(0.1, rngs=rngs)\n self.linear_out = nnx.Linear(dmid, dout, rngs=rngs)\n\n def __call__(self, x):\n x = nnx.relu(self.dropout(self.bn(self.linear(x))))\n return self.linear_out(x)\n\n model = Model(2, 64, 3, rngs=nnx.Rngs(0)) # eager initialization\n model.eval()\n\n y1 = model(jnp.ones((5, 2)))\n with tempfile.TemporaryDirectory() as tmpdir:\n path = f'{tmpdir}/model.pkl'\n with open(path, 'wb') as f:\n cloudpickle.dump(model, f)\n del model\n with open(path, 'rb') as f:\n model = pickle.load(f)\n\n self.assertIsInstance(model, Model)\n y2 = model(jnp.ones((5, 2)))\n np.testing.assert_allclose(y1, y2)\n\n def test_repr(self):\n class Block(nnx.Module):\n def __init__(self, din, dout, rngs: nnx.Rngs):\n self.linear = nnx.Linear(din, dout, rngs=rngs)\n self.bn = nnx.BatchNorm(dout, rngs=rngs)\n self.dropout = nnx.Dropout(0.2, rngs=rngs)\n\n def __call__(self, x):\n return nnx.relu(self.dropout(self.bn(self.linear(x))))\n\n class Foo(nnx.Module):\n def __init__(self, rngs: nnx.Rngs):\n self.block1 = Block(32, 128, rngs=rngs)\n self.block2 = Block(128, 10, rngs=rngs)\n\n def __call__(self, x):\n return self.block2(self.block1(x))\n\n obj = Foo(nnx.Rngs(0))\n\n leaves = nnx.to_flat_state(nnx.state(obj)).leaves\n\n expected_total = sum(int(np.prod(x.shape)) for x in leaves)\n expected_total_params = sum(\n int(np.prod(x.shape)) for x in leaves if isinstance(x, nnx.Param)\n )\n expected_total_batch_stats = sum(\n int(np.prod(x.shape)) for x in leaves if isinstance(x, nnx.BatchStat)\n )\n expected_total_rng_states = sum(\n int(np.prod(x.shape)) for x in leaves if isinstance(x, nnx.RngState)\n )\n\n foo_repr = repr(obj).replace(',', '').splitlines()\n\n self.assertIn(str(expected_total), foo_repr[0])\n self.assertIn(str(expected_total_params), foo_repr[0])\n self.assertIn(str(expected_total_batch_stats), foo_repr[0])\n self.assertIn(str(expected_total_rng_states), foo_repr[0])\n\n @parameterized.parameters(True, False)\n def test_view_info(self, graph):\n class Block(nnx.Module):\n def __init__(self, din, dout, *, rngs: nnx.Rngs):\n self.linear = nnx.Linear(din, dout, rngs=rngs)\n self.bn = nnx.BatchNorm(dout, rngs=rngs)\n self.dropout = nnx.Dropout(0.2, rngs=rngs)\n\n def __call__(self, x):\n return nnx.relu(self.dropout(self.bn(self.linear(x))))\n\n class Foo(nnx.Module):\n def __init__(self, rngs: nnx.Rngs):\n self.block1 = Block(32, 128, rngs=rngs)\n self.block2 = Block(128, 10, rngs=rngs)\n\n def __call__(self, x):\n return self.block2(self.block1(x))\n\n obj = Foo(rngs=nnx.Rngs(0))\n info_str = nnx.view_info(obj, graph=graph)\n self.assertEqual(info_str.count(\"BatchNorm:\"), 1)\n self.assertEqual(info_str.count(\"Dropout:\"), 1)\n\n @parameterized.parameters(True, False)\n def test_view_info_with_filter(self, graph):\n class Block(nnx.Module):\n def __init__(self, din, dout, *, rngs: nnx.Rngs):\n self.linear = nnx.Linear(din, dout, rngs=rngs)\n self.bn = nnx.BatchNorm(dout, rngs=rngs)\n self.dropout = nnx.Dropout(0.2, rngs=rngs)\n\n def __call__(self, x):\n return nnx.relu(self.dropout(self.bn(self.linear(x))))\n\n obj = Block(4, 8, rngs=nnx.Rngs(0))\n info_str = nnx.view_info(obj, only=nnx.Dropout, graph=graph)\n self.assertIn(\"Dropout:\", info_str)\n self.assertNotIn(\"BatchNorm:\", info_str)\n\n info_str = nnx.view_info(obj, only=nnx.MultiHeadAttention, graph=graph)\n self.assertEmpty(info_str)\n\n @parameterized.parameters(True, False)\n def test_view_info_with_custom_set_mode(self, graph):\n class Block(nnx.Module):\n def __init__(self, *, rngs: nnx.Rngs):\n pass\n\n def __call__(self, x):\n return x\n\n def set_view(self, arg1: bool | None = None, arg2: int | None = None, **kwargs) -> dict:\n \"\"\"Example set_view docstring. This follows Google style docstrings.\n\n Args:\n arg1: The first argument.\n arg2: The second argument.\n This has two lines.\n \"\"\"\n return kwargs\n\n obj = Block(rngs=nnx.Rngs(0))\n info_str = nnx.view_info(obj, graph=graph)\n self.assertEqual(f\"{obj.__class__.__qualname__}:\\n arg1: bool | None = None\\n The first argument.\\n arg2: int | None = None\\n The second argument.\\n This has two lines.\", info_str)\n\nclass TestModuleDataclass(absltest.TestCase):\n def test_basic(self):\n\n @dataclasses.dataclass\n class Foo(nnx.Module):\n a: int\n b: nnx.Variable[int]\n c: nnx.Param[int]\n d: nnx.Variable[int]\n e: nnx.Variable[int]\n f: int\n\n m = Foo(\n a=1, # graphdef\n b=nnx.Variable(2), # node\n c=nnx.Param(3), # param\n d=nnx.Variable(4), # var\n e=nnx.BatchStat(5), # var\n f=6, # graphdef int\n )\n\n graphdef, state = nnx.split(m)\n\n assert len(state) == 4\n assert state['b'].get_value() == 2\n assert isinstance(state['b'], nnx.Variable)\n assert state['c'].get_value() == 3\n assert isinstance(state['c'], nnx.Param)\n assert state['d'].get_value() == 4\n assert isinstance(state['d'], nnx.Variable)\n assert state['e'].get_value() == 5\n assert isinstance(state['e'], nnx.BatchStat)\n\n def test_field_specifiers(self):\n @nnx.dataclass\n class Foo(nnx.Pytree):\n a: int = nnx.static()\n b: jax.Array = nnx.data()\n\n m = Foo(a=1, b=jnp.array(2))\n\n leaves = jax.tree.leaves(m)\n assert len(leaves) == 1\n assert leaves[0] == jnp.array(2)\n\n def test_field_specifiers_forced(self):\n @nnx.dataclass\n class Bar(nnx.Pytree):\n a: int = nnx.data()\n\n m = Bar(a=1)\n\n leaves = jax.tree.leaves(m)\n assert len(leaves) == 1\n assert leaves[0] == 1\n\n def test_field_specifiers_with_defaults(self):\n @nnx.dataclass\n class Bar(nnx.Pytree):\n a: int = nnx.data(default=3)\n\n m = Bar()\n\n leaves = jax.tree.leaves(m)\n assert len(leaves) == 1\n assert leaves[0] == 3\n\n def test_field_specifiers_array_in_static(self):\n @nnx.dataclass\n class Bar(nnx.Pytree):\n a: jax.Array = nnx.static()\n\n with self.assertRaisesRegex(\n ValueError,\n 'Found unexpected data on value of type',\n ):\n m = Bar(a=jnp.array(3))\n\n def test_variable_in_static_list(self):\n @nnx.dataclass\n class Foo(nnx.Module):\n filters: list\n\n with self.assertRaisesRegex(\n ValueError,\n 'Found data on value of type',\n ):\n Foo([nnx.Variable(1)])\n\n def test_module_in_static_list(self):\n class Bar(nnx.Module):\n pass\n\n @nnx.dataclass\n class Foo(nnx.Module):\n filters: list\n\n with self.assertRaisesRegex(\n ValueError,\n 'Found data on value of type',\n ):\n Foo([Bar()])\n\n def test_post_init(self):\n\n @dataclasses.dataclass\n class DFoo(nnx.Module):\n din: int\n dout: int\n rngs: nnx.Rngs\n\n def __post_init__(self):\n self.bar = nnx.Linear(self.din, self.dout, rngs=self.rngs)\n\n def __call__(self, x):\n return self.bar(x)\n\n m = DFoo(1, 1, rngs=nnx.Rngs(0))\n\n assert hasattr(m, 'bar')\n\n\nclass TestModuleDef(parameterized.TestCase):\n def test_apply(self):\n class Foo(nnx.Module):\n def __init__(self, c: float, *, rngs: nnx.Rngs):\n self.w = nnx.Param(jax.random.uniform(rngs.params(), ()))\n self.c = c\n\n def __call__(self, x, *, rngs: nnx.Rngs):\n return self.w * x + rngs.e.normal(()) + self.c\n\n rngs = nnx.Rngs(0)\n foo = Foo(c=1.0, rngs=rngs)\n\n graphdef, states = nnx.split(foo)\n\n assert isinstance(states, nnx.State)\n assert isinstance(states['w'], nnx.Param)\n\n y, _updates = graphdef.apply(states)(x=2.0, rngs=nnx.Rngs(e=1))\n\n assert isinstance(y, jax.Array)\n\n def test_derefed_mod_apply(self):\n class Foo(nnx.Module):\n def __init__(self, c: float, *, rngs: nnx.Rngs):\n self.w = nnx.Param(\n jax.random.uniform(rngs.params(), ()),\n )\n self.c = nnx.Variable(c)\n\n def __call__(self, x, *, rngs: nnx.Rngs):\n return self.w * x + rngs.e.normal(()) + self.c\n\n foo = Foo(c=1.0, rngs=nnx.Rngs(0))\n\n graphdef, state = nnx.split(foo)\n\n assert isinstance(graphdef.nodes[0], nnx.graphlib.NodeDef | nnx.graphlib.NodeRef)\n assert isinstance(state, nnx.State)\n assert isinstance(state['w'], nnx.Param)\n assert isinstance(state['c'], nnx.Variable)\n\n y, (graphdef, state) = graphdef.apply(state)(x=2.0, rngs=nnx.Rngs(e=1))\n\n assert isinstance(y, jax.Array)\n\n def test_modules_iterator(self):\n class Foo(nnx.Module):\n def __init__(self, *, rngs: nnx.Rngs):\n self.submodules = nnx.data([\n {'a': nnx.Linear(1, 1, rngs=rngs)},\n {'b': nnx.Conv(1, 1, 1, rngs=rngs)},\n ])\n self.linear = nnx.Linear(1, 1, rngs=rngs)\n self.dropout = nnx.Dropout(0.5, rngs=rngs)\n\n module = Foo(rngs=nnx.Rngs(0))\n\n modules = list(nnx.iter_modules(module))\n\n assert len(modules) == 5\n assert modules[0][0] == ('dropout',)\n assert isinstance(modules[0][1], nnx.Dropout)\n assert modules[1][0] == ('linear',)\n assert isinstance(modules[1][1], nnx.Linear)\n assert modules[2][0] == ('submodules', 0, 'a')\n assert isinstance(modules[2][1], nnx.Linear)\n assert modules[3][0] == ('submodules', 1, 'b')\n assert isinstance(modules[3][1], nnx.Conv)\n assert modules[4][0] == ()\n assert isinstance(modules[4][1], Foo)\n\n @parameterized.parameters(True, False)\n def test_children_modules_iterator(self, graph):\n class Foo(nnx.Module):\n def __init__(self, *, rngs: nnx.Rngs):\n self.submodules = nnx.data([\n {'a': nnx.Linear(1, 1, rngs=rngs)},\n {'b': nnx.Conv(1, 1, 1, rngs=rngs)},\n ])\n self.linear = nnx.Linear(1, 1, rngs=rngs)\n self.dropout = nnx.Dropout(0.5, rngs=rngs)\n\n module = Foo(rngs=nnx.Rngs(0))\n\n modules = list(nnx.iter_children(module, graph=graph))\n\n assert len(modules) == 2\n assert modules[0][0] == 'dropout'\n assert isinstance(modules[0][1], nnx.Dropout)\n assert modules[1][0] == 'linear'\n assert isinstance(modules[1][1], nnx.Linear)\n\n def test_state_in_module(self):\n class Foo(nnx.Module):\n def __init__(self):\n self.a = nnx.data(nnx.State({'b': nnx.Param(jnp.array(1.0))}))\n\n foo = Foo()\n\n graphdef, state = nnx.split(foo)\n\n assert isinstance(state, nnx.State)\n assert isinstance(state['a'], nnx.State)\n\n foo2 = nnx.merge(graphdef, state)\n\n assert isinstance(foo2.a, nnx.State)\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/nnx/mutable_array_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport dataclasses\nfrom absl.testing import absltest, parameterized\nimport optax\nfrom flax import nnx\nimport flax.errors\nimport jax\nimport jax.numpy as jnp\nimport numpy as np\n\n\nclass TestPytree(absltest.TestCase):\n def test_pytree(self):\n class Foo(nnx.Module):\n def __init__(self):\n self.node = jnp.array(1)\n self.meta = 1\n\n m = nnx.vars_as(Foo(), ref=True)\n\n m = jax.tree.map(lambda x: x + 1, m)\n\n assert m.node == 2\n assert m.meta == 1\n\n def test_pytree_data_typehint(self):\n class Foo(nnx.Module):\n node: jax.Array = nnx.data()\n\n def __init__(self):\n self.node = jnp.array(1)\n self.meta = 1\n\n m = Foo()\n\n m = jax.tree.map(lambda x: x + 1, m)\n\n assert m.node == 2\n assert m.meta == 1\n\n def test_pytree_data_instance(self):\n class Foo(nnx.Module):\n def __init__(self):\n self.node = nnx.data(jnp.array(1))\n self.meta = 1\n\n m = Foo()\n\n m = jax.tree.map(lambda x: x + 1, m)\n\n assert m.node == 2\n assert m.meta == 1\n\n def test_pytree_dataclass(self):\n @nnx.dataclass\n class Foo(nnx.Module):\n node: jax.Array = nnx.data()\n meta: int\n meta2: int = 3\n meta3: int = 4\n meta4: int = 5\n node2: int = nnx.data(default=6)\n\n m = Foo(node=jnp.array(1), meta=1)\n\n m: Foo = jax.tree.map(lambda x: x + 1, m)\n\n assert m.node == 2\n\n assert m.meta == 1\n assert m.meta2 == 3\n assert m.meta3 == 4\n assert m.meta4 == 5\n assert m.node2 == 7\n\n def test_data_example(self):\n class Foo(nnx.Pytree):\n def __init__(self):\n self.data_attr = nnx.data(42) # pytree data\n self.static_attr = 'hello' # static attribute\n\n foo = Foo()\n\n self.assertEqual(jax.tree.leaves(foo), [42])\n\n def test_register_data_type(self):\n @dataclasses.dataclass(frozen=True)\n class MyType:\n value: int\n\n nnx.register_data_type(MyType)\n\n class Foo(nnx.Pytree):\n def __init__(self, a):\n self.a = MyType(a) # Automatically registered as data\n self.b = 'hello' # str not registered as data\n\n foo = Foo(42)\n\n self.assertTrue(nnx.is_data(foo.a))\n self.assertEqual(jax.tree.leaves(foo), [MyType(value=42)])\n\nclass TestVariableRefMode(absltest.TestCase):\n def test_split_mutable_array(self):\n m = jax.new_ref(1)\n graphdef, state = nnx.split(m)\n\n self.assertIs(m, state)\n\n m2 = nnx.merge(graphdef, state)\n\n self.assertIs(m2, m)\n\n def test_to_arrays_example(self):\n node = [nnx.Variable(1.0), nnx.Variable(2.0, mode='ref')]\n mutable_node = nnx.vars_as(node, ref=True)\n assert isinstance(mutable_node[0].get_raw_value(), jax.Ref)\n assert isinstance(mutable_node[1].get_raw_value(), jax.Ref)\n\n shared_array = nnx.Variable(1.0, mode='pytree')\n node = [shared_array, shared_array]\n with self.assertRaisesRegex(ValueError, 'Found duplicate at path'):\n nnx.vars_as(node, ref=True)\n\n node = [nnx.Variable(1.0), nnx.Variable(2.0)]\n mutable_node = nnx.vars_as(\n node, ref=True, only=lambda path, x: path[0] == 0\n )\n assert isinstance(mutable_node[0].get_raw_value(), jax.Ref)\n assert isinstance(mutable_node[1].get_raw_value(), float)\n\n def test_freeze_and_mutable_with_filter(self):\n class Foo(nnx.Module):\n def __init__(self):\n self.a = nnx.Param(1)\n self.b = nnx.BatchStat(2)\n\n m = nnx.vars_as(Foo(), hijax=True, ref=True)\n self.assertEqual(m.a.ref, True)\n self.assertEqual(m.b.ref, True)\n\n m2 = nnx.vars_as(m, hijax=False, only=nnx.BatchStat)\n self.assertEqual(m2.a.ref, True)\n self.assertEqual(m2.a.hijax, True)\n self.assertEqual(m2.b.ref, True)\n self.assertEqual(m2.b.hijax, False)\n self.assertIsNot(m, m2)\n\n m3 = nnx.vars_as(m2, hijax=True, only=nnx.BatchStat)\n self.assertEqual(m3.a.ref, True)\n self.assertEqual(m3.b.ref, True)\n self.assertEqual(m3.b.hijax, True)\n self.assertIsNot(m2, m3)\n self.assertIs(m.a, m3.a)\n\n def test_freeze_duplicate_error(self):\n class Foo(nnx.Module):\n def __init__(self):\n self.a = nnx.Param(1, mode='ref')\n self.b = self.a\n\n m = Foo()\n\n with self.assertRaisesRegex(ValueError, 'Found duplicate at path'):\n nnx.vars_as(m, ref=True)\n\n def test_mutable_array_split(self):\n class Foo(nnx.Module):\n def __init__(self):\n self.a = jax.new_ref(1)\n self.b = self.a\n\n m = Foo()\n\n ref_map = nnx.graphlib.RefMap()\n graphdef, state = nnx.graphlib.flatten(m, ref_index=ref_map, graph=True)\n self.assertLen(state, 1)\n self.assertLen(ref_map, 2) # 1 Foo + 1 ArrayRef\n\n m1 = nnx.merge(graphdef, state)\n self.assertIs(m1.a, m1.b)\n self.assertIsInstance(m1.a, jax.Ref)\n\n def test_mutable_array_split_merge_in_variable(self):\n class Foo(nnx.Module):\n def __init__(self):\n self.a = nnx.Param(1, ref=True)\n self.b = self.a\n\n m = Foo()\n\n ref_map = nnx.graphlib.RefMap()\n graphdef, state = nnx.graphlib.flatten(m, ref_index=ref_map, graph=True)\n self.assertLen(state, 1)\n self.assertLen(ref_map, 3) # 1 Foo + 1 Param + 1 Ref\n\n m1 = nnx.merge(graphdef, state)\n self.assertIs(m1.a, m1.b)\n self.assertIsInstance(m1.a, nnx.Param)\n\n def test_mutable_array_split_merge_in_variable_shared_array(self):\n class Foo(nnx.Module):\n def __init__(self):\n m_array = 1\n self.a = nnx.Param(m_array, ref=True)\n self.b = nnx.Param(m_array, ref=True)\n\n m = Foo()\n\n ref_map = nnx.graphlib.RefMap()\n graphdef, state = nnx.graphlib.flatten(m, ref_index=ref_map, graph=True)\n self.assertLen(state, 2)\n self.assertLen(ref_map, 5) # 1 Foo + 2 Param + 2 Ref\n\n m1 = nnx.merge(graphdef, state)\n # Each variable will own its own array and ref.\n self.assertIsInstance(m1.a, nnx.Param)\n\n def test_mutable_example(self):\n tree = [nnx.Variable(1.0), nnx.Variable(2.0, ref=True)]\n assert tree[0].ref == False\n assert tree[1].ref == True\n mutable_tree = nnx.vars_as(tree, ref=True)\n assert isinstance(mutable_tree[0].get_raw_value(), jax.Ref)\n assert isinstance(mutable_tree[1].get_raw_value(), jax.Ref)\n\n def test_mutable_array_split_freeze(self):\n class Foo(nnx.Module):\n def __init__(self):\n self.a = jax.new_ref(1)\n self.b = self.a\n\n m = Foo()\n\n ref_map = nnx.graphlib.RefMap()\n graphdef, state = nnx.graphlib.flatten(m, ref_index=ref_map, graph=True)\n state = nnx.vars_as(state, hijax=False)\n self.assertLen(state, 1)\n\n m1 = nnx.merge(graphdef, state)\n self.assertIs(m1.a, m1.b)\n self.assertIsInstance(m1.a, jax.Ref)\n\n def test_update_context(self):\n m1 = nnx.vars_as(nnx.Linear(1, 1, rngs=nnx.Rngs(0)), ref=True)\n with nnx.update_context('example'):\n with nnx.split_context('example') as ctx:\n graphdef, state = ctx.split(m1)\n\n with nnx.merge_context('example', True) as ctx:\n m2 = ctx.merge(graphdef, state)\n\n m_out1 = nnx.vars_as(nnx.Linear(1, 1, rngs=nnx.Rngs(0)), ref=True)\n\n with nnx.split_context('example') as ctx:\n graphdef_out, state_out = ctx.split((m2, m_out1, m2))\n\n self.assertIsInstance(\n state_out[0]['kernel'].get_value(), nnx.graphlib.NoUpdate\n )\n self.assertIsInstance(\n state_out[0]['bias'].get_value(), nnx.graphlib.NoUpdate\n )\n self.assertIsInstance(\n state_out[1]['kernel'].get_value(), nnx.graphlib.ArrayRefOutput\n )\n self.assertIsInstance(\n state_out[1]['bias'].get_value(), nnx.graphlib.ArrayRefOutput\n )\n # 2 ArrayRefOutput + 2 NoUpdate, however, NoUpdate are empty nodes\n self.assertLen(jax.tree.leaves(state_out), 2)\n\n with nnx.merge_context('example', False) as ctx:\n m3, m_out2, _ = ctx.merge(graphdef_out, state_out)\n\n self.assertIs(m3, m1)\n self.assertIsNot(m_out2, m_out1)\n\n def test_update_context_flatten(self):\n m1 = nnx.vars_as(nnx.Linear(1, 1, rngs=nnx.Rngs(0)), ref=True)\n with nnx.update_context('example'):\n with nnx.split_context('example') as ctx:\n graphdef, state = ctx.flatten(m1)\n\n with nnx.merge_context('example', True) as ctx:\n m2 = ctx.merge(graphdef, state)\n\n m_out1 = nnx.vars_as(nnx.Linear(1, 1, rngs=nnx.Rngs(0)), ref=True)\n\n with nnx.split_context('example') as ctx:\n graphdef_out, state_out = ctx.flatten((m2, m_out1, m2))\n\n state_out_dict = dict(state_out)\n\n self.assertIsInstance(\n state_out_dict[(0, 'kernel')].get_value(), nnx.graphlib.NoUpdate\n )\n self.assertIsInstance(\n state_out_dict[(0, 'bias')].get_value(), nnx.graphlib.NoUpdate\n )\n self.assertIsInstance(\n state_out_dict[(1, 'kernel')].get_value(), nnx.graphlib.ArrayRefOutput\n )\n self.assertIsInstance(\n state_out_dict[(1, 'bias')].get_value(), nnx.graphlib.ArrayRefOutput\n )\n # 2 ArrayRefOutput + 2 NoUpdate, however, NoUpdate are empty nodes\n self.assertLen(jax.tree.leaves(state_out), 2)\n\n with nnx.merge_context('example', False) as ctx:\n m3, m_out2, _ = ctx.merge(graphdef_out, state_out)\n\n self.assertIs(m3, m1)\n self.assertIsNot(m_out2, m_out1)\n\n def test_update_context_to_tree1(self):\n m1 = nnx.vars_as(nnx.Linear(1, 1, rngs=nnx.Rngs(0)), ref=True)\n with nnx.update_context('example'):\n m1_tree = nnx.to_tree((m1,), ctxtag='example')\n\n (m2,) = nnx.from_tree(m1_tree, ctxtag='example', is_inner=True)\n\n m_out1 = nnx.vars_as(nnx.Linear(1, 1, rngs=nnx.Rngs(0)), ref=True)\n\n # with nnx.split_context('example') as ctx:\n # graphdef_out, state_out = ctx.split((m2, m_out1))\n out_tree = nnx.to_tree(((m2,), m_out1, m2), ctxtag='example')\n\n self.assertIsInstance(\n out_tree[0][0].states[0]['kernel'].get_value(), nnx.graphlib.NoUpdate\n )\n self.assertIsInstance(\n out_tree[0][0].states[0]['bias'].get_value(), nnx.graphlib.NoUpdate\n )\n self.assertIsInstance(\n out_tree[1].states[0]['kernel'].get_value(), nnx.graphlib.ArrayRefOutput\n )\n self.assertIsInstance(\n out_tree[1].states[0]['bias'].get_value(), nnx.graphlib.ArrayRefOutput\n )\n self.assertEmpty(out_tree[2].states[0]) # Repeated m2 State\n\n # 2 ArrayRefOutput + 2 NoUpdate, however, NoUpdate are empty nodes\n self.assertLen(jax.tree.leaves(out_tree), 2)\n\n # with nnx.merge_context('example', False) as ctx:\n # m3, m_out2 = ctx.merge(graphdef_out, state_out)\n (m3,), m_out2, _ = nnx.from_tree(\n out_tree, ctxtag='example', is_inner=False\n )\n\n self.assertIs(m3, m1)\n self.assertIsNot(m_out2, m_out1)\n\n def test_update_context_to_tree2(self):\n m1 = nnx.vars_as(nnx.Linear(1, 1, rngs=nnx.Rngs(0)), ref=True)\n with nnx.update_context('example') as ctx:\n m1_tree = nnx.to_tree((m1,), ctxtag='example')\n\n (m2,) = nnx.from_tree(m1_tree, ctxtag='example', is_inner=True)\n\n m_out1 = nnx.vars_as(nnx.Linear(1, 1, rngs=nnx.Rngs(0)), ref=True)\n\n # with nnx.split_context('example') as ctx:\n # graphdef_out, state_out = ctx.split((m2, m_out1))\n out_tree = nnx.to_tree(((m2,), m_out1, m2), ctxtag='example')\n\n self.assertIsInstance(\n out_tree[0][0].states[0]['kernel'].get_value(), nnx.graphlib.NoUpdate\n )\n self.assertIsInstance(\n out_tree[0][0].states[0]['bias'].get_value(), nnx.graphlib.NoUpdate\n )\n self.assertIsInstance(\n out_tree[1].states[0]['kernel'].get_value(), nnx.graphlib.ArrayRefOutput\n )\n self.assertIsInstance(\n out_tree[1].states[0]['bias'].get_value(), nnx.graphlib.ArrayRefOutput\n )\n self.assertEmpty(out_tree[2].states[0]) # Repeated m2 State\n\n # 2 ArrayRefOutput + 2 NoUpdate, however, NoUpdate are empty nodes\n self.assertLen(jax.tree.leaves(out_tree), 2)\n\n # with nnx.merge_context('example', False) as ctx:\n # m3, m_out2 = ctx.merge(graphdef_out, state_out)\n (m3,), m_out2, _ = nnx.from_tree(\n out_tree, ctxtag='example', is_inner=False\n )\n\n self.assertIs(m3, m1)\n self.assertIsNot(m_out2, m_out1)\n\n def test_update_context_to_tree_trivial_prefix(self):\n m1 = nnx.vars_as(nnx.Linear(1, 1, rngs=nnx.Rngs(0)), ref=True)\n with nnx.update_context('example'):\n m1_tree = nnx.to_tree((m1,), ctxtag='example', prefix=0)\n\n (m2,) = nnx.from_tree(m1_tree, ctxtag='example', is_inner=True, prefix=0)\n\n m_out1 = nnx.vars_as(nnx.Linear(1, 1, rngs=nnx.Rngs(0)), ref=True)\n\n # with nnx.split_context('example') as ctx:\n # graphdef_out, state_out = ctx.split((m2, m_out1))\n out_tree = nnx.to_tree(((m2,), m_out1, m2), ctxtag='example', prefix=0)\n\n self.assertIsInstance(\n out_tree[0][0].states[0]['kernel'].get_value(), nnx.graphlib.NoUpdate\n )\n self.assertIsInstance(\n out_tree[0][0].states[0]['bias'].get_value(), nnx.graphlib.NoUpdate\n )\n self.assertIsInstance(\n out_tree[1].states[0]['kernel'].get_value(), nnx.graphlib.ArrayRefOutput\n )\n self.assertIsInstance(\n out_tree[1].states[0]['bias'].get_value(), nnx.graphlib.ArrayRefOutput\n )\n self.assertEmpty(out_tree[2].states[0]) # Repeated m2 State\n\n # 2 ArrayRefOutput + 2 NoUpdate, however, NoUpdate are empty nodes\n self.assertLen(jax.tree.leaves(out_tree), 2)\n\n # with nnx.merge_context('example', False) as ctx:\n # m3, m_out2 = ctx.merge(graphdef_out, state_out)\n (m3,), m_out2, _ = nnx.from_tree(\n out_tree, ctxtag='example', is_inner=False, prefix=0\n )\n\n self.assertIs(m3, m1)\n self.assertIsNot(m_out2, m_out1)\n\n def test_simple_jit(self):\n m1 = nnx.vars_as(nnx.Linear(1, 1, rngs=nnx.Rngs(0)), ref=True)\n m_out1 = None\n\n @nnx.jit\n def f(m2):\n nonlocal m_out1\n m_out1 = nnx.vars_as(nnx.Linear(1, 1, rngs=nnx.Rngs(0)), ref=True)\n return m_out1\n\n m_out2 = f(m1)\n\n self.assertIsNot(m_out1, m_out2)\n self.assertIsInstance(m_out2.kernel, nnx.Param)\n self.assertIsInstance(m_out2.kernel[...], jax.Array)\n\n def test_jit_mutable(self):\n @nnx.dataclass\n class Foo(nnx.Pytree):\n a: jax.Ref = nnx.data()\n\n m1 = Foo(a=jax.new_ref(1))\n\n @nnx.jit\n def f(m2: Foo):\n m2.a[...] += 1\n return m2\n\n m_out1 = f(m1)\n self.assertEqual(m_out1.a[...], 2)\n self.assertIs(m_out1, m1)\n self.assertIsInstance(m_out1.a, jax.Ref)\n\n def test_static(self):\n class C(nnx.Module):\n def __init__(self, meta):\n self.meta = meta\n\n n = 0\n\n @jax.jit\n def f(x):\n nonlocal n\n n += 1\n\n f(C(1))\n assert n == 1\n f(C(1))\n assert n == 1\n f(C(2))\n assert n == 2\n f(C(2))\n assert n == 2\n\n def test_variable_creation(self):\n v = nnx.Variable(jnp.array(1), ref=True)\n self.assertEqual(v[...], 1)\n self.assertTrue(v.ref)\n self.assertIsInstance(v.get_raw_value(), jax.Ref)\n\n def test_variable_metadata(self):\n v = nnx.Variable(jnp.array(1), a=2, b=3)\n self.assertEqual(v.a, 2)\n self.assertEqual(v.b, 3)\n\n def test_object(self):\n class Params(nnx.Pytree):\n def __init__(self, din: int, dout: int):\n self.w = nnx.Param(jnp.zeros((din, dout), jnp.float32))\n self.b = nnx.Param(jnp.zeros((dout,), jnp.float32))\n self.count = nnx.Variable(jnp.array(0))\n\n params = Params(3, 4)\n params = nnx.vars_as(params, ref=True)\n\n paths_leaves, treedef = jax.tree.flatten_with_path(params)\n paths, leaves = zip(*paths_leaves)\n\n self.assertLen(paths_leaves, 3)\n self.assertEqual(leaves[0].shape, (4,)) # b\n self.assertEqual(leaves[1].shape, ()) # count\n self.assertEqual(leaves[2].shape, (3, 4)) # w\n self.assertEqual(\n paths[0],\n (jax.tree_util.GetAttrKey('b'), jax.tree_util.GetAttrKey('value')),\n )\n self.assertEqual(\n paths[1],\n (jax.tree_util.GetAttrKey('count'), jax.tree_util.GetAttrKey('value')),\n )\n self.assertEqual(\n paths[2],\n (jax.tree_util.GetAttrKey('w'), jax.tree_util.GetAttrKey('value')),\n )\n\n params = jax.tree.unflatten(treedef, leaves)\n\n self.assertEqual(params.w.shape, (3, 4))\n self.assertEqual(params.b.shape, (4,))\n self.assertEqual(params.count.shape, ())\n self.assertIsInstance(params.w, nnx.Variable)\n self.assertIsInstance(params.w[...], jax.Array)\n self.assertIsInstance(params.b, nnx.Variable)\n self.assertIsInstance(params.b[...], jax.Array)\n self.assertIsInstance(params.count, nnx.Variable)\n self.assertIsInstance(params.count[...], jax.Array)\n\n @jax.jit\n def linear(params: Params, x: jax.Array):\n params.count[...] += 1\n return x @ params.w[...] + params.b[...][None]\n\n x = jnp.ones((1, 3))\n y = linear(params, x)\n self.assertEqual(y.shape, (1, 4))\n self.assertEqual(params.count[...], 1)\n y = linear(params, x)\n self.assertEqual(params.count[...], 2)\n\n def test_object_state(self):\n class Params(nnx.Pytree):\n def __init__(self, din: int, dout: int):\n self.w = jnp.zeros((din, dout), jnp.float32)\n self.b = jnp.zeros((dout,), jnp.float32)\n self.count = nnx.data(0)\n\n params = Params(3, 4)\n\n with self.assertRaises(flax.errors.TraceContextError):\n\n @jax.jit\n def f():\n params.count = 1\n\n f()\n\n @jax.jit\n def f(params: Params):\n params.count = 1\n return params\n\n params = f(params)\n self.assertEqual(params.count, 1)\n\n def test_rngs_create(self):\n rngs = nnx.Rngs(0)\n\n paths_leaves = jax.tree.leaves_with_path(rngs)\n paths, leaves = zip(*paths_leaves)\n self.assertLen(paths_leaves, 2)\n self.assertEqual(leaves[0].shape, ()) # key\n self.assertEqual(leaves[1].shape, ()) # count\n self.assertEqual(\n paths[0],\n (\n jax.tree_util.GetAttrKey('default'),\n jax.tree_util.GetAttrKey('count'),\n jax.tree_util.GetAttrKey('value'),\n ),\n )\n self.assertEqual(\n paths[1],\n (\n jax.tree_util.GetAttrKey('default'),\n jax.tree_util.GetAttrKey('key'),\n jax.tree_util.GetAttrKey('value'),\n ),\n )\n\n def test_rngs_call(self):\n rngs = nnx.Rngs(0)\n key = rngs()\n self.assertIsInstance(key, jax.Array)\n\n\nclass TestOptimizer(absltest.TestCase):\n def test_optimize_arrays(self):\n class Model(nnx.Module):\n def __init__(self, rngs):\n self.w = jax.random.uniform(rngs(), (2, 4))\n self.count = jnp.array(0)\n\n def __call__(self, x):\n self.count += 1\n return x @ self.w\n\n x = jax.random.normal(jax.random.key(0), (5, 2))\n y = jnp.ones((5, 4))\n\n wrt = lambda path, x: path[-1] == 'w'\n model = Model(nnx.Rngs(1))\n optimizer = nnx.Optimizer(model, tx=optax.adam(1e-3), wrt=wrt)\n\n @jax.jit\n def train_step(model, optimizer, x, y):\n graphdef, params, nondiff = nnx.split(model, wrt, ...)\n\n def loss_fn(params):\n model = nnx.merge(graphdef, params, nondiff)\n return jnp.mean((model(x) - y) ** 2), nnx.state(model, nnx.Not(wrt))\n\n (loss, updates), grads = jax.value_and_grad(loss_fn, has_aux=True)(params)\n nnx.update(model, updates)\n optimizer.update(model, grads)\n\n return loss, model, optimizer\n\n loss, model, optimizer = train_step(model, optimizer, x, y)\n\n self.assertNotEqual(loss, 0.0)\n self.assertEqual(model.count[...], 1)\n self.assertEqual(optimizer.step[...], 1)\n\n @nnx.var_defaults(hijax=True)\n def test_optimize_hijax(self):\n class Model(nnx.Module):\n def __init__(self, rngs):\n self.w = nnx.Variable(jax.random.uniform(rngs(), (2, 4)))\n self.count = nnx.Variable(jnp.array(0))\n\n def __call__(self, x):\n self.count[...] += 1\n return x @ self.w\n\n x = jax.random.normal(jax.random.key(0), (5, 2))\n y = jnp.ones((5, 4))\n\n wrt = lambda path, x: path[-1] == 'w'\n model = Model(nnx.Rngs(1))\n optimizer = nnx.Optimizer(model, tx=optax.adam(1e-3), wrt=wrt)\n\n @jax.jit\n def train_step(model, optimizer, x, y):\n graphdef, params, nondiff = nnx.split(model, wrt, ...)\n\n def loss_fn(params):\n model = nnx.merge(graphdef, params, nondiff)\n return jnp.mean((model(x) - y) ** 2)\n\n loss, grads = jax.value_and_grad(loss_fn)(nnx.vars_as(params, hijax=False))\n optimizer.update(params, grads)\n return loss\n\n loss = train_step(model, optimizer, x, y)\n\n self.assertNotEqual(loss, 0.0)\n\nclass TestHijaxVariables(parameterized.TestCase):\n def test_variable_to_hijax(self):\n v_low = nnx.Param(jnp.array(1), a='hi')\n v_hi = nnx.vars_as(v_low, hijax=True)\n\n self.assertTrue(v_hi.hijax)\n self.assertEqual(v_hi[...], 1)\n self.assertIsInstance(v_hi, nnx.Param)\n\n v_hi[...] = 2\n self.assertEqual(v_hi[...], 2)\n\n @jax.jit\n def set(v_hi, a):\n self.assertIsInstance(v_hi, nnx.Param)\n v_hi[...] = a\n self.assertEqual(v_hi.a, 'hi')\n self.assertTrue(v_hi.hijax)\n v_hi[...] += 5\n return v_hi + 2\n\n y = set(v_hi, 10)\n self.assertEqual(v_hi[...], 15)\n self.assertEqual(y, 17)\n\n v_low = nnx.vars_as(v_hi, hijax=False)\n self.assertIsInstance(v_low, nnx.Param)\n self.assertFalse(v_low.hijax)\n self.assertEqual(v_low[...], 15)\n\n def test_from_metadata(self):\n value = 1\n metadata = {\n 'a': 'hi',\n 'hijax': False,\n 'ref': False,\n }\n v_low = nnx.Param.from_metadata(value, metadata)\n self.assertIsInstance(v_low, nnx.Param)\n self.assertFalse(v_low.hijax)\n\n metadata['hijax'] = True\n v_hi = nnx.Param.from_metadata(value, metadata)\n self.assertIsInstance(v_hi, nnx.Param)\n self.assertTrue(v_hi.hijax)\n\n def test_variable_to_hijax_clean(self):\n v_low = nnx.Param(jnp.array([1]), tag='hello')\n print()\n print(v_low)\n assert not v_low.hijax\n v_hi = nnx.vars_as(v_low, hijax=True)\n v_hi[...] = jnp.array([2])\n assert v_hi.hijax\n print(v_hi)\n assert v_hi[...] == 2\n\n @jax.jit\n def set(v_hi, a):\n v_hi[...] = a\n print(v_hi)\n assert v_hi.tag == 'hello'\n\n set(v_hi, 10)\n\n assert v_hi[...] == 10\n\n v_low = nnx.vars_as(v_hi, hijax=False)\n\n assert not v_low.hijax\n assert v_low[...] == 10\n\n\n\n def test_pytree_value(self):\n v = nnx.Variable({'a': jnp.array(0), 'b': jnp.array(2)}, hijax=True)\n\n @jax.jit\n def inc_and_double(v):\n v['a'] += 1\n v['b'] *= 2\n\n inc_and_double(v)\n\n self.assertEqual(v['a'], 1)\n self.assertEqual(v['b'], 4)\n\n def test_hijax_dynamic_structure(self):\n x = jnp.ones((4, 5))\n metrics = nnx.Variable({}, hijax=True)\n\n @jax.jit\n def f(x, metrics: nnx.Variable):\n metrics['x_sum'] = jnp.sum(x)\n\n self.assertEmpty(metrics)\n f(x, metrics)\n self.assertIn('x_sum', metrics)\n self.assertEqual(metrics['x_sum'], 20)\n\n def test_hijax_and_pytree(self):\n class Foo(nnx.Pytree):\n def __init__(self, din, dout, rngs: nnx.Rngs):\n self.w = nnx.Param(rngs.uniform((din, dout)))\n self.b = nnx.Param(jnp.zeros((dout,)))\n self.count = nnx.Variable(0)\n\n foo = Foo(2, 4, nnx.Rngs(1))\n assert not foo.w.hijax\n assert not foo.b.hijax\n\n foo = nnx.vars_as(foo, hijax=True)\n\n assert foo.w.hijax\n assert foo.b.hijax\n\n @jax.jit\n def forward(foo, x):\n foo.count[...] += 1\n return x @ foo.w + foo.b[None]\n\n x = jnp.ones((1, 2))\n y = forward(foo, x)\n assert y.shape == (1, 4)\n assert foo.count[...] == 1\n\n def test_use_hijax(self):\n v_low = nnx.Param(1, a='hi')\n self.assertFalse(v_low.hijax)\n\n v_hi = nnx.Param(1, a='hi', hijax=True)\n self.assertTrue(v_hi.hijax)\n\n with nnx.var_defaults(hijax=True):\n v2 = nnx.Param(1, a='hi')\n self.assertIs(type(v2), nnx.variablelib.HijaxVariable)\n self.assertTrue(v2.hijax)\n\n @nnx.var_defaults(hijax=True)\n def test_hijax_rngs(self):\n rngs = nnx.Rngs(0)\n self.assertIs(type(rngs.default.key), nnx.variablelib.HijaxVariable)\n self.assertIs(type(rngs.default.count), nnx.variablelib.HijaxVariable)\n\n @jax.jit\n def f(rngs: nnx.Rngs):\n return rngs()\n\n k1 = f(rngs)\n k2 = f(rngs)\n\n assert k1 != k2\n\n @absltest.skip(reason='not yet supported')\n def test_return_hijax_from_transform(self):\n @jax.jit\n def create_var():\n return nnx.Param(1, hijax=True)\n\n v = create_var()\n self.assertTrue(v.hijax)\n\n @absltest.skip('not yet supported')\n @nnx.var_defaults(hijax=True)\n def test_lower(self):\n v = nnx.Param(jnp.ones((2, 3)))\n\n @jax.jit\n def f(v):\n v[...] += 1\n return v[...]\n\n e = f.lower(v)\n y = e.out_info[2]\n self.assertEqual(y.shape, ())\n\n @nnx.var_defaults(hijax=True)\n def test_eval_shape(self):\n v = nnx.Param(jnp.array(0))\n\n def f(v):\n v[...] += 1\n return v[...]\n\n y = jax.eval_shape(f, v)\n\n self.assertEqual(y.shape, ())\n\n @nnx.var_defaults(hijax=True)\n def test_no_qdd_grad(self):\n v = nnx.Param(jnp.array(3.0), hijax=False)\n\n self.assertFalse(v.hijax)\n\n def f(v):\n return v[...] ** 2\n\n grad = jax.grad(f)(v)\n\n self.assertIsInstance(grad, nnx.Param)\n self.assertEqual(grad[...], 6.0)\n\n @nnx.var_defaults(hijax=True)\n def test_no_qdd_grad_new(self):\n x = jnp.array(3.0)\n\n def f(x):\n v = nnx.Param(x, hijax=False)\n self.assertFalse(v.hijax)\n return v[...] ** 2\n\n grad = jax.grad(f)(x)\n\n self.assertIsInstance(grad, jax.Array)\n self.assertEqual(grad, 6.0)\n\n @parameterized.product(\n hijax=[True, False],\n ref=[True, False],\n )\n def test_variable_properties(self, hijax, ref):\n v = nnx.Variable(jnp.array(1), hijax=hijax, ref=ref)\n self.assertEqual(v.hijax, hijax)\n self.assertEqual(v.ref, ref)\n if hijax:\n self.assertIsInstance(v, nnx.variablelib.HijaxVariable)\n else:\n self.assertNotIsInstance(v, nnx.variablelib.HijaxVariable)\n if ref:\n self.assertIsInstance(v.get_raw_value(), jax.Ref)\n else:\n self.assertNotIsInstance(v.get_raw_value(), jax.Ref)\n\n @parameterized.product(\n hijax=[True, False],\n ref=[True, False],\n )\n def test_variable_copy_properties(self, hijax, ref):\n v_original = nnx.Variable(jnp.array(1))\n\n v = v_original.copy(hijax=hijax, ref=ref)\n self.assertEqual(v.hijax, hijax)\n self.assertEqual(v.ref, ref)\n if hijax:\n self.assertIsInstance(v, nnx.variablelib.HijaxVariable)\n else:\n self.assertNotIsInstance(v, nnx.variablelib.HijaxVariable)\n if ref:\n self.assertIsInstance(v.get_raw_value(), jax.Ref)\n else:\n self.assertNotIsInstance(v.get_raw_value(), jax.Ref)\n\n @parameterized.product(\n hijax=[True, False],\n ref=[True, False],\n )\n def test_variable_vars_as_properties(self, hijax, ref):\n v_original = nnx.Variable(jnp.array(1))\n\n v = nnx.vars_as(v_original, hijax=hijax, ref=ref)\n self.assertEqual(v.hijax, hijax)\n self.assertEqual(v.ref, ref)\n if hijax:\n self.assertIsInstance(v, nnx.variablelib.HijaxVariable)\n else:\n self.assertNotIsInstance(v, nnx.variablelib.HijaxVariable)\n if ref:\n self.assertIsInstance(v.get_raw_value(), jax.Ref)\n else:\n self.assertNotIsInstance(v.get_raw_value(), jax.Ref)\n\n\nclass TestVarDefaults(absltest.TestCase):\n def test_defaults(self):\n defaults = nnx.var_defaults()\n self.assertIsInstance(defaults, nnx.variablelib.VarDefaults)\n # Default values might depend on config/env, but generally hijax=False, ref=False initially\n self.assertFalse(defaults.hijax)\n self.assertFalse(defaults.ref)\n\n def test_context_manager_hijax(self):\n with nnx.var_defaults(hijax=True):\n self.assertTrue(nnx.var_defaults().hijax)\n v = nnx.Variable(1)\n self.assertTrue(v.hijax)\n\n self.assertFalse(nnx.var_defaults().hijax)\n v2 = nnx.Variable(1)\n self.assertFalse(v2.hijax)\n\n def test_context_manager_ref(self):\n with nnx.var_defaults(ref=True):\n self.assertTrue(nnx.var_defaults().ref)\n v = nnx.Variable(1)\n self.assertTrue(v.ref)\n\n self.assertFalse(nnx.var_defaults().ref)\n v2 = nnx.Variable(1)\n self.assertFalse(v2.ref)\n\n def test_context_manager_nested(self):\n with nnx.var_defaults(hijax=True, ref=False):\n self.assertTrue(nnx.var_defaults().hijax)\n self.assertFalse(nnx.var_defaults().ref)\n\n with nnx.var_defaults(ref=True):\n self.assertTrue(nnx.var_defaults().hijax)\n self.assertTrue(nnx.var_defaults().ref)\n\n with nnx.var_defaults(hijax=False):\n self.assertFalse(nnx.var_defaults().hijax)\n self.assertTrue(nnx.var_defaults().ref)\n\n self.assertTrue(nnx.var_defaults().hijax)\n self.assertFalse(nnx.var_defaults().ref)\n\n def test_mapping_protocol(self):\n defaults = nnx.var_defaults()\n self.assertIn('hijax', defaults)\n self.assertIn('ref', defaults)\n self.assertEqual(len(defaults), 2)\n self.assertEqual(list(defaults), ['hijax', 'ref'])\n self.assertEqual(defaults['hijax'], defaults.hijax)\n self.assertEqual(defaults['ref'], defaults.ref)\n\n def test_decorator(self):\n @nnx.var_defaults(hijax=True, ref=True)\n def f():\n return nnx.var_defaults()\n\n defaults = f()\n self.assertTrue(defaults.hijax)\n self.assertTrue(defaults.ref)\n self.assertFalse(nnx.var_defaults().hijax)\n\n def test_variable_init_override(self):\n with nnx.var_defaults(hijax=True):\n v = nnx.Variable(1, hijax=False)\n self.assertFalse(v.hijax)\n\n with nnx.var_defaults(ref=True):\n v = nnx.Variable(1, ref=False)\n self.assertFalse(v.ref)\n\nclass HijaxTransformCoverageTest(absltest.TestCase):\n # ------------\n # grad\n # ------------\n # with differentiable hijax arguments (immutable variable)\n def test_hitypes_as_grad_args(self):\n v = nnx.Variable((jnp.array(2.0), jnp.array(3.0)), hijax=False)\n\n def loss_fn(v):\n x = v[0]\n return x ** 2\n\n grads = jax.grad(loss_fn)(v)\n np.testing.assert_allclose(grads[0], 4.0)\n\n def test_hitypes_as_nondiff_grad_args(self):\n v = nnx.Variable((jnp.array(2.0), jnp.array(3.0)), hijax=False)\n x = jnp.array(3.0)\n\n def loss_fn(x, v):\n y = v[1]\n return x ** 2 + y\n\n grad = jax.grad(loss_fn)(x, v)\n np.testing.assert_allclose(grad, 6.0)\n\n def test_hitypes_as_captured_args(self):\n v = nnx.Variable((jnp.array(2.0), jnp.array(3.0)), hijax=False)\n\n def loss_fn(x):\n y = v[1]\n return x ** 2 + y\n\n grad = jax.grad(loss_fn)(jnp.array(4.0))\n np.testing.assert_allclose(grad, 8.0)\n\n # with differentiable mutable hijax arguments\n @absltest.skip(\"Not yet implemented\")\n def test_mutable_hitypes_as_grad_args(self):\n v = nnx.Variable(jnp.array(2.0), hijax=True)\n\n def loss_fn(v):\n return v[...] ** 2\n\n grads = jax.grad(loss_fn)(v)\n # NOTE: unclear what the tangent type will be here\n\n # with non-differentiable mutable hijax arguments\n def test_mutable_hitypes_as_nondiff_grad_args(self):\n v = nnx.Variable(jnp.array(2.0), hijax=True)\n x = jnp.array(3.0)\n\n def loss_fn(x, v):\n v[...] = jax.lax.stop_gradient(x * 2)\n return x ** 2 + v[...]\n\n grad = jax.grad(loss_fn)(x, v)\n np.testing.assert_allclose(v[...], 6.0)\n np.testing.assert_allclose(grad, 6.0)\n\n # with mutable hijax captured arguments\n def test_mutable_hitypes_as_captured_args(self):\n v = nnx.Variable(jnp.array(2.0), hijax=True)\n\n def loss_fn(x):\n v[...] = jax.lax.stop_gradient(x * 3)\n return x ** 2 + v[...]\n\n grad = jax.grad(loss_fn)(jnp.array(4.0))\n np.testing.assert_allclose(v[...], 12.0)\n np.testing.assert_allclose(grad, 8.0)\n\n #------------\n # scan\n #------------\n # with hijax carry arguments (immutable variable)\n @absltest.skip(\"scan not yet supported for hijax Variables\")\n def test_hitypes_as_scan_carry(self):\n v = nnx.Variable((jnp.array(1.0), jnp.array(2.0)), hijax=True, mutable=False)\n\n def body(v, _):\n x, y = v\n return nnx.Variable((x + 1.0, y + 2.0), hijax=True, mutable=False), None\n\n v_out, _ = jax.lax.scan(body, v, None, length=5)\n x, y = v_out[...]\n np.testing.assert_allclose(x, 6.0)\n np.testing.assert_allclose(y, 12.0)\n\n # with hijax extensive arguments (immutable variable)\n @absltest.skip(\"scan not yet supported for hijax Variables\")\n def test_hitypes_as_scan_extensive(self):\n v = nnx.Variable((jnp.arange(5), -jnp.arange(5)), hijax=True, mutable=False)\n\n def body(_, v_i):\n x, y = v_i\n v_i = nnx.Variable((x * 2, y * 2), hijax=True, mutable=False)\n return None, v_i\n\n _, v_out = jax.lax.scan(body, None, v)\n x, y = v_out\n np.testing.assert_allclose(x, jnp.arange(5) * 2)\n np.testing.assert_allclose(y, -jnp.arange(5) * 2)\n\n # with hijax captured arguments (immutable variable)\n @absltest.skip(\"scan not yet supported for hijax Variables\")\n def test_hitypes_as_scan_captured(self):\n v = nnx.Variable((jnp.array(3.0), jnp.array(4.0)), hijax=True, mutable=False)\n carry0 = jnp.array(1.0)\n xs = jnp.arange(5, dtype=jnp.float32)\n\n def body(carry, x):\n a, b = v\n carry = a * carry + b\n y = a * x + b\n return carry, nnx.Variable(y, hijax=True, mutable=False)\n\n carry, ys_v = jax.lax.scan(body, carry0, xs)\n ys = ys_v[...]\n np.testing.assert_allclose(carry, 727.0)\n np.testing.assert_allclose(ys, 3.0 * xs + 4.0)\n\n # with mutable hijax carry arguments\n @absltest.skip(\"has_qdd not yet supported for mutable Variable in scan carry\")\n def test_mutable_hitypes_as_scan_carry(self):\n v = nnx.Variable(jnp.array(1.0), hijax=True)\n\n def body(v, _):\n v[...] = v[...] * 2\n return v, None\n\n v_out, _ = jax.lax.scan(body, v, None, length=5)\n np.testing.assert_allclose(v_out[...], 32.0)\n\n # with mutable hijax extensive arguments\n @absltest.skip(\"Variable doesn't have shape attribute needed for scan extensive\")\n def test_mutable_hitypes_as_scan_extensive(self):\n vs = [nnx.Variable(jnp.float32(i), hijax=True) for i in range(5)]\n\n def body(_, v_i):\n val = v_i[...]\n v_i[...] = val * 2\n return None, v_i\n\n _, vs_out = jax.lax.scan(body, None, vs)\n for i, v in enumerate(vs_out):\n np.testing.assert_allclose(v[...], i * 2)\n\n # with mutable hijax captured arguments\n def test_mutable_hitypes_as_scan_captured(self):\n v = nnx.Variable(jnp.array(3.0), hijax=True)\n\n def body(_, __):\n v[...] = v[...] + 1.0\n return None, None\n\n jax.lax.scan(body, None, None, length=5)\n np.testing.assert_allclose(v[...], 8.0)\n\n def test_hijax_variable_in_jit_graph_updates_false(self):\n v = nnx.Variable(jnp.array(1.0), hijax=True)\n\n @nnx.jit(graph=True, graph_updates=False)\n def f(v, v2):\n self.assertIs(v, v2)\n v[...] += 1.0\n return v[...] * 2\n\n y = f(v, v)\n np.testing.assert_allclose(v[...], 2.0)\n np.testing.assert_allclose(y, 4.0)\n\n y = f(v, v)\n np.testing.assert_allclose(v[...], 3.0)\n np.testing.assert_allclose(y, 6.0)\n\n\nif __name__ == '__main__':\n absltest.main()\n\n" }, { "path": "tests/nnx/nn/attention_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport jax, jax.numpy as jnp\nfrom jax.lax import Precision\n\nfrom flax import linen\nfrom flax import nnx\nfrom flax.nnx.nn.attention import combine_masks\nfrom flax.typing import Dtype, PrecisionLike\n\nimport numpy as np\n\nimport typing as tp\nfrom absl.testing import parameterized\nfrom absl.testing import absltest\n\ntry:\n # JAX v0.8.0 and newer\n from jax import enable_x64\nexcept ImportError:\n from jax.experimental import enable_x64\n\n\nclass TestMultiHeadAttention(parameterized.TestCase):\n def test_basic(self):\n module = nnx.MultiHeadAttention(\n num_heads=2,\n in_features=3,\n qkv_features=6,\n out_features=6,\n rngs=nnx.Rngs(0),\n )\n y = module(jnp.ones((1, 7, 3)), decode=False)\n assert y.shape == (1, 7, 6)\n\n\n def test_multihead_sow_attention_weights(self):\n class Model(nnx.Module):\n attention_kwargs: dict\n\n def __init__(self, attention_kwargs, rng):\n self.attention_layers = nnx.data([\n nnx.MultiHeadAttention(**attention_kwargs, rngs=rng) for i in range(3)\n ])\n\n def __call__(self, x, sow_weights=False):\n x = self.attention_layers[0](x, sow_weights=sow_weights)\n x = self.attention_layers[1](x)\n x = self.attention_layers[2](x, sow_weights=sow_weights)\n return x\n\n rng = nnx.Rngs(0)\n x = jnp.ones((4, 6, 8))\n\n module = Model(\n dict(\n in_features=8,\n num_heads=8,\n kernel_init=nnx.initializers.ones_init(),\n bias_init=nnx.initializers.zeros_init(),\n deterministic=False,\n ),\n rng,\n )\n module.set_attributes(decode=False)\n\n _, intermediates = nnx.capture(module, nnx.Intermediate)(x, True)\n assert intermediates['attention_layers'][0]['attention_weights'][\n 0\n ].shape == (4, 8, 6, 6)\n assert 1 not in intermediates['attention_layers']\n assert intermediates['attention_layers'][2]['attention_weights'][\n 0\n ].shape == (4, 8, 6, 6)\n\n _, intermediates = nnx.capture(module, nnx.Intermediate)(x)\n assert not intermediates # empty\n\n def test_autoregressive_decode_with_x64(self):\n with enable_x64():\n x = jnp.ones((1, 4, 4))\n module = nnx.MultiHeadAttention(\n in_features=4,\n num_heads=2,\n qkv_features=4,\n decode=True,\n rngs=nnx.Rngs(0),\n )\n module.init_cache(x.shape, dtype=x.dtype)\n assert module.cached_key.shape == (1, 4, 2, 2)\n assert module.cached_value.shape == (1, 4, 2, 2)\n\n y1 = module(x[:, :1, :])\n y2 = module(x[:, 1:2, :])\n\n assert y1.shape == (1, 1, 4)\n assert y2.shape == (1, 1, 4)\n\n @parameterized.product(keep_rngs=[True, False])\n def test_keep_rngs(self, keep_rngs):\n rngs = nnx.Rngs(42)\n module = nnx.MultiHeadAttention(\n in_features=4,\n num_heads=2,\n qkv_features=4,\n decode=True,\n rngs=rngs,\n dropout_rate=0.5,\n keep_rngs=keep_rngs\n )\n if keep_rngs:\n assert module.rngs is not None\n else:\n assert module.rngs is None\n if keep_rngs:\n _, _, nondiff = nnx.split(module, nnx.Param, ...)\n assert isinstance(nondiff['rngs']['count'], nnx.RngCount)\n assert isinstance(nondiff['rngs']['key'], nnx.RngKey)\n else:\n nnx.split(module, nnx.Param)\n\n @parameterized.product(use_padding=[True, False], is_cross_attention=[True, False])\n def test_causal_mask_equivalence(\n self,\n use_padding: bool,\n is_cross_attention: bool\n ):\n batch_size = 1\n num_heads = 2\n q_len = 2\n kv_len = 4 if is_cross_attention else q_len\n head_dim = 4\n\n q = jax.random.normal(\n key=jax.random.key(0),\n shape=(batch_size, 1, q_len, num_heads, head_dim)\n )\n k = jax.random.normal(\n key=jax.random.key(1),\n shape=(batch_size, 1, kv_len, num_heads, head_dim)\n )\n v = jax.random.normal(\n key=jax.random.key(2),\n shape=(batch_size, 1, kv_len, num_heads, head_dim)\n )\n\n causal_mask = jnp.tril(jnp.ones(\n shape=(q_len, kv_len),\n dtype=jnp.bool_\n )\n )\n causal_mask = jnp.broadcast_to(\n array=causal_mask,\n shape=(batch_size, 1, num_heads, q_len, kv_len)\n )\n\n padding_mask = None\n\n if use_padding:\n padding_mask = jnp.ones(\n shape=(batch_size, 1, 1, q_len, kv_len),\n dtype=jnp.bool_,\n )\n padding_mask = padding_mask.at[..., -2:].set(False)\n\n manual_mask = combine_masks(padding_mask, causal_mask, dtype=q.dtype)\n\n # Jax.nn path with precombined mask and is_causal = False\n attn_jax = nnx.dot_product_attention(\n query=q,\n key=k,\n value=v,\n mask=manual_mask,\n is_causal=False,\n deterministic=True,\n module=None,\n )\n\n class DummyModule(nnx.Module):\n pass\n\n # nnx path with padding mask and is_causal = True (internally combines them)\n dummy = DummyModule()\n def _run(m):\n return nnx.dot_product_attention(\n query=q,\n key=k,\n value=v,\n mask=padding_mask,\n is_causal=True,\n deterministic=True,\n module=m,\n )\n attn_manual, _ = nnx.capture(_run, nnx.Intermediate)(dummy)\n\n np.testing.assert_allclose(attn_jax, attn_manual, atol=1e-6)\n\n# TODO: add all possible constructor argument values to parameterized.product\nclass TestLinenConsistency(parameterized.TestCase):\n @parameterized.product(\n use_bias=[True, False],\n dtype=[jnp.float32, jnp.float16],\n param_dtype=[jnp.float32, jnp.float16],\n precision=[Precision.DEFAULT, Precision.HIGH, Precision.HIGHEST],\n decode=[True, False],\n normalize_qk=[True, False],\n )\n def test_nnx_attention_equivalence(\n self,\n use_bias: bool,\n dtype: tp.Optional[Dtype],\n param_dtype: Dtype,\n precision: PrecisionLike,\n decode: bool,\n normalize_qk: bool,\n ):\n key = jax.random.key(42)\n rngs = nnx.Rngs(42)\n\n num_heads = 2\n in_features = 3\n qkv_features = 6\n out_features = 6\n\n x = jax.numpy.ones((1, in_features))\n model_nnx = nnx.MultiHeadAttention(\n num_heads=num_heads,\n in_features=in_features,\n qkv_features=qkv_features,\n out_features=out_features,\n use_bias=use_bias,\n dtype=dtype,\n param_dtype=param_dtype,\n precision=precision,\n decode=decode,\n normalize_qk=normalize_qk,\n rngs=rngs,\n )\n model = linen.MultiHeadDotProductAttention(\n num_heads=num_heads,\n qkv_features=qkv_features,\n out_features=out_features,\n use_bias=use_bias,\n dtype=dtype,\n param_dtype=param_dtype,\n precision=precision,\n decode=decode,\n normalize_qk=normalize_qk,\n )\n variables = model.init(key, x)\n\n for qkvo in ('query', 'key', 'value', 'out'):\n getattr(model_nnx, qkvo).kernel[...] = variables['params'][qkvo]['kernel']\n if use_bias:\n getattr(model_nnx, qkvo).bias[...] = variables['params'][qkvo]['bias']\n if decode:\n model_nnx.init_cache(x.shape, dtype=dtype)\n\n out_nnx = model_nnx(x)\n out, cache = model.apply(variables, x, mutable=['cache'])\n np.testing.assert_array_equal(out, out_nnx)\n\n\nclass TestKVFeatures(parameterized.TestCase):\n\n def test_varying_num_features(self):\n key = jax.random.key(42)\n rngs = nnx.Rngs(42)\n\n num_heads = 2\n in_features = 3\n in_kv_features = 4\n qkv_features = 6\n out_features = 6\n\n x = jax.numpy.ones((1, in_features))\n y = jax.random.normal(key, (1, in_kv_features))\n layer = nnx.MultiHeadAttention(\n num_heads=num_heads,\n in_features=in_features,\n qkv_features=qkv_features,\n out_features=out_features,\n in_kv_features=in_kv_features,\n rngs=rngs,\n decode=False\n )\n\n self.assertIsNotNone(layer(x, y))\n\nclass TestGQADotProductAttention(parameterized.TestCase):\n\n def test_gqa_shapes(self):\n B, T, S = 2, 4, 5\n D = 8\n num_heads_q = 6\n num_heads_kv = 3\n\n k1, k2, k3 = jax.random.split(jax.random.key(0), 3)\n query = jax.random.normal(k1, (B, T, num_heads_q, D))\n key = jax.random.normal(k2, (B, S, num_heads_kv, D))\n value = jax.random.normal(k3, (B, S, num_heads_kv, D))\n\n output = nnx.dot_product_attention(query, key, value)\n expected_shape = (B, T, num_heads_q, D)\n self.assertEqual(output.shape, expected_shape)\n\n def test_gqa_invalid_heads(self):\n B, T, D = 1, 4, 8\n query = jnp.ones((B, T, 5, D))\n key = jnp.ones((B, T, 2, D))\n value = key\n\n with self.assertRaisesRegex(ValueError, \"must be a multiple\"):\n nnx.dot_product_attention(query, key, value)\n\n def test_gqa_multihead_attention(self):\n in_feat = 128\n n_heads = 32\n n_kv_heads = 8\n qkv_feat = 2048\n head_dim = qkv_feat // n_heads\n\n model = nnx.MultiHeadAttention(\n num_heads=n_heads,\n in_features=in_feat,\n qkv_features=qkv_feat,\n num_kv_heads=n_kv_heads,\n rngs=nnx.Rngs(0),\n )\n\n assert model.query.kernel.shape == (in_feat, n_heads, head_dim)\n assert model.key.kernel.shape == (in_feat, n_kv_heads, head_dim)\n assert model.value.kernel.shape == (in_feat, n_kv_heads, head_dim)\n\n x = jnp.ones((1, 10, in_feat))\n y = model(x, decode=False)\n assert y.shape == (1, 10, in_feat)\n\n model.init_cache((1, 10, in_feat))\n assert model.cached_key.shape == (1, 10, n_kv_heads, head_dim)\n\n x_token = jnp.ones((1, 1, in_feat))\n y_token = model(x_token, decode=True)\n assert y_token.shape == (1, 1, in_feat)\n assert model.cache_index == 1\n\n def test_gqa_parity_with_jax(self):\n class DummyModule(nnx.Module):\n pass\n\n dummy_module = DummyModule()\n\n B, T, S, D = 2, 8, 8, 16\n num_heads_q = 4\n num_heads_kv = 2\n\n rng = jax.random.key(42)\n k1, k2, k3 = jax.random.split(rng, 3)\n\n query = jax.random.normal(k1, (B, T, num_heads_q, D))\n key = jax.random.normal(k2, (B, S, num_heads_kv, D))\n value = jax.random.normal(k3, (B, S, num_heads_kv, D))\n\n jax_out = jax.nn.dot_product_attention(query, key, value)\n\n # NNX should handle broadcasting internally\n def _run(m):\n return nnx.dot_product_attention(query, key, value, module=m)\n nnx_out, _ = nnx.capture(_run, nnx.Intermediate)(dummy_module)\n\n np.testing.assert_allclose(nnx_out, jax_out, atol=1e-3, rtol=1e-3)\n\n\nif __name__ == '__main__':\n absltest.main()\n\n\n" }, { "path": "tests/nnx/nn/conv_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom collections.abc import Sequence\nfrom functools import partial\nimport typing as tp\n\nimport jax\nfrom absl.testing import absltest\nfrom absl.testing import parameterized\nfrom jax import numpy as jnp\nfrom jax.lax import Precision\nimport numpy as np\n\nfrom flax import linen\nfrom flax import nnx\nfrom flax.typing import PaddingLike, Dtype, PrecisionLike\n\n\nclass TestConvLinenConsistency(parameterized.TestCase):\n @parameterized.product(\n strides=[None, (2, 3)],\n padding=['VALID', 'CIRCULAR', 'REFLECT', (4, 2)],\n input_dilation=[(2, 3)],\n kernel_dilation=[(2, 3)],\n feature_group_count=[3],\n use_bias=[True, False],\n use_mask=[False, True],\n dtype=[jnp.float32],\n param_dtype=[jnp.float16],\n precision=[Precision.HIGHEST],\n preferred_element_type=[None, jnp.float32],\n )\n def test_nnx_linen_conv_equivalence(\n self,\n strides: tp.Union[None, int, tp.Sequence[int]],\n padding: PaddingLike,\n input_dilation: tp.Union[None, int, tp.Sequence[int]],\n kernel_dilation: tp.Union[None, int, tp.Sequence[int]],\n feature_group_count: int,\n use_bias: bool,\n use_mask: bool,\n dtype: tp.Optional[Dtype],\n param_dtype: Dtype,\n precision: PrecisionLike,\n preferred_element_type: tp.Optional[Dtype],\n ):\n key = jax.random.key(42)\n rngs = nnx.Rngs(42)\n IN_FEATURES = 3\n OUT_FEATURES = 6\n INPUT_SHAPE = (24, 9, IN_FEATURES)\n kernel_size = (7, 4)\n if use_mask:\n mask = jnp.tril(jnp.ones((7, 4, 1, 6)))\n else:\n mask = None\n\n # Cannot use string padding specification for transpose conv\n if isinstance(input_dilation, Sequence) or (\n isinstance(input_dilation, int) and input_dilation > 1\n ):\n padding = (4, 2)\n\n x = jax.numpy.ones(INPUT_SHAPE)\n model_nnx = nnx.Conv(\n IN_FEATURES,\n OUT_FEATURES,\n kernel_size,\n strides,\n padding=padding,\n input_dilation=input_dilation,\n kernel_dilation=kernel_dilation,\n feature_group_count=feature_group_count,\n use_bias=use_bias,\n mask=mask,\n dtype=dtype,\n param_dtype=param_dtype,\n precision=precision,\n preferred_element_type=preferred_element_type,\n rngs=rngs,\n )\n if preferred_element_type is not None:\n conv_general_dilated = partial(\n jax.lax.conv_general_dilated,\n preferred_element_type=preferred_element_type,\n )\n else:\n conv_general_dilated = None\n model = linen.Conv(\n OUT_FEATURES,\n kernel_size=kernel_size,\n strides=strides,\n padding=padding,\n input_dilation=input_dilation,\n kernel_dilation=kernel_dilation,\n feature_group_count=feature_group_count,\n use_bias=use_bias,\n mask=mask,\n dtype=dtype,\n param_dtype=param_dtype,\n precision=precision,\n conv_general_dilated=conv_general_dilated,\n )\n variables = model.init(key, x)\n model_nnx.kernel[...] = variables['params']['kernel']\n if use_bias:\n model_nnx.bias[...] = variables['params']['bias']\n\n out_nnx = model_nnx(x)\n out = model.apply(variables, x)\n assert isinstance(out, jax.Array)\n np.testing.assert_array_equal(out, out_nnx)\n\n @parameterized.product(\n strides=[None, (2, 3)],\n padding=['VALID', 'CIRCULAR', (4, 2)],\n kernel_dilation=[(2, 3)],\n use_bias=[True, False],\n use_mask=[False, True],\n dtype=[jnp.float32],\n param_dtype=[jnp.float16],\n precision=[Precision.HIGHEST],\n preferred_element_type=[None, jnp.float32],\n )\n def test_nnx_linen_convtranspose_equivalence(\n self,\n strides: tp.Union[None, tp.Sequence[int]],\n padding: PaddingLike,\n kernel_dilation: tp.Union[None, tp.Sequence[int]],\n use_bias: bool,\n use_mask: bool,\n dtype: tp.Optional[Dtype],\n param_dtype: Dtype,\n precision: PrecisionLike,\n preferred_element_type: tp.Optional[Dtype],\n ):\n key = jax.random.key(42)\n rngs = nnx.Rngs(42)\n IN_FEATURES = 3\n OUT_FEATURES = 6\n INPUT_SHAPE = (24, 9, IN_FEATURES)\n kernel_size = (7, 4)\n if use_mask:\n mask = jnp.tril(jnp.ones((7, 4, 3, 6)))\n else:\n mask = None\n\n x = jax.numpy.ones(INPUT_SHAPE)\n model_nnx = nnx.ConvTranspose(\n IN_FEATURES,\n OUT_FEATURES,\n kernel_size,\n strides,\n padding=padding,\n kernel_dilation=kernel_dilation,\n use_bias=use_bias,\n mask=mask,\n dtype=dtype,\n param_dtype=param_dtype,\n precision=precision,\n preferred_element_type=preferred_element_type,\n rngs=rngs,\n )\n model = linen.ConvTranspose(\n OUT_FEATURES,\n kernel_size=kernel_size,\n strides=strides,\n padding=padding,\n kernel_dilation=kernel_dilation,\n use_bias=use_bias,\n mask=mask,\n dtype=dtype,\n param_dtype=param_dtype,\n precision=precision,\n preferred_element_type=preferred_element_type,\n )\n variables = model.init(key, x)\n model_nnx.kernel[...] = variables['params']['kernel']\n if use_bias:\n assert model_nnx.bias is not None\n model_nnx.bias[...] = variables['params']['bias']\n\n out_nnx = model_nnx(x)\n out = model.apply(variables, x)\n assert isinstance(out, jax.Array)\n np.testing.assert_array_equal(out, out_nnx)\n\n\nif __name__ == '__main__':\n absltest.main()" }, { "path": "tests/nnx/nn/embed_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport typing as tp\n\nimport jax\nfrom absl.testing import absltest\nfrom absl.testing import parameterized\nfrom jax import numpy as jnp\nimport numpy as np\n\nfrom flax import linen\nfrom flax import nnx\nfrom flax.typing import Dtype\n\n\nclass TestLinenConsistency(parameterized.TestCase):\n @parameterized.product(\n input_dtype=[jnp.int16, jnp.int32],\n num_embeddings=[1, 7],\n dtype=[jnp.float32, jnp.float16],\n param_dtype=[jnp.float32, jnp.float16],\n )\n def test_nnx_linen_equivalence(\n self,\n input_dtype: tp.Optional[Dtype],\n num_embeddings: int,\n dtype: tp.Optional[Dtype],\n param_dtype: Dtype,\n ):\n key = jax.random.key(42)\n rngs = nnx.Rngs(42)\n IN_FEATURES = 32\n NUM_EMBEDDINGS = num_embeddings\n\n x = jax.numpy.arange(NUM_EMBEDDINGS, dtype=input_dtype)\n model_nnx = nnx.eval_shape(\n lambda rngs: nnx.Embed(\n NUM_EMBEDDINGS,\n IN_FEATURES,\n dtype=dtype,\n param_dtype=param_dtype,\n rngs=rngs,\n ),\n rngs,\n )\n model = linen.Embed(\n NUM_EMBEDDINGS, IN_FEATURES, dtype=dtype, param_dtype=param_dtype\n )\n variables = model.init(key, x)\n model_nnx.embedding.set_value(variables['params']['embedding'])\n\n out_nnx = model_nnx(x)\n out = model.apply(variables, x)\n assert isinstance(out, jax.Array)\n np.testing.assert_array_equal(out, out_nnx)\n\n x = jax.numpy.ones((10,), dtype=input_dtype) * 10\n out_nnx = model_nnx(x)\n out = model.apply(variables, x)\n assert isinstance(out, jax.Array)\n np.testing.assert_array_equal(out, out_nnx)\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/nnx/nn/linear_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom functools import partial\nimport typing as tp\n\nimport jax\nimport jax.numpy as jnp\nfrom absl.testing import absltest\nfrom absl.testing import parameterized\nfrom jax.lax import Precision\nimport numpy as np\n\nfrom flax import linen\nfrom flax import nnx\nfrom flax.typing import Dtype, PrecisionLike, Shape\n\n\nclass TestLinearGeneral(parameterized.TestCase):\n @parameterized.product(\n dtype=[jnp.float32, jnp.float16],\n param_dtype=[jnp.float32, jnp.float16],\n precision=[Precision.DEFAULT, Precision.HIGH, Precision.HIGHEST],\n preferred_element_type=[None, jnp.float32],\n )\n def test_basic(\n self,\n dtype,\n param_dtype,\n precision,\n preferred_element_type,\n ):\n module = nnx.LinearGeneral(\n 2,\n 3,\n rngs=nnx.Rngs(0),\n dtype=dtype,\n param_dtype=param_dtype,\n precision=precision,\n preferred_element_type=preferred_element_type,\n )\n y = module(jnp.ones((1, 2)))\n\n assert y.shape == (1, 3)\n if preferred_element_type is not None:\n assert y.dtype == preferred_element_type\n assert module.kernel.shape == (2, 3)\n assert module.kernel.dtype == param_dtype\n assert module.bias is not None\n assert module.bias.shape == (3,)\n\n def test_basic_multi_features(self):\n module = nnx.LinearGeneral(2, (3, 4), rngs=nnx.Rngs(0))\n y = module(jnp.ones((1, 2)))\n\n assert y.shape == (1, 3, 4)\n assert module.kernel.shape == (2, 3, 4)\n assert module.bias is not None\n assert module.bias.shape == (3, 4)\n\n\nclass TestLinenConsistency(parameterized.TestCase):\n @parameterized.product(\n use_bias=[True, False],\n dtype=[jnp.float32, jnp.float16],\n param_dtype=[jnp.float32, jnp.float16],\n precision=[Precision.DEFAULT, Precision.HIGH, Precision.HIGHEST],\n preferred_element_type=[None, jnp.float32],\n )\n def test_nnx_linear_equivalence(\n self,\n use_bias: bool,\n dtype: tp.Optional[Dtype],\n param_dtype: Dtype,\n precision: PrecisionLike,\n preferred_element_type: tp.Optional[Dtype],\n ):\n key = jax.random.key(42)\n rngs = nnx.Rngs(42)\n IN_FEATURES = 32\n OUT_FEATURES = 64\n\n x = jax.numpy.ones((1, IN_FEATURES))\n model_nnx = nnx.eval_shape(\n lambda rngs: nnx.Linear(\n IN_FEATURES,\n OUT_FEATURES,\n use_bias=use_bias,\n dtype=dtype,\n param_dtype=param_dtype,\n precision=precision,\n preferred_element_type=preferred_element_type,\n rngs=rngs,\n ),\n rngs,\n )\n if preferred_element_type is not None:\n dot_general = partial(\n jax.lax.dot_general,\n preferred_element_type=preferred_element_type,\n )\n else:\n dot_general = None\n model = linen.Dense(\n OUT_FEATURES,\n use_bias=use_bias,\n dtype=dtype,\n param_dtype=param_dtype,\n precision=precision,\n dot_general=dot_general,\n )\n variables = model.init(key, x)\n model_nnx.kernel.set_value(variables['params']['kernel'])\n if use_bias:\n model_nnx.bias.set_value(variables['params']['bias'])\n\n out_nnx = model_nnx(x)\n out = model.apply(variables, x)\n assert isinstance(out, jax.Array)\n np.testing.assert_array_equal(out, out_nnx)\n\n @parameterized.product(\n einsum_str=['defab,bcef->adefc', 'd...ab,bc...->ad...c'],\n bias_shape=[None, (6, 7, 5)],\n dtype=[jnp.float32, jnp.float16],\n param_dtype=[jnp.float32, jnp.float16],\n precision=[Precision.DEFAULT, Precision.HIGH, Precision.HIGHEST],\n preferred_element_type=[None, jnp.float32],\n )\n def test_nnx_einsum_equivalence(\n self,\n einsum_str,\n bias_shape: tp.Optional[Shape],\n dtype: tp.Optional[Dtype],\n param_dtype: Dtype,\n precision: PrecisionLike,\n preferred_element_type: tp.Optional[Dtype],\n ):\n key = jax.random.key(42)\n rngs = nnx.Rngs(42)\n INPUT_SHAPE = (8, 6, 7, 3, 4)\n KERNEL_SHAPE = (4, 5, 6, 7)\n\n x = jax.random.normal(key, INPUT_SHAPE)\n model_nnx = nnx.Einsum(\n einsum_str,\n KERNEL_SHAPE,\n bias_shape,\n dtype=dtype,\n param_dtype=param_dtype,\n precision=precision,\n preferred_element_type=preferred_element_type,\n rngs=rngs,\n )\n model = linen.Einsum(\n KERNEL_SHAPE,\n einsum_str,\n use_bias=True if bias_shape is not None else False,\n dtype=dtype,\n param_dtype=param_dtype,\n precision=precision,\n preferred_element_type=preferred_element_type,\n )\n\n variables = model.init(key, x)\n variables['params']['kernel'] = model_nnx.kernel[...]\n if bias_shape is not None:\n assert model_nnx.bias is not None\n variables['params']['bias'] = model_nnx.bias[...]\n out_nnx = model_nnx(x)\n out = model.apply(variables, x)\n assert isinstance(out, jax.Array)\n np.testing.assert_array_equal(out, out_nnx)\n\n variables = model.init(key, x)\n model_nnx.kernel.set_value(variables['params']['kernel'])\n if bias_shape is not None:\n assert model_nnx.bias is not None\n model_nnx.bias.set_value(variables['params']['bias'])\n out_nnx = model_nnx(x)\n out = model.apply(variables, x)\n assert isinstance(out, jax.Array)\n np.testing.assert_array_equal(out, out_nnx)\n\n def test_einsum_op(self):\n def custom_einsum(*args, **kwargs):\n out = jnp.einsum(*args, **kwargs)\n return out.reshape((1, *out.shape))\n model = nnx.Einsum('ab,bc->ac', (3, 4), einsum_op=custom_einsum,\n rngs=nnx.Rngs(42))\n y = model(jnp.ones((2, 3)))\n assert y.shape == (1, 2, 4)\n\n\nclass TestPReLUConsistency(parameterized.TestCase):\n @parameterized.product(\n dtype=[jnp.float32, jnp.float16],\n param_dtype=[jnp.float32, jnp.float16],\n )\n def test_equivalence(self, dtype, param_dtype):\n key = jax.random.key(42)\n x = jnp.linspace(-10, 10, 20, dtype=dtype)\n negative_slope_init = 0.02\n nnx_prelu = nnx.PReLU(negative_slope_init=negative_slope_init, param_dtype=param_dtype)\n linen_prelu = linen.PReLU(negative_slope_init=negative_slope_init, param_dtype=param_dtype)\n\n variables = linen_prelu.init(key, x)\n expected = linen_prelu.apply(variables, x)\n output = nnx_prelu(x)\n np.testing.assert_array_equal(output, expected)\n\n # Check gradients\n @jax.jit\n def nnx_loss_function(model):\n y = model(x)\n return y.mean()\n\n @jax.jit\n def linen_loss_function(variables):\n y = linen_prelu.apply(variables, x)\n return y.mean()\n\n expected_loss, expected_grads = jax.value_and_grad(linen_loss_function)(variables)\n loss, grads = jax.value_and_grad(nnx_loss_function)(nnx_prelu)\n\n np.testing.assert_array_equal(loss, expected_loss)\n np.testing.assert_array_equal(\n expected_grads['params']['negative_slope'], grads.negative_slope[...]\n )\n\n\nclass TestLayersSameGraph(parameterized.TestCase):\n\n @parameterized.product(\n module_args_kwargs_initargs=[\n (nnx.LinearGeneral, (2, (3, 4)), (\"kernel_init\", \"bias_init\")),\n (nnx.Linear, (2, 4), (\"kernel_init\", \"bias_init\")),\n (nnx.Einsum, (\"ik,kj->ij\", (5, 4), 5), (\"kernel_init\", \"bias_init\")),\n (nnx.Conv, (2, 4, 3), (\"kernel_init\", \"bias_init\")),\n (nnx.ConvTranspose, (2, 4, 3), (\"kernel_init\", \"bias_init\")),\n (nnx.Embed, (2, 4), (\"embedding_init\",)),\n (\n nnx.MultiHeadAttention,\n (8, 5, 16),\n (\"kernel_init\", \"out_kernel_init\", \"bias_init\", \"out_bias_init\"),\n ),\n (nnx.BatchNorm, (3,), (\"scale_init\", \"bias_init\")),\n (nnx.LayerNorm, (3,), (\"scale_init\", \"bias_init\")),\n (nnx.RMSNorm, (3,), (\"scale_init\",)),\n (nnx.GroupNorm, (6, 3), (\"scale_init\", \"bias_init\")),\n (nnx.InstanceNorm, (6,), (\"scale_init\", \"bias_init\")),\n (\n nnx.LSTMCell,\n (4, 5),\n (\"kernel_init\", \"recurrent_kernel_init\", \"bias_init\"),\n ),\n (\n nnx.OptimizedLSTMCell,\n (4, 5),\n (\"kernel_init\", \"recurrent_kernel_init\", \"bias_init\"),\n ),\n (\n nnx.SimpleCell,\n (4, 5),\n (\"kernel_init\", \"recurrent_kernel_init\", \"bias_init\"),\n ),\n (\n nnx.GRUCell,\n (4, 5),\n (\"kernel_init\", \"recurrent_kernel_init\", \"bias_init\"),\n ),\n ],\n )\n def test(self, module_args_kwargs_initargs):\n module_cls, args, init_argnames = module_args_kwargs_initargs\n kwargs = {\"rngs\": nnx.Rngs(0)}\n init_zeros = nnx.initializers.zeros\n init_ones = nnx.initializers.ones\n init1_kwargs = {k: init_zeros for k in init_argnames}\n init2_kwargs = {k: init_ones for k in init_argnames}\n mod1 = module_cls(*args, **init1_kwargs, **kwargs)\n mod2 = module_cls(*args, **init2_kwargs, **kwargs)\n g1, g2 = nnx.graphdef(mod1), nnx.graphdef(mod2)\n assert g1 == g2\n\n\nclass TestLayersParamsMetadata(parameterized.TestCase):\n\n @parameterized.product(\n module_args_kwargs_initargs=[\n (nnx.LinearGeneral, (2, (3, 4)), ((\"kernel\", 2, ()), (\"bias\", 1, ()))),\n (nnx.Linear, (2, 4), ((\"kernel\", 2, ()), (\"bias\", 1, ()))),\n (nnx.Einsum, (\"ik,kj->ij\", (5, 4), 5), ((\"kernel\", 2, ()), (\"bias\", 1, ()))),\n (nnx.Conv, (2, 4, 3), ((\"kernel\", 2, ()), (\"bias\", 1, ()))),\n (nnx.ConvTranspose, (2, 4, 3), ((\"kernel\", 2, ()), (\"bias\", 1, ()))),\n (nnx.Embed, (2, 4), ((\"embedding\", 2, ()), )),\n (\n partial(nnx.MultiHeadAttention, normalize_qk=True),\n (8, 5, 16),\n (\n (\"kernel\", 2, ((\"query\", \"kernel\"), (\"key\", \"kernel\"), (\"value\", \"kernel\"))),\n (\"out_kernel\", 2, ((\"out\", \"kernel\"), )),\n (\"bias\", 1, ((\"query\", \"bias\"), (\"key\", \"bias\"), (\"value\", \"bias\"))),\n (\"out_bias\", 1, ((\"out\", \"bias\"), )),\n (\"query_ln_scale\", 1, ((\"query_ln\", \"scale\"), )),\n (\"key_ln_scale\", 1, ((\"key_ln\", \"scale\"), )),\n ),\n ),\n (nnx.BatchNorm, (3,), ((\"scale\", 1, ()), (\"bias\", 1, ()))),\n (nnx.LayerNorm, (3,), ((\"scale\", 1, ()), (\"bias\", 1, ()))),\n (nnx.RMSNorm, (3,), ((\"scale\", 1, ()), )),\n (nnx.GroupNorm, (6, 3), ((\"scale\", 1, ()), (\"bias\", 1, ()))),\n (nnx.InstanceNorm, (6,), ((\"scale\", 1, ()), (\"bias\", 1, ()))),\n (\n nnx.LoRA,\n (3, 2, 4),\n (\n (\"a\", 2, ((None, \"lora_a\"), )),\n (\"b\", 2, ((None, \"lora_b\"), )),\n )\n ),\n (\n partial(nnx.LoRALinear, lora_rank=4),\n (3, 2),\n (\n (\"a\", 2, ((\"lora\", \"lora_a\"), )),\n (\"b\", 2, ((\"lora\", \"lora_b\"), )),\n )\n ),\n (\n nnx.LSTMCell,\n (4, 5),\n (\n (\n \"kernel\",\n 2,\n ((name, \"kernel\") for name in [\"ii\", \"if_\", \"ig\", \"io\"])\n ),\n (\n \"recurrent_kernel\",\n 2,\n ((name, \"kernel\") for name in [\"hi\", \"hf\", \"hg\", \"ho\"])\n ),\n (\n \"bias\",\n 1,\n ((name, \"bias\") for name in [\"hi\", \"hf\", \"hg\", \"ho\"])\n ),\n ),\n ),\n (\n nnx.OptimizedLSTMCell,\n (4, 5),\n (\n (\"kernel\", 2, ((\"dense_i\", \"kernel\"), )),\n (\"recurrent_kernel\", 2, ((\"dense_h\", \"kernel\"), )),\n (\"bias\", 1, ((\"dense_h\", \"bias\"), )),\n )\n ),\n (\n nnx.SimpleCell,\n (4, 5),\n (\n (\"kernel\", 2, ((\"dense_i\", \"kernel\"), )),\n (\"bias\", 1, ((\"dense_i\", \"bias\"), )),\n (\"recurrent_kernel\", 2, ((\"dense_h\", \"kernel\"), )),\n )\n ),\n (\n nnx.GRUCell,\n (4, 5),\n (\n (\"kernel\", 2, ((\"dense_i\", \"kernel\"), )),\n (\"bias\", 1, ((\"dense_i\", \"bias\"), )),\n (\"recurrent_kernel\", 2, ((\"dense_h\", \"kernel\"), )),\n )\n ),\n ],\n )\n def test(self, module_args_kwargs_initargs):\n module_cls, args, metadata_argnames = module_args_kwargs_initargs\n kwargs = {\"rngs\": nnx.Rngs(0)}\n out_sharding = (\"din\", \"dout\")\n metadata_kwargs = {\n f\"{key}_metadata\": {\"out_sharding\": out_sharding[:le]}\n for key, le, _ in metadata_argnames\n }\n\n mesh = jax.make_mesh(\n (1, 1),\n out_sharding,\n axis_types=(jax.sharding.AxisType.Auto,) * len(out_sharding),\n )\n with jax.set_mesh(mesh):\n module = module_cls(*args, **metadata_kwargs, **kwargs)\n\n for key, le, attrs in metadata_argnames:\n attrs = attrs if attrs else ((None, key), )\n for attr_name, param_name in attrs:\n attr = getattr(module, attr_name) if attr_name is not None else module\n param = getattr(attr, param_name)\n self.assertIsNotNone(param.out_sharding)\n self.assertEqual(param.out_sharding, out_sharding[:le])\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/nnx/nn/lora_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport jax\nfrom absl.testing import absltest\nimport numpy as np\n\nfrom flax import nnx\nfrom jax import numpy as jnp\n\n\nclass TestLora(absltest.TestCase):\n def test_basic(self):\n module = nnx.LoRA(3, 2, 4, rngs=nnx.Rngs(0))\n x = jax.random.normal(jax.random.key(0), (1, 3))\n y = module(x)\n\n assert y.shape == (1, 4)\n assert module.lora_a.shape == (3, 2)\n assert module.lora_b.shape == (2, 4)\n np.testing.assert_allclose(y, x @ module.lora_a @ module.lora_b)\n\n def test_lora_base_module(self):\n rngs = nnx.Rngs(0)\n linear = nnx.Linear(3, 4, use_bias=False, rngs=rngs)\n module = nnx.LoRA(3, 2, 4, base_module=linear, rngs=rngs)\n x = jax.random.normal(jax.random.key(0), (1, 3))\n y = module(x)\n\n assert y.shape == (1, 4)\n assert module.base_module == linear\n assert module.base_module.kernel.shape == (3, 4)\n assert module.base_module.bias == None\n assert module.lora_a.shape == (3, 2)\n assert module.lora_b.shape == (2, 4)\n np.testing.assert_allclose(\n y, x @ linear.kernel + x @ module.lora_a @ module.lora_b\n )\n\n def test_layer_swap_lora(self):\n class MLP(nnx.Module):\n def __init__(self, dim, rngs: nnx.Rngs):\n self.linear1 = nnx.Linear(dim, dim, rngs=rngs)\n self.linear2 = nnx.Linear(dim, dim, rngs=rngs)\n\n def __call__(self, x):\n x = self.linear1(x)\n return self.linear2(x)\n\n rngs = nnx.Rngs(0)\n model = MLP(3, rngs=rngs)\n x = jax.random.normal(jax.random.key(0), (1, 3))\n y = model(x)\n\n # Replace one of the linear layers as LoRA linear layer.\n model.linear2 = nnx.LoRA(3, 4, 3, base_module=model.linear2, rngs=rngs)\n lora_y = model(x)\n\n assert y.shape == (1, 3)\n assert lora_y.shape == (1, 3)\n np.testing.assert_allclose(y, lora_y)\n a, b = model.linear2.lora_a[...], model.linear2.lora_b[...]\n np.testing.assert_allclose(y + model.linear1(x) @ a @ b, lora_y)\n\n def test_layer_swap_loralinear(self):\n class MLP(nnx.Module):\n def __init__(self, dim, rngs: nnx.Rngs):\n self.linear1 = nnx.Linear(dim, dim, rngs=rngs)\n self.linear2 = nnx.Linear(dim, dim, rngs=rngs)\n\n def __call__(self, x):\n x = self.linear1(x)\n return self.linear2(x)\n\n rngs = nnx.Rngs(0)\n model = MLP(3, rngs=rngs)\n x = jax.random.normal(jax.random.key(0), (1, 3))\n y = model(x)\n\n # Replace one of the linear layers as LoRA linear layer.\n _, state = nnx.split(\n model.linear2\n ) # To keep the kernel and bias of linear2\n model.linear2 = nnx.LoRALinear(3, 3, lora_rank=4, rngs=rngs)\n nnx.update(model.linear2, state)\n lora_y = model(x)\n\n assert y.shape == (1, 3)\n assert lora_y.shape == (1, 3)\n np.testing.assert_allclose(y, lora_y)\n a, b = model.linear2.lora.lora_a[...], model.linear2.lora.lora_b[...]\n np.testing.assert_allclose(y + model.linear1(x) @ a @ b, lora_y)\n\n def test_lora_param_type(self):\n rngs = nnx.Rngs(0)\n model = nnx.LoRA(3, 4, 2, lora_param_type=nnx.LoRAParam, rngs=rngs)\n _, lora_params, params = nnx.split(model, nnx.LoRAParam, nnx.Param)\n assert params == {}\n assert ('lora_a' in lora_params) and ('lora_b' in lora_params)\n np.testing.assert_allclose(lora_params['lora_a'][...], model.lora_a[...])\n\n model = nnx.LoRA(3, 4, 2, lora_param_type=nnx.Param, rngs=rngs)\n _, params, lora_params = nnx.split(model, nnx.Param, nnx.LoRAParam)\n assert ('lora_a' in params) and ('lora_b' in params)\n np.testing.assert_allclose(params['lora_a'][...], model.lora_a[...])\n assert lora_params == {}\n\n def test_dtype(self):\n rngs = nnx.Rngs(0)\n model = nnx.LoRA(3, 4, 2, dtype=jnp.float16, param_dtype=jnp.float32,\n rngs=rngs)\n assert model.lora_a.dtype == jnp.float32\n y = model(jnp.ones((1, 3)).astype(jnp.float32))\n assert y.dtype == jnp.float16\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/nnx/nn/normalization_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport typing as tp\n\nimport jax\nimport jax.numpy as jnp\nfrom absl.testing import absltest\nfrom absl.testing import parameterized\nimport numpy as np\n\nfrom flax import linen\nfrom flax import nnx\nfrom flax.typing import Dtype\n\n\nclass TestLinenConsistency(parameterized.TestCase):\n @parameterized.product(\n dtype=[jnp.float32, jnp.float16],\n param_dtype=[jnp.float32, jnp.float16],\n use_fast_variance=[True, False],\n mask=[None, np.array([True, False, True, False, True])],\n )\n def test_nnx_linen_batchnorm_equivalence(\n self,\n dtype: tp.Optional[Dtype],\n param_dtype: Dtype,\n use_fast_variance: bool,\n mask: tp.Optional[np.ndarray],\n ):\n class NNXModel(nnx.Module):\n def __init__(self, dtype, param_dtype, use_fast_variance, rngs):\n self.norm_layer = nnx.BatchNorm(\n 5,\n use_running_average=False,\n dtype=dtype,\n param_dtype=param_dtype,\n use_fast_variance=use_fast_variance,\n promote_dtype=lambda x, **kwargs: x, # ensure same behavior as Linen\n rngs=rngs,\n )\n self.linear = nnx.Linear(\n 5, 4, dtype=dtype, param_dtype=param_dtype, rngs=rngs\n )\n\n def __call__(self, x, *, mask=None):\n x = self.norm_layer(x, mask=mask)\n x = self.linear(x)\n return x\n\n class LinenModel(linen.Module):\n dtype: tp.Optional[Dtype] = None\n param_dtype: Dtype = jnp.float32\n use_fast_variance: bool = True\n\n def setup(self):\n self.norm_layer = linen.BatchNorm(\n use_running_average=False,\n dtype=self.dtype,\n param_dtype=self.param_dtype,\n use_fast_variance=use_fast_variance,\n )\n self.linear = linen.Dense(\n 4, dtype=self.dtype, param_dtype=self.param_dtype\n )\n\n def __call__(self, x, *, mask=None):\n x = self.norm_layer(x, mask=mask)\n x = self.linear(x)\n return x\n\n rngs = nnx.Rngs(42)\n x = jax.random.normal(jax.random.key(0), (10, 5))\n\n linen_model = LinenModel(\n dtype=dtype, param_dtype=param_dtype, use_fast_variance=use_fast_variance\n )\n variables: dict = linen_model.init(jax.random.key(1), x)\n linen_out, batch_stats = linen_model.apply(\n variables, x, mask=mask, mutable=['batch_stats']\n )\n\n nnx_model = NNXModel(\n dtype=dtype,\n param_dtype=param_dtype,\n use_fast_variance=use_fast_variance,\n rngs=rngs,\n )\n nnx_model.linear.kernel[...] = variables['params']['linear']['kernel']\n nnx_model.linear.bias[...] = variables['params']['linear']['bias']\n\n linen_out, updates = linen_model.apply(\n variables, x, mask=mask, mutable=['batch_stats']\n )\n variables.update(updates)\n nnx_out = nnx_model(x, mask=mask)\n np.testing.assert_array_equal(linen_out, nnx_out)\n # Compare BatchNorm parameters\n np.testing.assert_array_equal(\n variables['params']['norm_layer']['scale'],\n nnx_model.norm_layer.scale[...],\n )\n np.testing.assert_array_equal(\n variables['params']['norm_layer']['bias'], nnx_model.norm_layer.bias[...]\n )\n np.testing.assert_array_equal(\n variables['batch_stats']['norm_layer']['mean'],\n nnx_model.norm_layer.mean[...],\n )\n np.testing.assert_array_equal(\n variables['batch_stats']['norm_layer']['var'],\n nnx_model.norm_layer.var[...],\n )\n\n\n\n @parameterized.product(\n dtype=[jnp.float32, jnp.float16],\n param_dtype=[jnp.float32, jnp.float16],\n use_fast_variance=[True, False],\n mask=[None, np.array([True, False, True, False, True])],\n )\n def test_nnx_linen_layernorm_equivalence(\n self,\n dtype: tp.Optional[Dtype],\n param_dtype: Dtype,\n use_fast_variance: bool,\n mask: tp.Optional[np.ndarray],\n ):\n class NNXModel(nnx.Module):\n def __init__(self, dtype, param_dtype, use_fast_variance, rngs):\n self.norm_layer = nnx.LayerNorm(\n 5,\n dtype=dtype,\n param_dtype=param_dtype,\n use_fast_variance=use_fast_variance,\n promote_dtype=lambda x, **kwargs: x, # ensure same behavior as Linen\n rngs=rngs,\n )\n self.linear = nnx.Linear(\n 5, 4, dtype=dtype, param_dtype=param_dtype, rngs=rngs\n )\n\n def __call__(self, x, *, mask=None):\n x = self.norm_layer(x, mask=mask)\n x = self.linear(x)\n return x\n\n class LinenModel(linen.Module):\n dtype: tp.Optional[Dtype] = None\n param_dtype: Dtype = jnp.float32\n use_fast_variance: bool = True\n\n def setup(self):\n self.norm_layer = linen.LayerNorm(\n dtype=self.dtype,\n param_dtype=self.param_dtype,\n use_fast_variance=self.use_fast_variance,\n )\n self.linear = linen.Dense(\n 4, dtype=self.dtype, param_dtype=self.param_dtype\n )\n\n def __call__(self, x, *, mask=None):\n x = self.norm_layer(x, mask=mask)\n x = self.linear(x)\n return x\n\n rngs = nnx.Rngs(42)\n x = jax.random.normal(jax.random.key(0), (10, 5))\n\n linen_model = LinenModel(\n dtype=dtype, param_dtype=param_dtype, use_fast_variance=use_fast_variance\n )\n variables = linen_model.init(jax.random.key(1), x)\n linen_out = linen_model.apply(variables, x, mask=mask)\n assert isinstance(linen_out, jax.Array)\n\n nnx_model = NNXModel(\n dtype=dtype,\n param_dtype=param_dtype,\n use_fast_variance=use_fast_variance,\n rngs=rngs,\n )\n nnx_model.linear.kernel[...] = variables['params']['linear']['kernel']\n nnx_model.linear.bias[...] = variables['params']['linear']['bias']\n\n nnx_out = nnx_model(x, mask=mask)\n np.testing.assert_array_equal(linen_out, nnx_out)\n\n @parameterized.product(\n dtype=[jnp.float32, jnp.float16],\n param_dtype=[jnp.float32, jnp.float16],\n use_fast_variance=[True, False],\n mask=[None, np.array([True, False, True, False, True])],\n )\n def test_nnx_linen_rmsnorm_equivalence(\n self,\n dtype: tp.Optional[Dtype],\n param_dtype: Dtype,\n use_fast_variance: bool,\n mask: tp.Optional[np.ndarray],\n ):\n class NNXModel(nnx.Module):\n def __init__(self, dtype, param_dtype, use_fast_variance, rngs):\n self.norm_layer = nnx.RMSNorm(\n 5,\n dtype=dtype,\n param_dtype=param_dtype,\n use_fast_variance=use_fast_variance,\n promote_dtype=lambda x, **kwargs: x, # ensure same behavior as Linen\n rngs=rngs,\n )\n self.linear = nnx.Linear(\n 5, 4, dtype=dtype, param_dtype=param_dtype, rngs=rngs\n )\n\n def __call__(self, x, *, mask=None):\n x = self.norm_layer(x, mask=mask)\n x = self.linear(x)\n return x\n\n class LinenModel(linen.Module):\n dtype: tp.Optional[Dtype] = None\n param_dtype: Dtype = jnp.float32\n use_fast_variance: bool = True\n\n def setup(self):\n self.norm_layer = linen.RMSNorm(\n dtype=self.dtype,\n param_dtype=self.param_dtype,\n use_fast_variance=self.use_fast_variance,\n )\n self.linear = linen.Dense(\n 4, dtype=self.dtype, param_dtype=self.param_dtype\n )\n\n def __call__(self, x, *, mask=None):\n x = self.norm_layer(x, mask=mask)\n x = self.linear(x)\n return x\n\n rngs = nnx.Rngs(42)\n x = jax.random.normal(jax.random.key(0), (10, 5))\n\n linen_model = LinenModel(\n dtype=dtype, param_dtype=param_dtype, use_fast_variance=use_fast_variance\n )\n variables = linen_model.init(jax.random.key(1), x)\n linen_out = linen_model.apply(variables, x, mask=mask)\n\n nnx_model = NNXModel(\n dtype=dtype,\n param_dtype=param_dtype,\n use_fast_variance=use_fast_variance,\n rngs=rngs,\n )\n nnx_model.linear.kernel.set_value(variables['params']['linear']['kernel'])\n nnx_model.linear.bias.set_value(variables['params']['linear']['bias'])\n\n nnx_out = nnx_model(x, mask=mask)\n assert isinstance(linen_out, jax.Array)\n np.testing.assert_array_equal(linen_out, nnx_out)\n\n @parameterized.product(\n dtype=[jnp.float32, jnp.float16],\n param_dtype=[jnp.float32, jnp.float16],\n use_fast_variance=[True, False],\n mask=[None, np.array([True, False, True, False, True, False])],\n )\n def test_nnx_linen_groupnorm_equivalence(\n self,\n dtype: tp.Optional[Dtype],\n param_dtype: Dtype,\n use_fast_variance: bool,\n mask: tp.Optional[np.ndarray],\n ):\n class NNXModel(nnx.Module):\n def __init__(self, dtype, param_dtype, use_fast_variance, rngs):\n self.norm_layer = nnx.GroupNorm(\n 6,\n num_groups=3,\n dtype=dtype,\n param_dtype=param_dtype,\n use_fast_variance=use_fast_variance,\n promote_dtype=lambda x, **kwargs: x, # ensure same behavior as Linen\n rngs=rngs,\n )\n self.linear = nnx.Linear(\n 6, 4, dtype=dtype, param_dtype=param_dtype, rngs=rngs\n )\n\n def __call__(self, x, *, mask=None):\n x = self.norm_layer(x, mask=mask)\n x = self.linear(x)\n return x\n\n class LinenModel(linen.Module):\n dtype: tp.Optional[Dtype] = None\n param_dtype: Dtype = jnp.float32\n use_fast_variance: bool = True\n\n def setup(self):\n self.norm_layer = linen.GroupNorm(\n num_groups=3,\n dtype=self.dtype,\n param_dtype=self.param_dtype,\n use_fast_variance=self.use_fast_variance,\n )\n self.linear = linen.Dense(\n 4, dtype=self.dtype, param_dtype=self.param_dtype\n )\n\n def __call__(self, x, *, mask=None):\n x = self.norm_layer(x, mask=mask)\n x = self.linear(x)\n return x\n\n rngs = nnx.Rngs(42)\n x = jax.random.normal(jax.random.key(0), (10, 6))\n\n linen_model = LinenModel(\n dtype=dtype, param_dtype=param_dtype, use_fast_variance=use_fast_variance\n )\n variables = linen_model.init(jax.random.key(1), x)\n linen_out = linen_model.apply(variables, x, mask=mask)\n\n nnx_model = NNXModel(\n dtype=dtype,\n param_dtype=param_dtype,\n use_fast_variance=use_fast_variance,\n rngs=rngs,\n )\n nnx_model.linear.kernel[...] = variables['params']['linear']['kernel']\n nnx_model.linear.bias[...] = variables['params']['linear']['bias']\n\n nnx_out = nnx_model(x, mask=mask)\n assert isinstance(linen_out, jax.Array)\n np.testing.assert_array_equal(linen_out, nnx_out)\n\n @parameterized.product(\n dtype=[jnp.float32, jnp.float16],\n param_dtype=[jnp.float32, jnp.float16],\n scale_init=[\n nnx.initializers.ones,\n nnx.initializers.constant(10.0),\n nnx.initializers.constant(0.5),\n ],\n )\n def test_nnx_linen_weightnorm_equivalence(\n self,\n dtype: tp.Optional[Dtype],\n param_dtype: Dtype,\n scale_init: nnx.Initializer,\n ):\n class NNXModel(nnx.Module):\n def __init__(self, dtype, param_dtype, rngs):\n self.dense = nnx.Linear(\n 8, 4, dtype=dtype, param_dtype=param_dtype, rngs=rngs\n )\n self.normed = nnx.WeightNorm(\n self.dense,\n use_scale=True,\n scale_init=scale_init,\n feature_axes=-1,\n dtype=dtype,\n param_dtype=param_dtype,\n rngs=rngs,\n )\n\n def __call__(self, x, *, mask=None):\n return self.normed(x)\n\n class LinenModel(linen.Module):\n dtype: tp.Optional[Dtype] = None\n param_dtype: Dtype = jnp.float32\n\n def setup(self):\n self.dense = linen.Dense(\n 4, dtype=self.dtype, param_dtype=self.param_dtype\n )\n self.weight_norm = linen.WeightNorm(\n self.dense, variable_filter={'kernel'}, scale_init=scale_init\n )\n\n def __call__(self, x, *, mask=None):\n return self.weight_norm(x)\n\n rngs = nnx.Rngs(42)\n\n x = jax.random.normal(jax.random.key(0), (10, 8))\n\n linen_model = LinenModel(dtype=dtype, param_dtype=param_dtype)\n variables = linen_model.init(jax.random.key(1), x)\n\n nnx_model = NNXModel(dtype=dtype, param_dtype=param_dtype, rngs=rngs)\n nnx_model.dense.kernel.set_value(variables['params']['dense']['kernel'])\n nnx_model.dense.bias.set_value(variables['params']['dense']['bias'])\n\n linen_out = linen_model.apply(variables, x)\n nnx_out = nnx_model(x)\n\n np.testing.assert_array_equal(\n variables['params']['weight_norm']['dense/kernel/scale'],\n nnx_model.normed.scales[('kernel',)])\n np.testing.assert_array_equal(linen_out, nnx_out)\n\n @parameterized.product(\n dtype=[jnp.float32, jnp.float16],\n param_dtype=[jnp.float32, jnp.float16],\n use_fast_variance=[True, False],\n mask=[None, np.array([True, False, True, False, True, False])],\n )\n def test_nnx_linen_instancenorm_equivalence(\n self,\n dtype: tp.Optional[Dtype],\n param_dtype: Dtype,\n use_fast_variance: bool,\n mask: tp.Optional[np.ndarray],\n ):\n class NNXModel(nnx.Module):\n def __init__(self, dtype, param_dtype, use_fast_variance, rngs):\n self.norm_layer = nnx.InstanceNorm(\n 6,\n dtype=dtype,\n param_dtype=param_dtype,\n use_fast_variance=use_fast_variance,\n rngs=rngs,\n )\n self.linear = nnx.Linear(\n 6, 4, dtype=dtype, param_dtype=param_dtype, rngs=rngs\n )\n\n def __call__(self, x, *, mask=None):\n x = self.norm_layer(x, mask=mask)\n x = self.linear(x)\n return x\n\n class LinenModel(linen.Module):\n dtype: tp.Optional[Dtype] = None\n param_dtype: Dtype = jnp.float32\n use_fast_variance: bool = True\n\n def setup(self):\n self.norm_layer = linen.InstanceNorm(\n dtype=self.dtype,\n param_dtype=self.param_dtype,\n use_fast_variance=self.use_fast_variance,\n )\n self.linear = linen.Dense(\n 4, dtype=self.dtype, param_dtype=self.param_dtype\n )\n\n def __call__(self, x, *, mask=None):\n x = self.norm_layer(x, mask=mask)\n x = self.linear(x)\n return x\n\n rngs = nnx.Rngs(42)\n x = jax.random.normal(jax.random.key(0), (10, 6))\n\n linen_model = LinenModel(\n dtype=dtype, param_dtype=param_dtype, use_fast_variance=use_fast_variance\n )\n variables = linen_model.init(jax.random.key(1), x)\n linen_out = linen_model.apply(variables, x, mask=mask)\n\n nnx_model = NNXModel(\n dtype=dtype,\n param_dtype=param_dtype,\n use_fast_variance=use_fast_variance,\n rngs=rngs,\n )\n nnx_model.linear.kernel[...] = variables['params']['linear']['kernel']\n nnx_model.linear.bias[...] = variables['params']['linear']['bias']\n\n nnx_out = nnx_model(x, mask=mask)\n assert isinstance(linen_out, jax.Array)\n np.testing.assert_array_equal(linen_out, nnx_out)\n\n @parameterized.product(\n dtype=[jnp.float32, jnp.float16],\n param_dtype=[jnp.float32, jnp.float16],\n n_steps=[1, 10],\n update_stats=[True, False],\n )\n def test_nnx_linen_spectralnorm_equivalence(\n self,\n dtype: tp.Optional[Dtype],\n param_dtype: Dtype,\n n_steps: int,\n update_stats: bool,\n ):\n class NNXModel(nnx.Module):\n def __init__(self, dtype, param_dtype, rngs):\n self.seq = nnx.Sequential(\n nnx.Linear(\n 5, 4, dtype=dtype, param_dtype=param_dtype, rngs=rngs\n ),\n nnx.relu,\n nnx.BatchNorm(\n 4,\n dtype=dtype,\n param_dtype=param_dtype,\n rngs=rngs,\n use_running_average=not update_stats,\n ),\n )\n self.norm_layer = nnx.SpectralNorm(\n self.seq,\n n_steps=n_steps,\n dtype=dtype,\n param_dtype=param_dtype,\n update_stats=update_stats,\n rngs=rngs,\n )\n\n def __call__(self, x):\n return self.norm_layer(x)\n\n class LinenModel(linen.Module):\n dtype: tp.Optional[Dtype] = None\n param_dtype: Dtype = jnp.float32\n\n def setup(self):\n self.seq = linen.Sequential([\n linen.Dense(\n 4, dtype=self.dtype, param_dtype=self.param_dtype\n ),\n linen.relu,\n linen.BatchNorm(\n dtype=self.dtype,\n param_dtype=self.param_dtype,\n use_running_average=not update_stats,\n ),\n ])\n self.norm_layer = linen.SpectralNorm(self.seq, n_steps=n_steps)\n\n def __call__(self, x):\n return self.norm_layer(x, update_stats=update_stats)\n\n rngs = nnx.Rngs(42)\n x = jax.random.normal(jax.random.key(0), (10, 5))\n\n linen_model = LinenModel(dtype=dtype, param_dtype=param_dtype)\n variables = linen_model.init(jax.random.key(1), x)\n\n nnx_model = NNXModel(\n dtype=dtype, param_dtype=param_dtype, rngs=rngs\n )\n # Setup the same weights and batch stats\n var_params_seq_0 = variables['params']['seq']['layers_0']\n nnx_model.seq.layers[0].kernel.set_value(var_params_seq_0['kernel'])\n nnx_model.seq.layers[0].bias.set_value(var_params_seq_0['bias'])\n\n var_params_seq_2 = variables['params']['seq']['layers_2']\n nnx_model.seq.layers[2].scale.set_value(var_params_seq_2['scale'])\n nnx_model.seq.layers[2].bias.set_value(var_params_seq_0['bias'])\n\n var_norm_layer = variables['batch_stats']['norm_layer']\n nnx_model.norm_layer.batch_stats[\n ('layers', 0, 'kernel', 'u')\n ].set_value(var_norm_layer['seq/layers_0/kernel/u'])\n nnx_model.norm_layer.batch_stats[\n ('layers', 0, 'kernel', 'sigma')\n ].set_value(var_norm_layer['seq/layers_0/kernel/sigma'])\n\n linen_out = linen_model.apply(variables, x, mutable=['batch_stats'])\n nnx_out = nnx_model(x)\n np.testing.assert_array_equal(linen_out[0], nnx_out)\n\n np.testing.assert_array_equal(\n nnx_model.norm_layer.batch_stats[(\"layers\", 0, \"kernel\", \"u\")],\n linen_out[1]['batch_stats']['norm_layer']['seq/layers_0/kernel/u'],\n )\n np.testing.assert_array_equal(\n nnx_model.norm_layer.batch_stats[(\"layers\", 0, \"kernel\", \"sigma\")],\n linen_out[1]['batch_stats']['norm_layer']['seq/layers_0/kernel/sigma'],\n )\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/nnx/nn/recurrent_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport jax, jax.numpy as jnp\nfrom jax import random\n\nfrom flax import linen\nfrom flax import nnx\nfrom flax.nnx.nn import initializers\n\nimport numpy as np\n\nfrom absl.testing import absltest\n\n\nclass TestLSTMCell(absltest.TestCase):\n def test_basic(self):\n module = nnx.LSTMCell(\n in_features=3,\n hidden_features=4,\n rngs=nnx.Rngs(0),\n )\n x = jnp.ones((2, 3))\n carry = module.initialize_carry(x.shape, nnx.Rngs(0))\n new_carry, y = module(carry, x)\n self.assertEqual(y.shape, (2, 4))\n\n def test_lstm_sequence(self):\n \"\"\"Test LSTMCell over a sequence of inputs.\"\"\"\n module = nnx.LSTMCell(\n in_features=3,\n hidden_features=4,\n rngs=nnx.Rngs(0),\n )\n x = random.normal(random.PRNGKey(1), (5, 2, 3)) # seq_len, batch, feature\n carry = module.initialize_carry(x.shape[1:], nnx.Rngs(0))\n outputs = []\n for t in range(x.shape[0]):\n carry, y = module(carry, x[t])\n outputs.append(y)\n outputs = jnp.stack(outputs)\n self.assertEqual(outputs.shape, (5, 2, 4))\n\n def test_lstm_with_different_dtypes(self):\n \"\"\"Test LSTMCell with different data types.\"\"\"\n module = nnx.LSTMCell(\n in_features=3,\n hidden_features=4,\n dtype=jnp.bfloat16,\n param_dtype=jnp.bfloat16,\n rngs=nnx.Rngs(0),\n )\n x = jnp.ones((2, 3), dtype=jnp.bfloat16)\n carry = module.initialize_carry(x.shape, nnx.Rngs(0))\n new_carry, y = module(carry, x)\n self.assertEqual(y.dtype, jnp.bfloat16)\n self.assertEqual(y.shape, (2, 4))\n\n def test_lstm_with_custom_activations(self):\n \"\"\"Test LSTMCell with custom activation functions.\"\"\"\n module = nnx.LSTMCell(\n in_features=3,\n hidden_features=4,\n gate_fn=jax.nn.relu,\n activation_fn=jax.nn.elu,\n rngs=nnx.Rngs(0),\n )\n x = jnp.ones((1, 3))\n carry = module.initialize_carry(x.shape, nnx.Rngs(0))\n new_carry, y = module(carry, x)\n self.assertEqual(y.shape, (1, 4))\n\n def test_lstm_initialize_carry(self):\n \"\"\"Test the initialize_carry method.\"\"\"\n module = nnx.LSTMCell(\n in_features=3,\n hidden_features=4,\n rngs=nnx.Rngs(0),\n )\n x_shape = (1, 3)\n carry = module.initialize_carry(\n x_shape, nnx.Rngs(0), carry_init=initializers.ones\n )\n c, h = carry\n self.assertTrue(jnp.all(c == 1.0))\n self.assertTrue(jnp.all(h == 1.0))\n self.assertEqual(c.shape, (1, 4))\n self.assertEqual(h.shape, (1, 4))\n\n def test_lstm_with_variable_sequence_length(self):\n \"\"\"Test LSTMCell with variable sequence lengths.\"\"\"\n module = nnx.LSTMCell(in_features=3, hidden_features=4, rngs=nnx.Rngs(0))\n\n # Simulate a batch with variable sequence lengths\n x = jnp.array(\n [\n [[1, 2, 3], [4, 5, 6], [0, 0, 0]], # Sequence length 2\n [[7, 8, 9], [10, 11, 12], [13, 14, 15]], # Sequence length 3\n ]\n ) # Shape: (batch_size=2, max_seq_length=3, features=3)\n\n seq_lengths = jnp.array([2, 3]) # Actual lengths for each sequence\n batch_size = x.shape[0]\n max_seq_length = x.shape[1]\n carry = module.initialize_carry((batch_size, 3), nnx.Rngs(0))\n outputs = []\n for t in range(max_seq_length):\n input_t = x[:, t, :]\n carry, y = module(carry, input_t)\n outputs.append(y)\n outputs = jnp.stack(\n outputs, axis=1\n ) # Shape: (batch_size, max_seq_length, hidden_features)\n\n # Zero out outputs beyond the actual sequence lengths\n mask = jnp.arange(max_seq_length)[None, :] < seq_lengths[:, None]\n outputs = outputs * mask[:, :, None]\n self.assertEqual(outputs.shape, (2, 3, 4))\n\n def test_lstm_stateful(self):\n \"\"\"Test that LSTMCell maintains state across calls.\"\"\"\n module = nnx.LSTMCell(\n in_features=3,\n hidden_features=4,\n rngs=nnx.Rngs(0),\n )\n x1 = jnp.ones((1, 3))\n x2 = jnp.ones((1, 3)) * 2\n carry = module.initialize_carry(x1.shape, nnx.Rngs(0))\n carry, y1 = module(carry, x1)\n carry, y2 = module(carry, x2)\n self.assertEqual(y1.shape, (1, 4))\n self.assertEqual(y2.shape, (1, 4))\n\n def test_lstm_equivalence_with_flax_linen(self):\n \"\"\"Test that nnx.LSTMCell produces the same outputs as flax.linen.LSTMCell.\"\"\"\n in_features = 3\n hidden_features = 4\n key = random.PRNGKey(42)\n x = random.normal(key, (1, in_features))\n\n # Initialize nnx.LSTMCell\n rngs_nnx = nnx.Rngs(0)\n module_nnx = nnx.LSTMCell(\n in_features=in_features,\n hidden_features=hidden_features,\n rngs=rngs_nnx,\n )\n carry_nnx = module_nnx.initialize_carry(x.shape, rngs_nnx)\n # Initialize flax.linen.LSTMCell\n module_linen = linen.LSTMCell(\n features=hidden_features,\n )\n carry_linen = module_linen.initialize_carry(random.PRNGKey(0), x.shape)\n variables_linen = module_linen.init(random.PRNGKey(1), carry_linen, x)\n\n # Copy parameters from flax.linen.LSTMCell to nnx.LSTMCell\n params_linen = variables_linen['params']\n # Map the parameters from linen to nnx\n # Assuming the parameter names and shapes are compatible\n # For a precise mapping, you might need to adjust parameter names\n # Get the parameters from nnx module\n nnx_params = module_nnx.__dict__\n\n # Map parameters from linen to nnx\n for gate in ['i', 'f', 'g', 'o']:\n # Input kernels (input to gate)\n if gate == 'f':\n nnx_layer = getattr(module_nnx, f'if_')\n else:\n nnx_layer = getattr(module_nnx, f'i{gate}')\n linen_params = params_linen[f'i{gate}']\n nnx_layer.kernel[...] = linen_params['kernel']\n if nnx_layer.use_bias:\n nnx_layer.bias[...] = linen_params['bias']\n # Hidden kernels (hidden state to gate)\n nnx_layer = getattr(module_nnx, f'h{gate}')\n linen_params = params_linen[f'h{gate}']\n nnx_layer.kernel[...] = linen_params['kernel']\n if nnx_layer.use_bias:\n nnx_layer.bias[...] = linen_params['bias']\n\n # Run both modules\n new_carry_nnx, y_nnx = module_nnx(carry_nnx, x)\n new_carry_linen, y_linen = module_linen.apply(\n variables_linen, carry_linen, x\n )\n\n # Compare outputs\n np.testing.assert_allclose(y_nnx, y_linen, atol=1e-5)\n # Compare carries\n for c_nnx, c_linen in zip(new_carry_nnx, new_carry_linen):\n np.testing.assert_allclose(c_nnx, c_linen, atol=1e-5)\n\n\nclass TestRNN(absltest.TestCase):\n def test_rnn_with_lstm_cell(self):\n \"\"\"Test RNN module using LSTMCell.\"\"\"\n # Initialize the LSTMCell\n cell = nnx.LSTMCell(\n in_features=3,\n hidden_features=4,\n rngs=nnx.Rngs(0),\n )\n\n # Initialize the RNN module with the LSTMCell\n rnn = nnx.RNN(cell)\n\n # Create input data (batch_size=2, seq_length=5, features=3)\n x = jnp.ones((2, 5, 3))\n\n # Initialize the carry\n carry = cell.initialize_carry((2, 3), nnx.Rngs(0))\n\n # Run the RNN module\n outputs = rnn(x, initial_carry=carry)\n\n self.assertEqual(\n outputs.shape, (2, 5, 4)\n ) # Output features should match hidden_features\n\n def test_rnn_with_gru_cell(self):\n \"\"\"Test RNN module using GRUCell.\"\"\"\n # Initialize the GRUCell\n cell = nnx.GRUCell(\n in_features=3,\n hidden_features=4,\n rngs=nnx.Rngs(1),\n )\n\n # Initialize the RNN module with the GRUCell\n rnn = nnx.RNN(cell)\n\n # Create input data (batch_size=2, seq_length=5, features=3)\n x = jnp.ones((2, 5, 3))\n\n # Initialize the carry\n carry = cell.initialize_carry((2, 3), nnx.Rngs(1))\n\n # Run the RNN module\n outputs = rnn(x, initial_carry=carry)\n\n self.assertEqual(\n outputs.shape, (2, 5, 4)\n ) # Output features should match hidden_features\n\n def test_rnn_time_major(self):\n \"\"\"Test RNN module with time_major=True.\"\"\"\n # Initialize the LSTMCell\n cell = nnx.LSTMCell(\n in_features=3,\n hidden_features=4,\n rngs=nnx.Rngs(2),\n )\n\n # Initialize the RNN module with time_major=True\n rnn = nnx.RNN(cell, time_major=True)\n\n # Create input data (seq_length=5, batch_size=2, features=3)\n x = jnp.ones((5, 2, 3))\n\n # Initialize the carry\n carry = cell.initialize_carry(x.shape[1:2] + x.shape[2:], nnx.Rngs(2))\n\n # Run the RNN module\n outputs = rnn(x, initial_carry=carry)\n\n self.assertEqual(\n outputs.shape, (5, 2, 4)\n ) # Output features should match hidden_features\n\n def test_rnn_reverse(self):\n \"\"\"Test RNN module with reverse=True.\"\"\"\n # Initialize the LSTMCell\n cell = nnx.LSTMCell(\n in_features=3,\n hidden_features=4,\n rngs=nnx.Rngs(3),\n )\n\n # Initialize the RNN module with reverse=True\n rnn = nnx.RNN(cell, reverse=True)\n\n # Create input data (batch_size=2, seq_length=5, features=3)\n x = jnp.tile(jnp.arange(5), (2, 1)).reshape(\n 2, 5, 1\n ) # Distinct values to check reversal\n x = jnp.concatenate([x, x, x], axis=-1) # Shape: (2, 5, 3)\n\n # Run the RNN module\n outputs = rnn(x)\n\n # Check if the outputs are in reverse order\n outputs_reversed = outputs[:, ::-1, :]\n # Since we used distinct input values, we can compare outputs to check reversal\n # For simplicity, just check the shapes here\n self.assertEqual(outputs.shape, (2, 5, 4))\n self.assertEqual(outputs_reversed.shape, (2, 5, 4))\n\n def test_rnn_with_seq_lengths(self):\n \"\"\"Test RNN module with variable sequence lengths.\"\"\"\n # Initialize the LSTMCell\n cell = nnx.LSTMCell(\n in_features=3,\n hidden_features=4,\n rngs=nnx.Rngs(4),\n )\n\n # Initialize the RNN module\n rnn = nnx.RNN(cell, return_carry=True)\n\n # Create input data with padding (batch_size=2, seq_length=5, features=3)\n x = jnp.array(\n [\n [\n [1, 1, 1],\n [2, 2, 2],\n [3, 3, 3],\n [0, 0, 0],\n [0, 0, 0],\n ], # Sequence length 3\n [\n [4, 4, 4],\n [5, 5, 5],\n [6, 6, 6],\n [7, 7, 7],\n [8, 8, 8],\n ], # Sequence length 5\n ]\n ) # Shape: (2, 5, 3)\n\n seq_lengths = jnp.array([3, 5]) # Actual lengths for each sequence\n\n # Initialize the carry\n carry = cell.initialize_carry((2, 3), nnx.Rngs(4))\n\n # Run the RNN module\n final_carry, outputs = rnn(x, initial_carry=carry, seq_lengths=seq_lengths)\n\n self.assertEqual(outputs.shape, (2, 5, 4))\n\n self.assertEqual(\n final_carry[0].shape, (2, 4)\n ) # c: (batch_size, hidden_features)\n self.assertEqual(\n final_carry[1].shape, (2, 4)\n ) # h: (batch_size, hidden_features)\n\n # Todo: a better test by matching the outputs with the expected values\n\n def test_rnn_with_keep_order(self):\n \"\"\"Test RNN module with reverse=True and keep_order=True.\"\"\"\n # Initialize the LSTMCell\n cell = nnx.LSTMCell(\n in_features=3,\n hidden_features=4,\n rngs=nnx.Rngs(5),\n )\n\n # Initialize the RNN module with reverse=True and keep_order=True\n rnn = nnx.RNN(cell, reverse=True, keep_order=True)\n\n # Create input data (batch_size=2, seq_length=5, features=3)\n x = jnp.tile(jnp.arange(5), (2, 1)).reshape(\n 2, 5, 1\n ) # Distinct values to check reversal\n x = jnp.concatenate([x, x, x], axis=-1) # Shape: (2, 5, 3)\n\n # Initialize the carry\n carry = cell.initialize_carry((2, 3), nnx.Rngs(5))\n\n # Run the RNN module\n outputs = rnn(x, initial_carry=carry)\n\n # Check if the outputs are in the original order despite processing in reverse\n self.assertEqual(outputs.shape, (2, 5, 4))\n\n def test_rnn_equivalence_with_flax_linen(self):\n \"\"\"Test that nnx.RNN produces the same outputs as flax.linen.RNN.\"\"\"\n in_features = 3\n hidden_features = 4\n seq_length = 5\n batch_size = 2\n key = random.PRNGKey(42)\n\n # Create input data\n x = random.normal(key, (batch_size, seq_length, in_features))\n\n # Initialize nnx.LSTMCell and RNN\n rngs_nnx = nnx.Rngs(0)\n cell_nnx = nnx.LSTMCell(\n in_features=in_features,\n hidden_features=hidden_features,\n rngs=rngs_nnx,\n )\n rnn_nnx = nnx.RNN(cell_nnx)\n\n # Initialize flax.linen.LSTMCell and RNN\n cell_linen = linen.LSTMCell(features=hidden_features)\n rnn_linen = linen.RNN(cell_linen)\n carry_linen = cell_linen.initialize_carry(random.PRNGKey(0), x[:, 0].shape)\n variables_linen = rnn_linen.init(random.PRNGKey(1), x)\n\n # Copy parameters from flax.linen to nnx\n params_linen = variables_linen['params']['cell']\n # Copy cell parameters\n for gate in ['i', 'f', 'g', 'o']:\n # Input kernels\n if gate == 'f':\n nnx_layer = getattr(cell_nnx, f'if_')\n else:\n nnx_layer = getattr(cell_nnx, f'i{gate}')\n linen_params = params_linen[f'i{gate}']\n nnx_layer.kernel[...] = linen_params['kernel']\n if nnx_layer.use_bias:\n nnx_layer.bias[...] = linen_params['bias']\n # Hidden kernels\n nnx_layer = getattr(cell_nnx, f'h{gate}')\n linen_params = params_linen[f'h{gate}']\n nnx_layer.kernel[...] = linen_params['kernel']\n if nnx_layer.use_bias:\n nnx_layer.bias[...] = linen_params['bias']\n\n # Initialize carries\n carry_nnx = cell_nnx.initialize_carry((batch_size, in_features), rngs_nnx)\n\n # Run nnx.RNN\n outputs_nnx = rnn_nnx(x, initial_carry=carry_nnx)\n\n # Run flax.linen.RNN\n outputs_linen = rnn_linen.apply(\n variables_linen, x, initial_carry=carry_linen\n )\n\n # Compare outputs\n np.testing.assert_allclose(outputs_nnx, outputs_linen, atol=1e-5)\n\n def test_rnn_with_unroll(self):\n \"\"\"Test RNN module with unroll parameter.\"\"\"\n # Initialize the LSTMCell\n cell = nnx.LSTMCell(in_features=3, hidden_features=4, rngs=nnx.Rngs(6))\n\n # Initialize the RNN module with unroll=2\n rnn = nnx.RNN(cell, unroll=2)\n\n # Create input data (batch_size=2, seq_length=6, features=3)\n x = jnp.ones((2, 6, 3))\n\n # Initialize the carry\n carry = cell.initialize_carry((2, 3), nnx.Rngs(6))\n\n # Run the RNN module\n outputs = rnn(x, initial_carry=carry)\n\n self.assertEqual(\n outputs.shape, (2, 6, 4)\n ) # Output features should match hidden_features\n\n def test_rnn_with_custom_cell(self):\n \"\"\"Test RNN module with a custom RNN cell.\"\"\"\n\n class CustomRNNCell(nnx.Module):\n \"\"\"A simple custom RNN cell.\"\"\"\n\n in_features: int\n hidden_features: int\n\n def __init__(self, in_features, hidden_features, rngs):\n self.in_features = in_features\n self.hidden_features = hidden_features\n self.rngs = rngs\n self.dense = nnx.Linear(\n in_features=in_features + hidden_features,\n out_features=hidden_features,\n rngs=rngs,\n )\n\n def __call__(self, carry, inputs):\n h = carry\n x = jnp.concatenate([inputs, h], axis=-1)\n new_h = jax.nn.tanh(self.dense(x))\n return new_h, new_h\n\n def initialize_carry(self, input_shape, rngs):\n batch_size = input_shape[0]\n h = jnp.zeros((batch_size, self.hidden_features))\n return h\n\n @property\n def num_feature_axes(self) -> int:\n return 1\n\n # Initialize the custom RNN cell\n cell = CustomRNNCell(in_features=3, hidden_features=4, rngs=nnx.Rngs(7))\n\n # Initialize the RNN module\n rnn = nnx.RNN(cell)\n\n # Create input data (batch_size=2, seq_length=5, features=3)\n x = jnp.ones((2, 5, 3))\n\n # Initialize the carry\n carry = cell.initialize_carry((2, 3), cell.rngs)\n\n # Run the RNN module\n outputs = rnn(x, initial_carry=carry)\n\n self.assertEqual(\n outputs.shape, (2, 5, 4)\n ) # Output features should match hidden_features\n\n def test_rnn_with_different_dtypes(self):\n \"\"\"Test RNN module with different data types.\"\"\"\n # Initialize the LSTMCell with float16\n cell = nnx.LSTMCell(\n in_features=3,\n hidden_features=4,\n dtype=jnp.float16,\n param_dtype=jnp.float16,\n rngs=nnx.Rngs(8),\n )\n\n # Initialize the RNN module\n rnn = nnx.RNN(cell)\n\n # Create input data (batch_size=2, seq_length=5, features=3)\n x = jnp.ones((2, 5, 3), dtype=jnp.float16)\n\n # Initialize the carry\n carry = cell.initialize_carry((2, 3), nnx.Rngs(8))\n\n # Run the RNN module\n outputs = rnn(x, initial_carry=carry)\n\n self.assertEqual(outputs.dtype, jnp.float16)\n self.assertEqual(outputs.shape, (2, 5, 4))\n\n def test_rnn_with_variable_batch_size(self):\n \"\"\"Test RNN module with variable batch sizes.\"\"\"\n # Initialize the LSTMCell\n cell = nnx.LSTMCell(\n in_features=3,\n hidden_features=4,\n rngs=nnx.Rngs(9),\n )\n\n # Initialize the RNN module\n rnn = nnx.RNN(cell)\n\n for batch_size in [1, 2, 5]:\n # Create input data (batch_size, seq_length=5, features=3)\n x = jnp.ones((batch_size, 5, 3))\n\n # Initialize the carry\n carry = cell.initialize_carry((batch_size, 3), nnx.Rngs(9))\n\n # Run the RNN module\n outputs = rnn(x, initial_carry=carry)\n\n self.assertEqual(outputs.shape, (batch_size, 5, 4))\n\n def test_recurrent_dropout(self):\n class LSTMWithRecurrentDropout(nnx.OptimizedLSTMCell):\n def __init__(\n self,\n *,\n rngs: nnx.Rngs,\n in_features: int,\n hidden_features: int,\n dropout_rate: float,\n **kwargs,\n ):\n super().__init__(\n in_features=in_features,\n hidden_features=hidden_features,\n rngs=rngs,\n keep_rngs=True,\n **kwargs,\n )\n self.recurrent_dropout = nnx.Dropout(\n rate=dropout_rate, rng_collection='recurrent_dropout', rngs=rngs\n )\n\n def __call__(self, carry, x):\n h, c = carry\n new_h, new_c = super().__call__((h, c), x)\n new_h = jax.tree.map(self.recurrent_dropout, new_h)\n return new_h, new_c\n\n class RNNWithRecurrentDropout(nnx.Module):\n def __init__(\n self,\n *,\n rngs: nnx.Rngs,\n in_features: int,\n hidden_features: int = 32,\n dropout_rate: float = 0.5,\n recurrent_dropout_rate: float = 0.25,\n ):\n cell = LSTMWithRecurrentDropout(\n in_features=in_features,\n hidden_features=hidden_features,\n rngs=rngs,\n dropout_rate=recurrent_dropout_rate,\n )\n self.lstm = nnx.RNN(cell, broadcast_rngs='recurrent_dropout')\n self.dropout = nnx.Dropout(dropout_rate, rngs=rngs)\n self.dense = nnx.Linear(\n in_features=hidden_features, out_features=1, rngs=rngs\n )\n\n def __call__(self, x):\n x = self.lstm(x)\n x = self.dropout(x)\n x = x[:, -1, :] # Use only the final hidden state\n return self.dense(x)\n\n model = RNNWithRecurrentDropout(\n in_features=32,\n hidden_features=64,\n dropout_rate=0.2,\n recurrent_dropout_rate=0.1,\n rngs=nnx.Rngs(0, recurrent_dropout=1),\n )\n\n x = jnp.ones((8, 10, 32))\n self.assertEqual(model.lstm.cell.recurrent_dropout.rngs.count[...], 0)\n y = model(x)\n\n self.assertEqual(y.shape, (8, 1))\n self.assertEqual(model.lstm.cell.recurrent_dropout.rngs.count[...], 1)\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/nnx/nn/stochastic_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nimport jax.numpy as jnp\nimport numpy as np\n\nfrom flax import nnx\n\nimport pytest\n\n\nclass TestStochastic:\n def test_dropout_internal_rngs(self):\n n = 0\n m1 = nnx.Dropout(rate=0.5, deterministic=False, rngs=nnx.Rngs(dropout=0))\n m2 = nnx.Dropout(rate=0.5, deterministic=False)\n rngs2 = nnx.Rngs(dropout=0).fork()\n\n @nnx.jit\n def f(m, x, rngs=None):\n nonlocal n\n n += 1\n return m(x, rngs=rngs)\n\n x = jnp.ones((1, 10))\n assert m1.rngs is not None and m1.rngs.count[...] == 0\n\n y1 = f(m1, x)\n assert n == 1\n assert m1.rngs.count[...] == 1\n y2 = f(m2, x, rngs=rngs2)\n assert n == 2\n assert rngs2.dropout.count[...] == 1\n np.testing.assert_allclose(y1, y2)\n\n y1 = f(m1, x)\n assert m1.rngs.count[...] == 2\n y2 = f(m2, x, rngs=rngs2)\n assert rngs2.dropout.count[...] == 2\n np.testing.assert_allclose(y1, y2)\n\n assert n == 2\n\n def test_dropout_rng_override(self):\n m1 = nnx.Dropout(rate=0.5, deterministic=False, rngs=nnx.Rngs(dropout=0))\n m2 = nnx.Dropout(rate=0.5, deterministic=False, rngs=nnx.Rngs(dropout=1))\n x = jnp.ones((1, 10))\n\n y1 = m1(x)\n y2 = m2(x)\n with pytest.raises(AssertionError):\n np.testing.assert_allclose(y1, y2)\n\n y2 = m2(x, rngs=nnx.Rngs(dropout=0).fork())\n np.testing.assert_allclose(y1, y2)\n\n def test_dropout_arg_override(self):\n m = nnx.Dropout(rate=0.5)\n x = jnp.ones((1, 10))\n\n # deterministic call arg provided\n m(x, deterministic=True)\n # deterministic constructor arg provided\n m.set_attributes(deterministic=True)\n y = m(x)\n # both deterministic call and constructor arg provided\n with pytest.raises(AssertionError):\n np.testing.assert_allclose(\n y, m(x, deterministic=False, rngs=nnx.Rngs(dropout=0))\n )\n # no rng arg provided\n m.set_attributes(deterministic=False)\n with pytest.raises(\n ValueError,\n match='`deterministic` is False, but no `rngs` argument was provided to Dropout',\n ):\n m(x)\n\n def test_dropout_arg_override_view(self):\n m = nnx.Dropout(rate=0.5)\n x = jnp.ones((1, 10))\n\n # deterministic call arg provided\n m(x, deterministic=True)\n # deterministic constructor arg provided\n new_m = nnx.view(m, deterministic=True)\n y = new_m(x)\n # both deterministic call and constructor arg provided\n with pytest.raises(AssertionError):\n np.testing.assert_allclose(\n y, new_m(x, deterministic=False, rngs=nnx.Rngs(dropout=0))\n )\n # no rng arg provided\n new_m = nnx.view(m, deterministic=False)\n with pytest.raises(\n ValueError,\n match='`deterministic` is False, but no `rngs` argument was provided to Dropout',\n ):\n new_m(x)" }, { "path": "tests/nnx/optimizer_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom absl.testing import absltest\nfrom absl.testing import parameterized\nfrom flax import nnx\nimport jax\nimport jax.numpy as jnp\nimport numpy as np\nimport optax\n\ndef assert_equal(path, x, y):\n np.testing.assert_array_equal(x, y, err_msg=f'Mismatch at path: {path}')\n\n\ndef assert_not_equal(path, x, y):\n np.testing.assert_(\n np.any(np.not_equal(x, y)), msg=f'Unexpected match at path: {path}'\n )\n\n\nclass Model(nnx.Module):\n\n def __init__(self, in_features, out_features, rngs):\n self.linear1 = nnx.Linear(in_features, 3, rngs=rngs)\n self.linear2 = nnx.Linear(3, out_features, rngs=rngs)\n\n def __call__(self, x):\n return self.linear2(self.linear1(x))\n\n\nclass TestOptimizer(parameterized.TestCase):\n\n @parameterized.parameters(\n {'module_cls': nnx.Linear},\n {'module_cls': Model},\n )\n def test_split_merge(self, module_cls):\n x = jax.random.normal(jax.random.key(0), (1, 2))\n model = module_cls(2, 4, rngs=nnx.Rngs(0))\n tx = optax.adam(1e-3)\n optimizer = nnx.Optimizer(model, tx, wrt=nnx.Param)\n out = model(x)\n graphdef, optimizer = nnx.split(optimizer)\n optimizer = nnx.merge(graphdef, optimizer)\n np.testing.assert_allclose(out, model(x))\n\n def test_update(self):\n model = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n optimizer = nnx.Optimizer(model, optax.adamw(0.1), wrt=nnx.Param)\n\n def loss_fn(model):\n params = nnx.state(model)\n loss = sum(jnp.sum(x**2) for x in jax.tree.leaves(params))\n return loss\n\n grads = nnx.grad(loss_fn)(model)\n optimizer.update(model, grads)\n\n def test_sharding_propagation(self):\n with jax.set_mesh(\n jax.make_mesh(\n (1, 1),\n ('a', 'b'),\n axis_types=(jax.sharding.AxisType.Auto,) * len(('a', 'b')),\n )\n ):\n model = nnx.Linear(\n 2,\n 3,\n rngs=nnx.Rngs(0),\n kernel_init=nnx.with_partitioning(\n nnx.initializers.lecun_normal(),\n sharding=('a', 'b'),\n ),\n use_bias=False,\n )\n optimizer = nnx.Optimizer(model, optax.adamw(0.1), wrt=nnx.Param)\n\n state = nnx.state(optimizer)\n partition_spec = nnx.get_partition_spec(state)\n\n self.assertEqual(state['opt_state'][0]['mu']['kernel'].out_sharding, ('a', 'b'))\n self.assertEqual(\n partition_spec['opt_state'][0]['mu']['kernel'].get_value(),\n jax.sharding.PartitionSpec('a', 'b'),\n )\n\n @parameterized.product(\n module_cls=[nnx.Linear, Model],\n jit_decorator=[lambda f: f, nnx.jit, jax.jit],\n optimizer=[optax.sgd, optax.adam],\n )\n def test_jit(self, module_cls, jit_decorator, optimizer):\n x = jax.random.normal(jax.random.key(0), (1, 2))\n y = jnp.ones((1, 4))\n model = module_cls(2, 4, rngs=nnx.Rngs(0))\n tx = optimizer(\n 1e-3\n ) # TODO: this doesn't work with adam optimizer for some reason\n state = nnx.ModelAndOptimizer(model, tx)\n\n if jit_decorator == jax.jit:\n model_static, model_state = nnx.split(model)\n loss_fn = lambda graphdef, state, x, y: (\n (nnx.merge(graphdef, state)(x) - y) ** 2\n ).mean()\n initial_loss = loss_fn(model_static, model_state, x, y)\n\n def jax_jit_train_step(graphdef, state, x, y):\n state = nnx.merge(graphdef, state)\n model_static, model_state = nnx.split(model)\n grads = jax.grad(loss_fn, argnums=1)(model_static, model_state, x, y)\n state.update(grads)\n return nnx.split(state)\n\n graphdef, state = jit_decorator(jax_jit_train_step)(\n *nnx.split(state), x, y\n )\n state = nnx.merge(graphdef, state)\n new_loss = loss_fn(*nnx.split(state.model), x, y)\n\n else:\n loss_fn = lambda model, x, y: ((model(x) - y) ** 2).mean()\n initial_loss = loss_fn(state.model, x, y)\n\n def nnx_jit_train_step(optimizer: nnx.Optimizer, x, y):\n grads = nnx.grad(loss_fn)(optimizer.model, x, y)\n optimizer.update(grads)\n\n jit_decorator(nnx_jit_train_step)(state, x, y)\n new_loss = loss_fn(state.model, x, y)\n\n self.assertTrue(new_loss < initial_loss)\n\n @parameterized.product(\n module_cls=[nnx.Linear, Model],\n jit_decorator=[lambda f: f, nnx.jit, jax.jit],\n optimizer=[optax.lbfgs],\n )\n def test_jit_linesearch(self, module_cls, jit_decorator, optimizer):\n x = jax.random.normal(jax.random.key(0), (1, 2))\n y = jnp.ones((1, 4))\n model = module_cls(2, 4, rngs=nnx.Rngs(0))\n tx = optimizer(1e-3)\n state = nnx.ModelAndOptimizer(model, tx)\n\n if jit_decorator == jax.jit:\n model_static, model_state = nnx.split(state.model)\n loss_fn = lambda graphdef, state, x, y: (\n (nnx.merge(graphdef, state)(x) - y) ** 2\n ).mean()\n initial_loss = loss_fn(model_static, model_state, x, y)\n\n def jax_jit_train_step(graphdef, state, x, y):\n state = nnx.merge(graphdef, state)\n model_static, model_state = nnx.split(state.model)\n grads = jax.grad(loss_fn, argnums=1)(model_static, model_state, x, y)\n state.update(\n grads,\n grad=grads,\n value=initial_loss,\n value_fn=lambda state: loss_fn(model_static, state, x, y),\n )\n return nnx.split(state)\n\n graphdef, state = jit_decorator(jax_jit_train_step)(\n *nnx.split(state), x, y\n )\n state = nnx.merge(graphdef, state)\n new_loss = loss_fn(*nnx.split(state.model), x, y)\n\n else:\n graphdef = nnx.graphdef(model)\n loss_fn = lambda model, x, y: ((model(x) - y) ** 2).mean()\n\n loss_fn_split = lambda state: loss_fn(nnx.merge(graphdef, state), x, y)\n\n initial_loss = loss_fn(state.model, x, y)\n\n def nnx_jit_train_step(optimizer: nnx.Optimizer, x, y):\n grads = nnx.grad(loss_fn)(optimizer.model, x, y)\n optimizer.update(\n grads, grad=grads, value=initial_loss, value_fn=loss_fn_split\n )\n\n jit_decorator(nnx_jit_train_step)(state, x, y)\n new_loss = loss_fn(state.model, x, y)\n\n self.assertTrue(new_loss < initial_loss)\n\n @parameterized.product(\n module_cls=[nnx.Linear, Model],\n optimizer=[optax.sgd, optax.adam],\n )\n def test_metrics(self, module_cls, optimizer):\n class TrainState(nnx.ModelAndOptimizer):\n\n def __init__(self, model, tx, metrics):\n self.metrics = metrics\n super().__init__(model, tx)\n\n def update(self, *, grads, **updates): # type: ignore[signature-mismatch]\n self.metrics.update(**updates)\n super().update(grads)\n\n x = jax.random.normal(jax.random.key(0), (1, 2))\n y = jnp.ones((1, 4))\n model = module_cls(2, 4, rngs=nnx.Rngs(0))\n tx = optax.adam(1e-3)\n metrics = nnx.metrics.Average()\n state = TrainState(model, tx, metrics)\n\n loss_fn = lambda model: ((model(x) - y) ** 2).mean()\n grads = nnx.grad(loss_fn)(state.model)\n state.update(grads=grads, values=loss_fn(state.model))\n initial_loss = state.metrics.compute()\n state.update(grads=grads, values=loss_fn(state.model))\n self.assertTrue(state.metrics.compute() < initial_loss)\n\n @parameterized.parameters(\n {'variable': nnx.Param},\n {'variable': nnx.LoRAParam},\n {'variable': (nnx.Param, nnx.LoRAParam)},\n )\n def test_wrt_update(self, variable):\n in_features = 4\n out_features = 10\n model = nnx.LoRA(\n in_features=in_features,\n lora_rank=2,\n out_features=out_features,\n base_module=Model(\n in_features=in_features, out_features=out_features, rngs=nnx.Rngs(0)\n ),\n rngs=nnx.Rngs(1),\n )\n state = nnx.Optimizer(model, optax.adam(1e-3), wrt=variable)\n prev_variables, prev_other_variables = nnx.clone(nnx.state(model, variable, ...))\n\n x = jnp.ones((1, 4))\n y = jnp.ones((1, 10))\n loss_fn = lambda model, x, y: ((model(x) - y) ** 2).mean()\n grad_fn = nnx.grad(loss_fn, argnums=nnx.DiffState(0, variable))\n\n def step():\n grads = grad_fn(model, x, y)\n initial_loss = loss_fn(model, x, y)\n state.update(model, grads)\n self.assertTrue(loss_fn(model, x, y) < initial_loss)\n\n # Since lora_b is initialized to zeros by default, the gradient flow to lora_a\n # will be zeroed out in first call. Thus, run the step twice to make sure\n # lora_a is updated.\n for _ in range(2):\n step()\n\n # make sure only the Variable's filtered in `wrt` are changed, and the others are unchanged\n variables, other_variables = nnx.state(model, variable, ...)\n\n jax.tree.map_with_path(assert_not_equal, prev_variables, variables)\n\n if other_variables:\n jax.tree.map_with_path(\n assert_equal, prev_other_variables, other_variables\n )\n\n @parameterized.parameters(\n {'variable': nnx.Param},\n # {'variable': nnx.LoRAParam},\n {'variable': (nnx.Param, nnx.LoRAParam)},\n )\n def test_wrt_update_linesearch(self, variable):\n in_features = 4\n out_features = 10\n model = nnx.LoRA(\n in_features=in_features,\n lora_rank=2,\n out_features=out_features,\n base_module=Model(\n in_features=in_features, out_features=out_features, rngs=nnx.Rngs(0)\n ),\n rngs=nnx.Rngs(1),\n )\n state = nnx.Optimizer(model, optax.lbfgs(), wrt=variable)\n prev_variables, prev_other_variables = nnx.clone(nnx.state(model, variable, ...))\n\n x = jnp.ones((1, 4))\n y = jnp.ones((1, 10))\n loss_fn = lambda model, x, y: ((model(x) - y) ** 2).mean()\n\n grad_fn = nnx.grad(loss_fn, argnums=nnx.DiffState(0, variable))\n graphdef = nnx.graphdef(model)\n loss_fn_split = lambda state: loss_fn(nnx.merge(graphdef, state), x, y)\n\n def step():\n grads = grad_fn(model, x, y)\n initial_loss = loss_fn(model, x, y)\n state.update(\n model, grads, grad=grads, value_fn=loss_fn_split, value=initial_loss\n )\n self.assertTrue(loss_fn(model, x, y) < initial_loss)\n\n # Since lora_b is initialized to zeros by default, the gradient flow to lora_a\n # will be zeroed out in first call. Thus, run the step twice to make sure\n # lora_a is updated.\n for _ in range(2):\n step()\n\n # make sure only the Variable's filtered in `wrt` are changed, and the others are unchanged\n variables, other_variables = nnx.state(model, variable, ...)\n\n jax.tree.map_with_path(assert_not_equal, prev_variables, variables)\n\n if other_variables:\n jax.tree.map_with_path(\n assert_equal, prev_other_variables, other_variables\n )\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/nnx/partitioning_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom absl.testing import absltest\nfrom flax import nnx\nimport jax\nimport jax.numpy as jnp\n\n\nclass TestPartitioning(absltest.TestCase):\n\n def test_partition(self):\n m = nnx.Dict(\n a=nnx.List([nnx.Param(jnp.array(1)), nnx.BatchStat(jnp.array(2))]),\n b=nnx.Param(jnp.array(2)),\n c=100,\n )\n\n graphdef, params, rest = nnx.split(m, nnx.Param, ...)\n\n self.assertLen(params, 2)\n self.assertLen(rest, 1)\n\n # check params\n self.assertEqual(params['a'][0][...], m.a[0][...])\n self.assertEqual(params['b'][...], m.b[...])\n\n # check rest\n self.assertEqual(rest['a'][1][...], m.a[1][...])\n\n m2 = nnx.merge(graphdef, params, rest)\n\n self.assertEqual(m2.a[0][...], m.a[0][...])\n self.assertEqual(m2.a[1][...], m.a[1][...])\n self.assertEqual(m2.b[...], m.b[...])\n self.assertEqual(m2.c, 100)\n\n def test_complete_partitioning(self):\n m = nnx.Dict(\n a=nnx.List([nnx.Param(1), nnx.Param(2), nnx.Variable(3)]),\n b=nnx.Dict(c=nnx.Param(1), d=nnx.BatchStat(2)),\n )\n\n # no error\n nnx.split(m, nnx.Param, nnx.BatchStat, nnx.Variable)\n\n def test_complete_partitioning_plus_ellipsis(self):\n m = nnx.Dict(\n a=nnx.List([nnx.Param(1), nnx.Param(2), nnx.Variable(3)]),\n b=nnx.Dict(c=nnx.Param(1), d=nnx.BatchStat(2)),\n )\n\n # no error if additional ... is passed at the end\n nnx.split(m, nnx.Param, nnx.BatchStat, nnx.Variable, ...)\n\n def test_inclomplete_partition_error(self):\n m = nnx.Dict(\n a=nnx.List([nnx.Param(1), nnx.Param(2), nnx.Variable(3)]),\n b=nnx.Dict(c=nnx.Param(1), d=nnx.BatchStat(2)),\n )\n\n with self.assertRaisesRegex(\n ValueError, 'Non-exhaustive filters, got a non-empty remainder'\n ):\n nnx.split(m, nnx.Param)\n\n def test_ellipsis_not_last_error(self):\n m = nnx.Dict(\n a=nnx.List([nnx.Param(1), nnx.Param(2), nnx.Variable(3)]),\n b=nnx.Dict(c=nnx.Param(1), d=nnx.BatchStat(2)),\n )\n\n with self.assertRaisesRegex(\n ValueError, '`...` or `True` can only be used as the last filters'\n ):\n nnx.split(m, ..., nnx.Param)\n\n def test_update_from(self):\n m = nnx.Dict(\n a=nnx.List([nnx.Param(jnp.array(1)), nnx.BatchStat(jnp.array(3))]),\n b=nnx.Param(jnp.array(2)),\n c=100,\n )\n\n state = nnx.split(\n m,\n )[1]\n state = jax.tree.map(lambda x: x * 2, state)\n\n nnx.update(m, state)\n\n self.assertEqual(m.a[0][...], 2)\n self.assertEqual(m.a[1][...], 6)\n self.assertEqual(m.b[...], 4)\n self.assertEqual(m.c, 100)\n\n def test_update_from_with_array_leaf(self):\n m = nnx.Dict(\n a=nnx.List([nnx.Param(jnp.array(1)), nnx.BatchStat(jnp.array(3))]),\n b=nnx.Param(jnp.array(2)),\n c=nnx.Variable(jax.numpy.array(100)),\n )\n\n graphdef, state = nnx.split(m)\n state = jax.tree.map(lambda x: x * 2, state)\n\n nnx.update(m, state)\n\n self.assertEqual(m.a[0][...], 2)\n self.assertEqual(m.a[1][...], 6)\n self.assertEqual(m.b[...], 4)\n self.assertEqual(m.c[...], 200)\n\n def test_grad_example(self):\n m = nnx.Dict(\n a=nnx.List([nnx.Param(jnp.array(1.0)), nnx.BatchStat(jnp.array(-10))]),\n b=nnx.Param(jnp.array(2.0)),\n c=100,\n )\n\n params = nnx.state(m, nnx.Param)\n\n def loss(params):\n return sum(2 * p for p in jax.tree_util.tree_leaves(params))\n\n grads = jax.grad(loss)(params)\n nnx.update(m, grads)\n\n self.assertEqual(m.a[0][...], 2.0)\n self.assertEqual(m.a[1][...], -10)\n self.assertEqual(m.b[...], 2.0)\n self.assertEqual(m.c, 100)\n\n def test_get_paritition(self):\n m = nnx.Dict(\n a=nnx.List([nnx.Param(jnp.array(10.0)), nnx.Param(jnp.array(20.0))]),\n b=nnx.Param(jnp.array(10.0)),\n c=7,\n d=5.0,\n )\n\n state = nnx.state(m, nnx.Variable)\n self.assertEqual(state['a'][0][...], m.a[0][...])\n self.assertEqual(state['a'][1][...], m.a[1][...])\n self.assertEqual(state['b'][...], m.b[...])\n self.assertIsNot(state['b'], state['a'][0])\n self.assertLen(nnx.to_flat_state(state), 3)\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/nnx/rngs_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom functools import partial\nfrom typing import Any\n\nimport jax\nimport jax.numpy as jnp\nimport numpy as np\nfrom absl.testing import absltest, parameterized\n\nfrom flax import nnx\nfrom flax import errors\n\n\nclass TestRngs(parameterized.TestCase):\n def test_call(self):\n rngs = nnx.Rngs(0)\n key = rngs()\n\n def test_fallback(self):\n rngs = nnx.Rngs(0)\n key = rngs.dropout()\n\n def test_fallback_error_no_default(self):\n rngs = nnx.Rngs(some_name=0)\n with self.assertRaisesRegex(AttributeError, 'No RngStream named'):\n key = rngs.dropout()\n\n def test_rng_stream(self):\n key0 = jax.random.key(0)\n rngs = nnx.Rngs(params=key0)\n self.assertEqual(rngs.params.count[...], 0)\n\n key1 = rngs.params()\n self.assertEqual(rngs.params.count[...], 1)\n self.assertIs(rngs.params.key[...], key0)\n self.assertFalse(jnp.allclose(key0, key1))\n\n key2 = rngs.params()\n self.assertEqual(rngs.params.count[...], 2)\n self.assertIs(rngs.params.key[...], key0)\n self.assertFalse(jnp.allclose(key1, key2))\n\n def test_rng_trace_level_constraints(self):\n rngs = nnx.Rngs(0)\n\n @jax.jit\n def f():\n with self.assertRaisesRegex(\n errors.TraceContextError,\n 'Cannot mutate RngCount from a different trace level',\n ):\n rngs.params()\n\n f()\n\n rngs1: Any = None\n\n @jax.jit\n def h():\n nonlocal rngs1\n rngs1 = nnx.Rngs(1)\n\n h()\n\n self.assertIsInstance(rngs1, nnx.Rngs)\n with self.assertRaisesRegex(\n errors.TraceContextError,\n 'Cannot mutate RngCount from a different trace level',\n ):\n rngs1.params()\n\n def test_jit_updates(self):\n class Foo(nnx.Module):\n def __init__(self, not_rngs):\n rngs = not_rngs\n self.linear = nnx.Linear(2, 2, rngs=rngs)\n self.dropout = nnx.Dropout(0.5, deterministic=False)\n\n def __call__(self, x, rngs):\n x = self.linear(x)\n x = self.dropout(x, rngs=rngs)\n return x\n\n rngs = nnx.Rngs(0)\n m = Foo(rngs)\n\n # +1 for the Linear kernel, +1 for the Linear bias\n self.assertEqual(rngs['default'].count[...], 2)\n\n @nnx.jit\n def f(m: Foo, x: jax.Array, not_rngs: nnx.Rngs):\n rngs = not_rngs\n x = m(x, rngs)\n x = m(x, rngs)\n return x\n\n x = jnp.ones((2, 2))\n x = f(m, x, rngs)\n\n # +1 for the Dropout mask\n self.assertEqual(rngs['default'].count[...], 4)\n\n def test_lifting_rng_state(self):\n class Foo(nnx.Module):\n def __init__(self, rngs):\n self.rngs = rngs\n self.dropout = nnx.Dropout(0.5, deterministic=False)\n self.linear = nnx.Linear(2, 3, rngs=rngs)\n\n def __call__(self, x):\n x = self.linear(x)\n x = self.dropout(x, rngs=self.rngs)\n return x\n\n rngs = nnx.Rngs(params=0, dropout=1)\n m = Foo(rngs)\n graphdef, params, rng_counts, dropout_keys, param_keys = nnx.split(\n m, nnx.Param, nnx.RngCount, 'dropout', 'params'\n )\n\n self.assertEqual(m.rngs.params.count[...], 2)\n self.assertEqual(m.rngs['dropout'].count[...], 0)\n self.assertLen(nnx.to_flat_state(dropout_keys), 1)\n self.assertLen(nnx.to_flat_state(param_keys), 1)\n self.assertLen(nnx.to_flat_state(rng_counts), 2)\n\n # split dropout keys\n split_dropout_keys = jax.tree.map(\n lambda x: jax.random.split(x, 4), dropout_keys\n )\n # replicate params\n params = jax.tree.map(lambda x: jnp.stack([x] * 4, axis=0), params)\n\n @partial(\n jax.vmap,\n in_axes=(0, 0, None, None, 0),\n out_axes=(0, 0, None),\n )\n def f(params, dropout_keys, param_keys, rng_counts, x):\n m = nnx.merge(graphdef, params, dropout_keys, param_keys, rng_counts)\n y = m(x)\n _, params, rng_counts, dropout_keys, param_keys = nnx.split(\n m, nnx.Param, nnx.RngCount, 'dropout', 'params'\n )\n return y, params, rng_counts\n\n x = jnp.ones((4, 1, 2))\n y, params, rng_counts = f(\n params,\n split_dropout_keys,\n param_keys,\n rng_counts,\n x,\n )\n\n nnx.update(m, params, dropout_keys, param_keys, rng_counts)\n\n self.assertEqual(y.shape, (4, 1, 3))\n self.assertEqual(m.rngs.params.count[...], 2)\n self.assertEqual(m.rngs['dropout'].count[...], 1)\n\n @parameterized.parameters(True, False)\n def test_reseed(self, graph):\n class Model(nnx.Module):\n def __init__(self, rngs):\n self.linear = nnx.Linear(2, 3, rngs=rngs)\n self.dropout = nnx.Dropout(0.5, rngs=rngs)\n\n def __call__(self, x):\n return self.dropout(self.linear(x))\n\n model = Model(nnx.Rngs(params=0, dropout=42))\n x = jnp.ones((1, 2))\n\n y1 = model(x)\n\n nnx.reseed(model, graph=graph, dropout=42)\n y2 = model(x)\n\n np.testing.assert_allclose(y1, y2)\n\n @parameterized.parameters(True, False)\n def test_split_rngs(self, graph):\n rngs = nnx.Rngs(params=0, dropout=1)\n result = nnx.split_rngs(rngs, splits=5, graph=graph)\n if graph:\n self.assertEqual(rngs.params.key.shape, (5,))\n self.assertEqual(rngs['dropout'].key.shape, (5,))\n nnx.restore_rngs(result)\n self.assertEqual(rngs.params.key.shape, ())\n self.assertEqual(rngs['dropout'].key.shape, ())\n else:\n self.assertEqual(rngs.params.key.shape, ())\n self.assertEqual(rngs['dropout'].key.shape, ())\n self.assertEqual(result.params.key.shape, (5,))\n self.assertEqual(result['dropout'].key.shape, (5,))\n\n @parameterized.parameters(True, False)\n def test_fork_rngs(self, graph):\n rngs = nnx.Rngs(params=0, dropout=1)\n backups = nnx.fork_rngs(rngs, graph=graph)\n new_key = rngs.params.key.copy()\n nnx.restore_rngs(backups)\n self.assertNotEqual(rngs.params.key, new_key)\n\n def test_random_helpers(self):\n rngs = nnx.Rngs(0, params=1)\n\n x_nnx = rngs.normal((2, 3))\n x_jax = jax.random.normal(jax.random.fold_in(jax.random.key(0), 0), (2, 3))\n np.testing.assert_allclose(x_nnx, x_jax)\n\n x_nnx = rngs.params.uniform((2, 3))\n x_jax = jax.random.uniform(jax.random.fold_in(jax.random.key(1), 0), (2, 3))\n np.testing.assert_allclose(x_nnx, x_jax)\n\n x_nnx = rngs.lecun_normal()((2, 3))\n x_jax = jax.nn.initializers.lecun_normal()(\n jax.random.fold_in(jax.random.key(0), 1), (2, 3)\n )\n np.testing.assert_allclose(x_nnx, x_jax)\n\n x_nnx = rngs.params.lecun_uniform()((2, 3))\n x_jax = jax.nn.initializers.lecun_uniform()(\n jax.random.fold_in(jax.random.key(1), 1), (2, 3)\n )\n np.testing.assert_allclose(x_nnx, x_jax)\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/nnx/spmd_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nos.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=4'\n\nfrom absl.testing import absltest\nfrom absl.testing import parameterized\nfrom flax import nnx\nimport jax\nimport jax.numpy as jnp\nfrom jax.sharding import PartitionSpec as P, NamedSharding, AxisType, reshard\nfrom jax.experimental.layout import Format, Layout\n\nimport optax\n\n\nclass TestSPMD(parameterized.TestCase):\n\n def setUp(self):\n if jax.device_count() < 4:\n self.skipTest('At least 4 devices required')\n\n def test_init(self):\n class Foo(nnx.Module):\n def __init__(self):\n self.w = nnx.Param(\n nnx.with_partitioning(\n lambda: jnp.ones((8, 2)),\n sharding=('model', 'data'),\n )()\n )\n\n def __call__(self, x):\n return x @ self.w\n\n @jax.jit\n def create_module():\n return nnx.split(Foo())\n\n mesh = jax.make_mesh(\n (2, 2),\n ('model', 'data'),\n axis_types=(jax.sharding.AxisType.Auto,) * len(('model', 'data')),\n )\n\n with jax.set_mesh(mesh):\n m: Foo = nnx.merge(*create_module()) # type: ignore[invalid-annotation]\n x = jax.device_put(jnp.zeros((4, 8)), P(None, 'model'))\n y = m(x)\n\n assert m.w.shape == (8, 2)\n assert m.w.sharding.shard_shape(m.w.shape) == (4, 1)\n\n def test_init_all_devices(self):\n class Foo(nnx.Module):\n def __init__(self):\n self.w = nnx.Param(\n nnx.with_partitioning(\n lambda: jnp.ones((8, 2)),\n sharding=('model', 'data'),\n )()\n )\n\n def __call__(self, x):\n return x @ self.w\n\n @jax.jit\n def create_module():\n return nnx.split(Foo())\n\n mesh = jax.make_mesh(\n (1, 1),\n ('model', 'data'),\n axis_types=(jax.sharding.AxisType.Auto,) * len(('model', 'data')),\n )\n\n with jax.set_mesh(mesh):\n m: Foo = nnx.merge(*create_module()) # type: ignore[invalid-annotation]\n\n assert m.w.shape == (8, 2)\n assert m.w.sharding.shard_shape(m.w.shape) == (8, 2)\n\n def test_shard_optimizer_state(self):\n class Foo(nnx.Module):\n def __init__(self):\n self.w = nnx.Param(\n nnx.with_partitioning(\n lambda: jnp.ones((8, 2)),\n sharding=('row', 'col'),\n )()\n )\n\n def __call__(self, x):\n return x @ self.w\n\n mesh = jax.make_mesh(\n (2, 2),\n ('row', 'col'),\n axis_types=(jax.sharding.AxisType.Auto,) * len(('row', 'col')),\n )\n with jax.set_mesh(mesh):\n graphdef, params = nnx.split(Foo())\n state = nnx.TrainState.create(\n graphdef,\n params=params,\n tx=optax.adam(1e-3),\n )\n\n assert state.params['w'].sharding.is_equivalent_to(\n NamedSharding(mesh, P('row', 'col')), ndim=2)\n assert state.opt_state[0].mu['w'].sharding.is_equivalent_to(\n NamedSharding(mesh, P('row', 'col')), ndim=2)\n assert state.opt_state[0].nu['w'].sharding.is_equivalent_to(\n NamedSharding(mesh, P('row', 'col')), ndim=2)\n\n def test_add_remove_axis_in_transform(self):\n test = self\n kadds, kremoves, badds, bremoves = [], [], [], []\n class MLP(nnx.Module):\n\n @nnx.split_rngs(splits=5)\n @nnx.vmap(\n in_axes=(0, 0),\n transform_metadata={nnx.PARTITION_NAME: 'layers', 'nickname': 'nick'},\n )\n def __init__(self, rngs: nnx.Rngs):\n self.linear = nnx.Linear(\n 4,\n 4,\n kernel_init=nnx.with_metadata(\n nnx.initializers.lecun_normal(),\n out_sharding=('din', 'dout'),\n nickname=('in', 'out'),\n on_add_axis=lambda _, idx, name: kadds.append((idx, name)),\n on_remove_axis=lambda _, idx, name: kremoves.append((idx, name)),\n ),\n bias_init=nnx.with_metadata(\n nnx.initializers.zeros_init(), # no sharding annotation here!\n on_add_axis=lambda _, idx, name: badds.append((idx, name)),\n on_remove_axis=lambda _, idx, name: bremoves.append((idx, name)),\n ),\n rngs=rngs,\n )\n\n @nnx.scan(\n in_axes=(0, nnx.Carry),\n transform_metadata={nnx.PARTITION_NAME: 'layers'}\n )\n def __call__(self, x: jax.Array):\n x = self.linear(x)\n # test sharding layer axes is not present inside scan\n test.assertEqual(self.linear.kernel.shape, (4, 4))\n test.assertEqual(self.linear.kernel.out_sharding, ('din', 'dout'))\n # at least a remove_axis was already called to remove the layer axis\n test.assertEqual(kremoves[-1], (0, 'layers'))\n test.assertEqual(bremoves[-1], (0, 'layers'))\n return x, None\n\n mesh = jax.make_mesh(\n (1, 2, 2),\n ('layers', 'din', 'dout'),\n axis_types=(jax.sharding.AxisType.Auto,)\n * len(('layers', 'din', 'dout')),\n )\n with jax.set_mesh(mesh):\n m = MLP(rngs=nnx.Rngs(0))\n self.assertEqual(m.linear.kernel.shape, (5, 4, 4))\n self.assertEqual(m.linear.kernel.out_sharding, ('layers', 'din', 'dout'))\n self.assertEqual(m.linear.kernel.nickname, ('nick', 'in', 'out'))\n self.assertEqual(m.linear.bias.shape, (5, 4))\n # One add_axis called to add the `nnx.vmap` dimension\n self.assertEqual(kadds, [(0, 'layers')])\n self.assertEqual(kremoves, [])\n self.assertEqual(badds, [(0, 'layers')])\n self.assertEqual(bremoves, [])\n\n # One remove_axis and one add_axis called when in and out of `nnx.scan`\n with jax.set_mesh(mesh):\n _ = m(jnp.ones((5, 4)))\n self.assertEqual(kadds, [(0, 'layers'), (0, 'layers')])\n self.assertEqual(kremoves, [(0, 'layers')])\n self.assertEqual(badds, [(0, 'layers'), (0, 'layers')])\n self.assertEqual(bremoves, [(0, 'layers')])\n\n def test_transform_metadata_decorator(self):\n v = nnx.Param(\n jnp.array(0),\n out_sharding=('din', 'dout'),\n eager_sharding=False,\n )\n\n @nnx.transform_metadata(in_axes=0, out_axes=1, partition='din')\n def f(v):\n v[...] += 1\n self.assertEqual(v.out_sharding, ('dout',))\n v2 = nnx.Param(\n jnp.array(10),\n out_sharding=('dmid', 'dout'),\n eager_sharding=False,\n )\n return v2\n\n v2 = f(v)\n self.assertEqual(v.out_sharding, ('din', 'dout'))\n self.assertEqual(v[...], 1)\n self.assertEqual(v2.out_sharding, ('dmid', 'din', 'dout'))\n self.assertEqual(v2[...], 10)\n\n\n @parameterized.product(use_eager_sharding=[True, False])\n def test_eager_sharding_context(self, use_eager_sharding):\n rngs = nnx.Rngs(0)\n with nnx.use_eager_sharding(use_eager_sharding):\n mesh = jax.make_mesh(\n (2, 2),\n ('data', 'model'),\n axis_types=(jax.sharding.AxisType.Auto, jax.sharding.AxisType.Auto),\n )\n with jax.set_mesh(mesh):\n w = nnx.Param(\n rngs.lecun_normal()((4, 8)),\n out_sharding=(None, 'model'))\n if use_eager_sharding:\n assert has_sharding_spec(w)\n else:\n assert not has_sharding_spec(w)\n\n def test_out_sharding_linear_layers(self):\n mesh = jax.make_mesh((2, 2), (\"X\", \"Y\"), axis_types=(AxisType.Explicit, AxisType.Explicit))\n with jax.set_mesh(mesh):\n replicated_array = jnp.arange(4).reshape(2, 2)\n sharded_array = reshard(replicated_array, P(\"X\", None))\n layers = [\n nnx.Linear(2, 4, rngs=nnx.Rngs(0)),\n nnx.LinearGeneral(2, 4, rngs=nnx.Rngs(0)),\n nnx.Einsum('ab,bc->ac', (2, 4), (4,), rngs=nnx.Rngs(0)),\n ]\n for layer in layers:\n assert 'float32[2@X,4]' in str(jax.typeof(layer(sharded_array)))\n assert 'float32[2@X,4@Y]' in str(jax.typeof(layer(sharded_array, out_sharding=P(\"X\", \"Y\"))))\n\n def test_out_sharding_embed(self):\n mesh = jax.make_mesh((2, 2), (\"X\", \"Y\"), axis_types=(AxisType.Explicit, AxisType.Explicit))\n with jax.set_mesh(mesh):\n emb = nnx.Embed(num_embeddings=8, features=4, rngs=nnx.Rngs(0))\n emb = reshard(emb, P(\"X\"))\n sharded_array = reshard(jnp.arange(4), P(\"Y\"))\n self.assertRaises(Exception, emb, sharded_array)\n self.assertEqual('float32[4@X,4]',\n str(jax.typeof(emb(sharded_array, out_sharding=P(\"X\")))))\n\n def test_out_sharding_conv(self):\n mesh = jax.make_mesh((2, 2), (\"X\", \"Y\"), axis_types=(AxisType.Explicit, AxisType.Explicit))\n with jax.set_mesh(mesh):\n replicated_array = jnp.arange(32).reshape(2, 4, 4).astype(jnp.float32)\n sharded_array = reshard(replicated_array, P(\"X\", None, None))\n layer = nnx.Conv(4, 8, kernel_size=(3,), rngs=nnx.Rngs(0))\n assert 'float32[2@X,4,8]' in str(jax.typeof(layer(sharded_array)))\n assert 'float32[2@X,4@Y,8]' in str(jax.typeof(layer(sharded_array, out_sharding=P(\"X\", \"Y\", None))))\n\n def test_out_sharding_embed_attend(self):\n mesh = jax.make_mesh((2, 2), (\"X\", \"Y\"), axis_types=(AxisType.Explicit, AxisType.Explicit))\n with jax.set_mesh(mesh):\n replicated_array = jnp.arange(8).reshape(2, 4).astype(jnp.float32)\n sharded_array = reshard(replicated_array, P(\"X\", None))\n layer = nnx.Embed(num_embeddings=10, features=4, rngs=nnx.Rngs(0))\n assert 'float32[2@X,10]' in str(jax.typeof(layer.attend(sharded_array)))\n assert 'float32[2@X,10@Y]' in str(jax.typeof(layer.attend(sharded_array, out_sharding=P(\"X\", \"Y\"))))\n\n def test_out_sharding_dropout(self):\n mesh = jax.make_mesh((2, 2), (\"X\", \"Y\"), axis_types=(AxisType.Explicit, AxisType.Explicit))\n with jax.set_mesh(mesh):\n replicated_array = jnp.arange(8).reshape(2, 4).astype(jnp.float32)\n sharded_array = reshard(replicated_array, P(\"X\", None))\n layers = [\n nnx.Dropout(rate=0.5, rngs=nnx.Rngs(0)),\n nnx.Dropout(rate=0.5, broadcast_dims=(1,), rngs=nnx.Rngs(0)),\n ]\n for layer in layers:\n assert 'float32[2@X,4]' in str(jax.typeof(layer(sharded_array)))\n\n @jax.jit\n def func(x, rngs):\n return layer(x, rngs=rngs)\n\n assert 'float32[2@X,4]' in str(jax.typeof(func(sharded_array, nnx.Rngs(0))))\n\n @parameterized.product(use_hijax=[True, False])\n def test_logical_rules(self, use_hijax):\n self.enter_context(nnx.var_defaults(hijax=use_hijax))\n class Foo(nnx.Module):\n\n def __init__(self):\n self.w = nnx.Param(\n nnx.with_partitioning(\n lambda: jnp.ones((8, 2)),\n sharding=('row-alias', 'col-alias'),\n sharding_rules=(('row-alias', 'row'),),\n )()\n )\n self.b = nnx.Param(\n nnx.with_partitioning(\n lambda: jnp.zeros((2,)), sharding=('col-alias',)\n )()\n )\n\n def __call__(self, x):\n return x @ self.w + self.b\n\n mesh = jax.make_mesh(\n (1, 2, 2),\n ('layers', 'row', 'col'),\n axis_types=(jax.sharding.AxisType.Auto,)\n * len(('layers', 'row', 'col')),\n )\n with jax.set_mesh(mesh), nnx.logical_axis_rules((('col-alias', 'col'),)):\n model = Foo()\n optimizer = nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param)\n\n assert model.w.sharding.is_equivalent_to(\n NamedSharding(mesh, P('row', 'col')), ndim=2)\n assert optimizer.opt_state[0].mu['w'].sharding.is_equivalent_to(\n NamedSharding(mesh, P('row', 'col')), ndim=2)\n assert optimizer.opt_state[0].nu['w'].sharding.is_equivalent_to(\n NamedSharding(mesh, P('row', 'col')), ndim=2)\n\n def test_get_abstract_model(self):\n class Foo(nnx.Module):\n def __init__(self, rngs):\n self.linear = nnx.Linear(\n 8, 8, rngs=rngs, use_bias=False,\n kernel_init=nnx.with_partitioning(\n nnx.initializers.lecun_normal(), (None, 'model')))\n self.shared = self.linear.kernel\n\n mesh = jax.make_mesh(\n (2, 2),\n ('batch', 'model'),\n axis_types=(jax.sharding.AxisType.Auto,) * len(('batch', 'model')),\n )\n gdef, abs_state = nnx.get_abstract_model(lambda: Foo(nnx.Rngs(0)), mesh)\n assert len(jax.tree.leaves(abs_state)) == 1\n assert jax.tree.leaves(abs_state)[0].sharding.is_equivalent_to(\n NamedSharding(mesh, P(None, 'model')), ndim=2)\n\n @parameterized.parameters('auto', 'explicit', 'mixed')\n def test_sharding_axis_types(self, mode):\n if mode == 'auto':\n axis_types = (jax.sharding.AxisType.Auto, jax.sharding.AxisType.Auto)\n elif mode == 'explicit':\n axis_types = (jax.sharding.AxisType.Explicit, jax.sharding.AxisType.Explicit)\n else:\n axis_types = (jax.sharding.AxisType.Auto, jax.sharding.AxisType.Explicit)\n\n mesh = jax.make_mesh(\n (2, 2),\n ('row', 'col'),\n axis_types=axis_types,\n )\n if mode == 'mixed':\n with self.assertRaises(ValueError):\n nnx.Variable(\n jnp.ones((4, 4)),\n out_sharding=('row', 'col'),\n mesh=mesh,\n )\n else:\n v = nnx.Variable(\n jnp.ones((4, 4)),\n out_sharding=('row', 'col'),\n mesh=mesh,\n )\n self.assertEqual(v.sharding.mesh, mesh)\n self.assertEqual(v.sharding.spec, P('row', 'col'))\n\n def test_eval_shape_with_explicit_sharding(self):\n axis_types = (jax.sharding.AxisType.Explicit, jax.sharding.AxisType.Explicit)\n mesh1 = jax.make_mesh((2, 2), (\"a\", \"b\"), axis_types)\n class Model(nnx.Module):\n def __init__(self):\n self.p1 = nnx.Param(\n reshard(jnp.ones((4,4)), NamedSharding(mesh1, P('a', 'b'))),\n mesh=mesh1)\n\n abs_model = nnx.eval_shape(lambda: Model())\n self.assertEqual(abs_model.p1.sharding.spec, P('a', 'b'))\n\n def test_eval_shape_with_sharding0(self):\n # based on https://github.com/google/flax/issues/5110\n mesh1 = jax.make_mesh((2, 2), (\"a\", \"b\"), (jax.sharding.AxisType.Auto, jax.sharding.AxisType.Auto))\n mesh2 = jax.make_mesh((1, 4), (\"c\", \"d\"), (jax.sharding.AxisType.Auto, jax.sharding.AxisType.Auto))\n\n class Model(nnx.Module):\n def __init__(self):\n self.p1 = nnx.Linear(16, 16, rngs=nnx.Rngs(0), kernel_metadata={\"out_sharding\": (\"a\", \"b\"), \"mesh\": mesh1})\n self.p2 = nnx.Linear(16, 16, rngs=nnx.Rngs(0), kernel_metadata={\"out_sharding\": (\"c\", \"d\"), \"mesh\": mesh2})\n\n abs_model = nnx.eval_shape(lambda: Model())\n assert isinstance(abs_model.p1.kernel.sharding, jax.sharding.NamedSharding)\n assert abs_model.p1.kernel.sharding.mesh.axis_names == mesh1.axis_names\n assert abs_model.p1.kernel.sharding.spec == jax.P(\"a\", \"b\")\n assert isinstance(abs_model.p2.kernel.sharding, jax.sharding.NamedSharding)\n assert abs_model.p2.kernel.sharding.mesh.axis_names == mesh2.axis_names\n assert abs_model.p2.kernel.sharding.spec == jax.P(\"c\", \"d\")\n\n def test_eval_shape_with_sharding1(self):\n class Model(nnx.Module):\n def __init__(self):\n self.linear = nnx.Linear(10, 10, rngs=nnx.Rngs(0), kernel_metadata={\"out_sharding\": (\"a\", \"b\")})\n\n mesh = jax.make_mesh((2, 2), (\"a\", \"b\"), (jax.sharding.AxisType.Auto, jax.sharding.AxisType.Auto))\n with jax.set_mesh(mesh):\n abs_model = nnx.eval_shape(lambda: Model())\n assert isinstance(abs_model.linear.kernel.sharding, jax.sharding.NamedSharding)\n assert abs_model.linear.kernel.sharding.mesh.axis_names == mesh.axis_names\n assert abs_model.linear.kernel.sharding.spec == jax.P(\"a\", \"b\")\n\n @parameterized.product(axis_type_name=['auto', 'explicit'])\n def test_variable_out_sharding_types(self, axis_type_name):\n if axis_type_name == 'auto':\n axis_types = (jax.sharding.AxisType.Auto, jax.sharding.AxisType.Auto)\n else: # 'explicit'\n axis_types = (jax.sharding.AxisType.Explicit, jax.sharding.AxisType.Explicit)\n\n mesh = jax.make_mesh(\n (2, 2),\n ('data', 'model'),\n axis_types=axis_types,\n )\n\n with jax.set_mesh(mesh):\n value = jnp.ones((4, 4))\n\n # Test with PartitionSpec\n v_pspec = nnx.Variable(value, out_sharding=P('data', 'model'))\n self.assertEqual(v_pspec.sharding.spec, P('data', 'model'))\n\n # Test with NamedSharding\n ns = NamedSharding(mesh, P('data', None))\n v_namedsharding = nnx.Variable(value, out_sharding=ns)\n self.assertEqual(v_namedsharding.sharding, ns)\n\n # Test with Format\n if axis_type_name == 'auto':\n v_format = nnx.Variable(value, out_sharding=Format(Layout(major_to_minor=(1, 0)), ns))\n self.assertEqual(v_format.sharding, ns)\n\n def test_get_abstract_with_abstract_mesh(self):\n mesh = jax.make_mesh(\n (2, 2),\n ('a', 'b'),\n axis_types=(jax.sharding.AxisType.Auto,) * 2,\n )\n with jax.set_mesh(mesh):\n abs_model = nnx.eval_shape(\n lambda: nnx.Linear(\n 4,\n 8,\n rngs=nnx.Rngs(0),\n kernel_metadata={'out_sharding': ('a', 'b')},\n )\n )\n abs_model = nnx.abstract_with_sharding(abs_model)\n\n self.assertIsInstance(abs_model.kernel, nnx.Param)\n self.assertEqual(abs_model.kernel.sharding.spec, P('a', 'b'))\n self.assertEqual(\n abs_model.kernel.sharding.mesh.axis_names,\n mesh.axis_names,\n )\n\n def test_get_abstract_with_per_variable_mesh(self):\n mesh1 = jax.make_mesh(\n (2, 2),\n ('a', 'b'),\n axis_types=(jax.sharding.AxisType.Auto,) * 2,\n )\n mesh2 = jax.make_mesh(\n (1, 4),\n ('c', 'd'),\n axis_types=(jax.sharding.AxisType.Auto,) * 2,\n )\n\n class Model(nnx.Module):\n def __init__(self):\n self.p1 = nnx.Linear(\n 4,\n 8,\n rngs=nnx.Rngs(0),\n kernel_metadata={'out_sharding': ('a', 'b'), 'mesh': mesh1},\n )\n self.p2 = nnx.Linear(\n 4,\n 8,\n rngs=nnx.Rngs(0),\n kernel_metadata={'out_sharding': ('c', 'd'), 'mesh': mesh2},\n )\n\n abs_model = nnx.eval_shape(lambda: Model())\n abs_model = nnx.abstract_with_sharding(abs_model)\n\n self.assertEqual(abs_model.p1.kernel.sharding.spec, P('a', 'b'))\n self.assertEqual(abs_model.p1.kernel.sharding.mesh, mesh1)\n self.assertEqual(abs_model.p2.kernel.sharding.spec, P('c', 'd'))\n self.assertEqual(abs_model.p2.kernel.sharding.mesh, mesh2)\n\n def test_get_abstract_no_sharding_metadata(self):\n abs_model = nnx.eval_shape(lambda: nnx.Linear(4, 8, rngs=nnx.Rngs(0)))\n abs_model = nnx.abstract_with_sharding(abs_model)\n\n self.assertIsInstance(abs_model.kernel, nnx.Param)\n self.assertIsNone(\n getattr(abs_model.kernel.get_value(), 'sharding', None)\n )\n\ndef has_sharding_spec(array):\n sharding = array.sharding\n if hasattr(sharding, 'spec'):\n # For NamedSharding or PositionalSharding\n return sharding.spec is not None and any(\n s is not None for s in sharding.spec\n )\n return False\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/nnx/state_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom absl.testing import absltest\n\nfrom flax import nnx\nimport jax\nfrom jax import numpy as jnp\n\n\nclass StateTest(absltest.TestCase):\n def test_create_state(self):\n state = nnx.State(\n {'a': nnx.Param(jnp.array(1)), 'b': {'c': nnx.Param(jnp.array(2))}}\n )\n\n assert state['a'][...] == 1\n assert state['b']['c'][...] == 2\n\n def test_get_attr(self):\n state = nnx.State(\n {'a': nnx.Param(jnp.array(1)), 'b': {'c': nnx.Param(jnp.array(2))}}\n )\n\n assert state.a[...] == 1\n assert state.b.c[...] == 2\n\n def test_set_attr(self):\n state = nnx.State(\n {'a': nnx.Param(jnp.array(1)), 'b': {'c': nnx.Param(jnp.array(2))}}\n )\n\n state.a[...] = 3\n state.b.c[...] = 4\n\n assert state['a'][...] == 3\n assert state['b']['c'][...] == 4\n\n def test_set_attr_variables(self):\n state = nnx.State(\n {'a': nnx.Param(jnp.array(1)), 'b': {'c': nnx.Param(jnp.array(2))}}\n )\n\n state.a[...] = 3\n state.b.c[...] = 4\n\n assert isinstance(state.a, nnx.Param)\n assert state.a[...] == 3\n assert isinstance(state.b.c, nnx.Param)\n assert state.b.c[...] == 4\n\n def test_add_nested_attr(self):\n state = nnx.State(\n {'a': nnx.Param(jnp.array(1)), 'b': {'c': nnx.Param(jnp.array(2))}}\n )\n state.b.d = nnx.Param(jnp.array(5))\n\n assert state['b']['d'][...] == 5\n\n def test_delete_nested_attr(self):\n state = nnx.State(\n {'a': nnx.Param(jnp.array(1)), 'b': {'c': nnx.Param(jnp.array(2))}}\n )\n del state['b']['c']\n\n assert 'c' not in state['b']\n\n def test_integer_access(self):\n class Foo(nnx.Module):\n def __init__(self, *, rngs: nnx.Rngs):\n self.layers = nnx.List([\n nnx.Linear(1, 2, rngs=rngs),\n nnx.Linear(2, 3, rngs=rngs)\n ])\n\n module = Foo(rngs=nnx.Rngs(0))\n state = nnx.state(module)\n\n assert module.layers[0].kernel.shape == (1, 2)\n assert state.layers[0].kernel.shape == (1, 2)\n assert module.layers[1].kernel.shape == (2, 3)\n assert state.layers[1].kernel.shape == (2, 3)\n\n def test_pure_dict(self):\n module = nnx.Linear(4, 5, rngs=nnx.Rngs(0))\n state = nnx.state(module)\n pure_dict = nnx.to_pure_dict(state)\n assert isinstance(pure_dict, dict)\n assert isinstance(pure_dict['kernel'], jax.Array)\n assert isinstance(pure_dict['bias'], jax.Array)\n nnx.replace_by_pure_dict(state, jax.tree.map(jnp.zeros_like, pure_dict))\n assert isinstance(state, nnx.State)\n assert isinstance(state['kernel'], nnx.Variable)\n assert jnp.array_equal(state['kernel'][...], jnp.zeros((4, 5)))\n assert type(state['kernel']) == nnx.Param\n nnx.update(module, state)\n assert jnp.array_equal(module(jnp.ones((3, 4))), jnp.zeros((3, 5)))\n\n def test_diff(self):\n class MLPs(nnx.Module):\n def __init__(self, dim, rngs: nnx.Rngs, n=4):\n self.layers = nnx.List()\n for _ in range(n):\n self.layers.append(nnx.Linear(dim, dim, rngs=rngs, use_bias=False))\n\n def __call__(self, x):\n for layer in self.layers:\n x = layer(x)\n return x\n\n model1 = MLPs(4, rngs=nnx.Rngs(0), n=4)\n model2 = MLPs(4, rngs=nnx.Rngs(1), n=4)\n model3 = MLPs(4, rngs=nnx.Rngs(1), n=5)\n\n self.assertEqual(\n nnx.statelib.diff(nnx.state(model2), nnx.state(model1)),\n nnx.state({})\n )\n self.assertNotEqual(\n nnx.statelib.diff(nnx.state(model3), nnx.state(model1)),\n nnx.state({})\n )\n self.assertEqual(\n nnx.statelib.diff(nnx.state(model1), nnx.state(model3)),\n nnx.state({})\n )\n\n\nif __name__ == '__main__':\n absltest.main()" }, { "path": "tests/nnx/summary_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport jax.numpy as jnp\nfrom absl.testing import absltest\n\nfrom flax import nnx\n\nCONSOLE_TEST_KWARGS = dict(force_terminal=False, no_color=True, width=10_000)\n\n\nclass SummaryTest(absltest.TestCase):\n def test_tabulate(self):\n class Block(nnx.Module):\n def __init__(self, din, dout, rngs: nnx.Rngs):\n self.linear = nnx.Linear(din, dout, rngs=rngs)\n self.bn = nnx.BatchNorm(dout, rngs=rngs)\n self.dropout = nnx.Dropout(0.2, rngs=rngs)\n\n def forward(self, x):\n return nnx.relu(self.dropout(self.bn(self.linear(x))))\n\n class Foo(nnx.Module):\n def __init__(self, rngs: nnx.Rngs):\n self.block1 = Block(32, 128, rngs=rngs)\n self.block2 = Block(128, 10, rngs=rngs)\n\n def __call__(self, x):\n return self.block2.forward(self.block1.forward(x))\n\n foo = Foo(nnx.Rngs(0))\n x = jnp.ones((1, 32))\n table_repr_ = nnx.tabulate(\n foo, x, console_kwargs=CONSOLE_TEST_KWARGS\n )\n table_repr = table_repr_.splitlines()\n\n self.assertIn('Foo Summary', table_repr[0])\n self.assertIn('path', table_repr[2])\n self.assertIn('type', table_repr[2])\n self.assertIn('BatchStat', table_repr[2])\n self.assertIn('Param', table_repr[2])\n self.assertIn('block1/forward', table_repr[6])\n self.assertIn('Block', table_repr[6])\n self.assertIn('block1/linear', table_repr[8])\n self.assertIn('Linear', table_repr[8])\n self.assertIn('block1/bn', table_repr[13])\n self.assertIn('BatchNorm', table_repr[13])\n self.assertIn('block1/dropout', table_repr[18])\n self.assertIn('Dropout', table_repr[18])\n self.assertIn('block2/forward', table_repr[20])\n self.assertIn('Block', table_repr[20])\n self.assertIn('block2/linear', table_repr[22])\n self.assertIn('Linear', table_repr[22])\n self.assertIn('block2/bn', table_repr[27])\n self.assertIn('BatchNorm', table_repr[27])\n self.assertIn('block2/dropout', table_repr[32])\n self.assertIn('Dropout', table_repr[32])\n\n self.assertIn('Total', table_repr[34])\n self.assertIn('276 (1.1 KB)', table_repr[34])\n self.assertIn('5,790 (23.2 KB)', table_repr[34])\n self.assertIn('4 (24 B)', table_repr[34])\n self.assertIn('Total Parameters: 6,070 (24.3 KB)', table_repr[37])\n\n def test_multiple_inputs_and_outputs(self):\n class CustomMLP(nnx.Module):\n def __init__(self):\n self.weight = nnx.Param(jnp.ones((4, 8)))\n self.bias = nnx.Param(jnp.ones(8))\n\n def __call__(self, x, x2):\n y = x @ self.weight\n y += self.bias[None]\n y += x2\n return x, y, 2 * y\n\n cmlp = CustomMLP()\n x = jnp.ones((1, 4))\n x2 = jnp.ones((1, 8))\n table_repr = nnx.tabulate(\n cmlp, x, x2, console_kwargs=CONSOLE_TEST_KWARGS\n ).splitlines()\n\n self.assertIn('CustomMLP Summary', table_repr[0])\n self.assertIn('float32[1,4]', table_repr[4])\n self.assertIn('float32[1,8]', table_repr[5])\n self.assertIn('float32[1,8]', table_repr[6])\n\n def test_tabulate_empty_dict_first_arg(self):\n class Model(nnx.Module):\n def subroutine(self, foo, x):\n return x\n\n def __call__(self, x):\n return self.subroutine({}, x)\n\n model = Model()\n out = nnx.tabulate(\n model, jnp.zeros((1, 8)), depth=1, console_kwargs=CONSOLE_TEST_KWARGS\n )\n # Ensure empty dict argument is preserved and array input is shown\n self.assertIn('{}', out)\n self.assertIn('float32[1,8]', out)\n\n def test_tabulate_empty_dict_last_arg(self):\n class Model(nnx.Module):\n def subroutine(self, foo, x):\n return x\n\n def __call__(self, x):\n return self.subroutine(x, {})\n\n model = Model()\n out = nnx.tabulate(\n model, jnp.zeros((1, 8)), depth=1, console_kwargs=CONSOLE_TEST_KWARGS\n )\n # Ensure trailing empty dict is not dropped\n self.assertIn('{}', out)\n\n def test_tabulate_empty_dict_and_none_kwarg(self):\n class Model(nnx.Module):\n def subroutine(self, x, *, foo=None):\n return x\n\n def __call__(self, x):\n # One call with empty dict, one with None\n _ = self.subroutine(x, foo={})\n return self.subroutine(x, foo=None)\n\n model = Model()\n out = nnx.tabulate(\n model, jnp.zeros((1, 8)), depth=2, console_kwargs=CONSOLE_TEST_KWARGS\n )\n # Distinguish {} and None in output\n self.assertIn('{}', out)\n self.assertIn('None', out)\n\n def test_tabulate_empty_dict_property(self):\n class Model(nnx.Module):\n def __init__(self):\n self.foo = {}\n\n def subroutine(self, foo, x):\n return x\n\n def __call__(self, x):\n return self.subroutine(self.foo, x)\n\n model = Model()\n out = nnx.tabulate(\n model, jnp.zeros((1, 1024)), depth=1, console_kwargs=CONSOLE_TEST_KWARGS\n )\n # Should not crash and should show the empty dict argument\n self.assertIn('{}', out)\n\n\n def test_no_dup_flops(self):\n class Model(nnx.Module):\n def g(self, x):\n return x**2\n def __call__(self, x):\n return self.g(x)\n m = Model()\n x = jnp.ones(4)\n table_rep = nnx.tabulate(m, x, compute_flops=True)\n table_lines = table_rep.splitlines()\n self.assertEqual(sum(\" g \" in l for l in table_lines), 1)\n\n\n def test_flops(self):\n class Model(nnx.Module):\n def __init__(self):\n self.weight = nnx.Param(jnp.ones(4))\n\n def __call__(self, x1):\n return jnp.sum((x1 * self.weight)**2)\n m = Model()\n x = jnp.ones(4)\n table_repr1 = nnx.tabulate(\n m, x, compute_flops=True\n ).splitlines()\n self.assertIn('flops', table_repr1[2])\n self.assertNotIn('vjp_flops', table_repr1[2])\n table_repr2 = nnx.tabulate(\n m, x, compute_flops=True, compute_vjp_flops=True\n ).splitlines()\n self.assertIn('vjp_flops', table_repr2[2])\n\n def test_nested(self):\n class Block(nnx.Module):\n def __init__(self, rngs):\n self.linear = nnx.Linear(2, 2, rngs=rngs)\n self.bn = nnx.BatchNorm(2, rngs=rngs)\n\n def __call__(self, x):\n x = self.linear(x)\n x = self.bn(x)\n return nnx.relu(x)\n\n class Model(nnx.Module):\n def __init__(self, rngs):\n self.block1 = Block(rngs)\n self.block2 = Block(rngs)\n\n def __call__(self, x):\n x = self.block1(x)\n x = self.block2(x)\n return x\n\n m = Model(nnx.Rngs(0))\n x = jnp.ones((4, 2))\n table = nnx.tabulate(m, x, compute_flops=True, compute_vjp_flops=True)\n # We should see 3 calls per block, plus one overall call\n self.assertEqual(sum([s.startswith(\"├─\") for s in table.splitlines()]), 7)\n\n def test_time_complexity(self):\n counter = []\n\n class Block(nnx.Module):\n def __init__(self, rngs):\n self.linear = nnx.Linear(2, 2, rngs=rngs)\n\n def __call__(self, x):\n counter.append(1)\n return self.linear(x)\n\n class Model(nnx.Module):\n def __init__(self, rngs):\n for d in range(10):\n setattr(self, f\"linear{d}\", Block(rngs))\n\n def __call__(self, x):\n for d in range(10):\n x = getattr(self, f\"linear{d}\")(x)\n return x\n\n m = Model(nnx.Rngs(0))\n x = jnp.ones((4, 2))\n nnx.tabulate(m, x, compute_flops=True, compute_vjp_flops=False)\n self.assertEqual(len(counter), 10)\n\n def test_shared(self):\n class Block(nnx.Module):\n def __init__(self, linear: nnx.Linear, *, rngs):\n self.linear = linear\n self.bn = nnx.BatchNorm(2, rngs=rngs)\n\n def __call__(self, x):\n x = self.linear(x)\n x = self.bn(x)\n return nnx.relu(x)\n\n class Model(nnx.Module):\n def __init__(self, rngs):\n shared = nnx.Linear(2, 2, rngs=rngs)\n self.block1 = Block(shared, rngs=rngs)\n self.block2 = Block(shared, rngs=rngs)\n\n def __call__(self, x):\n x = self.block1(x)\n x = self.block2(x)\n return x\n\n m = Model(nnx.Rngs(0))\n x = jnp.ones((4, 2))\n table = nnx.tabulate(m, x, compute_vjp_flops=True)\n # We should see 3 calls per block, plus one overall call, minus the shared call\n self.assertEqual(sum([s.startswith(\"├─\") for s in table.splitlines()]), 6)\n\n def test_tabulate_with_variable_hooks(self):\n \"\"\"Test that tabulate works with Variables implementing hooks and custom metadata.\"\"\"\n\n class Custom:\n def __repr__(self):\n return \"\"\n\n class VarWithHooks(nnx.Variable):\n def on_get_value(self, value):\n return value\n\n def on_set_value(self, value):\n return value + 1.0\n\n class Model(nnx.Module):\n def __init__(self):\n # Variable with hooks\n self.hooked_param = VarWithHooks(value=jnp.ones((2, 3)))\n self.hooked_param.set_metadata('description', 'Custom parameter')\n self.hooked_param.set_metadata('trainable', True)\n\n # Variable with custom non-serializable metadata\n self.custom_param = nnx.Param(jnp.ones((2, 2)))\n self.custom_param.set_metadata('custom_obj', Custom())\n\n def __call__(self, x):\n return jnp.dot(x, self.hooked_param[...]) + self.custom_param.sum()\n\n module = Model()\n # Should not raise yaml.representer.RepresenterError\n table_repr = nnx.tabulate(module, jnp.ones((1, 2)), console_kwargs=CONSOLE_TEST_KWARGS)\n self.assertIsNotNone(table_repr)\n\n # Verify table contains expected content\n self.assertIn('Model Summary', table_repr)\n self.assertIn('hooked_param', table_repr)\n self.assertIn('on_set_value', table_repr)\n self.assertIn('', table_repr)\n\n # Verify metadata is preserved in the module\n self.assertEqual(module.hooked_param.get_metadata('description'), 'Custom parameter')\n self.assertEqual(module.hooked_param.get_metadata('trainable'), True)\n\n def test_tabulate_concrete_shape(self):\n class Net(nnx.Module):\n def __init__(self):\n self.rngs = nnx.Rngs(0)\n\n def __call__(self, x):\n return self.rngs.uniform((x.shape[0], 10))\n\n net = Net()\n x = jnp.zeros((4, 8))\n nnx.tabulate(net, x, console_kwargs={\"width\": 200})\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/nnx/test_traversals.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for flax.nnx.traversal.\"\"\"\nfrom absl.testing import absltest\nfrom flax.core import freeze\nfrom flax.nnx import traversals\nimport jax\n\n# Parse absl flags test_srcdir and test_tmpdir.\njax.config.parse_flags_with_absl()\n\n\nclass TraversalTest(absltest.TestCase):\n def test_flatten_mapping(self):\n xs = {'foo': 1, 'bar': {'a': 2, 'b': {}}}\n flat_xs = traversals.flatten_mapping(xs)\n self.assertEqual(\n flat_xs,\n {\n ('foo',): 1,\n ('bar', 'a'): 2,\n },\n )\n flat_xs = traversals.flatten_mapping(freeze(xs))\n self.assertEqual(\n flat_xs,\n {\n ('foo',): 1,\n ('bar', 'a'): 2,\n },\n )\n flat_xs = traversals.flatten_mapping(xs, sep='/')\n self.assertEqual(\n flat_xs,\n {\n 'foo': 1,\n 'bar/a': 2,\n },\n )\n\n def test_unflatten_mapping(self):\n expected_xs = {'foo': 1, 'bar': {'a': 2}}\n xs = traversals.unflatten_mapping(\n {\n ('foo',): 1,\n ('bar', 'a'): 2,\n }\n )\n self.assertEqual(xs, expected_xs)\n xs = traversals.unflatten_mapping(\n {\n 'foo': 1,\n 'bar/a': 2,\n },\n sep='/',\n )\n self.assertEqual(xs, expected_xs)\n\n def test_flatten_mapping_keep_empty(self):\n ys = {'a': {}}\n xs = {'foo': 1, 'bar': {'a': 2, 'b': {}}}\n flat_ys = traversals.flatten_mapping(ys, keep_empty_nodes=True)\n flat_xs = traversals.flatten_mapping(xs, keep_empty_nodes=True)\n empty_node = flat_ys[('a',)]\n self.assertEqual(\n flat_xs,\n {\n ('foo',): 1,\n ('bar', 'a'): 2,\n ('bar', 'b'): empty_node,\n },\n )\n xs_restore = traversals.unflatten_mapping(flat_xs)\n self.assertEqual(xs, xs_restore)\n\n def test_flatten_mapping_is_leaf(self):\n xs = {'foo': {'c': 4}, 'bar': {'a': 2, 'b': {}}}\n flat_xs = traversals.flatten_mapping(\n xs, is_leaf=lambda k, x: len(k) == 1 and len(x) == 2\n )\n self.assertEqual(\n flat_xs,\n {\n ('foo', 'c'): 4,\n ('bar',): {'a': 2, 'b': {}},\n },\n )\n xs_restore = traversals.unflatten_mapping(flat_xs)\n self.assertEqual(xs, xs_restore)\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/nnx/transforms_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nos.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=4'\n\nimport dataclasses\nfrom functools import partial\nimport typing as tp\n\nfrom absl.testing import absltest\nfrom absl.testing import parameterized\nfrom flax import nnx\nfrom flax.nnx.transforms.iteration import pure_jax_fancy_scan\nfrom flax.nnx.transforms import general\nimport jax\nfrom jax.experimental import checkify, mesh_utils\nimport jax.numpy as jnp\nimport numpy as np\nimport optax\nfrom flax import errors\n\n\nclass TestJIT(parameterized.TestCase):\n def test_jit(self):\n m = nnx.Dict(a=nnx.Param(1))\n\n @nnx.jit\n def g(m: nnx.Dict):\n m.a = 2\n return 1.0\n\n out = g(m)\n\n assert m.a == 2\n assert out == 1.0\n\n def test_mutable_array_input_output(self):\n m = jax.new_ref(jnp.array(1.0))\n\n @nnx.jit\n def f(m: jax.Ref):\n m[...] += 1.0\n m2 = jax.new_ref(jnp.array(10.0))\n return m2, m\n\n m2, m_out = f(m)\n\n self.assertEqual(m[...], 2.0)\n self.assertIs(m, m_out)\n self.assertIsInstance(m2, jax.Ref)\n\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_simple_double_call(self, graph_mode, graph_updates):\n n = 0\n m = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n\n @nnx.jit(graph=graph_mode, graph_updates=graph_updates)\n def f(m: nnx.Linear, x: jnp.ndarray) -> jnp.ndarray:\n nonlocal n\n n += 1\n return m(x)\n\n x = jnp.ones((1, 2))\n y = f(m, x)\n\n self.assertEqual(n, 1)\n self.assertEqual(y.shape, (1, 3))\n\n y = f(m, x)\n\n self.assertEqual(n, 1)\n self.assertEqual(y.shape, (1, 3))\n\n def test_jit_on_init(self):\n n = 0\n\n class Foo(nnx.Module):\n @nnx.jit(static_argnums=(1, 2))\n def __init__(self, din: int, dout: int, *, rngs: nnx.Rngs):\n nonlocal n\n n += 1\n\n key = rngs.params()\n self.w = nnx.Param(jax.random.normal(key, shape=(din, dout)))\n self.din = din\n self.dout = dout\n\n m = Foo(2, 3, rngs=nnx.Rngs(0))\n assert n == 1\n assert m.w.shape == (2, 3)\n assert m.din == 2\n assert m.dout == 3\n assert isinstance(m.din, int)\n assert isinstance(m.dout, int)\n assert isinstance(m.w[...], jax.Array)\n\n m = Foo(2, 3, rngs=nnx.Rngs(0))\n assert n == 1\n\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_jit_on_call(self, graph_mode, graph_updates):\n n = 0\n\n class Foo(nnx.Module):\n def __init__(self, din: int, dout: int, *, rngs: nnx.Rngs):\n key = rngs.params()\n self.w = nnx.Param(jax.random.normal(key, shape=(din, dout)))\n self.din = din\n self.dout = dout\n\n @nnx.jit(graph=graph_mode, graph_updates=graph_updates)\n def __call__(self, x: jax.Array) -> jax.Array:\n nonlocal n\n n += 1\n return jnp.dot(x, self.w)\n\n m = Foo(2, 3, rngs=nnx.Rngs(0))\n assert m.w.shape == (2, 3)\n assert m.din == 2\n assert m.dout == 3\n assert isinstance(m.din, int)\n assert isinstance(m.dout, int)\n assert isinstance(m.w[...], jax.Array)\n\n y = m(jnp.ones((1, 2)))\n assert y.shape == (1, 3)\n assert n == 1\n y = m(jnp.ones((1, 2)))\n assert n == 1\n\n def test_cached_unflatten(self):\n n = 0\n\n class Foo(nnx.Module):\n def __init__(self, *, rngs: nnx.Rngs):\n self.a = nnx.Linear(2, 2, rngs=rngs)\n self.b = nnx.BatchNorm(2, rngs=rngs)\n\n @nnx.jit\n def f(m: Foo):\n nonlocal n\n n += 1\n m.a, m.b = m.b, m.a # type: ignore\n\n m = Foo(rngs=nnx.Rngs(0))\n a = m.a\n b = m.b\n a_kernel = a.kernel[...]\n a_bias = a.bias[...]\n b_scale = b.scale[...]\n b_bias = b.bias[...]\n b_mean = b.mean[...]\n b_var = b.var[...]\n\n f(m)\n\n assert n == 1\n assert m.a is b\n assert m.b is a\n np.testing.assert_allclose(a_kernel, a.kernel[...])\n np.testing.assert_allclose(a_bias, a.bias[...])\n np.testing.assert_allclose(b_scale, b.scale[...])\n np.testing.assert_allclose(b_bias, b.bias[...])\n np.testing.assert_allclose(b_mean, b.mean[...])\n np.testing.assert_allclose(b_var, b.var[...])\n\n f(m)\n\n assert n == 2\n assert m.a is a\n assert m.b is b\n\n f(m)\n\n assert n == 2\n assert m.a is b\n assert m.b is a\n\n f(m)\n\n assert n == 2\n assert m.a is a\n assert m.b is b\n\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_jit_custom_vjp(self, graph_mode, graph_updates):\n @nnx.custom_vjp(graph=graph_mode, graph_updates=graph_updates)\n def f(x, y):\n return jnp.sin(x) * y\n\n def f_fwd(x, y):\n return f(x, y), (jnp.cos(x), jnp.sin(x), y)\n\n def f_bwd(res, g):\n cos_x, sin_x, y = res\n return (cos_x * g * y, sin_x * g)\n\n f.defvjp(f_fwd, f_bwd)\n\n nnx_out = nnx.jit(f, graph=graph_mode, graph_updates=graph_updates)(jnp.array([1.0, 2.0]), jnp.array([3.0, 4.0]))\n jax_out = jax.jit(f)(jnp.array([1.0, 2.0]), jnp.array([3.0, 4.0]))\n assert (nnx_out == jax_out).all()\n\n def test_cached_unflatten_same_type(self):\n n = 0\n\n class Foo(nnx.Module):\n def __init__(self, *, rngs: nnx.Rngs):\n self.a = nnx.Linear(2, 2, rngs=rngs)\n self.b = nnx.Linear(2, 2, rngs=rngs)\n\n @nnx.jit\n def f(m: Foo):\n nonlocal n\n n += 1\n m.a, m.b = m.b, m.a\n\n m = Foo(rngs=nnx.Rngs(0))\n a = m.a\n b = m.b\n\n f(m)\n\n assert n == 1\n assert m.a is b\n assert m.b is a\n\n f(m)\n\n assert n == 1\n assert m.a is a\n assert m.b is b\n\n def test_objects_in_pytree(self):\n n = 0\n\n class Foo(nnx.Module):\n def __init__(self, *, rngs: nnx.Rngs):\n self.a = nnx.Linear(2, 2, rngs=rngs)\n self.b = nnx.Linear(2, 2, rngs=rngs)\n\n class FooDict(tp.TypedDict):\n foo: Foo\n\n @nnx.jit\n def f(tree: tuple[FooDict]):\n nonlocal n\n n += 1\n m = tree[0]['foo']\n m.a, m.b = m.b, m.a\n\n m = Foo(rngs=nnx.Rngs(0))\n a = m.a\n b = m.b\n\n f(({'foo': m},))\n\n assert n == 1\n assert m.a is b\n assert m.b is a\n\n f(({'foo': m},))\n\n assert n == 1\n assert m.a is a\n assert m.b is b\n\n def test_cached_unflatten_swap_variables(self):\n class Foo(nnx.Module):\n def __init__(self):\n self.a = nnx.Param(1)\n self.b = nnx.Param(2)\n\n @nnx.jit\n def f(m: Foo):\n m.a, m.b = m.b, m.a\n\n m = Foo()\n a = m.a\n b = m.b\n\n f(m)\n\n assert m.a is b\n assert m.b is a\n\n def test_cached_unflatten_add_self_reference(self):\n n = 0\n\n class Foo(nnx.Module):\n def __init__(self):\n self.ref: tp.Optional[Foo] = nnx.data(None) # type: ignore[name-error]\n\n @nnx.jit\n def f(m: Foo):\n nonlocal n\n n += 1\n m.ref = m\n\n m = Foo()\n\n f(m)\n\n assert n == 1\n assert m.ref is m\n\n f(m)\n\n assert n == 2\n assert m.ref is m\n\n f(m)\n\n assert n == 2\n assert m.ref is m\n\n def test_cached_unflatten_ref_in_output(self):\n n = 0\n\n class Foo(nnx.Module):\n def __init__(self):\n self.ref: tp.Optional[Foo] = nnx.data(None) # type: ignore[name-error]\n\n @nnx.jit\n def f(m: Foo):\n nonlocal n\n n += 1\n m.ref = m\n return m\n\n m = Foo()\n\n m2 = f(m)\n\n assert n == 1\n assert m.ref is m\n assert m2 is m\n\n m2 = f(m)\n\n assert n == 2\n assert m.ref is m\n assert m2 is m\n\n m2 = f(m)\n\n assert n == 2\n assert m.ref is m\n assert m2 is m\n\n def test_apply_shardings(self):\n n_devices = max(jax.local_device_count() // 2, 1)\n devices = mesh_utils.create_device_mesh(\n (n_devices, jax.local_device_count() // n_devices)\n )\n mesh = jax.sharding.Mesh(devices, ('a', 'b'))\n\n def sharding(*args):\n return jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(*args))\n\n state_sharding = nnx.StateSharding(\n {\n nnx.PathContains('kernel'): sharding('a', 'b'),\n nnx.PathContains('bias'): sharding('b'),\n }\n )\n\n m = nnx.Linear(16, 32, rngs=nnx.Rngs(0))\n\n self.assertNotIsInstance(m.kernel.sharding, jax.sharding.NamedSharding)\n\n @nnx.jit(in_shardings=(state_sharding,))\n def constrain_object(m):\n pass\n\n constrain_object(m)\n\n self.assertIsInstance(m.kernel.sharding, jax.sharding.NamedSharding)\n\n def test_cache_args(self):\n m = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n\n @nnx.jit\n def f(cached_m: nnx.Linear, m: nnx.Linear):\n self.assertIsNot(cached_m, m)\n self.assertIs(cached_m.kernel, m.kernel)\n self.assertIs(cached_m.bias, m.bias)\n return cached_m\n\n cached_f = nnx.cached_partial(f, m)\n cached_m = cached_f(m)\n\n self.assertIsNot(m, cached_m)\n self.assertIs(m.kernel, cached_m.kernel)\n self.assertIs(m.bias, cached_m.bias)\n\n # test that cached m is reused\n cached_m2 = cached_f(m)\n self.assertIs(cached_m, cached_m2)\n\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_jit_wrapped(self, graph_mode, graph_updates):\n class Foo(nnx.Module):\n def __init__(self, *, rngs: nnx.Rngs):\n self.count = nnx.Variable(jnp.array(0))\n\n @nnx.jit(graph=graph_mode, graph_updates=graph_updates)\n def __call__(self, x: jax.Array) -> jax.Array:\n self.count[...] += 1\n return x * 2\n\n m = Foo(rngs=nnx.Rngs(0))\n x = jnp.array(3.0)\n\n @nnx.jit(graph=graph_mode, graph_updates=graph_updates)\n def f(m: nnx.Linear, x):\n return m(x)\n\n lowered = f.lower(m, x)\n compiled = lowered.compile()\n text = compiled.as_text()\n cost_analysis = compiled.cost_analysis()\n self.assertIsNotNone(cost_analysis)\n self.assertIsNotNone(text)\n\n y = compiled(m, x)\n np.testing.assert_allclose(y, 6.0)\n self.assertEqual(m.count[...], 1)\n y = compiled(m, x)\n self.assertEqual(m.count[...], 2)\n\n @parameterized.parameters(\n {'graph_mode': True, 'graph_updates': True, 'static_argnums': (2,), 'static_argnames': None},\n {'graph_mode': True, 'graph_updates': True, 'static_argnums': None, 'static_argnames': ('use_relu',)},\n {'graph_mode': True, 'graph_updates': False, 'static_argnums': (2,), 'static_argnames': None},\n {'graph_mode': True, 'graph_updates': False, 'static_argnums': None, 'static_argnames': ('use_relu',)},\n {'graph_mode': False, 'graph_updates': False, 'static_argnums': (2,), 'static_argnames': None},\n {'graph_mode': False, 'graph_updates': False, 'static_argnums': None, 'static_argnames': ('use_relu',)},\n )\n def test_jit_static_args_with_shardings(self, graph_mode, graph_updates, static_argnums, static_argnames):\n \"\"\"Test static arguments work correctly with in_shardings.\"\"\"\n n_devices = jax.local_device_count()\n devices = mesh_utils.create_device_mesh((n_devices,))\n mesh = jax.sharding.Mesh(devices, ('data',))\n\n def fn(x, scale, use_relu):\n y = x * scale\n if use_relu:\n y = jnp.maximum(y, 0)\n return y.sum()\n\n x = jnp.linspace(-1.0, 1.0, 16, dtype=jnp.float32).reshape(4, 4)\n x_sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec('data'))\n\n f = nnx.jit(fn, in_shardings=(x_sharding, None),\n static_argnums=static_argnums, static_argnames=static_argnames,\n graph=graph_mode, graph_updates=graph_updates)\n y_relu = f(x, 0.5, True)\n y_no_relu = f(x, 0.5, False)\n self.assertNotEqual(y_relu, y_no_relu)\n\n @parameterized.parameters(\n {\n 'static_args': {'static_argnums': (2, 3)},\n },\n {\n 'static_args': {'static_argnames': ('static_arg1', 'static_arg2')},\n },\n )\n def test_with_sharding_and_static_args(self, static_args):\n n_devices = max(jax.local_device_count() // 2, 1)\n devices = mesh_utils.create_device_mesh(\n (n_devices, jax.local_device_count() // n_devices)\n )\n mesh = jax.sharding.Mesh(devices, ('a', 'b'))\n\n def sharding(*args):\n return jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(*args))\n\n state_sharding = nnx.StateSharding(\n {\n nnx.PathContains('kernel'): sharding('a', 'b'),\n nnx.PathContains('bias'): sharding('b'),\n }\n )\n\n m = nnx.Linear(16, 32, rngs=nnx.Rngs(0))\n self.assertNotIsInstance(m.kernel.sharding, jax.sharding.NamedSharding)\n\n @nnx.jit(\n in_shardings=(state_sharding, None),\n **static_args,\n )\n def constrain_object(m, scale: float, static_arg1: bool, static_arg2: bool):\n new_sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec('b', 'a'))\n m.kernel = jax.lax.with_sharding_constraint(m.kernel, new_sharding)\n return None\n\n constrain_object(m, 0.5, True, True)\n self.assertEqual(m.kernel.sharding.spec, jax.sharding.PartitionSpec(\"a\", \"b\"))\n\n\nclass TestTreeJIT(parameterized.TestCase):\n @parameterized.parameters(\n (True, True),\n (True, False),\n (False, False),\n )\n def test_tree_jit_basic(self, graph, graph_updates):\n m = nnx.Dict(a=nnx.Param(jnp.array(1)))\n\n @nnx.jit(graph=graph, graph_updates=graph_updates)\n def g(m: nnx.Dict):\n m.a[...] = 2\n return 1.0\n\n out = g(m)\n\n assert m.a[...] == 2\n assert out == 1.0\n\n @parameterized.parameters(\n (True, True),\n (True, False),\n (False, False),\n )\n def test_tree_jit_module(self, graph, graph_updates):\n m = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n\n @nnx.jit(graph=graph, graph_updates=graph_updates)\n def f(m, x):\n return m(x)\n\n x = jnp.ones((1, 2))\n y = f(m, x)\n self.assertEqual(y.shape, (1, 3))\n\n def test_tree_jit_variable_update(self):\n class Foo(nnx.Module):\n def __init__(self):\n self.count = nnx.Variable(jnp.array(0))\n\n @nnx.jit(graph=False)\n def __call__(self, x):\n self.count[...] += 1\n return x * 2\n\n m = Foo()\n y = m(jnp.array(3.0))\n np.testing.assert_allclose(y, 6.0)\n self.assertEqual(m.count[...], 1)\n y = m(jnp.array(3.0))\n self.assertEqual(m.count[...], 2)\n\n @parameterized.parameters(\n (True, True),\n (True, False),\n (False, False),\n )\n def test_tree_jit_no_retrace(self, graph, graph_updates):\n n = 0\n m = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n\n @nnx.jit(graph=graph, graph_updates=graph_updates)\n def f(m, x):\n nonlocal n\n n += 1\n return m(x)\n\n x = jnp.ones((1, 2))\n y = f(m, x)\n self.assertEqual(n, 1)\n self.assertEqual(y.shape, (1, 3))\n\n y = f(m, x)\n self.assertEqual(n, 1)\n self.assertEqual(y.shape, (1, 3))\n\n def test_tree_jit_static_argnums(self):\n @nnx.jit(graph=False, static_argnums=(1,))\n def f(x, use_relu):\n if use_relu:\n return jnp.maximum(x, 0)\n return x\n\n x = jnp.array([-1.0, 2.0])\n y_relu = f(x, True)\n np.testing.assert_allclose(y_relu, jnp.array([0.0, 2.0]))\n y_no_relu = f(x, False)\n np.testing.assert_allclose(y_no_relu, x)\n\n def test_tree_jit_no_input_output_aliasing(self):\n v = nnx.Param(jnp.array(1.0))\n\n @nnx.jit(graph=False)\n def f(v):\n return v\n\n with self.assertRaisesRegex(ValueError, 'does not support returning input Variables as outputs'):\n f(v)\n\n def test_tree_jit_no_shared_variable_refs(self):\n v = nnx.Param(jnp.array(1.0))\n\n @nnx.jit(graph=False)\n def f(v1, v2):\n pass\n\n with self.assertRaisesRegex(\n ValueError, 'found at paths'\n ):\n f(v, v)\n\n def test_tree_jit_new_variable_output_ok(self):\n @nnx.jit(graph=False)\n def f(x):\n return nnx.Param(x + 1)\n\n v = f(jnp.array(1.0))\n self.assertIsInstance(v, nnx.Param)\n np.testing.assert_allclose(v[...], 2.0)\n\n def test_tree_jit_donate_argnums_unchanged_var(self):\n v = nnx.Param(jnp.array(1.0))\n\n @nnx.jit(graph=False, donate_argnums=(0,))\n def f(v):\n return v[...] + 1.0\n\n out = f(v)\n np.testing.assert_allclose(out, 2.0)\n np.testing.assert_allclose(v[...], 1.0)\n\n out = f(v)\n np.testing.assert_allclose(out, 2.0)\n np.testing.assert_allclose(v[...], 1.0)\n\n def test_tree_jit_donate_argnums_module(self):\n m = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n original_kernel = jnp.copy(m.kernel[...])\n\n @nnx.jit(graph=False, donate_argnums=(0,))\n def f(m, x):\n return m(x)\n\n x = jnp.ones((1, 2))\n y = f(m, x)\n self.assertEqual(y.shape, (1, 3))\n np.testing.assert_allclose(m.kernel[...], original_kernel)\n\n y = f(m, x)\n self.assertEqual(y.shape, (1, 3))\n np.testing.assert_allclose(m.kernel[...], original_kernel)\n\n def test_tree_jit_donate_argnums_with_mutation(self):\n v = nnx.Param(jnp.array(0.0))\n\n @nnx.jit(graph=False, donate_argnums=(0,))\n def f(v):\n v[...] += 1.0\n return None\n\n f(v)\n np.testing.assert_allclose(v[...], 1.0)\n f(v)\n np.testing.assert_allclose(v[...], 2.0)\n\n def test_tree_jit_donate_argnames(self):\n v = nnx.Param(jnp.array(1.0))\n\n @nnx.jit(graph=False, donate_argnames=('v',))\n def f(v):\n return v[...] + 1.0\n\n out = f(v=v)\n np.testing.assert_allclose(out, 2.0)\n np.testing.assert_allclose(v[...], 1.0)\n\n out = f(v=v)\n np.testing.assert_allclose(out, 2.0)\n np.testing.assert_allclose(v[...], 1.0)\n\n def test_tree_jit_donate_selective(self):\n donated = nnx.Param(jnp.array(1.0))\n not_donated = nnx.Param(jnp.array(2.0))\n\n @nnx.jit(graph=False, donate_argnums=(0,))\n def f(donated, not_donated):\n return donated[...] + not_donated[...]\n\n out = f(donated, not_donated)\n np.testing.assert_allclose(out, 3.0)\n np.testing.assert_allclose(donated[...], 1.0)\n np.testing.assert_allclose(not_donated[...], 2.0)\n\n out = f(donated, not_donated)\n np.testing.assert_allclose(out, 3.0)\n\n @parameterized.parameters(True, False)\n def test_jit_partial_basic(self, graph_mode):\n m = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n\n def f(m, x):\n return m(x)\n\n f_jit = nnx.jit_partial(f, m, graph=graph_mode, graph_updates=False)\n x = jnp.ones((1, 2))\n y = f_jit(x)\n self.assertEqual(y.shape, (1, 3))\n\n @parameterized.parameters(True, False)\n def test_jit_partial_lower_compile(self, graph_mode):\n class Foo(nnx.Module):\n def __init__(self):\n self.count = nnx.Variable(jnp.array(0))\n self.w = nnx.Param(jnp.ones((2, 3)))\n\n m = Foo()\n\n def f(m, x):\n m.count[...] += 1\n return x @ m.w[...]\n\n f_jit = nnx.jit_partial(f, m, graph=graph_mode, graph_updates=False)\n compiled = f_jit.lower(jnp.ones((1, 2))).compile()\n\n x = jnp.ones((1, 2))\n y = compiled(x)\n self.assertEqual(y.shape, (1, 3))\n np.testing.assert_allclose(y, jnp.ones((1, 3)) * 2)\n self.assertEqual(m.count[...], 1)\n\n y = compiled(x)\n self.assertEqual(m.count[...], 2)\n\n @parameterized.parameters(True, False)\n def test_jit_partial_variable_update(self, graph_mode):\n class Foo(nnx.Module):\n def __init__(self):\n self.count = nnx.Variable(jnp.array(0))\n\n m = Foo()\n\n def f(m, x):\n m.count[...] += 1\n return x * 2\n\n f_jit = nnx.jit_partial(f, m, graph=graph_mode, graph_updates=False)\n y = f_jit(jnp.array(3.0))\n np.testing.assert_allclose(y, 6.0)\n self.assertEqual(m.count[...], 1)\n y = f_jit(jnp.array(3.0))\n self.assertEqual(m.count[...], 2)\n\n @parameterized.parameters(True, False)\n def test_jit_partial_multiple_args(self, graph_mode):\n m1 = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n m2 = nnx.Linear(3, 4, rngs=nnx.Rngs(1))\n\n def f(m1, m2, x):\n return m2(m1(x))\n\n f_jit = nnx.jit_partial(f, m1, m2, graph=graph_mode, graph_updates=False)\n x = jnp.ones((1, 2))\n y = f_jit(x)\n self.assertEqual(y.shape, (1, 4))\n\n @parameterized.parameters(True, False)\n def test_jit_partial_no_retrace(self, graph_mode):\n n = 0\n m = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n\n def f(m, x):\n nonlocal n\n n += 1\n return m(x)\n\n f_jit = nnx.jit_partial(f, m, graph=graph_mode, graph_updates=False)\n x = jnp.ones((1, 2))\n f_jit(x)\n self.assertEqual(n, 1)\n f_jit(x)\n self.assertEqual(n, 1)\n\n @parameterized.parameters(True, False)\n def test_jit_partial_no_retrace_after_mutation(self, graph_mode):\n n = 0\n\n class Foo(nnx.Module):\n def __init__(self):\n self.w = nnx.Param(jnp.ones((2, 3)))\n\n m = Foo()\n\n def f(m, x):\n nonlocal n\n n += 1\n return x @ m.w[...]\n\n f_jit = nnx.jit_partial(f, m, graph=graph_mode, graph_updates=False)\n x = jnp.ones((1, 2))\n f_jit(x)\n self.assertEqual(n, 1)\n m.w[...] = jnp.zeros((2, 3))\n y = f_jit(x)\n self.assertEqual(n, 1)\n np.testing.assert_allclose(y, jnp.zeros((1, 3)))\n\n @parameterized.parameters(True, False)\n def test_jit_partial_no_partial_args(self, graph_mode):\n f_partial = nnx.jit_partial(lambda x: x * 2, graph=graph_mode, graph_updates=False)\n y = f_partial(jnp.array(3.0))\n np.testing.assert_allclose(y, 6.0)\n\n @parameterized.parameters((False,))\n def test_jit_partial_in_shardings_none_broadcast(self, graph_mode):\n n_devices = jax.local_device_count()\n devices = mesh_utils.create_device_mesh((n_devices,))\n mesh = jax.sharding.Mesh(devices, ('data',))\n\n m = nnx.Linear(4, 3, rngs=nnx.Rngs(0))\n\n def f(m, x):\n return m(x)\n\n f_jit = nnx.jit_partial(f, m, in_shardings=(None, None), graph=graph_mode, graph_updates=False)\n x = jnp.ones((n_devices, 4))\n y = f_jit(x)\n self.assertEqual(y.shape, (n_devices, 3))\n\n @parameterized.parameters((False,))\n def test_jit_partial_in_shardings_named(self, graph_mode):\n n_devices = jax.local_device_count()\n devices = mesh_utils.create_device_mesh((n_devices,))\n mesh = jax.sharding.Mesh(devices, ('data',))\n PS = jax.sharding.PartitionSpec\n\n v = nnx.Param(jnp.ones((n_devices, 4)))\n\n def f(v, x):\n return v[...] + x\n\n x_sharding = jax.sharding.NamedSharding(mesh, PS('data'))\n v_sharding = jax.sharding.NamedSharding(mesh, PS('data'))\n f_jit = nnx.jit_partial(\n f, v, in_shardings=(v_sharding, x_sharding), graph=graph_mode,\n graph_updates=False)\n x = jnp.ones((n_devices, 4))\n y = f_jit(x)\n self.assertEqual(y.shape, (n_devices, 4))\n\n @parameterized.parameters((False,))\n def test_jit_partial_mixed_shardings(self, graph_mode):\n n_devices = jax.local_device_count()\n devices = mesh_utils.create_device_mesh((n_devices,))\n mesh = jax.sharding.Mesh(devices, ('data',))\n PS = jax.sharding.PartitionSpec\n\n m1 = nnx.Linear(4, 3, rngs=nnx.Rngs(0))\n m2 = nnx.Linear(3, 2, rngs=nnx.Rngs(1))\n\n def f(m1, m2, x):\n return m2(m1(x))\n\n x_sharding = jax.sharding.NamedSharding(mesh, PS('data'))\n f_jit = nnx.jit_partial(\n f, m1, m2, in_shardings=(None, None, x_sharding), graph=graph_mode,\n graph_updates=False)\n x = jnp.ones((n_devices, 4))\n y = f_jit(x)\n self.assertEqual(y.shape, (n_devices, 2))\n\n @parameterized.parameters(True, False)\n def test_jit_partial_in_shardings_non_tuple(self, graph_mode):\n n_devices = jax.local_device_count()\n devices = mesh_utils.create_device_mesh((n_devices,))\n mesh = jax.sharding.Mesh(devices, ('data',))\n PS = jax.sharding.PartitionSpec\n\n m = nnx.Linear(4, 3, rngs=nnx.Rngs(0))\n\n def f(m, x):\n return m(x)\n\n sharding = jax.sharding.NamedSharding(mesh, PS())\n f_jit = nnx.jit_partial(f, m, in_shardings=sharding, graph=graph_mode, graph_updates=False)\n x = jnp.ones((n_devices, 4))\n y = f_jit(x)\n self.assertEqual(y.shape, (n_devices, 3))\n\n @parameterized.parameters(True, False)\n def test_jit_partial_train_step(self, graph_mode):\n model = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n optimizer = nnx.Optimizer(model, optax.adamw(1e-3), wrt=nnx.Param)\n\n def train_step(model, optimizer, x, y):\n def loss_fn(model):\n return jnp.mean((model(x) - y) ** 2)\n loss, grads = nnx.value_and_grad(loss_fn)(model)\n optimizer.update(model, grads)\n return loss\n\n train_step_fn = nnx.jit_partial(train_step, model, optimizer, graph=graph_mode, graph_updates=False)\n for _ in range(2):\n x, y = jnp.ones((10, 2)), jnp.ones((10, 3))\n loss = train_step_fn(x, y)\n self.assertIsInstance(loss, jax.Array)\n\n def test_jit_partial_shared_variable(self):\n v = nnx.Param(jnp.array(1.0))\n\n class Container(nnx.Module):\n def __init__(self, v):\n self.v = v\n\n c1 = Container(v)\n c2 = Container(v)\n\n def f(c1, c2, x):\n c1.v[...] += x\n return c1.v[...] + c2.v[...]\n\n f_jit = nnx.jit_partial(f, c1, c2, graph=True, graph_updates=False)\n y = f_jit(jnp.array(1.0))\n np.testing.assert_allclose(y, 4.0)\n np.testing.assert_allclose(v[...], 2.0)\n\n def test_jit_inconsistent_aliasing(self):\n v = nnx.Param(jnp.array(1.0))\n P = jax.sharding.PartitionSpec\n\n @nnx.jit(\n in_shardings=(P(), P('x')),\n graph=True, graph_updates=False,\n )\n def f(a, b):\n return a[...] + b[...]\n\n mesh = jax.sharding.Mesh(jax.devices(), ('x',))\n with mesh:\n with self.assertRaisesRegex(ValueError, 'Inconsistent aliasing'):\n f(v, v)\n\n\nclass TestEvalShape(parameterized.TestCase):\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_eval_shape(self, graph, graph_updates):\n abs_model = nnx.eval_shape(\n lambda: nnx.Linear(1, 2, rngs=nnx.Rngs(0)),\n graph=graph, graph_updates=graph_updates,\n )\n self.assertIsInstance(abs_model, nnx.Linear)\n self.assertIsInstance(abs_model.kernel.get_value(), jax.ShapeDtypeStruct)\n\n def test_eval_shape_mutable_array(self):\n with nnx.var_defaults(hijax=True):\n abs_model = nnx.eval_shape(lambda: nnx.Linear(1, 2, rngs=nnx.Rngs(0)))\n self.assertIsInstance(abs_model, nnx.Linear)\n self.assertIsInstance(abs_model.kernel.get_value(), jax.ShapeDtypeStruct)\n self.assertEqual(abs_model.kernel.shape, (1, 2))\n\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_eval_shape_with_module_input(self, graph, graph_updates):\n model = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n\n def f(m, x):\n return m(x)\n\n x = jnp.ones((4, 2))\n out = nnx.eval_shape(f, model, x, graph=graph, graph_updates=graph_updates)\n self.assertIsInstance(out, jax.ShapeDtypeStruct)\n self.assertEqual(out.shape, (4, 3))\n\n @parameterized.parameters(\n (True, False),\n (False, False),\n )\n def test_eval_shape_no_state_update(self, graph, graph_updates):\n count = nnx.Variable(jnp.array(0))\n\n def f(c):\n c[...] += 1\n return jnp.ones((3, 4)) * c[...]\n\n out = nnx.eval_shape(f, count, graph=graph, graph_updates=graph_updates)\n self.assertIsInstance(out, jax.ShapeDtypeStruct)\n self.assertEqual(out.shape, (3, 4))\n self.assertEqual(count[...], 0)\n\n @parameterized.parameters(\n (True, False),\n (False, False),\n )\n def test_eval_shape_no_input_output_aliasing(self, graph, graph_updates):\n v = nnx.Param(jnp.array(1.0))\n\n def f(v):\n return v\n\n with self.assertRaises(ValueError):\n nnx.eval_shape(f, v, graph=graph, graph_updates=graph_updates)\n\n @parameterized.parameters(\n (True, False),\n (False, False),\n )\n def test_eval_shape_no_shared_variable_refs(self, graph, graph_updates):\n v = nnx.Param(jnp.array(1.0))\n\n def f(v1, v2):\n v1[...] += 1\n return None\n\n with self.assertRaises(ValueError):\n nnx.eval_shape(f, v, v, graph=graph, graph_updates=graph_updates)\n\nclass TestShardMap(parameterized.TestCase):\n def test_basic_shardmap(self):\n n_devices = jax.local_device_count()\n devices = mesh_utils.create_device_mesh((n_devices,))\n mesh = jax.sharding.Mesh(devices, ('a',))\n PS = jax.sharding.PartitionSpec\n\n state_sharding = nnx.StateSharding(\n {\n nnx.PathContains('kernel'): PS(None, 'a'),\n nnx.PathContains('bias'): PS(),\n }\n )\n\n m = nnx.Linear(16, 32, rngs=nnx.Rngs(0))\n\n self.assertNotIsInstance(m.kernel.sharding, jax.sharding.NamedSharding)\n\n @nnx.shard_map(mesh=mesh, in_specs=(state_sharding,), out_specs=None)\n def f(m: nnx.Linear):\n self.assertEqual(\n m.kernel.shape, (m.in_features, m.out_features // n_devices)\n )\n self.assertEqual(m.bias.shape, (m.out_features,))\n\n f(m)\n\n self.assertIsInstance(m.kernel.sharding, jax.sharding.NamedSharding)\n\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_basic_shardmap_variables(self, graph, graph_updates):\n n_devices = jax.local_device_count()\n devices = mesh_utils.create_device_mesh((n_devices,))\n mesh = jax.sharding.Mesh(devices, ('a',))\n P = jax.sharding.PartitionSpec\n\n rngs = nnx.Rngs(0)\n w = nnx.Param(jax.random.normal(rngs.params(), (16, 32)))\n b = nnx.Param(jax.random.normal(rngs.params(), (32,)))\n count = nnx.BatchStat(jnp.array(0))\n\n self.assertNotIsInstance(w.sharding, jax.sharding.NamedSharding)\n\n @nnx.shard_map(\n mesh=mesh, in_specs=(P(None, 'a'), P(), P()), out_specs=None,\n graph=graph, graph_updates=graph_updates,\n )\n def f(w, b, count):\n count[...] += 1\n self.assertEqual(w.shape, (16, 32 // n_devices))\n self.assertEqual(b.shape, (32,))\n\n f(w, b, count)\n\n if graph and graph_updates:\n self.assertIsInstance(w.sharding, jax.sharding.NamedSharding)\n self.assertIsInstance(b.sharding, jax.sharding.NamedSharding)\n self.assertEqual(count[...], 1)\n\n def test_from_state(self):\n n_devices = jax.local_device_count()\n devices = mesh_utils.create_device_mesh((n_devices,))\n mesh = jax.sharding.Mesh(devices, ('a',))\n PS = jax.sharding.PartitionSpec\n\n state_spec = nnx.State(\n {\n 'kernel': PS(None, 'a'),\n 'bias': PS(),\n }\n )\n state_sharding = nnx.StateSharding(state_spec)\n\n m = nnx.Linear(16, 32, rngs=nnx.Rngs(0))\n\n self.assertNotIsInstance(m.kernel.sharding, jax.sharding.NamedSharding)\n\n @nnx.shard_map(mesh=mesh, in_specs=(state_sharding,), out_specs=None)\n def f(m: nnx.Linear):\n self.assertEqual(\n m.kernel.shape, (m.in_features, m.out_features // n_devices)\n )\n self.assertEqual(m.bias.shape, (m.out_features,))\n\n f(m)\n\n self.assertIsInstance(m.kernel.sharding, jax.sharding.NamedSharding)\n self.assertIsInstance(m.bias.sharding, jax.sharding.NamedSharding)\n\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_simple_data_parallel(self, graph, graph_updates):\n P = jax.sharding.PartitionSpec\n n_devices = jax.local_device_count()\n\n mesh = jax.sharding.Mesh(jax.local_devices(), ('data',))\n\n with jax.set_mesh(mesh):\n m = nnx.Linear(\n in_features=2,\n out_features=3,\n kernel_metadata={'out_sharding': jax.P(None)},\n bias_metadata={'out_sharding': jax.P(None)},\n rngs=nnx.Rngs(0),\n )\n x = jnp.ones((32, 2))\n\n @nnx.shard_map(\n mesh=mesh, in_specs=(P(None), P('data')), out_specs=P('data'),\n graph=graph, graph_updates=graph_updates,\n )\n def f(m, x):\n self.assertEqual(x.shape, (32 // n_devices, 2))\n return m(x)\n\n y = f(m, x)\n\n self.assertEqual(y.shape, (32, 3))\n self.assertIsInstance(y.sharding, jax.sharding.NamedSharding)\n if graph and graph_updates:\n self.assertIsInstance(m.kernel.sharding, jax.sharding.NamedSharding)\n self.assertIsInstance(m.bias.sharding, jax.sharding.NamedSharding)\n\n def test_simple_tensor_parallel(self):\n P = jax.sharding.PartitionSpec\n\n mesh = jax.sharding.Mesh(jax.local_devices(), ('model',))\n\n class MLP(nnx.Module):\n def __init__(self, din, dhidden, dout, *, rngs: nnx.Rngs):\n self.linear1 = nnx.Linear(din, dhidden, use_bias=False, rngs=rngs)\n self.linear2 = nnx.Linear(dhidden, dout, use_bias=False, rngs=rngs)\n\n def __call__(self, x):\n return self.linear2(jax.nn.relu(self.linear1(x)))\n\n m = MLP(2, 64, 3, rngs=nnx.Rngs(0))\n x = jnp.ones((32, 2))\n\n def path_ends_with(path_suffix):\n return lambda path, value: path[-len(path_suffix) :] == path_suffix\n\n model_sharding = nnx.StateSharding(\n {\n path_ends_with(('linear1', 'kernel')): P(None, 'model'),\n path_ends_with(('linear2', 'kernel')): P('model', None),\n }\n )\n\n @nnx.shard_map(\n mesh=mesh, in_specs=(model_sharding, P(None)), out_specs=P(None)\n )\n def f(m, x):\n y = m(x)\n return jax.lax.psum(y, 'model')\n\n y = f(m, x)\n\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_shardmap_with_sharding_names(self, graph, graph_updates):\n n_devices = jax.local_device_count()\n P = jax.sharding.PartitionSpec\n mesh = jax.sharding.Mesh(jax.local_devices(), ('data',))\n\n with jax.set_mesh(mesh):\n w = nnx.Param(jnp.ones((8, 4)), out_sharding=('data', None))\n b = nnx.Param(jnp.ones((4,)), out_sharding=(None,))\n\n self.assertIsInstance(w.get_raw_value().sharding, jax.sharding.NamedSharding)\n self.assertEqual(w.out_sharding, ('data', None))\n self.assertEqual(b.out_sharding, (None,))\n\n @nnx.shard_map(\n mesh=mesh, in_specs=(P('data', None), P(None)), out_specs=P('data', None),\n graph=graph, graph_updates=graph_updates,\n )\n def f(w, b):\n self.assertEqual(w.shape, (8 // n_devices, 4))\n self.assertEqual(b.shape, (4,))\n return w + b[None]\n\n y = f(w, b)\n self.assertEqual(y.shape, (8, 4))\n self.assertIsInstance(y.sharding, jax.sharding.NamedSharding)\n\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_shardmap_sharding_names_mutation(self, graph, graph_updates):\n n_devices = jax.local_device_count()\n P = jax.sharding.PartitionSpec\n mesh = jax.sharding.Mesh(jax.local_devices(), ('data',))\n\n with jax.set_mesh(mesh):\n w = nnx.Param(jnp.zeros((8, 4)), out_sharding=('data', None))\n count = nnx.BatchStat(jnp.array(0))\n\n @nnx.shard_map(\n mesh=mesh, in_specs=(P('data', None), P()), out_specs=P('data', None),\n graph=graph, graph_updates=graph_updates,\n )\n def f(w, count):\n count[...] += 1\n self.assertEqual(w.shape, (8 // n_devices, 4))\n return w + 1.0\n\n y = f(w, count)\n self.assertEqual(count[...], 1)\n self.assertEqual(y.shape, (8, 4))\n np.testing.assert_allclose(w[...], jnp.zeros((8, 4)))\n\n def test_shardmap_shared_variable(self):\n P = jax.sharding.PartitionSpec\n mesh = jax.sharding.Mesh(jax.local_devices(), ('data',))\n\n v = nnx.Param(jnp.array(1.0))\n\n class Container(nnx.Module):\n def __init__(self, v):\n self.v = v\n\n c1 = Container(v)\n c2 = Container(v)\n\n @nnx.shard_map(\n mesh=mesh, in_specs=(P(), P(), P()), out_specs=P(),\n graph=True, graph_updates=True,\n )\n def f(c1, c2, x):\n c1.v[...] += x\n return c1.v[...] + c2.v[...]\n\n y = f(c1, c2, jnp.array(1.0))\n np.testing.assert_allclose(y, 4.0)\n np.testing.assert_allclose(v[...], 2.0)\n\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_shardmap_module_variable_update(self, graph, graph_updates):\n P = jax.sharding.PartitionSpec\n mesh = jax.sharding.Mesh(jax.local_devices(), ('data',))\n\n class Foo(nnx.Module):\n def __init__(self):\n self.count = nnx.Variable(jnp.array(0))\n\n m = Foo()\n\n @nnx.shard_map(\n mesh=mesh, in_specs=(P(), P()), out_specs=P(),\n graph=graph, graph_updates=graph_updates,\n )\n def f(m, x):\n m.count[...] += 1\n return x * 2\n\n y = f(m, jnp.array(3.0))\n np.testing.assert_allclose(y, 6.0)\n self.assertEqual(m.count[...], 1)\n y = f(m, jnp.array(3.0))\n self.assertEqual(m.count[...], 2)\n\n def test_shard_map_inconsistent_aliasing(self):\n v = nnx.Param(jnp.array(1.0))\n P = jax.sharding.PartitionSpec\n mesh = jax.sharding.Mesh(jax.devices(), ('x',))\n\n @nnx.shard_map(\n mesh=mesh,\n in_specs=(P(), P('x')),\n out_specs=P(),\n graph=True, graph_updates=False,\n )\n def f(a, b):\n return a[...] + b[...]\n\n with self.assertRaisesRegex(ValueError, 'Inconsistent aliasing'):\n f(v, v)\n\n\nclass TestGrad(parameterized.TestCase):\n def test_grad(self):\n p1 = nnx.Param(10.0)\n p2 = nnx.Param(20.0)\n\n m = nnx.Dict(\n a=nnx.List([p1, p2]),\n b=p1,\n c=7,\n d=5.0,\n )\n\n @nnx.grad\n def f(m: nnx.Dict):\n # sum all params\n return m['a'][0][...] + m['a'][1][...] + m['b'][...]\n\n grads = f(m)\n\n assert m.a[0] is m.b\n assert isinstance(grads, nnx.State)\n assert grads['a'][0][...] == 2.0\n assert issubclass(type(grads['a'][0]), nnx.Variable)\n assert grads['a'][1][...] == 1.0\n assert issubclass(type(grads['a'][1]), nnx.Variable)\n assert len(nnx.to_flat_state(grads)) == 2\n\n nnx.update(m, grads)\n\n assert m['a'][0] is m.b\n assert m['a'][0][...] == 2.0\n assert m['a'][1][...] == 1.0\n assert m['b'][...] == 2.0\n assert m['c'] == 7\n assert m['d'] == 5.0\n\n def test_grad_with_multiple_ref_types(self):\n m = nnx.Dict(\n a=nnx.List([nnx.Param(jnp.array(10.0)), nnx.BatchStat(jnp.array(20.0))]),\n b=nnx.Param(jnp.array(10.0)),\n c=7,\n d=5.0,\n )\n\n @nnx.grad\n def f(m: nnx.Dict):\n # sum all params\n return m.a[0] + m.a[1] + m.b\n\n grads = f(m)\n\n assert isinstance(grads, nnx.State)\n assert grads['a'][0][...] == 1.0\n assert issubclass(type(grads['a'][0]), nnx.Param)\n assert len(grads) == 2\n\n nnx.update(m, grads)\n\n assert m.a[0][...] == 1.0\n assert m.a[1][...] == 20.0\n assert m.b[...] == 1.0\n assert m.c == 7\n assert m.d == 5.0\n\n def test_grad_with_type_predicate(self):\n m = nnx.Dict(\n a=nnx.List([nnx.Param(jnp.array(10.0)), nnx.BatchStat(jnp.array(20.0))]),\n b=nnx.Param(jnp.array(10.0)),\n c=7,\n d=5.0,\n )\n\n @nnx.grad(argnums=nnx.DiffState(0, nnx.BatchStat))\n def f(m: nnx.Dict):\n # sum all params\n return m.a[0] + m.a[1] + m.b\n\n grads = f(m)\n\n assert isinstance(grads, nnx.State)\n assert grads['a'][1][...] == 1.0\n assert issubclass(type(grads['a'][1]), nnx.BatchStat)\n assert len(grads) == 1\n\n nnx.update(m, grads)\n\n assert m.a[0][...] == 10.0\n assert m.a[1][...] == 1.0\n assert m.b[...] == 10.0\n assert m.c == 7\n assert m.d == 5.0\n\n def test_multiple_inputs(self):\n rngs = nnx.Rngs(0)\n m = nnx.Linear(2, 3, rngs=rngs)\n loss_fn = lambda m, x, y: jnp.mean((m(x) - y) ** 2)\n grad_fn = nnx.grad(loss_fn)\n x = jax.random.uniform(rngs(), (1, 2))\n y = jnp.ones((1, 3))\n grads = grad_fn(m, x, y)\n\n assert 'kernel' in grads\n assert grads['kernel'].shape == (2, 3)\n assert 'bias' in grads\n assert grads['bias'].shape == (3,)\n\n @parameterized.parameters(\n {\n 'loss_fn': lambda m1, m2, x, y: jnp.mean((m2(m1(x)) - y) ** 2),\n 'argnums': (0, 1),\n },\n {\n 'loss_fn': lambda x, m1, y, m2: jnp.mean((m2(m1(x)) - y) ** 2),\n 'argnums': (1, 3),\n },\n )\n def test_multiple_graph_nodes(self, loss_fn, argnums):\n rngs = nnx.Rngs(0)\n m1 = nnx.Linear(2, 3, rngs=rngs)\n m2 = nnx.Linear(3, 3, rngs=rngs)\n grad_fn = nnx.grad(loss_fn, argnums=argnums)\n x = jax.random.uniform(rngs(), (1, 2))\n y = jnp.ones((1, 3))\n inputs = [x, y]\n inputs.insert(argnums[0], m1)\n inputs.insert(argnums[1], m2)\n grads_m1, grads_m2 = grad_fn(*inputs)\n\n assert 'kernel' in grads_m1\n assert grads_m1['kernel'].shape == (2, 3)\n assert 'bias' in grads_m1\n assert grads_m1['bias'].shape == (3,)\n assert 'kernel' in grads_m2\n assert grads_m2['kernel'].shape == (3, 3)\n assert 'bias' in grads_m2\n assert grads_m2['bias'].shape == (3,)\n\n def test_multiple_args(self):\n m1 = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n m2 = nnx.Linear(2, 3, rngs=nnx.Rngs(1))\n\n m1_diffstate = nnx.DiffState(0, nnx.PathContains('kernel'))\n m2_diffstate = nnx.DiffState(1, nnx.PathContains('bias'))\n\n @nnx.grad(argnums=(m1_diffstate, m2_diffstate))\n def loss_fn(m1: nnx.Linear, m2: nnx.Linear):\n return jnp.mean(m1.kernel * m2.kernel) + jnp.mean(m1.bias * m2.bias)\n\n grads_m1, grads_m2 = loss_fn(m1, m2)\n\n self.assertIn('kernel', grads_m1)\n self.assertNotIn('bias', grads_m1)\n self.assertNotIn('kernel', grads_m2)\n self.assertIn('bias', grads_m2)\n\n def test_multiple_args_in_pytrees(self):\n m1 = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n m2 = nnx.Linear(2, 3, rngs=nnx.Rngs(1))\n\n m1_diffstate = nnx.DiffState(0, nnx.PathContains('kernel'))\n m2_diffstate = nnx.DiffState(1, nnx.PathContains('bias'))\n\n @nnx.grad(argnums=(m1_diffstate, m2_diffstate))\n def loss_fn(l1: list[nnx.Linear], l2: list[nnx.Linear]):\n return jnp.mean(l1[0].kernel * l2[0].kernel) + jnp.mean(\n l1[0].bias * l2[0].bias\n )\n\n grads_m1, grads_m2 = loss_fn([m1], [m2])\n\n self.assertIn('kernel', grads_m1[0])\n self.assertNotIn('bias', grads_m1[0])\n self.assertNotIn('kernel', grads_m2[0])\n self.assertIn('bias', grads_m2[0])\n\n def test_value_and_grad_multiple_args_in_pytrees(self):\n m1 = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n m2 = nnx.Linear(2, 3, rngs=nnx.Rngs(1))\n\n m1_diffstate = nnx.DiffState(0, nnx.PathContains('kernel'))\n m2_diffstate = nnx.DiffState(1, nnx.PathContains('bias'))\n\n @nnx.value_and_grad(argnums=(m1_diffstate, m2_diffstate))\n def loss_fn(l1: list[nnx.Linear], l2: list[nnx.Linear]):\n return jnp.mean(l1[0].kernel * l2[0].kernel) + jnp.mean(\n l1[0].bias * l2[0].bias\n )\n\n loss, (grads_m1, grads_m2) = loss_fn([m1], [m2])\n\n self.assertEqual(loss.shape, ())\n self.assertIn('kernel', grads_m1[0])\n self.assertNotIn('bias', grads_m1[0])\n self.assertNotIn('kernel', grads_m2[0])\n self.assertIn('bias', grads_m2[0])\n\n def test_value_and_grad_with_aux(self):\n m1 = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n m2 = nnx.Linear(2, 3, rngs=nnx.Rngs(1))\n\n m1_diffstate = nnx.DiffState(0, nnx.PathContains('kernel'))\n m2_diffstate = nnx.DiffState(1, nnx.PathContains('bias'))\n\n @nnx.value_and_grad(argnums=(m1_diffstate, m2_diffstate), has_aux=True)\n def loss_fn(l1: list[nnx.Linear], l2: list[nnx.Linear]):\n loss = jnp.mean(l1[0].kernel * l2[0].kernel) + jnp.mean(\n l1[0].bias * l2[0].bias\n )\n l1[0].kernel.set_value(jnp.array(-1.0))\n m3 = nnx.Linear(2, 3, rngs=nnx.Rngs(2))\n return loss, m3\n\n (loss, m3), (grads_m1, grads_m2) = loss_fn([m1], [m2])\n\n self.assertEqual(m1.kernel[...], -1.0)\n self.assertEqual(loss.shape, ())\n self.assertIsInstance(m3, nnx.Linear)\n self.assertIn('kernel', grads_m1[0])\n self.assertNotIn('bias', grads_m1[0])\n self.assertNotIn('kernel', grads_m2[0])\n self.assertIn('bias', grads_m2[0])\n\n def test_variables_in_grad(self):\n p1 = nnx.Param(10.0)\n p2 = nnx.Param(20.0)\n\n m = dict(a=[p1, p2], b=p1)\n\n @nnx.grad\n def f(m: dict):\n return m['a'][0] + m['a'][1] + m['b']\n\n grads = f(m)\n\n assert m['a'][0] is m['b']\n assert isinstance(grads, dict)\n assert issubclass(type(grads['a'][0]), nnx.Variable)\n assert grads['a'][1][...] == 1.0\n assert issubclass(type(grads['a'][1]), nnx.Variable)\n assert len(jax.tree.leaves(grads)) == 2\n\n jax.tree.map(\n nnx.update, m, grads, is_leaf=lambda x: isinstance(x, nnx.Variable)\n )\n\n assert m['a'][0] is m['b']\n assert m['a'][0][...] == 2.0\n assert m['a'][1][...] == 1.0\n assert m['b'][...] == 2.0\n\n @parameterized.parameters(\n (True, True),\n (True, False),\n (False, False),\n )\n def test_tree_mode_grad(self, graph, graph_updates):\n m = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n\n @nnx.grad(graph=graph, graph_updates=graph_updates)\n def loss_fn(m: nnx.Linear):\n return jnp.mean(m.kernel) + jnp.mean(m.bias)\n\n grads = loss_fn(m)\n\n grad_type = nnx.State if graph and graph_updates else nnx.Linear\n self.assertIsInstance(grads, grad_type)\n self.assertEqual(grads.kernel.shape, (2, 3))\n self.assertEqual(grads.bias.shape, (3,))\n\n @parameterized.parameters(\n (True, True),\n (True, False),\n (False, False),\n )\n def test_tree_mode_grad_multiple_inputs(self, graph, graph_updates):\n rngs = nnx.Rngs(0)\n m = nnx.Linear(2, 3, rngs=rngs)\n loss_fn = lambda m, x, y: jnp.mean((m(x) - y) ** 2)\n grad_fn = nnx.grad(loss_fn, graph=graph, graph_updates=graph_updates)\n x = jax.random.uniform(rngs(), (1, 2))\n y = jnp.ones((1, 3))\n grads = grad_fn(m, x, y)\n\n grad_type = nnx.State if graph and graph_updates else nnx.Linear\n self.assertIsInstance(grads, grad_type)\n self.assertEqual(grads.kernel.shape, (2, 3))\n self.assertEqual(grads.bias.shape, (3,))\n\n @parameterized.parameters(\n (True, True),\n (True, False),\n (False, False),\n )\n def test_tree_mode_grad_multiple_graph_nodes(self, graph, graph_updates):\n rngs = nnx.Rngs(0)\n m1 = nnx.Linear(2, 3, rngs=rngs)\n m2 = nnx.Linear(3, 3, rngs=rngs)\n loss_fn = lambda m1, m2, x, y: jnp.mean((m2(m1(x)) - y) ** 2)\n grad_fn = nnx.grad(loss_fn, argnums=(0, 1), graph=graph,\n graph_updates=graph_updates)\n x = jax.random.uniform(rngs(), (1, 2))\n y = jnp.ones((1, 3))\n grads_m1, grads_m2 = grad_fn(m1, m2, x, y)\n\n grad_type = nnx.State if graph and graph_updates else nnx.Linear\n self.assertIsInstance(grads_m1, grad_type)\n self.assertEqual(grads_m1.kernel.shape, (2, 3))\n self.assertEqual(grads_m1.bias.shape, (3,))\n self.assertIsInstance(grads_m2, grad_type)\n self.assertEqual(grads_m2.kernel.shape, (3, 3))\n self.assertEqual(grads_m2.bias.shape, (3,))\n\n @parameterized.parameters(\n (True, True),\n (True, False),\n (False, False),\n )\n def test_tree_mode_value_and_grad_with_aux(self, graph, graph_updates):\n m = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n\n @nnx.value_and_grad(has_aux=True, graph=graph, graph_updates=graph_updates)\n def loss_fn(m: nnx.Linear):\n loss = jnp.mean(m.kernel)\n m.kernel[...] = jnp.ones_like(m.kernel[...])\n return loss, {'aux': 1}\n\n (loss, aux), grads = loss_fn(m)\n\n self.assertEqual(loss.shape, ())\n self.assertEqual(aux, {'aux': 1})\n grad_type = nnx.State if graph and graph_updates else nnx.Linear\n self.assertIsInstance(grads, grad_type)\n self.assertEqual(grads.kernel.shape, (2, 3))\n np.testing.assert_allclose(m.kernel[...], jnp.ones((2, 3)))\n\n\nclass TestCustomVJP(parameterized.TestCase):\n @parameterized.parameters(\n (True, True),\n (True, False),\n (False, False),\n )\n def test_basic_call(self, graph, graph_updates):\n m1 = nnx.Linear(1, 1, rngs=nnx.Rngs(0))\n m2 = nnx.Linear(1, 1, rngs=nnx.Rngs(1))\n\n @nnx.custom_vjp(graph=graph, graph_updates=graph_updates)\n def f(m1: nnx.Linear, m2: nnx.Linear):\n y = m1.kernel * m2.kernel\n return y\n\n def f_fwd(m1, m2):\n y = f(m1, m2)\n return y, (m1, m2)\n\n def f_bwd(res, g):\n inputs_g, out_g = g\n m1, m2 = res\n return inputs_g\n\n f.defvjp(f_fwd, f_bwd)\n\n y = f(m1, m2)\n\n self.assertEqual(y.shape, (1, 1))\n\n @parameterized.parameters(\n (True,),\n (False,),\n )\n def test_basic_call_with_state(self, graph):\n m1 = nnx.Linear(1, 1, rngs=nnx.Rngs(0))\n m2 = nnx.Linear(1, 1, rngs=nnx.Rngs(1))\n state = nnx.BatchStat(jnp.array(0.0))\n\n @nnx.custom_vjp(\n nondiff_argnums=(2,), graph=graph, graph_updates=False,\n )\n def f(m1: nnx.Linear, m2: nnx.Linear, state):\n y = m1.kernel * m2.kernel\n state[...] = jnp.array(-1.0)\n return y\n\n def f_fwd(m1, m2, state):\n y = f(m1, m2, state)\n return y, (m1, m2)\n\n def f_bwd(state, res, g):\n inputs_g, out_g = g\n m1, m2 = res\n return inputs_g\n\n f.defvjp(f_fwd, f_bwd)\n\n y = f(m1, m2, state)\n\n self.assertEqual(state[...], -1.0)\n self.assertEqual(y.shape, (1, 1))\n\n @parameterized.parameters(\n (True, True),\n (True, False),\n (False, False),\n )\n def test_jax_example(self, graph, graph_updates):\n @dataclasses.dataclass\n class Foo(nnx.Module):\n x: nnx.Param[jax.Array]\n y: nnx.Param[jax.Array]\n z: int\n\n @nnx.custom_vjp(graph=graph, graph_updates=graph_updates)\n def f(m: Foo):\n m.z += 1\n return jnp.sin(m.x) * m.y # type: ignore\n\n def f_fwd(m: Foo):\n y = f(m)\n res = (jnp.cos(m.x), jnp.sin(m.x), m) # type: ignore\n return y, res\n\n def f_bwd(res, g):\n cos_x, sin_x, m = res\n if graph and graph_updates:\n (m_g,), out_g = g\n self.assertIsInstance(m_g, nnx.State)\n m_g['x'][...] = cos_x * out_g * m.y\n m_g['y'][...] = sin_x * out_g\n return (m_g,)\n else:\n out_g = g\n m_g = nnx.clone(m)\n m_g.x[...] = cos_x * out_g * m.y\n m_g.y[...] = sin_x * out_g\n return (m_g,)\n\n f.defvjp(f_fwd, f_bwd)\n\n m = Foo(nnx.Param(jnp.array(1.0)), nnx.Param(jnp.array(2.0)), 0)\n\n if graph and graph_updates:\n grads = nnx.grad(f, argnums=nnx.DiffState(0, ...))(m)\n self.assertIsInstance(grads, nnx.State)\n else:\n grads = nnx.grad(\n f, graph=graph, graph_updates=graph_updates,\n )(m)\n self.assertIsInstance(grads, Foo)\n np.testing.assert_allclose(grads.x[...], jnp.cos(1.0) * 2.0) # type: ignore\n np.testing.assert_allclose(grads.y[...], jnp.sin(1.0)) # type: ignore\n if graph and graph_updates:\n self.assertEqual(m.z, 1)\n else:\n self.assertEqual(m.z, 0)\n\n def test_diff_state(self):\n @dataclasses.dataclass\n class Foo(nnx.Module):\n x: nnx.Param[jax.Array]\n y: nnx.Param[jax.Array]\n z: int\n\n x_in_path = nnx.PathContains('x')\n diff_state = nnx.DiffState(0, x_in_path)\n\n @nnx.custom_vjp(nondiff_argnums=(diff_state,))\n def f(m: Foo):\n m.z += 1\n return jnp.sin(m.x) * m.y # type: ignore\n\n def f_fwd(m: Foo):\n y = f(m)\n res = (jnp.cos(m.x), m) # type: ignore\n return y, res\n\n def f_bwd(res, g):\n (m_g,), out_g = g\n cos_x, m = res\n\n self.assertIsInstance(m_g, nnx.State)\n self.assertEqual(out_g.shape, ())\n self.assertIsInstance(m, Foo)\n\n m_g['x'][...] = cos_x * out_g * m.y\n del m_g['y']\n return (m_g,)\n\n f.defvjp(f_fwd, f_bwd)\n\n m = Foo(nnx.Param(jnp.array(1.0)), nnx.Param(jnp.array(2.0)), 0)\n\n grad: nnx.State = nnx.grad(f, argnums=nnx.DiffState(0, x_in_path))(m)\n\n np.testing.assert_allclose(grad['x'][...], jnp.cos(1.0) * 2.0) # type: ignore\n self.assertEqual(m.z, 1)\n\n @parameterized.parameters(\n (True, True),\n (True, False),\n (False, False),\n )\n def test_jax_example_with_remat(self, graph, graph_updates):\n @dataclasses.dataclass\n class Foo(nnx.Module):\n x: nnx.Param[jax.Array]\n y: nnx.Param[jax.Array]\n z: int\n\n @nnx.custom_vjp(graph=graph, graph_updates=graph_updates)\n @nnx.remat(graph=graph, graph_updates=graph_updates)\n def f(m: Foo):\n m.z += 1\n return jnp.sin(m.x) * m.y # type: ignore\n\n def f_fwd(m: Foo):\n y = f(m)\n res = (jnp.cos(m.x), jnp.sin(m.x), m) # type: ignore\n return y, res\n\n def f_bwd(res, g):\n cos_x, sin_x, m = res\n if graph and graph_updates:\n (m_g,), out_g = g\n self.assertIsInstance(m_g, nnx.State)\n m_g['x'][...] = cos_x * out_g * m.y\n m_g['y'][...] = sin_x * out_g\n return (m_g,)\n else:\n out_g = g\n m_g = jax.tree.map(lambda x: x, m)\n m_g.x[...] = cos_x * out_g * m.y\n m_g.y[...] = sin_x * out_g\n return (m_g,)\n\n f.defvjp(f_fwd, f_bwd)\n\n m = Foo(nnx.Param(jnp.array(1.0)), nnx.Param(jnp.array(2.0)), 0)\n\n @nnx.jit(graph=graph, graph_updates=graph_updates)\n def loss_fn(m):\n return f(m)\n\n if graph and graph_updates:\n grads = nnx.grad(loss_fn, argnums=nnx.DiffState(0, ...))(m)\n self.assertIsInstance(grads, nnx.State)\n else:\n grads = nnx.grad(\n loss_fn, graph=graph, graph_updates=graph_updates,\n )(m)\n self.assertIsInstance(grads, Foo)\n np.testing.assert_allclose(grads.x[...], jnp.cos(1.0) * 2.0) # type: ignore\n np.testing.assert_allclose(grads.y[...], jnp.sin(1.0)) # type: ignore\n if graph and graph_updates:\n self.assertEqual(m.z, 1)\n else:\n self.assertEqual(m.z, 0)\n\n def test_two_args(self):\n @dataclasses.dataclass\n class Foo(nnx.Module):\n x: nnx.Param[jax.Array]\n y: nnx.Param[jax.Array]\n z: int\n\n @nnx.custom_vjp\n def f(m1: Foo, m2: Foo):\n m1.z += 1\n y = jnp.sin(m1.x) * m1.y # type: ignore\n return y, m2\n\n def f_fwd(m1: Foo, m2: Foo):\n y, m2 = f(m1, m2)\n res = (jnp.cos(m1.x), jnp.sin(m1.x), m1) # type: ignore\n return (y, m2), res\n\n def f_bwd(res, g):\n (m1_g, m2_g), (y_g, _) = g\n cos_x, sin_x, m = res\n\n self.assertIsInstance(m1_g, nnx.State)\n self.assertIsInstance(m2_g, nnx.State)\n self.assertEqual(y_g.shape, ())\n self.assertIsInstance(m, Foo)\n\n m1_g = nnx.State(dict(x=cos_x * y_g * m.y, y=sin_x * y_g))\n m2_g = nnx.State(dict(x=m2_g['x'], y=m2_g['y']))\n\n return m1_g, m2_g\n\n f.defvjp(f_fwd, f_bwd)\n\n m1 = Foo(nnx.Param(jnp.array(1.0)), nnx.Param(jnp.array(2.0)), 0)\n m2 = Foo(nnx.Param(jnp.array(3.0)), nnx.Param(jnp.array(4.0)), 0)\n\n def loss_fn(m1, m2):\n y, m2 = f(m1, m2)\n return y + m2.x * m2.y\n\n m1_grad: nnx.State\n m2_grad: nnx.State\n m1_grad, m2_grad = nnx.grad(\n loss_fn, argnums=(nnx.DiffState(0, ...), nnx.DiffState(1, ...))\n )(m1, m2)\n\n np.testing.assert_allclose(m1_grad['x'][...], jnp.cos(1.0) * 2.0) # type: ignore\n np.testing.assert_allclose(m1_grad['y'][...], jnp.sin(1.0)) # type: ignore\n self.assertEqual(m1.z, 1)\n np.testing.assert_allclose(m2_grad['x'][...], 4.0) # type: ignore\n np.testing.assert_allclose(m2_grad['y'][...], 3.0) # type: ignore\n\n @parameterized.parameters(\n (True, True),\n (True, False),\n (False, False),\n )\n def test_non_diff_args(self, graph, graph_updates):\n @dataclasses.dataclass\n class Foo(nnx.Module):\n x: nnx.Param[jax.Array]\n y: nnx.Param[jax.Array]\n z: int\n\n @nnx.custom_vjp(\n nondiff_argnums=(0, 2), graph=graph, graph_updates=graph_updates,\n )\n def f(a, m: Foo, b):\n self.assertEqual(a, 1)\n self.assertEqual(b, 2)\n m.z += 1\n return jnp.sin(m.x) * m.y # type: ignore\n\n def f_fwd(a, m: Foo, b):\n self.assertEqual(a, 1)\n self.assertEqual(b, 2)\n y = f(a, m, b)\n res = (jnp.cos(m.x), jnp.sin(m.x), m) # type: ignore\n return y, res\n\n def f_bwd(a, b, res, g):\n cos_x, sin_x, m = res\n self.assertEqual(a, 1)\n self.assertEqual(b, 2)\n if graph and graph_updates:\n (m_g,), out_g = g\n self.assertIsInstance(m_g, nnx.State)\n m_g['x'][...] = cos_x * out_g * m.y\n m_g['y'][...] = sin_x * out_g\n return (m_g,)\n else:\n out_g = g\n m_g = jax.tree.map(lambda x: x, m)\n m_g.x[...] = cos_x * out_g * m.y\n m_g.y[...] = sin_x * out_g\n return (m_g,)\n\n f.defvjp(f_fwd, f_bwd)\n\n m = Foo(nnx.Param(jnp.array(1.0)), nnx.Param(jnp.array(2.0)), 0)\n\n def loss_fn(m):\n a = 1\n b = 2\n return f(a, m, b)\n\n if graph and graph_updates:\n grads = nnx.grad(loss_fn, argnums=nnx.DiffState(0, ...))(m)\n self.assertIsInstance(grads, nnx.State)\n else:\n grads = nnx.grad(\n loss_fn, graph=graph, graph_updates=graph_updates,\n )(m)\n self.assertIsInstance(grads, Foo)\n np.testing.assert_allclose(grads.x[...], jnp.cos(1.0) * 2.0) # type: ignore\n np.testing.assert_allclose(grads.y[...], jnp.sin(1.0)) # type: ignore\n if graph and graph_updates:\n self.assertEqual(m.z, 1)\n else:\n self.assertEqual(m.z, 0)\n\n def test_docs_example(self):\n import jax.numpy as jnp\n from flax import nnx\n\n class Foo(nnx.Module):\n def __init__(self, x, y):\n self.x = nnx.Param(x)\n self.y = nnx.Param(y)\n\n @nnx.custom_vjp\n def f(m: Foo):\n return jnp.sin(m.x) * m.y # type: ignore\n\n def f_fwd(m: Foo):\n return f(m), (jnp.cos(m.x), jnp.sin(m.x), m) # type: ignore\n\n def f_bwd(res, g):\n ins_g, out_g = g\n cos_x, sin_x, m = res\n tangent_m = nnx.State(dict(x=cos_x * out_g * m.y, y=sin_x * out_g))\n return (tangent_m,)\n\n f.defvjp(f_fwd, f_bwd)\n\n m = Foo(x=jnp.array(1.0), y=jnp.array(2.0))\n grads = nnx.grad(f)(m)\n\n @parameterized.parameters(\n {'use_custom_vjp': False},\n {'use_custom_vjp': True},\n )\n def test_issue(self, use_custom_vjp: bool):\n class MyLinear(nnx.Module):\n def __init__(\n self, in_features: int, out_features: int, *, rngs: nnx.Rngs\n ):\n kernel_init = nnx.initializers.normal(in_features**-0.5)\n self.kernel = nnx.Param(\n kernel_init(rngs.params(), (in_features, out_features), jnp.float32)\n )\n self.bias = nnx.Param(jnp.zeros((out_features,), jnp.float32))\n self.n = nnx.BatchStat(jnp.array(0, jnp.uint32))\n\n def linear(m: MyLinear, x: jax.Array) -> jax.Array:\n m.n[...] += 1\n y = x @ m.kernel + m.bias\n return y\n\n def linear_fwd(m: MyLinear, x: jax.Array):\n return linear(m, x), (m, x)\n\n def linear_bwd(res, g):\n m, x = res\n (m_g, _x_grad), outputs_g = g\n kernel_grad = outputs_g[None, :] * x[:, None]\n bias_grad = outputs_g\n x_grad = m.kernel @ outputs_g\n assert x_grad.shape == x.shape, 'Shape mismatch for x'\n assert m.kernel.shape == kernel_grad.shape, 'Shape mismatch for kernel'\n assert m.bias.shape == bias_grad.shape, 'Shape mismatch for bias'\n return (m_g, x_grad)\n\n if use_custom_vjp:\n linear = nnx.custom_vjp(linear)\n linear.defvjp(linear_fwd, linear_bwd)\n\n @nnx.jit\n def loss_fn(x, mod):\n y = linear(mod, x)\n return y.mean()\n\n mod = MyLinear(10, 5, rngs=nnx.Rngs(0))\n self.assertEqual(mod.n[...], 0)\n x = jax.random.normal(jax.random.key(0), (10,))\n loss, grad = nnx.value_and_grad(loss_fn)(x, mod)\n self.assertEqual(loss.shape, ())\n self.assertEqual(grad.shape, (10,))\n self.assertEqual(mod.n[...], 1)\n\n @parameterized.parameters(\n (True, False),\n (False, False),\n )\n def test_tree_mode_basic_call(self, graph, graph_updates):\n m1 = nnx.Linear(1, 1, rngs=nnx.Rngs(0))\n m2 = nnx.Linear(1, 1, rngs=nnx.Rngs(1))\n state = nnx.BatchStat(jnp.array(0.0))\n\n @nnx.custom_vjp(\n nondiff_argnums=(2,), graph=graph, graph_updates=graph_updates,\n )\n def f(m1: nnx.Linear, m2: nnx.Linear, state):\n y = m1.kernel * m2.kernel\n state[...] = jnp.array(-1.0)\n return y\n\n def f_fwd(m1, m2, state):\n y = f(m1, m2, state)\n return y, (m1, m2)\n\n def f_bwd(state, res, g):\n m1, m2 = res\n return g, g\n\n f.defvjp(f_fwd, f_bwd)\n\n y = f(m1, m2, state)\n\n self.assertEqual(state[...], -1.0)\n self.assertEqual(y.shape, (1, 1))\n\n @parameterized.parameters(\n (True, False),\n (False, False),\n )\n def test_tree_mode_jax_example(self, graph, graph_updates):\n @dataclasses.dataclass\n class Foo(nnx.Module):\n x: nnx.Param[jax.Array]\n y: nnx.Param[jax.Array]\n\n @nnx.custom_vjp(graph=graph, graph_updates=graph_updates)\n def f(m: Foo):\n return jnp.sin(m.x) * m.y\n\n def f_fwd(m: Foo):\n y = f(m)\n res = (jnp.cos(m.x), jnp.sin(m.x), m)\n return y, res\n\n def f_bwd(res, g):\n cos_x, sin_x, m = res\n m_g = jax.tree.map(lambda x: x, m)\n m_g.x[...] = cos_x * g * m.y\n m_g.y[...] = sin_x * g\n return (m_g,)\n\n f.defvjp(f_fwd, f_bwd)\n\n m = Foo(nnx.Param(jnp.array(1.0)), nnx.Param(jnp.array(2.0)))\n\n grads = nnx.grad(f, graph=graph, graph_updates=graph_updates)(m)\n\n self.assertIsInstance(grads, Foo)\n np.testing.assert_allclose(grads.x[...], jnp.cos(1.0) * 2.0)\n np.testing.assert_allclose(grads.y[...], jnp.sin(1.0))\n\n @parameterized.parameters(\n (True, False),\n (False, False),\n )\n def test_tree_mode_with_remat(self, graph, graph_updates):\n @dataclasses.dataclass\n class Foo(nnx.Module):\n x: nnx.Param[jax.Array]\n y: nnx.Param[jax.Array]\n\n @nnx.custom_vjp(graph=graph, graph_updates=graph_updates)\n @nnx.remat(graph=graph, graph_updates=graph_updates)\n def f(m: Foo):\n return jnp.sin(m.x) * m.y\n\n def f_fwd(m: Foo):\n y = f(m)\n res = (jnp.cos(m.x), jnp.sin(m.x), m)\n return y, res\n\n def f_bwd(res, g):\n cos_x, sin_x, m = res\n m_g = jax.tree.map(lambda x: x, m)\n m_g.x[...] = cos_x * g * m.y\n m_g.y[...] = sin_x * g\n return (m_g,)\n\n f.defvjp(f_fwd, f_bwd)\n\n m = Foo(nnx.Param(jnp.array(1.0)), nnx.Param(jnp.array(2.0)))\n\n @nnx.jit(graph=graph, graph_updates=graph_updates)\n def loss_fn(m):\n return f(m)\n\n grads = nnx.grad(loss_fn, graph=graph, graph_updates=graph_updates)(m)\n\n self.assertIsInstance(grads, Foo)\n np.testing.assert_allclose(grads.x[...], jnp.cos(1.0) * 2.0)\n np.testing.assert_allclose(grads.y[...], jnp.sin(1.0))\n\n @parameterized.parameters(\n (True, False),\n (False, False),\n )\n def test_tree_mode_non_diff_args(self, graph, graph_updates):\n @dataclasses.dataclass\n class Foo(nnx.Module):\n x: nnx.Param[jax.Array]\n y: nnx.Param[jax.Array]\n\n @nnx.custom_vjp(nondiff_argnums=(0, 2), graph=graph,\n graph_updates=graph_updates)\n def f(a, m: Foo, b):\n self.assertEqual(a, 1)\n self.assertEqual(b, 2)\n return jnp.sin(m.x) * m.y\n\n def f_fwd(a, m: Foo, b):\n self.assertEqual(a, 1)\n self.assertEqual(b, 2)\n y = f(a, m, b)\n res = (jnp.cos(m.x), jnp.sin(m.x), m)\n return y, res\n\n def f_bwd(a, b, res, g):\n cos_x, sin_x, m = res\n self.assertEqual(a, 1)\n self.assertEqual(b, 2)\n m_g = jax.tree.map(lambda x: x, m)\n m_g.x[...] = cos_x * g * m.y\n m_g.y[...] = sin_x * g\n return (m_g,)\n\n f.defvjp(f_fwd, f_bwd)\n\n m = Foo(nnx.Param(jnp.array(1.0)), nnx.Param(jnp.array(2.0)))\n\n def loss_fn(m):\n a = 1\n b = 2\n return f(a, m, b)\n\n grads = nnx.grad(loss_fn, graph=graph, graph_updates=graph_updates)(m)\n\n self.assertIsInstance(grads, Foo)\n np.testing.assert_allclose(grads.x[...], jnp.cos(1.0) * 2.0)\n np.testing.assert_allclose(grads.y[...], jnp.sin(1.0))\n\n def test_tree_mode_diffstate_error(self):\n x_in_path = nnx.PathContains('x')\n diff_state = nnx.DiffState(0, x_in_path)\n with self.assertRaisesRegex(\n ValueError,\n r'`nondiff_argnums` cannot contain `DiffState` objects',\n ):\n nnx.custom_vjp(lambda m: m, nondiff_argnums=(diff_state,), graph=False)\n\n def test_grad_inconsistent_aliasing(self):\n v = nnx.Param(jnp.array(1.0))\n\n def f(v_diff, v_nondiff):\n return v_diff[...] + v_nondiff[...]\n\n with self.assertRaisesRegex(ValueError, 'Inconsistent aliasing'):\n nnx.grad(f, argnums=0, graph=True, graph_updates=False)(v, v)\n\n def test_custom_vjp_inconsistent_aliasing(self):\n v = nnx.Param(jnp.array(1.0))\n\n @nnx.custom_vjp(nondiff_argnums=(1,), graph=True, graph_updates=False)\n def f(v_diff, v_nondiff):\n return v_diff[...] + v_nondiff[...]\n\n def f_fwd(v_diff, v_nondiff):\n return f(v_diff, v_nondiff), ()\n\n def f_bwd(v_nondiff, res, g):\n return (nnx.clone(v_nondiff),)\n\n f.defvjp(f_fwd, f_bwd)\n\n with self.assertRaisesRegex(ValueError, 'Inconsistent aliasing'):\n f(v, v)\n\n\n def test_custom_vjp_diff_arg_mutation_error(self):\n @nnx.custom_vjp(graph=True, graph_updates=False)\n def f(m):\n m.x[...] += 1\n return m.x[...] * m.y[...]\n\n def f_fwd(m):\n return f(m), (m,)\n\n def f_bwd(res, g):\n (m,) = res\n m_g = nnx.clone(m)\n m_g.x[...] = g * m.y[...]\n m_g.y[...] = g * m.x[...]\n return (m_g,)\n\n f.defvjp(f_fwd, f_bwd)\n\n @dataclasses.dataclass\n class Foo(nnx.Module):\n x: nnx.Param[jax.Array]\n y: nnx.Param[jax.Array]\n\n m = Foo(nnx.Param(jnp.array(1.0)), nnx.Param(jnp.array(2.0)))\n with self.assertRaisesRegex(\n ValueError, 'Variables in differentiable argument'\n ):\n f(m)\n\n\nclass TestVjpJvp(parameterized.TestCase):\n\n @parameterized.parameters(\n (True, False),\n (False, False),\n )\n def test_vjp_basic(self, graph, graph_updates):\n @dataclasses.dataclass\n class Foo(nnx.Module):\n w: nnx.Param[jax.Array]\n\n def f(m: Foo, x):\n return jnp.sum(m.w * x)\n\n m = Foo(w=nnx.Param(jnp.array([1.0, 2.0, 3.0])))\n x = jnp.array([4.0, 5.0, 6.0])\n\n primals_out, vjp_fn = nnx.vjp(\n f, m, x, graph=graph, graph_updates=graph_updates,\n )\n\n np.testing.assert_allclose(primals_out, 1.0 * 4.0 + 2.0 * 5.0 + 3.0 * 6.0)\n\n m_grad, x_grad = vjp_fn(jnp.ones_like(primals_out))\n self.assertIsInstance(m_grad, Foo)\n np.testing.assert_allclose(m_grad.w[...], x)\n np.testing.assert_allclose(x_grad, m.w[...])\n\n @parameterized.parameters(\n (True, False),\n (False, False),\n )\n def test_vjp_has_aux(self, graph, graph_updates):\n @dataclasses.dataclass\n class Foo(nnx.Module):\n w: nnx.Param[jax.Array]\n\n def f(m: Foo, x):\n y = jnp.sum(m.w * x)\n return y, {'input': x}\n\n m = Foo(w=nnx.Param(jnp.array([1.0, 2.0])))\n x = jnp.array([3.0, 4.0])\n\n primals_out, vjp_fn, aux = nnx.vjp(\n f, m, x, has_aux=True, graph=graph, graph_updates=graph_updates,\n )\n\n np.testing.assert_allclose(primals_out, 1.0 * 3.0 + 2.0 * 4.0)\n np.testing.assert_allclose(aux['input'], x)\n\n @parameterized.parameters(\n (True, False),\n (False, False),\n )\n def test_vjp_state_propagation(self, graph, graph_updates):\n @dataclasses.dataclass\n class Foo(nnx.Module):\n w: nnx.Param[jax.Array]\n count: nnx.BatchStat[jax.Array]\n\n def f(m: Foo, x):\n m.count[...] += 1\n return jnp.sum(m.w * x)\n\n m = Foo(\n w=nnx.Param(jnp.array([1.0, 2.0])),\n count=nnx.BatchStat(jnp.array(0)),\n )\n x = jnp.array([3.0, 4.0])\n\n self.assertEqual(m.count[...], 0)\n primals_out, vjp_fn = nnx.vjp(\n f, m, x, graph=graph, graph_updates=graph_updates,\n )\n self.assertEqual(m.count[...], 1)\n\n @parameterized.parameters(\n (True, False),\n (False, False),\n )\n def test_vjp_matches_jax(self, graph, graph_updates):\n def f(w, x):\n return jnp.sum(w * x)\n\n w = jnp.array([1.0, 2.0, 3.0])\n x = jnp.array([4.0, 5.0, 6.0])\n\n jax_primals, jax_vjp_fn = jax.vjp(f, w, x)\n jax_grads = jax_vjp_fn(jnp.ones_like(jax_primals))\n\n nnx_primals, nnx_vjp_fn = nnx.vjp(\n f, w, x, graph=graph, graph_updates=graph_updates,\n )\n nnx_grads = nnx_vjp_fn(jnp.ones_like(nnx_primals))\n\n np.testing.assert_allclose(nnx_primals, jax_primals)\n for ng, jg in zip(nnx_grads, jax_grads):\n np.testing.assert_allclose(ng, jg)\n\n @parameterized.parameters(\n (True, False),\n (False, False),\n )\n def test_vjp_decorator(self, graph, graph_updates):\n @dataclasses.dataclass\n class Foo(nnx.Module):\n w: nnx.Param[jax.Array]\n\n @nnx.vjp(graph=graph, graph_updates=graph_updates)\n def f(m: Foo, x):\n return jnp.sum(m.w * x)\n\n m = Foo(w=nnx.Param(jnp.array([1.0, 2.0])))\n x = jnp.array([3.0, 4.0])\n\n primals_out, vjp_fn = f(m, x)\n np.testing.assert_allclose(primals_out, 11.0)\n m_grad, x_grad = vjp_fn(jnp.ones_like(primals_out))\n np.testing.assert_allclose(m_grad.w[...], x)\n\n @parameterized.parameters(\n (True, False),\n (False, False),\n )\n def test_jvp_basic(self, graph, graph_updates):\n @dataclasses.dataclass\n class Foo(nnx.Module):\n w: nnx.Param[jax.Array]\n\n def f(m: Foo, x):\n return jnp.sum(m.w * x)\n\n m = Foo(w=nnx.Param(jnp.array([1.0, 2.0, 3.0])))\n x = jnp.array([4.0, 5.0, 6.0])\n\n m_tangent = jax.tree.map(jnp.ones_like, m)\n x_tangent = jnp.ones_like(x)\n\n primals_out, tangent_out = nnx.jvp(\n f, (m, x), (m_tangent, x_tangent),\n graph=graph, graph_updates=graph_updates,\n )\n\n np.testing.assert_allclose(primals_out, 1.0 * 4.0 + 2.0 * 5.0 + 3.0 * 6.0)\n expected_tangent = jnp.sum(jnp.ones(3) * x + m.w[...] * jnp.ones(3))\n np.testing.assert_allclose(tangent_out, expected_tangent)\n\n @parameterized.parameters(\n (True, False),\n (False, False),\n )\n def test_jvp_has_aux(self, graph, graph_updates):\n @dataclasses.dataclass\n class Foo(nnx.Module):\n w: nnx.Param[jax.Array]\n\n def f(m: Foo, x):\n y = jnp.sum(m.w * x)\n return y, {'input': x}\n\n m = Foo(w=nnx.Param(jnp.array([1.0, 2.0])))\n x = jnp.array([3.0, 4.0])\n\n m_tangent = jax.tree.map(jnp.ones_like, m)\n x_tangent = jnp.ones_like(x)\n\n primals_out, tangent_out, aux = nnx.jvp(\n f, (m, x), (m_tangent, x_tangent), has_aux=True,\n graph=graph, graph_updates=graph_updates,\n )\n\n np.testing.assert_allclose(primals_out, 11.0)\n np.testing.assert_allclose(aux['input'], x)\n\n @parameterized.parameters(\n (True, False),\n (False, False),\n )\n def test_jvp_state_propagation(self, graph, graph_updates):\n @dataclasses.dataclass\n class Foo(nnx.Module):\n w: nnx.Param[jax.Array]\n count: nnx.BatchStat[jax.Array]\n\n def f(m: Foo, x):\n m.count[...] += 1\n return jnp.sum(m.w * x)\n\n m = Foo(\n w=nnx.Param(jnp.array([1.0, 2.0])),\n count=nnx.BatchStat(jnp.array(0.0)),\n )\n x = jnp.array([3.0, 4.0])\n\n m_tangent = jax.tree.map(jnp.zeros_like, m)\n x_tangent = jnp.zeros_like(x)\n\n self.assertEqual(m.count[...], 0.0)\n primals_out, tangent_out = nnx.jvp(\n f, (m, x), (m_tangent, x_tangent),\n graph=graph, graph_updates=graph_updates,\n )\n self.assertEqual(m.count[...], 1.0)\n\n @parameterized.parameters(\n (True, False),\n (False, False),\n )\n def test_jvp_matches_jax(self, graph, graph_updates):\n def f(w, x):\n return jnp.sum(w * x)\n\n w = jnp.array([1.0, 2.0, 3.0])\n x = jnp.array([4.0, 5.0, 6.0])\n\n w_tangent = jnp.ones_like(w)\n x_tangent = jnp.ones_like(x)\n\n jax_primals, jax_tangents = jax.jvp(f, (w, x), (w_tangent, x_tangent))\n nnx_primals, nnx_tangents = nnx.jvp(\n f, (w, x), (w_tangent, x_tangent),\n graph=graph, graph_updates=graph_updates,\n )\n\n np.testing.assert_allclose(nnx_primals, jax_primals)\n np.testing.assert_allclose(nnx_tangents, jax_tangents)\n\n @parameterized.parameters(\n (True, False),\n (False, False),\n )\n def test_jvp_decorator(self, graph, graph_updates):\n @dataclasses.dataclass\n class Foo(nnx.Module):\n w: nnx.Param[jax.Array]\n\n @nnx.jvp(graph=graph, graph_updates=graph_updates)\n def f(m: Foo, x):\n return jnp.sum(m.w * x)\n\n m = Foo(w=nnx.Param(jnp.array([1.0, 2.0])))\n x = jnp.array([3.0, 4.0])\n\n m_tangent = jax.tree.map(jnp.ones_like, m)\n x_tangent = jnp.ones_like(x)\n\n primals_out, tangent_out = f((m, x), (m_tangent, x_tangent))\n np.testing.assert_allclose(primals_out, 11.0)\n expected_tangent = jnp.sum(jnp.ones(2) * x + m.w[...] * jnp.ones(2))\n np.testing.assert_allclose(tangent_out, expected_tangent)\n\n\nclass TestScan(parameterized.TestCase):\n def test_basic(self):\n class Block(nnx.Module):\n def __init__(self, *, rngs: nnx.Rngs):\n self.linear = nnx.Linear(3, 3, rngs=rngs)\n\n def __call__(self, x: jax.Array):\n x = self.linear(x)\n x = nnx.gelu(x)\n return x\n\n @nnx.split_rngs(splits=5)\n @nnx.scan(in_axes=(nnx.Carry, 0), length=5)\n def create_block(_, rngs: nnx.Rngs):\n return None, Block(rngs=rngs)\n\n _, module = create_block(None, nnx.Rngs(0))\n\n assert module.linear.kernel.shape == (5, 3, 3)\n assert module.linear.bias.shape == (5, 3)\n\n @nnx.scan(in_axes=(nnx.Carry, 0, None), length=5)\n def forward_block(_, block: Block, x: jax.Array):\n return None, block(x)\n\n x = jnp.ones((1, 3))\n out, y = forward_block(None, module, x)\n\n assert y.shape == (5, 1, 3)\n assert out is None\n\n @parameterized.parameters(True, False)\n def test_variables_in_scan(self, graph_updates):\n def block_init(din, dout, rngs):\n w = nnx.Param(jax.random.normal(rngs.params(), (din, dout)))\n b = nnx.Param(jnp.zeros((dout,)))\n return w, b\n\n def block_forward(w, b, x):\n return nnx.gelu(x @ w + b[None])\n\n @nnx.split_rngs(splits=5)\n @nnx.scan(in_axes=0, out_axes=0, length=5, graph_updates=graph_updates)\n def create_block(rngs: nnx.Rngs):\n return block_init(3, 3, rngs)\n\n w, b = create_block(nnx.Rngs(0))\n\n assert w.shape == (5, 3, 3)\n assert b.shape == (5, 3)\n\n @nnx.scan(\n in_axes=(0, 0, nnx.Carry), out_axes=nnx.Carry,\n graph_updates=graph_updates,\n )\n def stack_forward(w, b, x):\n return block_forward(w, b, x)\n\n x = jnp.ones((1, 3))\n y = stack_forward(w, b, x)\n assert y.shape == (1, 3)\n\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_variables_as_carries_in_scan(self, graph, graph_updates):\n w = nnx.Param(jax.random.normal(jax.random.key(0), (3, 3)))\n b = nnx.Param(jnp.zeros((3,)))\n count = nnx.BatchStat(0)\n\n def block_forward(w, b, x):\n return nnx.gelu(x @ w + b[None])\n\n @nnx.scan(\n in_axes=(nnx.Carry, 0), out_axes=(nnx.Carry, 0),\n graph=graph, graph_updates=graph_updates,\n )\n def stack_forward(params, x):\n w, b, count = params\n y = block_forward(w, b, x)\n count[...] += 1\n return (w, b, count), y\n\n x = jnp.ones((5, 1, 3))\n (w, b, count), y = stack_forward((w, b, count), x)\n\n assert y.shape == (5, 1, 3)\n assert count[...] == 5\n\n def test_basic_no_carry(self):\n class Block(nnx.Module):\n def __init__(self, *, rngs: nnx.Rngs):\n self.linear = nnx.Linear(3, 3, rngs=rngs)\n\n def __call__(self, x: jax.Array):\n x = self.linear(x)\n x = nnx.gelu(x)\n return x\n\n @nnx.split_rngs(splits=5)\n @nnx.scan(in_axes=(0,), out_axes=0, length=5)\n def create_block(rngs: nnx.Rngs):\n return Block(rngs=rngs)\n\n module = create_block(nnx.Rngs(0))\n\n assert module.linear.kernel.shape == (5, 3, 3)\n assert module.linear.bias.shape == (5, 3)\n # assert module.node.shape == (2,)\n\n @nnx.scan(in_axes=(0, None), out_axes=0, length=5)\n def forward_block(block: Block, x: jax.Array):\n return block(x)\n\n x = jnp.ones((1, 3))\n y = forward_block(module, x)\n\n assert y.shape == (5, 1, 3)\n\n @parameterized.parameters(True, False)\n def test_all_carry(self, graph_updates):\n @dataclasses.dataclass\n class Foo(nnx.Module):\n n: nnx.BatchStat[int]\n\n foo = Foo(n=nnx.BatchStat(0))\n\n @nnx.scan(\n in_axes=nnx.Carry, out_axes=nnx.Carry, length=3,\n graph_updates=graph_updates,\n )\n def loop(foo: Foo):\n foo.n[...] += 1\n return foo\n\n foo2 = loop(foo)\n\n self.assertIs(foo2.n, foo.n)\n self.assertEqual(foo.n[...], 3)\n\n def test_all_carry_one_argument_error(self):\n @dataclasses.dataclass\n class Foo(nnx.Module):\n n: nnx.BatchStat[int]\n\n foo = Foo(n=nnx.BatchStat(0))\n\n @nnx.scan(in_axes=nnx.Carry, out_axes=nnx.Carry, length=3)\n def loop(foo: Foo, x): ...\n\n with self.assertRaisesRegex(\n ValueError,\n 'When in_axes=Carry, the function must take exactly one argument',\n ):\n loop(foo, 0)\n\n def test_all_carry_new_reference_error(self):\n class Foo(nnx.Module):\n def __init__(self, n: nnx.BatchStat[int]):\n self.n = n\n\n xs = jnp.arange(3)\n foo = Foo(n=nnx.BatchStat(0))\n\n @nnx.scan(in_axes=(nnx.Carry, 0), out_axes=(nnx.Carry, 0))\n def loop(foo: Foo, x):\n x = x + 1\n foo = Foo(nnx.BatchStat(foo.n[...] + 1)) # new reference\n return foo, x\n\n with self.assertRaisesRegex(\n ValueError,\n 'Carry references must be the same between iterations',\n ):\n loop(foo, xs)\n\n @parameterized.parameters(True, False)\n def test_all_scan(self, graph_updates):\n class Foo(nnx.Module):\n def __init__(self, n: nnx.BatchStat[jax.Array]):\n self.n = n\n\n xs = jnp.arange(3)\n foo = Foo(n=nnx.BatchStat(jnp.arange(3)))\n\n @nnx.scan(in_axes=0, out_axes=0, graph_updates=graph_updates)\n def loop(foo: Foo, x):\n x = x + 1\n foo.n[...] += 1\n return x\n\n ys = loop(foo, xs)\n\n np.testing.assert_allclose(ys, jnp.arange(1, 4))\n np.testing.assert_allclose(foo.n[...], jnp.arange(1, 4))\n\n def test_all_broadcast(self):\n class Foo(nnx.Module):\n def __init__(self, n: nnx.BatchStat[int]):\n self.n = n\n\n xs = jnp.array(1)\n foo = Foo(n=nnx.BatchStat(2))\n\n @nnx.scan(in_axes=None, out_axes=0, length=4)\n def loop(foo: Foo, x):\n return x + foo.n\n\n ys = loop(foo, xs)\n\n np.testing.assert_allclose(ys, 3)\n self.assertEqual(ys.shape, (4,))\n\n def test_input_output_carry_mismatch_error(self):\n with self.assertRaisesRegex(\n ValueError,\n 'If one of in_axes or out_axes has Carry, the other must also have Carry',\n ):\n\n @nnx.scan(in_axes=0, out_axes=(nnx.Carry, 0))\n def loop(a, b): ...\n\n with self.assertRaisesRegex(\n ValueError,\n 'If one of in_axes or out_axes has Carry, the other must also have Carry',\n ):\n\n @nnx.scan(in_axes=(nnx.Carry, 0), out_axes=0)\n def loop(a, b): ...\n\n def test_double_carry_error(self):\n with self.assertRaisesRegex(\n ValueError,\n 'Found multiple Carry definitions',\n ):\n\n @nnx.scan(in_axes=(nnx.Carry, nnx.Carry))\n def loop(a, b): ...\n\n def test_broadcast_in_output_error(self):\n with self.assertRaisesRegex(\n ValueError,\n 'Cannot broadcast output state',\n ):\n\n @nnx.scan(in_axes=(nnx.Carry, 0), out_axes=(nnx.Carry, None))\n def loop(a, b): ...\n\n with self.assertRaisesRegex(\n ValueError,\n 'Cannot broadcast output state. Got StateAxes',\n ):\n\n @nnx.scan(\n in_axes=(nnx.Carry, 0), out_axes=(nnx.Carry, nnx.StateAxes({...: None}))\n )\n def loop(a, b): ...\n\n @parameterized.parameters(\n (True, False), (False, False),\n )\n def test_scan_stateful(self, graph, graph_updates):\n count = nnx.Variable(jnp.array(0))\n\n @nnx.scan(graph=graph, graph_updates=graph_updates)\n def f(count, x):\n count[...] += 1\n return count, x + count[...]\n\n xs = jnp.arange(5)\n count_out, ys = f(count, xs)\n\n self.assertIs(count_out, count)\n self.assertEqual(count[...], 5)\n np.testing.assert_allclose(ys, jnp.array([1, 3, 5, 7, 9]))\n\n @parameterized.parameters(\n (True, False), (False, False),\n )\n def test_scan_carry_identity_error(self, graph, graph_updates):\n count = nnx.Variable(jnp.array(0))\n\n @nnx.scan(graph=graph, graph_updates=graph_updates)\n def f(count, x):\n new_count = nnx.Variable(count[...] + 1)\n return new_count, x\n\n with self.assertRaisesRegex(\n ValueError,\n 'scan Variable identity must be preserved',\n ):\n f(count, jnp.arange(3))\n\n def test_tree_mode_custom_axes(self):\n @nnx.scan(in_axes=nnx.Carry, out_axes=nnx.Carry, length=3, graph=False)\n def loop(x):\n return x\n\n result = loop(jnp.array(1.0))\n np.testing.assert_allclose(result, jnp.array(1.0))\n\n @parameterized.parameters(True, False)\n def test_only_carry(self, graph_updates):\n class Foo(nnx.Module):\n def __init__(self):\n self.c = nnx.BatchStat(jnp.array(0))\n\n @nnx.scan(in_axes=(nnx.Carry,), length=5, graph_updates=graph_updates)\n def loop(foo: Foo) -> tuple[Foo, jax.Array]:\n foo.c[...] += 1\n return foo, foo.c[...]\n\n foo = Foo()\n foo2, cs = loop(foo)\n self.assertIs(foo2.c, foo.c)\n np.testing.assert_allclose(cs, jnp.arange(1, 6))\n\n def test_out_axes(self):\n state_axes = nnx.StateAxes({(nnx.Param, nnx.RngState): 0, ...: None})\n\n class MLP(nnx.Module):\n @nnx.split_rngs(splits=5)\n @nnx.vmap(in_axes=(state_axes, state_axes), axis_size=5)\n def __init__(self, rngs: nnx.Rngs):\n self.linear = nnx.Linear(3, 3, rngs=rngs)\n self.node = nnx.BatchStat(jnp.ones((2,)))\n\n @nnx.scan(in_axes=(state_axes, nnx.Carry), out_axes=(nnx.Carry, 1, 2))\n def __call__(self, x: jax.Array):\n x = self.linear(x)\n x = nnx.gelu(x)\n return x, x, x\n\n module = MLP(rngs=nnx.Rngs(0))\n\n assert module.linear.kernel.shape == (5, 3, 3)\n assert module.linear.bias.shape == (5, 3)\n assert module.node.shape == (2,)\n\n x = jnp.ones((1, 3))\n c, y1, y2 = module(x)\n\n assert c.shape == (1, 3)\n assert y1.shape == (1, 5, 3)\n assert y2.shape == (1, 3, 5)\n\n def test_in_axes_simple(self):\n state_axes = nnx.StateAxes({(nnx.Param, nnx.RngState): 0, ...: None})\n\n class MLP(nnx.Module):\n @nnx.vmap(in_axes=(state_axes, 0))\n def __init__(self, key: jax.Array):\n rngs = nnx.Rngs(key)\n self.linear = nnx.Linear(3, 3, rngs=rngs)\n self.node = nnx.Variable(jnp.ones((2,)))\n\n @nnx.scan(in_axes=(state_axes, nnx.Carry), out_axes=nnx.Carry)\n def __call__(self, x: jax.Array):\n x = self.linear(x)\n x = nnx.gelu(x)\n return x\n\n key = jax.random.split(jax.random.key(0), 5)\n module = MLP(key=key)\n\n x = jnp.ones((1, 3))\n y = module(x)\n\n assert y.shape == (1, 3)\n\n def test_in_axes(self):\n state_axes = nnx.StateAxes({(nnx.Param, nnx.RngState, nnx.Intermediate): 0, ...: None})\n\n class MLP(nnx.Module):\n @nnx.split_rngs(splits=5)\n @nnx.vmap(in_axes=(state_axes, state_axes))\n def __init__(self, rngs: nnx.Rngs):\n self.linear = nnx.Linear(3, 3, rngs=rngs)\n self.node = nnx.Variable(jnp.ones((2,)))\n\n @nnx.scan(in_axes=(state_axes, nnx.Carry, 0))\n def __call__(\n self, x: jax.Array, a: jax.Array\n ) -> tp.Tuple[jax.Array, None]:\n assert x.shape == a.shape\n x = x + a\n x = self.linear(x)\n x = nnx.gelu(x)\n self.sow(nnx.Intermediate, \"data\", x)\n return x, None\n\n module = MLP(rngs=nnx.Rngs(0))\n\n assert module.linear.kernel.shape == (5, 3, 3)\n assert module.linear.bias.shape == (5, 3)\n assert module.node.shape == (2,)\n\n x = jnp.ones((1, 3))\n a = jnp.ones((5, 1, 3))\n (y, out), intermediates = nnx.capture(module, nnx.Intermediate)(x, a)\n\n assert y.shape == (1, 3)\n assert out is None\n\n assert intermediates['data'][0].shape == (5, 1, 3)\n\n def test_in_axes_broadcast(self):\n test = self\n state_axes = nnx.StateAxes({(nnx.Param, nnx.RngState): 0, ...: None})\n\n class MLP(nnx.Module):\n @nnx.split_rngs(splits=5)\n @nnx.vmap(in_axes=(state_axes, state_axes))\n def __init__(self, rngs: nnx.Rngs):\n self.linear = nnx.Linear(3, 3, rngs=rngs)\n self.node = nnx.BatchStat(jnp.ones((2,)))\n\n @nnx.scan(in_axes=(state_axes, nnx.Carry, 0, None))\n def __call__(\n self, x: jax.Array, a: jax.Array, b: jax.Array\n ) -> tp.Tuple[jax.Array, None]:\n test.assertEqual(x.shape, a.shape)\n test.assertEqual(x.shape, b.shape)\n x = x + a + b\n x = self.linear(x)\n x = nnx.gelu(x)\n return x, None\n\n module = MLP(rngs=nnx.Rngs(0))\n\n self.assertEqual(module.linear.kernel.shape, (5, 3, 3))\n self.assertEqual(module.linear.bias.shape, (5, 3))\n self.assertEqual(module.node.shape, (2,))\n\n x = jnp.ones((1, 3))\n a = jnp.ones((5, 1, 3))\n b = jnp.ones((1, 3))\n y, out = module(x, a, b)\n\n self.assertEqual(y.shape, (1, 3))\n self.assertIsNone(out)\n\n def test_complex(self):\n state_axes = nnx.StateAxes({(nnx.Param, nnx.RngState): 0, ...: None})\n\n class MLP(nnx.Module):\n @nnx.split_rngs(splits=5)\n @nnx.vmap(in_axes=(state_axes, state_axes))\n def __init__(self, rngs: nnx.Rngs):\n self.linear = nnx.Linear(3, 3, rngs=rngs)\n self.bn = nnx.BatchNorm(3, rngs=rngs)\n self.dropout = nnx.Dropout(0.5, rngs=rngs)\n self.node = nnx.Variable(jnp.ones((2,)))\n\n @nnx.scan(in_axes=(state_axes, nnx.Carry))\n def __call__(self, x: jax.Array):\n x = self.linear(x)\n x = self.bn(x)\n x = self.dropout(x)\n x = nnx.gelu(x)\n return x, None\n\n module = MLP(rngs=nnx.Rngs(0))\n module.set_attributes(deterministic=False, use_running_average=False)\n\n assert module.linear.kernel.shape == (5, 3, 3)\n assert module.linear.bias.shape == (5, 3)\n assert module.node.shape == (2,)\n\n x = jnp.ones((1, 3))\n y, _ = module(x)\n\n assert y.shape == (1, 3)\n\n def test_complex_view(self):\n state_axes = nnx.StateAxes({(nnx.Param, nnx.RngState): 0, ...: None})\n\n class MLP(nnx.Module):\n @nnx.split_rngs(splits=5)\n @nnx.vmap(in_axes=(state_axes, state_axes))\n def __init__(self, rngs: nnx.Rngs):\n self.linear = nnx.Linear(3, 3, rngs=rngs)\n self.bn = nnx.BatchNorm(3, rngs=rngs)\n self.dropout = nnx.Dropout(0.5, rngs=rngs)\n self.node = nnx.Variable(jnp.ones((2,)))\n\n @nnx.scan(in_axes=(state_axes, nnx.Carry))\n def __call__(self, x: jax.Array):\n x = self.linear(x)\n x = self.bn(x)\n x = self.dropout(x)\n x = nnx.gelu(x)\n return x, None\n\n module = MLP(rngs=nnx.Rngs(0))\n new_module = nnx.view(module, deterministic=False, use_running_average=False)\n\n assert new_module.linear.kernel.shape == (5, 3, 3)\n assert new_module.linear.bias.shape == (5, 3)\n assert new_module.node.shape == (2,)\n\n x = jnp.ones((1, 3))\n y, _ = new_module(x)\n\n assert y.shape == (1, 3)\n\n def test_complex_broadcast_dropout(self):\n state_axes = nnx.StateAxes({(nnx.Param, 'params'): 0, ...: None})\n\n class MLP(nnx.Module):\n @nnx.split_rngs(splits=5, only='params')\n @nnx.vmap(in_axes=(state_axes, state_axes))\n def __init__(self, rngs: nnx.Rngs):\n self.linear = nnx.Linear(3, 3, rngs=rngs)\n self.bn = nnx.BatchNorm(3, rngs=rngs)\n self.dropout = nnx.Dropout(0.5, rngs=rngs)\n self.node = nnx.Variable(jnp.ones((2,)))\n\n @nnx.split_rngs(splits=5, only='params')\n @nnx.scan(in_axes=(state_axes, nnx.Carry))\n def __call__(self, x: jax.Array):\n x = self.linear(x)\n x = self.bn(x)\n x = self.dropout(x)\n x = nnx.gelu(x)\n return x, None\n\n module = MLP(rngs=nnx.Rngs(params=0, dropout=1))\n module.set_attributes(deterministic=False, use_running_average=False)\n\n assert module.linear.kernel.shape == (5, 3, 3)\n assert module.linear.bias.shape == (5, 3)\n assert module.node.shape == (2,)\n\n x = jnp.ones((1, 3))\n y, _ = module(x)\n\n assert y.shape == (1, 3)\n\n def test_complex_broadcast_dropout_view(self):\n state_axes = nnx.StateAxes({(nnx.Param, 'params'): 0, ...: None})\n\n class MLP(nnx.Module):\n @nnx.split_rngs(splits=5, only='params')\n @nnx.vmap(in_axes=(state_axes, state_axes))\n def __init__(self, rngs: nnx.Rngs):\n self.linear = nnx.Linear(3, 3, rngs=rngs)\n self.bn = nnx.BatchNorm(3, rngs=rngs)\n self.dropout = nnx.Dropout(0.5, rngs=rngs)\n self.node = nnx.Variable(jnp.ones((2,)))\n\n @nnx.split_rngs(splits=5, only='params')\n @nnx.scan(in_axes=(state_axes, nnx.Carry))\n def __call__(self, x: jax.Array):\n x = self.linear(x)\n x = self.bn(x)\n x = self.dropout(x)\n x = nnx.gelu(x)\n return x, None\n\n module = MLP(rngs=nnx.Rngs(params=0, dropout=1))\n new_module = nnx.view(module, deterministic=False, use_running_average=False)\n\n assert new_module.linear.kernel.shape == (5, 3, 3)\n assert new_module.linear.bias.shape == (5, 3)\n assert new_module.node.shape == (2,)\n\n x = jnp.ones((1, 3))\n y, _ = new_module(x)\n\n assert y.shape == (1, 3)\n\n def test_complex_decorator(self):\n state_axes = nnx.StateAxes({(nnx.Param, nnx.RngState): 0, ...: None})\n\n class Block(nnx.Module):\n @nnx.split_rngs(splits=5)\n @nnx.vmap(in_axes=(state_axes, state_axes), axis_size=5)\n def __init__(self, rngs: nnx.Rngs):\n self.d = 3\n self.linear = nnx.Linear(3, 3, rngs=rngs)\n self.bn = nnx.BatchNorm(3, rngs=rngs)\n self.dropout = nnx.Dropout(0.5, rngs=rngs)\n self.node = nnx.Variable(jnp.ones((2,)))\n\n @nnx.scan(in_axes=(state_axes, nnx.Carry))\n def __call__(self, x: jax.Array):\n x = self.linear(x)\n x = self.bn(x)\n x = self.dropout(x)\n x = nnx.gelu(x)\n return x, None\n\n module = Block(rngs=nnx.Rngs(0))\n module.set_attributes(deterministic=False, use_running_average=False)\n\n assert module.d == 3\n assert module.linear.kernel.shape == (5, 3, 3)\n assert module.linear.bias.shape == (5, 3)\n assert module.node.shape == (2,)\n\n x = jnp.ones((1, 3))\n y, out = module(x)\n\n assert y.shape == (1, 3)\n assert out is None\n\n def test_complex_decorator_view(self):\n state_axes = nnx.StateAxes({(nnx.Param, nnx.RngState): 0, ...: None})\n\n class Block(nnx.Module):\n @nnx.split_rngs(splits=5)\n @nnx.vmap(in_axes=(state_axes, state_axes), axis_size=5)\n def __init__(self, rngs: nnx.Rngs):\n self.d = 3\n self.linear = nnx.Linear(3, 3, rngs=rngs)\n self.bn = nnx.BatchNorm(3, rngs=rngs)\n self.dropout = nnx.Dropout(0.5, rngs=rngs)\n self.node = nnx.Variable(jnp.ones((2,)))\n\n @nnx.scan(in_axes=(state_axes, nnx.Carry))\n def __call__(self, x: jax.Array):\n x = self.linear(x)\n x = self.bn(x)\n x = self.dropout(x)\n x = nnx.gelu(x)\n return x, None\n\n module = Block(rngs=nnx.Rngs(0))\n new_module = nnx.view(module, deterministic=False, use_running_average=False)\n\n assert new_module.d == 3\n assert new_module.linear.kernel.shape == (5, 3, 3)\n assert new_module.linear.bias.shape == (5, 3)\n assert new_module.node.shape == (2,)\n\n x = jnp.ones((1, 3))\n y, out = new_module(x)\n\n assert y.shape == (1, 3)\n assert out is None\n\n\n def test_scan_with_sharding(self):\n test = self\n state_axes = nnx.StateAxes({(nnx.Param, nnx.RngState): 0, ...: None})\n transform_metadata = {nnx.PARTITION_NAME: 'layers'}\n\n class MLP(nnx.Module):\n @nnx.split_rngs(splits=5)\n @nnx.vmap(\n in_axes=(state_axes, state_axes),\n transform_metadata=transform_metadata,\n )\n def __init__(self, rngs: nnx.Rngs):\n self.linear = nnx.Linear(\n 3,\n 3,\n kernel_init=nnx.with_metadata(\n nnx.initializers.lecun_normal(), out_sharding=('din', 'dout')\n ),\n bias_init=nnx.with_metadata(\n nnx.initializers.zeros_init(), out_sharding=('dout',)\n ),\n rngs=rngs,\n )\n\n @nnx.scan(\n in_axes=(state_axes, nnx.Carry), transform_metadata=transform_metadata\n )\n def __call__(self, x: jax.Array):\n x = self.linear(x)\n # test sharding layer axes is not present inside scan\n test.assertEqual(self.linear.kernel.shape, (3, 3))\n test.assertEqual(self.linear.kernel.out_sharding, ('din', 'dout'))\n test.assertEqual(self.linear.bias.shape, (3,))\n test.assertEqual(self.linear.bias.out_sharding, ('dout',))\n return x, None\n\n mesh = jax.make_mesh((1, 1, 1), ('layers', 'din', 'dout'), axis_types=(jax.sharding.AxisType.Auto,) * len(('layers', 'din', 'dout')))\n with jax.set_mesh(mesh):\n m = MLP(rngs=nnx.Rngs(0))\n\n # test sharding layers axes is set\n self.assertEqual(m.linear.kernel.shape, (5, 3, 3))\n self.assertEqual(m.linear.kernel.out_sharding, ('layers', 'din', 'dout'))\n self.assertEqual(m.linear.bias.shape, (5, 3))\n self.assertEqual(m.linear.bias.out_sharding, ('layers', 'dout'))\n\n x = jnp.ones((1, 3))\n with jax.set_mesh(mesh):\n y, out = m(x)\n\n # test sharding axes is preserved\n self.assertEqual(m.linear.kernel.shape, (5, 3, 3))\n self.assertEqual(m.linear.kernel.out_sharding, ('layers', 'din', 'dout'))\n self.assertEqual(m.linear.bias.shape, (5, 3))\n self.assertEqual(m.linear.bias.out_sharding, ('layers', 'dout'))\n\n def test_cache_tracing_simple(self):\n n = 0\n x = jnp.arange(5)\n count = jnp.array(0)\n\n @nnx.scan\n def f(count, x):\n nonlocal n\n n += 1\n return count + 1, x**2\n\n count, y = f(count, x)\n assert n == 1\n assert count == 5\n np.testing.assert_allclose(y, x**2)\n\n count, y = f(count, x)\n assert n == 1\n assert count == 10\n\n def test_cache_tracing_object(self):\n n = 0\n x = jnp.arange(5)\n count = jnp.array(0)\n\n class Foo(nnx.Pytree):\n @nnx.split_rngs(splits=5)\n @nnx.vmap(axis_size=5)\n def __init__(self, rngs: nnx.Rngs):\n self.x = nnx.Param(jax.random.normal(rngs(), shape=(3,)))\n\n foo = Foo(rngs=nnx.Rngs(0))\n assert foo.x.shape == (5, 3)\n\n @nnx.scan(in_axes=(nnx.Carry, 0, 0))\n def f(count, x, foo):\n nonlocal n\n n += 1\n assert foo.x.shape == (3,)\n return count + 1, x**2\n\n count, y = f(count, x, foo)\n assert n == 1\n assert count == 5\n np.testing.assert_allclose(y, x**2)\n\n count, y = f(count, x, foo)\n assert n == 1\n assert count == 10\n\n def test_scan_broadcast_keys(self):\n params_key = jax.random.split(jax.random.key(0), 3)\n rngs = nnx.Rngs(params=params_key, dropout=1)\n state_axes = nnx.StateAxes({'params': 0, ...: None})\n\n @nnx.scan(in_axes=(nnx.Carry, state_axes), length=3)\n def f(_, rngs: nnx.Rngs):\n param_key = rngs.params()\n dropout_key = rngs.dropout()\n return (), (param_key, dropout_key)\n\n _, (param_keys, dropout_keys) = f((), rngs)\n\n assert jnp.not_equal(param_keys[0], param_keys[1])\n assert jnp.not_equal(param_keys[1], param_keys[2])\n assert jnp.equal(dropout_keys[0], dropout_keys[1])\n assert jnp.equal(dropout_keys[1], dropout_keys[2])\n\n def test_rnn_example(self):\n class RNNCell(nnx.Module):\n def __init__(self, input_size, hidden_size, rngs):\n self.linear = nnx.Linear(\n hidden_size + input_size, hidden_size, rngs=rngs\n )\n self.drop = nnx.Dropout(0.1, rngs=rngs)\n self.hidden_size = hidden_size\n\n def __call__(self, carry, x) -> tuple[jax.Array, jax.Array]:\n carry = self.drop(carry) # recurrent dropout\n x = nnx.relu(self.linear(jnp.concatenate([carry, x], axis=-1)))\n return x, x\n\n def initial_state(self, batch_size: int):\n return jnp.zeros((batch_size, self.hidden_size))\n\n cell = RNNCell(20, 20, nnx.Rngs(params=0, dropout=1))\n\n state_axes = nnx.StateAxes({'dropout': None, ...: nnx.Carry})\n\n def rnn_forward(cell: RNNCell, x: jax.Array):\n carry = cell.initial_state(x.shape[0])\n\n @nnx.scan(in_axes=(state_axes, nnx.Carry, 1), out_axes=(nnx.Carry, 1))\n def unroll(cell: RNNCell, carry, x) -> tuple[jax.Array, jax.Array]:\n return cell(carry, x)\n\n _, y = unroll(cell, carry, x)\n return y\n\n x = jnp.ones((16, 10, 20))\n y = rnn_forward(cell, x)\n\n def test_carry_pytree_sow(self):\n class CarryAsPytree(nnx.Pytree):\n def __init__(self, data: jax.Array):\n self.data = data\n\n class Model(nnx.Module):\n def __init__(self, num_steps):\n self.fc = nnx.Linear(10, 10, rngs=nnx.Rngs(0))\n self.num_steps = num_steps\n\n def _step(self, state):\n new_data = state.data + 1\n self.sow(nnx.Intermediate, \"data\", new_data)\n state.data = new_data\n return state\n\n def _step2(self, state: tuple[CarryAsPytree, jax.Array, CarryAsPytree]):\n out = self.fc(state[1])\n\n new_data1 = state[0].data + 1\n new_data2 = state[2].data + 1\n\n self.sow(nnx.Intermediate, \"data1\", new_data1)\n self.sow(nnx.Intermediate, \"data2\", new_data2)\n\n state[0].data = new_data1\n state[2].data = new_data2\n return (state[0], out, state[2])\n\n @nnx.jit(static_argnames=(\"method\"))\n def __call__(self, state, method):\n state_axes = nnx.StateAxes({nnx.Intermediate: 0, ...: nnx.Carry})\n state_final = nnx.scan(\n method,\n in_axes=(state_axes, nnx.Carry),\n out_axes=nnx.Carry,\n length=self.num_steps,\n )(self, state)\n\n return state_final\n\n num_steps = 5\n model = Model(num_steps=num_steps)\n carry = CarryAsPytree(data=jnp.array(0.0))\n carry_final, intermediates = nnx.capture(model, nnx.Intermediate)(carry, method=Model._step)\n self.assertEqual(carry_final.data, num_steps)\n np.testing.assert_array_equal(\n intermediates['data'][0], 1.0 + jnp.arange(num_steps)\n )\n\n carry = (\n CarryAsPytree(data=jnp.array(0.0)),\n jnp.ones((10,)),\n CarryAsPytree(data=jnp.array(10.0))\n )\n\n carry_final, intermediates = nnx.capture(model, nnx.Intermediate)(carry, method=Model._step2)\n\n self.assertEqual(carry_final[0].data, num_steps)\n self.assertEqual(carry_final[2].data, 10 + num_steps)\n np.testing.assert_array_equal(\n intermediates['data1'][0], 1.0 + jnp.arange(num_steps)\n )\n np.testing.assert_array_equal(\n intermediates['data2'][0], 11.0 + jnp.arange(num_steps)\n )\n\n def test_broadcast_variable_mutation_rejected(self):\n v = nnx.Variable(jnp.array(1.0))\n\n @nnx.scan(\n in_axes=(None, nnx.Carry, 0), graph=False, graph_updates=False,\n )\n def fn(v, carry, x):\n v[...] = v[...] + 1.0\n return carry + x, carry\n\n with self.assertRaisesRegex(ValueError, 'Broadcast.*mutated'):\n fn(v, jnp.array(0.0), jnp.arange(3.0))\n\n def test_broadcast_out_axes_rejected(self):\n @nnx.scan(\n in_axes=(nnx.Carry, 0), out_axes=(nnx.Carry, None),\n graph=False, graph_updates=False,\n )\n def fn(carry, x):\n return carry + x, jnp.zeros(3)\n\n with self.assertRaisesRegex(ValueError, 'broadcast'):\n fn(jnp.array(0.0), jnp.arange(3.0))\n\n def test_scan_inconsistent_aliasing(self):\n v = nnx.Param(jnp.array(0.0))\n\n @nnx.scan(\n in_axes=(nnx.Carry, 0),\n out_axes=(nnx.Carry, 0),\n graph=True,\n graph_updates=False,\n )\n def f(carry, x):\n return carry, x[...]\n\n with self.assertRaisesRegex(ValueError, 'Inconsistent aliasing'):\n f(v, v)\n\n def test_scan_input_output_aliasing(self):\n v = nnx.Param(jnp.arange(5))\n\n @nnx.scan(in_axes=0, out_axes=0, graph=True, graph_updates=False)\n def f(carry):\n return carry\n\n with self.assertRaisesRegex(ValueError, 'does not support returning input Variables as outputs'):\n f(v)\n\n\nclass TestRemat(parameterized.TestCase):\n @parameterized.parameters(\n (True, True),\n (True, False),\n (False, False),\n )\n def test_remat_basic(self, graph, graph_updates):\n class RematLinear(nnx.Module):\n def __init__(self, din: int, dout: int, rngs: nnx.Rngs):\n self.linear = nnx.Linear(din, dout, rngs=rngs)\n\n @nnx.remat(graph=graph, graph_updates=graph_updates)\n def __call__(self, x: jax.Array) -> jax.Array:\n return self.linear(x)\n\n module = RematLinear(2, 3, nnx.Rngs(0))\n\n def loss_fn(module, x):\n y = module(x)\n return jnp.sum(y)\n\n grad_type = nnx.State if graph and graph_updates else RematLinear\n loss, grads = nnx.value_and_grad(\n loss_fn, graph=graph, graph_updates=graph_updates,\n )(module, jnp.ones((1, 2)))\n\n assert loss.shape == ()\n assert isinstance(grads, grad_type)\n\n @parameterized.parameters(\n (True, True),\n (True, False),\n (False, False),\n )\n def test_remat_variables(self, graph, graph_updates):\n rngs = nnx.Rngs(0)\n w = nnx.Param(jax.random.normal(rngs(), (2, 3)))\n b = nnx.Param(jax.random.normal(rngs(), (3,)))\n count = nnx.BatchStat(jnp.array(0))\n\n @nnx.remat(graph=graph, graph_updates=graph_updates)\n def linear(w, b, count, x):\n count[...] += 1\n return x @ w + b[None]\n\n def loss_fn(w, b, count, x):\n return jnp.sum(linear(w, b, count, x))\n\n x = jnp.ones((1, 2))\n loss, grads = nnx.value_and_grad(\n loss_fn, argnums=(0, 1), graph=graph, graph_updates=graph_updates,\n )(w, b, count, x)\n\n assert loss.shape == ()\n assert isinstance(grads, tuple)\n assert len(grads) == 2\n assert count[...] == 1\n\n def test_remat_with_scan_decorator(self):\n state_axes = nnx.StateAxes({(nnx.Param, nnx.RngState): 0, ...: None})\n\n class ScanLinear(nnx.Module):\n @nnx.split_rngs(splits=5)\n @nnx.vmap(in_axes=(state_axes, state_axes), axis_size=5)\n def __init__(self, rngs: nnx.Rngs):\n self.linear = nnx.Linear(3, 3, rngs=rngs)\n\n @nnx.scan(in_axes=(state_axes, nnx.Carry))\n @nnx.remat\n def __call__(self, x: jax.Array) -> tp.Tuple[jax.Array, None]:\n x = self.linear(x)\n return x, None\n\n m = ScanLinear(nnx.Rngs(0))\n\n assert m.linear.kernel.shape == (5, 3, 3)\n assert m.linear.bias.shape == (5, 3)\n\n y, _ = m(jnp.ones((1, 3)))\n assert y.shape == (1, 3)\n\n @parameterized.parameters(\n (True, False),\n (False, False),\n )\n def test_tree_mode_remat_basic(self, graph, graph_updates):\n model = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n\n @nnx.remat(graph=graph, graph_updates=graph_updates)\n def forward(model, x):\n return model(x)\n\n def loss_fn(model, x):\n y = forward(model, x)\n return jnp.sum(y)\n\n grads = nnx.grad(\n loss_fn, graph=graph, graph_updates=graph_updates,\n )(model, jnp.ones((1, 2)))\n assert grads.kernel.shape == (2, 3)\n assert grads.bias.shape == (3,)\n\n @parameterized.parameters(\n (True, False),\n (False, False),\n )\n def test_tree_mode_remat_stateful(self, graph, graph_updates):\n class Counter(nnx.Variable):\n pass\n\n class Linear(nnx.Module):\n def __init__(self, din, dout, *, rngs):\n key = rngs.params()\n self.w = nnx.Param(jax.random.uniform(key, (din, dout)))\n self.b = nnx.Param(jnp.zeros((dout,)))\n self.count = Counter(jnp.array(0))\n\n def __call__(self, x):\n self.count[...] += 1\n return x @ self.w + self.b[None]\n\n model = Linear(2, 3, rngs=nnx.Rngs(0))\n\n @nnx.remat(graph=graph, graph_updates=graph_updates)\n def forward(model, x):\n return model(x)\n\n y = forward(model, jnp.ones((1, 2)))\n assert y.shape == (1, 3)\n assert model.count[...] == 1\n\n\nclass TestVmap(parameterized.TestCase):\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_vmap_basic(self, graph, graph_updates):\n class LinearEnsemble(nnx.Module):\n def __init__(self, num, *, rngs):\n self.w = nnx.Param(jax.random.uniform(rngs(), (num, 2, 3)))\n\n model = LinearEnsemble(5, rngs=nnx.Rngs(0))\n x = jnp.ones((2,))\n\n @nnx.vmap(in_axes=(0, None), out_axes=0, graph=graph, graph_updates=graph_updates)\n def forward(model, x):\n return x @ model.w\n\n y = forward(model, x)\n assert y.shape == (5, 3)\n\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_vmap_stateful(self, graph, graph_updates):\n class Counter(nnx.Variable):\n pass\n\n class Linear(nnx.Module):\n def __init__(self, din, dout, *, rngs):\n key = rngs.params()\n self.w = nnx.Param(jax.random.uniform(key, (din, dout)))\n self.count = Counter(jnp.array(0))\n\n def __call__(self, x):\n self.count[...] += 1\n return x @ self.w\n\n model = Linear(2, 3, rngs=nnx.Rngs(0))\n\n @nnx.vmap(in_axes=(None, 0), out_axes=0, graph=graph, graph_updates=graph_updates)\n def forward(model, x):\n return model(x)\n\n x = jnp.ones((5, 2))\n y = forward(model, x)\n assert y.shape == (5, 3)\n assert model.count[...] == 1\n\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_vmap_variables(self, graph, graph_updates):\n rngs = nnx.Rngs(0)\n w = nnx.Param(jax.random.normal(rngs(), (5, 2, 3)))\n b = nnx.Param(jax.random.normal(rngs(), (5, 3)))\n\n @nnx.vmap(in_axes=(0, 0, 1), out_axes=1, graph=graph, graph_updates=graph_updates)\n def forward(w, b, x):\n return x @ w + b\n\n x = jax.random.uniform(rngs(), (2, 5))\n y = forward(w, b, x)\n assert y.shape == (3, 5)\n\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_vmap_ensemble_forward(self, graph, graph_updates):\n class Linear(nnx.Module):\n def __init__(self, din, dout, *, rngs):\n key = rngs.params()\n self.w = nnx.Param(jax.random.uniform(key, (din, dout)))\n self.b = nnx.Param(jnp.zeros((dout,)))\n\n def __call__(self, x):\n return x @ self.w + self.b[None]\n\n @nnx.vmap(in_axes=0, out_axes=0, graph=graph, graph_updates=graph_updates)\n def create_ensemble(keys):\n return Linear(2, 3, rngs=nnx.Rngs(keys))\n\n keys = jax.random.split(jax.random.key(0), 5)\n ensemble = create_ensemble(keys)\n\n assert ensemble.w.shape == (5, 2, 3)\n assert ensemble.b.shape == (5, 3)\n\n @nnx.vmap(in_axes=(0, None), out_axes=0, graph=graph, graph_updates=graph_updates)\n def forward(model, x):\n return model(x)\n\n x = jnp.ones((1, 2))\n y = forward(ensemble, x)\n assert y.shape == (5, 1, 3)\n\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_vmap_replicate(self, graph, graph_updates):\n model = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n\n @nnx.vmap(in_axes=(None, 0), out_axes=0, graph=graph, graph_updates=graph_updates)\n def forward(model, x):\n return model(x)\n\n x = jnp.ones((5, 1, 2))\n y = forward(model, x)\n assert y.shape == (5, 1, 3)\n\n def test_basic(self):\n @nnx.split_rngs(splits=5)\n @nnx.vmap(in_axes=0, out_axes=0, axis_size=5)\n def create_block(rngs: nnx.Rngs):\n return nnx.Linear(2, 3, rngs=rngs)\n\n rngs = nnx.Rngs(0)\n\n block = create_block(rngs)\n\n self.assertEqual(block.kernel.shape, (5, 2, 3))\n self.assertEqual(rngs.default.count[...], 1)\n\n @nnx.vmap(in_axes=(0, 1), out_axes=1)\n def forward(block: nnx.Linear, x):\n self.assertEqual(block.kernel.shape, (2, 3))\n self.assertEqual(block.bias.shape, (3,))\n self.assertEqual(x.shape, (2,))\n return block(x)\n\n x = jax.random.uniform(rngs(), (2, 5))\n y = forward(block, x)\n\n self.assertEqual(y.shape, (3, 5))\n\n def test_basic_variables(self):\n @nnx.split_rngs(splits=5)\n @nnx.vmap(in_axes=0, out_axes=0, axis_size=5)\n def create_block(rngs: nnx.Rngs):\n w = nnx.Param(jax.random.normal(rngs(), (2, 3)))\n b = nnx.Param(jax.random.normal(rngs(), (3,)))\n return w, b\n\n rngs = nnx.Rngs(0)\n w, b = create_block(rngs)\n\n self.assertEqual(w.shape, (5, 2, 3))\n self.assertEqual(b.shape, (5, 3))\n self.assertEqual(rngs.default.count[...], 1)\n\n @nnx.vmap(in_axes=(0, 0, 1), out_axes=1)\n def forward(w, b, x):\n self.assertEqual(w.shape, (2, 3))\n self.assertEqual(b.shape, (3,))\n self.assertEqual(x.shape, (2,))\n return x @ w + b\n\n x = jax.random.uniform(rngs(), (2, 5))\n y = forward(w, b, x)\n\n self.assertEqual(y.shape, (3, 5))\n\n def test_state_axes(self):\n class Block(nnx.Module):\n def __init__(self, rngs: nnx.Rngs):\n self.linear = nnx.Linear(3, 3, rngs=rngs)\n self.dropout = nnx.Dropout(0.5, deterministic=False, rngs=rngs)\n\n def __call__(self, x: jax.Array) -> jax.Array:\n x = self.linear(x)\n x = nnx.relu(x)\n x = self.dropout(x)\n return x\n\n @nnx.vmap(\n in_axes=0,\n out_axes=nnx.StateAxes({(nnx.Param, nnx.RngState): 0, ...: None}),\n )\n def create_block(rngs: nnx.Rngs):\n rngs = nnx.clone(rngs)\n return Block(rngs)\n\n rngs = nnx.Rngs(0)\n initial_key = rngs.default.key[...]\n\n backups = nnx.split_rngs(rngs, splits=5)\n module = create_block(rngs)\n nnx.restore_rngs(backups)\n\n assert rngs.default.count[...] == 1\n assert rngs.default.key[...] == initial_key\n assert not jnp.allclose(\n module.linear.kernel[0],\n module.linear.kernel[1],\n )\n assert module.linear.kernel.shape == (5, 3, 3)\n assert module.linear.bias.shape == (5, 3)\n\n x = jnp.ones((5, 1, 3))\n\n @nnx.vmap(\n in_axes=(nnx.StateAxes({(nnx.Param, nnx.RngState): 0, ...: None}), 0),\n )\n def forward_block(module, x):\n return module(x)\n\n backups = nnx.split_rngs(rngs, splits=5)\n y = forward_block(module, x)\n nnx.restore_rngs(backups)\n\n assert y.shape == (5, 1, 3)\n assert rngs.default.count[...] == 2\n assert rngs.default.key[...] == initial_key\n\n y2 = forward_block(module, x)\n\n assert not jnp.allclose(y, y2)\n\n def test_split_rngs_context_manager(self):\n class Block(nnx.Module):\n def __init__(self, rngs: nnx.Rngs):\n self.linear = nnx.Linear(3, 3, rngs=rngs)\n self.dropout = nnx.Dropout(0.5, deterministic=False, rngs=rngs)\n\n def __call__(self, x: jax.Array) -> jax.Array:\n x = self.linear(x)\n x = nnx.relu(x)\n x = self.dropout(x)\n return x\n\n state_axes = nnx.StateAxes({(nnx.Param, nnx.RngState): 0, ...: None})\n\n @nnx.vmap(in_axes=(state_axes,), out_axes=state_axes)\n def create_block(rngs: nnx.Rngs):\n return Block(rngs)\n\n rngs = nnx.Rngs(0)\n initial_key = rngs.default.key[...]\n\n module = create_block(rngs.split(5))\n\n assert rngs.default.count[...] == 1\n assert rngs.default.key[...] == initial_key\n assert not jnp.allclose(\n module.linear.kernel[0],\n module.linear.kernel[1],\n )\n assert module.linear.kernel.shape == (5, 3, 3)\n assert module.linear.bias.shape == (5, 3)\n\n x = jnp.ones((5, 1, 3))\n\n @nnx.vmap(in_axes=(state_axes, 0))\n def forward_block(module, x):\n return module(x)\n\n y = forward_block(module, x)\n\n assert y.shape == (5, 1, 3)\n assert rngs.default.key[...] == initial_key\n\n y2 = forward_block(module, x)\n\n assert not jnp.allclose(y, y2)\n\n def test_split_rngs_decorator(self):\n class Block(nnx.Module):\n def __init__(self, rngs: nnx.Rngs):\n self.linear = nnx.Linear(3, 3, rngs=rngs)\n self.dropout = nnx.Dropout(0.5, deterministic=False, rngs=rngs)\n\n def __call__(self, x: jax.Array) -> jax.Array:\n x = self.linear(x)\n x = nnx.relu(x)\n x = self.dropout(x)\n return x\n\n state_axes = nnx.StateAxes({(nnx.Param, nnx.RngState): 0, ...: None})\n\n @nnx.split_rngs(splits=5)\n @nnx.vmap(in_axes=(state_axes,), out_axes=state_axes)\n def create_block(rngs: nnx.Rngs):\n return Block(rngs)\n\n rngs = nnx.Rngs(0)\n initial_key = rngs.default.key[...]\n\n module = create_block(rngs)\n\n assert rngs.default.count[...] == 1\n assert rngs.default.key[...] == initial_key\n assert not jnp.allclose(\n module.linear.kernel[0],\n module.linear.kernel[1],\n )\n assert module.linear.kernel.shape == (5, 3, 3)\n assert module.linear.bias.shape == (5, 3)\n\n x = jnp.ones((5, 1, 3))\n\n @nnx.vmap(in_axes=(state_axes, 0))\n def forward_block(module, x):\n self.assertEqual(x.shape, (1, 3))\n return module(x)\n\n y = forward_block(module, x)\n\n assert y.shape == (5, 1, 3)\n assert rngs.default.key[...] == initial_key\n\n y2 = forward_block(module, x)\n\n assert not jnp.allclose(y, y2)\n\n def test_state_axes_simple(self):\n class Block(nnx.Module):\n def __init__(self, rngs: nnx.Rngs):\n self.linear = nnx.Linear(2, 3, rngs=rngs)\n self.bn = nnx.BatchNorm(3, rngs=rngs)\n self.dropout = nnx.Dropout(0.1, deterministic=False, rngs=rngs)\n\n def __call__(self, x: jax.Array) -> jax.Array:\n return nnx.relu(self.dropout(self.bn(self.linear(x))))\n\n state_axes = nnx.StateAxes({(nnx.BatchStat, 'dropout'): 0, ...: None})\n\n @nnx.vmap(in_axes=(state_axes,), out_axes=state_axes)\n def create_block(rngs: nnx.Rngs):\n return Block(rngs)\n\n rngs = nnx.Rngs(params=0, dropout=1)\n nnx.split_rngs(rngs, splits=5, only='dropout')\n\n module = create_block(rngs)\n\n assert module.linear.kernel.shape == (2, 3)\n assert module.bn.scale.shape == (3,)\n assert module.bn.mean.shape == (5, 3)\n\n @nnx.vmap(in_axes=(state_axes, 0), out_axes=0)\n def forward_block(module, x):\n return module(x)\n\n x = jnp.ones((5, 1, 2))\n y = forward_block(module, x)\n\n assert y.shape == (5, 1, 3)\n\n def test_split_rngs_decorator_simple(self):\n class Block(nnx.Module):\n def __init__(self, rngs: nnx.Rngs):\n self.linear = nnx.Linear(2, 3, rngs=rngs)\n self.bn = nnx.BatchNorm(3, rngs=rngs)\n self.dropout = nnx.Dropout(0.1, deterministic=False, rngs=rngs)\n\n def __call__(self, x: jax.Array) -> jax.Array:\n return nnx.relu(self.dropout(self.bn(self.linear(x))))\n\n state_axes = nnx.StateAxes({(nnx.BatchStat, 'dropout'): 0, ...: None})\n\n @nnx.split_rngs(splits=5, only='dropout')\n @nnx.vmap(in_axes=(state_axes,), out_axes=state_axes)\n def create_block(rngs: nnx.Rngs):\n return Block(rngs)\n\n rngs = nnx.Rngs(params=0, dropout=1)\n\n module = create_block(rngs)\n\n assert module.linear.kernel.shape == (2, 3)\n assert module.bn.scale.shape == (3,)\n assert module.bn.mean.shape == (5, 3)\n assert module.dropout.rngs is not None\n self.assertEqual(module.dropout.rngs.key.shape, (5,))\n\n @nnx.vmap(in_axes=(state_axes, 0), out_axes=0)\n def forward_block(module: Block, x):\n assert module.dropout.rngs is not None\n self.assertEqual(module.dropout.rngs.key.shape, ())\n return module(x)\n\n x = jnp.ones((5, 1, 2))\n y = forward_block(module, x)\n\n assert module.dropout.rngs is not None\n self.assertEqual(module.dropout.rngs.key.shape, (5,))\n assert y.shape == (5, 1, 3)\n\n def test_state_axes_super_simple(self):\n class Block(nnx.Module):\n def __init__(self, rngs: nnx.Rngs):\n self.linear = nnx.Linear(2, 3, rngs=rngs)\n self.bn = nnx.BatchNorm(3, rngs=rngs)\n self.dropout = nnx.Dropout(0.1, deterministic=False, rngs=rngs)\n\n def __call__(self, x: jax.Array) -> jax.Array:\n return nnx.relu(self.dropout(self.bn(self.linear(x))))\n\n @nnx.vmap(in_axes=0, out_axes=0)\n def create_block(rngs: nnx.Rngs):\n return Block(rngs)\n\n rngs = nnx.Rngs(0)\n nnx.split_rngs(rngs, splits=5)\n\n module = create_block(rngs)\n\n assert module.linear.kernel.shape == (5, 2, 3)\n assert module.bn.scale.shape == (5, 3)\n assert module.bn.mean.shape == (5, 3)\n\n @nnx.vmap(in_axes=(0, 0), out_axes=0)\n def forward_block(module, x):\n return module(x)\n\n x = jnp.ones((5, 1, 2))\n y = forward_block(module, x)\n\n assert y.shape == (5, 1, 3)\n\n def test_replicate(self):\n din = 3\n dout = 10\n\n class Block(nnx.Module):\n def __init__(self, rngs: nnx.Rngs):\n self.linear = nnx.Linear(din, dout, rngs=rngs)\n self.dropout = nnx.Dropout(0.5, deterministic=False, rngs=rngs)\n\n def __call__(self, x: jax.Array) -> jax.Array:\n return self.dropout(nnx.relu(self.linear(x)))\n\n def create_block(rngs: nnx.Rngs):\n return Block(rngs)\n\n state_axes = nnx.StateAxes({nnx.RngState: 0, ...: None})\n\n @nnx.split_rngs(splits=5)\n @partial(nnx.vmap, in_axes=(state_axes, 0), out_axes=0)\n def forward_block(module: Block, x):\n return module(x)\n\n rngs = nnx.Rngs(0)\n module = create_block(rngs)\n initial_key = module.dropout.rngs.key[...]\n\n assert module.dropout.rngs.count[...] == 0\n assert module.linear.kernel.shape == (din, dout)\n assert module.linear.bias.shape == (dout,)\n\n x = jnp.ones((5, 1, din))\n\n y = forward_block(module, x)\n\n assert y.shape == (5, 1, dout)\n assert module.dropout.rngs.count[...] == 1\n\n assert not jnp.allclose(y[0], y[1])\n\n y2 = forward_block(module, x)\n\n # dropout is working!\n assert not jnp.allclose(y, y2)\n\n assert module.dropout.rngs.key[...] == initial_key\n\n def test_consistent_aliasing_inputs(self):\n class Foo(nnx.Module):\n def __init__(self):\n self.a = nnx.Param(jnp.zeros((5, 5)))\n\n m = Foo()\n\n @nnx.vmap(in_axes=(0, 1))\n def f(m1, m2):\n pass\n\n with self.assertRaisesRegex(ValueError, 'Inconsistent aliasing detected'):\n f(m, m)\n\n def test_consistent_aliasing_input_output(self):\n class Foo(nnx.Module):\n def __init__(self):\n self.a = nnx.Param(jnp.zeros((2, 3)))\n\n m = Foo()\n\n @partial(nnx.vmap, in_axes=0, out_axes=1)\n def f(m):\n return m\n\n with self.assertRaisesRegex(ValueError, 'Inconsistent aliasing detected'):\n m2 = f(m)\n\n def test_consistent_aliasing_shared(self):\n class Shared(nnx.Module):\n def __init__(self):\n self.a = nnx.Param(jnp.zeros((3, 3)))\n\n class Foo(nnx.Module):\n def __init__(self, shared: Shared):\n self.a = shared\n\n shared = Shared()\n m1 = Foo(shared)\n m2 = Foo(shared)\n\n @nnx.vmap(in_axes=(0, 1))\n def f(m1, m2):\n pass\n\n with self.assertRaisesRegex(\n ValueError,\n r'Inconsistent aliasing detected([\\s\\S]*)Param([\\s\\S]*)a:'\n r' 0([\\s\\S]*)a: 1',\n ):\n f(m1, m2)\n\n def test_equivalent_state_axes_mapping(self):\n m = nnx.Linear(3, 3, rngs=nnx.Rngs(0))\n\n sa1 = nnx.StateAxes({...: 0})\n sa2 = nnx.StateAxes({nnx.Param: 0})\n\n @nnx.vmap(in_axes=(0, sa1, sa2))\n def f(m1, m2, m3):\n pass\n\n f(m, m, m)\n\n def test_equivalent_state_sharding_mapping(self):\n m = nnx.Linear(4, 4, rngs=nnx.Rngs(0))\n\n mesh = jax.sharding.Mesh(jax.devices(), ('mp',))\n sharding = jax.sharding.NamedSharding(\n mesh, jax.sharding.PartitionSpec('mp')\n )\n\n sa1 = nnx.StateSharding({...: sharding})\n sa2 = nnx.StateSharding({nnx.Param: sharding})\n\n @nnx.jit(in_shardings=(sharding, sa1, sa2))\n def f(m1, m2, m3):\n pass\n\n f(m, m, m)\n\n def test_captured_module_in_return_error(self):\n class Foo(nnx.Module):\n def __init__(self):\n self.a = jnp.zeros((4, 4))\n\n m = Foo()\n\n @nnx.vmap(in_axes=0, out_axes=0)\n def f(x):\n return x, m\n\n with self.assertRaisesRegex(\n errors.TraceContextError,\n r'Trying to extract graph node from different trace level.*Foo',\n ):\n x = jnp.zeros((4,))\n f(x)\n\n def test_vmap_and_cond_passthrough(self):\n class Broadcast(nnx.Variable[nnx.A]): ...\n\n class Vectorized(nnx.Variable[nnx.A]): ...\n\n class Env(nnx.Module):\n def __init__(self):\n self.broadcast = Broadcast(jnp.array(1))\n self.index = Vectorized(jnp.arange(8))\n self.step = Vectorized(jnp.zeros((8,), jnp.uint32))\n\n env = Env()\n\n @nnx.vmap(in_axes=(nnx.StateAxes({Broadcast: None, Vectorized: 0}),))\n def f(env: Env):\n self.assertEqual(env.step.shape, ())\n\n def increment(env: Env):\n env.step[...] += 1\n\n def no_nothing(env: Env):\n pass\n\n is_even = env.index % 2 == 0\n nnx.cond(is_even, increment, no_nothing, env)\n\n f(env)\n\n np.testing.assert_array_equal(env.step[...], [1, 0, 1, 0, 1, 0, 1, 0])\n\n def test_vmap_and_cond_passthrough_error(self):\n class Broadcast(nnx.Variable[nnx.A]): ...\n\n class Vectorized(nnx.Variable[nnx.A]): ...\n\n class Env(nnx.Module):\n def __init__(self):\n self.broadcast = Broadcast(jnp.array(1))\n self.index = Vectorized(jnp.arange(8))\n self.step = Vectorized(jnp.zeros((8,), jnp.uint32))\n\n env = Env()\n\n @nnx.vmap(in_axes=(nnx.StateAxes({Broadcast: None, Vectorized: 0}),))\n def f(env: Env):\n self.assertEqual(env.step.shape, ())\n\n def increment(env: Env):\n env.step[...] += 1\n env.broadcast[...] += 1\n\n def no_nothing(env: Env):\n pass\n\n is_even = env.index % 2 == 0\n nnx.cond(is_even, increment, no_nothing, env)\n\n with self.assertRaisesRegex(\n ValueError,\n r\"at vmap.*'broadcast'.*got axis spec None but output was batched on\"\n r' axis 0',\n ):\n f(env)\n\n def test_example(self):\n class Model(nnx.Module):\n def __init__(self, din, dout, *, rngs: nnx.Rngs):\n self.linear = nnx.Linear(din, dout, rngs=rngs)\n self.dropout = nnx.Dropout(0.5, rngs=rngs)\n self.bn = nnx.BatchNorm(dout, rngs=rngs)\n\n def __call__(self, x):\n return nnx.relu(self.dropout(self.bn(self.linear(x))))\n\n @nnx.vmap(in_axes=0, out_axes=0)\n def initialize_ensemble(key):\n rngs = nnx.Rngs(key)\n return Model(2, 3, rngs=rngs)\n\n keys = jax.random.split(jax.random.key(0), 5)\n ensemble = initialize_ensemble(keys)\n\n self.assertEqual(ensemble.linear.kernel.shape, (5, 2, 3))\n\n @nnx.vmap(in_axes=(0, None), out_axes=0)\n def forward(model, x):\n return model(x)\n\n x = jnp.ones((4, 2))\n y = forward(ensemble, x)\n self.assertEqual(y.shape, (5, 4, 3))\n\n def test_example_with_vectorization(self):\n class LinearEnsemble(nnx.Module):\n def __init__(self, num, rngs):\n self.w = nnx.Param(jax.random.uniform(rngs(), (num, 2, 3)))\n\n model = LinearEnsemble(5, rngs=nnx.Rngs(0))\n\n @nnx.vmap(in_axes=(0, None), out_axes=0)\n def forward(model, x):\n self.assertEqual(model.w.shape, (2, 3))\n return jnp.dot(x, model.w)\n\n x = jnp.ones((4, 2))\n y = forward(model, x)\n\n self.assertEqual(y.shape, (5, 4, 3))\n\n def test_metadata(self):\n @nnx.vmap(\n in_axes=(None,),\n out_axes=0,\n axis_size=5,\n transform_metadata={nnx.spmd.PARTITION_NAME: 'c'},\n )\n def create_block(rngs: nnx.Rngs):\n return nnx.Linear(\n 16,\n 32,\n rngs=rngs,\n kernel_init=nnx.with_partitioning(\n nnx.initializers.lecun_normal(), ('a', 'b')\n ),\n )\n\n\n mesh = jax.make_mesh((1, 1, 1), ('a', 'b', 'c'), axis_types=(jax.sharding.AxisType.Auto,) * len(('a', 'b', 'c')))\n with jax.set_mesh(mesh):\n m = create_block(nnx.Rngs(0))\n self.assertEqual(m.kernel.shape, (5, 16, 32))\n self.assertEqual(m.kernel.out_sharding, ('c', 'a', 'b'))\n\n def test_state_axes_from_state(self):\n class Model(nnx.Module):\n def __init__(self, din, dout, *, rngs: nnx.Rngs):\n self.linear = nnx.Linear(din, dout, rngs=rngs)\n self.bn = nnx.BatchNorm(dout, rngs=rngs)\n\n model = Model(2, 3, rngs=nnx.Rngs(0))\n state = nnx.state(model)\n\n state['linear']['kernel'] = 0\n state['linear']['bias'] = 1\n state['bn']['scale'] = 0\n state['bn']['mean'] = 1\n state['bn']['var'] = 0\n state['bn']['bias'] = None\n\n state_axes = nnx.StateAxes(state)\n\n self.assertEqual(state_axes.map_prefix(('linear', 'kernel'), None), 0)\n self.assertEqual(state_axes.map_prefix(('linear', 'bias'), None), 1)\n self.assertEqual(state_axes.map_prefix(('bn', 'scale'), None), 0)\n self.assertEqual(state_axes.map_prefix(('bn', 'mean'), None), 1)\n self.assertEqual(state_axes.map_prefix(('bn', 'var'), None), 0)\n self.assertEqual(state_axes.map_prefix(('bn', 'bias'), None), None)\n\n @nnx.vmap(out_axes=state_axes, axis_size=5)\n def create_block():\n return Model(2, 3, rngs=nnx.Rngs(0))\n\n model = create_block()\n\n self.assertEqual(model.linear.kernel.shape, (5, 2, 3))\n self.assertEqual(model.linear.bias.shape, (3, 5))\n self.assertEqual(model.bn.scale.shape, (5, 3))\n self.assertEqual(model.bn.mean.shape, (3, 5))\n self.assertEqual(model.bn.var.shape, (5, 3))\n self.assertEqual(model.bn.bias.shape, (3,))\n\n\n def test_vmap_inconsistent_aliasing(self):\n v = nnx.Param(jnp.arange(3.0))\n\n @nnx.vmap(in_axes=(0, None), graph=True, graph_updates=False)\n def f(v_mapped, v_broadcast):\n return v_mapped[...] + v_broadcast[...]\n\n with self.assertRaisesRegex(ValueError, 'Inconsistent aliasing'):\n f(v, v)\n\n\nclass TestPmap(parameterized.TestCase):\n def test_basic_single(self):\n class Block(nnx.Module):\n def __init__(self, rngs: nnx.Rngs):\n self.linear = nnx.Linear(3, 10, rngs=rngs)\n self.dropout = nnx.Dropout(0.5, deterministic=False, rngs=rngs)\n\n def __call__(self, x: jax.Array) -> jax.Array:\n x = self.linear(x)\n x = nnx.elu(x)\n x = self.dropout(x)\n return x\n\n state_axes = nnx.StateAxes({(nnx.Param, nnx.RngState): 0, ...: None})\n\n @nnx.split_rngs(splits=1)\n @nnx.pmap(in_axes=(state_axes,), out_axes=state_axes, axis_size=1,\n graph=True)\n def create_block(rngs: nnx.Rngs):\n return Block(rngs)\n\n rngs = nnx.Rngs(0)\n module = create_block(rngs)\n initial_key = module.dropout.rngs.key[...]\n\n assert module.dropout.rngs.count[0] == 0\n assert module.linear.kernel.shape == (1, 3, 10)\n assert module.linear.bias.shape == (1, 10)\n\n x = jnp.ones((1, 1, 3))\n\n @nnx.pmap(in_axes=(state_axes, 0), axis_size=1, graph=True)\n def forward_block(module, x):\n return module(x)\n\n y = forward_block(module, x)\n\n assert y.shape == (1, 1, 10)\n assert module.dropout.rngs.count[0] == 1\n assert module.dropout.rngs.key[...] == initial_key\n\n y2 = forward_block(module, x)\n\n assert not jnp.allclose(y, y2)\n\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_basic_demo_single(self, graph, graph_updates):\n class Block(nnx.Module):\n def __init__(self, rngs: nnx.Rngs):\n self.linear = nnx.Linear(20, 20, rngs=rngs)\n self.dropout = nnx.Dropout(0.2, deterministic=False, rngs=rngs)\n\n def __call__(self, x: jax.Array) -> jax.Array:\n return self.dropout(nnx.relu(self.linear(x)))\n\n @nnx.split_rngs(splits=1)\n @nnx.pmap(axis_size=1, graph=graph, graph_updates=graph_updates)\n def create_block(rngs: nnx.Rngs):\n return Block(rngs)\n\n @nnx.pmap(axis_size=1, graph=graph, graph_updates=graph_updates)\n def forward_block(module: Block, x):\n return module(x)\n\n rngs = nnx.Rngs(0)\n module = create_block(rngs)\n\n assert module.dropout.rngs.count[...] == 0\n assert module.linear.kernel.shape == (1, 20, 20)\n assert module.linear.bias.shape == (1, 20)\n\n x = jnp.ones((1, 10, 20))\n\n y = forward_block(module, x)\n\n assert y.shape == (1, 10, 20)\n assert module.dropout.rngs.count[...] == 1\n\n y2 = forward_block(module, x)\n\n # dropout is working!\n assert not jnp.allclose(y, y2)\n\n def test_replicate_single(self):\n din = 3\n dout = 10\n\n class Block(nnx.Module):\n def __init__(self, rngs: nnx.Rngs):\n self.linear = nnx.Linear(din, dout, rngs=rngs)\n self.dropout = nnx.Dropout(0.5, deterministic=False, rngs=rngs)\n\n def __call__(self, x: jax.Array) -> jax.Array:\n return self.dropout(nnx.relu(self.linear(x)))\n\n def create_block(rngs: nnx.Rngs):\n return Block(rngs)\n\n state_axes = nnx.StateAxes({nnx.RngState: 0, ...: None})\n\n @nnx.split_rngs(splits=1)\n @partial(nnx.pmap, in_axes=(state_axes, 0), out_axes=0, axis_size=1,\n graph=True)\n def forward_block(module: Block, x):\n return module(x)\n\n rngs = nnx.Rngs(0)\n module = create_block(rngs)\n initial_key = module.dropout.rngs.key[...]\n\n assert module.dropout.rngs.count[...] == 0\n assert module.linear.kernel.shape == (din, dout)\n assert module.linear.bias.shape == (dout,)\n\n x = jnp.ones((1, 5, din))\n\n y = forward_block(module, x)\n\n assert y.shape == (1, 5, dout)\n assert module.dropout.rngs.count[...] == 1\n\n y2 = forward_block(module, x)\n\n # dropout is working!\n assert not jnp.allclose(y, y2)\n\n assert module.dropout.rngs.key[...] == initial_key\n\n\nclass TestCond(parameterized.TestCase):\n def test_basic(self):\n class TimeStep(tp.NamedTuple):\n step: nnx.Variable[jax.Array]\n reward: nnx.Variable[jax.Array]\n\n @staticmethod\n def zero():\n return TimeStep(\n step=nnx.Variable(jnp.array(0)), reward=nnx.Variable(jnp.array(0.0))\n )\n\n @nnx.dataclass\n class Foo(nnx.Pytree):\n timestep: TimeStep = nnx.data()\n\n def update(self):\n def reward_2(self: Foo):\n self.timestep = TimeStep(\n step=nnx.Variable(self.timestep.step + 1),\n reward=nnx.Variable(jnp.array(2.0)),\n )\n\n def reward_0(self: Foo):\n self.timestep = TimeStep(\n step=nnx.Variable(self.timestep.step + 1),\n reward=nnx.Variable(jnp.array(0.0)),\n )\n\n nnx.cond(self.timestep.step % 2 == 0, reward_2, reward_0, self)\n\n foo = Foo(timestep=TimeStep.zero())\n foo.update()\n self.assertEqual(foo.timestep.step[...], 1)\n self.assertEqual(foo.timestep.reward[...], 2.0)\n foo.update()\n self.assertEqual(foo.timestep.step[...], 2)\n self.assertEqual(foo.timestep.reward[...], 0.0)\n foo.update()\n self.assertEqual(foo.timestep.step[...], 3)\n self.assertEqual(foo.timestep.reward[...], 2.0)\n foo.update()\n self.assertEqual(foo.timestep.step[...], 4)\n self.assertEqual(foo.timestep.reward[...], 0.0)\n\n @parameterized.parameters(\n (True, False),\n (False, False),\n )\n def test_basic_variable(self, graph, graph_updates):\n def collatz(x):\n def even(x):\n x[...] = x // 2\n\n def odd(x):\n x[...] = 3 * x + 1\n\n return nnx.cond(\n x % 2 == 0, even, odd, x,\n graph=graph, graph_updates=graph_updates,\n )\n\n x = nnx.Variable(jnp.array(8))\n collatz(x)\n self.assertEqual(x[...], 4)\n collatz(x)\n self.assertEqual(x[...], 2)\n collatz(x)\n self.assertEqual(x[...], 1)\n collatz(x)\n self.assertEqual(x[...], 4)\n\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_cond_and_vmap(self, graph, graph_updates):\n class Env(nnx.Pytree):\n def __init__(self):\n self.index = nnx.Variable(jnp.arange(8))\n self.step = nnx.Variable(jnp.zeros((8,), jnp.uint32))\n\n env = Env()\n model = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n\n @nnx.vmap(in_axes=(0, None), out_axes=None, graph=graph, graph_updates=graph_updates)\n def f(env: Env, model: nnx.Linear):\n self.assertEqual(env.index.shape, ())\n\n def increment(env: Env):\n env.step[...] += 1\n\n def no_nothing(env: Env):\n pass\n\n is_even = env.index % 2 == 0\n nnx.cond(\n is_even, increment, no_nothing, env,\n graph=graph, graph_updates=graph_updates,\n )\n\n f(env, model)\n\n np.testing.assert_array_equal(\n env.step[...], np.array([1, 0, 1, 0, 1, 0, 1, 0], np.uint32)\n )\n\n @parameterized.parameters(\n (True, False),\n (False, False),\n )\n def test_cond_different_variable_per_branch(self, graph, graph_updates):\n a = nnx.Variable(jnp.array(0))\n b = nnx.Variable(jnp.array(0))\n\n def update_a(a, b):\n a[...] += 1\n\n def update_b(a, b):\n b[...] += 10\n\n nnx.cond(\n True, update_a, update_b, a, b,\n graph=graph, graph_updates=graph_updates,\n )\n self.assertEqual(a[...], 1)\n self.assertEqual(b[...], 0)\n\n nnx.cond(\n False, update_a, update_b, a, b,\n graph=graph, graph_updates=graph_updates,\n )\n self.assertEqual(a[...], 1)\n self.assertEqual(b[...], 10)\n\n def test_cond_shared_references(self):\n\n @dataclasses.dataclass\n class Foo(nnx.Module):\n a: nnx.Variable\n b: nnx.Variable\n\n v = nnx.Variable(jnp.array(0))\n m = Foo(a=v, b=v)\n\n def true_fn(m):\n m.a[...] += 1\n\n def false_fn(m):\n m.b[...] += 2\n\n nnx.cond(True, true_fn, false_fn, m, graph=True, graph_updates=False)\n np.testing.assert_allclose(m.a[...], 1)\n np.testing.assert_allclose(m.b[...], 1)\n nnx.cond(False, true_fn, false_fn, m, graph=True, graph_updates=False)\n np.testing.assert_allclose(m.a[...], 3)\n np.testing.assert_allclose(m.b[...], 3)\n\n with self.assertRaises(ValueError):\n nnx.cond(True, true_fn, false_fn, m, graph=False)\n\nclass TestSwitch(parameterized.TestCase):\n @parameterized.parameters(\n (True, False),\n (False, False),\n )\n def test_basic(self, graph, graph_updates):\n class RoundTable(nnx.Module):\n def __init__(self):\n self.next_index = 0\n self.linear = nnx.Linear(10, 10, rngs=nnx.Rngs(0))\n self.linear.kernel[...] = jnp.identity(10)\n self.rounds_count = nnx.Variable(jnp.array(0))\n\n def __call__(self, x):\n def fn0(m, x):\n m.rounds_count[...] += 1\n return m.linear(x)\n def fn1(m, x):\n return m.linear(x) * 2\n def fn2(m, x):\n m.linear.kernel[...] = jnp.zeros((10, 10))\n return m.linear(x)\n\n y = nnx.switch(\n self.next_index, (fn0, fn1, fn2), self, x,\n graph=graph, graph_updates=graph_updates,\n )\n self.next_index = (self.next_index + 1) % 3\n return y\n\n model = RoundTable()\n x = jnp.ones((10,))\n np.testing.assert_array_equal(model(x), x)\n assert model.rounds_count[...] == 1\n assert model.next_index == 1\n np.testing.assert_array_equal(model(x), x * 2)\n assert model.rounds_count[...] == 1\n assert model.next_index == 2\n np.testing.assert_array_equal(model(x), jnp.zeros((10,)))\n assert model.rounds_count[...] == 1\n assert model.next_index == 0\n np.testing.assert_array_equal(model(x), jnp.zeros((10,)))\n assert model.rounds_count[...] == 2\n assert model.next_index == 1\n\n @parameterized.parameters(\n (True, False),\n (False, False),\n )\n def test_switch_variable(self, graph, graph_updates):\n def add_1(x):\n x[...] += 1\n\n def add_10(x):\n x[...] += 10\n\n def add_100(x):\n x[...] += 100\n\n x = nnx.Variable(jnp.array(0))\n nnx.switch(0, (add_1, add_10, add_100), x,\n graph=graph, graph_updates=graph_updates)\n self.assertEqual(x[...], 1)\n nnx.switch(1, (add_1, add_10, add_100), x,\n graph=graph, graph_updates=graph_updates)\n self.assertEqual(x[...], 11)\n nnx.switch(2, (add_1, add_10, add_100), x,\n graph=graph, graph_updates=graph_updates)\n self.assertEqual(x[...], 111)\n\n def test_switch_shared_references(self):\n @dataclasses.dataclass\n class Foo(nnx.Module):\n a: nnx.Variable\n b: nnx.Variable\n\n v = nnx.Variable(jnp.array(0))\n m = Foo(a=v, b=v)\n\n def add_a(m):\n m.a[...] += 1\n\n def add_b(m):\n m.b[...] += 10\n\n nnx.switch(0, (add_a, add_b), m, graph=True, graph_updates=False)\n np.testing.assert_allclose(m.a[...], 1)\n np.testing.assert_allclose(m.b[...], 1)\n\n nnx.switch(1, (add_a, add_b), m, graph=True, graph_updates=False)\n np.testing.assert_allclose(m.a[...], 11)\n np.testing.assert_allclose(m.b[...], 11)\n\n with self.assertRaises(ValueError):\n nnx.switch(0, (add_a, add_b), m, graph=False)\n\nclass TestWhileLoop(parameterized.TestCase):\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_basic(self, graph, graph_updates):\n def fwd_fn(input):\n m, x, c = input\n y = m(x)\n return m, y, c - 1.0\n\n module = nnx.Linear(10, 10, rngs=nnx.Rngs(0))\n module.kernel[...] = jnp.identity(10) * 2\n x = 1e1 * jax.random.normal(jax.random.key(0), (10,))\n\n _, y, _ = nnx.while_loop(\n lambda input: input[-1] > 0, fwd_fn, (module, x, 3.0),\n graph=graph, graph_updates=graph_updates)\n np.testing.assert_array_equal(y, x * 8)\n\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_multiple_objects(self, graph, graph_updates):\n def fwd_fn(input):\n m1, (w2,), x, c = input\n y = m1(x) @ w2\n return m1, (w2,), y, c - 1.0\n\n m1 = nnx.Linear(10, 10, rngs=nnx.Rngs(0))\n m1.kernel[...] = jnp.identity(10) * 2\n w2 = nnx.Variable(jnp.identity(10) * 0.5)\n x = 1e1 * jax.random.normal(jax.random.key(0), (10,))\n\n _, _, y, _ = nnx.while_loop(\n lambda input: input[-1] > 0, fwd_fn, (m1, (w2,), x, 3.0),\n graph=graph, graph_updates=graph_updates)\n np.testing.assert_allclose(y, x)\n\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_nested_module(self, graph, graph_updates):\n def fwd_fn(input):\n m, x, c = input\n y = m(x)\n return m, y, c - 1.0\n\n module = nnx.Linear(10, 10, rngs=nnx.Rngs(0))\n module.kernel[...] = jnp.identity(10) * 2\n module = nnx.Sequential(module)\n x = 1e1 * jax.random.normal(jax.random.key(0), (10,))\n\n _, y, _ = nnx.while_loop(\n lambda input: input[-1] > 0, fwd_fn, (module, x, 3.0),\n graph=graph, graph_updates=graph_updates)\n np.testing.assert_array_equal(y, x * 8)\n\n def test_shared_module(self):\n m1 = nnx.Linear(10, 10, rngs=nnx.Rngs(0))\n m2 = nnx.Linear(10, 10, use_bias=False, rngs=nnx.Rngs(0))\n m2.kernel = m1.kernel\n module = nnx.Sequential(m1, m2)\n self.assertLen(jax.tree.leaves(nnx.state(module)), 2) # only m1 params\n\n def fwd_fn(input):\n m, x, c = input\n y = m(x)\n m.layers[0].kernel[...] = jnp.zeros_like(m.layers[0].kernel[...])\n return m, y, c - 1.0\n\n x = 1e1 * jax.random.normal(jax.random.key(0), (10,))\n _, y, _ = nnx.while_loop(\n lambda input: input[-1] > 0, fwd_fn, (module, x, 2.0))\n self.assertLen(jax.tree.leaves(nnx.state(module)), 2) # only m1 params\n np.testing.assert_array_equal(\n m1.kernel[...],\n jnp.zeros((10, 10)),\n )\n np.testing.assert_array_equal(\n m2.kernel[...],\n jnp.zeros((10, 10)),\n )\n np.testing.assert_array_equal(y, jnp.zeros((10,)))\n\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_value_changed(self, graph, graph_updates):\n def fwd_fn(input):\n m, x, c = input\n m.kernel[...] = jnp.zeros_like(m.kernel)\n y = m(x)\n return m, y, c - 1.0\n\n module = nnx.Linear(10, 10, rngs=nnx.Rngs(0))\n x = 1e1 * jax.random.normal(jax.random.key(0), (10,))\n\n _, y, _ = nnx.while_loop(\n lambda input: input[-1] > 0, fwd_fn, (module, x, 3.0),\n graph=graph, graph_updates=graph_updates)\n np.testing.assert_array_equal(\n module.kernel[...],\n jnp.zeros((10, 10)),\n )\n np.testing.assert_array_equal(y, jnp.zeros((10,)))\n\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_ref_changed(self, graph, graph_updates):\n def fwd_fn(input):\n m, x, c = input\n y = m(x)\n m.kernel = nnx.Param(jnp.zeros_like(m.kernel))\n return m, y, c - 1.0\n\n module = nnx.Linear(10, 10, rngs=nnx.Rngs(0))\n x = 1e1 * jax.random.normal(jax.random.key(0), (10,))\n\n with self.assertRaises(ValueError):\n _, y, _ = nnx.while_loop(\n lambda input: input[-1] > 0, fwd_fn, (module, x, 2.0),\n graph=graph, graph_updates=graph_updates)\n\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_structure_changed(self, graph, graph_updates):\n def fwd_fn(input):\n m, x, c = input\n m = nnx.Linear(10, 10, use_bias=False, rngs=nnx.Rngs(1))\n m.kernel[...] = jnp.identity(10) * 2\n y = m(x)\n return m, y, c - 1.0\n\n module = nnx.Linear(10, 10, use_bias=True, rngs=nnx.Rngs(0))\n x = 1e1 * jax.random.normal(jax.random.key(0), (10,))\n\n with self.assertRaises((ValueError, TypeError)):\n _, y, _ = nnx.while_loop(\n lambda input: input[-1] > 0, fwd_fn, (module, x, 3.0),\n graph=graph, graph_updates=graph_updates)\n\n def test_repeated_object(self):\n m = nnx.Linear(10, 10, rngs=nnx.Rngs(0))\n\n def body_fn(val):\n count, m, _ = val\n return count + 1, m, m\n\n count, m, _ = nnx.while_loop(\n lambda val: val[0] < 2,\n body_fn,\n (0, m, m),\n )\n\n def test_immut_fori_loop(self):\n def immut_fn(i, carry):\n g_accum = carry\n grads = jax.tree.map(jnp.ones_like, g_accum)\n g_accum = jax.tree.map(lambda gm, g: gm + g, g_accum, grads)\n return g_accum\n\n model = nnx.Linear(10, 10, rngs=nnx.Rngs(0), use_bias=False)\n g_accum = jax.tree.map(jnp.zeros_like, nnx.state(model))\n nnx.fori_loop(0, 2, immut_fn, g_accum)\n\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_fori_loop_grad_accum(self, graph, graph_updates):\n accum = nnx.Variable(jnp.zeros((10, 10)))\n\n def accum_fn(i, accum):\n accum[...] += 1\n return accum\n\n accum = nnx.fori_loop(0, 3, accum_fn, accum,\n graph=graph, graph_updates=graph_updates)\n np.testing.assert_array_equal(accum[...], jnp.full((10, 10), 3.0))\n\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_fori_loop_basic(self, graph, graph_updates):\n def fwd_fn(i, input):\n m, x = input\n m.kernel[...] = jnp.identity(10) * i\n return m, m(x)\n\n module = nnx.Linear(10, 10, rngs=nnx.Rngs(0))\n x = jax.random.normal(jax.random.key(0), (10,))\n\n _, y = nnx.fori_loop(2, 4, fwd_fn, (module, x),\n graph=graph, graph_updates=graph_updates)\n np.testing.assert_array_equal(y, x * 2 * 3)\n\n def test_fori_loop_with_sharing(self):\n class A(nnx.Pytree):\n def __init__(self):\n self.params = nnx.Param(jnp.zeros((10,), dtype=int))\n\n class B(nnx.Pytree):\n def __init__(self, a: A):\n self.a = a\n\n class C(nnx.Pytree):\n def __init__(self, a: A):\n self.a = a\n\n class D(nnx.Pytree):\n def __init__(self):\n self.a = A()\n self.b = B(self.a)\n self.c = C(self.a)\n\n def increment(_, d: D) -> D:\n d.a.params[...] += 1\n return d\n\n @nnx.jit\n def rollout(d: D):\n nnx.fori_loop(0, 10, increment, d)\n\n d = D()\n rollout(d)\n\n np.testing.assert_array_equal(\n d.a.params[...], np.full((10,), 10, dtype=int)\n )\n\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_loops_multiple_modules(self, graph, graph_updates):\n class Foo(nnx.Module):\n def __init__(self):\n self.param = nnx.Param(jnp.zeros((1,)))\n def __call__(self, x):\n return self.param\n\n def loop_fn(inputs):\n return inputs\n while_loop_fn = lambda inputs: (*loop_fn(inputs[:-1]), inputs[-1]-1)\n fori_loop_fn = lambda i, inputs: loop_fn(inputs)\n a = Foo()\n b = Foo()\n nnx.while_loop(lambda input: input[-1] > 0, while_loop_fn, (a, b, 2),\n graph=graph, graph_updates=graph_updates)\n nnx.fori_loop(0, 2, fori_loop_fn, (a, b), graph=graph)\n\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_tree_mode_while_loop_stateful(self, graph, graph_updates):\n class Counter(nnx.Module):\n def __init__(self):\n self.count = nnx.Variable(jnp.array(0))\n\n counter = Counter()\n module = nnx.Linear(10, 10, rngs=nnx.Rngs(0))\n module.kernel[...] = jnp.identity(10) * 2\n x = jax.random.normal(jax.random.key(0), (10,))\n\n def body_fn(val):\n counter, module, x, i = val\n counter.count[...] += 1\n x = module(x)\n return counter, module, x, i - 1\n\n counter, module, y, _ = nnx.while_loop(\n lambda val: val[-1] > 0,\n body_fn,\n (counter, module, x, 3),\n graph=graph,\n graph_updates=graph_updates,\n )\n np.testing.assert_array_equal(counter.count[...], 3)\n np.testing.assert_array_equal(y, x * 8)\n\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_tree_mode_while_loop_inside_jit(self, graph, graph_updates):\n module = nnx.Linear(10, 10, rngs=nnx.Rngs(0))\n module.kernel[...] = jnp.identity(10) * 2\n x = jax.random.normal(jax.random.key(0), (10,))\n\n @nnx.jit(graph=graph, graph_updates=graph_updates)\n def f(module, x):\n def body_fn(val):\n m, x, c = val\n return m, m(x), c - 1.0\n _, y, _ = nnx.while_loop(\n lambda val: val[-1] > 0,\n body_fn,\n (module, x, 3.0),\n graph=graph,\n )\n return y\n\n y = f(module, x)\n np.testing.assert_array_equal(y, x * 8)\n\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_tree_mode_fori_loop_stateful(self, graph, graph_updates):\n class Counter(nnx.Module):\n def __init__(self):\n self.count = nnx.Variable(jnp.array(0))\n\n counter = Counter()\n module = nnx.Linear(10, 10, rngs=nnx.Rngs(0))\n module.kernel[...] = jnp.identity(10) * 2\n x = jax.random.normal(jax.random.key(0), (10,))\n\n def body_fn(i, val):\n counter, module, x = val\n counter.count[...] += 1\n x = module(x)\n return counter, module, x\n\n counter, module, y = nnx.fori_loop(\n 0, 3, body_fn, (counter, module, x),\n graph=graph, graph_updates=graph_updates,\n )\n np.testing.assert_array_equal(counter.count[...], 3)\n np.testing.assert_array_equal(y, x * 8)\n\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_tree_mode_fori_loop_inside_jit(self, graph, graph_updates):\n module = nnx.Linear(10, 10, rngs=nnx.Rngs(0))\n module.kernel[...] = jnp.identity(10) * 2\n x = jax.random.normal(jax.random.key(0), (10,))\n\n @nnx.jit(graph=graph, graph_updates=graph_updates)\n def f(module, x):\n def body_fn(i, val):\n m, x = val\n return m, m(x)\n _, y = nnx.fori_loop(\n 0, 3, body_fn, (module, x), graph=graph,\n )\n return y\n\n y = f(module, x)\n np.testing.assert_array_equal(y, x * 8)\n\nclass TestSplitMergeInputs(absltest.TestCase):\n def test_split_inputs(self):\n class StatefulLinear(nnx.Linear):\n def __init__(self, din: int, dout: int, rngs: nnx.Rngs):\n super().__init__(din, dout, rngs=rngs)\n self.counter = nnx.BatchStat(jnp.array(0, jnp.uint32))\n\n def __call__(self, x):\n self.counter[...] += 1\n return super().__call__(x)\n\n model = StatefulLinear(3, 4, rngs=nnx.Rngs(0))\n\n @general.split_inputs\n @jax.jit\n @general.merge_inputs\n def forward(model, x):\n return model(x)\n\n x = jnp.ones((2, 3))\n y = forward(model, x)\n\n self.assertEqual(model.counter[...], 1)\n\n def test_split_inputs_cond(self):\n class Counter(nnx.Linear):\n def __init__(self):\n self.count = nnx.BatchStat(jnp.array(0, jnp.uint32))\n\n def increment(self):\n self.count[...] += 1\n\n counter = Counter()\n\n @general.merge_inputs\n def increment(counter: Counter):\n counter.increment()\n\n @general.merge_inputs\n def no_nothing(counter: Counter):\n pass\n\n general.split_inputs(jax.lax.cond)(True, increment, no_nothing, counter)\n\n self.assertEqual(counter.count[...], 1)\n\n general.split_inputs(jax.lax.cond)(False, increment, no_nothing, counter)\n\n self.assertEqual(counter.count[...], 1)\n\n def test_split_inputs_vmap(self):\n class EnvState(nnx.Variable[nnx.A]):\n pass\n\n class Env(nnx.Pytree):\n def __init__(self):\n self.index = EnvState(jnp.arange(8))\n self.step = EnvState(jnp.zeros((8,), jnp.uint32))\n\n env = Env()\n model = nnx.Linear(2, 3, rngs=nnx.Rngs(0))\n\n # internally merge_inputs returns (args, out)\n in_axes = (0, None)\n out_axes = (in_axes, None)\n\n @general.split_inputs\n @partial(jax.vmap, in_axes=in_axes, out_axes=out_axes)\n @general.merge_inputs\n def f(env: Env, model: nnx.Linear):\n self.assertEqual(env.index.shape, ())\n\n @general.merge_inputs\n def increment(env: Env):\n env.step[...] += 1\n\n @general.merge_inputs\n def no_nothing(env: Env):\n pass\n\n is_even = env.index % 2 == 0\n general.split_inputs(jax.lax.cond)(is_even, increment, no_nothing, env)\n\n f(env, model)\n\n np.testing.assert_array_equal(\n env.step[...], np.array([1, 0, 1, 0, 1, 0, 1, 0], np.uint32)\n )\n\nclass TestCheckify(parameterized.TestCase):\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_basic(self, graph, graph_updates):\n\n @dataclasses.dataclass\n class Foo(nnx.Module):\n a: nnx.Param\n\n @nnx.jit(graph=graph, graph_updates=graph_updates)\n def f(m):\n y = jnp.sin(m.a) # error\n return m.a + y\n\n m = Foo(a=nnx.Param(jnp.inf))\n err, out = nnx.checkify(\n f, errors=checkify.float_checks,\n graph=graph, graph_updates=graph_updates,\n )(m)\n\n with self.assertRaisesRegex(ValueError, 'nan generated by primitive: sin'):\n err.throw()\n\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_checkify_stateful(self, graph, graph_updates):\n count = nnx.Variable(jnp.array(0))\n\n @nnx.jit(graph=graph, graph_updates=graph_updates)\n def f(c):\n c[...] += 1\n return c[...]\n\n err, out = nnx.checkify(\n f, graph=graph, graph_updates=graph_updates,\n )(count)\n self.assertEqual(count[...], 1)\n np.testing.assert_allclose(out, 1)\n\n\nclass TestBoundMethodTransforms(parameterized.TestCase):\n def test_remat_with_bound_method_raises(self):\n class M(nnx.Module):\n def __init__(self, *, rngs: nnx.Rngs):\n self.linear = nnx.Linear(3, 3, rngs=rngs)\n self.count = nnx.BatchStat(0)\n\n def block(self, x: jax.Array) -> jax.Array:\n self.count[...] += 1\n return self.linear(x)\n\n m = M(rngs=nnx.Rngs(0))\n with self.assertRaisesRegex(ValueError, 'bound methods'):\n nnx.remat(m.block)\n\n def test_jit_with_bound_method_raises(self):\n class M(nnx.Module):\n def __init__(self, *, rngs: nnx.Rngs):\n self.linear = nnx.Linear(2, 3, rngs=rngs)\n def apply(self, x: jax.Array, scale: int):\n return self.linear(x) * scale\n\n m = M(rngs=nnx.Rngs(0))\n with self.assertRaisesRegex(ValueError, 'bound methods'):\n nnx.jit(m.apply, static_argnums=1)\n\n def test_vmap_with_bound_method_raises(self):\n class M(nnx.Module):\n def __init__(self, *, rngs: nnx.Rngs):\n self.linear = nnx.Linear(2, 3, rngs=rngs)\n def __call__(self, x: jax.Array):\n return self.linear(x)\n\n m = M(rngs=nnx.Rngs(0))\n with self.assertRaisesRegex(ValueError, 'bound methods'):\n nnx.vmap(m.__call__, in_axes=(0,), out_axes=0)\n\n def test_eval_shape_with_bound_method_raises(self):\n class M(nnx.Module):\n def __init__(self, *, rngs: nnx.Rngs):\n self.linear = nnx.Linear(2, 3, rngs=rngs)\n def __call__(self, x: jax.Array):\n return self.linear(x)\n\n m = M(rngs=nnx.Rngs(0))\n x_spec = jax.ShapeDtypeStruct((1, 2), jnp.float32)\n with self.assertRaisesRegex(ValueError, 'bound methods'):\n nnx.eval_shape(m.__call__, x_spec)\n\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_grad_with_bound_method_raises(self, graph_mode, graph_updates):\n class M(nnx.Module):\n def __init__(self):\n self.w = nnx.Param(jnp.array(1.0))\n def loss(self, s: float):\n return (self.w * s) ** 2\n\n m = M()\n with self.assertRaisesRegex(ValueError, 'bound methods'):\n nnx.grad(m.loss, graph=graph_mode)\n\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_value_and_grad_with_bound_method_raises(self, graph_mode, graph_updates):\n class TestModel(nnx.Module):\n def __init__(self, *, rngs: nnx.Rngs):\n self.linear = nnx.Linear(2, 1, rngs=rngs)\n\n def loss_fn(self, x, y):\n pred = self.linear(x)\n return jnp.mean((pred - y) ** 2)\n\n model = TestModel(rngs=nnx.Rngs(0))\n with self.assertRaisesRegex(ValueError, 'bound methods'):\n nnx.value_and_grad(model.loss_fn, graph=graph_mode)\n\n def test_checkify_with_bound_method_raises(self):\n \"\"\"Test that checkify raises error for bound methods.\"\"\"\n class M(nnx.Module):\n def __call__(self, x: jax.Array):\n return x + 1\n\n m = M()\n with self.assertRaisesRegex(ValueError, 'bound methods'):\n nnx.checkify(m.__call__)\n\n def test_pmap_with_bound_method_raises(self):\n \"\"\"Test that pmap raises error for bound methods.\"\"\"\n class M(nnx.Module):\n def __call__(self, x: jax.Array):\n return x + 1\n\n m = M()\n with self.assertRaisesRegex(ValueError, 'bound methods'):\n nnx.pmap(m.__call__)\n\n def test_shard_map_with_bound_method_raises(self):\n \"\"\"Test that shard_map raises error for bound methods.\"\"\"\n class M(nnx.Module):\n def __call__(self, x: jax.Array):\n return x + 1\n\n m = M()\n mesh = jax.sharding.Mesh(jax.local_devices()[:1], ('data',))\n with self.assertRaisesRegex(ValueError, 'bound methods'):\n nnx.shard_map(m.__call__, mesh=mesh, in_specs=None, out_specs=None)\n\n def test_custom_vjp_with_bound_method_raises(self):\n \"\"\"Test that custom_vjp raises error for bound methods.\"\"\"\n class M(nnx.Module):\n def __call__(self, x: jax.Array):\n return x + 1\n\n m = M()\n with self.assertRaisesRegex(ValueError, 'bound methods'):\n nnx.custom_vjp(m.__call__)\n\n def test_scan_bound_method_raises(self):\n class M(nnx.Module):\n def __call__(self, x: jax.Array):\n return x + 1\n m = M()\n with self.assertRaisesRegex(ValueError, 'bound methods'):\n _ = nnx.scan(m.__call__, in_axes=(0,), out_axes=0)\n\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_tree_mode_pmap_basic(self, graph, graph_updates):\n class LinearEnsemble(nnx.Module):\n def __init__(self, num, *, rngs):\n self.w = nnx.Param(jax.random.uniform(rngs(), (num, 2, 3)))\n\n model = LinearEnsemble(1, rngs=nnx.Rngs(0))\n x = jnp.ones((2,))\n\n @nnx.pmap(in_axes=(0, None), out_axes=0, axis_size=1,\n graph=graph, graph_updates=graph_updates)\n def forward(model, x):\n return x @ model.w\n\n y = forward(model, x)\n assert y.shape == (1, 3)\n\n @parameterized.parameters(\n (True, True), (True, False), (False, False),\n )\n def test_tree_mode_pmap_stateful(self, graph, graph_updates):\n class Counter(nnx.Variable):\n pass\n\n class Linear(nnx.Module):\n def __init__(self, din, dout, *, rngs):\n key = rngs.params()\n self.w = nnx.Param(jax.random.uniform(key, (din, dout)))\n self.count = Counter(jnp.array(0))\n\n def __call__(self, x):\n self.count[...] += 1\n return x @ self.w\n\n model = Linear(2, 3, rngs=nnx.Rngs(0))\n\n @nnx.pmap(in_axes=(None, 0), out_axes=0, axis_size=1,\n graph=graph, graph_updates=graph_updates)\n def forward(model, x):\n return model(x)\n\n x = jnp.ones((1, 2))\n y = forward(model, x)\n assert y.shape == (1, 3)\n assert model.count.get_value() == 1\n\n def test_tree_mode_pmap_split_merge(self):\n class Block(nnx.Module):\n def __init__(self, rngs: nnx.Rngs):\n self.linear = nnx.Linear(3, 10, rngs=rngs)\n self.dropout = nnx.Dropout(0.5, deterministic=False, rngs=rngs)\n\n def __call__(self, x: jax.Array) -> jax.Array:\n x = self.linear(x)\n x = nnx.elu(x)\n x = self.dropout(x)\n return x\n\n rngs = nnx.Rngs(0)\n\n @nnx.split_rngs(splits=1, graph=False)\n @nnx.pmap(in_axes=0, out_axes=(None, 0, 0, None), axis_size=1, graph=False)\n def create_block(rngs):\n block = Block(rngs)\n graphdef, params_state, rng_state, rest_state = nnx.split(\n block, nnx.Param, nnx.RngState, ...,\n )\n return graphdef, params_state, rng_state, rest_state\n\n graphdef, params_state, rng_state, rest_state = create_block(rngs)\n\n assert rng_state.dropout.rngs.count[0] == 0\n assert params_state.linear.kernel.shape == (1, 3, 10)\n assert params_state.linear.bias.shape == (1, 10)\n\n x = jnp.ones((1, 1, 3))\n\n @nnx.pmap(in_axes=(0, 0, None, 0), axis_size=1, graph=False)\n def forward_block(params_state, rng_state, rest_state, x):\n return nnx.merge(graphdef, params_state, rng_state, rest_state)(x)\n\n y = forward_block(params_state, rng_state, rest_state, x)\n\n assert y.shape == (1, 1, 10)\n\n y2 = forward_block(params_state, rng_state, rest_state, x)\n\n assert not jnp.allclose(y, y2)\n\n def test_tree_mode_pmap_replicate(self):\n din = 3\n dout = 10\n\n class Block(nnx.Module):\n def __init__(self, rngs: nnx.Rngs):\n self.linear = nnx.Linear(din, dout, rngs=rngs)\n self.dropout = nnx.Dropout(0.5, deterministic=False, rngs=rngs)\n\n def __call__(self, x: jax.Array) -> jax.Array:\n return self.dropout(nnx.relu(self.linear(x)))\n\n rngs = nnx.Rngs(0)\n module = Block(rngs)\n\n assert module.dropout.rngs.count[...] == 0\n assert module.linear.kernel.shape == (din, dout)\n assert module.linear.bias.shape == (dout,)\n\n module = nnx.split_rngs(module, splits=1, graph=False)\n graphdef, rng_state, rest_state = nnx.split(\n module, nnx.RngState, ...,\n )\n\n @nnx.pmap(in_axes=(0, None, 0), out_axes=0, axis_size=1, graph=False)\n def forward_block(rng_state, rest_state, x):\n module = nnx.merge(graphdef, rng_state, rest_state)\n y = module(x)\n return y\n\n x = jnp.ones((1, 5, din))\n\n y = forward_block(rng_state, rest_state, x)\n\n assert y.shape == (1, 5, dout)\n assert module.dropout.rngs.count[0] == 1\n\n y2 = forward_block(rng_state, rest_state, x)\n\n assert module.dropout.rngs.count[0] == 2\n assert not jnp.allclose(y, y2)\n\n\nclass TestPureJaxFancyScan(absltest.TestCase):\n\n def test_carry_and_scan(self):\n def cumsum(carry, x):\n carry = carry + x\n return carry, carry\n\n final_carry, ys = pure_jax_fancy_scan(\n cumsum, jnp.array(0.0), jnp.arange(5.0),\n in_axes=(nnx.Carry, 0), out_axes=(nnx.Carry, 0),\n )\n np.testing.assert_allclose(final_carry, 10.0)\n np.testing.assert_allclose(ys, jnp.array([0., 1., 3., 6., 10.]))\n\n def test_carry_only_output(self):\n def sum_fn(carry, x):\n return carry + x\n\n result = pure_jax_fancy_scan(\n sum_fn, jnp.array(0.0), jnp.arange(5.0),\n in_axes=(nnx.Carry, 0), out_axes=nnx.Carry,\n )\n np.testing.assert_allclose(result, 10.0)\n\n def test_broadcast_args(self):\n def scale_cumsum(carry, scale, x):\n carry = carry + x * scale\n return carry, carry\n\n final_carry, _ = pure_jax_fancy_scan(\n scale_cumsum, jnp.array(0.0), jnp.array(2.0), jnp.arange(5.0),\n in_axes=(nnx.Carry, None, 0), out_axes=(nnx.Carry, 0),\n )\n np.testing.assert_allclose(final_carry, 20.0)\n\n def test_pytree_carry(self):\n def dict_scan(carry, x):\n carry = {'a': carry['a'] + x['a'], 'b': carry['b'] + x['b']}\n return carry, carry\n\n xs = {'a': jnp.arange(3.0), 'b': jnp.ones(3)}\n init = {'a': jnp.array(0.0), 'b': jnp.array(0.0)}\n final_carry, _ = pure_jax_fancy_scan(\n dict_scan, init, xs,\n in_axes=(nnx.Carry, 0), out_axes=(nnx.Carry, 0),\n )\n np.testing.assert_allclose(final_carry['a'], 3.0)\n np.testing.assert_allclose(final_carry['b'], 3.0)\n\n def test_no_carry_all_scanned(self):\n def double(x):\n return (x * 2,)\n\n (ys,) = pure_jax_fancy_scan(\n double, jnp.arange(5.0),\n in_axes=(0,), out_axes=(0,),\n )\n np.testing.assert_allclose(ys, jnp.arange(5.0) * 2)\n\n def test_reverse(self):\n def cumsum(carry, x):\n carry = carry + x\n return carry, carry\n\n final_carry, _ = pure_jax_fancy_scan(\n cumsum, jnp.array(0.0), jnp.arange(5.0),\n in_axes=(nnx.Carry, 0), out_axes=(nnx.Carry, 0), reverse=True,\n )\n np.testing.assert_allclose(final_carry, 10.0)\n\n def test_pytree_prefix_in_axes(self):\n def fn(carry, x):\n carry = carry + x['a'] + x['b']\n return carry, carry\n\n xs = {'a': jnp.arange(3.0), 'b': jnp.array(1.0)}\n final_carry, _ = pure_jax_fancy_scan(\n fn, jnp.array(0.0), xs,\n in_axes=(nnx.Carry, {'a': 0, 'b': None}), out_axes=(nnx.Carry, 0),\n )\n np.testing.assert_allclose(final_carry, 6.0)\n\n def test_nested_carry_rejected(self):\n with self.assertRaises(ValueError):\n pure_jax_fancy_scan(\n lambda x: x,\n {'a': jnp.array(1.0)},\n in_axes=({'a': nnx.Carry},), out_axes=nnx.Carry,\n )\n\n def test_broadcast_out_axes_rejected(self):\n with self.assertRaises(ValueError):\n pure_jax_fancy_scan(\n lambda c, x: (c, x),\n jnp.array(0.0), jnp.arange(3.0),\n in_axes=(nnx.Carry, 0), out_axes=(nnx.Carry, None),\n )\n\n def test_none_broadcast_input(self):\n def fn(carry, _unused, x):\n carry = carry + x\n return carry, carry\n\n final_carry, _ = pure_jax_fancy_scan(\n fn, jnp.array(0.0), None, jnp.arange(3.0),\n in_axes=(nnx.Carry, None, 0), out_axes=(nnx.Carry, 0),\n )\n np.testing.assert_allclose(final_carry, 3.0)\n\n def test_none_nested_in_arg(self):\n def fn(carry, x):\n carry = carry + x['a']\n return carry, carry\n\n xs = {'a': jnp.arange(3.0), 'b': None}\n final_carry, _ = pure_jax_fancy_scan(\n fn, jnp.array(0.0), xs,\n in_axes=(nnx.Carry, 0), out_axes=(nnx.Carry, 0),\n )\n np.testing.assert_allclose(final_carry, 3.0)\n\n def test_nested_carry_in_out_axes_rejected(self):\n with self.assertRaises(ValueError):\n pure_jax_fancy_scan(\n lambda c, x: (c, x),\n jnp.array(0.0), jnp.arange(3.0),\n in_axes=(nnx.Carry, 0), out_axes=({'a': nnx.Carry},),\n )\n\n def test_carry_in_in_axes_only_rejected(self):\n with self.assertRaises(ValueError):\n pure_jax_fancy_scan(\n lambda c, x: (c + x,),\n jnp.array(0.0), jnp.arange(3.0),\n in_axes=(nnx.Carry, 0), out_axes=(0,),\n )\n\n def test_carry_in_out_axes_only_rejected(self):\n with self.assertRaises(ValueError):\n pure_jax_fancy_scan(\n lambda x: x,\n jnp.arange(3.0),\n in_axes=(0,), out_axes=nnx.Carry,\n )\n\n def test_non_tuple_carry_only(self):\n def f(carry):\n return carry + 1.0\n\n result = pure_jax_fancy_scan(\n f, jnp.array(0.0),\n in_axes=nnx.Carry, out_axes=nnx.Carry, length=5,\n )\n np.testing.assert_allclose(result, 5.0)\n\n def test_non_tuple_scan_only(self):\n def f(x):\n return x * 2\n\n ys = pure_jax_fancy_scan(\n f, jnp.arange(5.0),\n in_axes=0, out_axes=0,\n )\n np.testing.assert_allclose(ys, jnp.arange(5.0) * 2)\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/nnx/variable_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport typing as tp\n\nimport jax\nimport jax.numpy as jnp\nimport numpy as np\nimport pytest\n\nfrom absl.testing import absltest, parameterized\nfrom flax import nnx\n\nA = tp.TypeVar('A')\n\nclass TestVariable(parameterized.TestCase):\n def test_pytree(self):\n r1 = nnx.Param(1)\n self.assertEqual(r1.get_value(), 1)\n\n r2 = jax.tree.map(lambda x: x + 1, r1)\n\n self.assertEqual(r1.get_value(), 1)\n self.assertEqual(r2.get_value(), 2)\n self.assertIsNot(r1, r2)\n\n def test_overloads_module(self):\n class Linear(nnx.Module):\n def __init__(self, din, dout, rngs: nnx.Rngs):\n key = rngs()\n self.w = nnx.Param(jax.random.normal(key, (din, dout)))\n self.b = nnx.Param(jax.numpy.zeros((dout,)))\n\n def __call__(self, x: jax.Array):\n return x @ self.w + self.b\n\n linear = Linear(3, 4, nnx.Rngs(0))\n x = jax.numpy.ones((3,))\n y = linear(x)\n self.assertEqual(y.shape, (4,))\n\n def test_jax_array(self):\n class Linear(nnx.Module):\n def __init__(self, din, dout, rngs: nnx.Rngs):\n key = rngs()\n self.w = nnx.Param(jax.random.normal(key, (din, dout)))\n self.b = nnx.Param(jax.numpy.zeros((dout,)))\n\n def __call__(self, x: jax.Array):\n return jnp.dot(x, self.w) + self.b # type: ignore[arg-type]\n\n linear = Linear(3, 4, nnx.Rngs(0))\n x = jax.numpy.ones((3,))\n y = linear(x)\n self.assertEqual(y.shape, (4,))\n\n def test_proxy_access(self):\n v = nnx.Param(jnp.ones((2, 3)))\n t = v.T\n\n self.assertEqual(t.shape, (3, 2))\n\n def test_proxy_call(self):\n class Callable(tp.NamedTuple):\n value: int\n\n def __call__(self, x):\n return self.value * x\n\n v = nnx.Param(Callable(2))\n result = v(3)\n\n self.assertEqual(result, 6)\n\n def test_binary_ops(self):\n v1 = nnx.Param(jnp.array(2))\n v2 = nnx.Param(jnp.array(3))\n\n result = v1 + v2\n\n self.assertEqual(result, 5)\n self.assertFalse(v1 == v2)\n\n v1[...] += v2\n\n self.assertEqual(v1[...], 5)\n\n @parameterized.product(\n v1=[np.array([1, 2]), np.array(2), 3],\n v2=[np.array([1, 2]), np.array(2), 3],\n )\n def test_eq_op(self, v1, v2):\n p1 = nnx.Param(jnp.asarray(v1) if isinstance(v1, np.ndarray) else v1)\n p2 = nnx.Param(jnp.asarray(v2) if isinstance(v2, np.ndarray) else v2)\n if isinstance(v1, np.ndarray) or isinstance(v2, np.ndarray):\n self.assertEqual((p1 == p2).all(), (v1 == v2).all())\n else:\n self.assertEqual(p1 == p2, v1 == v2)\n\n def test_mutable_array_context(self):\n initial_mode = nnx.var_defaults().hijax\n with nnx.var_defaults(hijax=False):\n v = nnx.Variable(jnp.array(1.0))\n self.assertEqual(nnx.var_defaults().hijax, False)\n self.assertNotIsInstance(v[...], jax.Ref)\n\n with nnx.var_defaults(hijax=True):\n v = nnx.Variable(jnp.array(1.0))\n self.assertEqual(nnx.var_defaults().hijax, True)\n self.assertIsInstance(v[...], jax.Array)\n\n v = nnx.Variable(jnp.array(2.0))\n self.assertIsInstance(v[...], jax.Array)\n self.assertEqual(nnx.var_defaults().hijax, False)\n\n nnx.var_defaults(hijax=True)\n\n v = nnx.Variable(jnp.array(0.0))\n self.assertEqual(nnx.var_defaults().hijax, True)\n self.assertIsInstance(v[...], jax.Array)\n\n v = nnx.Variable(jnp.array(1.0))\n self.assertEqual(nnx.var_defaults().hijax, initial_mode)\n self.assertIsInstance(v[...], jax.Array)\n\n def test_get_set_metadata(self):\n v = nnx.Variable(jnp.array(1.0))\n self.assertEqual(\n v.get_metadata(),\n {\n 'hijax': False,\n 'ref': False,\n 'eager_sharding': True,\n },\n )\n v.set_metadata(a=1, b=2)\n self.assertEqual(v.get_metadata('a'), 1)\n self.assertEqual(v.get_metadata('b'), 2)\n v.set_metadata({\n 'b': 3,\n 'c': 4,\n 'hijax': False,\n 'ref': False,\n 'eager_sharding': True,\n })\n self.assertEqual(\n v.get_metadata(),\n {\n 'b': 3,\n 'c': 4,\n 'hijax': False,\n 'ref': False,\n 'eager_sharding': True,\n },\n )\n self.assertEqual(v.get_metadata('b'), 3)\n self.assertEqual(v.get_metadata('c'), 4)\n c = v.get_metadata('c')\n self.assertEqual(c, 4)\n x = v.get_metadata('x', default=10)\n self.assertEqual(x, 10)\n\n def test_set_module_metadata(self):\n class Module(nnx.Module):\n def __init__(self):\n self.v = nnx.Variable(jnp.array(0.0))\n self.p = nnx.Param(jnp.array(1.0))\n\n m = Module()\n self.assertNotIn('foo', m.v.get_metadata())\n self.assertNotIn('foo', m.p.get_metadata())\n nnx.set_metadata(m, foo='bar')\n # Check that foo was added but the default metadata is still there\n v_metadata = m.v.get_metadata()\n p_metadata = m.p.get_metadata()\n self.assertEqual(v_metadata['foo'], 'bar')\n self.assertEqual(p_metadata['foo'], 'bar')\n # Check that default metadata is preserved\n self.assertIn('hijax', v_metadata)\n self.assertIn('ref', v_metadata)\n\n self.assertNotIn('differentiable', m.v.get_metadata())\n self.assertNotIn('differentiable', m.p.get_metadata())\n nnx.set_metadata(m, differentiable=False, only=nnx.Param)\n # Check that v still has foo but not differentiable\n v_metadata = m.v.get_metadata()\n self.assertEqual(v_metadata['foo'], 'bar')\n self.assertNotIn('differentiable', v_metadata)\n # Check that p has both foo and differentiable\n p_metadata = m.p.get_metadata()\n self.assertEqual(p_metadata['foo'], 'bar')\n self.assertEqual(p_metadata['differentiable'], False)\n\n @pytest.mark.skip(reason=\"Ref doesn't support broadcasting yet\")\n def test_broadcasting(self):\n v = nnx.Param(jnp.array([1.0, 2.0, 3.0]))\n x = v[None]\n self.assertEqual(x.shape, (1, 3))\n\n def test_set_metadata_out_sharding(self):\n v = nnx.Variable(jnp.array(1.0))\n\n v.set_metadata(out_sharding=jax.sharding.PartitionSpec(None))\n self.assertEqual(\n v.get_metadata('out_sharding'), jax.sharding.PartitionSpec(None)\n )\n\n v.set_metadata('out_sharding', jax.sharding.PartitionSpec('x'))\n self.assertEqual(\n v.get_metadata('out_sharding'), jax.sharding.PartitionSpec('x')\n )\n\n v.set_metadata({\n 'out_sharding': jax.sharding.PartitionSpec('y'),\n 'hijax': False,\n 'ref': False,\n 'eager_sharding': True,\n })\n self.assertEqual(\n v.get_metadata('out_sharding'), jax.sharding.PartitionSpec('y')\n )\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/pickle_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for flax.errors.\"\"\"\n\nfrom absl.testing import absltest\nfrom flax.errors import FlaxError, ScopeVariableNotFoundError\nimport pickle\n\nclass ErrorrsTest(absltest.TestCase):\n def test_exception_can_be_pickled(self):\n # tests the new __reduce__ method fixes bug reported in issue #4000\n ex = ScopeVariableNotFoundError('varname', 'collection', 'scope')\n pickled_ex = pickle.dumps(ex)\n unpicked_ex = pickle.loads(pickled_ex)\n self.assertIsInstance(unpicked_ex, FlaxError)\n self.assertIn('varname', str(unpicked_ex))\n self.assertIn('#flax.errors.ScopeVariableNotFoundError', str(unpicked_ex))\n self.assertNotIn('#flax.errors.FlaxError', str(unpicked_ex))\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/run_all_tests.sh", "content": "#!/bin/bash\n\n# export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python\nPYTEST_OPTS=\nRUN_DOCTEST=false\nRUN_MYPY=false\nRUN_PYTEST=false\nRUN_PYTYPE=false\nGH_VENV=false\n\nfor flag in \"$@\"; do\ncase $flag in\n --with-cov)\n PYTEST_OPTS+=\"--cov=flax --cov-report=xml --cov-report=term --cov-config=pyproject.toml\"\n ;;\n --help)\n echo \"Usage:\"\n echo \" --with-cov: Also generate pytest coverage.\"\n exit\n ;;\n --only-doctest)\n RUN_DOCTEST=true\n ;;\n --only-pytest)\n RUN_PYTEST=true\n ;;\n --only-pytype)\n RUN_PYTYPE=true\n ;;\n --only-mypy)\n RUN_MYPY=true\n ;;\n --use-venv)\n GH_VENV=true\n ;;\n *)\n echo \"Unknown flag: $flag\"\n exit 1\n ;;\nesac\ndone\n\n# if neither --only-doctest, --only-pytest, --only-pytype, --only-mypy is set, run all tests\nif ! $RUN_DOCTEST && ! $RUN_PYTEST && ! $RUN_PYTYPE && ! $RUN_MYPY; then\n RUN_DOCTEST=true\n RUN_PYTEST=true\n RUN_PYTYPE=true\n RUN_MYPY=true\nfi\n\n# Activate cached virtual env for github CI\nif $GH_VENV; then\n source $(dirname \"$0\")/../.venv/bin/activate\nfi\n\necho \"====== test config =======\"\necho \"PYTEST_OPTS: $PYTEST_OPTS\"\necho \"RUN_DOCTEST: $RUN_DOCTEST\"\necho \"RUN_PYTEST: $RUN_PYTEST\"\necho \"RUN_MYPY: $RUN_MYPY\"\necho \"RUN_PYTYPE: $RUN_PYTYPE\"\necho \"GH_VENV: $GH_VENV\"\necho \"WHICH PYTHON: $(which python)\"\necho \"jax: $(python -c 'import jax; print(jax.__version__)')\"\necho \"flax: $(python -c 'import flax; print(flax.__version__)')\"\necho \"==========================\"\necho \"\"\n\nsh $(dirname \"$0\")/download_dataset_metadata.sh || exit\n\n# Instead of using set -e, we have a manual error trap that\n# exits for any error code != 5 since pytest returns error code 5\n# for no found tests. (We may force minimal test coverage in examples\n# in the future!)\ntrap handle_errors ERR\nhandle_errors () {\n ret=\"$?\"\n if [[ \"$ret\" == 5 ]]; then\n echo \"error code $ret == no tests found in $egd\"\n else\n echo \"error code $ret\"\n exit 1\n fi\n}\n\n# Run embedded tests inside docs\nif $RUN_DOCTEST; then\n echo \"=== RUNNING DOCTESTS ===\"\n # test doctest\n sphinx-build -M doctest docs docs/_build -T\n sphinx-build -M doctest docs_nnx docs_nnx/_build -T\n # test build html\n sphinx-build -M html docs docs/_build -T\n sphinx-build -M html docs_nnx docs_nnx/_build -T\n # test docstrings\n pytest -n auto flax \\\n --doctest-modules \\\n --suppress-no-test-exit-code \\\n --ignore=flax/nnx/examples\nfi\n\n# check that flax is running on editable mode\n# (i.e. no notebook installed flax from pypi)\necho \"=== CHECKING FLAX IS EDITABLE ===\"\nassert_error=\"flax is not running on editable mode.\"\n(cd docs; python -c \"import flax; assert 'site-packages' not in flax.__file__, \\\"$assert_error\\\"\")\n\n# env vars must be set after doctest\nexport JAX_NUMPY_RANK_PROMOTION=raise\nexport FLAX_PROFILE=1\n\nif $RUN_PYTEST; then\n echo \"=== RUNNING PYTESTS ===\"\n # Run battery of core FLAX API tests.\n echo \"XLA_FLAGS='--xla_force_host_platform_device_count=4' pytest -n auto tests $PYTEST_OPTS\"\n XLA_FLAGS='--xla_force_host_platform_device_count=4' pytest -n auto tests $PYTEST_OPTS\n # Run nnx tests\n pytest -n auto docs/_ext/codediff_test.py $PYTEST_OPTS $PYTEST_IGNORE\n pytest -n auto docs_nnx/_ext/codediff_test.py $PYTEST_OPTS $PYTEST_IGNORE\n\n # Per-example tests.\n #\n # we apply pytest within each example to avoid pytest's annoying test-filename collision.\n # In pytest foo/bar/baz_test.py and baz/bleep/baz_test.py will collide and error out when\n # /foo/bar and /baz/bleep aren't set up as packages.\n for egd in $(find examples -maxdepth 1 -mindepth 1 -type d); do\n # skip if folder starts with \"_\"\n if [[ $egd == *\"_\"* ]]; then\n continue\n fi\n # skiping examples until tfds issue is resolved\n # pytest $egd\n done\nfi\n\nif $RUN_PYTYPE; then\n echo \"=== RUNNING PYTYPE ===\"\n # Validate types in examples.\n for egd in $(find examples -maxdepth 1 -mindepth 1 -type d); do\n # skip if folder starts with \"_\" or is \"nnx_toy_examples\"\n if [[ $egd == *\"_\"* ]] || [[ $egd == *\"nnx_toy_examples\"* ]]; then\n continue\n fi\n # use cd to make sure pytype cache lives in example dir and doesn't name clash\n # use *.py to avoid importing configs as a top-level import which leads to import errors\n # because config files use relative imports (e.g. from config import ...).\n (cd $egd ; pytype \"*.py\" --jobs auto --config ../../pyproject.toml)\n done\n # Validate types in library code.\n pytype --jobs auto --config pyproject.toml flax/\nfi\n\nif $RUN_MYPY; then\n echo \"=== RUNNING MYPY ===\"\n # Validate types in library code.\n mypy --config pyproject.toml flax/ --show-error-codes\nfi\n\n# Return error code 0 if no real failures happened.\necho \"finished all tests.\"\n" }, { "path": "tests/serialization_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for flax.struct and flax.serialization.\"\"\"\n\nimport collections\nimport platform\nfrom typing import Any, NamedTuple\n\nimport jax\nimport jax.numpy as jnp\nimport msgpack\nimport numpy as np\nimport optax\nimport pytest\nfrom absl.testing import absltest, parameterized\nfrom jax import random\nfrom jax.tree_util import Partial\n\nfrom flax import linen as nn\nfrom flax import serialization, struct\nfrom flax.core import freeze\nfrom flax.training import train_state\n\n# Parse absl flags test_srcdir and test_tmpdir.\njax.config.parse_flags_with_absl()\n\n\n@struct.dataclass\nclass Point:\n x: float\n y: float\n meta: Any = struct.field(pytree_node=False)\n\n\n@struct.dataclass\nclass Box:\n value: int\n\n\ndef to_state_dict(box: Box):\n return box.value\n\n\ndef from_state_dict(box: Box, state: Any):\n return box.replace(value=state)\n\n\nserialization.register_serialization_state(\n Box, to_state_dict, from_state_dict, override=True\n)\n\n\nclass OriginalTuple(NamedTuple):\n value: Any\n\n\nclass WrongTuple(NamedTuple):\n wrong_field: Any\n\n\nclass OriginalModule(nn.Module):\n @nn.compact\n def __call__(self, x):\n x = nn.Dense(10)(x)\n return x\n\n\nclass WrongModule(nn.Module):\n @nn.compact\n def __call__(self, x):\n x = nn.Dense(10)(x)\n x = nn.Dense(10)(x)\n return x\n\n\nclass SerializationTest(parameterized.TestCase):\n def test_dataclass_serialization(self):\n p = Point(x=1, y=2, meta={'dummy': True})\n state_dict = serialization.to_state_dict(p)\n self.assertEqual(\n state_dict,\n {\n 'x': 1,\n 'y': 2,\n },\n )\n restored_p = serialization.from_state_dict(p, {'x': 3, 'y': 4})\n expected_p = Point(x=3, y=4, meta={'dummy': True})\n self.assertEqual(restored_p, expected_p)\n\n with self.assertRaises(ValueError): # invalid field\n serialization.from_state_dict(p, {'z': 3})\n with self.assertRaises(ValueError): # missing field\n serialization.from_state_dict(p, {'x': 3})\n\n def test_pass_through_serialization(self):\n p = Box(value=123)\n state_dict = serialization.to_state_dict(p)\n self.assertEqual(\n state_dict,\n 123,\n )\n restored_box = serialization.from_state_dict(p, 123)\n expected_box = Box(value=123)\n self.assertEqual(restored_box, expected_box)\n\n def test_model_serialization(self):\n rng = random.key(0)\n module = nn.Dense(features=1, kernel_init=nn.initializers.ones_init())\n x = jnp.ones((1, 1), jnp.float32)\n initial_params = module.init(rng, x)\n state = serialization.to_state_dict(initial_params)\n self.assertEqual(\n state,\n {\n 'params': {\n 'kernel': np.ones((1, 1)),\n 'bias': np.zeros((1,)),\n }\n },\n )\n state = {\n 'params': {\n 'kernel': np.zeros((1, 1)),\n 'bias': np.zeros((1,)),\n }\n }\n restored_model = serialization.from_state_dict(initial_params, state)\n self.assertEqual(restored_model, freeze(state))\n\n def test_partial_serialization(self):\n add_one = Partial(jnp.add, 1)\n state = serialization.to_state_dict(add_one)\n self.assertEqual(state, {'args': {'0': 1}, 'keywords': {}})\n restored_add_one = serialization.from_state_dict(add_one, state)\n self.assertEqual(add_one.args, restored_add_one.args)\n\n def test_optimizer_serialization(self):\n rng = random.key(0)\n module = nn.Dense(features=1, kernel_init=nn.initializers.ones_init())\n x = jnp.ones((1, 1), jnp.float32)\n initial_params = module.init(rng, x)\n tx = optax.sgd(0.1, momentum=0.1)\n tx_state = tx.init(initial_params)\n state = serialization.to_state_dict(tx_state)\n expected_state = {\n '0': {\n 'trace': {\n 'params': {\n 'bias': np.array([0.0], dtype=jnp.float32),\n 'kernel': np.array([[0.0]], dtype=jnp.float32),\n }\n }\n },\n '1': {},\n }\n self.assertEqual(state, expected_state)\n state = jax.tree_util.tree_map(lambda x: x + 1, expected_state)\n restored_tx_state = serialization.from_state_dict(tx_state, state)\n tx_state_plus1 = jax.tree_util.tree_map(lambda x: x + 1, tx_state)\n self.assertEqual(restored_tx_state, tx_state_plus1)\n\n def test_collection_serialization(self):\n @struct.dataclass\n class DummyDataClass:\n x: float\n\n @classmethod\n def initializer(cls, shape):\n del shape\n return cls(x=0.0)\n\n class StatefulModule(nn.Module):\n @nn.compact\n def __call__(self):\n state = self.variable('state', 'dummy', DummyDataClass.initializer, ())\n state.value = state.value.replace(x=state.value.x + 1.0)\n\n initial_variables = StatefulModule().init(random.key(0))\n _, variables = StatefulModule().apply(initial_variables, mutable=['state'])\n serialized_state_dict = serialization.to_state_dict(variables)\n self.assertEqual(serialized_state_dict, {'state': {'dummy': {'x': 2.0}}})\n deserialized_state = serialization.from_state_dict(\n variables, serialized_state_dict\n )\n self.assertEqual(variables, deserialized_state)\n\n @parameterized.parameters([\n 'byte',\n 'b',\n 'ubyte',\n 'short',\n 'h',\n 'ushort',\n 'i',\n 'uint',\n 'intp',\n 'p',\n 'uintp',\n 'long',\n 'l',\n 'longlong',\n 'q',\n 'ulonglong',\n 'half',\n 'e',\n 'f',\n 'double',\n 'd',\n 'longdouble',\n 'g',\n 'cdouble',\n 'clongdouble',\n 'm',\n 'b1',\n 'int64',\n 'i8',\n 'uint64',\n 'u8',\n 'float16',\n 'f2',\n 'float32',\n 'f4',\n 'float64',\n 'f8',\n 'float128',\n 'f16',\n 'complex64',\n 'c8',\n 'complex128',\n 'c16',\n 'complex256',\n 'c32',\n 'm8',\n 'int32',\n 'i4',\n 'uint32',\n 'u4',\n 'int16',\n 'i2',\n 'uint16',\n 'u2',\n 'int8',\n 'i1',\n 'uint8',\n 'u1',\n 'single',\n 'csingle',\n 'intc',\n 'uintc',\n 'int',\n 'float',\n 'complex',\n 'bool',\n ])\n def test_numpy_serialization(self, dtype):\n np.random.seed(0)\n if (\n (dtype in {'float128', 'f16', 'complex256', 'c32'})\n and (platform.system() == 'Darwin')\n and (platform.machine() == 'arm64')\n ):\n pytest.skip(\n f'Mac M1 does not support dtype {dtype}'\n ) # skip testing these dtypes if user is on Mac M1\n\n v = np.random.uniform(-100, 100, size=()).astype(dtype)[()]\n restored_v = serialization.msgpack_restore(\n serialization.msgpack_serialize(v)\n )\n self.assertEqual(restored_v.dtype, v.dtype)\n np.testing.assert_array_equal(restored_v, v)\n\n for shape in [(), (5,), (10, 10), (1, 20, 30, 1)]:\n arr = np.random.uniform(-100, 100, size=shape).astype(dtype)\n restored_arr = serialization.msgpack_restore(\n serialization.msgpack_serialize(arr)\n )\n self.assertEqual(restored_arr.dtype, arr.dtype)\n np.testing.assert_array_equal(restored_arr, arr)\n\n def test_jax_numpy_serialization(self):\n jax_dtypes = [\n jnp.bool_,\n jnp.uint8,\n jnp.uint16,\n jnp.uint32,\n jnp.int8,\n jnp.int16,\n jnp.int32,\n jnp.bfloat16,\n jnp.float16,\n jnp.float32,\n jnp.complex64,\n ]\n for dtype in jax_dtypes:\n v = jnp.array(np.random.uniform(-100, 100, size=())).astype(dtype)[()]\n restored_v = serialization.msgpack_restore(\n serialization.msgpack_serialize(v)\n )\n self.assertEqual(restored_v.dtype, v.dtype)\n np.testing.assert_array_equal(restored_v, v)\n\n for shape in [(), (5,), (10, 10), (1, 20, 30, 1)]:\n arr = jnp.array(np.random.uniform(-100, 100, size=shape)).astype(dtype)\n restored_arr = serialization.msgpack_restore(\n serialization.msgpack_serialize(arr)\n )\n self.assertEqual(restored_arr.dtype, arr.dtype)\n np.testing.assert_array_equal(restored_arr, arr)\n\n def test_complex_serialization(self):\n for x in [1j, 1 + 2j]:\n restored_x = serialization.msgpack_restore(\n serialization.msgpack_serialize(x)\n )\n self.assertEqual(x, restored_x)\n\n def test_restore_chunked(self):\n old_chunksize = serialization.MAX_CHUNK_SIZE\n serialization.MAX_CHUNK_SIZE = 91 * 8\n try:\n tmp = np.random.uniform(-100, 100, size=(21, 37))\n serialized = serialization.to_bytes(tmp)\n restored = serialization.msgpack_restore(serialized)\n finally:\n serialization.MAX_CHUNK_SIZE = old_chunksize\n\n np.testing.assert_array_equal(restored, tmp)\n\n def test_restore_unchunked(self):\n \"\"\"Check if mgspack_restore works for unchunked inputs.\"\"\"\n\n def msgpack_serialize_legacy(pytree):\n \"\"\"Old implementation that was not chunking.\"\"\"\n return msgpack.packb(\n pytree, default=serialization._msgpack_ext_pack, strict_types=True\n )\n\n tmp = np.random.uniform(-100, 100, size=(21, 37))\n serialized = msgpack_serialize_legacy(tmp)\n old_chunksize = serialization.MAX_CHUNK_SIZE\n serialization.MAX_CHUNK_SIZE = 91 * 8\n try:\n restored = serialization.msgpack_restore(serialized)\n finally:\n serialization.MAX_CHUNK_SIZE = old_chunksize\n\n np.testing.assert_array_equal(restored, tmp)\n\n def test_namedtuple_serialization(self):\n foo_class = collections.namedtuple('Foo', 'a b c')\n x1 = foo_class(a=1, b=2, c=3)\n x1_serialized = serialization.to_bytes(x1)\n x2 = foo_class(a=0, b=0, c=0)\n restored_x1 = serialization.from_bytes(x2, x1_serialized)\n self.assertEqual(type(x1), type(restored_x1))\n self.assertEqual(x1, restored_x1)\n\n def test_namedtuple_restore_legacy(self):\n foo_class = collections.namedtuple('Foo', 'a b c')\n x1 = foo_class(a=1, b=2, c=3)\n legacy_encoding = {\n 'name': 'Foo',\n 'fields': {'0': 'a', '1': 'b', '2': 'c'},\n 'values': {'0': 1, '1': 2, '2': 3},\n }\n x2 = foo_class(a=0, b=0, c=0)\n restored_x1 = serialization.from_state_dict(x2, legacy_encoding)\n self.assertEqual(type(x1), type(restored_x1))\n self.assertEqual(x1, restored_x1)\n\n def test_model_serialization_to_bytes(self):\n rng = random.key(0)\n module = nn.Dense(features=1, kernel_init=nn.initializers.ones_init())\n initial_params = module.init(rng, jnp.ones((1, 1), jnp.float32))\n serialized_bytes = serialization.to_bytes(initial_params)\n restored_params = serialization.from_bytes(initial_params, serialized_bytes)\n self.assertEqual(restored_params, initial_params)\n\n def test_optimizer_serialization_to_bytes(self):\n rng = random.key(0)\n module = nn.Dense(features=1, kernel_init=nn.initializers.ones_init())\n initial_params = module.init(rng, jnp.ones((1, 1), jnp.float32))\n # model = nn.Model(module, initial_params)\n tx = optax.sgd(0.1, momentum=0.1)\n tx_state = tx.init(initial_params)\n serialized_bytes = serialization.to_bytes(tx_state)\n restored_tx_state = serialization.from_bytes(tx_state, serialized_bytes)\n self.assertEqual(restored_tx_state, tx_state)\n\n def test_serialization_chunking(self):\n old_chunksize = serialization.MAX_CHUNK_SIZE\n serialization.MAX_CHUNK_SIZE = 91 * 8\n try:\n tmp = {'a': np.ones((10, 10))}\n tmp = serialization._chunk_array_leaves_in_place(tmp)\n finally:\n serialization.MAX_CHUNK_SIZE = old_chunksize\n test = jax.tree_util.tree_map(jnp.shape, tmp)\n ref = {\n 'a': {\n '__msgpack_chunked_array__': (),\n 'chunks': {'0': (91,), '1': (9,)},\n 'shape': {'0': (), '1': ()},\n }\n }\n self.assertEqual(test, ref)\n\n def test_serialization_chunking2(self):\n old_chunksize = serialization.MAX_CHUNK_SIZE\n serialization.MAX_CHUNK_SIZE = 91 * 8\n try:\n tmp = {'a': np.ones((10, 10))}\n tmpbytes = serialization.to_bytes(tmp)\n newtmp = serialization.from_bytes(tmp, tmpbytes)\n finally:\n serialization.MAX_CHUNK_SIZE = old_chunksize\n jax.tree_util.tree_map(np.testing.assert_array_equal, tmp, newtmp)\n\n def test_serialization_chunking3(self):\n old_chunksize = serialization.MAX_CHUNK_SIZE\n serialization.MAX_CHUNK_SIZE = 91 * 8\n try:\n tmp = {'a': np.ones((10, 10))}\n tmpbytes = serialization.msgpack_serialize(tmp)\n newtmp = serialization.msgpack_restore(tmpbytes)\n finally:\n serialization.MAX_CHUNK_SIZE = old_chunksize\n\n jax.tree_util.tree_map(np.testing.assert_array_equal, tmp, newtmp)\n\n @parameterized.parameters(\n {\n 'target': [[[1, 2, 3], [4, 5]]],\n 'wrong_target': [[[1, 2, 3], [4]]],\n 'msg': (\n 'The size of the list and the state dict do not match,'\n ' got 1 and 2 at path ./0/1'\n ),\n },\n {\n 'target': (((1, 2, 3), (4, 5)),),\n 'wrong_target': (((1, 2, 3), (4,)),),\n 'msg': (\n 'The size of the list and the state dict do not match,'\n ' got 1 and 2 at path ./0/1'\n ),\n },\n {\n 'target': (((1, 2, 3), (OriginalTuple([4, 5]), 6)),),\n 'wrong_target': (((1, 2, 3), (WrongTuple([4, 5]), 6)),),\n 'msg': (\n 'The field names of the state dict and the named tuple do '\n \"not match, got {'value'} and {'wrong_field'} at path ./0/1/0\"\n ),\n },\n {\n 'target': {'a': {'b': {'c': [1, 2, 3], 'd': [4, 5]}}},\n 'wrong_target': {'a': {'b': {'c': [1, 2, 3], 'd': [4]}}},\n 'msg': (\n 'The size of the list and the state dict do not match,'\n ' got 1 and 2 at path ./a/b/d'\n ),\n },\n {\n 'target': {'a': {'b': {'c': [1, 2, 3], 'd': [4, 5]}}},\n 'wrong_target': {'a': {'b': {'c': [1, 2, 3], 'e': [4, 5]}}},\n 'msg': (\n 'The target dict keys and state dict keys do not match, target'\n \" dict contains keys {'e'} which are not present in state dict at\"\n ' path ./a/b'\n ),\n },\n {\n 'target': 'original_params',\n 'wrong_target': 'wrong_params',\n 'msg': (\n 'The target dict keys and state dict keys do not match, target'\n \" dict contains keys {'Dense_1'} which are not present in state\"\n ' dict at path ./params'\n ),\n },\n {\n 'target': 'original_train_state',\n 'wrong_target': 'wrong_train_state',\n 'msg': (\n 'The target dict keys and state dict keys do not match, target'\n \" dict contains keys {'Dense_1'} which are not present in state\"\n ' dict at path ./params/params'\n ),\n },\n )\n def test_serialization_errors(self, target, wrong_target, msg):\n if target == 'original_params':\n x = jnp.ones((1, 28, 28, 1))\n rng = jax.random.key(1)\n original_module = OriginalModule()\n target = original_module.init(rng, x)\n wrong_module = WrongModule()\n wrong_target = wrong_module.init(rng, x)\n\n elif target == 'original_train_state':\n x = jnp.ones((1, 28, 28, 1))\n rng = jax.random.key(1)\n original_module = OriginalModule()\n original_params = original_module.init(rng, x)\n wrong_module = WrongModule()\n wrong_params = wrong_module.init(rng, x)\n\n tx = optax.sgd(learning_rate=0.1, momentum=0.9)\n target = train_state.TrainState.create(\n apply_fn=original_module.apply, params=original_params, tx=tx\n )\n wrong_target = train_state.TrainState.create(\n apply_fn=wrong_module.apply, params=wrong_params, tx=tx\n )\n\n encoded_bytes = serialization.to_bytes(target)\n with self.assertRaisesWithLiteralMatch(ValueError, msg):\n serialization.from_bytes(wrong_target, encoded_bytes)\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/struct_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for flax.struct.\"\"\"\n\nimport dataclasses\nfrom typing import Any, Generic, TypeVar\n\nimport jax\nfrom absl.testing import absltest, parameterized\nfrom jax._src.tree_util import prefix_errors\n\nfrom flax import struct\n\n# Parse absl flags test_srcdir and test_tmpdir.\njax.config.parse_flags_with_absl()\n\n\n@struct.dataclass\nclass Point:\n x: float\n y: float\n meta: Any = struct.field(pytree_node=False)\n\n\nclass StructTest(parameterized.TestCase):\n def test_no_extra_fields(self):\n p = Point(x=1, y=2, meta={})\n with self.assertRaises(dataclasses.FrozenInstanceError):\n p.new_field = 1\n\n def test_mutation(self):\n p = Point(x=1, y=2, meta={})\n new_p = p.replace(x=3)\n self.assertEqual(new_p, Point(x=3, y=2, meta={}))\n with self.assertRaises(dataclasses.FrozenInstanceError):\n p.y = 3\n\n def test_slots(self):\n\n @struct.dataclass(frozen=False, slots=True)\n class SlotsPoint:\n x: float\n y: float\n p = SlotsPoint(x=1., y=2.)\n p.x = 3. # can assign to existing fields\n self.assertEqual(p, SlotsPoint(x=3., y=2.))\n with self.assertRaises(AttributeError):\n p.z = 0. # can't create new fields by accident.\n\n def test_pytree_nodes(self):\n p = Point(x=1, y=2, meta={'abc': True})\n leaves = jax.tree_util.tree_leaves(p)\n self.assertEqual(leaves, [1, 2])\n new_p = jax.tree_util.tree_map(lambda x: x + x, p)\n self.assertEqual(new_p, Point(x=2, y=4, meta={'abc': True}))\n\n def test_keypath_error(self):\n # TODO(mattjj): avoid using internal prefix_errors by testing vmap error msg\n (e,) = prefix_errors(Point(1.0, [2.0], meta={}), Point(1.0, 2.0, meta={}))\n with self.assertRaisesRegex(ValueError, r'in_axes\\.y'):\n raise e('in_axes')\n\n def test_double_wrap_no_op(self):\n class A:\n a: int\n\n self.assertFalse(hasattr(A, '_flax_dataclass'))\n\n A = struct.dataclass(A)\n self.assertTrue(hasattr(A, '_flax_dataclass'))\n\n A = struct.dataclass(A) # no-op\n self.assertTrue(hasattr(A, '_flax_dataclass'))\n\n def test_wrap_pytree_node_no_error(self):\n @struct.dataclass\n class A(struct.PyTreeNode):\n a: int\n\n @parameterized.parameters(\n {'mode': 'dataclass'},\n {'mode': 'pytreenode'},\n )\n def test_kw_only(self, mode):\n if mode == 'dataclass':\n @struct.dataclass\n class A:\n a: int = 1\n\n @struct.dataclass(kw_only=True)\n class B(A):\n b: int\n elif mode == 'pytreenode':\n class A(struct.PyTreeNode):\n a: int = 1\n\n class B(A, struct.PyTreeNode, kw_only=True):\n b: int\n\n obj = B(b=2)\n self.assertEqual(obj.a, 1)\n self.assertEqual(obj.b, 2)\n\n with self.assertRaisesRegex(TypeError, \"non-default argument 'b' follows default argument\"):\n if mode == 'dataclass':\n @struct.dataclass\n class B(A):\n b: int\n elif mode == 'pytreenode':\n class B(A, struct.PyTreeNode):\n b: int\n\n def test_metadata_pass_through(self):\n @struct.dataclass\n class A:\n foo: int = struct.field(default=9, metadata={'baz': 9})\n assert A.__dataclass_fields__['foo'].metadata == {'baz': 9, 'pytree_node': True}\n\n @parameterized.parameters(\n {'mode': 'dataclass'},\n {'mode': 'pytreenode'},\n )\n def test_mutable(self, mode):\n if mode == 'dataclass':\n @struct.dataclass\n class A:\n a: int = 1\n\n @struct.dataclass(frozen=False)\n class B:\n b: int = 1\n elif mode == 'pytreenode':\n class A(struct.PyTreeNode):\n a: int = 1\n\n class B(struct.PyTreeNode, frozen=False):\n b: int = 1\n\n obj = A()\n with self.assertRaisesRegex(dataclasses.FrozenInstanceError, \"cannot assign to field 'a'\"):\n obj.a = 2\n\n obj = B()\n obj.b = 2\n self.assertEqual(obj.b, 2)\n\n\n def test_generic_pytreenode_base_order(self):\n # PyTreeNode + Generic should work regardless of base order (#5233).\n T = TypeVar('T')\n U = TypeVar('U')\n\n # Generic after PyTreeNode.\n class A(struct.PyTreeNode, Generic[T, U]):\n x: int = 0\n\n self.assertEqual(A.__parameters__, (T, U))\n A[int, int] # should not raise\n\n # Generic before PyTreeNode.\n class B(Generic[T, U], struct.PyTreeNode):\n x: int = 0\n\n self.assertEqual(B.__parameters__, (T, U))\n B[int, int] # should not raise\n\n # Subclassing a parameterized generic PyTreeNode.\n class Base(struct.PyTreeNode, Generic[T, U]):\n x: int = 0\n\n class Sub(Base[int, str]):\n y: int = 1\n\n obj = Sub(x=1, y=2)\n self.assertEqual(obj.x, 1)\n self.assertEqual(obj.y, 2)\n leaves = jax.tree_util.tree_leaves(obj)\n self.assertEqual(leaves, [1, 2])\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/tensorboard_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for flax.metrics.tensorboard.\"\"\"\nimport itertools\nimport pathlib\nimport tempfile\n\nimport numpy as np\nimport tensorflow as tf\nfrom absl.testing import absltest\nfrom tensorboard.backend.event_processing import (\n directory_watcher,\n event_file_loader,\n)\nfrom tensorboard.util import tensor_util\n\nfrom flax.metrics.tensorboard import SummaryWriter, _flatten_dict\n\n\ndef _process_event(event):\n for value in event.summary.value:\n yield {'wall_time': event.wall_time, 'step': event.step, 'value': value}\n\n\ndef _disk_usage(path: pathlib.Path):\n \"\"\"Recursively computes the disk usage of a directory.\"\"\"\n if path.is_file():\n return path.stat().st_size\n elif path.is_dir():\n size_bytes = 0\n for file in path.iterdir():\n size_bytes += _disk_usage(file)\n return size_bytes\n else:\n raise NotImplementedError('What filetype is {file}?')\n\n\nclass TensorboardTest(absltest.TestCase):\n def parse_and_return_summary_value(self, path):\n \"\"\"Parse the event file in the given path and return the\n only summary value.\"\"\"\n event_value_list = []\n event_file_generator = directory_watcher.DirectoryWatcher(\n path, event_file_loader.EventFileLoader\n ).Load()\n event_values = itertools.chain.from_iterable(\n map(_process_event, event_file_generator)\n )\n for value_dict in event_values:\n event_value_list.append(value_dict)\n\n self.assertLen(event_value_list, 1)\n self.assertEqual(event_value_list[0]['step'], 1)\n self.assertGreater(event_value_list[0]['wall_time'], 0.0)\n return event_value_list[0]['value']\n\n def test_summarywriter_flush_after_close(self):\n log_dir = tempfile.mkdtemp()\n summary_writer = SummaryWriter(log_dir=log_dir)\n summary_writer.close()\n with self.assertRaises(AttributeError):\n summary_writer.flush()\n\n def test_summarywriter_scalar(self):\n log_dir = tempfile.mkdtemp()\n summary_writer = SummaryWriter(log_dir=log_dir)\n # Write the scalar and check if the event exists and check data.\n float_value = 99.1232\n summary_writer.scalar(tag='scalar_test', value=float_value, step=1)\n\n summary_value = self.parse_and_return_summary_value(path=log_dir)\n self.assertEqual(summary_value.tag, 'scalar_test')\n self.assertTrue(\n np.allclose(\n tensor_util.make_ndarray(summary_value.tensor).item(), float_value\n )\n )\n\n def test_summarywriter_text(self):\n log_dir = tempfile.mkdtemp()\n summary_writer = SummaryWriter(log_dir=log_dir)\n text = 'hello world.'\n summary_writer.text(tag='text_test', textdata=text, step=1)\n\n summary_value = self.parse_and_return_summary_value(path=log_dir)\n self.assertEqual(summary_value.tag, 'text_test')\n self.assertEqual(\n tensor_util.make_ndarray(summary_value.tensor).item().decode('utf-8'),\n text,\n )\n\n def test_summarywriter_image(self):\n log_dir = tempfile.mkdtemp()\n summary_writer = SummaryWriter(log_dir=log_dir)\n expected_img = np.random.uniform(low=0.0, high=255.0, size=(30, 30, 3))\n expected_img = expected_img.astype(np.uint8)\n summary_writer.image(tag='image_test', image=expected_img, step=1)\n summary_value = self.parse_and_return_summary_value(path=log_dir)\n\n self.assertEqual(summary_value.tag, 'image_test')\n actual_img = tf.image.decode_image(summary_value.tensor.string_val[2])\n self.assertTrue(np.allclose(actual_img, expected_img))\n\n def test_summarywriter_image_float_pixel_values(self):\n log_dir = tempfile.mkdtemp()\n summary_writer = SummaryWriter(log_dir=log_dir)\n expected_img = np.random.uniform(low=0.0, high=1.0, size=(30, 30, 3))\n summary_writer.image(tag='image_test', image=expected_img, step=1)\n summary_value = self.parse_and_return_summary_value(path=log_dir)\n\n # convert and scale expected_img appropriately to numpy uint8.\n expected_img = tf.image.convert_image_dtype(\n image=expected_img, dtype=np.uint8\n )\n\n self.assertEqual(summary_value.tag, 'image_test')\n actual_img = tf.image.decode_image(summary_value.tensor.string_val[2])\n self.assertTrue(np.allclose(actual_img, expected_img))\n\n def test_summarywriter_2dimage_scaled(self):\n log_dir = tempfile.mkdtemp()\n summary_writer = SummaryWriter(log_dir=log_dir)\n img = np.random.uniform(low=0.0, high=255.0, size=(30, 30))\n img = img.astype(np.uint8)\n summary_writer.image(tag='2dimage_test', image=img, step=1)\n summary_value = self.parse_and_return_summary_value(path=log_dir)\n\n self.assertEqual(summary_value.tag, '2dimage_test')\n actual_img = tf.image.decode_image(summary_value.tensor.string_val[2])\n # assert the image was increased in dimension\n self.assertEqual(actual_img.shape, (30, 30, 3))\n\n def test_summarywriter_single_channel_image_scaled(self):\n log_dir = tempfile.mkdtemp()\n summary_writer = SummaryWriter(log_dir=log_dir)\n img = np.random.uniform(low=0.0, high=255.0, size=(30, 30, 1))\n img = img.astype(np.uint8)\n summary_writer.image(tag='2dimage_1channel_test', image=img, step=1)\n summary_value = self.parse_and_return_summary_value(path=log_dir)\n\n self.assertEqual(summary_value.tag, '2dimage_1channel_test')\n actual_img = tf.image.decode_image(summary_value.tensor.string_val[2])\n # assert the image was increased in dimension\n self.assertEqual(actual_img.shape, (30, 30, 3))\n\n def test_summarywriter_multiple_images(self):\n log_dir = tempfile.mkdtemp()\n summary_writer = SummaryWriter(log_dir=log_dir)\n expected_img = np.random.uniform(low=0.0, high=255.0, size=(2, 30, 30, 3))\n expected_img = expected_img.astype(np.uint8)\n summary_writer.image(tag='multiple_images_test', image=expected_img, step=1)\n summary_value = self.parse_and_return_summary_value(path=log_dir)\n\n self.assertEqual(summary_value.tag, 'multiple_images_test')\n actual_imgs = [\n tf.image.decode_image(s) for s in summary_value.tensor.string_val[2:]\n ]\n self.assertTrue(np.allclose(np.stack(actual_imgs, axis=0), expected_img))\n\n def test_summarywriter_multiple_2dimages_scaled(self):\n log_dir = tempfile.mkdtemp()\n summary_writer = SummaryWriter(log_dir=log_dir)\n img = np.random.uniform(low=0.0, high=255.0, size=(2, 30, 30))\n img = img.astype(np.uint8)\n summary_writer.image(tag='multiple_2dimages_test', image=img, step=1)\n summary_value = self.parse_and_return_summary_value(path=log_dir)\n\n self.assertEqual(summary_value.tag, 'multiple_2dimages_test')\n actual_imgs = [\n tf.image.decode_image(s) for s in summary_value.tensor.string_val[2:]\n ]\n # assert the images were increased in dimension\n self.assertEqual(np.stack(actual_imgs, axis=0).shape, (2, 30, 30, 3))\n\n def test_summarywriter_audio(self):\n log_dir = tempfile.mkdtemp()\n summary_writer = SummaryWriter(log_dir=log_dir)\n expected_audio_samples = np.random.uniform(\n low=-1.0, high=1.0, size=(2, 48000, 2)\n )\n summary_writer.audio(\n tag='audio_test', audiodata=expected_audio_samples, step=1\n )\n\n summary_value = self.parse_and_return_summary_value(path=log_dir)\n self.assertEqual(summary_value.tag, 'audio_test')\n\n # Assert two audio files are parsed.\n self.assertLen(summary_value.tensor.string_val, 2)\n\n # Assert values.\n actual_audio_1 = tf.audio.decode_wav(\n summary_value.tensor.string_val[0]\n ).audio\n self.assertTrue(\n np.allclose(expected_audio_samples[0], actual_audio_1, atol=1e-04)\n )\n\n actual_audio_2 = tf.audio.decode_wav(\n summary_value.tensor.string_val[1]\n ).audio\n self.assertTrue(\n np.allclose(expected_audio_samples[1], actual_audio_2, atol=1e-04)\n )\n\n def test_summarywriter_audio_sampled_output(self):\n log_dir = tempfile.mkdtemp()\n summary_writer = SummaryWriter(log_dir=log_dir)\n expected_audio_samples = np.random.uniform(\n low=-1.0, high=1.0, size=(2, 48000, 2)\n )\n summary_writer.audio(\n tag='audio_test',\n audiodata=expected_audio_samples,\n step=1,\n max_outputs=1,\n )\n\n summary_value = self.parse_and_return_summary_value(path=log_dir)\n self.assertEqual(summary_value.tag, 'audio_test')\n\n # Assert only the first audio clip is available.\n self.assertLen(summary_value.tensor.string_val, 1)\n\n # Assert values.\n actual_audio = tf.audio.decode_wav(summary_value.tensor.string_val[0]).audio\n\n self.assertTrue(\n np.allclose(expected_audio_samples[0], actual_audio, atol=1e-04)\n )\n\n def test_summarywriter_clipped_audio(self):\n log_dir = tempfile.mkdtemp()\n summary_writer = SummaryWriter(log_dir=log_dir)\n expected_audio_samples = np.random.uniform(\n low=-2.0, high=2.0, size=(2, 48000, 2)\n )\n summary_writer.audio(\n tag='audio_test',\n audiodata=expected_audio_samples,\n step=1,\n max_outputs=1,\n )\n\n summary_value = self.parse_and_return_summary_value(path=log_dir)\n self.assertEqual(summary_value.tag, 'audio_test')\n\n # Assert one audio files is parsed.\n self.assertLen(summary_value.tensor.string_val, 1)\n\n # actual_audio is clipped.\n actual_audio = tf.audio.decode_wav(summary_value.tensor.string_val[0]).audio\n self.assertFalse(\n np.allclose(expected_audio_samples[0], actual_audio, atol=1e-04)\n )\n\n clipped_audio = np.clip(np.array(expected_audio_samples[0]), -1, 1)\n self.assertTrue(np.allclose(clipped_audio, actual_audio, atol=1e-04))\n\n def test_summarywriter_histogram_defaultbins(self):\n log_dir = tempfile.mkdtemp()\n summary_writer = SummaryWriter(log_dir=log_dir)\n histogram = np.arange(1000)\n # Histogram will be created for 30 (default) bins.\n summary_writer.histogram(tag='histogram_test', values=histogram, step=1)\n\n summary_value = self.parse_and_return_summary_value(path=log_dir)\n self.assertEqual(summary_value.tag, 'histogram_test')\n actual_histogram = tensor_util.make_ndarray(summary_value.tensor)\n self.assertTrue(actual_histogram.shape, (30, 3))\n self.assertTrue(\n np.allclose(actual_histogram[0], (0.0, 33.3, 34.0), atol=1e-01)\n )\n\n def test_summarywriter_histogram_2bins(self):\n log_dir = tempfile.mkdtemp()\n summary_writer = SummaryWriter(log_dir=log_dir)\n histogram = np.arange(1000)\n summary_writer.histogram(\n tag='histogram_test', values=histogram, step=1, bins=2\n )\n\n summary_value = self.parse_and_return_summary_value(path=log_dir)\n self.assertEqual(summary_value.tag, 'histogram_test')\n actual_histogram = tensor_util.make_ndarray(summary_value.tensor)\n self.assertTrue(actual_histogram.shape, (2, 3))\n self.assertTrue(\n np.allclose(actual_histogram[0], (0.0, 499.5, 500.0), atol=1e-01)\n )\n self.assertTrue(\n np.allclose(actual_histogram[1], (499.5, 999.0, 500.0), atol=1e-01)\n )\n\n def test_flatten_dict(self):\n # Valid types according to https://github.com/tensorflow/tensorboard/blob/1204566da5437af55109f7a4af18f9f8b7c4f864/tensorboard/plugins/hparams/summary_v2.py\n input_hparams = {\n # Example Invalid Types\n 'None': None,\n 'List': [1, 2, 3],\n 'Tuple': (1, 2, 3),\n 'Complex': complex('1+1j'),\n 'np.complex_': np.complex128('1+1j'),\n # Valid Python Types\n 'Bool': True,\n 'Int': 1,\n 'Float': 1.0,\n 'Str': 'test',\n # Valid Numpy Types\n 'np.bool_': np.bool_(1),\n 'np.integer': np.int_(1),\n 'np.floating': np.float64(1.0),\n 'np.character': np.str_('test'),\n # Nested dict to flatten\n 'Nested_Dict': {\n 'None': None,\n 'List': [1, 2, 3],\n 'Tuple': (1, 2, 3),\n 'Complex': complex('1+1j'),\n 'np.complex_': np.complex128('1+1j'),\n 'Bool': True,\n 'Int': 1,\n 'Float': 1.0,\n 'Str': 'test',\n 'np.bool_': np.bool_(1),\n 'np.integer': np.int_(1),\n 'np.floating': np.float64(1.0),\n 'np.character': np.str_('test'),\n },\n }\n\n result_hparams = _flatten_dict(input_hparams)\n\n expected_hparams = {\n 'None': 'None',\n 'List': '[1, 2, 3]',\n 'Tuple': '(1, 2, 3)',\n 'Complex': '(1+1j)',\n 'np.complex_': '(1+1j)',\n # Valid Python Types\n 'Bool': True,\n 'Int': 1,\n 'Float': 1.0,\n 'Str': 'test',\n # Valid Numpy Types\n 'np.bool_': np.bool_(1),\n 'np.integer': np.int_(1),\n 'np.floating': np.float64(1.0),\n 'np.character': np.str_('test'),\n # Nested Dict\n 'Nested_Dict.None': 'None',\n 'Nested_Dict.List': '[1, 2, 3]',\n 'Nested_Dict.Tuple': '(1, 2, 3)',\n 'Nested_Dict.Complex': '(1+1j)',\n 'Nested_Dict.np.complex_': '(1+1j)',\n 'Nested_Dict.Bool': True,\n 'Nested_Dict.Int': 1,\n 'Nested_Dict.Float': 1.0,\n 'Nested_Dict.Str': 'test',\n 'Nested_Dict.np.bool_': np.bool_(1),\n 'Nested_Dict.np.integer': np.int_(1),\n 'Nested_Dict.np.floating': np.float64(1.0),\n 'Nested_Dict.np.character': np.str_('test'),\n }\n\n self.assertDictEqual(result_hparams, expected_hparams)\n\n def test_auto_flush(self):\n tmp_dir = pathlib.Path(self.create_tempdir().full_path)\n summary_writer = SummaryWriter(tmp_dir, auto_flush=True)\n summary_writer.scalar('metric', 123, 1)\n filesize_before_flush = _disk_usage(tmp_dir)\n summary_writer.flush()\n filesize_after_flush = _disk_usage(tmp_dir)\n self.assertEqual(filesize_before_flush, filesize_after_flush)\n\n def test_no_auto_flush(self):\n tmp_dir = pathlib.Path(self.create_tempdir().full_path)\n summary_writer = SummaryWriter(tmp_dir, auto_flush=False)\n summary_writer.scalar('metric', 123, 1)\n filesize_before_flush = _disk_usage(tmp_dir)\n summary_writer.flush()\n filesize_after_flush = _disk_usage(tmp_dir)\n self.assertLess(filesize_before_flush, filesize_after_flush)\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/traceback_util_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for flax.traceback_util.\"\"\"\n\nimport contextlib\nimport sys\nimport traceback\n\nimport jax\nfrom absl.testing import absltest\nfrom jax import numpy as jnp\nfrom jax import random\nfrom jax._src import traceback_util as jax_traceback_util\n\nfrom flax import linen as nn\nfrom flax import traceback_util\n\n# pylint: disable=arguments-differ,protected-access, g-wrong-blank-lines\n\n# __tracebackhide__ is a python >=3.7 feature.\nTRACEBACKHIDE_SUPPORTED = tuple(sys.version_info)[:3] >= (3, 7, 0)\n\nEXPECTED_FILES = (__file__, contextlib.__spec__.origin)\n\n\nclass TracebackTest(absltest.TestCase):\n def test_exclusion_list(self):\n traceback_util.show_flax_in_tracebacks()\n exclusion_len_wo_flax = len(jax_traceback_util._exclude_paths)\n traceback_util.hide_flax_in_tracebacks()\n exclusion_len_w_flax = len(jax_traceback_util._exclude_paths)\n self.assertLen(\n traceback_util._flax_exclusions,\n exclusion_len_w_flax - exclusion_len_wo_flax,\n )\n\n def test_simple_exclusion_tracebackhide(self):\n if not TRACEBACKHIDE_SUPPORTED:\n return\n\n class Test1(nn.Module):\n @nn.remat\n @nn.compact\n def __call__(self, x):\n return Test2()(x)\n\n class Test2(nn.Module):\n @nn.jit\n @nn.compact\n def __call__(self, x):\n raise ValueError('error here.')\n return x # pylint: disable=unreachable\n\n traceback_util.hide_flax_in_tracebacks()\n jax.config.update('jax_traceback_filtering', 'tracebackhide')\n\n key = random.key(0)\n try:\n nn.jit(Test1)().init(key, jnp.ones((5, 3)))\n except ValueError as e:\n tb = e.__traceback__\n\n filtered_frames = 0\n unfiltered_frames = 0\n\n for f, _ in traceback.walk_tb(tb):\n if '__tracebackhide__' not in f.f_locals:\n self.assertIn(f.f_code.co_filename, EXPECTED_FILES)\n filtered_frames += 1\n unfiltered_frames += 1\n\n self.assertEqual(filtered_frames, 3)\n self.assertGreater(unfiltered_frames, filtered_frames)\n\n def test_simple_exclusion_remove_frames(self):\n class Test1(nn.Module):\n @nn.remat\n @nn.compact\n def __call__(self, x):\n return Test2()(x)\n\n class Test2(nn.Module):\n @nn.jit\n @nn.compact\n def __call__(self, x):\n raise ValueError('error here.')\n return x # pylint: disable=unreachable\n\n traceback_util.hide_flax_in_tracebacks()\n jax.config.update('jax_traceback_filtering', 'remove_frames')\n\n key = random.key(0)\n try:\n nn.jit(Test1)().init(key, jnp.ones((5, 3)))\n except ValueError as e:\n tb_filtered = e.__traceback__\n tb_unfiltered = e.__cause__.__traceback__\n e_cause = e.__cause__\n\n self.assertIsInstance(e_cause, jax_traceback_util.UnfilteredStackTrace)\n\n filtered_frames = 0\n for _, _ in traceback.walk_tb(tb_filtered):\n filtered_frames += 1\n\n unfiltered_frames = 0\n for _, _ in traceback.walk_tb(tb_unfiltered):\n unfiltered_frames += 1\n\n self.assertEqual(filtered_frames, 3)\n self.assertGreater(unfiltered_frames, filtered_frames)\n\n def test_dynamic_exclusion(self):\n if not TRACEBACKHIDE_SUPPORTED:\n return\n\n class Test1(nn.Module):\n @nn.remat\n @nn.compact\n def __call__(self, x):\n return Test2()(x)\n\n class Test2(nn.Module):\n @nn.jit\n @nn.compact\n def __call__(self, x):\n raise ValueError('error here.')\n return x # pylint: disable=unreachable\n\n key = random.key(0)\n\n traceback_util.show_flax_in_tracebacks()\n jax.config.update('jax_traceback_filtering', 'off')\n try:\n nn.jit(Test1)().init(key, jnp.ones((5, 3)))\n except ValueError as e:\n tb_all = e.__traceback__\n\n traceback_util.hide_flax_in_tracebacks()\n jax.config.update('jax_traceback_filtering', 'tracebackhide')\n try:\n nn.jit(Test1)().init(key, jnp.ones((5, 3)))\n except ValueError as e:\n tb_no_flax = e.__traceback__\n\n traceback_util.show_flax_in_tracebacks()\n jax.config.update('jax_traceback_filtering', 'tracebackhide')\n try:\n nn.jit(Test1)().init(key, jnp.ones((5, 3)))\n except ValueError as e:\n tb_w_flax = e.__traceback__\n\n filtered_frames_all = 0\n unfiltered_frames_all = 0\n for f, _ in traceback.walk_tb(tb_all):\n if '__tracebackhide__' not in f.f_locals:\n unfiltered_frames_all += 1\n else:\n filtered_frames_all += 1\n\n filtered_frames_no_flax = 0\n unfiltered_frames_no_flax = 0\n for f, _ in traceback.walk_tb(tb_no_flax):\n if '__tracebackhide__' not in f.f_locals:\n self.assertIn(f.f_code.co_filename, EXPECTED_FILES)\n unfiltered_frames_no_flax += 1\n else:\n filtered_frames_no_flax += 1\n\n filtered_frames_w_flax = 0\n unfiltered_frames_w_flax = 0\n for f, _ in traceback.walk_tb(tb_w_flax):\n if '__tracebackhide__' not in f.f_locals:\n unfiltered_frames_w_flax += 1\n else:\n filtered_frames_w_flax += 1\n\n self.assertEqual(\n unfiltered_frames_all + filtered_frames_all,\n unfiltered_frames_w_flax + filtered_frames_w_flax,\n )\n self.assertEqual(\n unfiltered_frames_all + filtered_frames_all,\n unfiltered_frames_no_flax + filtered_frames_no_flax,\n )\n self.assertEqual(unfiltered_frames_no_flax, 3)\n self.assertGreater(unfiltered_frames_all, unfiltered_frames_w_flax)\n self.assertGreater(unfiltered_frames_w_flax, unfiltered_frames_no_flax)\n\n\nif __name__ == '__main__':\n absltest.main()\n" }, { "path": "tests/traverse_util_test.py", "content": "# Copyright 2024 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for flax.traverse_util.\"\"\"\n\n\nimport collections\n\nimport jax\nimport jax.numpy as jnp\nimport numpy as np\nimport optax\nfrom absl.testing import absltest\n\nimport flax\nfrom flax import traverse_util\nfrom flax.core import freeze\n\n# Parse absl flags test_srcdir and test_tmpdir.\njax.config.parse_flags_with_absl()\n\n\nclass Foo:\n def __init__(self, foo, bar=None):\n self.foo = foo\n self.bar = bar\n\n def __eq__(self, other):\n return self.foo == other.foo and self.bar == other.bar\n\n\nPoint = collections.namedtuple('Point', ['x', 'y'])\n\n\nclass TraversalTest(absltest.TestCase):\n def test_traversal_id(self):\n x = 1\n traversal = traverse_util.t_identity\n self.assertEqual(list(traversal.iterate(x)), [1])\n y = traversal.update(lambda x: x + x, x)\n self.assertEqual(y, 2)\n\n def test_traverse_item(self):\n x = {'foo': 1}\n traversal = traverse_util.t_identity['foo']\n self.assertEqual(list(traversal.iterate(x)), [1])\n y = traversal.update(lambda x: x + x, x)\n self.assertEqual(y, {'foo': 2})\n\n def test_traverse_tuple_item(self):\n x = (1, 2, 3)\n traversal = traverse_util.t_identity[1]\n self.assertEqual(list(traversal.iterate(x)), [2])\n y = traversal.update(lambda x: x + x, x)\n self.assertEqual(y, (1, 4, 3))\n\n def test_traverse_tuple_items(self):\n x = (1, 2, 3, 4)\n traversal = traverse_util.t_identity[1:3]\n self.assertEqual(list(traversal.iterate(x)), [2, 3])\n y = traversal.update(lambda x: x + x, x)\n self.assertEqual(y, (1, 4, 6, 4))\n\n def test_traverse_namedtuple_item(self):\n x = Point(x=1, y=2)\n traversal = traverse_util.t_identity[1]\n self.assertEqual(list(traversal.iterate(x)), [2])\n y = traversal.update(lambda x: x + x, x)\n self.assertEqual(y, Point(x=1, y=4))\n\n def test_traverse_attr(self):\n x = Foo(foo=1)\n traversal = traverse_util.t_identity.foo\n self.assertEqual(list(traversal.iterate(x)), [1])\n y = traversal.update(lambda x: x + x, x)\n self.assertEqual(y, Foo(foo=2))\n\n def test_traverse_namedtuple_attr(self):\n x = Point(x=1, y=2)\n traversal = traverse_util.t_identity.y\n self.assertEqual(list(traversal.iterate(x)), [2])\n y = traversal.update(lambda x: x + x, x)\n self.assertEqual(y, Point(x=1, y=4))\n\n def test_traverse_dataclass_attr(self):\n x = Point(x=1, y=2)\n traversal = traverse_util.t_identity.y\n self.assertEqual(list(traversal.iterate(x)), [2])\n y = traversal.update(lambda x: x + x, x)\n self.assertEqual(y, Point(x=1, y=4))\n\n def test_traverse_merge(self):\n x = [{'foo': 1, 'bar': 2}, {'foo': 3, 'bar': 4}]\n traversal_base = traverse_util.t_identity.each()\n traversal = traversal_base.merge(\n traverse_util.TraverseItem('foo'), traverse_util.TraverseItem('bar')\n )\n self.assertEqual(list(traversal.iterate(x)), [1, 2, 3, 4])\n y = traversal.update(lambda x: x + x, x)\n self.assertEqual(y, [{'foo': 2, 'bar': 4}, {'foo': 6, 'bar': 8}])\n\n def test_traverse_each(self):\n x = [{'foo': 1}, {'foo': 2}]\n traversal = traverse_util.t_identity.each()['foo']\n self.assertEqual(list(traversal.iterate(x)), [1, 2])\n y = traversal.update(lambda x: x + x, x)\n self.assertEqual(y, [{'foo': 2}, {'foo': 4}])\n\n def test_traverse_each_dict(self):\n x = {'foo': 1, 'bar': 2}\n traversal = traverse_util.t_identity.each()\n self.assertEqual(list(traversal.iterate(x)), [1, 2])\n y = traversal.update(lambda x: x + x, x)\n self.assertEqual(y, {'foo': 2, 'bar': 4})\n\n def test_traverse_tree(self):\n x = [{'foo': 1}, {'bar': 2}]\n traversal = traverse_util.t_identity.tree()\n self.assertEqual(list(traversal.iterate(x)), [1, 2])\n y = traversal.update(lambda x: x + x, x)\n self.assertEqual(y, [{'foo': 2}, {'bar': 4}])\n\n def test_traverse_filter(self):\n x = [1, -2, 3, -4]\n traversal = traverse_util.t_identity.each().filter(lambda x: x < 0)\n self.assertEqual(list(traversal.iterate(x)), [-2, -4])\n y = traversal.update(lambda x: -x, x)\n self.assertEqual(y, [1, 2, 3, 4])\n\n def test_traversal_set(self):\n x = {'foo': [1, 2]}\n traversal = traverse_util.t_identity['foo'].each()\n y = traversal.set([3, 4], x)\n self.assertEqual(y, {'foo': [3, 4]})\n with self.assertRaises(ValueError):\n traversal.set([3], x) # too few values\n with self.assertRaises(ValueError):\n traversal.set([3, 4, 5], x) # too many values\n\n def test_flatten_dict(self):\n xs = {'foo': 1, 'bar': {'a': 2, 'b': {}}}\n flat_xs = traverse_util.flatten_dict(xs)\n self.assertEqual(\n flat_xs,\n {\n ('foo',): 1,\n ('bar', 'a'): 2,\n },\n )\n flat_xs = traverse_util.flatten_dict(freeze(xs))\n self.assertEqual(\n flat_xs,\n {\n ('foo',): 1,\n ('bar', 'a'): 2,\n },\n )\n flat_xs = traverse_util.flatten_dict(xs, sep='/')\n self.assertEqual(\n flat_xs,\n {\n 'foo': 1,\n 'bar/a': 2,\n },\n )\n\n def test_unflatten_dict(self):\n expected_xs = {'foo': 1, 'bar': {'a': 2}}\n xs = traverse_util.unflatten_dict(\n {\n ('foo',): 1,\n ('bar', 'a'): 2,\n }\n )\n self.assertEqual(xs, expected_xs)\n xs = traverse_util.unflatten_dict(\n {\n 'foo': 1,\n 'bar/a': 2,\n },\n sep='/',\n )\n self.assertEqual(xs, expected_xs)\n\n def test_flatten_dict_keep_empty(self):\n xs = {'foo': 1, 'bar': {'a': 2, 'b': {}}}\n flat_xs = traverse_util.flatten_dict(xs, keep_empty_nodes=True)\n self.assertEqual(\n flat_xs,\n {\n ('foo',): 1,\n ('bar', 'a'): 2,\n ('bar', 'b'): traverse_util.empty_node,\n },\n )\n xs_restore = traverse_util.unflatten_dict(flat_xs)\n self.assertEqual(xs, xs_restore)\n\n def test_flatten_dict_is_leaf(self):\n xs = {'foo': {'c': 4}, 'bar': {'a': 2, 'b': {}}}\n flat_xs = traverse_util.flatten_dict(\n xs, is_leaf=lambda k, x: len(k) == 1 and len(x) == 2\n )\n self.assertEqual(\n flat_xs,\n {\n ('foo', 'c'): 4,\n ('bar',): {'a': 2, 'b': {}},\n },\n )\n xs_restore = traverse_util.unflatten_dict(flat_xs)\n self.assertEqual(xs, xs_restore)\n\n\nclass ModelParamTraversalTest(absltest.TestCase):\n def test_only_works_on_model_params(self):\n traversal = traverse_util.ModelParamTraversal(lambda *_: True)\n with self.assertRaises(ValueError):\n list(traversal.iterate([]))\n\n def test_param_selection(self):\n params = {\n 'x': {\n 'kernel': 1,\n 'bias': 2,\n 'y': {\n 'kernel': 3,\n 'bias': 4,\n },\n 'z': {},\n },\n }\n expected_params = {\n 'x': {\n 'kernel': 2,\n 'bias': 2,\n 'y': {\n 'kernel': 6,\n 'bias': 4,\n },\n 'z': {},\n },\n }\n names = []\n\n def filter_fn(name, _):\n names.append(name) # track names passed to filter_fn for testing\n return 'kernel' in name\n\n traversal = traverse_util.ModelParamTraversal(filter_fn)\n\n values = list(traversal.iterate(params))\n configs = [\n (params, expected_params),\n (flax.core.FrozenDict(params), flax.core.FrozenDict(expected_params)),\n ]\n for model, expected_model in configs:\n self.assertEqual(values, [1, 3])\n self.assertEqual(\n set(names), {'/x/kernel', '/x/bias', '/x/y/kernel', '/x/y/bias'}\n )\n new_model = traversal.update(lambda x: x + x, model)\n self.assertEqual(new_model, expected_model)\n\n def test_path_value(self):\n params_in = {'a': {'b': 10, 'c': 2}}\n params_out = traverse_util.path_aware_map(\n lambda path, x: x + 1 if 'b' in path else -x, params_in\n )\n\n self.assertEqual(params_out, {'a': {'b': 11, 'c': -2}})\n\n def test_path_aware_map_with_multi_transform(self):\n params = {\n 'linear_1': {'w': jnp.zeros((5, 6)), 'b': jnp.zeros(5)},\n 'linear_2': {'w': jnp.zeros((6, 1)), 'b': jnp.zeros(1)},\n }\n gradients = jax.tree_util.tree_map(jnp.ones_like, params) # dummy gradients\n\n param_labels = traverse_util.path_aware_map(\n lambda path, x: 'kernel' if 'w' in path else 'bias', params\n )\n tx = optax.multi_transform(\n {'kernel': optax.sgd(1.0), 'bias': optax.set_to_zero()}, param_labels\n )\n state = tx.init(params)\n updates, new_state = tx.update(gradients, state, params)\n new_params = optax.apply_updates(params, updates)\n\n self.assertTrue(\n np.allclose(new_params['linear_1']['b'], params['linear_1']['b'])\n )\n self.assertTrue(\n np.allclose(new_params['linear_2']['b'], params['linear_2']['b'])\n )\n self.assertFalse(\n np.allclose(new_params['linear_1']['w'], params['linear_1']['w'])\n )\n self.assertFalse(\n np.allclose(new_params['linear_2']['w'], params['linear_2']['w'])\n )\n\n def test_path_aware_map_with_masked(self):\n params = {\n 'linear_1': {'w': jnp.zeros((5, 6)), 'b': jnp.zeros(5)},\n 'linear_2': {'w': jnp.zeros((6, 1)), 'b': jnp.zeros(1)},\n }\n gradients = jax.tree_util.tree_map(jnp.ones_like, params) # dummy gradients\n\n params_mask = traverse_util.path_aware_map(\n lambda path, x: 'w' in path, params\n )\n tx = optax.masked(optax.sgd(1.0), params_mask)\n state = tx.init(params)\n updates, new_state = tx.update(gradients, state, params)\n new_params = optax.apply_updates(params, updates)\n\n self.assertTrue(\n np.allclose(new_params['linear_1']['b'], gradients['linear_1']['b'])\n )\n self.assertTrue(\n np.allclose(new_params['linear_2']['b'], gradients['linear_2']['b'])\n )\n self.assertTrue(\n np.allclose(new_params['linear_1']['w'], -gradients['linear_1']['w'])\n )\n self.assertTrue(\n np.allclose(new_params['linear_2']['w'], -gradients['linear_2']['w'])\n )\n\n def test_path_aware_map_with_empty_nodes(self):\n params_in = {'a': {'b': 10, 'c': 2}, 'b': {}}\n params_out = traverse_util.path_aware_map(\n lambda path, x: x + 1 if 'b' in path else -x, params_in\n )\n\n self.assertEqual(params_out, {'a': {'b': 11, 'c': -2}, 'b': {}})\n\n\nif __name__ == '__main__':\n absltest.main()\n" } ]